145 lines
4.8 KiB
Python
145 lines
4.8 KiB
Python
"""HTTP client for the brain API.
|
|
|
|
Thin wrapper that the three scheduler tasks (feeder/auditor/watcher) use to
|
|
talk to didibrain-api. All calls have generous timeouts (some operations
|
|
internally trigger LLM calls + DB queries) and return None on failure so
|
|
the caller can decide whether to retry or just skip.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from typing import Any
|
|
|
|
import httpx
|
|
|
|
from scheduler.config import settings
|
|
from shared.logging import get_logger
|
|
|
|
log = get_logger(__name__)
|
|
|
|
|
|
class BrainClient:
|
|
"""Async client for the brain HTTP API.
|
|
|
|
Reuse one instance per task — httpx.AsyncClient pools connections.
|
|
"""
|
|
|
|
def __init__(self, *, base_url: str | None = None) -> None:
|
|
self._base = (base_url or settings.brain_api_url).rstrip("/")
|
|
self._http = httpx.AsyncClient(
|
|
base_url=self._base,
|
|
timeout=httpx.Timeout(
|
|
settings.brain_api_timeout_s, connect=10.0
|
|
),
|
|
headers={"Content-Type": "application/json"},
|
|
)
|
|
|
|
async def aclose(self) -> None:
|
|
await self._http.aclose()
|
|
|
|
async def __aenter__(self) -> BrainClient:
|
|
return self
|
|
|
|
async def __aexit__(self, *args: Any) -> None:
|
|
await self.aclose()
|
|
|
|
# --------------------------------------------------------------- ingest
|
|
async def ingest(
|
|
self,
|
|
*,
|
|
claim: str | None,
|
|
evidence: list[dict],
|
|
default_tags: list[str] | None = None,
|
|
run_extraction: bool = True,
|
|
) -> dict | None:
|
|
"""POST /v1/ingest. Returns response dict or None on error.
|
|
|
|
Each evidence item should have: url, title, summary, full_text (opt),
|
|
publisher (opt), published_at (ISO str, opt).
|
|
"""
|
|
body = {
|
|
"claim": claim,
|
|
"evidence": evidence,
|
|
"default_tags": default_tags or [],
|
|
"run_extraction": run_extraction,
|
|
}
|
|
try:
|
|
resp = await self._http.post("/v1/ingest", json=body)
|
|
if resp.status_code >= 400:
|
|
log.warning(
|
|
"brain_ingest_http_error",
|
|
status=resp.status_code,
|
|
body=resp.text[:200],
|
|
)
|
|
return None
|
|
return resp.json()
|
|
except httpx.HTTPError as e:
|
|
log.warning("brain_ingest_failed", error=str(e))
|
|
return None
|
|
|
|
# ----------------------------------------------------------- invalidate
|
|
async def invalidate(
|
|
self,
|
|
*,
|
|
topic_codes: list[str] | None = None,
|
|
entity_canonicals: list[str] | None = None,
|
|
claim_pattern: str | None = None,
|
|
since_iso: str | None = None,
|
|
invalidate_gold: bool = False,
|
|
dry_run: bool = False,
|
|
actor: str = "scheduler",
|
|
reason: str | None = None,
|
|
) -> dict | None:
|
|
"""POST /v1/cache/invalidate. Returns counts dict or None on error."""
|
|
body: dict[str, Any] = {
|
|
"invalidate_gold": invalidate_gold,
|
|
"dry_run": dry_run,
|
|
"actor": actor,
|
|
}
|
|
if topic_codes:
|
|
body["topic_codes"] = topic_codes
|
|
if entity_canonicals:
|
|
body["entity_canonicals"] = entity_canonicals
|
|
if claim_pattern:
|
|
body["claim_pattern"] = claim_pattern
|
|
if since_iso:
|
|
body["since"] = since_iso
|
|
if reason:
|
|
body["reason"] = reason
|
|
try:
|
|
resp = await self._http.post("/v1/cache/invalidate", json=body)
|
|
if resp.status_code >= 400:
|
|
log.warning(
|
|
"brain_invalidate_http_error",
|
|
status=resp.status_code,
|
|
body=resp.text[:200],
|
|
)
|
|
return None
|
|
return resp.json()
|
|
except httpx.HTTPError as e:
|
|
log.warning("brain_invalidate_failed", error=str(e))
|
|
return None
|
|
|
|
# ------------------------------------------------------------ canonicalize
|
|
async def canonicalize(
|
|
self, *, claim: str, current_date: str | None = None
|
|
) -> dict | None:
|
|
"""POST /v1/canonicalize. Returns response dict or None on error."""
|
|
body: dict[str, Any] = {"claim": claim}
|
|
if current_date:
|
|
body["current_date"] = current_date
|
|
try:
|
|
resp = await self._http.post("/v1/canonicalize", json=body)
|
|
if resp.status_code >= 400:
|
|
return None
|
|
return resp.json()
|
|
except httpx.HTTPError:
|
|
return None
|
|
|
|
# ------------------------------------------------------------------ misc
|
|
async def health(self) -> bool:
|
|
try:
|
|
resp = await self._http.get("/health")
|
|
return resp.status_code == 200
|
|
except httpx.HTTPError:
|
|
return False
|