veadk-python 0.2.16__py3-none-any.whl → 0.2.17__py3-none-any.whl

veadk/knowledgebase/backends/base_backend.py

@@ -18,8 +18,21 @@ from pydantic import BaseModel
 
 
 class BaseKnowledgebaseBackend(ABC, BaseModel):
+    """Base backend for knowledgebase.
+
+    Attributes:
+        index (str): Index or collection name of the vector storage.
+
+    Examples:
+        You can implement your own knowledgebase backend.
+
+        ```python
+        class CustomKnowledgebaseBackend(BaseKnowledgebaseBackend):
+            pass
+        ```
+    """
+
     index: str
-    """Index or collection name of the vector storage."""
 
     @abstractmethod
     def precheck_index_naming(self) -> None:

veadk/knowledgebase/backends/in_memory_backend.py

@@ -24,10 +24,19 @@ from veadk.knowledgebase.backends.utils import get_llama_index_splitter
 
 
 class InMemoryKnowledgeBackend(BaseKnowledgebaseBackend):
+    """A in-memory implementation backend for knowledgebase.
+
+    In-memory backend stores embedded text in a vector storage from Llama-index.
+
+    Attributes:
+        embedding_config (EmbeddingModelConfig):
+            Embedding config for text embedding and search.
+            Embedding config contains embedding model name and the corresponding dim.
+    """
+
     embedding_config: NormalEmbeddingModelConfig | EmbeddingModelConfig = Field(
         default_factory=EmbeddingModelConfig
     )
-    """Embedding model configs"""
 
     def model_post_init(self, __context: Any) -> None:
         self._embed_model = OpenAILikeEmbedding(

veadk/knowledgebase/backends/opensearch_backend.py

@@ -43,6 +43,32 @@ except ImportError:
 
 
 class OpensearchKnowledgeBackend(BaseKnowledgebaseBackend):
+    """Opensearch-based backend for knowledgebase.
+
+    Opensearch backend stores embedded text in a opensearch database by Llama-index.
+
+    Attributes:
+        opensearch_config (OpensearchConfig):
+            Opensearch database configurations.
+            Mainly contains opensearch host, port, username, password, etc.
+        embedding_config (EmbeddingModelConfig):
+            Embedding config for text embedding and search.
+            Embedding config contains embedding model name and the corresponding dim.
+
+    Examples:
+        Init a knowledgebase based on opensearch backend.
+
+        ```python
+        knowledgebase = Knowledgebase(backend="opensearch")
+        ```
+
+        With more configurations:
+
+        ```python
+        ...
+        ```
+    """
+
     opensearch_config: OpensearchConfig = Field(default_factory=OpensearchConfig)
     """Opensearch client configs"""
 

veadk/knowledgebase/backends/redis_backend.py

@@ -44,13 +44,40 @@ except ImportError:
 
 
 class RedisKnowledgeBackend(BaseKnowledgebaseBackend):
+    """Redis based backend for knowledgebase.
+
+    Redis backend stores embedded text in a redis database by Llama-index.
+
+    Attributes:
+        redis_config (RedisConfig):
+            Redis database configurations.
+            Mainly contains redis database host, port, etc.
+        embedding_config (EmbeddingModelConfig):
+            Embedding configurations for text embedding and search.
+            Embedding config contains embedding model name and the corresponding dim.
+
+    Notes:
+        Please ensure that your redis database supports Redisaearch stack.
+
+    Examples:
+        Init a knowledgebase based on redis backend.
+
+        ```python
+        knowledgebase = Knowledgebase(backend="redis")
+        ```
+
+        With more configurations:
+
+        ```python
+        ...
+        ```
+    """
+
     redis_config: RedisConfig = Field(default_factory=RedisConfig)
-    """Redis client configs"""
 
     embedding_config: EmbeddingModelConfig | NormalEmbeddingModelConfig = Field(
         default_factory=EmbeddingModelConfig
     )
-    """Embedding model configs"""
 
     def model_post_init(self, __context: Any) -> None:
         # We will use `from_url` to init Redis client once the

veadk/knowledgebase/backends/vikingdb_knowledge_backend.py

@@ -59,24 +59,62 @@ def get_files_in_directory(directory: str):
 
 
 class VikingDBKnowledgeBackend(BaseKnowledgebaseBackend):
+    """Volcengine Viking DB knowledgebase backend.
+
+    Volcegnine Viking DB knowledgebase provides powerful knowledgebase storage and search.
+
+    Attributes:
+        volcengine_access_key (str | None):
+            Access key for Volcengine. Loaded automatically from the
+            `VOLCENGINE_ACCESS_KEY` environment variable if not provided.
+
+        volcengine_secret_key (str | None):
+            Secret key for Volcengine. Loaded automatically from the
+            `VOLCENGINE_SECRET_KEY` environment variable if not provided.
+
+        session_token (str):
+            Optional session token for temporary credentials. Defaults to an empty string.
+
+        volcengine_project (str):
+            VikingDB knowledgebase project name in the Volcengine console platform.
+            Defaults to `"default"`.
+
+        region (str): Region of the VikingDB knowledgebase. Defaults to `"cn-beijing"`.
+
+        tos_config (TOSConfig | NormalTOSConfig):
+            TOS configuration used for uploading files to TOS (Volcengine’s Object Storage).
+            Defaults to a new instance of `TOSConfig`.
+
+    Notes:
+        Please make sure that you have created a bucket in your TOS.
+
+    Examples:
+        Init a knowledgebase based on VikingDB knowledgebase backend.
+
+        ```python
+        knowledgebase = Knowledgebase(backend="redis")
+        ```
+
+        With more configurations:
+
+        ```python
+        ...
+        ```
+    """
+
     volcengine_access_key: str | None = Field(
         default_factory=lambda: os.getenv("VOLCENGINE_ACCESS_KEY")
     )
-
     volcengine_secret_key: str | None = Field(
         default_factory=lambda: os.getenv("VOLCENGINE_SECRET_KEY")
     )
-
     session_token: str = ""
 
     volcengine_project: str = "default"
-    """VikingDB knowledgebase project in Volcengine console platform. Default by `default`"""
 
     region: str = "cn-beijing"
-    """VikingDB knowledgebase region"""
 
     tos_config: TOSConfig | NormalTOSConfig = Field(default_factory=TOSConfig)
-    """TOS config, used to upload files to TOS"""
 
     def model_post_init(self, __context: Any) -> None:
         self.precheck_index_naming()

veadk/knowledgebase/knowledgebase.py

@@ -56,6 +56,110 @@ def _get_backend_cls(backend: str) -> type[BaseKnowledgebaseBackend]:
 
 
 class KnowledgeBase(BaseModel):
+    """A knowledge base for storing user-related information.
+
+    This class represents a knowledge base used to store and retrieve user-specific data.
+    It supports multiple backend options, including in-memory, OpenSearch, Redis, and Volcengine's
+    VikingDB. The knowledge base allows for efficient document retrieval based on similarity,
+    with the ability to configure backend-specific settings.
+
+    Attributes:
+        name (str): The name of the knowledge base. Default is "user_knowledgebase".
+        description (str): A description of the knowledge base. Default is "This knowledgebase stores some user-related information."
+        backend (Union[Literal["local", "opensearch", "viking", "redis"], BaseKnowledgebaseBackend]):
+            The type of backend to use for storing and querying the knowledge base. Supported options include:
+            - 'local' for in-memory storage (data is lost when the program exits).
+            - 'opensearch' for OpenSearch (requires OpenSearch cluster).
+            - 'viking' for Volcengine VikingDB (requires VikingDB service).
+            - 'redis' for Redis with vector search capability (requires Redis).
+            Default is 'local'.
+        backend_config (dict): Configuration dictionary for the selected backend.
+        top_k (int): The number of top similar documents to retrieve during a search. Default is 10.
+        app_name (str): The name of the application associated with the knowledge base. If index is not provided, this value will be set to `index`.
+        index (str): The name of the knowledge base index.
+
+    Notes:
+        Please ensure that you have set the embedding-related configurations in environment variables.
+
+    Examples:
+        ### Simple backend
+
+        Create a local knowledgebase:
+
+        ```python
+        from veadk import Agent, Runner
+        from veadk.knowledgebase.knowledgebase import KnowledgeBase
+        from veadk.memory.short_term_memory import ShortTermMemory
+
+        app_name = "veadk_playground_app"
+        user_id = "veadk_playground_user"
+        session_id = "veadk_playground_session"
+
+
+        knowledgebase = KnowledgeBase(backend="opensearch", app_name=app_name)
+        knowledgebase.add_from_files(files=[knowledgebase_file])
+
+        agent = Agent(knowledgebase=knowledgebase)
+
+        runner = Runner(
+            agent=agent,
+            short_term_memory=ShortTermMemory(),
+            app_name=app_name,
+            user_id=user_id,
+        )
+
+        response = await runner.run(
+            messages="Tell me the secret of green.", session_id=session_id
+        )
+        print(response)
+        ```
+
+        ### Initialize knowledgebase with metadata
+
+        ```python
+        from veadk.knowledgebase import KnowledgeBase
+
+        knowledgebase = KnowledgeBase(
+            name="user_data",
+            description="A knowledgebase contains user hobbies.",
+            index="my_app",
+        )
+        ```
+
+        ### Initialize knowledgebase with backend instance
+
+        ```python
+        import veadk.config  # noqa
+
+        from veadk.knowledgebase import KnowledgeBase
+        from veadk.knowledgebase.backends.in_memory_backend import InMemoryKnowledgeBackend
+
+        backend = InMemoryKnowledgeBackend(
+            index="my_app",
+            embedding_config=...,
+        )
+
+        knowledgebase = KnowledgeBase(
+            name="user_data",
+            description="A knowledgebase contains user hobbies.",
+            backend=backend,
+        )
+        ```
+
+        ### Initialize knowledgebase with backend config
+
+        ```python
+        from veadk.knowledgebase import KnowledgeBase
+
+        knowledgebase = KnowledgeBase(
+            name="user_data",
+            description="A knowledgebase contains user hobbies.",
+            backend="local",
+            backend_config={"index": "user_app"},
+        )
+        ```
+    """
+
     name: str = "user_knowledgebase"
 
     description: str = "This knowledgebase stores some user-related information."
@@ -63,23 +167,14 @@ class KnowledgeBase(BaseModel):
     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
-    """Number of top similar documents to retrieve during search"""
 
     app_name: str = ""
 
     index: str = ""
-    """The name of the knowledgebase index. If not provided, it will be generated based on the `app_name`."""
 
     def model_post_init(self, __context: Any) -> None:
         if isinstance(self.backend, BaseKnowledgebaseBackend):
@@ -108,15 +203,81 @@ class KnowledgeBase(BaseModel):
         )
 
     def add_from_directory(self, directory: str, **kwargs) -> bool:
-        """Add knowledge from file path to knowledgebase"""
+        """Add knowledge from file path to knowledgebase.
+
+        Add the files in the directory to knowledgebase backend.
+
+        Args:
+            directory (str): The directory path that needs to store.
+
+        Returns:
+            bool: True if successfully store the knowledgebase, False otherwise.
+
+        Examples:
+            Store a directory to knowledgebase:
+
+            ```python
+            knowledgebase = Knowledgebase(backend="local")
+
+            if knowledgebase.add_from_directory("./knowledgebase"):
+                # add successfully
+                ...
+            else:
+                raise RuntimeError("Uploaded directory failed.")
+            ```
+        """
         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"""
+        """Add knowledge files to knowledgebase.
+
+        Add a list of files to knowledgebase backend.
+
+        Args:
+            files (str): The list of files.
+
+        Returns:
+            bool: True if successfully store the knowledgebase, False otherwise.
+
+        Examples:
+            Store files to knowledgebase:
+
+            ```python
+            knowledgebase = Knowledgebase(backend="local")
+
+            if knowledgebase.add_from_files("./knowledgebase"):
+                # add successfully
+                ...
+            else:
+                raise RuntimeError("Uploaded files failed.")
+            ```
+        """
         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"""
+        """Add a piece of text or a list of text to knowledgebase.
+
+        The `text` can be a string or a list of string. The text will be embedded and stored by the corresponding backend.
+
+        Args:
+            text (str | list[str]): The text string or a list of text strings.
+
+        Returns:
+            bool: True if successfully store the knowledgebase, False otherwise.
+
+        Examples:
+            Store a string or a list of string to knowledgebase:
+
+            ```python
+            knowledgebase = Knowledgebase(backend="local")
+
+            if knowledgebase.add_from_text("./knowledgebase"):
+                # add successfully
+                ...
+            else:
+                raise RuntimeError("Uploaded text failed.")
+            ```
+        """
         return self._backend.add_from_text(text=text, **kwargs)
 
     def search(self, query: str, top_k: int = 0, **kwargs) -> list[KnowledgebaseEntry]:

veadk/memory/long_term_memory.py

@@ -73,24 +73,111 @@ def _get_backend_cls(backend: str) -> type[BaseLongTermMemoryBackend]:
 
 
 class LongTermMemory(BaseMemoryService, BaseModel):
+    """Manages long-term memory storage and retrieval for applications.
+
+    This class provides an interface to store, retrieve, and manage long-term
+    contextual information using different backend types (e.g., OpenSearch, Redis).
+    It supports configuration of the backend service and retrieval behavior.
+
+    Attributes:
+        backend (Union[Literal["local", "opensearch", "redis", "viking", "viking_mem", "mem0"], BaseLongTermMemoryBackend]):
+            The type or instance of the long-term memory backend. Defaults to "opensearch".
+
+        backend_config (dict):
+            Configuration parameters for the selected backend. Defaults to an empty dictionary.
+
+        top_k (int):
+            The number of top similar documents to retrieve during search. Defaults to 5.
+
+        index (str):
+            The name of the index or collection used for storing memory items. Defaults to an empty string.
+
+        app_name (str):
+            The name of the application that owns this memory instance. Defaults to an empty string.
+
+        user_id (str):
+            Deprecated attribute. Retained for backward compatibility. Defaults to an empty string.
+
+    Notes:
+        Please ensure that you have set the embedding-related configurations in environment variables.
+
+    Examples:
+        ### Simple long-term memory
+
+        Once create a long-term memory withou any arguments, all configurations are come from **environment variables**.
+
+        ```python
+        import asyncio
+
+        from veadk import Agent, Runner
+        from veadk.memory.long_term_memory import LongTermMemory
+        from veadk.memory.short_term_memory import ShortTermMemory
+
+        app_name = "veadk_playground_app"
+        user_id = "veadk_playground_user"
+
+        long_term_memory = LongTermMemory(backend="local", app_name=app_name)
+
+        agent = Agent(long_term_memory=long_term_memory)
+
+        runner = Runner(
+            agent=agent,
+            app_name=app_name,
+            user_id=user_id,
+            short_term_memory=ShortTermMemory(),
+        )
+
+        # ===== add memory =====
+        session_id = "veadk_playground_session"
+        teaching_prompt = "I brought an ice-cream last week."
+
+        asyncio.run(runner.run(messages=teaching_prompt, session_id=session_id))
+        asyncio.run(
+            runner.save_session_to_long_term_memory(session_id=session_id)
+        )  # save session to long-term memory
+
+
+        # ===== check memory =====
+        session_id = "veadk_playground_session_2"  # use a new session
+        student_prompt = "What I brought last week?"
+
+        response = asyncio.run(runner.run(messages=student_prompt, session_id=session_id))
+
+        print(response)
+        ```
+
+        ### Create with a backend instance
+
+        ```python
+        from veadk.memory.long_term_memory import LongTermMemory
+        from veadk.memory.long_term_memory.backends import LongTermMemory
+
+        long_term_memory = LongTermMemory(backend=...)
+        ```
+
+        ### Create with backend configurations
+
+        ```python
+        from veadk.memory.long_term_memory import LongTermMemory
+
+        long_term_memory = LongTermMemory(backend="", backend_config={})
+        ```
+    """
+
     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."""
 
     index: str = ""
 
     app_name: str = ""
 
     user_id: str = ""
-    """Deprecated attribute"""
 
     def model_post_init(self, __context: Any) -> None:
         # Once user define a backend instance, use it directly
@@ -163,6 +250,33 @@ class LongTermMemory(BaseMemoryService, BaseModel):
         self,
         session: Session,
     ):
+        """Add a chat session's events to the long-term memory backend.
+
+        This method extracts and filters the events from a given `Session` object,
+        converts them into serialized strings, and stores them into the long-term
+        memory system. It is typically called after a chat session ends or when
+        important contextual data needs to be persisted for future retrieval.
+
+        Args:
+            session (Session):
+                The session object containing user ID and a list of events to persist.
+
+        Examples:
+            ```python
+            session = Session(
+                user_id="user_123",
+                events=[
+                    Event(role="user", content="I like Go and Rust."),
+                    Event(role="assistant", content="Got it! I'll remember that."),
+                ]
+            )
+
+            await memory_service.add_session_to_memory(session)
+            # Logs:
+            # Adding 2 events to long term memory: index=main
+            # Added 2 events to long term memory: index=main, user_id=user_123
+            ```
+        """
         user_id = session.user_id
         event_strings = self._filter_and_convert_events(session.events)
 
@@ -178,6 +292,36 @@ class LongTermMemory(BaseMemoryService, BaseModel):
     async def search_memory(
         self, *, app_name: str, user_id: str, query: str
     ) -> SearchMemoryResponse:
+        """Search memory entries for a given user and query.
+
+        This method queries the memory backend to retrieve the most relevant stored
+        memory chunks for a given user and text query. It then converts those raw
+        memory chunks into structured `MemoryEntry` objects to be returned to the caller.
+
+        Args:
+            app_name (str): Name of the application requesting the memory search.
+            user_id (str): Unique identifier for the user whose memory is being queried.
+            query (str): The text query to match against stored memory content.
+
+        Returns:
+            SearchMemoryResponse:
+                An object containing a list of `MemoryEntry` items representing
+                the retrieved memory snippets relevant to the query.
+
+        Examples:
+            ```python
+            response = await memory_service.search_memory(
+                app_name="chat_app",
+                user_id="user_123",
+                query="favorite programming language"
+            )
+
+            for memory in response.memories:
+                print(memory.content.parts[0].text)
+            # Output:
+            # User likes Python and TypeScript for backend development.
+            ```
+        """
         logger.info(f"Search memory with query={query}")
 
         memory_chunks = []

veadk/memory/long_term_memory_backends/mem0_backend.py

@@ -17,6 +17,7 @@ from typing import Any
 from pydantic import Field
 from typing_extensions import override
 
+from veadk.auth.veauth.viking_mem0_veauth import get_viking_mem0_token
 from veadk.configs.database_configs import Mem0Config
 from veadk.memory.long_term_memory_backends.base_backend import (
     BaseLongTermMemoryBackend,
@@ -43,6 +44,16 @@ class Mem0LTMBackend(BaseLongTermMemoryBackend):
     def model_post_init(self, __context: Any) -> None:
         """Initialize Mem0 client"""
 
+        if not self.mem0_config.api_key:
+            if not self.mem0_config.api_key_id and not self.mem0_config.project_id:
+                raise ValueError(
+                    "API Key not set, auto fetching api key needs `api_key_id` or `project_id`"
+                )
+            self.mem0_config.api_key = get_viking_mem0_token(
+                api_key_id=self.mem0_config.api_key_id,
+                memory_project_id=self.mem0_config.project_id,
+            )
+
         try:
             self._mem0_client = MemoryClient(
                 host=self.mem0_config.base_url,  # mem0 endpoint