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

veadk/knowledgebase/knowledgebase.py

@@ -11,168 +11,149 @@
 # 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 io
-import os.path
-from typing import Any, BinaryIO, Literal, TextIO
 
-from pydantic import BaseModel
+from typing import Any, Callable, Literal
 
-from veadk.database.database_adapter import get_knowledgebase_database_adapter
-from veadk.database.database_factory import DatabaseFactory
-from veadk.utils.misc import formatted_timestamp
+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__)
 
 
+def _get_backend_cls(backend: str) -> type[BaseKnowledgebaseBackend]:
+    match backend:
+        case "local":
+            from veadk.knowledgebase.backends.in_memory_backend import (
+                InMemoryKnowledgeBackend,
+            )
+
+            return InMemoryKnowledgeBackend
+        case "opensearch":
+            from veadk.knowledgebase.backends.opensearch_backend import (
+                OpensearchKnowledgeBackend,
+            )
+
+            return OpensearchKnowledgeBackend
+        case "viking":
+            from veadk.knowledgebase.backends.vikingdb_knowledge_backend import (
+                VikingDBKnowledgeBackend,
+            )
+
+            return VikingDBKnowledgeBackend
+        case "redis":
+            from veadk.knowledgebase.backends.redis_backend import (
+                RedisKnowledgeBackend,
+            )
+
+            return RedisKnowledgeBackend
+
+    raise ValueError(f"Unsupported knowledgebase backend: {backend}")
+
+
 def build_knowledgebase_index(app_name: str):
     return f"veadk_kb_{app_name}"
 
 
 class KnowledgeBase(BaseModel):
-    backend: Literal["local", "opensearch", "viking", "redis", "mysql"] = "local"
+    backend: Union[
+        Literal["local", "opensearch", "viking", "redis"], BaseKnowledgebaseBackend
+    ] = "local"
+    """Knowledgebase backend type. Supported backends are:
+    - local: In-memory knowledgebase, data will be lost when the program exits.
+    - opensearch: OpenSearch knowledgebase, requires an OpenSearch cluster.
+    - viking: Volcengine VikingDB knowledgebase, requires VikingDB service.
+    - redis: Redis knowledgebase, requires Redis with vector search capability.
+    Default is `local`."""
+
+    backend_config: dict = Field(default_factory=dict)
+    """Configuration for the backend"""
+
     top_k: int = 10
-    db_config: Any | None = None
+    """Number of top similar documents to retrieve during search.
 
-    def model_post_init(self, __context: Any) -> None:
-        logger.info(
-            f"Initializing knowledgebase: backend={self.backend} top_k={self.top_k}"
-        )
+    Default is 10."""
 
-        self._db_client = DatabaseFactory.create(
-            backend=self.backend, config=self.db_config
-        )
-        self._adapter = get_knowledgebase_database_adapter(self._db_client)
+    app_name: str = ""
 
-        logger.info(
-            f"Initialized knowledgebase: db_client={self._db_client.__class__.__name__} adapter={self._adapter}"
-        )
+    index: str = ""
+    """The name of the knowledgebase index. If not provided, it will be generated based on the `app_name`."""
 
-    def add(
-        self,
-        data: str | list[str] | TextIO | BinaryIO | bytes,
-        app_name: str,
-        **kwargs,
-    ):
-        """
-        Add documents to the vector database.
-        Args:
-            data (str | list[str] | TextIO | BinaryIO | bytes): The data to be added.
-                - str: A single file path. (viking only)
-                - list[str]: A list of file paths.
-                - TextIO: A file object (TextIO). (viking only) file descriptor
-                - BinaryIO: A file object (BinaryIO). (viking only) file descriptor
-                - bytes: Binary data. (viking only) binary data (f.read())
-            app_name: index name
-            **kwargs: Additional keyword arguments.
-                - file_name (str | list[str]): The file name or a list of file names (including suffix). (viking only)
-        """
-        if self.backend != "viking" and not (
-            isinstance(data, str) or isinstance(data, list)
-        ):
+    def model_post_init(self, __context: Any) -> None:
+        if isinstance(self.backend, BaseKnowledgebaseBackend):
+            self._backend = self.backend
+            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(
-                "Only vikingdb supports uploading files or file characters."
+                "Either `app_name` or `index` must be provided one of them."
             )
 
-        index = build_knowledgebase_index(app_name)
-        logger.info(f"Adding documents to knowledgebase: index={index}")
-
-        if self.backend == "viking":
-            # Case 1: Handling file paths or lists of file paths (str)
-            if isinstance(data, str) and os.path.isfile(data):
-                # Get the file name (including the suffix)
-                if "file_name" not in kwargs or not kwargs["file_name"]:
-                    kwargs["file_name"] = os.path.basename(data)
-                return self._adapter.add(data=data, index=index, **kwargs)
-            # Case 2: Handling when list[str] is a full path  (list[str])
-            if isinstance(data, list):
-                if all(isinstance(item, str) for item in data):
-                    all_paths = all(os.path.isfile(item) for item in data)
-                    all_not_paths = all(not os.path.isfile(item) for item in data)
-                    if all_paths:
-                        if "file_name" not in kwargs or not kwargs["file_name"]:
-                            kwargs["file_name"] = [
-                                os.path.basename(item) for item in data
-                            ]
-                        return self._adapter.add(data=data, index=index, **kwargs)
-                    elif (
-                        not all_not_paths
-                    ):  # Prevent the occurrence of non-existent paths
-                        # There is a mixture of paths and non-paths
-                        raise ValueError(
-                            "Mixed file paths and content strings in list are not allowed"
-                        )
-            # Case 3: Handling strings or string arrays (content)  (str or list[str])
-            if isinstance(data, str) or (
-                isinstance(data, list) and all(isinstance(item, str) for item in data)
-            ):
-                if "file_name" not in kwargs or not kwargs["file_name"]:
-                    if isinstance(data, str):
-                        kwargs["file_name"] = f"{formatted_timestamp()}.txt"
-                    else:  # list[str] without file_names
-                        prefix_file_name = formatted_timestamp()
-                        kwargs["file_name"] = [
-                            f"{prefix_file_name}_{i}.txt" for i in range(len(data))
-                        ]
-                return self._adapter.add(data=data, index=index, **kwargs)
-
-            # Case 4: Handling binary data (bytes)
-            if isinstance(data, bytes):
-                # user must give file_name
-                if "file_name" not in kwargs:
-                    raise ValueError("file_name must be provided for binary data")
-                return self._adapter.add(data=data, index=index, **kwargs)
-
-            # Case 5: Handling file objects TextIO or BinaryIO
-            if isinstance(data, (io.TextIOWrapper, io.BufferedReader)):
-                if not kwargs.get("file_name") and hasattr(data, "name"):
-                    kwargs["file_name"] = os.path.basename(data.name)
-                return self._adapter.add(data=data, index=index, **kwargs)
-            # Case6: Unsupported data type
-            raise TypeError(f"Unsupported data type: {type(data)}")
-
-        if not isinstance(data, list):
-            raise TypeError(
-                f"Unsupported data type: {type(data)}. Only viking support file_path and file bytes"
+        # 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."
             )
-        # not viking
-        return self._adapter.add(data=data, index=index, **kwargs)
 
-    def search(self, query: str, app_name: str, top_k: int | None = None) -> list[str]:
-        top_k = self.top_k if top_k is None else top_k
+        # 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})."
+            )
 
         logger.info(
-            f"Searching knowledgebase: app_name={app_name} query={query} top_k={top_k}"
+            f"Initializing knowledgebase: backend={self.backend} top_k={self.top_k}"
         )
-        index = build_knowledgebase_index(app_name)
-        result = self._adapter.query(query=query, index=index, top_k=top_k)
-        if len(result) == 0:
-            logger.warning(f"No documents found in knowledgebase. Query: {query}")
-        return result
-
-    def delete(self, app_name: str) -> bool:
-        index = build_knowledgebase_index(app_name)
-        return self._adapter.delete(index=index)
-
-    def delete_doc(self, app_name: str, id: str) -> bool:
-        index = build_knowledgebase_index(app_name)
-        return self._adapter.delete_doc(index=index, id=id)
-
-    def list_chunks(
-        self, app_name: str, offset: int = 0, limit: int = 100
-    ) -> list[dict]:
-        index = build_knowledgebase_index(app_name)
-        return self._adapter.list_chunks(index=index, offset=offset, limit=limit)
-
-    def list_docs(self, app_name: str, offset: int = 0, limit: int = 100) -> list[dict]:
-        if self.backend == "viking":
-            index = build_knowledgebase_index(app_name)
-            return self._adapter.list_docs(index=index, offset=offset, limit=limit)
-        else:
-            raise NotImplementedError(
-                f"list_docs not supported for {self.backend}, only viking support list_docs"
-            )
+        self._backend = _get_backend_cls(self.backend)(
+            index=self.index, **self.backend_config if self.backend_config else {}
+        )
+        logger.info(
+            f"Initialized knowledgebase with backend {self._backend.__class__.__name__}"
+        )
+
+    def add_from_directory(self, directory: str, **kwargs) -> bool:
+        """Add knowledge from file path to knowledgebase"""
+        return self._backend.add_from_directory(directory=directory, **kwargs)
+
+    def add_from_files(self, files: list[str], **kwargs) -> bool:
+        """Add knowledge (e.g, documents, strings, ...) to knowledgebase"""
+        return self._backend.add_from_files(files=files, **kwargs)
+
+    def add_from_text(self, text: str | list[str], **kwargs) -> bool:
+        """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[KnowledgebaseEntry]:
+        """Search knowledge from knowledgebase"""
+        if top_k == 0:
+            top_k = self.top_k
 
-    def exists(self, app_name: str) -> bool:
-        index = build_knowledgebase_index(app_name)
-        return self._adapter.index_exists(index=index)
+        _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)
+
+        For example, knowledgebase.delete(...) -> self._backend.delete(...)
+        """
+        return getattr(self._backend, name)

veadk/memory/__init__.py

@@ -11,3 +11,25 @@
 # 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 TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from veadk.memory.long_term_memory import LongTermMemory
+    from veadk.memory.short_term_memory import ShortTermMemory
+
+
+# Lazy loading for classes
+def __getattr__(name):
+    if name == "ShortTermMemory":
+        from veadk.memory.short_term_memory import ShortTermMemory
+
+        return ShortTermMemory
+    if name == "LongTeremMemory":
+        from veadk.memory.long_term_memory import LongTermMemory
+
+        return LongTermMemory
+    raise AttributeError(f"module 'veadk.memory' has no attribute '{name}'")
+
+
+__all__ = ["ShortTermMemory", "LongTermMemory"]

veadk/memory/long_term_memory.py

@@ -25,45 +25,110 @@ from google.adk.memory.base_memory_service import (
 from google.adk.memory.memory_entry import MemoryEntry
 from google.adk.sessions import Session
 from google.genai import types
-from pydantic import BaseModel
-from typing_extensions import override
+from pydantic import BaseModel, Field
+from typing_extensions import Union, override
 
-from veadk.database import DatabaseFactory
-from veadk.database.database_adapter import get_long_term_memory_database_adapter
+from veadk.memory.long_term_memory_backends.base_backend import (
+    BaseLongTermMemoryBackend,
+)
 from veadk.utils.logger import get_logger
 
 logger = get_logger(__name__)
 
 
+def _get_backend_cls(backend: str) -> type[BaseLongTermMemoryBackend]:
+    match backend:
+        case "local":
+            from veadk.memory.long_term_memory_backends.in_memory_backend import (
+                InMemoryLTMBackend,
+            )
+
+            return InMemoryLTMBackend
+        case "opensearch":
+            from veadk.memory.long_term_memory_backends.opensearch_backend import (
+                OpensearchLTMBackend,
+            )
+
+            return OpensearchLTMBackend
+        case "viking":
+            from veadk.memory.long_term_memory_backends.vikingdb_memory_backend import (
+                VikingDBLTMBackend,
+            )
+
+            return VikingDBLTMBackend
+        case "redis":
+            from veadk.memory.long_term_memory_backends.redis_backend import (
+                RedisLTMBackend,
+            )
+
+            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}")
+
+
 def build_long_term_memory_index(app_name: str, user_id: str):
     return f"{app_name}_{user_id}"
 
 
 class LongTermMemory(BaseMemoryService, BaseModel):
-    backend: Literal[
-        "local", "opensearch", "redis", "mysql", "viking", "viking_mem"
+    backend: Union[
+        Literal["local", "opensearch", "redis", "viking", "viking_mem", "mem0"],
+        BaseLongTermMemoryBackend,
     ] = "opensearch"
+    """Long term memory backend type"""
+
+    backend_config: dict = Field(default_factory=dict)
+    """Long term memory backend configuration"""
+
     top_k: int = 5
+    """Number of top similar documents to retrieve during search."""
+
+    app_name: str = ""
+
+    user_id: str = ""
 
     def model_post_init(self, __context: Any) -> None:
-        if self.backend == "viking":
+        if self.backend == "viking_mem":
             logger.warning(
-                "`viking` backend is deprecated, switching to `viking_mem` backend."
+                "The `viking_mem` backend is deprecated, please use `viking` instead."
             )
-            self.backend = "viking_mem"
+            self.backend = "viking"
 
-        logger.info(
-            f"Initializing long term memory: backend={self.backend} top_k={self.top_k}"
-        )
+        self._backend = None
 
-        self._db_client = DatabaseFactory.create(
-            backend=self.backend,
-        )
-        self._adapter = get_long_term_memory_database_adapter(self._db_client)
+        # Once user define a backend instance, use it directly
+        if isinstance(self.backend, BaseLongTermMemoryBackend):
+            self._backend = self.backend
+            logger.info(
+                f"Initialized long term memory with provided backend instance {self._backend.__class__.__name__}"
+            )
+            return
 
-        logger.info(
-            f"Initialized long term memory: db_client={self._db_client.__class__.__name__} adapter={self._adapter}"
-        )
+        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 {}
+            )
+        else:
+            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."
+            )
 
     def _filter_and_convert_events(self, events: list[Event]) -> list[str]:
         final_events = []
@@ -91,40 +156,58 @@ 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)
-        index = build_long_term_memory_index(session.app_name, session.user_id)
 
         logger.info(
-            f"Adding {len(event_strings)} events to long term memory: index={index}"
+            f"Adding {len(event_strings)} events to long term memory: index={self._index}"
         )
 
-        # check if viking memory database, should give a user id： if/else
-        if self.backend == "viking_mem":
-            self._adapter.add(data=event_strings, index=index, user_id=session.user_id)
-        else:
-            self._adapter.add(data=event_strings, index=index)
+        if self._backend:
+            self._backend.save_memory(event_strings=event_strings, user_id=user_id)
 
-        logger.info(
-            f"Added {len(event_strings)} events to long term memory: index={index}"
-        )
+            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."
+            )
 
     @override
     async def search_memory(self, *, app_name: str, user_id: str, query: str):
-        index = build_long_term_memory_index(app_name, user_id)
+        # prevent model invoke `load_memory` before add session to this memory
+        if not self._backend:
+            logger.error(
+                "Long term memory backend is not initialized, cannot search memory."
+            )
+            return SearchMemoryResponse(memories=[])
 
         logger.info(
-            f"Searching long term memory: query={query} index={index} top_k={self.top_k}"
+            f"Searching long term memory: query={query} index={self._index} top_k={self.top_k}"
         )
 
-        # user id if viking memory db
-        if self.backend == "viking_mem":
-            memory_chunks = self._adapter.query(
-                query=query, index=index, top_k=self.top_k, user_id=user_id
-            )
-        else:
-            memory_chunks = self._adapter.query(
-                query=query, index=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:
@@ -152,6 +235,6 @@ class LongTermMemory(BaseMemoryService, BaseModel):
             )
 
         logger.info(
-            f"Return {len(memory_events)} memory events for query: {query} index={index}"
+            f"Return {len(memory_events)} memory events for query: {query} index={self._index}"
         )
         return SearchMemoryResponse(memories=memory_events)

veadk/{database/base_database.py → memory/long_term_memory_backends/base_backend.py}

@@ -13,33 +13,21 @@
 # limitations under the License.
 
 from abc import ABC, abstractmethod
-from typing import Any
 
+from pydantic import BaseModel
 
-class DatabaseType:
-    LOCAL = "local"
-    RELATIONAL = "relational"
-    VECTOR = "vector"
-    KV = "kv"
 
+class BaseLongTermMemoryBackend(ABC, BaseModel):
+    index: str
 
-class BaseDatabase(ABC):
-    """Base class for database.
-
-    Args:
-        type: type of the database
-
-    Note:
-        No `update` function support currently.
-    """
-
-    def __init__(self):
-        pass
-
-    def add(self, texts: list[Any], **kwargs: Any): ...
+    @abstractmethod
+    def precheck_index_naming(self):
+        """Check the index name is valid or not"""
 
     @abstractmethod
-    def delete(self, **kwargs: Any): ...
+    def save_memory(self, event_strings: list[str], **kwargs) -> bool:
+        """Save memory to long term memory backend"""
 
     @abstractmethod
-    def query(self, query: str, **kwargs: Any) -> list[str]: ...
+    def search_memory(self, 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

@@ -0,0 +1,65 @@
+# 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 llama_index.core import Document, VectorStoreIndex
+from llama_index.core.schema import BaseNode
+from llama_index.embeddings.openai_like import OpenAILikeEmbedding
+from pydantic import Field
+from typing_extensions import Any, override
+
+from veadk.configs.model_configs import EmbeddingModelConfig
+from veadk.knowledgebase.backends.utils import get_llama_index_splitter
+from veadk.memory.long_term_memory_backends.base_backend import (
+    BaseLongTermMemoryBackend,
+)
+
+
+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,
+            api_key=self.embedding_config.api_key,
+            api_base=self.embedding_config.api_base,
+        )
+        self._vector_index = VectorStoreIndex([], embed_model=self._embed_model)
+
+    @override
+    def save_memory(self, event_strings: list[str], **kwargs) -> bool:
+        for event_string in event_strings:
+            document = Document(text=event_string)
+            nodes = self._split_documents([document])
+            self._vector_index.insert_nodes(nodes)
+        return True
+
+    @override
+    def search_memory(self, 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]
+
+    def _split_documents(self, documents: list[Document]) -> list[BaseNode]:
+        """Split document into chunks"""
+        nodes = []
+        for document in documents:
+            splitter = get_llama_index_splitter(document.metadata.get("file_path", ""))
+            _nodes = splitter.get_nodes_from_documents([document])
+            nodes.extend(_nodes)
+        return nodes