From fc8a9a34067bb1571bb02bf6b9dc308f89ba168e Mon Sep 17 00:00:00 2001
From: Jorge Yero Salazar <jyero96@gmail.com>
Date: Tue, 12 May 2026 10:25:01 -0500
Subject: [PATCH] Match litellm.completion supported model parameters with
 proxy model info (#27720)

* Use base_model for supported optional params

* Add test

* Formatting
---
 litellm/main.py                 |  6 ++++-
 tests/test_litellm/test_main.py | 39 +++++++++++++++++++++++++++++++++
 2 files changed, 44 insertions(+), 1 deletion(-)

diff --git a/litellm/main.py b/litellm/main.py
index 52a256fdd0..c3d1c2e05b 100644
--- a/litellm/main.py
+++ b/litellm/main.py
@@ -1528,7 +1528,11 @@ def completion(  # type: ignore # noqa: PLR0915
             "logit_bias": logit_bias,
             "user": user,
             # params to identify the model
-            "model": model,
+            "model": (
+                model_info.get("base_model")
+                if isinstance(model_info, dict) and model_info.get("base_model")
+                else model
+            ),
             "custom_llm_provider": custom_llm_provider,
             "response_format": response_format,
             "seed": seed,
diff --git a/tests/test_litellm/test_main.py b/tests/test_litellm/test_main.py
index 76336a91fc..b03579c2db 100644
--- a/tests/test_litellm/test_main.py
+++ b/tests/test_litellm/test_main.py
@@ -848,6 +848,45 @@ def test_gpt_5_4_responses_bridge_preserves_reasoning_summary_dict(
     }
 
 
+@pytest.mark.parametrize(
+    "model, model_info, expected_model_param",
+    [
+        ("gemini/gemini-3.1-pro", None, "gemini-3.1-pro"),
+        (
+            "gemini/gemini-3.1-pro",
+            {"base_model": "gemini-3.1-pro-preview"},
+            "gemini-3.1-pro-preview",
+        ),
+    ],
+)
+def test_completion_optional_params_base_model(
+    model: str,
+    model_info: dict | None,
+    expected_model_param: str,
+):
+    with patch("litellm.main.get_optional_params") as mock_get_optional_params:
+        mock_get_optional_params.return_value = MagicMock()
+
+        import litellm
+
+        kwargs = {
+            "model": model,
+            "messages": [{"role": "user", "content": "What is the capital of France?"}],
+            "api_key": "fake-key",
+            "mock_response": "Hey, how's it going?",
+        }
+        if model_info is not None:
+            kwargs["model_info"] = model_info
+
+        litellm.completion(**kwargs)
+
+        assert mock_get_optional_params.called is True
+        get_optional_params_model_param = mock_get_optional_params.call_args.kwargs[
+            "model"
+        ]
+        assert get_optional_params_model_param == expected_model_param
+
+
 @patch("litellm.completion_extras.responses_api_bridge.completion")
 def test_gpt_5_4_responses_bridge_merges_reasoning_summary_kwarg_without_tools(
     mock_responses_completion,