Use Gemini's reasoning features

src/prompto/apis/gemini/gemini.py

@@ -7,6 +7,7 @@
     HarmBlockThreshold,
     HarmCategory,
     SafetySetting,
+    ThinkingConfig,
 )
 
 from prompto.apis.base import AsyncAPI
@@ -338,10 +339,26 @@ async def _obtain_model_inputs(
                 f"parameters must be a dictionary, not {type(generation_config_params)}"
             )
 
+        # Derive the required ThinkingConfig from the parameters
+        # TBC - how do we get these values from the prompt_dict?
+
+        # Placeholder values for now
+        include_thoughts = False
+        thinking_budget = 9999
+
+        if include_thoughts is None and thinking_budget is None:
+            thinking_config = None
+        else:
+            thinking_config = ThinkingConfig(
+                include_thoughts=include_thoughts,
+                thinking_budget=thinking_budget,
+            )
+
         gen_content_config = GenerateContentConfig(
             **generation_config_params,
             safety_settings=safety_settings,
             system_instruction=system_instruction,
+            thinking_config=thinking_config,
         )
 
         return prompt, model_name, client, gen_content_config, None

tests/apis/gemini/test_gemini.py

@@ -8,6 +8,7 @@
     HarmBlockThreshold,
     HarmCategory,
     SafetySetting,
+    ThinkingConfig,
 )
 
 from prompto.apis.gemini import GeminiAPI
@@ -475,6 +476,74 @@ async def test_gemini_obtain_model_inputs(temporary_data_folders, monkeypatch):
         )
 
 
+@pytest.mark.xfail(
+    reason="This cannot work until we agree on the schema for the thinking config parameters"
+)
+@pytest.mark.asyncio
+async def test_gemini_obtain_model_inputs_thinking_config(
+    temporary_data_folders, monkeypatch
+):
+    settings = Settings(data_folder="data")
+    log_file = "log.txt"
+    monkeypatch.setenv("GEMINI_API_KEY", "DUMMY")
+    gemini_api = GeminiAPI(settings=settings, log_file=log_file)
+
+    # These test cases are very similar to the test above, so we will not repeat assertions for all
+    # attributes - only those that are relevant to the thinking config
+
+    # Case 1: test with *NO* thinking config parameters provided
+    test_case = await gemini_api._obtain_model_inputs(
+        {
+            "id": "gemini_id",
+            "api": "gemini",
+            "model_name": "gemini_model_name",
+            "prompt": "test prompt",
+            "parameters": {"temperature": 1, "max_output_tokens": 100},
+        }
+    )
+    assert isinstance(test_case, tuple)
+    assert isinstance(test_case[3], GenerateContentConfig)
+    assert test_case[3].thinking_config is None
+
+    # Case 2: test with thinking config parameters provided
+    # There are two possible ways we could provide the thinking config parameters.
+    # We need to select from one of these to options:
+    dummy_prompt_dicts = [
+        # Either within the parameters dictionary
+        {
+            "id": "gemini_id",
+            "api": "gemini",
+            "model_name": "gemini_model_name",
+            "prompt": "test prompt",
+            "parameters": {
+                "temperature": 1,
+                "max_output_tokens": 100,
+                "thinking_budget": 1234,
+                "include_thoughts": True,
+            },
+        },
+        # OR as top-level keys within the prompt dictionary
+        {
+            "id": "gemini_id",
+            "api": "gemini",
+            "model_name": "gemini_model_name",
+            "prompt": "test prompt",
+            "thinking_budget": 1234,
+            "include_thoughts": True,
+            "parameters": {"temperature": 1, "max_output_tokens": 100},
+        },
+    ]
+
+    for dummy_prompt_dict in dummy_prompt_dicts:
+        test_case = await gemini_api._obtain_model_inputs(dummy_prompt_dict)
+
+        assert isinstance(test_case, tuple)
+        assert isinstance(test_case[3], GenerateContentConfig)
+        assert isinstance(test_case[3].thinking_config, ThinkingConfig)
+        assert test_case[3].thinking_config.thinking_budget == 1234
+        assert test_case[3].thinking_config.include_thoughts is True
+
+
 @pytest.mark.asyncio
 async def test_gemini_obtain_model_inputs_safety_filters(
     temporary_data_folders, monkeypatch