ai-workspace-services/litellm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

(Bug Fix) - Accurate token counting for /anthropic/ API Routes on LiteLLM Proxy (#8880)

Browse Source 
* 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
This commit is contained in:
Ishaan Jaff 2025-02-27 15:43:03 -08:00 committed by GitHub
parent 24df2331ec
commit 047d1b1208
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
5 changed files with 140 additions and 28 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

13
litellm/llms/anthropic/chat/handler.py
View File

@ -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()

18
litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py
View File

@ -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(

87
tests/logging_callback_tests/test_token_counting.py
View File

@ -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
)

46
tests/pass_through_tests/test_anthropic_passthrough.py
View File

@ -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"]

4
tests/pass_through_unit_tests/test_unit_test_anthropic_pass_through.py
View File

@ -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