diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index 1ca6a6fd6d7..94387895e15 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -270,7 +270,12 @@ class ChunkProcessor: usage_chunk = chunk._hidden_params.get("usage", None) if usage_chunk is not None: if "prompt_tokens" in usage_chunk: - prompt_tokens = usage_chunk.get("prompt_tokens", 0) or 0 + usage_chunk_prompt_tokens = usage_chunk.get("prompt_tokens", 0) + if ( + usage_chunk_prompt_tokens is not None + and usage_chunk_prompt_tokens > 0 + ): + prompt_tokens = usage_chunk_prompt_tokens if "completion_tokens" in usage_chunk: completion_tokens = usage_chunk.get("completion_tokens", 0) or 0 if "cache_creation_input_tokens" in usage_chunk: diff --git a/litellm/llms/anthropic/chat/handler.py b/litellm/llms/anthropic/chat/handler.py index a30cd657091..da95ac075f0 100644 --- a/litellm/llms/anthropic/chat/handler.py +++ b/litellm/llms/anthropic/chat/handler.py @@ -769,6 +769,7 @@ class ModelResponseIterator: message=message, status_code=500, # it looks like Anthropic API does not return a status code in the chunk error - default to 500 ) + returned_chunk = GenericStreamingChunk( text=text, tool_use=tool_use, diff --git a/litellm/proxy/_new_secret_config.yaml b/litellm/proxy/_new_secret_config.yaml index acbac00107a..fcbc36d4d07 100644 --- a/litellm/proxy/_new_secret_config.yaml +++ b/litellm/proxy/_new_secret_config.yaml @@ -18,10 +18,10 @@ model_list: model: "openai/*" api_key: sk-1234 -# litellm_settings: -# fallbacks: [{ "claude-3-5-sonnet-20240620": ["claude-3-5-sonnet-aihubmix"] }] -# callbacks: ["otel", "prometheus"] -# default_redis_batch_cache_expiry: 10 +litellm_settings: + fallbacks: [{ "claude-3-5-sonnet-20240620": ["claude-3-5-sonnet-aihubmix"] }] + callbacks: ["prometheus"] + default_redis_batch_cache_expiry: 10 # litellm_settings: diff --git a/litellm/utils.py b/litellm/utils.py index 70f43e51210..62eae1aba79 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -5926,6 +5926,9 @@ class CustomStreamWrapper: ) self.messages = getattr(logging_obj, "messages", None) self.sent_stream_usage = False + self.send_stream_usage = ( + True if self.check_send_stream_usage(self.stream_options) else False + ) self.tool_call = False self.chunks: List = ( [] @@ -5938,6 +5941,12 @@ class CustomStreamWrapper: def __aiter__(self): return self + def check_send_stream_usage(self, stream_options: Optional[dict]): + return ( + stream_options is not None + and stream_options.get("include_usage", False) is True + ) + def check_is_function_call(self, logging_obj) -> bool: if hasattr(logging_obj, "optional_params") and isinstance( logging_obj.optional_params, dict @@ -6790,6 +6799,143 @@ class CustomStreamWrapper: is_empty = False return is_empty + def return_processed_chunk_logic( # noqa + self, completion_obj: dict, model_response: ModelResponse, response_obj: dict + ): + _streaming_choices = cast(StreamingChoices, model_response.choices[0]) + + if ( + "content" in completion_obj + and ( + isinstance(completion_obj["content"], str) + and len(completion_obj["content"]) > 0 + ) + or ( + "tool_calls" in completion_obj + and completion_obj["tool_calls"] is not None + and len(completion_obj["tool_calls"]) > 0 + ) + or ( + "function_call" in completion_obj + and completion_obj["function_call"] is not None + ) + ): # cannot set content of an OpenAI Object to be an empty string + self.safety_checker() + hold, model_response_str = self.check_special_tokens( + chunk=completion_obj["content"], + finish_reason=model_response.choices[0].finish_reason, + ) # filter out bos/eos tokens from openai-compatible hf endpoints + print_verbose(f"hold - {hold}, model_response_str - {model_response_str}") + if hold is False: + ## check if openai/azure chunk + original_chunk = response_obj.get("original_chunk", None) + if original_chunk: + model_response.id = original_chunk.id + self.response_id = original_chunk.id + if len(original_chunk.choices) > 0: + choices = [] + for idx, choice in enumerate(original_chunk.choices): + try: + if isinstance(choice, BaseModel): + try: + choice_json = choice.model_dump() + except Exception: + choice_json = choice.dict() + choice_json.pop( + "finish_reason", None + ) # for mistral etc. which return a value in their last chunk (not-openai compatible). + print_verbose(f"choice_json: {choice_json}") + choices.append(StreamingChoices(**choice_json)) + except Exception: + choices.append(StreamingChoices()) + print_verbose(f"choices in streaming: {choices}") + setattr(model_response, "choices", choices) + else: + return + model_response.system_fingerprint = ( + original_chunk.system_fingerprint + ) + setattr( + model_response, + "citations", + getattr(original_chunk, "citations", None), + ) + print_verbose(f"self.sent_first_chunk: {self.sent_first_chunk}") + if self.sent_first_chunk is False: + _streaming_choices.delta["role"] = "assistant" + self.sent_first_chunk = True + elif self.sent_first_chunk is True and hasattr( + _streaming_choices.delta, "role" + ): + _initial_delta = _streaming_choices.delta.model_dump() + _initial_delta.pop("role", None) + _streaming_choices.delta = Delta(**_initial_delta) + print_verbose( + f"_streaming_choices.delta: {_streaming_choices.delta}" + ) + else: + ## else + completion_obj["content"] = model_response_str + if self.sent_first_chunk is False: + completion_obj["role"] = "assistant" + self.sent_first_chunk = True + + _streaming_choices.delta = Delta(**completion_obj) + _index: Optional[int] = completion_obj.get("index") + if _index is not None: + model_response.choices[0].index = _index + print_verbose(f"returning model_response: {model_response}") + return model_response + else: + return + elif self.received_finish_reason is not None: + if self.sent_last_chunk is True: + # Bedrock returns the guardrail trace in the last chunk - we want to return this here + if self.custom_llm_provider == "bedrock" and "trace" in model_response: + return model_response + + # Default - return StopIteration + raise StopIteration + # flush any remaining holding chunk + if len(self.holding_chunk) > 0: + if _streaming_choices.delta.content is None: + _streaming_choices.delta.content = self.holding_chunk + else: + _streaming_choices.delta.content = ( + self.holding_chunk + _streaming_choices.delta.content + ) + self.holding_chunk = "" + # if delta is None + _is_delta_empty = self.is_delta_empty(delta=_streaming_choices.delta) + + if _is_delta_empty: + # get any function call arguments + 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 + elif ( + _streaming_choices.delta.tool_calls is not None + or _streaming_choices.delta.function_call is not None + ): + if self.sent_first_chunk is False: + _streaming_choices.delta["role"] = "assistant" + self.sent_first_chunk = True + return model_response + elif ( + len(model_response.choices) > 0 + and hasattr(_streaming_choices.delta, "audio") + and _streaming_choices.delta.audio is not None + ): + return model_response + else: + if hasattr(model_response, "usage"): + self.chunks.append(model_response) + return + def chunk_creator(self, chunk): # type: ignore # noqa: PLR0915 model_response = self.model_response_creator() response_obj = {} @@ -6843,6 +6989,7 @@ class CustomStreamWrapper: "provider_specific_fields" ].items(): setattr(model_response, key, value) + response_obj = anthropic_response_obj elif ( self.custom_llm_provider @@ -6910,7 +7057,7 @@ class CustomStreamWrapper: if self.sent_first_chunk is False: raise Exception("An unknown error occurred with the stream") self.received_finish_reason = "stop" - elif self.custom_llm_provider and (self.custom_llm_provider == "vertex_ai"): + elif self.custom_llm_provider == "vertex_ai": import proto # type: ignore if self.model.startswith("claude-3"): @@ -7293,145 +7440,11 @@ class CustomStreamWrapper: self.tool_call = True ## RETURN ARG - if ( - "content" in completion_obj - and ( - isinstance(completion_obj["content"], str) - and len(completion_obj["content"]) > 0 - ) - or ( - "tool_calls" in completion_obj - and completion_obj["tool_calls"] is not None - and len(completion_obj["tool_calls"]) > 0 - ) - or ( - "function_call" in completion_obj - and completion_obj["function_call"] is not None - ) - ): # cannot set content of an OpenAI Object to be an empty string - self.safety_checker() - hold, model_response_str = self.check_special_tokens( - chunk=completion_obj["content"], - finish_reason=model_response.choices[0].finish_reason, - ) # filter out bos/eos tokens from openai-compatible hf endpoints - print_verbose( - f"hold - {hold}, model_response_str - {model_response_str}" - ) - if hold is False: - ## check if openai/azure chunk - original_chunk = response_obj.get("original_chunk", None) - if original_chunk: - model_response.id = original_chunk.id - self.response_id = original_chunk.id - if len(original_chunk.choices) > 0: - choices = [] - for idx, choice in enumerate(original_chunk.choices): - try: - if isinstance(choice, BaseModel): - try: - choice_json = choice.model_dump() - except Exception: - choice_json = choice.dict() - choice_json.pop( - "finish_reason", None - ) # for mistral etc. which return a value in their last chunk (not-openai compatible). - print_verbose(f"choice_json: {choice_json}") - choices.append(StreamingChoices(**choice_json)) - except Exception: - choices.append(StreamingChoices()) - print_verbose(f"choices in streaming: {choices}") - model_response.choices = choices - else: - return - model_response.system_fingerprint = ( - original_chunk.system_fingerprint - ) - model_response.citations = getattr( - original_chunk, "citations", None - ) - print_verbose(f"self.sent_first_chunk: {self.sent_first_chunk}") - if self.sent_first_chunk is False: - model_response.choices[0].delta["role"] = "assistant" - self.sent_first_chunk = True - elif self.sent_first_chunk is True and hasattr( - model_response.choices[0].delta, "role" - ): - _initial_delta = model_response.choices[ - 0 - ].delta.model_dump() - _initial_delta.pop("role", None) - model_response.choices[0].delta = Delta(**_initial_delta) - print_verbose( - f"model_response.choices[0].delta: {model_response.choices[0].delta}" - ) - else: - ## else - completion_obj["content"] = model_response_str - if self.sent_first_chunk is False: - completion_obj["role"] = "assistant" - self.sent_first_chunk = True - - model_response.choices[0].delta = Delta(**completion_obj) - if completion_obj.get("index") is not None: - model_response.choices[0].index = completion_obj.get( - "index" - ) - print_verbose(f"returning model_response: {model_response}") - return model_response - else: - return - elif self.received_finish_reason is not None: - if self.sent_last_chunk is True: - # Bedrock returns the guardrail trace in the last chunk - we want to return this here - if ( - self.custom_llm_provider == "bedrock" - and "trace" in model_response - ): - return model_response - - # Default - return StopIteration - raise StopIteration - # flush any remaining holding chunk - if len(self.holding_chunk) > 0: - if model_response.choices[0].delta.content is None: - model_response.choices[0].delta.content = self.holding_chunk - else: - model_response.choices[0].delta.content = ( - self.holding_chunk + model_response.choices[0].delta.content - ) - self.holding_chunk = "" - # if delta is None - _is_delta_empty = self.is_delta_empty( - delta=model_response.choices[0].delta - ) - - if _is_delta_empty: - # get any function call arguments - 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 - elif ( - model_response.choices[0].delta.tool_calls is not None - or model_response.choices[0].delta.function_call is not None - ): - if self.sent_first_chunk is False: - model_response.choices[0].delta["role"] = "assistant" - self.sent_first_chunk = True - return model_response - elif ( - len(model_response.choices) > 0 - and hasattr(model_response.choices[0].delta, "audio") - and model_response.choices[0].delta.audio is not None - ): - return model_response - else: - if hasattr(model_response, "usage"): - self.chunks.append(model_response) - return + return self.return_processed_chunk_logic( + completion_obj=completion_obj, + model_response=model_response, + response_obj=response_obj, + ) except StopIteration: raise StopIteration except Exception as e: @@ -7577,27 +7590,24 @@ class CustomStreamWrapper: except StopIteration: if self.sent_last_chunk is True: - if ( - self.sent_stream_usage is False - and self.stream_options is not None - and self.stream_options.get("include_usage", False) is True - ): - # send the final chunk with stream options - complete_streaming_response = litellm.stream_chunk_builder( - chunks=self.chunks, messages=self.messages + complete_streaming_response = litellm.stream_chunk_builder( + chunks=self.chunks, messages=self.messages + ) + response = self.model_response_creator() + if complete_streaming_response is not None: + setattr( + response, + "usage", + getattr(complete_streaming_response, "usage"), ) - response = self.model_response_creator() - if complete_streaming_response is not None: - setattr( - response, - "usage", - getattr(complete_streaming_response, "usage"), - ) - ## LOGGING - threading.Thread( - target=self.logging_obj.success_handler, - args=(response, None, None, cache_hit), - ).start() # log response + + ## LOGGING + threading.Thread( + target=self.logging_obj.success_handler, + args=(response, None, None, cache_hit), + ).start() # log response + + if self.sent_stream_usage is False and self.send_stream_usage is True: self.sent_stream_usage = True return response raise # Re-raise StopIteration @@ -7715,10 +7725,7 @@ class CustomStreamWrapper: end_time=None, cache_hit=cache_hit, ) - # threading.Thread( - # target=self.logging_obj.success_handler, - # args=(processed_chunk, None, None, cache_hit), - # ).start() # log response + asyncio.create_task( self.logging_obj.async_success_handler( processed_chunk, cache_hit=cache_hit @@ -7799,82 +7806,33 @@ class CustomStreamWrapper: # RETURN RESULT self.chunks.append(processed_chunk) return processed_chunk - except StopAsyncIteration: + except (StopAsyncIteration, StopIteration): if self.sent_last_chunk is True: - if ( - self.sent_stream_usage is False - and self.stream_options is not None - and self.stream_options.get("include_usage", False) is True - ): - # send the final chunk with stream options - complete_streaming_response = litellm.stream_chunk_builder( - chunks=self.chunks, messages=self.messages + # log the final chunk with accurate streaming values + complete_streaming_response = litellm.stream_chunk_builder( + chunks=self.chunks, messages=self.messages + ) + response = self.model_response_creator() + if complete_streaming_response is not None: + setattr( + response, + "usage", + getattr(complete_streaming_response, "usage"), ) - response = self.model_response_creator() - if complete_streaming_response is not None: - setattr( - response, - "usage", - getattr(complete_streaming_response, "usage"), - ) - ## LOGGING - threading.Thread( - target=self.logging_obj.success_handler, - args=(response, None, None, cache_hit), - ).start() # log response - asyncio.create_task( - self.logging_obj.async_success_handler( - response, cache_hit=cache_hit - ) - ) - self.sent_stream_usage = True - return response - raise # Re-raise StopIteration - else: - self.sent_last_chunk = True - processed_chunk = self.finish_reason_handler() ## LOGGING threading.Thread( target=self.logging_obj.success_handler, - args=(processed_chunk, None, None, cache_hit), + args=(response, None, None, cache_hit), ).start() # log response asyncio.create_task( self.logging_obj.async_success_handler( - processed_chunk, cache_hit=cache_hit + response, cache_hit=cache_hit ) ) - return processed_chunk - except StopIteration: - if self.sent_last_chunk is True: - if ( - self.sent_stream_usage is False - and self.stream_options is not None - and self.stream_options.get("include_usage", False) is True - ): - # send the final chunk with stream options - complete_streaming_response = litellm.stream_chunk_builder( - chunks=self.chunks, messages=self.messages - ) - response = self.model_response_creator() - if complete_streaming_response is not None: - setattr( - response, - "usage", - getattr(complete_streaming_response, "usage"), - ) - ## LOGGING - threading.Thread( - target=self.logging_obj.success_handler, - args=(response, None, None, cache_hit), - ).start() # log response - asyncio.create_task( - self.logging_obj.async_success_handler( - response, cache_hit=cache_hit - ) - ) + if self.sent_stream_usage is False and self.send_stream_usage is True: self.sent_stream_usage = True return response - raise StopAsyncIteration + raise StopAsyncIteration # Re-raise StopIteration else: self.sent_last_chunk = True processed_chunk = self.finish_reason_handler() diff --git a/tests/local_testing/test_custom_callback_input.py b/tests/local_testing/test_custom_callback_input.py index 3ce3a618c4a..1744d389106 100644 --- a/tests/local_testing/test_custom_callback_input.py +++ b/tests/local_testing/test_custom_callback_input.py @@ -1562,3 +1562,65 @@ def test_logging_key_masking_gemini(): trimmed_key = key.split("key=")[1] trimmed_key = trimmed_key.replace("*", "") assert "PART" == trimmed_key + + +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.asyncio +async def test_standard_logging_payload_stream_usage(sync_mode): + """ + Even if stream_options is not provided, correct usage should be logged + """ + from litellm.types.utils import StandardLoggingPayload + from litellm.main import stream_chunk_builder + + stream = True + try: + # sync completion + customHandler = CompletionCustomHandler() + litellm.callbacks = [customHandler] + + if sync_mode: + patch_event = "log_success_event" + return_val = MagicMock() + else: + patch_event = "async_log_success_event" + return_val = AsyncMock() + + with patch.object(customHandler, patch_event, new=return_val) as mock_client: + if sync_mode: + resp = litellm.completion( + model="anthropic/claude-3-5-sonnet-20240620", + messages=[{"role": "user", "content": "Hey, how's it going?"}], + stream=stream, + ) + + chunks = [] + for chunk in resp: + chunks.append(chunk) + time.sleep(2) + else: + resp = await litellm.acompletion( + model="anthropic/claude-3-5-sonnet-20240620", + messages=[{"role": "user", "content": "Hey, how's it going?"}], + stream=stream, + ) + + chunks = [] + async for chunk in resp: + chunks.append(chunk) + await asyncio.sleep(2) + + mock_client.assert_called_once() + + standard_logging_object: StandardLoggingPayload = ( + mock_client.call_args.kwargs["kwargs"]["standard_logging_object"] + ) + + built_response = stream_chunk_builder(chunks=chunks) + assert ( + built_response.usage.total_tokens + != standard_logging_object["total_tokens"] + ) + print(f"standard_logging_object usage: {built_response.usage}") + except litellm.InternalServerError: + pass