From 2602102ce620f1bef6b5534677c4d0cca1d51a3d Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Sat, 2 Mar 2024 11:55:16 -0800 Subject: [PATCH 1/6] feat(proxy_server.py): enable `/organizations/new` endpoint allows admins to create organizations which can own teams --- litellm/proxy/_types.py | 38 +++++++++++++++ litellm/proxy/proxy_server.py | 90 ++++++++++++++++++++++++++++++++--- litellm/proxy/schema.prisma | 48 ++++++++++++++++--- schema.prisma | 43 +++++++++++++++-- 4 files changed, 202 insertions(+), 17 deletions(-) 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 From eedd446a6ac0ca5f24f78e0283eff03f31f532b2 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Sat, 2 Mar 2024 12:01:05 -0800 Subject: [PATCH 2/6] fix(proxy_server.py): fix pydantic versioning issue --- litellm/proxy/proxy_server.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 384d10961..e5af9186c 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -5451,9 +5451,7 @@ async def new_organization( 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) - ) + new_budget = prisma_client.jsonify_object(budget_row.json(exclude_none=True)) _budget = await prisma_client.db.litellm_budgettable.create(data={**new_budget}) # type: ignore @@ -5461,7 +5459,7 @@ async def new_organization( response = await prisma_client.db.litellm_organizationtable.create( data={ - **data.model_dump(exclude_none=True), # type: ignore + **data.json(exclude_none=True), # type: ignore "created_by": user_api_key_dict.user_id, "updated_by": user_api_key_dict.user_id, } From 6fb19c5d420b9f8ac475df075ce8cadd5fc2a6bc Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Sat, 2 Mar 2024 12:13:54 -0800 Subject: [PATCH 3/6] test(test_organizations.py): add testing for `/organization/new` endpoint --- litellm/proxy/schema.prisma | 4 ++-- schema.prisma | 8 ++++--- tests/test_organizations.py | 46 +++++++++++++++++++++++++++++++++++++ 3 files changed, 53 insertions(+), 5 deletions(-) create mode 100644 tests/test_organizations.py diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 4e0da8222..2607cf2b0 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -21,12 +21,12 @@ model LiteLLM_BudgetTable { created_by String updated_at DateTime @default(now()) @updatedAt @map("updated_at") updated_by String - organization LiteLLM_OrganizationTable[] + organization LiteLLM_OrganizationTable[] // multiple orgs can have the same budget } model LiteLLM_OrganizationTable { organization_id String @id @default(uuid()) - organization_alias String? + organization_alias String budget_id String metadata Json @default("{}") models String[] diff --git a/schema.prisma b/schema.prisma index f31fa130a..2607cf2b0 100644 --- a/schema.prisma +++ b/schema.prisma @@ -21,11 +21,12 @@ model LiteLLM_BudgetTable { created_by String updated_at DateTime @default(now()) @updatedAt @map("updated_at") updated_by String + organization LiteLLM_OrganizationTable[] // multiple orgs can have the same budget } model LiteLLM_OrganizationTable { organization_id String @id @default(uuid()) - organization_alias String? + organization_alias String budget_id String metadata Json @default("{}") models String[] @@ -35,7 +36,8 @@ model LiteLLM_OrganizationTable { 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]) + litellm_budget_table LiteLLM_BudgetTable? @relation(fields: [budget_id], references: [budget_id]) + teams LiteLLM_TeamTable[] } // Assign prod keys to groups, not individuals @@ -59,7 +61,7 @@ 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]) + litellm_organization_table LiteLLM_OrganizationTable? @relation(fields: [organization_id], references: [organization_id]) } // Track spend, rate limit, budget Users diff --git a/tests/test_organizations.py b/tests/test_organizations.py new file mode 100644 index 000000000..00e99cb66 --- /dev/null +++ b/tests/test_organizations.py @@ -0,0 +1,46 @@ +# What this tests ? +## Tests /organization endpoints. +import pytest +import asyncio +import aiohttp +import time, uuid +from openai import AsyncOpenAI + + +async def new_organization(session, i, organization_alias, max_budget=None): + url = "http://0.0.0.0:4000/organization/new" + headers = {"Authorization": "Bearer sk-1234", "Content-Type": "application/json"} + data = { + "organization_alias": organization_alias, + "models": ["azure-models"], + "max_budget": max_budget, + } + + async with session.post(url, headers=headers, json=data) as response: + status = response.status + response_text = await response.text() + + print(f"Response {i} (Status code: {status}):") + print(response_text) + print() + + if status != 200: + raise Exception(f"Request {i} did not return a 200 status code: {status}") + + return await response.json() + + +@pytest.mark.asyncio +async def test_organization_new(): + """ + Make 20 parallel calls to /user/new. Assert all worked. + """ + organization_alias = f"Organization: {uuid.uuid4()}" + async with aiohttp.ClientSession() as session: + tasks = [ + new_organization( + session=session, i=0, organization_alias=organization_alias + ) + for i in range(1, 20) + ] + await asyncio.gather(*tasks) From 8d22ed762efe1219d84ee57e50d6389b2d117942 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Sat, 2 Mar 2024 14:38:42 -0800 Subject: [PATCH 4/6] fix(proxy_server.py): enable admin to create new budget if none set for org --- litellm/proxy/_types.py | 34 ++++----- litellm/proxy/proxy_server.py | 137 +++++++++++++++++++++++++++++----- schema.prisma | 2 +- 3 files changed, 134 insertions(+), 39 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 6a73162ec..a962b6aa6 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -324,15 +324,6 @@ 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""" @@ -340,28 +331,33 @@ class LiteLLM_BudgetTable(LiteLLMBase): max_parallel_requests: Optional[int] = None tpm_limit: Optional[int] = None rpm_limit: Optional[int] = None - model_max_budget: dict + model_max_budget: Optional[dict] = None budget_duration: Optional[str] = None - budget_reset_at: Optional[datetime] = None - created_by: str - updated_by: str + + +class NewOrganizationRequest(LiteLLM_BudgetTable): + organization_alias: str + models: List = [] + budget_id: Optional[str] = None class LiteLLM_OrganizationTable(LiteLLMBase): + """Represents user-controllable params for a LiteLLM_OrganizationTable record""" - organization_id: str organization_alias: Optional[str] = None budget_id: str - metadata: dict + metadata: Optional[dict] = None models: List[str] - spend: float - model_spend: dict - created_at: datetime created_by: str - updated_at: datetime updated_by: str +class NewOrganizationResponse(LiteLLM_OrganizationTable): + organization_id: str + created_at: datetime + updated_at: datetime + + 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 e5af9186c..f4909a408 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -5422,46 +5422,133 @@ async def team_info( "/organization/new", tags=["organization management"], dependencies=[Depends(user_api_key_auth)], - response_model=LiteLLM_OrganizationTable, + response_model=NewOrganizationResponse, ) async def new_organization( data: NewOrganizationRequest, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): + """ + Allow orgs to own teams + + Set org level budgets + model access. + + Only admins can create orgs. + + # Parameters + + - `organization_alias`: *str* = The name of the organization. + - `models`: *List* = The models the organization has access to. + - `budget_id`: *Optional[str]* = The id for a budget (tpm/rpm/max budget) for the organization. + ### IF NO BUDGET - CREATE ONE WITH THESE PARAMS ### + - `max_budget`: *Optional[float]* = Max budget for org + - `tpm_limit`: *Optional[int]* = Max tpm limit for org + - `rpm_limit`: *Optional[int]* = Max rpm limit for org + - `model_max_budget`: *Optional[dict]* = Max budget for a specific model + - `budget_duration`: *Optional[str]* = Frequency of reseting org budget + + Case 1: Create new org **without** a budget_id + + ```bash + curl --location 'http://0.0.0.0:4000/organization/new' \ + + --header 'Authorization: Bearer sk-1234' \ + + --header 'Content-Type: application/json' \ + + --data '{ + "organization_alias": "my-secret-org", + "models": ["model1", "model2"], + "max_budget": 100 + }' + + + ``` + + Case 2: Create new org **with** a budget_id + + ```bash + curl --location 'http://0.0.0.0:4000/organization/new' \ + + --header 'Authorization: Bearer sk-1234' \ + + --header 'Content-Type: application/json' \ + + --data '{ + "organization_alias": "my-secret-org", + "models": ["model1", "model2"], + "budget_id": "428eeaa8-f3ac-4e85-a8fb-7dc8d7aa8689" + }' + ``` + """ global prisma_client if prisma_client is None: raise HTTPException(status_code=500, detail={"error": "No db connected"}) + if ( + user_api_key_dict.user_role is None + or user_api_key_dict.user_role != "proxy_admin" + ): + raise HTTPException( + status_code=401, + detail={ + "error": f"Only admins can create orgs. Your role is = {user_api_key_dict.user_role}" + }, + ) + if data.budget_id is None: """ Every organization needs a budget attached. - If none provided, create one based on user max + If none provided, create one based on provided values """ - 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, - ) + budget_row = LiteLLM_BudgetTable(**data.json(exclude_none=True)) new_budget = prisma_client.jsonify_object(budget_row.json(exclude_none=True)) - _budget = await prisma_client.db.litellm_budgettable.create(data={**new_budget}) # type: ignore + _budget = await prisma_client.db.litellm_budgettable.create( + data={ + **new_budget, # type: ignore + "created_by": user_api_key_dict.user_id or litellm_proxy_admin_name, + "updated_by": user_api_key_dict.user_id or litellm_proxy_admin_name, + } + ) # type: ignore data.budget_id = _budget.budget_id + """ + Ensure only models that user has access to, are given to org + """ + if len(user_api_key_dict.models) == 0: # user has access to all models + pass + else: + if len(data.models) == 0: + raise HTTPException( + status_code=400, + detail={ + "error": f"User not allowed to give access to all models. Select models you want org to have access to." + }, + ) + for m in data.models: + if m not in user_api_key_dict.models: + raise HTTPException( + status_code=400, + detail={ + "error": f"User not allowed to give access to model={m}. Models you have access to = {user_api_key_dict.models}" + }, + ) + organization_row = LiteLLM_OrganizationTable( + **data.json(exclude_none=True), + 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_organization_row = prisma_client.jsonify_object( + organization_row.json(exclude_none=True) + ) response = await prisma_client.db.litellm_organizationtable.create( data={ - **data.json(exclude_none=True), # type: ignore - "created_by": user_api_key_dict.user_id, - "updated_by": user_api_key_dict.user_id, + **new_organization_row, # type: ignore } ) @@ -5472,9 +5559,9 @@ async def new_organization( "/organization/update", tags=["organization management"], dependencies=[Depends(user_api_key_auth)], - response_model=LiteLLM_TeamTable, ) async def update_organization(): + """[TODO] Not Implemented yet. Let us know if you need this - https://github.com/BerriAI/litellm/issues""" pass @@ -5482,9 +5569,21 @@ async def update_organization(): "/organization/delete", tags=["organization management"], dependencies=[Depends(user_api_key_auth)], - response_model=LiteLLM_TeamTable, ) async def delete_organization(): + """[TODO] Not Implemented yet. Let us know if you need this - https://github.com/BerriAI/litellm/issues""" + pass + + +@router.post( + "/organization/info", + tags=["organization management"], + dependencies=[Depends(user_api_key_auth)], +) +async def info_organization(): + """ + Get the org specific information + """ pass diff --git a/schema.prisma b/schema.prisma index 2607cf2b0..e7932d634 100644 --- a/schema.prisma +++ b/schema.prisma @@ -14,7 +14,7 @@ model LiteLLM_BudgetTable { max_parallel_requests Int? tpm_limit BigInt? rpm_limit BigInt? - model_max_budget Json @default("{}") + model_max_budget Json? budget_duration String? budget_reset_at DateTime? created_at DateTime @default(now()) @map("created_at") From 8bb6897b461c1d49d2a1e05947d8ee5040cd80bf Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Sat, 2 Mar 2024 15:07:33 -0800 Subject: [PATCH 5/6] feat(proxy_server.py): exposes `/organization/info` and `/budget/info` endpoints --- litellm/proxy/_types.py | 8 ++++++ litellm/proxy/proxy_server.py | 52 +++++++++++++++++++++++++++++++++-- 2 files changed, 58 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index a962b6aa6..5cd289b66 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -358,6 +358,14 @@ class NewOrganizationResponse(LiteLLM_OrganizationTable): updated_at: datetime +class OrganizationRequest(LiteLLMBase): + organizations: List[str] + + +class BudgetRequest(LiteLLMBase): + budgets: List[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 f4909a408..bd3d111a7 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -5580,11 +5580,59 @@ async def delete_organization(): tags=["organization management"], dependencies=[Depends(user_api_key_auth)], ) -async def info_organization(): +async def info_organization(data: OrganizationRequest): """ Get the org specific information """ - pass + global prisma_client + + if prisma_client is None: + raise HTTPException(status_code=500, detail={"error": "No db connected"}) + + if len(data.organizations) == 0: + raise HTTPException( + status_code=400, + detail={ + "error": f"Specify list of organization id's to query. Passed in={data.organizations}" + }, + ) + response = await prisma_client.db.litellm_organizationtable.find_many( + where={"organization_id": {"in": data.organizations}}, + include={"litellm_budget_table": True}, + ) + + return response + + +#### BUDGET TABLE MANAGEMENT #### + + +@router.post( + "/budget/info", + tags=["organization management"], + dependencies=[Depends(user_api_key_auth)], +) +async def info_budget(data: BudgetRequest): + """ + Get the budget id specific information + """ + global prisma_client + + if prisma_client is None: + raise HTTPException(status_code=500, detail={"error": "No db connected"}) + + if len(data.budgets) == 0: + raise HTTPException( + status_code=400, + detail={ + "error": f"Specify list of budget id's to query. Passed in={data.budgets}" + }, + ) + response = await prisma_client.db.litellm_budgettable.find_many( + where={"budget_id": {"in": data.budgets}}, + ) + + return response #### MODEL MANAGEMENT #### From 96f11157c90a40286282a962bb4ded358996b17c Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Sat, 2 Mar 2024 15:32:43 -0800 Subject: [PATCH 6/6] fix(_types.py): have org id in newteamrequest pydantic object --- litellm/proxy/_types.py | 1 + 1 file changed, 1 insertion(+) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 5cd289b66..e981aef6d 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -250,6 +250,7 @@ class Member(LiteLLMBase): class NewTeamRequest(LiteLLMBase): team_alias: Optional[str] = None team_id: Optional[str] = None + organization_id: Optional[str] = None admins: list = [] members: list = [] members_with_roles: List[Member] = []