Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
95 changes: 95 additions & 0 deletions tests/test_vikingdb_memory_backend.py
Original file line number Diff line number Diff line change
Expand Up @@ -542,3 +542,98 @@ def test_direct_volcengine_viking_memory_client_uses_region_env(

assert client.get_host() == "api-knowledgebase.mlp.cn-shanghai.volces.com"
assert client.service_info.credentials.region == "cn-shanghai"


def test_viking_memory_search_returns_one_entry_per_result(
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""search_memory must not dump the whole result_list via str(list).

On main it returns a single JSON blob whose text is the Python repr of the
entire result_list, so callers (LongTermMemory → MemoryEntry) see one
unusable memory instead of each hit.
"""
module = _load_backend_module(monkeypatch)
captured: dict[str, Any] = {}

class FakeCollection:
def search_memory(self, **kwargs: Any) -> dict[str, Any]:
captured["search_kwargs"] = kwargs
return {
"code": 0,
"data": {
"result_list": [
{
"id": "m1",
"score": 0.9,
"memory_info": {"summary": "likes tea"},
},
{
"id": "m2",
"score": 0.8,
"memory_info": {"summary": "lives in Shanghai"},
},
]
},
}

class FakeSdkClient:
def get_collection(self, **kwargs: Any) -> FakeCollection:
return FakeCollection()

monkeypatch.setattr(
module.VikingDBLTMBackend,
"_collection_exist",
lambda self: True,
)
monkeypatch.setattr(
module.VikingDBLTMBackend,
"_get_sdk_client",
lambda self: FakeSdkClient(),
)

backend = module.VikingDBLTMBackend(index="agent_memory", api_key="mem-api-key")
results = backend.search_memory(user_id="user-42", query="tea", top_k=5)

assert captured["search_kwargs"]["filter"] == {
"user_id": "user-42",
"memory_type": backend.memory_type,
}
assert captured["search_kwargs"]["limit"] == 5
assert len(results) == 2

import json

parsed = [json.loads(item) for item in results]
texts = [item["parts"][0]["text"] for item in parsed]
assert texts == ["likes tea", "lives in Shanghai"]
# Must not be the Python list repr of the whole result_list.
assert not any(text.startswith("[{") for text in texts)


def test_viking_memory_search_returns_empty_when_no_hits(
monkeypatch: pytest.MonkeyPatch,
) -> None:
module = _load_backend_module(monkeypatch)

class FakeCollection:
def search_memory(self, **kwargs: Any) -> dict[str, Any]:
return {"code": 0, "data": {"result_list": []}}

class FakeSdkClient:
def get_collection(self, **kwargs: Any) -> FakeCollection:
return FakeCollection()

monkeypatch.setattr(
module.VikingDBLTMBackend,
"_collection_exist",
lambda self: True,
)
monkeypatch.setattr(
module.VikingDBLTMBackend,
"_get_sdk_client",
lambda self: FakeSdkClient(),
)

backend = module.VikingDBLTMBackend(index="agent_memory", api_key="mem-api-key")
assert backend.search_memory(user_id="user-42", query="tea", top_k=5) == []
59 changes: 50 additions & 9 deletions veadk/memory/long_term_memory_backends/vikingdb_memory_backend.py
Original file line number Diff line number Diff line change
Expand Up @@ -345,18 +345,59 @@ def search_memory(

logger.debug(f"Original response from Viking Memory: {response}")

result = response.get("data", {}).get("result_list", [])

return (
[
result_list = response.get("data", {}).get("result_list", [])
memories: list[str] = []
for item in result_list:
text = self._memory_item_to_text(item)
if not text:
continue
memories.append(
json.dumps(
{"role": "user", "parts": [{"text": str(result)}]},
{"role": "user", "parts": [{"text": text}]},
ensure_ascii=False,
)
]
if result
else []
)
)
return memories

@staticmethod
def _memory_item_to_text(item: Any) -> str:
"""Extract a readable memory string from one SearchMemory hit.

VikingDB returns each hit as an object with ``memory_info`` whose shape
depends on ``memory_type`` (event attributes, profile attributes, etc.).
Prefer common text fields; otherwise serialize ``memory_info``.
"""
if isinstance(item, str):
return item.strip()
if not isinstance(item, dict):
return str(item)

memory_info = item.get("memory_info")
if isinstance(memory_info, str):
return memory_info.strip()
if isinstance(memory_info, dict):
for key in (
"summary",
"user_profile",
"original_messages",
"question",
"answer",
"content",
"text",
):
value = memory_info.get(key)
if value and not isinstance(value, (dict, list)):
return str(value).strip()
if memory_info:
return json.dumps(memory_info, ensure_ascii=False)
elif memory_info is not None:
return str(memory_info).strip()

for key in ("summary", "content", "text"):
value = item.get(key)
if value and not isinstance(value, (dict, list)):
return str(value).strip()
return json.dumps(item, ensure_ascii=False) if item else ""

def get_user_profile(self, user_id: str) -> str:
logger.info(
Expand Down