scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-05 07:37:50 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 233

Merge 6210f4bbaa into c23cff7e45

Browse source
This commit is contained in:
ishaan-j-berri 2026-10-27 21:42:59 +00:00 • committed by GitHub
commit 673849fa7d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
17 changed files with 1311 additions and 88 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
litellm/proxy/_types.py
View file

@ -3927,6 +3927,8 @@ class ProxyErrorTypes(str, enum.Enum):
Team does not have access to the model
"""
team_selection_required = "team_selection_required"
user_model_access_denied = "user_model_access_denied"
"""
User does not have access to the model

209
litellm/proxy/auth/ui_session_models.py Normal file
View file

@ -0,0 +1,209 @@
"""Resolve dashboard model requests through the signed-in user's existing grants."""
from __future__ import annotations
from collections.abc import Sequence
from dataclasses import dataclass
from typing import TYPE_CHECKING, Final
from litellm.constants import UI_SESSION_TOKEN_TEAM_ID
from litellm.proxy._types import (
LiteLLM_TeamMembership,
LiteLLM_TeamTableCachedObj,
LiteLLM_UserTable,
LitellmUserRoles,
ProxyErrorTypes,
ProxyException,
UserAPIKeyAuth,
)
from litellm.proxy.auth.auth_checks import (
can_key_call_model,
can_team_access_model,
can_user_call_model,
check_team_member_model_access,
)
from litellm.proxy.auth.model_access_denied import model_access_denied_client_message
from litellm.proxy.auth.resolvers.grants import (
GrantResolver,
ResolvedGrants,
UserLookup,
user_models,
)
from litellm.proxy.auth.team_grants import team_grants
from litellm.proxy.mcp_server.ui_session_utils import is_ui_session_credential
from litellm.types.proxy.management_endpoints.model_management_endpoints import (
ModelGroupTeam,
)
if TYPE_CHECKING:
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.proxy.utils import PrismaClient, ProxyLogging
from litellm.router import Router
def is_model_ui_session(auth: UserAPIKeyAuth) -> bool:
return (
is_ui_session_credential(auth)
and auth.agent_id is None
and auth.user_role in (None, LitellmUserRoles.INTERNAL_USER, LitellmUserRoles.INTERNAL_USER_VIEW_ONLY)
)
@dataclass(frozen=True, slots=True)
class TeamModelSource:
team: LiteLLM_TeamTableCachedObj
membership: LiteLLM_TeamMembership | None
auth: UserAPIKeyAuth
@dataclass(frozen=True, slots=True)
class UISessionModelSources:
auth: UserAPIKeyAuth
user: LiteLLM_UserTable
teams: tuple[TeamModelSource, ...]
router: Router | None
prisma: PrismaClient
cache: UserApiKeyCache
logging: ProxyLogging
async def directly_allows(self, models: Sequence[str]) -> bool:
try:
for model in models:
await can_key_call_model(model, None, self.auth, self.router, self.prisma)
await can_user_call_model(model=model, user_object=self.user, llm_router=self.router)
except ProxyException:
return False
return True
async def team_allows(self, source: TeamModelSource, models: Sequence[str]) -> bool:
if source.team.blocked:
return False
try:
for model in models:
await can_team_access_model(
model=model,
team_object=source.team,
llm_router=self.router,
team_model_aliases=source.auth.team_model_aliases,
prisma_client=self.prisma,
)
await check_team_member_model_access(
model=model,
team_object=source.team,
valid_token=source.auth,
llm_router=self.router,
prisma_client=self.prisma,
user_api_key_cache=self.cache,
proxy_logging_obj=self.logging,
team_membership=source.membership,
team_membership_loaded=True,
)
except ProxyException:
return False
return True
async def available_teams(self, model: str) -> tuple[ModelGroupTeam, ...] | None:
if await self.directly_allows((model,)):
return ()
allowed: Final = tuple(
[
ModelGroupTeam(team_id=source.team.team_id, team_alias=source.team.team_alias)
for source in self.teams
if await self.team_allows(source, (model,))
]
)
return allowed or None
async def for_request(self, models: Sequence[str], selected_team_id: str | None) -> UserAPIKeyAuth:
if selected_team_id is None and await self.directly_allows(models):
return self.auth
allowed: Final = tuple([source for source in self.teams if await self.team_allows(source, models)])
if selected_team_id is not None:
selected: Final = next(
(source for source in allowed if source.team.team_id == selected_team_id),
None,
)
if selected is None:
raise ProxyException(
message="The selected team does not grant access to this model.",
type=ProxyErrorTypes.team_model_access_denied,
param="x-litellm-team-id",
code=403,
)
return selected.auth
if len(allowed) == 1:
return allowed[0].auth
if len(allowed) > 1:
raise ProxyException(
message="Choose a team for this model before sending the request.",
type=ProxyErrorTypes.team_selection_required,
param="x-litellm-team-id",
code=409,
)
raise ProxyException(
message=model_access_denied_client_message(model=list(models)),
type=ProxyErrorTypes.key_model_access_denied,
param="model",
code=403,
)
async def load_ui_session_model_sources(
auth: UserAPIKeyAuth,
router: Router | None,
prisma: PrismaClient | None,
cache: UserApiKeyCache,
logging: ProxyLogging,
) -> UISessionModelSources | None:
if not is_model_ui_session(auth):
return None
if prisma is None:
raise ProxyException(
message="Unable to resolve the signed-in user's model access.",
type=ProxyErrorTypes.no_db_connection,
param=None,
code=503,
)
resolver: Final = GrantResolver(prisma, cache, parent_otel_span=auth.parent_otel_span, proxy_logging_obj=logging)
identity: Final = await resolver.resolve(UserLookup(auth.user_id), team_id=None)
if not isinstance(identity, ResolvedGrants) or identity.user_object is None:
raise ProxyException(
message="Unable to resolve the signed-in user's model access.",
type=ProxyErrorTypes.auth_error,
param=None,
code=403,
)
user: Final = identity.user_object
if user.user_role not in (LitellmUserRoles.INTERNAL_USER, LitellmUserRoles.INTERNAL_USER_VIEW_ONLY):
return None
# Login copies the user model list into the UI key. Refresh that snapshot
# without changing the session's identity, restrictions, or budget.
live_auth: Final = auth.model_copy(update={"models": list(user_models(user))})
sources: Final = tuple(
[
await _team_model_source(resolver, auth, user, team_id)
for team_id in dict.fromkeys(user.teams)
if team_id != UI_SESSION_TOKEN_TEAM_ID
]
)
return UISessionModelSources(
live_auth, user, tuple(source for source in sources if source is not None), router, prisma, cache, logging
)
async def _team_model_source(
resolver: GrantResolver, auth: UserAPIKeyAuth, user: LiteLLM_UserTable, team_id: str
) -> TeamModelSource | None:
grant: Final = await resolver.resolve(UserLookup(user.user_id), team_id=team_id)
if not isinstance(grant, ResolvedGrants) or grant.team_object is None:
return None
team: Final = grant.team_object
scoped: Final = auth.model_copy(
update={
**team_grants(team, grant.team_membership, user.user_id),
"team_id": team.team_id,
"org_id": auth.org_id or team.organization_id,
"models": [],
}
)
return TeamModelSource(team, grant.team_membership, scoped)

25
litellm/proxy/auth/user_api_key_auth.py
View file

@ -2901,6 +2901,31 @@ async def validate_resolved_virtual_key( # noqa: C901 # Preserve ordering of e
param="None",
code=status.HTTP_401_UNAUTHORIZED,
)
if RouteChecks.is_llm_api_route(route=route):
from litellm.proxy.auth.ui_session_models import load_ui_session_model_sources
session_sources: Final = await load_ui_session_model_sources(
valid_token, llm_router, prisma_client, user_api_key_cache, proxy_logging_obj
)
if session_sources is not None:
requested_model: Final = _get_model_from_request_context(
request_data, route, request, llm_router, valid_token.team_id
)
requested_models: Final = (
(requested_model,) if isinstance(requested_model, str) else tuple(requested_model or ())
)
if requested_models:
valid_token = await session_sources.for_request( # rebind-ok: replace the local token with a scoped copy for all existing checks below
(
*requested_models,
*(
name
for target in iter_request_fallback_targets(request_data)
if (name := fallback_target_model_name(target)) is not None
),
),
request.headers.get("x-litellm-team-id"),
)
await enforce_key_and_fallback_model_access(
valid_token=valid_token,
request_data=request_data,

4
litellm/proxy/management_endpoints/key_management_endpoints.py
View file

@ -7901,7 +7901,7 @@ async def block_key(
--header "Authorization: Bearer $LITELLM_MASTER_KEY" \
--header 'Content-Type: application/json' \
--data '{
"key": "sk-Fn8Ej39NxjAXrvpUGKghGw"
"key": "YOUR_VIRTUAL_KEY"
}'
```
@ -8015,7 +8015,7 @@ async def unblock_key(
--header "Authorization: Bearer $LITELLM_MASTER_KEY" \
--header 'Content-Type: application/json' \
--data '{
"key": "sk-Fn8Ej39NxjAXrvpUGKghGw"
"key": "YOUR_VIRTUAL_KEY"
}'
```

25
litellm/proxy/proxy_server.py
View file

@ -1002,6 +1002,7 @@ from litellm.types.llms.openai import (
from litellm.types.proxy.control_plane_endpoints import WorkerRegistryEntry
from litellm.types.proxy.management_endpoints.model_management_endpoints import (
ModelGroupInfoProxy,
ModelGroupInfoResponse,
)
from litellm.types.proxy.management_endpoints.ui_sso import (
DefaultTeamSSOParams,
@ -17917,6 +17918,7 @@ _get_model_group_info: Final = get_model_group_info
"/model_group/info",
tags=["model management"],
dependencies=[Depends(user_api_key_auth)],
response_model=ModelGroupInfoResponse,
)
async def model_group_info(
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
@ -18120,13 +18122,24 @@ async def model_group_info(
append_agents_to_model_group,
)
visible_groups: Final = get_model_group_info(
llm_router=llm_router,
all_models_str=listed_group_names,
model_group=model_group,
team_id=user_api_key_dict.team_id,
)
from litellm.proxy.auth.ui_session_models import load_ui_session_model_sources
session_sources: Final = await load_ui_session_model_sources(
user_api_key_dict, llm_router, prisma_client, user_api_key_cache, proxy_logging_obj
)
if session_sources is not None:
for group in visible_groups:
group.available_teams = await session_sources.available_teams(group.model_group)
model_groups: Final = await append_agents_to_model_group(
model_groups=get_model_group_info(
llm_router=llm_router,
all_models_str=listed_group_names,
model_group=model_group,
team_id=user_api_key_dict.team_id,
),
model_groups=[
group for group in visible_groups if session_sources is None or group.available_teams is not None
],
user_api_key_dict=user_api_key_dict,
)
internal_to_public: Final = TeamModelNameTranslator.build_internal_to_public_map(llm_router, general_settings)

53
litellm/proxy/utils.py
View file

@ -9579,6 +9579,16 @@ async def get_available_models_for_user(
get_key_models,
get_team_models,
)
from litellm.proxy.auth.ui_session_models import load_ui_session_model_sources
session_sources: Final = (
await load_ui_session_model_sources(
user_api_key_dict, llm_router, prisma_client, user_api_key_cache, proxy_logging_obj
)
if user_api_key_cache is not None and proxy_logging_obj is not None and team_id is None
else None
)
model_auth: Final = session_sources.auth if session_sources is not None else user_api_key_dict
# Get proxy model list and access groups
if llm_router is None:
@ -9590,7 +9600,7 @@ async def get_available_models_for_user(
requested_team_object: Final = (
await _get_validated_team_object(
user_api_key_dict=user_api_key_dict,
user_api_key_dict=model_auth,
team_id=team_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
@ -9602,9 +9612,9 @@ async def get_available_models_for_user(
key_models: Final[Sequence[str]] = (
()
if requested_team_object is not None and requested_team_object.team_id != user_api_key_dict.team_id
if requested_team_object is not None and requested_team_object.team_id != model_auth.team_id
else get_key_models(
user_api_key_dict=user_api_key_dict,
user_api_key_dict=model_auth,
proxy_model_list=proxy_model_list,
model_access_groups=model_access_groups,
include_model_access_groups=include_model_access_groups,
@ -9612,19 +9622,17 @@ async def get_available_models_for_user(
)
team_models: Final = get_team_models(
team_models=(
requested_team_object.models if requested_team_object is not None else user_api_key_dict.team_models
),
team_models=(requested_team_object.models if requested_team_object is not None else model_auth.team_models),
proxy_model_list=proxy_model_list,
model_access_groups=model_access_groups,
include_model_access_groups=include_model_access_groups,
)
effective_team_id: Final = team_id or user_api_key_dict.team_id
effective_team_id: Final = team_id or model_auth.team_id
access_group_models: Final = (
await _get_access_group_models(
user_api_key_dict=user_api_key_dict,
user_api_key_dict=model_auth,
team_object=requested_team_object
or await _get_team_object_for_access_groups(
team_id=effective_team_id,
@ -9640,9 +9648,9 @@ async def get_available_models_for_user(
else _AccessGroupModels()
)
key_model_names: Final = cast("list[str]", user_api_key_dict.models) # cast-ok: key models are a string array
key_model_names: Final = cast("list[str]", model_auth.models) # cast-ok: key models are a string array
key_model_check_skipped: Final = (
bool(user_api_key_dict.config) or SpecialModelNames.all_team_models.value in key_model_names
bool(model_auth.config) or SpecialModelNames.all_team_models.value in key_model_names
)
granted_key_models: Final = (
(
@ -9675,19 +9683,38 @@ async def get_available_models_for_user(
list(
await personal_key_team_visible_models(
models=complete_models,
user_api_key_dict=user_api_key_dict,
user_api_key_dict=model_auth,
llm_router=llm_router,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
)
)
if personal_key_team_model_access_applies(user_api_key_dict)
if personal_key_team_model_access_applies(model_auth)
else complete_models
)
if session_sources is not None:
candidates: Final = list(all_models)
for source in session_sources.teams:
candidates.extend(
await get_available_models_for_user(
user_api_key_dict=source.auth,
llm_router=llm_router,
general_settings=general_settings,
user_model=user_model,
prisma_client=prisma_client,
proxy_logging_obj=proxy_logging_obj,
include_model_access_groups=include_model_access_groups,
only_model_access_groups=only_model_access_groups,
return_wildcard_routes=return_wildcard_routes,
user_api_key_cache=user_api_key_cache,
)
)
return [name for name in dict.fromkeys(candidates) if await session_sources.available_teams(name) is not None]
agent_visible: Final = await _agent_access_group_visible_models(
user_api_key_dict=user_api_key_dict,
user_api_key_dict=model_auth,
llm_router=llm_router,
include_model_access_groups=include_model_access_groups,
return_wildcard_routes=return_wildcard_routes,

14
litellm/types/proxy/management_endpoints/model_management_endpoints.py
View file

@ -34,6 +34,13 @@ class PatchModelRequest(updateDeployment):
return self
class ModelGroupTeam(LiteLLMBaseModel):
model_config = ConfigDict(frozen=True)
team_id: str
team_alias: str | None = None
class ModelGroupInfoProxy(ModelGroupInfo):
is_public_model_group: bool = Field(default=False)
health_status: str | None = Field(default=None)
@ -41,6 +48,13 @@ class ModelGroupInfoProxy(ModelGroupInfo):
health_checked_at: str | None = Field(default=None)
custom_tags: tuple[str, ...] | None = None
custom_metadata: Mapping[str, str | int | float | bool] | None = None
available_teams: tuple[ModelGroupTeam, ...] | None = None
class ModelGroupInfoResponse(LiteLLMBaseModel):
model_config = ConfigDict(frozen=True)
data: tuple[ModelGroupInfoProxy, ...]
class UpdateUsefulLinksRequest(LiteLLMBaseModel):

398
tests/unit/proxy/auth/test_ui_session_models.py Normal file
View file

@ -0,0 +1,398 @@
from typing import cast
from unittest.mock import AsyncMock
import pytest
from fastapi.encoders import jsonable_encoder
import litellm
from litellm.constants import UI_SESSION_TOKEN_TEAM_ID
from litellm.proxy import proxy_server
from litellm.proxy._types import (
LiteLLM_AccessGroupTable,
LiteLLM_BudgetTable,
LiteLLM_TeamMembership,
LiteLLM_TeamTableCachedObj,
LiteLLM_UserTable,
LitellmUserRoles,
ProxyException,
UserAPIKeyAuth,
)
from litellm.proxy.auth.auth_checks import can_key_call_model, team_max_budget_check
from litellm.proxy.auth.ui_session_models import load_ui_session_model_sources
from litellm.proxy.common_utils.user_api_key_cache import (
UserApiKeyCache,
team_membership_reservation_cache_key,
)
from litellm.proxy.utils import (
PrismaClient,
ProxyLogging,
get_available_models_for_user,
)
from litellm.types.proxy.management_endpoints.model_management_endpoints import (
ModelGroupInfoProxy,
ModelGroupInfoResponse,
ModelGroupTeam,
)
@pytest.fixture
def router():
return litellm.Router(
model_list=[
{
"model_name": name,
"litellm_params": {
"model": "openai/gpt-4.1-mini",
"api_key": "test-key",
},
}
for name in ("personal", "research", "shared", "forbidden")
]
)
@pytest.fixture
async def account(monkeypatch, router):
cache = UserApiKeyCache()
logging = ProxyLogging(cache)
# Every row is in the real auth cache; the database must never be touched.
prisma = cast(PrismaClient, object())
user = LiteLLM_UserTable(
user_id="researcher",
user_role=LitellmUserRoles.INTERNAL_USER_VIEW_ONLY,
models=["personal"],
teams=["research-team"],
max_budget=10,
)
team = LiteLLM_TeamTableCachedObj(
team_id="research-team",
team_alias="Research",
models=["research", "shared"],
max_budget=20,
)
member = LiteLLM_TeamMembership(user_id=user.user_id, team_id=team.team_id, spend=0)
auth = UserAPIKeyAuth(
token="session-hash",
user_id=user.user_id,
team_id=UI_SESSION_TOKEN_TEAM_ID,
user_role=LitellmUserRoles.INTERNAL_USER_VIEW_ONLY,
models=["personal"],
max_budget=3,
)
await cache.async_set_cache(key=user.user_id, value=user, model_type=LiteLLM_UserTable)
await cache.async_set_cache(key=f"team_id:{team.team_id}", value=team, model_type=LiteLLM_TeamTableCachedObj)
await cache.async_set_cache(
key=team_membership_reservation_cache_key(user.user_id, team.team_id),
value=member,
)
monkeypatch.setattr(proxy_server, "prisma_client", prisma)
monkeypatch.setattr(proxy_server, "user_api_key_cache", cache)
monkeypatch.setattr(proxy_server, "proxy_logging_obj", logging)
return auth, user, team, member, prisma, cache, logging
async def sources(account, router):
auth, _, _, _, prisma, cache, logging = account
result = await load_ui_session_model_sources(auth, router, prisma, cache, logging)
assert result is not None
return result
async def test_discovery_includes_team_models_without_exposing_unrelated_models(account, router):
auth, _, _, _, prisma, cache, logging = account
models = await get_available_models_for_user(auth, router, {}, None, prisma, logging, user_api_key_cache=cache)
access = await sources(account, router)
assert {model: await access.available_teams(model) for model in models} == {
"personal": (),
"research": (ModelGroupTeam(team_id="research-team", team_alias="Research"),),
"shared": (ModelGroupTeam(team_id="research-team", team_alias="Research"),),
}
async def test_sole_team_preserves_identity_role_and_budgets_without_changing_cached_session(account, router):
access = await sources(account, router)
scoped = await access.for_request(("research",), None)
assert (
scoped.team_id,
scoped.user_id,
scoped.user_role,
scoped.max_budget,
scoped.team_max_budget,
) == (
"research-team",
"researcher",
LitellmUserRoles.INTERNAL_USER_VIEW_ONLY,
3,
20,
)
assert account[0].team_id == UI_SESSION_TOKEN_TEAM_ID
assert account[0].models == ["personal"]
async def test_session_grants_resolve_before_the_auth_builder_attaches_the_user_role(account, router):
auth, _, _, _, prisma, cache, logging = account
access = await load_ui_session_model_sources(
auth.model_copy(update={"user_role": None}), router, prisma, cache, logging
)
assert access is not None
assert (await access.for_request(("research",), None)).team_id == "research-team"
async def test_direct_user_access_is_the_default_without_a_team_selection(account, router):
access = await sources(account, router)
assert await access.for_request(("personal",), None) == account[0]
async def test_session_model_access_fails_closed_without_a_database(account, router):
auth, _, _, _, _, cache, logging = account
with pytest.raises(ProxyException) as error:
await load_ui_session_model_sources(auth, router, None, cache, logging)
assert (error.value.code, error.value.type, error.value.param) == ("503", "no_db_connection", None)
async def test_deleted_user_cannot_reuse_model_grants_from_the_session(account, router, monkeypatch):
auth, _, _, _, _, cache, logging = account
prisma = PrismaClient(database_url="postgresql://localhost/session-tests", proxy_logging_obj=logging)
missing_user = AsyncMock(return_value=None)
monkeypatch.setattr("prisma.actions.LiteLLM_UserTableActions.find_unique", missing_user)
deleted_user_auth = auth.model_copy(update={"user_id": "deleted-session-user"})
with pytest.raises(ProxyException) as error:
await load_ui_session_model_sources(deleted_user_auth, router, prisma, cache, logging)
assert (error.value.code, error.value.type, error.value.param) == ("403", "auth_error", None)
missing_user.assert_awaited_once_with(
where={"user_id": "deleted-session-user"}, include={"organization_memberships": True}
)
async def test_deleted_team_is_not_discoverable_or_selectable_with_stale_membership(account, router, monkeypatch):
auth, _, team, _, _, cache, logging = account
prisma = PrismaClient(database_url="postgresql://localhost/session-tests", proxy_logging_obj=logging)
missing_team = AsyncMock(return_value=None)
monkeypatch.setattr("prisma.actions.LiteLLM_TeamTableActions.find_unique", missing_team)
await cache.async_delete_cache(f"team_id:{team.team_id}")
access = await load_ui_session_model_sources(auth, router, prisma, cache, logging)
assert access is not None
assert await access.available_teams("personal") == ()
assert await access.available_teams("research") is None
with pytest.raises(ProxyException) as error:
await access.for_request(("research",), team.team_id)
assert (error.value.code, error.value.type, error.value.param) == (
"403",
"team_model_access_denied",
"x-litellm-team-id",
)
missing_team.assert_awaited_once()
async def test_changed_live_user_role_stops_internal_session_grant_resolution(account, router):
auth, user, _, _, prisma, cache, logging = account
await cache.async_set_cache(
key=user.user_id,
value=user.model_copy(update={"user_role": LitellmUserRoles.PROXY_ADMIN}),
model_type=LiteLLM_UserTable,
)
assert await load_ui_session_model_sources(auth, router, prisma, cache, logging) is None
assert auth.user_role == LitellmUserRoles.INTERNAL_USER_VIEW_ONLY
@pytest.mark.parametrize("live_models", [["no-default-models"], ["shared"]])
async def test_removed_user_and_team_grants_deny_even_when_the_session_snapshot_allows(account, router, live_models):
auth, user, _, _, prisma, cache, logging = account
await cache.async_set_cache(
key=user.user_id,
value=user.model_copy(update={"models": live_models, "teams": []}),
model_type=LiteLLM_UserTable,
)
stale_auth = auth.model_copy(update={"models": ["personal", "research"]})
access = await load_ui_session_model_sources(stale_auth, router, prisma, cache, logging)
assert access is not None
for model in ("personal", "research"):
assert await access.available_teams(model) is None
with pytest.raises(ProxyException) as error:
await access.for_request((model,), None)
assert (error.value.code, error.value.type, error.value.param) == ("403", "key_model_access_denied", "model")
async def test_new_direct_grants_are_listed_and_usable_without_logging_in_again(account, router):
auth, user, _, _, prisma, cache, logging = account
await cache.async_set_cache(
key=user.user_id,
value=user.model_copy(update={"models": ["personal", "forbidden"]}),
model_type=LiteLLM_UserTable,
)
models = await get_available_models_for_user(auth, router, {}, None, prisma, logging, user_api_key_cache=cache)
access = await sources(account, router)
scoped = await access.for_request(("forbidden",), None)
assert set(models) == {"personal", "research", "shared", "forbidden"}
assert await access.available_teams("forbidden") == ()
assert await can_key_call_model("forbidden", None, scoped, router)
assert (scoped.team_id, scoped.user_id, scoped.max_budget) == (UI_SESSION_TOKEN_TEAM_ID, "researcher", 3)
assert auth.models == ["personal"]
async def test_explicit_team_is_not_silently_replaced_by_a_direct_user_grant(account, router):
_, _, team, _, _, cache, _ = account
await cache.async_set_cache(
key=f"team_id:{team.team_id}",
value=team.model_copy(update={"models": ["personal", "research"]}),
model_type=LiteLLM_TeamTableCachedObj,
)
access = await sources(account, router)
assert (await access.for_request(("personal",), "research-team")).team_id == "research-team"
with pytest.raises(ProxyException) as error:
await access.for_request(("personal",), "removed-team")
assert error.value.code == "403"
async def test_team_access_group_grants_are_discoverable_and_usable(account, router):
_, _, team, _, _, cache, _ = account
group = LiteLLM_AccessGroupTable(
access_group_id="research-group", access_group_name="Research models", access_model_names=["research"]
)
await cache.async_set_cache(key="access_group_id:research-group", value=group, model_type=LiteLLM_AccessGroupTable)
await cache.async_set_cache(
key=f"team_id:{team.team_id}",
value=team.model_copy(update={"models": ["shared"], "access_group_ids": ["research-group"]}),
model_type=LiteLLM_TeamTableCachedObj,
)
access = await sources(account, router)
assert await access.available_teams("research") == (ModelGroupTeam(team_id="research-team", team_alias="Research"),)
assert (await access.for_request(("research",), None)).team_id == "research-team"
async def test_multiple_granting_teams_require_explicit_choice(account, router):
auth, user, _, _, _, cache, _ = account
other = LiteLLM_TeamTableCachedObj(team_id="other-team", models=["research"])
await cache.async_set_cache(
key=user.user_id,
value=user.model_copy(update={"teams": ["research-team", "other-team"]}),
model_type=LiteLLM_UserTable,
)
await cache.async_set_cache(key="team_id:other-team", value=other, model_type=LiteLLM_TeamTableCachedObj)
await cache.async_set_cache(
key=team_membership_reservation_cache_key(user.user_id, other.team_id),
value=LiteLLM_TeamMembership(user_id=user.user_id, team_id=other.team_id),
)
access = await sources(account, router)
with pytest.raises(ProxyException) as error:
await access.for_request(("research",), None)
assert (error.value.code, error.value.type, error.value.param) == (
"409",
"team_selection_required",
"x-litellm-team-id",
)
assert (await access.for_request(("research",), "other-team")).team_id == "other-team"
assert auth.team_id == UI_SESSION_TOKEN_TEAM_ID
@pytest.mark.parametrize("reason", ["blocked", "removed", "member_scope"])
async def test_unavailable_team_grants_are_not_listed_or_selectable(account, router, reason):
_, user, team, member, _, cache, _ = account
if reason == "blocked":
await cache.async_set_cache(
key=f"team_id:{team.team_id}",
value=team.model_copy(update={"blocked": True}),
model_type=LiteLLM_TeamTableCachedObj,
)
elif reason == "removed":
await cache.async_set_cache(
key=user.user_id,
value=user.model_copy(update={"teams": []}),
model_type=LiteLLM_UserTable,
)
else:
await cache.async_set_cache(
key=team_membership_reservation_cache_key(user.user_id, team.team_id),
value=member.model_copy(update={"litellm_budget_table": LiteLLM_BudgetTable(allowed_models=["shared"])}),
)
access = await sources(account, router)
assert await access.available_teams("research") is None
with pytest.raises(ProxyException) as error:
await access.for_request(("research",), "research-team")
assert error.value.code == "403"
async def test_forbidden_model_and_fallback_keep_normal_key_denial(account, router):
access = await sources(account, router)
with pytest.raises(ProxyException) as error:
await access.for_request(("research", "forbidden"), None)
assert (error.value.code, error.value.type) == ("403", "key_model_access_denied")
async def test_selected_team_still_enforces_its_budget(account, router):
access = await sources(account, router)
scoped = await access.for_request(("research",), None)
team = account[2].model_copy(update={"max_budget": 1, "spend": 2})
with pytest.raises(litellm.BudgetExceededError):
await team_max_budget_check(team, scoped, account[6])
@pytest.mark.parametrize(
"auth",
[
UserAPIKeyAuth(
user_id="researcher",
team_id="research-team",
user_role=LitellmUserRoles.INTERNAL_USER,
),
UserAPIKeyAuth(user_id="researcher", user_role=LitellmUserRoles.INTERNAL_USER),
UserAPIKeyAuth(
user_id="admin",
team_id=UI_SESSION_TOKEN_TEAM_ID,
user_role=LitellmUserRoles.PROXY_ADMIN,
),
],
)
async def test_other_credentials_are_never_widened(auth, router, account):
assert await load_ui_session_model_sources(auth, router, account[4], account[5], account[6]) is None
def test_typed_response_preserves_existing_alias_and_arbitrary_metadata():
group = ModelGroupInfoProxy(
model_group="team-visible-alias",
providers=["openai"],
custom_tags=("research",),
custom_metadata={"n": 5, "flag": True, "label": "Research"},
supported_endpoints=("/v1/responses",),
)
assert ModelGroupInfoResponse(data=(group,)).model_dump(mode="json") == jsonable_encoder({"data": [group]})
async def test_team_public_model_preserves_capabilities_in_session_discovery(account, monkeypatch):
auth, user, team, _, _, cache, _ = account
router = litellm.Router(
model_list=[
{
"model_name": "model_name_research-team_private",
"litellm_params": {"model": "openai/gpt-4.1-mini", "api_key": "test-key"},
"model_info": {"team_id": "research-team", "team_public_model_name": "team-chat", "mode": "chat"},
}
]
)
await cache.async_set_cache(
key=f"team_id:{team.team_id}",
value=team.model_copy(update={"models": ["model_name_research-team_private"]}),
model_type=LiteLLM_TeamTableCachedObj,
)
await cache.async_set_cache(
key=user.user_id,
value=user.model_copy(update={"models": ["no-default-models"]}),
model_type=LiteLLM_UserTable,
)
monkeypatch.setattr(proxy_server, "llm_router", router)
monkeypatch.setattr(proxy_server, "llm_model_list", router.model_list)
monkeypatch.setattr(proxy_server, "general_settings", {})
monkeypatch.setattr(proxy_server, "user_model", None)
response = await proxy_server.model_group_info(user_api_key_dict=auth)
assert [(group.model_group, group.providers, group.mode, group.available_teams) for group in response["data"]] == [
(
"model_name_research-team_private",
["openai"],
"chat",
(ModelGroupTeam(team_id="research-team", team_alias="Research"),),
),
]

263
tests/unit/proxy/auth/test_user_api_key_auth.py
View file

@ -548,8 +548,6 @@ def test_get_api_key_from_custom_header_different_casing():
)
@pytest.mark.parametrize(
"user_role, auth_user_id, requested_user_id, expected_result",
[
@ -963,71 +961,60 @@ async def test_accepted_websocket_reauthorizes_the_frame_model_without_closing()
@pytest.mark.parametrize("enforce_rbac", [True, False])
@pytest.mark.asyncio
async def test_jwt_user_api_key_auth_builder_enforce_rbac(enforce_rbac, monkeypatch):
from litellm.proxy.auth.handle_jwt import JWTHandler, JWTAuthManager
from unittest.mock import patch, Mock
async def test_jwt_user_api_key_auth_builder_enforce_rbac(enforce_rbac: bool, monkeypatch: pytest.MonkeyPatch) -> None:
import jwt as pyjwt
from cryptography.hazmat.primitives.asymmetric import rsa
from litellm.proxy._types import LiteLLM_JWTAuth
from litellm.caching import DualCache
from litellm.proxy.auth.handle_jwt import JWTAuthManager, JWTHandler
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.proxy.utils import ProxyLogging
monkeypatch.setenv("JWT_PUBLIC_KEY_URL", "my-fake-url")
monkeypatch.setenv("JWT_PUBLIC_KEY_URL", "https://idp.example.com/keys")
monkeypatch.setenv("JWT_AUDIENCE", "api://LiteLLM_Proxy-dev")
monkeypatch.delenv("JWT_ISSUER", raising=False)
local_cache = DualCache()
keys = [
{
"kty": "RSA",
"use": "sig",
"kid": "z1rsYHHJ9-8mggt4HsZu8BKkBPw",
"x5t": "z1rsYHHJ9-8mggt4HsZu8BKkBPw",
"n": "pOe4GbleFDT1u5ioOQjNMmhvkDVoVD9cBKvX7AlErtWA_D6wc1w1iwkd6arYVCPObZbAB4vLSXrlpBSOuP6VYnXw_cTgniv_c82ra-mfqCpM-SbqzZ3sVqlcE_bwxvci_4PrxAW4R85ok12NXyZ2371H3yGevabi35AlVm-bQ24azo1hLK_0DzB6TxsAIOTOcKfIugOfqP-B2R4vR4u6pYftS8MWcxegr9iJ5JNtubI1X2JHpxJhkRoMVwKFna2GXmtzdxLi3yS_GffVCKfTbFMhalbJS1lSmLqhmLZZL-lrQZ6fansTl1vcGcoxnzPTwBkZMks0iVV4yfym_gKBXQ",
"e": "AQAB",
"x5c": [
"MIIC/TCCAeWgAwIBAgIIQk8Qok6pfXkwDQYJKoZIhvcNAQELBQAwLTErMCkGA1UEAxMiYWNjb3VudHMuYWNjZXNzY29udHJvbC53aW5kb3dzLm5ldDAeFw0yNDExMjcwOTA0MzlaFw0yOTExMjcwOTA0MzlaMC0xKzApBgNVBAMTImFjY291bnRzLmFjY2Vzc2NvbnRyb2wud2luZG93cy5uZXQwggEiMA0GCSqGSIb3DQEBAQUAA4IBDwAwggEKAoIBAQCk57gZuV4UNPW7mKg5CM0yaG+QNWhUP1wEq9fsCUSu1YD8PrBzXDWLCR3pqthUI85tlsAHi8tJeuWkFI64/pVidfD9xOCeK/9zzatr6Z+oKkz5JurNnexWqVwT9vDG9yL/g+vEBbhHzmiTXY1fJnbfvUffIZ69puLfkCVWb5tDbhrOjWEsr/QPMHpPGwAg5M5wp8i6A5+o/4HZHi9Hi7qlh+1LwxZzF6Cv2Inkk225sjVfYkenEmGRGgxXAoWdrYZea3N3EuLfJL8Z99UIp9NsUyFqVslLWVKYuqGYtlkv6WtBnp9qexOXW9wZyjGfM9PAGRkySzSJVXjJ/Kb+AoFdAgMBAAGjITAfMB0GA1UdDgQWBBSTO5FmUwwGS+1CNqg2uNgjxUjFijANBgkqhkiG9w0BAQsFAAOCAQEAok04z0ICMEHGqDTzx6eD7vvJP8itJTCSz8JcZcGVJofJpViGF3bNnyeSPa7vNDYP1Ps9XBvw3/n2s+yynZ8EwFxMyxCZRCSbLv0N+cAbH3rmZqGcgMJszZVwcFUtXQPTe1ZRyHtEyOB+PVFH7K7obysRVO/cC6EGqIF3pYWzez/dtMaXRAkdTNlz0ko62WoA4eMPwUFCITjW/Jxfxl0BNUbo82PXXKhaeVJb+EgFG5b/pWWPswWmBoQhmD5G1UODvEACHRl/cHsPPqe4YE+6D1/wMno/xqqyGltnk8v0d4TpNcQMn9oM19V+OGgrzWOvvXhvnhqUIVGMsRlyBGNHAw=="
],
"cloud_instance_name": "microsoftonline.com",
"issuer": "https://login.microsoftonline.com/bdfd79b3-8401-47fb-a764-6e595c455b05/v2.0",
}
]
local_cache.set_cache(
key="litellm_jwt_auth_keys_my-fake-url",
value=keys,
)
litellm_jwtauth = LiteLLM_JWTAuth(
**{
"admin_jwt_scope": "litellm_proxy_endpoints_access",
"admin_allowed_routes": ["openai_routes", "info_routes"],
"public_key_ttl": 600,
"enforce_rbac": enforce_rbac,
}
)
jwt_handler = JWTHandler()
private_key: Final = rsa.generate_private_key(public_exponent=65537, key_size=2048)
public_jwk: Final = {
**json.loads(pyjwt.algorithms.RSAAlgorithm.to_jwk(private_key.public_key())),
"kid": "synthetic-rbac-test",
}
claims: Final = {"sub": "synthetic-rbac-user", "aud": "api://LiteLLM_Proxy-dev"}
token: Final = pyjwt.encode(claims, private_key, algorithm="RS256", headers={"kid": "synthetic-rbac-test"})
local_cache: Final = UserApiKeyCache()
local_cache.set_cache(key="litellm_jwt_auth_keys_https://idp.example.com/keys", value=[public_jwk])
jwt_handler: Final = JWTHandler()
jwt_handler.update_environment(
prisma_client=None,
user_api_key_cache=local_cache,
litellm_jwtauth=litellm_jwtauth,
leeway=10000000000000,
litellm_jwtauth=LiteLLM_JWTAuth(
admin_jwt_scope="litellm_proxy_endpoints_access",
admin_allowed_routes=["openai_routes", "info_routes"],
public_key_ttl=600,
enforce_rbac=enforce_rbac,
),
)
args = {
"api_key": "eyJ0eXAiOiJKV1QiLCJhbGciOiJSUzI1NiIsIng1dCI6InoxcnNZSEhKOS04bWdndDRIc1p1OEJLa0JQdyIsImtpZCI6InoxcnNZSEhKOS04bWdndDRIc1p1OEJLa0JQdyJ9.eyJhdWQiOiJhcGk6Ly9MaXRlTExNX1Byb3h5LWRldiIsImlzcyI6Imh0dHBzOi8vc3RzLndpbmRvd3MubmV0L2JkZmQ3OWIzLTg0MDEtNDdmYi1hNzY0LTZlNTk1YzQ1NWIwNS8iLCJpYXQiOjE3MzcyNDE3ODEsIm5iZiI6MTczNzI0MTc4MSwiZXhwIjoxNzM3MjQ1NjgxLCJhaW8iOiJrMlJnWUpBNE5hZGg4MGJQdXlyRmxlV1o3dHZiQUE9PSIsImFwcGlkIjoiOGNjZjNkMDItMmNkNi00N2I5LTgxODUtMGVkYjI0YWJjZjY5IiwiYXBwaWRhY3IiOiIxIiwiaWRwIjoiaHR0cHM6Ly9zdHMud2luZG93cy5uZXQvYmRmZDc5YjMtODQwMS00N2ZiLWE3NjQtNmU1OTVjNDU1YjA1LyIsIm9pZCI6IjQ0YTg3YTYzLWFiNTUtNDc4NS1iMmFmLTMzNjllZWM4ZTEzOSIsInJoIjoiMS5BYjBBczNuOXZRR0UtMGVuWkc1WlhFVmJCY0VDbkl6NHJxaE9wZ2E0UGZSZjBsbTlBQUM5QUEuIiwic3ViIjoiNDRhODdhNjMtYWI1NS00Nzg1LWIyYWYtMzM2OWVlYzhlMTM5IiwidGlkIjoiYmRmZDc5YjMtODQwMS00N2ZiLWE3NjQtNmU1OTVjNDU1YjA1IiwidXRpIjoiY3ltNVhlcmhIMHVMSlNZU1JyQmhBQSIsInZlciI6IjEuMCJ9.UooJjM9pS-wgYsExqgHdrYyQhp7NbwAsr7au9dWJaLpsufXeyHJSg-Xd5VJ4RsDVJiDes3jkC7WeoAiaCfzEHpAum-p_aqqLYXf1QIYbi1hLC0m7y_klFcqMp11WbDa9TSTvg-o8q3x2Y5su8X23ymlFih4OP17b7JA6a4_2MybU5QkCEW1tQK6VspuuXzeDHvbfGeGYcIptHFyfttHMHHXRtX1o9bX7gOR_dwFITAXD18T4ZdAN_0y6f1OtVF9TMWQhMXhKU8ahn8TSg_CXmPl9T_1gV3ZWLvVtcdVrWs82fDz3-2lEw28z4bQEr1Z5xoAz7srhx1WEBu_ioAcQiA",
args: Final = {
"api_key": token,
"jwt_handler": jwt_handler,
"route": "/v1/chat/completions",
"prisma_client": None,
"user_api_key_cache": Mock(),
"user_api_key_cache": local_cache,
"parent_otel_span": None,
"proxy_logging_obj": Mock(),
"proxy_logging_obj": ProxyLogging(local_cache),
"request_data": {},
"general_settings": {},
}
if enforce_rbac:
with pytest.raises(HTTPException):
with pytest.raises(HTTPException) as denied:
await JWTAuthManager.auth_builder(**args)
assert denied.value.status_code == 403
assert "enforce_rbac is set to True" in denied.value.detail
else:
await JWTAuthManager.auth_builder(**args)
authorized: Final = await JWTAuthManager.auth_builder(**args)
assert authorized["jwt_claims"] == claims
assert authorized["is_proxy_admin"] is False
def test_user_api_key_auth_end_user_str():
@ -1138,7 +1125,7 @@ async def test_jwt_non_admin_team_route_access(monkeypatch):
"org_object": None,
"end_user_id": None,
"end_user_object": None,
"token": "eyJhbGciOiJSUzI1NiIsInR5cCIgOiAiSldUIiwia2lkIiA6ICJmR09YQTNhbHFObjByRzJ6OHJQT1FLZVVMSWxCNDFnVWl4VDJ5WE1QVG1ZIn0.eyJleHAiOjE3NDI2MDAzODIsImlhdCI6MTc0MjYwMDA4MiwianRpIjoiODRhNjZmZjAtMTE5OC00YmRkLTk1NzAtNWZhMjNhZjYxMmQyIiwiaXNzIjoiaHR0cDovL2xvY2FsaG9zdDo4MDgwL3JlYWxtcy9saXRlbGxtLXJlYWxtIiwiYXVkIjoiYWNjb3VudCIsInN1YiI6ImZmMGZjOGNiLWUyMjktNDkyYy05NzYwLWNlYzVhMDYxNmI2MyIsInR5cCI6IkJlYXJlciIsImF6cCI6ImxpdGVsbG0tdGVzdC1jbGllbnQtaWQiLCJzaWQiOiI4MTYwNjIxOC0yNmZmLTQwMjAtOWQxNy05Zjc0YmFlNTBkODUiLCJhY3IiOiIxIiwiYWxsb3dlZC1vcmlnaW5zIjpbImh0dHA6Ly9sb2NhbGhvc3Q6NDAwMC8qIl0sInJlYWxtX2FjY2VzcyI6eyJyb2xlcyI6WyJvZmZsaW5lX2FjY2VzcyIsImRlZmF1bHQtcm9sZXMtbGl0ZWxsbS1yZWFsbSIsInVtYV9hdXRob3JpemF0aW9uIl19LCJyZXNvdXJjZV9hY2Nlc3MiOnsiYWNjb3VudCI6eyJyb2xlcyI6WyJtYW5hZ2UtYWNjb3VudCIsIm1hbmFnZS1hY2NvdW50LWxpbmtzIiwidmlldy1wcm9maWxlIl19fSwic2NvcGUiOiJwcm9maWxlIGdyb3Vwcy1zY29wZSBlbWFpbCBsaXRlbGxtLmFwaS5jb25zdW1lciIsImVtYWlsX3ZlcmlmaWVkIjp0cnVlLCJuYW1lIjoiS3Jpc2ggRGhvbGFraWEiLCJncm91cHMiOlsiL28zX21pbmlfYWNjZXNzIl0sInByZWZlcnJlZF91c2VybmFtZSI6ImtycmlzaGRoMiIsImdpdmVuX25hbWUiOiJLcmlzaCIsImZhbWlseV9uYW1lIjoiRGhvbGFraWEiLCJlbWFpbCI6ImtycmlzaGRob2xha2lhMkBnbWFpbC5jb20ifQ.Fu2ErZhnfez-bhn_XmjkywcFdZHcFUSvzIzfdNiEowdA0soLmCyqf9731amP6m68shd9qk11e0mQhxFIAIxZPojViC1Csc9TBXLRRQ8ESMd6gPIj-DBkKVkQSZLJ1uibsh4Oo2RViGtqWVcEt32T8U_xhGdtdzNkJ8qy_e0fdNDsUnhmSaTQvmZJYarW0roIrkC-zYZrX3fftzbQfavSu9eqdfPf6wUttIrkaWThWUuORy-xaeZfSmvsGbEg027hh6QwlChiZTSF8R6bRxoqfPN3ZaGFFgbBXNRYZA_eYi2IevhIwJHi_r4o1UvtKAJyfPefm-M6hCfkN_6da4zsog",
"token": "fake.jwt.token",
}
# Create request
@ -2282,3 +2269,181 @@ async def test_non_object_metadata_is_served_like_unset_mappings_whatever_the_ba
)
assert mapped == unset
assert mapped[1] == "usecase-deployment"
@pytest.fixture
async def cached_chat_session(monkeypatch: pytest.MonkeyPatch) -> UserAPIKeyAuth:
from typing import cast
from litellm.constants import UI_SESSION_TOKEN_TEAM_ID
from litellm.proxy._types import LiteLLM_TeamMembership, LiteLLM_TeamTableCachedObj
from litellm.proxy.common_utils.user_api_key_cache import (
UserApiKeyCache,
team_membership_auth_cache_key,
team_membership_reservation_cache_key,
)
from litellm.proxy.utils import PrismaClient, ProxyLogging
cache: Final = UserApiKeyCache()
router: Final = litellm.Router(
model_list=[
{
"model_name": name,
"litellm_params": {"model": "openai/gpt-4.1-mini", "api_key": "test-key"},
}
for name in ("personal", "research", "clinical", "shared", "forbidden")
]
)
user: Final = LiteLLM_UserTable(
user_id="researcher",
user_role=LitellmUserRoles.INTERNAL_USER_VIEW_ONLY,
models=["personal"],
teams=["research-team", "clinical-team"],
max_budget=10,
)
await cache.async_set_cache(key=user.user_id, value=user, model_type=LiteLLM_UserTable)
for team in (
LiteLLM_TeamTableCachedObj(
team_id="research-team", team_alias="Research", models=["research", "shared"], max_budget=20
),
LiteLLM_TeamTableCachedObj(
team_id="clinical-team", team_alias="Clinical", models=["clinical", "shared"], max_budget=30
),
):
member: Final = LiteLLM_TeamMembership(user_id=user.user_id, team_id=team.team_id, spend=1)
await cache.async_set_cache(key=f"team_id:{team.team_id}", value=team, model_type=LiteLLM_TeamTableCachedObj)
await cache.async_set_cache(key=team_membership_reservation_cache_key(user.user_id, team.team_id), value=member)
await cache.async_set_cache(
key=team_membership_auth_cache_key(team_id=team.team_id, user_id=user.user_id),
value=member,
model_type=LiteLLM_TeamMembership,
)
monkeypatch.setattr(litellm.proxy.proxy_server, "prisma_client", cast(PrismaClient, object()))
monkeypatch.setattr(litellm.proxy.proxy_server, "user_api_key_cache", cache)
monkeypatch.setattr(litellm.proxy.proxy_server, "proxy_logging_obj", ProxyLogging(cache))
monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", router)
monkeypatch.setattr(litellm.proxy.proxy_server, "llm_model_list", router.model_list)
monkeypatch.setattr(litellm.proxy.proxy_server, "general_settings", {})
monkeypatch.setattr(litellm, "max_budget", 0)
return UserAPIKeyAuth(
token="session-hash",
user_id=user.user_id,
team_id=UI_SESSION_TOKEN_TEAM_ID,
user_role=LitellmUserRoles.INTERNAL_USER_VIEW_ONLY,
models=["personal"],
max_budget=3,
)
async def _validated_chat_request(
auth: UserAPIKeyAuth, request_data: dict[str, object], selected_team_id: str | None = None
) -> UserAPIKeyAuth:
from datetime import datetime, timezone
from fastapi import Request as FastAPIRequest
from litellm.proxy.auth.user_api_key_auth import validate_resolved_virtual_key
request: Final = FastAPIRequest(
{
"type": "http",
"method": "POST",
"path": "/v1/chat/completions",
"headers": ([(b"x-litellm-team-id", selected_team_id.encode())] if selected_team_id is not None else []),
"query_string": b"",
}
)
return await validate_resolved_virtual_key(
request=request,
request_data=request_data,
valid_token=auth,
api_key="session-hash",
route="/v1/chat/completions",
start_time=datetime(2026, 1, 1, tzinfo=timezone.utc),
parent_otel_span=None,
end_user_id=None,
end_user_params={},
_end_user_object=None,
)
@pytest.mark.asyncio
@pytest.mark.parametrize("selected_team_id", ["clinical-team", None], ids=["selected", "unique-fallback-grant"])
async def test_chat_session_request_binds_selected_team_and_preserves_user_and_key_budgets(
cached_chat_session: UserAPIKeyAuth,
selected_team_id: str | None,
) -> None:
from litellm.constants import UI_SESSION_TOKEN_TEAM_ID
authorized: Final = await _validated_chat_request(
cached_chat_session, {"model": "shared", "fallbacks": ["clinical"]}, selected_team_id
)
assert (
authorized.user_id,
authorized.team_id,
authorized.team_alias,
authorized.max_budget,
authorized.user_max_budget,
authorized.team_max_budget,
authorized.team_member_spend,
authorized.via_virtual_key,
) == ("researcher", "clinical-team", "Clinical", 3, 10, 30, 1, True)
assert (cached_chat_session.team_id, cached_chat_session.models) == (UI_SESSION_TOKEN_TEAM_ID, ["personal"])
@pytest.mark.asyncio
@pytest.mark.parametrize("fallback", ["research", {"model": "research"}])
async def test_chat_session_request_cannot_use_another_teams_fallback(
cached_chat_session: UserAPIKeyAuth, fallback: str | dict[str, str]
) -> None:
from litellm.proxy._types import ProxyException
with pytest.raises(ProxyException) as denied:
await _validated_chat_request(
cached_chat_session, {"model": "shared", "fallbacks": [fallback]}, "clinical-team"
)
assert (denied.value.code, denied.value.type, denied.value.param) == (
"403",
"team_model_access_denied",
"x-litellm-team-id",
)
@pytest.mark.asyncio
async def test_chat_session_request_requires_team_selection_for_shared_model(
cached_chat_session: UserAPIKeyAuth,
) -> None:
from litellm.proxy._types import ProxyException
with pytest.raises(ProxyException) as denied:
await _validated_chat_request(cached_chat_session, {"model": "shared"})
assert (denied.value.code, denied.value.type) == ("409", "team_selection_required")
@pytest.mark.asyncio
async def test_chat_session_request_still_enforces_original_key_budget_after_team_selection(
cached_chat_session: UserAPIKeyAuth,
) -> None:
exhausted: Final = cached_chat_session.model_copy(update={"spend": 4})
with pytest.raises(litellm.BudgetExceededError) as denied:
await _validated_chat_request(exhausted, {"model": "shared"}, "clinical-team")
assert (denied.value.current_cost, denied.value.max_budget) == (4, 3)
@pytest.mark.asyncio
async def test_chat_team_header_cannot_grant_an_ordinary_key_access_to_a_team_fallback(
cached_chat_session: UserAPIKeyAuth,
) -> None:
from litellm.proxy._types import ProxyException
ordinary_key: Final = cached_chat_session.model_copy(update={"team_id": None})
with pytest.raises(ProxyException) as denied:
await _validated_chat_request(ordinary_key, {"model": "personal", "fallbacks": ["clinical"]}, "clinical-team")
assert (denied.value.code, denied.value.type, denied.value.param) == ("403", "key_model_access_denied", "model")

63
ui/litellm-dashboard/src/app/chat/page.integration.test.tsx
View file

@ -1,6 +1,8 @@
import { fireEvent, screen, waitFor, within } from "@testing-library/react";
import { act, fireEvent, screen, waitFor, within } from "@testing-library/react";
import userEvent from "@testing-library/user-event";
import { beforeEach, describe, expect, it, vi } from "vitest";
import { useChatHistory } from "@/components/chat/useChatHistory";
import { fetchAvailableModels } from "@/components/llm_calls/fetch_models";
import ChatConversationPage from "./page";
import { renderWithProviders } from "@/../tests/test-utils";
import type { OnUrlUpdateFunction } from "nuqs/adapters/testing";
@ -64,10 +66,13 @@ vi.mock("@/contexts/ChatShellContext", () => ({
},
}));
const ONE_TURN_ARG_COUNT = 25;
const ONE_TURN_ARG_COUNT = 26;
const ON_TIMING_DATA_INDEX = 7;
const ON_USAGE_DATA_INDEX = 8;
const ON_TOTAL_LATENCY_INDEX = 24;
const CUSTOM_HEADERS_INDEX = 25;
const ON_RESPONSE_ID_INDEX = 15;
const PREVIOUS_RESPONSE_ID_INDEX = 14;
async function sendOneMessage(onUrlUpdate?: OnUrlUpdateFunction): Promise<void> {
renderWithProviders(<ChatConversationPage />, { onUrlUpdate });
@ -83,6 +88,7 @@ describe("/ui/chat request metrics", () => {
beforeEach(() => {
localStorage.clear();
mockMakeOpenAIResponsesRequest.mockReset();
vi.mocked(fetchAvailableModels).mockResolvedValue([{ model_group: "gpt-5.4-mini" }]);
shellState.storageUnavailable = false;
});
@ -109,7 +115,7 @@ describe("/ui/chat request metrics", () => {
expect(screen.getByLabelText("Cost: $0.000063")).toBeInTheDocument();
});
it("supplies the timing, usage and latency callbacks at the positional slots the Responses helper reads", async () => {
it("should preserve metric callback positions and omit a team header for a direct model grant", async () => {
mockMakeOpenAIResponsesRequest.mockResolvedValue(undefined);
await sendOneMessage();
@ -119,6 +125,57 @@ describe("/ui/chat request metrics", () => {
expect(typeof call[ON_TIMING_DATA_INDEX]).toBe("function");
expect(typeof call[ON_USAGE_DATA_INDEX]).toBe("function");
expect(typeof call[ON_TOTAL_LATENCY_INDEX]).toBe("function");
expect(call[CUSTOM_HEADERS_INDEX]).toBeUndefined();
});
it("should ignore a late old-team response ID after selecting a different billing team", async () => {
const user = userEvent.setup();
vi.mocked(fetchAvailableModels).mockResolvedValue([
{
model_group: "gpt-5.4-mini",
available_teams: [
{ team_id: "research", team_alias: "Research" },
{ team_id: "engineering", team_alias: "Engineering" },
],
},
]);
mockMakeOpenAIResponsesRequest.mockImplementation(async (...args: unknown[]) => {
const turn = mockMakeOpenAIResponsesRequest.mock.calls.length;
const updateText = args[1] as (role: string, delta: string) => void;
const onResponseId = args[ON_RESPONSE_ID_INDEX] as (id: string) => void;
updateText("assistant", `Reply ${turn}`);
onResponseId(`response-${turn}`);
});
renderWithProviders(<ChatConversationPage />);
await screen.findByRole("button", { name: "gpt-5.4-mini" });
await user.click(screen.getByRole("combobox", { name: "Team paying for this request" }));
await user.click(await screen.findByRole("option", { name: "Research" }));
async function send(text: string) {
act(() =>
fireEvent.change(screen.getByRole("textbox", { name: "Message", exact: true }), {
target: { value: text },
}),
);
await user.click(screen.getByRole("button", { name: "Send", exact: true }));
}
await send("First question");
await screen.findByText("Reply 1");
await waitFor(() => expect(screen.getByRole("combobox", { name: "Team paying for this request" })).toBeEnabled());
const oldResponseId = mockMakeOpenAIResponsesRequest.mock.calls[0][ON_RESPONSE_ID_INDEX] as (id: string) => void;
await user.click(screen.getByRole("combobox", { name: "Team paying for this request" }));
await user.click(await screen.findByRole("option", { name: "Engineering" }));
act(() => oldResponseId("late-research-response"));
await send("Continue with Engineering");
await screen.findByText("Reply 2");
const engineeringRequest = mockMakeOpenAIResponsesRequest.mock.calls[1];
expect(engineeringRequest[PREVIOUS_RESPONSE_ID_INDEX]).toBeNull();
expect(engineeringRequest[CUSTOM_HEADERS_INDEX]).toEqual({ "x-litellm-team-id": "engineering" });
await send("Continue again");
await screen.findByText("Reply 3");
expect(mockMakeOpenAIResponsesRequest.mock.calls[2][PREVIOUS_RESPONSE_ID_INDEX]).toBe("response-2");
});
it("puts the new conversation ID in the URL after the first message is sent", async () => {

177
ui/litellm-dashboard/src/app/chat/page.teams.integration.test.tsx Normal file
View file

@ -0,0 +1,177 @@
import { act, fireEvent, screen, waitFor } from "@testing-library/react";
import userEvent from "@testing-library/user-event";
import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
import { renderWithProviders, testQueryClient } from "@/../tests/test-utils";
import { ChatShellProvider } from "@/contexts/ChatShellContext";
import type { ModelGroup } from "@/components/llm_calls/fetch_models";
import ChatConversationPage from "./page";
vi.mock("next/navigation", () => ({ useRouter: () => ({ push: vi.fn(), replace: vi.fn() }) }));
const teams = [
{ team_id: "research", team_alias: "Research" },
{ team_id: "engineering", team_alias: "Engineering" },
];
const requests: { path: string; teamId: string | null; body: Record<string, unknown> }[] = [];
let models: ModelGroup[];
function renderChat() {
return renderWithProviders(
<ChatShellProvider
accessToken="session-token"
userId="team-user"
userEmail="member@example.com"
userRole="Internal User"
premiumUser={false}
>
<ChatConversationPage />
</ChatShellProvider>,
);
}
function enterPrompt(text = "Hello") {
act(() =>
fireEvent.change(screen.getByRole("textbox", { name: "Message", exact: true }), { target: { value: text } }),
);
}
beforeEach(() => {
localStorage.clear();
testQueryClient.clear();
requests.length = 0;
models = [{ model_group: "shared-chat", mode: "chat", providers: ["openai"], available_teams: teams }];
vi.stubGlobal(
"fetch",
vi.fn(async (input: RequestInfo | URL, init?: RequestInit) => {
const request =
input instanceof Request ? input : new Request(new URL(String(input), window.location.origin), init);
const path = new URL(request.url).pathname;
if (path.endsWith("/model_group/info")) return Response.json({ data: models });
if (path.endsWith("/responses") || path.endsWith("/images/generations")) {
requests.push({
path,
teamId: request.headers.get("x-litellm-team-id"),
body: JSON.parse(await request.text()),
});
if (path.endsWith("/images/generations"))
return Response.json({ created: 1, data: [{ b64_json: "aGVsbG8=" }] });
const events = [
{ type: "response.output_text.delta", delta: `Reply ${requests.length}` },
{ type: "response.completed", response: { id: `response-${requests.length}`, output: [] } },
];
return new Response(
events.map((event) => `event: ${event.type}\ndata: ${JSON.stringify(event)}\n\n`).join(""),
{ headers: { "Content-Type": "text/event-stream" } },
);
}
return Response.json({}, { status: 404 });
}),
);
});
afterEach(() => {
vi.unstubAllGlobals();
localStorage.clear();
});
describe("Chat session team selection", () => {
it("should render the model and require a billing choice only when multiple teams grant access", async () => {
renderChat();
await screen.findByRole("button", { name: "shared-chat" });
expect(screen.getByRole("combobox", { name: "Team paying for this request" })).toHaveTextContent("Choose team");
enterPrompt();
expect(screen.getByRole("button", { name: "Send", exact: true })).toBeDisabled();
expect(requests).toHaveLength(0);
});
it.each([
{ allowed: [], teamId: null },
{ allowed: [teams[0]], teamId: "research" },
])("should submit without a team selector when access is unambiguous ($teamId)", async ({ allowed, teamId }) => {
models[0].available_teams = allowed;
const user = userEvent.setup();
renderChat();
await screen.findByRole("button", { name: "shared-chat" });
expect(screen.queryByRole("combobox", { name: "Team paying for this request" })).not.toBeInTheDocument();
enterPrompt();
await user.click(screen.getByRole("button", { name: "Send", exact: true }));
await screen.findByText("Reply 1");
expect(requests[0].teamId).toBe(teamId);
});
it("should charge the selected team and replay history when switching teams", async () => {
const user = userEvent.setup();
renderChat();
await screen.findByRole("button", { name: "shared-chat" });
await user.click(screen.getByRole("combobox", { name: "Team paying for this request" }));
await user.click(await screen.findByRole("option", { name: "Research" }));
enterPrompt("First question");
await user.click(screen.getByRole("button", { name: "Send", exact: true }));
await screen.findByText("Reply 1");
await user.click(screen.getByRole("combobox", { name: "Team paying for this request" }));
await user.click(await screen.findByRole("option", { name: "Engineering" }));
enterPrompt("Follow up");
await user.click(screen.getByRole("button", { name: "Send", exact: true }));
await screen.findByText("Reply 2");
expect(requests.map((request) => request.teamId)).toEqual(["research", "engineering"]);
expect(requests[1].body).not.toHaveProperty("previous_response_id");
expect(requests[1].body.input).toEqual(
expect.arrayContaining([expect.objectContaining({ role: "user", content: "First question" })]),
);
});
it("should send image generation to the chosen team and clear the choice when changing models", async () => {
models.push({ model_group: "studio-image", mode: "image_generation", available_teams: teams });
const user = userEvent.setup();
renderChat();
await screen.findByRole("button", { name: "shared-chat" });
await user.click(screen.getByRole("combobox", { name: "Team paying for this request" }));
await user.click(await screen.findByRole("option", { name: "Research" }));
await user.click(screen.getByRole("button", { name: "shared-chat", exact: true }));
await user.click(await screen.findByRole("button", { name: /^studio-image/ }));
enterPrompt("Draw a circle");
expect(screen.getByRole("button", { name: "Generate image", exact: true })).toBeDisabled();
await user.click(screen.getByRole("combobox", { name: "Team paying for this request" }));
await user.click(await screen.findByRole("option", { name: "Engineering" }));
await user.click(screen.getByRole("button", { name: "Generate image", exact: true }));
await waitFor(() => expect(requests).toHaveLength(1));
expect(requests[0]).toMatchObject({
teamId: "engineering",
body: { model: "studio-image", prompt: "Draw a circle" },
});
});
it("should retain uploaded documents and replay them under the newly selected team", async () => {
const user = userEvent.setup();
renderChat();
await screen.findByRole("button", { name: "shared-chat" });
const file = new File(["team document"], "brief.txt", { type: "text/plain" });
await user.upload(screen.getByLabelText("Upload files"), file);
await screen.findByRole("button", { name: "Remove brief.txt" });
enterPrompt("Read the file");
expect(screen.getByRole("button", { name: "Send", exact: true })).toBeDisabled();
await user.click(screen.getByRole("combobox", { name: "Team paying for this request" }));
await user.click(await screen.findByRole("option", { name: "Research" }));
await user.click(screen.getByRole("button", { name: "Send", exact: true }));
await screen.findByText("Reply 1");
await user.click(screen.getByRole("combobox", { name: "Team paying for this request" }));
await user.click(await screen.findByRole("option", { name: "Engineering" }));
enterPrompt("Use that file again");
await user.click(screen.getByRole("button", { name: "Send", exact: true }));
await screen.findByText("Reply 2");
const document = {
type: "input_file",
filename: "brief.txt",
file_data: "data:text/plain;base64,dGVhbSBkb2N1bWVudA==",
};
const originalMessage = expect.objectContaining({
role: "user",
content: expect.arrayContaining([document]),
});
expect(requests.map((request) => request.teamId)).toEqual(["research", "engineering"]);
expect(requests[0].body.input).toEqual(expect.arrayContaining([originalMessage]));
expect(requests[1].body).not.toHaveProperty("previous_response_id");
expect(requests[1].body.input).toEqual(expect.arrayContaining([originalMessage]));
});
});

79
ui/litellm-dashboard/src/app/chat/page.tsx
View file

@ -12,7 +12,13 @@ import ChatMessages from "@/components/chat/ChatMessages";
import MCPConnectPicker from "@/components/chat/MCPConnectPicker";
import { fetchAvailableModels, type ModelGroup } from "@/components/llm_calls/fetch_models";
import type { AssistantMessageUpdate, ChatAttachment, ChatMessage } from "@/components/chat/types";
import ChatModelPicker, { isChatModel, modelChatMode, type ChatMode } from "@/components/chat/ChatModelPicker";
import ChatModelPicker, {
ChatTeamPicker,
chatTeamSelection,
isChatModel,
modelChatMode,
type ChatMode,
} from "@/components/chat/ChatModelPicker";
import MoyaiChatComposer from "@/components/chat/moyai/MoyaiChatComposer";
import MoyaiChatHome from "@/components/chat/moyai/MoyaiChatHome";
import { runChatRequest } from "@/components/chat/runChatRequest";
@ -88,6 +94,7 @@ export default function ChatConversationPage() {
} = useChatShell();
const [selectedModel, setSelectedModel] = useState<string | null>(null);
const [selectedTeamId, setSelectedTeamId] = useState<string | null>(null);
const [models, setModels] = useState<ModelGroup[]>([]);
const [mode, setMode] = useState<ChatMode>("chat");
const [attachments, setAttachments] = useState<ChatAttachment[]>([]);
@ -179,6 +186,7 @@ export default function ChatConversationPage() {
}
responsesSessionGenerationRef.current += 1;
setSelectedModel(model.model_group);
setSelectedTeamId(null);
setMode(modelChatMode(model));
setResponsesSessionId(null);
setAttachmentError(null);
@ -241,6 +249,14 @@ export default function ChatConversationPage() {
const model = generationModel ?? selectedModel;
if (!hasMessageContent(trimmed, images, requestMode) || !model) return;
if (isStreaming || abortControllerRef.current || readAttachmentsRef.current) return;
const teamSelection = chatTeamSelection(
models.find((item) => item.model_group === model),
selectedTeamId,
);
if (teamSelection.needsSelection) {
setAttachmentError("Choose which team should pay for this request.");
return;
}
const sessionGeneration = responsesSessionGenerationRef.current;
setInputText("");
@ -294,6 +310,7 @@ export default function ChatConversationPage() {
mode: requestMode,
prompt: trimmed,
accessToken,
teamId: teamSelection.teamId,
signal: requestController.signal,
messages: history,
previousResponseId,
@ -317,6 +334,8 @@ export default function ChatConversationPage() {
activeConversationId,
activeConversation,
selectedModel,
selectedTeamId,
models,
selectedMCPServers,
mode,
attachments,
@ -349,11 +368,31 @@ export default function ChatConversationPage() {
return;
}
if (!newContent.trim() && !original.attachments?.length) return;
const requestModel = original.imageGeneration?.model ?? selectedModel;
if (
chatTeamSelection(
models.find((item) => item.model_group === requestModel),
selectedTeamId,
).needsSelection
) {
setAttachmentError("Choose a team for this model before resending the message.");
return;
}
const priorMessages = msgs.slice(0, idx);
truncateFromMessage(activeConversationId, messageId);
handleSend(newContent, priorMessages, original.attachments ?? [], original.imageGeneration?.model);
},
[activeConversationId, isStreaming, activeConversation, truncateFromMessage, handleSend, mode, selectedModel],
[
activeConversationId,
isStreaming,
activeConversation,
truncateFromMessage,
handleSend,
mode,
selectedModel,
models,
selectedTeamId,
],
);
// Track scroll position to show/hide scroll-to-bottom button
@ -403,9 +442,21 @@ export default function ChatConversationPage() {
const showBlankState = !activeConversation || activeConversation.messages.length === 0;
const selectTeam = (teamId: string | null) => {
responsesSessionGenerationRef.current += 1;
setSelectedTeamId(teamId);
setResponsesSessionId(null);
setAttachmentError(null);
};
const isModelSelectionDisabled = isStreaming || isReadingAttachments;
const inputBar = (inConversation: boolean) => {
const preparing = isLoadingModels || isReadingAttachments;
const submitDisabled = !hasMessageContent(inputText, attachments, mode) || !selectedModel || preparing;
const model = models.find((item) => item.model_group === selectedModel);
const teamSelection = chatTeamSelection(model, selectedTeamId);
const canSubmit = hasMessageContent(inputText, attachments, mode) && selectedModel;
const submitDisabled = !canSubmit || preparing || teamSelection.needsSelection;
return (
<div
onDragOver={(event) => {
@ -449,13 +500,21 @@ export default function ChatConversationPage() {
/>
}
modelPicker={
<ChatModelPicker
models={models}
selectedModel={selectedModel}
loading={isLoadingModels}
disabled={isStreaming || isReadingAttachments}
onSelect={selectModel}
/>
<div className="flex min-w-0 flex-wrap items-center justify-end gap-1">
<ChatTeamPicker
model={model}
selectedTeamId={selectedTeamId}
disabled={isModelSelectionDisabled}
onSelect={selectTeam}
/>
<ChatModelPicker
models={models}
selectedModel={selectedModel}
loading={isLoadingModels}
disabled={isModelSelectionDisabled}
onSelect={selectModel}
/>
</div>
}
tools={
mode === "chat" ? (

46
ui/litellm-dashboard/src/components/chat/ChatModelPicker.tsx
View file

@ -4,6 +4,7 @@ import { Button } from "@/components/ui/button";
import { Input } from "@/components/ui/input";
import { Popover, PopoverContent, PopoverTrigger } from "@/components/ui/popover";
import { ScrollArea } from "@/components/ui/scroll-area";
import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select";
import { Skeleton } from "@/components/ui/skeleton";
import type { ModelGroup } from "@/components/llm_calls/fetch_models";
@ -17,6 +18,51 @@ export function isChatModel(model: ModelGroup): boolean {
return !model.mode || ["chat", "responses", "completion", "image_generation"].includes(model.mode);
}
export function chatTeamSelection(
model: ModelGroup | undefined,
selectedTeamId: string | null,
): { teamId: string | undefined; needsSelection: boolean } {
const teams = model?.available_teams ?? [];
const team = teams.length === 1 ? teams[0] : teams.find((item) => item.team_id === selectedTeamId);
return { teamId: team?.team_id, needsSelection: teams.length > 1 && !team };
}
export function ChatTeamPicker({
model,
selectedTeamId,
disabled,
onSelect,
}: {
model: ModelGroup | undefined;
selectedTeamId: string | null;
disabled: boolean;
onSelect: (teamId: string | null) => void;
}) {
const teams = model?.available_teams ?? [];
if (teams.length < 2) return null;
const selection = chatTeamSelection(model, selectedTeamId);
return (
<Select
value={selection.teamId ?? null}
items={teams.map((team) => ({ value: team.team_id, label: team.team_alias || team.team_id }))}
disabled={disabled}
onValueChange={onSelect}
>
<SelectTrigger aria-label="Team paying for this request" size="sm" className="max-w-40">
<SelectValue placeholder="Choose team" />
</SelectTrigger>
<SelectContent className="chat-theme w-max min-w-64 max-w-[calc(100vw-2rem)]">
{teams.map((team) => (
<SelectItem key={team.team_id} value={team.team_id}>
{team.team_alias || team.team_id}
</SelectItem>
))}
</SelectContent>
</Select>
);
}
export default function ChatModelPicker({
models,
selectedModel,

10
ui/litellm-dashboard/src/components/chat/imageHelpers.ts
View file

@ -42,13 +42,21 @@ export async function generateChatImages({
model,
accessToken,
signal,
teamId,
}: {
prompt: string;
model: string;
accessToken: string;
signal?: AbortSignal;
teamId?: string;
}): Promise<ChatImageAttachment[]> {
const client = new OpenAI({ apiKey: accessToken, baseURL: getGatewayBaseUrl(), dangerouslyAllowBrowser: true });
const clientOptions = {
apiKey: accessToken,
baseURL: getGatewayBaseUrl(),
dangerouslyAllowBrowser: true,
defaultHeaders: teamId ? { "x-litellm-team-id": teamId } : undefined,
};
const client = new OpenAI(clientOptions);
const response = await client.images.generate({ model, prompt }, { signal });
const mimeType = `image/${response.output_format ?? "png"}`;
const images = (response.data ?? []).flatMap((image, index): ChatImageAttachment[] => {

5
ui/litellm-dashboard/src/components/chat/runChatRequest.ts
View file

@ -9,6 +9,7 @@ interface ChatRequestOptions {
mode: "chat" | "image";
prompt: string;
accessToken: string;
teamId?: string;
signal: AbortSignal;
messages: Pick<
ChatMessage,
@ -40,6 +41,7 @@ export async function runChatRequest({
mode,
prompt,
accessToken,
teamId,
signal,
messages,
previousResponseId,
@ -54,7 +56,7 @@ export async function runChatRequest({
try {
signal.throwIfAborted();
if (mode === "image") {
const imageRequest = { prompt, model, accessToken, signal };
const imageRequest = { prompt, model, accessToken, signal, teamId };
const generatedImages = await generateChatImages(imageRequest);
signal.throwIfAborted();
onUpdate({ generatedImages });
@ -94,6 +96,7 @@ export async function runChatRequest({
undefined,
true,
(totalLatency) => onUpdate({ totalLatency }),
teamId ? { "x-litellm-team-id": teamId } : undefined,
);
signal.throwIfAborted();
} catch (error: unknown) {

6
ui/litellm-dashboard/src/components/llm_calls/fetch_models.tsx
View file

@ -2,9 +2,13 @@
import { excludeProxyWideSentinel } from "@/components/key_team_helpers/fetch_available_models_team_key";
import { modelAvailableCall, modelHubCall } from "@/components/networking";
import type { components } from "@/lib/http/schema";
export type ModelTeam = components["schemas"]["ModelGroupTeam"];
export interface ModelGroup {
model_group: string;
available_teams?: ModelTeam[];
providers?: string[];
mode?: string;
supports_reasoning?: boolean;
@ -14,6 +18,7 @@ export interface ModelGroup {
interface AvailableModel {
model_group?: string | null;
available_teams?: ModelTeam[] | null;
model_name?: string | null;
id?: string | null;
providers?: string[] | null;
@ -27,6 +32,7 @@ const toModelGroup = (item: AvailableModel): ModelGroup => {
const groupName = (item.model_group || item.id || item.model_name) ?? "";
return {
model_group: groupName,
...(item.available_teams && { available_teams: item.available_teams }),
...(item.providers && { providers: item.providers }),
...(item.mode && { mode: item.mode }),
...(item.supports_reasoning === true && { supports_reasoning: true }),

20
ui/litellm-dashboard/src/lib/http/schema.d.ts generated vendored
View file

@ -8325,7 +8325,7 @@ export interface paths {
* Example:
* ```bash
* curl --location 'http://0.0.0.0:4000/key/block' --header "Authorization: Bearer $LITELLM_MASTER_KEY" --header 'Content-Type: application/json' --data '{
* "key": "sk-Fn8Ej39NxjAXrvpUGKghGw"
* "key": "YOUR_VIRTUAL_KEY"
* }'
* ```
*
@ -8884,7 +8884,7 @@ export interface paths {
* Example:
* ```bash
* curl --location 'http://0.0.0.0:4000/key/unblock' --header "Authorization: Bearer $LITELLM_MASTER_KEY" --header 'Content-Type: application/json' --data '{
* "key": "sk-Fn8Ej39NxjAXrvpUGKghGw"
* "key": "YOUR_VIRTUAL_KEY"
* }'
* ```
*
@ -38426,6 +38426,8 @@ export interface components {
};
/** ModelGroupInfoProxy */
ModelGroupInfoProxy: {
/** Available Teams */
available_teams?: components["schemas"]["ModelGroupTeam"][] | null;
/** Configurable Clientside Auth Params */
configurable_clientside_auth_params?: (string | components["schemas"]["ConfigurableClientsideParamsCustomAuth-Output"])[] | null;
/** Custom Metadata */
@ -38519,6 +38521,18 @@ export interface components {
/** Tpm */
tpm?: number | null;
};
/** ModelGroupInfoResponse */
ModelGroupInfoResponse: {
/** Data */
data: components["schemas"]["ModelGroupInfoProxy"][];
};
/** ModelGroupTeam */
ModelGroupTeam: {
/** Team Alias */
team_alias?: string | null;
/** Team Id */
team_id: string;
};
/** ModelInfoDelete */
ModelInfoDelete: {
/** Id */
@ -65757,7 +65771,7 @@ export interface operations {
[name: string]: unknown;
};
content: {
"application/json": unknown;
"application/json": components["schemas"]["ModelGroupInfoResponse"];
};
};
/** @description Validation Error */