307 lines
10 KiB
Python
307 lines
10 KiB
Python
"""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()
|