phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-26 03:04:13 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

fix(main.py): fix pydantic warning for usage dict

Browse source
This commit is contained in:
Krrish Dholakia 2023-12-02 20:02:55 -08:00
parent f0d8a87c48
commit 6c0eec4ff4
2 changed files with 29 additions and 32 deletions
Download patch file Download diff file Expand all files Collapse all files

9
litellm/main.py
View file

@ -30,7 +30,8 @@ from litellm.utils import (
get_api_key,
mock_completion_streaming_obj,
convert_to_model_response_object,
token_counter
token_counter,
Usage
)
from .llms import (
anthropic,
@ -1288,11 +1289,7 @@ def completion(
model_response["model"] = "ollama/" + model
prompt_tokens = len(encoding.encode(prompt)) # type: ignore
completion_tokens = len(encoding.encode(response_string))
model_response["usage"] = {
"prompt_tokens": prompt_tokens,
"completion_tokens": completion_tokens,
"total_tokens": prompt_tokens + completion_tokens,
}
model_response["usage"] = Usage(prompt_tokens=prompt_tokens, completion_tokens=completion_tokens, total_tokens=prompt_tokens + completion_tokens)
response = model_response
elif (
custom_llm_provider == "baseten"