From 075f7ebb5fc4614b7896b7150f11f4ab4b6b6622 Mon Sep 17 00:00:00 2001 From: YutaSaito <36355491+uc4w6c@users.noreply.github.com> Date: Wed, 14 Jan 2026 21:20:16 +0900 Subject: [PATCH 1/8] feat: contextual gap checks, word-form digits (#18301) Co-authored-by: Krish Dholakia --- .../litellm_content_filter/content_filter.py | 213 ++++++++++++++++-- .../litellm_content_filter/patterns.json | 23 +- .../litellm_content_filter/patterns.py | 22 +- .../content_filter/test_content_filter.py | 1 - 4 files changed, 226 insertions(+), 33 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py index a04e438f48..c9bd0135a0 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py +++ b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py @@ -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"(? (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) diff --git a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/patterns.json b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/patterns.json index d8ec22f81a..f2427b5b92 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/patterns.json +++ b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/patterns.json @@ -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 @@ } ] } - - diff --git a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/patterns.py b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/patterns.py index 776cf5bd8d..d3a66690a9 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/patterns.py +++ b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/patterns.py @@ -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. diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/content_filter/test_content_filter.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/content_filter/test_content_filter.py index 474d2a3003..265bc530dc 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/content_filter/test_content_filter.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/content_filter/test_content_filter.py @@ -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, ) From e7cc53f217d4d2d4412af6c013aab178a556ef2a Mon Sep 17 00:00:00 2001 From: Harshit Jain <48647625+Harshit28j@users.noreply.github.com> Date: Wed, 14 Jan 2026 22:12:04 +0530 Subject: [PATCH 2/8] fix(dynamic_rate_limiter_v3): fix TPM 25% limiting by ensuring priority logic only runs when configured (#19092) --- .../proxy/hooks/dynamic_rate_limiter_v3.py | 35 +++++++++++-------- 1 file changed, 21 insertions(+), 14 deletions(-) diff --git a/litellm/proxy/hooks/dynamic_rate_limiter_v3.py b/litellm/proxy/hooks/dynamic_rate_limiter_v3.py index 755f5fdc20..a659d62e3e 100644 --- a/litellm/proxy/hooks/dynamic_rate_limiter_v3.py +++ b/litellm/proxy/hooks/dynamic_rate_limiter_v3.py @@ -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}" From e8c4cad8851088103177aee1b68a7f1dbb3301ab Mon Sep 17 00:00:00 2001 From: Harshit Jain <48647625+Harshit28j@users.noreply.github.com> Date: Wed, 14 Jan 2026 22:14:48 +0530 Subject: [PATCH 3/8] feat(proxy): cleanup spend logs cron verification, fix, and docs (#19085) --- .../docs/proxy/spend_logs_deletion.md | 12 ++ litellm/proxy/proxy_server.py | 115 ++++++++++++------ .../proxy/test_spend_log_cleanup.py | 108 ++++++++++++++++ 3 files changed, 199 insertions(+), 36 deletions(-) diff --git a/docs/my-website/docs/proxy/spend_logs_deletion.md b/docs/my-website/docs/proxy/spend_logs_deletion.md index 05627c0774..b021457173 100644 --- a/docs/my-website/docs/proxy/spend_logs_deletion.md +++ b/docs/my-website/docs/proxy/spend_logs_deletion.md @@ -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) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 1dfaf78cb5..1b01ac3a30 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -3240,20 +3240,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): @@ -3266,11 +3268,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: @@ -3278,37 +3284,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( @@ -4675,27 +4695,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: @@ -9885,7 +9926,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 = [] """ diff --git a/tests/test_litellm/proxy/test_spend_log_cleanup.py b/tests/test_litellm/proxy/test_spend_log_cleanup.py index 6aa18c560c..1ffbb83cae 100644 --- a/tests/test_litellm/proxy/test_spend_log_cleanup.py +++ b/tests/test_litellm/proxy/test_spend_log_cleanup.py @@ -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 From 1391e419166b7ebf07e880fee53e1ed6aae323d9 Mon Sep 17 00:00:00 2001 From: Kris Xia Date: Thu, 15 Jan 2026 00:47:43 +0800 Subject: [PATCH 4/8] fix(vertex_ai): improve passthrough endpoint url parsing and construction (#17402) (#17526) * fix(vertex_ai): improve passthrough endpoint url parsing and construction (#17402) * test(proxy): add test for vertex passthrough load balancing Add a test that verifies _base_vertex_proxy_route uses get_available_deployment for proper load balancing instead of get_model_list. This ensures the correct deployment is selected from the router and vertex credentials are properly fetched. Also refactor the implementation to: - Use get_available_deployment instead of get_model_list - Add error handling for deployment retrieval - Improve code structure with try-except block * feat(proxy): add pass-through deployment filtering methods Add dedicated methods to filter and select deployments for pass-through endpoints: - Implement get_available_deployment_for_pass_through() to ensure only deployments with use_in_pass_through=True are considered - Implement async_get_available_deployment_for_pass_through() for async operations - Add _filter_pass_through_deployments() helper method to filter by use_in_pass_through flag - Update vertex pass-through route to use the new dedicated method This ensures pass-through endpoints respect the use_in_pass_through configuration and apply proper load balancing strategy only to configured deployments. Add comprehensive tests to verify filtering and load balancing behavior. --- litellm/llms/vertex_ai/common_utils.py | 19 + .../llm_passthrough_endpoints.py | 20 ++ litellm/router.py | 339 ++++++++++++++++++ .../vertex_ai/test_vertex_ai_common_utils.py | 57 ++- .../test_vertex_passthrough_load_balancing.py | 222 ++++++++++++ 5 files changed, 650 insertions(+), 7 deletions(-) create mode 100644 tests/test_litellm/proxy/pass_through_endpoints/test_vertex_passthrough_load_balancing.py diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py index 2aa6a00c72..5ccbb8cd08 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/litellm/llms/vertex_ai/common_utils.py @@ -771,6 +771,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: @@ -820,6 +830,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 ) diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 5299b30b52..e48fd22bc8 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -1555,6 +1555,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, ) @@ -1584,6 +1585,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, diff --git a/litellm/router.py b/litellm/router.py index 638df49ac0..6523b5513a 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -7994,6 +7994,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=".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, @@ -8146,6 +8294,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]: @@ -8168,6 +8479,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 ): diff --git a/tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py b/tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py index b5637db3e5..12e35a4728 100644 --- a/tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py +++ b/tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py @@ -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", diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_vertex_passthrough_load_balancing.py b/tests/test_litellm/proxy/pass_through_endpoints/test_vertex_passthrough_load_balancing.py new file mode 100644 index 0000000000..ceb231eb4c --- /dev/null +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_vertex_passthrough_load_balancing.py @@ -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 + From d92a0168cc419d8fbbdcbb81e03f9b2e500fe6ed Mon Sep 17 00:00:00 2001 From: Rayan Pal <90289028+theonlypal@users.noreply.github.com> Date: Wed, 14 Jan 2026 09:28:05 -0800 Subject: [PATCH 5/8] fix: keep type field in Gemini schema when properties is empty (#18979) --- litellm/llms/vertex_ai/common_utils.py | 4 ++-- .../vertex_ai/test_gemini_empty_properties.py | 16 ++++++++++++++++ .../vertex_ai/test_vertex_ai_common_utils.py | 4 ++-- 3 files changed, 20 insertions(+), 4 deletions(-) create mode 100644 tests/test_litellm/llms/vertex_ai/test_gemini_empty_properties.py diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py index 5ccbb8cd08..704e45e301 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/litellm/llms/vertex_ai/common_utils.py @@ -660,11 +660,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(): diff --git a/tests/test_litellm/llms/vertex_ai/test_gemini_empty_properties.py b/tests/test_litellm/llms/vertex_ai/test_gemini_empty_properties.py new file mode 100644 index 0000000000..1a4e4d35ca --- /dev/null +++ b/tests/test_litellm/llms/vertex_ai/test_gemini_empty_properties.py @@ -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 diff --git a/tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py b/tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py index 12e35a4728..19d8f17493 100644 --- a/tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py +++ b/tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py @@ -1297,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" From 377d4ce02015ee96f07b081600f669c6e42e2187 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Fri, 16 Jan 2026 11:19:41 +0530 Subject: [PATCH 6/8] Add medium value support for detail param for gemini --- litellm/llms/vertex_ai/gemini/transformation.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/litellm/llms/vertex_ai/gemini/transformation.py b/litellm/llms/vertex_ai/gemini/transformation.py index 22042f7d64..8f1338db92 100644 --- a/litellm/llms/vertex_ai/gemini/transformation.py +++ b/litellm/llms/vertex_ai/gemini/transformation.py @@ -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 From f8e25aa0166dcfdf2d47378b93943a2bb956d5be Mon Sep 17 00:00:00 2001 From: Yuta Saito Date: Fri, 16 Jan 2026 18:23:01 +0900 Subject: [PATCH 7/8] chore: add ALLOWED_CVES. Because Wolfi glibc still flagged even on 2.42-r5. --- ci_cd/security_scans.sh | 1 + 1 file changed, 1 insertion(+) diff --git a/ci_cd/security_scans.sh b/ci_cd/security_scans.sh index 42ae25026d..9931730b7a 100755 --- a/ci_cd/security_scans.sh +++ b/ci_cd/security_scans.sh @@ -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 From 09fb1581cbd44913afce8e7052024aa7a57eb5c2 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Fri, 16 Jan 2026 16:37:44 +0530 Subject: [PATCH 8/8] Fix:add async_get_available_deployment_for_pass_through in code tests --- .../test_router_get_deployments.py | 202 ++++++++++++++++++ 1 file changed, 202 insertions(+) diff --git a/tests/local_testing/test_router_get_deployments.py b/tests/local_testing/test_router_get_deployments.py index 358ed74f55..8df04b4f1d 100644 --- a/tests/local_testing/test_router_get_deployments.py +++ b/tests/local_testing/test_router_get_deployments.py @@ -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}")