diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index bceedb41aa..dedb37d6dd 100644 --- a/.pre-commit-config.yaml +++ b/.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 - # - id: flake8 - # name: flake8 (router.py function length) - # files: ^litellm/router\.py$ - # args: [--max-function-length=40] - # # additional_dependencies: [flake8-functions] + files: (litellm/|litellm_proxy_extras/).*\.py - 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 - # - 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/ \ No newline at end of file + language: system \ No newline at end of file diff --git a/ci_cd/publish-proxy-extras.sh b/ci_cd/publish-proxy-extras.sh new file mode 100644 index 0000000000..6c83d1f921 --- /dev/null +++ b/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" \ No newline at end of file diff --git a/ci_cd/run_migration.py b/ci_cd/run_migration.py index 69b4308e57..b11a38395c 100644 --- a/ci_cd/run_migration.py +++ b/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 = deploy_dir / "migrations" + migrations_dir = root_dir / "litellm-proxy-extras" / "litellm_proxy_extras" / "migrations" schema_path = root_dir / "schema.prisma" # Create temporary PostgreSQL database diff --git a/docs/my-website/docs/proxy/config_settings.md b/docs/my-website/docs/proxy/config_settings.md index 4a62184df7..05b3e0be37 100644 --- a/docs/my-website/docs/proxy/config_settings.md +++ b/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 diff --git a/litellm-proxy-extras/LICENSE b/litellm-proxy-extras/LICENSE new file mode 100644 index 0000000000..3bfef5bae9 --- /dev/null +++ b/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. diff --git a/litellm-proxy-extras/README.md b/litellm-proxy-extras/README.md new file mode 100644 index 0000000000..29453f65ba --- /dev/null +++ b/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 +``` + diff --git a/litellm-proxy-extras/dist/litellm_proxy_extras-0.1.0-py3-none-any.whl b/litellm-proxy-extras/dist/litellm_proxy_extras-0.1.0-py3-none-any.whl new file mode 100644 index 0000000000..1aff64ef58 Binary files /dev/null and b/litellm-proxy-extras/dist/litellm_proxy_extras-0.1.0-py3-none-any.whl differ diff --git a/litellm-proxy-extras/dist/litellm_proxy_extras-0.1.0.tar.gz b/litellm-proxy-extras/dist/litellm_proxy_extras-0.1.0.tar.gz new file mode 100644 index 0000000000..0bdf828163 Binary files /dev/null and b/litellm-proxy-extras/dist/litellm_proxy_extras-0.1.0.tar.gz differ diff --git a/litellm-proxy-extras/dist/litellm_proxy_extras-0.1.1-py3-none-any.whl b/litellm-proxy-extras/dist/litellm_proxy_extras-0.1.1-py3-none-any.whl new file mode 100644 index 0000000000..e2583935a4 Binary files /dev/null and b/litellm-proxy-extras/dist/litellm_proxy_extras-0.1.1-py3-none-any.whl differ diff --git a/litellm-proxy-extras/dist/litellm_proxy_extras-0.1.1.tar.gz b/litellm-proxy-extras/dist/litellm_proxy_extras-0.1.1.tar.gz new file mode 100644 index 0000000000..c9111dd9c3 Binary files /dev/null and b/litellm-proxy-extras/dist/litellm_proxy_extras-0.1.1.tar.gz differ diff --git a/litellm-proxy-extras/litellm_proxy_extras/__init__.py b/litellm-proxy-extras/litellm_proxy_extras/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/litellm-proxy-extras/litellm_proxy_extras/_logging.py b/litellm-proxy-extras/litellm_proxy_extras/_logging.py new file mode 100644 index 0000000000..118caecf48 --- /dev/null +++ b/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) diff --git a/deploy/migrations/20250326162113_baseline/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250326162113_baseline/migration.sql similarity index 100% rename from deploy/migrations/20250326162113_baseline/migration.sql rename to litellm-proxy-extras/litellm_proxy_extras/migrations/20250326162113_baseline/migration.sql diff --git a/deploy/migrations/20250326171002_add_daily_user_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250326171002_add_daily_user_table/migration.sql similarity index 100% rename from deploy/migrations/20250326171002_add_daily_user_table/migration.sql rename to litellm-proxy-extras/litellm_proxy_extras/migrations/20250326171002_add_daily_user_table/migration.sql diff --git a/deploy/migrations/20250327180120_add_api_requests_to_daily_user_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250327180120_add_api_requests_to_daily_user_table/migration.sql similarity index 100% rename from deploy/migrations/20250327180120_add_api_requests_to_daily_user_table/migration.sql rename to litellm-proxy-extras/litellm_proxy_extras/migrations/20250327180120_add_api_requests_to_daily_user_table/migration.sql diff --git a/deploy/migrations/20250329084805_new_cron_job_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250329084805_new_cron_job_table/migration.sql similarity index 100% rename from deploy/migrations/20250329084805_new_cron_job_table/migration.sql rename to litellm-proxy-extras/litellm_proxy_extras/migrations/20250329084805_new_cron_job_table/migration.sql diff --git a/deploy/migrations/migration_lock.toml b/litellm-proxy-extras/litellm_proxy_extras/migrations/migration_lock.toml similarity index 100% rename from deploy/migrations/migration_lock.toml rename to litellm-proxy-extras/litellm_proxy_extras/migrations/migration_lock.toml diff --git a/litellm-proxy-extras/litellm_proxy_extras/utils.py b/litellm-proxy-extras/litellm_proxy_extras/utils.py new file mode 100644 index 0000000000..894ae34122 --- /dev/null +++ b/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 diff --git a/litellm-proxy-extras/pyproject.toml b/litellm-proxy-extras/pyproject.toml new file mode 100644 index 0000000000..c130a7fa9b --- /dev/null +++ b/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 = \"" +] \ No newline at end of file diff --git a/litellm-proxy-extras/tests/__init__.py b/litellm-proxy-extras/tests/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/litellm/proxy/db/prisma_client.py b/litellm/proxy/db/prisma_client.py index 4e38321e91..c81dc35e0c 100644 --- a/litellm/proxy/db/prisma_client.py +++ b/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( - ["prisma", "migrate", "deploy"], - timeout=60, - check=True, - capture_output=True, - text=True, + from litellm_proxy_extras.utils import ProxyExtrasDBManager + except ImportError as e: + verbose_proxy_logger.error( + f"\033[1;31mLiteLLM: Failed to import proxy extras. Got {e}\033[0m" ) - verbose_proxy_logger.info("prisma migrate deploy completed") + return False - # Resolve all migrations in the migrations directory - migrations_dir = os.path.join(prisma_dir, "migrations") - PrismaManager._resolve_all_migrations(migrations_dir) + prisma_dir = PrismaManager._get_prisma_dir() + schema_path = prisma_dir + "/schema.prisma" - return True - except subprocess.CalledProcessError as e: - 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 + return ProxyExtrasDBManager.setup_database( + schema_path=schema_path, use_migrate=use_migrate + ) else: # Use prisma db push with increased timeout subprocess.run( diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py index 8196eb597e..de78247daf 100644 --- a/litellm/proxy/proxy_cli.py +++ b/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/"): diff --git a/poetry.lock b/poetry.lock index d659d66952..157291f5f3 100644 --- a/poetry.lock +++ b/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" diff --git a/pyproject.toml b/pyproject.toml index 5eb4a71160..34564bd049 100644 --- a/pyproject.toml +++ b/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 = [ diff --git a/requirements.txt b/requirements.txt index c28a2e5cf2..b3f0f1d073 100644 --- a/requirements.txt +++ b/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 diff --git a/tests/code_coverage_tests/liccheck.ini b/tests/code_coverage_tests/liccheck.ini index efd7dffe1c..c5ecc898f2 100644 --- a/tests/code_coverage_tests/liccheck.ini +++ b/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 \ No newline at end of file +jinja2: >=3.1.4 # BSD 3-Clause License +litellm-proxy-extras: >=0.1.1 # MIT License \ No newline at end of file diff --git a/tests/proxy_unit_tests/test_db_schema_migration.py b/tests/proxy_unit_tests/test_db_schema_migration.py index 089c50ebf8..b317818375 100644 --- a/tests/proxy_unit_tests/test_db_schema_migration.py +++ b/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")