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