Skip to content

Commit

Permalink
test(test_keys.py): separate streaming key info test from normal comp…
Browse files Browse the repository at this point in the history
…letion key info test
  • Loading branch information
krrishdholakia committed Feb 3, 2024
1 parent 2d59331 commit 8988353
Show file tree
Hide file tree
Showing 2 changed files with 15 additions and 1 deletion.
4 changes: 3 additions & 1 deletion litellm/proxy/proxy_server.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down
12 changes: 12 additions & 0 deletions tests/test_keys.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"]
Expand All @@ -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
)
Expand Down

0 comments on commit 8988353

Please sign in to comment.