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

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,7 +23,7 @@ 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,
 )
@@ -35,17 +36,24 @@ 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)
@@ -56,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 = VikingDBMemoryClient(
-            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:
@@ -101,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"],
+            "memory_type": ["sys_event_v1", "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

@@ -117,18 +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
-        ):
+            return session
+        else:
             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/runner.py

@@ -34,7 +34,7 @@ from veadk.evaluation import EvalSetRecorder
 from veadk.memory.short_term_memory import ShortTermMemory
 from veadk.types import MediaMessage
 from veadk.utils.logger import get_logger
-from veadk.utils.misc import formatted_timestamp, read_png_to_bytes
+from veadk.utils.misc import formatted_timestamp, read_file_to_bytes
 
 logger = get_logger(__name__)
 
@@ -50,11 +50,7 @@ RunnerMessage = Union[
 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 (
-                part.inline_data
-                and part.inline_data.mime_type == "image/png"
-                and self.upload_inline_data_to_tos
-            ):
+            if part.inline_data and self.upload_inline_data_to_tos:
                 await process_func(
                     part,
                     self.app_name,
@@ -105,9 +101,20 @@ def _convert_messages(
     if isinstance(messages, str):
         _messages = [types.Content(role="user", parts=[types.Part(text=messages)])]
     elif isinstance(messages, MediaMessage):
-        assert messages.media.endswith(".png"), (
-            "The MediaMessage only supports PNG format file for now."
+        import filetype
+
+        file_data = read_file_to_bytes(messages.media)
+
+        kind = filetype.guess(file_data)
+        if kind is None:
+            raise ValueError("Unsupported or unknown file type.")
+
+        mime_type = kind.mime
+
+        assert mime_type.startswith(("image/", "video/")), (
+            f"Unsupported media type: {mime_type}"
         )
+
         _messages = [
             types.Content(
                 role="user",
@@ -116,8 +123,8 @@ def _convert_messages(
                     types.Part(
                         inline_data=Blob(
                             display_name=messages.media,
-                            data=read_png_to_bytes(messages.media),
-                            mime_type="image/png",
+                            data=file_data,
+                            mime_type=mime_type,
                         )
                     ),
                 ],
@@ -277,7 +284,8 @@ class Runner(ADKRunner):
                         and event.content.parts[0].text is not None
                         and len(event.content.parts[0].text.strip()) > 0
                     ):
-                        final_output += event.content.parts[0].text
+                        final_output = event.content.parts[0].text
+                        logger.debug(f"Event output: {final_output}")
             except LlmCallsLimitExceededError as e:
                 logger.warning(f"Max number of llm calls limit exceeded: {e}")
                 final_output = ""