mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-11 22:51:28 +00:00
Merge remote-tracking branch 'origin/litellm_internal_staging' into litellm_lit_7015_migration_job_node_selector
This commit is contained in:
mateo-berri
commit
3b9568bcdd
13 changed files with 543 additions and 52 deletions
|
|
@ -19,7 +19,7 @@ import httpx
|
|||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm._uuid import uuid
|
||||
from litellm.constants import REDACTED_BY_LITELLM
|
||||
from litellm.constants import REDACTED_BY_LITELLM, REDACTED_BY_LITELM_STRING
|
||||
from litellm.integrations.custom_batch_logger import CustomBatchLogger
|
||||
from litellm.integrations.datadog.datadog_handler import (
|
||||
get_datadog_base_url_from_env,
|
||||
|
|
@ -46,9 +46,10 @@ from litellm.llms.custom_httpx.http_handler import (
|
|||
from litellm.proxy.spend_tracking.savings import extract_cache_creation_tokens, extract_cache_read_tokens
|
||||
from litellm.types.integrations.datadog_llm_obs import *
|
||||
from litellm.types.utils import (
|
||||
AUDIT_GUARDRAIL_FIELDS,
|
||||
PROMPT_CARRYING_GUARDRAIL_FIELDS,
|
||||
PROMPT_QUOTING_ROUTING_DECISION_FIELDS,
|
||||
CallTypes,
|
||||
StandardLoggingGuardrailInformation,
|
||||
StandardLoggingPayload,
|
||||
StandardLoggingPayloadErrorInformation,
|
||||
)
|
||||
|
|
@ -60,6 +61,8 @@ _SAFE_REDACTED_MESSAGE_ROLES: Final = frozenset(
|
|||
{"agent", "assistant", "developer", "function", "model", "system", "tool", "user"}
|
||||
)
|
||||
|
||||
_CLASSIFIED_GUARDRAIL_FIELDS: Final = AUDIT_GUARDRAIL_FIELDS | PROMPT_CARRYING_GUARDRAIL_FIELDS
|
||||
|
||||
_PROMPT_CARRYING_METADATA_FIELDS: Final = frozenset(
|
||||
{
|
||||
"routing_decision",
|
||||
|
|
@ -108,6 +111,49 @@ def _router_span_fields(
|
|||
)
|
||||
|
||||
|
||||
def _guardrail_entries(guardrail_information: object) -> tuple[Mapping[str, object], ...]:
|
||||
"""The guardrail records as a sequence, whatever shape the payload carries.
|
||||
|
||||
`guardrail_information` is typed as a list, but a guardrail that writes the metadata key itself
|
||||
can leave a single record there; Prometheus normalizes the same shape at
|
||||
`_guardrail_overhead_seconds`.
|
||||
"""
|
||||
if isinstance(guardrail_information, Mapping):
|
||||
return (guardrail_information,)
|
||||
if isinstance(guardrail_information, (list, tuple)):
|
||||
return tuple(entry for entry in guardrail_information if isinstance(entry, Mapping))
|
||||
return ()
|
||||
|
||||
|
||||
def _guardrail_entry_without_prompt_carriers(entry: Mapping[str, object]) -> Mapping[str, object]:
|
||||
"""One guardrail record kept as its audit fields, with the prompt-quoting ones marked redacted.
|
||||
|
||||
Built as an allow-list rather than a deny-list: a key neither set classifies is dropped, so a
|
||||
guardrail that records its own extra detail cannot put the caller's prompt on a redacted span.
|
||||
"""
|
||||
return { # mutable-ok: a fresh record built per entry, handed straight to the span serializer
|
||||
field: REDACTED_BY_LITELM_STRING if field in PROMPT_CARRYING_GUARDRAIL_FIELDS else value
|
||||
for field, value in entry.items()
|
||||
if field in _CLASSIFIED_GUARDRAIL_FIELDS
|
||||
}
|
||||
|
||||
|
||||
def _guardrail_information_without_prompt_carriers(
|
||||
guardrail_information: object,
|
||||
) -> tuple[Mapping[str, object], ...] | None:
|
||||
"""The guardrail records reduced to what a redacted span may carry.
|
||||
|
||||
Redaction removes the prompt, not the record that a guardrail ran: the name, mode, status,
|
||||
timings and masked-entity counts are what an operator reads to answer whether a guardrail
|
||||
caught anything on a request, and none of them reproduce the prompt. Field-level rather than
|
||||
dropping the list, which is what `_sanitize_guardrail_information_for_spend_logs` already does
|
||||
for spend logs.
|
||||
"""
|
||||
if guardrail_information is None:
|
||||
return None
|
||||
return tuple(_guardrail_entry_without_prompt_carriers(entry) for entry in _guardrail_entries(guardrail_information))
|
||||
|
||||
|
||||
def _metadata_without_prompt_carriers(standard_logging_metadata: Mapping[str, Any]) -> Mapping[str, Any]:
|
||||
"""The metadata minus the records that quote prompts, tool arguments, tool results, or retrieved text."""
|
||||
return MappingProxyType(
|
||||
|
|
@ -872,7 +918,9 @@ class DataDogLLMObsLogger(CustomBatchLogger):
|
|||
"cache_key": standard_logging_payload.get("cache_key", "unknown"),
|
||||
"saved_cache_cost": standard_logging_payload.get("saved_cache_cost", 0),
|
||||
"guardrail_information": (
|
||||
None if redact_prompt_text else standard_logging_payload.get("guardrail_information", None)
|
||||
_guardrail_information_without_prompt_carriers(standard_logging_payload.get("guardrail_information"))
|
||||
if redact_prompt_text
|
||||
else standard_logging_payload.get("guardrail_information", None)
|
||||
),
|
||||
"is_streamed_request": self._get_stream_value_from_payload(standard_logging_payload),
|
||||
"latency_metrics": dict(self._get_latency_metrics(standard_logging_payload)),
|
||||
|
|
@ -904,14 +952,12 @@ class DataDogLLMObsLogger(CustomBatchLogger):
|
|||
latency_metrics["litellm_overhead_time_ms"] = litellm_overhead_ms
|
||||
|
||||
# Guardrail overhead latency
|
||||
guardrail_info: Final[list[StandardLoggingGuardrailInformation] | None] = standard_logging_payload.get(
|
||||
"guardrail_information"
|
||||
)
|
||||
if guardrail_info is not None:
|
||||
guardrail_info: Final = _guardrail_entries(standard_logging_payload.get("guardrail_information"))
|
||||
if guardrail_info:
|
||||
total_duration = 0.0
|
||||
for info in guardrail_info:
|
||||
_guardrail_duration_seconds: float | None = info.get("duration")
|
||||
if _guardrail_duration_seconds is not None:
|
||||
_guardrail_duration_seconds = info.get("duration")
|
||||
if isinstance(_guardrail_duration_seconds, (int, float, str)):
|
||||
total_duration += float(_guardrail_duration_seconds)
|
||||
|
||||
if total_duration > 0:
|
||||
|
|
|
|||
|
|
@ -2,11 +2,11 @@ import json
|
|||
import re
|
||||
from collections.abc import Collection, Mapping
|
||||
from types import MappingProxyType, UnionType
|
||||
from typing import Any, Final, Union, get_args, get_origin
|
||||
from typing import Annotated, Any, Final, Union, get_args, get_origin
|
||||
|
||||
import orjson
|
||||
from fastapi import Request, UploadFile, status
|
||||
from typing_extensions import ReadOnly
|
||||
from typing_extensions import NotRequired, ReadOnly, Required
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.constants import MAX_REQUEST_BODY_SIZE_TO_REPAIR_MB
|
||||
|
|
@ -18,6 +18,8 @@ from litellm.types.router import Deployment
|
|||
|
||||
_FORM_CONTENT_TYPES: Final[frozenset[str]] = frozenset({"application/x-www-form-urlencoded", "multipart/form-data"})
|
||||
|
||||
_ANNOTATION_QUALIFIERS: Final[frozenset[object]] = frozenset({Annotated, NotRequired, ReadOnly, Required})
|
||||
|
||||
|
||||
def _normalize_media_type(content_type: str) -> str:
|
||||
"""Return the bare media type per RFC 7231: strip params, trim, lowercase."""
|
||||
|
|
@ -42,9 +44,17 @@ def _is_json_content_type(content_type: str) -> bool:
|
|||
return _normalize_media_type(content_type) == "application/json"
|
||||
|
||||
|
||||
def _unqualified(annotation: object) -> object:
|
||||
"""Which qualifiers ``get_type_hints`` already stripped varies by interpreter version, so peel them all."""
|
||||
if get_origin(annotation) not in _ANNOTATION_QUALIFIERS:
|
||||
return annotation
|
||||
qualified: Final[tuple[object, ...]] = get_args(annotation)
|
||||
return _unqualified(qualified[0])
|
||||
|
||||
|
||||
def _numeric_form_type(annotation: object) -> type[int] | type[float] | None:
|
||||
"""The scalar to parse an ``int``/``float``-typed field as, else ``None``."""
|
||||
unwrapped: Final = get_args(annotation)[0] if get_origin(annotation) is ReadOnly else annotation
|
||||
unwrapped: Final = _unqualified(annotation)
|
||||
candidates: Final = (
|
||||
tuple(arg for arg in get_args(unwrapped) if arg is not type(None))
|
||||
if get_origin(unwrapped) in (Union, UnionType)
|
||||
|
|
|
|||
|
|
@ -38,6 +38,7 @@ from litellm.proxy.common_utils.timezone_utils import (
|
|||
get_budget_reset_settings,
|
||||
)
|
||||
from litellm.proxy.common_utils.user_api_key_cache import (
|
||||
end_user_cache_key,
|
||||
model_access_group_cache_key,
|
||||
model_access_group_spend_counter_key,
|
||||
tag_cache_key,
|
||||
|
|
@ -177,6 +178,21 @@ def _model_access_group_cache_keys(row: _ModelAccessGroupRow) -> tuple[str, ...]
|
|||
return (model_access_group_cache_key(row.access_group_name),)
|
||||
|
||||
|
||||
def _enduser_counter_key(row: _EndUserRow) -> str:
|
||||
return f"spend:end_user:{row.user_id}"
|
||||
|
||||
|
||||
def _enduser_cache_keys(row: _EndUserRow) -> tuple[str, ...]:
|
||||
return (end_user_cache_key(row.user_id),)
|
||||
|
||||
|
||||
def _enduser_carried_spend(row: _EndUserRow, caps: Mapping[str, float]) -> float:
|
||||
if not caps:
|
||||
return 0.0
|
||||
effective_budget_id: Final[str | None] = row.budget_id or litellm.max_end_user_budget_id
|
||||
return _carried_spend(row.spend, caps.get(effective_budget_id) if effective_budget_id is not None else None)
|
||||
|
||||
|
||||
def _budget_link_where(
|
||||
budget_ids: Sequence[str],
|
||||
extra: Mapping[str, object] = MappingProxyType({}),
|
||||
|
|
@ -650,6 +666,7 @@ class ResetBudgetJob:
|
|||
if _rollover_enabled()
|
||||
else {} # mutable-ok: empty sentinel immediately frozen by MappingProxyType
|
||||
)
|
||||
endusers: Final[tuple[_EndUserRow, ...]] = await self._collect_endusers_to_reset(budget_ids)
|
||||
return _BudgetCascade(
|
||||
budgets=tuple(budgets_to_reset),
|
||||
budget_ids=budget_ids,
|
||||
|
|
@ -661,7 +678,7 @@ class ResetBudgetJob:
|
|||
for b in budgets_to_reset
|
||||
if b.budget_id is not None and b.budget_duration is not None
|
||||
),
|
||||
endusers=await self._collect_endusers_to_reset(budget_ids),
|
||||
endusers=endusers,
|
||||
counter_resets=(
|
||||
*(
|
||||
(_team_membership_counter_key(row), _row_carried_spend(row, rollover_caps))
|
||||
|
|
@ -674,6 +691,7 @@ class ResetBudgetJob:
|
|||
(_model_access_group_counter_key(row), _row_carried_spend(row, rollover_caps))
|
||||
for row in model_access_groups
|
||||
),
|
||||
*((_enduser_counter_key(row), _enduser_carried_spend(row, rollover_caps)) for row in endusers),
|
||||
),
|
||||
rollover_caps=rollover_caps,
|
||||
cache_keys=(
|
||||
|
|
@ -682,6 +700,7 @@ class ResetBudgetJob:
|
|||
*(key for row in orgs for key in _org_cache_keys(row)),
|
||||
*(key for row in tags for key in _tag_cache_keys(row)),
|
||||
*(key for row in model_access_groups for key in _model_access_group_cache_keys(row)),
|
||||
*(key for row in endusers for key in _enduser_cache_keys(row)),
|
||||
),
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -26,6 +26,7 @@ from litellm.proxy._types import Litellm_EntityType
|
|||
from litellm.repositories.organization_repository import OrganizationRepository
|
||||
from litellm.repositories.table_repositories import (
|
||||
BudgetWindowSpendRepository,
|
||||
EndUserRepository,
|
||||
SpendLogsRepository,
|
||||
TeamMembershipRepository,
|
||||
)
|
||||
|
|
@ -36,6 +37,8 @@ from litellm.repositories.verification_token_repository import (
|
|||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from prisma.types import LiteLLM_EndUserTableWhereUniqueInput
|
||||
|
||||
from litellm.caching.dual_cache import DualCache
|
||||
from litellm.proxy.utils import PrismaClient
|
||||
|
||||
|
|
@ -47,6 +50,8 @@ _WINDOW_SPEND_ENTITY_TYPES: Final[Mapping[str, str]] = MappingProxyType(
|
|||
}
|
||||
)
|
||||
|
||||
END_USER_COUNTER_PREFIX: Final = "spend:end_user:"
|
||||
|
||||
_WINDOW_SPEND_LOG_FIELDS: Final[Mapping[str, str]] = MappingProxyType(
|
||||
{
|
||||
"Key": "api_key",
|
||||
|
|
@ -74,6 +79,10 @@ class SpendCounterReseed:
|
|||
End-user and tag spend counters intentionally do not reseed here. Their
|
||||
auth paths already load the corresponding objects via get_end_user_object()
|
||||
and get_tag_objects_batch(); callers pass those values as fallback_spend.
|
||||
end_user_from_db is the one end-user read, used only as the budget floor when
|
||||
a counter sits below that cached spend: a worker that did not run the budget
|
||||
reset still caches the pre-reset end-user object, and LiteLLM_EndUserTable
|
||||
is the row the reset zeroed.
|
||||
"""
|
||||
|
||||
_locks: ClassVar["OrderedDict[str, asyncio.Lock]"] = OrderedDict()
|
||||
|
|
@ -129,7 +138,7 @@ class SpendCounterReseed:
|
|||
elif counter_key.startswith("spend:user:"):
|
||||
user_id = counter_key[len("spend:user:") :]
|
||||
row = await UserRepository(prisma_client).table.find_unique(where={"user_id": user_id})
|
||||
elif counter_key.startswith("spend:end_user:") or counter_key.startswith("spend:tag:"):
|
||||
elif counter_key.startswith(END_USER_COUNTER_PREFIX) or counter_key.startswith("spend:tag:"):
|
||||
return None
|
||||
elif counter_key.startswith("spend:org:"):
|
||||
org_id: Final = counter_key[len("spend:org:") :]
|
||||
|
|
@ -143,6 +152,20 @@ class SpendCounterReseed:
|
|||
return None
|
||||
return float(getattr(row, "spend", 0.0) or 0.0)
|
||||
|
||||
@staticmethod
|
||||
async def end_user_from_db(prisma_client: Optional["PrismaClient"], counter_key: str) -> float | None:
|
||||
if prisma_client is None or not counter_key.startswith(END_USER_COUNTER_PREFIX):
|
||||
return None
|
||||
where: Final[LiteLLM_EndUserTableWhereUniqueInput] = {"user_id": counter_key[len(END_USER_COUNTER_PREFIX) :]}
|
||||
try:
|
||||
row: Final = await EndUserRepository(prisma_client).table.find_unique(where=where)
|
||||
except Exception: # noqa: BLE001 # a failed floor read falls back to the cached spend, like from_db
|
||||
verbose_proxy_logger.exception("SpendCounterReseed.end_user_from_db: failed for %s", counter_key)
|
||||
return None
|
||||
if row is None:
|
||||
return None
|
||||
return float(row.spend or 0.0)
|
||||
|
||||
@staticmethod
|
||||
def _is_key_or_team_window_counter(counter_key: str) -> bool:
|
||||
for prefix in ("spend:key:", "spend:team:"):
|
||||
|
|
|
|||
|
|
@ -423,7 +423,7 @@ from litellm.proxy.db.proxy_worker_heartbeat import (
|
|||
PROXY_WORKER_HEARTBEAT_INTERVAL_SECONDS,
|
||||
ProxyWorkerHeartbeat,
|
||||
)
|
||||
from litellm.proxy.db.spend_counter_reseed import SpendCounterReseed
|
||||
from litellm.proxy.db.spend_counter_reseed import END_USER_COUNTER_PREFIX, SpendCounterReseed
|
||||
from litellm.proxy.discovery_endpoints import ui_discovery_endpoints_router
|
||||
from litellm.proxy.fine_tuning_endpoints.endpoints import router as fine_tuning_router
|
||||
from litellm.proxy.fine_tuning_endpoints.endpoints import set_fine_tuning_config
|
||||
|
|
@ -2477,7 +2477,8 @@ async def get_current_spend(
|
|||
authoritative source depends on the counter: primary key/team/user/org
|
||||
counters read the DB row; per-window counters (``window_start`` supplied)
|
||||
read the maintained window-spend row and only aggregate spend logs when
|
||||
that row is missing or stale; end-user/tag counters have no DB row, so the caller's
|
||||
that row is missing or stale; end-user counters read ``LiteLLM_EndUserTable``, the
|
||||
row the budget reset zeroes; tag counters have no DB row, so the caller's
|
||||
``fallback_spend`` (loaded fresh in auth) is authoritative. The DB read is
|
||||
skipped for healthy primary counters (counter at or above recorded spend)
|
||||
and cached in-process for a few seconds, so a persistently stale counter
|
||||
|
|
@ -2511,8 +2512,8 @@ async def get_current_spend(
|
|||
await _repair_stale_spend_counter(counter_key=counter_key, db_spend=authoritative)
|
||||
return authoritative
|
||||
elif fallback_spend > current:
|
||||
# end-user / tag counters have no DB row; fallback_spend is the
|
||||
# authoritative recorded value loaded in auth.
|
||||
# nothing to read (tag counters, an end user without a row or a DB client, a
|
||||
# failed read); fallback_spend is the authoritative recorded value loaded in auth.
|
||||
return fallback_spend
|
||||
|
||||
# Opt-in hard guarantee: when the spend backing this admit decision came
|
||||
|
|
@ -2580,6 +2581,29 @@ async def reseed_spend_counter_from_db(counter_key: str) -> None:
|
|||
await _repair_stale_spend_counter(counter_key=counter_key, db_spend=db_spend)
|
||||
|
||||
|
||||
async def _floor_spend_from_db(
|
||||
counter_key: str,
|
||||
window_entity_type: str | None,
|
||||
window_entity_id: str | None,
|
||||
window_duration: str | None,
|
||||
window_start: datetime | None,
|
||||
) -> float | None:
|
||||
if counter_key.startswith(END_USER_COUNTER_PREFIX):
|
||||
return await SpendCounterReseed.end_user_from_db(prisma_client=prisma_client, counter_key=counter_key)
|
||||
entity_spend: Final = await SpendCounterReseed.from_db(prisma_client=prisma_client, counter_key=counter_key)
|
||||
if entity_spend is not None:
|
||||
return entity_spend
|
||||
if window_entity_type is None or window_entity_id is None or window_start is None:
|
||||
return None
|
||||
return await SpendCounterReseed.window_from_db(
|
||||
prisma_client=prisma_client,
|
||||
entity_type=window_entity_type,
|
||||
entity_id=window_entity_id,
|
||||
window_duration=window_duration,
|
||||
window_start=window_start,
|
||||
)
|
||||
|
||||
|
||||
async def _authoritative_floor_spend(
|
||||
counter_key: str,
|
||||
window_entity_type: str | None = None,
|
||||
|
|
@ -2592,20 +2616,13 @@ async def _authoritative_floor_spend(
|
|||
if cached is not None:
|
||||
return float(cached)
|
||||
|
||||
db_spend = await SpendCounterReseed.from_db(prisma_client=prisma_client, counter_key=counter_key)
|
||||
if (
|
||||
db_spend is None
|
||||
and window_entity_type is not None
|
||||
and window_entity_id is not None
|
||||
and window_start is not None
|
||||
):
|
||||
db_spend = await SpendCounterReseed.window_from_db(
|
||||
prisma_client=prisma_client,
|
||||
entity_type=window_entity_type,
|
||||
entity_id=window_entity_id,
|
||||
window_duration=window_duration,
|
||||
window_start=window_start,
|
||||
)
|
||||
db_spend: Final = await _floor_spend_from_db(
|
||||
counter_key=counter_key,
|
||||
window_entity_type=window_entity_type,
|
||||
window_entity_id=window_entity_id,
|
||||
window_duration=window_duration,
|
||||
window_start=window_start,
|
||||
)
|
||||
if db_spend is None:
|
||||
return None
|
||||
|
||||
|
|
|
|||
|
|
@ -37,6 +37,7 @@ from litellm.proxy._types import SpendLogsMetadata, SpendLogsPayload, SpendLogsR
|
|||
from litellm.proxy.spend_tracking.spend_log_error_logger import spend_log_error
|
||||
from litellm.proxy.utils import PrismaClient, hash_token
|
||||
from litellm.types.utils import (
|
||||
PROMPT_CARRYING_GUARDRAIL_FIELDS,
|
||||
CallTypes,
|
||||
CostBreakdown,
|
||||
StandardLoggingGuardrailInformation,
|
||||
|
|
@ -1073,13 +1074,6 @@ def _sanitize_guardrail_information_for_spend_logs(
|
|||
return [_redact_prompt_fields_in_guardrail_entry(entry) for entry in entries if isinstance(entry, dict)]
|
||||
|
||||
|
||||
_PROMPT_CARRYING_GUARDRAIL_FIELDS: Final = (
|
||||
"guardrail_request",
|
||||
"guardrail_response",
|
||||
"match_details",
|
||||
"classification",
|
||||
)
|
||||
|
||||
_NUMERIC_COMPRESSION_STAT_KEYS: Final = (
|
||||
"tokens_before",
|
||||
"tokens_after",
|
||||
|
|
@ -1114,7 +1108,7 @@ def _redact_prompt_fields_in_guardrail_entry(
|
|||
preserved_stats: Final = _numeric_compression_stats_from_guardrail_response(entry.get("guardrail_response"))
|
||||
redacted: Final[StandardLoggingGuardrailInformation] = {
|
||||
**entry,
|
||||
**{key: REDACTED_BY_LITELM_STRING for key in _PROMPT_CARRYING_GUARDRAIL_FIELDS if key in entry},
|
||||
**{key: REDACTED_BY_LITELM_STRING for key in PROMPT_CARRYING_GUARDRAIL_FIELDS if key in entry},
|
||||
}
|
||||
if preserved_stats is None:
|
||||
return redacted
|
||||
|
|
|
|||
|
|
@ -3080,6 +3080,50 @@ class GuardrailMode(TypedDict, total=False):
|
|||
|
||||
GuardrailStatus = Literal["success", "guardrail_intervened", "guardrail_failed_to_respond", "not_run"]
|
||||
|
||||
# Fields on a guardrail record whose values can quote the caller's prompt: the payload sent to the
|
||||
# guardrail, the provider response that echoes it back, and the two first-party hooks that inline
|
||||
# prompt substrings (``block_code_execution`` and ``litellm_content_filter``). Every other field
|
||||
# reports what the guardrail decided without reproducing the prompt, so redaction replaces these
|
||||
# four and keeps the rest of the record.
|
||||
PROMPT_CARRYING_GUARDRAIL_FIELDS: Final[frozenset[str]] = frozenset(
|
||||
{
|
||||
"guardrail_request",
|
||||
"guardrail_response",
|
||||
"match_details",
|
||||
"classification",
|
||||
}
|
||||
)
|
||||
|
||||
# The rest of the record: what the guardrail is, what it decided, how long it took and what it cost.
|
||||
# None of these reproduce the prompt, so a redacted record keeps them and stays explainable.
|
||||
# `test_every_guardrail_field_is_classified` fails if a field is added to the record without being
|
||||
# placed in one set or the other, so a new field is dropped from redacted records rather than
|
||||
# shipped unexamined.
|
||||
AUDIT_GUARDRAIL_FIELDS: Final[frozenset[str]] = frozenset(
|
||||
{
|
||||
"guardrail_name",
|
||||
"guardrail_provider",
|
||||
"guardrail_mode",
|
||||
"guardrail_status",
|
||||
"start_time",
|
||||
"end_time",
|
||||
"duration",
|
||||
"masked_entity_count",
|
||||
"guardrail_id",
|
||||
"policy_template",
|
||||
"detection_method",
|
||||
"confidence_score",
|
||||
"patterns_checked",
|
||||
"alert_recipients",
|
||||
"risk_score",
|
||||
"violation_categories",
|
||||
"guardrail_action",
|
||||
"guardrail_usage",
|
||||
"guardrail_cost",
|
||||
"guardrail_cost_in_spend",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
class StandardLoggingGuardrailInformation(TypedDict, total=False):
|
||||
guardrail_name: str | None
|
||||
|
|
|
|||
|
|
@ -66,6 +66,7 @@ IGNORE_FUNCTIONS = [
|
|||
"_json_safe", # max depth set (_MAX_DEPTH) plus a seen-ids cycle guard for self-referential input.
|
||||
"_redact_agent_params_tree", # max depth set (default 10), same shape as _redact_sensitive_litellm_params.
|
||||
"_restore_redacted_nested_value", # max depth set (default 10), mirrors _redact_agent_params_tree on the write side.
|
||||
"_unqualified", # bounded by the qualifier depth of a static TypedDict annotation (Annotated, Required/NotRequired, ReadOnly around one type, no cycles possible).
|
||||
]
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -12,7 +12,7 @@ spelling of prompt-cache counts (`prompt_tokens_details.cached_tokens`).
|
|||
import json
|
||||
import os
|
||||
from datetime import datetime, timedelta
|
||||
from typing import Any
|
||||
from typing import Any, Final
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
|
|
@ -20,6 +20,8 @@ import pytest
|
|||
import litellm
|
||||
from litellm.integrations.datadog.datadog_llm_obs import DataDogLLMObsLogger
|
||||
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
|
||||
from litellm.types.guardrails import GuardrailEventHooks
|
||||
from litellm.types.utils import StandardLoggingGuardrailInformation
|
||||
|
||||
TOOL_DEFINITION: dict[str, Any] = {
|
||||
"type": "function",
|
||||
|
|
@ -631,10 +633,133 @@ def test_redaction_drops_every_prompt_carrying_metadata_record(logger: DataDogLL
|
|||
for record in sensitive_metadata:
|
||||
assert record not in redacted["meta"]["metadata"]
|
||||
assert record in unredacted["meta"]["metadata"]
|
||||
assert redacted["meta"]["metadata"]["guardrail_information"] is None
|
||||
assert redacted["meta"]["metadata"]["guardrail_information"] == [
|
||||
{"guardrail_name": "g", "guardrail_request": "REDACTED_BY_LITELM"}
|
||||
] # the record survives; only the field quoting the prompt is replaced
|
||||
assert unredacted["meta"]["metadata"]["guardrail_information"] is not None
|
||||
|
||||
|
||||
_AUDIT_RECORD: Final[StandardLoggingGuardrailInformation] = StandardLoggingGuardrailInformation(
|
||||
guardrail_name="bedrock-pii",
|
||||
guardrail_provider="bedrock",
|
||||
guardrail_mode=GuardrailEventHooks.pre_call,
|
||||
guardrail_status="guardrail_intervened",
|
||||
guardrail_response={"action": "MASK", "match": "alice@acme.com"},
|
||||
match_details=[{"pattern": "email", "match": "alice@acme.com"}],
|
||||
classification="the user asked for alice@acme.com",
|
||||
masked_entity_count={"EMAIL": 2},
|
||||
violation_categories=["pii"],
|
||||
duration=0.01,
|
||||
)
|
||||
|
||||
|
||||
def _payload_with_guardrail_record(guardrail_information: object) -> dict[str, Any]:
|
||||
payload = build_payload()
|
||||
payload["standard_logging_object"]["guardrail_information"] = guardrail_information
|
||||
return payload
|
||||
|
||||
|
||||
def test_redaction_keeps_the_guardrail_audit_record(logger: DataDogLLMObsLogger) -> None:
|
||||
"""Redaction removes the prompt, not the operator's record that a guardrail intervened."""
|
||||
redacted = _span_json(
|
||||
_redacting_logger(turn_off_message_logging=True),
|
||||
_payload_with_guardrail_record([dict(_AUDIT_RECORD)]),
|
||||
)
|
||||
record = redacted["meta"]["metadata"]["guardrail_information"][0]
|
||||
|
||||
for field in ("guardrail_request", "guardrail_response", "match_details", "classification"):
|
||||
assert record.get(field, "REDACTED_BY_LITELM") == "REDACTED_BY_LITELM"
|
||||
assert record["guardrail_name"] == "bedrock-pii"
|
||||
assert record["guardrail_provider"] == "bedrock"
|
||||
assert record["guardrail_mode"] == "pre_call"
|
||||
assert record["guardrail_status"] == "guardrail_intervened"
|
||||
assert record["masked_entity_count"] == {"EMAIL": 2}
|
||||
assert record["violation_categories"] == ["pii"]
|
||||
assert record["duration"] == 0.01
|
||||
assert "alice@acme.com" not in safe_dumps(redacted["meta"]["metadata"])
|
||||
|
||||
|
||||
def test_a_caller_supplied_redaction_header_cannot_blank_the_guardrail_record(
|
||||
logger: DataDogLLMObsLogger,
|
||||
) -> None:
|
||||
"""Any key may redact its own prompts with the header; none may erase what a guardrail caught."""
|
||||
payload = _payload_with_guardrail_record([dict(_AUDIT_RECORD)])
|
||||
payload["litellm_params"] = {"metadata": {"headers": {"x-litellm-enable-message-redaction": "true"}}}
|
||||
|
||||
span = _span_json(logger, payload)
|
||||
record = span["meta"]["metadata"]["guardrail_information"][0]
|
||||
|
||||
assert span["meta"]["input"]["messages"] == [{"role": "user", "content": "redacted-by-litellm"}]
|
||||
assert record["guardrail_status"] == "guardrail_intervened"
|
||||
assert record["masked_entity_count"] == {"EMAIL": 2}
|
||||
assert "alice@acme.com" not in safe_dumps(span["meta"]["metadata"])
|
||||
|
||||
|
||||
def test_a_guardrails_own_extra_field_never_reaches_a_redacted_span(logger: DataDogLLMObsLogger) -> None:
|
||||
"""A guardrail may record whatever it likes; only classified fields survive redaction."""
|
||||
span = _span_json(
|
||||
_redacting_logger(turn_off_message_logging=True),
|
||||
_payload_with_guardrail_record([{**_AUDIT_RECORD, "matched_text": "the caller asked about alice@acme.com"}]),
|
||||
)
|
||||
record = span["meta"]["metadata"]["guardrail_information"][0]
|
||||
|
||||
assert "matched_text" not in record
|
||||
assert record["guardrail_status"] == "guardrail_intervened"
|
||||
assert "alice@acme.com" not in safe_dumps(span["meta"]["metadata"])
|
||||
|
||||
|
||||
def test_a_lone_guardrail_record_survives_redaction(logger: DataDogLLMObsLogger) -> None:
|
||||
"""A guardrail that writes the metadata key itself leaves one record, not a list of them."""
|
||||
span = _span_json(
|
||||
_redacting_logger(turn_off_message_logging=True),
|
||||
_payload_with_guardrail_record(dict(_AUDIT_RECORD)),
|
||||
)
|
||||
metadata = span["meta"]["metadata"]
|
||||
|
||||
assert metadata["guardrail_information"] == [
|
||||
{
|
||||
"guardrail_name": "bedrock-pii",
|
||||
"guardrail_provider": "bedrock",
|
||||
"guardrail_mode": "pre_call",
|
||||
"guardrail_status": "guardrail_intervened",
|
||||
"guardrail_response": "REDACTED_BY_LITELM",
|
||||
"match_details": "REDACTED_BY_LITELM",
|
||||
"classification": "REDACTED_BY_LITELM",
|
||||
"masked_entity_count": {"EMAIL": 2},
|
||||
"violation_categories": ["pii"],
|
||||
"duration": 0.01,
|
||||
}
|
||||
]
|
||||
assert metadata["latency_metrics"]["guardrail_overhead_time_ms"] == 10.0
|
||||
|
||||
|
||||
@pytest.mark.parametrize("guardrail_information", [None, [], 5, "abc", [None, "x"], {}])
|
||||
def test_odd_guardrail_shapes_still_produce_a_span(
|
||||
guardrail_information: object,
|
||||
) -> None:
|
||||
"""The redacted branch replaced an expression that could not fail, so it must not start failing."""
|
||||
span = _span_json(
|
||||
_redacting_logger(turn_off_message_logging=True),
|
||||
_payload_with_guardrail_record(guardrail_information),
|
||||
)
|
||||
|
||||
assert span["meta"]["input"]["messages"] == [{"role": "user", "content": "redacted-by-litellm"}]
|
||||
assert span["meta"]["metadata"]["guardrail_information"] in (None, [], [{}])
|
||||
|
||||
|
||||
def test_a_redacted_span_carries_every_declared_guardrail_field() -> None:
|
||||
"""A field added to the record without a redaction decision would be dropped, so it fails here."""
|
||||
declared = dict.fromkeys(StandardLoggingGuardrailInformation.__annotations__, "alice@acme.com")
|
||||
payload = _payload_with_guardrail_record([{**declared, "duration": 0.01}])
|
||||
|
||||
span = _span_json(_redacting_logger(turn_off_message_logging=True), payload)
|
||||
record = span["meta"]["metadata"]["guardrail_information"][0]
|
||||
|
||||
assert set(record) == set(declared)
|
||||
for field in ("guardrail_request", "guardrail_response", "match_details", "classification"):
|
||||
assert record[field] == "REDACTED_BY_LITELM"
|
||||
|
||||
|
||||
def test_tool_definitions_accept_the_bare_anthropic_shape(logger: DataDogLLMObsLogger) -> None:
|
||||
"""The Anthropic surface declares tools unwrapped, with input_schema instead of parameters."""
|
||||
payload = build(
|
||||
|
|
|
|||
|
|
@ -1053,6 +1053,8 @@ class TestNumericFormFields:
|
|||
read_only: ReadOnly[int | None]
|
||||
not_required: NotRequired[ReadOnly[int]]
|
||||
required: Required[ReadOnly[Annotated[float, "meta"]]]
|
||||
read_only_not_required: ReadOnly[NotRequired[int]]
|
||||
read_only_required: ReadOnly[Required[float]]
|
||||
|
||||
assert dict(numeric_form_fields(get_type_hints(Schema))) == {
|
||||
"plain": int,
|
||||
|
|
@ -1061,6 +1063,22 @@ class TestNumericFormFields:
|
|||
"read_only": int,
|
||||
"not_required": int,
|
||||
"required": float,
|
||||
"read_only_not_required": int,
|
||||
"read_only_required": float,
|
||||
}
|
||||
|
||||
def test_qualifiers_are_unwrapped_when_get_type_hints_keeps_extras(self):
|
||||
from typing_extensions import Annotated, NotRequired, ReadOnly, Required, TypedDict
|
||||
|
||||
class Schema(TypedDict, total=False):
|
||||
annotated: ReadOnly[Annotated[int, "meta"]]
|
||||
not_required: NotRequired[ReadOnly[int]]
|
||||
required: Required[ReadOnly[Annotated[float, "meta"]]]
|
||||
|
||||
assert dict(numeric_form_fields(get_type_hints(Schema, include_extras=True))) == {
|
||||
"annotated": int,
|
||||
"not_required": int,
|
||||
"required": float,
|
||||
}
|
||||
|
||||
def test_non_scalar_and_bool_fields_are_skipped(self):
|
||||
|
|
|
|||
|
|
@ -4,7 +4,7 @@ import sys
|
|||
import types
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from datetime import time as dt_time
|
||||
from typing import Any, Dict, List
|
||||
from typing import Any, Dict, Final, List
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import httpx
|
||||
|
|
@ -1495,6 +1495,32 @@ def test_budget_table_reset_invalidates_every_tag_not_just_the_first(reset_budge
|
|||
assert deleted == {"tag:tenant-a", "tag:tenant-b", "tag:tenant-c"}
|
||||
|
||||
|
||||
def test_budget_table_reset_invalidates_enduser_counter_and_cache(reset_budget_job, mock_prisma_client, monkeypatch):
|
||||
"""When an end user's budget resets, its Redis spend counter is zeroed and its management cache is evicted."""
|
||||
counter_cache: Final = _make_counter_invalidation_job(monkeypatch)
|
||||
budget: Final = _budget_row(budget_id="budget-1")
|
||||
mock_prisma_client.data["budget"] = [budget]
|
||||
test_enduser: Final = type(
|
||||
"LiteLLM_EndUserTable",
|
||||
(),
|
||||
{
|
||||
"spend": 20.0,
|
||||
"litellm_budget_table": budget,
|
||||
"budget_id": "budget-1",
|
||||
"user_id": "customer-42",
|
||||
},
|
||||
)
|
||||
mock_prisma_client.data["enduser"] = [test_enduser]
|
||||
|
||||
asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table())
|
||||
|
||||
counter_cache.in_memory_cache.set_cache.assert_any_call(key="spend:end_user:customer-42", value=0.0, ttl=60)
|
||||
counter_cache.redis_cache.async_set_cache.assert_any_await(key="spend:end_user:customer-42", value=0.0, ttl=60)
|
||||
deleted: Final = {call.kwargs.get("key") for call in counter_cache.user_api_key_cache.async_delete_cache.await_args_list}
|
||||
assert "end_user_id:customer-42" in deleted
|
||||
|
||||
|
||||
|
||||
def test_budget_table_reset_commits_even_when_cache_eviction_fails(reset_budget_job, mock_prisma_client, monkeypatch):
|
||||
"""Eviction runs after the commit, so a broken cache cannot undo the write."""
|
||||
counter_cache = _make_counter_invalidation_job(monkeypatch)
|
||||
|
|
@ -3028,6 +3054,38 @@ def test_budget_cascade_carries_enduser_overage_when_rollover_enabled(
|
|||
} in enduser_writes
|
||||
|
||||
|
||||
def test_budget_cascade_carries_default_tier_enduser_counter_when_rollover_enabled(
|
||||
rollover_enabled, reset_budget_job, mock_prisma_client, monkeypatch
|
||||
):
|
||||
"""An end user on the default budget (no budget_id on its row) 5 over the cap
|
||||
keeps a counter of 5 in the next window and loses its cached object."""
|
||||
import litellm
|
||||
|
||||
counter_cache: Final = _make_counter_invalidation_job(monkeypatch)
|
||||
monkeypatch.setattr(litellm, "max_end_user_budget_id", "default-enduser-budget")
|
||||
mock_prisma_client.data["budget"] = [
|
||||
_budget_row(budget_id="default-enduser-budget", budget_duration="1d", max_budget=10.0)
|
||||
]
|
||||
implicit_enduser: Final = type(
|
||||
"EndUserRow",
|
||||
(),
|
||||
{
|
||||
"spend": 15.0,
|
||||
"user_id": "enduser-implicit",
|
||||
"budget_id": None,
|
||||
"model_dump": lambda self=None: {"spend": 15.0, "user_id": "enduser-implicit", "budget_id": None, "blocked": False},
|
||||
},
|
||||
)
|
||||
mock_prisma_client.db.litellm_endusertable.set_find_many_results([implicit_enduser])
|
||||
|
||||
asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table())
|
||||
|
||||
counter_cache.in_memory_cache.set_cache.assert_any_call(key="spend:end_user:enduser-implicit", value=5.0, ttl=60)
|
||||
counter_cache.redis_cache.async_set_cache.assert_any_await(key="spend:end_user:enduser-implicit", value=5.0, ttl=60)
|
||||
deleted: Final = {call.kwargs.get("key") for call in counter_cache.user_api_key_cache.async_delete_cache.await_args_list}
|
||||
assert "end_user_id:enduser-implicit" in deleted
|
||||
|
||||
|
||||
def _replay_spend_writes(writes, spend):
|
||||
"""Apply the queued update_many statements in order, the way the DB
|
||||
transaction executes them, and return the row's final spend."""
|
||||
|
|
|
|||
|
|
@ -9,6 +9,7 @@ from __future__ import annotations
|
|||
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from types import SimpleNamespace
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
|
||||
|
|
@ -18,7 +19,7 @@ from litellm.proxy.db.spend_counter_reseed import SpendCounterReseed
|
|||
WINDOW_START = datetime(2026, 8, 1, tzinfo=timezone.utc)
|
||||
|
||||
|
||||
class _FakeWindowSpendTable:
|
||||
class _FakeFindUniqueTable:
|
||||
def __init__(self, row: SimpleNamespace | None, error: Exception | None = None) -> None:
|
||||
self._row = row
|
||||
self._error = error
|
||||
|
|
@ -47,10 +48,13 @@ class _FakePrismaClient:
|
|||
row: SimpleNamespace | None = None,
|
||||
spend_logs_total: float = 0.0,
|
||||
error: Exception | None = None,
|
||||
end_user_row: SimpleNamespace | None = None,
|
||||
end_user_error: Exception | None = None,
|
||||
) -> None:
|
||||
self.db = SimpleNamespace(
|
||||
litellm_budgetwindowspend=_FakeWindowSpendTable(row=row, error=error),
|
||||
litellm_budgetwindowspend=_FakeFindUniqueTable(row=row, error=error),
|
||||
litellm_spendlogs=_FakeSpendLogsTable(total=spend_logs_total),
|
||||
litellm_endusertable=_FakeFindUniqueTable(row=end_user_row, error=end_user_error),
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -248,3 +252,65 @@ async def test_coalesced_window_seeds_a_cold_counter_from_the_row():
|
|||
assert result == 4.5
|
||||
assert cache.in_memory_cache.get_cache(key=counter_key) == 4.5
|
||||
assert prisma.db.litellm_spendlogs.call_count == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_end_user_from_db_reads_the_end_user_row_by_user_id():
|
||||
prisma: Final = _FakePrismaClient(end_user_row=SimpleNamespace(user_id="customer-42", spend=0.0))
|
||||
|
||||
result: Final = await SpendCounterReseed.end_user_from_db(
|
||||
prisma_client=prisma, counter_key="spend:end_user:customer-42"
|
||||
)
|
||||
|
||||
assert result == 0.0
|
||||
assert prisma.db.litellm_endusertable.where_clauses == [{"user_id": "customer-42"}]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_end_user_from_db_returns_the_recorded_spend():
|
||||
prisma: Final = _FakePrismaClient(end_user_row=SimpleNamespace(user_id="customer-42", spend=12.5))
|
||||
|
||||
assert (
|
||||
await SpendCounterReseed.end_user_from_db(prisma_client=prisma, counter_key="spend:end_user:customer-42")
|
||||
== 12.5
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("counter_key", ["spend:key:hashed", "spend:team:t1", "spend:tag:t1"])
|
||||
async def test_end_user_from_db_ignores_other_counter_kinds_without_touching_the_db(counter_key):
|
||||
prisma: Final = _FakePrismaClient(end_user_row=SimpleNamespace(user_id="x", spend=5.0))
|
||||
|
||||
assert await SpendCounterReseed.end_user_from_db(prisma_client=prisma, counter_key=counter_key) is None
|
||||
assert prisma.db.litellm_endusertable.where_clauses == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_end_user_from_db_returns_none_without_a_row_a_client_or_on_db_error():
|
||||
assert (
|
||||
await SpendCounterReseed.end_user_from_db(prisma_client=None, counter_key="spend:end_user:customer-42")
|
||||
is None
|
||||
)
|
||||
assert (
|
||||
await SpendCounterReseed.end_user_from_db(
|
||||
prisma_client=_FakePrismaClient(end_user_row=None), counter_key="spend:end_user:customer-42"
|
||||
)
|
||||
is None
|
||||
)
|
||||
assert (
|
||||
await SpendCounterReseed.end_user_from_db(
|
||||
prisma_client=_FakePrismaClient(end_user_error=RuntimeError("db down")),
|
||||
counter_key="spend:end_user:customer-42",
|
||||
)
|
||||
is None
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_from_db_still_never_reads_the_end_user_row():
|
||||
"""A cold end-user counter keeps seeding from the cached end-user object the auth
|
||||
path already loaded; the row is read only as the budget floor."""
|
||||
prisma: Final = _FakePrismaClient(end_user_row=SimpleNamespace(user_id="customer-42", spend=5.0))
|
||||
|
||||
assert await SpendCounterReseed.from_db(prisma_client=prisma, counter_key="spend:end_user:customer-42") is None
|
||||
assert prisma.db.litellm_endusertable.where_clauses == []
|
||||
|
|
|
|||
|
|
@ -22,6 +22,7 @@ from __future__ import annotations
|
|||
|
||||
import asyncio
|
||||
from datetime import datetime
|
||||
from typing import Final
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
|
|
@ -222,16 +223,19 @@ async def test_get_current_spend_floor_caches_db_read(monkeypatch):
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_current_spend_floors_end_user_tag_against_fallback(monkeypatch):
|
||||
"""End-user and tag counters have no DB row (from_db returns None). When the
|
||||
counter is stale-low, enforcement falls back to the caller's recorded spend
|
||||
(loaded fresh in auth) instead of trusting the stale counter."""
|
||||
fake_cache = _make_spend_counter_cache(redis_get_value=2.0)
|
||||
@pytest.mark.parametrize("counter_key", ("spend:end_user:e1", "spend:tag:t1"))
|
||||
async def test_get_current_spend_floors_end_user_tag_against_fallback(monkeypatch, counter_key):
|
||||
"""Tag counters have no DB row (from_db returns None), and an end-user counter has
|
||||
none to read without a DB client. When such a counter is stale-low, enforcement
|
||||
falls back to the caller's recorded spend (loaded fresh in auth) instead of
|
||||
trusting the stale counter."""
|
||||
fake_cache: Final = _make_spend_counter_cache(redis_get_value=2.0)
|
||||
monkeypatch.setattr(ps, "spend_counter_cache", fake_cache)
|
||||
monkeypatch.setattr(ps, "prisma_client", None)
|
||||
monkeypatch.setattr(ps.SpendCounterReseed, "from_db", AsyncMock(return_value=None))
|
||||
|
||||
result = await ps.get_current_spend(
|
||||
counter_key="spend:end_user:e1",
|
||||
result: Final = await ps.get_current_spend(
|
||||
counter_key=counter_key,
|
||||
fallback_spend=20.0,
|
||||
max_budget=10.0,
|
||||
)
|
||||
|
|
@ -241,6 +245,72 @@ async def test_get_current_spend_floors_end_user_tag_against_fallback(monkeypatc
|
|||
fake_cache.redis_cache.async_set_max.assert_not_called()
|
||||
|
||||
|
||||
def _make_prisma_with_end_user_row(spend: float | None):
|
||||
prisma: Final = MagicMock()
|
||||
prisma.db.litellm_endusertable.find_unique = AsyncMock(
|
||||
return_value=None if spend is None else MagicMock(spend=spend)
|
||||
)
|
||||
return prisma
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_current_spend_end_user_floor_admits_after_a_reset_on_a_stale_worker(monkeypatch):
|
||||
"""The reset job zeroes LiteLLM_EndUserTable.spend and the shared counter, but it
|
||||
evicts the cached end-user object only on the worker that ran the reset. Every
|
||||
other worker still passes the pre-reset spend as fallback_spend, and that stale
|
||||
copy must not out-vote the reset row."""
|
||||
fake_cache: Final = _make_spend_counter_cache(redis_get_value=0.0)
|
||||
monkeypatch.setattr(ps, "spend_counter_cache", fake_cache)
|
||||
prisma: Final = _make_prisma_with_end_user_row(spend=0.0)
|
||||
monkeypatch.setattr(ps, "prisma_client", prisma)
|
||||
|
||||
result = await ps.get_current_spend(
|
||||
counter_key="spend:end_user:customer-42",
|
||||
fallback_spend=0.000032,
|
||||
max_budget=0.00003,
|
||||
fallback_authoritative=True,
|
||||
)
|
||||
|
||||
assert result == 0.0
|
||||
prisma.db.litellm_endusertable.find_unique.assert_awaited_once_with(where={"user_id": "customer-42"})
|
||||
fake_cache.redis_cache.async_set_max.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_current_spend_end_user_floor_repairs_a_stale_low_counter(monkeypatch):
|
||||
"""After a Redis restart the end-user counter can sit below the recorded spend;
|
||||
the row wins and the shared counter is raised so other workers stop admitting on
|
||||
the stale value."""
|
||||
fake_cache: Final = _make_spend_counter_cache(redis_get_value=2.0)
|
||||
monkeypatch.setattr(ps, "spend_counter_cache", fake_cache)
|
||||
monkeypatch.setattr(ps, "prisma_client", _make_prisma_with_end_user_row(spend=12.0))
|
||||
|
||||
result: Final = await ps.get_current_spend(
|
||||
counter_key="spend:end_user:customer-42",
|
||||
fallback_spend=12.0,
|
||||
max_budget=10.0,
|
||||
)
|
||||
|
||||
assert result == 12.0
|
||||
fake_cache.redis_cache.async_set_max.assert_awaited_once_with(key="spend:end_user:customer-42", value=12.0)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_current_spend_end_user_without_a_row_keeps_the_cached_spend(monkeypatch):
|
||||
fake_cache: Final = _make_spend_counter_cache(redis_get_value=0.0)
|
||||
monkeypatch.setattr(ps, "spend_counter_cache", fake_cache)
|
||||
monkeypatch.setattr(ps, "prisma_client", _make_prisma_with_end_user_row(spend=None))
|
||||
|
||||
result: Final = await ps.get_current_spend(
|
||||
counter_key="spend:end_user:customer-42",
|
||||
fallback_spend=20.0,
|
||||
max_budget=10.0,
|
||||
)
|
||||
|
||||
assert result == 20.0
|
||||
fake_cache.redis_cache.async_set_max.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_current_spend_floors_window_against_spend_logs(monkeypatch):
|
||||
"""Per-window counters have no DB row but aggregate from spend logs. A
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue