mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-24 00:52:24 +00:00
test(integration): assert recount pins and unique fixture request ids
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
kerry
parent
a841750d46
commit
2edea0be08
5 changed files with 86 additions and 65 deletions
|
|
@ -12,6 +12,7 @@ from pathlib import Path
|
|||
from queue import SimpleQueue
|
||||
import struct
|
||||
from typing import Final, cast
|
||||
import uuid
|
||||
import zlib
|
||||
|
||||
import httpx
|
||||
|
|
@ -79,10 +80,15 @@ def _aws_str_header(name: str, value: str) -> bytes:
|
|||
)
|
||||
|
||||
|
||||
def _aws_event_frame(event_type: str, payload: Mapping[str, JsonValue], scenario_id: str) -> bytes:
|
||||
def _aws_event_frame(
|
||||
event_type: str,
|
||||
payload: Mapping[str, JsonValue],
|
||||
scenario_id: str,
|
||||
unique_id: str,
|
||||
) -> bytes:
|
||||
payload_bytes: Final = json.dumps(payload, separators=(",", ":")).replace(
|
||||
"$REQUEST_ID", scenario_id
|
||||
).encode()
|
||||
).replace("$UNIQUE_ID", unique_id).encode()
|
||||
headers_bytes: Final = (
|
||||
_aws_str_header(":event-type", event_type)
|
||||
+ _aws_str_header(":content-type", "application/json")
|
||||
|
|
@ -209,11 +215,14 @@ class Provider:
|
|||
|
||||
@staticmethod
|
||||
def _response(response: StoredResponse, scenario_id: str) -> Response:
|
||||
unique_id: Final = f"{scenario_id}-{uuid.uuid4().hex[:8]}"
|
||||
match response:
|
||||
case JsonResponse():
|
||||
return Response(
|
||||
content=json.dumps(response.body, separators=(",", ":")).replace(
|
||||
"$REQUEST_ID", scenario_id
|
||||
).replace(
|
||||
"$UNIQUE_ID", unique_id
|
||||
).encode(),
|
||||
media_type=response.content_type,
|
||||
status_code=response.status,
|
||||
|
|
@ -227,13 +236,15 @@ class Provider:
|
|||
if response.frame_delay_ms > 0:
|
||||
async def stream() -> AsyncIterator[bytes]:
|
||||
for frame in response.frames:
|
||||
yield f"{frame.replace('$REQUEST_ID', scenario_id)}\n\n".encode()
|
||||
yield (
|
||||
f"{frame.replace('$REQUEST_ID', scenario_id).replace('$UNIQUE_ID', unique_id)}\n\n"
|
||||
).encode()
|
||||
await asyncio.sleep(response.frame_delay_ms / 1000)
|
||||
|
||||
return StreamingResponse(stream(), media_type=response.content_type)
|
||||
stream_body: Final = ("\n\n".join(response.frames) + "\n\n").replace(
|
||||
"$REQUEST_ID", scenario_id
|
||||
)
|
||||
).replace("$UNIQUE_ID", unique_id)
|
||||
return Response(content=stream_body.encode(), media_type=response.content_type)
|
||||
case EventStreamResponse():
|
||||
events: Final = (
|
||||
|
|
@ -244,6 +255,7 @@ class Provider:
|
|||
"bytes": base64.b64encode(
|
||||
json.dumps(event.payload, separators=(",", ":"))
|
||||
.replace("$REQUEST_ID", scenario_id)
|
||||
.replace("$UNIQUE_ID", unique_id)
|
||||
.encode()
|
||||
).decode(),
|
||||
},
|
||||
|
|
@ -254,7 +266,7 @@ class Provider:
|
|||
else response.events
|
||||
)
|
||||
event_body: Final = b"".join(
|
||||
_aws_event_frame(event.event_type, event.payload, scenario_id) for event in events
|
||||
_aws_event_frame(event.event_type, event.payload, scenario_id, unique_id) for event in events
|
||||
)
|
||||
return Response(content=event_body, media_type=response.content_type)
|
||||
|
||||
|
|
|
|||
|
|
@ -62,12 +62,12 @@ class FailureRow(BaseModel):
|
|||
|
||||
|
||||
class DailySpend(BaseModel):
|
||||
model_config = ConfigDict(extra="ignore")
|
||||
model_config = ConfigDict(frozen=True, extra="forbid")
|
||||
|
||||
spend: float | None = None
|
||||
prompt_tokens: int | None = None
|
||||
completion_tokens: int | None = None
|
||||
api_requests: int | None = None
|
||||
spend: float
|
||||
prompt_tokens: int
|
||||
completion_tokens: int
|
||||
api_requests: int
|
||||
|
||||
|
||||
class Rollups(BaseModel):
|
||||
|
|
@ -245,7 +245,6 @@ def register_scenario_deployment(
|
|||
*,
|
||||
response: StoredResponse | None = None,
|
||||
marker_suffix: str = "",
|
||||
model_name: str | None = None,
|
||||
) -> RegisteredDeployment:
|
||||
control_url: Final = os.environ["INTEGRATION_UPSTREAM_URL"].rstrip("/")
|
||||
run_marker: Final = sha256(key.encode()).hexdigest()[:12]
|
||||
|
|
@ -254,7 +253,7 @@ def register_scenario_deployment(
|
|||
case.response if response is None else response,
|
||||
)
|
||||
scenario.cleanups.callback(delete_scenario, handle)
|
||||
registered_model_name: Final = model_name or f"cost-{marker}{marker_suffix}-{run_marker}"
|
||||
registered_model_name: Final = f"cost-{marker}{marker_suffix}-{run_marker}"
|
||||
parameters: Final = {
|
||||
"model": case.litellm_model,
|
||||
"api_key": case.api_key,
|
||||
|
|
|
|||
|
|
@ -177,6 +177,7 @@ class RecountExpected(BaseModel):
|
|||
recount: RecountRates
|
||||
prompt_tokens: int | None = None
|
||||
completion_tokens: int | None = None
|
||||
min_completion_tokens: int | None = None
|
||||
|
||||
|
||||
class FailureDetails(BaseModel):
|
||||
|
|
@ -463,6 +464,23 @@ def data_errors() -> tuple[str, ...]:
|
|||
or not isinstance(case.expected, RecountExpected)
|
||||
)
|
||||
)
|
||||
invalid_rollup_ids: Final = sorted(
|
||||
case.name
|
||||
for case in CASES
|
||||
if isinstance(case.expected, ExactExpected)
|
||||
and case.expected.rollups
|
||||
and "$UNIQUE_ID" not in case.response.model_dump_json()
|
||||
)
|
||||
invalid_pinned_tool_ids: Final = sorted(
|
||||
case.name
|
||||
for case in CASES
|
||||
if isinstance(case.expected, RecountExpected)
|
||||
and (case.expected.prompt_tokens is not None or case.expected.completion_tokens is not None)
|
||||
and any(
|
||||
marker in case.response.model_dump_json()
|
||||
for marker in ('"id": "call_$REQUEST_ID"', '"id": "toolu_$REQUEST_ID"')
|
||||
)
|
||||
)
|
||||
return tuple(
|
||||
message
|
||||
for message in (
|
||||
|
|
@ -478,6 +496,10 @@ def data_errors() -> tuple[str, ...]:
|
|||
f"invalid passthrough opt-outs: {invalid_opt_outs}" if invalid_opt_outs else None,
|
||||
f"invalid fallback responses: {invalid_fallbacks}" if invalid_fallbacks else None,
|
||||
f"invalid disconnect cases: {invalid_disconnects}" if invalid_disconnects else None,
|
||||
f"rollup responses lack $UNIQUE_ID: {invalid_rollup_ids}" if invalid_rollup_ids else None,
|
||||
f"pinned tool IDs contain $REQUEST_ID: {invalid_pinned_tool_ids}"
|
||||
if invalid_pinned_tool_ids
|
||||
else None,
|
||||
)
|
||||
if message is not None
|
||||
)
|
||||
|
|
|
|||
|
|
@ -6930,7 +6930,7 @@
|
|||
"response": {
|
||||
"content_type": "application/json",
|
||||
"body": {
|
||||
"id": "msg_$REQUEST_ID",
|
||||
"id": "msg_$UNIQUE_ID",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"model": "claude-sonnet-5",
|
||||
|
|
@ -7690,7 +7690,7 @@
|
|||
"content_type": "text/event-stream",
|
||||
"frames": [
|
||||
"event: message_start\ndata: {\"type\": \"message_start\", \"message\": {\"id\": \"msg_$REQUEST_ID\", \"type\": \"message\", \"role\": \"assistant\", \"model\": \"claude-sonnet-5\", \"content\": [], \"stop_reason\": null}}",
|
||||
"event: content_block_start\ndata: {\"type\": \"content_block_start\", \"index\": 0, \"content_block\": {\"type\": \"tool_use\", \"id\": \"toolu_$REQUEST_ID\", \"name\": \"get_weather\", \"input\": {}}}",
|
||||
"event: content_block_start\ndata: {\"type\": \"content_block_start\", \"index\": 0, \"content_block\": {\"type\": \"tool_use\", \"id\": \"call_fixture_0001\", \"name\": \"get_weather\", \"input\": {}}}",
|
||||
"event: content_block_delta\ndata: {\"type\": \"content_block_delta\", \"index\": 0, \"delta\": {\"type\": \"input_json_delta\", \"partial_json\": \"{\\\"city\\\": \\\"Berlin\\\", \\\"days\\\": 7, \\\"units\\\": \\\"metric\\\", \\\"notes\\\": \\\"filler filler filler filler fil\"}}",
|
||||
"event: content_block_delta\ndata: {\"type\": \"content_block_delta\", \"index\": 0, \"delta\": {\"type\": \"input_json_delta\", \"partial_json\": \"ler filler filler filler filler filler filler filler filler filler filler filler filler fi\"}}",
|
||||
"event: content_block_delta\ndata: {\"type\": \"content_block_delta\", \"index\": 0, \"delta\": {\"type\": \"input_json_delta\", \"partial_json\": \"ller filler filler filler filler filler filler filler filler filler filler filler filler \\\"}\"}}",
|
||||
|
|
@ -7704,8 +7704,7 @@
|
|||
"input_cost_per_token": 3e-06,
|
||||
"output_cost_per_token": 1.5e-05
|
||||
},
|
||||
"prompt_tokens": 49,
|
||||
"completion_tokens": 111
|
||||
"min_completion_tokens": 60
|
||||
}
|
||||
},
|
||||
{
|
||||
|
|
@ -12877,8 +12876,7 @@
|
|||
"input_cost_per_token": 5.2e-07,
|
||||
"output_cost_per_token": 3.12e-06
|
||||
},
|
||||
"prompt_tokens": 46,
|
||||
"completion_tokens": 88
|
||||
"min_completion_tokens": 60
|
||||
}
|
||||
},
|
||||
{
|
||||
|
|
@ -20752,7 +20750,7 @@
|
|||
"response": {
|
||||
"content_type": "application/json",
|
||||
"body": {
|
||||
"id": "chatcmpl-$REQUEST_ID",
|
||||
"id": "chatcmpl-$UNIQUE_ID",
|
||||
"object": "chat.completion",
|
||||
"created": 1789788262,
|
||||
"model": "gpt-5.6",
|
||||
|
|
@ -21610,7 +21608,7 @@
|
|||
"content_type": "text/event-stream",
|
||||
"frames": [
|
||||
"data: {\"id\": \"chatcmpl-$REQUEST_ID\", \"object\": \"chat.completion.chunk\", \"created\": 1789788263, \"model\": \"gpt-5.6\", \"choices\": [{\"index\": 0, \"delta\": {\"role\": \"assistant\"}, \"finish_reason\": null}], \"usage\": null}",
|
||||
"data: {\"id\": \"chatcmpl-$REQUEST_ID\", \"object\": \"chat.completion.chunk\", \"created\": 1789788263, \"model\": \"gpt-5.6\", \"choices\": [{\"index\": 0, \"delta\": {\"role\": \"assistant\", \"tool_calls\": [{\"index\": 0, \"id\": \"call_$REQUEST_ID\", \"type\": \"function\", \"function\": {\"name\": \"get_weather\", \"arguments\": \"\"}}]}, \"finish_reason\": null}], \"usage\": null}",
|
||||
"data: {\"id\": \"chatcmpl-$REQUEST_ID\", \"object\": \"chat.completion.chunk\", \"created\": 1789788263, \"model\": \"gpt-5.6\", \"choices\": [{\"index\": 0, \"delta\": {\"role\": \"assistant\", \"tool_calls\": [{\"index\": 0, \"id\": \"call_fixture_0001\", \"type\": \"function\", \"function\": {\"name\": \"get_weather\", \"arguments\": \"\"}}]}, \"finish_reason\": null}], \"usage\": null}",
|
||||
"data: {\"id\": \"chatcmpl-$REQUEST_ID\", \"object\": \"chat.completion.chunk\", \"created\": 1789788263, \"model\": \"gpt-5.6\", \"choices\": [{\"index\": 0, \"delta\": {\"tool_calls\": [{\"index\": 0, \"function\": {\"arguments\": \"{\\\"city\\\": \\\"Berlin\\\", \\\"days\\\": 7, \\\"units\\\": \\\"metric\\\", \\\"notes\\\": \\\"filler filler filler filler fil\"}}]}, \"finish_reason\": null}], \"usage\": null}",
|
||||
"data: {\"id\": \"chatcmpl-$REQUEST_ID\", \"object\": \"chat.completion.chunk\", \"created\": 1789788263, \"model\": \"gpt-5.6\", \"choices\": [{\"index\": 0, \"delta\": {\"tool_calls\": [{\"index\": 0, \"function\": {\"arguments\": \"ler filler filler filler filler filler filler filler filler filler filler filler filler fi\"}}]}, \"finish_reason\": null}], \"usage\": null}",
|
||||
"data: {\"id\": \"chatcmpl-$REQUEST_ID\", \"object\": \"chat.completion.chunk\", \"created\": 1789788263, \"model\": \"gpt-5.6\", \"choices\": [{\"index\": 0, \"delta\": {\"tool_calls\": [{\"index\": 0, \"function\": {\"arguments\": \"ller filler filler filler filler filler filler filler filler filler filler filler filler \\\"}\"}}]}, \"finish_reason\": null}], \"usage\": null}",
|
||||
|
|
@ -21623,8 +21621,7 @@
|
|||
"input_cost_per_token": 1.75e-06,
|
||||
"output_cost_per_token": 1.4e-05
|
||||
},
|
||||
"prompt_tokens": 44,
|
||||
"completion_tokens": 105
|
||||
"min_completion_tokens": 60
|
||||
}
|
||||
},
|
||||
{
|
||||
|
|
|
|||
|
|
@ -209,6 +209,35 @@ def _assert_exact(
|
|||
assert_total_is_sum_of_components(row, breakdown, case.name)
|
||||
|
||||
|
||||
def _assert_recount(case: CostTrackingTestCase, expected: RecountExpected, row: CostRow) -> None:
|
||||
assert row.prompt_tokens is not None and row.prompt_tokens > 0, (
|
||||
f"{case.name}: recount case counted no input tokens: prompt_tokens={row.prompt_tokens}"
|
||||
)
|
||||
assert row.completion_tokens is not None and row.completion_tokens > 0, (
|
||||
f"{case.name}: recount case counted no output tokens: completion_tokens={row.completion_tokens}"
|
||||
)
|
||||
if expected.prompt_tokens is not None:
|
||||
assert row.prompt_tokens == expected.prompt_tokens, (
|
||||
f"{case.name}: prompt_tokens {row.prompt_tokens} != pinned {expected.prompt_tokens}"
|
||||
)
|
||||
if expected.completion_tokens is not None:
|
||||
assert row.completion_tokens == expected.completion_tokens, (
|
||||
f"{case.name}: completion_tokens {row.completion_tokens} != pinned {expected.completion_tokens}"
|
||||
)
|
||||
if expected.min_completion_tokens is not None:
|
||||
assert row.completion_tokens >= expected.min_completion_tokens, (
|
||||
f"{case.name}: completion_tokens {row.completion_tokens} < minimum {expected.min_completion_tokens}"
|
||||
)
|
||||
recount: Final = row.prompt_tokens * expected.recount.input_cost_per_token + (
|
||||
row.completion_tokens * expected.recount.output_cost_per_token
|
||||
)
|
||||
assert row.spend is not None and approx_equal(row.spend, recount), (
|
||||
f"{case.name}: spend {row.spend} != recount {recount} at map rates"
|
||||
)
|
||||
assert row.breakdown is not None, f"{case.name}: no cost_breakdown persisted"
|
||||
assert_total_is_sum_of_components(row, row.breakdown, case.name)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("case", _CASES)
|
||||
def test_case_bills_expected_cost(gateway: Gateway, case: CostTrackingTestCase) -> None:
|
||||
marker: Final = sha256(case.name.encode()).hexdigest()[:12]
|
||||
|
|
@ -251,20 +280,6 @@ def test_case_bills_expected_cost(gateway: Gateway, case: CostTrackingTestCase)
|
|||
if case.fallback_from is not None
|
||||
else None
|
||||
)
|
||||
if isinstance(expected, ExactExpected) and expected.rollups:
|
||||
assert deployment is not None
|
||||
rollup_deployments: Final = tuple(
|
||||
register_scenario_deployment(
|
||||
scenario,
|
||||
case,
|
||||
marker,
|
||||
key,
|
||||
marker_suffix=f"-r{index}",
|
||||
model_name=deployment.model_name,
|
||||
)
|
||||
for index in (2, 3)
|
||||
)
|
||||
assert len(rollup_deployments) == 2
|
||||
model_name: Final = (
|
||||
case.model
|
||||
if passthrough_provider in {"gemini", "anthropic"}
|
||||
|
|
@ -334,18 +349,8 @@ def test_case_bills_expected_cost(gateway: Gateway, case: CostTrackingTestCase)
|
|||
assert len(frames) == case.disconnect_after_frames
|
||||
row: Final = poll_cost_row(key)
|
||||
assert isinstance(expected, RecountExpected)
|
||||
if expected.prompt_tokens is not None:
|
||||
assert row.prompt_tokens == expected.prompt_tokens
|
||||
if expected.completion_tokens is not None:
|
||||
assert row.completion_tokens == expected.completion_tokens
|
||||
assert row.prompt_tokens is not None and row.prompt_tokens > 0
|
||||
assert row.completion_tokens is not None and row.completion_tokens > 0
|
||||
recount: Final = row.prompt_tokens * expected.recount.input_cost_per_token + (
|
||||
row.completion_tokens * expected.recount.output_cost_per_token
|
||||
)
|
||||
assert row.spend is not None and approx_equal(row.spend, recount)
|
||||
assert row.breakdown is not None
|
||||
assert_total_is_sum_of_components(row, row.breakdown, case.name)
|
||||
assert row.status == "success", f"{case.name}: disconnect row status was {row.status}"
|
||||
_assert_recount(case, expected, row)
|
||||
return
|
||||
responses: Final = tuple(
|
||||
(
|
||||
|
|
@ -374,21 +379,7 @@ def test_case_bills_expected_cost(gateway: Gateway, case: CostTrackingTestCase)
|
|||
rows: Final = poll_rows(key) if len(responses) > 1 else (poll_cost_row(key),)
|
||||
if isinstance(expected, RecountExpected):
|
||||
row: Final = rows[0]
|
||||
assert row.prompt_tokens is not None and row.prompt_tokens > 0, (
|
||||
f"{case.name}: recount case counted no input tokens: prompt_tokens={row.prompt_tokens}"
|
||||
)
|
||||
assert row.completion_tokens is not None and row.completion_tokens > 0, (
|
||||
f"{case.name}: recount case counted no output tokens: completion_tokens={row.completion_tokens}"
|
||||
)
|
||||
recount: Final = row.prompt_tokens * case.expected.recount.input_cost_per_token + (
|
||||
row.completion_tokens * case.expected.recount.output_cost_per_token
|
||||
)
|
||||
assert row.spend is not None and approx_equal(row.spend, recount), (
|
||||
f"{case.name}: spend {row.spend} != recount {recount} at map rates"
|
||||
)
|
||||
breakdown: Final = row.breakdown
|
||||
assert breakdown is not None, f"{case.name}: no cost_breakdown persisted"
|
||||
assert_total_is_sum_of_components(row, breakdown, case.name)
|
||||
_assert_recount(case, expected, row)
|
||||
return
|
||||
assert isinstance(expected, ExactExpected)
|
||||
if fallback_deployment is not None:
|
||||
|
|
@ -423,8 +414,8 @@ def test_case_bills_expected_cost(gateway: Gateway, case: CostTrackingTestCase)
|
|||
assert approx_equal(rollups.team_spend, target_spend)
|
||||
assert approx_equal(rollups.user_spend, target_spend)
|
||||
assert approx_equal(rollups.end_user_spend, target_spend)
|
||||
assert approx_equal(rollups.daily_user.spend or 0.0, target_spend)
|
||||
assert approx_equal(rollups.daily_team.spend or 0.0, target_spend)
|
||||
assert approx_equal(rollups.daily_user.spend, target_spend)
|
||||
assert approx_equal(rollups.daily_team.spend, target_spend)
|
||||
assert rollups.daily_user.prompt_tokens == expected.prompt_tokens * 3
|
||||
assert rollups.daily_user.completion_tokens == expected.completion_tokens * 3
|
||||
assert rollups.daily_user.api_requests == 3
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue