forked from phoenix/litellm-mirror
test(test_keys.py): separate streaming key info test from normal completion key info test
This commit is contained in:
parent
3aab719521
commit
bfe2faa454
2 changed files with 15 additions and 1 deletions
|
@ -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"
|
f"Model not in litellm model cost map. Add custom pricing - https://docs.litellm.ai/docs/proxy/custom_pricing"
|
||||||
)
|
)
|
||||||
except Exception as e:
|
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(
|
async def update_database(
|
||||||
|
|
|
@ -318,6 +318,17 @@ async def test_key_info_spend_values():
|
||||||
rounded_response_cost = round(response_cost, 8)
|
rounded_response_cost = round(response_cost, 8)
|
||||||
rounded_key_info_spend = round(key_info["info"]["spend"], 8)
|
rounded_key_info_spend = round(key_info["info"]["spend"], 8)
|
||||||
assert rounded_response_cost == rounded_key_info_spend
|
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
|
## streaming - azure
|
||||||
key_gen = await generate_key(session=session, i=0)
|
key_gen = await generate_key(session=session, i=0)
|
||||||
new_key = key_gen["key"]
|
new_key = key_gen["key"]
|
||||||
|
@ -332,6 +343,7 @@ async def test_key_info_spend_values():
|
||||||
)
|
)
|
||||||
response_cost = prompt_cost + completion_cost
|
response_cost = prompt_cost + completion_cost
|
||||||
await asyncio.sleep(5) # allow db log to be updated
|
await asyncio.sleep(5) # allow db log to be updated
|
||||||
|
print(f"new_key: {new_key}")
|
||||||
key_info = await get_key_info(
|
key_info = await get_key_info(
|
||||||
session=session, get_key=new_key, call_key=new_key
|
session=session, get_key=new_key, call_key=new_key
|
||||||
)
|
)
|
||||||
|
|
Loading…
Add table
Add a link
Reference in a new issue