mirror of
https://github.com/microsoft/agent-framework.git
synced 2026-06-16 21:04:09 +08:00
Python: [BREAKING] Refactor middleware layering and split Anthropic raw client (#4746)
* [BREAKING] Refactor middleware layering and raw clients Reorder chat client layers so function invocation wraps chat middleware, and chat middleware stays outside telemetry while still running for each inner model call. Add middleware pipeline caching, refresh docs and samples, and split Anthropic into raw and public clients to match the standard layering model. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Tighten typing ignores in ancillary modules Add targeted typing ignores in workflow visualization and lab modules so pyright stays clean alongside the middleware refactor work. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Fix categorize_middleware to unpack tuple/Sequence and use relative MRO assertions - Broaden isinstance check in categorize_middleware from list to Sequence so tuples and other Sequence types are properly unpacked instead of being appended as a single item. - Replace fragile hardcoded MRO index assertions in anthropic test with relative ordering via mro.index(). - Add regression tests for categorize_middleware with tuple, list, and None inputs. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Fix middleware string decomposition, add middleware param to FunctionInvocationLayer, and add tests (#4710) - Guard categorize_middleware Sequence check against str/bytes to prevent character-by-character decomposition of accidentally passed strings - Add explicit middleware parameter to FunctionInvocationLayer.get_response and merge it into client_kwargs before categorization, fixing the inconsistency where only OpenAIChatClient supported this parameter - Add assertions that RawAnthropicClient does not inherit convenience layers - Add chat middleware cache test with non-empty base middleware - Add tests for single unwrapped middleware item and string input Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Apply pre-commit auto-fixes * Apply pre-commit auto-fixes * Address review feedback for #4710: review comment fixes --------- Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Co-authored-by: Copilot <copilot@github.com>
This commit is contained in:
committed by
GitHub
Unverified
parent
cefda44283
commit
0cd40f8354
@@ -966,16 +966,7 @@ def _apply_get_response_docstrings() -> None:
|
||||
from .observability import ChatTelemetryLayer
|
||||
|
||||
apply_layered_docstring(ChatTelemetryLayer.get_response, BaseChatClient.get_response)
|
||||
apply_layered_docstring(
|
||||
FunctionInvocationLayer.get_response,
|
||||
ChatTelemetryLayer.get_response,
|
||||
extra_keyword_args={
|
||||
"function_middleware": """
|
||||
Optional per-call function middleware.
|
||||
When omitted, middleware configured on the client or forwarded from higher layers is used.
|
||||
""",
|
||||
},
|
||||
)
|
||||
apply_layered_docstring(FunctionInvocationLayer.get_response, ChatTelemetryLayer.get_response)
|
||||
apply_layered_docstring(
|
||||
ChatMiddlewareLayer.get_response,
|
||||
FunctionInvocationLayer.get_response,
|
||||
|
||||
@@ -742,12 +742,17 @@ class AgentMiddlewarePipeline(BaseMiddlewarePipeline):
|
||||
middleware: The list of agent middleware to include in the pipeline.
|
||||
"""
|
||||
super().__init__()
|
||||
self._source_middleware: tuple[AgentMiddlewareTypes, ...] = tuple(middleware)
|
||||
self._middleware: list[AgentMiddleware] = []
|
||||
|
||||
if middleware:
|
||||
for mdlware in middleware:
|
||||
self._register_middleware(mdlware)
|
||||
|
||||
def matches(self, middleware: Sequence[AgentMiddlewareTypes]) -> bool:
|
||||
"""Return whether this pipeline was built from the provided middleware sequence."""
|
||||
return self._source_middleware == tuple(middleware)
|
||||
|
||||
def _register_middleware(self, middleware: AgentMiddlewareTypes) -> None:
|
||||
"""Register an agent middleware item.
|
||||
|
||||
@@ -824,12 +829,17 @@ class FunctionMiddlewarePipeline(BaseMiddlewarePipeline):
|
||||
middleware: The list of function middleware to include in the pipeline.
|
||||
"""
|
||||
super().__init__()
|
||||
self._source_middleware: tuple[FunctionMiddlewareTypes, ...] = tuple(middleware)
|
||||
self._middleware: list[FunctionMiddleware] = []
|
||||
|
||||
if middleware:
|
||||
for mdlware in middleware:
|
||||
self._register_middleware(mdlware)
|
||||
|
||||
def matches(self, middleware: Sequence[FunctionMiddlewareTypes]) -> bool:
|
||||
"""Return whether this pipeline was built from the provided middleware sequence."""
|
||||
return self._source_middleware == tuple(middleware)
|
||||
|
||||
def _register_middleware(self, middleware: FunctionMiddlewareTypes) -> None:
|
||||
"""Register a function middleware item.
|
||||
|
||||
@@ -892,12 +902,17 @@ class ChatMiddlewarePipeline(BaseMiddlewarePipeline):
|
||||
middleware: The list of chat middleware to include in the pipeline.
|
||||
"""
|
||||
super().__init__()
|
||||
self._source_middleware: tuple[ChatMiddlewareTypes, ...] = tuple(middleware)
|
||||
self._middleware: list[ChatMiddleware] = []
|
||||
|
||||
if middleware:
|
||||
for mdlware in middleware:
|
||||
self._register_middleware(mdlware)
|
||||
|
||||
def matches(self, middleware: Sequence[ChatMiddlewareTypes]) -> bool:
|
||||
"""Return whether this pipeline was built from the provided middleware sequence."""
|
||||
return self._source_middleware == tuple(middleware)
|
||||
|
||||
def _register_middleware(self, middleware: ChatMiddlewareTypes) -> None:
|
||||
"""Register a chat middleware item.
|
||||
|
||||
@@ -980,16 +995,26 @@ class ChatMiddlewareLayer(Generic[OptionsCoT]):
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
middleware: Sequence[ChatAndFunctionMiddlewareTypes] | None = None,
|
||||
middleware: Sequence[ChatMiddlewareTypes] | None = None,
|
||||
**kwargs: Any,
|
||||
) -> None:
|
||||
middleware_list = categorize_middleware(*(middleware or []))
|
||||
self.chat_middleware = middleware_list["chat"]
|
||||
if "function_middleware" in kwargs and middleware_list["function"]:
|
||||
raise ValueError("Cannot specify 'function_middleware' and 'middleware' at the same time.")
|
||||
kwargs["function_middleware"] = middleware_list["function"]
|
||||
self.chat_middleware = list(middleware) if middleware else []
|
||||
self._cached_chat_middleware_pipeline: ChatMiddlewarePipeline | None = None
|
||||
super().__init__(**kwargs)
|
||||
|
||||
def _get_chat_middleware_pipeline(
|
||||
self,
|
||||
middleware: Sequence[ChatMiddlewareTypes],
|
||||
) -> ChatMiddlewarePipeline:
|
||||
effective_middleware = [*self.chat_middleware, *middleware]
|
||||
if self._cached_chat_middleware_pipeline is not None and self._cached_chat_middleware_pipeline.matches(
|
||||
effective_middleware
|
||||
):
|
||||
return self._cached_chat_middleware_pipeline
|
||||
|
||||
self._cached_chat_middleware_pipeline = ChatMiddlewarePipeline(*effective_middleware)
|
||||
return self._cached_chat_middleware_pipeline
|
||||
|
||||
@overload
|
||||
def get_response(
|
||||
self,
|
||||
@@ -1052,14 +1077,8 @@ class ChatMiddlewareLayer(Generic[OptionsCoT]):
|
||||
kwargs["tokenizer"] = tokenizer
|
||||
|
||||
effective_client_kwargs = dict(client_kwargs) if client_kwargs is not None else {}
|
||||
call_middleware = kwargs.pop("middleware", effective_client_kwargs.pop("middleware", []))
|
||||
middleware = categorize_middleware(call_middleware)
|
||||
effective_client_kwargs["function_middleware"] = middleware["function"]
|
||||
|
||||
pipeline = ChatMiddlewarePipeline(
|
||||
*self.chat_middleware,
|
||||
*middleware["chat"],
|
||||
)
|
||||
call_middleware = effective_client_kwargs.pop("middleware", [])
|
||||
pipeline = self._get_chat_middleware_pipeline(call_middleware) # type: ignore[reportUnknownArgumentType]
|
||||
if not pipeline.has_middlewares:
|
||||
return super_get_response( # type: ignore[no-any-return]
|
||||
messages=messages,
|
||||
@@ -1134,12 +1153,25 @@ class AgentMiddlewareLayer:
|
||||
) -> None:
|
||||
middleware_list = categorize_middleware(middleware)
|
||||
self.agent_middleware = middleware_list["agent"]
|
||||
self._cached_agent_middleware_pipeline: AgentMiddlewarePipeline | None = None
|
||||
# Pass middleware to super so BaseAgent can store it for dynamic rebuild
|
||||
super().__init__(*args, middleware=middleware, **kwargs) # type: ignore[call-arg]
|
||||
# Note: We intentionally don't extend client's middleware lists here.
|
||||
# Chat and function middleware is passed to the chat client at runtime via kwargs
|
||||
# in AgentMiddlewareLayer.run(), where it's properly combined with run-level middleware.
|
||||
|
||||
def _get_agent_middleware_pipeline(
|
||||
self,
|
||||
middleware: Sequence[AgentMiddlewareTypes],
|
||||
) -> AgentMiddlewarePipeline:
|
||||
if self._cached_agent_middleware_pipeline is not None and self._cached_agent_middleware_pipeline.matches(
|
||||
middleware
|
||||
):
|
||||
return self._cached_agent_middleware_pipeline
|
||||
|
||||
self._cached_agent_middleware_pipeline = AgentMiddlewarePipeline(*middleware)
|
||||
return self._cached_agent_middleware_pipeline
|
||||
|
||||
@overload
|
||||
def run(
|
||||
self,
|
||||
@@ -1210,7 +1242,7 @@ class AgentMiddlewareLayer:
|
||||
)
|
||||
base_middleware_list = categorize_middleware(base_middleware)
|
||||
run_middleware_list = categorize_middleware(middleware)
|
||||
pipeline = AgentMiddlewarePipeline(*base_middleware_list["agent"], *run_middleware_list["agent"])
|
||||
pipeline = self._get_agent_middleware_pipeline([*base_middleware_list["agent"], *run_middleware_list["agent"]])
|
||||
|
||||
# Combine base and run-level function/chat middleware for forwarding to chat client
|
||||
combined_function_chat_middleware = (
|
||||
@@ -1392,7 +1424,7 @@ def categorize_middleware(
|
||||
all_middleware: list[Any] = []
|
||||
for source in middleware_sources:
|
||||
if source:
|
||||
if isinstance(source, list):
|
||||
if isinstance(source, Sequence) and not isinstance(source, (str, bytes)):
|
||||
all_middleware.extend(source) # type: ignore
|
||||
else:
|
||||
all_middleware.append(source)
|
||||
|
||||
@@ -63,7 +63,12 @@ if TYPE_CHECKING:
|
||||
from ._clients import SupportsChatGetResponse
|
||||
from ._compaction import CompactionStrategy, TokenizerProtocol
|
||||
from ._mcp import MCPTool
|
||||
from ._middleware import FunctionInvocationContext, FunctionMiddlewarePipeline, FunctionMiddlewareTypes
|
||||
from ._middleware import (
|
||||
ChatAndFunctionMiddlewareTypes,
|
||||
FunctionInvocationContext,
|
||||
FunctionMiddlewarePipeline,
|
||||
FunctionMiddlewareTypes,
|
||||
)
|
||||
from ._sessions import AgentSession
|
||||
from ._types import (
|
||||
ChatOptions,
|
||||
@@ -2024,18 +2029,37 @@ class FunctionInvocationLayer(Generic[OptionsCoT]):
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
function_middleware: Sequence[FunctionMiddlewareTypes] | None = None,
|
||||
middleware: Sequence[ChatAndFunctionMiddlewareTypes] | None = None,
|
||||
function_invocation_configuration: FunctionInvocationConfiguration | None = None,
|
||||
**kwargs: Any,
|
||||
) -> None:
|
||||
self.function_middleware: list[FunctionMiddlewareTypes] = (
|
||||
list(function_middleware) if function_middleware else []
|
||||
)
|
||||
from ._middleware import categorize_middleware
|
||||
|
||||
middleware_list = categorize_middleware(middleware)
|
||||
self.function_middleware: list[FunctionMiddlewareTypes] = list(middleware_list["function"])
|
||||
self._cached_function_middleware_pipeline: FunctionMiddlewarePipeline | None = None
|
||||
self.function_invocation_configuration = normalize_function_invocation_configuration(
|
||||
function_invocation_configuration
|
||||
)
|
||||
if (chat_middleware := (middleware_list["chat"] or None)) is not None:
|
||||
kwargs["middleware"] = chat_middleware
|
||||
super().__init__(**kwargs)
|
||||
|
||||
def _get_function_middleware_pipeline(
|
||||
self,
|
||||
middleware: Sequence[FunctionMiddlewareTypes],
|
||||
) -> FunctionMiddlewarePipeline:
|
||||
from ._middleware import FunctionMiddlewarePipeline
|
||||
|
||||
effective_middleware = [*self.function_middleware, *middleware]
|
||||
if self._cached_function_middleware_pipeline is not None and self._cached_function_middleware_pipeline.matches(
|
||||
effective_middleware
|
||||
):
|
||||
return self._cached_function_middleware_pipeline
|
||||
|
||||
self._cached_function_middleware_pipeline = FunctionMiddlewarePipeline(*effective_middleware)
|
||||
return self._cached_function_middleware_pipeline
|
||||
|
||||
@overload
|
||||
def get_response(
|
||||
self,
|
||||
@@ -2043,6 +2067,7 @@ class FunctionInvocationLayer(Generic[OptionsCoT]):
|
||||
*,
|
||||
stream: Literal[False] = ...,
|
||||
options: ChatOptions[ResponseModelBoundT],
|
||||
middleware: Sequence[ChatAndFunctionMiddlewareTypes] | None = None,
|
||||
compaction_strategy: CompactionStrategy | None = None,
|
||||
tokenizer: TokenizerProtocol | None = None,
|
||||
function_invocation_kwargs: Mapping[str, Any] | None = None,
|
||||
@@ -2057,6 +2082,7 @@ class FunctionInvocationLayer(Generic[OptionsCoT]):
|
||||
*,
|
||||
stream: Literal[False] = ...,
|
||||
options: OptionsCoT | ChatOptions[None] | None = None,
|
||||
middleware: Sequence[ChatAndFunctionMiddlewareTypes] | None = None,
|
||||
compaction_strategy: CompactionStrategy | None = None,
|
||||
tokenizer: TokenizerProtocol | None = None,
|
||||
function_invocation_kwargs: Mapping[str, Any] | None = None,
|
||||
@@ -2071,6 +2097,7 @@ class FunctionInvocationLayer(Generic[OptionsCoT]):
|
||||
*,
|
||||
stream: Literal[True],
|
||||
options: OptionsCoT | ChatOptions[Any] | None = None,
|
||||
middleware: Sequence[ChatAndFunctionMiddlewareTypes] | None = None,
|
||||
compaction_strategy: CompactionStrategy | None = None,
|
||||
tokenizer: TokenizerProtocol | None = None,
|
||||
function_invocation_kwargs: Mapping[str, Any] | None = None,
|
||||
@@ -2084,14 +2111,14 @@ class FunctionInvocationLayer(Generic[OptionsCoT]):
|
||||
*,
|
||||
stream: bool = False,
|
||||
options: OptionsCoT | ChatOptions[Any] | None = None,
|
||||
function_middleware: Sequence[FunctionMiddlewareTypes] | None = None,
|
||||
middleware: Sequence[ChatAndFunctionMiddlewareTypes] | None = None,
|
||||
compaction_strategy: CompactionStrategy | None = None,
|
||||
tokenizer: TokenizerProtocol | None = None,
|
||||
function_invocation_kwargs: Mapping[str, Any] | None = None,
|
||||
client_kwargs: Mapping[str, Any] | None = None,
|
||||
**kwargs: Any,
|
||||
) -> Awaitable[ChatResponse[Any]] | ResponseStream[ChatResponseUpdate, ChatResponse[Any]]:
|
||||
from ._middleware import FunctionMiddlewarePipeline
|
||||
from ._middleware import categorize_middleware
|
||||
from ._types import (
|
||||
ChatResponse,
|
||||
ChatResponseUpdate,
|
||||
@@ -2109,16 +2136,21 @@ class FunctionInvocationLayer(Generic[OptionsCoT]):
|
||||
)
|
||||
|
||||
effective_client_kwargs = dict(client_kwargs) if client_kwargs is not None else {}
|
||||
effective_function_middleware = function_middleware
|
||||
if effective_function_middleware is None:
|
||||
middleware_from_client_kwargs = effective_client_kwargs.pop("function_middleware", None)
|
||||
if middleware_from_client_kwargs is not None:
|
||||
effective_function_middleware = cast(Sequence[Any], middleware_from_client_kwargs)
|
||||
if middleware is not None:
|
||||
existing = effective_client_kwargs.get("middleware", [])
|
||||
effective_client_kwargs["middleware"] = [
|
||||
*(
|
||||
existing
|
||||
if isinstance(existing, Sequence) and not isinstance(existing, (str, bytes))
|
||||
else [existing]
|
||||
),
|
||||
*middleware,
|
||||
]
|
||||
runtime_middleware = categorize_middleware(effective_client_kwargs.pop("middleware", []))
|
||||
|
||||
# ChatMiddleware adds this kwarg
|
||||
function_middleware_pipeline = FunctionMiddlewarePipeline(
|
||||
*(self.function_middleware), *(effective_function_middleware or [])
|
||||
)
|
||||
function_middleware_pipeline = self._get_function_middleware_pipeline(runtime_middleware["function"])
|
||||
if runtime_middleware["chat"]:
|
||||
effective_client_kwargs["middleware"] = runtime_middleware["chat"]
|
||||
max_errors = self.function_invocation_configuration.get(
|
||||
"max_consecutive_errors_per_request", DEFAULT_MAX_CONSECUTIVE_ERRORS_PER_REQUEST
|
||||
)
|
||||
|
||||
@@ -109,7 +109,7 @@ class WorkflowViz:
|
||||
|
||||
# Create a temporary graphviz Source object
|
||||
dot_content = self.to_digraph(include_internal_executors=include_internal_executors)
|
||||
source = graphviz.Source(dot_content)
|
||||
source = graphviz.Source(dot_content) # type: ignore[reportUnknownVariableType]
|
||||
|
||||
try:
|
||||
if filename:
|
||||
@@ -131,7 +131,7 @@ class WorkflowViz:
|
||||
|
||||
source.render(base_name, format=format, cleanup=True) # type: ignore
|
||||
return f"{base_name}.{format}"
|
||||
except graphviz.backend.execute.ExecutableNotFound as e:
|
||||
except graphviz.backend.execute.ExecutableNotFound as e: # type: ignore
|
||||
raise ImportError(
|
||||
"The graphviz executables are not found. The graphviz Python package is installed, but the "
|
||||
"graphviz executables (dot, neato, etc.) are not available on your system's PATH. "
|
||||
|
||||
@@ -152,8 +152,8 @@ AzureOpenAIChatClientT = TypeVar("AzureOpenAIChatClientT", bound="AzureOpenAICha
|
||||
|
||||
class AzureOpenAIChatClient( # type: ignore[misc]
|
||||
AzureOpenAIConfigMixin,
|
||||
ChatMiddlewareLayer[AzureOpenAIChatOptionsT],
|
||||
FunctionInvocationLayer[AzureOpenAIChatOptionsT],
|
||||
ChatMiddlewareLayer[AzureOpenAIChatOptionsT],
|
||||
ChatTelemetryLayer[AzureOpenAIChatOptionsT],
|
||||
RawOpenAIChatClient[AzureOpenAIChatOptionsT],
|
||||
Generic[AzureOpenAIChatOptionsT],
|
||||
|
||||
@@ -51,8 +51,8 @@ AzureOpenAIResponsesOptionsT = TypeVar(
|
||||
|
||||
class AzureOpenAIResponsesClient( # type: ignore[misc]
|
||||
AzureOpenAIConfigMixin,
|
||||
ChatMiddlewareLayer[AzureOpenAIResponsesOptionsT],
|
||||
FunctionInvocationLayer[AzureOpenAIResponsesOptionsT],
|
||||
ChatMiddlewareLayer[AzureOpenAIResponsesOptionsT],
|
||||
ChatTelemetryLayer[AzureOpenAIResponsesOptionsT],
|
||||
RawOpenAIResponsesClient[AzureOpenAIResponsesOptionsT],
|
||||
Generic[AzureOpenAIResponsesOptionsT],
|
||||
|
||||
@@ -362,11 +362,15 @@ def _create_otlp_exporters(
|
||||
if protocol == "grpc":
|
||||
# Import all gRPC exporters
|
||||
try:
|
||||
from opentelemetry.exporter.otlp.proto.grpc._log_exporter import OTLPLogExporter as GRPCLogExporter
|
||||
from opentelemetry.exporter.otlp.proto.grpc.metric_exporter import (
|
||||
OTLPMetricExporter as GRPCMetricExporter,
|
||||
from opentelemetry.exporter.otlp.proto.grpc._log_exporter import ( # type: ignore[reportMissingImports]
|
||||
OTLPLogExporter as GRPCLogExporter, # type: ignore[reportUnknownVariableType]
|
||||
)
|
||||
from opentelemetry.exporter.otlp.proto.grpc.metric_exporter import ( # type: ignore[reportMissingImports]
|
||||
OTLPMetricExporter as GRPCMetricExporter, # type: ignore[reportUnknownVariableType]
|
||||
)
|
||||
from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import ( # type: ignore[reportMissingImports]
|
||||
OTLPSpanExporter as GRPCSpanExporter, # type: ignore[reportUnknownVariableType]
|
||||
)
|
||||
from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter as GRPCSpanExporter
|
||||
except ImportError as exc:
|
||||
raise ImportError(
|
||||
"opentelemetry-exporter-otlp-proto-grpc is required for OTLP gRPC exporters. "
|
||||
@@ -375,21 +379,21 @@ def _create_otlp_exporters(
|
||||
|
||||
if actual_logs_endpoint:
|
||||
exporters.append(
|
||||
GRPCLogExporter(
|
||||
GRPCLogExporter( # type: ignore[reportUnknownArgumentType]
|
||||
endpoint=actual_logs_endpoint,
|
||||
headers=actual_logs_headers if actual_logs_headers else None,
|
||||
)
|
||||
)
|
||||
if actual_traces_endpoint:
|
||||
exporters.append(
|
||||
GRPCSpanExporter(
|
||||
GRPCSpanExporter( # type: ignore[reportUnknownArgumentType]
|
||||
endpoint=actual_traces_endpoint,
|
||||
headers=actual_traces_headers if actual_traces_headers else None,
|
||||
)
|
||||
)
|
||||
if actual_metrics_endpoint:
|
||||
exporters.append(
|
||||
GRPCMetricExporter(
|
||||
GRPCMetricExporter( # type: ignore[reportUnknownArgumentType]
|
||||
endpoint=actual_metrics_endpoint,
|
||||
headers=actual_metrics_headers if actual_metrics_headers else None,
|
||||
)
|
||||
|
||||
@@ -210,8 +210,8 @@ OpenAIAssistantsOptionsT = TypeVar(
|
||||
|
||||
class OpenAIAssistantsClient( # type: ignore[misc]
|
||||
OpenAIConfigMixin,
|
||||
ChatMiddlewareLayer[OpenAIAssistantsOptionsT],
|
||||
FunctionInvocationLayer[OpenAIAssistantsOptionsT],
|
||||
ChatMiddlewareLayer[OpenAIAssistantsOptionsT],
|
||||
ChatTelemetryLayer[OpenAIAssistantsOptionsT],
|
||||
BaseChatClient[OpenAIAssistantsOptionsT],
|
||||
Generic[OpenAIAssistantsOptionsT],
|
||||
|
||||
@@ -31,7 +31,7 @@ from pydantic import BaseModel
|
||||
|
||||
from .._clients import BaseChatClient
|
||||
from .._docstrings import apply_layered_docstring
|
||||
from .._middleware import ChatAndFunctionMiddlewareTypes, ChatMiddlewareLayer, FunctionMiddlewareTypes
|
||||
from .._middleware import ChatAndFunctionMiddlewareTypes, ChatMiddlewareLayer
|
||||
from .._settings import load_settings
|
||||
from .._tools import (
|
||||
FunctionInvocationConfiguration,
|
||||
@@ -156,9 +156,9 @@ class RawOpenAIChatClient( # type: ignore[misc]
|
||||
you should consider which additional layers to apply. There is a defined ordering that
|
||||
you should follow:
|
||||
|
||||
1. **ChatMiddlewareLayer** - Should be applied first as it also prepares function middleware
|
||||
2. **FunctionInvocationLayer** - Handles tool/function calling loop
|
||||
3. **ChatTelemetryLayer** - Must be inside the function calling loop for correct per-call telemetry
|
||||
1. **FunctionInvocationLayer** - Owns the tool/function calling loop and routes function middleware
|
||||
2. **ChatMiddlewareLayer** - Applies chat middleware per model call and stays outside telemetry
|
||||
3. **ChatTelemetryLayer** - Must stay inside chat middleware for correct per-call telemetry
|
||||
|
||||
Use ``OpenAIChatClient`` instead for a fully-featured client with all layers applied.
|
||||
"""
|
||||
@@ -776,8 +776,8 @@ class RawOpenAIChatClient( # type: ignore[misc]
|
||||
|
||||
class OpenAIChatClient( # type: ignore[misc]
|
||||
OpenAIConfigMixin,
|
||||
ChatMiddlewareLayer[OpenAIChatOptionsT],
|
||||
FunctionInvocationLayer[OpenAIChatOptionsT],
|
||||
ChatMiddlewareLayer[OpenAIChatOptionsT],
|
||||
ChatTelemetryLayer[OpenAIChatOptionsT],
|
||||
RawOpenAIChatClient[OpenAIChatOptionsT],
|
||||
Generic[OpenAIChatOptionsT],
|
||||
@@ -791,7 +791,6 @@ class OpenAIChatClient( # type: ignore[misc]
|
||||
*,
|
||||
stream: Literal[False] = ...,
|
||||
options: ChatOptions[ResponseModelBoundT],
|
||||
function_middleware: Sequence[FunctionMiddlewareTypes] | None = None,
|
||||
function_invocation_kwargs: Mapping[str, Any] | None = None,
|
||||
client_kwargs: Mapping[str, Any] | None = None,
|
||||
middleware: Sequence[ChatAndFunctionMiddlewareTypes] | None = None,
|
||||
@@ -805,7 +804,6 @@ class OpenAIChatClient( # type: ignore[misc]
|
||||
*,
|
||||
stream: Literal[False] = ...,
|
||||
options: OpenAIChatOptionsT | ChatOptions[None] | None = None,
|
||||
function_middleware: Sequence[FunctionMiddlewareTypes] | None = None,
|
||||
function_invocation_kwargs: Mapping[str, Any] | None = None,
|
||||
client_kwargs: Mapping[str, Any] | None = None,
|
||||
middleware: Sequence[ChatAndFunctionMiddlewareTypes] | None = None,
|
||||
@@ -819,7 +817,6 @@ class OpenAIChatClient( # type: ignore[misc]
|
||||
*,
|
||||
stream: Literal[True],
|
||||
options: OpenAIChatOptionsT | ChatOptions[Any] | None = None,
|
||||
function_middleware: Sequence[FunctionMiddlewareTypes] | None = None,
|
||||
function_invocation_kwargs: Mapping[str, Any] | None = None,
|
||||
client_kwargs: Mapping[str, Any] | None = None,
|
||||
middleware: Sequence[ChatAndFunctionMiddlewareTypes] | None = None,
|
||||
@@ -833,7 +830,6 @@ class OpenAIChatClient( # type: ignore[misc]
|
||||
*,
|
||||
stream: bool = False,
|
||||
options: OpenAIChatOptionsT | ChatOptions[Any] | None = None,
|
||||
function_middleware: Sequence[FunctionMiddlewareTypes] | None = None,
|
||||
function_invocation_kwargs: Mapping[str, Any] | None = None,
|
||||
client_kwargs: Mapping[str, Any] | None = None,
|
||||
middleware: Sequence[ChatAndFunctionMiddlewareTypes] | None = None,
|
||||
@@ -844,14 +840,15 @@ class OpenAIChatClient( # type: ignore[misc]
|
||||
"Callable[..., Awaitable[ChatResponse[Any]] | ResponseStream[ChatResponseUpdate, ChatResponse[Any]]]",
|
||||
super().get_response, # type: ignore[misc]
|
||||
)
|
||||
effective_client_kwargs = dict(client_kwargs) if client_kwargs is not None else {}
|
||||
if middleware is not None:
|
||||
effective_client_kwargs["middleware"] = middleware
|
||||
return super_get_response( # type: ignore[no-any-return]
|
||||
messages=messages,
|
||||
stream=stream,
|
||||
options=options,
|
||||
function_middleware=function_middleware,
|
||||
function_invocation_kwargs=function_invocation_kwargs,
|
||||
client_kwargs=client_kwargs,
|
||||
middleware=middleware,
|
||||
client_kwargs=effective_client_kwargs,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
@@ -967,10 +964,6 @@ def _apply_openai_chat_client_docstrings() -> None:
|
||||
OpenAIChatClient.get_response,
|
||||
RawOpenAIChatClient.get_response,
|
||||
extra_keyword_args={
|
||||
"function_middleware": """
|
||||
Optional per-call function middleware.
|
||||
When omitted, middleware configured on the client or forwarded from higher layers is used.
|
||||
""",
|
||||
"middleware": """
|
||||
Optional per-call chat and function middleware.
|
||||
This is merged with any middleware configured on the client for the current request.
|
||||
|
||||
@@ -249,9 +249,9 @@ class RawOpenAIResponsesClient( # type: ignore[misc]
|
||||
you should consider which additional layers to apply. There is a defined ordering that
|
||||
you should follow:
|
||||
|
||||
1. **ChatMiddlewareLayer** - Should be applied first as it also prepares function middleware
|
||||
2. **FunctionInvocationLayer** - Handles tool/function calling loop
|
||||
3. **ChatTelemetryLayer** - Must be inside the function calling loop for correct per-call telemetry
|
||||
1. **FunctionInvocationLayer** - Owns the tool/function calling loop and routes function middleware
|
||||
2. **ChatMiddlewareLayer** - Applies chat middleware per model call and stays outside telemetry
|
||||
3. **ChatTelemetryLayer** - Must stay inside chat middleware for correct per-call telemetry
|
||||
|
||||
Use ``OpenAIResponsesClient`` instead for a fully-featured client with all layers applied.
|
||||
"""
|
||||
@@ -2259,8 +2259,8 @@ class RawOpenAIResponsesClient( # type: ignore[misc]
|
||||
|
||||
class OpenAIResponsesClient( # type: ignore[misc]
|
||||
OpenAIConfigMixin,
|
||||
ChatMiddlewareLayer[OpenAIResponsesOptionsT],
|
||||
FunctionInvocationLayer[OpenAIResponsesOptionsT],
|
||||
ChatMiddlewareLayer[OpenAIResponsesOptionsT],
|
||||
ChatTelemetryLayer[OpenAIResponsesOptionsT],
|
||||
RawOpenAIResponsesClient[OpenAIResponsesOptionsT],
|
||||
Generic[OpenAIResponsesOptionsT],
|
||||
|
||||
@@ -128,8 +128,8 @@ class MockChatClient:
|
||||
|
||||
|
||||
class MockBaseChatClient(
|
||||
ChatMiddlewareLayer[OptionsCoT],
|
||||
FunctionInvocationLayer[OptionsCoT],
|
||||
ChatMiddlewareLayer[OptionsCoT],
|
||||
ChatTelemetryLayer[OptionsCoT],
|
||||
BaseChatClient[OptionsCoT],
|
||||
Generic[OptionsCoT],
|
||||
@@ -137,7 +137,7 @@ class MockBaseChatClient(
|
||||
"""Mock implementation of a full-featured ChatClient."""
|
||||
|
||||
def __init__(self, **kwargs: Any):
|
||||
super().__init__(function_middleware=[], **kwargs)
|
||||
super().__init__(middleware=[], **kwargs)
|
||||
self.run_responses: list[ChatResponse] = []
|
||||
self.streaming_responses: list[list[ChatResponseUpdate]] = []
|
||||
self.call_count: int = 0
|
||||
|
||||
@@ -74,8 +74,8 @@ def test_openai_chat_client_get_response_docstring_surfaces_layered_runtime_docs
|
||||
assert docstring is not None
|
||||
assert "Get a response from a chat client." in docstring
|
||||
assert "function_invocation_kwargs" in docstring
|
||||
assert "function_middleware: Optional per-call function middleware." in docstring
|
||||
assert "middleware: Optional per-call chat and function middleware." in docstring
|
||||
assert "function_middleware: Optional per-call function middleware." not in docstring
|
||||
|
||||
|
||||
def test_openai_chat_client_get_response_is_defined_on_openai_class() -> None:
|
||||
@@ -84,7 +84,6 @@ def test_openai_chat_client_get_response_is_defined_on_openai_class() -> None:
|
||||
signature = inspect.signature(OpenAIChatClient.get_response)
|
||||
|
||||
assert OpenAIChatClient.get_response.__qualname__ == "OpenAIChatClient.get_response"
|
||||
assert "function_middleware" in signature.parameters
|
||||
assert "middleware" in signature.parameters
|
||||
|
||||
|
||||
|
||||
@@ -3226,7 +3226,7 @@ async def test_terminate_loop_single_function_call(chat_client_base: SupportsCha
|
||||
response = await chat_client_base.get_response(
|
||||
"hello",
|
||||
options={"tool_choice": "auto", "tools": [ai_func]},
|
||||
middleware=[TerminateLoopMiddleware()],
|
||||
client_kwargs={"middleware": [TerminateLoopMiddleware()]},
|
||||
)
|
||||
|
||||
# Function should NOT have been executed - middleware intercepted it
|
||||
@@ -3292,7 +3292,7 @@ async def test_terminate_loop_multiple_function_calls_one_terminates(chat_client
|
||||
response = await chat_client_base.get_response(
|
||||
"hello",
|
||||
options={"tool_choice": "auto", "tools": [normal_func, terminating_func]},
|
||||
middleware=[SelectiveTerminateMiddleware()],
|
||||
client_kwargs={"middleware": [SelectiveTerminateMiddleware()]},
|
||||
)
|
||||
|
||||
# normal_function should have executed (middleware calls next_handler)
|
||||
@@ -3345,7 +3345,7 @@ async def test_terminate_loop_streaming_single_function_call(chat_client_base: S
|
||||
async for update in chat_client_base.get_response(
|
||||
"hello",
|
||||
options={"tool_choice": "auto", "tools": [ai_func]},
|
||||
middleware=[TerminateLoopMiddleware()],
|
||||
client_kwargs={"middleware": [TerminateLoopMiddleware()]},
|
||||
stream=True,
|
||||
):
|
||||
updates.append(update)
|
||||
@@ -3389,12 +3389,12 @@ async def test_conversation_id_updated_in_options_between_tool_iterations():
|
||||
conversation_ids_received: list[str | None] = []
|
||||
|
||||
class TrackingChatClient(
|
||||
ChatMiddlewareLayer,
|
||||
FunctionInvocationLayer,
|
||||
ChatMiddlewareLayer,
|
||||
BaseChatClient,
|
||||
):
|
||||
def __init__(self) -> None:
|
||||
super().__init__(function_middleware=[])
|
||||
super().__init__(middleware=[])
|
||||
self.run_responses: list[ChatResponse] = []
|
||||
self.streaming_responses: list[list[ChatResponseUpdate]] = []
|
||||
self.call_count: int = 0
|
||||
|
||||
@@ -84,8 +84,8 @@ class _MockBaseChatClient(BaseChatClient[Any]):
|
||||
|
||||
|
||||
class FunctionInvokingMockClient(
|
||||
ChatMiddlewareLayer[Any],
|
||||
FunctionInvocationLayer[Any],
|
||||
ChatMiddlewareLayer[Any],
|
||||
ChatTelemetryLayer[Any],
|
||||
_MockBaseChatClient,
|
||||
):
|
||||
|
||||
@@ -28,6 +28,7 @@ from agent_framework._middleware import (
|
||||
FunctionMiddleware,
|
||||
FunctionMiddlewarePipeline,
|
||||
MiddlewareTermination,
|
||||
categorize_middleware,
|
||||
)
|
||||
from agent_framework._tools import FunctionTool
|
||||
|
||||
@@ -1681,3 +1682,49 @@ def mock_chat_client() -> Any:
|
||||
client = MagicMock(spec=SupportsChatGetResponse)
|
||||
client.service_url = MagicMock(return_value="mock://test")
|
||||
return client
|
||||
|
||||
|
||||
class TestCategorizeMiddleware:
|
||||
"""Test cases for categorize_middleware."""
|
||||
|
||||
def test_categorize_middleware_with_tuple(self) -> None:
|
||||
"""Test that tuple middleware sources are unpacked, not appended as a single item."""
|
||||
chat_mw = TestChatMiddleware()
|
||||
function_mw = TestFunctionMiddleware()
|
||||
agent_mw = TestAgentMiddleware()
|
||||
result = categorize_middleware((chat_mw, function_mw, agent_mw))
|
||||
assert result["chat"] == [chat_mw]
|
||||
assert result["function"] == [function_mw]
|
||||
assert result["agent"] == [agent_mw]
|
||||
|
||||
def test_categorize_middleware_with_list(self) -> None:
|
||||
"""Test that list middleware sources are unpacked correctly."""
|
||||
chat_mw = TestChatMiddleware()
|
||||
function_mw = TestFunctionMiddleware()
|
||||
result = categorize_middleware([chat_mw, function_mw])
|
||||
assert result["chat"] == [chat_mw]
|
||||
assert result["function"] == [function_mw]
|
||||
assert result["agent"] == []
|
||||
|
||||
def test_categorize_middleware_with_none(self) -> None:
|
||||
"""Test that None middleware sources are handled."""
|
||||
result = categorize_middleware(None)
|
||||
assert result["chat"] == []
|
||||
assert result["function"] == []
|
||||
assert result["agent"] == []
|
||||
|
||||
def test_categorize_middleware_with_single_item(self) -> None:
|
||||
"""Test that a single unwrapped middleware item is appended correctly."""
|
||||
chat_mw = TestChatMiddleware()
|
||||
result = categorize_middleware(chat_mw)
|
||||
assert result["chat"] == [chat_mw]
|
||||
assert result["function"] == []
|
||||
assert result["agent"] == []
|
||||
|
||||
def test_categorize_middleware_with_string_does_not_decompose(self) -> None:
|
||||
"""Test that a string is not decomposed character-by-character."""
|
||||
result = categorize_middleware("not_a_middleware")
|
||||
# String should be treated as a single item, not decomposed into characters
|
||||
total_items = len(result["chat"]) + len(result["function"]) + len(result["agent"])
|
||||
assert total_items == 1
|
||||
assert result["agent"] == ["not_a_middleware"]
|
||||
|
||||
@@ -697,6 +697,26 @@ class TestChatAgentFunctionMiddlewareWithTools:
|
||||
assert function_calls[0].name == "sample_tool_function"
|
||||
assert function_results[0].call_id == function_calls[0].call_id
|
||||
|
||||
def test_agent_middleware_pipeline_cache_reuses_matching_middleware(self) -> None:
|
||||
"""Test that identical agent middleware sets reuse the cached pipeline."""
|
||||
|
||||
@agent_middleware
|
||||
async def first_middleware(context: AgentContext, call_next: Callable[[], Awaitable[None]]) -> None:
|
||||
await call_next()
|
||||
|
||||
@agent_middleware
|
||||
async def second_middleware(context: AgentContext, call_next: Callable[[], Awaitable[None]]) -> None:
|
||||
await call_next()
|
||||
|
||||
agent = Agent(client=MockBaseChatClient())
|
||||
|
||||
first_pipeline = agent._get_agent_middleware_pipeline([first_middleware])
|
||||
second_pipeline = agent._get_agent_middleware_pipeline([first_middleware])
|
||||
third_pipeline = agent._get_agent_middleware_pipeline([second_middleware])
|
||||
|
||||
assert first_pipeline is second_pipeline
|
||||
assert third_pipeline is not first_pipeline
|
||||
|
||||
async def test_function_middleware_can_access_and_override_custom_kwargs(
|
||||
self, chat_client_base: "MockBaseChatClient"
|
||||
) -> None:
|
||||
@@ -1969,6 +1989,77 @@ class TestChatAgentChatMiddleware:
|
||||
"agent_middleware_after",
|
||||
]
|
||||
|
||||
async def test_combined_middleware_with_tool_loop(self) -> None:
|
||||
"""Test Agent middleware ordering when tool calls trigger multiple chat rounds."""
|
||||
execution_order: list[str] = []
|
||||
chat_round = 0
|
||||
client = MockBaseChatClient()
|
||||
client.run_responses = [
|
||||
ChatResponse(
|
||||
messages=[
|
||||
Message(
|
||||
role="assistant",
|
||||
contents=[
|
||||
Content.from_function_call(
|
||||
call_id="call_123",
|
||||
name="sample_tool_function",
|
||||
arguments='{"location": "Seattle"}',
|
||||
)
|
||||
],
|
||||
)
|
||||
]
|
||||
),
|
||||
ChatResponse(messages=[Message(role="assistant", text="Final response")]),
|
||||
]
|
||||
|
||||
async def tracking_agent_middleware(
|
||||
context: AgentContext,
|
||||
call_next: Callable[[], Awaitable[None]],
|
||||
) -> None:
|
||||
execution_order.append("agent_middleware_before")
|
||||
await call_next()
|
||||
execution_order.append("agent_middleware_after")
|
||||
|
||||
async def tracking_chat_middleware(
|
||||
context: ChatContext,
|
||||
call_next: Callable[[], Awaitable[None]],
|
||||
) -> None:
|
||||
nonlocal chat_round
|
||||
chat_round += 1
|
||||
execution_order.append(f"chat_middleware_before_{chat_round}")
|
||||
await call_next()
|
||||
execution_order.append(f"chat_middleware_after_{chat_round}")
|
||||
|
||||
async def tracking_function_middleware(
|
||||
context: FunctionInvocationContext,
|
||||
call_next: Callable[[], Awaitable[None]],
|
||||
) -> None:
|
||||
execution_order.append("function_middleware_before")
|
||||
await call_next()
|
||||
execution_order.append("function_middleware_after")
|
||||
|
||||
agent = Agent(
|
||||
client=client,
|
||||
middleware=[tracking_chat_middleware, tracking_function_middleware, tracking_agent_middleware],
|
||||
tools=[sample_tool_function],
|
||||
)
|
||||
|
||||
response = await agent.run([Message(role="user", text="test")])
|
||||
|
||||
assert response is not None
|
||||
assert client.call_count == 2
|
||||
assert response.messages[-1].text == "Final response"
|
||||
assert execution_order == [
|
||||
"agent_middleware_before",
|
||||
"chat_middleware_before_1",
|
||||
"chat_middleware_after_1",
|
||||
"function_middleware_before",
|
||||
"function_middleware_after",
|
||||
"chat_middleware_before_2",
|
||||
"chat_middleware_after_2",
|
||||
"agent_middleware_after",
|
||||
]
|
||||
|
||||
async def test_agent_middleware_can_access_and_override_custom_kwargs(self) -> None:
|
||||
"""Test that agent middleware can access and override custom parameters like temperature."""
|
||||
captured_kwargs: dict[str, Any] = {}
|
||||
|
||||
@@ -274,7 +274,10 @@ class TestChatMiddleware:
|
||||
|
||||
# First call with run-level middleware
|
||||
messages = [Message(role="user", text="first message")]
|
||||
response1 = await chat_client_base.get_response(messages, middleware=[counting_middleware])
|
||||
response1 = await chat_client_base.get_response(
|
||||
messages,
|
||||
client_kwargs={"middleware": [counting_middleware]},
|
||||
)
|
||||
assert response1 is not None
|
||||
assert execution_count["count"] == 1
|
||||
|
||||
@@ -286,7 +289,10 @@ class TestChatMiddleware:
|
||||
|
||||
# Third call with run-level middleware again - should execute
|
||||
messages = [Message(role="user", text="third message")]
|
||||
response3 = await chat_client_base.get_response(messages, middleware=[counting_middleware])
|
||||
response3 = await chat_client_base.get_response(
|
||||
messages,
|
||||
client_kwargs={"middleware": [counting_middleware]},
|
||||
)
|
||||
assert response3 is not None
|
||||
assert execution_count["count"] == 2 # Should be 2 now
|
||||
|
||||
@@ -335,6 +341,81 @@ class TestChatMiddleware:
|
||||
assert modified_kwargs["new_param"] == "added_by_middleware"
|
||||
assert modified_kwargs["custom_param"] == "test_value" # Should still be there
|
||||
|
||||
def test_chat_middleware_pipeline_cache_reuses_matching_middleware(
|
||||
self,
|
||||
chat_client_base: "MockBaseChatClient",
|
||||
) -> None:
|
||||
"""Test that identical chat middleware sets reuse the cached pipeline."""
|
||||
|
||||
@chat_middleware
|
||||
async def first_middleware(context: ChatContext, call_next: Callable[[], Awaitable[None]]) -> None:
|
||||
await call_next()
|
||||
|
||||
@chat_middleware
|
||||
async def second_middleware(context: ChatContext, call_next: Callable[[], Awaitable[None]]) -> None:
|
||||
await call_next()
|
||||
|
||||
first_pipeline = chat_client_base._get_chat_middleware_pipeline([first_middleware])
|
||||
second_pipeline = chat_client_base._get_chat_middleware_pipeline([first_middleware])
|
||||
third_pipeline = chat_client_base._get_chat_middleware_pipeline([second_middleware])
|
||||
|
||||
assert first_pipeline is second_pipeline
|
||||
assert third_pipeline is not first_pipeline
|
||||
|
||||
def test_chat_middleware_pipeline_cache_includes_base_middleware(
|
||||
self,
|
||||
chat_client_base: "MockBaseChatClient",
|
||||
) -> None:
|
||||
"""Test that chat middleware cache key includes base middleware to prevent incorrect reuse."""
|
||||
|
||||
@chat_middleware
|
||||
async def base_middleware(context: ChatContext, call_next: Callable[[], Awaitable[None]]) -> None:
|
||||
await call_next()
|
||||
|
||||
@chat_middleware
|
||||
async def runtime_middleware(context: ChatContext, call_next: Callable[[], Awaitable[None]]) -> None:
|
||||
await call_next()
|
||||
|
||||
# Without base middleware
|
||||
pipeline_no_base = chat_client_base._get_chat_middleware_pipeline([runtime_middleware])
|
||||
|
||||
# With base middleware
|
||||
chat_client_base.chat_middleware = [base_middleware]
|
||||
pipeline_with_base = chat_client_base._get_chat_middleware_pipeline([runtime_middleware])
|
||||
|
||||
assert pipeline_with_base is not pipeline_no_base
|
||||
|
||||
def test_function_middleware_pipeline_cache_reuses_matching_middleware(
|
||||
self,
|
||||
chat_client_base: "MockBaseChatClient",
|
||||
) -> None:
|
||||
"""Test that identical function middleware sets reuse the cached pipeline."""
|
||||
|
||||
@function_middleware
|
||||
async def base_middleware(context: FunctionInvocationContext, call_next: Callable[[], Awaitable[None]]) -> None:
|
||||
await call_next()
|
||||
|
||||
@function_middleware
|
||||
async def first_runtime_middleware(
|
||||
context: FunctionInvocationContext, call_next: Callable[[], Awaitable[None]]
|
||||
) -> None:
|
||||
await call_next()
|
||||
|
||||
@function_middleware
|
||||
async def second_runtime_middleware(
|
||||
context: FunctionInvocationContext, call_next: Callable[[], Awaitable[None]]
|
||||
) -> None:
|
||||
await call_next()
|
||||
|
||||
chat_client_base.function_middleware = [base_middleware]
|
||||
|
||||
first_pipeline = chat_client_base._get_function_middleware_pipeline([first_runtime_middleware])
|
||||
second_pipeline = chat_client_base._get_function_middleware_pipeline([first_runtime_middleware])
|
||||
third_pipeline = chat_client_base._get_function_middleware_pipeline([second_runtime_middleware])
|
||||
|
||||
assert first_pipeline is second_pipeline
|
||||
assert third_pipeline is not first_pipeline
|
||||
|
||||
async def test_function_middleware_registration_on_chat_client(
|
||||
self, chat_client_base: "MockBaseChatClient"
|
||||
) -> None:
|
||||
@@ -450,7 +531,9 @@ class TestChatMiddleware:
|
||||
# Execute the chat client directly with run-level middleware and tools
|
||||
messages = [Message(role="user", text="What's the weather in New York?")]
|
||||
response = await client.get_response(
|
||||
messages, options={"tools": [sample_tool_wrapped]}, middleware=[run_level_function_middleware]
|
||||
messages,
|
||||
options={"tools": [sample_tool_wrapped]},
|
||||
client_kwargs={"middleware": [run_level_function_middleware]},
|
||||
)
|
||||
|
||||
# Verify response
|
||||
@@ -463,3 +546,156 @@ class TestChatMiddleware:
|
||||
"run_level_function_middleware_before",
|
||||
"run_level_function_middleware_after",
|
||||
]
|
||||
|
||||
async def test_run_level_chat_and_function_middleware_split_per_function_loop_round(self) -> None:
|
||||
"""Test mixed run-level middleware is split so chat middleware runs per model call."""
|
||||
execution_order: list[str] = []
|
||||
chat_round = 0
|
||||
|
||||
@chat_middleware
|
||||
async def run_level_chat_middleware(
|
||||
context: ChatContext,
|
||||
call_next: Callable[[], Awaitable[None]],
|
||||
) -> None:
|
||||
nonlocal chat_round
|
||||
chat_round += 1
|
||||
execution_order.append(f"chat_middleware_before_{chat_round}")
|
||||
await call_next()
|
||||
execution_order.append(f"chat_middleware_after_{chat_round}")
|
||||
|
||||
@function_middleware
|
||||
async def run_level_function_middleware(
|
||||
context: FunctionInvocationContext,
|
||||
call_next: Callable[[], Awaitable[None]],
|
||||
) -> None:
|
||||
execution_order.append("function_middleware_before")
|
||||
await call_next()
|
||||
execution_order.append("function_middleware_after")
|
||||
|
||||
def sample_tool(location: str) -> str:
|
||||
"""Get weather for a location."""
|
||||
return f"Weather in {location}: sunny"
|
||||
|
||||
sample_tool_wrapped = FunctionTool(
|
||||
func=sample_tool,
|
||||
name="sample_tool",
|
||||
description="Get weather for a location",
|
||||
approval_mode="never_require",
|
||||
)
|
||||
|
||||
client = MockBaseChatClient()
|
||||
client.run_responses = [
|
||||
ChatResponse(
|
||||
messages=[
|
||||
Message(
|
||||
role="assistant",
|
||||
contents=[
|
||||
Content.from_function_call(
|
||||
call_id="call_3",
|
||||
name="sample_tool",
|
||||
arguments={"location": "Seattle"},
|
||||
)
|
||||
],
|
||||
)
|
||||
]
|
||||
),
|
||||
ChatResponse(messages=[Message(role="assistant", text="Based on the weather data, it's sunny!")]),
|
||||
]
|
||||
|
||||
response = await client.get_response(
|
||||
[Message(role="user", text="What's the weather in Seattle?")],
|
||||
options={"tools": [sample_tool_wrapped]},
|
||||
client_kwargs={"middleware": [run_level_chat_middleware, run_level_function_middleware]},
|
||||
)
|
||||
|
||||
assert response is not None
|
||||
assert client.call_count == 2
|
||||
assert response.messages[-1].text == "Based on the weather data, it's sunny!"
|
||||
assert execution_order == [
|
||||
"chat_middleware_before_1",
|
||||
"chat_middleware_after_1",
|
||||
"function_middleware_before",
|
||||
"function_middleware_after",
|
||||
"chat_middleware_before_2",
|
||||
"chat_middleware_after_2",
|
||||
]
|
||||
|
||||
async def test_run_level_chat_and_function_middleware_split_per_function_loop_round_streaming(self) -> None:
|
||||
"""Test mixed run-level middleware is split so chat middleware runs per model call in streaming mode."""
|
||||
execution_order: list[str] = []
|
||||
chat_round = 0
|
||||
|
||||
@chat_middleware
|
||||
async def run_level_chat_middleware(
|
||||
context: ChatContext,
|
||||
call_next: Callable[[], Awaitable[None]],
|
||||
) -> None:
|
||||
nonlocal chat_round
|
||||
chat_round += 1
|
||||
execution_order.append(f"chat_middleware_before_{chat_round}")
|
||||
await call_next()
|
||||
execution_order.append(f"chat_middleware_after_{chat_round}")
|
||||
|
||||
@function_middleware
|
||||
async def run_level_function_middleware(
|
||||
context: FunctionInvocationContext,
|
||||
call_next: Callable[[], Awaitable[None]],
|
||||
) -> None:
|
||||
execution_order.append("function_middleware_before")
|
||||
await call_next()
|
||||
execution_order.append("function_middleware_after")
|
||||
|
||||
def sample_tool(location: str) -> str:
|
||||
"""Get weather for a location."""
|
||||
return f"Weather in {location}: sunny"
|
||||
|
||||
sample_tool_wrapped = FunctionTool(
|
||||
func=sample_tool,
|
||||
name="sample_tool",
|
||||
description="Get weather for a location",
|
||||
approval_mode="never_require",
|
||||
)
|
||||
|
||||
client = MockBaseChatClient()
|
||||
client.streaming_responses = [
|
||||
[
|
||||
ChatResponseUpdate(
|
||||
contents=[
|
||||
Content.from_function_call(
|
||||
call_id="call_3",
|
||||
name="sample_tool",
|
||||
arguments='{"location": "Seattle"}',
|
||||
)
|
||||
],
|
||||
role="assistant",
|
||||
finish_reason="tool_calls",
|
||||
),
|
||||
],
|
||||
[
|
||||
ChatResponseUpdate(
|
||||
contents=[Content.from_text("Based on the weather data, it's sunny!")],
|
||||
role="assistant",
|
||||
finish_reason="stop",
|
||||
),
|
||||
],
|
||||
]
|
||||
|
||||
updates: list[ChatResponseUpdate] = []
|
||||
async for update in client.get_response(
|
||||
[Message(role="user", text="What's the weather in Seattle?")],
|
||||
options={"tools": [sample_tool_wrapped]},
|
||||
client_kwargs={"middleware": [run_level_chat_middleware, run_level_function_middleware]},
|
||||
stream=True,
|
||||
):
|
||||
updates.append(update)
|
||||
|
||||
assert client.call_count == 2
|
||||
assert len(updates) > 0
|
||||
assert execution_order == [
|
||||
"chat_middleware_before_1",
|
||||
"chat_middleware_after_1",
|
||||
"function_middleware_before",
|
||||
"function_middleware_after",
|
||||
"chat_middleware_before_2",
|
||||
"chat_middleware_after_2",
|
||||
]
|
||||
|
||||
@@ -2437,7 +2437,7 @@ def test_capture_response(span_exporter: InMemorySpanExporter):
|
||||
async def test_layer_ordering_span_sequence_with_function_calling(span_exporter: InMemorySpanExporter):
|
||||
"""Test that with correct layer ordering, spans appear in the expected sequence.
|
||||
|
||||
When using the correct layer ordering (ChatMiddlewareLayer, FunctionInvocationLayer,
|
||||
When using the correct layer ordering (FunctionInvocationLayer, ChatMiddlewareLayer,
|
||||
ChatTelemetryLayer, BaseChatClient), the spans should appear in this order:
|
||||
1. First 'chat' span (initial LLM call that returns function call)
|
||||
2. 'execute_tool' span (function invocation)
|
||||
@@ -2454,11 +2454,11 @@ async def test_layer_ordering_span_sequence_with_function_calling(span_exporter:
|
||||
def get_weather(location: str) -> str:
|
||||
return f"The weather in {location} is sunny."
|
||||
|
||||
# Correct layer ordering: FunctionInvocationLayer BEFORE ChatTelemetryLayer
|
||||
# This ensures each inner LLM call gets its own telemetry span
|
||||
# Correct layer ordering: FunctionInvocationLayer BEFORE ChatMiddlewareLayer BEFORE ChatTelemetryLayer
|
||||
# This ensures each inner LLM call traverses chat middleware and still gets its own telemetry span
|
||||
class MockChatClientWithLayers(
|
||||
ChatMiddlewareLayer,
|
||||
FunctionInvocationLayer,
|
||||
ChatMiddlewareLayer,
|
||||
ChatTelemetryLayer,
|
||||
BaseChatClient,
|
||||
):
|
||||
|
||||
Reference in New Issue
Block a user