完善文档页文本向量与重排序调用支持

2026-05-31 21:18:41 +08:00
parent 8ee7a7969e
commit 644a6f9d17
24 changed files with 1945 additions and 71 deletions
@@ -1,8 +1,45 @@
 import type { GatewayTask } from '@easyai-ai-gateway/contracts';
-import { createChatTask, createImageEditTask, createImageGenerationTask } from '../api';
+import { createCompatibleChatCompletion, createEmbedding, createImageEditTask, createImageGenerationTask, createRerank } from '../api';
 import type { TaskForm } from '../types';

-export function runTask(token: string, task: TaskForm): Promise<{ task: GatewayTask; next: Record<string, string> }> {
+export interface RunTaskResponse {
+  localOnly?: boolean;
+  next?: Record<string, string>;
+  task: GatewayTask;
+}
+
+export async function runTask(token: string, task: TaskForm): Promise<RunTaskResponse> {
+  if (task.kind === 'chat.completions') {
+    const result = await createCompatibleChatCompletion(token, {
+      model: task.model,
+      runMode: 'simulation',
+      simulation: true,
+      stream: false,
+      messages: [{ role: 'user', content: task.prompt }],
+    });
+    return { localOnly: true, task: compatibleTask(task, result) };
+  }
+  if (task.kind === 'embeddings') {
+    const result = await createEmbedding(token, {
+      model: task.model,
+      input: embeddingInput(task.prompt),
+      dimensions: task.dimensions,
+      runMode: 'simulation',
+      simulation: true,
+    });
+    return { localOnly: true, task: compatibleTask(task, result) };
+  }
+  if (task.kind === 'reranks') {
+    const result = await createRerank(token, {
+      model: task.model,
+      query: task.prompt,
+      documents: rerankDocuments(task.documents),
+      top_n: task.topN,
+      runMode: 'simulation',
+      simulation: true,
+    });
+    return { localOnly: true, task: compatibleTask(task, result) };
+  }
  if (task.kind === 'images.generations') {
    return createImageGenerationTask(token, {
      model: task.model,
@@ -23,10 +60,78 @@ export function runTask(token: string, task: TaskForm): Promise<{ task: GatewayT
      simulation: true,
    });
  }
-  return createChatTask(token, {
+  throw new Error(`Unsupported task kind: ${task.kind}`);
+}
+
+function compatibleTask(task: TaskForm, result: Record<string, unknown>): GatewayTask {
+  const now = new Date().toISOString();
+  return {
+    id: `docs-${task.kind}-${Date.now()}`,
+    asyncMode: false,
+    createdAt: now,
+    finishedAt: now,
+    kind: task.kind,
    model: task.model,
+    modelType: modelTypeForKind(task.kind),
+    request: requestSnapshot(task),
+    result,
+    runMode: 'simulation',
+    status: 'succeeded',
+    updatedAt: now,
+    userId: 'docs-runner',
+  };
+}
+
+function requestSnapshot(task: TaskForm): Record<string, unknown> {
+  if (task.kind === 'embeddings') {
+    return {
+      model: task.model,
+      input: embeddingInput(task.prompt),
+      dimensions: task.dimensions,
+      runMode: 'simulation',
+      simulation: true,
+    };
+  }
+  if (task.kind === 'reranks') {
+    return {
+      model: task.model,
+      query: task.prompt,
+      documents: rerankDocuments(task.documents),
+      top_n: task.topN,
+      runMode: 'simulation',
+      simulation: true,
+    };
+  }
+  return {
+    model: task.model,
+    messages: [{ role: 'user', content: task.prompt }],
    runMode: 'simulation',
    simulation: true,
-    messages: [{ role: 'user', content: task.prompt }],
-  });
+    stream: false,
+  };
+}
+
+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';
+  return 'text_generate';
 }