"""Async clients for OpenViking and EverOS used by the lightweight API.""" from __future__ import annotations from datetime import datetime, timezone from typing import Any import httpx from .config import get_config from .store import OpenVikingUserKeyStore class OpenVikingMemorySystemClient: def __init__(self, store: OpenVikingUserKeyStore | None = None) -> None: config = get_config() self.base_url = config.openviking.url.rstrip("/") self.root_key = config.openviking.api_key or "your-secret-root-key" self.timeout = config.openviking.timeout self.verify_ssl = config.openviking.verify_ssl self.store = store or OpenVikingUserKeyStore(config.storage.sqlite_path) async def health(self) -> dict[str, Any]: async with self._client(self.root_key) as client: response = await client.get("/health") response.raise_for_status() return response.json() async def ensure_user(self, user_id: str) -> str: existing = self.store.get_user_key(user_id) if existing: return existing async with self._client(self.root_key) as client: response = await client.post( "/api/v1/admin/accounts", json={"account_id": user_id, "admin_user_id": user_id}, ) response.raise_for_status() data = response.json() user_key = self._extract_user_key(data) if not user_key: raise RuntimeError("OpenViking did not return user_key") self.store.save_user_key(user_id, user_key) return user_key async def ensure_session(self, user_key: str, session_id: str) -> dict[str, Any]: async with self._client(user_key) as client: response = await client.post("/api/v1/sessions", json={"session_id": session_id}) if response.status_code in {409, 422}: return {"session_id": session_id, "status": "exists"} response.raise_for_status() return response.json() async def append_message(self, user_key: str, session_id: str, role: str, content: str) -> dict[str, Any]: async with self._client(user_key) as client: response = await client.post( f"/api/v1/sessions/{session_id}/messages", json={"role": role, "content": content}, ) response.raise_for_status() return response.json() async def commit_session(self, user_key: str, session_id: str) -> dict[str, Any]: async with self._client(user_key) as client: response = await client.post(f"/api/v1/sessions/{session_id}/commit") response.raise_for_status() return response.json() async def extract_session(self, user_key: str, session_id: str) -> dict[str, Any]: async with self._client(user_key) as client: response = await client.post(f"/api/v1/sessions/{session_id}/extract") response.raise_for_status() return response.json() async def get_task(self, user_key: str, task_id: str) -> dict[str, Any]: async with self._client(user_key) as client: response = await client.get(f"/api/v1/tasks/{task_id}") response.raise_for_status() return response.json() async def find(self, user_key: str, user_id: str, query: str, limit: int) -> dict[str, Any]: async with self._client(user_key) as client: response = await client.post( "/api/v1/search/find", json={ "query": query, "target_uri": f"viking://user/{user_id}/memories/", "limit": limit, }, ) response.raise_for_status() return response.json() async def search(self, user_key: str, session_id: str | None, query: str, limit: int) -> dict[str, Any]: payload: dict[str, Any] = {"query": query, "limit": limit} if session_id: payload["session_id"] = session_id async with self._client(user_key) as client: response = await client.post("/api/v1/search/search", json=payload) response.raise_for_status() return response.json() def _client(self, api_key: str) -> httpx.AsyncClient: return httpx.AsyncClient( base_url=self.base_url, headers={"X-API-Key": api_key, "Content-Type": "application/json"}, timeout=self.timeout, verify=self.verify_ssl, ) def _extract_user_key(self, data: dict[str, Any]) -> str | None: result = data.get("result") if isinstance(data.get("result"), dict) else data value = result.get("user_key") if isinstance(result, dict) else None return str(value) if value else None class EverOSMemorySystemClient: def __init__(self) -> None: config = get_config() self.base_url = config.everos.url.rstrip("/") self.api_key = config.everos.api_key self.timeout = config.everos.timeout self.verify_ssl = config.everos.verify_ssl self.health_path = config.everos.health_path async def health(self) -> dict[str, Any]: async with self._client() as client: response = await client.get(self.health_path) response.raise_for_status() return response.json() async def append_message(self, user_id: str, session_id: str, role: str, content: str) -> dict[str, Any]: async with self._client() as client: response = await client.post( "/api/v1/memories", json=self.build_message_payload(user_id=user_id, session_id=session_id, role=role, content=content), ) response.raise_for_status() return response.json() def build_message_payload(self, user_id: str, session_id: str, role: str, content: str) -> dict[str, Any]: everos_role = "assistant" if role == "assistant" else "user" sender_id = "assistant" if everos_role == "assistant" else user_id timestamp = int(datetime.now(timezone.utc).timestamp() * 1000) return { "user_id": user_id, "session_id": session_id, "messages": [ { "message_id": f"msg_{timestamp}", "timestamp": timestamp, "sender_id": sender_id, "sender_name": sender_id, "role": everos_role, "content": content, } ], } async def flush(self, user_id: str, session_id: str) -> dict[str, Any]: async with self._client() as client: response = await client.post("/api/v1/memories/flush", json={"user_id": user_id, "session_id": session_id}) response.raise_for_status() return response.json() async def search(self, user_id: str, session_id: str | None, query: str, method: str, limit: int) -> dict[str, Any]: filters: dict[str, Any] = {"user_id": user_id} if session_id: filters["session_id"] = session_id async with self._client() as client: response = await client.post( "/api/v1/memories/search", json={ "query": query, "method": method, "memory_types": ["episodic_memory", "profile", "raw_message"], "filters": filters, "top_k": limit, "include_original_data": True, }, ) response.raise_for_status() return response.json() async def get_profile(self, user_id: str) -> dict[str, Any]: async with self._client() as client: response = await client.post( "/api/v1/memories/get", json={ "memory_type": "profile", "filters": {"user_id": user_id}, "page": 1, "page_size": 20, }, ) response.raise_for_status() return response.json() def _client(self) -> httpx.AsyncClient: headers = {"Content-Type": "application/json"} if self.api_key: headers["X-API-Key"] = self.api_key headers["Authorization"] = f"Bearer {self.api_key}" return httpx.AsyncClient( base_url=self.base_url, headers=headers, timeout=self.timeout, verify=self.verify_ssl, )