scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-09 03:18:44 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 109

fix(e2e): share pricing request models

Browse source
This commit is contained in:
Yuneng Jiang 2026-09-17 09:45:21 -07:00
parent 90775f4f76
commit 3e4f1ea439
No known key found for this signature in database
2 changed files with 62 additions and 54 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

42
tests/e2e/models.py
View file

@ -905,6 +905,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

74
tests/e2e/quota_management/spend_tracking/test_pricing_config_spend_e2e.py
View file

@ -3,7 +3,7 @@ from __future__ import annotations
import os
import time
from collections.abc import Iterator
from typing import Final, Literal
from typing import Final
import pytest
from cost_rows import CostRow, approx_equal, poll_cost_row
@ -11,15 +11,23 @@ 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 pydantic import BaseModel, RootModel
from spend_e2e_client import SpendClient
pytestmark = pytest.mark.e2e
@ -31,48 +39,6 @@ MARGIN_PERCENT: Final = 0.1
MARGIN_FIXED: Final = 0.0005
class _ConfigPatchResponse(BaseModel):
status: str
values: dict[str, float | dict[str, float]]
class _DiscountConfig(RootModel[dict[str, float]]):
pass
class _MarginConfig(RootModel[dict[str, float | dict[str, float]]]):
pass
class _DiscountConfigResponse(BaseModel):
values: dict[str, float]
class _MarginConfigResponse(BaseModel):
values: dict[str, float | dict[str, float]]
class _BedrockParams(BaseModel):
guardrail: Literal["bedrock"] = "bedrock"
mode: Literal["pre_call"] = "pre_call"
default_on: bool = False
guardrailIdentifier: str
guardrailVersion: str
class _GuardrailSpec(BaseModel):
guardrail_name: str
litellm_params: _BedrockParams
class _GuardrailCreate(BaseModel):
guardrail: _GuardrailSpec
class _GuardrailCreateResponse(BaseModel):
guardrail_id: str
def _register_bedrock_guardrail(
client: SpendClient,
resources: ResourceManager,
@ -84,16 +50,16 @@ def _register_bedrock_guardrail(
client.proxy.transport.post(
"/guardrails",
headers=client.proxy.transport.master,
json=_GuardrailCreate(
guardrail=_GuardrailSpec(
json=GuardrailCreateBody(
guardrail=BedrockGuardrailSpec(
guardrail_name=name,
litellm_params=_BedrockParams(
litellm_params=BedrockGuardrailParams(
guardrailIdentifier=identifier,
guardrailVersion=version,
),
)
),
response_type=_GuardrailCreateResponse,
response_type=GuardrailCreateResponse,
)
).guardrail_id
settle_propagation(time.monotonic())
@ -129,8 +95,8 @@ def _set_discount(client: SpendClient, values: dict[str, float]) -> float:
client.proxy.transport.patch(
"/config/cost_discount_config",
headers=client.proxy.transport.master,
json=_DiscountConfig(values),
response_type=_ConfigPatchResponse,
json=CostDiscountConfig(values),
response_type=ConfigPatchResponse,
)
)
return time.monotonic()
@ -141,8 +107,8 @@ def _set_margin(client: SpendClient, values: dict[str, float | dict[str, float]]
client.proxy.transport.patch(
"/config/cost_margin_config",
headers=client.proxy.transport.master,
json=_MarginConfig(values),
response_type=_ConfigPatchResponse,
json=CostMarginConfig(values),
response_type=ConfigPatchResponse,
)
)
return time.monotonic()
@ -154,7 +120,7 @@ def _get_discount(client: SpendClient) -> dict[str, float]:
"/config/cost_discount_config",
headers=client.proxy.transport.master,
params=NoBody(),
response_type=_DiscountConfigResponse,
response_type=CostDiscountConfigResponse,
)
).values
@ -165,7 +131,7 @@ def _get_margin(client: SpendClient) -> dict[str, float | dict[str, float]]:
"/config/cost_margin_config",
headers=client.proxy.transport.master,
params=NoBody(),
response_type=_MarginConfigResponse,
response_type=CostMarginConfigResponse,
)
).values