Refine memory system user-key flow and search output
This commit is contained in:
@ -5,13 +5,34 @@ from memory_system_api.clients import EverOSMemorySystemClient, OpenVikingMemory
|
||||
|
||||
class FakeStore:
|
||||
def __init__(self):
|
||||
self.saved = {}
|
||||
self.users = {}
|
||||
self.accounts = {}
|
||||
self.sessions = []
|
||||
self.tasks = []
|
||||
|
||||
def get_user_key(self, user_id: str) -> str | None:
|
||||
return self.saved.get(user_id)
|
||||
return self.users.get(user_id)
|
||||
|
||||
def save_user_key(self, user_id: str, user_key: str) -> None:
|
||||
self.saved[user_id] = user_key
|
||||
self.users[user_id] = user_key
|
||||
|
||||
def get_account_key(self, account_id: str) -> str | None:
|
||||
return self.accounts.get(account_id)
|
||||
|
||||
def save_account_key(self, account_id: str, admin_user_id: str, account_key: str) -> None:
|
||||
self.accounts[account_id] = account_key
|
||||
|
||||
def account_key_matches(self, account_id: str, account_key: str) -> bool:
|
||||
return self.accounts.get(account_id) == account_key
|
||||
|
||||
def user_key_matches(self, user_id: str, user_key: str) -> bool:
|
||||
return self.users.get(user_id) == user_key
|
||||
|
||||
def save_session(self, user_id: str, session_id: str) -> None:
|
||||
self.sessions.append((user_id, session_id))
|
||||
|
||||
def save_task(self, user_id: str, session_id: str, task_id: str, archive_uri: str | None) -> None:
|
||||
self.tasks.append((user_id, session_id, task_id, archive_uri))
|
||||
|
||||
|
||||
class FakeResponse:
|
||||
@ -44,13 +65,52 @@ class FakeAsyncClient:
|
||||
self.calls.append(("post", self.api_key, self.headers, path, json))
|
||||
return self.responses.pop(0)
|
||||
|
||||
async def get(self, path: str) -> FakeResponse:
|
||||
self.calls.append(("get", self.api_key, self.headers, path, None))
|
||||
return self.responses.pop(0)
|
||||
|
||||
def test_openviking_uses_root_identity_when_account_already_exists():
|
||||
|
||||
def test_openviking_rejects_unknown_user_credentials():
|
||||
store = FakeStore()
|
||||
client = OpenVikingMemorySystemClient(store=store)
|
||||
|
||||
try:
|
||||
client.credential_for_user("tom", "missing-key")
|
||||
except PermissionError as exc:
|
||||
assert "Invalid user key" in str(exc)
|
||||
else:
|
||||
raise AssertionError("expected PermissionError")
|
||||
|
||||
|
||||
def test_openviking_accepts_matching_user_credentials():
|
||||
store = FakeStore()
|
||||
store.save_user_key("tom", "tom-key")
|
||||
client = OpenVikingMemorySystemClient(store=store)
|
||||
client.root_key = "root-key"
|
||||
|
||||
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.user_id == "tom"
|
||||
assert credential.agent_id == "sess-1"
|
||||
|
||||
|
||||
def test_openviking_create_user_initializes_admin_workspace_first():
|
||||
store = FakeStore()
|
||||
client = OpenVikingMemorySystemClient(store=store)
|
||||
client.root_key = "root-key"
|
||||
calls = []
|
||||
responses = [FakeResponse(409, {"status": "error", "error": {"code": "CONFLICT"}})]
|
||||
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"}},
|
||||
),
|
||||
]
|
||||
client._client = lambda api_key, extra_headers=None: FakeAsyncClient( # type: ignore[method-assign]
|
||||
calls,
|
||||
responses,
|
||||
@ -58,15 +118,64 @@ def test_openviking_uses_root_identity_when_account_already_exists():
|
||||
extra_headers or {},
|
||||
)
|
||||
|
||||
credential = asyncio.run(client.ensure_user("tom"))
|
||||
result = asyncio.run(client.create_user("userA"))
|
||||
|
||||
assert credential.api_key == "root-key"
|
||||
assert credential.account_id == "tom"
|
||||
assert credential.user_id == "tom"
|
||||
assert store.saved == {}
|
||||
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 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"},
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
def test_openviking_root_identity_headers_are_sent_for_session_create():
|
||||
def test_openviking_create_user_reuses_existing_admin_workspace():
|
||||
store = FakeStore()
|
||||
store.save_account_key("admin", "admin", "admin-key")
|
||||
client = OpenVikingMemorySystemClient(store=store)
|
||||
client.root_key = "root-key"
|
||||
calls = []
|
||||
responses = [
|
||||
FakeResponse(
|
||||
200,
|
||||
{"status": "ok", "result": {"account_id": "admin", "user_id": "userB", "user_key": "userB-key"}},
|
||||
)
|
||||
]
|
||||
client._client = lambda api_key, extra_headers=None: FakeAsyncClient( # type: ignore[method-assign]
|
||||
calls,
|
||||
responses,
|
||||
api_key,
|
||||
extra_headers or {},
|
||||
)
|
||||
|
||||
result = asyncio.run(client.create_user("userB"))
|
||||
|
||||
assert result == {"status": "ok", "result": {"account_id": "admin", "user_id": "userB", "user_key": "userB-key"}}
|
||||
assert store.users == {"userB": "userB-key"}
|
||||
assert calls == [
|
||||
(
|
||||
"post",
|
||||
"root-key",
|
||||
{},
|
||||
"/api/v1/admin/accounts/admin/users",
|
||||
{"user_id": "userB", "role": "user"},
|
||||
)
|
||||
]
|
||||
|
||||
|
||||
def test_openviking_user_key_auth_is_used_for_session_create():
|
||||
client = OpenVikingMemorySystemClient(store=FakeStore())
|
||||
client.root_key = "root-key"
|
||||
calls = []
|
||||
@ -77,22 +186,121 @@ def test_openviking_root_identity_headers_are_sent_for_session_create():
|
||||
api_key,
|
||||
extra_headers or {},
|
||||
)
|
||||
credential = client.root_credential("tom")
|
||||
credential = client.user_credential("tom-key", "tom", agent_id="sess-2")
|
||||
|
||||
result = asyncio.run(client.ensure_session(credential, "sess-2"))
|
||||
|
||||
assert result == {"status": "ok", "result": {"session_id": "sess-2"}}
|
||||
assert client.store.sessions == [("tom", "sess-2")]
|
||||
assert calls == [
|
||||
(
|
||||
"post",
|
||||
"root-key",
|
||||
{"X-OpenViking-Account": "tom", "X-OpenViking-User": "tom"},
|
||||
"tom-key",
|
||||
{},
|
||||
"/api/v1/sessions",
|
||||
{"session_id": "sess-2"},
|
||||
)
|
||||
]
|
||||
|
||||
|
||||
def test_openviking_find_uses_current_identity_memory_scope():
|
||||
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.find(credential, "咖啡", 5))
|
||||
|
||||
assert result == {"status": "ok", "result": {"memories": []}}
|
||||
assert calls == [
|
||||
(
|
||||
"post",
|
||||
"tom-key",
|
||||
{},
|
||||
"/api/v1/search/find",
|
||||
{"query": "咖啡", "target_uri": "viking://user/tom/memories/", "limit": 5},
|
||||
)
|
||||
]
|
||||
|
||||
|
||||
def test_openviking_search_uses_session_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, "sess-1", "咖啡", 5))
|
||||
|
||||
assert result == {"status": "ok", "result": {"memories": []}}
|
||||
assert calls == [
|
||||
(
|
||||
"post",
|
||||
"tom-key",
|
||||
{},
|
||||
"/api/v1/search/search",
|
||||
{"query": "咖啡", "limit": 5, "session_id": "sess-1", "target_uri": "viking://user/tom/sess-1"},
|
||||
)
|
||||
]
|
||||
|
||||
|
||||
def test_openviking_commit_keeps_no_recent_live_messages():
|
||||
client = OpenVikingMemorySystemClient(store=FakeStore())
|
||||
calls = []
|
||||
responses = [
|
||||
FakeResponse(
|
||||
200,
|
||||
{
|
||||
"status": "ok",
|
||||
"result": {
|
||||
"status": "accepted",
|
||||
"task_id": "task-1",
|
||||
"archive_uri": "viking://session/tom/sess-1/history/archive_001",
|
||||
},
|
||||
},
|
||||
)
|
||||
]
|
||||
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.commit_session(credential, "sess-1"))
|
||||
|
||||
assert result == {
|
||||
"status": "ok",
|
||||
"result": {
|
||||
"status": "accepted",
|
||||
"task_id": "task-1",
|
||||
"archive_uri": "viking://session/tom/sess-1/history/archive_001",
|
||||
},
|
||||
}
|
||||
assert client.store.tasks == [("tom", "sess-1", "task-1", "viking://session/tom/sess-1/history/archive_001")]
|
||||
assert calls == [
|
||||
(
|
||||
"post",
|
||||
"tom-key",
|
||||
{},
|
||||
"/api/v1/sessions/sess-1/commit",
|
||||
{"keep_recent_count": 0},
|
||||
)
|
||||
]
|
||||
|
||||
|
||||
def test_everos_assistant_payload_does_not_use_user_id_as_sender():
|
||||
client = EverOSMemorySystemClient()
|
||||
|
||||
|
||||
Reference in New Issue
Block a user