mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-25 01:02:15 +00:00
test(integration): completed batch spend row records reasoning tokens and error file failures (Pylon #7928)
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
kerry
Devin AI
parent
22d86f4fd3
commit
1ac0f47f4e
1 changed files with 202 additions and 0 deletions
202
tests/integration/spend/test_batch_completion_accounting.py
Normal file
202
tests/integration/spend/test_batch_completion_accounting.py
Normal file
|
|
@ -0,0 +1,202 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import uuid
|
||||
from hashlib import sha256
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
from integration._support.client import JSON_OBJECT, Gateway, eventually, string_value
|
||||
from integration._support.database import read_rows
|
||||
from integration._support.upstream import delete_scenario, register_scenario
|
||||
from integration.cost_calculation.cost_tracking_case import JsonResponse, RoutedResponse, TextResponse
|
||||
from pydantic import JsonValue
|
||||
|
||||
FIRST_LINE: Final = {"prompt_tokens": 10, "completion_tokens": 7, "reasoning_tokens": 4}
|
||||
SECOND_LINE: Final = {"prompt_tokens": 5, "completion_tokens": 3, "reasoning_tokens": 2}
|
||||
ERROR_FILE_LINES: Final = 2
|
||||
|
||||
|
||||
def _succeeded_line(index: int, model: str, prompt_tokens: int, completion_tokens: int, reasoning_tokens: int) -> str:
|
||||
return json.dumps(
|
||||
{
|
||||
"id": f"batch_req_{index}",
|
||||
"custom_id": f"r{index}",
|
||||
"response": {
|
||||
"status_code": 200,
|
||||
"request_id": f"$REQUEST_ID-{index}",
|
||||
"body": {
|
||||
"id": f"chatcmpl-$REQUEST_ID-{index}",
|
||||
"object": "chat.completion",
|
||||
"model": model,
|
||||
"choices": [
|
||||
{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}
|
||||
],
|
||||
"usage": {
|
||||
"prompt_tokens": prompt_tokens,
|
||||
"completion_tokens": completion_tokens,
|
||||
"total_tokens": prompt_tokens + completion_tokens,
|
||||
"completion_tokens_details": {"reasoning_tokens": reasoning_tokens},
|
||||
},
|
||||
},
|
||||
},
|
||||
"error": None,
|
||||
},
|
||||
separators=(",", ":"),
|
||||
)
|
||||
|
||||
|
||||
def _failed_line(index: int) -> str:
|
||||
return json.dumps(
|
||||
{
|
||||
"id": f"batch_req_{index}",
|
||||
"custom_id": f"r{index}",
|
||||
"response": {
|
||||
"status_code": 400,
|
||||
"request_id": f"$REQUEST_ID-{index}",
|
||||
"body": {"error": {"message": "rejected line", "type": "invalid_request_error", "code": "400"}},
|
||||
},
|
||||
"error": {"code": "bad_request", "message": "rejected line"},
|
||||
},
|
||||
separators=(",", ":"),
|
||||
)
|
||||
|
||||
|
||||
def _batch_routes(model: str) -> RoutedResponse:
|
||||
output_lines: Final = (
|
||||
_succeeded_line(1, model, **FIRST_LINE),
|
||||
_succeeded_line(2, model, **SECOND_LINE),
|
||||
_failed_line(3),
|
||||
)
|
||||
error_lines: Final = tuple(_failed_line(index) for index in range(4, 4 + ERROR_FILE_LINES))
|
||||
completed: Final = {
|
||||
"id": "batch-$REQUEST_ID",
|
||||
"object": "batch",
|
||||
"endpoint": "/v1/chat/completions",
|
||||
"errors": None,
|
||||
"input_file_id": "file-in-$REQUEST_ID",
|
||||
"completion_window": "24h",
|
||||
"status": "completed",
|
||||
"output_file_id": "file-out-$REQUEST_ID",
|
||||
"error_file_id": "file-err-$REQUEST_ID",
|
||||
"created_at": 1,
|
||||
"in_progress_at": 1,
|
||||
"completed_at": 1,
|
||||
"expires_at": 1,
|
||||
"request_counts": {"total": 5, "completed": 2, "failed": 3},
|
||||
"metadata": None,
|
||||
}
|
||||
return RoutedResponse(
|
||||
content_type="application/x-routed",
|
||||
routes={
|
||||
"POST /files": JsonResponse(
|
||||
content_type="application/json",
|
||||
body={
|
||||
"id": "file-in-$REQUEST_ID",
|
||||
"object": "file",
|
||||
"purpose": "batch",
|
||||
"bytes": 100,
|
||||
"created_at": 1,
|
||||
"filename": "in.jsonl",
|
||||
"status": "processed",
|
||||
},
|
||||
),
|
||||
"POST /batches": JsonResponse(
|
||||
content_type="application/json",
|
||||
body={**completed, "status": "validating", "output_file_id": None, "error_file_id": None},
|
||||
),
|
||||
"GET /batches/batch-$REQUEST_ID": JsonResponse(content_type="application/json", body=completed),
|
||||
"GET /files/file-out-$REQUEST_ID/content": TextResponse(
|
||||
content_type="application/jsonl", body="\n".join(output_lines) + "\n"
|
||||
),
|
||||
"GET /files/file-err-$REQUEST_ID/content": TextResponse(
|
||||
content_type="application/jsonl", body="\n".join(error_lines) + "\n"
|
||||
),
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def _input_file(model: str) -> bytes:
|
||||
return (
|
||||
"\n".join(
|
||||
json.dumps(
|
||||
{
|
||||
"custom_id": f"r{index}",
|
||||
"method": "POST",
|
||||
"url": "/v1/chat/completions",
|
||||
"body": {"model": model, "messages": [{"role": "user", "content": "batch accounting"}]},
|
||||
},
|
||||
separators=(",", ":"),
|
||||
)
|
||||
for index in range(1, 6)
|
||||
)
|
||||
+ "\n"
|
||||
).encode()
|
||||
|
||||
|
||||
def _metadata(value: object) -> dict[str, JsonValue]:
|
||||
return JSON_OBJECT.validate_json(value) if isinstance(value, str) else JSON_OBJECT.validate_python(value)
|
||||
|
||||
|
||||
@pytest.mark.covers("quota_management.spend_tracking.batch_costs.reasoning_tokens_and_error_file_failures_recorded")
|
||||
def test_completed_batch_spend_row_records_reasoning_tokens_and_error_file_failures(gateway: Gateway) -> None:
|
||||
with gateway.scenario() as scenario:
|
||||
key: Final = scenario.key()
|
||||
scenario_id: Final = f"batch-accounting-{uuid.uuid4().hex[:12]}"
|
||||
handle: Final = register_scenario(scenario_id, _batch_routes("gpt-4o-mini"))
|
||||
scenario.cleanups.callback(delete_scenario, handle)
|
||||
model: Final = scenario.model(api_base=handle.api_base())
|
||||
file_response: Final = gateway.request_multipart(
|
||||
"/v1/files",
|
||||
{"purpose": "batch", "model": model},
|
||||
{"file": ("in.jsonl", _input_file(model), "application/jsonl")},
|
||||
key=key,
|
||||
)
|
||||
assert file_response.status_code == 200, file_response.text
|
||||
batch_response: Final = gateway.request(
|
||||
"POST",
|
||||
"/v1/batches",
|
||||
{
|
||||
"input_file_id": string_value(JSON_OBJECT.validate_json(file_response.content)["id"]),
|
||||
"endpoint": "/v1/chat/completions",
|
||||
"completion_window": "24h",
|
||||
"model": model,
|
||||
},
|
||||
key=key,
|
||||
)
|
||||
assert batch_response.status_code == 200, batch_response.text
|
||||
batch_id: Final = string_value(JSON_OBJECT.validate_json(batch_response.content)["id"])
|
||||
retrieval: Final = gateway.request("GET", f"/v1/batches/{batch_id}", key=key)
|
||||
assert retrieval.status_code == 200, retrieval.text
|
||||
assert retrieval.json()["status"] == "completed", retrieval.text
|
||||
rows: Final = eventually(
|
||||
lambda: read_rows(
|
||||
'SELECT status, prompt_tokens, completion_tokens, metadata FROM "LiteLLM_SpendLogs" '
|
||||
"WHERE api_key=%s AND call_type='aretrieve_batch'",
|
||||
(sha256(key.encode()).hexdigest(),),
|
||||
),
|
||||
lambda values: len(values) == 1,
|
||||
seconds=70,
|
||||
)
|
||||
row: Final = rows[0]
|
||||
metadata: Final = _metadata(row["metadata"])
|
||||
prompt_tokens: Final = FIRST_LINE["prompt_tokens"] + SECOND_LINE["prompt_tokens"]
|
||||
completion_tokens: Final = FIRST_LINE["completion_tokens"] + SECOND_LINE["completion_tokens"]
|
||||
reasoning_tokens: Final = FIRST_LINE["reasoning_tokens"] + SECOND_LINE["reasoning_tokens"]
|
||||
assert row["status"] == "success", retrieval.text
|
||||
assert (row["prompt_tokens"], row["completion_tokens"]) == (prompt_tokens, completion_tokens), retrieval.text
|
||||
assert (metadata["batch_successful_requests"], metadata["batch_failed_requests"]) == (
|
||||
2,
|
||||
1 + ERROR_FILE_LINES,
|
||||
), json.dumps(metadata)
|
||||
usage: Final = JSON_OBJECT.validate_python(metadata["usage_object"])
|
||||
details: Final = JSON_OBJECT.validate_python(usage["completion_tokens_details"])
|
||||
assert (usage["prompt_tokens"], usage["completion_tokens"], usage["total_tokens"]) == (
|
||||
prompt_tokens,
|
||||
completion_tokens,
|
||||
prompt_tokens + completion_tokens,
|
||||
), json.dumps(metadata)
|
||||
assert {name: value for name, value in details.items() if value is not None} == {
|
||||
"reasoning_tokens": reasoning_tokens,
|
||||
"text_tokens": completion_tokens - reasoning_tokens,
|
||||
}, json.dumps(metadata)
|
||||
Loading…
Add table
Reference in a new issue