litellm-proxy-extras 0.1.8__py3-none-any.whl → 0.1.19__py3-none-any.whl

litellm_proxy_extras/migrations/20250415151647_add_cache_read_write_tokens_daily_spend_transactions/migration.sql

@@ -0,0 +1,4 @@
+-- AlterTable
+ALTER TABLE "LiteLLM_DailyUserSpend" ADD COLUMN     "cache_creation_input_tokens" INTEGER NOT NULL DEFAULT 0,
+ADD COLUMN     "cache_read_input_tokens" INTEGER NOT NULL DEFAULT 0;
+

litellm_proxy_extras/migrations/20250416115320_add_tag_table_to_db/migration.sql

@@ -0,0 +1,45 @@
+-- AlterTable
+ALTER TABLE "LiteLLM_DailyTeamSpend" ADD COLUMN     "cache_creation_input_tokens" INTEGER NOT NULL DEFAULT 0,
+ADD COLUMN     "cache_read_input_tokens" INTEGER NOT NULL DEFAULT 0;
+
+-- CreateTable
+CREATE TABLE "LiteLLM_DailyTagSpend" (
+    "id" TEXT NOT NULL,
+    "tag" TEXT NOT NULL,
+    "date" TEXT NOT NULL,
+    "api_key" TEXT NOT NULL,
+    "model" TEXT NOT NULL,
+    "model_group" TEXT,
+    "custom_llm_provider" TEXT,
+    "prompt_tokens" INTEGER NOT NULL DEFAULT 0,
+    "completion_tokens" INTEGER NOT NULL DEFAULT 0,
+    "cache_read_input_tokens" INTEGER NOT NULL DEFAULT 0,
+    "cache_creation_input_tokens" INTEGER NOT NULL DEFAULT 0,
+    "spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0,
+    "api_requests" INTEGER NOT NULL DEFAULT 0,
+    "successful_requests" INTEGER NOT NULL DEFAULT 0,
+    "failed_requests" INTEGER NOT NULL DEFAULT 0,
+    "created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
+    "updated_at" TIMESTAMP(3) NOT NULL,
+
+    CONSTRAINT "LiteLLM_DailyTagSpend_pkey" PRIMARY KEY ("id")
+);
+
+-- CreateIndex
+CREATE UNIQUE INDEX "LiteLLM_DailyTagSpend_tag_key" ON "LiteLLM_DailyTagSpend"("tag");
+
+-- CreateIndex
+CREATE INDEX "LiteLLM_DailyTagSpend_date_idx" ON "LiteLLM_DailyTagSpend"("date");
+
+-- CreateIndex
+CREATE INDEX "LiteLLM_DailyTagSpend_tag_idx" ON "LiteLLM_DailyTagSpend"("tag");
+
+-- CreateIndex
+CREATE INDEX "LiteLLM_DailyTagSpend_api_key_idx" ON "LiteLLM_DailyTagSpend"("api_key");
+
+-- CreateIndex
+CREATE INDEX "LiteLLM_DailyTagSpend_model_idx" ON "LiteLLM_DailyTagSpend"("model");
+
+-- CreateIndex
+CREATE UNIQUE INDEX "LiteLLM_DailyTagSpend_tag_date_api_key_model_custom_llm_pro_key" ON "LiteLLM_DailyTagSpend"("tag", "date", "api_key", "model", "custom_llm_provider");
+

litellm_proxy_extras/migrations/20250416151339_drop_tag_uniqueness_requirement/migration.sql

@@ -0,0 +1,3 @@
+-- DropIndex
+DROP INDEX "LiteLLM_DailyTagSpend_tag_key";
+

litellm_proxy_extras/migrations/20250416185146_add_allowed_routes_litellm_verification_token/migration.sql

@@ -0,0 +1,3 @@
+-- AlterTable
+ALTER TABLE "LiteLLM_VerificationToken" ADD COLUMN     "allowed_routes" TEXT[] DEFAULT ARRAY[]::TEXT[];
+

litellm_proxy_extras/migrations/20250425182129_add_session_id/migration.sql

@@ -0,0 +1,4 @@
+-- AlterTable
+ALTER TABLE "LiteLLM_SpendLogs" ADD COLUMN     "proxy_server_request" JSONB DEFAULT '{}',
+ADD COLUMN     "session_id" TEXT;
+

litellm_proxy_extras/migrations/20250430193429_add_managed_vector_stores/migration.sql

@@ -0,0 +1,14 @@
+-- CreateTable
+CREATE TABLE "LiteLLM_ManagedVectorStoresTable" (
+    "vector_store_id" TEXT NOT NULL,
+    "custom_llm_provider" TEXT NOT NULL,
+    "vector_store_name" TEXT,
+    "vector_store_description" TEXT,
+    "vector_store_metadata" JSONB,
+    "created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
+    "updated_at" TIMESTAMP(3) NOT NULL,
+    "litellm_credential_name" TEXT,
+
+    CONSTRAINT "LiteLLM_ManagedVectorStoresTable_pkey" PRIMARY KEY ("vector_store_id")
+);
+

litellm_proxy_extras/migrations/20250507161526_add_mcp_table_to_db/migration.sql

@@ -0,0 +1,17 @@
+-- CreateTable
+CREATE TABLE "LiteLLM_MCPServerTable" (
+    "server_id" TEXT NOT NULL,
+    "alias" TEXT,
+    "description" TEXT,
+    "url" TEXT NOT NULL,
+    "transport" TEXT NOT NULL DEFAULT 'sse',
+    "spec_version" TEXT NOT NULL DEFAULT '2025-03-26',
+    "auth_type" TEXT,
+    "created_at" TIMESTAMP(3) DEFAULT CURRENT_TIMESTAMP,
+    "created_by" TEXT,
+    "updated_at" TIMESTAMP(3) DEFAULT CURRENT_TIMESTAMP,
+    "updated_by" TEXT,
+
+    CONSTRAINT "LiteLLM_MCPServerTable_pkey" PRIMARY KEY ("server_id")
+);
+

litellm_proxy_extras/migrations/20250507184818_add_mcp_key_team_permission_mgmt/migration.sql

@@ -0,0 +1,32 @@
+-- AlterTable
+ALTER TABLE "LiteLLM_OrganizationTable" ADD COLUMN     "object_permission_id" TEXT;
+
+-- AlterTable
+ALTER TABLE "LiteLLM_TeamTable" ADD COLUMN     "object_permission_id" TEXT;
+
+-- AlterTable
+ALTER TABLE "LiteLLM_UserTable" ADD COLUMN     "object_permission_id" TEXT;
+
+-- AlterTable
+ALTER TABLE "LiteLLM_VerificationToken" ADD COLUMN     "object_permission_id" TEXT;
+
+-- CreateTable
+CREATE TABLE "LiteLLM_ObjectPermissionTable" (
+    "object_permission_id" TEXT NOT NULL,
+    "mcp_servers" TEXT[] DEFAULT ARRAY[]::TEXT[],
+
+    CONSTRAINT "LiteLLM_ObjectPermissionTable_pkey" PRIMARY KEY ("object_permission_id")
+);
+
+-- AddForeignKey
+ALTER TABLE "LiteLLM_OrganizationTable" ADD CONSTRAINT "LiteLLM_OrganizationTable_object_permission_id_fkey" FOREIGN KEY ("object_permission_id") REFERENCES "LiteLLM_ObjectPermissionTable"("object_permission_id") ON DELETE SET NULL ON UPDATE CASCADE;
+
+-- AddForeignKey
+ALTER TABLE "LiteLLM_TeamTable" ADD CONSTRAINT "LiteLLM_TeamTable_object_permission_id_fkey" FOREIGN KEY ("object_permission_id") REFERENCES "LiteLLM_ObjectPermissionTable"("object_permission_id") ON DELETE SET NULL ON UPDATE CASCADE;
+
+-- AddForeignKey
+ALTER TABLE "LiteLLM_UserTable" ADD CONSTRAINT "LiteLLM_UserTable_object_permission_id_fkey" FOREIGN KEY ("object_permission_id") REFERENCES "LiteLLM_ObjectPermissionTable"("object_permission_id") ON DELETE SET NULL ON UPDATE CASCADE;
+
+-- AddForeignKey
+ALTER TABLE "LiteLLM_VerificationToken" ADD CONSTRAINT "LiteLLM_VerificationToken_object_permission_id_fkey" FOREIGN KEY ("object_permission_id") REFERENCES "LiteLLM_ObjectPermissionTable"("object_permission_id") ON DELETE SET NULL ON UPDATE CASCADE;
+

litellm_proxy_extras/migrations/20250508072103_add_status_to_spendlogs/migration.sql

@@ -0,0 +1,3 @@
+-- AlterTable
+ALTER TABLE "LiteLLM_SpendLogs" ADD COLUMN     "status" TEXT;
+

litellm_proxy_extras/migrations/20250509141545_use_big_int_for_daily_spend_tables/migration.sql

@@ -0,0 +1,27 @@
+-- AlterTable
+ALTER TABLE "LiteLLM_DailyTagSpend" ALTER COLUMN "prompt_tokens" SET DATA TYPE BIGINT,
+ALTER COLUMN "completion_tokens" SET DATA TYPE BIGINT,
+ALTER COLUMN "cache_read_input_tokens" SET DATA TYPE BIGINT,
+ALTER COLUMN "cache_creation_input_tokens" SET DATA TYPE BIGINT,
+ALTER COLUMN "api_requests" SET DATA TYPE BIGINT,
+ALTER COLUMN "successful_requests" SET DATA TYPE BIGINT,
+ALTER COLUMN "failed_requests" SET DATA TYPE BIGINT;
+
+-- AlterTable
+ALTER TABLE "LiteLLM_DailyTeamSpend" ALTER COLUMN "prompt_tokens" SET DATA TYPE BIGINT,
+ALTER COLUMN "completion_tokens" SET DATA TYPE BIGINT,
+ALTER COLUMN "api_requests" SET DATA TYPE BIGINT,
+ALTER COLUMN "successful_requests" SET DATA TYPE BIGINT,
+ALTER COLUMN "failed_requests" SET DATA TYPE BIGINT,
+ALTER COLUMN "cache_creation_input_tokens" SET DATA TYPE BIGINT,
+ALTER COLUMN "cache_read_input_tokens" SET DATA TYPE BIGINT;
+
+-- AlterTable
+ALTER TABLE "LiteLLM_DailyUserSpend" ALTER COLUMN "prompt_tokens" SET DATA TYPE BIGINT,
+ALTER COLUMN "completion_tokens" SET DATA TYPE BIGINT,
+ALTER COLUMN "api_requests" SET DATA TYPE BIGINT,
+ALTER COLUMN "failed_requests" SET DATA TYPE BIGINT,
+ALTER COLUMN "successful_requests" SET DATA TYPE BIGINT,
+ALTER COLUMN "cache_creation_input_tokens" SET DATA TYPE BIGINT,
+ALTER COLUMN "cache_read_input_tokens" SET DATA TYPE BIGINT;
+

litellm_proxy_extras/migrations/20250510142544_add_session_id_index_spend_logs/migration.sql

@@ -0,0 +1,3 @@
+-- CreateIndex
+CREATE INDEX "LiteLLM_SpendLogs_session_id_idx" ON "LiteLLM_SpendLogs"("session_id");
+

litellm_proxy_extras/schema.prisma

@@ -61,6 +61,7 @@ model LiteLLM_OrganizationTable {
     models     String[]
     spend      Float    @default(0.0)
     model_spend      Json @default("{}")
+    object_permission_id String?
     created_at    DateTime               @default(now()) @map("created_at")
     created_by String
     updated_at    DateTime               @default(now()) @updatedAt @map("updated_at")
@@ -70,6 +71,7 @@ model LiteLLM_OrganizationTable {
     users LiteLLM_UserTable[]
     keys LiteLLM_VerificationToken[]
     members LiteLLM_OrganizationMembership[] @relation("OrganizationToMembership")
+    object_permission LiteLLM_ObjectPermissionTable?   @relation(fields: [object_permission_id], references: [object_permission_id])
 }
 
 // Model info for teams, just has model aliases for now.
@@ -89,6 +91,7 @@ model LiteLLM_TeamTable {
 		team_id    String @id @default(uuid())
     team_alias  String? 
     organization_id String?
+    object_permission_id String?
     admins String[]
     members String[]
     members_with_roles Json @default("{}")
@@ -110,6 +113,7 @@ model LiteLLM_TeamTable {
     model_id Int? @unique // id for LiteLLM_ModelTable -> stores team-level model aliases
     litellm_organization_table LiteLLM_OrganizationTable?   @relation(fields: [organization_id], references: [organization_id])
     litellm_model_table LiteLLM_ModelTable? @relation(fields: [model_id], references: [id])
+    object_permission LiteLLM_ObjectPermissionTable?   @relation(fields: [object_permission_id], references: [object_permission_id])
 }
 
 // Track spend, rate limit, budget Users
@@ -119,6 +123,7 @@ model LiteLLM_UserTable {
     team_id    String?
     sso_user_id String? @unique
     organization_id String?
+    object_permission_id String?
     password  String?
     teams    String[] @default([])
     user_role  String?
@@ -144,6 +149,32 @@ model LiteLLM_UserTable {
     invitations_created LiteLLM_InvitationLink[] @relation("CreatedBy")
     invitations_updated LiteLLM_InvitationLink[] @relation("UpdatedBy")
     invitations_user    LiteLLM_InvitationLink[] @relation("UserId")
+    object_permission LiteLLM_ObjectPermissionTable?   @relation(fields: [object_permission_id], references: [object_permission_id])
+}
+
+model LiteLLM_ObjectPermissionTable {
+  object_permission_id  String         @id @default(uuid())
+  mcp_servers           String[]       @default([])
+  
+  teams                 LiteLLM_TeamTable[]
+  verification_tokens   LiteLLM_VerificationToken[]
+  organizations         LiteLLM_OrganizationTable[]
+  users                 LiteLLM_UserTable[]
+}
+
+// Holds the MCP server configuration 
+model LiteLLM_MCPServerTable {
+  server_id    String         @id @default(uuid())
+  alias        String?
+  description  String?
+  url          String
+  transport    String         @default("sse")
+  spec_version String         @default("2025-03-26")
+  auth_type    String?        
+  created_at   DateTime?      @default(now()) @map("created_at")
+  created_by   String?
+  updated_at   DateTime?      @default(now()) @updatedAt @map("updated_at")
+  updated_by   String?
 }
 
 // Generate Tokens for Proxy
@@ -169,16 +200,19 @@ model LiteLLM_VerificationToken {
     budget_duration String? 
     budget_reset_at DateTime?
     allowed_cache_controls String[] @default([])
+    allowed_routes   String[] @default([])
     model_spend      Json @default("{}")
     model_max_budget Json @default("{}")
     budget_id String?
     organization_id String?
+    object_permission_id String?
     created_at      DateTime?               @default(now()) @map("created_at")
     created_by String?
     updated_at      DateTime?               @default(now()) @updatedAt @map("updated_at")
     updated_by String?
     litellm_budget_table LiteLLM_BudgetTable?   @relation(fields: [budget_id], references: [budget_id])
     litellm_organization_table LiteLLM_OrganizationTable?   @relation(fields: [organization_id], references: [organization_id])
+    object_permission LiteLLM_ObjectPermissionTable?   @relation(fields: [object_permission_id], references: [object_permission_id])
 }
 
 model LiteLLM_EndUserTable {
@@ -225,8 +259,12 @@ model LiteLLM_SpendLogs {
   requester_ip_address String?
   messages            Json?     @default("{}")
   response            Json?     @default("{}")
+  session_id          String?
+  status              String?
+  proxy_server_request Json?     @default("{}")
   @@index([startTime])
   @@index([end_user])
+  @@index([session_id])
 }
 
 // View spend, model, api_key per request
@@ -324,12 +362,14 @@ model LiteLLM_DailyUserSpend {
   model               String   
   model_group         String?  
   custom_llm_provider String?  
-  prompt_tokens       Int      @default(0)
-  completion_tokens   Int      @default(0)
+  prompt_tokens       BigInt      @default(0)
+  completion_tokens   BigInt      @default(0)
+  cache_read_input_tokens     BigInt      @default(0)
+  cache_creation_input_tokens BigInt      @default(0)
   spend               Float    @default(0.0)
-  api_requests        Int      @default(0)
-  successful_requests Int      @default(0)
-  failed_requests     Int      @default(0)
+  api_requests        BigInt      @default(0)
+  successful_requests BigInt      @default(0)
+  failed_requests     BigInt      @default(0)
   created_at          DateTime @default(now())
   updated_at          DateTime @updatedAt
 
@@ -340,6 +380,60 @@ model LiteLLM_DailyUserSpend {
   @@index([model])
 }
 
+// Track daily team spend metrics per model and key
+model LiteLLM_DailyTeamSpend {
+  id                  String   @id @default(uuid())
+  team_id             String
+  date                String
+  api_key             String   
+  model               String   
+  model_group         String?  
+  custom_llm_provider String?  
+  prompt_tokens       BigInt      @default(0)
+  completion_tokens   BigInt      @default(0)
+  cache_read_input_tokens     BigInt      @default(0)
+  cache_creation_input_tokens BigInt      @default(0)
+  spend               Float    @default(0.0)
+  api_requests        BigInt      @default(0)
+  successful_requests BigInt      @default(0)
+  failed_requests     BigInt      @default(0)
+  created_at          DateTime @default(now())
+  updated_at          DateTime @updatedAt
+
+  @@unique([team_id, date, api_key, model, custom_llm_provider])
+  @@index([date])
+  @@index([team_id])
+  @@index([api_key])
+  @@index([model])
+}
+
+// Track daily team spend metrics per model and key
+model LiteLLM_DailyTagSpend {
+  id                  String   @id @default(uuid())
+  tag                 String   
+  date                String
+  api_key             String   
+  model               String   
+  model_group         String?  
+  custom_llm_provider String?  
+  prompt_tokens       BigInt      @default(0)
+  completion_tokens   BigInt      @default(0)
+  cache_read_input_tokens     BigInt      @default(0)
+  cache_creation_input_tokens BigInt      @default(0)
+  spend               Float    @default(0.0)
+  api_requests        BigInt      @default(0)
+  successful_requests BigInt      @default(0)
+  failed_requests     BigInt      @default(0)
+  created_at          DateTime @default(now())
+  updated_at          DateTime @updatedAt
+
+  @@unique([tag, date, api_key, model, custom_llm_provider])
+  @@index([date])
+  @@index([tag])
+  @@index([api_key])
+  @@index([model])
+}
+
 
 // Track the status of cron jobs running. Only allow one pod to run the job at a time
 model LiteLLM_CronJob {
@@ -366,3 +460,14 @@ model LiteLLM_ManagedFileTable {
   @@index([unified_file_id])
 }
 
+
+model LiteLLM_ManagedVectorStoresTable {
+  vector_store_id String @id
+  custom_llm_provider String
+  vector_store_name String?
+  vector_store_description String?
+  vector_store_metadata Json?
+  created_at DateTime @default(now())
+  updated_at DateTime @updatedAt
+  litellm_credential_name String?
+}

litellm_proxy_extras/utils.py

@@ -1,8 +1,11 @@
 import glob
 import os
 import random
+import re
+import shutil
 import subprocess
 import time
+from datetime import datetime
 from pathlib import Path
 from typing import Optional
 
@@ -18,9 +21,30 @@ def str_to_bool(value: Optional[str]) -> bool:
 class ProxyExtrasDBManager:
     @staticmethod
     def _get_prisma_dir() -> str:
-        """Get the path to the migrations directory"""
-        migrations_dir = os.path.dirname(__file__)
-        return migrations_dir
+        """
+        Get the path to the migrations directory
+
+        Set os.environ["LITELLM_MIGRATION_DIR"] to a custom migrations directory, to support baselining db in read-only fs.
+        """
+        custom_migrations_dir = os.getenv("LITELLM_MIGRATION_DIR")
+        pkg_migrations_dir = os.path.dirname(__file__)
+        if custom_migrations_dir:
+            # If migrations_dir exists, copy contents
+            if os.path.exists(custom_migrations_dir):
+                # Copy contents instead of directory itself
+                for item in os.listdir(pkg_migrations_dir):
+                    src_path = os.path.join(pkg_migrations_dir, item)
+                    dst_path = os.path.join(custom_migrations_dir, item)
+                    if os.path.isdir(src_path):
+                        shutil.copytree(src_path, dst_path, dirs_exist_ok=True)
+                    else:
+                        shutil.copy2(src_path, dst_path)
+            else:
+                # If directory doesn't exist, create it and copy everything
+                shutil.copytree(pkg_migrations_dir, custom_migrations_dir)
+            return custom_migrations_dir
+
+        return pkg_migrations_dir
 
     @staticmethod
     def _create_baseline_migration(schema_path: str) -> bool:
@@ -32,27 +56,29 @@ class ProxyExtrasDBManager:
         # Create migrations/0_init directory
         init_dir.mkdir(parents=True, exist_ok=True)
 
-        # Generate migration SQL file
-        migration_file = init_dir / "migration.sql"
+        database_url = os.getenv("DATABASE_URL")
 
         try:
-            # Generate migration diff with increased timeout
+            # 1. Generate migration SQL file by comparing empty state to current db state
+            logger.info("Generating baseline migration...")
+            migration_file = init_dir / "migration.sql"
             subprocess.run(
                 [
                     "prisma",
                     "migrate",
                     "diff",
                     "--from-empty",
-                    "--to-schema-datamodel",
-                    str(schema_path),
+                    "--to-url",
+                    database_url,
                     "--script",
                 ],
                 stdout=open(migration_file, "w"),
                 check=True,
                 timeout=30,
-            )  # 30 second timeout
+            )
 
-            # Mark migration as applied with increased timeout
+            # 3. Mark the migration as applied since it represents current state
+            logger.info("Marking baseline migration as applied...")
             subprocess.run(
                 [
                     "prisma",
@@ -72,8 +98,10 @@ class ProxyExtrasDBManager:
             )
             return False
         except subprocess.CalledProcessError as e:
-            logger.warning(f"Error creating baseline migration: {e}")
-            return False
+            logger.warning(
+                f"Error creating baseline migration: {e}, {e.stderr}, {e.stdout}"
+            )
+            raise e
 
     @staticmethod
     def _get_migration_names(migrations_dir: str) -> list:
@@ -83,8 +111,105 @@ class ProxyExtrasDBManager:
         return [Path(p).parent.name for p in migration_paths]
 
     @staticmethod
-    def _resolve_all_migrations(migrations_dir: str):
-        """Mark all existing migrations as applied"""
+    def _roll_back_migration(migration_name: str):
+        """Mark a specific migration as rolled back"""
+        subprocess.run(
+            ["prisma", "migrate", "resolve", "--rolled-back", migration_name],
+            timeout=60,
+            check=True,
+            capture_output=True,
+        )
+
+    @staticmethod
+    def _resolve_specific_migration(migration_name: str):
+        """Mark a specific migration as applied"""
+        subprocess.run(
+            ["prisma", "migrate", "resolve", "--applied", migration_name],
+            timeout=60,
+            check=True,
+            capture_output=True,
+        )
+
+    @staticmethod
+    def _resolve_all_migrations(migrations_dir: str, schema_path: str):
+        """
+        1. Compare the current database state to schema.prisma and generate a migration for the diff.
+        2. Run prisma migrate deploy to apply any pending migrations.
+        3. Mark all existing migrations as applied.
+        """
+        database_url = os.getenv("DATABASE_URL")
+        diff_dir = (
+            Path(migrations_dir)
+            / "migrations"
+            / f"{datetime.now().strftime('%Y%m%d%H%M%S')}_baseline_diff"
+        )
+        try:
+            diff_dir.mkdir(parents=True, exist_ok=True)
+        except Exception as e:
+            if "Permission denied" in str(e):
+                logger.warning(
+                    f"Permission denied - {e}\nunable to baseline db. Set LITELLM_MIGRATION_DIR environment variable to a writable directory to enable migrations."
+                )
+                return
+            raise e
+        diff_sql_path = diff_dir / "migration.sql"
+
+        # 1. Generate migration SQL for the diff between DB and schema
+        try:
+            logger.info("Generating migration diff between DB and schema.prisma...")
+            with open(diff_sql_path, "w") as f:
+                subprocess.run(
+                    [
+                        "prisma",
+                        "migrate",
+                        "diff",
+                        "--from-url",
+                        database_url,
+                        "--to-schema-datamodel",
+                        schema_path,
+                        "--script",
+                    ],
+                    check=True,
+                    timeout=60,
+                    stdout=f,
+                )
+        except subprocess.CalledProcessError as e:
+            logger.warning(f"Failed to generate migration diff: {e.stderr}")
+        except subprocess.TimeoutExpired:
+            logger.warning("Migration diff generation timed out.")
+
+        # check if the migration was created
+        if not diff_sql_path.exists():
+            logger.warning("Migration diff was not created")
+            return
+        logger.info(f"Migration diff created at {diff_sql_path}")
+
+        # 2. Run prisma db execute to apply the migration
+        try:
+            logger.info("Running prisma db execute to apply the migration diff...")
+            result = subprocess.run(
+                [
+                    "prisma",
+                    "db",
+                    "execute",
+                    "--file",
+                    str(diff_sql_path),
+                    "--schema",
+                    schema_path,
+                ],
+                timeout=60,
+                check=True,
+                capture_output=True,
+                text=True,
+            )
+            logger.info(f"prisma db execute stdout: {result.stdout}")
+            logger.info("✅ Migration diff applied successfully")
+        except subprocess.CalledProcessError as e:
+            logger.warning(f"Failed to apply migration diff: {e.stderr}")
+        except subprocess.TimeoutExpired:
+            logger.warning("Migration diff application timed out.")
+
+        # 3. Mark all migrations as applied
         migration_names = ProxyExtrasDBManager._get_migration_names(migrations_dir)
         logger.info(f"Resolving {len(migration_names)} migrations")
         for migration_name in migration_names:
@@ -105,7 +230,7 @@ class ProxyExtrasDBManager:
                     )
 
     @staticmethod
-    def setup_database(schema_path: str, use_migrate: bool = False) -> bool:
+    def setup_database(use_migrate: bool = False) -> bool:
         """
         Set up the database using either prisma migrate or prisma db push
         Uses migrations from litellm-proxy-extras package
@@ -117,6 +242,7 @@ class ProxyExtrasDBManager:
         Returns:
             bool: True if setup was successful, False otherwise
         """
+        schema_path = ProxyExtrasDBManager._get_prisma_dir() + "/schema.prisma"
         use_migrate = str_to_bool(os.getenv("USE_PRISMA_MIGRATE")) or use_migrate
         for attempt in range(4):
             original_dir = os.getcwd()
@@ -141,7 +267,34 @@ class ProxyExtrasDBManager:
                         return True
                     except subprocess.CalledProcessError as e:
                         logger.info(f"prisma db error: {e.stderr}, e: {e.stdout}")
-                        if (
+                        if "P3009" in e.stderr:
+                            # Extract the failed migration name from the error message
+                            migration_match = re.search(
+                                r"`(\d+_.*)` migration", e.stderr
+                            )
+                            if migration_match:
+                                failed_migration = migration_match.group(1)
+                                logger.info(
+                                    f"Found failed migration: {failed_migration}, marking as rolled back"
+                                )
+                                # Mark the failed migration as rolled back
+                                subprocess.run(
+                                    [
+                                        "prisma",
+                                        "migrate",
+                                        "resolve",
+                                        "--rolled-back",
+                                        failed_migration,
+                                    ],
+                                    timeout=60,
+                                    check=True,
+                                    capture_output=True,
+                                    text=True,
+                                )
+                                logger.info(
+                                    f"✅ Migration {failed_migration} marked as rolled back... retrying"
+                                )
+                        elif (
                             "P3005" in e.stderr
                             and "database schema is not empty" in e.stderr
                         ):
@@ -152,9 +305,34 @@ class ProxyExtrasDBManager:
                             logger.info(
                                 "Baseline migration created, resolving all migrations"
                             )
-                            ProxyExtrasDBManager._resolve_all_migrations(migrations_dir)
+                            ProxyExtrasDBManager._resolve_all_migrations(
+                                migrations_dir, schema_path
+                            )
                             logger.info("✅ All migrations resolved.")
                             return True
+                        elif (
+                            "P3018" in e.stderr
+                        ):  # PostgreSQL error code for duplicate column
+                            logger.info(
+                                "Migration already exists, resolving specific migration"
+                            )
+                            # Extract the migration name from the error message
+                            migration_match = re.search(
+                                r"Migration name: (\d+_.*)", e.stderr
+                            )
+                            if migration_match:
+                                migration_name = migration_match.group(1)
+                                logger.info(f"Rolling back migration {migration_name}")
+                                ProxyExtrasDBManager._roll_back_migration(
+                                    migration_name
+                                )
+                                logger.info(
+                                    f"Resolving migration {migration_name} that failed due to existing columns"
+                                )
+                                ProxyExtrasDBManager._resolve_specific_migration(
+                                    migration_name
+                                )
+                                logger.info("✅ Migration resolved.")
                 else:
                     # Use prisma db push with increased timeout
                     subprocess.run(

{litellm_proxy_extras-0.1.8.dist-info → litellm_proxy_extras-0.1.19.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: litellm-proxy-extras
-Version: 0.1.8
+Version: 0.1.19
 Summary: Additional files for the LiteLLM Proxy. Reduces the size of the main litellm package.
 Author: BerriAI
 Requires-Python: >=3.8, !=2.7.*, !=3.0.*, !=3.1.*, !=3.2.*, !=3.3.*, !=3.4.*, !=3.5.*, !=3.6.*, !=3.7.*

litellm_proxy_extras-0.1.19.dist-info/RECORD

@@ -0,0 +1,28 @@
+litellm_proxy_extras/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+litellm_proxy_extras/_logging.py,sha256=7KoWerTOol5IPNyNdbZvLuSlpQbEGE235VgzwpgafKQ,393
+litellm_proxy_extras/migrations/20250326162113_baseline/migration.sql,sha256=9aRWmBbLf7EWbCMXifDl5zL9bAw0uPXJut1AXNKrSTE,13383
+litellm_proxy_extras/migrations/20250326171002_add_daily_user_table/migration.sql,sha256=dY-dNCLosWmXNli2B9wqX4hZpp3s0DL3IwEPtTTC134,1179
+litellm_proxy_extras/migrations/20250327180120_add_api_requests_to_daily_user_table/migration.sql,sha256=or5TaEgH4cHwR5kDvVjZvcAj1OwzTxjqwO-lxXe3FXk,110
+litellm_proxy_extras/migrations/20250329084805_new_cron_job_table/migration.sql,sha256=eZNDwrzKtWFXkTqOKb9JS4hzum1dI-VTXvMqzhryfxc,404
+litellm_proxy_extras/migrations/20250331215456_track_success_and_failed_requests_daily_agg_table/migration.sql,sha256=tyeLY6u8KFyw71osCBM-sdjhIgvHoFCK88cIx8dExNY,178
+litellm_proxy_extras/migrations/20250411215431_add_managed_file_table/migration.sql,sha256=Yu2K37Q90LDhxsFo_64sH0PXdSQ3sHs45Lqzxv2t_20,625
+litellm_proxy_extras/migrations/20250412081753_team_member_permissions/migration.sql,sha256=v3vDx5lb6SLCzXCe_A2NZj7zzmucRXM08aQun_G_MkE,120
+litellm_proxy_extras/migrations/20250415151647_add_cache_read_write_tokens_daily_spend_transactions/migration.sql,sha256=pXUhTLpyXLhGaDnT-epbyW-YSdmO1z5_dxOsETD53Os,194
+litellm_proxy_extras/migrations/20250415191926_add_daily_team_table/migration.sql,sha256=M8DFQOxIPjGcJXIswfHQs2LBTFdJyQBpRbpOT5RIHmw,1330
+litellm_proxy_extras/migrations/20250416115320_add_tag_table_to_db/migration.sql,sha256=meXcoLzZI8-huUGfeJxB9bcOFLTC83mlnpN5LvYbd7I,1720
+litellm_proxy_extras/migrations/20250416151339_drop_tag_uniqueness_requirement/migration.sql,sha256=ezPzK5yd-RDziBLt6X-BEoWYmNYHUublg19OobGuuEs,58
+litellm_proxy_extras/migrations/20250416185146_add_allowed_routes_litellm_verification_token/migration.sql,sha256=XGyDLGakqBt30xuGS-RT5tWhk_Ki7teenxpvKDkIo2E,119
+litellm_proxy_extras/migrations/20250425182129_add_session_id/migration.sql,sha256=dSjU5QavHMM60siy6f_69DZpFCA4lUfWDL0AUEPEt2g,139
+litellm_proxy_extras/migrations/20250430193429_add_managed_vector_stores/migration.sql,sha256=SmXedKWYWwYAZjg-usw-905nXxyU9SOws9teFATezUA,479
+litellm_proxy_extras/migrations/20250507161526_add_mcp_table_to_db/migration.sql,sha256=C_0iU6F-XUST0BekRSrFGMW69_pwBLSS7Pu5HvGNJ7c,511
+litellm_proxy_extras/migrations/20250507184818_add_mcp_key_team_permission_mgmt/migration.sql,sha256=G2NONS9XAMDdkH-0tpaPimK4ezs94hkwteRK1m7oA1Y,1672
+litellm_proxy_extras/migrations/20250508072103_add_status_to_spendlogs/migration.sql,sha256=V4mIhoi3n8sUzV8_YY9AA_wkgvk65m5ZOAU3_DtoUo8,77
+litellm_proxy_extras/migrations/20250509141545_use_big_int_for_daily_spend_tables/migration.sql,sha256=UPaAXsryGvVZT_SnOf-1yHu5bKPTVr2gP3zY64FKR7o,1331
+litellm_proxy_extras/migrations/20250510142544_add_session_id_index_spend_logs/migration.sql,sha256=M_y4xLnlGApiUIovLv0xA7lL5gJBZWrcFe7b5MSUOik,102
+litellm_proxy_extras/migrations/migration_lock.toml,sha256=HbF6jQUaoTYRBzZ1LF4fi37ZK26o6AMRL7viSXBHwhA,24
+litellm_proxy_extras/schema.prisma,sha256=k3TXUCJaMkZ1q_9Je2V0aoeYYw2osqyMFUTSxHRLfEI,18573
+litellm_proxy_extras/utils.py,sha256=QImDkrfX_NGCp0tJ4T9IZpmXyzVZqey4jnVVEO0Hcrs,15050
+litellm_proxy_extras-0.1.19.dist-info/LICENSE,sha256=sXDWv46INd01fgEWgdsCj01R4vsOqJIFj1bgH7ObgnM,1419
+litellm_proxy_extras-0.1.19.dist-info/METADATA,sha256=hauDdaBeDu-id2S09n3YVOUleHcaYahlF9YPhob2ssM,1268
+litellm_proxy_extras-0.1.19.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+litellm_proxy_extras-0.1.19.dist-info/RECORD,,

litellm_proxy_extras-0.1.8.dist-info/RECORD

@@ -1,17 +0,0 @@
-litellm_proxy_extras/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-litellm_proxy_extras/_logging.py,sha256=7KoWerTOol5IPNyNdbZvLuSlpQbEGE235VgzwpgafKQ,393
-litellm_proxy_extras/migrations/20250326162113_baseline/migration.sql,sha256=9aRWmBbLf7EWbCMXifDl5zL9bAw0uPXJut1AXNKrSTE,13383
-litellm_proxy_extras/migrations/20250326171002_add_daily_user_table/migration.sql,sha256=dY-dNCLosWmXNli2B9wqX4hZpp3s0DL3IwEPtTTC134,1179
-litellm_proxy_extras/migrations/20250327180120_add_api_requests_to_daily_user_table/migration.sql,sha256=or5TaEgH4cHwR5kDvVjZvcAj1OwzTxjqwO-lxXe3FXk,110
-litellm_proxy_extras/migrations/20250329084805_new_cron_job_table/migration.sql,sha256=eZNDwrzKtWFXkTqOKb9JS4hzum1dI-VTXvMqzhryfxc,404
-litellm_proxy_extras/migrations/20250331215456_track_success_and_failed_requests_daily_agg_table/migration.sql,sha256=tyeLY6u8KFyw71osCBM-sdjhIgvHoFCK88cIx8dExNY,178
-litellm_proxy_extras/migrations/20250411215431_add_managed_file_table/migration.sql,sha256=Yu2K37Q90LDhxsFo_64sH0PXdSQ3sHs45Lqzxv2t_20,625
-litellm_proxy_extras/migrations/20250412081753_team_member_permissions/migration.sql,sha256=v3vDx5lb6SLCzXCe_A2NZj7zzmucRXM08aQun_G_MkE,120
-litellm_proxy_extras/migrations/20250415191926_add_daily_team_table/migration.sql,sha256=M8DFQOxIPjGcJXIswfHQs2LBTFdJyQBpRbpOT5RIHmw,1330
-litellm_proxy_extras/migrations/migration_lock.toml,sha256=HbF6jQUaoTYRBzZ1LF4fi37ZK26o6AMRL7viSXBHwhA,24
-litellm_proxy_extras/schema.prisma,sha256=AJs-oTl488I539rej87DQw73l2x28IhLNGw2gZ6T8AA,14579
-litellm_proxy_extras/utils.py,sha256=2CQEBosLTjXpGEBwx3sGYglSD4QNy609iV4Ppo9AJdo,7087
-litellm_proxy_extras-0.1.8.dist-info/LICENSE,sha256=sXDWv46INd01fgEWgdsCj01R4vsOqJIFj1bgH7ObgnM,1419
-litellm_proxy_extras-0.1.8.dist-info/METADATA,sha256=gnPQXPGAp6L1crea8OX5ucn3kuqPRhjOPaAW-LlxGBU,1267
-litellm_proxy_extras-0.1.8.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-litellm_proxy_extras-0.1.8.dist-info/RECORD,,