test_retry_outer_transform_inner() — langchain Function Reference

Source Code

libs/langchain_v1/tests/unit_tests/agents/middleware/core/test_wrap_model_call.py lines 875–924

    def test_retry_outer_transform_inner(self) -> None:
        """Test retry as outer middleware with transform as inner."""
        call_count = {"value": 0}

        class FailOnceThenSucceed(GenericFakeChatModel):
            @override
            def _generate(
                self,
                messages: list[BaseMessage],
                stop: list[str] | None = None,
                run_manager: CallbackManagerForLLMRun | None = None,
                **kwargs: Any,
            ) -> ChatResult:
                call_count["value"] += 1
                if call_count["value"] == 1:
                    msg = "First call fails"
                    raise ValueError(msg)
                return super()._generate(messages, **kwargs)

        class RetryMiddleware(AgentMiddleware):
            def wrap_model_call(
                self,
                request: ModelRequest,
                handler: Callable[[ModelRequest], ModelResponse],
            ) -> ModelCallResult:
                try:
                    return handler(request)
                except Exception:
                    return handler(request)

        class UppercaseMiddleware(AgentMiddleware):
            def wrap_model_call(
                self,
                request: ModelRequest,
                handler: Callable[[ModelRequest], ModelResponse],
            ) -> ModelCallResult:
                result = handler(request)
                # result is ModelResponse, extract AIMessage from it
                ai_message = result.result[0]
                assert isinstance(ai_message.content, str)
                return AIMessage(content=ai_message.content.upper())

        model = FailOnceThenSucceed(messages=iter([AIMessage(content="success")]))
        # Retry outer, Uppercase inner
        agent = create_agent(model=model, middleware=[RetryMiddleware(), UppercaseMiddleware()])

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

        # Should retry and uppercase the result
        assert result["messages"][1].content == "SUCCESS"