test_multiple_yields_in_retry_loop() — langchain Function Reference

Source Code

libs/langchain_v1/tests/unit_tests/agents/middleware/core/test_composition.py lines 312–349

    def test_multiple_yields_in_retry_loop(self) -> None:
        """Test handler that retries multiple times."""
        call_count = {"value": 0}

        def outer_counts_calls(
            request: ModelRequest, handler: Callable[[ModelRequest], ModelResponse]
        ) -> ModelResponse:
            call_count["value"] += 1
            return handler(request)

        def inner_retries(
            request: ModelRequest, handler: Callable[[ModelRequest], ModelResponse]
        ) -> ModelResponse:
            try:
                return handler(request)
            except ValueError:
                # Retry once on error
                return handler(request)

        composed = _chain_model_call_handlers([outer_counts_calls, inner_retries])
        assert composed is not None

        attempt = {"value": 0}

        def mock_base_handler(req: ModelRequest) -> ModelResponse:
            attempt["value"] += 1
            if attempt["value"] == 1:
                msg = "fail"
                raise ValueError(msg)
            return ModelResponse(result=[AIMessage(content="ok")], structured_response=None)

        result = composed(create_test_request(), mock_base_handler)

        # Outer called once, inner retried so base handler called twice
        assert call_count["value"] == 1
        assert attempt["value"] == 2
        assert isinstance(result, _ComposedExtendedModelResponse)
        assert result.model_response.result[0].content == "ok"