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

Refine memory system user-key flow and search output

Browse Source
This commit is contained in:
tomtan
2026-05-22 16:30:42 +08:00
parent 92632553ab
commit d73f59f38d
16 changed files with 1888 additions and 255 deletions
Download Patch File Download Diff File Expand all files Collapse all files

142
memory_system_api/clients.py
View File

@ -8,7 +8,7 @@ from typing import Any
import httpx
from .config import get_config
from .store import OpenVikingUserKeyStore
from .store import ADMIN_ACCOUNT_ID, ADMIN_USER_ID, OpenVikingUserKeyStore
@dataclass(frozen=True)
@ -16,6 +16,8 @@ class OpenVikingCredential:
api_key: str
account_id: str | None = None
user_id: str | None = None
agent_id: str | None = None
user_key_auth: bool = False
class OpenVikingMemorySystemClient:
@ -33,36 +35,83 @@ class OpenVikingMemorySystemClient:
response.raise_for_status()
return response.json()
async def ensure_user(self, user_id: str) -> OpenVikingCredential:
existing = self.store.get_user_key(user_id)
if existing:
return OpenVikingCredential(api_key=existing)
async def create_account(self, account_id: str = ADMIN_ACCOUNT_ID, admin_user_id: str = ADMIN_USER_ID) -> dict[str, Any]:
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},
json=self._create_account_payload(account_id, admin_user_id),
)
if response.status_code == 409:
return self.root_credential(user_id)
return response.json()
response.raise_for_status()
data = response.json()
user_key = self._extract_user_key(data)
if not user_key:
return self.root_credential(user_id)
self.store.save_user_key(user_id, user_key)
return OpenVikingCredential(api_key=user_key)
if user_key:
self.store.save_account_key(account_id, admin_user_id, user_key)
self.store.save_user_key(admin_user_id, user_key)
return data
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_admin_workspace(self) -> None:
if self.store.get_account_key(ADMIN_ACCOUNT_ID):
return
await self.create_account(ADMIN_ACCOUNT_ID, ADMIN_USER_ID)
async def create_user(self, user_id: str, role: str = "user") -> dict[str, Any]:
existing = self.store.get_user_key(user_id)
if existing:
return {"status": "ok", "result": {"account_id": ADMIN_ACCOUNT_ID, "user_id": user_id, "user_key": existing}}
await self.ensure_admin_workspace()
if user_id == ADMIN_USER_ID:
user_key = self.store.get_user_key(user_id)
return {"status": "ok", "result": {"account_id": ADMIN_ACCOUNT_ID, "user_id": user_id, "user_key": user_key}}
async with self._client(self.root_key) as client:
response = await client.post(
f"/api/v1/admin/accounts/{ADMIN_ACCOUNT_ID}/users",
json={"user_id": user_id, "role": role},
)
response.raise_for_status()
data = response.json()
user_key = self._extract_user_key(data)
if user_key:
self.store.save_user_key(user_id, user_key)
return data
def credential_for_user(
self,
user_id: str,
user_key: str,
agent_id: str | None = None,
) -> OpenVikingCredential:
if not self.store.user_key_matches(user_id, user_key):
raise PermissionError("Invalid user key")
return self.user_credential(user_key, user_id, agent_id=agent_id)
def user_credential(
self,
user_key: str,
user_id: str,
agent_id: str | None = None,
) -> OpenVikingCredential:
return OpenVikingCredential(
api_key=user_key,
account_id=ADMIN_ACCOUNT_ID,
user_id=user_id,
agent_id=agent_id,
user_key_auth=True,
)
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}:
self._save_session(credential, session_id)
return {"session_id": session_id, "status": "exists"}
response.raise_for_status()
self._save_session(credential, session_id)
return response.json()
async def append_message(
@ -78,9 +127,14 @@ class OpenVikingMemorySystemClient:
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 = await client.post(
f"/api/v1/sessions/{session_id}/commit",
json={"keep_recent_count": 0},
)
response.raise_for_status()
return response.json()
data = response.json()
self._save_commit_metadata(credential, session_id, data)
return data
async def extract_session(self, credential: OpenVikingCredential | str, session_id: str) -> dict[str, Any]:
async with self._credential_client(credential) as client:
@ -94,15 +148,15 @@ class OpenVikingMemorySystemClient:
response.raise_for_status()
return response.json()
async def find(
self, credential: OpenVikingCredential | str, user_id: str, query: str, limit: int
) -> dict[str, Any]:
async def find(self, credential: OpenVikingCredential | str, query: str, limit: int) -> dict[str, Any]:
user_id = credential.user_id if isinstance(credential, OpenVikingCredential) else None
target_uri = f"viking://user/{user_id}/memories/" if user_id else "viking://user/memories/"
async with self._credential_client(credential) as client:
response = await client.post(
"/api/v1/search/find",
json={
"query": query,
"target_uri": f"viking://user/{user_id}/memories/",
"target_uri": target_uri,
"limit": limit,
},
)
@ -115,6 +169,12 @@ class OpenVikingMemorySystemClient:
payload: dict[str, Any] = {"query": query, "limit": limit}
if session_id:
payload["session_id"] = session_id
if isinstance(credential, OpenVikingCredential) and credential.user_id:
payload["target_uri"] = (
f"viking://user/{credential.user_id}/{session_id}"
if session_id
else f"viking://user/{credential.user_id}/memories/"
)
async with self._credential_client(credential) as client:
response = await client.post("/api/v1/search/search", json=payload)
response.raise_for_status()
@ -123,15 +183,23 @@ class OpenVikingMemorySystemClient:
def _credential_client(self, credential: OpenVikingCredential | str) -> httpx.AsyncClient:
if isinstance(credential, str):
return self._client(credential)
if credential.user_key_auth:
return self._client(credential.api_key)
headers = {}
if credential.account_id:
headers["X-OpenViking-Account"] = credential.account_id
if credential.user_id:
headers["X-OpenViking-User"] = credential.user_id
if credential.agent_id:
headers["X-OpenViking-Agent"] = credential.agent_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"}
headers = {"Content-Type": "application/json"}
if api_key == self.root_key:
headers["X-API-Key"] = api_key
else:
headers["Authorization"] = f"Bearer {api_key}"
if extra_headers:
headers.update(extra_headers)
return httpx.AsyncClient(
@ -146,6 +214,38 @@ class OpenVikingMemorySystemClient:
value = result.get("user_key") if isinstance(result, dict) else None
return str(value) if value else None
def _create_account_payload(self, account_id: str, admin_user_id: str) -> dict[str, Any]:
return {
"account_id": account_id,
"admin_user_id": admin_user_id,
}
def _save_session(self, credential: OpenVikingCredential | str, session_id: str) -> None:
if isinstance(credential, OpenVikingCredential) and credential.user_id:
self.store.save_session(credential.user_id, session_id)
def _save_commit_metadata(
self,
credential: OpenVikingCredential | str,
session_id: str,
data: dict[str, Any],
) -> None:
if not isinstance(credential, OpenVikingCredential) or not credential.user_id:
return
result = data.get("result") if isinstance(data.get("result"), dict) else data
if not isinstance(result, dict):
return
task_id = result.get("task_id")
if not task_id:
return
archive_uri = result.get("archive_uri")
self.store.save_task(
user_id=credential.user_id,
session_id=session_id,
task_id=str(task_id),
archive_uri=str(archive_uri) if archive_uri else None,
)
class EverOSMemorySystemClient:
def __init__(self) -> None: