fix: omit None config values from Gemini GenerateContentConfig (#5528)
This commit is contained in:
+9
-6
@@ -157,12 +157,15 @@ class GeminiLLM(LLMBase):
|
||||
# Extract system instruction and reformat messages
|
||||
system_instruction, contents = self._reformat_messages(messages)
|
||||
|
||||
# Prepare generation config
|
||||
config_params = {
|
||||
"temperature": self.config.temperature,
|
||||
"max_output_tokens": self.config.max_tokens,
|
||||
"top_p": self.config.top_p,
|
||||
}
|
||||
# Prepare generation config — only include non-None values so the
|
||||
# Gemini SDK uses its own defaults instead of rejecting None.
|
||||
config_params = {}
|
||||
if self.config.temperature is not None:
|
||||
config_params["temperature"] = self.config.temperature
|
||||
if self.config.max_tokens is not None:
|
||||
config_params["max_output_tokens"] = self.config.max_tokens
|
||||
if self.config.top_p is not None:
|
||||
config_params["top_p"] = self.config.top_p
|
||||
|
||||
# Add system instruction to config if present
|
||||
if system_instruction:
|
||||
|
||||
@@ -187,3 +187,45 @@ def test_parse_response_empty_parts_with_tools(mock_gemini_client: Mock):
|
||||
|
||||
result = llm._parse_response(mock_response, tools=[{"function": {"name": "test"}}])
|
||||
assert result == {"content": None, "tool_calls": []}
|
||||
|
||||
|
||||
def test_none_config_values_omitted_from_generation_config(mock_gemini_client: Mock):
|
||||
"""When temperature/max_tokens/top_p are None, they must not be passed
|
||||
to GenerateContentConfig (verified via model_fields_set)."""
|
||||
config = BaseLlmConfig(model="gemini-2.0-flash", temperature=None, max_tokens=None, top_p=None)
|
||||
llm = GeminiLLM(config)
|
||||
|
||||
mock_part = Mock(text="ok")
|
||||
mock_content = Mock(parts=[mock_part])
|
||||
mock_candidate = Mock(content=mock_content)
|
||||
mock_response = Mock(candidates=[mock_candidate])
|
||||
mock_gemini_client.models.generate_content.return_value = mock_response
|
||||
|
||||
llm.generate_response([{"role": "user", "content": "hi"}])
|
||||
|
||||
config_arg = mock_gemini_client.models.generate_content.call_args.kwargs["config"]
|
||||
assert "temperature" not in config_arg.model_fields_set
|
||||
assert "max_output_tokens" not in config_arg.model_fields_set
|
||||
assert "top_p" not in config_arg.model_fields_set
|
||||
|
||||
|
||||
def test_explicit_config_values_passed_to_generation_config(mock_gemini_client: Mock):
|
||||
"""When temperature/max_tokens/top_p are explicitly set, they must appear."""
|
||||
config = BaseLlmConfig(model="gemini-2.0-flash", temperature=0.5, max_tokens=200, top_p=0.9)
|
||||
llm = GeminiLLM(config)
|
||||
|
||||
mock_part = Mock(text="ok")
|
||||
mock_content = Mock(parts=[mock_part])
|
||||
mock_candidate = Mock(content=mock_content)
|
||||
mock_response = Mock(candidates=[mock_candidate])
|
||||
mock_gemini_client.models.generate_content.return_value = mock_response
|
||||
|
||||
llm.generate_response([{"role": "user", "content": "hi"}])
|
||||
|
||||
config_arg = mock_gemini_client.models.generate_content.call_args.kwargs["config"]
|
||||
assert "temperature" in config_arg.model_fields_set
|
||||
assert config_arg.temperature == 0.5
|
||||
assert "max_output_tokens" in config_arg.model_fields_set
|
||||
assert config_arg.max_output_tokens == 200
|
||||
assert "top_p" in config_arg.model_fields_set
|
||||
assert config_arg.top_p == 0.9
|
||||
|
||||
Reference in New Issue
Block a user