Fix tests for dashscope

Change-Id: Ie958a04905738463e0d913816251142c1cf03d06
Co-developed-by: Cursor <[email protected]>

Author: Cirilla-zmh

instrumentation-loongsuite/loongsuite-instrumentation-dashscope/src/loongsuite/instrumentation/dashscope/patch.py

@@ -502,23 +502,11 @@ def wrap_image_synthesis_async_call(
             # Execute the wrapped call (submit task)
             result = wrapped(*args, **kwargs)
 
-            # Extract task_id and update invocation
-            task_id = None
-            if result and hasattr(result, "output"):
-                if hasattr(result.output, "get"):
-                    task_id = result.output.get("task_id")
-                elif hasattr(result.output, "task_id"):
-                    task_id = getattr(result.output, "task_id", None)
-
-            if task_id:
-                invocation.attributes["gen_ai.response.id"] = task_id
-                invocation.attributes["dashscope.task_id"] = task_id
-
-            # Note: Span linking is not currently supported by ExtendedTelemetryHandler.
-            # If span linking is needed in the future, it should be implemented in the handler.
-            # For now, we skip storing span context for linking.
-
-            # Update invocation with async response data (task_id, task_status)
+            # TODO: Implement span linking between async_call and wait spans.
+            # This would allow wait() span to be linked to async_call() span for better trace visualization.
+            # Consider using a task_id-based mapping to store span contexts across threads.
+
+            # Update invocation with async response data (request_id, task_id)
             _update_invocation_from_image_synthesis_async_response(
                 invocation, result
             )
@@ -566,19 +554,19 @@ def wrap_image_synthesis_wait(wrapped, instance, args, kwargs, handler=None):
             # If cannot extract task_id, skip instrumentation
             return wrapped(*args, **kwargs)
 
-        # Note: Span linking is not currently supported by ExtendedTelemetryHandler.
-        # If span linking is needed in the future, it should be implemented in the handler.
+        # TODO: Implement span linking between async_call and wait spans.
+        # This would allow wait() span to be linked to async_call() span for better trace visualization.
+        # Consider using a task_id-based mapping to store span contexts across threads.
 
         # Create invocation object (wait phase doesn't know model, use "unknown")
+        # Use "wait generate_content" as operation_name to make span name clearer
         invocation = _create_invocation_from_image_synthesis({}, "unknown")
-        invocation.operation_name = "generate_content"
+        # TODO: Add semantic conventions for wait operations
+        invocation.operation_name = "wait generate_content"
         invocation.attributes["gen_ai.request.async"] = True
-        invocation.attributes["gen_ai.response.id"] = task_id
-        invocation.attributes["dashscope.task_id"] = task_id
-        invocation.attributes["dashscope.operation"] = "wait"
-
-        # Note: Span linking is not currently supported by ExtendedTelemetryHandler.
-        # If span linking is needed in the future, it should be implemented in the handler.
+        # Note: response_id will be set from response.output.task_id in _update_invocation_from_image_synthesis_response
+        # We set task_id here as a fallback
+        invocation.response_id = task_id
 
         # Start LLM invocation (creates span)
         handler.start_llm(invocation)

instrumentation-loongsuite/loongsuite-instrumentation-dashscope/src/loongsuite/instrumentation/dashscope/utils.py

@@ -26,6 +26,7 @@
     ToolCall,
     ToolCallResponse,
     ToolDefinition,
+    Uri,
 )
 
 
@@ -738,29 +739,6 @@ def _create_invocation_from_image_synthesis(
                     InputMessage(role="user", parts=parts)
                 ]
 
-    # Extract negative_prompt (as attribute)
-    negative_prompt = kwargs.get("negative_prompt")
-    if negative_prompt:
-        if isinstance(negative_prompt, str):
-            invocation.attributes["dashscope.negative_prompt"] = (
-                negative_prompt
-            )
-
-    # Extract size (image dimensions)
-    size = kwargs.get("size")
-    if size:
-        invocation.attributes["dashscope.image.size"] = size
-
-    # Extract n (number of images to generate)
-    n = kwargs.get("n")
-    if n is not None:
-        invocation.attributes["dashscope.image.n"] = n
-
-    # Extract similarity parameter (if available)
-    similarity = kwargs.get("similarity")
-    if similarity is not None:
-        invocation.attributes["dashscope.image.similarity"] = similarity
-
     return invocation
 
 
@@ -790,16 +768,8 @@ def _update_invocation_from_image_synthesis_response(
         except (KeyError, AttributeError):
             pass
 
-        # Extract request ID (if available)
-        # Note: For ImageSynthesis, request_id is the main identifier, not task_id
-        try:
-            request_id = getattr(response, "request_id", None)
-            if request_id:
-                invocation.response_id = request_id
-        except (KeyError, AttributeError):
-            pass
-
-        # Extract task_id and task_status from output
+        # Extract task_id from output and set as response_id
+        # Note: For ImageSynthesis, response_id should be task_id, not request_id
         try:
             output = getattr(response, "output", None)
             if output:
@@ -811,54 +781,56 @@ def _update_invocation_from_image_synthesis_response(
                     task_id = getattr(output, "task_id", None)
 
                 if task_id:
-                    # Store task_id in attributes
-                    # Note: gen_ai.response.id should be request_id, not task_id
-                    # task_id is stored separately in dashscope.task_id
-                    invocation.attributes["dashscope.task_id"] = task_id
-                    # Don't set gen_ai.response.id to task_id, as it should be request_id
-                    # Only set response_id to task_id if request_id is not available
-                    if not invocation.response_id:
-                        invocation.response_id = task_id
-                        invocation.attributes["gen_ai.response.id"] = task_id
-
-                # Extract task_status
-                task_status = None
-                if hasattr(output, "get"):
-                    task_status = output.get("task_status")
-                elif hasattr(output, "task_status"):
-                    task_status = getattr(output, "task_status", None)
-
-                if task_status:
-                    invocation.attributes["dashscope.task_status"] = (
-                        task_status
-                    )
+                    invocation.response_id = task_id
 
-                # Extract image URLs from results
-                # TODO: If returned as files or binary data, handle accordingly
+                # Extract image URLs from results and add as Uri MessageParts
                 results = None
                 if hasattr(output, "get"):
                     results = output.get("results")
                 elif hasattr(output, "results"):
                     results = getattr(output, "results", None)
 
                 if results and isinstance(results, list):
-                    image_urls = []
+                    image_uris = []
                     for result in results:
                         if isinstance(result, dict):
                             url = result.get("url")
                             if url:
-                                image_urls.append(url)
+                                image_uris.append(
+                                    Uri(
+                                        uri=url,
+                                        modality="image",
+                                        mime_type=None,
+                                        type="uri",
+                                    )
+                                )
                         elif hasattr(result, "url"):
                             url = getattr(result, "url", None)
                             if url:
-                                image_urls.append(url)
-                    if image_urls:
-                        # Store first image URL as attribute (or all if needed)
-                        invocation.attributes["dashscope.image.url"] = (
-                            image_urls[0]
-                            if len(image_urls) == 1
-                            else str(image_urls)
-                        )
+                                image_uris.append(
+                                    Uri(
+                                        uri=url,
+                                        modality="image",
+                                        mime_type=None,
+                                        type="uri",
+                                    )
+                                )
+                    if image_uris:
+                        # Add image URIs to output messages
+                        # If output_messages is empty, create a new one
+                        if not invocation.output_messages:
+                            invocation.output_messages = [
+                                OutputMessage(
+                                    role="assistant",
+                                    parts=image_uris,
+                                    finish_reason="stop",
+                                )
+                            ]
+                        else:
+                            # Append URIs to the last output message
+                            invocation.output_messages[-1].parts.extend(
+                                image_uris
+                            )
         except (KeyError, AttributeError):
             pass
     except (KeyError, AttributeError):
@@ -872,7 +844,7 @@ def _update_invocation_from_image_synthesis_async_response(
     """Update LLMInvocation with ImageSynthesis async_call response data.
 
     This is called when async_call() returns, before wait() is called.
-    Only extracts task_id and task_status (usually PENDING).
+    Extracts task_id and sets it as response_id.
 
     Args:
         invocation: LLMInvocation to update
@@ -882,28 +854,16 @@ def _update_invocation_from_image_synthesis_async_response(
         return
 
     try:
-        # Extract task_id and task_status from output
+        # Extract task_id from output and set as response_id
         output = getattr(response, "output", None)
         if output:
-            # Extract task_id
             task_id = None
             if hasattr(output, "get"):
                 task_id = output.get("task_id")
             elif hasattr(output, "task_id"):
                 task_id = getattr(output, "task_id", None)
 
             if task_id:
-                invocation.attributes["gen_ai.response.id"] = task_id
-                invocation.attributes["dashscope.task_id"] = task_id
-
-            # Extract task_status (usually PENDING for async_call)
-            task_status = None
-            if hasattr(output, "get"):
-                task_status = output.get("task_status")
-            elif hasattr(output, "task_status"):
-                task_status = getattr(output, "task_status", None)
-
-            if task_status:
-                invocation.attributes["dashscope.task_status"] = task_status
+                invocation.response_id = task_id
     except (KeyError, AttributeError):
         pass