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

fix(utils.py): fix streaming cost tracking

Browse source
This commit is contained in:
Krrish Dholakia 2024-01-23 14:39:35 -08:00
parent 7eb96e46a4
commit afada01ffc
4 changed files with 39 additions and 13 deletions
Download patch file Download diff file Expand all files Collapse all files

32
litellm/utils.py
View file

@ -1067,10 +1067,15 @@ class Logging:
## if model in model cost map - log the response cost
## else set cost to None
verbose_logger.debug(f"Model={self.model}; result={result}")
if result is not None and (
isinstance(result, ModelResponse)
or isinstance(result, EmbeddingResponse)
):
verbose_logger.debug(f"self.stream: {self.stream}")
if (
result is not None
and (
isinstance(result, ModelResponse)
or isinstance(result, EmbeddingResponse)
)
and self.stream != True
): # handle streaming separately
try:
self.model_call_details["response_cost"] = litellm.completion_cost(
completion_response=result,
@ -1125,7 +1130,7 @@ class Logging:
else:
self.sync_streaming_chunks.append(result)
if complete_streaming_response:
if complete_streaming_response is not None:
verbose_logger.debug(
f"Logging Details LiteLLM-Success Call streaming complete"
)
@ -1418,11 +1423,23 @@ class Logging:
complete_streaming_response = None
else:
self.streaming_chunks.append(result)
if complete_streaming_response:
if complete_streaming_response is not None:
print_verbose("Async success callbacks: Got a complete streaming response")
self.model_call_details[
"complete_streaming_response"
] = complete_streaming_response
try:
self.model_call_details["response_cost"] = litellm.completion_cost(
completion_response=complete_streaming_response,
)
verbose_logger.debug(
f"Model={self.model}; cost={self.model_call_details['response_cost']}"
)
except litellm.NotFoundError as e:
verbose_logger.debug(
f"Model={self.model} not found in completion cost map."
)
self.model_call_details["response_cost"] = None
for callback in litellm._async_success_callback:
try:
@ -2867,6 +2884,9 @@ def cost_per_token(
if model in model_cost_ref:
verbose_logger.debug(f"Success: model={model} in model_cost_map")
verbose_logger.debug(
f"prompt_tokens={prompt_tokens}; completion_tokens={completion_tokens}"
)
if (
model_cost_ref[model].get("input_cost_per_token", None) is not None
and model_cost_ref[model].get("output_cost_per_token", None) is not None