phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-25 18:54:30 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

fix(proxy_server.py): enable spend tracking for team-based jwt auth

Browse source
This commit is contained in:
Krrish Dholakia 2024-03-28 20:16:07 -07:00
parent 792999d756
commit c15ba368e7
3 changed files with 25 additions and 5 deletions
Download patch file Download diff file Expand all files Collapse all files

7
litellm/proxy/_new_secret_config.yaml
View file

@ -7,4 +7,9 @@ model_list:
general_settings:
master_key: sk-1234
proxy_batch_write_at: 5
proxy_batch_write_at: 5 # 👈 Frequency of batch writing logs to server (in seconds)
enable_jwt_auth: True
litellm_jwtauth:
admin_jwt_scope: "litellm_proxy_admin"
team_jwt_scope: "litellm_team"
public_key_ttl: 600

12
litellm/proxy/proxy_server.py
View file

@ -293,7 +293,7 @@ litellm_proxy_admin_name = "default_user_id"
ui_access_mode: Literal["admin", "all"] = "all"
proxy_budget_rescheduler_min_time = 597
proxy_budget_rescheduler_max_time = 605
proxy_batch_write_at = 60 # in seconds
proxy_batch_write_at = 10 # in seconds
litellm_master_key_hash = None
disable_spend_logs = False
jwt_handler = JWTHandler()
@ -1161,7 +1161,7 @@ async def _PROXY_track_cost_callback(
verbose_proxy_logger.debug(
f"user_api_key {user_api_key}, prisma_client: {prisma_client}, custom_db_client: {custom_db_client}"
)
if user_api_key is not None:
if user_api_key is not None or user_id is not None or team_id is not None:
## UPDATE DATABASE
await update_database(
token=user_api_key,
@ -1182,7 +1182,9 @@ async def _PROXY_track_cost_callback(
response_cost=response_cost,
)
else:
raise Exception("User API key missing from custom callback.")
raise Exception(
"User API key and team id and user id missing from custom callback."
)
else:
if kwargs["stream"] != True or (
kwargs["stream"] == True and "complete_streaming_response" in kwargs
@ -1223,7 +1225,7 @@ async def update_database(
verbose_proxy_logger.info(
f"Enters prisma db call, response_cost: {response_cost}, token: {token}; user_id: {user_id}; team_id: {team_id}"
)
if isinstance(token, str) and token.startswith("sk-"):
if token is not None and isinstance(token, str) and token.startswith("sk-"):
hashed_token = hash_token(token=token)
else:
hashed_token = token
@ -1337,6 +1339,8 @@ async def update_database(
verbose_proxy_logger.debug(
f"adding spend to key db. Response cost: {response_cost}. Token: {hashed_token}."
)
if hashed_token is None:
return
if prisma_client is not None:
prisma_client.key_list_transactons[hashed_token] = (
response_cost

11
litellm/proxy/utils.py
View file

@ -1901,6 +1901,7 @@ async def update_spend(
spend_logs: list,
"""
n_retry_times = 3
verbose_proxy_logger.debug("INSIDE UPDATE SPEND")
### UPDATE USER TABLE ###
if len(prisma_client.user_list_transactons.keys()) > 0:
for i in range(n_retry_times + 1):
@ -1989,6 +1990,11 @@ async def update_spend(
raise e
### UPDATE TEAM TABLE ###
verbose_proxy_logger.debug(
"Team Spend transactions: {}".format(
len(prisma_client.team_list_transactons.keys())
)
)
if len(prisma_client.team_list_transactons.keys()) > 0:
for i in range(n_retry_times + 1):
try:
@ -2000,6 +2006,11 @@ async def update_spend(
team_id,
response_cost,
) in prisma_client.team_list_transactons.items():
verbose_proxy_logger.debug(
"Updating spend for team id={} by {}".format(
team_id, response_cost
)
)
batcher.litellm_teamtable.update_many( # 'update_many' prevents error from being raised if no row exists
where={"team_id": team_id},
data={"spend": {"increment": response_cost}},