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

feat - add validation for existing customers

Browse source
This commit is contained in:
Ishaan Jaff 2024-05-27 18:29:09 -07:00
parent cdf32ebf0e
commit 24f0b82755
1 changed files with 64 additions and 38 deletions
Download patch file Download diff file Expand all files Collapse all files

102
litellm/proxy/proxy_server.py
View file

@ -7833,51 +7833,77 @@ async def new_end_user(
status_code=500, status_code=500,
detail={"error": CommonProxyErrors.db_not_connected_error.value}, detail={"error": CommonProxyErrors.db_not_connected_error.value},
) )
try:
## VALIDATION ## ## VALIDATION ##
if data.default_model is not None: if data.default_model is not None:
if llm_router is None: if llm_router is None:
raise HTTPException( raise HTTPException(
status_code=422, detail={"error": CommonProxyErrors.no_llm_router.value} status_code=422,
) detail={"error": CommonProxyErrors.no_llm_router.value},
elif data.default_model not in llm_router.get_model_names(): )
raise HTTPException( elif data.default_model not in llm_router.get_model_names():
status_code=422, raise HTTPException(
detail={ status_code=422,
"error": "Default Model not on proxy. Configure via `/model/new` or config.yaml. Default_model={}, proxy_model_names={}".format( detail={
data.default_model, set(llm_router.get_model_names()) "error": "Default Model not on proxy. Configure via `/model/new` or config.yaml. Default_model={}, proxy_model_names={}".format(
) data.default_model, set(llm_router.get_model_names())
}, )
},
)
new_end_user_obj: Dict = {}
## CREATE BUDGET ## if set
if data.max_budget is not None:
budget_record = await prisma_client.db.litellm_budgettable.create(
data={
"max_budget": data.max_budget,
"created_by": user_api_key_dict.user_id or litellm_proxy_admin_name, # type: ignore
"updated_by": user_api_key_dict.user_id or litellm_proxy_admin_name,
}
) )
new_end_user_obj: Dict = {} new_end_user_obj["budget_id"] = budget_record.budget_id
elif data.budget_id is not None:
new_end_user_obj["budget_id"] = data.budget_id
## CREATE BUDGET ## if set _user_data = data.dict(exclude_none=True)
if data.max_budget is not None:
budget_record = await prisma_client.db.litellm_budgettable.create( for k, v in _user_data.items():
data={ if k != "max_budget" and k != "budget_id":
"max_budget": data.max_budget, new_end_user_obj[k] = v
"created_by": user_api_key_dict.user_id or litellm_proxy_admin_name, # type: ignore
"updated_by": user_api_key_dict.user_id or litellm_proxy_admin_name, ## WRITE TO DB ##
} end_user_record = await prisma_client.db.litellm_endusertable.create(
data=new_end_user_obj # type: ignore
) )
new_end_user_obj["budget_id"] = budget_record.budget_id return end_user_record
elif data.budget_id is not None: except Exception as e:
new_end_user_obj["budget_id"] = data.budget_id if "Unique constraint failed on the fields: (`user_id`)" in str(e):
raise ProxyException(
message=f"Customer already exists, passed user_id={data.user_id}. Please pass a new user_id.",
type="bad_request",
code=400,
param="user_id",
)
_user_data = data.dict(exclude_none=True) if isinstance(e, HTTPException):
raise ProxyException(
for k, v in _user_data.items(): message=getattr(e, "detail", f"Internal Server Error({str(e)})"),
if k != "max_budget" and k != "budget_id": type="internal_error",
new_end_user_obj[k] = v param=getattr(e, "param", "None"),
code=getattr(e, "status_code", status.HTTP_500_INTERNAL_SERVER_ERROR),
## WRITE TO DB ## )
end_user_record = await prisma_client.db.litellm_endusertable.create( elif isinstance(e, ProxyException):
data=new_end_user_obj # type: ignore raise e
) raise ProxyException(
message="Internal Server Error, " + str(e),
return end_user_record type="internal_error",
param=getattr(e, "param", "None"),
code=status.HTTP_500_INTERNAL_SERVER_ERROR,
)
@router.get( @router.get(