scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-08 22:21:35 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 78

fix(cost_calculator.py): handle openai usage pydantic object

Browse source 
Fixes https://github.com/BerriAI/litellm/issues/5165
This commit is contained in:
Krrish Dholakia 2024-08-12 15:45:21 -07:00
parent fdd9a07051
commit 22e2840daa
2 changed files with 15 additions and 2 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
litellm/cost_calculator.py
View file

@ -490,6 +490,10 @@ def completion_cost(
isinstance(completion_response, BaseModel)
or isinstance(completion_response, dict)
): # tts returns a custom class
if isinstance(completion_response, BaseModel) and not isinstance(
completion_response, litellm.Usage
):
completion_response = litellm.Usage(**completion_response.model_dump())
# get input/output tokens from completion_response
prompt_tokens = completion_response.get("usage", {}).get("prompt_tokens", 0)
completion_tokens = completion_response.get("usage", {}).get(

13
litellm/tests/test_completion_cost.py
View file

@ -918,9 +918,18 @@ def test_vertex_ai_llama_predict_cost():
assert predictive_cost == 0
def test_vertex_ai_mistral_predict_cost():
@pytest.mark.parametrize("usage", ["litellm_usage", "openai_usage"])
def test_vertex_ai_mistral_predict_cost(usage):
from litellm.types.utils import Choices, Message, ModelResponse, Usage
if usage == "litellm_usage":
response_usage = Usage(prompt_tokens=32, completion_tokens=55, total_tokens=87)
else:
from openai.types.completion_usage import CompletionUsage
response_usage = CompletionUsage(
prompt_tokens=32, completion_tokens=55, total_tokens=87
)
response_object = ModelResponse(
id="26c0ef045020429d9c5c9b078c01e564",
choices=[
@ -939,7 +948,7 @@ def test_vertex_ai_mistral_predict_cost():
model="vertex_ai/mistral-large",
object="chat.completion",
system_fingerprint=None,
usage=Usage(prompt_tokens=32, completion_tokens=55, total_tokens=87),
usage=response_usage,
)
model = "mistral-large@2407"
messages = [{"role": "user", "content": "Hey, hows it going???"}]