veadk-python 0.2.4__py3-none-any.whl → 0.2.6__py3-none-any.whl

veadk/integrations/ve_tos/ve_tos.py

@@ -0,0 +1,238 @@
+# 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 os
+from veadk.config import getenv
+from veadk.utils.logger import get_logger
+import asyncio
+from typing import Union
+from pydantic import BaseModel, Field
+from typing import Any
+from urllib.parse import urlparse
+from datetime import datetime
+
+# Initialize logger before using it
+logger = get_logger(__name__)
+
+# Try to import tos module, and provide helpful error message if it fails
+try:
+    import tos
+except ImportError as e:
+    logger.error(
+        "Failed to import 'tos' module. Please install it using: pip install tos\n"
+    )
+    raise ImportError(
+        "Missing 'tos' module. Please install it using: pip install tos\n"
+    ) from e
+
+
+class TOSConfig(BaseModel):
+    region: str = Field(
+        default_factory=lambda: getenv("DATABASE_TOS_REGION"),
+        description="TOS region",
+    )
+    ak: str = Field(
+        default_factory=lambda: getenv("VOLCENGINE_ACCESS_KEY"),
+        description="Volcengine access key",
+    )
+    sk: str = Field(
+        default_factory=lambda: getenv("VOLCENGINE_SECRET_KEY"),
+        description="Volcengine secret key",
+    )
+    bucket_name: str = Field(
+        default_factory=lambda: getenv("DATABASE_TOS_BUCKET"),
+        description="TOS bucket name",
+    )
+
+
+class VeTOS(BaseModel):
+    config: TOSConfig = Field(default_factory=TOSConfig)
+
+    def model_post_init(self, __context: Any) -> None:
+        try:
+            self._client = tos.TosClientV2(
+                self.config.ak,
+                self.config.sk,
+                endpoint=f"tos-{self.config.region}.volces.com",
+                region=self.config.region,
+            )
+            logger.info("Connected to TOS successfully.")
+        except Exception as e:
+            logger.error(f"Client initialization failed:{e}")
+            self._client = None
+
+    def _refresh_client(self):
+        try:
+            if self._client:
+                self._client.close()
+            self._client = tos.TosClientV2(
+                self.config.ak,
+                self.config.sk,
+                endpoint=f"tos-{self.config.region}.volces.com",
+                region=self.config.region,
+            )
+            logger.info("refreshed client successfully.")
+        except Exception as e:
+            logger.error(f"Failed to refresh client: {str(e)}")
+            self._client = None
+
+    def create_bucket(self) -> bool:
+        """If the bucket does not exist, create it and set CORS rules"""
+        if not self._client:
+            logger.error("TOS client is not initialized")
+            return False
+        try:
+            self._client.head_bucket(self.config.bucket_name)
+            logger.info(f"Bucket {self.config.bucket_name} already exists")
+        except tos.exceptions.TosServerError as e:
+            if e.status_code == 404:
+                try:
+                    self._client.create_bucket(
+                        bucket=self.config.bucket_name,
+                        storage_class=tos.StorageClassType.Storage_Class_Standard,
+                        acl=tos.ACLType.ACL_Public_Read,  # 公开读
+                    )
+                    logger.info(
+                        f"Bucket {self.config.bucket_name} created successfully"
+                    )
+                    self._refresh_client()
+                except Exception as create_error:
+                    logger.error(f"Bucket creation failed: {str(create_error)}")
+                    return False
+            else:
+                logger.error(f"Bucket check failed: {str(e)}")
+                return False
+        except Exception as e:
+            logger.error(f"Bucket check failed: {str(e)}")
+            return False
+
+        # 确保在所有路径上返回布尔值
+        return self._set_cors_rules()
+
+    def _set_cors_rules(self) -> bool:
+        if not self._client:
+            logger.error("TOS client is not initialized")
+            return False
+        try:
+            rule = tos.models2.CORSRule(
+                allowed_origins=["*"],
+                allowed_methods=["GET", "HEAD"],
+                allowed_headers=["*"],
+                max_age_seconds=1000,
+            )
+            self._client.put_bucket_cors(self.config.bucket_name, [rule])
+            logger.info(
+                f"CORS rules for bucket {self.config.bucket_name} set successfully"
+            )
+            return True
+        except Exception as e:
+            logger.error(
+                f"Failed to set CORS rules for bucket {self.config.bucket_name}: {str(e)}"
+            )
+            return False
+
+    def build_tos_url(
+        self, user_id: str, app_name: str, session_id: str, data_path: str
+    ) -> tuple[str, str]:
+        """generate TOS object key"""
+        parsed_url = urlparse(data_path)
+
+        if parsed_url.scheme and parsed_url.scheme in ("http", "https", "ftp", "ftps"):
+            file_name = os.path.basename(parsed_url.path)
+        else:
+            file_name = os.path.basename(data_path)
+
+        timestamp: str = datetime.now().strftime("%Y%m%d%H%M%S%f")[:-3]
+        object_key: str = f"{app_name}-{user_id}-{session_id}/{timestamp}-{file_name}"
+        tos_url: str = f"https://{self.config.bucket_name}.tos-{self.config.region}.volces.com/{object_key}"
+
+        return object_key, tos_url
+
+    def upload(
+        self,
+        object_key: str,
+        data: Union[str, bytes],
+    ):
+        if isinstance(data, str):
+            # data is a file path
+            return asyncio.to_thread(self._do_upload_file, object_key, data)
+        elif isinstance(data, bytes):
+            # data is bytes content
+            return asyncio.to_thread(self._do_upload_bytes, object_key, data)
+        else:
+            error_msg = f"Upload failed: data type error. Only str (file path) and bytes are supported, got {type(data)}"
+            logger.error(error_msg)
+            raise ValueError(error_msg)
+
+    def _do_upload_bytes(self, object_key: str, data: bytes) -> None:
+        try:
+            if not self._client:
+                return
+            if not self.create_bucket():
+                return
+            self._client.put_object(
+                bucket=self.config.bucket_name, key=object_key, content=data
+            )
+            logger.debug(f"Upload success, object_key: {object_key}")
+            self._close()
+            return
+        except Exception as e:
+            logger.error(f"Upload failed: {e}")
+            self._close()
+            return
+
+    def _do_upload_file(self, object_key: str, file_path: str) -> None:
+        try:
+            if not self._client:
+                return
+            if not self.create_bucket():
+                return
+            self._client.put_object_from_file(
+                bucket=self.config.bucket_name, key=object_key, file_path=file_path
+            )
+            self._close()
+            logger.debug(f"Upload success, object_key: {object_key}")
+            return
+        except Exception as e:
+            logger.error(f"Upload failed: {e}")
+            self._close()
+            return
+
+    def download(self, object_key: str, save_path: str) -> bool:
+        """download image from TOS"""
+        if not self._client:
+            logger.error("TOS client is not initialized")
+            return False
+        try:
+            object_stream = self._client.get_object(self.config.bucket_name, object_key)
+
+            save_dir = os.path.dirname(save_path)
+            if save_dir and not os.path.exists(save_dir):
+                os.makedirs(save_dir, exist_ok=True)
+
+            with open(save_path, "wb") as f:
+                for chunk in object_stream:
+                    f.write(chunk)
+
+            logger.debug(f"Image download success, saved to: {save_path}")
+            return True
+
+        except Exception as e:
+            logger.error(f"Image download failed: {str(e)}")
+
+            return False
+
+    def _close(self):
+        if self._client:
+            self._client.close()

veadk/knowledgebase/knowledgebase.py

@@ -80,3 +80,11 @@ class KnowledgeBase:
         if len(result) == 0:
             logger.warning(f"No documents found in knowledgebase. Query: {query}")
         return result
+
+    def delete_doc(self, app_name: str, id: str) -> bool:
+        index = build_knowledgebase_index(app_name)
+        return self.adapter.delete_doc(index=index, id=id)
+
+    def list_docs(self, app_name: str, offset: int = 0, limit: int = 100) -> list[dict]:
+        index = build_knowledgebase_index(app_name)
+        return self.adapter.list_docs(index=index, offset=offset, limit=limit)

veadk/runner.py

@@ -11,9 +11,11 @@
 # 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 asyncio
 from typing import Union
 
 from google.adk.agents import RunConfig
+from google.adk.agents.invocation_context import LlmCallsLimitExceededError
 from google.adk.agents.run_config import StreamingMode
 from google.adk.plugins.base_plugin import BasePlugin
 from google.adk.runners import Runner as ADKRunner
@@ -25,6 +27,7 @@ from veadk.agent import Agent
 from veadk.agents.loop_agent import LoopAgent
 from veadk.agents.parallel_agent import ParallelAgent
 from veadk.agents.sequential_agent import SequentialAgent
+from veadk.config import getenv
 from veadk.evaluation import EvalSetRecorder
 from veadk.memory.short_term_memory import ShortTermMemory
 from veadk.types import MediaMessage
@@ -49,20 +52,25 @@ class Runner:
     def __init__(
         self,
         agent: VeAgent,
-        short_term_memory: ShortTermMemory,
+        short_term_memory: ShortTermMemory | None = None,
         plugins: list[BasePlugin] | None = None,
         app_name: str = "veadk_default_app",
         user_id: str = "veadk_default_user",
     ):
-        # basic settings
         self.app_name = app_name
         self.user_id = user_id
 
-        # agent settings
         self.agent = agent
 
-        self.short_term_memory = short_term_memory
-        self.session_service = short_term_memory.session_service
+        if not short_term_memory:
+            logger.info(
+                "No short term memory provided, using a in-memory memory by default."
+            )
+            self.short_term_memory = ShortTermMemory()
+        else:
+            self.short_term_memory = short_term_memory
+
+        self.session_service = self.short_term_memory.session_service
 
         # prevent VeRemoteAgent has no long-term memory attr
         if isinstance(self.agent, Agent):
@@ -78,13 +86,37 @@ class Runner:
             plugins=plugins,
         )
 
-    def _convert_messages(self, messages) -> list:
+    def _convert_messages(
+        self, messages, session_id, upload_inline_data_to_tos
+    ) -> list:
         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."
             )
+            data = read_png_to_bytes(messages.media)
+            tos_url = "<tos_url>"
+            if upload_inline_data_to_tos:
+                try:
+                    from veadk.integrations.ve_tos.ve_tos import VeTOS
+
+                    ve_tos = VeTOS()
+                    object_key, tos_url = ve_tos.build_tos_url(
+                        self.user_id, self.app_name, session_id, messages.media
+                    )
+                    upload_task = ve_tos.upload(object_key, data)
+                    if upload_task is not None:
+                        asyncio.create_task(upload_task)
+                except Exception as e:
+                    logger.error(f"Upload to TOS failed: {e}")
+                    tos_url = None
+
+            else:
+                logger.warning(
+                    "Loss of multimodal data may occur in the tracing process."
+                )
+
             messages = [
                 types.Content(
                     role="user",
@@ -92,8 +124,8 @@ class Runner:
                         types.Part(text=messages.text),
                         types.Part(
                             inline_data=Blob(
-                                display_name=messages.media,
-                                data=read_png_to_bytes(messages.media),
+                                display_name=tos_url,
+                                data=data,
                                 mime_type="image/png",
                             )
                         ),
@@ -103,7 +135,11 @@ class Runner:
         elif isinstance(messages, list):
             converted_messages = []
             for message in messages:
-                converted_messages.extend(self._convert_messages(message))
+                converted_messages.extend(
+                    self._convert_messages(
+                        message, session_id, upload_inline_data_to_tos
+                    )
+                )
             messages = converted_messages
         else:
             raise ValueError(f"Unknown message type: {type(messages)}")
@@ -114,35 +150,51 @@ class Runner:
         self,
         session_id: str,
         message: types.Content,
+        run_config: RunConfig | None = None,
         stream: bool = False,
     ):
         stream_mode = StreamingMode.SSE if stream else StreamingMode.NONE
 
-        async def event_generator():
-            async for event in self.runner.run_async(
-                user_id=self.user_id,
-                session_id=session_id,
-                new_message=message,
-                run_config=RunConfig(streaming_mode=stream_mode),
-            ):
-                if event.get_function_calls():
-                    for function_call in event.get_function_calls():
-                        logger.debug(f"Function call: {function_call}")
-                elif (
-                    event.content is not None
-                    and event.content.parts
-                    and event.content.parts[0].text is not None
-                    and len(event.content.parts[0].text.strip()) > 0
-                ):
-                    yield event.content.parts[0].text
+        if run_config is not None:
+            stream_mode = run_config.streaming_mode
+        else:
+            run_config = RunConfig(
+                streaming_mode=stream_mode,
+                max_llm_calls=int(getenv("MODEL_AGENT_MAX_LLM_CALLS", 100)),
+            )
 
-        final_output = ""
-        async for chunk in event_generator():
+        logger.info(f"Run config: {run_config}")
+
+        try:
+
+            async def event_generator():
+                async for event in self.runner.run_async(
+                    user_id=self.user_id,
+                    session_id=session_id,
+                    new_message=message,
+                    run_config=run_config,
+                ):
+                    if event.get_function_calls():
+                        for function_call in event.get_function_calls():
+                            logger.debug(f"Function call: {function_call}")
+                    elif (
+                        event.content is not None
+                        and event.content.parts
+                        and event.content.parts[0].text is not None
+                        and len(event.content.parts[0].text.strip()) > 0
+                    ):
+                        yield event.content.parts[0].text
+
+            final_output = ""
+            async for chunk in event_generator():
+                if stream:
+                    print(chunk, end="", flush=True)
+                final_output += chunk
             if stream:
-                print(chunk, end="", flush=True)
-            final_output += chunk
-        if stream:
-            print()  # end with a new line
+                print()  # end with a new line
+        except LlmCallsLimitExceededError as e:
+            logger.warning(f"Max number of llm calls limit exceeded: {e}")
+            final_output = ""
 
         return final_output
 
@@ -151,9 +203,13 @@ class Runner:
         messages: RunnerMessage,
         session_id: str,
         stream: bool = False,
+        run_config: RunConfig | None = None,
         save_tracing_data: bool = False,
+        upload_inline_data_to_tos: bool = False,
     ):
-        converted_messages: list = self._convert_messages(messages)
+        converted_messages: list = self._convert_messages(
+            messages, session_id, upload_inline_data_to_tos
+        )
 
         await self.short_term_memory.create_session(
             app_name=self.app_name, user_id=self.user_id, session_id=session_id
@@ -163,7 +219,9 @@ class Runner:
 
         final_output = ""
         for converted_message in converted_messages:
-            final_output = await self._run(session_id, converted_message, stream)
+            final_output = await self._run(
+                session_id, converted_message, run_config, stream
+            )
 
         # try to save tracing file
         if save_tracing_data:
@@ -193,6 +251,54 @@ class Runner:
             logger.warning(f"Get tracer id failed as {e}")
             return "<unknown_trace_id>"
 
+    async def run_with_raw_message(
+        self,
+        message: types.Content,
+        session_id: str,
+        run_config: RunConfig | None = None,
+    ):
+        run_config = (
+            RunConfig(max_llm_calls=int(getenv("MODEL_AGENT_MAX_LLM_CALLS", 100)))
+            if not run_config
+            else run_config
+        )
+
+        logger.info(f"Run config: {run_config}")
+
+        await self.short_term_memory.create_session(
+            app_name=self.app_name, user_id=self.user_id, session_id=session_id
+        )
+
+        try:
+
+            async def event_generator():
+                async for event in self.runner.run_async(
+                    user_id=self.user_id,
+                    session_id=session_id,
+                    new_message=message,
+                    run_config=run_config,
+                ):
+                    if event.get_function_calls():
+                        for function_call in event.get_function_calls():
+                            logger.debug(f"Function call: {function_call}")
+                    elif (
+                        event.content is not None
+                        and event.content.parts
+                        and event.content.parts[0].text is not None
+                        and len(event.content.parts[0].text.strip()) > 0
+                    ):
+                        yield event.content.parts[0].text
+
+            final_output = ""
+
+            async for chunk in event_generator():
+                final_output += chunk
+        except LlmCallsLimitExceededError as e:
+            logger.warning(f"Max number of llm calls limit exceeded: {e}")
+            final_output = ""
+
+        return final_output
+
     def _print_trace_id(self) -> None:
         if not isinstance(self.agent, Agent):
             logger.warning(