fix(proxy): resolve team-alias models in the stream usage support gate

Team-scoped models store an internal model_name_{team_id}_{uuid} name
with the public alias only in team_public_model_name, so resolving them
through get_model_list without team_id returned no deployments and the
gate skipped injection, leaving those streams on tiktoken estimates.
Thread user_api_key_dict.team_id through the gate.
This commit is contained in:
mateo-berri 2026-07-30 17:46:37 -07:00
parent 0732122536
commit c0de87d08d
2 changed files with 24 additions and 5 deletions

View file

@ -258,7 +258,7 @@ def _litellm_model_supports_stream_options(litellm_model: str) -> bool:
return supported_params is not None and "stream_options" in supported_params
def _deployment_litellm_model(deployment: Mapping[str, object]) -> Optional[str]:
def _deployment_litellm_model(deployment: Mapping[str, object]) -> str | None:
litellm_params = deployment.get("litellm_params")
if isinstance(litellm_params, Mapping):
litellm_model = litellm_params.get("model")
@ -269,11 +269,12 @@ def _deployment_litellm_model(deployment: Mapping[str, object]) -> Optional[str]
def _model_deployments_support_stream_options(
model: object,
llm_router: Optional[Router],
llm_router: Router | None,
team_id: str | None,
) -> bool:
if not isinstance(model, str):
return False
deployments = llm_router.get_model_list(model_name=model) if llm_router is not None else None
deployments = llm_router.get_model_list(model_name=model, team_id=team_id) if llm_router is not None else None
deployment_models = tuple(
litellm_model
for deployment in deployments or ()
@ -1309,6 +1310,7 @@ class ProxyBaseLLMRequestProcessing:
supports_stream_options=lambda: _model_deployments_support_stream_options(
model=self.data.get("model"),
llm_router=llm_router,
team_id=user_api_key_dict.team_id,
),
)
)

View file

@ -5265,12 +5265,12 @@ class TestApplyStreamUsageTracking:
class TestModelDeploymentsSupportStreamOptions:
def _support(self, model, llm_router=None) -> bool:
def _support(self, model, llm_router=None, team_id=None) -> bool:
from litellm.proxy.common_request_processing import (
_model_deployments_support_stream_options,
)
return _model_deployments_support_stream_options(model=model, llm_router=llm_router)
return _model_deployments_support_stream_options(model=model, llm_router=llm_router, team_id=team_id)
def test_openai_compatible_deployment_supports_stream_options(self):
router = litellm.Router(
@ -5335,5 +5335,22 @@ class TestModelDeploymentsSupportStreamOptions:
def test_unmapped_model_name_is_not_injected(self):
assert self._support("some-unmapped-public-alias", None) is False
def test_team_alias_model_resolves_with_team_id(self):
router = litellm.Router(
model_list=[
{
"model_name": "model_name_team-1_8b6a0b3f",
"litellm_params": {"model": "azure/gpt-5.4-nano", "api_key": "fake"},
"model_info": {
"team_id": "team-1",
"team_public_model_name": "team-gpt",
},
}
]
)
assert self._support("team-gpt", router, team_id="team-1") is True
assert self._support("team-gpt", router, team_id=None) is False
def test_non_string_model_is_not_injected(self):
assert self._support(None, None) is False