save work

sdk/ai/azure-ai-agents/azure/ai/agents/telemetry/_ai_agents_instrumentor.py

@@ -25,6 +25,7 @@
 from azure.ai.agents.models import (
     MessageAttachment,
     MessageDeltaChunk,
+    MessageContent,
     MessageIncompleteDetails,
     MessageInputContentBlock,
     MessageInputTextBlock,
@@ -375,7 +376,8 @@ def add_thread_message_event(
         message: Union[ThreadMessage, ThreadMessageOptions],
         usage: Optional[_models.RunStepCompletionUsage] = None,
     ) -> None:
-        content_body = {}
+
+        content_body: Optional[Union[str, Dict[str, Any]]] = None
         if _trace_agents_content:
             if isinstance(message, ThreadMessage):
                 for content in message.content:
@@ -385,14 +387,15 @@ def add_thread_message_event(
                         annotations = self._get_field(typed_content, "annotations")
                         if annotations:
                             content_details["annotations"] = [a.as_dict() for a in annotations]
+                        content_body = {}
                         content_body[content.type] = content_details
             elif isinstance(message, ThreadMessageOptions):
                 if isinstance(message.content, str):
-                    content_body = {"text": {"value": message.content}}
+                    content_body = message.content
                 elif isinstance(message.content, List):
-                    for content2 in message.content:
-                        if isinstance(content2, MessageInputTextBlock):
-                            content_body[content2.type] = {"value": content2.text}
+                    for block in message.content:
+                        if isinstance(block, MessageInputTextBlock):
+                            content_body = block.text
                             break
 
         self._add_message_event(

sdk/ai/azure-ai-agents/tests/test_ai_agents_instrumentor.py

@@ -5,7 +5,6 @@
 # ------------------------------------
 # cSpell:disable# cSpell:disable
 import pytest
-from enum import Enum, IntEnum
 import os
 import json
 import jsonref
@@ -19,7 +18,6 @@
     McpTool,
     MessageDeltaChunk,
     MessageDeltaTextContent,
-    MessageInputContentBlock,
     MessageInputTextBlock,
     OpenApiAnonymousAuthDetails,
     OpenApiTool,
@@ -52,7 +50,7 @@
 )
 
 from test_agents_client_base import agentClientPreparer
-from test_ai_instrumentor_base import TestAiAgentsInstrumentorBase
+from test_ai_instrumentor_base import TestAiAgentsInstrumentorBase, MessageCreationMode
 
 CONTENT_TRACING_ENV_VARIABLE = "OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT"
 settings.tracing_implementation = "OpenTelemetry"
@@ -228,33 +226,28 @@ def set_env_var(var_name, value):
             self.cleanup()  # This also undefines CONTENT_TRACING_ENV_VARIABLE
             os.environ.pop(OLD_CONTENT_TRACING_ENV_VARIABLE, None)
 
-    class MessageCreationMode(IntEnum):
-        MESSAGE_CREATE_STR = 1 # Test calls `client.messages.create(content="...", ...)` to create the messages in a dedicated call
-        MESSAGE_CREATE_INPUT_TEXT_BLOCK = 2 # Test calls `client.messages.create(content=[MessageInputTextBlock(...)], ...)` to create the message in a dedicated call
-        THREAD_CREATE_STR = 3 # Test calls `client.threads.create(messages=[ThreadMessageOptions(...)])`.
-        THREAD_CREATE_INPUT_TEXT_BLOCK = 4 # Test calls `client.threads.create(messages=[ThreadMessageOptions(content=[MessageInputTextBlock(...)])])`.
-
     @pytest.mark.usefixtures("instrument_with_content")
     @agentClientPreparer()
     @recorded_by_proxy
     @pytest.mark.parametrize("message_creation_mode", [m.value for m in MessageCreationMode])
-    def test_agent_chat_with_tracing_content_recording_enabled(self, message_creation_mode: MessageCreationMode = MessageCreationMode.THREAD_CREATE_INPUT_TEXT_BLOCK, **kwargs):
+    def test_agent_chat_with_tracing_content_recording_enabled(self, message_creation_mode: MessageCreationMode, **kwargs):
         assert True == AIAgentsInstrumentor().is_content_recording_enabled()
         assert True == AIAgentsInstrumentor().is_instrumented()
 
         client = self.create_client(**kwargs)
         agent = client.create_agent(model="gpt-4o-mini", name="my-agent", instructions="You are helpful agent")
         user_content = "Hello, tell me a joke"
 
-        if message_creation_mode == self.MessageCreationMode.MESSAGE_CREATE_STR:
+        # Test 4 different patterns of thread & message creation
+        if message_creation_mode == MessageCreationMode.MESSAGE_CREATE_STR:
             thread = client.threads.create()
             client.messages.create(thread_id=thread.id, role="user", content=user_content)
-        elif message_creation_mode == self.MessageCreationMode.MESSAGE_CREATE_INPUT_TEXT_BLOCK:
+        elif message_creation_mode == MessageCreationMode.MESSAGE_CREATE_INPUT_TEXT_BLOCK:
             thread = client.threads.create()
             client.messages.create(thread_id=thread.id, role="user", content=[MessageInputTextBlock(text=user_content)])
-        elif message_creation_mode == self.MessageCreationMode.THREAD_CREATE_STR:
+        elif message_creation_mode == MessageCreationMode.THREAD_CREATE_STR:
             thread = client.threads.create(messages=[ThreadMessageOptions(role="user", content=user_content)])
-        elif message_creation_mode == self.MessageCreationMode.THREAD_CREATE_INPUT_TEXT_BLOCK:
+        elif message_creation_mode == MessageCreationMode.THREAD_CREATE_INPUT_TEXT_BLOCK:
             thread = client.threads.create(messages=[ThreadMessageOptions(role="user", content=[MessageInputTextBlock(text=user_content)])])
         else:
             assert False, f"Unknown message creation mode: {message_creation_mode}"
@@ -316,21 +309,21 @@ def test_agent_chat_with_tracing_content_recording_enabled(self, message_creatio
         attributes_match = GenAiTraceVerifier().check_span_attributes(span, expected_attributes)
         assert attributes_match == True
 
-        if message_creation_mode in (self.MessageCreationMode.THREAD_CREATE_STR, self.MessageCreationMode.THREAD_CREATE_INPUT_TEXT_BLOCK):
+        if message_creation_mode in (MessageCreationMode.THREAD_CREATE_STR, MessageCreationMode.THREAD_CREATE_INPUT_TEXT_BLOCK):
             expected_events = [
                 {
                     "name": "gen_ai.user.message",
                     "attributes": {
                         "gen_ai.system": "az.ai.agents",
-                        "gen_ai.event.content": '{"content": {"text": {"value": "Hello, tell me a joke"}}, "role": "user"}',
+                        "gen_ai.event.content": '{"content": "Hello, tell me a joke", "role": "user"}',
                     },
                 }
             ]
             events_match = GenAiTraceVerifier().check_span_events(span, expected_events)
             assert events_match == True
 
         # ------------------------- Validate "create_message" span ---------------------------------
-        if message_creation_mode in (self.MessageCreationMode.MESSAGE_CREATE_STR, self.MessageCreationMode.MESSAGE_CREATE_INPUT_TEXT_BLOCK):
+        if message_creation_mode in (MessageCreationMode.MESSAGE_CREATE_STR, MessageCreationMode.MESSAGE_CREATE_INPUT_TEXT_BLOCK):
             spans = self.exporter.get_spans_by_name("create_message")
             assert len(spans) == 1
             span = spans[0]

sdk/ai/azure-ai-agents/tests/test_ai_instrumentor_base.py

@@ -2,6 +2,7 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
 # ------------------------------------
+from enum import IntEnum
 from typing import Any, Dict, List, Optional
 import json
 import os
@@ -21,6 +22,14 @@
 
 CONTENT_TRACING_ENV_VARIABLE = "OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT"
 
+# We test different ways to create a thread & message, to cover all tracing code paths
+class MessageCreationMode(IntEnum):
+    MESSAGE_CREATE_STR = 1 # Test calls `client.messages.create(content="...", ...)` to create the messages in a dedicated call
+    MESSAGE_CREATE_INPUT_TEXT_BLOCK = 2 # Test calls `client.messages.create(content=[MessageInputTextBlock(...)], ...)` to create the message in a dedicated call
+    THREAD_CREATE_STR = 3 # Test calls `client.threads.create(messages=[ThreadMessageOptions(...)])`.
+    THREAD_CREATE_INPUT_TEXT_BLOCK = 4 # Test calls `client.threads.create(messages=[ThreadMessageOptions(content=[MessageInputTextBlock(...)])])`.
+
+
 class TestAiAgentsInstrumentorBase(TestAgentClientBase):
     """The utility methods, used by AI Instrumentor test."""