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

veadk/memory/long_term_memory_backends/mem0_backend.py

@@ -13,12 +13,11 @@
 # limitations under the License.
 
 from typing import Any
-from typing_extensions import override
+
 from pydantic import Field
+from typing_extensions import override
 
 from veadk.configs.database_configs import Mem0Config
-
-
 from veadk.memory.long_term_memory_backends.base_backend import (
     BaseLongTermMemoryBackend,
 )
@@ -46,12 +45,17 @@ class Mem0LTMBackend(BaseLongTermMemoryBackend):
 
         try:
             self._mem0_client = MemoryClient(
-                # base_url=self.mem0_config.base_url,        # mem0 endpoint
+                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: {str(e)}")
+            logger.error(
+                f"Failed to initialize Mem0 client for host {self.mem0_config.base_url} : {str(e)}"
+            )
             raise
 
     def precheck_index_naming(self):
@@ -61,7 +65,9 @@ class Mem0LTMBackend(BaseLongTermMemoryBackend):
         pass
 
     @override
-    def save_memory(self, event_strings: list[str], **kwargs) -> bool:
+    def save_memory(
+        self, event_strings: list[str], user_id: str = "default_user", **kwargs
+    ) -> bool:
         """Save memory to Mem0
 
         Args:
@@ -71,8 +77,6 @@ class Mem0LTMBackend(BaseLongTermMemoryBackend):
         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}"
@@ -84,6 +88,7 @@ class Mem0LTMBackend(BaseLongTermMemoryBackend):
                     [{"role": "user", "content": event_string}],
                     user_id=user_id,
                     output_format="v1.1",
+                    async_mode=True,
                 )
                 logger.debug(f"Saved memory result: {result}")
 
@@ -94,7 +99,9 @@ class Mem0LTMBackend(BaseLongTermMemoryBackend):
             return False
 
     @override
-    def search_memory(self, query: str, top_k: int, **kwargs) -> list[str]:
+    def search_memory(
+        self, query: str, top_k: int, user_id: str = "default_user", **kwargs
+    ) -> list[str]:
         """Search memory from Mem0
 
         Args:
@@ -105,7 +112,6 @@ class Mem0LTMBackend(BaseLongTermMemoryBackend):
         Returns:
             list[str]: List of memory strings
         """
-        user_id = kwargs.get("user_id", "default_user")
 
         try:
             logger.info(
@@ -116,7 +122,16 @@ class Mem0LTMBackend(BaseLongTermMemoryBackend):
                 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:

veadk/memory/long_term_memory_backends/opensearch_backend.py

@@ -14,11 +14,7 @@
 
 import re
 
-from llama_index.core import (
-    Document,
-    StorageContext,
-    VectorStoreIndex,
-)
+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
@@ -31,6 +27,7 @@ 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 (
@@ -42,6 +39,8 @@ except ImportError:
         "Please install VeADK extensions\npip install veadk-python[extensions]"
     )
 
+logger = get_logger(__name__)
+
 
 class OpensearchLTMBackend(BaseLongTermMemoryBackend):
     opensearch_config: OpensearchConfig = Field(default_factory=OpensearchConfig)
@@ -52,19 +51,30 @@ class OpensearchLTMBackend(BaseLongTermMemoryBackend):
     )
     """Embedding model configs"""
 
-    def precheck_index_naming(self):
+    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(self.index, str)
-            and not self.index.startswith(("_", "-"))
-            and self.index.islower()
-            and re.match(r"^[a-z0-9_\-.]+$", self.index)
+            isinstance(index, str)
+            and not index.startswith(("_", "-"))
+            and index.islower()
+            and re.match(r"^[a-z0-9_\-.]+$", index)
         ):
             raise ValueError(
-                "The index name does not conform to the naming rules of OpenSearch"
+                f"The index name {index} does not conform to the naming rules of OpenSearch"
             )
 
-    def model_post_init(self, __context: Any) -> None:
-        self._opensearch_client = OpensearchVectorClient(
+    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=(
@@ -74,39 +84,33 @@ class OpensearchLTMBackend(BaseLongTermMemoryBackend):
             use_ssl=True,
             verify_certs=False,
             dim=self.embedding_config.dim,
-            index=self.index,  # collection name
+            index=index,
         )
-
-        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,
+        vector_store = OpensearchVectorStore(client=opensearch_client)
+        return VectorStoreIndex.from_vector_store(
+            vector_store=vector_store, embed_model=self._embed_model
         )
-        self._retriever = self._vector_index.as_retriever()
 
     @override
-    def save_memory(self, event_strings: list[str], **kwargs) -> bool:
+    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])
-            self._vector_index.insert_nodes(nodes)
+            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)
+    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]
 

veadk/memory/long_term_memory_backends/redis_backend.py

@@ -12,11 +12,7 @@
 # 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 import Document, VectorStoreIndex
 from llama_index.core.schema import BaseNode
 from llama_index.embeddings.openai_like import OpenAILikeEmbedding
 from pydantic import Field
@@ -29,21 +25,22 @@ 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 llama_index.vector_stores.redis.schema import (
-        RedisIndexInfo,
-        RedisVectorStoreSchema,
-    )
     from redis import Redis
-    from redisvl.schema.fields import BaseVectorFieldAttributes
+    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"""
@@ -53,67 +50,83 @@ class RedisLTMBackend(BaseLongTermMemoryBackend):
     )
     """Embedding model configs"""
 
-    def precheck_index_naming(self):
+    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 model_post_init(self, __context: Any) -> None:
+    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=...)
-
-        self._redis_client = Redis(
+        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,
+        # 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)
 
-        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
+        logger.info(
+            f"Create vector store done, index_name={vector_store.index_name} prefix={vector_store.schema.index.prefix}"
         )
 
-        self._vector_index = VectorStoreIndex.from_documents(
-            documents=[],
-            storage_context=self._storage_context,
-            embed_model=self._embed_model,
+        return VectorStoreIndex.from_vector_store(
+            vector_store=vector_store, embed_model=self._embed_model
         )
 
     @override
-    def save_memory(self, event_strings: list[str], **kwargs) -> bool:
+    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])
-            self._vector_index.insert_nodes(nodes)
+            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)
+    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]
 

veadk/memory/long_term_memory_backends/vikingdb_memory_backend.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import json
+import os
 import re
 import time
 import uuid
@@ -22,34 +23,37 @@ from pydantic import Field
 from typing_extensions import override
 
 import veadk.config  # noqa E401
-from veadk.config import getenv
+from veadk.auth.veauth.utils import get_credential_from_vefaas_iam
+from veadk.integrations.ve_viking_db_memory.ve_viking_db_memory import (
+    VikingDBMemoryClient,
+)
 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_access_key: str | None = Field(
+        default_factory=lambda: os.getenv("VOLCENGINE_ACCESS_KEY")
     )
 
-    volcengine_secret_key: str = Field(
-        default_factory=lambda: getenv("VOLCENGINE_SECRET_KEY")
+    volcengine_secret_key: str | None = Field(
+        default_factory=lambda: os.getenv("VOLCENGINE_SECRET_KEY")
     )
 
+    session_token: str = ""
+
     region: str = "cn-beijing"
     """VikingDB memory region"""
 
+    def model_post_init(self, __context: Any) -> None:
+        # check whether collection exist, if not, create it
+        if not self._collection_exist():
+            self._create_collection()
+
     def precheck_index_naming(self):
         if not (
             isinstance(self.index, str)
@@ -60,37 +64,39 @@ class VikingDBLTMBackend(BaseLongTermMemoryBackend):
                 "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)
+            client = self._get_client()
+            client.get_collection(collection_name=self.index)
             return True
         except Exception:
             return False
 
     def _create_collection(self) -> None:
-        response = self._client.create_collection(
+        client = self._get_client()
+        response = client.create_collection(
             collection_name=self.index,
             description="Created by Volcengine Agent Development Kit VeADK",
             builtin_event_types=["sys_event_v1"],
         )
         return response
 
+    def _get_client(self) -> VikingDBMemoryClient:
+        if not (self.volcengine_access_key and self.volcengine_secret_key):
+            cred = get_credential_from_vefaas_iam()
+            self.volcengine_access_key = cred.access_key_id
+            self.volcengine_secret_key = cred.secret_access_key
+            self.session_token = cred.session_token
+
+        return VikingDBMemoryClient(
+            ak=self.volcengine_access_key,
+            sk=self.volcengine_secret_key,
+            sts_token=self.session_token,
+            region=self.region,
+        )
+
     @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")
+    def save_memory(self, user_id: str, event_strings: list[str], **kwargs) -> bool:
         session_id = str(uuid.uuid1())
         messages = []
         for raw_events in event_strings:
@@ -105,31 +111,46 @@ class VikingDBLTMBackend(BaseLongTermMemoryBackend):
             "default_assistant_id": "assistant",
             "time": int(time.time() * 1000),
         }
-        response = self._client.add_messages(
+
+        logger.debug(
+            f"Request for add {len(messages)} memory to VikingDB: collection_name={self.index}, metadata={metadata}, session_id={session_id}"
+        )
+
+        client = self._get_client()
+        response = client.add_messages(
             collection_name=self.index,
             messages=messages,
             metadata=metadata,
             session_id=session_id,
         )
 
+        logger.debug(f"Response from add memory to VikingDB: {response}")
+
         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")
+    def search_memory(
+        self, user_id: str, query: str, top_k: int, **kwargs
+    ) -> list[str]:
         filter = {
             "user_id": user_id,
             "memory_type": ["sys_event_v1"],
         }
-        response = self._client.search_memory(
+
+        logger.debug(
+            f"Request for search memory in VikingDB: filter={filter}, collection_name={self.index}, query={query}, limit={top_k}"
+        )
+
+        client = self._get_client()
+        response = client.search_memory(
             collection_name=self.index, query=query, filter=filter, limit=top_k
         )
 
+        logger.debug(f"Response from search memory in VikingDB: {response}")
+
         if not response.get("code") == 0:
             raise ValueError(f"Search VikingDB memory error: {response}")
 

veadk/memory/short_term_memory.py

@@ -19,6 +19,7 @@ from google.adk.sessions import (
     BaseSessionService,
     DatabaseSessionService,
     InMemorySessionService,
+    Session,
 )
 from pydantic import BaseModel, Field, PrivateAttr
 
@@ -106,7 +107,7 @@ class ShortTermMemory(BaseModel):
         app_name: str,
         user_id: str,
         session_id: str,
-    ) -> None:
+    ) -> Session | None:
         if isinstance(self._session_service, DatabaseSessionService):
             list_sessions_response = await self._session_service.list_sessions(
                 app_name=app_name, user_id=user_id
@@ -116,13 +117,16 @@ class ShortTermMemory(BaseModel):
                 f"Loaded {len(list_sessions_response.sessions)} sessions from db {self.db_url}."
             )
 
-        if (
-            await self._session_service.get_session(
-                app_name=app_name, user_id=user_id, session_id=session_id
+        session = await self._session_service.get_session(
+            app_name=app_name, user_id=user_id, session_id=session_id
+        )
+
+        if session:
+            logger.info(
+                f"Session {session_id} already exists with app_name={app_name} user_id={user_id}."
             )
-            is None
-        ):
-            # create a new session for this running
-            await self._session_service.create_session(
+            return session
+        else:
+            return await self._session_service.create_session(
                 app_name=app_name, user_id=user_id, session_id=session_id
             )

veadk/memory/short_term_memory_backends/postgresql_backend.py

@@ -14,6 +14,7 @@
 
 from functools import cached_property
 from typing import Any
+from venv import logger
 
 from google.adk.sessions import (
     BaseSessionService,
@@ -33,7 +34,8 @@ class PostgreSqlSTMBackend(BaseShortTermMemoryBackend):
     postgresql_config: PostgreSqlConfig = Field(default_factory=PostgreSqlConfig)
 
     def model_post_init(self, context: Any) -> None:
-        self._db_url = f"postgresql+psycopg2://{self.postgresql_config.user}:{self.postgresql_config.password}@{self.postgresql_config.host}:{self.postgresql_config.port}/{self.postgresql_config.database}"
+        self._db_url = f"postgresql://{self.postgresql_config.user}:{self.postgresql_config.password}@{self.postgresql_config.host}:{self.postgresql_config.port}/{self.postgresql_config.database}"
+        logger.debug(self._db_url)
 
     @cached_property
     @override