scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-10 22:41:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 83

Merge pull request #19461 from BerriAI/litellm_staging_01_20_2026

Browse source 
Litellm staging 01 20 2026
This commit is contained in:
Sameer Kankute 2026-01-21 18:31:14 +05:30 • committed by GitHub
commit b0499841b0
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
18 changed files with 1452 additions and 247 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

196
docs/my-website/docs/providers/gemini.md
View file

@ -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:**
<Tabs>
<TabItem value="images" label="Images">
```python
from litellm import completion
@ -1604,10 +1609,193 @@ response = completion(
)
```
</TabItem>
<TabItem value="videos" label="Videos with Files">
```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,
)
```
</TabItem>
</Tabs>
:::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:**
<Tabs>
<TabItem value="basic" label="Basic Video Metadata">
```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)
```
</TabItem>
<TabItem value="combined" label="Combined with Detail">
```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)
```
</TabItem>
<TabItem value="proxy" label="PROXY">
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 <YOUR-LITELLM-KEY>" \
-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"
}
}
}
]
}
]
}'
```
</TabItem>
</Tabs>
## Sample Usage
```python
import os

238
docs/my-website/docs/providers/vertex.md
View file

@ -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:**
<Tabs>
<TabItem value="images" label="Images">
```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,
)
```
</TabItem>
<TabItem value="videos" label="Videos with Files">
```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,
)
```
</TabItem>
</Tabs>
:::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:**
<Tabs>
<TabItem value="basic" label="Basic Video Metadata">
```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)
```
</TabItem>
<TabItem value="combined" label="Combined with Detail">
```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)
```
</TabItem>
<TabItem value="proxy" label="PROXY">
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 <YOUR-LITELLM-KEY>" \
-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"
}
}
}
]
}
]
}'
```
</TabItem>
</Tabs>
## Usage - PDF / Videos / Audio etc. Files

2
litellm/integrations/opentelemetry.py
View file

@ -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

105
litellm/llms/vertex_ai/gemini/transformation.py
View file

@ -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:

43
litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py
View file

@ -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(

12
litellm/model_prices_and_context_window_backup.json
View file

@ -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",

101
litellm/proxy/common_utils/key_rotation_manager.py
View file

@ -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,
)

3
litellm/proxy/hooks/key_management_event_hooks.py
View file

@ -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,
)

182
litellm/proxy/prompts/prompt_endpoints.py
View file

@ -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

33
litellm/proxy/prompts/prompt_registry.py
View file

@ -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()
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()

66
litellm/proxy/proxy_server.py
View file

@ -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

2
litellm/types/llms/openai.py
View file

@ -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):

12
model_prices_and_context_window.json
View file

@ -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",

8
tests/test_litellm/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py
View file

@ -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

335
tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py
View file

@ -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"

28
tests/test_litellm/llms/vertex_ai/test_vertex.py
View file

@ -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

144
tests/test_litellm/proxy/common_utils/test_key_rotation_integration.py Normal file
View file

@ -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"

189
tests/test_litellm/proxy/prompts/test_prompt_endpoints_crud.py Normal file
View file

@ -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