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)