Refactor OpenVikingMemorySystemClient to use OpenVikingCredential and update session handling; increase EverOS timeout to 180 seconds; enhance tests for credential handling and vector removal in search results.

memory_system_api/clients.py

@@ -1,6 +1,7 @@
 """Async clients for OpenViking and EverOS used by the lightweight API."""
 from __future__ import annotations
 
+from dataclasses import dataclass
 from datetime import datetime, timezone
 from typing import Any
 
@@ -10,6 +11,13 @@ from .config import get_config
 from .store import OpenVikingUserKeyStore
 
 
+@dataclass(frozen=True)
+class OpenVikingCredential:
+    api_key: str
+    account_id: str | None = None
+    user_id: str | None = None
+
+
 class OpenVikingMemorySystemClient:
     def __init__(self, store: OpenVikingUserKeyStore | None = None) -> None:
         config = get_config()
@@ -25,35 +33,42 @@ class OpenVikingMemorySystemClient:
             response.raise_for_status()
             return response.json()
 
-    async def ensure_user(self, user_id: str) -> str:
+    async def ensure_user(self, user_id: str) -> OpenVikingCredential:
         existing = self.store.get_user_key(user_id)
         if existing:
-            return existing
+            return OpenVikingCredential(api_key=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},
             )
+            if response.status_code == 409:
+                return self.root_credential(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")
+            return self.root_credential(user_id)
         self.store.save_user_key(user_id, user_key)
-        return user_key
+        return OpenVikingCredential(api_key=user_key)
 
-    async def ensure_session(self, user_key: str, session_id: str) -> dict[str, Any]:
-        async with self._client(user_key) as client:
+    def root_credential(self, user_id: str) -> OpenVikingCredential:
+        return OpenVikingCredential(api_key=self.root_key, account_id=user_id, user_id=user_id)
+
+    async def ensure_session(self, credential: OpenVikingCredential | str, session_id: str) -> dict[str, Any]:
+        async with self._credential_client(credential) 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:
+    async def append_message(
+        self, credential: OpenVikingCredential | str, session_id: str, role: str, content: str
+    ) -> dict[str, Any]:
+        async with self._credential_client(credential) as client:
             response = await client.post(
                 f"/api/v1/sessions/{session_id}/messages",
                 json={"role": role, "content": content},
@@ -61,26 +76,28 @@ class OpenVikingMemorySystemClient:
             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:
+    async def commit_session(self, credential: OpenVikingCredential | str, session_id: str) -> dict[str, Any]:
+        async with self._credential_client(credential) 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:
+    async def extract_session(self, credential: OpenVikingCredential | str, session_id: str) -> dict[str, Any]:
+        async with self._credential_client(credential) 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:
+    async def get_task(self, credential: OpenVikingCredential | str, task_id: str) -> dict[str, Any]:
+        async with self._credential_client(credential) 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:
+    async def find(
+        self, credential: OpenVikingCredential | str, user_id: str, query: str, limit: int
+    ) -> dict[str, Any]:
+        async with self._credential_client(credential) as client:
             response = await client.post(
                 "/api/v1/search/find",
                 json={
@@ -92,19 +109,34 @@ class OpenVikingMemorySystemClient:
             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]:
+    async def search(
+        self, credential: OpenVikingCredential | 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:
+        async with self._credential_client(credential) 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:
+    def _credential_client(self, credential: OpenVikingCredential | str) -> httpx.AsyncClient:
+        if isinstance(credential, str):
+            return self._client(credential)
+        headers = {}
+        if credential.account_id:
+            headers["X-OpenViking-Account"] = credential.account_id
+        if credential.user_id:
+            headers["X-OpenViking-User"] = credential.user_id
+        return self._client(credential.api_key, headers)
+
+    def _client(self, api_key: str, extra_headers: dict[str, str] | None = None) -> httpx.AsyncClient:
+        headers = {"X-API-Key": api_key, "Content-Type": "application/json"}
+        if extra_headers:
+            headers.update(extra_headers)
         return httpx.AsyncClient(
             base_url=self.base_url,
-            headers={"X-API-Key": api_key, "Content-Type": "application/json"},
+            headers=headers,
             timeout=self.timeout,
             verify=self.verify_ssl,
         )