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
@@ -0,0 +1,69 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from collections.abc import Generator
|
||||
from typing import Any
|
||||
from unittest.mock import patch
|
||||
|
||||
from opentelemetry.sdk.trace.export import SimpleSpanProcessor, SpanExporter
|
||||
from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
|
||||
from pytest import fixture
|
||||
|
||||
|
||||
@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 True
|
||||
|
||||
|
||||
@fixture(autouse=True)
|
||||
def span_exporter(monkeypatch, enable_otel: bool, enable_sensitive_data: bool) -> Generator[SpanExporter]:
|
||||
"""Fixture to remove environment variables for OtelSettings."""
|
||||
|
||||
env_vars = [
|
||||
"ENABLE_OTEL",
|
||||
"ENABLE_SENSITIVE_DATA",
|
||||
"OTLP_ENDPOINT",
|
||||
"APPLICATIONINSIGHTS_CONNECTION_STRING",
|
||||
"APPLICATIONINSIGHTS_LIVE_METRICS",
|
||||
]
|
||||
|
||||
for key in env_vars:
|
||||
monkeypatch.delenv(key, raising=False) # type: ignore
|
||||
monkeypatch.setenv("ENABLE_OTEL", str(enable_otel)) # type: ignore
|
||||
if not enable_otel:
|
||||
# we overwrite sensitive data for tests
|
||||
enable_sensitive_data = False
|
||||
monkeypatch.setenv("ENABLE_SENSITIVE_DATA", str(enable_sensitive_data)) # type: ignore
|
||||
import importlib
|
||||
|
||||
from opentelemetry import trace
|
||||
|
||||
import agent_framework.observability as observability
|
||||
|
||||
# Reload the module to ensure a clean state for tests, then create a
|
||||
# fresh OtelSettings instance and patch the module attribute.
|
||||
importlib.reload(observability)
|
||||
|
||||
# recreate otel settings with values from above and no file.
|
||||
otel = observability.OtelSettings(env_file_path="test.env")
|
||||
otel.setup_observability()
|
||||
monkeypatch.setattr(observability, "OTEL_SETTINGS", otel, raising=False) # type: ignore
|
||||
exporter = InMemorySpanExporter()
|
||||
with (
|
||||
patch("agent_framework.observability.OTEL_SETTINGS", otel),
|
||||
patch("agent_framework.observability.setup_observability"),
|
||||
):
|
||||
if enable_otel or enable_sensitive_data:
|
||||
trace.get_tracer_provider().add_span_processor(
|
||||
SimpleSpanProcessor(exporter) # type: ignore[func-returns-value]
|
||||
)
|
||||
|
||||
yield exporter
|
||||
# Clean up
|
||||
exporter.clear()
|
||||
@@ -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():
|
||||
|
||||
@@ -3,8 +3,6 @@ from typing import Any
|
||||
|
||||
from pytest import fixture
|
||||
|
||||
from agent_framework.telemetry import OtelSettings, setup_telemetry
|
||||
|
||||
|
||||
# region Connector Settings fixtures
|
||||
@fixture
|
||||
@@ -51,26 +49,3 @@ def openai_unit_test_env(monkeypatch, exclude_list, override_env_param_dict): #
|
||||
monkeypatch.setenv(key, value) # type: ignore
|
||||
|
||||
return env_vars
|
||||
|
||||
|
||||
@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
|
||||
|
||||
@@ -783,7 +783,7 @@ def test_create_streaming_response_content_with_mcp_approval_request() -> None:
|
||||
|
||||
@pytest.mark.parametrize("enable_otel", [False], indirect=True)
|
||||
@pytest.mark.parametrize("enable_sensitive_data", [False], indirect=True)
|
||||
def test_end_to_end_mcp_approval_flow(otel_settings) -> None:
|
||||
def test_end_to_end_mcp_approval_flow() -> None:
|
||||
"""End-to-end mocked test:
|
||||
model issues an mcp_approval_request, user approves, client sends mcp_approval_response.
|
||||
"""
|
||||
|
||||
@@ -0,0 +1 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
@@ -5,9 +5,6 @@ from typing import Any
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
from opentelemetry.sdk.trace import TracerProvider
|
||||
from opentelemetry.sdk.trace.export import SimpleSpanProcessor
|
||||
from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
|
||||
|
||||
from agent_framework import (
|
||||
Executor,
|
||||
@@ -27,54 +24,7 @@ from agent_framework._workflow._edge import (
|
||||
SwitchCaseEdgeGroupDefault,
|
||||
)
|
||||
from agent_framework._workflow._edge_runner import create_edge_runner
|
||||
from agent_framework._workflow._telemetry import EdgeGroupDeliveryStatus, workflow_tracer
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def tracing_enabled():
|
||||
"""Enable tracing for tests."""
|
||||
import os
|
||||
|
||||
original_value = os.environ.get("AGENT_FRAMEWORK_WORKFLOW_ENABLE_OTEL_DIAGNOSTICS")
|
||||
os.environ["AGENT_FRAMEWORK_WORKFLOW_ENABLE_OTEL_DIAGNOSTICS"] = "true"
|
||||
|
||||
# Force reload the settings to pick up the environment variable
|
||||
from agent_framework._workflow._telemetry import WorkflowDiagnosticSettings
|
||||
|
||||
workflow_tracer.settings = WorkflowDiagnosticSettings()
|
||||
|
||||
yield
|
||||
|
||||
# Restore original value
|
||||
if original_value is None:
|
||||
os.environ.pop("AGENT_FRAMEWORK_WORKFLOW_ENABLE_OTEL_DIAGNOSTICS", None)
|
||||
else:
|
||||
os.environ["AGENT_FRAMEWORK_WORKFLOW_ENABLE_OTEL_DIAGNOSTICS"] = original_value
|
||||
|
||||
# Reload settings again
|
||||
workflow_tracer.settings = WorkflowDiagnosticSettings()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def span_exporter(tracing_enabled):
|
||||
"""Set up OpenTelemetry test infrastructure."""
|
||||
|
||||
# Use the built-in InMemorySpanExporter for better compatibility
|
||||
exporter = InMemorySpanExporter()
|
||||
tracer_provider = TracerProvider()
|
||||
tracer_provider.add_span_processor(SimpleSpanProcessor(exporter))
|
||||
|
||||
# Store original tracer
|
||||
original_tracer = workflow_tracer.tracer
|
||||
|
||||
# Set up our test tracer
|
||||
workflow_tracer.tracer = tracer_provider.get_tracer("agent_framework")
|
||||
|
||||
yield exporter
|
||||
|
||||
# Clean up
|
||||
exporter.clear()
|
||||
workflow_tracer.tracer = original_tracer
|
||||
from agent_framework.observability import EdgeGroupDeliveryStatus
|
||||
|
||||
|
||||
@dataclass
|
||||
|
||||
+33
-90
@@ -1,66 +1,22 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import os
|
||||
from collections.abc import Generator
|
||||
from typing import Any, cast
|
||||
|
||||
import pytest
|
||||
from opentelemetry import trace
|
||||
from opentelemetry.sdk.trace import TracerProvider
|
||||
from opentelemetry.sdk.trace.export import SimpleSpanProcessor
|
||||
from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
|
||||
|
||||
from agent_framework import WorkflowBuilder
|
||||
from agent_framework._workflow._executor import Executor, handler
|
||||
from agent_framework._workflow._runner_context import InProcRunnerContext, Message
|
||||
from agent_framework._workflow._shared_state import SharedState
|
||||
from agent_framework._workflow._telemetry import WorkflowTracer, workflow_tracer
|
||||
from agent_framework._workflow._workflow import Workflow
|
||||
from agent_framework._workflow._workflow_context import WorkflowContext
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def tracing_enabled() -> Generator[None, None, None]:
|
||||
"""Enable tracing for tests."""
|
||||
original_value = os.environ.get("AGENT_FRAMEWORK_WORKFLOW_ENABLE_OTEL")
|
||||
os.environ["AGENT_FRAMEWORK_WORKFLOW_ENABLE_OTEL"] = "true"
|
||||
|
||||
# Force reload the settings to pick up the environment variable
|
||||
from agent_framework._workflow._telemetry import WorkflowDiagnosticSettings
|
||||
|
||||
workflow_tracer.settings = WorkflowDiagnosticSettings()
|
||||
|
||||
yield
|
||||
|
||||
# Restore original value
|
||||
if original_value is None:
|
||||
os.environ.pop("AGENT_FRAMEWORK_WORKFLOW_ENABLE_OTEL", None)
|
||||
else:
|
||||
os.environ["AGENT_FRAMEWORK_WORKFLOW_ENABLE_OTEL"] = original_value
|
||||
|
||||
# Reload settings again
|
||||
workflow_tracer.settings = WorkflowDiagnosticSettings()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def span_exporter(tracing_enabled: Any) -> Generator[InMemorySpanExporter, None, None]:
|
||||
"""Set up OpenTelemetry test infrastructure."""
|
||||
# Use the built-in InMemorySpanExporter for better compatibility
|
||||
exporter = InMemorySpanExporter()
|
||||
tracer_provider = TracerProvider()
|
||||
tracer_provider.add_span_processor(SimpleSpanProcessor(exporter))
|
||||
|
||||
# Store original tracer
|
||||
original_tracer = workflow_tracer.tracer
|
||||
|
||||
# Set up our test tracer
|
||||
workflow_tracer.tracer = tracer_provider.get_tracer("agent_framework")
|
||||
|
||||
yield exporter
|
||||
|
||||
# Clean up
|
||||
exporter.clear()
|
||||
workflow_tracer.tracer = original_tracer
|
||||
from agent_framework.observability import (
|
||||
OtelAttr,
|
||||
create_processing_span,
|
||||
create_workflow_span,
|
||||
)
|
||||
|
||||
|
||||
class MockExecutor(Executor):
|
||||
@@ -142,34 +98,7 @@ class FanInAggregator(Executor):
|
||||
return self._processed_messages
|
||||
|
||||
|
||||
async def test_workflow_tracer_configuration() -> None:
|
||||
"""Test that workflow tracer can be enabled and disabled."""
|
||||
# Test disabled by default
|
||||
tracer = WorkflowTracer()
|
||||
assert not tracer.enabled
|
||||
|
||||
# Test enabled with environment variable
|
||||
original_value = os.environ.get("AGENT_FRAMEWORK_WORKFLOW_ENABLE_OTEL")
|
||||
os.environ["AGENT_FRAMEWORK_WORKFLOW_ENABLE_OTEL"] = "true"
|
||||
|
||||
# Force reload the settings to pick up the environment variable
|
||||
from agent_framework._workflow._telemetry import WorkflowDiagnosticSettings
|
||||
|
||||
tracer.settings = WorkflowDiagnosticSettings()
|
||||
|
||||
assert tracer.enabled
|
||||
|
||||
# Restore original value
|
||||
if original_value is None:
|
||||
os.environ.pop("AGENT_FRAMEWORK_WORKFLOW_ENABLE_OTEL", None)
|
||||
else:
|
||||
os.environ["AGENT_FRAMEWORK_WORKFLOW_ENABLE_OTEL"] = original_value
|
||||
|
||||
# Reload settings again
|
||||
tracer.settings = WorkflowDiagnosticSettings()
|
||||
|
||||
|
||||
async def test_span_creation_and_attributes(tracing_enabled: Any, span_exporter: InMemorySpanExporter) -> None:
|
||||
async def test_span_creation_and_attributes(span_exporter: InMemorySpanExporter) -> None:
|
||||
"""Test creation and attributes of all span types (workflow, processing, sending)."""
|
||||
# Create a mock workflow object
|
||||
mock_workflow = cast(
|
||||
@@ -186,12 +115,22 @@ async def test_span_creation_and_attributes(tracing_enabled: Any, span_exporter:
|
||||
)
|
||||
|
||||
# Test all span types in nested context
|
||||
with workflow_tracer.create_workflow_run_span(mock_workflow) as workflow_span:
|
||||
workflow_tracer.add_workflow_event("workflow.started")
|
||||
|
||||
with create_workflow_span(
|
||||
OtelAttr.WORKFLOW_RUN_SPAN,
|
||||
{
|
||||
OtelAttr.WORKFLOW_ID: mock_workflow.id,
|
||||
},
|
||||
) as workflow_span:
|
||||
workflow_span.add_event(OtelAttr.WORKFLOW_STARTED)
|
||||
sending_attributes = {
|
||||
OtelAttr.MESSAGE_TYPE: "ResponseMessage",
|
||||
OtelAttr.MESSAGE_DESTINATION_EXECUTOR_ID: "target-789",
|
||||
}
|
||||
with (
|
||||
workflow_tracer.create_processing_span("executor-456", "TestExecutor", "TestMessage") as processing_span,
|
||||
workflow_tracer.create_sending_span("ResponseMessage", "target-789") as sending_span,
|
||||
create_processing_span("executor-456", "TestExecutor", "TestMessage") as processing_span,
|
||||
create_workflow_span(
|
||||
OtelAttr.MESSAGE_SEND_SPAN, sending_attributes, kind=trace.SpanKind.PRODUCER
|
||||
) as sending_span,
|
||||
):
|
||||
# Verify all spans are recording
|
||||
assert workflow_span is not None and workflow_span.is_recording()
|
||||
@@ -205,7 +144,7 @@ async def test_span_creation_and_attributes(tracing_enabled: Any, span_exporter:
|
||||
workflow_span = next(s for s in spans if s.name == "workflow.run")
|
||||
assert workflow_span.kind == trace.SpanKind.INTERNAL
|
||||
assert workflow_span.attributes is not None
|
||||
assert workflow_span.attributes.get("workflow.id") == "test-workflow-123"
|
||||
assert workflow_span.attributes.get(OtelAttr.WORKFLOW_ID) == "test-workflow-123"
|
||||
assert workflow_span.events is not None
|
||||
event_names = [event.name for event in workflow_span.events]
|
||||
assert "workflow.started" in event_names
|
||||
@@ -226,12 +165,14 @@ async def test_span_creation_and_attributes(tracing_enabled: Any, span_exporter:
|
||||
assert sending_span.attributes.get("message.destination_executor_id") == "target-789"
|
||||
|
||||
|
||||
async def test_trace_context_handling(tracing_enabled: Any, span_exporter: InMemorySpanExporter) -> None:
|
||||
async def test_trace_context_handling(span_exporter: InMemorySpanExporter) -> None:
|
||||
"""Test trace context propagation and handling in messages and executors."""
|
||||
shared_state = SharedState()
|
||||
ctx = InProcRunnerContext()
|
||||
executor = MockExecutor("test-executor")
|
||||
|
||||
span_exporter.clear()
|
||||
|
||||
# Test trace context propagation in messages
|
||||
workflow_ctx: WorkflowContext[str] = WorkflowContext(
|
||||
"test-executor",
|
||||
@@ -273,7 +214,8 @@ async def test_trace_context_handling(tracing_enabled: Any, span_exporter: InMem
|
||||
assert processing_span.attributes.get("message.type") == "str"
|
||||
|
||||
|
||||
async def test_trace_context_disabled_when_tracing_disabled() -> None:
|
||||
@pytest.mark.parametrize("enable_otel", [False], indirect=True)
|
||||
async def test_trace_context_disabled_when_tracing_disabled(enable_otel, span_exporter: InMemorySpanExporter) -> None:
|
||||
"""Test that no trace context is added when tracing is disabled."""
|
||||
# Tracing should be disabled by default
|
||||
shared_state = SharedState()
|
||||
@@ -298,7 +240,7 @@ async def test_trace_context_disabled_when_tracing_disabled() -> None:
|
||||
assert message.source_span_id is None
|
||||
|
||||
|
||||
async def test_end_to_end_workflow_tracing(tracing_enabled: Any, span_exporter: InMemorySpanExporter) -> None:
|
||||
async def test_end_to_end_workflow_tracing(span_exporter: InMemorySpanExporter) -> None:
|
||||
"""Test end-to-end tracing including workflow build, execution, and span linking with fan-in edges."""
|
||||
# Create executors for fan-in scenario
|
||||
executor1 = MockExecutor("executor1")
|
||||
@@ -321,7 +263,7 @@ async def test_end_to_end_workflow_tracing(tracing_enabled: Any, span_exporter:
|
||||
|
||||
build_span = build_spans[0]
|
||||
assert build_span.attributes is not None
|
||||
assert build_span.attributes.get("workflow.id") == workflow.id
|
||||
assert build_span.attributes.get(OtelAttr.WORKFLOW_ID) == workflow.id
|
||||
assert build_span.attributes.get("workflow.definition") is not None
|
||||
definition = build_span.attributes.get("workflow.definition")
|
||||
assert definition == workflow.model_dump_json(by_alias=True)
|
||||
@@ -422,7 +364,7 @@ async def test_end_to_end_workflow_tracing(tracing_enabled: Any, span_exporter:
|
||||
assert len(aggregator_span.links) >= 2, f"Expected at least 2 links, got {len(aggregator_span.links)}"
|
||||
|
||||
|
||||
async def test_workflow_error_handling_in_tracing(tracing_enabled: Any, span_exporter: InMemorySpanExporter) -> None:
|
||||
async def test_workflow_error_handling_in_tracing(span_exporter: InMemorySpanExporter) -> None:
|
||||
"""Test that workflow errors are properly recorded in traces."""
|
||||
|
||||
class FailingExecutor(Executor):
|
||||
@@ -457,7 +399,8 @@ async def test_workflow_error_handling_in_tracing(tracing_enabled: Any, span_exp
|
||||
assert workflow_span.status.status_code.name == "ERROR"
|
||||
|
||||
|
||||
async def test_message_trace_context_serialization() -> None:
|
||||
@pytest.mark.parametrize("enable_otel", [False], indirect=True)
|
||||
async def test_message_trace_context_serialization(span_exporter: InMemorySpanExporter) -> None:
|
||||
"""Test that message trace context is properly serialized/deserialized."""
|
||||
ctx = InProcRunnerContext()
|
||||
|
||||
@@ -491,7 +434,7 @@ async def test_message_trace_context_serialization() -> None:
|
||||
assert restored_msg.source_span_ids == ["span123"] # Test new format
|
||||
|
||||
|
||||
async def test_workflow_build_error_tracing(tracing_enabled: Any, span_exporter: InMemorySpanExporter) -> None:
|
||||
async def test_workflow_build_error_tracing(span_exporter: InMemorySpanExporter) -> None:
|
||||
"""Test that build errors are properly recorded in build spans."""
|
||||
|
||||
# Test validation error by not setting start executor
|
||||
Reference in New Issue
Block a user