veadk-python 0.2.9__py3-none-any.whl → 0.2.11__py3-none-any.whl

veadk/integrations/ve_viking_db_memory/ve_viking_db_memory.py

@@ -0,0 +1,293 @@
+# 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 volcengine.ApiInfo import ApiInfo
+from volcengine.auth.SignerV4 import SignerV4
+from volcengine.base.Service import Service
+from volcengine.Credentials import Credentials
+from volcengine.ServiceInfo import ServiceInfo
+
+
+class VikingDBMemoryException(Exception):
+    def __init__(self, code, request_id, message=None):
+        self.code = code
+        self.request_id = request_id
+        self.message = "{}, code:{}，request_id:{}".format(
+            message, self.code, self.request_id
+        )
+
+    def __str__(self):
+        return self.message
+
+
+class VikingDBMemoryClient(Service):
+    _instance_lock = threading.Lock()
+
+    def __new__(cls, *args, **kwargs):
+        if not hasattr(VikingDBMemoryClient, "_instance"):
+            with VikingDBMemoryClient._instance_lock:
+                if not hasattr(VikingDBMemoryClient, "_instance"):
+                    VikingDBMemoryClient._instance = object.__new__(cls)
+        return VikingDBMemoryClient._instance
+
+    def __init__(
+        self,
+        host="api-knowledgebase.mlp.cn-beijing.volces.com",
+        region="cn-beijing",
+        ak="",
+        sk="",
+        sts_token="",
+        scheme="http",
+        connection_timeout=30,
+        socket_timeout=30,
+    ):
+        self.service_info = VikingDBMemoryClient.get_service_info(
+            host, region, scheme, connection_timeout, socket_timeout
+        )
+        self.api_info = VikingDBMemoryClient.get_api_info()
+        super(VikingDBMemoryClient, self).__init__(self.service_info, self.api_info)
+        if ak:
+            self.set_ak(ak)
+        if sk:
+            self.set_sk(sk)
+        if sts_token:
+            self.set_session_token(session_token=sts_token)
+        try:
+            self.get_body("Ping", {}, json.dumps({}))
+        except Exception as e:
+            raise VikingDBMemoryException(
+                1000028, "missed", "host or region is incorrect: {}".format(str(e))
+            ) from None
+
+    def setHeader(self, header):
+        api_info = VikingDBMemoryClient.get_api_info()
+        for key in api_info:
+            for item in header:
+                api_info[key].header[item] = header[item]
+        self.api_info = api_info
+
+    @staticmethod
+    def get_service_info(host, region, scheme, connection_timeout, socket_timeout):
+        service_info = ServiceInfo(
+            host,
+            {"Host": host},
+            Credentials("", "", "air", region),
+            connection_timeout,
+            socket_timeout,
+            scheme=scheme,
+        )
+        return service_info
+
+    @staticmethod
+    def get_api_info():
+        api_info = {
+            "CreateCollection": ApiInfo(
+                "POST",
+                "/api/memory/collection/create",
+                {},
+                {},
+                {"Accept": "application/json", "Content-Type": "application/json"},
+            ),
+            "GetCollection": ApiInfo(
+                "POST",
+                "/api/memory/collection/info",
+                {},
+                {},
+                {"Accept": "application/json", "Content-Type": "application/json"},
+            ),
+            "DropCollection": ApiInfo(
+                "POST",
+                "/api/memory/collection/delete",
+                {},
+                {},
+                {"Accept": "application/json", "Content-Type": "application/json"},
+            ),
+            "UpdateCollection": ApiInfo(
+                "POST",
+                "/api/memory/collection/update",
+                {},
+                {},
+                {"Accept": "application/json", "Content-Type": "application/json"},
+            ),
+            "SearchMemory": ApiInfo(
+                "POST",
+                "/api/memory/search",
+                {},
+                {},
+                {"Accept": "application/json", "Content-Type": "application/json"},
+            ),
+            "AddMessages": ApiInfo(
+                "POST",
+                "/api/memory/messages/add",
+                {},
+                {},
+                {"Accept": "application/json", "Content-Type": "application/json"},
+            ),
+            "Ping": ApiInfo(
+                "GET",
+                "/api/memory/ping",
+                {},
+                {},
+                {"Accept": "application/json", "Content-Type": "application/json"},
+            ),
+        }
+        return api_info
+
+    def get_body(self, api, params, body):
+        if api not in self.api_info:
+            raise Exception("no such api")
+        api_info = self.api_info[api]
+        r = self.prepare_request(api_info, params)
+        r.headers["Content-Type"] = "application/json"
+        r.headers["Traffic-Source"] = "SDK"
+        r.body = body
+
+        SignerV4.sign(r, self.service_info.credentials)
+
+        url = r.build()
+        resp = self.session.get(
+            url,
+            headers=r.headers,
+            data=r.body,
+            timeout=(
+                self.service_info.connection_timeout,
+                self.service_info.socket_timeout,
+            ),
+        )
+        if resp.status_code == 200:
+            return json.dumps(resp.json())
+        else:
+            raise Exception(resp.text.encode("utf-8"))
+
+    def get_body_exception(self, api, params, body):
+        try:
+            res = self.get_body(api, params, body)
+        except Exception as e:
+            try:
+                res_json = json.loads(e.args[0].decode("utf-8"))
+            except Exception as e:
+                raise VikingDBMemoryException(
+                    1000028, "missed", "json load res error, res:{}".format(str(e))
+                ) from None
+            code = res_json.get("code", 1000028)
+            request_id = res_json.get("request_id", 1000028)
+            message = res_json.get("message", None)
+
+            raise VikingDBMemoryException(code, request_id, message)
+
+        if res == "":
+            raise VikingDBMemoryException(
+                1000028,
+                "missed",
+                "empty response due to unknown error, please contact customer service",
+            ) from None
+        return res
+
+    def get_exception(self, api, params):
+        try:
+            res = self.get(api, params)
+        except Exception as e:
+            try:
+                res_json = json.loads(e.args[0].decode("utf-8"))
+            except Exception as e:
+                raise VikingDBMemoryException(
+                    1000028, "missed", "json load res error, res:{}".format(str(e))
+                ) from None
+            code = res_json.get("code", 1000028)
+            request_id = res_json.get("request_id", 1000028)
+            message = res_json.get("message", None)
+            raise VikingDBMemoryException(code, request_id, message)
+        if res == "":
+            raise VikingDBMemoryException(
+                1000028,
+                "missed",
+                "empty response due to unknown error, please contact customer service",
+            ) from None
+        return res
+
+    def create_collection(
+        self,
+        collection_name,
+        description="",
+        custom_event_type_schemas=[],
+        custom_entity_type_schemas=[],
+        builtin_event_types=[],
+        builtin_entity_types=[],
+    ):
+        params = {
+            "CollectionName": collection_name,
+            "Description": description,
+            "CustomEventTypeSchemas": custom_event_type_schemas,
+            "CustomEntityTypeSchemas": custom_entity_type_schemas,
+            "BuiltinEventTypes": builtin_event_types,
+            "BuiltinEntityTypes": builtin_entity_types,
+        }
+        res = self.json("CreateCollection", {}, json.dumps(params))
+        return json.loads(res)
+
+    def get_collection(self, collection_name):
+        params = {"CollectionName": collection_name}
+        res = self.json("GetCollection", {}, json.dumps(params))
+        return json.loads(res)
+
+    def drop_collection(self, collection_name):
+        params = {"CollectionName": collection_name}
+        res = self.json("DropCollection", {}, json.dumps(params))
+        return json.loads(res)
+
+    def update_collection(
+        self,
+        collection_name,
+        custom_event_type_schemas=[],
+        custom_entity_type_schemas=[],
+        builtin_event_types=[],
+        builtin_entity_types=[],
+    ):
+        params = {
+            "CollectionName": collection_name,
+            "CustomEventTypeSchemas": custom_event_type_schemas,
+            "CustomEntityTypeSchemas": custom_entity_type_schemas,
+            "BuiltinEventTypes": builtin_event_types,
+            "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)

veadk/knowledgebase/knowledgebase.py

@@ -12,10 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Any, Callable, Literal
+from __future__ import annotations
+
+from typing import Any, Callable, Literal, Union
 
 from pydantic import BaseModel, Field
-from typing_extensions import Union
 
 from veadk.knowledgebase.backends.base_backend import BaseKnowledgebaseBackend
 from veadk.knowledgebase.entry import KnowledgebaseEntry
@@ -54,11 +55,11 @@ def _get_backend_cls(backend: str) -> type[BaseKnowledgebaseBackend]:
     raise ValueError(f"Unsupported knowledgebase backend: {backend}")
 
 
-def build_knowledgebase_index(app_name: str):
-    return f"veadk_kb_{app_name}"
+class KnowledgeBase(BaseModel):
+    name: str = "user_knowledgebase"
 
+    description: str = "This knowledgebase stores some user-related information."
 
-class KnowledgeBase(BaseModel):
     backend: Union[
         Literal["local", "opensearch", "viking", "redis"], BaseKnowledgebaseBackend
     ] = "local"
@@ -73,9 +74,7 @@ class KnowledgeBase(BaseModel):
     """Configuration for the backend"""
 
     top_k: int = 10
-    """Number of top similar documents to retrieve during search.
-
-    Default is 10."""
+    """Number of top similar documents to retrieve during search"""
 
     app_name: str = ""
 
@@ -85,38 +84,27 @@ class KnowledgeBase(BaseModel):
     def model_post_init(self, __context: Any) -> None:
         if isinstance(self.backend, BaseKnowledgebaseBackend):
             self._backend = self.backend
+            self.index = self._backend.index
             logger.info(
                 f"Initialized knowledgebase with provided backend instance {self._backend.__class__.__name__}"
             )
             return
 
-        # must provide at least one of them
-        if not self.app_name and not self.index:
-            raise ValueError(
-                "Either `app_name` or `index` must be provided one of them."
-            )
-
-        # priority use index
-        if self.app_name and self.index:
-            logger.warning(
-                "`app_name` and `index` are both provided, using `index` as the knowledgebase index name."
-            )
+        # Once user define backend config, use it directly
+        if self.backend_config:
+            self._backend = _get_backend_cls(self.backend)(**self.backend_config)
+            return
 
-        # generate index name if `index` not provided but `app_name` is provided
-        if self.app_name and not self.index:
-            self.index = build_knowledgebase_index(self.app_name)
-            logger.info(
-                f"Knowledgebase index is set to {self.index} (generated by the app_name: {self.app_name})."
-            )
+        self.index = self.index or self.app_name
+        if not self.index:
+            raise ValueError("Either `index` or `app_name` must be provided.")
 
         logger.info(
-            f"Initializing knowledgebase: backend={self.backend} top_k={self.top_k}"
-        )
-        self._backend = _get_backend_cls(self.backend)(
-            index=self.index, **self.backend_config if self.backend_config else {}
+            f"Initializing knowledgebase: backend={self.backend} index={self.index} top_k={self.top_k}"
         )
+        self._backend = _get_backend_cls(self.backend)(index=self.index)
         logger.info(
-            f"Initialized knowledgebase with backend {self._backend.__class__.__name__}"
+            f"Initialized knowledgebase with backend {self.backend.__class__.__name__}"
         )
 
     def add_from_directory(self, directory: str, **kwargs) -> bool:
@@ -133,8 +121,7 @@ class KnowledgeBase(BaseModel):
 
     def search(self, query: str, top_k: int = 0, **kwargs) -> list[KnowledgebaseEntry]:
         """Search knowledge from knowledgebase"""
-        if top_k == 0:
-            top_k = self.top_k
+        top_k = top_k if top_k != 0 else self.top_k
 
         _entries = self._backend.search(query=query, top_k=top_k, **kwargs)
 

veadk/memory/__init__.py

@@ -25,7 +25,7 @@ def __getattr__(name):
         from veadk.memory.short_term_memory import ShortTermMemory
 
         return ShortTermMemory
-    if name == "LongTeremMemory":
+    if name == "LongTermMemory":
         from veadk.memory.long_term_memory import LongTermMemory
 
         return LongTermMemory

veadk/memory/long_term_memory.py

@@ -72,10 +72,6 @@ def _get_backend_cls(backend: str) -> type[BaseLongTermMemoryBackend]:
     raise ValueError(f"Unsupported long term memory backend: {backend}")
 
 
-def build_long_term_memory_index(app_name: str, user_id: str):
-    return f"{app_name}_{user_id}"
-
-
 class LongTermMemory(BaseMemoryService, BaseModel):
     backend: Union[
         Literal["local", "opensearch", "redis", "viking", "viking_mem", "mem0"],
@@ -89,46 +85,48 @@ class LongTermMemory(BaseMemoryService, BaseModel):
     top_k: int = 5
     """Number of top similar documents to retrieve during search."""
 
+    index: str = ""
+
     app_name: str = ""
 
     user_id: str = ""
+    """Deprecated attribute"""
 
     def model_post_init(self, __context: Any) -> None:
-        if self.backend == "viking_mem":
-            logger.warning(
-                "The `viking_mem` backend is deprecated, please use `viking` instead."
-            )
-            self.backend = "viking"
-
-        self._backend = None
-
         # Once user define a backend instance, use it directly
         if isinstance(self.backend, BaseLongTermMemoryBackend):
             self._backend = self.backend
+            self.index = self._backend.index
             logger.info(
-                f"Initialized long term memory with provided backend instance {self._backend.__class__.__name__}"
+                f"Initialized long term memory with provided backend instance {self._backend.__class__.__name__}, index={self.index}"
             )
             return
 
+        # Once user define backend config, use it directly
         if self.backend_config:
-            logger.warning(
-                f"Initialized long term memory backend {self.backend} with config. We will ignore `app_name` and `user_id` if provided."
-            )
             self._backend = _get_backend_cls(self.backend)(**self.backend_config)
             return
 
-        if self.app_name and self.user_id:
-            self._index = build_long_term_memory_index(
-                app_name=self.app_name, user_id=self.user_id
-            )
-            logger.info(f"Long term memory index set to {self._index}.")
-            self._backend = _get_backend_cls(self.backend)(
-                index=self._index, **self.backend_config if self.backend_config else {}
+        # Check index
+        self.index = self.index or self.app_name
+        if not self.index:
+            logger.warning(
+                "Attribute `index` or `app_name` not provided, use `default_app` instead."
             )
-        else:
+            self.index = "default_app"
+
+        # Forward compliance
+        if self.backend == "viking_mem":
             logger.warning(
-                "Neither `backend_instance`, `backend_config`, nor (`app_name`/`user_id`) is provided, the long term memory storage will initialize when adding a session."
+                "The `viking_mem` backend is deprecated, change to `viking` instead."
             )
+            self.backend = "viking"
+
+        self._backend = _get_backend_cls(self.backend)(index=self.index)
+
+        logger.info(
+            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]:
         final_events = []
@@ -156,58 +154,32 @@ class LongTermMemory(BaseMemoryService, BaseModel):
         self,
         session: Session,
     ):
-        app_name = session.app_name
         user_id = session.user_id
-
-        if not self._backend and isinstance(self.backend, str):
-            self._index = build_long_term_memory_index(app_name, user_id)
-            self._backend = _get_backend_cls(self.backend)(
-                index=self._index, **self.backend_config if self.backend_config else {}
-            )
-            logger.info(
-                f"Initialize long term memory backend now, index is {self._index}"
-            )
-
-        if not self._index and self._index != build_long_term_memory_index(
-            app_name, user_id
-        ):
-            logger.warning(
-                f"The `app_name` or `user_id` is different from the initialized one, skip add session to memory. Initialized index: {self._index}, current built index: {build_long_term_memory_index(app_name, user_id)}"
-            )
-            return
         event_strings = self._filter_and_convert_events(session.events)
 
         logger.info(
-            f"Adding {len(event_strings)} events to long term memory: index={self._index}"
+            f"Adding {len(event_strings)} events to long term memory: index={self.index}"
+        )
+        self._backend.save_memory(user_id=user_id, event_strings=event_strings)
+        logger.info(
+            f"Added {len(event_strings)} events to long term memory: index={self.index}, user_id={user_id}"
         )
 
-        if self._backend:
-            self._backend.save_memory(event_strings=event_strings, user_id=user_id)
+    @override
+    async def search_memory(
+        self, *, app_name: str, user_id: str, query: str
+    ) -> SearchMemoryResponse:
+        logger.info(f"Search memory with query={query}")
 
-            logger.info(
-                f"Added {len(event_strings)} events to long term memory: index={self._index}"
-            )
-        else:
-            logger.error(
-                "Long term memory backend initialize failed, cannot add session to memory."
+        memory_chunks = []
+        try:
+            memory_chunks = self._backend.search_memory(
+                query=query, top_k=self.top_k, user_id=user_id
             )
-
-    @override
-    async def search_memory(self, *, app_name: str, user_id: str, query: str):
-        # prevent model invoke `load_memory` before add session to this memory
-        if not self._backend:
+        except Exception as e:
             logger.error(
-                "Long term memory backend is not initialized, cannot search memory."
+                f"Exception orrcus during memory search: {e}. Return empty memory chunks"
             )
-            return SearchMemoryResponse(memories=[])
-
-        logger.info(
-            f"Searching long term memory: query={query} index={self._index} top_k={self.top_k}"
-        )
-
-        memory_chunks = self._backend.search_memory(
-            query=query, top_k=self.top_k, user_id=user_id
-        )
 
         memory_events = []
         for memory in memory_chunks:
@@ -235,6 +207,6 @@ class LongTermMemory(BaseMemoryService, BaseModel):
             )
 
         logger.info(
-            f"Return {len(memory_events)} memory events for query: {query} index={self._index}"
+            f"Return {len(memory_events)} memory events for query: {query} index={self.index} user_id={user_id}"
         )
         return SearchMemoryResponse(memories=memory_events)

veadk/memory/long_term_memory_backends/base_backend.py

@@ -25,9 +25,11 @@ class BaseLongTermMemoryBackend(ABC, BaseModel):
         """Check the index name is valid or not"""
 
     @abstractmethod
-    def save_memory(self, event_strings: list[str], **kwargs) -> bool:
+    def save_memory(self, user_id: str, event_strings: list[str], **kwargs) -> bool:
         """Save memory to long term memory backend"""
 
     @abstractmethod
-    def search_memory(self, query: str, top_k: int, **kwargs) -> list[str]:
+    def search_memory(
+        self, user_id: str, query: str, top_k: int, **kwargs
+    ) -> list[str]:
         """Retrieve memory from long term memory backend"""

veadk/memory/long_term_memory_backends/in_memory_backend.py

@@ -29,10 +29,6 @@ class InMemoryLTMBackend(BaseLongTermMemoryBackend):
     embedding_config: EmbeddingModelConfig = Field(default_factory=EmbeddingModelConfig)
     """Embedding model configs"""
 
-    def precheck_index_naming(self):
-        # no checking
-        pass
-
     def model_post_init(self, __context: Any) -> None:
         self._embed_model = OpenAILikeEmbedding(
             model_name=self.embedding_config.name,
@@ -41,8 +37,12 @@ class InMemoryLTMBackend(BaseLongTermMemoryBackend):
         )
         self._vector_index = VectorStoreIndex([], embed_model=self._embed_model)
 
+    def precheck_index_naming(self):
+        # no checking
+        pass
+
     @override
-    def save_memory(self, event_strings: list[str], **kwargs) -> bool:
+    def save_memory(self, user_id: str, event_strings: list[str], **kwargs) -> bool:
         for event_string in event_strings:
             document = Document(text=event_string)
             nodes = self._split_documents([document])
@@ -50,7 +50,9 @@ class InMemoryLTMBackend(BaseLongTermMemoryBackend):
         return True
 
     @override
-    def search_memory(self, query: str, top_k: int, **kwargs) -> list[str]:
+    def search_memory(
+        self, user_id: str, query: str, top_k: int, **kwargs
+    ) -> list[str]:
         _retriever = self._vector_index.as_retriever(similarity_top_k=top_k)
         retrieved_nodes = _retriever.retrieve(query)
         return [node.text for node in retrieved_nodes]