diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index ce4de2d14..6a73162ec 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -324,6 +324,44 @@ class TeamRequest(LiteLLMBase): teams: List[str] +class NewOrganizationRequest(LiteLLMBase): + organization_alias: Optional[str] = None + models: List = [] + budget_id: Optional[str] = None + tpm_limit: Optional[int] = None + rpm_limit: Optional[int] = None + max_budget: Optional[float] = None + + +class LiteLLM_BudgetTable(LiteLLMBase): + """Represents user-controllable params for a LiteLLM_BudgetTable record""" + + max_budget: Optional[float] = None + max_parallel_requests: Optional[int] = None + tpm_limit: Optional[int] = None + rpm_limit: Optional[int] = None + model_max_budget: dict + budget_duration: Optional[str] = None + budget_reset_at: Optional[datetime] = None + created_by: str + updated_by: str + + +class LiteLLM_OrganizationTable(LiteLLMBase): + + organization_id: str + organization_alias: Optional[str] = None + budget_id: str + metadata: dict + models: List[str] + spend: float + model_spend: dict + created_at: datetime + created_by: str + updated_at: datetime + updated_by: str + + class KeyManagementSystem(enum.Enum): GOOGLE_KMS = "google_kms" AZURE_KEY_VAULT = "azure_key_vault" diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 869de6dde..384d10961 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -239,6 +239,7 @@ health_check_interval = None health_check_results = {} queue: List = [] litellm_proxy_budget_name = "litellm-proxy-budget" +litellm_proxy_admin_name = "default_user_id" ui_access_mode: Literal["admin", "all"] = "all" proxy_budget_rescheduler_min_time = 597 proxy_budget_rescheduler_max_time = 605 @@ -335,7 +336,11 @@ async def user_api_key_auth( # note: never string compare api keys, this is vulenerable to a time attack. Use secrets.compare_digest instead is_master_key_valid = secrets.compare_digest(api_key, master_key) if is_master_key_valid: - return UserAPIKeyAuth(api_key=master_key, user_role="proxy_admin") + return UserAPIKeyAuth( + api_key=master_key, + user_role="proxy_admin", + user_id=litellm_proxy_admin_name, + ) if isinstance( api_key, str ): # if generated token, make sure it starts with sk-. @@ -360,7 +365,6 @@ async def user_api_key_auth( valid_token = await prisma_client.get_data( token=api_key, table_name="combined_view" ) - elif custom_db_client is not None: try: valid_token = await custom_db_client.get_data( @@ -2213,7 +2217,7 @@ def parse_cache_control(cache_control): @router.on_event("startup") async def startup_event(): - global prisma_client, master_key, use_background_health_checks, llm_router, llm_model_list, general_settings, proxy_budget_rescheduler_min_time, proxy_budget_rescheduler_max_time + global prisma_client, master_key, use_background_health_checks, llm_router, llm_model_list, general_settings, proxy_budget_rescheduler_min_time, proxy_budget_rescheduler_max_time, litellm_proxy_admin_name import json ### LOAD MASTER KEY ### @@ -2260,9 +2264,8 @@ async def startup_event(): if prisma_client is not None and master_key is not None: # add master key to db - user_id = "default_user_id" if os.getenv("PROXY_ADMIN_ID", None) is not None: - user_id = os.getenv("PROXY_ADMIN_ID") + litellm_proxy_admin_name = os.getenv("PROXY_ADMIN_ID") asyncio.create_task( generate_key_helper_fn( @@ -2272,7 +2275,7 @@ async def startup_event(): config={}, spend=0, token=master_key, - user_id=user_id, + user_id=litellm_proxy_admin_name, user_role="proxy_admin", query_type="update_data", update_key_values={ @@ -5412,6 +5415,81 @@ async def team_info( ) +#### ORGANIZATION MANAGEMENT #### + + +@router.post( + "/organization/new", + tags=["organization management"], + dependencies=[Depends(user_api_key_auth)], + response_model=LiteLLM_OrganizationTable, +) +async def new_organization( + data: NewOrganizationRequest, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + global prisma_client + + if prisma_client is None: + raise HTTPException(status_code=500, detail={"error": "No db connected"}) + + if data.budget_id is None: + """ + Every organization needs a budget attached. + + If none provided, create one based on user max + """ + budget_row = LiteLLM_BudgetTable( + max_budget=user_api_key_dict.max_budget, + max_parallel_requests=user_api_key_dict.max_parallel_requests, + model_max_budget=user_api_key_dict.model_max_budget, + tpm_limit=user_api_key_dict.tpm_limit, + rpm_limit=user_api_key_dict.rpm_limit, + budget_duration=user_api_key_dict.budget_duration, + budget_reset_at=user_api_key_dict.budget_reset_at, + created_by=user_api_key_dict.user_id or litellm_proxy_admin_name, + updated_by=user_api_key_dict.user_id or litellm_proxy_admin_name, + ) + + new_budget = prisma_client.jsonify_object( + budget_row.model_dump(exclude_none=True) + ) + + _budget = await prisma_client.db.litellm_budgettable.create(data={**new_budget}) # type: ignore + + data.budget_id = _budget.budget_id + + response = await prisma_client.db.litellm_organizationtable.create( + data={ + **data.model_dump(exclude_none=True), # type: ignore + "created_by": user_api_key_dict.user_id, + "updated_by": user_api_key_dict.user_id, + } + ) + + return response + + +@router.post( + "/organization/update", + tags=["organization management"], + dependencies=[Depends(user_api_key_auth)], + response_model=LiteLLM_TeamTable, +) +async def update_organization(): + pass + + +@router.post( + "/organization/delete", + tags=["organization management"], + dependencies=[Depends(user_api_key_auth)], + response_model=LiteLLM_TeamTable, +) +async def delete_organization(): + pass + + #### MODEL MANAGEMENT #### diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 6a9b72728..4e0da8222 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -7,10 +7,44 @@ generator client { provider = "prisma-client-py" } +// Budget / Rate Limits for an org +model LiteLLM_BudgetTable { + budget_id String @id @default(uuid()) + max_budget Float? + max_parallel_requests Int? + tpm_limit BigInt? + rpm_limit BigInt? + model_max_budget Json @default("{}") + budget_duration String? + budget_reset_at DateTime? + created_at DateTime @default(now()) @map("created_at") + created_by String + updated_at DateTime @default(now()) @updatedAt @map("updated_at") + updated_by String + organization LiteLLM_OrganizationTable[] +} + +model LiteLLM_OrganizationTable { + organization_id String @id @default(uuid()) + organization_alias String? + budget_id String + metadata Json @default("{}") + models String[] + spend Float @default(0.0) + model_spend Json @default("{}") + created_at DateTime @default(now()) @map("created_at") + created_by String + updated_at DateTime @default(now()) @updatedAt @map("updated_at") + updated_by String + litellm_budget_table LiteLLM_BudgetTable? @relation(fields: [budget_id], references: [budget_id]) + teams LiteLLM_TeamTable[] +} + // Assign prod keys to groups, not individuals model LiteLLM_TeamTable { - team_id String @unique + team_id String @id @default(uuid()) team_alias String? + organization_id String? admins String[] members String[] members_with_roles Json @default("{}") @@ -27,11 +61,12 @@ model LiteLLM_TeamTable { updated_at DateTime @default(now()) @updatedAt @map("updated_at") model_spend Json @default("{}") model_max_budget Json @default("{}") + litellm_organization_table LiteLLM_OrganizationTable? @relation(fields: [organization_id], references: [organization_id]) } // Track spend, rate limit, budget Users model LiteLLM_UserTable { - user_id String @unique + user_id String @id team_id String? teams String[] @default([]) user_role String? @@ -51,7 +86,7 @@ model LiteLLM_UserTable { // Generate Tokens for Proxy model LiteLLM_VerificationToken { - token String @unique + token String @id key_name String? key_alias String? spend Float @default(0.0) @@ -82,7 +117,7 @@ model LiteLLM_Config { // View spend, model, api_key per request model LiteLLM_SpendLogs { - request_id String @unique + request_id String @id call_type String api_key String @default ("") spend Float @default(0.0) @@ -98,11 +133,12 @@ model LiteLLM_SpendLogs { cache_key String @default("") request_tags Json @default("[]") team_id String? - end_user String? + end_user String? } + // Beta - allow team members to request access to a model model LiteLLM_UserNotifications { - request_id String @unique + request_id String @id user_id String models String[] justification String diff --git a/schema.prisma b/schema.prisma index d08295e10..f31fa130a 100644 --- a/schema.prisma +++ b/schema.prisma @@ -7,10 +7,42 @@ generator client { provider = "prisma-client-py" } +// Budget / Rate Limits for an org +model LiteLLM_BudgetTable { + budget_id String @id @default(uuid()) + max_budget Float? + max_parallel_requests Int? + tpm_limit BigInt? + rpm_limit BigInt? + model_max_budget Json @default("{}") + budget_duration String? + budget_reset_at DateTime? + created_at DateTime @default(now()) @map("created_at") + created_by String + updated_at DateTime @default(now()) @updatedAt @map("updated_at") + updated_by String +} + +model LiteLLM_OrganizationTable { + organization_id String @id @default(uuid()) + organization_alias String? + budget_id String + metadata Json @default("{}") + models String[] + spend Float @default(0.0) + model_spend Json @default("{}") + created_at DateTime @default(now()) @map("created_at") + created_by String + updated_at DateTime @default(now()) @updatedAt @map("updated_at") + updated_by String + litellm_budget_table LiteLLM_BudgetTable @relation(fields: [budget_id], references: [budget_id]) +} + // Assign prod keys to groups, not individuals model LiteLLM_TeamTable { - team_id String @unique + team_id String @id @default(uuid()) team_alias String? + organization_id String? admins String[] members String[] members_with_roles Json @default("{}") @@ -27,11 +59,12 @@ model LiteLLM_TeamTable { updated_at DateTime @default(now()) @updatedAt @map("updated_at") model_spend Json @default("{}") model_max_budget Json @default("{}") + litellm_organization_table LiteLLM_OrganizationTable @relation(fields: [organization_id], references: [organization_id]) } // Track spend, rate limit, budget Users model LiteLLM_UserTable { - user_id String @unique + user_id String @id team_id String? teams String[] @default([]) user_role String? @@ -51,7 +84,7 @@ model LiteLLM_UserTable { // Generate Tokens for Proxy model LiteLLM_VerificationToken { - token String @unique + token String @id key_name String? key_alias String? spend Float @default(0.0) @@ -82,7 +115,7 @@ model LiteLLM_Config { // View spend, model, api_key per request model LiteLLM_SpendLogs { - request_id String @unique + request_id String @id call_type String api_key String @default ("") spend Float @default(0.0) @@ -103,7 +136,7 @@ model LiteLLM_SpendLogs { // Beta - allow team members to request access to a model model LiteLLM_UserNotifications { - request_id String @unique + request_id String @id user_id String models String[] justification String