Simplify to memory system api

memory_system_api/clients.py

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