mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-01 07:03:02 +00:00
feat(model_info): add supports_continuation capability flag for Gemini long decode continuation
This commit is contained in:
parent
580b9266a1
commit
b04f3e2d1e
12 changed files with 98 additions and 0 deletions
|
|
@ -2373,6 +2373,7 @@ if TYPE_CHECKING:
|
|||
supports_web_search: Callable[..., bool]
|
||||
supports_tool_search: Callable[..., bool]
|
||||
supports_url_context: Callable[..., bool]
|
||||
supports_continuation: Callable[..., bool]
|
||||
supports_response_schema: Callable[..., bool]
|
||||
supports_parallel_function_calling: Callable[..., bool]
|
||||
supports_vision: Callable[..., bool]
|
||||
|
|
|
|||
|
|
@ -33,6 +33,7 @@ UTILS_NAMES: Final = (
|
|||
"supports_web_search",
|
||||
"supports_tool_search",
|
||||
"supports_url_context",
|
||||
"supports_continuation",
|
||||
"supports_response_schema",
|
||||
"supports_parallel_function_calling",
|
||||
"supports_vision",
|
||||
|
|
@ -488,6 +489,7 @@ _UTILS_IMPORT_MAP: Final = {
|
|||
"supports_web_search": (".utils", "supports_web_search"),
|
||||
"supports_tool_search": (".utils", "supports_tool_search"),
|
||||
"supports_url_context": (".utils", "supports_url_context"),
|
||||
"supports_continuation": (".utils", "supports_continuation"),
|
||||
"supports_response_schema": (".utils", "supports_response_schema"),
|
||||
"supports_parallel_function_calling": (
|
||||
".utils",
|
||||
|
|
|
|||
|
|
@ -76755,6 +76755,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,
|
||||
|
|
@ -76814,6 +76815,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,
|
||||
|
|
|
|||
|
|
@ -13269,6 +13269,8 @@ class PythonCore(DeprecatedRouterMembers):
|
|||
and model_info["supports_url_context"] is True
|
||||
):
|
||||
model_group_info.supports_url_context = True
|
||||
if model_info.get("supports_continuation") is True:
|
||||
model_group_info.supports_continuation = True
|
||||
|
||||
if model_info.get("supports_reasoning", None) is not None and model_info["supports_reasoning"] is True:
|
||||
model_group_info.supports_reasoning = True
|
||||
|
|
|
|||
|
|
@ -1061,6 +1061,7 @@ class ModelGroupInfo(LiteLLMBaseModel):
|
|||
supports_vision: bool = Field(default=False)
|
||||
supports_web_search: bool = Field(default=False)
|
||||
supports_url_context: bool = Field(default=False)
|
||||
supports_continuation: bool = Field(default=False)
|
||||
supports_reasoning: bool = Field(default=False)
|
||||
supports_function_calling: bool = Field(default=False)
|
||||
supports_fast_mode: bool = Field(default=False)
|
||||
|
|
|
|||
|
|
@ -227,6 +227,7 @@ class ProviderSpecificModelInfo(TypedDict, total=False):
|
|||
supports_tool_search: bool | None
|
||||
supports_mid_conversation_system: bool | None
|
||||
supports_url_context: bool | None
|
||||
supports_continuation: ReadOnly[bool | None]
|
||||
supports_none_reasoning_effort: bool | None
|
||||
supports_minimal_reasoning_effort: bool | None
|
||||
supports_low_reasoning_effort: bool | None
|
||||
|
|
|
|||
|
|
@ -3183,6 +3183,14 @@ def supports_url_context(model: str, custom_llm_provider: str | None = None) ->
|
|||
)
|
||||
|
||||
|
||||
def supports_continuation(model: str, custom_llm_provider: str | None = None) -> bool:
|
||||
return supports_factory(
|
||||
model=model,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
key="supports_continuation",
|
||||
)
|
||||
|
||||
|
||||
def supports_native_streaming(model: str, custom_llm_provider: str | None) -> bool:
|
||||
"""
|
||||
Check if the given model supports native streaming and return a boolean value.
|
||||
|
|
@ -7360,6 +7368,7 @@ def get_model_info_helper(
|
|||
),
|
||||
supports_web_search=_model_info.get("supports_web_search", None),
|
||||
supports_url_context=_model_info.get("supports_url_context", None),
|
||||
supports_continuation=_model_info.get("supports_continuation", None),
|
||||
supports_reasoning=_model_info.get("supports_reasoning", None),
|
||||
supports_adaptive_thinking=_model_info.get("supports_adaptive_thinking", None),
|
||||
supports_legacy_thinking=_model_info.get("supports_legacy_thinking", None),
|
||||
|
|
@ -7516,6 +7525,7 @@ def get_model_info(
|
|||
supports_pdf_input: Optional[bool]
|
||||
supports_web_search: Optional[bool]
|
||||
supports_url_context: Optional[bool]
|
||||
supports_continuation: Optional[bool]
|
||||
supports_reasoning: Optional[bool]
|
||||
Raises:
|
||||
Exception: If the model is not mapped yet.
|
||||
|
|
|
|||
|
|
@ -76755,6 +76755,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,
|
||||
|
|
@ -76814,6 +76815,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,
|
||||
|
|
|
|||
|
|
@ -845,6 +845,19 @@ def _search_free_requests_rows(path: Path) -> Mapping[str, dict[str, JsonValue]]
|
|||
)
|
||||
|
||||
|
||||
def _supports_continuation_rows(path: Path) -> frozenset[str]:
|
||||
prices: Final[dict[str, JsonValue]] = TypeAdapter(dict[str, JsonValue]).validate_json(path.read_bytes())
|
||||
return frozenset(
|
||||
key for key, row in prices.items() if isinstance(row, dict) and row.get("supports_continuation") is True
|
||||
)
|
||||
|
||||
|
||||
def test_gemini_4_argon_google_rows_carry_supports_continuation_in_both_cost_map_copies() -> None:
|
||||
main_rows: Final = _supports_continuation_rows(PRICES_PATH)
|
||||
assert {"gemini/gemini-4-argon", "vertex_ai/gemini-4-argon"} <= main_rows, sorted(main_rows)
|
||||
assert _supports_continuation_rows(BACKUP_PRICES_PATH) == main_rows
|
||||
|
||||
|
||||
@pytest.mark.parametrize("path", (PRICES_PATH, BACKUP_PRICES_PATH), ids=("main", "backup"))
|
||||
def test_search_free_requests_rows_are_grounding_rows_with_one_allowance_per_group(path: Path) -> None:
|
||||
rows = _search_free_requests_rows(path)
|
||||
|
|
|
|||
|
|
@ -1497,6 +1497,36 @@ def test_set_model_group_info(model_list, user_facing_model_group_name):
|
|||
assert resp.model_group == user_facing_model_group_name
|
||||
|
||||
|
||||
def test_set_model_group_info_reports_supports_continuation_when_any_deployment_has_it():
|
||||
router: Final = Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "long-decode",
|
||||
"litellm_params": {"model": "openai/gpt-4.1", "api_key": "fake"},
|
||||
},
|
||||
{
|
||||
"model_name": "long-decode",
|
||||
"litellm_params": {"model": "openai/gpt-4.1-mini", "api_key": "fake"},
|
||||
"model_info": {"supports_continuation": True},
|
||||
},
|
||||
{
|
||||
"model_name": "long-decode",
|
||||
"litellm_params": {"model": "openai/gpt-4.1-nano", "api_key": "fake"},
|
||||
},
|
||||
{
|
||||
"model_name": "plain",
|
||||
"litellm_params": {"model": "openai/gpt-4.1", "api_key": "fake"},
|
||||
},
|
||||
]
|
||||
)
|
||||
long_decode: Final = router._set_model_group_info(
|
||||
model_group="long-decode", user_facing_model_group_name="long-decode"
|
||||
)
|
||||
plain: Final = router._set_model_group_info(model_group="plain", user_facing_model_group_name="plain")
|
||||
assert long_decode is not None and long_decode.supports_continuation is True
|
||||
assert plain is not None and plain.supports_continuation is False
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_set_response_headers(model_list):
|
||||
"""Test if the 'set_response_headers' function is working correctly"""
|
||||
|
|
|
|||
|
|
@ -476,6 +476,34 @@ 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("continuation", [True, False])
|
||||
def test_supports_continuation_reads_its_own_cost_map_flag(
|
||||
local_model_cost_map: None, monkeypatch: pytest.MonkeyPatch, continuation: bool
|
||||
) -> None:
|
||||
model: Final = "gemini/continuation-flag-model"
|
||||
monkeypatch.setitem(
|
||||
litellm.model_cost,
|
||||
model,
|
||||
{
|
||||
"litellm_provider": "gemini",
|
||||
"mode": "chat",
|
||||
"supports_continuation": continuation,
|
||||
"supports_url_context": not continuation,
|
||||
},
|
||||
)
|
||||
assert litellm.supports_continuation(model) is continuation
|
||||
assert litellm.get_model_info(model).get("supports_continuation") is continuation
|
||||
|
||||
|
||||
def test_supports_continuation_is_unset_when_the_cost_map_omits_it(
|
||||
local_model_cost_map: None, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
model: Final = "gemini/no-continuation-flag-model"
|
||||
monkeypatch.setitem(litellm.model_cost, model, {"litellm_provider": "gemini", "mode": "chat"})
|
||||
assert litellm.supports_continuation(model) is False
|
||||
assert litellm.get_model_info(model).get("supports_continuation") is None
|
||||
|
||||
|
||||
def test_get_optional_params_image_gen():
|
||||
from litellm.llms.azure.image_generation import AzureGPTImageGenerationConfig
|
||||
|
||||
|
|
@ -1189,6 +1217,7 @@ _COST_MAP_ROW_SCHEMA: Final = {
|
|||
"supports_bedrock_runtime_chat_completions_inline_reasoning": {"type": "boolean"},
|
||||
"supports_bedrock_converse_tool_result_images": {"type": "boolean"},
|
||||
"supports_url_context": {"type": "boolean"},
|
||||
"supports_continuation": {"type": "boolean"},
|
||||
"supports_multimodal": {"type": "boolean"},
|
||||
"uses_embed_content": {"type": "boolean"},
|
||||
"supports_reasoning": {"type": "boolean"},
|
||||
|
|
|
|||
5
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
5
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
|
|
@ -38881,6 +38881,11 @@ export interface components {
|
|||
* @default false
|
||||
*/
|
||||
supports_background_chat: boolean;
|
||||
/**
|
||||
* Supports Continuation
|
||||
* @default false
|
||||
*/
|
||||
supports_continuation: boolean;
|
||||
/**
|
||||
* Supports Fast Mode
|
||||
* @default false
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue