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/agent_framework/_agents.py
+12 -4
View File
@@ -9,11 +9,12 @@ from uuid import uuid4
from pydantic import BaseModel, Field, PrivateAttr
from ._clients import ChatClientProtocol
from ._clients import BaseChatClient, ChatClientProtocol
from ._logging import get_logger
from ._mcp import MCPTool
from ._pydantic import AFBaseModel
from ._threads import AgentThread, ChatMessageStore, deserialize_thread_state, thread_on_new_messages
from ._tools import ToolProtocol
from ._tools import FUNCTION_INVOKING_CHAT_CLIENT_MARKER, ToolProtocol
from ._types import (
AgentRunResponse,
AgentRunResponseUpdate,
@@ -32,6 +33,8 @@ if sys.version_info >= (3, 11):
else:
from typing_extensions import Self # pragma: no cover
logger = get_logger("agent_framework")
TThreadType = TypeVar("TThreadType", bound="AgentThread")
__all__ = ["AgentProtocol", "BaseAgent", "ChatAgent"]
@@ -248,6 +251,11 @@ class ChatAgent(BaseAgent):
kwargs: any additional keyword arguments.
Unused, can be used by subclasses of this Agent.
"""
if not hasattr(chat_client, FUNCTION_INVOKING_CHAT_CLIENT_MARKER) and isinstance(chat_client, BaseChatClient):
logger.warning(
"The provided chat client does not support function invoking, this might limit agent capabilities."
)
kwargs.update(additional_properties or {})
# We ignore the MCP Servers here and store them separately,
@@ -317,8 +325,8 @@ class ChatAgent(BaseAgent):
should check if there is already a agent name defined, and if not
set it to this value.
"""
if hasattr(self.chat_client, "_update_agent_name") and callable(self.chat_client._update_agent_name): # type: ignore[reportAttributeAccessIssue]
self.chat_client._update_agent_name(self.name) # type: ignore[reportAttributeAccessIssue]
if hasattr(self.chat_client, "_update_agent_name") and callable(self.chat_client._update_agent_name): # type: ignore[reportAttributeAccessIssue, attr-defined]
self.chat_client._update_agent_name(self.name) # type: ignore[reportAttributeAccessIssue, attr-defined]
async def run(
self,
python/packages/main/agent_framework/_clients.py
+34 -230
View File
@@ -2,32 +2,29 @@
import asyncio
from abc import ABC, abstractmethod
from collections.abc import AsyncIterable, Awaitable, Callable, MutableMapping, MutableSequence, Sequence
from functools import wraps
from collections.abc import AsyncIterable, Callable, MutableMapping, MutableSequence, Sequence
from typing import TYPE_CHECKING, Any, Generic, Literal, Protocol, TypeVar, runtime_checkable
from pydantic import BaseModel
from pydantic import BaseModel, Field
from ._logging import get_logger
from ._mcp import MCPTool
from ._pydantic import AFBaseModel
from ._threads import ChatMessageStore
from ._tools import AIFunction, ToolProtocol
from ._tools import ToolProtocol
from ._types import (
ChatMessage,
ChatOptions,
ChatResponse,
ChatResponseUpdate,
ChatToolMode,
Contents,
FunctionCallContent,
FunctionResultContent,
GeneratedEmbeddings,
)
if TYPE_CHECKING:
from ._agents import ChatAgent
TInput = TypeVar("TInput", contravariant=True)
TEmbedding = TypeVar("TEmbedding")
TBaseChatClient = TypeVar("TBaseChatClient", bound="BaseChatClient")
@@ -38,215 +35,8 @@ __all__ = [
"BaseChatClient",
"ChatClientProtocol",
"EmbeddingGenerator",
"use_tool_calling",
]
# region Tool Calling Functions and Decorators
async def _auto_invoke_function(
function_call_content: FunctionCallContent,
custom_args: dict[str, Any] | None = None,
*,
tool_map: dict[str, AIFunction[BaseModel, Any]],
sequence_index: int | None = None,
request_index: int | None = None,
) -> Contents:
"""Invoke a function call requested by the agent, applying filters that are defined in the agent."""
tool: AIFunction[BaseModel, Any] | None = tool_map.get(function_call_content.name)
if tool is None:
raise KeyError(f"No tool or function named '{function_call_content.name}'")
parsed_args: dict[str, Any] = dict(function_call_content.parse_arguments() or {})
# Merge with user-supplied args; right-hand side dominates, so parsed args win on conflicts.
merged_args: dict[str, Any] = (custom_args or {}) | parsed_args
args = tool.input_model.model_validate(merged_args)
exception = None
try:
function_result = await tool.invoke(arguments=args, tool_call_id=function_call_content.call_id)
except Exception as ex:
exception = ex
function_result = None
return FunctionResultContent(
call_id=function_call_content.call_id,
exception=exception,
result=function_result,
)
def _tool_call_non_streaming(
func: Callable[..., Awaitable["ChatResponse"]],
) -> Callable[..., Awaitable["ChatResponse"]]:
"""Decorate the internal _inner_get_response method to enable tool calls."""
@wraps(func)
async def wrapper(
self: "BaseChatClient",
*,
messages: MutableSequence[ChatMessage],
chat_options: ChatOptions,
**kwargs: Any,
) -> ChatResponse:
response: ChatResponse | None = None
fcc_messages: list[ChatMessage] = []
for attempt_idx in range(getattr(self, "__maximum_iterations_per_request", 10)):
response = await func(self, messages=messages, chat_options=chat_options, **kwargs)
# if there are function calls, we will handle them first
function_results = {
it.call_id for it in response.messages[0].contents if isinstance(it, FunctionResultContent)
}
function_calls = [
it
for it in response.messages[0].contents
if isinstance(it, FunctionCallContent) and it.call_id not in function_results
]
if function_calls:
# Run all function calls concurrently
results = await asyncio.gather(*[
_auto_invoke_function(
function_call,
custom_args=kwargs,
tool_map={t.name: t for t in chat_options.tools or [] if isinstance(t, AIFunction)}, # type: ignore[reportPrivateUsage]
sequence_index=seq_idx,
request_index=attempt_idx,
)
for seq_idx, function_call in enumerate(function_calls)
])
# add a single ChatMessage to the response with the results
result_message = ChatMessage(role="tool", contents=results)
response.messages.append(result_message)
# response should contain 2 messages after this,
# one with function call contents
# and one with function result contents
# the amount and call_id's should match
# this runs in every but the first run
# we need to keep track of all function call messages
fcc_messages.extend(response.messages)
# and add them as additional context to the messages
if chat_options.store:
messages.clear()
messages.append(result_message)
else:
messages.extend(response.messages)
continue
# If we reach this point, it means there were no function calls to handle,
# we'll add the previous function call and responses
# to the front of the list, so that the final response is the last one
# TODO (eavanvalkenburg): control this behavior?
if fcc_messages:
for msg in reversed(fcc_messages):
response.messages.insert(0, msg)
return response
# Failsafe: give up on tools, ask model for plain answer
chat_options.tool_choice = "none"
self._prepare_tool_choice(chat_options=chat_options) # type: ignore[reportPrivateUsage]
response = await func(self, messages=messages, chat_options=chat_options, **kwargs)
if fcc_messages:
for msg in reversed(fcc_messages):
response.messages.insert(0, msg)
return response
return wrapper
def _tool_call_streaming(
func: Callable[..., AsyncIterable["ChatResponseUpdate"]],
) -> Callable[..., AsyncIterable["ChatResponseUpdate"]]:
"""Decorate the internal _inner_get_response method to enable tool calls."""
@wraps(func)
async def wrapper(
self: "BaseChatClient",
*,
messages: MutableSequence[ChatMessage],
chat_options: ChatOptions,
**kwargs: Any,
) -> AsyncIterable[ChatResponseUpdate]:
"""Wrap the inner get streaming response method to handle tool calls."""
for attempt_idx in range(getattr(self, "__maximum_iterations_per_request", 10)):
function_call_returned = False
all_messages: list[ChatResponseUpdate] = []
async for update in func(self, messages=messages, chat_options=chat_options, **kwargs):
if update.contents and any(isinstance(item, FunctionCallContent) for item in update.contents):
all_messages.append(update)
function_call_returned = True
yield update
if not function_call_returned:
return
# There is one FunctionCallContent response stream in the messages, combining now to create
# the full completion depending on the prompt, the message may contain both function call
# content and others
response: ChatResponse = ChatResponse.from_chat_response_updates(all_messages)
# add the single assistant response message to the history
messages.append(response.messages[0])
function_calls = [item for item in response.messages[0].contents if isinstance(item, FunctionCallContent)]
# When conversation id is present, it means that messages are hosted on the server.
# In this case, we need to update ChatOptions with conversation id and also clear messages
if response.conversation_id is not None:
chat_options.conversation_id = response.conversation_id
messages = []
if function_calls:
# Run all function calls concurrently
results = await asyncio.gather(*[
_auto_invoke_function(
function_call,
custom_args=kwargs,
tool_map={t.name: t for t in chat_options.tools or [] if isinstance(t, AIFunction)}, # type: ignore[reportPrivateUsage]
sequence_index=seq_idx,
request_index=attempt_idx,
)
for seq_idx, function_call in enumerate(function_calls)
])
yield ChatResponseUpdate(contents=results, role="tool")
function_result_msg = ChatMessage(role="tool", contents=results)
response.messages.append(function_result_msg)
messages.append(function_result_msg)
continue
# Failsafe: give up on tools, ask model for plain answer
chat_options.tool_choice = "none"
self._prepare_tool_choice(chat_options=chat_options) # type: ignore[reportPrivateUsage]
async for update in func(self, messages=messages, chat_options=chat_options, **kwargs):
yield update
return wrapper
def use_tool_calling(cls: type[TBaseChatClient]) -> type[TBaseChatClient]:
"""Class decorator that enables tool calling for a chat client.
Remarks:
This only works on classes that derive from BaseChatClient
and the `_inner_get_response`
and `_inner_get_streaming_response` methods.
It also sets a `__maximum_iterations_per_request` attribute on the class.
if you want to expose this to end_users, do a version of this:
@property
def maximum_iterations_per_request(self):
return getattr(self, "__maximum_iterations_per_request", 10)
@maximum_iterations_per_request.setter
def maximum_iterations_per_request(self, value: int) -> None:
setattr(self, "__maximum_iterations_per_request", value)
"""
setattr(cls, "__maximum_iterations_per_request", 10)
if inner_response := getattr(cls, "_inner_get_response", None):
cls._inner_get_response = _tool_call_non_streaming(inner_response) # type: ignore
if inner_streaming_response := getattr(cls, "_inner_get_streaming_response", None):
cls._inner_get_streaming_response = _tool_call_streaming(inner_streaming_response) # type: ignore
return cls
# region ChatClientProtocol Protocol
@@ -255,6 +45,11 @@ def use_tool_calling(cls: type[TBaseChatClient]) -> type[TBaseChatClient]:
class ChatClientProtocol(Protocol):
"""A protocol for a chat client that can generate responses."""
@property
def additional_properties(self) -> dict[str, Any]:
"""Get additional properties associated with the client."""
...
async def get_response(
self,
messages: str | ChatMessage | list[str] | list[ChatMessage],
@@ -371,26 +166,35 @@ class ChatClientProtocol(Protocol):
...
# region ChatClientBase
def prepare_messages(messages: str | ChatMessage | list[str] | list[ChatMessage]) -> list[ChatMessage]:
"""Turn the allowed input into a list of chat messages."""
if isinstance(messages, str):
return [ChatMessage(role="user", text=messages)]
if isinstance(messages, ChatMessage):
return [messages]
return_messages: list[ChatMessage] = []
for msg in messages:
if isinstance(msg, str):
msg = ChatMessage(role="user", text=msg)
return_messages.append(msg)
return return_messages
class BaseChatClient(AFBaseModel, ABC):
"""Base class for chat clients."""
MODEL_PROVIDER_NAME: str = "unknown"
additional_properties: dict[str, Any] = Field(default_factory=dict)
OTEL_PROVIDER_NAME: str = "unknown"
# This is used for OTel setup, should be overridden in subclasses
def _prepare_messages(
def prepare_messages(
self, messages: str | ChatMessage | list[str] | list[ChatMessage]
) -> MutableSequence[ChatMessage]:
"""Turn the allowed input into a list of chat messages."""
if isinstance(messages, str):
return [ChatMessage(role="user", text=messages)]
if isinstance(messages, ChatMessage):
return [messages]
return_messages: list[ChatMessage] = []
for msg in messages:
if isinstance(msg, str):
msg = ChatMessage(role="user", text=msg)
return_messages.append(msg)
return return_messages
return prepare_messages(messages)
@staticmethod
def _normalize_tools(
@@ -537,7 +341,7 @@ class BaseChatClient(AFBaseModel, ABC):
user=user,
additional_properties=additional_properties or {},
)
prepped_messages = self._prepare_messages(messages)
prepped_messages = self.prepare_messages(messages)
self._prepare_tool_choice(chat_options=chat_options)
return await self._inner_get_response(messages=prepped_messages, chat_options=chat_options, **kwargs)
@@ -617,7 +421,7 @@ class BaseChatClient(AFBaseModel, ABC):
additional_properties=additional_properties or {},
**kwargs,
)
prepped_messages = self._prepare_messages(messages)
prepped_messages = self.prepare_messages(messages)
self._prepare_tool_choice(chat_options=chat_options)
async for update in self._inner_get_streaming_response(
messages=prepped_messages, chat_options=chat_options, **kwargs
@@ -640,13 +444,13 @@ class BaseChatClient(AFBaseModel, ABC):
else:
chat_options.tool_choice = chat_tool_mode.mode
def service_url(self) -> str | None:
def service_url(self) -> str:
"""Get the URL of the service.
Override this in the subclass to return the proper URL.
If the service does not have a URL, return None.
"""
return None
return "Unknown"
def create_agent(
self,
python/packages/main/agent_framework/_tools.py
+382 -31
View File
@@ -1,13 +1,16 @@
# Copyright (c) Microsoft. All rights reserved.
import asyncio
import inspect
import sys
from collections.abc import Awaitable, Callable, Collection
from collections.abc import AsyncIterable, Awaitable, Callable, Collection, MutableMapping, Sequence
from functools import wraps
from time import perf_counter
from time import perf_counter, time_ns
from typing import (
TYPE_CHECKING,
Annotated,
Any,
Final,
Generic,
Literal,
Protocol,
@@ -17,27 +20,39 @@ from typing import (
runtime_checkable,
)
from opentelemetry import metrics, trace
from opentelemetry import metrics
from pydantic import AnyUrl, BaseModel, Field, PrivateAttr, ValidationError, create_model, field_validator
from ._logging import get_logger
from ._pydantic import AFBaseModel
from .exceptions import ToolException
from .telemetry import GenAIAttributes, start_as_current_span
from .exceptions import ChatClientInitializationError, ToolException
from .telemetry import (
OPERATION_DURATION_BUCKET_BOUNDARIES,
OtelAttr,
_capture_exception, # type: ignore
get_function_span,
meter,
)
if TYPE_CHECKING:
from ._types import Contents
from ._clients import ChatClientProtocol
from ._types import (
ChatMessage,
ChatResponse,
ChatResponseUpdate,
Contents,
FunctionCallContent,
)
if sys.version_info >= (3, 12):
from typing import TypedDict # pragma: no cover
else:
from typing_extensions import TypedDict # pragma: no cover
tracer: trace.Tracer = trace.get_tracer("agent_framework")
meter: metrics.Meter = metrics.get_meter_provider().get_meter("agent_framework")
logger = get_logger()
__all__ = [
"FUNCTION_INVOKING_CHAT_CLIENT_MARKER",
"AIFunction",
"HostedCodeInterpreterTool",
"HostedFileSearchTool",
@@ -46,9 +61,17 @@ __all__ = [
"HostedWebSearchTool",
"ToolProtocol",
"ai_function",
"use_function_invocation",
]
logger = get_logger()
FUNCTION_INVOKING_CHAT_CLIENT_MARKER: Final[str] = "__function_invoking_chat_client__"
DEFAULT_MAX_ITERATIONS: Final[int] = 10
TChatClient = TypeVar("TChatClient", bound="ChatClientProtocol")
# region Helpers
def _parse_inputs(
inputs: "Contents | dict[str, Any] | str | list[Contents | dict[str, Any] | str] | None",
) -> list["Contents"]:
@@ -91,6 +114,7 @@ def _parse_inputs(
return parsed_inputs
# region Tools
@runtime_checkable
class ToolProtocol(Protocol):
"""Represents a generic tool that can be specified to an AI service.
@@ -337,7 +361,7 @@ class HostedFileSearchTool(BaseTool):
class AIFunction(BaseTool, Generic[ArgsT, ReturnT]):
"""A ToolProtocol that is callable as code.
"""A AITool that is callable as code.
Args:
name: The name of the function.
@@ -351,9 +375,10 @@ class AIFunction(BaseTool, Generic[ArgsT, ReturnT]):
input_model: type[ArgsT]
_invocation_duration_histogram: metrics.Histogram = PrivateAttr(
default_factory=lambda: meter.create_histogram(
GenAIAttributes.MEASUREMENT_FUNCTION_INVOCATION_DURATION.value,
unit="s",
name=OtelAttr.MEASUREMENT_FUNCTION_INVOCATION_DURATION,
unit=OtelAttr.DURATION_UNIT,
description="Measures the duration of a function's execution",
explicit_bucket_boundaries_advisory=OPERATION_DURATION_BUCKET_BOUNDARIES,
)
)
@@ -371,40 +396,60 @@ class AIFunction(BaseTool, Generic[ArgsT, ReturnT]):
Args:
arguments: A Pydantic model instance containing the arguments for the function.
kwargs: keyword arguments to pass to the function, will not be used if `args` is provided.
otel_settings: Optional model diagnostics settings to override the default settings.
kwargs: keyword arguments to pass to the function, will not be used if `arguments` is provided.
"""
global OTEL_SETTINGS
from .telemetry import OTEL_SETTINGS, setup_telemetry
tool_call_id = kwargs.pop("tool_call_id", None)
if arguments is not None:
if not isinstance(arguments, self.input_model):
raise TypeError(f"Expected {self.input_model.__name__}, got {type(arguments).__name__}")
kwargs = arguments.model_dump(exclude_none=True)
logger.info(f"Function name: {self.name}")
logger.debug(f"Function arguments: {kwargs}")
with start_as_current_span(
tracer, self, metadata={"tool_call_id": tool_call_id, "kwargs": kwargs}
) as current_span:
attributes: dict[str, Any] = {
GenAIAttributes.MEASUREMENT_FUNCTION_TAG_NAME.value: self.name,
GenAIAttributes.TOOL_CALL_ID.value: tool_call_id,
if not OTEL_SETTINGS.ENABLED: # type: ignore
logger.info(f"Function name: {self.name}")
logger.debug(f"Function arguments: {kwargs}")
res = self.__call__(**kwargs)
result = await res if inspect.isawaitable(res) else res
logger.info(f"Function {self.name} succeeded.")
logger.debug(f"Function result: {result or 'None'}")
return result # type: ignore[reportReturnType]
setup_telemetry()
with get_function_span(
function=self,
tool_call_id=tool_call_id,
) as span:
hist_attributes: dict[str, Any] = {
OtelAttr.MEASUREMENT_FUNCTION_TAG_NAME: self.name,
OtelAttr.TOOL_CALL_ID: tool_call_id or "unknown",
}
starting_time_stamp = perf_counter()
logger.info(f"Function name: {self.name}")
if OTEL_SETTINGS.SENSITIVE_DATA_ENABLED: # type: ignore
logger.debug(f"Function arguments: {kwargs}")
start_time_stamp = perf_counter()
end_time_stamp: float | None = None
try:
res = self.__call__(**kwargs)
result = await res if inspect.isawaitable(res) else res
logger.info(f"Function {self.name} succeeded.")
logger.debug(f"Function result: {result or 'None'}")
return result # type: ignore[reportReturnType]
end_time_stamp = perf_counter()
except Exception as exception:
attributes[GenAIAttributes.ERROR_TYPE.value] = type(exception).__name__
current_span.record_exception(exception)
current_span.set_attribute(GenAIAttributes.ERROR_TYPE.value, type(exception).__name__)
current_span.set_status(trace.StatusCode.ERROR, description=str(exception))
end_time_stamp = perf_counter()
hist_attributes[OtelAttr.ERROR_TYPE] = type(exception).__name__
_capture_exception(span=span, exception=exception, timestamp=time_ns())
logger.error(f"Function failed. Error: {exception}")
raise
else:
logger.info(f"Function {self.name} succeeded.")
if OTEL_SETTINGS.SENSITIVE_DATA_ENABLED: # type: ignore
logger.debug(f"Function result: {result or 'None'}")
return result # type: ignore[reportReturnType]
finally:
duration = perf_counter() - starting_time_stamp
self._invocation_duration_histogram.record(duration, attributes=attributes)
logger.info("Function completed. Duration: %fs", duration)
duration = (end_time_stamp or perf_counter()) - start_time_stamp
span.set_attribute(OtelAttr.MEASUREMENT_FUNCTION_INVOCATION_DURATION, duration)
self._invocation_duration_histogram.record(duration, attributes=hist_attributes)
logger.info("Function duration: %fs", duration)
def parameters(self) -> dict[str, Any]:
"""Create the json schema of the parameters."""
@@ -422,6 +467,9 @@ class AIFunction(BaseTool, Generic[ArgsT, ReturnT]):
}
# region AI Function Decorator
def _parse_annotation(annotation: Any) -> Any:
"""Parse a type annotation and return the corresponding type.
@@ -499,3 +547,306 @@ def ai_function(
return wrapper(func)
return decorator(func) if func else decorator # type: ignore[reportReturnType, return-value]
# region Function Invoking Chat Client
async def _auto_invoke_function(
function_call_content: "FunctionCallContent",
custom_args: dict[str, Any] | None = None,
*,
tool_map: dict[str, AIFunction[BaseModel, Any]],
sequence_index: int | None = None,
request_index: int | None = None,
) -> "Contents":
"""Invoke a function call requested by the agent, applying filters that are defined in the agent."""
from ._types import FunctionResultContent
tool: AIFunction[BaseModel, Any] | None = tool_map.get(function_call_content.name)
if tool is None:
raise KeyError(f"No tool or function named '{function_call_content.name}'")
parsed_args: dict[str, Any] = dict(function_call_content.parse_arguments() or {})
# Merge with user-supplied args; right-hand side dominates, so parsed args win on conflicts.
merged_args: dict[str, Any] = (custom_args or {}) | parsed_args
args = tool.input_model.model_validate(merged_args)
exception = None
try:
function_result = await tool.invoke(
arguments=args,
tool_call_id=function_call_content.call_id,
) # type: ignore[arg-type]
except Exception as ex:
exception = ex
function_result = None
return FunctionResultContent(
call_id=function_call_content.call_id,
exception=exception,
result=function_result,
)
def _get_tool_map(
tools: "ToolProtocol \
| Callable[..., Any] \
| MutableMapping[str, Any] \
| list[ToolProtocol | Callable[..., Any] | MutableMapping[str, Any]]",
) -> dict[str, AIFunction[Any, Any]]:
ai_function_list: dict[str, AIFunction[Any, Any]] = {}
for tool in tools if isinstance(tools, list) else [tools]:
if isinstance(tool, AIFunction):
ai_function_list[tool.name] = tool
continue
if callable(tool):
# Convert to AITool if it's a function or callable
ai_tool = ai_function(tool)
ai_function_list[ai_tool.name] = ai_tool
return ai_function_list
async def execute_function_calls(
custom_args: dict[str, Any],
attempt_idx: int,
function_calls: Sequence["FunctionCallContent"],
tools: "ToolProtocol \
| Callable[..., Any] \
| MutableMapping[str, Any] \
| list[ToolProtocol | Callable[..., Any] | MutableMapping[str, Any]]",
) -> list["Contents"]:
tool_map = _get_tool_map(tools)
# Run all function calls concurrently
return await asyncio.gather(*[
_auto_invoke_function(
function_call_content=function_call,
custom_args=custom_args,
tool_map=tool_map,
sequence_index=seq_idx,
request_index=attempt_idx,
)
for seq_idx, function_call in enumerate(function_calls)
])
def update_conversation_id(kwargs: dict[str, Any], conversation_id: str | None) -> None:
"""Update kwargs with conversation id."""
if conversation_id is None:
return
if "chat_options" in kwargs:
kwargs["chat_options"].conversation_id = conversation_id
else:
kwargs["conversation_id"] = conversation_id
def _handle_function_calls_response(
func: Callable[..., Awaitable["ChatResponse"]],
*,
max_iterations: int = 10,
) -> Callable[..., Awaitable["ChatResponse"]]:
"""Decorate the get_response method to enable function calls.
Args:
func: The get_response method to decorate.
max_iterations: The maximum number of function call iterations to perform.
"""
def decorator(
func: Callable[..., Awaitable["ChatResponse"]],
) -> Callable[..., Awaitable["ChatResponse"]]:
"""Inner decorator."""
@wraps(func)
async def function_invocation_wrapper(
self: "ChatClientProtocol",
messages: "str | ChatMessage | list[str] | list[ChatMessage]",
**kwargs: Any,
) -> "ChatResponse":
from ._clients import prepare_messages
from ._types import ChatMessage, ChatOptions, FunctionCallContent, FunctionResultContent
prepped_messages = prepare_messages(messages)
response: "ChatResponse | None" = None
fcc_messages: "list[ChatMessage]" = []
for attempt_idx in range(max_iterations):
response = await func(self, messages=prepped_messages, **kwargs)
# if there are function calls, we will handle them first
function_results = {
it.call_id for it in response.messages[0].contents if isinstance(it, FunctionResultContent)
}
function_calls = [
it
for it in response.messages[0].contents
if isinstance(it, FunctionCallContent) and it.call_id not in function_results
]
if response.conversation_id is not None:
update_conversation_id(kwargs, response.conversation_id)
prepped_messages = []
tools = kwargs.get("tools")
if not tools and (chat_options := kwargs.get("chat_options")) and isinstance(chat_options, ChatOptions):
tools = chat_options.tools
if function_calls and tools:
function_results = await execute_function_calls(
custom_args=kwargs,
attempt_idx=attempt_idx,
function_calls=function_calls,
tools=tools, # type: ignore
)
# add a single ChatMessage to the response with the results
result_message = ChatMessage(role="tool", contents=function_results) # type: ignore[call-overload]
response.messages.append(result_message)
# response should contain 2 messages after this,
# one with function call contents
# and one with function result contents
# the amount and call_id's should match
# this runs in every but the first run
# we need to keep track of all function call messages
fcc_messages.extend(response.messages)
# and add them as additional context to the messages
if kwargs.get("store"):
prepped_messages.clear()
prepped_messages.append(result_message)
else:
prepped_messages.extend(response.messages)
continue
# If we reach this point, it means there were no function calls to handle,
# we'll add the previous function call and responses
# to the front of the list, so that the final response is the last one
# TODO (eavanvalkenburg): control this behavior?
if fcc_messages:
for msg in reversed(fcc_messages):
response.messages.insert(0, msg)
return response
# Failsafe: give up on tools, ask model for plain answer
kwargs["tool_choice"] = "none"
response = await func(self, messages=prepped_messages, **kwargs)
if fcc_messages:
for msg in reversed(fcc_messages):
response.messages.insert(0, msg)
return response
return function_invocation_wrapper # type: ignore
return decorator(func)
def _handle_function_calls_streaming_response(
func: Callable[..., AsyncIterable["ChatResponseUpdate"]],
*,
max_iterations: int = 10,
) -> Callable[..., AsyncIterable["ChatResponseUpdate"]]:
"""Decorate the get_streaming_response method to handle function calls.
Args:
func: The get_streaming_response method to decorate.
max_iterations: The maximum number of function call iterations to perform.
"""
def decorator(
func: Callable[..., AsyncIterable["ChatResponseUpdate"]],
) -> Callable[..., AsyncIterable["ChatResponseUpdate"]]:
"""Inner decorator."""
@wraps(func)
async def streaming_function_invocation_wrapper(
self: "ChatClientProtocol",
messages: "str | ChatMessage | list[str] | list[ChatMessage]",
**kwargs: Any,
) -> AsyncIterable["ChatResponseUpdate"]:
"""Wrap the inner get streaming response method to handle tool calls."""
from ._clients import prepare_messages
from ._types import ChatMessage, ChatOptions, ChatResponse, ChatResponseUpdate, FunctionCallContent
prepped_messages = prepare_messages(messages)
for attempt_idx in range(max_iterations):
all_updates: list["ChatResponseUpdate"] = []
async for update in func(self, messages=prepped_messages, **kwargs):
all_updates.append(update)
yield update
# efficient check for FunctionCallContent in the updates
# if there is at least one, this stops and continuous
# if there are no FCC's then it returns
if not any(isinstance(item, FunctionCallContent) for upd in all_updates for item in upd.contents):
return
# Now combining the updates to create the full response.
# Depending on the prompt, the message may contain both function call
# content and others
response: "ChatResponse" = ChatResponse.from_chat_response_updates(all_updates)
# add the response message to the previous messages
prepped_messages.append(response.messages[0])
# get the fccs
function_calls = [
item for item in response.messages[0].contents if isinstance(item, FunctionCallContent)
]
# When conversation id is present, it means that messages are hosted on the server.
# In this case, we need to update kwargs with conversation id and also clear messages
if response.conversation_id is not None:
update_conversation_id(kwargs, response.conversation_id)
prepped_messages = []
tools = kwargs.get("tools")
if not tools and (chat_options := kwargs.get("chat_options")) and isinstance(chat_options, ChatOptions):
tools = chat_options.tools
if function_calls and tools:
function_results = await execute_function_calls(
custom_args=kwargs,
attempt_idx=attempt_idx,
function_calls=function_calls,
tools=tools, # type: ignore[reportArgumentType]
)
function_result_msg = ChatMessage(role="tool", contents=function_results)
yield ChatResponseUpdate(contents=function_results, role="tool")
response.messages.append(function_result_msg)
prepped_messages.append(function_result_msg)
continue
# Failsafe: give up on tools, ask model for plain answer
kwargs["tool_choice"] = "none"
async for update in func(self, messages=prepped_messages, **kwargs):
yield update
return streaming_function_invocation_wrapper
return decorator(func)
def use_function_invocation(
chat_client: type[TChatClient],
) -> type[TChatClient]:
"""Class decorator that enables tool calling for a chat client."""
if getattr(chat_client, FUNCTION_INVOKING_CHAT_CLIENT_MARKER, False):
return chat_client
max_iterations = DEFAULT_MAX_ITERATIONS
try:
chat_client.get_response = _handle_function_calls_response( # type: ignore
func=chat_client.get_response, # type: ignore
max_iterations=max_iterations,
)
except AttributeError as ex:
raise ChatClientInitializationError(
f"Chat client {chat_client.__name__} does not have a get_response method, cannot apply function invocation."
) from ex
try:
chat_client.get_streaming_response = _handle_function_calls_streaming_response( # type: ignore
func=chat_client.get_streaming_response,
max_iterations=max_iterations,
)
except AttributeError as ex:
raise ChatClientInitializationError(
f"Chat client {chat_client.__name__} does not have a get_streaming_response method, "
"cannot apply function invocation."
) from ex
setattr(chat_client, FUNCTION_INVOKING_CHAT_CLIENT_MARKER, True)
return chat_client
python/packages/main/agent_framework/_types.py
+1 -1
View File
@@ -933,7 +933,7 @@ class FunctionCallContent(BaseContent):
if not isinstance(other, FunctionCallContent):
raise TypeError("Incompatible type")
if other.call_id and self.call_id != other.call_id:
raise AdditionItemMismatch
raise AdditionItemMismatch("", log_level=None)
if not self.arguments:
arguments = other.arguments
elif not other.arguments:
python/packages/main/agent_framework/exceptions.py
+31 -10
View File
@@ -1,7 +1,7 @@
# Copyright (c) Microsoft. All rights reserved.
import logging
from typing import Any
from typing import Any, Literal
logger = logging.getLogger("agent_framework")
@@ -12,12 +12,20 @@ class AgentFrameworkException(Exception):
Automatically logs the message as debug.
"""
def __init__(self, message: str, inner_exception: Exception | None = None, *args: Any):
def __init__(
self,
message: str,
inner_exception: Exception | None = None,
log_level: Literal[0] | Literal[10] | Literal[20] | Literal[30] | Literal[40] | Literal[50] | None = 10,
*args: Any,
**kwargs: Any,
):
"""Create an AgentFrameworkException.
This emits a debug log, with the inner_exception if provided.
This emits a debug log (by default), with the inner_exception if provided.
"""
logger.debug(message, exc_info=inner_exception)
if log_level is not None:
logger.log(log_level, message, exc_info=inner_exception)
if inner_exception:
super().__init__(message, inner_exception, *args) # type: ignore
super().__init__(message, *args) # type: ignore
@@ -35,6 +43,24 @@ class AgentExecutionException(AgentException):
pass
class AgentInitializationError(AgentException):
"""An error occurred while initializing the agent."""
pass
class ChatClientException(AgentFrameworkException):
"""An error occurred while dealing with a chat client."""
pass
class ChatClientInitializationError(ChatClientException):
"""An error occurred while initializing the chat client."""
pass
# region Service Exceptions
@@ -101,9 +127,4 @@ class ToolExecutionException(ToolException):
class AdditionItemMismatch(AgentFrameworkException):
"""An error occurred while adding two types."""
def __init__(self) -> None:
"""Create an AdditionItemMismatch.
Unlike the AgentFrameworkException, this does not log the message automatically,
"""
pass
pass
python/packages/main/agent_framework/openai/_assistants_client.py
+6 -4
View File
@@ -20,8 +20,8 @@ from openai.types.beta.threads.run_submit_tool_outputs_params import ToolOutput
from openai.types.beta.threads.runs import RunStep
from pydantic import Field, PrivateAttr, SecretStr, ValidationError
from .._clients import BaseChatClient, use_tool_calling
from .._tools import AIFunction, HostedCodeInterpreterTool, HostedFileSearchTool
from .._clients import BaseChatClient
from .._tools import AIFunction, HostedCodeInterpreterTool, HostedFileSearchTool, use_function_invocation
from .._types import (
ChatMessage,
ChatOptions,
@@ -50,8 +50,8 @@ else:
__all__ = ["OpenAIAssistantsClient"]
@use_function_invocation
@use_telemetry
@use_tool_calling
class OpenAIAssistantsClient(OpenAIConfigMixin, BaseChatClient):
"""OpenAI Assistants client."""
@@ -166,7 +166,9 @@ class OpenAIAssistantsClient(OpenAIConfigMixin, BaseChatClient):
# Get the thread ID
thread_id: str | None = (
chat_options.conversation_id if chat_options.conversation_id is not None else self.thread_id
chat_options.conversation_id
if chat_options.conversation_id is not None
else run_options.get("conversation_id", self.thread_id)
)
if thread_id is None and tool_results is not None:
python/packages/main/agent_framework/openai/_chat_client.py
+32 -10
View File
@@ -1,6 +1,7 @@
# Copyright (c) Microsoft. All rights reserved.
import json
import sys
from collections.abc import AsyncIterable, Mapping, MutableMapping, MutableSequence, Sequence
from datetime import datetime
from itertools import chain
@@ -15,9 +16,9 @@ from openai.types.chat.chat_completion_chunk import Choice as ChunkChoice
from openai.types.chat.chat_completion_message_custom_tool_call import ChatCompletionMessageCustomToolCall
from pydantic import BaseModel, SecretStr, ValidationError
from .._clients import BaseChatClient, use_tool_calling
from .._clients import BaseChatClient
from .._logging import get_logger
from .._tools import AIFunction, HostedWebSearchTool, ToolProtocol
from .._tools import AIFunction, HostedWebSearchTool, ToolProtocol, use_function_invocation
from .._types import (
ChatMessage,
ChatOptions,
@@ -41,14 +42,17 @@ from ..telemetry import use_telemetry
from ._exceptions import OpenAIContentFilterException
from ._shared import OpenAIBase, OpenAIConfigMixin, OpenAISettings, prepare_function_call_results
if sys.version_info >= (3, 12):
from typing import override # type: ignore # pragma: no cover
else:
from typing_extensions import override # type: ignore[import] # pragma: no cover
__all__ = ["OpenAIChatClient"]
logger = get_logger("agent_framework.openai")
# region Base Client
@use_telemetry
@use_tool_calling
class OpenAIBaseChatClient(OpenAIBase, BaseChatClient):
"""OpenAI Chat completion class."""
@@ -233,11 +237,26 @@ class OpenAIBaseChatClient(OpenAIBase, BaseChatClient):
)
def _usage_details_from_openai(self, usage: CompletionUsage) -> UsageDetails:
return UsageDetails(
prompt_tokens=usage.prompt_tokens,
completion_tokens=usage.completion_tokens,
total_tokens=usage.total_tokens,
details = UsageDetails(
input_token_count=usage.prompt_tokens,
output_token_count=usage.completion_tokens,
total_token_count=usage.total_tokens,
)
if usage.completion_tokens_details:
if tokens := usage.completion_tokens_details.accepted_prediction_tokens:
details["completion/accepted_prediction_tokens"] = tokens
if tokens := usage.completion_tokens_details.audio_tokens:
details["completion/audio_tokens"] = tokens
if tokens := usage.completion_tokens_details.reasoning_tokens:
details["completion/reasoning_tokens"] = tokens
if tokens := usage.completion_tokens_details.rejected_prediction_tokens:
details["completion/rejected_prediction_tokens"] = tokens
if usage.prompt_tokens_details:
if tokens := usage.prompt_tokens_details.audio_tokens:
details["prompt/audio_tokens"] = tokens
if tokens := usage.prompt_tokens_details.cached_tokens:
details["prompt/cached_tokens"] = tokens
return details
def _parse_text_from_choice(self, choice: Choice | ChunkChoice) -> TextContent | None:
"""Parse the choice into a TextContent object."""
@@ -362,13 +381,14 @@ class OpenAIBaseChatClient(OpenAIBase, BaseChatClient):
case _:
return content.model_dump(exclude_none=True)
def service_url(self) -> str | None:
@override
def service_url(self) -> str:
"""Get the URL of the service.
Override this in the subclass to return the proper URL.
If the service does not have a URL, return None.
"""
return str(self.client.base_url) if self.client else None
return str(self.client.base_url) if self.client else "Unknown"
# region Public client
@@ -376,6 +396,8 @@ class OpenAIBaseChatClient(OpenAIBase, BaseChatClient):
TOpenAIChatClient = TypeVar("TOpenAIChatClient", bound="OpenAIChatClient")
@use_function_invocation
@use_telemetry
class OpenAIChatClient(OpenAIConfigMixin, OpenAIBaseChatClient):
"""OpenAI Chat completion class."""
python/packages/main/agent_framework/openai/_responses_client.py
+4 -3
View File
@@ -25,7 +25,7 @@ from openai.types.responses.web_search_tool_param import UserLocation as WebSear
from openai.types.responses.web_search_tool_param import WebSearchToolParam
from pydantic import BaseModel, SecretStr, ValidationError
from .._clients import BaseChatClient, use_tool_calling
from .._clients import BaseChatClient
from .._logging import get_logger
from .._tools import (
AIFunction,
@@ -34,6 +34,7 @@ from .._tools import (
HostedMCPTool,
HostedWebSearchTool,
ToolProtocol,
use_function_invocation,
)
from .._types import (
ChatMessage,
@@ -406,7 +407,7 @@ class OpenAIBaseResponsesClient(OpenAIBase, BaseChatClient):
tool_args["file_ids"] = []
for tool_input in tool.inputs:
if isinstance(tool_input, HostedFileContent):
tool_args["file_ids"].append(tool_input.file_id)
tool_args["file_ids"].append(tool_input.file_id) # type: ignore[attr-defined]
if not tool_args["file_ids"]:
tool_args.pop("file_ids")
response_tools.append(
@@ -1040,8 +1041,8 @@ class OpenAIBaseResponsesClient(OpenAIBase, BaseChatClient):
TOpenAIResponsesClient = TypeVar("TOpenAIResponsesClient", bound="OpenAIResponsesClient")
@use_function_invocation
@use_telemetry
@use_tool_calling
class OpenAIResponsesClient(OpenAIConfigMixin, OpenAIBaseResponsesClient):
"""OpenAI Responses client class."""
python/packages/main/agent_framework/openai/_shared.py
+2 -2
View File
@@ -60,7 +60,7 @@ def prepare_function_call_results(content: Contents | Any | list[Contents | Any]
results.extend(res)
else:
results.append(res)
return results[0] if len(results) == 1 else results
return results[0] if len(results) == 1 else json.dumps(results)
if isinstance(content, BaseModel):
return content.model_dump_json(exclude_none=True, exclude={"raw_representation", "additional_properties"})
# fallback
@@ -127,7 +127,7 @@ class OpenAIBase(AFBaseModel):
class OpenAIConfigMixin(OpenAIBase):
"""Internal class for configuring a connection to an OpenAI service."""
MODEL_PROVIDER_NAME: ClassVar[str] = "openai" # type: ignore[reportIncompatibleVariableOverride, misc]
OTEL_PROVIDER_NAME: ClassVar[str] = "openai" # type: ignore[reportIncompatibleVariableOverride, misc]
@validate_call(config=ConfigDict(arbitrary_types_allowed=True))
def __init__(
python/packages/main/agent_framework/telemetry.py
+882 -540
View File
File diff suppressed because it is too large Load Diff
python/packages/main/pyproject.toml
+3
View File
@@ -30,6 +30,9 @@ dependencies = [
"opentelemetry-api ~= 1.24",
"opentelemetry-sdk ~= 1.24",
"mcp>=1.12",
"azure-monitor-opentelemetry>=1.7.0",
"azure-monitor-opentelemetry-exporter>=1.0.0b41",
"opentelemetry-exporter-otlp-proto-grpc>=1.36.0",
]
[project.optional-dependencies]
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.
"""