test(test_cost_calc.py): fix test to handle llm api errors

This commit is contained in:
Krrish Dholakia 2024-12-24 16:49:02 -08:00
parent 78fe124c14
commit 48d5acf678

View file

@ -63,7 +63,7 @@ def test_run(model: str):
"""
Relevant issue - https://github.com/BerriAI/litellm/issues/4965
"""
# litellm.set_verbose = True
litellm.set_verbose = True
prompt = "Hi"
kwargs = dict(
model=model,
@ -84,6 +84,7 @@ def test_run(model: str):
print(f"Non-stream output: {non_stream_output}")
print(f"Non-stream usage : {response.usage}") # type: ignore
non_stream_usage = response.usage
try:
print(
f"Non-stream cost : {response._hidden_params['response_cost'] * 100:.4f}"
@ -95,6 +96,11 @@ def test_run(model: str):
response = router.completion(**kwargs, stream=True, stream_options={"include_usage": True}) # type: ignore
response = stream_chunk_builder(list(response), messages=kwargs["messages"]) # type: ignore
output = response.choices[0].message.content.replace("\n", "") # type: ignore
if response.usage.completion_tokens != non_stream_usage.completion_tokens:
pytest.skip(
"LLM API returning inconsistent usage"
) # handles transient openai errors
streaming_cost_calc = completion_cost(response) * 100
print(f"Stream output : {output}")