install prisma migration files - connects litellm proxy to litellm's prisma migration files (#9637)

* build(README.md): initial commit adding a separate folder for additional proxy files. Meant to reduce size of core package

* build(litellm-proxy-extras/): new pip package for storing migration files

allows litellm proxy to use migration files, without adding them to core repo

* build(litellm-proxy-extras/): cleanup pyproject.toml

* build: move prisma migration files inside new proxy extras package

* build(run_migration.py): update script to write to correct folder

* build(proxy_cli.py): load in migration files from litellm-proxy-extras

Closes https://github.com/BerriAI/litellm/issues/9558

* build: add MIT license to litellm-proxy-extras

* test: update test

* fix: fix schema

* bump: version 0.1.0 → 0.1.1

* build(publish-proxy-extras.sh): add script for publishing new proxy-extras version

* build(liccheck.ini): add litellm-proxy-extras to authorized packages

* fix(litellm-proxy-extras/utils.py): move prisma migrate logic inside extra proxy pkg

easier since migrations folder already there

* build(pre-commit-config.yaml): add litellm_proxy_extras to ci tests

* docs(config_settings.md): document new env var

* build(pyproject.toml): bump relevant files when litellm-proxy-extras version changed

* build(pre-commit-config.yaml): run poetry check on litellm-proxy-extras as well

.pre-commit-config.yaml

@@ -6,46 +6,35 @@ repos:
         entry: pyright
         language: system
         types: [python]
-        files: ^litellm/
+        files: ^(litellm/|litellm_proxy_extras/)
     -   id: isort
         name: isort
         entry: isort
         language: system
         types: [python]
-        files: litellm/.*\.py
+        files: (litellm/|litellm_proxy_extras/).*\.py
         exclude: ^litellm/__init__.py$
     -   id: black
         name: black
         entry: poetry run black
         language: system
         types: [python]
-        files: litellm/.*\.py
+        files: (litellm/|litellm_proxy_extras/).*\.py
 -   repo: https://github.com/pycqa/flake8
     rev: 7.0.0  # The version of flake8 to use
     hooks:
     -  id: flake8
        exclude: ^litellm/tests/|^litellm/proxy/tests/|^litellm/tests/litellm/|^tests/litellm/
        additional_dependencies: [flake8-print]
-       files: litellm/.*\.py
+       files: (litellm/|litellm_proxy_extras/).*\.py
-    # -  id: flake8
-    #    name: flake8 (router.py function length)
-    #    files: ^litellm/router\.py$
-    #    args: [--max-function-length=40]
-    # #    additional_dependencies: [flake8-functions]
 -   repo: https://github.com/python-poetry/poetry
     rev: 1.8.0
     hooks:
       - id: poetry-check
+        files: ^(pyproject.toml|litellm-proxy-extras/pyproject.toml)$
 -   repo: local
     hooks:
     -   id: check-files-match
         name: Check if files match
         entry: python3 ci_cd/check_files_match.py
-        language: system
+        language: system
-    # -   id: check-file-length
-    #     name: Check file length
-    #     entry: python check_file_length.py
-    #     args: ["10000"]  # set your desired maximum number of lines
-    #     language: python
-    #     files: litellm/.*\.py
-    #     exclude: ^litellm/tests/

ci_cd/publish-proxy-extras.sh

@@ -0,0 +1,19 @@
+#!/bin/bash
+
+# Exit on error
+set -e
+
+echo "🚀 Building and publishing litellm-proxy-extras"
+
+# Navigate to litellm-proxy-extras directory
+cd "$(dirname "$0")/../litellm-proxy-extras"
+
+# Build the package
+echo "📦 Building package..."
+poetry build
+
+# Publish to PyPI
+echo "🌎 Publishing to PyPI..."
+poetry publish
+
+echo "✅ Done! Package published successfully"

ci_cd/run_migration.py

@@ -8,7 +8,7 @@ import shutil
 
 def create_migration(migration_name: str = None):
     """
-    Create a new migration SQL file in deploy/migrations directory by comparing
+    Create a new migration SQL file in the migrations directory by comparing
     current database state with schema
 
     Args:
@@ -17,8 +17,7 @@ def create_migration(migration_name: str = None):
     try:
         # Get paths
         root_dir = Path(__file__).parent.parent
-        deploy_dir = root_dir / "deploy"
+        migrations_dir = root_dir / "litellm-proxy-extras" / "litellm_proxy_extras" / "migrations"
-        migrations_dir = deploy_dir / "migrations"
         schema_path = root_dir / "schema.prisma"
 
         # Create temporary PostgreSQL database

docs/my-website/docs/proxy/config_settings.md

@@ -515,4 +515,5 @@ router_settings:
 | UPSTREAM_LANGFUSE_RELEASE | Release version identifier for upstream Langfuse
 | UPSTREAM_LANGFUSE_SECRET_KEY | Secret key for upstream Langfuse authentication
 | USE_AWS_KMS | Flag to enable AWS Key Management Service for encryption
+| USE_PRISMA_MIGRATE | Flag to use prisma migrate instead of prisma db push. Recommended for production environments.
 | WEBHOOK_URL | URL for receiving webhooks from external services

litellm-proxy-extras/LICENSE

@@ -0,0 +1,26 @@
+Portions of this software are licensed as follows:
+
+* All content that resides under the "enterprise/" directory of this repository, if that directory exists, is licensed under the license defined in "enterprise/LICENSE".
+* Content outside of the above mentioned directories or restrictions above is available under the MIT license as defined below.
+---
+MIT License
+
+Copyright (c) 2023 Berri AI
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

litellm-proxy-extras/README.md

@@ -0,0 +1,21 @@
+Additional files for the proxy. Reduces the size of the main litellm package.
+
+Currently, only stores the migration.sql files for litellm-proxy.
+
+To install, run:
+
+```bash
+pip install litellm-proxy-extras
+```
+OR 
+
+```bash
+pip install litellm[proxy] # installs litellm-proxy-extras and other proxy dependencies.
+```
+
+To use the migrations, run:
+
+```bash
+litellm --use_prisma_migrate
+```
+

litellm-proxy-extras/litellm_proxy_extras/_logging.py

@@ -0,0 +1,12 @@
+import logging
+
+# Set up package logger
+logger = logging.getLogger("litellm_proxy_extras")
+if not logger.handlers:  # Only add handler if none exists
+    handler = logging.StreamHandler()
+    formatter = logging.Formatter(
+        "%(asctime)s - %(name)s - %(levelname)s - %(message)s"
+    )
+    handler.setFormatter(formatter)
+    logger.addHandler(handler)
+    logger.setLevel(logging.INFO)

litellm-proxy-extras/litellm_proxy_extras/utils.py

@@ -0,0 +1,80 @@
+import os
+import random
+import subprocess
+import time
+from typing import Optional
+
+from litellm_proxy_extras._logging import logger
+
+
+def str_to_bool(value: Optional[str]) -> bool:
+    if value is None:
+        return False
+    return value.lower() in ("true", "1", "t", "y", "yes")
+
+
+class ProxyExtrasDBManager:
+    @staticmethod
+    def setup_database(schema_path: str, use_migrate: bool = False) -> bool:
+        """
+        Set up the database using either prisma migrate or prisma db push
+        Uses migrations from litellm-proxy-extras package
+
+        Args:
+            schema_path (str): Path to the Prisma schema file
+            use_migrate (bool): Whether to use prisma migrate instead of db push
+
+        Returns:
+            bool: True if setup was successful, False otherwise
+        """
+        use_migrate = str_to_bool(os.getenv("USE_PRISMA_MIGRATE")) or use_migrate
+        for attempt in range(4):
+            original_dir = os.getcwd()
+            schema_dir = os.path.dirname(schema_path)
+            os.chdir(schema_dir)
+
+            try:
+                if use_migrate:
+                    logger.info("Running prisma migrate deploy")
+                    try:
+                        # Set migrations directory for Prisma
+                        subprocess.run(
+                            ["prisma", "migrate", "deploy"],
+                            timeout=60,
+                            check=True,
+                            capture_output=True,
+                            text=True,
+                        )
+                        logger.info("prisma migrate deploy completed")
+                        return True
+                    except subprocess.CalledProcessError as e:
+                        logger.info(f"prisma db error: {e.stderr}, e: {e.stdout}")
+                        if (
+                            "P3005" in e.stderr
+                            and "database schema is not empty" in e.stderr
+                        ):
+                            logger.info("Error: Database schema is not empty")
+                            return False
+                else:
+                    # Use prisma db push with increased timeout
+                    subprocess.run(
+                        ["prisma", "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))
+            except subprocess.CalledProcessError as e:
+                attempts_left = 3 - attempt
+                retry_msg = (
+                    f" Retrying... ({attempts_left} attempts left)"
+                    if attempts_left > 0
+                    else ""
+                )
+                logger.info(f"The process failed to execute. Details: {e}.{retry_msg}")
+                time.sleep(random.randrange(5, 15))
+            finally:
+                os.chdir(original_dir)
+        return False

litellm-proxy-extras/pyproject.toml

@@ -0,0 +1,30 @@
+[tool.poetry]
+name = "litellm-proxy-extras"
+version = "0.1.1"
+description = "Additional files for the LiteLLM Proxy. Reduces the size of the main litellm package."
+authors = ["BerriAI"]
+readme = "README.md"
+
+
+[tool.poetry.urls]
+homepage = "https://litellm.ai"
+Homepage = "https://litellm.ai"
+repository = "https://github.com/BerriAI/litellm"
+Repository = "https://github.com/BerriAI/litellm"
+documentation = "https://docs.litellm.ai"
+Documentation = "https://docs.litellm.ai"
+
+[tool.poetry.dependencies]
+python = ">=3.8.1,<4.0, !=3.9.7"
+
+[build-system]
+requires = ["poetry-core"]
+build-backend = "poetry.core.masonry.api"
+
+[tool.commitizen]
+version = "0.1.1"
+version_files = [
+    "pyproject.toml:version",
+    "../requirements.txt:litellm-proxy-extras==",
+    "../pyproject.toml:litellm-proxy-extras = {version = \""
+]

litellm/proxy/db/prisma_client.py

@@ -179,42 +179,6 @@ class PrismaManager:
             verbose_proxy_logger.warning(f"Error creating baseline migration: {e}")
             return False
 
-    @staticmethod
-    def _copy_spend_tracking_migrations(prisma_dir: str) -> bool:
-        import shutil
-        from pathlib import Path
-
-        """
-        Check for and copy over spend tracking migrations if they exist in the deploy directory.
-        Returns True if migrations were found and copied, False otherwise.
-        """
-        try:
-            # Get the current file's directory
-            current_dir = Path(__file__).parent
-
-            # Check for migrations in the deploy directory (../../deploy/migrations)
-            deploy_migrations_dir = (
-                current_dir.parent.parent.parent / "deploy" / "migrations"
-            )
-
-            # Local migrations directory
-            local_migrations_dir = Path(prisma_dir + "/migrations")
-
-            if deploy_migrations_dir.exists():
-                # Create local migrations directory if it doesn't exist
-                local_migrations_dir.mkdir(parents=True, exist_ok=True)
-
-                # Copy all migration files
-                # Copy entire migrations folder recursively
-                shutil.copytree(
-                    deploy_migrations_dir, local_migrations_dir, dirs_exist_ok=True
-                )
-
-                return True
-            return False
-        except Exception:
-            return False
-
     @staticmethod
     def _get_migration_names(migrations_dir: str) -> list:
         """Get all migration directory names from the migrations folder"""
@@ -251,6 +215,7 @@ class PrismaManager:
             bool: True if setup was successful, False otherwise
         """
 
+        use_migrate = str_to_bool(os.getenv("USE_PRISMA_MIGRATE")) or use_migrate
         for attempt in range(4):
             original_dir = os.getcwd()
             prisma_dir = PrismaManager._get_prisma_dir()
@@ -258,44 +223,20 @@ class PrismaManager:
             os.chdir(prisma_dir)
             try:
                 if use_migrate:
-                    PrismaManager._copy_spend_tracking_migrations(
-                        prisma_dir
-                    )  # place a migration in the migrations directory
-                    verbose_proxy_logger.info("Running prisma migrate deploy")
                     try:
-                        subprocess.run(
+                        from litellm_proxy_extras.utils import ProxyExtrasDBManager
-                            ["prisma", "migrate", "deploy"],
+                    except ImportError as e:
-                            timeout=60,
+                        verbose_proxy_logger.error(
-                            check=True,
+                            f"\033[1;31mLiteLLM: Failed to import proxy extras. Got {e}\033[0m"
-                            capture_output=True,
-                            text=True,
                         )
-                        verbose_proxy_logger.info("prisma migrate deploy completed")
+                        return False
 
-                        # Resolve all migrations in the migrations directory
+                    prisma_dir = PrismaManager._get_prisma_dir()
-                        migrations_dir = os.path.join(prisma_dir, "migrations")
+                    schema_path = prisma_dir + "/schema.prisma"
-                        PrismaManager._resolve_all_migrations(migrations_dir)
 
-                        return True
+                    return ProxyExtrasDBManager.setup_database(
-                    except subprocess.CalledProcessError as e:
+                        schema_path=schema_path, use_migrate=use_migrate
-                        verbose_proxy_logger.warning(
+                    )
-                            f"prisma db error: {e.stderr}, e: {e.stdout}"
-                        )
-                        if (
-                            "P3005" in e.stderr
-                            and "database schema is not empty" in e.stderr
-                        ):
-                            verbose_proxy_logger.info("Creating baseline migration")
-                            if PrismaManager._create_baseline_migration(schema_path):
-                                verbose_proxy_logger.info(
-                                    "Resolving all migrations after baseline"
-                                )
-
-                                # Resolve all migrations after baseline
-                                migrations_dir = os.path.join(prisma_dir, "migrations")
-                                PrismaManager._resolve_all_migrations(migrations_dir)
-
-                                return True
                 else:
                     # Use prisma db push with increased timeout
                     subprocess.run(

litellm/proxy/proxy_cli.py

@@ -672,7 +672,7 @@ def run_server(  # noqa: PLR0915
                 LiteLLMDatabaseConnectionPool.database_connection_pool_limit.value,
             )
             db_connection_timeout = general_settings.get(
-                "database_connection_timeout",
+                "database_connection_pool_timeout",
                 LiteLLMDatabaseConnectionPool.database_connection_pool_timeout.value,
             )
             if database_url and database_url.startswith("os.environ/"):

poetry.lock

@@ -1676,6 +1676,17 @@ files = [
 importlib-resources = {version = ">=1.4.0", markers = "python_version < \"3.9\""}
 referencing = ">=0.31.0"
 
+[[package]]
+name = "litellm-proxy-extras"
+version = "0.1.1"
+description = "Additional files for the LiteLLM Proxy. Reduces the size of the main litellm package."
+optional = true
+python-versions = "!=2.7.*,!=3.0.*,!=3.1.*,!=3.2.*,!=3.3.*,!=3.4.*,!=3.5.*,!=3.6.*,!=3.7.*,>=3.8"
+files = [
+    {file = "litellm_proxy_extras-0.1.1-py3-none-any.whl", hash = "sha256:2b3c4c5474bacbde2424c1cd13b21f85c65e9c4346f6159badd49a210eedef5c"},
+    {file = "litellm_proxy_extras-0.1.1.tar.gz", hash = "sha256:a1eb911ad2e3742238863d314a8bd6d02dd0cc213ba040b2c0593f132fbf3117"},
+]
+
 [[package]]
 name = "markupsafe"
 version = "2.1.5"
@@ -4119,9 +4130,9 @@ type = ["pytest-mypy"]
 
 [extras]
 extra-proxy = ["azure-identity", "azure-keyvault-secrets", "google-cloud-kms", "prisma", "redisvl", "resend"]
-proxy = ["PyJWT", "apscheduler", "backoff", "boto3", "cryptography", "fastapi", "fastapi-sso", "gunicorn", "mcp", "orjson", "pynacl", "python-multipart", "pyyaml", "rq", "uvicorn", "uvloop", "websockets"]
+proxy = ["PyJWT", "apscheduler", "backoff", "boto3", "cryptography", "fastapi", "fastapi-sso", "gunicorn", "litellm-proxy-extras", "mcp", "orjson", "pynacl", "python-multipart", "pyyaml", "rq", "uvicorn", "uvloop", "websockets"]
 
 [metadata]
 lock-version = "2.0"
 python-versions = ">=3.8.1,<4.0, !=3.9.7"
-content-hash = "36792478ff4afec5c8e748caf9b2ae6bebf3dd223e78bea2626b6589ef3277e4"
+content-hash = "16cbf20784776377805f5e33c6bc97dce76303132aa3d81c7e6fe743f0ee3fc1"

pyproject.toml

@@ -55,6 +55,7 @@ websockets = {version = "^13.1.0", optional = true}
 boto3 = {version = "1.34.34", optional = true}
 redisvl = {version = "^0.4.1", optional = true, markers = "python_version >= '3.9' and python_version < '3.14'"}
 mcp = {version = "1.5.0", optional = true, python = ">=3.10"}
+litellm-proxy-extras = {version = "0.1.1", optional = true}
 
 [tool.poetry.extras]
 proxy = [
@@ -74,7 +75,8 @@ proxy = [
     "pynacl",
     "websockets",
     "boto3",
-    "mcp"
+    "mcp",
+    "litellm-proxy-extras"
 ]
 
 extra_proxy = [

requirements.txt

@@ -38,6 +38,7 @@ sentry_sdk==2.21.0 # for sentry error handling
 detect-secrets==1.5.0 # Enterprise - secret detection / masking in LLM requests
 cryptography==43.0.1
 tzdata==2025.1 # IANA time zone database
+litellm-proxy-extras==0.1.1 # for proxy extras - e.g. prisma migrations
 
 ### LITELLM PACKAGE DEPENDENCIES
 python-dotenv==1.0.0 # for env 

tests/code_coverage_tests/liccheck.ini

@@ -85,4 +85,5 @@ anthropic: >=0.21.3 # MIT License
 detect-secrets: >=1.5.0 # MIT License
 importlib-metadata: >=6.8.0 # Apache 2.0 License
 tokenizers: >=0.20.2 # Apache 2.0 License
-jinja2: >=3.1.4 # BSD 3-Clause License
+jinja2: >=3.1.4 # BSD 3-Clause License
+litellm-proxy-extras: >=0.1.1 # MIT License

tests/proxy_unit_tests/test_db_schema_migration.py

@@ -24,7 +24,7 @@ def test_aaaasschema_migration_check(schema_setup, monkeypatch):
     # test_db_url = "postgresql://neondb_owner:npg_JiZPS0DAhRn4@ep-delicate-wave-a55cvbuc.us-east-2.aws.neon.tech/neondb?sslmode=require"
     monkeypatch.setenv("DATABASE_URL", test_db_url)
 
-    deploy_dir = Path("./deploy")
+    deploy_dir = Path("./litellm-proxy-extras/litellm_proxy_extras")
     source_migrations_dir = deploy_dir / "migrations"
     schema_path = Path("./schema.prisma")