scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-07 08:26:10 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 36

Merge pull request #1603 from BerriAI/litellm_global_spend_updates

Browse source 
Litellm global spend updates
This commit is contained in:
Krish Dholakia 2024-01-24 20:20:27 -08:00 • committed by GitHub
commit fe592aa7ec
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
7 changed files with 356 additions and 50 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

3
.circleci/config.yml
View file

@ -152,7 +152,8 @@ jobs:
my-app:latest \
--config /app/config.yaml \
--port 4000 \
--num_workers 8
--num_workers 8 \
--debug
- run:
name: Install curl and dockerize
command: |

115
litellm/proxy/proxy_server.py
View file

@ -197,6 +197,7 @@ use_queue = False
health_check_interval = None
health_check_results = {}
queue: List = []
litellm_proxy_budget_name = "litellm-proxy-budget"
### INITIALIZE GLOBAL LOGGING OBJECT ###
proxy_logging_obj = ProxyLogging(user_api_key_cache=user_api_key_cache)
### REDIS QUEUE ###
@ -370,30 +371,62 @@ async def user_api_key_auth(
)
# Check 2. If user_id for this token is in budget
## Check 2.5 If global proxy is in budget
if valid_token.user_id is not None:
if prisma_client is not None:
user_id_information = await prisma_client.get_data(
user_id=valid_token.user_id, table_name="user"
user_id_list=[valid_token.user_id, litellm_proxy_budget_name],
table_name="user",
query_type="find_all",
)
if custom_db_client is not None:
user_id_information = await custom_db_client.get_data(
key=valid_token.user_id, table_name="user"
)
verbose_proxy_logger.debug(
f"user_id_information: {user_id_information}"
)
# Token exists, not expired now check if its in budget for the user
if valid_token.spend is not None and valid_token.user_id is not None:
user_max_budget = getattr(user_id_information, "max_budget", None)
user_current_spend = getattr(user_id_information, "spend", None)
if user_id_information is not None:
if isinstance(user_id_information, list):
## Check if user in budget
for _user in user_id_information:
if _user is None:
continue
assert isinstance(_user, dict)
# Token exists, not expired now check if its in budget for the user
user_max_budget = _user.get("max_budget", None)
user_current_spend = _user.get("spend", None)
if user_max_budget is not None and user_current_spend is not None:
if user_current_spend > user_max_budget:
raise Exception(
f"ExceededBudget: User {valid_token.user_id} has exceeded their budget. Current spend: {user_current_spend}; Max Budget: {user_max_budget}"
verbose_proxy_logger.debug(
f"user_max_budget: {user_max_budget}; user_current_spend: {user_current_spend}"
)
if (
user_max_budget is not None
and user_current_spend is not None
):
if user_current_spend > user_max_budget:
raise Exception(
f"ExceededBudget: User {valid_token.user_id} has exceeded their budget. Current spend: {user_current_spend}; Max Budget: {user_max_budget}"
)
else:
# Token exists, not expired now check if its in budget for the user
user_max_budget = getattr(
user_id_information, "max_budget", None
)
user_current_spend = getattr(user_id_information, "spend", None)
if (
user_max_budget is not None
and user_current_spend is not None
):
if user_current_spend > user_max_budget:
raise Exception(
f"ExceededBudget: User {valid_token.user_id} has exceeded their budget. Current spend: {user_current_spend}; Max Budget: {user_max_budget}"
)
# Check 3. If token is expired
if valid_token.expires is not None:
current_time = datetime.now(timezone.utc)
@ -642,29 +675,40 @@ async def update_database(
### UPDATE USER SPEND ###
async def _update_user_db():
if user_id is None:
return
if prisma_client is not None:
existing_spend_obj = await prisma_client.get_data(user_id=user_id)
elif custom_db_client is not None:
existing_spend_obj = await custom_db_client.get_data(
key=user_id, table_name="user"
)
if existing_spend_obj is None:
existing_spend = 0
else:
existing_spend = existing_spend_obj.spend
"""
- Update that user's row
- Update litellm-proxy-budget row (global proxy spend)
"""
user_ids = [user_id, litellm_proxy_budget_name]
data_list = []
for id in user_ids:
if id is None:
continue
if prisma_client is not None:
existing_spend_obj = await prisma_client.get_data(user_id=id)
elif custom_db_client is not None:
existing_spend_obj = await custom_db_client.get_data(
key=id, table_name="user"
)
if existing_spend_obj is None:
existing_spend = 0
existing_spend = LiteLLM_UserTable(user_id=id, spend=0)
else:
existing_spend = existing_spend_obj.spend
# Calculate the new cost by adding the existing cost and response_cost
new_spend = existing_spend + response_cost
# Calculate the new cost by adding the existing cost and response_cost
existing_spend_obj.spend = existing_spend + response_cost
verbose_proxy_logger.debug(f"new cost: {existing_spend_obj.spend}")
data_list.append(existing_spend_obj)
verbose_proxy_logger.debug(f"new cost: {new_spend}")
# Update the cost column for the given user id
if prisma_client is not None:
await prisma_client.update_data(
user_id=user_id, data={"spend": new_spend}
data_list=data_list, query_type="update_many", table_name="user"
)
elif custom_db_client is not None:
elif custom_db_client is not None and user_id is not None:
new_spend = data_list[0].spend
await custom_db_client.update_data(
key=user_id, value={"spend": new_spend}, table_name="user"
)
@ -1161,6 +1205,7 @@ async def generate_key_helper_fn(
tpm_limit: Optional[int] = None,
rpm_limit: Optional[int] = None,
query_type: Literal["insert_data", "update_data"] = "insert_data",
update_key_values: Optional[dict] = None,
):
global prisma_client, custom_db_client
@ -1261,7 +1306,9 @@ async def generate_key_helper_fn(
key_data["models"] = user_row.models
elif query_type == "update_data":
user_row = await prisma_client.update_data(
data=user_data, table_name="user"
data=user_data,
table_name="user",
update_key_values=update_key_values,
)
## CREATE KEY
@ -1567,7 +1614,13 @@ async def startup_event():
if prisma_client is not None and master_key is not None:
# add master key to db
await generate_key_helper_fn(
duration=None, models=[], aliases={}, config={}, spend=0, token=master_key
duration=None,
models=[],
aliases={},
config={},
spend=0,
token=master_key,
user_id="default_user_id",
)
if (
@ -1577,7 +1630,7 @@ async def startup_event():
):
# add proxy budget to db in the user table
await generate_key_helper_fn(
user_id="litellm-proxy-budget",
user_id=litellm_proxy_budget_name,
duration=None,
models=[],
aliases={},
@ -1586,6 +1639,10 @@ async def startup_event():
max_budget=litellm.max_budget,
budget_duration=litellm.budget_duration,
query_type="update_data",
update_key_values={
"max_budget": litellm.max_budget,
"budget_duration": litellm.budget_duration,
},
)
verbose_proxy_logger.debug(

42
litellm/proxy/utils.py
View file

@ -392,6 +392,7 @@ class PrismaClient:
self,
token: Optional[str] = None,
user_id: Optional[str] = None,
user_id_list: Optional[list] = None,
key_val: Optional[dict] = None,
table_name: Optional[Literal["user", "key", "config", "spend"]] = None,
query_type: Literal["find_unique", "find_all"] = "find_unique",
@ -473,6 +474,17 @@ class PrismaClient:
"budget_reset_at": {"lt": reset_at},
}
)
elif query_type == "find_all" and user_id_list is not None:
user_id_values = str(tuple(user_id_list))
sql_query = f"""
SELECT *
FROM "LiteLLM_UserTable"
WHERE "user_id" IN {user_id_values}
"""
# Execute the raw query
# The asterisk before `user_id_list` unpacks the list into separate arguments
response = await self.db.query_raw(sql_query)
elif query_type == "find_all":
response = await self.db.litellm_usertable.find_many( # type: ignore
order={"spend": "desc"},
@ -616,6 +628,7 @@ class PrismaClient:
user_id: Optional[str] = None,
query_type: Literal["update", "update_many"] = "update",
table_name: Optional[Literal["user", "key", "config", "spend"]] = None,
update_key_values: Optional[dict] = None,
):
"""
Update existing data
@ -642,29 +655,23 @@ class PrismaClient:
user_id is not None
or (table_name is not None and table_name == "user")
and query_type == "update"
and update_key_values is not None
):
"""
If data['spend'] + data['user'], update the user table with spend info as well
"""
if user_id is None:
user_id = db_data["user_id"]
update_user_row = await self.db.litellm_usertable.update(
update_user_row = await self.db.litellm_usertable.upsert(
where={"user_id": user_id}, # type: ignore
data={**db_data}, # type: ignore
data={
"create": {**db_data}, # type: ignore
"update": {
**update_key_values # type: ignore
}, # just update user-specified values, if it already exists
},
)
if update_user_row is None:
# if the provided user does not exist, STILL Track this!
# make a new user with {"user_id": user_id, "spend": data['spend']}
db_data["user_id"] = user_id
update_user_row = await self.db.litellm_usertable.upsert(
where={"user_id": user_id}, # type: ignore
data={
"create": {**db_data}, # type: ignore
"update": {}, # don't do anything if it already exists
},
)
print_verbose(
verbose_proxy_logger.info(
"\033[91m"
+ f"DB User Table - update succeeded {update_user_row}"
+ "\033[0m"
@ -708,6 +715,7 @@ class PrismaClient:
Batch write update queries
"""
batcher = self.db.batch_()
verbose_proxy_logger.debug(f"data list for user table: {data_list}")
for idx, user in enumerate(data_list):
try:
data_json = self.jsonify_object(data=user.model_dump())
@ -718,8 +726,8 @@ class PrismaClient:
data={**data_json}, # type: ignore
)
await batcher.commit()
print_verbose(
"\033[91m" + f"DB User Table update succeeded" + "\033[0m"
verbose_proxy_logger.info(
"\033[91m" + f"DB User Table Batch update succeeded" + "\033[0m"
)
except Exception as e:
asyncio.create_task(

174
litellm/tests/test_key_generate_prisma.py
View file

@ -24,7 +24,7 @@ from fastapi import Request
from datetime import datetime
load_dotenv()
import os, io
import os, io, time
# this file is to test litellm/proxy
@ -83,6 +83,7 @@ def prisma_client():
# Reset litellm.proxy.proxy_server.prisma_client to None
litellm.proxy.proxy_server.custom_db_client = None
litellm.proxy.proxy_server.litellm_proxy_budget_name = "litellm-proxy-budget"
return prisma_client
@ -282,6 +283,90 @@ def test_call_with_user_over_budget(prisma_client):
print(vars(e))
def test_call_with_proxy_over_budget(prisma_client):
# 5.1 Make a call with a proxy over budget, expect to fail
setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client)
setattr(litellm.proxy.proxy_server, "master_key", "sk-1234")
litellm_proxy_budget_name = f"litellm-proxy-budget-{time.time()}"
setattr(
litellm.proxy.proxy_server,
"litellm_proxy_budget_name",
litellm_proxy_budget_name,
)
try:
async def test():
await litellm.proxy.proxy_server.prisma_client.connect()
## CREATE PROXY + USER BUDGET ##
request = NewUserRequest(
max_budget=0.00001, user_id=litellm_proxy_budget_name
)
await new_user(request)
request = NewUserRequest()
key = await new_user(request)
print(key)
generated_key = key.key
user_id = key.user_id
bearer_token = "Bearer " + generated_key
request = Request(scope={"type": "http"})
request._url = URL(url="/chat/completions")
# use generated key to auth in
result = await user_api_key_auth(request=request, api_key=bearer_token)
print("result from user auth with new key", result)
# update spend using track_cost callback, make 2nd request, it should fail
from litellm.proxy.proxy_server import track_cost_callback
from litellm import ModelResponse, Choices, Message, Usage
resp = ModelResponse(
id="chatcmpl-e41836bb-bb8b-4df2-8e70-8f3e160155ac",
choices=[
Choices(
finish_reason=None,
index=0,
message=Message(
content=" Sure! Here is a short poem about the sky:\n\nA canvas of blue, a",
role="assistant",
),
)
],
model="gpt-35-turbo", # azure always has model written like this
usage=Usage(prompt_tokens=210, completion_tokens=200, total_tokens=410),
)
await track_cost_callback(
kwargs={
"stream": False,
"litellm_params": {
"metadata": {
"user_api_key": generated_key,
"user_api_key_user_id": user_id,
}
},
"response_cost": 0.00002,
},
completion_response=resp,
start_time=datetime.now(),
end_time=datetime.now(),
)
# use generated key to auth in
result = await user_api_key_auth(request=request, api_key=bearer_token)
print("result from user auth with new key", result)
pytest.fail(f"This should have failed!. They key crossed it's budget")
asyncio.run(test())
except Exception as e:
if hasattr(e, "message"):
error_detail = e.message
else:
error_detail = traceback.format_exc()
assert "Authentication Error, ExceededBudget:" in error_detail
print(vars(e))
def test_call_with_user_over_budget_stream(prisma_client):
# 6. Make a call with a key over budget, expect to fail
setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client)
@ -358,6 +443,93 @@ def test_call_with_user_over_budget_stream(prisma_client):
print(vars(e))
def test_call_with_proxy_over_budget_stream(prisma_client):
# 6.1 Make a call with a global proxy over budget, expect to fail
setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client)
setattr(litellm.proxy.proxy_server, "master_key", "sk-1234")
litellm_proxy_budget_name = f"litellm-proxy-budget-{time.time()}"
setattr(
litellm.proxy.proxy_server,
"litellm_proxy_budget_name",
litellm_proxy_budget_name,
)
from litellm._logging import verbose_proxy_logger
import logging
litellm.set_verbose = True
verbose_proxy_logger.setLevel(logging.DEBUG)
try:
async def test():
await litellm.proxy.proxy_server.prisma_client.connect()
## CREATE PROXY + USER BUDGET ##
request = NewUserRequest(
max_budget=0.00001, user_id=litellm_proxy_budget_name
)
await new_user(request)
request = NewUserRequest()
key = await new_user(request)
print(key)
generated_key = key.key
user_id = key.user_id
bearer_token = "Bearer " + generated_key
request = Request(scope={"type": "http"})
request._url = URL(url="/chat/completions")
# use generated key to auth in
result = await user_api_key_auth(request=request, api_key=bearer_token)
print("result from user auth with new key", result)
# update spend using track_cost callback, make 2nd request, it should fail
from litellm.proxy.proxy_server import track_cost_callback
from litellm import ModelResponse, Choices, Message, Usage
resp = ModelResponse(
id="chatcmpl-e41836bb-bb8b-4df2-8e70-8f3e160155ac",
choices=[
Choices(
finish_reason=None,
index=0,
message=Message(
content=" Sure! Here is a short poem about the sky:\n\nA canvas of blue, a",
role="assistant",
),
)
],
model="gpt-35-turbo", # azure always has model written like this
usage=Usage(prompt_tokens=210, completion_tokens=200, total_tokens=410),
)
await track_cost_callback(
kwargs={
"stream": True,
"complete_streaming_response": resp,
"litellm_params": {
"metadata": {
"user_api_key": generated_key,
"user_api_key_user_id": user_id,
}
},
"response_cost": 0.00002,
},
completion_response=ModelResponse(),
start_time=datetime.now(),
end_time=datetime.now(),
)
# use generated key to auth in
result = await user_api_key_auth(request=request, api_key=bearer_token)
print("result from user auth with new key", result)
pytest.fail(f"This should have failed!. They key crossed it's budget")
asyncio.run(test())
except Exception as e:
error_detail = e.message
assert "Authentication Error, ExceededBudget:" in error_detail
print(vars(e))
def test_generate_and_call_with_valid_key_never_expires(prisma_client):
# 7. Make a call with an key that never expires, expect to pass

7
litellm/utils.py
View file

@ -1090,7 +1090,12 @@ class Logging:
else: # streaming chunks + image gen.
self.model_call_details["response_cost"] = None
if litellm.max_budget and self.stream:
if (
litellm.max_budget
and self.stream
and result is not None
and "content" in result
):
time_diff = (end_time - start_time).total_seconds()
float_diff = float(time_diff)
litellm._current_cost += litellm.completion_cost(

2
proxy_server_config.yaml
View file

@ -41,6 +41,8 @@ model_list:
litellm_settings:
drop_params: True
max_budget: 100
budget_duration: 30d
general_settings:
master_key: sk-1234 # [OPTIONAL] Only use this if you to require all calls to contain this key (Authorization: Bearer sk-1234)
# database_url: "postgresql://<user>:<password>@<host>:<port>/<dbname>" # [OPTIONAL] use for token-based auth to proxy

63
tests/test_users.py
View file

@ -4,6 +4,7 @@ import pytest
import asyncio
import aiohttp
import time
from openai import AsyncOpenAI
async def new_user(session, i, user_id=None, budget=None, budget_duration=None):
@ -105,7 +106,7 @@ async def test_user_update():
@pytest.mark.asyncio
async def test_users_with_budgets():
async def test_users_budgets_reset():
"""
- Create key with budget and 5s duration
- Get 'reset_at' value
@ -128,3 +129,63 @@ async def test_users_with_budgets():
)
reset_at_new_value = user_info["user_info"]["budget_reset_at"]
assert reset_at_init_value != reset_at_new_value
async def chat_completion(session, key, model="gpt-4"):
client = AsyncOpenAI(api_key=key, base_url="http://0.0.0.0:4000")
messages = [
{"role": "system", "content": "You are a helpful assistant"},
{"role": "user", "content": f"Hello! {time.time()}"},
]
data = {
"model": model,
"messages": messages,
}
response = await client.chat.completions.create(**data)
async def chat_completion_streaming(session, key, model="gpt-4"):
client = AsyncOpenAI(api_key=key, base_url="http://0.0.0.0:4000")
messages = [
{"role": "system", "content": "You are a helpful assistant"},
{"role": "user", "content": f"Hello! {time.time()}"},
]
data = {"model": model, "messages": messages, "stream": True}
response = await client.chat.completions.create(**data)
async for chunk in response:
continue
@pytest.mark.asyncio
async def test_global_proxy_budget_update():
"""
- Get proxy current spend
- Make chat completion call (normal)
- Assert spend increased
- Make chat completion call (streaming)
- Assert spend increased
"""
get_user = f"litellm-proxy-budget"
async with aiohttp.ClientSession() as session:
user_info = await get_user_info(
session=session, get_user=get_user, call_user="sk-1234"
)
original_spend = user_info["user_info"]["spend"]
await chat_completion(session=session, key="sk-1234")
await asyncio.sleep(5) # let db update
user_info = await get_user_info(
session=session, get_user=get_user, call_user="sk-1234"
)
new_spend = user_info["user_info"]["spend"]
print(f"new_spend: {new_spend}; original_spend: {original_spend}")
assert new_spend > original_spend
await chat_completion_streaming(session=session, key="sk-1234")
await asyncio.sleep(5) # let db update
user_info = await get_user_info(
session=session, get_user=get_user, call_user="sk-1234"
)
new_new_spend = user_info["user_info"]["spend"]
print(f"new_spend: {new_spend}; original_spend: {original_spend}")
assert new_new_spend > new_spend