diff --git a/src/api/client.ts b/src/api/client.ts index 94de7e5..009dbfe 100644 --- a/src/api/client.ts +++ b/src/api/client.ts @@ -1,5 +1,6 @@ import type { ResolvedModel } from "../config/types"; import type { CommandType } from "../core/types"; +import { isAiRemoverModel } from "../utils/model-routing.ts"; import { type ApiEnvelope, endpoints, type TaskData } from "./types.ts"; async function httpJson( @@ -121,7 +122,7 @@ export function buildSubmitTarget( const modelRef = model.modelName ?? model.id; const suffix = - model.submitMode === "canonical" + model.submitMode === "canonical" || (command === "edit" && isAiRemoverModel(model)) ? "" : { generate: "", diff --git a/src/config/models.ts b/src/config/models.ts index 212775a..8bd3e82 100644 --- a/src/config/models.ts +++ b/src/config/models.ts @@ -1,4 +1,5 @@ import { ModelCache } from "../cache"; +import { inferApiModelType } from "../utils/model-routing.ts"; import { ConfigError } from "./load"; import { getRegistryModel } from "./registry"; import type { ResolvedModel, ResolvedModelSummary, WavespeedConfig } from "./types"; @@ -8,7 +9,7 @@ import type { ResolvedModel, ResolvedModelSummary, WavespeedConfig } from "./typ * Used to validate model IDs against the API without tight coupling */ export interface ApiModelCache { - models: Array<{ model_id: string }>; + models: Array<{ model_id: string; type?: string }>; } const BUILTIN_MODEL_ID = "seedream-v4"; @@ -46,13 +47,75 @@ export type ModelCommandName = "generate" | "edit" | "generate-sequential" | "ed * Minimal cache interface shared by CLI and MCP model resolution. */ export interface ApiModelCacheProvider { - getCachedModels(): Promise>; + getCachedModels(): Promise>; } -function toApiModelCache(models: Array<{ model_id: string }>): ApiModelCache | undefined { +function toApiModelCache( + models: Array<{ model_id: string; type?: string }>, +): ApiModelCache | undefined { return models.length > 0 ? { models } : undefined; } +function stripCanonicalModelSuffix(modelRef?: string): string | undefined { + if (!modelRef) { + return undefined; + } + + for (const suffix of CANONICAL_MODEL_SUFFIXES) { + if (modelRef.endsWith(suffix)) { + return modelRef.slice(0, -suffix.length); + } + } + + return modelRef; +} + +function normalizeModelRefForType( + modelRef: string | undefined, + apiModelType?: string, +): string | undefined { + return apiModelType === "ai-remover" ? stripCanonicalModelSuffix(modelRef) : modelRef; +} + +function findCachedModelType( + apiCache: ApiModelCache | undefined, + ...refs: Array +): string | undefined { + for (const ref of refs) { + const candidates = [ref, stripCanonicalModelSuffix(ref)].filter( + (candidate, index, values): candidate is string => + Boolean(candidate) && values.indexOf(candidate) === index, + ); + + if (candidates.length === 0) { + continue; + } + + for (const candidate of candidates) { + const match = apiCache?.models.find((model) => model.model_id === candidate); + if (match?.type) { + return match.type; + } + } + } + + return undefined; +} + +function resolveApiModelType( + apiCache: ApiModelCache | undefined, + id: string, + modelName?: string, + configuredType?: string, +): string | undefined { + return ( + configuredType ?? + findCachedModelType(apiCache, modelName, id) ?? + inferApiModelType(modelName ?? id) ?? + inferApiModelType(stripCanonicalModelSuffix(modelName ?? id)) + ); +} + /** * Resolve the requested model using cached API metadata when available, without * forcing a network fetch during normal CLI execution. @@ -142,43 +205,73 @@ function resolveModelId( ): ResolvedModel { const modelConfig = config?.models?.[modelId]; if (modelConfig) { - return resolveFromConfigModel(modelId, modelConfig); + const apiModelType = resolveApiModelType( + apiCache, + modelId, + modelConfig.modelName, + modelConfig.apiModelType, + ); + const normalizedModelRef = normalizeModelRefForType( + modelConfig.modelName ?? modelId, + apiModelType, + ); + + return resolveFromConfigModel( + modelId, + { + ...modelConfig, + modelName: apiModelType === "ai-remover" ? normalizedModelRef : modelConfig.modelName, + }, + inferSubmitMode(normalizedModelRef), + apiModelType, + ); } const registryModel = getRegistryModel(modelId); if (registryModel) { + const apiModelType = resolveApiModelType(apiCache, modelId, registryModel.modelName); + const normalizedModelRef = normalizeModelRefForType(registryModel.modelName, apiModelType); + return resolveFromConfigModel( modelId, { provider: registryModel.provider, apiBaseUrl: registryModel.apiBaseUrl, - modelName: registryModel.modelName, + modelName: normalizedModelRef, apiKeyEnv: "WAVESPEED_API_KEY", }, - "base", + inferSubmitMode(normalizedModelRef), + apiModelType, ); } const cachedApiModel = apiCache?.models.find((model) => model.model_id === modelId); if (cachedApiModel) { + const normalizedModelRef = normalizeModelRefForType(modelId, cachedApiModel.type); + return resolveFromConfigModel( modelId, { provider: "wavespeed", - modelName: modelId, + modelName: normalizedModelRef, }, - "canonical", + inferSubmitMode(normalizedModelRef), + cachedApiModel.type, ); } if (modelId.includes("/")) { + const apiModelType = resolveApiModelType(apiCache, modelId, modelId); + const normalizedModelRef = normalizeModelRefForType(modelId, apiModelType); + return resolveFromConfigModel( modelId, { provider: "wavespeed", - modelName: modelId, + modelName: normalizedModelRef, }, - "canonical", + inferSubmitMode(normalizedModelRef), + apiModelType, ); } @@ -212,10 +305,12 @@ function resolveFromConfigModel( apiBaseUrl?: string; apiKeyEnv?: string; modelName?: string; + apiModelType?: string; type?: "image" | "chat" | "completion"; requestDefaults?: ResolvedModel["requestDefaults"]; }, submitMode: ResolvedModel["submitMode"] = inferSubmitMode(model.modelName ?? id), + apiModelType: string | undefined = model.apiModelType, ): ResolvedModel { const provider = model.provider; @@ -254,6 +349,7 @@ function resolveFromConfigModel( apiKey, apiKeyEnv, modelName: model.modelName, + apiModelType, type, requestDefaults, isFromConfig: true, diff --git a/src/config/types.ts b/src/config/types.ts index 6d51d22..0c640af 100644 --- a/src/config/types.ts +++ b/src/config/types.ts @@ -13,6 +13,7 @@ export interface ModelConfig { apiBaseUrl?: string; apiKeyEnv?: string; modelName?: string; + apiModelType?: string; type?: "image" | "chat" | "completion"; requestDefaults?: RequestDefaults; } @@ -41,6 +42,7 @@ export interface ResolvedModel { apiKey: string; apiKeyEnv: string; modelName?: string; + apiModelType?: string; type: "image" | "chat" | "completion"; requestDefaults: RequestDefaults; isFromConfig: boolean; diff --git a/src/core/operations.ts b/src/core/operations.ts index ab142bb..4437136 100644 --- a/src/core/operations.ts +++ b/src/core/operations.ts @@ -1,4 +1,5 @@ import { submitTask } from "../api/client"; +import { buildEditPayload } from "../utils/model-routing.ts"; import { pollUntilDone } from "../utils/polling"; import type { EditParams, @@ -71,26 +72,12 @@ export async function generateImage(params: GenerateParams): Promise { - const { - prompt, - images, - size = "2048*2048", - base64Output = false, - syncMode = true, - model, - } = params; - - const payload: Record = { - prompt, - images, - size, - enable_base64_output: base64Output, - enable_sync_mode: syncMode, - }; + const { images, model } = params; try { + const payload = buildEditPayload(params); console.error( - `[DEBUG] editImage: Submitting task with model=${model.id}, size=${size}, images=${images.length}, syncMode=${syncMode}`, + `[DEBUG] editImage: Submitting task with model=${model.id}, images=${images.length}, aiRemover=${"image" in payload}, payloadKeys=${Object.keys(payload).join(",")}`, ); const created = await submitTask(model, "edit", payload); console.error(`[DEBUG] editImage: Task submitted, id=${created.id}, status=${created.status}`); diff --git a/src/utils/model-routing.ts b/src/utils/model-routing.ts new file mode 100644 index 0000000..c06037c --- /dev/null +++ b/src/utils/model-routing.ts @@ -0,0 +1,72 @@ +import type { ResolvedModel } from "../config/types.ts"; +import type { EditParams } from "../core/types.ts"; + +const AI_REMOVER_MODEL_HINTS = [ + "remove-background", + "background-remover", + "image-eraser", + "image-text-remover", + "image-watermark-remover", +] as const; + +function getModelRef(model: ResolvedModel): string { + return model.modelName ?? model.id; +} + +export function inferApiModelType(modelRef?: string): string | undefined { + if (!modelRef) { + return undefined; + } + + const lower = modelRef.toLowerCase(); + return AI_REMOVER_MODEL_HINTS.some((hint) => lower.includes(hint)) ? "ai-remover" : undefined; +} + +export function isAiRemoverModel(model: ResolvedModel): boolean { + if (model.apiModelType) { + return model.apiModelType === "ai-remover"; + } + + return inferApiModelType(getModelRef(model)) === "ai-remover"; +} + +function supportsPromptDrivenAiRemoval(model: ResolvedModel): boolean { + return getModelRef(model).toLowerCase().includes("image-eraser"); +} + +export function buildEditPayload(params: EditParams): Record { + const { + prompt, + images, + size = "2048*2048", + base64Output = false, + syncMode = true, + model, + } = params; + + if (!isAiRemoverModel(model)) { + return { + prompt, + images, + size, + enable_base64_output: base64Output, + enable_sync_mode: syncMode, + }; + } + + if (images.length !== 1) { + throw new Error(`Model '${model.id}' accepts exactly 1 input image`); + } + + const payload: Record = { + image: images[0], + enable_base64_output: base64Output, + enable_sync_mode: syncMode, + }; + + if (supportsPromptDrivenAiRemoval(model) && prompt.trim()) { + payload.prompt = prompt; + } + + return payload; +} diff --git a/tests/api/client.test.ts b/tests/api/client.test.ts index 9163e6d..8bb447e 100644 --- a/tests/api/client.test.ts +++ b/tests/api/client.test.ts @@ -41,6 +41,34 @@ const configuredCanonicalAlias: ResolvedModel = { submitMode: "canonical", }; +const aiRemoverAliasModel: ResolvedModel = { + id: "background-remover", + provider: "wavespeed", + apiBaseUrl: "https://api.wavespeed.ai", + apiKeyEnv: "WAVESPEED_API_KEY", + apiKey: "test-api-key", + modelName: "wavespeed-ai/image-background-remover", + apiModelType: "ai-remover", + type: "image", + requestDefaults: {}, + isFromConfig: true, + submitMode: "base", +}; + +const explicitNonRemoverModel: ResolvedModel = { + id: "background-remover-edit", + provider: "wavespeed", + apiBaseUrl: "https://api.wavespeed.ai", + apiKeyEnv: "WAVESPEED_API_KEY", + apiKey: "test-api-key", + modelName: "vendor/background-remover", + apiModelType: "image-to-image", + type: "image", + requestDefaults: {}, + isFromConfig: true, + submitMode: "base", +}; + describe("API Client", () => { const originalEnv = process.env.WAVESPEED_API_KEY; const originalFetch = globalThis.fetch; @@ -274,6 +302,20 @@ describe("API Client", () => { path: "/api/v3/google/nano-banana-2/edit", }); }); + + it("routes ai-remover aliases through the base model path", () => { + expect(buildSubmitTarget(aiRemoverAliasModel, "edit")).toEqual({ + model: "wavespeed-ai/image-background-remover", + path: "/api/v3/wavespeed-ai/image-background-remover", + }); + }); + + it("honors explicit non-remover model types before heuristic matches", () => { + expect(buildSubmitTarget(explicitNonRemoverModel, "edit")).toEqual({ + model: "vendor/background-remover/edit", + path: "/api/v3/vendor/background-remover/edit", + }); + }); }); describe("Static endpoints", () => { diff --git a/tests/commands/cli.test.ts b/tests/commands/cli.test.ts index 00f6218..367b574 100644 --- a/tests/commands/cli.test.ts +++ b/tests/commands/cli.test.ts @@ -85,7 +85,7 @@ describe("CLI Integration Tests", () => { const runCLI = async ( args: string[], ): Promise<{ stdout: string; stderr: string; exitCode: number }> => { - const process = spawn(["bun", "run", cliEntryPath, ...args], { + const spawnedProcess = spawn([process.execPath, "run", cliEntryPath, ...args], { cwd: tempDir, env: { ...Bun.env, WAVESPEED_API_KEY: "test-cli-key" }, stdout: "pipe", @@ -93,16 +93,16 @@ describe("CLI Integration Tests", () => { }); const [stdout, stderr] = await Promise.all([ - new Response(process.stdout).text(), - new Response(process.stderr).text(), + new Response(spawnedProcess.stdout).text(), + new Response(spawnedProcess.stderr).text(), ]); - await process.exited; + await spawnedProcess.exited; return { stdout, stderr, - exitCode: process.exitCode || 0, + exitCode: spawnedProcess.exitCode || 0, }; }; diff --git a/tests/config/models.test.ts b/tests/config/models.test.ts index 5569411..bea88e8 100644 --- a/tests/config/models.test.ts +++ b/tests/config/models.test.ts @@ -25,10 +25,10 @@ function makeConfig(partial: Partial): WavespeedConfig { }; } -function makeCacheProvider(...modelIds: string[]) { +function makeCacheProvider(...models: Array) { return { async getCachedModels() { - return modelIds.map((model_id) => ({ model_id })); + return models.map((model) => (typeof model === "string" ? { model_id: model } : model)); }, }; } @@ -291,6 +291,29 @@ describe("config/models.resolveModel", () => { expect(ce.exitCode).toBe(3); } }); + + it("infers ai-remover routing for configured aliases that point at remover models", () => { + process.env.WAVESPEED_API_KEY = "test-key"; + + const resolved = resolveModel( + "edit", + "background-remover", + makeConfig({ + models: { + "background-remover": { + provider: "wavespeed", + modelName: "wavespeed-ai/image-background-remover", + }, + }, + }), + ); + + expect(resolved.apiModelType).toBe("ai-remover"); + expect(buildSubmitTarget(resolved, "edit")).toEqual({ + model: "wavespeed-ai/image-background-remover", + path: "/api/v3/wavespeed-ai/image-background-remover", + }); + }); }); describe("config/models.resolveModelForRequest", () => { @@ -380,6 +403,62 @@ describe("config/models.resolveModelForRequest", () => { expect(resolved.submitMode).toBe("canonical"); }); + it("preserves cached ai-remover metadata for configured aliases", async () => { + process.env.WAVESPEED_API_KEY = "test-key"; + + const resolved = await resolveModelForRequest( + "edit", + "background-remover", + makeConfig({ + models: { + "background-remover": { + provider: "wavespeed", + modelName: "wavespeed-ai/image-background-remover", + }, + }, + }), + makeCacheProvider({ + model_id: "wavespeed-ai/image-background-remover", + type: "ai-remover", + }), + ); + + expect(resolved.apiModelType).toBe("ai-remover"); + expect(buildSubmitTarget(resolved, "edit")).toEqual({ + model: "wavespeed-ai/image-background-remover", + path: "/api/v3/wavespeed-ai/image-background-remover", + }); + }); + + it("normalizes configured canonical remover refs to the base model path", async () => { + process.env.WAVESPEED_API_KEY = "test-key"; + + const resolved = await resolveModelForRequest( + "edit", + "background-remover", + makeConfig({ + models: { + "background-remover": { + provider: "wavespeed", + modelName: "wavespeed-ai/image-background-remover/edit", + }, + }, + }), + makeCacheProvider({ + model_id: "wavespeed-ai/image-background-remover", + type: "ai-remover", + }), + ); + + expect(resolved.modelName).toBe("wavespeed-ai/image-background-remover"); + expect(resolved.submitMode).toBe("base"); + expect(resolved.apiModelType).toBe("ai-remover"); + expect(buildSubmitTarget(resolved, "edit")).toEqual({ + model: "wavespeed-ai/image-background-remover", + path: "/api/v3/wavespeed-ai/image-background-remover", + }); + }); + it("accepts raw API model ids when cache is cold", async () => { process.env.WAVESPEED_API_KEY = "test-key"; @@ -395,6 +474,44 @@ describe("config/models.resolveModelForRequest", () => { expect(resolved.submitMode).toBe("canonical"); }); + it("infers ai-remover metadata for raw API model ids when cache is cold", async () => { + process.env.WAVESPEED_API_KEY = "test-key"; + + const resolved = await resolveModelForRequest( + "edit", + "wavespeed-ai/image-background-remover", + undefined, + makeCacheProvider(), + ); + + expect(resolved.id).toBe("wavespeed-ai/image-background-remover"); + expect(resolved.modelName).toBe("wavespeed-ai/image-background-remover"); + expect(resolved.submitMode).toBe("base"); + expect(resolved.apiModelType).toBe("ai-remover"); + }); + + it("uses cached base-model typing for canonical remover raw ids", async () => { + process.env.WAVESPEED_API_KEY = "test-key"; + + const resolved = await resolveModelForRequest( + "edit", + "acme/remove-bg/edit", + undefined, + makeCacheProvider({ + model_id: "acme/remove-bg", + type: "ai-remover", + }), + ); + + expect(resolved.modelName).toBe("acme/remove-bg"); + expect(resolved.submitMode).toBe("base"); + expect(resolved.apiModelType).toBe("ai-remover"); + expect(buildSubmitTarget(resolved, "edit")).toEqual({ + model: "acme/remove-bg", + path: "/api/v3/acme/remove-bg", + }); + }); + it("rejects unknown plain model ids when cache is cold", async () => { process.env.WAVESPEED_API_KEY = "test-key"; diff --git a/tests/core/operations.test.ts b/tests/core/operations.test.ts index c71eb2b..ee83f4b 100644 --- a/tests/core/operations.test.ts +++ b/tests/core/operations.test.ts @@ -22,6 +22,26 @@ const aliasedModel: ResolvedModel = { submitMode: "base", }; +const aiRemoverAliasModel: ResolvedModel = { + id: "background-remover", + provider: "wavespeed", + apiBaseUrl: "https://api.test.local", + apiKeyEnv: "WAVESPEED_API_KEY", + apiKey: "test-api-key", + modelName: "wavespeed-ai/image-background-remover", + apiModelType: "ai-remover", + type: "image", + requestDefaults: {}, + isFromConfig: true, + submitMode: "base", +}; + +const aiRemoverEraserModel: ResolvedModel = { + ...aiRemoverAliasModel, + id: "image-eraser", + modelName: "wavespeed-ai/image-eraser", +}; + afterEach(() => { globalThis.fetch = originalFetch; }); @@ -135,4 +155,130 @@ describe("Core operations", () => { expect(postRequest?.body?.model).toBe(testCase.expectedModel); }); } + + it("submits ai-remover edits with a singular image payload on the base route", async () => { + const requests: Array<{ method: string; url: string; body?: Record }> = []; + + globalThis.fetch = async (url: string | URL, options?: RequestInit) => { + const method = options?.method || "GET"; + const urlStr = url.toString(); + const body = + method === "POST" + ? (JSON.parse((options?.body as string | undefined) || "{}") as Record) + : undefined; + + requests.push({ method, url: urlStr, body }); + + if (method === "POST") { + return new Response( + JSON.stringify({ + data: { + id: "task-123", + status: "created", + outputs: [], + }, + }), + { + status: 200, + headers: { "content-type": "application/json" }, + }, + ); + } + + return new Response( + JSON.stringify({ + data: { + id: "task-123", + status: "completed", + outputs: [], + }, + }), + { + status: 200, + headers: { "content-type": "application/json" }, + }, + ); + }; + + const result = await editImage({ + prompt: "remove background", + images: ["https://example.com/source.png"], + size: "1024*1024", + syncMode: false, + model: aiRemoverAliasModel, + }); + + expect(result.success).toBe(true); + + const postRequest = requests.find((request) => request.method === "POST"); + expect(postRequest?.url).toBe( + "https://api.test.local/api/v3/wavespeed-ai/image-background-remover", + ); + expect(postRequest?.body).toEqual({ + image: "https://example.com/source.png", + enable_base64_output: false, + enable_sync_mode: false, + model: "wavespeed-ai/image-background-remover", + }); + }); + + it("preserves prompt forwarding for prompt-driven ai-remover models", async () => { + let capturedBody: Record | undefined; + + globalThis.fetch = async (_url: string | URL, options?: RequestInit) => { + const method = options?.method || "GET"; + if (method === "POST") { + capturedBody = JSON.parse((options?.body as string | undefined) || "{}"); + return new Response( + JSON.stringify({ + data: { + id: "task-123", + status: "created", + outputs: [], + }, + }), + { + status: 200, + headers: { "content-type": "application/json" }, + }, + ); + } + + return new Response( + JSON.stringify({ + data: { + id: "task-123", + status: "completed", + outputs: [], + }, + }), + { + status: 200, + headers: { "content-type": "application/json" }, + }, + ); + }; + + const result = await editImage({ + prompt: "remove the logo", + images: ["https://example.com/source.png"], + syncMode: false, + model: aiRemoverEraserModel, + }); + + expect(result.success).toBe(true); + expect(capturedBody?.prompt).toBe("remove the logo"); + }); + + it("fails fast when an ai-remover edit receives multiple images", async () => { + const result = await editImage({ + prompt: "remove background", + images: ["https://example.com/source.png", "https://example.com/second.png"], + syncMode: false, + model: aiRemoverAliasModel, + }); + + expect(result.success).toBe(false); + expect(result.error).toContain("exactly 1 input image"); + }); });