From 3072cbff1c6a85ee357c7a7a4b3f264c256f0df8 Mon Sep 17 00:00:00 2001 From: Adriana Knight Date: Fri, 23 Oct 2026 16:09:13 -0400 Subject: [PATCH] Add max_retires passthrough arg to built-in handler factories --- .../launchdarkly_ai_claude_messages.txt | 2 +- .../handler.py | 32 +++-- .../claude-messages/tests/test_handler.py | 62 ++++++++++ .../launchdarkly_ai_langchain_messages.txt | 2 +- .../handler.py | 66 ++++++++-- .../langchain-messages/tests/test_handler.py | 113 ++++++++++++++++++ .../launchdarkly_ai_litellm_messages.txt | 2 +- .../handler.py | 31 ++++- .../litellm-messages/tests/test_handler.py | 54 +++++++++ .../launchdarkly_ai_openai_messages.txt | 2 +- .../handler.py | 32 +++-- .../openai-messages/tests/test_handler.py | 62 ++++++++++ 12 files changed, 425 insertions(+), 35 deletions(-) diff --git a/packages/claude-messages/api-report/launchdarkly_ai_claude_messages.txt b/packages/claude-messages/api-report/launchdarkly_ai_claude_messages.txt index b051ead8..8a4b4624 100644 --- a/packages/claude-messages/api-report/launchdarkly_ai_claude_messages.txt +++ b/packages/claude-messages/api-report/launchdarkly_ai_claude_messages.txt @@ -1,3 +1,3 @@ # Public API of launchdarkly_ai_claude_messages. Generated by `make api-report`; do not edit. claude_messages(config_key: str, user_input: str, context: LDContext, **kwargs: Any) -> Any -create_claude_messages_handler(*, capture_content: bool | None = None) -> ProviderHandler +create_claude_messages_handler(*, capture_content: bool | None = None, max_retries: int | None = None) -> ProviderHandler diff --git a/packages/claude-messages/src/launchdarkly_ai_claude_messages/handler.py b/packages/claude-messages/src/launchdarkly_ai_claude_messages/handler.py index ebd17db4..8b009481 100644 --- a/packages/claude-messages/src/launchdarkly_ai_claude_messages/handler.py +++ b/packages/claude-messages/src/launchdarkly_ai_claude_messages/handler.py @@ -529,7 +529,7 @@ def _is_coroutine(fn: Any) -> bool: def create_claude_messages_handler( - *, capture_content: bool | None = None + *, capture_content: bool | None = None, max_retries: int | None = None ) -> ProviderHandler: """Creates a ``ProviderHandler`` for Anthropic Claude (messages API). @@ -541,22 +541,35 @@ def create_claude_messages_handler( for this handler. Conversation content is PII, so a run emits only metadata, meaning models, token counts, timings and tool names, until capture is turned on. Turning this on sends the text of every request and response to whatever collector the SDK is pointed at. + + Set *max_retries* to change how many times the ``AsyncAnthropic`` client retries a failed + request (a rate limit, an overload, a timeout, a 5xx) before the error reaches the handler. + ``None`` (the default) leaves the client's own default of 2 in place. Raise it when many calls + share one rate-limit window, as an evaluation run with many concurrent workers does: two + retries is thin when a whole batch hits the same window at once. The retries happen inside the + client, so a retried turn is still one ``chat`` span. """ report_usage( "claude-messages.createClaudeMessagesHandler", PACKAGE_NAME, __version__ ) - return _create_claude_messages_handler(capture_content=capture_content) + return _create_claude_messages_handler( + capture_content=capture_content, max_retries=max_retries + ) def _create_claude_messages_handler( - *, capture_content: bool | None = None + *, capture_content: bool | None = None, max_retries: int | None = None ) -> ProviderHandler: """Non-reporting :func:`create_claude_messages_handler`, used by this package's wrappers.""" content_override = capture_content import importlib anthropic_mod = importlib.import_module("anthropic") - client = anthropic_mod.AsyncAnthropic() + # Passed only when set, so an omitted value keeps the vendor default rather than a copy of it. + client_options: dict[str, Any] = {} + if max_retries is not None: + client_options["max_retries"] = max_retries + client = anthropic_mod.AsyncAnthropic(**client_options) async def _call_impl( config: AiConfigRep, @@ -950,14 +963,17 @@ def claude_messages( **kwargs: Any, ) -> Any: """Convenience wrapper: creates a handler and calls config(...).invoke().""" - # Both are lifted out of kwargs: capture_content configures the handler, variables belong to - # the invocation. Leaving either in would pass it to config(), which takes neither, so a caller - # asking for content on spans got a TypeError instead of content. + # All three are lifted out of kwargs: capture_content and max_retries configure the handler, + # variables belong to the invocation. Leaving any in would pass it to config(), which takes + # none of them, so a caller asking for content on spans got a TypeError instead of content. report_usage("claude-messages.claudeMessages", PACKAGE_NAME, __version__) variables = kwargs.pop("variables", None) capture_content = kwargs.pop("capture_content", None) + max_retries = kwargs.pop("max_retries", None) return _config( key=config_key, - handler=_create_claude_messages_handler(capture_content=capture_content), + handler=_create_claude_messages_handler( + capture_content=capture_content, max_retries=max_retries + ), **kwargs, ).invoke(user_input, context, variables=variables) diff --git a/packages/claude-messages/tests/test_handler.py b/packages/claude-messages/tests/test_handler.py index 3732a860..f7e8a4f0 100644 --- a/packages/claude-messages/tests/test_handler.py +++ b/packages/claude-messages/tests/test_handler.py @@ -3127,3 +3127,65 @@ async def test_effort_still_lands_in_output_config_when_output_config_is_malform await create_claude_messages_handler()(config, "q", {}, {}) kwargs = mock_anthropic.messages.create.call_args.kwargs assert kwargs["output_config"] == {"effort": "low"} + + +class TestMaxRetries: + """`max_retries` reaches the `AsyncAnthropic` client the factory builds, and only when set. + + The client's own retries are the right layer for rate limits, so an omitted value must leave the + vendor default in place rather than replace it, and nothing may ever turn retries off. + """ + + def _build( + self, monkeypatch: pytest.MonkeyPatch, **kwargs: Any + ) -> tuple[Any, MagicMock]: + import anthropic + + from launchdarkly_ai_claude_messages import create_claude_messages_handler + + monkeypatch.setenv("ANTHROPIC_API_KEY", "test-key") + real_client = anthropic.AsyncAnthropic + built: list[Any] = [] + + def _construct(**options: Any) -> Any: + built.append(real_client(**options)) + return built[-1] + + with patch("anthropic.AsyncAnthropic", side_effect=_construct) as ctor: + create_claude_messages_handler(**kwargs) + return built[0], ctor + + def test_value_reaches_the_constructed_client( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + client, _ctor = self._build(monkeypatch, max_retries=7) + assert client.max_retries == 7 + + def test_omitted_value_keeps_the_vendor_default( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + import anthropic + + client, ctor = self._build(monkeypatch) + assert "max_retries" not in ctor.call_args.kwargs + assert client.max_retries == anthropic.DEFAULT_MAX_RETRIES + + def test_explicit_zero_is_the_callers_choice( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + client, _ctor = self._build(monkeypatch, max_retries=0) + assert client.max_retries == 0 + + def test_convenience_wrapper_forwards_it_to_the_factory(self) -> None: + import launchdarkly_ai_claude_messages.handler as handler_mod + + factory = MagicMock() + fake_config = MagicMock() + with ( + patch.object(handler_mod, "_create_claude_messages_handler", factory), + patch.object(handler_mod, "_config", fake_config), + ): + handler_mod.claude_messages("k", "q", {}, max_retries=6) + assert factory.call_args.kwargs["max_retries"] == 6 + # config() takes no such argument, so leaving it in kwargs would raise TypeError. + assert "max_retries" not in fake_config.call_args.kwargs diff --git a/packages/langchain-messages/api-report/launchdarkly_ai_langchain_messages.txt b/packages/langchain-messages/api-report/launchdarkly_ai_langchain_messages.txt index c9602d48..7bf2aaa8 100644 --- a/packages/langchain-messages/api-report/launchdarkly_ai_langchain_messages.txt +++ b/packages/langchain-messages/api-report/launchdarkly_ai_langchain_messages.txt @@ -1,3 +1,3 @@ # Public API of launchdarkly_ai_langchain_messages. Generated by `make api-report`; do not edit. -create_langchain_messages_handler(llm: Any = None, *, capture_content: bool | None = None) -> ProviderHandler +create_langchain_messages_handler(llm: Any = None, *, capture_content: bool | None = None, max_retries: int | None = None) -> ProviderHandler langchain_messages(config_key: str, user_input: str, context: LDContext, llm: Any = None, **kwargs: Any) -> Any diff --git a/packages/langchain-messages/src/launchdarkly_ai_langchain_messages/handler.py b/packages/langchain-messages/src/launchdarkly_ai_langchain_messages/handler.py index 78f36622..0da211b4 100644 --- a/packages/langchain-messages/src/launchdarkly_ai_langchain_messages/handler.py +++ b/packages/langchain-messages/src/launchdarkly_ai_langchain_messages/handler.py @@ -591,6 +591,7 @@ def _model_constructor_kwargs( *, mapping_keys: frozenset[str] = frozenset(), renames: Mapping[str, str] | None = None, + max_retries: int | None = None, ) -> dict[str, Any]: params = model_parameters(config) for alias, name in (renames or {}).items(): @@ -601,6 +602,10 @@ def _model_constructor_kwargs( params, forwarded_keys, mapping_keys=mapping_keys ) parameters["model"] = _resolved_model_name(config, fallback_name) + # From the handler's options, never from model.parameters, which excludes it above. Set only + # when given, so an omitted value keeps each chat model's own default. + if max_retries is not None: + parameters["max_retries"] = max_retries return parameters @@ -609,13 +614,16 @@ def _is_model_factory(llm: Any) -> bool: return callable(llm) and not hasattr(llm, "invoke") and not hasattr(llm, "ainvoke") -def _make_default_chat_model(config: AiConfigRep, importlib: Any) -> Any: +def _make_default_chat_model( + config: AiConfigRep, importlib: Any, max_retries: int | None = None +) -> Any: """ Instantiate the appropriate LangChain chat model based on ``config.provider.name``. Falls back to ``ChatOpenAI`` when the provider is not recognised. Requires the matching ``langchain-`` integration package to be installed. ``model.parameters`` are filtered to that class's forwarded list first. ``tools`` is never on - that list, since tools are bound separately from ``config["tools"]``. + that list, since tools are bound separately from ``config["tools"]``. *max_retries*, when set, + is passed to whichever class is built. """ provider = config.get("provider", {}).get("name", "openai").lower() if provider == "anthropic": @@ -629,6 +637,7 @@ def _make_default_chat_model(config: AiConfigRep, importlib: Any) -> Any: _CHAT_ANTHROPIC_FORWARDED_KEYS, mapping_keys=_CHAT_ANTHROPIC_MAPPING_KEYS, renames=_CHAT_ANTHROPIC_RENAMES, + max_retries=max_retries, ) ) ) @@ -648,6 +657,7 @@ def _make_default_chat_model(config: AiConfigRep, importlib: Any) -> Any: _CHAT_BEDROCK_CONVERSE_FORWARDED_KEYS, mapping_keys=_CHAT_BEDROCK_CONVERSE_MAPPING_KEYS, renames=_CHAT_BEDROCK_CONVERSE_RENAMES, + max_retries=max_retries, ) ) lc_openai = importlib.import_module("langchain_openai") @@ -657,14 +667,17 @@ def _make_default_chat_model(config: AiConfigRep, importlib: Any) -> Any: "gpt-4o", _CHAT_OPENAI_FORWARDED_KEYS, mapping_keys=_CHAT_OPENAI_MAPPING_KEYS, + max_retries=max_retries, ) ) -async def _resolve_base_model(config: AiConfigRep, llm: Any, importlib: Any) -> Any: +async def _resolve_base_model( + config: AiConfigRep, llm: Any, importlib: Any, max_retries: int | None = None +) -> Any: invocation = _config_for_model_call(config) if llm is None: - return _make_default_chat_model(invocation, importlib) + return _make_default_chat_model(invocation, importlib, max_retries) if _is_model_factory(llm): model = llm(invocation) if asyncio.iscoroutine(model): @@ -762,7 +775,10 @@ async def _run_structured_turn( def create_langchain_messages_handler( - llm: Any = None, *, capture_content: bool | None = None + llm: Any = None, + *, + capture_content: bool | None = None, + max_retries: int | None = None, ) -> ProviderHandler: """ Creates a ``ProviderHandler`` for LangChain (chat models). @@ -777,15 +793,36 @@ def create_langchain_messages_handler( which is off until ``configure`` sets it. ``True`` or ``False`` overrides that default for this handler. Conversation content is PII, so a run emits only metadata, meaning models, token counts, timings and tool names, until capture is turned on. + + Set *max_retries* to change how many times the default chat model retries a failed request (a + rate limit, a timeout, a 5xx) before the error reaches the handler. It is passed to the + ``ChatOpenAI``, ``ChatAnthropic`` or ``ChatBedrockConverse`` the handler builds. ``None`` (the + default) leaves that class's own default in place, which is 2 for OpenAI and Anthropic. Raise + it when many calls share one rate-limit window, as an evaluation run with many concurrent + workers does: two retries is thin when a whole batch hits the same window at once. It cannot + be combined with *llm*, because the handler does not build that model: set ``max_retries`` on + the model instead. """ report_usage("langchain-messages.createLangChainHandler", PACKAGE_NAME, __version__) - return _create_langchain_messages_handler(llm, capture_content=capture_content) + return _create_langchain_messages_handler( + llm, capture_content=capture_content, max_retries=max_retries + ) def _create_langchain_messages_handler( - llm: Any = None, *, capture_content: bool | None = None + llm: Any = None, + *, + capture_content: bool | None = None, + max_retries: int | None = None, ) -> ProviderHandler: """Non-reporting :func:`create_langchain_messages_handler`, used by this package's wrappers.""" + # Refused rather than ignored: the handler never constructs a model it is given, so the value + # would go nowhere and the caller would believe their retries had been raised. + if llm is not None and max_retries is not None: + raise ValueError( + "max_retries applies only to the default chat model. Set max_retries on the model " + "passed as llm instead." + ) content_override = capture_content async def _call_impl( @@ -826,7 +863,7 @@ async def _call_impl( system_instructions=system_instructions, messages=span_messages, ) - base_model = await _resolve_base_model(config, llm, importlib) + base_model = await _resolve_base_model(config, llm, importlib, max_retries) tool_defs = _build_tools(config.get("tools") or {}) output_format = config.get("outputFormat") @@ -1097,6 +1134,7 @@ def _stream_impl( variables or {}, history, capture_content=effective_capture_content(content_override), + max_retries=max_retries, ) return _create_handler( @@ -1116,6 +1154,7 @@ async def _stream_gen( history: list[dict[str, Any]] | None = None, *, capture_content: bool = False, + max_retries: int | None = None, ) -> AsyncGenerator[dict[str, Any], None]: """Streams the run, emitting the same span tree as the blocking path. @@ -1138,7 +1177,7 @@ async def _stream_gen( """ import importlib - base_model = await _resolve_base_model(config, llm, importlib) + base_model = await _resolve_base_model(config, llm, importlib, max_retries) span = start_root_span(config, variables) parent = parent_context_of(span) @@ -1501,16 +1540,17 @@ def langchain_messages( **kwargs: Any, ) -> Any: """Convenience wrapper: creates a handler and calls config(...).invoke().""" - # Both are lifted out of kwargs: capture_content configures the handler, variables belong to - # the invocation. Leaving either in would pass it to config(), which takes neither, so a caller - # asking for content on spans got a TypeError instead of content. + # All three are lifted out of kwargs: capture_content and max_retries configure the handler, + # variables belong to the invocation. Leaving any in would pass it to config(), which takes + # none of them, so a caller asking for content on spans got a TypeError instead of content. report_usage("langchain-messages.langchainMessages", PACKAGE_NAME, __version__) variables = kwargs.pop("variables", None) capture_content = kwargs.pop("capture_content", None) + max_retries = kwargs.pop("max_retries", None) return _config( key=config_key, handler=_create_langchain_messages_handler( - llm, capture_content=capture_content + llm, capture_content=capture_content, max_retries=max_retries ), **kwargs, ).invoke(user_input, context, variables=variables) diff --git a/packages/langchain-messages/tests/test_handler.py b/packages/langchain-messages/tests/test_handler.py index d18aed5f..c79dc0d4 100644 --- a/packages/langchain-messages/tests/test_handler.py +++ b/packages/langchain-messages/tests/test_handler.py @@ -3014,3 +3014,116 @@ def test_api_key_and_base_url_are_never_forwarded( assert "api_key" not in kwargs assert "base_url" not in kwargs assert kwargs["temperature"] == 0.3 + + +class TestMaxRetries: + """`max_retries` reaches the chat model the handler builds, and only when set. + + The provider client's own retries are the right layer for rate limits, so an omitted value must + leave each class's default in place rather than replace it. + """ + + @pytest.fixture(autouse=True) + def _api_keys(self, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("OPENAI_API_KEY", "test-key") + monkeypatch.setenv("ANTHROPIC_API_KEY", "test-key") + + def test_value_reaches_the_openai_client(self) -> None: + import importlib + + import launchdarkly_ai_langchain_messages.handler as handler_mod + + model = handler_mod._make_default_chat_model(CONFIG, importlib, max_retries=7) + assert isinstance(model, _REAL_CHAT_OPENAI) + assert model.root_async_client.max_retries == 7 + + def test_value_reaches_the_anthropic_client(self) -> None: + import importlib + + import launchdarkly_ai_langchain_messages.handler as handler_mod + + cfg = { + **CONFIG, + "provider": {"name": "Anthropic"}, + "model": {"name": "claude-sonnet-4-5"}, + } + model = handler_mod._make_default_chat_model(cfg, importlib, max_retries=7) + assert isinstance(model, _REAL_CHAT_ANTHROPIC) + assert model._async_client.max_retries == 7 + + def test_value_reaches_the_bedrock_constructor(self) -> None: + import launchdarkly_ai_langchain_messages.handler as handler_mod + + ctor = MagicMock() + ctor.model_fields = _FAKE_BEDROCK_FIELDS + imports = MagicMock() + imports.import_module.return_value = MagicMock(ChatBedrockConverse=ctor) + cfg = { + **CONFIG, + "provider": {"name": "Bedrock"}, + "model": {"name": "anthropic.claude-sonnet-4-5"}, + } + handler_mod._make_default_chat_model(cfg, imports, max_retries=7) + assert ctor.call_args.kwargs["max_retries"] == 7 + + def test_omitted_value_keeps_the_vendor_default(self) -> None: + import importlib + + import openai + + import launchdarkly_ai_langchain_messages.handler as handler_mod + + model = handler_mod._make_default_chat_model(CONFIG, importlib) + assert model.root_async_client.max_retries == openai.DEFAULT_MAX_RETRIES + + def test_config_cannot_set_it(self) -> None: + import importlib + + import launchdarkly_ai_langchain_messages.handler as handler_mod + + cfg = { + **CONFIG, + "model": {"name": "gpt-4o", "parameters": {"max_retries": 50}}, + } + model = handler_mod._make_default_chat_model(cfg, importlib, max_retries=3) + assert model.root_async_client.max_retries == 3 + + @pytest.mark.asyncio + async def test_handler_passes_it_on_invoke_and_stream(self) -> None: + ctx, _rec = _recording() + from launchdarkly_ai_langchain_messages import create_langchain_messages_handler + + llm = _make_llm("default-openai") + ctor = MagicMock(return_value=llm) + ctor.model_fields = _REAL_CHAT_OPENAI.model_fields + handler = create_langchain_messages_handler(max_retries=7) + with ( + ctx, + patch.dict(sys.modules, {"langchain_openai": MagicMock(ChatOpenAI=ctor)}), + ): + await handler(CONFIG, "q", {}, {}) + assert ctor.call_args.kwargs["max_retries"] == 7 + ctor.reset_mock() + async for _event in await handler.stream(CONFIG, "q", {}, {}): + pass + assert ctor.call_args.kwargs["max_retries"] == 7 + + def test_refused_alongside_a_caller_supplied_model(self) -> None: + from launchdarkly_ai_langchain_messages import create_langchain_messages_handler + + with pytest.raises(ValueError, match="max_retries"): + create_langchain_messages_handler(_make_llm("mine"), max_retries=7) + + def test_convenience_wrapper_forwards_it_to_the_factory(self) -> None: + import launchdarkly_ai_langchain_messages.handler as handler_mod + + factory = MagicMock() + fake_config = MagicMock() + with ( + patch.object(handler_mod, "_create_langchain_messages_handler", factory), + patch.object(handler_mod, "_config", fake_config), + ): + handler_mod.langchain_messages("k", "q", {}, max_retries=6) + assert factory.call_args.kwargs["max_retries"] == 6 + # config() takes no such argument, so leaving it in kwargs would raise TypeError. + assert "max_retries" not in fake_config.call_args.kwargs diff --git a/packages/litellm-messages/api-report/launchdarkly_ai_litellm_messages.txt b/packages/litellm-messages/api-report/launchdarkly_ai_litellm_messages.txt index c61f206d..dfedab34 100644 --- a/packages/litellm-messages/api-report/launchdarkly_ai_litellm_messages.txt +++ b/packages/litellm-messages/api-report/launchdarkly_ai_litellm_messages.txt @@ -1,3 +1,3 @@ # Public API of launchdarkly_ai_litellm_messages. Generated by `make api-report`; do not edit. -create_litellm_messages_handler(*, completion: Completion | None = None, capture_content: bool | None = None, base_url: str | None = None, api_key: str | None = None) -> ProviderHandler +create_litellm_messages_handler(*, completion: Completion | None = None, capture_content: bool | None = None, base_url: str | None = None, api_key: str | None = None, max_retries: int | None = None) -> ProviderHandler litellm_messages(config_key: str, user_input: str, context: LDContext, **kwargs: Any) -> Any diff --git a/packages/litellm-messages/src/launchdarkly_ai_litellm_messages/handler.py b/packages/litellm-messages/src/launchdarkly_ai_litellm_messages/handler.py index 6ded1a1b..88040d4b 100644 --- a/packages/litellm-messages/src/launchdarkly_ai_litellm_messages/handler.py +++ b/packages/litellm-messages/src/launchdarkly_ai_litellm_messages/handler.py @@ -49,7 +49,8 @@ def _json_schema(output_format: dict[str, Any]) -> dict[str, Any]: #: The cross-SDK list for LiteLLM Messages (TESTING.md section 1.12): the ``model.parameters`` #: keys forwarded to ``litellm.acompletion``. Every other key is dropped. The handler sets #: ``model``, ``messages``, ``stream``, ``stream_options``, ``tools`` and ``response_format`` itself, -#: and ``api_key`` (with ``api_base``) comes from the handler's options, never from the config. +#: and ``api_key`` (with ``api_base``) and ``max_retries`` come from the handler's options, never +#: from the config. #: ``TestAcompletionAcceptsExactlyTheseKeys`` in this package's tests asserts every named parameter #: of ``acompletion`` is classified as forwarded, handler-owned or excluded. _COMPLETION_FORWARDED_KEYS = frozenset( @@ -244,6 +245,7 @@ def _request( stream: bool, base_url: str | None = None, api_key: str | None = None, + max_retries: int | None = None, ) -> dict[str, Any]: parameters = select_forwarded_parameters( model_parameters(config_value), _COMPLETION_FORWARDED_KEYS @@ -260,6 +262,10 @@ def _request( parameters["api_base"] = base_url if api_key is not None: parameters["api_key"] = api_key + # LiteLLM hands max_retries to the backing provider's client. Not num_retries, which would also + # start LiteLLM's own retry loop around the call, so a failure would be retried at two layers. + if max_retries is not None: + parameters["max_retries"] = max_retries if tools: parameters["tools"] = tools if stream: @@ -426,7 +432,18 @@ def create_litellm_messages_handler( capture_content: bool | None = None, base_url: str | None = None, api_key: str | None = None, + max_retries: int | None = None, ) -> ProviderHandler: + """Creates a ``ProviderHandler`` for any model LiteLLM can route to. + + Set *max_retries* to change how many times the backing provider's client retries a failed + request (a rate limit, a timeout, a 5xx) before the error reaches the handler. It is passed to + ``acompletion`` on every model turn. ``None`` (the default) leaves LiteLLM's own default in + place, which depends on the backing provider and is 2 for OpenAI-compatible ones. Raise it + when many calls share one rate-limit window, as an evaluation run with many concurrent workers + does: two retries is thin when a whole batch hits the same window at once. A custom + *completion* receives it as a keyword argument too. + """ report_usage( "litellm-messages.createLiteLLMMessagesHandler", PACKAGE_NAME, __version__ ) @@ -435,6 +452,7 @@ def create_litellm_messages_handler( capture_content=capture_content, base_url=base_url, api_key=api_key, + max_retries=max_retries, ) @@ -444,6 +462,7 @@ def _create_litellm_messages_handler( capture_content: bool | None = None, base_url: str | None = None, api_key: str | None = None, + max_retries: int | None = None, ) -> ProviderHandler: """Non-reporting :func:`create_litellm_messages_handler`, used by this package's wrappers.""" content_override = capture_content @@ -495,6 +514,7 @@ async def call( stream=False, base_url=base_url, api_key=api_key, + max_retries=max_retries, ) ) turn_usage = _usage(response) @@ -581,6 +601,7 @@ def stream( capture_content=effective_capture_content(content_override), base_url=base_url, api_key=api_key, + max_retries=max_retries, ) return _create_handler( @@ -602,6 +623,7 @@ async def _stream( capture_content: bool, base_url: str | None = None, api_key: str | None = None, + max_retries: int | None = None, ) -> AsyncGenerator[dict[str, Any], None]: messages = _build_messages(config_value, user_input, variables, history) tools = _build_tools(config_value, handlers) @@ -641,6 +663,7 @@ async def _stream( stream=True, base_url=base_url, api_key=api_key, + max_retries=max_retries, ) ) fragments: dict[int, dict[str, Any]] = {} @@ -764,11 +787,15 @@ def litellm_messages( capture_content = kwargs.pop("capture_content", None) base_url = kwargs.pop("base_url", None) api_key = kwargs.pop("api_key", None) + max_retries = kwargs.pop("max_retries", None) report_usage("litellm-messages.litellmMessages", PACKAGE_NAME, __version__) return _config( key=config_key, handler=_create_litellm_messages_handler( - capture_content=capture_content, base_url=base_url, api_key=api_key + capture_content=capture_content, + base_url=base_url, + api_key=api_key, + max_retries=max_retries, ), **kwargs, ).invoke(user_input, context, variables) diff --git a/packages/litellm-messages/tests/test_handler.py b/packages/litellm-messages/tests/test_handler.py index 5f99d772..dbd76af9 100644 --- a/packages/litellm-messages/tests/test_handler.py +++ b/packages/litellm-messages/tests/test_handler.py @@ -742,3 +742,57 @@ async def call(parameters: dict[str, Any]) -> object: return completion.call_args.kwargs assert await probe_forwarded_keys(self._candidates(), call) == LITELLM_MESSAGES + + +class TestMaxRetries: + """`max_retries` reaches every `acompletion` call, and only when set. + + LiteLLM hands it to the backing provider's client, which is the right layer for rate-limit + retries. An omitted value must leave LiteLLM's default in place rather than replace it. + """ + + async def test_value_reaches_acompletion(self) -> None: + completion = AsyncMock(return_value=_response()) + await create_litellm_messages_handler(completion=completion, max_retries=7)( + CONFIG, "question" + ) + assert completion.call_args.kwargs["max_retries"] == 7 + # The provider-client knob, not LiteLLM's own retry loop as well. + assert "num_retries" not in completion.call_args.kwargs + + async def test_value_reaches_every_streamed_turn(self) -> None: + completion = AsyncMock(return_value=FakeStream([_chunk(text="hi")])) + await _events( + await create_litellm_messages_handler( + completion=completion, max_retries=7 + ).stream(CONFIG, "q") + ) + assert completion.call_args.kwargs["max_retries"] == 7 + + async def test_omitted_value_keeps_the_litellm_default(self) -> None: + completion = AsyncMock(return_value=_response()) + await create_litellm_messages_handler(completion=completion)(CONFIG, "question") + assert "max_retries" not in completion.call_args.kwargs + + async def test_config_cannot_set_it(self) -> None: + completion = AsyncMock(return_value=_response()) + config = { + **CONFIG, + "model": {**CONFIG["model"], "parameters": {"max_retries": 50}}, + } + await create_litellm_messages_handler(completion=completion, max_retries=3)( + config, "question" + ) + assert completion.call_args.kwargs["max_retries"] == 3 + + def test_convenience_wrapper_forwards_it_to_the_handler(self) -> None: + factory = MagicMock() + config = MagicMock() + with ( + patch.object(handler_mod, "_create_litellm_messages_handler", factory), + patch.object(handler_mod, "_config", config), + ): + litellm_messages("k", "q", {}, max_retries=6) + assert factory.call_args.kwargs["max_retries"] == 6 + # config() takes no such argument, so leaving it in kwargs would raise TypeError. + assert "max_retries" not in config.call_args.kwargs diff --git a/packages/openai-messages/api-report/launchdarkly_ai_openai_messages.txt b/packages/openai-messages/api-report/launchdarkly_ai_openai_messages.txt index 23e9598e..d2846f52 100644 --- a/packages/openai-messages/api-report/launchdarkly_ai_openai_messages.txt +++ b/packages/openai-messages/api-report/launchdarkly_ai_openai_messages.txt @@ -1,3 +1,3 @@ # Public API of launchdarkly_ai_openai_messages. Generated by `make api-report`; do not edit. -create_openai_messages_handler(*, capture_content: bool | None = None) -> ProviderHandler +create_openai_messages_handler(*, capture_content: bool | None = None, max_retries: int | None = None) -> ProviderHandler openai_messages(config_key: str, user_input: str, context: LDContext, **kwargs: Any) -> Any diff --git a/packages/openai-messages/src/launchdarkly_ai_openai_messages/handler.py b/packages/openai-messages/src/launchdarkly_ai_openai_messages/handler.py index 5e804e41..0c6544c3 100644 --- a/packages/openai-messages/src/launchdarkly_ai_openai_messages/handler.py +++ b/packages/openai-messages/src/launchdarkly_ai_openai_messages/handler.py @@ -378,7 +378,7 @@ async def _run_model_turn( def create_openai_messages_handler( - *, capture_content: bool | None = None + *, capture_content: bool | None = None, max_retries: int | None = None ) -> ProviderHandler: """ Creates a ``ProviderHandler`` for OpenAI (responses API). @@ -389,20 +389,33 @@ def create_openai_messages_handler( which is off until ``configure`` sets it. ``True`` or ``False`` overrides that default for this handler. Conversation content is PII, so a run emits only metadata, meaning models, token counts, timings and tool names, until capture is turned on. + + Set *max_retries* to change how many times the ``AsyncOpenAI`` client retries a failed + request (a rate limit, a timeout, a 5xx) before the error reaches the handler. ``None`` (the + default) leaves the client's own default of 2 in place. Raise it when many calls share one + rate-limit window, as an evaluation run with many concurrent workers does: two retries is thin + when a whole batch hits the same window at once. The retries happen inside the client, so a + retried turn is still one ``chat`` span. """ report_usage("openai-messages.createOpenAIHandler", PACKAGE_NAME, __version__) - return _create_openai_messages_handler(capture_content=capture_content) + return _create_openai_messages_handler( + capture_content=capture_content, max_retries=max_retries + ) def _create_openai_messages_handler( - *, capture_content: bool | None = None + *, capture_content: bool | None = None, max_retries: int | None = None ) -> ProviderHandler: """Non-reporting :func:`create_openai_messages_handler`, used by this package's wrappers.""" content_override = capture_content import importlib openai_mod = importlib.import_module("openai") - client = openai_mod.AsyncOpenAI() + # Passed only when set, so an omitted value keeps the vendor default rather than a copy of it. + client_options: dict[str, Any] = {} + if max_retries is not None: + client_options["max_retries"] = max_retries + client = openai_mod.AsyncOpenAI(**client_options) async def _call_impl( config: AiConfigRep, @@ -847,14 +860,17 @@ def openai_messages( **kwargs: Any, ) -> Any: """Convenience wrapper: creates a handler and calls config(...).invoke().""" - # Both are lifted out of kwargs: capture_content configures the handler, variables belong to - # the invocation. Leaving either in would pass it to config(), which takes neither, so a caller - # asking for content on spans got a TypeError instead of content. + # All three are lifted out of kwargs: capture_content and max_retries configure the handler, + # variables belong to the invocation. Leaving any in would pass it to config(), which takes + # none of them, so a caller asking for content on spans got a TypeError instead of content. report_usage("openai-messages.openaiMessages", PACKAGE_NAME, __version__) variables = kwargs.pop("variables", None) capture_content = kwargs.pop("capture_content", None) + max_retries = kwargs.pop("max_retries", None) return _config( key=config_key, - handler=_create_openai_messages_handler(capture_content=capture_content), + handler=_create_openai_messages_handler( + capture_content=capture_content, max_retries=max_retries + ), **kwargs, ).invoke(user_input, context, variables=variables) diff --git a/packages/openai-messages/tests/test_handler.py b/packages/openai-messages/tests/test_handler.py index b4bdb69a..576f3b12 100644 --- a/packages/openai-messages/tests/test_handler.py +++ b/packages/openai-messages/tests/test_handler.py @@ -3029,3 +3029,65 @@ async def test_stream(self, mock_openai: MagicMock) -> None: kwargs = mock_openai.responses.stream.call_args.kwargs assert "reasoning" not in kwargs assert kwargs["top_p"] == 0.4 + + +class TestMaxRetries: + """`max_retries` reaches the `AsyncOpenAI` client the factory builds, and only when set. + + The client's own retries are the right layer for rate limits, so an omitted value must leave the + vendor default in place rather than replace it, and nothing may ever turn retries off. + """ + + def _build( + self, monkeypatch: pytest.MonkeyPatch, **kwargs: Any + ) -> tuple[Any, MagicMock]: + import openai + + from launchdarkly_ai_openai_messages import create_openai_messages_handler + + monkeypatch.setenv("OPENAI_API_KEY", "test-key") + real_client = openai.AsyncOpenAI + built: list[Any] = [] + + def _construct(**options: Any) -> Any: + built.append(real_client(**options)) + return built[-1] + + with patch("openai.AsyncOpenAI", side_effect=_construct) as ctor: + create_openai_messages_handler(**kwargs) + return built[0], ctor + + def test_value_reaches_the_constructed_client( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + client, _ctor = self._build(monkeypatch, max_retries=7) + assert client.max_retries == 7 + + def test_omitted_value_keeps_the_vendor_default( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + import openai + + client, ctor = self._build(monkeypatch) + assert "max_retries" not in ctor.call_args.kwargs + assert client.max_retries == openai.DEFAULT_MAX_RETRIES + + def test_explicit_zero_is_the_callers_choice( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + client, _ctor = self._build(monkeypatch, max_retries=0) + assert client.max_retries == 0 + + def test_convenience_wrapper_forwards_it_to_the_factory(self) -> None: + import launchdarkly_ai_openai_messages.handler as handler_mod + + factory = MagicMock() + fake_config = MagicMock() + with ( + patch.object(handler_mod, "_create_openai_messages_handler", factory), + patch.object(handler_mod, "_config", fake_config), + ): + handler_mod.openai_messages("k", "q", {}, max_retries=6) + assert factory.call_args.kwargs["max_retries"] == 6 + # config() takes no such argument, so leaving it in kwargs would raise TypeError. + assert "max_retries" not in fake_config.call_args.kwargs