From 784c783c3f2c5ee27819648fa1328d5b47d63c2f Mon Sep 17 00:00:00 2001 From: Stan Fateev Date: Thu, 30 Jul 2026 12:16:33 +0400 Subject: [PATCH] fix: reject multiple response candidates --- src/google/adk/models/apigee_llm.py | 2 + src/google/adk/models/base_llm.py | 15 ++++++++ src/google/adk/models/google_llm.py | 2 + src/google/adk/models/lite_llm.py | 2 + .../models/test_completions_http_client.py | 24 +++++++++++- tests/unittests/models/test_google_llm.py | 20 ++++++++++ tests/unittests/models/test_litellm.py | 37 +++++++++++++++++++ tests/unittests/models/test_llm_response.py | 23 ++++++++++++ 8 files changed, 123 insertions(+), 2 deletions(-) diff --git a/src/google/adk/models/apigee_llm.py b/src/google/adk/models/apigee_llm.py index 84d41f6af17..6897a048300 100644 --- a/src/google/adk/models/apigee_llm.py +++ b/src/google/adk/models/apigee_llm.py @@ -36,6 +36,7 @@ from typing_extensions import override from ..utils.env_utils import is_enterprise_mode_enabled +from .base_llm import _validate_candidate_count from .google_llm import Gemini from .llm_response import LlmResponse @@ -633,6 +634,7 @@ def _map_config_parameters( self, config: types.GenerateContentConfig, payload: dict[str, Any] ) -> None: """Maps configuration parameters to the payload.""" + _validate_candidate_count(config) if config.temperature is not None: payload['temperature'] = config.temperature if config.top_p is not None: diff --git a/src/google/adk/models/base_llm.py b/src/google/adk/models/base_llm.py index b2a770a3d17..d359c893580 100644 --- a/src/google/adk/models/base_llm.py +++ b/src/google/adk/models/base_llm.py @@ -29,6 +29,21 @@ from .llm_response import LlmResponse +def _validate_candidate_count( + config: types.GenerateContentConfig | None, +) -> None: + """Validates that the request fits ADK's single-candidate response model.""" + if ( + config + and config.candidate_count is not None + and config.candidate_count > 1 + ): + raise ValueError( + 'ADK supports only one response candidate; candidate_count must be 1' + ' or unset.' + ) + + class BaseLlm(BaseModel): """The BaseLLM class.""" diff --git a/src/google/adk/models/google_llm.py b/src/google/adk/models/google_llm.py index 2c240386ac7..a5c7de8f804 100644 --- a/src/google/adk/models/google_llm.py +++ b/src/google/adk/models/google_llm.py @@ -40,6 +40,7 @@ from ..utils.context_utils import Aclosing from ..utils.streaming_utils import StreamingResponseAggregator from ..utils.variant_utils import GoogleLLMVariant +from .base_llm import _validate_candidate_count from .base_llm import BaseLlm from .base_llm_connection import BaseLlmConnection from .gemini_llm_connection import GeminiLlmConnection @@ -194,6 +195,7 @@ async def generate_content_async( Yields: LlmResponse: The model response. """ + _validate_candidate_count(llm_request.config) await self._preprocess_request(llm_request) self._maybe_append_user_content(llm_request) diff --git a/src/google/adk/models/lite_llm.py b/src/google/adk/models/lite_llm.py index b5dd232bb9c..253dc02e29c 100644 --- a/src/google/adk/models/lite_llm.py +++ b/src/google/adk/models/lite_llm.py @@ -53,6 +53,7 @@ from typing_extensions import override from ..utils._google_client_headers import merge_tracking_headers +from .base_llm import _validate_candidate_count from .base_llm import BaseLlm from .llm_request import LlmRequest from .llm_response import LlmResponse @@ -2348,6 +2349,7 @@ async def _get_completion_inputs( The litellm inputs (message list, tool dictionary, response format, generation params, and tool_choice). """ + _validate_candidate_count(llm_request.config) _ensure_litellm_imported() # Determine provider for file handling diff --git a/tests/unittests/models/test_completions_http_client.py b/tests/unittests/models/test_completions_http_client.py index 5022862cfea..415842c968c 100644 --- a/tests/unittests/models/test_completions_http_client.py +++ b/tests/unittests/models/test_completions_http_client.py @@ -79,7 +79,7 @@ async def test_construct_payload_with_config(client, llm_request): frequency_penalty=0.5, presence_penalty=0.5, seed=42, - candidate_count=2, + candidate_count=1, response_mime_type='application/json', ) @@ -107,10 +107,30 @@ async def test_construct_payload_with_config(client, llm_request): assert payload['frequency_penalty'] == 0.5 assert payload['presence_penalty'] == 0.5 assert payload['seed'] == 42 - assert payload['n'] == 2 + assert payload['n'] == 1 assert payload['response_format'] == {'type': 'json_object'} +@pytest.mark.asyncio +async def test_construct_payload_rejects_multiple_candidates( + client, llm_request +): + llm_request.config = types.GenerateContentConfig(candidate_count=2) + + with mock.patch.object(httpx.AsyncClient, 'post') as mock_post: + with pytest.raises( + ValueError, match='supports only one response candidate' + ): + _ = [ + response + async for response in client.generate_content_async( + llm_request, stream=False + ) + ] + + mock_post.assert_not_called() + + @pytest.mark.asyncio async def test_construct_payload_with_tools(client, llm_request): tool = types.Tool( diff --git a/tests/unittests/models/test_google_llm.py b/tests/unittests/models/test_google_llm.py index a7ec360c031..7b3b8f07ebd 100644 --- a/tests/unittests/models/test_google_llm.py +++ b/tests/unittests/models/test_google_llm.py @@ -396,6 +396,26 @@ async def mock_coro(): mock_client.aio.models.generate_content.assert_called_once() +@pytest.mark.asyncio +async def test_generate_content_async_rejects_multiple_candidates( + gemini_llm, llm_request +): + llm_request.config.candidate_count = 2 + + with mock.patch.object(gemini_llm, "api_client") as mock_client: + with pytest.raises( + ValueError, match="supports only one response candidate" + ): + _ = [ + response + async for response in gemini_llm.generate_content_async( + llm_request, stream=False + ) + ] + + mock_client.aio.models.generate_content.assert_not_called() + + @pytest.mark.asyncio async def test_generate_content_async_stream(gemini_llm, llm_request): with mock.patch.object(gemini_llm, "api_client") as mock_client: diff --git a/tests/unittests/models/test_litellm.py b/tests/unittests/models/test_litellm.py index 899c453e862..90c0ef7cf45 100644 --- a/tests/unittests/models/test_litellm.py +++ b/tests/unittests/models/test_litellm.py @@ -2612,6 +2612,30 @@ def test_model_response_to_generate_content_response_reasoning_content(): assert response.content.parts[1].text == "Answer" +def test_model_response_to_generate_content_response_uses_first_choice(): + """Test LiteLLM conversion follows the single-candidate contract.""" + model_response = ModelResponse( + model="test-model", + choices=[ + { + "message": {"role": "assistant", "content": "First"}, + "finish_reason": "stop", + }, + { + "message": {"role": "assistant", "content": "Second"}, + "finish_reason": "stop", + }, + ], + ) + + response = _model_response_to_generate_content_response(model_response) + + assert len(model_response.choices) == 2 + assert [ + part.text for part in response.content.parts if part.text is not None + ] == ["First"] + + def test_message_to_generate_content_response_reasoning_field(): """Test that the 'reasoning' field is supported (LM Studio, vLLM).""" message = { @@ -4779,6 +4803,19 @@ async def test_get_completion_inputs_generation_params(): assert "stop_sequences" not in generation_params +@pytest.mark.asyncio +async def test_get_completion_inputs_rejects_multiple_candidates(): + req = LlmRequest( + contents=[ + types.Content(role="user", parts=[types.Part.from_text(text="hi")]), + ], + config=types.GenerateContentConfig(candidate_count=2), + ) + + with pytest.raises(ValueError, match="supports only one response candidate"): + await _get_completion_inputs(req, model="gpt-4o-mini") + + @pytest.mark.asyncio async def test_get_completion_inputs_empty_generation_params(): # Test that generation_params is None when no generation parameters are set diff --git a/tests/unittests/models/test_llm_response.py b/tests/unittests/models/test_llm_response.py index f4fbd6b8467..57a6e1934fc 100644 --- a/tests/unittests/models/test_llm_response.py +++ b/tests/unittests/models/test_llm_response.py @@ -64,6 +64,29 @@ def test_llm_response_create_without_logprobs(): assert response.content.parts[0].text == 'Response text' +def test_llm_response_create_uses_first_candidate(): + """Test LlmResponse.create() follows the single-candidate contract.""" + generate_content_response = types.GenerateContentResponse( + candidates=[ + types.Candidate( + content=types.Content(parts=[types.Part(text='First')]), + finish_reason=types.FinishReason.STOP, + ), + types.Candidate( + content=types.Content(parts=[types.Part(text='Second')]), + finish_reason=types.FinishReason.STOP, + ), + ] + ) + + response = LlmResponse.create(generate_content_response) + + assert len(generate_content_response.candidates) == 2 + assert [ + part.text for part in response.content.parts if part.text is not None + ] == ['First'] + + def test_llm_response_create_error_case_with_logprobs(): """Test LlmResponse.create() includes logprobs in error cases.""" avg_logprobs = -2.1