test_max_attempts_limit() — langchain Function Reference

Source Code

libs/langchain_v1/tests/unit_tests/agents/middleware/core/test_wrap_model_call.py lines 235–293

    def test_max_attempts_limit(self) -> None:
        """Test that middleware controls termination via retry limits."""

        class AlwaysFailingModel(BaseChatModel):
            """Model that always fails."""

            @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)

            @property
            def _llm_type(self) -> str:
                return "always_failing"

        class LimitedRetryMiddleware(AgentMiddleware):
            """Middleware that limits its own retries."""

            def __init__(self, max_retries: int = 10):
                super().__init__()
                self.max_retries = max_retries
                self.attempt_count = 0

            def wrap_model_call(
                self,
                request: ModelRequest,
                handler: Callable[[ModelRequest], ModelResponse],
            ) -> ModelCallResult:
                last_exception = None
                for _attempt in range(self.max_retries):
                    self.attempt_count += 1
                    try:
                        return handler(request)
                    except Exception as e:
                        last_exception = e
                        # Continue to retry

                # All retries exhausted, re-raise the last error
                if last_exception:
                    raise last_exception
                pytest.fail("Should have raised an exception")

        model = AlwaysFailingModel()
        middleware = LimitedRetryMiddleware(max_retries=10)

        agent = create_agent(model=model, middleware=[middleware])

        # Should fail with the model's error after middleware stops retrying
        with pytest.raises(ValueError, match="Always fails"):
            agent.invoke({"messages": [HumanMessage("Test")]})

        # Should have attempted exactly 10 times as configured
        assert middleware.attempt_count == 10