diff --git a/litellm/router.py b/litellm/router.py index 6a43a6a05c0..474c284c389 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -2857,6 +2857,17 @@ class Router: MID_STREAM_CONTINUATION_MARKER, ) + emitted_tokens: Final = int( + getattr(complete_response_object_usage, "completion_tokens", 0) or 0 + ) + reduced_ceilings: Final = self._continuation_output_ceilings(initial_kwargs, emitted_tokens) + if reduced_ceilings is None: + # The caller's output allowance is already spent; surface the + # error rather than grant a fresh allowance on this fallback hop. + if e.original_exception is not None: + raise e.original_exception from e + raise + initial_kwargs.update(reduced_ceilings) initial_kwargs["messages"] = self._build_completion_continuation_input( messages, e.generated_content ) @@ -3076,6 +3087,28 @@ class Router: return False return True + @staticmethod + def _continuation_output_ceilings( + request_kwargs: Mapping[str, object], + emitted_tokens: int, + ) -> Mapping[str, int] | None: + """The max_tokens / max_completion_tokens a continuation should carry, each + reduced by the tokens already emitted so the whole answer stays within the + caller's original allowance instead of getting a fresh one on every fallback + hop. Reductions compound across hops because the trimmed ceiling is what the + next hop sees. Returns None when the allowance is already exhausted, so the + stream must not be continued.""" + ceilings: Final = MappingProxyType( + { + key: value - emitted_tokens + for key in ("max_tokens", "max_completion_tokens") + if isinstance(value := request_kwargs.get(key), int) + } + ) + if ceilings and min(ceilings.values()) <= 0: + return None + return ceilings + @staticmethod def _build_completion_continuation_input( messages: Sequence[Mapping[str, str]], 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 ccd454bc207..c640374858a 100644 --- a/litellm/router_utils/pre_call_checks/continuation_prefill_check.py +++ b/litellm/router_utils/pre_call_checks/continuation_prefill_check.py @@ -46,6 +46,7 @@ def _deployment_supports_prefill(deployment: object) -> bool: if isinstance(declared, bool): return declared except ValidationError: + # No usable model_info override; fall through to the cost-map lookup. pass try: litellm_params: Final = _STR_KEYED_DICT_ADAPTER.validate_python(deployment_map.get("litellm_params")) diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index c37a23df7d1..987b96b3228 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -2363,6 +2363,26 @@ def test_build_completion_continuation_input_folds_into_existing_prefill(): assert twice[-1] == {"role": "assistant", "content": "part one part two", "prefix": True} +def test_continuation_output_ceilings_reduces_by_emitted_tokens(): + """A continuation must complete within the caller's original allowance, so each + output ceiling is reduced by the tokens already emitted.""" + assert litellm.Router._continuation_output_ceilings({"max_tokens": 100}, 30) == {"max_tokens": 70} + assert litellm.Router._continuation_output_ceilings({"max_tokens": 100, "max_completion_tokens": 40}, 25) == { + "max_tokens": 75, + "max_completion_tokens": 15, + } + # no ceiling configured -> nothing to reduce, continuation proceeds as before + assert litellm.Router._continuation_output_ceilings({}, 50) == {} + + +def test_continuation_output_ceilings_none_when_allowance_exhausted(): + """When the emitted tokens already meet or exceed a ceiling, there is no budget + left to continue, so the helper signals a decline rather than a fresh allowance.""" + assert litellm.Router._continuation_output_ceilings({"max_tokens": 20}, 20) is None + assert litellm.Router._continuation_output_ceilings({"max_tokens": 20}, 25) is None + assert litellm.Router._continuation_output_ceilings({"max_tokens": 100, "max_completion_tokens": 10}, 10) is None + + def test_mid_stream_continuation_eligible_allows_text_response_format(): """response_format={"type": "text"} is the unconstrained default and must stay eligible, unlike json_object / json_schema."""