mirror of
https://github.com/microsoft/agent-framework.git
synced 2026-06-16 21:04:09 +08:00
Python: Improved telemetry setup (#421)
* 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
committed by
GitHub
Unverified
GitHub
parent
6aa746d891
commit
82ca4065cb
@@ -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()
|
||||
|
||||
@@ -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"
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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")))
|
||||
|
||||
@@ -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():
|
||||
|
||||
Reference in New Issue
Block a user