Merge pull request #19214 from BerriAI/main

merge main in malformed tool call PR
This commit is contained in:
Sameer Kankute
2026-01-16 17:04:44 +05:30
committed by GitHub
17 changed files with 1321 additions and 94 deletions
+1
View File
@@ -129,6 +129,7 @@ run_grype_scans() {
"CVE-2025-13836" # Python 3.13 HTTP response reading OOM/DoS - no fix available in base image
"CVE-2025-12084" # Python 3.13 xml.dom.minidom quadratic algorithm - no fix available in base image
"CVE-2025-60876" # BusyBox wget HTTP request splitting - no fix available in Chainguard Wolfi base image
"CVE-2026-0861" # Wolfi glibc still flagged even on 2.42-r5; upstream patched build unavailable yet
"CVE-2010-4756" # glibc glob DoS - awaiting patched Wolfi glibc build
"CVE-2019-1010022" # glibc stack guard bypass - awaiting patched Wolfi glibc build
"CVE-2019-1010023" # glibc ldd remap issue - awaiting patched Wolfi glibc build
@@ -30,6 +30,9 @@ general_settings:
# Optional: set how frequently cleanup should run - default is daily
maximum_spend_logs_retention_interval: "1d" # Run cleanup daily
# Optional: set exact time for cleanup (Cron syntax)
maximum_spend_logs_cleanup_cron: "0 4 * * *" # Run at 04:00 AM daily
litellm_settings:
cache: true
cache_params:
@@ -51,6 +54,15 @@ How long logs should be kept before deletion. Supported formats:
How often the cleanup job should run. Uses the same format as above. If not set, cleanup will run every 24 hours if and only if `maximum_spend_logs_retention_period` is set.
#### `maximum_spend_logs_cleanup_cron` (optional)
Schedule the cleanup using standard cron syntax. This takes precedence over `maximum_spend_logs_retention_interval`.
Examples:
- `"0 4 * * *"` Run at 04:00 AM daily
- `"0 0 * * 0"` Run at midnight every Sunday
- `"*/30 * * * *"` Run every 30 minutes
## How it works
### Step 1. Lock Acquisition (Optional with Redis)
+21 -2
View File
@@ -665,11 +665,11 @@ def add_object_type(schema):
if "required" in schema and schema["required"] is None:
schema.pop("required", None)
# Gemini doesn't accept empty properties for object types
# If properties is empty, remove it and the type field
# If properties is empty, remove it but keep type as object
if not properties:
schema.pop("properties", None)
schema.pop("type", None)
schema.pop("required", None)
schema["type"] = "object"
else:
schema["type"] = "object"
for name, value in properties.items():
@@ -776,6 +776,16 @@ def get_vertex_location_from_url(url: str) -> Optional[str]:
return match.group(1) if match else None
def get_vertex_model_id_from_url(url: str) -> Optional[str]:
"""
Get the vertex model id from the url
`https://${LOCATION}-aiplatform.googleapis.com/v1/projects/${PROJECT_ID}/locations/${LOCATION}/publishers/google/models/${MODEL_ID}:streamGenerateContent`
"""
match = re.search(r"/models/([^/:]+)", url)
return match.group(1) if match else None
def replace_project_and_location_in_route(
requested_route: str, vertex_project: str, vertex_location: str
) -> str:
@@ -825,6 +835,15 @@ def construct_target_url(
if "cachedContent" in requested_route:
vertex_version = "v1beta1"
# Check if the requested route starts with a version
# e.g. /v1beta1/publishers/google/models/gemini-3-pro-preview:streamGenerateContent
if requested_route.startswith("/v1/"):
vertex_version = "v1"
requested_route = requested_route.replace("/v1/", "/", 1)
elif requested_route.startswith("/v1beta1/"):
vertex_version = "v1beta1"
requested_route = requested_route.replace("/v1beta1/", "/", 1)
base_requested_route = "{}/projects/{}/locations/{}".format(
vertex_version, vertex_project, vertex_location
)
@@ -68,6 +68,8 @@ def _convert_detail_to_media_resolution_enum(
) -> Optional[Dict[str, str]]:
if detail == "low":
return {"level": "MEDIA_RESOLUTION_LOW"}
elif detail == "medium":
return {"level": "MEDIA_RESOLUTION_MEDIUM"}
elif detail == "high":
return {"level": "MEDIA_RESOLUTION_HIGH"}
return None
@@ -50,8 +50,32 @@ from litellm.types.proxy.guardrails.guardrail_hooks.litellm_content_filter impor
ContentFilterDetection,
PatternDetection,
)
from .patterns import PATTERN_EXTRA_CONFIG, get_compiled_pattern
from .patterns import get_compiled_pattern
MAX_KEYWORD_VALUE_GAP_WORDS = 1
GAP_WORD_TOKENIZER = re.compile(r"\b\w+\b")
WORD_NUMBER_MAP = {
"zero": "0",
"oh": "0",
"one": "1",
"two": "2",
"three": "3",
"four": "4",
"five": "5",
"six": "6",
"seven": "7",
"eight": "8",
"nine": "9",
}
WORD_NUMBER_TOKEN_REGEX = "|".join(WORD_NUMBER_MAP.keys())
WORD_NUMBER_SEQUENCE_PATTERN = re.compile(
rf"(?<![A-Za-z])(?:{WORD_NUMBER_TOKEN_REGEX})(?:[\s\-]+(?:{WORD_NUMBER_TOKEN_REGEX}))+(?![A-Za-z])",
re.IGNORECASE,
)
WORD_NUMBER_TOKEN_FINDER = re.compile(rf"(?:{WORD_NUMBER_TOKEN_REGEX})", re.IGNORECASE)
# Helper data structure for category-based detection
@@ -144,9 +168,9 @@ class ContentFilterGuardrail(CustomGuardrail):
self.image_model = image_model
# Store loaded categories
self.loaded_categories: Dict[str, CategoryConfig] = {}
self.category_keywords: Dict[str, Tuple[str, str, ContentFilterAction]] = (
{}
) # keyword -> (category, severity, action)
self.category_keywords: Dict[
str, Tuple[str, str, ContentFilterAction]
] = {} # keyword -> (category, severity, action)
# Load categories if provided
if categories:
@@ -170,7 +194,7 @@ class ContentFilterGuardrail(CustomGuardrail):
normalized_blocked_words.append(word)
# Compile regex patterns
self.compiled_patterns: List[Tuple[Pattern, str, ContentFilterAction]] = []
self.compiled_patterns: List[Dict[str, Any]] = []
for pattern_config in normalized_patterns:
self._add_pattern(pattern_config)
@@ -323,11 +347,13 @@ class ContentFilterGuardrail(CustomGuardrail):
pattern_config: ContentFilterPattern configuration
"""
try:
extra_config: Dict[str, Any] = {}
if pattern_config.pattern_type == "prebuilt":
if not pattern_config.pattern_name:
raise ValueError("pattern_name is required for prebuilt patterns")
compiled = get_compiled_pattern(pattern_config.pattern_name)
pattern_name = pattern_config.pattern_name
extra_config = PATTERN_EXTRA_CONFIG.get(pattern_name, {}) or {}
elif pattern_config.pattern_type == "regex":
if not pattern_config.pattern:
raise ValueError("pattern is required for regex patterns")
@@ -336,8 +362,20 @@ class ContentFilterGuardrail(CustomGuardrail):
else:
raise ValueError(f"Unknown pattern_type: {pattern_config.pattern_type}")
keyword_regex: Optional[Pattern] = None
if extra_config.get("keyword_pattern"):
keyword_regex = re.compile(
extra_config["keyword_pattern"], re.IGNORECASE
)
self.compiled_patterns.append(
(compiled, pattern_name, pattern_config.action)
{
"regex": compiled,
"pattern_name": pattern_name,
"action": pattern_config.action,
"keyword_regex": keyword_regex,
"allow_word_numbers": bool(extra_config.get("allow_word_numbers")),
}
)
verbose_proxy_logger.debug(
f"Added pattern: {pattern_name} with action {pattern_config.action}"
@@ -395,6 +433,130 @@ class ContentFilterGuardrail(CustomGuardrail):
except Exception as e:
raise Exception(f"Error loading blocked words file {file_path}: {str(e)}")
def _find_pattern_spans(
self, text: str, pattern_entry: Dict[str, Any]
) -> List[Tuple[int, int]]:
"""Return all match spans for a pattern, applying contextual rules if required."""
regex: Pattern = pattern_entry["regex"]
keyword_regex: Optional[Pattern] = pattern_entry.get("keyword_regex")
allow_word_numbers: bool = pattern_entry.get("allow_word_numbers", False)
keyword_matches: Optional[List[re.Match]] = None
if keyword_regex is not None:
keyword_matches = list(keyword_regex.finditer(text))
if not keyword_matches:
return []
match_spans: List[Tuple[int, int]] = []
for match in regex.finditer(text):
if keyword_matches is not None and not self._match_near_keyword(
match.start(), match.end(), keyword_matches, text
):
continue
match_spans.append((match.start(), match.end()))
if allow_word_numbers:
for word_match in WORD_NUMBER_SEQUENCE_PATTERN.finditer(text):
digits = self._convert_word_number_sequence(word_match.group())
if not digits:
continue
if not regex.fullmatch(digits):
continue
if keyword_matches is not None and not self._match_near_keyword(
word_match.start(), word_match.end(), keyword_matches, text
):
continue
match_spans.append((word_match.start(), word_match.end()))
return self._merge_spans(match_spans)
def _match_near_keyword(
self,
value_start: int,
value_end: int,
keyword_matches: List[re.Match],
text: str,
) -> bool:
"""Check if a value is separated from a keyword by an allowed gap."""
for keyword_match in keyword_matches:
keyword_start = keyword_match.start()
keyword_end = keyword_match.end()
if value_start >= keyword_end:
gap_text = text[keyword_end:value_start]
elif keyword_start >= value_end:
gap_text = text[value_end:keyword_start]
else:
return True # overlapping
if self._gap_text_allowed(gap_text):
return True
return False
def _gap_text_allowed(self, gap_text: str) -> bool:
"""Return True if the gap between keyword and value meets word-count rules."""
if not gap_text.strip():
return True
if any(char.isdigit() for char in gap_text):
return False
words = GAP_WORD_TOKENIZER.findall(gap_text)
return len(words) <= MAX_KEYWORD_VALUE_GAP_WORDS
def _merge_spans(self, spans: List[Tuple[int, int]]) -> List[Tuple[int, int]]:
"""Merge overlapping spans to avoid double-masking."""
if not spans:
return []
spans.sort(key=lambda item: item[0])
merged: List[Tuple[int, int]] = [spans[0]]
for start, end in spans[1:]:
last_start, last_end = merged[-1]
if start <= last_end:
merged[-1] = (last_start, max(last_end, end))
else:
merged.append((start, end))
return merged
def _mask_spans(
self, text: str, spans: List[Tuple[int, int]], redaction: str
) -> str:
"""Apply masking for the provided spans using the given redaction tag."""
if not spans:
return text
result_parts: List[str] = []
previous_end = 0
for start, end in spans:
result_parts.append(text[previous_end:start])
result_parts.append(redaction)
previous_end = end
result_parts.append(text[previous_end:])
return "".join(result_parts)
def _convert_word_number_sequence(self, sequence: str) -> Optional[str]:
"""Convert a spelled-out digit sequence (e.g., 'One-Two') into digits."""
tokens = WORD_NUMBER_TOKEN_FINDER.findall(sequence)
if not tokens:
return None
digits: List[str] = []
for token in tokens:
digit = WORD_NUMBER_MAP.get(token.lower())
if digit is None:
return None
digits.append(digit)
return "".join(digits) if digits else None
def _check_patterns(
self, text: str
) -> Optional[Tuple[str, str, ContentFilterAction]]:
@@ -407,10 +569,13 @@ class ContentFilterGuardrail(CustomGuardrail):
Returns:
Tuple of (matched_text, pattern_name, action) if match found, None otherwise
"""
for compiled_pattern, pattern_name, action in self.compiled_patterns:
match = compiled_pattern.search(text)
if match:
matched_text = match.group(0)
for pattern_entry in self.compiled_patterns:
spans = self._find_pattern_spans(text, pattern_entry)
if spans:
start, end = spans[0]
matched_text = text[start:end]
pattern_name = pattern_entry["pattern_name"]
action = pattern_entry["action"]
verbose_proxy_logger.debug(
f"Pattern '{pattern_name}' matched: {matched_text[:20]}..."
)
@@ -582,11 +747,13 @@ class ContentFilterGuardrail(CustomGuardrail):
)
# Check regex patterns - process ALL patterns, not just first match
for compiled_pattern, pattern_name, action in self.compiled_patterns:
match = compiled_pattern.search(text)
if not match:
for pattern_entry in self.compiled_patterns:
spans = self._find_pattern_spans(text, pattern_entry)
if not spans:
continue
pattern_name = pattern_entry["pattern_name"]
action = pattern_entry["action"]
if detections is not None:
# Don't log matched_text to avoid exposing sensitive content (emails, credit cards, etc.)
pattern_detection: PatternDetection = {
@@ -604,11 +771,10 @@ class ContentFilterGuardrail(CustomGuardrail):
detail={"error": error_msg, "pattern": pattern_name},
)
elif action == ContentFilterAction.MASK:
# Replace ALL matches of this pattern with redaction tag
redaction_tag = self.pattern_redaction_format.format(
pattern_name=pattern_name.upper()
)
text = compiled_pattern.sub(redaction_tag, text)
text = self._mask_spans(text, spans, redaction_tag)
verbose_proxy_logger.info(
f"Masked all {pattern_name} matches in content"
)
@@ -924,19 +1090,28 @@ class ContentFilterGuardrail(CustomGuardrail):
if pattern_match:
matched_text, pattern_name, action = pattern_match
if action == ContentFilterAction.BLOCK:
error_msg = f"Content blocked: {pattern_name} pattern detected"
error_msg = (
f"Content blocked: {pattern_name} pattern detected"
)
verbose_proxy_logger.warning(error_msg)
raise HTTPException(
status_code=403,
detail={"error": error_msg, "pattern": pattern_name},
detail={
"error": error_msg,
"pattern": pattern_name,
},
)
# Check blocked words
blocked_word_match = self._check_blocked_words(accumulated_content)
blocked_word_match = self._check_blocked_words(
accumulated_content
)
if blocked_word_match:
keyword, action, description = blocked_word_match
if action == ContentFilterAction.BLOCK:
error_msg = f"Content blocked: keyword '{keyword}' detected"
error_msg = (
f"Content blocked: keyword '{keyword}' detected"
)
if description:
error_msg += f" ({description})"
verbose_proxy_logger.warning(error_msg)
@@ -120,11 +120,11 @@
"description": "Detects URLs (http/https)"
},
{
"name": "passport_us",
"display_name": "Passport (US)",
"pattern": "\\b[0-9]{9}\\b",
"category": "PII Patterns",
"description": "US passport numbers (9 digits)"
"name": "passport_us",
"display_name": "Passport (US)",
"pattern": "\\b[0-9]{9}\\b",
"category": "PII Patterns",
"description": "US passport numbers (9 digits)"
},
{
"name": "passport_uk",
@@ -203,7 +203,6 @@
"category": "Protected Class - Fair Lending",
"description": "Detects race, ethnicity and national origin terms - protected under ECOA and Fair Housing Act"
},
{
"name": "religion",
"display_name": "Religion & Creed (Protected Class)",
@@ -236,7 +235,7 @@
"name": "military_status",
"display_name": "Military Status (Protected Class)",
"pattern": "\\b(veteran|military|armed\\s+forces|army|navy|air\\s+force|marine(s|\\s+corps)?|coast\\s+guard|national\\s+guard|reserve(s|ist)?|active\\s+duty|deployment|deployed|enlisted|commissioned|honorable\\s+discharge|dishonorable\\s+discharge|VA\\s+benefits|GI\\s+bill|military\\s+service|service\\s+member|servicemember|SCRA|MLA|military\\s+lending)\\b",
"category": "Protected Class - Fair Lending",
"category": "Protected Class - Fair Lending",
"description": "Detects military status terms - protected under SCRA and MLA"
},
{
@@ -245,7 +244,7 @@
"pattern": "\\b(welfare|public\\s+assistance|food\\s+stamps|SNAP|WIC|TANF|medicaid|section\\s+8|housing\\s+voucher|subsidized\\s+housing|public\\s+housing|government\\s+benefits|social\\s+services|unemployment\\s+(benefits|insurance)|UI\\s+benefits|EBT|benefit\\s+recipient)\\b",
"category": "Protected Class - Fair Lending",
"description": "Detects public assistance terms - protected under ECOA"
} ,
},
{
"name": "weapons_firearms",
"display_name": "Weapons & Firearms",
@@ -313,10 +312,12 @@
{
"name": "nl_bsn_contextual",
"display_name": "BSN (Dutch Citizen Service Number)",
"pattern": "\\b(?:BSN|B\\.S\\.N\\.|burgerservicenummer|burger\\s*service\\s*nummer|sofi\\s*nummer|sofinummer|persoonsnummer|identificatienummer|citizen\\s*service\\s*number)[:\\s]*[0-9]{9}\\b|\\b[0-9]{9}\\b(?=\\s*(?:BSN|burgerservicenummer|sofinummer))",
"pattern": "\\b[0-9]{9}\\b",
"category": "PII Patterns",
"action": "MASK",
"description": "Detects Dutch BSN numbers with contextual keywords"
"description": "Detects Dutch BSN numbers with contextual keywords",
"keyword_pattern": "(?:\\b(?:BSN|B\\.S\\.N\\.|burgerservicenummer|burger\\s*service\\s*nummer|sofi\\s*nummer|sofinummer|persoonsnummer|identificatienummer|citizen\\s*service\\s*number)\\b|8\\s*5\\s*\\|\\\\\\|)",
"allow_word_numbers": true
},
{
"name": "br_cpf",
@@ -369,5 +370,3 @@
}
]
}
@@ -9,7 +9,7 @@ import json
import os
import re
from enum import Enum
from typing import Dict, List, Pattern
from typing import Any, Dict, List, Pattern
def _load_patterns_from_json() -> Dict:
@@ -41,6 +41,26 @@ PREBUILT_PATTERNS: Dict[str, str] = {
}
# Capture any extra configuration declared per pattern (e.g., contextual keywords)
KNOWN_PATTERN_KEYS = {
"name",
"display_name",
"pattern",
"category",
"action",
"description",
}
PATTERN_EXTRA_CONFIG: Dict[str, Dict[str, Any]] = {}
for pattern_data in _PATTERNS_DATA["patterns"]:
extra_config = {
key: value
for key, value in pattern_data.items()
if key not in KNOWN_PATTERN_KEYS
}
PATTERN_EXTRA_CONFIG[pattern_data["name"]] = extra_config
def get_compiled_pattern(pattern_name: str) -> Pattern:
"""
Get a compiled regex pattern by name.
+21 -14
View File
@@ -114,25 +114,25 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
) -> Optional[str]:
"""
Get priority from user_api_key_dict.
Checks team metadata first (takes precedence), then falls back to key metadata.
Args:
user_api_key_dict: User authentication info
Returns:
Priority string if found, None otherwise
"""
priority: Optional[str] = None
# Check team metadata first (takes precedence)
if user_api_key_dict.team_metadata is not None:
priority = user_api_key_dict.team_metadata.get("priority", None)
# Fall back to key metadata
if priority is None:
priority = user_api_key_dict.metadata.get("priority", None)
return priority
def _normalize_priority_weights(
@@ -299,10 +299,13 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
"""
descriptors: List[RateLimitDescriptor] = []
if litellm.priority_reservation is None:
return descriptors
# Get model group info
model_group_info: Optional[ModelGroupInfo] = (
self.llm_router.get_model_group_info(model_group=model)
)
model_group_info: Optional[
ModelGroupInfo
] = self.llm_router.get_model_group_info(model_group=model)
if model_group_info is None:
return descriptors
@@ -577,9 +580,9 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
)
# Get model configuration
model_group_info: Optional[ModelGroupInfo] = (
self.llm_router.get_model_group_info(model_group=model)
)
model_group_info: Optional[
ModelGroupInfo
] = self.llm_router.get_model_group_info(model_group=model)
if model_group_info is None:
verbose_proxy_logger.debug(
f"No model group info for {model}, allowing request"
@@ -703,7 +706,9 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
# Get priority from user_api_key_auth_metadata in standard_logging_metadata
# This is where user_api_key_dict.metadata is stored during pre-call
user_api_key_auth_metadata = standard_logging_metadata.get("user_api_key_auth_metadata") or {}
user_api_key_auth_metadata = (
standard_logging_metadata.get("user_api_key_auth_metadata") or {}
)
key_priority: Optional[str] = user_api_key_auth_metadata.get("priority")
# Get total tokens from response
@@ -775,7 +780,9 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
# Only log 'priority' if it's known safe; otherwise, redact.
SAFE_PRIORITIES = {"low", "medium", "high", "default"}
logged_priority = key_priority if key_priority in SAFE_PRIORITIES else "REDACTED"
logged_priority = (
key_priority if key_priority in SAFE_PRIORITIES else "REDACTED"
)
verbose_proxy_logger.debug(
f"[Dynamic Rate Limiter] Incremented tokens by {total_tokens} for "
f"model={model_group}, priority={logged_priority}"
@@ -1554,6 +1554,7 @@ async def _base_vertex_proxy_route(
from litellm.llms.vertex_ai.common_utils import (
construct_target_url,
get_vertex_location_from_url,
get_vertex_model_id_from_url,
get_vertex_project_id_from_url,
)
@@ -1583,6 +1584,25 @@ async def _base_vertex_proxy_route(
vertex_location=vertex_location,
)
if vertex_project is None or vertex_location is None:
# Check if model is in router config
model_id = get_vertex_model_id_from_url(endpoint)
if model_id:
from litellm.proxy.proxy_server import llm_router
if llm_router:
try:
# Use the dedicated pass-through deployment selection method to automatically filter use_in_pass_through=True
deployment = llm_router.get_available_deployment_for_pass_through(model=model_id)
if deployment:
litellm_params = deployment.get("litellm_params", {})
vertex_project = litellm_params.get("vertex_project")
vertex_location = litellm_params.get("vertex_location")
except Exception as e:
verbose_proxy_logger.debug(
f"Error getting available deployment for model {model_id}: {e}"
)
vertex_credentials = passthrough_endpoint_router.get_vertex_credentials(
project_id=vertex_project,
location=vertex_location,
+79 -36
View File
@@ -3253,20 +3253,22 @@ class ProxyConfig:
) -> Optional[dict]:
"""
Get router_settings in priority order: Key > Team > Global
Returns:
dict: Combined router_settings, or None if no settings found
"""
if prisma_client is None:
return None
import json
import yaml
# 1. Try key-level router_settings
if user_api_key_dict is not None:
# Check if router_settings is available on the key object
key_router_settings_value = getattr(user_api_key_dict, "router_settings", None)
key_router_settings_value = getattr(
user_api_key_dict, "router_settings", None
)
if key_router_settings_value is not None:
key_router_settings = None
if isinstance(key_router_settings_value, str):
@@ -3279,11 +3281,15 @@ class ProxyConfig:
pass
elif isinstance(key_router_settings_value, dict):
key_router_settings = key_router_settings_value
# If key has router_settings (non-empty dict), use it
if key_router_settings is not None and isinstance(key_router_settings, dict) and key_router_settings:
if (
key_router_settings is not None
and isinstance(key_router_settings, dict)
and key_router_settings
):
return key_router_settings
# 2. Try team-level router_settings
if user_api_key_dict is not None and user_api_key_dict.team_id is not None:
try:
@@ -3291,37 +3297,51 @@ class ProxyConfig:
where={"team_id": user_api_key_dict.team_id}
)
if team_obj is not None:
team_router_settings_value = getattr(team_obj, "router_settings", None)
team_router_settings_value = getattr(
team_obj, "router_settings", None
)
if team_router_settings_value is not None:
team_router_settings = None
if isinstance(team_router_settings_value, str):
try:
team_router_settings = yaml.safe_load(team_router_settings_value)
team_router_settings = yaml.safe_load(
team_router_settings_value
)
except (yaml.YAMLError, json.JSONDecodeError):
try:
team_router_settings = json.loads(team_router_settings_value)
team_router_settings = json.loads(
team_router_settings_value
)
except json.JSONDecodeError:
pass
elif isinstance(team_router_settings_value, dict):
team_router_settings = team_router_settings_value
# If team has router_settings (non-empty dict), use it
if team_router_settings is not None and isinstance(team_router_settings, dict) and team_router_settings:
if (
team_router_settings is not None
and isinstance(team_router_settings, dict)
and team_router_settings
):
return team_router_settings
except Exception:
# If team lookup fails, continue to global settings
pass
# 3. Try global router_settings
try:
db_router_settings = await prisma_client.db.litellm_config.find_first(
where={"param_name": "router_settings"}
)
if db_router_settings is not None and isinstance(db_router_settings.param_value, dict) and db_router_settings.param_value:
if (
db_router_settings is not None
and isinstance(db_router_settings.param_value, dict)
and db_router_settings.param_value
):
return db_router_settings.param_value
except Exception:
pass
return None
async def _add_router_settings_from_db_config(
@@ -4688,27 +4708,48 @@ class ProxyStartupEvent:
### SPEND LOG CLEANUP ###
if general_settings.get("maximum_spend_logs_retention_period") is not None:
spend_log_cleanup = SpendLogCleanup()
# Get the interval from config or default to 1 day
retention_interval = general_settings.get(
"maximum_spend_logs_retention_interval", "1d"
)
try:
interval_seconds = duration_in_seconds(retention_interval)
scheduler.add_job(
spend_log_cleanup.cleanup_old_spend_logs,
"interval",
seconds=interval_seconds
+ random.randint(0, 60), # Add small random offset
# REMOVED jitter parameter - major cause of memory leak
args=[prisma_client],
id="spend_log_cleanup_job",
replace_existing=True,
misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME,
)
except ValueError:
verbose_proxy_logger.error(
"Invalid maximum_spend_logs_retention_interval value"
cleanup_cron = general_settings.get("maximum_spend_logs_cleanup_cron")
if cleanup_cron:
from apscheduler.triggers.cron import CronTrigger
try:
cron_trigger = CronTrigger.from_crontab(cleanup_cron)
scheduler.add_job(
spend_log_cleanup.cleanup_old_spend_logs,
cron_trigger,
args=[prisma_client],
id="spend_log_cleanup_job",
replace_existing=True,
misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME,
)
verbose_proxy_logger.info(
f"Spend log cleanup scheduled with cron: {cleanup_cron}"
)
except ValueError:
verbose_proxy_logger.error(
f"Invalid maximum_spend_logs_cleanup_cron value: {cleanup_cron}"
)
else:
# Interval-based scheduling (existing behavior)
retention_interval = general_settings.get(
"maximum_spend_logs_retention_interval", "1d"
)
try:
interval_seconds = duration_in_seconds(retention_interval)
scheduler.add_job(
spend_log_cleanup.cleanup_old_spend_logs,
"interval",
seconds=interval_seconds + random.randint(0, 60),
args=[prisma_client],
id="spend_log_cleanup_job",
replace_existing=True,
misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME,
)
except ValueError:
verbose_proxy_logger.error(
"Invalid maximum_spend_logs_retention_interval value"
)
### CHECK BATCH COST ###
if llm_router is not None:
try:
@@ -9922,7 +9963,9 @@ async def get_config(): # noqa: PLR0915
_success_callbacks = normalize_callback(_success_callbacks)
_failure_callbacks = normalize_callback(_failure_callbacks)
_success_and_failure_callbacks = normalize_callback(_success_and_failure_callbacks)
_success_and_failure_callbacks = normalize_callback(
_success_and_failure_callbacks
)
_data_to_return = []
"""
+339
View File
@@ -8032,6 +8032,154 @@ class Router:
)
raise e
async def async_get_available_deployment_for_pass_through(
self,
model: str,
request_kwargs: Dict,
messages: Optional[List[Dict[str, str]]] = None,
input: Optional[Union[str, List]] = None,
specific_deployment: Optional[bool] = False,
):
"""
Async version of get_available_deployment_for_pass_through
Only returns deployments configured with use_in_pass_through=True
"""
try:
parent_otel_span = _get_parent_otel_span_from_kwargs(request_kwargs)
# 1. Execute pre-routing hook
pre_routing_hook_response = await self.async_pre_routing_hook(
model=model,
request_kwargs=request_kwargs,
messages=messages,
input=input,
specific_deployment=specific_deployment,
)
if pre_routing_hook_response is not None:
model = pre_routing_hook_response.model
messages = pre_routing_hook_response.messages
# 2. Get healthy deployments
healthy_deployments = await self.async_get_healthy_deployments(
model=model,
request_kwargs=request_kwargs,
messages=messages,
input=input,
specific_deployment=specific_deployment,
parent_otel_span=parent_otel_span,
)
# 3. If specific deployment returned, verify if it supports pass-through
if isinstance(healthy_deployments, dict):
litellm_params = healthy_deployments.get("litellm_params", {})
if litellm_params.get("use_in_pass_through"):
return healthy_deployments
else:
raise litellm.BadRequestError(
message=f"Deployment {healthy_deployments.get('model_info', {}).get('id')} does not support pass-through endpoint (use_in_pass_through=False)",
model=model,
llm_provider="",
)
# 4. Filter deployments that support pass-through
pass_through_deployments = self._filter_pass_through_deployments(
healthy_deployments=healthy_deployments
)
if len(pass_through_deployments) == 0:
raise litellm.BadRequestError(
message=f"Model {model} has no deployments configured with use_in_pass_through=True. Please add use_in_pass_through: true to the deployment configuration",
model=model,
llm_provider="",
)
# 5. Apply load balancing strategy
start_time = time.perf_counter()
if (
self.routing_strategy == "usage-based-routing-v2"
and self.lowesttpm_logger_v2 is not None
):
deployment = (
await self.lowesttpm_logger_v2.async_get_available_deployments(
model_group=model,
healthy_deployments=pass_through_deployments, # type: ignore
messages=messages,
input=input,
)
)
elif (
self.routing_strategy == "latency-based-routing"
and self.lowestlatency_logger is not None
):
deployment = (
await self.lowestlatency_logger.async_get_available_deployments(
model_group=model,
healthy_deployments=pass_through_deployments, # type: ignore
messages=messages,
input=input,
request_kwargs=request_kwargs,
)
)
elif self.routing_strategy == "simple-shuffle":
return simple_shuffle(
llm_router_instance=self,
healthy_deployments=pass_through_deployments,
model=model,
)
elif (
self.routing_strategy == "least-busy"
and self.leastbusy_logger is not None
):
deployment = (
await self.leastbusy_logger.async_get_available_deployments(
model_group=model,
healthy_deployments=pass_through_deployments, # type: ignore
)
)
else:
deployment = None
if deployment is None:
exception = await async_raise_no_deployment_exception(
litellm_router_instance=self,
model=model,
parent_otel_span=parent_otel_span,
)
raise exception
verbose_router_logger.info(
f"async_get_available_deployment_for_pass_through model: {model}, selected deployment: {self.print_deployment(deployment)}"
)
end_time = time.perf_counter()
_duration = end_time - start_time
asyncio.create_task(
self.service_logger_obj.async_service_success_hook(
service=ServiceTypes.ROUTER,
duration=_duration,
call_type="<routing_strategy>.async_get_available_deployments",
parent_otel_span=parent_otel_span,
start_time=start_time,
end_time=end_time,
)
)
return deployment
except Exception as e:
traceback_exception = traceback.format_exc()
if request_kwargs is not None:
logging_obj = request_kwargs.get("litellm_logging_obj", None)
if logging_obj is not None:
threading.Thread(
target=logging_obj.failure_handler,
args=(e, traceback_exception),
).start()
asyncio.create_task(
logging_obj.async_failure_handler(e, traceback_exception) # type: ignore
)
raise e
async def async_pre_routing_hook(
self,
model: str,
@@ -8184,6 +8332,169 @@ class Router:
)
return deployment
def get_available_deployment_for_pass_through(
self,
model: str,
messages: Optional[List[Dict[str, str]]] = None,
input: Optional[Union[str, List]] = None,
specific_deployment: Optional[bool] = False,
request_kwargs: Optional[Dict] = None,
):
"""
Returns deployments available for pass-through endpoints (based on load balancing strategy)
Similar to get_available_deployment, but only returns deployments with use_in_pass_through=True
Args:
model: Model name
messages: Optional list of messages
input: Optional input data
specific_deployment: Whether to find a specific deployment
request_kwargs: Optional request parameters
Returns:
Dict: Selected deployment configuration
Raises:
BadRequestError: If no deployment is configured with use_in_pass_through=True
RouterRateLimitError: If no pass-through deployments are available
"""
# 1. Perform common checks to get healthy deployments list
model, healthy_deployments = self._common_checks_available_deployment(
model=model,
messages=messages,
input=input,
specific_deployment=specific_deployment,
)
# 2. If the returned is a specific deployment (Dict), verify and return directly
if isinstance(healthy_deployments, dict):
litellm_params = healthy_deployments.get("litellm_params", {})
if litellm_params.get("use_in_pass_through"):
return healthy_deployments
else:
# Specific deployment does not support pass-through
raise litellm.BadRequestError(
message=f"Deployment {healthy_deployments.get('model_info', {}).get('id')} does not support pass-through endpoint (use_in_pass_through=False)",
model=model,
llm_provider="",
)
# 3. Filter deployments that support pass-through
pass_through_deployments = self._filter_pass_through_deployments(
healthy_deployments=healthy_deployments
)
if len(pass_through_deployments) == 0:
# No deployments support pass-through
raise litellm.BadRequestError(
message=f"Model {model} has no deployment configured with use_in_pass_through=True. Please add use_in_pass_through: true in the deployment configuration",
model=model,
llm_provider="",
)
# 4. Apply cooldown filtering
parent_otel_span: Optional[Span] = _get_parent_otel_span_from_kwargs(
request_kwargs
)
cooldown_deployments = _get_cooldown_deployments(
litellm_router_instance=self, parent_otel_span=parent_otel_span
)
pass_through_deployments = self._filter_cooldown_deployments(
healthy_deployments=pass_through_deployments,
cooldown_deployments=cooldown_deployments,
)
# 5. Apply pre-call checks (if enabled)
if self.enable_pre_call_checks and messages is not None:
pass_through_deployments = self._pre_call_checks(
model=model,
healthy_deployments=pass_through_deployments,
messages=messages,
request_kwargs=request_kwargs,
)
if len(pass_through_deployments) == 0:
model_ids = self.get_model_ids(model_name=model)
_cooldown_time = self.cooldown_cache.get_min_cooldown(
model_ids=model_ids, parent_otel_span=parent_otel_span
)
_cooldown_list = _get_cooldown_deployments(
litellm_router_instance=self, parent_otel_span=parent_otel_span
)
raise RouterRateLimitError(
model=model,
cooldown_time=_cooldown_time,
enable_pre_call_checks=self.enable_pre_call_checks,
cooldown_list=_cooldown_list,
)
# 6. Apply load balancing strategy
if self.routing_strategy == "least-busy" and self.leastbusy_logger is not None:
deployment = self.leastbusy_logger.get_available_deployments(
model_group=model, healthy_deployments=pass_through_deployments # type: ignore
)
elif self.routing_strategy == "simple-shuffle":
return simple_shuffle(
llm_router_instance=self,
healthy_deployments=pass_through_deployments,
model=model,
)
elif (
self.routing_strategy == "latency-based-routing"
and self.lowestlatency_logger is not None
):
deployment = self.lowestlatency_logger.get_available_deployments(
model_group=model,
healthy_deployments=pass_through_deployments, # type: ignore
request_kwargs=request_kwargs,
)
elif (
self.routing_strategy == "usage-based-routing"
and self.lowesttpm_logger is not None
):
deployment = self.lowesttpm_logger.get_available_deployments(
model_group=model,
healthy_deployments=pass_through_deployments, # type: ignore
messages=messages,
input=input,
)
elif (
self.routing_strategy == "usage-based-routing-v2"
and self.lowesttpm_logger_v2 is not None
):
deployment = self.lowesttpm_logger_v2.get_available_deployments(
model_group=model,
healthy_deployments=pass_through_deployments, # type: ignore
messages=messages,
input=input,
)
else:
deployment = None
if deployment is None:
verbose_router_logger.info(
f"get_available_deployment_for_pass_through model: {model}, no available deployments"
)
model_ids = self.get_model_ids(model_name=model)
_cooldown_time = self.cooldown_cache.get_min_cooldown(
model_ids=model_ids, parent_otel_span=parent_otel_span
)
_cooldown_list = _get_cooldown_deployments(
litellm_router_instance=self, parent_otel_span=parent_otel_span
)
raise RouterRateLimitError(
model=model,
cooldown_time=_cooldown_time,
enable_pre_call_checks=self.enable_pre_call_checks,
cooldown_list=_cooldown_list,
)
verbose_router_logger.info(
f"get_available_deployment_for_pass_through model: {model}, selected deployment: {self.print_deployment(deployment)}"
)
return deployment
def _filter_cooldown_deployments(
self, healthy_deployments: List[Dict], cooldown_deployments: List[str]
) -> List[Dict]:
@@ -8206,6 +8517,34 @@ class Router:
if deployment["model_info"]["id"] not in cooldown_set
]
def _filter_pass_through_deployments(
self, healthy_deployments: List[Dict]
) -> List[Dict]:
"""
Filter out deployments configured with use_in_pass_through=True
Args:
healthy_deployments: List of healthy deployments
Returns:
List[Dict]: Only includes a list of deployments that support pass-through
"""
verbose_router_logger.debug(
f"Filter pass-through deployments from {len(healthy_deployments)} healthy deployments"
)
pass_through_deployments = [
deployment
for deployment in healthy_deployments
if deployment.get("litellm_params", {}).get("use_in_pass_through", False)
]
verbose_router_logger.debug(
f"Found {len(pass_through_deployments)} deployments with pass-through enabled"
)
return pass_through_deployments
def _track_deployment_metrics(
self, deployment, parent_otel_span: Optional[Span], response=None
):
@@ -592,3 +592,205 @@ async def test_weighted_selection_router_async(rpm_list, tpm_list):
except Exception as e:
traceback.print_exc()
pytest.fail(f"Error occurred: {e}")
def test_get_available_deployment_for_pass_through():
"""
Test get_available_deployment_for_pass_through function
- Tests that only deployments with use_in_pass_through=True are returned
- Tests that BadRequestError is raised when no pass-through deployments exist
"""
try:
litellm.set_verbose = False
model_list = [
{
"model_name": "gpt-3.5-turbo",
"litellm_params": {
"model": "gpt-3.5-turbo",
"api_key": os.getenv("OPENAI_API_KEY"),
"use_in_pass_through": True,
},
},
{
"model_name": "gpt-3.5-turbo",
"litellm_params": {
"model": "azure/gpt-4.1-mini",
"api_key": os.getenv("AZURE_API_KEY"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_version": os.getenv("AZURE_API_VERSION"),
"use_in_pass_through": False,
},
},
]
router = Router(
model_list=model_list,
)
# Test that only pass-through deployment is returned
selected_model = router.get_available_deployment_for_pass_through(
"gpt-3.5-turbo"
)
assert selected_model["litellm_params"]["model"] == "gpt-3.5-turbo"
assert selected_model["litellm_params"]["use_in_pass_through"] is True
router.reset()
except Exception as e:
traceback.print_exc()
pytest.fail(f"Error occurred: {e}")
def test_get_available_deployment_for_pass_through_no_deployments():
"""
Test get_available_deployment_for_pass_through raises BadRequestError
when no deployments have use_in_pass_through=True
"""
try:
litellm.set_verbose = False
model_list = [
{
"model_name": "gpt-3.5-turbo",
"litellm_params": {
"model": "gpt-3.5-turbo",
"api_key": os.getenv("OPENAI_API_KEY"),
"use_in_pass_through": False,
},
},
{
"model_name": "gpt-3.5-turbo",
"litellm_params": {
"model": "azure/gpt-4.1-mini",
"api_key": os.getenv("AZURE_API_KEY"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_version": os.getenv("AZURE_API_VERSION"),
"use_in_pass_through": False,
},
},
]
router = Router(
model_list=model_list,
)
# Test that BadRequestError is raised when no pass-through deployments exist
try:
router.get_available_deployment_for_pass_through("gpt-3.5-turbo")
pytest.fail(
"Expected BadRequestError when no pass-through deployments exist"
)
except litellm.BadRequestError as e:
assert "use_in_pass_through=True" in str(e)
router.reset()
except Exception as e:
if isinstance(e, litellm.BadRequestError):
pass # Expected error
else:
traceback.print_exc()
pytest.fail(f"Error occurred: {e}")
@pytest.mark.asyncio
async def test_async_get_available_deployment_for_pass_through():
"""
Test async_get_available_deployment_for_pass_through function
- Tests that only deployments with use_in_pass_through=True are returned
- Tests async version works correctly
"""
try:
litellm.set_verbose = False
model_list = [
{
"model_name": "gpt-3.5-turbo",
"litellm_params": {
"model": "gpt-3.5-turbo",
"api_key": os.getenv("OPENAI_API_KEY"),
"use_in_pass_through": True,
},
},
{
"model_name": "gpt-3.5-turbo",
"litellm_params": {
"model": "azure/gpt-4.1-mini",
"api_key": os.getenv("AZURE_API_KEY"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_version": os.getenv("AZURE_API_VERSION"),
"use_in_pass_through": False,
},
},
]
router = Router(
model_list=model_list,
)
# Test that only pass-through deployment is returned
selected_model = await router.async_get_available_deployment_for_pass_through(
model="gpt-3.5-turbo", request_kwargs={}
)
assert selected_model["litellm_params"]["model"] == "gpt-3.5-turbo"
assert selected_model["litellm_params"]["use_in_pass_through"] is True
router.reset()
except Exception as e:
traceback.print_exc()
pytest.fail(f"Error occurred: {e}")
def test_filter_pass_through_deployments():
"""
Test _filter_pass_through_deployments function
- Tests that it correctly filters deployments with use_in_pass_through=True
"""
try:
litellm.set_verbose = False
model_list = [
{
"model_name": "gpt-3.5-turbo",
"litellm_params": {
"model": "gpt-3.5-turbo",
"api_key": os.getenv("OPENAI_API_KEY"),
"use_in_pass_through": True,
},
},
{
"model_name": "gpt-3.5-turbo",
"litellm_params": {
"model": "azure/gpt-4.1-mini",
"api_key": os.getenv("AZURE_API_KEY"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_version": os.getenv("AZURE_API_VERSION"),
"use_in_pass_through": False,
},
},
{
"model_name": "gpt-3.5-turbo",
"litellm_params": {
"model": "azure/gpt-35-turbo",
"api_key": os.getenv("AZURE_API_KEY"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_version": os.getenv("AZURE_API_VERSION"),
"use_in_pass_through": True,
},
},
]
router = Router(
model_list=model_list,
)
# Get all healthy deployments
healthy_deployments = router.get_model_list()
# Filter pass-through deployments
pass_through_deployments = router._filter_pass_through_deployments(
healthy_deployments
)
# Should only have 2 deployments with use_in_pass_through=True
assert len(pass_through_deployments) == 2
# Verify all returned deployments have use_in_pass_through=True
for deployment in pass_through_deployments:
assert deployment["litellm_params"]["use_in_pass_through"] is True
router.reset()
except Exception as e:
traceback.print_exc()
pytest.fail(f"Error occurred: {e}")
@@ -0,0 +1,16 @@
"""Test for Gemini schema handling with empty properties."""
import os
import sys
sys.path.insert(0, os.path.abspath("../../../.."))
from litellm.llms.vertex_ai.common_utils import add_object_type
def test_add_object_type_empty_properties_keeps_type():
"""Gemini requires type: object even when properties is empty."""
schema = {"properties": {}, "type": "object"}
add_object_type(schema)
assert schema.get("type") == "object"
assert "properties" not in schema
@@ -1,7 +1,6 @@
import os
import sys
from typing import Any, Dict
from unittest.mock import MagicMock, call, patch
from unittest.mock import patch
import pytest
@@ -11,7 +10,6 @@ sys.path.insert(
0, os.path.abspath("../../..")
) # Adds the parent directory to the system path
import litellm
from litellm.llms.vertex_ai.common_utils import (
_get_vertex_url,
convert_anyof_null_to_nullable,
@@ -798,9 +796,54 @@ def test_fix_enum_empty_strings():
assert "mobile" in enum_values
assert "tablet" in enum_values
# 3. Other properties preserved
assert input_schema["properties"]["user_agent_type"]["type"] == "string"
assert input_schema["properties"]["user_agent_type"]["description"] == "Device type for user agent"
def test_get_vertex_model_id_from_url():
"""Test get_vertex_model_id_from_url with various URLs"""
from litellm.llms.vertex_ai.common_utils import get_vertex_model_id_from_url
# Test with valid URL
url = "https://us-central1-aiplatform.googleapis.com/v1/projects/test-project/locations/us-central1/publishers/google/models/gemini-pro:streamGenerateContent"
model_id = get_vertex_model_id_from_url(url)
assert model_id == "gemini-pro"
# Test with invalid URL
url = "https://invalid-url.com"
model_id = get_vertex_model_id_from_url(url)
assert model_id is None
def test_construct_target_url_with_version_prefix():
"""Test construct_target_url with version prefixes"""
from litellm.llms.vertex_ai.common_utils import construct_target_url
# Test with /v1/ prefix
url = "/v1/publishers/google/models/gemini-pro:streamGenerateContent"
vertex_project = "test-project"
vertex_location = "us-central1"
base_url = "https://us-central1-aiplatform.googleapis.com"
target_url = construct_target_url(
base_url=base_url,
requested_route=url,
vertex_project=vertex_project,
vertex_location=vertex_location,
)
expected_url = "https://us-central1-aiplatform.googleapis.com/v1/projects/test-project/locations/us-central1/publishers/google/models/gemini-pro:streamGenerateContent"
assert str(target_url) == expected_url
# Test with /v1beta1/ prefix
url = "/v1beta1/publishers/google/models/gemini-pro:streamGenerateContent"
target_url = construct_target_url(
base_url=base_url,
requested_route=url,
vertex_project=vertex_project,
vertex_location=vertex_location,
)
expected_url = "https://us-central1-aiplatform.googleapis.com/v1beta1/projects/test-project/locations/us-central1/publishers/google/models/gemini-pro:streamGenerateContent"
assert str(target_url) == expected_url
def test_fix_enum_types():
@@ -862,7 +905,7 @@ def test_fix_enum_types():
"truncateMode": {
"enum": ["auto", "none", "start", "end"], # Kept - string type
"type": "string",
"description": "How to truncate content"
"description": "How to truncate content",
},
"maxLength": { # enum removed
"type": "integer",
@@ -1254,8 +1297,8 @@ def test_build_vertex_schema_empty_properties():
# Verify empty properties was removed
assert "properties" not in go_back_schema, "Empty properties should be removed"
# Verify type was also removed (since object without properties is invalid in Gemini)
assert "type" not in go_back_schema, "Type should be removed when properties is empty"
# Verify type is kept as object (Gemini requires type: object even without properties)
assert go_back_schema.get("type") == "object", "Type should be kept as object when properties is empty"
# Verify required was also removed
assert "required" not in go_back_schema, "Required should be removed when properties is empty"
@@ -14,7 +14,6 @@ sys.path.insert(
from fastapi import HTTPException
import litellm
from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.content_filter import (
ContentFilterGuardrail,
)
@@ -0,0 +1,222 @@
import pytest
from unittest.mock import MagicMock, AsyncMock, patch
from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import _base_vertex_proxy_route
from litellm.types.router import DeploymentTypedDict
@pytest.mark.asyncio
async def test_vertex_passthrough_load_balancing():
"""
Test that _base_vertex_proxy_route uses llm_router.get_available_deployment_for_pass_through
instead of get_model_list to ensure load balancing works with pass-through filtering.
"""
# Setup mocks
mock_request = MagicMock()
mock_response = MagicMock()
mock_handler = MagicMock()
# Mock the router
mock_router = MagicMock()
mock_deployment = {
"litellm_params": {
"model": "vertex_ai/gemini-pro",
"vertex_project": "test-project-lb",
"vertex_location": "us-central1-lb",
"use_in_pass_through": True
}
}
mock_router.get_available_deployment_for_pass_through.return_value = mock_deployment
# Mock get_vertex_model_id_from_url to return a model ID
with patch("litellm.llms.vertex_ai.common_utils.get_vertex_model_id_from_url", return_value="gemini-pro"), \
patch("litellm.proxy.proxy_server.llm_router", mock_router), \
patch("litellm.llms.vertex_ai.common_utils.get_vertex_project_id_from_url", return_value=None), \
patch("litellm.llms.vertex_ai.common_utils.get_vertex_location_from_url", return_value=None), \
patch("litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.passthrough_endpoint_router") as mock_pt_router, \
patch("litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints._prepare_vertex_auth_headers", new_callable=AsyncMock) as mock_prep_headers, \
patch("litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.create_pass_through_route") as mock_create_route, \
patch("litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.user_api_key_auth", new_callable=AsyncMock) as mock_auth:
# Setup additional mocks to avoid side effects
mock_pt_router.get_vertex_credentials.return_value = MagicMock()
mock_prep_headers.return_value = ({}, "https://test.url", False, "test-project-lb", "us-central1-lb")
mock_endpoint_func = AsyncMock()
mock_create_route.return_value = mock_endpoint_func
mock_auth.return_value = {}
# Execute
await _base_vertex_proxy_route(
endpoint="https://us-central1-aiplatform.googleapis.com/v1/projects/my-project/locations/us-central1/publishers/google/models/gemini-pro:streamGenerateContent",
request=mock_request,
fastapi_response=mock_response,
get_vertex_pass_through_handler=mock_handler
)
# Verify
# 1. Check that get_available_deployment_for_pass_through was called with the correct model ID
mock_router.get_available_deployment_for_pass_through.assert_called_once_with(model="gemini-pro")
# 2. Check that get_model_list was NOT called (this ensures we aren't doing the old logic)
mock_router.get_model_list.assert_not_called()
# 3. Verify that the project and location from the deployment were used (passed to _prepare_vertex_auth_headers)
# The args are: request, vertex_credentials, router_credentials, vertex_project, vertex_location, ...
# We check the 4th and 5th args (index 3 and 4)
call_args = mock_prep_headers.call_args
assert call_args[1]['vertex_project'] == "test-project-lb"
assert call_args[1]['vertex_location'] == "us-central1-lb"
def test_get_available_deployment_for_pass_through_filters_correctly():
"""
Test that get_available_deployment_for_pass_through filters deployments correctly
"""
from litellm.router import Router
# Configure router with both pass-through and non-pass-through deployments
model_list = [
{
"model_name": "gemini-pro",
"litellm_params": {
"model": "vertex_ai/gemini-pro",
"vertex_project": "project-1",
"vertex_location": "us-central1",
"use_in_pass_through": True, # Supports pass-through
}
},
{
"model_name": "gemini-pro",
"litellm_params": {
"model": "vertex_ai/gemini-pro",
"vertex_project": "project-2",
"vertex_location": "us-west1",
"use_in_pass_through": False, # Does not support pass-through
}
},
{
"model_name": "gemini-pro",
"litellm_params": {
"model": "vertex_ai/gemini-pro",
"vertex_project": "project-3",
"vertex_location": "us-east1",
# use_in_pass_through not set (defaults to False)
}
},
]
router = Router(model_list=model_list, routing_strategy="simple-shuffle")
# Test: Should only return project-1 (use_in_pass_through=True)
deployment = router.get_available_deployment_for_pass_through(model="gemini-pro")
assert deployment is not None
assert deployment["litellm_params"]["vertex_project"] == "project-1"
assert deployment["litellm_params"]["use_in_pass_through"] is True
def test_get_available_deployment_for_pass_through_no_deployments():
"""
Test that correct error is thrown when there are no pass-through deployments
"""
import litellm
from litellm.router import Router
model_list = [
{
"model_name": "gemini-pro",
"litellm_params": {
"model": "vertex_ai/gemini-pro",
"vertex_project": "project-1",
"vertex_location": "us-central1",
"use_in_pass_through": False, # Does not support pass-through
}
}
]
router = Router(model_list=model_list)
# Should throw BadRequestError
with pytest.raises(litellm.BadRequestError) as exc_info:
router.get_available_deployment_for_pass_through(model="gemini-pro")
assert "use_in_pass_through=True" in str(exc_info.value)
def test_get_available_deployment_for_pass_through_load_balancing():
"""
Test load balancing for pass-through deployments
"""
from litellm.router import Router
model_list = [
{
"model_name": "gemini-pro",
"litellm_params": {
"model": "vertex_ai/gemini-pro",
"vertex_project": "project-1",
"vertex_location": "us-central1",
"use_in_pass_through": True,
"rpm": 100,
}
},
{
"model_name": "gemini-pro",
"litellm_params": {
"model": "vertex_ai/gemini-pro",
"vertex_project": "project-2",
"vertex_location": "us-west1",
"use_in_pass_through": True,
"rpm": 200, # Higher RPM should be selected more frequently
}
},
]
router = Router(
model_list=model_list,
routing_strategy="simple-shuffle"
)
# Call multiple times and track selected deployments
selections = {"project-1": 0, "project-2": 0}
for _ in range(100):
deployment = router.get_available_deployment_for_pass_through(model="gemini-pro")
project = deployment["litellm_params"]["vertex_project"]
selections[project] += 1
# Due to rpm weight, project-2 should be selected more times
assert selections["project-2"] > selections["project-1"]
@pytest.mark.asyncio
async def test_async_get_available_deployment_for_pass_through():
"""
Test the async version of get_available_deployment_for_pass_through
"""
from litellm.router import Router
model_list = [
{
"model_name": "gemini-pro",
"litellm_params": {
"model": "vertex_ai/gemini-pro",
"vertex_project": "project-1",
"vertex_location": "us-central1",
"use_in_pass_through": True,
}
}
]
router = Router(
model_list=model_list,
routing_strategy="simple-shuffle"
)
deployment = await router.async_get_available_deployment_for_pass_through(
model="gemini-pro",
request_kwargs={}
)
assert deployment is not None
assert deployment["litellm_params"]["use_in_pass_through"] is True
@@ -10,6 +10,114 @@ import pytest
from litellm.proxy.db.db_transaction_queue.spend_log_cleanup import SpendLogCleanup
def test_spend_log_cleanup_cron_scheduling():
"""Test that cron expressions are correctly parsed for spend log cleanup scheduling"""
from apscheduler.triggers.cron import CronTrigger
# Valid cron expressions
cron_expr = "0 4 * * *" # 4:00 AM daily
trigger = CronTrigger.from_crontab(cron_expr)
assert trigger is not None
# Every minute (useful for testing)
trigger_minute = CronTrigger.from_crontab("*/1 * * * *")
assert trigger_minute is not None
# Specific day and hour
trigger_weekly = CronTrigger.from_crontab("0 3 * * 0") # 3 AM every Sunday
assert trigger_weekly is not None
# Invalid cron expression should raise ValueError
with pytest.raises(ValueError):
CronTrigger.from_crontab("invalid cron")
with pytest.raises(ValueError):
CronTrigger.from_crontab("60 25 * * *") # Invalid minute and hour
def test_spend_log_cleanup_cron_scheduler_integration():
"""
Integration test: Verify the proxy_server scheduler logic correctly adds
cron-based cleanup job when maximum_spend_logs_cleanup_cron is configured.
This tests the logic in proxy_server.py lines 4671-4717 without requiring
a real database connection.
"""
from unittest.mock import MagicMock
from apscheduler.triggers.cron import CronTrigger
# Mock scheduler
mock_scheduler = MagicMock()
mock_prisma_client = MagicMock()
mock_cleanup_instance = MagicMock()
# Test Case 1: Cron-based scheduling
general_settings_cron = {
"maximum_spend_logs_retention_period": "7d",
"maximum_spend_logs_cleanup_cron": "0 4 * * *", # 4 AM daily
}
cleanup_cron = general_settings_cron.get("maximum_spend_logs_cleanup_cron")
assert cleanup_cron is not None
# Simulate the scheduler logic from proxy_server.py
cron_trigger = CronTrigger.from_crontab(cleanup_cron)
mock_scheduler.add_job(
mock_cleanup_instance.cleanup_old_spend_logs,
cron_trigger,
args=[mock_prisma_client],
id="spend_log_cleanup_job",
replace_existing=True,
misfire_grace_time=3600,
)
# Verify scheduler was called correctly
mock_scheduler.add_job.assert_called_once()
call_args = mock_scheduler.add_job.call_args
# Verify the trigger is a CronTrigger
assert isinstance(call_args[0][1], CronTrigger)
# Verify job ID
assert call_args[1]["id"] == "spend_log_cleanup_job"
assert call_args[1]["replace_existing"] is True
# Test Case 2: Interval-based scheduling (fallback)
mock_scheduler.reset_mock()
general_settings_interval = {
"maximum_spend_logs_retention_period": "7d",
# No cron, so it should fall back to interval
}
cleanup_cron_fallback = general_settings_interval.get(
"maximum_spend_logs_cleanup_cron"
)
assert cleanup_cron_fallback is None # No cron configured
# Simulate interval-based scheduling fallback
retention_interval = general_settings_interval.get(
"maximum_spend_logs_retention_interval", "1d"
)
from litellm.litellm_core_utils.duration_parser import duration_in_seconds
interval_seconds = duration_in_seconds(retention_interval)
mock_scheduler.add_job(
mock_cleanup_instance.cleanup_old_spend_logs,
"interval",
seconds=interval_seconds,
args=[mock_prisma_client],
id="spend_log_cleanup_job",
replace_existing=True,
)
# Verify interval scheduling was called
mock_scheduler.add_job.assert_called_once()
interval_call_args = mock_scheduler.add_job.call_args
assert interval_call_args[0][1] == "interval"
assert interval_call_args[1]["seconds"] == 86400 # 1 day in seconds
@pytest.mark.asyncio
async def test_should_delete_spend_logs():
# Test case 1: No retention set