import type { GatewayTask } from '@easyai-ai-gateway/contracts'; import { createCompatibleChatCompletion, createEmbedding, createImageEditTask, createImageGenerationTask, createRerank, createResponse, createVideoGenerationTask, getAPITask, } from '../api'; import type { TaskForm, TaskSubmissionMode } from '../types'; export interface RunTaskResponse { localOnly?: boolean; next?: Record; submissionMode: TaskSubmissionMode; task: GatewayTask; } export interface RunTaskOptions { requestBody?: Record; submissionMode?: TaskSubmissionMode; } const simulationParameterKeys = ['runMode', 'run_mode', 'simulation', 'testMode', 'test_mode'] as const; export function applyTaskSubmissionMode( input: Record, submissionMode: TaskSubmissionMode, ): Record { const body = { ...input }; for (const key of simulationParameterKeys) delete body[key]; if (submissionMode === 'simulation') { body.runMode = 'simulation'; body.simulation = true; } return body; } export async function runTask(token: string, task: TaskForm, options: RunTaskOptions = {}): Promise { const submissionMode = options.submissionMode ?? 'simulation'; const requestBody = task.kind === 'tasks.retrieve' ? { taskId: task.taskId } : applyTaskSubmissionMode(options.requestBody ?? defaultRequestBody(task), submissionMode); if (task.kind === 'chat.completions') { const result = await createCompatibleChatCompletion( token, requestBody as Parameters[1], ); return { localOnly: true, submissionMode, task: compatibleTask(task, result, requestBody, submissionMode) }; } if (task.kind === 'responses') { const result = await createResponse(token, requestBody as Parameters[1]); return { localOnly: true, submissionMode, task: compatibleTask(task, result, requestBody, submissionMode) }; } if (task.kind === 'embeddings') { const result = await createEmbedding(token, requestBody as Parameters[1]); return { localOnly: true, submissionMode, task: compatibleTask(task, result, requestBody, submissionMode) }; } if (task.kind === 'reranks') { const result = await createRerank(token, requestBody as Parameters[1]); return { localOnly: true, submissionMode, task: compatibleTask(task, result, requestBody, submissionMode) }; } if (task.kind === 'images.generations') { const response = await createImageGenerationTask( token, requestBody as Parameters[1], ); return { ...response, submissionMode }; } if (task.kind === 'images.edits') { const response = await createImageEditTask(token, requestBody as Parameters[1]); return { ...response, submissionMode }; } if (task.kind === 'videos.generations') { const response = await createVideoGenerationTask( token, requestBody as unknown as Parameters[1], ); return { ...response, submissionMode }; } if (task.kind === 'tasks.retrieve') { const taskId = task.taskId?.trim(); if (!taskId) throw new Error('请输入要取回的 Task ID'); const result = await getAPITask(token, taskId); return { localOnly: true, submissionMode: 'production', task: compatibleTask(task, result as unknown as Record, requestBody, 'production'), }; } throw new Error(`Unsupported task kind: ${task.kind}`); } function compatibleTask( task: TaskForm, result: Record, requestBody: Record, submissionMode: TaskSubmissionMode, ): GatewayTask { const now = new Date().toISOString(); return { id: `docs-${task.kind}-${Date.now()}`, asyncMode: false, createdAt: now, finishedAt: now, kind: task.kind, model: typeof requestBody.model === 'string' ? requestBody.model : task.model, modelType: modelTypeForKind(task.kind), request: requestBody, result, runMode: submissionMode, status: 'succeeded', updatedAt: now, userId: 'docs-runner', }; } function defaultRequestBody(task: TaskForm): Record { if (task.kind === 'chat.completions') { return { model: task.model, messages: [{ role: 'user', content: task.prompt }], stream: false, }; } if (task.kind === 'responses') { return { model: task.model, input: task.prompt, instructions: task.instructions, previous_response_id: task.previousResponseId, store: true, stream: false, }; } if (task.kind === 'embeddings') { return { model: task.model, input: embeddingInput(task.prompt), dimensions: task.dimensions, }; } if (task.kind === 'reranks') { return { model: task.model, query: task.prompt, documents: rerankDocuments(task.documents), top_n: task.topN, }; } if (task.kind === 'images.generations') { return { model: task.model, prompt: task.prompt, quality: 'medium', size: '1024x1024', }; } if (task.kind === 'images.edits') { return { model: task.model, prompt: task.prompt, image: task.image, mask: task.mask, }; } if (task.kind === 'videos.generations') { return { model: task.model, content: [{ type: 'text', text: task.prompt }], aspect_ratio: task.aspectRatio ?? '16:9', resolution: task.resolution ?? '720p', duration: task.duration ?? 5, audio: task.outputAudio ?? true, }; } if (task.kind === 'tasks.retrieve') return { taskId: task.taskId }; return { model: task.model }; } function embeddingInput(prompt: string) { const lines = splitLines(prompt); return lines.length > 1 ? lines : (lines[0] ?? prompt); } function rerankDocuments(value?: string) { const documents = splitLines(value ?? ''); return documents.length ? documents : ['AI Gateway 提供 OpenAI 兼容接口。', '图片生成任务支持异步队列。']; } function splitLines(value: string) { return value .split(/\n+/) .map((item) => item.trim()) .filter(Boolean); } function modelTypeForKind(kind: TaskForm['kind']) { if (kind === 'embeddings') return 'text_embedding'; if (kind === 'reranks') return 'text_rerank'; if (kind === 'images.generations') return 'image_generate'; if (kind === 'images.edits') return 'image_edit'; if (kind === 'videos.generations') return 'video_generate'; if (kind === 'tasks.retrieve') return 'task'; return 'text_generate'; }