veadk-python 0.2.27__py3-none-any.whl

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/memory/long_term_memory_backends/base_backend.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 abc import ABC, abstractmethod
+
+from pydantic import BaseModel
+
+
+class BaseLongTermMemoryBackend(ABC, BaseModel):
+    index: str
+
+    @abstractmethod
+    def precheck_index_naming(self):
+        """Check the index name is valid or not"""
+
+    @abstractmethod
+    def save_memory(self, user_id: str, event_strings: list[str], **kwargs) -> bool:
+        """Save memory to long term memory backend"""
+
+    @abstractmethod
+    def search_memory(
+        self, user_id: str, 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,67 @@
+# 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 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)
+
+    def precheck_index_naming(self):
+        # no checking
+        pass
+
+    @override
+    def save_memory(self, user_id: str, 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, user_id: str, 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/mem0_backend.py

@@ -0,0 +1,155 @@
+# 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 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,
+)
+from veadk.utils.logger import get_logger
+
+logger = get_logger(__name__)
+
+try:
+    from mem0 import MemoryClient
+
+except ImportError:
+    logger.error(
+        "Failed to import mem0 or dotenv. Please install them with 'pip install mem0 '"
+    )
+    raise ImportError("Required packages not installed: mem0")
+
+
+class Mem0LTMBackend(BaseLongTermMemoryBackend):
+    """Mem0 long term memory backend implementation"""
+
+    mem0_config: Mem0Config = Field(default_factory=Mem0Config)
+
+    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
+                api_key=self.mem0_config.api_key,  # mem0 API key
+            )
+            logger.info(
+                f"Initialized Mem0 client for host: {self.mem0_config.base_url}"
+            )
+            logger.info(f"Initialized Mem0 client for index: {self.index}")
+        except Exception as e:
+            logger.error(
+                f"Failed to initialize Mem0 client for host {self.mem0_config.base_url} : {str(e)}"
+            )
+            raise
+
+    def precheck_index_naming(self):
+        """Check if the index name is valid
+        For Mem0, there are no specific naming constraints
+        """
+        pass
+
+    @override
+    def save_memory(
+        self, event_strings: list[str], user_id: str = "default_user", **kwargs
+    ) -> bool:
+        """Save memory to Mem0
+
+        Args:
+            event_strings: List of event strings to save
+            **kwargs: Additional parameters, including 'user_id' for Mem0
+
+        Returns:
+            bool: True if saved successfully, False otherwise
+        """
+        try:
+            logger.info(
+                f"Saving {len(event_strings)} events to Mem0 for user: {user_id}"
+            )
+
+            for event_string in event_strings:
+                # Save event string to Mem0
+                result = self._mem0_client.add(
+                    [{"role": "user", "content": event_string}],
+                    user_id=user_id,
+                    output_format="v1.1",
+                    async_mode=True,
+                )
+                logger.debug(f"Saved memory result: {result}")
+
+            logger.info(f"Successfully saved {len(event_strings)} events to Mem0")
+            return True
+        except Exception as e:
+            logger.error(f"Failed to save memory to Mem0: {str(e)}")
+            return False
+
+    @override
+    def search_memory(
+        self, query: str, top_k: int, user_id: str = "default_user", **kwargs
+    ) -> list[str]:
+        """Search memory from Mem0
+
+        Args:
+            query: Search query
+            top_k: Number of results to return
+            **kwargs: Additional parameters, including 'user_id' for Mem0
+
+        Returns:
+            list[str]: List of memory strings
+        """
+
+        try:
+            logger.info(
+                f"Searching Mem0 for query: {query}, user: {user_id}, top_k: {top_k}"
+            )
+
+            memories = self._mem0_client.search(
+                query, user_id=user_id, output_format="v1.1", top_k=top_k
+            )
+
+            logger.debug(f"return relevant memories: {memories}")
+
+            memory_list = []
+            # 如果 memories 是列表，直接返回
+            if isinstance(memories, list):
+                for mem in memories:
+                    if "memory" in mem:
+                        memory_list.append(mem["memory"])
+                return memory_list
+
+            if memories.get("results", []):
+                for mem in memories["results"]:
+                    if "memory" in mem:
+                        memory_list.append(mem["memory"])
+
+            logger.info(f"Found {len(memory_list)} memories matching query: {query}")
+            return memory_list
+        except Exception as e:
+            logger.error(f"Failed to search memory from Mem0: {str(e)}")
+            return []

veadk/memory/long_term_memory_backends/opensearch_backend.py

@@ -0,0 +1,124 @@
+# 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, 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,
+)
+from veadk.utils.logger import get_logger
+
+try:
+    from llama_index.vector_stores.opensearch import (
+        OpensearchVectorClient,
+        OpensearchVectorStore,
+    )
+except ImportError:
+    raise ImportError(
+        "Please install VeADK extensions\npip install veadk-python[extensions]"
+    )
+
+logger = get_logger(__name__)
+
+
+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 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,
+        )
+
+    def precheck_index_naming(self, index: str):
+        if not (
+            isinstance(index, str)
+            and not index.startswith(("_", "-"))
+            and index.islower()
+            and re.match(r"^[a-z0-9_\-.]+$", index)
+        ):
+            raise ValueError(
+                f"The index name {index} does not conform to the naming rules of OpenSearch"
+            )
+
+    def _create_vector_index(self, index: str) -> VectorStoreIndex:
+        logger.info(f"Create OpenSearch vector index with index={index}")
+
+        self.precheck_index_naming(index)
+
+        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=index,
+        )
+        vector_store = OpensearchVectorStore(client=opensearch_client)
+        return VectorStoreIndex.from_vector_store(
+            vector_store=vector_store, embed_model=self._embed_model
+        )
+
+    @override
+    def save_memory(self, user_id: str, event_strings: list[str], **kwargs) -> bool:
+        index = f"{self.index}_{user_id}"
+        vector_index = self._create_vector_index(index)
+
+        for event_string in event_strings:
+            document = Document(text=event_string)
+            nodes = self._split_documents([document])
+            vector_index.insert_nodes(nodes)
+        return True
+
+    @override
+    def search_memory(
+        self, user_id: str, query: str, top_k: int, **kwargs
+    ) -> list[str]:
+        index = f"{self.index}_{user_id}"
+
+        vector_index = self._create_vector_index(index)
+
+        _retriever = 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,140 @@
+# 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
+
+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,
+)
+from veadk.utils.logger import get_logger
+
+try:
+    from llama_index.vector_stores.redis import RedisVectorStore
+    from redis import Redis
+    from redisvl.schema import IndexSchema
+
+except ImportError:
+    raise ImportError(
+        "Please install VeADK extensions\npip install veadk-python[extensions]"
+    )
+
+
+logger = get_logger(__name__)
+
+
+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 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,
+        )
+
+    def precheck_index_naming(self, index: str):
+        # no checking
+        pass
+
+    def _create_vector_index(self, index: str) -> VectorStoreIndex:
+        logger.info(f"Create Redis vector index with index={index}")
+
+        self.precheck_index_naming(index)
+
+        # We will use `from_url` to init Redis client once the
+        # AK/SK -> STS token is ready.
+        # self._redis_client = Redis.from_url(url=...)
+        redis_client = Redis(
+            host=self.redis_config.host,
+            port=self.redis_config.port,
+            db=self.redis_config.db,
+            password=self.redis_config.password,
+        )
+
+        # Create an index for each user
+        # Should be Optimized in the future
+        schema = IndexSchema.from_dict(
+            {
+                "index": {"name": index, "prefix": index, "key_separator": "_"},
+                "fields": [
+                    {"name": "id", "type": "tag", "attrs": {"sortable": False}},
+                    {"name": "doc_id", "type": "tag", "attrs": {"sortable": False}},
+                    {"name": "text", "type": "text", "attrs": {"weight": 1.0}},
+                    {
+                        "name": "vector",
+                        "type": "vector",
+                        "attrs": {
+                            "dims": self.embedding_config.dim,
+                            "algorithm": "flat",
+                            "distance_metric": "cosine",
+                        },
+                    },
+                ],
+            }
+        )
+        vector_store = RedisVectorStore(schema=schema, redis_client=redis_client)
+
+        logger.info(
+            f"Create vector store done, index_name={vector_store.index_name} prefix={vector_store.schema.index.prefix}"
+        )
+
+        return VectorStoreIndex.from_vector_store(
+            vector_store=vector_store, embed_model=self._embed_model
+        )
+
+    @override
+    def save_memory(self, user_id: str, event_strings: list[str], **kwargs) -> bool:
+        index = f"veadk-ltm/{self.index}/{user_id}"
+        vector_index = self._create_vector_index(index)
+
+        for event_string in event_strings:
+            document = Document(text=event_string)
+            nodes = self._split_documents([document])
+            vector_index.insert_nodes(nodes)
+
+        return True
+
+    @override
+    def search_memory(
+        self, user_id: str, query: str, top_k: int, **kwargs
+    ) -> list[str]:
+        index = f"veadk-ltm/{self.index}/{user_id}"
+        vector_index = self._create_vector_index(index)
+
+        _retriever = 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