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

veadk/memory/long_term_memory.py

@@ -13,8 +13,9 @@
 # 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 Literal
+from typing import Any, Literal
 
 from google.adk.events.event import Event
 from google.adk.memory.base_memory_service import (
@@ -24,47 +25,98 @@ 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 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
+
+    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):
-    def __init__(
-        self,
-        backend: Literal[
-            "local", "opensearch", "redis", "mysql", "viking", "viking_mem"
-        ] = "opensearch",
-        top_k: int = 5,
-    ):
-        if backend == "viking":
-            logger.warning(
-                "`viking` backend is deprecated, switching to `viking_mem` backend."
-            )
-            backend = "viking_mem"
-        self.top_k = top_k
-        self.backend = backend
+class LongTermMemory(BaseMemoryService, BaseModel):
+    backend: Union[
+        Literal["local", "opensearch", "redis", "viking", "viking_mem"],
+        BaseLongTermMemoryBackend,
+    ] = "opensearch"
+    """Long term memory backend type"""
 
-        logger.info(
-            f"Initializing long term memory: backend={self.backend} top_k={self.top_k}"
-        )
+    backend_config: dict = Field(default_factory=dict)
+    """Long term memory backend configuration"""
 
-        self.db_client = DatabaseFactory.create(
-            backend=backend,
-        )
-        self.adapter = get_long_term_memory_database_adapter(self.db_client)
+    top_k: int = 5
+    """Number of top similar documents to retrieve during search."""
 
-        logger.info(
-            f"Initialized long term memory: db_client={self.db_client.__class__.__name__} adapter={self.adapter}"
-        )
+    app_name: str = ""
+
+    user_id: str = ""
+
+    def model_post_init(self, __context: Any) -> None:
+        self._backend = None
+
+        # 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
+
+        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 = []
@@ -92,40 +144,57 @@ class LongTermMemory(BaseMemoryService):
         self,
         session: Session,
     ):
+        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)(
+                index=self._index, **self.backend_config if self.backend_config else {}
+            )
+            logger.info(
+                f"Initialize long term memory backend now, index is {self._index}"
+            )
+
         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:
@@ -153,6 +222,6 @@ class LongTermMemory(BaseMemoryService):
             )
 
         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/memory/long_term_memory_backends/__init__.py

@@ -0,0 +1,13 @@
+# 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.

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

veadk/memory/long_term_memory_backends/opensearch_backend.py

@@ -0,0 +1,120 @@
+# 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.
+
+import re
+
+from llama_index.core import (
+    Document,
+    StorageContext,
+    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
+
+import veadk.config  # noqa E401
+from veadk.configs.database_configs import OpensearchConfig
+from veadk.configs.model_configs import EmbeddingModelConfig, NormalEmbeddingModelConfig
+from veadk.knowledgebase.backends.utils import get_llama_index_splitter
+from veadk.memory.long_term_memory_backends.base_backend import (
+    BaseLongTermMemoryBackend,
+)
+
+try:
+    from llama_index.vector_stores.opensearch import (
+        OpensearchVectorClient,
+        OpensearchVectorStore,
+    )
+except ImportError:
+    raise ImportError(
+        "Please install VeADK extensions\npip install veadk-python[extensions]"
+    )
+
+
+class OpensearchLTMBackend(BaseLongTermMemoryBackend):
+    opensearch_config: OpensearchConfig = Field(default_factory=OpensearchConfig)
+    """Opensearch client configs"""
+
+    embedding_config: EmbeddingModelConfig | NormalEmbeddingModelConfig = Field(
+        default_factory=EmbeddingModelConfig
+    )
+    """Embedding model configs"""
+
+    def precheck_index_naming(self):
+        if not (
+            isinstance(self.index, str)
+            and not self.index.startswith(("_", "-"))
+            and self.index.islower()
+            and re.match(r"^[a-z0-9_\-.]+$", self.index)
+        ):
+            raise ValueError(
+                "The index name does not conform to the naming rules of OpenSearch"
+            )
+
+    def model_post_init(self, __context: Any) -> None:
+        self._opensearch_client = OpensearchVectorClient(
+            endpoint=self.opensearch_config.host,
+            port=self.opensearch_config.port,
+            http_auth=(
+                self.opensearch_config.username,
+                self.opensearch_config.password,
+            ),
+            use_ssl=True,
+            verify_certs=False,
+            dim=self.embedding_config.dim,
+            index=self.index,  # collection name
+        )
+
+        self._vector_store = OpensearchVectorStore(client=self._opensearch_client)
+
+        self._storage_context = StorageContext.from_defaults(
+            vector_store=self._vector_store
+        )
+
+        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.from_documents(
+            documents=[],
+            storage_context=self._storage_context,
+            embed_model=self._embed_model,
+        )
+        self._retriever = self._vector_index.as_retriever()
+
+    @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

veadk/memory/long_term_memory_backends/redis_backend.py

@@ -0,0 +1,127 @@
+# 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,
+    StorageContext,
+    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
+
+import veadk.config  # noqa E401
+from veadk.configs.database_configs import RedisConfig
+from veadk.configs.model_configs import EmbeddingModelConfig, NormalEmbeddingModelConfig
+from veadk.knowledgebase.backends.utils import get_llama_index_splitter
+from veadk.memory.long_term_memory_backends.base_backend import (
+    BaseLongTermMemoryBackend,
+)
+
+try:
+    from llama_index.vector_stores.redis import RedisVectorStore
+    from llama_index.vector_stores.redis.schema import (
+        RedisIndexInfo,
+        RedisVectorStoreSchema,
+    )
+    from redis import Redis
+    from redisvl.schema.fields import BaseVectorFieldAttributes
+except ImportError:
+    raise ImportError(
+        "Please install VeADK extensions\npip install veadk-python[extensions]"
+    )
+
+
+class RedisLTMBackend(BaseLongTermMemoryBackend):
+    redis_config: RedisConfig = Field(default_factory=RedisConfig)
+    """Redis client configs"""
+
+    embedding_config: EmbeddingModelConfig | NormalEmbeddingModelConfig = Field(
+        default_factory=EmbeddingModelConfig
+    )
+    """Embedding model configs"""
+
+    def precheck_index_naming(self):
+        # no checking
+        pass
+
+    def model_post_init(self, __context: Any) -> None:
+        # We will use `from_url` to init Redis client once the
+        # AK/SK -> STS token is ready.
+        # self._redis_client = Redis.from_url(url=...)
+
+        self._redis_client = Redis(
+            host=self.redis_config.host,
+            port=self.redis_config.port,
+            db=self.redis_config.db,
+            password=self.redis_config.password,
+        )
+
+        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._schema = RedisVectorStoreSchema(
+            index=RedisIndexInfo(name=self.index),
+        )
+        if "vector" in self._schema.fields:
+            vector_field = self._schema.fields["vector"]
+            if (
+                vector_field
+                and vector_field.attrs
+                and isinstance(vector_field.attrs, BaseVectorFieldAttributes)
+            ):
+                vector_field.attrs.dims = self.embedding_config.dim
+        self._vector_store = RedisVectorStore(
+            schema=self._schema,
+            redis_client=self._redis_client,
+            overwrite=True,
+            collection_name=self.index,
+        )
+
+        self._storage_context = StorageContext.from_defaults(
+            vector_store=self._vector_store
+        )
+
+        self._vector_index = VectorStoreIndex.from_documents(
+            documents=[],
+            storage_context=self._storage_context,
+            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