From 3e53742a4fa0e694dcebf6e137c2766ef7c19a17 Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Wed, 24 Jan 2024 11:10:13 -0800 Subject: [PATCH] (test) model in SpendLogs table --- litellm/tests/test_key_generate_prisma.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/litellm/tests/test_key_generate_prisma.py b/litellm/tests/test_key_generate_prisma.py index 08b5d1f93..fe5f8fa56 100644 --- a/litellm/tests/test_key_generate_prisma.py +++ b/litellm/tests/test_key_generate_prisma.py @@ -825,6 +825,7 @@ def test_call_with_key_over_budget_stream(prisma_client): ) await track_cost_callback( kwargs={ + "model": "sagemaker-chatgpt-v-2", "stream": True, "complete_streaming_response": resp, "litellm_params": { @@ -850,6 +851,7 @@ def test_call_with_key_over_budget_stream(prisma_client): assert spend_log.request_id == request_id assert spend_log.spend == float("5e-05") + assert spend_log.model == "sagemaker-chatgpt-v-2" # use generated key to auth in result = await user_api_key_auth(request=request, api_key=bearer_token)