agent/agent-framework
1
0
Fork 0
mirror of https://github.com/microsoft/agent-framework.git synced 2026-06-16 21:04:09 +08:00
Code Issues Packages Projects Releases Wiki Activity

Python: Improved telemetry setup (#421)

Browse Source 
* test with stack and simplified names

* quick demo of agent decorator

* moved builder to protocol to enhance functionality

* undid chatclientAgent -> agent rename

* one more

* reverted AIAgent rename

* final reverts

* fixed foundry import

* revert changes

* streamlined otel and fcc decorators

* cleanup of telemetry

* further refinement

* lots of updates

* fixed typing

* fix for mypy

* added input and output atttributes

* fix import

* initial work on baking in otel

* major update to telemetry

* final fixes after rename

* fix

* fix test

* updated tests

* fix for tests

* fixes for tests

* updated based on comments

* removed agent decorator

* fix for Python: ServiceResponseException when using multiple tools
Fixes #649

* addressed comments

* fix tests

* fix tests

* fix tools tests

* fix for conversation_id in assistants client

* fix responses test

* fix tests and mypy

* updated test

* foundry fix

---------

Co-authored-by: Chris <66376200+crickman@users.noreply.github.com>
This commit is contained in:
Eduard van Valkenburg
2025-09-10 16:52:42 +02:00
committed by GitHub
Unverified
parent 6aa746d891
commit 82ca4065cb
45 changed files with 3246 additions and 2984 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/packages/main/tests/main/conftest.py
+222 -18
View File
@@ -1,11 +1,64 @@
# Copyright (c) Microsoft. All rights reserved.
from typing import Any
from pydantic import BaseModel
import asyncio
import logging
import sys
from collections.abc import AsyncIterable, MutableSequence
from typing import Any
from unittest.mock import patch
from uuid import uuid4
from pydantic import BaseModel, Field
from pytest import fixture
from agent_framework import ChatMessage, ToolProtocol, ai_function
from agent_framework.telemetry import ModelDiagnosticSettings
from agent_framework import (
AgentProtocol,
AgentRunResponse,
AgentRunResponseUpdate,
AgentThread,
BaseChatClient,
ChatMessage,
ChatOptions,
ChatResponse,
ChatResponseUpdate,
Role,
TextContent,
ToolProtocol,
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
else:
from typing_extensions import override # type: ignore[import]
# region Chat History
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")
@@ -13,13 +66,16 @@ def chat_history() -> list[ChatMessage]:
return []
# region Tools
@fixture
def ai_tool() -> ToolProtocol:
"""Returns a generic ToolProtocol."""
class GenericTool(BaseModel):
name: str
description: str | None = None
description: str
additional_properties: dict[str, Any] | None = None
def parameters(self) -> dict[str, Any]:
@@ -43,17 +99,165 @@ def ai_function_tool() -> ToolProtocol:
return simple_function
# region Chat Clients
class MockChatClient:
"""Simple implementation of a chat client."""
def __init__(self) -> None:
self.additional_properties: dict[str, Any] = {}
async def get_response(
self,
messages: str | ChatMessage | list[str] | list[ChatMessage],
**kwargs: Any,
) -> ChatResponse:
logger.debug(f"Running custom chat client, with: {messages=}, {kwargs=}")
return ChatResponse(messages=ChatMessage(role="assistant", text="test response"))
async def get_streaming_response(
self,
messages: str | ChatMessage | list[str] | list[ChatMessage],
**kwargs: Any,
) -> AsyncIterable[ChatResponseUpdate]:
logger.debug(f"Running custom chat client stream, with: {messages=}, {kwargs=}")
yield ChatResponseUpdate(text=TextContent(text="test streaming response "), role="assistant")
yield ChatResponseUpdate(contents=[TextContent(text="another update")], role="assistant")
class MockBaseChatClient(BaseChatClient):
"""Mock implementation of the BaseChatClient."""
run_responses: list[ChatResponse] = Field(default_factory=list)
streaming_responses: list[list[ChatResponseUpdate]] = Field(default_factory=list)
@override
async def _inner_get_response(
self,
*,
messages: MutableSequence[ChatMessage],
chat_options: ChatOptions,
**kwargs: Any,
) -> ChatResponse:
"""Send a chat request to the AI service.
Args:
messages: The chat messages to send.
chat_options: The options for the request.
kwargs: Any additional keyword arguments.
Returns:
The chat response contents representing the response(s).
"""
logger.debug(f"Running base chat client inner, with: {messages=}, {chat_options=}, {kwargs=}")
if not self.run_responses:
return ChatResponse(messages=ChatMessage(role="assistant", text=f"test response - {messages[0].text}"))
if chat_options.tool_choice == "none":
return ChatResponse(
messages=ChatMessage(role="assistant", text="I broke out of the function invocation loop...")
)
return self.run_responses.pop(0)
@override
async def _inner_get_streaming_response(
self,
*,
messages: MutableSequence[ChatMessage],
chat_options: ChatOptions,
**kwargs: Any,
) -> AsyncIterable[ChatResponseUpdate]:
logger.debug(f"Running base chat client inner stream, with: {messages=}, {chat_options=}, {kwargs=}")
if not self.streaming_responses:
yield ChatResponseUpdate(text=f"update - {messages[0].text}", role="assistant")
return
if chat_options.tool_choice == "none":
yield ChatResponseUpdate(text="I broke out of the function invocation loop...", role="assistant")
return
response = self.streaming_responses.pop(0)
for update in response:
yield update
await asyncio.sleep(0)
@fixture
def model_diagnostic_settings(monkeypatch, request) -> ModelDiagnosticSettings:
"""Fixture to set environment variables for ModelDiagnosticSettings."""
enabled = getattr(request, "param", (None, None))[0]
sensitive = getattr(request, "param", (None, None))[1]
if enabled is None:
monkeypatch.delenv("AGENT_FRAMEWORK_GENAI_ENABLE_OTEL_DIAGNOSTICS", raising=False)
else:
monkeypatch.setenv("AGENT_FRAMEWORK_GENAI_ENABLE_OTEL_DIAGNOSTICS", str(enabled).lower())
if sensitive is None:
monkeypatch.delenv("AGENT_FRAMEWORK_GENAI_ENABLE_OTEL_DIAGNOSTICS_SENSITIVE", raising=False)
else:
monkeypatch.setenv("AGENT_FRAMEWORK_GENAI_ENABLE_OTEL_DIAGNOSTICS_SENSITIVE", str(sensitive).lower())
return ModelDiagnosticSettings(env_file_path="test.env")
def enable_function_calling(request: Any) -> bool:
return request.param if hasattr(request, "param") else True
@fixture
def max_iterations(request: Any) -> int:
return request.param if hasattr(request, "param") else 2
@fixture
def chat_client(enable_function_calling: bool, max_iterations: int) -> MockChatClient:
if enable_function_calling:
with patch("agent_framework._tools.DEFAULT_MAX_ITERATIONS", max_iterations):
return use_function_invocation(MockChatClient)()
return MockChatClient()
@fixture
def chat_client_base(enable_function_calling: bool, max_iterations: int) -> MockBaseChatClient:
if enable_function_calling:
with patch("agent_framework._tools.DEFAULT_MAX_ITERATIONS", max_iterations):
return use_function_invocation(MockBaseChatClient)()
return MockBaseChatClient()
# region Agents
class MockAgentThread(AgentThread):
pass
# Mock Agent implementation for testing
class MockAgent(AgentProtocol):
@property
def id(self) -> str:
return str(uuid4())
@property
def name(self) -> str | None:
"""Returns the name of the agent."""
return "Name"
@property
def display_name(self) -> str:
"""Returns the name of the agent."""
return "Display Name"
@property
def description(self) -> str | None:
return "Description"
async def run(
self,
messages: str | ChatMessage | list[str] | list[ChatMessage] | None = None,
*,
thread: AgentThread | None = None,
**kwargs: Any,
) -> AgentRunResponse:
logger.debug(f"Running mock agent, with: {messages=}, {thread=}, {kwargs=}")
return AgentRunResponse(messages=[ChatMessage(role=Role.ASSISTANT, contents=[TextContent("Response")])])
async def run_stream(
self,
messages: str | ChatMessage | list[str] | list[ChatMessage] | None = None,
*,
thread: AgentThread | None = None,
**kwargs: Any,
) -> AsyncIterable[AgentRunResponseUpdate]:
logger.debug(f"Running mock agent stream, with: {messages=}, {thread=}, {kwargs=}")
yield AgentRunResponseUpdate(contents=[TextContent("Response")])
def get_new_thread(self) -> AgentThread:
return MockAgentThread()
@fixture
def agent_thread() -> AgentThread:
return MockAgentThread()
@fixture
def agent() -> AgentProtocol:
return MockAgent()
python/packages/main/tests/main/test_agents.py
+19 -111
View File
@@ -1,122 +1,27 @@
# Copyright (c) Microsoft. All rights reserved.
from collections.abc import AsyncIterable, MutableSequence
from typing import Any
from collections.abc import AsyncIterable
from uuid import uuid4
from pytest import fixture, raises
from pytest import raises
from agent_framework import (
AgentProtocol,
AgentRunResponse,
AgentRunResponseUpdate,
AgentThread,
BaseChatClient,
ChatAgent,
ChatClientProtocol,
ChatMessage,
ChatMessageList,
ChatOptions,
ChatResponse,
ChatResponseUpdate,
HostedCodeInterpreterTool,
Role,
TextContent,
)
from agent_framework.exceptions import AgentExecutionException
# Mock AgentThread implementation for testing
class MockAgentThread(AgentThread):
pass
# Mock Agent implementation for testing
class MockAgent(AgentProtocol):
@property
def id(self) -> str:
return str(uuid4())
@property
def name(self) -> str | None:
"""Returns the name of the agent."""
return "Name"
@property
def display_name(self) -> str:
"""Returns the name of the agent."""
return "Display Name"
@property
def description(self) -> str | None:
return "Description"
async def run(
self,
messages: str | ChatMessage | list[str] | list[ChatMessage] | None = None,
*,
thread: AgentThread | None = None,
**kwargs: Any,
) -> AgentRunResponse:
return AgentRunResponse(messages=[ChatMessage(role=Role.ASSISTANT, contents=[TextContent("Response")])])
async def run_stream(
self,
messages: str | ChatMessage | list[str] | list[ChatMessage] | None = None,
*,
thread: AgentThread | None = None,
**kwargs: Any,
) -> AsyncIterable[AgentRunResponseUpdate]:
yield AgentRunResponseUpdate(contents=[TextContent("Response")])
def get_new_thread(self) -> AgentThread:
return MockAgentThread()
# Mock ChatClientProtocol implementation for testing
class MockChatClient(BaseChatClient):
_mock_response: ChatResponse | None = None
def __init__(self, mock_response: ChatResponse | None = None) -> None:
self._mock_response = mock_response
async def _inner_get_response(
self,
*,
messages: MutableSequence[ChatMessage],
chat_options: ChatOptions,
**kwargs: Any,
) -> ChatResponse:
return (
self._mock_response
if self._mock_response
else ChatResponse(messages=ChatMessage(role=Role.ASSISTANT, text="test response"))
)
async def _inner_get_streaming_response(
self,
*,
messages: MutableSequence[ChatMessage],
chat_options: ChatOptions,
**kwargs: Any,
) -> AsyncIterable[ChatResponseUpdate]:
yield ChatResponseUpdate(role=Role.ASSISTANT, text=TextContent(text="test streaming response"))
@fixture
def agent_thread() -> AgentThread:
return MockAgentThread()
@fixture
def agent() -> AgentProtocol:
return MockAgent()
@fixture
def chat_client() -> BaseChatClient:
return MockChatClient()
def test_agent_thread_type(agent_thread: AgentThread) -> None:
assert isinstance(agent_thread, AgentThread)
@@ -178,7 +83,7 @@ async def test_chat_client_agent_run_streaming(chat_client: ChatClientProtocol)
result = await AgentRunResponse.from_agent_response_generator(agent.run_stream("Hello"))
assert result.text == "test streaming response"
assert result.text == "test streaming response another update"
async def test_chat_client_agent_get_new_thread(chat_client: ChatClientProtocol) -> None:
@@ -203,14 +108,16 @@ async def test_chat_client_agent_prepare_thread_and_messages(chat_client: ChatCl
assert result_messages[1].text == "Test"
async def test_chat_client_agent_update_thread_id() -> None:
chat_client = MockChatClient(
mock_response=ChatResponse(
messages=[ChatMessage(role=Role.ASSISTANT, contents=[TextContent("test response")])],
conversation_id="123",
)
async def test_chat_client_agent_update_thread_id(chat_client_base: ChatClientProtocol) -> None:
mock_response = ChatResponse(
messages=[ChatMessage(role=Role.ASSISTANT, contents=[TextContent("test response")])],
conversation_id="123",
)
chat_client_base.run_responses = [mock_response]
agent = ChatAgent(
chat_client=chat_client_base,
tools=HostedCodeInterpreterTool(),
)
agent = ChatAgent(chat_client=chat_client)
thread = agent.get_new_thread()
result = await agent.run("Hello", thread=thread)
@@ -263,15 +170,16 @@ async def test_chat_client_agent_author_name_as_agent_name(chat_client: ChatClie
assert result.messages[0].author_name == "TestAgent"
async def test_chat_client_agent_author_name_is_used_from_response() -> None:
chat_client = MockChatClient(
mock_response=ChatResponse(
async def test_chat_client_agent_author_name_is_used_from_response(chat_client_base: ChatClientProtocol) -> None:
chat_client_base.run_responses = [
ChatResponse(
messages=[
ChatMessage(role=Role.ASSISTANT, contents=[TextContent("test response")], author_name="TestAuthor")
]
)
)
agent = ChatAgent(chat_client=chat_client)
]
agent = ChatAgent(chat_client=chat_client_base, tools=HostedCodeInterpreterTool())
result = await agent.run("Hello")
assert result.text == "test response"
python/packages/main/tests/main/test_clients.py
+32 -134
View File
@@ -1,18 +1,15 @@
# Copyright (c) Microsoft. All rights reserved.
import asyncio
import sys
from collections.abc import AsyncIterable, MutableSequence, Sequence
from collections.abc import Sequence
from typing import Any
from pydantic import Field
from pytest import fixture
from agent_framework import (
BaseChatClient,
ChatClientProtocol,
ChatMessage,
ChatOptions,
ChatResponse,
ChatResponseUpdate,
EmbeddingGenerator,
@@ -22,81 +19,12 @@ from agent_framework import (
Role,
TextContent,
ai_function,
use_tool_calling,
)
if sys.version_info >= (3, 12):
from typing import override # type: ignore
pass # type: ignore
else:
from typing_extensions import override # type: ignore[import]
class MockChatClient:
"""Simple implementation of a chat client."""
async def get_response(
self,
messages: ChatMessage | Sequence[ChatMessage],
**kwargs: Any,
) -> ChatResponse:
# Implement the method
return ChatResponse(messages=ChatMessage(role="assistant", text="test response"))
async def get_streaming_response(
self,
messages: ChatMessage | Sequence[ChatMessage],
**kwargs: Any,
) -> AsyncIterable[ChatResponseUpdate]:
# Implement the method
yield ChatResponseUpdate(text=TextContent(text="test streaming response"), role="assistant")
yield ChatResponseUpdate(contents=[TextContent(text="another update")], role="assistant")
@use_tool_calling
class MockBaseChatClient(BaseChatClient):
"""Mock implementation of the BaseChatClient."""
run_responses: list[ChatResponse] = Field(default_factory=list)
streaming_responses: list[list[ChatResponseUpdate]] = Field(default_factory=list)
@override
async def _inner_get_response(
self,
*,
messages: MutableSequence[ChatMessage],
chat_options: ChatOptions,
**kwargs: Any,
) -> ChatResponse:
"""Send a chat request to the AI service.
Args:
messages: The chat messages to send.
chat_options: The options for the request.
kwargs: Any additional keyword arguments.
Returns:
The chat response contents representing the response(s).
"""
if not self.run_responses or chat_options.tool_choice == "none":
return ChatResponse(messages=ChatMessage(role="assistant", text=f"test response - {messages[0].text}"))
return self.run_responses.pop(0)
@override
async def _inner_get_streaming_response(
self,
*,
messages: MutableSequence[ChatMessage],
chat_options: ChatOptions,
**kwargs: Any,
) -> AsyncIterable[ChatResponseUpdate]:
if not self.streaming_responses or chat_options.tool_choice == "none":
yield ChatResponseUpdate(text=f"update - {messages[0].text}", role="assistant")
return
response = self.streaming_responses.pop(0)
for update in response:
yield update
await asyncio.sleep(0)
pass # type: ignore[import]
class MockEmbeddingGenerator:
@@ -114,35 +42,25 @@ class MockEmbeddingGenerator:
return embeddings
@fixture
def chat_client() -> MockChatClient:
return MockChatClient()
@fixture
def chat_client_base() -> MockBaseChatClient:
return MockBaseChatClient()
@fixture
def embedding_generator() -> MockEmbeddingGenerator:
gen: EmbeddingGenerator[str, list[float]] = MockEmbeddingGenerator()
return gen
def test_chat_client_type(chat_client: MockChatClient):
def test_chat_client_type(chat_client: ChatClientProtocol):
assert isinstance(chat_client, ChatClientProtocol)
async def test_chat_client_get_response(chat_client: MockChatClient):
async def test_chat_client_get_response(chat_client: ChatClientProtocol):
response = await chat_client.get_response(ChatMessage(role="user", text="Hello"))
assert response.text == "test response"
assert response.messages[0].role == Role.ASSISTANT
async def test_chat_client_get_streaming_response(chat_client: MockChatClient):
async def test_chat_client_get_streaming_response(chat_client: ChatClientProtocol):
async for update in chat_client.get_streaming_response(ChatMessage(role="user", text="Hello")):
assert update.text == "test streaming response" or update.text == "another update"
assert update.text == "test streaming response " or update.text == "another update"
assert update.role == Role.ASSISTANT
@@ -158,23 +76,23 @@ async def test_embedding_generator_generate(embedding_generator: MockEmbeddingGe
assert len(emb) == 5
def test_base_client(chat_client_base: MockBaseChatClient):
def test_base_client(chat_client_base: ChatClientProtocol):
assert isinstance(chat_client_base, BaseChatClient)
assert isinstance(chat_client_base, ChatClientProtocol)
async def test_base_client_get_response(chat_client_base: MockBaseChatClient):
async def test_base_client_get_response(chat_client_base: ChatClientProtocol):
response = await chat_client_base.get_response(ChatMessage(role="user", text="Hello"))
assert response.messages[0].role == Role.ASSISTANT
assert response.messages[0].text == "test response - Hello"
async def test_base_client_get_streaming_response(chat_client_base: MockBaseChatClient):
async def test_base_client_get_streaming_response(chat_client_base: ChatClientProtocol):
async for update in chat_client_base.get_streaming_response(ChatMessage(role="user", text="Hello")):
assert update.text == "update - Hello" or update.text == "another update"
async def test_base_client_with_function_calling(chat_client_base: MockBaseChatClient):
async def test_base_client_with_function_calling(chat_client_base: ChatClientProtocol):
exec_counter = 0
@ai_function(name="test_function")
@@ -208,8 +126,7 @@ async def test_base_client_with_function_calling(chat_client_base: MockBaseChatC
assert response.messages[2].text == "done"
async def test_base_client_with_function_calling_disabled(chat_client_base: MockBaseChatClient):
chat_client_base.__maximum_iterations_per_request = 0
async def test_base_client_with_function_calling_resets(chat_client_base: ChatClientProtocol):
exec_counter = 0
@ai_function(name="test_function")
@@ -225,16 +142,32 @@ async def test_base_client_with_function_calling_disabled(chat_client_base: Mock
contents=[FunctionCallContent(call_id="1", name="test_function", arguments='{"arg1": "value1"}')],
)
),
ChatResponse(
messages=ChatMessage(
role="assistant",
contents=[FunctionCallContent(call_id="2", name="test_function", arguments='{"arg1": "value1"}')],
)
),
ChatResponse(messages=ChatMessage(role="assistant", text="done")),
]
response = await chat_client_base.get_response("hello", tool_choice="auto", tools=[ai_func])
assert exec_counter == 0
assert len(response.messages) == 1
assert exec_counter == 2
assert len(response.messages) == 5
assert response.messages[0].role == Role.ASSISTANT
assert response.messages[0].text == "test response - hello"
assert response.messages[1].role == Role.TOOL
assert response.messages[2].role == Role.ASSISTANT
assert response.messages[3].role == Role.TOOL
assert response.messages[4].role == Role.ASSISTANT
assert isinstance(response.messages[0].contents[0], FunctionCallContent)
assert isinstance(response.messages[1].contents[0], FunctionResultContent)
assert isinstance(response.messages[2].contents[0], FunctionCallContent)
assert isinstance(response.messages[3].contents[0], FunctionResultContent)
# after these two responses, it would try another regular call, but since max_iterations is 1, it stops and calls
assert isinstance(response.messages[4].contents[0], TextContent)
assert response.text == "I broke out of the function invocation loop..."
async def test_base_client_with_streaming_function_calling(chat_client_base: MockBaseChatClient):
async def test_base_client_with_streaming_function_calling(chat_client_base: ChatClientProtocol):
exec_counter = 0
@ai_function(name="test_function")
@@ -270,38 +203,3 @@ async def test_base_client_with_streaming_function_calling(chat_client_base: Moc
assert updates[2].contents[0].call_id == "1"
assert updates[3].text == "Processed value1"
assert exec_counter == 1
async def test_base_client_with_streaming_function_calling_disabled(chat_client_base: MockBaseChatClient):
chat_client_base.__maximum_iterations_per_request = 0
exec_counter = 0
@ai_function(name="test_function")
def ai_func(arg1: str) -> str:
nonlocal exec_counter
exec_counter += 1
return f"Processed {arg1}"
chat_client_base.streaming_responses = [
[
ChatResponseUpdate(
contents=[FunctionCallContent(call_id="1", name="test_function", arguments='{"arg1":')],
role="assistant",
),
ChatResponseUpdate(
contents=[FunctionCallContent(call_id="1", name="test_function", arguments='"value1"}')],
role="assistant",
),
],
[
ChatResponseUpdate(
contents=[TextContent(text="Processed value1")],
role="assistant",
)
],
]
updates = []
async for update in chat_client_base.get_streaming_response("hello", tool_choice="auto", tools=[ai_func]):
updates.append(update)
assert len(updates) == 1
assert exec_counter == 0
python/packages/main/tests/main/test_telemetry.py
+154 -578
View File
@@ -1,14 +1,19 @@
# Copyright (c) Microsoft. All rights reserved.
import logging
from collections.abc import AsyncIterable, MutableSequence
from collections.abc import MutableSequence
from typing import Any
from unittest.mock import Mock, patch
from unittest.mock import MagicMock, Mock, patch
import pytest
from opentelemetry.semconv_ai import SpanAttributes
from opentelemetry.trace import StatusCode
from agent_framework import (
AgentProtocol,
AgentRunResponse,
AgentThread,
BaseChatClient,
ChatMessage,
ChatOptions,
ChatResponse,
@@ -16,15 +21,19 @@ from agent_framework import (
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,
TELEMETRY_DISABLED_ENV_VAR,
USER_AGENT_KEY,
USER_AGENT_TELEMETRY_DISABLED_ENV_VAR,
ChatMessageListTimestampFilter,
GenAIAttributes,
OtelAttr,
get_function_span,
prepend_agent_framework_to_user_agent,
start_as_current_span,
use_agent_telemetry,
use_telemetry,
)
@@ -33,7 +42,7 @@ from agent_framework.telemetry import (
def test_telemetry_disabled_env_var():
"""Test that the telemetry disabled environment variable is correctly defined."""
assert TELEMETRY_DISABLED_ENV_VAR == "AZURE_TELEMETRY_DISABLED"
assert USER_AGENT_TELEMETRY_DISABLED_ENV_VAR == "AGENT_FRAMEWORK_USER_AGENT_DISABLED"
def test_user_agent_key():
@@ -78,20 +87,20 @@ def test_app_info_when_telemetry_disabled():
def test_role_event_map():
"""Test that ROLE_EVENT_MAP contains expected mappings."""
assert ROLE_EVENT_MAP["system"] == GenAIAttributes.SYSTEM_MESSAGE.value
assert ROLE_EVENT_MAP["user"] == GenAIAttributes.USER_MESSAGE.value
assert ROLE_EVENT_MAP["assistant"] == GenAIAttributes.ASSISTANT_MESSAGE.value
assert ROLE_EVENT_MAP["tool"] == GenAIAttributes.TOOL_MESSAGE.value
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 GenAIAttributes enum has expected values."""
assert GenAIAttributes.OPERATION.value == "gen_ai.operation.name"
assert GenAIAttributes.SYSTEM.value == "gen_ai.system"
assert GenAIAttributes.MODEL.value == "gen_ai.request.model"
assert GenAIAttributes.CHAT_COMPLETION_OPERATION.value == "chat"
assert GenAIAttributes.TOOL_EXECUTION_OPERATION.value == "execute_tool"
assert GenAIAttributes.AGENT_INVOKE_OPERATION.value == "invoke_agent"
"""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
@@ -135,47 +144,6 @@ def test_modifies_original_dict():
assert "User-Agent" in headers
# region ModelDiagnosticSettings tests
@pytest.mark.parametrize("model_diagnostic_settings", [(None, None)], indirect=True)
def test_default_values(model_diagnostic_settings):
"""Test default values for ModelDiagnosticSettings."""
assert not model_diagnostic_settings.ENABLED
assert not model_diagnostic_settings.SENSITIVE_EVENTS_ENABLED
@pytest.mark.parametrize("model_diagnostic_settings", [(False, False)], indirect=True)
def test_disabled(model_diagnostic_settings):
"""Test default values for ModelDiagnosticSettings."""
assert not model_diagnostic_settings.ENABLED
assert not model_diagnostic_settings.SENSITIVE_EVENTS_ENABLED
@pytest.mark.parametrize("model_diagnostic_settings", [(True, False)], indirect=True)
def test_non_sensitive_events_enabled(model_diagnostic_settings):
"""Test loading model_diagnostic_settings from environment variables."""
assert model_diagnostic_settings.ENABLED
assert not model_diagnostic_settings.SENSITIVE_EVENTS_ENABLED
@pytest.mark.parametrize("model_diagnostic_settings", [(True, True)], indirect=True)
def test_sensitive_events_enabled(model_diagnostic_settings):
"""Test loading model_diagnostic_settings from environment variables."""
assert model_diagnostic_settings.ENABLED
assert model_diagnostic_settings.SENSITIVE_EVENTS_ENABLED
@pytest.mark.parametrize("model_diagnostic_settings", [(False, True)], indirect=True)
def test_sensitive_events_enabled_only(model_diagnostic_settings):
"""Test loading sensitive events setting from environment.
But when sensitive events are enabled, diagnostics are also enabled.
"""
assert model_diagnostic_settings.ENABLED
assert model_diagnostic_settings.SENSITIVE_EVENTS_ENABLED
# region Test ChatMessageListTimestampFilter
@@ -213,88 +181,74 @@ def test_filter_with_index_key():
def test_index_key_constant():
"""Test that INDEX_KEY constant is correctly defined."""
assert ChatMessageListTimestampFilter.INDEX_KEY == "CHAT_MESSAGE_INDEX"
assert ChatMessageListTimestampFilter.INDEX_KEY == "chat_message_index"
# region Test start_as_current_span
# region Test get_function_span
def test_start_span_basic():
"""Test starting a span with basic function info."""
mock_tracer = Mock()
mock_span = Mock()
mock_tracer.start_as_current_span.return_value = mock_span
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"
# Create a mock function
mock_function = Mock()
mock_function.name = "test_function"
mock_function.description = "Test function description"
result = start_as_current_span(mock_tracer, mock_function)
result = get_function_span(mock_function)
assert result == mock_span
mock_tracer.start_as_current_span.assert_called_once()
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[0][0] == "execute_tool test_function"
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[GenAIAttributes.OPERATION.value] == GenAIAttributes.TOOL_EXECUTION_OPERATION.value
assert attributes[GenAIAttributes.TOOL_NAME.value] == "test_function"
assert attributes[GenAIAttributes.TOOL_DESCRIPTION.value] == "Test function description"
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_metadata():
"""Test starting a span with metadata containing tool_call_id."""
def test_start_span_with_tool_call_id():
"""Test starting a span with tool_call_id."""
mock_tracer = Mock()
mock_span = Mock()
mock_tracer.start_as_current_span.return_value = mock_span
with patch("agent_framework.telemetry.tracer", mock_tracer):
mock_span = Mock()
mock_tracer.start_as_current_span.return_value = mock_span
mock_function = Mock()
mock_function.name = "test_function"
mock_function.description = "Test function"
mock_function = Mock()
mock_function.name = "test_function"
mock_function.description = "Test function"
metadata = {"tool_call_id": "test_call_123"}
tool_call_id = "test_call_123"
_ = start_as_current_span(mock_tracer, mock_function, metadata)
_ = get_function_span(mock_function, tool_call_id)
call_args = mock_tracer.start_as_current_span.call_args
attributes = call_args[1]["attributes"]
assert attributes[GenAIAttributes.TOOL_CALL_ID.value] == "test_call_123"
call_args = mock_tracer.start_as_current_span.call_args
attributes = call_args[1]["attributes"]
assert attributes[OtelAttr.TOOL_CALL_ID] == "test_call_123"
def test_start_span_without_description():
"""Test starting a span when function has no description."""
mock_tracer = Mock()
mock_span = Mock()
mock_tracer.start_as_current_span.return_value = mock_span
with patch("agent_framework.telemetry.tracer", mock_tracer):
mock_span = Mock()
mock_tracer.start_as_current_span.return_value = mock_span
mock_function = Mock()
mock_function.name = "test_function"
mock_function.description = None
mock_function = Mock()
mock_function.name = "test_function"
mock_function.description = None
start_as_current_span(mock_tracer, mock_function)
get_function_span(mock_function)
call_args = mock_tracer.start_as_current_span.call_args
attributes = call_args[1]["attributes"]
assert GenAIAttributes.TOOL_DESCRIPTION.value not in attributes
def test_start_span_empty_metadata():
"""Test starting a span with empty metadata."""
mock_tracer = Mock()
mock_span = Mock()
mock_tracer.start_as_current_span.return_value = mock_span
mock_function = Mock()
mock_function.name = "test_function"
mock_function.description = "Test function"
start_as_current_span(mock_tracer, mock_function, {})
call_args = mock_tracer.start_as_current_span.call_args
attributes = call_args[1]["attributes"]
assert GenAIAttributes.TOOL_CALL_ID.value not in attributes
call_args = mock_tracer.start_as_current_span.call_args
attributes = call_args[1]["attributes"]
assert OtelAttr.TOOL_DESCRIPTION not in attributes
# region Test use_telemetry decorator
@@ -305,12 +259,10 @@ def test_decorator_with_valid_class():
# Create a mock class with the required methods
class MockChatClient:
MODEL_PROVIDER_NAME = "test_provider"
async def _inner_get_response(self, *, messages, chat_options, **kwargs):
async def get_response(self, messages, **kwargs):
return Mock()
async def _inner_get_streaming_response(self, *, messages, chat_options, **kwargs):
async def get_streaming_response(self, messages, **kwargs):
async def gen():
yield Mock()
@@ -318,39 +270,31 @@ def test_decorator_with_valid_class():
# Apply the decorator
decorated_class = use_telemetry(MockChatClient)
# Check that the methods were wrapped
assert hasattr(decorated_class._inner_get_response, "__model_diagnostics_chat_client__")
assert hasattr(decorated_class._inner_get_streaming_response, "__model_diagnostics_streaming_chat_completion__")
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:
MODEL_PROVIDER_NAME = "test_provider"
OTEL_PROVIDER_NAME = "test_provider"
# Apply the decorator - should not raise an error
decorated_class = use_telemetry(MockChatClient)
# Class should be returned unchanged
assert decorated_class is MockChatClient
with pytest.raises(ChatClientInitializationError):
use_telemetry(MockChatClient)
def test_decorator_with_partial_methods():
"""Test decorator when only one method is present."""
class MockChatClient:
MODEL_PROVIDER_NAME = "test_provider"
OTEL_PROVIDER_NAME = "test_provider"
async def _inner_get_response(self, *, messages, chat_options, **kwargs):
async def get_response(self, messages, **kwargs):
return Mock()
decorated_class = use_telemetry(MockChatClient)
# Only the present method should be wrapped
assert hasattr(decorated_class._inner_get_response, "__model_diagnostics_chat_client__")
assert not hasattr(decorated_class, "_inner_get_streaming_response")
with pytest.raises(ChatClientInitializationError):
use_telemetry(MockChatClient)
# region Test telemetry decorator with mock client
@@ -360,12 +304,7 @@ def test_decorator_with_partial_methods():
def mock_chat_client():
"""Create a mock chat client for testing."""
class MockChatClient:
MODEL_PROVIDER_NAME = "test_provider"
def __init__(self):
self.ai_model_id = "test-model"
class MockChatClient(BaseChatClient):
def service_url(self):
return "https://test.example.com"
@@ -384,223 +323,77 @@ def mock_chat_client():
yield ChatResponseUpdate(text="Hello", role=Role.ASSISTANT)
yield ChatResponseUpdate(text=" world", role=Role.ASSISTANT)
return MockChatClient()
return MockChatClient
@pytest.mark.parametrize("model_diagnostic_settings", [(False, False)], indirect=True)
async def test_telemetry_disabled_bypasses_instrumentation(mock_chat_client, model_diagnostic_settings):
"""Test that when diagnostics are disabled, telemetry is bypassed."""
decorated_class = use_telemetry(type(mock_chat_client))
client = decorated_class()
messages = [ChatMessage(role=Role.USER, text="Test message")]
chat_options = ChatOptions()
with (
patch("agent_framework.telemetry.MODEL_DIAGNOSTICS_SETTINGS", model_diagnostic_settings),
patch("agent_framework.telemetry.use_span") as mock_use_span,
):
# This should not create any spans
response = await client._inner_get_response(messages=messages, chat_options=chat_options)
assert response is not None
mock_use_span.assert_not_called()
@pytest.mark.parametrize("model_diagnostic_settings", [(True, True)], indirect=True)
async def test_instrumentation_enabled(mock_chat_client, model_diagnostic_settings):
@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."""
decorated_class = use_telemetry(type(mock_chat_client))
client = decorated_class()
client = use_telemetry(mock_chat_client)()
messages = [ChatMessage(role=Role.USER, text="Test message")]
chat_options = ChatOptions()
with (
patch("agent_framework.telemetry.MODEL_DIAGNOSTICS_SETTINGS", model_diagnostic_settings),
patch("agent_framework.telemetry.use_span") as mock_use_span,
patch("agent_framework.telemetry.logger") as mock_logger,
patch("agent_framework.telemetry._get_span") as mock_response_span,
patch("agent_framework.telemetry._capture_messages") as mock_log_messages,
):
response = await client._inner_get_response(messages=messages, chat_options=chat_options)
response = await client.get_response(messages=messages, chat_options=chat_options)
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
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("model_diagnostic_settings", [(True, False)], indirect=True)
async def test_streaming_response_with_diagnostics_enabled_via_decorator(mock_chat_client, model_diagnostic_settings):
@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."""
decorated_class = use_telemetry(type(mock_chat_client))
client = decorated_class()
client = use_telemetry(mock_chat_client)()
messages = [ChatMessage(role=Role.USER, text="Test")]
chat_options = ChatOptions()
with (
patch("agent_framework.telemetry.MODEL_DIAGNOSTICS_SETTINGS", model_diagnostic_settings),
patch("agent_framework.telemetry.use_span") as mock_use_span,
patch("agent_framework.telemetry._get_chat_response_span") as mock_get_span,
patch("agent_framework.telemetry._set_chat_response_input") as mock_set_input,
patch("agent_framework.telemetry._set_chat_response_output") as mock_set_output,
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,
):
mock_span = Mock()
mock_use_span.return_value.__enter__.return_value = mock_span
mock_use_span.return_value.__exit__.return_value = None
# We can't easily mock ChatResponse.from_chat_response_updates since it's imported locally,
# but we can verify telemetry calls were made
# Collect all yielded updates
updates = []
async for update in client._inner_get_streaming_response(messages=messages, chat_options=chat_options):
async for update in client.get_streaming_response(messages=messages, chat_options=chat_options):
updates.append(update)
# Verify we got the expected updates
# Verify we got the expected updates, this shouldn't be dependent on otel
assert len(updates) == 2
# Verify telemetry calls were made
mock_get_span.assert_called_once()
mock_set_input.assert_called_once_with("test_provider", messages)
mock_set_output.assert_called_once()
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()
@pytest.mark.parametrize("model_diagnostic_settings", [(True, False)], indirect=True)
async def test_streaming_response_with_exception_via_decorator(mock_chat_client, model_diagnostic_settings):
"""Test streaming telemetry exception handling through decorator."""
async def _inner_get_streaming_response(
self, *, messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
) -> AsyncIterable[ChatResponseUpdate]:
yield ChatResponseUpdate(text="Partial", role=Role.ASSISTANT)
raise ValueError("Test streaming error")
type(mock_chat_client)._inner_get_streaming_response = _inner_get_streaming_response
decorated_class = use_telemetry(type(mock_chat_client))
client = decorated_class()
messages = [ChatMessage(role=Role.USER, text="Test")]
chat_options = ChatOptions()
with (
patch("agent_framework.telemetry.MODEL_DIAGNOSTICS_SETTINGS", model_diagnostic_settings),
patch("agent_framework.telemetry.use_span") as mock_use_span,
patch("agent_framework.telemetry._get_chat_response_span"),
patch("agent_framework.telemetry._set_chat_response_input"),
patch("agent_framework.telemetry._set_error") as mock_set_error,
):
mock_span = Mock()
mock_use_span.return_value.__enter__.return_value = mock_span
mock_use_span.return_value.__exit__.return_value = None
# Should raise the exception and call error handler
with pytest.raises(ValueError, match="Test streaming error"):
async for _ in client._inner_get_streaming_response(messages=messages, chat_options=chat_options):
pass
# Verify error was recorded
mock_set_error.assert_called_once()
assert isinstance(mock_set_error.call_args[0][1], ValueError)
@pytest.mark.parametrize("model_diagnostic_settings", [(False, False)], indirect=True)
async def test_streaming_response_diagnostics_disabled_via_decorator(model_diagnostic_settings):
"""Test streaming response when diagnostics are disabled."""
from agent_framework import ChatResponseUpdate
class MockStreamingClientNoDiagnostics:
MODEL_PROVIDER_NAME = "test_provider"
async def _inner_get_streaming_response(
self, *, messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
) -> AsyncIterable[ChatResponseUpdate]:
yield ChatResponseUpdate(text="Test", role=Role.ASSISTANT)
decorated_class = use_telemetry(MockStreamingClientNoDiagnostics)
client = decorated_class()
messages = [ChatMessage(role=Role.USER, text="Test")]
chat_options = ChatOptions()
with (
patch("agent_framework.telemetry.MODEL_DIAGNOSTICS_SETTINGS", model_diagnostic_settings),
patch("agent_framework.telemetry._get_chat_response_span") as mock_get_span,
):
# Should not create spans when diagnostics are disabled
updates = []
async for update in client._inner_get_streaming_response(messages=messages, chat_options=chat_options):
updates.append(update)
assert len(updates) == 1
# Should not have called telemetry functions
mock_get_span.assert_not_called()
# region Test empty streaming response handling
@pytest.mark.parametrize("model_diagnostic_settings", [(True, False)], indirect=True)
async def test_empty_streaming_response_via_decorator(model_diagnostic_settings):
"""Test streaming wrapper with empty response."""
class MockEmptyStreamingClient:
MODEL_PROVIDER_NAME = "test_provider"
def __init__(self):
self.ai_model_id = "test_model"
def service_url(self) -> str:
return "https://test.com"
async def _inner_get_streaming_response(
self, *, messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
) -> AsyncIterable[ChatResponseUpdate]:
# Return empty stream
return
yield # This will never be reached
decorated_class = use_telemetry(MockEmptyStreamingClient)
client = decorated_class()
messages = [ChatMessage(role=Role.USER, text="Test")]
chat_options = ChatOptions()
with (
patch("agent_framework.telemetry.MODEL_DIAGNOSTICS_SETTINGS", model_diagnostic_settings),
patch("agent_framework.telemetry.use_span") as mock_use_span,
patch("agent_framework.telemetry._get_chat_response_span"),
patch("agent_framework.telemetry._set_chat_response_input"),
patch("agent_framework.telemetry._set_chat_response_output") as mock_set_output,
):
mock_span = Mock()
mock_use_span.return_value.__enter__.return_value = mock_span
mock_use_span.return_value.__exit__.return_value = None
# Should handle empty stream gracefully
updates = []
async for update in client._inner_get_streaming_response(messages=messages, chat_options=chat_options):
updates.append(update)
assert len(updates) == 0
# Should still call telemetry
mock_set_output.assert_called_once()
def test_start_as_current_span_with_none_metadata():
"""Test start_as_current_span with None metadata."""
"""Test get_function_span with None metadata."""
mock_tracer = Mock()
mock_span = Mock()
mock_tracer.start_as_current_span.return_value = mock_span
with patch("agent_framework.telemetry.tracer", mock_tracer):
mock_span = Mock()
mock_tracer.start_as_current_span.return_value = mock_span
mock_function = Mock()
mock_function.name = "test_function"
mock_function.description = "Test description"
mock_function = Mock()
mock_function.name = "test_function"
mock_function.description = "Test description"
result = start_as_current_span(mock_tracer, mock_function, None)
result = get_function_span(mock_function, None)
assert result == mock_span
call_args = mock_tracer.start_as_current_span.call_args
attributes = call_args[1]["attributes"]
assert GenAIAttributes.TOOL_CALL_ID.value not in attributes
assert result == mock_span
call_args = mock_tracer.start_as_current_span.call_args
attributes = call_args[1]["attributes"]
assert attributes[OtelAttr.TOOL_CALL_ID] == "unknown"
def test_prepend_user_agent_with_none_value():
@@ -618,7 +411,6 @@ def test_prepend_user_agent_with_none_value():
def test_agent_decorator_with_valid_class():
"""Test that agent decorator works with a valid ChatAgent-like class."""
from agent_framework.telemetry import use_agent_telemetry
# Create a mock class with the required methods
class MockChatClientAgent:
@@ -639,33 +431,31 @@ def test_agent_decorator_with_valid_class():
return gen()
def get_new_thread(self) -> AgentThread:
return AgentThread()
# Apply the decorator
decorated_class = use_agent_telemetry(MockChatClientAgent)
# Check that the methods were wrapped
assert hasattr(decorated_class.run, "__model_diagnostics_agent_run__")
assert hasattr(decorated_class.run_stream, "__model_diagnostics_streaming_agent_run__")
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."""
from agent_framework.telemetry import use_agent_telemetry
class MockChatClientAgent:
class MockAgent:
AGENT_SYSTEM_NAME = "test_agent_system"
# Apply the decorator - should not raise an error
decorated_class = use_agent_telemetry(MockChatClientAgent)
# Class should be returned unchanged
assert decorated_class is MockChatClientAgent
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 MockChatClientAgent:
class MockAgent:
AGENT_SYSTEM_NAME = "test_agent_system"
def __init__(self):
@@ -676,11 +466,8 @@ def test_agent_decorator_with_partial_methods():
async def run(self, messages=None, *, thread=None, **kwargs):
return Mock()
decorated_class = use_agent_telemetry(MockChatClientAgent)
# Only the present method should be wrapped
assert hasattr(decorated_class.run, "__model_diagnostics_agent_run__")
assert not hasattr(decorated_class, "run_stream")
with pytest.raises(AgentInitializationError):
use_agent_telemetry(MockAgent)
# region Test agent telemetry decorator with mock agent
@@ -689,7 +476,6 @@ def test_agent_decorator_with_partial_methods():
@pytest.fixture
def mock_chat_client_agent():
"""Create a mock chat client agent for testing."""
from agent_framework import AgentRunResponse, ChatMessage, Role, UsageDetails
class MockChatClientAgent:
AGENT_SYSTEM_NAME = "test_agent_system"
@@ -714,37 +500,16 @@ def mock_chat_client_agent():
yield AgentRunResponseUpdate(text="Hello", role=Role.ASSISTANT)
yield AgentRunResponseUpdate(text=" from agent", role=Role.ASSISTANT)
return MockChatClientAgent()
return MockChatClientAgent
@pytest.mark.parametrize("model_diagnostic_settings", [(False, False)], indirect=True)
async def test_agent_telemetry_disabled_bypasses_instrumentation(mock_chat_client_agent, model_diagnostic_settings):
"""Test that when agent diagnostics are disabled, telemetry is bypassed."""
from agent_framework.telemetry import use_agent_telemetry
decorated_class = use_agent_telemetry(type(mock_chat_client_agent))
agent = decorated_class()
with (
patch("agent_framework.telemetry.MODEL_DIAGNOSTICS_SETTINGS", model_diagnostic_settings),
patch("agent_framework.telemetry.use_span") as mock_use_span,
):
# This should not create any spans
response = await agent.run("Test message")
assert response is not None
mock_use_span.assert_not_called()
@pytest.mark.parametrize("model_diagnostic_settings", [(True, True)], indirect=True)
async def test_agent_instrumentation_enabled(mock_chat_client_agent, model_diagnostic_settings):
@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."""
from agent_framework.telemetry import use_agent_telemetry
decorated_class = use_agent_telemetry(type(mock_chat_client_agent))
agent = decorated_class()
agent = use_agent_telemetry(mock_chat_client_agent)()
with (
patch("agent_framework.telemetry.MODEL_DIAGNOSTICS_SETTINGS", model_diagnostic_settings),
patch("agent_framework.telemetry.use_span") as mock_use_span,
patch("agent_framework.telemetry.logger") as mock_logger,
):
@@ -752,30 +517,21 @@ async def test_agent_instrumentation_enabled(mock_chat_client_agent, model_diagn
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
assert mock_logger.info.call_count == (2 if otel_settings.enable_sensitive_data else 0)
@pytest.mark.parametrize("model_diagnostic_settings", [(True, False)], indirect=True)
@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, model_diagnostic_settings
mock_chat_client_agent: AgentProtocol, otel_settings
):
"""Test agent streaming telemetry through the use_agent_telemetry decorator."""
from agent_framework.telemetry import use_agent_telemetry
decorated_class = use_agent_telemetry(type(mock_chat_client_agent))
agent = decorated_class()
agent = use_agent_telemetry(mock_chat_client_agent)()
with (
patch("agent_framework.telemetry.MODEL_DIAGNOSTICS_SETTINGS", model_diagnostic_settings),
patch("agent_framework.telemetry.use_span") as mock_use_span,
patch("agent_framework.telemetry._get_agent_run_span") as mock_get_span,
patch("agent_framework.telemetry._set_agent_run_input") as mock_set_input,
patch("agent_framework.telemetry._set_agent_run_output") as mock_set_output,
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,
):
mock_span = Mock()
mock_use_span.return_value.__enter__.return_value = mock_span
mock_use_span.return_value.__exit__.return_value = None
# Collect all yielded updates
updates = []
async for update in agent.run_stream("Test message"):
@@ -786,219 +542,39 @@ async def test_agent_streaming_response_with_diagnostics_enabled_via_decorator(
# Verify telemetry calls were made
mock_get_span.assert_called_once()
mock_set_input.assert_called_once_with("test_agent_system", "Test message")
mock_set_output.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()
@pytest.mark.parametrize("model_diagnostic_settings", [(True, False)], indirect=True)
async def test_agent_streaming_response_with_exception_via_decorator(mock_chat_client_agent, model_diagnostic_settings):
"""Test agent streaming telemetry exception handling through decorator."""
from agent_framework.telemetry import use_agent_telemetry
async def run_stream(self, messages=None, *, thread=None, **kwargs):
from agent_framework import AgentRunResponseUpdate, Role
yield AgentRunResponseUpdate(text="Partial", role=Role.ASSISTANT)
raise ValueError("Test agent streaming error")
type(mock_chat_client_agent).run_stream = run_stream
decorated_class = use_agent_telemetry(type(mock_chat_client_agent))
agent = decorated_class()
with (
patch("agent_framework.telemetry.MODEL_DIAGNOSTICS_SETTINGS", model_diagnostic_settings),
patch("agent_framework.telemetry.use_span") as mock_use_span,
patch("agent_framework.telemetry._get_agent_run_span"),
patch("agent_framework.telemetry._set_agent_run_input"),
patch("agent_framework.telemetry._set_error") as mock_set_error,
):
mock_span = Mock()
mock_use_span.return_value.__enter__.return_value = mock_span
mock_use_span.return_value.__exit__.return_value = None
# Should raise the exception and call error handler
with pytest.raises(ValueError, match="Test agent streaming error"):
async for _ in agent.run_stream("Test message"):
pass
# Verify error was recorded
mock_set_error.assert_called_once()
assert isinstance(mock_set_error.call_args[0][1], ValueError)
@pytest.mark.parametrize("model_diagnostic_settings", [(False, False)], indirect=True)
async def test_agent_streaming_response_diagnostics_disabled_via_decorator(model_diagnostic_settings):
"""Test agent streaming response when diagnostics are disabled."""
from agent_framework import AgentRunResponseUpdate, Role
from agent_framework.telemetry import use_agent_telemetry
class MockStreamingAgentNoDiagnostics:
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_stream(self, messages=None, *, thread=None, **kwargs):
yield AgentRunResponseUpdate(text="Test", role=Role.ASSISTANT)
decorated_class = use_agent_telemetry(MockStreamingAgentNoDiagnostics)
agent = decorated_class()
with (
patch("agent_framework.telemetry.MODEL_DIAGNOSTICS_SETTINGS", model_diagnostic_settings),
patch("agent_framework.telemetry._get_agent_run_span") as mock_get_span,
):
# Should not create spans when diagnostics are disabled
updates = []
async for update in agent.run_stream("Test message"):
updates.append(update)
assert len(updates) == 1
# Should not have called telemetry functions
mock_get_span.assert_not_called()
@pytest.mark.parametrize("model_diagnostic_settings", [(True, False)], indirect=True)
async def test_agent_empty_streaming_response_via_decorator(model_diagnostic_settings):
"""Test agent streaming wrapper with empty response."""
from agent_framework.telemetry import use_agent_telemetry
class MockEmptyStreamingAgent:
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_stream(self, messages=None, *, thread=None, **kwargs):
# Return empty stream
return
yield # This will never be reached
decorated_class = use_agent_telemetry(MockEmptyStreamingAgent)
agent = decorated_class()
with (
patch("agent_framework.telemetry.MODEL_DIAGNOSTICS_SETTINGS", model_diagnostic_settings),
patch("agent_framework.telemetry.use_span") as mock_use_span,
patch("agent_framework.telemetry._get_agent_run_span"),
patch("agent_framework.telemetry._set_agent_run_input"),
patch("agent_framework.telemetry._set_agent_run_output") as mock_set_output,
):
mock_span = Mock()
mock_use_span.return_value.__enter__.return_value = mock_span
mock_use_span.return_value.__exit__.return_value = None
# Should handle empty stream gracefully
updates = []
async for update in agent.run_stream("Test message"):
updates.append(update)
assert len(updates) == 0
# Should still call telemetry
mock_set_output.assert_called_once()
@pytest.mark.parametrize("model_diagnostic_settings", [(True, True)], indirect=True)
async def test_agent_run_with_thread_and_kwargs(mock_chat_client_agent, model_diagnostic_settings):
"""Test agent run with thread and additional kwargs."""
from agent_framework.telemetry import use_agent_telemetry
decorated_class = use_agent_telemetry(type(mock_chat_client_agent))
agent = decorated_class()
# Mock thread
mock_thread = Mock()
mock_thread.id = "test_thread_id"
with (
patch("agent_framework.telemetry.MODEL_DIAGNOSTICS_SETTINGS", model_diagnostic_settings),
patch("agent_framework.telemetry.use_span") as mock_use_span,
patch("agent_framework.telemetry._get_agent_run_span") as mock_get_span,
):
mock_span = Mock()
mock_use_span.return_value.__enter__.return_value = mock_span
mock_use_span.return_value.__exit__.return_value = None
# Test with thread and additional kwargs
response = await agent.run(
"Test message", thread=mock_thread, temperature=0.7, max_tokens=100, model="test-model"
)
assert response is not None
# Verify the span was created with the correct parameters
mock_get_span.assert_called_once()
call_kwargs = mock_get_span.call_args[1]
assert call_kwargs["agent"] == agent
assert call_kwargs["thread"] == mock_thread
assert call_kwargs["temperature"] == 0.7
assert call_kwargs["max_tokens"] == 100
assert call_kwargs["model"] == "test-model"
@pytest.mark.parametrize("model_diagnostic_settings", [(True, False)], indirect=True)
async def test_agent_run_with_list_messages(mock_chat_client_agent, model_diagnostic_settings):
"""Test agent run with list of messages."""
from agent_framework import ChatMessage, Role
from agent_framework.telemetry import use_agent_telemetry
decorated_class = use_agent_telemetry(type(mock_chat_client_agent))
agent = decorated_class()
messages = [
ChatMessage(role=Role.USER, text="First message"),
ChatMessage(role=Role.ASSISTANT, text="Response"),
ChatMessage(role=Role.USER, text="Second message"),
]
with (
patch("agent_framework.telemetry.MODEL_DIAGNOSTICS_SETTINGS", model_diagnostic_settings),
patch("agent_framework.telemetry.use_span") as mock_use_span,
patch("agent_framework.telemetry._set_agent_run_input") as mock_set_input,
):
mock_span = Mock()
mock_use_span.return_value.__enter__.return_value = mock_span
mock_use_span.return_value.__exit__.return_value = None
response = await agent.run(messages)
assert response is not None
# Verify input was set with the list of messages
mock_set_input.assert_called_once_with("test_agent_system", messages)
@pytest.mark.parametrize("model_diagnostic_settings", [(True, False)], indirect=True)
async def test_agent_run_with_exception_handling(mock_chat_client_agent, model_diagnostic_settings):
async def test_agent_run_with_exception_handling(mock_chat_client_agent: AgentProtocol):
"""Test agent run with exception handling."""
from agent_framework.telemetry import use_agent_telemetry
async def run_with_error(self, messages=None, *, thread=None, **kwargs):
raise RuntimeError("Agent run error")
type(mock_chat_client_agent).run = run_with_error
mock_chat_client_agent.run = run_with_error
decorated_class = use_agent_telemetry(type(mock_chat_client_agent))
agent = decorated_class()
agent = use_agent_telemetry(mock_chat_client_agent)()
from opentelemetry.trace import Span
with (
patch("agent_framework.telemetry.MODEL_DIAGNOSTICS_SETTINGS", model_diagnostic_settings),
patch("agent_framework.telemetry.use_span") as mock_use_span,
patch("agent_framework.telemetry._get_span") as mock_get_span,
):
mock_span = Mock()
mock_use_span.return_value.__enter__.return_value = mock_span
mock_use_span.return_value.__exit__.return_value = None
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_once_with(
GenAIAttributes.ERROR_TYPE.value, str(type(RuntimeError("Agent run error")))
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"))
)
mock_span.set_status.assert_called_once_with(StatusCode.ERROR, repr(RuntimeError("Agent run error")))
python/packages/main/tests/main/test_tools.py
+45 -68
View File
@@ -15,7 +15,7 @@ from agent_framework import (
)
from agent_framework._tools import _parse_inputs
from agent_framework.exceptions import ToolException
from agent_framework.telemetry import GenAIAttributes
from agent_framework.telemetry import OtelAttr
# region AIFunction and ai_function decorator tests
@@ -83,19 +83,23 @@ async def test_ai_function_decorator_with_async():
assert (await async_test_tool(1, 2)) == 3
# Telemetry tests for AIFunction
async def test_ai_function_invoke_telemetry_enabled():
@pytest.mark.parametrize("otel_settings", [(True, True)], indirect=True)
async def test_ai_function_invoke_telemetry_enabled(otel_settings):
"""Test the ai_function invoke method with telemetry enabled."""
@ai_function(name="telemetry_test_tool", description="A test tool for telemetry")
@ai_function(
name="telemetry_test_tool",
description="A test tool for telemetry",
additional_properties={"otel_settings": otel_settings},
)
def telemetry_test_tool(x: int, y: int) -> int:
"""A function that adds two numbers for telemetry testing."""
return x + y
# Mock the tracer and span
with (
patch("agent_framework._tools.tracer") as mock_tracer,
patch("agent_framework._tools.start_as_current_span") as mock_start_span,
patch("agent_framework.telemetry.tracer"),
patch("agent_framework._tools.get_function_span") as mock_start_span,
):
mock_span = Mock()
mock_context_manager = Mock()
@@ -114,23 +118,26 @@ async def test_ai_function_invoke_telemetry_enabled():
assert result == 3
# Verify telemetry calls
mock_start_span.assert_called_once_with(
mock_tracer, telemetry_test_tool, metadata={"tool_call_id": "test_call_id", "kwargs": {"x": 1, "y": 2}}
)
mock_start_span.assert_called_once_with(function=telemetry_test_tool, 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[GenAIAttributes.MEASUREMENT_FUNCTION_TAG_NAME.value] == "telemetry_test_tool"
assert attributes[GenAIAttributes.TOOL_CALL_ID.value] == "test_call_id"
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():
@pytest.mark.parametrize("otel_settings", [(True, True)], indirect=True)
async def test_ai_function_invoke_telemetry_with_pydantic_args(otel_settings):
"""Test the ai_function invoke method with Pydantic model arguments."""
@ai_function(name="pydantic_test_tool", description="A test tool with Pydantic args")
@ai_function(
name="pydantic_test_tool",
description="A test tool with Pydantic args",
additional_properties={"otel_settings": otel_settings},
)
def pydantic_test_tool(x: int, y: int) -> int:
"""A function that adds two numbers using Pydantic args."""
return x + y
@@ -139,8 +146,8 @@ async def test_ai_function_invoke_telemetry_with_pydantic_args():
args_model = pydantic_test_tool.input_model(x=5, y=10)
with (
patch("agent_framework._tools.tracer") as mock_tracer,
patch("agent_framework._tools.start_as_current_span") as mock_start_span,
patch("agent_framework.telemetry.tracer"),
patch("agent_framework._tools.get_function_span") as mock_start_span,
):
mock_span = Mock()
mock_context_manager = Mock()
@@ -159,21 +166,27 @@ async def test_ai_function_invoke_telemetry_with_pydantic_args():
# Verify telemetry calls
mock_start_span.assert_called_once_with(
mock_tracer, pydantic_test_tool, metadata={"tool_call_id": "pydantic_call", "kwargs": {"x": 5, "y": 10}}
function=pydantic_test_tool,
tool_call_id="pydantic_call",
)
async def test_ai_function_invoke_telemetry_with_exception():
@pytest.mark.parametrize("otel_settings", [(True, True)], indirect=True)
async def test_ai_function_invoke_telemetry_with_exception(otel_settings):
"""Test the ai_function invoke method with telemetry when an exception occurs."""
@ai_function(name="exception_test_tool", description="A test tool that raises an exception")
@ai_function(
name="exception_test_tool",
description="A test tool that raises an exception",
additional_properties={"otel_settings": otel_settings},
)
def exception_test_tool(x: int, y: int) -> int:
"""A function that raises an exception for telemetry testing."""
raise ValueError("Test exception for telemetry")
with (
patch("agent_framework._tools.tracer"),
patch("agent_framework._tools.start_as_current_span") as mock_start_span,
patch("agent_framework.telemetry.tracer"),
patch("agent_framework._tools.get_function_span") as mock_start_span,
):
mock_span = Mock()
mock_context_manager = Mock()
@@ -200,20 +213,25 @@ async def test_ai_function_invoke_telemetry_with_exception():
mock_histogram.record.assert_called_once()
call_args = mock_histogram.record.call_args
attributes = call_args[1]["attributes"]
assert attributes[GenAIAttributes.ERROR_TYPE.value] == "ValueError"
assert attributes[OtelAttr.ERROR_TYPE] == ValueError.__name__
async def test_ai_function_invoke_telemetry_async_function():
@pytest.mark.parametrize("otel_settings", [(True, True)], indirect=True)
async def test_ai_function_invoke_telemetry_async_function(otel_settings):
"""Test the ai_function invoke method with telemetry on async function."""
@ai_function(name="async_telemetry_test", description="An async test tool for telemetry")
@ai_function(
name="async_telemetry_test",
description="An async test tool for telemetry",
additional_properties={"otel_settings": otel_settings},
)
async def async_telemetry_test(x: int, y: int) -> int:
"""An async function for telemetry testing."""
return x * y
with (
patch("agent_framework._tools.tracer") as mock_tracer,
patch("agent_framework._tools.start_as_current_span") as mock_start_span,
patch("agent_framework.telemetry.tracer"),
patch("agent_framework._tools.get_function_span") as mock_start_span,
):
mock_span = Mock()
mock_context_manager = Mock()
@@ -231,54 +249,13 @@ async def test_ai_function_invoke_telemetry_async_function():
assert result == 12
# Verify telemetry calls
mock_start_span.assert_called_once_with(
mock_tracer, async_telemetry_test, metadata={"tool_call_id": "async_call", "kwargs": {"x": 3, "y": 4}}
)
mock_start_span.assert_called_once_with(function=async_telemetry_test, tool_call_id="async_call")
# Verify histogram recording
mock_histogram.record.assert_called_once()
call_args = mock_histogram.record.call_args
attributes = call_args[1]["attributes"]
assert attributes[GenAIAttributes.MEASUREMENT_FUNCTION_TAG_NAME.value] == "async_telemetry_test"
async def test_ai_function_invoke_telemetry_no_tool_call_id():
"""Test the ai_function invoke method with telemetry when no tool_call_id is provided."""
@ai_function(name="no_id_test_tool", description="A test tool without tool_call_id")
def no_id_test_tool(x: int) -> int:
"""A function for testing without tool_call_id."""
return x * 2
with (
patch("agent_framework._tools.tracer") as mock_tracer,
patch("agent_framework._tools.start_as_current_span") 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()
no_id_test_tool._invocation_duration_histogram = mock_histogram
# Call invoke without tool_call_id
result = await no_id_test_tool.invoke(x=5)
# Verify result
assert result == 10
# Verify telemetry calls
mock_start_span.assert_called_once_with(
mock_tracer, no_id_test_tool, metadata={"tool_call_id": None, "kwargs": {"x": 5}}
)
# Verify histogram attributes
mock_histogram.record.assert_called_once()
call_args = mock_histogram.record.call_args
attributes = call_args[1]["attributes"]
assert attributes[GenAIAttributes.TOOL_CALL_ID.value] is None
assert attributes[OtelAttr.MEASUREMENT_FUNCTION_TAG_NAME] == "async_telemetry_test"
async def test_ai_function_invoke_invalid_pydantic_args():
python/packages/main/tests/openai/conftest.py
+25
View File
@@ -3,6 +3,8 @@ from typing import Any
from pytest import fixture
from agent_framework.telemetry import OtelSettings, setup_telemetry
# region Connector Settings fixtures
@fixture
@@ -49,3 +51,26 @@ 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
python/packages/main/tests/openai/test_openai_responses_client.py
+14 -17
View File
@@ -373,7 +373,7 @@ def test_chat_message_parsing_with_function_calls() -> None:
asyncio.run(client.get_response(messages=messages))
def test_response_format_parse_path() -> None:
async def test_response_format_parse_path() -> None:
"""Test get_response response_format parsing path."""
client = OpenAIResponsesClient(ai_model_id="test-model", api_key="test-key")
@@ -386,19 +386,18 @@ def test_response_format_parse_path() -> None:
mock_parsed_response.metadata = {}
mock_parsed_response.output_parsed = None
mock_parsed_response.usage = None
mock_parsed_response.finish_reason = None
with patch.object(client.client.responses, "parse", return_value=mock_parsed_response):
response = asyncio.run(
client.get_response(
messages=[ChatMessage(role="user", text="Test message")], response_format=OutputStruct, store=True
)
response = await client.get_response(
messages=[ChatMessage(role="user", text="Test message")], response_format=OutputStruct, store=True
)
assert response.conversation_id == "parsed_response_123"
assert response.ai_model_id == "test-model"
def test_bad_request_error_non_content_filter() -> None:
async def test_bad_request_error_non_content_filter() -> None:
"""Test get_response BadRequestError without content_filter."""
client = OpenAIResponsesClient(ai_model_id="test-model", api_key="test-key")
@@ -412,10 +411,8 @@ def test_bad_request_error_non_content_filter() -> None:
with patch.object(client.client.responses, "parse", side_effect=mock_error):
with pytest.raises(ServiceResponseException) as exc_info:
asyncio.run(
client.get_response(
messages=[ChatMessage(role="user", text="Test message")], response_format=OutputStruct
)
await client.get_response(
messages=[ChatMessage(role="user", text="Test message")], response_format=OutputStruct
)
assert "failed to complete the prompt" in str(exc_info.value)
@@ -440,11 +437,13 @@ async def test_streaming_content_filter_exception_handling() -> None:
break
def test_get_streaming_response_with_all_parameters() -> None:
@skip_if_openai_integration_tests_disabled
async def test_get_streaming_response_with_all_parameters() -> None:
"""Test get_streaming_response with all possible parameters."""
client = OpenAIResponsesClient(ai_model_id="test-model", api_key="test-key")
async def run_streaming_test():
# Should fail due to invalid API key
with pytest.raises(ServiceResponseException):
response = client.get_streaming_response(
messages=[ChatMessage(role="user", text="Test streaming")],
include=["file_search_call.results"],
@@ -471,10 +470,6 @@ def test_get_streaming_response_with_all_parameters() -> None:
async for _ in response:
break
# Should fail due to invalid API key
with pytest.raises(ServiceResponseException):
asyncio.run(run_streaming_test())
def test_response_content_creation_with_annotations() -> None:
"""Test _create_response_content with different annotation types."""
@@ -731,7 +726,9 @@ def test_create_streaming_response_content_with_mcp_approval_request() -> None:
assert fa.function_call.name == "do_stream_action"
def test_end_to_end_mcp_approval_flow() -> 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:
"""End-to-end mocked test:
model issues an mcp_approval_request, user approves, client sends mcp_approval_response.
"""