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

(feat) use user_notification table

Browse source
This commit is contained in:
ishaan-jaff 2024-02-19 16:53:40 -08:00
parent a13565dd95
commit 6985880506
4 changed files with 19 additions and 12 deletions
Download patch file Download diff file Expand all files Collapse all files

3
litellm/proxy/proxy_server.py
View file

@ -4149,9 +4149,10 @@ async def user_request_model(request: Request):
"models": new_models, "models": new_models,
"justification": justification, "justification": justification,
"user_id": user_id, "user_id": user_id,
"status": "pending",
"request_id": str(uuid.uuid4()), "request_id": str(uuid.uuid4()),
}, },
table_name="model_request", table_name="user_notification",
) )
return {"status": "success"} return {"status": "success"}
# update based on remaining passed in values # update based on remaining passed in values

3
litellm/proxy/schema.prisma
View file

@ -98,9 +98,10 @@ model LiteLLM_SpendLogs {
} }
// Beta - allow team members to request access to a model // Beta - allow team members to request access to a model
model LiteLLM_ModelRequests { model LiteLLM_UserNotifications {
request_id String @unique request_id String @unique
user_id String user_id String
models String[] models String[]
justification String justification String
status String // approved, disapproved, pending
} }

12
litellm/proxy/utils.py
View file

@ -698,7 +698,9 @@ class PrismaClient:
async def insert_data( async def insert_data(
self, self,
data: dict, data: dict,
table_name: Literal["user", "key", "config", "spend", "team", "model_request"], table_name: Literal[
"user", "key", "config", "spend", "team", "user_notification"
],
): ):
""" """
Add a key to the database. If it already exists, do nothing. Add a key to the database. If it already exists, do nothing.
@ -780,17 +782,19 @@ class PrismaClient:
) )
verbose_proxy_logger.info(f"Data Inserted into Spend Table") verbose_proxy_logger.info(f"Data Inserted into Spend Table")
return new_spend_row return new_spend_row
elif table_name == "model_request": elif table_name == "user_notification":
db_data = self.jsonify_object(data=data) db_data = self.jsonify_object(data=data)
new_model_request_row = await self.db.litellm_modelrequests.upsert( new_user_notification_row = (
await self.db.litellm_usernotifications.upsert(
where={"request_id": data["request_id"]}, where={"request_id": data["request_id"]},
data={ data={
"create": {**db_data}, # type: ignore "create": {**db_data}, # type: ignore
"update": {}, # don't do anything if it already exists "update": {}, # don't do anything if it already exists
}, },
) )
)
verbose_proxy_logger.info(f"Data Inserted into Model Request Table") verbose_proxy_logger.info(f"Data Inserted into Model Request Table")
return new_model_request_row return new_user_notification_row
except Exception as e: except Exception as e:
print_verbose(f"LiteLLM Prisma Client Exception: {e}") print_verbose(f"LiteLLM Prisma Client Exception: {e}")

3
schema.prisma
View file

@ -98,9 +98,10 @@ model LiteLLM_SpendLogs {
} }
// Beta - allow team members to request access to a model // Beta - allow team members to request access to a model
model LiteLLM_ModelRequests { model LiteLLM_UserNotifications {
request_id String @unique request_id String @unique
user_id String user_id String
models String[] models String[]
justification String justification String
status String // approved, disapproved, pending
} }