add image url in coze

Author: ZQlQZ

tests/test_tos.py

@@ -14,7 +14,7 @@
 
 import pytest
 from unittest import mock
-import veadk.database.tos.tos_client as tos_mod
+import veadk.integrations.ve_tos.ve_tos as tos_mod
 
 # 使用 pytest-asyncio
 pytest_plugins = ("pytest_asyncio",)
@@ -47,7 +47,7 @@ def __init__(self, msg):
 
 @pytest.fixture
 def tos_client(mock_client):
-    return tos_mod.TOSClient()
+    return tos_mod.VeTOS()
 
 
 def test_create_bucket_exists(tos_client, mock_client):
@@ -75,6 +75,7 @@ async def test_upload_bytes_success(tos_client, mock_client):
     result = await tos_client.upload("obj-key", data)
     assert result is True
     mock_client.put_object.assert_called_once()
+    mock_client.close.assert_called_once()
 
 
 @pytest.mark.asyncio
@@ -86,6 +87,7 @@ async def test_upload_file_success(tmp_path, tos_client, mock_client):
     result = await tos_client.upload("obj-key", str(file_path))
     assert result is True
     mock_client.put_object_from_file.assert_called_once()
+    mock_client.close.assert_called_once()
 
 
 def test_download_success(tmp_path, tos_client, mock_client):
@@ -101,8 +103,3 @@ def test_download_fail(tos_client, mock_client):
     mock_client.get_object.side_effect = Exception("boom")
     result = tos_client.download("obj-key", "somewhere.txt")
     assert result is False
-
-
-def test_close(tos_client, mock_client):
-    tos_client.close()
-    mock_client.close.assert_called_once()

veadk/database/tos/tos_client.py veadk/integrations/ve_tos/ve_tos.pyveadk/database/tos/tos_client.py renamed to veadk/integrations/ve_tos/ve_tos.py

@@ -20,6 +20,8 @@
 from typing import Union
 from pydantic import BaseModel, Field
 from typing import Any
+from urllib.parse import urlparse
+from datetime import datetime
 
 logger = get_logger(__name__)
 
@@ -43,7 +45,7 @@ class TOSConfig(BaseModel):
     )
 
 
-class TOSClient(BaseModel):
+class VeTOS(BaseModel):
     config: TOSConfig = Field(default_factory=TOSConfig)
 
     def model_post_init(self, __context: Any) -> None:
@@ -78,6 +80,23 @@ def create_bucket(self) -> bool:
             logger.error(f"Bucket creation failed: {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,
@@ -106,9 +125,11 @@ def _do_upload_bytes(self, object_key: str, bytes: bytes) -> bool:
                 bucket=self.config.bucket_name, key=object_key, content=bytes
             )
             logger.debug(f"Upload success, object_key: {object_key}")
+            self._close()
             return True
         except Exception as e:
             logger.error(f"Upload failed: {e}")
+            self._close()
             return False
 
     def _do_upload_file(self, object_key: str, file_path: str) -> bool:
@@ -121,10 +142,12 @@ def _do_upload_file(self, object_key: str, file_path: str) -> bool:
             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
         except Exception as e:
             logger.error(f"Upload failed: {e}")
+            self._close()
             return False
 
     def download(self, object_key: str, save_path: str) -> bool:
@@ -148,6 +171,6 @@ def download(self, object_key: str, save_path: str) -> bool:
 
             return False
 
-    def close(self):
+    def _close(self):
         if self._client:
             self._client.close()

veadk/runner.py

@@ -11,11 +11,8 @@
 # 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
 import asyncio
 from typing import Union
-from urllib.parse import urlparse
-from datetime import datetime
 
 from google.adk.agents import RunConfig
 from google.adk.agents.invocation_context import LlmCallsLimitExceededError
@@ -35,7 +32,7 @@
 from veadk.types import MediaMessage
 from veadk.utils.logger import get_logger
 from veadk.utils.misc import read_png_to_bytes
-from veadk.database.tos.tos_client import TOSClient
+from veadk.integrations.ve_tos.ve_tos import VeTOS
 
 logger = get_logger(__name__)
 
@@ -89,27 +86,6 @@ def __init__(
             plugins=plugins,
         )
 
-    def _build_tos_object_key(
-        self, user_id: str, app_name: str, session_id: str, data_path: 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}"
-        return object_key
-
-    def _upload_to_tos(self, data: Union[str, bytes], object_key: str):
-        tos_client = TOSClient()
-        asyncio.create_task(tos_client.upload(object_key, data))
-        tos_client.close()
-        return
-
     def _convert_messages(self, messages, session_id) -> list:
         if isinstance(messages, str):
             messages = [types.Content(role="user", parts=[types.Part(text=messages)])]
@@ -118,14 +94,16 @@ def _convert_messages(self, messages, session_id) -> list:
                 "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:
-                object_key = self._build_tos_object_key(
-                    self.user_id, self.app_name, session_id, messages.media
-                )
-                self._upload_to_tos(data, object_key)
+                asyncio.create_task(ve_tos.upload(object_key, data))
             except Exception as e:
                 logger.error(f"Upload to TOS failed: {e}")
-                object_key = None
+                tos_url = None
 
             messages = [
                 types.Content(
@@ -134,7 +112,7 @@ def _convert_messages(self, messages, session_id) -> list:
                         types.Part(text=messages.text),
                         types.Part(
                             inline_data=Blob(
-                                display_name=object_key,
+                                display_name=tos_url,
                                 data=data,
                                 mime_type="image/png",
                             )

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

@@ -137,6 +137,15 @@ def llm_gen_ai_prompt(params: LLMAttributesParams) -> ExtractorResponse:
                         if part.function_call.args
                         else json.dumps({})
                     )
+                # image
+                if part.inline_data:
+                    message[f"gen_ai.prompt.{idx}.type"] = "image_url"
+                    message[f"gen_ai.prompt.{idx}.image_url.name"] = (
+                        part.inline_data.display_name.split("/")[-1]
+                    )
+                    message[f"gen_ai.prompt.{idx}.image_url.url"] = (
+                        part.inline_data.display_name
+                    )
 
                 if message:
                     messages.append(message)
@@ -234,6 +243,14 @@ def llm_gen_ai_user_message(params: LLMAttributesParams) -> ExtractorResponse:
                             message_part[f"parts.{idx}.content"] = str(
                                 part.function_response
                             )
+                        if part.inline_data:
+                            message_part[f"parts.{idx}.type"] = "image_url"
+                            message_part[f"parts.{idx}.image_url.name"] = (
+                                part.inline_data.display_name.split("/")[-1]
+                            )
+                            message_part[f"parts.{idx}.image_url.url"] = (
+                                part.inline_data.display_name
+                            )
 
                     message_parts.append(message_part)
 

veadk/tracing/telemetry/telemetry.py

@@ -87,6 +87,17 @@ def _set_agent_input_attribute(
                     "gen_ai.user.message",
                     {f"parts.{idx}.type": "text", f"parts.{idx}.content": part.text},
                 )
+            if part.inline_data:
+                span.add_event(
+                    "gen_ai.user.message",
+                    {
+                        f"parts.{idx}.type": "image_url",
+                        f"parts.{idx}.image_url.name": part.inline_data.display_name.split(
+                            "/"
+                        )[-1],
+                        f"parts.{idx}.image_url.url": part.inline_data.display_name,
+                    },
+                )
 
 
 def _set_agent_output_attribute(span: _Span, llm_response: LlmResponse) -> None: