From 6885d483d64759941d8a52e120593d67799ffea8 Mon Sep 17 00:00:00 2001 From: "fuyi.0" Date: Fri, 5 Dec 2025 11:52:37 +0800 Subject: [PATCH 1/3] feat: integrate vikingdb-python-sdk for memory backend --- pyproject.toml | 1 + .../ve_viking_db_memory.py | 27 +---------------- .../vikingdb_memory_backend.py | 30 ++++++++++++++----- 3 files changed, 25 insertions(+), 33 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 9816c11c..af2df1f0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -38,6 +38,7 @@ dependencies = [ "aiomysql>=0.3.2", # For async MySQL database (short term memory) "opensearch-py==2.8.0", "filetype>=1.2.0", + "vikingdb-python-sdk>=0.1.3", ] [project.scripts] diff --git a/veadk/integrations/ve_viking_db_memory/ve_viking_db_memory.py b/veadk/integrations/ve_viking_db_memory/ve_viking_db_memory.py index 0207cd32..43e5d1ae 100644 --- a/veadk/integrations/ve_viking_db_memory/ve_viking_db_memory.py +++ b/veadk/integrations/ve_viking_db_memory/ve_viking_db_memory.py @@ -280,29 +280,4 @@ def update_collection( "BuiltinEntityTypes": builtin_entity_types, } res = self.json("UpdateCollection", {}, json.dumps(params)) - return json.loads(res) - - def search_memory(self, collection_name, query, filter, limit=10): - params = { - "collection_name": collection_name, - "limit": limit, - "filter": filter, - } - if query: - params["query"] = query - res = self.json("SearchMemory", {}, json.dumps(params)) - return json.loads(res) - - def add_messages( - self, collection_name, session_id, messages, metadata, entities=None - ): - params = { - "collection_name": collection_name, - "session_id": session_id, - "messages": messages, - "metadata": metadata, - } - if entities is not None: - params["entities"] = entities - res = self.json("AddMessages", {}, json.dumps(params)) - return json.loads(res) + return json.loads(res) \ No newline at end of file diff --git a/veadk/memory/long_term_memory_backends/vikingdb_memory_backend.py b/veadk/memory/long_term_memory_backends/vikingdb_memory_backend.py index 0656d63c..fefc11e2 100644 --- a/veadk/memory/long_term_memory_backends/vikingdb_memory_backend.py +++ b/veadk/memory/long_term_memory_backends/vikingdb_memory_backend.py @@ -30,6 +30,9 @@ from veadk.memory.long_term_memory_backends.base_backend import ( BaseLongTermMemoryBackend, ) +from vikingdb import IAM +from vikingdb.memory import VikingMem + from veadk.utils.logger import get_logger logger = get_logger(__name__) @@ -118,6 +121,16 @@ def _get_client(self) -> VikingDBMemoryClient: region=self.region, ) + def _get_sdk_client(self) -> VikingMem: + return VikingMem( + host="api-knowledgebase.mlp.cn-beijing.volces.com", + region="cn-beijing", + auth=IAM( + ak=self.volcengine_access_key, + sk=self.volcengine_secret_key, + ), + ) + @override def save_memory(self, user_id: str, event_strings: list[str], **kwargs) -> bool: session_id = str(uuid.uuid1()) @@ -139,12 +152,12 @@ def save_memory(self, user_id: str, event_strings: list[str], **kwargs) -> bool: f"Request for add {len(messages)} memory to VikingDB: collection_name={self.index}, metadata={metadata}, session_id={session_id}" ) - client = self._get_client() - response = client.add_messages( - collection_name=self.index, + client = self._get_sdk_client() + collection = client.get_collection(collection_name=self.index) + response =collection.add_session( + session_id=session_id, messages=messages, metadata=metadata, - session_id=session_id, ) logger.debug(f"Response from add memory to VikingDB: {response}") @@ -164,9 +177,12 @@ def search_memory( f"Request for search memory in VikingDB: filter={filter}, collection_name={self.index}, query={query}, limit={top_k}" ) - client = self._get_client() - response = client.search_memory( - collection_name=self.index, query=query, filter=filter, limit=top_k + client = self._get_sdk_client() + collection = client.get_collection(collection_name=self.index) + response = collection.search_memory( + query=query, + filter=filter, + limit=top_k, ) logger.debug(f"Response from search memory in VikingDB: {response}") From 91e96b79374b0a9efa55109d74a9dae4e6739d24 Mon Sep 17 00:00:00 2001 From: "fuyi.0" Date: Mon, 8 Dec 2025 16:43:21 +0800 Subject: [PATCH 2/3] fix: add sts token in viking memory sdk client --- .../integrations/ve_viking_db_memory/ve_viking_db_memory.py | 3 +++ .../long_term_memory_backends/vikingdb_memory_backend.py | 6 ++++-- 2 files changed, 7 insertions(+), 2 deletions(-) diff --git a/veadk/integrations/ve_viking_db_memory/ve_viking_db_memory.py b/veadk/integrations/ve_viking_db_memory/ve_viking_db_memory.py index 43e5d1ae..4b2cad0c 100644 --- a/veadk/integrations/ve_viking_db_memory/ve_viking_db_memory.py +++ b/veadk/integrations/ve_viking_db_memory/ve_viking_db_memory.py @@ -95,6 +95,9 @@ def setHeader(self, header): api_info[key].header[item] = header[item] self.api_info = api_info + def get_host(self): + return self.service_info.host + @staticmethod def get_service_info(host, region, scheme, connection_timeout, socket_timeout): service_info = ServiceInfo( diff --git a/veadk/memory/long_term_memory_backends/vikingdb_memory_backend.py b/veadk/memory/long_term_memory_backends/vikingdb_memory_backend.py index fefc11e2..773b04cf 100644 --- a/veadk/memory/long_term_memory_backends/vikingdb_memory_backend.py +++ b/veadk/memory/long_term_memory_backends/vikingdb_memory_backend.py @@ -122,13 +122,15 @@ def _get_client(self) -> VikingDBMemoryClient: ) def _get_sdk_client(self) -> VikingMem: + client = self._get_client() return VikingMem( - host="api-knowledgebase.mlp.cn-beijing.volces.com", - region="cn-beijing", + host=client.get_host(), + region=self.region, auth=IAM( ak=self.volcengine_access_key, sk=self.volcengine_secret_key, ), + sts_token=self.session_token, ) @override From 4176c4aa08fcad91ad957ead9c9a10ef49704090 Mon Sep 17 00:00:00 2001 From: "fuyi.0" Date: Thu, 11 Dec 2025 14:00:08 +0800 Subject: [PATCH 3/3] fix: pre-commit check --- veadk/integrations/ve_viking_db_memory/ve_viking_db_memory.py | 2 +- .../memory/long_term_memory_backends/vikingdb_memory_backend.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/veadk/integrations/ve_viking_db_memory/ve_viking_db_memory.py b/veadk/integrations/ve_viking_db_memory/ve_viking_db_memory.py index 4b2cad0c..04d1f481 100644 --- a/veadk/integrations/ve_viking_db_memory/ve_viking_db_memory.py +++ b/veadk/integrations/ve_viking_db_memory/ve_viking_db_memory.py @@ -283,4 +283,4 @@ def update_collection( "BuiltinEntityTypes": builtin_entity_types, } res = self.json("UpdateCollection", {}, json.dumps(params)) - return json.loads(res) \ No newline at end of file + return json.loads(res) diff --git a/veadk/memory/long_term_memory_backends/vikingdb_memory_backend.py b/veadk/memory/long_term_memory_backends/vikingdb_memory_backend.py index e63d3ecf..89df8088 100644 --- a/veadk/memory/long_term_memory_backends/vikingdb_memory_backend.py +++ b/veadk/memory/long_term_memory_backends/vikingdb_memory_backend.py @@ -157,7 +157,7 @@ def save_memory(self, user_id: str, event_strings: list[str], **kwargs) -> bool: client = self._get_sdk_client() collection = client.get_collection(collection_name=self.index) - response =collection.add_session( + response = collection.add_session( session_id=session_id, messages=messages, metadata=metadata,