test_decorator_complex_retry_logic() — langchain Function Reference

Source Code

libs/langchain_v1/tests/unit_tests/agents/middleware/core/test_wrap_model_call.py lines 1303–1346

    def test_decorator_complex_retry_logic(self) -> None:
        """Test decorator with complex retry logic and backoff."""
        attempts = []
        call_count = {"value": 0}

        class UnreliableModel(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"] <= 2:
                    msg = f"Attempt {call_count['value']} failed"
                    raise ValueError(msg)
                return super()._generate(messages, **kwargs)

        @wrap_model_call
        def retry_with_tracking(
            request: ModelRequest,
            handler: Callable[[ModelRequest], ModelResponse],
        ) -> ModelCallResult:
            max_retries = 3
            for attempt in range(max_retries):
                attempts.append(attempt + 1)
                try:
                    return handler(request)
                except Exception:
                    # On error, continue to next attempt
                    if attempt < max_retries - 1:
                        continue  # Retry
                    raise  # All retries failed
            pytest.fail("Should have raised an exception")

        model = UnreliableModel(messages=iter([AIMessage(content="Finally worked")]))
        agent = create_agent(model=model, middleware=[retry_with_tracking])

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

        assert attempts == [1, 2, 3]
        assert result["messages"][1].content == "Finally worked"