Add thinking_config to GeminiModel (#1609)

Author: Kludex

pydantic_ai_slim/pydantic_ai/models/anthropic.py

@@ -90,7 +90,7 @@
 """
 
 
-class AnthropicModelSettings(ModelSettings):
+class AnthropicModelSettings(ModelSettings, total=False):
     """Settings used for an Anthropic model request.
 
     ALL FIELDS MUST BE `anthropic_` PREFIXED SO YOU CAN MERGE THEM WITH OTHER MODELS.

pydantic_ai_slim/pydantic_ai/models/cohere.py

@@ -78,7 +78,7 @@
 """
 
 
-class CohereModelSettings(ModelSettings):
+class CohereModelSettings(ModelSettings, total=False):
     """Settings used for a Cohere model request.
 
     ALL FIELDS MUST BE `cohere_` PREFIXED SO YOU CAN MERGE THEM WITH OTHER MODELS.

pydantic_ai_slim/pydantic_ai/models/gemini.py

@@ -73,14 +73,26 @@
 """
 
 
-class GeminiModelSettings(ModelSettings):
+class GeminiModelSettings(ModelSettings, total=False):
     """Settings used for a Gemini model request.
 
     ALL FIELDS MUST BE `gemini_` PREFIXED SO YOU CAN MERGE THEM WITH OTHER MODELS.
     """
 
     gemini_safety_settings: list[GeminiSafetySettings]
 
+    gemini_thinking_config: ThinkingConfig
+    """Thinking is "on" by default in both the API and AI Studio.
+
+    Being on by default doesn't mean the model will send back thoughts. For that, you would need to set `include_thoughts`
+    to `True`, but since end of January 2025, `thoughts` are not returned anymore, and are only displayed in the Google
+    AI Studio. See https://discuss.ai.google.dev/t/thoughts-are-missing-cot-not-included-anymore/63653 for more details.
+
+    If you want to avoid the model spending any tokens on thinking, you can set `thinking_budget` to `0`.
+
+    See more about it on <https://ai.google.dev/gemini-api/docs/thinking>.
+    """
+
 
 @dataclass(init=False)
 class GeminiModel(Model):
@@ -223,7 +235,9 @@ async def _make_request(
                 generation_config['presence_penalty'] = presence_penalty
             if (frequency_penalty := model_settings.get('frequency_penalty')) is not None:
                 generation_config['frequency_penalty'] = frequency_penalty
-            if (gemini_safety_settings := model_settings.get('gemini_safety_settings')) != []:
+            if (thinkingConfig := model_settings.get('gemini_thinking_config')) is not None:
+                generation_config['thinking_config'] = thinkingConfig  # pragma: no cover
+            if (gemini_safety_settings := model_settings.get('gemini_safety_settings')) is not None:
                 request_data['safetySettings'] = gemini_safety_settings
         if generation_config:
             request_data['generationConfig'] = generation_config
@@ -497,6 +511,16 @@ class GeminiSafetySettings(TypedDict):
     """
 
 
+class ThinkingConfig(TypedDict, total=False):
+    """The thinking features configuration."""
+
+    include_thoughts: Annotated[bool, pydantic.Field(alias='includeThoughts')]
+    """Indicates whether to include thoughts in the response. If true, thoughts are returned only if the model supports thought and thoughts are available."""
+
+    thinking_budget: Annotated[int, pydantic.Field(alias='thinkingBudget')]
+    """Indicates the thinking budget in tokens."""
+
+
 class _GeminiGenerationConfig(TypedDict, total=False):
     """Schema for an API request to the Gemini API.
 
@@ -511,6 +535,7 @@ class _GeminiGenerationConfig(TypedDict, total=False):
     presence_penalty: float
     frequency_penalty: float
     stop_sequences: list[str]
+    thinking_config: ThinkingConfig
 
 
 class _GeminiContent(TypedDict):

pydantic_ai_slim/pydantic_ai/models/groq.py

@@ -82,7 +82,7 @@
 """
 
 
-class GroqModelSettings(ModelSettings):
+class GroqModelSettings(ModelSettings, total=False):
     """Settings used for a Groq model request.
 
     ALL FIELDS MUST BE `groq_` PREFIXED SO YOU CAN MERGE THEM WITH OTHER MODELS.

pydantic_ai_slim/pydantic_ai/models/mistral.py

@@ -91,7 +91,7 @@
 """
 
 
-class MistralModelSettings(ModelSettings):
+class MistralModelSettings(ModelSettings, total=False):
     """Settings used for a Mistral model request.
 
     ALL FIELDS MUST BE `mistral_` PREFIXED SO YOU CAN MERGE THEM WITH OTHER MODELS.