chore: add vision model request headers (#360)

* chore: add vision model request headers

* update image edit

Author: yaozheng-fang

veadk/tools/builtin_tools/generate_image.py

@@ -12,7 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from veadk.tools.builtin_tools.image_generate import image_generate  # noqa: F401
+from veadk.tools.builtin_tools.image_generate import (
+    image_generate,  # noqa: F401
+)
 from veadk.utils.logger import get_logger
 
 logger = get_logger(__name__)

veadk/tools/builtin_tools/image_edit.py

@@ -12,27 +12,30 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import base64
+import json
+import traceback
 from typing import Dict
+
 from google.adk.tools import ToolContext
+from opentelemetry import trace
+from opentelemetry.trace import Span
 from volcenginesdkarkruntime import Ark
+
 from veadk.config import getenv, settings
 from veadk.consts import (
     DEFAULT_IMAGE_EDIT_MODEL_API_BASE,
     DEFAULT_IMAGE_EDIT_MODEL_NAME,
 )
-import base64
-from opentelemetry import trace
-import traceback
-import json
-from veadk.version import VERSION
-from opentelemetry.trace import Span
 from veadk.utils.logger import get_logger
+from veadk.version import VERSION
 
 logger = get_logger(__name__)
 
 client = Ark(
     api_key=getenv(
-        "MODEL_EDIT_API_KEY", getenv("MODEL_AGENT_API_KEY", settings.model.api_key)
+        "MODEL_EDIT_API_KEY",
+        getenv("MODEL_AGENT_API_KEY", settings.model.api_key),
     ),
     base_url=getenv("MODEL_EDIT_API_BASE", DEFAULT_IMAGE_EDIT_MODEL_API_BASE),
 )
@@ -135,6 +138,14 @@ async def image_edit(
                 response = client.images.generate(
                     model=getenv("MODEL_EDIT_NAME", DEFAULT_IMAGE_EDIT_MODEL_NAME),
                     **inputs,
+                    extra_headers={
+                        "veadk-source": "veadk",
+                        "veadk-version": VERSION,
+                        "User-Agent": f"VeADK/{VERSION}",
+                        "X-Client-Request-Id": getenv(
+                            "MODEL_AGENT_CLIENT_REQ_ID", f"veadk/{VERSION}"
+                        ),
+                    },
                 )
                 output_part = None
                 if response.data and len(response.data) > 0:
@@ -154,7 +165,8 @@ async def image_edit(
                             image_bytes = base64.b64decode(image)
 
                             tos_url = _upload_image_to_tos(
-                                image_bytes=image_bytes, object_key=f"{image_name}.png"
+                                image_bytes=image_bytes,
+                                object_key=f"{image_name}.png",
                             )
                             if tos_url:
                                 tool_context.state[f"{image_name}_url"] = tos_url
@@ -277,10 +289,11 @@ def add_span_attributes(
 
 def _upload_image_to_tos(image_bytes: bytes, object_key: str) -> None:
     try:
-        from veadk.integrations.ve_tos.ve_tos import VeTOS
         import os
         from datetime import datetime
 
+        from veadk.integrations.ve_tos.ve_tos import VeTOS
+
         timestamp: str = datetime.now().strftime("%Y%m%d%H%M%S%f")[:-3]
         object_key = f"{timestamp}-{object_key}"
         bucket_name = os.getenv("DATABASE_TOS_BUCKET")

veadk/tools/builtin_tools/image_generate.py

@@ -26,7 +26,9 @@
 from opentelemetry import trace
 from opentelemetry.trace import Span
 from volcenginesdkarkruntime import Ark
-from volcenginesdkarkruntime.types.images.images import SequentialImageGenerationOptions
+from volcenginesdkarkruntime.types.images.images import (
+    SequentialImageGenerationOptions,
+)
 
 from veadk.config import getenv, settings
 from veadk.consts import (
@@ -41,7 +43,8 @@
 
 client = Ark(
     api_key=getenv(
-        "MODEL_IMAGE_API_KEY", getenv("MODEL_AGENT_API_KEY", settings.model.api_key)
+        "MODEL_IMAGE_API_KEY",
+        getenv("MODEL_AGENT_API_KEY", settings.model.api_key),
     ),
     base_url=getenv("MODEL_IMAGE_API_BASE", DEFAULT_IMAGE_GENERATE_MODEL_API_BASE),
 )
@@ -119,11 +122,24 @@ def handle_single_task_sync(
                 and sequential_image_generation == "auto"
                 and max_images
             ):
-                response = client.images.generate(
-                    model=getenv("MODEL_IMAGE_NAME", DEFAULT_IMAGE_GENERATE_MODEL_NAME),
-                    **inputs,
-                    sequential_image_generation_options=SequentialImageGenerationOptions(
-                        max_images=max_images
+                response = (
+                    client.images.generate(
+                        model=getenv(
+                            "MODEL_IMAGE_NAME",
+                            DEFAULT_IMAGE_GENERATE_MODEL_NAME,
+                        ),
+                        **inputs,
+                        sequential_image_generation_options=SequentialImageGenerationOptions(
+                            max_images=max_images
+                        ),
+                        extra_headers={
+                            "veadk-source": "veadk",
+                            "veadk-version": VERSION,
+                            "User-Agent": f"VeADK/{VERSION}",
+                            "X-Client-Request-Id": getenv(
+                                "MODEL_AGENT_CLIENT_REQ_ID", f"veadk/{VERSION}"
+                            ),
+                        },
                     ),
                 )
             else:
@@ -157,7 +173,8 @@ def handle_single_task_sync(
                             continue
                         image_bytes = base64.b64decode(b64)
                         image_url = _upload_image_to_tos(
-                            image_bytes=image_bytes, object_key=f"{image_name}.png"
+                            image_bytes=image_bytes,
+                            object_key=f"{image_name}.png",
                         )
                         if not image_url:
                             logger.error(f"Upload image to TOS failed: {image_name}")
@@ -367,7 +384,11 @@ def make_task(idx, item):
     logger.debug(
         f"image_generate success_list: {success_list}\nerror_list: {error_list}"
     )
-    return {"status": "success", "success_list": success_list, "error_list": error_list}
+    return {
+        "status": "success",
+        "success_list": success_list,
+        "error_list": error_list,
+    }
 
 
 def add_span_attributes(

veadk/tools/builtin_tools/video_generate.py

@@ -34,7 +34,8 @@
 
 client = Ark(
     api_key=getenv(
-        "MODEL_VIDEO_API_KEY", getenv("MODEL_AGENT_API_KEY", settings.model.api_key)
+        "MODEL_VIDEO_API_KEY",
+        getenv("MODEL_AGENT_API_KEY", settings.model.api_key),
     ),
     base_url=getenv("MODEL_VIDEO_API_BASE", DEFAULT_VIDEO_MODEL_API_BASE),
 )
@@ -48,6 +49,14 @@ async def generate(prompt, first_frame_image=None, last_frame_image=None):
                 content=[
                     {"type": "text", "text": prompt},
                 ],
+                extra_headers={
+                    "veadk-source": "veadk",
+                    "veadk-version": VERSION,
+                    "User-Agent": f"VeADK/{VERSION}",
+                    "X-Client-Request-Id": getenv(
+                        "MODEL_AGENT_CLIENT_REQ_ID", f"veadk/{VERSION}"
+                    ),
+                },
             )
         elif last_frame_image is None:
             response = client.content_generation.tasks.create(
@@ -62,6 +71,14 @@ async def generate(prompt, first_frame_image=None, last_frame_image=None):
                         },
                     ],
                 ),
+                extra_headers={
+                    "veadk-source": "veadk",
+                    "veadk-version": VERSION,
+                    "User-Agent": f"VeADK/{VERSION}",
+                    "X-Client-Request-Id": getenv(
+                        "MODEL_AGENT_CLIENT_REQ_ID", f"veadk/{VERSION}"
+                    ),
+                },
             )
         else:
             response = client.content_generation.tasks.create(
@@ -79,6 +96,14 @@ async def generate(prompt, first_frame_image=None, last_frame_image=None):
                         "role": "last_frame",
                     },
                 ],
+                extra_headers={
+                    "veadk-source": "veadk",
+                    "veadk-version": VERSION,
+                    "User-Agent": f"VeADK/{VERSION}",
+                    "X-Client-Request-Id": getenv(
+                        "MODEL_AGENT_CLIENT_REQ_ID", f"veadk/{VERSION}"
+                    ),
+                },
             )
     except:
         traceback.print_exc()