diff --git a/haystack/components/generators/chat/utils.py b/haystack/components/generators/chat/utils.py new file mode 100644 index 00000000000..618816bb47a --- /dev/null +++ b/haystack/components/generators/chat/utils.py @@ -0,0 +1,53 @@ +# SPDX-FileCopyrightText: 2022-present deepset GmbH +# +# SPDX-License-Identifier: Apache-2.0 + +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 +# `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", + "AzureOpenAIChatGenerator": "max_completion_tokens", + "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", + "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. + """ + 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 new file mode 100644 index 00000000000..97ab0cdd00e --- /dev/null +++ b/test/components/generators/chat/test_utils.py @@ -0,0 +1,97 @@ +# 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 +from haystack.dataclasses import ChatMessage + + +def integration_generator(class_name: str) -> ChatGenerator: + """ + Return an object named `class_name` that satisfies `ChatGenerator`. + + 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]: + 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"), + [ + ("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"), + ("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_an_unlisted_subclass_is_not_recognized(self): + class ProviderChatGenerator(OpenAIChatGenerator): + pass + + 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