Python: Input & output attributes in invoke_agent span (#12834)

### Motivation and Context

<!-- Thank you for your contribution to the semantic-kernel repo!
Please help reviewers and future users, providing the following
information:
  1. Why is this change required?
  2. What problem does it solve?
  3. What scenario does it contribute to?
  4. If it fixes an open issue, please link to the issue here.
-->
The OTel workgroup has been thinking about enhancing tracing for agents.
This is a proposal from our tracing team:
https://github.com/open-telemetry/semantic-conventions/pull/2528. This
PR implements part of the proposal so that when the proposal gets
merged, we can quickly follow.

### Description

<!-- Describe your changes, the overall approach, the underlying design.
These notes will help understanding how your code works. Thanks! -->
Two new attributes are added to the `invoke_agent` span, namely the
`gen_ai.agent.invocation_input` and `gen_ai.agent.invocation_output`.

<img width="747" height="583" alt="image"
src="https://github.com/user-attachments/assets/d8b81f68-915b-48ee-8f86-66a0d4f686cc"
/>


### Contribution Checklist

<!-- Before submitting this PR, please make sure: -->

- [x] The code builds clean without any errors or warnings
- [x] The PR follows the [SK Contribution
Guidelines](https://github.com/microsoft/semantic-kernel/blob/main/CONTRIBUTING.md)
and the [pre-submission formatting
script](https://github.com/microsoft/semantic-kernel/blob/main/CONTRIBUTING.md#development-scripts)
raises no violations
- [x] All unit tests pass, and I have added new tests where possible
- [x] I didn't break anyone 😄

---------

Co-authored-by: Evan Mattson <[email protected]>

python/samples/getting_started_with_agents/multi_agent_orchestration/observability.py

@@ -0,0 +1,94 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import logging
+import os
+
+from azure.monitor.opentelemetry.exporter import AzureMonitorLogExporter, AzureMonitorTraceExporter
+from dotenv import load_dotenv
+from opentelemetry import trace
+from opentelemetry._logs import set_logger_provider
+from opentelemetry.sdk._logs import LoggerProvider, LoggingHandler
+from opentelemetry.sdk._logs.export import BatchLogRecordProcessor
+from opentelemetry.sdk.resources import Resource
+from opentelemetry.sdk.trace import TracerProvider
+from opentelemetry.sdk.trace.export import BatchSpanProcessor
+from opentelemetry.semconv.resource import ResourceAttributes
+from opentelemetry.trace import set_tracer_provider
+from opentelemetry.trace.span import format_trace_id
+
+load_dotenv()
+
+APPINSIGHTS_CONNECTION_STRING = os.getenv("APPINSIGHTS_CONNECTION_STRING")
+
+resource = Resource.create({ResourceAttributes.SERVICE_NAME: "multi_agent_orchestration_sample"})
+
+
+def set_up_logging():
+    class KernelFilter(logging.Filter):
+        """A filter to not process records from semantic_kernel."""
+
+        # These are the namespaces that we want to exclude from logging for the purposes of this demo.
+        namespaces_to_exclude: list[str] = [
+            "semantic_kernel.functions.kernel_plugin",
+            "semantic_kernel.prompt_template.kernel_prompt_template",
+        ]
+
+        def filter(self, record):
+            return not any([record.name.startswith(namespace) for namespace in self.namespaces_to_exclude])
+
+    exporters = []
+    exporters.append(AzureMonitorLogExporter(connection_string=APPINSIGHTS_CONNECTION_STRING))
+
+    # Create and set a global logger provider for the application.
+    logger_provider = LoggerProvider(resource=resource)
+    # Log processors are initialized with an exporter which is responsible
+    # for sending the telemetry data to a particular backend.
+    for log_exporter in exporters:
+        logger_provider.add_log_record_processor(BatchLogRecordProcessor(log_exporter))
+    # Sets the global default logger provider
+    set_logger_provider(logger_provider)
+
+    # Create a logging handler to write logging records, in OTLP format, to the exporter.
+    handler = LoggingHandler()
+    # Add filters to the handler to only process records from semantic_kernel.
+    handler.addFilter(logging.Filter("semantic_kernel"))
+    handler.addFilter(KernelFilter())
+    # Attach the handler to the root logger. `getLogger()` with no arguments returns the root logger.
+    # Events from all child loggers will be processed by this handler.
+    logger = logging.getLogger()
+    logger.addHandler(handler)
+    # Set the logging level to NOTSET to allow all records to be processed by the handler.
+    logger.setLevel(logging.NOTSET)
+
+
+def set_up_tracing():
+    exporters = []
+    exporters.append(AzureMonitorTraceExporter(connection_string=APPINSIGHTS_CONNECTION_STRING))
+
+    # Initialize a trace provider for the application. This is a factory for creating tracers.
+    tracer_provider = TracerProvider(resource=resource)
+    # Span processors are initialized with an exporter which is responsible
+    # for sending the telemetry data to a particular backend.
+    for exporter in exporters:
+        tracer_provider.add_span_processor(BatchSpanProcessor(exporter))
+    # Sets the global default tracer provider
+    set_tracer_provider(tracer_provider)
+
+
+def enable_observability(func):
+    """A decorator to enable observability for the samples."""
+
+    async def wrapper(*args, **kwargs):
+        if not APPINSIGHTS_CONNECTION_STRING:
+            # If the connection string is not set, skip observability setup.
+            return await func(*args, **kwargs)
+
+        set_up_logging()
+        set_up_tracing()
+
+        tracer = trace.get_tracer(__name__)
+        with tracer.start_as_current_span("main") as current_span:
+            print(f"Trace ID: {format_trace_id(current_span.get_span_context().trace_id)}")
+            return await func(*args, **kwargs)
+
+    return wrapper

python/samples/getting_started_with_agents/multi_agent_orchestration/step4c_handoff_mix_agent_types.py

@@ -5,6 +5,7 @@ import asyncio
 from azure.ai.projects.aio import AIProjectClient
 from azure.identity.aio import DefaultAzureCredential
 
+from samples.getting_started_with_agents.multi_agent_orchestration.observability import enable_observability
 from semantic_kernel.agents import (
     Agent,
     AzureAIAgent,
@@ -180,6 +181,7 @@ def human_response_function() -> ChatMessageContent:
     return ChatMessageContent(role=AuthorRole.USER, content=user_input)
 
 
+@enable_observability
 async def main():
     """Main function to run the agents."""
     # 0. Initialize the Azure AI agent clients

python/semantic_kernel/agents/azure_ai/azure_ai_agent.py

@@ -58,6 +58,7 @@ from semantic_kernel.utils.naming import generate_random_ascii_name
 from semantic_kernel.utils.telemetry.agent_diagnostics.decorators import (
     trace_agent_get_response,
     trace_agent_invocation,
+    trace_agent_streaming_invocation,
 )
 from semantic_kernel.utils.telemetry.user_agent import APP_INFO, SEMANTIC_KERNEL_USER_AGENT
 
@@ -832,7 +833,7 @@ class AzureAIAgent(DeclarativeSpecMixin, Agent):
                 # Emit tool-related messages only via callback
                 await on_intermediate_message(message)
 
-    @trace_agent_invocation
+    @trace_agent_streaming_invocation
     @override
     async def invoke_stream(
         self,

python/semantic_kernel/agents/bedrock/bedrock_agent.py

@@ -44,6 +44,7 @@ from semantic_kernel.utils.feature_stage_decorator import experimental
 from semantic_kernel.utils.telemetry.agent_diagnostics.decorators import (
     trace_agent_get_response,
     trace_agent_invocation,
+    trace_agent_streaming_invocation,
 )
 
 logger = logging.getLogger(__name__)
@@ -458,7 +459,7 @@ class BedrockAgent(BedrockAgentBase):
             "Failed to get a response from the agent. Please consider increasing the auto invoke attempts."
         )
 
-    @trace_agent_invocation
+    @trace_agent_streaming_invocation
     @override
     async def invoke_stream(
         self,

python/semantic_kernel/agents/chat_completion/chat_completion_agent.py

@@ -34,6 +34,7 @@ from semantic_kernel.prompt_template.prompt_template_config import PromptTemplat
 from semantic_kernel.utils.telemetry.agent_diagnostics.decorators import (
     trace_agent_get_response,
     trace_agent_invocation,
+    trace_agent_streaming_invocation,
 )
 
 if TYPE_CHECKING:
@@ -370,7 +371,7 @@ class ChatCompletionAgent(DeclarativeSpecMixin, Agent):
         ):
             yield AgentResponseItem(message=response, thread=thread)
 
-    @trace_agent_invocation
+    @trace_agent_streaming_invocation
     @override
     async def invoke_stream(
         self,

python/semantic_kernel/agents/copilot_studio/copilot_studio_agent.py

@@ -7,21 +7,10 @@ from os import environ, path
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, ClassVar, Literal
 
-from microsoft.agents.copilotstudio.client import (
-    AgentType,
-    CopilotClient,
-    PowerPlatformCloud,
-)
+from microsoft.agents.copilotstudio.client import AgentType, CopilotClient, PowerPlatformCloud
 from microsoft.agents.core.models import ActivityTypes
-from msal import (
-    ConfidentialClientApplication,
-    PublicClientApplication,
-)
-from msal_extensions import (
-    FilePersistence,
-    PersistedTokenCache,
-    build_encrypted_persistence,
-)
+from msal import ConfidentialClientApplication, PublicClientApplication
+from msal_extensions import FilePersistence, PersistedTokenCache, build_encrypted_persistence
 from pydantic import ValidationError
 
 from semantic_kernel.agents import Agent
@@ -47,6 +36,7 @@ from semantic_kernel.utils.naming import generate_random_ascii_name
 from semantic_kernel.utils.telemetry.agent_diagnostics.decorators import (
     trace_agent_get_response,
     trace_agent_invocation,
+    trace_agent_streaming_invocation,
 )
 
 if sys.version_info >= (3, 12):
@@ -509,6 +499,7 @@ class CopilotStudioAgent(Agent):
         ):
             yield AgentResponseItem(message=response, thread=thread)
 
+    @trace_agent_streaming_invocation
     @override
     async def invoke_stream(
         self,

python/semantic_kernel/agents/open_ai/openai_assistant_agent.py

@@ -54,6 +54,7 @@ from semantic_kernel.utils.naming import generate_random_ascii_name
 from semantic_kernel.utils.telemetry.agent_diagnostics.decorators import (
     trace_agent_get_response,
     trace_agent_invocation,
+    trace_agent_streaming_invocation,
 )
 from semantic_kernel.utils.telemetry.user_agent import APP_INFO, prepend_semantic_kernel_to_user_agent
 
@@ -947,7 +948,7 @@ class OpenAIAssistantAgent(DeclarativeSpecMixin, Agent):
                 # Emit tool-related messages only via callback
                 await on_intermediate_message(message)
 
-    @trace_agent_invocation
+    @trace_agent_streaming_invocation
     @override
     async def invoke_stream(
         self,

python/semantic_kernel/agents/open_ai/openai_responses_agent.py

@@ -46,6 +46,7 @@ from semantic_kernel.utils.naming import generate_random_ascii_name
 from semantic_kernel.utils.telemetry.agent_diagnostics.decorators import (
     trace_agent_get_response,
     trace_agent_invocation,
+    trace_agent_streaming_invocation,
 )
 from semantic_kernel.utils.telemetry.user_agent import APP_INFO, prepend_semantic_kernel_to_user_agent
 
@@ -1040,7 +1041,7 @@ class OpenAIResponsesAgent(DeclarativeSpecMixin, Agent):
                 # Emit tool-related messages only via callback
                 await on_intermediate_message(message)
 
-    @trace_agent_invocation
+    @trace_agent_streaming_invocation
     @override
     async def invoke_stream(
         self,

python/semantic_kernel/utils/telemetry/agent_diagnostics/decorators.py

@@ -1,74 +1,240 @@
 # Copyright (c) Microsoft. All rights reserved.
 
 import functools
+import json
 from collections.abc import AsyncIterable, Awaitable, Callable
-from typing import ParamSpec, TypeVar, cast
+from functools import reduce
+from typing import ParamSpec, cast
 
-from opentelemetry.trace import get_tracer
+from opentelemetry.trace import Span, StatusCode, get_tracer
 
+from semantic_kernel.agents.agent import Agent, AgentResponseItem
+from semantic_kernel.contents.chat_message_content import ChatMessageContent
+from semantic_kernel.contents.streaming_chat_message_content import StreamingChatMessageContent
+from semantic_kernel.contents.utils.author_role import AuthorRole
 from semantic_kernel.utils.feature_stage_decorator import experimental
 from semantic_kernel.utils.telemetry.agent_diagnostics import gen_ai_attributes
+from semantic_kernel.utils.telemetry.model_diagnostics.model_diagnostics_settings import ModelDiagnosticSettings
+
+# Module to instrument GenAI agents using OpenTelemetry and OpenTelemetry Semantic Conventions.
+# These are experimental features and may change in the future.
+
+# To enable these features, set one of the following environment variables to true:
+#    SEMANTICKERNEL_EXPERIMENTAL_GENAI_ENABLE_OTEL_DIAGNOSTICS
+#    SEMANTICKERNEL_EXPERIMENTAL_GENAI_ENABLE_OTEL_DIAGNOSTICS_SENSITIVE
+# We are re-using the model diagnostic settings to control the instrumentation of agents
+# because it makes sense to have a system wide setting for diagnostics. The name "model"
+# is a legacy name because the concept of agent was not yet introduced when the settings were created.
+MODEL_DIAGNOSTICS_SETTINGS = ModelDiagnosticSettings()
 
 P = ParamSpec("P")
-T = TypeVar("T")
 
 # Creates a tracer from the global tracer provider
 tracer = get_tracer(__name__)
 
-
-@experimental
-def trace_agent_invocation(invoke_func: Callable[P, AsyncIterable[T]]) -> Callable[P, AsyncIterable[T]]:
-    """Decorator to trace agent invocation."""
-    OPERATION_NAME = "invoke_agent"
-
-    @functools.wraps(invoke_func)
-    async def wrapper_decorator(*args: P.args, **kwargs: P.kwargs) -> AsyncIterable[T]:
-        from semantic_kernel.agents.agent import Agent
-
-        agent = cast(Agent, args[0])
-        with tracer.start_as_current_span(f"{OPERATION_NAME} {agent.name}") as span:
-            span.set_attributes({
-                gen_ai_attributes.OPERATION: OPERATION_NAME,
-                gen_ai_attributes.AGENT_ID: agent.id,
-                gen_ai_attributes.AGENT_NAME: agent.name,
-            })
-
-            if agent.description:
-                span.set_attribute(gen_ai_attributes.AGENT_DESCRIPTION, agent.description)
-
-            async for response in invoke_func(*args, **kwargs):
-                yield response
-
-    # Mark the wrapper decorator as an agent diagnostics decorator
-    wrapper_decorator.__agent_diagnostics__ = True  # type: ignore
-
-    return wrapper_decorator
+OPERATION_NAME = "invoke_agent"
 
 
 @experimental
-def trace_agent_get_response(get_response_func: Callable[P, Awaitable[T]]) -> Callable[P, Awaitable[T]]:
+def are_model_diagnostics_enabled() -> bool:
+    """Check if model diagnostics are enabled.
+
+    Model diagnostics are enabled if either diagnostic is enabled or diagnostic with sensitive events is enabled.
+    """
+    return (
+        MODEL_DIAGNOSTICS_SETTINGS.enable_otel_diagnostics
+        or MODEL_DIAGNOSTICS_SETTINGS.enable_otel_diagnostics_sensitive
+    )
+
+
+@experimental
+def are_sensitive_events_enabled() -> bool:
+    """Check if sensitive events are enabled.
+
+    Sensitive events are enabled if the diagnostic with sensitive events is enabled.
+    """
+    return MODEL_DIAGNOSTICS_SETTINGS.enable_otel_diagnostics_sensitive
+
+
+@experimental
+def trace_agent_get_response(
+    get_response_func: Callable[P, Awaitable[AgentResponseItem[ChatMessageContent]]],
+) -> Callable[P, Awaitable[AgentResponseItem[ChatMessageContent]]]:
     """Decorator to trace agent invocation."""
-    OPERATION_NAME = "invoke_agent"
 
     @functools.wraps(get_response_func)
-    async def wrapper_decorator(*args: P.args, **kwargs: P.kwargs) -> T:
-        from semantic_kernel.agents.agent import Agent
+    async def wrapper_decorator(
+        *args: P.args,
+        **kwargs: P.kwargs,
+    ) -> AgentResponseItem[ChatMessageContent]:
+        if not are_model_diagnostics_enabled():
+            # If model diagnostics are not enabled, just return the responses
+            return await get_response_func(*args, **kwargs)
 
         agent = cast(Agent, args[0])
+        messages = args[1] if len(args) > 1 else None
 
-        with tracer.start_as_current_span(f"{OPERATION_NAME} {agent.name}") as span:
-            span.set_attributes({
-                gen_ai_attributes.OPERATION: OPERATION_NAME,
-                gen_ai_attributes.AGENT_ID: agent.id,
-                gen_ai_attributes.AGENT_NAME: agent.name,
-            })
-
-            if agent.description:
-                span.set_attribute(gen_ai_attributes.AGENT_DESCRIPTION, agent.description)
-
-            return await get_response_func(*args, **kwargs)
+        with _start_as_current_span(agent) as span:
+            try:
+                _set_agent_invocation_input(span, messages)  # type: ignore
+                response = await get_response_func(*args, **kwargs)
+                _set_agent_invocation_output(span, [response.message])
+                return response
+            except Exception as e:
+                _set_agent_invocation_error(span, e)
+                raise
 
     # Mark the wrapper decorator as an agent diagnostics decorator
     wrapper_decorator.__agent_diagnostics__ = True  # type: ignore
 
     return wrapper_decorator
+
+
+@experimental
+def trace_agent_invocation(
+    invoke_func: Callable[P, AsyncIterable[AgentResponseItem[ChatMessageContent]]],
+) -> Callable[P, AsyncIterable[AgentResponseItem[ChatMessageContent]]]:
+    """Decorator to trace agent invocation."""
+
+    @functools.wraps(invoke_func)
+    async def wrapper_decorator(
+        *args: P.args,
+        **kwargs: P.kwargs,
+    ) -> AsyncIterable[AgentResponseItem[ChatMessageContent]]:
+        if not are_model_diagnostics_enabled():
+            # If model diagnostics are not enabled, just return the responses
+            async for response in invoke_func(*args, **kwargs):
+                yield response
+            return
+
+        agent = cast(Agent, args[0])
+        messages = args[1] if len(args) > 1 else None
+
+        with _start_as_current_span(agent) as current_span:
+            _set_agent_invocation_input(current_span, messages)  # type: ignore
+            try:
+                responses: list[ChatMessageContent] = []
+                async for response in invoke_func(*args, **kwargs):
+                    responses.append(response.message)
+                    yield response
+                _set_agent_invocation_output(current_span, responses)
+            except Exception as e:
+                _set_agent_invocation_error(current_span, e)
+                raise
+
+    # Mark the wrapper decorator as an agent diagnostics decorator
+    wrapper_decorator.__agent_diagnostics__ = True  # type: ignore
+
+    return wrapper_decorator
+
+
+@experimental
+def trace_agent_streaming_invocation(
+    invoke_func: Callable[P, AsyncIterable[AgentResponseItem[StreamingChatMessageContent]]],
+) -> Callable[P, AsyncIterable[AgentResponseItem[StreamingChatMessageContent]]]:
+    """Decorator to trace agent streaming invocation."""
+
+    @functools.wraps(invoke_func)
+    async def wrapper_decorator(
+        *args: P.args,
+        **kwargs: P.kwargs,
+    ) -> AsyncIterable[AgentResponseItem[StreamingChatMessageContent]]:
+        if not are_model_diagnostics_enabled():
+            # If model diagnostics are not enabled, just return the responses
+            async for chunk in invoke_func(*args, **kwargs):
+                yield chunk
+            return
+
+        agent = cast(Agent, args[0])
+        messages = args[1] if len(args) > 1 else None
+
+        with _start_as_current_span(agent) as current_span:
+            _set_agent_invocation_input(current_span, messages)  # type: ignore
+            try:
+                chunks: list[StreamingChatMessageContent] = []
+                async for chunk in invoke_func(*args, **kwargs):
+                    chunks.append(chunk.message)
+                    yield chunk
+                # Concatenate the streaming chunks
+                if chunks:
+                    response = reduce(lambda x, y: x + y, chunks)
+                    _set_agent_invocation_output(current_span, [response])
+                else:
+                    _set_agent_invocation_output(current_span, [])
+            except Exception as e:
+                _set_agent_invocation_error(current_span, e)
+                raise
+
+    # Mark the wrapper decorator as an agent diagnostics decorator
+    wrapper_decorator.__agent_diagnostics__ = True  # type: ignore
+
+    return wrapper_decorator
+
+
+def _start_as_current_span(agent: Agent):
+    """Starts a span for the given agent.
+
+    Args:
+        agent (Agent): The agent for which to start the span.
+
+    Returns:
+        Span: The started span as a context manager.
+    """
+    attributes = {
+        gen_ai_attributes.OPERATION: OPERATION_NAME,
+        gen_ai_attributes.AGENT_ID: agent.id,
+        gen_ai_attributes.AGENT_NAME: agent.name,
+    }
+
+    if agent.description:
+        attributes[gen_ai_attributes.AGENT_DESCRIPTION] = agent.description
+
+    return tracer.start_as_current_span(f"{OPERATION_NAME} {agent.name}", attributes=attributes)
+
+
+def _set_agent_invocation_input(
+    current_span: Span,
+    messages: str | ChatMessageContent | list[str | ChatMessageContent] | None,
+) -> None:
+    """Set the agent input attributes in the span."""
+    if are_sensitive_events_enabled():
+        parsed_messages = _parse_agent_invocation_messages(messages)
+        current_span.set_attribute(
+            gen_ai_attributes.AGENT_INVOCATION_INPUT,
+            json.dumps([message.to_dict() for message in parsed_messages]),
+        )
+
+
+def _set_agent_invocation_output(current_span: Span, response: list[ChatMessageContent]) -> None:
+    """Set the agent output attributes in the span."""
+    if are_sensitive_events_enabled():
+        current_span.set_attribute(
+            gen_ai_attributes.AGENT_INVOCATION_OUTPUT,
+            json.dumps([message.to_dict() for message in response]),
+        )
+
+
+def _set_agent_invocation_error(current_span: Span, error: Exception) -> None:
+    """Set the agent error attributes in the span."""
+    current_span.set_attribute(gen_ai_attributes.ERROR_TYPE, type(error).__name__)
+    current_span.set_status(StatusCode.ERROR, repr(error))
+
+
+def _parse_agent_invocation_messages(
+    messages: str | ChatMessageContent | list[str | ChatMessageContent] | None,
+) -> list[ChatMessageContent]:
+    """Parse the agent invocation messages into a list of ChatMessageContent."""
+    if not messages:
+        return []
+
+    if isinstance(messages, str):
+        return [ChatMessageContent(role=AuthorRole.USER, content=messages)]
+    if isinstance(messages, ChatMessageContent):
+        return [messages]
+    if isinstance(messages, list):
+        return [
+            msg if isinstance(msg, ChatMessageContent) else ChatMessageContent(role=AuthorRole.USER, content=msg)
+            for msg in messages
+        ]
+
+    return []

python/semantic_kernel/utils/telemetry/agent_diagnostics/gen_ai_attributes.py

@@ -10,3 +10,6 @@ OPERATION = "gen_ai.operation.name"
 AGENT_ID = "gen_ai.agent.id"
 AGENT_NAME = "gen_ai.agent.name"
 AGENT_DESCRIPTION = "gen_ai.agent.description"
+AGENT_INVOCATION_INPUT = "gen_ai.agent.invocation_input"
+AGENT_INVOCATION_OUTPUT = "gen_ai.agent.invocation_output"
+ERROR_TYPE = "error.type"

python/tests/unit/utils/agent_diagnostics/conftest.py

@@ -0,0 +1,22 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import pytest
+
+import semantic_kernel
+from semantic_kernel.utils.telemetry.model_diagnostics.model_diagnostics_settings import ModelDiagnosticSettings
+
+
+@pytest.fixture()
+def model_diagnostics_unit_test_env(monkeypatch):
+    """Fixture to set environment variables for Model Diagnostics Unit Tests."""
+    env_vars = {
+        "SEMANTICKERNEL_EXPERIMENTAL_GENAI_ENABLE_OTEL_DIAGNOSTICS": "true",
+        "SEMANTICKERNEL_EXPERIMENTAL_GENAI_ENABLE_OTEL_DIAGNOSTICS_SENSITIVE": "true",
+    }
+
+    for key, value in env_vars.items():
+        monkeypatch.setenv(key, value)
+
+    # Need to reload the settings to pick up the new environment variables since the
+    # settings are loaded at import time and this fixture is called after the import
+    semantic_kernel.utils.telemetry.agent_diagnostics.decorators.MODEL_DIAGNOSTICS_SETTINGS = ModelDiagnosticSettings()

python/tests/unit/utils/agent_diagnostics/test_trace_chat_completion_agent.py

@@ -10,7 +10,29 @@ from semantic_kernel.exceptions.kernel_exceptions import KernelServiceNotFoundEr
 
 
 @patch("semantic_kernel.utils.telemetry.agent_diagnostics.decorators.tracer")
-async def test_chat_completion_agent_invoke(mock_tracer, chat_history):
+async def test_chat_completion_agent_get_response(
+    mock_tracer,
+    chat_history,
+    model_diagnostics_unit_test_env,
+):
+    # Arrange
+    chat_completion_agent = ChatCompletionAgent()
+    thread = ChatHistoryAgentThread(chat_history=chat_history)
+    # Act
+    with pytest.raises(KernelServiceNotFoundError):
+        await chat_completion_agent.get_response(messages="test", thread=thread)
+    # Assert
+    mock_tracer.start_as_current_span.assert_called_once()
+    args, _ = mock_tracer.start_as_current_span.call_args
+    assert args[0] == f"invoke_agent {chat_completion_agent.name}"
+
+
+@patch("semantic_kernel.utils.telemetry.agent_diagnostics.decorators.tracer")
+async def test_chat_completion_agent_invoke(
+    mock_tracer,
+    chat_history,
+    model_diagnostics_unit_test_env,
+):
     # Arrange
     chat_completion_agent = ChatCompletionAgent()
     thread = ChatHistoryAgentThread(chat_history=chat_history)
@@ -19,11 +41,17 @@ async def test_chat_completion_agent_invoke(mock_tracer, chat_history):
         async for _ in chat_completion_agent.invoke(messages="test", thread=thread):
             pass
     # Assert
-    mock_tracer.start_as_current_span.assert_called_once_with(f"invoke_agent {chat_completion_agent.name}")
+    mock_tracer.start_as_current_span.assert_called_once()
+    args, _ = mock_tracer.start_as_current_span.call_args
+    assert args[0] == f"invoke_agent {chat_completion_agent.name}"
 
 
 @patch("semantic_kernel.utils.telemetry.agent_diagnostics.decorators.tracer")
-async def test_chat_completion_agent_invoke_stream(mock_tracer, chat_history):
+async def test_chat_completion_agent_invoke_stream(
+    mock_tracer,
+    chat_history,
+    model_diagnostics_unit_test_env,
+):
     # Arrange
     chat_completion_agent = ChatCompletionAgent()
     thread = ChatHistoryAgentThread(chat_history=chat_history)
@@ -32,4 +60,6 @@ async def test_chat_completion_agent_invoke_stream(mock_tracer, chat_history):
         async for _ in chat_completion_agent.invoke_stream(messages="test", thread=thread):
             pass
     # Assert
-    mock_tracer.start_as_current_span.assert_called_once_with(f"invoke_agent {chat_completion_agent.name}")
+    mock_tracer.start_as_current_span.assert_called_once()
+    args, _ = mock_tracer.start_as_current_span.call_args
+    assert args[0] == f"invoke_agent {chat_completion_agent.name}"

python/tests/unit/utils/agent_diagnostics/test_trace_open_ai_assistant_agent.py

@@ -13,7 +13,50 @@ from semantic_kernel.contents.utils.author_role import AuthorRole
 
 
 @patch("semantic_kernel.utils.telemetry.agent_diagnostics.decorators.tracer")
-async def test_open_ai_assistant_agent_invoke(mock_tracer, chat_history, openai_unit_test_env):
+async def test_open_ai_assistant_agent_get_response(
+    mock_tracer,
+    chat_history,
+    openai_unit_test_env,
+    model_diagnostics_unit_test_env,
+):
+    # Arrange
+    client = AsyncMock(spec=AsyncOpenAI)
+    definition = AsyncMock(spec=Assistant)
+    definition.name = "agentName"
+    definition.description = "agentDescription"
+    definition.id = "agentId"
+    definition.instructions = "agentInstructions"
+    definition.tools = []
+    definition.model = "agentModel"
+    definition.temperature = 1.0
+    definition.top_p = 1.0
+    definition.metadata = {}
+    openai_assistant_agent = OpenAIAssistantAgent(client=client, definition=definition)
+
+    thread = AsyncMock(spec=AssistantAgentThread)
+
+    async def fake_invoke(*args, **kwargs):
+        yield True, ChatMessageContent(role=AuthorRole.ASSISTANT, content="content")
+
+    # Act
+    with patch(
+        "semantic_kernel.agents.open_ai.assistant_thread_actions.AssistantThreadActions.invoke",
+        side_effect=fake_invoke,
+    ):
+        await openai_assistant_agent.get_response(messages="message", thread=thread)
+    # Assert
+    mock_tracer.start_as_current_span.assert_called_once()
+    args, _ = mock_tracer.start_as_current_span.call_args
+    assert args[0] == f"invoke_agent {openai_assistant_agent.name}"
+
+
+@patch("semantic_kernel.utils.telemetry.agent_diagnostics.decorators.tracer")
+async def test_open_ai_assistant_agent_invoke(
+    mock_tracer,
+    chat_history,
+    openai_unit_test_env,
+    model_diagnostics_unit_test_env,
+):
     # Arrange
     client = AsyncMock(spec=AsyncOpenAI)
     definition = AsyncMock(spec=Assistant)
@@ -41,11 +84,18 @@ async def test_open_ai_assistant_agent_invoke(mock_tracer, chat_history, openai_
         async for item in openai_assistant_agent.invoke(messages="message", thread=thread):
             pass
     # Assert
-    mock_tracer.start_as_current_span.assert_called_once_with(f"invoke_agent {openai_assistant_agent.name}")
+    mock_tracer.start_as_current_span.assert_called_once()
+    args, _ = mock_tracer.start_as_current_span.call_args
+    assert args[0] == f"invoke_agent {openai_assistant_agent.name}"
 
 
 @patch("semantic_kernel.utils.telemetry.agent_diagnostics.decorators.tracer")
-async def test_open_ai_assistant_agent_invoke_stream(mock_tracer, chat_history, openai_unit_test_env):
+async def test_open_ai_assistant_agent_invoke_stream(
+    mock_tracer,
+    chat_history,
+    openai_unit_test_env,
+    model_diagnostics_unit_test_env,
+):
     # Arrange
     client = AsyncMock(spec=AsyncOpenAI)
     definition = AsyncMock(spec=Assistant)
@@ -73,4 +123,6 @@ async def test_open_ai_assistant_agent_invoke_stream(mock_tracer, chat_history,
         async for item in openai_assistant_agent.invoke_stream(messages="message", thread=thread):
             pass
     # Assert
-    mock_tracer.start_as_current_span.assert_called_once_with(f"invoke_agent {openai_assistant_agent.name}")
+    mock_tracer.start_as_current_span.assert_called_once()
+    args, _ = mock_tracer.start_as_current_span.call_args
+    assert args[0] == f"invoke_agent {openai_assistant_agent.name}"