Skip to content
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,21 @@ class GoogleAIPromptExecutionSettings(PromptExecutionSettings):
temperature: Annotated[float | None, Field(ge=0.0, le=2.0)] = None
top_p: float | None = None
top_k: int | None = None
thinking_level: Literal["minimal", "low", "medium", "high"] | str | None = None
Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Personally I prefer the | str option for forward-compatibility, while still enabling completions from the Literal, but happy to defer to maintainer style on this one.


Comment thread
markmcd marked this conversation as resolved.
def prepare_settings_dict(self, **kwargs) -> dict[str, Any]:
"""Prepare the settings as a dictionary for sending to the AI service.

This method extracts the thinking_level from the settings dictionary
and maps it to the thinking_config expected by the Google AI SDK.
"""
settings_dict = super().prepare_settings_dict(**kwargs)

thinking_level = settings_dict.pop("thinking_level", None)
if thinking_level:
settings_dict["thinking_config"] = {"thinking_level": thinking_level}

return settings_dict


class GoogleAITextPromptExecutionSettings(GoogleAIPromptExecutionSettings):
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -131,3 +131,24 @@ def test_default_google_ai_embedding_prompt_execution_settings():
settings = GoogleAIEmbeddingPromptExecutionSettings()

assert settings.output_dimensionality is None


def test_google_ai_prompt_execution_settings_thinking_level():
settings = GoogleAIPromptExecutionSettings(thinking_level="high")
assert settings.thinking_level == "high"


def test_google_ai_prompt_execution_settings_thinking_level_mapping():
settings = GoogleAIPromptExecutionSettings(thinking_level="medium")
settings_dict = settings.prepare_settings_dict()

assert "thinking_level" not in settings_dict
assert settings_dict["thinking_config"] == {"thinking_level": "medium"}
Comment thread
markmcd marked this conversation as resolved.
Comment thread
markmcd marked this conversation as resolved.


def test_google_ai_prompt_execution_settings_no_thinking_level():
settings = GoogleAIPromptExecutionSettings(temperature=0.5)
settings_dict = settings.prepare_settings_dict()

assert "thinking_level" not in settings_dict
assert "thinking_level" not in settings_dict.get("thinking_config", {})
Loading