feat(runner): update more image formats (#234)

* Update _convert_message

* Update _convert_message

---------

Co-authored-by: wangjiaju <wangjiaju.716@bytedance.com>

Author: doraemonlove

pyproject.toml

@@ -35,6 +35,7 @@ dependencies = [
     "psycopg2-binary>=2.9.10", # For PostgreSQL database (short term memory)
     "pymysql>=1.1.1", # For MySQL database (short term memory)
     "opensearch-py==2.8.0",
+    "filetype>=1.2.0",
 ]
 
 [project.scripts]

veadk/runner.py

@@ -34,7 +34,7 @@
 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 @@
 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,
                         )
                     ),
                 ],

veadk/tools/builtin_tools/generate_image.py

@@ -28,7 +28,7 @@
 from veadk.config import getenv
 from veadk.consts import DEFAULT_IMAGE_GENERATE_MODEL_NAME, DEFAULT_MODEL_AGENT_API_BASE
 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
 from veadk.version import VERSION
 
 logger = get_logger(__name__)
@@ -299,7 +299,7 @@ async def image_generate(
                         artifact=Part(
                             inline_data=Blob(
                                 display_name=filename,
-                                data=read_png_to_bytes(image_tos_url),
+                                data=read_file_to_bytes(image_tos_url),
                                 mime_type=mimetypes.guess_type(image_tos_url)[0],
                             )
                         ),

veadk/utils/misc.py

@@ -36,18 +36,14 @@ def formatted_timestamp() -> str:
     return time.strftime("%Y%m%d%H%M%S", time.localtime())
 
 
-def read_png_to_bytes(png_path: str) -> bytes:
-    # Determine whether it is a local file or a network file
-    if png_path.startswith(("http://", "https://")):
-        # Network file: Download via URL and return bytes
-        response = requests.get(png_path)
-        response.raise_for_status()  # Check if the HTTP request is successful
+def read_file_to_bytes(file_path: str) -> bytes:
+    if file_path.startswith(("http://", "https://")):
+        response = requests.get(file_path)
+        response.raise_for_status()
         return response.content
     else:
-        # Local file
-        with open(png_path, "rb") as f:
-            data = f.read()
-    return data
+        with open(file_path, "rb") as f:
+            return f.read()
 
 
 def load_module_from_file(module_name: str, file_path: str) -> types.ModuleType: