diff --git a/autogpt_platform/backend/backend/api/features/admin/llm_routes.py b/autogpt_platform/backend/backend/api/features/admin/llm_routes.py index dd87547cee..c0837ce355 100644 --- a/autogpt_platform/backend/backend/api/features/admin/llm_routes.py +++ b/autogpt_platform/backend/backend/api/features/admin/llm_routes.py @@ -22,7 +22,7 @@ async def _refresh_runtime_state() -> None: try: # Refresh registry from database await llm_registry.refresh_llm_registry() - refresh_llm_costs() + await refresh_llm_costs() # Clear block schema caches so they're regenerated with updated model options from backend.data.block import BlockSchema diff --git a/autogpt_platform/backend/backend/api/rest_api.py b/autogpt_platform/backend/backend/api/rest_api.py index 86f77a2c71..da125ad048 100644 --- a/autogpt_platform/backend/backend/api/rest_api.py +++ b/autogpt_platform/backend/backend/api/rest_api.py @@ -120,7 +120,7 @@ async def lifespan_context(app: fastapi.FastAPI): # Refresh LLM registry before initializing blocks so blocks can use registry data await llm_registry.refresh_llm_registry() - refresh_llm_costs() + await refresh_llm_costs() # Clear block schema caches so they're regenerated with updated discriminator_mapping from backend.data.block import BlockSchema diff --git a/autogpt_platform/backend/backend/data/block.py b/autogpt_platform/backend/backend/data/block.py index bf01addf13..9be954c5ff 100644 --- a/autogpt_platform/backend/backend/data/block.py +++ b/autogpt_platform/backend/backend/data/block.py @@ -911,7 +911,7 @@ async def initialize_blocks() -> None: if is_connected(): await llm_registry.refresh_llm_registry() - refresh_llm_costs() + await refresh_llm_costs() logger.info("LLM registry refreshed during block initialization") else: logger.warning( diff --git a/autogpt_platform/backend/backend/data/block_cost_config.py b/autogpt_platform/backend/backend/data/block_cost_config.py index 69a4e633dc..3b2dd8534a 100644 --- a/autogpt_platform/backend/backend/data/block_cost_config.py +++ b/autogpt_platform/backend/backend/data/block_cost_config.py @@ -1,6 +1,8 @@ import logging from typing import Type +import prisma.models + from backend.blocks.ai_image_customizer import AIImageCustomizerBlock, GeminiImageModel from backend.blocks.ai_image_generator_block import AIImageGeneratorBlock, ImageGenModel from backend.blocks.ai_music_generator import AIMusicGeneratorBlock @@ -75,8 +77,40 @@ PROVIDER_CREDENTIALS = { LLM_COST: list[BlockCost] = [] -def _build_llm_costs_from_registry() -> list[BlockCost]: - """Build BlockCost list from all models in the LLM registry.""" +async def _build_llm_costs_from_registry() -> list[BlockCost]: + """ + Build BlockCost list from all models in the LLM registry. + + This function checks for active model migrations with customCreditCost overrides. + When a model has been migrated with a custom price, that price is used instead + of the target model's default cost. + """ + # Query active migrations with custom pricing overrides + migration_overrides: dict[str, int] = {} + try: + active_migrations = await prisma.models.LlmModelMigration.prisma().find_many( + where={ + "isReverted": False, + "customCreditCost": {"not": None}, + } + ) + migration_overrides = { + migration.sourceModelSlug: migration.customCreditCost + for migration in active_migrations + if migration.customCreditCost is not None + } + if migration_overrides: + logger.info( + "Found %d active model migrations with custom pricing overrides", + len(migration_overrides), + ) + except Exception as exc: + logger.warning( + "Failed to query model migration overrides: %s. Proceeding with default costs.", + exc, + exc_info=True, + ) + costs: list[BlockCost] = [] for model in llm_registry.iter_dynamic_models(): for cost in model.costs: @@ -88,6 +122,18 @@ def _build_llm_costs_from_registry() -> list[BlockCost]: cost.credential_provider, ) continue + + # Check if this model has a custom cost override from migration + cost_amount = migration_overrides.get(model.slug, cost.credit_cost) + + if model.slug in migration_overrides: + logger.debug( + "Applying custom cost override for model %s: %d credits (default: %d)", + model.slug, + cost_amount, + cost.credit_cost, + ) + cost_filter = { "model": model.slug, "credentials": { @@ -100,16 +146,21 @@ def _build_llm_costs_from_registry() -> list[BlockCost]: BlockCost( cost_type=BlockCostType.RUN, cost_filter=cost_filter, - cost_amount=cost.credit_cost, + cost_amount=cost_amount, ) ) return costs -def refresh_llm_costs() -> None: - """Refresh LLM costs from the registry. All costs now come from the database.""" +async def refresh_llm_costs() -> None: + """ + Refresh LLM costs from the registry. All costs now come from the database. + + This function also checks for active model migrations with custom pricing overrides + and applies them to ensure accurate billing. + """ LLM_COST.clear() - LLM_COST.extend(_build_llm_costs_from_registry()) + LLM_COST.extend(await _build_llm_costs_from_registry()) # Initial load will happen after registry is refreshed at startup diff --git a/autogpt_platform/backend/backend/executor/llm_registry_init.py b/autogpt_platform/backend/backend/executor/llm_registry_init.py index 0779e932da..e55795d8af 100644 --- a/autogpt_platform/backend/backend/executor/llm_registry_init.py +++ b/autogpt_platform/backend/backend/executor/llm_registry_init.py @@ -47,7 +47,7 @@ async def refresh_registry_on_notification() -> None: # Refresh registry and costs await llm_registry.refresh_llm_registry() - refresh_llm_costs() + await refresh_llm_costs() # Clear block schema caches so they regenerate with new model options BlockSchema.clear_all_schema_caches() diff --git a/autogpt_platform/backend/migrations/20260212000000_add_llm_migration_composite_index/migration.sql b/autogpt_platform/backend/migrations/20260212000000_add_llm_migration_composite_index/migration.sql new file mode 100644 index 0000000000..5f975e4c19 --- /dev/null +++ b/autogpt_platform/backend/migrations/20260212000000_add_llm_migration_composite_index/migration.sql @@ -0,0 +1,6 @@ +-- Add composite index on LlmModelMigration for optimized active migration queries +-- This index improves performance when querying for non-reverted migrations by model slug +-- Used by the billing system to apply customCreditCost overrides + +-- CreateIndex +CREATE INDEX "LlmModelMigration_sourceModelSlug_isReverted_idx" ON "LlmModelMigration"("sourceModelSlug", "isReverted"); diff --git a/autogpt_platform/backend/migrations/20260212000001_sync_llm_models_with_dev/migration.sql b/autogpt_platform/backend/migrations/20260212000001_sync_llm_models_with_dev/migration.sql new file mode 100644 index 0000000000..d8b1fda5fe --- /dev/null +++ b/autogpt_platform/backend/migrations/20260212000001_sync_llm_models_with_dev/migration.sql @@ -0,0 +1,61 @@ +-- Sync LLM models with latest dev branch changes +-- This migration adds new models and removes deprecated ones + +-- Remove models that were deleted from dev +DELETE FROM "LlmModelCost" WHERE "llmModelId" IN ( + SELECT "id" FROM "LlmModel" WHERE "slug" IN ('o3', 'o3-mini', 'claude-3-7-sonnet-20250219') +); + +DELETE FROM "LlmModel" WHERE "slug" IN ('o3', 'o3-mini', 'claude-3-7-sonnet-20250219'); + +-- Add new models from dev +WITH provider_ids AS ( + SELECT "id", "name" FROM "LlmProvider" +) +INSERT INTO "LlmModel" ("id", "slug", "displayName", "description", "providerId", "contextWindow", "maxOutputTokens", "isEnabled", "capabilities", "metadata") +SELECT + gen_random_uuid(), + model_slug, + model_display_name, + NULL, + p."id", + context_window, + max_output_tokens, + true, + '{}'::jsonb, + '{}'::jsonb +FROM (VALUES + -- New OpenAI model + ('gpt-5.2-2025-12-11', 'GPT 5.2', 'openai', 400000, 128000), + -- New Anthropic model + ('claude-opus-4-6', 'Claude 4.6 Opus', 'anthropic', 200000, 64000) +) AS models(model_slug, model_display_name, provider_name, context_window, max_output_tokens) +JOIN provider_ids p ON p."name" = models.provider_name +ON CONFLICT ("slug") DO NOTHING; + +-- Add costs for new models +WITH model_ids AS ( + SELECT "id", "slug", "providerId" FROM "LlmModel" +), +provider_ids AS ( + SELECT "id", "name" FROM "LlmProvider" +) +INSERT INTO "LlmModelCost" ("id", "unit", "creditCost", "credentialProvider", "credentialId", "credentialType", "currency", "metadata", "llmModelId") +SELECT + gen_random_uuid(), + 'RUN'::"LlmCostUnit", + cost, + p."name", + NULL, + 'api_key', + NULL, + '{}'::jsonb, + m."id" +FROM (VALUES + -- New model costs (estimate based on similar models) + ('gpt-5.2-2025-12-11', 5), -- Similar to GPT 5.1 + ('claude-opus-4-6', 21) -- Similar to other Opus 4.x models +) AS costs(model_slug, cost) +JOIN model_ids m ON m."slug" = costs.model_slug +JOIN provider_ids p ON p."id" = m."providerId" +ON CONFLICT ("llmModelId", "credentialProvider", "unit") DO NOTHING;