diff --git a/.github/workflows/test-linting.yml b/.github/workflows/test-linting.yml index 09406d77634..8d2b2c2f972 100644 --- a/.github/workflows/test-linting.yml +++ b/.github/workflows/test-linting.yml @@ -104,6 +104,7 @@ jobs: - name: Check basedpyright budget (delta vs base) env: BASE_SHA: ${{ github.event.pull_request.base.sha }} + NODE_OPTIONS: --max-old-space-size=12288 run: | (uv run --no-sync basedpyright --outputjson || true) | uv run --no-sync python scripts/type_check_gate.py --base "$BASE_SHA" diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index 75d4d13eb71..28602fc235f 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -1,9 +1,9 @@ { "reportAny": { - "limit": 37484 + "limit": 34906 }, "reportArgumentType": { - "limit": 2704 + "limit": 2701 }, "reportAssignmentType": { "limit": 330 @@ -12,7 +12,7 @@ "limit": 516 }, "reportCallIssue": { - "limit": 124 + "limit": 123 }, "reportConstantRedefinition": { "limit": 59 @@ -24,7 +24,7 @@ "limit": 42 }, "reportExplicitAny": { - "limit": 10389 + "limit": 10230 }, "reportFunctionMemberAccess": { "limit": 11 @@ -54,10 +54,10 @@ "limit": 0 }, "reportMissingParameterType": { - "limit": 5900 + "limit": 5893 }, "reportMissingTypeArgument": { - "limit": 15903 + "limit": 15886 }, "reportMissingTypeStubs": { "limit": 41 @@ -99,31 +99,31 @@ "limit": 0 }, "reportUnknownArgumentType": { - "limit": 45894 + "limit": 45870 }, "reportUnknownLambdaType": { "limit": 113 }, "reportUnknownMemberType": { - "limit": 40539 + "limit": 40525 }, "reportUnknownParameterType": { - "limit": 20403 + "limit": 20384 }, "reportUnknownVariableType": { - "limit": 32141 + "limit": 32099 }, "reportUnnecessaryCast": { "limit": 177 }, "reportUnnecessaryComparison": { - "limit": 1025 + "limit": 1023 }, "reportUnnecessaryContains": { "limit": 7 }, "reportUnnecessaryIsInstance": { - "limit": 1209 + "limit": 1206 }, "reportUntypedBaseClass": { "limit": 165 diff --git a/db_scripts/backfill_daily_tool_spend.sql b/db_scripts/backfill_daily_tool_spend.sql new file mode 100644 index 00000000000..309b9dbe0ff --- /dev/null +++ b/db_scripts/backfill_daily_tool_spend.sql @@ -0,0 +1,46 @@ +-- One-shot backfill of the LiteLLM_DailyToolSpend rollup from the per-request +-- LiteLLM_SpendLogToolIndex x LiteLLM_SpendLogs tables. +-- +-- This is an opt-in, manual operation. New deployments do not need it: the +-- rollup is written at request time from the moment the release is deployed. +-- Run it only if you want the Cost Optimization "Spend by tool" card to show +-- history from before the deploy, and only once. +-- +-- IMPORTANT caveats before running: +-- +-- 1. Pre-deploy index rows may include tools that were merely DECLARED in a +-- request body but never invoked (the release this ships with stops +-- recording those). For agentic clients that declare many tools per +-- request, backfilled history attributes each request's full spend to +-- every declared tool, overstating per-tool spend. Post-deploy rows do not +-- have this problem. If your traffic is mostly such clients, consider not +-- backfilling. +-- +-- 2. Coverage is bounded by spend-log retention: rows older than +-- maximum_spend_logs_retention_period are already gone. +-- +-- 3. Replace the cutover timestamp below with the time you deployed the +-- release, so backfilled per-request rows cannot double-count on top of +-- rollup rows the new writer already created. ON CONFLICT DO NOTHING is a +-- second guard for (date, tool_name) buckets the writer already touched: +-- such buckets keep the writer's numbers and skip the backfill's. +-- +-- Usage: +-- psql "$DATABASE_URL" -v cutover="'2026-07-25T00:00:00Z'" -f db_scripts/backfill_daily_tool_spend.sql + +SET TIME ZONE 'UTC'; + +INSERT INTO "LiteLLM_DailyToolSpend" (date, tool_name, spend, total_tokens, request_count, created_at, updated_at) +SELECT + to_char(ti.start_time, 'YYYY-MM-DD') AS date, + ti.tool_name, + COALESCE(SUM(sl.spend), 0) AS spend, + COALESCE(SUM(sl.total_tokens), 0) AS total_tokens, + COUNT(*) AS request_count, + now() AS created_at, + now() AS updated_at +FROM "LiteLLM_SpendLogToolIndex" ti +JOIN "LiteLLM_SpendLogs" sl ON sl.request_id = ti.request_id +WHERE ti.start_time < :cutover::timestamptz +GROUP BY 1, 2 +ON CONFLICT (date, tool_name) DO NOTHING; diff --git a/enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py b/enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py index a057df65500..9d668985eb8 100644 --- a/enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py +++ b/enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py @@ -11,7 +11,8 @@ Endpoints for /project operations #### PROJECT MANAGEMENT #### import json -from typing import List, Optional, Union +from collections.abc import Mapping, Sequence +from typing import TYPE_CHECKING from fastapi import APIRouter, Depends, HTTPException, Request @@ -25,15 +26,24 @@ from litellm.proxy.management_helpers.utils import ( ) from litellm.proxy.utils import PrismaClient, handle_exception_on_proxy +if TYPE_CHECKING: + from prisma import models as prisma_models + from prisma.actions import LiteLLM_TeamTableActions + router = APIRouter() +def _team_table(prisma_client: PrismaClient) -> "LiteLLM_TeamTableActions[prisma_models.LiteLLM_TeamTable]": + team_table: LiteLLM_TeamTableActions[prisma_models.LiteLLM_TeamTable] = prisma_client.db.litellm_teamtable + return team_table + + async def _check_user_permission_for_project( user_api_key_dict: UserAPIKeyAuth, - team_id: Optional[str], + team_id: str | None, prisma_client: PrismaClient, require_admin: bool = False, - team_object: Optional[LiteLLM_TeamTable] = None, + team_object: LiteLLM_TeamTable | None = None, ) -> bool: """ Check if user has permission to manage a project. @@ -57,9 +67,7 @@ async def _check_user_permission_for_project( team = team_object if team is None: - team = await prisma_client.db.litellm_teamtable.find_unique( - where={"team_id": team_id} - ) + team = await _team_table(prisma_client).find_unique(where={"team_id": team_id}) if team and team.admins: return user_api_key_dict.user_id in team.admins @@ -70,9 +78,9 @@ async def _check_user_permission_for_project( async def _validate_team_exists( team_id: str, prisma_client: PrismaClient, -): +) -> "prisma_models.LiteLLM_TeamTable": """Validate that a team exists. Returns the team row.""" - team = await prisma_client.db.litellm_teamtable.find_unique( + team = await _team_table(prisma_client).find_unique( where={"team_id": team_id}, ) @@ -89,7 +97,7 @@ async def _validate_team_exists( def _check_team_project_limits( team_object: LiteLLM_TeamTable, - data: Union[NewProjectRequest, UpdateProjectRequest], + data: NewProjectRequest | UpdateProjectRequest, ) -> None: """ Check that project limits respect its parent Team's limits. @@ -108,16 +116,12 @@ def _check_team_project_limits( if data.max_budget is not None and data.max_budget < 0: raise HTTPException( status_code=400, - detail={ - "error": f"max_budget cannot be negative. Received: {data.max_budget}" - }, + detail={"error": f"max_budget cannot be negative. Received: {data.max_budget}"}, ) if data.soft_budget is not None and data.soft_budget < 0: raise HTTPException( status_code=400, - detail={ - "error": f"soft_budget cannot be negative. Received: {data.soft_budget}" - }, + detail={"error": f"soft_budget cannot be negative. Received: {data.soft_budget}"}, ) # --- soft_budget < max_budget --- @@ -131,7 +135,7 @@ def _check_team_project_limits( ) # --- Validate project models are a subset of team models --- - project_models = getattr(data, "models", None) + project_models = data.models team_models = team_object.models or [] if project_models and len(team_models) > 0: # If team has 'all-proxy-models', skip validation as it allows all models @@ -148,11 +152,7 @@ def _check_team_project_limits( # --- Validate project max_budget <= team max_budget --- # Team stores budget fields directly (max_budget, tpm_limit, rpm_limit) # unlike Project which uses a separate LiteLLM_BudgetTable relation - if ( - data.max_budget is not None - and team_object.max_budget is not None - and data.max_budget > team_object.max_budget - ): + if data.max_budget is not None and team_object.max_budget is not None and data.max_budget > team_object.max_budget: raise HTTPException( status_code=400, detail={ @@ -161,11 +161,7 @@ def _check_team_project_limits( ) # --- Validate project tpm_limit <= team tpm_limit --- - if ( - data.tpm_limit is not None - and team_object.tpm_limit is not None - and data.tpm_limit > team_object.tpm_limit - ): + if data.tpm_limit is not None and team_object.tpm_limit is not None and data.tpm_limit > team_object.tpm_limit: raise HTTPException( status_code=400, detail={ @@ -174,11 +170,7 @@ def _check_team_project_limits( ) # --- Validate project rpm_limit <= team rpm_limit --- - if ( - data.rpm_limit is not None - and team_object.rpm_limit is not None - and data.rpm_limit > team_object.rpm_limit - ): + if data.rpm_limit is not None and team_object.rpm_limit is not None and data.rpm_limit > team_object.rpm_limit: raise HTTPException( status_code=400, detail={ @@ -189,19 +181,19 @@ def _check_team_project_limits( async def _create_budget_for_project( data: NewProjectRequest, - user_id: Optional[str], + user_id: str | None, litellm_proxy_admin_name: str, prisma_client: PrismaClient, ) -> str: """Create a budget for the project and return budget_id.""" budget_params = LiteLLM_BudgetTable.model_fields.keys() - _json_data = data.json(exclude_none=True) + _json_data: Mapping[str, object] = data.json(exclude_none=True) _budget_data = {k: v for k, v in _json_data.items() if k in budget_params} - budget_row = LiteLLM_BudgetTable(**_budget_data) + budget_row = LiteLLM_BudgetTable.model_validate(_budget_data) new_budget = prisma_client.jsonify_object(budget_row.json(exclude_none=True)) - _budget = await prisma_client.db.litellm_budgettable.create( + _budget: prisma_models.LiteLLM_BudgetTable = await prisma_client.db.litellm_budgettable.create( data={ **new_budget, "created_by": user_id or litellm_proxy_admin_name, @@ -214,8 +206,8 @@ async def _create_budget_for_project( async def _set_project_object_permission( data: NewProjectRequest, - prisma_client: Optional[PrismaClient], -) -> Optional[str]: + prisma_client: PrismaClient | None, +) -> str | None: """ Creates the LiteLLM_ObjectPermissionTable record for the project. Returns the object_permission_id if created, otherwise None. @@ -224,7 +216,7 @@ async def _set_project_object_permission( return None if data.object_permission is not None: - created_object_permission = ( + created_object_permission: prisma_models.LiteLLM_ObjectPermissionTable = ( await prisma_client.db.litellm_objectpermissiontable.create( data=data.object_permission.model_dump(exclude_none=True), ) @@ -344,8 +336,7 @@ async def new_project( raise HTTPException( status_code=403, detail={ - "error": "Only premium users can add tags to projects. " - + CommonProxyErrors.not_premium_user.value + "error": "Only premium users can add tags to projects. " + CommonProxyErrors.not_premium_user.value }, ) @@ -353,8 +344,7 @@ async def new_project( raise HTTPException( status_code=403, detail={ - "error": "Project management is an enterprise feature. " - + CommonProxyErrors.not_premium_user.value + "error": "Project management is an enterprise feature. " + CommonProxyErrors.not_premium_user.value }, ) @@ -375,13 +365,11 @@ async def new_project( ) # Validate team exists and get team object with budget - team_object = await _validate_team_exists( - team_id=data.team_id, prisma_client=prisma_client - ) + team_object = await _validate_team_exists(team_id=data.team_id, prisma_client=prisma_client) # Validate project limits against team limits _check_team_project_limits( - team_object=LiteLLM_TeamTable(**team_object.model_dump()), + team_object=LiteLLM_TeamTable.model_validate(team_object.model_dump()), data=data, ) @@ -391,7 +379,7 @@ async def new_project( user_api_key_dict=user_api_key_dict, team_id=data.team_id, prisma_client=prisma_client, - team_object=LiteLLM_TeamTable(**team_object.model_dump()), + team_object=LiteLLM_TeamTable.model_validate(team_object.model_dump()), ) if not has_permission: @@ -449,17 +437,13 @@ async def new_project( value=getattr(data, field), ) - new_project_row = prisma_client.jsonify_object( - project_row.json(exclude_none=True) - ) + new_project_row = prisma_client.jsonify_object(project_row.json(exclude_none=True)) # Remove budget fields (following organization_endpoints.py pattern) new_project_row = _remove_budget_fields_from_project_data(new_project_row) - verbose_proxy_logger.info( - f"new_project_row: {json.dumps(new_project_row, indent=2)}" - ) - response = await prisma_client.db.litellm_projecttable.create( + verbose_proxy_logger.info(f"new_project_row: {json.dumps(new_project_row, indent=2)}") + response: prisma_models.LiteLLM_ProjectTable = await prisma_client.db.litellm_projecttable.create( data={ **new_project_row, # type: ignore }, @@ -469,9 +453,7 @@ async def new_project( return response except Exception as e: verbose_proxy_logger.exception( - "litellm.proxy.management_endpoints.project_endpoints.new_project(): Exception occured - {}".format( - str(e) - ) + "litellm.proxy.management_endpoints.project_endpoints.new_project(): Exception occured - {}".format(str(e)) ) raise handle_exception_on_proxy(e) @@ -539,8 +521,7 @@ async def update_project( raise HTTPException( status_code=403, detail={ - "error": "Only premium users can add tags to projects. " - + CommonProxyErrors.not_premium_user.value + "error": "Only premium users can add tags to projects. " + CommonProxyErrors.not_premium_user.value }, ) @@ -548,8 +529,7 @@ async def update_project( raise HTTPException( status_code=403, detail={ - "error": "Project management is an enterprise feature. " - + CommonProxyErrors.not_premium_user.value + "error": "Project management is an enterprise feature. " + CommonProxyErrors.not_premium_user.value }, ) @@ -576,9 +556,9 @@ async def update_project( ) # Fetch existing project - existing_project = await prisma_client.db.litellm_projecttable.find_unique( - where={"project_id": data.project_id} - ) + existing_project: ( + prisma_models.LiteLLM_ProjectTable | None + ) = await prisma_client.db.litellm_projecttable.find_unique(where={"project_id": data.project_id}) if existing_project is None: raise ProxyException( @@ -595,9 +575,7 @@ async def update_project( target_team_id = data.team_id or existing_project.team_id target_team_obj = None if target_team_id is not None: - target_team_obj = await _validate_team_exists( - team_id=target_team_id, prisma_client=prisma_client - ) + target_team_obj = await _validate_team_exists(team_id=target_team_id, prisma_client=prisma_client) has_permission = await _check_user_permission_for_project( user_api_key_dict=user_api_key_dict, @@ -620,32 +598,26 @@ async def update_project( team_id=data.team_id, prisma_client=prisma_client, team_object=( - LiteLLM_TeamTable(**target_team_obj.model_dump()) - if target_team_obj - else None + LiteLLM_TeamTable.model_validate(target_team_obj.model_dump()) if target_team_obj else None ), ) if not can_assign_to_target: raise HTTPException( status_code=403, - detail={ - "error": "Cannot reassign project to a team you are not an admin of" - }, + detail={"error": "Cannot reassign project to a team you are not an admin of"}, ) # Validate project limits against team limits if target_team_obj is not None: _check_team_project_limits( - team_object=LiteLLM_TeamTable(**target_team_obj.model_dump()), + team_object=LiteLLM_TeamTable.model_validate(target_team_obj.model_dump()), data=data, ) # Prepare update data update_data = data.json(exclude_none=True, exclude={"project_id"}) update_data = prisma_client.jsonify_object(update_data) - update_data["updated_by"] = ( - user_api_key_dict.user_id or litellm_proxy_admin_name - ) + update_data["updated_by"] = user_api_key_dict.user_id or litellm_proxy_admin_name # Handle budget updates budget_fields = LiteLLM_BudgetTable.model_fields.keys() @@ -671,21 +643,17 @@ async def update_project( if existing_project.object_permission_id: # Update existing permission await prisma_client.db.litellm_objectpermissiontable.update( - where={ - "object_permission_id": existing_project.object_permission_id - }, + where={"object_permission_id": existing_project.object_permission_id}, data=object_permission_data, ) else: # Create new permission - created_permission = ( + created_permission: prisma_models.LiteLLM_ObjectPermissionTable = ( await prisma_client.db.litellm_objectpermissiontable.create( data=object_permission_data, ) ) - update_data["object_permission_id"] = ( - created_permission.object_permission_id - ) + update_data["object_permission_id"] = created_permission.object_permission_id # Handle metadata fields for field in LiteLLM_ManagementEndpoint_MetadataFields: @@ -698,7 +666,7 @@ async def update_project( update_data = _remove_budget_fields_from_project_data(update_data) # Update project - updated_project = await prisma_client.db.litellm_projecttable.update( + updated_project: prisma_models.LiteLLM_ProjectTable | None = await prisma_client.db.litellm_projecttable.update( where={"project_id": data.project_id}, data=update_data, include={"litellm_budget_table": True, "object_permission": True}, @@ -718,7 +686,7 @@ async def update_project( "/project/delete", tags=["project management"], dependencies=[Depends(user_api_key_auth)], - response_model=List[LiteLLM_ProjectTable], + response_model=list[LiteLLM_ProjectTable], ) @management_endpoint_wrapper async def delete_project( @@ -749,8 +717,7 @@ async def delete_project( raise HTTPException( status_code=403, detail={ - "error": "Project management is an enterprise feature. " - + CommonProxyErrors.not_premium_user.value + "error": "Project management is an enterprise feature. " + CommonProxyErrors.not_premium_user.value }, ) @@ -778,9 +745,7 @@ async def delete_project( for project_id in data.project_ids: # Check if project exists - existing_project = await prisma_client.db.litellm_projecttable.find_unique( - where={"project_id": project_id} - ) + existing_project = await prisma_client.db.litellm_projecttable.find_unique(where={"project_id": project_id}) if existing_project is None: raise ProxyException( @@ -791,11 +756,9 @@ async def delete_project( ) # Check if there are any keys associated with this project - associated_keys = ( - await prisma_client.db.litellm_verificationtoken.find_many( - where={"project_id": project_id} - ) - ) + associated_keys: Sequence[ + prisma_models.LiteLLM_VerificationToken + ] = await prisma_client.db.litellm_verificationtoken.find_many(where={"project_id": project_id}) if len(associated_keys) > 0: raise ProxyException( @@ -806,9 +769,9 @@ async def delete_project( ) # Delete the project - deleted_project = await prisma_client.db.litellm_projecttable.delete( - where={"project_id": project_id} - ) + deleted_project: ( + prisma_models.LiteLLM_ProjectTable | None + ) = await prisma_client.db.litellm_projecttable.delete(where={"project_id": project_id}) deleted_projects.append(deleted_project) @@ -854,7 +817,7 @@ async def project_info( ) # Fetch project - project = await prisma_client.db.litellm_projecttable.find_unique( + project: prisma_models.LiteLLM_ProjectTable | None = await prisma_client.db.litellm_projecttable.find_unique( where={"project_id": project_id}, include={"litellm_budget_table": True, "object_permission": True}, ) @@ -872,17 +835,11 @@ async def project_info( is_team_member = False if project.team_id and user_api_key_dict.user_id: - team = await prisma_client.db.litellm_teamtable.find_unique( - where={"team_id": project.team_id} - ) + team = await _team_table(prisma_client).find_unique(where={"team_id": project.team_id}) if team: caller_user_id = user_api_key_dict.user_id for m in team.members_with_roles or []: - m_user_id = ( - m.get("user_id") - if isinstance(m, dict) - else getattr(m, "user_id", None) - ) + m_user_id = m.get("user_id") if isinstance(m, dict) else getattr(m, "user_id", None) if m_user_id == caller_user_id: is_team_member = True break @@ -896,9 +853,7 @@ async def project_info( return project except Exception as e: verbose_proxy_logger.exception( - "litellm.proxy.management_endpoints.project_endpoints.project_info(): Exception occured - {}".format( - str(e) - ) + "litellm.proxy.management_endpoints.project_endpoints.project_info(): Exception occured - {}".format(str(e)) ) raise handle_exception_on_proxy(e) @@ -907,7 +862,7 @@ async def project_info( "/project/list", tags=["project management"], dependencies=[Depends(user_api_key_auth)], - response_model=List[LiteLLM_ProjectTable], + response_model=list[LiteLLM_ProjectTable], ) async def list_projects( user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), @@ -932,21 +887,19 @@ async def list_projects( # If proxy admin, get all projects if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN: - projects = await prisma_client.db.litellm_projecttable.find_many( + projects: Sequence[ + prisma_models.LiteLLM_ProjectTable + ] = await prisma_client.db.litellm_projecttable.find_many( include={"litellm_budget_table": True, "object_permission": True} ) else: # Look up the user's team memberships via the reverse-index on # LiteLLM_UserTable.teams (maintained by team_member_add alongside # members_with_roles). This avoids a full scan of all team rows. - user_record = await prisma_client.db.litellm_usertable.find_unique( + user_record: prisma_models.LiteLLM_UserTable | None = await prisma_client.db.litellm_usertable.find_unique( where={"user_id": user_api_key_dict.user_id}, ) - user_team_ids = ( - user_record.teams - if user_record is not None and user_record.teams - else [] - ) + user_team_ids: Sequence[str] = user_record.teams if user_record is not None and user_record.teams else [] projects = await prisma_client.db.litellm_projecttable.find_many( where={"team_id": {"in": user_team_ids}}, diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260725000000_add_daily_tool_spend/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260725000000_add_daily_tool_spend/migration.sql new file mode 100644 index 00000000000..e02ed01a554 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260725000000_add_daily_tool_spend/migration.sql @@ -0,0 +1,12 @@ +-- CreateTable +CREATE TABLE IF NOT EXISTS "LiteLLM_DailyToolSpend" ( + "date" TEXT NOT NULL, + "tool_name" TEXT NOT NULL, + "spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0, + "total_tokens" BIGINT NOT NULL DEFAULT 0, + "request_count" BIGINT NOT NULL DEFAULT 0, + "created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + "updated_at" TIMESTAMP(3) NOT NULL, + + CONSTRAINT "LiteLLM_DailyToolSpend_pkey" PRIMARY KEY ("date","tool_name") +); diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 314c59d3560..ed4f7d4f0a0 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -1102,6 +1102,19 @@ model LiteLLM_SpendLogToolIndex { @@index([start_time]) } +// Daily tool spend rollup (one row per tool per day) – the Cost Optimization card reads this, never SpendLogs +model LiteLLM_DailyToolSpend { + date String + tool_name String + spend Float @default(0.0) + total_tokens BigInt @default(0) + request_count BigInt @default(0) + created_at DateTime @default(now()) + updated_at DateTime @updatedAt + + @@id([date, tool_name]) +} + // Prompt table for storing prompt configurations model LiteLLM_PromptTable { id String @id @default(uuid()) diff --git a/litellm/completion_extras/litellm_responses_transformation/handler.py b/litellm/completion_extras/litellm_responses_transformation/handler.py index 8f12d855880..15f5b28e30e 100644 --- a/litellm/completion_extras/litellm_responses_transformation/handler.py +++ b/litellm/completion_extras/litellm_responses_transformation/handler.py @@ -209,7 +209,15 @@ class ResponsesToCompletionBridgeHandler: json_mode=kwargs.get("json_mode"), ) elif isinstance(result, ModelResponse): - return result + if not stream: + return result + return self._completed_response_as_stream( + response=result, + model=model, + custom_llm_provider=custom_llm_provider, + logging_obj=logging_obj, + json_mode=kwargs.get("json_mode"), + ) elif not stream: responses_api_response = self._collect_response_from_stream(result) return self.transformation_handler.transform_response( @@ -299,7 +307,15 @@ class ResponsesToCompletionBridgeHandler: json_mode=kwargs.get("json_mode"), ) elif isinstance(result, ModelResponse): - return result + if not stream: + return result + return self._completed_response_as_stream( + response=result, + model=model, + custom_llm_provider=custom_llm_provider, + logging_obj=logging_obj, + json_mode=kwargs.get("json_mode"), + ) elif not stream: responses_api_response = await self._collect_response_from_stream_async(result) return self.transformation_handler.transform_response( @@ -331,6 +347,25 @@ class ResponsesToCompletionBridgeHandler: ) return self._apply_post_stream_processing(streamwrapper, model, custom_llm_provider) + def _completed_response_as_stream( + self, + response: "ModelResponse", + model: str, + custom_llm_provider: str, + logging_obj: "LiteLLMLoggingObj", + json_mode: bool | None, + ) -> "CustomStreamWrapper": + from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper + from litellm.llms.base_llm.base_model_iterator import MockResponseIterator + + streamwrapper = CustomStreamWrapper( + completion_stream=MockResponseIterator(model_response=response, json_mode=json_mode), + model=model, + custom_llm_provider=custom_llm_provider, + logging_obj=logging_obj, + ) + return self._apply_post_stream_processing(streamwrapper, model, custom_llm_provider) + @staticmethod def _apply_post_stream_processing( stream: "CustomStreamWrapper", diff --git a/litellm/completion_extras/litellm_responses_transformation/transformation.py b/litellm/completion_extras/litellm_responses_transformation/transformation.py index 3e50fe66039..89a44fcdeef 100644 --- a/litellm/completion_extras/litellm_responses_transformation/transformation.py +++ b/litellm/completion_extras/litellm_responses_transformation/transformation.py @@ -1077,6 +1077,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator): def __init__(self, streaming_response, sync_stream: bool, json_mode: Optional[bool] = False): super().__init__(streaming_response, sync_stream, json_mode) + self._chat_completion_id: str | None = None def _handle_string_chunk( self, str_line: Union[str, "BaseModel"] @@ -1384,4 +1385,13 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator): ModelResponseStream: OpenAI-formatted streaming chunk """ verbose_logger.debug(f"Chat provider: transform_streaming_response called with chunk: {chunk}") - return OpenAiResponsesToChatCompletionStreamIterator.translate_responses_chunk_to_openai_stream(chunk) + return self._with_stream_scoped_id( + OpenAiResponsesToChatCompletionStreamIterator.translate_responses_chunk_to_openai_stream(chunk) + ) + + def _with_stream_scoped_id(self, chunk: "ModelResponseStream") -> "ModelResponseStream": + if self._chat_completion_id is None: + self._chat_completion_id = chunk.id + else: + chunk.id = self._chat_completion_id + return chunk diff --git a/litellm/constants.py b/litellm/constants.py index 963c7376c09..167e1426ba1 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1462,7 +1462,7 @@ SPEND_LOG_CLEANUP_MAX_CONSECUTIVE_BATCH_FAILURES = int(os.getenv("SPEND_LOG_CLEA SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS = float( os.getenv("SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS", 0.5) ) -TOOL_SPEND_MAX_WINDOW_DAYS = 30 +TOOL_SPEND_TOP_TOOLS = 100 SPEND_LOG_PARTITION_INTERVAL = os.getenv("SPEND_LOG_PARTITION_INTERVAL", "day") SPEND_LOG_PARTITION_PRECREATE_AHEAD = int(os.getenv("SPEND_LOG_PARTITION_PRECREATE_AHEAD", 7)) SPEND_LOG_QUEUE_SIZE_THRESHOLD = int(os.getenv("SPEND_LOG_QUEUE_SIZE_THRESHOLD", 100)) diff --git a/litellm/integrations/langsmith.py b/litellm/integrations/langsmith.py index 18c4baccd51..565ea833768 100644 --- a/litellm/integrations/langsmith.py +++ b/litellm/integrations/langsmith.py @@ -133,6 +133,15 @@ class LangsmithLogger(CustomBatchLogger): "dotted_order": metadata.get("dotted_order", None), } + def _redact_metadata(self, metadata: dict) -> dict: + # helper is shallow; also scrub nested requester_metadata since + # LangSmith forwards the whole dict into the run + redacted = redact_user_api_key_info(metadata=dict(metadata)) + nested = redacted.get("requester_metadata") + if isinstance(nested, dict): + redacted["requester_metadata"] = redact_user_api_key_info(metadata=nested) + return redacted + def _build_extra_metadata(self, metadata: Dict): extra_metadata = dict(metadata) requester_metadata = extra_metadata.get("requester_metadata") @@ -141,13 +150,7 @@ class LangsmithLogger(CustomBatchLogger): if key in requester_metadata and key not in extra_metadata: extra_metadata[key] = requester_metadata[key] - # helper is shallow; also scrub nested requester_metadata since - # LangSmith forwards the whole dict into `extra` - extra_metadata = redact_user_api_key_info(metadata=extra_metadata) - nested = extra_metadata.get("requester_metadata") - if isinstance(nested, dict): - extra_metadata["requester_metadata"] = redact_user_api_key_info(metadata=nested) - return extra_metadata + return self._redact_metadata(extra_metadata) def _build_outputs_with_usage(self, payload: StandardLoggingPayload) -> Dict[str, Any]: response = payload["response"] @@ -200,12 +203,13 @@ class LangsmithLogger(CustomBatchLogger): metadata = payload["metadata"] extra_metadata = self._build_extra_metadata(dict(metadata)) + inputs = {**payload, "metadata": self._redact_metadata(dict(metadata))} outputs = self._build_outputs_with_usage(payload) data = { "name": fields["run_name"], "run_type": "llm", - "inputs": payload, + "inputs": inputs, "outputs": outputs, "session_name": fields["project_name"], "start_time": payload["startTime"], diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index 64d4dd578b2..24597c02ea2 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -16,6 +16,7 @@ from typing import ( Dict, List, Literal, + Mapping, Optional, Sequence, Tuple, @@ -1449,6 +1450,8 @@ class PrometheusLogger(CustomLogger): prompt_details = usage_object.get("prompt_tokens_details") or {} completion_details = usage_object.get("completion_tokens_details") or {} + cache_creation_detail_tokens = PrometheusLogger._resolve_cache_write_tokens(prompt_details) + detail_metrics: List[Tuple[Any, DEFINED_PROMETHEUS_METRICS, Any]] = [ ( self.litellm_input_cached_tokens_metric, @@ -1458,7 +1461,7 @@ class PrometheusLogger(CustomLogger): ( self.litellm_input_cache_creation_tokens_metric, "litellm_input_cache_creation_tokens_metric", - (prompt_details.get("cache_creation_tokens") if isinstance(prompt_details, dict) else None), + cache_creation_detail_tokens, ), ( self.litellm_input_audio_tokens_metric, @@ -1597,27 +1600,12 @@ class PrometheusLogger(CustomLogger): ) # Provider prompt caching metrics are independent of LiteLLM cache_hit. - provider_cache_read_tokens = 0 - provider_cache_creation_tokens = 0 usage_obj = (standard_logging_payload.get("metadata", {}) or {}).get("usage_object") if isinstance(usage_obj, dict): - # Prefer explicit provider cache fields when available. - _read = usage_obj.get("cache_read_input_tokens") - _write = usage_obj.get("cache_creation_input_tokens") - - if isinstance(_read, int): - provider_cache_read_tokens = _read - if isinstance(_write, int): - provider_cache_creation_tokens = _write - - # Fallback to prompt_tokens_details.cached_tokens (common normalization point). - # Only fallback when the explicit field is genuinely absent (None). - if _read is None: - prompt_details = usage_obj.get("prompt_tokens_details") - if isinstance(prompt_details, dict): - cached_tokens = prompt_details.get("cached_tokens") - if isinstance(cached_tokens, int): - provider_cache_read_tokens = cached_tokens + ( + provider_cache_read_tokens, + provider_cache_creation_tokens, + ) = PrometheusLogger._resolve_provider_cache_tokens(usage_obj) if provider_cache_read_tokens > 0: PrometheusLogger._inc_labeled_counter( @@ -1639,6 +1627,40 @@ class PrometheusLogger(CustomLogger): amount=float(provider_cache_creation_tokens), ) + @staticmethod + def _resolve_provider_cache_tokens(usage_obj: Mapping[str, object]) -> tuple[int, int]: + # Prefer explicit provider cache fields when available. + _read = usage_obj.get("cache_read_input_tokens") + _write = usage_obj.get("cache_creation_input_tokens") + + provider_cache_read_tokens = _read if isinstance(_read, int) else 0 + provider_cache_creation_tokens = _write if isinstance(_write, int) else 0 + + # Fallback to prompt_tokens_details (common normalization point). + # Only fallback when the explicit field is genuinely absent (None). + prompt_details = usage_obj.get("prompt_tokens_details") + if _read is None and isinstance(prompt_details, dict): + cached_tokens = prompt_details.get("cached_tokens") + if isinstance(cached_tokens, int): + provider_cache_read_tokens = cached_tokens + + if _write is None: + write_tokens = PrometheusLogger._resolve_cache_write_tokens(prompt_details) + if write_tokens is not None: + provider_cache_creation_tokens = write_tokens + + return provider_cache_read_tokens, provider_cache_creation_tokens + + @staticmethod + def _resolve_cache_write_tokens(prompt_details: object) -> int | None: + if not isinstance(prompt_details, dict): + return None + for key in ("cache_write_tokens", "cache_creation_tokens"): + value = prompt_details.get(key) + if isinstance(value, int) and not isinstance(value, bool): + return value + return None + def _increment_mcp_tool_call_metrics( self, standard_logging_payload: StandardLoggingPayload, diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 0a4ba411e77..b624108d670 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -5593,18 +5593,6 @@ def scrub_sensitive_keys_in_metadata(litellm_params: Optional[dict]): litellm_params["_langfuse_masking_function"] = masking_fn litellm_params["metadata"] = metadata - ## check user_api_key_metadata for sensitive logging keys - cleaned_user_api_key_metadata = {} - if "user_api_key_metadata" in metadata and isinstance(metadata["user_api_key_metadata"], dict): - for k, v in metadata["user_api_key_metadata"].items(): - if k == "logging": # prevent logging user logging keys - cleaned_user_api_key_metadata[k] = "scrubbed_by_litellm_for_sensitive_keys" - else: - cleaned_user_api_key_metadata[k] = v - - metadata["user_api_key_metadata"] = cleaned_user_api_key_metadata - litellm_params["metadata"] = metadata - return litellm_params diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index f7ff4d6b16f..4e3d94e2ab3 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -6,6 +6,7 @@ import mimetypes import re import xml.etree.ElementTree as ET from enum import Enum +from collections.abc import Mapping from typing import Any, Dict, List, Optional, Set, Tuple, TypedDict, Union, cast, overload from jinja2.sandbox import ImmutableSandboxedEnvironment @@ -5350,7 +5351,9 @@ def prompt_factory( def get_attribute_or_key(tool_or_function, attribute, default=None): if hasattr(tool_or_function, attribute): return getattr(tool_or_function, attribute) - return tool_or_function.get(attribute, default) + if isinstance(tool_or_function, Mapping): + return tool_or_function.get(attribute, default) + return default class NormalizedToolCall(TypedDict): @@ -5379,14 +5382,18 @@ def _parse_tool_call_arguments(raw: Any, tool_name: Optional[str], context: str) return parsed if isinstance(parsed, dict) else {} -def _tool_calls_from_chat_completion_response(response: Any) -> list[NormalizedToolCall]: +def _tool_calls_from_chat_completion_response( + response: Any, include_all_choices: bool = False +) -> list[NormalizedToolCall]: choices = get_attribute_or_key(response, "choices", None) if not (isinstance(choices, list) and choices): return [] - message = get_attribute_or_key(choices[0], "message", None) - tool_calls = get_attribute_or_key(message, "tool_calls", None) if message else None - if not isinstance(tool_calls, list): - return [] + tool_calls: list[Any] = [] + for choice in choices if include_all_choices else choices[:1]: + message = get_attribute_or_key(choice, "message", None) + choice_tool_calls = get_attribute_or_key(message, "tool_calls", None) if message else None + if isinstance(choice_tool_calls, list): + tool_calls.extend(choice_tool_calls) result: list[NormalizedToolCall] = [] for tc in tool_calls: fn = get_attribute_or_key(tc, "function", None) @@ -5449,7 +5456,7 @@ def _tool_calls_from_anthropic_messages_response(response: Any) -> list[Normaliz return result -def get_tool_calls_from_response(response: Any) -> list[NormalizedToolCall]: +def get_tool_calls_from_response(response: Any, include_all_choices: bool = False) -> list[NormalizedToolCall]: """ Extract tool/function calls from a response object into a normalized ``{"id", "name", "arguments"}`` shape, regardless of which API surface @@ -5457,11 +5464,20 @@ def get_tool_calls_from_response(response: Any) -> list[NormalizedToolCall]: the Responses API (``output`` items of type ``function_call``), or the Anthropic Messages API (``content`` blocks of type ``tool_use``). + ``include_all_choices`` decides the chat-completions scope: the default + reads only ``choices[0]``, which is what consumers that act on THE reply + (e.g. guardrails rebuilding the primary assistant message) want; usage + accounting passes True because every choice of an ``n>1`` request costs + money and its tool calls really ran. The other surfaces have a single + output, so the flag has no effect on them. + Callers that only care about a specific tool should filter the result by ``name`` themselves -- this returns every tool call found. """ + chat_tool_calls = _tool_calls_from_chat_completion_response(response, include_all_choices=include_all_choices) + if chat_tool_calls: + return chat_tool_calls for extractor in ( - _tool_calls_from_chat_completion_response, _tool_calls_from_responses_api_response, _tool_calls_from_anthropic_messages_response, ): diff --git a/litellm/llms/azure/batches/handler.py b/litellm/llms/azure/batches/handler.py index 808fb3d9600..4a064756295 100644 --- a/litellm/llms/azure/batches/handler.py +++ b/litellm/llms/azure/batches/handler.py @@ -2,7 +2,8 @@ Azure Batches API Handler """ -from typing import Any, Coroutine, Optional, Union, cast +from collections.abc import Coroutine +from typing import cast import httpx from openai import AsyncOpenAI, OpenAI @@ -33,32 +34,30 @@ class AzureBatchesAPI(BaseAzureLLM): async def acreate_batch( self, create_batch_data: CreateBatchRequest, - azure_client: Union[AsyncAzureOpenAI, AsyncOpenAI], + azure_client: AsyncAzureOpenAI | AsyncOpenAI, ) -> LiteLLMBatch: response = await azure_client.batches.create(**create_batch_data) # type: ignore[arg-type] - return LiteLLMBatch(**response.model_dump()) + return LiteLLMBatch.model_validate(response.model_dump()) def create_batch( self, _is_async: bool, create_batch_data: CreateBatchRequest, - api_key: Optional[str], - api_base: Optional[str], - api_version: Optional[str], - timeout: Union[float, httpx.Timeout], - max_retries: Optional[int], - client: Optional[Union[AzureOpenAI, AsyncAzureOpenAI, OpenAI, AsyncOpenAI]] = None, - litellm_params: Optional[dict] = None, - ) -> Union[LiteLLMBatch, Coroutine[Any, Any, LiteLLMBatch]]: - azure_client: Optional[Union[AzureOpenAI, AsyncAzureOpenAI, OpenAI, AsyncOpenAI]] = ( - self.get_azure_openai_client( - api_key=api_key, - api_base=api_base, - api_version=api_version, - client=client, - _is_async=_is_async, - litellm_params=litellm_params or {}, - ) + api_key: str | None, + api_base: str | None, + api_version: str | None, + timeout: float | httpx.Timeout, + max_retries: int | None, + client: AzureOpenAI | AsyncAzureOpenAI | OpenAI | AsyncOpenAI | None = None, + litellm_params: dict | None = None, + ) -> LiteLLMBatch | Coroutine[object, object, LiteLLMBatch]: + azure_client: AzureOpenAI | AsyncAzureOpenAI | OpenAI | AsyncOpenAI | None = self.get_azure_openai_client( + api_key=api_key, + api_base=api_base, + api_version=api_version, + client=client, + _is_async=_is_async, + litellm_params=litellm_params or {}, ) if azure_client is None: raise ValueError( @@ -73,38 +72,36 @@ class AzureBatchesAPI(BaseAzureLLM): return self.acreate_batch( # type: ignore create_batch_data=create_batch_data, azure_client=azure_client ) - response = cast(Union[AzureOpenAI, OpenAI], azure_client).batches.create(**create_batch_data) # type: ignore[arg-type] - return LiteLLMBatch(**response.model_dump()) + response = cast(AzureOpenAI | OpenAI, azure_client).batches.create(**create_batch_data) # type: ignore[arg-type] + return LiteLLMBatch.model_validate(response.model_dump()) async def aretrieve_batch( self, retrieve_batch_data: RetrieveBatchRequest, - client: Union[AsyncAzureOpenAI, AsyncOpenAI], + client: AsyncAzureOpenAI | AsyncOpenAI, ) -> LiteLLMBatch: response = await client.batches.retrieve(**retrieve_batch_data) # type: ignore[arg-type] - return LiteLLMBatch(**response.model_dump()) + return LiteLLMBatch.model_validate(response.model_dump()) def retrieve_batch( self, _is_async: bool, retrieve_batch_data: RetrieveBatchRequest, - api_key: Optional[str], - api_base: Optional[str], - api_version: Optional[str], - timeout: Union[float, httpx.Timeout], - max_retries: Optional[int], - client: Optional[Union[AzureOpenAI, AsyncAzureOpenAI, OpenAI, AsyncOpenAI]] = None, - litellm_params: Optional[dict] = None, + api_key: str | None, + api_base: str | None, + api_version: str | None, + timeout: float | httpx.Timeout, + max_retries: int | None, + client: AzureOpenAI | AsyncAzureOpenAI | OpenAI | AsyncOpenAI | None = None, + litellm_params: dict | None = None, ): - azure_client: Optional[Union[AzureOpenAI, AsyncAzureOpenAI, OpenAI, AsyncOpenAI]] = ( - self.get_azure_openai_client( - api_key=api_key, - api_base=api_base, - api_version=api_version, - client=client, - _is_async=_is_async, - litellm_params=litellm_params or {}, - ) + azure_client: AzureOpenAI | AsyncAzureOpenAI | OpenAI | AsyncOpenAI | None = self.get_azure_openai_client( + api_key=api_key, + api_base=api_base, + api_version=api_version, + client=client, + _is_async=_is_async, + litellm_params=litellm_params or {}, ) if azure_client is None: raise ValueError( @@ -119,38 +116,36 @@ class AzureBatchesAPI(BaseAzureLLM): return self.aretrieve_batch( # type: ignore retrieve_batch_data=retrieve_batch_data, client=azure_client ) - response = cast(Union[AzureOpenAI, OpenAI], azure_client).batches.retrieve(**retrieve_batch_data) - return LiteLLMBatch(**response.model_dump()) + response = cast(AzureOpenAI | OpenAI, azure_client).batches.retrieve(**retrieve_batch_data) + return LiteLLMBatch.model_validate(response.model_dump()) async def acancel_batch( self, cancel_batch_data: CancelBatchRequest, - client: Union[AsyncAzureOpenAI, AsyncOpenAI], + client: AsyncAzureOpenAI | AsyncOpenAI, ) -> LiteLLMBatch: response = await client.batches.cancel(**cancel_batch_data) - return LiteLLMBatch(**response.model_dump()) + return LiteLLMBatch.model_validate(response.model_dump()) def cancel_batch( self, _is_async: bool, cancel_batch_data: CancelBatchRequest, - api_key: Optional[str], - api_base: Optional[str], - api_version: Optional[str], - timeout: Union[float, httpx.Timeout], - max_retries: Optional[int], - client: Optional[Union[AzureOpenAI, AsyncAzureOpenAI, OpenAI, AsyncOpenAI]] = None, - litellm_params: Optional[dict] = None, + api_key: str | None, + api_base: str | None, + api_version: str | None, + timeout: float | httpx.Timeout, + max_retries: int | None, + client: AzureOpenAI | AsyncAzureOpenAI | OpenAI | AsyncOpenAI | None = None, + litellm_params: dict | None = None, ): - azure_client: Optional[Union[AzureOpenAI, AsyncAzureOpenAI, OpenAI, AsyncOpenAI]] = ( - self.get_azure_openai_client( - api_key=api_key, - api_base=api_base, - api_version=api_version, - client=client, - _is_async=_is_async, - litellm_params=litellm_params or {}, - ) + azure_client: AzureOpenAI | AsyncAzureOpenAI | OpenAI | AsyncOpenAI | None = self.get_azure_openai_client( + api_key=api_key, + api_base=api_base, + api_version=api_version, + client=client, + _is_async=_is_async, + litellm_params=litellm_params or {}, ) if azure_client is None: raise ValueError( @@ -172,13 +167,13 @@ class AzureBatchesAPI(BaseAzureLLM): "Azure client is not an instance of AzureOpenAI or OpenAI. Make sure you passed a sync client." ) response = azure_client.batches.cancel(**cancel_batch_data) - return LiteLLMBatch(**response.model_dump()) + return LiteLLMBatch.model_validate(response.model_dump()) async def alist_batches( self, - client: Union[AsyncAzureOpenAI, AsyncOpenAI], - after: Optional[str] = None, - limit: Optional[int] = None, + client: AsyncAzureOpenAI | AsyncOpenAI, + after: str | None = None, + limit: int | None = None, ): response = await client.batches.list(after=after, limit=limit) # type: ignore return response @@ -186,25 +181,23 @@ class AzureBatchesAPI(BaseAzureLLM): def list_batches( self, _is_async: bool, - api_key: Optional[str], - api_base: Optional[str], - api_version: Optional[str], - timeout: Union[float, httpx.Timeout], - max_retries: Optional[int], - after: Optional[str] = None, - limit: Optional[int] = None, - client: Optional[Union[AzureOpenAI, AsyncAzureOpenAI, OpenAI, AsyncOpenAI]] = None, - litellm_params: Optional[dict] = None, + api_key: str | None, + api_base: str | None, + api_version: str | None, + timeout: float | httpx.Timeout, + max_retries: int | None, + after: str | None = None, + limit: int | None = None, + client: AzureOpenAI | AsyncAzureOpenAI | OpenAI | AsyncOpenAI | None = None, + litellm_params: dict | None = None, ): - azure_client: Optional[Union[AzureOpenAI, AsyncAzureOpenAI, OpenAI, AsyncOpenAI]] = ( - self.get_azure_openai_client( - api_key=api_key, - api_base=api_base, - api_version=api_version, - client=client, - _is_async=_is_async, - litellm_params=litellm_params or {}, - ) + azure_client: AzureOpenAI | AsyncAzureOpenAI | OpenAI | AsyncOpenAI | None = self.get_azure_openai_client( + api_key=api_key, + api_base=api_base, + api_version=api_version, + client=client, + _is_async=_is_async, + litellm_params=litellm_params or {}, ) if azure_client is None: raise ValueError( diff --git a/litellm/llms/azure_ai/anthropic/count_tokens/transformation.py b/litellm/llms/azure_ai/anthropic/count_tokens/transformation.py index 5e1fb69f40d..ba930f40059 100644 --- a/litellm/llms/azure_ai/anthropic/count_tokens/transformation.py +++ b/litellm/llms/azure_ai/anthropic/count_tokens/transformation.py @@ -4,8 +4,6 @@ Azure AI Anthropic CountTokens API transformation logic. Extends the base Anthropic CountTokens transformation with Azure authentication. """ -from typing import Any, Dict, Optional - from litellm.constants import ANTHROPIC_TOKEN_COUNTING_BETA_VERSION from litellm.llms.anthropic.count_tokens.transformation import ( AnthropicCountTokensConfig, @@ -25,8 +23,8 @@ class AzureAIAnthropicCountTokensConfig(AnthropicCountTokensConfig): def get_required_headers( self, api_key: str, - litellm_params: Optional[Dict[str, Any]] = None, - ) -> Dict[str, str]: + litellm_params: dict[str, object] | None = None, + ) -> dict[str, str]: """ Get the required headers for the Azure AI Anthropic CountTokens API. @@ -53,7 +51,7 @@ class AzureAIAnthropicCountTokensConfig(AnthropicCountTokensConfig): if "api_key" not in litellm_params: litellm_params["api_key"] = api_key - litellm_params_obj = GenericLiteLLMParams(**litellm_params) + litellm_params_obj = GenericLiteLLMParams.model_validate(litellm_params) # Get Azure auth headers (api-key or Authorization) azure_headers = BaseAzureLLM._base_validate_azure_environment(headers={}, litellm_params=litellm_params_obj) diff --git a/litellm/llms/openai/evals/transformation.py b/litellm/llms/openai/evals/transformation.py index 8a55fec58a6..1ccaed72f26 100644 --- a/litellm/llms/openai/evals/transformation.py +++ b/litellm/llms/openai/evals/transformation.py @@ -2,7 +2,7 @@ OpenAI Evals API configuration and transformations """ -from typing import Any, Dict, Optional, Tuple +from collections.abc import Mapping import httpx @@ -31,6 +31,10 @@ from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import LlmProviders +def _parsed_response_json(raw_response: httpx.Response) -> Mapping[str, object]: + return raw_response.json() + + class OpenAIEvalsConfig(BaseEvalsAPIConfig): """OpenAI-specific Evals API configuration""" @@ -38,7 +42,7 @@ class OpenAIEvalsConfig(BaseEvalsAPIConfig): def custom_llm_provider(self) -> LlmProviders: return LlmProviders.OPENAI - def validate_environment(self, headers: dict, litellm_params: Optional[GenericLiteLLMParams]) -> dict: + def validate_environment(self, headers: dict, litellm_params: GenericLiteLLMParams | None) -> dict: """Add OpenAI-specific headers""" import litellm from litellm.secret_managers.main import get_secret_str @@ -61,9 +65,9 @@ class OpenAIEvalsConfig(BaseEvalsAPIConfig): def get_complete_url( self, - api_base: Optional[str], + api_base: str | None, endpoint: str, - eval_id: Optional[str] = None, + eval_id: str | None = None, ) -> str: """Get complete URL for OpenAI Evals API""" if api_base is None: @@ -79,7 +83,7 @@ class OpenAIEvalsConfig(BaseEvalsAPIConfig): create_request: CreateEvalRequest, litellm_params: GenericLiteLLMParams, headers: dict, - ) -> Dict: + ) -> dict: """Transform create eval request for OpenAI""" verbose_logger.debug("Transforming create eval request: %s", create_request) @@ -94,17 +98,17 @@ class OpenAIEvalsConfig(BaseEvalsAPIConfig): logging_obj: LiteLLMLoggingObj, ) -> Eval: """Transform OpenAI response to Eval object""" - response_json = raw_response.json() + response_json = _parsed_response_json(raw_response) verbose_logger.debug("Transforming create eval response: %s", response_json) - return Eval(**response_json) + return Eval.model_validate(response_json) def transform_list_evals_request( self, list_params: ListEvalsParams, litellm_params: GenericLiteLLMParams, headers: dict, - ) -> Tuple[str, Dict]: + ) -> tuple[str, dict]: """Transform list evals request for OpenAI""" api_base = "https://api.openai.com" if litellm_params and litellm_params.api_base: @@ -113,7 +117,7 @@ class OpenAIEvalsConfig(BaseEvalsAPIConfig): url = self.get_complete_url(api_base=api_base, endpoint="evals") # Build query parameters - query_params: Dict[str, Any] = {} + query_params: dict[str, object] = {} if "limit" in list_params and list_params["limit"]: query_params["limit"] = list_params["limit"] if "after" in list_params and list_params["after"]: @@ -138,10 +142,10 @@ class OpenAIEvalsConfig(BaseEvalsAPIConfig): logging_obj: LiteLLMLoggingObj, ) -> ListEvalsResponse: """Transform OpenAI response to ListEvalsResponse""" - response_json = raw_response.json() + response_json = _parsed_response_json(raw_response) verbose_logger.debug("Transforming list evals response: %s", response_json) - return ListEvalsResponse(**response_json) + return ListEvalsResponse.model_validate(response_json) def transform_get_eval_request( self, @@ -149,7 +153,7 @@ class OpenAIEvalsConfig(BaseEvalsAPIConfig): api_base: str, litellm_params: GenericLiteLLMParams, headers: dict, - ) -> Tuple[str, Dict]: + ) -> tuple[str, dict]: """Transform get eval request for OpenAI""" url = self.get_complete_url(api_base=api_base, endpoint="evals", eval_id=eval_id) @@ -163,10 +167,10 @@ class OpenAIEvalsConfig(BaseEvalsAPIConfig): logging_obj: LiteLLMLoggingObj, ) -> Eval: """Transform OpenAI response to Eval object""" - response_json = raw_response.json() + response_json = _parsed_response_json(raw_response) verbose_logger.debug("Transforming get eval response: %s", response_json) - return Eval(**response_json) + return Eval.model_validate(response_json) def transform_update_eval_request( self, @@ -175,7 +179,7 @@ class OpenAIEvalsConfig(BaseEvalsAPIConfig): api_base: str, litellm_params: GenericLiteLLMParams, headers: dict, - ) -> Tuple[str, Dict, Dict]: + ) -> tuple[str, dict, dict]: """Transform update eval request for OpenAI""" url = self.get_complete_url(api_base=api_base, endpoint="evals", eval_id=eval_id) @@ -192,10 +196,10 @@ class OpenAIEvalsConfig(BaseEvalsAPIConfig): logging_obj: LiteLLMLoggingObj, ) -> Eval: """Transform OpenAI response to Eval object""" - response_json = raw_response.json() + response_json = _parsed_response_json(raw_response) verbose_logger.debug("Transforming update eval response: %s", response_json) - return Eval(**response_json) + return Eval.model_validate(response_json) def transform_delete_eval_request( self, @@ -203,7 +207,7 @@ class OpenAIEvalsConfig(BaseEvalsAPIConfig): api_base: str, litellm_params: GenericLiteLLMParams, headers: dict, - ) -> Tuple[str, Dict]: + ) -> tuple[str, dict]: """Transform delete eval request for OpenAI""" url = self.get_complete_url(api_base=api_base, endpoint="evals", eval_id=eval_id) @@ -217,10 +221,10 @@ class OpenAIEvalsConfig(BaseEvalsAPIConfig): logging_obj: LiteLLMLoggingObj, ) -> DeleteEvalResponse: """Transform OpenAI response to DeleteEvalResponse""" - response_json = raw_response.json() + response_json = _parsed_response_json(raw_response) verbose_logger.debug("Transforming delete eval response: %s", response_json) - return DeleteEvalResponse(**response_json) + return DeleteEvalResponse.model_validate(response_json) def transform_cancel_eval_request( self, @@ -228,12 +232,12 @@ class OpenAIEvalsConfig(BaseEvalsAPIConfig): api_base: str, litellm_params: GenericLiteLLMParams, headers: dict, - ) -> Tuple[str, Dict, Dict]: + ) -> tuple[str, dict, dict]: """Transform cancel eval request for OpenAI""" url = f"{self.get_complete_url(api_base=api_base, endpoint='evals', eval_id=eval_id)}/cancel" # Empty body for cancel request - request_body: Dict[str, Any] = {} + request_body: dict[str, object] = {} verbose_logger.debug("Cancel eval request - URL: %s", url) @@ -245,10 +249,10 @@ class OpenAIEvalsConfig(BaseEvalsAPIConfig): logging_obj: LiteLLMLoggingObj, ) -> CancelEvalResponse: """Transform OpenAI response to CancelEvalResponse""" - response_json = raw_response.json() + response_json = _parsed_response_json(raw_response) verbose_logger.debug("Transforming cancel eval response: %s", response_json) - return CancelEvalResponse(**response_json) + return CancelEvalResponse.model_validate(response_json) # Run API Transformations def transform_create_run_request( @@ -257,7 +261,7 @@ class OpenAIEvalsConfig(BaseEvalsAPIConfig): create_request: CreateRunRequest, litellm_params: GenericLiteLLMParams, headers: dict, - ) -> Tuple[str, Dict]: + ) -> tuple[str, dict]: """Transform create run request for OpenAI""" api_base = "https://api.openai.com" if litellm_params and litellm_params.api_base: @@ -279,10 +283,10 @@ class OpenAIEvalsConfig(BaseEvalsAPIConfig): logging_obj: LiteLLMLoggingObj, ) -> Run: """Transform OpenAI response to Run object""" - response_json = raw_response.json() + response_json = _parsed_response_json(raw_response) verbose_logger.debug("Transforming create run response: %s", response_json) - return Run(**response_json) + return Run.model_validate(response_json) def transform_list_runs_request( self, @@ -290,7 +294,7 @@ class OpenAIEvalsConfig(BaseEvalsAPIConfig): list_params: ListRunsParams, litellm_params: GenericLiteLLMParams, headers: dict, - ) -> Tuple[str, Dict]: + ) -> tuple[str, dict]: """Transform list runs request for OpenAI""" api_base = "https://api.openai.com" if litellm_params and litellm_params.api_base: @@ -300,7 +304,7 @@ class OpenAIEvalsConfig(BaseEvalsAPIConfig): url = f"{api_base}/v1/evals/{encoded_eval_id}/runs" # Build query parameters - query_params: Dict[str, Any] = {} + query_params: dict[str, object] = {} if "limit" in list_params and list_params["limit"]: query_params["limit"] = list_params["limit"] if "after" in list_params and list_params["after"]: @@ -323,10 +327,10 @@ class OpenAIEvalsConfig(BaseEvalsAPIConfig): logging_obj: LiteLLMLoggingObj, ) -> ListRunsResponse: """Transform OpenAI response to ListRunsResponse""" - response_json = raw_response.json() + response_json = _parsed_response_json(raw_response) verbose_logger.debug("Transforming list runs response: %s", response_json) - return ListRunsResponse(**response_json) + return ListRunsResponse.model_validate(response_json) def transform_get_run_request( self, @@ -335,7 +339,7 @@ class OpenAIEvalsConfig(BaseEvalsAPIConfig): api_base: str, litellm_params: GenericLiteLLMParams, headers: dict, - ) -> Tuple[str, Dict]: + ) -> tuple[str, dict]: """Transform get run request for OpenAI""" encoded_eval_id = encode_url_path_segment(eval_id, field_name="eval_id") encoded_run_id = encode_url_path_segment(run_id, field_name="run_id") @@ -351,10 +355,10 @@ class OpenAIEvalsConfig(BaseEvalsAPIConfig): logging_obj: LiteLLMLoggingObj, ) -> Run: """Transform OpenAI response to Run object""" - response_json = raw_response.json() + response_json = _parsed_response_json(raw_response) verbose_logger.debug("Transforming get run response: %s", response_json) - return Run(**response_json) + return Run.model_validate(response_json) def transform_cancel_run_request( self, @@ -363,14 +367,14 @@ class OpenAIEvalsConfig(BaseEvalsAPIConfig): api_base: str, litellm_params: GenericLiteLLMParams, headers: dict, - ) -> Tuple[str, Dict, Dict]: + ) -> tuple[str, dict, dict]: """Transform cancel run request for OpenAI""" encoded_eval_id = encode_url_path_segment(eval_id, field_name="eval_id") encoded_run_id = encode_url_path_segment(run_id, field_name="run_id") url = f"{api_base}/v1/evals/{encoded_eval_id}/runs/{encoded_run_id}/cancel" # Empty body for cancel request - request_body: Dict[str, Any] = {} + request_body: dict[str, object] = {} verbose_logger.debug("Cancel run request - URL: %s", url) @@ -382,10 +386,10 @@ class OpenAIEvalsConfig(BaseEvalsAPIConfig): logging_obj: LiteLLMLoggingObj, ) -> CancelRunResponse: """Transform OpenAI response to CancelRunResponse""" - response_json = raw_response.json() + response_json = _parsed_response_json(raw_response) verbose_logger.debug("Transforming cancel run response: %s", response_json) - return CancelRunResponse(**response_json) + return CancelRunResponse.model_validate(response_json) def transform_delete_run_request( self, @@ -394,14 +398,14 @@ class OpenAIEvalsConfig(BaseEvalsAPIConfig): api_base: str, litellm_params: GenericLiteLLMParams, headers: dict, - ) -> Tuple[str, Dict, Dict]: + ) -> tuple[str, dict, dict]: """Transform delete run request for OpenAI""" encoded_eval_id = encode_url_path_segment(eval_id, field_name="eval_id") encoded_run_id = encode_url_path_segment(run_id, field_name="run_id") url = f"{api_base}/v1/evals/{encoded_eval_id}/runs/{encoded_run_id}" # Empty body for delete request - request_body: Dict[str, Any] = {} + request_body: dict[str, object] = {} verbose_logger.debug("Delete run request - URL: %s", url) @@ -413,7 +417,7 @@ class OpenAIEvalsConfig(BaseEvalsAPIConfig): logging_obj: LiteLLMLoggingObj, ) -> RunDeleteResponse: """Transform OpenAI response to RunDeleteResponse""" - response_json = raw_response.json() + response_json = _parsed_response_json(raw_response) verbose_logger.debug("Transforming delete run response: %s", response_json) - return RunDeleteResponse(**response_json) + return RunDeleteResponse.model_validate(response_json) diff --git a/litellm/llms/volcengine/responses/transformation.py b/litellm/llms/volcengine/responses/transformation.py index 56950151969..4b20962e100 100644 --- a/litellm/llms/volcengine/responses/transformation.py +++ b/litellm/llms/volcengine/responses/transformation.py @@ -1,11 +1,9 @@ +from collections.abc import Callable, Mapping, Sequence from typing import ( TYPE_CHECKING, Any, - Dict, - List, Literal, - Optional, - Tuple, + Protocol, Union, get_args, get_origin, @@ -17,10 +15,10 @@ from pydantic import fields as pyd_fields import litellm from litellm._logging import verbose_logger from litellm.litellm_core_utils.core_helpers import process_response_headers -from litellm.litellm_core_utils.url_utils import encode_url_path_segment from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import ( _safe_convert_created_field, ) +from litellm.litellm_core_utils.url_utils import encode_url_path_segment from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import ( @@ -47,8 +45,15 @@ else: LiteLLMLoggingObj = Any +class _EventModelClass(Protocol): + @property + def model_fields(self) -> Mapping[str, pyd_fields.FieldInfo]: ... + + def model_validate(self, obj: Mapping[str, object]) -> ResponsesAPIStreamingResponse: ... + + class VolcEngineResponsesAPIConfig(OpenAIResponsesAPIConfig): - _SUPPORTED_OPTIONAL_PARAMS: List[str] = [ + _SUPPORTED_OPTIONAL_PARAMS: list[str] = [ # Doc-listed knobs "instructions", "max_output_tokens", @@ -89,9 +94,7 @@ class VolcEngineResponsesAPIConfig(OpenAIResponsesAPIConfig): supported.remove("metadata") return supported - def get_error_class( - self, error_message: str, status_code: int, headers: Union[dict, httpx.Headers] - ) -> VolcEngineError: + def get_error_class(self, error_message: str, status_code: int, headers: dict | httpx.Headers) -> VolcEngineError: typed_headers: httpx.Headers = headers if isinstance(headers, httpx.Headers) else httpx.Headers(headers or {}) return VolcEngineError( status_code=status_code, @@ -99,14 +102,14 @@ class VolcEngineResponsesAPIConfig(OpenAIResponsesAPIConfig): headers=typed_headers, ) - def validate_environment(self, headers: dict, model: str, litellm_params: Optional[GenericLiteLLMParams]) -> dict: + def validate_environment(self, headers: dict, model: str, litellm_params: GenericLiteLLMParams | None) -> dict: """ Build auth headers for Volcengine Responses API. """ if litellm_params is None: litellm_params = GenericLiteLLMParams() elif isinstance(litellm_params, dict): - litellm_params = GenericLiteLLMParams(**litellm_params) + litellm_params = GenericLiteLLMParams.model_validate(litellm_params) api_key = ( litellm_params.api_key @@ -122,7 +125,7 @@ class VolcEngineResponsesAPIConfig(OpenAIResponsesAPIConfig): def get_complete_url( self, - api_base: Optional[str], + api_base: str | None, litellm_params: dict, ) -> str: """ @@ -149,7 +152,7 @@ class VolcEngineResponsesAPIConfig(OpenAIResponsesAPIConfig): response_api_optional_params: ResponsesAPIOptionalRequestParams, model: str, drop_params: bool, - ) -> Dict: + ) -> dict: """ Volcengine Responses API aligns with OpenAI parameters. Remove parameters not supported by the public docs. @@ -173,11 +176,11 @@ class VolcEngineResponsesAPIConfig(OpenAIResponsesAPIConfig): def transform_responses_api_request( self, model: str, - input: Union[str, ResponseInputParam], - response_api_optional_request_params: Dict, + input: str | ResponseInputParam, + response_api_optional_request_params: dict, litellm_params: GenericLiteLLMParams, headers: dict, - ) -> Dict: + ) -> dict: """ Volcengine rejects any undocumented fields (including extra_body). Fail fast with clear errors and re-filter with the documented whitelist before delegating @@ -210,7 +213,7 @@ class VolcEngineResponsesAPIConfig(OpenAIResponsesAPIConfig): def transform_streaming_response( self, model: str, - parsed_chunk: dict, + parsed_chunk: Mapping[str, object], logging_obj: LiteLLMLoggingObj, ) -> ResponsesAPIStreamingResponse: """ @@ -222,18 +225,19 @@ class VolcEngineResponsesAPIConfig(OpenAIResponsesAPIConfig): if isinstance(chunk, dict): resp = chunk.get("response") if isinstance(resp, dict) and "output" not in resp: + resp_items: Mapping[str, object] = resp patched_chunk = dict(chunk) - patched_resp = dict(resp) + patched_resp = dict(resp_items) patched_resp["output"] = [] patched_chunk["response"] = patched_resp chunk = patched_chunk event_type = str(chunk.get("type")) if isinstance(chunk, dict) else None - event_pydantic_model = OpenAIResponsesAPIConfig.get_event_model_class(event_type=event_type) + event_pydantic_model: _EventModelClass = OpenAIResponsesAPIConfig.get_event_model_class(event_type=event_type) patched_chunk = self._fill_missing_fields(chunk, event_pydantic_model) - return event_pydantic_model(**patched_chunk) + return event_pydantic_model.model_validate(patched_chunk) def transform_response_api_response( self, @@ -246,7 +250,7 @@ class VolcEngineResponsesAPIConfig(OpenAIResponsesAPIConfig): original_response=raw_response.text, additional_args={"complete_input_dict": {}}, ) - raw_response_json = raw_response.json() + raw_response_json = self._parsed_response_body(raw_response) if "created_at" in raw_response_json: raw_response_json["created_at"] = _safe_convert_created_field(raw_response_json["created_at"]) except Exception: @@ -256,10 +260,11 @@ class VolcEngineResponsesAPIConfig(OpenAIResponsesAPIConfig): processed_headers = process_response_headers(raw_response_headers) try: - response = ResponsesAPIResponse(**raw_response_json) + response = ResponsesAPIResponse.model_validate(raw_response_json) except Exception: verbose_logger.debug("Volcengine Responses API: falling back to model_construct for response parsing.") - response = ResponsesAPIResponse.model_construct(**raw_response_json) + construct_response: Callable[..., ResponsesAPIResponse] = ResponsesAPIResponse.model_construct + response = construct_response(**raw_response_json) response._hidden_params["additional_headers"] = processed_headers response._hidden_params["headers"] = raw_response_headers @@ -274,10 +279,10 @@ class VolcEngineResponsesAPIConfig(OpenAIResponsesAPIConfig): api_base: str, litellm_params: GenericLiteLLMParams, headers: dict, - ) -> Tuple[str, Dict]: + ) -> tuple[str, dict]: encoded_response_id = encode_url_path_segment(response_id, field_name="response_id") url = f"{api_base}/{encoded_response_id}" - data: Dict = {} + data: dict = {} return url, data def transform_delete_response_api_response( @@ -286,16 +291,17 @@ class VolcEngineResponsesAPIConfig(OpenAIResponsesAPIConfig): logging_obj: LiteLLMLoggingObj, ) -> DeleteResponseResult: try: - raw_response_json = raw_response.json() + raw_response_json = self._parsed_response_body(raw_response) except Exception: raise VolcEngineError(message=raw_response.text, status_code=raw_response.status_code) try: - return DeleteResponseResult(**raw_response_json) + return DeleteResponseResult.model_validate(raw_response_json) except Exception: verbose_logger.debug( "Volcengine Responses API: falling back to model_construct for delete response parsing." ) - return DeleteResponseResult.model_construct(**raw_response_json) + construct_delete_result: Callable[..., DeleteResponseResult] = DeleteResponseResult.model_construct + return construct_delete_result(**raw_response_json) ######################################################### ########## GET RESPONSE API TRANSFORMATION ############### @@ -306,10 +312,10 @@ class VolcEngineResponsesAPIConfig(OpenAIResponsesAPIConfig): api_base: str, litellm_params: GenericLiteLLMParams, headers: dict, - ) -> Tuple[str, Dict]: + ) -> tuple[str, dict]: encoded_response_id = encode_url_path_segment(response_id, field_name="response_id") url = f"{api_base}/{encoded_response_id}" - data: Dict = {} + data: dict = {} return url, data def transform_get_response_api_response( @@ -318,14 +324,14 @@ class VolcEngineResponsesAPIConfig(OpenAIResponsesAPIConfig): logging_obj: LiteLLMLoggingObj, ) -> ResponsesAPIResponse: try: - raw_response_json = raw_response.json() + raw_response_json = self._parsed_response_body(raw_response) except Exception: raise VolcEngineError(message=raw_response.text, status_code=raw_response.status_code) raw_response_headers = dict(raw_response.headers) processed_headers = process_response_headers(raw_response_headers) - response = ResponsesAPIResponse(**raw_response_json) + response = ResponsesAPIResponse.model_validate(raw_response_json) response._hidden_params["additional_headers"] = processed_headers response._hidden_params["headers"] = raw_response_headers return response @@ -339,15 +345,15 @@ class VolcEngineResponsesAPIConfig(OpenAIResponsesAPIConfig): api_base: str, litellm_params: GenericLiteLLMParams, headers: dict, - after: Optional[str] = None, - before: Optional[str] = None, - include: Optional[List[str]] = None, + after: str | None = None, + before: str | None = None, + include: list[str] | None = None, limit: int = 20, order: Literal["asc", "desc"] = "desc", - ) -> Tuple[str, Dict]: + ) -> tuple[str, dict]: encoded_response_id = encode_url_path_segment(response_id, field_name="response_id") url = f"{api_base}/{encoded_response_id}/input_items" - params: Dict[str, Any] = {} + params: dict[str, str | int] = {} if after is not None: params["after"] = after if before is not None: @@ -364,9 +370,9 @@ class VolcEngineResponsesAPIConfig(OpenAIResponsesAPIConfig): self, raw_response: httpx.Response, logging_obj: LiteLLMLoggingObj, - ) -> Dict: + ) -> dict: try: - return raw_response.json() + return self._parsed_response_body(raw_response) except Exception: raise VolcEngineError(message=raw_response.text, status_code=raw_response.status_code) @@ -379,10 +385,10 @@ class VolcEngineResponsesAPIConfig(OpenAIResponsesAPIConfig): api_base: str, litellm_params: GenericLiteLLMParams, headers: dict, - ) -> Tuple[str, Dict]: + ) -> tuple[str, dict]: encoded_response_id = encode_url_path_segment(response_id, field_name="response_id") url = f"{api_base}/{encoded_response_id}/cancel" - data: Dict = {} + data: dict = {} return url, data def transform_cancel_response_api_response( @@ -391,23 +397,23 @@ class VolcEngineResponsesAPIConfig(OpenAIResponsesAPIConfig): logging_obj: LiteLLMLoggingObj, ) -> ResponsesAPIResponse: try: - raw_response_json = raw_response.json() + raw_response_json = self._parsed_response_body(raw_response) except Exception: raise VolcEngineError(message=raw_response.text, status_code=raw_response.status_code) raw_response_headers = dict(raw_response.headers) processed_headers = process_response_headers(raw_response_headers) - response = ResponsesAPIResponse(**raw_response_json) + response = ResponsesAPIResponse.model_validate(raw_response_json) response._hidden_params["additional_headers"] = processed_headers response._hidden_params["headers"] = raw_response_headers return response def should_fake_stream( self, - model: Optional[str], - stream: Optional[bool], - custom_llm_provider: Optional[str] = None, + model: str | None, + stream: bool | None, + custom_llm_provider: str | None = None, ) -> bool: """ Volcengine Responses API supports native streaming; never fall back to fake stream. @@ -415,7 +421,24 @@ class VolcEngineResponsesAPIConfig(OpenAIResponsesAPIConfig): return False @staticmethod - def _fill_missing_fields(chunk: Any, event_model: Any) -> Dict[str, Any]: + def _parsed_response_body(raw_response: httpx.Response) -> dict[str, object]: + return raw_response.json() + + @staticmethod + def _annotation_origin(annotation: object) -> object: + return get_origin(annotation) + + @staticmethod + def _annotation_args(annotation: object) -> tuple[object, ...]: + return get_args(annotation) + + @staticmethod + def _field_annotation(field: pyd_fields.FieldInfo) -> object: + annotation: object = field.annotation + return annotation + + @staticmethod + def _fill_missing_fields(chunk: Mapping[str, object], event_model: object | None) -> Mapping[str, object]: """ Heuristically fill missing required fields with safe defaults based on the event model's field annotations. This keeps parsing tolerant of providers that @@ -424,31 +447,37 @@ class VolcEngineResponsesAPIConfig(OpenAIResponsesAPIConfig): if not isinstance(chunk, dict) or event_model is None: return chunk - patched: Dict[str, Any] = dict(chunk) - fields_map = getattr(event_model, "model_fields", {}) or {} + patched = dict(chunk) + fields_map: Mapping[str, pyd_fields.FieldInfo] = getattr(event_model, "model_fields", {}) or {} for name, field in fields_map.items(): if name in patched: - patched[name] = VolcEngineResponsesAPIConfig._maybe_fill_nested(patched[name], field.annotation) + patched[name] = VolcEngineResponsesAPIConfig._maybe_fill_nested( + patched[name], VolcEngineResponsesAPIConfig._field_annotation(field) + ) continue # Explicit default or factory - if field.default is not pyd_fields.PydanticUndefined and field.default is not None: - patched[name] = field.default + field_default: object = field.default + if field_default is not pyd_fields.PydanticUndefined and field_default is not None: + patched[name] = field_default continue - if field.default_factory is not None and field.default_factory is not pyd_fields.PydanticUndefined: - patched[name] = field.default_factory() + default_factory: Callable[..., object] | None = field.default_factory + if default_factory is not None and default_factory is not pyd_fields.PydanticUndefined: + patched[name] = default_factory() continue # Heuristic defaults for missing required fields - patched[name] = VolcEngineResponsesAPIConfig._default_for_annotation(field.annotation) + patched[name] = VolcEngineResponsesAPIConfig._default_for_annotation( + VolcEngineResponsesAPIConfig._field_annotation(field) + ) return patched @staticmethod - def _default_for_annotation(annotation: Any) -> Any: - origin = get_origin(annotation) - args = get_args(annotation) + def _default_for_annotation(annotation: object) -> object: + origin = VolcEngineResponsesAPIConfig._annotation_origin(annotation) + args = VolcEngineResponsesAPIConfig._annotation_args(annotation) if annotation is int: return 0 @@ -456,7 +485,7 @@ class VolcEngineResponsesAPIConfig(OpenAIResponsesAPIConfig): return [] if origin is Union: # Prefer empty list when any option is a list - if any((arg is list or get_origin(arg) is list) for arg in args): + if any((arg is list or VolcEngineResponsesAPIConfig._annotation_origin(arg) is list) for arg in args): return [] if type(None) in args: return None @@ -467,53 +496,51 @@ class VolcEngineResponsesAPIConfig(OpenAIResponsesAPIConfig): return None @staticmethod - def _maybe_fill_nested(value: Any, annotation: Any) -> Any: + def _maybe_fill_nested(value: object, annotation: object) -> object: """ Recursively fill nested dict/list structures based on the annotated model. """ model_cls = VolcEngineResponsesAPIConfig._pick_model_class(annotation, value) - args = get_args(annotation) + args = VolcEngineResponsesAPIConfig._annotation_args(annotation) if isinstance(value, dict) and model_cls is not None: - return VolcEngineResponsesAPIConfig._fill_missing_fields(value, model_cls) + nested_items: Mapping[str, object] = value + return VolcEngineResponsesAPIConfig._fill_missing_fields(nested_items, model_cls) if isinstance(value, list): # Attempt to fill list elements if we know the element annotation - elem_ann: Any = args[0] if args else None + elem_ann: object = args[0] if args else None if elem_ann is not None: - return [VolcEngineResponsesAPIConfig._maybe_fill_nested(v, elem_ann) for v in value] + nested_elements: Sequence[object] = value + return [VolcEngineResponsesAPIConfig._maybe_fill_nested(v, elem_ann) for v in nested_elements] return value @staticmethod - def _pick_model_class(annotation: Any, value: Any) -> Optional[Any]: + def _pick_model_class(annotation: object, value: object) -> object | None: """ Choose the best-matching Pydantic model class for a nested dict. """ - candidates: List[Any] = [] - origin = get_origin(annotation) - - if hasattr(annotation, "model_fields"): - candidates.append(annotation) - if origin is Union: - for arg in get_args(annotation): - if hasattr(arg, "model_fields"): - candidates.append(arg) + origin = VolcEngineResponsesAPIConfig._annotation_origin(annotation) + union_args = VolcEngineResponsesAPIConfig._annotation_args(annotation) if origin is Union else () + candidates = tuple(candidate for candidate in (annotation, *union_args) if hasattr(candidate, "model_fields")) if not candidates: return None # Try to match by literal "type" field when available if isinstance(value, dict): - v_type = value.get("type") + value_items: Mapping[str, object] = value + v_type = value_items.get("type") for candidate in candidates: try: - type_field = candidate.model_fields.get("type") + candidate_fields: Mapping[str, pyd_fields.FieldInfo] = getattr(candidate, "model_fields") + type_field = candidate_fields.get("type") if type_field is None: continue - literal_ann = type_field.annotation - if get_origin(literal_ann) is Literal: - literal_values = get_args(literal_ann) + literal_ann = VolcEngineResponsesAPIConfig._field_annotation(type_field) + if VolcEngineResponsesAPIConfig._annotation_origin(literal_ann) is Literal: + literal_values = VolcEngineResponsesAPIConfig._annotation_args(literal_ann) if v_type in literal_values: return candidate except Exception: diff --git a/litellm/ocr/main.py b/litellm/ocr/main.py index 38f3f804e10..f53f32eecfa 100644 --- a/litellm/ocr/main.py +++ b/litellm/ocr/main.py @@ -7,9 +7,10 @@ import base64 import mimetypes import os import re +from collections.abc import Callable, Coroutine, Mapping from dataclasses import dataclass from io import IOBase -from typing import Any, Callable, Coroutine, Union, cast +from typing import Any, cast import httpx @@ -42,7 +43,7 @@ class _PreparedOCRRequest: provider_config: BaseOCRConfig optional_params: dict[str, object] litellm_params: dict[str, object] - effective_timeout: Union[float, httpx.Timeout] + effective_timeout: float | httpx.Timeout litellm_logging_obj: LiteLLMLoggingObj @@ -63,13 +64,13 @@ _RUST_OCR_PROVIDERS = { def _prepare_ocr_request( model: str, - document: dict[str, Any], + document: Mapping[str, object], api_key: str | None, api_base: str | None, - timeout: Union[float, httpx.Timeout] | None, + timeout: float | httpx.Timeout | None, custom_llm_provider: str | None, - extra_headers: dict[str, Any] | None, - kwargs: dict[str, Any], + extra_headers: dict[str, object] | None, + kwargs: dict[str, object], ) -> _PreparedOCRRequest: litellm_logging_obj = cast(LiteLLMLoggingObj, kwargs.pop("litellm_logging_obj")) litellm_call_id = cast(str | None, kwargs.get("litellm_call_id", None)) @@ -120,7 +121,7 @@ def _prepare_ocr_request( verbose_logger.debug(f"OCR call - model: {model}, provider: {custom_llm_provider}") - litellm_params = GenericLiteLLMParams(**kwargs) + litellm_params = GenericLiteLLMParams.model_validate(kwargs) supported_params = ocr_provider_config.get_supported_ocr_params(model=model) non_default_params = {} @@ -155,7 +156,7 @@ def _prepare_ocr_request( api_key=api_key, api_base=api_base, custom_llm_provider=custom_llm_provider, - extra_headers=cast(dict[str, object] | None, extra_headers), + extra_headers=extra_headers, provider_config=ocr_provider_config, optional_params=cast(dict[str, object], optional_params), litellm_params=dict(litellm_params), @@ -305,13 +306,13 @@ async def _run_rust_aocr( @client async def aocr( model: str, - document: dict[str, Any], + document: Mapping[str, object], api_key: str | None = None, api_base: str | None = None, - timeout: Union[float, httpx.Timeout] | None = None, + timeout: float | httpx.Timeout | None = None, custom_llm_provider: str | None = None, - extra_headers: dict[str, Any] | None = None, - **kwargs, + extra_headers: dict[str, object] | None = None, + **kwargs: object, ) -> OCRResponse: """ Async OCR function. @@ -567,14 +568,14 @@ def convert_file_document_to_url_document(document: dict[str, Any]) -> dict[str, @client def ocr( model: str, - document: dict[str, Any], + document: Mapping[str, object], api_key: str | None = None, api_base: str | None = None, - timeout: Union[float, httpx.Timeout] | None = None, + timeout: float | httpx.Timeout | None = None, custom_llm_provider: str | None = None, - extra_headers: dict[str, Any] | None = None, - **kwargs, -) -> Union[OCRResponse, Coroutine[Any, Any, OCRResponse]]: + extra_headers: dict[str, object] | None = None, + **kwargs: object, +) -> OCRResponse | Coroutine[object, object, OCRResponse]: """ Synchronous OCR function. diff --git a/litellm/proxy/_experimental/mcp_server/db.py b/litellm/proxy/_experimental/mcp_server/db.py index aeba74ca3ad..3221f3b8dd4 100644 --- a/litellm/proxy/_experimental/mcp_server/db.py +++ b/litellm/proxy/_experimental/mcp_server/db.py @@ -2,8 +2,14 @@ import base64 import binascii import hashlib import json +from collections.abc import Awaitable, Callable, Iterable, Mapping, Sequence from datetime import datetime, timedelta, timezone -from typing import TYPE_CHECKING, Any, Awaitable, Callable, Dict, Iterable, List, Optional, Set, Union, cast +from typing import ( + TYPE_CHECKING, + Any, + TypedDict, + cast, +) from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid @@ -13,8 +19,8 @@ from litellm.proxy._experimental.mcp_server.oauth_utils import build_upstream_oa from litellm.proxy._types import ( LiteLLM_MCPServerTable, LiteLLM_ObjectPermissionTable, - LiteLLM_TeamTable, MCPApprovalStatus, + MCPEnvVar, MCPEnvVarScope, MCPSubmissionsSummary, NewMCPServerRequest, @@ -42,9 +48,13 @@ from litellm.types.llms.custom_http import httpxSpecialProvider from litellm.types.mcp import MCPCredentials if TYPE_CHECKING: + from prisma import models as prisma_db_models + from prisma import types as prisma_db_types + from prisma.actions import LiteLLM_MCPUserCredentialsActions, LiteLLM_MCPUserEnvVarsActions + from litellm.types.mcp_server.mcp_server_manager import MCPServer -_AUTH_FLOW_SCOPED_FIELDS: frozenset = frozenset( +_AUTH_FLOW_SCOPED_FIELDS: "frozenset[str]" = frozenset( { "issuer", "authorization_url", @@ -60,7 +70,7 @@ _AUTH_FLOW_SCOPED_FIELDS: frozenset = frozenset( ) -def _blank_to_none(value: Optional[str]) -> Optional[str]: +def _blank_to_none(value: str | None) -> str | None: if not isinstance(value, str): return None return value.strip() or None @@ -73,7 +83,7 @@ def _blank_to_none(value: Optional[str]) -> Optional[str]: # the current code has never written — a cleared column can then never be # silently resurrected by a stale blob copy. These keys are stored plaintext # (endpoints/identifiers, not secrets), so values lift as-is. -_TOKEN_EXCHANGE_COLUMN_FIELDS: frozenset = frozenset( +_TOKEN_EXCHANGE_COLUMN_FIELDS: "frozenset[str]" = frozenset( { "token_exchange_endpoint", "audience", @@ -86,13 +96,33 @@ _TOKEN_EXCHANGE_COLUMN_FIELDS: frozenset = frozenset( # OAuth app (client_id/client_secret) plus the same authorize relay, and neither mints anything the # gateway keeps. So a switch WITHIN this class must preserve the stored app, unlike a cross-class # switch (e.g. an oauth2 row whose client may be DCR-minted and is not reusable elsewhere). -_CLIENT_FORWARDED_AUTH_TYPES: frozenset = frozenset({"true_passthrough", "oauth_delegate"}) +_CLIENT_FORWARDED_AUTH_TYPES: "frozenset[str]" = frozenset({"true_passthrough", "oauth_delegate"}) # Minted token material that must never survive a client rotation on a persisted row. -_MINTED_TOKEN_CREDENTIAL_FIELDS: frozenset = frozenset({"access_token", "refresh_token", "expires_in"}) +_MINTED_TOKEN_CREDENTIAL_FIELDS: "frozenset[str]" = frozenset({"access_token", "refresh_token", "expires_in"}) -def _credential_auth_class(auth_type: Optional[str]) -> Optional[str]: +class _OAuthCredentialAccessToken(TypedDict): + access_token: str + + +class OAuthCredentialPayload(_OAuthCredentialAccessToken, total=False): + type: str + refresh_token: str + expires_at: str + connected_at: str + scopes: list[str] + server_id: str + + +class _OAuthTokenRefreshResponse(TypedDict, total=False): + access_token: str + refresh_token: str + expires_in: int + scope: str + + +def _credential_auth_class(auth_type: str | None) -> str | None: """Collapse the client-forwarded modes to one credential class; every other auth_type is its own class. Used so credential handling keys off whether the stored-credential shape actually changed, not off a raw auth_type inequality that treats true_passthrough<->oauth_delegate as a full reset.""" @@ -101,7 +131,7 @@ def _credential_auth_class(auth_type: Optional[str]) -> Optional[str]: return auth_type -def _drop_stale_minted_on_client_rotation(merged: Dict[str, Any], new_creds: Dict[str, Any]) -> Dict[str, Any]: +def _drop_stale_minted_on_client_rotation(merged: dict[str, object], new_creds: dict[str, object]) -> dict[str, object]: """When the update rotates the client, drop stale minted token keys it did not itself set, so an old app's access/refresh token never rides forward under the new client. A no-op when no client key changed.""" if "client_id" not in new_creds and "client_secret" not in new_creds: @@ -111,13 +141,13 @@ def _drop_stale_minted_on_client_rotation(merged: Dict[str, Any], new_creds: Dic } -def _is_global_env_var_scope(scope: Any) -> bool: +def _is_global_env_var_scope(scope: object) -> bool: """``scope="user"`` entries are placeholders the user fills in; everything else (including a missing scope) is an admin-supplied global value.""" return scope != MCPEnvVarScope.user and scope != "user" -def _encrypt_global_env_var_values(env_vars: Iterable[Dict[str, Any]]) -> None: +def _encrypt_global_env_var_values(env_vars: Iterable[dict[str, str]]) -> None: """Encrypt ``scope="global"`` env var values in place before persisting. Global values hold admin-supplied secrets (API keys, passwords) that get @@ -133,7 +163,7 @@ def _encrypt_global_env_var_values(env_vars: Iterable[Dict[str, Any]]) -> None: entry["value"] = encrypt_value_helper(value) -def decrypt_global_env_var_values(env_vars: Optional[Iterable[Any]]) -> None: +def decrypt_global_env_var_values(env_vars: Iterable[MCPEnvVar | dict[str, str]] | None) -> None: """Decrypt ``scope="global"`` env var values in place after reading the DB. Accepts ``MCPEnvVar`` models (``LiteLLM_MCPServerTable``) or plain dicts @@ -172,7 +202,7 @@ def decrypt_global_env_var_values(env_vars: Optional[Iterable[Any]]) -> None: entry.value = decrypted -def _decrypt_env_vars_on_returned_row(row: Any) -> None: +def _decrypt_env_vars_on_returned_row(row: object) -> None: """Decrypt ``scope="global"`` env var values on a row returned by Prisma create/update. Prisma may hand back ``env_vars`` either as a parsed list (the common case for @@ -202,8 +232,8 @@ def _decrypt_env_vars_on_returned_row(row: Any) -> None: def _reencrypt_global_env_var_values( - env_vars: Optional[Iterable[Any]], new_encryption_key: str -) -> Optional[List[Dict[str, Any]]]: + env_vars: str | Iterable[Mapping[str, str]] | None, new_encryption_key: str +) -> list[dict[str, str]] | None: """Re-encrypt ``scope="global"`` env var values for master-key rotation. Each global value is decrypted with the current salt key and re-encrypted @@ -214,14 +244,17 @@ def _reencrypt_global_env_var_values( """ if not env_vars: return None + entries: Iterable[Mapping[str, str]] if isinstance(env_vars, str): try: - env_vars = json.loads(env_vars) + entries = json.loads(env_vars) except (json.JSONDecodeError, TypeError): return None - if not env_vars: + if not entries: return None - rebuilt = [dict(v) for v in env_vars] + else: + entries = env_vars + rebuilt = [dict(v) for v in entries] rotated = False for entry in rebuilt: if not _is_global_env_var_scope(entry.get("scope")): @@ -247,10 +280,10 @@ def _reencrypt_global_env_var_values( def _prepare_mcp_server_data( - data: Union[NewMCPServerRequest, UpdateMCPServerRequest], + data: NewMCPServerRequest | UpdateMCPServerRequest, exclude_unset: bool = False, - fields_set: Optional[Set[str]] = None, -) -> Dict[str, Any]: + fields_set: set[str] | None = None, +) -> dict[str, Any]: """ Helper function to prepare MCP server data for database operations. Handles JSON field serialization for mcp_info and env fields. @@ -326,7 +359,7 @@ def _prepare_mcp_server_data( # column so the exclude_unset filter is respected: a partial update that # omits env_vars never overwrites the stored value. Global values are # encrypted at rest before serialization. - env_vars = data_dict.get("env_vars") + env_vars: Sequence[Mapping[str, str]] | None = data_dict.get("env_vars") if env_vars is not None: serialized_env_vars = [dict(v) for v in env_vars] _encrypt_global_env_var_values(serialized_env_vars) @@ -353,7 +386,7 @@ def _prepare_mcp_server_data( return data_dict -def encrypt_credentials(credentials: MCPCredentials, encryption_key: Optional[str]) -> MCPCredentials: +def encrypt_credentials(credentials: MCPCredentials, encryption_key: str | None) -> MCPCredentials: auth_value = credentials.get("auth_value") if auth_value is not None: credentials["auth_value"] = encrypt_value_helper( @@ -401,6 +434,98 @@ def encrypt_credentials(credentials: MCPCredentials, encryption_key: Optional[st return credentials +def _credentials_blob_to_mutable_dict(blob: str | Mapping[str, object]) -> dict[str, object]: + parsed_blob: dict[str, object] = json.loads(blob) if isinstance(blob, str) else dict(blob) + return parsed_blob + + +async def _db_find_mcp_server_rows( + prisma_client: PrismaClient, + where: "prisma_db_types.LiteLLM_MCPServerTableWhereInput | None" = None, +) -> "list[prisma_db_models.LiteLLM_MCPServerTable]": + rows: list[prisma_db_models.LiteLLM_MCPServerTable] = await MCPServerRepository(prisma_client).table.find_many( + where=where + ) + return rows + + +async def _db_find_mcp_server_row( + prisma_client: PrismaClient, server_id: str +) -> "prisma_db_models.LiteLLM_MCPServerTable | None": + row: prisma_db_models.LiteLLM_MCPServerTable | None = await MCPServerRepository(prisma_client).table.find_unique( + where={"server_id": server_id} + ) + return row + + +async def _db_update_mcp_server_row( + prisma_client: PrismaClient, + server_id: str, + data: "prisma_db_types.LiteLLM_MCPServerTableUpdateInput", +) -> "prisma_db_models.LiteLLM_MCPServerTable": + row: prisma_db_models.LiteLLM_MCPServerTable = await MCPServerRepository(prisma_client).table.update( + where={"server_id": server_id}, + data=data, + ) + return row + + +def _user_credential_actions( + prisma_client: PrismaClient, +) -> "LiteLLM_MCPUserCredentialsActions[prisma_db_models.LiteLLM_MCPUserCredentials]": + table: LiteLLM_MCPUserCredentialsActions[prisma_db_models.LiteLLM_MCPUserCredentials] = ( + MCPUserCredentialsRepository(prisma_client).table + ) + return table + + +def _user_env_var_actions( + prisma_client: PrismaClient, +) -> "LiteLLM_MCPUserEnvVarsActions[prisma_db_models.LiteLLM_MCPUserEnvVars]": + table: LiteLLM_MCPUserEnvVarsActions[prisma_db_models.LiteLLM_MCPUserEnvVars] = ( + prisma_client.db.litellm_mcpuserenvvars + ) + return table + + +async def _db_find_user_credential_row( + prisma_client: PrismaClient, user_id: str, server_id: str +) -> "prisma_db_models.LiteLLM_MCPUserCredentials | None": + return await _user_credential_actions(prisma_client).find_unique( + where={"user_id_server_id": {"user_id": user_id, "server_id": server_id}} + ) + + +async def _db_find_user_credential_rows( + prisma_client: PrismaClient, + where: "prisma_db_types.LiteLLM_MCPUserCredentialsWhereInput | None" = None, +) -> "list[prisma_db_models.LiteLLM_MCPUserCredentials]": + return await _user_credential_actions(prisma_client).find_many(where=where) + + +async def _db_upsert_user_credential_row( + prisma_client: PrismaClient, user_id: str, server_id: str, credential_b64: str +) -> None: + await MCPUserCredentialsRepository(prisma_client).table.upsert( + where={"user_id_server_id": {"user_id": user_id, "server_id": server_id}}, + data={ + "create": { + "user_id": user_id, + "server_id": server_id, + "credential_b64": credential_b64, + }, + "update": {"credential_b64": credential_b64}, + }, + ) + + +async def _db_find_user_env_var_rows( + prisma_client: PrismaClient, + where: "prisma_db_types.LiteLLM_MCPUserEnvVarsWhereInput | None" = None, +) -> "list[prisma_db_models.LiteLLM_MCPUserEnvVars]": + return await _user_env_var_actions(prisma_client).find_many(where=where) + + def decrypt_credentials( credentials: MCPCredentials, ) -> MCPCredentials: @@ -428,19 +553,19 @@ def decrypt_credentials( async def get_all_mcp_servers( prisma_client: PrismaClient, - approval_status: Optional[str] = None, -) -> List[LiteLLM_MCPServerTable]: + approval_status: str | None = None, +) -> list[LiteLLM_MCPServerTable]: """ Returns mcp servers from the db, optionally filtered by approval_status. Pass approval_status=None to return all servers regardless of approval state. """ try: - where: Dict[str, Any] = {} + where: prisma_db_types.LiteLLM_MCPServerTableWhereInput = {} if approval_status is not None: where["approval_status"] = approval_status - mcp_servers = await MCPServerRepository(prisma_client).table.find_many(where=where if where else {}) + mcp_servers = await _db_find_mcp_server_rows(prisma_client, where if where else {}) - tables = [LiteLLM_MCPServerTable(**mcp_server.model_dump()) for mcp_server in mcp_servers] + tables = [LiteLLM_MCPServerTable.model_validate(mcp_server.model_dump()) for mcp_server in mcp_servers] for table in tables: decrypt_global_env_var_values(table.env_vars) return tables @@ -451,45 +576,45 @@ async def get_all_mcp_servers( return [] -async def get_mcp_server(prisma_client: PrismaClient, server_id: str) -> Optional[LiteLLM_MCPServerTable]: +async def get_mcp_server(prisma_client: PrismaClient, server_id: str) -> LiteLLM_MCPServerTable | None: """ Returns the matching mcp server from the db iff exists """ - mcp_server: Optional[LiteLLM_MCPServerTable] = await MCPServerRepository(prisma_client).table.find_unique( - where={ - "server_id": server_id, - } - ) + mcp_server = await _db_find_mcp_server_row(prisma_client, server_id) if mcp_server is None: return None - table = LiteLLM_MCPServerTable(**mcp_server.model_dump()) + table = LiteLLM_MCPServerTable.model_validate(mcp_server.model_dump()) decrypt_global_env_var_values(table.env_vars) return table -async def get_mcp_servers(prisma_client: PrismaClient, server_ids: Iterable[str]) -> List[LiteLLM_MCPServerTable]: +async def get_mcp_servers(prisma_client: PrismaClient, server_ids: Iterable[str]) -> list[LiteLLM_MCPServerTable]: """ Returns the matching mcp servers from the db with the server_ids """ - _mcp_servers: List[LiteLLM_MCPServerTable] = await MCPServerRepository(prisma_client).table.find_many( + _mcp_servers: list[prisma_db_models.LiteLLM_MCPServerTable] = await MCPServerRepository( + prisma_client + ).table.find_many( where={ "server_id": {"in": server_ids}, } ) - final_mcp_servers: List[LiteLLM_MCPServerTable] = [] + final_mcp_servers: list[LiteLLM_MCPServerTable] = [] for _mcp_server in _mcp_servers: - table = LiteLLM_MCPServerTable(**_mcp_server.model_dump()) + table = LiteLLM_MCPServerTable.model_validate(_mcp_server.model_dump()) decrypt_global_env_var_values(table.env_vars) final_mcp_servers.append(table) return final_mcp_servers -async def get_mcp_servers_by_verificationtoken(prisma_client: PrismaClient, token: str) -> List[str]: +async def get_mcp_servers_by_verificationtoken(prisma_client: PrismaClient, token: str) -> list[str]: """ Returns the mcp servers from the db for the verification token """ - verification_token_record: LiteLLM_TeamTable = await VerificationTokenRepository(prisma_client).table.find_unique( + verification_token_record: prisma_db_models.LiteLLM_VerificationToken | None = await VerificationTokenRepository( + prisma_client + ).table.find_unique( where={ "token": token, }, @@ -498,17 +623,17 @@ async def get_mcp_servers_by_verificationtoken(prisma_client: PrismaClient, toke }, ) - mcp_servers: Optional[List[str]] = [] + mcp_servers: list[str] | None = [] if verification_token_record is not None and verification_token_record.object_permission is not None: mcp_servers = verification_token_record.object_permission.mcp_servers return mcp_servers or [] -async def get_mcp_servers_by_team(prisma_client: PrismaClient, team_id: str) -> List[str]: +async def get_mcp_servers_by_team(prisma_client: PrismaClient, team_id: str) -> list[str]: """ Returns the mcp servers from the db for the team id """ - team_record: LiteLLM_TeamTable = await TeamRepository(prisma_client).table.find_unique( + team_record: prisma_db_models.LiteLLM_TeamTable | None = await TeamRepository(prisma_client).table.find_unique( where={ "team_id": team_id, }, @@ -517,7 +642,7 @@ async def get_mcp_servers_by_team(prisma_client: PrismaClient, team_id: str) -> }, ) - mcp_servers: Optional[List[str]] = [] + mcp_servers: list[str] | None = [] if team_record is not None and team_record.object_permission is not None: mcp_servers = team_record.object_permission.mcp_servers return mcp_servers or [] @@ -526,14 +651,14 @@ async def get_mcp_servers_by_team(prisma_client: PrismaClient, team_id: str) -> async def get_all_mcp_servers_for_user( prisma_client: PrismaClient, user: UserAPIKeyAuth, -) -> List[LiteLLM_MCPServerTable]: +) -> list[LiteLLM_MCPServerTable]: """ Get all the mcp servers filtered by the given user has access to. Following Least-Privilege Principle - the requestor should only be able to see the mcp servers that they have access to. """ - mcp_server_ids: Set[str] = set() + mcp_server_ids: set[str] = set() mcp_servers = [] # Get the mcp servers for the key @@ -554,11 +679,13 @@ async def get_all_mcp_servers_for_user( async def get_objectpermissions_for_mcp_server( prisma_client: PrismaClient, mcp_server_id: str -) -> List[LiteLLM_ObjectPermissionTable]: +) -> list[LiteLLM_ObjectPermissionTable]: """ Get all the object permissions records and the associated team and verficiationtoken records that have access to the mcp server """ - object_permission_records = await ObjectPermissionRepository(prisma_client).table.find_many( + object_permission_records: list[LiteLLM_ObjectPermissionTable] = await ObjectPermissionRepository( + prisma_client + ).table.find_many( where={ "mcp_servers": {"has": mcp_server_id}, }, @@ -571,11 +698,15 @@ async def get_objectpermissions_for_mcp_server( return object_permission_records -async def get_virtualkeys_for_mcp_server(prisma_client: PrismaClient, server_id: str) -> List: +async def get_virtualkeys_for_mcp_server( + prisma_client: PrismaClient, server_id: str +) -> "list[prisma_db_models.LiteLLM_VerificationToken]": """ Get all the virtual keys that have access to the mcp server """ - virtual_keys = await VerificationTokenRepository(prisma_client).table.find_many( + virtual_keys: list[prisma_db_models.LiteLLM_VerificationToken] | None = await VerificationTokenRepository( + prisma_client + ).table.find_many( where={ "mcp_servers": {"has": server_id}, }, @@ -603,8 +734,8 @@ async def delete_mcp_server_from_virtualkey(): async def delete_mcp_server( prisma_client: PrismaClient, server_id: str, - invalidate_token_cache: Optional[Callable[[str, str], Awaitable[None]]] = None, -) -> Optional[LiteLLM_MCPServerTable]: + invalidate_token_cache: Callable[[str, str], Awaitable[None]] | None = None, +) -> LiteLLM_MCPServerTable | None: """ Delete the mcp server from the db by server_id @@ -629,11 +760,11 @@ async def delete_mcp_server( }, ) if deleted_server is not None: - credential_user_ids: List[str] = [] + credential_user_ids: list[str] = [] try: - credential_rows = await prisma_client.db.litellm_mcpusercredentials.find_many( - where={"server_id": server_id} - ) + credential_rows: Sequence[ + prisma_db_models.LiteLLM_MCPUserCredentials + ] = await prisma_client.db.litellm_mcpusercredentials.find_many(where={"server_id": server_id}) credential_user_ids = [row.user_id for row in credential_rows] except Exception as e: # noqa: BLE001 - enumeration is best-effort; cached tokens expire by TTL verbose_proxy_logger.warning( @@ -684,7 +815,7 @@ async def create_mcp_server( data_dict["created_by"] = touched_by data_dict["updated_by"] = touched_by - new_mcp_server = await MCPServerRepository(prisma_client).table.create( + new_mcp_server: LiteLLM_MCPServerTable = await MCPServerRepository(prisma_client).table.create( data=data_dict # type: ignore ) @@ -696,13 +827,11 @@ async def update_mcp_server( prisma_client: PrismaClient, data: UpdateMCPServerRequest, touched_by: str, - fields_set: Optional[Set[str]] = None, + fields_set: set[str] | None = None, ) -> LiteLLM_MCPServerTable: """ Update a new mcp server record in the db """ - import json - from litellm.litellm_core_utils.safe_json_dumps import safe_dumps # Use helper to prepare data with proper JSON serialization. @@ -720,7 +849,7 @@ async def update_mcp_server( url_provided = "url" in data_dict and data_dict["url"] is not None issuer_provided = "issuer" in data_dict if data.auth_type or has_credentials or explicit_te_write or url_provided or issuer_provided: - existing = await MCPServerRepository(prisma_client).table.find_unique(where={"server_id": data.server_id}) + existing = await _db_find_mcp_server_row(prisma_client, data.server_id) auth_type_changed = bool( data.auth_type @@ -760,9 +889,7 @@ async def update_mcp_server( # repopulate the column the admin just cleared. (When credentials ARE in the # update, the merge below performs the same migration.) if explicit_te_write and "credentials" not in data_dict and existing is not None and existing.credentials: - existing_creds = ( - json.loads(existing.credentials) if isinstance(existing.credentials, str) else dict(existing.credentials) - ) + existing_creds = _credentials_blob_to_mutable_dict(existing.credentials) if _TOKEN_EXCHANGE_COLUMN_FIELDS & existing_creds.keys(): for te_field in _TOKEN_EXCHANGE_COLUMN_FIELDS: legacy_value = existing_creds.pop(te_field, None) @@ -781,16 +908,8 @@ async def update_mcp_server( # within the client-forwarded class (true_passthrough ↔ oauth_delegate) keeps # the same declared app and so must merge, not replace. if not auth_type_changed: - existing_creds = ( - json.loads(existing.credentials) - if isinstance(existing.credentials, str) - else dict(existing.credentials) - ) - new_creds = ( - json.loads(data_dict["credentials"]) - if isinstance(data_dict["credentials"], str) - else dict(data_dict["credentials"]) - ) + existing_creds = _credentials_blob_to_mutable_dict(existing.credentials) + new_creds = _credentials_blob_to_mutable_dict(data_dict["credentials"]) # New values override existing; existing keys not in update are preserved. A client # rotation additionally drops the previous app's stale minted token keys. merged = _drop_stale_minted_on_client_rotation({**existing_creds, **new_creds}, new_creds) @@ -820,7 +939,7 @@ async def update_mcp_server( data_dict["credentials"] = Json(None) - updated_mcp_server = await MCPServerRepository(prisma_client).table.update( + updated_mcp_server: LiteLLM_MCPServerTable = await MCPServerRepository(prisma_client).table.update( where={"server_id": data.server_id}, data=data_dict, # type: ignore ) @@ -835,7 +954,9 @@ async def get_mcp_server_oauth_client_credentials(prisma_client: PrismaClient, s LiteLLM_MCPServerTable row, so their dynamically registered client lives here keyed by server_id. The returned value is the raw credentials blob for ``_get_persisted_dcr_credentials`` to parse.""" - row = await MCPServerOAuthClientRepository(prisma_client).table.find_unique(where={"server_id": server_id}) + row: prisma_db_models.LiteLLM_MCPServerOAuthClient | None = await MCPServerOAuthClientRepository( + prisma_client + ).table.find_unique(where={"server_id": server_id}) if row is None: return None return row.credentials @@ -851,7 +972,7 @@ async def upsert_mcp_server_oauth_client_credentials( same way regardless of which store a server's client came from.""" from litellm.litellm_core_utils.safe_json_dumps import safe_dumps - encrypted = encrypt_credentials(credentials=dict(credentials), encryption_key=_get_salt_key()) + encrypted = encrypt_credentials(credentials=MCPCredentials(**credentials), encryption_key=_get_salt_key()) blob = safe_dumps(encrypted) await MCPServerOAuthClientRepository(prisma_client).table.upsert( where={"server_id": server_id}, @@ -862,7 +983,9 @@ async def upsert_mcp_server_oauth_client_credentials( ) -def _reencrypt_mcp_credentials_blob(credentials: object, new_master_key: str) -> str | None: +def _reencrypt_mcp_credentials_blob( + credentials: "str | Mapping[str, object] | None", new_master_key: str +) -> str | None: """Decrypt an at-rest MCP credentials blob with the current key and re-encrypt it under new_master_key, returning the serialized blob or None when there is nothing to rotate. Shared by every table that stores an encrypted MCP credentials blob so a master-key rotation covers them @@ -871,7 +994,7 @@ def _reencrypt_mcp_credentials_blob(credentials: object, new_master_key: str) -> return None from litellm.litellm_core_utils.safe_json_dumps import safe_dumps # noqa: PLC0415 # avoids circular import - creds_dict = json.loads(credentials) if isinstance(credentials, str) else dict(credentials) + creds_dict = _credentials_blob_to_mutable_dict(credentials) decrypted = decrypt_credentials(credentials=cast(MCPCredentials, creds_dict)) encrypted = encrypt_credentials(credentials=decrypted, encryption_key=new_master_key) return safe_dumps(encrypted) @@ -880,11 +1003,11 @@ def _reencrypt_mcp_credentials_blob(credentials: object, new_master_key: str) -> async def rotate_mcp_server_credentials_master_key(prisma_client: PrismaClient, touched_by: str, new_master_key: str): from litellm.litellm_core_utils.safe_json_dumps import safe_dumps # noqa: PLC0415 # avoids circular import - mcp_servers = await MCPServerRepository(prisma_client).table.find_many() + mcp_servers = await _db_find_mcp_server_rows(prisma_client) updated = 0 for mcp_server in mcp_servers: - update_data: Dict[str, Any] = {} + update_data: dict[str, str] = {} rotated_credentials = _reencrypt_mcp_credentials_blob(mcp_server.credentials, new_master_key) if rotated_credentials is not None: @@ -904,7 +1027,9 @@ async def rotate_mcp_server_credentials_master_key(prisma_client: PrismaClient, ) updated += 1 - oauth_clients = await MCPServerOAuthClientRepository(prisma_client).table.find_many() + oauth_clients: list[prisma_db_models.LiteLLM_MCPServerOAuthClient] = await MCPServerOAuthClientRepository( + prisma_client + ).table.find_many() oauth_updated = 0 for oauth_client in oauth_clients: rotated_credentials = _reencrypt_mcp_credentials_blob(oauth_client.credentials, new_master_key) @@ -923,7 +1048,7 @@ async def rotate_mcp_server_credentials_master_key(prisma_client: PrismaClient, ) -def _decode_user_credential(stored: str) -> Optional[str]: +def _decode_user_credential(stored: str) -> str | None: """Read back a value persisted in ``LiteLLM_MCPUserCredentials.credential_b64``. Tries nacl decryption first (current write format). Falls back to a @@ -945,7 +1070,7 @@ def _decode_user_credential(stored: str) -> Optional[str]: return None -def _decode_oauth_payload(stored: str) -> Optional[Dict[str, Any]]: +def _decode_oauth_payload(stored: str) -> OAuthCredentialPayload | None: """Return the OAuth2 payload dict if ``stored`` holds one, else ``None``. A row is considered an OAuth2 credential iff its decoded value parses as @@ -955,6 +1080,7 @@ def _decode_oauth_payload(stored: str) -> Optional[Dict[str, Any]]: decoded = _decode_user_credential(stored) if decoded is None: return None + parsed: OAuthCredentialPayload | None try: parsed = json.loads(decoded) except (ValueError, TypeError): @@ -972,7 +1098,7 @@ async def rotate_mcp_user_credentials_master_key(prisma_client: PrismaClient, ne under the new master key. Rows that are unreadable under both paths are logged and skipped so one corrupt row does not abort the rotation. """ - rows = await MCPUserCredentialsRepository(prisma_client).table.find_many() + rows = await _db_find_user_credential_rows(prisma_client) rotated = 0 skipped = 0 for row in rows: @@ -987,7 +1113,7 @@ async def rotate_mcp_user_credentials_master_key(prisma_client: PrismaClient, ne skipped += 1 continue re_encrypted = encrypt_value_helper(plaintext, new_encryption_key=new_master_key) - await MCPUserCredentialsRepository(prisma_client).table.update( + await _user_credential_actions(prisma_client).update( where={ "user_id_server_id": { "user_id": row.user_id, @@ -1012,7 +1138,7 @@ async def rotate_mcp_user_env_vars_master_key(prisma_client: PrismaClient, new_m skipped so one corrupt row does not abort the rotation nor overwrite values that may still be recoverable. """ - rows = await prisma_client.db.litellm_mcpuserenvvars.find_many() + rows = await _db_find_user_env_var_rows(prisma_client) rotated = 0 skipped = 0 for row in rows: @@ -1031,7 +1157,7 @@ async def rotate_mcp_user_env_vars_master_key(prisma_client: PrismaClient, new_m skipped += 1 continue re_encrypted = encrypt_value_helper(plaintext, new_encryption_key=new_master_key) - await prisma_client.db.litellm_mcpuserenvvars.update( + await _user_env_var_actions(prisma_client).update( where={ "user_id_server_id": { "user_id": row.user_id, @@ -1057,29 +1183,17 @@ async def store_user_credential( """Store a user credential for a BYOK MCP server.""" encoded = encrypt_value_helper(credential) - await MCPUserCredentialsRepository(prisma_client).table.upsert( - where={"user_id_server_id": {"user_id": user_id, "server_id": server_id}}, - data={ - "create": { - "user_id": user_id, - "server_id": server_id, - "credential_b64": encoded, - }, - "update": {"credential_b64": encoded}, - }, - ) + await _db_upsert_user_credential_row(prisma_client, user_id, server_id, encoded) async def get_user_credential( prisma_client: PrismaClient, user_id: str, server_id: str, -) -> Optional[str]: +) -> str | None: """Return credential for a user+server pair, or None.""" - row = await MCPUserCredentialsRepository(prisma_client).table.find_unique( - where={"user_id_server_id": {"user_id": user_id, "server_id": server_id}} - ) + row = await _db_find_user_credential_row(prisma_client, user_id, server_id) if row is None: return None return _decode_user_credential(row.credential_b64) @@ -1091,9 +1205,7 @@ async def has_user_credential( server_id: str, ) -> bool: """Return True if the user has a stored credential for this server.""" - row = await MCPUserCredentialsRepository(prisma_client).table.find_unique( - where={"user_id_server_id": {"user_id": user_id, "server_id": server_id}} - ) + row = await _db_find_user_credential_row(prisma_client, user_id, server_id) return row is not None @@ -1103,7 +1215,7 @@ async def delete_user_credential( server_id: str, ) -> None: """Delete the user's stored credential for a BYOK MCP server.""" - await MCPUserCredentialsRepository(prisma_client).table.delete( + await _user_credential_actions(prisma_client).delete( where={"user_id_server_id": {"user_id": user_id, "server_id": server_id}} ) @@ -1116,9 +1228,9 @@ async def store_user_oauth_credential( user_id: str, server_id: str, access_token: str, - refresh_token: Optional[str] = None, - expires_in: Optional[int] = None, - scopes: Optional[List[str]] = None, + refresh_token: str | None = None, + expires_in: int | None = None, + scopes: list[str] | None = None, skip_byok_guard: bool = False, ) -> None: """Persist an OAuth2 access token for a user+server pair. @@ -1128,11 +1240,11 @@ async def store_user_oauth_credential( differentiates it from plain BYOK API keys. """ - expires_at: Optional[str] = None + expires_at: str | None = None if expires_in is not None: expires_at = (datetime.now(timezone.utc) + timedelta(seconds=expires_in)).isoformat() - payload: Dict[str, Any] = { + payload: OAuthCredentialPayload = { "type": "oauth2", "access_token": access_token, "connected_at": datetime.now(timezone.utc).isoformat(), @@ -1148,9 +1260,7 @@ async def store_user_oauth_credential( # Skip the guard when the caller knows the row is already an OAuth2 credential # (e.g. during token refresh), saving an extra DB round-trip. if not skip_byok_guard: - existing = await MCPUserCredentialsRepository(prisma_client).table.find_unique( - where={"user_id_server_id": {"user_id": user_id, "server_id": server_id}} - ) + existing = await _db_find_user_credential_row(prisma_client, user_id, server_id) if existing is not None and _decode_oauth_payload(existing.credential_b64) is None: # Existing row is either a BYOK secret or an OAuth2 row that no # longer decrypts (e.g. after a salt-key rotation). In either @@ -1163,20 +1273,10 @@ async def store_user_oauth_credential( ) encoded = encrypt_value_helper(json.dumps(payload)) - await MCPUserCredentialsRepository(prisma_client).table.upsert( - where={"user_id_server_id": {"user_id": user_id, "server_id": server_id}}, - data={ - "create": { - "user_id": user_id, - "server_id": server_id, - "credential_b64": encoded, - }, - "update": {"credential_b64": encoded}, - }, - ) + await _db_upsert_user_credential_row(prisma_client, user_id, server_id, encoded) -def is_oauth_credential_expired(cred: Dict[str, Any], buffer_seconds: int = 0) -> bool: +def is_oauth_credential_expired(cred: OAuthCredentialPayload, buffer_seconds: int = 0) -> bool: """Return True if the OAuth2 credential's access_token has expired. Checks the ``expires_at`` ISO-format string stored in the credential payload. @@ -1201,12 +1301,10 @@ async def get_user_oauth_credential( prisma_client: PrismaClient, user_id: str, server_id: str, -) -> Optional[Dict[str, Any]]: +) -> OAuthCredentialPayload | None: """Return the decoded OAuth2 payload dict for a user+server pair, or None.""" - row = await MCPUserCredentialsRepository(prisma_client).table.find_unique( - where={"user_id_server_id": {"user_id": user_id, "server_id": server_id}} - ) + row = await _db_find_user_credential_row(prisma_client, user_id, server_id) if row is None: return None return _decode_oauth_payload(row.credential_b64) @@ -1215,11 +1313,11 @@ async def get_user_oauth_credential( async def list_user_oauth_credentials( prisma_client: PrismaClient, user_id: str, -) -> List[Dict[str, Any]]: +) -> list[OAuthCredentialPayload]: """Return all OAuth2 credential payloads for a user, tagged with server_id.""" - rows = await MCPUserCredentialsRepository(prisma_client).table.find_many(where={"user_id": user_id}) - results: List[Dict[str, Any]] = [] + rows = await _db_find_user_credential_rows(prisma_client, {"user_id": user_id}) + results: list[OAuthCredentialPayload] = [] for row in rows: payload = _decode_oauth_payload(row.credential_b64) if payload is None: @@ -1229,7 +1327,7 @@ async def list_user_oauth_credentials( return results -def _decrypted_credential_field(creds: Dict[str, object], field: str) -> object: +def _decrypted_credential_field(creds: dict[str, object], field: str) -> object: """Return one credential field decrypted with the global salt key; non-string and legacy plaintext values come back unchanged (decrypt_value_helper returns the original on failure).""" value = creds.get(field) @@ -1258,12 +1356,12 @@ def mcp_oauth_token_identity(server: object) -> tuple[object, ...]: creds = getattr(server, "credentials", None) if isinstance(creds, str): try: - parsed: object = json.loads(creds) + parsed: dict[str, object] | None = json.loads(creds) except ValueError: parsed = None else: parsed = creds - creds_dict: Dict[str, object] = parsed if isinstance(parsed, dict) else {} + creds_dict: dict[str, object] = parsed if isinstance(parsed, dict) else {} return ( getattr(server, "url", None), getattr(server, "spec_path", None), @@ -1283,7 +1381,7 @@ def mcp_oauth_token_identity(server: object) -> tuple[object, ...]: async def purge_user_oauth_credentials_for_server( prisma_client: PrismaClient, server_id: str, - invalidate_token_cache: Optional[Callable[[str, str], Awaitable[None]]] = None, + invalidate_token_cache: Callable[[str, str], Awaitable[None]] | None = None, ) -> int: """Delete every stored per-user OAuth token for a server and invalidate each user's cached token everywhere it can be served from (the legacy per-user token cache and the v2 per-user OAuth @@ -1301,12 +1399,11 @@ async def purge_user_oauth_credentials_for_server( invalidate_token_cache is injectable for tests; it defaults to the manager's shared invalidate_user_oauth_token_cache, the single invalidation point for per-user tokens.""" - repo = MCPUserCredentialsRepository(prisma_client) - rows = await repo.table.find_many(where={"server_id": server_id}) + rows = await _db_find_user_credential_rows(prisma_client, {"server_id": server_id}) oauth_rows = [row for row in rows if _decode_oauth_payload(row.credential_b64) is not None] if not oauth_rows: return 0 - deleted_count = await repo.table.delete_many( + deleted_count = await _user_credential_actions(prisma_client).delete_many( where={"server_id": server_id, "user_id": {"in": [row.user_id for row in oauth_rows]}} ) if invalidate_token_cache is None: @@ -1332,9 +1429,9 @@ async def purge_user_oauth_credentials_for_server( async def refresh_user_oauth_token( prisma_client: PrismaClient, user_id: str, - server: Any, - cred: Dict[str, Any], -) -> Optional[Dict[str, Any]]: + server: "MCPServer", + cred: OAuthCredentialPayload, +) -> OAuthCredentialPayload | None: """Attempt to refresh a per-user OAuth2 token using its stored refresh_token. POSTs to ``server.token_url`` with ``grant_type=refresh_token``. @@ -1345,11 +1442,11 @@ async def refresh_user_oauth_token( warning and returns ``None`` — the caller is responsible for clearing the stale credential and triggering re-authentication. """ - refresh_token: Optional[str] = cred.get("refresh_token") - token_url: Optional[str] = getattr(server, "token_url", None) + refresh_token: str | None = cred.get("refresh_token") + token_url: str | None = getattr(server, "token_url", None) server_id: str = getattr(server, "server_id", "") - client_id: Optional[str] = getattr(server, "client_id", None) - client_secret: Optional[str] = getattr(server, "client_secret", None) + client_id: str | None = getattr(server, "client_id", None) + client_secret: str | None = getattr(server, "client_secret", None) if not refresh_token: verbose_proxy_logger.debug( @@ -1372,7 +1469,7 @@ async def refresh_user_oauth_token( client_id=client_id, client_secret=client_secret, ) - token_data: Dict[str, str] = { + token_data: dict[str, str] = { "grant_type": "refresh_token", "refresh_token": refresh_token, **token_request.body, @@ -1384,7 +1481,7 @@ async def refresh_user_oauth_token( data=token_data, ) response.raise_for_status() - body: Dict[str, Any] = response.json() + body: _OAuthTokenRefreshResponse = response.json() except Exception as exc: verbose_proxy_logger.warning( "refresh_user_oauth_token: refresh request failed for user=%s server=%s: %s", @@ -1394,7 +1491,7 @@ async def refresh_user_oauth_token( ) return None - access_token: Optional[str] = body.get("access_token") + access_token: str | None = body.get("access_token") if not access_token: verbose_proxy_logger.warning( "refresh_user_oauth_token: token response missing access_token for user=%s server=%s", @@ -1403,7 +1500,7 @@ async def refresh_user_oauth_token( ) return None - expires_in: Optional[int] = None + expires_in: int | None = None raw_expires = body.get("expires_in") try: expires_in = int(raw_expires) if raw_expires is not None else None @@ -1411,10 +1508,10 @@ async def refresh_user_oauth_token( pass # Rotate refresh token when the provider returns a new one - new_refresh_token: Optional[str] = body.get("refresh_token") or refresh_token + new_refresh_token: str | None = body.get("refresh_token") or refresh_token raw_scope = body.get("scope") - scopes: Optional[List[str]] = (raw_scope.split() if isinstance(raw_scope, str) and raw_scope else None) or cred.get( + scopes: list[str] | None = (raw_scope.split() if isinstance(raw_scope, str) and raw_scope else None) or cred.get( "scopes" ) @@ -1439,10 +1536,10 @@ async def refresh_user_oauth_token( async def resolve_valid_user_oauth_token( user_id: str, - server: Any, - cred: Optional[Dict[str, Any]], - prisma_client: Optional[PrismaClient] = None, -) -> Optional[Dict[str, Any]]: + server: "MCPServer", + cred: OAuthCredentialPayload | None, + prisma_client: PrismaClient | None = None, +) -> OAuthCredentialPayload | None: """Return an OAuth2 credential whose access_token is good for the next request. Returns the credential unchanged while its token is valid for at least @@ -1480,7 +1577,7 @@ async def resolve_valid_user_oauth_token( async def resolve_user_oauth_access_token( user_id: str | None, server: "MCPServer", - prefetched_creds: dict[str, dict[str, object]] | None = None, + prefetched_creds: Mapping[str, OAuthCredentialPayload] | None = None, ) -> str | None: """Resolve a user's valid OAuth2 access token for a server: Redis cache, else DB + refresh. @@ -1491,7 +1588,7 @@ async def resolve_user_oauth_access_token( usable token; any error is swallowed to ``None`` so a transient failure reads as "not authorized" rather than raising. """ - server_id = getattr(server, "server_id", None) + server_id: str | None = getattr(server, "server_id", None) if not user_id or not server_id: return None try: @@ -1568,8 +1665,9 @@ async def get_active_submitted_mcp_server_ids_for_user( if not user_id: return [] - rows = await MCPServerRepository(prisma_client).table.find_many( - where={ + rows = await _db_find_mcp_server_rows( + prisma_client, + { "submitted_by": user_id, "approval_status": MCPApprovalStatus.active, }, @@ -1584,15 +1682,16 @@ async def approve_mcp_server( ) -> LiteLLM_MCPServerTable: """Set approval_status=active and record reviewed_at.""" now = datetime.now(timezone.utc) - updated = await MCPServerRepository(prisma_client).table.update( - where={"server_id": server_id}, - data={ + updated = await _db_update_mcp_server_row( + prisma_client, + server_id, + { "approval_status": MCPApprovalStatus.active, "reviewed_at": now, "updated_by": touched_by, }, ) - table = LiteLLM_MCPServerTable(**updated.model_dump()) + table = LiteLLM_MCPServerTable.model_validate(updated.model_dump()) decrypt_global_env_var_values(table.env_vars) return table @@ -1601,22 +1700,19 @@ async def reject_mcp_server( prisma_client: PrismaClient, server_id: str, touched_by: str, - review_notes: Optional[str] = None, + review_notes: str | None = None, ) -> LiteLLM_MCPServerTable: """Set approval_status=rejected, record reviewed_at and review_notes.""" now = datetime.now(timezone.utc) - data: Dict[str, Any] = { + data: prisma_db_types.LiteLLM_MCPServerTableUpdateInput = { "approval_status": MCPApprovalStatus.rejected, "reviewed_at": now, "updated_by": touched_by, } if review_notes is not None: data["review_notes"] = review_notes - updated = await MCPServerRepository(prisma_client).table.update( - where={"server_id": server_id}, - data=data, - ) - table = LiteLLM_MCPServerTable(**updated.model_dump()) + updated = await _db_update_mcp_server_row(prisma_client, server_id, data) + table = LiteLLM_MCPServerTable.model_validate(updated.model_dump()) decrypt_global_env_var_values(table.env_vars) return table @@ -1629,12 +1725,12 @@ async def get_mcp_submissions( along with a summary count breakdown by approval_status. Mirrors get_guardrail_submissions() from guardrail_endpoints.py. """ - rows = await MCPServerRepository(prisma_client).table.find_many( + rows: list[prisma_db_models.LiteLLM_MCPServerTable] = await MCPServerRepository(prisma_client).table.find_many( where={"submitted_at": {"not": None}}, order={"submitted_at": "desc"}, take=500, # safety cap; paginate if needed in a future iteration ) - items = [LiteLLM_MCPServerTable(**r.model_dump()) for r in rows] + items = [LiteLLM_MCPServerTable.model_validate(r.model_dump()) for r in rows] for item in items: decrypt_global_env_var_values(item.env_vars) @@ -1654,7 +1750,7 @@ async def get_mcp_submissions( # ── Per-user MCP environment variables ──────────────────────────────────── -def _decode_user_env_vars(stored: str) -> Dict[str, str]: +def _decode_user_env_vars(stored: str) -> dict[str, str]: """Decrypt a ``values_b64`` blob and parse it as a flat ``{name: value}`` dict.""" decrypted = decrypt_value_helper( value=stored, @@ -1670,6 +1766,7 @@ def _decode_user_env_vars(stored: str) -> Dict[str, str]: "re-enter them rather than silently forwarding ciphertext" ) return {} + parsed: dict[str, object] | None try: parsed = json.loads(decrypted) except (ValueError, TypeError): @@ -1683,9 +1780,9 @@ async def get_user_env_vars( prisma_client: PrismaClient, user_id: str, server_id: str, -) -> Dict[str, str]: +) -> dict[str, str]: """Return the calling user's env var dict for ``server_id`` (empty if none).""" - row = await prisma_client.db.litellm_mcpuserenvvars.find_unique( + row = await _user_env_var_actions(prisma_client).find_unique( where={"user_id_server_id": {"user_id": user_id, "server_id": server_id}} ) if row is None: @@ -1697,7 +1794,7 @@ async def get_user_env_vars_bulk( prisma_client: PrismaClient, user_id: str, server_ids: Iterable[str], -) -> Dict[str, Dict[str, str]]: +) -> dict[str, dict[str, str]]: """Return ``{server_id: {var_name: value}}`` for one user across many servers. Servers with no stored row are simply absent from the result. @@ -1705,7 +1802,7 @@ async def get_user_env_vars_bulk( ids = list(server_ids) if not ids: return {} - rows = await prisma_client.db.litellm_mcpuserenvvars.find_many(where={"user_id": user_id, "server_id": {"in": ids}}) + rows = await _db_find_user_env_var_rows(prisma_client, {"user_id": user_id, "server_id": {"in": ids}}) return {row.server_id: _decode_user_env_vars(row.values_b64) for row in rows} @@ -1713,9 +1810,9 @@ async def merge_user_env_vars( prisma_client: PrismaClient, user_id: str, server_id: str, - updates: Dict[str, str], + updates: dict[str, str], allowed_names: Iterable[str], -) -> Dict[str, str]: +) -> dict[str, str]: """Merge ``updates`` into the user's stored env vars for ``server_id`` and return the resulting set. @@ -1732,7 +1829,7 @@ async def merge_user_env_vars( ) async with prisma_client.db.tx() as tx: await tx.execute_raw("SELECT pg_advisory_xact_lock($1::bigint)", lock_key) - row = await tx.litellm_mcpuserenvvars.find_unique( + row: prisma_db_models.LiteLLM_MCPUserEnvVars | None = await tx.litellm_mcpuserenvvars.find_unique( where={"user_id_server_id": {"user_id": user_id, "server_id": server_id}} ) existing = _decode_user_env_vars(row.values_b64) if row is not None else {} @@ -1762,4 +1859,4 @@ async def delete_user_env_vars( Uses ``delete_many`` so a missing row is a no-op; real DB errors still propagate to the caller instead of being silently swallowed. """ - await prisma_client.db.litellm_mcpuserenvvars.delete_many(where={"user_id": user_id, "server_id": server_id}) + await _user_env_var_actions(prisma_client).delete_many(where={"user_id": user_id, "server_id": server_id}) diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/per_user_oauth_store.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/per_user_oauth_store.py index 21001c09f25..3a2c748bb82 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/per_user_oauth_store.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/per_user_oauth_store.py @@ -11,7 +11,7 @@ collaborators acquire their globals per call, mirroring v1's lazy-import pattern from __future__ import annotations import asyncio -from collections.abc import Callable +from collections.abc import Callable, Mapping from typing import TYPE_CHECKING from litellm._logging import verbose_logger @@ -54,7 +54,7 @@ ServerLookup = Callable[[str], "MCPServer | None"] StoreBuilder = Callable[[ServerLookup], tuple[InvalidatableOAuthTokenStore, bool]] -async def _read_credential(user_id: str, server_id: str) -> dict[str, object] | None: +async def _read_credential(user_id: str, server_id: str) -> Mapping[str, object] | None: from litellm.proxy._experimental.mcp_server.db import ( # noqa: PLC0415 get_user_oauth_credential, ) diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/v2_token_store.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/v2_token_store.py index f1b68042c94..eefeec84bfa 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/v2_token_store.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/v2_token_store.py @@ -10,14 +10,14 @@ injected, so the DB/decoding plumbing stays testable and out of this seam. from __future__ import annotations -from collections.abc import Awaitable, Callable +from collections.abc import Awaitable, Callable, Mapping from datetime import datetime, timezone from litellm.proxy._experimental.mcp_server.outbound_credentials.oauth_token_store import ( OAuthToken, ) -CredentialReader = Callable[[str, str], Awaitable["dict[str, object] | None"]] +CredentialReader = Callable[[str, str], Awaitable["Mapping[str, object] | None"]] def _iso_to_epoch(expires_at: str) -> float | None: @@ -39,7 +39,7 @@ def _to_scopes(raw: object) -> tuple[str, ...]: return () -def _to_oauth_token(payload: dict[str, object]) -> OAuthToken | None: +def _to_oauth_token(payload: Mapping[str, object]) -> OAuthToken | None: access_token = payload.get("access_token") if not isinstance(access_token, str): return None diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index 26e4176e09b..af3d966c95b 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -1,18 +1,11 @@ import asyncio import importlib +from collections.abc import Awaitable, Callable, Mapping from datetime import datetime from typing import ( + TYPE_CHECKING, Any, - Awaitable, - Callable, - Dict, - List, Literal, - Mapping, - Optional, - Set, - Tuple, - Union, ) import httpx @@ -38,6 +31,9 @@ from litellm.proxy._experimental.mcp_server.utils import ( from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.auth.ip_address_utils import IPAddressUtils from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + +if TYPE_CHECKING: + from litellm.proxy._experimental.mcp_server.db import OAuthCredentialPayload from litellm.proxy.common_utils.http_parsing_utils import _safe_get_request_headers from litellm.types.mcp import MCPAuth from litellm.types.utils import CallTypes @@ -97,12 +93,12 @@ if MCP_AVAILABLE: ######################################################## ############ MCP Server REST API Routes ################# async def _safe_fire_mcp_tool_call_logging( - logging_obj: Optional[Any], + logging_obj: Any | None, result: Any, start_time: datetime, end_time: datetime, - user_api_key_auth: Optional[UserAPIKeyAuth] = None, - request_data: Optional[Mapping[str, object]] = None, + user_api_key_auth: UserAPIKeyAuth | None = None, + request_data: Mapping[str, object] | None = None, ) -> None: if logging_obj is None: return @@ -134,7 +130,7 @@ if MCP_AVAILABLE: async def _handle_virtual_mcp_tool( request: Request, - data: Dict[str, Any], + data: dict[str, Any], tool_name: str, user_api_key_dict: UserAPIKeyAuth, ) -> Any: @@ -212,9 +208,9 @@ if MCP_AVAILABLE: def _get_server_auth_header( server, - mcp_server_auth_headers: Optional[Dict[str, Dict[str, str]]], - mcp_auth_header: Optional[str], - ) -> Optional[Union[Dict[str, str], str]]: + mcp_server_auth_headers: dict[str, dict[str, str]] | None, + mcp_auth_header: str | None, + ) -> dict[str, str] | str | None: """Helper function to get server-specific auth header with case-insensitive matching.""" from litellm.proxy._experimental.mcp_server.utils import ( lookup_mcp_server_auth_in_headers, @@ -230,7 +226,7 @@ if MCP_AVAILABLE: return server_auth return mcp_auth_header - def _is_v1_resolved_oauth2_server(server: Optional[MCPServer]) -> bool: + def _is_v1_resolved_oauth2_server(server: MCPServer | None) -> bool: """Whether this server's per-user OAuth2 token is still resolved by v1. A server the v2 resolver owns reads its stored token from the resolver at connect @@ -246,7 +242,7 @@ if MCP_AVAILABLE: return False return to_server_spec(server) is None - def _v1_resolved_oauth2_server_ids(allowed_server_ids: List[str]) -> Set[str]: + def _v1_resolved_oauth2_server_ids(allowed_server_ids: list[str]) -> set[str]: """Return the subset of *allowed_server_ids* whose per-user OAuth2 token is still resolved by v1. @@ -260,10 +256,10 @@ if MCP_AVAILABLE: } async def _get_user_oauth_extra_headers( - server, + server: MCPServer, user_api_key_dict: UserAPIKeyAuth, - prefetched_creds: Optional[Dict[str, Dict[str, Any]]] = None, - ) -> Optional[Dict[str, str]]: + prefetched_creds: dict[str, "OAuthCredentialPayload"] | None = None, + ) -> dict[str, str] | None: """ For OAuth2 servers, look up the user's stored access token and return it as extra_headers {"Authorization": "Bearer "} so that it reaches @@ -315,7 +311,7 @@ if MCP_AVAILABLE: async def _prefetch_user_oauth_creds( user_api_key_dict: UserAPIKeyAuth, - ) -> Dict[str, Dict[str, Any]]: + ) -> dict[str, "OAuthCredentialPayload"]: """Fetch all OAuth2 credentials for the user in a single DB query. Returns a dict keyed by server_id. Used to avoid N+1 DB queries when @@ -379,8 +375,8 @@ if MCP_AVAILABLE: def _resolve_mcp_server_id_for_rest( server_id: str, - allowed_server_ids: Union[Set[str], List[str]], - client_ip: Optional[str] = None, + allowed_server_ids: set[str] | list[str], + client_ip: str | None = None, ) -> str: """ Map REST ``server_id`` (UUID, server_name, or alias) to canonical server_id. @@ -400,7 +396,7 @@ if MCP_AVAILABLE: request: Request, user_api_key_dict: UserAPIKeyAuth, server_id: str, - ) -> Tuple[List[MCPServer], str]: + ) -> tuple[list[MCPServer], str]: """ Resolve allowed MCP servers for a tool call with IP filtering. @@ -471,7 +467,7 @@ if MCP_AVAILABLE: ) # Build allowed_mcp_servers list (only include allowed servers) - allowed_mcp_servers: List[MCPServer] = [] + allowed_mcp_servers: list[MCPServer] = [] for allowed_server_id in allowed_server_ids_set: server = global_mcp_server_manager.get_mcp_server_by_id(allowed_server_id) if server is not None: @@ -482,9 +478,9 @@ if MCP_AVAILABLE: async def _get_tools_for_single_server( server, server_auth_header, - raw_headers: Optional[Dict[str, str]] = None, - user_api_key_auth: Optional[UserAPIKeyAuth] = None, - extra_headers: Optional[Dict[str, str]] = None, + raw_headers: dict[str, str] | None = None, + user_api_key_auth: UserAPIKeyAuth | None = None, + extra_headers: dict[str, str] | None = None, apply_tool_filters: bool = True, ): """Helper function to get tools for a single server. @@ -530,7 +526,7 @@ if MCP_AVAILABLE: async def _resolve_allowed_mcp_servers_for_tool_call( user_api_key_dict: UserAPIKeyAuth, server_id: str, - ) -> List[MCPServer]: + ) -> list[MCPServer]: """Resolve allowed MCP servers for the given user and validate server_id access.""" auth_contexts = await build_effective_auth_contexts(user_api_key_dict) allowed_server_ids_set = set() @@ -545,7 +541,7 @@ if MCP_AVAILABLE: "message": f"The key is not allowed to access server {server_id}", }, ) - allowed_mcp_servers: List[MCPServer] = [] + allowed_mcp_servers: list[MCPServer] = [] for allowed_server_id in allowed_server_ids_set: server = global_mcp_server_manager.get_mcp_server_by_id(allowed_server_id) if server is not None: @@ -554,10 +550,10 @@ if MCP_AVAILABLE: async def _list_tools_for_single_server( server_id: str, - allowed_server_ids: List[str], - rest_client_ip: Optional[str], + allowed_server_ids: list[str], + rest_client_ip: str | None, mcp_server_auth_headers: dict, - mcp_auth_header: Optional[str], + mcp_auth_header: str | None, raw_headers_from_request: dict, user_api_key_dict: UserAPIKeyAuth, apply_tool_filters: bool = True, @@ -644,12 +640,12 @@ if MCP_AVAILABLE: "message": "Successfully retrieved tools", } - def _as_query_str(value: Any) -> Optional[str]: + def _as_query_str(value: Any) -> str | None: """Coerce an Optional[str] Query param to str|None, dropping unresolved FastAPI defaults.""" return value if isinstance(value, str) else None async def _resolve_toolset_scope( - toolset_name: Optional[str], + toolset_name: str | None, user_api_key_dict: UserAPIKeyAuth, ) -> UserAPIKeyAuth: """Resolve ``toolset_name`` to its scoped ``UserAPIKeyAuth``, or return unchanged.""" @@ -670,11 +666,9 @@ if MCP_AVAILABLE: @router.get("/tools/list", dependencies=[Depends(user_api_key_auth)]) async def list_tool_rest_api( request: Request, - server_id: Optional[str] = Query(None, description="The server id to list tools for"), - mcp_server_name: Optional[str] = Query( - None, description="Filter tools to a single MCP server by name or alias" - ), - toolset_name: Optional[str] = Query(None, description="Filter tools to a single toolset by name"), + server_id: str | None = Query(None, description="The server id to list tools for"), + mcp_server_name: str | None = Query(None, description="Filter tools to a single MCP server by name or alias"), + toolset_name: str | None = Query(None, description="Filter tools to a single toolset by name"), include_disabled_tools: bool = Query( False, description=( @@ -981,7 +975,7 @@ if MCP_AVAILABLE: ) = await _resolve_allowed_mcp_servers_with_ip_filter(request, user_api_key_dict, server_id) # Look up per-user OAuth headers for this server (mirrors list_tool_rest_api). - user_oauth_extra_headers: Optional[Dict[str, str]] = None + user_oauth_extra_headers: dict[str, str] | None = None target_server = next( (s for s in allowed_mcp_servers if s.server_id == canonical_server_id), None, @@ -1094,18 +1088,18 @@ if MCP_AVAILABLE: (client_id, client_secret, scopes) — any value may be ``None``. """ creds = request.credentials if isinstance(request.credentials, dict) else {} - client_id: Optional[str] = creds.get("client_id") - client_secret: Optional[str] = creds.get("client_secret") + client_id: str | None = creds.get("client_id") + client_secret: str | None = creds.get("client_secret") scopes_raw = creds.get("scopes") - scopes: Optional[List[str]] = scopes_raw if isinstance(scopes_raw, list) else None + scopes: list[str] | None = scopes_raw if isinstance(scopes_raw, list) else None return client_id, client_secret, scopes async def _execute_with_mcp_client( request: NewMCPServerRequest, operation: Callable[..., Awaitable[Any]], - mcp_auth_header: Optional[Union[str, Dict[str, str]]] = None, - oauth2_headers: Optional[Dict[str, str]] = None, - raw_headers: Optional[Dict[str, str]] = None, + mcp_auth_header: str | dict[str, str] | None = None, + oauth2_headers: dict[str, str] | None = None, + raw_headers: dict[str, str] | None = None, ) -> dict: """ Create a temporary MCP client from *request*, run *operation*, and return the result. @@ -1128,7 +1122,7 @@ if MCP_AVAILABLE: try: client_id, client_secret, scopes = _extract_credentials(request) - _oauth2_flow: Optional[Literal["client_credentials", "authorization_code"]] = request.oauth2_flow or ( + _oauth2_flow: Literal["client_credentials", "authorization_code"] | None = request.oauth2_flow or ( "client_credentials" if client_id and client_secret and request.token_url else None ) # client_credentials requires token_url to fetch a token; without it the @@ -1244,7 +1238,7 @@ if MCP_AVAILABLE: spec = await load_openapi_spec_async(spec_path) paths = spec.get("paths", {}) components = spec.get("components", {}) - tools: List[dict] = [] + tools: list[dict] = [] used_names: set = set() for path, path_item in paths.items(): for method in ("get", "post", "put", "delete", "patch"): @@ -1351,7 +1345,7 @@ if MCP_AVAILABLE: headers = request.headers - mcp_auth_header: Optional[str] = None + mcp_auth_header: str | None = None if new_mcp_server_request.auth_type in { MCPAuth.api_key, MCPAuth.bearer_token, @@ -1365,7 +1359,7 @@ if MCP_AVAILABLE: # Authorization doubles as the admission fallback (LITELLM_API_KEY_HEADER_NAME_SECONDARY): # when the primary x-litellm-api-key header is absent, the Authorization value is the # caller's LiteLLM key, not an upstream token, and must never be forwarded upstream. - oauth2_headers: Optional[Dict[str, str]] = None + oauth2_headers: dict[str, str] | None = None if new_mcp_server_request.auth_type in _UPSTREAM_OAUTH_DISCOVERY_AUTH_TYPES and headers.get( MCPRequestHandler.LITELLM_API_KEY_HEADER_NAME_PRIMARY ): @@ -1376,8 +1370,8 @@ if MCP_AVAILABLE: return await session.list_tools() list_tools_response = await client.run_with_session(_list_tools_session_operation) - list_tools_result: List[MCPTool] = list_tools_response.tools - model_dumped_tools: List[dict] = [tool.model_dump() for tool in list_tools_result] + list_tools_result: list[MCPTool] = list_tools_response.tools + model_dumped_tools: list[dict] = [tool.model_dump() for tool in list_tools_result] return { "tools": model_dumped_tools, "error": None, diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 9a1fffd5a67..14673cf12c1 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -13,18 +13,11 @@ import time import traceback import types import uuid +from collections.abc import AsyncIterator, Callable, Mapping from datetime import datetime from typing import ( + TYPE_CHECKING, Any, - AsyncIterator, - Callable, - Dict, - List, - Mapping, - Optional, - Set, - Tuple, - Union, cast, ) @@ -86,11 +79,14 @@ from litellm.types.mcp_server.mcp_server_manager import MCPInfo, MCPServer from litellm.types.utils import CallTypes, StandardLoggingMCPToolCall from litellm.utils import Rules, client, function_setup +if TYPE_CHECKING: + from litellm.proxy._experimental.mcp_server.db import OAuthCredentialPayload + # Short-lived in-memory cache for BYOK credentials. # Keyed by (user_id, server_id); value is (credential_or_None, monotonic_timestamp). # Storing the credential value (not just a bool) means _get_byok_credential and # _check_byok_credential share a single DB round-trip per TTL window. -_byok_cred_cache: Dict[Tuple[str, str], Tuple[Optional[str], float]] = {} +_byok_cred_cache: dict[tuple[str, str], tuple[str | None, float]] = {} _BYOK_CRED_CACHE_TTL = 60 # seconds _BYOK_CRED_CACHE_MAX_SIZE = 4096 # cap to prevent unbounded growth _STATEFUL_SESSION_IDLE_TIMEOUT_SECONDS = 30 * 60 @@ -120,7 +116,7 @@ def _invalidate_byok_cred_cache(user_id: str, server_id: str) -> None: _byok_cred_cache.pop((user_id, server_id), None) -def _write_byok_cred_cache(user_id: str, server_id: str, credential: Optional[str]) -> None: +def _write_byok_cred_cache(user_id: str, server_id: str, credential: str | None) -> None: """Write a credential value to the cache, evicting all entries if at capacity.""" if len(_byok_cred_cache) >= _BYOK_CRED_CACHE_MAX_SIZE: _byok_cred_cache.clear() @@ -150,7 +146,7 @@ try: # Robust auth lookup keyed by session_object. _session_obj_auth_storage: "weakref.WeakKeyDictionary[Any, MCPAuthenticatedUser]" = weakref.WeakKeyDictionary() - active_mcp_session_var: contextvars.ContextVar[Optional[_McpServerSession]] = contextvars.ContextVar( + active_mcp_session_var: contextvars.ContextVar[_McpServerSession | None] = contextvars.ContextVar( "active_mcp_session", default=None ) except ImportError as e: @@ -175,8 +171,8 @@ _INITIALIZATION_LOCK = asyncio.Lock() def _mcp_session_id_from_headers( - raw_headers: Optional[Dict[str, str]], -) -> Optional[str]: + raw_headers: dict[str, str] | None, +) -> str | None: """The ``mcp-session-id`` of a stateful MCP session, read case-insensitively from the request headers. ``None`` for stateless calls (no such header).""" if not raw_headers: @@ -201,10 +197,10 @@ def _jsonrpc_text_has_top_level_method(text: str) -> bool: depth = 0 in_string = False escaped = False - in_object: List[bool] = [] + in_object: list[bool] = [] reading_key = False expect_key = False - key_chars: List[str] = [] + key_chars: list[str] = [] for ch in text: if in_string: if escaped: @@ -241,7 +237,7 @@ def _jsonrpc_text_has_top_level_method(text: str) -> bool: return False -def _mcp_meta_trace_carrier(req_ctx: object) -> Optional[dict[str, str]]: +def _mcp_meta_trace_carrier(req_ctx: object) -> dict[str, str] | None: """The W3C trace context (``traceparent``/``tracestate``) the MCP client propagated in the request's ``params._meta`` (SEP-414), or ``None``. @@ -264,7 +260,7 @@ def _mcp_meta_trace_carrier(req_ctx: object) -> Optional[dict[str, str]]: return carrier or None -def _otel_set_mcp_trace_carrier(carrier: Optional[dict[str, str]]) -> object: +def _otel_set_mcp_trace_carrier(carrier: dict[str, str] | None) -> object: """Stash ``carrier`` for the otel_v2 MCP span and return a reset token, or ``None`` when otel_v2 is unavailable. Lazily imported so opentelemetry stays an optional dependency.""" @@ -462,12 +458,12 @@ if MCP_AVAILABLE: Object returned by the /tools/list REST API route. """ - mcp_info: Optional[MCPInfo] = None + mcp_info: MCPInfo | None = None model_config = ConfigDict(arbitrary_types_allowed=True) - def _normalize_resource_contents(contents: list) -> List[ReadResourceContents]: + def _normalize_resource_contents(contents: list) -> list[ReadResourceContents]: """Normalize ResourceContents to ReadResourceContents, preserving meta (MCP 1.26.0+).""" - normalized: List[ReadResourceContents] = [] + normalized: list[ReadResourceContents] = [] for content in contents: meta = getattr(content, "meta", None) if meta is None and hasattr(content, "model_dump"): @@ -495,15 +491,15 @@ if MCP_AVAILABLE: def _gateway_create_initialization_options( self, - notification_options: Optional[NotificationOptions] = None, - experimental_capabilities: Optional[Dict[str, Dict[str, Any]]] = None, + notification_options: NotificationOptions | None = None, + experimental_capabilities: dict[str, dict[str, Any]] | None = None, ) -> InitializationOptions: opts = Server.create_initialization_options( self, notification_options=notification_options, experimental_capabilities=experimental_capabilities or {}, ) - updates: Dict[str, Any] = {} + updates: dict[str, Any] = {} merged = _mcp_gateway_initialize_instructions.get() if merged is not None: updates["instructions"] = merged @@ -538,21 +534,21 @@ if MCP_AVAILABLE: json_response=False, # enables SSE streaming stateless=False, ) - _stateful_session_auth_contexts: Dict[str, MCPAuthenticatedUser] = {} - _stateful_session_auth_context_last_seen: Dict[str, float] = {} + _stateful_session_auth_contexts: dict[str, MCPAuthenticatedUser] = {} + _stateful_session_auth_context_last_seen: dict[str, float] = {} # Maps session_id -> owner identifier (hashed API key/token) so we can # reject requests that supply a session_id created by a different caller. # Without this, a leaked mcp-session-id could be driven (or terminated) # by any other authenticated proxy user. - _stateful_session_owners: Dict[str, str] = {} + _stateful_session_owners: dict[str, str] = {} # Per-session lock that serializes ``handle_request`` for the same # mcp-session-id. The stored ``MCPAuthenticatedUser`` is mutated in place # by ``_update_auth_context`` each request; without this lock, two # concurrent requests on the same session would clobber each other's # auth headers / mcp_servers / oauth state while in-flight callbacks are # still reading the shared object. - _stateful_session_locks: Dict[str, asyncio.Lock] = {} - _stateful_session_active_request_counts: Dict[str, int] = {} + _stateful_session_locks: dict[str, asyncio.Lock] = {} + _stateful_session_active_request_counts: dict[str, int] = {} def _remove_stateful_session_tracking(session_id: str) -> None: _stateful_session_auth_contexts.pop(session_id, None) @@ -576,10 +572,10 @@ if MCP_AVAILABLE: _session_manager_cm = None _session_manager_stateful_cm = None _sse_session_manager_cm = None - _stateful_auth_context_cleanup_task: Optional[asyncio.Task] = None + _stateful_auth_context_cleanup_task: asyncio.Task | None = None async def _purge_expired_stateful_session_auth_contexts( - now: Optional[float] = None, + now: float | None = None, ) -> None: """Terminate expired stateful sessions and drop their auth contexts.""" now = time.monotonic() if now is None else now @@ -626,7 +622,7 @@ if MCP_AVAILABLE: """ server_instances = getattr(session_manager_stateful, "_server_instances", {}) - def _owned_live_session_ids() -> List[str]: + def _owned_live_session_ids() -> list[str]: return [ session_id for session_id, session_owner in _stateful_session_owners.items() @@ -736,7 +732,7 @@ if MCP_AVAILABLE: ######################################################## @server.list_tools() - async def handle_list_tools() -> "ListToolsResult | List[Tool]": + async def handle_list_tools() -> "ListToolsResult | list[Tool]": """ List all available tools, with each server's listing outcome attached to the result's ``_meta`` (SERVER_OUTCOMES_META_KEY) so a broken upstream is distinguishable from a healthy @@ -816,7 +812,7 @@ if MCP_AVAILABLE: if _session_reset_token is not None: active_mcp_session_var.reset(_session_reset_token) - def _capture_host_progress_callback(host_server) -> Optional[Callable]: + def _capture_host_progress_callback(host_server) -> Callable | None: """Return a progress-forwarding callback bound to the host MCP session. Returns ``None`` when the host did not supply a progress token. @@ -834,7 +830,7 @@ if MCP_AVAILABLE: return None host_session = host_ctx.session - async def forward_progress(progress: float, total: Optional[float]): + async def forward_progress(progress: float, total: float | None): """Forward progress notifications from external MCP to Host""" try: await host_session.send_progress_notification( @@ -853,7 +849,7 @@ if MCP_AVAILABLE: name: str, arguments: dict[str, Any], user_api_key_auth: UserAPIKeyAuth, - ) -> Optional[LiteLLMLoggingObj]: + ) -> LiteLLMLoggingObj | None: """Run the pre-call pipeline (guardrails + logging setup) for a virtual mcp_tool_call so the SSE path spend-logs like the REST path.""" from fastapi import Request @@ -889,15 +885,15 @@ if MCP_AVAILABLE: async def _dispatch_virtual_mcp_tool( name: str, - arguments: Optional[dict[str, Any]], - user_api_key_auth: Optional[UserAPIKeyAuth], - client_ip: Optional[str], - mcp_servers: Optional[list[str]] = None, - mcp_auth_header: Optional[str] = None, - mcp_server_auth_headers: Optional[dict[str, dict[str, str]]] = None, - oauth2_headers: Optional[dict[str, str]] = None, - raw_headers: Optional[dict[str, str]] = None, - ) -> Optional[CallToolResult]: + arguments: dict[str, Any] | None, + user_api_key_auth: UserAPIKeyAuth | None, + client_ip: str | None, + mcp_servers: list[str] | None = None, + mcp_auth_header: str | None = None, + mcp_server_auth_headers: dict[str, dict[str, str]] | None = None, + oauth2_headers: dict[str, str] | None = None, + raw_headers: dict[str, str] | None = None, + ) -> CallToolResult | None: """Handle the mcp_tool_search / mcp_tool_call virtual tools. Returns a CallToolResult when ``name`` is a virtual tool, else ``None`` so @@ -961,7 +957,7 @@ if MCP_AVAILABLE: ) @server.call_tool() - async def mcp_server_tool_call(name: str, arguments: Dict[str, Any] | None) -> CallToolResult: + async def mcp_server_tool_call(name: str, arguments: dict[str, Any] | None) -> CallToolResult: """ Call a specific tool with the provided arguments Args: @@ -1134,7 +1130,7 @@ if MCP_AVAILABLE: active_mcp_session_var.reset(_session_reset_token) @server.list_prompts() - async def list_prompts() -> List[Prompt]: + async def list_prompts() -> list[Prompt]: """ List all available prompts """ @@ -1183,7 +1179,7 @@ if MCP_AVAILABLE: active_mcp_session_var.reset(_session_reset_token) @server.get_prompt() - async def get_prompt(name: str, arguments: Optional[Dict[str, str]]) -> GetPromptResult: + async def get_prompt(name: str, arguments: dict[str, str] | None) -> GetPromptResult: """ Get a specific prompt with the provided arguments @@ -1230,7 +1226,7 @@ if MCP_AVAILABLE: active_mcp_session_var.reset(_session_reset_token) @server.list_resources() - async def list_resources() -> List[Resource]: + async def list_resources() -> list[Resource]: """List all available resources.""" from mcp.server.lowlevel.server import request_ctx @@ -1273,7 +1269,7 @@ if MCP_AVAILABLE: active_mcp_session_var.reset(_session_reset_token) @server.list_resource_templates() - async def list_resource_templates() -> List[ResourceTemplate]: + async def list_resource_templates() -> list[ResourceTemplate]: """List all available resource templates.""" from mcp.server.lowlevel.server import request_ctx @@ -1361,9 +1357,9 @@ if MCP_AVAILABLE: ######################################################## async def _get_allowed_mcp_servers_from_mcp_server_names( - mcp_servers: Optional[List[str]], - allowed_mcp_servers: List[MCPServer], - ) -> List[MCPServer]: + mcp_servers: list[str] | None, + allowed_mcp_servers: list[MCPServer], + ) -> list[MCPServer]: """ Get the filtered MCP servers from the MCP server names. @@ -1418,7 +1414,7 @@ if MCP_AVAILABLE: return allowed_mcp_servers - def _tool_name_matches(tool_name: str, filter_list: List[str]) -> bool: + def _tool_name_matches(tool_name: str, filter_list: list[str]) -> bool: """ Check if a tool name matches any name in the filter list. @@ -1448,9 +1444,9 @@ if MCP_AVAILABLE: return unprefixed_name.lower() in filter_list_lower def filter_tools_by_allowed_tools( - tools: List[MCPTool], + tools: list[MCPTool], mcp_server: MCPServer, - ) -> List[MCPTool]: + ) -> list[MCPTool]: """ Filter tools by allowed/disallowed tools configuration. @@ -1486,9 +1482,9 @@ if MCP_AVAILABLE: return tools_to_return def apply_tool_overrides( - tools: List[MCPTool], + tools: list[MCPTool], mcp_server: MCPServer, - ) -> List[MCPTool]: + ) -> list[MCPTool]: """Apply admin-configured display name/description overrides to tools. Overrides are keyed by the unprefixed tool name, same convention as @@ -1508,7 +1504,7 @@ if MCP_AVAILABLE: tool.description = description_map[lookup_key] return tools - def _get_client_ip_from_context() -> Optional[str]: + def _get_client_ip_from_context() -> str | None: """ Extract client_ip from auth context. Returns None if context not set (caller should handle this as "no IP filtering"). @@ -1522,10 +1518,10 @@ if MCP_AVAILABLE: return None async def _get_allowed_mcp_servers( - user_api_key_auth: Optional[UserAPIKeyAuth], - mcp_servers: Optional[List[str]], - client_ip: Optional[str] = None, - ) -> List[MCPServer]: + user_api_key_auth: UserAPIKeyAuth | None, + mcp_servers: list[str] | None, + client_ip: str | None = None, + ) -> list[MCPServer]: """Return allowed MCP servers for a request after applying filters. Args: @@ -1566,7 +1562,7 @@ if MCP_AVAILABLE: _ip_blocked, client_ip, ) - allowed_mcp_servers: List[MCPServer] = [] + allowed_mcp_servers: list[MCPServer] = [] for allowed_mcp_server_id in allowed_mcp_server_ids: mcp_server = global_mcp_server_manager.get_mcp_server_by_id(allowed_mcp_server_id) if mcp_server is not None: @@ -1584,7 +1580,7 @@ if MCP_AVAILABLE: def _client_has_per_server_auth_header( server: MCPServer, - mcp_server_auth_headers: Optional[Dict[str, Dict[str, str]]], + mcp_server_auth_headers: dict[str, dict[str, str]] | None, ) -> bool: """True if the request carries a per-server ``x-mcp-{alias}-authorization`` header for this server. This is the multi-server binding: it names one @@ -1613,8 +1609,8 @@ if MCP_AVAILABLE: def _client_has_passthrough_authorization( server: MCPServer, - oauth2_headers: Optional[Dict[str, str]], - mcp_server_auth_headers: Optional[Dict[str, Dict[str, str]]], + oauth2_headers: dict[str, str] | None, + mcp_server_auth_headers: dict[str, dict[str, str]] | None, ) -> bool: """True if the incoming request already carries an ``Authorization`` header the gateway will forward to this pass-through server. @@ -1632,9 +1628,9 @@ if MCP_AVAILABLE: async def _get_user_oauth_extra_headers_from_db( server: MCPServer, - user_api_key_auth: Optional[UserAPIKeyAuth], - prefetched_creds: Optional[Dict[str, Dict[str, Any]]] = None, - ) -> Optional[Dict[str, str]]: + user_api_key_auth: UserAPIKeyAuth | None, + prefetched_creds: dict[str, dict[str, Any]] | None = None, + ) -> dict[str, str] | None: """Stored OAuth2 token for (user, server) as an ``Authorization: Bearer`` header, or None. Thin wrapper over ``resolve_user_oauth_access_token`` (Redis cache, else DB + refresh); @@ -1652,8 +1648,8 @@ if MCP_AVAILABLE: return {"Authorization": f"Bearer {token}"} if token else None async def _prefetch_oauth_creds_for_user( - user_api_key_auth: Optional[UserAPIKeyAuth], - ) -> Dict[str, Dict[str, Any]]: + user_api_key_auth: UserAPIKeyAuth | None, + ) -> dict[str, "OAuthCredentialPayload"]: """Fetch all OAuth2 credentials for the user in one DB query. Returns a dict keyed by server_id to avoid N+1 queries in asyncio.gather loops. @@ -1678,13 +1674,13 @@ if MCP_AVAILABLE: def _prepare_mcp_server_headers( server: MCPServer, - mcp_server_auth_headers: Optional[Dict[str, Dict[str, str]]], - mcp_auth_header: Optional[str], - oauth2_headers: Optional[Dict[str, str]], - raw_headers: Optional[Dict[str, str]], - user_api_key_auth: Optional[UserAPIKeyAuth] = None, - scope_servers: Optional[list[MCPServer]] = None, - ) -> Tuple[Optional[Union[Dict[str, str], str]], Optional[Dict[str, str]]]: + mcp_server_auth_headers: dict[str, dict[str, str]] | None, + mcp_auth_header: str | None, + oauth2_headers: dict[str, str] | None, + raw_headers: dict[str, str] | None, + user_api_key_auth: UserAPIKeyAuth | None = None, + scope_servers: list[MCPServer] | None = None, + ) -> tuple[dict[str, str] | str | None, dict[str, str] | None]: """Build auth and extra headers for a server. ``scope_servers`` is the full server list a fan-out handler iterates. Passing it lets the @@ -1693,7 +1689,7 @@ if MCP_AVAILABLE: explicitly-addressed operations leave it None. Per-server ``x-mcp-{alias}-authorization`` headers are unaffected — they bind one token to one server and are the multi-server shape. """ - server_auth_header: Optional[Union[Dict[str, str], str]] = None + server_auth_header: dict[str, str] | str | None = None if mcp_server_auth_headers: from litellm.proxy._experimental.mcp_server.utils import ( lookup_mcp_server_auth_in_headers, @@ -1705,7 +1701,7 @@ if MCP_AVAILABLE: server_name=server.server_name, ) - extra_headers: Optional[Dict[str, str]] = None + extra_headers: dict[str, str] | None = None is_client_forwarded_mode = server.is_true_passthrough or server.is_oauth_delegate # In a multi-server listing scope the request-wide Authorization can only carry one token, # so it is withheld from a client-forwarded server when another server in scope also consumes @@ -1781,13 +1777,13 @@ if MCP_AVAILABLE: return server_auth_header, extra_headers def _merge_gateway_initialize_instructions( - allowed_mcp_servers: List[MCPServer], - ) -> Optional[str]: + allowed_mcp_servers: list[MCPServer], + ) -> str | None: """YAML/DB override, else upstream text (prefetch on init, or list_tools / health_check / call_tool cache).""" if not allowed_mcp_servers: return None - texts: List[Tuple[str, str]] = [] + texts: list[tuple[str, str]] = [] for server in allowed_mcp_servers: label = server.alias or server.server_name or server.name or server.server_id or "mcp" if server.instructions and server.instructions.strip(): @@ -1807,9 +1803,9 @@ if MCP_AVAILABLE: @contextlib.asynccontextmanager async def _gateway_initialize_instructions_request_scope( - user_api_key_auth: Optional[UserAPIKeyAuth], - mcp_servers: Optional[List[str]], - client_ip: Optional[str], + user_api_key_auth: UserAPIKeyAuth | None, + mcp_servers: list[str] | None, + client_ip: str | None, scoped_server_endpoint: bool = False, ) -> AsyncIterator[None]: allowed = await _get_allowed_mcp_servers( @@ -1852,17 +1848,17 @@ if MCP_AVAILABLE: return get_server_prefix(server) or "unknown" async def _get_tools_from_mcp_servers( - user_api_key_auth: Optional[UserAPIKeyAuth], - mcp_auth_header: Optional[str], - mcp_servers: Optional[List[str]], - mcp_server_auth_headers: Optional[Dict[str, Dict[str, str]]] = None, - oauth2_headers: Optional[Dict[str, str]] = None, - raw_headers: Optional[Dict[str, str]] = None, + user_api_key_auth: UserAPIKeyAuth | None, + mcp_auth_header: str | None, + mcp_servers: list[str] | None, + mcp_server_auth_headers: dict[str, dict[str, str]] | None = None, + oauth2_headers: dict[str, str] | None = None, + raw_headers: dict[str, str] | None = None, log_list_tools_to_spendlogs: bool = False, - list_tools_log_source: Optional[str] = None, - litellm_trace_id: Optional[str] = None, - request_tags: Optional[list[str]] = None, - client_ip: Optional[str] = None, + list_tools_log_source: str | None = None, + litellm_trace_id: str | None = None, + request_tags: list[str] | None = None, + client_ip: str | None = None, ) -> AggregateToolListing: """ Helper method to fetch tools from MCP servers based on server filtering criteria. @@ -1882,8 +1878,8 @@ if MCP_AVAILABLE: return AggregateToolListing(tools=[], outcomes={}) list_tools_start_time = datetime.now() - litellm_logging_obj: Optional[LiteLLMLoggingObj] = None - list_tools_request_data: Dict[str, Any] = {} + litellm_logging_obj: LiteLLMLoggingObj | None = None + list_tools_request_data: dict[str, Any] = {} if log_list_tools_to_spendlogs: # This is intentionally minimal: only async_success_handler / post_call_failure_hook @@ -1891,7 +1887,7 @@ if MCP_AVAILABLE: list_tools_call_id = str(uuid.uuid4()) # Derive trace_id from raw_headers when not explicitly passed (same as A2A / MCP call_tool) effective_litellm_trace_id = litellm_trace_id or get_chain_id_from_headers(raw_headers) - spend_logs_metadata: Dict[str, Any] = { + spend_logs_metadata: dict[str, Any] = { "mcp_operation": "list_tools", } if isinstance(list_tools_log_source, str): @@ -1964,7 +1960,7 @@ if MCP_AVAILABLE: async def _fetch_and_filter_server_tools( server: MCPServer, - ) -> "tuple[List[MCPTool], ServerOutcome]": + ) -> "tuple[list[MCPTool], ServerOutcome]": """Fetch and filter tools from a single server, classifying any failure into that server's outcome so the aggregate can keep serving the healthy subset without a broken server masquerading as an empty one.""" @@ -2058,8 +2054,8 @@ if MCP_AVAILABLE: results = await asyncio.gather(*tasks) # Flatten results into single list - all_tools: List[MCPTool] = [tool for tools, _ in results for tool in tools] - server_outcomes: Dict[str, ServerOutcome] = { + all_tools: list[MCPTool] = [tool for tools, _ in results for tool in tools] + server_outcomes: dict[str, ServerOutcome] = { _aggregate_server_key(server): outcome for server, (_, outcome) in zip(allowed_mcp_servers, results) if server is not None @@ -2067,7 +2063,7 @@ if MCP_AVAILABLE: # If logging is enabled, enrich spend_logs_metadata with counts if litellm_logging_obj: - per_server_tool_counts: Dict[str, int] = { + per_server_tool_counts: dict[str, int] = { _aggregate_server_key(server): len(server_tools) for server, (server_tools, _) in zip(allowed_mcp_servers, results) if server is not None @@ -2126,13 +2122,13 @@ if MCP_AVAILABLE: raise async def _get_prompts_from_mcp_servers( - user_api_key_auth: Optional[UserAPIKeyAuth], - mcp_auth_header: Optional[str], - mcp_servers: Optional[List[str]], - mcp_server_auth_headers: Optional[Dict[str, Dict[str, str]]] = None, - oauth2_headers: Optional[Dict[str, str]] = None, - raw_headers: Optional[Dict[str, str]] = None, - ) -> List[Prompt]: + user_api_key_auth: UserAPIKeyAuth | None, + mcp_auth_header: str | None, + mcp_servers: list[str] | None, + mcp_server_auth_headers: dict[str, dict[str, str]] | None = None, + oauth2_headers: dict[str, str] | None = None, + raw_headers: dict[str, str] | None = None, + ) -> list[Prompt]: """ Helper method to fetch prompt from MCP servers based on server filtering criteria. @@ -2191,13 +2187,13 @@ if MCP_AVAILABLE: return all_prompts async def _get_resources_from_mcp_servers( - user_api_key_auth: Optional[UserAPIKeyAuth], - mcp_auth_header: Optional[str], - mcp_servers: Optional[List[str]], - mcp_server_auth_headers: Optional[Dict[str, Dict[str, str]]] = None, - oauth2_headers: Optional[Dict[str, str]] = None, - raw_headers: Optional[Dict[str, str]] = None, - ) -> List[Resource]: + user_api_key_auth: UserAPIKeyAuth | None, + mcp_auth_header: str | None, + mcp_servers: list[str] | None, + mcp_server_auth_headers: dict[str, dict[str, str]] | None = None, + oauth2_headers: dict[str, str] | None = None, + raw_headers: dict[str, str] | None = None, + ) -> list[Resource]: """Fetch resources from allowed MCP servers.""" if not MCP_AVAILABLE: @@ -2208,7 +2204,7 @@ if MCP_AVAILABLE: mcp_servers=mcp_servers, ) - all_resources: List[Resource] = [] + all_resources: list[Resource] = [] for server in allowed_mcp_servers: if server is None: continue @@ -2242,13 +2238,13 @@ if MCP_AVAILABLE: return all_resources async def _get_resource_templates_from_mcp_servers( - user_api_key_auth: Optional[UserAPIKeyAuth], - mcp_auth_header: Optional[str], - mcp_servers: Optional[List[str]], - mcp_server_auth_headers: Optional[Dict[str, Dict[str, str]]] = None, - oauth2_headers: Optional[Dict[str, str]] = None, - raw_headers: Optional[Dict[str, str]] = None, - ) -> List[ResourceTemplate]: + user_api_key_auth: UserAPIKeyAuth | None, + mcp_auth_header: str | None, + mcp_servers: list[str] | None, + mcp_server_auth_headers: dict[str, dict[str, str]] | None = None, + oauth2_headers: dict[str, str] | None = None, + raw_headers: dict[str, str] | None = None, + ) -> list[ResourceTemplate]: """Fetch resource templates from allowed MCP servers.""" if not MCP_AVAILABLE: @@ -2259,7 +2255,7 @@ if MCP_AVAILABLE: mcp_servers=mcp_servers, ) - all_resource_templates: List[ResourceTemplate] = [] + all_resource_templates: list[ResourceTemplate] = [] for server in allowed_mcp_servers: if server is None: continue @@ -2303,10 +2299,10 @@ if MCP_AVAILABLE: return all_resource_templates async def filter_tools_by_key_team_permissions( - tools: List[MCPTool], + tools: list[MCPTool], server_id: str, - user_api_key_auth: Optional[UserAPIKeyAuth], - ) -> List[MCPTool]: + user_api_key_auth: UserAPIKeyAuth | None, + ) -> list[MCPTool]: """ Filter tools based on key/team mcp_tool_permissions. @@ -2329,15 +2325,15 @@ if MCP_AVAILABLE: return [t for t in tools if strip_known_server_prefix(t.name, server) in allowed_tool_names] async def _list_mcp_tools( - user_api_key_auth: Optional[UserAPIKeyAuth] = None, - mcp_auth_header: Optional[str] = None, - mcp_servers: Optional[List[str]] = None, - mcp_server_auth_headers: Optional[Dict[str, Dict[str, str]]] = None, - oauth2_headers: Optional[Dict[str, str]] = None, - raw_headers: Optional[Dict[str, str]] = None, + user_api_key_auth: UserAPIKeyAuth | None = None, + mcp_auth_header: str | None = None, + mcp_servers: list[str] | None = None, + mcp_server_auth_headers: dict[str, dict[str, str]] | None = None, + oauth2_headers: dict[str, str] | None = None, + raw_headers: dict[str, str] | None = None, log_list_tools_to_spendlogs: bool = False, - list_tools_log_source: Optional[str] = None, - client_ip: Optional[str] = None, + list_tools_log_source: str | None = None, + client_ip: str | None = None, ) -> AggregateToolListing: """ List all available MCP tools. @@ -2376,13 +2372,13 @@ if MCP_AVAILABLE: return AggregateToolListing(tools=[], outcomes={}) async def _list_mcp_prompts( - user_api_key_auth: Optional[UserAPIKeyAuth] = None, - mcp_auth_header: Optional[str] = None, - mcp_servers: Optional[List[str]] = None, - mcp_server_auth_headers: Optional[Dict[str, Dict[str, str]]] = None, - oauth2_headers: Optional[Dict[str, str]] = None, - raw_headers: Optional[Dict[str, str]] = None, - ) -> List[Prompt]: + user_api_key_auth: UserAPIKeyAuth | None = None, + mcp_auth_header: str | None = None, + mcp_servers: list[str] | None = None, + mcp_server_auth_headers: dict[str, dict[str, str]] | None = None, + oauth2_headers: dict[str, str] | None = None, + raw_headers: dict[str, str] | None = None, + ) -> list[Prompt]: """ List all available MCP prompts. @@ -2416,19 +2412,19 @@ if MCP_AVAILABLE: return managed_prompts async def _list_mcp_resources( - user_api_key_auth: Optional[UserAPIKeyAuth] = None, - mcp_auth_header: Optional[str] = None, - mcp_servers: Optional[List[str]] = None, - mcp_server_auth_headers: Optional[Dict[str, Dict[str, str]]] = None, - oauth2_headers: Optional[Dict[str, str]] = None, - raw_headers: Optional[Dict[str, str]] = None, - ) -> List[Resource]: + user_api_key_auth: UserAPIKeyAuth | None = None, + mcp_auth_header: str | None = None, + mcp_servers: list[str] | None = None, + mcp_server_auth_headers: dict[str, dict[str, str]] | None = None, + oauth2_headers: dict[str, str] | None = None, + raw_headers: dict[str, str] | None = None, + ) -> list[Resource]: """List all available MCP resources.""" if not MCP_AVAILABLE: return [] - managed_resources: List[Resource] = [] + managed_resources: list[Resource] = [] try: managed_resources = await _get_resources_from_mcp_servers( user_api_key_auth=user_api_key_auth, @@ -2445,19 +2441,19 @@ if MCP_AVAILABLE: return managed_resources async def _list_mcp_resource_templates( - user_api_key_auth: Optional[UserAPIKeyAuth] = None, - mcp_auth_header: Optional[str] = None, - mcp_servers: Optional[List[str]] = None, - mcp_server_auth_headers: Optional[Dict[str, Dict[str, str]]] = None, - oauth2_headers: Optional[Dict[str, str]] = None, - raw_headers: Optional[Dict[str, str]] = None, - ) -> List[ResourceTemplate]: + user_api_key_auth: UserAPIKeyAuth | None = None, + mcp_auth_header: str | None = None, + mcp_servers: list[str] | None = None, + mcp_server_auth_headers: dict[str, dict[str, str]] | None = None, + oauth2_headers: dict[str, str] | None = None, + raw_headers: dict[str, str] | None = None, + ) -> list[ResourceTemplate]: """List all available MCP resource templates.""" if not MCP_AVAILABLE: return [] - managed_resource_templates: List[ResourceTemplate] = [] + managed_resource_templates: list[ResourceTemplate] = [] try: managed_resource_templates = await _get_resource_templates_from_mcp_servers( user_api_key_auth=user_api_key_auth, @@ -2481,7 +2477,7 @@ if MCP_AVAILABLE: def _resolve_display_name_to_original( name: str, - allowed_mcp_servers: List[MCPServer], + allowed_mcp_servers: list[MCPServer], ) -> str: """Translate a display-name override back to the original prefixed tool name. @@ -2499,8 +2495,8 @@ if MCP_AVAILABLE: async def _get_byok_credential( mcp_server: MCPServer, - user_api_key_auth: Optional[UserAPIKeyAuth], - ) -> Optional[str]: + user_api_key_auth: UserAPIKeyAuth | None, + ) -> str | None: """Retrieve the stored BYOK credential for a user+server pair. Uses the shared _byok_cred_cache to avoid a DB round-trip on every @@ -2534,7 +2530,7 @@ if MCP_AVAILABLE: async def _check_byok_credential( mcp_server: MCPServer, - user_api_key_auth: Optional[UserAPIKeyAuth], + user_api_key_auth: UserAPIKeyAuth | None, ) -> None: """ If the MCP server is BYOK-enabled, verify that the requesting user has a @@ -2622,15 +2618,15 @@ if MCP_AVAILABLE: async def execute_mcp_tool( name: str, - arguments: Dict[str, Any], - allowed_mcp_servers: List[MCPServer], + arguments: dict[str, Any], + allowed_mcp_servers: list[MCPServer], start_time: datetime, - user_api_key_auth: Optional[UserAPIKeyAuth] = None, - mcp_auth_header: Optional[str] = None, - mcp_server_auth_headers: Optional[Dict[str, Dict[str, str]]] = None, - oauth2_headers: Optional[Dict[str, str]] = None, - raw_headers: Optional[Dict[str, str]] = None, - host_progress_callback: Optional[Callable] = None, + user_api_key_auth: UserAPIKeyAuth | None = None, + mcp_auth_header: str | None = None, + mcp_server_auth_headers: dict[str, dict[str, str]] | None = None, + oauth2_headers: dict[str, str] | None = None, + raw_headers: dict[str, str] | None = None, + host_progress_callback: Callable | None = None, **kwargs: Any, ) -> CallToolResult: """ @@ -2654,8 +2650,8 @@ if MCP_AVAILABLE: CallToolResult: Tool execution result """ # Track resolved MCP server for both permission checks and dispatch - mcp_server: Optional[MCPServer] = None - requested_server_id: Optional[str] = kwargs.get("requested_server_id") + mcp_server: MCPServer | None = None + requested_server_id: str | None = kwargs.get("requested_server_id") # If the client called with a display-name override (e.g. "Get Pet"), # translate it back to the original prefixed name before any routing. @@ -2664,7 +2660,7 @@ if MCP_AVAILABLE: # Remove prefix from tool name for logging and processing original_tool_name, server_name = split_server_prefix_from_name(name) - requested_server: Optional[MCPServer] = None + requested_server: MCPServer | None = None if requested_server_id: requested_server = next( (s for s in allowed_mcp_servers if s.server_id == requested_server_id), @@ -2673,7 +2669,7 @@ if MCP_AVAILABLE: name_is_prefixed = False if requested_server is not None and MCP_TOOL_PREFIX_SEPARATOR in name: - all_registry_prefixes: Set[str] = set() + all_registry_prefixes: set[str] = set() for registry_server in global_mcp_server_manager.get_registry().values(): for known_prefix in iter_known_server_prefixes(registry_server): all_registry_prefixes.add(normalize_server_name(known_prefix)) @@ -2736,7 +2732,7 @@ if MCP_AVAILABLE: server_name=server_name, session_id=_mcp_session_id_from_headers(raw_headers), ) - litellm_logging_obj: Optional[LiteLLMLoggingObj] = kwargs.get("litellm_logging_obj", None) + litellm_logging_obj: LiteLLMLoggingObj | None = kwargs.get("litellm_logging_obj", None) if litellm_logging_obj: litellm_logging_obj.model_call_details["mcp_tool_call_metadata"] = standard_logging_mcp_tool_call litellm_logging_obj.model = f"MCP: {name}" @@ -2829,7 +2825,7 @@ if MCP_AVAILABLE: # because the tool function has headers baked into its closure. # Pre-format the full Authorization header value using the server's # configured auth_type so the generator doesn't need to know the prefix. - auth_header_value: Optional[str] = None + auth_header_value: str | None = None if mcp_auth_header: server_auth_type = getattr(mcp_server, "auth_type", None) if mcp_server else None if server_auth_type == MCPAuth.api_key: @@ -2845,7 +2841,7 @@ if MCP_AVAILABLE: # OpenAPI/local path agrees with the managed paths: M2M and the resolver-owned modes # (token_exchange's raw subject token, authorization_code's stored token) must never # have the caller's Authorization forwarded verbatim upstream. - forwarded_headers: Optional[Dict[str, str]] = None + forwarded_headers: dict[str, str] | None = None if mcp_server and mcp_server.extra_headers and raw_headers: normalized_raw = {str(k).lower(): v for k, v in raw_headers.items() if isinstance(k, str)} skip_caller_authorization = _should_strip_caller_authorization( @@ -2931,8 +2927,8 @@ if MCP_AVAILABLE: result: Any, start_time: datetime, end_time: datetime, - user_api_key_auth: Optional[UserAPIKeyAuth] = None, - request_data: Optional[Mapping[str, object]] = None, + user_api_key_auth: UserAPIKeyAuth | None = None, + request_data: Mapping[str, object] | None = None, ) -> None: """Fire post-call logging for an executed MCP tool call. @@ -2987,20 +2983,20 @@ if MCP_AVAILABLE: @client async def call_mcp_tool( name: str, - arguments: Optional[Dict[str, Any]] = None, - user_api_key_auth: Optional[UserAPIKeyAuth] = None, - mcp_auth_header: Optional[str] = None, - mcp_servers: Optional[List[str]] = None, - mcp_server_auth_headers: Optional[Dict[str, Dict[str, str]]] = None, - oauth2_headers: Optional[Dict[str, str]] = None, - raw_headers: Optional[Dict[str, str]] = None, + arguments: dict[str, Any] | None = None, + user_api_key_auth: UserAPIKeyAuth | None = None, + mcp_auth_header: str | None = None, + mcp_servers: list[str] | None = None, + mcp_server_auth_headers: dict[str, dict[str, str]] | None = None, + oauth2_headers: dict[str, str] | None = None, + raw_headers: dict[str, str] | None = None, **kwargs: Any, ) -> CallToolResult: """ Call a specific tool with the provided arguments (handles prefixed tool names). """ start_time = datetime.now() - litellm_logging_obj: Optional[LiteLLMLoggingObj] = kwargs.get("litellm_logging_obj", None) + litellm_logging_obj: LiteLLMLoggingObj | None = kwargs.get("litellm_logging_obj", None) try: if arguments is None: @@ -3011,7 +3007,7 @@ if MCP_AVAILABLE: user_api_key_auth=user_api_key_auth, ) - allowed_mcp_servers: List[MCPServer] = [] + allowed_mcp_servers: list[MCPServer] = [] for allowed_mcp_server_id in allowed_mcp_server_ids: allowed_server = global_mcp_server_manager.get_mcp_server_by_id(allowed_mcp_server_id) if allowed_server is not None: @@ -3078,13 +3074,13 @@ if MCP_AVAILABLE: async def mcp_get_prompt( name: str, - arguments: Optional[Dict[str, Any]] = None, - user_api_key_auth: Optional[UserAPIKeyAuth] = None, - mcp_auth_header: Optional[str] = None, - mcp_servers: Optional[List[str]] = None, - mcp_server_auth_headers: Optional[Dict[str, Dict[str, str]]] = None, - oauth2_headers: Optional[Dict[str, str]] = None, - raw_headers: Optional[Dict[str, str]] = None, + arguments: dict[str, Any] | None = None, + user_api_key_auth: UserAPIKeyAuth | None = None, + mcp_auth_header: str | None = None, + mcp_servers: list[str] | None = None, + mcp_server_auth_headers: dict[str, dict[str, str]] | None = None, + oauth2_headers: dict[str, str] | None = None, + raw_headers: dict[str, str] | None = None, ) -> GetPromptResult: """ Fetch a specific MCP prompt, handling both prefixed and unprefixed names. @@ -3130,12 +3126,12 @@ if MCP_AVAILABLE: async def mcp_read_resource( url: AnyUrl, - user_api_key_auth: Optional[UserAPIKeyAuth] = None, - mcp_auth_header: Optional[str] = None, - mcp_servers: Optional[List[str]] = None, - mcp_server_auth_headers: Optional[Dict[str, Dict[str, str]]] = None, - oauth2_headers: Optional[Dict[str, str]] = None, - raw_headers: Optional[Dict[str, str]] = None, + user_api_key_auth: UserAPIKeyAuth | None = None, + mcp_auth_header: str | None = None, + mcp_servers: list[str] | None = None, + mcp_server_auth_headers: dict[str, dict[str, str]] | None = None, + oauth2_headers: dict[str, str] | None = None, + raw_headers: dict[str, str] | None = None, ) -> ReadResourceResult: """Read resource contents from upstream MCP servers.""" @@ -3179,9 +3175,9 @@ if MCP_AVAILABLE: def _get_standard_logging_mcp_tool_call( name: str, - arguments: Dict[str, Any], - server_name: Optional[str], - session_id: Optional[str] = None, + arguments: dict[str, Any], + server_name: str | None, + session_id: str | None = None, ) -> StandardLoggingMCPToolCall: mcp_server = global_mcp_server_manager._get_mcp_server_from_tool_name(name) namespaced_tool_name = f"{server_name}/{name}" if server_name else name @@ -3208,14 +3204,14 @@ if MCP_AVAILABLE: async def _handle_managed_mcp_tool( server_name: str, name: str, - arguments: Dict[str, Any], - user_api_key_auth: Optional[UserAPIKeyAuth] = None, - mcp_auth_header: Optional[str] = None, - mcp_server_auth_headers: Optional[Dict[str, Dict[str, str]]] = None, - oauth2_headers: Optional[Dict[str, str]] = None, - raw_headers: Optional[Dict[str, str]] = None, - litellm_logging_obj: Optional[Any] = None, - host_progress_callback: Optional[Callable] = None, + arguments: dict[str, Any], + user_api_key_auth: UserAPIKeyAuth | None = None, + mcp_auth_header: str | None = None, + mcp_server_auth_headers: dict[str, dict[str, str]] | None = None, + oauth2_headers: dict[str, str] | None = None, + raw_headers: dict[str, str] | None = None, + litellm_logging_obj: Any | None = None, + host_progress_callback: Callable | None = None, ) -> CallToolResult: """Handle tool execution for managed server tools""" # Import here to avoid circular import @@ -3237,8 +3233,8 @@ if MCP_AVAILABLE: return call_tool_result async def _handle_local_mcp_tool( - name: str, arguments: Dict[str, Any] - ) -> List[Union[TextContent, ImageContent, EmbeddedResource]]: + name: str, arguments: dict[str, Any] + ) -> list[TextContent | ImageContent | EmbeddedResource]: """ Handle tool execution for local registry tools Note: Local tools don't use prefixes, so we use the original name @@ -3260,13 +3256,13 @@ if MCP_AVAILABLE: verbose_logger.exception(f"Error executing local tool {name}: {str(e)}") return [TextContent(text=f"Error: {str(e)}", type="text")] - def _get_mcp_servers_in_path(path: str) -> Optional[List[str]]: + def _get_mcp_servers_in_path(path: str) -> list[str] | None: """ Get the MCP servers from the path """ import re - mcp_servers_from_path: Optional[List[str]] = None + mcp_servers_from_path: list[str] | None = None segments = [s for s in path.split("/") if s] if len(segments) >= 2 and segments[1] == "mcp" and segments[0] != "mcp": return [segments[0]] @@ -3338,7 +3334,7 @@ if MCP_AVAILABLE: raw_headers, ) - def _get_session_id_from_scope(scope: Scope) -> Optional[str]: + def _get_session_id_from_scope(scope: Scope) -> str | None: """ Extract mcp-session-id from ASGI scope headers. Returns None if not present. @@ -3350,9 +3346,9 @@ if MCP_AVAILABLE: return None def _owner_fingerprint_for( - user_api_key_auth: Optional[UserAPIKeyAuth], - oauth2_headers: Optional[Dict[str, str]] = None, - client_ip: Optional[str] = None, + user_api_key_auth: UserAPIKeyAuth | None, + oauth2_headers: dict[str, str] | None = None, + client_ip: str | None = None, ) -> str: """ Stable, non-reversible identifier for the caller used to bind an @@ -3377,7 +3373,7 @@ if MCP_AVAILABLE: is best-effort in that mode. """ - def _bytes_for_hash(value: Any) -> Optional[bytes]: + def _bytes_for_hash(value: Any) -> bytes | None: """Only hash str/bytes secrets; skip mocks and other unexpected types.""" if value is None: return None @@ -3420,7 +3416,7 @@ if MCP_AVAILABLE: async def _read_request_body_for_routing( receive: Receive, - ) -> Tuple[List[Message], bytes]: + ) -> tuple[list[Message], bytes]: """ Read just enough of the request body to decide whether this is a JSON-RPC ``initialize`` call. Returns the consumed ASGI messages so @@ -3434,8 +3430,8 @@ if MCP_AVAILABLE: force the proxy to buffer an arbitrarily large payload just to make a routing decision. """ - consumed_messages: List[Message] = [] - body_chunks: List[bytes] = [] + consumed_messages: list[Message] = [] + body_chunks: list[bytes] = [] peeked_bytes = 0 while True: @@ -3490,14 +3486,14 @@ if MCP_AVAILABLE: _mcp_session_header = b"mcp-session-id" _headers = scope.get("headers", []) - def _normalize_header_name(header_name: Any) -> Optional[bytes]: + def _normalize_header_name(header_name: Any) -> bytes | None: if isinstance(header_name, bytes): return header_name.lower() if isinstance(header_name, str): return header_name.lower().encode("utf-8", errors="replace") return None - _session_id: Optional[str] = None + _session_id: str | None = None for header_name, header_value in _headers: if _normalize_header_name(header_name) == _mcp_session_header: if isinstance(header_value, bytes): @@ -3641,12 +3637,12 @@ if MCP_AVAILABLE: async def _raise_preemptive_401_for_unauthenticated_servers( scope: Scope, - mcp_servers: Optional[List[str]], - oauth2_headers: Optional[Dict[str, str]], - mcp_server_auth_headers: Optional[Dict[str, Dict[str, str]]], - user_api_key_auth: Optional[UserAPIKeyAuth], - client_ip: Optional[str], - allowed_server_ids: Optional[Set[str]] = None, + mcp_servers: list[str] | None, + oauth2_headers: dict[str, str] | None, + mcp_server_auth_headers: dict[str, dict[str, str]] | None, + user_api_key_auth: UserAPIKeyAuth | None, + client_ip: str | None, + allowed_server_ids: set[str] | None = None, ) -> None: """Fail fast with HTTP 401 for MCP servers that need user auth but didn't receive it on this request. Covers both gateway-managed OAuth2 @@ -3825,7 +3821,7 @@ if MCP_AVAILABLE: headers={"www-authenticate": upstream_www_authenticate}, ) - def _get_authorization_header_from_scope(scope: Scope) -> Optional[str]: + def _get_authorization_header_from_scope(scope: Scope) -> str | None: """First ``Authorization`` header value in the ASGI scope, or None.""" for key, value in scope.get("headers", []): if key.lower() == b"authorization": @@ -3835,7 +3831,7 @@ if MCP_AVAILABLE: def _scope_has_authorization_header(scope: Scope) -> bool: return _get_authorization_header_from_scope(scope) is not None - def _get_forwarded_auth_from_scope(scope: Scope) -> Optional[str]: + def _get_forwarded_auth_from_scope(scope: Scope) -> str | None: """Return the upstream-bound ``Authorization`` header value, or None. Only returns the ``Authorization`` header when ``x-litellm-api-key`` is @@ -3869,7 +3865,7 @@ if MCP_AVAILABLE: url: str, auth_header: str, timeout: float = 5.0, - ) -> tuple[int, Optional[str]]: + ) -> tuple[int, str | None]: """JSON-RPC initialize-probe the upstream URL to check whether the token is accepted. Uses POST so StreamableHTTP MCP servers run the same auth path as a @@ -3921,9 +3917,9 @@ if MCP_AVAILABLE: async def _check_passthrough_upstream_auth( scope: Scope, - user_api_key_auth: Optional[UserAPIKeyAuth], - mcp_servers: Optional[List[str]], - client_ip: Optional[str], + user_api_key_auth: UserAPIKeyAuth | None, + mcp_servers: list[str] | None, + client_ip: str | None, ) -> None: """Probe pass-through and delegate-auth upstream servers in parallel before the MCP session starts. @@ -3978,7 +3974,7 @@ if MCP_AVAILABLE: mcp_servers=mcp_servers, client_ip=client_ip, ) - passthrough_targets: Tuple[Tuple[MCPServer, str, str], ...] = ( + passthrough_targets: tuple[tuple[MCPServer, str, str], ...] = ( tuple( (srv, forwarded_auth, srv.name) for srv in allowed_servers @@ -3997,7 +3993,7 @@ if MCP_AVAILABLE: ) # Probe the admission-resolved delegate server only when the caller is actually # authorized for it (present in the IP-filtered allowed set), keyed by server_id. - delegate_targets: Tuple[Tuple[MCPServer, str, str], ...] = ( + delegate_targets: tuple[tuple[MCPServer, str, str], ...] = ( tuple( (srv, delegate_auth, requested_single_target) for srv in allowed_servers @@ -4065,7 +4061,7 @@ if MCP_AVAILABLE: # Apply toolset scope if set server-side via ContextVar (set by # /toolset/{name}/mcp and /{name}/mcp route handlers in proxy_server.py). active_toolset_id = _mcp_active_toolset_id.get() - toolset_allowed_server_ids: Optional[Set[str]] = None + toolset_allowed_server_ids: set[str] | None = None if active_toolset_id and user_api_key_auth is not None: user_api_key_auth = await _apply_toolset_scope(user_api_key_auth, active_toolset_id) op = user_api_key_auth.object_permission @@ -4116,7 +4112,7 @@ if MCP_AVAILABLE: # - No session ID + other → stateless (curl, Inspector, Notion) session_id = _get_session_id_from_scope(scope) is_initialize = False - consumed_messages: List[Message] = [] + consumed_messages: list[Message] = [] # Owner-binding: a live stateful session may only be driven by the # caller that created it. Reject mismatches with 403 so a leaked @@ -4250,11 +4246,11 @@ if MCP_AVAILABLE: "top-level key scan, skipping session lock to avoid deadlock" ) - session_lock: Optional[asyncio.Lock] = None + session_lock: asyncio.Lock | None = None if use_stateful and session_id and request_method in ("POST", "DELETE") and not is_jsonrpc_response: session_lock = _stateful_session_locks.setdefault(session_id, asyncio.Lock()) - active_request_session_ids: List[str] = [] + active_request_session_ids: list[str] = [] def _increment_active_request_session(session_id_to_track: str) -> None: if session_id_to_track in active_request_session_ids: @@ -4387,7 +4383,7 @@ if MCP_AVAILABLE: # downstream probe list matches the fully-authorized server set # (mirrors the streamable HTTP handler). active_toolset_id = _mcp_active_toolset_id.get() - toolset_allowed_server_ids: Optional[Set[str]] = None + toolset_allowed_server_ids: set[str] | None = None if active_toolset_id and user_api_key_auth is not None: user_api_key_auth = await _apply_toolset_scope(user_api_key_auth, active_toolset_id) op = user_api_key_auth.object_permission @@ -4483,7 +4479,7 @@ if MCP_AVAILABLE: "/enabled", description="Returns if the MCP server is enabled", ) - def get_mcp_server_enabled() -> Dict[str, bool]: + def get_mcp_server_enabled() -> dict[str, bool]: """ Returns if the MCP server is enabled """ @@ -4502,13 +4498,13 @@ if MCP_AVAILABLE: def _update_auth_context( auth_user: MCPAuthenticatedUser, - user_api_key_auth: Optional[UserAPIKeyAuth], - mcp_auth_header: Optional[str] = None, - mcp_servers: Optional[List[str]] = None, - mcp_server_auth_headers: Optional[Dict[str, Dict[str, str]]] = None, - oauth2_headers: Optional[Dict[str, str]] = None, - raw_headers: Optional[Dict[str, str]] = None, - client_ip: Optional[str] = None, + user_api_key_auth: UserAPIKeyAuth | None, + mcp_auth_header: str | None = None, + mcp_servers: list[str] | None = None, + mcp_server_auth_headers: dict[str, dict[str, str]] | None = None, + oauth2_headers: dict[str, str] | None = None, + raw_headers: dict[str, str] | None = None, + client_ip: str | None = None, ) -> None: auth_user.user_api_key_auth = user_api_key_auth auth_user.mcp_auth_header = mcp_auth_header @@ -4519,13 +4515,13 @@ if MCP_AVAILABLE: auth_user.client_ip = client_ip def set_auth_context( - user_api_key_auth: Optional[UserAPIKeyAuth], - mcp_auth_header: Optional[str] = None, - mcp_servers: Optional[List[str]] = None, - mcp_server_auth_headers: Optional[Dict[str, Dict[str, str]]] = None, - oauth2_headers: Optional[Dict[str, str]] = None, - raw_headers: Optional[Dict[str, str]] = None, - client_ip: Optional[str] = None, + user_api_key_auth: UserAPIKeyAuth | None, + mcp_auth_header: str | None = None, + mcp_servers: list[str] | None = None, + mcp_server_auth_headers: dict[str, dict[str, str]] | None = None, + oauth2_headers: dict[str, str] | None = None, + raw_headers: dict[str, str] | None = None, + client_ip: str | None = None, ) -> MCPAuthenticatedUser: """ Set the UserAPIKeyAuth in the auth context variable. @@ -4550,14 +4546,14 @@ if MCP_AVAILABLE: return auth_user def _set_or_update_auth_context( - user_api_key_auth: Optional[UserAPIKeyAuth], - mcp_auth_header: Optional[str] = None, - mcp_servers: Optional[List[str]] = None, - mcp_server_auth_headers: Optional[Dict[str, Dict[str, str]]] = None, - oauth2_headers: Optional[Dict[str, str]] = None, - raw_headers: Optional[Dict[str, str]] = None, - client_ip: Optional[str] = None, - session_id: Optional[str] = None, + user_api_key_auth: UserAPIKeyAuth | None, + mcp_auth_header: str | None = None, + mcp_servers: list[str] | None = None, + mcp_server_auth_headers: dict[str, dict[str, str]] | None = None, + oauth2_headers: dict[str, str] | None = None, + raw_headers: dict[str, str] | None = None, + client_ip: str | None = None, + session_id: str | None = None, touch_last_seen: bool = True, copy_existing_session_auth_context: bool = False, ) -> MCPAuthenticatedUser: @@ -4601,7 +4597,7 @@ if MCP_AVAILABLE: send: Send, auth_user: MCPAuthenticatedUser, owner_fingerprint: str, - on_session_registered: Optional[Callable[[str], None]] = None, + on_session_registered: Callable[[str], None] | None = None, ) -> Send: async def wrapped_send(message: Message) -> None: if message.get("type") == "http.response.start": @@ -4620,14 +4616,14 @@ if MCP_AVAILABLE: return wrapped_send - def get_auth_context() -> Tuple[ - Optional[UserAPIKeyAuth], - Optional[str], - Optional[List[str]], - Optional[Dict[str, Dict[str, str]]], - Optional[Dict[str, str]], - Optional[Dict[str, str]], - Optional[str], + def get_auth_context() -> tuple[ + UserAPIKeyAuth | None, + str | None, + list[str] | None, + dict[str, dict[str, str]] | None, + dict[str, str] | None, + dict[str, str] | None, + str | None, ]: """ Get the UserAPIKeyAuth from the auth context variable. @@ -4681,12 +4677,12 @@ if MCP_AVAILABLE: "session identity — session object is unhashable" ) - def _recover_auth_from_session() -> Optional[MCPAuthenticatedUser]: + def _recover_auth_from_session() -> MCPAuthenticatedUser | None: session = _get_current_session() if session is None: return None - stored: Optional[MCPAuthenticatedUser] = None + stored: MCPAuthenticatedUser | None = None try: stored = _session_obj_auth_storage.get(session) except TypeError: @@ -4698,14 +4694,14 @@ if MCP_AVAILABLE: return stored - async def get_or_extract_auth_context() -> Tuple[ - Optional[UserAPIKeyAuth], - Optional[str], - Optional[List[str]], - Optional[Dict[str, Dict[str, str]]], - Optional[Dict[str, str]], - Optional[Dict[str, str]], - Optional[str], + async def get_or_extract_auth_context() -> tuple[ + UserAPIKeyAuth | None, + str | None, + list[str] | None, + dict[str, dict[str, str]] | None, + dict[str, str] | None, + dict[str, str] | None, + str | None, ]: """ Get auth context from ContextVar first, then fall back to session @@ -4744,14 +4740,14 @@ if MCP_AVAILABLE: _client_ip, ) - def get_active_mcp_session() -> Optional[_McpServerSession]: + def get_active_mcp_session() -> _McpServerSession | None: """Return the active MCP session captured during handler execution.""" session = active_mcp_session_var.get() if session is not None: return session return _get_current_session() - def get_active_auth_context() -> Optional[MCPAuthenticatedUser]: + def get_active_auth_context() -> MCPAuthenticatedUser | None: """Return auth context from ContextVar or session storage.""" auth = auth_context_var.get() if auth and isinstance(auth, MCPAuthenticatedUser): diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 972c831073f..96f6ee89d56 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -26858,12 +26858,6 @@ } ], "title": "Start Date" - }, - "total_spend": { - "default": 0.0, - "description": "Deduplicated spend of every request that called at least one tool in the window; less than the sum of per-tool attributed spend whenever multi-tool requests exist", - "title": "Total Spend", - "type": "number" } }, "title": "ToolSpendResponse", @@ -27417,7 +27411,7 @@ }, "/v1/tool/spend": { "get": { - "description": "Spend attributed to each tool over a date range, for the Cost Optimization dashboard.\n\nJoins ``LiteLLM_SpendLogToolIndex`` (which tool names ran on which request) to\n``LiteLLM_SpendLogs`` (what the request cost). A request that used multiple tools\ncounts its full spend toward each of those tools, so per-tool numbers are\nattributions. ``total_spend`` is the deduplicated spend of every request that\ncalled at least one tool in the window, so it never double counts.\n\n``start_date`` is clamped to at most 30 days before ``end_date`` (serving up to\n31 calendar dates inclusive, the same width as the endpoint's default window):\na wider requested range is clamped, and the response's ``start_date`` reflects\nthe effective window actually served.", + "description": "Spend attributed to each tool over a date range, for the Cost Optimization dashboard.\n\nReads the ``LiteLLM_DailyToolSpend`` rollup, written at request time from invoked\ntools only (MCP tool calls and response tool_calls; declaring a tool without\ninvoking it does not count). A request that invoked multiple tools counts its\nfull spend toward each of them, so per-tool numbers are attributions and do not\nsum to a deduplicated total.\n\n``by_tool`` is the top ``TOOL_SPEND_TOP_TOOLS`` tools by spend, aggregated in\nSQL, and ``daily`` covers only those tools, so the response is bounded by\ndays x TOOL_SPEND_TOP_TOOLS regardless of the requested range or how many\ndistinct tool names exist.", "operationId": "get_tool_spend_v1_tool_spend_get", "parameters": [ { @@ -27588,7 +27582,7 @@ }, "/v1/tool/{tool_name}/logs": { "get": { - "description": "Return paginated spend logs for requests that used this tool (from SpendLogToolIndex).", + "description": "Return paginated spend logs for requests that invoked this tool (from SpendLogToolIndex).\nDeclaring a tool in a request body without the model invoking it does not create an entry.", "operationId": "get_tool_usage_logs_v1_tool__tool_name__logs_get", "parameters": [ { diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index fe91ea27b35..d873ad477de 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -3654,6 +3654,13 @@ DB_CONNECTION_ERROR_TYPES = ( httpx.ReadTimeout, ) +# What a NON-IDEMPOTENT write (increment upsert) may retry: only ConnectError +# proves the statements never reached the database. Post-send errors are +# ambiguous; a stalled statement can leave its transaction open on the pooled +# connection, where a retry stacks a second increment set into the same commit. +# Idempotent writes (create_many with skip_duplicates) may retry the full tuple. +DB_RETRY_SAFE_ERROR_TYPES = (httpx.ConnectError,) + class SSOUserDefinedValues(TypedDict): models: List[str] diff --git a/litellm/proxy/agent_endpoints/agent_registry.py b/litellm/proxy/agent_endpoints/agent_registry.py index 1373d055d4f..43139b18162 100644 --- a/litellm/proxy/agent_endpoints/agent_registry.py +++ b/litellm/proxy/agent_endpoints/agent_registry.py @@ -1,7 +1,8 @@ import hashlib import json +from collections.abc import Iterator, Mapping, Sequence from datetime import datetime, timezone -from typing import Any, Dict, List, Optional +from typing import Any, Protocol, TypedDict import litellm from litellm.litellm_core_utils.safe_json_dumps import safe_dumps @@ -13,9 +14,81 @@ from litellm.repositories.table_repositories import AgentsRepository from litellm.types.agents import AgentConfig, AgentResponse, PatchAgentRequest +class AgentObjectPermissionRecord(Protocol): + def model_dump(self) -> dict[str, object]: ... + + def dict(self) -> dict[str, object]: ... + + +class AgentRecordDump(TypedDict): + agent_id: str + agent_name: str + litellm_params: dict[str, object] | None + agent_card_params: dict[str, object] + static_headers: dict[str, str] | None + extra_headers: list[str] | None + object_permission: dict[str, object] | None + spend: float + tpm_limit: int | None + rpm_limit: int | None + session_tpm_limit: int | None + session_rpm_limit: int | None + created_at: datetime + updated_at: datetime + created_by: str | None + updated_by: str | None + + +class AgentRecord(Protocol): + agent_id: str + agent_name: str + object_permission_id: str | None + object_permission: AgentObjectPermissionRecord | None + spend: float + + def model_dump(self) -> AgentRecordDump: ... + + def __iter__(self) -> Iterator[tuple[str, object]]: ... + + +class AgentTableClient(Protocol): + async def create( + self, + data: Mapping[str, object], + include: Mapping[str, bool] | None = None, + ) -> AgentRecord: ... + + async def find_unique( + self, + where: Mapping[str, object], + include: Mapping[str, bool] | None = None, + ) -> AgentRecord | None: ... + + async def find_many( + self, + where: Mapping[str, object] | None = None, + order: Mapping[str, str] | None = None, + include: Mapping[str, bool] | None = None, + ) -> Sequence[AgentRecord]: ... + + async def update( + self, + where: Mapping[str, object], + data: Mapping[str, object], + include: Mapping[str, bool] | None = None, + ) -> AgentRecord: ... + + async def delete(self, where: Mapping[str, object]) -> AgentRecord: ... + + +def agents_table(prisma_client: PrismaClient) -> AgentTableClient: + table: AgentTableClient = AgentsRepository(prisma_client).table + return table + + class AgentRegistry: def __init__(self): - self.agent_list: List[AgentResponse] = [] + self.agent_list: list[AgentResponse] = [] def reset_agent_list(self): self.agent_list = [] @@ -26,13 +99,13 @@ class AgentRegistry: def deregister_agent(self, agent_name: str): self.agent_list = [agent for agent in self.agent_list if agent.agent_name != agent_name] - def get_agent_list(self, agent_names: Optional[List[str]] = None): + def get_agent_list(self, agent_names: Sequence[str] | None = None): if agent_names is not None: return [agent for agent in self.agent_list if agent.agent_name in agent_names] return self.agent_list - def get_public_agent_list(self) -> List[AgentResponse]: - public_agent_list: List[AgentResponse] = [] + def get_public_agent_list(self) -> list[AgentResponse]: + public_agent_list: list[AgentResponse] = [] if litellm.public_agent_groups is None: return public_agent_list for agent in self.agent_list: @@ -43,7 +116,7 @@ class AgentRegistry: def _create_agent_id(self, agent_config: AgentConfig) -> str: return hashlib.sha256(json.dumps(agent_config, sort_keys=True).encode()).hexdigest() - def load_agents_from_config(self, agent_config: Optional[List[AgentConfig]] = None): + def load_agents_from_config(self, agent_config: Sequence[AgentConfig] | None = None): if agent_config is None: return None @@ -63,8 +136,8 @@ class AgentRegistry: def load_agents_from_db_and_config( self, - agent_config: Optional[List[AgentConfig]] = None, - db_agents: Optional[List[Dict[str, Any]]] = None, + agent_config: Sequence[AgentConfig] | None = None, + db_agents: list[dict[str, Any]] | None = None, ): self.reset_agent_list() @@ -96,7 +169,7 @@ class AgentRegistry: agent: AgentConfig, prisma_client: PrismaClient, created_by: str, - agent_id: Optional[str] = None, + agent_id: str | None = None, ) -> AgentResponse: """ Add an agent to the database. @@ -126,18 +199,18 @@ class AgentRegistry: agent_card_params: str = safe_dumps(agent_card_params_dict) # Handle object_permission (MCP tool access for agent) - object_permission_id: Optional[str] = None + object_permission_id: str | None = None if agent.get("object_permission") is not None: agent_copy = dict(agent) object_permission_id = await handle_update_object_permission_common(agent_copy, None, prisma_client) # Serialize static_headers static_headers_obj = agent.get("static_headers") - static_headers_val: Optional[str] = safe_dumps(dict(static_headers_obj)) if static_headers_obj else None + static_headers_val: str | None = safe_dumps(dict(static_headers_obj)) if static_headers_obj else None - extra_headers_val: Optional[List[str]] = agent.get("extra_headers") + extra_headers_val = agent.get("extra_headers") - create_data: Dict[str, Any] = { + create_data: dict[str, object] = { "agent_name": agent_name, "litellm_params": litellm_params, "agent_card_params": agent_card_params, @@ -166,7 +239,7 @@ class AgentRegistry: create_data[rate_field] = _val # Create agent in DB - created_agent = await AgentsRepository(prisma_client).table.create( + created_agent = await agents_table(prisma_client).create( data=create_data, include={"object_permission": True}, ) @@ -181,12 +254,12 @@ class AgentRegistry: except Exception as e: raise Exception(f"Error adding agent to DB: {str(e)}") - async def delete_agent_from_db(self, agent_id: str, prisma_client: PrismaClient) -> Dict[str, Any]: + async def delete_agent_from_db(self, agent_id: str, prisma_client: PrismaClient) -> Mapping[str, object]: """ Delete an agent from the database """ try: - deleted_agent = await AgentsRepository(prisma_client).table.delete(where={"agent_id": agent_id}) + deleted_agent = await agents_table(prisma_client).delete(where={"agent_id": agent_id}) return dict(deleted_agent) except Exception as e: raise Exception(f"Error deleting agent from DB: {str(e)}") @@ -221,7 +294,7 @@ class AgentRegistry: raise Exception(f"Agent with ID {agent_id} not found") augment_agent = {**existing_agent, **agent} - update_data: Dict[str, Any] = {} + update_data: dict[str, Any] = {} if augment_agent.get("agent_name"): update_data["agent_name"] = augment_agent.get("agent_name") if augment_agent.get("litellm_params"): @@ -254,7 +327,7 @@ class AgentRegistry: if object_permission_id is not None: update_data["object_permission_id"] = object_permission_id # Patch agent in DB - patched_agent = await AgentsRepository(prisma_client).table.update( + patched_agent = await agents_table(prisma_client).update( where={"agent_id": agent_id}, data={ **update_data, @@ -307,9 +380,9 @@ class AgentRegistry: static_headers_val_u: str = ( safe_dumps(dict(static_headers_obj_u)) if static_headers_obj_u is not None else safe_dumps({}) ) - extra_headers_val_u: List[str] = agent.get("extra_headers") or [] + extra_headers_val_u = agent.get("extra_headers") or [] - update_data: Dict[str, Any] = { + update_data: dict[str, object] = { "agent_name": agent_name, "litellm_params": litellm_params, "agent_card_params": agent_card_params, @@ -330,7 +403,7 @@ class AgentRegistry: update_data[rate_field] = _val if agent.get("object_permission") is not None: - existing_agent = await AgentsRepository(prisma_client).table.find_unique(where={"agent_id": agent_id}) + existing_agent = await agents_table(prisma_client).find_unique(where={"agent_id": agent_id}) existing_object_permission_id = ( existing_agent.object_permission_id if existing_agent is not None else None ) @@ -344,7 +417,7 @@ class AgentRegistry: update_data["object_permission_id"] = object_permission_id # Update agent in DB - updated_agent = await AgentsRepository(prisma_client).table.update( + updated_agent = await agents_table(prisma_client).update( where={"agent_id": agent_id}, data=update_data, include={"object_permission": True}, @@ -363,17 +436,17 @@ class AgentRegistry: @staticmethod async def get_all_agents_from_db( prisma_client: PrismaClient, - ) -> List[Dict[str, Any]]: + ) -> list[dict[str, object]]: """ Get all agents from the database """ try: - agents_from_db = await AgentsRepository(prisma_client).table.find_many( + agents_from_db = await agents_table(prisma_client).find_many( order={"created_at": "desc"}, include={"object_permission": True}, ) - agents: List[Dict[str, Any]] = [] + agents: list[dict[str, object]] = [] for agent in agents_from_db: agent_dict = dict(agent) # object_permission is eagerly loaded via include above @@ -391,7 +464,7 @@ class AgentRegistry: def get_agent_by_id( self, agent_id: str, - ) -> Optional[AgentResponse]: + ) -> AgentResponse | None: """ Get an agent by its ID from the database """ @@ -404,7 +477,7 @@ class AgentRegistry: except Exception as e: raise Exception(f"Error getting agent from DB: {str(e)}") - def get_agent_by_name(self, agent_name: str) -> Optional[AgentResponse]: + def get_agent_by_name(self, agent_name: str) -> AgentResponse | None: """ Get an agent by its name from the database """ diff --git a/litellm/proxy/agent_endpoints/endpoints.py b/litellm/proxy/agent_endpoints/endpoints.py index 2421f270974..c3308bbfa8c 100644 --- a/litellm/proxy/agent_endpoints/endpoints.py +++ b/litellm/proxy/agent_endpoints/endpoints.py @@ -11,9 +11,11 @@ Follows the A2A Spec. import asyncio import os import uuid -from typing import Any, Dict, List, Mapping +from collections.abc import Mapping, Sequence +from typing import TypedDict from fastapi import APIRouter, Depends, HTTPException, Query, Request +from typing_extensions import Required import litellm from litellm._logging import verbose_proxy_logger @@ -30,6 +32,7 @@ from litellm.proxy.common_utils.rbac_utils import check_feature_access_for_user from litellm.proxy.management_endpoints.common_daily_activity import get_daily_activity from litellm.proxy.utils import get_custom_url from litellm.types.agents import ( + AgentCard, AgentConfig, AgentKeySummary, AgentMakePublicResponse, @@ -49,7 +52,7 @@ def _proxy_base_url(http_request: Request) -> str: return get_custom_url(str(http_request.base_url), route=None) -def _validate_protocol_version(upstream_card: Mapping[str, Any] | None) -> None: +def _validate_protocol_version(upstream_card: AgentCard | None) -> None: """Reject an agent card pinning an unsupported A2A protocol version.""" version = upstream_card.get("protocolVersion") if upstream_card else None if version is not None and normalize_protocol_version(version) is None: @@ -63,12 +66,12 @@ def _validate_protocol_version(upstream_card: Mapping[str, Any] | None) -> None: def _build_merged_agent_card( - upstream_card: Mapping[str, Any] | None, + upstream_card: AgentCard | None, *, agent_id: str, http_request: Request, agent_name: str | None = None, -) -> Dict[str, Any]: +) -> dict[str, object]: """Apply the LiteLLM-fronting merge to ``upstream_card`` for ``agent_id``.""" proxy_base = _proxy_base_url(http_request) _validate_protocol_version(upstream_card) @@ -88,7 +91,7 @@ def _build_merged_agent_card( router = APIRouter() -async def _attach_keys_to_agents(agents: list[AgentResponse], prisma_client) -> None: +async def _attach_keys_to_agents(agents: Sequence[AgentResponse], prisma_client) -> None: """Attach each agent's virtual keys, derived from the key table's agent_id foreign key. Mirrors how spend is joined into the agent response so the UI never has to cross-reference a full key dump client-side. Only non-secret @@ -113,7 +116,7 @@ async def _attach_keys_to_agents(agents: list[AgentResponse], prisma_client) -> def _redact_sensitive_agent_fields( - agents: list[AgentResponse], + agents: Sequence[AgentResponse], ) -> list[AgentResponse]: """ Return copies of the given agents with sensitive configuration fields @@ -156,9 +159,15 @@ AGENT_HEALTH_CHECK_TIMEOUT_SECONDS = float(os.environ.get("LITELLM_AGENT_HEALTH_ AGENT_HEALTH_CHECK_GATHER_TIMEOUT_SECONDS = float(os.environ.get("LITELLM_AGENT_HEALTH_CHECK_GATHER_TIMEOUT", "30.0")) +class _AgentHealthResult(TypedDict, total=False): + agent_id: Required[str] + healthy: Required[bool] + error: str + + async def _check_agent_url_health( agent: AgentResponse, -) -> Dict[str, Any]: +) -> _AgentHealthResult: """ Perform a GET request against the agent's URL and return the health result. @@ -194,7 +203,7 @@ async def _check_agent_url_health( "/v1/agents", tags=["[beta] A2A Agents"], dependencies=[Depends(user_api_key_auth)], - response_model=List[AgentResponse], + response_model=list[AgentResponse], ) async def get_agents( request: Request, @@ -230,7 +239,7 @@ async def get_agents( ) try: - returned_agents: List[AgentResponse] = [] + returned_agents: list[AgentResponse] = [] # Admin users get all agents if ( @@ -256,7 +265,7 @@ async def get_agents( if prisma_client is not None: agent_ids = [agent.agent_id for agent in returned_agents] if agent_ids: - db_agents = await AgentsRepository(prisma_client).table.find_many( + db_agents = await agents_table(prisma_client).find_many( where={"agent_id": {"in": agent_ids}}, ) spend_map = {a.agent_id: a.spend for a in db_agents} @@ -285,7 +294,7 @@ async def get_agents( agents_with_url = [agent for agent in returned_agents if (agent.agent_card_params or {}).get("url")] agents_without_url = [agent for agent in returned_agents if not (agent.agent_card_params or {}).get("url")] try: - health_results = await asyncio.wait_for( + health_results: Sequence[_AgentHealthResult] = await asyncio.wait_for( asyncio.gather(*[_check_agent_url_health(agent) for agent in agents_with_url]), timeout=AGENT_HEALTH_CHECK_GATHER_TIMEOUT_SECONDS, ) @@ -317,10 +326,12 @@ async def get_agents( #### CRUD ENDPOINTS FOR AGENTS #### +from litellm.proxy.agent_endpoints.agent_registry import ( + agents_table, +) from litellm.proxy.agent_endpoints.agent_registry import ( global_agent_registry as AGENT_REGISTRY, ) -from litellm.repositories.table_repositories import AgentsRepository @router.post( @@ -487,7 +498,7 @@ async def get_agent_by_id( try: agent = AGENT_REGISTRY.get_agent_by_id(agent_id=agent_id) if agent is None: - agent_row = await AgentsRepository(prisma_client).table.find_unique( + agent_row = await agents_table(prisma_client).find_unique( where={"agent_id": agent_id}, include={"object_permission": True}, ) @@ -501,7 +512,7 @@ async def get_agent_by_id( agent = AgentResponse(**agent_dict) # type: ignore else: # Agent found in memory — refresh spend from DB - db_row = await AgentsRepository(prisma_client).table.find_unique(where={"agent_id": agent_id}) + db_row = await agents_table(prisma_client).find_unique(where={"agent_id": agent_id}) if db_row is not None: agent.spend = db_row.spend @@ -578,7 +589,7 @@ async def update_agent( try: # Check if agent exists - existing_agent = await AgentsRepository(prisma_client).table.find_unique(where={"agent_id": agent_id}) + existing_agent = await agents_table(prisma_client).find_unique(where={"agent_id": agent_id}) if existing_agent is not None: existing_agent = dict(existing_agent) @@ -680,7 +691,7 @@ async def patch_agent( try: # Check if agent exists - existing_agent = await AgentsRepository(prisma_client).table.find_unique(where={"agent_id": agent_id}) + existing_agent = await agents_table(prisma_client).find_unique(where={"agent_id": agent_id}) if existing_agent is not None: existing_agent = dict(existing_agent) @@ -767,9 +778,9 @@ async def delete_agent( try: # Check if agent exists - existing_agent = await AgentsRepository(prisma_client).table.find_unique(where={"agent_id": agent_id}) + existing_agent = await agents_table(prisma_client).find_unique(where={"agent_id": agent_id}) if existing_agent is not None: - existing_agent = dict[Any, Any](existing_agent) + existing_agent = dict[str, object](existing_agent) if existing_agent is None: raise HTTPException(status_code=404, detail=f"Agent with ID {agent_id} not found in DB.") @@ -849,7 +860,7 @@ async def make_agent_public( agent = AGENT_REGISTRY.get_agent_by_id(agent_id=agent_id) if agent is None: # check if agent exists in DB - agent = await AgentsRepository(prisma_client).table.find_unique(where={"agent_id": agent_id}) + agent = await agents_table(prisma_client).find_unique(where={"agent_id": agent_id}) if agent is not None: agent = AgentResponse(**agent.model_dump()) # type: ignore @@ -966,7 +977,7 @@ async def make_agents_public( agent = AGENT_REGISTRY.get_agent_by_id(agent_id=agent_id) if agent is None: # check if agent exists in DB - agent = await AgentsRepository(prisma_client).table.find_unique(where={"agent_id": agent_id}) + agent = await agents_table(prisma_client).find_unique(where={"agent_id": agent_id}) if agent is not None: agent = AgentResponse(**agent.model_dump()) # type: ignore @@ -1031,7 +1042,7 @@ async def get_agent_daily_activity( ) agent_ids_list = agent_ids.split(",") if agent_ids else None - exclude_agent_ids_list: List[str] | None = None + exclude_agent_ids_list: list[str] | None = None if exclude_agent_ids: exclude_agent_ids_list = exclude_agent_ids.split(",") if exclude_agent_ids else None @@ -1044,7 +1055,7 @@ async def get_agent_daily_activity( ) from litellm.proxy.management_endpoints.common_utils import _user_has_admin_view - where_condition: Dict[str, Any] = {} + where_condition: dict[str, object] = {} if not _user_has_admin_view(user_api_key_dict): permitted_agent_ids = await AgentRequestHandler.get_allowed_agents(user_api_key_auth=user_api_key_dict) # `get_allowed_agents` returns an empty list when the caller's key @@ -1058,7 +1069,7 @@ async def get_agent_daily_activity( if user_api_key_dict.user_id is None: permitted_agent_ids = [] else: - owned_records = await AgentsRepository(prisma_client).table.find_many( + owned_records = await agents_table(prisma_client).find_many( where={"created_by": user_api_key_dict.user_id} ) permitted_agent_ids = [a.agent_id for a in owned_records] @@ -1093,8 +1104,10 @@ async def get_agent_daily_activity( if agent_ids_list: where_condition["agent_id"] = {"in": list(agent_ids_list)} - agent_records = await AgentsRepository(prisma_client).table.find_many(where=where_condition) - agent_metadata = {agent.agent_id: {"agent_name": agent.agent_name} for agent in agent_records} + agent_records = await agents_table(prisma_client).find_many(where=where_condition) + agent_metadata: Mapping[str, dict[str, object]] = { + agent.agent_id: {"agent_name": agent.agent_name} for agent in agent_records + } return await get_daily_activity( prisma_client=prisma_client, diff --git a/litellm/proxy/common_utils/callback_utils.py b/litellm/proxy/common_utils/callback_utils.py index 33bca782e0b..8eb122f23af 100644 --- a/litellm/proxy/common_utils/callback_utils.py +++ b/litellm/proxy/common_utils/callback_utils.py @@ -31,6 +31,10 @@ _EXTRA_SENSITIVE_CALLBACK_KEYS = {"gcs_path_service_account"} # already-encrypted input cheaply (no decrypt-attempt round trip) and # avoid double-encrypting if `LITELLM_SALT_KEY` is rotated between writes. _CALLBACK_VAR_ENCRYPTED_PREFIX = "litellm_enc::" +# Metadata slots that hold operator-configured callback setup (and therefore +# integration credentials). Resolved from UserAPIKeyAuth during pre-call setup, +# never read back off the copies stamped into request metadata. +_CALLBACK_CONFIG_SLOTS = frozenset({"logging", "callback_settings"}) blue_color_code = "\033[94m" reset_color_code = "\033[0m" @@ -547,6 +551,13 @@ def normalize_callback_names(callbacks: Iterable[Any]) -> List[Any]: return [c.lower() if isinstance(c, str) else c for c in callbacks] +def strip_callback_config(metadata: dict[str, Any] | None) -> dict[str, Any] | None: + """Return key/team metadata without the slots that carry callback credentials.""" + if not isinstance(metadata, dict): + return metadata + return {k: v for k, v in metadata.items() if k not in _CALLBACK_CONFIG_SLOTS} + + def encrypt_callback_vars(metadata: Any) -> Any: """Return a deep copy of metadata with callback_vars values encrypted at rest. diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 2262141f426..fd8132fef22 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -34,7 +34,7 @@ from litellm.constants import ( ) from litellm.litellm_core_utils.safe_json_loads import safe_json_loads from litellm.proxy._types import ( - DB_CONNECTION_ERROR_TYPES, + DB_RETRY_SAFE_ERROR_TYPES, BaseDailySpendTransaction, DailyAgentSpendTransaction, DailyEndUserSpendTransaction, @@ -182,6 +182,12 @@ class DBSpendUpdateWriter: payload=payload, prisma_client=prisma_client, ) + await self._enqueue_tool_usage_transaction( + payload=payload, + completion_response=completion_response, + prisma_client=prisma_client, + kwargs=kwargs, + ) else: verbose_proxy_logger.debug( "disable_spend_logs=True. Skipping writing spend logs to db. Other spend updates - Key/User/Team table will still occur." @@ -223,6 +229,36 @@ class DBSpendUpdateWriter: end_user_id, ) + async def _enqueue_tool_usage_transaction( + self, + payload: SpendLogsPayload, + completion_response: "litellm.ModelResponse | Any | Exception | None", + prisma_client: "PrismaClient | None", + kwargs: "dict | None" = None, + ) -> None: + try: + if prisma_client is None: + return + from litellm.proxy.db.spend_log_tool_index import ( + build_tool_usage_transaction, + ) + + transaction = build_tool_usage_transaction( + request_id=payload["request_id"], + start_time_iso=str(payload["startTime"]), + mcp_namespaced_tool_name=payload.get("mcp_namespaced_tool_name"), + spend=payload["spend"], + total_tokens=payload["total_tokens"], + completion_response=completion_response, + realtime_tool_calls=(kwargs or {}).get("realtime_tool_calls"), + ) + if transaction is None: + return + async with prisma_client._tool_usage_transactions_lock: + prisma_client.tool_usage_transactions.append(transaction) + except Exception as e: + verbose_proxy_logger.debug("_enqueue_tool_usage_transaction error (non-blocking): %s", e) + def _enqueue_tool_registry_upsert( self, kwargs: Optional[dict], @@ -299,21 +335,10 @@ class DBSpendUpdateWriter: _enqueue(name) # --- Response tool_calls (OpenAI format; Anthropic pass-through converts tool_use here) --- - if completion_response is not None and hasattr(completion_response, "choices"): - for choice in completion_response.choices or []: - message = getattr(choice, "message", None) - if message is None: - continue - tool_calls = getattr(message, "tool_calls", None) - if not tool_calls: - continue - for tc in tool_calls: - fn = getattr(tc, "function", None) - if fn is None: - continue - tool_name = getattr(fn, "name", None) - if tool_name: - _enqueue(tool_name) + from litellm.proxy.db.spend_log_tool_index import response_tool_call_names + + for tool_name in response_tool_call_names(completion_response): + _enqueue(tool_name) except Exception as e: verbose_proxy_logger.debug("_enqueue_tool_registry_upsert error (non-blocking): %s", e) @@ -1096,7 +1121,7 @@ class DBSpendUpdateWriter: data={"spend": {"increment": response_cost}}, ) break - except DB_CONNECTION_ERROR_TYPES as e: + except DB_RETRY_SAFE_ERROR_TYPES as e: if i >= n_retry_times: # If we've reached the maximum number of retries _raise_failed_update_spend_exception( e=e, @@ -1139,7 +1164,7 @@ class DBSpendUpdateWriter: }, ) break - except DB_CONNECTION_ERROR_TYPES as e: + except DB_RETRY_SAFE_ERROR_TYPES as e: if i >= n_retry_times: # If we've reached the maximum number of retries _raise_failed_update_spend_exception( e=e, @@ -1172,7 +1197,7 @@ class DBSpendUpdateWriter: data={"spend": {"increment": response_cost}}, ) break - except DB_CONNECTION_ERROR_TYPES as e: + except DB_RETRY_SAFE_ERROR_TYPES as e: if i >= n_retry_times: # If we've reached the maximum number of retries _raise_failed_update_spend_exception( e=e, @@ -1219,7 +1244,7 @@ class DBSpendUpdateWriter: ) # Transaction succeeded, break out of retry loop break - except DB_CONNECTION_ERROR_TYPES as e: + except DB_RETRY_SAFE_ERROR_TYPES as e: if i >= n_retry_times: # If we've reached the maximum number of retries _raise_failed_update_spend_exception( e=e, @@ -1261,7 +1286,7 @@ class DBSpendUpdateWriter: data={"spend": {"increment": response_cost}}, ) break - except DB_CONNECTION_ERROR_TYPES as e: + except DB_RETRY_SAFE_ERROR_TYPES as e: if i >= n_retry_times: # If we've reached the maximum number of retries _raise_failed_update_spend_exception( e=e, @@ -1347,7 +1372,7 @@ class DBSpendUpdateWriter: data={"spend": {"increment": response_cost}}, ) break - except DB_CONNECTION_ERROR_TYPES as e: + except DB_RETRY_SAFE_ERROR_TYPES as e: if i >= n_retry_times: _raise_failed_update_spend_exception( e=e, @@ -1644,7 +1669,7 @@ class DBSpendUpdateWriter: break - except DB_CONNECTION_ERROR_TYPES as e: + except DB_RETRY_SAFE_ERROR_TYPES as e: if i >= n_retry_times: _raise_failed_update_spend_exception( e=e, diff --git a/litellm/proxy/db/spend_log_tool_index.py b/litellm/proxy/db/spend_log_tool_index.py index 80036e235f7..802e893d473 100644 --- a/litellm/proxy/db/spend_log_tool_index.py +++ b/litellm/proxy/db/spend_log_tool_index.py @@ -1,140 +1,150 @@ """ -Track tool usage for the dashboard: insert into SpendLogToolIndex when spend logs -are written, so "last N requests for tool X" and "how is this tool called in production" -queries are fast. +Tool usage tracking for the dashboard. + +At request time the spend writer builds one ToolUsageTransaction per request that +invoked tools (MCP namespaced tool name plus response tool_calls; declared-but-not- +invoked tools are excluded) and queues it on the prisma client. The spend-log flush +job drains the queue into LiteLLM_SpendLogToolIndex (per-request drill-down) and +LiteLLM_DailyToolSpend (the per-day rollup the Cost Optimization card reads) in a +single transaction, so a failed flush never leaves a partial rollup increment. """ +from __future__ import annotations + +import asyncio +import random +from dataclasses import dataclass from datetime import datetime, timezone -from typing import Any, Dict, List, Set +from itertools import groupby +from typing import TYPE_CHECKING, Any, Sequence -from litellm._logging import verbose_proxy_logger -from litellm.litellm_core_utils.safe_json_loads import safe_json_loads -from litellm.proxy.utils import PrismaClient -from litellm.repositories.table_repositories import SpendLogToolIndexRepository +from litellm.proxy._types import DB_RETRY_SAFE_ERROR_TYPES + +if TYPE_CHECKING: + from litellm.proxy.utils import PrismaClient -def _add_tool_calls_to_set(tool_calls: Any, out: Set[str]) -> None: - """Extract tool names from OpenAI-style tool_calls list into out.""" - if not isinstance(tool_calls, list): - return - for tc in tool_calls: - if not isinstance(tc, dict): - continue - fn = tc.get("function") - if isinstance(fn, dict): - name = fn.get("name") - if name and isinstance(name, str) and name.strip(): - out.add(name.strip()) +@dataclass(frozen=True, slots=True) +class ToolUsageTransaction: + request_id: str + date: str + start_time: datetime + tool_names: tuple[str, ...] + spend: float + total_tokens: int -def _parse_tool_names_from_payload(payload: Dict[str, Any]) -> Set[str]: - """ - Extract deduplicated tool names from a spend log payload. - Sources: mcp_namespaced_tool_name, response (tool_calls), proxy_server_request (tools). - """ - tool_names: Set[str] = set() +def response_tool_call_names(completion_response: Any) -> tuple[str, ...]: + """Tool names invoked in a completion response, in call order, for any response + surface get_tool_calls_from_response understands (chat completions, Responses + API output items, Anthropic Messages tool_use blocks). Reads every choice of + an ``n>1`` chat response: each choice cost money and its tool calls ran.""" + if completion_response is None or isinstance(completion_response, Exception): + return () + from litellm.litellm_core_utils.prompt_templates.factory import ( + get_tool_calls_from_response, + ) - # Top-level MCP tool name (single tool per request for that flow) - mcp_name = payload.get("mcp_namespaced_tool_name") - if mcp_name and isinstance(mcp_name, str) and mcp_name.strip(): - tool_names.add(mcp_name.strip()) - - # Response: OpenAI-style tool_calls[].function.name or choices[0].message.tool_calls - response_raw = payload.get("response") - if response_raw: - response_obj = safe_json_loads(response_raw, default=None) if isinstance(response_raw, str) else response_raw - if isinstance(response_obj, dict): - _add_tool_calls_to_set(response_obj.get("tool_calls"), tool_names) - choices = response_obj.get("choices") - if isinstance(choices, list) and choices: - msg = choices[0].get("message") if isinstance(choices[0], dict) else None - if isinstance(msg, dict): - _add_tool_calls_to_set(msg.get("tool_calls"), tool_names) - - # Request body: tools[].function.name - request_raw = payload.get("proxy_server_request") - if request_raw: - request_obj = safe_json_loads(request_raw, default=None) if isinstance(request_raw, str) else request_raw - if isinstance(request_obj, dict): - body = request_obj.get("body", request_obj) - if isinstance(body, dict): - request_obj = body - if isinstance(request_obj, dict): - tools = request_obj.get("tools") - if isinstance(tools, list): - for t in tools: - if isinstance(t, dict): - fn = t.get("function") - if isinstance(fn, dict): - name = fn.get("name") - if name and isinstance(name, str) and name.strip(): - tool_names.add(name.strip()) - - return tool_names + return tuple( + stripped + for tool_call in get_tool_calls_from_response(completion_response, include_all_choices=True) + if isinstance(name := tool_call.get("name"), str) and (stripped := name.strip()) + ) -async def process_spend_logs_tool_usage( - prisma_client: PrismaClient, - logs_to_process: List[Dict[str, Any]], -) -> None: - """ - After spend logs are written: insert SpendLogToolIndex rows from each payload. - Extracts tool names from mcp_namespaced_tool_name, response tool_calls, and - proxy_server_request tools. - """ - if not logs_to_process: - return - - index_rows: List[Dict[str, Any]] = [] - - for payload in logs_to_process: - request_id = payload.get("request_id") - start_time = payload.get("startTime") - if not request_id or not start_time: - continue - if isinstance(start_time, str): - try: - start_time = datetime.fromisoformat(start_time.replace("Z", "+00:00")) - except (ValueError, TypeError): - continue - if start_time.tzinfo is None: - start_time = start_time.replace(tzinfo=timezone.utc) - - tool_names = _parse_tool_names_from_payload(payload) - for tool_name in tool_names: - index_rows.append( - { - "request_id": request_id, - "tool_name": tool_name, - "start_time": start_time, - } - ) - - if not index_rows: - return - +def build_tool_usage_transaction( + request_id: str, + start_time_iso: str, + mcp_namespaced_tool_name: str | None, + spend: float, + total_tokens: int, + completion_response: Any, + realtime_tool_calls: Any = None, +) -> ToolUsageTransaction | None: + """None when the request invoked no tools. Realtime sessions carry invoked + tools in kwargs["realtime_tool_calls"] (OpenAI tool_calls shape) rather than + on a response object, so they are normalized through the same owner by + wrapping them in the chat-completion shape. Date derivation must match the + daily spend writer's ``startTime.split("T")[0]`` so rollup rows land in the + same UTC day bucket as LiteLLM_DailyUserSpend.""" + mcp_names = ( + (mcp_namespaced_tool_name.strip(),) if mcp_namespaced_tool_name and mcp_namespaced_tool_name.strip() else () + ) + realtime_names = ( + response_tool_call_names({"choices": [{"message": {"tool_calls": realtime_tool_calls}}]}) + if realtime_tool_calls + else () + ) + tool_names = tuple(dict.fromkeys(mcp_names + response_tool_call_names(completion_response) + realtime_names)) + if not tool_names: + return None try: - index_data = [] - for r in index_rows: - st = r["start_time"] - if isinstance(st, str): - try: - st = datetime.fromisoformat(st.replace("Z", "+00:00")) - except (ValueError, TypeError): - continue - if st.tzinfo is None: - st = st.replace(tzinfo=timezone.utc) - index_data.append( - { - "request_id": r["request_id"], - "tool_name": r["tool_name"], - "start_time": st, - } - ) - if index_data: - await SpendLogToolIndexRepository(prisma_client).table.create_many( - data=index_data, - skip_duplicates=True, - ) - except Exception as e: - verbose_proxy_logger.warning("Tool usage tracking (SpendLogToolIndex) failed (non-fatal): %s", e) + start_time = datetime.fromisoformat(start_time_iso.replace("Z", "+00:00")) + except ValueError: + return None + return ToolUsageTransaction( + request_id=request_id, + date=start_time_iso.split("T")[0], + start_time=start_time if start_time.tzinfo else start_time.replace(tzinfo=timezone.utc), + tool_names=tool_names, + spend=spend, + total_tokens=total_tokens, + ) + + +async def flush_tool_usage_transactions( + prisma_client: PrismaClient, + transactions: Sequence[ToolUsageTransaction], + n_retry_times: int = 3, +) -> None: + """Write index rows and rollup upserts for a drained queue batch in one + transaction. Retries only ConnectError, the one failure that proves the + statements never reached the database. Post-send failures (Read timeouts + and errors) are ambiguous and are NOT retried: the engine can abandon the + transaction open on the pooled connection, so a retry's statements stack + into the same transaction and one commit applies both increment sets. + Ambiguous failures drop the batch; the caller logs it at error. Callers + must not add their own retry around this function.""" + if not transactions: + return + + index_rows = [ + {"request_id": txn.request_id, "tool_name": tool_name, "start_time": txn.start_time} + for txn in transactions + for tool_name in txn.tool_names + ] + per_tool_day = sorted( + ((txn.date, tool_name, txn.spend, txn.total_tokens) for txn in transactions for tool_name in txn.tool_names), + key=lambda entry: (entry[0], entry[1]), + ) + + for attempt in range(n_retry_times + 1): + try: + async with prisma_client.db.batch_() as batcher: + batcher.litellm_spendlogtoolindex.create_many(data=index_rows, skip_duplicates=True) + for (date_key, tool_name), grouped in groupby(per_tool_day, key=lambda entry: (entry[0], entry[1])): + entries = tuple(grouped) + spend = sum(entry[2] for entry in entries) + total_tokens = sum(entry[3] for entry in entries) + batcher.litellm_dailytoolspend.upsert( + where={"date_tool_name": {"date": date_key, "tool_name": tool_name}}, + data={ + "create": { + "date": date_key, + "tool_name": tool_name, + "spend": spend, + "total_tokens": total_tokens, + "request_count": len(entries), + }, + "update": { + "spend": {"increment": spend}, + "total_tokens": {"increment": total_tokens}, + "request_count": {"increment": len(entries)}, + }, + }, + ) + return + except DB_RETRY_SAFE_ERROR_TYPES: + if attempt >= n_retry_times: + raise + await asyncio.sleep(2**attempt + random.uniform(0, 1)) diff --git a/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py b/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py index 7b166185865..2735acd7787 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py +++ b/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py @@ -28,6 +28,11 @@ from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, # pyright: ignore[reportUnknownVariableType] httpxSpecialProvider, ) +from litellm.proxy.guardrails.guardrail_hooks.content_text import ( + content_to_text, + is_all_text_parts, + merge_rewritten_text_parts, +) from litellm.secret_managers.main import get_secret_str from litellm.types.guardrails import GuardrailEventHooks, Mode from litellm.types.integrations.custom_logger import AgenticLoopPlan, AgenticLoopRequestPatch @@ -51,6 +56,60 @@ def _is_object_list(value: object) -> TypeGuard[list[object]]: # guard-ok: isin return isinstance(value, list) +def _flatten_messages_for_compression(messages: list[dict[str, object]]) -> list[dict[str, object]]: + """Collapse all-text list-of-parts content to plain strings for /v1/compress. + + The compression service's transforms only rewrite string content and skip + the OpenAI list-of-parts shape, which is what every Anthropic-format + request translates to. Only rows whose parts are ALL text are flattened: + cache_control breakpoints are positional (each caches the prefix ending + at its part), so merging text across a non-text part would move a later + breakpoint to the other side of it. Rows with non-text parts are sent + unchanged and pass through the service untouched. + """ + flattened: list[dict[str, object]] = [] + for msg in messages: + content = msg.get("content") + if is_all_text_parts(content): + text = content_to_text(content) + if text: + flattened.append({**msg, "content": text}) + continue + flattened.append(msg) + return flattened + + +def _restore_content_shapes( + originals: list[dict[str, object]], returned: list[dict[str, object]] +) -> list[dict[str, object]]: + """Write compressed text back into each original row's content shape. + + Rows are matched positionally; the pairing is only trusted when the + service kept the row count and every role lines up. If it restructured + the conversation (e.g. dropped rows), its output is adopted as-is, which + is the pre-flattening behavior. + """ + if len(returned) != len(originals): + return returned + for orig, ret in zip(originals, returned): + if orig.get("role") != ret.get("role"): + return returned + restored: list[dict[str, object]] = [] + for orig, ret in zip(originals, returned): + orig_content = orig.get("content") + ret_content = ret.get("content") + if isinstance(orig_content, list) and isinstance(ret_content, str): + if ret_content == content_to_text(orig_content): + # Untouched row: keep the exact original parts, including + # per-part fields like cache_control on later text parts. + restored.append({**ret, "content": orig_content}) + else: + restored.append({**ret, "content": merge_rewritten_text_parts(orig_content, ret_content)}) + else: + restored.append(ret) + return restored + + def extract_hashes_from_messages(messages: list[dict[str, object]]) -> list[str]: hashes: list[str] = [] for msg in messages: @@ -491,10 +550,11 @@ class HeadroomGuardrail(CustomGuardrail): model = self.headroom_model or request_data.get("model") start_time = time.time() compressed, compression_succeeded, stats = await self._call_compress( - messages=messages, + messages=_flatten_messages_for_compression(messages), model=model if isinstance(model, str) else None, ) end_time = time.time() + compressed = _restore_content_shapes(originals=messages, returned=compressed) from litellm.proxy.common_utils.callback_utils import ( add_guardrail_to_applied_guardrails_header, diff --git a/litellm/proxy/guardrails/usage_endpoints.py b/litellm/proxy/guardrails/usage_endpoints.py index f56b22ddd49..603d3b096d3 100644 --- a/litellm/proxy/guardrails/usage_endpoints.py +++ b/litellm/proxy/guardrails/usage_endpoints.py @@ -4,11 +4,13 @@ GET /guardrails/usage/overview, /guardrails/usage/detail/:id, /guardrails/usage/ """ import json +from collections.abc import Mapping, Sequence from datetime import datetime, timedelta, timezone -from typing import Any, Dict, List, Optional +from typing import TYPE_CHECKING, Any, Literal, Union, overload from fastapi import APIRouter, Depends, Query from pydantic import BaseModel +from typing_extensions import NotRequired, TypedDict from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth @@ -21,12 +23,65 @@ from litellm.repositories.table_repositories import ( SpendLogsRepository, ) +if TYPE_CHECKING: + from prisma import models as prisma_models + from prisma import types as prisma_types + from prisma.actions import LiteLLM_GuardrailsTableActions, LiteLLM_PolicyTableActions + + from litellm.proxy.utils import PrismaClient + from litellm.types.guardrails import Guardrail + + _DbOrConfigGuardrail = Union[prisma_models.LiteLLM_GuardrailsTable, Guardrail] + _DailyMetricsRow = Union[prisma_models.LiteLLM_DailyGuardrailMetrics, prisma_models.LiteLLM_DailyPolicyMetrics] + router = APIRouter() +def _guardrails_table( + prisma_client: "PrismaClient", +) -> "LiteLLM_GuardrailsTableActions[prisma_models.LiteLLM_GuardrailsTable]": + guardrails_table: LiteLLM_GuardrailsTableActions[prisma_models.LiteLLM_GuardrailsTable] = GuardrailsRepository( + prisma_client + ).table + return guardrails_table + + +def _policies_table( + prisma_client: "PrismaClient", +) -> "LiteLLM_PolicyTableActions[prisma_models.LiteLLM_PolicyTable]": + policies_table: LiteLLM_PolicyTableActions[prisma_models.LiteLLM_PolicyTable] = PolicyRepository( + prisma_client + ).table + return policies_table + + # --- Response models --- +class UsageChartPoint(TypedDict): + date: str + passed: int + blocked: int + score: NotRequired[float | None] + + +class _MetricTotals(TypedDict): + requests: int + passed: int + blocked: int + flagged: int + + +class _PrevPeriodCounts(TypedDict): + req: int + blocked: int + + +class _DailyPassBlocked(TypedDict): + passed: int + blocked: int + + class UsageOverviewRow(BaseModel): id: str name: str @@ -34,15 +89,15 @@ class UsageOverviewRow(BaseModel): provider: str requestsEvaluated: int failRate: float - avgScore: Optional[float] - avgLatency: Optional[float] + avgScore: float | None + avgLatency: float | None status: str # healthy | warning | critical trend: str # up | down | stable class UsageOverviewResponse(BaseModel): - rows: List[UsageOverviewRow] - chart: List[Dict[str, Any]] # [{ date, passed, blocked }] + rows: list[UsageOverviewRow] + chart: list[UsageChartPoint] # [{ date, passed, blocked }] totalRequests: int totalBlocked: int passRate: float @@ -55,28 +110,28 @@ class UsageDetailResponse(BaseModel): provider: str requestsEvaluated: int failRate: float - avgScore: Optional[float] - avgLatency: Optional[float] + avgScore: float | None + avgLatency: float | None status: str trend: str - description: Optional[str] - time_series: List[Dict[str, Any]] + description: str | None + time_series: list[UsageChartPoint] class UsageLogEntry(BaseModel): id: str timestamp: str action: str # blocked | passed | flagged - score: Optional[float] - latency_ms: Optional[float] - model: Optional[str] - input_snippet: Optional[str] - output_snippet: Optional[str] - reason: Optional[str] + score: float | None + latency_ms: float | None + model: str | None + input_snippet: str | None + output_snippet: str | None + reason: str | None class UsageLogsResponse(BaseModel): - logs: List[UsageLogEntry] + logs: list[UsageLogEntry] total: int page: int page_size: int @@ -101,10 +156,10 @@ def _trend_from_comparison(current_fail: float, previous_fail: float) -> str: return "stable" -def _aggregate_daily_metrics(metrics: Any, id_attr: str) -> Dict[str, Dict[str, Any]]: - agg: Dict[str, Dict[str, Any]] = {} +def _aggregate_daily_metrics(metrics: "Sequence[_DailyMetricsRow]", id_attr: str) -> Mapping[str, _MetricTotals]: + agg: dict[str, _MetricTotals] = {} for m in metrics: - gid = getattr(m, id_attr) + gid: str = getattr(m, id_attr) if gid not in agg: agg[gid] = {"requests": 0, "passed": 0, "blocked": 0, "flagged": 0} agg[gid]["requests"] += int(m.requests_evaluated or 0) @@ -114,10 +169,10 @@ def _aggregate_daily_metrics(metrics: Any, id_attr: str) -> Dict[str, Dict[str, return agg -def _prev_fail_rates(metrics_prev: Any, id_attr: str) -> Dict[str, float]: - prev_agg_raw: Dict[str, Dict[str, int]] = {} +def _prev_fail_rates(metrics_prev: "Sequence[_DailyMetricsRow]", id_attr: str) -> Mapping[str, float]: + prev_agg_raw: dict[str, _PrevPeriodCounts] = {} for m in metrics_prev: - gid = getattr(m, id_attr) + gid: str = getattr(m, id_attr) r, b = int(m.requests_evaluated or 0), int(m.blocked_count or 0) if gid not in prev_agg_raw: prev_agg_raw[gid] = {"req": 0, "blocked": 0} @@ -126,8 +181,8 @@ def _prev_fail_rates(metrics_prev: Any, id_attr: str) -> Dict[str, float]: return {gid: (100.0 * v["blocked"] / v["req"]) if v["req"] else 0.0 for gid, v in prev_agg_raw.items()} -def _chart_from_metrics(metrics: Any) -> List[Dict[str, Any]]: - chart_by_date: Dict[str, Dict[str, int]] = {} +def _chart_from_metrics(metrics: "Sequence[_DailyMetricsRow]") -> list[UsageChartPoint]: + chart_by_date: dict[str, _DailyPassBlocked] = {} for m in metrics: d = m.date if d not in chart_by_date: @@ -137,14 +192,26 @@ def _chart_from_metrics(metrics: Any) -> List[Dict[str, Any]]: return [{"date": d, "passed": v["passed"], "blocked": v["blocked"]} for d, v in sorted(chart_by_date.items())] -def _get_guardrail_field(g: Any, field: str) -> Any: +_GuardrailStrField = Literal["guardrail_id", "guardrail_name"] +_GuardrailObjectField = Literal["litellm_params", "guardrail_info"] + + +@overload +def _get_guardrail_field(g: "_DbOrConfigGuardrail", field: _GuardrailStrField) -> str | None: ... + + +@overload +def _get_guardrail_field(g: "_DbOrConfigGuardrail", field: _GuardrailObjectField) -> object: ... + + +def _get_guardrail_field(g: "_DbOrConfigGuardrail", field: _GuardrailStrField | _GuardrailObjectField) -> object: """Read `field` off a guardrail whether it's a Prisma row (attr) or a dict/TypedDict (key).""" if isinstance(g, dict): return g.get(field) return getattr(g, field, None) -def _to_dict(value: Any) -> Dict[str, Any]: +def _to_dict(value: object) -> dict[str, Any]: """Coerce a pydantic model (e.g. LitellmParams) / dict value into a plain dict.""" if isinstance(value, BaseModel): return value.model_dump(exclude_none=True) @@ -153,7 +220,7 @@ def _to_dict(value: Any) -> Dict[str, Any]: return {} -def _get_guardrail_attrs(g: Any) -> tuple[Any, str]: +def _get_guardrail_attrs(g: "_DbOrConfigGuardrail") -> tuple[Any, str]: """Get (guardrail_id, display_name) from guardrail - handles Prisma model or dict.""" gid = _get_guardrail_field(g, "guardrail_id") name = _get_guardrail_field(g, "guardrail_name") @@ -161,18 +228,18 @@ def _get_guardrail_attrs(g: Any) -> tuple[Any, str]: def _guardrail_overview_rows( - guardrails: Any, - agg: Dict[str, Dict[str, Any]], - prev_agg: Dict[str, float], -) -> List[UsageOverviewRow]: - rows: List[UsageOverviewRow] = [] - covered_keys: set = set() + guardrails: "Sequence[_DbOrConfigGuardrail]", + agg: Mapping[str, _MetricTotals], + prev_agg: Mapping[str, float], +) -> list[UsageOverviewRow]: + rows: list[UsageOverviewRow] = [] + covered_keys: set[str] = set() for g in guardrails: gid, display_name = _get_guardrail_attrs(g) # Metrics are keyed by logical name from spend log metadata; guardrails table uses UUID - lookup_keys = [k for k in (display_name, gid) if k] + lookup_keys: Sequence[str] = [k for k in (display_name, gid) if k] covered_keys.update(lookup_keys) - a = {"requests": 0, "passed": 0, "blocked": 0, "flagged": 0} + a: _MetricTotals = {"requests": 0, "passed": 0, "blocked": 0, "flagged": 0} for k in lookup_keys: if k in agg: a = agg[k] @@ -229,11 +296,11 @@ def _guardrail_overview_rows( def _policy_overview_rows( - policies: Any, - agg: Dict[str, Dict[str, Any]], - prev_agg: Dict[str, float], -) -> List[UsageOverviewRow]: - rows: List[UsageOverviewRow] = [] + policies: "Sequence[prisma_models.LiteLLM_PolicyTable]", + agg: Mapping[str, _MetricTotals], + prev_agg: Mapping[str, float], +) -> list[UsageOverviewRow]: + rows: list[UsageOverviewRow] = [] for p in policies: pid = p.policy_id a = agg.get(pid, {"requests": 0, "passed": 0, "blocked": 0, "flagged": 0}) @@ -264,8 +331,8 @@ def _policy_overview_rows( response_model=UsageOverviewResponse, ) async def guardrails_usage_overview( - start_date: Optional[str] = Query(None, description="YYYY-MM-DD"), - end_date: Optional[str] = Query(None, description="YYYY-MM-DD"), + start_date: str | None = Query(None, description="YYYY-MM-DD"), + end_date: str | None = Query(None, description="YYYY-MM-DD"), user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """Return guardrail performance overview for the dashboard.""" @@ -281,23 +348,23 @@ async def guardrails_usage_overview( from litellm.proxy.guardrails.guardrail_registry import IN_MEMORY_GUARDRAIL_HANDLER try: - db_guardrails = await GuardrailsRepository(prisma_client).table.find_many() + db_guardrails = await _guardrails_table(prisma_client).find_many() seen_ids = {gid for g in db_guardrails if (gid := _get_guardrail_field(g, "guardrail_id")) is not None} config_guardrails = [ g for g in IN_MEMORY_GUARDRAIL_HANDLER.list_config_guardrails() if g.get("guardrail_id") not in seen_ids ] - guardrails: List[Any] = [*db_guardrails, *config_guardrails] + guardrails: Sequence[_DbOrConfigGuardrail] = [*db_guardrails, *config_guardrails] # Daily metrics in range - metrics = await DailyGuardrailMetricsRepository(prisma_client).table.find_many( - where={"date": {"gte": start, "lte": end}} - ) + metrics: Sequence[prisma_models.LiteLLM_DailyGuardrailMetrics] = await DailyGuardrailMetricsRepository( + prisma_client + ).table.find_many(where={"date": {"gte": start, "lte": end}}) # Previous period for trend start_prev = (datetime.strptime(start, "%Y-%m-%d") - timedelta(days=7)).strftime("%Y-%m-%d") - metrics_prev = await DailyGuardrailMetricsRepository(prisma_client).table.find_many( - where={"date": {"gte": start_prev, "lt": start}} - ) + metrics_prev: Sequence[prisma_models.LiteLLM_DailyGuardrailMetrics] = await DailyGuardrailMetricsRepository( + prisma_client + ).table.find_many(where={"date": {"gte": start_prev, "lt": start}}) agg = _aggregate_daily_metrics(metrics, "guardrail_id") prev_agg = _prev_fail_rates(metrics_prev, "guardrail_id") @@ -327,8 +394,8 @@ async def guardrails_usage_overview( ) async def guardrails_usage_detail( guardrail_id: str, - start_date: Optional[str] = Query(None), - end_date: Optional[str] = Query(None), + start_date: str | None = Query(None), + end_date: str | None = Query(None), user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """Return single guardrail usage metrics and time series.""" @@ -345,7 +412,7 @@ async def guardrails_usage_detail( from litellm.proxy.guardrails.guardrail_registry import IN_MEMORY_GUARDRAIL_HANDLER - guardrail: Any = await GuardrailsRepository(prisma_client).table.find_unique(where={"guardrail_id": guardrail_id}) + guardrail = await _guardrails_table(prisma_client).find_unique(where={"guardrail_id": guardrail_id}) if guardrail is None: guardrail = IN_MEMORY_GUARDRAIL_HANDLER.get_config_guardrail_by_id(guardrail_id=guardrail_id) if guardrail is None: @@ -357,13 +424,17 @@ async def guardrails_usage_detail( logical_id = _get_guardrail_field(guardrail, "guardrail_name") metric_ids = [i for i in (logical_id, guardrail_id) if i] - metrics = await DailyGuardrailMetricsRepository(prisma_client).table.find_many( + metrics: Sequence[prisma_models.LiteLLM_DailyGuardrailMetrics] = await DailyGuardrailMetricsRepository( + prisma_client + ).table.find_many( where={ "guardrail_id": {"in": metric_ids}, "date": {"gte": start, "lte": end}, } ) - metrics_prev = await DailyGuardrailMetricsRepository(prisma_client).table.find_many( + metrics_prev: Sequence[prisma_models.LiteLLM_DailyGuardrailMetrics] = await DailyGuardrailMetricsRepository( + prisma_client + ).table.find_many( where={ "guardrail_id": {"in": metric_ids}, "date": {"lt": start}, @@ -380,14 +451,14 @@ async def guardrails_usage_detail( trend = _trend_from_comparison(fail_rate, prev_fail) # Aggregate by date in case metrics exist under both UUID and logical name - ts_by_date: Dict[str, Dict[str, Any]] = {} + ts_by_date: dict[str, _DailyPassBlocked] = {} for m in metrics: d = m.date if d not in ts_by_date: ts_by_date[d] = {"passed": 0, "blocked": 0} ts_by_date[d]["passed"] += int(m.passed_count or 0) ts_by_date[d]["blocked"] += int(m.blocked_count or 0) - time_series = [ + time_series: list[UsageChartPoint] = [ {"date": d, "passed": v["passed"], "blocked": v["blocked"], "score": None} for d, v in sorted(ts_by_date.items()) ] @@ -412,18 +483,18 @@ async def guardrails_usage_detail( def _build_usage_logs_where( - guardrail_ids: Optional[List[str]], - policy_id: Optional[str], - start_date: Optional[str], - end_date: Optional[str], -) -> Dict[str, Any]: - where: Dict[str, Any] = {} + guardrail_ids: list[str] | None, + policy_id: str | None, + start_date: str | None, + end_date: str | None, +) -> "prisma_types.LiteLLM_SpendLogGuardrailIndexWhereInput": + where: prisma_types.LiteLLM_SpendLogGuardrailIndexWhereInput = {} if guardrail_ids: where["guardrail_id"] = {"in": guardrail_ids} if len(guardrail_ids) > 1 else guardrail_ids[0] if policy_id: where["policy_id"] = policy_id if start_date or end_date: - st_filter: Dict[str, Any] = {} + st_filter: prisma_types.DateTimeFilter = {} if start_date: sd = start_date.replace("Z", "+00:00").strip() if "T" not in sd: @@ -438,7 +509,9 @@ def _build_usage_logs_where( return where -def _usage_log_entry_from_row(r: Any, sl: Any, action_filter: Optional[str]) -> Optional[UsageLogEntry]: +def _usage_log_entry_from_row( + r: "prisma_models.LiteLLM_SpendLogGuardrailIndex", sl: Any, action_filter: str | None +) -> UsageLogEntry | None: meta = sl.metadata if isinstance(meta, str): try: @@ -488,7 +561,7 @@ def _usage_log_entry_from_row(r: Any, sl: Any, action_filter: Optional[str]) -> ) -def _snippet(text: Any, max_len: int = 200) -> Optional[str]: +def _snippet(text: Any, max_len: int = 200) -> str | None: if text is None: return None if isinstance(text, str): @@ -510,7 +583,7 @@ def _snippet(text: Any, max_len: int = 200) -> Optional[str]: return result -def _input_snippet_for_log(sl: Any) -> Optional[str]: +def _input_snippet_for_log(sl: "prisma_models.LiteLLM_SpendLogs") -> str | None: """Snippet for request input: prefer messages, fall back to proxy_server_request (same as drawer).""" out = _snippet(sl.messages) if out: @@ -541,13 +614,13 @@ def _input_snippet_for_log(sl: Any) -> Optional[str]: response_model=UsageLogsResponse, ) async def guardrails_usage_logs( - guardrail_id: Optional[str] = Query(None), - policy_id: Optional[str] = Query(None), + guardrail_id: str | None = Query(None), + policy_id: str | None = Query(None), page: int = Query(1, ge=1), page_size: int = Query(50, ge=1, le=100), - action: Optional[str] = Query(None), - start_date: Optional[str] = Query(None), - end_date: Optional[str] = Query(None), + action: str | None = Query(None), + start_date: str | None = Query(None), + end_date: str | None = Query(None), user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """Return paginated run logs for a guardrail (or policy) from SpendLogs via index.""" @@ -562,13 +635,11 @@ async def guardrails_usage_logs( try: # Index rows may store either guardrail_id (UUID) or guardrail_name from metadata. # Query by both so we match regardless of which was written. - effective_guardrail_ids: List[str] = [guardrail_id] if guardrail_id else [] + effective_guardrail_ids: list[str] = [guardrail_id] if guardrail_id else [] if guardrail_id: from litellm.proxy.guardrails.guardrail_registry import IN_MEMORY_GUARDRAIL_HANDLER - guardrail: Any = await GuardrailsRepository(prisma_client).table.find_unique( - where={"guardrail_id": guardrail_id} - ) + guardrail = await _guardrails_table(prisma_client).find_unique(where={"guardrail_id": guardrail_id}) if guardrail is None: guardrail = IN_MEMORY_GUARDRAIL_HANDLER.get_config_guardrail_by_id(guardrail_id=guardrail_id) if guardrail: @@ -577,19 +648,23 @@ async def guardrails_usage_logs( effective_guardrail_ids.append(logical_name) where = _build_usage_logs_where(effective_guardrail_ids or None, policy_id, start_date, end_date) - index_rows = await SpendLogGuardrailIndexRepository(prisma_client).table.find_many( + index_rows: Sequence[prisma_models.LiteLLM_SpendLogGuardrailIndex] = await SpendLogGuardrailIndexRepository( + prisma_client + ).table.find_many( where=where, order={"start_time": "desc"}, skip=(page - 1) * page_size, take=page_size + 1, ) - total = await SpendLogGuardrailIndexRepository(prisma_client).table.count(where=where) + total: int = await SpendLogGuardrailIndexRepository(prisma_client).table.count(where=where) request_ids = [r.request_id for r in index_rows[:page_size]] if not request_ids: return UsageLogsResponse(logs=[], total=total, page=page, page_size=page_size) - spend_logs = await SpendLogsRepository(prisma_client).table.find_many(where={"request_id": {"in": request_ids}}) + spend_logs: Sequence[prisma_models.LiteLLM_SpendLogs] = await SpendLogsRepository( + prisma_client + ).table.find_many(where={"request_id": {"in": request_ids}}) log_by_id = {s.request_id: s for s in spend_logs} - logs_out: List[UsageLogEntry] = [] + logs_out: list[UsageLogEntry] = [] for r in index_rows[:page_size]: sl = log_by_id.get(r.request_id) if not sl: @@ -614,8 +689,8 @@ async def guardrails_usage_logs( response_model=UsageOverviewResponse, ) async def policies_usage_overview( - start_date: Optional[str] = Query(None, description="YYYY-MM-DD"), - end_date: Optional[str] = Query(None, description="YYYY-MM-DD"), + start_date: str | None = Query(None, description="YYYY-MM-DD"), + end_date: str | None = Query(None, description="YYYY-MM-DD"), user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """Return policy performance overview for the dashboard.""" @@ -629,11 +704,13 @@ async def policies_usage_overview( start = start_date or (now - timedelta(days=7)).strftime("%Y-%m-%d") try: - policies = await PolicyRepository(prisma_client).table.find_many() - metrics = await DailyPolicyMetricsRepository(prisma_client).table.find_many( - where={"date": {"gte": start, "lte": end}} - ) - metrics_prev = await DailyPolicyMetricsRepository(prisma_client).table.find_many( + policies = await _policies_table(prisma_client).find_many() + metrics: Sequence[prisma_models.LiteLLM_DailyPolicyMetrics] = await DailyPolicyMetricsRepository( + prisma_client + ).table.find_many(where={"date": {"gte": start, "lte": end}}) + metrics_prev: Sequence[prisma_models.LiteLLM_DailyPolicyMetrics] = await DailyPolicyMetricsRepository( + prisma_client + ).table.find_many( where={ "date": { "gte": (datetime.strptime(start, "%Y-%m-%d") - timedelta(days=7)).strftime("%Y-%m-%d"), diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 31e4e1694b0..d6d06d2e745 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -35,6 +35,7 @@ from litellm.proxy._types import ( from litellm.proxy.common_utils.callback_utils import ( decrypt_callback_vars, get_metadata_variable_name_from_kwargs, + strip_callback_config, ) from litellm.proxy.common_utils.http_parsing_utils import _safe_get_request_headers @@ -1274,7 +1275,7 @@ class LiteLLMProxyRequestSetup: user_api_key_budget_reset_at=( user_api_key_dict.budget_reset_at.isoformat() if user_api_key_dict.budget_reset_at else None ), - user_api_key_auth_metadata=user_api_key_dict.metadata, + user_api_key_auth_metadata=strip_callback_config(user_api_key_dict.metadata), ) return user_api_key_logged_metadata @@ -1912,8 +1913,8 @@ async def add_litellm_data_to_request( data[_metadata_variable_name]["user_api_key_user_spend"] = user_api_key_dict.user_spend data[_metadata_variable_name]["user_api_key_user_max_budget"] = user_api_key_dict.user_max_budget - data[_metadata_variable_name]["user_api_key_metadata"] = user_api_key_dict.metadata - data[_metadata_variable_name]["user_api_key_team_metadata"] = user_api_key_dict.team_metadata + data[_metadata_variable_name]["user_api_key_metadata"] = strip_callback_config(user_api_key_dict.metadata) + data[_metadata_variable_name]["user_api_key_team_metadata"] = strip_callback_config(user_api_key_dict.team_metadata) data[_metadata_variable_name]["user_api_key_object_permission_id"] = getattr( user_api_key_dict, "object_permission_id", None ) diff --git a/litellm/proxy/management_endpoints/common_daily_activity.py b/litellm/proxy/management_endpoints/common_daily_activity.py index a5ecf4e7f93..9b756d14815 100644 --- a/litellm/proxy/management_endpoints/common_daily_activity.py +++ b/litellm/proxy/management_endpoints/common_daily_activity.py @@ -1,9 +1,15 @@ import asyncio +from collections.abc import Awaitable, Callable, Mapping, Sequence from datetime import datetime from types import SimpleNamespace -from typing import Any, Awaitable, Callable, Dict, List, Optional, Set, Tuple, Union +from typing import ( + TYPE_CHECKING, + Protocol, + Union, +) from fastapi import HTTPException, status +from typing_extensions import TypedDict from litellm._logging import verbose_proxy_logger from litellm.proxy._types import CommonProxyErrors @@ -16,6 +22,7 @@ from litellm.types.proxy.management_endpoints.common_daily_activity import ( BreakdownMetrics, DailySpendData, DailySpendMetadata, + GroupedData, KeyMetadata, KeyMetricWithMetadata, MetricWithMetadata, @@ -23,8 +30,16 @@ from litellm.types.proxy.management_endpoints.common_daily_activity import ( SpendMetrics, ) +if TYPE_CHECKING: + from prisma.models import ( + LiteLLM_DeletedVerificationToken as PrismaDeletedVerificationToken, + ) + from prisma.models import ( + LiteLLM_VerificationToken as PrismaVerificationToken, + ) + # Mapping from Prisma accessor names to actual PostgreSQL table names. -_PRISMA_TO_PG_TABLE: Dict[str, str] = { +_PRISMA_TO_PG_TABLE: Mapping[str, str] = { "litellm_dailyuserspend": "LiteLLM_DailyUserSpend", "litellm_dailyteamspend": "LiteLLM_DailyTeamSpend", "litellm_dailyorganizationspend": "LiteLLM_DailyOrganizationSpend", @@ -34,7 +49,98 @@ _PRISMA_TO_PG_TABLE: Dict[str, str] = { } -def update_metrics(existing_metrics: SpendMetrics, record: Any) -> SpendMetrics: +class DailySpendRecord(Protocol): + @property + def date(self) -> str: ... + + @property + def api_key(self) -> str: ... + + @property + def model(self) -> str | None: ... + + @property + def model_group(self) -> str | None: ... + + @property + def custom_llm_provider(self) -> str | None: ... + + @property + def mcp_namespaced_tool_name(self) -> str | None: ... + + @property + def endpoint(self) -> str | None: ... + + @property + def prompt_tokens(self) -> int: ... + + @property + def completion_tokens(self) -> int: ... + + @property + def spend(self) -> float: ... + + @property + def cache_read_input_tokens(self) -> int: ... + + @property + def cache_creation_input_tokens(self) -> int: ... + + @property + def compression_saved_tokens(self) -> int: ... + + @property + def compression_savings_spend(self) -> float: ... + + @property + def prompt_caching_savings_spend(self) -> float: ... + + @property + def api_requests(self) -> int: ... + + @property + def successful_requests(self) -> int: ... + + @property + def failed_requests(self) -> int: ... + + +class _KeyMetadataDict(TypedDict, total=False): + key_alias: str | None + team_id: str | None + + +_WhereValue = Union[str, dict[str, object]] + + +class _AggregatedSpendData(TypedDict): + results: list[DailySpendData] + totals: SpendMetrics + + +class _GroupingSetsRow(SimpleNamespace): + date: str + api_key: str | None + model: str | None + model_group: str | None + custom_llm_provider: str | None + mcp_namespaced_tool_name: str | None + endpoint: str | None + group_level: int + spend: float | None + prompt_tokens: int | None + completion_tokens: int | None + cache_read_input_tokens: int | None + cache_creation_input_tokens: int | None + compression_saved_tokens: int | None + compression_savings_spend: float | None + prompt_caching_savings_spend: float | None + api_requests: int | None + successful_requests: int | None + failed_requests: int | None + + +def update_metrics(existing_metrics: SpendMetrics, record: DailySpendRecord) -> SpendMetrics: """Update metrics with new record data. Rollup rows can carry None for numeric fields when SUM() spans zero rows @@ -58,7 +164,7 @@ def update_metrics(existing_metrics: SpendMetrics, record: Any) -> SpendMetrics: return existing_metrics -def _is_user_agent_tag(tag: Optional[str]) -> bool: +def _is_user_agent_tag(tag: str | None) -> bool: """Determine whether a tag should be treated as a User-Agent tag.""" if not tag: return False @@ -66,15 +172,15 @@ def _is_user_agent_tag(tag: Optional[str]) -> bool: return normalized_tag.startswith("user-agent:") or normalized_tag.startswith("user agent:") -def compute_tag_metadata_totals(records: List[Any]) -> SpendMetrics: +def compute_tag_metadata_totals(records: Sequence[DailySpendRecord]) -> SpendMetrics: """ Deduplicate spend metrics for tags using request_id, ignoring User-Agent prefixed tags. Each unique request_id contributes at most one record (the tag with max spend) to metadata. """ - deduped_records: Dict[str, Any] = {} + deduped_records: dict[str, DailySpendRecord] = {} for record in records: - request_id = getattr(record, "request_id", None) + request_id: str | None = getattr(record, "request_id", None) if not request_id: continue @@ -94,12 +200,12 @@ def compute_tag_metadata_totals(records: List[Any]) -> SpendMetrics: def update_breakdown_metrics( breakdown: BreakdownMetrics, - record: Any, - model_metadata: Dict[str, Dict[str, Any]], - provider_metadata: Dict[str, Dict[str, Any]], - api_key_metadata: Dict[str, Dict[str, Any]], - entity_id_field: Optional[str] = None, - entity_metadata_field: Optional[Dict[str, dict]] = None, + record: DailySpendRecord, + model_metadata: Mapping[str, dict[str, object]], + provider_metadata: Mapping[str, dict[str, object]], + api_key_metadata: Mapping[str, _KeyMetadataDict], + entity_id_field: str | None = None, + entity_metadata_field: Mapping[str, dict[str, object]] | None = None, ) -> BreakdownMetrics: """Updates breakdown metrics for a single record using the existing update_metrics function""" @@ -269,23 +375,27 @@ def update_breakdown_metrics( async def get_api_key_metadata( prisma_client: PrismaClient, - api_keys: Set[str], -) -> Dict[str, Dict[str, Any]]: + api_keys: set[str], +) -> dict[str, _KeyMetadataDict]: """Get api key metadata, falling back to deleted keys table for keys not found in active table. This ensures that key_alias and team_id are preserved in historical activity logs even after a key is deleted or regenerated. """ - key_records = await VerificationTokenRepository(prisma_client).table.find_many( + key_records: list[PrismaVerificationToken] = await VerificationTokenRepository(prisma_client).table.find_many( where={"token": {"in": list(api_keys)}} ) - result = {k.token: {"key_alias": k.key_alias, "team_id": k.team_id} for k in key_records} + result: dict[str, _KeyMetadataDict] = { + k.token: {"key_alias": k.key_alias, "team_id": k.team_id} for k in key_records + } # For any keys not found in the active table, check the deleted keys table missing_keys = api_keys - set(result.keys()) if missing_keys: try: - deleted_key_records = await DeletedVerificationTokenRepository(prisma_client).table.find_many( + deleted_key_records: list[PrismaDeletedVerificationToken] = await DeletedVerificationTokenRepository( + prisma_client + ).table.find_many( where={"token": {"in": list(missing_keys)}}, order={"deleted_at": "desc"}, ) @@ -309,8 +419,8 @@ async def get_api_key_metadata( def _adjust_dates_for_timezone( start_date: str, end_date: str, - timezone_offset_minutes: Optional[int], -) -> Tuple[str, str]: + timezone_offset_minutes: int | None, +) -> tuple[str, str]: """ Pass-through for the local date range; the timezone offset is intentionally ignored here. @@ -335,19 +445,19 @@ def _adjust_dates_for_timezone( def _build_where_conditions( *, entity_id_field: str, - entity_id: Optional[Union[str, List[str]]], + entity_id: str | list[str] | None, start_date: str, end_date: str, - model: Optional[str], - api_key: Optional[Union[str, List[str]]], - exclude_entity_ids: Optional[List[str]] = None, - timezone_offset_minutes: Optional[int] = None, -) -> Dict[str, Any]: + model: str | None, + api_key: str | list[str] | None, + exclude_entity_ids: list[str] | None = None, + timezone_offset_minutes: int | None = None, +) -> dict[str, "_WhereValue"]: """Build prisma where clause for daily activity queries.""" # Adjust dates for timezone if provided adjusted_start, adjusted_end = _adjust_dates_for_timezone(start_date, end_date, timezone_offset_minutes) - where_conditions: Dict[str, Any] = { + where_conditions: dict[str, _WhereValue] = { "date": { "gte": adjusted_start, "lte": adjusted_end, @@ -369,7 +479,7 @@ def _build_where_conditions( where_conditions[entity_id_field] = {"equals": entity_id} if exclude_entity_ids: - current = where_conditions.get(entity_id_field, {}) + current: _WhereValue = where_conditions.get(entity_id_field, {}) if isinstance(current, str): current = {"equals": current} current["not"] = {"in": exclude_entity_ids} @@ -382,14 +492,14 @@ def _build_aggregated_sql_query( *, table_name: str, entity_id_field: str, - entity_id: Optional[Union[str, List[str]]], + entity_id: str | list[str] | None, start_date: str, end_date: str, - model: Optional[str], - api_key: Optional[str], - exclude_entity_ids: Optional[List[str]] = None, - timezone_offset_minutes: Optional[int] = None, -) -> Tuple[str, List[Any]]: + model: str | None, + api_key: str | None, + exclude_entity_ids: list[str] | None = None, + timezone_offset_minutes: int | None = None, +) -> tuple[str, list[str]]: """Build a parameterized SQL GROUP BY query for aggregated daily activity. Groups by (date, api_key, model, model_group, custom_llm_provider, @@ -406,8 +516,8 @@ def _build_aggregated_sql_query( adjusted_start, adjusted_end = _adjust_dates_for_timezone(start_date, end_date, timezone_offset_minutes) - sql_conditions: List[str] = [] - sql_params: List[Any] = [] + sql_conditions: list[str] = [] + sql_params: list[str] = [] p = 1 # parameter index (1-based for PostgreSQL $N placeholders) # Date range (always present) @@ -506,17 +616,17 @@ def _build_aggregated_sql_query( def _aggregate_spend_records_sync( *, - records: List[Any], - api_key_metadata: Dict[str, Dict[str, Any]], - entity_id_field: Optional[str], - entity_metadata_field: Optional[Dict[str, dict]], -) -> Dict[str, Any]: - model_metadata: Dict[str, Dict[str, Any]] = {} - provider_metadata: Dict[str, Dict[str, Any]] = {} + records: Sequence[DailySpendRecord], + api_key_metadata: Mapping[str, _KeyMetadataDict], + entity_id_field: str | None, + entity_metadata_field: Mapping[str, dict[str, object]] | None, +) -> _AggregatedSpendData: + model_metadata: dict[str, dict[str, object]] = {} + provider_metadata: dict[str, dict[str, object]] = {} - results: List[DailySpendData] = [] + results: list[DailySpendData] = [] total_metrics = SpendMetrics() - grouped_data: Dict[str, Dict[str, Any]] = {} + grouped_data: dict[str, GroupedData] = {} for record in records: date_str = record.date @@ -557,18 +667,18 @@ def _aggregate_spend_records_sync( async def _aggregate_spend_records( *, prisma_client: PrismaClient, - records: List[Any], - entity_id_field: Optional[str], - entity_metadata_field: Optional[Dict[str, dict]], -) -> Dict[str, Any]: + records: Sequence[DailySpendRecord], + entity_id_field: str | None, + entity_metadata_field: Mapping[str, dict[str, object]] | None, +) -> _AggregatedSpendData: """Aggregate rows into DailySpendData list and total metrics. The per-row loop is offloaded to a worker thread via asyncio.to_thread so a large result set doesn't peg the event loop. """ - api_keys: Set[str] = {record.api_key for record in records if record.api_key} + api_keys: set[str] = {record.api_key for record in records if record.api_key} - api_key_metadata: Dict[str, Dict[str, Any]] = {} + api_key_metadata: dict[str, _KeyMetadataDict] = {} if api_keys: api_key_metadata = await get_api_key_metadata(prisma_client, api_keys) @@ -603,7 +713,7 @@ _GROUP_DATE_ENDPOINT = 62 # 0b0111110 _GROUP_DATE_ENDPOINT_API_KEY = 30 # 0b0011110 -def _record_to_spend_metrics(record: Any) -> SpendMetrics: +def _record_to_spend_metrics(record: _GroupingSetsRow) -> SpendMetrics: """Build a SpendMetrics directly from one already-aggregated rollup row. SUM() over zero rows is SQL NULL, so rollup rows (notably the grand-total @@ -627,16 +737,16 @@ def _record_to_spend_metrics(record: Any) -> SpendMetrics: ) -def _key_metadata(api_key_metadata: Dict[str, Dict[str, Any]], api_key: str) -> KeyMetadata: +def _key_metadata(api_key_metadata: Mapping[str, _KeyMetadataDict], api_key: str) -> KeyMetadata: meta = api_key_metadata.get(api_key, {}) return KeyMetadata(key_alias=meta.get("key_alias"), team_id=meta.get("team_id")) def _aggregate_grouping_sets_records_sync( *, - records: List[Any], - api_key_metadata: Dict[str, Dict[str, Any]], -) -> Dict[str, Any]: + records: Sequence[_GroupingSetsRow], + api_key_metadata: Mapping[str, _KeyMetadataDict], +) -> _AggregatedSpendData: """Build the response from rollup rows produced by the GROUPING SETS query. Each row carries a `group_level` bitmask (from Postgres GROUPING()) that @@ -645,16 +755,16 @@ def _aggregate_grouping_sets_records_sync( summing in Python and no nested update_metrics calls. """ total_metrics = SpendMetrics() - grouped_data: Dict[str, Dict[str, Any]] = {} + grouped_data: dict[str, GroupedData] = {} - def ensure_date(date_str: str) -> Dict[str, Any]: - bucket = grouped_data.get(date_str) + def ensure_date(date_str: str) -> GroupedData: + bucket: GroupedData | None = grouped_data.get(date_str) if bucket is None: bucket = {"metrics": SpendMetrics(), "breakdown": BreakdownMetrics()} grouped_data[date_str] = bucket return bucket - def assign_metric_with_metadata(target: Dict[str, MetricWithMetadata], key: str, metrics: SpendMetrics) -> None: + def assign_metric_with_metadata(target: dict[str, MetricWithMetadata], key: str, metrics: SpendMetrics) -> None: existing = target.get(key) if existing is None: target[key] = MetricWithMetadata(metrics=metrics, metadata={}) @@ -662,7 +772,7 @@ def _aggregate_grouping_sets_records_sync( existing.metrics = metrics def assign_api_key_breakdown( - target: Dict[str, MetricWithMetadata], + target: dict[str, MetricWithMetadata], parent_key: str, api_key: str, metrics: SpendMetrics, @@ -753,12 +863,12 @@ def _aggregate_grouping_sets_records_sync( async def _aggregate_grouping_sets_records( *, prisma_client: PrismaClient, - records: List[Any], -) -> Dict[str, Any]: + records: Sequence[_GroupingSetsRow], +) -> _AggregatedSpendData: """Async wrapper: fetch api_key_metadata, then dispatch on a worker thread.""" - api_keys: Set[str] = {r.api_key for r in records if r.api_key} + api_keys: set[str] = {r.api_key for r in records if r.api_key} - api_key_metadata: Dict[str, Dict[str, Any]] = {} + api_key_metadata: dict[str, _KeyMetadataDict] = {} if api_keys: api_key_metadata = await get_api_key_metadata(prisma_client, api_keys) @@ -770,21 +880,22 @@ async def _aggregate_grouping_sets_records( async def get_daily_activity( - prisma_client: Optional[PrismaClient], + prisma_client: PrismaClient | None, table_name: str, entity_id_field: str, - entity_id: Optional[Union[str, List[str]]], - entity_metadata_field: Optional[Dict[str, dict]], - start_date: Optional[str], - end_date: Optional[str], - model: Optional[str], - api_key: Optional[Union[str, List[str]]], + entity_id: str | list[str] | None, + entity_metadata_field: Mapping[str, dict[str, object]] | None, + start_date: str | None, + end_date: str | None, + model: str | None, + api_key: str | list[str] | None, page: int, page_size: int, - exclude_entity_ids: Optional[List[str]] = None, - metadata_metrics_func: Optional[Callable[[List[Any]], SpendMetrics]] = None, - timezone_offset_minutes: Optional[int] = None, - resolve_entity_metadata: Optional[Callable[[list[Any]], Awaitable[dict[str, dict]]]] = None, + exclude_entity_ids: list[str] | None = None, + metadata_metrics_func: Callable[[Sequence[DailySpendRecord]], SpendMetrics] | None = None, + timezone_offset_minutes: int | None = None, + resolve_entity_metadata: Callable[[Sequence[DailySpendRecord]], Awaitable[dict[str, dict[str, object]]]] + | None = None, ) -> SpendAnalyticsPaginatedResponse: """Common function to get daily activity for any entity type. @@ -819,7 +930,7 @@ async def get_daily_activity( ) # Get total count for pagination - total_count = await getattr(prisma_client.db, table_name).count(where=where_conditions) + total_count: int = await getattr(prisma_client.db, table_name).count(where=where_conditions) # Fetch paginated results. # ``date`` alone is not a unique sort key -- a busy tenant has many @@ -831,7 +942,7 @@ async def get_daily_activity( # total. Adding ``id`` (the row's UUID primary key, present on both # LiteLLM_DailyUserSpend and LiteLLM_DailyTeamSpend) as a tiebreaker # gives every page a stable cursor (#30164). - daily_spend_data = await getattr(prisma_client.db, table_name).find_many( + daily_spend_data: Sequence[DailySpendRecord] = await getattr(prisma_client.db, table_name).find_many( where=where_conditions, order=[ {"date": "desc"}, @@ -889,17 +1000,17 @@ async def get_daily_activity( async def get_daily_activity_aggregated( - prisma_client: Optional[PrismaClient], + prisma_client: PrismaClient | None, table_name: str, entity_id_field: str, - entity_id: Optional[Union[str, List[str]]], - entity_metadata_field: Optional[Dict[str, dict]], - start_date: Optional[str], - end_date: Optional[str], - model: Optional[str], - api_key: Optional[str], - exclude_entity_ids: Optional[List[str]] = None, - timezone_offset_minutes: Optional[int] = None, + entity_id: str | list[str] | None, + entity_metadata_field: Mapping[str, dict[str, object]] | None, + start_date: str | None, + end_date: str | None, + model: str | None, + api_key: str | None, + exclude_entity_ids: list[str] | None = None, + timezone_offset_minutes: int | None = None, ) -> SpendAnalyticsPaginatedResponse: """Aggregated variant that returns the full result set (no pagination). @@ -939,7 +1050,7 @@ async def get_daily_activity_aggregated( if rows is None: rows = [] - records = [SimpleNamespace(**row) for row in rows] + records = [_GroupingSetsRow(**row) for row in rows] # The grouping-sets dispatcher places each row directly in its bucket # using the row's GROUPING() bitmask. No Python-side summing needed. diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index a1592d512f5..a2c16e88839 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -15,8 +15,9 @@ These are members of a Team on LiteLLM import asyncio import json import traceback +from collections.abc import Sequence from datetime import datetime, timezone -from typing import Any, Dict, List, Optional, Union, cast +from typing import Any, Optional, cast import fastapi from fastapi import APIRouter, Depends, Header, HTTPException, Request, status @@ -29,6 +30,7 @@ from litellm.proxy.auth.auth_checks import get_team_object, get_user_object from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.hooks.user_management_event_hooks import UserManagementEventHooks from litellm.proxy.management_endpoints.common_daily_activity import ( + DailySpendRecord, get_daily_activity, get_daily_activity_aggregated, ) @@ -59,17 +61,17 @@ from litellm.repositories.verification_token_repository import ( from litellm.types.proxy.management_endpoints.common_daily_activity import ( SpendAnalyticsPaginatedResponse, ) -from litellm.types.proxy.management_endpoints.scim_v2 import ( - SCIM_ENTERPRISE_METADATA_KEY, - SCIM_ENTITLEMENTS_METADATA_KEY, - SCIM_ROLES_METADATA_KEY, -) from litellm.types.proxy.management_endpoints.internal_user_endpoints import ( BulkUpdateUserRequest, BulkUpdateUserResponse, UserListResponse, UserUpdateResult, ) +from litellm.types.proxy.management_endpoints.scim_v2 import ( + SCIM_ENTERPRISE_METADATA_KEY, + SCIM_ENTITLEMENTS_METADATA_KEY, + SCIM_ROLES_METADATA_KEY, +) if TYPE_CHECKING: from litellm.proxy.proxy_server import PrismaClient @@ -127,11 +129,11 @@ def _update_internal_new_user_params(data_json: dict, data: NewUserRequest) -> d async def _check_duplicate_user_field( field_name: str, - field_value: Optional[str], + field_value: str | None, prisma_client: Any, *, case_insensitive: bool = False, - label: Optional[str] = None, + label: str | None = None, ) -> None: """ Helper function to check if a field already exists in the user table. @@ -167,7 +169,7 @@ async def _check_duplicate_user_field( ) -async def _check_duplicate_user_email(user_email: Optional[str], prisma_client: Any) -> None: +async def _check_duplicate_user_email(user_email: str | None, prisma_client: Any) -> None: """ Helper function to check if a user email already exists in the database. """ @@ -180,7 +182,7 @@ async def _check_duplicate_user_email(user_email: Optional[str], prisma_client: ) -async def _check_duplicate_user_id(user_id: Optional[str], prisma_client: Any) -> None: +async def _check_duplicate_user_id(user_id: str | None, prisma_client: Any) -> None: """ Helper function to check if a user id already exists in the database. """ @@ -194,7 +196,7 @@ async def _check_duplicate_user_id(user_id: Optional[str], prisma_client: Any) - async def _add_user_to_organizations( user_id: str, - organizations: List[str], + organizations: list[str], prisma_client: "PrismaClient", user_api_key_dict: UserAPIKeyAuth, ): @@ -231,8 +233,8 @@ async def _add_user_to_team( user_id: str, team_id: str, user_api_key_dict: UserAPIKeyAuth, - user_email: Optional[str] = None, - max_budget_in_team: Optional[float] = None, + user_email: str | None = None, + max_budget_in_team: float | None = None, user_role: Literal["user", "admin"] = "user", ): from litellm.proxy.management_endpoints.team_endpoints import team_member_add @@ -258,10 +260,12 @@ async def _add_user_to_team( ) ) else: - verbose_proxy_logger.debug( - "litellm.proxy.management_endpoints.internal_user_endpoints.new_user(): Exception occured - {}".format( - str(e) - ) + verbose_proxy_logger.error( + "litellm.proxy.management_endpoints.internal_user_endpoints._add_user_to_team(): " + "failed to add user %s to team %s - %s", + user_id, + team_id, + str(e), ) except Exception as e: if "already exists" in str(e) or "doesn't exist" in str(e): @@ -277,10 +281,17 @@ async def _add_user_to_team( ) ) else: + verbose_proxy_logger.error( + "litellm.proxy.management_endpoints.internal_user_endpoints._add_user_to_team(): " + "failed to add user %s to team %s - %s", + user_id, + team_id, + str(e), + ) raise e -def check_if_default_team_set() -> Optional[Union[List[str], List[NewUserRequestTeam]]]: +def check_if_default_team_set() -> list[str] | list[NewUserRequestTeam] | None: if litellm.default_internal_user_params is None: return None teams = litellm.default_internal_user_params.get("teams") @@ -306,9 +317,9 @@ def check_if_default_team_set() -> Optional[Union[List[str], List[NewUserRequest async def add_new_user_to_default_team( user_id: str, - user_email: Optional[str], + user_email: str | None, user_api_key_dict: UserAPIKeyAuth, - teams: Union[List[str], List[NewUserRequestTeam]], + teams: list[str] | list[NewUserRequestTeam], prisma_client: "PrismaClient", ): tasks = [] @@ -459,7 +470,7 @@ async def new_user( teams = data.teams if teams is None: teams = check_if_default_team_set() - organization_ids = cast(Optional[List[str]], data_json.pop("organizations", None)) + organization_ids = cast(list[str] | None, data_json.pop("organizations", None)) response = await generate_key_helper_fn(request_type="user", **data_json) # Admin UI Logic @@ -484,7 +495,7 @@ async def new_user( prisma_client=prisma_client, ) - user_id = cast(Optional[str], response.get("user_id", None)) + user_id = cast(str | None, response.get("user_id", None)) if organization_ids is not None and user_id is not None: await _add_user_to_organizations( @@ -560,9 +571,9 @@ async def ui_get_available_role( def get_team_from_list( - team_list: Optional[Union[List[LiteLLM_TeamTable], List[TeamListResponseObject]]], + team_list: list[LiteLLM_TeamTable] | list[TeamListResponseObject] | None, team_id: str, -) -> Optional[Union[LiteLLM_TeamTable, LiteLLM_TeamMembership]]: +) -> LiteLLM_TeamTable | LiteLLM_TeamMembership | None: if team_list is None: return None @@ -584,12 +595,12 @@ def _is_valid_user_id(user_id: str) -> bool: return True -def get_user_id_from_request(request: Request) -> Optional[str]: +def get_user_id_from_request(request: Request) -> str | None: """ Get the user id from the request """ # Get the raw query string and parse it properly to handle + characters - user_id: Optional[str] = None + user_id: str | None = None query_string = str(request.url.query) if "user_id=" in query_string: # Extract the user_id value from the raw query string @@ -605,14 +616,14 @@ def get_user_id_from_request(request: Request) -> Optional[str]: return user_id -def _normalize_user_info_user_id(request: Request, user_id: Optional[str]) -> Optional[str]: +def _normalize_user_info_user_id(request: Request, user_id: str | None) -> str | None: """Normalize URL-decoded user_id while preserving '+' characters.""" if user_id is not None and " " in user_id: return get_user_id_from_request(request=request) return user_id -def _enforce_user_info_access(user_id: Optional[str], user_api_key_dict: UserAPIKeyAuth) -> None: +def _enforce_user_info_access(user_id: str | None, user_api_key_dict: UserAPIKeyAuth) -> None: """Re-validate that the caller may read the resolved ``user_id`` after URL-decoding has been finalized. @@ -645,10 +656,10 @@ def _enforce_user_info_access(user_id: Optional[str], user_api_key_dict: UserAPI async def _get_user_info_teams( prisma_client: Any, - user_id: Optional[str], - user_info: Optional[Any], + user_id: str | None, + user_info: Any | None, user_api_key_dict: UserAPIKeyAuth, -) -> tuple[list[Any], Optional[list[Any]]]: +) -> tuple[list[Any], list[Any] | None]: """Fetch and merge teams from membership + user.teams field.""" from litellm.proxy.management_endpoints.team_endpoints import list_team @@ -667,7 +678,7 @@ async def _get_user_info_teams( team_list = teams_1 team_id_list = [team.team_id for team in teams_1] - teams_2: Optional[list[Any]] = None + teams_2: list[Any] | None = None target_team_ids = getattr(user_info, "teams", None) if target_team_ids and isinstance(target_team_ids, list): @@ -701,8 +712,8 @@ _SCIM_DIRECTORY_METADATA_KEYS = frozenset( def _redact_scim_enterprise_metadata( - metadata: Optional[Dict[str, Any]], -) -> Optional[Dict[str, Any]]: + metadata: dict[str, Any] | None, +) -> dict[str, Any] | None: """SCIM enterprise attributes, entitlements, and roles are persisted in user metadata so reporting can group on them, but they are directory-only fields that generic user-info endpoints must not surface; SCIM clients read them @@ -713,11 +724,11 @@ def _redact_scim_enterprise_metadata( def _build_user_info_response( - user_id: Optional[str], - user_info: Optional[Any], - keys: Optional[List[LiteLLM_VerificationToken]], + user_id: str | None, + user_info: Any | None, + keys: list[LiteLLM_VerificationToken] | None, team_list: list[Any], - teams_1: Optional[list[Any]], + teams_1: list[Any] | None, ) -> UserInfoResponse: """Create UserInfoResponse while filtering sensitive fields.""" if user_info is None and keys is not None: @@ -749,7 +760,7 @@ def _build_user_info_response( @management_endpoint_wrapper async def user_info( request: Request, - user_id: Optional[str] = fastapi.Query(default=None, description="User ID in the request parameters"), + user_id: str | None = fastapi.Query(default=None, description="User ID in the request parameters"), user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """ @@ -886,7 +897,7 @@ async def _check_user_info_v2_access( @management_endpoint_wrapper async def user_info_v2( request: Request, - user_id: Optional[str] = fastapi.Query(default=None, description="User ID in the request parameters"), + user_id: str | None = fastapi.Query(default=None, description="User ID in the request parameters"), user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """ @@ -996,7 +1007,7 @@ async def _get_user_info_for_proxy_admin(user_api_key_dict: UserAPIKeyAuth): verbose_proxy_logger.debug("results_keys: %s", results) - _keys_in_db: List = results[0]["keys"] or [] + _keys_in_db: list = results[0]["keys"] or [] # cast all keys to LiteLLM_VerificationToken keys_in_db = [] for key in _keys_in_db: @@ -1005,7 +1016,7 @@ async def _get_user_info_for_proxy_admin(user_api_key_dict: UserAPIKeyAuth): keys_in_db.append(LiteLLM_VerificationToken(**key)) # cast all teams to LiteLLM_TeamTable - _teams_in_db: List = results[0]["teams"] or [] + _teams_in_db: list = results[0]["teams"] or [] _teams_in_db = [LiteLLM_TeamTable(**team) for team in _teams_in_db] _teams_in_db.sort(key=lambda x: getattr(x, "team_alias", "") or "") returned_keys = _process_keys_for_user_info(keys=keys_in_db, all_teams=_teams_in_db) @@ -1032,8 +1043,8 @@ async def _get_user_info_for_proxy_admin(user_api_key_dict: UserAPIKeyAuth): def _process_keys_for_user_info( - keys: Optional[List[LiteLLM_VerificationToken]], - all_teams: Optional[Union[List[LiteLLM_TeamTable], List[TeamListResponseObject]]], + keys: list[LiteLLM_VerificationToken] | None, + all_teams: list[LiteLLM_TeamTable] | list[TeamListResponseObject] | None, ): from litellm.constants import UI_SESSION_TOKEN_TEAM_ID from litellm.proxy.proxy_server import general_settings, litellm_master_key_hash @@ -1073,9 +1084,7 @@ def _process_keys_for_user_info( return returned_keys -def _update_internal_user_params( - data_json: dict, data: Union[UpdateUserRequest, UpdateUserRequestNoUserIDorEmail] -) -> dict: +def _update_internal_user_params(data_json: dict, data: UpdateUserRequest | UpdateUserRequestNoUserIDorEmail) -> dict: non_default_values = {} fields_set = data.fields_set() if hasattr(data, "fields_set") else set() @@ -1124,11 +1133,11 @@ def _update_internal_user_params( async def _schedule_user_update_audit_log( - response: Dict[str, Any], - existing_user_row: Optional[BaseModel], - litellm_changed_by: Optional[str], + response: dict[str, Any], + existing_user_row: BaseModel | None, + litellm_changed_by: str | None, user_api_key_dict: UserAPIKeyAuth, - litellm_proxy_admin_name: Optional[str], + litellm_proxy_admin_name: str | None, ) -> None: from litellm.proxy.proxy_server import prisma_client @@ -1156,7 +1165,7 @@ async def _schedule_user_update_audit_log( def _check_user_update_authz( user_request: UpdateUserRequest, user_api_key_dict: UserAPIKeyAuth, - existing_user_row: Optional[BaseModel], + existing_user_row: BaseModel | None, ) -> None: """Authorization checks for /user/update — raises HTTPException on failure.""" if user_request.user_role is not None and user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value: @@ -1201,8 +1210,8 @@ async def _invalidate_user_spend_counter_if_changed( async def _update_single_user_helper( user_request: UpdateUserRequest, user_api_key_dict: UserAPIKeyAuth, - litellm_changed_by: Optional[str] = None, -) -> Dict[str, Any]: + litellm_changed_by: str | None = None, +) -> dict[str, Any]: """ Helper function to update a single user. Used by both user_update and bulk_user_update endpoints. @@ -1226,7 +1235,7 @@ async def _update_single_user_helper( non_default_values = _update_internal_user_params(data_json=data_json, data=user_request) _hash_password_in_dict(non_default_values) - existing_user_row: Optional[BaseModel] = None + existing_user_row: BaseModel | None = None if user_request.user_id: existing_user_row = await UserRepository(prisma_client).table.find_first( where={"user_id": user_request.user_id} @@ -1261,7 +1270,7 @@ async def _update_single_user_helper( ) existing_metadata = ( - cast(Dict, getattr(existing_user_row, "metadata", {}) or {}) if existing_user_row is not None else {} + cast(dict, getattr(existing_user_row, "metadata", {}) or {}) if existing_user_row is not None else {} ) non_default_values = prepare_metadata_fields( @@ -1274,7 +1283,7 @@ async def _update_single_user_helper( validate_finite_spend(non_default_values.get("spend")) # Perform the update - response: Optional[Dict[str, Any]] = None + response: dict[str, Any] | None = None if user_request.user_id and len(user_request.user_id) > 0: non_default_values["user_id"] = user_request.user_id @@ -1434,11 +1443,11 @@ async def user_update( async def bulk_update_processed_users( - users_to_update: List[UpdateUserRequest], + users_to_update: list[UpdateUserRequest], user_api_key_dict: UserAPIKeyAuth, - litellm_changed_by: Optional[str] = None, + litellm_changed_by: str | None = None, ) -> BulkUpdateUserResponse: - results: List[UserUpdateResult] = [] + results: list[UserUpdateResult] = [] successful_updates = 0 failed_updates = 0 @@ -1502,7 +1511,7 @@ async def bulk_update_processed_users( async def bulk_user_update( data: BulkUpdateUserRequest, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), - litellm_changed_by: Optional[str] = Header( + litellm_changed_by: str | None = Header( None, description="The litellm-changed-by header enables tracking of actions performed by authorized users on behalf of other users, providing an audit trail for accountability", ), @@ -1578,7 +1587,7 @@ async def bulk_user_update( ) # Determine the list of users to update - users_to_update: Union[List[UpdateUserRequest], List[UpdateUserRequestNoUserIDorEmail]] = [] + users_to_update: list[UpdateUserRequest] | list[UpdateUserRequestNoUserIDorEmail] = [] if data.all_users and data.user_updates: # Only proxy admins can update all users at once @@ -1616,7 +1625,7 @@ async def bulk_user_update( successful_updates = 0 failed_updates = 0 - results: List[UserUpdateResult] = [] + results: list[UserUpdateResult] = [] try: # Perform bulk database update @@ -1696,7 +1705,7 @@ async def bulk_user_update( ) return await bulk_update_processed_users( - users_to_update=cast(List[UpdateUserRequest], users_to_update), + users_to_update=cast(list[UpdateUserRequest], users_to_update), user_api_key_dict=user_api_key_dict, litellm_changed_by=litellm_changed_by, ) @@ -1704,7 +1713,7 @@ async def bulk_user_update( async def get_user_key_counts( prisma_client, - user_ids: Optional[List[str]] = None, + user_ids: list[str] | None = None, ): """ Helper function to get the count of keys for each user using Prisma's count method. @@ -1739,8 +1748,8 @@ async def get_user_key_counts( return result -def _validate_sort_params(sort_by: Optional[str], sort_order: str) -> Optional[Dict[str, str]]: - order_by: Dict[str, str] = {} +def _validate_sort_params(sort_by: str | None, sort_order: str) -> dict[str, str] | None: + order_by: dict[str, str] = {} if sort_by is None: return None @@ -1773,11 +1782,11 @@ def _validate_sort_params(sort_by: Optional[str], sort_order: str) -> Optional[D async def _authorize_user_list_request( user_api_key_dict: UserAPIKeyAuth, - organization_ids: Optional[str], + organization_ids: str | None, prisma_client: Any, user_api_key_cache: Any, proxy_logging_obj: Any, -) -> Optional[str]: +) -> str | None: """ Authorize the /user/list request and return the (possibly scoped) organization_ids string. @@ -1844,19 +1853,19 @@ async def _authorize_user_list_request( response_model=UserListResponse, ) async def get_users( - role: Optional[str] = fastapi.Query(default=None, description="Filter users by role"), - user_ids: Optional[str] = fastapi.Query(default=None, description="Get list of users by user_ids"), - sso_user_ids: Optional[str] = fastapi.Query(default=None, description="Get list of users by sso_user_id"), - user_email: Optional[str] = fastapi.Query(default=None, description="Filter users by partial email match"), - team: Optional[str] = fastapi.Query(default=None, description="Filter users by team id"), + role: str | None = fastapi.Query(default=None, description="Filter users by role"), + user_ids: str | None = fastapi.Query(default=None, description="Get list of users by user_ids"), + sso_user_ids: str | None = fastapi.Query(default=None, description="Get list of users by sso_user_id"), + user_email: str | None = fastapi.Query(default=None, description="Filter users by partial email match"), + team: str | None = fastapi.Query(default=None, description="Filter users by team id"), page: int = fastapi.Query(default=1, ge=1, description="Page number"), page_size: int = fastapi.Query(default=25, ge=1, le=100, description="Number of items per page"), - sort_by: Optional[str] = fastapi.Query( + sort_by: str | None = fastapi.Query( default=None, description="Column to sort by (e.g. 'user_id', 'user_email', 'created_at', 'spend')", ), sort_order: str = fastapi.Query(default="asc", description="Sort order ('asc' or 'desc')"), - organization_ids: Optional[str] = fastapi.Query( + organization_ids: str | None = fastapi.Query( default=None, description="Filter users by organization membership. Comma-separated list of org IDs.", ), @@ -1914,7 +1923,7 @@ async def get_users( skip = (page - 1) * page_size # Build where conditions based on provided parameters - where_conditions: Dict[str, Any] = {} + where_conditions: dict[str, Any] = {} if role: where_conditions["user_role"] = role @@ -1958,7 +1967,7 @@ async def get_users( # Build order_by conditions - order_by: Optional[Dict[str, str]] = ( + order_by: dict[str, str] | None = ( _validate_sort_params(sort_by, sort_order) if sort_by is not None and isinstance(sort_by, str) else None ) @@ -1984,7 +1993,7 @@ async def get_users( total_pages = -(-total_count // page_size) # Ceiling division # Prepare response - user_list: List[LiteLLM_UserTableWithKeyCount] = [] + user_list: list[LiteLLM_UserTableWithKeyCount] = [] if users is not None: for user in users: user_dump = user.model_dump() @@ -2011,7 +2020,7 @@ async def get_users( async def delete_user( data: DeleteUserRequest, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), - litellm_changed_by: Optional[str] = Header( + litellm_changed_by: str | None = Header( None, description="The litellm-changed-by header enables tracking of actions performed by authorized users on behalf of other users, providing an audit trail for accountability", ), @@ -2080,7 +2089,7 @@ async def delete_user( # Batch-fetch target memberships once before the per-user loop. Avoids # an N+1 DB call when delete_user is called with a large user_ids list. - target_org_ids_by_user: Dict[str, set] = {} + target_org_ids_by_user: dict[str, set] = {} if not caller_is_proxy_admin: all_target_memberships = await OrganizationMembershipRepository(prisma_client).table.find_many( where={"user_id": {"in": data.user_ids}} @@ -2156,7 +2165,7 @@ async def delete_user( ), ) if is_member_in_team: - _db_new_team_members: List[dict] = [m.model_dump() for m in new_team_members] + _db_new_team_members: list[dict] = [m.model_dump() for m in new_team_members] team.members_with_roles = json.dumps(_db_new_team_members) teams_to_update.append(team) @@ -2241,11 +2250,11 @@ async def add_internal_user_to_organization( async def _resolve_org_filter_for_user_search( user_api_key_dict: UserAPIKeyAuth, - team_id: Optional[str], + team_id: str | None, prisma_client: Any, user_api_key_cache: Any, proxy_logging_obj: Any, -) -> Optional[List[str]]: +) -> list[str] | None: """ Return a list of org IDs to filter by, or ``None`` for no filter. @@ -2279,7 +2288,7 @@ async def _resolve_org_filter_for_user_search( # Collect org IDs from ALL org memberships (any role, not just ORG_ADMIN). # This allows team admins who are org members to search users in their org. - member_org_ids: List[str] = [] + member_org_ids: list[str] = [] if caller_user is not None: member_org_ids = [m.organization_id for m in (caller_user.organization_memberships or [])] @@ -2311,7 +2320,7 @@ async def _resolve_team_org_filter( prisma_client: Any, user_api_key_cache: Any, proxy_logging_obj: Any, -) -> List[str]: +) -> list[str]: """Look up the team and return its org as a filter list, or raise 403.""" from litellm.proxy.management_endpoints.common_utils import _is_user_team_admin @@ -2351,13 +2360,13 @@ async def _resolve_team_org_filter( dependencies=[Depends(user_api_key_auth)], include_in_schema=False, responses={ - 200: {"model": List[LiteLLM_UserTableFiltered]}, + 200: {"model": list[LiteLLM_UserTableFiltered]}, }, ) async def ui_view_users( - user_id: Optional[str] = fastapi.Query(default=None, description="User ID in the request parameters"), - user_email: Optional[str] = fastapi.Query(default=None, description="User email in the request parameters"), - team_id: Optional[str] = fastapi.Query( + user_id: str | None = fastapi.Query(default=None, description="User ID in the request parameters"), + user_email: str | None = fastapi.Query(default=None, description="User email in the request parameters"), + team_id: str | None = fastapi.Query( default=None, description="Team ID — used when a team admin searches for users to add to their team", ), @@ -2400,7 +2409,7 @@ async def ui_view_users( skip = (page - 1) * page_size # Build where conditions based on provided parameters - where_conditions: Dict[str, Any] = {} + where_conditions: dict[str, Any] = {} if user_id: where_conditions["user_id"] = { @@ -2419,7 +2428,7 @@ async def ui_view_users( where_conditions["organization_memberships"] = {"some": {"organization_id": {"in": org_filter_ids}}} # Query users with pagination and filters - users: Optional[List[BaseModel]] = await UserRepository(prisma_client).table.find_many( + users: list[BaseModel] | None = await UserRepository(prisma_client).table.find_many( where=where_conditions, skip=skip, take=page_size, @@ -2441,10 +2450,14 @@ async def ui_view_users( # Using shared metric helper implementations from common_daily_activity -async def _resolve_user_email_metadata(prisma_client: "PrismaClient", records: list[Any]) -> dict[str, dict]: +async def _resolve_user_email_metadata( + prisma_client: "PrismaClient", records: Sequence[DailySpendRecord] +) -> dict[str, dict]: """Map each user_id on the page to its email/alias so the Usage dashboard can label the 'Spend Per User' chart with the email instead of the raw UUID.""" - user_ids = {record.user_id for record in records if getattr(record, "user_id", None)} + user_ids = { + user_id for record in records if isinstance(user_id := getattr(record, "user_id", None), str) and user_id + } if not user_ids: return {} users = await UserRepository(prisma_client).table.find_many(where={"user_id": {"in": list(user_ids)}}) @@ -2459,29 +2472,29 @@ async def _resolve_user_email_metadata(prisma_client: "PrismaClient", records: l ) @management_endpoint_wrapper async def get_user_daily_activity( - start_date: Optional[str] = fastapi.Query( + start_date: str | None = fastapi.Query( default=None, description="Start date in YYYY-MM-DD format", ), - end_date: Optional[str] = fastapi.Query( + end_date: str | None = fastapi.Query( default=None, description="End date in YYYY-MM-DD format", ), - model: Optional[str] = fastapi.Query( + model: str | None = fastapi.Query( default=None, description="Filter by specific model", ), - api_key: Optional[str] = fastapi.Query( + api_key: str | None = fastapi.Query( default=None, description="Filter by specific API key", ), - user_id: Optional[str] = fastapi.Query( + user_id: str | None = fastapi.Query( default=None, description="Filter by specific user ID. Admins can filter by any user or omit for global view. Non-admins must provide their own user_id.", ), page: int = fastapi.Query(default=1, description="Page number for pagination", ge=1), page_size: int = fastapi.Query(default=50, description="Items per page", ge=1, le=1000), - timezone: Optional[int] = fastapi.Query( + timezone: int | None = fastapi.Query( default=None, description="Timezone offset in minutes from UTC (e.g., 480 for PST). " "Matches JavaScript's Date.getTimezoneOffset() convention.", @@ -2568,27 +2581,27 @@ async def get_user_daily_activity( ) @management_endpoint_wrapper async def get_user_daily_activity_aggregated( - start_date: Optional[str] = fastapi.Query( + start_date: str | None = fastapi.Query( default=None, description="Start date in YYYY-MM-DD format", ), - end_date: Optional[str] = fastapi.Query( + end_date: str | None = fastapi.Query( default=None, description="End date in YYYY-MM-DD format", ), - model: Optional[str] = fastapi.Query( + model: str | None = fastapi.Query( default=None, description="Filter by specific model", ), - api_key: Optional[str] = fastapi.Query( + api_key: str | None = fastapi.Query( default=None, description="Filter by specific API key", ), - user_id: Optional[str] = fastapi.Query( + user_id: str | None = fastapi.Query( default=None, description="Filter by specific user ID. Admins can filter by any user or omit for global view. Non-admins must provide their own user_id.", ), - timezone: Optional[int] = fastapi.Query( + timezone: int | None = fastapi.Query( default=None, description="Timezone offset in minutes from UTC (e.g., 480 for PST). " "Matches JavaScript's Date.getTimezoneOffset() convention.", diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index f591e855a81..282184d6495 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -19,9 +19,10 @@ import functools import importlib import json import os +from collections.abc import Iterable from dataclasses import dataclass from datetime import datetime, timedelta, timezone -from typing import Any, Dict, Iterable, List, Literal, Optional, Set +from typing import Any, Literal from fastapi import ( APIRouter, @@ -47,10 +48,10 @@ from litellm._logging import verbose_logger, verbose_proxy_logger from litellm._uuid import uuid from litellm.constants import LITELLM_PROXY_ADMIN_NAME from litellm.proxy._experimental.mcp_server.utils import ( - build_env_var_setup_url, - collect_env_var_references, LITELLM_MCP_SERVER_DESCRIPTION, LITELLM_MCP_SERVER_NAME, + build_env_var_setup_url, + collect_env_var_references, get_server_prefix, parse_admin_env_vars, ) @@ -194,7 +195,7 @@ if MCP_AVAILABLE: expires_at: datetime def _validate_mcp_server_name_fields(payload: Any) -> None: - candidates: List[tuple[str, Optional[str]]] = [] + candidates: list[tuple[str, str | None]] = [] server_name = getattr(payload, "server_name", None) alias = getattr(payload, "alias", None) @@ -260,7 +261,7 @@ if MCP_AVAILABLE: general_settings as proxy_general_settings, ) - required_fields: Optional[List[str]] = proxy_general_settings.get("mcp_required_fields") + required_fields: list[str] | None = proxy_general_settings.get("mcp_required_fields") if not required_fields: return @@ -320,7 +321,7 @@ if MCP_AVAILABLE: return server.server_name return server.server_id - def _build_mcp_registry_entry_for_server(server: MCPServer, base_url: str) -> Dict[str, Any]: + def _build_mcp_registry_entry_for_server(server: MCPServer, base_url: str) -> dict[str, Any]: server_name = _build_mcp_registry_server_name(server) title = server_name description = server_name @@ -344,7 +345,7 @@ if MCP_AVAILABLE: ], } - def _build_builtin_registry_entry(base_url: str) -> Dict[str, Any]: + def _build_builtin_registry_entry(base_url: str) -> dict[str, Any]: remote_url = _build_registry_remote_url(base_url, "/mcp") return { "name": LITELLM_MCP_SERVER_NAME, @@ -359,7 +360,7 @@ if MCP_AVAILABLE: ], } - _temporary_mcp_servers: Dict[str, _TemporaryMCPServerEntry] = {} + _temporary_mcp_servers: dict[str, _TemporaryMCPServerEntry] = {} def _prune_expired_temporary_mcp_servers() -> None: if not _temporary_mcp_servers: @@ -391,7 +392,7 @@ if MCP_AVAILABLE: if cache_backend is None or not hasattr(cache_backend, "async_set_cache"): return - payload: Dict[str, Any] = server.model_dump(mode="json") + payload: dict[str, Any] = server.model_dump(mode="json") payload_json = json.dumps(payload) try: encrypted_payload = encrypt_value_helper(payload_json) @@ -414,7 +415,7 @@ if MCP_AVAILABLE: async def _get_temporary_mcp_server_from_redis( server_id: str, - ) -> Optional[MCPServer]: + ) -> MCPServer | None: """ Best-effort read from Redis shared cache. Returns None on miss/errors. @@ -455,7 +456,7 @@ if MCP_AVAILABLE: return None if not isinstance(loaded, dict): return None - payload_dict: Dict[str, Any] = loaded + payload_dict: dict[str, Any] = loaded try: return MCPServer(**payload_dict) @@ -465,7 +466,7 @@ if MCP_AVAILABLE: async def get_cached_temporary_mcp_server( server_id: str, - ) -> Optional[MCPServer]: + ) -> MCPServer | None: _prune_expired_temporary_mcp_servers() entry = _temporary_mcp_servers.get(server_id) if entry is None: @@ -520,7 +521,7 @@ if MCP_AVAILABLE: def _redact_mcp_credentials_list( mcp_servers: Iterable[LiteLLM_MCPServerTable], - ) -> List[LiteLLM_MCPServerTable]: + ) -> list[LiteLLM_MCPServerTable]: return [_redact_mcp_credentials(server) for server in mcp_servers] def _user_is_full_admin(user_api_key_dict: UserAPIKeyAuth) -> bool: @@ -587,7 +588,7 @@ if MCP_AVAILABLE: def _sanitize_mcp_server_list_for_non_admin( mcp_servers: Iterable[LiteLLM_MCPServerTable], - ) -> List[LiteLLM_MCPServerTable]: + ) -> list[LiteLLM_MCPServerTable]: return [_sanitize_mcp_server_for_non_admin(s) for s in mcp_servers] def _sanitize_mcp_server_for_virtual_key( @@ -644,7 +645,7 @@ if MCP_AVAILABLE: def _sanitize_mcp_server_list_for_virtual_key( mcp_servers: Iterable[LiteLLM_MCPServerTable], - ) -> List[LiteLLM_MCPServerTable]: + ) -> list[LiteLLM_MCPServerTable]: return [_sanitize_mcp_server_for_virtual_key(server) for server in mcp_servers] # (server attribute, credentials key) a session server inherits from the server it derives from. @@ -697,7 +698,7 @@ if MCP_AVAILABLE: except AttributeError: pass - payload_dict: Dict[str, Any] + payload_dict: dict[str, Any] try: payload_dict = payload.model_dump() # type: ignore[attr-defined] except AttributeError: @@ -707,7 +708,7 @@ if MCP_AVAILABLE: def _build_temporary_mcp_server_record( payload: NewMCPServerRequest, - created_by: Optional[str], + created_by: str | None, ) -> LiteLLM_MCPServerTable: now = datetime.utcnow() server_id = payload.server_id or str(uuid.uuid4()) @@ -848,7 +849,7 @@ if MCP_AVAILABLE: verbose_proxy_logger.debug("MCP registry request from IP=%s", client_ip) base_url = get_request_base_url(request) - registry_servers: List[Dict[str, Any]] = [] + registry_servers: list[dict[str, Any]] = [] registry_servers.append({"server": _build_builtin_registry_entry(base_url)}) # Centralized IP-based filtering: external callers only see public servers @@ -881,7 +882,7 @@ if MCP_AVAILABLE: async def _get_team_scoped_mcp_server_list( team_id: str, - ) -> List[LiteLLM_MCPServerTable]: + ) -> list[LiteLLM_MCPServerTable]: """ Return MCP servers scoped to a team: team's allowed servers + allow_all_keys servers. Used by the Create Key UI to populate the MCP server dropdown. @@ -908,7 +909,7 @@ if MCP_AVAILABLE: return [] # Collect servers from registry - servers: List[LiteLLM_MCPServerTable] = [] + servers: list[LiteLLM_MCPServerTable] = [] for server_id in all_allowed_ids: server = global_mcp_server_manager.get_mcp_server_by_id(server_id) if server is not None: @@ -919,7 +920,7 @@ if MCP_AVAILABLE: async def _resolve_accessible_mcp_servers( user_api_key_dict: UserAPIKeyAuth, - ) -> List[LiteLLM_MCPServerTable]: + ) -> list[LiteLLM_MCPServerTable]: """The server set the dashboard grid shows (GET /v1/mcp/server, no team filter), returned unredacted. Callers that surface this to a client must apply their own redaction; the per-user env-var status endpoint relies on @@ -932,7 +933,7 @@ if MCP_AVAILABLE: if _get_user_mcp_management_mode() == "view_all" and not _is_restricted_virtual_key_request(user_api_key_dict): return await global_mcp_server_manager.get_all_mcp_servers_unfiltered() - aggregated: Dict[str, LiteLLM_MCPServerTable] = {} + aggregated: dict[str, LiteLLM_MCPServerTable] = {} for auth_context in await build_effective_auth_contexts(user_api_key_dict): for server in await global_mcp_server_manager.get_all_allowed_mcp_servers(user_api_key_auth=auth_context): aggregated.setdefault(server.server_id, server) @@ -942,11 +943,11 @@ if MCP_AVAILABLE: "/server", description="Returns the mcp server list with associated teams", dependencies=[Depends(user_api_key_auth)], - response_model=List[LiteLLM_MCPServerTable], + response_model=list[LiteLLM_MCPServerTable], ) async def fetch_all_mcp_servers( user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), - team_id: Optional[str] = Query( + team_id: str | None = Query( None, description="Filter MCP servers by team scope. When provided, returns only " "servers the team has access to plus globally available (allow_all_keys) servers. " @@ -1048,7 +1049,7 @@ if MCP_AVAILABLE: dependencies=[Depends(user_api_key_auth)], ) async def health_check_servers( - server_ids: Optional[List[str]] = Query( + server_ids: list[str] | None = Query( None, description="Server IDs to check. If not provided, checks all accessible servers.", ), @@ -1081,7 +1082,7 @@ if MCP_AVAILABLE: auth_contexts = await build_effective_auth_contexts(user_api_key_dict) - server_status_map: Dict[str, Optional[Literal["healthy", "unhealthy", "unknown"]]] = {} + server_status_map: dict[str, Literal["healthy", "unhealthy", "unknown"] | None] = {} for auth_context in auth_contexts: servers = await global_mcp_server_manager.get_all_mcp_servers_with_health_and_teams( user_api_key_auth=auth_context, @@ -1399,7 +1400,7 @@ if MCP_AVAILABLE: async def add_mcp_server( payload: NewMCPServerRequest, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), - litellm_changed_by: Optional[str] = Header( + litellm_changed_by: str | None = Header( None, description="The litellm-changed-by header enables tracking of actions performed by authorized users on behalf of other users, providing an audit trail for accountability", ), @@ -1489,7 +1490,7 @@ if MCP_AVAILABLE: async def add_session_mcp_server( payload: NewMCPServerRequest, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), - litellm_changed_by: Optional[str] = Header( + litellm_changed_by: str | None = Header( None, description="The litellm-changed-by header enables tracking of actions performed by authorized users on behalf of other users, providing an audit trail for accountability", ), @@ -1647,7 +1648,7 @@ if MCP_AVAILABLE: async def _get_cached_temporary_mcp_server_or_404( server_id: str, user_api_key_dict: UserAPIKeyAuth, - request: Optional[Request] = None, + request: Request | None = None, ) -> MCPServer: server = await get_cached_temporary_mcp_server(server_id) resolved_from_temp_cache = server is not None @@ -1677,7 +1678,7 @@ if MCP_AVAILABLE: status_code=status.HTTP_403_FORBIDDEN, detail={"error": f"Access denied to MCP server {server_id}"}, ) - allowed_server_ids: Set[str] = set() + allowed_server_ids: set[str] = set() for auth_context in await build_effective_auth_contexts(user_api_key_dict): allowed_server_ids.update(await global_mcp_server_manager.get_allowed_mcp_servers(auth_context)) if server.server_id not in allowed_server_ids: @@ -1696,13 +1697,13 @@ if MCP_AVAILABLE: request: Request, server_id: str, user_api_key_dict: UserAPIKeyAuth = Depends(_mcp_oauth_user_api_key_auth), - client_id: Optional[str] = None, + client_id: str | None = None, redirect_uri: str = Query(...), state: str = "", - code_challenge: Optional[str] = None, - code_challenge_method: Optional[str] = None, - response_type: Optional[str] = None, - scope: Optional[str] = None, + code_challenge: str | None = None, + code_challenge_method: str | None = None, + response_type: str | None = None, + scope: str | None = None, ): mcp_server = await _get_cached_temporary_mcp_server_or_404(server_id, user_api_key_dict, request=request) _raise_if_not_oauth2(mcp_server) @@ -1756,13 +1757,13 @@ if MCP_AVAILABLE: server_id: str, user_api_key_dict: UserAPIKeyAuth = Depends(_mcp_oauth_user_api_key_auth), grant_type: str = Form(...), - code: Optional[str] = Form(None), - redirect_uri: Optional[str] = Form(None), - client_id: Optional[str] = Form(None), - client_secret: Optional[str] = Form(None), - code_verifier: Optional[str] = Form(None), - refresh_token: Optional[str] = Form(None), - scope: Optional[str] = Form(None), + code: str | None = Form(None), + redirect_uri: str | None = Form(None), + client_id: str | None = Form(None), + client_secret: str | None = Form(None), + code_verifier: str | None = Form(None), + refresh_token: str | None = Form(None), + scope: str | None = Form(None), ): mcp_server = await _get_cached_temporary_mcp_server_or_404(server_id, user_api_key_dict, request=request) _raise_if_not_oauth2(mcp_server) @@ -1844,7 +1845,7 @@ if MCP_AVAILABLE: async def remove_mcp_server( server_id: str, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), - litellm_changed_by: Optional[str] = Header( + litellm_changed_by: str | None = Header( None, description="The litellm-changed-by header enables tracking of actions performed by authorized users on behalf of other users, providing an audit trail for accountability", ), @@ -2007,7 +2008,7 @@ if MCP_AVAILABLE: # expires_at rather than recomputing it here (which could diverge by # milliseconds or if the storage logic ever adds a grace period). stored = await get_user_oauth_credential(prisma_client, user_id, server_id) - expires_at: Optional[str] = stored.get("expires_at") if stored else None + expires_at: str | None = stored.get("expires_at") if stored else None return MCPOAuthUserCredentialStatus( server_id=server_id, has_credential=True, @@ -2076,7 +2077,7 @@ if MCP_AVAILABLE: cred = await get_user_oauth_credential(prisma_client, user_id, server_id) if cred is None: return MCPOAuthUserCredentialStatus(server_id=server_id, has_credential=False, is_expired=False) - expires_at: Optional[str] = cred.get("expires_at") + expires_at: str | None = cred.get("expires_at") is_expired = False if expires_at: try: @@ -2096,7 +2097,7 @@ if MCP_AVAILABLE: "/user-credentials", description="List all OAuth2 MCP credentials stored for the calling user", dependencies=[Depends(user_api_key_auth)], - response_model=List[MCPUserCredentialListItem], + response_model=list[MCPUserCredentialListItem], ) @management_endpoint_wrapper async def list_mcp_user_credentials( @@ -2114,13 +2115,15 @@ if MCP_AVAILABLE: if not oauth_creds: return [] # Fetch server metadata for display names — single batch query instead of N+1. - server_ids = [c["server_id"] for c in oauth_creds] + server_ids = [c["server_id"] for c in oauth_creds if "server_id" in c] servers = {srv.server_id: srv for srv in await get_mcp_servers(prisma_client, server_ids)} - items: List[MCPUserCredentialListItem] = [] + items: list[MCPUserCredentialListItem] = [] for cred in oauth_creds: + if "server_id" not in cred: + continue sid = cred["server_id"] srv = servers.get(sid) - expires_at: Optional[str] = cred.get("expires_at") + expires_at: str | None = cred.get("expires_at") items.append( MCPUserCredentialListItem( server_id=sid, @@ -2182,7 +2185,7 @@ if MCP_AVAILABLE: def _compute_user_env_var_status( *, server: LiteLLM_MCPServerTable, - stored_values: Dict[str, str], + stored_values: dict[str, str], ) -> MCPUserEnvVarsStatus: """Build a status object for one server given the user's stored values. @@ -2211,7 +2214,7 @@ if MCP_AVAILABLE: user_var_names = {spec["name"] for spec in user_specs} blocking = {name for name in (referenced & user_var_names) if name not in global_values} - required: List[MCPUserEnvVarSpec] = [] + required: list[MCPUserEnvVarSpec] = [] missing_count = 0 for spec in user_specs: name = spec["name"] @@ -2334,12 +2337,12 @@ if MCP_AVAILABLE: description="Per-user MCP env var status across every server the user can access. " "Used by the dashboard to highlight servers with missing per-user vars.", dependencies=[Depends(user_api_key_auth)], - response_model=List[MCPUserEnvVarsStatus], + response_model=list[MCPUserEnvVarsStatus], ) @management_endpoint_wrapper async def list_mcp_user_env_var_status( user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), - ) -> List[MCPUserEnvVarsStatus]: + ) -> list[MCPUserEnvVarsStatus]: prisma_client = get_prisma_client_or_throw("Database not connected. Connect a database to your proxy") user_id = user_api_key_dict.user_id or "" if not user_id: @@ -2349,7 +2352,7 @@ if MCP_AVAILABLE: return [] server_ids = [s.server_id for s in accessible] stored_bulk = await get_user_env_vars_bulk(prisma_client, user_id, server_ids) - statuses: List[MCPUserEnvVarsStatus] = [] + statuses: list[MCPUserEnvVarsStatus] = [] for server in accessible: stored = stored_bulk.get(server.server_id, {}) status_obj = _compute_user_env_var_status(server=server, stored_values=stored) @@ -2368,7 +2371,7 @@ if MCP_AVAILABLE: async def edit_mcp_server( payload: UpdateMCPServerRequest, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), - litellm_changed_by: Optional[str] = Header( + litellm_changed_by: str | None = Header( None, description="The litellm-changed-by header enables tracking of actions performed by authorized users on behalf of other users, providing an audit trail for accountability", ), @@ -2564,16 +2567,16 @@ if MCP_AVAILABLE: "mcp_registry.json", ) - _mcp_registry_cache: Optional[Dict[str, Any]] = None + _mcp_registry_cache: dict[str, Any] | None = None - def _load_mcp_registry() -> Dict[str, Any]: + def _load_mcp_registry() -> dict[str, Any]: """Load the curated MCP registry from disk. Cached after first read.""" global _mcp_registry_cache if _mcp_registry_cache is not None: return _mcp_registry_cache try: with open(_MCP_REGISTRY_PATH, "r") as f: - data: Dict[str, Any] = json.load(f) + data: dict[str, Any] = json.load(f) except Exception as e: verbose_proxy_logger.warning(f"Failed to load MCP registry from {_MCP_REGISTRY_PATH}: {e}") data = {"servers": []} @@ -2586,8 +2589,8 @@ if MCP_AVAILABLE: dependencies=[Depends(user_api_key_auth)], ) async def discover_mcp_servers( - query: Optional[str] = Query(None, description="Search filter for server names and descriptions"), - category: Optional[str] = Query(None, description="Filter by category"), + query: str | None = Query(None, description="Search filter for server names and descriptions"), + category: str | None = Query(None, description="Filter by category"), user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """ @@ -2641,9 +2644,9 @@ if MCP_AVAILABLE: ) @functools.lru_cache(maxsize=1) - def _load_openapi_registry() -> Dict[str, Any]: + def _load_openapi_registry() -> dict[str, Any]: with open(_OPENAPI_REGISTRY_PATH, "r") as f: - data: Dict[str, Any] = json.load(f) + data: dict[str, Any] = json.load(f) return data @router.get( @@ -2694,7 +2697,7 @@ if MCP_AVAILABLE: async def add_mcp_toolset( payload: NewMCPToolsetRequest, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), - litellm_changed_by: Optional[str] = Header(None), + litellm_changed_by: str | None = Header(None), ): """Create a named toolset — a curated selection of {server_id, tool_name} pairs.""" prisma_client = get_prisma_client_or_throw("Database not connected. Connect a database to your proxy") @@ -2783,7 +2786,7 @@ if MCP_AVAILABLE: async def edit_mcp_toolset( payload: UpdateMCPToolsetRequest, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), - litellm_changed_by: Optional[str] = Header(None), + litellm_changed_by: str | None = Header(None), ): prisma_client = get_prisma_client_or_throw("Database not connected. Connect a database to your proxy") if LitellmUserRoles.PROXY_ADMIN != user_api_key_dict.user_role: @@ -2833,7 +2836,7 @@ if MCP_AVAILABLE: async def remove_mcp_toolset( toolset_id: str, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), - litellm_changed_by: Optional[str] = Header(None), + litellm_changed_by: str | None = Header(None), ): prisma_client = get_prisma_client_or_throw("Database not connected. Connect a database to your proxy") if LitellmUserRoles.PROXY_ADMIN != user_api_key_dict.user_role: diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index d458d0f7c4a..91f0b9b2790 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -1040,22 +1040,6 @@ class ModelManagementAuthChecks: return True -def _deployment_name_and_model(deployment: Optional[Union[Deployment, Dict[str, object]]]) -> Tuple[Optional[str], str]: - """Return (model_name, litellm_params.model) for a deployment. - - delete_deployment is annotated to return a Deployment but hands back the raw - model_list dict at runtime, so both shapes are handled; the model defaults to "". - """ - if deployment is None: - return None, "" - if isinstance(deployment, dict): - name = deployment.get("model_name") - params = deployment.get("litellm_params") - model = params.get("model") if isinstance(params, dict) else None - return (name if isinstance(name, str) else None), (model if isinstance(model, str) else "") - return deployment.model_name, str(getattr(deployment.litellm_params, "model", "") or "") - - #### [BETA] - This is a beta endpoint, format might change based on user feedback. - https://github.com/BerriAI/litellm/issues/964 @router.post( "/model/delete", @@ -1127,19 +1111,7 @@ async def delete_model( ## DELETE FROM ROUTER ## if llm_router is not None: - deleted_deployment = llm_router.delete_deployment(id=model_info.id) - # delete_deployment only drops the deployment from model_list; the auto/ - # complexity router registries are keyed by model_name and would otherwise - # retain a stale (now unbacked) entry, so evict it here too. Guard on the - # auto_router/ prefix (as clear_cache does): a regular DB model that merely - # shares a model_name with a config-defined router must not evict that router, - # since add_deployment never restores config-defined routers. - deleted_name, deleted_model = _deployment_name_and_model(deleted_deployment) - if deleted_name is not None and deleted_model.startswith("auto_router/"): - llm_router.auto_routers.pop(deleted_name, None) - llm_router.complexity_routers.pop(deleted_name, None) - llm_router.adaptive_routers.pop(deleted_name, None) - llm_router.quality_routers.pop(deleted_name, None) + llm_router.delete_deployment(id=model_info.id) # Runs after the row delete so the sibling check sees post-delete state. if model_params.model_info.team_id is not None: diff --git a/litellm/proxy/management_endpoints/organization_endpoints.py b/litellm/proxy/management_endpoints/organization_endpoints.py index df3c9c3c17b..1a4129ec5ee 100644 --- a/litellm/proxy/management_endpoints/organization_endpoints.py +++ b/litellm/proxy/management_endpoints/organization_endpoints.py @@ -13,7 +13,13 @@ Endpoints for /organization operations #### ORGANIZATION MANAGEMENT #### -from typing import Annotated, Any, Dict, List, Mapping, Optional, Tuple +from collections.abc import Mapping, Sequence +from typing import ( + TYPE_CHECKING, + Annotated, + Protocol, + overload, +) import fastapi from fastapi import APIRouter, Depends, HTTPException, Request, status @@ -57,9 +63,162 @@ from litellm.types.proxy.management_endpoints.common_daily_activity import ( ) from litellm.utils import _update_dictionary +if TYPE_CHECKING: + from types import TracebackType + + from prisma.models import LiteLLM_BudgetTable as PrismaBudgetTable + from prisma.models import ( + LiteLLM_ObjectPermissionTable as PrismaObjectPermissionTable, + ) + from prisma.models import ( + LiteLLM_OrganizationMembership as PrismaOrganizationMembership, + ) + from prisma.models import LiteLLM_OrganizationTable as PrismaOrganizationTable + from prisma.models import LiteLLM_UserTable as PrismaUserTable + router = APIRouter() +class _UserTableClient(Protocol): + async def find_unique(self, where: Mapping[str, object]) -> "PrismaUserTable | None": ... + + +class _BudgetTableClient(Protocol): + async def create(self, data: Mapping[str, object]) -> "PrismaBudgetTable": ... + + +class _ObjectPermissionTableClient(Protocol): + async def create(self, data: Mapping[str, object]) -> "PrismaObjectPermissionTable": ... + + +class _OrganizationTableClient(Protocol): + async def create( + self, data: Mapping[str, object], include: Mapping[str, object] | None = None + ) -> "PrismaOrganizationTable": ... + + async def find_unique( + self, where: Mapping[str, object], include: Mapping[str, object] | None = None + ) -> "PrismaOrganizationTable | None": ... + + async def find_many( + self, + where: Mapping[str, object] | None = None, + include: Mapping[str, object] | None = None, + ) -> "Sequence[PrismaOrganizationTable]": ... + + async def update( + self, + where: Mapping[str, object], + data: Mapping[str, object], + include: Mapping[str, object] | None = None, + ) -> "PrismaOrganizationTable": ... + + async def delete( + self, where: Mapping[str, object], include: Mapping[str, object] | None = None + ) -> "PrismaOrganizationTable | None": ... + + +class _OrganizationMembershipTableClient(Protocol): + async def create(self, data: Mapping[str, object]) -> "PrismaOrganizationMembership": ... + + async def find_unique( + self, where: Mapping[str, object], include: Mapping[str, object] | None = None + ) -> "PrismaOrganizationMembership | None": ... + + async def find_many( + self, where: Mapping[str, object] | None = None + ) -> "Sequence[PrismaOrganizationMembership]": ... + + async def update( + self, where: Mapping[str, object], data: Mapping[str, object] + ) -> "PrismaOrganizationMembership": ... + + async def delete(self, where: Mapping[str, object]) -> "PrismaOrganizationMembership | None": ... + + async def delete_many(self, where: Mapping[str, object]) -> int: ... + + +class _TeamTableClient(Protocol): + async def delete_many(self, where: Mapping[str, object]) -> int: ... + + +class _VerificationTokenTableClient(Protocol): + async def delete_many(self, where: Mapping[str, object]) -> int: ... + + +class _ObjectPermissionTxClient(Protocol): + async def upsert( + self, where: Mapping[str, object], data: Mapping[str, object] + ) -> "PrismaObjectPermissionTable": ... + + +class _BudgetTxClient(Protocol): + async def update(self, where: Mapping[str, object], data: Mapping[str, object]) -> "PrismaBudgetTable | None": ... + + +class _TransactionTables(Protocol): + @property + def litellm_objectpermissiontable(self) -> "_ObjectPermissionTxClient": ... + + @property + def litellm_budgettable(self) -> "_BudgetTxClient": ... + + @property + def litellm_organizationtable(self) -> "_OrganizationTableClient": ... + + +class _TransactionManager(Protocol): + async def __aenter__(self) -> "_TransactionTables": ... + + async def __aexit__( + self, + exc_type: type[BaseException] | None, + exc_value: BaseException | None, + traceback: "TracebackType | None", + ) -> bool | None: ... + + +@overload +def _table(repository: BudgetRepository) -> "_BudgetTableClient": ... + + +@overload +def _table(repository: ObjectPermissionRepository) -> "_ObjectPermissionTableClient": ... + + +@overload +def _table(repository: OrganizationRepository) -> "_OrganizationTableClient": ... + + +@overload +def _table(repository: OrganizationMembershipRepository) -> "_OrganizationMembershipTableClient": ... + + +@overload +def _table(repository: TeamRepository) -> "_TeamTableClient": ... + + +@overload +def _table(repository: UserRepository) -> "_UserTableClient": ... + + +@overload +def _table(repository: VerificationTokenRepository) -> "_VerificationTokenTableClient": ... + + +def _table( + repository: BudgetRepository + | ObjectPermissionRepository + | OrganizationRepository + | OrganizationMembershipRepository + | TeamRepository + | UserRepository + | VerificationTokenRepository, +) -> object: + prisma_table: object = repository.table + return prisma_table + + async def _verify_org_access( organization_id: str, user_api_key_dict: UserAPIKeyAuth, @@ -265,14 +424,15 @@ async def new_organization( detail={"error": f"soft_budget must be a non-negative finite number. Received: {data.soft_budget}"}, ) - user_object_correct_type: Optional[LiteLLM_UserTable] = None + user_object_correct_type: LiteLLM_UserTable | None = None if user_api_key_dict.user_id is not None: try: - user_object = await UserRepository(prisma_client).table.find_unique( + user_object = await _table(UserRepository(prisma_client)).find_unique( where={"user_id": user_api_key_dict.user_id} ) - user_object_correct_type = LiteLLM_UserTable(**user_object.model_dump()) + if user_object is not None: + user_object_correct_type = LiteLLM_UserTable.model_validate(user_object.model_dump()) except Exception: pass @@ -285,19 +445,21 @@ async def new_organization( budget_params = LiteLLM_BudgetTable.model_fields.keys() # Only include Budget Params when creating an entry in litellm_budgettable - _json_data = data.json(exclude_none=True) + _json_data = _STR_OBJECT_DICT_ADAPTER.validate_python(data.json(exclude_none=True)) _budget_data = {k: v for k, v in _json_data.items() if k in budget_params} - budget_row = LiteLLM_BudgetTable(**_budget_data) + budget_row = LiteLLM_BudgetTable.model_validate(_budget_data) - new_budget = prisma_client.jsonify_object(budget_row.json(exclude_none=True)) + new_budget = _STR_OBJECT_DICT_ADAPTER.validate_python( + prisma_client.jsonify_object(budget_row.json(exclude_none=True)) + ) - _budget = await BudgetRepository(prisma_client).table.create( + _budget = await _table(BudgetRepository(prisma_client)).create( data={ - **new_budget, # type: ignore + **new_budget, "created_by": user_api_key_dict.user_id or litellm_proxy_admin_name, "updated_by": user_api_key_dict.user_id or litellm_proxy_admin_name, } - ) # type: ignore + ) data.budget_id = _budget.budget_id @@ -339,11 +501,13 @@ async def new_organization( value=getattr(data, field), ) - new_organization_row = prisma_client.jsonify_object(organization_row.json(exclude_none=True)) + new_organization_row = _STR_OBJECT_DICT_ADAPTER.validate_python( + prisma_client.jsonify_object(organization_row.json(exclude_none=True)) + ) verbose_proxy_logger.info(f"new_organization_row: {json.dumps(new_organization_row, indent=2)}") - response = await OrganizationRepository(prisma_client).table.create( + response = await _table(OrganizationRepository(prisma_client)).create( data={ - **new_organization_row, # type: ignore + **new_organization_row, }, include={"litellm_budget_table": True}, ) @@ -357,14 +521,14 @@ async def new_organization( tags=["organization management"], ) async def get_organization_daily_activity( - organization_ids: Optional[str] = None, - start_date: Optional[str] = None, - end_date: Optional[str] = None, - model: Optional[str] = None, - api_key: Optional[str] = None, + organization_ids: str | None = None, + start_date: str | None = None, + end_date: str | None = None, + model: str | None = None, + api_key: str | None = None, page: int = 1, page_size: int = 10, - exclude_organization_ids: Optional[str] = None, + exclude_organization_ids: str | None = None, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """ @@ -382,13 +546,13 @@ async def get_organization_daily_activity( # Parse comma-separated ids org_ids_list = organization_ids.split(",") if organization_ids else None - exclude_org_ids_list: Optional[List[str]] = None + exclude_org_ids_list: list[str] | None = None if exclude_organization_ids: exclude_org_ids_list = exclude_organization_ids.split(",") if exclude_organization_ids else None # Restrict non-proxy-admins to only organizations where they are org_admin if not _user_has_admin_view(user_api_key_dict): - memberships = await OrganizationMembershipRepository(prisma_client).table.find_many( + memberships = await _table(OrganizationMembershipRepository(prisma_client)).find_many( where={"user_id": user_api_key_dict.user_id} ) admin_org_ids = [m.organization_id for m in memberships if m.user_role == LitellmUserRoles.ORG_ADMIN.value] @@ -405,11 +569,10 @@ async def get_organization_daily_activity( ) # Fetch organization aliases for metadata - where_condition = {} + where_condition = _STR_OBJECT_DICT_ADAPTER.validate_python({}) if org_ids_list: where_condition["organization_id"] = {"in": list(org_ids_list)} - org_aliases = await OrganizationRepository(prisma_client).table.find_many(where=where_condition) - org_alias_metadata = {o.organization_id: {"organization_alias": o.organization_alias} for o in org_aliases} + org_aliases = await _table(OrganizationRepository(prisma_client)).find_many(where=where_condition) # Query daily activity for organizations return await get_daily_activity( @@ -417,7 +580,7 @@ async def get_organization_daily_activity( table_name="litellm_dailyorganizationspend", entity_id_field="organization_id", entity_id=org_ids_list, - entity_metadata_field=org_alias_metadata, + entity_metadata_field={o.organization_id: {"organization_alias": o.organization_alias} for o in org_aliases}, exclude_entity_ids=exclude_org_ids_list, start_date=start_date, end_date=end_date, @@ -430,8 +593,8 @@ async def get_organization_daily_activity( async def _set_object_permission( data: NewOrganizationRequest, - prisma_client: Optional[PrismaClient], -) -> Optional[str]: + prisma_client: PrismaClient | None, +) -> str | None: """ Creates the LiteLLM_ObjectPermissionTable record for the organization. - Handles permissions for vector stores and mcp servers. @@ -442,7 +605,7 @@ async def _set_object_permission( return None if data.object_permission is not None: - created_object_permission = await ObjectPermissionRepository(prisma_client).table.create( + created_object_permission = await _table(ObjectPermissionRepository(prisma_client)).create( data=data.object_permission.model_dump(exclude_none=True), ) del data.object_permission @@ -534,10 +697,14 @@ async def update_organization( if updated_organization_row_json.get("metadata") is not None: existing_metadata = existing_organization_row.metadata or {} updated_metadata = updated_organization_row_json.get("metadata", {}) - merged_metadata = _update_dictionary(existing_dict=existing_metadata.copy(), new_dict=updated_metadata) + merged_metadata: Mapping[str, object] = _update_dictionary( + existing_dict=existing_metadata.copy(), new_dict=updated_metadata + ) updated_organization_row_json["metadata"] = merged_metadata - updated_organization_row = prisma_client.jsonify_object(updated_organization_row_json) + updated_organization_row = _STR_OBJECT_DICT_ADAPTER.validate_python( + prisma_client.jsonify_object(updated_organization_row_json) + ) if data.object_permission is not None: updated_organization_row = await handle_update_object_permission( data_json=updated_organization_row, @@ -559,7 +726,7 @@ async def update_organization( for field in LiteLLM_BudgetTable.model_fields.keys(): updated_organization_row.pop(field, None) - response = await OrganizationRepository(prisma_client).table.update( + response = await _table(OrganizationRepository(prisma_client)).update( where={"organization_id": data.organization_id}, data=updated_organization_row, include={"members": True, "teams": True, "litellm_budget_table": True}, @@ -569,9 +736,9 @@ async def update_organization( async def handle_update_object_permission( - data_json: dict, + data_json: dict[str, object], existing_organization_row: LiteLLM_OrganizationTable, -) -> dict: +) -> dict[str, object]: """ Handle the update of object permission for an organization. @@ -677,7 +844,7 @@ async def update_organization_v2( prisma_client=prisma_client, ) - existing_organization_row = await OrganizationRepository(prisma_client).table.find_unique( + existing_organization_row = await _table(OrganizationRepository(prisma_client)).find_unique( where={"organization_id": organization_id}, ) if existing_organization_row is None: @@ -711,15 +878,18 @@ async def update_organization_v2( else ({"object_permission_id": None} if object_permission_cleared else {}) ) - organization_write_data = prisma_client.jsonify_object( - { - **org_column_updates, - **object_permission_write, - "updated_by": user_api_key_dict.user_id, - } + organization_write_data = _STR_OBJECT_DICT_ADAPTER.validate_python( + prisma_client.jsonify_object( + { + **org_column_updates, + **object_permission_write, + "updated_by": user_api_key_dict.user_id, + } + ) ) - async with prisma_client.db.tx() as tx: + tx_manager: _TransactionManager = prisma_client.db.tx() + async with tx_manager as tx: if object_permission_upsert is not None: await tx.litellm_objectpermissiontable.upsert( where={"object_permission_id": object_permission_upsert.object_permission_id}, @@ -729,11 +899,12 @@ async def update_organization_v2( }, ) if budget_updates: + budget_write_data = _STR_OBJECT_DICT_ADAPTER.validate_python( + prisma_client.jsonify_object(dict(build_budget_write_data(budget_updates, user_api_key_dict.user_id))) + ) await tx.litellm_budgettable.update( where={"budget_id": existing_organization_row.budget_id}, - data=prisma_client.jsonify_object( - dict(build_budget_write_data(budget_updates, user_api_key_dict.user_id)) - ), + data=budget_write_data, ) response = await tx.litellm_organizationtable.update( where={"organization_id": organization_id}, @@ -748,7 +919,7 @@ async def update_organization_v2( "/organization/delete", tags=["organization management"], dependencies=[Depends(user_api_key_auth)], - response_model=List[LiteLLM_OrganizationTableWithMembers], + response_model=list[LiteLLM_OrganizationTableWithMembers], ) async def delete_organization( data: DeleteOrganizationRequest, @@ -778,15 +949,15 @@ async def delete_organization( deleted_orgs = [] for organization_id in data.organization_ids: # delete all teams in the organization - await TeamRepository(prisma_client).table.delete_many(where={"organization_id": organization_id}) + await _table(TeamRepository(prisma_client)).delete_many(where={"organization_id": organization_id}) # delete all members in the organization - await OrganizationMembershipRepository(prisma_client).table.delete_many( + await _table(OrganizationMembershipRepository(prisma_client)).delete_many( where={"organization_id": organization_id} ) # delete all keys in the organization - await VerificationTokenRepository(prisma_client).table.delete_many(where={"organization_id": organization_id}) + await _table(VerificationTokenRepository(prisma_client)).delete_many(where={"organization_id": organization_id}) # delete the organization - deleted_org = await OrganizationRepository(prisma_client).table.delete( + deleted_org = await _table(OrganizationRepository(prisma_client)).delete( where={"organization_id": organization_id}, include={"members": True, "teams": True, "litellm_budget_table": True}, ) @@ -804,13 +975,11 @@ async def delete_organization( "/organization/list", tags=["organization management"], dependencies=[Depends(user_api_key_auth)], - response_model=List[LiteLLM_OrganizationTableWithMembers], + response_model=list[LiteLLM_OrganizationTableWithMembers], ) async def list_organization( - org_id: Optional[str] = fastapi.Query( - default=None, description="Filter organizations by exact organization_id match" - ), - org_alias: Optional[str] = fastapi.Query( + org_id: str | None = fastapi.Query(default=None, description="Filter organizations by exact organization_id match"), + org_alias: str | None = fastapi.Query( default=None, description="Filter organizations by partial organization_alias match. Supports case-insensitive search.", ), @@ -849,7 +1018,7 @@ async def list_organization( ) # Build where conditions based on provided filters - where_conditions: Dict[str, Any] = {} + where_conditions: dict[str, object] = {} if org_id: where_conditions["organization_id"] = org_id @@ -862,13 +1031,13 @@ async def list_organization( # if proxy admin or admin viewer - get all orgs (with optional filters) if _user_has_admin_view(user_api_key_dict): - response = await OrganizationRepository(prisma_client).table.find_many( + response = await _table(OrganizationRepository(prisma_client)).find_many( where=where_conditions if where_conditions else None, include={"litellm_budget_table": True, "members": True, "teams": True}, ) # if internal user - get orgs they are a member of (with optional filters) else: - org_memberships = await OrganizationMembershipRepository(prisma_client).table.find_many( + org_memberships = await _table(OrganizationMembershipRepository(prisma_client)).find_many( where={"user_id": user_api_key_dict.user_id} ) membership_org_ids = [membership.organization_id for membership in org_memberships] @@ -882,7 +1051,7 @@ async def list_organization( response = [] else: where_conditions["organization_id"] = org_id - response = await OrganizationRepository(prisma_client).table.find_many( + response = await _table(OrganizationRepository(prisma_client)).find_many( where=where_conditions, include={ "litellm_budget_table": True, @@ -893,7 +1062,7 @@ async def list_organization( else: # Filter by membership and any additional filters where_conditions["organization_id"] = {"in": membership_org_ids} - response = await OrganizationRepository(prisma_client).table.find_many( + response = await _table(OrganizationRepository(prisma_client)).find_many( where=where_conditions, include={ "litellm_budget_table": True, @@ -933,9 +1102,7 @@ async def info_organization( prisma_client=prisma_client, ) - response: Optional[LiteLLM_OrganizationTableWithMembers] = await OrganizationRepository( - prisma_client - ).table.find_unique( + response = await _table(OrganizationRepository(prisma_client)).find_unique( where={"organization_id": organization_id}, include={ "litellm_budget_table": True, @@ -952,7 +1119,7 @@ async def info_organization( if response is None: raise HTTPException(status_code=404, detail={"error": "Organization not found"}) - response_pydantic_obj = LiteLLM_OrganizationTableWithMembers(**response.model_dump()) + response_pydantic_obj = LiteLLM_OrganizationTableWithMembers.model_validate(response.model_dump()) return response_pydantic_obj @@ -988,7 +1155,7 @@ async def deprecated_info_organization( prisma_client=prisma_client, ) - response = await OrganizationRepository(prisma_client).table.find_many( + response = await _table(OrganizationRepository(prisma_client)).find_many( where={"organization_id": {"in": data.organizations}}, include={"litellm_budget_table": True}, ) @@ -1065,7 +1232,7 @@ async def organization_member_add( ) # Check if organization exists - existing_organization_row = await OrganizationRepository(prisma_client).table.find_unique( + existing_organization_row = await _table(OrganizationRepository(prisma_client)).find_unique( where={"organization_id": data.organization_id} ) if existing_organization_row is None: @@ -1076,14 +1243,14 @@ async def organization_member_add( }, ) - members: List[OrgMember] - if isinstance(data.member, List): + members: Sequence[OrgMember] + if isinstance(data.member, list): members = data.member else: members = [data.member] - updated_users: List[LiteLLM_UserTable] = [] - updated_organization_memberships: List[LiteLLM_OrganizationMembershipTable] = [] + updated_users: list[LiteLLM_UserTable] = [] + updated_organization_memberships: list[LiteLLM_OrganizationMembershipTable] = [] for member in members: ( @@ -1138,7 +1305,7 @@ async def find_member_if_email(user_email: str, prisma_client: PrismaClient) -> "error": f"Unique user not found for user_email={user_email}. Potential duplicate OR non-existent user_email in LiteLLM_UserTable. Use 'user_id' instead." }, ) - existing_user_email_row_pydantic = LiteLLM_UserTable(**existing_user_email_row.model_dump()) + existing_user_email_row_pydantic = LiteLLM_UserTable.model_validate(existing_user_email_row.model_dump()) return existing_user_email_row_pydantic @@ -1176,7 +1343,7 @@ async def organization_member_update( ) # Check if organization exists - existing_organization_row = await OrganizationRepository(prisma_client).table.find_unique( + existing_organization_row = await _table(OrganizationRepository(prisma_client)).find_unique( where={"organization_id": data.organization_id} ) if existing_organization_row is None: @@ -1193,7 +1360,9 @@ async def organization_member_update( data.user_id = existing_user_email_row.user_id try: - existing_organization_membership = await OrganizationMembershipRepository(prisma_client).table.find_unique( + existing_organization_membership = await _table( + OrganizationMembershipRepository(prisma_client) + ).find_unique( where={ "user_id_organization_id": { "user_id": data.user_id, @@ -1218,7 +1387,7 @@ async def organization_member_update( # org-scoped operations. An org-admin of any org could otherwise # alter a PROXY_ADMIN user's per-org role, which has downstream # effects on admin UI filtering and scope derivation. - target_user_row = await UserRepository(prisma_client).table.find_unique(where={"user_id": data.user_id}) + target_user_row = await _table(UserRepository(prisma_client)).find_unique(where={"user_id": data.user_id}) if target_user_row is not None and getattr(target_user_row, "user_role", None) in ( LitellmUserRoles.PROXY_ADMIN.value, LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY.value, @@ -1235,7 +1404,7 @@ async def organization_member_update( # Update member role if data.role is not None: - await OrganizationMembershipRepository(prisma_client).table.update( + await _table(OrganizationMembershipRepository(prisma_client)).update( where={ "user_id_organization_id": { "user_id": data.user_id, @@ -1258,7 +1427,7 @@ async def organization_member_update( ) # update organization membership with new budget_id - await OrganizationMembershipRepository(prisma_client).table.update( + await _table(OrganizationMembershipRepository(prisma_client)).update( where={ "user_id_organization_id": { "user_id": data.user_id, @@ -1267,9 +1436,7 @@ async def organization_member_update( }, data={"budget_id": budget_id}, ) - final_organization_membership: Optional[BaseModel] = await OrganizationMembershipRepository( - prisma_client - ).table.find_unique( + final_organization_membership = await _table(OrganizationMembershipRepository(prisma_client)).find_unique( where={ "user_id_organization_id": { "user_id": data.user_id, @@ -1285,8 +1452,8 @@ async def organization_member_update( detail={"error": f"Member not found in organization={data.organization_id} for user_id={data.user_id}"}, ) - final_organization_membership_pydantic = LiteLLM_OrganizationMembershipTable( - **final_organization_membership.model_dump(exclude_none=True) + final_organization_membership_pydantic = LiteLLM_OrganizationMembershipTable.model_validate( + final_organization_membership.model_dump(exclude_none=True) ) return final_organization_membership_pydantic except Exception as e: @@ -1328,7 +1495,7 @@ async def organization_member_delete( existing_user_email_row = await find_member_if_email(data.user_email, prisma_client) data.user_id = existing_user_email_row.user_id - member_to_delete = await OrganizationMembershipRepository(prisma_client).table.delete( + member_to_delete = await _table(OrganizationMembershipRepository(prisma_client)).delete( where={ "user_id_organization_id": { "user_id": data.user_id, @@ -1347,7 +1514,7 @@ async def add_member_to_organization( member: OrgMember, organization_id: str, prisma_client: PrismaClient, -) -> Tuple[LiteLLM_UserTable, LiteLLM_OrganizationMembershipTable]: +) -> tuple[LiteLLM_UserTable, LiteLLM_OrganizationMembershipTable]: """ Add a member to an organization @@ -1357,12 +1524,12 @@ async def add_member_to_organization( """ try: - user_object: Optional[LiteLLM_UserTable] = None + user_object: LiteLLM_UserTable | None = None existing_user_id_row = None existing_user_email_row = None ## Check if user exists in LiteLLM_UserTable - user exists - either the user_id or user_email is in LiteLLM_UserTable if member.user_id is not None: - existing_user_id_row = await UserRepository(prisma_client).table.find_unique( + existing_user_id_row = await _table(UserRepository(prisma_client)).find_unique( where={"user_id": member.user_id} ) @@ -1387,16 +1554,16 @@ async def add_member_to_organization( _returned_user = await prisma_client.insert_data(data=new_user_defaults, table_name="user") # type: ignore if _returned_user is not None: - user_object = LiteLLM_UserTable(**_returned_user.model_dump()) + user_object = LiteLLM_UserTable.model_validate(_returned_user.model_dump()) elif existing_user_email_row is not None and len(existing_user_email_row) > 1: raise HTTPException( status_code=400, detail={"error": "Multiple users with this email found in db. Please use 'user_id' instead."}, ) elif existing_user_email_row is not None: - user_object = LiteLLM_UserTable(**existing_user_email_row.model_dump()) + user_object = LiteLLM_UserTable.model_validate(existing_user_email_row.model_dump()) elif existing_user_id_row is not None: - user_object = LiteLLM_UserTable(**existing_user_id_row.model_dump()) + user_object = LiteLLM_UserTable.model_validate(existing_user_id_row.model_dump()) else: raise HTTPException( status_code=404, @@ -1409,14 +1576,16 @@ async def add_member_to_organization( ) # Add user to organization - _organization_membership = await OrganizationMembershipRepository(prisma_client).table.create( + _organization_membership = await _table(OrganizationMembershipRepository(prisma_client)).create( data={ "organization_id": organization_id, "user_id": user_object.user_id, "user_role": member.role, } ) - organization_membership = LiteLLM_OrganizationMembershipTable(**_organization_membership.model_dump()) + organization_membership = LiteLLM_OrganizationMembershipTable.model_validate( + _organization_membership.model_dump() + ) return user_object, organization_membership except Exception as e: diff --git a/litellm/proxy/management_endpoints/tag_management_endpoints.py b/litellm/proxy/management_endpoints/tag_management_endpoints.py index 3cf933ee84c..47a8670e26f 100644 --- a/litellm/proxy/management_endpoints/tag_management_endpoints.py +++ b/litellm/proxy/management_endpoints/tag_management_endpoints.py @@ -12,8 +12,14 @@ All /tag management endpoints import asyncio import json +from collections.abc import Mapping, Sequence from datetime import datetime -from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union +from typing import ( + TYPE_CHECKING, + Protocol, + TypedDict, + overload, +) from fastapi import APIRouter, Depends, HTTPException, Query @@ -42,16 +48,101 @@ from litellm.types.tag_management import ( ) if TYPE_CHECKING: + from prisma.models import LiteLLM_BudgetTable as PrismaBudgetTable + from prisma.models import LiteLLM_ProxyModelTable as PrismaProxyModelTable + from prisma.models import LiteLLM_TagTable as PrismaTagTable + from prisma.models import LiteLLM_VerificationToken as PrismaVerificationToken + from litellm import Router + from litellm.proxy.utils import PrismaClient from litellm.types.router import Deployment router = APIRouter() +class _TagRecord(Protocol): + tag_name: str + description: str | None + models: Sequence[str] + model_info: object + budget_id: str | None + created_at: datetime + updated_at: datetime + created_by: str | None + litellm_budget_table: "PrismaBudgetTable | None" + + +class _TagTableClient(Protocol): + async def find_unique(self, where: Mapping[str, object]) -> "_TagRecord | None": ... + + async def find_many( + self, + where: Mapping[str, object] | None = None, + include: Mapping[str, object] | None = None, + ) -> "Sequence[_TagRecord]": ... + + async def create(self, data: Mapping[str, object]) -> "PrismaTagTable": ... + + async def update(self, where: Mapping[str, object], data: Mapping[str, object]) -> "PrismaTagTable": ... + + async def delete(self, where: Mapping[str, object]) -> "PrismaTagTable | None": ... + + +class _ModelTableClient(Protocol): + async def find_many(self, where: Mapping[str, object] | None = None) -> "Sequence[PrismaProxyModelTable]": ... + + +class _VerificationTokenTableClient(Protocol): + async def find_many( + self, + where: Mapping[str, object] | None = None, + select: Mapping[str, object] | None = None, + ) -> "Sequence[PrismaVerificationToken]": ... + + +class _DailyTagSpendGroupByRow(TypedDict): + tag: str | None + _min: Mapping[str, object] + _max: Mapping[str, object] + + +class _DailyTagSpendTableClient(Protocol): + async def group_by( + self, + by: Sequence[str], + where: Mapping[str, object] | None = None, + min: Mapping[str, object] | None = None, + max: Mapping[str, object] | None = None, + ) -> "Sequence[_DailyTagSpendGroupByRow]": ... + + +@overload +def _table(repository: DailyTagSpendRepository) -> "_DailyTagSpendTableClient": ... + + +@overload +def _table(repository: ModelRepository) -> "_ModelTableClient": ... + + +@overload +def _table(repository: TagRepository) -> "_TagTableClient": ... + + +@overload +def _table(repository: VerificationTokenRepository) -> "_VerificationTokenTableClient": ... + + +def _table( + repository: DailyTagSpendRepository | ModelRepository | TagRepository | VerificationTokenRepository, +) -> object: + prisma_table: object = repository.table + return prisma_table + + async def _get_internal_user_api_keys( - prisma_client, + prisma_client: "PrismaClient", user_api_key_dict: UserAPIKeyAuth, -) -> List[str]: +) -> list[str]: user_role = user_api_key_dict.user_role if user_role is None or not user_role.is_internal_user_role: return [] @@ -64,7 +155,7 @@ async def _get_internal_user_api_keys( if user_id is None: return sorted(user_api_keys) - key_records = await VerificationTokenRepository(prisma_client).table.find_many( + key_records = await _table(VerificationTokenRepository(prisma_client)).find_many( where={"user_id": user_id}, select={"token": True}, ) @@ -74,9 +165,9 @@ async def _get_internal_user_api_keys( async def _get_tag_list_scope( - prisma_client, + prisma_client: "PrismaClient", user_api_key_dict: UserAPIKeyAuth, -) -> Optional[Dict[str, dict]]: +) -> Mapping[str, Mapping[str, Sequence[str]]] | None: user_role = user_api_key_dict.user_role if user_api_key_has_admin_view(user_api_key_dict) or (user_role is None or not user_role.is_internal_user_role): return None @@ -89,10 +180,10 @@ async def _get_tag_list_scope( async def _get_tag_daily_activity_api_key_filter( - prisma_client, + prisma_client: "PrismaClient", user_api_key_dict: UserAPIKeyAuth, - requested_api_key: Optional[str], -) -> Optional[Union[str, List[str]]]: + requested_api_key: str | None, +) -> str | list[str] | None: user_role = user_api_key_dict.user_role if user_api_key_has_admin_view(user_api_key_dict) or (user_role is None or not user_role.is_internal_user_role): return requested_api_key @@ -106,17 +197,17 @@ async def _get_tag_daily_activity_api_key_filter( return scoped_api_keys -async def _get_model_names(prisma_client, model_ids: list) -> Dict[str, str]: +async def _get_model_names(prisma_client: "PrismaClient", model_ids: Sequence[str]) -> dict[str, str]: """Helper function to get model names from model IDs""" try: - models = await ModelRepository(prisma_client).table.find_many(where={"model_id": {"in": model_ids}}) + models = await _table(ModelRepository(prisma_client)).find_many(where={"model_id": {"in": model_ids}}) return {model.model_id: model.model_name for model in models} except Exception as e: verbose_proxy_logger.error(f"Error getting model names: {str(e)}") return {} -async def get_deployments_by_model(model: str, llm_router: "Router") -> List["Deployment"]: +async def get_deployments_by_model(model: str, llm_router: "Router") -> list["Deployment"]: """ Get all deployments by model """ @@ -181,7 +272,7 @@ async def new_tag( raise HTTPException(status_code=500, detail=CommonProxyErrors.no_llm_router.value) try: # Check if tag already exists - existing_tag = await TagRepository(prisma_client).table.find_unique(where={"tag_name": tag.name}) + existing_tag = await _table(TagRepository(prisma_client)).find_unique(where={"tag_name": tag.name}) if existing_tag is not None: raise HTTPException(status_code=400, detail=f"Tag {tag.name} already exists") @@ -198,7 +289,7 @@ async def new_tag( model_info = await _get_model_names(prisma_client, tag.models or []) # Create new tag in database - new_tag_record = await TagRepository(prisma_client).table.create( + new_tag_record = await _table(TagRepository(prisma_client)).create( data={ "tag_name": tag.name, "description": tag.description, @@ -321,7 +412,7 @@ async def update_tag( try: # Check if tag exists - existing_tag = await TagRepository(prisma_client).table.find_unique(where={"tag_name": tag.name}) + existing_tag = await _table(TagRepository(prisma_client)).find_unique(where={"tag_name": tag.name}) if existing_tag is None: raise HTTPException(status_code=404, detail=f"Tag {tag.name} not found") @@ -351,7 +442,7 @@ async def update_tag( update_data["budget_id"] = budget_id # Update tag in database - updated_tag_record = await TagRepository(prisma_client).table.update( + updated_tag_record = await _table(TagRepository(prisma_client)).update( where={"tag_name": tag.name}, data=update_data, ) @@ -398,7 +489,7 @@ async def info_tag( try: # Query tags from database with budget info - tag_records = await TagRepository(prisma_client).table.find_many( + tag_records = await _table(TagRepository(prisma_client)).find_many( where={"tag_name": {"in": data.names}}, include={"litellm_budget_table": True}, ) @@ -413,7 +504,7 @@ async def info_tag( requested_tags = {} for tag_record in tag_records: # Parse model_info from JSON - model_info = {} + model_info: object = {} if tag_record.model_info: if isinstance(tag_record.model_info, str): model_info = json.loads(tag_record.model_info) @@ -441,7 +532,7 @@ async def info_tag( raise HTTPException(status_code=500, detail=str(e)) -def _validate_tag_list_date_range(start_date: Optional[str], end_date: Optional[str]) -> None: +def _validate_tag_list_date_range(start_date: str | None, end_date: str | None) -> None: """Require both dates together, and enforce YYYY-MM-DD format with start <= end.""" if (start_date is None) != (end_date is None): raise HTTPException( @@ -472,7 +563,7 @@ def _validate_tag_list_date_range(start_date: Optional[str], end_date: Optional[ ) async def list_tags( user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), - start_date: Optional[str] = Query( + start_date: str | None = Query( None, description=( "Optional start date (YYYY-MM-DD). When provided together with " @@ -480,7 +571,7 @@ async def list_tags( "Stored tags are always returned." ), ), - end_date: Optional[str] = Query( + end_date: str | None = Query( None, description="Optional end date (YYYY-MM-DD). Must be given with start_date.", ), @@ -506,13 +597,13 @@ async def list_tags( # Prisma's distinct fetches all columns for all rows and deduplicates # in application code, which is extremely slow on large tables. # See: https://www.prisma.io/docs/orm/prisma-client/queries/aggregation-grouping-summarizing#distinct-under-the-hood - dynamic_tag_where: Dict[str, Any] = {"tag": {"not": None}} + dynamic_tag_where: dict[str, object] = {"tag": {"not": None}} if tag_scope: dynamic_tag_where = {**dynamic_tag_where, **tag_scope} if start_date is not None and end_date is not None: dynamic_tag_where["date"] = {"gte": start_date, "lte": end_date} - dynamic_tag_rows = await DailyTagSpendRepository(prisma_client).table.group_by( + dynamic_tag_rows = await _table(DailyTagSpendRepository(prisma_client)).group_by( by=["tag"], where=dynamic_tag_where, min={"created_at": True}, @@ -526,7 +617,7 @@ async def list_tags( stored_tag_where = {"tag_name": {"in": used_tag_names}} if tag_scope is not None else None ## QUERY STORED TAGS ## - tag_records = await TagRepository(prisma_client).table.find_many( + tag_records = await _table(TagRepository(prisma_client)).find_many( where=stored_tag_where, include={"litellm_budget_table": True}, ) @@ -536,7 +627,7 @@ async def list_tags( for tag_record in tag_records: stored_tag_names.add(tag_record.tag_name) # Parse model_info from JSON - model_info = {} + model_info: object = {} if tag_record.model_info: if isinstance(tag_record.model_info, str): model_info = json.loads(tag_record.model_info) @@ -598,12 +689,12 @@ async def delete_tag( try: # Check if tag exists - existing_tag = await TagRepository(prisma_client).table.find_unique(where={"tag_name": data.name}) + existing_tag = await _table(TagRepository(prisma_client)).find_unique(where={"tag_name": data.name}) if existing_tag is None: raise HTTPException(status_code=404, detail=f"Tag {data.name} not found") # Delete tag from database - await TagRepository(prisma_client).table.delete(where={"tag_name": data.name}) + await _table(TagRepository(prisma_client)).delete(where={"tag_name": data.name}) return {"message": f"Tag {data.name} deleted successfully"} except Exception as e: @@ -617,11 +708,11 @@ async def delete_tag( dependencies=[Depends(user_api_key_auth)], ) async def get_tag_daily_activity( - tags: Optional[str] = None, - start_date: Optional[str] = None, - end_date: Optional[str] = None, - model: Optional[str] = None, - api_key: Optional[str] = None, + tags: str | None = None, + start_date: str | None = None, + end_date: str | None = None, + model: str | None = None, + api_key: str | None = None, page: int = 1, page_size: int = 10, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), diff --git a/litellm/proxy/management_endpoints/tool_management_endpoints.py b/litellm/proxy/management_endpoints/tool_management_endpoints.py index b6a445ef327..ad0b4f7444c 100644 --- a/litellm/proxy/management_endpoints/tool_management_endpoints.py +++ b/litellm/proxy/management_endpoints/tool_management_endpoints.py @@ -11,21 +11,21 @@ POST /v1/tool/policy - Update the input_policy / output_policy for a import uuid from datetime import datetime, timedelta, timezone -from itertools import groupby from typing import TYPE_CHECKING, Annotated, Any, List, Optional from fastapi import APIRouter, Depends, HTTPException, Query -from pydantic import BaseModel, TypeAdapter +from pydantic import BaseModel, Field, TypeAdapter if TYPE_CHECKING: from litellm.proxy.utils import PrismaClient from litellm._logging import verbose_proxy_logger -from litellm.constants import TOOL_SPEND_MAX_WINDOW_DAYS +from litellm.constants import TOOL_SPEND_TOP_TOOLS from litellm.proxy._types import CommonProxyErrors, LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.repositories.object_permission_repository import ObjectPermissionRepository from litellm.repositories.table_repositories import ( + DailyToolSpendRepository, SpendLogsRepository, SpendLogToolIndexRepository, ) @@ -142,53 +142,18 @@ def _parse_day_start(value: str | None) -> datetime | None: ) -class _ToolSpendRow(BaseModel): - date: str +class _ToolSpendSums(BaseModel): + spend: float = 0.0 + total_tokens: int = 0 + request_count: int = 0 + + +class _TopToolRow(BaseModel): tool_name: str - call_count: int - spend: float - total_tokens: int + sums: _ToolSpendSums = Field(alias="_sum") -class _RequestTotalRow(BaseModel): - total_spend: float - - -_TOOL_SPEND_ROWS = TypeAdapter(list[_ToolSpendRow]) -_REQUEST_TOTAL_ROWS = TypeAdapter(list[_RequestTotalRow]) - - -def _summarize_tool(name: str, grp: tuple[_ToolSpendRow, ...]) -> ToolSpendEntry: - return ToolSpendEntry( - tool_name=name, - spend=sum(r.spend for r in grp), - call_count=sum(r.call_count for r in grp), - total_tokens=sum(r.total_tokens for r in grp), - ) - - -def _build_tool_spend_response( - rows: list[_ToolSpendRow], - total_spend: float, - start_date: str, - end_date: str, -) -> ToolSpendResponse: - daily = [ - ToolSpendDailyEntry(date=r.date, tool_name=r.tool_name, spend=r.spend, call_count=r.call_count) for r in rows - ] - grouped = groupby(sorted(rows, key=lambda r: r.tool_name), key=lambda r: r.tool_name) - by_tool = sorted( - (_summarize_tool(name, tuple(grp)) for name, grp in grouped), - key=lambda e: e.spend, - reverse=True, - ) - return ToolSpendResponse( - by_tool=by_tool, - daily=daily, - total_spend=total_spend, - start_date=start_date, - end_date=end_date, - ) +_TOP_TOOL_ROWS = TypeAdapter(list[_TopToolRow]) @router.get( @@ -205,16 +170,16 @@ async def get_tool_spend( """ Spend attributed to each tool over a date range, for the Cost Optimization dashboard. - Joins ``LiteLLM_SpendLogToolIndex`` (which tool names ran on which request) to - ``LiteLLM_SpendLogs`` (what the request cost). A request that used multiple tools - counts its full spend toward each of those tools, so per-tool numbers are - attributions. ``total_spend`` is the deduplicated spend of every request that - called at least one tool in the window, so it never double counts. + Reads the ``LiteLLM_DailyToolSpend`` rollup, written at request time from invoked + tools only (MCP tool calls and response tool_calls; declaring a tool without + invoking it does not count). A request that invoked multiple tools counts its + full spend toward each of them, so per-tool numbers are attributions and do not + sum to a deduplicated total. - ``start_date`` is clamped to at most 30 days before ``end_date`` (serving up to - 31 calendar dates inclusive, the same width as the endpoint's default window): - a wider requested range is clamped, and the response's ``start_date`` reflects - the effective window actually served. + ``by_tool`` is the top ``TOOL_SPEND_TOP_TOOLS`` tools by spend, aggregated in + SQL, and ``daily`` covers only those tools, so the response is bounded by + days x TOOL_SPEND_TOP_TOOLS regardless of the requested range or how many + distinct tool names exist. """ from litellm.proxy.proxy_server import prisma_client @@ -230,64 +195,46 @@ async def get_tool_spend( if prisma_client is None: raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value) - now = datetime.now(timezone.utc) - end_day = _parse_day_start(end_date) - # Anchor the floor to a midnight so the clamp compares dates with dates: - # parsed start_dates are midnight-aligned, and a floor carrying now's - # time-of-day would invisibly truncate an explicit start_date to mid-day. - today = now.replace(hour=0, minute=0, second=0, microsecond=0) - window_floor = (end_day or today) - timedelta(days=TOOL_SPEND_MAX_WINDOW_DAYS) - start_dt = _parse_day_start(start_date) or window_floor - if start_dt < window_floor: - start_dt = window_floor - end_exclusive = (end_day + timedelta(days=1)) if end_day else now + end_day = _parse_day_start(end_date) or datetime.now(timezone.utc) + start_day = _parse_day_start(start_date) or end_day - timedelta(days=30) + start_str = start_day.strftime("%Y-%m-%d") + end_str = end_day.strftime("%Y-%m-%d") + date_window = {"date": {"gte": start_str, "lte": end_str}} - # ti.start_time defines the window in both queries; the sl."startTime" bounds - # exist only so the planner can use the SpendLogs startTime index, and carry a - # 1s margin because the two writers can disagree by ~1ms on the same request. - rows = await prisma_client.db.query_raw( - """ - SELECT to_char(ti.start_time, 'YYYY-MM-DD') AS date, - ti.tool_name AS tool_name, - COUNT(*)::int AS call_count, - COALESCE(SUM(sl.spend), 0)::double precision AS spend, - COALESCE(SUM(sl.total_tokens), 0)::bigint AS total_tokens - FROM "LiteLLM_SpendLogToolIndex" ti - JOIN "LiteLLM_SpendLogs" sl ON sl.request_id = ti.request_id - WHERE ti.start_time >= ($1::timestamptz AT TIME ZONE 'UTC') - AND ti.start_time < ($2::timestamptz AT TIME ZONE 'UTC') - AND sl."startTime" >= ($1::timestamptz AT TIME ZONE 'UTC') - interval '1 second' - AND sl."startTime" < ($2::timestamptz AT TIME ZONE 'UTC') + interval '1 second' - GROUP BY date, ti.tool_name - ORDER BY date ASC, spend DESC - """, - start_dt.isoformat(), - end_exclusive.isoformat(), - ) - totals = await prisma_client.db.query_raw( - """ - SELECT COALESCE(SUM(sl.spend), 0)::double precision AS total_spend - FROM "LiteLLM_SpendLogs" sl - WHERE sl."startTime" >= ($1::timestamptz AT TIME ZONE 'UTC') - interval '1 second' - AND sl."startTime" < ($2::timestamptz AT TIME ZONE 'UTC') + interval '1 second' - AND EXISTS ( - SELECT 1 - FROM "LiteLLM_SpendLogToolIndex" ti - WHERE ti.request_id = sl.request_id - AND ti.start_time >= ($1::timestamptz AT TIME ZONE 'UTC') - AND ti.start_time < ($2::timestamptz AT TIME ZONE 'UTC') + table = DailyToolSpendRepository(prisma_client).table + top_tools = _TOP_TOOL_ROWS.validate_python( + await table.group_by( + by=["tool_name"], + sum={"spend": True, "total_tokens": True, "request_count": True}, + where=date_window, + order={"_sum": {"spend": "desc"}}, + take=TOOL_SPEND_TOP_TOOLS, ) - """, - start_dt.isoformat(), - end_exclusive.isoformat(), + or [] ) - total_rows = _REQUEST_TOTAL_ROWS.validate_python(totals or []) - return _build_tool_spend_response( - rows=_TOOL_SPEND_ROWS.validate_python(rows or []), - total_spend=total_rows[0].total_spend if total_rows else 0.0, - start_date=start_dt.strftime("%Y-%m-%d"), - end_date=(end_day or now).strftime("%Y-%m-%d"), + by_tool = [ + ToolSpendEntry( + tool_name=row.tool_name, + spend=row.sums.spend, + call_count=row.sums.request_count, + total_tokens=row.sums.total_tokens, + ) + for row in top_tools + ] + + daily_rows = ( + await table.find_many( + where={**date_window, "tool_name": {"in": [row.tool_name for row in top_tools]}}, + order=[{"date": "asc"}, {"spend": "desc"}], + ) + if top_tools + else [] ) + daily = [ + ToolSpendDailyEntry(date=row.date, tool_name=row.tool_name, spend=row.spend, call_count=row.request_count) + for row in daily_rows + ] + return ToolSpendResponse(by_tool=by_tool, daily=daily, start_date=start_str, end_date=end_str) @router.get( @@ -388,7 +335,8 @@ async def get_tool_usage_logs( user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """ - Return paginated spend logs for requests that used this tool (from SpendLogToolIndex). + Return paginated spend logs for requests that invoked this tool (from SpendLogToolIndex). + Declaring a tool in a request body without the model invoking it does not create an entry. """ from litellm.proxy.proxy_server import prisma_client diff --git a/litellm/proxy/policy_engine/policy_registry.py b/litellm/proxy/policy_engine/policy_registry.py index 0dec93251f8..e1afbf2f5f2 100644 --- a/litellm/proxy/policy_engine/policy_registry.py +++ b/litellm/proxy/policy_engine/policy_registry.py @@ -8,8 +8,16 @@ by policy_attachments (see AttachmentRegistry). """ import json +from collections.abc import Mapping, Sequence from datetime import datetime, timezone -from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple +from typing import ( + TYPE_CHECKING, + Any, + Optional, + Protocol, + TypedDict, + Union, +) from litellm._logging import verbose_proxy_logger from litellm.repositories.table_repositories import PolicyRepository @@ -33,7 +41,89 @@ if TYPE_CHECKING: POLICY_VERSION_ID_PREFIX = "policy_" -def _row_to_policy_db_response(row: Any) -> PolicyDBResponse: +class _RawPipelineStep(TypedDict): + guardrail: str + + +class _RawPipelineConfig(TypedDict, total=False): + mode: str + steps: Sequence[Union[PipelineStep, "_RawPipelineStep"]] + + +class _PolicyRow(Protocol): + policy_id: str + policy_name: str + version_number: int + version_status: str + parent_version_id: str | None + is_latest: bool + published_at: datetime | None + production_at: datetime | None + inherit: str | None + description: str | None + guardrails_add: list[str] | None + guardrails_remove: list[str] | None + condition: dict[str, object] | None + pipeline: dict[str, object] | None + created_at: datetime + updated_at: datetime + created_by: str | None + updated_by: str | None + + +class _PolicyVersionSourceRow(Protocol): + policy_id: str + policy_name: str + version_number: int + inherit: str | None + description: str | None + guardrails_add: Sequence[str] | None + guardrails_remove: Sequence[str] | None + condition: Mapping[str, object] | str | None + pipeline: Mapping[str, object] | str | None + + +class _PolicyTableClient(Protocol): + async def create(self, data: Mapping[str, object]) -> _PolicyRow: ... + + async def find_unique(self, where: Mapping[str, object]) -> _PolicyRow | None: ... + + async def find_many( + self, + where: Mapping[str, object] | None = None, + order: Mapping[str, str] | None = None, + ) -> Sequence[_PolicyRow]: ... + + async def update(self, where: Mapping[str, object], data: Mapping[str, object]) -> _PolicyRow: ... + + async def update_many(self, where: Mapping[str, object], data: Mapping[str, object]) -> int: ... + + async def delete(self, where: Mapping[str, object]) -> _PolicyRow | None: ... + + async def delete_many(self, where: Mapping[str, object]) -> int: ... + + +class _PolicyVersionSourceTableClient(Protocol): + async def find_unique(self, where: Mapping[str, object]) -> _PolicyVersionSourceRow | None: ... + + async def find_first( + self, + where: Mapping[str, object], + order: Mapping[str, str] | None = None, + ) -> _PolicyVersionSourceRow | None: ... + + +def _policy_table(prisma_client: "PrismaClient") -> _PolicyTableClient: + table: _PolicyTableClient = PolicyRepository(prisma_client).table + return table + + +def _policy_version_source_table(prisma_client: "PrismaClient") -> _PolicyVersionSourceTableClient: + table: _PolicyVersionSourceTableClient = PolicyRepository(prisma_client).table + return table + + +def _row_to_policy_db_response(row: _PolicyRow) -> PolicyDBResponse: """Build PolicyDBResponse from a Prisma LiteLLM_PolicyTable row.""" return PolicyDBResponse( policy_id=row.policy_id, @@ -71,11 +161,11 @@ class PolicyRegistry: """ def __init__(self): - self._policies: Dict[str, Policy] = {} - self._policies_by_id: Dict[str, Tuple[str, Policy]] = {} + self._policies: dict[str, Policy] = {} + self._policies_by_id: dict[str, tuple[str, Policy]] = {} self._initialized: bool = False - def load_policies(self, policies_config: Dict[str, Any]) -> None: + def load_policies(self, policies_config: Mapping[str, dict[str, object]]) -> None: """ Load policies from a configuration dictionary. @@ -98,7 +188,7 @@ class PolicyRegistry: self._initialized = True verbose_proxy_logger.info(f"Loaded {len(self._policies)} policies") - def _parse_policy(self, policy_name: str, policy_data: Dict[str, Any]) -> Policy: + def _parse_policy(self, policy_name: str, policy_data: dict[str, Any]) -> Policy: """ Parse a policy from raw configuration data. @@ -139,13 +229,13 @@ class PolicyRegistry: @staticmethod def _parse_pipeline( - pipeline_data: Optional[Dict[str, Any]], - ) -> Optional[GuardrailPipeline]: + pipeline_data: Optional["_RawPipelineConfig"], + ) -> GuardrailPipeline | None: """Parse a pipeline configuration from raw data.""" if pipeline_data is None: return None - steps_data = pipeline_data.get("steps", []) + steps_data: Sequence[PipelineStep | _RawPipelineStep] = pipeline_data.get("steps", []) steps = [PipelineStep(**step_data) if isinstance(step_data, dict) else step_data for step_data in steps_data] return GuardrailPipeline( @@ -153,7 +243,7 @@ class PolicyRegistry: steps=steps, ) - def get_policy(self, policy_name: str) -> Optional[Policy]: + def get_policy(self, policy_name: str) -> Policy | None: """ Get a policy by name. @@ -165,7 +255,7 @@ class PolicyRegistry: """ return self._policies.get(policy_name) - def get_all_policies(self) -> Dict[str, Policy]: + def get_all_policies(self) -> dict[str, Policy]: """ Get all loaded policies. @@ -174,7 +264,7 @@ class PolicyRegistry: """ return self._policies.copy() - def get_policy_names(self) -> List[str]: + def get_policy_names(self) -> list[str]: """ Get list of all policy names. @@ -247,7 +337,7 @@ class PolicyRegistry: self, policy_request: PolicyCreateRequest, prisma_client: "PrismaClient", - created_by: Optional[str] = None, + created_by: str | None = None, ) -> PolicyDBResponse: """ Add a policy to the database. @@ -263,7 +353,7 @@ class PolicyRegistry: try: now = datetime.now(timezone.utc) # Build data dict; new policy is v1 production - data: Dict[str, Any] = { + data: dict[str, object] = { "policy_name": policy_request.policy_name, "version_number": 1, "version_status": "production", @@ -289,7 +379,7 @@ class PolicyRegistry: validated_pipeline = GuardrailPipeline(**policy_request.pipeline) data["pipeline"] = json.dumps(validated_pipeline.model_dump()) - created_policy = await PolicyRepository(prisma_client).table.create(data=data) + created_policy = await _policy_table(prisma_client).create(data=data) # Also add to in-memory registry policy = self._parse_policy( @@ -317,7 +407,7 @@ class PolicyRegistry: policy_id: str, policy_request: PolicyUpdateRequest, prisma_client: "PrismaClient", - updated_by: Optional[str] = None, + updated_by: str | None = None, ) -> PolicyDBResponse: """ Update a policy in the database. Only draft versions can be updated. @@ -335,7 +425,7 @@ class PolicyRegistry: Exception: If policy is not in draft status (only drafts are editable). """ try: - existing = await PolicyRepository(prisma_client).table.find_unique(where={"policy_id": policy_id}) + existing = await _policy_table(prisma_client).find_unique(where={"policy_id": policy_id}) if existing is None: raise Exception(f"Policy with ID {policy_id} not found") version_status = getattr(existing, "version_status", "production") @@ -343,7 +433,7 @@ class PolicyRegistry: raise Exception(f"Only draft versions can be updated. This policy has status '{version_status}'.") # Build update data - only include fields that are set - update_data: Dict[str, Any] = { + update_data: dict[str, object] = { "updated_at": datetime.now(timezone.utc), "updated_by": updated_by, } @@ -364,7 +454,7 @@ class PolicyRegistry: validated_pipeline = GuardrailPipeline(**policy_request.pipeline) update_data["pipeline"] = json.dumps(validated_pipeline.model_dump()) - updated_policy = await PolicyRepository(prisma_client).table.update( + updated_policy = await _policy_table(prisma_client).update( where={"policy_id": policy_id}, data=update_data, ) @@ -380,7 +470,7 @@ class PolicyRegistry: self, policy_id: str, prisma_client: "PrismaClient", - ) -> Dict[str, Any]: + ) -> Mapping[str, str]: """ Delete a policy version from the database. @@ -395,7 +485,7 @@ class PolicyRegistry: Dict with "message" and optional "warning" if production was deleted. """ try: - policy = await PolicyRepository(prisma_client).table.find_unique(where={"policy_id": policy_id}) + policy = await _policy_table(prisma_client).find_unique(where={"policy_id": policy_id}) if policy is None: raise Exception(f"Policy with ID {policy_id} not found") @@ -404,9 +494,9 @@ class PolicyRegistry: policy_name = policy.policy_name # Delete from DB - await PolicyRepository(prisma_client).table.delete(where={"policy_id": policy_id}) + await _policy_table(prisma_client).delete(where={"policy_id": policy_id}) - result: Dict[str, Any] = {"message": f"Policy {policy_id} deleted successfully"} + result: dict[str, str] = {"message": f"Policy {policy_id} deleted successfully"} # Remove from in-memory registry only if this was the production version if version_status == "production": @@ -425,7 +515,7 @@ class PolicyRegistry: self, policy_id: str, prisma_client: "PrismaClient", - ) -> Optional[PolicyDBResponse]: + ) -> PolicyDBResponse | None: """ Get a policy by ID from the database. @@ -437,7 +527,7 @@ class PolicyRegistry: PolicyDBResponse if found, None otherwise """ try: - policy = await PolicyRepository(prisma_client).table.find_unique(where={"policy_id": policy_id}) + policy = await _policy_table(prisma_client).find_unique(where={"policy_id": policy_id}) if policy is None: return None @@ -447,7 +537,7 @@ class PolicyRegistry: verbose_proxy_logger.exception(f"Error getting policy from DB: {e}") raise Exception(f"Error getting policy from DB: {str(e)}") - def get_policy_by_id_for_request(self, policy_id: str) -> Optional[Tuple[str, Policy]]: + def get_policy_by_id_for_request(self, policy_id: str) -> tuple[str, Policy] | None: """ Return a policy version by ID from in-memory cache (no DB access). @@ -466,8 +556,8 @@ class PolicyRegistry: async def get_all_policies_from_db( self, prisma_client: "PrismaClient", - version_status: Optional[str] = None, - ) -> List[PolicyDBResponse]: + version_status: str | None = None, + ) -> list[PolicyDBResponse]: """ Get all policies from the database, optionally filtered by version_status. @@ -480,11 +570,11 @@ class PolicyRegistry: List of PolicyDBResponse objects """ try: - where: Dict[str, Any] = {} + where: dict[str, str] = {} if version_status is not None: where["version_status"] = version_status - policies = await PolicyRepository(prisma_client).table.find_many( + policies = await _policy_table(prisma_client).find_many( where=where if where else None, order={"created_at": "desc"}, ) @@ -524,7 +614,7 @@ class PolicyRegistry: self.add_policy(policy_response.policy_name, policy) self._policies_by_id = {} - non_production = await PolicyRepository(prisma_client).table.find_many( + non_production = await _policy_table(prisma_client).find_many( where={"version_status": {"in": ["draft", "published"]}}, order={"created_at": "desc"}, ) @@ -557,7 +647,7 @@ class PolicyRegistry: self, policy_name: str, prisma_client: "PrismaClient", - ) -> List[str]: + ) -> list[str]: """ Resolve all guardrails for a policy from the database. @@ -622,7 +712,7 @@ class PolicyRegistry: PolicyVersionListResponse with policy_name and list of versions """ try: - rows = await PolicyRepository(prisma_client).table.find_many( + rows = await _policy_table(prisma_client).find_many( where={"policy_name": policy_name}, order={"version_number": "desc"}, ) @@ -640,8 +730,8 @@ class PolicyRegistry: self, policy_name: str, prisma_client: "PrismaClient", - source_policy_id: Optional[str] = None, - created_by: Optional[str] = None, + source_policy_id: str | None = None, + created_by: str | None = None, ) -> PolicyDBResponse: """ Create a new draft version of a policy. Copies all fields from the source. @@ -658,14 +748,16 @@ class PolicyRegistry: """ try: if source_policy_id is not None: - source = await PolicyRepository(prisma_client).table.find_unique(where={"policy_id": source_policy_id}) + source = await _policy_version_source_table(prisma_client).find_unique( + where={"policy_id": source_policy_id} + ) if source is None: raise Exception(f"Source policy {source_policy_id} not found") if source.policy_name != policy_name: raise Exception(f"Source policy name '{source.policy_name}' does not match '{policy_name}'") else: # Find current production version for this policy_name - prod = await PolicyRepository(prisma_client).table.find_first( + prod = await _policy_version_source_table(prisma_client).find_first( where={ "policy_name": policy_name, "version_status": "production", @@ -676,7 +768,7 @@ class PolicyRegistry: source = prod # Next version number - latest = await PolicyRepository(prisma_client).table.find_first( + latest = await _policy_version_source_table(prisma_client).find_first( where={"policy_name": policy_name}, order={"version_number": "desc"}, ) @@ -684,12 +776,12 @@ class PolicyRegistry: now = datetime.now(timezone.utc) # Set is_latest=False on all existing versions for this policy_name - await PolicyRepository(prisma_client).table.update_many( + await _policy_table(prisma_client).update_many( where={"policy_name": policy_name}, data={"is_latest": False}, ) - data: Dict[str, Any] = { + data: dict[str, object] = { "policy_name": policy_name, "version_number": next_num, "version_status": "draft", @@ -714,7 +806,7 @@ class PolicyRegistry: if source.pipeline is not None: data["pipeline"] = json.dumps(source.pipeline) if isinstance(source.pipeline, dict) else source.pipeline - created = await PolicyRepository(prisma_client).table.create(data=data) + created = await _policy_table(prisma_client).create(data=data) return _row_to_policy_db_response(created) except Exception as e: verbose_proxy_logger.exception(f"Error creating new version: {e}") @@ -725,7 +817,7 @@ class PolicyRegistry: policy_id: str, new_status: str, prisma_client: "PrismaClient", - updated_by: Optional[str] = None, + updated_by: str | None = None, ) -> PolicyDBResponse: """ Update a policy version's status. Valid transitions: @@ -748,7 +840,7 @@ class PolicyRegistry: if new_status not in ("published", "production"): raise Exception(f"Invalid status '{new_status}'. Use 'published' or 'production'.") - row = await PolicyRepository(prisma_client).table.find_unique(where={"policy_id": policy_id}) + row = await _policy_table(prisma_client).find_unique(where={"policy_id": policy_id}) if row is None: raise Exception(f"Policy with ID {policy_id} not found") @@ -759,7 +851,7 @@ class PolicyRegistry: if new_status == "published": if current != "draft": raise Exception(f"Only draft versions can be published. Current status: '{current}'.") - updated = await PolicyRepository(prisma_client).table.update( + updated = await _policy_table(prisma_client).update( where={"policy_id": policy_id}, data={ "version_status": "published", @@ -780,7 +872,7 @@ class PolicyRegistry: raise Exception("Cannot promote draft directly to production. Publish the version first.") # Demote current production to published - await PolicyRepository(prisma_client).table.update_many( + await _policy_table(prisma_client).update_many( where={ "policy_name": policy_name, "version_status": "production", @@ -793,7 +885,7 @@ class PolicyRegistry: ) # Promote this version to production - updated = await PolicyRepository(prisma_client).table.update( + updated = await _policy_table(prisma_client).update( where={"policy_id": policy_id}, data={ "version_status": "production", @@ -843,8 +935,8 @@ class PolicyRegistry: PolicyVersionCompareResponse with both versions and field_diffs """ try: - a = await PolicyRepository(prisma_client).table.find_unique(where={"policy_id": policy_id_a}) - b = await PolicyRepository(prisma_client).table.find_unique(where={"policy_id": policy_id_b}) + a = await _policy_table(prisma_client).find_unique(where={"policy_id": policy_id_a}) + b = await _policy_table(prisma_client).find_unique(where={"policy_id": policy_id_b}) if a is None: raise Exception(f"Policy {policy_id_a} not found") if b is None: @@ -854,15 +946,15 @@ class PolicyRegistry: resp_b = _row_to_policy_db_response(b) # Compare fields that are part of policy content (not metadata) - compare_fields = [ + compare_fields = ( "inherit", "description", "guardrails_add", "guardrails_remove", "condition", "pipeline", - ] - field_diffs: Dict[str, Dict[str, Any]] = {} + ) + field_diffs: dict[str, dict[str, object]] = {} for field in compare_fields: val_a = getattr(resp_a, field) val_b = getattr(resp_b, field) @@ -882,7 +974,7 @@ class PolicyRegistry: self, policy_name: str, prisma_client: "PrismaClient", - ) -> Dict[str, str]: + ) -> Mapping[str, str]: """ Delete all versions of a policy. Also removes from in-memory registry. @@ -894,7 +986,7 @@ class PolicyRegistry: Dict with success message """ try: - await PolicyRepository(prisma_client).table.delete_many(where={"policy_name": policy_name}) + await _policy_table(prisma_client).delete_many(where={"policy_name": policy_name}) self.remove_policy(policy_name) return {"message": f"All versions of policy '{policy_name}' deleted successfully"} except Exception as e: @@ -903,7 +995,7 @@ class PolicyRegistry: # Global singleton instance -_policy_registry: Optional[PolicyRegistry] = None +_policy_registry: PolicyRegistry | None = None def get_policy_registry() -> PolicyRegistry: diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index e0365b7d621..4486cd7de59 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -109,6 +109,7 @@ from litellm.proxy.common_utils.callback_utils import ( is_sensitive_callback_key, normalize_callback_names, process_callback, + strip_callback_config, ) from litellm.proxy.common_utils.realtime_utils import _realtime_request_body from litellm.router_utils.add_retry_fallback_headers import ( @@ -13408,7 +13409,7 @@ async def async_queue_request( # extra_body); see above for the same guard upstream. data["metadata"] = {} data["metadata"]["user_api_key"] = user_api_key_dict.api_key - data["metadata"]["user_api_key_metadata"] = user_api_key_dict.metadata + data["metadata"]["user_api_key_metadata"] = strip_callback_config(user_api_key_dict.metadata) _headers = _safe_get_request_headers(request).copy() _headers.pop("authorization", None) # do not store the original `sk-..` api key in the db data["metadata"]["headers"] = _headers diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 314c59d3560..ed4f7d4f0a0 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -1102,6 +1102,19 @@ model LiteLLM_SpendLogToolIndex { @@index([start_time]) } +// Daily tool spend rollup (one row per tool per day) – the Cost Optimization card reads this, never SpendLogs +model LiteLLM_DailyToolSpend { + date String + tool_name String + spend Float @default(0.0) + total_tokens BigInt @default(0) + request_count BigInt @default(0) + created_at DateTime @default(now()) + updated_at DateTime @updatedAt + + @@id([date, tool_name]) +} + // Prompt table for storing prompt configurations model LiteLLM_PromptTable { id String @id @default(uuid()) diff --git a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py index 10c71c00110..8f3f8ad1bfc 100644 --- a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py +++ b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py @@ -2,6 +2,7 @@ import asyncio import json import os +from collections import Counter from collections.abc import Mapping from typing import Any, Dict, List, Optional, Set, Tuple, Type, Union from urllib.parse import urlparse @@ -25,6 +26,7 @@ from litellm.repositories.table_repositories import ( SSOConfigRepository, UISettingsRepository, ) +from litellm.repositories.team_repository import TeamRepository from litellm.types.proxy.management_endpoints.ui_sso import ( DefaultTeamSSOParams, SSOConfig, @@ -598,6 +600,51 @@ async def get_default_team_settings(): ) +def _default_team_ids(teams: list[str] | list[NewUserRequestTeam]) -> tuple[str, ...]: + return tuple(team if isinstance(team, str) else team.team_id for team in teams) + + +async def _validate_default_teams_exist(teams: list[str] | list[NewUserRequestTeam]) -> None: + """Reject default teams that cannot be assigned. + + New users are added to these teams long after the settings are saved, and that + consume path swallows the resulting 404, so an unknown team id would silently + drop every future user's team assignment unless it is caught here. + """ + team_ids = _default_team_ids(teams) + if not team_ids: + return + + duplicate_ids = tuple(team_id for team_id, count in Counter(team_ids).items() if count > 1) + if duplicate_ids: + raise HTTPException( + status_code=400, + detail={ + "error": f"Duplicate default team id(s): {', '.join(duplicate_ids)}. List each default team only once." + }, + ) + + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + raise HTTPException( + status_code=500, + detail={"error": "Database not connected. Please connect a database."}, + ) + + existing_teams = await TeamRepository(prisma_client).find_many(where={"team_id": {"in": list(team_ids)}}) + existing_team_ids = {team.team_id for team in existing_teams} + missing_ids = tuple(team_id for team_id in team_ids if team_id not in existing_team_ids) + if missing_ids: + raise HTTPException( + status_code=400, + detail={ + "error": f"Team(s) not found: {', '.join(missing_ids)}. " + "A team must exist before it can be set as a default team for new users." + }, + ) + + async def update_default_team_member_budget(teams: List[NewUserRequestTeam], user_api_key_dict: UserAPIKeyAuth): """ 1. Update the max member budget for the team @@ -706,6 +753,9 @@ async def update_internal_user_settings( Update the default internal user parameters for SSO users. These settings will be applied to new users who sign in via SSO. """ + if settings.teams is not None: + await _validate_default_teams_exist(settings.teams) + if settings.teams is not None and all(isinstance(team, NewUserRequestTeam) for team in settings.teams): await update_default_team_member_budget( settings.teams, diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index e85ccf150d2..924189fed4b 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -41,6 +41,7 @@ from litellm.constants import ( ) from litellm.proxy._types import ( DB_CONNECTION_ERROR_TYPES, + DB_RETRY_SAFE_ERROR_TYPES, CommonProxyErrors, ProxyErrorTypes, ProxyException, @@ -175,6 +176,7 @@ if TYPE_CHECKING: from prisma.client import TransactionManager from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.proxy.db.spend_log_tool_index import ToolUsageTransaction Span = Union[_Span, Any] else: @@ -2917,6 +2919,8 @@ async def prefetch_config_params(prisma_client: Any, param_names: List[str]) -> class PrismaClient: spend_log_transactions: List = [] _spend_log_transactions_lock = asyncio.Lock() + tool_usage_transactions: List["ToolUsageTransaction"] = [] + _tool_usage_transactions_lock = asyncio.Lock() def __init__( self, @@ -5334,7 +5338,7 @@ class ProxyUpdateSpend: ) break - except DB_CONNECTION_ERROR_TYPES as e: + except DB_RETRY_SAFE_ERROR_TYPES as e: if i >= n_retry_times: # If we've reached the maximum number of retries _raise_failed_update_spend_exception( e=e, start_time=start_time, proxy_logging_obj=proxy_logging_obj @@ -5473,12 +5477,15 @@ async def update_spend( queue_size = len(prisma_client.spend_log_transactions) verbose_proxy_logger.debug("Spend Logs transactions: {}".format(queue_size)) + async with prisma_client._tool_usage_transactions_lock: + tool_usage_queue_size = len(prisma_client.tool_usage_transactions) + # Process spend log transactions when called directly. # This keeps backwards compatibility with the old behavior. # See update_spend_logs_job and _monitor_spend_logs_queue for the new behavior. # Safe to keep: under high concurrency this can take up to ~30s to run, # so it's unlikely to overlap with monitor_spend_logs_queue. - if queue_size > 0: + if queue_size > 0 or tool_usage_queue_size > 0: await update_spend_logs_job( prisma_client=prisma_client, db_writer_client=db_writer_client, @@ -5545,10 +5552,14 @@ async def update_spend_logs_job( n_retry_times = 3 MAX_LOGS_PER_INTERVAL = 10000 - # Atomically pop batch from queue + # Atomically pop batch from queue. The tool usage queue counts toward the + # emptiness check: a spend-log write failure aborts a run before the tool + # drain below, and those entries must not strand once the spend queue drains. async with prisma_client._spend_log_transactions_lock: queue_size = len(prisma_client.spend_log_transactions) - if queue_size == 0: + async with prisma_client._tool_usage_transactions_lock: + tool_queue_size = len(prisma_client.tool_usage_transactions) + if queue_size == 0 and tool_queue_size == 0: return async with prisma_client._spend_log_transactions_lock: @@ -5579,17 +5590,23 @@ async def update_spend_logs_job( guardrail_tracking_err, ) - # Tool usage tracking (same batch): SpendLogToolIndex for "last N requests for tool X" + # Tool usage tracking: drain the request-time queue into the tool index and the + # LiteLLM_DailyToolSpend rollup. Never retried; a dropped batch is permanently + # absent from the rollup, so failures log at error. + async with prisma_client._tool_usage_transactions_lock: + tool_usage_to_process = prisma_client.tool_usage_transactions[:MAX_LOGS_PER_INTERVAL] + prisma_client.tool_usage_transactions = prisma_client.tool_usage_transactions[len(tool_usage_to_process) :] try: - from litellm.proxy.db.spend_log_tool_index import process_spend_logs_tool_usage + from litellm.proxy.db.spend_log_tool_index import flush_tool_usage_transactions - await process_spend_logs_tool_usage( + await flush_tool_usage_transactions( prisma_client=prisma_client, - logs_to_process=logs_to_process, + transactions=tool_usage_to_process, ) except Exception as tool_tracking_err: - verbose_proxy_logger.warning( - "Spend tracking - tool usage tracking failed (non-fatal): %s", + verbose_proxy_logger.error( + "Spend tracking - tool usage flush failed; %s tool usage transactions dropped: %s", + len(tool_usage_to_process), tool_tracking_err, ) @@ -5625,9 +5642,13 @@ async def _monitor_spend_logs_queue( while True: try: - # Check queue size with lock protection + # Check queue sizes with lock protection; the tool usage queue keeps + # the monitor firing when a prior failed run left it nonempty. async with prisma_client._spend_log_transactions_lock: - queue_size = len(prisma_client.spend_log_transactions) + spend_queue_size = len(prisma_client.spend_log_transactions) + async with prisma_client._tool_usage_transactions_lock: + tool_queue_size = len(prisma_client.tool_usage_transactions) + queue_size = spend_queue_size + tool_queue_size if queue_size > 0: if queue_size >= threshold: diff --git a/litellm/repositories/__init__.py b/litellm/repositories/__init__.py index 4451f0865da..29c953e06cf 100644 --- a/litellm/repositories/__init__.py +++ b/litellm/repositories/__init__.py @@ -23,6 +23,7 @@ from litellm.repositories.table_repositories import ( DailyGuardrailMetricsRepository, DailyPolicyMetricsRepository, DailyTagSpendRepository, + DailyToolSpendRepository, DeletedTeamRepository, DeletedVerificationTokenRepository, DeprecatedVerificationTokenRepository, @@ -104,6 +105,7 @@ __all__ = [ "ManagedVectorStoreIndexRepository", "WorkflowMessageRepository", "DailyTagSpendRepository", + "DailyToolSpendRepository", "SpendLogToolIndexRepository", "SpendLogGuardrailIndexRepository", "UserNotificationsRepository", diff --git a/litellm/repositories/table_repositories.py b/litellm/repositories/table_repositories.py index dc2a7d25259..54008c0950c 100644 --- a/litellm/repositories/table_repositories.py +++ b/litellm/repositories/table_repositories.py @@ -181,6 +181,10 @@ class SpendLogToolIndexRepository(PrismaTableRepository): table_name = "litellm_spendlogtoolindex" +class DailyToolSpendRepository(PrismaTableRepository): + table_name = "litellm_dailytoolspend" + + class SpendLogGuardrailIndexRepository(PrismaTableRepository): table_name = "litellm_spendlogguardrailindex" diff --git a/litellm/repositories/verification_token_repository.py b/litellm/repositories/verification_token_repository.py index 3ea5f32629b..19352c1b3c4 100644 --- a/litellm/repositories/verification_token_repository.py +++ b/litellm/repositories/verification_token_repository.py @@ -3,18 +3,37 @@ VerificationToken repository for database operations on LiteLLM_VerificationToke """ import json +from collections.abc import Iterator, Mapping from datetime import datetime -from typing import Any, Dict, List, Optional, Type +from typing import TYPE_CHECKING, Any, Protocol from litellm.models.verification_token import ( LiteLLM_VerificationToken, ) from litellm.repositories.base_repository import BaseRepository +if TYPE_CHECKING: + from prisma.models import ( + LiteLLM_VerificationToken as PrismaVerificationToken, + ) + + from litellm.proxy.utils import PrismaClient + + +class _DictConvertible(Protocol): + def dict(self) -> dict[str, object]: ... + + def __iter__(self) -> Iterator[tuple[str, object]]: ... + class VerificationTokenRepository(BaseRepository[LiteLLM_VerificationToken]): """Repository for verification token (API key) database operations.""" + @property + def prisma_client(self) -> "PrismaClient": + prisma_client: PrismaClient = super().prisma_client + return prisma_client + @property def table(self) -> Any: return self.prisma_client.db.litellm_verificationtoken @@ -24,10 +43,10 @@ class VerificationTokenRepository(BaseRepository[LiteLLM_VerificationToken]): return self.prisma_client.db.litellm_deletedverificationtoken @property - def model_class(self) -> Type[LiteLLM_VerificationToken]: + def model_class(self) -> type[LiteLLM_VerificationToken]: return LiteLLM_VerificationToken - def _to_model(self, record: Any) -> Optional[LiteLLM_VerificationToken]: + def _to_model(self, record: _DictConvertible | None) -> LiteLLM_VerificationToken | None: """Convert a database record to a VerificationToken model.""" if record is None: return None @@ -46,42 +65,43 @@ class VerificationTokenRepository(BaseRepository[LiteLLM_VerificationToken]): "litellm_budget_table", ] for field in json_fields: - if isinstance(data.get(field), str): - data[field] = json.loads(data[field]) + value = data.get(field) + if isinstance(value, str): + data[field] = json.loads(value) if data.get("org_id") is None and data.get("organization_id") is not None: data["org_id"] = data["organization_id"] - return LiteLLM_VerificationToken(**data) + return LiteLLM_VerificationToken.model_validate(data) - async def find_by_id(self, token: str, id_field: str = "token") -> Optional[LiteLLM_VerificationToken]: + async def find_by_id(self, token: str, id_field: str = "token") -> LiteLLM_VerificationToken | None: return await super().find_by_id(token, id_field) - async def find_by_alias(self, key_alias: str) -> Optional[LiteLLM_VerificationToken]: + async def find_by_alias(self, key_alias: str) -> LiteLLM_VerificationToken | None: """Find a token by key alias.""" - records = await self.table.find_many(where={"key_alias": key_alias}) + records: list[PrismaVerificationToken] = await self.table.find_many(where={"key_alias": key_alias}) if records: return self._to_model(records[0]) return None - async def find_by_user_id(self, user_id: str) -> List[LiteLLM_VerificationToken]: + async def find_by_user_id(self, user_id: str) -> list[LiteLLM_VerificationToken]: """Find all tokens belonging to a user.""" - records = await self.table.find_many(where={"user_id": user_id}) + records: list[PrismaVerificationToken] = await self.table.find_many(where={"user_id": user_id}) return self._to_model_list(records) - async def find_by_team_id(self, team_id: str) -> List[LiteLLM_VerificationToken]: + async def find_by_team_id(self, team_id: str) -> list[LiteLLM_VerificationToken]: """Find all tokens belonging to a team.""" - records = await self.table.find_many(where={"team_id": team_id}) + records: list[PrismaVerificationToken] = await self.table.find_many(where={"team_id": team_id}) return self._to_model_list(records) - async def find_by_project_id(self, project_id: str) -> List[LiteLLM_VerificationToken]: + async def find_by_project_id(self, project_id: str) -> list[LiteLLM_VerificationToken]: """Find all tokens belonging to a project.""" - records = await self.table.find_many(where={"project_id": project_id}) + records: list[PrismaVerificationToken] = await self.table.find_many(where={"project_id": project_id}) return self._to_model_list(records) - async def find_active_tokens(self) -> List[LiteLLM_VerificationToken]: + async def find_active_tokens(self) -> list[LiteLLM_VerificationToken]: """Find all active (non-expired, non-blocked) tokens.""" - records = await self.table.find_many( + records: list[PrismaVerificationToken] = await self.table.find_many( where={ "blocked": {"not": True}, "OR": [{"expires": None}, {"expires": {"gt": datetime.utcnow()}}], @@ -92,31 +112,31 @@ class VerificationTokenRepository(BaseRepository[LiteLLM_VerificationToken]): def _build_token_data( self, token: str, - key_name: Optional[str] = None, - key_alias: Optional[str] = None, - max_budget: Optional[float] = None, - expires: Optional[datetime] = None, - models: Optional[List[str]] = None, - aliases: Optional[Dict[str, str]] = None, - config: Optional[Dict[str, Any]] = None, - user_id: Optional[str] = None, - team_id: Optional[str] = None, - agent_id: Optional[str] = None, - project_id: Optional[str] = None, - max_parallel_requests: Optional[int] = None, - metadata: Optional[Dict[str, Any]] = None, - tpm_limit: Optional[int] = None, - rpm_limit: Optional[int] = None, - budget_duration: Optional[str] = None, - allowed_cache_controls: Optional[List[str]] = None, - allowed_routes: Optional[List[str]] = None, - permissions: Optional[Dict[str, Any]] = None, - org_id: Optional[str] = None, - created_by: Optional[str] = None, - object_permission_id: Optional[str] = None, - access_group_ids: Optional[List[str]] = None, - budget_id: Optional[str] = None, - ) -> Dict[str, Any]: + key_name: str | None = None, + key_alias: str | None = None, + max_budget: float | None = None, + expires: datetime | None = None, + models: list[str] | None = None, + aliases: dict[str, str] | None = None, + config: Mapping[str, object] | None = None, + user_id: str | None = None, + team_id: str | None = None, + agent_id: str | None = None, + project_id: str | None = None, + max_parallel_requests: int | None = None, + metadata: Mapping[str, object] | None = None, + tpm_limit: int | None = None, + rpm_limit: int | None = None, + budget_duration: str | None = None, + allowed_cache_controls: list[str] | None = None, + allowed_routes: list[str] | None = None, + permissions: Mapping[str, object] | None = None, + org_id: str | None = None, + created_by: str | None = None, + object_permission_id: str | None = None, + access_group_ids: list[str] | None = None, + budget_id: str | None = None, + ) -> dict[str, object]: """Build data dictionary for token creation.""" json_fields = { "aliases": aliases, @@ -145,7 +165,7 @@ class VerificationTokenRepository(BaseRepository[LiteLLM_VerificationToken]): "access_group_ids": access_group_ids, "budget_id": budget_id, } - data: Dict[str, Any] = {k: v for k, v in simple_fields.items() if v is not None} + data: dict[str, object] = {k: v for k, v in simple_fields.items() if v is not None} for key, val in json_fields.items(): if val is not None: data[key] = json.dumps(val) @@ -159,30 +179,30 @@ class VerificationTokenRepository(BaseRepository[LiteLLM_VerificationToken]): async def create_token( self, token: str, - key_name: Optional[str] = None, - key_alias: Optional[str] = None, - max_budget: Optional[float] = None, - expires: Optional[datetime] = None, - models: Optional[List[str]] = None, - aliases: Optional[Dict[str, str]] = None, - config: Optional[Dict[str, Any]] = None, - user_id: Optional[str] = None, - team_id: Optional[str] = None, - agent_id: Optional[str] = None, - project_id: Optional[str] = None, - max_parallel_requests: Optional[int] = None, - metadata: Optional[Dict[str, Any]] = None, - tpm_limit: Optional[int] = None, - rpm_limit: Optional[int] = None, - budget_duration: Optional[str] = None, - allowed_cache_controls: Optional[List[str]] = None, - allowed_routes: Optional[List[str]] = None, - permissions: Optional[Dict[str, Any]] = None, - org_id: Optional[str] = None, - created_by: Optional[str] = None, - object_permission_id: Optional[str] = None, - access_group_ids: Optional[List[str]] = None, - budget_id: Optional[str] = None, + key_name: str | None = None, + key_alias: str | None = None, + max_budget: float | None = None, + expires: datetime | None = None, + models: list[str] | None = None, + aliases: dict[str, str] | None = None, + config: Mapping[str, object] | None = None, + user_id: str | None = None, + team_id: str | None = None, + agent_id: str | None = None, + project_id: str | None = None, + max_parallel_requests: int | None = None, + metadata: Mapping[str, object] | None = None, + tpm_limit: int | None = None, + rpm_limit: int | None = None, + budget_duration: str | None = None, + allowed_cache_controls: list[str] | None = None, + allowed_routes: list[str] | None = None, + permissions: Mapping[str, object] | None = None, + org_id: str | None = None, + created_by: str | None = None, + object_permission_id: str | None = None, + access_group_ids: list[str] | None = None, + budget_id: str | None = None, ) -> LiteLLM_VerificationToken: """Create a new verification token.""" data = self._build_token_data( @@ -217,28 +237,28 @@ class VerificationTokenRepository(BaseRepository[LiteLLM_VerificationToken]): async def update_token( self, token: str, - updated_by: Optional[str] = None, - key_name: Optional[str] = None, - key_alias: Optional[str] = None, - max_budget: Optional[float] = None, - expires: Optional[datetime] = None, - models: Optional[List[str]] = None, - aliases: Optional[Dict[str, str]] = None, - config: Optional[Dict[str, Any]] = None, - max_parallel_requests: Optional[int] = None, - metadata: Optional[Dict[str, Any]] = None, - tpm_limit: Optional[int] = None, - rpm_limit: Optional[int] = None, - budget_duration: Optional[str] = None, - allowed_cache_controls: Optional[List[str]] = None, - allowed_routes: Optional[List[str]] = None, - permissions: Optional[Dict[str, Any]] = None, - blocked: Optional[bool] = None, - object_permission_id: Optional[str] = None, - access_group_ids: Optional[List[str]] = None, - ) -> Optional[LiteLLM_VerificationToken]: + updated_by: str | None = None, + key_name: str | None = None, + key_alias: str | None = None, + max_budget: float | None = None, + expires: datetime | None = None, + models: list[str] | None = None, + aliases: dict[str, str] | None = None, + config: Mapping[str, object] | None = None, + max_parallel_requests: int | None = None, + metadata: Mapping[str, object] | None = None, + tpm_limit: int | None = None, + rpm_limit: int | None = None, + budget_duration: str | None = None, + allowed_cache_controls: list[str] | None = None, + allowed_routes: list[str] | None = None, + permissions: Mapping[str, object] | None = None, + blocked: bool | None = None, + object_permission_id: str | None = None, + access_group_ids: list[str] | None = None, + ) -> LiteLLM_VerificationToken | None: """Update a verification token.""" - data: Dict[str, Any] = {} + data: dict[str, object] = {} if updated_by is not None: data["updated_by"] = updated_by if key_name is not None: @@ -283,10 +303,10 @@ class VerificationTokenRepository(BaseRepository[LiteLLM_VerificationToken]): async def delete_token( self, token: str, - deleted_by: Optional[str] = None, - deleted_by_api_key: Optional[str] = None, - litellm_changed_by: Optional[str] = None, - ) -> Optional[LiteLLM_VerificationToken]: + deleted_by: str | None = None, + deleted_by_api_key: str | None = None, + litellm_changed_by: str | None = None, + ) -> LiteLLM_VerificationToken | None: """Delete a token and archive it to the deleted tokens table. Uses a transaction to ensure atomicity of the archive-then-delete operation. @@ -307,14 +327,14 @@ class VerificationTokenRepository(BaseRepository[LiteLLM_VerificationToken]): return token_record - def _build_archive_data(self, token: LiteLLM_VerificationToken) -> Dict[str, Any]: + def _build_archive_data(self, token: LiteLLM_VerificationToken) -> dict[str, object]: """Build archive data with only columns present in LiteLLM_DeletedVerificationToken. Serializes JSON columns to strings (the archive table stores them as JSON columns the same way the live table does) and maps ``org_id`` onto the ``organization_id`` column so the foreign key is preserved. """ - data = token.model_dump(exclude_none=True) + data: dict[str, object] = token.model_dump(exclude_none=True) for field in ("object_permission", "litellm_budget_table", "budget_limits"): data.pop(field, None) @@ -336,24 +356,24 @@ class VerificationTokenRepository(BaseRepository[LiteLLM_VerificationToken]): data[field] = json.dumps(data[field]) return data - async def update_spend(self, token: str, spend: float) -> Optional[LiteLLM_VerificationToken]: + async def update_spend(self, token: str, spend: float) -> LiteLLM_VerificationToken | None: """Update token spend.""" return await self.update(token, {"spend": spend}, id_field="token") - async def update_last_active(self, token: str) -> Optional[LiteLLM_VerificationToken]: + async def update_last_active(self, token: str) -> LiteLLM_VerificationToken | None: """Update the last_active timestamp.""" return await self.update(token, {"last_active": datetime.utcnow()}, id_field="token") - async def block_token(self, token: str, updated_by: Optional[str] = None) -> Optional[LiteLLM_VerificationToken]: + async def block_token(self, token: str, updated_by: str | None = None) -> LiteLLM_VerificationToken | None: """Block a token.""" - data: Dict[str, Any] = {"blocked": True} + data: dict[str, object] = {"blocked": True} if updated_by is not None: data["updated_by"] = updated_by return await self.update(token, data, id_field="token") - async def unblock_token(self, token: str, updated_by: Optional[str] = None) -> Optional[LiteLLM_VerificationToken]: + async def unblock_token(self, token: str, updated_by: str | None = None) -> LiteLLM_VerificationToken | None: """Unblock a token.""" - data: Dict[str, Any] = {"blocked": False} + data: dict[str, object] = {"blocked": False} if updated_by is not None: data["updated_by"] = updated_by return await self.update(token, data, id_field="token") diff --git a/litellm/router.py b/litellm/router.py index 78fe3ff025e..487d6a31226 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -7702,6 +7702,21 @@ class Router: """True when this deployment opts in via the `auto_router/adaptive_router` model prefix.""" return litellm_params.model.startswith("auto_router/adaptive_router") + def _deployment_participates_in_adaptive_routing(self, litellm_params: LiteLLM_Params) -> bool: + """True when this deployment owns an `adaptive_routers` entry once finalized: + a dedicated adaptive router, or a complexity router whose config enables the + adaptive companion. Mirrors the two arms of + `_finalize_adaptive_router_if_configured`, which is the registry's only writer.""" + if self._is_adaptive_router_deployment(litellm_params=litellm_params): + return True + if not self._is_complexity_router_deployment(litellm_params=litellm_params): + return False + config = litellm_params.complexity_router_config + if not config: + return False + adaptive_flag: object = config.get("adaptive") + return bool(adaptive_flag) + @staticmethod def _has_registered_strategy( registry: dict[str, list[TaggedPreRoutingStrategy[_PreRoutingStrategyT]]], @@ -7734,20 +7749,56 @@ class Router: TaggedPreRoutingStrategy(tags=tags, strategy=strategy), ] + @staticmethod + def _unregister_pre_routing_strategy( + registry: dict[str, list[TaggedPreRoutingStrategy[_PreRoutingStrategyT]]], + model_name: str, + tags: tuple[str, ...], + ) -> bool: + """Drop the strategy registered for this exact (model_name, tags) pair, leaving + strategies registered under the same name with different tags in place. Returns + whether anything was actually dropped.""" + existing = registry.get(model_name, []) + remaining = [entry for entry in existing if entry.tags != tags] + if len(remaining) == len(existing): + return False + if remaining: + registry[model_name] = remaining + else: + registry.pop(model_name, None) + return True + + def _unregister_pre_routing_strategy_for_deployment(self, deployment: Deployment) -> None: + """ + Release the pre-routing strategy a deployment holds, so removing it from the + model_list also frees its (model_name, tags) slot. + + Without this, re-adding the deployment (an edit arriving via upsert_deployment, + or a router recreated under a name that was deleted earlier) hits the + "already exists" guard in `_register_pre_routing_strategy`, which + `ignore_invalid_deployments` swallows - the deployment then silently never + makes it back into the model_list. + + Released from every registry rather than the first match, because registration is + one-to-many: a complexity router configured with `adaptive` is also registered in + `adaptive_routers` under the same (model_name, tags) by the deferred finalize pass. + Guarded on the auto_router/ prefix so removing a *regular* deployment can't evict a + router that merely shares its model_name. + """ + if not deployment.litellm_params.model.startswith("auto_router/"): + return + model_name = deployment.model_name + tags = self._deployment_tags(deployment) + for registry in (self.auto_routers, self.complexity_routers, self.quality_routers): + self._unregister_pre_routing_strategy(registry, model_name, tags) + if self._unregister_pre_routing_strategy(self.adaptive_routers, model_name, tags): + self._sync_adaptive_router_hooks() + def _finalize_adaptive_router_if_configured(self) -> None: """Locate every adaptive-router deployment in the finalized model_list and build an AdaptiveRouter for each. Safe no-op when none are configured. Idempotent: skips any deployment whose (model_name, tags) pair is already initialized, so hot-reloads don't rebuild routers that would lose state.""" - # Drop any adaptive-router hooks left over from a previous Router - # instance (e.g. after `/config/reload` replaced `llm_router`). Without - # this, stale AdaptiveRouterPostCallHook callbacks from the old Router - # remain wired up in `litellm.callbacks` and double-fire signal - # recording for every request. - from litellm.router_strategy.adaptive_router.hooks import ( - AdaptiveRouterPostCallHook, - ) - for entry in self.model_list or []: lp = entry.get("litellm_params") if isinstance(entry, dict) else entry.litellm_params lp_model = (lp.get("model") if isinstance(lp, dict) else lp.model) if lp else None @@ -7779,6 +7830,16 @@ class Router: TaggedPreRoutingStrategy(tags=tagged.tags, strategy=adaptive_router), ] + self._sync_adaptive_router_hooks() + + def _sync_adaptive_router_hooks(self) -> None: + """Rebuild the AdaptiveRouterPostCallHook set so it is exactly one hook per + currently registered adaptive router. Run at every point the adaptive registry + changes, otherwise a released router keeps recording turns through its hook.""" + from litellm.router_strategy.adaptive_router.hooks import ( + AdaptiveRouterPostCallHook, + ) + for callback in litellm.logging_callback_manager.get_custom_loggers_for_type(AdaptiveRouterPostCallHook): litellm.logging_callback_manager.remove_callback_from_all_lists(callback) for tagged_adaptive_routers in self.adaptive_routers.values(): @@ -8401,13 +8462,29 @@ class Router: self._invalidate_access_groups_cache() self._update_deployment_indices_after_removal(model_id=deployment_id, removal_idx=removal_idx) + # Free the outgoing deployment's pre-routing strategy slot (keyed by the + # OLD model_name/tags) before the re-add below re-registers it. + self._unregister_pre_routing_strategy_for_deployment(deployment=_deployment_on_router) + # if the model_id is not in router self.add_deployment(deployment=deployment) + # add_deployment() builds every strategy EXCEPT the adaptive one, which + # set_model_list() defers until the whole model_list is visible. Re-run that + # deferred pass so an adaptive router whose slot was just released above is + # rebuilt rather than left unregistered. + if self._deployment_participates_in_adaptive_routing(litellm_params=deployment.litellm_params) or ( + _deployment_on_router is not None + and self._deployment_participates_in_adaptive_routing( + litellm_params=_deployment_on_router.litellm_params + ) + ): + self._finalize_adaptive_router_if_configured() return deployment except Exception as e: if self.ignore_invalid_deployments: - verbose_router_logger.debug( - f"Error upserting deployment: {e}, ignoring and continuing with other deployments." + verbose_router_logger.warning( + f"Error upserting deployment {deployment.model_name} (id={deployment.model_info.id}): {e}. " + "Dropping it and continuing with other deployments." ) return None else: @@ -8436,6 +8513,16 @@ class Router: _budget_limiter = self._get_router_deployment_budget_limiter() if _budget_limiter is not None: _budget_limiter.unregister_deployment_budget(model_id=id) + try: + self._unregister_pre_routing_strategy_for_deployment( + deployment=item if isinstance(item, Deployment) else Deployment(**item) + ) + except Exception: + verbose_router_logger.exception( + "delete_deployment: could not release pre-routing strategies for model_id=%s; " + "the deployment is out of the model_list and its indices are repaired", + id, + ) return item else: return None diff --git a/litellm/types/tool_management.py b/litellm/types/tool_management.py index 71ec412e8ef..ccf4b7dbc9f 100644 --- a/litellm/types/tool_management.py +++ b/litellm/types/tool_management.py @@ -124,12 +124,5 @@ class ToolSpendDailyEntry(BaseModel): class ToolSpendResponse(BaseModel): by_tool: List[ToolSpendEntry] = Field(default_factory=list) daily: List[ToolSpendDailyEntry] = Field(default_factory=list) - total_spend: float = Field( - 0.0, - description=( - "Deduplicated spend of every request that called at least one tool in the window; " - "less than the sum of per-tool attributed spend whenever multi-tool requests exist" - ), - ) start_date: str | None = None end_date: str | None = None diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index 437d09b726b..addee5fc68a 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -1,12 +1,12 @@ { "ANN001": { - "limit": 3152 + "limit": 3142 }, "ANN002": { "limit": 69 }, "ANN003": { - "limit": 835 + "limit": 831 }, "ANN201": { "limit": 2138 @@ -24,7 +24,7 @@ "limit": 130 }, "ANN401": { - "limit": 2074 + "limit": 2015 }, "ASYNC230": { "limit": 14 @@ -123,7 +123,7 @@ "limit": 52 }, "I001": { - "limit": 273 + "limit": 270 }, "LOG015": { "limit": 8 @@ -135,7 +135,7 @@ "limit": 30 }, "PERF401": { - "limit": 146 + "limit": 144 }, "PERF402": { "limit": 9 @@ -222,7 +222,7 @@ "limit": 38 }, "RET504": { - "limit": 719 + "limit": 717 }, "RUF010": { "limit": 874 @@ -237,7 +237,7 @@ "limit": 41 }, "RUF022": { - "limit": 84 + "limit": 85 }, "RUF023": { "limit": 5 @@ -306,7 +306,7 @@ "limit": 9 }, "TID251": { - "limit": 2700 + "limit": 2652 }, "TRY002": { "limit": 548 @@ -324,10 +324,10 @@ "limit": 883 }, "UP006": { - "limit": 12789 + "limit": 12147 }, "UP007": { - "limit": 2570 + "limit": 2526 }, "UP008": { "limit": 5 @@ -354,7 +354,7 @@ "limit": 4 }, "UP035": { - "limit": 2284 + "limit": 2232 }, "UP036": { "limit": 4 @@ -363,6 +363,6 @@ "limit": 105 }, "UP045": { - "limit": 18458 + "limit": 17824 } } diff --git a/schema.prisma b/schema.prisma index 314c59d3560..ed4f7d4f0a0 100644 --- a/schema.prisma +++ b/schema.prisma @@ -1102,6 +1102,19 @@ model LiteLLM_SpendLogToolIndex { @@index([start_time]) } +// Daily tool spend rollup (one row per tool per day) – the Cost Optimization card reads this, never SpendLogs +model LiteLLM_DailyToolSpend { + date String + tool_name String + spend Float @default(0.0) + total_tokens BigInt @default(0) + request_count BigInt @default(0) + created_at DateTime @default(now()) + updated_at DateTime @updatedAt + + @@id([date, tool_name]) +} + // Prompt table for storing prompt configurations model LiteLLM_PromptTable { id String @id @default(uuid()) diff --git a/tests/e2e/a2a/a2a_client.py b/tests/e2e/a2a/a2a_client.py index 916ef623d3a..e83897025a3 100644 --- a/tests/e2e/a2a/a2a_client.py +++ b/tests/e2e/a2a/a2a_client.py @@ -11,12 +11,13 @@ here because only this suite uses them. from __future__ import annotations +import time import warnings from dataclasses import dataclass from pydantic import BaseModel, ConfigDict, Field -from e2e_http import NoBody, Result, get_external, is_ok +from e2e_http import NoBody, Result, Success, get_external, is_ok from proxy_client import ProxyClient @@ -290,12 +291,46 @@ class A2AClient: proxy: ProxyClient def register_agent(self, body: AgentRegisterBody) -> Result[AgentResponse]: - return self.proxy.transport.post( + """Register an agent and, on success, wait until the data plane serves it. + + /v1/agents is a control-plane route; the /a2a/{agent_id} routes that serve + the card and run message/send are data plane, and only see the agent after + the next DB reload. A card read or message/send issued the instant this + returns can therefore 404 on the agent it just created. Waiting here keeps + every caller from having to poll, the same way ProxyClient.create_model + waits for a new model to become servable. + """ + result = self.proxy.transport.post( "/v1/agents", headers=self.proxy.transport.master, json=body, response_type=AgentResponse, ) + if isinstance(result, Success): + self._await_agent_servable(result.data.agent_id) + return result + + def _await_agent_servable(self, agent_id: str) -> None: + """Block until the data plane serves `agent_id`'s card, or fail loudly at + poll_timeout (a real propagation problem, surfaced here rather than as a + downstream 404 on whichever /a2a call the test happened to make first).""" + deadline = time.monotonic() + self.proxy.poll_timeout + while True: + result = self.proxy.transport.get( + f"/a2a/{agent_id}/.well-known/agent-card.json", + headers=self.proxy.transport.master, + params=NoBody(), + response_type=ServedAgentCard, + ) + if isinstance(result, Success): + return + if time.monotonic() >= deadline: + raise AssertionError( + f"agent {agent_id!r} was registered but never became servable on the " + f"data plane within {self.proxy.poll_timeout}s of POST /v1/agents " + f"(control/data-plane propagation issue); last card read: {result}" + ) + time.sleep(self.proxy.poll_interval) def get_agent(self, agent_id: str) -> Result[AgentResponse]: return self.proxy.transport.get( diff --git a/tests/e2e/coverage_registry/llm_conversational.yaml b/tests/e2e/coverage_registry/llm_conversational.yaml index 6cb60c19247..e8fc8067ee0 100644 --- a/tests/e2e/coverage_registry/llm_conversational.yaml +++ b/tests/e2e/coverage_registry/llm_conversational.yaml @@ -5,6 +5,9 @@ - {id: llm.chat_completions.openai.passthrough.nonstream.cost_logged, module: llm, tier: P1, subject_endpoint: chat_completions, route: openai, capability: basic, streaming: nonstream, assertions: [works, cost_logged], source: "test_passthrough_e2e.py", rationale: "OpenAI-format chat via the raw /openai/{endpoint} passthrough (/openai/v1/chat/completions); proxy swaps in OPENAI_API_KEY and still logs a costed pass_through_endpoint row (LIT-4752)"} - {id: llm.chat_completions.openai.tool_use.nonstream.works, module: llm, tier: P0, subject_endpoint: chat_completions, route: openai, capability: tool_use, streaming: nonstream, assertions: [works], source: "model_prices json", rationale: "OpenAI function_calling; high usage"} - {id: llm.chat_completions.openai.tool_use.stream.works, module: llm, tier: P0, subject_endpoint: chat_completions, route: openai, capability: tool_use, streaming: stream, assertions: [works], source: "model_prices json", rationale: "Tool calls over streaming"} +- {id: llm.chat_completions.openai.basic.stream.bridge_shares_chunk_id, module: llm, tier: P0, subject_endpoint: chat_completions, route: openai, capability: basic, streaming: stream, assertions: [stable_chunk_id], source: "completion_extras/litellm_responses_transformation/transformation.py", rationale: "A responses-only model served over /chat/completions must stream every chunk under one chat completion id; per-chunk ids make id-accumulating SDKs drop the response", fail_before_fix: proven} +- {id: llm.chat_completions.openai.basic.stream.bridge_streams_sse, module: llm, tier: P0, subject_endpoint: chat_completions, route: openai, capability: basic, streaming: stream, assertions: [works], source: "completion_extras/litellm_responses_transformation/handler.py", rationale: "The Responses bridge must answer a streaming chat request with real SSE (content deltas, finish_reason, [DONE]), never a completed response the SSE generator cannot iterate"} +- {id: llm.chat_completions.openai.tool_use.stream.bridge_streams_tool_call, module: llm, tier: P0, subject_endpoint: chat_completions, route: openai, capability: tool_use, streaming: stream, assertions: [works], source: "completion_extras/litellm_responses_transformation/transformation.py", rationale: "Tool calls translated from Responses events must reassemble into one named call with parseable argument JSON over the bridged stream"} - {id: llm.chat_completions.openai.vision.nonstream.works, module: llm, tier: P0, subject_endpoint: chat_completions, route: openai, capability: vision, streaming: nonstream, assertions: [works], source: "model_prices json", rationale: "gpt-4o vision; high usage"} - {id: llm.chat_completions.openai.prompt_cache_5m.nonstream.works, module: llm, tier: P0, subject_endpoint: chat_completions, route: openai, capability: prompt_cache_5m, streaming: nonstream, assertions: [works], source: "model_prices json", rationale: "Prompt caching cost optimization"} - {id: llm.chat_completions.openai.service_tier.nonstream.works, module: llm, tier: P1, subject_endpoint: chat_completions, route: openai, capability: service_tier, streaming: nonstream, assertions: [works], source: "OpenAI service_tier param", rationale: "OpenAI scale-tier request option is forwarded and echoed"} diff --git a/tests/e2e/e2e_http.py b/tests/e2e/e2e_http.py index d16e84bd754..74e57f86b88 100644 --- a/tests/e2e/e2e_http.py +++ b/tests/e2e/e2e_http.py @@ -137,6 +137,7 @@ class StreamingResponse(BaseModel): # quota) arrive as SSE error events inside an otherwise-successful response; # the consumed body is elided, so this is the only place they surface. stream_error: str | None = None + stream_done: bool = False @property def ok(self) -> bool: @@ -408,6 +409,7 @@ def _streaming_outcome(resp: requests.Response, stream: bool) -> StreamingRespon chunks = 0 stream_error: str | None = None stream_events: list[str] = [] + stream_done = False for line in lines: if not line: continue @@ -415,7 +417,9 @@ def _streaming_outcome(resp: requests.Response, stream: bool) -> StreamingRespon decoded_line = line.decode(errors="replace") if decoded_line.startswith("data: "): payload = decoded_line.removeprefix("data: ") - if payload != "[DONE]": + if payload == "[DONE]": + stream_done = True + else: stream_events.append(payload) if stream_error is None and ( line.startswith(b"event: error") @@ -433,6 +437,7 @@ def _streaming_outcome(resp: requests.Response, stream: bool) -> StreamingRespon body="", chunks=chunks, stream_events=stream_events, + stream_done=stream_done, stream_error=stream_error, ) diff --git a/tests/e2e/guardrails/guardrails_client.py b/tests/e2e/guardrails/guardrails_client.py index d56e4e9311a..5a54a4f0bbc 100644 --- a/tests/e2e/guardrails/guardrails_client.py +++ b/tests/e2e/guardrails/guardrails_client.py @@ -5,6 +5,7 @@ and chat through them on the shared ProxyClient so resources.defer cleans up. from __future__ import annotations import time +from collections.abc import Callable from dataclasses import dataclass from typing import Literal @@ -287,3 +288,24 @@ class GuardrailsClient: def build_client(proxy: ProxyClient) -> GuardrailsClient: return GuardrailsClient(proxy=proxy) + + +def poll_until_blocked(call: Callable[[], Result[ChatResponse]]) -> Result[ChatResponse]: + """Retry a call that a guardrail should reject until it is, returning the last result. + + Registering a guardrail is a control-plane write; the data-plane worker that + serves /chat/completions picks it up only on its next periodic DB sync (~30s in + proxy_server.py). A call issued right after the create therefore runs against a + worker that has no guardrail yet and is allowed through, which is in-flight + propagation rather than a guardrail that failed to block. Polling to the deadline + waits that out so the assertions judge the synced state; a guardrail that never + blocks still fails, on the last allowed result. + """ + deadline = time.monotonic() + POLL_TIMEOUT + last = call() + while time.monotonic() < deadline: + if not isinstance(last, Success): + return last + time.sleep(POLL_INTERVAL) + last = call() + return last diff --git a/tests/e2e/guardrails/test_bedrock_guardrail_e2e.py b/tests/e2e/guardrails/test_bedrock_guardrail_e2e.py index ba3c5071cbb..dd61e630d7d 100644 --- a/tests/e2e/guardrails/test_bedrock_guardrail_e2e.py +++ b/tests/e2e/guardrails/test_bedrock_guardrail_e2e.py @@ -18,7 +18,7 @@ import pytest from e2e_config import unique_marker from e2e_http import UnknownApiError -from guardrails_client import GuardrailsClient +from guardrails_client import GuardrailsClient, poll_until_blocked from lifecycle import ResourceManager pytestmark = pytest.mark.e2e @@ -50,7 +50,9 @@ class TestBedrockGuardrail: # Selected per request rather than registered default_on, so an upstream # ApplyGuardrail failure surfaces here instead of 403ing every other suite # running against this proxy. - result = client.chat(scoped_key, MODEL, BLOCKED_PROMPT, guardrails=[name]) + result = poll_until_blocked( + lambda: client.chat(scoped_key, MODEL, BLOCKED_PROMPT, guardrails=[name]) + ) match result: case UnknownApiError(status_code=status, body=body): diff --git a/tests/e2e/guardrails/test_block_code_execution_guardrail_e2e.py b/tests/e2e/guardrails/test_block_code_execution_guardrail_e2e.py index de087b190d0..7cf4c195424 100644 --- a/tests/e2e/guardrails/test_block_code_execution_guardrail_e2e.py +++ b/tests/e2e/guardrails/test_block_code_execution_guardrail_e2e.py @@ -14,9 +14,11 @@ the shared proxy, and the chat backend is a gemini deployment created for the te from __future__ import annotations +import time + import pytest -from e2e_config import unique_marker +from e2e_config import POLL_INTERVAL, POLL_TIMEOUT, unique_marker from e2e_http import unwrap from guardrails_client import BlockCodeExecutionParamsBody, GuardrailsClient from lifecycle import ResourceManager @@ -54,7 +56,18 @@ class TestBlockCodeExecutionGuardrail: ) resources.defer(lambda: client.delete_guardrail(guardrail_id)) + # This guardrail replaces the reply rather than erroring, so wait for the + # block marker to appear instead of for a non-success status. The data-plane + # worker only picks a new guardrail up on its next DB sync (~30s), so the + # first call after the create is served without it. + deadline = time.monotonic() + POLL_TIMEOUT blocked = unwrap(client.chat(scoped_key, model, EXECUTION_REQUEST, guardrails=[name])) + while time.monotonic() < deadline: + if _BLOCK_MARKER in _first_content(blocked).lower(): + break + time.sleep(POLL_INTERVAL) + blocked = unwrap(client.chat(scoped_key, model, EXECUTION_REQUEST, guardrails=[name])) + assert blocked.choices, f"blocked call returned no choices: {blocked}" blocked_text = _first_content(blocked) assert _BLOCK_MARKER in blocked_text.lower(), ( diff --git a/tests/e2e/guardrails/test_openai_moderation_guardrail_e2e.py b/tests/e2e/guardrails/test_openai_moderation_guardrail_e2e.py index 39950259fb5..d117832221d 100644 --- a/tests/e2e/guardrails/test_openai_moderation_guardrail_e2e.py +++ b/tests/e2e/guardrails/test_openai_moderation_guardrail_e2e.py @@ -16,7 +16,11 @@ import pytest from e2e_config import unique_marker from e2e_http import UnknownApiError, unwrap -from guardrails_client import GuardrailsClient, OpenAIModerationParamsBody +from guardrails_client import ( + GuardrailsClient, + OpenAIModerationParamsBody, + poll_until_blocked, +) from lifecycle import ResourceManager pytestmark = pytest.mark.e2e @@ -45,7 +49,9 @@ class TestOpenAIModerationGuardrail: ) resources.defer(lambda: client.delete_guardrail(guardrail_id)) - blocked = client.chat(scoped_key, model, FLAGGED_PROMPT, guardrails=[name]) + blocked = poll_until_blocked( + lambda: client.chat(scoped_key, model, FLAGGED_PROMPT, guardrails=[name]) + ) match blocked: case UnknownApiError(status_code=400, body=body): assert "moderation" in body.lower(), ( diff --git a/tests/e2e/guardrails/test_presidio_guardrail_e2e.py b/tests/e2e/guardrails/test_presidio_guardrail_e2e.py index d103714b1dd..9742dfc6ae7 100644 --- a/tests/e2e/guardrails/test_presidio_guardrail_e2e.py +++ b/tests/e2e/guardrails/test_presidio_guardrail_e2e.py @@ -1,8 +1,8 @@ -"""Live e2e: the built-in Presidio PII guardrail masks PII on the request, on the -model output, and in what the proxy logs. +"""Live e2e: the built-in Presidio PII guardrail masks PII on the request and on +the model output. Presidio replaces detected PII with `` placeholders (e.g. -``) via a real analyzer + anonymizer. Three modes are checked +``) via a real analyzer + anonymizer. Two modes are checked independently, each opted into per request (default_on=False) so it never touches unrelated traffic: @@ -10,32 +10,31 @@ unrelated traffic: repeat-verbatim request comes back with the placeholder, never the raw email - post_call (apply_to_output): PII the model itself emits is masked on the way out, so the caller never receives the raw value the model produced -- logging_only: the call is not blocked, and the request the proxy records is - masked. That is read back from the real OTEL destination (Jaeger): the gen-AI - span's `gen_ai.input.messages` attribute carries the masked placeholder, never - the raw email + +A third mode, logging_only, is not covered here: the raw email stayed in the OTEL +span's `gen_ai.input.messages` on every attempt over a full poll deadline while +these two modes masked correctly, so that cell is tracked in LIT-4841 rather than +asserted against known-failing behavior. Analyzer/anonymizer bases come from PRESIDIO_ANALYZER_API_BASE / PRESIDIO_ANONYMIZER_API_BASE (compose provides the in-network hosts; point them at -locally published container ports for a host run). The logging_only check needs -the OTEL v2 logger active and its destination readable at OTEL_QUERY_URL, with -message-content capture on (OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT). -The chat backend is a gemini deployment created for the test. +locally published container ports for a host run). The chat backend is a gemini +deployment created for the test. """ from __future__ import annotations import os import time +from collections.abc import Callable import pytest from e2e_config import POLL_INTERVAL, POLL_TIMEOUT, unique_marker -from e2e_http import NoBody, require_successful_call, unwrap +from e2e_http import unwrap from guardrails_client import GuardrailMode, GuardrailsClient, PresidioParamsBody from lifecycle import ResourceManager -from models import ChatBody, ChatMessage, ChatResponse, ReadinessDetailsResponse -from otel_client import JaegerSpan, OtelReader, build_otel_reader +from models import ChatResponse pytestmark = pytest.mark.e2e @@ -44,10 +43,6 @@ PLACEHOLDER = "" ECHO_REQUEST = f"Repeat the following text back exactly, verbatim, with no changes: My email is {RAW_EMAIL}" EMIT_REQUEST = f"Output exactly this one line and nothing else: Please contact {RAW_EMAIL} today" -LOG_REQUEST = f"Say hello and include this email once verbatim: {RAW_EMAIL}" - -OTEL_V2_LOGGER = "OpenTelemetryV2" -INPUT_MESSAGES_TAG = "gen_ai.input.messages" def _content(response: ChatResponse) -> str: @@ -57,35 +52,6 @@ def _content(response: ChatResponse) -> str: return (message.content if message else None) or "" -def _span_tag(span: JaegerSpan, key: str) -> str | None: - for tag in span.tags: - if tag.key == key and isinstance(tag.value, str): - return tag.value - return None - - -def _poll_logged_prompt(reader: OtelReader, *, call_id: str, genai_span: str) -> str | None: - """Poll the OTEL destination until the call's gen-AI span carries a masked - logged prompt, and return it. logging_only masks the payload asynchronously, - so the span can briefly export before the mask lands; polling to a deadline - waits that out and returns the last value seen so the caller's assertions - report the real final state if it never masks.""" - deadline = time.monotonic() + POLL_TIMEOUT - last: str | None = None - while time.monotonic() < deadline: - for trace in reader.traces_for_call(call_id): - for span in trace.spans: - if span.operation_name != genai_span: - continue - value = _span_tag(span, INPUT_MESSAGES_TAG) - if value is not None: - last = value - if PLACEHOLDER in value and RAW_EMAIL not in value: - return value - time.sleep(POLL_INTERVAL) - return last - - def _presidio_params( mode: GuardrailMode, *, apply_to_output: bool = False, logging_only: bool = False ) -> PresidioParamsBody: @@ -101,19 +67,25 @@ def _presidio_params( ) -def _require_otel_v2_active(client: GuardrailsClient) -> None: - details = unwrap( - client.proxy.transport.get( - "/health/readiness/details", - headers=client.proxy.transport.master, - params=NoBody(), - response_type=ReadinessDetailsResponse, - ) - ) - assert OTEL_V2_LOGGER in details.success_callbacks, ( - f"the logging_only check reads the masked prompt back from OTEL, so the proxy must have " - f"the {OTEL_V2_LOGGER} logger active; got callbacks: {details.success_callbacks}" - ) +def _poll_until_masked(call: Callable[[], str]) -> str: + """Retry a call until the guardrail masks its PII, returning the last content. + + Registering a guardrail is a control-plane write; the data-plane worker that + serves /chat/completions only picks it up on its next periodic DB sync (~30s + in proxy_server.py), so a call issued the instant after the create runs + against a worker that has no guardrail yet and passes the raw value through. + That is in-flight propagation, not a masking failure. Polling to the deadline + waits it out, so the assertions that follow judge the synced state; if the + mask never lands the last unmasked content is returned and they still fail. + """ + deadline = time.monotonic() + POLL_TIMEOUT + last = call() + while time.monotonic() < deadline: + if PLACEHOLDER in last and RAW_EMAIL not in last: + return last + time.sleep(POLL_INTERVAL) + last = call() + return last class TestPresidioGuardrail: @@ -129,8 +101,10 @@ class TestPresidioGuardrail: guardrail_id = client.register(name, _presidio_params("pre_call")) resources.defer(lambda: client.delete_guardrail(guardrail_id)) - echoed = _content( - unwrap(client.chat(scoped_key, model, ECHO_REQUEST, guardrails=[name], max_tokens=128)) + echoed = _poll_until_masked( + lambda: _content( + unwrap(client.chat(scoped_key, model, ECHO_REQUEST, guardrails=[name], max_tokens=128)) + ) ) assert RAW_EMAIL not in echoed, ( "pre_call masking must strip the raw email before the model sees it, but the " @@ -153,8 +127,10 @@ class TestPresidioGuardrail: guardrail_id = client.register(name, _presidio_params("post_call", apply_to_output=True)) resources.defer(lambda: client.delete_guardrail(guardrail_id)) - out = _content( - unwrap(client.chat(scoped_key, model, EMIT_REQUEST, guardrails=[name], max_tokens=128)) + out = _poll_until_masked( + lambda: _content( + unwrap(client.chat(scoped_key, model, EMIT_REQUEST, guardrails=[name], max_tokens=128)) + ) ) assert RAW_EMAIL not in out, ( "post_call masking must strip PII the model emitted, but the raw email reached the " @@ -163,46 +139,3 @@ class TestPresidioGuardrail: assert PLACEHOLDER in out, ( f"the masked placeholder should replace the model's PII output, got: {out[:300]!r}" ) - - @pytest.mark.covers( - "guardrail.presidio.logging_only.masks", - exercised_on=["chat_completions"], - ) - def test_logging_only_masks_the_logged_prompt( - self, client: GuardrailsClient, resources: ResourceManager, scoped_key: str - ) -> None: - _require_otel_v2_active(client) - reader = build_otel_reader() - - model = client.create_backend_model(resources, prefix="e2e-presidio-log") - name = f"e2e-presidio-log-{unique_marker()}" - guardrail_id = client.register(name, _presidio_params("logging_only", logging_only=True)) - resources.defer(lambda: client.delete_guardrail(guardrail_id)) - - outcome = client.proxy.transport.send( - "/chat/completions", - headers=client.proxy.transport.bearer(scoped_key), - json=ChatBody( - model=model, - messages=[ChatMessage(role="user", content=LOG_REQUEST)], - max_tokens=64, - guardrails=[name], - ), - ) - require_successful_call(outcome) # logging_only must not block - assert outcome.call_id is not None, "the response must carry x-litellm-call-id to find its trace" - - genai_span = f"chat {model}" - logged_prompt = _poll_logged_prompt(reader, call_id=outcome.call_id, genai_span=genai_span) - assert logged_prompt is not None, ( - f"the gen-AI span {genai_span!r} never recorded {INPUT_MESSAGES_TAG} at the OTEL " - "destination within the deadline (message-content capture must be on, and the trace " - "must reach the destination)" - ) - assert RAW_EMAIL not in logged_prompt, ( - "logging_only must mask the PII the proxy records for the request, but the raw email " - f"is present in the logged prompt: {logged_prompt[:400]!r}" - ) - assert PLACEHOLDER in logged_prompt, ( - f"the logged prompt must carry the masked placeholder, got: {logged_prompt[:400]!r}" - ) diff --git a/tests/e2e/llm_translation/test_chat_completions_regression_e2e.py b/tests/e2e/llm_translation/test_chat_completions_regression_e2e.py index 6a69384d31a..655d426c28d 100644 --- a/tests/e2e/llm_translation/test_chat_completions_regression_e2e.py +++ b/tests/e2e/llm_translation/test_chat_completions_regression_e2e.py @@ -84,12 +84,6 @@ OPENAI_VISION_BACKEND = "openai/gpt-4o" # OpenAI caches a shared prompt prefix once it exceeds ~1024 tokens; this is well # past that, so a repeat call reports cached prompt tokens. -CACHE_PREFIX = ( - "You are a meticulous assistant. Follow these standing instructions exactly. " - * 300 -) - - def _vision_messages() -> list[ChatMessage]: return [ ChatMessage( @@ -582,36 +576,6 @@ class TestOpenAIChatCompletions: response = unwrap(client.proxy.chat(key, ChatBody(model=model, messages=_vision_messages(), max_tokens=32))) _assert_describes_cat(response) - @pytest.mark.covers( - "llm.chat_completions.openai.prompt_cache_5m.nonstream.works", - exercised_on=["chat_completions"], - ) - def test_openai_chat_prompt_cache_hits_on_repeat( - self, client: PassthroughClient, resources: ResourceManager - ) -> None: - model = f"e2e-openai-cache-{unique_marker()}" - model_id = client.proxy.create_model( - model, LiteLLMParamsBody(model=OPENAI_BACKEND, api_key="os.environ/OPENAI_API_KEY") - ) - resources.defer(lambda: client.proxy.delete_model(model_id)) - key = resources.key() - - body = ChatBody( - model=model, - messages=[ - ChatMessage(role="system", content=CACHE_PREFIX), - ChatMessage(role="user", content="Reply with the single word pong."), - ], - max_tokens=16, - ) - unwrap(client.proxy.chat(key, body)) - second = unwrap(client.proxy.chat(key, body)) - - details = second.usage.prompt_tokens_details if second.usage else None - assert details and details.cached_tokens and details.cached_tokens > 0, ( - f"a repeated large-prefix prompt must report cached prompt tokens, got usage={second.usage}" - ) - @pytest.mark.covers( "llm.chat_completions.openai.tool_use.stream.works", exercised_on=["chat_completions"], diff --git a/tests/e2e/llm_translation/test_passthrough_e2e.py b/tests/e2e/llm_translation/test_passthrough_e2e.py index b7d4d7cd668..ed5c657d23e 100644 --- a/tests/e2e/llm_translation/test_passthrough_e2e.py +++ b/tests/e2e/llm_translation/test_passthrough_e2e.py @@ -160,18 +160,6 @@ def test_anthropic_passthrough_tool_call_logs_cost( assert row.custom_llm_provider == "anthropic" -@pytest.mark.covers("llm.chat_completions.openai.passthrough.nonstream.cost_logged") -def test_openai_passthrough_nonstreaming_logs_cost( - client: PassthroughClient, scoped_key: str -) -> None: - result = client.openai_chat(scoped_key, "gpt-5.4-mini", "Say hello in one word") - require_successful_call(result) - - row = _fetch_cost_breakdown(client, result) - assert row.custom_llm_provider == "openai" - assert "gpt-5" in (row.model or "") - - class TestPassthroughModelAllowlist: """A passthrough route must honor the calling key's model allow-list. diff --git a/tests/e2e/llm_translation/test_passthrough_headers_e2e.py b/tests/e2e/llm_translation/test_passthrough_headers_e2e.py index 045988334d5..95d5d0c3f6a 100644 --- a/tests/e2e/llm_translation/test_passthrough_headers_e2e.py +++ b/tests/e2e/llm_translation/test_passthrough_headers_e2e.py @@ -13,6 +13,8 @@ this specific request's header - not a stale or cached one - got there. from __future__ import annotations +import time + import pytest from pydantic import BaseModel, Field @@ -78,9 +80,39 @@ def _create_passthrough(client: PassthroughClient, *, path: str) -> PassThroughE assert created.endpoints, "create returned no endpoints" endpoint = created.endpoints[0] assert endpoint.id, "created pass-through endpoint has no id" + _await_route_serving(client, path=path) return endpoint +def _await_route_serving(client: PassthroughClient, *, path: str) -> None: + """Block until the data plane routes `path`, instead of 404ing on it. + + POST /config/pass_through_endpoint is a control-plane write; the worker that + serves the route only registers it on its next config reload, so a call issued + right after the create gets a bare 404 that looks like a broken route rather + than in-flight propagation. Measured at ~18s on a live proxy. + """ + deadline = time.monotonic() + client.proxy.poll_timeout + while True: + # Any non-404 means the route is registered; this probe deliberately sends + # no anthropic-version so it is rejected upstream rather than billing a + # real completion on every poll. + result = client.proxy.transport.send( + path, + headers=client.proxy.transport.master, + json=_messages_body(), + ) + if result.status_code != 404: + return + if time.monotonic() >= deadline: + raise AssertionError( + f"pass-through route {path!r} was created but never became routable on the " + f"data plane within {client.proxy.poll_timeout}s (config reload issue); " + f"last status {result.status_code}: {result.body[:200]}" + ) + time.sleep(client.proxy.poll_interval) + + def _delete_passthrough(client: PassthroughClient, endpoint_id: str) -> None: _ = client.proxy.transport.delete( "/config/pass_through_endpoint", diff --git a/tests/e2e/llm_translation/test_responses_bridge_streaming_e2e.py b/tests/e2e/llm_translation/test_responses_bridge_streaming_e2e.py new file mode 100644 index 00000000000..9a45743a0cd --- /dev/null +++ b/tests/e2e/llm_translation/test_responses_bridge_streaming_e2e.py @@ -0,0 +1,173 @@ +"""Live /chat/completions streaming through the Responses API bridge. + +Responses-only models (gpt-5.3-codex here, the same shape as the GPT-5.6 models +customers reach over bedrock_mantle) cannot serve /chat/completions natively, so the +proxy translates the request to /v1/responses and translates each Responses event back +into a chat completion chunk. Two customer-visible contracts only hold on that path: + +- every chunk of one stream carries the same ``id`` (#32854). The bridge builds a chunk + per Responses event, so a regression there hands each chunk a fresh ``chatcmpl-`` + and SDKs that accumulate by id (openai-go's ChatCompletionAccumulator) silently drop + everything after the first chunk while the HTTP response still looks healthy +- the bridge always answers a streaming request with a real SSE stream (#33154). When it + hands back an already-completed response instead, the proxy's SSE generator dies with + "'async for' requires an object with __aiter__ method" mid-stream +""" + +from __future__ import annotations + +import pytest +from pydantic import BaseModel + +from e2e_config import unique_marker +from e2e_http import StreamingResponse +from lifecycle import ResourceManager +from models import ChatBody, ChatMessage, ChatTool, ChatToolFunction, LiteLLMParamsBody +from passthrough_client import PassthroughClient + +pytestmark = pytest.mark.e2e + +RESPONSES_ONLY_BACKEND = "openai/gpt-5.3-codex" + + +class _BridgeToolCallFunction(BaseModel): + name: str | None = None + arguments: str | None = None + + +class _BridgeToolCall(BaseModel): + function: _BridgeToolCallFunction = _BridgeToolCallFunction() + + +class _BridgeDelta(BaseModel): + content: str | None = None + tool_calls: list[_BridgeToolCall] | None = None + + +class _BridgeChoice(BaseModel): + delta: _BridgeDelta = _BridgeDelta() + finish_reason: str | None = None + + +class _BridgeChunk(BaseModel): + id: str + choices: list[_BridgeChoice] = [] + + +class _WeatherArgs(BaseModel): + location: str + + +_WEATHER_TOOL = ChatTool( + function=ChatToolFunction( + name="get_weather", + description="Get the current weather for a location", + parameters={ + "type": "object", + "properties": {"location": {"type": "string"}}, + "required": ["location"], + }, + ) +) + + +def _bridge_chunks(result: StreamingResponse) -> list[_BridgeChunk]: + """Parse the SSE events of a bridged stream, failing loudly on a stream that never + established, carried an error event, or delivered no chunks.""" + assert result.ok and result.is_streaming, f"bridged stream was not established: {result}" + assert result.stream_error is None, f"bridged stream carried an error event: {result.stream_error}" + chunks = [_BridgeChunk.model_validate_json(event) for event in result.stream_events] + assert chunks, f"bridged stream delivered no chunks: {result.body[:500]}" + return chunks + + +class TestResponsesBridgeChatCompletionsStreaming: + @pytest.fixture + def bridged_model(self, client: PassthroughClient, resources: ResourceManager) -> str: + model = f"e2e-bridge-stream-{unique_marker()}" + model_id = client.proxy.create_model( + model, + LiteLLMParamsBody(model=RESPONSES_ONLY_BACKEND, api_key="os.environ/OPENAI_API_KEY"), + ) + resources.defer(lambda: client.proxy.delete_model(model_id)) + return model + + @pytest.mark.covers( + "llm.chat_completions.openai.basic.stream.bridge_shares_chunk_id", + exercised_on=["chat_completions"], + ) + def test_bridged_stream_shares_one_chunk_id( + self, client: PassthroughClient, resources: ResourceManager, bridged_model: str + ) -> None: + result = client.proxy.chat_stream( + resources.key(), + ChatBody( + model=bridged_model, + messages=[ChatMessage(role="user", content=f"Count from 1 to 5, one number per line. {unique_marker()}")], + max_tokens=64, + stream=True, + ), + ) + + chunks = _bridge_chunks(result) + ids = {chunk.id for chunk in chunks} + assert len(ids) == 1, f"bridged stream used {len(ids)} different chunk ids: {sorted(ids)[:5]}" + assert ids.pop().startswith("chatcmpl-"), f"bridged chunk id is not chat-completion shaped: {chunks[0].id}" + + @pytest.mark.covers( + "llm.chat_completions.openai.basic.stream.bridge_streams_sse", + exercised_on=["chat_completions"], + ) + def test_bridged_stream_delivers_content_finish_reason_and_done( + self, client: PassthroughClient, resources: ResourceManager, bridged_model: str + ) -> None: + result = client.proxy.chat_stream( + resources.key(), + ChatBody( + model=bridged_model, + messages=[ChatMessage(role="user", content=f"Reply with the single word pong. {unique_marker()}")], + max_tokens=32, + stream=True, + ), + ) + + chunks = _bridge_chunks(result) + content = "".join(choice.delta.content or "" for chunk in chunks for choice in chunk.choices) + assert content.strip(), f"bridged stream completed with no content deltas: {result.stream_events[:3]}" + assert any( + choice.finish_reason for chunk in chunks for choice in chunk.choices + ), f"bridged stream never emitted a finish_reason: {result.stream_events[-3:]}" + assert result.stream_done, f"bridged stream did not terminate with [DONE]: {result.stream_events[-2:]}" + + @pytest.mark.covers( + "llm.chat_completions.openai.tool_use.stream.bridge_streams_tool_call", + exercised_on=["chat_completions"], + ) + def test_bridged_stream_reassembles_tool_call( + self, client: PassthroughClient, resources: ResourceManager, bridged_model: str + ) -> None: + result = client.proxy.chat_stream( + resources.key(), + ChatBody( + model=bridged_model, + messages=[ + ChatMessage( + role="user", + content="What is the weather in San Francisco? Use the get_weather tool.", + ) + ], + tools=[_WEATHER_TOOL], + tool_choice="required", + max_tokens=256, + stream=True, + ), + ) + + chunks = _bridge_chunks(result) + calls = [call for chunk in chunks for choice in chunk.choices for call in (choice.delta.tool_calls or [])] + assert calls, f"bridged stream returned no tool call for a tool-forced prompt: {result.stream_events[:5]}" + name = "".join(call.function.name or "" for call in calls) + arguments = "".join(call.function.arguments or "" for call in calls) + assert name == "get_weather", f"bridged stream streamed the wrong tool name: {name!r}" + args = _WeatherArgs.model_validate_json(arguments) + assert args.location.strip(), f"bridged tool call arguments missing location: {arguments!r}" diff --git a/tests/e2e/llm_translation/test_responses_metadata_e2e.py b/tests/e2e/llm_translation/test_responses_metadata_e2e.py deleted file mode 100644 index df854dcfa19..00000000000 --- a/tests/e2e/llm_translation/test_responses_metadata_e2e.py +++ /dev/null @@ -1,122 +0,0 @@ -"""Live e2e: /v1/responses with store + metadata (LIT-1201 customer path). - -Customers attach metadata and store=true, then continue with previous_response_id. -Both turns must succeed, and any Redis keys written for the session must carry a -positive TTL (not unbounded). -""" - -from __future__ import annotations - -import os -import socket -import time - -import pytest -from pydantic import BaseModel, ConfigDict - -from e2e_config import unique_marker -from e2e_http import require_successful_call -from endpoints_client import EndpointsClient, ResponsesResult -from lifecycle import ResourceManager -from models import LiteLLMParamsBody - -pytestmark = pytest.mark.e2e - - -class ResponsesMetadataBody(BaseModel): - model: str - input: str - store: bool = True - metadata: dict[str, str] - previous_response_id: str | None = None - instructions: str | None = "You are a helpful assistant." - - -class RedisKeyInfo(BaseModel): - model_config = ConfigDict(frozen=True) - - key: str - ttl: int - - -def _redis_scan(marker: str) -> tuple[RedisKeyInfo, ...]: - import redis - - host = os.environ["REDIS_HOST"] - port = int((os.environ.get("REDIS_PORT") or "6379").strip() or "6379") - try: - with socket.create_connection((host, port), timeout=3): - pass - except OSError as exc: - raise AssertionError( - f"REDIS_HOST={host!r}:{port} unreachable ({exc}); " - "LIT-1201 TTL check needs Redis the proxy writes to." - ) from exc - - client = redis.Redis(host=host, port=port, decode_responses=True, socket_timeout=5) - found: list[RedisKeyInfo] = [] - for key in client.scan_iter(match=f"*{marker}*", count=200): - found.append(RedisKeyInfo(key=str(key), ttl=int(client.ttl(key)))) - return tuple(found) - - -class TestResponsesMetadata: - @pytest.mark.covers( - "llm.responses.openai.basic.nonstream.works", - "other.config.responses.metadata_redis_ttl_bounded", - exercised_on=["responses"], - ) - def test_store_metadata_continues_and_redis_keys_have_ttl( - self, endpoints_client: EndpointsClient, resources: ResourceManager - ) -> None: - # Anthropic avoids OpenAI/Gemini quota flakes; Responses translation still - # exercises store + metadata + previous_response_id on the proxy. - marker = unique_marker() - model = f"e2e-resp-meta-{marker}" - model_id = endpoints_client.create_model( - model, - LiteLLMParamsBody( - model="anthropic/claude-haiku-4-5-20251001", - api_key="os.environ/ANTHROPIC_API_KEY", - ), - ) - resources.defer(lambda: endpoints_client.delete_model(model_id)) - key = resources.key() - - first = endpoints_client.proxy.transport.send( - "/v1/responses", - headers=endpoints_client.proxy.transport.bearer(key), - json=ResponsesMetadataBody( - model=model, - input=f"Remember marker {marker}. Reply with one word.", - metadata={"session_id": marker, "customer": "e2e"}, - ), - ) - require_successful_call(first) - parsed = ResponsesResult.model_validate_json(first.body) - assert parsed.id, f"responses must return an id: {first.body[:300]}" - assert parsed.text.strip(), f"responses returned empty text: {first.body[:300]}" - - second = endpoints_client.proxy.transport.send( - "/v1/responses", - headers=endpoints_client.proxy.transport.bearer(key), - json=ResponsesMetadataBody( - model=model, - input="Reply with the single word ok.", - previous_response_id=parsed.id, - metadata={"session_id": marker, "turn": "2"}, - ), - ) - require_successful_call(second) - second_parsed = ResponsesResult.model_validate_json(second.body) - assert second_parsed.text.strip(), ( - f"previous_response_id follow-up returned empty text: {second.body[:300]}" - ) - - time.sleep(1.0) - keys = _redis_scan(marker) - unbounded = tuple(k for k in keys if k.ttl == -1) - assert not unbounded, ( - "responses metadata must not leave Redis keys without TTL (LIT-1201); " - f"unbounded={unbounded}" - ) diff --git a/tests/e2e/mcp/mcp_client.py b/tests/e2e/mcp/mcp_client.py index f758a41cae6..4b1725bb205 100644 --- a/tests/e2e/mcp/mcp_client.py +++ b/tests/e2e/mcp/mcp_client.py @@ -11,12 +11,13 @@ request/response bodies are co-located here because only this suite speaks MCP. from __future__ import annotations +import time from collections.abc import Mapping from dataclasses import dataclass from pydantic import BaseModel, ConfigDict, Field, RootModel -from e2e_http import Headers, NoBody, Result, unwrap +from e2e_http import Headers, NoBody, Result, Success, unwrap from models import KeyGenerateBody, ObjectPermission from proxy_client import ProxyClient @@ -223,6 +224,31 @@ class McpClient: response_type=McpToolsListResponse, ) + def await_tool(self, key: str, server_id: str, needle: str) -> str: + """Poll tools/list until `server_id` serves a tool matching `needle`, and + return its fully-qualified name. Fails at poll_timeout. + + /v1/mcp/server returns as soon as the DB row is written, but the gateway + runs the initialize + tools/list handshake against the upstream lazily on + the first request that needs it, and reports a server it has not + discovered yet exactly like a dead one: an empty tool list. Waiting is + what separates the two. + """ + deadline = time.monotonic() + self.proxy.poll_timeout + while True: + result = self.list_tools(key) + if isinstance(result, Success): + tool_name = result.data.tool_name_containing(server_id, needle) + if tool_name is not None: + return tool_name + if time.monotonic() >= deadline: + raise AssertionError( + f"server {server_id} never served a tool matching {needle!r} within " + f"{self.proxy.poll_timeout}s of registration (upstream unreachable, or " + f"the key's grant was not applied); last tools/list: {result}" + ) + time.sleep(self.proxy.poll_interval) + def register_mcp_content_filter(self, *, name: str, blocked_keyword: str) -> str: """Register a default-on content-filter guardrail that runs on the MCP tool-call hook (pre_mcp_call) and blocks a single keyword. The keyword is diff --git a/tests/e2e/mcp/test_mcp_datadog_e2e.py b/tests/e2e/mcp/test_mcp_datadog_e2e.py index c772c4d3899..8a539b86bff 100644 --- a/tests/e2e/mcp/test_mcp_datadog_e2e.py +++ b/tests/e2e/mcp/test_mcp_datadog_e2e.py @@ -77,12 +77,7 @@ class TestDatadogMcpRoundTrip: "within the poll deadline; MCP search would have nothing to find" ) - tools = unwrap(client.list_tools(key)) - tool_name = tools.tool_name_containing(server_id, SEARCH_LOGS_TOOL) - assert tool_name is not None, ( - f"granted key never saw {SEARCH_LOGS_TOOL} on server {server_id}; " - f"tools={tools.tool_names_for_server(server_id)}" - ) + tool_name = client.await_tool(key, server_id, SEARCH_LOGS_TOOL) call = unwrap( client.call_tool( diff --git a/tests/e2e/mcp/test_mcp_guardrail_e2e.py b/tests/e2e/mcp/test_mcp_guardrail_e2e.py index 63239444454..dcab235465f 100644 --- a/tests/e2e/mcp/test_mcp_guardrail_e2e.py +++ b/tests/e2e/mcp/test_mcp_guardrail_e2e.py @@ -22,7 +22,7 @@ import pytest from datadog_mcp import SEARCH_LOGS_TOOL, assert_dd_mcp_creds, register_datadog_mcp from e2e_config import DD_SEARCH_FROM, unique_marker -from e2e_http import Result, Success, UnknownApiError, unwrap +from e2e_http import Result, Success, UnknownApiError from lifecycle import ResourceManager from mcp_client import McpCallToolResponse, McpClient, McpToolArguments @@ -75,12 +75,7 @@ class TestMcpToolCallGuardrail: key = client.generate_key(user_id=f"e2e-mcp-guard-{marker}", mcp_servers=[server_id]) resources.defer(lambda: client.proxy.delete_key(key)) - tools = unwrap(client.list_tools(key)) - tool_name = tools.tool_name_containing(server_id, SEARCH_LOGS_TOOL) - assert tool_name is not None, ( - f"granted key never saw {SEARCH_LOGS_TOOL} on server {server_id}; " - f"tools={tools.tool_names_for_server(server_id)}" - ) + tool_name = client.await_tool(key, server_id, SEARCH_LOGS_TOOL) def search(query: str) -> Result[McpCallToolResponse]: arguments: McpToolArguments = { diff --git a/tests/e2e/mcp/test_mcp_key_access_e2e.py b/tests/e2e/mcp/test_mcp_key_access_e2e.py index 412b33d244a..4aeb811a64f 100644 --- a/tests/e2e/mcp/test_mcp_key_access_e2e.py +++ b/tests/e2e/mcp/test_mcp_key_access_e2e.py @@ -48,12 +48,7 @@ class TestMcpKeyWithoutAccessIsDenied: permitted_key = _key(client, resources, mcp_servers=[server_id]) denied_key = _key(client, resources, mcp_servers=None) - permitted = unwrap(client.list_tools(permitted_key)) - tool_name = permitted.tool_name_containing(server_id, SEARCH_LOGS_TOOL) - assert tool_name is not None, ( - f"granted key did not see {SEARCH_LOGS_TOOL} (upstream dead or grant not applied): " - f"{permitted.tool_names_for_server(server_id)}" - ) + _ = client.await_tool(permitted_key, server_id, SEARCH_LOGS_TOOL) denied_tools = unwrap(client.list_tools(denied_key)).tool_names_for_server(server_id) assert denied_tools == frozenset(), ( @@ -73,12 +68,7 @@ class TestMcpKeyWithoutAccessIsDenied: permitted_key = _key(client, resources, mcp_servers=[server_id]) denied_key = _key(client, resources, mcp_servers=None) - permitted = unwrap(client.list_tools(permitted_key)) - tool_name = permitted.tool_name_containing(server_id, SEARCH_LOGS_TOOL) - assert tool_name is not None, ( - f"granted key did not discover {SEARCH_LOGS_TOOL} (upstream dead or grant not applied): " - f"{permitted.tool_names_for_server(server_id)}" - ) + tool_name = client.await_tool(permitted_key, server_id, SEARCH_LOGS_TOOL) search_args = { "query": "service:litellm", diff --git a/tests/proxy_unit_tests/test_update_spend.py b/tests/proxy_unit_tests/test_update_spend.py index e2dca0a0f81..131f46a3e21 100644 --- a/tests/proxy_unit_tests/test_update_spend.py +++ b/tests/proxy_unit_tests/test_update_spend.py @@ -28,11 +28,13 @@ class MockPrismaClient: # Initialize transaction lists self.spend_log_transactions = [] self.daily_user_spend_transactions = {} + self.tool_usage_transactions = [] - # Add lock for spend_log_transactions (matches real PrismaClient) + # Add locks for the transaction queues (matches real PrismaClient) import asyncio self._spend_log_transactions_lock = asyncio.Lock() + self._tool_usage_transactions_lock = asyncio.Lock() def jsonify_object(self, obj): return obj diff --git a/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_handler.py b/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_handler.py index a5bc01c2b74..8ecb7f4c6f0 100644 --- a/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_handler.py +++ b/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_handler.py @@ -203,3 +203,65 @@ async def test_acompletion_preserves_top_level_stream_flag_in_responses_request( assert result is stream assert transform_request.call_args.kwargs["optional_params"]["stream"] is True + + +def _completed_chat_response() -> ModelResponse: + return ModelResponse( + id="chatcmpl-completed", + model="gpt-5.4", + choices=[ + { + "index": 0, + "message": {"role": "assistant", "content": "pong"}, + "finish_reason": "stop", + } + ], + ) + + +@pytest.mark.asyncio +async def test_acompletion_streams_completed_model_response(): + """A streaming request whose bridge call comes back already completed must still be + handed back as an async-iterable stream. Returning the bare ModelResponse crashed the + proxy's SSE generator with "'async for' requires an object with __aiter__ method". + Regression for #33154.""" + completed = _completed_chat_response() + bridge = ResponsesToCompletionBridgeHandler() + + with ( + patch.object( + bridge.transformation_handler, + "transform_request", + return_value={"model": "gpt-5.4", "input": "hi"}, + ), + patch("litellm.aresponses", new=AsyncMock(return_value=completed)), + ): + result = await bridge.acompletion(**_bridge_kwargs(stream=True)) + + assert isinstance(result, CustomStreamWrapper), f"streaming request got {type(result)}" + chunks = [chunk async for chunk in result] + assert "".join( + chunk.choices[0].delta.content or "" for chunk in chunks + ) == "pong", f"completed response did not stream its content: {chunks}" + assert [c for c in chunks if c.choices[0].finish_reason], "stream never emitted a finish_reason" + + +def test_completion_streams_completed_model_response(): + completed = _completed_chat_response() + bridge = ResponsesToCompletionBridgeHandler() + + with ( + patch.object( + bridge.transformation_handler, + "transform_request", + return_value={"model": "gpt-5.4", "input": "hi"}, + ), + patch("litellm.responses", return_value=completed), + ): + result = bridge.completion(**_bridge_kwargs(stream=True)) + + assert isinstance(result, CustomStreamWrapper), f"streaming request got {type(result)}" + chunks = list(result) + assert "".join(chunk.choices[0].delta.content or "" for chunk in chunks) == "pong", ( + f"completed response did not stream its content: {chunks}" + ) diff --git a/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py b/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py index 6907e4d0d02..a111b932f2c 100644 --- a/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py +++ b/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py @@ -2855,6 +2855,41 @@ def test_streaming_function_call_tool_id_for_degenerate_call_id(): assert stream_tool_id("fc_2", "call_tokyo") == "call_tokyo" +def test_streaming_chunks_share_one_chat_completion_id(): + """Every chunk of one streamed chat completion must carry the same ``id``, per the + OpenAI spec. The bridge builds a fresh ``ModelResponseStream`` per Responses event, + so without a stream-scoped id each chunk got a new ``chatcmpl-`` and clients + that validate id consistency (openai-go's ChatCompletionAccumulator) silently + dropped every chunk after the first. Regression for #32854.""" + from litellm.completion_extras.litellm_responses_transformation.transformation import ( + OpenAiResponsesToChatCompletionStreamIterator, + ) + + iterator = OpenAiResponsesToChatCompletionStreamIterator( + streaming_response=None, sync_stream=True + ) + events = [ + {"type": "response.created", "response": {"id": "resp_abc", "output": []}}, + {"type": "response.output_text.delta", "delta": "Hel"}, + {"type": "response.output_text.delta", "delta": "lo"}, + { + "type": "response.completed", + "response": {"id": "resp_abc", "output": [{"type": "message"}]}, + }, + ] + + ids = [iterator.chunk_parser(event).id for event in events] + + assert len(set(ids)) == 1, f"streamed chunks carried different ids: {ids}" + assert ids[0], "streamed chunks carried an empty id" + + other_stream = OpenAiResponsesToChatCompletionStreamIterator( + streaming_response=None, sync_stream=True + ) + assert ( + other_stream.chunk_parser(events[1]).id != ids[0] + ), "a separate stream must get its own id, not a process-wide one" + @pytest.mark.asyncio @pytest.mark.parametrize( diff --git a/tests/test_litellm/integrations/test_langsmith_init.py b/tests/test_litellm/integrations/test_langsmith_init.py index 5d6b7c74690..129dda4abde 100644 --- a/tests/test_litellm/integrations/test_langsmith_init.py +++ b/tests/test_litellm/integrations/test_langsmith_init.py @@ -347,3 +347,90 @@ class TestLangsmithRedactUserApiKeyInfo: assert "user_api_key_user_id" not in nested assert nested["session_id"] == "sess-1" assert extra["session_id"] == "sess-1" + + def test_redact_enabled_strips_user_api_key_info_from_inputs(self, reset_redact_flag): + """ + Regression (LIT-4306): `inputs` is the whole StandardLoggingPayload, so + `redact_user_api_key_info` has to cover `inputs.metadata` the same way it + covers `extra` - including the nested `requester_metadata` copy. Before + the fix `extra` was redacted and `inputs` shipped every user_api_key_* + field verbatim. + """ + litellm.redact_user_api_key_info = True + logger = self._logger() + metadata = self._metadata_with_user_api_key_fields() + metadata["user_api_key_auth_metadata"] = {"priority": "high"} + payload = { + "id": "run-1", + "response": {"choices": []}, + "metadata": metadata, + "startTime": 1.0, + "endTime": 2.0, + "request_tags": [], + "error_str": None, + "status": "success", + "response_cost": 0.0, + "prompt_tokens": 1, + "completion_tokens": 1, + "total_tokens": 2, + } + credentials = { + "LANGSMITH_API_KEY": "test-key", + "LANGSMITH_PROJECT": "test-project", + "LANGSMITH_BASE_URL": "https://api.smith.langchain.com", + } + + data = logger._prepare_log_data( + kwargs={"litellm_params": {"metadata": metadata}, "standard_logging_object": payload}, + response_obj=None, + start_time=1.0, + end_time=2.0, + credentials=credentials, + ) + + inputs_metadata = data["inputs"]["metadata"] + assert [k for k in inputs_metadata if k.startswith("user_api_key")] == [] + assert [k for k in inputs_metadata["requester_metadata"] if k.startswith("user_api_key")] == [] + # inputs and extra must agree - they go through the same redaction now + assert [k for k in data["extra"] if k.startswith("user_api_key")] == [] + # non-identity payload is untouched + assert inputs_metadata["model"] == "gpt-4" + assert inputs_metadata["requester_metadata"]["session_id"] == "sess-1" + assert data["inputs"]["total_tokens"] == 2 + # the shared standard_logging_object other loggers read is not mutated + assert "user_api_key_hash" in payload["metadata"] + assert "user_api_key_user_id" in payload["metadata"]["requester_metadata"] + + def test_redact_disabled_keeps_user_api_key_info_in_inputs(self, reset_redact_flag): + """Flag off: the identity fields stay. The flag governs them, not this fix.""" + litellm.redact_user_api_key_info = False + logger = self._logger() + metadata = self._metadata_with_user_api_key_fields() + payload = { + "id": "run-1", + "response": {"choices": []}, + "metadata": metadata, + "startTime": 1.0, + "endTime": 2.0, + "request_tags": [], + "error_str": None, + "status": "success", + "response_cost": 0.0, + "prompt_tokens": 1, + "completion_tokens": 1, + "total_tokens": 2, + } + + data = logger._prepare_log_data( + kwargs={"litellm_params": {"metadata": metadata}, "standard_logging_object": payload}, + response_obj=None, + start_time=1.0, + end_time=2.0, + credentials={ + "LANGSMITH_API_KEY": "test-key", + "LANGSMITH_PROJECT": "test-project", + "LANGSMITH_BASE_URL": "https://api.smith.langchain.com", + }, + ) + + assert data["inputs"]["metadata"]["user_api_key_hash"] == "abc123" diff --git a/tests/test_litellm/integrations/test_prometheus_cache_metrics.py b/tests/test_litellm/integrations/test_prometheus_cache_metrics.py index 6c9923322fd..aa031bb813b 100644 --- a/tests/test_litellm/integrations/test_prometheus_cache_metrics.py +++ b/tests/test_litellm/integrations/test_prometheus_cache_metrics.py @@ -258,6 +258,158 @@ class TestPrometheusCacheMetrics: # Should not emit read metric, because explicit provider value is zero. mock_logger.litellm_provider_cache_read_input_tokens_metric.labels.assert_not_called() + def test_provider_cache_creation_fallback_to_cache_write_tokens( + self, sample_enum_values + ): + """OpenAI-style usage (prompt_tokens_details.cache_write_tokens, no top-level + cache_creation_input_tokens) must populate the provider cache creation metric.""" + mock_logger = MagicMock() + + from litellm.integrations.prometheus import PrometheusLogger + + standard_logging_payload = { + "cache_hit": False, + "total_tokens": 12100, + "prompt_tokens": 12000, + "completion_tokens": 100, + "model_group": "openai", + "request_tags": [], + "metadata": { + "usage_object": { + "prompt_tokens_details": { + "cached_tokens": 0, + "cache_write_tokens": 800, + }, + } + }, + } + + mock_logger.litellm_cache_hits_metric = MagicMock() + mock_logger.litellm_cache_misses_metric = MagicMock() + mock_logger.litellm_cached_tokens_metric = MagicMock() + mock_logger.litellm_provider_cache_read_input_tokens_metric = MagicMock() + mock_logger.litellm_provider_cache_creation_input_tokens_metric = MagicMock() + mock_logger.get_labels_for_metric = MagicMock( + return_value=[ + "model", + "hashed_api_key", + "api_key_alias", + "team", + "team_alias", + "end_user", + "user", + ] + ) + + PrometheusLogger._increment_cache_metrics( + mock_logger, + standard_logging_payload=standard_logging_payload, + enum_values=sample_enum_values, + ) + + mock_logger.litellm_provider_cache_creation_input_tokens_metric.labels().inc.assert_called_once_with( + 800 + ) + + def test_provider_cache_creation_fallback_to_cache_creation_tokens( + self, sample_enum_values + ): + """Normalized litellm usage dumps carry cache_creation_tokens in + prompt_tokens_details; the fallback must read it when cache_write_tokens is absent.""" + mock_logger = MagicMock() + + from litellm.integrations.prometheus import PrometheusLogger + + standard_logging_payload = { + "cache_hit": False, + "total_tokens": 100, + "prompt_tokens": 50, + "completion_tokens": 50, + "model_group": "openai", + "request_tags": [], + "metadata": { + "usage_object": { + "prompt_tokens_details": {"cache_creation_tokens": 42}, + } + }, + } + + mock_logger.litellm_cache_hits_metric = MagicMock() + mock_logger.litellm_cache_misses_metric = MagicMock() + mock_logger.litellm_cached_tokens_metric = MagicMock() + mock_logger.litellm_provider_cache_read_input_tokens_metric = MagicMock() + mock_logger.litellm_provider_cache_creation_input_tokens_metric = MagicMock() + mock_logger.get_labels_for_metric = MagicMock( + return_value=[ + "model", + "hashed_api_key", + "api_key_alias", + "team", + "team_alias", + "end_user", + "user", + ] + ) + + PrometheusLogger._increment_cache_metrics( + mock_logger, + standard_logging_payload=standard_logging_payload, + enum_values=sample_enum_values, + ) + + mock_logger.litellm_provider_cache_creation_input_tokens_metric.labels().inc.assert_called_once_with( + 42 + ) + + def test_provider_cache_creation_does_not_fallback_on_explicit_zero( + self, sample_enum_values + ): + """Explicit cache_creation_input_tokens=0 must not trigger fallback to + prompt_tokens_details, mirroring the cache-read semantics.""" + mock_logger = MagicMock() + + from litellm.integrations.prometheus import PrometheusLogger + + standard_logging_payload = { + "cache_hit": False, + "total_tokens": 100, + "prompt_tokens": 50, + "completion_tokens": 50, + "model_group": "openai", + "request_tags": [], + "metadata": { + "usage_object": { + "cache_creation_input_tokens": 0, + "prompt_tokens_details": {"cache_write_tokens": 800}, + } + }, + } + + mock_logger.litellm_cache_hits_metric = MagicMock() + mock_logger.litellm_cache_misses_metric = MagicMock() + mock_logger.litellm_cached_tokens_metric = MagicMock() + mock_logger.litellm_provider_cache_read_input_tokens_metric = MagicMock() + mock_logger.litellm_provider_cache_creation_input_tokens_metric = MagicMock() + mock_logger.get_labels_for_metric = MagicMock( + return_value=[ + "model", + "hashed_api_key", + "api_key_alias", + "team", + "team_alias", + "end_user", + "user", + ] + ) + + PrometheusLogger._increment_cache_metrics( + mock_logger, + standard_logging_payload=standard_logging_payload, + enum_values=sample_enum_values, + ) + + mock_logger.litellm_provider_cache_creation_input_tokens_metric.labels.assert_not_called() + def test_increment_cache_metrics_when_cache_hit_is_none(self, sample_enum_values): """Test that no metrics are incremented when cache_hit is None""" # Create mock for PrometheusLogger instance diff --git a/tests/test_litellm/integrations/test_prometheus_token_detail_metrics.py b/tests/test_litellm/integrations/test_prometheus_token_detail_metrics.py index 72a4e80717b..5e3846d6fa2 100644 --- a/tests/test_litellm/integrations/test_prometheus_token_detail_metrics.py +++ b/tests/test_litellm/integrations/test_prometheus_token_detail_metrics.py @@ -150,6 +150,57 @@ class TestIncrementTokenDetailMetrics: 10.0 ) + def test_cache_creation_falls_back_to_cache_write_tokens(self, sample_enum_values): + logger = _make_mock_logger() + payload = { + "metadata": { + "usage_object": { + "prompt_tokens": 12000, + "completion_tokens": 100, + "total_tokens": 12100, + "prompt_tokens_details": { + "cached_tokens": 0, + "cache_write_tokens": 800, + }, + } + }, + } + + PrometheusLogger._increment_token_detail_metrics( + logger, + standard_logging_payload=payload, + enum_values=sample_enum_values, + ) + + logger.litellm_input_cache_creation_tokens_metric.labels().inc.assert_called_once_with( + 800.0 + ) + + def test_cache_write_tokens_takes_precedence_over_cache_creation_tokens( + self, sample_enum_values + ): + logger = _make_mock_logger() + payload = { + "metadata": { + "usage_object": { + "prompt_tokens_details": { + "cache_creation_tokens": 25, + "cache_write_tokens": 800, + }, + } + }, + } + + PrometheusLogger._increment_token_detail_metrics( + logger, + standard_logging_payload=payload, + enum_values=sample_enum_values, + ) + + logger.litellm_input_cache_creation_tokens_metric.labels().inc.assert_called_once_with( + 800.0 + ) + def test_skips_metrics_when_value_is_zero(self, sample_enum_values): logger = _make_mock_logger() payload = { diff --git a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py index bcda88ea609..9565de1139c 100644 --- a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py +++ b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py @@ -3166,3 +3166,34 @@ async def test_bedrock_converse_message_level_cache_point_preserves_ttl_async(): ) assert _collect_cache_points(result) == [{"type": "default", "ttl": "1h"}] + + +def _n_choices_response(*names_per_choice): + from types import SimpleNamespace + + choices = [ + SimpleNamespace( + message=SimpleNamespace( + tool_calls=[SimpleNamespace(id=f"c{i}", function=SimpleNamespace(name=name, arguments="{}"))] + ) + ) + for i, name in enumerate(names_per_choice) + ] + return SimpleNamespace(choices=choices) + + +def test_get_tool_calls_from_response_defaults_to_primary_choice_only(): + from litellm.litellm_core_utils.prompt_templates.factory import get_tool_calls_from_response + + response = _n_choices_response("tool_alpha", "tool_beta") + + assert [tc["name"] for tc in get_tool_calls_from_response(response)] == ["tool_alpha"] + + +def test_get_tool_calls_from_response_include_all_choices_reads_every_choice(): + from litellm.litellm_core_utils.prompt_templates.factory import get_tool_calls_from_response + + response = _n_choices_response("tool_alpha", "tool_beta") + + names = [tc["name"] for tc in get_tool_calls_from_response(response, include_all_choices=True)] + assert names == ["tool_alpha", "tool_beta"] diff --git a/tests/test_litellm/llms/openai/evals/test_openai_evals_transformation.py b/tests/test_litellm/llms/openai/evals/test_openai_evals_transformation.py index f39be511b97..9a30ca0ee60 100644 --- a/tests/test_litellm/llms/openai/evals/test_openai_evals_transformation.py +++ b/tests/test_litellm/llms/openai/evals/test_openai_evals_transformation.py @@ -252,9 +252,7 @@ def test_transform_cancel_eval_response(config: OpenAIEvalsConfig): "object": "eval", "status": "cancelled", }, - request=httpx.Request( - "POST", "https://api.openai.com/v1/evals/eval_123/cancel" - ), + request=httpx.Request("POST", "https://api.openai.com/v1/evals/eval_123/cancel"), ) result = config.transform_cancel_eval_response( @@ -276,8 +274,169 @@ def test_transform_run_requests_encode_eval_and_run_ids(config: OpenAIEvalsConfi headers={}, ) - assert ( - url - == "https://api.openai.com/v1/evals/..%2F..%2Fevals%3Fx%3D1%23frag/runs/..%2Fruns%23other/cancel" - ) + assert url == "https://api.openai.com/v1/evals/..%2F..%2Fevals%3Fx%3D1%23frag/runs/..%2Fruns%23other/cancel" assert request_body == {} + + +def _eval_json_response(url: str, method: str = "GET") -> httpx.Response: + return httpx.Response( + status_code=200, + json={ + "id": "eval_123", + "object": "eval", + "created_at": 1234567890, + "name": "Test Eval", + "data_source_config": {"type": "stored_completions"}, + "testing_criteria": [], + }, + request=httpx.Request(method, url), + ) + + +def _run_json(run_id: str = "evalrun_123", status: str = "queued") -> dict: + return { + "id": run_id, + "object": "eval.run", + "created_at": 1234567890, + "status": status, + "data_source": {"type": "completions"}, + "eval_id": "eval_123", + } + + +def test_transform_get_eval_response(config: OpenAIEvalsConfig): + result = config.transform_get_eval_response( + raw_response=_eval_json_response("https://api.openai.com/v1/evals/eval_123"), + logging_obj=None, + ) + + assert result.id == "eval_123" + assert result.object == "eval" + assert result.name == "Test Eval" + + +def test_transform_update_eval_response(config: OpenAIEvalsConfig): + result = config.transform_update_eval_response( + raw_response=_eval_json_response("https://api.openai.com/v1/evals/eval_123", method="POST"), + logging_obj=None, + ) + + assert result.id == "eval_123" + assert result.name == "Test Eval" + + +def test_transform_create_run_response(config: OpenAIEvalsConfig): + response = httpx.Response( + status_code=200, + json=_run_json(), + request=httpx.Request("POST", "https://api.openai.com/v1/evals/eval_123/runs"), + ) + + result = config.transform_create_run_response( + raw_response=response, + logging_obj=None, + ) + + assert result.id == "evalrun_123" + assert result.status == "queued" + assert result.eval_id == "eval_123" + + +def test_transform_list_runs_request(config: OpenAIEvalsConfig): + url, query_params = config.transform_list_runs_request( + eval_id="eval_123", + list_params={"limit": 5, "after": "evalrun_1", "order": "asc"}, + litellm_params=GenericLiteLLMParams(api_base="https://api.openai.com"), + headers={}, + ) + + assert url == "https://api.openai.com/v1/evals/eval_123/runs" + assert query_params == {"limit": 5, "after": "evalrun_1", "order": "asc"} + + +def test_transform_list_runs_response(config: OpenAIEvalsConfig): + response = httpx.Response( + status_code=200, + json={ + "object": "list", + "data": [_run_json()], + "first_id": "evalrun_123", + "last_id": "evalrun_123", + "has_more": False, + }, + request=httpx.Request("GET", "https://api.openai.com/v1/evals/eval_123/runs"), + ) + + result = config.transform_list_runs_response( + raw_response=response, + logging_obj=None, + ) + + assert result.object == "list" + assert len(result.data) == 1 + assert result.data[0].id == "evalrun_123" + assert result.has_more is False + + +def test_transform_get_run_response(config: OpenAIEvalsConfig): + response = httpx.Response( + status_code=200, + json=_run_json(status="completed"), + request=httpx.Request("GET", "https://api.openai.com/v1/evals/eval_123/runs/evalrun_123"), + ) + + result = config.transform_get_run_response( + raw_response=response, + logging_obj=None, + ) + + assert result.id == "evalrun_123" + assert result.status == "completed" + + +def test_transform_cancel_run_response(config: OpenAIEvalsConfig): + response = httpx.Response( + status_code=200, + json={"id": "evalrun_123", "object": "eval.run", "status": "cancelled"}, + request=httpx.Request( + "POST", + "https://api.openai.com/v1/evals/eval_123/runs/evalrun_123/cancel", + ), + ) + + result = config.transform_cancel_run_response( + raw_response=response, + logging_obj=None, + ) + + assert result.id == "evalrun_123" + assert result.status == "cancelled" + + +def test_transform_delete_run_request(config: OpenAIEvalsConfig): + url, headers, request_body = config.transform_delete_run_request( + eval_id="eval_123", + run_id="evalrun_123", + api_base="https://api.openai.com", + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert url == "https://api.openai.com/v1/evals/eval_123/runs/evalrun_123" + assert request_body == {} + + +def test_transform_delete_run_response(config: OpenAIEvalsConfig): + response = httpx.Response( + status_code=200, + json={"run_id": "evalrun_123", "object": "eval.run.deleted", "deleted": True}, + request=httpx.Request("DELETE", "https://api.openai.com/v1/evals/eval_123/runs/evalrun_123"), + ) + + result = config.transform_delete_run_response( + raw_response=response, + logging_obj=None, + ) + + assert result.run_id == "evalrun_123" + assert result.deleted is True diff --git a/tests/test_litellm/llms/volcengine/responses/test_volcengine_responses_transformation.py b/tests/test_litellm/llms/volcengine/responses/test_volcengine_responses_transformation.py index 13571e63c7d..4581f4af7b6 100644 --- a/tests/test_litellm/llms/volcengine/responses/test_volcengine_responses_transformation.py +++ b/tests/test_litellm/llms/volcengine/responses/test_volcengine_responses_transformation.py @@ -4,9 +4,11 @@ Tests for Volcengine Responses API transformation. import os import sys +from typing import List, Literal, Optional, Union import httpx import pytest +from pydantic import BaseModel, Field sys.path.insert(0, os.path.abspath("../../../../..")) @@ -32,12 +34,10 @@ class TestVolcengineResponsesAPITransformation: ) assert config is not None, "Config should not be None for Volcengine provider" - assert isinstance( - config, VolcEngineResponsesAPIConfig - ), f"Expected VolcEngineResponsesAPIConfig, got {type(config)}" - assert ( - config.custom_llm_provider == LlmProviders.VOLCENGINE - ), "custom_llm_provider should be VOLCENGINE" + assert isinstance(config, VolcEngineResponsesAPIConfig), ( + f"Expected VolcEngineResponsesAPIConfig, got {type(config)}" + ) + assert config.custom_llm_provider == LlmProviders.VOLCENGINE, "custom_llm_provider should be VOLCENGINE" def test_parallel_tool_calls_dropped(self): """Volcengine does not list parallel_tool_calls; ensure it is removed.""" @@ -54,9 +54,7 @@ class TestVolcengineResponsesAPITransformation: drop_params=False, ) - assert ( - "parallel_tool_calls" not in mapped - ), "parallel_tool_calls must be dropped" + assert "parallel_tool_calls" not in mapped, "parallel_tool_calls must be dropped" assert mapped.get("temperature") == 0.5 assert "metadata" not in mapped, "Undocumented params should not be included" @@ -91,14 +89,10 @@ class TestVolcengineResponsesAPITransformation: default_url = config.get_complete_url(api_base=None, litellm_params={}) assert default_url == "https://ark.cn-beijing.volces.com/api/v3/responses" - api_base_with_api = config.get_complete_url( - api_base="https://custom.volc.com/api/v3", litellm_params={} - ) + api_base_with_api = config.get_complete_url(api_base="https://custom.volc.com/api/v3", litellm_params={}) assert api_base_with_api == "https://custom.volc.com/api/v3/responses" - api_base_full = config.get_complete_url( - api_base="https://custom.volc.com/api/v3/responses", litellm_params={} - ) + api_base_full = config.get_complete_url(api_base="https://custom.volc.com/api/v3/responses", litellm_params={}) assert api_base_full == "https://custom.volc.com/api/v3/responses" def test_response_id_path_requests_encode_response_id(self): @@ -112,10 +106,7 @@ class TestVolcengineResponsesAPITransformation: headers={}, ) - assert ( - url - == "https://custom.volc.com/api/v3/responses/..%2F..%2Fresponses%2Fother%3Fx%3D1%23frag/cancel" - ) + assert url == "https://custom.volc.com/api/v3/responses/..%2F..%2Fresponses%2Fother%3Fx%3D1%23frag/cancel" assert params == {} @pytest.mark.parametrize( @@ -125,9 +116,7 @@ class TestVolcengineResponsesAPITransformation: (GenericLiteLLMParams(api_key="attr-key"), "attr-key"), ], ) - def test_validate_environment_uses_api_key( - self, monkeypatch, litellm_params, expected_key - ): + def test_validate_environment_uses_api_key(self, monkeypatch, litellm_params, expected_key): """validate_environment should pull api key from params/env and attach headers.""" config = VolcEngineResponsesAPIConfig() @@ -135,9 +124,7 @@ class TestVolcengineResponsesAPITransformation: monkeypatch.delenv("ARK_API_KEY", raising=False) monkeypatch.delenv("VOLCENGINE_API_KEY", raising=False) - headers = config.validate_environment( - headers={}, model="volcengine/demo-model", litellm_params=litellm_params - ) + headers = config.validate_environment(headers={}, model="volcengine/demo-model", litellm_params=litellm_params) assert headers.get("Authorization") == f"Bearer {expected_key}" assert headers.get("Content-Type") == "application/json" @@ -151,9 +138,7 @@ class TestVolcengineResponsesAPITransformation: monkeypatch.delenv("VOLCENGINE_API_KEY", raising=False) with pytest.raises(ValueError): - config.validate_environment( - headers={}, model="volcengine/demo", litellm_params={} - ) + config.validate_environment(headers={}, model="volcengine/demo", litellm_params={}) def test_unsupported_params_are_dropped_with_extra_body(self): """Unknown fields (including extra_body) should be dropped before send.""" @@ -240,9 +225,7 @@ class TestVolcengineResponsesAPITransformation: # Use class name comparison instead of isinstance to avoid issues with # module reloading during parallel test execution (conftest reloads litellm) - assert ( - type(error).__name__ == "VolcEngineError" - ), f"Expected VolcEngineError, got {type(error).__name__}" + assert type(error).__name__ == "VolcEngineError", f"Expected VolcEngineError, got {type(error).__name__}" assert error.status_code == 400 assert error.message == "bad request" assert error.headers.get("x") == "y" @@ -296,3 +279,206 @@ class TestVolcengineResponsesAPITransformation: assert isinstance(result, DeleteResponseResult) assert result.deleted is True + + def test_transform_streaming_response_fills_missing_required_fields(self): + config = VolcEngineResponsesAPIConfig() + + event = config.transform_streaming_response( + model="volcengine/demo-model", + parsed_chunk={"type": "response.completed", "response": {"id": "resp_1"}}, + logging_obj=None, + ) + + assert type(event).__name__ == "ResponseCompletedEvent" + assert event.type == "response.completed" + assert event.response.id == "resp_1" + assert event.response.output == [] + assert event.response.created_at == 0 + + def test_transform_response_api_response_falls_back_to_model_construct(self): + config = VolcEngineResponsesAPIConfig() + http_response = httpx.Response( + status_code=200, + json={"id": "resp_fallback", "created_at": 123, "output": "not-a-list"}, + request=httpx.Request("POST", "https://example.com/responses"), + headers={"x-test": "1"}, + ) + + result = config.transform_response_api_response( + model="volcengine/demo-model", + raw_response=http_response, + logging_obj=type( + "Logger", + (), + {"post_call": staticmethod(lambda **kwargs: None)}, + ), + ) + + assert result.id == "resp_fallback" + assert result.output == "not-a-list" + assert result._hidden_params["headers"].get("x-test") == "1" + + def test_transform_delete_response_api_request_builds_url(self): + config = VolcEngineResponsesAPIConfig() + + url, data = config.transform_delete_response_api_request( + response_id="resp_123", + api_base="https://custom.volc.com/api/v3/responses", + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert url == "https://custom.volc.com/api/v3/responses/resp_123" + assert data == {} + + def test_transform_get_response_api_request_and_response(self): + config = VolcEngineResponsesAPIConfig() + + url, data = config.transform_get_response_api_request( + response_id="resp 123", + api_base="https://custom.volc.com/api/v3/responses", + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert url == "https://custom.volc.com/api/v3/responses/resp%20123" + assert data == {} + + http_response = httpx.Response( + status_code=200, + json={ + "id": "resp_123", + "object": "response", + "created_at": 123, + "status": "completed", + "output": [], + "model": "demo-model", + }, + request=httpx.Request("GET", url), + headers={"x-test": "1"}, + ) + + result = config.transform_get_response_api_response( + raw_response=http_response, + logging_obj=None, + ) + + assert result.id == "resp_123" + assert result._hidden_params["headers"].get("x-test") == "1" + + def test_transform_cancel_response_api_response_parses_json(self): + config = VolcEngineResponsesAPIConfig() + http_response = httpx.Response( + status_code=200, + json={ + "id": "resp_123", + "object": "response", + "created_at": 123, + "status": "cancelled", + "output": [], + "model": "demo-model", + }, + request=httpx.Request("POST", "https://example.com/responses/resp_123/cancel"), + headers={"x-test": "1"}, + ) + + result = config.transform_cancel_response_api_response( + raw_response=http_response, + logging_obj=None, + ) + + assert result.id == "resp_123" + assert result.status == "cancelled" + assert result._hidden_params["headers"].get("x-test") == "1" + + def test_transform_list_input_items_request_builds_query_params(self): + config = VolcEngineResponsesAPIConfig() + + url, params = config.transform_list_input_items_request( + response_id="resp_123", + api_base="https://custom.volc.com/api/v3/responses", + litellm_params=GenericLiteLLMParams(), + headers={}, + after="item_a", + before="item_b", + include=["metadata", "usage"], + limit=5, + order="asc", + ) + + assert url == "https://custom.volc.com/api/v3/responses/resp_123/input_items" + assert params == { + "after": "item_a", + "before": "item_b", + "include": "metadata,usage", + "limit": 5, + "order": "asc", + } + + def test_transform_list_input_items_response_returns_parsed_body(self): + config = VolcEngineResponsesAPIConfig() + payload = {"object": "list", "data": [{"id": "item_1"}]} + http_response = httpx.Response( + status_code=200, + json=payload, + request=httpx.Request("GET", "https://example.com/responses/resp_123/input_items"), + ) + + result = config.transform_list_input_items_response( + raw_response=http_response, + logging_obj=None, + ) + + assert result == payload + + +class _FillWidget(BaseModel): + type: Literal["widget"] + count: int + parts: List[str] + label: Optional[str] + + +class _FillGadget(BaseModel): + type: Literal["gadget"] + name: str + + +class _FillEnvelope(BaseModel): + kind: str = "envelope" + tags: List[str] = Field(default_factory=lambda: ["default-tag"]) + payload: Union[_FillWidget, _FillGadget] + entries: List[_FillWidget] + note: Optional[str] + values: Union[List[str], str] + + +class TestVolcengineStreamingFieldFill: + def test_fill_uses_defaults_factories_and_heuristics(self): + filled = VolcEngineResponsesAPIConfig._fill_missing_fields( + {"payload": {"type": "gadget", "name": "g"}, "entries": [{"type": "widget"}]}, + _FillEnvelope, + ) + + assert filled["kind"] == "envelope" + assert filled["tags"] == ["default-tag"] + assert filled["note"] is None + assert filled["values"] == [] + + validated = _FillEnvelope.model_validate(filled) + assert isinstance(validated.payload, _FillGadget) + assert validated.entries[0].count == 0 + assert validated.entries[0].parts == [] + assert validated.entries[0].label is None + + def test_fill_selects_union_member_by_type_literal(self): + filled = VolcEngineResponsesAPIConfig._fill_missing_fields( + {"payload": {"type": "widget"}, "entries": []}, + _FillEnvelope, + ) + + validated = _FillEnvelope.model_validate(filled) + assert isinstance(validated.payload, _FillWidget) + assert validated.payload.count == 0 + assert validated.payload.parts == [] + assert validated.payload.label is None diff --git a/tests/test_litellm/proxy/common_utils/test_callback_utils.py b/tests/test_litellm/proxy/common_utils/test_callback_utils.py index 8f390c096d7..bfd4ffe1593 100644 --- a/tests/test_litellm/proxy/common_utils/test_callback_utils.py +++ b/tests/test_litellm/proxy/common_utils/test_callback_utils.py @@ -18,6 +18,7 @@ from litellm.proxy.common_utils.callback_utils import ( get_remaining_tokens_and_requests_from_request_data, normalize_callback_names, sanitize_openai_provider_metadata, + strip_callback_config, ) import litellm @@ -452,3 +453,41 @@ def test_initialize_callbacks_on_proxy_non_dict_callback_specific_params_root( ) finally: litellm.callbacks = original_callbacks + + +def test_strip_callback_config_drops_credential_bearing_slots(): + """ + `logging` and `callback_settings` hold operator-configured integration + credentials. Both must be dropped from the key/team metadata the proxy + stamps into request metadata, while every other field survives untouched + (`priority` is read back by the dynamic rate limiter, `guardrails` by the + guardrail hooks). + """ + metadata = { + "logging": [ + { + "callback_name": "langsmith", + "callback_vars": {"langsmith_api_key": "litellm_enc::ciphertext"}, + } + ], + "callback_settings": {"callback_vars": {"langfuse_secret_key": "litellm_enc::other"}}, + "priority": "high", + "guardrails": ["presidio"], + "langsmith_provisioning": {"api_key_id": "prov-1"}, + } + + stripped = strip_callback_config(metadata) + + assert "logging" not in stripped + assert "callback_settings" not in stripped + assert stripped["priority"] == "high" + assert stripped["guardrails"] == ["presidio"] + assert stripped["langsmith_provisioning"] == {"api_key_id": "prov-1"} + # the caller's dict (UserAPIKeyAuth.metadata) is shared state - never mutate it + assert "logging" in metadata + assert "callback_settings" in metadata + + +@pytest.mark.parametrize("value", [None, "not-a-dict", 42]) +def test_strip_callback_config_passes_through_non_dicts(value): + assert strip_callback_config(value) is value diff --git a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py index 8149cf90e70..191080e3a48 100644 --- a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py +++ b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py @@ -76,6 +76,162 @@ async def test_daily_spend_tracking_with_disabled_spend_logs(): assert call_args["payload"]["custom_llm_provider"] == "openai" +def _tool_call_response(*names: str) -> object: + from types import SimpleNamespace + + tool_calls = [SimpleNamespace(function=SimpleNamespace(name=name)) for name in names] + return SimpleNamespace(choices=[SimpleNamespace(message=SimpleNamespace(tool_calls=tool_calls))]) + + +def _tool_usage_prisma() -> MagicMock: + prisma = MagicMock() + prisma.tool_usage_transactions = [] + prisma._tool_usage_transactions_lock = asyncio.Lock() + prisma.spend_log_transactions = [] + prisma._spend_log_transactions_lock = asyncio.Lock() + return prisma + + +def _minimal_spend_payload() -> dict: + return { + "request_id": "req-tool-1", + "startTime": datetime(2026, 7, 25, 10, 0, tzinfo=timezone.utc), + "endTime": datetime(2026, 7, 25, 10, 0, 1, tzinfo=timezone.utc), + "spend": 0.0, + "total_tokens": 42, + "mcp_namespaced_tool_name": None, + } + + +@pytest.mark.asyncio +async def test_update_database_enqueues_tool_usage_for_invoked_tools(): + db_writer = DBSpendUpdateWriter() + db_writer._insert_spend_log_to_db = AsyncMock() + db_writer._batch_database_updates = AsyncMock() + prisma = _tool_usage_prisma() + + with ( + patch("litellm.proxy.proxy_server.disable_spend_logs", False), + patch("litellm.proxy.proxy_server.prisma_client", prisma), + patch("litellm.proxy.proxy_server.litellm_proxy_budget_name", "test-budget"), + patch( + "litellm.proxy.spend_tracking.spend_tracking_utils.get_logging_payload", + return_value=_minimal_spend_payload(), + ), + ): + await db_writer.update_database( + token="test-token", + user_id="test-user", + end_user_id=None, + team_id=None, + org_id=None, + kwargs={"model": "gpt-4"}, + completion_response=_tool_call_response("get_weather"), + start_time=datetime.now(timezone.utc), + end_time=datetime.now(timezone.utc), + response_cost=0.1, + ) + await asyncio.sleep(0) + + assert len(prisma.tool_usage_transactions) == 1 + transaction = prisma.tool_usage_transactions[0] + assert transaction.request_id == "req-tool-1" + assert transaction.tool_names == ("get_weather",) + assert transaction.spend == 0.1 + assert transaction.total_tokens == 42 + assert transaction.date == "2026-07-25" + + +@pytest.mark.asyncio +async def test_update_database_enqueues_realtime_tool_usage(): + db_writer = DBSpendUpdateWriter() + db_writer._insert_spend_log_to_db = AsyncMock() + db_writer._batch_database_updates = AsyncMock() + prisma = _tool_usage_prisma() + + with ( + patch("litellm.proxy.proxy_server.disable_spend_logs", False), + patch("litellm.proxy.proxy_server.prisma_client", prisma), + patch("litellm.proxy.proxy_server.litellm_proxy_budget_name", "test-budget"), + patch( + "litellm.proxy.spend_tracking.spend_tracking_utils.get_logging_payload", + return_value=_minimal_spend_payload(), + ), + ): + await db_writer.update_database( + token="test-token", + user_id="test-user", + end_user_id=None, + team_id=None, + org_id=None, + kwargs={ + "model": "gpt-realtime", + "realtime_tool_calls": [ + {"id": "call_1", "type": "function", "function": {"name": "rt_tool", "arguments": "{}"}} + ], + }, + completion_response=None, + start_time=datetime.now(timezone.utc), + end_time=datetime.now(timezone.utc), + response_cost=0.2, + ) + await asyncio.sleep(0) + + assert len(prisma.tool_usage_transactions) == 1 + assert prisma.tool_usage_transactions[0].tool_names == ("rt_tool",) + + +def test_enqueue_tool_registry_upsert_reads_every_choice(): + from types import SimpleNamespace as NS + + db_writer = DBSpendUpdateWriter() + db_writer.tool_discovery_queue = MagicMock() + response = NS( + choices=[ + NS(message=NS(tool_calls=[NS(function=NS(name="tool_alpha"))])), + NS(message=NS(tool_calls=[NS(function=NS(name="tool_beta"))])), + ] + ) + + db_writer._enqueue_tool_registry_upsert(kwargs={}, completion_response=response) + + enqueued = [call.args[0]["tool_name"] for call in db_writer.tool_discovery_queue.add_update.call_args_list] + assert enqueued == ["tool_alpha", "tool_beta"] + + +@pytest.mark.asyncio +async def test_update_database_skips_tool_usage_when_spend_logs_disabled(): + db_writer = DBSpendUpdateWriter() + db_writer._insert_spend_log_to_db = AsyncMock() + db_writer._batch_database_updates = AsyncMock() + prisma = _tool_usage_prisma() + + with ( + patch("litellm.proxy.proxy_server.disable_spend_logs", True), + patch("litellm.proxy.proxy_server.prisma_client", prisma), + patch("litellm.proxy.proxy_server.litellm_proxy_budget_name", "test-budget"), + patch( + "litellm.proxy.spend_tracking.spend_tracking_utils.get_logging_payload", + return_value=_minimal_spend_payload(), + ), + ): + await db_writer.update_database( + token="test-token", + user_id="test-user", + end_user_id=None, + team_id=None, + org_id=None, + kwargs={"model": "gpt-4"}, + completion_response=_tool_call_response("get_weather"), + start_time=datetime.now(timezone.utc), + end_time=datetime.now(timezone.utc), + response_cost=0.1, + ) + await asyncio.sleep(0) + + assert prisma.tool_usage_transactions == [] + + @pytest.mark.asyncio async def test_update_daily_spend_with_null_entity_id(): """ @@ -152,6 +308,84 @@ async def test_update_daily_spend_with_null_entity_id(): assert create_data["failed_requests"] == 0 +def _daily_txn(user_id: str = "user1") -> dict: + return { + "user_id": user_id, + "date": "2024-01-01", + "api_key": "test-api-key", + "model": "gpt-4", + "custom_llm_provider": "openai", + "prompt_tokens": 10, + "completion_tokens": 20, + "spend": 0.1, + "api_requests": 1, + "successful_requests": 1, + "failed_requests": 0, + } + + +@pytest.mark.asyncio +async def test_update_daily_spend_does_not_retry_post_send_ambiguous_errors(): + # Regression for the double-apply hazard: a ReadTimeout means the batch was + # sent and its outcome is unknown; the engine can leave the transaction open + # on the pooled connection, so retrying stacks a second set of increments + # into it and one commit applies both. Post-send failures must drop the + # batch (loudly), never retry it. + import httpx + + mock_prisma_client = MagicMock() + mock_prisma_client.db.batch_ = MagicMock(side_effect=httpx.ReadTimeout("ambiguous")) + proxy_logging = MagicMock() + proxy_logging.failure_handler = AsyncMock() + + with pytest.raises(httpx.ReadTimeout): + await DBSpendUpdateWriter._update_daily_spend( + n_retry_times=3, + prisma_client=mock_prisma_client, + proxy_logging_obj=proxy_logging, + daily_spend_transactions={"k1": _daily_txn()}, + entity_type="user", + entity_id_field="user_id", + table_name="litellm_dailyuserspend", + unique_constraint_name="user_id_date_api_key_model_custom_llm_provider_mcp_namespaced_tool_name_endpoint", + ) + + mock_prisma_client.db.batch_.assert_called_once() + + +@pytest.mark.asyncio +async def test_update_daily_spend_retries_connect_errors(monkeypatch): + # ConnectError proves the statements never reached the database, so it is + # the one failure the writer may retry. + import httpx + + mock_batcher = MagicMock() + good_ctx = MagicMock() + good_ctx.__aenter__ = AsyncMock(return_value=mock_batcher) + good_ctx.__aexit__ = AsyncMock(return_value=None) + mock_prisma_client = MagicMock() + mock_prisma_client.db.batch_ = MagicMock(side_effect=[httpx.ConnectError("down"), good_ctx]) + proxy_logging = MagicMock() + proxy_logging.failure_handler = AsyncMock() + + async def fake_sleep(seconds: float) -> None: + return None + + monkeypatch.setattr("litellm.proxy.db.db_spend_update_writer.asyncio.sleep", fake_sleep) + await DBSpendUpdateWriter._update_daily_spend( + n_retry_times=3, + prisma_client=mock_prisma_client, + proxy_logging_obj=proxy_logging, + daily_spend_transactions={"k1": _daily_txn()}, + entity_type="user", + entity_id_field="user_id", + table_name="litellm_dailyuserspend", + unique_constraint_name="user_id_date_api_key_model_custom_llm_provider_mcp_namespaced_tool_name_endpoint", + ) + + assert mock_prisma_client.db.batch_.call_count == 2 + + @pytest.mark.asyncio async def test_update_daily_spend_sorting(): """ diff --git a/tests/test_litellm/proxy/db/test_spend_log_tool_index.py b/tests/test_litellm/proxy/db/test_spend_log_tool_index.py new file mode 100644 index 00000000000..71073fd216e --- /dev/null +++ b/tests/test_litellm/proxy/db/test_spend_log_tool_index.py @@ -0,0 +1,348 @@ +""" +Tests for the tool usage writer: ToolUsageTransaction construction (invoked tools +only) and the flush that writes LiteLLM_SpendLogToolIndex plus the +LiteLLM_DailyToolSpend rollup in one transaction. +""" + +from types import SimpleNamespace +from typing import Any +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from litellm.proxy.db.spend_log_tool_index import ( + ToolUsageTransaction, + build_tool_usage_transaction, + flush_tool_usage_transactions, + response_tool_call_names, +) + + +def _response_with_tool_calls(*names: str) -> SimpleNamespace: + tool_calls = [SimpleNamespace(function=SimpleNamespace(name=name)) for name in names] + return SimpleNamespace(choices=[SimpleNamespace(message=SimpleNamespace(tool_calls=tool_calls))]) + + +class _FakeBatcher: + def __init__(self) -> None: + self.litellm_spendlogtoolindex = MagicMock() + self.litellm_dailytoolspend = MagicMock() + + async def __aenter__(self) -> "_FakeBatcher": + return self + + async def __aexit__(self, *args: Any) -> None: + return None + + +def _prisma_with_batcher() -> tuple[MagicMock, _FakeBatcher]: + batcher = _FakeBatcher() + prisma = MagicMock() + prisma.db.batch_ = MagicMock(return_value=batcher) + return prisma, batcher + + +class TestBuildToolUsageTransaction: + def test_declared_tools_never_reach_the_transaction(self): + # Regression for the inflation bug: the builder's only non-MCP source is + # the response's tool_calls, so a request declaring N tools while the + # model invokes one produces exactly one attribution. + transaction = build_tool_usage_transaction( + request_id="r1", + start_time_iso="2026-07-25T10:00:00+00:00", + mcp_namespaced_tool_name=None, + spend=0.5, + total_tokens=100, + completion_response=_response_with_tool_calls("get_weather"), + ) + assert transaction is not None + assert transaction.tool_names == ("get_weather",) + + def test_no_invoked_tools_returns_none(self): + assert ( + build_tool_usage_transaction( + request_id="r1", + start_time_iso="2026-07-25T10:00:00+00:00", + mcp_namespaced_tool_name=None, + spend=0.5, + total_tokens=100, + completion_response=SimpleNamespace(choices=[SimpleNamespace(message=SimpleNamespace(tool_calls=None))]), + ) + is None + ) + + def test_mcp_name_and_response_names_dedupe(self): + transaction = build_tool_usage_transaction( + request_id="r1", + start_time_iso="2026-07-25T10:00:00+00:00", + mcp_namespaced_tool_name="srv/tool_a", + spend=0.5, + total_tokens=100, + completion_response=_response_with_tool_calls("srv/tool_a", "tool_b", "tool_b"), + ) + assert transaction is not None + assert transaction.tool_names == ("srv/tool_a", "tool_b") + + def test_date_matches_daily_spend_writer_derivation(self): + # The daily spend writer derives its date bucket as + # payload["startTime"].split("T")[0] (db_spend_update_writer.py), i.e. the + # timestamp's own calendar date, NOT the astimezone-UTC date. A non-UTC + # isoformat pins the difference: 2026-07-25T22:00:00-07:00 is 2026-07-26 + # in UTC but must bucket as 2026-07-25 to match LiteLLM_DailyUserSpend. + start_time_iso = "2026-07-25T22:00:00-07:00" + transaction = build_tool_usage_transaction( + request_id="r1", + start_time_iso=start_time_iso, + mcp_namespaced_tool_name="srv/tool_a", + spend=0.5, + total_tokens=100, + completion_response=None, + ) + assert transaction is not None + assert transaction.date == start_time_iso.split("T")[0] == "2026-07-25" + + def test_realtime_tool_calls_reach_the_transaction(self): + # Realtime sessions carry invoked tools in kwargs["realtime_tool_calls"] + # (OpenAI tool_calls dict shape, built in realtime_streaming.py), not on a + # response object; they must land in the rollup like any other invocation. + realtime_tool_calls = [ + {"id": "call_1", "type": "function", "function": {"name": "rt_get_weather", "arguments": "{}"}}, + ] + transaction = build_tool_usage_transaction( + request_id="r1", + start_time_iso="2026-07-25T10:00:00+00:00", + mcp_namespaced_tool_name=None, + spend=0.5, + total_tokens=100, + completion_response=None, + realtime_tool_calls=realtime_tool_calls, + ) + assert transaction is not None + assert transaction.tool_names == ("rt_get_weather",) + + def test_realtime_names_dedupe_against_response_names(self): + realtime_tool_calls = [{"type": "function", "function": {"name": "get_weather", "arguments": "{}"}}] + transaction = build_tool_usage_transaction( + request_id="r1", + start_time_iso="2026-07-25T10:00:00+00:00", + mcp_namespaced_tool_name=None, + spend=0.5, + total_tokens=100, + completion_response=_response_with_tool_calls("get_weather"), + realtime_tool_calls=realtime_tool_calls, + ) + assert transaction is not None + assert transaction.tool_names == ("get_weather",) + + def test_n_greater_than_one_tools_from_every_choice_reach_the_transaction(self): + # Regression: an n>1 request pays for every choice, and a tool invoked + # only in a later choice really ran; it must not be dropped because the + # extractor read choices[0] alone. + from types import SimpleNamespace as NS + + response = NS( + choices=[ + NS(message=NS(tool_calls=[NS(function=NS(name="tool_alpha"))])), + NS(message=NS(tool_calls=[NS(function=NS(name="tool_beta"))])), + ] + ) + transaction = build_tool_usage_transaction( + request_id="r1", + start_time_iso="2026-07-25T10:00:00+00:00", + mcp_namespaced_tool_name=None, + spend=0.5, + total_tokens=100, + completion_response=response, + ) + assert transaction is not None + assert transaction.tool_names == ("tool_alpha", "tool_beta") + + def test_unparseable_start_time_returns_none(self): + assert ( + build_tool_usage_transaction( + request_id="r1", + start_time_iso="not-a-timestamp", + mcp_namespaced_tool_name="srv/tool_a", + spend=0.5, + total_tokens=100, + completion_response=None, + ) + is None + ) + + +class TestResponseToolCallNames: + def test_unrecognized_shapes_yield_nothing(self): + assert response_tool_call_names(None) == () + assert response_tool_call_names(SimpleNamespace()) == () + assert response_tool_call_names(ValueError("boom")) == () + + def test_blank_names_are_dropped(self): + assert response_tool_call_names(_response_with_tool_calls(" ", "real_tool")) == ("real_tool",) + + def test_responses_api_output_function_calls(self): + # Regression: /v1/responses carries invocations in output[] items of + # type function_call, not in choices; they must reach the rollup. + response = SimpleNamespace( + output=[ + SimpleNamespace(type="function_call", name="get_weather", call_id="c1", arguments="{}"), + SimpleNamespace(type="message", name=None, call_id=None, arguments=None), + ] + ) + assert response_tool_call_names(response) == ("get_weather",) + + def test_anthropic_messages_tool_use_blocks(self): + response = { + "content": [ + {"type": "text", "text": "checking"}, + {"type": "tool_use", "id": "t1", "name": "ant_get_weather", "input": {"city": "Paris"}}, + ] + } + assert response_tool_call_names(response) == ("ant_get_weather",) + + +def _transaction( + request_id: str, + date: str = "2026-07-25", + tool_names: tuple = ("tool_a",), + spend: float = 1.0, + total_tokens: int = 10, +) -> ToolUsageTransaction: + from datetime import datetime, timezone + + return ToolUsageTransaction( + request_id=request_id, + date=date, + start_time=datetime(2026, 7, 25, 10, 0, tzinfo=timezone.utc), + tool_names=tool_names, + spend=spend, + total_tokens=total_tokens, + ) + + +class TestFlushToolUsageTransactions: + @pytest.mark.asyncio + async def test_multi_tool_request_attributes_full_spend_to_each_tool(self): + prisma, batcher = _prisma_with_batcher() + await flush_tool_usage_transactions( + prisma_client=prisma, + transactions=[_transaction("r1", tool_names=("tool_a", "tool_b"), spend=0.10, total_tokens=100)], + ) + index_rows = batcher.litellm_spendlogtoolindex.create_many.call_args.kwargs["data"] + assert [(r["request_id"], r["tool_name"]) for r in index_rows] == [("r1", "tool_a"), ("r1", "tool_b")] + assert batcher.litellm_spendlogtoolindex.create_many.call_args.kwargs["skip_duplicates"] is True + + upserts = { + c.kwargs["where"]["date_tool_name"]["tool_name"]: c.kwargs["data"] + for c in batcher.litellm_dailytoolspend.upsert.call_args_list + } + assert set(upserts) == {"tool_a", "tool_b"} + for data in upserts.values(): + assert data["create"]["spend"] == 0.10 + assert data["create"]["request_count"] == 1 + assert data["update"]["spend"] == {"increment": 0.10} + assert data["update"]["request_count"] == {"increment": 1} + + @pytest.mark.asyncio + async def test_same_day_same_tool_aggregates_within_batch(self): + prisma, batcher = _prisma_with_batcher() + await flush_tool_usage_transactions( + prisma_client=prisma, + transactions=[ + _transaction("r1", spend=0.10, total_tokens=100), + _transaction("r2", spend=0.30, total_tokens=200), + ], + ) + assert batcher.litellm_dailytoolspend.upsert.call_count == 1 + data = batcher.litellm_dailytoolspend.upsert.call_args.kwargs["data"] + assert data["create"] == { + "date": "2026-07-25", + "tool_name": "tool_a", + "spend": pytest.approx(0.40), + "total_tokens": 300, + "request_count": 2, + } + assert data["update"]["spend"] == {"increment": pytest.approx(0.40)} + assert data["update"]["total_tokens"] == {"increment": 300} + assert data["update"]["request_count"] == {"increment": 2} + + @pytest.mark.asyncio + async def test_index_rows_and_rollup_share_one_transaction(self): + # Both writes go through the same batch_() so a failed flush cannot leave + # index rows without their rollup increments (or vice versa); increments + # are not idempotent, so partial states must be unreachable. + prisma, batcher = _prisma_with_batcher() + await flush_tool_usage_transactions( + prisma_client=prisma, + transactions=[_transaction("r1")], + ) + prisma.db.batch_.assert_called_once() + batcher.litellm_spendlogtoolindex.create_many.assert_called_once() + batcher.litellm_dailytoolspend.upsert.assert_called_once() + + @pytest.mark.asyncio + async def test_empty_batch_touches_nothing(self): + prisma, _ = _prisma_with_batcher() + await flush_tool_usage_transactions(prisma_client=prisma, transactions=[]) + prisma.db.batch_.assert_not_called() + + @pytest.mark.asyncio + async def test_connection_errors_retry_and_succeed(self, monkeypatch): + # A failed batch commits nothing, so retrying a connection error cannot + # double-count; the flush must retry rather than drop the batch. + import httpx + + batcher = _FakeBatcher() + prisma = MagicMock() + prisma.db.batch_ = MagicMock(side_effect=[httpx.ConnectError("down"), batcher]) + sleeps: list[float] = [] + + async def fake_sleep(seconds: float) -> None: + sleeps.append(seconds) + + monkeypatch.setattr("litellm.proxy.db.spend_log_tool_index.asyncio.sleep", fake_sleep) + await flush_tool_usage_transactions(prisma_client=prisma, transactions=[_transaction("r1")]) + assert prisma.db.batch_.call_count == 2 + assert len(sleeps) == 1 + batcher.litellm_dailytoolspend.upsert.assert_called_once() + + @pytest.mark.asyncio + async def test_connection_errors_exhaust_retries_then_raise(self, monkeypatch): + import httpx + + prisma = MagicMock() + prisma.db.batch_ = MagicMock(side_effect=httpx.ConnectError("down")) + + async def fake_sleep(seconds: float) -> None: + return None + + monkeypatch.setattr("litellm.proxy.db.spend_log_tool_index.asyncio.sleep", fake_sleep) + with pytest.raises(httpx.ConnectError): + await flush_tool_usage_transactions( + prisma_client=prisma, transactions=[_transaction("r1")], n_retry_times=2 + ) + assert prisma.db.batch_.call_count == 3 + + @pytest.mark.asyncio + async def test_non_connection_errors_do_not_retry(self): + prisma = MagicMock() + prisma.db.batch_ = MagicMock(side_effect=ValueError("bad data")) + with pytest.raises(ValueError): + await flush_tool_usage_transactions(prisma_client=prisma, transactions=[_transaction("r1")]) + prisma.db.batch_.assert_called_once() + + @pytest.mark.asyncio + @pytest.mark.parametrize("ambiguous_error", ["ReadTimeout", "ReadError"]) + async def test_post_send_ambiguous_errors_drop_without_retry(self, ambiguous_error): + # A ReadTimeout means the statements were sent and the outcome is + # unknown; the engine can leave the transaction open on the pooled + # connection, so a retry's statements would stack into it and one + # commit would apply both increment sets. These must never retry. + import httpx + + error = getattr(httpx, ambiguous_error)("ambiguous") + prisma = MagicMock() + prisma.db.batch_ = MagicMock(side_effect=error) + with pytest.raises((httpx.ReadTimeout, httpx.ReadError)): + await flush_tool_usage_transactions(prisma_client=prisma, transactions=[_transaction("r1")]) + prisma.db.batch_.assert_called_once() diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py index 4dc527ca45d..248893ed153 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py @@ -1551,3 +1551,234 @@ async def test_apply_guardrail_litellm_timeout_fail_open_forwards_uncompressed() ) assert result["structured_messages"] == ORIGINAL_MESSAGES + + + + +# --------------------------------------------------------------------------- +# Content-parts flattening (LIT-4795) +# +# Anthropic-format requests translate to messages whose content is a list of +# part dicts. The compression service only rewrites string content, so the +# guardrail flattens ALL-TEXT part lists on the wire and restores the +# original shapes afterwards. Rows with non-text parts are never flattened: +# cache_control breakpoints are positional, and merging text across a +# non-text part would move a later breakpoint to the other side of it. +# --------------------------------------------------------------------------- + +PARTS_MESSAGES = [ + { + "role": "system", + "content": [ + {"type": "text", "text": "You are Claude Code.", "cache_control": {"type": "ephemeral"}}, + { + "type": "text", + "text": "Second system block. " + "B" * 5000, + "cache_control": {"type": "ephemeral", "ttl": "1h"}, + }, + ], + }, + { + "role": "user", + "content": [ + {"type": "text", "text": "Mixed row text."}, + {"type": "image_url", "image_url": {"url": "https://example.com/x.png"}}, + ], + }, + {"role": "tool", "content": "tool output " + "C" * 500}, +] + +FLATTENED_SYSTEM_TEXT = "You are Claude Code.\n\nSecond system block. " + "B" * 5000 + + +def _parts_copy() -> list: + return json.loads(json.dumps(PARTS_MESSAGES)) + + +def _echo_wire_view() -> list: + """What the service receives (and echoes back when it changes nothing).""" + return [ + {"role": "system", "content": FLATTENED_SYSTEM_TEXT}, + json.loads(json.dumps(PARTS_MESSAGES[1])), + {"role": "tool", "content": "tool output " + "C" * 500}, + ] + + +@pytest.mark.asyncio +async def test_apply_guardrail_flattens_all_text_rows_only( + guardrail: HeadroomGuardrail, +): + inputs = GenericGuardrailAPIInputs( + texts=["B" * 5000], + structured_messages=_parts_copy(), + ) + mock_response = _make_compress_response(_echo_wire_view()) + + with patch.object( + guardrail.async_handler, + "post", + new_callable=AsyncMock, + return_value=mock_response, + ) as mock_post: + await guardrail.apply_guardrail( + inputs=inputs, + request_data={"model": "claude-fable-5"}, + input_type="request", + ) + + wire_messages = mock_post.call_args.kwargs["json"]["messages"] + assert wire_messages[0]["content"] == FLATTENED_SYSTEM_TEXT + # Mixed text+image row is never flattened: merging its text would move a + # later cache_control breakpoint across the image part. + assert isinstance(wire_messages[1]["content"], list) + assert wire_messages[2]["content"] == "tool output " + "C" * 500 + + +@pytest.mark.asyncio +async def test_apply_guardrail_restores_rewritten_all_text_row( + guardrail: HeadroomGuardrail, +): + inputs = GenericGuardrailAPIInputs( + texts=["B" * 5000], + structured_messages=_parts_copy(), + ) + compressed = _echo_wire_view() + compressed[0]["content"] = "compressed system. Retrieve more: hash=b573993006976af767214fac" + mock_response = _make_compress_response(compressed) + + with patch.object( + guardrail.async_handler, + "post", + new_callable=AsyncMock, + return_value=mock_response, + ): + result = await guardrail.apply_guardrail( + inputs=inputs, + request_data={"model": "claude-fable-5"}, + input_type="request", + ) + + messages = result["structured_messages"] + system_content = messages[0]["content"] + # Rewritten all-text row collapses to one part carrying the LAST declared + # breakpoint: an Anthropic breakpoint caches the prefix ending at its + # part, so after the merge the last one (and its TTL) still describes the + # row. + assert isinstance(system_content, list) + assert len(system_content) == 1 + assert system_content[0]["text"] == "compressed system. Retrieve more: hash=b573993006976af767214fac" + assert system_content[0]["cache_control"] == {"type": "ephemeral", "ttl": "1h"} + # Mixed row passes through byte-identical. + assert messages[1]["content"] == PARTS_MESSAGES[1]["content"] + # Hashes inside restored parts still drive retrieve-tool injection. + assert has_headroom_retrieve_tool(result.get("tools") or []) + + +@pytest.mark.asyncio +async def test_apply_guardrail_keeps_originals_when_service_echoes_unchanged( + guardrail: HeadroomGuardrail, +): + inputs = GenericGuardrailAPIInputs( + texts=["B" * 5000], + structured_messages=_parts_copy(), + ) + mock_response = _make_compress_response(_echo_wire_view()) + + with patch.object( + guardrail.async_handler, + "post", + new_callable=AsyncMock, + return_value=mock_response, + ): + result = await guardrail.apply_guardrail( + inputs=inputs, + request_data={"model": "claude-fable-5"}, + input_type="request", + ) + + messages = result["structured_messages"] + assert [m["content"] for m in messages] == [m["content"] for m in PARTS_MESSAGES] + + +@pytest.mark.asyncio +async def test_apply_guardrail_adopts_service_output_when_rows_dropped( + guardrail: HeadroomGuardrail, +): + inputs = GenericGuardrailAPIInputs( + texts=["B" * 5000], + structured_messages=_parts_copy(), + ) + dropped = [ + {"role": "system", "content": FLATTENED_SYSTEM_TEXT}, + {"role": "user", "content": "B" * 50}, + ] + mock_response = _make_compress_response(dropped) + + with patch.object( + guardrail.async_handler, + "post", + new_callable=AsyncMock, + return_value=mock_response, + ): + result = await guardrail.apply_guardrail( + inputs=inputs, + request_data={"model": "claude-fable-5"}, + input_type="request", + ) + + assert result["structured_messages"] == dropped + + +@pytest.mark.asyncio +async def test_apply_guardrail_sends_textless_parts_rows_unflattened( + guardrail: HeadroomGuardrail, +): + image_only = [ + {"role": "user", "content": [{"type": "image_url", "image_url": {"url": "https://example.com/x.png"}}]}, + {"role": "user", "content": "D" * 5000}, + ] + inputs = GenericGuardrailAPIInputs( + texts=["D" * 5000], + structured_messages=json.loads(json.dumps(image_only)), + ) + mock_response = _make_compress_response(json.loads(json.dumps(image_only))) + + with patch.object( + guardrail.async_handler, + "post", + new_callable=AsyncMock, + return_value=mock_response, + ) as mock_post: + await guardrail.apply_guardrail( + inputs=inputs, + request_data={"model": "claude-fable-5"}, + input_type="request", + ) + + wire_messages = mock_post.call_args.kwargs["json"]["messages"] + assert isinstance(wire_messages[0]["content"], list) + assert wire_messages[1]["content"] == "D" * 5000 + + +@pytest.mark.asyncio +async def test_fail_open_returns_original_parts_shapes(): + guardrail = _make_guardrail(unreachable_fallback="fail_open") + inputs = GenericGuardrailAPIInputs( + texts=["B" * 5000], + structured_messages=_parts_copy(), + ) + + with patch.object( + guardrail.async_handler, + "post", + new_callable=AsyncMock, + side_effect=httpx.ConnectError("boom"), + ): + result = await guardrail.apply_guardrail( + inputs=inputs, + request_data={}, + input_type="request", + ) + + messages = result["structured_messages"] + assert [m["content"] for m in messages] == [m["content"] for m in PARTS_MESSAGES] diff --git a/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py index be0267d69c5..5cbc3e72d83 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py @@ -3605,3 +3605,65 @@ async def test_add_new_user_to_default_team_string_teams_have_no_member_budget(m assert mock_add.call_args.kwargs["max_budget_in_team"] is None assert mock_add.call_args.kwargs["team_id"] == "string-team" + + +@pytest.mark.asyncio +async def test_add_user_to_team_logs_unknown_team_at_error(mocker, caplog): + """A default team that no longer exists makes every membership write 404. + + The failure is swallowed so user creation still succeeds, so the log line is + the only signal an operator gets; it must be ERROR and name the team. + """ + import logging + + from fastapi import HTTPException + + from litellm.proxy.management_endpoints.internal_user_endpoints import ( + _add_user_to_team, + ) + + mocker.patch( + "litellm.proxy.management_endpoints.team_endpoints.team_member_add", + new_callable=mocker.AsyncMock, + side_effect=HTTPException(status_code=404, detail={"error": "Team not found"}), + ) + + with caplog.at_level(logging.DEBUG, logger="LiteLLM Proxy"): + await _add_user_to_team( + user_id="sso-user", + team_id="deleted-team", + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + ) + + errors = [r.getMessage() for r in caplog.records if r.levelno >= logging.ERROR] + assert len(errors) == 1, f"expected exactly one ERROR log, got {errors}" + assert "deleted-team" in errors[0] + assert "sso-user" in errors[0] + + +@pytest.mark.asyncio +async def test_add_user_to_team_keeps_already_a_member_quiet(mocker, caplog): + """Re-adding an existing member is expected on every login and must not + produce an ERROR, otherwise the real failures above are lost in the noise.""" + import logging + + from fastapi import HTTPException + + from litellm.proxy.management_endpoints.internal_user_endpoints import ( + _add_user_to_team, + ) + + mocker.patch( + "litellm.proxy.management_endpoints.team_endpoints.team_member_add", + new_callable=mocker.AsyncMock, + side_effect=HTTPException(status_code=400, detail={"error": "User already exists in team"}), + ) + + with caplog.at_level(logging.DEBUG, logger="LiteLLM Proxy"): + await _add_user_to_team( + user_id="sso-user", + team_id="existing-team", + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + ) + + assert [r.getMessage() for r in caplog.records if r.levelno >= logging.ERROR] == [] diff --git a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py index f3e5e2c9b71..bd5eda0197b 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py @@ -636,14 +636,33 @@ class TestDeleteModelClearsRouterRegistry: not just from model_list, or a stale (now unbacked) router entry lingers until restart. """ + @staticmethod + def _complexity_router_deployment(model_id: str, tags: list | None = None) -> dict: + return { + "model_name": "smart-router", + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_config": {"tiers": {"SIMPLE": "gpt-4o-mini", "MEDIUM": "gpt-4o"}}, + "complexity_router_default_model": "gpt-4o", + **({"tags": tags} if tags else {}), + }, + "model_info": {"id": model_id, "db_model": True}, + } + @pytest.mark.asyncio - async def test_delete_model_pops_router_registries(self): + async def test_delete_model_releases_only_the_deleted_routers_slot(self): + """Deleting one tagged router must release its own slot and leave a sibling + sharing the model_name registered. A blanket pop(model_name) here would take + both down, and nothing reloads on the delete path to restore the survivor. + """ + import litellm + from litellm.proxy.management_endpoints.model_management_endpoints import ModelInfoDelete from litellm.proxy.management_endpoints.model_management_endpoints import ( delete_model as delete_model_endpoint, ) - from litellm.proxy.management_endpoints.model_management_endpoints import ModelInfoDelete model_id = "router-del-1" + surviving_id = "router-del-2" admin_user = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN) db_row = LiteLLM_ProxyModelTable( model_id=model_id, @@ -660,16 +679,16 @@ class TestDeleteModelClearsRouterRegistry: mock_prisma.db.litellm_proxymodeltable.find_unique = AsyncMock(return_value=db_row) mock_prisma.db.litellm_proxymodeltable.delete = AsyncMock(return_value=db_row) - mock_router = MagicMock() - mock_router.delete_deployment = MagicMock( - return_value={ - "model_name": "smart-router", - "litellm_params": {"model": "auto_router/complexity_router"}, - "model_info": {"id": model_id}, - } + real_router = litellm.Router( + model_list=[ + {"model_name": "gpt-4o", "litellm_params": {"model": "gpt-4o"}}, + {"model_name": "gpt-4o-mini", "litellm_params": {"model": "gpt-4o-mini"}}, + self._complexity_router_deployment(model_id, tags=["team-a"]), + self._complexity_router_deployment(surviving_id, tags=["team-b"]), + ], + ignore_invalid_deployments=True, ) - mock_router.auto_routers = {"smart-router": MagicMock()} - mock_router.complexity_routers = {"smart-router": MagicMock()} + assert len(real_router.complexity_routers["smart-router"]) == 2 _PS = "litellm.proxy.proxy_server" with ( @@ -679,16 +698,17 @@ class TestDeleteModelClearsRouterRegistry: patch(f"{_PS}.proxy_logging_obj", MagicMock()), patch(f"{_PS}.general_settings", {}), patch(f"{_PS}.premium_user", True), - patch(f"{_PS}.llm_router", mock_router), + patch(f"{_PS}.llm_router", real_router), ): await delete_model_endpoint( model_info=ModelInfoDelete(id=model_id), user_api_key_dict=admin_user, ) - mock_router.delete_deployment.assert_called_once_with(id=model_id) - assert "smart-router" not in mock_router.auto_routers - assert "smart-router" not in mock_router.complexity_routers + assert model_id not in [m["model_info"]["id"] for m in real_router.model_list] + surviving = real_router.complexity_routers["smart-router"] + assert len(surviving) == 1 + assert surviving[0].tags == ("team-b",) @pytest.mark.asyncio async def test_delete_regular_model_preserves_config_router_sharing_name(self): diff --git a/tests/test_litellm/proxy/management_endpoints/test_tool_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_tool_management_endpoints.py index c908250fa64..45c3c6c2466 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_tool_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_tool_management_endpoints.py @@ -19,11 +19,7 @@ from fastapi.testclient import TestClient sys.path.insert(0, os.path.abspath("../../..")) -from litellm.proxy.management_endpoints.tool_management_endpoints import ( - _build_tool_spend_response, - _ToolSpendRow, - router, -) +from litellm.proxy.management_endpoints.tool_management_endpoints import router from litellm.types.tool_management import LiteLLM_ToolTableRow # --- helpers --- @@ -64,6 +60,30 @@ def _override_auth(): _MOCK_PRISMA = MagicMock() +def _rollup_row(date: str, tool_name: str, spend: float, request_count: int, total_tokens: int) -> MagicMock: + row = MagicMock() + row.date = date + row.tool_name = tool_name + row.spend = spend + row.request_count = request_count + row.total_tokens = total_tokens + return row + + +def _group_row(tool_name: str, spend: float, request_count: int, total_tokens: int) -> dict: + return {"tool_name": tool_name, "_sum": {"spend": spend, "total_tokens": total_tokens, "request_count": request_count}} + + +def _rollup_prisma(group_rows: list, daily_rows: list | None = None) -> MagicMock: + prisma = MagicMock() + prisma.db.query_raw = AsyncMock(return_value=[]) + prisma.db.litellm_spendlogs.find_many = AsyncMock(return_value=[]) + prisma.db.litellm_spendlogtoolindex.find_many = AsyncMock(return_value=[]) + prisma.db.litellm_dailytoolspend.group_by = AsyncMock(return_value=group_rows) + prisma.db.litellm_dailytoolspend.find_many = AsyncMock(return_value=daily_rows or []) + return prisma + + # --- test class --- @@ -154,21 +174,23 @@ class TestToolManagementEndpoints: assert resp.status_code == 422 def test_tool_spend_route_not_shadowed_by_get_tool(self): - prisma = MagicMock() - prisma.db.query_raw = AsyncMock(return_value=[]) + prisma = _rollup_prisma([]) with patch("litellm.proxy.proxy_server.prisma_client", prisma): resp = self.client.get("/v1/tool/spend") assert resp.status_code == 200 assert resp.json()["by_tool"] == [] - def test_tool_spend_aggregates_and_sorts(self): - rows = [ - {"date": "2026-07-01", "tool_name": "search", "call_count": 2, "spend": 1.0, "total_tokens": 100}, - {"date": "2026-07-02", "tool_name": "search", "call_count": 1, "spend": 4.0, "total_tokens": 50}, - {"date": "2026-07-01", "tool_name": "read_file", "call_count": 3, "spend": 2.0, "total_tokens": 300}, + def test_tool_spend_serves_sql_aggregates_and_daily_series(self): + group_rows = [ + _group_row("search", spend=5.0, request_count=3, total_tokens=150), + _group_row("read_file", spend=2.0, request_count=3, total_tokens=300), ] - prisma = MagicMock() - prisma.db.query_raw = AsyncMock(side_effect=[rows, [{"total_spend": 5.5}]]) + daily_rows = [ + _rollup_row("2026-07-01", "search", spend=1.0, request_count=2, total_tokens=100), + _rollup_row("2026-07-01", "read_file", spend=2.0, request_count=3, total_tokens=300), + _rollup_row("2026-07-02", "search", spend=4.0, request_count=1, total_tokens=50), + ] + prisma = _rollup_prisma(group_rows, daily_rows) with patch("litellm.proxy.proxy_server.prisma_client", prisma): resp = self.client.get("/v1/tool/spend?start_date=2026-07-01&end_date=2026-07-02") assert resp.status_code == 200 @@ -179,94 +201,89 @@ class TestToolManagementEndpoints: assert search["call_count"] == 3 assert search["total_tokens"] == 150 assert len(body["daily"]) == 3 + assert body["daily"][0]["call_count"] == 2 assert body["start_date"] == "2026-07-01" assert body["end_date"] == "2026-07-02" - assert body["total_spend"] == 5.5 + + def test_tool_spend_coerces_bigint_string_sums(self): + # prisma group_by returns BigInt sums as strings ("808"); the response + # must coerce them to ints rather than 500 on validation. + group_rows = [{"tool_name": "search", "_sum": {"spend": 0.5, "total_tokens": "808", "request_count": "3"}}] + prisma = _rollup_prisma(group_rows) + with patch("litellm.proxy.proxy_server.prisma_client", prisma): + resp = self.client.get("/v1/tool/spend?start_date=2026-07-01&end_date=2026-07-02") + assert resp.status_code == 200 + assert resp.json()["by_tool"][0]["total_tokens"] == 808 + assert resp.json()["by_tool"][0]["call_count"] == 3 + + def test_tool_spend_daily_restricted_to_top_tools_and_capped(self): + from litellm.constants import TOOL_SPEND_TOP_TOOLS + + group_rows = [_group_row("search", spend=5.0, request_count=1, total_tokens=10)] + prisma = _rollup_prisma(group_rows) + with patch("litellm.proxy.proxy_server.prisma_client", prisma): + resp = self.client.get("/v1/tool/spend?start_date=2026-07-01&end_date=2026-07-02") + assert resp.status_code == 200 + group_kwargs = prisma.db.litellm_dailytoolspend.group_by.await_args.kwargs + assert group_kwargs["take"] == TOOL_SPEND_TOP_TOOLS + assert group_kwargs["order"] == {"_sum": {"spend": "desc"}} + daily_where = prisma.db.litellm_dailytoolspend.find_many.await_args.kwargs["where"] + assert daily_where["tool_name"] == {"in": ["search"]} + + def test_tool_spend_skips_daily_query_when_no_tools(self): + prisma = _rollup_prisma([]) + with patch("litellm.proxy.proxy_server.prisma_client", prisma): + resp = self.client.get("/v1/tool/spend?start_date=2026-07-01&end_date=2026-07-02") + assert resp.status_code == 200 + prisma.db.litellm_dailytoolspend.find_many.assert_not_awaited() @patch("litellm.proxy.proxy_server.prisma_client", None) def test_tool_spend_no_db_returns_500(self): resp = self.client.get("/v1/tool/spend") assert resp.status_code == 500 - def test_tool_spend_end_date_is_inclusive_via_exclusive_next_day_bound(self): - prisma = MagicMock() - prisma.db.query_raw = AsyncMock(return_value=[]) + def test_tool_spend_reads_rollup_only_never_spendlogs(self): + # Regression for the GA blocker: the dashboard aggregate must be served + # entirely from LiteLLM_DailyToolSpend; any query_raw or SpendLogs table + # access on this path reintroduces the per-request scan. + prisma = _rollup_prisma([]) with patch("litellm.proxy.proxy_server.prisma_client", prisma): resp = self.client.get("/v1/tool/spend?start_date=2026-07-01&end_date=2026-07-02") assert resp.status_code == 200 - expected_binds = ( - datetime(2026, 7, 1, tzinfo=timezone.utc).isoformat(), - datetime(2026, 7, 3, tzinfo=timezone.utc).isoformat(), - ) - assert prisma.db.query_raw.await_count == 2 - for call in prisma.db.query_raw.await_args_list: - assert tuple(call.args[1:]) == expected_binds + prisma.db.query_raw.assert_not_awaited() + prisma.db.litellm_spendlogs.find_many.assert_not_awaited() + prisma.db.litellm_spendlogtoolindex.find_many.assert_not_awaited() + prisma.db.litellm_dailytoolspend.group_by.assert_awaited_once() + + def test_tool_spend_windows_rollup_by_inclusive_date_strings(self): + prisma = _rollup_prisma([]) + with patch("litellm.proxy.proxy_server.prisma_client", prisma): + resp = self.client.get("/v1/tool/spend?start_date=2026-07-01&end_date=2026-07-02") + assert resp.status_code == 200 + where = prisma.db.litellm_dailytoolspend.group_by.await_args.kwargs["where"] + assert where == {"date": {"gte": "2026-07-01", "lte": "2026-07-02"}} assert resp.json()["end_date"] == "2026-07-02" - def test_tool_spend_start_clamped_to_30_days_before_end(self): - # Clamped floor is end_date minus 30 days, serving up to 31 calendar dates - # inclusive: deliberately the same width as the endpoint's default window, - # so the dashboard's default range never triggers the clamp. - prisma = MagicMock() - prisma.db.query_raw = AsyncMock(return_value=[]) + def test_tool_spend_wide_range_served_fully(self): + # Regression: the 30-day clamp is gone; a 182-day request is served as + # requested because the rollup read is O(tools x dates). + prisma = _rollup_prisma([]) with patch("litellm.proxy.proxy_server.prisma_client", prisma): resp = self.client.get("/v1/tool/spend?start_date=2026-01-01&end_date=2026-07-01") assert resp.status_code == 200 - expected_binds = ( - datetime(2026, 6, 1, tzinfo=timezone.utc).isoformat(), - datetime(2026, 7, 2, tzinfo=timezone.utc).isoformat(), - ) - assert prisma.db.query_raw.await_count == 2 - for call in prisma.db.query_raw.await_args_list: - assert tuple(call.args[1:]) == expected_binds - assert resp.json()["start_date"] == "2026-06-01" + where = prisma.db.litellm_dailytoolspend.group_by.await_args.kwargs["where"] + assert where == {"date": {"gte": "2026-01-01", "lte": "2026-07-01"}} + assert resp.json()["start_date"] == "2026-01-01" assert resp.json()["end_date"] == "2026-07-01" - def test_tool_spend_range_within_cap_is_not_clamped(self): - prisma = MagicMock() - prisma.db.query_raw = AsyncMock(return_value=[]) + def test_tool_spend_defaults_to_trailing_30_days(self): + prisma = _rollup_prisma([]) with patch("litellm.proxy.proxy_server.prisma_client", prisma): - resp = self.client.get("/v1/tool/spend?start_date=2026-06-25&end_date=2026-07-01") + resp = self.client.get("/v1/tool/spend") assert resp.status_code == 200 - for call in prisma.db.query_raw.await_args_list: - assert call.args[1] == datetime(2026, 6, 25, tzinfo=timezone.utc).isoformat() - assert resp.json()["start_date"] == "2026-06-25" - - def test_tool_spend_start_honored_when_end_date_omitted(self): - # Regression: with end_date omitted the floor anchors to today's UTC - # midnight, not now's time-of-day, so an explicit start_date exactly 30 - # days back is served from midnight rather than truncated to mid-day. - prisma = MagicMock() - prisma.db.query_raw = AsyncMock(return_value=[]) - floor_day = datetime.now(timezone.utc).replace(hour=0, minute=0, second=0, microsecond=0) - timedelta(days=30) - with patch("litellm.proxy.proxy_server.prisma_client", prisma): - resp = self.client.get(f"/v1/tool/spend?start_date={floor_day.strftime('%Y-%m-%d')}") - assert resp.status_code == 200 - for call in prisma.db.query_raw.await_args_list: - assert call.args[1] == floor_day.isoformat() - assert resp.json()["start_date"] == floor_day.strftime("%Y-%m-%d") - - def test_tool_spend_clamp_without_end_date_lands_on_midnight(self): - prisma = MagicMock() - prisma.db.query_raw = AsyncMock(return_value=[]) - floor_day = datetime.now(timezone.utc).replace(hour=0, minute=0, second=0, microsecond=0) - timedelta(days=30) - with patch("litellm.proxy.proxy_server.prisma_client", prisma): - resp = self.client.get("/v1/tool/spend?start_date=2020-01-01") - assert resp.status_code == 200 - for call in prisma.db.query_raw.await_args_list: - assert call.args[1] == floor_day.isoformat() - assert resp.json()["start_date"] == floor_day.strftime("%Y-%m-%d") - - def test_tool_spend_total_query_bounds_outer_spendlogs_scan(self): - prisma = MagicMock() - prisma.db.query_raw = AsyncMock(return_value=[]) - with patch("litellm.proxy.proxy_server.prisma_client", prisma): - resp = self.client.get("/v1/tool/spend?start_date=2026-07-01&end_date=2026-07-02") - assert resp.status_code == 200 - for call in prisma.db.query_raw.await_args_list: - sql = call.args[0] - assert 'sl."startTime" >=' in sql - assert 'sl."startTime" <' in sql + today = datetime.now(timezone.utc) + assert resp.json()["end_date"] == today.strftime("%Y-%m-%d") + assert resp.json()["start_date"] == (today - timedelta(days=30)).strftime("%Y-%m-%d") @pytest.mark.parametrize( "query", @@ -279,13 +296,12 @@ class TestToolManagementEndpoints: ], ) def test_tool_spend_malformed_date_returns_400(self, query: str): - prisma = MagicMock() - prisma.db.query_raw = AsyncMock(return_value=[]) + prisma = _rollup_prisma([]) with patch("litellm.proxy.proxy_server.prisma_client", prisma): resp = self.client.get(f"/v1/tool/spend?{query}") assert resp.status_code == 400 assert "Invalid date format" in resp.json()["detail"] - prisma.db.query_raw.assert_not_awaited() + prisma.db.litellm_dailytoolspend.group_by.assert_not_awaited() def test_tool_spend_non_admin_returns_403(self): from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth @@ -296,38 +312,8 @@ class TestToolManagementEndpoints: api_key="sk-user", user_id="u1", user_role=LitellmUserRoles.INTERNAL_USER ) client = TestClient(app, raise_server_exceptions=True) - prisma = MagicMock() - prisma.db.query_raw = AsyncMock(return_value=[]) + prisma = _rollup_prisma([]) with patch("litellm.proxy.proxy_server.prisma_client", prisma): resp = client.get("/v1/tool/spend") assert resp.status_code == 403 - prisma.db.query_raw.assert_not_awaited() - - -def _spend_row(date: str, tool_name: str, spend: float, call_count: int = 1, total_tokens: int = 10) -> _ToolSpendRow: - return _ToolSpendRow(date=date, tool_name=tool_name, call_count=call_count, spend=spend, total_tokens=total_tokens) - - -class TestBuildToolSpendResponse: - def test_multi_tool_attribution_double_counts_per_tool_but_not_total(self): - rows = [ - _spend_row("2026-07-01", "a", spend=3.0), - _spend_row("2026-07-01", "b", spend=3.0), - ] - resp = _build_tool_spend_response(rows, total_spend=3.0, start_date="2026-07-01", end_date="2026-07-01") - by_tool = {t.tool_name: t.spend for t in resp.by_tool} - assert by_tool == {"a": 3.0, "b": 3.0} - assert resp.total_spend == 3.0 - - def test_groups_across_days_and_sorts_by_spend(self): - rows = [ - _spend_row("2026-07-01", "b", spend=1.0, call_count=2, total_tokens=100), - _spend_row("2026-07-02", "b", spend=4.0, call_count=1, total_tokens=50), - _spend_row("2026-07-01", "a", spend=2.0, call_count=3, total_tokens=300), - ] - resp = _build_tool_spend_response(rows, total_spend=7.0, start_date="2026-07-01", end_date="2026-07-02") - assert [(t.tool_name, t.spend, t.call_count, t.total_tokens) for t in resp.by_tool] == [ - ("b", 5.0, 3, 150), - ("a", 2.0, 3, 300), - ] - assert len(resp.daily) == 3 + prisma.db.litellm_dailytoolspend.group_by.assert_not_awaited() diff --git a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py index 570a25e840a..1e516cba31d 100644 --- a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py +++ b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py @@ -5913,3 +5913,39 @@ async def test_overwrite_user_with_key_hash_rejects_alias_without_marker(monkeyp ) assert updated_data["user"] == "caller-chosen-id" + + +def test_get_sanitized_user_information_from_key_drops_callback_config(): + """ + Regression (LIT-4306): `user_api_key_auth_metadata` lands in the + StandardLoggingPayload every integration receives, so the per-key callback + config (and the integration credentials inside it) must not ride along. + Everything else - notably `priority`, which the dynamic rate limiter reads + back off this exact field - has to survive. + """ + user_api_key_dict = UserAPIKeyAuth( + api_key="test-key-hash", + metadata={ + "logging": [ + { + "callback_name": "langsmith", + "callback_vars": {"langsmith_api_key": "litellm_enc::ciphertext"}, + } + ], + "callback_settings": {"callback_vars": {"langfuse_secret_key": "litellm_enc::other"}}, + "priority": "high", + }, + ) + + result = LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key( + user_api_key_dict=user_api_key_dict + ) + + auth_metadata = result["user_api_key_auth_metadata"] + assert "logging" not in auth_metadata + assert "callback_settings" not in auth_metadata + assert "litellm_enc::" not in json.dumps(auth_metadata) + assert auth_metadata["priority"] == "high" + # UserAPIKeyAuth is the live auth object; the per-key callbacks are resolved + # from it during pre-call, so it must not be mutated by building the log view + assert "logging" in (user_api_key_dict.metadata or {}) diff --git a/tests/test_litellm/proxy/test_spend_log_cleanup.py b/tests/test_litellm/proxy/test_spend_log_cleanup.py index f969b040a0d..2ba9257e1da 100644 --- a/tests/test_litellm/proxy/test_spend_log_cleanup.py +++ b/tests/test_litellm/proxy/test_spend_log_cleanup.py @@ -193,6 +193,12 @@ async def test_cleanup_old_spend_logs_batch_deletion(): tool_index_sql = mock_db.execute_raw.call_args_list[3][0][0] assert 'DELETE FROM "LiteLLM_SpendLogToolIndex"' in tool_index_sql + # The LiteLLM_DailyToolSpend rollup must outlive spend-log retention: it is + # the only copy of tool spend history once its per-request sources expire, + # so spend-log cleanup must never touch it. + for call in mock_db.execute_raw.call_args_list: + assert "LiteLLM_DailyToolSpend" not in call[0][0] + @pytest.mark.asyncio async def test_cleanup_old_spend_logs_retention_period_cutoff(): diff --git a/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py b/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py index 20451f5d0ac..d4fd5bc2dce 100644 --- a/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py +++ b/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py @@ -2680,6 +2680,142 @@ def test_update_ui_settings_writes_audit_log(monkeypatch): app.dependency_overrides.pop(user_api_key_auth, None) +@pytest.fixture +def mock_team_lookup(monkeypatch): + """Back /update/internal_user_settings with a fake team table. + + Yields the set of team ids that exist; the test mutates it before the call. + Also exposes the find_many mock so a test can assert the lookup was skipped. + """ + from unittest.mock import AsyncMock, MagicMock + + import litellm + import litellm.proxy.proxy_server as proxy_server_module + + existing_team_ids: set = set() + + async def _find_many(where): + requested = where["team_id"]["in"] + return [{"team_id": team_id} for team_id in requested if team_id in existing_team_ids] + + find_many = AsyncMock(side_effect=_find_many) + fake_prisma = MagicMock() + fake_prisma.db.litellm_teamtable.find_many = find_many + + member_budget_update = AsyncMock() + + monkeypatch.setattr(proxy_server_module, "prisma_client", fake_prisma) + monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", True) + monkeypatch.setattr(litellm, "default_internal_user_params", {}) + monkeypatch.setattr( + "litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints.update_default_team_member_budget", + member_budget_update, + ) + + return { + "existing_team_ids": existing_team_ids, + "find_many": find_many, + "member_budget_update": member_budget_update, + } + + +def test_update_internal_user_settings_rejects_unknown_team_object(mock_proxy_config, mock_auth, mock_team_lookup): + """Regression: saving a default team that doesn't exist used to return 200, + then silently fail for every SSO user because the membership write 404s.""" + mock_team_lookup["existing_team_ids"].add("real-team") + + resp = client.patch( + "/update/internal_user_settings", + json={ + "max_budget": 10.0, + "teams": [ + {"team_id": "real-team", "max_budget_in_team": 5.0}, + {"team_id": "ghost-team"}, + ], + }, + ) + + assert resp.status_code == 400, resp.text + assert "ghost-team" in resp.json()["detail"]["error"] + assert "real-team" not in resp.json()["detail"]["error"] + assert mock_proxy_config["save_call_count"]() == 0 + assert mock_team_lookup["member_budget_update"].await_count == 0, ( + "per-member budgets must not be written before the team ids are validated" + ) + + import litellm + + assert litellm.default_internal_user_params == {} + + +def test_update_internal_user_settings_rejects_unknown_team_string(mock_proxy_config, mock_auth, mock_team_lookup): + """The bare-string team shape must be validated too.""" + resp = client.patch( + "/update/internal_user_settings", + json={"teams": ["ghost-team"]}, + ) + + assert resp.status_code == 400, resp.text + assert "ghost-team" in resp.json()["detail"]["error"] + assert mock_proxy_config["save_call_count"]() == 0 + + +def test_update_internal_user_settings_rejects_duplicate_team_ids(mock_proxy_config, mock_auth, mock_team_lookup): + """Listing a team twice makes its per-member budget a race between the two + entries, so the payload is rejected rather than silently resolved.""" + mock_team_lookup["existing_team_ids"].add("real-team") + + resp = client.patch( + "/update/internal_user_settings", + json={ + "teams": [ + {"team_id": "real-team", "max_budget_in_team": 5.0}, + {"team_id": "real-team", "max_budget_in_team": 50.0}, + ] + }, + ) + + assert resp.status_code == 400, resp.text + assert "real-team" in resp.json()["detail"]["error"] + assert mock_proxy_config["save_call_count"]() == 0 + + +def test_update_internal_user_settings_saves_when_all_teams_exist(mock_proxy_config, mock_auth, mock_team_lookup): + """Valid team ids still save, and still reach the per-member budget update.""" + mock_team_lookup["existing_team_ids"].update({"team-a", "team-b"}) + + resp = client.patch( + "/update/internal_user_settings", + json={ + "max_budget": 10.0, + "teams": [ + {"team_id": "team-a", "max_budget_in_team": 5.0}, + {"team_id": "team-b"}, + ], + }, + ) + + assert resp.status_code == 200, resp.text + assert [team["team_id"] for team in resp.json()["settings"]["teams"]] == [ + "team-a", + "team-b", + ] + assert mock_proxy_config["save_call_count"]() == 1 + mock_team_lookup["member_budget_update"].assert_awaited_once() + + +def test_update_internal_user_settings_without_teams_skips_team_lookup(mock_proxy_config, mock_auth, mock_team_lookup): + """Settings changes that don't touch teams must not pay for a DB round trip.""" + resp = client.patch( + "/update/internal_user_settings", + json={"max_budget": 10.0}, + ) + + assert resp.status_code == 200, resp.text + mock_team_lookup["find_many"].assert_not_awaited() + assert mock_proxy_config["save_call_count"]() == 1 + + def test_update_mcp_semantic_filter_settings_requires_proxy_admin(monkeypatch): """Non-admin callers must not mutate global MCP semantic filter settings.""" from litellm.proxy._types import UserAPIKeyAuth diff --git a/tests/test_litellm/proxy/utils/prisma_and_spend/conftest.py b/tests/test_litellm/proxy/utils/prisma_and_spend/conftest.py index af62b7eef62..74c9abd9978 100644 --- a/tests/test_litellm/proxy/utils/prisma_and_spend/conftest.py +++ b/tests/test_litellm/proxy/utils/prisma_and_spend/conftest.py @@ -128,6 +128,8 @@ def mock_prisma_client() -> MagicMock: client.proxy_logging_obj.failure_handler = AsyncMock() client.spend_log_transactions = [] client._spend_log_transactions_lock = asyncio.Lock() + client.tool_usage_transactions = [] + client._tool_usage_transactions_lock = asyncio.Lock() client.jsonify_object = lambda data: dict(data) client.db.is_connected = MagicMock(return_value=False) client.db.connect = AsyncMock() diff --git a/tests/test_litellm/proxy/utils/prisma_and_spend/test_proxy_update_spend.py b/tests/test_litellm/proxy/utils/prisma_and_spend/test_proxy_update_spend.py index d5d4de7f2cf..f075acc7307 100644 --- a/tests/test_litellm/proxy/utils/prisma_and_spend/test_proxy_update_spend.py +++ b/tests/test_litellm/proxy/utils/prisma_and_spend/test_proxy_update_spend.py @@ -68,12 +68,12 @@ async def test_update_end_user_spend_upserts_each_end_user( @pytest.mark.asyncio -async def test_update_end_user_spend_retries_on_connection_error( +async def test_update_end_user_spend_retries_on_connect_error( mock_prisma_client: Any, monkeypatch: pytest.MonkeyPatch ) -> None: - """``DB_CONNECTION_ERROR_TYPES`` failures should be retried with backoff; - once retries are exhausted, ``_raise_failed_update_spend_exception`` is - invoked and the original exception bubbles up. + """``DB_RETRY_SAFE_ERROR_TYPES`` (ConnectError, statements provably never + sent) retries with backoff; once retries are exhausted the original + exception bubbles up via ``_raise_failed_update_spend_exception``. """ import httpx import litellm.proxy.utils as utils_mod @@ -85,11 +85,11 @@ async def test_update_end_user_spend_retries_on_connection_error( monkeypatch.setattr(utils_mod.asyncio, "sleep", _fake_sleep) - err = httpx.ReadError("conn reset") + err = httpx.ConnectError("down") mock_prisma_client.db.tx = MagicMock(side_effect=err) proxy_logging = MagicMock() proxy_logging.failure_handler = AsyncMock() - with pytest.raises(httpx.ReadError): + with pytest.raises(httpx.ConnectError): await ProxyUpdateSpend.update_end_user_spend( n_retry_times=1, prisma_client=mock_prisma_client, @@ -99,6 +99,29 @@ async def test_update_end_user_spend_retries_on_connection_error( assert sleeps == [1.0] +@pytest.mark.asyncio +@pytest.mark.parametrize("ambiguous_error_name", ["ReadTimeout", "ReadError"]) +async def test_update_end_user_spend_does_not_retry_post_send_ambiguous_errors( + mock_prisma_client: Any, ambiguous_error_name: str +) -> None: + """Post-send errors are ambiguous and retrying can double-apply increments + (see DB_RETRY_SAFE_ERROR_TYPES); they must raise on the first attempt.""" + import httpx + + err = getattr(httpx, ambiguous_error_name)("ambiguous") + mock_prisma_client.db.tx = MagicMock(side_effect=err) + proxy_logging = MagicMock() + proxy_logging.failure_handler = AsyncMock() + with pytest.raises((httpx.ReadTimeout, httpx.ReadError)): + await ProxyUpdateSpend.update_end_user_spend( + n_retry_times=3, + prisma_client=mock_prisma_client, + proxy_logging_obj=proxy_logging, + end_user_list_transactions={"u": 1.0}, + ) + mock_prisma_client.db.tx.assert_called_once() + + @pytest.mark.asyncio async def test_update_end_user_spend_non_connection_error_raises_immediately( mock_prisma_client: Any, diff --git a/tests/test_litellm/proxy/utils/prisma_and_spend/test_spend_functions.py b/tests/test_litellm/proxy/utils/prisma_and_spend/test_spend_functions.py index a0b3af54750..d9eeb168611 100644 --- a/tests/test_litellm/proxy/utils/prisma_and_spend/test_spend_functions.py +++ b/tests/test_litellm/proxy/utils/prisma_and_spend/test_spend_functions.py @@ -188,6 +188,35 @@ async def test_update_spend_logs_job_skips_when_queue_empty( assert mock_prisma_client.db.litellm_spendlogs.create_many.await_count == 0 +@pytest.mark.asyncio +async def test_update_spend_logs_job_drains_tool_queue_when_spend_queue_empty( + mock_prisma_client: Any, monkeypatch: pytest.MonkeyPatch +) -> None: + # Regression: a spend-log write failure aborts a run before the tool drain, + # so tool transactions can outlive the spend queue; the job must still run + # for them instead of early-returning on the empty spend queue. + import litellm.proxy.db.spend_log_tool_index as tool_mod + import litellm.proxy.guardrails.usage_tracking as guard_mod + + proxy_logging = MagicMock() + proxy_logging.failure_handler = AsyncMock() + mock_prisma_client.spend_log_transactions = [] + mock_prisma_client.tool_usage_transactions = [MagicMock()] + mock_prisma_client.db.litellm_spendlogs.create_many = AsyncMock() + monkeypatch.setattr(guard_mod, "process_spend_logs_guardrail_usage", AsyncMock(), raising=False) + flush_stub = AsyncMock() + monkeypatch.setattr(tool_mod, "flush_tool_usage_transactions", flush_stub, raising=False) + + await update_spend_logs_job( + prisma_client=mock_prisma_client, + db_writer_client=None, + proxy_logging_obj=proxy_logging, + ) + + assert len(flush_stub.await_args.kwargs["transactions"]) == 1 + assert mock_prisma_client.tool_usage_transactions == [] + + @pytest.mark.asyncio async def test_update_spend_logs_job_processes_and_clears_queue( mock_prisma_client: Any, make_spend_log_row: Any, monkeypatch: pytest.MonkeyPatch @@ -208,7 +237,7 @@ async def test_update_spend_logs_job_processes_and_clears_queue( guard_mod, "process_spend_logs_guardrail_usage", AsyncMock(), raising=False ) monkeypatch.setattr( - tool_mod, "process_spend_logs_tool_usage", AsyncMock(), raising=False + tool_mod, "flush_tool_usage_transactions", AsyncMock(), raising=False ) await update_spend_logs_job( diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index a9e5b3316e0..ad4e430c603 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -5936,3 +5936,446 @@ async def test_acreate_batch_request_bedrock_tags_override_deployment_tags(): bedrock_tags=request_tags, ) assert mock_sign.call_args.kwargs["data"]["tags"] == request_tags + + +class TestPreRoutingStrategyRegistryLifecycle: + """ + Regression tests: a deployment leaving the model_list must release the + pre-routing strategy slot it holds in `auto_routers` / `complexity_routers` / + `adaptive_routers` / `quality_routers`. + + Before this fix, editing an auto-router-family model (a UI save, which reaches + every other pod as an `upsert_deployment` from the periodic DB reload) popped + the deployment out of the model_list and then failed to re-add it: registration + raised "already exists" against the stale registry entry, and + `ignore_invalid_deployments=True` swallowed the error. The router vanished from + the Models page and stayed gone until a proxy restart, while the DB row and the + "saved successfully" response both looked fine. + """ + + @staticmethod + def _complexity_router_params(default_model: str, tags=None) -> dict: + return { + "model": "auto_router/complexity_router", + "complexity_router_config": { + "tiers": {"SIMPLE": "gpt-4o-mini", "MEDIUM": "gpt-4o", "COMPLEX": "gpt-4o"} + }, + "complexity_router_default_model": default_model, + **({"tags": tags} if tags else {}), + } + + @classmethod + def _router_with_complexity_router(cls, default_model: str = "gpt-4o") -> "litellm.Router": + return litellm.Router( + model_list=[ + {"model_name": "gpt-4o", "litellm_params": {"model": "gpt-4o"}}, + {"model_name": "gpt-4o-mini", "litellm_params": {"model": "gpt-4o-mini"}}, + { + "model_name": "smart-router", + "litellm_params": cls._complexity_router_params(default_model), + "model_info": {"id": "router-1", "db_model": True}, + }, + ], + ignore_invalid_deployments=True, + ) + + @staticmethod + def _model_names(router: "litellm.Router") -> list: + return [model["model_name"] for model in router.model_list] + + def test_upsert_of_edited_router_keeps_it_routable(self): + from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo + + router = self._router_with_complexity_router() + + router.upsert_deployment( + deployment=Deployment( + model_name="smart-router", + litellm_params=LiteLLM_Params(**self._complexity_router_params("gpt-4o-mini")), + model_info=ModelInfo(id="router-1", db_model=True), + ) + ) + + assert "smart-router" in self._model_names(router) + registered = router.complexity_routers["smart-router"] + assert len(registered) == 1 + # the surviving strategy is the edited one, not the pre-edit leftover + assert registered[0].strategy.config.default_model == "gpt-4o-mini" + + def test_unchanged_upsert_leaves_router_untouched(self): + from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo + + router = self._router_with_complexity_router() + strategy_before = router.complexity_routers["smart-router"][0].strategy + + for _ in range(3): + router.upsert_deployment( + deployment=Deployment( + model_name="smart-router", + litellm_params=LiteLLM_Params(**self._complexity_router_params("gpt-4o")), + model_info=ModelInfo(id="router-1", db_model=True), + ) + ) + + assert "smart-router" in self._model_names(router) + assert router.complexity_routers["smart-router"][0].strategy is strategy_before + + def test_delete_frees_the_name_for_a_new_router(self): + from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo + + router = self._router_with_complexity_router() + + router.delete_deployment(id="router-1") + assert "smart-router" not in router.complexity_routers + + router.add_deployment( + deployment=Deployment( + model_name="smart-router", + litellm_params=LiteLLM_Params(**self._complexity_router_params("gpt-4o-mini")), + model_info=ModelInfo(id="router-2", db_model=True), + ) + ) + + assert "smart-router" in self._model_names(router) + assert router.complexity_routers["smart-router"][0].strategy.config.default_model == "gpt-4o-mini" + + def test_delete_only_frees_the_matching_tag_slot(self): + router = litellm.Router( + model_list=[ + {"model_name": "gpt-4o", "litellm_params": {"model": "gpt-4o"}}, + {"model_name": "gpt-4o-mini", "litellm_params": {"model": "gpt-4o-mini"}}, + { + "model_name": "shared-router", + "litellm_params": self._complexity_router_params("gpt-4o", tags=["team-a"]), + "model_info": {"id": "router-a"}, + }, + { + "model_name": "shared-router", + "litellm_params": self._complexity_router_params("gpt-4o-mini", tags=["team-b"]), + "model_info": {"id": "router-b"}, + }, + ], + ignore_invalid_deployments=True, + ) + assert len(router.complexity_routers["shared-router"]) == 2 + + router.delete_deployment(id="router-a") + + remaining = router.complexity_routers["shared-router"] + assert len(remaining) == 1 + assert remaining[0].tags == ("team-b",) + + def test_delete_of_regular_model_preserves_router_sharing_its_name(self): + router = litellm.Router( + model_list=[ + {"model_name": "gpt-4o", "litellm_params": {"model": "gpt-4o"}}, + {"model_name": "gpt-4o-mini", "litellm_params": {"model": "gpt-4o-mini"}}, + { + "model_name": "shared-name", + "litellm_params": self._complexity_router_params("gpt-4o"), + "model_info": {"id": "router-1"}, + }, + { + "model_name": "shared-name", + "litellm_params": {"model": "openai/gpt-4o"}, + "model_info": {"id": "regular-1"}, + }, + ], + ignore_invalid_deployments=True, + ) + strategy = router.complexity_routers["shared-name"][0].strategy + + router.delete_deployment(id="regular-1") + + assert router.complexity_routers["shared-name"][0].strategy is strategy + + def test_upsert_of_edited_adaptive_router_rebuilds_it(self): + """Adaptive routers are built by set_model_list()'s deferred pass, not by + add_deployment(), so releasing the slot on edit must be paired with a rebuild - + otherwise the edit silently turns adaptive routing off.""" + from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo + + def adaptive_params(available_models: list) -> dict: + return { + "model": "auto_router/adaptive_router", + "adaptive_router_config": {"available_models": available_models}, + } + + router = litellm.Router( + model_list=[ + {"model_name": "gpt-4o", "litellm_params": {"model": "openai/gpt-4o"}}, + {"model_name": "gpt-4o-mini", "litellm_params": {"model": "openai/gpt-4o-mini"}}, + { + "model_name": "adaptive-router", + "litellm_params": adaptive_params(["gpt-4o-mini"]), + "model_info": {"id": "router-1", "db_model": True}, + }, + ], + ignore_invalid_deployments=True, + ) + assert "adaptive-router" in router.adaptive_routers + + router.upsert_deployment( + deployment=Deployment( + model_name="adaptive-router", + litellm_params=LiteLLM_Params(**adaptive_params(["gpt-4o", "gpt-4o-mini"])), + model_info=ModelInfo(id="router-1", db_model=True), + ) + ) + + assert "adaptive-router" in self._model_names(router) + registered = router.adaptive_routers["adaptive-router"] + assert len(registered) == 1 + assert set(registered[0].strategy.config.available_models) == {"gpt-4o", "gpt-4o-mini"} + + def test_delete_repairs_indices_even_when_strategy_release_fails(self): + """Structural removal and strategy release are not equally critical. Once the entry + leaves model_list the index maps must be repaired no matter what, so releasing the + registry slot runs after that repair and cannot abandon the router half-updated.""" + router = self._router_with_complexity_router() + idx = router.model_id_to_deployment_index_map["router-1"] + router.model_list[idx] = {"model_name": "smart-router", "litellm_params": None} + + returned = router.delete_deployment(id="router-1") + + assert returned is not None + assert "router-1" not in router.model_id_to_deployment_index_map + assert all(entry.get("model_info", {}).get("id") != "router-1" for entry in router.model_list) + assert router.get_deployment(model_id="router-1") is None + assert "gpt-4o" in self._model_names(router) + + def test_delete_of_adaptive_enabled_complexity_router_frees_both_registries(self): + """A complexity router with adaptive set is registered in BOTH complexity_routers + and adaptive_routers under the same (model_name, tags). Releasing only the first + match leaves the adaptive strategy live, so a deleted alias stays routable and its + post-call hook keeps recording.""" + import litellm as litellm_module + from litellm.router_strategy.adaptive_router.hooks import AdaptiveRouterPostCallHook + + params = { + "model": "auto_router/complexity_router", + "complexity_router_config": { + "tiers": {"SIMPLE": "gpt-4o-mini", "MEDIUM": "gpt-4o"}, + "adaptive": True, + }, + "complexity_router_default_model": "gpt-4o", + } + router = litellm.Router( + model_list=[ + {"model_name": "gpt-4o", "litellm_params": {"model": "openai/gpt-4o"}}, + {"model_name": "gpt-4o-mini", "litellm_params": {"model": "openai/gpt-4o-mini"}}, + { + "model_name": "hybrid-router", + "litellm_params": params, + "model_info": {"id": "router-1", "db_model": True}, + }, + ], + ignore_invalid_deployments=True, + ) + assert "hybrid-router" in router.complexity_routers + assert "hybrid-router" in router.adaptive_routers + hooks = litellm_module.logging_callback_manager.get_custom_loggers_for_type(AdaptiveRouterPostCallHook) + assert len(hooks) == 1 + + router.delete_deployment(id="router-1") + + assert "hybrid-router" not in router.complexity_routers + assert "hybrid-router" not in router.adaptive_routers + remaining_hooks = litellm_module.logging_callback_manager.get_custom_loggers_for_type( + AdaptiveRouterPostCallHook + ) + assert remaining_hooks == [] + + def test_upsert_of_edited_quality_router_keeps_it_routable(self): + """_unregister_pre_routing_strategy_for_deployment dispatches on four prefixes; + quality_router is one of them and would otherwise go unexercised.""" + from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo + + def quality_params(default_model: str) -> dict: + return { + "model": "auto_router/quality_router", + "quality_router_default_model": default_model, + } + + router = litellm.Router( + model_list=[ + {"model_name": "gpt-4o", "litellm_params": {"model": "openai/gpt-4o"}}, + {"model_name": "gpt-4o-mini", "litellm_params": {"model": "openai/gpt-4o-mini"}}, + { + "model_name": "quality-router", + "litellm_params": quality_params("gpt-4o"), + "model_info": {"id": "router-1", "db_model": True}, + }, + ], + ignore_invalid_deployments=True, + ) + assert "quality-router" in router.quality_routers + + router.upsert_deployment( + deployment=Deployment( + model_name="quality-router", + litellm_params=LiteLLM_Params(**quality_params("gpt-4o-mini")), + model_info=ModelInfo(id="router-1", db_model=True), + ) + ) + + assert "quality-router" in self._model_names(router) + registered = router.quality_routers["quality-router"] + assert len(registered) == 1 + assert registered[0].strategy.config.default_model == "gpt-4o-mini" + + @staticmethod + def _hybrid_router_params(tiers: dict) -> dict: + return { + "model": "auto_router/complexity_router", + "complexity_router_config": {"tiers": tiers, "adaptive": True}, + "complexity_router_default_model": "gpt-4o", + } + + @classmethod + def _router_with_hybrid_router(cls) -> "litellm.Router": + return litellm.Router( + model_list=[ + {"model_name": "gpt-4o", "litellm_params": {"model": "openai/gpt-4o"}}, + {"model_name": "gpt-4o-mini", "litellm_params": {"model": "openai/gpt-4o-mini"}}, + { + "model_name": "hybrid-router", + "litellm_params": cls._hybrid_router_params({"SIMPLE": "gpt-4o-mini", "MEDIUM": "gpt-4o"}), + "model_info": {"id": "router-1", "db_model": True}, + }, + ], + ignore_invalid_deployments=True, + ) + + def test_upsert_of_edited_hybrid_complexity_router_relinks_adaptive(self): + """Editing an adaptive-enabled complexity router releases its adaptive companion + along with the complexity slot; the finalize re-run must fire for it (not just for + `auto_router/adaptive_router` deployments) or the rebuilt complexity router keeps + routing while bandit recording, DB persistence and /adaptive_router/state all + silently stop until the next full reload.""" + import litellm as litellm_module + from litellm.router_strategy.adaptive_router.hooks import AdaptiveRouterPostCallHook + from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo + + router = self._router_with_hybrid_router() + assert "hybrid-router" in router.adaptive_routers + + router.upsert_deployment( + deployment=Deployment( + model_name="hybrid-router", + litellm_params=LiteLLM_Params( + **self._hybrid_router_params( + {"SIMPLE": "gpt-4o-mini", "MEDIUM": "gpt-4o", "COMPLEX": "gpt-4o"} + ) + ), + model_info=ModelInfo(id="router-1", db_model=True), + ) + ) + + assert "hybrid-router" in self._model_names(router) + assert "hybrid-router" in router.complexity_routers + assert "hybrid-router" in router.adaptive_routers + rebuilt = router.complexity_routers["hybrid-router"][0].strategy + assert router.adaptive_routers["hybrid-router"][0].strategy is rebuilt.adaptive_router + hooks = litellm_module.logging_callback_manager.get_custom_loggers_for_type(AdaptiveRouterPostCallHook) + assert len(hooks) == 1 + + def test_upsert_turning_adaptive_on_builds_the_companion(self): + """An edit that flips `adaptive: true` on an existing complexity router must + register the companion immediately; neither side of the old prefix-only gate + matches a complexity deployment, so the flip was a silent no-op until restart.""" + from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo + + router = self._router_with_complexity_router() + assert "smart-router" not in router.adaptive_routers + + params = self._complexity_router_params("gpt-4o") + params["complexity_router_config"] = {**params["complexity_router_config"], "adaptive": True} + router.upsert_deployment( + deployment=Deployment( + model_name="smart-router", + litellm_params=LiteLLM_Params(**params), + model_info=ModelInfo(id="router-1", db_model=True), + ) + ) + + assert "smart-router" in router.adaptive_routers + + def test_unregister_pre_routing_strategy_scopes_the_drop_by_tags(self): + """The bool return drives the hook re-sync; a tag mismatch must report False and + leave the registry untouched, and dropping the last entry must free the key.""" + from litellm.types.router import TaggedPreRoutingStrategy + + registry = { + "m": [ + TaggedPreRoutingStrategy(tags=("team-a",), strategy=object()), + TaggedPreRoutingStrategy(tags=(), strategy=object()), + ] + } + + assert litellm.Router._unregister_pre_routing_strategy(registry, "m", ("team-b",)) is False + assert len(registry["m"]) == 2 + + assert litellm.Router._unregister_pre_routing_strategy(registry, "m", ("team-a",)) is True + assert [entry.tags for entry in registry["m"]] == [()] + + assert litellm.Router._unregister_pre_routing_strategy(registry, "m", ()) is True + assert "m" not in registry + + def test_unregister_for_deployment_ignores_non_router_deployments(self): + """Direct twin of the endpoint-level test: a regular deployment that shares a + router's model_name must not evict the router's registry slot.""" + from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo + + router = self._router_with_complexity_router() + + router._unregister_pre_routing_strategy_for_deployment( + deployment=Deployment( + model_name="smart-router", + litellm_params=LiteLLM_Params(model="openai/gpt-4o"), + model_info=ModelInfo(id="plain-1", db_model=True), + ) + ) + + assert "smart-router" in router.complexity_routers + + def test_sync_adaptive_router_hooks_keeps_one_hook_per_registered_router(self): + """Re-syncing must replace, not accumulate: a duplicated hook double-fires + bandit signal recording for every request.""" + import litellm as litellm_module + from litellm.router_strategy.adaptive_router.hooks import AdaptiveRouterPostCallHook + + router = self._router_with_hybrid_router() + + router._sync_adaptive_router_hooks() + router._sync_adaptive_router_hooks() + + hooks = litellm_module.logging_callback_manager.get_custom_loggers_for_type(AdaptiveRouterPostCallHook) + assert len(hooks) == 1 + + def test_deployment_participates_in_adaptive_routing_matrix(self): + """The upsert finalize re-run keys off this predicate for both the incoming and + outgoing deployment; a false negative silently strands the adaptive companion.""" + from litellm.types.router import LiteLLM_Params + + router = self._router_with_complexity_router() + + cases = [ + ({"model": "auto_router/adaptive_router", "adaptive_router_config": {}}, True), + (self._hybrid_router_params({"SIMPLE": "gpt-4o-mini"}), True), + (self._complexity_router_params("gpt-4o"), False), + ( + { + "model": "auto_router/complexity_router", + "complexity_router_config": {"tiers": {"SIMPLE": "gpt-4o-mini"}, "adaptive": False}, + "complexity_router_default_model": "gpt-4o", + }, + False, + ), + ({"model": "openai/gpt-4o"}, False), + ] + for params, expected in cases: + actual = router._deployment_participates_in_adaptive_routing( + litellm_params=LiteLLM_Params(**params) + ) + assert actual is expected, params["model"] diff --git a/type-discipline-budget.json b/type-discipline-budget.json index 410bb8d9250..d56d5a6e305 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,9 +1,9 @@ { "LIT001": { - "limit": 23408 + "limit": 23287 }, "LIT002": { - "limit": 27511 + "limit": 27473 }, "LIT003": { "limit": 292 @@ -15,7 +15,7 @@ "limit": 0 }, "LIT006": { - "limit": 1111 + "limit": 1109 }, "LIT007": { "limit": 0 @@ -24,6 +24,6 @@ "limit": 1004 }, "LIT009": { - "limit": 2501 + "limit": 2495 } } diff --git a/ui/litellm-dashboard/eslint-suppressions.json b/ui/litellm-dashboard/eslint-suppressions.json index 4a89c7d4387..66e2055f579 100644 --- a/ui/litellm-dashboard/eslint-suppressions.json +++ b/ui/litellm-dashboard/eslint-suppressions.json @@ -225,11 +225,6 @@ "count": 1 } }, - "src/app/(dashboard)/cost-optimization/_components/UsageTab.tsx": { - "no-restricted-imports": { - "count": 1 - } - }, "src/app/(dashboard)/cost-tracking/_components/add_margin_form.tsx": { "local/filename-pascal-case": { "count": 1 @@ -1867,11 +1862,6 @@ "count": 1 } }, - "src/app/(dashboard)/users/_components/DefaultUserSettings.tsx": { - "no-restricted-imports": { - "count": 2 - } - }, "src/app/(dashboard)/users/_components/edit_user.tsx": { "local/filename-pascal-case": { "count": 1 @@ -3368,10 +3358,10 @@ "count": 5 }, "no-restricted-syntax": { - "count": 154 + "count": 153 }, "prefer-const": { - "count": 33 + "count": 32 } }, "src/components/object_permissions_view.tsx": { @@ -4339,11 +4329,6 @@ "count": 1 } }, - "src/lib/http/client.ts": { - "no-nested-ternary": { - "count": 1 - } - }, "src/utils/dataUtils.test.ts": { "max-nested-callbacks": { "count": 1 @@ -4365,4 +4350,4 @@ "count": 1 } } -} +} \ No newline at end of file diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CostOptimizationView.activity.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CostOptimizationView.activity.test.tsx index 27261768b8d..363525c48af 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CostOptimizationView.activity.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CostOptimizationView.activity.test.tsx @@ -5,7 +5,7 @@ const mockUserDailyActivityCall = vi.fn(); vi.mock("@/components/networking", () => ({ userDailyActivityCall: (...args: unknown[]) => mockUserDailyActivityCall(...args), - getToolSpend: vi.fn().mockResolvedValue({ by_tool: [], daily: [], total_spend: 0, start_date: null, end_date: null }), + getToolSpend: vi.fn().mockResolvedValue({ by_tool: [], daily: [], start_date: null, end_date: null }), getGeneralSettingsCall: vi.fn().mockResolvedValue([]), })); @@ -19,7 +19,7 @@ vi.mock("@/components/shared/charts", () => ({ DonutChart: () =>
, BarChart: () =>
, CustomLegend: () =>
, - DEFAULT_COLOR_CYCLE: ["emerald"], + SEQUENTIAL_COLOR_RAMP: ["indigo"], })); vi.mock("@/app/(dashboard)/router-settings/_components/general_settings", () => ({ diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/UsageTab.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/UsageTab.test.tsx index f84167f5a82..5c26ac30477 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/UsageTab.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/UsageTab.test.tsx @@ -23,18 +23,37 @@ vi.mock("@/components/shared/charts", () => ({ DonutChart: ({ data, label }: { data: unknown; label: string }) => (
), - BarChart: ({ data, categories }: { data: unknown; categories: string[] }) => ( -
+ BarChart: ({ + data, + categories, + colors, + showLegend, + maxBarSize, + }: { + data: unknown; + categories: string[]; + colors?: readonly string[]; + showLegend?: boolean; + maxBarSize?: number; + }) => ( +
), CustomLegend: ({ categories }: { categories: readonly string[] }) => (
{categories.join(",")}
), - DEFAULT_COLOR_CYCLE: ["emerald", "blue", "violet", "amber"], + SEQUENTIAL_COLOR_RAMP: ["indigo", "blue", "sky", "cyan"], })); import UsageTab from "./UsageTab"; -const emptyToolSpend: ToolSpendResponse = { by_tool: [], daily: [], total_spend: 0, start_date: null, end_date: null }; +const emptyToolSpend: ToolSpendResponse = { by_tool: [], daily: [], start_date: null, end_date: null }; const baseMetrics = (overrides: Partial): SpendMetrics => ({ spend: 0, @@ -216,7 +235,6 @@ describe("UsageTab", () => { { tool_name: "read_file", spend: 1.0, call_count: 2, total_tokens: 50 }, ], daily: [{ date: "2026-07-12", tool_name: "search", spend: 4.0, call_count: 3 }], - total_spend: 5.0, start_date: "2026-07-12", end_date: "2026-07-12", }; @@ -225,32 +243,29 @@ describe("UsageTab", () => { const bars = await findAllByTestId("bar-chart"); const series = JSON.parse(bars[0].getAttribute("data-series") ?? "[]"); expect(series[0]).toMatchObject({ tool_name: "search", spend: 4.0 }); + // The 64px bar cap is this card's opt-in; the shared BarChart must not cap + // by default (other consumers keep their pre-existing geometry). + expect(bars[0].getAttribute("data-max-bar-size")).toBe("64"); }); - it("notes the 30-day cap when the server clamps the tool spend window", async () => { + it("renders the tool legend once outside the charts, with both charts sharing the tool colors", async () => { const toolSpend = { - by_tool: [{ tool_name: "search", spend: 4.0, call_count: 3, total_tokens: 150 }], + by_tool: [ + { tool_name: "search", spend: 4.0, call_count: 3, total_tokens: 150 }, + { tool_name: "read_file", spend: 1.0, call_count: 2, total_tokens: 50 }, + ], daily: [{ date: "2026-07-12", tool_name: "search", spend: 4.0, call_count: 3 }], - total_spend: 4.0, - start_date: "2026-07-05", - end_date: "2026-07-14", + start_date: "2026-07-12", + end_date: "2026-07-12", }; - const { findByText } = renderWith([day("2026-07-12", {})], { toolSpend }); + const { findAllByTestId, getAllByTestId } = renderWith([day("2026-07-12", {})], { toolSpend }); - expect(await findByText(/capped at 30 days before the end of the selected range/)).toBeInTheDocument(); - }); + const bars = await findAllByTestId("bar-chart"); + const [totalByTool, dailyByTool] = bars.slice(-2); + expect(dailyByTool.getAttribute("data-show-legend")).toBe("false"); + expect(totalByTool.getAttribute("data-colors")).toBe(dailyByTool.getAttribute("data-colors")); - it("shows no cap note when the served window matches the request", async () => { - const toolSpend = { - by_tool: [{ tool_name: "search", spend: 4.0, call_count: 3, total_tokens: 150 }], - daily: [{ date: "2026-07-12", tool_name: "search", spend: 4.0, call_count: 3 }], - total_spend: 4.0, - start_date: "2026-07-01", - end_date: "2026-07-14", - }; - const { findAllByTestId, queryByText } = renderWith([day("2026-07-12", {})], { toolSpend }); - - await findAllByTestId("bar-chart"); - expect(queryByText(/capped at 30 days before the end of the selected range/)).not.toBeInTheDocument(); + const toolLegends = getAllByTestId("chart-legend").filter((legend) => legend.textContent === "search,read_file"); + expect(toolLegends).toHaveLength(1); }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/UsageTab.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/UsageTab.tsx index 15ce84b8445..ec37418e0b5 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/UsageTab.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/UsageTab.tsx @@ -3,7 +3,7 @@ import React, { useEffect, useMemo, useState } from "react"; import { Info } from "lucide-react"; -import { AreaChart, BarChart, CustomLegend, DonutChart, DEFAULT_COLOR_CYCLE } from "@/components/shared/charts"; +import { AreaChart, BarChart, CustomLegend, DonutChart, SEQUENTIAL_COLOR_RAMP } from "@/components/shared/charts"; import AdvancedDatePicker from "@/components/shared/advanced_date_picker"; import { Card, CardContent, CardHeader, CardTitle } from "@/components/ui/card"; import { Popover, PopoverContent, PopoverTrigger } from "@/components/ui/popover"; @@ -34,7 +34,6 @@ interface UsageTabProps { const EMPTY_TOOL_SPEND: ToolSpendResponse = { by_tool: [], daily: [], - total_spend: 0, start_date: null, end_date: null, }; @@ -103,7 +102,6 @@ const UsageTab: React.FC = ({ accessToken, activity }) => { const toolSpend = toolSpendState?.key === rangeKey ? toolSpendState.data : null; const toolSpendLoading = toolSpendEnabled && toolSpend === null; - const toolSpendWindowClamped = !!toolSpend?.start_date && !!startTime && toolSpend.start_date > isoDay(startTime); const compressionTotal = useMemo(() => results.reduce((sum, d) => sum + compressionOf(d.metrics), 0), [results]); const cachingTotal = useMemo(() => results.reduce((sum, d) => sum + cachingOf(d.metrics), 0), [results]); @@ -168,7 +166,7 @@ const UsageTab: React.FC = ({ accessToken, activity }) => { })), [toolSpend, topToolNames], ); - const toolColors = useMemo(() => DEFAULT_COLOR_CYCLE.slice(0, Math.max(topToolNames.length, 1)), [topToolNames]); + const toolColors = useMemo(() => SEQUENTIAL_COLOR_RAMP.slice(0, Math.max(topToolNames.length, 1)), [topToolNames]); return (
@@ -262,15 +260,10 @@ const UsageTab: React.FC = ({ accessToken, activity }) => { Spend by tool

- Spend on requests that called each tool (MCP and client-side tools). A request that used multiple tools - counts its full spend toward each, so this attributes rather than partitions spend. + Spend on requests that invoked each tool (MCP and client-side tools); declaring a tool without invoking it + does not count. A request that invoked multiple tools counts its full spend toward each, so this attributes + rather than partitions spend.

- {toolSpendWindowClamped && ( -

- Tool spend is capped at 30 days before the end of the selected range; showing spend since{" "} - {toolSpend?.start_date}. -

- )} {topTools.length === 0 ? ( @@ -285,22 +278,27 @@ const UsageTab: React.FC = ({ accessToken, activity }) => { data={topToolsChart} index="tool_name" categories={["spend"]} - colors={["emerald"]} + colors={toolColors} + colorByDatum layout="vertical" yAxisWidth={140} + maxBarSize={64} showLegend={false} valueFormatter={usd} />

Daily spend by tool

+
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/DefaultUserSettings.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/DefaultUserSettings.test.tsx deleted file mode 100644 index 06dafcfcffd..00000000000 --- a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/DefaultUserSettings.test.tsx +++ /dev/null @@ -1,153 +0,0 @@ -import { act, fireEvent, render, screen, waitFor } from "@testing-library/react"; -import { beforeEach, describe, expect, it, vi } from "vitest"; -import DefaultUserSettings from "./DefaultUserSettings"; -import * as networking from "@/components/networking"; - -vi.mock("@/components/networking", () => ({ - getInternalUserSettings: vi.fn(), - updateInternalUserSettings: vi.fn(), - modelAvailableCall: vi.fn(), -})); - -vi.mock("@/components/common_components/budget_duration_dropdown", () => ({ - default: ({ value, onChange }: { value: string | null; onChange: (value: string | null) => void }) => ( - - ), - getBudgetDurationLabel: (value: string) => value, -})); - -vi.mock("@/components/key_team_helpers/fetch_available_models_team_key", () => ({ - getModelDisplayName: (model: string) => model, -})); - -describe("DefaultUserSettings", () => { - const mockGetInternalUserSettings = vi.mocked(networking.getInternalUserSettings); - const mockUpdateInternalUserSettings = vi.mocked(networking.updateInternalUserSettings); - const mockModelAvailableCall = vi.mocked(networking.modelAvailableCall); - - const defaultProps = { - accessToken: "test-token", - userID: "user-123", - userRole: "Admin", - possibleUIRoles: { - internal_user_admin: { - ui_label: "Admin", - description: "Full access", - }, - internal_user_viewer: { - ui_label: "Viewer", - description: "Read-only access", - }, - }, - }; - - const mockSettings = { - values: { - user_role: "internal_user_admin", - budget_duration: "monthly", - max_budget: 1000, - teams: [], - }, - field_schema: { - description: "Default user settings", - properties: { - user_role: { - type: "string", - description: "User role", - }, - budget_duration: { - type: "string", - description: "Budget duration", - }, - max_budget: { - type: "number", - description: "Maximum budget", - }, - teams: { - type: "array", - description: "Teams", - }, - }, - }, - }; - - beforeEach(() => { - mockGetInternalUserSettings.mockClear(); - mockUpdateInternalUserSettings.mockClear(); - mockModelAvailableCall.mockClear(); - mockModelAvailableCall.mockResolvedValue({ - data: [{ id: "gpt-4" }, { id: "gpt-3.5-turbo" }], - }); - }); - - it("should render", async () => { - mockGetInternalUserSettings.mockResolvedValue(mockSettings); - - render(); - - await waitFor(() => { - expect(mockGetInternalUserSettings).toHaveBeenCalled(); - }); - - expect(screen.getByText("Default User Settings")).toBeInTheDocument(); - }); - - it("should toggle edit mode when edit button is clicked", async () => { - mockGetInternalUserSettings.mockResolvedValue(mockSettings); - - render(); - - await waitFor(() => { - expect(screen.getByText("Edit Settings")).toBeInTheDocument(); - }); - - const editButton = screen.getByText("Edit Settings"); - act(() => { - fireEvent.click(editButton); - }); - - expect(screen.getByText("Cancel")).toBeInTheDocument(); - expect(screen.getByText("Save Changes")).toBeInTheDocument(); - expect(screen.queryByText("Edit Settings")).not.toBeInTheDocument(); - }); - - it("should save settings when save button is clicked", async () => { - mockGetInternalUserSettings.mockResolvedValue(mockSettings); - mockUpdateInternalUserSettings.mockResolvedValue({ - settings: { - ...mockSettings.values, - max_budget: 2000, - }, - }); - - render(); - - await waitFor(() => { - expect(screen.getByText("Edit Settings")).toBeInTheDocument(); - }); - - const editButton = screen.getByText("Edit Settings"); - act(() => { - fireEvent.click(editButton); - }); - - await waitFor(() => { - expect(screen.getByText("Save Changes")).toBeInTheDocument(); - }); - - const saveButton = screen.getByText("Save Changes"); - act(() => { - fireEvent.click(saveButton); - }); - - await waitFor(() => { - expect(mockUpdateInternalUserSettings).toHaveBeenCalled(); - }); - - expect(screen.getByText("Edit Settings")).toBeInTheDocument(); - }); -}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/DefaultUserSettings.tsx b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/DefaultUserSettings.tsx deleted file mode 100644 index 7fee2e14b27..00000000000 --- a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/DefaultUserSettings.tsx +++ /dev/null @@ -1,492 +0,0 @@ -import React, { useState, useEffect } from "react"; -import { Card, Title, Text, Divider, TextInput } from "@tremor/react"; -import { Button, Typography, Spin, Switch, Select, InputNumber } from "antd"; -import { PlusOutlined, DeleteOutlined } from "@ant-design/icons"; -import { getInternalUserSettings, updateInternalUserSettings, modelAvailableCall } from "@/components/networking"; -import BudgetDurationDropdown, { - getBudgetDurationLabel, -} from "@/components/common_components/budget_duration_dropdown"; -import { getModelDisplayName } from "@/components/key_team_helpers/fetch_available_models_team_key"; -import { formatNumberWithCommas } from "@/utils/dataUtils"; -import NotificationManager from "@/components/molecules/notifications_manager"; - -interface DefaultUserSettingsProps { - accessToken: string | null; - possibleUIRoles?: Record> | null; - userID: string; - userRole: string; -} - -interface TeamEntry { - team_id: string; - max_budget_in_team?: number; - user_role: "user" | "admin"; -} - -const DefaultUserSettings: React.FC = ({ - accessToken, - possibleUIRoles, - userID, - userRole, -}) => { - const [loading, setLoading] = useState(true); - const [settings, setSettings] = useState(null); - const [isEditing, setIsEditing] = useState(false); - const [editedValues, setEditedValues] = useState({}); - const [saving, setSaving] = useState(false); - const [availableModels, setAvailableModels] = useState([]); - const { Paragraph } = Typography; - const { Option } = Select; - - useEffect(() => { - const fetchSSOSettings = async () => { - if (!accessToken) { - setLoading(false); - return; - } - - try { - const data = await getInternalUserSettings(accessToken); - setSettings(data); - setEditedValues(data.values || {}); - - // Fetch available models - if (accessToken) { - try { - const modelResponse = await modelAvailableCall(accessToken, userID, userRole); - if (modelResponse && modelResponse.data) { - const modelNames = modelResponse.data.map((model: { id: string }) => model.id); - setAvailableModels(modelNames); - } - } catch (error) { - console.error("Error fetching available models:", error); - } - } - } catch (error) { - console.error("Error fetching SSO settings:", error); - NotificationManager.fromBackend("Failed to fetch SSO settings"); - } finally { - setLoading(false); - } - }; - - fetchSSOSettings(); - }, [accessToken]); - - const handleSaveSettings = async () => { - if (!accessToken) return; - - setSaving(true); - try { - // Convert empty strings to null - const processedValues = Object.entries(editedValues).reduce( - (acc, [key, value]) => { - acc[key] = value === "" ? null : value; - return acc; - }, - {} as Record, - ); - - const updatedSettings = await updateInternalUserSettings(accessToken, processedValues); - setSettings({ ...settings, values: updatedSettings.settings }); - setIsEditing(false); - } catch (error) { - console.error("Error updating SSO settings:", error); - NotificationManager.fromBackend("Failed to update settings: " + error); - } finally { - setSaving(false); - } - }; - - const handleTextInputChange = (key: string, value: any) => { - setEditedValues((prev: Record) => ({ - ...prev, - [key]: value, - })); - }; - - // Helper function to normalize teams array to consistent format - const normalizeTeams = (teams: any[]): TeamEntry[] => { - if (!teams || !Array.isArray(teams)) return []; - - return teams.map((team) => { - if (typeof team === "string") { - return { - team_id: team, - user_role: "user" as const, - }; - } else if (typeof team === "object" && team.team_id) { - return { - team_id: team.team_id, - max_budget_in_team: team.max_budget_in_team, - user_role: team.user_role || "user", - }; - } - return { - team_id: "", - user_role: "user" as const, - }; - }); - }; - - // Teams editor component - const renderTeamsEditor = (teams: any[]) => { - const normalizedTeams = normalizeTeams(teams); - - const updateTeam = (index: number, field: keyof TeamEntry, value: any) => { - const updatedTeams = [...normalizedTeams]; - updatedTeams[index] = { - ...updatedTeams[index], - [field]: value, - }; - handleTextInputChange("teams", updatedTeams); - }; - - const addTeam = () => { - const newTeam: TeamEntry = { - team_id: "", - user_role: "user", - }; - handleTextInputChange("teams", [...normalizedTeams, newTeam]); - }; - - const removeTeam = (index: number) => { - const updatedTeams = normalizedTeams.filter((_, i) => i !== index); - handleTextInputChange("teams", updatedTeams); - }; - - return ( -
- {normalizedTeams.map((team, index) => ( -
-
- Team {index + 1} - -
- -
-
- Team ID - updateTeam(index, "team_id", e.target.value)} - placeholder="Enter team ID" - /> -
- -
- Max Budget in Team - updateTeam(index, "max_budget_in_team", value)} - placeholder="Optional" - min={0} - step={0.01} - precision={2} - /> -
- -
- User Role - -
-
-
- ))} - - -
- ); - }; - - const renderEditableField = (key: string, property: any, value: any) => { - const type = property.type; - - if (key === "teams") { - return
{renderTeamsEditor(editedValues[key] || [])}
; - } else if (key === "user_role" && possibleUIRoles) { - return ( - - ); - } else if (key === "budget_duration") { - return ( - handleTextInputChange(key, value)} - className="mt-2" - /> - ); - } else if (type === "boolean") { - return ( -
- handleTextInputChange(key, checked)} /> -
- ); - } else if (type === "array" && property.items?.enum) { - return ( - - ); - } else if (key === "models") { - return ( - - ); - } else if (type === "string" && property.enum) { - return ( - - ); - } else { - return ( - handleTextInputChange(key, e.target.value)} - placeholder={property.description || ""} - className="mt-2" - /> - ); - } - }; - - const renderValue = (key: string, value: any): JSX.Element => { - if (value === null || value === undefined) return Not set; - - if (key === "teams" && Array.isArray(value)) { - if (value.length === 0) return No teams assigned; - - const normalizedTeams = normalizeTeams(value); - - return ( -
- {normalizedTeams.map((team, index) => ( -
-
-
- Team ID: -

{team.team_id || "Not specified"}

-
-
- Max Budget: -

- {team.max_budget_in_team !== undefined - ? `$${formatNumberWithCommas(team.max_budget_in_team, 4)}` - : "No limit"} -

-
-
- Role: -

{team.user_role}

-
-
-
- ))} -
- ); - } - - if (key === "user_role" && possibleUIRoles && possibleUIRoles[value]) { - const { ui_label, description } = possibleUIRoles[value]; - return ( -
- {ui_label} - {description &&

{description}

} -
- ); - } - - if (key === "budget_duration") { - return {getBudgetDurationLabel(value)}; - } - - if (typeof value === "boolean") { - return {value ? "Enabled" : "Disabled"}; - } - - if (key === "models" && Array.isArray(value)) { - if (value.length === 0) return None; - - return ( -
- {value.map((model, index) => ( - - {getModelDisplayName(model)} - - ))} -
- ); - } - - if (typeof value === "object") { - if (Array.isArray(value)) { - if (value.length === 0) return None; - - return ( -
- {value.map((item, index) => ( - - {typeof item === "object" ? JSON.stringify(item) : String(item)} - - ))} -
- ); - } - - return ( - 
{JSON.stringify(value, null, 2)}
- ); - } - - return {String(value)}; - }; - - if (loading) { - return ( -
- -
- ); - } - - if (!settings) { - return ( - - No settings available or you do not have permission to view them. - - ); - } - - // Dynamically render settings based on the schema - const renderSettings = () => { - const { values, field_schema } = settings; - - if (!field_schema || !field_schema.properties) { - return No schema information available; - } - - return Object.entries(field_schema.properties).map(([key, property]: [string, any]) => { - const value = values[key]; - const displayName = key.replace(/_/g, " ").replace(/\b\w/g, (l) => l.toUpperCase()); - - return ( -
- {displayName} - - {property.description || "No description available"} - - - {isEditing ? ( -
{renderEditableField(key, property, value)}
- ) : ( -
{renderValue(key, value)}
- )} -
- ); - }); - }; - - return ( - -
- Default User Settings - {!loading && - settings && - (isEditing ? ( -
- - -
- ) : ( - - ))} -
- - {settings?.field_schema?.description && ( - {settings.field_schema.description} - )} - - -
{renderSettings()}
- - ); -}; - -export default DefaultUserSettings; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/default-user-settings/DefaultUserSettingsForm.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/default-user-settings/DefaultUserSettingsForm.test.tsx new file mode 100644 index 00000000000..bfdcc70fb0c --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/default-user-settings/DefaultUserSettingsForm.test.tsx @@ -0,0 +1,296 @@ +import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; +import { render, screen, waitFor } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import React from "react"; +import { beforeEach, describe, expect, it, vi } from "vitest"; + +vi.mock("@/app/(dashboard)/hooks/teams/useTeams", () => ({ + useInfiniteTeams: () => ({ + data: { + pages: [ + { + teams: [ + { team_id: "team-alpha", team_alias: "Alpha" }, + { team_id: "team-beta", team_alias: "Beta" }, + ], + }, + ], + }, + fetchNextPage: vi.fn(), + hasNextPage: false, + isFetchingNextPage: false, + isLoading: false, + }), +})); + +vi.mock("@/components/ModelSelect/ModelSelect", async (importOriginal) => { + const actual = await importOriginal(); + return { + MODEL_SENTINEL_OPTIONS: actual.MODEL_SENTINEL_OPTIONS, + ModelSelect: ({ onChange }: { onChange: (values: string[]) => void }) => ( + + ), + }; +}); + +import NotificationsManager from "@/components/molecules/notifications_manager"; + +import { DefaultUserSettingsForm } from "./DefaultUserSettingsForm"; +import type { InternalUserSettings } from "./mapper"; + +const POSSIBLE_UI_ROLES = { + internal_user: { ui_label: "Internal User", description: "create and view own keys" }, + internal_user_viewer: { ui_label: "Internal Viewer", description: "view own keys" }, + proxy_admin: { ui_label: "Admin", description: "all permissions" }, +}; + +const SETTINGS: InternalUserSettings = { + values: { + user_role: "internal_user", + max_budget: 100, + budget_duration: "30d", + models: ["gpt-5.2"], + teams: [{ team_id: "team-alpha", max_budget_in_team: 25, user_role: "user" }], + }, + field_schema: {}, +}; + +const SAVED_BODY = { + user_role: "internal_user", + max_budget: 100, + budget_duration: "30d", + models: ["gpt-5.2"], + teams: [{ team_id: "team-alpha", max_budget_in_team: 25, user_role: "user" }], +}; + +const renderForm = (overrides?: { + fetchSettings?: ReturnType; + updateSettings?: ReturnType; +}) => { + const fetchSettings = overrides?.fetchSettings ?? vi.fn().mockResolvedValue(SETTINGS); + const updateSettings = overrides?.updateSettings ?? vi.fn().mockResolvedValue(undefined); + const queryClient = new QueryClient({ defaultOptions: { queries: { retry: false } } }); + + render( + + + , + ); + + return { fetchSettings, updateSettings }; +}; + +const saveButton = async () => await screen.findByRole("button", { name: "Save Changes" }); + +const enterEditMode = async (user: ReturnType) => { + await user.click(await screen.findByRole("button", { name: "Edit Settings" })); +}; + +describe("DefaultUserSettingsForm", () => { + beforeEach(() => { + vi.clearAllMocks(); + }); + + it("shows a read-only summary until Edit Settings is clicked", async () => { + renderForm(); + + expect(await screen.findByText("Internal User")).toBeInTheDocument(); + expect(screen.getByText("100")).toBeInTheDocument(); + expect(screen.getByText("monthly")).toBeInTheDocument(); + expect(screen.getByText("gpt-5.2")).toBeInTheDocument(); + expect(screen.getByText(/team-alpha/)).toBeInTheDocument(); + expect(screen.queryByRole("button", { name: "Save Changes" })).not.toBeInTheDocument(); + expect(screen.queryByLabelText("Max Budget (USD)")).not.toBeInTheDocument(); + }); + + it("labels model sentinels in the read-only summary", async () => { + renderForm({ + fetchSettings: vi + .fn() + .mockResolvedValue({ ...SETTINGS, values: { ...SETTINGS.values, models: ["all-proxy-models"] } }), + }); + + expect(await screen.findByText("All Proxy Models")).toBeInTheDocument(); + }); + + it("disables Save until the loaded settings are edited", async () => { + const user = userEvent.setup(); + renderForm(); + + await enterEditMode(user); + + expect(await saveButton()).toBeDisabled(); + }); + + it("shows an error instead of the form when the settings cannot be loaded", async () => { + renderForm({ fetchSettings: vi.fn().mockRejectedValue(new Error("nope")) }); + + expect(await screen.findByRole("alert")).toHaveTextContent("Could not load the default user settings."); + expect(screen.queryByRole("button", { name: "Edit Settings" })).not.toBeInTheDocument(); + expect(screen.queryByRole("button", { name: "Save Changes" })).not.toBeInTheDocument(); + }); + + it("sends every field on save, not only the edited one", async () => { + const user = userEvent.setup(); + const { updateSettings } = renderForm(); + + await enterEditMode(user); + await user.clear(await screen.findByLabelText("Max Budget (USD)")); + await user.type(screen.getByLabelText("Max Budget (USD)"), "250"); + await user.click(await saveButton()); + + await waitFor(() => expect(updateSettings).toHaveBeenCalledTimes(1)); + expect(updateSettings).toHaveBeenCalledWith({ ...SAVED_BODY, max_budget: 250 }); + }); + + it("clears an emptied budget with null", async () => { + const user = userEvent.setup(); + const { updateSettings } = renderForm(); + + await enterEditMode(user); + await user.clear(await screen.findByLabelText("Max Budget (USD)")); + await user.click(await saveButton()); + + await waitFor(() => expect(updateSettings).toHaveBeenCalledTimes(1)); + expect(updateSettings).toHaveBeenCalledWith({ ...SAVED_BODY, max_budget: null }); + }); + + it("sends the models selection through unchanged, sentinel values included", async () => { + const user = userEvent.setup(); + const { updateSettings } = renderForm(); + + await enterEditMode(user); + await user.click(await screen.findByRole("button", { name: "set-models" })); + await user.click(await saveButton()); + + await waitFor(() => expect(updateSettings).toHaveBeenCalledTimes(1)); + expect(updateSettings).toHaveBeenCalledWith({ ...SAVED_BODY, models: ["all-proxy-models"] }); + }); + + it("saves a team that was picked from the searchable list", async () => { + const user = userEvent.setup(); + const { updateSettings } = renderForm(); + + await enterEditMode(user); + await user.click(await screen.findByRole("button", { name: "Add Team" })); + await user.click(screen.getAllByLabelText("Team")[1]); + await user.click(await screen.findByText("Beta")); + await user.click(await saveButton()); + + await waitFor(() => expect(updateSettings).toHaveBeenCalledTimes(1)); + expect(updateSettings).toHaveBeenCalledWith({ + ...SAVED_BODY, + teams: [ + { team_id: "team-alpha", max_budget_in_team: 25, user_role: "user" }, + { team_id: "team-beta", max_budget_in_team: null, user_role: "user" }, + ], + }); + }); + + it("never turns a team id typed into the picker into a saved team", async () => { + const user = userEvent.setup(); + const { updateSettings } = renderForm(); + + await enterEditMode(user); + await user.click(await screen.findByRole("button", { name: "Add Team" })); + await user.type(screen.getAllByLabelText("Team")[1], "team-alhpa"); + await user.keyboard("{Escape}"); + await user.click(await saveButton()); + + expect(await screen.findByText("Select a team")).toBeInTheDocument(); + expect(screen.getAllByLabelText("Team")[1]).toHaveValue(""); + expect(updateSettings).not.toHaveBeenCalled(); + }); + + it("blocks saving the same default team twice", async () => { + const user = userEvent.setup(); + const { updateSettings } = renderForm(); + + await enterEditMode(user); + await user.click(await screen.findByRole("button", { name: "Add Team" })); + await user.click(screen.getAllByLabelText("Team")[1]); + await user.click(await screen.findByText("Alpha")); + await user.click(await saveButton()); + + expect(await screen.findByText("This team is already listed")).toBeInTheDocument(); + expect(updateSettings).not.toHaveBeenCalled(); + }); + + it("drops a removed team row from the saved settings", async () => { + const user = userEvent.setup(); + const { updateSettings } = renderForm(); + + await enterEditMode(user); + await user.click(await screen.findByRole("button", { name: "Remove" })); + await user.click(await saveButton()); + + await waitFor(() => expect(updateSettings).toHaveBeenCalledTimes(1)); + expect(updateSettings).toHaveBeenCalledWith({ ...SAVED_BODY, teams: null }); + }); + + it("returns to the read-only view showing the new values after a successful save", async () => { + const user = userEvent.setup(); + const updated = { ...SETTINGS, values: { ...SETTINGS.values, max_budget: 250 } }; + const { updateSettings } = renderForm({ + fetchSettings: vi.fn().mockResolvedValueOnce(SETTINGS).mockResolvedValue(updated), + }); + + await enterEditMode(user); + await user.clear(await screen.findByLabelText("Max Budget (USD)")); + await user.type(screen.getByLabelText("Max Budget (USD)"), "250"); + await user.click(await saveButton()); + + await waitFor(() => expect(updateSettings).toHaveBeenCalledTimes(1)); + expect(await screen.findByRole("button", { name: "Edit Settings" })).toBeInTheDocument(); + expect(await screen.findByText("250")).toBeInTheDocument(); + expect(screen.queryByRole("button", { name: "Save Changes" })).not.toBeInTheDocument(); + expect(NotificationsManager.success).toHaveBeenCalledWith("Default user settings updated successfully"); + + await enterEditMode(user); + expect(await saveButton()).toBeDisabled(); + }); + + it("keeps the edit and surfaces the backend error when the save fails", async () => { + const user = userEvent.setup(); + const { updateSettings } = renderForm({ + updateSettings: vi.fn().mockRejectedValue(new Error("Team(s) not found: team-alhpa.")), + }); + + await enterEditMode(user); + await user.clear(await screen.findByLabelText("Max Budget (USD)")); + await user.type(screen.getByLabelText("Max Budget (USD)"), "250"); + await user.click(await saveButton()); + + await waitFor(() => expect(updateSettings).toHaveBeenCalledTimes(1)); + await waitFor(() => + expect(NotificationsManager.fromBackend).toHaveBeenCalledWith("Team(s) not found: team-alhpa."), + ); + expect(await saveButton()).toBeEnabled(); + expect(screen.getByLabelText("Max Budget (USD)")).toHaveValue(250); + }); + + it("discards edits and returns to the read-only view when Cancel is pressed", async () => { + const user = userEvent.setup(); + const { updateSettings } = renderForm(); + + await enterEditMode(user); + await user.clear(await screen.findByLabelText("Max Budget (USD)")); + await user.type(screen.getByLabelText("Max Budget (USD)"), "250"); + await user.click(screen.getByRole("button", { name: "Cancel" })); + + expect(await screen.findByRole("button", { name: "Edit Settings" })).toBeInTheDocument(); + expect(screen.getByText("100")).toBeInTheDocument(); + expect(screen.queryByRole("button", { name: "Save Changes" })).not.toBeInTheDocument(); + expect(updateSettings).not.toHaveBeenCalled(); + + await enterEditMode(user); + expect(screen.getByLabelText("Max Budget (USD)")).toHaveValue(100); + expect(await saveButton()).toBeDisabled(); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/default-user-settings/DefaultUserSettingsForm.tsx b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/default-user-settings/DefaultUserSettingsForm.tsx new file mode 100644 index 00000000000..b1474e7cd0c --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/default-user-settings/DefaultUserSettingsForm.tsx @@ -0,0 +1,433 @@ +"use client"; + +import { useMutation, useQuery, useQueryClient } from "@tanstack/react-query"; +import * as React from "react"; +import { useFieldArray, type Control } from "react-hook-form"; + +import { useInfiniteTeams } from "@/app/(dashboard)/hooks/teams/useTeams"; +import { ModelSelect, MODEL_SENTINEL_OPTIONS } from "@/components/ModelSelect/ModelSelect"; +import NotificationsManager from "@/components/molecules/notifications_manager"; +import { PaginatedSearchSelect } from "@/components/shared/PaginatedSearchSelect"; +import { FieldGroup } from "@/components/shared/form/field"; +import { FormField } from "@/components/shared/form/FormField"; +import type { SearchSelectOption } from "@/components/shared/SearchSelect"; +import { Button } from "@/components/ui/button"; +import { Card, CardAction, CardContent, CardDescription, CardHeader, CardTitle } from "@/components/ui/card"; +import { Input } from "@/components/ui/input"; +import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; +import { Skeleton } from "@/components/ui/skeleton"; +import { useZodForm } from "@/lib/forms/useZodForm"; +import { fetchClient } from "@/lib/http/api"; + +import { buildBody, settingsToForm, type DefaultInternalUserParams, type InternalUserSettings } from "./mapper"; +import { defaultUserSettingsSchema, EMPTY_TEAM_ROW, type DefaultUserSettingsFormValues } from "./schema"; + +const NO_RESET = "never"; + +const BUDGET_DURATION_OPTIONS = [ + { value: NO_RESET, label: "No reset" }, + { value: "1h", label: "hourly" }, + { value: "24h", label: "daily" }, + { value: "7d", label: "weekly" }, + { value: "30d", label: "monthly" }, +] as const; + +const TEAM_ROLE_OPTIONS = [ + { value: "user", label: "User" }, + { value: "admin", label: "Admin" }, +] as const; + +const MODEL_SENTINEL_LABELS: ReadonlyMap = new Map( + MODEL_SENTINEL_OPTIONS.map(({ value, label }) => [value, label]), +); + +const TEAMS_PAGE_SIZE = 50; + +const SETTINGS_QUERY_KEY = ["internalUserSettings"] as const; + +const defaultFetchSettings = async (): Promise => { + const { data } = await fetchClient.GET("/get/internal_user_settings"); + if (data === undefined) { + throw new Error("Failed to load default user settings"); + } + return data; +}; + +const defaultUpdateSettings = async (body: DefaultInternalUserParams): Promise => { + await fetchClient.PATCH("/update/internal_user_settings", { body }); +}; + +interface RoleOption { + value: string; + label: string; + description: string; +} + +type SettingsControl = Control; + +const TeamPickerField = ({ control, index }: { control: SettingsControl; index: number }) => { + const [search, setSearch] = React.useState(""); + const { data, fetchNextPage, hasNextPage, isFetchingNextPage, isLoading } = useInfiniteTeams( + TEAMS_PAGE_SIZE, + search === "" ? undefined : search, + ); + + const options = React.useMemo( + () => + (data?.pages ?? []).flatMap((page) => + page.teams.map((team) => ({ + label: team.team_alias || team.team_id, + value: team.team_id, + sublabel: team.team_id, + })), + ), + [data], + ); + + return ( + + {({ id, value, onChange, "aria-invalid": ariaInvalid, "aria-describedby": ariaDescribedBy }) => ( + void fetchNextPage()} + hasNextPage={hasNextPage} + isLoading={isLoading} + isFetchingNextPage={isFetchingNextPage} + placeholder="Search a team" + emptyText="No teams found" + inputId={id} + aria-invalid={ariaInvalid} + aria-describedby={ariaDescribedBy} + /> + )} + + ); +}; + +const TeamsField = ({ control }: { control: SettingsControl }) => { + const { fields, append, remove } = useFieldArray({ control, name: "teams" }); + + return ( +
+
+

Default Teams

+

+ New users are added to these teams. Only teams that already exist can be selected. +

+
+ + {fields.map((field, index) => ( +
+
+

Team {index + 1}

+ +
+ +
+ + + + {({ ref, ...budgetField }) => ( + + )} + + + + {({ id, value, onChange, "aria-invalid": ariaInvalid, "aria-describedby": ariaDescribedBy }) => ( + + )} + +
+
+ ))} + + +
+ ); +}; + +const ViewRow = ({ label, children }: { label: string; children: React.ReactNode }) => ( +
+

{label}

+

{children}

+
+); + +interface SettingsViewProps { + values: DefaultUserSettingsFormValues; + roleOptions: readonly RoleOption[]; +} + +const SettingsView = ({ values, roleOptions }: SettingsViewProps) => { + const roleLabel = roleOptions.find((option) => option.value === values.user_role)?.label ?? values.user_role; + const durationValue = values.budget_duration === "" ? NO_RESET : values.budget_duration; + const durationLabel = + BUDGET_DURATION_OPTIONS.find((option) => option.value === durationValue)?.label ?? values.budget_duration; + + return ( +
+ {roleLabel === "" ? "Not set" : roleLabel} + {values.max_budget === "" ? "Not set" : values.max_budget} + {durationLabel} + + {values.models.length === 0 + ? "Not set" + : values.models.map((model) => MODEL_SENTINEL_LABELS.get(model) ?? model).join(", ")} + +
+

Default Teams

+ {values.teams.length === 0 ? ( +

None

+ ) : ( + values.teams.map((team) => ( +

+ {team.team_id} + {team.max_budget_in_team !== "" && <> · ${team.max_budget_in_team} max budget} + <> · {team.user_role} +

+ )) + )} +
+
+ ); +}; + +interface SettingsFormProps { + initialValues: DefaultUserSettingsFormValues; + roleOptions: readonly RoleOption[]; + updateSettings: (body: DefaultInternalUserParams) => Promise; + onCancel: () => void; + onSaved: () => void; +} + +const SettingsForm = ({ initialValues, roleOptions, updateSettings, onCancel, onSaved }: SettingsFormProps) => { + const queryClient = useQueryClient(); + const form = useZodForm(defaultUserSettingsSchema, { defaultValues: initialValues }); + const { isDirty } = form.formState; + + const mutation = useMutation({ + mutationFn: (values: DefaultUserSettingsFormValues) => updateSettings(buildBody(values)), + onSuccess: (_result, values) => { + NotificationsManager.success("Default user settings updated successfully"); + queryClient.invalidateQueries({ queryKey: SETTINGS_QUERY_KEY }); + form.reset(values); + onSaved(); + }, + onError: (error: unknown) => + NotificationsManager.fromBackend( + error instanceof Error ? error.message : "Failed to update default user settings", + ), + }); + + const onSubmit = form.handleSubmit((values) => mutation.mutate(values)); + + return ( +
+ + + {({ id, value, onChange, "aria-invalid": ariaInvalid, "aria-describedby": ariaDescribedBy }) => ( + + )} + + + + {({ ref, ...field }) => } + + + + {({ id, value, onChange, "aria-invalid": ariaInvalid, "aria-describedby": ariaDescribedBy }) => ( + + )} + + + + {(field) => ( + + )} + + + + + +
+ + +
+
+ ); +}; + +const SettingsCard = ({ action, children }: { action?: React.ReactNode; children: React.ReactNode }) => ( + + + Default User Settings + + Applied to every new internal user created through SSO or the user management APIs. + + {action !== undefined && {action}} + + {children} + +); + +export interface DefaultUserSettingsFormProps { + possibleUIRoles?: Record> | null; + fetchSettings?: () => Promise; + updateSettings?: (body: DefaultInternalUserParams) => Promise; +} + +export const DefaultUserSettingsForm = ({ + possibleUIRoles, + fetchSettings = defaultFetchSettings, + updateSettings = defaultUpdateSettings, +}: DefaultUserSettingsFormProps) => { + const [isEditing, setIsEditing] = React.useState(false); + const { data, isPending, isError } = useQuery({ queryKey: SETTINGS_QUERY_KEY, queryFn: fetchSettings }); + + const roleOptions = React.useMemo( + () => + Object.entries(possibleUIRoles ?? {}) + .filter(([role]) => role.includes("internal_user")) + .map(([role, meta]) => ({ value: role, label: meta.ui_label || role, description: meta.description ?? "" })), + [possibleUIRoles], + ); + + const initialValues = React.useMemo(() => (data === undefined ? undefined : settingsToForm(data.values)), [data]); + + if (isPending) { + return ( + + + + ); + } + + if (isError || initialValues === undefined) { + return ( + +

Could not load the default user settings.

+ + ); + } + + return ( + setIsEditing(true)}> + Edit Settings + + ) + } + > + {isEditing ? ( + setIsEditing(false)} + onSaved={() => setIsEditing(false)} + /> + ) : ( + + )} + + ); +}; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/default-user-settings/mapper.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/default-user-settings/mapper.test.ts new file mode 100644 index 00000000000..e8b350332c8 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/default-user-settings/mapper.test.ts @@ -0,0 +1,116 @@ +import { describe, expect, it } from "vitest"; + +import { buildBody, settingsToForm } from "./mapper"; +import type { DefaultUserSettingsFormValues } from "./schema"; + +const CONFIGURED_SETTINGS = { + user_role: "internal_user", + max_budget: 100.5, + budget_duration: "30d", + models: ["gpt-5.2"], + teams: [{ team_id: "team-1", max_budget_in_team: 25, user_role: "admin" }], +}; + +const UNCONFIGURED_SETTINGS = { + user_role: "internal_user_viewer", + max_budget: null, + budget_duration: null, + models: null, + teams: null, +}; + +const CONFIGURED_FORM = { + user_role: "internal_user", + max_budget: "100.5", + budget_duration: "30d", + models: ["gpt-5.2"], + teams: [{ team_id: "team-1", max_budget_in_team: "25", user_role: "admin" }], +}; + +const UNCONFIGURED_FORM = { + user_role: "internal_user_viewer", + max_budget: "", + budget_duration: "", + models: [], + teams: [], +}; + +describe("settingsToForm", () => { + it("maps a fully populated settings blob onto widget-space strings", () => { + expect(settingsToForm(CONFIGURED_SETTINGS)).toStrictEqual(CONFIGURED_FORM); + }); + + it("maps an unconfigured settings blob onto empty widget state", () => { + expect(settingsToForm(UNCONFIGURED_SETTINGS)).toStrictEqual(UNCONFIGURED_FORM); + }); + + it("hydrates the legacy list-of-team-ids shape as full team rows", () => { + expect(settingsToForm({ teams: ["team-1", "team-2"] }).teams).toStrictEqual([ + { team_id: "team-1", max_budget_in_team: "", user_role: "user" }, + { team_id: "team-2", max_budget_in_team: "", user_role: "user" }, + ]); + }); + + it("defaults a team row's role to user and leaves an absent in-team budget blank", () => { + expect(settingsToForm({ teams: [{ team_id: "team-1" }] }).teams).toStrictEqual([ + { team_id: "team-1", max_budget_in_team: "", user_role: "user" }, + ]); + }); + + it("degrades an unrecognisable team entry to a blank row instead of throwing", () => { + expect(settingsToForm({ teams: [{ max_budget_in_team: 5 }, 7] }).teams).toStrictEqual([ + { team_id: "", max_budget_in_team: "", user_role: "user" }, + { team_id: "", max_budget_in_team: "", user_role: "user" }, + ]); + }); +}); + +const formValues = (overrides: Partial = {}): DefaultUserSettingsFormValues => ({ + user_role: "internal_user", + max_budget: "100", + budget_duration: "30d", + models: ["gpt-5.2"], + teams: [{ team_id: "team-1", max_budget_in_team: "25", user_role: "admin" }], + ...overrides, +}); + +const SAVED_BODY = { + user_role: "internal_user", + max_budget: 100, + budget_duration: "30d", + models: ["gpt-5.2"], + teams: [{ team_id: "team-1", max_budget_in_team: 25, user_role: "admin" }], +}; + +const CLEARED_FORM = { user_role: "", max_budget: "", budget_duration: "", models: [], teams: [] }; + +const CLEARED_BODY = { + user_role: null, + max_budget: null, + budget_duration: null, + models: null, + teams: null, +}; + +describe("buildBody", () => { + it("sends every field, because the endpoint replaces the whole settings object", () => { + expect(buildBody(formValues())).toStrictEqual(SAVED_BODY); + }); + + it("clears emptied fields with null so the backend drops them", () => { + expect(buildBody(formValues(CLEARED_FORM))).toStrictEqual(CLEARED_BODY); + }); + + it("sends teams as objects and nulls an in-team budget that was left blank", () => { + expect( + buildBody(formValues({ teams: [{ team_id: "team-9", max_budget_in_team: "", user_role: "user" }] })), + ).toStrictEqual({ + ...SAVED_BODY, + teams: [{ team_id: "team-9", max_budget_in_team: null, user_role: "user" }], + }); + }); + + it("refuses to send a role the backend does not accept", () => { + expect(buildBody(formValues({ user_role: "made_up_role" })).user_role).toBeNull(); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/default-user-settings/mapper.ts b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/default-user-settings/mapper.ts new file mode 100644 index 00000000000..50365081afb --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/default-user-settings/mapper.ts @@ -0,0 +1,75 @@ +import { z } from "zod/v4"; + +import type { components } from "@/lib/http/schema"; + +import { EMPTY_TEAM_ROW, type DefaultTeamRowValues, type DefaultUserSettingsFormValues } from "./schema"; + +export type InternalUserSettings = components["schemas"]["InternalUserSettingsResponse"]; +export type DefaultInternalUserParams = components["schemas"]["DefaultInternalUserParams"]; +type DefaultTeamBody = components["schemas"]["NewUserRequestTeam"]; + +const teamRowFromServer = z + .union([ + z.string().transform((teamId): DefaultTeamRowValues => ({ ...EMPTY_TEAM_ROW, team_id: teamId })), + z + .object({ + team_id: z.string(), + max_budget_in_team: z.number().nullish(), + user_role: z.enum(["user", "admin"]).catch("user"), + }) + .transform( + (team): DefaultTeamRowValues => ({ + team_id: team.team_id, + max_budget_in_team: team.max_budget_in_team?.toString() ?? "", + user_role: team.user_role, + }), + ), + ]) + .catch(EMPTY_TEAM_ROW); + +const serverValuesShape = { + user_role: z.string().nullish().catch(null), + max_budget: z.number().nullish().catch(null), + budget_duration: z.string().nullish().catch(null), + models: z.array(z.string()).nullish().catch(null), + teams: z.array(teamRowFromServer).nullish().catch(null), +}; + +const serverValuesSchema = z.object(serverValuesShape); + +export const settingsToForm = (values: InternalUserSettings["values"]): DefaultUserSettingsFormValues => { + const parsed = serverValuesSchema.parse(values); + + return { + user_role: parsed.user_role ?? "", + max_budget: parsed.max_budget?.toString() ?? "", + budget_duration: parsed.budget_duration ?? "", + models: parsed.models ?? [], + teams: parsed.teams ?? [], + }; +}; + +const DEFAULT_USER_ROLES = ["internal_user", "internal_user_viewer", "proxy_admin", "proxy_admin_viewer"] as const; + +const asDefaultUserRole = (raw: string): DefaultInternalUserParams["user_role"] => + DEFAULT_USER_ROLES.find((role) => role === raw) ?? null; + +const numberOrNull = (raw: string): number | null => (raw.trim() === "" ? null : Number(raw)); + +const textOrNull = (raw: string): string | null => (raw.trim() === "" ? null : raw); + +const listOrNull = (items: readonly T[]): T[] | null => (items.length === 0 ? null : [...items]); + +const toTeamBody = (team: DefaultTeamRowValues): DefaultTeamBody => ({ + team_id: team.team_id, + max_budget_in_team: numberOrNull(team.max_budget_in_team), + user_role: team.user_role, +}); + +export const buildBody = (values: DefaultUserSettingsFormValues): DefaultInternalUserParams => ({ + user_role: asDefaultUserRole(values.user_role), + max_budget: numberOrNull(values.max_budget), + budget_duration: textOrNull(values.budget_duration), + models: listOrNull(values.models), + teams: listOrNull(values.teams.map(toTeamBody)), +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/default-user-settings/schema.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/default-user-settings/schema.test.ts new file mode 100644 index 00000000000..889f87f2203 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/default-user-settings/schema.test.ts @@ -0,0 +1,64 @@ +import { describe, expect, it } from "vitest"; + +import { defaultUserSettingsSchema, type DefaultUserSettingsFormValues } from "./schema"; + +const values = (overrides: Partial = {}): DefaultUserSettingsFormValues => ({ + user_role: "internal_user", + max_budget: "", + budget_duration: "", + models: [], + teams: [], + ...overrides, +}); + +const issuesFor = (input: DefaultUserSettingsFormValues) => { + const result = defaultUserSettingsSchema.safeParse(input); + return result.success ? [] : result.error.issues.map((issue) => ({ path: issue.path, message: issue.message })); +}; + +describe("defaultUserSettingsSchema", () => { + it("accepts settings with no default teams", () => { + expect(defaultUserSettingsSchema.safeParse(values()).success).toBe(true); + }); + + it("rejects a team row that has no team selected", () => { + expect(issuesFor(values({ teams: [{ team_id: "", max_budget_in_team: "", user_role: "user" }] }))).toStrictEqual([ + { path: ["teams", 0, "team_id"], message: "Select a team" }, + ]); + }); + + it("rejects the same team appearing twice", () => { + expect( + issuesFor( + values({ + teams: [ + { team_id: "team-1", max_budget_in_team: "", user_role: "user" }, + { team_id: "team-1", max_budget_in_team: "", user_role: "admin" }, + ], + }), + ), + ).toStrictEqual([{ path: ["teams", 1, "team_id"], message: "This team is already listed" }]); + }); + + it("does not treat two blank rows as duplicates of each other", () => { + const issues = issuesFor( + values({ + teams: [ + { team_id: "", max_budget_in_team: "", user_role: "user" }, + { team_id: "", max_budget_in_team: "", user_role: "user" }, + ], + }), + ); + + expect(issues.map((issue) => issue.message)).toStrictEqual(["Select a team", "Select a team"]); + }); + + it("rejects non-numeric budgets on the form and on a team row", () => { + expect(issuesFor(values({ max_budget: "lots" }))).toStrictEqual([ + { path: ["max_budget"], message: "Must be a non-negative number" }, + ]); + expect( + issuesFor(values({ teams: [{ team_id: "team-1", max_budget_in_team: "-5", user_role: "user" }] })), + ).toStrictEqual([{ path: ["teams", 0, "max_budget_in_team"], message: "Must be a non-negative number" }]); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/default-user-settings/schema.ts b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/default-user-settings/schema.ts new file mode 100644 index 00000000000..7309e3745da --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/default-user-settings/schema.ts @@ -0,0 +1,44 @@ +import { z } from "zod/v4"; + +const isBlank = (value: string): boolean => value.trim() === ""; + +const amountOrEmpty = z + .string() + .refine( + (value) => isBlank(value) || (Number.isFinite(Number(value)) && Number(value) >= 0), + "Must be a non-negative number", + ); + +const defaultTeamRowSchema = z.object({ + team_id: z.string().min(1, "Select a team"), + max_budget_in_team: amountOrEmpty, + user_role: z.enum(["user", "admin"]), +}); + +export type DefaultTeamRowValues = z.output; + +export const EMPTY_TEAM_ROW: DefaultTeamRowValues = { team_id: "", max_budget_in_team: "", user_role: "user" }; + +const defaultUserSettingsShape = { + user_role: z.string(), + max_budget: amountOrEmpty, + budget_duration: z.string(), + models: z.array(z.string()), + teams: z.array(defaultTeamRowSchema), +}; + +export const defaultUserSettingsSchema = z.object(defaultUserSettingsShape).superRefine((values, ctx) => { + const repeatedRows = values.teams.flatMap((team, index) => + team.team_id !== "" && values.teams.findIndex((other) => other.team_id === team.team_id) < index ? [index] : [], + ); + + repeatedRows.forEach((index) => + ctx.addIssue({ + code: "custom", + message: "This team is already listed", + path: ["teams", index, "team_id"], + }), + ); +}); + +export type DefaultUserSettingsFormValues = z.output; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users.test.tsx index 8dc11babd72..5fcc55c1e98 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users.test.tsx @@ -27,7 +27,6 @@ vi.mock("@/components/networking", () => ({ DEFAULT_TEAM_DISABLED: false, SSO_ENABLED: false, }), - getInternalUserSettings: vi.fn().mockResolvedValue({}), })); // The detail view has its own test; stub it so this file covers the parent's swap. diff --git a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users.tsx b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users.tsx index ce912c09373..2c1d28d82f8 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users.tsx @@ -30,7 +30,7 @@ import { import DeleteResourceModal from "@/components/common_components/DeleteResourceModal"; import NotificationsManager from "@/components/molecules/notifications_manager"; import { modelAvailableCall, userDeleteCall } from "@/components/networking"; -import DefaultUserSettings from "./DefaultUserSettings"; +import { DefaultUserSettingsForm } from "./default-user-settings/DefaultUserSettingsForm"; import { UsersTable } from "./view_users/UsersTable"; import UserInfoView from "./view_users/user_info_view"; import { UserInfo } from "@/components/networking"; @@ -412,12 +412,7 @@ const ViewUserDashboard: React.FC = ({
) : ( - + )} diff --git a/ui/litellm-dashboard/src/components/ModelSelect/ModelSelect.tsx b/ui/litellm-dashboard/src/components/ModelSelect/ModelSelect.tsx index 898374d2a61..0965683c241 100644 --- a/ui/litellm-dashboard/src/components/ModelSelect/ModelSelect.tsx +++ b/ui/litellm-dashboard/src/components/ModelSelect/ModelSelect.tsx @@ -16,7 +16,7 @@ const MODEL_SELECT_NO_DEFAULT_MODELS_SPECIAL_VALUE = { value: "no-default-models", } as const; -const MODEL_SELECT_SPECIAL_VALUES_ARRAY = [ +export const MODEL_SENTINEL_OPTIONS = [ MODEL_SELECT_ALL_PROXY_MODELS_SPECIAL_VALUE, MODEL_SELECT_NO_DEFAULT_MODELS_SPECIAL_VALUE, ] as const; @@ -100,7 +100,7 @@ export const ModelSelect = (props: ModelSelectProps) => { const { data: organization, isLoading: isLoadingOrganization } = useOrganization(organizationID); const { data: currentUser, isLoading: isCurrentUserLoading } = useCurrentUser(); - const isSpecialOption = (value: string) => MODEL_SELECT_SPECIAL_VALUES_ARRAY.some((sv) => sv.value === value); + const isSpecialOption = (value: string) => MODEL_SENTINEL_OPTIONS.some((sv) => sv.value === value); const hasSpecialOptionSelected = value.some(isSpecialOption); const isLoading = isLoadingAllProxyModels || isLoadingTeam || isLoadingOrganization || isCurrentUserLoading; const organizationHasAllProxyModels = diff --git a/ui/litellm-dashboard/src/components/ToolDetail.tsx b/ui/litellm-dashboard/src/components/ToolDetail.tsx index 6a8457a559f..06f14638141 100644 --- a/ui/litellm-dashboard/src/components/ToolDetail.tsx +++ b/ui/litellm-dashboard/src/components/ToolDetail.tsx @@ -430,7 +430,7 @@ export function ToolDetail({ toolName, onBack, accessToken }: ToolDetailProps) {

- Recent logs + Recent invocations

{ - try { - const data = await apiClient.get(`/get/internal_user_settings`, { accessToken }); - return data; - } catch (error) { - console.error("Failed to fetch SSO settings:", error); - throw error; - } -}; - -export const updateInternalUserSettings = async (accessToken: string, settings: Record) => { - try { - // Construct base URL - let url = proxyBaseUrl ? `${proxyBaseUrl}/update/internal_user_settings` : `/update/internal_user_settings`; - - const response = await fetch(url, { - method: "PATCH", - headers: { - [globalLitellmHeaderName]: `Bearer ${accessToken}`, - "Content-Type": "application/json", - }, - body: JSON.stringify(settings), - }); - - if (!response.ok) { - const errorData = await response.text(); - handleError(errorData); - throw new Error(errorData); - } - - const data = await response.json(); - NotificationsManager.success("Internal user settings updated successfully"); - return data; - } catch (error) { - console.error("Failed to update internal user settings:", error); - throw error; - } -}; - export const fetchOpenAPIRegistry = async (accessToken: string) => { try { const url = proxyBaseUrl ? `${proxyBaseUrl}/v1/mcp/openapi-registry` : `/v1/mcp/openapi-registry`; @@ -7583,7 +7544,6 @@ export interface ToolSpendDailyEntry { export interface ToolSpendResponse { by_tool: ToolSpendEntry[]; daily: ToolSpendDailyEntry[]; - total_spend: number; start_date: string | null; end_date: string | null; } diff --git a/ui/litellm-dashboard/src/components/shared/PaginatedSearchSelect.tsx b/ui/litellm-dashboard/src/components/shared/PaginatedSearchSelect.tsx index b59cd1263ea..fde29fd5362 100644 --- a/ui/litellm-dashboard/src/components/shared/PaginatedSearchSelect.tsx +++ b/ui/litellm-dashboard/src/components/shared/PaginatedSearchSelect.tsx @@ -34,6 +34,9 @@ interface PaginatedSearchSelectProps { loadingText?: string; disabled?: boolean; className?: string; + inputId?: string; + "aria-invalid"?: true | undefined; + "aria-describedby"?: string; } export function PaginatedSearchSelect({ @@ -50,6 +53,9 @@ export function PaginatedSearchSelect({ loadingText = "Loading…", disabled = false, className, + inputId, + "aria-invalid": ariaInvalid, + "aria-describedby": ariaDescribedBy, }: PaginatedSearchSelectProps) { const selected = useMemo(() => { if (value === undefined || value === "") return null; @@ -90,6 +96,9 @@ export function PaginatedSearchSelect({ disabled={disabled} > { expect(container.querySelector("style")).toBeNull(); }); + it("colors each bar by its datum when colorByDatum is set, instead of one fill for the series", () => { + const singleCategory = [ + { tool: "alpha", spend: 3 }, + { tool: "beta", spend: 2 }, + { tool: "gamma", spend: 1 }, + ]; + + const { container, rerender } = render( + , + ); + const sharedFills = Array.from(container.querySelectorAll("path.recharts-rectangle")).map((rect) => + rect.getAttribute("fill"), + ); + expect(new Set(sharedFills).size).toBe(1); + + rerender( + , + ); + const perDatumFills = Array.from(container.querySelectorAll("path.recharts-rectangle")).map((rect) => + rect.getAttribute("fill"), + ); + expect(perDatumFills).toEqual([ + "var(--color-blue-500, #3b82f6)", + "var(--color-cyan-500, #06b6d4)", + "var(--color-violet-500, #8b5cf6)", + ]); + }); + it("stacks bars into a single column per index when stack is set", () => { const { container } = render( , diff --git a/ui/litellm-dashboard/src/components/shared/charts/bar_chart.tsx b/ui/litellm-dashboard/src/components/shared/charts/bar_chart.tsx index 7069ececb70..ab2cc66eaf4 100644 --- a/ui/litellm-dashboard/src/components/shared/charts/bar_chart.tsx +++ b/ui/litellm-dashboard/src/components/shared/charts/bar_chart.tsx @@ -1,7 +1,7 @@ "use client"; import * as React from "react"; -import { Bar, BarChart as RechartsBarChart, CartesianGrid, XAxis, YAxis } from "recharts"; +import { Bar, BarChart as RechartsBarChart, CartesianGrid, Cell, XAxis, YAxis } from "recharts"; import { ChartContainer, ChartLegend, ChartLegendContent, ChartTooltip, type ChartConfig } from "@/components/ui/chart"; import { cn } from "@/lib/cva.config"; import { ValueTooltip, type ChartTooltipComponent } from "./chart_tooltip"; @@ -12,6 +12,8 @@ export type BarChartProps> = { index: string; categories: readonly string[]; colors?: readonly ChartColor[]; + colorByDatum?: boolean; + maxBarSize?: number; valueFormatter?: (value: number) => string; stack?: boolean; layout?: "horizontal" | "vertical"; @@ -32,6 +34,8 @@ export function BarChart>({ index, categories, colors, + colorByDatum = false, + maxBarSize, valueFormatter, stack = false, layout = "horizontal", @@ -57,7 +61,7 @@ export function BarChart>({ ); } - const fills = categoryFills(categories.length, colors); + const fills = categoryFills(colorByDatum ? data.length : categories.length, colors); const config: ChartConfig = Object.fromEntries(categories.map((category) => [category, { label: category }])); const vertical = layout === "vertical"; const TooltipContent = customTooltip ?? ValueTooltip; @@ -115,6 +119,7 @@ export function BarChart>({ fill={fills[i]} stackId={stack ? "stack" : undefined} isAnimationActive={false} + maxBarSize={maxBarSize} onClick={ onValueChange ? (item: { payload?: TDatum }) => { @@ -122,7 +127,9 @@ export function BarChart>({ } : undefined } - /> + > + {colorByDatum && data.map((_, dataIndex) => )} + ))} diff --git a/ui/litellm-dashboard/src/components/shared/charts/chart_legend.test.tsx b/ui/litellm-dashboard/src/components/shared/charts/chart_legend.test.tsx index 889927aca43..28afe5faf9c 100644 --- a/ui/litellm-dashboard/src/components/shared/charts/chart_legend.test.tsx +++ b/ui/litellm-dashboard/src/components/shared/charts/chart_legend.test.tsx @@ -21,6 +21,14 @@ describe("CustomLegend", () => { expect(dots[1]?.getAttribute("style")).toContain("--color-green-500"); }); + it("wraps onto multiple lines instead of overflowing when there are many categories", () => { + const { container } = render( + `metrics.tool_${i}`)} colors={["blue", "green"]} />, + ); + + expect(container.firstElementChild?.className).toContain("flex-wrap"); + }); + it("cycles colors when there are more categories than colors", () => { const { container } = render( , diff --git a/ui/litellm-dashboard/src/components/shared/charts/chart_legend.tsx b/ui/litellm-dashboard/src/components/shared/charts/chart_legend.tsx index da252d8bf63..1551f3d0e39 100644 --- a/ui/litellm-dashboard/src/components/shared/charts/chart_legend.tsx +++ b/ui/litellm-dashboard/src/components/shared/charts/chart_legend.tsx @@ -11,7 +11,7 @@ export const CustomLegend = ({ categories: readonly string[]; colors: readonly ChartColor[]; }) => ( -
+
{categories.map((category, idx) => (
`var(--color-${color}-500, ${CHART_COLOR_HEX[color]})`; +export const SEQUENTIAL_COLOR_RAMP: readonly ChartColor[] = [ + "#1e3a8a", + "#1d4ed8", + "#2563eb", + "#3b82f6", + "#60a5fa", + "#93c5fd", + "#bfdbfe", + "#dbeafe", +]; + +const NAMED_COLOR_HEX: Readonly> = CHART_COLOR_HEX; + +export const chartColorValue = (color: ChartColor): string => + color in NAMED_COLOR_HEX ? `var(--color-${color}-500, ${NAMED_COLOR_HEX[color]})` : color; export const categoryFills = (count: number, colors?: readonly ChartColor[]): readonly string[] => { const cycle = colors && colors.length > 0 ? colors : DEFAULT_COLOR_CYCLE; diff --git a/ui/litellm-dashboard/src/components/shared/charts/index.ts b/ui/litellm-dashboard/src/components/shared/charts/index.ts index 8383c767064..69edd3fb13f 100644 --- a/ui/litellm-dashboard/src/components/shared/charts/index.ts +++ b/ui/litellm-dashboard/src/components/shared/charts/index.ts @@ -8,6 +8,13 @@ export { type ChartTooltipComponent, type ChartTooltipProps, } from "./chart_tooltip"; -export { CHART_COLOR_HEX, DEFAULT_COLOR_CYCLE, categoryFills, chartColorValue, type ChartColor } from "./colors"; +export { + CHART_COLOR_HEX, + DEFAULT_COLOR_CYCLE, + SEQUENTIAL_COLOR_RAMP, + categoryFills, + chartColorValue, + type ChartColor, +} from "./colors"; export { DonutChart, type DonutChartProps } from "./donut_chart"; export { LineChart, type LineChartCurveType, type LineChartProps } from "./line_chart"; diff --git a/ui/litellm-dashboard/src/components/ui/chart.tsx b/ui/litellm-dashboard/src/components/ui/chart.tsx index 14e10b9f06f..f0310819097 100644 --- a/ui/litellm-dashboard/src/components/ui/chart.tsx +++ b/ui/litellm-dashboard/src/components/ui/chart.tsx @@ -264,7 +264,11 @@ const ChartLegendContent = React.forwardRef< return (
{payload .filter((item) => item.type !== "none") diff --git a/ui/litellm-dashboard/src/lib/http/client.test.ts b/ui/litellm-dashboard/src/lib/http/client.test.ts index e0b5a73d11a..ac1bc86f871 100644 --- a/ui/litellm-dashboard/src/lib/http/client.test.ts +++ b/ui/litellm-dashboard/src/lib/http/client.test.ts @@ -119,4 +119,10 @@ describe("deriveErrorMessage", () => { it("falls back to a string detail field", () => { expect(deriveErrorMessage({ detail: "detail text" })).toBe("detail text"); }); + + it("unwraps the HTTPException detail.error shape management endpoints raise", () => { + expect(deriveErrorMessage({ detail: { error: "Team(s) not found: ghost-team" } })).toBe( + "Team(s) not found: ghost-team", + ); + }); }); diff --git a/ui/litellm-dashboard/src/lib/http/client.ts b/ui/litellm-dashboard/src/lib/http/client.ts index e2b47f7b354..1370d3e9273 100644 --- a/ui/litellm-dashboard/src/lib/http/client.ts +++ b/ui/litellm-dashboard/src/lib/http/client.ts @@ -44,13 +44,15 @@ export class ApiError extends Error { * Lives here because error parsing is the client's job; networking.tsx re-exports * it so existing `@/components/networking` import paths keep working. */ +const deriveDetailMessage = (detail: any): string | undefined => { + if (Array.isArray(detail)) return detail.map((d: any) => d?.msg || JSON.stringify(d)).join("; "); + if (typeof detail === "string") return detail; + if (typeof detail?.error === "string") return detail.error; + return undefined; +}; + export const deriveErrorMessage = (errorData: any): string => { - const detail = errorData?.detail; - const detailStr = Array.isArray(detail) - ? detail.map((d: any) => d?.msg || JSON.stringify(d)).join("; ") - : typeof detail === "string" - ? detail - : undefined; + const detailStr = deriveDetailMessage(errorData?.detail); return ( (errorData?.error && (errorData.error.message || (typeof errorData.error === "string" ? errorData.error : undefined))) || diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index da447f8c19b..01d79f5b898 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -18023,16 +18023,16 @@ export interface paths { * Get Tool Spend * @description Spend attributed to each tool over a date range, for the Cost Optimization dashboard. * - * Joins ``LiteLLM_SpendLogToolIndex`` (which tool names ran on which request) to - * ``LiteLLM_SpendLogs`` (what the request cost). A request that used multiple tools - * counts its full spend toward each of those tools, so per-tool numbers are - * attributions. ``total_spend`` is the deduplicated spend of every request that - * called at least one tool in the window, so it never double counts. + * Reads the ``LiteLLM_DailyToolSpend`` rollup, written at request time from invoked + * tools only (MCP tool calls and response tool_calls; declaring a tool without + * invoking it does not count). A request that invoked multiple tools counts its + * full spend toward each of them, so per-tool numbers are attributions and do not + * sum to a deduplicated total. * - * ``start_date`` is clamped to at most 30 days before ``end_date`` (serving up to - * 31 calendar dates inclusive, the same width as the endpoint's default window): - * a wider requested range is clamped, and the response's ``start_date`` reflects - * the effective window actually served. + * ``by_tool`` is the top ``TOOL_SPEND_TOP_TOOLS`` tools by spend, aggregated in + * SQL, and ``daily`` covers only those tools, so the response is bounded by + * days x TOOL_SPEND_TOP_TOOLS regardless of the requested range or how many + * distinct tool names exist. */ get: operations["get_tool_spend_v1_tool_spend_get"]; put?: never; @@ -18092,7 +18092,8 @@ export interface paths { }; /** * Get Tool Usage Logs - * @description Return paginated spend logs for requests that used this tool (from SpendLogToolIndex). + * @description Return paginated spend logs for requests that invoked this tool (from SpendLogToolIndex). + * Declaring a tool in a request body without the model invoking it does not create an entry. */ get: operations["get_tool_usage_logs_v1_tool__tool_name__logs_get"]; put?: never; @@ -32261,12 +32262,6 @@ export interface components { end_date?: string | null; /** Start Date */ start_date?: string | null; - /** - * Total Spend - * @description Deduplicated spend of every request that called at least one tool in the window; less than the sum of per-tool attributed spend whenever multi-tool requests exist - * @default 0 - */ - total_spend: number; }; /** * ToolUsageLogEntry