test_structured_output_retry.py — langchain Source File

Source Code

"""Tests for StructuredOutputRetryMiddleware functionality."""

from collections.abc import Callable

import pytest
from langchain_core.messages import HumanMessage
from langchain_core.tools import tool
from langgraph.checkpoint.memory import InMemorySaver
from pydantic import BaseModel

from langchain.agents import create_agent
from langchain.agents.middleware.types import (
    AgentMiddleware,
    ModelRequest,
    ModelResponse,
)
from langchain.agents.structured_output import StructuredOutputError, ToolStrategy
from tests.unit_tests.agents.model import FakeToolCallingModel


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

    def __init__(self, max_retries: int) -> None:
        """Initialize the structured output retry middleware.

        Args:
            max_retries: Maximum number of retry attempts.
        """
        self.max_retries = max_retries

    def wrap_model_call(
        self, request: ModelRequest, handler: Callable[[ModelRequest], ModelResponse]
    ) -> ModelResponse:
        """Intercept and control model execution via handler callback.

        Args:
            request: The model request containing messages and configuration.
            handler: The function to call the model.

        Returns:
            The model response.

        Raises:
            StructuredOutputError: If max retries exceeded without success.
        """
        for attempt in range(self.max_retries + 1):
            try:
                return handler(request)
            except StructuredOutputError as exc:
                if attempt == self.max_retries:
                    raise

                # Include both the AI message and error in a single human message
                # to maintain valid chat history alternation
                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."
                )
// ... (310 more lines)