ai-chatbot/lib/ai/specialists/index.ts

import type { AnalysisResult } from "@/lib/ai/smart-analyzer";
import type { AgentThreadState } from "@/lib/ai/runtime/agent-thread-state";
import { refinePlanWithImageEditSubagent } from "@/lib/ai/image-edit-subagent";
import { refinePlanWithImageGenerationSubagent } from "@/lib/ai/image-generation-subagent";
import { refinePlanWithStoryboardSubagent } from "@/lib/video-agent/storyboard-subagent";
import type {
  AgentExecutionPlan,
  AgentPlanNextAction,
  AgentPlanResponseMode,
  AgentPlannedToolCall,
  AgentPreferredCapability,
  AgentTargetResolutionMode,
} from "@/lib/ai/auto-agent-graph";
import type {
  AgentProjectToolDefinition,
  AgentProjectToolName,
} from "@/lib/ai/tools-registry";

export type AgentSpecialistDomain =
  | "conversation"
  | "vision"
  | "image_edit"
  | "image_edit_memory"
  | "image_generate"
  | "image_generation_memory"
  | "video_generate"
  | "storyboard"
  | "canvas";

export type AgentSpecialistId =
  | "conversation-coordinator"
  | "vision-analyst"
  | "image-edit-specialist"
  | "image-edit-subagent"
  | "image-generate-specialist"
  | "image-generation-subagent"
  | "video-specialist"
  | "storyboard-subagent"
  | "canvas-specialist";

export interface AgentSpecialistRoute {
  domain: AgentSpecialistDomain;
  specialistId: AgentSpecialistId;
  reason: string;
  planningHints: string[];
  toolAllowlist?: AgentProjectToolName[];
}

export interface AgentSpecialistContext {
  text: string;
  analysis: AnalysisResult;
  hasImage: boolean;
  hasCanvasSelection: boolean;
  targetResolutionMode: AgentTargetResolutionMode;
  mediaSource?: "current_message" | "canvas_selection" | "history" | "none";
  selectedImageCount?: number;
  threadState?: AgentThreadState | null;
}

export interface AgentPlanRefineInput extends AgentSpecialistContext {
  plan: AgentExecutionPlan;
  availableTools: AgentProjectToolDefinition[];
}

function hasTool(
  availableTools: AgentProjectToolDefinition[],
  name: AgentProjectToolName,
) {
  return availableTools.some((tool) => tool.name === name);
}

function upsertTool(
  toolCalls: AgentPlannedToolCall[],
  name: AgentProjectToolName,
  reason: string,
  options?: {
    prepend?: boolean;
    args?: Record<string, unknown>;
  },
) {
  const existing = toolCalls.find((tool) => tool.name === name);
  if (existing) {
    if (!existing.reason) {
      existing.reason = reason;
    }
    if (!existing.args && options?.args) {
      existing.args = options.args;
    }
    return toolCalls;
  }

  const nextTool = {
    name,
    reason,
    args: options?.args,
  } satisfies AgentPlannedToolCall;

  return options?.prepend ? [nextTool, ...toolCalls] : [...toolCalls, nextTool];
}

function forcePlanShape(params: {
  plan: AgentExecutionPlan;
  nextAction: AgentPlanNextAction;
  intent: AnalysisResult["intent"];
  responseMode?: AgentPlanResponseMode;
  preferredCapability?: AgentPreferredCapability;
  subIntent?: string;
  toolCalls: AgentPlannedToolCall[];
  shouldUseStructuredReply?: boolean;
  targetSource?: AgentTargetResolutionMode;
}): AgentExecutionPlan {
  return {
    ...params.plan,
    nextAction: params.nextAction,
    intent: params.intent,
    responseMode:
      params.responseMode ?? params.plan.responseMode,
    preferredCapability:
      params.preferredCapability ?? params.plan.preferredCapability,
    subIntent:
      params.intent === "image_edit"
        ? params.subIntent ?? params.plan.subIntent
        : undefined,
    toolCalls: params.toolCalls,
    shouldUseStructuredReply:
      params.shouldUseStructuredReply ?? params.plan.shouldUseStructuredReply,
    targetSource: params.targetSource ?? params.plan.targetSource,
  };
}

function prependTargetResolverTool(params: {
  toolCalls: AgentPlannedToolCall[];
  availableTools: AgentProjectToolDefinition[];
  targetSource: AgentTargetResolutionMode;
  hasCanvasSelection?: boolean;
}) {
  if (params.targetSource === "none" || params.targetSource === "clarification") {
    return params.toolCalls;
  }

  let source: "selected_canvas" | "latest_result" | "original_upload" | "auto" = "auto";
  if (params.targetSource === "explicit_selection" || params.hasCanvasSelection) {
    source = "selected_canvas";
  } else if (params.targetSource === "current_upload") {
    source = "original_upload";
  } else if (
    params.targetSource === "latest_result" ||
    params.targetSource === "working_memory"
  ) {
    source = "latest_result";
  }

  if (hasTool(params.availableTools, "resolveTargetImage")) {
    return upsertTool(
      params.toolCalls,
      "resolveTargetImage",
      "先锁定当前要处理的目标图。",
      {
        prepend: true,
        args: { source },
      },
    );
  }

  if (source === "selected_canvas" && hasTool(params.availableTools, "useSelectedCanvasImage")) {
    return upsertTool(
      params.toolCalls,
      "useSelectedCanvasImage",
      "先使用当前选中的画布图片作为目标图。",
      { prepend: true },
    );
  }
  if (source === "latest_result" && hasTool(params.availableTools, "useLatestResultImage")) {
    return upsertTool(
      params.toolCalls,
      "useLatestResultImage",
      "先使用最近结果图作为目标图。",
      { prepend: true },
    );
  }
  if (source === "original_upload" && hasTool(params.availableTools, "useOriginalUploadImage")) {
    return upsertTool(
      params.toolCalls,
      "useOriginalUploadImage",
      "先使用原上传图作为目标图。",
      { prepend: true },
    );
  }

  return params.toolCalls;
}

function appendCanvasPublishTool(
  toolCalls: AgentPlannedToolCall[],
  availableTools: AgentProjectToolDefinition[],
) {
  if (hasTool(availableTools, "publishResultToCanvas")) {
    return upsertTool(
      toolCalls,
      "publishResultToCanvas",
      "执行结果需要同步回无限画布。",
    );
  }
  if (hasTool(availableTools, "publishToCanvas")) {
    return upsertTool(
      toolCalls,
      "publishToCanvas",
      "执行结果需要同步回无限画布。",
    );
  }
  return toolCalls;
}

function normalizeEditCapability(
  analysis: AnalysisResult,
  text: string,
): {
  capability: AgentPreferredCapability;
  toolName: AgentProjectToolName;
  subIntent?: string;
} {
  const normalizedText = String(text || "").toLowerCase();
  const subIntent = String(analysis.subIntent || "").trim().toLowerCase();

  if (subIntent === "remove_background" || /抠图|去背|去背景|透明底/.test(normalizedText)) {
    return {
      capability: "remove_background",
      toolName: "removeBackground",
      subIntent: "remove_background",
    };
  }
  if (subIntent === "extend" || /扩图|拓图|补边|延展|加宽|加高/.test(normalizedText)) {
    return {
      capability: "extend",
      toolName: "extendImage",
      subIntent: "extend",
    };
  }
  if (subIntent === "resize" || /尺寸|分辨率|改成.*x|改为.*x|比例|画布/.test(normalizedText)) {
    return {
      capability: "resize",
      toolName: "resizeImage",
      subIntent: "resize",
    };
  }
  if (subIntent === "box_cutout" || /框选抠图|框选/.test(normalizedText)) {
    return {
      capability: "box_cutout",
      toolName: "boxCutout",
      subIntent: "box_cutout",
    };
  }
  if (subIntent === "enhance" || /增强|放大|高清|超清|锐化|修复/.test(normalizedText)) {
    return {
      capability: "enhance",
      toolName: "enhanceImage",
      subIntent: "enhance",
    };
  }
  return {
    capability: "generic_edit",
    toolName: "editImage",
    subIntent: analysis.subIntent || undefined,
  };
}

export function refinePlanWithSpecialist(
  specialist: AgentSpecialistRoute,
  input: AgentPlanRefineInput,
): AgentExecutionPlan {
  const { plan, availableTools, analysis, text } = input;
  let toolCalls = [...plan.toolCalls];

  if (specialist.domain === "conversation") {
    const shouldClarify = plan.nextAction === "clarify";
    if (hasTool(availableTools, "respondDirect")) {
      toolCalls = [{ name: "respondDirect", reason: "纯对话由对话协调专家直接回复" }];
    }
    return forcePlanShape({
      plan,
      nextAction: shouldClarify ? "clarify" : "respond",
      intent: "chat",
      responseMode: shouldClarify ? "clarify" : "direct_answer",
      preferredCapability: "respond",
      toolCalls,
      shouldUseStructuredReply: false,
      targetSource: "none",
    });
  }

  if (specialist.domain === "vision") {
    if (hasTool(availableTools, "analyzeTargetImage")) {
      toolCalls = [{ name: "analyzeTargetImage", reason: "视觉分析任务由看图专家处理" }];
    } else if (hasTool(availableTools, "analyzeImage")) {
      toolCalls = [{ name: "analyzeImage", reason: "视觉分析任务由看图专家处理" }];
    } else if (hasTool(availableTools, "respondDirect")) {
      toolCalls = [{ name: "respondDirect", reason: "缺少看图工具时退回直接回答" }];
    }
    toolCalls = prependTargetResolverTool({
      toolCalls,
      availableTools,
      targetSource: plan.targetSource,
      hasCanvasSelection: input.hasCanvasSelection,
    });
    return forcePlanShape({
      plan,
      nextAction: "analyze_image",
      intent: "vision",
      responseMode: "tool_then_answer",
      preferredCapability: "analyze",
      toolCalls,
      shouldUseStructuredReply: false,
    });
  }

  if (specialist.domain === "image_edit") {
    const editTarget = normalizeEditCapability(analysis, text);
    if (hasTool(availableTools, editTarget.toolName)) {
      toolCalls = [{ name: editTarget.toolName, reason: "修图专家选择最匹配的编辑能力" }];
    } else if (hasTool(availableTools, "editTargetImage")) {
      toolCalls = [{ name: "editTargetImage", reason: "修图专家回退到目标图编辑能力" }];
    } else if (hasTool(availableTools, "editImage")) {
      toolCalls = [{ name: "editImage", reason: "修图专家回退到通用编辑能力" }];
    }
    toolCalls = prependTargetResolverTool({
      toolCalls,
      availableTools,
      targetSource: plan.targetSource,
      hasCanvasSelection: input.hasCanvasSelection,
    });
    return forcePlanShape({
      plan,
      nextAction: "edit_image",
      intent: "image_edit",
      subIntent: editTarget.subIntent,
      responseMode: "tool_then_execute",
      preferredCapability: editTarget.capability,
      toolCalls,
      shouldUseStructuredReply: true,
    });
  }

  if (specialist.domain === "image_edit_memory") {
    return refinePlanWithImageEditSubagent({
      text,
      plan,
      availableTools,
      threadState: input.threadState,
      hasImage: input.hasImage,
    });
  }

  if (specialist.domain === "image_generate") {
    if (hasTool(availableTools, "generateNewImage")) {
      toolCalls = [{ name: "generateNewImage", reason: "生图专家执行图片生成" }];
    } else if (hasTool(availableTools, "generateImage")) {
      toolCalls = [{ name: "generateImage", reason: "生图专家执行图片生成" }];
    }
    toolCalls = prependTargetResolverTool({
      toolCalls,
      availableTools,
      targetSource: plan.targetSource,
      hasCanvasSelection: input.hasCanvasSelection,
    });
    return forcePlanShape({
      plan,
      nextAction: "generate_image",
      intent: "image_gen",
      responseMode: "tool_then_execute",
      preferredCapability: "generate_image",
      toolCalls,
      shouldUseStructuredReply: true,
    });
  }

  if (specialist.domain === "image_generation_memory") {
    return refinePlanWithImageGenerationSubagent({
      text,
      plan,
      availableTools,
      threadState: input.threadState,
      hasImage: input.hasImage,
    });
  }

  if (specialist.domain === "video_generate") {
    if (hasTool(availableTools, "generateVideo")) {
      toolCalls = [{ name: "generateVideo", reason: "视频专家执行视频生成" }];
    }
    return forcePlanShape({
      plan,
      nextAction: "generate_video",
      intent: "video_gen",
      responseMode: "tool_then_execute",
      preferredCapability: "generate_video",
      toolCalls,
      shouldUseStructuredReply: true,
    });
  }

  if (specialist.domain === "storyboard") {
    return refinePlanWithStoryboardSubagent({
      text,
      plan,
      availableTools,
      threadState: input.threadState,
    });
  }

  if (specialist.domain === "canvas") {
    const isExecutionAction =
      plan.nextAction === "edit_image" ||
      plan.nextAction === "generate_image";
    const normalizedTargetSource =
      input.hasCanvasSelection && plan.targetSource === "none"
        ? "explicit_selection"
        : plan.targetSource;

    if (
      normalizedTargetSource !== "none" &&
      normalizedTargetSource !== "clarification"
    ) {
      toolCalls = prependTargetResolverTool({
        toolCalls,
        availableTools,
        targetSource: normalizedTargetSource,
        hasCanvasSelection: input.hasCanvasSelection,
      });
    }

    if (
      isExecutionAction &&
      (hasTool(availableTools, "publishResultToCanvas") ||
        hasTool(availableTools, "publishToCanvas"))
    ) {
      toolCalls = appendCanvasPublishTool(toolCalls, availableTools);
    }

    return {
      ...plan,
      toolCalls,
      targetSource: normalizedTargetSource,
    };
  }

  return plan;
}