in the meantime
instead of
@wrap_model_call
def select_format(
request: ModelRequest,
handler: Callable[[ModelRequest], ModelResponse],
) -> ModelResponse:
key = request.runtime.context.response_format_key
request.response_format = ToolStrategy(SCHEMAS[key]) # subset of the Union
return handler(request)
try this
class EnergyMiddleware(AgentMiddleware):
def wrap_model_call(self, request, handler):
key = request.runtime.context.response_format_key
return handler(
request.override(response_format=ToolStrategy(RESPONSE_FORMATS[key]))
)
Two changes: .override(…) returns a new ModelRequest, and you must forward that new request to handler. If you mutate in place and forward the original, the factory may still see the Union because of dataclass aliasing plus the deprecation path.
Also worth double-checking: that every invoke() passes context=ResponseFormatContext(response_format_key=“…”), and that the key actually exists in RESPONSE_FORMATS. Full known-good snippet is in the doc.