mirror of
https://github.com/BerriAI/litellm.git
synced 2026-08-28 05:25:59 +00:00
revert: remove unrelated changes from HttpOnly cookie fix branch
Reset files not related to the login cookie fix back to main: - prometheus.py, bedrock converse, guardrail handler - auth_checks.py, reset_budget_job.py, audit_logs.py - test_user_api_key_auth.py Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
Hendrik Jaks
parent
866405f443
commit
0684a1e275
7 changed files with 264 additions and 1125 deletions
|
|
@ -65,17 +65,6 @@ def _get_cached_end_user_id_for_cost_tracking():
|
|||
|
||||
class PrometheusLogger(CustomLogger):
|
||||
# Class variables or attributes
|
||||
|
||||
@staticmethod
|
||||
def get_instance() -> Optional["PrometheusLogger"]:
|
||||
"""Find the PrometheusLogger instance from litellm.callbacks, if registered."""
|
||||
import litellm
|
||||
|
||||
for cb in litellm.callbacks:
|
||||
if isinstance(cb, PrometheusLogger):
|
||||
return cb
|
||||
return None
|
||||
|
||||
def __init__( # noqa: PLR0915
|
||||
self,
|
||||
**kwargs,
|
||||
|
|
@ -191,31 +180,6 @@ class PrometheusLogger(CustomLogger):
|
|||
),
|
||||
)
|
||||
|
||||
# Remaining Budget for Org
|
||||
self.litellm_remaining_org_budget_metric = self._gauge_factory(
|
||||
"litellm_remaining_org_budget_metric",
|
||||
"Remaining budget for org",
|
||||
labelnames=self.get_labels_for_metric(
|
||||
"litellm_remaining_org_budget_metric"
|
||||
),
|
||||
)
|
||||
|
||||
# Max Budget for Org
|
||||
self.litellm_org_max_budget_metric = self._gauge_factory(
|
||||
"litellm_org_max_budget_metric",
|
||||
"Maximum budget set for org",
|
||||
labelnames=self.get_labels_for_metric("litellm_org_max_budget_metric"),
|
||||
)
|
||||
|
||||
# Org Budget Reset At
|
||||
self.litellm_org_budget_remaining_hours_metric = self._gauge_factory(
|
||||
"litellm_org_budget_remaining_hours_metric",
|
||||
"Remaining hours for org budget to be reset",
|
||||
labelnames=self.get_labels_for_metric(
|
||||
"litellm_org_budget_remaining_hours_metric"
|
||||
),
|
||||
)
|
||||
|
||||
# Remaining Budget for API Key
|
||||
self.litellm_remaining_api_key_budget_metric = self._gauge_factory(
|
||||
"litellm_remaining_api_key_budget_metric",
|
||||
|
|
@ -476,76 +440,6 @@ class PrometheusLogger(CustomLogger):
|
|||
labelnames=[],
|
||||
)
|
||||
|
||||
########################################
|
||||
# Managed Batch Metrics
|
||||
########################################
|
||||
self.litellm_managed_batch_created_total = self._counter_factory(
|
||||
name="litellm_managed_batch_created_total",
|
||||
documentation="Total number of managed batches created",
|
||||
labelnames=[
|
||||
"model",
|
||||
"api_provider",
|
||||
"user",
|
||||
"user_email",
|
||||
"api_key_alias",
|
||||
],
|
||||
)
|
||||
|
||||
self.litellm_managed_file_size_bytes = self._gauge_factory(
|
||||
"litellm_managed_file_size_bytes",
|
||||
"Size of the most recent managed batch file in bytes (last-seen value per label combination)",
|
||||
labelnames=["purpose", "file_type", "model", "api_provider", "user"],
|
||||
)
|
||||
|
||||
self.litellm_managed_batch_duration_seconds = self._histogram_factory(
|
||||
"litellm_managed_batch_duration_seconds",
|
||||
"Duration of completed managed batches in seconds (completed_at - created_at)",
|
||||
labelnames=["model", "api_provider"],
|
||||
buckets=BATCH_DURATION_BUCKETS,
|
||||
)
|
||||
|
||||
self.litellm_managed_file_created_total = self._counter_factory(
|
||||
name="litellm_managed_file_created_total",
|
||||
documentation="Total number of managed files created",
|
||||
labelnames=[
|
||||
"model",
|
||||
"api_provider",
|
||||
"user",
|
||||
"user_email",
|
||||
"api_key_alias",
|
||||
],
|
||||
)
|
||||
|
||||
self.litellm_managed_file_deleted_total = self._counter_factory(
|
||||
name="litellm_managed_file_deleted_total",
|
||||
documentation="Total number of managed file deletions (success or blocked)",
|
||||
labelnames=["result"],
|
||||
)
|
||||
|
||||
self.litellm_check_batch_cost_jobs_polled = self._gauge_factory(
|
||||
"litellm_check_batch_cost_jobs_polled",
|
||||
"Number of unprocessed batches found by the last CheckBatchCost poll",
|
||||
labelnames=[],
|
||||
)
|
||||
|
||||
self.litellm_check_batch_cost_jobs_processed_total = self._counter_factory(
|
||||
name="litellm_check_batch_cost_jobs_processed_total",
|
||||
documentation="Total number of batches successfully cost-tracked by CheckBatchCost",
|
||||
labelnames=["model", "api_provider"],
|
||||
)
|
||||
|
||||
self.litellm_check_batch_cost_errors_total = self._counter_factory(
|
||||
name="litellm_check_batch_cost_errors_total",
|
||||
documentation="Total number of errors in CheckBatchCost by error type",
|
||||
labelnames=["error_type"],
|
||||
)
|
||||
|
||||
self.litellm_check_batch_cost_last_run_timestamp = self._gauge_factory(
|
||||
"litellm_check_batch_cost_last_run_timestamp",
|
||||
"Unix timestamp of the last CheckBatchCost job run",
|
||||
labelnames=[],
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
print_verbose(f"Got exception on init prometheus client {str(e)}")
|
||||
raise e
|
||||
|
|
@ -1028,9 +922,6 @@ class PrometheusLogger(CustomLogger):
|
|||
user_api_team_alias = standard_logging_payload["metadata"][
|
||||
"user_api_key_team_alias"
|
||||
]
|
||||
user_api_key_org_id = standard_logging_payload["metadata"].get(
|
||||
"user_api_key_org_id"
|
||||
)
|
||||
output_tokens = standard_logging_payload["completion_tokens"]
|
||||
tokens_used = standard_logging_payload["total_tokens"]
|
||||
response_cost = standard_logging_payload["response_cost"]
|
||||
|
|
@ -1040,14 +931,10 @@ class PrometheusLogger(CustomLogger):
|
|||
user_api_key_auth_metadata: Optional[dict] = standard_logging_payload[
|
||||
"metadata"
|
||||
].get("user_api_key_auth_metadata")
|
||||
spend_logs_metadata: Optional[dict] = standard_logging_payload["metadata"].get(
|
||||
"spend_logs_metadata"
|
||||
)
|
||||
|
||||
combined_metadata: Dict[str, Any] = {
|
||||
**(_requester_metadata if _requester_metadata else {}),
|
||||
**(user_api_key_auth_metadata if user_api_key_auth_metadata else {}),
|
||||
**(spend_logs_metadata if spend_logs_metadata else {}),
|
||||
}
|
||||
if standard_logging_payload is not None and isinstance(
|
||||
standard_logging_payload, dict
|
||||
|
|
@ -1087,11 +974,9 @@ class PrometheusLogger(CustomLogger):
|
|||
),
|
||||
client_ip=standard_logging_payload["metadata"].get("requester_ip_address"),
|
||||
user_agent=standard_logging_payload["metadata"].get("user_agent"),
|
||||
stream=(
|
||||
str(standard_logging_payload.get("stream"))
|
||||
if litellm.prometheus_emit_stream_label
|
||||
else None
|
||||
),
|
||||
stream=str(standard_logging_payload.get("stream"))
|
||||
if litellm.prometheus_emit_stream_label
|
||||
else None,
|
||||
)
|
||||
|
||||
if (
|
||||
|
|
@ -1141,7 +1026,6 @@ class PrometheusLogger(CustomLogger):
|
|||
litellm_params=litellm_params,
|
||||
response_cost=response_cost,
|
||||
user_id=user_id,
|
||||
user_api_key_org_id=user_api_key_org_id,
|
||||
)
|
||||
|
||||
# set proxy virtual key rpm/tpm metrics
|
||||
|
|
@ -1297,7 +1181,6 @@ class PrometheusLogger(CustomLogger):
|
|||
litellm_params: dict,
|
||||
response_cost: float,
|
||||
user_id: Optional[str] = None,
|
||||
user_api_key_org_id: Optional[str] = None,
|
||||
):
|
||||
_metadata = litellm_params.get("metadata") or {}
|
||||
_team_spend = _metadata.get("user_api_key_team_spend", None)
|
||||
|
|
@ -1330,16 +1213,12 @@ class PrometheusLogger(CustomLogger):
|
|||
user_max_budget=_user_max_budget,
|
||||
response_cost=response_cost,
|
||||
),
|
||||
self._set_org_budget_metrics_after_api_request(
|
||||
org_id=user_api_key_org_id,
|
||||
response_cost=response_cost,
|
||||
),
|
||||
return_exceptions=True,
|
||||
)
|
||||
for i, r in enumerate(results):
|
||||
if isinstance(r, Exception):
|
||||
verbose_logger.debug(
|
||||
f"[Non-Blocking] Prometheus: Budget metric lookup {['key', 'team', 'user', 'org'][i]} failed: {r}"
|
||||
f"[Non-Blocking] Prometheus: Budget metric lookup {['key', 'team', 'user'][i]} failed: {r}"
|
||||
)
|
||||
|
||||
def _increment_top_level_request_and_spend_metrics(
|
||||
|
|
@ -1528,9 +1407,6 @@ class PrometheusLogger(CustomLogger):
|
|||
user_api_team_alias = standard_logging_payload["metadata"][
|
||||
"user_api_key_team_alias"
|
||||
]
|
||||
user_api_key_org_id = standard_logging_payload["metadata"].get(
|
||||
"user_api_key_org_id"
|
||||
)
|
||||
|
||||
try:
|
||||
self.litellm_llm_api_failed_requests_metric.labels(
|
||||
|
|
@ -1546,10 +1422,6 @@ class PrometheusLogger(CustomLogger):
|
|||
),
|
||||
).inc()
|
||||
self.set_llm_deployment_failure_metrics(kwargs)
|
||||
await self._set_org_budget_metrics_after_api_request(
|
||||
org_id=user_api_key_org_id,
|
||||
response_cost=0,
|
||||
)
|
||||
except Exception as e:
|
||||
verbose_logger.exception(
|
||||
"prometheus Layer Error(): Exception occured - {}".format(str(e))
|
||||
|
|
@ -1757,11 +1629,9 @@ class PrometheusLogger(CustomLogger):
|
|||
client_ip=_metadata.get("requester_ip_address"),
|
||||
user_agent=_metadata.get("user_agent"),
|
||||
model_id=model_id,
|
||||
stream=(
|
||||
str(request_data.get("stream"))
|
||||
if litellm.prometheus_emit_stream_label
|
||||
else None
|
||||
),
|
||||
stream=str(request_data.get("stream"))
|
||||
if litellm.prometheus_emit_stream_label
|
||||
else None,
|
||||
)
|
||||
_labels = prometheus_label_factory(
|
||||
supported_enum_labels=self.get_labels_for_metric(
|
||||
|
|
@ -2284,127 +2154,6 @@ class PrometheusLogger(CustomLogger):
|
|||
except Exception as e:
|
||||
verbose_logger.debug(f"Error recording guardrail metrics: {str(e)}")
|
||||
|
||||
########################################
|
||||
# Managed Batch Metric Recording Methods
|
||||
########################################
|
||||
|
||||
def record_managed_batch_created(
|
||||
self,
|
||||
model: Optional[str],
|
||||
api_provider: Optional[str],
|
||||
user: Optional[str],
|
||||
user_email: Optional[str],
|
||||
api_key_alias: Optional[str],
|
||||
):
|
||||
try:
|
||||
self.litellm_managed_batch_created_total.labels(
|
||||
model=model,
|
||||
api_provider=api_provider,
|
||||
user=user,
|
||||
user_email=user_email,
|
||||
api_key_alias=api_key_alias,
|
||||
).inc()
|
||||
except Exception as e:
|
||||
verbose_logger.warning(f"Error recording batch created metric: {e}")
|
||||
|
||||
def record_managed_file_size(
|
||||
self,
|
||||
size_bytes: int,
|
||||
purpose: str,
|
||||
file_type: str,
|
||||
model: Optional[str] = None,
|
||||
api_provider: Optional[str] = None,
|
||||
user: Optional[str] = None,
|
||||
):
|
||||
"""Record the size of a managed file. Uses a gauge (last-seen value per label combination)."""
|
||||
try:
|
||||
self.litellm_managed_file_size_bytes.labels(
|
||||
purpose=purpose,
|
||||
file_type=file_type,
|
||||
model=model or "",
|
||||
api_provider=api_provider or "",
|
||||
user=user or "",
|
||||
).set(size_bytes)
|
||||
except Exception as e:
|
||||
verbose_logger.warning(f"Error recording file size metric: {e}")
|
||||
|
||||
def record_managed_batch_duration(
|
||||
self,
|
||||
duration_seconds: float,
|
||||
model: Optional[str] = None,
|
||||
api_provider: Optional[str] = None,
|
||||
):
|
||||
try:
|
||||
self.litellm_managed_batch_duration_seconds.labels(
|
||||
model=model or "",
|
||||
api_provider=api_provider or "",
|
||||
).observe(duration_seconds)
|
||||
except Exception as e:
|
||||
verbose_logger.warning(f"Error recording batch duration metric: {e}")
|
||||
|
||||
def record_managed_file_created(
|
||||
self,
|
||||
model: Optional[str],
|
||||
api_provider: Optional[str],
|
||||
user: Optional[str],
|
||||
user_email: Optional[str],
|
||||
api_key_alias: Optional[str],
|
||||
):
|
||||
try:
|
||||
self.litellm_managed_file_created_total.labels(
|
||||
model=model,
|
||||
api_provider=api_provider,
|
||||
user=user,
|
||||
user_email=user_email,
|
||||
api_key_alias=api_key_alias,
|
||||
).inc()
|
||||
except Exception as e:
|
||||
verbose_logger.warning(f"Error recording file created metric: {e}")
|
||||
|
||||
def record_managed_file_deleted(self, result: str):
|
||||
"""Record a managed file deletion attempt. result is 'success' or 'blocked'."""
|
||||
try:
|
||||
self.litellm_managed_file_deleted_total.labels(result=result).inc()
|
||||
except Exception as e:
|
||||
verbose_logger.warning(f"Error recording file deleted metric: {e}")
|
||||
|
||||
def record_check_batch_cost_run(
|
||||
self,
|
||||
jobs_polled: int,
|
||||
processed_models: Optional[List[Tuple[Optional[str], Optional[str]]]] = None,
|
||||
):
|
||||
"""
|
||||
Record CheckBatchCost polling metrics.
|
||||
|
||||
Args:
|
||||
jobs_polled: Number of unprocessed batches found
|
||||
processed_models: List of (model, api_provider) tuples for processed jobs
|
||||
"""
|
||||
import time
|
||||
|
||||
try:
|
||||
self.litellm_check_batch_cost_last_run_timestamp.set(time.time())
|
||||
self.litellm_check_batch_cost_jobs_polled.set(jobs_polled)
|
||||
|
||||
if processed_models:
|
||||
for model, api_provider in processed_models:
|
||||
self.litellm_check_batch_cost_jobs_processed_total.labels(
|
||||
model=model or "",
|
||||
api_provider=api_provider or "",
|
||||
).inc()
|
||||
except Exception as e:
|
||||
verbose_logger.warning(f"Error recording check batch cost metrics: {e}")
|
||||
|
||||
def record_check_batch_cost_error(self, error_type: str):
|
||||
try:
|
||||
self.litellm_check_batch_cost_errors_total.labels(
|
||||
error_type=error_type,
|
||||
).inc()
|
||||
except Exception as e:
|
||||
verbose_logger.warning(
|
||||
f"Error recording check batch cost error metric: {e}"
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _get_exception_class_name(exception: Exception) -> str:
|
||||
exception_class_name = ""
|
||||
|
|
@ -2781,35 +2530,6 @@ class PrometheusLogger(CustomLogger):
|
|||
data_type="users",
|
||||
)
|
||||
|
||||
async def _initialize_org_budget_metrics(self):
|
||||
"""
|
||||
Initialize org budget metrics by reusing the generic pagination logic.
|
||||
"""
|
||||
from litellm.proxy.proxy_server import prisma_client
|
||||
|
||||
if prisma_client is None:
|
||||
verbose_logger.debug(
|
||||
"Prometheus: skipping org metrics initialization, DB not initialized"
|
||||
)
|
||||
return
|
||||
|
||||
async def fetch_orgs(page_size: int, page: int) -> Tuple[list, Optional[int]]:
|
||||
skip = (page - 1) * page_size
|
||||
orgs = await prisma_client.db.litellm_organizationtable.find_many(
|
||||
skip=skip,
|
||||
take=page_size,
|
||||
order={"created_at": "desc"},
|
||||
include={"litellm_budget_table": True},
|
||||
)
|
||||
total_count = await prisma_client.db.litellm_organizationtable.count()
|
||||
return orgs, total_count
|
||||
|
||||
await self._initialize_budget_metrics(
|
||||
data_fetch_function=fetch_orgs,
|
||||
set_metrics_function=self._set_org_list_budget_metrics,
|
||||
data_type="orgs",
|
||||
)
|
||||
|
||||
async def initialize_remaining_budget_metrics(self):
|
||||
"""
|
||||
Handler for initializing remaining budget metrics for all teams to avoid metric discrepancies.
|
||||
|
|
@ -2844,11 +2564,10 @@ class PrometheusLogger(CustomLogger):
|
|||
"""
|
||||
Helper to initialize remaining budget metrics for all teams, API keys, and users.
|
||||
"""
|
||||
verbose_logger.debug("Emitting key, team, user, org budget metrics....")
|
||||
verbose_logger.debug("Emitting key, team, user budget metrics....")
|
||||
await self._initialize_team_budget_metrics()
|
||||
await self._initialize_api_key_budget_metrics()
|
||||
await self._initialize_user_budget_metrics()
|
||||
await self._initialize_org_budget_metrics()
|
||||
await self._initialize_user_and_team_count_metrics()
|
||||
|
||||
async def _initialize_user_and_team_count_metrics(self):
|
||||
|
|
@ -2904,22 +2623,6 @@ class PrometheusLogger(CustomLogger):
|
|||
for user in users:
|
||||
self._set_user_budget_metrics(user)
|
||||
|
||||
async def _set_org_list_budget_metrics(self, orgs: list):
|
||||
"""Helper function to set budget metrics for a list of orgs"""
|
||||
for org in orgs:
|
||||
budget_table = getattr(org, "litellm_budget_table", None)
|
||||
self._set_org_budget_metrics(
|
||||
org_id=org.organization_id or "",
|
||||
org_alias=org.organization_alias or "",
|
||||
spend=org.spend or 0.0,
|
||||
max_budget=budget_table.max_budget if budget_table else None,
|
||||
budget_reset_at=(
|
||||
getattr(budget_table, "budget_reset_at", None)
|
||||
if budget_table
|
||||
else None
|
||||
),
|
||||
)
|
||||
|
||||
async def _set_team_budget_metrics_after_api_request(
|
||||
self,
|
||||
user_api_team: Optional[str],
|
||||
|
|
@ -3041,113 +2744,6 @@ class PrometheusLogger(CustomLogger):
|
|||
)
|
||||
)
|
||||
|
||||
async def _set_org_budget_metrics_after_api_request(
|
||||
self,
|
||||
org_id: Optional[str],
|
||||
response_cost: float,
|
||||
):
|
||||
"""
|
||||
Set org budget metrics after an LLM API request
|
||||
|
||||
- Fetches org info via cache (get_org_object)
|
||||
- Sets org budget metrics
|
||||
"""
|
||||
if not org_id:
|
||||
return
|
||||
|
||||
from litellm.proxy.auth.auth_checks import get_org_object
|
||||
from litellm.proxy.proxy_server import prisma_client, user_api_key_cache
|
||||
|
||||
if prisma_client is None:
|
||||
return
|
||||
|
||||
try:
|
||||
org_info = await get_org_object(
|
||||
org_id=org_id,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
include_budget_table=True,
|
||||
)
|
||||
except Exception as e:
|
||||
verbose_logger.debug(
|
||||
f"[Non-Blocking] Prometheus: Error getting org info: {str(e)}"
|
||||
)
|
||||
return
|
||||
|
||||
if org_info is None:
|
||||
return
|
||||
|
||||
org_alias = org_info.organization_alias or ""
|
||||
_total_org_spend = (org_info.spend or 0.0) + response_cost
|
||||
budget_table = org_info.litellm_budget_table
|
||||
max_budget = budget_table.max_budget if budget_table else None
|
||||
budget_reset_at = (
|
||||
getattr(budget_table, "budget_reset_at", None) if budget_table else None
|
||||
)
|
||||
|
||||
self._set_org_budget_metrics(
|
||||
org_id=org_id,
|
||||
org_alias=org_alias,
|
||||
spend=_total_org_spend,
|
||||
max_budget=max_budget,
|
||||
budget_reset_at=budget_reset_at,
|
||||
)
|
||||
|
||||
def _set_org_budget_metrics(
|
||||
self,
|
||||
org_id: str,
|
||||
org_alias: str,
|
||||
spend: float,
|
||||
max_budget: Optional[float],
|
||||
budget_reset_at: Optional[datetime],
|
||||
):
|
||||
"""
|
||||
Set org budget metrics for a single org
|
||||
|
||||
- Remaining Budget
|
||||
- Max Budget
|
||||
- Budget Reset At
|
||||
"""
|
||||
enum_values = UserAPIKeyLabelValues(
|
||||
org_id=org_id,
|
||||
org_alias=org_alias,
|
||||
)
|
||||
|
||||
_labels = prometheus_label_factory(
|
||||
supported_enum_labels=self.get_labels_for_metric(
|
||||
metric_name="litellm_remaining_org_budget_metric"
|
||||
),
|
||||
enum_values=enum_values,
|
||||
)
|
||||
self.litellm_remaining_org_budget_metric.labels(**_labels).set(
|
||||
self._safe_get_remaining_budget(
|
||||
max_budget=max_budget,
|
||||
spend=spend,
|
||||
)
|
||||
)
|
||||
|
||||
if max_budget is not None:
|
||||
_labels = prometheus_label_factory(
|
||||
supported_enum_labels=self.get_labels_for_metric(
|
||||
metric_name="litellm_org_max_budget_metric"
|
||||
),
|
||||
enum_values=enum_values,
|
||||
)
|
||||
self.litellm_org_max_budget_metric.labels(**_labels).set(max_budget)
|
||||
|
||||
if budget_reset_at is not None:
|
||||
_labels = prometheus_label_factory(
|
||||
supported_enum_labels=self.get_labels_for_metric(
|
||||
metric_name="litellm_org_budget_remaining_hours_metric"
|
||||
),
|
||||
enum_values=enum_values,
|
||||
)
|
||||
self.litellm_org_budget_remaining_hours_metric.labels(**_labels).set(
|
||||
self._get_remaining_hours_for_budget_reset(
|
||||
budget_reset_at=budget_reset_at
|
||||
)
|
||||
)
|
||||
|
||||
def _set_key_budget_metrics(self, user_api_key_dict: UserAPIKeyAuth):
|
||||
"""
|
||||
Set virtual key budget metrics
|
||||
|
|
|
|||
|
|
@ -91,6 +91,34 @@ UNSUPPORTED_BEDROCK_CONVERSE_BETA_PATTERNS = [
|
|||
"compact-2026-01-12", # The compact beta feature is not currently supported on the Converse and ConverseStream APIs
|
||||
]
|
||||
|
||||
# Models that support Bedrock's native structured outputs API (outputConfig.textFormat)
|
||||
# Uses substring matching against the Bedrock model ID
|
||||
# Ref: https://docs.aws.amazon.com/bedrock/latest/userguide/structured-output.html
|
||||
BEDROCK_NATIVE_STRUCTURED_OUTPUT_MODELS = {
|
||||
# Anthropic Claude 4.5+
|
||||
"claude-haiku-4-5",
|
||||
"claude-sonnet-4-5",
|
||||
"claude-opus-4-5",
|
||||
"claude-opus-4-6",
|
||||
# Qwen3
|
||||
"qwen3",
|
||||
# DeepSeek
|
||||
"deepseek-v3.1",
|
||||
# Gemma 3
|
||||
"gemma-3",
|
||||
# MiniMax
|
||||
"minimax-m2",
|
||||
# Mistral (magistral-small excluded: broken constrained decoding on Bedrock)
|
||||
"ministral",
|
||||
"mistral-large-3",
|
||||
"voxtral",
|
||||
# Moonshot
|
||||
"kimi-k2",
|
||||
# NVIDIA
|
||||
"nemotron-nano",
|
||||
# OpenAI (gpt-oss excluded: broken constrained decoding, works via tool-call fallback)
|
||||
}
|
||||
|
||||
|
||||
class AmazonConverseConfig(BaseConfig):
|
||||
"""
|
||||
|
|
@ -465,7 +493,8 @@ class AmazonConverseConfig(BaseConfig):
|
|||
budget = thinking.get("budget_tokens")
|
||||
if isinstance(budget, int) and budget < BEDROCK_MIN_THINKING_BUDGET_TOKENS:
|
||||
verbose_logger.debug(
|
||||
"Bedrock requires thinking.budget_tokens >= %d, got %d. Clamping to minimum.",
|
||||
"Bedrock requires thinking.budget_tokens >= %d, got %d. "
|
||||
"Clamping to minimum.",
|
||||
BEDROCK_MIN_THINKING_BUDGET_TOKENS,
|
||||
budget,
|
||||
)
|
||||
|
|
@ -734,20 +763,10 @@ class AmazonConverseConfig(BaseConfig):
|
|||
return _tool
|
||||
|
||||
@staticmethod
|
||||
def _supports_native_structured_outputs(
|
||||
model: str, custom_llm_provider: Optional[str] = None
|
||||
) -> bool:
|
||||
"""Check if the Bedrock model supports native structured outputs (outputConfig.textFormat).
|
||||
|
||||
Delegates to the standard ``supports_native_structured_output`` utility
|
||||
which looks up the flag in ``litellm.model_cost`` via
|
||||
``_get_model_info_helper``.
|
||||
Ref: https://docs.aws.amazon.com/bedrock/latest/userguide/structured-output.html
|
||||
"""
|
||||
from litellm.utils import supports_native_structured_output
|
||||
|
||||
return supports_native_structured_output(
|
||||
model=model, custom_llm_provider=custom_llm_provider
|
||||
def _supports_native_structured_outputs(model: str) -> bool:
|
||||
"""Check if the Bedrock model supports native structured outputs (outputConfig.textFormat)."""
|
||||
return any(
|
||||
substring in model for substring in BEDROCK_NATIVE_STRUCTURED_OUTPUT_MODELS
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
|
|
@ -894,9 +913,7 @@ class AmazonConverseConfig(BaseConfig):
|
|||
)
|
||||
if param == "tool_choice":
|
||||
_tool_choice_value = self.map_tool_choice_values(
|
||||
model=model,
|
||||
tool_choice=value,
|
||||
drop_params=drop_params, # type: ignore
|
||||
model=model, tool_choice=value, drop_params=drop_params # type: ignore
|
||||
)
|
||||
if _tool_choice_value is not None:
|
||||
optional_params["tool_choice"] = _tool_choice_value
|
||||
|
|
@ -989,10 +1006,7 @@ class AmazonConverseConfig(BaseConfig):
|
|||
if "type" in value and value["type"] == "text":
|
||||
return optional_params
|
||||
|
||||
if (
|
||||
self._supports_native_structured_outputs(model, self.custom_llm_provider)
|
||||
and json_schema is not None
|
||||
):
|
||||
if self._supports_native_structured_outputs(model) and json_schema is not None:
|
||||
# Use Bedrock's native structured outputs API (outputConfig.textFormat)
|
||||
# No synthetic tool injection, no fake_stream needed.
|
||||
# Requires an explicit schema — json_object with no schema falls through
|
||||
|
|
@ -1432,16 +1446,6 @@ class AmazonConverseConfig(BaseConfig):
|
|||
original_tools, model, headers, additional_request_params
|
||||
)
|
||||
|
||||
# Append cachePoint to tools if cache_control_injection_points has tool_config
|
||||
cache_injection_points = additional_request_params.pop(
|
||||
"cache_control_injection_points", None
|
||||
)
|
||||
if cache_injection_points and len(bedrock_tools) > 0:
|
||||
for point in cache_injection_points:
|
||||
if point.get("location") == "tool_config":
|
||||
bedrock_tools.append({"cachePoint": {"type": "default"}})
|
||||
break
|
||||
|
||||
bedrock_tool_config: Optional[ToolConfigBlock] = None
|
||||
if len(bedrock_tools) > 0:
|
||||
tool_choice_values: ToolChoiceValuesBlock = inference_params.pop(
|
||||
|
|
|
|||
|
|
@ -260,7 +260,6 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
|
|||
guardrail_to_apply: "CustomGuardrail",
|
||||
litellm_logging_obj: Optional[Any] = None,
|
||||
user_api_key_dict: Optional[Any] = None,
|
||||
request_data: Optional[dict] = None,
|
||||
) -> Any:
|
||||
"""
|
||||
Process output response by applying guardrails to text content.
|
||||
|
|
@ -309,21 +308,15 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
|
|||
|
||||
# Step 2: Apply guardrail to all texts and tool calls in batch
|
||||
if texts_to_check or tool_calls_to_check:
|
||||
# Use the real request_data if provided (proxy path), otherwise
|
||||
# create a standalone dict (SDK / direct-call path).
|
||||
if request_data is None:
|
||||
request_data = {"response": response}
|
||||
else:
|
||||
if "response" not in request_data:
|
||||
request_data["response"] = response
|
||||
# Create a request_data dict with response info and user API key metadata
|
||||
request_data: dict = {"response": response}
|
||||
|
||||
# Add user API key metadata with prefixed keys
|
||||
if "litellm_metadata" not in request_data:
|
||||
user_metadata = self.transform_user_api_key_dict_to_metadata(
|
||||
user_api_key_dict
|
||||
)
|
||||
if user_metadata:
|
||||
request_data["litellm_metadata"] = user_metadata
|
||||
user_metadata = self.transform_user_api_key_dict_to_metadata(
|
||||
user_api_key_dict
|
||||
)
|
||||
if user_metadata:
|
||||
request_data["litellm_metadata"] = user_metadata
|
||||
|
||||
inputs = GenericGuardrailAPIInputs(texts=texts_to_check)
|
||||
if images_to_check:
|
||||
|
|
@ -371,7 +364,6 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
|
|||
guardrail_to_apply: "CustomGuardrail",
|
||||
litellm_logging_obj: Optional[Any] = None,
|
||||
user_api_key_dict: Optional[Any] = None,
|
||||
request_data: Optional[dict] = None,
|
||||
) -> List["ModelResponseStream"]:
|
||||
"""
|
||||
Process output streaming responses by applying guardrails to text content.
|
||||
|
|
@ -410,7 +402,6 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
|
|||
guardrail_to_apply=guardrail_to_apply,
|
||||
litellm_logging_obj=litellm_logging_obj,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
request_data=request_data,
|
||||
)
|
||||
|
||||
return responses_so_far
|
||||
|
|
@ -445,21 +436,15 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
|
|||
|
||||
# Step 3: Apply guardrail to all combined texts in batch
|
||||
if texts_to_check:
|
||||
# Use the real request_data if provided (proxy path), otherwise
|
||||
# create a standalone dict (SDK / direct-call path).
|
||||
if request_data is None:
|
||||
request_data = {"responses": responses_so_far}
|
||||
else:
|
||||
if "responses" not in request_data:
|
||||
request_data["responses"] = responses_so_far
|
||||
# Create a request_data dict with response info and user API key metadata
|
||||
request_data: dict = {"responses": responses_so_far}
|
||||
|
||||
# Add user API key metadata with prefixed keys
|
||||
if "litellm_metadata" not in request_data:
|
||||
user_metadata = self.transform_user_api_key_dict_to_metadata(
|
||||
user_api_key_dict
|
||||
)
|
||||
if user_metadata:
|
||||
request_data["litellm_metadata"] = user_metadata
|
||||
user_metadata = self.transform_user_api_key_dict_to_metadata(
|
||||
user_api_key_dict
|
||||
)
|
||||
if user_metadata:
|
||||
request_data["litellm_metadata"] = user_metadata
|
||||
|
||||
inputs = GenericGuardrailAPIInputs(texts=texts_to_check)
|
||||
if images_to_check:
|
||||
|
|
|
|||
|
|
@ -8,7 +8,6 @@ Run checks for:
|
|||
2. If user is in budget
|
||||
3. If end_user ('user' passed to /chat/completions, /embeddings endpoint) is in budget
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import re
|
||||
import time
|
||||
|
|
@ -30,7 +29,6 @@ from litellm.constants import (
|
|||
DEFAULT_MAX_RECURSE_DEPTH,
|
||||
EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE,
|
||||
)
|
||||
from litellm.litellm_core_utils.dd_tracing import tracer
|
||||
from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
|
||||
from litellm.proxy._types import (
|
||||
RBAC_ROLES,
|
||||
|
|
@ -40,7 +38,6 @@ from litellm.proxy._types import (
|
|||
LiteLLM_EndUserTable,
|
||||
Litellm_EntityType,
|
||||
LiteLLM_JWTAuth,
|
||||
LiteLLM_ManagedVectorStoresTable,
|
||||
LiteLLM_ObjectPermissionTable,
|
||||
LiteLLM_OrganizationMembershipTable,
|
||||
LiteLLM_OrganizationTable,
|
||||
|
|
@ -405,26 +402,23 @@ async def common_checks( # noqa: PLR0915
|
|||
# 1. If team is blocked
|
||||
if team_object is not None and team_object.blocked is True:
|
||||
raise Exception(
|
||||
f"Team={team_object.team_id} is blocked. Update via `/team/unblock` if you're an admin."
|
||||
f"Team={team_object.team_id} is blocked. Update via `/team/unblock` if your admin."
|
||||
)
|
||||
|
||||
# 2. If team can call model
|
||||
if _model and team_object:
|
||||
with tracer.trace("litellm.proxy.auth.common_checks.can_team_access_model"):
|
||||
if not await can_team_access_model(
|
||||
model=_model,
|
||||
team_object=team_object,
|
||||
llm_router=llm_router,
|
||||
team_model_aliases=(
|
||||
valid_token.team_model_aliases if valid_token else None
|
||||
),
|
||||
):
|
||||
raise ProxyException(
|
||||
message=f"Team not allowed to access model. Team={team_object.team_id}, Model={_model}. Allowed team models = {team_object.models}",
|
||||
type=ProxyErrorTypes.team_model_access_denied,
|
||||
param="model",
|
||||
code=status.HTTP_401_UNAUTHORIZED,
|
||||
)
|
||||
if not await can_team_access_model(
|
||||
model=_model,
|
||||
team_object=team_object,
|
||||
llm_router=llm_router,
|
||||
team_model_aliases=valid_token.team_model_aliases if valid_token else None,
|
||||
):
|
||||
raise ProxyException(
|
||||
message=f"Team not allowed to access model. Team={team_object.team_id}, Model={_model}. Allowed team models = {team_object.models}",
|
||||
type=ProxyErrorTypes.team_model_access_denied,
|
||||
param="model",
|
||||
code=status.HTTP_401_UNAUTHORIZED,
|
||||
)
|
||||
|
||||
# Require trace id for agent keys when agent has require_trace_id_on_calls_by_agent
|
||||
if valid_token is not None and valid_token.agent_id:
|
||||
|
|
@ -449,62 +443,54 @@ async def common_checks( # noqa: PLR0915
|
|||
|
||||
## 2.1 If user can call model (if personal key)
|
||||
if _model and team_object is None and user_object is not None:
|
||||
with tracer.trace("litellm.proxy.auth.common_checks.can_user_call_model"):
|
||||
await can_user_call_model(
|
||||
model=_model,
|
||||
llm_router=llm_router,
|
||||
user_object=user_object,
|
||||
)
|
||||
await can_user_call_model(
|
||||
model=_model,
|
||||
llm_router=llm_router,
|
||||
user_object=user_object,
|
||||
)
|
||||
|
||||
# 1.1 - 2.2 - 3.0.2 - 3.0.3: Project checks (blocked, model access, budget)
|
||||
with tracer.trace("litellm.proxy.auth.common_checks.run_project_checks"):
|
||||
await _run_project_checks(
|
||||
project_object=project_object,
|
||||
_model=_model,
|
||||
llm_router=llm_router,
|
||||
skip_budget_checks=skip_budget_checks,
|
||||
valid_token=valid_token,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
await _run_project_checks(
|
||||
project_object=project_object,
|
||||
_model=_model,
|
||||
llm_router=llm_router,
|
||||
skip_budget_checks=skip_budget_checks,
|
||||
valid_token=valid_token,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
|
||||
# If this is a free model, skip all budget checks
|
||||
if not skip_budget_checks:
|
||||
# 3. If team is in budget
|
||||
with tracer.trace("litellm.proxy.auth.common_checks.team_max_budget_check"):
|
||||
await _team_max_budget_check(
|
||||
team_object=team_object,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
valid_token=valid_token,
|
||||
)
|
||||
await _team_max_budget_check(
|
||||
team_object=team_object,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
valid_token=valid_token,
|
||||
)
|
||||
|
||||
# 3.0.5. If team is over soft budget (alert only, doesn't block)
|
||||
with tracer.trace("litellm.proxy.auth.common_checks.team_soft_budget_check"):
|
||||
await _team_soft_budget_check(
|
||||
team_object=team_object,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
valid_token=valid_token,
|
||||
)
|
||||
await _team_soft_budget_check(
|
||||
team_object=team_object,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
valid_token=valid_token,
|
||||
)
|
||||
|
||||
# 3.1. If organization is in budget
|
||||
with tracer.trace(
|
||||
"litellm.proxy.auth.common_checks.organization_max_budget_check"
|
||||
):
|
||||
await _organization_max_budget_check(
|
||||
valid_token=valid_token,
|
||||
team_object=team_object,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
await _organization_max_budget_check(
|
||||
valid_token=valid_token,
|
||||
team_object=team_object,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
|
||||
with tracer.trace("litellm.proxy.auth.common_checks.tag_max_budget_check"):
|
||||
await _tag_max_budget_check(
|
||||
request_body=request_body,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
valid_token=valid_token,
|
||||
)
|
||||
await _tag_max_budget_check(
|
||||
request_body=request_body,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
valid_token=valid_token,
|
||||
)
|
||||
|
||||
# 4. If user is in budget
|
||||
## 4.1 check personal budget, if personal key
|
||||
|
|
@ -522,15 +508,14 @@ async def common_checks( # noqa: PLR0915
|
|||
)
|
||||
|
||||
## 4.2 check team member budget, if team key
|
||||
with tracer.trace("litellm.proxy.auth.common_checks.check_team_member_budget"):
|
||||
await _check_team_member_budget(
|
||||
team_object=team_object,
|
||||
user_object=user_object,
|
||||
valid_token=valid_token,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
await _check_team_member_budget(
|
||||
team_object=team_object,
|
||||
user_object=user_object,
|
||||
valid_token=valid_token,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
|
||||
# 5. If end_user ('user' passed to /chat/completions, /embeddings endpoint) is in budget
|
||||
if (
|
||||
|
|
@ -569,21 +554,19 @@ async def common_checks( # noqa: PLR0915
|
|||
)
|
||||
|
||||
# 11. [OPTIONAL] Vector store checks - is the object allowed to access the vector store
|
||||
with tracer.trace("litellm.proxy.auth.common_checks.vector_store_access_check"):
|
||||
await vector_store_access_check(
|
||||
request_body=request_body,
|
||||
team_object=team_object,
|
||||
valid_token=valid_token,
|
||||
)
|
||||
await vector_store_access_check(
|
||||
request_body=request_body,
|
||||
team_object=team_object,
|
||||
valid_token=valid_token,
|
||||
)
|
||||
|
||||
# 12. [OPTIONAL] Tool allowlist - key/team allowed_tools (no DB in hot path)
|
||||
with tracer.trace("litellm.proxy.auth.common_checks.check_tools_allowlist"):
|
||||
await check_tools_allowlist(
|
||||
request_body=request_body,
|
||||
valid_token=valid_token,
|
||||
team_object=team_object,
|
||||
route=route,
|
||||
)
|
||||
await check_tools_allowlist(
|
||||
request_body=request_body,
|
||||
valid_token=valid_token,
|
||||
team_object=team_object,
|
||||
route=route,
|
||||
)
|
||||
|
||||
return True
|
||||
|
||||
|
|
@ -2296,71 +2279,6 @@ async def get_object_permission(
|
|||
return None
|
||||
|
||||
|
||||
@log_db_metrics
|
||||
async def get_managed_vector_store_rows_by_uuids(
|
||||
uuids: List[str],
|
||||
prisma_client: Optional[PrismaClient],
|
||||
user_api_key_cache: DualCache,
|
||||
parent_otel_span: Optional[Span] = None,
|
||||
proxy_logging_obj: Optional[ProxyLogging] = None,
|
||||
) -> List[LiteLLM_ManagedVectorStoresTable]:
|
||||
"""
|
||||
Fetch managed vector store rows by their internal UUIDs.
|
||||
|
||||
Follows the get_team_object / get_key_object / get_object_permission pattern:
|
||||
cache-first lookup (in-memory / Redis), DB fallback only on cache miss.
|
||||
Critical-path DB access must go through this helper to avoid raw Prisma
|
||||
calls on the hot request path.
|
||||
"""
|
||||
if not uuids or prisma_client is None:
|
||||
return []
|
||||
|
||||
result: List[LiteLLM_ManagedVectorStoresTable] = []
|
||||
cache_misses: List[str] = []
|
||||
|
||||
for uuid in uuids:
|
||||
key = "managed_vector_store_id:{}".format(uuid)
|
||||
cached = await user_api_key_cache.async_get_cache(key=key)
|
||||
if cached is not None:
|
||||
if isinstance(cached, dict):
|
||||
result.append(LiteLLM_ManagedVectorStoresTable(**cached))
|
||||
elif isinstance(cached, LiteLLM_ManagedVectorStoresTable):
|
||||
result.append(cached)
|
||||
else:
|
||||
cache_misses.append(uuid)
|
||||
else:
|
||||
cache_misses.append(uuid)
|
||||
|
||||
if not cache_misses:
|
||||
return result
|
||||
|
||||
rows = await prisma_client.db.litellm_managedvectorstorestable.find_many(
|
||||
where={"vector_store_id": {"in": cache_misses}},
|
||||
take=len(cache_misses),
|
||||
)
|
||||
|
||||
for row in rows:
|
||||
row_dict = (
|
||||
row.model_dump()
|
||||
if hasattr(row, "model_dump")
|
||||
else (row.dict() if hasattr(row, "dict") else None)
|
||||
)
|
||||
if not isinstance(row_dict, dict) or not row_dict:
|
||||
row_dict = dict(row) if hasattr(row, "__dict__") else {}
|
||||
if not row_dict:
|
||||
continue
|
||||
cached_obj = LiteLLM_ManagedVectorStoresTable(**row_dict)
|
||||
key = "managed_vector_store_id:{}".format(cached_obj.vector_store_id)
|
||||
await user_api_key_cache.async_set_cache(
|
||||
key=key,
|
||||
value=row_dict,
|
||||
ttl=DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL,
|
||||
)
|
||||
result.append(cached_obj)
|
||||
|
||||
return result
|
||||
|
||||
|
||||
@log_db_metrics
|
||||
async def get_org_object(
|
||||
org_id: str,
|
||||
|
|
@ -2877,15 +2795,7 @@ async def _virtual_key_max_budget_check(
|
|||
Triggers a budget alert if the token is over it's max budget.
|
||||
|
||||
"""
|
||||
if valid_token.max_budget is not None:
|
||||
from litellm.proxy.proxy_server import get_current_spend
|
||||
|
||||
# Read spend from cross-pod counter (Redis-first) or cached object (fallback)
|
||||
spend = await get_current_spend(
|
||||
counter_key=f"spend:key:{valid_token.token}",
|
||||
fallback_spend=valid_token.spend or 0.0,
|
||||
)
|
||||
|
||||
if valid_token.spend is not None and valid_token.max_budget is not None:
|
||||
####################################
|
||||
# collect information for alerting #
|
||||
####################################
|
||||
|
|
@ -2897,7 +2807,7 @@ async def _virtual_key_max_budget_check(
|
|||
|
||||
call_info = CallInfo(
|
||||
token=valid_token.token,
|
||||
spend=spend,
|
||||
spend=valid_token.spend,
|
||||
max_budget=valid_token.max_budget,
|
||||
soft_budget=valid_token.soft_budget,
|
||||
user_id=valid_token.user_id,
|
||||
|
|
@ -2918,9 +2828,9 @@ async def _virtual_key_max_budget_check(
|
|||
# collect information for alerting #
|
||||
####################################
|
||||
|
||||
if spend >= valid_token.max_budget:
|
||||
if valid_token.spend >= valid_token.max_budget:
|
||||
raise litellm.BudgetExceededError(
|
||||
current_cost=spend,
|
||||
current_cost=valid_token.spend,
|
||||
max_budget=valid_token.max_budget,
|
||||
)
|
||||
|
||||
|
|
@ -3051,14 +2961,6 @@ async def _check_team_member_budget(
|
|||
team_member_budget = team_membership.litellm_budget_table.max_budget
|
||||
team_member_spend = team_membership.spend or 0.0
|
||||
|
||||
# Read from cross-pod counter (Redis-first) if available
|
||||
from litellm.proxy.proxy_server import get_current_spend
|
||||
|
||||
team_member_spend = await get_current_spend(
|
||||
counter_key=f"spend:team_member:{valid_token.user_id}:{team_object.team_id}",
|
||||
fallback_spend=team_member_spend,
|
||||
)
|
||||
|
||||
if team_member_spend >= team_member_budget:
|
||||
raise litellm.BudgetExceededError(
|
||||
current_cost=team_member_spend,
|
||||
|
|
@ -3079,39 +2981,35 @@ async def _team_max_budget_check(
|
|||
BudgetExceededError if the team is over it's max budget.
|
||||
Triggers a budget alert if the team is over it's max budget.
|
||||
"""
|
||||
if team_object is not None and team_object.max_budget is not None:
|
||||
from litellm.proxy.proxy_server import get_current_spend
|
||||
|
||||
# Read spend from cross-pod counter (Redis-first) or cached object (fallback)
|
||||
spend = await get_current_spend(
|
||||
counter_key=f"spend:team:{team_object.team_id}",
|
||||
fallback_spend=team_object.spend or 0.0,
|
||||
)
|
||||
|
||||
if spend > team_object.max_budget:
|
||||
if valid_token:
|
||||
call_info = CallInfo(
|
||||
token=valid_token.token,
|
||||
spend=spend,
|
||||
max_budget=team_object.max_budget,
|
||||
user_id=valid_token.user_id,
|
||||
team_id=valid_token.team_id,
|
||||
team_alias=valid_token.team_alias,
|
||||
organization_id=valid_token.org_id,
|
||||
event_group=Litellm_EntityType.TEAM,
|
||||
)
|
||||
asyncio.create_task(
|
||||
proxy_logging_obj.budget_alerts(
|
||||
type="team_budget",
|
||||
user_info=call_info,
|
||||
)
|
||||
)
|
||||
|
||||
raise litellm.BudgetExceededError(
|
||||
current_cost=spend,
|
||||
if (
|
||||
team_object is not None
|
||||
and team_object.max_budget is not None
|
||||
and team_object.spend is not None
|
||||
and team_object.spend > team_object.max_budget
|
||||
):
|
||||
if valid_token:
|
||||
call_info = CallInfo(
|
||||
token=valid_token.token,
|
||||
spend=team_object.spend,
|
||||
max_budget=team_object.max_budget,
|
||||
message=f"Budget has been exceeded! Team={team_object.team_id} Current cost: {spend}, Max budget: {team_object.max_budget}",
|
||||
user_id=valid_token.user_id,
|
||||
team_id=valid_token.team_id,
|
||||
team_alias=valid_token.team_alias,
|
||||
organization_id=valid_token.org_id,
|
||||
event_group=Litellm_EntityType.TEAM,
|
||||
)
|
||||
asyncio.create_task(
|
||||
proxy_logging_obj.budget_alerts(
|
||||
type="team_budget",
|
||||
user_info=call_info,
|
||||
)
|
||||
)
|
||||
|
||||
raise litellm.BudgetExceededError(
|
||||
current_cost=team_object.spend,
|
||||
max_budget=team_object.max_budget,
|
||||
message=f"Budget has been exceeded! Team={team_object.team_id} Current cost: {team_object.spend}, Max budget: {team_object.max_budget}",
|
||||
)
|
||||
|
||||
|
||||
async def _team_soft_budget_check(
|
||||
|
|
|
|||
|
|
@ -54,47 +54,16 @@ class ResetBudgetJob:
|
|||
"""
|
||||
Resets the budget for all LiteLLM Team Members if their budget has expired
|
||||
"""
|
||||
budget_ids = [
|
||||
budget.budget_id
|
||||
for budget in budgets_to_reset
|
||||
if budget.budget_id is not None
|
||||
]
|
||||
|
||||
# Reset spend counters for affected team members.
|
||||
# Reset Redis directly so a transient failure doesn't leave stale
|
||||
# counters that get_current_spend would read as authoritative.
|
||||
try:
|
||||
from litellm.proxy.proxy_server import spend_counter_cache
|
||||
|
||||
memberships = await self.prisma_client.db.litellm_teammembership.find_many(
|
||||
where={"budget_id": {"in": budget_ids}}
|
||||
)
|
||||
for m in memberships:
|
||||
counter_key = f"spend:team_member:{m.user_id}:{m.team_id}"
|
||||
# Always reset in-memory
|
||||
spend_counter_cache.in_memory_cache.set_cache(
|
||||
key=counter_key, value=0.0
|
||||
)
|
||||
# Explicitly reset Redis with warning on failure
|
||||
if spend_counter_cache.redis_cache is not None:
|
||||
try:
|
||||
await spend_counter_cache.redis_cache.async_set_cache(
|
||||
key=counter_key, value=0.0
|
||||
)
|
||||
except Exception as redis_err:
|
||||
verbose_proxy_logger.warning(
|
||||
"Failed to reset team member spend counter in Redis %s: %s. "
|
||||
"Budget may be over-enforced until counter expires.",
|
||||
counter_key,
|
||||
redis_err,
|
||||
)
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.warning(
|
||||
"Failed to reset team member spend counters: %s", e
|
||||
)
|
||||
|
||||
return await self.prisma_client.db.litellm_teammembership.update_many(
|
||||
where={"budget_id": {"in": budget_ids}},
|
||||
where={
|
||||
"budget_id": {
|
||||
"in": [
|
||||
budget.budget_id
|
||||
for budget in budgets_to_reset
|
||||
if budget.budget_id is not None
|
||||
]
|
||||
}
|
||||
},
|
||||
data={
|
||||
"spend": 0,
|
||||
},
|
||||
|
|
@ -562,43 +531,6 @@ class ResetBudgetJob:
|
|||
"""
|
||||
try:
|
||||
item.spend = 0.0
|
||||
|
||||
# Reset the cross-pod spend counter.
|
||||
# Reset Redis directly (not via DualCache) so a Redis failure
|
||||
# doesn't silently leave a stale counter that get_current_spend
|
||||
# would read as authoritative, permanently blocking the user.
|
||||
from litellm.proxy.proxy_server import spend_counter_cache
|
||||
|
||||
counter_key = None
|
||||
if item_type == "key" and hasattr(item, "token") and item.token is not None:
|
||||
counter_key = f"spend:key:{item.token}"
|
||||
elif (
|
||||
item_type == "team"
|
||||
and hasattr(item, "team_id")
|
||||
and item.team_id is not None
|
||||
):
|
||||
counter_key = f"spend:team:{item.team_id}"
|
||||
|
||||
if counter_key is not None:
|
||||
# Always reset in-memory (local fallback)
|
||||
spend_counter_cache.in_memory_cache.set_cache(
|
||||
key=counter_key, value=0.0
|
||||
)
|
||||
# Explicitly reset Redis with warning on failure
|
||||
if spend_counter_cache.redis_cache is not None:
|
||||
try:
|
||||
await spend_counter_cache.redis_cache.async_set_cache(
|
||||
key=counter_key, value=0.0
|
||||
)
|
||||
except Exception as redis_err:
|
||||
verbose_proxy_logger.warning(
|
||||
"Failed to reset spend counter in Redis for %s key=%s: %s. "
|
||||
"Budget may be over-enforced until counter expires.",
|
||||
item_type,
|
||||
counter_key,
|
||||
redis_err,
|
||||
)
|
||||
|
||||
if hasattr(item, "budget_duration") and item.budget_duration is not None:
|
||||
# Get standardized reset time based on budget duration
|
||||
from litellm.proxy.common_utils.timezone_utils import (
|
||||
|
|
|
|||
|
|
@ -2,15 +2,12 @@
|
|||
Functions to create audit logs for LiteLLM Proxy
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
from litellm._uuid import uuid
|
||||
from datetime import datetime, timezone
|
||||
from typing import Dict
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm._uuid import uuid
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.proxy._types import (
|
||||
AUDIT_ACTIONS,
|
||||
LiteLLM_AuditLogs,
|
||||
|
|
@ -18,99 +15,6 @@ from litellm.proxy._types import (
|
|||
Optional,
|
||||
UserAPIKeyAuth,
|
||||
)
|
||||
from litellm.types.utils import StandardAuditLogPayload
|
||||
|
||||
_audit_log_callback_cache: Dict[str, CustomLogger] = {}
|
||||
|
||||
|
||||
def _resolve_audit_log_callback(name: str) -> Optional[CustomLogger]:
|
||||
"""Resolve a string callback name to a CustomLogger instance, with caching."""
|
||||
if name in _audit_log_callback_cache:
|
||||
return _audit_log_callback_cache[name]
|
||||
|
||||
from litellm.litellm_core_utils.litellm_logging import (
|
||||
_init_custom_logger_compatible_class,
|
||||
)
|
||||
|
||||
instance = _init_custom_logger_compatible_class(
|
||||
logging_integration=name, # type: ignore
|
||||
internal_usage_cache=None,
|
||||
llm_router=None,
|
||||
)
|
||||
|
||||
if instance is not None:
|
||||
_audit_log_callback_cache[name] = instance
|
||||
return instance
|
||||
|
||||
|
||||
def _build_audit_log_payload(
|
||||
request_data: LiteLLM_AuditLogs,
|
||||
) -> StandardAuditLogPayload:
|
||||
"""Convert LiteLLM_AuditLogs to StandardAuditLogPayload for callback dispatch."""
|
||||
updated_at = ""
|
||||
if request_data.updated_at is not None:
|
||||
updated_at = request_data.updated_at.isoformat()
|
||||
|
||||
table_name_str: str = (
|
||||
request_data.table_name.value
|
||||
if isinstance(request_data.table_name, LitellmTableNames)
|
||||
else str(request_data.table_name)
|
||||
)
|
||||
|
||||
return StandardAuditLogPayload(
|
||||
id=request_data.id,
|
||||
updated_at=updated_at,
|
||||
changed_by=request_data.changed_by or "",
|
||||
changed_by_api_key=request_data.changed_by_api_key or "",
|
||||
action=request_data.action,
|
||||
table_name=table_name_str,
|
||||
object_id=request_data.object_id,
|
||||
before_value=request_data.before_value,
|
||||
updated_values=request_data.updated_values,
|
||||
)
|
||||
|
||||
|
||||
def _audit_log_task_done_callback(task: asyncio.Task) -> None:
|
||||
"""Log exceptions from audit log callback tasks so they don't slip through silently."""
|
||||
try:
|
||||
exc = task.exception()
|
||||
except asyncio.CancelledError:
|
||||
return
|
||||
if exc is not None:
|
||||
verbose_proxy_logger.error(
|
||||
"Audit log callback task failed: %s", exc, exc_info=exc
|
||||
)
|
||||
|
||||
|
||||
async def _dispatch_audit_log_to_callbacks(
|
||||
request_data: LiteLLM_AuditLogs,
|
||||
) -> None:
|
||||
"""Dispatch audit log to all registered audit_log_callbacks."""
|
||||
if not litellm.audit_log_callbacks:
|
||||
return
|
||||
|
||||
payload = _build_audit_log_payload(request_data)
|
||||
|
||||
for callback in litellm.audit_log_callbacks:
|
||||
try:
|
||||
resolved: Optional[CustomLogger] = (
|
||||
callback if isinstance(callback, CustomLogger) else None
|
||||
)
|
||||
if isinstance(callback, str):
|
||||
resolved = _resolve_audit_log_callback(callback)
|
||||
if resolved is None:
|
||||
verbose_proxy_logger.warning(
|
||||
"Could not resolve audit log callback: %s", callback
|
||||
)
|
||||
continue
|
||||
|
||||
if isinstance(resolved, CustomLogger):
|
||||
task = asyncio.create_task(resolved.async_log_audit_log_event(payload))
|
||||
task.add_done_callback(_audit_log_task_done_callback)
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.error(
|
||||
"Failed dispatching audit log to callback: %s", e
|
||||
)
|
||||
|
||||
|
||||
async def create_object_audit_log(
|
||||
|
|
@ -136,22 +40,20 @@ async def create_object_audit_log(
|
|||
"""
|
||||
from litellm.secret_managers.main import get_secret_bool
|
||||
|
||||
_store_audit_logs: Optional[bool] = litellm.store_audit_logs or get_secret_bool(
|
||||
store_audit_logs = litellm.store_audit_logs or get_secret_bool(
|
||||
"LITELLM_STORE_AUDIT_LOGS"
|
||||
)
|
||||
|
||||
if _store_audit_logs is not True:
|
||||
if store_audit_logs is not True:
|
||||
return
|
||||
|
||||
_changed_by = (
|
||||
litellm_changed_by or user_api_key_dict.user_id or litellm_proxy_admin_name
|
||||
)
|
||||
|
||||
await create_audit_log_for_update(
|
||||
request_data=LiteLLM_AuditLogs(
|
||||
id=str(uuid.uuid4()),
|
||||
updated_at=datetime.now(timezone.utc),
|
||||
changed_by=_changed_by,
|
||||
changed_by=litellm_changed_by
|
||||
or user_api_key_dict.user_id
|
||||
or litellm_proxy_admin_name,
|
||||
changed_by_api_key=user_api_key_dict.api_key,
|
||||
table_name=table_name,
|
||||
object_id=object_id,
|
||||
|
|
@ -168,10 +70,10 @@ async def create_audit_log_for_update(request_data: LiteLLM_AuditLogs):
|
|||
"""
|
||||
from litellm.secret_managers.main import get_secret_bool
|
||||
|
||||
_store_audit_logs: Optional[bool] = litellm.store_audit_logs or get_secret_bool(
|
||||
store_audit_logs = litellm.store_audit_logs or get_secret_bool(
|
||||
"LITELLM_STORE_AUDIT_LOGS"
|
||||
)
|
||||
if _store_audit_logs is not True:
|
||||
if store_audit_logs is not True:
|
||||
return
|
||||
|
||||
from litellm.proxy.proxy_server import premium_user, prisma_client
|
||||
|
|
@ -179,6 +81,9 @@ async def create_audit_log_for_update(request_data: LiteLLM_AuditLogs):
|
|||
if premium_user is not True:
|
||||
return
|
||||
|
||||
if prisma_client is None:
|
||||
raise Exception("prisma_client is None, no DB connected")
|
||||
|
||||
verbose_proxy_logger.debug("creating audit log for %s", request_data)
|
||||
|
||||
if isinstance(request_data.updated_values, dict):
|
||||
|
|
@ -187,15 +92,6 @@ async def create_audit_log_for_update(request_data: LiteLLM_AuditLogs):
|
|||
if isinstance(request_data.before_value, dict):
|
||||
request_data.before_value = json.dumps(request_data.before_value)
|
||||
|
||||
# Dispatch to external audit log callbacks regardless of DB availability
|
||||
await _dispatch_audit_log_to_callbacks(request_data)
|
||||
|
||||
if prisma_client is None:
|
||||
verbose_proxy_logger.error(
|
||||
"prisma_client is None, cannot write audit log to DB"
|
||||
)
|
||||
return
|
||||
|
||||
_request_data = request_data.model_dump(exclude_none=True)
|
||||
|
||||
try:
|
||||
|
|
@ -207,3 +103,5 @@ async def create_audit_log_for_update(request_data: LiteLLM_AuditLogs):
|
|||
except Exception as e:
|
||||
# [Non-Blocking Exception. Do not allow blocking LLM API call]
|
||||
verbose_proxy_logger.error(f"Failed Creating audit log {e}")
|
||||
|
||||
return
|
||||
|
|
|
|||
|
|
@ -72,7 +72,7 @@ def test_get_api_key_with_custom_litellm_key_header(
|
|||
def test_team_metadata_with_tags_flows_through_jwt_auth():
|
||||
"""
|
||||
Test that team_metadata (specifically tags) flows through JWT authentication.
|
||||
|
||||
|
||||
This is a regression test for the issue where JWT auth was not populating
|
||||
team_metadata, causing team-level tags to be missing in litellm_pre_call_utils.py
|
||||
"""
|
||||
|
|
@ -87,7 +87,7 @@ def test_team_metadata_with_tags_flows_through_jwt_auth():
|
|||
rpm_limit=100,
|
||||
models=["gpt-4", "gpt-3.5-turbo"],
|
||||
)
|
||||
|
||||
|
||||
# Simulate constructing UserAPIKeyAuth like we do in JWT auth
|
||||
# This is the pattern from user_api_key_auth.py lines 552-587
|
||||
user_api_key_auth = UserAPIKeyAuth(
|
||||
|
|
@ -100,16 +100,14 @@ def test_team_metadata_with_tags_flows_through_jwt_auth():
|
|||
user_role="internal_user",
|
||||
user_id="test-user",
|
||||
)
|
||||
|
||||
|
||||
# Verify team_metadata is set
|
||||
assert (
|
||||
user_api_key_auth.team_metadata is not None
|
||||
), "team_metadata should be populated"
|
||||
assert user_api_key_auth.team_metadata is not None, "team_metadata should be populated"
|
||||
assert user_api_key_auth.team_metadata == team_object.metadata, (
|
||||
f"team_metadata not correctly mapped. "
|
||||
f"Expected: {team_object.metadata}, Got: {user_api_key_auth.team_metadata}"
|
||||
)
|
||||
|
||||
|
||||
# Specifically verify tags are present
|
||||
assert "tags" in user_api_key_auth.team_metadata, "tags should be in team_metadata"
|
||||
assert user_api_key_auth.team_metadata["tags"] == ["production", "high-priority"], (
|
||||
|
|
@ -120,7 +118,7 @@ def test_team_metadata_with_tags_flows_through_jwt_auth():
|
|||
|
||||
def test_route_checks_is_llm_api_route():
|
||||
"""Test RouteChecks.is_llm_api_route() correctly identifies LLM API routes including passthrough endpoints"""
|
||||
|
||||
|
||||
# Test OpenAI routes
|
||||
openai_routes = [
|
||||
"/v1/chat/completions",
|
||||
|
|
@ -144,22 +142,18 @@ def test_route_checks_is_llm_api_route():
|
|||
"/v1/realtime",
|
||||
"/realtime",
|
||||
]
|
||||
|
||||
|
||||
for route in openai_routes:
|
||||
assert RouteChecks.is_llm_api_route(
|
||||
route=route
|
||||
), f"Route {route} should be identified as LLM API route"
|
||||
assert RouteChecks.is_llm_api_route(route=route), f"Route {route} should be identified as LLM API route"
|
||||
|
||||
# Test Anthropic routes
|
||||
anthropic_routes = [
|
||||
"/v1/messages",
|
||||
"/v1/messages/count_tokens",
|
||||
]
|
||||
|
||||
|
||||
for route in anthropic_routes:
|
||||
assert RouteChecks.is_llm_api_route(
|
||||
route=route
|
||||
), f"Route {route} should be identified as LLM API route"
|
||||
assert RouteChecks.is_llm_api_route(route=route), f"Route {route} should be identified as LLM API route"
|
||||
|
||||
# Test passthrough routes (this is the key improvement over the old route checking)
|
||||
passthrough_routes = [
|
||||
|
|
@ -177,11 +171,9 @@ def test_route_checks_is_llm_api_route():
|
|||
"/vllm/v1/chat/completions",
|
||||
"/mistral/v1/chat/completions",
|
||||
]
|
||||
|
||||
|
||||
for route in passthrough_routes:
|
||||
assert RouteChecks.is_llm_api_route(
|
||||
route=route
|
||||
), f"Route {route} should be identified as LLM API route"
|
||||
assert RouteChecks.is_llm_api_route(route=route), f"Route {route} should be identified as LLM API route"
|
||||
|
||||
# Test MCP routes
|
||||
mcp_routes = [
|
||||
|
|
@ -189,11 +181,9 @@ def test_route_checks_is_llm_api_route():
|
|||
"/mcp/",
|
||||
"/mcp/test",
|
||||
]
|
||||
|
||||
|
||||
for route in mcp_routes:
|
||||
assert RouteChecks.is_llm_api_route(
|
||||
route=route
|
||||
), f"Route {route} should be identified as LLM API route"
|
||||
assert RouteChecks.is_llm_api_route(route=route), f"Route {route} should be identified as LLM API route"
|
||||
|
||||
# Test LiteLLM native RAG routes
|
||||
rag_routes = [
|
||||
|
|
@ -203,9 +193,7 @@ def test_route_checks_is_llm_api_route():
|
|||
"/v1/rag/query",
|
||||
]
|
||||
for route in rag_routes:
|
||||
assert RouteChecks.is_llm_api_route(
|
||||
route=route
|
||||
), f"Route {route} should be identified as LLM API route"
|
||||
assert RouteChecks.is_llm_api_route(route=route), f"Route {route} should be identified as LLM API route"
|
||||
|
||||
# Test routes with placeholders
|
||||
placeholder_routes = [
|
||||
|
|
@ -218,11 +206,9 @@ def test_route_checks_is_llm_api_route():
|
|||
"/v1/batches/batch_123",
|
||||
"/batches/batch_123",
|
||||
]
|
||||
|
||||
|
||||
for route in placeholder_routes:
|
||||
assert RouteChecks.is_llm_api_route(
|
||||
route=route
|
||||
), f"Route {route} should be identified as LLM API route"
|
||||
assert RouteChecks.is_llm_api_route(route=route), f"Route {route} should be identified as LLM API route"
|
||||
|
||||
# Test Azure OpenAI routes
|
||||
azure_routes = [
|
||||
|
|
@ -231,11 +217,9 @@ def test_route_checks_is_llm_api_route():
|
|||
"/engines/gpt-4/chat/completions",
|
||||
"/engines/gpt-3.5-turbo/completions",
|
||||
]
|
||||
|
||||
|
||||
for route in azure_routes:
|
||||
assert RouteChecks.is_llm_api_route(
|
||||
route=route
|
||||
), f"Route {route} should be identified as LLM API route"
|
||||
assert RouteChecks.is_llm_api_route(route=route), f"Route {route} should be identified as LLM API route"
|
||||
|
||||
# Test non-LLM routes (should return False)
|
||||
non_llm_routes = [
|
||||
|
|
@ -252,11 +236,9 @@ def test_route_checks_is_llm_api_route():
|
|||
"/debug",
|
||||
"/test",
|
||||
]
|
||||
|
||||
|
||||
for route in non_llm_routes:
|
||||
assert not RouteChecks.is_llm_api_route(
|
||||
route=route
|
||||
), f"Route {route} should NOT be identified as LLM API route"
|
||||
assert not RouteChecks.is_llm_api_route(route=route), f"Route {route} should NOT be identified as LLM API route"
|
||||
|
||||
# Test invalid inputs
|
||||
invalid_inputs = [
|
||||
|
|
@ -266,11 +248,9 @@ def test_route_checks_is_llm_api_route():
|
|||
{},
|
||||
"",
|
||||
]
|
||||
|
||||
|
||||
for invalid_input in invalid_inputs:
|
||||
assert not RouteChecks.is_llm_api_route(
|
||||
route=invalid_input
|
||||
), f"Invalid input {invalid_input} should return False"
|
||||
assert not RouteChecks.is_llm_api_route(route=invalid_input), f"Invalid input {invalid_input} should return False"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -279,7 +259,7 @@ async def test_proxy_admin_expired_key_from_cache():
|
|||
Test that PROXY_ADMIN keys retrieved from cache are checked for expiration
|
||||
before being returned. This prevents expired keys from bypassing expiration checks
|
||||
when retrieved from cache (which normally happens at lines 1014-1036).
|
||||
|
||||
|
||||
Regression test for issue where PROXY_ADMIN keys from cache skipped expiration check.
|
||||
"""
|
||||
from datetime import datetime, timedelta, timezone
|
||||
|
|
@ -300,42 +280,39 @@ async def test_proxy_admin_expired_key_from_cache():
|
|||
api_key = "sk-test-proxy-admin-key"
|
||||
hashed_key = hash_token(api_key)
|
||||
expired_time = datetime.now(timezone.utc) - timedelta(hours=1) # Expired 1 hour ago
|
||||
|
||||
|
||||
expired_token = UserAPIKeyAuth(
|
||||
api_key=api_key,
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN,
|
||||
expires=expired_time,
|
||||
token=hashed_key,
|
||||
)
|
||||
|
||||
|
||||
# Mock cache to return the expired token
|
||||
mock_cache = AsyncMock()
|
||||
mock_cache.async_get_cache = AsyncMock(return_value=expired_token)
|
||||
mock_cache.delete_cache = MagicMock()
|
||||
|
||||
|
||||
# Mock proxy_logging_obj
|
||||
mock_proxy_logging_obj = MagicMock()
|
||||
mock_proxy_logging_obj.internal_usage_cache = MagicMock()
|
||||
mock_proxy_logging_obj.internal_usage_cache.dual_cache = AsyncMock()
|
||||
mock_proxy_logging_obj.internal_usage_cache.dual_cache.async_delete_cache = (
|
||||
AsyncMock()
|
||||
)
|
||||
mock_proxy_logging_obj.internal_usage_cache.dual_cache.async_delete_cache = AsyncMock()
|
||||
# Mock post_call_failure_hook as async function returning None (no transformation)
|
||||
mock_proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None)
|
||||
|
||||
|
||||
# Mock prisma_client
|
||||
mock_prisma_client = MagicMock()
|
||||
|
||||
|
||||
# Mock get_key_object to return expired token from cache
|
||||
with patch(
|
||||
"litellm.proxy.auth.user_api_key_auth.get_key_object",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_get_key_object, patch(
|
||||
"litellm.proxy.auth.user_api_key_auth._delete_cache_key_object",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_delete_cache:
|
||||
) as mock_get_key_object, \
|
||||
patch("litellm.proxy.auth.user_api_key_auth._delete_cache_key_object", new_callable=AsyncMock) as mock_delete_cache:
|
||||
|
||||
mock_get_key_object.return_value = expired_token
|
||||
|
||||
|
||||
# Set attributes on proxy_server module (these are imported inside _user_api_key_auth_builder)
|
||||
import litellm.proxy.proxy_server as _proxy_server_mod
|
||||
|
||||
|
|
@ -354,12 +331,14 @@ async def test_proxy_admin_expired_key_from_cache():
|
|||
"litellm_proxy_admin_name": "admin",
|
||||
}
|
||||
_original_values = {
|
||||
attr: getattr(_proxy_server_mod, attr, None) for attr in _attrs_to_set
|
||||
attr: getattr(_proxy_server_mod, attr, None)
|
||||
for attr in _attrs_to_set
|
||||
}
|
||||
try:
|
||||
for attr, val in _attrs_to_set.items():
|
||||
setattr(_proxy_server_mod, attr, val)
|
||||
|
||||
|
||||
# Create a mock request
|
||||
request = Request(scope={"type": "http"})
|
||||
request._url = URL(url="/chat/completions")
|
||||
|
|
@ -379,41 +358,38 @@ async def test_proxy_admin_expired_key_from_cache():
|
|||
)
|
||||
|
||||
# Verify that ProxyException was raised with expired_key type
|
||||
assert hasattr(
|
||||
exc_info.value, "type"
|
||||
), "Exception should have 'type' attribute"
|
||||
assert (
|
||||
exc_info.value.type == ProxyErrorTypes.expired_key
|
||||
), f"Expected expired_key error type, got {exc_info.value.type}"
|
||||
assert "Expired Key" in str(
|
||||
exc_info.value.message
|
||||
), f"Exception message should mention 'Expired Key', got: {exc_info.value.message}"
|
||||
assert hasattr(exc_info.value, "type"), "Exception should have 'type' attribute"
|
||||
assert exc_info.value.type == ProxyErrorTypes.expired_key, (
|
||||
f"Expected expired_key error type, got {exc_info.value.type}"
|
||||
)
|
||||
assert "Expired Key" in str(exc_info.value.message), (
|
||||
f"Exception message should mention 'Expired Key', got: {exc_info.value.message}"
|
||||
)
|
||||
|
||||
# Verify that the param field does NOT leak the full API key (Issue #18731)
|
||||
# The param should be abbreviated like "sk-...XXXX" not the full plaintext key
|
||||
assert (
|
||||
exc_info.value.param is not None
|
||||
), "Exception should have 'param' attribute"
|
||||
assert exc_info.value.param is not None, "Exception should have 'param' attribute"
|
||||
assert exc_info.value.param != api_key, (
|
||||
f"SECURITY: Full API key should NOT be in param field! "
|
||||
f"Got: {exc_info.value.param}, Expected abbreviated format like 'sk-...XXXX'"
|
||||
)
|
||||
assert exc_info.value.param.startswith(
|
||||
"sk-..."
|
||||
), f"Param should be abbreviated to 'sk-...XXXX' format. Got: {exc_info.value.param}"
|
||||
assert exc_info.value.param.startswith("sk-..."), (
|
||||
f"Param should be abbreviated to 'sk-...XXXX' format. Got: {exc_info.value.param}"
|
||||
)
|
||||
|
||||
# Verify that cache deletion was called
|
||||
mock_delete_cache.assert_called_once()
|
||||
call_args = mock_delete_cache.call_args
|
||||
assert (
|
||||
call_args[1]["hashed_token"] == hashed_key
|
||||
), "Cache deletion should be called with the hashed key"
|
||||
assert call_args[1]["hashed_token"] == hashed_key, (
|
||||
"Cache deletion should be called with the hashed key"
|
||||
)
|
||||
finally:
|
||||
# Restore all module-level attributes so subsequent tests are not affected
|
||||
for attr, val in _original_values.items():
|
||||
setattr(_proxy_server_mod, attr, val)
|
||||
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_return_user_api_key_auth_obj_user_spend_and_budget():
|
||||
"""
|
||||
|
|
@ -424,7 +400,7 @@ async def test_return_user_api_key_auth_obj_user_spend_and_budget():
|
|||
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.auth.user_api_key_auth import _return_user_api_key_auth_obj
|
||||
|
||||
|
||||
user_obj = type(
|
||||
"LiteLLM_UserTable",
|
||||
(),
|
||||
|
|
@ -437,7 +413,7 @@ async def test_return_user_api_key_auth_obj_user_spend_and_budget():
|
|||
"user_role": "internal_user",
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
api_key = "sk-test-key"
|
||||
valid_token_dict = {
|
||||
"user_id": "test-user",
|
||||
|
|
@ -445,10 +421,10 @@ async def test_return_user_api_key_auth_obj_user_spend_and_budget():
|
|||
}
|
||||
route = "/chat/completions"
|
||||
start_time = datetime.now()
|
||||
|
||||
|
||||
mock_service_logger = MagicMock()
|
||||
mock_service_logger.async_service_success_hook = AsyncMock()
|
||||
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.auth.user_api_key_auth.user_api_key_service_logger_obj",
|
||||
new=mock_service_logger,
|
||||
|
|
@ -462,7 +438,7 @@ async def test_return_user_api_key_auth_obj_user_spend_and_budget():
|
|||
start_time=start_time,
|
||||
user_role=None,
|
||||
)
|
||||
|
||||
|
||||
assert isinstance(result, UserAPIKeyAuth)
|
||||
assert result.user_spend == 250.0
|
||||
assert result.user_max_budget == 1000.0
|
||||
|
|
@ -494,7 +470,9 @@ def test_proxy_admin_jwt_auth_includes_identity_fields():
|
|||
user_role=LitellmUserRoles.PROXY_ADMIN,
|
||||
user_id="user-abc",
|
||||
team_id="team-123",
|
||||
team_alias=(team_object.team_alias if team_object is not None else None),
|
||||
team_alias=(
|
||||
team_object.team_alias if team_object is not None else None
|
||||
),
|
||||
team_metadata=team_object.metadata if team_object is not None else None,
|
||||
org_id="org-456",
|
||||
end_user_id="end-user-789",
|
||||
|
|
@ -525,7 +503,9 @@ def test_proxy_admin_jwt_auth_handles_no_team_object():
|
|||
user_role=LitellmUserRoles.PROXY_ADMIN,
|
||||
user_id="admin-user",
|
||||
team_id=None,
|
||||
team_alias=(team_object.team_alias if team_object is not None else None),
|
||||
team_alias=(
|
||||
team_object.team_alias if team_object is not None else None
|
||||
),
|
||||
team_metadata=team_object.metadata if team_object is not None else None,
|
||||
org_id=None,
|
||||
end_user_id=None,
|
||||
|
|
@ -554,10 +534,7 @@ class TestJWTOAuth2Coexistence:
|
|||
def test_is_jwt_detects_jwt_tokens(self):
|
||||
"""JWT tokens have 3 dot-separated parts."""
|
||||
assert JWTHandler.is_jwt("header.payload.signature") is True
|
||||
assert (
|
||||
JWTHandler.is_jwt("eyJhbGciOiJSUzI1NiJ9.eyJzdWIiOiJ1c2VyMSJ9.sig123")
|
||||
is True
|
||||
)
|
||||
assert JWTHandler.is_jwt("eyJhbGciOiJSUzI1NiJ9.eyJzdWIiOiJ1c2VyMSJ9.sig123") is True
|
||||
|
||||
def test_is_jwt_rejects_opaque_tokens(self):
|
||||
"""Opaque OAuth2 tokens do not have 3 dot-separated parts."""
|
||||
|
|
@ -566,10 +543,6 @@ class TestJWTOAuth2Coexistence:
|
|||
assert JWTHandler.is_jwt("Bearer token") is False
|
||||
assert JWTHandler.is_jwt("two.parts") is False
|
||||
|
||||
def test_is_jwt_returns_false_for_none(self):
|
||||
"""None token (missing Authorization header) should not be treated as JWT."""
|
||||
assert JWTHandler.is_jwt(None) is False
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_both_enabled_opaque_token_uses_oauth2(self):
|
||||
"""
|
||||
|
|
@ -594,20 +567,13 @@ class TestJWTOAuth2Coexistence:
|
|||
mock_request.headers = {"authorization": f"Bearer {opaque_token}"}
|
||||
mock_request.query_params = {}
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.proxy_server.general_settings", general_settings
|
||||
), patch("litellm.proxy.proxy_server.premium_user", True), patch(
|
||||
"litellm.proxy.proxy_server.master_key", "sk-master"
|
||||
), patch(
|
||||
"litellm.proxy.proxy_server.prisma_client", None
|
||||
), patch(
|
||||
"litellm.proxy.auth.user_api_key_auth.Oauth2Handler.check_oauth2_token",
|
||||
new_callable=AsyncMock,
|
||||
return_value=mock_oauth2_response,
|
||||
) as mock_oauth2, patch(
|
||||
"litellm.proxy.auth.user_api_key_auth.JWTAuthManager.auth_builder",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_jwt_auth:
|
||||
with patch("litellm.proxy.proxy_server.general_settings", general_settings), \
|
||||
patch("litellm.proxy.proxy_server.premium_user", True), \
|
||||
patch("litellm.proxy.proxy_server.master_key", "sk-master"), \
|
||||
patch("litellm.proxy.proxy_server.prisma_client", None), \
|
||||
patch("litellm.proxy.auth.user_api_key_auth.Oauth2Handler.check_oauth2_token", new_callable=AsyncMock, return_value=mock_oauth2_response) as mock_oauth2, \
|
||||
patch("litellm.proxy.auth.user_api_key_auth.JWTAuthManager.auth_builder", new_callable=AsyncMock) as mock_jwt_auth:
|
||||
|
||||
litellm.proxy.proxy_server.jwt_handler.update_environment(
|
||||
prisma_client=None,
|
||||
user_api_key_cache=DualCache(),
|
||||
|
|
@ -658,20 +624,13 @@ class TestJWTOAuth2Coexistence:
|
|||
mock_request.headers = {"authorization": f"Bearer {jwt_token}"}
|
||||
mock_request.query_params = {}
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.proxy_server.general_settings", general_settings
|
||||
), patch("litellm.proxy.proxy_server.premium_user", True), patch(
|
||||
"litellm.proxy.proxy_server.master_key", "sk-master"
|
||||
), patch(
|
||||
"litellm.proxy.proxy_server.prisma_client", None
|
||||
), patch(
|
||||
"litellm.proxy.auth.user_api_key_auth.Oauth2Handler.check_oauth2_token",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_oauth2, patch(
|
||||
"litellm.proxy.auth.user_api_key_auth.JWTAuthManager.auth_builder",
|
||||
new_callable=AsyncMock,
|
||||
return_value=mock_jwt_result,
|
||||
) as mock_jwt_auth:
|
||||
with patch("litellm.proxy.proxy_server.general_settings", general_settings), \
|
||||
patch("litellm.proxy.proxy_server.premium_user", True), \
|
||||
patch("litellm.proxy.proxy_server.master_key", "sk-master"), \
|
||||
patch("litellm.proxy.proxy_server.prisma_client", None), \
|
||||
patch("litellm.proxy.auth.user_api_key_auth.Oauth2Handler.check_oauth2_token", new_callable=AsyncMock) as mock_oauth2, \
|
||||
patch("litellm.proxy.auth.user_api_key_auth.JWTAuthManager.auth_builder", new_callable=AsyncMock, return_value=mock_jwt_result) as mock_jwt_auth:
|
||||
|
||||
litellm.proxy.proxy_server.jwt_handler.update_environment(
|
||||
prisma_client=None,
|
||||
user_api_key_cache=DualCache(),
|
||||
|
|
@ -712,17 +671,12 @@ class TestJWTOAuth2Coexistence:
|
|||
mock_request.headers = {"authorization": f"Bearer {jwt_like_token}"}
|
||||
mock_request.query_params = {}
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.proxy_server.general_settings", general_settings
|
||||
), patch("litellm.proxy.proxy_server.premium_user", True), patch(
|
||||
"litellm.proxy.proxy_server.master_key", "sk-master"
|
||||
), patch(
|
||||
"litellm.proxy.proxy_server.prisma_client", None
|
||||
), patch(
|
||||
"litellm.proxy.auth.user_api_key_auth.Oauth2Handler.check_oauth2_token",
|
||||
new_callable=AsyncMock,
|
||||
return_value=mock_oauth2_response,
|
||||
) as mock_oauth2:
|
||||
with patch("litellm.proxy.proxy_server.general_settings", general_settings), \
|
||||
patch("litellm.proxy.proxy_server.premium_user", True), \
|
||||
patch("litellm.proxy.proxy_server.master_key", "sk-master"), \
|
||||
patch("litellm.proxy.proxy_server.prisma_client", None), \
|
||||
patch("litellm.proxy.auth.user_api_key_auth.Oauth2Handler.check_oauth2_token", new_callable=AsyncMock, return_value=mock_oauth2_response) as mock_oauth2:
|
||||
|
||||
result = await user_api_key_auth(
|
||||
request=mock_request,
|
||||
api_key=f"Bearer {jwt_like_token}",
|
||||
|
|
@ -731,131 +685,3 @@ class TestJWTOAuth2Coexistence:
|
|||
# OAuth2 should handle it since JWT auth is disabled
|
||||
mock_oauth2.assert_called_once_with(token=jwt_like_token)
|
||||
assert result.user_id == "oauth2-user"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_user_api_key_auth_builder_no_blocking_calls():
|
||||
"""
|
||||
_user_api_key_auth_builder must never call any synchronous DualCache method
|
||||
(set_cache, get_cache, batch_get_cache, increment_cache, delete_cache) on
|
||||
the hot auth path — those methods call Redis synchronously and block the
|
||||
event loop. Only async_* variants are allowed.
|
||||
"""
|
||||
from starlette.datastructures import URL
|
||||
from starlette.requests import Request
|
||||
|
||||
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
|
||||
from litellm.proxy.auth.user_api_key_auth import _user_api_key_auth_builder
|
||||
|
||||
_blocking_methods = [
|
||||
"set_cache",
|
||||
"get_cache",
|
||||
"batch_get_cache",
|
||||
"increment_cache",
|
||||
"delete_cache",
|
||||
]
|
||||
|
||||
api_key = "sk-test-no-blocking-cache"
|
||||
valid_token = UserAPIKeyAuth(
|
||||
api_key=api_key,
|
||||
token=api_key,
|
||||
user_role=LitellmUserRoles.INTERNAL_USER,
|
||||
team_id="team-abc",
|
||||
)
|
||||
|
||||
mock_cache = AsyncMock()
|
||||
mock_cache.async_get_cache = AsyncMock(return_value=valid_token)
|
||||
mock_cache.async_set_cache = AsyncMock(return_value=None)
|
||||
# Wire sync methods on the instance as plain MagicMocks (no side_effect) so
|
||||
# calls are recorded but not raised — the function's broad except Exception
|
||||
# would swallow a raised error. We assert not_called() after the run instead.
|
||||
for _m in _blocking_methods:
|
||||
setattr(mock_cache, _m, MagicMock())
|
||||
|
||||
mock_proxy_logging_obj = MagicMock()
|
||||
mock_proxy_logging_obj.internal_usage_cache = MagicMock()
|
||||
mock_proxy_logging_obj.internal_usage_cache.dual_cache = AsyncMock()
|
||||
mock_proxy_logging_obj.internal_usage_cache.dual_cache.async_delete_cache = (
|
||||
AsyncMock()
|
||||
)
|
||||
mock_proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None)
|
||||
|
||||
import litellm.proxy.proxy_server as _proxy_server_mod
|
||||
|
||||
_attrs = {
|
||||
"prisma_client": MagicMock(),
|
||||
"user_api_key_cache": mock_cache,
|
||||
"proxy_logging_obj": mock_proxy_logging_obj,
|
||||
"master_key": "sk-master-key",
|
||||
"general_settings": {},
|
||||
"llm_model_list": [],
|
||||
"llm_router": None,
|
||||
"open_telemetry_logger": None,
|
||||
"model_max_budget_limiter": MagicMock(),
|
||||
"user_custom_auth": None,
|
||||
"jwt_handler": None,
|
||||
"litellm_proxy_admin_name": "admin",
|
||||
}
|
||||
_originals = {k: getattr(_proxy_server_mod, k, None) for k in _attrs}
|
||||
|
||||
try:
|
||||
for k, v in _attrs.items():
|
||||
setattr(_proxy_server_mod, k, v)
|
||||
|
||||
request = Request(scope={"type": "http"})
|
||||
request._url = URL(url="/chat/completions")
|
||||
|
||||
import contextlib
|
||||
|
||||
from litellm.caching.dual_cache import DualCache
|
||||
|
||||
blocking_patches = [
|
||||
patch.object(
|
||||
DualCache,
|
||||
m,
|
||||
MagicMock(
|
||||
side_effect=AssertionError(
|
||||
f"Blocking DualCache.{m}() called on async hot path — use async_{m}() instead"
|
||||
)
|
||||
),
|
||||
)
|
||||
for m in _blocking_methods
|
||||
]
|
||||
|
||||
with contextlib.ExitStack() as stack:
|
||||
for p in blocking_patches:
|
||||
stack.enter_context(p)
|
||||
stack.enter_context(
|
||||
patch(
|
||||
"litellm.proxy.auth.user_api_key_auth.get_key_object",
|
||||
new_callable=AsyncMock,
|
||||
return_value=valid_token,
|
||||
)
|
||||
)
|
||||
stack.enter_context(
|
||||
patch(
|
||||
"litellm.proxy.auth.user_api_key_auth.get_team_object",
|
||||
new_callable=AsyncMock,
|
||||
return_value=None,
|
||||
)
|
||||
)
|
||||
await _user_api_key_auth_builder(
|
||||
request=request,
|
||||
api_key=f"Bearer {api_key}",
|
||||
azure_api_key_header="",
|
||||
anthropic_api_key_header=None,
|
||||
google_ai_studio_api_key_header=None,
|
||||
azure_apim_header=None,
|
||||
request_data={},
|
||||
)
|
||||
|
||||
for _m in _blocking_methods:
|
||||
mock = getattr(mock_cache, _m)
|
||||
assert mock.call_count == 0, (
|
||||
f"Blocking DualCache.{_m}() was called {mock.call_count} time(s) "
|
||||
f"on the async hot path — use async_{_m}() instead"
|
||||
)
|
||||
|
||||
finally:
|
||||
for k, v in _originals.items():
|
||||
setattr(_proxy_server_mod, k, v)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue