Update src/rotator_library/providers/gemini_provider.py

src/rotator_library/providers/gemini_provider.py

@@ -14,6 +14,7 @@ if not lib_logger.handlers:
 class GeminiProvider(ProviderInterface):
     """
     Provider implementation for the Google Gemini API.
+    # TEST EDIT
     """
     async def get_models(self, api_key: str, client: httpx.AsyncClient) -> List[str]:
         """
@@ -88,54 +89,103 @@ class GeminiProvider(ProviderInterface):
 
         return gemini_settings
 
+    def _is_gemini_3(self, model: str) -> bool:
+        """
+        Helper to identify Gemini 3 models.
+        """
+        model_name = model.split("/")[-1].replace(":thinking", "")
+        return model_name.startswith("gemini-3-")
+
     def handle_thinking_parameter(self, payload: Dict[str, Any], model: str):
         """
-        Handles reasoning parameters for Gemini models, with three distinct paths:
-        1. Applies a non-standard, high-value token budget if 'custom_reasoning_budget' is true.
-        2. Leaves the 'reasoning_effort' parameter alone for LiteLLM to handle if it's present
-           without the custom flag.
-        3. Applies a default 'thinking' value for specific models if no other reasoning
-           parameters are provided, ensuring they 'think' by default.
+        Handles reasoning parameters for Gemini models (Gemini 2.5 and Gemini 3).
+        Maps 'reasoning_effort' to the appropriate Gemini-specific format for LiteLLM.
         """
         # Set default temperature to 1 if not provided
         if "temperature" not in payload:
             payload["temperature"] = 1
 
-        custom_reasoning_budget = payload.get("custom_reasoning_budget", False)
-        reasoning_effort = payload.get("reasoning_effort")
+        reasoning_effort = payload.pop("reasoning_effort", None)
+        custom_reasoning_budget = payload.pop("custom_reasoning_budget", False)
 
         # If 'thinking' is already explicitly set, do nothing to avoid overriding it.
         if "thinking" in payload:
             return
 
-        # Path 1: Custom budget is explicitly requested.
-        if custom_reasoning_budget:
-            # Case 1a: Both params are present, so we can apply the custom budget.
-            if reasoning_effort:
-                if "gemini-2.5-pro" in model:
-                    budgets = {"low": 8192, "medium": 16384, "high": 32768}
-                elif "gemini-2.5-flash" in model:
-                    budgets = {"low": 6144, "medium": 12288, "high": 24576}
-                else: # Fallback for other models if the custom flag is still used
-                    budgets = {"low": 1024, "medium": 2048, "high": 4096}
-                
-                budget = budgets.get(reasoning_effort)
-                if budget is not None:
-                    payload["thinking"] = {"type": "enabled", "budget_tokens": budget}
-                elif reasoning_effort == "disable":
-                    payload["thinking"] = {"type": "enabled", "budget_tokens": 0}
-                
-                # Clean up the handled 'reasoning_effort' parameter.
-                payload.pop("reasoning_effort", None)
-
-            # Case 1b: In all cases where the custom flag was present, remove it
-            # as it's not a standard LiteLLM parameter.
-            payload.pop("custom_reasoning_budget", None)
+        is_gemini_25 = "gemini-2.5" in model
+        is_gemini_3 = self._is_gemini_3(model)
+        is_gemini_3_flash = "gemini-3-flash" in model
+
+        if not (is_gemini_25 or is_gemini_3):
+            return
+
+        # Normalize and validate effort
+        if reasoning_effort is None:
+            effort = "auto"
+        elif isinstance(reasoning_effort, str):
+            effort = reasoning_effort.strip().lower() or "auto"
+        else:
+            effort = "auto"
+
+        valid_efforts = {
+            "auto", "disable", "off", "none", "minimal", "low", 
+            "low_medium", "medium", "medium_high", "high"
+        }
+        if effort not in valid_efforts:
+            effort = "auto"
+
+        # Path 1: Gemini 3 Logic (Native Level Strings)
+        if is_gemini_3:
+            if is_gemini_3_flash:
+                # minimal/low/medium/high
+                if effort in ("disable", "off", "none"):
+                    level = "minimal"
+                elif effort in ("minimal", "low"):
+                    level = "low"
+                elif effort in ("low_medium", "medium"):
+                    level = "medium"
+                else: # auto, medium_high, high
+                    level = "high"
+            else: # Gemini 3 Pro (usually only low/high)
+                if effort in ("disable", "off", "none", "minimal", "low", "low_medium"):
+                    level = "low"
+                else: # auto, medium, medium_high, high
+                    level = "high"
+            
+            payload["thinking"] = {"type": "enabled", "thinking_level": level}
+            return
+
+        # Path 2: Gemini 2.5 Logic (Token Budget)
+        if effort in ("disable", "off", "none"):
+            payload["thinking"] = {"type": "enabled", "budget_tokens": 0}
+            return
+
+        # If auto is used without custom budget, LiteLLM handles with -1
+        if effort == "auto" and not custom_reasoning_budget:
+            payload["thinking"] = {"type": "enabled", "budget_tokens": -1}
             return
 
-        # Path 2: No custom budget. Now check for standard or default behavior.
-        # If 'reasoning_effort' is present, we do nothing, allowing LiteLLM to handle it.
-        # If 'reasoning_effort' is NOT present, then we apply the default thinking behavior.
-        if not reasoning_effort:
-            if "gemini-2.5-pro" in model or "gemini-2.5-flash" in model:
-                payload["thinking"] = {"type": "enabled", "budget_tokens": -1}
+        # Model-specific budgets
+        if "gemini-2.5-flash" in model:
+            budgets = {
+                "minimal": 3072,
+                "low": 6144,
+                "low_medium": 9216,
+                "medium": 12288,
+                "medium_high": 18432,
+                "high": 24576,
+                "auto": 12288 if custom_reasoning_budget else -1
+            }
+        else: # Gemini 2.5 Pro
+            budgets = {
+                "minimal": 4096,
+                "low": 8192,
+                "low_medium": 12288,
+                "medium": 16384,
+                "medium_high": 24576,
+                "high": 32768,
+                "auto": 16384 if custom_reasoning_budget else -1
+            }
+        
+        budget = budgets.get(effort, budgets["auto"])
+        payload["thinking"] = {"type": "enabled", "budget_tokens": budget}