Spaces:
Running
Running
| from __future__ import annotations | |
| import time | |
| from dataclasses import dataclass, field | |
| from typing import Any | |
| import httpx | |
| from longmemeval.config import TERMINAL_TASK_STATUSES, Settings | |
| class TimedResult: | |
| data: Any | |
| latency_ms: float | |
| class BrainAPIClient: | |
| settings: Settings | |
| timeout: float = 120.0 | |
| max_retries: int = 5 | |
| _client: httpx.Client = field(init=False, repr=False) | |
| def __post_init__(self) -> None: | |
| self.settings.require_brainapi() | |
| self._client = httpx.Client( | |
| base_url=self.settings.brainapi_url, | |
| headers={ | |
| "Authorization": f"Bearer {self.settings.brainpat_token}", | |
| "Content-Type": "application/json", | |
| }, | |
| timeout=self.timeout, | |
| ) | |
| def close(self) -> None: | |
| self._client.close() | |
| def __enter__(self) -> BrainAPIClient: | |
| return self | |
| def __exit__(self, *args: object) -> None: | |
| self.close() | |
| def _request( | |
| self, | |
| method: str, | |
| path: str, | |
| *, | |
| json: dict | None = None, | |
| params: dict | None = None, | |
| ) -> TimedResult: | |
| last_error: Exception | None = None | |
| for attempt in range(self.max_retries): | |
| started = time.perf_counter() | |
| try: | |
| response = self._client.request( | |
| method, path, json=json, params=params | |
| ) | |
| latency_ms = (time.perf_counter() - started) * 1000 | |
| if response.status_code in (429, 500, 502, 503, 504): | |
| wait = min(2**attempt * 0.5, 30) | |
| time.sleep(wait) | |
| last_error = httpx.HTTPStatusError( | |
| f"HTTP {response.status_code}", | |
| request=response.request, | |
| response=response, | |
| ) | |
| continue | |
| response.raise_for_status() | |
| if response.status_code == 204 or not response.content: | |
| return TimedResult(data=None, latency_ms=latency_ms) | |
| return TimedResult(data=response.json(), latency_ms=latency_ms) | |
| except (httpx.TransportError, httpx.TimeoutException) as exc: | |
| last_error = exc | |
| wait = min(2**attempt * 0.5, 30) | |
| time.sleep(wait) | |
| raise RuntimeError(f"Request failed after retries: {method} {path}: {last_error}") | |
| def ingest_text( | |
| self, | |
| text: str, | |
| brain_id: str, | |
| *, | |
| observate_for: list[str] | None = None, | |
| preferred_extraction_entities: list[str] | None = None, | |
| source_timestamp: str | None = None, | |
| ) -> TimedResult: | |
| body = { | |
| "data": {"data_type": "text", "text_data": text}, | |
| "brain_id": brain_id, | |
| "observate_for": observate_for | |
| or [ | |
| "facts about each speaker", | |
| "events and when they happened", | |
| "relationships between people", | |
| ], | |
| "preferred_extraction_entities": preferred_extraction_entities | |
| or ["Person", "Event", "Location", "Date"], | |
| } | |
| if source_timestamp: | |
| body["source_timestamp"] = source_timestamp | |
| return self._request("POST", "/ingest/", json=body) | |
| def get_task(self, task_id: str, brain_id: str) -> TimedResult: | |
| return self._request( | |
| "GET", f"/tasks/{task_id}", params={"brain_id": brain_id} | |
| ) | |
| def wait_for_task( | |
| self, | |
| task_id: str, | |
| brain_id: str, | |
| *, | |
| timeout_s: float = 600.0, | |
| poll_interval_s: float = 2.0, | |
| ) -> TimedResult: | |
| started = time.perf_counter() | |
| while True: | |
| result = self.get_task(task_id, brain_id) | |
| status = (result.data or {}).get("status", "unknown") | |
| if status in TERMINAL_TASK_STATUSES: | |
| return TimedResult( | |
| data=result.data, | |
| latency_ms=(time.perf_counter() - started) * 1000, | |
| ) | |
| if time.perf_counter() - started > timeout_s: | |
| raise TimeoutError( | |
| f"Task {task_id} did not finish within {timeout_s}s " | |
| f"(last status={status})" | |
| ) | |
| time.sleep(poll_interval_s) | |
| def retrieve_context( | |
| self, | |
| text: str, | |
| brain_id: str, | |
| *, | |
| historical_limit: int = 10, | |
| max_passages: int = 8, | |
| max_facts: int = 40, | |
| apply_fact_filter: bool = True, | |
| use_ppr: bool = False, | |
| sufficiency_retry: bool = False, | |
| cross_event_bridges: int = 3, | |
| ) -> TimedResult: | |
| body: dict[str, Any] = { | |
| "text": text, | |
| "brain_id": brain_id, | |
| "historical_limit": historical_limit, | |
| "max_passages": max_passages, | |
| "max_facts": max_facts, | |
| "apply_fact_filter": apply_fact_filter, | |
| "use_ppr": use_ppr, | |
| "sufficiency_retry": sufficiency_retry, | |
| "cross_event_bridges": cross_event_bridges, | |
| } | |
| return self._request("POST", "/retrieve/context", json=body) | |
| def entity_status(self, target: str, brain_id: str) -> TimedResult: | |
| return self._request( | |
| "GET", | |
| "/retrieve/entity/status", | |
| params={"target": target, "brain_id": brain_id}, | |
| ) | |
| def list_entities(self, brain_id: str, *, limit: int = 5) -> TimedResult: | |
| return self._request( | |
| "GET", | |
| "/retrieve/entities", | |
| params={"brain_id": brain_id, "limit": limit}, | |
| ) | |