From 047d1b120801037f258bfd0ea40ee6b8c36f97fa Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Thu, 27 Feb 2025 15:43:03 -0800 Subject: [PATCH] (Bug Fix) - Accurate token counting for `/anthropic/` API Routes on LiteLLM Proxy (#8880) * fix _create_anthropic_response_logging_payload * fix - pass through don't create standard logging payload * fix logged key hash * test_init_kwargs_for_pass_through_endpoint_basic * test_unit_test_anthropic_pass_through * fix anthropic pass through logging handler * test_stream_token_counting_anthropic_with_include_usage * convert_str_chunk_to_generic_chunk * _build_complete_streaming_response * test_anthropic_basic_completion_with_headers * test_anthropic_streaming_with_headers * improve test for pass through token counting --- litellm/llms/anthropic/chat/handler.py | 13 +-- .../anthropic_passthrough_logging_handler.py | 18 ++-- .../test_token_counting.py | 87 +++++++++++++++++++ .../test_anthropic_passthrough.py | 46 ++++++++-- .../test_unit_test_anthropic_pass_through.py | 4 + 5 files changed, 140 insertions(+), 28 deletions(-) diff --git a/litellm/llms/anthropic/chat/handler.py b/litellm/llms/anthropic/chat/handler.py index 4d8d8767a4..46c8edae03 100644 --- a/litellm/llms/anthropic/chat/handler.py +++ b/litellm/llms/anthropic/chat/handler.py @@ -810,9 +810,7 @@ class ModelResponseIterator: except ValueError as e: raise RuntimeError(f"Error parsing chunk: {e},\nReceived chunk: {chunk}") - def convert_str_chunk_to_generic_chunk( - self, chunk: str - ) -> Union[GenericStreamingChunk, ModelResponseStream]: + def convert_str_chunk_to_generic_chunk(self, chunk: str) -> ModelResponseStream: """ Convert a string chunk to a GenericStreamingChunk @@ -832,11 +830,4 @@ class ModelResponseIterator: data_json = json.loads(str_line[5:]) return self.chunk_parser(chunk=data_json) else: - return GenericStreamingChunk( - text="", - is_finished=False, - finish_reason="", - usage=None, - index=0, - tool_use=None, - ) + return ModelResponseStream() diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py index 17ff358bdb..d3496540ed 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py @@ -196,23 +196,19 @@ class AnthropicPassthroughLoggingHandler: streaming_response=None, sync_stream=False, ) - litellm_custom_stream_wrapper = litellm.CustomStreamWrapper( - completion_stream=anthropic_model_response_iterator, - model=model, - logging_obj=litellm_logging_obj, - custom_llm_provider="anthropic", - ) all_openai_chunks = [] for _chunk_str in all_chunks: try: - generic_chunk = anthropic_model_response_iterator.convert_str_chunk_to_generic_chunk( + transformed_openai_chunk = anthropic_model_response_iterator.convert_str_chunk_to_generic_chunk( chunk=_chunk_str ) - litellm_chunk = litellm_custom_stream_wrapper.chunk_creator( - chunk=generic_chunk + if transformed_openai_chunk is not None: + all_openai_chunks.append(transformed_openai_chunk) + + verbose_proxy_logger.debug( + "all openai chunks= %s", + json.dumps(all_openai_chunks, indent=4, default=str), ) - if litellm_chunk is not None: - all_openai_chunks.append(litellm_chunk) except (StopIteration, StopAsyncIteration): break complete_streaming_response = litellm.stream_chunk_builder( diff --git a/tests/logging_callback_tests/test_token_counting.py b/tests/logging_callback_tests/test_token_counting.py index bce938a670..341ef2a545 100644 --- a/tests/logging_callback_tests/test_token_counting.py +++ b/tests/logging_callback_tests/test_token_counting.py @@ -157,3 +157,90 @@ async def test_stream_token_counting_with_redaction(): actual_usage.completion_tokens == custom_logger.recorded_usage.completion_tokens ) assert actual_usage.total_tokens == custom_logger.recorded_usage.total_tokens + + +@pytest.mark.asyncio +async def test_stream_token_counting_anthropic_with_include_usage(): + """ """ + from anthropic import Anthropic + + anthropic_client = Anthropic(api_key=os.getenv("ANTHROPIC_API_KEY")) + litellm._turn_on_debug() + + custom_logger = TestCustomLogger() + litellm.logging_callback_manager.add_litellm_callback(custom_logger) + + input_text = "Respond in just 1 word. Say ping" + + response = await litellm.acompletion( + model="claude-3-5-sonnet-20240620", + messages=[{"role": "user", "content": input_text}], + max_tokens=4096, + stream=True, + ) + + actual_usage = None + output_text = "" + async for chunk in response: + output_text += chunk["choices"][0]["delta"]["content"] or "" + pass + + await asyncio.sleep(1) + + print("\n\n\n\n\n") + print( + "recorded_usage", + json.dumps(custom_logger.recorded_usage, indent=4, default=str), + ) + print("\n\n\n\n\n") + + # print making the same request with anthropic client + anthropic_response = anthropic_client.messages.create( + model="claude-3-5-sonnet-20240620", + max_tokens=4096, + messages=[{"role": "user", "content": input_text}], + stream=True, + ) + usage = None + all_anthropic_usage_chunks = [] + for chunk in anthropic_response: + print("chunk", json.dumps(chunk, indent=4, default=str)) + if hasattr(chunk, "message"): + if chunk.message.usage: + print( + "USAGE BLOCK", + json.dumps(chunk.message.usage, indent=4, default=str), + ) + all_anthropic_usage_chunks.append(chunk.message.usage) + elif hasattr(chunk, "usage"): + print("USAGE BLOCK", json.dumps(chunk.usage, indent=4, default=str)) + all_anthropic_usage_chunks.append(chunk.usage) + + print( + "all_anthropic_usage_chunks", + json.dumps(all_anthropic_usage_chunks, indent=4, default=str), + ) + + input_tokens_anthropic_api = sum( + [getattr(usage, "input_tokens", 0) for usage in all_anthropic_usage_chunks] + ) + output_tokens_anthropic_api = sum( + [getattr(usage, "output_tokens", 0) for usage in all_anthropic_usage_chunks] + ) + print("input_tokens_anthropic_api", input_tokens_anthropic_api) + print("output_tokens_anthropic_api", output_tokens_anthropic_api) + + print("input_tokens_litellm", custom_logger.recorded_usage.prompt_tokens) + print("output_tokens_litellm", custom_logger.recorded_usage.completion_tokens) + + ## Assert Accuracy of token counting + # input tokens should be exactly the same + assert input_tokens_anthropic_api == custom_logger.recorded_usage.prompt_tokens + + # output tokens can have at max abs diff of 10. We can't guarantee the response from two api calls will be exactly the same + assert ( + abs( + output_tokens_anthropic_api - custom_logger.recorded_usage.completion_tokens + ) + <= 10 + ) diff --git a/tests/pass_through_tests/test_anthropic_passthrough.py b/tests/pass_through_tests/test_anthropic_passthrough.py index a6a1c9c0ed..ba729a485c 100644 --- a/tests/pass_through_tests/test_anthropic_passthrough.py +++ b/tests/pass_through_tests/test_anthropic_passthrough.py @@ -6,6 +6,7 @@ import pytest import anthropic import aiohttp import asyncio +import json client = anthropic.Anthropic( base_url="http://0.0.0.0:4000/anthropic", api_key="sk-1234" @@ -78,6 +79,13 @@ async def test_anthropic_basic_completion_with_headers(): response_json = await response.json() response_headers = response.headers + print( + "non-streaming response", + json.dumps(response_json, indent=4, default=str), + ) + reported_usage = response_json.get("usage", None) + anthropic_api_input_tokens = reported_usage.get("input_tokens", None) + anthropic_api_output_tokens = reported_usage.get("output_tokens", None) litellm_call_id = response_headers.get("x-litellm-call-id") print(f"LiteLLM Call ID: {litellm_call_id}") @@ -121,10 +129,12 @@ async def test_anthropic_basic_completion_with_headers(): log_entry["spend"], (int, float) ), "Spend should be a number" assert log_entry["total_tokens"] > 0, "Should have some tokens" - assert log_entry["prompt_tokens"] > 0, "Should have prompt tokens" assert ( - log_entry["completion_tokens"] > 0 - ), "Should have completion tokens" + log_entry["prompt_tokens"] == anthropic_api_input_tokens + ), f"Should have prompt tokens matching anthropic api. Expected {anthropic_api_input_tokens} but got {log_entry['prompt_tokens']}" + assert ( + log_entry["completion_tokens"] == anthropic_api_output_tokens + ), f"Should have completion tokens matching anthropic api. Expected {anthropic_api_output_tokens} but got {log_entry['completion_tokens']}" assert ( log_entry["total_tokens"] == log_entry["prompt_tokens"] + log_entry["completion_tokens"] @@ -197,9 +207,30 @@ async def test_anthropic_streaming_with_headers(): collected_output.append(text[6:]) # Remove 'data: ' prefix print("Collected output:", "".join(collected_output)) + anthropic_api_usage_chunks = [] + for chunk in collected_output: + chunk_json = json.loads(chunk) + if "usage" in chunk_json: + anthropic_api_usage_chunks.append(chunk_json["usage"]) + elif "message" in chunk_json and "usage" in chunk_json["message"]: + anthropic_api_usage_chunks.append(chunk_json["message"]["usage"]) + + print( + "anthropic_api_usage_chunks", + json.dumps(anthropic_api_usage_chunks, indent=4, default=str), + ) + + anthropic_api_input_tokens = sum( + [usage.get("input_tokens", 0) for usage in anthropic_api_usage_chunks] + ) + anthropic_api_output_tokens = max( + [usage.get("output_tokens", 0) for usage in anthropic_api_usage_chunks] + ) + print("anthropic_api_input_tokens", anthropic_api_input_tokens) + print("anthropic_api_output_tokens", anthropic_api_output_tokens) # Wait for spend to be logged - await asyncio.sleep(20) + await asyncio.sleep(10) # Check spend logs for this specific request async with session.get( @@ -236,8 +267,11 @@ async def test_anthropic_streaming_with_headers(): ), "Spend should be a number" assert log_entry["total_tokens"] > 0, "Should have some tokens" assert ( - log_entry["completion_tokens"] > 0 - ), "Should have completion tokens" + log_entry["prompt_tokens"] == anthropic_api_input_tokens + ), f"Should have prompt tokens matching anthropic api. Expected {anthropic_api_input_tokens} but got {log_entry['prompt_tokens']}" + assert ( + log_entry["completion_tokens"] == anthropic_api_output_tokens + ), f"Should have completion tokens matching anthropic api. Expected {anthropic_api_output_tokens} but got {log_entry['completion_tokens']}" assert ( log_entry["total_tokens"] == log_entry["prompt_tokens"] + log_entry["completion_tokens"] diff --git a/tests/pass_through_unit_tests/test_unit_test_anthropic_pass_through.py b/tests/pass_through_unit_tests/test_unit_test_anthropic_pass_through.py index 5404c3ec88..bcd93de0bb 100644 --- a/tests/pass_through_unit_tests/test_unit_test_anthropic_pass_through.py +++ b/tests/pass_through_unit_tests/test_unit_test_anthropic_pass_through.py @@ -353,6 +353,7 @@ def test_handle_logging_anthropic_collected_chunks(all_chunks): ) assert isinstance(result["result"], ModelResponse) + print("result=", json.dumps(result, indent=4, default=str)) def test_build_complete_streaming_response(all_chunks): @@ -370,3 +371,6 @@ def test_build_complete_streaming_response(all_chunks): ) assert isinstance(result, ModelResponse) + assert result.usage.prompt_tokens == 17 + assert result.usage.completion_tokens == 249 + assert result.usage.total_tokens == 266