test_middle_retry_middleware() — langchain Function Reference

Source Code

libs/langchain_v1/tests/unit_tests/agents/middleware/core/test_wrap_model_call.py lines 926–1002

    def test_middle_retry_middleware(self) -> None:
        """Test that middle middleware doing retry causes inner to execute twice."""
        execution_order = []
        model_calls = []

        class OuterMiddleware(AgentMiddleware):
            def wrap_model_call(
                self,
                request: ModelRequest,
                handler: Callable[[ModelRequest], ModelResponse],
            ) -> ModelCallResult:
                execution_order.append("outer-before")
                result = handler(request)
                execution_order.append("outer-after")
                return result

        class MiddleRetryMiddleware(AgentMiddleware):
            def wrap_model_call(
                self,
                request: ModelRequest,
                handler: Callable[[ModelRequest], ModelResponse],
            ) -> ModelCallResult:
                execution_order.append("middle-before")
                # Always retry once (call handler twice)
                result = handler(request)
                execution_order.append("middle-retry")
                result = handler(request)
                execution_order.append("middle-after")
                return result

        class InnerMiddleware(AgentMiddleware):
            def wrap_model_call(
                self,
                request: ModelRequest,
                handler: Callable[[ModelRequest], ModelResponse],
            ) -> ModelCallResult:
                execution_order.append("inner-before")
                result = handler(request)
                execution_order.append("inner-after")
                return result

        class TrackingModel(GenericFakeChatModel):
            @override
            def _generate(
                self,
                messages: list[BaseMessage],
                stop: list[str] | None = None,
                run_manager: CallbackManagerForLLMRun | None = None,
                **kwargs: Any,
            ) -> ChatResult:
                model_calls.append(len(messages))
                return super()._generate(messages, **kwargs)

        model = TrackingModel(
            messages=iter([AIMessage(content="Response 1"), AIMessage(content="Response 2")])
        )
        agent = create_agent(
            model=model,
            middleware=[OuterMiddleware(), MiddleRetryMiddleware(), InnerMiddleware()],
        )

        agent.invoke({"messages": [HumanMessage("Test")]})

        # Middle yields twice, so inner runs twice
        assert execution_order == [
            "outer-before",
            "middle-before",
            "inner-before",  # First execution
            "inner-after",
            "middle-retry",  # Middle yields again
            "inner-before",  # Second execution
            "inner-after",
            "middle-after",
            "outer-after",
        ]
        # Model should be called twice
        assert len(model_calls) == 2