mirror of
https://github.com/tiennm99/litellm.git
synced 2026-07-11 01:05:19 +00:00
Merge pull request #19214 from BerriAI/main
merge main in malformed tool call PR
This commit is contained in:
@@ -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)
|
||||
|
||||
@@ -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
|
||||
|
||||
+194
-19
@@ -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.
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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 = []
|
||||
"""
|
||||
|
||||
@@ -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"
|
||||
|
||||
-1
@@ -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,
|
||||
)
|
||||
|
||||
+222
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user