ai-chatbot/tests/helpers/agent-image-provider-harness.ts

27 lines
2.7 KiB
TypeScript

import { serverLoader } from "./canvas-download-server-loader";
export function agentImageProviderHarness(options: { generate?: (id: string, count: number, call: number) => string; models?: any[] } = {}) {
const calls: any[] = [], paused = new Set<string>();
const models = options.models || ["primary", "backup"].map((id, i) => ({ model: { id, modelId: id, displayName: id, isEnabled: true, supportsImageGen: true, supportsVision: true, isDefault: !i, imageGenConfig: { maxImages: 12 }, capabilities: { kind: "image" } }, provider: { id: "same-provider", type: "openai-image", apiKey: "synthetic", isEnabled: true } }));
const mocks: Record<string, any> = {
"@/lib/db/queries": { getEnabledAiModelsWithProvider: async () => models },
"@/lib/ai/rmbg/service": {},
"@/lib/net/url-access": { filterReachableRemoteUrls: async (urls: string[]) => ({ kept: urls, dropped: [] }) },
"@/lib/canvas/image-concurrency": { acquireCanvasImageModelSlot: async () => ({ waitMs: 0, maxConcurrent: 999, release: async () => {} }), normalizeCanvasImageModelKey: (v: string) => v, recordCanvasImageModelExecutionDuration() {} },
"@/lib/ai/runtime/model-health": { recordAiModelHealthSuccess: async () => {}, recordAiModelHealthFailure: async () => {}, classifyAiModelHealthFailure: () => ({ reason: "provider_error" }), shouldFailoverAiModelError: () => false },
"@/lib/ai/runtime/agent-model-circuit-store": { agentImageCircuitStore: { isPaused: async (e: any) => paused.has(e.model.id), fail: async (e: any) => { paused.add(e.model.id); } } },
"@/lib/ai/image-model-pricing": { getModelGenerationCreditCost: ({ imageGenConfig }: any) => imageGenConfig?.creditCost || 0 },
"@/lib/storage/service": {},
};
for (const name of ["gemini-image", "alibaba-image", "bfl-image", "ruxa-image", "fal-image", "snapgen-image", "openai-image", "bananaimage", "visionary", "atlascloud", "vip98k", "stepfun-image", "midjourney", "runway", "krea", "volcano-vision"]) mocks[`@/lib/ai/providers/${name}`] = {};
mocks["@/lib/ai/providers/openai-image"] = { createOpenAIImage: () => ({ chat: (id: string) => ({ doGenerate: async (params: any) => {
calls.push({ id, ...params });
const text = options.generate?.(id, params.imageParams.imageCount, calls.length) || `![image](https://fixture.invalid/${id}-${calls.length}.png)`;
return { text };
} }) }) };
const generate = serverLoader(mocks)("lib/canvas/generate-image.ts").generateCanvasImage;
const request = { prompt: "a cat", model: "primary", imageCount: 1, imageSize: "2K", aspectRatio: "3:4", agentImageMode: "auto", allowFailover: true, referenceImages: ["https://fixture.invalid/source.png"], negativePrompt: "no text", providerImageParams: { seed: 0, watermark: false } };
return { generate, request, calls, paused };
}