diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260214185341_object_permissions_for_end_users/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260214185341_object_permissions_for_end_users/migration.sql new file mode 100644 index 00000000000..5c5dc6fd6f1 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260214185341_object_permissions_for_end_users/migration.sql @@ -0,0 +1,6 @@ +-- AlterTable +ALTER TABLE "LiteLLM_EndUserTable" ADD COLUMN "object_permission_id" TEXT; + +-- AddForeignKey +ALTER TABLE "LiteLLM_EndUserTable" ADD CONSTRAINT "LiteLLM_EndUserTable_object_permission_id_fkey" FOREIGN KEY ("object_permission_id") REFERENCES "LiteLLM_ObjectPermissionTable"("object_permission_id") ON DELETE SET NULL ON UPDATE CASCADE; + diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index c2fca8705cb..af52d998949 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -233,6 +233,7 @@ model LiteLLM_ObjectPermissionTable { verification_tokens LiteLLM_VerificationToken[] organizations LiteLLM_OrganizationTable[] users LiteLLM_UserTable[] + end_users LiteLLM_EndUserTable[] } // Holds the MCP server configuration @@ -390,7 +391,9 @@ model LiteLLM_EndUserTable { allowed_model_region String? // require all user requests to use models in this specific region default_model String? // use along with 'allowed_model_region'. if no available model in region, default to this model. budget_id String? + object_permission_id String? litellm_budget_table LiteLLM_BudgetTable? @relation(fields: [budget_id], references: [budget_id]) + object_permission LiteLLM_ObjectPermissionTable? @relation(fields: [object_permission_id], references: [object_permission_id]) blocked Boolean @default(false) } diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index d549338972c..b71ae9fde60 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1,60 +1,40 @@ import enum import json from datetime import datetime -from typing import TYPE_CHECKING, Any, Callable, Dict, List, Literal, Optional, Union +from typing import (TYPE_CHECKING, Any, Callable, Dict, List, Literal, + Optional, Union) import httpx -from pydantic import ( - BaseModel, - ConfigDict, - Field, - Json, - field_validator, - model_validator, -) +from pydantic import (BaseModel, ConfigDict, Field, Json, field_validator, + model_validator) from typing_extensions import Required, TypedDict from litellm._uuid import uuid from litellm.types.integrations.slack_alerting import AlertType -from litellm.types.llms.openai import ( - AllMessageValues, - OpenAIFileObject, - ResponsesAPIResponse, -) -from litellm.types.mcp import ( - MCPAuth, - MCPAuthType, - MCPCredentials, - MCPTransport, - MCPTransportType, -) +from litellm.types.llms.openai import (AllMessageValues, OpenAIFileObject, + ResponsesAPIResponse) +from litellm.types.mcp import (MCPAuth, MCPAuthType, MCPCredentials, + MCPTransport, MCPTransportType) from litellm.types.mcp_server.mcp_server_manager import MCPInfo from litellm.types.router import RouterErrors, UpdateRouterConfig from litellm.types.secret_managers.main import KeyManagementSystem -from litellm.types.utils import ( - CallTypes, - CostBreakdown, - EmbeddingResponse, - GenericBudgetConfigType, - ImageResponse, - LiteLLMBatch, - LiteLLMFineTuningJob, - LiteLLMPydanticObjectBase, - ModelResponse, - ProviderField, - StandardCallbackDynamicParams, - StandardLoggingGuardrailInformation, - StandardLoggingMCPToolCall, - StandardLoggingModelInformation, - StandardLoggingPayloadErrorInformation, - StandardLoggingPayloadStatus, - StandardLoggingVectorStoreRequest, - StandardPassThroughResponseObject, - TextCompletionResponse, -) +from litellm.types.utils import (CallTypes, CostBreakdown, EmbeddingResponse, + GenericBudgetConfigType, ImageResponse, + LiteLLMBatch, LiteLLMFineTuningJob, + LiteLLMPydanticObjectBase, ModelResponse, + ProviderField, StandardCallbackDynamicParams, + StandardLoggingGuardrailInformation, + StandardLoggingMCPToolCall, + StandardLoggingModelInformation, + StandardLoggingPayloadErrorInformation, + StandardLoggingPayloadStatus, + StandardLoggingVectorStoreRequest, + StandardPassThroughResponseObject, + TextCompletionResponse) from litellm.types.videos.main import VideoObject -from .types_utils.utils import get_instance_fn, validate_custom_validate_return_type +from .types_utils.utils import (get_instance_fn, + validate_custom_validate_return_type) if TYPE_CHECKING: from opentelemetry.trace import Span as _Span @@ -1412,6 +1392,7 @@ class NewCustomerRequest(BudgetNewRequest): default_model: Optional[str] = ( None # if no equivalent model in allowed region - default all requests to this model ) + object_permission: Optional[LiteLLM_ObjectPermissionBase] = None @model_validator(mode="before") @classmethod @@ -1439,6 +1420,7 @@ class UpdateCustomerRequest(LiteLLMPydanticObjectBase): default_model: Optional[str] = ( None # if no equivalent model in allowed region - default all requests to this model ) + object_permission: Optional[LiteLLM_ObjectPermissionBase] = None class DeleteCustomerRequest(LiteLLMPydanticObjectBase): @@ -2336,7 +2318,8 @@ class UserAPIKeyAuth( This is used to track number of requests/spend for health check calls. """ - from litellm.constants import LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME + from litellm.constants import \ + LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME return cls( api_key=LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME, @@ -2368,7 +2351,8 @@ class UserAPIKeyAuth( This is used to track actions performed by automated system jobs. """ - from litellm.constants import LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME + from litellm.constants import \ + LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME return cls( api_key=LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME, @@ -2528,6 +2512,8 @@ class LiteLLM_EndUserTable(LiteLLMPydanticObjectBase): allowed_model_region: Optional[AllowedModelRegion] = None default_model: Optional[str] = None litellm_budget_table: Optional[LiteLLM_BudgetTable] = None + object_permission_id: Optional[str] = None + object_permission: Optional[LiteLLM_ObjectPermissionTable] = None @model_validator(mode="before") @classmethod @@ -2637,7 +2623,8 @@ class LiteLLM_AuditLogs(LiteLLMPydanticObjectBase): @model_validator(mode="after") def mask_api_keys(self): - from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker + from litellm.litellm_core_utils.sensitive_data_masker import \ + SensitiveDataMasker masker = SensitiveDataMasker(sensitive_patterns={"key"}) diff --git a/litellm/proxy/management_endpoints/customer_endpoints.py b/litellm/proxy/management_endpoints/customer_endpoints.py index 9ff0fe6e590..025b8194803 100644 --- a/litellm/proxy/management_endpoints/customer_endpoints.py +++ b/litellm/proxy/management_endpoints/customer_endpoints.py @@ -19,11 +19,13 @@ import litellm from litellm._logging import verbose_proxy_logger from litellm.proxy._types import * from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.management_endpoints.common_daily_activity import \ + get_daily_activity +from litellm.proxy.management_helpers.object_permission_utils import ( + _set_object_permission, handle_update_object_permission_common) from litellm.proxy.utils import handle_exception_on_proxy -from litellm.types.proxy.management_endpoints.common_daily_activity import ( - SpendAnalyticsPaginatedResponse, -) -from litellm.proxy.management_endpoints.common_daily_activity import get_daily_activity +from litellm.types.proxy.management_endpoints.common_daily_activity import \ + SpendAnalyticsPaginatedResponse router = APIRouter() @@ -107,9 +109,8 @@ async def unblock_user(data: BlockUsers): ``` """ try: - from enterprise.enterprise_hooks.blocked_user_list import ( - _ENTERPRISE_BlockedUserList, - ) + from enterprise.enterprise_hooks.blocked_user_list import \ + _ENTERPRISE_BlockedUserList except ImportError: raise HTTPException( status_code=400, @@ -200,6 +201,7 @@ async def new_end_user( - soft_budget: Optional[float] - [Not Implemented Yet] Get alerts when customer crosses given budget, doesn't block requests. - spend: Optional[float] - Specify initial spend for a given customer. - budget_reset_at: Optional[str] - Specify the date and time when the budget should be reset. + - object_permission: Optional[LiteLLM_ObjectPermissionBase] - Customer-specific object permission. Example - {"mcp_servers": ["server_1", "server_2"], "vector_stores": ["vector_store_1"], "agents": ["agent_1"]}. IF null or {} then no object permission. - Allow specifying allowed regions @@ -233,11 +235,8 @@ async def new_end_user( - end-user object - currently allowed models """ - from litellm.proxy.proxy_server import ( - litellm_proxy_admin_name, - llm_router, - prisma_client, - ) + from litellm.proxy.proxy_server import (litellm_proxy_admin_name, + llm_router, prisma_client) if prisma_client is None: raise HTTPException( @@ -289,10 +288,16 @@ async def new_end_user( if k not in BudgetNewRequest.model_fields.keys(): new_end_user_obj[k] = v + ## Handle Object Permission - MCP Servers, Vector Stores etc. + new_end_user_obj = await _set_object_permission( + data_json=new_end_user_obj, + prisma_client=prisma_client, + ) + ## WRITE TO DB ## end_user_record = await prisma_client.db.litellm_endusertable.create( data=new_end_user_obj, # type: ignore - include={"litellm_budget_table": True}, + include={"litellm_budget_table": True, "object_permission": True}, ) return end_user_record @@ -351,7 +356,7 @@ async def end_user_info( ) user_info = await prisma_client.db.litellm_endusertable.find_first( - where={"user_id": end_user_id}, include={"litellm_budget_table": True} + where={"user_id": end_user_id}, include={"litellm_budget_table": True, "object_permission": True} ) if user_info is None: @@ -401,6 +406,7 @@ async def update_end_user( - default_model: Optional[str] = ( None # if no equivalent model in allowed region - default all requests to this model ) + - object_permission: Optional[LiteLLM_ObjectPermissionBase] = Customer-specific object permission. Example - {"mcp_servers": ["server_1"], "vector_stores": ["vector_store_1"]}. IF null or {} then no object permission. Example curl: ``` @@ -416,7 +422,8 @@ async def update_end_user( ``` """ - from litellm.proxy.proxy_server import litellm_proxy_admin_name, prisma_client + from litellm.proxy.proxy_server import (litellm_proxy_admin_name, + prisma_client) try: data_json: dict = data.json() @@ -467,6 +474,21 @@ async def update_end_user( elif k in LiteLLM_EndUserTable.model_fields.keys(): update_end_user_table_data[k] = v + ## Handle object permission updates (MCP servers, vector stores, etc.) + if "object_permission" in non_default_values: + existing_object_permission_id = ( + end_user_table_data_typed.object_permission_id + if end_user_table_data_typed is not None + else None + ) + object_permission_id = await handle_update_object_permission_common( + data_json=non_default_values, + existing_object_permission_id=existing_object_permission_id, + prisma_client=prisma_client, + ) + if object_permission_id is not None: + update_end_user_table_data["object_permission_id"] = object_permission_id + ## Check if we need to create a new budget (only if budget fields are provided, not just budget_id) ## if budget_table_data: if end_user_budget_table is None: @@ -502,7 +524,7 @@ async def update_end_user( update_end_user_table_data["user_id"] = data.user_id # type: ignore verbose_proxy_logger.debug("In update customer, user_id condition block.") response = await prisma_client.db.litellm_endusertable.update( - where={"user_id": data.user_id}, data=update_end_user_table_data, include={"litellm_budget_table": True} # type: ignore + where={"user_id": data.user_id}, data=update_end_user_table_data, include={"litellm_budget_table": True, "object_permission": True} # type: ignore ) if response is None: raise ValueError( @@ -663,7 +685,7 @@ async def list_end_user( ) response = await prisma_client.db.litellm_endusertable.find_many( - include={"litellm_budget_table": True} + include={"litellm_budget_table": True, "object_permission": True} ) returned_response: List[LiteLLM_EndUserTable] = [] @@ -706,9 +728,7 @@ async def get_customer_daily_activity( """ Get daily activity for specific organizations or all accessible organizations. """ - from litellm.proxy.proxy_server import ( - prisma_client, - ) + from litellm.proxy.proxy_server import prisma_client if prisma_client is None: raise HTTPException( diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index c2fca8705cb..af52d998949 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -233,6 +233,7 @@ model LiteLLM_ObjectPermissionTable { verification_tokens LiteLLM_VerificationToken[] organizations LiteLLM_OrganizationTable[] users LiteLLM_UserTable[] + end_users LiteLLM_EndUserTable[] } // Holds the MCP server configuration @@ -390,7 +391,9 @@ model LiteLLM_EndUserTable { allowed_model_region String? // require all user requests to use models in this specific region default_model String? // use along with 'allowed_model_region'. if no available model in region, default to this model. budget_id String? + object_permission_id String? litellm_budget_table LiteLLM_BudgetTable? @relation(fields: [budget_id], references: [budget_id]) + object_permission LiteLLM_ObjectPermissionTable? @relation(fields: [object_permission_id], references: [object_permission_id]) blocked Boolean @default(false) } diff --git a/schema.prisma b/schema.prisma index c2fca8705cb..af52d998949 100644 --- a/schema.prisma +++ b/schema.prisma @@ -233,6 +233,7 @@ model LiteLLM_ObjectPermissionTable { verification_tokens LiteLLM_VerificationToken[] organizations LiteLLM_OrganizationTable[] users LiteLLM_UserTable[] + end_users LiteLLM_EndUserTable[] } // Holds the MCP server configuration @@ -390,7 +391,9 @@ model LiteLLM_EndUserTable { allowed_model_region String? // require all user requests to use models in this specific region default_model String? // use along with 'allowed_model_region'. if no available model in region, default to this model. budget_id String? + object_permission_id String? litellm_budget_table LiteLLM_BudgetTable? @relation(fields: [budget_id], references: [budget_id]) + object_permission LiteLLM_ObjectPermissionTable? @relation(fields: [object_permission_id], references: [object_permission_id]) blocked Boolean @default(false) }