phoenix/litellm
1
0
Fork 0
forked from phoenix/litellm-mirror
Code Pull requests Activity

Merge pull request #2907 from BerriAI/litellm_map_user_to_org

Browse source 
fix(proxy_server.py): allow mapping a user to an org
This commit is contained in:
Krish Dholakia 2024-04-08 20:50:47 -07:00 • committed by GitHub
commit 763e92a03e
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 43 additions and 8 deletions
Download patch file Download diff file Expand all files Collapse all files

2
litellm/proxy/_types.py
View file

@ -363,6 +363,8 @@ class NewUserRequest(GenerateKeyRequest):
max_budget: Optional[float] = None
user_email: Optional[str] = None
user_role: Optional[str] = None
teams: Optional[list] = None
organization_id: Optional[str] = None
auto_create_key: bool = (
True # flag used for returning a key as part of the /user/new response
)

9
litellm/proxy/proxy_server.py
View file

@ -2553,6 +2553,8 @@ async def generate_key_helper_fn(
allowed_cache_controls: Optional[list] = [],
permissions: Optional[dict] = {},
model_max_budget: Optional[dict] = {},
teams: Optional[list] = None,
organization_id: Optional[str] = None,
table_name: Optional[Literal["key", "user"]] = None,
):
global prisma_client, custom_db_client, user_api_key_cache, litellm_proxy_admin_name
@ -2600,6 +2602,7 @@ async def generate_key_helper_fn(
"user_email": user_email,
"user_id": user_id,
"team_id": team_id,
"organization_id": organization_id,
"user_role": user_role,
"spend": spend,
"models": models,
@ -2610,6 +2613,8 @@ async def generate_key_helper_fn(
"budget_reset_at": reset_at,
"allowed_cache_controls": allowed_cache_controls,
}
if teams is not None:
user_data["teams"] = teams
key_data = {
"token": token,
"key_alias": key_alias,
@ -2710,6 +2715,8 @@ async def generate_key_helper_fn(
await custom_db_client.insert_data(value=key_data, table_name="key")
except Exception as e:
traceback.print_exc()
if isinstance(e, HTTPException):
raise e
raise HTTPException(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
detail={"error": "Internal Server Error."},
@ -5470,6 +5477,8 @@ async def new_user(data: NewUserRequest):
Parameters:
- user_id: Optional[str] - Specify a user id. If not set, a unique id will be generated.
- user_alias: Optional[str] - A descriptive name for you to know who this user id refers to.
- teams: Optional[list] - specify a list of team id's a user belongs to.
- organization_id: Optional[str] - specify the org a user belongs to.
- user_email: Optional[str] - Specify a user email.
- user_role: Optional[str] - Specify a user role - "admin", "app_owner", "app_user"
- max_budget: Optional[float] - Specify max budget for a given user.

3
litellm/proxy/schema.prisma
View file

@ -53,6 +53,7 @@ model LiteLLM_OrganizationTable {
updated_by String
litellm_budget_table LiteLLM_BudgetTable? @relation(fields: [budget_id], references: [budget_id])
teams LiteLLM_TeamTable[]
users LiteLLM_UserTable[]
}
// Model info for teams, just has model aliases for now.
@ -99,6 +100,7 @@ model LiteLLM_UserTable {
user_id String @id
user_alias String?
team_id String?
organization_id String?
teams String[] @default([])
user_role String?
max_budget Float?
@ -113,6 +115,7 @@ model LiteLLM_UserTable {
allowed_cache_controls String[] @default([])
model_spend Json @default("{}")
model_max_budget Json @default("{}")
litellm_organization_table LiteLLM_OrganizationTable? @relation(fields: [organization_id], references: [organization_id])
}
// Generate Tokens for Proxy

34
litellm/proxy/utils.py
View file

@ -461,7 +461,12 @@ class ProxyLogging:
"""
### ALERTING ###
if isinstance(original_exception, HTTPException):
error_message = original_exception.detail
if isinstance(original_exception.detail, str):
error_message = original_exception.detail
elif isinstance(original_exception.detail, dict):
error_message = json.dumps(original_exception.detail)
else:
error_message = str(original_exception)
else:
error_message = str(original_exception)
if isinstance(traceback_str, str):
@ -1159,13 +1164,26 @@ class PrismaClient:
return new_verification_token
elif table_name == "user":
db_data = self.jsonify_object(data=data)
new_user_row = await self.db.litellm_usertable.upsert(
where={"user_id": data["user_id"]},
data={
"create": {**db_data}, # type: ignore
"update": {}, # don't do anything if it already exists
},
)
try:
new_user_row = await self.db.litellm_usertable.upsert(
where={"user_id": data["user_id"]},
data={
"create": {**db_data}, # type: ignore
"update": {}, # don't do anything if it already exists
},
)
except Exception as e:
if (
"Foreign key constraint failed on the field: `LiteLLM_UserTable_organization_id_fkey (index)`"
in str(e)
):
raise HTTPException(
status_code=400,
detail={
"error": f"Foreign Key Constraint failed. Organization ID={db_data['organization_id']} does not exist in LiteLLM_OrganizationTable. Create via `/organization/new`."
},
)
raise e
verbose_proxy_logger.info("Data Inserted into User Table")
return new_user_row
elif table_name == "team":

3
schema.prisma
View file

@ -53,6 +53,7 @@ model LiteLLM_OrganizationTable {
updated_by String
litellm_budget_table LiteLLM_BudgetTable? @relation(fields: [budget_id], references: [budget_id])
teams LiteLLM_TeamTable[]
users LiteLLM_UserTable[]
}
// Model info for teams, just has model aliases for now.
@ -99,6 +100,7 @@ model LiteLLM_UserTable {
user_id String @id
user_alias String?
team_id String?
organization_id String?
teams String[] @default([])
user_role String?
max_budget Float?
@ -113,6 +115,7 @@ model LiteLLM_UserTable {
allowed_cache_controls String[] @default([])
model_spend Json @default("{}")
model_max_budget Json @default("{}")
litellm_organization_table LiteLLM_OrganizationTable? @relation(fields: [organization_id], references: [organization_id])
}
// Generate Tokens for Proxy