diff --git a/app/api/admin/vip98k-models/route.ts b/app/api/admin/vip98k-models/route.ts index 145884a3b..11929f657 100644 --- a/app/api/admin/vip98k-models/route.ts +++ b/app/api/admin/vip98k-models/route.ts @@ -2,10 +2,15 @@ import { and, eq } from "drizzle-orm"; import { NextResponse } from "next/server"; import { fetchVip98kModelPresets, + hasVip98kTextToImageMode, type Vip98kModelPreset, } from "@/lib/ai/vip98k-models"; import { db } from "@/lib/db"; -import { getAiProviderById } from "@/lib/db/queries"; +import { + deleteAiModel, + getAiModelsByProviderId, + getAiProviderById, +} from "@/lib/db/queries"; import { aiModel, type AiModel } from "@/lib/db/schema"; import { requireAdmin } from "@/lib/middleware/admin"; @@ -97,8 +102,35 @@ export async function POST(request: Request) { const addedModels: AiModel[] = []; const updatedModels: AiModel[] = []; + const removedModels: AiModel[] = []; const errors: Array<{ modelId: string; error: string }> = []; + const existingModels = await getAiModelsByProviderId({ providerId }); + const textToImageModels = existingModels.filter( + (model) => + model.supportsImageGen && + hasVip98kTextToImageMode({ + modelId: model.modelId, + capabilities: model.capabilities, + defaultParams: model.defaultParams, + }), + ); + for (const model of textToImageModels) { + try { + const deleted = await deleteAiModel({ id: model.id }); + removedModels.push(...deleted); + } catch (error) { + console.error( + `[VIP98K Models API] Failed to remove text-to-image model ${model.modelId}:`, + error, + ); + errors.push({ + modelId: model.modelId, + error: error instanceof Error ? error.message : "Removal failed", + }); + } + } + for (const preset of presets) { try { const existing = await db.query.aiModel.findFirst({ @@ -161,8 +193,10 @@ export async function POST(request: Request) { success: true, addedCount: addedModels.length, updatedCount: updatedModels.length, + removedCount: removedModels.length, addedModels, updatedModels, + removedModels, errors: errors.length > 0 ? errors : undefined, counts: { requested: presets.length, diff --git a/lib/ai/vip98k-models.ts b/lib/ai/vip98k-models.ts index 9bc8fca6b..2557c9603 100644 --- a/lib/ai/vip98k-models.ts +++ b/lib/ai/vip98k-models.ts @@ -355,6 +355,40 @@ function getCapabilities(model: Vip98kCatalogModel) { ); } +function normalizeMode(value: unknown) { + return String(value || "").trim().toLowerCase().replace(/_/g, "-"); +} + +function getStoredSupportedModes(value: unknown) { + const record = toRecord(value); + return parseStringArray( + record?.supportedModes || record?.supported_modes, + ).map(normalizeMode); +} + +export function hasVip98kTextToImageMode(input: { + modelId?: string | null; + capabilities?: unknown; + defaultParams?: unknown; +}) { + const modes = new Set([ + ...getStoredSupportedModes(input.capabilities), + ...getStoredSupportedModes(input.defaultParams), + ]); + return ( + modes.has("text-to-image") || + normalizeMode(input.modelId).includes("text-to-image") + ); +} + +export function isVip98kTextToImageCatalogModel( + model: Vip98kCatalogModel, +) { + if (resolvePresetType(model) !== "image") return false; + const capabilities = getCapabilities(model); + return capabilities.length === 0 || capabilities.includes("text-to-image"); +} + function getInputModes(model: Vip98kCatalogModel) { return parseStringArray(model.inputModes || model.input_modes).map((item) => item.trim().toLowerCase(), @@ -771,6 +805,7 @@ export function buildVip98kModelPresets( .map((model, index) => { const modelId = getModelId(model); if (!modelId) return null; + if (isVip98kTextToImageCatalogModel(model)) return null; const type = resolvePresetType(model); if (type === "image") return createImagePreset(model, index); if (type === "video") return createVideoPreset(model, index); diff --git a/tests/routes/model-classification.test.ts b/tests/routes/model-classification.test.ts index 010b1351a..1bfa3074d 100644 --- a/tests/routes/model-classification.test.ts +++ b/tests/routes/model-classification.test.ts @@ -385,35 +385,69 @@ test.describe("AI model media classification", () => { modelId: "black-forest-labs/flux-dev", displayName: "Flux Dev", type: "image", - capabilities: ["text_to_image"], + capabilities: ["image_to_image"], }, { - modelId: "google/nano-banana-2/text-to-image", + modelId: "google/nano-banana-2/edit", displayName: "Nano Banana 2", type: "image", - capabilities: ["text_to_image"], + capabilities: ["image_to_image"], }, { modelId: "future-image-with-date", displayName: "Future image", type: "image", - capabilities: ["text_to_image"], + capabilities: ["image_to_image"], releasedAt: "2027-02-01", }, { modelId: "old-image-with-date", displayName: "Old image", type: "image", - capabilities: ["text_to_image"], + capabilities: ["image_to_image"], releasedAt: "2025-12-15", }, ]); expect(presets.find((item) => item.modelId === "black-forest-labs/flux-dev")?.isHidden).toBe(true); expect( - presets.find((item) => item.modelId === "google/nano-banana-2/text-to-image")?.isHidden, + presets.find((item) => item.modelId === "google/nano-banana-2/edit")?.isHidden, ).toBe(false); expect(presets.find((item) => item.modelId === "future-image-with-date")?.isHidden).toBe(false); expect(presets.find((item) => item.modelId === "old-image-with-date")?.isHidden).toBe(true); }); + + test("excludes VIP98K text-to-image models without applying a price limit", () => { + const presets = buildVip98kModelPresets([ + { + modelId: "pure-text-to-image", + type: "image", + capabilities: ["text_to_image"], + pricing: { effective_price: 0.02 }, + }, + { + modelId: "mixed-image-model", + type: "image", + capabilities: ["text_to_image", "image_to_image"], + pricing: { effective_price: 0.2 }, + }, + { + modelId: "expensive-image-edit", + type: "image", + capabilities: ["image_to_image"], + pricing: { effective_price: 2.5 }, + }, + { + modelId: "remove-background", + type: "image_tool", + capabilities: ["image_processing"], + pricing: { effective_price: 1.2 }, + }, + ]); + + expect(presets.map((item) => item.modelId)).toEqual([ + "expensive-image-edit", + "remove-background", + ]); + }); });