test_max_retries() — langchain Function Reference

Source Code

libs/langchain_v1/tests/unit_tests/agents/middleware/core/test_wrap_model_call.py lines 153–199

    def test_max_retries(self) -> None:
        """Test middleware with maximum retry limit."""

        class AlwaysFailModel(GenericFakeChatModel):
            @override
            def _generate(
                self,
                messages: list[BaseMessage],
                stop: list[str] | None = None,
                run_manager: CallbackManagerForLLMRun | None = None,
                **kwargs: Any,
            ) -> ChatResult:
                msg = "Always fails"
                raise ValueError(msg)

        class MaxRetriesMiddleware(AgentMiddleware):
            def __init__(self, max_retries: int = 3):
                super().__init__()
                self.max_retries = max_retries
                self.attempts = []

            def wrap_model_call(
                self,
                request: ModelRequest,
                handler: Callable[[ModelRequest], ModelResponse],
            ) -> ModelCallResult:
                last_exception = None
                for attempt in range(self.max_retries):
                    self.attempts.append(attempt + 1)
                    try:
                        return handler(request)
                    except Exception as e:
                        last_exception = e
                        continue
                # Re-raise the last exception
                if last_exception:
                    raise last_exception
                pytest.fail("Should have raised an exception")

        retry_middleware = MaxRetriesMiddleware(max_retries=3)
        model = AlwaysFailModel(messages=iter([]))
        agent = create_agent(model=model, middleware=[retry_middleware])

        with pytest.raises(ValueError, match="Always fails"):
            agent.invoke({"messages": [HumanMessage("Test")]})

        assert retry_middleware.attempts == [1, 2, 3]