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

feat(proxy_server.py): support model access groups

Browse source
This commit is contained in:
Krrish Dholakia 2024-01-17 15:45:31 -08:00
parent 70c4227f5e
commit 98b83fa780
6 changed files with 56 additions and 19 deletions
Download patch file Download diff file Expand all files Collapse all files

3
litellm/proxy/_types.py
View file

@ -307,9 +307,12 @@ class LiteLLM_UserTable(LiteLLMBase):
max_budget: Optional[float]
spend: float = 0.0
user_email: Optional[str]
models: list = []
@root_validator(pre=True)
def set_model_info(cls, values):
if values.get("spend") is None:
values.update({"spend": 0.0})
if values.get("models") is None:
values.update({"models", []})
return values

2
litellm/proxy/db/dynamo_db.py
View file

@ -171,7 +171,7 @@ class DynamoDBWrapper(CustomDB):
if isinstance(v, datetime):
value[k] = v.isoformat()
await table.put_item(item=value)
return await table.put_item(item=value)
async def get_data(self, key: str, table_name: Literal["user", "key", "config"]):
from aiodynamo.client import Client

48
litellm/proxy/proxy_server.py
View file

@ -325,7 +325,28 @@ async def user_api_key_auth(
model = data.get("model", None)
if model in litellm.model_alias_map:
model = litellm.model_alias_map[model]
if model and model not in valid_token.models:
## check if model in allowed model names
verbose_proxy_logger.debug(
f"LLM Model List pre access group check: {llm_model_list}"
)
access_groups = []
for m in llm_model_list:
for group in m.get("model_info", {}).get("access_groups", []):
access_groups.append((m["model_name"], group))
allowed_models = valid_token.models
if (
len(access_groups) > 0
): # check if token contains any model access groups
for m in valid_token.models:
for model_name, group in access_groups:
if m == group:
allowed_models.append(model_name)
verbose_proxy_logger.debug(
f"model: {model}; allowed_models: {allowed_models}"
)
if model is not None and model not in allowed_models:
raise ValueError(
f"API Key not allowed to access model. This token can only access models={valid_token.models}. Tried to access {model}"
)
@ -1057,6 +1078,7 @@ async def generate_key_helper_fn(
"user_email": user_email,
"user_id": user_id,
"spend": spend,
"models": models,
}
key_data = {
"token": token,
@ -1070,14 +1092,28 @@ async def generate_key_helper_fn(
"metadata": metadata_json,
}
if prisma_client is not None:
verification_token_data = dict(key_data)
verification_token_data.update(user_data)
verbose_proxy_logger.debug("PrismaClient: Before Insert Data")
await prisma_client.insert_data(data=verification_token_data)
## CREATE USER (If necessary)
verbose_proxy_logger.debug(f"CustomDBClient: Creating User={user_data}")
user_row = await prisma_client.insert_data(
data=user_data, table_name="user"
)
## use default user model list if no key-specific model list provided
if len(user_row.models) > 0 and len(key_data["models"]) == 0: # type: ignore
key_data["models"] = user_row.models
## CREATE KEY
verbose_proxy_logger.debug(f"CustomDBClient: Creating Key={key_data}")
await prisma_client.insert_data(data=key_data, table_name="key")
elif custom_db_client is not None:
## CREATE USER (If necessary)
verbose_proxy_logger.debug(f"CustomDBClient: Creating User={user_data}")
await custom_db_client.insert_data(value=user_data, table_name="user")
user_row = await custom_db_client.insert_data(
value=user_data, table_name="user"
)
## use default user model list if no key-specific model list provided
if len(user_row["models"]) > 0 and len(key_data["models"]) == 0: # type: ignore
key_data["models"] = user_row["models"]
## CREATE KEY
verbose_proxy_logger.debug(f"CustomDBClient: Creating Key={key_data}")
await custom_db_client.insert_data(value=key_data, table_name="key")

3
litellm/proxy/schema.prisma
View file

@ -12,6 +12,7 @@ model LiteLLM_UserTable {
max_budget Float?
spend Float @default(0.0)
user_email String?
models String[] @default([])
}
// required for token gen
@ -19,7 +20,7 @@ model LiteLLM_VerificationToken {
token String @unique
spend Float @default(0.0)
expires DateTime?
models String[]
models String[] @default([])
aliases Json @default("{}")
config Json @default("{}")
user_id String?

18
litellm/proxy/utils.py
View file

@ -409,19 +409,17 @@ class PrismaClient:
on_backoff=on_backoff, # specifying the function to call on backoff
)
async def insert_data(
self, data: dict, table_name: Literal["user+key", "config"] = "user+key"
self, data: dict, table_name: Literal["user", "key", "config"]
):
"""
Add a key to the database. If it already exists, do nothing.
"""
try:
if table_name == "user+key":
if table_name == "key":
token = data["token"]
hashed_token = self.hash_token(token=token)
db_data = self.jsonify_object(data=data)
db_data["token"] = hashed_token
max_budget = db_data.pop("max_budget", None)
user_email = db_data.pop("user_email", None)
print_verbose(
"PrismaClient: Before upsert into litellm_verificationtoken"
)
@ -434,19 +432,17 @@ class PrismaClient:
"update": {}, # don't do anything if it already exists
},
)
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": {
"user_id": data["user_id"],
"max_budget": max_budget,
"user_email": user_email,
},
"create": {**db_data}, # type: ignore
"update": {}, # don't do anything if it already exists
},
)
return new_verification_token
return new_user_row
elif table_name == "config":
"""
For each param,

1
schema.prisma
View file

@ -12,6 +12,7 @@ model LiteLLM_UserTable {
max_budget Float?
spend Float @default(0.0)
user_email String?
models String[]
}
// required for token gen