From 2c6228ba8bbbd13999e71d1dfc304baa2d9e0352 Mon Sep 17 00:00:00 2001 From: Adriana Knight Date: Fri, 23 Oct 2026 16:09:13 -0400 Subject: [PATCH] Add max_retries passthrough arg to built-in handler factories --- .../launchdarkly_ai_claude_messages.txt | 2 +- .../handler.py | 35 ++- .../claude-messages/tests/test_handler.py | 121 ++++++++++ .../api-report/launchdarkly_ai_server.txt | 1 + .../src/launchdarkly_ai_server/__init__.py | 2 + .../src/launchdarkly_ai_server/utils.py | 22 ++ packages/client/tests/test_public_surface.py | 1 + packages/client/tests/test_utils.py | 19 ++ .../launchdarkly_ai_langchain_messages.txt | 2 +- .../handler.py | 68 ++++-- .../langchain-messages/tests/test_handler.py | 211 ++++++++++++++++++ .../launchdarkly_ai_litellm_messages.txt | 2 +- .../handler.py | 59 ++++- .../litellm-messages/tests/test_handler.py | 180 +++++++++++++++ .../launchdarkly_ai_openai_messages.txt | 2 +- .../handler.py | 35 ++- .../openai-messages/tests/test_handler.py | 121 ++++++++++ 17 files changed, 848 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 5222f989..74e308e1 100644 --- a/packages/claude-messages/src/launchdarkly_ai_claude_messages/handler.py +++ b/packages/claude-messages/src/launchdarkly_ai_claude_messages/handler.py @@ -26,6 +26,7 @@ set_tool_call_content_attributes, to_semconv_finish_reason, tool_result_string, + validate_max_retries, ) from launchdarkly_ai_server.agui import TurnEmitter from launchdarkly_ai_server.lifecycle import effective_capture_content @@ -530,7 +531,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). @@ -542,22 +543,37 @@ def create_claude_messages_handler( for this handler. Conversation content is PII. An explicit ``False`` emits only metadata, meaning models, token counts, timings and tool names. 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. Raises ``TypeError`` or ``ValueError`` + when it is not a non-negative int. """ 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.""" + max_retries = validate_max_retries(max_retries) 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, @@ -951,14 +967,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 26f8ab9d..8be37915 100644 --- a/packages/claude-messages/tests/test_handler.py +++ b/packages/claude-messages/tests/test_handler.py @@ -3156,3 +3156,124 @@ 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, 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: + handler = create_claude_messages_handler(**kwargs) + return handler, built[0], ctor + + def test_value_reaches_the_constructed_client( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + _handler, 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 + + _handler, 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: + _handler, client, _ctor = self._build(monkeypatch, max_retries=0) + assert client.max_retries == 0 + + @pytest.mark.parametrize("value", [True, 2.5, "3"]) + def test_non_integer_is_refused_at_construction(self, value: Any) -> None: + from launchdarkly_ai_claude_messages import create_claude_messages_handler + + with patch("anthropic.AsyncAnthropic") as ctor: + with pytest.raises(TypeError, match="max_retries"): + create_claude_messages_handler(max_retries=value) + ctor.assert_not_called() + + def test_negative_is_refused_at_construction(self) -> None: + from launchdarkly_ai_claude_messages import create_claude_messages_handler + + with patch("anthropic.AsyncAnthropic") as ctor: + with pytest.raises(ValueError, match="max_retries"): + create_claude_messages_handler(max_retries=-1) + ctor.assert_not_called() + + async def test_config_cannot_set_it(self, monkeypatch: pytest.MonkeyPatch) -> None: + import launchdarkly_ai_claude_messages.spans as spans_mod + + handler, client, _ctor = self._build(monkeypatch, max_retries=3) + config = { + **CONFIG, + "model": {**CONFIG["model"], "parameters": {"max_retries": 50}}, + } + create = AsyncMock(return_value=_anthropic_response([_text_block("hi")])) + with ( + patch.object(spans_mod, "_HAS_OTEL", False), + patch.object(client.messages, "create", create), + ): + await handler(config, "q", {}, {}) + assert client.max_retries == 3 + # A per-request max_retries would override the client's, so it must not be forwarded. + assert "max_retries" not in create.call_args.kwargs + + async def test_invoke_and_stream_share_the_configured_client( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + import launchdarkly_ai_claude_messages.spans as spans_mod + + handler, client, ctor = self._build(monkeypatch, max_retries=7) + create = AsyncMock(return_value=_anthropic_response([_text_block("hi")])) + stream = MagicMock(return_value=_make_stream_context(["hi"])[0]) + with ( + patch.object(spans_mod, "_HAS_OTEL", False), + patch.object(client.messages, "create", create), + patch.object(client.messages, "stream", stream), + ): + await handler(CONFIG, "q", {}, {}) + async for _event in await handler.stream(CONFIG, "q"): + pass + create.assert_awaited_once() + assert "max_retries" not in create.call_args.kwargs + stream.assert_called_once() + assert "max_retries" not in stream.call_args.kwargs + # One client, built once, serves both paths, so both get its retry count. + assert ctor.call_count == 1 + assert client.max_retries == 7 + + 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/client/api-report/launchdarkly_ai_server.txt b/packages/client/api-report/launchdarkly_ai_server.txt index 4ac0e30a..f07e227d 100644 --- a/packages/client/api-report/launchdarkly_ai_server.txt +++ b/packages/client/api-report/launchdarkly_ai_server.txt @@ -412,6 +412,7 @@ UsageDict [class, @dataclass] input_details: InputTokenDetails | None = None output: int = 0 total: int = 0 +validate_max_retries(max_retries: int | None) -> int | None VariationMeta = dict[str, Any] watch_skills(skills: Sequence[Skill | SkillReference | str] | str | SkillSource, root: str | os.PathLike[str], *, prune: bool = True, timeout: float = 10.0, on_unavailable: OnUnavailable = 'keep', debounce: float = DEFAULT_DEBOUNCE_SECONDS, on_reconcile: Callable[[ReconcileReport], Any] | None = None) -> tuple[ReconcileReport, SkillWatcher] [async] wrap_tool_handlers(tool_handlers: Mapping[str, Callable[..., Any] | NativeTool] | None, user_context: LDContext, track_data: TrackData) -> dict[str, Callable[..., Any]] diff --git a/packages/client/src/launchdarkly_ai_server/__init__.py b/packages/client/src/launchdarkly_ai_server/__init__.py index 48355f8c..2c753801 100644 --- a/packages/client/src/launchdarkly_ai_server/__init__.py +++ b/packages/client/src/launchdarkly_ai_server/__init__.py @@ -163,6 +163,7 @@ set_usage_span_attributes, to_ld_context, to_usage_dict, + validate_max_retries, ) from .utils import _create_handler as _create_handler @@ -242,6 +243,7 @@ "parse_usage", "set_ld_span_attributes", "to_ld_context", + "validate_max_retries", # history "compose_history", "content_to_text", diff --git a/packages/client/src/launchdarkly_ai_server/utils.py b/packages/client/src/launchdarkly_ai_server/utils.py index 5e1fabba..3b343bf4 100644 --- a/packages/client/src/launchdarkly_ai_server/utils.py +++ b/packages/client/src/launchdarkly_ai_server/utils.py @@ -72,6 +72,28 @@ def model_parameters(config: AiConfigRep) -> dict[str, Any]: return dict(parameters) if isinstance(parameters, dict) else {} +def validate_max_retries(max_retries: int | None) -> int | None: + """Checks a handler factory's ``max_retries`` option and returns it unchanged. + + ``None`` means the option was not given, so the vendor default applies. Any other value must be + a non-negative ``int``; ``0`` is a valid count. ``bool`` is refused even though it subclasses + ``int``, since ``max_retries=True`` is a mistake rather than one retry. Factories call this at + construction, so a bad value fails when the handler is built rather than on its first call. + + Raises ``TypeError`` for a non-integer and ``ValueError`` for a negative integer. Both messages + name the option. + """ + if max_retries is None: + return None + if isinstance(max_retries, bool) or not isinstance(max_retries, int): + raise TypeError( + f"max_retries must be a non-negative int or None, got {type(max_retries).__name__}" + ) + if max_retries < 0: + raise ValueError(f"max_retries must be non-negative, got {max_retries}") + return max_retries + + def collapse_messages_to_instructions(config: AiConfigRep) -> AiConfigRep: """ When only an agent handler is available for a messages-mode config, collapse diff --git a/packages/client/tests/test_public_surface.py b/packages/client/tests/test_public_surface.py index 29177a3a..449c9f15 100644 --- a/packages/client/tests/test_public_surface.py +++ b/packages/client/tests/test_public_surface.py @@ -116,6 +116,7 @@ "to_semconv_finish_reason", "to_usage_dict", "tool_result_string", + "validate_max_retries", "wrap_tool_handlers", # Agent Skills "set_skill_store", diff --git a/packages/client/tests/test_utils.py b/packages/client/tests/test_utils.py index dd6035e8..21640135 100644 --- a/packages/client/tests/test_utils.py +++ b/packages/client/tests/test_utils.py @@ -16,6 +16,7 @@ parse_json_with_possible_fences, parse_template, parse_usage, + validate_max_retries, ) from launchdarkly_ai_server.utils import make_graph_track_data, model_parameters @@ -434,3 +435,21 @@ def test_never_reads_custom(self) -> None: result = model_parameters(config) assert "secret" not in result assert "custom" not in result + + +class TestValidateMaxRetries: + """TESTING.md A.22: a factory refuses a value that is not a non-negative int, naming the option.""" + + @pytest.mark.parametrize("value", [None, 0, 1, 7, 50]) + def test_accepts_none_and_non_negative_ints(self, value: int | None) -> None: + assert validate_max_retries(value) == value + + @pytest.mark.parametrize("value", [True, False, 2.0, "3", [3]]) + def test_refuses_non_integers_with_type_error(self, value: Any) -> None: + with pytest.raises(TypeError, match="max_retries"): + validate_max_retries(value) + + @pytest.mark.parametrize("value", [-1, -50]) + def test_refuses_negatives_with_value_error(self, value: int) -> None: + with pytest.raises(ValueError, match="max_retries"): + validate_max_retries(value) 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 0ecaca1b..5cd254ad 100644 --- a/packages/langchain-messages/src/launchdarkly_ai_langchain_messages/handler.py +++ b/packages/langchain-messages/src/launchdarkly_ai_langchain_messages/handler.py @@ -33,6 +33,7 @@ set_output_content_attributes, set_tool_call_content_attributes, tool_result_string, + validate_max_retries, ) from launchdarkly_ai_server.agui import TurnEmitter from launchdarkly_ai_server.lifecycle import effective_capture_content @@ -592,6 +593,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(): @@ -602,6 +604,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 @@ -610,13 +616,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": @@ -636,6 +645,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, ) ) ) @@ -655,6 +665,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") @@ -664,14 +675,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): @@ -769,7 +783,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). @@ -784,15 +801,37 @@ def create_langchain_messages_handler( which is on until ``configure`` turns it off. ``True`` or ``False`` overrides that default for this handler. Conversation content is PII. An explicit ``False`` emits only metadata, meaning models, token counts, timings and tool names. + + 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. Raises ``TypeError`` or ``ValueError`` when it is not a non-negative int. """ 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.""" + max_retries = validate_max_retries(max_retries) + # 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( @@ -833,7 +872,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") @@ -1104,6 +1143,7 @@ def _stream_impl( variables or {}, history, capture_content=effective_capture_content(content_override), + max_retries=max_retries, ) return _create_handler( @@ -1123,6 +1163,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. @@ -1145,7 +1186,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) @@ -1508,16 +1549,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 7bab34d2..56be574a 100644 --- a/packages/langchain-messages/tests/test_handler.py +++ b/packages/langchain-messages/tests/test_handler.py @@ -3053,3 +3053,214 @@ 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_value_reaches_the_real_bedrock_client( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + import importlib + + pytest.importorskip("langchain_aws") + import launchdarkly_ai_langchain_messages.handler as handler_mod + + monkeypatch.setenv("AWS_DEFAULT_REGION", "us-east-1") + monkeypatch.setenv("AWS_ACCESS_KEY_ID", "test-key") + monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "test-secret") + cfg = { + **CONFIG, + "provider": {"name": "Bedrock"}, + "model": {"name": "anthropic.claude-sonnet-4-5"}, + } + model = handler_mod._make_default_chat_model(cfg, importlib, max_retries=7) + # botocore counts attempts, not retries: the first try plus seven retries. + assert model.client.meta.config.retries["total_max_attempts"] == 8 + + @pytest.mark.parametrize( + ("provider", "module", "cls", "fields"), + [ + ( + "OpenAI", + "langchain_openai", + "ChatOpenAI", + _REAL_CHAT_OPENAI.model_fields, + ), + ( + "Anthropic", + "langchain_anthropic", + "ChatAnthropic", + _REAL_CHAT_ANTHROPIC.model_fields, + ), + ("Bedrock", "langchain_aws", "ChatBedrockConverse", _FAKE_BEDROCK_FIELDS), + ], + ) + def test_omitted_value_is_not_passed_to_any_constructor( + self, provider: str, module: str, cls: str, fields: Any + ) -> None: + import launchdarkly_ai_langchain_messages.handler as handler_mod + + ctor = MagicMock() + ctor.model_fields = fields + imports = MagicMock() + imports.import_module.return_value = MagicMock(**{cls: ctor}) + cfg = {**CONFIG, "provider": {"name": provider}, "model": {"name": "m"}} + handler_mod._make_default_chat_model(cfg, imports) + imports.import_module.assert_called_with(module) + # Hard-coding the vendor default would freeze it when the vendor changes it. + assert "max_retries" not in ctor.call_args.kwargs + + def test_omitted_value_keeps_the_anthropic_default(self) -> None: + import importlib + + import anthropic + + 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) + assert model._async_client.max_retries == anthropic.DEFAULT_MAX_RETRIES + + def test_explicit_zero_is_the_callers_choice(self) -> None: + import importlib + + import launchdarkly_ai_langchain_messages.handler as handler_mod + + model = handler_mod._make_default_chat_model(CONFIG, importlib, max_retries=0) + assert model.root_async_client.max_retries == 0 + + def test_refused_alongside_a_caller_supplied_model_factory(self) -> None: + from launchdarkly_ai_langchain_messages import create_langchain_messages_handler + + def factory(_config: Any) -> Any: + return _make_llm("mine") + + with pytest.raises(ValueError, match="max_retries"): + create_langchain_messages_handler(factory, max_retries=7) + + @pytest.mark.parametrize("value", [True, 2.5, "3"]) + def test_non_integer_is_refused_at_construction(self, value: Any) -> None: + from launchdarkly_ai_langchain_messages import create_langchain_messages_handler + + with pytest.raises(TypeError, match="max_retries"): + create_langchain_messages_handler(max_retries=value) + + def test_negative_is_refused_at_construction(self) -> None: + from launchdarkly_ai_langchain_messages import create_langchain_messages_handler + + with pytest.raises(ValueError, match="max_retries"): + create_langchain_messages_handler(max_retries=-1) + + 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 00ca89f2..97b0914d 100644 --- a/packages/litellm-messages/src/launchdarkly_ai_litellm_messages/handler.py +++ b/packages/litellm-messages/src/launchdarkly_ai_litellm_messages/handler.py @@ -30,6 +30,7 @@ set_usage_span_attributes, text_message, to_semconv_finish_reason, + validate_max_retries, ) from launchdarkly_ai_server.lifecycle import effective_capture_content from launchdarkly_ai_server.parameter_forwarding import select_forwarded_parameters @@ -50,7 +51,8 @@ def _json_schema(output_format: dict[str, Any]) -> dict[str, Any]: #: The cross-SDK list for LiteLLM Messages: 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 the retry count (``max_retries`` or ``num_retries``, +#: by route) 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( @@ -135,6 +137,31 @@ def _model_name(config_value: AiConfigRep) -> str: return str((config_value.get("model") or {}).get("name") or "") +#: LiteLLM routes that build an OpenAI SDK client, the only code in LiteLLM that reads +#: ``max_retries`` (TESTING.md A.22, checked against litellm 1.106.2). Every other route ignores it +#: and gets ``num_retries`` instead. Adding a provider here is a spec change. +_OPENAI_CLIENT_PROVIDERS = frozenset({"openai", "azure"}) + + +def _retry_option(model: str) -> str: + """The ``acompletion`` keyword that carries the retry count on *model*'s route. + + ``max_retries`` on an OpenAI-client route, where the client retries. ``num_retries`` everywhere + else, where LiteLLM's own retry loop is the only retry layer. Never both: on an OpenAI-client + route ``completion()`` copies ``num_retries`` into the client's ``max_retries``, so a failure + would be retried at two layers. + """ + import litellm + + # A model LiteLLM cannot resolve (a proxy alias behind a custom completion, say) raises here. + # It gets LiteLLM's own loop, the retry layer that works whatever the route turns out to be. + try: + provider = litellm.get_llm_provider(model)[1] + except Exception: + return "num_retries" + return "max_retries" if provider in _OPENAI_CLIENT_PROVIDERS else "num_retries" + + def _provider_name(config_value: AiConfigRep) -> str: return str((config_value.get("provider") or {}).get("name") or "litellm").lower() @@ -246,6 +273,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 @@ -262,6 +290,9 @@ def _request( parameters["api_base"] = base_url if api_key is not None: parameters["api_key"] = api_key + # Chosen per turn because the route comes from the variation's model, not from the factory. + if max_retries is not None: + parameters[_retry_option(parameters["model"])] = max_retries if tools: parameters["tools"] = tools if stream: @@ -428,7 +459,20 @@ 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 a failed request (a rate limit, a timeout, a 5xx) + is retried before the error reaches the handler. It is passed to ``acompletion`` on every model + turn, as ``max_retries`` when the model routes through LiteLLM's OpenAI client (``openai`` and + ``azure``) and as ``num_retries``, LiteLLM's own retry loop, for every other provider. + ``None`` (the default) passes neither, so LiteLLM's own default applies. 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 the same keyword the route would get. Raises ``TypeError`` or ``ValueError`` when it + is not a non-negative int. + """ report_usage( "litellm-messages.createLiteLLMMessagesHandler", PACKAGE_NAME, __version__ ) @@ -437,6 +481,7 @@ def create_litellm_messages_handler( capture_content=capture_content, base_url=base_url, api_key=api_key, + max_retries=max_retries, ) @@ -446,8 +491,10 @@ 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.""" + max_retries = validate_max_retries(max_retries) content_override = capture_content if completion is None: from litellm import acompletion @@ -497,6 +544,7 @@ async def call( stream=False, base_url=base_url, api_key=api_key, + max_retries=max_retries, ) ) turn_usage = _usage(response) @@ -583,6 +631,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( @@ -604,6 +653,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) @@ -643,6 +693,7 @@ async def _stream( stream=True, base_url=base_url, api_key=api_key, + max_retries=max_retries, ) ) fragments: dict[int, dict[str, Any]] = {} @@ -766,11 +817,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 839d6d66..740975bb 100644 --- a/packages/litellm-messages/tests/test_handler.py +++ b/packages/litellm-messages/tests/test_handler.py @@ -742,3 +742,183 @@ 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 under the keyword the model's route reads. + + TESTING.md A.22: ``max_retries`` on LiteLLM's OpenAI-client routes (``openai``, ``azure``), + ``num_retries`` on every other route, never both, and neither when the option is omitted. + """ + + @staticmethod + def _config(model: str, **extra: Any) -> dict[str, Any]: + return {**CONFIG, "model": {"name": model}, **extra} + + @pytest.mark.parametrize( + ("model", "sent", "withheld"), + [ + ("openai/gpt-4o", "max_retries", "num_retries"), + ("gpt-4o", "max_retries", "num_retries"), + ("azure/my-deployment", "max_retries", "num_retries"), + ("anthropic/claude-sonnet-4-5", "num_retries", "max_retries"), + ("bedrock/anthropic.claude-sonnet-4-5", "num_retries", "max_retries"), + ("groq/llama-3.1-8b-instant", "num_retries", "max_retries"), + ], + ) + async def test_keyword_follows_the_route( + self, model: str, sent: str, withheld: str + ) -> None: + completion = AsyncMock(return_value=_response()) + await create_litellm_messages_handler(completion=completion, max_retries=8)( + self._config(model), "question" + ) + assert completion.call_args.kwargs[sent] == 8 + # Both at once would retry each failure at two layers. + assert withheld not in completion.call_args.kwargs + + async def test_unresolvable_model_gets_litellms_own_loop(self) -> None: + completion = AsyncMock(return_value=_response()) + await create_litellm_messages_handler(completion=completion, max_retries=8)( + self._config("my-proxy-alias"), "question" + ) + assert completion.call_args.kwargs["num_retries"] == 8 + assert "max_retries" not in completion.call_args.kwargs + + async def test_value_reaches_every_tool_loop_turn(self) -> None: + call = _tool_call("weather", "call-1", '{"city":"Austin"}') + completion = AsyncMock( + side_effect=[_response(None, tool_calls=[call]), _response("sunny")] + ) + config = self._config( + "openai/gpt-4o", + tools={"weather": {"description": "weather", "parameters": {}}}, + ) + await create_litellm_messages_handler(completion=completion, max_retries=7)( + config, "q", {"weather": AsyncMock(return_value="sunny")} + ) + assert completion.await_count == 2 + assert [c.kwargs["max_retries"] for c in completion.call_args_list] == [7, 7] + + async def test_value_reaches_every_streamed_tool_loop_turn(self) -> None: + first = FakeStream( + [ + _chunk( + tool_calls=[ + SimpleNamespace( + index=0, + id="call-1", + function=SimpleNamespace( + name="weather", arguments='{"city":"Oslo"}' + ), + ) + ] + ), + ] + ) + second = FakeStream([_chunk(text="cold")]) + completion = AsyncMock(side_effect=[first, second]) + config = { + **CONFIG, + "tools": {"weather": {"description": "x", "parameters": {}}}, + } + await _events( + await create_litellm_messages_handler( + completion=completion, max_retries=7 + ).stream(config, "q", {"weather": AsyncMock(return_value="cold")}) + ) + assert completion.await_count == 2 + assert [c.kwargs["num_retries"] for c in completion.call_args_list] == [7, 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 + assert "num_retries" not in completion.call_args.kwargs + + async def test_explicit_zero_is_the_callers_choice(self) -> None: + completion = AsyncMock(return_value=_response()) + await create_litellm_messages_handler(completion=completion, max_retries=0)( + CONFIG, "question" + ) + assert completion.call_args.kwargs["num_retries"] == 0 + + async def test_config_cannot_set_it(self) -> None: + completion = AsyncMock(return_value=_response()) + config = self._config( + "openai/gpt-4o", + ) + config["model"]["parameters"] = {"max_retries": 50, "num_retries": 50} + await create_litellm_messages_handler(completion=completion, max_retries=3)( + config, "question" + ) + assert completion.call_args.kwargs["max_retries"] == 3 + assert "num_retries" not in completion.call_args.kwargs + + @pytest.mark.parametrize("value", [True, 2.5, "3"]) + def test_non_integer_is_refused_at_construction(self, value: Any) -> None: + with pytest.raises(TypeError, match="max_retries"): + create_litellm_messages_handler(completion=AsyncMock(), max_retries=value) + + def test_negative_is_refused_at_construction(self) -> None: + with pytest.raises(ValueError, match="max_retries"): + create_litellm_messages_handler(completion=AsyncMock(), max_retries=-1) + + @pytest.fixture + def always_500(self) -> Any: + """A local endpoint that fails every request with a 500, counting them. + + A 500 rather than a 429: LiteLLM backs off exponentially, up to 10 s a wait, only after a + rate-limit error, so a 500 keeps the retries immediate. + """ + import threading + from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer + + hits: list[str] = [] + + class _Fail(BaseHTTPRequestHandler): + def do_POST(self) -> None: + hits.append(self.path) + self.rfile.read(int(self.headers.get("content-length", 0))) + self.send_response(500) + self.send_header("content-type", "application/json") + self.end_headers() + self.wfile.write( + b'{"type":"error","error":{"type":"api_error","message":"boom"}}' + ) + + def log_message(self, *_args: Any) -> None: + pass + + server = ThreadingHTTPServer(("127.0.0.1", 0), _Fail) + threading.Thread(target=server.serve_forever, daemon=True).start() + try: + yield f"http://127.0.0.1:{server.server_address[1]}", hits + finally: + server.shutdown() + server.server_close() + + async def test_real_litellm_retries_a_non_openai_route( + self, always_500: Any + ) -> None: + # The completion= stub cannot show LiteLLM honours the keyword, so this runs the real + # acompletion. On this route max_retries alone would be ignored: one request, not four. + base_url, hits = always_500 + handler = create_litellm_messages_handler( + base_url=base_url, api_key="test-key", max_retries=3 + ) + with pytest.raises(Exception): # noqa: B017 - LiteLLM's error type is not the point + await handler(self._config("anthropic/claude-sonnet-4-5"), "q") + assert len(hits) == 4 + + 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 0b0d63f9..8bd81147 100644 --- a/packages/openai-messages/src/launchdarkly_ai_openai_messages/handler.py +++ b/packages/openai-messages/src/launchdarkly_ai_openai_messages/handler.py @@ -28,6 +28,7 @@ set_output_content_attributes, set_tool_call_content_attributes, tool_result_string, + validate_max_retries, ) from launchdarkly_ai_server.agui import TurnEmitter from launchdarkly_ai_server.lifecycle import effective_capture_content @@ -379,7 +380,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). @@ -390,20 +391,35 @@ def create_openai_messages_handler( which is on until ``configure`` turns it off. ``True`` or ``False`` overrides that default for this handler. Conversation content is PII. An explicit ``False`` emits only metadata, meaning models, token counts, timings and tool names. + + 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. Raises ``TypeError`` or ``ValueError`` when it is not + a non-negative int. """ 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.""" + max_retries = validate_max_retries(max_retries) 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, @@ -848,14 +864,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 44f30794..2f5ae956 100644 --- a/packages/openai-messages/tests/test_handler.py +++ b/packages/openai-messages/tests/test_handler.py @@ -3059,3 +3059,124 @@ 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, 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: + handler = create_openai_messages_handler(**kwargs) + return handler, built[0], ctor + + def test_value_reaches_the_constructed_client( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + _handler, 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 + + _handler, 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: + _handler, client, _ctor = self._build(monkeypatch, max_retries=0) + assert client.max_retries == 0 + + @pytest.mark.parametrize("value", [True, 2.5, "3"]) + def test_non_integer_is_refused_at_construction(self, value: Any) -> None: + from launchdarkly_ai_openai_messages import create_openai_messages_handler + + with patch("openai.AsyncOpenAI") as ctor: + with pytest.raises(TypeError, match="max_retries"): + create_openai_messages_handler(max_retries=value) + ctor.assert_not_called() + + def test_negative_is_refused_at_construction(self) -> None: + from launchdarkly_ai_openai_messages import create_openai_messages_handler + + with patch("openai.AsyncOpenAI") as ctor: + with pytest.raises(ValueError, match="max_retries"): + create_openai_messages_handler(max_retries=-1) + ctor.assert_not_called() + + async def test_config_cannot_set_it(self, monkeypatch: pytest.MonkeyPatch) -> None: + import launchdarkly_ai_openai_messages.spans as spans_mod + + handler, client, _ctor = self._build(monkeypatch, max_retries=3) + config = { + **CONFIG, + "model": {**CONFIG["model"], "parameters": {"max_retries": 50}}, + } + create = AsyncMock(return_value=_make_response()) + with ( + patch.object(spans_mod, "_HAS_OTEL", False), + patch.object(client.responses, "create", create), + ): + await handler(config, "q", {}, {}) + assert client.max_retries == 3 + # A per-request max_retries would override the client's, so it must not be forwarded. + assert "max_retries" not in create.call_args.kwargs + + async def test_invoke_and_stream_share_the_configured_client( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + import launchdarkly_ai_openai_messages.spans as spans_mod + + handler, client, ctor = self._build(monkeypatch, max_retries=7) + create = AsyncMock(return_value=_make_response()) + stream = MagicMock(return_value=_make_openai_stream_context(["hi"])) + with ( + patch.object(spans_mod, "_HAS_OTEL", False), + patch.object(client.responses, "create", create), + patch.object(client.responses, "stream", stream), + ): + await handler(CONFIG, "q", {}, {}) + async for _event in await handler.stream(CONFIG, "q"): + pass + create.assert_awaited_once() + assert "max_retries" not in create.call_args.kwargs + stream.assert_called_once() + assert "max_retries" not in stream.call_args.kwargs + # One client, built once, serves both paths, so both get its retry count. + assert ctor.call_count == 1 + assert client.max_retries == 7 + + 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