Livrare LOT 1 - Didi

ai_platform/modules/didi_brain/scheduler/brain_client.py

@@ -0,0 +1,145 @@
+"""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