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

veadk/knowledgebase/entry.py

@@ -0,0 +1,25 @@
+# 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 pydantic import BaseModel
+
+
+class KnowledgebaseEntry(BaseModel):
+    """Represents a single entry in the knowledgebase."""
+
+    # The main content of the knowledgebase entry.
+    content: str
+
+    # Optional metadata associated with the entry.
+    metadata: dict | None = None

veadk/knowledgebase/knowledgebase.py

@@ -18,6 +18,7 @@ from pydantic import BaseModel, Field
 from typing_extensions import Union
 
 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__)
@@ -72,8 +73,8 @@ class KnowledgeBase(BaseModel):
     """Configuration for the backend"""
 
     top_k: int = 10
-    """Number of top similar documents to retrieve during search. 
-    
+    """Number of top similar documents to retrieve during search.
+
     Default is 10."""
 
     app_name: str = ""
@@ -130,11 +131,25 @@ class KnowledgeBase(BaseModel):
         """Add knowledge from text to knowledgebase"""
         return self._backend.add_from_text(text=text, **kwargs)
 
-    def search(self, query: str, top_k: int = 0, **kwargs) -> list[str]:
+    def search(self, query: str, top_k: int = 0, **kwargs) -> list[KnowledgebaseEntry]:
         """Search knowledge from knowledgebase"""
         if top_k == 0:
             top_k = self.top_k
-        return self._backend.search(query=query, top_k=top_k, **kwargs)
+
+        _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)

veadk/memory/__init__.py

@@ -25,7 +25,7 @@ def __getattr__(name):
         from veadk.memory.short_term_memory import ShortTermMemory
 
         return ShortTermMemory
-    if name == "LongTeremMemory":
+    if name == "LongTermMemory":
         from veadk.memory.long_term_memory import LongTermMemory
 
         return LongTermMemory

veadk/memory/long_term_memory.py

@@ -62,6 +62,12 @@ def _get_backend_cls(backend: str) -> type[BaseLongTermMemoryBackend]:
             )
 
             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}")
 
@@ -72,7 +78,7 @@ def build_long_term_memory_index(app_name: str, user_id: str):
 
 class LongTermMemory(BaseMemoryService, BaseModel):
     backend: Union[
-        Literal["local", "opensearch", "redis", "viking", "viking_mem"],
+        Literal["local", "opensearch", "redis", "viking", "viking_mem", "mem0"],
         BaseLongTermMemoryBackend,
     ] = "opensearch"
     """Long term memory backend type"""
@@ -88,6 +94,12 @@ class LongTermMemory(BaseMemoryService, BaseModel):
     user_id: str = ""
 
     def model_post_init(self, __context: Any) -> None:
+        if self.backend == "viking_mem":
+            logger.warning(
+                "The `viking_mem` backend is deprecated, please use `viking` instead."
+            )
+            self.backend = "viking"
+
         self._backend = None
 
         # Once user define a backend instance, use it directly
@@ -103,6 +115,14 @@ class LongTermMemory(BaseMemoryService, BaseModel):
                 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)
+            _index = self.backend_config.get("index", None)
+            if _index:
+                self._index = _index
+                logger.info(f"Long term memory index set to {self._index}.")
+            else:
+                logger.warning(
+                    "Cannot find index via backend_config, please set `index` parameter."
+                )
             return
 
         if self.app_name and self.user_id:
@@ -147,12 +167,6 @@ class LongTermMemory(BaseMemoryService, BaseModel):
         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)(
@@ -162,6 +176,13 @@ class LongTermMemory(BaseMemoryService, BaseModel):
                 f"Initialize long term memory backend now, index is {self._index}"
             )
 
+        if not self._index and 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
         event_strings = self._filter_and_convert_events(session.events)
 
         logger.info(
@@ -182,6 +203,23 @@ class LongTermMemory(BaseMemoryService, BaseModel):
     @override
     async def search_memory(self, *, app_name: str, user_id: str, query: str):
         # prevent model invoke `load_memory` before add session to this memory
+        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}"
+            )
+
+        if not self._index and 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. Initialized index: {self._index}, current built index: {build_long_term_memory_index(app_name, user_id)}. Search memory return empty list."
+            )
+            return SearchMemoryResponse(memories=[])
+
         if not self._backend:
             logger.error(
                 "Long term memory backend is not initialized, cannot search memory."

veadk/memory/long_term_memory_backends/mem0_backend.py

@@ -0,0 +1,144 @@
+# 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(
+                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], **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",
+                    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, **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
+            )
+
+            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/vikingdb_memory_backend.py

@@ -23,18 +23,14 @@ from typing_extensions import override
 
 import veadk.config  # noqa E401
 from veadk.config import getenv
+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__)
 
 
@@ -61,7 +57,7 @@ class VikingDBLTMBackend(BaseLongTermMemoryBackend):
             )
 
     def model_post_init(self, __context: Any) -> None:
-        self._client = VikingDBMemoryService(
+        self._client = VikingDBMemoryClient(
             ak=self.volcengine_access_key,
             sk=self.volcengine_secret_key,
             region=self.region,

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
@@ -122,7 +123,12 @@ class ShortTermMemory(BaseModel):
             )
             is None
         ):
-            # create a new session for this running
-            await self._session_service.create_session(
+            return await self._session_service.create_session(
                 app_name=app_name, user_id=user_id, session_id=session_id
             )
+        else:
+            logger.info(
+                f"Session {session_id} already exists with app_name={app_name} user_id={user_id}."
+            )
+
+            return None

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

veadk/runner.py

@@ -12,8 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import asyncio
 import functools
+import os
 from types import MethodType
 from typing import Union
 
@@ -47,7 +47,7 @@ RunnerMessage = Union[
 ]
 
 
-def pre_run_process(self, process_func, new_message, user_id, session_id):
+async def pre_run_process(self, process_func, new_message, user_id, session_id):
     if new_message.parts:
         for part in new_message.parts:
             if (
@@ -55,13 +55,12 @@ def pre_run_process(self, process_func, new_message, user_id, session_id):
                 and part.inline_data.mime_type == "image/png"
                 and self.upload_inline_data_to_tos
             ):
-                process_func(
+                await process_func(
                     part,
                     self.app_name,
                     user_id,
                     session_id,
                 )
-    return
 
 
 def post_run_process(self):
@@ -79,7 +78,7 @@ def intercept_new_message(process_func):
             new_message: types.Content,
             **kwargs,
         ):
-            pre_run_process(self, process_func, new_message, user_id, session_id)
+            await pre_run_process(self, process_func, new_message, user_id, session_id)
 
             async for event in func(
                 user_id=user_id,
@@ -137,27 +136,21 @@ def _convert_messages(
     return _messages
 
 
-def _upload_image_to_tos(
+async def _upload_image_to_tos(
     part: genai.types.Part, app_name: str, user_id: str, session_id: str
 ) -> None:
     try:
         if part.inline_data and part.inline_data.display_name and part.inline_data.data:
             from veadk.integrations.ve_tos.ve_tos import VeTOS
 
+            filename = os.path.basename(part.inline_data.display_name)
+            object_key = f"{app_name}/{user_id}-{session_id}-{filename}"
             ve_tos = VeTOS()
-
-            object_key, tos_url = ve_tos.build_tos_url(
-                user_id=user_id,
-                app_name=app_name,
-                session_id=session_id,
-                data_path=part.inline_data.display_name,
+            tos_url = ve_tos.build_tos_signed_url(object_key=object_key)
+            await ve_tos.async_upload_bytes(
+                object_key=object_key,
+                data=part.inline_data.data,
             )
-
-            upload_task = ve_tos.upload(object_key, part.inline_data.data)
-
-            if upload_task is not None:
-                asyncio.create_task(upload_task)
-
             part.inline_data.display_name = tos_url
     except Exception as e:
         logger.error(f"Upload to TOS failed: {e}")
@@ -226,7 +219,7 @@ class Runner(ADKRunner):
         )
 
         self.run_async = MethodType(
-            intercept_new_message(_upload_image_to_tos)(self.run_async), self
+            intercept_new_message(_upload_image_to_tos)(super().run_async), self
         )
 
     async def run(
@@ -256,15 +249,21 @@ class Runner(ADKRunner):
         )
 
         if self.short_term_memory:
-            await self.short_term_memory.create_session(
-                app_name=self.app_name, user_id=self.user_id, session_id=session_id
+            session = await self.short_term_memory.create_session(
+                app_name=self.app_name, user_id=user_id, session_id=session_id
+            )
+            assert session, (
+                f"Failed to create session with app_name={self.app_name}, user_id={user_id}, session_id={session_id}, "
+            )
+            logger.debug(
+                f"Auto create session: {session.id}, user_id: {session.user_id}, app_name: {self.app_name}"
             )
 
         final_output = ""
         for converted_message in converted_messages:
             try:
                 async for event in self.run_async(
-                    user_id=self.user_id,
+                    user_id=user_id,
                     session_id=session_id,
                     new_message=converted_message,
                     run_config=run_config,
@@ -366,19 +365,28 @@ class Runner(ADKRunner):
         )
         return eval_set_path
 
-    async def save_session_to_long_term_memory(self, session_id: str) -> None:
+    async def save_session_to_long_term_memory(
+        self, session_id: str, user_id: str = "", app_name: str = ""
+    ) -> None:
         if not self.long_term_memory:
             logger.warning("Long-term memory is not enabled. Failed to save session.")
             return
 
+        if not user_id:
+            user_id = self.user_id
+
+        if not app_name:
+            app_name = self.app_name
+
         session = await self.session_service.get_session(
-            app_name=self.app_name,
-            user_id=self.user_id,
+            app_name=app_name,
+            user_id=user_id,
             session_id=session_id,
         )
+
         if not session:
             logger.error(
-                f"Session {session_id} not found in session service, cannot save to long-term memory."
+                f"Session {session_id} (app_name={app_name}, user_id={user_id}) not found in session service, cannot save to long-term memory."
             )
             return