diff --git a/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.58-py3-none-any.whl b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.58-py3-none-any.whl new file mode 100644 index 0000000000..756080f854 Binary files /dev/null and b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.58-py3-none-any.whl differ diff --git a/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.58.tar.gz b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.58.tar.gz new file mode 100644 index 0000000000..8a52b05840 Binary files /dev/null and b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.58.tar.gz differ diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260318140652_add_index_to_team_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260318140652_add_index_to_team_table/migration.sql new file mode 100644 index 0000000000..494aaf6238 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260318140652_add_index_to_team_table/migration.sql @@ -0,0 +1,9 @@ +-- CreateIndex +CREATE INDEX "LiteLLM_TeamTable_organization_id_idx" ON "LiteLLM_TeamTable"("organization_id"); + +-- CreateIndex +CREATE INDEX "LiteLLM_TeamTable_team_alias_idx" ON "LiteLLM_TeamTable"("team_alias"); + +-- CreateIndex +CREATE INDEX "LiteLLM_TeamTable_created_at_idx" ON "LiteLLM_TeamTable"("created_at"); + diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index ce79c2b3d5..a2c8329540 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -146,6 +146,10 @@ model LiteLLM_TeamTable { litellm_model_table LiteLLM_ModelTable? @relation(fields: [model_id], references: [id]) object_permission LiteLLM_ObjectPermissionTable? @relation(fields: [object_permission_id], references: [object_permission_id]) projects LiteLLM_ProjectTable[] + + @@index([organization_id]) + @@index([team_alias]) + @@index([created_at]) } // Projects sit between teams and keys for use-case management diff --git a/litellm-proxy-extras/pyproject.toml b/litellm-proxy-extras/pyproject.toml index 006aad9480..8aa9e70e43 100644 --- a/litellm-proxy-extras/pyproject.toml +++ b/litellm-proxy-extras/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "litellm-proxy-extras" -version = "0.4.57" +version = "0.4.58" description = "Additional files for the LiteLLM Proxy. Reduces the size of the main litellm package." authors = ["BerriAI"] readme = "README.md" @@ -22,7 +22,7 @@ requires = ["poetry-core"] build-backend = "poetry.core.masonry.api" [tool.commitizen] -version = "0.4.57" +version = "0.4.58" version_files = [ "pyproject.toml:version", "../requirements.txt:litellm-proxy-extras==", diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index ceb3d5d277..b033d3fce0 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -101,6 +101,7 @@ from litellm.types.proxy.management_endpoints.team_endpoints import ( BulkTeamMemberAddRequest, BulkTeamMemberAddResponse, GetTeamMemberPermissionsResponse, + TeamListItem, TeamListResponse, TeamMemberAddResult, UpdateTeamMemberPermissionsRequest, @@ -3212,6 +3213,40 @@ async def list_available_teams( return available_teams_correct_type +async def _get_org_admin_org_ids( + user_id: str, + prisma_client: Any, + user_api_key_cache: Any, + proxy_logging_obj: Any, +) -> Optional[List[str]]: + """ + Return the list of organization IDs where the user is an org admin. + Returns None if the user is not an org admin of any organization or if + the user cannot be found. + """ + try: + caller_user = await get_user_object( + user_id=user_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + user_id_upsert=False, + proxy_logging_obj=proxy_logging_obj, + ) + except ValueError: + # get_user_object raises ValueError when the user doesn't exist + return None + + if caller_user is None: + return None + + org_ids = [ + m.organization_id + for m in (caller_user.organization_memberships or []) + if m.user_role == LitellmUserRoles.ORG_ADMIN.value + ] + return org_ids if org_ids else None + + async def _build_team_list_where_conditions( prisma_client: PrismaClient, team_id: Optional[str], @@ -3219,8 +3254,16 @@ async def _build_team_list_where_conditions( organization_id: Optional[str], user_id: Optional[str], use_deleted_table: bool, -) -> Dict[str, Any]: - """Build where conditions for team list query.""" + org_admin_org_ids: Optional[List[str]] = None, + user_api_key_cache: Optional[Any] = None, + proxy_logging_obj: Optional[Any] = None, +) -> Optional[Dict[str, Any]]: + """ + Build where conditions for team list query. + + Returns None when the query is guaranteed to yield no results (e.g. user + has no team memberships), allowing the caller to skip the DB round-trip. + """ where_conditions: Dict[str, Any] = {} if team_id: @@ -3234,58 +3277,79 @@ async def _build_team_list_where_conditions( if organization_id: where_conditions["organization_id"] = organization_id + elif org_admin_org_ids is not None and not user_id: + # Org admin without explicit org or user filter: scope to their orgs. + # NOTE: when user_id is provided, no org filter is applied — the + # query returns all teams the target user belongs to across all + # organisations. This matches the legacy /team/list behaviour in + # _authorize_and_filter_teams which fetches direct-membership teams + # without an org constraint. + where_conditions["organization_id"] = {"in": org_admin_org_ids} if user_id: try: - user_object = await prisma_client.db.litellm_usertable.find_unique( - where={"user_id": user_id} + user_object_correct_type = await get_user_object( + user_id=user_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + user_id_upsert=False, + proxy_logging_obj=proxy_logging_obj, ) - except Exception: + except ValueError: raise HTTPException( status_code=404, detail={"error": f"User not found, passed user_id={user_id}"}, ) - if user_object is None: + if user_object_correct_type is None: raise HTTPException( status_code=404, detail={"error": f"User not found, passed user_id={user_id}"}, ) - user_object_correct_type = LiteLLM_UserTable(**user_object.model_dump()) + user_team_ids = user_object_correct_type.teams or [] if use_deleted_table: where_conditions["members"] = {"has": user_id} else: - if team_id is None: - where_conditions["team_id"] = {"in": user_object_correct_type.teams} - elif team_id in user_object_correct_type.teams: - where_conditions["team_id"] = team_id + # When user_id is provided, filter by that user's direct team + # memberships. For org admins the access control gate in + # list_team_v2 already verified the caller's authority — the + # filter logic is the same as for regular users. + if not user_team_ids: + return None # no memberships — skip the DB query + elif team_id is not None: + # team_id exact-match already in where_conditions; verify membership + if team_id not in user_team_ids: + raise HTTPException( + status_code=404, + detail={"error": f"User is not a member of team_id={team_id}"}, + ) else: - raise HTTPException( - status_code=404, - detail={"error": f"User is not a member of team_id={team_id}"}, - ) + where_conditions["team_id"] = {"in": user_team_ids} return where_conditions -def _convert_teams_to_response( - teams: List[Any], use_deleted_table: bool -) -> List[Union[LiteLLM_TeamTable, LiteLLM_DeletedTeamTable]]: - """Convert Prisma models to Pydantic models.""" - team_list: List[Union[LiteLLM_TeamTable, LiteLLM_DeletedTeamTable]] = [] - if teams: - for team in teams: - # Convert Prisma model to dict (supports both Pydantic v1 and v2) - try: - team_dict = team.model_dump() - except Exception: - # Fallback for Pydantic v1 compatibility - team_dict = team.dict() - if use_deleted_table: - # Use deleted team type to preserve deleted_at, deleted_by, etc. - team_list.append(LiteLLM_DeletedTeamTable(**team_dict)) - else: - team_list.append(LiteLLM_TeamTable(**team_dict)) +def _convert_teams_to_response_models( + teams: list, + use_deleted_table: bool, +) -> List[Union[TeamListItem, LiteLLM_TeamTable, LiteLLM_DeletedTeamTable]]: + """Convert raw Prisma team rows to response models.""" + team_list: List[Union[TeamListItem, LiteLLM_TeamTable, LiteLLM_DeletedTeamTable]] = [] + for team in teams: + try: + team_dict = team.model_dump() + except Exception: + team_dict = team.dict() + + if use_deleted_table: + team_list.append(LiteLLM_DeletedTeamTable(**team_dict)) + else: + members_with_roles = team_dict.get("members_with_roles") + if not isinstance(members_with_roles, list): + members_with_roles = [] + team_dict["members_with_roles"] = members_with_roles + members_count = len(members_with_roles) + team_list.append(TeamListItem(**team_dict, members_count=members_count)) return team_list @@ -3353,7 +3417,11 @@ async def list_team_v2( status: Optional[str] Filter by status. Currently supports "deleted" to query deleted teams. """ - from litellm.proxy.proxy_server import prisma_client + from litellm.proxy.proxy_server import ( + prisma_client, + proxy_logging_obj, + user_api_key_cache, + ) if prisma_client is None: raise HTTPException( @@ -3361,20 +3429,55 @@ async def list_team_v2( detail={"error": f"No db connected. prisma client={prisma_client}"}, ) - if not allowed_route_check_inside_route( - user_api_key_dict=user_api_key_dict, requested_user_id=user_id - ): - raise HTTPException( - status_code=401, - detail={ - "error": "Only admin users can query all teams/other teams. Your user role={}".format( - user_api_key_dict.user_role - ) - }, - ) + # --- Access control --- + # Proxy admins and admin viewers can query any teams. + # Org admins can query teams within their organizations. + # Regular users can only query their own teams. + is_proxy_admin = _user_has_admin_view(user_api_key_dict) + org_admin_org_ids: Optional[List[str]] = None - if user_id is None and user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: - user_id = user_api_key_dict.user_id + if not is_proxy_admin: + # Always check org admin status so that even own-queries see + # the full set of organisation teams, not just direct memberships. + if user_api_key_dict.user_id: + org_admin_org_ids = await _get_org_admin_org_ids( + user_id=user_api_key_dict.user_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + + if org_admin_org_ids is not None: + # Org admin: validate org_id filter if provided + if organization_id and organization_id not in org_admin_org_ids: + raise HTTPException( + status_code=403, + detail={ + "error": "You can only view teams within your organizations." + }, + ) + verbose_proxy_logger.debug( + "list_team_v2: org admin access for user=%s, org_ids=%s, user_id_filter=%s", + user_api_key_dict.user_id, + org_admin_org_ids, + user_id, + ) + else: + # Not an org admin — fall back to standard route check + if not allowed_route_check_inside_route( + user_api_key_dict=user_api_key_dict, requested_user_id=user_id + ): + raise HTTPException( + status_code=401, + detail={ + "error": "Only admin users can query all teams/other teams. Your user role={}".format( + user_api_key_dict.user_role + ) + }, + ) + # Regular user — auto-inject caller's user_id + if user_id is None: + user_id = user_api_key_dict.user_id if status is not None and status != "deleted": raise HTTPException( @@ -3389,7 +3492,8 @@ async def list_team_v2( # Calculate skip and take for pagination skip = (page - 1) * page_size - # Build where conditions based on provided parameters + # Build where conditions based on provided parameters. + # Returns None when the query is guaranteed to yield no results. where_conditions = await _build_team_list_where_conditions( prisma_client=prisma_client, team_id=team_id, @@ -3397,8 +3501,20 @@ async def list_team_v2( organization_id=organization_id, user_id=user_id, use_deleted_table=use_deleted_table, + org_admin_org_ids=org_admin_org_ids, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, ) + if where_conditions is None: + return { + "teams": [], + "total": 0, + "page": page, + "page_size": page_size, + "total_pages": 0, + } + # Build order_by conditions valid_sort_columns = ["team_id", "team_alias", "created_at"] order_by = None @@ -3434,8 +3550,8 @@ async def list_team_v2( # Calculate total pages total_pages = -(-total_count // page_size) # Ceiling division - # Convert Prisma models to Pydantic models, preserving deleted fields when applicable - team_list = _convert_teams_to_response(teams, use_deleted_table) + # Convert Prisma models to response models with members_count + team_list = _convert_teams_to_response_models(teams, use_deleted_table) return { "teams": team_list, diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index b68872e2ed..46be6b31e1 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -146,6 +146,10 @@ model LiteLLM_TeamTable { litellm_model_table LiteLLM_ModelTable? @relation(fields: [model_id], references: [id]) object_permission LiteLLM_ObjectPermissionTable? @relation(fields: [object_permission_id], references: [object_permission_id]) projects LiteLLM_ProjectTable[] + + @@index([organization_id]) + @@index([team_alias]) + @@index([created_at]) } // Projects sit between teams and keys for use-case management diff --git a/litellm/types/proxy/management_endpoints/team_endpoints.py b/litellm/types/proxy/management_endpoints/team_endpoints.py index 77816fa78c..5055a65783 100644 --- a/litellm/types/proxy/management_endpoints/team_endpoints.py +++ b/litellm/types/proxy/management_endpoints/team_endpoints.py @@ -43,10 +43,16 @@ class UpdateTeamMemberPermissionsRequest(BaseModel): team_member_permissions: List[str] +class TeamListItem(LiteLLM_TeamTable): + """A team item in the paginated list response, enriched with computed fields.""" + + members_count: int = 0 + + class TeamListResponse(BaseModel): """Response to get the list of teams""" - teams: List[Union[LiteLLM_TeamTable, LiteLLM_DeletedTeamTable]] + teams: List[Union[TeamListItem, LiteLLM_TeamTable, LiteLLM_DeletedTeamTable]] total: int page: int page_size: int diff --git a/poetry.lock b/poetry.lock index 86d5b1e4c9..e3b083d778 100644 --- a/poetry.lock +++ b/poetry.lock @@ -3222,15 +3222,15 @@ files = [ [[package]] name = "litellm-proxy-extras" -version = "0.4.57" +version = "0.4.58" description = "Additional files for the LiteLLM Proxy. Reduces the size of the main litellm package." optional = true python-versions = "!=2.7.*,!=3.0.*,!=3.1.*,!=3.2.*,!=3.3.*,!=3.4.*,!=3.5.*,!=3.6.*,!=3.7.*,>=3.8" groups = ["main"] markers = "extra == \"proxy\"" files = [ - {file = "litellm_proxy_extras-0.4.57-py3-none-any.whl", hash = "sha256:04538223cd80318a72d70c6e10f701598e58c763368296a6503c674c92fbdb62"}, - {file = "litellm_proxy_extras-0.4.57.tar.gz", hash = "sha256:ef9b95dc42237614216833bd5d46ebf9dea1caa5ea14ea1a66d7f7842b224ec2"}, + {file = "litellm_proxy_extras-0.4.58-py3-none-any.whl", hash = "sha256:8863e70de833c0e35119a1cbbf583619bdebe52222efd5654586519175ba403b"}, + {file = "litellm_proxy_extras-0.4.58.tar.gz", hash = "sha256:84a67483329eced8be4fc61c4e43f117287aa4e3deeb8ddf8fe8cdc9a8508836"}, ] [[package]] @@ -8018,4 +8018,4 @@ utils = ["numpydoc"] [metadata] lock-version = "2.1" python-versions = ">=3.9,<4.0" -content-hash = "0002021a7733b370a9855b26198e8e6dc49a62b67f1eface6f2fbe406ff3c3ac" +content-hash = "eda34dfd8b35474beffee18893d6782c7b3d0d3d2c610f66237eb97176f43527" diff --git a/pyproject.toml b/pyproject.toml index 55a825d831..3fb05cc8b3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -61,7 +61,7 @@ boto3 = { version = "^1.40.76", optional = true } redisvl = {version = "^0.4.1", optional = true, markers = "python_version >= '3.9' and python_version < '3.14'"} mcp = {version = ">=1.25.0,<2.0.0", optional = true, python = ">=3.10"} a2a-sdk = {version = "^0.3.22", optional = true, python = ">=3.10"} -litellm-proxy-extras = {version = "^0.4.57", optional = true} +litellm-proxy-extras = {version = "^0.4.58", optional = true} rich = {version = "^13.7.1", optional = true} litellm-enterprise = {version = "^0.1.33", optional = true} diskcache = {version = "^5.6.1", optional = true} diff --git a/requirements.txt b/requirements.txt index 827986487f..d420f4ac60 100644 --- a/requirements.txt +++ b/requirements.txt @@ -57,7 +57,7 @@ grpcio>=1.75.0; python_version >= "3.14" sentry_sdk==2.21.0 # for sentry error handling detect-secrets==1.5.0 # Enterprise - secret detection / masking in LLM requests tzdata==2025.1 # IANA time zone database -litellm-proxy-extras==0.4.57 # for proxy extras - e.g. prisma migrations +litellm-proxy-extras==0.4.58 # for proxy extras - e.g. prisma migrations llm-sandbox==0.3.31 # for skill execution in sandbox ### LITELLM PACKAGE DEPENDENCIES python-dotenv==1.0.1 # for env diff --git a/schema.prisma b/schema.prisma index 939f1eb0f4..fde9a466a2 100644 --- a/schema.prisma +++ b/schema.prisma @@ -146,6 +146,10 @@ model LiteLLM_TeamTable { litellm_model_table LiteLLM_ModelTable? @relation(fields: [model_id], references: [id]) object_permission LiteLLM_ObjectPermissionTable? @relation(fields: [object_permission_id], references: [object_permission_id]) projects LiteLLM_ProjectTable[] + + @@index([organization_id]) + @@index([team_alias]) + @@index([created_at]) } // Projects sit between teams and keys for use-case management diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index 41a724c271..c325b0b6fc 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -2062,7 +2062,14 @@ async def test_list_team_v2_security_check_non_admin_user(): user_id="non_admin_user_123", ) - with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma_client: + with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma_client, \ + patch("litellm.proxy.proxy_server.user_api_key_cache"), \ + patch("litellm.proxy.proxy_server.proxy_logging_obj"), \ + patch( + "litellm.proxy.management_endpoints.team_endpoints.get_user_object", + new_callable=AsyncMock, + return_value=None, + ): mock_prisma_client.return_value = MagicMock() # Mock non-None prisma client # Should raise HTTPException with 401 status @@ -2103,7 +2110,14 @@ async def test_list_team_v2_security_check_non_admin_user_other_user(): user_id="non_admin_user_123", ) - with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma_client: + with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma_client, \ + patch("litellm.proxy.proxy_server.user_api_key_cache"), \ + patch("litellm.proxy.proxy_server.proxy_logging_obj"), \ + patch( + "litellm.proxy.management_endpoints.team_endpoints.get_user_object", + new_callable=AsyncMock, + return_value=None, + ): mock_prisma_client.return_value = MagicMock() # Mock non-None prisma client # Should raise HTTPException with 401 status @@ -2142,19 +2156,21 @@ async def test_list_team_v2_security_check_non_admin_user_own_teams(): user_id="non_admin_user_123", ) - with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma_client: + with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma_client, \ + patch("litellm.proxy.proxy_server.user_api_key_cache"), \ + patch("litellm.proxy.proxy_server.proxy_logging_obj"): # Mock prisma client and database operations mock_db = Mock() mock_prisma_client.db = mock_db - - # Mock user lookup - mock_user_object = Mock() - mock_user_object.model_dump.return_value = { - "user_id": "non_admin_user_123", - "teams": ["team_1", "team_2"], - } - mock_db.litellm_usertable.find_unique = AsyncMock(return_value=mock_user_object) - + + # Mock get_user_object to return a user with teams + from litellm.proxy._types import LiteLLM_UserTable + + mock_user = LiteLLM_UserTable( + user_id="non_admin_user_123", + teams=["team_1", "team_2"], + ) + # Mock team lookup mock_teams = [ Mock(model_dump=lambda: {"team_id": "team_1", "team_alias": "Team 1"}), @@ -2163,21 +2179,26 @@ async def test_list_team_v2_security_check_non_admin_user_own_teams(): mock_db.litellm_teamtable.find_many = AsyncMock(return_value=mock_teams) mock_db.litellm_teamtable.count = AsyncMock(return_value=2) - # Should NOT raise an exception - result = await list_team_v2( - http_request=mock_request, - user_id="non_admin_user_123", # Non-admin querying their own teams - user_api_key_dict=mock_user_api_key_dict_non_admin, - team_id=None, - page=1, - page_size=10, - status=None, - ) + with patch( + "litellm.proxy.management_endpoints.team_endpoints.get_user_object", + new_callable=AsyncMock, + return_value=mock_user, + ): + # Should NOT raise an exception + result = await list_team_v2( + http_request=mock_request, + user_id="non_admin_user_123", # Non-admin querying their own teams + user_api_key_dict=mock_user_api_key_dict_non_admin, + team_id=None, + page=1, + page_size=10, + status=None, + ) - # Should return results without error - assert "teams" in result - assert "total" in result - assert result["total"] == 2 + # Should return results without error + assert "teams" in result + assert "total" in result + assert result["total"] == 2 @pytest.mark.asyncio @@ -2293,27 +2314,280 @@ async def test_list_team_v2_with_status_deleted(): assert len(result["teams"]) == 2 +@pytest.mark.asyncio +async def test_list_team_v2_org_admin_sees_org_teams(): + """ + Test that an org admin (internal_user role with org_admin membership) + can list teams scoped to their organisations without getting a 401. + """ + from datetime import datetime + from unittest.mock import AsyncMock, Mock, patch + + from fastapi import Request + + from litellm.proxy._types import ( + LiteLLM_OrganizationMembershipTable, + LiteLLM_UserTable, + LitellmUserRoles, + UserAPIKeyAuth, + ) + from litellm.proxy.management_endpoints.team_endpoints import list_team_v2 + + mock_request = Mock(spec=Request) + mock_user_api_key_dict = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + user_id="org_admin_user", + ) + + mock_user = LiteLLM_UserTable( + user_id="org_admin_user", + teams=[], + organization_memberships=[ + LiteLLM_OrganizationMembershipTable( + user_id="org_admin_user", + organization_id="org_A", + user_role="org_admin", + spend=0.0, + created_at=datetime.now(), + updated_at=datetime.now(), + ), + ], + ) + + with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, \ + patch("litellm.proxy.proxy_server.user_api_key_cache"), \ + patch("litellm.proxy.proxy_server.proxy_logging_obj"), \ + patch( + "litellm.proxy.management_endpoints.team_endpoints.get_user_object", + new_callable=AsyncMock, + return_value=mock_user, + ): + mock_db = Mock() + mock_prisma.db = mock_db + + mock_team = Mock() + mock_team.model_dump.return_value = { + "team_id": "team_in_org_A", + "team_alias": "Org A Team", + "organization_id": "org_A", + "members_with_roles": [{"user_id": "u1", "role": "user"}], + } + mock_db.litellm_teamtable.find_many = AsyncMock(return_value=[mock_team]) + mock_db.litellm_teamtable.count = AsyncMock(return_value=1) + + result = await list_team_v2( + http_request=mock_request, + user_id=None, + organization_id=None, + team_id=None, + team_alias=None, + user_api_key_dict=mock_user_api_key_dict, + page=1, + page_size=10, + sort_by=None, + sort_order="asc", + status=None, + ) + + assert result["total"] == 1 + assert len(result["teams"]) == 1 + assert result["teams"][0].members_count == 1 + + # Verify org-scoped where clause + where = mock_db.litellm_teamtable.find_many.call_args.kwargs["where"] + assert where["organization_id"] == {"in": ["org_A"]} + + +@pytest.mark.asyncio +async def test_list_team_v2_org_admin_cannot_view_other_orgs(): + """ + Test that an org admin is rejected with 403 when filtering by an + organisation they do not administer. + """ + from datetime import datetime + from unittest.mock import AsyncMock, Mock, patch + + from fastapi import HTTPException, Request + + from litellm.proxy._types import ( + LiteLLM_OrganizationMembershipTable, + LiteLLM_UserTable, + LitellmUserRoles, + UserAPIKeyAuth, + ) + from litellm.proxy.management_endpoints.team_endpoints import list_team_v2 + + mock_request = Mock(spec=Request) + mock_user_api_key_dict = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + user_id="org_admin_user", + ) + + mock_user = LiteLLM_UserTable( + user_id="org_admin_user", + teams=[], + organization_memberships=[ + LiteLLM_OrganizationMembershipTable( + user_id="org_admin_user", + organization_id="org_A", + user_role="org_admin", + spend=0.0, + created_at=datetime.now(), + updated_at=datetime.now(), + ), + ], + ) + + with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, \ + patch("litellm.proxy.proxy_server.user_api_key_cache"), \ + patch("litellm.proxy.proxy_server.proxy_logging_obj"), \ + patch( + "litellm.proxy.management_endpoints.team_endpoints.get_user_object", + new_callable=AsyncMock, + return_value=mock_user, + ): + mock_prisma.db = Mock() + + with pytest.raises(HTTPException) as exc_info: + await list_team_v2( + http_request=mock_request, + user_id=None, + organization_id="org_B", # not their org + team_id=None, + team_alias=None, + user_api_key_dict=mock_user_api_key_dict, + page=1, + page_size=10, + sort_by=None, + sort_order="asc", + status=None, + ) + + assert exc_info.value.status_code == 403 + assert "only view teams within your organizations" in str( + exc_info.value.detail + ).lower() + + +@pytest.mark.asyncio +async def test_list_team_v2_org_admin_with_user_id_returns_user_teams(): + """ + Test that an org admin passing user_id gets that user's direct team + memberships (not all org teams). + """ + from datetime import datetime + from unittest.mock import AsyncMock, Mock, patch + + from fastapi import Request + + from litellm.proxy._types import ( + LiteLLM_OrganizationMembershipTable, + LiteLLM_UserTable, + LitellmUserRoles, + UserAPIKeyAuth, + ) + from litellm.proxy.management_endpoints.team_endpoints import list_team_v2 + + mock_request = Mock(spec=Request) + mock_user_api_key_dict = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + user_id="org_admin_user", + ) + + mock_org_admin = LiteLLM_UserTable( + user_id="org_admin_user", + teams=["team_1"], + organization_memberships=[ + LiteLLM_OrganizationMembershipTable( + user_id="org_admin_user", + organization_id="org_A", + user_role="org_admin", + spend=0.0, + created_at=datetime.now(), + updated_at=datetime.now(), + ), + ], + ) + + # The target user whose teams we want to list + mock_target_user = LiteLLM_UserTable( + user_id="target_user", + teams=["team_X", "team_Y"], + ) + + call_count = 0 + + async def mock_get_user_object(**kwargs): + nonlocal call_count + call_count += 1 + # First call: org admin lookup in list_team_v2 + # Second call: target user lookup in _build_team_list_where_conditions + if call_count == 1: + return mock_org_admin + return mock_target_user + + with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, \ + patch("litellm.proxy.proxy_server.user_api_key_cache"), \ + patch("litellm.proxy.proxy_server.proxy_logging_obj"), \ + patch( + "litellm.proxy.management_endpoints.team_endpoints.get_user_object", + side_effect=mock_get_user_object, + ): + mock_db = Mock() + mock_prisma.db = mock_db + + mock_team = Mock() + mock_team.model_dump.return_value = { + "team_id": "team_X", + "team_alias": "Target Team", + "members_with_roles": [{"user_id": "target_user", "role": "user"}], + } + mock_db.litellm_teamtable.find_many = AsyncMock(return_value=[mock_team]) + mock_db.litellm_teamtable.count = AsyncMock(return_value=1) + + result = await list_team_v2( + http_request=mock_request, + user_id="target_user", + organization_id=None, + team_id=None, + team_alias=None, + user_api_key_dict=mock_user_api_key_dict, + page=1, + page_size=10, + sort_by=None, + sort_order="asc", + status=None, + ) + + assert result["total"] == 1 + + # Verify the where clause filters by user's teams, not org scope + where = mock_db.litellm_teamtable.find_many.call_args.kwargs["where"] + assert where["team_id"] == {"in": ["team_X", "team_Y"]} + assert "organization_id" not in where + + @pytest.mark.asyncio async def test_list_team_v2_with_invalid_status(): """ Test that invalid status parameter raises HTTPException. """ from unittest.mock import Mock, patch - + from fastapi import HTTPException, Request - + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.management_endpoints.team_endpoints import list_team_v2 - + # Mock request mock_request = Mock(spec=Request) - + # Mock admin user mock_user_api_key_dict_admin = UserAPIKeyAuth( user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin_user_123", ) - + mock_prisma_client = Mock() # Mock prisma_client to be non-None diff --git a/ui/litellm-dashboard/src/components/DebugWarningBanner.test.tsx b/ui/litellm-dashboard/src/components/DebugWarningBanner.test.tsx new file mode 100644 index 0000000000..2412a84da3 --- /dev/null +++ b/ui/litellm-dashboard/src/components/DebugWarningBanner.test.tsx @@ -0,0 +1,41 @@ +import { renderWithProviders, screen } from "../../tests/test-utils"; +import { vi } from "vitest"; +import { DebugWarningBanner } from "./DebugWarningBanner"; + +vi.mock("@/app/(dashboard)/hooks/healthReadiness/useHealthReadiness", () => ({ + useHealthReadiness: vi.fn(), +})); + +import { useHealthReadiness } from "@/app/(dashboard)/hooks/healthReadiness/useHealthReadiness"; + +describe("DebugWarningBanner", () => { + it("should render", () => { + vi.mocked(useHealthReadiness).mockReturnValue({ data: { is_detailed_debug: true } } as any); + renderWithProviders(); + expect(screen.getByRole("alert")).toBeInTheDocument(); + }); + + it("should show warning when detailed debug mode is active", () => { + vi.mocked(useHealthReadiness).mockReturnValue({ data: { is_detailed_debug: true } } as any); + renderWithProviders(); + expect(screen.getByText(/Performance Warning: Detailed Debug Mode Active/i)).toBeInTheDocument(); + }); + + it("should mention LITELLM_LOG=DEBUG in the description", () => { + vi.mocked(useHealthReadiness).mockReturnValue({ data: { is_detailed_debug: true } } as any); + renderWithProviders(); + expect(screen.getByText("LITELLM_LOG=DEBUG")).toBeInTheDocument(); + }); + + it("should render nothing when is_detailed_debug is false", () => { + vi.mocked(useHealthReadiness).mockReturnValue({ data: { is_detailed_debug: false } } as any); + const { container } = renderWithProviders(); + expect(container).toBeEmptyDOMElement(); + }); + + it("should render nothing when health data is undefined", () => { + vi.mocked(useHealthReadiness).mockReturnValue({ data: undefined } as any); + const { container } = renderWithProviders(); + expect(container).toBeEmptyDOMElement(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/EntityUsageExport/ExportFormatSelector.test.tsx b/ui/litellm-dashboard/src/components/EntityUsageExport/ExportFormatSelector.test.tsx new file mode 100644 index 0000000000..8447df4a6b --- /dev/null +++ b/ui/litellm-dashboard/src/components/EntityUsageExport/ExportFormatSelector.test.tsx @@ -0,0 +1,20 @@ +import { renderWithProviders, screen } from "../../../tests/test-utils"; +import { vi } from "vitest"; +import ExportFormatSelector from "./ExportFormatSelector"; + +describe("ExportFormatSelector", () => { + it("should render", () => { + renderWithProviders(); + expect(screen.getByText("Format")).toBeInTheDocument(); + }); + + it("should display the current value", () => { + renderWithProviders(); + expect(screen.getByText("CSV (Excel, Google Sheets)")).toBeInTheDocument(); + }); + + it("should display JSON label when json is selected", () => { + renderWithProviders(); + expect(screen.getByText("JSON (includes metadata)")).toBeInTheDocument(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/EntityUsageExport/ExportSummary.test.tsx b/ui/litellm-dashboard/src/components/EntityUsageExport/ExportSummary.test.tsx new file mode 100644 index 0000000000..8ebeb33b1e --- /dev/null +++ b/ui/litellm-dashboard/src/components/EntityUsageExport/ExportSummary.test.tsx @@ -0,0 +1,59 @@ +import { renderWithProviders, screen } from "../../../tests/test-utils"; +import ExportSummary from "./ExportSummary"; + +describe("ExportSummary", () => { + it("should render", () => { + const dateRange = { + from: new Date("2024-01-01"), + to: new Date("2024-01-31"), + }; + const { container } = renderWithProviders( + + ); + expect(container).not.toBeEmptyDOMElement(); + }); + + it("should display the date range", () => { + const from = new Date(2024, 0, 1); + const to = new Date(2024, 0, 31); + const dateRange = { from, to }; + renderWithProviders( + + ); + expect(screen.getByText(new RegExp(from.toLocaleDateString()))).toBeInTheDocument(); + expect(screen.getByText(new RegExp(to.toLocaleDateString()))).toBeInTheDocument(); + }); + + it("should show filter count when filters are selected", () => { + const dateRange = { + from: new Date("2024-01-01"), + to: new Date("2024-01-31"), + }; + renderWithProviders( + + ); + expect(screen.getByText(/3 filters/)).toBeInTheDocument(); + }); + + it("should show singular 'filter' for one filter", () => { + const dateRange = { + from: new Date("2024-01-01"), + to: new Date("2024-01-31"), + }; + renderWithProviders( + + ); + expect(screen.getByText(/1 filter$/)).toBeInTheDocument(); + }); + + it("should not show filter count when no filters selected", () => { + const dateRange = { + from: new Date("2024-01-01"), + to: new Date("2024-01-31"), + }; + renderWithProviders( + + ); + expect(screen.queryByText(/filter/)).not.toBeInTheDocument(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/EntityUsageExport/ExportTypeSelector.test.tsx b/ui/litellm-dashboard/src/components/EntityUsageExport/ExportTypeSelector.test.tsx new file mode 100644 index 0000000000..36d2e12b71 --- /dev/null +++ b/ui/litellm-dashboard/src/components/EntityUsageExport/ExportTypeSelector.test.tsx @@ -0,0 +1,46 @@ +import { renderWithProviders, screen } from "../../../tests/test-utils"; +import userEvent from "@testing-library/user-event"; +import { vi } from "vitest"; +import ExportTypeSelector from "./ExportTypeSelector"; + +describe("ExportTypeSelector", () => { + it("should render", () => { + renderWithProviders( + + ); + expect(screen.getByText("Export type")).toBeInTheDocument(); + }); + + it("should display entity type in radio labels", () => { + renderWithProviders( + + ); + expect(screen.getByText(/Day-by-day breakdown by team$/)).toBeInTheDocument(); + expect(screen.getByText(/Day-by-day breakdown by team and key/)).toBeInTheDocument(); + expect(screen.getByText(/Day-by-day by team and model/)).toBeInTheDocument(); + }); + + it("should display the correct entity type for different entities", () => { + renderWithProviders( + + ); + expect(screen.getByText(/Day-by-day breakdown by organization$/)).toBeInTheDocument(); + }); + + it("should call onChange when a radio option is selected", async () => { + const user = userEvent.setup(); + const onChange = vi.fn(); + renderWithProviders( + + ); + await user.click(screen.getByRole("radio", { name: /Day-by-day breakdown by team and key/i })); + expect(onChange).toHaveBeenCalledWith("daily_with_keys"); + }); + + it("should have the correct radio checked", () => { + renderWithProviders( + + ); + expect(screen.getByRole("radio", { name: /Day-by-day by team and model/i })).toBeChecked(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/GuardrailsMonitor/MetricCard.test.tsx b/ui/litellm-dashboard/src/components/GuardrailsMonitor/MetricCard.test.tsx new file mode 100644 index 0000000000..d30df1a006 --- /dev/null +++ b/ui/litellm-dashboard/src/components/GuardrailsMonitor/MetricCard.test.tsx @@ -0,0 +1,49 @@ +import { renderWithProviders, screen } from "../../../tests/test-utils"; +import React from "react"; +import { MetricCard } from "./MetricCard"; + +describe("MetricCard", () => { + it("should render", () => { + renderWithProviders(); + expect(screen.getByText("Total Requests")).toBeInTheDocument(); + }); + + it("should display the label and value", () => { + renderWithProviders(); + expect(screen.getByText("Success Rate")).toBeInTheDocument(); + expect(screen.getByText("98.5%")).toBeInTheDocument(); + }); + + it("should display numeric values", () => { + renderWithProviders(); + expect(screen.getByText("42")).toBeInTheDocument(); + }); + + it("should render icon when provided", () => { + renderWithProviders( + icon} + /> + ); + expect(screen.getByTestId("test-icon")).toBeInTheDocument(); + }); + + it("should not render icon container when no icon provided", () => { + renderWithProviders(); + expect(screen.queryByTestId("test-icon")).not.toBeInTheDocument(); + }); + + it("should render subtitle when provided", () => { + renderWithProviders( + + ); + expect(screen.getByText("Last 24 hours")).toBeInTheDocument(); + }); + + it("should not render subtitle when not provided", () => { + renderWithProviders(); + expect(screen.queryByText("Last 24 hours")).not.toBeInTheDocument(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/HelpLink.test.tsx b/ui/litellm-dashboard/src/components/HelpLink.test.tsx index 72033bd93a..76c42b8159 100644 --- a/ui/litellm-dashboard/src/components/HelpLink.test.tsx +++ b/ui/litellm-dashboard/src/components/HelpLink.test.tsx @@ -23,6 +23,11 @@ describe("HelpLink", () => { expect(screen.getByText("Custom docs link")).toBeInTheDocument(); }); + it("should have the correct href", () => { + renderWithProviders(); + expect(screen.getByRole("link")).toHaveAttribute("href", "https://docs.example.com/test"); + }); + it("should include a screen-reader-only label for accessibility", () => { renderWithProviders(); @@ -46,7 +51,21 @@ describe("HelpIcon", () => { expect(screen.getByText("Tooltip help text")).toBeInTheDocument(); }); + it("should hide tooltip content when not hovered", () => { + renderWithProviders(); + expect(screen.queryByText("Hidden tooltip")).not.toBeInTheDocument(); + }); + it("should show learn more link when learnMoreHref is provided", async () => { + const user = userEvent.setup(); + renderWithProviders( + + ); + await user.hover(screen.getByRole("button", { name: /help information/i })); + expect(screen.getByText("Learn more")).toBeInTheDocument(); + }); + + it("should use custom learn more text when provided", async () => { const user = userEvent.setup(); renderWithProviders( { expect(screen.getByRole("button", { name: /docs/i })).toBeInTheDocument(); }); + it("should hide menu items initially", () => { + renderWithProviders(); + expect(screen.queryByText("Custom pricing")).not.toBeInTheDocument(); + }); + it("should show menu items when button is clicked", async () => { const user = userEvent.setup(); renderWithProviders(); diff --git a/ui/litellm-dashboard/src/components/ToolPolicies/PolicySelect.test.tsx b/ui/litellm-dashboard/src/components/ToolPolicies/PolicySelect.test.tsx new file mode 100644 index 0000000000..cde59fe9c6 --- /dev/null +++ b/ui/litellm-dashboard/src/components/ToolPolicies/PolicySelect.test.tsx @@ -0,0 +1,84 @@ +import { renderWithProviders, screen } from "../../../tests/test-utils"; +import { vi } from "vitest"; +import { + PolicySelect, + policyStyle, + INPUT_POLICY_OPTIONS, + OUTPUT_POLICY_OPTIONS, +} from "./PolicySelect"; + +describe("policyStyle", () => { + it("should return the matching option for a known policy", () => { + expect(policyStyle("trusted")).toEqual(INPUT_POLICY_OPTIONS[1]); + }); + + it("should return the matching option for blocked", () => { + expect(policyStyle("blocked")).toEqual(INPUT_POLICY_OPTIONS[2]); + }); + + it("should return the first option as fallback for unknown policy", () => { + expect(policyStyle("unknown")).toEqual(INPUT_POLICY_OPTIONS[0]); + }); +}); + +describe("PolicySelect", () => { + it("should render", () => { + renderWithProviders( + + ); + expect(screen.getByText("untrusted")).toBeInTheDocument(); + }); + + it("should show the current policy value", () => { + renderWithProviders( + + ); + expect(screen.getByText("trusted")).toBeInTheDocument(); + }); + + it("should be disabled when saving is true", () => { + renderWithProviders( + + ); + expect(screen.getByRole("combobox")).toHaveAttribute("aria-expanded", "false"); + expect(screen.getByRole("combobox").closest(".ant-select")).toHaveClass("ant-select-disabled"); + }); + + it("should not be disabled when saving is false", () => { + renderWithProviders( + + ); + expect(screen.getByRole("combobox").closest(".ant-select")).not.toHaveClass("ant-select-disabled"); + }); +}); + +describe("Policy option constants", () => { + it("should have 3 input policy options", () => { + expect(INPUT_POLICY_OPTIONS).toHaveLength(3); + }); + + it("should have 2 output policy options (no blocked)", () => { + expect(OUTPUT_POLICY_OPTIONS).toHaveLength(2); + expect(OUTPUT_POLICY_OPTIONS.map((o) => o.value)).toEqual(["untrusted", "trusted"]); + }); +}); diff --git a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx new file mode 100644 index 0000000000..59d8721b14 --- /dev/null +++ b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx @@ -0,0 +1,82 @@ +import { renderWithProviders, screen } from "../../../tests/test-utils"; +import { vi } from "vitest"; +import ComplexityRouterConfig from "./ComplexityRouterConfig"; + +const mockModelInfo = [ + { model_group: "gpt-4" }, + { model_group: "gpt-3.5-turbo" }, + { model_group: "claude-3-opus" }, +] as any[]; + +const defaultTiers = { + SIMPLE: "gpt-3.5-turbo", + MEDIUM: "gpt-3.5-turbo", + COMPLEX: "gpt-4", + REASONING: "claude-3-opus", +}; + +describe("ComplexityRouterConfig", () => { + it("should render", () => { + renderWithProviders( + + ); + expect(screen.getByText("Complexity Tier Configuration")).toBeInTheDocument(); + }); + + it("should display all four tier labels", () => { + renderWithProviders( + + ); + expect(screen.getByText("Simple Tier")).toBeInTheDocument(); + expect(screen.getByText("Medium Tier")).toBeInTheDocument(); + expect(screen.getByText("Complex Tier")).toBeInTheDocument(); + expect(screen.getByText("Reasoning Tier")).toBeInTheDocument(); + }); + + it("should show example queries for each tier", () => { + renderWithProviders( + + ); + expect(screen.getByText(/Hello!/)).toBeInTheDocument(); + expect(screen.getByText(/Explain how REST APIs work/)).toBeInTheDocument(); + expect(screen.getByText(/Design a microservices architecture/)).toBeInTheDocument(); + expect(screen.getByText(/Think step by step/)).toBeInTheDocument(); + }); + + it("should display the how classification works section", () => { + renderWithProviders( + + ); + expect(screen.getByText("How Classification Works")).toBeInTheDocument(); + }); + + it("should show score thresholds in the classification section", () => { + renderWithProviders( + + ); + expect(screen.getByText(/Score < 0.15/)).toBeInTheDocument(); + expect(screen.getByText(/Score 0.15 - 0.35/)).toBeInTheDocument(); + expect(screen.getByText(/Score 0.35 - 0.60/)).toBeInTheDocument(); + expect(screen.getByText(/Score > 0.60/)).toBeInTheDocument(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/agents/agent_card_grid.test.tsx b/ui/litellm-dashboard/src/components/agents/agent_card_grid.test.tsx new file mode 100644 index 0000000000..c8a21d1b6e --- /dev/null +++ b/ui/litellm-dashboard/src/components/agents/agent_card_grid.test.tsx @@ -0,0 +1,90 @@ +import { renderWithProviders, screen } from "../../../tests/test-utils"; +import { vi } from "vitest"; +import AgentCardGrid from "./agent_card_grid"; +import type { Agent, AgentKeyInfo } from "./types"; + +vi.mock("./agent_card", () => ({ + default: ({ agent, onAgentClick }: any) => ( +
onAgentClick(agent.agent_id)}> + {agent.agent_name} +
+ ), +})); + +const mockAgents: Agent[] = [ + { + agent_id: "agent-1", + agent_name: "Test Agent 1", + litellm_params: { model: "gpt-4" }, + agent_card_params: { description: "First agent" }, + }, + { + agent_id: "agent-2", + agent_name: "Test Agent 2", + litellm_params: { model: "claude-3" }, + agent_card_params: { description: "Second agent" }, + }, +]; + +const mockKeyInfoMap: Record = { + "agent-1": { has_key: true, key_alias: "key-1" }, + "agent-2": { has_key: false }, +}; + +const defaultProps = { + agentsList: mockAgents, + keyInfoMap: mockKeyInfoMap, + isLoading: false, + onDeleteClick: vi.fn(), + accessToken: "test-token", + onAgentUpdated: vi.fn(), + isAdmin: true, + onAgentClick: vi.fn(), +}; + +describe("AgentCardGrid", () => { + it("should render", () => { + renderWithProviders(); + expect(screen.getByText("Test Agent 1")).toBeInTheDocument(); + }); + + it("should render all agent cards", () => { + renderWithProviders(); + expect(screen.getByText("Test Agent 1")).toBeInTheDocument(); + expect(screen.getByText("Test Agent 2")).toBeInTheDocument(); + }); + + it("should show loading skeletons when isLoading is true", () => { + renderWithProviders(); + expect(screen.queryByText("Test Agent 1")).not.toBeInTheDocument(); + }); + + it("should show admin empty state message when no agents and isAdmin", () => { + renderWithProviders( + + ); + expect( + screen.getByText("No agents found. Create one to get started.") + ).toBeInTheDocument(); + }); + + it("should show non-admin empty state message when no agents and not admin", () => { + renderWithProviders( + + ); + expect( + screen.getByText("No agents found. Contact an admin to create agents.") + ).toBeInTheDocument(); + }); + + it("should call onAgentClick when a card is clicked", async () => { + const onAgentClick = vi.fn(); + renderWithProviders( + + ); + const { default: userEvent } = await import("@testing-library/user-event"); + const user = userEvent.setup(); + await user.click(screen.getByTestId("agent-card-agent-1")); + expect(onAgentClick).toHaveBeenCalledWith("agent-1"); + }); +}); diff --git a/ui/litellm-dashboard/src/components/common_components/RateLimitTypeFormItem.test.tsx b/ui/litellm-dashboard/src/components/common_components/RateLimitTypeFormItem.test.tsx new file mode 100644 index 0000000000..49b8028c51 --- /dev/null +++ b/ui/litellm-dashboard/src/components/common_components/RateLimitTypeFormItem.test.tsx @@ -0,0 +1,61 @@ +import { renderWithProviders, screen } from "../../../tests/test-utils"; +import userEvent from "@testing-library/user-event"; +import { vi } from "vitest"; +import { Form } from "antd"; +import React from "react"; +import { RateLimitTypeFormItem } from "./RateLimitTypeFormItem"; + +const Wrapper: React.FC<{ children: React.ReactNode }> = ({ children }) => ( +
{children}
+); + +describe("RateLimitTypeFormItem", () => { + it("should render", () => { + renderWithProviders( + + + + ); + expect(screen.getByText(/TPM Rate Limit Type/)).toBeInTheDocument(); + }); + + it("should display TPM label for tpm type", () => { + renderWithProviders( + + + + ); + expect(screen.getByText(/TPM Rate Limit Type/)).toBeInTheDocument(); + }); + + it("should display RPM label for rpm type", () => { + renderWithProviders( + + + + ); + expect(screen.getByText(/RPM Rate Limit Type/)).toBeInTheDocument(); + }); + + it("should show the select placeholder by default", () => { + renderWithProviders( + + + + ); + expect(screen.getByText("Select rate limit type")).toBeInTheDocument(); + }); + + it("should call onChange when provided", async () => { + const user = userEvent.setup(); + const onChange = vi.fn(); + renderWithProviders( + + + + ); + await user.click(screen.getByRole("combobox")); + await user.click(screen.getByText("Guaranteed throughput")); + expect(onChange).toHaveBeenCalledWith("guaranteed_throughput"); + }); +}); diff --git a/ui/litellm-dashboard/src/components/playground/chat_ui/ChatUI.tsx b/ui/litellm-dashboard/src/components/playground/chat_ui/ChatUI.tsx index eba9d37669..bc0d52cf58 100644 --- a/ui/litellm-dashboard/src/components/playground/chat_ui/ChatUI.tsx +++ b/ui/litellm-dashboard/src/components/playground/chat_ui/ChatUI.tsx @@ -60,6 +60,7 @@ import CodeInterpreterOutput from "./CodeInterpreterOutput"; import CodeInterpreterTool from "./CodeInterpreterTool"; import { generateCodeSnippet } from "./CodeSnippets"; import EndpointSelector from "./EndpointSelector"; +import FilePreviewCard from "./FilePreviewCard"; import MCPEventsDisplay from "./MCPEventsDisplay"; import type { MCPEvent } from "../../mcp_tools/types"; import { EndpointType, getEndpointType } from "./mode_endpoint_mapping"; @@ -2231,67 +2232,19 @@ const ChatUI: React.FC = ({ {/* Show file previews above input when files are uploaded */} {endpointType === EndpointType.RESPONSES && responsesUploadedImage && ( -
-
-
- {responsesUploadedImage.name.toLowerCase().endsWith(".pdf") ? ( -
- -
- ) : ( - Upload preview - )} -
-
-
{responsesUploadedImage.name}
-
- {responsesUploadedImage.name.toLowerCase().endsWith(".pdf") ? "PDF" : "Image"} -
-
- -
-
+ )} {endpointType === EndpointType.CHAT && chatUploadedImage && ( -
-
-
- {chatUploadedImage.name.toLowerCase().endsWith(".pdf") ? ( -
- -
- ) : ( - Upload preview - )} -
-
-
{chatUploadedImage.name}
-
- {chatUploadedImage.name.toLowerCase().endsWith(".pdf") ? "PDF" : "Image"} -
-
- -
-
+ )} {/* Code Interpreter indicator and sample prompts when enabled */} diff --git a/ui/litellm-dashboard/src/components/playground/chat_ui/FilePreviewCard.test.tsx b/ui/litellm-dashboard/src/components/playground/chat_ui/FilePreviewCard.test.tsx new file mode 100644 index 0000000000..1014df5d16 --- /dev/null +++ b/ui/litellm-dashboard/src/components/playground/chat_ui/FilePreviewCard.test.tsx @@ -0,0 +1,70 @@ +import { render, screen } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import { describe, expect, it, vi } from "vitest"; +import FilePreviewCard from "./FilePreviewCard"; + +function makeFile(name: string): File { + return new File(["dummy"], name, { type: "application/octet-stream" }); +} + +describe("FilePreviewCard", () => { + it("should render", () => { + render( + + ); + expect(screen.getByText("photo.png")).toBeInTheDocument(); + }); + + it("should display the file name", () => { + render( + + ); + expect(screen.getByText("my-screenshot.jpg")).toBeInTheDocument(); + }); + + it("should show 'Image' label for non-PDF files", () => { + render( + + ); + expect(screen.getByText("Image")).toBeInTheDocument(); + }); + + it("should show 'PDF' label for PDF files", () => { + render( + + ); + expect(screen.getByText("PDF")).toBeInTheDocument(); + }); + + it("should render an image preview when the file is not a PDF", () => { + render( + + ); + expect(screen.getByAltText("Upload preview")).toBeInTheDocument(); + }); + + it("should not render an image preview when the file is a PDF", () => { + render( + + ); + expect(screen.queryByAltText("Upload preview")).not.toBeInTheDocument(); + }); + + it("should call onRemove when the remove button is clicked", async () => { + const onRemove = vi.fn(); + const user = userEvent.setup(); + render( + + ); + await user.click(screen.getByRole("button")); + expect(onRemove).toHaveBeenCalledOnce(); + }); + + it("should treat .PDF (uppercase) as a PDF file", () => { + render( + + ); + expect(screen.getByText("PDF")).toBeInTheDocument(); + expect(screen.queryByAltText("Upload preview")).not.toBeInTheDocument(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/playground/chat_ui/FilePreviewCard.tsx b/ui/litellm-dashboard/src/components/playground/chat_ui/FilePreviewCard.tsx new file mode 100644 index 0000000000..26bdf07eda --- /dev/null +++ b/ui/litellm-dashboard/src/components/playground/chat_ui/FilePreviewCard.tsx @@ -0,0 +1,45 @@ +import { DeleteOutlined, FilePdfOutlined } from "@ant-design/icons"; + +interface FilePreviewCardProps { + file: File; + previewUrl: string | null; + onRemove: () => void; +} + +function FilePreviewCard({ file, previewUrl, onRemove }: FilePreviewCardProps) { + const isPdf = file.name.toLowerCase().endsWith(".pdf"); + + return ( +
+
+
+ {isPdf ? ( +
+ +
+ ) : ( + Upload preview + )} +
+
+
{file.name}
+
+ {isPdf ? "PDF" : "Image"} +
+
+ +
+
+ ); +} + +export default FilePreviewCard;