chore: add vision model request headers

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_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()