606 lines
18 KiB
Python
606 lines
18 KiB
Python
import asyncio
|
|
|
|
from memory_system_api.clients import EverOSMemorySystemClient, OpenVikingMemorySystemClient
|
|
|
|
|
|
class FakeStore:
|
|
def __init__(self):
|
|
self.users = {}
|
|
self.accounts = {}
|
|
self.sessions = []
|
|
self.tasks = []
|
|
|
|
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, account_id: str = "admin") -> None:
|
|
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:
|
|
def __init__(self, status_code: int, data: dict):
|
|
self.status_code = status_code
|
|
self._data = data
|
|
|
|
def json(self) -> dict:
|
|
return self._data
|
|
|
|
def raise_for_status(self) -> None:
|
|
if self.status_code >= 400:
|
|
raise AssertionError(f"unexpected status {self.status_code}")
|
|
|
|
|
|
class FakeAsyncClient:
|
|
def __init__(self, calls: list, responses: list[FakeResponse], api_key: str, headers: dict):
|
|
self.calls = calls
|
|
self.responses = responses
|
|
self.api_key = api_key
|
|
self.headers = headers
|
|
|
|
async def __aenter__(self):
|
|
return self
|
|
|
|
async def __aexit__(self, exc_type, exc, tb):
|
|
return False
|
|
|
|
async def post(self, path: str, json: dict | None = None, files: dict | None = None) -> FakeResponse:
|
|
if files and "file" in files:
|
|
uploaded = files["file"]
|
|
files = {"file": uploaded[0] if isinstance(uploaded, tuple) else uploaded}
|
|
self.calls.append(("post", self.api_key, self.headers, path, json, files))
|
|
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)
|
|
|
|
async def delete(self, path: str, params: dict | None = None) -> FakeResponse:
|
|
self.calls.append(("delete", self.api_key, self.headers, path, params))
|
|
return self.responses.pop(0)
|
|
|
|
|
|
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 == "tom_account"
|
|
assert credential.user_id == "tom"
|
|
assert credential.agent_id == "sess-1"
|
|
|
|
|
|
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"
|
|
calls = []
|
|
responses = [
|
|
FakeResponse(
|
|
200,
|
|
{
|
|
"status": "ok",
|
|
"result": {
|
|
"account_id": "userA_account",
|
|
"admin_user_id": "userA",
|
|
"user_key": "userA-key",
|
|
},
|
|
},
|
|
),
|
|
]
|
|
client._client = lambda api_key, extra_headers=None, json_content_type=True: FakeAsyncClient( # type: ignore[method-assign]
|
|
calls,
|
|
responses,
|
|
api_key,
|
|
extra_headers or {},
|
|
)
|
|
|
|
result = asyncio.run(client.create_user("userA"))
|
|
|
|
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": "userA_account", "admin_user_id": "userA"},
|
|
None,
|
|
),
|
|
]
|
|
|
|
|
|
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)
|
|
client.root_key = "root-key"
|
|
calls = []
|
|
responses = [
|
|
FakeResponse(
|
|
200,
|
|
{
|
|
"status": "ok",
|
|
"result": {
|
|
"account_id": "userB_account",
|
|
"admin_user_id": "userB",
|
|
"user_key": "userB-key",
|
|
},
|
|
},
|
|
)
|
|
]
|
|
client._client = lambda api_key, extra_headers=None, json_content_type=True: 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": "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",
|
|
{"account_id": "userB_account", "admin_user_id": "userB"},
|
|
None,
|
|
)
|
|
]
|
|
|
|
|
|
def test_openviking_user_key_auth_is_used_for_session_create():
|
|
client = OpenVikingMemorySystemClient(store=FakeStore())
|
|
client.root_key = "root-key"
|
|
calls = []
|
|
responses = [FakeResponse(200, {"status": "ok", "result": {"session_id": "sess-2"}})]
|
|
client._client = lambda api_key, extra_headers=None, json_content_type=True: FakeAsyncClient( # type: ignore[method-assign]
|
|
calls,
|
|
responses,
|
|
api_key,
|
|
extra_headers or {},
|
|
)
|
|
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",
|
|
"tom-key",
|
|
{},
|
|
"/api/v1/sessions",
|
|
{"session_id": "sess-2"},
|
|
None,
|
|
)
|
|
]
|
|
|
|
|
|
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, json_content_type=True: 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},
|
|
None,
|
|
)
|
|
]
|
|
|
|
|
|
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": []}})]
|
|
client._client = lambda api_key, extra_headers=None, json_content_type=True: 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))
|
|
|
|
assert result == {"status": "ok", "result": {"memories": []}}
|
|
assert calls == [
|
|
(
|
|
"post",
|
|
"tom-key",
|
|
{},
|
|
"/api/v1/search/search",
|
|
{
|
|
"query": "咖啡",
|
|
"target_uri": "viking://user/memories",
|
|
"limit": 5,
|
|
"level": 3,
|
|
"score_threshold": 0.7,
|
|
},
|
|
None,
|
|
)
|
|
]
|
|
|
|
|
|
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, json_content_type=True: 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,
|
|
},
|
|
None,
|
|
)
|
|
]
|
|
|
|
|
|
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, json_content_type=True: 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",
|
|
},
|
|
None,
|
|
)
|
|
]
|
|
|
|
|
|
def test_openviking_get_session_context_uses_user_key_auth():
|
|
client = OpenVikingMemorySystemClient(store=FakeStore())
|
|
calls = []
|
|
responses = [
|
|
FakeResponse(
|
|
200,
|
|
{
|
|
"status": "ok",
|
|
"result": {
|
|
"latest_archive_overview": "# Working Memory",
|
|
"messages": [],
|
|
},
|
|
},
|
|
)
|
|
]
|
|
client._client = lambda api_key, extra_headers=None, json_content_type=True: 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.get_session_context(credential, "sess-1"))
|
|
|
|
assert result == {
|
|
"status": "ok",
|
|
"result": {
|
|
"latest_archive_overview": "# Working Memory",
|
|
"messages": [],
|
|
},
|
|
}
|
|
assert calls == [
|
|
(
|
|
"get",
|
|
"tom-key",
|
|
{},
|
|
"/api/v1/sessions/sess-1/context",
|
|
None,
|
|
)
|
|
]
|
|
|
|
|
|
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, json_content_type=True: 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},
|
|
None,
|
|
)
|
|
]
|
|
|
|
|
|
def test_openviking_upload_temp_file_posts_multipart(tmp_path):
|
|
path = tmp_path / "report.pdf"
|
|
path.write_bytes(b"pdf-bytes")
|
|
client = OpenVikingMemorySystemClient(store=FakeStore())
|
|
calls = []
|
|
responses = [
|
|
FakeResponse(
|
|
200,
|
|
{"status": "ok", "result": {"temp_file_id": "upload_report.pdf"}},
|
|
)
|
|
]
|
|
client._client = lambda api_key, extra_headers=None, json_content_type=True: FakeAsyncClient( # type: ignore[method-assign]
|
|
calls,
|
|
responses,
|
|
api_key,
|
|
extra_headers or {},
|
|
)
|
|
credential = client.user_credential("tom-key", "tom")
|
|
|
|
result = asyncio.run(client.upload_temp_file(credential, path))
|
|
|
|
assert result == {"status": "ok", "result": {"temp_file_id": "upload_report.pdf"}}
|
|
assert calls == [
|
|
(
|
|
"post",
|
|
"tom-key",
|
|
{},
|
|
"/api/v1/resources/temp_upload",
|
|
None,
|
|
{"file": "report.pdf"},
|
|
)
|
|
]
|
|
|
|
|
|
def test_openviking_add_resource_posts_url_payload():
|
|
client = OpenVikingMemorySystemClient(store=FakeStore())
|
|
calls = []
|
|
responses = [FakeResponse(200, {"status": "ok", "result": {"uri": "viking://resources/tom/images/photo.png"}})]
|
|
client._client = lambda api_key, extra_headers=None, json_content_type=True: FakeAsyncClient( # type: ignore[method-assign]
|
|
calls,
|
|
responses,
|
|
api_key,
|
|
extra_headers or {},
|
|
)
|
|
credential = client.user_credential("tom-key", "tom")
|
|
|
|
result = asyncio.run(
|
|
client.add_resource(
|
|
credential,
|
|
path="https://example.com/photo.png",
|
|
to="viking://resources/tom/images/photo.png",
|
|
reason="上传远程图片",
|
|
wait=True,
|
|
directly_upload_media=True,
|
|
)
|
|
)
|
|
|
|
assert result == {"status": "ok", "result": {"uri": "viking://resources/tom/images/photo.png"}}
|
|
assert calls == [
|
|
(
|
|
"post",
|
|
"tom-key",
|
|
{},
|
|
"/api/v1/resources",
|
|
{
|
|
"path": "https://example.com/photo.png",
|
|
"to": "viking://resources/tom/images/photo.png",
|
|
"reason": "上传远程图片",
|
|
"wait": True,
|
|
"directly_upload_media": True,
|
|
},
|
|
None,
|
|
)
|
|
]
|
|
|
|
|
|
def test_openviking_delete_resource_sends_uri_and_recursive_flag():
|
|
client = OpenVikingMemorySystemClient(store=FakeStore())
|
|
calls = []
|
|
responses = [FakeResponse(200, {"status": "ok", "result": {"estimated_deleted_count": 4}})]
|
|
client._client = lambda api_key, extra_headers=None, json_content_type=True: FakeAsyncClient( # type: ignore[method-assign]
|
|
calls,
|
|
responses,
|
|
api_key,
|
|
extra_headers or {},
|
|
)
|
|
credential = client.user_credential("tom-key", "tom")
|
|
|
|
result = asyncio.run(
|
|
client.delete_resource(
|
|
credential,
|
|
uri="viking://resources/tom/files/report.pdf",
|
|
recursive=True,
|
|
)
|
|
)
|
|
|
|
assert result == {"status": "ok", "result": {"estimated_deleted_count": 4}}
|
|
assert calls == [
|
|
(
|
|
"delete",
|
|
"tom-key",
|
|
{},
|
|
"/api/v1/fs",
|
|
{"uri": "viking://resources/tom/files/report.pdf", "recursive": "true"},
|
|
)
|
|
]
|
|
|
|
|
|
def test_everos_assistant_payload_does_not_use_user_id_as_sender():
|
|
client = EverOSMemorySystemClient()
|
|
|
|
payload = client.build_message_payload(
|
|
user_id="tom",
|
|
session_id="sess-1",
|
|
role="assistant",
|
|
content="我记住了",
|
|
)
|
|
|
|
message = payload["messages"][0]
|
|
assert message["role"] == "assistant"
|
|
assert message["sender_id"] != "tom"
|
|
assert message["sender_name"] != "tom"
|
|
|
|
|
|
def test_everos_user_payload_uses_user_id_as_sender():
|
|
client = EverOSMemorySystemClient()
|
|
|
|
payload = client.build_message_payload(
|
|
user_id="tom",
|
|
session_id="sess-1",
|
|
role="user",
|
|
content="我喜欢拿铁",
|
|
)
|
|
|
|
message = payload["messages"][0]
|
|
assert message["role"] == "user"
|
|
assert message["sender_id"] == "tom"
|
|
assert message["sender_name"] == "tom"
|