veadk-python 0.2.7__py3-none-any.whl → 0.2.9__py3-none-any.whl

veadk/memory/long_term_memory_backends/mem0_backend.py

@@ -0,0 +1,129 @@
+# 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 typing_extensions import override
+from pydantic import Field
+
+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"""
+
+        try:
+            self._mem0_client = MemoryClient(
+                # base_url=self.mem0_config.base_url,        # mem0 endpoint
+                api_key=self.mem0_config.api_key,  # mem0 API key
+            )
+            logger.info(f"Initialized Mem0 client for index: {self.index}")
+        except Exception as e:
+            logger.error(f"Failed to initialize Mem0 client: {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], **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
+        """
+        user_id = kwargs.get("user_id", "default_user")
+
+        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",
+                )
+                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, **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
+        """
+        user_id = kwargs.get("user_id", "default_user")
+
+        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
+            )
+
+            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,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

veadk/memory/long_term_memory_backends/vikingdb_memory_backend.py

@@ -0,0 +1,148 @@
+# 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 json
+import re
+import time
+import uuid
+from typing import Any
+
+from pydantic import Field
+from typing_extensions import override
+
+import veadk.config  # noqa E401
+from veadk.config import getenv
+from veadk.memory.long_term_memory_backends.base_backend import (
+    BaseLongTermMemoryBackend,
+)
+from veadk.utils.logger import get_logger
+
+try:
+    from mcp_server_vikingdb_memory.common.memory_client import VikingDBMemoryService
+except ImportError:
+    raise ImportError(
+        "Please install VeADK extensions\npip install veadk-python[extensions]"
+    )
+
+logger = get_logger(__name__)
+
+
+class VikingDBLTMBackend(BaseLongTermMemoryBackend):
+    volcengine_access_key: str = Field(
+        default_factory=lambda: getenv("VOLCENGINE_ACCESS_KEY")
+    )
+
+    volcengine_secret_key: str = Field(
+        default_factory=lambda: getenv("VOLCENGINE_SECRET_KEY")
+    )
+
+    region: str = "cn-beijing"
+    """VikingDB memory region"""
+
+    def precheck_index_naming(self):
+        if not (
+            isinstance(self.index, str)
+            and 1 <= len(self.index) <= 128
+            and re.fullmatch(r"^[a-zA-Z][a-zA-Z0-9_]*$", self.index)
+        ):
+            raise ValueError(
+                "The index name does not conform to the rules: it must start with an English letter, contain only letters, numbers, and underscores, and have a length of 1-128."
+            )
+
+    def model_post_init(self, __context: Any) -> None:
+        self._client = VikingDBMemoryService(
+            ak=self.volcengine_access_key,
+            sk=self.volcengine_secret_key,
+            region=self.region,
+        )
+
+        # check whether collection exist, if not, create it
+        if not self._collection_exist():
+            self._create_collection()
+
+    def _collection_exist(self) -> bool:
+        try:
+            self._client.get_collection(collection_name=self.index)
+            return True
+        except Exception:
+            return False
+
+    def _create_collection(self) -> None:
+        response = self._client.create_collection(
+            collection_name=self.index,
+            description="Created by Volcengine Agent Development Kit VeADK",
+            builtin_event_types=["sys_event_v1"],
+        )
+        return response
+
+    @override
+    def save_memory(self, event_strings: list[str], **kwargs) -> bool:
+        user_id = kwargs.get("user_id")
+        if user_id is None:
+            raise ValueError("user_id is required")
+        session_id = str(uuid.uuid1())
+        messages = []
+        for raw_events in event_strings:
+            event = json.loads(raw_events)
+            content = event["parts"][0]["text"]
+            role = (
+                "user" if event["role"] == "user" else "assistant"
+            )  # field 'role': viking memory only allow 'assistant','system','user',
+            messages.append({"role": role, "content": content})
+        metadata = {
+            "default_user_id": user_id,
+            "default_assistant_id": "assistant",
+            "time": int(time.time() * 1000),
+        }
+        response = self._client.add_messages(
+            collection_name=self.index,
+            messages=messages,
+            metadata=metadata,
+            session_id=session_id,
+        )
+
+        if not response.get("code") == 0:
+            raise ValueError(f"Save VikingDB memory error: {response}")
+
+        return True
+
+    @override
+    def search_memory(self, query: str, top_k: int, **kwargs) -> list[str]:
+        user_id = kwargs.get("user_id")
+        if user_id is None:
+            raise ValueError("user_id is required")
+        filter = {
+            "user_id": user_id,
+            "memory_type": ["sys_event_v1"],
+        }
+        response = self._client.search_memory(
+            collection_name=self.index, query=query, filter=filter, limit=top_k
+        )
+
+        if not response.get("code") == 0:
+            raise ValueError(f"Search VikingDB memory error: {response}")
+
+        result = response.get("data", {}).get("result_list", [])
+        if result:
+            return [
+                json.dumps(
+                    {
+                        "role": "user",
+                        "parts": [{"text": r.get("memory_info").get("summary")}],
+                    },
+                    ensure_ascii=False,
+                )
+                for r in result
+            ]
+        return []