Multiple response formats when creating agents?

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.