mirror of
https://github.com/microsoft/agent-framework.git
synced 2026-06-16 21:04:09 +08:00
Python: Telemetry and observability follow-up (#833)
* updated telemetry work * updated telemetry * slight improvement * updated tests * fixes for telemetry * fixes for mypy * added settings setup to runner to avoid error * streamline usage * updated tests * updated tests * further refinement * fix dumped item for otel * removed enable_workflow_otel * final fixes * final fixes * updated samples * removed exporters * fix tests * fixed last import' * fixed devui
This commit is contained in:
Eduard van Valkenburg
committed by
GitHub
Unverified
GitHub
parent
f93f16a9ad
commit
2576e7a091
@@ -27,7 +27,6 @@ from agent_framework import (
|
||||
ai_function,
|
||||
use_function_invocation,
|
||||
)
|
||||
from agent_framework.telemetry import OtelSettings, setup_telemetry
|
||||
|
||||
if sys.version_info >= (3, 12):
|
||||
from typing import override # type: ignore
|
||||
@@ -38,29 +37,6 @@ else:
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@fixture
|
||||
def enable_otel(request: Any) -> bool:
|
||||
"""Fixture that returns a boolean indicating if Otel is enabled."""
|
||||
return request.param if hasattr(request, "param") else True
|
||||
|
||||
|
||||
@fixture
|
||||
def enable_sensitive_data(request: Any) -> bool:
|
||||
"""Fixture that returns a boolean indicating if sensitive data is enabled."""
|
||||
return request.param if hasattr(request, "param") else False
|
||||
|
||||
|
||||
@fixture
|
||||
def otel_settings(enable_otel: bool, enable_sensitive_data: bool) -> OtelSettings:
|
||||
"""Fixture to set environment variables for OtelSettings."""
|
||||
|
||||
from agent_framework.telemetry import OTEL_SETTINGS
|
||||
|
||||
setup_telemetry(enable_otel=enable_otel, enable_sensitive_data=enable_sensitive_data)
|
||||
|
||||
return OTEL_SETTINGS
|
||||
|
||||
|
||||
@fixture(scope="function")
|
||||
def chat_history() -> list[ChatMessage]:
|
||||
return []
|
||||
|
||||
@@ -0,0 +1,469 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import logging
|
||||
from collections.abc import MutableSequence
|
||||
from typing import Any
|
||||
from unittest.mock import MagicMock, Mock, patch
|
||||
|
||||
import pytest
|
||||
from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
|
||||
from opentelemetry.semconv_ai import SpanAttributes
|
||||
from opentelemetry.trace import StatusCode
|
||||
|
||||
from agent_framework import (
|
||||
AGENT_FRAMEWORK_USER_AGENT,
|
||||
AgentProtocol,
|
||||
AgentRunResponse,
|
||||
AgentThread,
|
||||
BaseChatClient,
|
||||
ChatMessage,
|
||||
ChatOptions,
|
||||
ChatResponse,
|
||||
ChatResponseUpdate,
|
||||
Role,
|
||||
UsageDetails,
|
||||
prepend_agent_framework_to_user_agent,
|
||||
)
|
||||
from agent_framework.exceptions import AgentInitializationError, ChatClientInitializationError
|
||||
from agent_framework.observability import (
|
||||
OPEN_TELEMETRY_AGENT_MARKER,
|
||||
OPEN_TELEMETRY_CHAT_CLIENT_MARKER,
|
||||
ROLE_EVENT_MAP,
|
||||
ChatMessageListTimestampFilter,
|
||||
OtelAttr,
|
||||
get_function_span,
|
||||
use_agent_observability,
|
||||
use_observability,
|
||||
)
|
||||
|
||||
# region Test constants
|
||||
|
||||
|
||||
def test_role_event_map():
|
||||
"""Test that ROLE_EVENT_MAP contains expected mappings."""
|
||||
assert ROLE_EVENT_MAP["system"] == OtelAttr.SYSTEM_MESSAGE
|
||||
assert ROLE_EVENT_MAP["user"] == OtelAttr.USER_MESSAGE
|
||||
assert ROLE_EVENT_MAP["assistant"] == OtelAttr.ASSISTANT_MESSAGE
|
||||
assert ROLE_EVENT_MAP["tool"] == OtelAttr.TOOL_MESSAGE
|
||||
|
||||
|
||||
def test_enum_values():
|
||||
"""Test that OtelAttr enum has expected values."""
|
||||
assert OtelAttr.OPERATION == "gen_ai.operation.name"
|
||||
assert SpanAttributes.LLM_SYSTEM == "gen_ai.system"
|
||||
assert SpanAttributes.LLM_REQUEST_MODEL == "gen_ai.request.model"
|
||||
assert OtelAttr.CHAT_COMPLETION_OPERATION == "chat"
|
||||
assert OtelAttr.TOOL_EXECUTION_OPERATION == "execute_tool"
|
||||
assert OtelAttr.AGENT_INVOKE_OPERATION == "invoke_agent"
|
||||
|
||||
|
||||
# region Test ChatMessageListTimestampFilter
|
||||
|
||||
|
||||
def test_filter_without_index_key():
|
||||
"""Test filter method when record doesn't have INDEX_KEY."""
|
||||
log_filter = ChatMessageListTimestampFilter()
|
||||
record = logging.LogRecord(
|
||||
name="test", level=logging.INFO, pathname="", lineno=0, msg="test message", args=(), exc_info=None
|
||||
)
|
||||
original_created = record.created
|
||||
|
||||
result = log_filter.filter(record)
|
||||
|
||||
assert result is True
|
||||
assert record.created == original_created
|
||||
|
||||
|
||||
def test_filter_with_index_key():
|
||||
"""Test filter method when record has INDEX_KEY."""
|
||||
log_filter = ChatMessageListTimestampFilter()
|
||||
record = logging.LogRecord(
|
||||
name="test", level=logging.INFO, pathname="", lineno=0, msg="test message", args=(), exc_info=None
|
||||
)
|
||||
original_created = record.created
|
||||
|
||||
# Add the index key
|
||||
setattr(record, ChatMessageListTimestampFilter.INDEX_KEY, 5)
|
||||
|
||||
result = log_filter.filter(record)
|
||||
|
||||
assert result is True
|
||||
# Should increment by 5 microseconds (5 * 1e-6)
|
||||
assert record.created == original_created + 5 * 1e-6
|
||||
|
||||
|
||||
def test_index_key_constant():
|
||||
"""Test that INDEX_KEY constant is correctly defined."""
|
||||
assert ChatMessageListTimestampFilter.INDEX_KEY == "chat_message_index"
|
||||
|
||||
|
||||
# region Test get_function_span
|
||||
|
||||
|
||||
def test_start_span_basic(span_exporter: InMemorySpanExporter):
|
||||
"""Test starting a span with basic function info."""
|
||||
# Create a mock function
|
||||
mock_function = Mock()
|
||||
mock_function.name = "test_function"
|
||||
mock_function.description = "Test function description"
|
||||
attributes = {
|
||||
OtelAttr.OPERATION: OtelAttr.TOOL_EXECUTION_OPERATION,
|
||||
OtelAttr.TOOL_NAME: "test_function",
|
||||
OtelAttr.TOOL_DESCRIPTION: "Test function description",
|
||||
OtelAttr.TOOL_TYPE: "function",
|
||||
}
|
||||
span_exporter.clear()
|
||||
with get_function_span(attributes) as function_span:
|
||||
assert function_span is not None
|
||||
function_span.set_attribute("test_attr", "test_value")
|
||||
|
||||
spans = span_exporter.get_finished_spans()
|
||||
assert len(spans) == 1
|
||||
span = spans[0]
|
||||
assert span.name == "execute_tool test_function"
|
||||
assert span.attributes["test_attr"] == "test_value"
|
||||
assert span.attributes[OtelAttr.OPERATION.value] == OtelAttr.TOOL_EXECUTION_OPERATION
|
||||
assert span.attributes[OtelAttr.TOOL_NAME] == "test_function"
|
||||
assert span.attributes[OtelAttr.TOOL_DESCRIPTION] == "Test function description"
|
||||
|
||||
|
||||
def test_start_span_with_tool_call_id(span_exporter: InMemorySpanExporter):
|
||||
"""Test starting a span with tool_call_id."""
|
||||
|
||||
tool_call_id = "test_call_123"
|
||||
attributes = {
|
||||
OtelAttr.OPERATION: OtelAttr.TOOL_EXECUTION_OPERATION,
|
||||
OtelAttr.TOOL_NAME: "test_function",
|
||||
OtelAttr.TOOL_DESCRIPTION: "Test function",
|
||||
OtelAttr.TOOL_TYPE: "function",
|
||||
OtelAttr.TOOL_CALL_ID: tool_call_id,
|
||||
}
|
||||
|
||||
span_exporter.clear()
|
||||
with get_function_span(attributes) as function_span:
|
||||
assert function_span is not None
|
||||
function_span.set_attribute("test_attr", "test_value")
|
||||
spans = span_exporter.get_finished_spans()
|
||||
assert len(spans) == 1
|
||||
span = spans[0]
|
||||
assert span.name == "execute_tool test_function"
|
||||
assert span.attributes["test_attr"] == "test_value"
|
||||
assert span.attributes[OtelAttr.TOOL_CALL_ID] == tool_call_id
|
||||
# Verify all attributes
|
||||
assert span.attributes[OtelAttr.OPERATION.value] == OtelAttr.TOOL_EXECUTION_OPERATION
|
||||
assert span.attributes[OtelAttr.TOOL_NAME] == "test_function"
|
||||
assert span.attributes[OtelAttr.TOOL_DESCRIPTION] == "Test function"
|
||||
assert span.attributes[OtelAttr.TOOL_TYPE] == "function"
|
||||
|
||||
|
||||
# region Test use_observability decorator
|
||||
|
||||
|
||||
def test_decorator_with_valid_class():
|
||||
"""Test that decorator works with a valid BaseChatClient-like class."""
|
||||
|
||||
# Create a mock class with the required methods
|
||||
class MockChatClient:
|
||||
async def get_response(self, messages, **kwargs):
|
||||
return Mock()
|
||||
|
||||
async def get_streaming_response(self, messages, **kwargs):
|
||||
async def gen():
|
||||
yield Mock()
|
||||
|
||||
return gen()
|
||||
|
||||
# Apply the decorator
|
||||
decorated_class = use_observability(MockChatClient)
|
||||
assert hasattr(decorated_class, OPEN_TELEMETRY_CHAT_CLIENT_MARKER)
|
||||
|
||||
|
||||
def test_decorator_with_missing_methods():
|
||||
"""Test that decorator handles classes missing required methods gracefully."""
|
||||
|
||||
class MockChatClient:
|
||||
OTEL_PROVIDER_NAME = "test_provider"
|
||||
|
||||
# Apply the decorator - should not raise an error
|
||||
with pytest.raises(ChatClientInitializationError):
|
||||
use_observability(MockChatClient)
|
||||
|
||||
|
||||
def test_decorator_with_partial_methods():
|
||||
"""Test decorator when only one method is present."""
|
||||
|
||||
class MockChatClient:
|
||||
OTEL_PROVIDER_NAME = "test_provider"
|
||||
|
||||
async def get_response(self, messages, **kwargs):
|
||||
return Mock()
|
||||
|
||||
with pytest.raises(ChatClientInitializationError):
|
||||
use_observability(MockChatClient)
|
||||
|
||||
|
||||
# region Test telemetry decorator with mock client
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_chat_client():
|
||||
"""Create a mock chat client for testing."""
|
||||
|
||||
class MockChatClient(BaseChatClient):
|
||||
def service_url(self):
|
||||
return "https://test.example.com"
|
||||
|
||||
async def _inner_get_response(
|
||||
self, *, messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
|
||||
):
|
||||
return ChatResponse(
|
||||
messages=[ChatMessage(role=Role.ASSISTANT, text="Test response")],
|
||||
usage_details=UsageDetails(input_token_count=10, output_token_count=20),
|
||||
finish_reason=None,
|
||||
)
|
||||
|
||||
async def _inner_get_streaming_response(
|
||||
self, *, messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
|
||||
):
|
||||
yield ChatResponseUpdate(text="Hello", role=Role.ASSISTANT)
|
||||
yield ChatResponseUpdate(text=" world", role=Role.ASSISTANT)
|
||||
|
||||
return MockChatClient
|
||||
|
||||
|
||||
@pytest.mark.parametrize("enable_sensitive_data", [True, False], indirect=True)
|
||||
async def test_chat_client_observability(mock_chat_client, span_exporter: InMemorySpanExporter, enable_sensitive_data):
|
||||
"""Test that when diagnostics are enabled, telemetry is applied."""
|
||||
client = use_observability(mock_chat_client)()
|
||||
|
||||
messages = [ChatMessage(role=Role.USER, text="Test message")]
|
||||
span_exporter.clear()
|
||||
response = await client.get_response(messages=messages, ai_model_id="Test")
|
||||
assert response is not None
|
||||
spans = span_exporter.get_finished_spans()
|
||||
assert len(spans) == 1
|
||||
span = spans[0]
|
||||
assert span.name == "chat Test"
|
||||
assert span.attributes[OtelAttr.OPERATION.value] == OtelAttr.CHAT_COMPLETION_OPERATION
|
||||
assert span.attributes[SpanAttributes.LLM_REQUEST_MODEL] == "Test"
|
||||
assert span.attributes[OtelAttr.INPUT_TOKENS] == 10
|
||||
assert span.attributes[OtelAttr.OUTPUT_TOKENS] == 20
|
||||
if enable_sensitive_data:
|
||||
assert span.attributes[OtelAttr.INPUT_MESSAGES] is not None
|
||||
assert span.attributes[OtelAttr.OUTPUT_MESSAGES] is not None
|
||||
|
||||
|
||||
@pytest.mark.parametrize("enable_sensitive_data", [True, False], indirect=True)
|
||||
async def test_chat_client_streaming_observability(
|
||||
mock_chat_client, span_exporter: InMemorySpanExporter, enable_sensitive_data
|
||||
):
|
||||
"""Test streaming telemetry through the use_observability decorator."""
|
||||
client = use_observability(mock_chat_client)()
|
||||
messages = [ChatMessage(role=Role.USER, text="Test")]
|
||||
span_exporter.clear()
|
||||
# Collect all yielded updates
|
||||
updates = []
|
||||
async for update in client.get_streaming_response(messages=messages, ai_model_id="Test"):
|
||||
updates.append(update)
|
||||
|
||||
# Verify we got the expected updates, this shouldn't be dependent on otel
|
||||
assert len(updates) == 2
|
||||
spans = span_exporter.get_finished_spans()
|
||||
assert len(spans) == 1
|
||||
span = spans[0]
|
||||
assert span.name == "chat Test"
|
||||
assert span.attributes[OtelAttr.OPERATION.value] == OtelAttr.CHAT_COMPLETION_OPERATION
|
||||
assert span.attributes[SpanAttributes.LLM_REQUEST_MODEL] == "Test"
|
||||
if enable_sensitive_data:
|
||||
assert span.attributes[OtelAttr.INPUT_MESSAGES] is not None
|
||||
assert span.attributes[OtelAttr.OUTPUT_MESSAGES] is not None
|
||||
|
||||
|
||||
def test_prepend_user_agent_with_none_value():
|
||||
"""Test prepend user agent with None value in headers."""
|
||||
headers = {"User-Agent": None}
|
||||
result = prepend_agent_framework_to_user_agent(headers)
|
||||
|
||||
# Should handle None gracefully
|
||||
assert "User-Agent" in result
|
||||
assert AGENT_FRAMEWORK_USER_AGENT in str(result["User-Agent"])
|
||||
|
||||
|
||||
# region Test use_agent_observability decorator
|
||||
|
||||
|
||||
def test_agent_decorator_with_valid_class():
|
||||
"""Test that agent decorator works with a valid ChatAgent-like class."""
|
||||
|
||||
# Create a mock class with the required methods
|
||||
class MockChatClientAgent:
|
||||
AGENT_SYSTEM_NAME = "test_agent_system"
|
||||
|
||||
def __init__(self):
|
||||
self.id = "test_agent_id"
|
||||
self.name = "test_agent"
|
||||
self.display_name = "Test Agent"
|
||||
self.description = "Test agent description"
|
||||
|
||||
async def run(self, messages=None, *, thread=None, **kwargs):
|
||||
return Mock()
|
||||
|
||||
async def run_stream(self, messages=None, *, thread=None, **kwargs):
|
||||
async def gen():
|
||||
yield Mock()
|
||||
|
||||
return gen()
|
||||
|
||||
def get_new_thread(self) -> AgentThread:
|
||||
return AgentThread()
|
||||
|
||||
# Apply the decorator
|
||||
decorated_class = use_agent_observability(MockChatClientAgent)
|
||||
|
||||
assert hasattr(decorated_class, OPEN_TELEMETRY_AGENT_MARKER)
|
||||
|
||||
|
||||
def test_agent_decorator_with_missing_methods():
|
||||
"""Test that agent decorator handles classes missing required methods gracefully."""
|
||||
|
||||
class MockAgent:
|
||||
AGENT_SYSTEM_NAME = "test_agent_system"
|
||||
|
||||
# Apply the decorator - should not raise an error
|
||||
with pytest.raises(AgentInitializationError):
|
||||
use_agent_observability(MockAgent)
|
||||
|
||||
|
||||
def test_agent_decorator_with_partial_methods():
|
||||
"""Test agent decorator when only one method is present."""
|
||||
from agent_framework.observability import use_agent_observability
|
||||
|
||||
class MockAgent:
|
||||
AGENT_SYSTEM_NAME = "test_agent_system"
|
||||
|
||||
def __init__(self):
|
||||
self.id = "test_agent_id"
|
||||
self.name = "test_agent"
|
||||
self.display_name = "Test Agent"
|
||||
|
||||
async def run(self, messages=None, *, thread=None, **kwargs):
|
||||
return Mock()
|
||||
|
||||
with pytest.raises(AgentInitializationError):
|
||||
use_agent_observability(MockAgent)
|
||||
|
||||
|
||||
# region Test agent telemetry decorator with mock agent
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_chat_agent():
|
||||
"""Create a mock chat client agent for testing."""
|
||||
|
||||
class MockChatClientAgent:
|
||||
AGENT_SYSTEM_NAME = "test_agent_system"
|
||||
|
||||
def __init__(self):
|
||||
self.id = "test_agent_id"
|
||||
self.name = "test_agent"
|
||||
self.display_name = "Test Agent"
|
||||
self.description = "Test agent description"
|
||||
|
||||
async def run(self, messages=None, *, thread=None, **kwargs):
|
||||
return AgentRunResponse(
|
||||
messages=[ChatMessage(role=Role.ASSISTANT, text="Agent response")],
|
||||
usage_details=UsageDetails(input_token_count=15, output_token_count=25),
|
||||
response_id="test_response_id",
|
||||
raw_representation=Mock(finish_reason=Mock(value="stop")),
|
||||
)
|
||||
|
||||
async def run_stream(self, messages=None, *, thread=None, **kwargs):
|
||||
from agent_framework import AgentRunResponseUpdate
|
||||
|
||||
yield AgentRunResponseUpdate(text="Hello", role=Role.ASSISTANT)
|
||||
yield AgentRunResponseUpdate(text=" from agent", role=Role.ASSISTANT)
|
||||
|
||||
return MockChatClientAgent
|
||||
|
||||
|
||||
@pytest.mark.parametrize("enable_sensitive_data", [True, False], indirect=True)
|
||||
async def test_agent_instrumentation_enabled(
|
||||
mock_chat_agent: AgentProtocol, span_exporter: InMemorySpanExporter, enable_sensitive_data
|
||||
):
|
||||
"""Test that when agent diagnostics are enabled, telemetry is applied."""
|
||||
|
||||
agent = use_agent_observability(mock_chat_agent)()
|
||||
|
||||
span_exporter.clear()
|
||||
response = await agent.run("Test message")
|
||||
assert response is not None
|
||||
spans = span_exporter.get_finished_spans()
|
||||
assert len(spans) == 1
|
||||
span = spans[0]
|
||||
assert span.name == "invoke_agent Test Agent"
|
||||
assert span.attributes[OtelAttr.OPERATION.value] == OtelAttr.AGENT_INVOKE_OPERATION
|
||||
assert span.attributes[OtelAttr.AGENT_ID] == "test_agent_id"
|
||||
assert span.attributes[OtelAttr.AGENT_NAME] == "Test Agent"
|
||||
assert span.attributes[OtelAttr.AGENT_DESCRIPTION] == "Test agent description"
|
||||
assert span.attributes[SpanAttributes.LLM_REQUEST_MODEL] == "unknown"
|
||||
assert span.attributes[OtelAttr.INPUT_TOKENS] == 15
|
||||
assert span.attributes[OtelAttr.OUTPUT_TOKENS] == 25
|
||||
if enable_sensitive_data:
|
||||
assert span.attributes[OtelAttr.OUTPUT_MESSAGES] is not None
|
||||
|
||||
|
||||
@pytest.mark.parametrize("enable_sensitive_data", [True, False], indirect=True)
|
||||
async def test_agent_streaming_response_with_diagnostics_enabled_via_decorator(
|
||||
mock_chat_agent: AgentProtocol, span_exporter: InMemorySpanExporter, enable_sensitive_data
|
||||
):
|
||||
"""Test agent streaming telemetry through the use_agent_observability decorator."""
|
||||
agent = use_agent_observability(mock_chat_agent)()
|
||||
span_exporter.clear()
|
||||
updates = []
|
||||
async for update in agent.run_stream("Test message"):
|
||||
updates.append(update)
|
||||
|
||||
# Verify we got the expected updates
|
||||
assert len(updates) == 2
|
||||
spans = span_exporter.get_finished_spans()
|
||||
assert len(spans) == 1
|
||||
span = spans[0]
|
||||
assert span.name == "invoke_agent Test Agent"
|
||||
assert span.attributes[OtelAttr.OPERATION.value] == OtelAttr.AGENT_INVOKE_OPERATION
|
||||
assert span.attributes[OtelAttr.AGENT_ID] == "test_agent_id"
|
||||
assert span.attributes[OtelAttr.AGENT_NAME] == "Test Agent"
|
||||
assert span.attributes[OtelAttr.AGENT_DESCRIPTION] == "Test agent description"
|
||||
assert span.attributes[SpanAttributes.LLM_REQUEST_MODEL] == "unknown"
|
||||
if enable_sensitive_data:
|
||||
assert span.attributes.get(OtelAttr.OUTPUT_MESSAGES) is not None # Streaming, so no usage yet
|
||||
|
||||
|
||||
async def test_agent_run_with_exception_handling(mock_chat_agent: AgentProtocol):
|
||||
"""Test agent run with exception handling."""
|
||||
|
||||
async def run_with_error(self, messages=None, *, thread=None, **kwargs):
|
||||
raise RuntimeError("Agent run error")
|
||||
|
||||
mock_chat_agent.run = run_with_error
|
||||
|
||||
agent = use_agent_observability(mock_chat_agent)()
|
||||
|
||||
from opentelemetry.trace import Span
|
||||
|
||||
with (
|
||||
patch("agent_framework.observability._get_span") as mock_get_span,
|
||||
):
|
||||
mock_span = MagicMock(spec=Span)
|
||||
# Ensure the patched context manager returns mock_span when entered
|
||||
mock_get_span.return_value.__enter__.return_value = mock_span
|
||||
# Should raise the exception and call error handler
|
||||
with pytest.raises(RuntimeError, match="Agent run error"):
|
||||
await agent.run("Test message")
|
||||
|
||||
# Verify error was recorded
|
||||
# Check that both error attributes were set on the span
|
||||
mock_span.set_attribute.assert_called_with(OtelAttr.ERROR_TYPE, "RuntimeError")
|
||||
mock_span.record_exception.assert_called_once()
|
||||
mock_span.set_status.assert_called_once_with(
|
||||
status=StatusCode.ERROR, description=repr(RuntimeError("Agent run error"))
|
||||
)
|
||||
@@ -1,44 +1,14 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import logging
|
||||
from collections.abc import MutableSequence
|
||||
from typing import Any
|
||||
from unittest.mock import MagicMock, Mock, patch
|
||||
|
||||
import pytest
|
||||
from opentelemetry.semconv_ai import SpanAttributes
|
||||
from opentelemetry.trace import StatusCode
|
||||
from unittest.mock import patch
|
||||
|
||||
from agent_framework import (
|
||||
AgentProtocol,
|
||||
AgentRunResponse,
|
||||
AgentThread,
|
||||
BaseChatClient,
|
||||
ChatMessage,
|
||||
ChatOptions,
|
||||
ChatResponse,
|
||||
ChatResponseUpdate,
|
||||
Role,
|
||||
UsageDetails,
|
||||
)
|
||||
from agent_framework.exceptions import AgentInitializationError, ChatClientInitializationError
|
||||
from agent_framework.telemetry import (
|
||||
AGENT_FRAMEWORK_USER_AGENT,
|
||||
OPEN_TELEMETRY_AGENT_MARKER,
|
||||
OPEN_TELEMETRY_CHAT_CLIENT_MARKER,
|
||||
ROLE_EVENT_MAP,
|
||||
USER_AGENT_KEY,
|
||||
USER_AGENT_TELEMETRY_DISABLED_ENV_VAR,
|
||||
ChatMessageListTimestampFilter,
|
||||
OtelAttr,
|
||||
get_function_span,
|
||||
prepend_agent_framework_to_user_agent,
|
||||
use_agent_telemetry,
|
||||
use_telemetry,
|
||||
)
|
||||
|
||||
from .utils import CopyingMock
|
||||
|
||||
# region Test constants
|
||||
|
||||
|
||||
@@ -59,13 +29,13 @@ def test_agent_framework_user_agent_format():
|
||||
|
||||
def test_app_info_when_telemetry_enabled():
|
||||
"""Test that APP_INFO is set when telemetry is enabled."""
|
||||
with patch("agent_framework.telemetry.IS_TELEMETRY_ENABLED", True):
|
||||
with patch("agent_framework._telemetry.IS_TELEMETRY_ENABLED", True):
|
||||
import importlib
|
||||
|
||||
import agent_framework.telemetry
|
||||
import agent_framework._telemetry
|
||||
|
||||
importlib.reload(agent_framework.telemetry)
|
||||
from agent_framework.telemetry import APP_INFO
|
||||
importlib.reload(agent_framework._telemetry)
|
||||
from agent_framework import APP_INFO
|
||||
|
||||
assert APP_INFO is not None
|
||||
assert "agent-framework-version" in APP_INFO
|
||||
@@ -75,7 +45,7 @@ def test_app_info_when_telemetry_enabled():
|
||||
def test_app_info_when_telemetry_disabled():
|
||||
"""Test that APP_INFO is None when telemetry is disabled."""
|
||||
# Test the logic directly since APP_INFO is set at module import time
|
||||
with patch("agent_framework.telemetry.IS_TELEMETRY_ENABLED", False):
|
||||
with patch("agent_framework._telemetry.IS_TELEMETRY_ENABLED", False):
|
||||
# Simulate the module's logic for APP_INFO
|
||||
test_app_info = (
|
||||
{
|
||||
@@ -87,24 +57,6 @@ def test_app_info_when_telemetry_disabled():
|
||||
assert test_app_info is None
|
||||
|
||||
|
||||
def test_role_event_map():
|
||||
"""Test that ROLE_EVENT_MAP contains expected mappings."""
|
||||
assert ROLE_EVENT_MAP["system"] == OtelAttr.SYSTEM_MESSAGE
|
||||
assert ROLE_EVENT_MAP["user"] == OtelAttr.USER_MESSAGE
|
||||
assert ROLE_EVENT_MAP["assistant"] == OtelAttr.ASSISTANT_MESSAGE
|
||||
assert ROLE_EVENT_MAP["tool"] == OtelAttr.TOOL_MESSAGE
|
||||
|
||||
|
||||
def test_enum_values():
|
||||
"""Test that OtelAttr enum has expected values."""
|
||||
assert OtelAttr.OPERATION == "gen_ai.operation.name"
|
||||
assert SpanAttributes.LLM_SYSTEM == "gen_ai.system"
|
||||
assert SpanAttributes.LLM_REQUEST_MODEL == "gen_ai.request.model"
|
||||
assert OtelAttr.CHAT_COMPLETION_OPERATION == "chat"
|
||||
assert OtelAttr.TOOL_EXECUTION_OPERATION == "execute_tool"
|
||||
assert OtelAttr.AGENT_INVOKE_OPERATION == "invoke_agent"
|
||||
|
||||
|
||||
# region Test prepend_agent_framework_to_user_agent
|
||||
|
||||
|
||||
@@ -144,415 +96,3 @@ def test_modifies_original_dict():
|
||||
|
||||
assert result is headers # Same object
|
||||
assert "User-Agent" in headers
|
||||
|
||||
|
||||
# region Test ChatMessageListTimestampFilter
|
||||
|
||||
|
||||
def test_filter_without_index_key():
|
||||
"""Test filter method when record doesn't have INDEX_KEY."""
|
||||
log_filter = ChatMessageListTimestampFilter()
|
||||
record = logging.LogRecord(
|
||||
name="test", level=logging.INFO, pathname="", lineno=0, msg="test message", args=(), exc_info=None
|
||||
)
|
||||
original_created = record.created
|
||||
|
||||
result = log_filter.filter(record)
|
||||
|
||||
assert result is True
|
||||
assert record.created == original_created
|
||||
|
||||
|
||||
def test_filter_with_index_key():
|
||||
"""Test filter method when record has INDEX_KEY."""
|
||||
log_filter = ChatMessageListTimestampFilter()
|
||||
record = logging.LogRecord(
|
||||
name="test", level=logging.INFO, pathname="", lineno=0, msg="test message", args=(), exc_info=None
|
||||
)
|
||||
original_created = record.created
|
||||
|
||||
# Add the index key
|
||||
setattr(record, ChatMessageListTimestampFilter.INDEX_KEY, 5)
|
||||
|
||||
result = log_filter.filter(record)
|
||||
|
||||
assert result is True
|
||||
# Should increment by 5 microseconds (5 * 1e-6)
|
||||
assert record.created == original_created + 5 * 1e-6
|
||||
|
||||
|
||||
def test_index_key_constant():
|
||||
"""Test that INDEX_KEY constant is correctly defined."""
|
||||
assert ChatMessageListTimestampFilter.INDEX_KEY == "chat_message_index"
|
||||
|
||||
|
||||
# region Test get_function_span
|
||||
|
||||
|
||||
def test_start_span_basic():
|
||||
"""Test starting a span with basic function info."""
|
||||
mock_tracer = Mock()
|
||||
with patch("agent_framework.telemetry.tracer", mock_tracer):
|
||||
mock_span = Mock()
|
||||
mock_tracer.start_as_current_span.return_value = mock_span
|
||||
|
||||
# Create a mock function
|
||||
mock_function = Mock()
|
||||
mock_function.name = "test_function"
|
||||
mock_function.description = "Test function description"
|
||||
attributes = {
|
||||
OtelAttr.OPERATION: OtelAttr.TOOL_EXECUTION_OPERATION,
|
||||
OtelAttr.TOOL_NAME: "test_function",
|
||||
OtelAttr.TOOL_DESCRIPTION: "Test function description",
|
||||
OtelAttr.TOOL_TYPE: "function",
|
||||
}
|
||||
|
||||
result = get_function_span(attributes)
|
||||
|
||||
assert result == mock_span
|
||||
mock_tracer.start_as_current_span.assert_called_once()
|
||||
|
||||
call_args = mock_tracer.start_as_current_span.call_args
|
||||
assert call_args[1]["name"] == "execute_tool test_function"
|
||||
|
||||
attributes = call_args[1]["attributes"]
|
||||
assert attributes[OtelAttr.OPERATION.value] == OtelAttr.TOOL_EXECUTION_OPERATION
|
||||
assert attributes[OtelAttr.TOOL_NAME] == "test_function"
|
||||
assert attributes[OtelAttr.TOOL_DESCRIPTION] == "Test function description"
|
||||
|
||||
|
||||
def test_start_span_with_tool_call_id():
|
||||
"""Test starting a span with tool_call_id."""
|
||||
mock_tracer = Mock()
|
||||
with patch("agent_framework.telemetry.tracer", mock_tracer):
|
||||
mock_span = CopyingMock()
|
||||
mock_tracer.start_as_current_span.return_value = mock_span
|
||||
|
||||
mock_function = Mock()
|
||||
mock_function.name = "test_function"
|
||||
mock_function.description = "Test function"
|
||||
|
||||
tool_call_id = "test_call_123"
|
||||
attributes = {
|
||||
OtelAttr.OPERATION: OtelAttr.TOOL_EXECUTION_OPERATION,
|
||||
OtelAttr.TOOL_NAME: "test_function",
|
||||
OtelAttr.TOOL_DESCRIPTION: "Test function",
|
||||
OtelAttr.TOOL_TYPE: "function",
|
||||
OtelAttr.TOOL_CALL_ID: tool_call_id,
|
||||
}
|
||||
|
||||
_ = get_function_span(attributes)
|
||||
|
||||
call_args = mock_tracer.start_as_current_span.call_args
|
||||
attributes = call_args[1]["attributes"]
|
||||
assert attributes[OtelAttr.TOOL_CALL_ID] == "test_call_123"
|
||||
|
||||
|
||||
# region Test use_telemetry decorator
|
||||
|
||||
|
||||
def test_decorator_with_valid_class():
|
||||
"""Test that decorator works with a valid BaseChatClient-like class."""
|
||||
|
||||
# Create a mock class with the required methods
|
||||
class MockChatClient:
|
||||
async def get_response(self, messages, **kwargs):
|
||||
return Mock()
|
||||
|
||||
async def get_streaming_response(self, messages, **kwargs):
|
||||
async def gen():
|
||||
yield Mock()
|
||||
|
||||
return gen()
|
||||
|
||||
# Apply the decorator
|
||||
decorated_class = use_telemetry(MockChatClient)
|
||||
assert hasattr(decorated_class, OPEN_TELEMETRY_CHAT_CLIENT_MARKER)
|
||||
|
||||
|
||||
def test_decorator_with_missing_methods():
|
||||
"""Test that decorator handles classes missing required methods gracefully."""
|
||||
|
||||
class MockChatClient:
|
||||
OTEL_PROVIDER_NAME = "test_provider"
|
||||
|
||||
# Apply the decorator - should not raise an error
|
||||
with pytest.raises(ChatClientInitializationError):
|
||||
use_telemetry(MockChatClient)
|
||||
|
||||
|
||||
def test_decorator_with_partial_methods():
|
||||
"""Test decorator when only one method is present."""
|
||||
|
||||
class MockChatClient:
|
||||
OTEL_PROVIDER_NAME = "test_provider"
|
||||
|
||||
async def get_response(self, messages, **kwargs):
|
||||
return Mock()
|
||||
|
||||
with pytest.raises(ChatClientInitializationError):
|
||||
use_telemetry(MockChatClient)
|
||||
|
||||
|
||||
# region Test telemetry decorator with mock client
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_chat_client():
|
||||
"""Create a mock chat client for testing."""
|
||||
|
||||
class MockChatClient(BaseChatClient):
|
||||
def service_url(self):
|
||||
return "https://test.example.com"
|
||||
|
||||
async def _inner_get_response(
|
||||
self, *, messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
|
||||
):
|
||||
return ChatResponse(
|
||||
messages=[ChatMessage(role=Role.ASSISTANT, text="Test response")],
|
||||
usage_details=UsageDetails(input_token_count=10, output_token_count=20),
|
||||
finish_reason=None,
|
||||
)
|
||||
|
||||
async def _inner_get_streaming_response(
|
||||
self, *, messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
|
||||
):
|
||||
yield ChatResponseUpdate(text="Hello", role=Role.ASSISTANT)
|
||||
yield ChatResponseUpdate(text=" world", role=Role.ASSISTANT)
|
||||
|
||||
return MockChatClient
|
||||
|
||||
|
||||
@pytest.mark.parametrize("enable_sensitive_data", [True, False], indirect=True)
|
||||
async def test_instrumentation_enabled(mock_chat_client, otel_settings):
|
||||
"""Test that when diagnostics are enabled, telemetry is applied."""
|
||||
client = use_telemetry(mock_chat_client)()
|
||||
|
||||
messages = [ChatMessage(role=Role.USER, text="Test message")]
|
||||
chat_options = ChatOptions()
|
||||
|
||||
with (
|
||||
patch("agent_framework.telemetry._get_span") as mock_response_span,
|
||||
patch("agent_framework.telemetry._capture_messages") as mock_log_messages,
|
||||
):
|
||||
response = await client.get_response(messages=messages, chat_options=chat_options)
|
||||
assert response is not None
|
||||
mock_response_span.assert_called_once()
|
||||
|
||||
# Check that log messages was called only if sensitive events are enabled
|
||||
assert mock_log_messages.call_count == (2 if otel_settings.enable_sensitive_data else 0)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("enable_sensitive_data", [True, False], indirect=True)
|
||||
async def test_streaming_response_with_otel(mock_chat_client, otel_settings):
|
||||
"""Test streaming telemetry through the use_telemetry decorator."""
|
||||
client = use_telemetry(mock_chat_client)()
|
||||
messages = [ChatMessage(role=Role.USER, text="Test")]
|
||||
chat_options = ChatOptions()
|
||||
|
||||
with (
|
||||
patch("agent_framework.telemetry._get_span") as mock_response_span,
|
||||
patch("agent_framework.telemetry._capture_messages") as mock_log_messages,
|
||||
patch("agent_framework.telemetry._capture_response") as mock_set_output,
|
||||
):
|
||||
# Collect all yielded updates
|
||||
updates = []
|
||||
async for update in client.get_streaming_response(messages=messages, chat_options=chat_options):
|
||||
updates.append(update)
|
||||
|
||||
# Verify we got the expected updates, this shouldn't be dependent on otel
|
||||
assert len(updates) == 2
|
||||
|
||||
# Verify telemetry calls were made
|
||||
mock_response_span.assert_called_once()
|
||||
if otel_settings.enable_sensitive_data:
|
||||
mock_log_messages.assert_called()
|
||||
assert mock_log_messages.call_count == 2 # One for input, one for output
|
||||
else:
|
||||
mock_log_messages.assert_not_called()
|
||||
|
||||
mock_set_output.assert_called_once()
|
||||
|
||||
|
||||
def test_prepend_user_agent_with_none_value():
|
||||
"""Test prepend user agent with None value in headers."""
|
||||
headers = {"User-Agent": None}
|
||||
result = prepend_agent_framework_to_user_agent(headers)
|
||||
|
||||
# Should handle None gracefully
|
||||
assert "User-Agent" in result
|
||||
assert AGENT_FRAMEWORK_USER_AGENT in str(result["User-Agent"])
|
||||
|
||||
|
||||
# region Test use_agent_telemetry decorator
|
||||
|
||||
|
||||
def test_agent_decorator_with_valid_class():
|
||||
"""Test that agent decorator works with a valid ChatAgent-like class."""
|
||||
|
||||
# Create a mock class with the required methods
|
||||
class MockChatClientAgent:
|
||||
AGENT_SYSTEM_NAME = "test_agent_system"
|
||||
|
||||
def __init__(self):
|
||||
self.id = "test_agent_id"
|
||||
self.name = "test_agent"
|
||||
self.display_name = "Test Agent"
|
||||
self.description = "Test agent description"
|
||||
|
||||
async def run(self, messages=None, *, thread=None, **kwargs):
|
||||
return Mock()
|
||||
|
||||
async def run_stream(self, messages=None, *, thread=None, **kwargs):
|
||||
async def gen():
|
||||
yield Mock()
|
||||
|
||||
return gen()
|
||||
|
||||
def get_new_thread(self) -> AgentThread:
|
||||
return AgentThread()
|
||||
|
||||
# Apply the decorator
|
||||
decorated_class = use_agent_telemetry(MockChatClientAgent)
|
||||
|
||||
assert hasattr(decorated_class, OPEN_TELEMETRY_AGENT_MARKER)
|
||||
|
||||
|
||||
def test_agent_decorator_with_missing_methods():
|
||||
"""Test that agent decorator handles classes missing required methods gracefully."""
|
||||
|
||||
class MockAgent:
|
||||
AGENT_SYSTEM_NAME = "test_agent_system"
|
||||
|
||||
# Apply the decorator - should not raise an error
|
||||
with pytest.raises(AgentInitializationError):
|
||||
use_agent_telemetry(MockAgent)
|
||||
|
||||
|
||||
def test_agent_decorator_with_partial_methods():
|
||||
"""Test agent decorator when only one method is present."""
|
||||
from agent_framework.telemetry import use_agent_telemetry
|
||||
|
||||
class MockAgent:
|
||||
AGENT_SYSTEM_NAME = "test_agent_system"
|
||||
|
||||
def __init__(self):
|
||||
self.id = "test_agent_id"
|
||||
self.name = "test_agent"
|
||||
self.display_name = "Test Agent"
|
||||
|
||||
async def run(self, messages=None, *, thread=None, **kwargs):
|
||||
return Mock()
|
||||
|
||||
with pytest.raises(AgentInitializationError):
|
||||
use_agent_telemetry(MockAgent)
|
||||
|
||||
|
||||
# region Test agent telemetry decorator with mock agent
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_chat_client_agent():
|
||||
"""Create a mock chat client agent for testing."""
|
||||
|
||||
class MockChatClientAgent:
|
||||
AGENT_SYSTEM_NAME = "test_agent_system"
|
||||
|
||||
def __init__(self):
|
||||
self.id = "test_agent_id"
|
||||
self.name = "test_agent"
|
||||
self.display_name = "Test Agent"
|
||||
self.description = "Test agent description"
|
||||
|
||||
async def run(self, messages=None, *, thread=None, **kwargs):
|
||||
return AgentRunResponse(
|
||||
messages=[ChatMessage(role=Role.ASSISTANT, text="Agent response")],
|
||||
usage_details=UsageDetails(input_token_count=15, output_token_count=25),
|
||||
response_id="test_response_id",
|
||||
raw_representation=Mock(finish_reason=Mock(value="stop")),
|
||||
)
|
||||
|
||||
async def run_stream(self, messages=None, *, thread=None, **kwargs):
|
||||
from agent_framework import AgentRunResponseUpdate
|
||||
|
||||
yield AgentRunResponseUpdate(text="Hello", role=Role.ASSISTANT)
|
||||
yield AgentRunResponseUpdate(text=" from agent", role=Role.ASSISTANT)
|
||||
|
||||
return MockChatClientAgent
|
||||
|
||||
|
||||
@pytest.mark.parametrize("enable_sensitive_data", [True, False], indirect=True)
|
||||
async def test_agent_instrumentation_enabled(mock_chat_client_agent: AgentProtocol, otel_settings):
|
||||
"""Test that when agent diagnostics are enabled, telemetry is applied."""
|
||||
|
||||
agent = use_agent_telemetry(mock_chat_client_agent)()
|
||||
|
||||
with (
|
||||
patch("agent_framework.telemetry.use_span") as mock_use_span,
|
||||
patch("agent_framework.telemetry.logger") as mock_logger,
|
||||
):
|
||||
response = await agent.run("Test message")
|
||||
assert response is not None
|
||||
mock_use_span.assert_called_once()
|
||||
# Check that logger.info was called (telemetry logs input/output)
|
||||
assert mock_logger.info.call_count == (2 if otel_settings.enable_sensitive_data else 0)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("enable_sensitive_data", [True, False], indirect=True)
|
||||
async def test_agent_streaming_response_with_diagnostics_enabled_via_decorator(
|
||||
mock_chat_client_agent: AgentProtocol, otel_settings
|
||||
):
|
||||
"""Test agent streaming telemetry through the use_agent_telemetry decorator."""
|
||||
agent = use_agent_telemetry(mock_chat_client_agent)()
|
||||
|
||||
with (
|
||||
patch("agent_framework.telemetry._get_span") as mock_get_span,
|
||||
patch("agent_framework.telemetry._capture_messages") as mock_capture_messages,
|
||||
patch("agent_framework.telemetry._capture_response") as mock_capture_response,
|
||||
):
|
||||
# Collect all yielded updates
|
||||
updates = []
|
||||
async for update in agent.run_stream("Test message"):
|
||||
updates.append(update)
|
||||
|
||||
# Verify we got the expected updates
|
||||
assert len(updates) == 2
|
||||
|
||||
# Verify telemetry calls were made
|
||||
mock_get_span.assert_called_once()
|
||||
mock_capture_response.assert_called_once()
|
||||
if otel_settings.enable_sensitive_data:
|
||||
mock_capture_messages.assert_called()
|
||||
else:
|
||||
mock_capture_messages.assert_not_called()
|
||||
|
||||
|
||||
async def test_agent_run_with_exception_handling(mock_chat_client_agent: AgentProtocol):
|
||||
"""Test agent run with exception handling."""
|
||||
|
||||
async def run_with_error(self, messages=None, *, thread=None, **kwargs):
|
||||
raise RuntimeError("Agent run error")
|
||||
|
||||
mock_chat_client_agent.run = run_with_error
|
||||
|
||||
agent = use_agent_telemetry(mock_chat_client_agent)()
|
||||
|
||||
from opentelemetry.trace import Span
|
||||
|
||||
with (
|
||||
patch("agent_framework.telemetry._get_span") as mock_get_span,
|
||||
):
|
||||
mock_span = MagicMock(spec=Span)
|
||||
# Ensure the patched context manager returns mock_span when entered
|
||||
mock_get_span.return_value.__enter__.return_value = mock_span
|
||||
# Should raise the exception and call error handler
|
||||
with pytest.raises(RuntimeError, match="Agent run error"):
|
||||
await agent.run("Test message")
|
||||
|
||||
# Verify error was recorded
|
||||
# Check that both error attributes were set on the span
|
||||
mock_span.set_attribute.assert_called_with(OtelAttr.ERROR_TYPE, "RuntimeError")
|
||||
mock_span.record_exception.assert_called_once()
|
||||
mock_span.set_status.assert_called_once_with(
|
||||
status=StatusCode.ERROR, description=repr(RuntimeError("Agent run error"))
|
||||
)
|
||||
|
||||
@@ -1,9 +1,10 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from typing import Any
|
||||
from unittest.mock import Mock, patch
|
||||
from unittest.mock import Mock
|
||||
|
||||
import pytest
|
||||
from opentelemetry import trace
|
||||
from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
|
||||
from pydantic import BaseModel
|
||||
|
||||
from agent_framework import (
|
||||
@@ -15,9 +16,7 @@ from agent_framework import (
|
||||
)
|
||||
from agent_framework._tools import _parse_inputs
|
||||
from agent_framework.exceptions import ToolException
|
||||
from agent_framework.telemetry import OtelAttr
|
||||
|
||||
from .utils import CopyingMock
|
||||
from agent_framework.observability import OtelAttr
|
||||
|
||||
# region AIFunction and ai_function decorator tests
|
||||
|
||||
@@ -85,8 +84,7 @@ async def test_ai_function_decorator_with_async():
|
||||
assert (await async_test_tool(1, 2)) == 3
|
||||
|
||||
|
||||
@pytest.mark.parametrize("enable_sensitive_data", [True], indirect=True)
|
||||
async def test_ai_function_invoke_telemetry_enabled(otel_settings):
|
||||
async def test_ai_function_invoke_telemetry_enabled(span_exporter: InMemorySpanExporter):
|
||||
"""Test the ai_function invoke method with telemetry enabled."""
|
||||
|
||||
@ai_function(
|
||||
@@ -97,52 +95,83 @@ async def test_ai_function_invoke_telemetry_enabled(otel_settings):
|
||||
"""A function that adds two numbers for telemetry testing."""
|
||||
return x + y
|
||||
|
||||
# Mock the tracer and span
|
||||
with (
|
||||
patch("agent_framework.telemetry.tracer"),
|
||||
# the span creation uses a form of deepcopy, so need to mock that way
|
||||
patch("agent_framework._tools.get_function_span", new_callable=CopyingMock) as mock_start_span,
|
||||
):
|
||||
mock_span = Mock()
|
||||
mock_context_manager = Mock()
|
||||
mock_context_manager.__enter__ = Mock(return_value=mock_span)
|
||||
mock_context_manager.__exit__ = Mock(return_value=None)
|
||||
mock_start_span.return_value = mock_context_manager
|
||||
# Mock the histogram
|
||||
mock_histogram = Mock()
|
||||
telemetry_test_tool._invocation_duration_histogram = mock_histogram
|
||||
span_exporter.clear()
|
||||
# Call invoke
|
||||
result = await telemetry_test_tool.invoke(x=1, y=2, tool_call_id="test_call_id")
|
||||
|
||||
# Mock the histogram
|
||||
mock_histogram = Mock()
|
||||
telemetry_test_tool._invocation_duration_histogram = mock_histogram
|
||||
# Verify result
|
||||
assert result == 3
|
||||
|
||||
# Call invoke
|
||||
result = await telemetry_test_tool.invoke(x=1, y=2, tool_call_id="test_call_id")
|
||||
# Verify telemetry calls
|
||||
spans = span_exporter.get_finished_spans()
|
||||
assert len(spans) == 1
|
||||
span = spans[0]
|
||||
assert OtelAttr.TOOL_EXECUTION_OPERATION.value in span.name
|
||||
assert "telemetry_test_tool" in span.name
|
||||
assert span.attributes[OtelAttr.TOOL_NAME] == "telemetry_test_tool"
|
||||
assert span.attributes[OtelAttr.TOOL_CALL_ID] == "test_call_id"
|
||||
assert span.attributes[OtelAttr.TOOL_TYPE] == "function"
|
||||
assert span.attributes[OtelAttr.TOOL_DESCRIPTION] == "A test tool for telemetry"
|
||||
assert span.attributes[OtelAttr.TOOL_ARGUMENTS] == '{"x": 1, "y": 2}'
|
||||
assert span.attributes[OtelAttr.TOOL_RESULT] == "3"
|
||||
|
||||
# Verify result
|
||||
assert result == 3
|
||||
|
||||
# Verify telemetry calls
|
||||
mock_start_span.assert_called_once_with(
|
||||
attributes={
|
||||
OtelAttr.OPERATION: OtelAttr.TOOL_EXECUTION_OPERATION,
|
||||
OtelAttr.TOOL_NAME: "telemetry_test_tool",
|
||||
OtelAttr.TOOL_CALL_ID: "test_call_id",
|
||||
OtelAttr.TOOL_TYPE: "function",
|
||||
OtelAttr.TOOL_DESCRIPTION: "A test tool for telemetry",
|
||||
OtelAttr.TOOL_ARGUMENTS: '{"x": 1, "y": 2}',
|
||||
}
|
||||
)
|
||||
assert mock_span.set_attribute.call_count == 2
|
||||
|
||||
# Verify histogram was called with correct attributes
|
||||
mock_histogram.record.assert_called_once()
|
||||
call_args = mock_histogram.record.call_args
|
||||
assert call_args[0][0] > 0 # duration should be positive
|
||||
attributes = call_args[1]["attributes"]
|
||||
assert attributes[OtelAttr.MEASUREMENT_FUNCTION_TAG_NAME] == "telemetry_test_tool"
|
||||
assert attributes[OtelAttr.TOOL_CALL_ID] == "test_call_id"
|
||||
# Verify histogram was called with correct attributes
|
||||
mock_histogram.record.assert_called_once()
|
||||
call_args = mock_histogram.record.call_args
|
||||
assert call_args[0][0] > 0 # duration should be positive
|
||||
attributes = call_args[1]["attributes"]
|
||||
assert attributes[OtelAttr.MEASUREMENT_FUNCTION_TAG_NAME] == "telemetry_test_tool"
|
||||
assert attributes[OtelAttr.TOOL_CALL_ID] == "test_call_id"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("enable_sensitive_data", [True], indirect=True)
|
||||
async def test_ai_function_invoke_telemetry_with_pydantic_args(otel_settings):
|
||||
@pytest.mark.parametrize("enable_sensitive_data", [False], indirect=True)
|
||||
async def test_ai_function_invoke_telemetry_sensitive_disabled(span_exporter: InMemorySpanExporter):
|
||||
"""Test the ai_function invoke method with telemetry enabled."""
|
||||
|
||||
@ai_function(
|
||||
name="telemetry_test_tool",
|
||||
description="A test tool for telemetry",
|
||||
)
|
||||
def telemetry_test_tool(x: int, y: int) -> int:
|
||||
"""A function that adds two numbers for telemetry testing."""
|
||||
return x + y
|
||||
|
||||
# Mock the histogram
|
||||
mock_histogram = Mock()
|
||||
telemetry_test_tool._invocation_duration_histogram = mock_histogram
|
||||
span_exporter.clear()
|
||||
# Call invoke
|
||||
result = await telemetry_test_tool.invoke(x=1, y=2, tool_call_id="test_call_id")
|
||||
|
||||
# Verify result
|
||||
assert result == 3
|
||||
|
||||
# Verify telemetry calls
|
||||
spans = span_exporter.get_finished_spans()
|
||||
assert len(spans) == 1
|
||||
span = spans[0]
|
||||
assert OtelAttr.TOOL_EXECUTION_OPERATION.value in span.name
|
||||
assert "telemetry_test_tool" in span.name
|
||||
assert span.attributes[OtelAttr.TOOL_NAME] == "telemetry_test_tool"
|
||||
assert span.attributes[OtelAttr.TOOL_CALL_ID] == "test_call_id"
|
||||
assert span.attributes[OtelAttr.TOOL_TYPE] == "function"
|
||||
assert span.attributes[OtelAttr.TOOL_DESCRIPTION] == "A test tool for telemetry"
|
||||
assert OtelAttr.TOOL_ARGUMENTS not in span.attributes
|
||||
assert OtelAttr.TOOL_RESULT not in span.attributes
|
||||
|
||||
# Verify histogram was called with correct attributes
|
||||
mock_histogram.record.assert_called_once()
|
||||
call_args = mock_histogram.record.call_args
|
||||
assert call_args[0][0] > 0 # duration should be positive
|
||||
attributes = call_args[1]["attributes"]
|
||||
assert attributes[OtelAttr.MEASUREMENT_FUNCTION_TAG_NAME] == "telemetry_test_tool"
|
||||
assert attributes[OtelAttr.TOOL_CALL_ID] == "test_call_id"
|
||||
|
||||
|
||||
async def test_ai_function_invoke_telemetry_with_pydantic_args(span_exporter: InMemorySpanExporter):
|
||||
"""Test the ai_function invoke method with Pydantic model arguments."""
|
||||
|
||||
@ai_function(
|
||||
@@ -156,42 +185,27 @@ async def test_ai_function_invoke_telemetry_with_pydantic_args(otel_settings):
|
||||
# Create arguments as Pydantic model instance
|
||||
args_model = pydantic_test_tool.input_model(x=5, y=10)
|
||||
|
||||
with (
|
||||
patch("agent_framework.telemetry.tracer"),
|
||||
# the span creation uses a form of deepcopy, so need to mock that way
|
||||
patch("agent_framework._tools.get_function_span", new_callable=CopyingMock) as mock_start_span,
|
||||
):
|
||||
mock_span = Mock()
|
||||
mock_context_manager = Mock()
|
||||
mock_context_manager.__enter__ = Mock(return_value=mock_span)
|
||||
mock_context_manager.__exit__ = Mock(return_value=None)
|
||||
mock_start_span.return_value = mock_context_manager
|
||||
mock_histogram = Mock()
|
||||
pydantic_test_tool._invocation_duration_histogram = mock_histogram
|
||||
span_exporter.clear()
|
||||
# Call invoke with Pydantic model
|
||||
result = await pydantic_test_tool.invoke(arguments=args_model, tool_call_id="pydantic_call")
|
||||
|
||||
mock_histogram = Mock()
|
||||
pydantic_test_tool._invocation_duration_histogram = mock_histogram
|
||||
|
||||
# Call invoke with Pydantic model
|
||||
result = await pydantic_test_tool.invoke(arguments=args_model, tool_call_id="pydantic_call")
|
||||
|
||||
# Verify result
|
||||
assert result == 15
|
||||
|
||||
# Verify telemetry calls
|
||||
mock_start_span.assert_called_once_with(
|
||||
attributes={
|
||||
OtelAttr.OPERATION: OtelAttr.TOOL_EXECUTION_OPERATION,
|
||||
OtelAttr.TOOL_NAME: "pydantic_test_tool",
|
||||
OtelAttr.TOOL_CALL_ID: "pydantic_call",
|
||||
OtelAttr.TOOL_TYPE: "function",
|
||||
OtelAttr.TOOL_DESCRIPTION: "A test tool with Pydantic args",
|
||||
OtelAttr.TOOL_ARGUMENTS: '{"x":5,"y":10}',
|
||||
}
|
||||
)
|
||||
assert mock_span.set_attribute.call_count == 2
|
||||
# Verify result
|
||||
assert result == 15
|
||||
spans = span_exporter.get_finished_spans()
|
||||
assert len(spans) == 1
|
||||
span = spans[0]
|
||||
assert OtelAttr.TOOL_EXECUTION_OPERATION.value in span.name
|
||||
assert "pydantic_test_tool" in span.name
|
||||
assert span.attributes[OtelAttr.TOOL_NAME] == "pydantic_test_tool"
|
||||
assert span.attributes[OtelAttr.TOOL_CALL_ID] == "pydantic_call"
|
||||
assert span.attributes[OtelAttr.TOOL_TYPE] == "function"
|
||||
assert span.attributes[OtelAttr.TOOL_DESCRIPTION] == "A test tool with Pydantic args"
|
||||
assert span.attributes[OtelAttr.TOOL_ARGUMENTS] == '{"x":5,"y":10}'
|
||||
|
||||
|
||||
@pytest.mark.parametrize("otel_settings", [(True, True)], indirect=True)
|
||||
async def test_ai_function_invoke_telemetry_with_exception(otel_settings):
|
||||
async def test_ai_function_invoke_telemetry_with_exception(span_exporter: InMemorySpanExporter):
|
||||
"""Test the ai_function invoke method with telemetry when an exception occurs."""
|
||||
|
||||
@ai_function(
|
||||
@@ -202,41 +216,33 @@ async def test_ai_function_invoke_telemetry_with_exception(otel_settings):
|
||||
"""A function that raises an exception for telemetry testing."""
|
||||
raise ValueError("Test exception for telemetry")
|
||||
|
||||
with (
|
||||
patch("agent_framework.telemetry.tracer"),
|
||||
# the span creation uses a form of deepcopy, so need to mock that way
|
||||
patch("agent_framework._tools.get_function_span", new_callable=CopyingMock) as mock_start_span,
|
||||
):
|
||||
mock_span = Mock()
|
||||
mock_context_manager = Mock()
|
||||
mock_context_manager.__enter__ = Mock(return_value=mock_span)
|
||||
mock_context_manager.__exit__ = Mock(return_value=None)
|
||||
mock_start_span.return_value = mock_context_manager
|
||||
mock_histogram = Mock()
|
||||
exception_test_tool._invocation_duration_histogram = mock_histogram
|
||||
span_exporter.clear()
|
||||
# Call invoke and expect exception
|
||||
with pytest.raises(ValueError, match="Test exception for telemetry"):
|
||||
await exception_test_tool.invoke(x=1, y=2, tool_call_id="exception_call")
|
||||
spans = span_exporter.get_finished_spans()
|
||||
assert len(spans) == 1
|
||||
span = spans[0]
|
||||
assert OtelAttr.TOOL_EXECUTION_OPERATION.value in span.name
|
||||
assert "exception_test_tool" in span.name
|
||||
assert span.attributes[OtelAttr.TOOL_NAME] == "exception_test_tool"
|
||||
assert span.attributes[OtelAttr.TOOL_CALL_ID] == "exception_call"
|
||||
assert span.attributes[OtelAttr.TOOL_TYPE] == "function"
|
||||
assert span.attributes[OtelAttr.TOOL_DESCRIPTION] == "A test tool that raises an exception"
|
||||
assert span.attributes[OtelAttr.TOOL_ARGUMENTS] == '{"x": 1, "y": 2}'
|
||||
assert span.attributes[OtelAttr.ERROR_TYPE] == ValueError.__name__
|
||||
assert span.status.status_code == trace.StatusCode.ERROR
|
||||
|
||||
mock_histogram = Mock()
|
||||
exception_test_tool._invocation_duration_histogram = mock_histogram
|
||||
|
||||
# Call invoke and expect exception
|
||||
with pytest.raises(ValueError, match="Test exception for telemetry"):
|
||||
await exception_test_tool.invoke(x=1, y=2, tool_call_id="exception_call")
|
||||
|
||||
# Verify telemetry calls
|
||||
mock_start_span.assert_called_once()
|
||||
|
||||
# Verify span exception recording
|
||||
mock_span.record_exception.assert_called_once()
|
||||
mock_span.set_attribute.assert_called()
|
||||
mock_span.set_status.assert_called_once()
|
||||
|
||||
# Verify histogram was called with error attributes
|
||||
mock_histogram.record.assert_called_once()
|
||||
call_args = mock_histogram.record.call_args
|
||||
attributes = call_args[1]["attributes"]
|
||||
assert attributes[OtelAttr.ERROR_TYPE] == ValueError.__name__
|
||||
# Verify histogram was called with error attributes
|
||||
mock_histogram.record.assert_called_once()
|
||||
call_args = mock_histogram.record.call_args
|
||||
attributes = call_args[1]["attributes"]
|
||||
assert attributes[OtelAttr.ERROR_TYPE] == ValueError.__name__
|
||||
|
||||
|
||||
@pytest.mark.parametrize("enable_sensitive_data", [True], indirect=True)
|
||||
async def test_ai_function_invoke_telemetry_async_function(otel_settings):
|
||||
async def test_ai_function_invoke_telemetry_async_function(span_exporter: InMemorySpanExporter):
|
||||
"""Test the ai_function invoke method with telemetry on async function."""
|
||||
|
||||
@ai_function(
|
||||
@@ -247,44 +253,30 @@ async def test_ai_function_invoke_telemetry_async_function(otel_settings):
|
||||
"""An async function for telemetry testing."""
|
||||
return x * y
|
||||
|
||||
with (
|
||||
patch("agent_framework.telemetry.tracer"),
|
||||
# the span creation uses a form of deepcopy, so need to mock that way
|
||||
patch("agent_framework._tools.get_function_span", new_callable=CopyingMock) as mock_start_span,
|
||||
):
|
||||
mock_span = Mock()
|
||||
mock_context_manager = Mock()
|
||||
mock_context_manager.__enter__ = Mock(return_value=mock_span)
|
||||
mock_context_manager.__exit__ = Mock(return_value=None)
|
||||
mock_start_span.return_value = mock_context_manager
|
||||
mock_histogram = Mock()
|
||||
async_telemetry_test._invocation_duration_histogram = mock_histogram
|
||||
span_exporter.clear()
|
||||
# Call invoke
|
||||
result = await async_telemetry_test.invoke(x=3, y=4, tool_call_id="async_call")
|
||||
|
||||
mock_histogram = Mock()
|
||||
async_telemetry_test._invocation_duration_histogram = mock_histogram
|
||||
# Verify result
|
||||
assert result == 12
|
||||
spans = span_exporter.get_finished_spans()
|
||||
assert len(spans) == 1
|
||||
span = spans[0]
|
||||
assert OtelAttr.TOOL_EXECUTION_OPERATION.value in span.name
|
||||
assert "async_telemetry_test" in span.name
|
||||
assert span.attributes[OtelAttr.TOOL_NAME] == "async_telemetry_test"
|
||||
assert span.attributes[OtelAttr.TOOL_CALL_ID] == "async_call"
|
||||
assert span.attributes[OtelAttr.TOOL_TYPE] == "function"
|
||||
assert span.attributes[OtelAttr.TOOL_DESCRIPTION] == "An async test tool for telemetry"
|
||||
assert span.attributes[OtelAttr.TOOL_ARGUMENTS] == '{"x": 3, "y": 4}'
|
||||
|
||||
# Call invoke
|
||||
result = await async_telemetry_test.invoke(x=3, y=4, tool_call_id="async_call")
|
||||
|
||||
# Verify result
|
||||
assert result == 12
|
||||
|
||||
# Verify telemetry calls
|
||||
mock_start_span.assert_called_once_with(
|
||||
attributes={
|
||||
OtelAttr.OPERATION: OtelAttr.TOOL_EXECUTION_OPERATION,
|
||||
OtelAttr.TOOL_NAME: "async_telemetry_test",
|
||||
OtelAttr.TOOL_CALL_ID: "async_call",
|
||||
OtelAttr.TOOL_TYPE: "function",
|
||||
OtelAttr.TOOL_DESCRIPTION: "An async test tool for telemetry",
|
||||
OtelAttr.TOOL_ARGUMENTS: '{"x": 3, "y": 4}',
|
||||
}
|
||||
)
|
||||
assert mock_span.set_attribute.call_count == 2
|
||||
|
||||
# Verify histogram recording
|
||||
mock_histogram.record.assert_called_once()
|
||||
call_args = mock_histogram.record.call_args
|
||||
attributes = call_args[1]["attributes"]
|
||||
assert attributes[OtelAttr.MEASUREMENT_FUNCTION_TAG_NAME] == "async_telemetry_test"
|
||||
# Verify histogram recording
|
||||
mock_histogram.record.assert_called_once()
|
||||
call_args = mock_histogram.record.call_args
|
||||
attributes = call_args[1]["attributes"]
|
||||
assert attributes[OtelAttr.MEASUREMENT_FUNCTION_TAG_NAME] == "async_telemetry_test"
|
||||
|
||||
|
||||
async def test_ai_function_invoke_invalid_pydantic_args():
|
||||
|
||||
Reference in New Issue
Block a user