veadk-python 0.2.27__py3-none-any.whl

veadk/knowledgebase/knowledgebase.py

@@ -0,0 +1,307 @@
+# 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 __future__ import annotations
+
+from typing import Any, Callable, Literal, Union
+
+from pydantic import BaseModel, Field
+
+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}")
+
+
+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."
+
+    backend: Union[
+        Literal["local", "opensearch", "viking", "redis"], BaseKnowledgebaseBackend
+    ] = "local"
+
+    backend_config: dict = Field(default_factory=dict)
+
+    top_k: int = 10
+
+    app_name: str = ""
+
+    index: str = ""
+
+    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
+
+        # Once user define backend config, use it directly
+        if self.backend_config:
+            self._backend = _get_backend_cls(self.backend)(**self.backend_config)
+            return
+
+        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} 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__}"
+        )
+
+    def add_from_directory(self, directory: str, **kwargs) -> bool:
+        """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 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 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]:
+        """Search knowledge from knowledgebase"""
+        top_k = top_k if top_k != 0 else self.top_k
+
+        _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

@@ -0,0 +1,35 @@
+# 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 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 == "LongTermMemory":
+        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

@@ -0,0 +1,365 @@
+# 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.
+
+# adapted from Google ADK memory service adk-python/src/google/adk/memory/vertex_ai_memory_bank_service.py at 0a9e67dbca67789247e882d16b139dbdc76a329a · google/adk-python
+
+import json
+from typing import Any, Literal
+
+from google.adk.events.event import Event
+from google.adk.memory.base_memory_service import (
+    BaseMemoryService,
+    SearchMemoryResponse,
+)
+from google.adk.memory.memory_entry import MemoryEntry
+from google.adk.sessions import Session
+from google.genai import types
+from pydantic import BaseModel, Field
+from typing_extensions import Union, override
+
+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}")
+
+
+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"
+
+    backend_config: dict = Field(default_factory=dict)
+
+    top_k: int = 5
+
+    index: str = ""
+
+    app_name: str = ""
+
+    user_id: str = ""
+
+    def model_post_init(self, __context: Any) -> 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__}, index={self.index}"
+            )
+            return
+
+        # Once user define backend config, use it directly
+        if self.backend_config:
+            if "index" not in self.backend_config:
+                logger.warning(
+                    "Attribute `index` not provided in backend_config, use `index` or `app_name` instead."
+                )
+                self.backend_config["index"] = self.index or self.app_name
+
+            logger.debug(
+                f"Init {self.backend}, Use provided backend config: {self.backend_config}"
+            )
+            self._backend = _get_backend_cls(self.backend)(**self.backend_config)
+            return
+
+        # 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."
+            )
+            self.index = "default_app"
+
+        # Forward compliance
+        if self.backend == "viking_mem":
+            logger.warning(
+                "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 = []
+        for event in events:
+            # filter: bad event
+            if not event.content or not event.content.parts:
+                continue
+
+            # filter: only add user event to memory to enhance retrieve performance
+            if not event.author == "user":
+                continue
+
+            # filter: discard function call and function response
+            if not event.content.parts[0].text:
+                continue
+
+            # convert: to string-format for storage
+            message = event.content.model_dump(exclude_none=True, mode="json")
+
+            final_events.append(json.dumps(message, ensure_ascii=False))
+        return final_events
+
+    @override
+    async def add_session_to_memory(
+        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)
+
+        logger.info(
+            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}"
+        )
+
+    @override
+    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 = []
+        try:
+            memory_chunks = self._backend.search_memory(
+                query=query, top_k=self.top_k, user_id=user_id
+            )
+        except Exception as e:
+            logger.error(
+                f"Exception orrcus during memory search: {e}. Return empty memory chunks"
+            )
+
+        memory_events = []
+        for memory in memory_chunks:
+            try:
+                memory_dict = json.loads(memory)
+                try:
+                    text = memory_dict["parts"][0]["text"]
+                    role = memory_dict["role"]
+                except KeyError as _:
+                    # prevent not a standard text-based event
+                    logger.warning(
+                        f"Memory content: {memory_dict}. Skip return this memory."
+                    )
+                    continue
+            except json.JSONDecodeError:
+                # prevent the memory string is not dumped by `Event` class
+                text = memory
+                role = "user"
+
+            memory_events.append(
+                MemoryEntry(
+                    author="user",
+                    content=types.Content(parts=[types.Part(text=text)], role=role),
+                )
+            )
+
+        logger.info(
+            f"Return {len(memory_events)} memory events for query: {query} index={self.index} user_id={user_id}"
+        )
+        return SearchMemoryResponse(memories=memory_events)