diff --git a/pyproject.toml b/pyproject.toml index 720dbc1f..6ff78efe 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -48,6 +48,7 @@ dependencies = [ "volcengine>=1.0.0", # For AgentKit Runtime API access "python-frontmatter==1.1.0", "tos>=2.8.4", # For TOS storage and Viking DB + "openviking-sdk>=0.1.3", ] [project.scripts] diff --git a/tests/test_openviking_long_term_memory.py b/tests/test_openviking_long_term_memory.py new file mode 100644 index 00000000..f0ecbd77 --- /dev/null +++ b/tests/test_openviking_long_term_memory.py @@ -0,0 +1,860 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import json +import threading +from typing import Any + +import pytest +from google.adk.events import Event +from google.adk.sessions import Session +from google.genai import types +from pydantic import Field + +from veadk.memory.long_term_memory import LongTermMemory +from veadk.memory.long_term_memory_backends.base_backend import ( + BaseLongTermMemoryBackend, +) +from veadk.memory.long_term_memory_backends.openviking_backend import ( + OpenVikingLTMBackend, +) + + +def _install_fake_openviking_sdk( + monkeypatch: pytest.MonkeyPatch, + *, + responses: dict[str, dict[str, Any]] | None = None, +): + calls = [] + responses = responses or {} + + class FakeSyncHTTPClient: + def __init__(self, **kwargs): + self.kwargs = kwargs + + def initialize(self): + pass + + def close(self): + pass + + def create_session(self, *, session_id=None, memory_policy=None, **kwargs): + calls.append( + { + "method": "create_session", + "actor_peer_id": self.kwargs.get("actor_peer_id"), + "payload": { + "session_id": session_id, + "memory_policy": memory_policy, + **kwargs, + }, + } + ) + return responses.get("create_session", {}) + + def add_message( + self, session_id, *, role, content=None, peer_id=None, **kwargs + ): + calls.append( + { + "method": "add_message", + "actor_peer_id": self.kwargs.get("actor_peer_id"), + "payload": { + "session_id": session_id, + "role": role, + "content": content, + "peer_id": peer_id, + **kwargs, + }, + } + ) + return responses.get("add_message", {}) + + def commit_session(self, *, session_id, keep_recent_count=0, **kwargs): + calls.append( + { + "method": "commit_session", + "actor_peer_id": self.kwargs.get("actor_peer_id"), + "payload": { + "session_id": session_id, + "keep_recent_count": keep_recent_count, + **kwargs, + }, + } + ) + return responses.get("commit_session", {}) + + def find(self, *, query, target_uri="", limit=10, context_type=None, **kwargs): + calls.append( + { + "method": "find", + "actor_peer_id": self.kwargs.get("actor_peer_id"), + "payload": { + "query": query, + "target_uri": target_uri, + "context_type": context_type, + "limit": limit, + **kwargs, + }, + } + ) + return responses.get("find", {}) + + monkeypatch.setattr( + "veadk.memory.long_term_memory_backends.openviking_backend.SyncHTTPClient", + FakeSyncHTTPClient, + ) + return calls + + +def _install_loop_sensitive_openviking_sdk( + monkeypatch: pytest.MonkeyPatch, + *, + event_loop_thread_id: int, + responses: dict[str, dict[str, Any]] | None = None, +): + calls = [] + responses = responses or {} + + class LoopSensitiveSyncHTTPClient: + def __init__(self, **kwargs): + self.kwargs = kwargs + + def _record(self, method: str, payload: dict[str, Any] | None = None): + thread_id = threading.get_ident() + if thread_id == event_loop_thread_id: + raise RuntimeError(f"{method} ran on the event loop thread") + calls.append( + { + "method": method, + "thread_id": thread_id, + "actor_peer_id": self.kwargs.get("actor_peer_id"), + "payload": payload or {}, + } + ) + + def initialize(self): + self._record("initialize") + + def close(self): + self._record("close") + + def create_session(self, *, session_id=None, memory_policy=None, **kwargs): + self._record( + "create_session", + { + "session_id": session_id, + "memory_policy": memory_policy, + **kwargs, + }, + ) + return responses.get("create_session", {}) + + def add_message( + self, session_id, *, role, content=None, peer_id=None, **kwargs + ): + self._record( + "add_message", + { + "session_id": session_id, + "role": role, + "content": content, + "peer_id": peer_id, + **kwargs, + }, + ) + return responses.get("add_message", {}) + + def commit_session(self, *, session_id, keep_recent_count=0, **kwargs): + self._record( + "commit_session", + { + "session_id": session_id, + "keep_recent_count": keep_recent_count, + **kwargs, + }, + ) + return responses.get("commit_session", {}) + + def find(self, *, query, target_uri="", limit=10, context_type=None, **kwargs): + self._record( + "find", + { + "query": query, + "target_uri": target_uri, + "context_type": context_type, + "limit": limit, + **kwargs, + }, + ) + return responses.get("find", {}) + + monkeypatch.setattr( + "veadk.memory.long_term_memory_backends.openviking_backend.SyncHTTPClient", + LoopSensitiveSyncHTTPClient, + ) + return calls + + +def test_openviking_backend_writes_peer_messages_and_commits(monkeypatch): + calls = _install_fake_openviking_sdk(monkeypatch) + backend = OpenVikingLTMBackend( + index="support_app", + url="http://openviking.test", + api_key="owner-key", + ) + + assert backend.save_memory( + user_id="alice", + event_strings=[ + json.dumps( + {"role": "user", "parts": [{"text": "我喜欢简短直接的回答"}]}, + ensure_ascii=False, + ) + ], + app_name="support_app", + session_id="session_001", + ) + + assert [call["method"] for call in calls] == [ + "create_session", + "add_message", + "commit_session", + ] + openviking_session_id = calls[0]["payload"]["session_id"] + assert openviking_session_id != "session_001" + assert openviking_session_id.startswith("veadk__") + assert "support_app" in openviking_session_id + assert "alice" in openviking_session_id + assert calls[0] == { + "method": "create_session", + "actor_peer_id": None, + "payload": { + "session_id": openviking_session_id, + "memory_policy": { + "self": {"enabled": False}, + "peer": {"enabled": True}, + "memory_types": ["entities", "events", "preferences"], + }, + }, + } + assert calls[1] == { + "method": "add_message", + "actor_peer_id": None, + "payload": { + "session_id": openviking_session_id, + "role": "user", + "content": "我喜欢简短直接的回答", + "peer_id": "alice", + }, + } + assert calls[2] == { + "method": "commit_session", + "actor_peer_id": None, + "payload": {"session_id": openviking_session_id, "keep_recent_count": 0}, + } + + +def test_openviking_backend_sdk_create_payload(monkeypatch): + calls = _install_fake_openviking_sdk(monkeypatch) + backend = OpenVikingLTMBackend( + index="support_app", + url="http://openviking.test", + api_key="owner-key", + ) + client = backend._new_client() + + backend._create_session(client=client, session_id="sa1") + + assert calls == [ + { + "method": "create_session", + "actor_peer_id": None, + "payload": { + "session_id": "sa1", + "memory_policy": { + "self": {"enabled": False}, + "peer": {"enabled": True}, + "memory_types": ["entities", "events", "preferences"], + }, + }, + } + ] + + +def test_openviking_backend_sdk_message_payload(monkeypatch): + calls = _install_fake_openviking_sdk(monkeypatch) + backend = OpenVikingLTMBackend( + index="support_app", + url="http://openviking.test", + api_key="owner-key", + ) + client = backend._new_client() + + backend._add_message( + client=client, + session_id="sa1", + role="user", + content="请记住：我的专属暗号是紫竹A1。", + peer_id="a1", + ) + + assert calls == [ + { + "method": "add_message", + "actor_peer_id": None, + "payload": { + "session_id": "sa1", + "role": "user", + "content": "请记住：我的专属暗号是紫竹A1。", + "peer_id": "a1", + }, + } + ] + + +def test_openviking_backend_search_uses_find_with_actor_peer(monkeypatch): + calls = _install_fake_openviking_sdk( + monkeypatch, + responses={"find": {"memories": [{"abstract": "用户喜欢简短直接的回答"}]}}, + ) + backend = OpenVikingLTMBackend( + index="support_app", + url="http://openviking.test", + api_key="owner-key", + ) + + memories = backend.search_memory( + user_id="alice", + query="用户偏好", + top_k=3, + app_name="support_app", + session_id="session_002", + ) + + assert memories == [ + json.dumps({"abstract": "用户喜欢简短直接的回答"}, ensure_ascii=False) + ] + assert calls[0]["actor_peer_id"] == "alice" + assert [call["method"] for call in calls] == ["find"] + assert "session_id" not in calls[0]["payload"] + assert calls[0] == { + "method": "find", + "actor_peer_id": "alice", + "payload": { + "query": "用户偏好", + "target_uri": "viking://user/peers/alice/memories", + "context_type": "memory", + "limit": 3, + }, + } + + +def test_openviking_backend_search_falls_back_to_find_without_session(monkeypatch): + calls = _install_fake_openviking_sdk( + monkeypatch, + responses={"find": {"memories": []}}, + ) + backend = OpenVikingLTMBackend( + index="support_app", + url="http://openviking.test", + api_key="owner-key", + ) + + assert ( + backend.search_memory( + user_id="alice", + query="用户偏好", + top_k=3, + app_name="support_app", + ) + == [] + ) + + assert calls == [ + { + "method": "find", + "actor_peer_id": "alice", + "payload": { + "query": "用户偏好", + "target_uri": "viking://user/peers/alice/memories", + "context_type": "memory", + "limit": 3, + }, + } + ] + + +def test_openviking_backend_sdk_find_payload(monkeypatch): + calls = _install_fake_openviking_sdk( + monkeypatch, + responses={"find": {"memories": []}}, + ) + backend = OpenVikingLTMBackend( + index="support_app", + url="http://openviking.test", + api_key="owner-key", + ) + + backend._search_with_actor_client( + peer_id="a1", + query="我的专属暗号和回答风格偏好是什么？", + top_k=10, + ) + + assert calls == [ + { + "method": "find", + "actor_peer_id": "a1", + "payload": { + "query": "我的专属暗号和回答风格偏好是什么？", + "target_uri": "viking://user/peers/a1/memories", + "context_type": "memory", + "limit": 10, + }, + } + ] + + +def test_openviking_backend_sdk_commit_payload(monkeypatch): + calls = _install_fake_openviking_sdk(monkeypatch) + backend = OpenVikingLTMBackend( + index="support_app", + url="http://openviking.test", + api_key="owner-key", + ) + client = backend._new_client() + + backend._commit_session(client=client, session_id="sa1") + + assert calls == [ + { + "method": "commit_session", + "actor_peer_id": None, + "payload": {"session_id": "sa1", "keep_recent_count": 0}, + } + ] + + +def test_openviking_backend_extracts_top_level_adk_memory_items(monkeypatch): + _install_fake_openviking_sdk( + monkeypatch, + responses={ + "find": { + "memories": [ + { + "author": "user", + "content": { + "parts": [ + "{'text': \"'# 用户偏好记录\\n- Result: 成功记录用户偏好'\"}" + ], + "role": "user", + }, + "custom_metadata": {}, + } + ] + } + }, + ) + backend = OpenVikingLTMBackend( + index="support_app", + url="http://openviking.test", + api_key="owner-key", + ) + + memories = backend.search_memory( + user_id="alice", + query="用户偏好", + top_k=3, + app_name="support_app", + ) + + assert json.loads(memories[0]) == { + "author": "user", + "content": { + "parts": ["{'text': \"'# 用户偏好记录\\n- Result: 成功记录用户偏好'\"}"], + "role": "user", + }, + "custom_metadata": {}, + } + + +def test_openviking_backend_deduplicates_memories_by_uri_and_keeps_best_score(): + backend = OpenVikingLTMBackend( + index="support_app", + url="http://openviking.test", + api_key="owner-key", + ) + + memories = [ + json.loads(memory) + for memory in backend._extract_memories( + { + "memories": [ + { + "uri": "viking://user/default/peers/aa01/memories/a.md", + "score": 0.1, + "abstract": "low score", + }, + { + "uri": "viking://user/default/peers/aa01/memories/b.md", + "score": 0.3, + "abstract": "other memory", + }, + { + "uri": "viking://user/default/peers/aa01/memories/a.md", + "score": 0.9, + "abstract": "high score", + }, + ] + } + ) + ] + + assert len(memories) == 2 + assert memories[0]["abstract"] == "high score" + assert memories[0]["score"] == 0.9 + assert memories[1]["abstract"] == "other memory" + + +@pytest.mark.parametrize( + "peer_id", + ["app/alice", "app:alice", "alice+1", ".", "..", "alice bob"], +) +def test_openviking_backend_rejects_unsafe_peer_id(peer_id): + backend = OpenVikingLTMBackend( + index="support_app", + url="http://openviking.test", + api_key="owner-key", + peer_id_resolver=lambda app_name, user_id: peer_id, + ) + + with pytest.raises(ValueError, match="safe single path segment"): + backend.search_memory( + user_id="alice", + query="用户偏好", + top_k=3, + app_name="support_app", + ) + + +def test_openviking_backend_namespaces_same_session_by_peer(monkeypatch): + calls = _install_fake_openviking_sdk(monkeypatch) + backend = OpenVikingLTMBackend( + index="support_app", + url="http://openviking.test", + api_key="owner-key", + ) + + for user_id in ("alice", "bob"): + assert backend.save_memory( + user_id=user_id, + event_strings=[json.dumps({"role": "user", "parts": [{"text": user_id}]})], + app_name="support_app", + session_id="session", + ) + + created_sessions = [ + call["payload"]["session_id"] + for call in calls + if call["method"] == "create_session" + ] + message_calls = [call for call in calls if call["method"] == "add_message"] + + assert len(created_sessions) == 2 + assert created_sessions[0] != created_sessions[1] + assert "alice" in created_sessions[0] + assert "bob" in created_sessions[1] + assert message_calls[0]["payload"]["peer_id"] == "alice" + assert message_calls[0]["payload"]["session_id"] == created_sessions[0] + assert message_calls[1]["payload"]["peer_id"] == "bob" + assert message_calls[1]["payload"]["session_id"] == created_sessions[1] + + +@pytest.mark.asyncio +async def test_long_term_memory_openviking_search_runs_sdk_off_event_loop_thread( + monkeypatch, +): + event_loop_thread_id = threading.get_ident() + calls = _install_loop_sensitive_openviking_sdk( + monkeypatch, + event_loop_thread_id=event_loop_thread_id, + responses={ + "find": { + "memories": [ + { + "uri": "viking://user/peers/alice/memories/preferences/a.md", + "abstract": "用户喜欢简短直接的回答", + } + ] + } + }, + ) + memory = LongTermMemory( + backend="openviking", + backend_config={ + "index": "support_app", + "url": "http://openviking.test", + "api_key": "owner-key", + }, + top_k=2, + ) + + response = await memory.search_memory( + app_name="support_app", + user_id="alice", + query="用户偏好", + ) + + assert response.memories[0].content.parts[0].text == "用户喜欢简短直接的回答" + assert [call["method"] for call in calls] == ["initialize", "find", "close"] + assert {call["thread_id"] for call in calls} != {event_loop_thread_id} + assert calls[1]["actor_peer_id"] == "alice" + assert calls[1]["payload"] == { + "query": "用户偏好", + "target_uri": "viking://user/peers/alice/memories", + "context_type": "memory", + "limit": 2, + } + + +@pytest.mark.asyncio +async def test_long_term_memory_openviking_save_runs_sdk_off_event_loop_thread( + monkeypatch, +): + event_loop_thread_id = threading.get_ident() + calls = _install_loop_sensitive_openviking_sdk( + monkeypatch, + event_loop_thread_id=event_loop_thread_id, + ) + memory = LongTermMemory( + backend="openviking", + backend_config={ + "index": "support_app", + "url": "http://openviking.test", + "api_key": "owner-key", + }, + ) + session = Session( + id="session_001", + app_name="support_app", + user_id="alice", + state={}, + events=[ + Event( + author="user", + content=types.Content( + parts=[types.Part(text="记住我喜欢短回答")], + role="user", + ), + ) + ], + ) + + await memory.add_session_to_memory(session) + + assert [call["method"] for call in calls] == [ + "initialize", + "create_session", + "add_message", + "commit_session", + "close", + ] + assert {call["thread_id"] for call in calls} != {event_loop_thread_id} + assert len({call["thread_id"] for call in calls}) == 1 + assert calls[1]["payload"]["session_id"].startswith("veadk__") + assert calls[2]["payload"]["peer_id"] == "alice" + assert calls[2]["payload"]["content"] == "记住我喜欢短回答" + + +class _RecordingBackend(BaseLongTermMemoryBackend): + saved_call: dict[str, Any] = Field(default_factory=dict) + search_call: dict[str, Any] = Field(default_factory=dict) + + def precheck_index_naming(self): + pass + + def save_memory(self, user_id: str, event_strings: list[str], **kwargs) -> bool: + self.saved_call = { + "user_id": user_id, + "event_strings": event_strings, + "kwargs": kwargs, + } + return True + + def search_memory( + self, user_id: str, query: str, top_k: int, **kwargs + ) -> list[str]: + self.search_call = { + "user_id": user_id, + "query": query, + "top_k": top_k, + "kwargs": kwargs, + } + return [ + json.dumps( + { + "author": "assistant", + "role": "model", + "parts": [{"text": "remembered"}], + } + ) + ] + + +class _RawOpenVikingResponseBackend(BaseLongTermMemoryBackend): + def precheck_index_naming(self): + pass + + def save_memory(self, user_id: str, event_strings: list[str], **kwargs) -> bool: + return True + + def search_memory( + self, user_id: str, query: str, top_k: int, **kwargs + ) -> list[str]: + return [ + json.dumps( + { + "memories": [ + { + "author": "user", + "content": { + "parts": [ + "{'text': \"'# 用户偏好记录\\n- Result: 成功记录用户偏好'\"}" + ], + "role": "user", + }, + "custom_metadata": {"source": "openviking"}, + } + ] + }, + ensure_ascii=False, + ) + ] + + +@pytest.mark.asyncio +async def test_long_term_memory_passes_session_and_app_to_backend(): + backend = _RecordingBackend(index="support_app") + memory = LongTermMemory(backend=backend, app_name="support_app") + session = Session( + id="session_001", + app_name="support_app", + user_id="alice", + state={}, + events=[ + Event( + author="user", + content=types.Content( + parts=[types.Part(text="记住我喜欢短回答")], + role="user", + ), + ) + ], + ) + + await memory.add_session_to_memory(session, source="manual") + + assert backend.saved_call["user_id"] == "alice" + assert backend.saved_call["kwargs"] == { + "session_id": "session_001", + "app_name": "support_app", + "source": "manual", + } + + +@pytest.mark.asyncio +async def test_long_term_memory_does_not_pass_session_to_search_backend(): + backend = _RecordingBackend(index="support_app") + memory = LongTermMemory(backend=backend, app_name="support_app", top_k=2) + + response = await memory.search_memory( + app_name="support_app", + user_id="alice", + query="用户偏好", + ) + + assert response.memories[0].author == "assistant" + assert response.memories[0].content.role == "model" + assert response.memories[0].content.parts[0].text == "remembered" + assert backend.search_call == { + "user_id": "alice", + "query": "用户偏好", + "top_k": 2, + "kwargs": { + "app_name": "support_app", + }, + } + + +def test_long_term_memory_parses_openviking_abstract_and_metadata(): + backend = _RecordingBackend(index="support_app") + memory = LongTermMemory(backend=backend, app_name="support_app", top_k=2) + + entries = memory._convert_memory_chunk_to_entries( + json.dumps( + { + "context_type": "memory", + "uri": "viking://user/default/peers/aa01/memories/entities/人物/用户A.md", + "level": 2, + "score": 0.40563082695007324, + "category": "", + "match_reason": "", + "relations": [], + "abstract": "# 用户A\n- 专属暗号：紫竹A001_manual_001", + "overview": None, + }, + ensure_ascii=False, + ) + ) + + assert len(entries) == 1 + assert ( + entries[0].id + == "viking://user/default/peers/aa01/memories/entities/人物/用户A.md" + ) + assert ( + entries[0].content.parts[0].text == "# 用户A\n- 专属暗号：紫竹A001_manual_001" + ) + assert entries[0].custom_metadata == { + "context_type": "memory", + "uri": "viking://user/default/peers/aa01/memories/entities/人物/用户A.md", + "level": 2, + "score": 0.40563082695007324, + "category": "", + "match_reason": "", + "relations": [], + "overview": None, + } + + +@pytest.mark.asyncio +async def test_long_term_memory_parses_raw_openviking_memory_response(): + backend = _RawOpenVikingResponseBackend(index="support_app") + memory = LongTermMemory(backend=backend, app_name="support_app", top_k=2) + + response = await memory.search_memory( + app_name="support_app", + user_id="alice", + query="用户偏好", + ) + + assert len(response.memories) == 1 + assert response.memories[0].author == "user" + assert response.memories[0].content.role == "user" + assert ( + response.memories[0].content.parts[0].text + == "# 用户偏好记录\n- Result: 成功记录用户偏好" + ) + assert response.memories[0].custom_metadata == {"source": "openviking"} diff --git a/veadk/configs/database_configs.py b/veadk/configs/database_configs.py index ed8c09f7..a82eb113 100644 --- a/veadk/configs/database_configs.py +++ b/veadk/configs/database_configs.py @@ -102,6 +102,16 @@ class Mem0Config(BaseSettings): base_url: str = "" # "https://api.mem0.ai/v1" +class OpenVikingConfig(BaseSettings): + model_config = SettingsConfigDict(env_prefix="DATABASE_OPENVIKING_") + + url: str = "" + """OpenViking server base URL""" + + api_key: str = "" + """OpenViking service owner user key""" + + class VikingKnowledgebaseConfig(BaseSettings): model_config = SettingsConfigDict(env_prefix="DATABASE_VIKING_") diff --git a/veadk/memory/long_term_memory.py b/veadk/memory/long_term_memory.py index df5b4577..cbb7e313 100644 --- a/veadk/memory/long_term_memory.py +++ b/veadk/memory/long_term_memory.py @@ -14,7 +14,10 @@ # adapted from Google ADK memory service adk-python/src/google/adk/memory/vertex_ai_memory_bank_service.py at 0a9e67dbca67789247e882d16b139dbdc76a329a · google/adk-python +import ast +import asyncio import json +from collections.abc import Iterable from typing import Any, Literal from google.adk.events.event import Event @@ -69,6 +72,12 @@ def _get_backend_cls(backend: str) -> type[BaseLongTermMemoryBackend]: ) return Mem0LTMBackend + case "openviking": + from veadk.memory.long_term_memory_backends.openviking_backend import ( + OpenVikingLTMBackend, + ) + + return OpenVikingLTMBackend case _: raise ValueError(f"Unsupported long term memory backend: {backend}") except ImportError as e: @@ -88,7 +97,7 @@ class LongTermMemory(BaseMemoryService, BaseModel): It supports configuration of the backend service and retrieval behavior. Attributes: - backend (Union[Literal["local", "opensearch", "redis", "viking", "viking_mem", "mem0"], BaseLongTermMemoryBackend]): + backend (Union[Literal["local", "opensearch", "redis", "viking", "viking_mem", "mem0", "openviking"], BaseLongTermMemoryBackend]): The type or instance of the long-term memory backend. Defaults to "opensearch". backend_config (dict): @@ -111,7 +120,15 @@ class LongTermMemory(BaseMemoryService, BaseModel): """ backend: Union[ - Literal["local", "opensearch", "redis", "viking", "viking_mem", "mem0"], + Literal[ + "local", + "opensearch", + "redis", + "viking", + "viking_mem", + "mem0", + "openviking", + ], BaseLongTermMemoryBackend, ] = "opensearch" @@ -170,7 +187,9 @@ def model_post_init(self, __context: Any) -> None: f"Initialized long term memory with provided backend instance {self._backend.__class__.__name__}, index={self.index}" ) - def _filter_and_convert_events(self, events: list[Event]) -> list[str]: + def _filter_and_convert_events( + self, events: Iterable[Event], *, include_assistant: bool = False + ) -> list[str]: final_events = [] for event in events: # filter: bad event @@ -178,7 +197,7 @@ def _filter_and_convert_events(self, events: list[Event]) -> list[str]: continue # filter: only add user event to memory to enhance retrieve performance - if not event.author == "user": + if not include_assistant and not event.author == "user": continue # filter: discard function call and function response @@ -187,10 +206,19 @@ def _filter_and_convert_events(self, events: list[Event]) -> list[str]: # convert: to string-format for storage message = event.content.model_dump(exclude_none=True, mode="json") + message["role"] = self._normalize_event_role(event, message) final_events.append(json.dumps(message, ensure_ascii=False)) return final_events + def _normalize_event_role(self, event: Event, message: dict[str, Any]) -> str: + role = message.get("role") + if event.author == "user" or role == "user": + return "user" + if role == "system": + return "system" + return "assistant" + @override async def add_session_to_memory( self, @@ -225,17 +253,34 @@ async def add_session_to_memory( ``` """ user_id = session.user_id - event_strings = self._filter_and_convert_events(session.events) + save_kwargs = dict(kwargs) + nested_kwargs = save_kwargs.pop("kwargs", None) + if isinstance(nested_kwargs, dict): + save_kwargs.update(nested_kwargs) + app_name = self.app_name or getattr(session, "app_name", "") + include_assistant = ( + self.backend == "openviking" + or self._backend.__class__.__name__ == "OpenVikingLTMBackend" + ) + event_strings = self._filter_and_convert_events( + session.events, + include_assistant=include_assistant, + ) logger.info( f"Adding {len(event_strings)} events to long term memory: index={self.index}" ) - if self.backend == "viking": - self._backend.save_memory( - user_id=user_id, event_strings=event_strings, **kwargs - ) + save_call = { + "user_id": user_id, + "event_strings": event_strings, + "session_id": session.id, + "app_name": app_name, + **save_kwargs, + } + if self._uses_openviking_backend(): + await asyncio.to_thread(self._backend.save_memory, **save_call) else: - self._backend.save_memory(user_id=user_id, event_strings=event_strings) + self._backend.save_memory(**save_call) logger.info( f"Added {len(event_strings)} events to long term memory: index={self.index}, user_id={user_id}" ) @@ -264,9 +309,20 @@ async def search_memory( memory_chunks = [] try: - memory_chunks = self._backend.search_memory( - query=query, top_k=self.top_k, user_id=user_id - ) + search_kwargs = {"app_name": app_name} + search_call = { + "query": query, + "top_k": self.top_k, + "user_id": user_id, + **search_kwargs, + } + if self._uses_openviking_backend(): + memory_chunks = await asyncio.to_thread( + self._backend.search_memory, + **search_call, + ) + else: + memory_chunks = self._backend.search_memory(**search_call) except Exception as e: logger.error( f"Exception orrcus during memory search: {e}. Return empty memory chunks" @@ -274,34 +330,154 @@ async def search_memory( memory_events = [] for memory in memory_chunks: - try: - memory_dict = json.loads(memory) - try: - text = memory_dict["parts"][0]["text"] - role = memory_dict["role"] - except KeyError as _: - # prevent not a standard text-based event - logger.warning( - f"Memory content: {memory_dict}. Skip return this memory." - ) - continue - except json.JSONDecodeError: - # prevent the memory string is not dumped by `Event` class - text = memory - role = "user" - - memory_events.append( - MemoryEntry( - author="user", - content=types.Content(parts=[types.Part(text=text)], role=role), - ) - ) + memory_events.extend(self._convert_memory_chunk_to_entries(memory)) logger.info( f"Return {len(memory_events)} memory events for query: {query} index={self.index} user_id={user_id}" ) return SearchMemoryResponse(memories=memory_events) + def _uses_openviking_backend(self) -> bool: + return ( + self.backend == "openviking" + or self._backend.__class__.__name__ == "OpenVikingLTMBackend" + ) + + def _convert_memory_chunk_to_entries(self, memory: str) -> list[MemoryEntry]: + try: + memory_dict = json.loads(memory) + except json.JSONDecodeError: + return [ + MemoryEntry( + author="user", + content=types.Content( + parts=[types.Part(text=memory)], + role="user", + ), + ) + ] + + if not isinstance(memory_dict, dict): + return [ + MemoryEntry( + author="user", + content=types.Content( + parts=[types.Part(text=str(memory_dict))], + role="user", + ), + ) + ] + + memories = memory_dict.get("memories") + if isinstance(memories, list): + entries = [] + for item in memories: + if isinstance(item, dict): + entry = self._convert_memory_dict_to_entry(item) + if entry: + entries.append(entry) + else: + entries.extend(self._convert_memory_chunk_to_entries(str(item))) + return entries + + entry = self._convert_memory_dict_to_entry(memory_dict) + return [entry] if entry else [] + + def _convert_memory_dict_to_entry( + self, memory_dict: dict[str, Any] + ) -> MemoryEntry | None: + content = memory_dict.get("content") + if isinstance(content, dict): + role = str(content.get("role") or memory_dict.get("role") or "user") + text = self._extract_memory_parts_text(content.get("parts") or []) + else: + role = str(memory_dict.get("role") or "user") + text = self._extract_memory_parts_text(memory_dict.get("parts") or []) + if not text: + text = self._extract_memory_text_field(memory_dict) + + if not text: + logger.warning(f"Memory content: {memory_dict}. Skip return this memory.") + return None + + custom_metadata = memory_dict.get("custom_metadata") + if not isinstance(custom_metadata, dict): + custom_metadata = self._extract_memory_custom_metadata(memory_dict) + + return MemoryEntry( + author=memory_dict.get("author", "user"), + content=types.Content(parts=[types.Part(text=text)], role=role), + custom_metadata=custom_metadata, + id=memory_dict.get("id") or memory_dict.get("uri"), + timestamp=memory_dict.get("timestamp"), + ) + + def _extract_memory_custom_metadata( + self, memory_dict: dict[str, Any] + ) -> dict[str, Any]: + metadata_keys = ( + "context_type", + "uri", + "level", + "score", + "category", + "match_reason", + "relations", + "overview", + ) + return {key: memory_dict[key] for key in metadata_keys if key in memory_dict} + + def _extract_memory_parts_text(self, parts: list[Any]) -> str: + text_parts = [] + for part in parts: + text = self._extract_memory_part_text(part) + if text: + text_parts.append(text) + return "\n".join(text_parts) + + def _extract_memory_part_text(self, part: Any) -> str: + if isinstance(part, dict): + if "text" in part: + return self._clean_memory_text(part["text"]) + return json.dumps(part, ensure_ascii=False) + if isinstance(part, str): + try: + parsed = ast.literal_eval(part) + except (ValueError, SyntaxError): + return self._clean_memory_text(part) + if isinstance(parsed, dict) and "text" in parsed: + return self._clean_memory_text(parsed["text"]) + if isinstance(parsed, str): + return self._clean_memory_text(parsed) + return json.dumps(parsed, ensure_ascii=False) + return str(part) + + def _extract_memory_text_field(self, memory_dict: dict[str, Any]) -> str: + for key in ("text", "abstract", "summary", "content"): + value = memory_dict.get(key) + if value and not isinstance(value, dict): + return self._clean_memory_text(value) + return "" + + def _clean_memory_text(self, value: Any) -> str: + text = str(value) + for _ in range(2): + stripped = text.strip() + if not ( + len(stripped) >= 2 + and stripped[0] == stripped[-1] + and stripped[0] in {"'", '"'} + ): + return stripped + try: + parsed = ast.literal_eval(stripped) + except (ValueError, SyntaxError): + return stripped[1:-1] + if not isinstance(parsed, str): + return str(parsed) + text = parsed + return text.strip() + def get_user_profile(self, user_id: str) -> str: logger.info(f"Get user profile for user_id={user_id}") if self.backend == "viking": diff --git a/veadk/memory/long_term_memory_backends/openviking_backend.py b/veadk/memory/long_term_memory_backends/openviking_backend.py new file mode 100644 index 00000000..c0fc40e6 --- /dev/null +++ b/veadk/memory/long_term_memory_backends/openviking_backend.py @@ -0,0 +1,322 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +import hashlib +import json +import re +import uuid +from collections.abc import Callable +from typing import Any, ClassVar + +from openviking_sdk import SyncHTTPClient +from pydantic import Field +from typing_extensions import override + +from veadk.configs.database_configs import OpenVikingConfig +from veadk.memory.long_term_memory_backends.base_backend import ( + BaseLongTermMemoryBackend, +) +from veadk.utils.logger import get_logger + +logger = get_logger(__name__) + +_SAFE_ID_PATTERN = re.compile(r"^[A-Za-z0-9_.@-]+$") +_UNSAFE_ID_CHARS = re.compile(r"[^A-Za-z0-9_.@-]+") + + +def default_peer_id_resolver(app_name: str, user_id: str) -> str: + del app_name + return user_id + + +class OpenVikingLTMBackend(BaseLongTermMemoryBackend): + """OpenViking long term memory backend using the OpenViking SDK.""" + + openviking_config: OpenVikingConfig = Field(default_factory=OpenVikingConfig) + url: str = "" + api_key: str = "" + peer_id_resolver: Callable[[str, str], str] | None = None + timeout: float = 30 + + _PEER_MEMORY_TYPES: ClassVar[list[str]] = ["entities", "events", "preferences"] + _MEMORY_POLICY: ClassVar[dict[str, Any]] = { + "self": {"enabled": False}, + "peer": {"enabled": True}, + "memory_types": _PEER_MEMORY_TYPES, + } + + def model_post_init(self, __context: Any) -> None: + self.url = (self.url or self.openviking_config.url).rstrip("/") + self.api_key = self.api_key or self.openviking_config.api_key + if not self.url: + raise ValueError( + "OpenViking URL is required. Set DATABASE_OPENVIKING_URL or pass url." + ) + if not self.api_key: + raise ValueError( + "OpenViking API key is required. Set DATABASE_OPENVIKING_API_KEY or pass api_key." + ) + if not self.peer_id_resolver: + self.peer_id_resolver = default_peer_id_resolver + + def precheck_index_naming(self): + if not self.index: + raise ValueError("OpenViking backend index/app_name must not be empty") + + @override + def save_memory( + self, + user_id: str, + event_strings: list[str], + **kwargs, + ) -> bool: + if not event_strings: + return True + + app_name = kwargs.get("app_name") or self.index + session_id = kwargs.get("session_id") or str(uuid.uuid4()) + peer_id = self._resolve_peer_id(app_name=app_name, user_id=user_id) + openviking_session_id = self._openviking_session_id( + app_name=app_name, + peer_id=peer_id, + session_id=session_id, + ) + + client = self._new_client() + try: + client.initialize() + self._create_session(client=client, session_id=openviking_session_id) + for event_string in event_strings: + role, content = self._parse_event_string(event_string) + self._add_message( + client=client, + session_id=openviking_session_id, + role=role, + content=content, + peer_id=peer_id, + ) + self._commit_session(client=client, session_id=openviking_session_id) + finally: + client.close() + return True + + @override + def search_memory( + self, + user_id: str, + query: str, + top_k: int, + **kwargs, + ) -> list[str]: + app_name = kwargs.get("app_name") or self.index + peer_id = self._resolve_peer_id(app_name=app_name, user_id=user_id) + + response = self._search_with_actor_client( + peer_id=peer_id, + query=query, + top_k=top_k, + ) + return self._extract_memories(response) + + def _resolve_peer_id(self, *, app_name: str, user_id: str) -> str: + assert self.peer_id_resolver is not None + peer_id = str(self.peer_id_resolver(app_name, user_id) or "").strip() + if not peer_id: + raise ValueError("OpenViking peer_id must not be empty") + if peer_id in {".", ".."} or not _SAFE_ID_PATTERN.match(peer_id): + raise ValueError( + "OpenViking peer_id must be a safe single path segment. " + "Allowed characters: letters, digits, '.', '_', '@', '-'." + ) + return peer_id + + def _peer_memory_target_uri(self, *, peer_id: str) -> str: + return f"viking://user/peers/{peer_id}/memories" + + def _openviking_session_id( + self, + *, + app_name: str, + peer_id: str, + session_id: str, + ) -> str: + return "__".join( + [ + "veadk", + self._safe_identifier_part(app_name, fallback="app"), + self._safe_identifier_part(peer_id, fallback="peer"), + self._safe_identifier_part(session_id, fallback="session"), + ] + ) + + def _safe_identifier_part(self, value: Any, *, fallback: str) -> str: + raw = str(value or "").strip() + normalized = _UNSAFE_ID_CHARS.sub("_", raw).strip("._-") + if not normalized or normalized in {".", ".."}: + normalized = fallback + digest = hashlib.sha256(raw.encode("utf-8")).hexdigest()[:12] + return f"{normalized[:48]}_{digest}" + + def _new_client(self, *, actor_peer_id: str | None = None) -> SyncHTTPClient: + return SyncHTTPClient( + url=self.url, + api_key=self.api_key, + actor_peer_id=actor_peer_id, + timeout=self.timeout, + ) + + def _create_session(self, *, client: SyncHTTPClient, session_id: str) -> None: + try: + client.create_session( + session_id=session_id, + memory_policy=self._MEMORY_POLICY, + ) + except Exception as e: + if self._is_existing_session_error(e): + logger.debug(f"OpenViking session already exists, continue: {e}") + return + raise + + def _add_message( + self, + *, + client: SyncHTTPClient, + session_id: str, + role: str, + content: str, + peer_id: str, + ) -> None: + client.add_message( + session_id=session_id, + role=role, + content=content, + peer_id=peer_id, + ) + + def _commit_session(self, *, client: SyncHTTPClient, session_id: str) -> None: + client.commit_session(session_id=session_id, keep_recent_count=0) + + def _search_with_actor_client( + self, *, peer_id: str, query: str, top_k: int + ) -> dict[str, Any]: + target_uri = self._peer_memory_target_uri(peer_id=peer_id) + client = self._new_client(actor_peer_id=peer_id) + try: + client.initialize() + return client.find( + query=query, + target_uri=target_uri, + context_type="memory", + limit=top_k, + ) + finally: + client.close() + + def _is_existing_session_error(self, error: Exception) -> bool: + code = str(getattr(error, "code", "")).replace("_", "").upper() + message = str(error).lower() + return ( + code in {"ALREADYEXISTS", "ALREADYEXISTSERROR", "CONFLICT"} + or "already" in message + ) + + def _parse_event_string(self, event_string: str) -> tuple[str, str]: + try: + event = json.loads(event_string) + except json.JSONDecodeError: + return "user", event_string + + role = self._normalize_role(event.get("role")) + content = self._extract_text_from_event(event) + if not content: + content = event_string + return role, content + + def _normalize_role(self, role: Any) -> str: + if role == "user": + return "user" + if role == "system": + return "system" + return "assistant" + + def _extract_text_from_event(self, event: dict[str, Any]) -> str: + parts = event.get("parts") or [] + text_parts = [] + for part in parts: + if not isinstance(part, dict): + continue + text = part.get("text") + if text: + text_parts.append(str(text)) + return "\n".join(text_parts) + + def _extract_memories(self, response: dict[str, Any]) -> list[str]: + result = response.get("result") + result = result if isinstance(result, dict) else response + memories = result.get("memories") or response.get("memories") or [] + return [ + memory + for memory in ( + self._memory_to_text(item) + for item in self._deduplicate_memory_items(memories) + ) + if memory + ] + + def _deduplicate_memory_items(self, memories: list[Any]) -> list[Any]: + deduplicated: list[Any] = [] + key_to_index: dict[str, int] = {} + for item in memories: + key = self._memory_dedupe_key(item) + if not key: + deduplicated.append(item) + continue + + existing_index = key_to_index.get(key) + if existing_index is None: + key_to_index[key] = len(deduplicated) + deduplicated.append(item) + continue + + existing = deduplicated[existing_index] + if self._memory_score(item) > self._memory_score(existing): + deduplicated[existing_index] = item + return deduplicated + + def _memory_dedupe_key(self, item: Any) -> str: + if isinstance(item, dict): + uri = str(item.get("uri") or "").strip() + if uri: + return f"uri:{uri}" + return f"dict:{json.dumps(item, ensure_ascii=False, sort_keys=True)}" + if isinstance(item, str): + return f"str:{item}" + return "" + + def _memory_score(self, item: Any) -> float: + if not isinstance(item, dict): + return 0 + try: + return float(item.get("score") or 0) + except (TypeError, ValueError): + return 0 + + def _memory_to_text(self, item: Any) -> str: + if isinstance(item, str): + return item + if not isinstance(item, dict): + return str(item) + return json.dumps(item, ensure_ascii=False) diff --git a/veadk/runner.py b/veadk/runner.py index c9518bfc..55cce3b4 100644 --- a/veadk/runner.py +++ b/veadk/runner.py @@ -757,5 +757,5 @@ async def save_session_to_long_term_memory( ) return - await self.long_term_memory.add_session_to_memory(session, kwargs=kwargs) + await self.long_term_memory.add_session_to_memory(session, **kwargs) logger.info(f"Add session `{session.id}` to long term memory.")