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

fix(proxy_server.py): increment cached global proxy spend object

Browse source
This commit is contained in:
Krrish Dholakia 2024-03-29 20:02:31 -07:00
parent 49e2624240
commit 9fc8b8a215
3 changed files with 25 additions and 9 deletions
Download patch file Download diff file Expand all files Collapse all files

15
litellm/tests/test_key_generate_prisma.py
View file

@ -409,18 +409,20 @@ def test_call_with_proxy_over_budget(prisma_client):
litellm_proxy_budget_name = f"litellm-proxy-budget-{time.time()}"
setattr(
litellm.proxy.proxy_server,
"litellm_proxy_budget_name",
"litellm_proxy_admin_name",
litellm_proxy_budget_name,
)
setattr(litellm, "max_budget", 0.00001)
from litellm.proxy.proxy_server import user_api_key_cache
user_api_key_cache.set_cache(
key="{}:spend".format(litellm_proxy_budget_name), value=0
)
setattr(litellm.proxy.proxy_server, "user_api_key_cache", user_api_key_cache)
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)
@ -472,6 +474,7 @@ def test_call_with_proxy_over_budget(prisma_client):
start_time=datetime.now(),
end_time=datetime.now(),
)
await asyncio.sleep(5)
# use generated key to auth in
result = await user_api_key_auth(request=request, api_key=bearer_token)