scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-20 00:11:50 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 75

fix(pass_through): inject cost into Anthropic streaming chunks + fix SSE parsing in tests

Browse source 
- streaming_handler.py: add EndpointType.ANTHROPIC branch so cost gets
  injected into message_delta chunks on the passthrough path
- test_anthropic_passthrough.py: fix SSE parsing — AnthropicResponsesStreamWrapper
  yields full multi-line SSE frames as single bytes objects, so split each
  chunk by \n before checking for 'data: ' prefix; remove @pytest.mark.skip
  from test_anthropic_messages_openai_model_streaming_cost_injection
- test_image_variation.py: add .name attr to BytesIO + OPENAI_API_KEY skipif
- test_claude_agent_sdk.py + test_config.yaml: nova-premier → nova-pro
This commit is contained in:
Ishaan Jaffer 2026-03-07 17:23:03 -08:00
parent 9adbbe2cae
commit 50b5b7eb62
5 changed files with 87 additions and 62 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

8
litellm/proxy/pass_through_endpoints/streaming_handler.py
View file

@ -67,6 +67,14 @@ class PassThroughStreamingHandler:
)
if modified_chunk is not None:
chunk = modified_chunk
elif endpoint_type == EndpointType.ANTHROPIC:
modified_chunk = (
ProxyBaseLLMRequestProcessing._process_chunk_with_cost_injection(
chunk, model_name
)
)
if modified_chunk is not None:
chunk = modified_chunk
yield chunk

5
tests/image_gen_tests/test_image_variation.py
View file

@ -45,10 +45,15 @@ def image_url():
# Load the image into a file-like object
image_file = BytesIO(response.content)
image_file.name = "litellm_logo.png"
return image_file
@pytest.mark.skipif(
not os.environ.get("OPENAI_API_KEY"),
reason="OPENAI_API_KEY not set",
)
def test_openai_image_variation_openai_sdk(image_url):
from openai import OpenAI

122
tests/pass_through_tests/test_anthropic_passthrough.py
View file

@ -337,41 +337,46 @@ async def test_anthropic_messages_streaming_cost_injection():
async with aiohttp.ClientSession() as session:
async with session.post(
"http://0.0.0.0:4000/v1/messages",
json=payload,
headers=headers
"http://0.0.0.0:4000/v1/messages",
json=payload,
headers=headers,
) as response:
assert response.status == 200
# Collect all SSE events
# Collect all SSE events.
# Split each chunk by newlines to handle both:
# - Anthropic direct path: chunks arrive as individual lines
# - OpenAI/Responses API path: chunks are full multi-line SSE events
events = []
async for line in response.content:
line_str = line.decode('utf-8').strip()
if line_str.startswith('data: '):
try:
data = json.loads(line_str[6:]) # Remove 'data: ' prefix
events.append(data)
except json.JSONDecodeError:
continue
async for chunk in response.content:
chunk_str = chunk.decode("utf-8")
for line in chunk_str.split("\n"):
line = line.strip()
if line.startswith("data: "):
try:
data = json.loads(line[6:]) # Remove 'data: ' prefix
events.append(data)
except json.JSONDecodeError:
continue
# Find message_delta event with usage
message_delta_events = [
event for event in events
if event.get('type') == 'message_delta' and 'usage' in event
event for event in events
if event.get("type") == "message_delta" and "usage" in event
]
assert len(message_delta_events) > 0, "No message_delta events with usage found"
# Check that cost is included in usage
for event in message_delta_events:
usage = event.get('usage', {})
assert 'cost' in usage, f"Cost not found in usage: {usage}"
assert isinstance(usage['cost'], (int, float)), f"Cost should be numeric: {usage['cost']}"
assert usage['cost'] >= 0, f"Cost should be non-negative: {usage['cost']}"
print(f"✅ Found message_delta with cost: {usage}")
print(f"✅ Test passed: Found {len(message_delta_events)} message_delta events with cost")
usage = event.get("usage", {})
assert "cost" in usage, f"Cost not found in usage: {usage}"
assert isinstance(usage["cost"], (int, float)), f"Cost should be numeric: {usage['cost']}"
assert usage["cost"] >= 0, f"Cost should be non-negative: {usage['cost']}"
print(f"Found message_delta with cost: {usage}")
print(f"Test passed: Found {len(message_delta_events)} message_delta events with cost")
@pytest.mark.asyncio
@ -381,54 +386,61 @@ async def test_anthropic_messages_openai_model_streaming_cost_injection():
Test that cost is injected into message_delta usage for OpenAI model via Anthropic Messages API
"""
print("Testing cost injection in Anthropic Messages API with OpenAI model")
headers = {
"Authorization": "Bearer sk-1234",
"Content-Type": "application/json",
"anthropic-version": "2023-06-01",
}
payload = {
"model": "openai/gpt-4o",
"max_tokens": 10,
"stream": True,
"messages": [{"role": "user", "content": "Say 'Hi'"}],
}
async with aiohttp.ClientSession() as session:
async with session.post(
"http://0.0.0.0:4000/v1/messages",
json=payload,
headers=headers
"http://0.0.0.0:4000/v1/messages",
json=payload,
headers=headers,
) as response:
assert response.status == 200
# Collect all SSE events
# Collect all SSE events.
# Split each chunk by newlines to handle both:
# - Direct API paths: chunks arrive as individual lines
# - OpenAI/Responses API path: AnthropicResponsesStreamWrapper yields
# full multi-line SSE events as single bytes objects, so a naive
# startswith('data: ') check on the whole chunk misses them.
events = []
async for line in response.content:
line_str = line.decode('utf-8').strip()
if line_str.startswith('data: '):
try:
data = json.loads(line_str[6:]) # Remove 'data: ' prefix
events.append(data)
except json.JSONDecodeError:
continue
async for chunk in response.content:
chunk_str = chunk.decode("utf-8")
for line in chunk_str.split("\n"):
line = line.strip()
if line.startswith("data: "):
try:
data = json.loads(line[6:]) # Remove 'data: ' prefix
events.append(data)
except json.JSONDecodeError:
continue
# Find message_delta event with usage
message_delta_events = [
event for event in events
if event.get('type') == 'message_delta' and 'usage' in event
event for event in events
if event.get("type") == "message_delta" and "usage" in event
]
assert len(message_delta_events) > 0, "No message_delta events with usage found"
# Check that cost is included in usage
for event in message_delta_events:
usage = event.get('usage', {})
assert 'cost' in usage, f"Cost not found in usage: {usage}"
assert isinstance(usage['cost'], (int, float)), f"Cost should be numeric: {usage['cost']}"
assert usage['cost'] >= 0, f"Cost should be non-negative: {usage['cost']}"
print(f"✅ Found message_delta with cost: {usage}")
print(f"✅ Test passed: Found {len(message_delta_events)} message_delta events with cost")
usage = event.get("usage", {})
assert "cost" in usage, f"Cost not found in usage: {usage}"
assert isinstance(usage["cost"], (int, float)), f"Cost should be numeric: {usage['cost']}"
assert usage["cost"] >= 0, f"Cost should be non-negative: {usage['cost']}"
print(f"Found message_delta with cost: {usage}")
print(f"Test passed: Found {len(message_delta_events)} message_delta events with cost")

10
tests/proxy_e2e_anthropic_messages_tests/test_claude_agent_sdk.py
View file

@ -4,7 +4,7 @@ E2E tests for Claude Agent SDK with LiteLLM Proxy using Bedrock models.
Tests streaming messages across different Bedrock models:
- Regular Bedrock Claude Sonnet 4.5
- Bedrock Converse Claude Sonnet 4.5
- AWS Nova Premier
- AWS Nova Pro
"""
import os
@ -14,14 +14,14 @@ from claude_agent_sdk import ClaudeSDKClient, ClaudeAgentOptions
# Test models from test_config.yaml
# Note: bedrock-converse-claude-sonnet-4.5 removed temporarily as the Bedrock Converse API
# Note: bedrock-converse-claude-sonnet-4.5 removed temporarily as the Bedrock Converse API
# for Claude Sonnet 4.5 may not be available in all regions/accounts
# Note: bedrock-nova-premier requires an inference profile for on-demand throughput
# https://docs.aws.amazon.com/bedrock/latest/userguide/inference-profiles.html
# Note: bedrock-nova-premier requires provisioned throughput (not standard cross-region
# inference profile) and is not reliably available in CI accounts. Using nova-pro instead.
TEST_MODELS = [
("bedrock-claude-sonnet-4.5", "Bedrock Invoke API"),
("bedrock-converse-claude-sonnet-4.5", "Bedrock Converse API"),
("bedrock-nova-premier", "AWS Nova Premier"),
("bedrock-nova-pro", "AWS Nova Pro"),
]

4
tests/proxy_e2e_anthropic_messages_tests/test_config.yaml
View file

@ -24,9 +24,9 @@ model_list:
model: "bedrock/us.anthropic.claude-opus-4-5-20251101-v1:0"
aws_region_name: "us-east-1"
- model_name: bedrock-nova-premier
- model_name: bedrock-nova-pro
litellm_params:
model: "bedrock/us.amazon.nova-premier-v1:0"
model: "bedrock/us.amazon.nova-pro-v1:0"
aws_region_name: "us-east-1"
# Converse API models