Python: Add agent invocation spans (#10255)

### 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.
-->
Address feature request:
#10174

### Description

<!-- Describe your changes, the overall approach, the underlying design.
These notes will help understanding how your code works. Thanks! -->
This pull request introduces a new decorator for tracing agent
invocations and applies it to various methods in the
`ChatCompletionAgent` and `OpenAIAssistantAgent` classes. Additionally,
it includes new unit tests to ensure the decorator is correctly applied
and functions as expected.

### Key Changes:

**Decorator Implementation:**
* Added a new `trace_agent_invocation` decorator in
`python/semantic_kernel/utils/telemetry/agent_diagnostics/decorators.py`
to trace agent invocations using OpenTelemetry.

**Decorator Application:**
* Applied the `trace_agent_invocation` decorator to the `invoke` and
`invoke_stream` methods in `ChatCompletionAgent`
(`python/semantic_kernel/agents/chat_completion/chat_completion_agent.py`).
[[1]](diffhunk://#diff-0f0a27c107368504c4347c88528b7b4234dcead1919005bcae13f5d16d6cf26dR19)
[[2]](diffhunk://#diff-0f0a27c107368504c4347c88528b7b4234dcead1919005bcae13f5d16d6cf26dR80)
[[3]](diffhunk://#diff-0f0a27c107368504c4347c88528b7b4234dcead1919005bcae13f5d16d6cf26dR134)
* Applied the `trace_agent_invocation` decorator to the `invoke` and
`invoke_stream` methods in `OpenAIAssistantAgent`
(`python/semantic_kernel/agents/open_ai/open_ai_assistant_base.py`).
[[1]](diffhunk://#diff-70e75c57136d3a90d045af4b8d59fb8e8101251d7c5126458bb5d9e6c36556a4R46)
[[2]](diffhunk://#diff-70e75c57136d3a90d045af4b8d59fb8e8101251d7c5126458bb5d9e6c36556a4R608)
[[3]](diffhunk://#diff-70e75c57136d3a90d045af4b8d59fb8e8101251d7c5126458bb5d9e6c36556a4R861)

### 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/semantic_kernel/agents/chat_completion/chat_completion_agent.py

@@ -16,6 +16,7 @@
 from semantic_kernel.contents.utils.author_role import AuthorRole
 from semantic_kernel.exceptions import KernelServiceNotFoundError
 from semantic_kernel.utils.experimental_decorator import experimental_class
+from semantic_kernel.utils.telemetry.agent_diagnostics.decorators import trace_agent_invocation
 
 if TYPE_CHECKING:
     from semantic_kernel.kernel import Kernel
@@ -76,6 +77,7 @@ def __init__(
             args["kernel"] = kernel
         super().__init__(**args)
 
+    @trace_agent_invocation
     async def invoke(self, history: ChatHistory) -> AsyncIterable[ChatMessageContent]:
         """Invoke the chat history handler.
 
@@ -129,6 +131,7 @@ async def invoke(self, history: ChatHistory) -> AsyncIterable[ChatMessageContent
             message.name = self.name
             yield message
 
+    @trace_agent_invocation
     async def invoke_stream(self, history: ChatHistory) -> AsyncIterable[StreamingChatMessageContent]:
         """Invoke the chat history handler in streaming mode.
 

python/semantic_kernel/agents/open_ai/open_ai_assistant_base.py

@@ -43,6 +43,7 @@
     AgentInvokeException,
 )
 from semantic_kernel.utils.experimental_decorator import experimental_class
+from semantic_kernel.utils.telemetry.agent_diagnostics.decorators import trace_agent_invocation
 
 if TYPE_CHECKING:
     from semantic_kernel.contents.chat_history import ChatHistory
@@ -604,6 +605,7 @@ async def delete_vector_store(self, vector_store_id: str) -> None:
 
     # region Agent Invoke Methods
 
+    @trace_agent_invocation
     async def invoke(
         self,
         thread_id: str,
@@ -856,6 +858,7 @@ def sort_key(step: RunStep):
                             yield True, content
                 processed_step_ids.add(completed_step.id)
 
+    @trace_agent_invocation
     async def invoke_stream(
         self,
         thread_id: str,

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

@@ -0,0 +1,34 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import functools
+from collections.abc import AsyncIterable, Callable
+from typing import TYPE_CHECKING, Any
+
+from opentelemetry.trace import get_tracer
+
+from semantic_kernel.utils.experimental_decorator import experimental_function
+
+if TYPE_CHECKING:
+    from semantic_kernel.agents.agent import Agent
+
+
+# Creates a tracer from the global tracer provider
+tracer = get_tracer(__name__)
+
+
+@experimental_function
+def trace_agent_invocation(invoke_func: Callable) -> Callable:
+    """Decorator to trace agent invocation."""
+
+    @functools.wraps(invoke_func)
+    async def wrapper_decorator(*args: Any, **kwargs: Any) -> AsyncIterable:
+        agent: "Agent" = args[0]
+
+        with tracer.start_as_current_span(agent.name):
+            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

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

@@ -0,0 +1,43 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import pytest
+
+from semantic_kernel.agents.chat_completion.chat_completion_agent import ChatCompletionAgent
+from semantic_kernel.agents.open_ai.open_ai_assistant_base import OpenAIAssistantBase
+
+pytestmark = pytest.mark.parametrize(
+    "decorated_method, expected_attribute",
+    [
+        # region ChatCompletionAgent
+        pytest.param(
+            ChatCompletionAgent.invoke,
+            "__agent_diagnostics__",
+            id="ChatCompletionAgent.invoke",
+        ),
+        pytest.param(
+            ChatCompletionAgent.invoke_stream,
+            "__agent_diagnostics__",
+            id="ChatCompletionAgent.invoke_stream",
+        ),
+        # endregion
+        # region OpenAIAssistantAgent
+        pytest.param(
+            OpenAIAssistantBase.invoke,
+            "__agent_diagnostics__",
+            id="OpenAIAssistantBase.invoke",
+        ),
+        pytest.param(
+            OpenAIAssistantBase.invoke_stream,
+            "__agent_diagnostics__",
+            id="OpenAIAssistantBase.invoke_stream",
+        ),
+        # endregion
+    ],
+)
+
+
+def test_decorated(decorated_method, expected_attribute):
+    """Test that the connectors are being decorated properly with the agent diagnostics decorators."""
+    assert hasattr(decorated_method, expected_attribute) and getattr(decorated_method, expected_attribute), (
+        f"{decorated_method} should be decorated with the appropriate agent diagnostics decorator."
+    )

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

@@ -0,0 +1,33 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+
+from unittest.mock import patch
+
+import pytest
+
+from semantic_kernel.agents.chat_completion.chat_completion_agent import ChatCompletionAgent
+from semantic_kernel.exceptions.kernel_exceptions import KernelServiceNotFoundError
+
+
+@patch("semantic_kernel.utils.telemetry.agent_diagnostics.decorators.tracer")
+async def test_chat_completion_agent_invoke(mock_tracer, chat_history):
+    # Arrange
+    chat_completion_agent = ChatCompletionAgent()
+    # Act
+    with pytest.raises(KernelServiceNotFoundError):
+        async for _ in chat_completion_agent.invoke(chat_history):
+            pass
+    # Assert
+    mock_tracer.start_as_current_span.assert_called_once_with(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):
+    # Arrange
+    chat_completion_agent = ChatCompletionAgent()
+    # Act
+    with pytest.raises(KernelServiceNotFoundError):
+        async for _ in chat_completion_agent.invoke_stream(chat_history):
+            pass
+    # Assert
+    mock_tracer.start_as_current_span.assert_called_once_with(chat_completion_agent.name)

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

@@ -0,0 +1,33 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+
+from unittest.mock import patch
+
+import pytest
+
+from semantic_kernel.agents.open_ai.open_ai_assistant_agent import OpenAIAssistantAgent
+from semantic_kernel.exceptions.agent_exceptions import AgentInitializationException
+
+
+@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):
+    # Arrange
+    open_ai_assistant_agent = OpenAIAssistantAgent()
+    # Act
+    with pytest.raises(AgentInitializationException):
+        async for _ in open_ai_assistant_agent.invoke(chat_history):
+            pass
+    # Assert
+    mock_tracer.start_as_current_span.assert_called_once_with(open_ai_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):
+    # Arrange
+    open_ai_assistant_agent = OpenAIAssistantAgent()
+    # Act
+    with pytest.raises(AgentInitializationException):
+        async for _ in open_ai_assistant_agent.invoke_stream(chat_history):
+            pass
+    # Assert
+    mock_tracer.start_as_current_span.assert_called_once_with(open_ai_assistant_agent.name)