diff --git a/litellm/models/end_user.py b/litellm/models/end_user.py index 15fd03ec2ca..9bf895b9447 100644 --- a/litellm/models/end_user.py +++ b/litellm/models/end_user.py @@ -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 diff --git a/litellm/proxy/management_endpoints/customer_endpoints.py b/litellm/proxy/management_endpoints/customer_endpoints.py index 7c8a9b88191..84f67bdc3bc 100644 --- a/litellm/proxy/management_endpoints/customer_endpoints.py +++ b/litellm/proxy/management_endpoints/customer_endpoints.py @@ -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( diff --git a/litellm/types/proxy/management_endpoints/customer_endpoints.py b/litellm/types/proxy/management_endpoints/customer_endpoints.py new file mode 100644 index 00000000000..e7653360d63 --- /dev/null +++ b/litellm/types/proxy/management_endpoints/customer_endpoints.py @@ -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 diff --git a/tests/test_litellm/proxy/management_endpoints/test_customer_budget.py b/tests/test_litellm/proxy/management_endpoints/test_customer_budget.py index 41f43c75f7d..0beca0c15e8 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_customer_budget.py +++ b/tests/test_litellm/proxy/management_endpoints/test_customer_budget.py @@ -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 ) diff --git a/tests/test_litellm/proxy/management_endpoints/test_customer_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_customer_endpoints.py index 6c5ccd3562f..d4089b23e81 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_customer_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_customer_endpoints.py @@ -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": []}. 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": , "message": }. + 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']", + } diff --git a/ui/litellm-dashboard/scripts/gen-api-types.mjs b/ui/litellm-dashboard/scripts/gen-api-types.mjs index 3c9373ec547..6b9f8581292 100644 --- a/ui/litellm-dashboard/scripts/gen-api-types.mjs +++ b/ui/litellm-dashboard/scripts/gen-api-types.mjs @@ -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 { diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 7d5d617c824..f15eaf9ea1f 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -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: {