feat: integrate vikingdb-python-sdk for memory backend

Author: ethfoo

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]

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)

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}")