tomtan/memory-gateway
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

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.

Browse Source
This commit is contained in:
tomtan
2026-05-19 12:30:35 +08:00
parent e689b13e4a
commit ff514ad1f5
5 changed files with 222 additions and 28 deletions
Download Patch File Download Diff File Expand all files Collapse all files

72
memory_system_api/clients.py
View File

@ -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,
)

2
memory_system_api/config.py
View File

@ -25,7 +25,7 @@ class OpenVikingConfig(BaseModel):
class EverOSConfig(BaseModel):
url: str = "http://127.0.0.1:1995"
api_key: str = ""
timeout: int = 30
timeout: int = 180
verify_ssl: bool = True
health_path: str = "/health"

26
memory_system_api/service.py
View File

@ -27,16 +27,14 @@ class MemorySystemService:
if not messages:
raise ValueError("at least one message is required")
user_key = await self.openviking.ensure_user(request.user_id)
await self.openviking.ensure_session(user_key, request.session_id)
print("user_key:", user_key) # Debugging line to check the user_key value
credential = await self.openviking.ensure_user(request.user_id)
await self.openviking.ensure_session(credential, request.session_id)
async def write_openviking() -> list[dict[str, Any]]:
results = []
for message in messages:
results.append(
await self.openviking.append_message(user_key, request.session_id, message["role"], message["content"])
await self.openviking.append_message(credential, request.session_id, message["role"], message["content"])
)
return results
@ -97,6 +95,7 @@ class MemorySystemService:
)
backends = await self._run_backends(openviking=search_openviking, everos=search_everos)
backends = self._remove_vectors_from_backends(backends)
items = self._merge_search_items(backends)
return SearchResponse(status=self._aggregate_status(backends), items=items[: request.limit], backends=backends)
@ -126,7 +125,9 @@ class MemorySystemService:
try:
return BackendStatus(status="success", result=await call())
except Exception as exc: # noqa: BLE001
return BackendStatus(status="failed", error=str(exc))
message = str(exc)
error = f"{type(exc).__name__}: {message}" if message else type(exc).__name__
return BackendStatus(status="failed", error=error)
def _aggregate_status(self, backends: dict[str, BackendStatus]) -> str:
statuses = {backend.status for backend in backends.values()}
@ -163,3 +164,16 @@ class MemorySystemService:
if "source_backend" in item:
return item
return {"source_backend": backend_name, **item}
def _remove_vectors_from_backends(self, backends: dict[str, BackendStatus]) -> dict[str, BackendStatus]:
return {
name: backend.model_copy(update={"result": self._remove_vectors(backend.result)})
for name, backend in backends.items()
}
def _remove_vectors(self, value: Any) -> Any:
if isinstance(value, dict):
return {key: self._remove_vectors(item) for key, item in value.items() if key != "vector"}
if isinstance(value, list):
return [self._remove_vectors(item) for item in value]
return value