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

veadk/knowledgebase/backends/base_backend.py

@@ -29,19 +29,19 @@ class BaseKnowledgebaseBackend(ABC, BaseModel):
         """
 
     @abstractmethod
-    def add_from_directory(self, directory: str, **kwargs) -> bool:
+    def add_from_directory(self, directory: str, *args, **kwargs) -> bool:
         """Add knowledge from file path to knowledgebase"""
 
     @abstractmethod
-    def add_from_files(self, files: list[str], **kwargs) -> bool:
+    def add_from_files(self, files: list[str], *args, **kwargs) -> bool:
         """Add knowledge (e.g, documents, strings, ...) to knowledgebase"""
 
     @abstractmethod
-    def add_from_text(self, text: str | list[str], **kwargs) -> bool:
+    def add_from_text(self, text: str | list[str], *args, **kwargs) -> bool:
         """Add knowledge from text to knowledgebase"""
 
     @abstractmethod
-    def search(self, **kwargs) -> list:
+    def search(self, *args, **kwargs) -> list:
         """Search knowledge from knowledgebase"""
 
     # Optional methods for future use:

veadk/knowledgebase/backends/vikingdb_knowledge_backend.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import asyncio
+import json
 import re
 from pathlib import Path
 from typing import Any, Literal
@@ -24,9 +25,9 @@ from typing_extensions import override
 import veadk.config  # noqa E401
 from veadk.config import getenv
 from veadk.configs.database_configs import NormalTOSConfig, TOSConfig
-from veadk.consts import DEFAULT_TOS_BUCKET_NAME
 from veadk.knowledgebase.backends.base_backend import BaseKnowledgebaseBackend
 from veadk.knowledgebase.backends.utils import build_vikingdb_knowledgebase_request
+from veadk.knowledgebase.entry import KnowledgebaseEntry
 from veadk.utils.logger import get_logger
 from veadk.utils.misc import formatted_timestamp
 
@@ -48,13 +49,6 @@ def _read_file_to_bytes(file_path: str) -> tuple[bytes, str]:
     return file_content, file_name
 
 
-def _extract_tos_attributes(**kwargs) -> tuple[str, str]:
-    """Extract TOS attributes from kwargs"""
-    tos_bucket_name = kwargs.get("tos_bucket_name", DEFAULT_TOS_BUCKET_NAME)
-    tos_bucket_path = kwargs.get("tos_bucket_path", "knowledgebase")
-    return tos_bucket_name, tos_bucket_path
-
-
 def get_files_in_directory(directory: str):
     dir_path = Path(directory)
     if not dir_path.is_dir():
@@ -109,15 +103,27 @@ class VikingDBKnowledgeBackend(BaseKnowledgebaseBackend):
         )
 
     @override
-    def add_from_directory(self, directory: str, **kwargs) -> bool:
-        """
+    def add_from_directory(
+        self,
+        directory: str,
+        tos_bucket_name: str | None = None,
+        tos_bucket_path: str = "knowledgebase",
+        metadata: dict | None = None,
+        **kwargs,
+    ) -> bool:
+        """Add knowledge from a directory to the knowledgebase.
+
         Args:
-            directory: str, the directory to add to knowledgebase
-            **kwargs:
-                - tos_bucket_name: str, the bucket name of TOS
-                - tos_bucket_path: str, the path of TOS bucket
+            directory (str): The directory to add to knowledgebase.
+            tos_bucket_name (str | None, optional): The bucket name of TOS. Defaults to None.
+            tos_bucket_path (str, optional): The path of TOS bucket. Defaults to "knowledgebase".
+            metadata (dict | None, optional): The metadata of the files. Defaults to None.
+            **kwargs: Additional keyword arguments.
+
+        Returns:
+            bool: True if successful, False otherwise.
         """
-        tos_bucket_name, tos_bucket_path = _extract_tos_attributes(**kwargs)
+        tos_bucket_name = tos_bucket_name or self.tos_config.bucket
         files = get_files_in_directory(directory=directory)
         for _file in files:
             content, file_name = _read_file_to_bytes(_file)
@@ -125,40 +131,64 @@ class VikingDBKnowledgeBackend(BaseKnowledgebaseBackend):
                 content,
                 tos_bucket_name=tos_bucket_name,
                 object_key=f"{tos_bucket_path}/{file_name}",
+                metadata=metadata,
             )
             self._add_doc(tos_url=tos_url)
         return True
 
     @override
-    def add_from_files(self, files: list[str], **kwargs) -> bool:
-        """
+    def add_from_files(
+        self,
+        files: list[str],
+        tos_bucket_name: str | None = None,
+        tos_bucket_path: str = "knowledgebase",
+        metadata: dict | None = None,
+        **kwargs,
+    ) -> bool:
+        """Add knowledge from a directory to the knowledgebase.
+
         Args:
-            files:  list[str], the files to add to knowledgebase
-            **kwargs:
-                - tos_bucket_name: str, the bucket name of TOS
-                - tos_bucket_path: str, the path of TOS bucket
+            files (list[str]): The files to add to knowledgebase.
+            tos_bucket_name (str | None, optional): The bucket name of TOS. Defaults to None.
+            tos_bucket_path (str, optional): The path of TOS bucket. Defaults to "knowledgebase".
+            metadata (dict | None, optional): The metadata of the files. Defaults to None.
+            **kwargs: Additional keyword arguments.
+
+        Returns:
+            bool: True if successful, False otherwise.
         """
-        tos_bucket_name, tos_bucket_path = _extract_tos_attributes(**kwargs)
+        tos_bucket_name = tos_bucket_name or self.tos_config.bucket
         for _file in files:
             content, file_name = _read_file_to_bytes(_file)
             tos_url = self._upload_bytes_to_tos(
                 content,
                 tos_bucket_name=tos_bucket_name,
                 object_key=f"{tos_bucket_path}/{file_name}",
+                metadata=metadata,
             )
             self._add_doc(tos_url=tos_url)
         return True
 
     @override
-    def add_from_text(self, text: str | list[str], **kwargs) -> bool:
-        """
+    def add_from_text(
+        self,
+        text: str | list[str],
+        tos_bucket_name: str | None = None,
+        tos_bucket_path: str = "knowledgebase",
+        metadata: dict | None = None,
+        **kwargs,
+    ) -> bool:
+        """Add knowledge from text to the knowledgebase.
+
         Args:
-            text:   str or list[str], the text to add to knowledgebase
-            **kwargs:
-                - tos_bucket_name: str, the bucket name of TOS
-                - tos_bucket_path: str, the path of TOS bucket
+            text (str | list[str]): The text to add to knowledgebase.
+            tos_bucket_name (str | None, optional): The bucket name of TOS. Defaults to None.
+            tos_bucket_path (str, optional): The path of TOS bucket. Defaults to "knowledgebase".
+
+        Returns:
+            bool: True if successful, False otherwise.
         """
-        tos_bucket_name, tos_bucket_path = _extract_tos_attributes(**kwargs)
+        tos_bucket_name = tos_bucket_name or self.tos_config.bucket
         if isinstance(text, list):
             object_keys = kwargs.get(
                 "tos_object_keys",
@@ -170,7 +200,7 @@ class VikingDBKnowledgeBackend(BaseKnowledgebaseBackend):
             for _text, _object_key in zip(text, object_keys):
                 _content = _text.encode("utf-8")
                 tos_url = self._upload_bytes_to_tos(
-                    _content, tos_bucket_name, _object_key
+                    _content, tos_bucket_name, _object_key, metadata=metadata
                 )
                 self._add_doc(tos_url=tos_url)
             return True
@@ -179,26 +209,42 @@ class VikingDBKnowledgeBackend(BaseKnowledgebaseBackend):
             object_key = kwargs.get(
                 "object_key", f"veadk/knowledgebase/{formatted_timestamp()}.txt"
             )
-            tos_url = self._upload_bytes_to_tos(content, tos_bucket_name, object_key)
+            tos_url = self._upload_bytes_to_tos(
+                content, tos_bucket_name, object_key, metadata=metadata
+            )
             self._add_doc(tos_url=tos_url)
         else:
             raise ValueError("text must be str or list[str]")
         return True
 
-    def add_from_bytes(self, content: bytes, file_name: str, **kwargs) -> bool:
-        """
+    def add_from_bytes(
+        self,
+        content: bytes,
+        file_name: str,
+        tos_bucket_name: str | None = None,
+        tos_bucket_path: str = "knowledgebase",
+        metadata: dict | None = None,
+        **kwargs,
+    ) -> bool:
+        """Add knowledge from bytes to the knowledgebase.
+
         Args:
-            content: bytes, the content to add to knowledgebase, bytes
-            file_name: str, the file name of the content
-            **kwargs:
-                - tos_bucket_name: str, the bucket name of TOS
-                - tos_bucket_path: str, the path of TOS bucket
+            content (bytes): The content to add to knowledgebase.
+            file_name (str): The file name of the content.
+            tos_bucket_name (str | None, optional): The bucket name of TOS. Defaults to None.
+            tos_bucket_path (str, optional): The path of TOS bucket. Defaults to "knowledgebase".
+            metadata (dict | None, optional): The metadata of the files. Defaults to None.
+            **kwargs: Additional keyword arguments.
+
+        Returns:
+            bool: True if successful, False otherwise.
         """
-        tos_bucket_name, tos_bucket_path = _extract_tos_attributes(**kwargs)
+        tos_bucket_name = tos_bucket_name or self.tos_config.bucket
         tos_url = self._upload_bytes_to_tos(
             content,
             tos_bucket_name=tos_bucket_name,
             object_key=f"{tos_bucket_path}/{file_name}",
+            metadata=metadata,
         )
         response = self._add_doc(tos_url=tos_url)
         if response["code"] == 0:
@@ -206,8 +252,16 @@ class VikingDBKnowledgeBackend(BaseKnowledgebaseBackend):
         return False
 
     @override
-    def search(self, query: str, top_k: int = 5) -> list:
-        return self._search_knowledge(query=query, top_k=top_k)
+    def search(
+        self,
+        query: str,
+        top_k: int = 5,
+        metadata: dict | None = None,
+        rerank: bool = True,
+    ) -> list:
+        return self._search_knowledge(
+            query=query, top_k=top_k, metadata=metadata, rerank=rerank
+        )
 
     def delete_collection(self) -> bool:
         DELETE_COLLECTION_PATH = "/api/knowledge/collection/delete"
@@ -330,7 +384,7 @@ class VikingDBKnowledgeBackend(BaseKnowledgebaseBackend):
         response = self._do_request(
             body={
                 "name": self.index,
-                "project": "default",
+                "project": self.volcengine_project,
                 "description": "Created by Volcengine Agent Development Kit (VeADK).",
             },
             path=CREATE_COLLECTION_PATH,
@@ -343,10 +397,27 @@ class VikingDBKnowledgeBackend(BaseKnowledgebaseBackend):
             )
 
     def _upload_bytes_to_tos(
-        self, content: bytes, tos_bucket_name: str, object_key: str
+        self,
+        content: bytes,
+        tos_bucket_name: str,
+        object_key: str,
+        metadata: dict | None = None,
     ) -> str:
+        # Here, we set the metadata via the TOS object, ref: https://www.volcengine.com/docs/84313/1254624
         self._tos_client.bucket_name = tos_bucket_name
-        asyncio.run(self._tos_client.upload(object_key=object_key, data=content))
+        coro = self._tos_client.upload(
+            object_key=object_key,
+            bucket_name=tos_bucket_name,
+            data=content,
+            metadata=metadata,
+        )
+        try:
+            loop = asyncio.get_running_loop()
+            loop.run_until_complete(
+                coro
+            ) if not loop.is_running() else asyncio.ensure_future(coro)
+        except RuntimeError:
+            asyncio.run(coro)
         return f"{self._tos_client.bucket_name}/{object_key}"
 
     def _add_doc(self, tos_url: str) -> Any:
@@ -355,7 +426,7 @@ class VikingDBKnowledgeBackend(BaseKnowledgebaseBackend):
         response = self._do_request(
             body={
                 "collection_name": self.index,
-                "project": "default",
+                "project": self.volcengine_project,
                 "add_type": "tos",
                 "tos_path": tos_url,
             },
@@ -364,14 +435,43 @@ class VikingDBKnowledgeBackend(BaseKnowledgebaseBackend):
         )
         return response
 
-    def _search_knowledge(self, query: str, top_k: int = 5) -> list[str]:
+    def _search_knowledge(
+        self,
+        query: str,
+        top_k: int = 5,
+        metadata: dict | None = None,
+        rerank: bool = True,
+        chunk_diffusion_count: int | None = 3,
+    ) -> list[KnowledgebaseEntry]:
         SEARCH_KNOWLEDGE_PATH = "/api/knowledge/collection/search_knowledge"
 
+        query_param = (
+            {
+                "doc_filter": {
+                    "op": "and",
+                    "conds": [
+                        {"op": "must", "field": str(k), "conds": [str(v)]}
+                        for k, v in metadata.items()
+                    ],
+                }
+            }
+            if metadata
+            else None
+        )
+
+        post_precessing = {
+            "rerank_swich": rerank,
+            "chunk_diffusion_count": chunk_diffusion_count,
+        }
+
         response = self._do_request(
             body={
                 "name": self.index,
+                "project": self.volcengine_project,
                 "query": query,
                 "limit": top_k,
+                "query_param": query_param,
+                "post_processing": post_precessing,
             },
             path=SEARCH_KNOWLEDGE_PATH,
             method="POST",
@@ -382,11 +482,19 @@ class VikingDBKnowledgeBackend(BaseKnowledgebaseBackend):
                 f"Error during knowledge search: {response.get('code')}, message: {response.get('message')}"
             )
 
-        search_result_list = response.get("data", {}).get("result_list", [])
+        entries = []
+        for result in response.get("data", {}).get("result_list", []):
+            doc_meta_raw_str = result.get("doc_info", {}).get("doc_meta")
+            doc_meta_list = json.loads(doc_meta_raw_str) if doc_meta_raw_str else []
+            metadata = {}
+            for meta in doc_meta_list:
+                metadata[meta["field_name"]] = meta["field_value"]
 
-        return [
-            search_result.get("content", "") for search_result in search_result_list
-        ]
+            entries.append(
+                KnowledgebaseEntry(content=result.get("content", ""), metadata=metadata)
+            )
+
+        return entries
 
     def _do_request(
         self,
@@ -409,4 +517,8 @@ class VikingDBKnowledgeBackend(BaseKnowledgebaseBackend):
             headers=request.headers,
             data=request.body,
         )
+        if not response.ok:
+            logger.error(
+                f"VikingDBKnowledgeBackend error during request: {response.json()}"
+            )
         return response.json()

veadk/knowledgebase/entry.py

@@ -0,0 +1,25 @@
+# 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 pydantic import BaseModel
+
+
+class KnowledgebaseEntry(BaseModel):
+    """Represents a single entry in the knowledgebase."""
+
+    # The main content of the knowledgebase entry.
+    content: str
+
+    # Optional metadata associated with the entry.
+    metadata: dict | None = None

veadk/knowledgebase/knowledgebase.py

@@ -18,6 +18,7 @@ from pydantic import BaseModel, Field
 from typing_extensions import Union
 
 from veadk.knowledgebase.backends.base_backend import BaseKnowledgebaseBackend
+from veadk.knowledgebase.entry import KnowledgebaseEntry
 from veadk.utils.logger import get_logger
 
 logger = get_logger(__name__)
@@ -72,8 +73,8 @@ class KnowledgeBase(BaseModel):
     """Configuration for the backend"""
 
     top_k: int = 10
-    """Number of top similar documents to retrieve during search. 
-    
+    """Number of top similar documents to retrieve during search.
+
     Default is 10."""
 
     app_name: str = ""
@@ -130,11 +131,25 @@ class KnowledgeBase(BaseModel):
         """Add knowledge from text to knowledgebase"""
         return self._backend.add_from_text(text=text, **kwargs)
 
-    def search(self, query: str, top_k: int = 0, **kwargs) -> list[str]:
+    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
-        return self._backend.search(query=query, top_k=top_k, **kwargs)
+
+        _entries = self._backend.search(query=query, top_k=top_k, **kwargs)
+
+        entries = []
+        for entry in _entries:
+            if isinstance(entry, KnowledgebaseEntry):
+                entries.append(entry)
+            elif isinstance(entry, str):
+                entries.append(KnowledgebaseEntry(content=entry))
+            else:
+                logger.error(
+                    f"Unsupported entry type from backend search method: {type(entry)} with {entry}. Expected `KnowledgebaseEntry` or `str`. Skip for this entry."
+                )
+
+        return entries
 
     def __getattr__(self, name) -> Callable:
         """In case of knowledgebase have no backends' methods (`delete`, `list_chunks`, etc)

veadk/memory/long_term_memory.py

@@ -62,6 +62,12 @@ def _get_backend_cls(backend: str) -> type[BaseLongTermMemoryBackend]:
             )
 
             return RedisLTMBackend
+        case "mem0":
+            from veadk.memory.long_term_memory_backends.mem0_backend import (
+                Mem0LTMBackend,
+            )
+
+            return Mem0LTMBackend
 
     raise ValueError(f"Unsupported long term memory backend: {backend}")
 
@@ -72,7 +78,7 @@ def build_long_term_memory_index(app_name: str, user_id: str):
 
 class LongTermMemory(BaseMemoryService, BaseModel):
     backend: Union[
-        Literal["local", "opensearch", "redis", "viking", "viking_mem"],
+        Literal["local", "opensearch", "redis", "viking", "viking_mem", "mem0"],
         BaseLongTermMemoryBackend,
     ] = "opensearch"
     """Long term memory backend type"""
@@ -88,6 +94,12 @@ class LongTermMemory(BaseMemoryService, BaseModel):
     user_id: str = ""
 
     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
@@ -147,12 +159,6 @@ class LongTermMemory(BaseMemoryService, BaseModel):
         app_name = session.app_name
         user_id = session.user_id
 
-        if 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
-
         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)(
@@ -162,6 +168,13 @@ class LongTermMemory(BaseMemoryService, BaseModel):
                 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(

veadk/memory/long_term_memory_backends/mem0_backend.py

@@ -0,0 +1,129 @@
+# 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 typing import Any
+from typing_extensions import override
+from pydantic import Field
+
+from veadk.configs.database_configs import Mem0Config
+
+
+from veadk.memory.long_term_memory_backends.base_backend import (
+    BaseLongTermMemoryBackend,
+)
+from veadk.utils.logger import get_logger
+
+logger = get_logger(__name__)
+
+try:
+    from mem0 import MemoryClient
+
+except ImportError:
+    logger.error(
+        "Failed to import mem0 or dotenv. Please install them with 'pip install mem0 '"
+    )
+    raise ImportError("Required packages not installed: mem0")
+
+
+class Mem0LTMBackend(BaseLongTermMemoryBackend):
+    """Mem0 long term memory backend implementation"""
+
+    mem0_config: Mem0Config = Field(default_factory=Mem0Config)
+
+    def model_post_init(self, __context: Any) -> None:
+        """Initialize Mem0 client"""
+
+        try:
+            self._mem0_client = MemoryClient(
+                # base_url=self.mem0_config.base_url,        # mem0 endpoint
+                api_key=self.mem0_config.api_key,  # mem0 API key
+            )
+            logger.info(f"Initialized Mem0 client for index: {self.index}")
+        except Exception as e:
+            logger.error(f"Failed to initialize Mem0 client: {str(e)}")
+            raise
+
+    def precheck_index_naming(self):
+        """Check if the index name is valid
+        For Mem0, there are no specific naming constraints
+        """
+        pass
+
+    @override
+    def save_memory(self, event_strings: list[str], **kwargs) -> bool:
+        """Save memory to Mem0
+
+        Args:
+            event_strings: List of event strings to save
+            **kwargs: Additional parameters, including 'user_id' for Mem0
+
+        Returns:
+            bool: True if saved successfully, False otherwise
+        """
+        user_id = kwargs.get("user_id", "default_user")
+
+        try:
+            logger.info(
+                f"Saving {len(event_strings)} events to Mem0 for user: {user_id}"
+            )
+
+            for event_string in event_strings:
+                # Save event string to Mem0
+                result = self._mem0_client.add(
+                    [{"role": "user", "content": event_string}],
+                    user_id=user_id,
+                    output_format="v1.1",
+                )
+                logger.debug(f"Saved memory result: {result}")
+
+            logger.info(f"Successfully saved {len(event_strings)} events to Mem0")
+            return True
+        except Exception as e:
+            logger.error(f"Failed to save memory to Mem0: {str(e)}")
+            return False
+
+    @override
+    def search_memory(self, query: str, top_k: int, **kwargs) -> list[str]:
+        """Search memory from Mem0
+
+        Args:
+            query: Search query
+            top_k: Number of results to return
+            **kwargs: Additional parameters, including 'user_id' for Mem0
+
+        Returns:
+            list[str]: List of memory strings
+        """
+        user_id = kwargs.get("user_id", "default_user")
+
+        try:
+            logger.info(
+                f"Searching Mem0 for query: {query}, user: {user_id}, top_k: {top_k}"
+            )
+
+            memories = self._mem0_client.search(
+                query, user_id=user_id, output_format="v1.1", top_k=top_k
+            )
+
+            memory_list = []
+            if memories.get("results", []):
+                for mem in memories["results"]:
+                    if "memory" in mem:
+                        memory_list.append(mem["memory"])
+
+            logger.info(f"Found {len(memory_list)} memories matching query: {query}")
+            return memory_list
+        except Exception as e:
+            logger.error(f"Failed to search memory from Mem0: {str(e)}")
+            return []