mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-24 00:52:24 +00:00
Merge pull request #42028 from BerriAI/litellm_cost_shard_passthrough
test(integration): passthrough route cost cases
This commit is contained in:
kerry-berri
GitHub
commit
b5f02f72fe
7 changed files with 515 additions and 32 deletions
|
|
@ -121,6 +121,10 @@ start_proxy() {
|
|||
"LITELLM_MODEL_COST_MAP_URL=$INTEGRATION_UPSTREAM_URL/_cost_map"
|
||||
"MODEL_COST_MAP_MIN_MODEL_COUNT=1"
|
||||
"MODEL_COST_MAP_MAX_SHRINK_RATIO=0"
|
||||
"GEMINI_API_BASE=$INTEGRATION_UPSTREAM_URL"
|
||||
"ANTHROPIC_API_BASE=$INTEGRATION_UPSTREAM_URL"
|
||||
"GEMINI_API_KEY=sk-scripted-provider"
|
||||
"ANTHROPIC_API_KEY=sk-scripted-provider"
|
||||
)
|
||||
else
|
||||
cost_map_env=("LITELLM_LOCAL_MODEL_COST_MAP=True")
|
||||
|
|
|
|||
|
|
@ -58,13 +58,18 @@ class Gateway:
|
|||
*,
|
||||
key: str | None = None,
|
||||
params: Mapping[str, str] | None = None,
|
||||
headers: Mapping[str, str] | None = None,
|
||||
) -> httpx.Response:
|
||||
request_headers: Final = {
|
||||
"Authorization": f"Bearer {self.key if key is None else key}",
|
||||
**(headers or {}),
|
||||
}
|
||||
return self.client.request(
|
||||
method,
|
||||
path,
|
||||
json=body,
|
||||
params=params,
|
||||
headers={"Authorization": f"Bearer {self.key if key is None else key}"},
|
||||
headers=request_headers,
|
||||
)
|
||||
|
||||
def request_multipart(
|
||||
|
|
|
|||
|
|
@ -194,9 +194,11 @@ class Provider:
|
|||
|
||||
async def scripted(self, request: Request) -> Response:
|
||||
segments: Final = tuple(segment for segment in cast(str, request.path_params["path"]).split("/") if segment)
|
||||
if not segments:
|
||||
return JSONResponse({"error": "Unknown scenario"}, status_code=404)
|
||||
scenario_id: Final = segments[0].split(":", 1)[0]
|
||||
scenario_id: Final = (
|
||||
segments[0].split(":", 1)[0]
|
||||
if segments and self.scenario_store.get(segments[0].split(":", 1)[0]) is not None
|
||||
else request.headers.get("x-scripted-scenario", "")
|
||||
)
|
||||
response: Final = self.scenario_store.get(scenario_id)
|
||||
if response is None:
|
||||
return JSONResponse({"error": "Unknown scenario"}, status_code=404)
|
||||
|
|
|
|||
|
|
@ -1516,6 +1516,24 @@
|
|||
"tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[anthropic.claude-sonnet-5-v1:0-messages_cache_read]": [
|
||||
"quota_management.spend_tracking.cost_matrix.logs_cost"
|
||||
],
|
||||
"tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[gemini-3.1-pro-passthrough-generate_content_priced_via_gemini_key]": [
|
||||
"quota_management.spend_tracking.cost_matrix.logs_cost"
|
||||
],
|
||||
"tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[gemini-3.1-pro-passthrough-stream_generate_content_priced_via_vertex_key]": [
|
||||
"quota_management.spend_tracking.scripted_wire.logs_cost"
|
||||
],
|
||||
"tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[claude-sonnet-5-passthrough-messages]": [
|
||||
"quota_management.spend_tracking.cost_matrix.logs_cost"
|
||||
],
|
||||
"tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[claude-sonnet-5-passthrough-messages_cache_read]": [
|
||||
"quota_management.spend_tracking.cost_matrix.logs_cost"
|
||||
],
|
||||
"tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[anthropic.claude-sonnet-5-v1:0-passthrough-converse]": [
|
||||
"quota_management.spend_tracking.cost_matrix.logs_cost"
|
||||
],
|
||||
"tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[anthropic.claude-sonnet-5-v1:0-passthrough-converse_stream]": [
|
||||
"quota_management.spend_tracking.scripted_wire.logs_cost"
|
||||
],
|
||||
"tests/integration/mcp/test_mcp_lifecycle.py::test_health_intersects_route_restricted_key_grants_in_both_management_modes": [
|
||||
"other.mcp.health.restricted_keys_intersect_grants_in_both_modes"
|
||||
],
|
||||
|
|
|
|||
|
|
@ -147,6 +147,7 @@ class ExactExpected(BaseModel):
|
|||
reasoning_cost: float | None = None
|
||||
tool_usage_cost: float | None = None
|
||||
breakdown_persisted: bool = True
|
||||
cost_header: bool = True
|
||||
|
||||
|
||||
class RecountRates(BaseModel):
|
||||
|
|
@ -183,19 +184,22 @@ class CostTrackingTestCase(BaseModel):
|
|||
name: str
|
||||
covers: str
|
||||
model: str
|
||||
endpoint: Literal[
|
||||
"/v1/chat/completions",
|
||||
"/v1/responses",
|
||||
"/v1/messages",
|
||||
"/v1/embeddings",
|
||||
"/v1/rerank",
|
||||
"/v1/completions",
|
||||
"/v1/moderations",
|
||||
"/v1/audio/transcriptions",
|
||||
"/v1/audio/speech",
|
||||
"/v1/images/generations",
|
||||
"/v1/images/edits",
|
||||
] = "/v1/chat/completions"
|
||||
endpoint: (
|
||||
Literal[
|
||||
"/v1/chat/completions",
|
||||
"/v1/responses",
|
||||
"/v1/messages",
|
||||
"/v1/embeddings",
|
||||
"/v1/rerank",
|
||||
"/v1/completions",
|
||||
"/v1/moderations",
|
||||
"/v1/audio/transcriptions",
|
||||
"/v1/audio/speech",
|
||||
"/v1/images/generations",
|
||||
"/v1/images/edits",
|
||||
]
|
||||
| Annotated[str, Field(pattern=r"^/(gemini|anthropic|bedrock)/")]
|
||||
) = "/v1/chat/completions"
|
||||
deployment: Deployment | None = None
|
||||
upload: Upload | None = None
|
||||
request: dict[str, JsonValue]
|
||||
|
|
@ -239,6 +243,17 @@ class CostTrackingTestCase(BaseModel):
|
|||
def base_model(self) -> str | None:
|
||||
return self.deployment.base_model if self.deployment else None
|
||||
|
||||
@property
|
||||
def passthrough_provider(self) -> Literal["gemini", "anthropic", "bedrock"] | None:
|
||||
provider: Final = self.endpoint.removeprefix("/").split("/", 1)[0]
|
||||
if provider == "gemini":
|
||||
return "gemini"
|
||||
if provider == "anthropic":
|
||||
return "anthropic"
|
||||
if provider == "bedrock":
|
||||
return "bedrock"
|
||||
return None
|
||||
|
||||
|
||||
class _CasesFile(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="forbid")
|
||||
|
|
@ -374,6 +389,19 @@ def data_errors() -> tuple[str, ...]:
|
|||
and case.response.status != 200
|
||||
)
|
||||
)
|
||||
invalid_opt_outs: Final = sorted(
|
||||
case.name
|
||||
for case in CASES
|
||||
if isinstance(case.expected, ExactExpected)
|
||||
and (
|
||||
(
|
||||
not case.expected.breakdown_persisted
|
||||
and case.passthrough_provider is None
|
||||
and case.rates.mode != "image_generation"
|
||||
)
|
||||
or (not case.expected.cost_header and case.passthrough_provider is None)
|
||||
)
|
||||
)
|
||||
return tuple(
|
||||
message
|
||||
for message in (
|
||||
|
|
@ -386,6 +414,7 @@ def data_errors() -> tuple[str, ...]:
|
|||
f"failure response statuses are inconsistent: {failure_response_mismatches}"
|
||||
if failure_response_mismatches
|
||||
else None,
|
||||
f"invalid passthrough opt-outs: {invalid_opt_outs}" if invalid_opt_outs else None,
|
||||
)
|
||||
if message is not None
|
||||
)
|
||||
|
|
|
|||
|
|
@ -28397,6 +28397,379 @@
|
|||
"completion_tokens": 380,
|
||||
"cache_read_cost": 0.00405504
|
||||
}
|
||||
},
|
||||
{
|
||||
"name": "gemini-3.1-pro-passthrough-generate_content_priced_via_gemini_key",
|
||||
"covers": "quota_management.spend_tracking.cost_matrix.logs_cost",
|
||||
"model": "gemini/gemini-3.1-pro",
|
||||
"request": {
|
||||
"model": "$MODEL",
|
||||
"messages": [
|
||||
{
|
||||
"role": "system",
|
||||
"content": [
|
||||
{
|
||||
"type": "text",
|
||||
"text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line."
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "text",
|
||||
"text": "5fdf6b7dd9b9 summarize the attached material in one line and name the city weather"
|
||||
}
|
||||
]
|
||||
}
|
||||
],
|
||||
"stream": false,
|
||||
"allowed_openai_params": []
|
||||
},
|
||||
"response": {
|
||||
"content_type": "application/json",
|
||||
"body": {
|
||||
"candidates": [
|
||||
{
|
||||
"content": {
|
||||
"parts": [
|
||||
{
|
||||
"text": "scripted answer 5fdf6b7dd9b9"
|
||||
}
|
||||
],
|
||||
"role": "model"
|
||||
},
|
||||
"finishReason": "STOP",
|
||||
"index": 0
|
||||
}
|
||||
],
|
||||
"usageMetadata": {
|
||||
"promptTokenCount": 1840,
|
||||
"candidatesTokenCount": 412,
|
||||
"totalTokenCount": 2252,
|
||||
"promptTokensDetails": [
|
||||
{
|
||||
"modality": "TEXT",
|
||||
"tokenCount": 1840
|
||||
}
|
||||
]
|
||||
},
|
||||
"modelVersion": "gemini-3.1-pro"
|
||||
}
|
||||
},
|
||||
"expected": {
|
||||
"spend": 0.008624,
|
||||
"input_cost": 0.00368,
|
||||
"output_cost": 0.004944,
|
||||
"prompt_tokens": 1840,
|
||||
"completion_tokens": 412,
|
||||
"breakdown_persisted": false,
|
||||
"cost_header": false
|
||||
},
|
||||
"endpoint": "/gemini/v1beta/models/$MODEL:generateContent"
|
||||
},
|
||||
{
|
||||
"name": "gemini-3.1-pro-passthrough-stream_generate_content_priced_via_vertex_key",
|
||||
"covers": "quota_management.spend_tracking.scripted_wire.logs_cost",
|
||||
"model": "gemini-3.1-pro",
|
||||
"request": {
|
||||
"model": "$MODEL",
|
||||
"messages": [
|
||||
{
|
||||
"role": "system",
|
||||
"content": [
|
||||
{
|
||||
"type": "text",
|
||||
"text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line."
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "text",
|
||||
"text": "52b6a80ff038 summarize the attached material in one line and name the city weather"
|
||||
}
|
||||
]
|
||||
}
|
||||
],
|
||||
"stream": true,
|
||||
"stream_options": {
|
||||
"include_usage": true
|
||||
},
|
||||
"allowed_openai_params": []
|
||||
},
|
||||
"response": {
|
||||
"content_type": "text/event-stream",
|
||||
"frames": [
|
||||
"data: {\"candidates\": [{\"content\": {\"parts\": [{\"text\": \"scripted answer 52b6a80ff038\"}], \"role\": \"model\"}, \"finishReason\": \"STOP\", \"index\": 0}], \"modelVersion\": \"gemini-3.1-pro\"}",
|
||||
"data: {\"candidates\": [], \"usageMetadata\": {\"promptTokenCount\": 1840, \"candidatesTokenCount\": 412, \"totalTokenCount\": 2252, \"promptTokensDetails\": [{\"modality\": \"TEXT\", \"tokenCount\": 1840}]}, \"modelVersion\": \"gemini-3.1-pro\"}"
|
||||
]
|
||||
},
|
||||
"expected": {
|
||||
"spend": 0.0090552,
|
||||
"input_cost": 0.003864,
|
||||
"output_cost": 0.0051912,
|
||||
"prompt_tokens": 1840,
|
||||
"completion_tokens": 412,
|
||||
"breakdown_persisted": false,
|
||||
"cost_header": false
|
||||
},
|
||||
"endpoint": "/gemini/v1beta/models/$MODEL:streamGenerateContent?alt=sse"
|
||||
},
|
||||
{
|
||||
"name": "claude-sonnet-5-passthrough-messages",
|
||||
"covers": "quota_management.spend_tracking.cost_matrix.logs_cost",
|
||||
"model": "claude-sonnet-5",
|
||||
"endpoint": "/anthropic/v1/messages",
|
||||
"request": {
|
||||
"model": "$MODEL",
|
||||
"max_tokens": 412,
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": "summarize the attached material in one line"
|
||||
}
|
||||
]
|
||||
},
|
||||
"response": {
|
||||
"content_type": "application/json",
|
||||
"body": {
|
||||
"id": "msg_$REQUEST_ID",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"model": "claude-sonnet-5",
|
||||
"content": [
|
||||
{
|
||||
"type": "text",
|
||||
"text": "scripted response"
|
||||
}
|
||||
],
|
||||
"stop_reason": "end_turn",
|
||||
"stop_sequence": null,
|
||||
"usage": {
|
||||
"input_tokens": 1840,
|
||||
"output_tokens": 412
|
||||
}
|
||||
}
|
||||
},
|
||||
"expected": {
|
||||
"spend": 0.0117,
|
||||
"input_cost": 0.00552,
|
||||
"output_cost": 0.00618,
|
||||
"prompt_tokens": 1840,
|
||||
"completion_tokens": 412,
|
||||
"breakdown_persisted": false,
|
||||
"cost_header": false
|
||||
}
|
||||
},
|
||||
{
|
||||
"name": "claude-sonnet-5-passthrough-messages_cache_read",
|
||||
"covers": "quota_management.spend_tracking.cost_matrix.logs_cost",
|
||||
"model": "claude-sonnet-5",
|
||||
"endpoint": "/anthropic/v1/messages",
|
||||
"request": {
|
||||
"model": "$MODEL",
|
||||
"max_tokens": 412,
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "text",
|
||||
"text": "cached text",
|
||||
"cache_control": {
|
||||
"type": "ephemeral"
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
"response": {
|
||||
"content_type": "application/json",
|
||||
"body": {
|
||||
"id": "msg_$REQUEST_ID",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"model": "claude-sonnet-5",
|
||||
"content": [
|
||||
{
|
||||
"type": "text",
|
||||
"text": "scripted response"
|
||||
}
|
||||
],
|
||||
"stop_reason": "end_turn",
|
||||
"stop_sequence": null,
|
||||
"usage": {
|
||||
"input_tokens": 640,
|
||||
"output_tokens": 380,
|
||||
"cache_read_input_tokens": 12288
|
||||
}
|
||||
}
|
||||
},
|
||||
"expected": {
|
||||
"spend": 0.0113064,
|
||||
"input_cost": 0.0056064,
|
||||
"output_cost": 0.0057,
|
||||
"prompt_tokens": 12928,
|
||||
"completion_tokens": 380,
|
||||
"cache_read_cost": 0.0036864,
|
||||
"breakdown_persisted": false,
|
||||
"cost_header": false
|
||||
}
|
||||
},
|
||||
{
|
||||
"name": "anthropic.claude-sonnet-5-v1:0-passthrough-converse",
|
||||
"covers": "quota_management.spend_tracking.cost_matrix.logs_cost",
|
||||
"model": "anthropic.claude-sonnet-5-v1:0",
|
||||
"request": {
|
||||
"model": "$MODEL",
|
||||
"messages": [
|
||||
{
|
||||
"role": "system",
|
||||
"content": [
|
||||
{
|
||||
"type": "text",
|
||||
"text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line."
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "text",
|
||||
"text": "9aad4de0556c summarize the attached material in one line and name the city weather"
|
||||
}
|
||||
]
|
||||
}
|
||||
],
|
||||
"stream": false,
|
||||
"allowed_openai_params": []
|
||||
},
|
||||
"response": {
|
||||
"content_type": "application/json",
|
||||
"body": {
|
||||
"output": {
|
||||
"message": {
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{
|
||||
"text": "scripted answer 9aad4de0556c"
|
||||
}
|
||||
]
|
||||
}
|
||||
},
|
||||
"stopReason": "end_turn",
|
||||
"usage": {
|
||||
"inputTokens": 1840,
|
||||
"outputTokens": 412,
|
||||
"totalTokens": 2252
|
||||
},
|
||||
"metrics": {
|
||||
"latencyMs": 42
|
||||
}
|
||||
}
|
||||
},
|
||||
"expected": {
|
||||
"spend": 0.01287,
|
||||
"input_cost": 0.006072,
|
||||
"output_cost": 0.006798,
|
||||
"prompt_tokens": 1840,
|
||||
"completion_tokens": 412,
|
||||
"cost_header": false
|
||||
},
|
||||
"endpoint": "/bedrock/model/$MODEL/converse"
|
||||
},
|
||||
{
|
||||
"name": "anthropic.claude-sonnet-5-v1:0-passthrough-converse_stream",
|
||||
"covers": "quota_management.spend_tracking.scripted_wire.logs_cost",
|
||||
"model": "anthropic.claude-sonnet-5-v1:0",
|
||||
"request": {
|
||||
"model": "$MODEL",
|
||||
"messages": [
|
||||
{
|
||||
"role": "system",
|
||||
"content": [
|
||||
{
|
||||
"type": "text",
|
||||
"text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line."
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "text",
|
||||
"text": "a9257967d38a summarize the attached material in one line and name the city weather"
|
||||
}
|
||||
]
|
||||
}
|
||||
],
|
||||
"stream": true,
|
||||
"stream_options": {
|
||||
"include_usage": true
|
||||
},
|
||||
"allowed_openai_params": []
|
||||
},
|
||||
"response": {
|
||||
"content_type": "application/vnd.amazon.eventstream",
|
||||
"events": [
|
||||
{
|
||||
"event_type": "messageStart",
|
||||
"payload": {
|
||||
"role": "assistant"
|
||||
}
|
||||
},
|
||||
{
|
||||
"event_type": "contentBlockDelta",
|
||||
"payload": {
|
||||
"delta": {
|
||||
"text": "scripted answer a9257967d38a"
|
||||
},
|
||||
"contentBlockIndex": 0
|
||||
}
|
||||
},
|
||||
{
|
||||
"event_type": "contentBlockStop",
|
||||
"payload": {
|
||||
"contentBlockIndex": 0
|
||||
}
|
||||
},
|
||||
{
|
||||
"event_type": "messageStop",
|
||||
"payload": {
|
||||
"stopReason": "end_turn"
|
||||
}
|
||||
},
|
||||
{
|
||||
"event_type": "metadata",
|
||||
"payload": {
|
||||
"usage": {
|
||||
"inputTokens": 1840,
|
||||
"outputTokens": 412,
|
||||
"totalTokens": 2252
|
||||
},
|
||||
"metrics": {
|
||||
"latencyMs": 42
|
||||
}
|
||||
}
|
||||
}
|
||||
]
|
||||
},
|
||||
"expected": {
|
||||
"spend": 0.01287,
|
||||
"input_cost": 0.006072,
|
||||
"output_cost": 0.006798,
|
||||
"prompt_tokens": 1840,
|
||||
"completion_tokens": 412,
|
||||
"cost_header": false
|
||||
},
|
||||
"endpoint": "/bedrock/model/$MODEL/converse-stream"
|
||||
}
|
||||
]
|
||||
}
|
||||
|
|
|
|||
|
|
@ -4,16 +4,16 @@ from __future__ import annotations
|
|||
|
||||
import io
|
||||
import json
|
||||
from hashlib import sha256
|
||||
import struct
|
||||
from typing import Final, cast
|
||||
import wave
|
||||
import zlib
|
||||
from hashlib import sha256
|
||||
from typing import Final, cast
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
from integration._support.client import JSON_OBJECT, Gateway
|
||||
from integration._support.upstream import delete_scenario, register_scenario
|
||||
from integration.cost_calculation.conftest import (
|
||||
CostBreakdown,
|
||||
approx_equal,
|
||||
|
|
@ -23,14 +23,15 @@ from integration.cost_calculation.conftest import (
|
|||
register_scenario_deployment,
|
||||
)
|
||||
from integration.cost_calculation.cost_tracking_case import (
|
||||
BinaryResponse,
|
||||
CASES,
|
||||
BinaryResponse,
|
||||
CostTrackingTestCase,
|
||||
ExactExpected,
|
||||
FailureExpected,
|
||||
RecountExpected,
|
||||
data_errors,
|
||||
)
|
||||
from pydantic import JsonValue
|
||||
|
||||
if _data_errors := data_errors():
|
||||
raise ValueError("\n".join(_data_errors))
|
||||
|
|
@ -96,6 +97,16 @@ def _assert_stream_has_no_error(response_text: str) -> None:
|
|||
), f"stream carried an error event: {parsed}"
|
||||
|
||||
|
||||
def _replace_model(value: JsonValue, model_name: str) -> JsonValue:
|
||||
if isinstance(value, str):
|
||||
return value.replace("$MODEL", model_name)
|
||||
if isinstance(value, list):
|
||||
return [_replace_model(item, model_name) for item in value]
|
||||
if isinstance(value, dict):
|
||||
return {key: _replace_model(item, model_name) for key, item in value.items()}
|
||||
return value
|
||||
|
||||
|
||||
def _assert_breakdown(
|
||||
case: CostTrackingTestCase,
|
||||
expected: ExactExpected,
|
||||
|
|
@ -139,12 +150,12 @@ def _assert_breakdown(
|
|||
assert actual_component is not None and approx_equal(actual_component, expected_component), (
|
||||
f"{case.name}: {field} {actual_component} != expected {expected_component}"
|
||||
)
|
||||
if case.response.content_type == "application/json":
|
||||
if expected.cost_header and case.response.content_type == "application/json":
|
||||
header: Final = response.headers.get(header_name)
|
||||
assert header is not None and approx_equal(float(header), expected_component), (
|
||||
f"{case.name}: {header_name} {header} != expected {expected_component}"
|
||||
)
|
||||
if case.response.content_type == "application/json" and any(
|
||||
if expected.cost_header and case.response.content_type == "application/json" and any(
|
||||
component is not None
|
||||
for component in (
|
||||
expected.cache_read_cost,
|
||||
|
|
@ -171,11 +182,51 @@ def test_case_bills_expected_cost(gateway: Gateway, case: CostTrackingTestCase)
|
|||
marker: Final = sha256(case.name.encode()).hexdigest()[:12]
|
||||
with gateway.scenario() as scenario:
|
||||
key: Final = scenario.key()
|
||||
model_name: Final = register_scenario_deployment(scenario, case, marker, key)
|
||||
passthrough_provider: Final = case.passthrough_provider
|
||||
scenario_id: Final = f"sc-{marker}-{sha256(key.encode()).hexdigest()[:12]}"
|
||||
scenario_handle: Final = (
|
||||
register_scenario(scenario_id, case.response)
|
||||
if passthrough_provider in {"gemini", "anthropic"}
|
||||
else None
|
||||
)
|
||||
if scenario_handle is not None:
|
||||
scenario.cleanups.callback(delete_scenario, scenario_handle)
|
||||
model_name: Final = (
|
||||
case.model
|
||||
if passthrough_provider in {"gemini", "anthropic"}
|
||||
else register_scenario_deployment(scenario, case, marker, key)
|
||||
)
|
||||
request_model: Final = (
|
||||
case.model.rsplit("/", 1)[-1]
|
||||
if passthrough_provider in {"gemini", "anthropic"}
|
||||
else model_name
|
||||
)
|
||||
request_body: Final = JSON_OBJECT.validate_python(
|
||||
_replace_model(case.request, request_model)
|
||||
if passthrough_provider is not None
|
||||
else {**case.request, "model": model_name}
|
||||
)
|
||||
request_headers: Final = (
|
||||
{
|
||||
"x-pass-x-scripted-scenario": scenario_id,
|
||||
**(
|
||||
{"x-goog-api-key": key}
|
||||
if passthrough_provider == "gemini"
|
||||
else {}
|
||||
),
|
||||
}
|
||||
if passthrough_provider is not None
|
||||
else {}
|
||||
)
|
||||
request_path: Final = (
|
||||
case.endpoint.replace("$MODEL", request_model)
|
||||
if passthrough_provider is not None
|
||||
else case.endpoint
|
||||
)
|
||||
response: Final = (
|
||||
_multipart_request(gateway, case, model_name, key)
|
||||
if case.upload is not None
|
||||
else gateway.request("POST", case.endpoint, {**case.request, "model": model_name}, key=key)
|
||||
else gateway.request("POST", request_path, request_body, key=key, headers=request_headers)
|
||||
)
|
||||
if isinstance(case.expected, FailureExpected):
|
||||
assert response.status_code == case.expected.failure.status, (
|
||||
|
|
@ -220,13 +271,14 @@ def test_case_bills_expected_cost(gateway: Gateway, case: CostTrackingTestCase)
|
|||
)
|
||||
elif case.response.content_type == "application/json":
|
||||
header: Final = cast(str | None, response.headers.get("x-litellm-response-cost"))
|
||||
assert (
|
||||
(header is None or approx_equal(float(header), 0.0))
|
||||
if expected.spend == 0
|
||||
else (header is not None and approx_equal(float(header), expected.spend))
|
||||
), (
|
||||
f"{case.name}: x-litellm-response-cost {header} != expected {expected.spend}"
|
||||
)
|
||||
if expected.cost_header and expected.spend != 0:
|
||||
assert header is not None and approx_equal(float(header), expected.spend), (
|
||||
f"{case.name}: x-litellm-response-cost {header} != expected {expected.spend}"
|
||||
)
|
||||
elif header is not None:
|
||||
assert approx_equal(float(header), expected.spend), (
|
||||
f"{case.name}: x-litellm-response-cost {header} != expected {expected.spend}"
|
||||
)
|
||||
assert row.spend is not None and approx_equal(row.spend, expected.spend), (
|
||||
f"{case.name}: spend {row.spend} != expected {expected.spend} "
|
||||
f"(breakdown {row.breakdown.model_dump() if row.breakdown is not None else None})"
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue