diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 5d9661c2b8..295c809416 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -380,6 +380,11 @@ async def user_api_key_auth( # 3. If 'user' passed to /chat/completions, /embeddings endpoint is in budget # 4. If token is expired # 5. If token spend is under Budget for the token + # 6. If token spend per model is under budget per model + + request_data = await _read_request_body( + request=request + ) # request data, used across all checks. Making this easily available # Check 1. If token can call model litellm.model_alias_map = valid_token.aliases @@ -454,7 +459,6 @@ async def user_api_key_auth( if ( litellm.max_user_budget is not None ): # Check if 'user' passed in /chat/completions is in budget, only checked if litellm.max_user_budget is set - request_data = await _read_request_body(request=request) user_passed_to_chat_completions = request_data.get("user", None) if user_passed_to_chat_completions is not None: user_id_list.append(user_passed_to_chat_completions) @@ -587,6 +591,25 @@ async def user_api_key_auth( f"ExceededTokenBudget: Current spend for token: {valid_token.spend}; Max Budget for Token: {valid_token.max_budget}" ) + # Check 5. Token Model Spend is under Model budget + max_budget_per_model = valid_token.model_max_budget + spend_per_model = valid_token.model_spend + + if max_budget_per_model is not None and spend_per_model is not None: + current_model = request_data.get("model") + if current_model is not None: + current_model_spend = spend_per_model.get(current_model, None) + current_model_budget = max_budget_per_model.get(current_model, None) + + if ( + current_model_spend is not None + and current_model_budget is not None + ): + if current_model_spend > current_model_budget: + raise Exception( + f"ExceededModelBudget: Current spend for model: {current_model_spend}; Max Budget for Model: {current_model_budget}" + ) + # Token passed all checks api_key = valid_token.token diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 51e7c3c15c..1cdacf76fb 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -1379,19 +1379,22 @@ async def _read_request_body(request): """ import ast, json - request_data = {} - if request is None: - return request_data - body = await request.body() - - if body == b"" or body is None: - return request_data - body_str = body.decode() try: - request_data = ast.literal_eval(body_str) + request_data = {} + if request is None: + return request_data + body = await request.body() + + if body == b"" or body is None: + return request_data + body_str = body.decode() + try: + request_data = ast.literal_eval(body_str) + except: + request_data = json.loads(body_str) + return request_data except: - request_data = json.loads(body_str) - return request_data + return {} def _is_valid_team_configs(team_id=None, team_config=None, request_data=None): diff --git a/litellm/tests/test_key_generate_prisma.py b/litellm/tests/test_key_generate_prisma.py index 89d4f4d3eb..1f5507d3f9 100644 --- a/litellm/tests/test_key_generate_prisma.py +++ b/litellm/tests/test_key_generate_prisma.py @@ -1101,6 +1101,116 @@ def test_call_with_key_over_budget(prisma_client): print(vars(e)) +def test_call_with_key_over_model_budget(prisma_client): + # 12. Make a call with a key over budget, expect to fail + setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client) + setattr(litellm.proxy.proxy_server, "master_key", "sk-1234") + try: + + async def test(): + await litellm.proxy.proxy_server.prisma_client.connect() + + # set budget for chatgpt-v-2 to 0.000001, expect the next request to fail + request = GenerateKeyRequest( + max_budget=1000, + model_max_budget={ + "chatgpt-v-2": 0.000001, + }, + metadata={"user_api_key": 0.0001}, + ) + key = await generate_key_fn(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") + + async def return_body(): + return b'{"model": "chatgpt-v-2"}' + + request.body = return_body + + # 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 ( + _PROXY_track_cost_callback as track_cost_callback, + ) + from litellm import ModelResponse, Choices, Message, Usage + from litellm.caching import Cache + + litellm.cache = Cache() + import time + + request_id = f"chatcmpl-e41836bb-bb8b-4df2-8e70-8f3e160155ac{time.time()}" + + resp = ModelResponse( + id=request_id, + 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={ + "model": "chatgpt-v-2", + "stream": False, + "litellm_params": { + "metadata": { + "user_api_key": hash_token(generated_key), + "user_api_key_user_id": user_id, + } + }, + "response_cost": 0.00002, + }, + completion_response=resp, + start_time=datetime.now(), + end_time=datetime.now(), + ) + await asyncio.sleep(10) + # test spend_log was written and we can read it + spend_logs = await view_spend_logs(request_id=request_id) + + print("read spend logs", spend_logs) + assert len(spend_logs) == 1 + + spend_log = spend_logs[0] + + assert spend_log.request_id == request_id + assert spend_log.spend == float("2e-05") + assert spend_log.model == "chatgpt-v-2" + assert ( + spend_log.cache_key + == "a61ae14fe4a8b8014a61e6ae01a100c8bc6770ac37c293242afed954bc69207d" + ) + + # 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: + # print(f"Error - {str(e)}") + traceback.print_exc() + error_detail = e.message + assert "Authentication Error, ExceededModelBudget:" in error_detail + print(vars(e)) + + @pytest.mark.asyncio() async def test_call_with_key_never_over_budget(prisma_client): # Make a call with a key with budget=None, it should never fail