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

veadk/integrations/ve_tos/ve_tos.py

@@ -15,7 +15,6 @@
 import os
 from veadk.config import getenv
 from veadk.utils.logger import get_logger
-import tos
 import asyncio
 from typing import Union
 from pydantic import BaseModel, Field
@@ -23,8 +22,20 @@ 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(
@@ -59,25 +70,76 @@ class VeTOS(BaseModel):
             logger.info("Connected to TOS successfully.")
         except Exception as e:
             logger.error(f"Client initialization failed:{e}")
-            return None
+            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"""
+        """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")
-            return True
         except tos.exceptions.TosServerError as e:
             if e.status_code == 404:
-                self._client.create_bucket(
-                    bucket=self.config.bucket_name,
-                    storage_class=tos.StorageClassType.Storage_Class_Standard,
-                    acl=tos.ACLType.ACL_Private,
-                )
-                logger.info(f"Bucket {self.config.bucket_name} created successfully")
-                return True
+                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"Bucket creation failed: {str(e)}")
+            logger.error(
+                f"Failed to set CORS rules for bucket {self.config.bucket_name}: {str(e)}"
+            )
             return False
 
     def build_tos_url(
@@ -103,55 +165,55 @@ class VeTOS(BaseModel):
         data: Union[str, bytes],
     ):
         if isinstance(data, str):
-            data_type = "file"
+            # data is a file path
+            return asyncio.to_thread(self._do_upload_file, object_key, data)
         elif isinstance(data, bytes):
-            data_type = "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)
-        if data_type == "file":
-            return asyncio.to_thread(self._do_upload_file, object_key, data)
-        elif data_type == "bytes":
-            return asyncio.to_thread(self._do_upload_bytes, object_key, data)
 
-    def _do_upload_bytes(self, object_key: str, bytes: bytes) -> bool:
+    def _do_upload_bytes(self, object_key: str, data: bytes) -> None:
         try:
             if not self._client:
-                return False
+                return
             if not self.create_bucket():
-                return False
+                return
             self._client.put_object(
-                bucket=self.config.bucket_name, key=object_key, content=bytes
+                bucket=self.config.bucket_name, key=object_key, content=data
             )
             logger.debug(f"Upload success, object_key: {object_key}")
             self._close()
-            return True
+            return
         except Exception as e:
             logger.error(f"Upload failed: {e}")
             self._close()
-            return False
+            return
 
-    def _do_upload_file(self, object_key: str, file_path: str) -> bool:
+    def _do_upload_file(self, object_key: str, file_path: str) -> None:
         try:
             if not self._client:
-                return False
+                return
             if not self.create_bucket():
-                return False
-
+                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 True
+            return
         except Exception as e:
             logger.error(f"Upload failed: {e}")
             self._close()
-            return False
+            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)
 

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

@@ -27,12 +27,12 @@ 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
 from veadk.utils.logger import get_logger
 from veadk.utils.misc import read_png_to_bytes
-from veadk.integrations.ve_tos.ve_tos import VeTOS
 
 logger = get_logger(__name__)
 
@@ -86,7 +86,9 @@ class Runner:
             plugins=plugins,
         )
 
-    def _convert_messages(self, messages, session_id) -> 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):
@@ -94,16 +96,26 @@ class Runner:
                 "The MediaMessage only supports PNG format file for now."
             )
             data = read_png_to_bytes(messages.media)
-
-            ve_tos = VeTOS()
-            object_key, tos_url = ve_tos.build_tos_url(
-                self.user_id, self.app_name, session_id, messages.media
-            )
-            try:
-                asyncio.create_task(ve_tos.upload(object_key, data))
-            except Exception as e:
-                logger.error(f"Upload to TOS failed: {e}")
-                tos_url = None
+            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(
@@ -123,7 +135,11 @@ class Runner:
         elif isinstance(messages, list):
             converted_messages = []
             for message in messages:
-                converted_messages.extend(self._convert_messages(message, session_id))
+                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)}")
@@ -142,7 +158,13 @@ class Runner:
         if run_config is not None:
             stream_mode = run_config.streaming_mode
         else:
-            run_config = RunConfig(streaming_mode=stream_mode)
+            run_config = RunConfig(
+                streaming_mode=stream_mode,
+                max_llm_calls=int(getenv("MODEL_AGENT_MAX_LLM_CALLS", 100)),
+            )
+
+        logger.info(f"Run config: {run_config}")
+
         try:
 
             async def event_generator():
@@ -172,6 +194,7 @@ class Runner:
                 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
 
@@ -182,8 +205,11 @@ class Runner:
         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, session_id)
+        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
@@ -231,7 +257,13 @@ class Runner:
         session_id: str,
         run_config: RunConfig | None = None,
     ):
-        run_config = RunConfig() if not run_config else run_config
+        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
@@ -263,6 +295,7 @@ class Runner:
                 final_output += chunk
         except LlmCallsLimitExceededError as e:
             logger.warning(f"Max number of llm calls limit exceeded: {e}")
+            final_output = ""
 
         return final_output
 

veadk/tracing/telemetry/attributes/extractors/common_attributes_extractors.py

@@ -49,6 +49,10 @@ def common_cozeloop_report_source(**kwargs) -> str:
     return "veadk"
 
 
+def common_cozeloop_call_type(**kwargs) -> str:
+    return kwargs.get("call_type")
+
+
 def llm_openinference_instrumentation_veadk(**kwargs) -> str:
     return VERSION
 
@@ -68,4 +72,5 @@ COMMON_ATTRIBUTES = {
     "user.id": common_gen_ai_user_id,  # CozeLoop / TLS required
     "session.id": common_gen_ai_session_id,  # CozeLoop / TLS required
     "cozeloop.report.source": common_cozeloop_report_source,  # CozeLoop required
+    "cozeloop.call_type": common_cozeloop_call_type,  # CozeLoop required
 }