diff --git a/config.yaml.full b/config.yaml.full index daab5fe0..aaf601d0 100644 --- a/config.yaml.full +++ b/config.yaml.full @@ -16,6 +16,18 @@ model: dim: 2560 api_base: https://ark.cn-beijing.volces.com/api/v3/embeddings api_key: + video: + name: doubao-seedance-1-0-pro-250528 + api_base: https://ark.cn-beijing.volces.com/api/v3/ + api_key: + image: + name: doubao-seedream-3-0-t2i-250415 + api_base: https://ark.cn-beijing.volces.com/api/v3/ + api_key: + edit: + name: doubao-seededit-3-0-i2i-250628 + api_base: https://ark.cn-beijing.volces.com/api/v3/ + api_key: volcengine: # [optional] for Viking DB and `web_search` tool diff --git a/veadk/tools/builtin_tools/image_edit.py b/veadk/tools/builtin_tools/image_edit.py index 92d52f9a..3f126d48 100644 --- a/veadk/tools/builtin_tools/image_edit.py +++ b/veadk/tools/builtin_tools/image_edit.py @@ -11,84 +11,226 @@ # 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. + from typing import Dict from google.adk.tools import ToolContext from google.genai import types from volcenginesdkarkruntime import Ark from veadk.config import getenv 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 logger = get_logger(__name__) client = Ark( - api_key=getenv("MODEL_IMAGE_API_KEY"), - base_url=getenv("MODEL_IMAGE_API_BASE"), + api_key=getenv("MODEL_EDIT_API_KEY"), + base_url=getenv("MODEL_EDIT_API_BASE"), ) async def image_edit( - origin_image: str, - image_name: str, - image_prompt: str, - response_format: str, - guidance_scale: float, - watermark: bool, - seed: int, + params: list, tool_context: ToolContext, ) -> Dict: - """Edit an image accoding to the prompt. + """ + Edit images in batch according to prompts and optional settings. + + Each item in `params` describes a single image-edit request. Args: - origin_image: The url or the base64 string of the edited image. - image_name: The name of the generated image. - image_prompt: The prompt that describes the image. - response_format: str, b64_json or url, default url. - guidance_scale: default 2.5. - watermark: default True. - seed: default -1. + params (list[dict]): + A list of image editing requests. Each item supports: + + Required: + - origin_image (str): + The URL or Base64 string of the original image to edit. + Example: + * URL: "https://example.com/image.png" + * Base64: "data:image/png;base64," + + - prompt (str): + The textual description/instruction for editing the image. + Supports English and Chinese. + + Optional: + - image_name (str): + Name/identifier for the generated image. + - response_format (str): + Format of the returned image. + * "url": JPEG link (default) + * "b64_json": Base64 string in JSON + + - guidance_scale (float): + How strongly the prompt affects the result. + Range: [1.0, 10.0], default 2.5. + + - watermark (bool): + Whether to add watermark. + Default: True. + + - seed (int): + Random seed for reproducibility. + Range: [-1, 2^31-1], default -1 (random). + + Returns: + Dict: API response containing generated image metadata. + Example: + { + "status": "success", + "success_list": [{"image_name": ""}], + "error_list": [{}] + } + + Notes: + - Uses SeedEdit 3.0 model. + - Provide the same `seed` for consistent outputs across runs. + - A high `guidance_scale` enforces stricter adherence to text prompt. """ - try: - response = client.images.generate( - model=getenv("MODEL_EDIT_NAME"), - image=origin_image, - prompt=image_prompt, - response_format=response_format, - guidance_scale=guidance_scale, - watermark=watermark, - seed=seed, - ) - - if response.data and len(response.data) > 0: - for item in response.data: - if response_format == "url": - image = item.url - tool_context.state["generated_image_url"] = image - - elif response_format == "b64_json": - image = item.b64_json - image_bytes = base64.b64decode(image) - - tool_context.state["generated_image_url"] = ( - f"data:image/jpeg;base64,{image}" - ) - - report_artifact = types.Part.from_bytes( - data=image_bytes, mime_type="image/png" - ) - await tool_context.save_artifact(image_name, report_artifact) - logger.debug(f"Image saved as ADK artifact: {image_name}") - - return {"status": "success", "image_name": image_name, "image": image} - else: - error_details = f"No images returned by Doubao model: {response}" + success_list = [] + error_list = [] + for idx, item in enumerate(params): + image_name = item.get("image_name", f"generated_image_{idx}") + prompt = item.get("prompt") + origin_image = item.get("origin_image") + response_format = item.get("response_format", "url") + guidance_scale = item.get("guidance_scale", 2.5) + watermark = item.get("watermark", True) + seed = item.get("seed", -1) + + try: + tracer = trace.get_tracer("gcp.vertex.agent") + with tracer.start_as_current_span("call_llm") as span: + inputs = { + "prompt": prompt, + "image": origin_image, + "response_format": response_format, + "guidance_scale": guidance_scale, + "watermark": watermark, + "seed": seed, + } + input_part = { + "role": "user", + "content": json.dumps(inputs, ensure_ascii=False), + } + response = client.images.generate( + model=getenv("MODEL_EDIT_NAME"), **inputs + ) + output_part = None + if response.data and len(response.data) > 0: + for item in response.data: + if response_format == "url": + image = item.url + tool_context.state[f"{image_name}_url"] = image + output_part = { + "message.role": "model", + "message.content": image, + } + elif response_format == "b64_json": + image = item.b64_json + image_bytes = base64.b64decode(image) + + tool_context.state[f"{image_name}_url"] = ( + f"data:image/jpeg;base64,{image}" + ) + + report_artifact = types.Part.from_bytes( + data=image_bytes, mime_type="image/png" + ) + await tool_context.save_artifact( + image_name, report_artifact + ) + logger.debug(f"Image saved as ADK artifact: {image_name}") + + success_list.append({image_name: image}) + else: + error_details = f"No images returned by Doubao model: {response}" + logger.error(error_details) + error_list.append(image_name) + + add_span_attributes( + span, + tool_context, + input_part=input_part, + output_part=output_part, + output_tokens=response.usage.output_tokens, + total_tokens=response.usage.total_tokens, + request_model=getenv("MODEL_EDIT_NAME"), + response_model=getenv("MODEL_EDIT_NAME"), + ) + + except Exception as e: + error_details = f"No images returned by Doubao model: {e}" logger.error(error_details) - return {"status": "error", "message": error_details} + traceback.print_exc() + error_list.append(image_name) - except Exception as e: + if len(success_list) == 0: return { "status": "error", - "message": f"Doubao image generation failed: {str(e)}", + "success_list": success_list, + "error_list": error_list, + } + else: + return { + "status": "success", + "success_list": success_list, + "error_list": error_list, } + + +def add_span_attributes( + span: Span, + tool_context: ToolContext, + input_part: dict = None, + output_part: dict = None, + input_tokens: int = None, + output_tokens: int = None, + total_tokens: int = None, + request_model: str = None, + response_model: str = None, +): + try: + # common attributes + app_name = tool_context._invocation_context.app_name + user_id = tool_context._invocation_context.user_id + agent_name = tool_context.agent_name + session_id = tool_context._invocation_context.session.id + span.set_attribute("gen_ai.agent.name", agent_name) + span.set_attribute("openinference.instrumentation.veadk", VERSION) + span.set_attribute("gen_ai.app.name", app_name) + span.set_attribute("gen_ai.user.id", user_id) + span.set_attribute("gen_ai.session.id", session_id) + span.set_attribute("agent_name", agent_name) + span.set_attribute("agent.name", agent_name) + span.set_attribute("app_name", app_name) + span.set_attribute("app.name", app_name) + span.set_attribute("user.id", user_id) + span.set_attribute("session.id", session_id) + span.set_attribute("cozeloop.report.source", "veadk") + + # llm attributes + span.set_attribute("gen_ai.system", "openai") + span.set_attribute("gen_ai.operation.name", "chat") + if request_model: + span.set_attribute("gen_ai.request.model", request_model) + if response_model: + span.set_attribute("gen_ai.response.model", response_model) + if total_tokens: + span.set_attribute("gen_ai.usage.total_tokens", total_tokens) + if output_tokens: + span.set_attribute("gen_ai.usage.output_tokens", output_tokens) + if input_tokens: + span.set_attribute("gen_ai.usage.input_tokens", input_tokens) + if input_part: + span.add_event("gen_ai.user.message", input_part) + if output_part: + span.add_event("gen_ai.choice", output_part) + + except Exception: + traceback.print_exc() diff --git a/veadk/tools/builtin_tools/image_generate.py b/veadk/tools/builtin_tools/image_generate.py index b069078e..392a9e74 100644 --- a/veadk/tools/builtin_tools/image_generate.py +++ b/veadk/tools/builtin_tools/image_generate.py @@ -11,6 +11,7 @@ # 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. + from typing import Dict from google.genai import types @@ -18,7 +19,11 @@ from veadk.config import getenv import base64 from volcenginesdkarkruntime import Ark - +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 logger = get_logger(__name__) @@ -30,66 +35,202 @@ async def image_generate( - image_name: str, - image_prompt: str, - response_format: str, - size: str, - guidance_scale: float, - watermark: bool, - seed: int, + params: list, tool_context: ToolContext, ) -> Dict: - """Generate an image accoding to the prompt. + """ + Generate images in batch according to prompts and optional settings. + + Each item in `params` describes a single image-generation request. Args: - image_name: The name of the generated image. - image_prompt: The prompt that describes the image. - response_format: str, b64_json or url, default url. - size: default 1024x1024. - guidance_scale: default 2.5. - watermark: default True. - seed: default -1. + params (list[dict]): + A list of image generation requests. Each item supports: + + Required: + - prompt (str): + The textual description of the desired image. + Supports English and Chinese. + + Optional: + - image_name (str): + Name/identifier for the generated image. + + - response_format (str): + Format of the returned image. + * "url": JPEG link (default) + * "b64_json": Base64 string in JSON + - size (str): + Resolution of the generated image. + Default: "1024x1024". + Must be within [512x512, 2048x2048]. + Common options: 1024x1024, 864x1152, 1280x720, etc. + + - guidance_scale (float): + How strongly the prompt affects the result. + Range: [1.0, 10.0], default 2.5. + + - watermark (bool): + Whether to add watermark. + Default: True. + + - seed (int): + Random seed for reproducibility. + Range: [-1, 2^31-1], default -1 (random). + + Returns: + Dict: API response containing generated image metadata. + Example: + { + "status": "success", + "success_list": [{"image_name": ""}], + "error_list": [{}] + } + + Notes: + - Best suited for creating original images from text. + - Use a fixed `seed` for reproducibility. + - Choose appropriate `size` for desired aspect ratio. """ - try: - response = client.images.generate( - model=getenv("MODEL_IMAGE_NAME"), - prompt=image_prompt, - response_format=response_format, - size=size, - guidance_scale=guidance_scale, - watermark=watermark, - seed=seed, - ) - - if response.data and len(response.data) > 0: - for item in response.data: - if response_format == "url": - image = item.url - tool_context.state["generated_image_url"] = image - - elif response_format == "b64_json": - image = item.b64_json - image_bytes = base64.b64decode(image) - - tool_context.state["generated_image_url"] = ( - f"data:image/jpeg;base64,{image}" - ) - - report_artifact = types.Part.from_bytes( - data=image_bytes, mime_type="image/png" - ) - await tool_context.save_artifact(image_name, report_artifact) - logger.debug(f"Image saved as ADK artifact: {image_name}") - - return {"status": "success", "image_name": image_name, "image": image} - else: - error_details = f"No images returned by Doubao model: {response}" + success_list = [] + error_list = [] + for idx, item in enumerate(params): + prompt = item.get("prompt", "") + image_name = item.get("image_name", f"generated_image_{idx}") + response_format = item.get("response_format", "url") + size = item.get("size", "1024x1024") + guidance_scale = item.get("guidance_scale", 2.5) + watermark = item.get("watermark", True) + seed = item.get("seed", -1) + + try: + tracer = trace.get_tracer("gcp.vertex.agent") + with tracer.start_as_current_span("call_llm") as span: + inputs = { + "prompt": prompt, + "response_format": response_format, + "size": size, + "guidance_scale": guidance_scale, + "watermark": watermark, + "seed": seed, + } + input_part = { + "role": "user", + "content": json.dumps(inputs, ensure_ascii=False), + } + response = client.images.generate( + model=getenv("MODEL_IMAGE_NAME"), **inputs + ) + output_part = None + if response.data and len(response.data) > 0: + for item in response.data: + if response_format == "url": + image = item.url + tool_context.state[f"{image_name}_url"] = image + output_part = { + "message.role": "model", + "message.content": image, + } + elif response_format == "b64_json": + image = item.b64_json + image_bytes = base64.b64decode(image) + + tool_context.state[f"{image_name}_url"] = ( + f"data:image/jpeg;base64,{image}" + ) + + report_artifact = types.Part.from_bytes( + data=image_bytes, mime_type="image/png" + ) + await tool_context.save_artifact( + image_name, report_artifact + ) + logger.debug(f"Image saved as ADK artifact: {image_name}") + + success_list.append({image_name: image}) + else: + error_details = f"No images returned by Doubao model: {response}" + logger.error(error_details) + error_list.append(image_name) + + add_span_attributes( + span, + tool_context, + input_part=input_part, + output_part=output_part, + output_tokens=response.usage.output_tokens, + total_tokens=response.usage.total_tokens, + request_model=getenv("MODEL_IMAGE_NAME"), + response_model=getenv("MODEL_IMAGE_NAME"), + ) + + except Exception as e: + error_details = f"No images returned by Doubao model: {e}" logger.error(error_details) - return {"status": "error", "message": error_details} + error_list.append(image_name) - except Exception as e: + if len(success_list) == 0: return { "status": "error", - "message": f"Doubao image generation failed: {str(e)}", + "success_list": success_list, + "error_list": error_list, + } + else: + return { + "status": "success", + "success_list": success_list, + "error_list": error_list, } + + +def add_span_attributes( + span: Span, + tool_context: ToolContext, + input_part: dict = None, + output_part: dict = None, + input_tokens: int = None, + output_tokens: int = None, + total_tokens: int = None, + request_model: str = None, + response_model: str = None, +): + try: + # common attributes + app_name = tool_context._invocation_context.app_name + user_id = tool_context._invocation_context.user_id + agent_name = tool_context.agent_name + session_id = tool_context._invocation_context.session.id + span.set_attribute("gen_ai.agent.name", agent_name) + span.set_attribute("openinference.instrumentation.veadk", VERSION) + span.set_attribute("gen_ai.app.name", app_name) + span.set_attribute("gen_ai.user.id", user_id) + span.set_attribute("gen_ai.session.id", session_id) + span.set_attribute("agent_name", agent_name) + span.set_attribute("agent.name", agent_name) + span.set_attribute("app_name", app_name) + span.set_attribute("app.name", app_name) + span.set_attribute("user.id", user_id) + span.set_attribute("session.id", session_id) + span.set_attribute("cozeloop.report.source", "veadk") + + # llm attributes + span.set_attribute("gen_ai.system", "openai") + span.set_attribute("gen_ai.operation.name", "chat") + if request_model: + span.set_attribute("gen_ai.request.model", request_model) + if response_model: + span.set_attribute("gen_ai.response.model", response_model) + if total_tokens: + span.set_attribute("gen_ai.usage.total_tokens", total_tokens) + if output_tokens: + span.set_attribute("gen_ai.usage.output_tokens", output_tokens) + if input_tokens: + span.set_attribute("gen_ai.usage.input_tokens", input_tokens) + if input_part: + span.add_event("gen_ai.user.message", input_part) + if output_part: + span.add_event("gen_ai.choice", output_part) + + except Exception: + traceback.print_exc() diff --git a/veadk/tools/builtin_tools/video_generate.py b/veadk/tools/builtin_tools/video_generate.py index d3212f85..b4ff88df 100644 --- a/veadk/tools/builtin_tools/video_generate.py +++ b/veadk/tools/builtin_tools/video_generate.py @@ -11,12 +11,17 @@ # 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. + from typing import Dict from google.adk.tools import ToolContext from volcenginesdkarkruntime import Ark from veadk.config import getenv import time import traceback +import json +from veadk.version import VERSION +from opentelemetry import trace +from opentelemetry.trace import Span from veadk.utils.logger import get_logger @@ -28,7 +33,7 @@ ) -async def generate(tool_context, prompt, first_frame_image=None, last_frame_image=None): +async def generate(prompt, first_frame_image=None, last_frame_image=None): try: if first_frame_image is None: logger.debug("text generation") @@ -75,67 +80,247 @@ async def generate(tool_context, prompt, first_frame_image=None, last_frame_imag async def video_generate(params: list, tool_context: ToolContext) -> Dict: - """Generate video in batch according to the prompt. + """ + Generate videos in **batch** from text prompts, optionally guided by a first/last frame, + and fine-tuned via *model text commands* (a.k.a. `parameters` appended to the prompt). + + This API creates video-generation tasks. Each item in `params` describes a single video. + The function submits all items in one call and returns task metadata for tracking. Args: - params: - video_name: The name of the generated video. - first_frame: The first frame of the video, url or base64 string, or None. - last_frame:The last frame of the video, url or base64 string, or None. - prompt:The prompt of the video. + params (list[dict]): + A list of video generation requests. Each item supports the fields below. + + Required per item: + - video_name (str): + Name/identifier of the output video file. + + - prompt (str): + Text describing the video to generate. Supports zh/EN. + You may append **model text commands** after the prompt to control resolution, + aspect ratio, duration, fps, watermark, seed, camera lock, etc. + Format: `... --rs --rt --dur --fps --wm --seed --cf ` + Example: + "小猫骑着滑板穿过公园。 --rs 720p --rt 16:9 --dur 5 --fps 24 --wm true --seed 11 --cf false" + + Optional per item: + - first_frame (str | None): + URL or Base64 string (data URL) for the **first frame** (role = `first_frame`). + Use when you want the clip to start from a specific image. + + - last_frame (str | None): + URL or Base64 string (data URL) for the **last frame** (role = `last_frame`). + Use when you want the clip to end on a specific image. + + Notes on first/last frame: + * When both frames are provided, **match width/height** to avoid cropping; if they differ, + the tail frame may be auto-cropped to fit. + * If you only need one guided frame, provide either `first_frame` or `last_frame` (not both). + + Image input constraints (for first/last frame): + - Formats: jpeg, png, webp, bmp, tiff, gif + - Aspect ratio (宽:高): 0.4–2.5 + - Width/Height (px): 300–6000 + - Size: < 30 MB + - Base64 data URL example: `data:image/png;base64,` + + Model text commands (append after the prompt; unsupported keys are ignored by some models): + --rs / --resolution Video resolution. Common values: 480p, 720p, 1080p. + Default depends on model (e.g., doubao-seedance-1-0-pro: 1080p, + some others default 720p). + + --rt / --ratio Aspect ratio. Typical: 16:9 (default), 9:16, 4:3, 3:4, 1:1, 2:1, 21:9. + Some models support `keep_ratio` (keep source image ratio) or `adaptive` + (auto choose suitable ratio). + + --dur / --duration Clip length in seconds. Seedance supports **3–12 s**; + Wan2.1 仅支持 5 s。Default varies by model. + + --fps / --framespersecond Frame rate. Common: 16 or 24 (model-dependent; e.g., seaweed=24, wan2.1=16). + + --wm / --watermark Whether to add watermark. Default: **false** (per doc). + + --seed Random seed in [-1, 2^32-1]. Default **-1** = auto seed. + Same seed may yield similar (not guaranteed identical) results across runs. + + --cf / --camerafixed Lock camera movement. Some models support this flag. + true: try to keep camera fixed; false: allow movement. Default: **false**. + + Returns: + Dict: + API response containing task creation results for each input item. A typical shape is: + { + "status": "success", + "success_list": [{"video_name": "video_url"}], + "error_list": [] + } + + Constraints & Tips: + - Keep prompt concise and focused (建议 ≤ 500 字); too many details may distract the model. + - If using first/last frames, ensure their **aspect ratio matches** your chosen `--rt` to minimize cropping. + - If you must reproduce results, specify an explicit `--seed`. + - Unsupported parameters are ignored silently or may cause validation errors (model-specific). + + Minimal examples: + 1) Text-only batch of two 5-second clips at 720p, 16:9, 24 fps: + params = [ + { + "video_name": "cat_park.mp4", + "prompt": "小猫骑着滑板穿过公园。 --rs 720p --rt 16:9 --dur 5 --fps 24 --wm false" + }, + { + "video_name": "city_night.mp4", + "prompt": "霓虹灯下的城市延时摄影风。 --rs 720p --rt 16:9 --dur 5 --fps 24 --seed 7" + }, + ] + + 2) With guided first/last frame (square, 6 s, camera fixed): + params = [ + { + "video_name": "logo_reveal.mp4", + "first_frame": "https://cdn.example.com/brand/logo_start.png", + "last_frame": "https://cdn.example.com/brand/logo_end.png", + "prompt": "品牌 Logo 从线稿到上色的变化。 --rs 1080p --rt 1:1 --dur 6 --fps 24 --cf true" + } + ] """ batch_size = 10 success_list = [] error_list = [] - for start_idx in range(0, len(params), batch_size): - batch = params[start_idx : start_idx + batch_size] - task_dict = {} - for item in batch: - video_name = item["video_name"] - first_frame = item["first_frame"] - last_frame = item["last_frame"] - prompt = item["prompt"] - try: - if not first_frame: - response = await generate(tool_context, prompt) - elif not last_frame: - response = await generate(tool_context, prompt, first_frame) - else: - response = await generate( - tool_context, prompt, first_frame, last_frame - ) - task_dict[response.id] = video_name - except Exception: - traceback.print_exc() - while True: - task_list = list(task_dict.keys()) - if len(task_list) == 0: - break - for task_id in task_list: - result = client.content_generation.tasks.get(task_id=task_id) - status = result.status - if status == "succeeded": - logger.debug("----- task succeeded -----") - tool_context.state[f"{task_dict[task_id]}_video_url"] = ( - result.content.video_url - ) - success_list.append({task_dict[task_id]: result.content.video_url}) - task_dict.pop(task_id, None) - elif status == "failed": - logger.debug("----- task failed -----") - logger.debug(f"Error: {result.error}") - error_list.append(task_dict[task_id]) - task_dict.pop(task_id, None) - else: - logger.debug( - f"Current status: {status}, Retrying after 10 seconds..." - ) - time.sleep(10) - - if len(success_list) == 0: - return {"status": "error", "message": f"Following videos failed: {error_list}"} - else: - return { - "status": "success", - "message": f"Following videos generated: {success_list}\nFollowing videos failed: {error_list}", - } + tracer = trace.get_tracer("gcp.vertex.agent") + with tracer.start_as_current_span("call_llm") as span: + input_part = {"role": "user"} + output_part = {"message.role": "model"} + + for idx, item in enumerate(params): + input_part[f"parts.{idx}.type"] = "text" + input_part[f"parts.{idx}.text"] = json.dumps(item, ensure_ascii=False) + + for start_idx in range(0, len(params), batch_size): + batch = params[start_idx : start_idx + batch_size] + task_dict = {} + for idx, item in enumerate(batch): + video_name = item["video_name"] + prompt = item["prompt"] + first_frame = item.get("first_frame", None) + last_frame = item.get("last_frame", None) + try: + if not first_frame: + response = await generate(prompt) + elif not last_frame: + response = await generate(prompt, first_frame) + else: + response = await generate(prompt, first_frame, last_frame) + task_dict[response.id] = video_name + except Exception as e: + logger.error(f"Error: {e}") + error_list.append(video_name) + + total_tokens = 0 + while True: + task_list = list(task_dict.keys()) + if len(task_list) == 0: + break + for idx, task_id in enumerate(task_list): + result = client.content_generation.tasks.get(task_id=task_id) + status = result.status + if status == "succeeded": + logger.debug("----- task succeeded -----") + tool_context.state[f"{task_dict[task_id]}_video_url"] = ( + result.content.video_url + ) + total_tokens += result.usage.completion_tokens + output_part[f"message.parts.{idx}.type"] = "text" + output_part[f"message.parts.{idx}.text"] = ( + f"{task_dict[task_id]}: {result.content.video_url}" + ) + success_list.append( + {task_dict[task_id]: result.content.video_url} + ) + task_dict.pop(task_id, None) + elif status == "failed": + logger.error("----- task failed -----") + logger.error(f"Error: {result.error}") + error_list.append(task_dict[task_id]) + task_dict.pop(task_id, None) + else: + logger.debug( + f"Current status: {status}, Retrying after 10 seconds..." + ) + time.sleep(10) + + add_span_attributes( + span, + tool_context, + input_part=input_part, + output_part=output_part, + output_tokens=total_tokens, + total_tokens=total_tokens, + request_model=getenv("MODEL_VIDEO_NAME"), + response_model=getenv("MODEL_VIDEO_NAME"), + ) + + if len(success_list) == 0: + return { + "status": "error", + "success_list": success_list, + "error_list": error_list, + } + else: + return { + "status": "success", + "success_list": success_list, + "error_list": error_list, + } + + +def add_span_attributes( + span: Span, + tool_context: ToolContext, + input_part: dict = None, + output_part: dict = None, + input_tokens: int = None, + output_tokens: int = None, + total_tokens: int = None, + request_model: str = None, + response_model: str = None, +): + try: + # common attributes + app_name = tool_context._invocation_context.app_name + user_id = tool_context._invocation_context.user_id + agent_name = tool_context.agent_name + session_id = tool_context._invocation_context.session.id + span.set_attribute("gen_ai.agent.name", agent_name) + span.set_attribute("openinference.instrumentation.veadk", VERSION) + span.set_attribute("gen_ai.app.name", app_name) + span.set_attribute("gen_ai.user.id", user_id) + span.set_attribute("gen_ai.session.id", session_id) + span.set_attribute("agent_name", agent_name) + span.set_attribute("agent.name", agent_name) + span.set_attribute("app_name", app_name) + span.set_attribute("app.name", app_name) + span.set_attribute("user.id", user_id) + span.set_attribute("session.id", session_id) + span.set_attribute("cozeloop.report.source", "veadk") + + # llm attributes + span.set_attribute("gen_ai.system", "openai") + span.set_attribute("gen_ai.operation.name", "chat") + if request_model: + span.set_attribute("gen_ai.request.model", request_model) + if response_model: + span.set_attribute("gen_ai.response.model", response_model) + if total_tokens: + span.set_attribute("gen_ai.usage.total_tokens", total_tokens) + if output_tokens: + span.set_attribute("gen_ai.usage.output_tokens", output_tokens) + if input_tokens: + span.set_attribute("gen_ai.usage.input_tokens", input_tokens) + if input_part: + span.add_event("gen_ai.user.message", input_part) + if output_part: + span.add_event("gen_ai.choice", output_part) + + except Exception: + traceback.print_exc()