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

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

veadk/memory/short_term_memory.py

@@ -50,20 +50,87 @@ def wrap_get_session_with_callbacks(obj, callback_fn: Callable):
 
 
 class ShortTermMemory(BaseModel):
+    """Short term memory for agent execution.
+
+    The short term memory represents the context of the agent model. All content in the short term memory will be sent to agent model directly, including the system prompt, historical user prompt, and historical model responses.
+
+    Attributes:
+        backend (Literal["local", "mysql", "sqlite", "postgresql", "database"]):
+            The backend of short term memory:
+            - `local` for in-memory storage
+            - `mysql` for mysql / PostgreSQL storage
+            - `sqlite` for locally sqlite storage
+        backend_configs (dict): Configuration dict for init short term memory backend.
+        db_url (str):
+            Database connection url for init short term memory backend.
+            For example, `sqlite:///./test.db`. Once set, it will override the `backend` parameter.
+        local_database_path (str):
+            Local database path, only used when `backend` is `sqlite`.
+            Default to `/tmp/veadk_local_database.db`.
+        after_load_memory_callback (Callable | None):
+            A callback to be called after loading memory from the backend. The callback function should accept `Session` as an input.
+
+    Examples:
+        ### In-memory simple memory
+
+        You can initialize a short term memory with in-memory storage:
+
+        ```python
+        from veadk import Agent, Runner
+        from veadk.memory.short_term_memory import ShortTermMemory
+        import asyncio
+
+        session_id = "veadk_playground_session"
+
+        agent = Agent()
+        short_term_memory = ShortTermMemory(backend="local")
+
+        runner = Runner(
+            agent=agent, short_term_memory=short_term_memory)
+
+        # This invocation will be stored in short-term memory
+        response = asyncio.run(runner.run(
+            messages="My name is VeADK", session_id=session_id
+        ))
+        print(response)
+
+        # The history invocation can be fetched by model
+        response = asyncio.run(runner.run(
+            messages="Do you remember my name?", session_id=session_id # keep the same `session_id`
+        ))
+        print(response)
+        ```
+
+        ### Memory with a Database URL
+
+        Also you can use a databasae connection URL to initialize a short-term memory:
+
+        ```python
+        from veadk.memory.short_term_memory import ShortTermMemory
+
+        short_term_memory = ShortTermMemory(db_url="...")
+        ```
+
+        ### Memory with SQLite
+
+        Once you want to start the short term memory with a local SQLite, you can specify the backend to `sqlite`. It will create a local database in `local_database_path`:
+
+        ```python
+        from veadk.memory.short_term_memory import ShortTermMemory
+
+        short_term_memory = ShortTermMemory(backend="sqlite", local_database_path="")
+        ```
+    """
+
     backend: Literal["local", "mysql", "sqlite", "postgresql", "database"] = "local"
-    """Short term memory backend. `Local` for in-memory storage, `mysql` for mysql / PostgreSQL storage. `sqlite` for sqlite storage."""
 
     backend_configs: dict = Field(default_factory=dict)
-    """Backend specific configurations."""
 
     db_url: str = ""
-    """Database connection URL, e.g. `sqlite:///./test.db`. Once set, it will override the `backend` parameter."""
 
     local_database_path: str = "/tmp/veadk_local_database.db"
-    """Local database path, only used when `backend` is `sqlite`. Default to `/tmp/veadk_local_database.db`."""
 
     after_load_memory_callback: Callable | None = None
-    """A callback to be called after loading memory from the backend. The callback function should accept `Session` as an input."""
 
     _session_service: BaseSessionService = PrivateAttr()
 
@@ -108,6 +175,53 @@ class ShortTermMemory(BaseModel):
         user_id: str,
         session_id: str,
     ) -> Session | None:
+        """Create or retrieve a user session.
+
+        Short term memory can attempt to create a new session for a given application and user. If a session with the same `session_id` already exists, it will be returned instead of creating a new one.
+
+        If the underlying session service is backed by a database (`DatabaseSessionService`), the method first lists all existing sessions for the given `app_name` and `user_id` and logs the number of sessions found. It then checks whether a session with the specified `session_id` already exists:
+        - If it exists → returns the existing session.
+        - If it does not exist → creates and returns a new session.
+
+        Args:
+            app_name (str): The name of the application associated with the session.
+            user_id (str): The unique identifier of the user.
+            session_id (str): The unique identifier of the session to be created or retrieved.
+
+        Returns:
+            Session | None: The retrieved or newly created `Session` object, or `None` if the session creation failed.
+
+        Examples:
+            Create a new session manually:
+
+            ```python
+            import asyncio
+
+            from veadk.memory import ShortTermMemory
+
+            app_name = "app_name"
+            user_id = "user_id"
+            session_id = "session_id"
+
+            short_term_memory = ShortTermMemory()
+
+            session = asyncio.run(
+                short_term_memory.create_session(
+                    app_name=app_name, user_id=user_id, session_id=session_id
+                )
+            )
+
+            print(session)
+
+            session = asyncio.run(
+                short_term_memory.session_service.get_session(
+                    app_name=app_name, user_id=user_id, session_id=session_id
+                )
+            )
+
+            print(session)
+            ```
+        """
         if isinstance(self._session_service, DatabaseSessionService):
             list_sessions_response = await self._session_service.list_sessions(
                 app_name=app_name, user_id=user_id