Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
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