Livrare LOT 1 - Didi
This commit is contained in:
commit
5380c3fc63
990 changed files with 133308 additions and 0 deletions
307
ai_platform/modules/didi_brain/shared/atomic_api.py
Normal file
307
ai_platform/modules/didi_brain/shared/atomic_api.py
Normal file
|
|
@ -0,0 +1,307 @@
|
|||
"""Typed async client for the Atomic REST API.
|
||||
|
||||
Atomic is the brain — this client is how every other DidiBrain component
|
||||
(bootstrap, scraper, extractor, lint, didi_client) talks to it.
|
||||
|
||||
Design notes:
|
||||
- One AtomicClient instance per long-running process. Reusable httpx.AsyncClient.
|
||||
- Auth: Bearer token loaded from `settings.atomic_token`. If empty, public
|
||||
endpoints (/health, setup) still work.
|
||||
- Errors: any non-2xx is wrapped in AtomicApiError with the response body.
|
||||
- Pagination: helpers like `iter_atoms()` (added later) yield pages.
|
||||
|
||||
Only the endpoints we actually need are wrapped here. Add more as required;
|
||||
do not pre-emptively wrap the full ~78-route surface.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import AsyncIterator, Sequence
|
||||
from contextlib import asynccontextmanager
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
from tenacity import (
|
||||
AsyncRetrying,
|
||||
retry_if_exception_type,
|
||||
stop_after_attempt,
|
||||
wait_exponential,
|
||||
)
|
||||
|
||||
from shared.config import settings
|
||||
from shared.logging import get_logger
|
||||
|
||||
log = get_logger(__name__)
|
||||
|
||||
|
||||
class AtomicApiError(Exception):
|
||||
"""Raised on any non-2xx response from atomic-server."""
|
||||
|
||||
def __init__(self, message: str, *, status: int, body: str = "", url: str = ""):
|
||||
super().__init__(message)
|
||||
self.status = status
|
||||
self.body = body
|
||||
self.url = url
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class AtomSummary:
|
||||
"""Minimal atom view returned by list/search endpoints."""
|
||||
|
||||
id: str
|
||||
content: str
|
||||
source_url: str | None
|
||||
embedding_status: str
|
||||
tagging_status: str
|
||||
created_at: str
|
||||
updated_at: str
|
||||
tags: list[dict[str, Any]]
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, d: dict[str, Any]) -> AtomSummary:
|
||||
return cls(
|
||||
id=d["id"],
|
||||
content=d.get("content", ""),
|
||||
source_url=d.get("source_url"),
|
||||
embedding_status=d.get("embedding_status", "unknown"),
|
||||
tagging_status=d.get("tagging_status", "unknown"),
|
||||
created_at=d.get("created_at", ""),
|
||||
updated_at=d.get("updated_at", ""),
|
||||
tags=d.get("tags", []),
|
||||
)
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class SearchHit:
|
||||
"""One result from POST /api/search."""
|
||||
|
||||
atom_id: str
|
||||
similarity: float
|
||||
matching_chunk_content: str | None
|
||||
snippet: str | None
|
||||
source_url: str | None
|
||||
tags: list[dict[str, Any]]
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, d: dict[str, Any]) -> SearchHit:
|
||||
return cls(
|
||||
atom_id=d.get("id") or d.get("atom_id", ""),
|
||||
similarity=float(d.get("similarity_score") or d.get("similarity", 0.0)),
|
||||
matching_chunk_content=d.get("matching_chunk_content"),
|
||||
snippet=d.get("snippet"),
|
||||
source_url=d.get("source_url"),
|
||||
tags=d.get("tags", []),
|
||||
)
|
||||
|
||||
|
||||
class AtomicClient:
|
||||
"""Async REST client for atomic-server. Use as `async with AtomicClient() as a:`."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
base_url: str | None = None,
|
||||
token: str | None = None,
|
||||
timeout: float = 60.0,
|
||||
max_attempts: int = 3,
|
||||
):
|
||||
self._base_url = (base_url or settings.atomic_url).rstrip("/")
|
||||
self._token = token if token is not None else settings.atomic_token
|
||||
self._max_attempts = max_attempts
|
||||
|
||||
headers = {"Content-Type": "application/json"}
|
||||
if self._token:
|
||||
headers["Authorization"] = f"Bearer {self._token}"
|
||||
|
||||
self._http = httpx.AsyncClient(
|
||||
base_url=self._base_url,
|
||||
headers=headers,
|
||||
timeout=httpx.Timeout(timeout, connect=10.0),
|
||||
)
|
||||
|
||||
async def aclose(self) -> None:
|
||||
await self._http.aclose()
|
||||
|
||||
async def __aenter__(self) -> AtomicClient:
|
||||
return self
|
||||
|
||||
async def __aexit__(self, *args: Any) -> None:
|
||||
await self.aclose()
|
||||
|
||||
@property
|
||||
def base_url(self) -> str:
|
||||
return self._base_url
|
||||
|
||||
@property
|
||||
def has_token(self) -> bool:
|
||||
return bool(self._token)
|
||||
|
||||
# ---------------------------------------------------------------- internals
|
||||
async def _request(
|
||||
self,
|
||||
method: str,
|
||||
path: str,
|
||||
*,
|
||||
json: Any = None,
|
||||
params: dict[str, Any] | None = None,
|
||||
retry_on_5xx: bool = True,
|
||||
) -> Any:
|
||||
async for attempt in AsyncRetrying(
|
||||
stop=stop_after_attempt(self._max_attempts),
|
||||
wait=wait_exponential(multiplier=1, min=1, max=10),
|
||||
retry=retry_if_exception_type((httpx.HTTPError, httpx.TimeoutException)),
|
||||
reraise=True,
|
||||
):
|
||||
with attempt:
|
||||
resp = await self._http.request(method, path, json=json, params=params)
|
||||
if resp.status_code >= 500 and retry_on_5xx:
|
||||
raise httpx.HTTPError(f"5xx from {path}: {resp.status_code}")
|
||||
if resp.status_code >= 400:
|
||||
raise AtomicApiError(
|
||||
f"HTTP {resp.status_code} {method} {path}",
|
||||
status=resp.status_code,
|
||||
body=resp.text[:1000],
|
||||
url=str(resp.url),
|
||||
)
|
||||
if resp.status_code == 204 or not resp.content:
|
||||
return None
|
||||
try:
|
||||
return resp.json()
|
||||
except ValueError:
|
||||
return resp.text
|
||||
raise AtomicApiError("retry loop exited without result", status=0) # unreachable
|
||||
|
||||
# ----------------------------------------------------------------- health
|
||||
async def health(self) -> dict[str, Any]:
|
||||
return await self._request("GET", "/health")
|
||||
|
||||
async def setup_status(self) -> dict[str, Any]:
|
||||
"""Whether the instance still needs initial token claim."""
|
||||
return await self._request("GET", "/api/setup/status")
|
||||
|
||||
# ----------------------------------------------------------------- settings
|
||||
async def get_settings(self) -> dict[str, Any]:
|
||||
return await self._request("GET", "/api/settings")
|
||||
|
||||
async def set_setting(self, key: str, value: str) -> Any:
|
||||
"""Set a single setting. Use set_settings() for bulk."""
|
||||
return await self._request(
|
||||
"PUT", f"/api/settings/{key}", json={"value": value}
|
||||
)
|
||||
|
||||
async def set_settings(self, items: dict[str, str]) -> dict[str, Any]:
|
||||
"""Apply many settings sequentially. Returns map of key → response."""
|
||||
out: dict[str, Any] = {}
|
||||
for k, v in items.items():
|
||||
out[k] = await self.set_setting(k, v)
|
||||
return out
|
||||
|
||||
# -------------------------------------------------------------------- tags
|
||||
async def list_tags(self, *, min_count: int = 0) -> list[dict[str, Any]]:
|
||||
return await self._request("GET", "/api/tags", params={"min_count": min_count})
|
||||
|
||||
async def create_tag(
|
||||
self, name: str, *, parent_id: str | None = None
|
||||
) -> dict[str, Any]:
|
||||
body: dict[str, Any] = {"name": name}
|
||||
if parent_id:
|
||||
body["parent_id"] = parent_id
|
||||
return await self._request("POST", "/api/tags", json=body)
|
||||
|
||||
# ------------------------------------------------------------------- atoms
|
||||
async def create_atom(
|
||||
self,
|
||||
*,
|
||||
content: str,
|
||||
source_url: str | None = None,
|
||||
tag_ids: Sequence[str] | None = None,
|
||||
published_at: str | None = None,
|
||||
) -> dict[str, Any]:
|
||||
body: dict[str, Any] = {"content": content, "tag_ids": list(tag_ids or [])}
|
||||
if source_url:
|
||||
body["source_url"] = source_url
|
||||
if published_at:
|
||||
body["published_at"] = published_at
|
||||
return await self._request("POST", "/api/atoms", json=body)
|
||||
|
||||
async def get_atom(self, atom_id: str) -> dict[str, Any]:
|
||||
return await self._request("GET", f"/api/atoms/{atom_id}")
|
||||
|
||||
async def delete_atom(self, atom_id: str) -> None:
|
||||
await self._request("DELETE", f"/api/atoms/{atom_id}")
|
||||
|
||||
async def get_atom_by_source_url(self, source_url: str) -> dict[str, Any] | None:
|
||||
try:
|
||||
# Atomic's GetAtomBySourceUrlQuery uses `url` not `source_url`.
|
||||
return await self._request(
|
||||
"GET", "/api/atoms/by-source-url", params={"url": source_url}
|
||||
)
|
||||
except AtomicApiError as e:
|
||||
if e.status == 404:
|
||||
return None
|
||||
raise
|
||||
|
||||
async def get_embedding_status(self, atom_id: str) -> dict[str, Any]:
|
||||
return await self._request("GET", f"/api/atoms/{atom_id}/embedding-status")
|
||||
|
||||
async def list_atoms(
|
||||
self,
|
||||
*,
|
||||
limit: int = 50,
|
||||
offset: int = 0,
|
||||
tag_id: str | None = None,
|
||||
) -> dict[str, Any]:
|
||||
params: dict[str, Any] = {"limit": limit, "offset": offset}
|
||||
if tag_id:
|
||||
params["tag_id"] = tag_id
|
||||
return await self._request("GET", "/api/atoms", params=params)
|
||||
|
||||
# ------------------------------------------------------------------ search
|
||||
async def search(
|
||||
self,
|
||||
query: str,
|
||||
*,
|
||||
mode: str = "semantic",
|
||||
limit: int = 20,
|
||||
threshold: float | None = None,
|
||||
) -> list[SearchHit]:
|
||||
body: dict[str, Any] = {"query": query, "mode": mode, "limit": limit}
|
||||
if threshold is not None:
|
||||
body["threshold"] = threshold
|
||||
result = await self._request("POST", "/api/search", json=body)
|
||||
# Atomic returns either a list directly or {"results": [...]}
|
||||
if isinstance(result, dict):
|
||||
items = result.get("results") or result.get("data") or []
|
||||
else:
|
||||
items = result or []
|
||||
return [SearchHit.from_dict(item) for item in items]
|
||||
|
||||
async def find_similar(
|
||||
self, atom_id: str, *, threshold: float = 0.5, limit: int = 20
|
||||
) -> list[SearchHit]:
|
||||
params = {"threshold": threshold, "limit": limit}
|
||||
result = await self._request(
|
||||
"GET", f"/api/atoms/{atom_id}/similar", params=params
|
||||
)
|
||||
items = result if isinstance(result, list) else result.get("results", [])
|
||||
return [SearchHit.from_dict(item) for item in items]
|
||||
|
||||
# ------------------------------------------------------------- embeddings
|
||||
async def get_pipeline_status(self) -> dict[str, Any]:
|
||||
return await self._request("GET", "/api/embeddings/status")
|
||||
|
||||
async def process_pending_embeddings(self) -> dict[str, Any]:
|
||||
return await self._request("POST", "/api/embeddings/process-pending")
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def atomic_client(
|
||||
*, token: str | None = None
|
||||
) -> AsyncIterator[AtomicClient]:
|
||||
"""`async with atomic_client() as a:` for short-lived scripts."""
|
||||
client = AtomicClient(token=token)
|
||||
try:
|
||||
yield client
|
||||
finally:
|
||||
await client.aclose()
|
||||
Loading…
Add table
Add a link
Reference in a new issue