From bfe2faa45455e30b0d8d1652680cabd874a46aeb Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Fri, 2 Feb 2024 18:57:19 -0800 Subject: [PATCH] test(test_keys.py): separate streaming key info test from normal completion key info test --- litellm/proxy/proxy_server.py | 4 +++- tests/test_keys.py | 12 ++++++++++++ 2 files changed, 15 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 925f8dfc0..6af2a9d48 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -743,7 +743,9 @@ async def _PROXY_track_cost_callback( f"Model not in litellm model cost map. Add custom pricing - https://docs.litellm.ai/docs/proxy/custom_pricing" ) except Exception as e: - verbose_proxy_logger.debug(f"error in tracking cost callback - {str(e)}") + verbose_proxy_logger.debug( + f"error in tracking cost callback - {traceback.format_exc}" + ) async def update_database( diff --git a/tests/test_keys.py b/tests/test_keys.py index ee8789f38..52eb7ad24 100644 --- a/tests/test_keys.py +++ b/tests/test_keys.py @@ -318,6 +318,17 @@ async def test_key_info_spend_values(): rounded_response_cost = round(response_cost, 8) rounded_key_info_spend = round(key_info["info"]["spend"], 8) assert rounded_response_cost == rounded_key_info_spend + + +@pytest.mark.asyncio +async def test_key_info_spend_values_streaming(): + """ + Test to ensure spend is correctly calculated. + - create key + - make completion call + - assert cost is expected value + """ + async with aiohttp.ClientSession() as session: ## streaming - azure key_gen = await generate_key(session=session, i=0) new_key = key_gen["key"] @@ -332,6 +343,7 @@ async def test_key_info_spend_values(): ) response_cost = prompt_cost + completion_cost await asyncio.sleep(5) # allow db log to be updated + print(f"new_key: {new_key}") key_info = await get_key_info( session=session, get_key=new_key, call_key=new_key )