scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-08-28 05:25:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 36

feat(proxy): type Customer Management response_model for OpenAPI coverage (#31043)

Browse source 
* feat(proxy): type Customer Management response_model for OpenAPI coverage

Add response_model to the five remaining untyped /customer operations
(block, unblock, new, update, delete) so the generated OpenAPI schema
documents a concrete response body. new/update reuse the canonical
LiteLLM_EndUserTable (matching info/list); block, unblock, and delete
get small dedicated models in
litellm/types/proxy/management_endpoints/customer_endpoints.py.

Together with the already-typed info/list/daily-activity routes this
brings the Customer Management group to full response_model coverage.
Regression tests assert each public /customer/* route declares the
expected response_model and that /customer/new surfaces a typed schema
in app.openapi(), so dropping a response_model fails CI.

* fix(proxy): keep budget_id in typed customer responses

Address review feedback on the Customer Management response_model typing.
Greptile flagged that response_model=LiteLLM_EndUserTable on /customer/new
and /customer/update silently drops fields the raw Prisma model_dump()
echoed. Checking the schema, budget_id is the only such scalar column that
was missing from the Pydantic model (created_at/updated_at/tpm_limit do not
exist on litellm_endusertable), so add budget_id to LiteLLM_EndUserTable.

This restores budget_id on new/update and also fixes the pre-existing gap
where /customer/info and /customer/list (already typed) dropped it, which
the UI Customer type expects. A regression test pins budget_id surviving the
response_model filter on /customer/update.

Also document UnblockUsersResponse.blocked_users via a Field description: it
holds the users that remain blocked after the call. The key name predates
this PR and is kept to avoid a backwards-incompatible rename on a beta route.

* fix(proxy): keep nested budget fields in customer responses

response_model=LiteLLM_EndUserTable nests the budget as the narrow write
allowlist LiteLLM_BudgetTable, which silently drops the server-managed
fields the customer endpoints used to return (budget_reset_at, created_at).

Introduce CustomerResponse, a thin response model that nests
LiteLLM_BudgetTableFull (the repo's budget response model), and apply it on
/customer/new, /customer/update, /customer/info and /customer/list. list
also builds CustomerResponse so its budget isn't narrowed at construction
time. created_by/updated_at/updated_by remain omitted, matching how budgets
are returned elsewhere.

The shared LiteLLM_EndUserTable is left untouched: it's constructed in many
places that pass narrow budget instances, and pydantic v2 won't coerce a
budget instance into a wider nested model. Typing only at the response
boundary (where the handler hands FastAPI a dict) sidesteps that. A
regression test pins budget_reset_at + created_at through the filter and
asserts the internal audit fields stay out.

* test(proxy): add golden-master characterization tests for customer responses

Lock the exact JSON body each customer-object endpoint (info/list/new/update)
and delete return today, so the upcoming type-safety refactor of the handlers
is only allowed to land if it reproduces these byte for byte. Pins null-field
inclusion, the nested budget shape (server fields kept, audit fields dropped),
and object_permission reverse-relation stripping. Green against current code.

* refactor(proxy): make the customer response flow type-safe

Replace the untyped dict + bolt-on response_model pattern on the customer
object endpoints with explicit typed construction. A single mapper,
_to_customer_response, validates a DB row into CustomerResponse at one
Any -> typed seam; new/update/info/list now return it (or a list of it) and
carry real -> CustomerResponse / -> List[CustomerResponse] return
annotations, and delete returns DeleteCustomersResponse. basedpyright now
verifies the handlers' return shapes instead of a runtime filter doing it
silently.

This also deletes the four copy-pasted object_permission reverse-relation
cleanup loops: pydantic's extra=ignore drops those undeclared fields during
validation, so the loops were dead code (proven by the golden-master tests,
which stay byte-for-byte green). basedpyright errors on the file drop from
140 to 116, all from removed dict plumbing.

CustomerResponse stays a thin subclass of LiteLLM_EndUserTable so it inherits
the existing validators/config unchanged (behavior preservation); only the
nested budget type is widened.

* refactor(proxy): annotate customer response mapper param as BaseModel

Address review nit: the mapper's untyped `record` added an ANN001 violation.
The incoming rows are pydantic v2 models, so type the param as BaseModel
rather than object (object has no model_dump, which would just move the
problem to basedpyright). This clears the ANN001 and also drops three
basedpyright unknown-type violations the untyped param was adding.

* style(test): ruff format customer endpoint tests

* test(proxy): give customer budget test update mocks a valid model_dump

The type-safe response refactor validates the update result via
_to_customer_response (CustomerResponse.model_validate(record.model_dump())).
These budget tests mocked the end-user update to return a bare MagicMock,
so model_dump() yielded a MagicMock that fails validation. Give each update
mock a minimal valid dict; the tests assert on the prisma calls, not the body.

* chore(ui): regenerate API types from proxy OpenAPI spec

* fix(ui): make generated API types stable across Python versions

Python 3.13 strips a docstring's common leading indentation at compile
time while 3.12 keeps it, so app.openapi() emits differently-indented
description strings depending on the interpreter. The dashboard type
generator ran locally on 3.13 and in CI on 3.12, so schema.d.ts drifted
and the "Verify schema.d.ts matches the proxy OpenAPI spec" check failed

Normalize every description through inspect.cleandoc in the spec dump so
the output is identical regardless of interpreter, then regenerate
This commit is contained in:
ryan-crabbe-berri 2026-06-30 09:58:01 -07:00 • committed by GitHub
parent 0965a4d1f4
commit 3dce3daff6
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
7 changed files with 495 additions and 150 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
litellm/models/end_user.py
View file

@ -21,6 +21,7 @@ class LiteLLM_EndUserTable(LiteLLMPydanticObjectBase):
spend: float = 0.0
allowed_model_region: Optional[Literal["eu", "us"]] = None
default_model: Optional[str] = None
budget_id: Optional[str] = None
litellm_budget_table: Optional[LiteLLM_BudgetTable] = None
object_permission_id: Optional[str] = None
object_permission: Optional[LiteLLM_ObjectPermissionTable] = None

105
litellm/proxy/management_endpoints/customer_endpoints.py
View file

@ -15,6 +15,7 @@ from typing import List, Optional
import fastapi
from fastapi import APIRouter, Depends, HTTPException, Request
from pydantic import BaseModel
import litellm
from litellm._logging import verbose_proxy_logger
@ -32,10 +33,26 @@ from litellm.repositories.table_repositories import EndUserRepository
from litellm.types.proxy.management_endpoints.common_daily_activity import (
SpendAnalyticsPaginatedResponse,
)
from litellm.types.proxy.management_endpoints.customer_endpoints import (
BlockUsersResponse,
CustomerResponse,
DeleteCustomersResponse,
UnblockUsersResponse,
)
router = APIRouter()
def _to_customer_response(record: BaseModel) -> CustomerResponse:
"""Validate a raw end-user DB row into the typed customer response.
object_permission reverse relations and the budget's audit fields are
dropped here by the response model's field set, so callers need no manual
cleanup.
"""
return CustomerResponse.model_validate(record.model_dump())
@router.post(
"/end_user/block",
tags=["Customer Management"],
@ -46,6 +63,7 @@ router = APIRouter()
"/customer/block",
tags=["Customer Management"],
dependencies=[Depends(user_api_key_auth)],
response_model=BlockUsersResponse,
)
async def block_user(data: BlockUsers):
"""
@ -100,6 +118,7 @@ async def block_user(data: BlockUsers):
"/customer/unblock",
tags=["Customer Management"],
dependencies=[Depends(user_api_key_auth)],
response_model=UnblockUsersResponse,
)
async def unblock_user(data: BlockUsers):
"""
@ -213,11 +232,12 @@ async def _handle_customer_object_permission_update(
"/customer/new",
tags=["Customer Management"],
dependencies=[Depends(user_api_key_auth)],
response_model=CustomerResponse,
)
async def new_end_user(
data: NewCustomerRequest,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
) -> CustomerResponse:
"""
Allow creating a new Customer
@ -370,20 +390,7 @@ async def new_end_user(
include={"litellm_budget_table": True, "object_permission": True},
)
# Convert to dict and clean up recursive fields
response_dict = end_user_record.model_dump()
if response_dict.get("object_permission"):
# Remove reverse relations from object_permission
for field in [
"teams",
"verification_tokens",
"organizations",
"users",
"end_users",
]:
response_dict["object_permission"].pop(field, None)
return response_dict
return _to_customer_response(end_user_record)
except Exception as e:
verbose_proxy_logger.exception(
"litellm.proxy.management_endpoints.customer_endpoints.new_end_user(): Exception occured - {}".format(
@ -404,7 +411,7 @@ async def new_end_user(
"/customer/info",
tags=["Customer Management"],
dependencies=[Depends(user_api_key_auth)],
response_model=LiteLLM_EndUserTable,
response_model=CustomerResponse,
)
@router.get(
"/end_user/info",
@ -414,7 +421,7 @@ async def new_end_user(
)
async def end_user_info(
end_user_id: str = fastapi.Query(description="End User ID in the request parameters"),
):
) -> CustomerResponse:
"""
Get information about an end-user. An `end_user` is a customer (external user) of the proxy.
@ -449,20 +456,7 @@ async def end_user_info(
param="end_user_id",
)
# Convert to dict and clean up recursive fields
response_dict = user_info.model_dump(exclude_none=True)
if response_dict.get("object_permission"):
# Remove reverse relations from object_permission
for field in [
"teams",
"verification_tokens",
"organizations",
"users",
"end_users",
]:
response_dict["object_permission"].pop(field, None)
return response_dict
return _to_customer_response(user_info)
except Exception as e:
verbose_proxy_logger.exception(
@ -477,6 +471,7 @@ async def end_user_info(
"/customer/update",
tags=["Customer Management"],
dependencies=[Depends(user_api_key_auth)],
response_model=CustomerResponse,
)
@router.post(
"/end_user/update",
@ -487,7 +482,7 @@ async def end_user_info(
async def update_end_user(
data: UpdateCustomerRequest,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
) -> CustomerResponse:
"""
Example curl
@ -641,20 +636,7 @@ async def update_end_user(
raise ValueError(f"Failed updating customer data. User ID does not exist passed user_id={data.user_id}")
verbose_proxy_logger.debug(f"received response from updating prisma client. response={response}")
# Convert to dict and clean up recursive fields
response_dict = response.model_dump()
if response_dict.get("object_permission"):
# Remove reverse relations from object_permission
for field in [
"teams",
"verification_tokens",
"organizations",
"users",
"end_users",
]:
response_dict["object_permission"].pop(field, None)
return response_dict
return _to_customer_response(response)
else:
raise ValueError(f"user_id is required, passed user_id = {data.user_id}")
@ -671,6 +653,7 @@ async def update_end_user(
"/customer/delete",
tags=["Customer Management"],
dependencies=[Depends(user_api_key_auth)],
response_model=DeleteCustomersResponse,
)
@router.post(
"/end_user/delete",
@ -681,7 +664,7 @@ async def update_end_user(
async def delete_end_user(
data: DeleteCustomerRequest,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
) -> DeleteCustomersResponse:
"""
Delete multiple end-users.
@ -728,10 +711,10 @@ async def delete_end_user(
where={"user_id": {"in": data.user_ids}}
)
verbose_proxy_logger.debug(f"received response from updating prisma client. response={response}")
return {
"deleted_customers": response,
"message": "Successfully deleted customers with ids: " + str(data.user_ids),
}
return DeleteCustomersResponse(
deleted_customers=response,
message="Successfully deleted customers with ids: " + str(data.user_ids),
)
else:
raise ValueError(f"user_id is required, passed user_id = {data.user_ids}")
@ -747,7 +730,7 @@ async def delete_end_user(
"/customer/list",
tags=["Customer Management"],
dependencies=[Depends(user_api_key_auth)],
response_model=List[LiteLLM_EndUserTable],
response_model=List[CustomerResponse],
)
@router.get(
"/end_user/list",
@ -758,7 +741,7 @@ async def delete_end_user(
async def list_end_user(
http_request: Request,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
) -> List[CustomerResponse]:
"""
[Admin-only] List all available customers
@ -791,21 +774,7 @@ async def list_end_user(
include={"litellm_budget_table": True, "object_permission": True}
)
returned_response: List[LiteLLM_EndUserTable] = []
for item in response:
item_dict = item.model_dump()
# Remove reverse relations from object_permission
if item_dict.get("object_permission"):
for field in [
"teams",
"verification_tokens",
"organizations",
"users",
"end_users",
]:
item_dict["object_permission"].pop(field, None)
returned_response.append(LiteLLM_EndUserTable(**item_dict))
return returned_response
return [_to_customer_response(item) for item in response]
except Exception as e:
verbose_proxy_logger.exception(

30
litellm/types/proxy/management_endpoints/customer_endpoints.py Normal file
View file

@ -0,0 +1,30 @@
from typing import List, Optional
from pydantic import BaseModel, Field
from litellm.models.budget import LiteLLM_BudgetTableFull
from litellm.models.end_user import LiteLLM_EndUserTable
class CustomerResponse(LiteLLM_EndUserTable):
"""Customer object returned by the /customer read+write endpoints.
Nests the full budget response model so server-managed budget fields
(budget_reset_at, created_at) survive response_model filtering, rather than
the narrow write-allowlist shape LiteLLM_EndUserTable carries for internal use.
"""
litellm_budget_table: Optional[LiteLLM_BudgetTableFull] = None # pyright: ignore
class BlockUsersResponse(BaseModel):
blocked_users: List[LiteLLM_EndUserTable]
class UnblockUsersResponse(BaseModel):
blocked_users: List[str] = Field(description="User IDs that remain blocked after this unblock call")
class DeleteCustomersResponse(BaseModel):
deleted_customers: int
message: str

13
tests/test_litellm/proxy/management_endpoints/test_customer_budget.py
View file

@ -134,8 +134,10 @@ async def test_update_customer_creates_budget_with_proper_relations(
)
# Mock end user update
mock_updated_user = MagicMock()
mock_updated_user.model_dump.return_value = {"user_id": "test-user", "blocked": False}
mock_prisma_client.db.litellm_endusertable.update = AsyncMock(
return_value=MagicMock()
return_value=mock_updated_user
)
# Create update request with budget creation fields (not just budget_id)
@ -190,8 +192,10 @@ async def test_update_customer_creates_budget_with_required_fields(
)
# Mock end user update
mock_updated_user = MagicMock()
mock_updated_user.model_dump.return_value = {"user_id": "test-user", "blocked": False}
mock_prisma_client.db.litellm_endusertable.update = AsyncMock(
return_value=MagicMock()
return_value=mock_updated_user
)
# Create update request with budget creation fields
@ -253,8 +257,10 @@ async def test_update_customer_budget_creation_with_fallback_admin(
)
# Mock end user update
mock_updated_user = MagicMock()
mock_updated_user.model_dump.return_value = {"user_id": "test-user", "blocked": False}
mock_prisma_client.db.litellm_endusertable.update = AsyncMock(
return_value=MagicMock()
return_value=mock_updated_user
)
# Create update request with budget creation fields
@ -309,6 +315,7 @@ async def test_update_customer_with_budget_id_and_creation_fields(
# Mock end user update
mock_updated_user = MagicMock()
mock_updated_user.model_dump.return_value = {"user_id": "test-user", "blocked": False}
mock_prisma_client.db.litellm_endusertable.update = AsyncMock(
return_value=mock_updated_user
)

380
tests/test_litellm/proxy/management_endpoints/test_customer_endpoints.py
View file

@ -1,18 +1,28 @@
from typing import List
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from fastapi import FastAPI, HTTPException, Request, status
from fastapi.responses import JSONResponse
from fastapi.routing import APIRoute
from fastapi.testclient import TestClient
from litellm.proxy._types import (
LiteLLM_BudgetTable,
LiteLLM_EndUserTable,
LitellmUserRoles,
ProxyException,
)
from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth, user_api_key_auth
from litellm.proxy.management_endpoints.customer_endpoints import router
from litellm.types.proxy.management_endpoints.common_daily_activity import (
SpendAnalyticsPaginatedResponse,
)
from litellm.types.proxy.management_endpoints.customer_endpoints import (
BlockUsersResponse,
CustomerResponse,
DeleteCustomersResponse,
UnblockUsersResponse,
)
app = FastAPI()
@ -22,9 +32,7 @@ async def openai_exception_handler(request: Request, exc: ProxyException):
headers = exc.headers
error_dict = exc.to_dict()
return JSONResponse(
status_code=(
int(exc.code) if exc.code else status.HTTP_500_INTERNAL_SERVER_ERROR
),
status_code=(int(exc.code) if exc.code else status.HTTP_500_INTERNAL_SERVER_ERROR),
content={"error": error_dict},
headers=headers,
)
@ -54,30 +62,20 @@ def mock_user_api_key_auth():
def test_update_customer_success(mock_prisma_client, mock_user_api_key_auth):
# Mock the database responses
mock_end_user = LiteLLM_EndUserTable(
user_id="test-user-1", alias="Test User", blocked=False
)
updated_mock_end_user = LiteLLM_EndUserTable(
user_id="test-user-1", alias="Updated Test User", blocked=False
)
mock_end_user = LiteLLM_EndUserTable(user_id="test-user-1", alias="Test User", blocked=False)
updated_mock_end_user = LiteLLM_EndUserTable(user_id="test-user-1", alias="Updated Test User", blocked=False)
# Mock the find_first response
mock_prisma_client.db.litellm_endusertable.find_first = AsyncMock(
return_value=mock_end_user
)
mock_prisma_client.db.litellm_endusertable.find_first = AsyncMock(return_value=mock_end_user)
# Mock the update response
mock_prisma_client.db.litellm_endusertable.update = AsyncMock(
return_value=updated_mock_end_user
)
mock_prisma_client.db.litellm_endusertable.update = AsyncMock(return_value=updated_mock_end_user)
# Test data
test_data = {"user_id": "test-user-1", "alias": "Updated Test User"}
# Make the request
response = client.post(
"/customer/update", json=test_data, headers={"Authorization": "Bearer test-key"}
)
response = client.post("/customer/update", json=test_data, headers={"Authorization": "Bearer test-key"})
# Assert response
assert response.status_code == 200
@ -106,10 +104,7 @@ def test_update_customer_not_found(mock_prisma_client, mock_user_api_key_auth):
assert response.status_code == 404
response_json = response.json()
assert "error" in response_json
assert (
response_json["error"]["message"]
== "End User Id=non-existent-user does not exist in db"
)
assert response_json["error"]["message"] == "End User Id=non-existent-user does not exist in db"
assert response_json["error"]["type"] == "not_found"
assert response_json["error"]["param"] == "user_id"
assert response_json["error"]["code"] == "404"
@ -132,10 +127,7 @@ def test_info_customer_not_found(mock_prisma_client, mock_user_api_key_auth):
assert response.status_code == 404
response_json = response.json()
assert "error" in response_json
assert (
response_json["error"]["message"]
== "End User Id=non-existent-user does not exist in db"
)
assert response_json["error"]["message"] == "End User Id=non-existent-user does not exist in db"
assert response_json["error"]["type"] == "not_found"
assert response_json["error"]["param"] == "end_user_id"
assert response_json["error"]["code"] == "404"
@ -220,11 +212,6 @@ def test_error_schema_consistency(mock_prisma_client, mock_user_api_key_auth):
assert error["code"] == "404"
# Test /customer/new - duplicate user error
from unittest.mock import MagicMock
mock_end_user = LiteLLM_EndUserTable(
user_id="existing-user", alias="Existing User", blocked=False
)
mock_prisma_client.db.litellm_endusertable.create = AsyncMock(
side_effect=Exception("Unique constraint failed on the fields: (`user_id`)")
)
@ -238,9 +225,7 @@ def test_error_schema_consistency(mock_prisma_client, mock_user_api_key_auth):
assert error["code"] == "400"
def test_customer_endpoints_error_schema_consistency(
mock_prisma_client, mock_user_api_key_auth
):
def test_customer_endpoints_error_schema_consistency(mock_prisma_client, mock_user_api_key_auth):
"""
Test the exact scenarios from the curl examples provided.
@ -307,9 +292,7 @@ def test_customer_endpoints_error_schema_consistency(
assert "Customer already exists" in error2["message"]
# Verify both errors have the same schema structure
assert set(error1.keys()) == set(
error2.keys()
), "Both errors should have the same top-level keys"
assert set(error1.keys()) == set(error2.keys()), "Both errors should have the same top-level keys"
# Both should have string values for all fields
for key in ["message", "type", "code"]:
@ -317,6 +300,153 @@ def test_customer_endpoints_error_schema_consistency(
assert isinstance(error2[key], str), f"error2[{key}] should be a string"
EXPECTED_RESPONSE_MODELS = {
"/customer/block": BlockUsersResponse,
"/customer/unblock": UnblockUsersResponse,
"/customer/new": CustomerResponse,
"/customer/update": CustomerResponse,
"/customer/delete": DeleteCustomersResponse,
"/customer/info": CustomerResponse,
"/customer/list": List[CustomerResponse],
"/customer/daily/activity": SpendAnalyticsPaginatedResponse,
}
@pytest.mark.parametrize("path, expected_model", EXPECTED_RESPONSE_MODELS.items())
def test_customer_routes_declare_response_model(path, expected_model):
"""
Every public /customer/* operation must declare a typed response_model so
the generated OpenAPI schema documents the response body. Regression for the
OpenAPI response-type coverage goal: drop a response_model and this fails.
"""
route = next(r for r in router.routes if isinstance(r, APIRoute) and r.path == path)
assert route.response_model == expected_model
def test_customer_new_documented_in_openapi_schema():
"""
The response_model must surface in the OpenAPI schema as a concrete ref, not
an empty/default response. This is what the coverage metric measures.
"""
schema = app.openapi()["paths"]["/customer/new"]["post"]
json_schema = schema["responses"]["200"]["content"]["application/json"]["schema"]
assert json_schema["$ref"].endswith("/CustomerResponse")
def test_update_customer_response_preserves_budget_id(mock_prisma_client, mock_user_api_key_auth):
"""
Regression for the response_model field-stripping concern: budget_id is a real
column on the end-user table that /customer/update echoes. response_model=
LiteLLM_EndUserTable must NOT drop it, so budget_id stays in LiteLLM_EndUserTable.
"""
existing = LiteLLM_EndUserTable(user_id="cust-1", blocked=False)
updated = LiteLLM_EndUserTable(user_id="cust-1", blocked=False, budget_id="budget-123")
mock_prisma_client.db.litellm_endusertable.find_first = AsyncMock(return_value=existing)
mock_prisma_client.db.litellm_endusertable.update = AsyncMock(return_value=updated)
response = client.post(
"/customer/update",
json={"user_id": "cust-1", "budget_id": "budget-123"},
headers={"Authorization": "Bearer test-key"},
)
assert response.status_code == 200
assert response.json()["budget_id"] == "budget-123"
def test_update_customer_response_keeps_nested_budget_server_fields(mock_prisma_client, mock_user_api_key_auth):
"""
Faithfulness regression: /customer/update embeds the full budget row. The
response_model must keep the server-managed budget fields the endpoint used
to return (budget_reset_at, created_at) instead of the narrow write-allowlist
shape. The intentionally-internal audit fields (created_by/updated_by) stay out.
"""
existing = LiteLLM_EndUserTable(user_id="cust-1", blocked=False)
raw_row = MagicMock()
raw_row.model_dump.return_value = {
"user_id": "cust-1",
"blocked": False,
"alias": "renamed",
"spend": 0.0,
"allowed_model_region": None,
"default_model": None,
"budget_id": "b-1",
"object_permission_id": None,
"object_permission": None,
"litellm_budget_table": {
"budget_id": "b-1",
"max_budget": 10.0,
"budget_duration": "30d",
"budget_reset_at": "2024-02-01T00:00:00",
"created_at": "2024-01-01T00:00:00",
"created_by": "admin",
"updated_at": "2024-01-02T00:00:00",
"updated_by": "admin",
},
}
mock_prisma_client.db.litellm_endusertable.find_first = AsyncMock(return_value=existing)
mock_prisma_client.db.litellm_endusertable.update = AsyncMock(return_value=raw_row)
response = client.post(
"/customer/update",
json={"user_id": "cust-1", "alias": "renamed"},
headers={"Authorization": "Bearer test-key"},
)
assert response.status_code == 200
budget = response.json()["litellm_budget_table"]
assert budget["budget_reset_at"] == "2024-02-01T00:00:00"
assert budget["created_at"] == "2024-01-01T00:00:00"
assert "created_by" not in budget
assert "updated_by" not in budget
def test_block_customer_success_serializes_through_response_model(mock_prisma_client, mock_user_api_key_auth):
"""
/customer/block returns {"blocked_users": [<end user rows>]}. With
response_model=BlockUsersResponse, a shape mismatch would raise a 500
ResponseValidationError, so a clean 200 proves the model matches runtime output.
"""
blocked_row = LiteLLM_EndUserTable(user_id="blocked-1", blocked=True)
mock_prisma_client.db.litellm_endusertable.upsert = AsyncMock(return_value=blocked_row)
response = client.post(
"/customer/block",
json={"user_ids": ["blocked-1"]},
headers={"Authorization": "Bearer test-key"},
)
assert response.status_code == 200
body = response.json()
assert body["blocked_users"][0]["user_id"] == "blocked-1"
assert body["blocked_users"][0]["blocked"] is True
def test_delete_customer_success_serializes_through_response_model(mock_prisma_client, mock_user_api_key_auth):
"""
/customer/delete returns {"deleted_customers": <int>, "message": <str>}.
response_model=DeleteCustomersResponse enforces that exact shape.
"""
existing = [
LiteLLM_EndUserTable(user_id="u1", blocked=False),
LiteLLM_EndUserTable(user_id="u2", blocked=False),
]
mock_prisma_client.db.litellm_endusertable.find_many = AsyncMock(return_value=existing)
mock_prisma_client.db.litellm_endusertable.delete_many = AsyncMock(return_value=2)
response = client.post(
"/customer/delete",
json={"user_ids": ["u1", "u2"]},
headers={"Authorization": "Bearer test-key"},
)
assert response.status_code == 200
assert response.json() == {
"deleted_customers": 2,
"message": "Successfully deleted customers with ids: ['u1', 'u2']",
}
@pytest.mark.asyncio
async def test_get_customer_daily_activity_admin_param_passing(monkeypatch):
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
@ -331,9 +461,7 @@ async def test_get_customer_daily_activity_admin_param_passing(monkeypatch):
mocked_response = MagicMock(name="SpendAnalyticsPaginatedResponse")
get_daily_activity_mock = AsyncMock(return_value=mocked_response)
monkeypatch.setattr(
customer_endpoints, "get_daily_activity", get_daily_activity_mock
)
monkeypatch.setattr(customer_endpoints, "get_daily_activity", get_daily_activity_mock)
auth = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin1")
result = await get_customer_daily_activity(
@ -380,16 +508,12 @@ async def test_get_customer_daily_activity_with_end_user_aliases(monkeypatch):
mock_end_user2.user_id = "end-user-2"
mock_end_user2.alias = "Customer Two"
mock_prisma_client.db.litellm_endusertable.find_many = AsyncMock(
return_value=[mock_end_user1, mock_end_user2]
)
mock_prisma_client.db.litellm_endusertable.find_many = AsyncMock(return_value=[mock_end_user1, mock_end_user2])
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
mocked_response = MagicMock(name="SpendAnalyticsPaginatedResponse")
get_daily_activity_mock = AsyncMock(return_value=mocked_response)
monkeypatch.setattr(
customer_endpoints, "get_daily_activity", get_daily_activity_mock
)
monkeypatch.setattr(customer_endpoints, "get_daily_activity", get_daily_activity_mock)
auth = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin1")
await get_customer_daily_activity(
@ -436,9 +560,7 @@ async def test_get_customer_daily_activity_non_admin_is_rejected(monkeypatch):
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
get_daily_activity_mock = AsyncMock()
monkeypatch.setattr(
customer_endpoints, "get_daily_activity", get_daily_activity_mock
)
monkeypatch.setattr(customer_endpoints, "get_daily_activity", get_daily_activity_mock)
non_admin_key = UserAPIKeyAuth(
user_id="regular-user-abc",
@ -482,9 +604,7 @@ async def test_get_customer_daily_activity_service_account_key_is_rejected(monke
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
get_daily_activity_mock = AsyncMock()
monkeypatch.setattr(
customer_endpoints, "get_daily_activity", get_daily_activity_mock
)
monkeypatch.setattr(customer_endpoints, "get_daily_activity", get_daily_activity_mock)
service_account_key = UserAPIKeyAuth(
user_id=None,
@ -507,3 +627,157 @@ async def test_get_customer_daily_activity_service_account_key_is_rejected(monke
assert exc_info.value.status_code == 401
assert "Admin-only endpoint" in str(exc_info.value.detail)
get_daily_activity_mock.assert_not_called()
# ---------------------------------------------------------------------------
# Characterization (golden-master) tests.
#
# These lock the EXACT JSON body every customer-object endpoint returns today,
# so a type-safety refactor of the handlers is only allowed to land if it
# reproduces these byte for byte. The input below is what a Prisma row's
# .model_dump() yields (full nested budget incl. audit fields + object_permission
# incl. reverse relations); the expected output is what the live endpoint emits.
# ---------------------------------------------------------------------------
_FULL_DB_ROW = {
"user_id": "c1",
"blocked": False,
"alias": "Acme",
"spend": 1.5,
"allowed_model_region": None,
"default_model": None,
"budget_id": "b1",
"object_permission_id": "p1",
"litellm_budget_table": {
"budget_id": "b1",
"max_budget": 10.0,
"soft_budget": None,
"max_parallel_requests": None,
"tpm_limit": None,
"rpm_limit": None,
"model_max_budget": None,
"budget_duration": "30d",
"allowed_models": [],
"budget_reset_at": "2024-02-01T00:00:00",
"created_at": "2024-01-01T00:00:00",
"created_by": "admin",
"updated_at": "2024-01-02T00:00:00",
"updated_by": "admin",
},
"object_permission": {
"object_permission_id": "p1",
"mcp_servers": ["s1"],
"mcp_access_groups": [],
"mcp_tool_permissions": None,
"vector_stores": [],
"agents": [],
"agent_access_groups": [],
"models": [],
"mcp_toolsets": None,
"blocked_tools": [],
"search_tools": [],
"teams": [{"team_id": "t1"}],
"users": [{"user_id": "x"}],
"end_users": [],
"organizations": [],
"verification_tokens": [],
},
}
_EXPECTED_CUSTOMER = {
"user_id": "c1",
"blocked": False,
"alias": "Acme",
"spend": 1.5,
"allowed_model_region": None,
"default_model": None,
"budget_id": "b1",
"litellm_budget_table": {
"budget_id": "b1",
"soft_budget": None,
"max_budget": 10.0,
"max_parallel_requests": None,
"tpm_limit": None,
"rpm_limit": None,
"model_max_budget": None,
"budget_duration": "30d",
"allowed_models": [],
"budget_reset_at": "2024-02-01T00:00:00",
"created_at": "2024-01-01T00:00:00",
},
"object_permission_id": "p1",
"object_permission": {
"object_permission_id": "p1",
"mcp_servers": ["s1"],
"mcp_access_groups": [],
"mcp_tool_permissions": None,
"vector_stores": [],
"agents": [],
"agent_access_groups": [],
"models": [],
"mcp_toolsets": None,
"blocked_tools": [],
"search_tools": [],
},
}
def _row(dump: dict) -> MagicMock:
row = MagicMock()
row.model_dump.return_value = dump
return row
def test_char_info_body(mock_prisma_client, mock_user_api_key_auth):
mock_prisma_client.db.litellm_endusertable.find_first = AsyncMock(return_value=_row(_FULL_DB_ROW))
response = client.get("/customer/info?end_user_id=c1", headers={"Authorization": "Bearer k"})
assert response.status_code == 200
assert response.json() == _EXPECTED_CUSTOMER
def test_char_list_body(mock_prisma_client, mock_user_api_key_auth):
mock_prisma_client.db.litellm_endusertable.find_many = AsyncMock(return_value=[_row(_FULL_DB_ROW)])
response = client.get("/customer/list", headers={"Authorization": "Bearer k"})
assert response.status_code == 200
assert response.json() == [_EXPECTED_CUSTOMER]
def test_char_new_body(mock_prisma_client, mock_user_api_key_auth):
mock_prisma_client.db.litellm_endusertable.create = AsyncMock(return_value=_row(_FULL_DB_ROW))
response = client.post("/customer/new", json={"user_id": "c1"}, headers={"Authorization": "Bearer k"})
assert response.status_code == 200
assert response.json() == _EXPECTED_CUSTOMER
def test_char_update_body(mock_prisma_client, mock_user_api_key_auth):
mock_prisma_client.db.litellm_endusertable.find_first = AsyncMock(
return_value=_row({"user_id": "c1", "blocked": False})
)
mock_prisma_client.db.litellm_endusertable.update = AsyncMock(return_value=_row(_FULL_DB_ROW))
response = client.post(
"/customer/update",
json={"user_id": "c1", "alias": "Acme"},
headers={"Authorization": "Bearer k"},
)
assert response.status_code == 200
assert response.json() == _EXPECTED_CUSTOMER
def test_char_delete_body(mock_prisma_client, mock_user_api_key_auth):
mock_prisma_client.db.litellm_endusertable.find_many = AsyncMock(
return_value=[
LiteLLM_EndUserTable(user_id="c1", blocked=False),
LiteLLM_EndUserTable(user_id="c2", blocked=False),
]
)
mock_prisma_client.db.litellm_endusertable.delete_many = AsyncMock(return_value=2)
response = client.post(
"/customer/delete",
json={"user_ids": ["c1", "c2"]},
headers={"Authorization": "Bearer k"},
)
assert response.status_code == 200
assert response.json() == {
"deleted_customers": 2,
"message": "Successfully deleted customers with ids: ['c1', 'c2']",
}

15
ui/litellm-dashboard/scripts/gen-api-types.mjs
View file

@ -26,15 +26,26 @@ const python = (process.env.LITELLM_PYTHON ?? "python3").split(" ");
// The dashboard calls internal UI routes that the public /openapi.json hides via
// include_in_schema=False. Force them in so they get typed here; this mutates a
// throwaway interpreter, so the spec the proxy actually serves is unchanged.
// Python 3.13 strips a docstring's common leading indentation at compile time
// while 3.12 keeps it, so the same model yields differently-indented descriptions
// depending on the interpreter — enough to make this output non-reproducible
// across CI and contributors. inspect.cleandoc normalizes every description to one
// canonical form regardless of interpreter, so the generated file is stable.
const dumpSpec = [
"import json, sys",
"import inspect, json, sys",
"from litellm.proxy.proxy_server import app",
"from fastapi.routing import APIRoute",
"for route in app.routes:",
" if isinstance(route, APIRoute):",
" route.include_in_schema = True",
"app.openapi_schema = None",
"with open(sys.argv[1], 'w') as f: json.dump(app.openapi(), f, sort_keys=True)",
"def normalize(node):",
" if isinstance(node, dict):",
" return {k: inspect.cleandoc(v) if k == 'description' and isinstance(v, str) else normalize(v) for k, v in node.items()}",
" if isinstance(node, list):",
" return [normalize(v) for v in node]",
" return node",
"with open(sys.argv[1], 'w') as f: json.dump(normalize(app.openapi()), f, sort_keys=True)",
].join("\n");
try {

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

@ -3027,8 +3027,8 @@ export interface paths {
/**
* Get Active Tasks Stats
* @description Returns:
* total_active_tasks: int
* by_name: { coroutine_name: count }
* total_active_tasks: int
* by_name: { coroutine_name: count }
*/
get: operations["get_active_tasks_stats_debug_asyncio_tasks_get"];
put?: never;
@ -21003,6 +21003,11 @@ export interface components {
/** User Ids */
user_ids: string[];
};
/** BlockUsersResponse */
BlockUsersResponse: {
/** Blocked Users */
blocked_users: components["schemas"]["LiteLLM_EndUserTable"][];
};
/**
* BlockedWord
* @description Represents a blocked word with its action and optional description
@ -22651,10 +22656,10 @@ export interface components {
/**
* ContentFilterCategoryConfig
* @description category: "harmful_self_harm"
* enabled: true
* action: "BLOCK"
* severity_threshold: "medium"
* category_file: "/path/to/custom_file.yaml" # optional override
* enabled: true
* action: "BLOCK"
* severity_threshold: "medium"
* category_file: "/path/to/custom_file.yaml" # optional override
*/
ContentFilterCategoryConfig: {
/**
@ -22879,6 +22884,37 @@ export interface components {
[key: string]: unknown;
};
};
/**
* CustomerResponse
* @description Customer object returned by the /customer read+write endpoints.
*
* Nests the full budget response model so server-managed budget fields
* (budget_reset_at, created_at) survive response_model filtering, rather than
* the narrow write-allowlist shape LiteLLM_EndUserTable carries for internal use.
*/
CustomerResponse: {
/** Alias */
alias?: string | null;
/** Allowed Model Region */
allowed_model_region?: ("eu" | "us") | null;
/** Blocked */
blocked: boolean;
/** Budget Id */
budget_id?: string | null;
/** Default Model */
default_model?: string | null;
litellm_budget_table?: components["schemas"]["LiteLLM_BudgetTableFull"] | null;
object_permission?: components["schemas"]["LiteLLM_ObjectPermissionTable"] | null;
/** Object Permission Id */
object_permission_id?: string | null;
/**
* Spend
* @default 0
*/
spend: number;
/** User Id */
user_id: string;
};
/** DailySpendData */
DailySpendData: {
breakdown?: components["schemas"]["BreakdownMetrics"];
@ -23043,6 +23079,13 @@ export interface components {
/** User Ids */
user_ids: string[];
};
/** DeleteCustomersResponse */
DeleteCustomersResponse: {
/** Deleted Customers */
deleted_customers: number;
/** Message */
message: string;
};
/**
* DeleteEvalResponse
* @description Response from deleting an evaluation
@ -24792,6 +24835,8 @@ export interface components {
allowed_model_region?: ("eu" | "us") | null;
/** Blocked */
blocked: boolean;
/** Budget Id */
budget_id?: string | null;
/** Default Model */
default_model?: string | null;
litellm_budget_table?: components["schemas"]["LiteLLM_BudgetTable"] | null;
@ -31477,6 +31522,14 @@ export interface components {
*/
workers: components["schemas"]["WorkerRegistryEntry"][];
};
/** UnblockUsersResponse */
UnblockUsersResponse: {
/**
* Blocked Users
* @description User IDs that remain blocked after this unblock call
*/
blocked_users: string[];
};
/**
* UpdateCustomerRequest
* @description Update a Customer, use this to update customer budgets etc
@ -37482,7 +37535,7 @@ export interface operations {
[name: string]: unknown;
};
content: {
"application/json": unknown;
"application/json": components["schemas"]["BlockUsersResponse"];
};
};
/** @description Validation Error */
@ -37553,7 +37606,7 @@ export interface operations {
[name: string]: unknown;
};
content: {
"application/json": unknown;
"application/json": components["schemas"]["DeleteCustomersResponse"];
};
};
/** @description Validation Error */
@ -37585,7 +37638,7 @@ export interface operations {
[name: string]: unknown;
};
content: {
"application/json": components["schemas"]["LiteLLM_EndUserTable"];
"application/json": components["schemas"]["CustomerResponse"];
};
};
/** @description Validation Error */
@ -37614,7 +37667,7 @@ export interface operations {
[name: string]: unknown;
};
content: {
"application/json": components["schemas"]["LiteLLM_EndUserTable"][];
"application/json": components["schemas"]["CustomerResponse"][];
};
};
};
@ -37638,7 +37691,7 @@ export interface operations {
[name: string]: unknown;
};
content: {
"application/json": unknown;
"application/json": components["schemas"]["CustomerResponse"];
};
};
/** @description Validation Error */
@ -37671,7 +37724,7 @@ export interface operations {
[name: string]: unknown;
};
content: {
"application/json": unknown;
"application/json": components["schemas"]["UnblockUsersResponse"];
};
};
/** @description Validation Error */
@ -37704,7 +37757,7 @@ export interface operations {
[name: string]: unknown;
};
content: {
"application/json": unknown;
"application/json": components["schemas"]["CustomerResponse"];
};
};
/** @description Validation Error */
@ -38130,7 +38183,7 @@ export interface operations {
[name: string]: unknown;
};
content: {
"application/json": unknown;
"application/json": components["schemas"]["DeleteCustomersResponse"];
};
};
/** @description Validation Error */
@ -38162,7 +38215,7 @@ export interface operations {
[name: string]: unknown;
};
content: {
"application/json": unknown;
"application/json": components["schemas"]["CustomerResponse"];
};
};
/** @description Validation Error */
@ -38191,7 +38244,7 @@ export interface operations {
[name: string]: unknown;
};
content: {
"application/json": unknown;
"application/json": components["schemas"]["CustomerResponse"][];
};
};
};
@ -38215,7 +38268,7 @@ export interface operations {
[name: string]: unknown;
};
content: {
"application/json": unknown;
"application/json": components["schemas"]["CustomerResponse"];
};
};
/** @description Validation Error */
@ -38281,7 +38334,7 @@ export interface operations {
[name: string]: unknown;
};
content: {
"application/json": unknown;
"application/json": components["schemas"]["CustomerResponse"];
};
};
/** @description Validation Error */
@ -43890,13 +43943,13 @@ export interface operations {
/**
* @description Unified rate-limit error.
*
* Every rate-limit condition surfaced by litellm — whether it originated from
* an upstream LLM provider, a vendor batch endpoint, or one of litellm's own
* proxy-side limiters (parallel-requests, dynamic-rate, batch-rate, budget,
* max-iterations, etc.) — is raised as an instance of this class.
* Every rate-limit condition surfaced by litellm — whether it originated from
* an upstream LLM provider, a vendor batch endpoint, or one of litellm's own
* proxy-side limiters (parallel-requests, dynamic-rate, batch-rate, budget,
* max-iterations, etc.) — is raised as an instance of this class.
*
* The :attr:`category` attribute lets callers distinguish the source. See
* :class:`RateLimitErrorCategory` for the available values.
* The :attr:`category` attribute lets callers distinguish the source. See
* :class:`RateLimitErrorCategory` for the available values.
*/
429: {
headers: {