From 75ee0d126c957fd226bba259ee62962cc31cce2c Mon Sep 17 00:00:00 2001 From: Harshit Jain <48647625+Harshit28j@users.noreply.github.com> Date: Tue, 20 Jan 2026 23:23:16 +0530 Subject: [PATCH 01/10] Fix/prisma schema permission (#19391) * fix: add prisma permission issue * Add test case for prisma generate --- .../litellm_proxy_extras/utils.py | 95 +++++++++++---- litellm/proxy/db/prisma_client.py | 38 +++++- litellm/proxy/prisma_migration.py | 27 +++-- litellm/proxy/proxy_cli.py | 6 +- .../proxy/test_migration_failure_handling.py | 114 ++++++++++++++++++ 5 files changed, 238 insertions(+), 42 deletions(-) create mode 100644 tests/test_litellm/proxy/test_migration_failure_handling.py diff --git a/litellm-proxy-extras/litellm_proxy_extras/utils.py b/litellm-proxy-extras/litellm_proxy_extras/utils.py index 7ffbe95be13..1aed555c5a9 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/utils.py +++ b/litellm-proxy-extras/litellm_proxy_extras/utils.py @@ -18,14 +18,15 @@ def str_to_bool(value: Optional[str]) -> bool: return value.lower() in ("true", "1", "t", "y", "yes") - def _get_prisma_env() -> dict: """Get environment variables for Prisma, handling offline mode if configured.""" prisma_env = os.environ.copy() if str_to_bool(os.getenv("PRISMA_OFFLINE_MODE")): # These env vars prevent Prisma from attempting downloads prisma_env["NPM_CONFIG_PREFER_OFFLINE"] = "true" - prisma_env["NPM_CONFIG_CACHE"] = os.getenv("NPM_CONFIG_CACHE", "/app/.cache/npm") + prisma_env["NPM_CONFIG_CACHE"] = os.getenv( + "NPM_CONFIG_CACHE", "/app/.cache/npm" + ) return prisma_env @@ -34,29 +35,28 @@ def _get_prisma_command() -> str: if str_to_bool(os.getenv("PRISMA_OFFLINE_MODE")): # Primary location where Prisma Python package installs the CLI default_cli_path = "/app/.cache/prisma-python/binaries/node_modules/.bin/prisma" - + # Check if custom path is provided (for flexibility) custom_cli_path = os.getenv("PRISMA_CLI_PATH") if custom_cli_path and os.path.exists(custom_cli_path): logger.info(f"Using custom Prisma CLI at {custom_cli_path}") return custom_cli_path - + # Check the default location if os.path.exists(default_cli_path): logger.info(f"Using cached Prisma CLI at {default_cli_path}") return default_cli_path - + # If not found, log warning and fall back logger.warning( f"Prisma CLI not found at {default_cli_path}. " "Falling back to Python wrapper (may attempt downloads)" ) - + # Fall back to the Python wrapper (will work in online mode) return "prisma" - class ProxyExtrasDBManager: @staticmethod def _get_prisma_dir() -> str: @@ -119,7 +119,7 @@ class ProxyExtrasDBManager: stdout=open(migration_file, "w"), check=True, timeout=30, - env=prisma_env + env=prisma_env, ) # 3. Mark the migration as applied since it represents current state @@ -134,7 +134,7 @@ class ProxyExtrasDBManager: ], check=True, timeout=30, - env=prisma_env + env=prisma_env, ) return True @@ -159,14 +159,20 @@ class ProxyExtrasDBManager: @staticmethod def _roll_back_migration(migration_name: str): """Mark a specific migration as rolled back""" - # Set up environment for offline mode if configured + # Set up environment for offline mode if configured prisma_env = _get_prisma_env() subprocess.run( - [_get_prisma_command(), "migrate", "resolve", "--rolled-back", migration_name], + [ + _get_prisma_command(), + "migrate", + "resolve", + "--rolled-back", + migration_name, + ], timeout=60, check=True, capture_output=True, - env=prisma_env + env=prisma_env, ) @staticmethod @@ -178,7 +184,7 @@ class ProxyExtrasDBManager: timeout=60, check=True, capture_output=True, - env=prisma_env + env=prisma_env, ) @staticmethod @@ -248,7 +254,7 @@ class ProxyExtrasDBManager: if not database_url: logger.error("DATABASE_URL not set") return - + diff_dir = ( Path(migrations_dir) / "migrations" @@ -283,7 +289,7 @@ class ProxyExtrasDBManager: check=True, timeout=60, stdout=f, - env=_get_prisma_env() + env=_get_prisma_env(), ) except subprocess.CalledProcessError as e: logger.warning(f"Failed to generate migration diff: {e.stderr}") @@ -313,7 +319,7 @@ class ProxyExtrasDBManager: check=True, capture_output=True, text=True, - env=_get_prisma_env() + env=_get_prisma_env(), ) logger.info(f"prisma db execute stdout: {result.stdout}") logger.info("✅ Migration diff applied successfully") @@ -331,12 +337,18 @@ class ProxyExtrasDBManager: try: logger.info(f"Resolving migration: {migration_name}") subprocess.run( - [_get_prisma_command(), "migrate", "resolve", "--applied", migration_name], + [ + _get_prisma_command(), + "migrate", + "resolve", + "--applied", + migration_name, + ], timeout=60, check=True, capture_output=True, text=True, - env=_get_prisma_env() + env=_get_prisma_env(), ) logger.debug(f"Resolved migration: {migration_name}") except subprocess.CalledProcessError as e: @@ -375,7 +387,7 @@ class ProxyExtrasDBManager: check=True, capture_output=True, text=True, - env=_get_prisma_env() + env=_get_prisma_env(), ) logger.info(f"prisma migrate deploy stdout: {result.stdout}") @@ -413,7 +425,7 @@ class ProxyExtrasDBManager: check=True, capture_output=True, text=True, - env=_get_prisma_env() + env=_get_prisma_env(), ) logger.info( f"✅ Migration {failed_migration} marked as rolled back... retrying" @@ -509,12 +521,43 @@ class ProxyExtrasDBManager: raise else: # Use prisma db push with increased timeout - subprocess.run( - [_get_prisma_command(), "db", "push", "--accept-data-loss"], - timeout=60, - check=True, - ) - return True + try: + subprocess.run( + [_get_prisma_command(), "db", "push", "--accept-data-loss"], + timeout=60, + check=True, + capture_output=True, # capture output to check for errors + text=True, + env=_get_prisma_env(), + ) + return True + except subprocess.CalledProcessError as e: + if ( + "Permission denied" in e.stderr + and "schema.prisma" in e.stderr + ): + logger.warning( + f"Permission denied during prisma generate: {e.stderr}. Retrying with --skip-generate..." + ) + # Retry with --skip-generate + subprocess.run( + [ + _get_prisma_command(), + "db", + "push", + "--accept-data-loss", + "--skip-generate", + ], + timeout=60, + check=True, + capture_output=True, + text=True, + env=_get_prisma_env(), + ) + logger.info("✅ prisma db push --skip-generate completed") + return True + else: + raise e except subprocess.TimeoutExpired: logger.info(f"Attempt {attempt + 1} timed out") time.sleep(random.randrange(5, 15)) diff --git a/litellm/proxy/db/prisma_client.py b/litellm/proxy/db/prisma_client.py index c9c0cfe8f68..95800e96589 100644 --- a/litellm/proxy/db/prisma_client.py +++ b/litellm/proxy/db/prisma_client.py @@ -386,11 +386,39 @@ class PrismaManager: return ProxyExtrasDBManager.setup_database(use_migrate=use_migrate) else: # Use prisma db push with increased timeout - subprocess.run( - ["prisma", "db", "push", "--accept-data-loss"], - timeout=60, - check=True, - ) + try: + subprocess.run( + ["prisma", "db", "push", "--accept-data-loss"], + timeout=60, + check=True, + capture_output=True, + text=True, + ) + except subprocess.CalledProcessError as e: + if ( + "Permission denied" in e.stderr + and "schema.prisma" in e.stderr + ): + verbose_proxy_logger.warning( + f"Permission denied during prisma generate: {e.stderr}. Retrying with --skip-generate..." + ) + # Retry with --skip-generate + subprocess.run( + [ + "prisma", + "db", + "push", + "--accept-data-loss", + "--skip-generate", + ], + timeout=60, + check=True, + capture_output=True, + text=True, + ) + return True + else: + raise e return True except subprocess.TimeoutExpired: verbose_proxy_logger.warning(f"Attempt {attempt + 1} timed out") diff --git a/litellm/proxy/prisma_migration.py b/litellm/proxy/prisma_migration.py index 251d1e56287..2fe12b1439c 100644 --- a/litellm/proxy/prisma_migration.py +++ b/litellm/proxy/prisma_migration.py @@ -15,14 +15,21 @@ from litellm.proxy.proxy_cli import run_server # Call the Click command with standalone_mode=False run_server(["--skip_server_startup"], standalone_mode=False) -# run prisma generate +# Run prisma generate verbose_proxy_logger.info("Running 'prisma generate'...") -result = subprocess.run(["prisma", "generate"], capture_output=True, text=True) -verbose_proxy_logger.info(f"'prisma generate' stdout: {result.stdout}") # Log stdout -exit_code = result.returncode - -if exit_code != 0: - verbose_proxy_logger.info(f"'prisma generate' failed with exit code {exit_code}.") - verbose_proxy_logger.error( - f"'prisma generate' stderr: {result.stderr}" - ) # Log stderr +try: + result = subprocess.run(["prisma", "generate"], capture_output=True, text=True) + if result.returncode != 0: + if "Permission denied" in result.stderr: + verbose_proxy_logger.warning( + f"Permission denied during 'prisma generate'. Skipping generation, assuming client is pre-generated. Error: {result.stderr}" + ) + else: + verbose_proxy_logger.info( + f"'prisma generate' failed with exit code {result.returncode}." + ) + verbose_proxy_logger.error( + f"'prisma generate' stderr: {result.stderr}" + ) # Log stderr +except Exception as e: + verbose_proxy_logger.error(f"Error running prisma generate: {e}") diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py index 2059246674b..028aad9a2e1 100644 --- a/litellm/proxy/proxy_cli.py +++ b/litellm/proxy/proxy_cli.py @@ -797,7 +797,11 @@ def run_server( # noqa: PLR0915 ): check_prisma_schema_diff(db_url=None) else: - PrismaManager.setup_database(use_migrate=not use_prisma_db_push) + if not PrismaManager.setup_database( + use_migrate=not use_prisma_db_push + ): + print("LiteLLM: Database setup failed. Exiting...") # noqa + sys.exit(1) else: print( # noqa f"Unable to connect to DB. DATABASE_URL found in environment, but prisma package not found." # noqa diff --git a/tests/test_litellm/proxy/test_migration_failure_handling.py b/tests/test_litellm/proxy/test_migration_failure_handling.py new file mode 100644 index 00000000000..426049007b1 --- /dev/null +++ b/tests/test_litellm/proxy/test_migration_failure_handling.py @@ -0,0 +1,114 @@ +import sys +import os +import subprocess +from unittest.mock import MagicMock, patch +from click.testing import CliRunner + +# Add parent directory to path to allow importing litellm +sys.path.insert(0, os.path.abspath("../../..")) + +from litellm.proxy.db.prisma_client import PrismaManager +from litellm.proxy.proxy_cli import run_server + + +class TestMigrationFailureHandling: + @patch("subprocess.run") + def test_prisma_client_permission_error_retry(self, mock_subprocess_run): + """ + Regression Test: Verifies that PrismaManager.setup_database + catches PermissionError during 'prisma db push' and retries with '--skip-generate'. + """ + # Mock behavior: + # call 1: raises CalledProcessError with "Permission denied" and "schema.prisma" + # call 2 (retry): succeeds + + error_output = "Error: Permission denied writing to ... schema.prisma" + + mock_process_error = subprocess.CalledProcessError( + returncode=1, cmd=["prisma", "db", "push"], stderr=error_output + ) + + mock_subprocess_run.side_effect = [ + mock_process_error, # 1st attempt fails with permission error + MagicMock(returncode=0), # 2nd attempt (retry) succeeds + ] + + # Ensure we run the 'db push' path (use_migrate=False) + # We also need to mock should_update_prisma_schema to return True + + with patch( + "litellm.proxy.db.prisma_client.should_update_prisma_schema", + return_value=True, + ): + # Run setup_database with use_migrate=False to trigger 'prisma db push' path + result = PrismaManager.setup_database(use_migrate=False) + + # Assert success + assert result is True + + # Verify calls + assert mock_subprocess_run.call_count == 2 + + # Check 1st call arguments (standard push) + args1, _ = mock_subprocess_run.call_args_list[0] + assert "push" in args1[0] + assert "--skip-generate" not in args1[0] + + # Check 2nd call arguments (retry with skip-generate) + args2, _ = mock_subprocess_run.call_args_list[1] + assert "push" in args2[0] + assert "--skip-generate" in args2[0] + + def test_proxy_cli_exit_on_migration_fail(self): + """ + Regression Test: Verifies that proxy_cli.run_server exits with NON-ZERO status + if PrismaManager.setup_database returns False. + """ + runner = CliRunner() + + # Mock setup_database to return False (Simulating failure) + # Mock should_update_prisma_schema to return True (Ensure we hit the DB setup logic) + with patch( + "litellm.proxy.db.prisma_client.PrismaManager.setup_database", + return_value=False, + ), patch( + "litellm.proxy.db.prisma_client.should_update_prisma_schema", + return_value=True, + ): + # Mock dependencies to prevent actual server startup and handle imports + mock_app = MagicMock() + mock_proxy_config = MagicMock() + + # Patch sys.modules to prevent ImportErrors for proxy_server + with patch.dict( + "sys.modules", + { + "proxy_server": MagicMock( + app=mock_app, ProxyConfig=mock_proxy_config + ) + }, + ): + with patch( + "litellm.proxy.proxy_cli.ProxyInitializationHelpers._get_default_unvicorn_init_args" + ) as mock_get_args: + mock_get_args.return_value = { + "app": "app", + "host": "localhost", + "port": 8000, + } + + # Set DATABASE_URL to trigger DB logic + with patch.dict( + os.environ, + {"DATABASE_URL": "postgresql://user:pass@localhost:5432/db"}, + ): + # Execute: Run server with --local and --skip_server_startup + result = runner.invoke( + run_server, ["--local", "--skip_server_startup"] + ) + + # Assert: Exit code should be non-zero (failure) + assert ( + result.exit_code != 0 + ), f"Expected non-zero exit code, got {result.exit_code}. Output: {result.output}" + assert "Database setup failed. Exiting..." in result.output From 1c8bf19f1efe35b807768a1373f8e413761f6dd8 Mon Sep 17 00:00:00 2001 From: Harshit Jain <48647625+Harshit28j@users.noreply.github.com> Date: Tue, 20 Jan 2026 23:25:09 +0530 Subject: [PATCH 02/10] fix(proxy_server): pass search_tools to Router during DB-triggered initialization (#19388) --- litellm/proxy/proxy_server.py | 55 ++++++++++++++++++----------------- 1 file changed, 28 insertions(+), 27 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index c3a4de314e5..bbf434aead1 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,11 +3095,12 @@ 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: if llm_router is None and master_key is not None: verbose_proxy_logger.debug(f"len new_models: {len(new_models)}") @@ -3114,6 +3115,7 @@ 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}") @@ -3134,7 +3136,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 +3945,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 +4275,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 +9730,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 From 56bf6001e9c4a624687e6a608ac1e289b4789dc2 Mon Sep 17 00:00:00 2001 From: Kris Xia Date: Wed, 21 Jan 2026 03:36:55 +0800 Subject: [PATCH 03/10] Supports setting media_resolution and fps parameters on each video file, when using Gemini video understanding. (#19273) * feat: add gemini video metadata and detail support Implement support for video_metadata and enhanced detail parameter for Gemini 3.0+ models: - Add video_metadata field to ChatCompletionFileObjectFile type - Supports fps, start_offset, and end_offset parameters - Properly converts snake_case to camelCase for Gemini API - Extend detail parameter to support medium and ultra_high levels - Maps to MEDIA_RESOLUTION_MEDIUM and MEDIA_RESOLUTION_ULTRA_HIGH - Update _process_gemini_image to handle video metadata transformation - Add version gating to only apply features for Gemini 3+ models - Add comprehensive test coverage (6 new test cases) - Test detail parameter with file objects - Test video_metadata fields (fps, start_offset, end_offset) - Test combined detail + video_metadata usage - Test new detail levels (medium, ultra_high) - Test version gating (Gemini 1.5 vs 3.0) Note: video_metadata is only supported for video files but error handling is delegated to Vertex AI for other media types. * refactor: rename _process_gemini_image to _process_gemini_media The function handles multiple media types (images, audio, video, PDF), not just images. Renamed to better reflect its actual purpose. - Update function name in transformation.py - Update all function calls and references - Update test names and imports to match - Improve docstring to clarify it handles all media types * docs: add video metadata and media resolution control documentation Add comprehensive documentation for Gemini 3+ video processing features: - Document media resolution control (detail parameter) for images and videos - Add video_metadata field documentation (fps, start_offset, end_offset) - Include usage examples with tabs for basic, combined, and proxy scenarios - Update both Gemini and Vertex AI provider documentation - Clarify snake_case to camelCase field conversion for Gemini API Signed-off-by: Kris Xia * refactor(gemini): extract metadata application into helper function Extract duplicated Gemini 3+ media_resolution and video_metadata application logic from _process_gemini_media into a dedicated _apply_gemini_3_metadata helper function to improve code maintainability. --------- Signed-off-by: Kris Xia --- docs/my-website/docs/providers/gemini.md | 196 ++++++++++++- docs/my-website/docs/providers/vertex.md | 238 +++++++++++++++ .../llms/vertex_ai/gemini/transformation.py | 105 ++++--- litellm/types/llms/openai.py | 2 + .../test_vertex_ai_gemini_transformation.py | 8 +- ...test_vertex_and_google_ai_studio_gemini.py | 270 ++++++++++++++++++ .../llms/vertex_ai/test_vertex.py | 28 +- 7 files changed, 792 insertions(+), 55 deletions(-) diff --git a/docs/my-website/docs/providers/gemini.md b/docs/my-website/docs/providers/gemini.md index 32dea2069b7..11866e68d15 100644 --- a/docs/my-website/docs/providers/gemini.md +++ b/docs/my-website/docs/providers/gemini.md @@ -1547,16 +1547,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 @@ -1593,10 +1598,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 be2bf86ab10..c47d7d914a8 100644 --- a/docs/my-website/docs/providers/vertex.md +++ b/docs/my-website/docs/providers/vertex.md @@ -1957,6 +1957,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/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/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/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 f44e8640105..810769023bd 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 @@ -2653,3 +2653,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 From 20323feecc8d8d0b955a368b7c8d30a6da995d6c Mon Sep 17 00:00:00 2001 From: Harshit Jain <48647625+Harshit28j@users.noreply.github.com> Date: Wed, 21 Jan 2026 01:58:34 +0530 Subject: [PATCH 04/10] fix(prompts): fix prompt info lookup and delete using correct IDs (#19358) * fix(prompts): fix prompt info lookup and delete using correct IDs * add regression tests cases --- litellm/proxy/prompts/prompt_endpoints.py | 182 +++++++++-------- litellm/proxy/prompts/prompt_registry.py | 33 ++- .../prompts/test_prompt_endpoints_crud.py | 189 ++++++++++++++++++ 3 files changed, 315 insertions(+), 89 deletions(-) create mode 100644 tests/test_litellm/proxy/prompts/test_prompt_endpoints_crud.py 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/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 From b36e704e06298184e44b720e489c7f93a1c5692b Mon Sep 17 00:00:00 2001 From: Harshit Jain <48647625+Harshit28j@users.noreply.github.com> Date: Wed, 21 Jan 2026 08:00:36 +0530 Subject: [PATCH 05/10] fix: ensure auto-rotation updates existing AWS secret instead of creating new one (#19455) --- .../common_utils/key_rotation_manager.py | 101 +++++++----- .../proxy/hooks/key_management_event_hooks.py | 3 +- .../test_key_rotation_integration.py | 144 ++++++++++++++++++ 3 files changed, 207 insertions(+), 41 deletions(-) create mode 100644 tests/test_litellm/proxy/common_utils/test_key_rotation_integration.py 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/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" From b4ed387d24b55b8ff7da0a7f7b77a666e1b89bc4 Mon Sep 17 00:00:00 2001 From: Cesar Garcia <128240629+Chesars@users.noreply.github.com> Date: Wed, 21 Jan 2026 00:31:25 -0300 Subject: [PATCH 06/10] fix(vertex_ai): handle reasoning_effort as dict from OpenAI Agents SDK (#19419) The OpenAI Agents SDK (v0.6.9+) now passes reasoning_effort as a dict when summary is specified: {"effort": "high", "summary": "auto"} This change extracts the "effort" value from the dict for Vertex AI, which only supports thinkingLevel (not summary). Before: reasoning_effort={"effort": "high"} was silently ignored After: reasoning_effort={"effort": "high"} correctly maps to thinkingLevel Fixes #19411 --- .../vertex_and_google_ai_studio_gemini.py | 43 +++++++----- ...test_vertex_and_google_ai_studio_gemini.py | 65 +++++++++++++++++++ 2 files changed, 91 insertions(+), 17 deletions(-) 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 f65a19ac46f..46ea78ac9ef 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 @@ -988,25 +988,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/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 810769023bd..fe262b8d840 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 @@ -1887,6 +1887,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 ( From b810a68f89971647502c852d1e71986770e57453 Mon Sep 17 00:00:00 2001 From: Connor Luebbehusen Date: Wed, 21 Jan 2026 05:58:38 -0500 Subject: [PATCH 07/10] fix: correct gemini-2.5-flash-lite audio input and cache read pricing --- litellm/model_prices_and_context_window_backup.json | 12 ++++++------ model_prices_and_context_window.json | 12 ++++++------ 2 files changed, 12 insertions(+), 12 deletions(-) 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/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", From d879dcdcef6e1baebde3b1cd3b8834cb1b8da217 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Wed, 21 Jan 2026 17:19:32 +0530 Subject: [PATCH 08/10] Revert "Fix/prisma schema permission (#19391)" This reverts commit 75ee0d126c957fd226bba259ee62962cc31cce2c. --- .../litellm_proxy_extras/utils.py | 95 ++++----------- litellm/proxy/db/prisma_client.py | 38 +----- litellm/proxy/prisma_migration.py | 27 ++--- litellm/proxy/proxy_cli.py | 6 +- .../proxy/test_migration_failure_handling.py | 114 ------------------ 5 files changed, 42 insertions(+), 238 deletions(-) delete mode 100644 tests/test_litellm/proxy/test_migration_failure_handling.py diff --git a/litellm-proxy-extras/litellm_proxy_extras/utils.py b/litellm-proxy-extras/litellm_proxy_extras/utils.py index 1aed555c5a9..7ffbe95be13 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/utils.py +++ b/litellm-proxy-extras/litellm_proxy_extras/utils.py @@ -18,15 +18,14 @@ def str_to_bool(value: Optional[str]) -> bool: return value.lower() in ("true", "1", "t", "y", "yes") + def _get_prisma_env() -> dict: """Get environment variables for Prisma, handling offline mode if configured.""" prisma_env = os.environ.copy() if str_to_bool(os.getenv("PRISMA_OFFLINE_MODE")): # These env vars prevent Prisma from attempting downloads prisma_env["NPM_CONFIG_PREFER_OFFLINE"] = "true" - prisma_env["NPM_CONFIG_CACHE"] = os.getenv( - "NPM_CONFIG_CACHE", "/app/.cache/npm" - ) + prisma_env["NPM_CONFIG_CACHE"] = os.getenv("NPM_CONFIG_CACHE", "/app/.cache/npm") return prisma_env @@ -35,28 +34,29 @@ def _get_prisma_command() -> str: if str_to_bool(os.getenv("PRISMA_OFFLINE_MODE")): # Primary location where Prisma Python package installs the CLI default_cli_path = "/app/.cache/prisma-python/binaries/node_modules/.bin/prisma" - + # Check if custom path is provided (for flexibility) custom_cli_path = os.getenv("PRISMA_CLI_PATH") if custom_cli_path and os.path.exists(custom_cli_path): logger.info(f"Using custom Prisma CLI at {custom_cli_path}") return custom_cli_path - + # Check the default location if os.path.exists(default_cli_path): logger.info(f"Using cached Prisma CLI at {default_cli_path}") return default_cli_path - + # If not found, log warning and fall back logger.warning( f"Prisma CLI not found at {default_cli_path}. " "Falling back to Python wrapper (may attempt downloads)" ) - + # Fall back to the Python wrapper (will work in online mode) return "prisma" + class ProxyExtrasDBManager: @staticmethod def _get_prisma_dir() -> str: @@ -119,7 +119,7 @@ class ProxyExtrasDBManager: stdout=open(migration_file, "w"), check=True, timeout=30, - env=prisma_env, + env=prisma_env ) # 3. Mark the migration as applied since it represents current state @@ -134,7 +134,7 @@ class ProxyExtrasDBManager: ], check=True, timeout=30, - env=prisma_env, + env=prisma_env ) return True @@ -159,20 +159,14 @@ class ProxyExtrasDBManager: @staticmethod def _roll_back_migration(migration_name: str): """Mark a specific migration as rolled back""" - # Set up environment for offline mode if configured + # Set up environment for offline mode if configured prisma_env = _get_prisma_env() subprocess.run( - [ - _get_prisma_command(), - "migrate", - "resolve", - "--rolled-back", - migration_name, - ], + [_get_prisma_command(), "migrate", "resolve", "--rolled-back", migration_name], timeout=60, check=True, capture_output=True, - env=prisma_env, + env=prisma_env ) @staticmethod @@ -184,7 +178,7 @@ class ProxyExtrasDBManager: timeout=60, check=True, capture_output=True, - env=prisma_env, + env=prisma_env ) @staticmethod @@ -254,7 +248,7 @@ class ProxyExtrasDBManager: if not database_url: logger.error("DATABASE_URL not set") return - + diff_dir = ( Path(migrations_dir) / "migrations" @@ -289,7 +283,7 @@ class ProxyExtrasDBManager: check=True, timeout=60, stdout=f, - env=_get_prisma_env(), + env=_get_prisma_env() ) except subprocess.CalledProcessError as e: logger.warning(f"Failed to generate migration diff: {e.stderr}") @@ -319,7 +313,7 @@ class ProxyExtrasDBManager: check=True, capture_output=True, text=True, - env=_get_prisma_env(), + env=_get_prisma_env() ) logger.info(f"prisma db execute stdout: {result.stdout}") logger.info("✅ Migration diff applied successfully") @@ -337,18 +331,12 @@ class ProxyExtrasDBManager: try: logger.info(f"Resolving migration: {migration_name}") subprocess.run( - [ - _get_prisma_command(), - "migrate", - "resolve", - "--applied", - migration_name, - ], + [_get_prisma_command(), "migrate", "resolve", "--applied", migration_name], timeout=60, check=True, capture_output=True, text=True, - env=_get_prisma_env(), + env=_get_prisma_env() ) logger.debug(f"Resolved migration: {migration_name}") except subprocess.CalledProcessError as e: @@ -387,7 +375,7 @@ class ProxyExtrasDBManager: check=True, capture_output=True, text=True, - env=_get_prisma_env(), + env=_get_prisma_env() ) logger.info(f"prisma migrate deploy stdout: {result.stdout}") @@ -425,7 +413,7 @@ class ProxyExtrasDBManager: check=True, capture_output=True, text=True, - env=_get_prisma_env(), + env=_get_prisma_env() ) logger.info( f"✅ Migration {failed_migration} marked as rolled back... retrying" @@ -521,43 +509,12 @@ class ProxyExtrasDBManager: raise else: # Use prisma db push with increased timeout - try: - subprocess.run( - [_get_prisma_command(), "db", "push", "--accept-data-loss"], - timeout=60, - check=True, - capture_output=True, # capture output to check for errors - text=True, - env=_get_prisma_env(), - ) - return True - except subprocess.CalledProcessError as e: - if ( - "Permission denied" in e.stderr - and "schema.prisma" in e.stderr - ): - logger.warning( - f"Permission denied during prisma generate: {e.stderr}. Retrying with --skip-generate..." - ) - # Retry with --skip-generate - subprocess.run( - [ - _get_prisma_command(), - "db", - "push", - "--accept-data-loss", - "--skip-generate", - ], - timeout=60, - check=True, - capture_output=True, - text=True, - env=_get_prisma_env(), - ) - logger.info("✅ prisma db push --skip-generate completed") - return True - else: - raise e + subprocess.run( + [_get_prisma_command(), "db", "push", "--accept-data-loss"], + timeout=60, + check=True, + ) + return True except subprocess.TimeoutExpired: logger.info(f"Attempt {attempt + 1} timed out") time.sleep(random.randrange(5, 15)) diff --git a/litellm/proxy/db/prisma_client.py b/litellm/proxy/db/prisma_client.py index 95800e96589..c9c0cfe8f68 100644 --- a/litellm/proxy/db/prisma_client.py +++ b/litellm/proxy/db/prisma_client.py @@ -386,39 +386,11 @@ class PrismaManager: return ProxyExtrasDBManager.setup_database(use_migrate=use_migrate) else: # Use prisma db push with increased timeout - try: - subprocess.run( - ["prisma", "db", "push", "--accept-data-loss"], - timeout=60, - check=True, - capture_output=True, - text=True, - ) - except subprocess.CalledProcessError as e: - if ( - "Permission denied" in e.stderr - and "schema.prisma" in e.stderr - ): - verbose_proxy_logger.warning( - f"Permission denied during prisma generate: {e.stderr}. Retrying with --skip-generate..." - ) - # Retry with --skip-generate - subprocess.run( - [ - "prisma", - "db", - "push", - "--accept-data-loss", - "--skip-generate", - ], - timeout=60, - check=True, - capture_output=True, - text=True, - ) - return True - else: - raise e + subprocess.run( + ["prisma", "db", "push", "--accept-data-loss"], + timeout=60, + check=True, + ) return True except subprocess.TimeoutExpired: verbose_proxy_logger.warning(f"Attempt {attempt + 1} timed out") diff --git a/litellm/proxy/prisma_migration.py b/litellm/proxy/prisma_migration.py index 2fe12b1439c..251d1e56287 100644 --- a/litellm/proxy/prisma_migration.py +++ b/litellm/proxy/prisma_migration.py @@ -15,21 +15,14 @@ from litellm.proxy.proxy_cli import run_server # Call the Click command with standalone_mode=False run_server(["--skip_server_startup"], standalone_mode=False) -# Run prisma generate +# run prisma generate verbose_proxy_logger.info("Running 'prisma generate'...") -try: - result = subprocess.run(["prisma", "generate"], capture_output=True, text=True) - if result.returncode != 0: - if "Permission denied" in result.stderr: - verbose_proxy_logger.warning( - f"Permission denied during 'prisma generate'. Skipping generation, assuming client is pre-generated. Error: {result.stderr}" - ) - else: - verbose_proxy_logger.info( - f"'prisma generate' failed with exit code {result.returncode}." - ) - verbose_proxy_logger.error( - f"'prisma generate' stderr: {result.stderr}" - ) # Log stderr -except Exception as e: - verbose_proxy_logger.error(f"Error running prisma generate: {e}") +result = subprocess.run(["prisma", "generate"], capture_output=True, text=True) +verbose_proxy_logger.info(f"'prisma generate' stdout: {result.stdout}") # Log stdout +exit_code = result.returncode + +if exit_code != 0: + verbose_proxy_logger.info(f"'prisma generate' failed with exit code {exit_code}.") + verbose_proxy_logger.error( + f"'prisma generate' stderr: {result.stderr}" + ) # Log stderr diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py index 028aad9a2e1..2059246674b 100644 --- a/litellm/proxy/proxy_cli.py +++ b/litellm/proxy/proxy_cli.py @@ -797,11 +797,7 @@ def run_server( # noqa: PLR0915 ): check_prisma_schema_diff(db_url=None) else: - if not PrismaManager.setup_database( - use_migrate=not use_prisma_db_push - ): - print("LiteLLM: Database setup failed. Exiting...") # noqa - sys.exit(1) + PrismaManager.setup_database(use_migrate=not use_prisma_db_push) else: print( # noqa f"Unable to connect to DB. DATABASE_URL found in environment, but prisma package not found." # noqa diff --git a/tests/test_litellm/proxy/test_migration_failure_handling.py b/tests/test_litellm/proxy/test_migration_failure_handling.py deleted file mode 100644 index 426049007b1..00000000000 --- a/tests/test_litellm/proxy/test_migration_failure_handling.py +++ /dev/null @@ -1,114 +0,0 @@ -import sys -import os -import subprocess -from unittest.mock import MagicMock, patch -from click.testing import CliRunner - -# Add parent directory to path to allow importing litellm -sys.path.insert(0, os.path.abspath("../../..")) - -from litellm.proxy.db.prisma_client import PrismaManager -from litellm.proxy.proxy_cli import run_server - - -class TestMigrationFailureHandling: - @patch("subprocess.run") - def test_prisma_client_permission_error_retry(self, mock_subprocess_run): - """ - Regression Test: Verifies that PrismaManager.setup_database - catches PermissionError during 'prisma db push' and retries with '--skip-generate'. - """ - # Mock behavior: - # call 1: raises CalledProcessError with "Permission denied" and "schema.prisma" - # call 2 (retry): succeeds - - error_output = "Error: Permission denied writing to ... schema.prisma" - - mock_process_error = subprocess.CalledProcessError( - returncode=1, cmd=["prisma", "db", "push"], stderr=error_output - ) - - mock_subprocess_run.side_effect = [ - mock_process_error, # 1st attempt fails with permission error - MagicMock(returncode=0), # 2nd attempt (retry) succeeds - ] - - # Ensure we run the 'db push' path (use_migrate=False) - # We also need to mock should_update_prisma_schema to return True - - with patch( - "litellm.proxy.db.prisma_client.should_update_prisma_schema", - return_value=True, - ): - # Run setup_database with use_migrate=False to trigger 'prisma db push' path - result = PrismaManager.setup_database(use_migrate=False) - - # Assert success - assert result is True - - # Verify calls - assert mock_subprocess_run.call_count == 2 - - # Check 1st call arguments (standard push) - args1, _ = mock_subprocess_run.call_args_list[0] - assert "push" in args1[0] - assert "--skip-generate" not in args1[0] - - # Check 2nd call arguments (retry with skip-generate) - args2, _ = mock_subprocess_run.call_args_list[1] - assert "push" in args2[0] - assert "--skip-generate" in args2[0] - - def test_proxy_cli_exit_on_migration_fail(self): - """ - Regression Test: Verifies that proxy_cli.run_server exits with NON-ZERO status - if PrismaManager.setup_database returns False. - """ - runner = CliRunner() - - # Mock setup_database to return False (Simulating failure) - # Mock should_update_prisma_schema to return True (Ensure we hit the DB setup logic) - with patch( - "litellm.proxy.db.prisma_client.PrismaManager.setup_database", - return_value=False, - ), patch( - "litellm.proxy.db.prisma_client.should_update_prisma_schema", - return_value=True, - ): - # Mock dependencies to prevent actual server startup and handle imports - mock_app = MagicMock() - mock_proxy_config = MagicMock() - - # Patch sys.modules to prevent ImportErrors for proxy_server - with patch.dict( - "sys.modules", - { - "proxy_server": MagicMock( - app=mock_app, ProxyConfig=mock_proxy_config - ) - }, - ): - with patch( - "litellm.proxy.proxy_cli.ProxyInitializationHelpers._get_default_unvicorn_init_args" - ) as mock_get_args: - mock_get_args.return_value = { - "app": "app", - "host": "localhost", - "port": 8000, - } - - # Set DATABASE_URL to trigger DB logic - with patch.dict( - os.environ, - {"DATABASE_URL": "postgresql://user:pass@localhost:5432/db"}, - ): - # Execute: Run server with --local and --skip_server_startup - result = runner.invoke( - run_server, ["--local", "--skip_server_startup"] - ) - - # Assert: Exit code should be non-zero (failure) - assert ( - result.exit_code != 0 - ), f"Expected non-zero exit code, got {result.exit_code}. Output: {result.output}" - assert "Database setup failed. Exiting..." in result.output From c8d065656e51e0ee4c2f028cfe927ea6961e94a9 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Wed, 21 Jan 2026 17:33:53 +0530 Subject: [PATCH 09/10] Fix litellm_staging_01_20_2026 mypy issues --- litellm/integrations/opentelemetry.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 From 22158f8f0340f3ddf207608df9b76c77b571dac8 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Wed, 21 Jan 2026 17:35:28 +0530 Subject: [PATCH 10/10] Fix litellm_staging_01_20_2026 mypy issues --- litellm/proxy/proxy_server.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index bbf434aead1..eef2af89799 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -3102,11 +3102,12 @@ class ProxyConfig: 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}") @@ -3120,12 +3121,12 @@ class ProxyConfig: ) 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(