feat: support structured output retry middleware (#33663)

* attach the latest `AIMessage` to all `StructuredOutputError`s so that
relevant middleware can use as desired
* raise `StructuredOutputError` from `ProviderStrategy` logic in case of
failed parsing (so that we can retry from middleware)
* added a test suite w/ example custom middleware that retries for tool
+ provider strategy

Long term, we could add our own opinionated structured output retry
middleware, but this at least unblocks folks who want to use custom
retry logic in the short term :)

```py
class StructuredOutputRetryMiddleware(AgentMiddleware):
    """Retries model calls when structured output parsing fails."""

    def __init__(self, max_retries: int) -> None:
        self.max_retries = max_retries

    def wrap_model_call(
        self, request: ModelRequest, handler: Callable[[ModelRequest], ModelResponse]
    ) -> ModelResponse:
        for attempt in range(self.max_retries + 1):
            try:
                return handler(request)
            except StructuredOutputError as exc:
                if attempt == self.max_retries:
                    raise

                ai_content = exc.ai_message.content
                error_message = (
                    f"Your previous response was:\n{ai_content}\n\n"
                    f"Error: {exc}. Please try again with a valid response."
                )
                request.messages.append(HumanMessage(content=error_message))
```

Author: sydney-runkle

libs/langchain_v1/langchain/agents/factory.py

@@ -33,6 +33,7 @@
     ProviderStrategy,
     ProviderStrategyBinding,
     ResponseFormat,
+    StructuredOutputError,
     StructuredOutputValidationError,
     ToolStrategy,
 )
@@ -797,8 +798,16 @@ def _handle_model_output(
                 provider_strategy_binding = ProviderStrategyBinding.from_schema_spec(
                     effective_response_format.schema_spec
                 )
-                structured_response = provider_strategy_binding.parse(output)
-                return {"messages": [output], "structured_response": structured_response}
+                try:
+                    structured_response = provider_strategy_binding.parse(output)
+                except Exception as exc:  # noqa: BLE001
+                    schema_name = getattr(
+                        effective_response_format.schema_spec.schema, "__name__", "response_format"
+                    )
+                    validation_error = StructuredOutputValidationError(schema_name, exc, output)
+                    raise validation_error
+                else:
+                    return {"messages": [output], "structured_response": structured_response}
             return {"messages": [output]}
 
         # Handle structured output with tool strategy
@@ -812,11 +821,11 @@ def _handle_model_output(
             ]
 
             if structured_tool_calls:
-                exception: Exception | None = None
+                exception: StructuredOutputError | None = None
                 if len(structured_tool_calls) > 1:
                     # Handle multiple structured outputs error
                     tool_names = [tc["name"] for tc in structured_tool_calls]
-                    exception = MultipleStructuredOutputsError(tool_names)
+                    exception = MultipleStructuredOutputsError(tool_names, output)
                     should_retry, error_message = _handle_structured_output_error(
                         exception, effective_response_format
                     )
@@ -858,7 +867,7 @@ def _handle_model_output(
                         "structured_response": structured_response,
                     }
                 except Exception as exc:  # noqa: BLE001
-                    exception = StructuredOutputValidationError(tool_call["name"], exc)
+                    exception = StructuredOutputValidationError(tool_call["name"], exc, output)
                     should_retry, error_message = _handle_structured_output_error(
                         exception, effective_response_format
                     )

libs/langchain_v1/langchain/agents/structured_output.py

@@ -34,17 +34,21 @@
 class StructuredOutputError(Exception):
     """Base class for structured output errors."""
 
+    ai_message: AIMessage
+
 
 class MultipleStructuredOutputsError(StructuredOutputError):
     """Raised when model returns multiple structured output tool calls when only one is expected."""
 
-    def __init__(self, tool_names: list[str]) -> None:
+    def __init__(self, tool_names: list[str], ai_message: AIMessage) -> None:
         """Initialize `MultipleStructuredOutputsError`.
 
         Args:
             tool_names: The names of the tools called for structured output.
+            ai_message: The AI message that contained the invalid multiple tool calls.
         """
         self.tool_names = tool_names
+        self.ai_message = ai_message
 
         super().__init__(
             "Model incorrectly returned multiple structured responses "
@@ -55,15 +59,17 @@ def __init__(self, tool_names: list[str]) -> None:
 class StructuredOutputValidationError(StructuredOutputError):
     """Raised when structured output tool call arguments fail to parse according to the schema."""
 
-    def __init__(self, tool_name: str, source: Exception) -> None:
+    def __init__(self, tool_name: str, source: Exception, ai_message: AIMessage) -> None:
         """Initialize `StructuredOutputValidationError`.
 
         Args:
             tool_name: The name of the tool that failed.
             source: The exception that occurred.
+            ai_message: The AI message that contained the invalid structured output.
         """
         self.tool_name = tool_name
         self.source = source
+        self.ai_message = ai_message
         super().__init__(f"Failed to parse structured output for tool '{tool_name}': {source}.")