Python: Agent and Function middleware (#770)

* Initial middleware implementation

* Small fixes

* Small updates

* Small updates in samples

* Moved middleware functionality to decorator

* Removed obsolete file

* Renamed AgentInvocationContext to AzureRunContext

* Added unit tests

* Small settings update for test discovery in VS Code

* Added unit tests

* Reverted changes in environment settings

* Added context result override

* Renaming and updates to logic

* Added more samples

* Updated DEV_SETUP.md

* Addressed PR feedback

* Addressed PR feedback

* Removed unused parameter

* Small fix

* Small fix in telemetry logic

* Revert "Small fix in telemetry logic"

This reverts commit 6f82660d2d.

* Small fix

---------

Co-authored-by: Chris <66376200+crickman@users.noreply.github.com>
This commit is contained in:
Dmytro Struk
2025-09-18 16:30:05 -07:00
committed by GitHub
Unverified
parent 538be4c149
commit 99860a5d07
16 changed files with 3071 additions and 40 deletions
+45 -7
View File
@@ -576,6 +576,7 @@ async def _auto_invoke_function(
tool_map: dict[str, AIFunction[BaseModel, Any]],
sequence_index: int | None = None,
request_index: int | None = None,
middleware_pipeline: Any = None, # Optional MiddlewarePipeline
) -> "Contents":
"""Invoke a function call requested by the agent, applying filters that are defined in the agent."""
from ._types import FunctionResultContent
@@ -590,14 +591,43 @@ async def _auto_invoke_function(
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(
# Execute through middleware pipeline if available
if middleware_pipeline and hasattr(middleware_pipeline, "has_middlewares") and middleware_pipeline.has_middlewares:
from ._middleware import FunctionInvocationContext
middleware_context = FunctionInvocationContext(
function=tool,
arguments=args,
tool_call_id=function_call_content.call_id,
) # type: ignore[arg-type]
except Exception as ex:
exception = ex
function_result = None
)
async def final_function_handler(context_obj: Any) -> Any:
return await tool.invoke(
arguments=context_obj.arguments,
tool_call_id=function_call_content.call_id,
)
try:
function_result = await middleware_pipeline.execute(
function=tool,
arguments=args,
context=middleware_context,
final_handler=final_function_handler,
)
except Exception as ex:
exception = ex
function_result = None
else:
# No middleware - execute directly
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,
@@ -631,6 +661,7 @@ async def execute_function_calls(
| Callable[..., Any] \
| MutableMapping[str, Any] \
| list[ToolProtocol | Callable[..., Any] | MutableMapping[str, Any]]",
middleware_pipeline: Any = None, # Optional MiddlewarePipeline to avoid circular imports
) -> list["Contents"]:
tool_map = _get_tool_map(tools)
# Run all function calls concurrently
@@ -641,6 +672,7 @@ async def execute_function_calls(
tool_map=tool_map,
sequence_index=seq_idx,
request_index=attempt_idx,
middleware_pipeline=middleware_pipeline,
)
for seq_idx, function_call in enumerate(function_calls)
])
@@ -706,11 +738,14 @@ def _handle_function_calls_response(
if not tools and (chat_options := kwargs.get("chat_options")) and isinstance(chat_options, ChatOptions):
tools = chat_options.tools
if function_calls and tools:
# Extract function middleware pipeline from kwargs if available
middleware_pipeline = kwargs.get("_function_middleware_pipeline")
function_results = await execute_function_calls(
custom_args=kwargs,
attempt_idx=attempt_idx,
function_calls=function_calls,
tools=tools, # type: ignore
middleware_pipeline=middleware_pipeline,
)
# add a single ChatMessage to the response with the results
result_message = ChatMessage(role="tool", contents=function_results) # type: ignore[call-overload]
@@ -815,11 +850,14 @@ def _handle_function_calls_streaming_response(
tools = chat_options.tools
if function_calls and tools:
# Extract function middleware pipeline from kwargs if available
middleware_pipeline = kwargs.get("_function_middleware_pipeline")
function_results = await execute_function_calls(
custom_args=kwargs,
attempt_idx=attempt_idx,
function_calls=function_calls,
tools=tools, # type: ignore[reportArgumentType]
middleware_pipeline=middleware_pipeline,
)
function_result_msg = ChatMessage(role="tool", contents=function_results)
yield ChatResponseUpdate(contents=function_results, role="tool")