From 6ecfdcfda92ec466e8a554bc9f022afb9e903cad Mon Sep 17 00:00:00 2001 From: Tao Chen Date: Tue, 21 Jan 2025 23:26:18 -0800 Subject: [PATCH] Python: Add agent invocation spans (#10255) ### Motivation and Context Address feature request: https://github.com/microsoft/semantic-kernel/issues/10174 ### Description 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 - [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 :smile: --------- Co-authored-by: Evan Mattson <35585003+moonbox3@users.noreply.github.com> --- .../chat_completion/chat_completion_agent.py | 3 ++ .../agents/open_ai/open_ai_assistant_base.py | 3 ++ .../telemetry/agent_diagnostics/__init__.py | 0 .../telemetry/agent_diagnostics/decorators.py | 34 +++++++++++++++ .../agent_diagnostics/test_agent_decorated.py | 43 +++++++++++++++++++ .../test_trace_chat_completion_agent.py | 33 ++++++++++++++ .../test_trace_open_ai_assistant_agent.py | 33 ++++++++++++++ ...corated.py => test_connector_decorated.py} | 0 8 files changed, 149 insertions(+) create mode 100644 python/semantic_kernel/utils/telemetry/agent_diagnostics/__init__.py create mode 100644 python/semantic_kernel/utils/telemetry/agent_diagnostics/decorators.py create mode 100644 python/tests/unit/utils/agent_diagnostics/test_agent_decorated.py create mode 100644 python/tests/unit/utils/agent_diagnostics/test_trace_chat_completion_agent.py create mode 100644 python/tests/unit/utils/agent_diagnostics/test_trace_open_ai_assistant_agent.py rename python/tests/unit/utils/model_diagnostics/{test_decorated.py => test_connector_decorated.py} (100%) diff --git a/python/semantic_kernel/agents/chat_completion/chat_completion_agent.py b/python/semantic_kernel/agents/chat_completion/chat_completion_agent.py index fede96bccd48..352787e81d8c 100644 --- a/python/semantic_kernel/agents/chat_completion/chat_completion_agent.py +++ b/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. diff --git a/python/semantic_kernel/agents/open_ai/open_ai_assistant_base.py b/python/semantic_kernel/agents/open_ai/open_ai_assistant_base.py index 32fc549ca36a..3b072043751c 100644 --- a/python/semantic_kernel/agents/open_ai/open_ai_assistant_base.py +++ b/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, diff --git a/python/semantic_kernel/utils/telemetry/agent_diagnostics/__init__.py b/python/semantic_kernel/utils/telemetry/agent_diagnostics/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/python/semantic_kernel/utils/telemetry/agent_diagnostics/decorators.py b/python/semantic_kernel/utils/telemetry/agent_diagnostics/decorators.py new file mode 100644 index 000000000000..91e5f28197f4 --- /dev/null +++ b/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 diff --git a/python/tests/unit/utils/agent_diagnostics/test_agent_decorated.py b/python/tests/unit/utils/agent_diagnostics/test_agent_decorated.py new file mode 100644 index 000000000000..747bf38c52b1 --- /dev/null +++ b/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." + ) diff --git a/python/tests/unit/utils/agent_diagnostics/test_trace_chat_completion_agent.py b/python/tests/unit/utils/agent_diagnostics/test_trace_chat_completion_agent.py new file mode 100644 index 000000000000..3c1df16efa14 --- /dev/null +++ b/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) diff --git a/python/tests/unit/utils/agent_diagnostics/test_trace_open_ai_assistant_agent.py b/python/tests/unit/utils/agent_diagnostics/test_trace_open_ai_assistant_agent.py new file mode 100644 index 000000000000..d4a7ae6134ef --- /dev/null +++ b/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) diff --git a/python/tests/unit/utils/model_diagnostics/test_decorated.py b/python/tests/unit/utils/model_diagnostics/test_connector_decorated.py similarity index 100% rename from python/tests/unit/utils/model_diagnostics/test_decorated.py rename to python/tests/unit/utils/model_diagnostics/test_connector_decorated.py