diff --git a/litellm/main.py b/litellm/main.py index 6c6edcc787c..67b935a55cb 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -5196,17 +5196,24 @@ def stream_chunk_builder( prompt_tokens = 0 completion_tokens = 0 for chunk in chunks: + usage_chunk: Optional[Usage] = None if "usage" in chunk: - if "prompt_tokens" in chunk["usage"]: - prompt_tokens = chunk["usage"].get("prompt_tokens", 0) or 0 - if "completion_tokens" in chunk["usage"]: - completion_tokens = chunk["usage"].get("completion_tokens", 0) or 0 + usage_chunk = chunk.usage + elif hasattr(chunk, "_hidden_params") and "usage" in chunk._hidden_params: + usage_chunk = chunk._hidden_params["usage"] + if usage_chunk is not None: + if "prompt_tokens" in usage_chunk: + prompt_tokens = usage_chunk.get("prompt_tokens", 0) or 0 + if "completion_tokens" in usage_chunk: + completion_tokens = usage_chunk.get("completion_tokens", 0) or 0 try: response["usage"]["prompt_tokens"] = prompt_tokens or token_counter( model=model, messages=messages ) - except: # don't allow this failing to block a complete streaming response from being returned - print_verbose(f"token_counter failed, assuming prompt tokens is 0") + except ( + Exception + ): # don't allow this failing to block a complete streaming response from being returned + print_verbose("token_counter failed, assuming prompt tokens is 0") response["usage"]["prompt_tokens"] = 0 response["usage"]["completion_tokens"] = completion_tokens or token_counter( model=model, diff --git a/litellm/tests/test_cost_calc.py b/litellm/tests/test_cost_calc.py new file mode 100644 index 00000000000..39d3c28fd75 --- /dev/null +++ b/litellm/tests/test_cost_calc.py @@ -0,0 +1,105 @@ +import os +import sys +import traceback + +from dotenv import load_dotenv + +load_dotenv() +import io +import os + +sys.path.insert( + 0, os.path.abspath("../..") +) # Adds the parent directory to the system path +from typing import Literal + +import pytest +from pydantic import BaseModel, ConfigDict + +import litellm +from litellm import Router, completion_cost, stream_chunk_builder + +models = [ + dict( + model_name="openai/gpt-3.5-turbo", + ), + dict( + model_name="anthropic/claude-3-haiku-20240307", + ), + dict( + model_name="together_ai/meta-llama/Llama-2-7b-chat-hf", + ), +] + +router = Router( + model_list=[ + { + "model_name": m["model_name"], + "litellm_params": { + "model": m.get("model", m["model_name"]), + }, + } + for m in models + ], + routing_strategy="simple-shuffle", + num_retries=3, + retry_after=1, + timeout=60.0, + allowed_fails=2, + cooldown_time=0, + debug_level="INFO", +) + + +@pytest.mark.parametrize( + "model", + [ + "openai/gpt-3.5-turbo", + "anthropic/claude-3-haiku-20240307", + "together_ai/meta-llama/Llama-2-7b-chat-hf", + ], +) +def test_run(model: str): + """ + Relevant issue - https://github.com/BerriAI/litellm/issues/4965 + """ + prompt = "Hi" + kwargs = dict( + model=model, + messages=[{"role": "user", "content": prompt}], + temperature=0.001, + top_p=0.001, + max_tokens=20, + input_cost_per_token=2, + output_cost_per_token=2, + ) + + print(f"--------- {model} ---------") + print(f"Prompt: {prompt}") + + response = router.completion(**kwargs) # type: ignore + non_stream_output = response.choices[0].message.content.replace("\n", "") # type: ignore + non_stream_cost_calc = response._hidden_params["response_cost"] * 100 + + print(f"Non-stream output: {non_stream_output}") + print(f"Non-stream usage : {response.usage}") # type: ignore + try: + print( + f"Non-stream cost : {response._hidden_params['response_cost'] * 100:.4f}" + ) + except TypeError: + print("Non-stream cost : NONE") + print(f"Non-stream cost : {completion_cost(response) * 100:.4f} (response)") + + response = router.completion(**kwargs, stream=True) # type: ignore + response = stream_chunk_builder(list(response), messages=kwargs["messages"]) # type: ignore + output = response.choices[0].message.content.replace("\n", "") # type: ignore + streaming_cost_calc = completion_cost(response) * 100 + print(f"Stream output : {output}") + + if output == non_stream_output: + # assert cost is the same + assert streaming_cost_calc == non_stream_cost_calc + print(f"Stream usage : {response.usage}") # type: ignore + print(f"Stream cost : {streaming_cost_calc} (response)") + print("") diff --git a/litellm/tests/test_streaming.py b/litellm/tests/test_streaming.py index bd2d889e37f..9c53d5cfbcf 100644 --- a/litellm/tests/test_streaming.py +++ b/litellm/tests/test_streaming.py @@ -3096,6 +3096,7 @@ def test_completion_claude_3_function_call_with_streaming(): elif idx == 1 and chunk.choices[0].finish_reason is None: validate_second_streaming_function_calling_chunk(chunk=chunk) elif chunk.choices[0].finish_reason is not None: # last chunk + assert "usage" in chunk._hidden_params validate_final_streaming_function_calling_chunk(chunk=chunk) idx += 1 # raise Exception("it worked!") diff --git a/litellm/utils.py b/litellm/utils.py index 778527fd4ab..90b70c87e74 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -8381,6 +8381,28 @@ def get_secret( ######## Streaming Class ############################ # wraps the completion stream to return the correct format for the model # replicate/anthropic/cohere + + +def calculate_total_usage(chunks: List[ModelResponse]) -> Usage: + """Assume most recent usage chunk has total usage uptil then.""" + prompt_tokens: int = 0 + completion_tokens: int = 0 + for chunk in chunks: + if "usage" in chunk: + if "prompt_tokens" in chunk["usage"]: + prompt_tokens = chunk["usage"].get("prompt_tokens", 0) or 0 + if "completion_tokens" in chunk["usage"]: + completion_tokens = chunk["usage"].get("completion_tokens", 0) or 0 + + returned_usage_chunk = Usage( + prompt_tokens=prompt_tokens, + completion_tokens=completion_tokens, + total_tokens=prompt_tokens + completion_tokens, + ) + + return returned_usage_chunk + + class CustomStreamWrapper: def __init__( self, @@ -9270,7 +9292,9 @@ class CustomStreamWrapper: verbose_logger.debug(traceback.format_exc()) return "" - def model_response_creator(self, chunk: Optional[dict] = None): + def model_response_creator( + self, chunk: Optional[dict] = None, hidden_params: Optional[dict] = None + ): _model = self.model _received_llm_provider = self.custom_llm_provider _logging_obj_llm_provider = self.logging_obj.model_call_details.get("custom_llm_provider", None) # type: ignore @@ -9284,6 +9308,7 @@ class CustomStreamWrapper: else: # pop model keyword chunk.pop("model", None) + model_response = ModelResponse( stream=True, model=_model, stream_options=self.stream_options, **chunk ) @@ -9293,6 +9318,8 @@ class CustomStreamWrapper: self.response_id = model_response.id # type: ignore if self.system_fingerprint is not None: model_response.system_fingerprint = self.system_fingerprint + if hidden_params is not None: + model_response._hidden_params = hidden_params model_response._hidden_params["custom_llm_provider"] = _logging_obj_llm_provider model_response._hidden_params["created_at"] = time.time() @@ -9347,11 +9374,7 @@ class CustomStreamWrapper: "finish_reason" ] - if ( - self.stream_options - and self.stream_options.get("include_usage", False) is True - and anthropic_response_obj["usage"] is not None - ): + if anthropic_response_obj["usage"] is not None: model_response.usage = litellm.Usage( prompt_tokens=anthropic_response_obj["usage"]["prompt_tokens"], completion_tokens=anthropic_response_obj["usage"][ @@ -9674,11 +9697,7 @@ class CustomStreamWrapper: print_verbose(f"completion obj content: {completion_obj['content']}") if response_obj["is_finished"]: self.received_finish_reason = response_obj["finish_reason"] - if ( - self.stream_options - and self.stream_options.get("include_usage", False) == True - and response_obj["usage"] is not None - ): + if response_obj["usage"] is not None: model_response.usage = litellm.Usage( prompt_tokens=response_obj["usage"].prompt_tokens, completion_tokens=response_obj["usage"].completion_tokens, @@ -9692,11 +9711,7 @@ class CustomStreamWrapper: print_verbose(f"completion obj content: {completion_obj['content']}") if response_obj["is_finished"]: self.received_finish_reason = response_obj["finish_reason"] - if ( - self.stream_options - and self.stream_options.get("include_usage", False) == True - and response_obj["usage"] is not None - ): + if response_obj["usage"] is not None: model_response.usage = litellm.Usage( prompt_tokens=response_obj["usage"].prompt_tokens, completion_tokens=response_obj["usage"].completion_tokens, @@ -9764,16 +9779,26 @@ class CustomStreamWrapper: if response_obj["logprobs"] is not None: model_response.choices[0].logprobs = response_obj["logprobs"] - if ( - self.stream_options is not None - and self.stream_options["include_usage"] == True - and response_obj["usage"] is not None - ): - model_response.usage = litellm.Usage( - prompt_tokens=response_obj["usage"].prompt_tokens, - completion_tokens=response_obj["usage"].completion_tokens, - total_tokens=response_obj["usage"].total_tokens, - ) + if response_obj["usage"] is not None: + if isinstance(response_obj["usage"], dict): + model_response.usage = litellm.Usage( + prompt_tokens=response_obj["usage"].get( + "prompt_tokens", None + ) + or None, + completion_tokens=response_obj["usage"].get( + "completion_tokens", None + ) + or None, + total_tokens=response_obj["usage"].get("total_tokens", None) + or None, + ) + elif isinstance(response_obj["usage"], BaseModel): + model_response.usage = litellm.Usage( + prompt_tokens=response_obj["usage"].prompt_tokens, + completion_tokens=response_obj["usage"].completion_tokens, + total_tokens=response_obj["usage"].total_tokens, + ) model_response.model = self.model print_verbose( @@ -9887,19 +9912,6 @@ class CustomStreamWrapper: ## RETURN ARG if ( - "content" in completion_obj - and isinstance(completion_obj["content"], str) - and len(completion_obj["content"]) == 0 - and hasattr(model_response, "usage") - and hasattr(model_response.usage, "prompt_tokens") - ): - if self.sent_first_chunk is False: - completion_obj["role"] = "assistant" - self.sent_first_chunk = True - model_response.choices[0].delta = Delta(**completion_obj) - print_verbose(f"returning model_response: {model_response}") - return model_response - elif ( "content" in completion_obj and ( isinstance(completion_obj["content"], str) @@ -9994,6 +10006,7 @@ class CustomStreamWrapper: model_response.choices[0].finish_reason = map_finish_reason( finish_reason=self.received_finish_reason ) # ensure consistent output to openai + self.sent_last_chunk = True return model_response @@ -10006,6 +10019,8 @@ class CustomStreamWrapper: self.sent_first_chunk = True return model_response else: + if hasattr(model_response, "usage"): + self.chunks.append(model_response) return except StopIteration: raise StopIteration @@ -10122,17 +10137,22 @@ class CustomStreamWrapper: del obj_dict["usage"] # Create a new object without the removed attribute - response = self.model_response_creator(chunk=obj_dict) - + response = self.model_response_creator( + chunk=obj_dict, hidden_params=response._hidden_params + ) + # add usage as hidden param + if self.sent_last_chunk is True and self.stream_options is None: + usage = calculate_total_usage(chunks=self.chunks) + response._hidden_params["usage"] = usage # RETURN RESULT return response except StopIteration: if self.sent_last_chunk is True: if ( - self.sent_stream_usage == False + self.sent_stream_usage is False and self.stream_options is not None - and self.stream_options.get("include_usage", False) == True + and self.stream_options.get("include_usage", False) is True ): # send the final chunk with stream options complete_streaming_response = litellm.stream_chunk_builder( @@ -10140,6 +10160,7 @@ class CustomStreamWrapper: ) response = self.model_response_creator() response.usage = complete_streaming_response.usage # type: ignore + response._hidden_params["usage"] = complete_streaming_response.usage # type: ignore ## LOGGING threading.Thread( target=self.logging_obj.success_handler, @@ -10151,6 +10172,9 @@ class CustomStreamWrapper: else: self.sent_last_chunk = True processed_chunk = self.finish_reason_handler() + if self.stream_options is None: # add usage as hidden param + usage = calculate_total_usage(chunks=self.chunks) + setattr(processed_chunk, "usage", usage) ## LOGGING threading.Thread( target=self.logging_obj.success_handler, diff --git a/tests/test_debug_warning.py b/tests/test_debug_warning.py new file mode 100644 index 00000000000..e69de29bb2d