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

Refactor OpenViking Memory API and User Management

Browse Source 
- Updated API authentication headers to use `X-API-Key` for both admin and user APIs.
- Modified the account creation process to directly create user-specific accounts without requiring an admin workspace.
- Enhanced user creation to return account-specific details, including `admin_user_id`.
- Introduced new endpoints for retrieving task status and user profiles, allowing for more flexible user data management.
- Updated search functionality to include additional parameters such as `level` and `score_threshold`.
- Improved the handling of user keys in the storage layer to associate them with specific accounts.
- Added tests to validate the new user account creation process and search functionalities, ensuring proper integration with the OpenViking service.
- Included new documentation to reflect changes in API usage and expected request/response formats.
This commit is contained in:
tomtan
2026-05-27 16:09:28 +08:00
parent a89807b174
commit 70cda923b2
13 changed files with 543 additions and 165 deletions
Download Patch File Download Diff File Expand all files Collapse all files

91
memory_system_api/clients.py
View File

@ -49,36 +49,24 @@ class OpenVikingMemorySystemClient:
user_key = self._extract_user_key(data)
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)
self.store.save_user_key(admin_user_id, user_key, account_id=account_id)
return data
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)
account_id = self._account_id_for_user(user_id)
if existing:
return {"status": "ok", "result": {"account_id": ADMIN_ACCOUNT_ID, "user_id": user_id, "user_key": existing}}
return {
"status": "ok",
"result": {
"account_id": account_id,
"admin_user_id": user_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
return await self.create_account(account_id, user_id)
def credential_for_user(
self,
@ -98,7 +86,7 @@ class OpenVikingMemorySystemClient:
) -> OpenVikingCredential:
return OpenVikingCredential(
api_key=user_key,
account_id=ADMIN_ACCOUNT_ID,
account_id=self._account_id_for_user(user_id),
user_id=user_id,
agent_id=agent_id,
user_key_auth=True,
@ -164,22 +152,46 @@ class OpenVikingMemorySystemClient:
return response.json()
async def search(
self, credential: OpenVikingCredential | str, session_id: str | None, query: str, limit: int
self,
credential: OpenVikingCredential | str,
query: str,
limit: int,
level: int = 2,
score_threshold: float = 0.8,
target_uri: str = "viking://user/memories",
) -> dict[str, Any]:
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/"
)
payload: dict[str, Any] = {
"query": query,
"target_uri": target_uri,
"limit": limit,
"level": level,
"score_threshold": score_threshold,
}
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()
async def search_profile_memories(
self,
credential: OpenVikingCredential | str,
query: str,
limit: int,
level: int,
) -> dict[str, Any]:
async with self._credential_client(credential) as client:
response = await client.post(
"/api/v1/search/search",
json={
"query": query,
"limit": limit,
"level": level,
"target_uri": "viking://user/memories",
},
)
response.raise_for_status()
return response.json()
async def get_session_context(self, credential: OpenVikingCredential | str, session_id: str) -> dict[str, Any]:
async with self._credential_client(credential) as client:
response = await client.get(f"/api/v1/sessions/{session_id}/context")
@ -201,11 +213,7 @@ class OpenVikingMemorySystemClient:
return self._client(credential.api_key, headers)
def _client(self, api_key: str, extra_headers: dict[str, str] | None = None) -> httpx.AsyncClient:
headers = {"Content-Type": "application/json"}
if api_key == self.root_key:
headers["X-API-Key"] = api_key
else:
headers["Authorization"] = f"Bearer {api_key}"
headers = {"Content-Type": "application/json", "X-API-Key": api_key}
if extra_headers:
headers.update(extra_headers)
return httpx.AsyncClient(
@ -226,6 +234,9 @@ class OpenVikingMemorySystemClient:
"admin_user_id": admin_user_id,
}
def _account_id_for_user(self, user_id: str) -> str:
return f"{user_id}_account"
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)