From 736fd9cba6166ce3287d87e39c31ce5d858f6e8a Mon Sep 17 00:00:00 2001 From: mrinal-berri Date: Fri, 30 Oct 2026 15:47:01 -0700 Subject: [PATCH] feat(gemini): expose long decode continuation capability --- litellm/__init__.py | 1 + litellm/_lazy_imports_registry.py | 2 ++ ...odel_prices_and_context_window_backup.json | 2 ++ litellm/types/utils.py | 1 + litellm/utils.py | 7 +++++ model_prices_and_context_window.json | 2 ++ ...test_proxy_server_endpoints_and_startup.py | 29 +++++++++++++++++++ .../test_model_prices_and_context_window.py | 10 +++++++ tests/unit/test_utils.py | 15 ++++++++++ 9 files changed, 69 insertions(+) diff --git a/litellm/__init__.py b/litellm/__init__.py index f64531f86fb..1679c45dfc7 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -2379,6 +2379,7 @@ if TYPE_CHECKING: supports_audio_input: Callable[..., bool] supports_video_input: Callable[..., bool] supports_box_questions: Callable[..., bool] + supports_continuation: Callable[..., bool] supports_audio_output: Callable[..., bool] supports_system_messages: Callable[..., bool] supports_reasoning: Callable[..., bool] diff --git a/litellm/_lazy_imports_registry.py b/litellm/_lazy_imports_registry.py index c1f484f8d64..ff8df0ac114 100644 --- a/litellm/_lazy_imports_registry.py +++ b/litellm/_lazy_imports_registry.py @@ -39,6 +39,7 @@ UTILS_NAMES: Final = ( "supports_audio_input", "supports_video_input", "supports_box_questions", + "supports_continuation", "supports_audio_output", "supports_system_messages", "supports_reasoning", @@ -497,6 +498,7 @@ _UTILS_IMPORT_MAP: Final = { "supports_audio_input": (".utils", "supports_audio_input"), "supports_video_input": (".utils", "supports_video_input"), "supports_box_questions": (".utils", "supports_box_questions"), + "supports_continuation": (".utils", "supports_continuation"), "supports_audio_output": (".utils", "supports_audio_output"), "supports_system_messages": (".utils", "supports_system_messages"), "supports_reasoning": (".utils", "supports_reasoning"), diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 843ab574da4..7c39061f8b8 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -76646,6 +76646,7 @@ "supports_audio_input": true, "supports_audio_output": false, "supports_computer_use": true, + "supports_continuation": true, "supports_function_calling": true, "supports_minimal_reasoning_effort": false, "supports_native_streaming": true, @@ -76705,6 +76706,7 @@ ], "supports_audio_input": true, "supports_computer_use": true, + "supports_continuation": true, "supports_function_calling": true, "supports_minimal_reasoning_effort": false, "supports_native_streaming": true, diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 0c3440e88f7..4757dc0f881 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -210,6 +210,7 @@ class ProviderSpecificModelInfo(TypedDict, total=False): supports_audio_input: bool | None supports_video_input: ReadOnly[bool | None] supports_box_questions: ReadOnly[bool | None] + supports_continuation: ReadOnly[bool | None] supports_embedding_image_input: bool | None supports_audio_output: bool | None supports_pdf_input: bool | None diff --git a/litellm/utils.py b/litellm/utils.py index c2633d7eb81..78222e0ecb2 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -3458,6 +3458,11 @@ def is_explicitly_disabled_factory(model: str, custom_llm_provider: str | None, return False +def supports_continuation(model: str, custom_llm_provider: str | None = None) -> bool: + """Check if the model declares support for long-decode continuation.""" + return supports_factory(model=model, custom_llm_provider=custom_llm_provider, key="supports_continuation") + + def supports_audio_input(model: str, custom_llm_provider: str | None = None) -> bool: """Check if a given model supports audio input in a chat completion call""" return supports_factory(model=model, custom_llm_provider=custom_llm_provider, key="supports_audio_input") @@ -7343,6 +7348,7 @@ def get_model_info_helper( supports_audio_input=_model_info.get("supports_audio_input", None), supports_video_input=_model_info.get("supports_video_input", None), supports_box_questions=_model_info.get("supports_box_questions", None), + supports_continuation=_model_info.get("supports_continuation", None), supports_audio_output=_model_info.get("supports_audio_output", None), supports_pdf_input=_model_info.get("supports_pdf_input", None), supports_embedding_image_input=_model_info.get("supports_embedding_image_input", None), @@ -7505,6 +7511,7 @@ def get_model_info( supports_audio_input: Optional[bool] supports_video_input: Optional[bool] supports_box_questions: Optional[bool] + supports_continuation: Optional[bool] supports_audio_output: Optional[bool] supports_pdf_input: Optional[bool] supports_web_search: Optional[bool] diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 843ab574da4..7c39061f8b8 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -76646,6 +76646,7 @@ "supports_audio_input": true, "supports_audio_output": false, "supports_computer_use": true, + "supports_continuation": true, "supports_function_calling": true, "supports_minimal_reasoning_effort": false, "supports_native_streaming": true, @@ -76705,6 +76706,7 @@ ], "supports_audio_input": true, "supports_computer_use": true, + "supports_continuation": true, "supports_function_calling": true, "supports_minimal_reasoning_effort": false, "supports_native_streaming": true, diff --git a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py index bc3a228321a..5245ca07ade 100644 --- a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py +++ b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py @@ -6226,6 +6226,35 @@ def test_model_info_v1_list_skips_fastapi_jsonable_encoder(monkeypatch): assert encoder_spy.call_count == 0 + +@pytest.mark.parametrize("route", ("/model/info", "/v1/model/info")) +def test_model_info_exposes_continuation_and_deployment_override( + route: str, local_model_cost_map: None, monkeypatch: pytest.MonkeyPatch +) -> None: + deployments: Final = [ + {"model_name": "default", "litellm_params": {"model": "gemini/gemini-4-argon"}}, + { + "model_name": "override", + "litellm_params": {"model": "vertex_ai/gemini-4-argon"}, + "model_info": {"supports_continuation": False}, + }, + ] + router: Final = litellm.Router(model_list=deployments) + monkeypatch.setattr(proxy_server_module, "llm_router", router) + monkeypatch.setattr(proxy_server_module, "llm_model_list", router.model_list) + monkeypatch.setattr(proxy_server_module, "prisma_client", None) + monkeypatch.setattr(proxy_server_module, "general_settings", {"infer_model_from_keys": False}) + monkeypatch.setattr(proxy_server_module, "user_model", None) + monkeypatch.setitem( + app.dependency_overrides, + user_api_key_auth, + lambda: UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, api_key=MASTER_KEY, models=[], team_models=[]), + ) + response: Final = TestClient(app).get(route) + assert response.status_code == 200 + flags: Final = {row["model_name"]: row["model_info"]["supports_continuation"] for row in response.json()["data"]} + assert flags == {"default": True, "override": False} + def test_model_info_v1_cli_model_returns_single_deployment_as_json(monkeypatch): """ A proxy started with `litellm --model ` answers /model/info with one deployment diff --git a/tests/unit/test_model_prices_and_context_window.py b/tests/unit/test_model_prices_and_context_window.py index 8ca678a94ad..ac66fbe66cd 100644 --- a/tests/unit/test_model_prices_and_context_window.py +++ b/tests/unit/test_model_prices_and_context_window.py @@ -44,6 +44,16 @@ def test_azure_audio_snapshot_capabilities_match_declared_modalities(model: str, assert (litellm.supports_audio_input(model), litellm.supports_audio_output(model)) == expected + +@pytest.mark.parametrize("model", ("gemini/gemini-4-argon", "vertex_ai/gemini-4-argon")) +def test_gemini_long_decode_capability_is_exposed(model: str, local_model_cost_map: None) -> None: + # Google long decode documentation, verified 2026-10-30: https://ai.google.dev/gemini-api/docs/generate-content/long-decode + catalog: Final = TypeAdapter(dict[str, JsonValue]).validate_json(PRICES_PATH.read_bytes()) + backup: Final = TypeAdapter(dict[str, JsonValue]).validate_json(BACKUP_PRICES_PATH.read_bytes()) + assert catalog[model] == backup[model] + assert litellm.get_model_info(model)["supports_continuation"] is True + assert litellm.supports_continuation(model) is True + class _CopilotEndpointRow(BaseModel): model_config = ConfigDict(frozen=True, extra="ignore") diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index 5cf872f4ea7..7b7ff858189 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -476,6 +476,20 @@ def test_supports_video_input_reads_its_own_cost_map_flag( assert litellm.get_model_info(model)["supports_video_input"] is video_input + +@pytest.mark.parametrize("declared", (True, False, None)) +def test_continuation_capability_reads_declared_model_info( + local_model_cost_map: None, monkeypatch: pytest.MonkeyPatch, declared: bool | None +) -> None: + model: Final = "openai/continuation-capability-test" + monkeypatch.setitem( + litellm.model_cost, + model, + {"litellm_provider": "openai", "mode": "chat", "supports_continuation": declared}, + ) + assert litellm.get_model_info(model)["supports_continuation"] is declared + assert litellm.supports_continuation(model) is (declared is True) + def test_get_optional_params_image_gen(): from litellm.llms.azure.image_generation import AzureGPTImageGenerationConfig @@ -1175,6 +1189,7 @@ _COST_MAP_ROW_SCHEMA: Final = { "supports_tool_search": {"type": "boolean"}, "supports_video_input": {"type": "boolean"}, "supports_box_questions": {"type": "boolean"}, + "supports_continuation": {"type": "boolean"}, "supports_vision": {"type": "boolean"}, "supports_web_search": {"type": "boolean"}, "token_count_framing": {"type": "integer"},