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

160
tests/test_memory_system_clients.py
View File

@ -13,7 +13,7 @@ class FakeStore:
def get_user_key(self, user_id: str) -> str | None:
return self.users.get(user_id)
def save_user_key(self, user_id: str, user_key: str) -> None:
def save_user_key(self, user_id: str, user_key: str, account_id: str = "admin") -> None:
self.users[user_id] = user_key
def get_account_key(self, account_id: str) -> str | None:
@ -91,12 +91,24 @@ def test_openviking_accepts_matching_user_credentials():
credential = client.credential_for_user("tom", "tom-key", agent_id="sess-1")
assert credential.api_key == "tom-key"
assert credential.account_id == "admin"
assert credential.account_id == "tom_account"
assert credential.user_id == "tom"
assert credential.agent_id == "sess-1"
def test_openviking_create_user_initializes_admin_workspace_first():
def test_openviking_client_uses_x_api_key_for_user_keys():
client = OpenVikingMemorySystemClient(store=FakeStore())
client.root_key = "root-key"
http_client = client._client("tom-key")
try:
assert http_client.headers["X-API-Key"] == "tom-key"
assert "Authorization" not in http_client.headers
finally:
asyncio.run(http_client.aclose())
def test_openviking_create_user_creates_isolated_admin_account():
store = FakeStore()
client = OpenVikingMemorySystemClient(store=store)
client.root_key = "root-key"
@ -104,11 +116,14 @@ def test_openviking_create_user_initializes_admin_workspace_first():
responses = [
FakeResponse(
200,
{"status": "ok", "result": {"account_id": "admin", "admin_user_id": "admin", "user_key": "admin-key"}},
),
FakeResponse(
200,
{"status": "ok", "result": {"account_id": "admin", "user_id": "userA", "user_key": "userA-key"}},
{
"status": "ok",
"result": {
"account_id": "userA_account",
"admin_user_id": "userA",
"user_key": "userA-key",
},
},
),
]
client._client = lambda api_key, extra_headers=None: FakeAsyncClient( # type: ignore[method-assign]
@ -120,28 +135,28 @@ def test_openviking_create_user_initializes_admin_workspace_first():
result = asyncio.run(client.create_user("userA"))
assert result == {"status": "ok", "result": {"account_id": "admin", "user_id": "userA", "user_key": "userA-key"}}
assert store.accounts == {"admin": "admin-key"}
assert store.users == {"admin": "admin-key", "userA": "userA-key"}
assert result == {
"status": "ok",
"result": {
"account_id": "userA_account",
"admin_user_id": "userA",
"user_key": "userA-key",
},
}
assert store.accounts == {"userA_account": "userA-key"}
assert store.users == {"userA": "userA-key"}
assert calls == [
(
"post",
"root-key",
{},
"/api/v1/admin/accounts",
{"account_id": "admin", "admin_user_id": "admin"},
),
(
"post",
"root-key",
{},
"/api/v1/admin/accounts/admin/users",
{"user_id": "userA", "role": "user"},
{"account_id": "userA_account", "admin_user_id": "userA"},
),
]
def test_openviking_create_user_reuses_existing_admin_workspace():
def test_openviking_create_user_creates_account_even_when_admin_workspace_exists():
store = FakeStore()
store.save_account_key("admin", "admin", "admin-key")
client = OpenVikingMemorySystemClient(store=store)
@ -150,7 +165,14 @@ def test_openviking_create_user_reuses_existing_admin_workspace():
responses = [
FakeResponse(
200,
{"status": "ok", "result": {"account_id": "admin", "user_id": "userB", "user_key": "userB-key"}},
{
"status": "ok",
"result": {
"account_id": "userB_account",
"admin_user_id": "userB",
"user_key": "userB-key",
},
},
)
]
client._client = lambda api_key, extra_headers=None: FakeAsyncClient( # type: ignore[method-assign]
@ -162,15 +184,23 @@ def test_openviking_create_user_reuses_existing_admin_workspace():
result = asyncio.run(client.create_user("userB"))
assert result == {"status": "ok", "result": {"account_id": "admin", "user_id": "userB", "user_key": "userB-key"}}
assert result == {
"status": "ok",
"result": {
"account_id": "userB_account",
"admin_user_id": "userB",
"user_key": "userB-key",
},
}
assert store.accounts == {"admin": "admin-key", "userB_account": "userB-key"}
assert store.users == {"userB": "userB-key"}
assert calls == [
(
"post",
"root-key",
{},
"/api/v1/admin/accounts/admin/users",
{"user_id": "userB", "role": "user"},
"/api/v1/admin/accounts",
{"account_id": "userB_account", "admin_user_id": "userB"},
)
]
@ -229,7 +259,7 @@ def test_openviking_find_uses_current_identity_memory_scope():
]
def test_openviking_search_uses_session_target_uri():
def test_openviking_search_uses_fixed_user_memory_target_with_level_and_score_threshold():
client = OpenVikingMemorySystemClient(store=FakeStore())
calls = []
responses = [FakeResponse(200, {"status": "ok", "result": {"memories": []}})]
@ -241,7 +271,7 @@ def test_openviking_search_uses_session_target_uri():
)
credential = client.user_credential("tom-key", "tom", agent_id="sess-1")
result = asyncio.run(client.search(credential, "sess-1", "咖啡", 5))
result = asyncio.run(client.search(credential, "咖啡", 5, level=3, score_threshold=0.7))
assert result == {"status": "ok", "result": {"memories": []}}
assert calls == [
@ -250,7 +280,83 @@ def test_openviking_search_uses_session_target_uri():
"tom-key",
{},
"/api/v1/search/search",
{"query": "咖啡", "limit": 5, "session_id": "sess-1", "target_uri": "viking://user/tom/sess-1"},
{
"query": "咖啡",
"target_uri": "viking://user/memories",
"limit": 5,
"level": 3,
"score_threshold": 0.7,
},
)
]
def test_openviking_search_accepts_custom_target_uri():
client = OpenVikingMemorySystemClient(store=FakeStore())
calls = []
responses = [FakeResponse(200, {"status": "ok", "result": {"memories": []}})]
client._client = lambda api_key, extra_headers=None: FakeAsyncClient( # type: ignore[method-assign]
calls,
responses,
api_key,
extra_headers or {},
)
credential = client.user_credential("tom-key", "tom", agent_id="sess-1")
result = asyncio.run(client.search(
credential,
"咖啡",
5,
level=3,
score_threshold=0.7,
target_uri="viking://user/custom/memories",
))
assert result == {"status": "ok", "result": {"memories": []}}
assert calls == [
(
"post",
"tom-key",
{},
"/api/v1/search/search",
{
"query": "咖啡",
"target_uri": "viking://user/custom/memories",
"limit": 5,
"level": 3,
"score_threshold": 0.7,
},
)
]
def test_openviking_profile_search_uses_user_memory_target_and_level():
client = OpenVikingMemorySystemClient(store=FakeStore())
calls = []
responses = [FakeResponse(200, {"status": "ok", "result": {"memories": []}})]
client._client = lambda api_key, extra_headers=None: FakeAsyncClient( # type: ignore[method-assign]
calls,
responses,
api_key,
extra_headers or {},
)
credential = client.user_credential("tom-key", "tom")
result = asyncio.run(client.search_profile_memories(credential, "我想喝东西", 10, 2))
assert result == {"status": "ok", "result": {"memories": []}}
assert calls == [
(
"post",
"tom-key",
{},
"/api/v1/search/search",
{
"query": "我想喝东西",
"limit": 10,
"level": 2,
"target_uri": "viking://user/memories",
},
)
]