feat: prepare ARR for controlled public deployment
This commit is contained in:
Wyndham ARR
335
agent_integration/client.py
Normal file
335
agent_integration/client.py
Normal file
@@ -0,0 +1,335 @@
|
||||
"""Synchronous HTTP client for the DeerFlow Agent Profile Open API."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import secrets
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Dict, Iterator, List, Literal, Optional
|
||||
from urllib.parse import quote
|
||||
|
||||
import httpx
|
||||
|
||||
|
||||
AuthMode = Literal["bearer", "x-api-key"]
|
||||
|
||||
|
||||
class OpenAgentAPIError(RuntimeError):
|
||||
"""Raised when the API returns a non-success HTTP response."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
status_code: int,
|
||||
detail: Any,
|
||||
response_text: str,
|
||||
*,
|
||||
request_id: Optional[str] = None,
|
||||
) -> None:
|
||||
self.status_code = status_code
|
||||
self.detail = detail
|
||||
self.response_text = response_text
|
||||
self.request_id = request_id
|
||||
suffix = f" (request_id={request_id})" if request_id else ""
|
||||
super().__init__(f"Open Agent API request failed with HTTP {status_code}: {detail}{suffix}")
|
||||
|
||||
@property
|
||||
def retryable(self) -> bool:
|
||||
return self.status_code in {408, 429} or self.status_code >= 500
|
||||
|
||||
@property
|
||||
def active_run_conflict(self) -> bool:
|
||||
return self.status_code == 409
|
||||
|
||||
|
||||
class OpenAgentTransportError(RuntimeError):
|
||||
"""Raised when the HTTP exchange fails before a valid API response is received."""
|
||||
|
||||
|
||||
class OpenAgentProtocolError(RuntimeError):
|
||||
"""Raised when a successful API response does not match the documented shape."""
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class OpenAgentEvent:
|
||||
event: Optional[str]
|
||||
data: Any
|
||||
event_id: Optional[str] = None
|
||||
retry_ms: Optional[int] = None
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
value: Dict[str, Any] = {"event": self.event, "data": self.data}
|
||||
if self.event_id is not None:
|
||||
value["id"] = self.event_id
|
||||
if self.retry_ms is not None:
|
||||
value["retry"] = self.retry_ms
|
||||
return value
|
||||
|
||||
|
||||
class OpenAgentAPIClient:
|
||||
"""Small synchronous client matching the supplied Agent Open API contract."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
base_url: str,
|
||||
api_key: str,
|
||||
auth_mode: AuthMode = "bearer",
|
||||
csrf_token: Optional[str] = None,
|
||||
timeout: float = 60.0,
|
||||
connect_timeout: float = 10.0,
|
||||
http_client: Optional[httpx.Client] = None,
|
||||
) -> None:
|
||||
if auth_mode not in ("bearer", "x-api-key"):
|
||||
raise ValueError("auth_mode must be 'bearer' or 'x-api-key'")
|
||||
if not api_key or not api_key.strip():
|
||||
raise ValueError("api_key must not be empty")
|
||||
if timeout <= 0 or connect_timeout <= 0:
|
||||
raise ValueError("timeouts must be greater than zero")
|
||||
|
||||
self.base_url = base_url.rstrip("/")
|
||||
self.api_key = api_key.strip()
|
||||
self.auth_mode = auth_mode
|
||||
self.csrf_token = csrf_token or secrets.token_urlsafe(32)
|
||||
self.timeout = httpx.Timeout(timeout, connect=connect_timeout)
|
||||
self._client = http_client or httpx.Client(timeout=self.timeout)
|
||||
self._owns_client = http_client is None
|
||||
|
||||
def close(self) -> None:
|
||||
if self._owns_client:
|
||||
self._client.close()
|
||||
|
||||
def __enter__(self) -> "OpenAgentAPIClient":
|
||||
return self
|
||||
|
||||
def __exit__(self, exc_type: object, exc: object, traceback: object) -> None:
|
||||
self.close()
|
||||
|
||||
def create_session(
|
||||
self,
|
||||
*,
|
||||
external_subject_id: Optional[str] = None,
|
||||
idempotency_key: Optional[str] = None,
|
||||
metadata: Optional[Dict[str, Any]] = None,
|
||||
) -> Dict[str, Any]:
|
||||
payload = self._compact_payload(
|
||||
{
|
||||
"external_subject_id": external_subject_id,
|
||||
"idempotency_key": idempotency_key,
|
||||
"metadata": metadata,
|
||||
}
|
||||
)
|
||||
return self._request_json("POST", "/api/open/agent-sessions", json_body=payload)
|
||||
|
||||
def send_message(
|
||||
self,
|
||||
session_id: str,
|
||||
message: str,
|
||||
*,
|
||||
idempotency_key: Optional[str] = None,
|
||||
metadata: Optional[Dict[str, Any]] = None,
|
||||
) -> Dict[str, Any]:
|
||||
payload = self._message_payload(message, idempotency_key=idempotency_key, metadata=metadata)
|
||||
return self._request_json(
|
||||
"POST",
|
||||
f"/api/open/agent-sessions/{self._path_segment(session_id)}/messages",
|
||||
json_body=payload,
|
||||
)
|
||||
|
||||
def stream_message(
|
||||
self,
|
||||
session_id: str,
|
||||
message: str,
|
||||
*,
|
||||
idempotency_key: Optional[str] = None,
|
||||
metadata: Optional[Dict[str, Any]] = None,
|
||||
) -> Iterator[OpenAgentEvent]:
|
||||
payload = self._message_payload(message, idempotency_key=idempotency_key, metadata=metadata)
|
||||
path = f"/api/open/agent-sessions/{self._path_segment(session_id)}/messages/stream"
|
||||
try:
|
||||
with self._client.stream(
|
||||
"POST",
|
||||
self._url(path),
|
||||
headers=self._headers(include_csrf=True, accept="text/event-stream"),
|
||||
json=payload,
|
||||
timeout=self.timeout,
|
||||
) as response:
|
||||
self._raise_for_error(response)
|
||||
yield from self._iter_sse_events(response)
|
||||
except OpenAgentAPIError:
|
||||
raise
|
||||
except httpx.RequestError as exc:
|
||||
raise OpenAgentTransportError(f"Open Agent API stream failed: {exc}") from exc
|
||||
|
||||
def get_run(self, session_id: str, run_id: str) -> Dict[str, Any]:
|
||||
path = (
|
||||
f"/api/open/agent-sessions/{self._path_segment(session_id)}"
|
||||
f"/runs/{self._path_segment(run_id)}"
|
||||
)
|
||||
return self._request_json("GET", path)
|
||||
|
||||
def cancel_run(self, session_id: str, run_id: str) -> Dict[str, Any]:
|
||||
path = (
|
||||
f"/api/open/agent-sessions/{self._path_segment(session_id)}"
|
||||
f"/runs/{self._path_segment(run_id)}/cancel"
|
||||
)
|
||||
return self._request_json("POST", path)
|
||||
|
||||
def _request_json(
|
||||
self,
|
||||
method: str,
|
||||
path: str,
|
||||
*,
|
||||
json_body: Optional[Dict[str, Any]] = None,
|
||||
) -> Dict[str, Any]:
|
||||
try:
|
||||
response = self._client.request(
|
||||
method,
|
||||
self._url(path),
|
||||
headers=self._headers(
|
||||
include_csrf=method.upper() in {"POST", "PUT", "PATCH", "DELETE"},
|
||||
accept="application/json",
|
||||
),
|
||||
json=json_body,
|
||||
timeout=self.timeout,
|
||||
)
|
||||
except httpx.RequestError as exc:
|
||||
raise OpenAgentTransportError(f"Open Agent API request failed: {exc}") from exc
|
||||
|
||||
self._raise_for_error(response)
|
||||
if not response.content:
|
||||
return {}
|
||||
try:
|
||||
body = response.json()
|
||||
except ValueError as exc:
|
||||
raise OpenAgentProtocolError("Open Agent API returned invalid JSON") from exc
|
||||
if not isinstance(body, dict):
|
||||
raise OpenAgentProtocolError("Open Agent API JSON response must be an object")
|
||||
return body
|
||||
|
||||
def _headers(self, *, include_csrf: bool, accept: str) -> Dict[str, str]:
|
||||
headers = {
|
||||
"Accept": accept,
|
||||
"User-Agent": "arr-open-agent-integration/0.1",
|
||||
}
|
||||
if self.auth_mode == "x-api-key":
|
||||
headers["X-DeerFlow-Open-API-Key"] = self.api_key
|
||||
else:
|
||||
headers["Authorization"] = f"Bearer {self.api_key}"
|
||||
if include_csrf:
|
||||
headers["X-CSRF-Token"] = self.csrf_token
|
||||
headers["Cookie"] = f"csrf_token={self.csrf_token}"
|
||||
return headers
|
||||
|
||||
def _url(self, path: str) -> str:
|
||||
return f"{self.base_url}{path}"
|
||||
|
||||
@staticmethod
|
||||
def _path_segment(value: str) -> str:
|
||||
return quote(str(value), safe="")
|
||||
|
||||
@staticmethod
|
||||
def _compact_payload(payload: Dict[str, Any]) -> Dict[str, Any]:
|
||||
return {key: value for key, value in payload.items() if value is not None}
|
||||
|
||||
def _message_payload(
|
||||
self,
|
||||
message: str,
|
||||
*,
|
||||
idempotency_key: Optional[str],
|
||||
metadata: Optional[Dict[str, Any]],
|
||||
) -> Dict[str, Any]:
|
||||
return self._compact_payload(
|
||||
{
|
||||
"message": message,
|
||||
"idempotency_key": idempotency_key,
|
||||
"metadata": metadata,
|
||||
}
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _iter_sse_events(response: httpx.Response) -> Iterator[OpenAgentEvent]:
|
||||
event_name: Optional[str] = None
|
||||
event_id: Optional[str] = None
|
||||
retry_ms: Optional[int] = None
|
||||
data_lines: List[str] = []
|
||||
|
||||
def finish_event() -> Optional[OpenAgentEvent]:
|
||||
nonlocal event_name, event_id, retry_ms, data_lines
|
||||
if event_name is None and event_id is None and retry_ms is None and not data_lines:
|
||||
return None
|
||||
data_text = "\n".join(data_lines)
|
||||
if not data_text:
|
||||
data: Any = None
|
||||
else:
|
||||
try:
|
||||
data = json.loads(data_text)
|
||||
except ValueError:
|
||||
data = data_text
|
||||
event = OpenAgentEvent(
|
||||
event=event_name,
|
||||
data=data,
|
||||
event_id=event_id,
|
||||
retry_ms=retry_ms,
|
||||
)
|
||||
event_name = None
|
||||
event_id = None
|
||||
retry_ms = None
|
||||
data_lines = []
|
||||
return event
|
||||
|
||||
for line in response.iter_lines():
|
||||
if line == "":
|
||||
event = finish_event()
|
||||
if event is not None:
|
||||
yield event
|
||||
continue
|
||||
if line.startswith(":"):
|
||||
continue
|
||||
|
||||
field, separator, value = line.partition(":")
|
||||
if not separator:
|
||||
value = ""
|
||||
elif value.startswith(" "):
|
||||
value = value[1:]
|
||||
|
||||
if field == "event":
|
||||
event_name = value
|
||||
elif field == "data":
|
||||
data_lines.append(value)
|
||||
elif field == "id":
|
||||
event_id = value
|
||||
elif field == "retry":
|
||||
try:
|
||||
parsed_retry = int(value)
|
||||
except ValueError:
|
||||
continue
|
||||
if parsed_retry >= 0:
|
||||
retry_ms = parsed_retry
|
||||
|
||||
event = finish_event()
|
||||
if event is not None:
|
||||
yield event
|
||||
|
||||
@staticmethod
|
||||
def _raise_for_error(response: httpx.Response) -> None:
|
||||
if response.status_code < 400:
|
||||
return
|
||||
try:
|
||||
response_text = response.text
|
||||
except httpx.ResponseNotRead:
|
||||
response.read()
|
||||
response_text = response.text
|
||||
try:
|
||||
body = response.json()
|
||||
except ValueError:
|
||||
detail: Any = response_text or response.reason_phrase
|
||||
else:
|
||||
detail = body.get("detail", body) if isinstance(body, dict) else body
|
||||
request_id = response.headers.get("x-request-id") or response.headers.get("request-id")
|
||||
raise OpenAgentAPIError(
|
||||
response.status_code,
|
||||
detail,
|
||||
response_text,
|
||||
request_id=request_id,
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user