From de94aff776f037a37581f51c1f1d3f24a1b91202 Mon Sep 17 00:00:00 2001 From: Sebastian Husch Lee Date: Tue, 11 Aug 2026 13:29:54 +0200 Subject: [PATCH 1/4] Initial version of max token resolution --- haystack/components/generators/chat/utils.py | 78 ++++++++++++ test/components/generators/chat/test_utils.py | 119 ++++++++++++++++++ 2 files changed, 197 insertions(+) create mode 100644 haystack/components/generators/chat/utils.py create mode 100644 test/components/generators/chat/test_utils.py diff --git a/haystack/components/generators/chat/utils.py b/haystack/components/generators/chat/utils.py new file mode 100644 index 00000000000..cdbdd8dab55 --- /dev/null +++ b/haystack/components/generators/chat/utils.py @@ -0,0 +1,78 @@ +# SPDX-FileCopyrightText: 2022-present deepset GmbH +# +# SPDX-License-Identifier: Apache-2.0 + +from typing import Any + +from haystack.components.generators.chat.types import ChatGenerator + +# The `generation_kwargs` key that caps a reply's length, per Chat Generator. +# +# Providers do not agree on a name for this setting and the `ChatGenerator` protocol does not standardize it, so a +# caller that wants to bound a reply has to know which key the generator expects. Generators are matched by class name +# rather than by `isinstance`, because most of the entries below live in `haystack-core-integrations` and are not +# importable from here. Only the most derived match is used, so an unlisted subclass resolves through the generator it +# inherits from: that is what covers the many OpenAI-compatible integrations without naming any of them. +_OUTPUT_TOKEN_LIMIT_KEYS = { + # Haystack + "OpenAIChatGenerator": "max_completion_tokens", + "AzureOpenAIChatGenerator": "max_completion_tokens", + "OpenAIResponsesChatGenerator": "max_output_tokens", + "AzureOpenAIResponsesChatGenerator": "max_output_tokens", + # haystack-core-integrations + "AmazonBedrockChatGenerator": "maxTokens", + "AnthropicChatGenerator": "max_tokens", + "GoogleAIGeminiChatGenerator": "max_output_tokens", + "GoogleGenAIChatGenerator": "max_output_tokens", + "HuggingFaceAPIChatGenerator": "max_tokens", + "LiteLLMChatGenerator": "max_tokens", + "LlamaCppChatGenerator": "max_tokens", + "OllamaChatGenerator": "num_predict", + "TransformersChatGenerator": "max_new_tokens", + "VertexAIGeminiChatGenerator": "max_output_tokens", + "VLLMChatGenerator": "max_tokens", + "WatsonxChatGenerator": "max_new_tokens", +} + + +def _generator_output_token_limit_key(chat_generator: ChatGenerator) -> str | None: + """ + Return the `generation_kwargs` key that limits a Chat Generator's output length. + + :param chat_generator: The generator to look up. + :returns: The key the generator expects, or None when the generator is not recognized. + """ + # Walk from the most derived class, so a provider's own entry wins over the generator it inherits from. + for cls in type(chat_generator).__mro__: + key = _OUTPUT_TOKEN_LIMIT_KEYS.get(cls.__name__) + if key is not None: + return key + return None + + +def _resolve_output_token_limit(chat_generator: ChatGenerator, default_limit: int) -> tuple[int, dict[str, Any] | None]: + """ + Resolve an effective output-token limit and the runtime kwargs that hold a Chat Generator to it. + + A recognized limit already configured on the generator wins and is not repeated at runtime. When a recognized + generator has no limit configured, the default is returned as that provider's runtime setting. An unrecognized + generator receives no runtime setting, because the `ChatGenerator` protocol guarantees nothing beyond + `run(messages)` and guessing a key would raise a `TypeError` in the generator. + + :param chat_generator: The generator whose output should be limited. + :param default_limit: The positive fallback output-token limit. + :returns: The effective limit, and the `generation_kwargs` to pass at runtime, or None when there are none. A + caller that gets None can still send the limit as prompt guidance and measure the reply itself. + """ + limit_key = _generator_output_token_limit_key(chat_generator=chat_generator) + if limit_key is None: + return default_limit, None + + configured = getattr(chat_generator, "generation_kwargs", None) + if isinstance(configured, dict) and limit_key in configured: + value = configured[limit_key] + if isinstance(value, int) and not isinstance(value, bool) and value > 0: + return value, None + # The generator owns this setting. Do not silently replace an invalid value; let it report the problem. + return default_limit, None + return default_limit, {limit_key: default_limit} diff --git a/test/components/generators/chat/test_utils.py b/test/components/generators/chat/test_utils.py new file mode 100644 index 00000000000..9c7fcadffed --- /dev/null +++ b/test/components/generators/chat/test_utils.py @@ -0,0 +1,119 @@ +# SPDX-FileCopyrightText: 2022-present deepset GmbH +# +# SPDX-License-Identifier: Apache-2.0 + +from typing import Any + +import pytest + +from haystack.components.generators.chat import ( + AzureOpenAIChatGenerator, + AzureOpenAIResponsesChatGenerator, + OpenAIChatGenerator, + OpenAIResponsesChatGenerator, +) +from haystack.components.generators.chat.types import ChatGenerator +from haystack.components.generators.chat.utils import _generator_output_token_limit_key, _resolve_output_token_limit +from haystack.dataclasses import ChatMessage + + +def integration_generator(class_name: str) -> ChatGenerator: + """ + Stand in for a Chat Generator that lives in `haystack-core-integrations`. + + Those packages are not importable from here, which is why the lookup matches on class name. A stand-in named the + same way is therefore an accurate test of the mechanism, but it cannot catch a rename on the integration side. + """ + + def run(self: Any, messages: list[ChatMessage], **kwargs: Any) -> dict[str, Any]: + return {"replies": []} + + generator: ChatGenerator = type(class_name, (), {"run": run})() + return generator + + +class TestGeneratorOutputTokenLimitKey: + @pytest.mark.parametrize( + ("generator", "expected"), + [ + pytest.param(OpenAIChatGenerator(), "max_completion_tokens", id="openai"), + pytest.param( + AzureOpenAIChatGenerator(azure_endpoint="https://test.openai.azure.com"), + "max_completion_tokens", + id="azure-openai", + ), + pytest.param(OpenAIResponsesChatGenerator(), "max_output_tokens", id="openai-responses"), + pytest.param( + AzureOpenAIResponsesChatGenerator(azure_endpoint="https://test.openai.azure.com"), + "max_output_tokens", + id="azure-openai-responses", + ), + ], + ) + def test_recognizes_built_in_generators(self, generator, expected): + assert _generator_output_token_limit_key(chat_generator=generator) == expected + + @pytest.mark.parametrize( + ("class_name", "expected"), + [ + ("AmazonBedrockChatGenerator", "maxTokens"), + ("AnthropicChatGenerator", "max_tokens"), + ("GoogleAIGeminiChatGenerator", "max_output_tokens"), + ("GoogleGenAIChatGenerator", "max_output_tokens"), + ("HuggingFaceAPIChatGenerator", "max_tokens"), + ("LiteLLMChatGenerator", "max_tokens"), + ("LlamaCppChatGenerator", "max_tokens"), + ("OllamaChatGenerator", "num_predict"), + ("TransformersChatGenerator", "max_new_tokens"), + ("VertexAIGeminiChatGenerator", "max_output_tokens"), + ("VLLMChatGenerator", "max_tokens"), + ("WatsonxChatGenerator", "max_new_tokens"), + ], + ) + def test_recognizes_integration_generators(self, class_name, expected): + assert _generator_output_token_limit_key(chat_generator=integration_generator(class_name)) == expected + + def test_a_subclass_resolves_through_the_generator_it_inherits_from(self): + # This is what covers the OpenAI-compatible integrations, such as Mistral, OpenRouter, and TogetherAI, without + # naming any of them. + class ProviderChatGenerator(OpenAIChatGenerator): + pass + + assert _generator_output_token_limit_key(chat_generator=ProviderChatGenerator()) == "max_completion_tokens" + + def test_the_most_derived_entry_wins(self): + # No shipped generator depends on this today, since every subclass that is listed agrees with its base. It is + # pinned so that adding an entry for a subclass cannot be silently overridden by the base it inherits from. + base = type("OpenAIChatGenerator", (), {}) + derived = type("OpenAIResponsesChatGenerator", (base,), {}) + + assert _generator_output_token_limit_key(chat_generator=derived()) == "max_output_tokens" + + def test_unknown_generator_is_not_recognized(self): + assert _generator_output_token_limit_key(chat_generator=integration_generator("MysteryChatGenerator")) is None + + +class TestResolveOutputTokenLimit: + def test_sends_the_default_as_the_providers_runtime_setting(self): + assert _resolve_output_token_limit(chat_generator=OpenAIChatGenerator(), default_limit=100) == ( + 100, + {"max_completion_tokens": 100}, + ) + + def test_a_configured_limit_wins_and_is_not_repeated_at_runtime(self): + generator = OpenAIChatGenerator(generation_kwargs={"temperature": 0, "max_completion_tokens": 23}) + original = dict(generator.generation_kwargs) + + assert _resolve_output_token_limit(chat_generator=generator, default_limit=100) == (23, None) + assert generator.generation_kwargs == original + + @pytest.mark.parametrize("value", [0, -1, True, "512", None], ids=["zero", "negative", "bool", "string", "none"]) + def test_an_invalid_configured_limit_is_left_for_the_generator_to_report(self, value): + generator = OpenAIChatGenerator(generation_kwargs={"max_completion_tokens": value}) + # The generator owns the setting, so it is neither used nor silently overwritten at runtime. + assert _resolve_output_token_limit(chat_generator=generator, default_limit=100) == (100, None) + + def test_an_unknown_generator_receives_no_guessed_setting(self): + # The protocol only guarantees `run(messages)`, so passing a guessed kwarg would raise inside the generator. + generator = integration_generator("MysteryChatGenerator") + assert _resolve_output_token_limit(chat_generator=generator, default_limit=100) == (100, None) From fbb8b8f142f4b28dfe661b89538e043740520eb3 Mon Sep 17 00:00:00 2001 From: Sebastian Husch Lee Date: Tue, 11 Aug 2026 13:37:03 +0200 Subject: [PATCH 2/4] add more docs --- haystack/components/generators/chat/utils.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/haystack/components/generators/chat/utils.py b/haystack/components/generators/chat/utils.py index cdbdd8dab55..4748fde054b 100644 --- a/haystack/components/generators/chat/utils.py +++ b/haystack/components/generators/chat/utils.py @@ -54,9 +54,13 @@ def _resolve_output_token_limit(chat_generator: ChatGenerator, default_limit: in """ Resolve an effective output-token limit and the runtime kwargs that hold a Chat Generator to it. - A recognized limit already configured on the generator wins and is not repeated at runtime. When a recognized - generator has no limit configured, the default is returned as that provider's runtime setting. An unrecognized - generator receives no runtime setting, because the `ChatGenerator` protocol guarantees nothing beyond + A recognized limit already configured on the generator wins and is not repeated at runtime. This counts a limit the + generator set for itself: `HuggingFaceAPIChatGenerator` and `TransformersChatGenerator` both leave a default of 512 + in their `generation_kwargs`, and reading the dict cannot tell that apart from a deliberate choice, so a caller + asking for more than that gets the generator's number back. + + When a recognized generator has no limit configured, the default is returned as that provider's runtime setting. An + unrecognized generator receives no runtime setting, because the `ChatGenerator` protocol guarantees nothing beyond `run(messages)` and guessing a key would raise a `TypeError` in the generator. :param chat_generator: The generator whose output should be limited. From a11e3c7d0e4e65e43bc646ef34b60fbf79e1cb6c Mon Sep 17 00:00:00 2001 From: Sebastian Husch Lee Date: Tue, 11 Aug 2026 15:28:20 +0200 Subject: [PATCH 3/4] fix docstrings --- haystack/components/generators/chat/utils.py | 44 +++++------------ test/components/generators/chat/test_utils.py | 48 +++++++++++-------- 2 files changed, 39 insertions(+), 53 deletions(-) diff --git a/haystack/components/generators/chat/utils.py b/haystack/components/generators/chat/utils.py index 4748fde054b..0408a9ad55f 100644 --- a/haystack/components/generators/chat/utils.py +++ b/haystack/components/generators/chat/utils.py @@ -6,13 +6,9 @@ from haystack.components.generators.chat.types import ChatGenerator -# The `generation_kwargs` key that caps a reply's length, per Chat Generator. -# -# Providers do not agree on a name for this setting and the `ChatGenerator` protocol does not standardize it, so a -# caller that wants to bound a reply has to know which key the generator expects. Generators are matched by class name -# rather than by `isinstance`, because most of the entries below live in `haystack-core-integrations` and are not -# importable from here. Only the most derived match is used, so an unlisted subclass resolves through the generator it -# inherits from: that is what covers the many OpenAI-compatible integrations without naming any of them. +# The `generation_kwargs` key that caps a reply's length, per Chat Generator. Providers do not agree on a name and the +# `ChatGenerator` protocol does not standardize it. Keyed by class name because most of these live in +# `haystack-core-integrations` and cannot be imported here. _OUTPUT_TOKEN_LIMIT_KEYS = { # Haystack "OpenAIChatGenerator": "max_completion_tokens", @@ -42,7 +38,8 @@ def _generator_output_token_limit_key(chat_generator: ChatGenerator) -> str | No :param chat_generator: The generator to look up. :returns: The key the generator expects, or None when the generator is not recognized. """ - # Walk from the most derived class, so a provider's own entry wins over the generator it inherits from. + # `__mro__` is the class followed by its bases, so an unlisted generator still matches through one it inherits + # from, such as the integrations built on `OpenAIChatGenerator`. Most derived first, so its own entry wins. for cls in type(chat_generator).__mro__: key = _OUTPUT_TOKEN_LIMIT_KEYS.get(cls.__name__) if key is not None: @@ -50,33 +47,16 @@ def _generator_output_token_limit_key(chat_generator: ChatGenerator) -> str | No return None -def _resolve_output_token_limit(chat_generator: ChatGenerator, default_limit: int) -> tuple[int, dict[str, Any] | None]: +def _run_kwargs_with_output_limit(chat_generator: ChatGenerator, limit: int) -> dict[str, Any]: """ - Resolve an effective output-token limit and the runtime kwargs that hold a Chat Generator to it. - - A recognized limit already configured on the generator wins and is not repeated at runtime. This counts a limit the - generator set for itself: `HuggingFaceAPIChatGenerator` and `TransformersChatGenerator` both leave a default of 512 - in their `generation_kwargs`, and reading the dict cannot tell that apart from a deliberate choice, so a caller - asking for more than that gets the generator's number back. - - When a recognized generator has no limit configured, the default is returned as that provider's runtime setting. An - unrecognized generator receives no runtime setting, because the `ChatGenerator` protocol guarantees nothing beyond - `run(messages)` and guessing a key would raise a `TypeError` in the generator. + Return the `run` kwargs that set a Chat Generator's output-token limit to `limit`. :param chat_generator: The generator whose output should be limited. - :param default_limit: The positive fallback output-token limit. - :returns: The effective limit, and the `generation_kwargs` to pass at runtime, or None when there are none. A - caller that gets None can still send the limit as prompt guidance and measure the reply itself. + :param limit: The positive output-token limit to set. + :returns: The kwargs, or an empty dict when the generator is not recognized. """ limit_key = _generator_output_token_limit_key(chat_generator=chat_generator) if limit_key is None: - return default_limit, None - - configured = getattr(chat_generator, "generation_kwargs", None) - if isinstance(configured, dict) and limit_key in configured: - value = configured[limit_key] - if isinstance(value, int) and not isinstance(value, bool) and value > 0: - return value, None - # The generator owns this setting. Do not silently replace an invalid value; let it report the problem. - return default_limit, None - return default_limit, {limit_key: default_limit} + # Nothing at all rather than an empty `generation_kwargs`, which an unrecognized generator need not accept. + return {} + return {"generation_kwargs": {limit_key: limit}} diff --git a/test/components/generators/chat/test_utils.py b/test/components/generators/chat/test_utils.py index 9c7fcadffed..1a7fa344b81 100644 --- a/test/components/generators/chat/test_utils.py +++ b/test/components/generators/chat/test_utils.py @@ -13,16 +13,16 @@ OpenAIResponsesChatGenerator, ) from haystack.components.generators.chat.types import ChatGenerator -from haystack.components.generators.chat.utils import _generator_output_token_limit_key, _resolve_output_token_limit +from haystack.components.generators.chat.utils import _generator_output_token_limit_key, _run_kwargs_with_output_limit from haystack.dataclasses import ChatMessage def integration_generator(class_name: str) -> ChatGenerator: """ - Stand in for a Chat Generator that lives in `haystack-core-integrations`. + Return an object named `class_name` that satisfies `ChatGenerator`. - Those packages are not importable from here, which is why the lookup matches on class name. A stand-in named the - same way is therefore an accurate test of the mechanism, but it cannot catch a rename on the integration side. + It stands in for a generator from `haystack-core-integrations`, which cannot be imported here. Lookup is by class + name, so the name is the only part that has to match; a rename on the integration side goes unnoticed. """ def run(self: Any, messages: list[ChatMessage], **kwargs: Any) -> dict[str, Any]: @@ -93,27 +93,33 @@ def test_unknown_generator_is_not_recognized(self): assert _generator_output_token_limit_key(chat_generator=integration_generator("MysteryChatGenerator")) is None -class TestResolveOutputTokenLimit: - def test_sends_the_default_as_the_providers_runtime_setting(self): - assert _resolve_output_token_limit(chat_generator=OpenAIChatGenerator(), default_limit=100) == ( - 100, - {"max_completion_tokens": 100}, - ) +class TestRunKwargsWithOutputLimit: + def test_passes_the_limit_as_the_providers_runtime_setting(self): + kwargs = _run_kwargs_with_output_limit(chat_generator=OpenAIChatGenerator(), limit=100) - def test_a_configured_limit_wins_and_is_not_repeated_at_runtime(self): + assert kwargs == {"generation_kwargs": {"max_completion_tokens": 100}} + + def test_overrides_a_limit_the_generator_already_configures(self): + # A caller that reserves room for a reply of a given size needs that size honored, so the runtime value wins. + generator = OpenAIChatGenerator(generation_kwargs={"temperature": 0, "max_completion_tokens": 23}) + + kwargs = _run_kwargs_with_output_limit(chat_generator=generator, limit=100) + + assert kwargs == {"generation_kwargs": {"max_completion_tokens": 100}} + # Only this call is affected; the generator keeps its own settings, including the ones it is not asked about. + assert generator.generation_kwargs == {"temperature": 0, "max_completion_tokens": 23} + + def test_the_override_reaches_the_generator(self): + # Generators merge runtime kwargs over configured ones, which is what makes the override above take effect. generator = OpenAIChatGenerator(generation_kwargs={"temperature": 0, "max_completion_tokens": 23}) - original = dict(generator.generation_kwargs) + kwargs = _run_kwargs_with_output_limit(chat_generator=generator, limit=100) - assert _resolve_output_token_limit(chat_generator=generator, default_limit=100) == (23, None) - assert generator.generation_kwargs == original + merged = {**generator.generation_kwargs, **kwargs["generation_kwargs"]} - @pytest.mark.parametrize("value", [0, -1, True, "512", None], ids=["zero", "negative", "bool", "string", "none"]) - def test_an_invalid_configured_limit_is_left_for_the_generator_to_report(self, value): - generator = OpenAIChatGenerator(generation_kwargs={"max_completion_tokens": value}) - # The generator owns the setting, so it is neither used nor silently overwritten at runtime. - assert _resolve_output_token_limit(chat_generator=generator, default_limit=100) == (100, None) + assert merged == {"temperature": 0, "max_completion_tokens": 100} def test_an_unknown_generator_receives_no_guessed_setting(self): - # The protocol only guarantees `run(messages)`, so passing a guessed kwarg would raise inside the generator. + # The protocol only guarantees `run(messages)`, so even an empty `generation_kwargs` could raise a TypeError. generator = integration_generator("MysteryChatGenerator") - assert _resolve_output_token_limit(chat_generator=generator, default_limit=100) == (100, None) + + assert _run_kwargs_with_output_limit(chat_generator=generator, limit=100) == {} From 56df26964d3bfdc5187c0bbbd3ad06cfbd51c9ea Mon Sep 17 00:00:00 2001 From: Sebastian Husch Lee Date: Wed, 12 Aug 2026 09:24:31 +0200 Subject: [PATCH 4/4] simplify --- haystack/components/generators/chat/utils.py | 39 +++++------- test/components/generators/chat/test_utils.py | 62 +++++-------------- 2 files changed, 32 insertions(+), 69 deletions(-) diff --git a/haystack/components/generators/chat/utils.py b/haystack/components/generators/chat/utils.py index 0408a9ad55f..618816bb47a 100644 --- a/haystack/components/generators/chat/utils.py +++ b/haystack/components/generators/chat/utils.py @@ -2,8 +2,6 @@ # # SPDX-License-Identifier: Apache-2.0 -from typing import Any - from haystack.components.generators.chat.types import ChatGenerator # The `generation_kwargs` key that caps a reply's length, per Chat Generator. Providers do not agree on a name and the @@ -16,14 +14,28 @@ "OpenAIResponsesChatGenerator": "max_output_tokens", "AzureOpenAIResponsesChatGenerator": "max_output_tokens", # haystack-core-integrations + "AIMLAPIChatGenerator": "max_tokens", "AmazonBedrockChatGenerator": "maxTokens", "AnthropicChatGenerator": "max_tokens", + "AnthropicFoundryChatGenerator": "max_tokens", + "AnthropicVertexChatGenerator": "max_tokens", + "CohereChatGenerator": "max_tokens", + "CometAPIChatGenerator": "max_tokens", + "EdenAIChatGenerator": "max_tokens", "GoogleAIGeminiChatGenerator": "max_output_tokens", "GoogleGenAIChatGenerator": "max_output_tokens", "HuggingFaceAPIChatGenerator": "max_tokens", "LiteLLMChatGenerator": "max_tokens", "LlamaCppChatGenerator": "max_tokens", + "LlamaStackChatGenerator": "max_tokens", + "MistralChatGenerator": "max_tokens", + "NvidiaChatGenerator": "max_tokens", "OllamaChatGenerator": "num_predict", + "OpenRouterChatGenerator": "max_tokens", + "OrcaRouterChatGenerator": "max_tokens", + "PerplexityChatGenerator": "max_output_tokens", + "STACKITChatGenerator": "max_tokens", + "TogetherAIChatGenerator": "max_tokens", "TransformersChatGenerator": "max_new_tokens", "VertexAIGeminiChatGenerator": "max_output_tokens", "VLLMChatGenerator": "max_tokens", @@ -38,25 +50,4 @@ def _generator_output_token_limit_key(chat_generator: ChatGenerator) -> str | No :param chat_generator: The generator to look up. :returns: The key the generator expects, or None when the generator is not recognized. """ - # `__mro__` is the class followed by its bases, so an unlisted generator still matches through one it inherits - # from, such as the integrations built on `OpenAIChatGenerator`. Most derived first, so its own entry wins. - for cls in type(chat_generator).__mro__: - key = _OUTPUT_TOKEN_LIMIT_KEYS.get(cls.__name__) - if key is not None: - return key - return None - - -def _run_kwargs_with_output_limit(chat_generator: ChatGenerator, limit: int) -> dict[str, Any]: - """ - Return the `run` kwargs that set a Chat Generator's output-token limit to `limit`. - - :param chat_generator: The generator whose output should be limited. - :param limit: The positive output-token limit to set. - :returns: The kwargs, or an empty dict when the generator is not recognized. - """ - limit_key = _generator_output_token_limit_key(chat_generator=chat_generator) - if limit_key is None: - # Nothing at all rather than an empty `generation_kwargs`, which an unrecognized generator need not accept. - return {} - return {"generation_kwargs": {limit_key: limit}} + return _OUTPUT_TOKEN_LIMIT_KEYS.get(type(chat_generator).__name__) diff --git a/test/components/generators/chat/test_utils.py b/test/components/generators/chat/test_utils.py index 1a7fa344b81..97ab0cdd00e 100644 --- a/test/components/generators/chat/test_utils.py +++ b/test/components/generators/chat/test_utils.py @@ -13,7 +13,7 @@ OpenAIResponsesChatGenerator, ) from haystack.components.generators.chat.types import ChatGenerator -from haystack.components.generators.chat.utils import _generator_output_token_limit_key, _run_kwargs_with_output_limit +from haystack.components.generators.chat.utils import _generator_output_token_limit_key from haystack.dataclasses import ChatMessage @@ -56,14 +56,28 @@ def test_recognizes_built_in_generators(self, generator, expected): @pytest.mark.parametrize( ("class_name", "expected"), [ + ("AIMLAPIChatGenerator", "max_tokens"), ("AmazonBedrockChatGenerator", "maxTokens"), ("AnthropicChatGenerator", "max_tokens"), + ("AnthropicFoundryChatGenerator", "max_tokens"), + ("AnthropicVertexChatGenerator", "max_tokens"), + ("CohereChatGenerator", "max_tokens"), + ("CometAPIChatGenerator", "max_tokens"), + ("EdenAIChatGenerator", "max_tokens"), ("GoogleAIGeminiChatGenerator", "max_output_tokens"), ("GoogleGenAIChatGenerator", "max_output_tokens"), ("HuggingFaceAPIChatGenerator", "max_tokens"), ("LiteLLMChatGenerator", "max_tokens"), ("LlamaCppChatGenerator", "max_tokens"), + ("LlamaStackChatGenerator", "max_tokens"), + ("MistralChatGenerator", "max_tokens"), + ("NvidiaChatGenerator", "max_tokens"), ("OllamaChatGenerator", "num_predict"), + ("OpenRouterChatGenerator", "max_tokens"), + ("OrcaRouterChatGenerator", "max_tokens"), + ("PerplexityChatGenerator", "max_output_tokens"), + ("STACKITChatGenerator", "max_tokens"), + ("TogetherAIChatGenerator", "max_tokens"), ("TransformersChatGenerator", "max_new_tokens"), ("VertexAIGeminiChatGenerator", "max_output_tokens"), ("VLLMChatGenerator", "max_tokens"), @@ -73,53 +87,11 @@ def test_recognizes_built_in_generators(self, generator, expected): def test_recognizes_integration_generators(self, class_name, expected): assert _generator_output_token_limit_key(chat_generator=integration_generator(class_name)) == expected - def test_a_subclass_resolves_through_the_generator_it_inherits_from(self): - # This is what covers the OpenAI-compatible integrations, such as Mistral, OpenRouter, and TogetherAI, without - # naming any of them. + def test_an_unlisted_subclass_is_not_recognized(self): class ProviderChatGenerator(OpenAIChatGenerator): pass - assert _generator_output_token_limit_key(chat_generator=ProviderChatGenerator()) == "max_completion_tokens" - - def test_the_most_derived_entry_wins(self): - # No shipped generator depends on this today, since every subclass that is listed agrees with its base. It is - # pinned so that adding an entry for a subclass cannot be silently overridden by the base it inherits from. - base = type("OpenAIChatGenerator", (), {}) - derived = type("OpenAIResponsesChatGenerator", (base,), {}) - - assert _generator_output_token_limit_key(chat_generator=derived()) == "max_output_tokens" + assert _generator_output_token_limit_key(chat_generator=ProviderChatGenerator()) is None def test_unknown_generator_is_not_recognized(self): assert _generator_output_token_limit_key(chat_generator=integration_generator("MysteryChatGenerator")) is None - - -class TestRunKwargsWithOutputLimit: - def test_passes_the_limit_as_the_providers_runtime_setting(self): - kwargs = _run_kwargs_with_output_limit(chat_generator=OpenAIChatGenerator(), limit=100) - - assert kwargs == {"generation_kwargs": {"max_completion_tokens": 100}} - - def test_overrides_a_limit_the_generator_already_configures(self): - # A caller that reserves room for a reply of a given size needs that size honored, so the runtime value wins. - generator = OpenAIChatGenerator(generation_kwargs={"temperature": 0, "max_completion_tokens": 23}) - - kwargs = _run_kwargs_with_output_limit(chat_generator=generator, limit=100) - - assert kwargs == {"generation_kwargs": {"max_completion_tokens": 100}} - # Only this call is affected; the generator keeps its own settings, including the ones it is not asked about. - assert generator.generation_kwargs == {"temperature": 0, "max_completion_tokens": 23} - - def test_the_override_reaches_the_generator(self): - # Generators merge runtime kwargs over configured ones, which is what makes the override above take effect. - generator = OpenAIChatGenerator(generation_kwargs={"temperature": 0, "max_completion_tokens": 23}) - kwargs = _run_kwargs_with_output_limit(chat_generator=generator, limit=100) - - merged = {**generator.generation_kwargs, **kwargs["generation_kwargs"]} - - assert merged == {"temperature": 0, "max_completion_tokens": 100} - - def test_an_unknown_generator_receives_no_guessed_setting(self): - # The protocol only guarantees `run(messages)`, so even an empty `generation_kwargs` could raise a TypeError. - generator = integration_generator("MysteryChatGenerator") - - assert _run_kwargs_with_output_limit(chat_generator=generator, limit=100) == {}