mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-09 03:18:44 +00:00
Merge 3e4f1ea439 into 4b67a2b845
This commit is contained in:
yuneng-jiang
GitHub
commit
0af63b0c81
4 changed files with 358 additions and 0 deletions
|
|
@ -44,6 +44,9 @@
|
|||
- {id: quota_management.spend_tracking.messages_bridge.logs_cost, module: quota_management, tier: P1, behavior: spend_tracking, variant: messages_bridge, assertions: [logs_cost], exercised_on: [messages], source: "llms/anthropic/pass_through/responses_adapters/handler.py", rationale: "A streaming /v1/messages request served by an openai-provider model is bridged through the anthropic-messages -> Responses adapter and must aggregate the consumed SSE stream into one spend row with nonzero cost and token counts, attributed to custom_llm_provider openai under call_type anthropic_messages"}
|
||||
- {id: quota_management.spend_tracking.embeddings.logs_cost, module: quota_management, tier: P1, behavior: spend_tracking, variant: embeddings, assertions: [logs_cost], exercised_on: [embeddings], source: "proxy/spend_tracking/spend_tracking_utils.py", rationale: "Embedding calls write nonzero spend rows"}
|
||||
- {id: quota_management.spend_tracking.cache_hit.zero_cost, module: quota_management, tier: P1, behavior: spend_tracking, variant: cache_hit, assertions: [zero_cost], exercised_on: [chat_completions], source: "proxy/spend_tracking/spend_tracking_utils.py", rationale: "A response-cache hit logs at zero cost with the cache-hit marker"}
|
||||
- {id: quota_management.spend_tracking.discount_config.logs_cost, module: quota_management, tier: P1, behavior: spend_tracking, variant: discount_config, assertions: [logs_cost], exercised_on: [chat_completions], source: "proxy/management_endpoints/cost_tracking_settings.py + cost_calculator.py", rationale: "A persisted provider discount changes the real request's spend row"}
|
||||
- {id: quota_management.spend_tracking.margin_config.logs_cost, module: quota_management, tier: P1, behavior: spend_tracking, variant: margin_config, assertions: [logs_cost], exercised_on: [chat_completions], source: "proxy/management_endpoints/cost_tracking_settings.py + cost_calculator.py", rationale: "A persisted provider margin is added to the real request's spend row"}
|
||||
- {id: quota_management.spend_tracking.guardrail_cost.logs_cost, module: quota_management, tier: P1, behavior: spend_tracking, variant: guardrail_cost, assertions: [logs_cost], exercised_on: [chat_completions], source: "litellm_core_utils/llm_cost_calc/guardrail_cost.py + proxy/spend_tracking/spend_tracking_utils.py", rationale: "A priced Bedrock guardrail charge reaches the persisted spend row"}
|
||||
- {id: quota_management.spend_tracking.key_rollup.matches_sum_of_logs, module: quota_management, tier: P1, behavior: spend_tracking, variant: key_rollup, assertions: [matches_sum_of_logs], exercised_on: [chat_completions], source: "proxy/db/db_spend_update_writer.py", rationale: "A key's rolled-up spend equals the sum of its log rows"}
|
||||
- {id: quota_management.spend_tracking.concurrent_burst.loses_no_spend, module: quota_management, tier: P1, behavior: spend_tracking, variant: concurrent_burst, assertions: [loses_no_spend], exercised_on: [chat_completions], source: "proxy/db/db_spend_update_writer.py", rationale: "Concurrent calls all land as spend; no row lost to write contention"}
|
||||
- {id: quota_management.spend_tracking.surface_consistency.matches_every_surface, module: quota_management, tier: P1, behavior: spend_tracking, variant: surface_consistency, assertions: [matches_every_surface], exercised_on: [chat_completions], source: "proxy/db/db_spend_update_writer.py", rationale: "One priced request lands the same response_cost on the spend log row, /key/info, /team/info, the usage export's /user/daily/activity/aggregated row, and the litellm_spend_metric Prometheus sample; each is a separate writer, so a rounding, dropped, or double-counted write on one drifts it from the rest (LIT-3620, LIT-5045)"}
|
||||
|
|
|
|||
|
|
@ -1204,6 +1204,48 @@ class CostMap(RootModel[dict[str, CostMapEntry]]):
|
|||
pass
|
||||
|
||||
|
||||
class ConfigPatchResponse(BaseModel):
|
||||
status: str
|
||||
values: dict[str, float | dict[str, float]]
|
||||
|
||||
|
||||
class CostDiscountConfig(RootModel[dict[str, float]]):
|
||||
pass
|
||||
|
||||
|
||||
class CostMarginConfig(RootModel[dict[str, float | dict[str, float]]]):
|
||||
pass
|
||||
|
||||
|
||||
class CostDiscountConfigResponse(BaseModel):
|
||||
values: dict[str, float]
|
||||
|
||||
|
||||
class CostMarginConfigResponse(BaseModel):
|
||||
values: dict[str, float | dict[str, float]]
|
||||
|
||||
|
||||
class BedrockGuardrailParams(BaseModel):
|
||||
guardrail: Literal["bedrock"] = "bedrock"
|
||||
mode: Literal["pre_call"] = "pre_call"
|
||||
default_on: bool = False
|
||||
guardrailIdentifier: str
|
||||
guardrailVersion: str
|
||||
|
||||
|
||||
class BedrockGuardrailSpec(BaseModel):
|
||||
guardrail_name: str
|
||||
litellm_params: BedrockGuardrailParams
|
||||
|
||||
|
||||
class GuardrailCreateBody(BaseModel):
|
||||
guardrail: BedrockGuardrailSpec
|
||||
|
||||
|
||||
class GuardrailCreateResponse(BaseModel):
|
||||
guardrail_id: str
|
||||
|
||||
|
||||
class FileEntry(BaseModel):
|
||||
id: str
|
||||
|
||||
|
|
|
|||
|
|
@ -49,6 +49,13 @@ class CostBreakdownRow(BaseModel):
|
|||
reasoning_cost: float | None = None
|
||||
tool_usage_cost: float | None = None
|
||||
total_cost: float | None = None
|
||||
original_cost: float | None = None
|
||||
discount_percent: float | None = None
|
||||
discount_amount: float | None = None
|
||||
margin_percent: float | None = None
|
||||
margin_fixed_amount: float | None = None
|
||||
margin_total_amount: float | None = None
|
||||
guardrail_cost: float | None = None
|
||||
service_tier: str | None = None
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -0,0 +1,306 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import time
|
||||
from collections.abc import Iterator
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
from cost_rows import CostRow, approx_equal, poll_cost_row
|
||||
from e2e_config import settle_propagation, unique_marker
|
||||
from e2e_http import NoBody, StreamingResponse, unwrap
|
||||
from lifecycle import ResourceManager
|
||||
from models import (
|
||||
BedrockGuardrailParams,
|
||||
BedrockGuardrailSpec,
|
||||
ChatBody,
|
||||
ChatMessage,
|
||||
ChatResponse,
|
||||
ConfigPatchResponse,
|
||||
CostDiscountConfig,
|
||||
CostDiscountConfigResponse,
|
||||
CostMarginConfig,
|
||||
CostMarginConfigResponse,
|
||||
GuardrailCreateBody,
|
||||
GuardrailCreateResponse,
|
||||
KeyDeleteBody,
|
||||
KeyGenerateBody,
|
||||
LiteLLMParamsBody,
|
||||
ModelDeleteBody,
|
||||
)
|
||||
from spend_e2e_client import SpendClient
|
||||
|
||||
pytestmark = pytest.mark.e2e
|
||||
|
||||
INPUT_RATE: Final = 0.00005
|
||||
OUTPUT_RATE: Final = 0.0001
|
||||
DISCOUNT: Final = 0.25
|
||||
MARGIN_PERCENT: Final = 0.1
|
||||
MARGIN_FIXED: Final = 0.0005
|
||||
|
||||
|
||||
def _register_bedrock_guardrail(
|
||||
client: SpendClient,
|
||||
resources: ResourceManager,
|
||||
name: str,
|
||||
identifier: str,
|
||||
version: str,
|
||||
) -> None:
|
||||
guardrail_id: Final = unwrap(
|
||||
client.proxy.transport.post(
|
||||
"/guardrails",
|
||||
headers=client.proxy.transport.master,
|
||||
json=GuardrailCreateBody(
|
||||
guardrail=BedrockGuardrailSpec(
|
||||
guardrail_name=name,
|
||||
litellm_params=BedrockGuardrailParams(
|
||||
guardrailIdentifier=identifier,
|
||||
guardrailVersion=version,
|
||||
),
|
||||
)
|
||||
),
|
||||
response_type=GuardrailCreateResponse,
|
||||
)
|
||||
).guardrail_id
|
||||
settle_propagation(time.monotonic())
|
||||
resources.defer(lambda: _delete_guardrail(client, guardrail_id))
|
||||
|
||||
|
||||
def _delete_guardrail(client: SpendClient, guardrail_id: str) -> None:
|
||||
unwrap(
|
||||
client.proxy.transport.delete(
|
||||
f"/guardrails/{guardrail_id}",
|
||||
headers=client.proxy.transport.master,
|
||||
json=NoBody(),
|
||||
response_type=NoBody,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def _guarded_chat(client: SpendClient, key: str, model: str, name: str) -> StreamingResponse:
|
||||
return client.proxy.transport.send(
|
||||
"/chat/completions",
|
||||
headers=client.proxy.transport.bearer(key),
|
||||
json=ChatBody(
|
||||
model=model,
|
||||
messages=[ChatMessage(role="user", content=f"reply with a short greeting {unique_marker()}")],
|
||||
max_tokens=16,
|
||||
guardrails=[name],
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def _set_discount(client: SpendClient, values: dict[str, float]) -> float:
|
||||
unwrap(
|
||||
client.proxy.transport.patch(
|
||||
"/config/cost_discount_config",
|
||||
headers=client.proxy.transport.master,
|
||||
json=CostDiscountConfig(values),
|
||||
response_type=ConfigPatchResponse,
|
||||
)
|
||||
)
|
||||
return time.monotonic()
|
||||
|
||||
|
||||
def _set_margin(client: SpendClient, values: dict[str, float | dict[str, float]]) -> float:
|
||||
unwrap(
|
||||
client.proxy.transport.patch(
|
||||
"/config/cost_margin_config",
|
||||
headers=client.proxy.transport.master,
|
||||
json=CostMarginConfig(values),
|
||||
response_type=ConfigPatchResponse,
|
||||
)
|
||||
)
|
||||
return time.monotonic()
|
||||
|
||||
|
||||
def _get_discount(client: SpendClient) -> dict[str, float]:
|
||||
return unwrap(
|
||||
client.proxy.transport.get(
|
||||
"/config/cost_discount_config",
|
||||
headers=client.proxy.transport.master,
|
||||
params=NoBody(),
|
||||
response_type=CostDiscountConfigResponse,
|
||||
)
|
||||
).values
|
||||
|
||||
|
||||
def _get_margin(client: SpendClient) -> dict[str, float | dict[str, float]]:
|
||||
return unwrap(
|
||||
client.proxy.transport.get(
|
||||
"/config/cost_margin_config",
|
||||
headers=client.proxy.transport.master,
|
||||
params=NoBody(),
|
||||
response_type=CostMarginConfigResponse,
|
||||
)
|
||||
).values
|
||||
|
||||
|
||||
def _register_model(client: SpendClient, resources: ResourceManager, prefix: str) -> str:
|
||||
model: Final = f"{prefix}-{unique_marker()}"
|
||||
model_id: Final = client.proxy.create_model(
|
||||
model,
|
||||
LiteLLMParamsBody(
|
||||
model="openai/gpt-4o-mini",
|
||||
api_key="os.environ/OPENAI_API_KEY",
|
||||
input_cost_per_token=INPUT_RATE,
|
||||
output_cost_per_token=OUTPUT_RATE,
|
||||
),
|
||||
)
|
||||
resources.defer(lambda: _delete_model(client, model_id))
|
||||
return model
|
||||
|
||||
|
||||
def _delete_model(client: SpendClient, model_id: str) -> None:
|
||||
unwrap(
|
||||
client.proxy.transport.post(
|
||||
"/model/delete",
|
||||
headers=client.proxy.transport.master,
|
||||
json=ModelDeleteBody(id=model_id),
|
||||
response_type=NoBody,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def _delete_key(client: SpendClient, key: str) -> None:
|
||||
unwrap(
|
||||
client.proxy.transport.post(
|
||||
"/key/delete",
|
||||
headers=client.proxy.transport.master,
|
||||
json=KeyDeleteBody(keys=[key]),
|
||||
response_type=NoBody,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def _base_cost(row: CostRow, prompt_tokens: int, completion_tokens: int) -> float:
|
||||
assert prompt_tokens > 0 and completion_tokens >= 0
|
||||
assert row.prompt_tokens == prompt_tokens and row.completion_tokens == completion_tokens
|
||||
base_cost: Final = prompt_tokens * INPUT_RATE + completion_tokens * OUTPUT_RATE
|
||||
assert row.breakdown.original_cost is not None
|
||||
assert approx_equal(row.breakdown.original_cost, base_cost)
|
||||
return base_cost
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def strict_resources(client: SpendClient) -> Iterator[ResourceManager]:
|
||||
manager: Final = ResourceManager(client=client.proxy, strict_cleanup=True)
|
||||
manager.init()
|
||||
yield manager
|
||||
manager.teardown()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def scoped_key(client: SpendClient, strict_resources: ResourceManager) -> str:
|
||||
key: Final = client.proxy.generate_key(KeyGenerateBody(user_id="e2e-test-user"))
|
||||
strict_resources.defer(lambda: _delete_key(client, key))
|
||||
return key
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def restored_pricing_config(client: SpendClient) -> Iterator[None]:
|
||||
discount: Final = _get_discount(client)
|
||||
margin: Final = _get_margin(client)
|
||||
try:
|
||||
_set_discount(client, {})
|
||||
reset_at: Final = _set_margin(client, {})
|
||||
settle_propagation(reset_at)
|
||||
yield
|
||||
finally:
|
||||
try:
|
||||
_set_discount(client, discount)
|
||||
finally:
|
||||
restored_at: Final = _set_margin(client, margin)
|
||||
settle_propagation(restored_at)
|
||||
|
||||
|
||||
class TestPricingConfigSpend:
|
||||
@pytest.mark.covers(
|
||||
"quota_management.spend_tracking.discount_config.logs_cost",
|
||||
exercised_on=["chat_completions"],
|
||||
)
|
||||
def test_configured_discount_reaches_persisted_spend_row(
|
||||
self,
|
||||
client: SpendClient,
|
||||
strict_resources: ResourceManager,
|
||||
scoped_key: str,
|
||||
restored_pricing_config: None,
|
||||
) -> None:
|
||||
discount_set_at: Final = _set_discount(client, {"openai": DISCOUNT})
|
||||
settle_propagation(discount_set_at)
|
||||
model: Final = _register_model(client, strict_resources, "discount-priced")
|
||||
chat: Final = unwrap(client.chat(scoped_key, model, f"reply with one word {unique_marker()}", max_tokens=16))
|
||||
assert chat.id and chat.usage
|
||||
assert chat.usage.prompt_tokens is not None and chat.usage.prompt_tokens > 0
|
||||
assert chat.usage.completion_tokens is not None and chat.usage.completion_tokens >= 0
|
||||
|
||||
row: Final = poll_cost_row(client.proxy, chat.id)
|
||||
assert row is not None
|
||||
base_cost: Final = _base_cost(row, chat.usage.prompt_tokens, chat.usage.completion_tokens)
|
||||
breakdown: Final = row.breakdown
|
||||
assert breakdown.discount_percent is not None and approx_equal(breakdown.discount_percent, DISCOUNT)
|
||||
assert breakdown.discount_amount is not None and approx_equal(breakdown.discount_amount, base_cost * DISCOUNT)
|
||||
assert row.spend is not None and approx_equal(row.spend, base_cost * (1 - DISCOUNT))
|
||||
|
||||
@pytest.mark.covers(
|
||||
"quota_management.spend_tracking.margin_config.logs_cost",
|
||||
exercised_on=["chat_completions"],
|
||||
)
|
||||
def test_configured_margin_reaches_persisted_spend_row(
|
||||
self,
|
||||
client: SpendClient,
|
||||
strict_resources: ResourceManager,
|
||||
scoped_key: str,
|
||||
restored_pricing_config: None,
|
||||
) -> None:
|
||||
margin_set_at: Final = _set_margin(
|
||||
client, {"openai": {"percentage": MARGIN_PERCENT, "fixed_amount": MARGIN_FIXED}}
|
||||
)
|
||||
settle_propagation(margin_set_at)
|
||||
model: Final = _register_model(client, strict_resources, "margin-priced")
|
||||
chat: Final = unwrap(client.chat(scoped_key, model, f"reply with one word {unique_marker()}", max_tokens=16))
|
||||
assert chat.id and chat.usage
|
||||
assert chat.usage.prompt_tokens is not None and chat.usage.prompt_tokens > 0
|
||||
assert chat.usage.completion_tokens is not None and chat.usage.completion_tokens >= 0
|
||||
|
||||
row: Final = poll_cost_row(client.proxy, chat.id)
|
||||
assert row is not None
|
||||
base_cost: Final = _base_cost(row, chat.usage.prompt_tokens, chat.usage.completion_tokens)
|
||||
breakdown: Final = row.breakdown
|
||||
expected_margin: Final = base_cost * MARGIN_PERCENT + MARGIN_FIXED
|
||||
assert breakdown.margin_percent is not None and approx_equal(breakdown.margin_percent, MARGIN_PERCENT)
|
||||
assert breakdown.margin_fixed_amount is not None and approx_equal(breakdown.margin_fixed_amount, MARGIN_FIXED)
|
||||
assert breakdown.margin_total_amount is not None and approx_equal(breakdown.margin_total_amount, expected_margin)
|
||||
assert row.spend is not None and approx_equal(row.spend, base_cost + expected_margin)
|
||||
|
||||
@pytest.mark.covers(
|
||||
"quota_management.spend_tracking.guardrail_cost.logs_cost",
|
||||
exercised_on=["chat_completions"],
|
||||
)
|
||||
def test_bedrock_guardrail_cost_reaches_persisted_spend_row(
|
||||
self,
|
||||
client: SpendClient,
|
||||
strict_resources: ResourceManager,
|
||||
scoped_key: str,
|
||||
restored_pricing_config: None,
|
||||
) -> None:
|
||||
identifier: Final = os.environ["BEDROCK_GUARDRAIL_IDENTIFIER"]
|
||||
version: Final = os.environ["BEDROCK_GUARDRAIL_VERSION"]
|
||||
name: Final = f"e2e-bedrock-cost-{unique_marker()}"
|
||||
_register_bedrock_guardrail(client, strict_resources, name, identifier, version)
|
||||
model: Final = _register_model(client, strict_resources, "guardrail-priced")
|
||||
result: Final = _guarded_chat(client, scoped_key, model, name)
|
||||
assert result.ok, f"guarded request failed with {result.status_code}: {result.body[:400]}"
|
||||
assert name in {value.strip() for value in result.headers.get("x-litellm-applied-guardrails", "").split(",")}
|
||||
|
||||
chat: Final = ChatResponse.model_validate_json(result.body)
|
||||
assert chat.id and chat.usage
|
||||
assert chat.usage.prompt_tokens is not None and chat.usage.prompt_tokens > 0
|
||||
assert chat.usage.completion_tokens is not None and chat.usage.completion_tokens >= 0
|
||||
row: Final = poll_cost_row(client.proxy, chat.id)
|
||||
assert row is not None
|
||||
base_cost: Final = _base_cost(row, chat.usage.prompt_tokens, chat.usage.completion_tokens)
|
||||
breakdown: Final = row.breakdown
|
||||
assert breakdown.guardrail_cost is not None and breakdown.guardrail_cost > 0
|
||||
assert breakdown.total_cost is not None and approx_equal(breakdown.total_cost, base_cost + breakdown.guardrail_cost)
|
||||
assert row.spend is not None and approx_equal(row.spend, breakdown.total_cost)
|
||||
Loading…
Add table
Reference in a new issue