diff --git a/litellm/router_utils/pre_call_checks/continuation_prefill_check.py b/litellm/router_utils/pre_call_checks/continuation_prefill_check.py index f9750b4933b..ccd454bc207 100644 --- a/litellm/router_utils/pre_call_checks/continuation_prefill_check.py +++ b/litellm/router_utils/pre_call_checks/continuation_prefill_check.py @@ -35,6 +35,19 @@ _STR_KEYED_DICT_ADAPTER: Final = TypeAdapter(dict[str, object]) def _deployment_supports_prefill(deployment: object) -> bool: try: deployment_map: Final = _STR_KEYED_DICT_ADAPTER.validate_python(deployment) + except ValidationError: + return False + # A per-deployment model_info override wins, so a model that is not in the cost + # map (or is registered generically) can still opt in or out explicitly with + # `model_info: {"supports_assistant_prefill": true|false}`. + try: + model_info: Final = _STR_KEYED_DICT_ADAPTER.validate_python(deployment_map.get("model_info")) + declared: Final = model_info.get("supports_assistant_prefill") + if isinstance(declared, bool): + return declared + except ValidationError: + pass + try: litellm_params: Final = _STR_KEYED_DICT_ADAPTER.validate_python(deployment_map.get("litellm_params")) except ValidationError: return False diff --git a/tests/test_litellm/router_utils/pre_call_checks/test_continuation_prefill_check.py b/tests/test_litellm/router_utils/pre_call_checks/test_continuation_prefill_check.py index 2e1c41bf6ab..1d6b667464e 100644 --- a/tests/test_litellm/router_utils/pre_call_checks/test_continuation_prefill_check.py +++ b/tests/test_litellm/router_utils/pre_call_checks/test_continuation_prefill_check.py @@ -20,6 +20,25 @@ def test_deployment_supports_prefill_reads_capability(): assert _deployment_supports_prefill(_deployment(NON_PREFILL_MODEL, "b")) is False +def test_deployment_model_info_override_wins_over_cost_map(): + # model_info True opts in a model that is not in the cost map + assert ( + _deployment_supports_prefill( + {"litellm_params": {"model": "vendor/custom-model"}, "model_info": {"supports_assistant_prefill": True}} + ) + is True + ) + # model_info False opts out a model the cost map would otherwise allow + assert ( + _deployment_supports_prefill( + {"litellm_params": {"model": PREFILL_MODEL}, "model_info": {"supports_assistant_prefill": False}} + ) + is False + ) + # model_info without the key falls through to the cost map + assert _deployment_supports_prefill(_deployment(PREFILL_MODEL, "z")) is True + + def test_deployment_supports_prefill_rejects_malformed_deployments(): assert _deployment_supports_prefill({}) is False assert _deployment_supports_prefill({"litellm_params": {}}) is False