scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-15 23:31:29 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 165

merge: resolve conflict with litellm_internal_staging (ruff budget ceilings)

Browse source
This commit is contained in:
Yucheng Zhu 2026-07-27 16:51:24 -07:00
commit 1042b56d2f
121 changed files with 7164 additions and 3239 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
.github/workflows/test-linting.yml vendored
View file

@ -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"

24
basedpyright-code-budget.json
View file

@ -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

46
db_scripts/backfill_daily_tool_spend.sql Normal file
View file

@ -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;

191
enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py
View file

@ -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}},

12
litellm-proxy-extras/litellm_proxy_extras/migrations/20260725000000_add_daily_tool_spend/migration.sql Normal file
View file

@ -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")
);

13
litellm-proxy-extras/litellm_proxy_extras/schema.prisma
View file

@ -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())

39
litellm/completion_extras/litellm_responses_transformation/handler.py
View file

@ -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",

12
litellm/completion_extras/litellm_responses_transformation/transformation.py
View file

@ -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

2
litellm/constants.py
View file

@ -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))

20
litellm/integrations/langsmith.py
View file

@ -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"],

62
litellm/integrations/prometheus.py
View file

@ -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,

12
litellm/litellm_core_utils/litellm_logging.py
View file

@ -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

32
litellm/litellm_core_utils/prompt_templates/factory.py
View file

@ -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,
):

157
litellm/llms/azure/batches/handler.py
View file

@ -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(

8
litellm/llms/azure_ai/anthropic/count_tokens/transformation.py
View file

@ -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)

88
litellm/llms/openai/evals/transformation.py
View file

@ -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)

185
litellm/llms/volcengine/responses/transformation.py
View file

@ -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:

35
litellm/ocr/main.py
View file

@ -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.

467
litellm/proxy/_experimental/mcp_server/db.py
View file

File diff suppressed because it is too large Load diff

4
litellm/proxy/_experimental/mcp_server/outbound_credentials/per_user_oauth_store.py
View file

@ -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,
)

6
litellm/proxy/_experimental/mcp_server/outbound_credentials/v2_token_store.py
View file

@ -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

102
litellm/proxy/_experimental/mcp_server/rest_endpoints.py
View file

@ -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 <token>"} 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,

580
litellm/proxy/_experimental/mcp_server/server.py
View file

File diff suppressed because it is too large Load diff

10
litellm/proxy/_lazy_openapi_snapshot.json
View file

@ -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": [
{

7
litellm/proxy/_types.py
View file

@ -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]

127
litellm/proxy/agent_endpoints/agent_registry.py
View file

@ -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
"""

63
litellm/proxy/agent_endpoints/endpoints.py
View file

@ -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,

11
litellm/proxy/common_utils/callback_utils.py
View file

@ -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.

71
litellm/proxy/db/db_spend_update_writer.py
View file

@ -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,

264
litellm/proxy/db/spend_log_tool_index.py
View file

@ -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))

62
litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py
View file

@ -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,

251
litellm/proxy/guardrails/usage_endpoints.py
View file

@ -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"),

7
litellm/proxy/litellm_pre_call_utils.py
View file

@ -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
)

287
litellm/proxy/management_endpoints/common_daily_activity.py
View file

@ -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.

225
litellm/proxy/management_endpoints/internal_user_endpoints.py
View file

@ -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.",

133
litellm/proxy/management_endpoints/mcp_management_endpoints.py
View file

@ -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:

30
litellm/proxy/management_endpoints/model_management_endpoints.py
View file

@ -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:

347
litellm/proxy/management_endpoints/organization_endpoints.py
View file

@ -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:

155
litellm/proxy/management_endpoints/tag_management_endpoints.py
View file

@ -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),

170
litellm/proxy/management_endpoints/tool_management_endpoints.py
View file

@ -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

200
litellm/proxy/policy_engine/policy_registry.py
View file

@ -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:

3
litellm/proxy/proxy_server.py
View file

@ -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

13
litellm/proxy/schema.prisma
View file

@ -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())

50
litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py
View file

@ -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,

45
litellm/proxy/utils.py
View file

@ -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:

2
litellm/repositories/__init__.py
View file

@ -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",

4
litellm/repositories/table_repositories.py
View file

@ -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"

220
litellm/repositories/verification_token_repository.py
View file

@ -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")

109
litellm/router.py
View file

@ -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

7
litellm/types/tool_management.py
View file

@ -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

24
ruff-strict-budget.json
View file

@ -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
}
}

13
schema.prisma
View file

@ -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())

39
tests/e2e/a2a/a2a_client.py
View file

@ -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(

3
tests/e2e/coverage_registry/llm_conversational.yaml
View file

@ -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"}

7
tests/e2e/e2e_http.py
View file

@ -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="<streamed>",
chunks=chunks,
stream_events=stream_events,
stream_done=stream_done,
stream_error=stream_error,
)

22
tests/e2e/guardrails/guardrails_client.py
View file

@ -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

6
tests/e2e/guardrails/test_bedrock_guardrail_e2e.py
View file

@ -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):

15
tests/e2e/guardrails/test_block_code_execution_guardrail_e2e.py
View file

@ -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(), (

10
tests/e2e/guardrails/test_openai_moderation_guardrail_e2e.py
View file

@ -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(), (

147
tests/e2e/guardrails/test_presidio_guardrail_e2e.py
View file

@ -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 `<ENTITY_TYPE>` placeholders (e.g.
`<EMAIL_ADDRESS>`) via a real analyzer + anonymizer. Three modes are checked
`<EMAIL_ADDRESS>`) 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 = "<EMAIL_ADDRESS>"
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}"
)

36
tests/e2e/llm_translation/test_chat_completions_regression_e2e.py
View file

@ -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"],

12
tests/e2e/llm_translation/test_passthrough_e2e.py
View file

@ -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.

32
tests/e2e/llm_translation/test_passthrough_headers_e2e.py
View file

@ -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",

173
tests/e2e/llm_translation/test_responses_bridge_streaming_e2e.py Normal file
View file

@ -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-<uuid>``
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}"

122
tests/e2e/llm_translation/test_responses_metadata_e2e.py
View file

@ -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}"
)

28
tests/e2e/mcp/mcp_client.py
View file

@ -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

7
tests/e2e/mcp/test_mcp_datadog_e2e.py
View file

@ -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(

9
tests/e2e/mcp/test_mcp_guardrail_e2e.py
View file

@ -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 = {

14
tests/e2e/mcp/test_mcp_key_access_e2e.py
View file

@ -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",

4
tests/proxy_unit_tests/test_update_spend.py
View file

@ -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

62
tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_handler.py
View file

@ -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}"
)

35
tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py
View file

@ -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-<uuid>`` 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(

87
tests/test_litellm/integrations/test_langsmith_init.py
View file

@ -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"

152
tests/test_litellm/integrations/test_prometheus_cache_metrics.py
View file

@ -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

51
tests/test_litellm/integrations/test_prometheus_token_detail_metrics.py
View file

@ -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 = {

31
tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py
View file

@ -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"]

173
tests/test_litellm/llms/openai/evals/test_openai_evals_transformation.py
View file

@ -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

248
tests/test_litellm/llms/volcengine/responses/test_volcengine_responses_transformation.py
View file

@ -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

39
tests/test_litellm/proxy/common_utils/test_callback_utils.py
View file

@ -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

234
tests/test_litellm/proxy/db/test_db_spend_update_writer.py
View file

@ -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():
"""

348
tests/test_litellm/proxy/db/test_spend_log_tool_index.py Normal file
View file

@ -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()

231
tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py
View file

@ -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]

62
tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py
View file

@ -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] == []

50
tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py
View file

@ -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):

220
tests/test_litellm/proxy/management_endpoints/test_tool_management_endpoints.py
View file

@ -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()

36
tests/test_litellm/proxy/test_litellm_pre_call_utils.py
View file

@ -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 {})

6
tests/test_litellm/proxy/test_spend_log_cleanup.py
View file

@ -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():

136
tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py
View file

@ -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

2
tests/test_litellm/proxy/utils/prisma_and_spend/conftest.py
View file

@ -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()

35
tests/test_litellm/proxy/utils/prisma_and_spend/test_proxy_update_spend.py
View file

@ -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,

31
tests/test_litellm/proxy/utils/prisma_and_spend/test_spend_functions.py
View file

@ -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(

443
tests/test_litellm/test_router.py
View file

@ -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"]

8
type-discipline-budget.json
View file

@ -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
}
}

21
ui/litellm-dashboard/eslint-suppressions.json
View file

@ -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
}
}
}
}

4
ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CostOptimizationView.activity.test.tsx
View file

@ -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: () => <div />,
BarChart: () => <div />,
CustomLegend: () => <div />,
DEFAULT_COLOR_CYCLE: ["emerald"],
SEQUENTIAL_COLOR_RAMP: ["indigo"],
}));
vi.mock("@/app/(dashboard)/router-settings/_components/general_settings", () => ({

65
ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/UsageTab.test.tsx
View file

@ -23,18 +23,37 @@ vi.mock("@/components/shared/charts", () => ({
DonutChart: ({ data, label }: { data: unknown; label: string }) => (
<div data-testid="donut-chart" data-label={label} data-slices={JSON.stringify(data)} />
),
BarChart: ({ data, categories }: { data: unknown; categories: string[] }) => (
<div data-testid="bar-chart" data-categories={categories.join(",")} data-series={JSON.stringify(data)} />
BarChart: ({
data,
categories,
colors,
showLegend,
maxBarSize,
}: {
data: unknown;
categories: string[];
colors?: readonly string[];
showLegend?: boolean;
maxBarSize?: number;
}) => (
<div
data-testid="bar-chart"
data-categories={categories.join(",")}
data-colors={(colors ?? []).join(",")}
data-show-legend={String(showLegend ?? true)}
data-max-bar-size={maxBarSize === undefined ? "" : String(maxBarSize)}
data-series={JSON.stringify(data)}
/>
),
CustomLegend: ({ categories }: { categories: readonly string[] }) => (
<div data-testid="chart-legend">{categories.join(",")}</div>
),
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>): 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);
});
});

24
ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/UsageTab.tsx
View file

@ -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<UsageTabProps> = ({ 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<UsageTabProps> = ({ 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 (
<div className="w-full space-y-6">
@ -262,15 +260,10 @@ const UsageTab: React.FC<UsageTabProps> = ({ accessToken, activity }) => {
<CardHeader>
<CardTitle>Spend by tool</CardTitle>
<p className="text-sm text-muted-foreground">
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.
</p>
{toolSpendWindowClamped && (
<p className="text-xs text-muted-foreground">
Tool spend is capped at 30 days before the end of the selected range; showing spend since{" "}
{toolSpend?.start_date}.
</p>
)}
</CardHeader>
<CardContent>
{topTools.length === 0 ? (
@ -285,22 +278,27 @@ const UsageTab: React.FC<UsageTabProps> = ({ accessToken, activity }) => {
data={topToolsChart}
index="tool_name"
categories={["spend"]}
colors={["emerald"]}
colors={toolColors}
colorByDatum
layout="vertical"
yAxisWidth={140}
maxBarSize={64}
showLegend={false}
valueFormatter={usd}
/>
</div>
<div>
<p className="mb-2 text-sm font-medium text-muted-foreground">Daily spend by tool</p>
<CustomLegend categories={topToolNames} colors={toolColors} />
<BarChart
data={dailyToolSeries}
index="date"
categories={topToolNames}
colors={toolColors}
stack
maxBarSize={64}
valueFormatter={usd}
showLegend={false}
/>
</div>
</div>

153
ui/litellm-dashboard/src/app/(dashboard)/users/_components/DefaultUserSettings.test.tsx
View file

@ -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 }) => (
<select data-testid="budget-duration" value={value || ""} onChange={(e) => onChange(e.target.value || null)}>
<option value="">Select duration</option>
<option value="daily">Daily</option>
<option value="monthly">Monthly</option>
</select>
),
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(<DefaultUserSettings {...defaultProps} />);
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(<DefaultUserSettings {...defaultProps} />);
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(<DefaultUserSettings {...defaultProps} />);
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();
});
});

492
ui/litellm-dashboard/src/app/(dashboard)/users/_components/DefaultUserSettings.tsx
View file

@ -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<string, Record<string, string>> | null;
userID: string;
userRole: string;
}
interface TeamEntry {
team_id: string;
max_budget_in_team?: number;
user_role: "user" | "admin";
}
const DefaultUserSettings: React.FC<DefaultUserSettingsProps> = ({
accessToken,
possibleUIRoles,
userID,
userRole,
}) => {
const [loading, setLoading] = useState<boolean>(true);
const [settings, setSettings] = useState<any>(null);
const [isEditing, setIsEditing] = useState<boolean>(false);
const [editedValues, setEditedValues] = useState<any>({});
const [saving, setSaving] = useState<boolean>(false);
const [availableModels, setAvailableModels] = useState<string[]>([]);
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<string, any>,
);
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<string, any>) => ({
...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 (
<div className="space-y-3">
{normalizedTeams.map((team, index) => (
<div key={index} className="border rounded-lg p-4 bg-gray-50">
<div className="flex items-center justify-between mb-3">
<Text className="font-medium">Team {index + 1}</Text>
<Button size="small" danger icon={<DeleteOutlined />} onClick={() => removeTeam(index)}>
Remove
</Button>
</div>
<div className="grid grid-cols-1 md:grid-cols-3 gap-3">
<div>
<Text className="text-sm font-medium mb-1">Team ID</Text>
<TextInput
value={team.team_id}
onChange={(e) => updateTeam(index, "team_id", e.target.value)}
placeholder="Enter team ID"
/>
</div>
<div>
<Text className="text-sm font-medium mb-1">Max Budget in Team</Text>
<InputNumber
style={{ width: "100%" }}
value={team.max_budget_in_team}
onChange={(value) => updateTeam(index, "max_budget_in_team", value)}
placeholder="Optional"
min={0}
step={0.01}
precision={2}
/>
</div>
<div>
<Text className="text-sm font-medium mb-1">User Role</Text>
<Select
style={{ width: "100%" }}
value={team.user_role}
onChange={(value) => updateTeam(index, "user_role", value)}
>
<Option value="user">User</Option>
<Option value="admin">Admin</Option>
</Select>
</div>
</div>
</div>
))}
<Button icon={<PlusOutlined />} onClick={addTeam} className="w-full">
Add Team
</Button>
</div>
);
};
const renderEditableField = (key: string, property: any, value: any) => {
const type = property.type;
if (key === "teams") {
return <div className="mt-2">{renderTeamsEditor(editedValues[key] || [])}</div>;
} else if (key === "user_role" && possibleUIRoles) {
return (
<Select
style={{ width: "100%" }}
value={editedValues[key] || ""}
onChange={(value) => handleTextInputChange(key, value)}
className="mt-2"
>
{Object.entries(possibleUIRoles)
.filter(([role]) => role.includes("internal_user"))
.map(([role, { ui_label, description }]) => (
<Option key={role} value={role}>
<div className="flex items-center">
<span>{ui_label}</span>
<span className="ml-2 text-xs text-gray-500">{description}</span>
</div>
</Option>
))}
</Select>
);
} else if (key === "budget_duration") {
return (
<BudgetDurationDropdown
value={editedValues[key] || null}
onChange={(value) => handleTextInputChange(key, value)}
className="mt-2"
/>
);
} else if (type === "boolean") {
return (
<div className="mt-2">
<Switch checked={!!editedValues[key]} onChange={(checked) => handleTextInputChange(key, checked)} />
</div>
);
} else if (type === "array" && property.items?.enum) {
return (
<Select
mode="multiple"
style={{ width: "100%" }}
value={editedValues[key] || []}
onChange={(value) => handleTextInputChange(key, value)}
className="mt-2"
>
{property.items.enum.map((option: string) => (
<Option key={option} value={option}>
{option}
</Option>
))}
</Select>
);
} else if (key === "models") {
return (
<Select
mode="multiple"
style={{ width: "100%" }}
value={editedValues[key] || []}
onChange={(value) => handleTextInputChange(key, value)}
className="mt-2"
>
<Option key="no-default-models" value="no-default-models">
No Default Models
</Option>
<Option key="all-proxy-models" value="all-proxy-models">
All Proxy Models
</Option>
{availableModels.map((model: string) => (
<Option key={model} value={model}>
{getModelDisplayName(model)}
</Option>
))}
</Select>
);
} else if (type === "string" && property.enum) {
return (
<Select
style={{ width: "100%" }}
value={editedValues[key] || ""}
onChange={(value) => handleTextInputChange(key, value)}
className="mt-2"
>
{property.enum.map((option: string) => (
<Option key={option} value={option}>
{option}
</Option>
))}
</Select>
);
} else {
return (
<TextInput
value={editedValues[key] !== undefined ? String(editedValues[key]) : ""}
onChange={(e) => 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 <span className="text-gray-400">Not set</span>;
if (key === "teams" && Array.isArray(value)) {
if (value.length === 0) return <span className="text-gray-400">No teams assigned</span>;
const normalizedTeams = normalizeTeams(value);
return (
<div className="space-y-2 mt-1">
{normalizedTeams.map((team, index) => (
<div key={index} className="border rounded-lg p-3 bg-white">
<div className="grid grid-cols-1 md:grid-cols-3 gap-2 text-sm">
<div>
<span className="font-medium text-gray-600">Team ID:</span>
<p className="text-gray-900">{team.team_id || "Not specified"}</p>
</div>
<div>
<span className="font-medium text-gray-600">Max Budget:</span>
<p className="text-gray-900">
{team.max_budget_in_team !== undefined
? `$${formatNumberWithCommas(team.max_budget_in_team, 4)}`
: "No limit"}
</p>
</div>
<div>
<span className="font-medium text-gray-600">Role:</span>
<p className="text-gray-900 capitalize">{team.user_role}</p>
</div>
</div>
</div>
))}
</div>
);
}
if (key === "user_role" && possibleUIRoles && possibleUIRoles[value]) {
const { ui_label, description } = possibleUIRoles[value];
return (
<div>
<span className="font-medium">{ui_label}</span>
{description && <p className="text-xs text-gray-500 mt-1">{description}</p>}
</div>
);
}
if (key === "budget_duration") {
return <span>{getBudgetDurationLabel(value)}</span>;
}
if (typeof value === "boolean") {
return <span>{value ? "Enabled" : "Disabled"}</span>;
}
if (key === "models" && Array.isArray(value)) {
if (value.length === 0) return <span className="text-gray-400">None</span>;
return (
<div className="flex flex-wrap gap-2 mt-1">
{value.map((model, index) => (
<span key={index} className="px-2 py-1 bg-blue-100 rounded-sm text-xs">
{getModelDisplayName(model)}
</span>
))}
</div>
);
}
if (typeof value === "object") {
if (Array.isArray(value)) {
if (value.length === 0) return <span className="text-gray-400">None</span>;
return (
<div className="flex flex-wrap gap-2 mt-1">
{value.map((item, index) => (
<span key={index} className="px-2 py-1 bg-blue-100 rounded-sm text-xs">
{typeof item === "object" ? JSON.stringify(item) : String(item)}
</span>
))}
</div>
);
}
return (
<pre className="bg-gray-100 p-2 rounded-sm text-xs overflow-auto mt-1">{JSON.stringify(value, null, 2)}</pre>
);
}
return <span>{String(value)}</span>;
};
if (loading) {
return (
<div className="flex justify-center items-center h-64">
<Spin size="large" />
</div>
);
}
if (!settings) {
return (
<Card>
<Text>No settings available or you do not have permission to view them.</Text>
</Card>
);
}
// Dynamically render settings based on the schema
const renderSettings = () => {
const { values, field_schema } = settings;
if (!field_schema || !field_schema.properties) {
return <Text>No schema information available</Text>;
}
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 (
<div key={key} className="mb-6 pb-6 border-b border-gray-200 last:border-0">
<Text className="font-medium text-lg">{displayName}</Text>
<Paragraph className="text-sm text-gray-500 mt-1">
{property.description || "No description available"}
</Paragraph>
{isEditing ? (
<div className="mt-2">{renderEditableField(key, property, value)}</div>
) : (
<div className="mt-1 p-2 bg-gray-50 rounded-sm">{renderValue(key, value)}</div>
)}
</div>
);
});
};
return (
<Card>
<div className="flex justify-between items-center mb-4">
<Title>Default User Settings</Title>
{!loading &&
settings &&
(isEditing ? (
<div className="flex gap-2">
<Button
onClick={() => {
setIsEditing(false);
setEditedValues(settings.values || {});
}}
disabled={saving}
>
Cancel
</Button>
<Button type="primary" onClick={handleSaveSettings} loading={saving}>
Save Changes
</Button>
</div>
) : (
<Button type="primary" onClick={() => setIsEditing(true)}>
Edit Settings
</Button>
))}
</div>
{settings?.field_schema?.description && (
<Paragraph className="mb-4">{settings.field_schema.description}</Paragraph>
)}
<Divider />
<div className="mt-4 space-y-4">{renderSettings()}</div>
</Card>
);
};
export default DefaultUserSettings;

296
ui/litellm-dashboard/src/app/(dashboard)/users/_components/default-user-settings/DefaultUserSettingsForm.test.tsx Normal file
View file

@ -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<typeof import("@/components/ModelSelect/ModelSelect")>();
return {
MODEL_SENTINEL_OPTIONS: actual.MODEL_SENTINEL_OPTIONS,
ModelSelect: ({ onChange }: { onChange: (values: string[]) => void }) => (
<button type="button" onClick={() => onChange(["all-proxy-models"])}>
set-models
</button>
),
};
});
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<typeof vi.fn>;
updateSettings?: ReturnType<typeof vi.fn>;
}) => {
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(
<QueryClientProvider client={queryClient}>
<DefaultUserSettingsForm
possibleUIRoles={POSSIBLE_UI_ROLES}
fetchSettings={fetchSettings}
updateSettings={updateSettings}
/>
</QueryClientProvider>,
);
return { fetchSettings, updateSettings };
};
const saveButton = async () => await screen.findByRole("button", { name: "Save Changes" });
const enterEditMode = async (user: ReturnType<typeof userEvent.setup>) => {
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();
});
});

Some files were not shown because too many files have changed in this diff Show more