diff --git a/docs/my-website/docs/providers/gemini.md b/docs/my-website/docs/providers/gemini.md index 110e3f3f090..23a02f7365c 100644 --- a/docs/my-website/docs/providers/gemini.md +++ b/docs/my-website/docs/providers/gemini.md @@ -1558,16 +1558,21 @@ LiteLLM Supports the following image types passed in `url` - Images with direct links - https://storage.googleapis.com/github-repo/img/gemini/intro/landmark3.jpg - Image in local storage - ./localimage.jpeg -## Image Resolution Control (Gemini 3+) +## Media Resolution Control (Images & Videos) -For Gemini 3+ models, LiteLLM supports per-part media resolution control using OpenAI's `detail` parameter. This allows you to specify different resolution levels for individual images in your request. +For Gemini 3+ models, LiteLLM supports per-part media resolution control using OpenAI's `detail` parameter. This allows you to specify different resolution levels for individual images and videos in your request, whether using `image_url` or `file` content types. **Supported `detail` values:** - `"low"` - Maps to `media_resolution: "low"` (280 tokens for images, 70 tokens per frame for videos) +- `"medium"` - Maps to `media_resolution: "medium"` - `"high"` - Maps to `media_resolution: "high"` (1120 tokens for images) +- `"ultra_high"` - Maps to `media_resolution: "ultra_high"` - `"auto"` or `None` - Model decides optimal resolution (no `media_resolution` set) -**Usage Example:** +**Usage Examples:** + + + ```python from litellm import completion @@ -1604,10 +1609,193 @@ response = completion( ) ``` + + + +```python +from litellm import completion + +messages = [ + { + "role": "user", + "content": [ + { + "type": "text", + "text": "Analyze this video" + }, + { + "type": "file", + "file": { + "file_id": "gs://my-bucket/video.mp4", + "format": "video/mp4", + "detail": "high" # High resolution for detailed video analysis + } + } + ] + } +] + +response = completion( + model="gemini/gemini-3-pro-preview", + messages=messages, +) +``` + + + + :::info -**Per-Part Resolution:** Each image in your request can have its own `detail` setting, allowing mixed-resolution requests (e.g., a high-res chart alongside a low-res icon). This feature is only available for Gemini 3+ models. +**Per-Part Resolution:** Each image or video in your request can have its own `detail` setting, allowing mixed-resolution requests (e.g., a high-res chart alongside a low-res icon). This feature works with both `image_url` and `file` content types, and is only available for Gemini 3+ models. ::: +## Video Metadata Control + +For Gemini 3+ models, LiteLLM supports fine-grained video processing control through the `video_metadata` field. This allows you to specify frame extraction rates and time ranges for video analysis. + +**Supported `video_metadata` parameters:** + +| Parameter | Type | Description | Example | +|-----------|------|-------------|---------| +| `fps` | Number | Frame extraction rate (frames per second) | `5` | +| `start_offset` | String | Start time for video clip processing | `"10s"` | +| `end_offset` | String | End time for video clip processing | `"60s"` | + +:::note +**Field Name Conversion:** LiteLLM automatically converts snake_case field names to camelCase for the Gemini API: +- `start_offset` → `startOffset` +- `end_offset` → `endOffset` +- `fps` remains unchanged +::: + +:::warning +- **Gemini 3+ Only:** This feature is only available for Gemini 3.0 and newer models +- **Video Files Recommended:** While `video_metadata` is designed for video files, error handling for other media types is delegated to the Vertex AI API +- **File Formats Supported:** Works with `gs://`, `https://`, and base64-encoded video files +::: + +**Usage Examples:** + + + + +```python +from litellm import completion + +response = completion( + model="gemini/gemini-3-pro-preview", + messages=[ + { + "role": "user", + "content": [ + {"type": "text", "text": "Analyze this video clip"}, + { + "type": "file", + "file": { + "file_id": "gs://my-bucket/video.mp4", + "format": "video/mp4", + "video_metadata": { + "fps": 5, # Extract 5 frames per second + "start_offset": "10s", # Start from 10 seconds + "end_offset": "60s" # End at 60 seconds + } + } + } + ] + } + ] +) + +print(response.choices[0].message.content) +``` + + + + +```python +from litellm import completion + +response = completion( + model="gemini/gemini-3-pro-preview", + messages=[ + { + "role": "user", + "content": [ + {"type": "text", "text": "Provide detailed analysis of this video segment"}, + { + "type": "file", + "file": { + "file_id": "https://example.com/presentation.mp4", + "format": "video/mp4", + "detail": "high", # High resolution for detailed analysis + "video_metadata": { + "fps": 10, # Extract 10 frames per second + "start_offset": "30s", # Start from 30 seconds + "end_offset": "90s" # End at 90 seconds + } + } + } + ] + } + ] +) + +print(response.choices[0].message.content) +``` + + + + +1. Setup config.yaml + +```yaml +model_list: + - model_name: gemini-3-pro + litellm_params: + model: gemini/gemini-3-pro-preview + api_key: os.environ/GEMINI_API_KEY +``` + +2. Start proxy + +```bash +litellm --config /path/to/config.yaml +``` + +3. Make request + +```bash +curl http://0.0.0.0:4000/v1/chat/completions \ + -H "Content-Type: application/json" \ + -H "Authorization: Bearer " \ + -d '{ + "model": "gemini-3-pro", + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Analyze this video clip"}, + { + "type": "file", + "file": { + "file_id": "gs://my-bucket/video.mp4", + "format": "video/mp4", + "detail": "high", + "video_metadata": { + "fps": 5, + "start_offset": "10s", + "end_offset": "60s" + } + } + } + ] + } + ] + }' +``` + + + + ## Sample Usage ```python import os diff --git a/docs/my-website/docs/providers/vertex.md b/docs/my-website/docs/providers/vertex.md index 5647b5292ef..63e4dceec00 100644 --- a/docs/my-website/docs/providers/vertex.md +++ b/docs/my-website/docs/providers/vertex.md @@ -1968,6 +1968,244 @@ assert isinstance( ``` +## Media Resolution Control (Images & Videos) + +For Gemini 3+ models, LiteLLM supports per-part media resolution control using OpenAI's `detail` parameter. This allows you to specify different resolution levels for individual images and videos in your request, whether using `image_url` or `file` content types. + +**Supported `detail` values:** +- `"low"` - Maps to `media_resolution: "low"` (280 tokens for images, 70 tokens per frame for videos) +- `"medium"` - Maps to `media_resolution: "medium"` +- `"high"` - Maps to `media_resolution: "high"` (1120 tokens for images) +- `"ultra_high"` - Maps to `media_resolution: "ultra_high"` +- `"auto"` or `None` - Model decides optimal resolution (no `media_resolution` set) + +**Usage Examples:** + + + + +```python +from litellm import completion + +messages = [ + { + "role": "user", + "content": [ + { + "type": "image_url", + "image_url": { + "url": "https://example.com/chart.png", + "detail": "high" # High resolution for detailed chart analysis + } + }, + { + "type": "text", + "text": "Analyze this chart" + }, + { + "type": "image_url", + "image_url": { + "url": "https://example.com/icon.png", + "detail": "low" # Low resolution for simple icon + } + } + ] + } +] + +response = completion( + model="vertex_ai/gemini-3-pro-preview", + messages=messages, +) +``` + + + + +```python +from litellm import completion + +messages = [ + { + "role": "user", + "content": [ + { + "type": "text", + "text": "Analyze this video" + }, + { + "type": "file", + "file": { + "file_id": "gs://my-bucket/video.mp4", + "format": "video/mp4", + "detail": "high" # High resolution for detailed video analysis + } + } + ] + } +] + +response = completion( + model="vertex_ai/gemini-3-pro-preview", + messages=messages, +) +``` + + + + +:::info +**Per-Part Resolution:** Each image or video in your request can have its own `detail` setting, allowing mixed-resolution requests (e.g., a high-res chart alongside a low-res icon). This feature works with both `image_url` and `file` content types, and is only available for Gemini 3+ models. +::: + +## Video Metadata Control + +For Gemini 3+ models, LiteLLM supports fine-grained video processing control through the `video_metadata` field. This allows you to specify frame extraction rates and time ranges for video analysis. + +**Supported `video_metadata` parameters:** + +| Parameter | Type | Description | Example | +|-----------|------|-------------|---------| +| `fps` | Number | Frame extraction rate (frames per second) | `5` | +| `start_offset` | String | Start time for video clip processing | `"10s"` | +| `end_offset` | String | End time for video clip processing | `"60s"` | + +:::note +**Field Name Conversion:** LiteLLM automatically converts snake_case field names to camelCase for the Gemini API: +- `start_offset` → `startOffset` +- `end_offset` → `endOffset` +- `fps` remains unchanged +::: + +:::warning +- **Gemini 3+ Only:** This feature is only available for Gemini 3.0 and newer models +- **Video Files Recommended:** While `video_metadata` is designed for video files, error handling for other media types is delegated to the Vertex AI API +- **File Formats Supported:** Works with `gs://`, `https://`, and base64-encoded video files +::: + +**Usage Examples:** + + + + +```python +from litellm import completion + +response = completion( + model="vertex_ai/gemini-3-pro-preview", + messages=[ + { + "role": "user", + "content": [ + {"type": "text", "text": "Analyze this video clip"}, + { + "type": "file", + "file": { + "file_id": "gs://my-bucket/video.mp4", + "format": "video/mp4", + "video_metadata": { + "fps": 5, # Extract 5 frames per second + "start_offset": "10s", # Start from 10 seconds + "end_offset": "60s" # End at 60 seconds + } + } + } + ] + } + ] +) + +print(response.choices[0].message.content) +``` + + + + +```python +from litellm import completion + +response = completion( + model="vertex_ai/gemini-3-pro-preview", + messages=[ + { + "role": "user", + "content": [ + {"type": "text", "text": "Provide detailed analysis of this video segment"}, + { + "type": "file", + "file": { + "file_id": "https://example.com/presentation.mp4", + "format": "video/mp4", + "detail": "high", # High resolution for detailed analysis + "video_metadata": { + "fps": 10, # Extract 10 frames per second + "start_offset": "30s", # Start from 30 seconds + "end_offset": "90s" # End at 90 seconds + } + } + } + ] + } + ] +) + +print(response.choices[0].message.content) +``` + + + + +1. Setup config.yaml + +```yaml +model_list: + - model_name: gemini-3-pro + litellm_params: + model: vertex_ai/gemini-3-pro-preview + vertex_project: your-project + vertex_location: us-central1 +``` + +2. Start proxy + +```bash +litellm --config /path/to/config.yaml +``` + +3. Make request + +```bash +curl http://0.0.0.0:4000/v1/chat/completions \ + -H "Content-Type: application/json" \ + -H "Authorization: Bearer " \ + -d '{ + "model": "gemini-3-pro", + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Analyze this video clip"}, + { + "type": "file", + "file": { + "file_id": "gs://my-bucket/video.mp4", + "format": "video/mp4", + "detail": "high", + "video_metadata": { + "fps": 5, + "start_offset": "10s", + "end_offset": "60s" + } + } + } + ] + } + ] + }' +``` + + + ## Usage - PDF / Videos / Audio etc. Files diff --git a/litellm/integrations/opentelemetry.py b/litellm/integrations/opentelemetry.py index 9fbccac68dd..0f8c2238d4b 100644 --- a/litellm/integrations/opentelemetry.py +++ b/litellm/integrations/opentelemetry.py @@ -988,7 +988,7 @@ class OpenTelemetry(CustomLogger): from opentelemetry._logs import SeverityNumber, get_logger, get_logger_provider try: - from opentelemetry.sdk._logs import LogRecord as SdkLogRecord # OTEL < 1.39.0 + from opentelemetry.sdk._logs import LogRecord as SdkLogRecord # type: ignore[attr-defined] # OTEL < 1.39.0 except ImportError: from opentelemetry.sdk._logs._internal import LogRecord as SdkLogRecord # OTEL >= 1.39.0 diff --git a/litellm/llms/vertex_ai/gemini/transformation.py b/litellm/llms/vertex_ai/gemini/transformation.py index 8f1338db92e..96e0963a920 100644 --- a/litellm/llms/vertex_ai/gemini/transformation.py +++ b/litellm/llms/vertex_ai/gemini/transformation.py @@ -72,17 +72,64 @@ def _convert_detail_to_media_resolution_enum( return {"level": "MEDIA_RESOLUTION_MEDIUM"} elif detail == "high": return {"level": "MEDIA_RESOLUTION_HIGH"} + elif detail == "ultra_high": + return {"level": "MEDIA_RESOLUTION_ULTRA_HIGH"} return None -def _process_gemini_image( - image_url: str, +def _apply_gemini_3_metadata( + part: PartType, + model: Optional[str], + media_resolution_enum: Optional[Dict[str, str]], + video_metadata: Optional[Dict[str, Any]], +) -> PartType: + """ + Apply the unique media_resolution and video_metadata parameters of Gemini 3+ + """ + if model is None: + return part + + from .vertex_and_google_ai_studio_gemini import VertexGeminiConfig + + if not VertexGeminiConfig._is_gemini_3_or_newer(model): + return part + + part_dict = dict(part) + + if media_resolution_enum is not None: + part_dict["media_resolution"] = media_resolution_enum + + if video_metadata is not None: + gemini_video_metadata = {} + if "fps" in video_metadata: + gemini_video_metadata["fps"] = video_metadata["fps"] + if "start_offset" in video_metadata: + gemini_video_metadata["startOffset"] = video_metadata["start_offset"] + if "end_offset" in video_metadata: + gemini_video_metadata["endOffset"] = video_metadata["end_offset"] + if gemini_video_metadata: + part_dict["video_metadata"] = gemini_video_metadata + + return cast(PartType, part_dict) + + +def _process_gemini_media( + image_url: str, format: Optional[str] = None, media_resolution_enum: Optional[Dict[str, str]] = None, model: Optional[str] = None, + video_metadata: Optional[Dict[str, Any]] = None, ) -> PartType: """ - Given an image URL, return the appropriate PartType for Gemini + Given a media URL (image, audio, or video), return the appropriate PartType for Gemini + By the way, actually video_metadata can only be used with videos; it cannot be used with images, audio, or files. However, I haven't made any special handling because vertex returns a parameter error. + + Args: + image_url: The URL or base64 string of the media (image, audio, or video) + format: The MIME type of the media + media_resolution_enum: Media resolution level (for Gemini 3+) + model: The model name (to check version compatibility) + video_metadata: Video-specific metadata (fps, start_offset, end_offset) """ try: @@ -104,14 +151,9 @@ def _process_gemini_image( mime_type = format file_data = FileDataType(mime_type=mime_type, file_uri=image_url) part: PartType = {"file_data": file_data} - - if media_resolution_enum is not None and model is not None: - from .vertex_and_google_ai_studio_gemini import VertexGeminiConfig - if VertexGeminiConfig._is_gemini_3_or_newer(model): - part_dict = dict(part) - part_dict["media_resolution"] = media_resolution_enum - return cast(PartType, part_dict) - return part + return _apply_gemini_3_metadata( + part, model, media_resolution_enum, video_metadata + ) elif ( "https://" in image_url and (image_type := format or _get_image_mime_type_from_url(image_url)) @@ -119,27 +161,16 @@ def _process_gemini_image( ): file_data = FileDataType(mime_type=image_type, file_uri=image_url) part = {"file_data": file_data} - - if media_resolution_enum is not None and model is not None: - from .vertex_and_google_ai_studio_gemini import VertexGeminiConfig - if VertexGeminiConfig._is_gemini_3_or_newer(model): - part_dict = dict(part) - part_dict["media_resolution"] = media_resolution_enum - return cast(PartType, part_dict) - return part + return _apply_gemini_3_metadata( + part, model, media_resolution_enum, video_metadata + ) elif "http://" in image_url or "https://" in image_url or "base64" in image_url: image = convert_to_anthropic_image_obj(image_url, format=format) _blob: BlobType = {"data": image["data"], "mime_type": image["media_type"]} - part = {"inline_data": cast(BlobType, _blob)} - - if media_resolution_enum is not None and model is not None: - from .vertex_and_google_ai_studio_gemini import VertexGeminiConfig - if VertexGeminiConfig._is_gemini_3_or_newer(model): - part_dict = dict(part) - part_dict["media_resolution"] = media_resolution_enum - return cast(PartType, part_dict) - return part + return _apply_gemini_3_metadata( + part, model, media_resolution_enum, video_metadata + ) raise Exception("Invalid image received - {}".format(image_url)) except Exception as e: raise e @@ -253,8 +284,8 @@ def _gemini_convert_messages_with_history( # noqa: PLR0915 media_resolution_enum = _convert_detail_to_media_resolution_enum(detail) else: image_url = img_element["image_url"] - _part = _process_gemini_image( - image_url=image_url, + _part = _process_gemini_media( + image_url=image_url, format=format, media_resolution_enum=media_resolution_enum, model=model, @@ -279,7 +310,7 @@ def _gemini_convert_messages_with_history( # noqa: PLR0915 ) ) ) - _part = _process_gemini_image( + _part = _process_gemini_media( image_url=openai_image_str, format=audio_format_modified, model=model, @@ -290,16 +321,24 @@ def _gemini_convert_messages_with_history( # noqa: PLR0915 file_id = file_element["file"].get("file_id") format = file_element["file"].get("format") file_data = file_element["file"].get("file_data") + detail = file_element["file"].get("detail") + video_metadata = file_element["file"].get("video_metadata") passed_file = file_id or file_data if passed_file is None: raise Exception( "Unknown file type. Please pass in a file_id or file_data" ) + + # Convert detail to media_resolution_enum + media_resolution_enum = _convert_detail_to_media_resolution_enum(detail) + try: - _part = _process_gemini_image( - image_url=passed_file, + _part = _process_gemini_media( + image_url=passed_file, format=format, model=model, + media_resolution_enum=media_resolution_enum, + video_metadata=video_metadata, ) _parts.append(_part) except Exception: diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py index dd34fbd772b..2d2e07e74db 100644 --- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py +++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py @@ -1018,25 +1018,34 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): optional_params["parallel_tool_calls"] = value elif param == "seed": optional_params["seed"] = value - elif param == "reasoning_effort" and isinstance(value, str): - # Validate no conflict with thinking_level - VertexGeminiConfig._validate_thinking_config_conflicts( - optional_params=optional_params, - param_name="reasoning_effort", - param_description="thinking_budget", - ) - if VertexGeminiConfig._is_gemini_3_or_newer(model): - optional_params["thinkingConfig"] = ( - VertexGeminiConfig._map_reasoning_effort_to_thinking_level( - value, model - ) + elif param == "reasoning_effort": + # Extract effort value - handle both string and dict formats + # Dict format comes from OpenAI Agents SDK: {"effort": "high", "summary": "auto"} + effort_value: Optional[str] = None + if isinstance(value, str): + effort_value = value + elif isinstance(value, dict): + effort_value = value.get("effort") + + if effort_value is not None: + # Validate no conflict with thinking_level + VertexGeminiConfig._validate_thinking_config_conflicts( + optional_params=optional_params, + param_name="reasoning_effort", + param_description="thinking_budget", ) - else: - optional_params["thinkingConfig"] = ( - VertexGeminiConfig._map_reasoning_effort_to_thinking_budget( - value, model + if VertexGeminiConfig._is_gemini_3_or_newer(model): + optional_params["thinkingConfig"] = ( + VertexGeminiConfig._map_reasoning_effort_to_thinking_level( + effort_value, model + ) + ) + else: + optional_params["thinkingConfig"] = ( + VertexGeminiConfig._map_reasoning_effort_to_thinking_budget( + effort_value, model + ) ) - ) elif param == "thinking": # Validate no conflict with thinking_level VertexGeminiConfig._validate_thinking_config_conflicts( diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 135b0d46ed0..0a6f6271739 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -12696,8 +12696,8 @@ "supports_web_search": true }, "gemini-2.5-flash-lite": { - "cache_read_input_token_cost": 2.5e-08, - "input_cost_per_audio_token": 5e-07, + "cache_read_input_token_cost": 1e-08, + "input_cost_per_audio_token": 3e-07, "input_cost_per_token": 1e-07, "litellm_provider": "vertex_ai-language-models", "max_audio_length_hours": 8.4, @@ -12741,7 +12741,7 @@ "supports_web_search": true }, "gemini-2.5-flash-lite-preview-09-2025": { - "cache_read_input_token_cost": 2.5e-08, + "cache_read_input_token_cost": 1e-08, "input_cost_per_audio_token": 3e-07, "input_cost_per_token": 1e-07, "litellm_provider": "vertex_ai-language-models", @@ -14532,8 +14532,8 @@ "supports_web_search": true }, "gemini/gemini-2.5-flash-lite": { - "cache_read_input_token_cost": 2.5e-08, - "input_cost_per_audio_token": 5e-07, + "cache_read_input_token_cost": 1e-08, + "input_cost_per_audio_token": 3e-07, "input_cost_per_token": 1e-07, "litellm_provider": "gemini", "max_audio_length_hours": 8.4, @@ -14579,7 +14579,7 @@ "tpm": 250000 }, "gemini/gemini-2.5-flash-lite-preview-09-2025": { - "cache_read_input_token_cost": 2.5e-08, + "cache_read_input_token_cost": 1e-08, "input_cost_per_audio_token": 3e-07, "input_cost_per_token": 1e-07, "litellm_provider": "gemini", diff --git a/litellm/proxy/common_utils/key_rotation_manager.py b/litellm/proxy/common_utils/key_rotation_manager.py index 3c367eafbc1..13bbf2272f7 100644 --- a/litellm/proxy/common_utils/key_rotation_manager.py +++ b/litellm/proxy/common_utils/key_rotation_manager.py @@ -26,97 +26,119 @@ class KeyRotationManager: """ Manages automated key rotation based on individual key rotation schedules. """ - + def __init__(self, prisma_client: PrismaClient): self.prisma_client = prisma_client - + async def process_rotations(self): """ Main entry point - find and rotate keys that are due for rotation """ try: verbose_proxy_logger.info("Starting scheduled key rotation check...") - + # Find keys that are due for rotation keys_to_rotate = await self._find_keys_needing_rotation() - + if not keys_to_rotate: verbose_proxy_logger.debug("No keys are due for rotation at this time") return - - verbose_proxy_logger.info(f"Found {len(keys_to_rotate)} keys due for rotation") - + + verbose_proxy_logger.info( + f"Found {len(keys_to_rotate)} keys due for rotation" + ) + # Rotate each key for key in keys_to_rotate: try: await self._rotate_key(key) - key_identifier = key.key_name or (key.token[:8] + "..." if key.token else "unknown") - verbose_proxy_logger.info(f"Successfully rotated key: {key_identifier}") + key_identifier = key.key_name or ( + key.token[:8] + "..." if key.token else "unknown" + ) + verbose_proxy_logger.info( + f"Successfully rotated key: {key_identifier}" + ) except Exception as e: - key_identifier = key.key_name or (key.token[:8] + "..." if key.token else "unknown") - verbose_proxy_logger.error(f"Failed to rotate key {key_identifier}: {e}") - + key_identifier = key.key_name or ( + key.token[:8] + "..." if key.token else "unknown" + ) + verbose_proxy_logger.error( + f"Failed to rotate key {key_identifier}: {e}" + ) + except Exception as e: verbose_proxy_logger.error(f"Key rotation process failed: {e}") - + async def _find_keys_needing_rotation(self) -> List[LiteLLM_VerificationToken]: """ Find keys that are due for rotation based on their key_rotation_at timestamp. - + Logic: - Key has auto_rotate = true - key_rotation_at is null (needs initial setup) OR key_rotation_at <= now """ now = datetime.now(timezone.utc) - - keys_with_rotation = await self.prisma_client.db.litellm_verificationtoken.find_many( - where={ - "auto_rotate": True, # Only keys marked for auto rotation - "OR": [ - {"key_rotation_at": None}, # Keys that need initial rotation time setup - {"key_rotation_at": {"lte": now}} # Keys where rotation time has passed - ] - } + + keys_with_rotation = ( + await self.prisma_client.db.litellm_verificationtoken.find_many( + where={ + "auto_rotate": True, # Only keys marked for auto rotation + "OR": [ + { + "key_rotation_at": None + }, # Keys that need initial rotation time setup + { + "key_rotation_at": {"lte": now} + }, # Keys where rotation time has passed + ], + } + ) ) - + return keys_with_rotation - + def _should_rotate_key(self, key: LiteLLM_VerificationToken, now: datetime) -> bool: """ Determine if a key should be rotated based on key_rotation_at timestamp. """ if not key.rotation_interval: return False - + # If key_rotation_at is not set, rotate immediately (and set it) if key.key_rotation_at is None: return True - + # Check if the rotation time has passed return now >= key.key_rotation_at - + async def _rotate_key(self, key: LiteLLM_VerificationToken): """ Rotate a single key using existing regenerate_key_fn and call the rotation hook """ - # Create regenerate request + # Create regenerate request regenerate_request = RegenerateKeyRequest( - key=key.token or "" + key=key.token or "", + key_alias=key.key_alias, # Pass key alias to ensure correct secret is updated in AWS Secrets Manager ) - + # Create a system user for key rotation from litellm.proxy._types import UserAPIKeyAuth + system_user = UserAPIKeyAuth.get_litellm_internal_jobs_user_api_key_auth() - + # Use existing regenerate key function response = await regenerate_key_fn( data=regenerate_request, user_api_key_dict=system_user, - litellm_changed_by=LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME + litellm_changed_by=LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME, ) - + # Update the NEW key with rotation info (regenerate_key_fn creates a new token) - if isinstance(response, GenerateKeyResponse) and response.token_id and key.rotation_interval: + if ( + isinstance(response, GenerateKeyResponse) + and response.token_id + and key.rotation_interval + ): # Calculate next rotation time using helper function now = datetime.now(timezone.utc) next_rotation_time = _calculate_key_rotation_time(key.rotation_interval) @@ -125,10 +147,10 @@ class KeyRotationManager: data={ "rotation_count": (key.rotation_count or 0) + 1, "last_rotation_at": now, - "key_rotation_at": next_rotation_time - } + "key_rotation_at": next_rotation_time, + }, ) - + # Call the existing rotation hook for notifications, audit logs, etc. if isinstance(response, GenerateKeyResponse): await KeyManagementEventHooks.async_key_rotated_hook( @@ -136,6 +158,5 @@ class KeyRotationManager: existing_key_row=key, response=response, user_api_key_dict=system_user, - litellm_changed_by=LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME + litellm_changed_by=LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME, ) - \ No newline at end of file diff --git a/litellm/proxy/hooks/key_management_event_hooks.py b/litellm/proxy/hooks/key_management_event_hooks.py index 9263bca100c..50f8b2a3ded 100644 --- a/litellm/proxy/hooks/key_management_event_hooks.py +++ b/litellm/proxy/hooks/key_management_event_hooks.py @@ -152,7 +152,8 @@ class KeyManagementEventHooks: ) await KeyManagementEventHooks._rotate_virtual_key_in_secret_manager( current_secret_name=initial_secret_name, - new_secret_name=data.key_alias + new_secret_name=response.key_alias + or data.key_alias or f"virtual-key-{response.token_id}", new_secret_value=response.key, ) diff --git a/litellm/proxy/prompts/prompt_endpoints.py b/litellm/proxy/prompts/prompt_endpoints.py index 0c77b6f8510..73e0ece3e2c 100644 --- a/litellm/proxy/prompts/prompt_endpoints.py +++ b/litellm/proxy/prompts/prompt_endpoints.py @@ -36,13 +36,13 @@ router = APIRouter() def get_base_prompt_id(prompt_id: str) -> str: """ Extract the base prompt ID by stripping the version suffix if present. - + Args: prompt_id: Prompt ID that may include version suffix (e.g., "jack_success.v1" or "jack_success_v1") - + Returns: Base prompt ID without version suffix (e.g., "jack_success") - + Examples: >>> get_base_prompt_id("jack_success.v1") "jack_success" @@ -63,13 +63,13 @@ def get_base_prompt_id(prompt_id: str) -> str: def get_version_number(prompt_id: str) -> int: """ Extract the version number from a versioned prompt ID. - + Args: prompt_id: Prompt ID that may include version suffix (e.g., "jack_success.v2" or "jack_success_v2") - + Returns: Version number (defaults to 1 if no version suffix or invalid format) - + Examples: >>> get_version_number("jack_success.v2") 2 @@ -85,7 +85,7 @@ def get_version_number(prompt_id: str) -> int: return int(version_str) except ValueError: pass - + # Try underscore separator (_v) if "_v" in prompt_id: version_str = prompt_id.split("_v")[1] @@ -93,21 +93,21 @@ def get_version_number(prompt_id: str) -> int: return int(version_str) except ValueError: pass - + return 1 def construct_versioned_prompt_id(prompt_id: str, version: Optional[int] = None) -> str: """ Construct a versioned prompt ID from a base prompt_id and version number. - + Args: prompt_id: Base prompt ID (e.g., "jack_success") version: Version number (if None, returns the base prompt_id unchanged) - + Returns: Versioned prompt ID (e.g., "jack_success.v4") - + Examples: >>> construct_versioned_prompt_id("jack_success", 4) "jack_success.v4" @@ -118,7 +118,7 @@ def construct_versioned_prompt_id(prompt_id: str, version: Optional[int] = None) """ if version is None: return prompt_id - + # Strip any existing version suffix first base_id = get_base_prompt_id(prompt_id) return f"{base_id}.v{version}" @@ -127,14 +127,14 @@ def construct_versioned_prompt_id(prompt_id: str, version: Optional[int] = None) def get_latest_version_prompt_id(prompt_id: str, all_prompt_ids: Dict[str, Any]) -> str: """ Find the latest version of a prompt from available prompt IDs. - + Args: prompt_id: Base prompt ID or versioned prompt ID (e.g., "jack_success" or "jack_success.v2") all_prompt_ids: Dictionary of all available prompt IDs (keys are prompt IDs) - + Returns: The prompt ID with the highest version number, or the original prompt_id if no versions exist - + Examples: >>> all_ids = {"jack.v1": {}, "jack.v2": {}, "jack.v3": {}} >>> get_latest_version_prompt_id("jack", all_ids) @@ -146,14 +146,14 @@ def get_latest_version_prompt_id(prompt_id: str, all_prompt_ids: Dict[str, Any]) "simple" """ base_id = get_base_prompt_id(prompt_id=prompt_id) - + # Find all versions of this prompt matching_versions = [] for stored_prompt_id in all_prompt_ids.keys(): if get_base_prompt_id(prompt_id=stored_prompt_id) == base_id: version_num = get_version_number(prompt_id=stored_prompt_id) matching_versions.append((version_num, stored_prompt_id)) - + # Use the highest version number if matching_versions: matching_versions.sort(reverse=True) @@ -166,45 +166,47 @@ def get_latest_version_prompt_id(prompt_id: str, all_prompt_ids: Dict[str, Any]) def get_latest_prompt_versions(prompts: List[PromptSpec]) -> List[PromptSpec]: """ Filter a list of prompts to return only the latest version of each unique prompt. - + Args: prompts: List of PromptSpec objects - + Returns: List of PromptSpec objects with only the latest version of each prompt """ latest_prompts: Dict[str, PromptSpec] = {} - + for prompt in prompts: base_id = get_base_prompt_id(prompt_id=prompt.prompt_id) version = get_version_number(prompt_id=prompt.prompt_id) - + # Keep the prompt with the highest version number if base_id not in latest_prompts: latest_prompts[base_id] = prompt else: - existing_version = get_version_number(prompt_id=latest_prompts[base_id].prompt_id) + existing_version = get_version_number( + prompt_id=latest_prompts[base_id].prompt_id + ) if version > existing_version: latest_prompts[base_id] = prompt - + return list(latest_prompts.values()) async def get_next_version_for_prompt(prisma_client, prompt_id: str) -> int: """ Get the next version number for a prompt. - + Args: prisma_client: Prisma database client prompt_id: Base prompt ID - + Returns: Next version number (1 if no versions exist, max_version + 1 otherwise) """ existing_prompts = await prisma_client.db.litellm_prompttable.find_many( where={"prompt_id": prompt_id} ) - + if existing_prompts: max_version = max(p.version for p in existing_prompts) return max_version + 1 @@ -215,27 +217,27 @@ async def get_next_version_for_prompt(prisma_client, prompt_id: str) -> int: def create_versioned_prompt_spec(db_prompt) -> PromptSpec: """ Helper function to create a PromptSpec with versioned prompt_id from a DB prompt entry. - + Args: db_prompt: The DB prompt object (from prisma) - + Returns: PromptSpec with versioned prompt_id (e.g., "chat_prompt.v1") """ import json from litellm.types.prompts.init_prompts import PromptLiteLLMParams - + prompt_dict = db_prompt.model_dump() base_prompt_id = prompt_dict["prompt_id"] version = prompt_dict.get("version", 1) - + # Parse litellm_params litellm_params_data = prompt_dict.get("litellm_params") if isinstance(litellm_params_data, str): litellm_params_data = json.loads(litellm_params_data) litellm_params = PromptLiteLLMParams(**litellm_params_data) - + # Parse prompt_info prompt_info_data = prompt_dict.get("prompt_info") if prompt_info_data: @@ -244,10 +246,10 @@ def create_versioned_prompt_spec(db_prompt) -> PromptSpec: prompt_info = PromptInfo(**prompt_info_data) else: prompt_info = PromptInfo(prompt_type="db") - + # Create versioned prompt_id versioned_prompt_id = f"{base_prompt_id}.v{version}" - + return PromptSpec( prompt_id=versioned_prompt_id, litellm_params=litellm_params, @@ -319,10 +321,14 @@ async def list_prompts( prompt_list = [] for prompt_id in prompts: if prompt_id in IN_MEMORY_PROMPT_REGISTRY.IN_MEMORY_PROMPTS: - original_prompt = IN_MEMORY_PROMPT_REGISTRY.IN_MEMORY_PROMPTS[prompt_id] + original_prompt = IN_MEMORY_PROMPT_REGISTRY.IN_MEMORY_PROMPTS[ + prompt_id + ] # Create a copy with base prompt_id (without version suffix) prompt_copy = PromptSpec( - prompt_id=get_base_prompt_id(prompt_id=original_prompt.prompt_id), + prompt_id=get_base_prompt_id( + prompt_id=original_prompt.prompt_id + ), litellm_params=original_prompt.litellm_params, prompt_info=original_prompt.prompt_info, created_at=original_prompt.created_at, @@ -407,32 +413,33 @@ async def get_prompt_versions( raise HTTPException( status_code=403, detail="Only proxy admins can view prompt versions" ) - + # Strip version suffix if provided (e.g., "jack_success.v1" -> "jack_success") base_prompt_id = get_base_prompt_id(prompt_id=prompt_id) - + # Get all prompts and filter by base_prompt_id all_prompts = list(IN_MEMORY_PROMPT_REGISTRY.IN_MEMORY_PROMPTS.values()) prompt_versions = [ - prompt for prompt in all_prompts + prompt + for prompt in all_prompts if get_base_prompt_id(prompt_id=prompt.prompt_id) == base_prompt_id ] - + if not prompt_versions: raise HTTPException( status_code=404, detail=f"No versions found for prompt ID {base_prompt_id}" ) - + # Create response with explicit version field for each prompt versioned_prompts = [] for prompt in prompt_versions: # Extract version number from the root prompt_id which has version suffix # (e.g., "jack-sparrow.v3" -> 3) version_number = get_version_number(prompt_id=prompt.prompt_id) - + # Strip version from prompt_id for clean display base_prompt_id = get_base_prompt_id(prompt_id=prompt.prompt_id) - + # Create a copy with explicit version field and clean prompt_id versioned_prompt = PromptSpec( prompt_id=base_prompt_id, # Clean ID without version (e.g., "jack-sparrow") @@ -443,10 +450,10 @@ async def get_prompt_versions( version=version_number, # Explicit version field (e.g., 3) ) versioned_prompts.append(versioned_prompt) - + # Sort by version number (descending - newest first) versioned_prompts.sort(key=lambda p: p.version or 1, reverse=True) - + return ListPromptsResponse(prompts=versioned_prompts) @@ -518,21 +525,21 @@ async def get_prompt_info( # Try to get prompt directly first prompt_spec = IN_MEMORY_PROMPT_REGISTRY.get_prompt_by_id(prompt_id) - + # If not found, try to find the latest version if prompt_spec is None: latest_prompt_id = get_latest_version_prompt_id( prompt_id=prompt_id, - all_prompt_ids=IN_MEMORY_PROMPT_REGISTRY.IN_MEMORY_PROMPTS + all_prompt_ids=IN_MEMORY_PROMPT_REGISTRY.IN_MEMORY_PROMPTS, ) prompt_spec = IN_MEMORY_PROMPT_REGISTRY.get_prompt_by_id(latest_prompt_id) - + if prompt_spec is None: raise HTTPException(status_code=400, detail=f"Prompt {prompt_id} not found") # Extract version number from the prompt_id version_number = get_version_number(prompt_id=prompt_spec.prompt_id) - + # Create a copy of the prompt spec with the base prompt ID (stripped of version) # and explicit version field for consistency with list_prompts and versions endpoints prompt_spec_response = PromptSpec( @@ -547,7 +554,9 @@ async def get_prompt_info( # Get prompt content from the callback prompt_template: Optional[PromptTemplateBase] = None try: - prompt_callback = IN_MEMORY_PROMPT_REGISTRY.get_prompt_callback_by_id(prompt_id) + prompt_callback = IN_MEMORY_PROMPT_REGISTRY.get_prompt_callback_by_id( + prompt_spec.prompt_id + ) if prompt_callback is not None: # Extract content based on integration type integration_name = prompt_callback.integration_name @@ -723,12 +732,12 @@ async def update_prompt( try: # Strip version suffix from prompt_id if present (e.g., "jack_success.v1" -> "jack_success") base_prompt_id = get_base_prompt_id(prompt_id=prompt_id) - + # Check if any version exists existing_prompts = await prisma_client.db.litellm_prompttable.find_many( where={"prompt_id": base_prompt_id} ) - + if not existing_prompts: raise HTTPException( status_code=404, detail=f"Prompt with ID {base_prompt_id} not found" @@ -736,7 +745,10 @@ async def update_prompt( # Check if it's a config prompt existing_in_memory = IN_MEMORY_PROMPT_REGISTRY.get_prompt_by_id(prompt_id) - if existing_in_memory and existing_in_memory.prompt_info.prompt_type == "config": + if ( + existing_in_memory + and existing_in_memory.prompt_info.prompt_type == "config" + ): raise HTTPException( status_code=400, detail="Cannot update config prompts.", @@ -828,17 +840,19 @@ async def delete_prompt( try: # Try to get prompt directly first existing_prompt = IN_MEMORY_PROMPT_REGISTRY.get_prompt_by_id(prompt_id) - + # If not found, try to find the latest version if existing_prompt is None: latest_prompt_id = get_latest_version_prompt_id( prompt_id=prompt_id, - all_prompt_ids=IN_MEMORY_PROMPT_REGISTRY.IN_MEMORY_PROMPTS + all_prompt_ids=IN_MEMORY_PROMPT_REGISTRY.IN_MEMORY_PROMPTS, + ) + existing_prompt = IN_MEMORY_PROMPT_REGISTRY.get_prompt_by_id( + latest_prompt_id ) - existing_prompt = IN_MEMORY_PROMPT_REGISTRY.get_prompt_by_id(latest_prompt_id) # Use the resolved prompt_id for deletion prompt_id = latest_prompt_id - + if existing_prompt is None: raise HTTPException( status_code=404, detail=f"Prompt with ID {prompt_id} not found" @@ -850,17 +864,18 @@ async def delete_prompt( detail="Cannot delete config prompts.", ) - # Delete the prompt from the database + # Get the base prompt ID (without version suffix) for database deletion + base_prompt_id = get_base_prompt_id(prompt_id=prompt_id) + + # Delete all versions of the prompt from the database await prisma_client.db.litellm_prompttable.delete_many( - where={"prompt_id": prompt_id} + where={"prompt_id": base_prompt_id} ) - # Remove the prompt from memory - del IN_MEMORY_PROMPT_REGISTRY.IN_MEMORY_PROMPTS[prompt_id] - if prompt_id in IN_MEMORY_PROMPT_REGISTRY.prompt_id_to_custom_prompt: - del IN_MEMORY_PROMPT_REGISTRY.prompt_id_to_custom_prompt[prompt_id] + # Remove all versions of the prompt from memory + IN_MEMORY_PROMPT_REGISTRY.delete_prompts_by_base_id(base_prompt_id) - return {"message": f"Prompt {prompt_id} deleted successfully"} + return {"message": f"Prompt {base_prompt_id} deleted successfully"} except HTTPException as e: raise e @@ -1036,68 +1051,66 @@ async def test_prompt( user_temperature, version, ) - + try: # Parse the dotprompt content and create PromptTemplate prompt_manager = PromptManager() frontmatter, template_content = prompt_manager._parse_frontmatter( content=request.dotprompt_content ) - + # Create PromptTemplate to leverage existing parameter extraction logic template = PromptTemplate( - content=template_content, - metadata=frontmatter, - template_id="test_prompt" + content=template_content, metadata=frontmatter, template_id="test_prompt" ) - + # Extract model from template if not template.model: raise HTTPException( - status_code=400, - detail="Model is required in dotprompt metadata" + status_code=400, detail="Model is required in dotprompt metadata" ) - + # Always render the template to extract system messages and other metadata variables = request.prompt_variables or {} rendered_content = prompt_manager.jinja_env.from_string( template_content ).render(**variables) - + # Convert rendered content to messages using DotpromptManager's method dotprompt_manager = DotpromptManager() rendered_messages = dotprompt_manager._convert_to_messages( rendered_content=rendered_content ) - + if not rendered_messages: raise HTTPException( - status_code=400, - detail="No messages found in rendered prompt" + status_code=400, detail="No messages found in rendered prompt" ) - + # If conversation history is provided, use it but preserve system messages if request.conversation_history: # Extract system messages from rendered prompt - system_messages = [msg for msg in rendered_messages if msg.get("role") == "system"] + system_messages = [ + msg for msg in rendered_messages if msg.get("role") == "system" + ] # Use conversation history for user/assistant messages messages = system_messages + request.conversation_history else: messages = rendered_messages # type: ignore[assignment] - + # Use PromptTemplate's optional_params which already extracts all parameters optional_params = template.optional_params.copy() - + # Always stream the response optional_params["stream"] = True - + # Build request data for chat completion data = { "model": template.model, "messages": messages, } data.update(optional_params) - + # Use ProxyBaseLLMRequestProcessing to go through all proxy logic base_llm_response_processor = ProxyBaseLLMRequestProcessing(data=data) result = await base_llm_response_processor.base_process_llm_request( @@ -1118,12 +1131,12 @@ async def test_prompt( user_api_base=user_api_base, version=version, ) - + if isinstance(result, BaseModel): return result.model_dump(exclude_none=True, exclude_unset=True) else: return result - + except HTTPException as e: raise e except Exception as e: @@ -1192,4 +1205,3 @@ async def convert_prompt_file_to_json( temp_file_path.parent.rmdir() except OSError: pass # Directory not empty or other error - diff --git a/litellm/proxy/prompts/prompt_registry.py b/litellm/proxy/prompts/prompt_registry.py index b4717687704..58df60a42cb 100644 --- a/litellm/proxy/prompts/prompt_registry.py +++ b/litellm/proxy/prompts/prompt_registry.py @@ -97,9 +97,9 @@ class InMemoryPromptRegistry: Prompt id to Prompt object mapping """ - self.prompt_id_to_custom_prompt: Dict[str, Optional[CustomPromptManagement]] = ( - {} - ) + self.prompt_id_to_custom_prompt: Dict[ + str, Optional[CustomPromptManagement] + ] = {} """ Guardrail id to CustomGuardrail object mapping """ @@ -174,5 +174,30 @@ class InMemoryPromptRegistry: """ return self.prompt_id_to_custom_prompt.get(prompt_id) + def delete_prompts_by_base_id(self, base_prompt_id: str) -> list[str]: + """ + Delete all prompts matching the given base prompt ID from memory. -IN_MEMORY_PROMPT_REGISTRY = InMemoryPromptRegistry() \ No newline at end of file + Args: + base_prompt_id: The base prompt ID (without version suffix) + + Returns: + List of prompt IDs that were deleted + """ + from litellm.proxy.prompts.prompt_endpoints import get_base_prompt_id + + prompts_to_delete = [ + pid + for pid in self.IN_MEMORY_PROMPTS.keys() + if get_base_prompt_id(prompt_id=pid) == base_prompt_id + ] + + for pid in prompts_to_delete: + del self.IN_MEMORY_PROMPTS[pid] + if pid in self.prompt_id_to_custom_prompt: + del self.prompt_id_to_custom_prompt[pid] + + return prompts_to_delete + + +IN_MEMORY_PROMPT_REGISTRY = InMemoryPromptRegistry() diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index c3a4de314e5..eef2af89799 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -548,9 +548,9 @@ except ImportError: server_root_path = os.getenv("SERVER_ROOT_PATH", "") _license_check = LicenseCheck() premium_user: bool = _license_check.is_premium() -premium_user_data: Optional[ - "EnterpriseLicenseData" -] = _license_check.airgapped_license_data +premium_user_data: Optional["EnterpriseLicenseData"] = ( + _license_check.airgapped_license_data +) global_max_parallel_request_retries_env: Optional[str] = os.getenv( "LITELLM_GLOBAL_MAX_PARALLEL_REQUEST_RETRIES" ) @@ -899,7 +899,7 @@ def get_openapi_schema(): from litellm.proxy.common_utils.custom_openapi_spec import CustomOpenAPISpec openapi_schema = CustomOpenAPISpec.add_llm_api_request_schema_body(openapi_schema) - + # Fix Swagger UI execute path error when server_root_path is set if server_root_path: openapi_schema["servers"] = [{"url": "/" + server_root_path.strip("/")}] @@ -925,7 +925,7 @@ def custom_openapi(): from litellm.proxy.common_utils.custom_openapi_spec import CustomOpenAPISpec openapi_schema = CustomOpenAPISpec.add_llm_api_request_schema_body(openapi_schema) - + # Fix Swagger UI execute path error when server_root_path is set if server_root_path: openapi_schema["servers"] = [{"url": "/" + server_root_path.strip("/")}] @@ -1203,9 +1203,9 @@ master_key: Optional[str] = None config_agents: Optional[List[AgentConfig]] = None otel_logging = False prisma_client: Optional[PrismaClient] = None -shared_aiohttp_session: Optional[ - "ClientSession" -] = None # Global shared session for connection reuse +shared_aiohttp_session: Optional["ClientSession"] = ( + None # Global shared session for connection reuse +) user_api_key_cache = DualCache( default_in_memory_ttl=UserAPIKeyCacheTTLEnum.in_memory_cache_ttl.value ) @@ -1213,9 +1213,9 @@ model_max_budget_limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter( dual_cache=user_api_key_cache ) litellm.logging_callback_manager.add_litellm_callback(model_max_budget_limiter) -redis_usage_cache: Optional[ - RedisCache -] = None # redis cache used for tracking spend, tpm/rpm limits +redis_usage_cache: Optional[RedisCache] = ( + None # redis cache used for tracking spend, tpm/rpm limits +) polling_via_cache_enabled: Union[Literal["all"], List[str], bool] = False polling_cache_ttl: int = 3600 # Default 1 hour TTL for polling cache user_custom_auth = None @@ -1554,9 +1554,9 @@ async def update_cache( # noqa: PLR0915 _id = "team_id:{}".format(team_id) try: # Fetch the existing cost for the given user - existing_spend_obj: Optional[ - LiteLLM_TeamTable - ] = await user_api_key_cache.async_get_cache(key=_id) + existing_spend_obj: Optional[LiteLLM_TeamTable] = ( + await user_api_key_cache.async_get_cache(key=_id) + ) if existing_spend_obj is None: # do nothing if team not in api key cache return @@ -3095,17 +3095,19 @@ class ProxyConfig: async def _update_llm_router( self, - new_models: list, + new_models: Optional[Json], proxy_logging_obj: ProxyLogging, ): global llm_router, llm_model_list, master_key, general_settings - + config_data = await proxy_config.get_config() + search_tools = self.parse_search_tools(config_data) try: + models_list: list = new_models if isinstance(new_models, list) else [] if llm_router is None and master_key is not None: - verbose_proxy_logger.debug(f"len new_models: {len(new_models)}") + verbose_proxy_logger.debug(f"len new_models: {len(models_list)}") _model_list: list = self.decrypt_model_list_from_db( - new_models=new_models + new_models=models_list ) if len(_model_list) > 0: verbose_proxy_logger.debug(f"_model_list: {_model_list}") @@ -3114,16 +3116,17 @@ class ProxyConfig: router_general_settings=RouterGeneralSettings( async_only_mode=True # only init async clients ), + search_tools=search_tools, ignore_invalid_deployments=True, ) verbose_proxy_logger.debug(f"updated llm_router: {llm_router}") else: - verbose_proxy_logger.debug(f"len new_models: {len(new_models)}") + verbose_proxy_logger.debug(f"len new_models: {len(models_list)}") ## DELETE MODEL LOGIC - await self._delete_deployment(db_models=new_models) + await self._delete_deployment(db_models=models_list) ## ADD MODEL LOGIC - self._add_deployment(db_models=new_models) + self._add_deployment(db_models=models_list) except Exception as e: verbose_proxy_logger.exception( @@ -3134,7 +3137,6 @@ class ProxyConfig: llm_model_list = llm_router.get_model_list() # check if user set any callbacks in Config Table - config_data = await proxy_config.get_config() self._add_callbacks_from_db_config(config_data) # router settings @@ -3944,10 +3946,10 @@ class ProxyConfig: ) try: - guardrails_in_db: List[ - Guardrail - ] = await GuardrailRegistry.get_all_guardrails_from_db( - prisma_client=prisma_client + guardrails_in_db: List[Guardrail] = ( + await GuardrailRegistry.get_all_guardrails_from_db( + prisma_client=prisma_client + ) ) verbose_proxy_logger.debug( "guardrails from the DB %s", str(guardrails_in_db) @@ -4274,9 +4276,9 @@ async def initialize( # noqa: PLR0915 user_api_base = api_base dynamic_config[user_model]["api_base"] = api_base if api_version: - os.environ[ - "AZURE_API_VERSION" - ] = api_version # set this for azure - litellm can read this from the env + os.environ["AZURE_API_VERSION"] = ( + api_version # set this for azure - litellm can read this from the env + ) if max_tokens: # model-specific param dynamic_config[user_model]["max_tokens"] = max_tokens if temperature: # model-specific param @@ -9729,9 +9731,9 @@ async def get_config_list( hasattr(sub_field_info, "description") and sub_field_info.description is not None ): - nested_fields[ - idx - ].field_description = sub_field_info.description + nested_fields[idx].field_description = ( + sub_field_info.description + ) idx += 1 _stored_in_db = None diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index 467c57c33d5..84185b6eec0 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -654,6 +654,8 @@ class ChatCompletionFileObjectFile(TypedDict, total=False): file_id: str filename: str format: str + detail: str # For video/image resolution control (low, medium, high, ultra_high) + video_metadata: Dict[str, Any] # For video-specific metadata (fps, start_offset, end_offset) class ChatCompletionFileObject(TypedDict): diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 135b0d46ed0..0a6f6271739 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -12696,8 +12696,8 @@ "supports_web_search": true }, "gemini-2.5-flash-lite": { - "cache_read_input_token_cost": 2.5e-08, - "input_cost_per_audio_token": 5e-07, + "cache_read_input_token_cost": 1e-08, + "input_cost_per_audio_token": 3e-07, "input_cost_per_token": 1e-07, "litellm_provider": "vertex_ai-language-models", "max_audio_length_hours": 8.4, @@ -12741,7 +12741,7 @@ "supports_web_search": true }, "gemini-2.5-flash-lite-preview-09-2025": { - "cache_read_input_token_cost": 2.5e-08, + "cache_read_input_token_cost": 1e-08, "input_cost_per_audio_token": 3e-07, "input_cost_per_token": 1e-07, "litellm_provider": "vertex_ai-language-models", @@ -14532,8 +14532,8 @@ "supports_web_search": true }, "gemini/gemini-2.5-flash-lite": { - "cache_read_input_token_cost": 2.5e-08, - "input_cost_per_audio_token": 5e-07, + "cache_read_input_token_cost": 1e-08, + "input_cost_per_audio_token": 3e-07, "input_cost_per_token": 1e-07, "litellm_provider": "gemini", "max_audio_length_hours": 8.4, @@ -14579,7 +14579,7 @@ "tpm": 250000 }, "gemini/gemini-2.5-flash-lite-preview-09-2025": { - "cache_read_input_token_cost": 2.5e-08, + "cache_read_input_token_cost": 1e-08, "input_cost_per_audio_token": 3e-07, "input_cost_per_token": 1e-07, "litellm_provider": "gemini", diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py index 70e6e9452e5..ebda37bb633 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py @@ -735,13 +735,13 @@ def test_file_data_field_order(): Related issue: Gemini API returns 400 INVALID_ARGUMENT when fields are in wrong order. """ import json - from litellm.llms.vertex_ai.gemini.transformation import _process_gemini_image + from litellm.llms.vertex_ai.gemini.transformation import _process_gemini_media # Test with HTTPS URL and explicit format (audio file) file_url = "https://generativelanguage.googleapis.com/v1beta/files/test123" format = "audio/mpeg" - result = _process_gemini_image(image_url=file_url, format=format) + result = _process_gemini_media(image_url=file_url, format=format) # Verify the result has file_data assert "file_data" in result @@ -770,12 +770,12 @@ def test_file_data_field_order(): def test_file_data_field_order_gcs_urls(): """Test that GCS URLs also maintain correct field order.""" import json - from litellm.llms.vertex_ai.gemini.transformation import _process_gemini_image + from litellm.llms.vertex_ai.gemini.transformation import _process_gemini_media # Test with GCS URL gcs_url = "gs://bucket/audio.mp3" - result = _process_gemini_image(image_url=gcs_url) + result = _process_gemini_media(image_url=gcs_url) # Verify the result has file_data assert "file_data" in result diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py index 969199f6ad0..5be080b53fa 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py +++ b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py @@ -1980,6 +1980,71 @@ def test_reasoning_effort_maps_to_thinking_level_gemini_3(): assert result["thinkingConfig"]["includeThoughts"] is False +def test_reasoning_effort_dict_format_gemini_3(): + """ + Test that reasoning_effort works when passed as dict format from OpenAI Agents SDK. + + The OpenAI Agents SDK passes reasoning_effort as {"effort": "high", "summary": "auto"} + instead of just a string. This test verifies that we correctly extract the effort value. + + Related issue: https://github.com/BerriAI/litellm/issues/19411 + """ + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( + VertexGeminiConfig, + ) + + v = VertexGeminiConfig() + model = "gemini-3-pro-preview" + + # Test dict format with effort="high" (OpenAI Agents SDK format) + optional_params = {} + non_default_params = {"reasoning_effort": {"effort": "high", "summary": "auto"}} + result = v.map_openai_params( + non_default_params=non_default_params, + optional_params=optional_params, + model=model, + drop_params=False, + ) + assert result["thinkingConfig"]["thinkingLevel"] == "high" + assert result["thinkingConfig"]["includeThoughts"] is True + + # Test dict format with effort="low" + optional_params = {} + non_default_params = {"reasoning_effort": {"effort": "low"}} + result = v.map_openai_params( + non_default_params=non_default_params, + optional_params=optional_params, + model=model, + drop_params=False, + ) + assert result["thinkingConfig"]["thinkingLevel"] == "low" + assert result["thinkingConfig"]["includeThoughts"] is True + + # Test dict format with effort="medium" + optional_params = {} + non_default_params = {"reasoning_effort": {"effort": "medium"}} + result = v.map_openai_params( + non_default_params=non_default_params, + optional_params=optional_params, + model=model, + drop_params=False, + ) + assert result["thinkingConfig"]["thinkingLevel"] == "high" + assert result["thinkingConfig"]["includeThoughts"] is True + + # Test dict format without effort key - should fall back to Gemini 3 default (low) + optional_params = {} + non_default_params = {"reasoning_effort": {"summary": "auto"}} + result = v.map_openai_params( + non_default_params=non_default_params, + optional_params=optional_params, + model=model, + drop_params=False, + ) + # Gemini 3 defaults to thinkingLevel="low" when no explicit effort is set + assert result["thinkingConfig"]["thinkingLevel"] == "low" + + def test_temperature_default_for_gemini_3(): """Test that temperature defaults to 1.0 for Gemini 3+ models when not specified""" from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( @@ -2746,3 +2811,273 @@ def test_gemini_image_gen_usage_metadata_prompt_vs_completion_separation(): # candidatesTokenCount (1290) - image_tokens (1290) = 0 assert result.completion_tokens_details.text_tokens == 0, \ "Completion text tokens should be 0 (image-only response)" + + +def test_file_object_detail_parameter(): + """Test that detail parameter works for type: file objects (Issue #19026)""" + from litellm.llms.vertex_ai.gemini.transformation import ( + _gemini_convert_messages_with_history, + ) + + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "What's in this video?"}, + { + "type": "file", + "file": { + "file_id": "https://example.com/video.mp4", + "format": "video/mp4", + "detail": "low" + } + } + ] + } + ] + + contents = _gemini_convert_messages_with_history( + messages=messages, model="gemini-3-pro-preview" + ) + + # Verify media_resolution is set for file objects + assert len(contents) == 1 + assert len(contents[0]["parts"]) == 2 # text + file + + # Find the file part + file_part = None + for part in contents[0]["parts"]: + if "file_data" in part: + file_part = part + break + + assert file_part is not None, "File part should exist" + assert "media_resolution" in file_part, "media_resolution should be set for file objects" + assert file_part["media_resolution"] == {"level": "MEDIA_RESOLUTION_LOW"} + + +def test_video_metadata_fps(): + """Test fps parameter in video_metadata (Issue #19026)""" + from litellm.llms.vertex_ai.gemini.transformation import ( + _gemini_convert_messages_with_history, + ) + + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Analyze this video"}, + { + "type": "file", + "file": { + "file_id": "gs://bucket/video.mp4", + "format": "video/mp4", + "video_metadata": {"fps": 5} + } + } + ] + } + ] + + contents = _gemini_convert_messages_with_history( + messages=messages, model="gemini-3-pro-preview" + ) + + # Find the file part + file_part = None + for part in contents[0]["parts"]: + if "file_data" in part: + file_part = part + break + + assert file_part is not None + assert "video_metadata" in file_part, "video_metadata should be present" + assert file_part["video_metadata"]["fps"] == 5 + + +def test_video_metadata_complete(): + """Test all video_metadata fields: fps, start_offset, end_offset (Issue #19026)""" + from litellm.llms.vertex_ai.gemini.transformation import ( + _gemini_convert_messages_with_history, + ) + + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Analyze this video clip"}, + { + "type": "file", + "file": { + "file_id": "gs://bucket/video.mp4", + "format": "video/mp4", + "video_metadata": { + "start_offset": "10s", + "end_offset": "60s", + "fps": 5 + } + } + } + ] + } + ] + + contents = _gemini_convert_messages_with_history( + messages=messages, model="gemini-3-pro-preview" + ) + + # Find the file part + file_part = None + for part in contents[0]["parts"]: + if "file_data" in part: + file_part = part + break + + assert file_part is not None + assert "video_metadata" in file_part + + # Verify field name conversion: snake_case -> camelCase + vm = file_part["video_metadata"] + assert vm["startOffset"] == "10s", "start_offset should be converted to startOffset" + assert vm["endOffset"] == "60s", "end_offset should be converted to endOffset" + assert vm["fps"] == 5, "fps should remain unchanged" + + +def test_detail_and_video_metadata_combined(): + """Test using both detail and video_metadata together (Issue #19026)""" + from litellm.llms.vertex_ai.gemini.transformation import ( + _gemini_convert_messages_with_history, + ) + + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Analyze video"}, + { + "type": "file", + "file": { + "file_id": "https://example.com/video.mp4", + "format": "video/mp4", + "detail": "high", + "video_metadata": {"fps": 10} + } + } + ] + } + ] + + contents = _gemini_convert_messages_with_history( + messages=messages, model="gemini-3-pro-preview" + ) + + # Find the file part + file_part = None + for part in contents[0]["parts"]: + if "file_data" in part: + file_part = part + break + + assert file_part is not None + assert "media_resolution" in file_part + assert file_part["media_resolution"] == {"level": "MEDIA_RESOLUTION_HIGH"} + assert "video_metadata" in file_part + assert file_part["video_metadata"]["fps"] == 10 + + +def test_new_detail_levels(): + """Test new detail levels: medium and ultra_high (Issue #19026)""" + from litellm.llms.vertex_ai.gemini.transformation import ( + _convert_detail_to_media_resolution_enum, + _gemini_convert_messages_with_history, + ) + + # Test mapping function + assert _convert_detail_to_media_resolution_enum("low") == {"level": "MEDIA_RESOLUTION_LOW"} + assert _convert_detail_to_media_resolution_enum("medium") == {"level": "MEDIA_RESOLUTION_MEDIUM"} + assert _convert_detail_to_media_resolution_enum("high") == {"level": "MEDIA_RESOLUTION_HIGH"} + assert _convert_detail_to_media_resolution_enum("ultra_high") == {"level": "MEDIA_RESOLUTION_ULTRA_HIGH"} + + # Test with actual message transformation + messages = [ + { + "role": "user", + "content": [ + { + "type": "file", + "file": { + "file_id": "https://example.com/video.mp4", + "format": "video/mp4", + "detail": "medium" + } + } + ] + } + ] + + contents = _gemini_convert_messages_with_history( + messages=messages, model="gemini-3-pro-preview" + ) + + file_part = None + for part in contents[0]["parts"]: + if "file_data" in part: + file_part = part + break + + assert file_part is not None + assert file_part["media_resolution"] == {"level": "MEDIA_RESOLUTION_MEDIUM"} + + +def test_video_metadata_only_for_gemini_3(): + """Test that video_metadata is only applied for Gemini 3+ models (Issue #19026)""" + from litellm.llms.vertex_ai.gemini.transformation import ( + _gemini_convert_messages_with_history, + ) + + messages = [ + { + "role": "user", + "content": [ + { + "type": "file", + "file": { + "file_id": "https://example.com/video.mp4", + "format": "video/mp4", + "detail": "high", + "video_metadata": {"fps": 5} + } + } + ] + } + ] + + # Test with Gemini 1.5 (should not have video_metadata or media_resolution) + contents_1_5 = _gemini_convert_messages_with_history( + messages=messages, model="gemini-1.5-pro" + ) + + file_part_1_5 = None + for part in contents_1_5[0]["parts"]: + if "file_data" in part: + file_part_1_5 = part + break + + assert file_part_1_5 is not None + assert "media_resolution" not in file_part_1_5, "Gemini 1.5 should not have media_resolution" + assert "video_metadata" not in file_part_1_5, "Gemini 1.5 should not have video_metadata" + + # Test with Gemini 3 (should have both) + contents_3 = _gemini_convert_messages_with_history( + messages=messages, model="gemini-3-pro-preview" + ) + + file_part_3 = None + for part in contents_3[0]["parts"]: + if "file_data" in part: + file_part_3 = part + break + + assert file_part_3 is not None + assert "media_resolution" in file_part_3, "Gemini 3 should have media_resolution" + assert "video_metadata" in file_part_3, "Gemini 3 should have video_metadata" diff --git a/tests/test_litellm/llms/vertex_ai/test_vertex.py b/tests/test_litellm/llms/vertex_ai/test_vertex.py index 39ed09f81be..fdba86af4a7 100644 --- a/tests/test_litellm/llms/vertex_ai/test_vertex.py +++ b/tests/test_litellm/llms/vertex_ai/test_vertex.py @@ -19,7 +19,7 @@ import pytest import litellm from litellm import get_optional_params -from litellm.llms.vertex_ai.gemini.transformation import _process_gemini_image +from litellm.llms.vertex_ai.gemini.transformation import _process_gemini_media from litellm.types.llms.vertex_ai import BlobType @@ -1191,46 +1191,46 @@ def test_logprobs(): assert resp.choices[0].logprobs is not None -def test_process_gemini_image(): - """Test the _process_gemini_image function for different image sources""" - from litellm.llms.vertex_ai.gemini.transformation import _process_gemini_image +def test_process_gemini_media(): + """Test the _process_gemini_media function for different image sources""" + from litellm.llms.vertex_ai.gemini.transformation import _process_gemini_media from litellm.types.llms.vertex_ai import FileDataType # Test GCS URI - gcs_result = _process_gemini_image("gs://bucket/image.png") + gcs_result = _process_gemini_media("gs://bucket/image.png") assert gcs_result["file_data"] == FileDataType( mime_type="image/png", file_uri="gs://bucket/image.png" ) # Test gs url with format specified - gcs_result = _process_gemini_image("gs://bucket/image", format="image/jpeg") + gcs_result = _process_gemini_media("gs://bucket/image", format="image/jpeg") assert gcs_result["file_data"] == FileDataType( mime_type="image/jpeg", file_uri="gs://bucket/image" ) # Test HTTPS JPG URL - https_result = _process_gemini_image("https://example.com/image.jpg") + https_result = _process_gemini_media("https://example.com/image.jpg") print("https_result JPG", https_result) assert https_result["file_data"] == FileDataType( mime_type="image/jpeg", file_uri="https://example.com/image.jpg" ) # Test HTTPS PNG URL - https_result = _process_gemini_image("https://example.com/image.png") + https_result = _process_gemini_media("https://example.com/image.png") print("https_result PNG", https_result) assert https_result["file_data"] == FileDataType( mime_type="image/png", file_uri="https://example.com/image.png" ) # Test HTTPS VIDEO URL - https_result = _process_gemini_image("https://cloud-samples-data/video/animals.mp4") + https_result = _process_gemini_media("https://cloud-samples-data/video/animals.mp4") print("https_result PNG", https_result) assert https_result["file_data"] == FileDataType( mime_type="video/mp4", file_uri="https://cloud-samples-data/video/animals.mp4" ) # Test HTTPS PDF URL - https_result = _process_gemini_image("https://cloud-samples-data/pdf/animals.pdf") + https_result = _process_gemini_media("https://cloud-samples-data/pdf/animals.pdf") print("https_result PDF", https_result) assert https_result["file_data"] == FileDataType( mime_type="application/pdf", @@ -1239,7 +1239,7 @@ def test_process_gemini_image(): # Test base64 image base64_image = "data:image/jpeg;base64,/9j/4AAQSkZJRg..." - base64_result = _process_gemini_image(base64_image) + base64_result = _process_gemini_media(base64_image) print("base64_result", base64_result) assert base64_result["inline_data"]["mime_type"] == "image/jpeg" assert base64_result["inline_data"]["data"] == "/9j/4AAQSkZJRg..." @@ -1368,11 +1368,11 @@ def mock_blob(): "http://subdomain.domain.com/path/to/image.png", ], ) -def test_process_gemini_image_http_url( +def test_process_gemini_media_http_url( http_url: str, mock_convert_url_to_base64: Mock, mock_blob: Mock ) -> None: """ - Test that _process_gemini_image correctly handles HTTP URLs. + Test that _process_gemini_media correctly handles HTTP URLs. Args: http_url: Test HTTP URL @@ -1384,7 +1384,7 @@ def test_process_gemini_image_http_url( expected_image_data = "data:image/jpeg;base64,/9j/4AAQSkZJRg..." mock_convert_url_to_base64.return_value = expected_image_data # Act - result = _process_gemini_image(http_url) + result = _process_gemini_media(http_url) # assert result["file_data"]["file_uri"] == http_url diff --git a/tests/test_litellm/proxy/common_utils/test_key_rotation_integration.py b/tests/test_litellm/proxy/common_utils/test_key_rotation_integration.py new file mode 100644 index 00000000000..308c8cdbce1 --- /dev/null +++ b/tests/test_litellm/proxy/common_utils/test_key_rotation_integration.py @@ -0,0 +1,144 @@ +""" +Regression test for AWS Secrets Manager Auto-Rotation Bug Fix + +This test verifies that KeyRotationManager correctly passes key_alias +when calling regenerate_key_fn, ensuring the secret is rotated at the +correct location in AWS Secrets Manager. + +Bug Fixed: Key alias was not passed during auto-rotation, causing +secrets to be created at a new location instead of updating in-place. +""" +import os +import sys +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +sys.path.insert(0, os.path.abspath("../../../..")) + +from litellm.proxy._types import ( + GenerateKeyResponse, + LiteLLM_VerificationToken, + RegenerateKeyRequest, +) +from litellm.proxy.common_utils.key_rotation_manager import KeyRotationManager + + +class TestKeyRotationManagerPassesKeyAlias: + """ + Regression tests to ensure KeyRotationManager passes key_alias + to regenerate_key_fn during auto-rotation. + """ + + @pytest.mark.asyncio + async def test_rotate_key_passes_key_alias_to_regenerate_request(self): + """ + Verify that _rotate_key includes key_alias in the RegenerateKeyRequest. + + This is the core fix: previously, key_alias was NOT passed, causing + the secret manager hook to use a generated name instead of the alias. + """ + # Create a mock key with an alias + test_alias = "tenant1/my-important-key" + test_token = "sk-test-token-hash-12345" + + mock_key = MagicMock(spec=LiteLLM_VerificationToken) + mock_key.token = test_token + mock_key.key_alias = test_alias + mock_key.key_name = "sk-...1234" + mock_key.rotation_interval = "30d" + mock_key.rotation_count = 0 + + # Create mock prisma client + mock_prisma = MagicMock() + mock_prisma.db.litellm_verificationtoken.update = AsyncMock( + return_value=mock_key + ) + + # Create mock response + mock_response = GenerateKeyResponse( + key="sk-new-key-value", + token_id="new-token-hash", + key_alias=test_alias, + ) + + # Capture the RegenerateKeyRequest passed to regenerate_key_fn + captured_request = None + + async def capture_regenerate_key_fn( + data, user_api_key_dict, litellm_changed_by + ): + nonlocal captured_request + captured_request = data + return mock_response + + # Patch regenerate_key_fn to capture the request + with patch( + "litellm.proxy.common_utils.key_rotation_manager.regenerate_key_fn", + side_effect=capture_regenerate_key_fn, + ): + with patch( + "litellm.proxy.common_utils.key_rotation_manager.KeyManagementEventHooks.async_key_rotated_hook", + new_callable=AsyncMock, + ): + rotation_manager = KeyRotationManager(mock_prisma) + await rotation_manager._rotate_key(mock_key) + + # CRITICAL ASSERTION: key_alias must be passed + assert captured_request is not None, "regenerate_key_fn should have been called" + assert isinstance(captured_request, RegenerateKeyRequest) + assert captured_request.key == test_token, "Token should be passed correctly" + assert captured_request.key_alias == test_alias, ( + f"key_alias should be '{test_alias}' but was '{captured_request.key_alias}'. " + "This is the bug we fixed - key_alias was not being passed!" + ) + + @pytest.mark.asyncio + async def test_rotate_key_passes_none_alias_when_key_has_no_alias(self): + """ + Verify that _rotate_key handles keys without an alias gracefully. + """ + test_token = "sk-test-token-hash-67890" + + mock_key = MagicMock(spec=LiteLLM_VerificationToken) + mock_key.token = test_token + mock_key.key_alias = None # No alias set + mock_key.key_name = "sk-...5678" + mock_key.rotation_interval = "30d" + mock_key.rotation_count = 0 + + mock_prisma = MagicMock() + mock_prisma.db.litellm_verificationtoken.update = AsyncMock( + return_value=mock_key + ) + + mock_response = GenerateKeyResponse( + key="sk-new-key-value", + token_id="new-token-hash", + ) + + captured_request = None + + async def capture_regenerate_key_fn( + data, user_api_key_dict, litellm_changed_by + ): + nonlocal captured_request + captured_request = data + return mock_response + + with patch( + "litellm.proxy.common_utils.key_rotation_manager.regenerate_key_fn", + side_effect=capture_regenerate_key_fn, + ): + with patch( + "litellm.proxy.common_utils.key_rotation_manager.KeyManagementEventHooks.async_key_rotated_hook", + new_callable=AsyncMock, + ): + rotation_manager = KeyRotationManager(mock_prisma) + await rotation_manager._rotate_key(mock_key) + + assert captured_request is not None + assert captured_request.key == test_token + assert ( + captured_request.key_alias is None + ), "key_alias should be None for keys without alias" diff --git a/tests/test_litellm/proxy/prompts/test_prompt_endpoints_crud.py b/tests/test_litellm/proxy/prompts/test_prompt_endpoints_crud.py new file mode 100644 index 00000000000..2c5bc1bf87d --- /dev/null +++ b/tests/test_litellm/proxy/prompts/test_prompt_endpoints_crud.py @@ -0,0 +1,189 @@ +import pytest +from unittest.mock import MagicMock, AsyncMock, patch +from litellm.proxy._types import UserAPIKeyAuth, LitellmUserRoles +from litellm.types.prompts.init_prompts import ( + PromptSpec, + PromptLiteLLMParams, + PromptInfo, +) + + +@pytest.mark.asyncio +async def test_delete_prompt_success(): + """ + Test that delete_prompt correctly identifies the base prompt ID + and deletes all versions from DB and memory. + """ + from litellm.proxy.prompts.prompt_endpoints import delete_prompt + + # Mock user auth + mock_user_auth = UserAPIKeyAuth( + api_key="sk-1234", user_role=LitellmUserRoles.PROXY_ADMIN + ) + + # Mock DB Client + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_prompttable.delete_many = AsyncMock(return_value=None) + + # Mock In-Memory Registry + with patch( + "litellm.proxy.prompts.prompt_registry.IN_MEMORY_PROMPT_REGISTRY" + ) as mock_registry: + # User passes "test_prompt.v2" + # We simulate that get_prompt_by_id returns the prompt spec for v2 + prompt_spec = PromptSpec( + prompt_id="test_prompt.v2", + litellm_params=PromptLiteLLMParams( + prompt_id="test_prompt", prompt_integration="dotprompt" + ), + prompt_info=PromptInfo(prompt_type="db"), + ) + mock_registry.get_prompt_by_id.return_value = prompt_spec + + # Patch the prisma client in the endpoint module + with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client): + response = await delete_prompt( + prompt_id="test_prompt.v2", user_api_key_dict=mock_user_auth + ) + + # Assertions + expected_base_id = "test_prompt" + + # 1. DB deletion should use base ID + mock_prisma_client.db.litellm_prompttable.delete_many.assert_called_once_with( + where={"prompt_id": expected_base_id} + ) + + # 2. Memory deletion should use base ID + mock_registry.delete_prompts_by_base_id.assert_called_once_with( + expected_base_id + ) + + assert response == { + "message": f"Prompt {expected_base_id} deleted successfully" + } + + +@pytest.mark.asyncio +async def test_delete_prompt_by_base_id_success(): + """ + Test that delete_prompt works when passed a base ID directly, + finding the latest version to confirm existence, then deleting. + """ + from litellm.proxy.prompts.prompt_endpoints import delete_prompt + + # Mock user auth + mock_user_auth = UserAPIKeyAuth( + api_key="sk-1234", user_role=LitellmUserRoles.PROXY_ADMIN + ) + + # Mock DB Client + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_prompttable.delete_many = AsyncMock(return_value=None) + + # Mock In-Memory Registry + with patch( + "litellm.proxy.prompts.prompt_registry.IN_MEMORY_PROMPT_REGISTRY" + ) as mock_registry: + # User passes "test_prompt" (base ID) + # 1. get_prompt_by_id("test_prompt") -> None (if it's not registered as base) + # 2. It calls get_latest_version_prompt_id -> returns "test_prompt.v3" + # 3. get_prompt_by_id("test_prompt.v3") -> returns Spec + + # Setup mocks behavior + def get_prompt_side_effect(prompt_id): + if prompt_id == "test_prompt": + return None + if prompt_id == "test_prompt.v3": + return PromptSpec( + prompt_id="test_prompt.v3", + litellm_params=PromptLiteLLMParams( + prompt_id="test_prompt", prompt_integration="dotprompt" + ), + prompt_info=PromptInfo(prompt_type="db"), + ) + return None + + mock_registry.get_prompt_by_id.side_effect = get_prompt_side_effect + mock_registry.IN_MEMORY_PROMPTS = { + "test_prompt.v1": {}, + "test_prompt.v2": {}, + "test_prompt.v3": {}, + } + + # Patch the prisma client in the endpoint module + with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client): + response = await delete_prompt( + prompt_id="test_prompt", user_api_key_dict=mock_user_auth + ) + + # Assertions + expected_base_id = "test_prompt" + + # 1. DB deletion should use base ID + mock_prisma_client.db.litellm_prompttable.delete_many.assert_called_once_with( + where={"prompt_id": expected_base_id} + ) + + # 2. Memory deletion should use base ID + mock_registry.delete_prompts_by_base_id.assert_called_once_with( + expected_base_id + ) + + assert response == { + "message": f"Prompt {expected_base_id} deleted successfully" + } + + +@pytest.mark.asyncio +async def test_get_prompt_info_by_base_id(): + """ + Test that get_prompt_info correctly resolves a base ID to the latest version. + """ + from litellm.proxy.prompts.prompt_endpoints import get_prompt_info + + # Mock user auth + mock_user_auth = UserAPIKeyAuth( + api_key="sk-1234", user_role=LitellmUserRoles.PROXY_ADMIN + ) + + # Mock In-Memory Registry + with patch( + "litellm.proxy.prompts.prompt_registry.IN_MEMORY_PROMPT_REGISTRY" + ) as mock_registry: + # Setup mocks behavior + prompt_spec_v3 = PromptSpec( + prompt_id="test_prompt.v3", + litellm_params=PromptLiteLLMParams( + prompt_id="test_prompt", prompt_integration="dotprompt" + ), + prompt_info=PromptInfo(prompt_type="db"), + ) + + # When get_prompt_by_id is called with "test_prompt", return None (so it searches versions) + # When called with "test_prompt.v3", return the spec + def get_prompt_side_effect(prompt_id): + if prompt_id == "test_prompt": + return None + if prompt_id == "test_prompt.v3": + return prompt_spec_v3 + return None + + mock_registry.get_prompt_by_id.side_effect = get_prompt_side_effect + mock_registry.IN_MEMORY_PROMPTS = { + "test_prompt.v1": {}, + "test_prompt.v2": {}, + "test_prompt.v3": {}, + } + + # We also need to mock get_prompt_callback_by_id to avoid content extraction errors/logic + mock_registry.get_prompt_callback_by_id.return_value = None + + response = await get_prompt_info( + prompt_id="test_prompt", user_api_key_dict=mock_user_auth + ) + + assert ( + response.prompt_spec.prompt_id == "test_prompt" + ) # Should return base ID in spec response + assert response.prompt_spec.version == 3 # Should identify it as version 3