ai-chatbot/lib/ai/image-edit-subagent.ts

import "server-only";

import type { AgentThreadState } from "@/lib/ai/runtime/agent-thread-state";
import type { AgentExecutionPlan, AgentPlannedToolCall } from "@/lib/ai/auto-agent-graph";
import type { AgentSpecialistRoute } from "@/lib/ai/specialists";
import type { AgentProjectToolDefinition } from "@/lib/ai/tools-registry";
import { getImageWorkflowPlanningHints } from "@/lib/ai/image-workflow-memory";

function normalizeCompactText(text: string) {
  return String(text || "").toLowerCase().replace(/\s+/g, "").trim();
}

export function mentionsImageEditFollowup(text: string): boolean {
  const normalized = normalizeCompactText(text);
  if (!normalized) return false;

  return (
    /继续改|继续调整|再改一下|再调一下|上一版|上一张结果|这张继续改|保留.*不变|只改|继续优化这张图|继续修这张图/.test(
      normalized
    ) ||
    /^(图|图片|这张图|这张|这个图|上传的图|上传图片|刚才那张图|那张图)$/.test(
      normalized
    )
  );
}

function hasImageEditWorkflow(threadState?: AgentThreadState | null) {
  return threadState?.workflowState?.type === "image_edit_followup";
}

export function buildImageEditSpecialistRoute(params: {
  text: string;
  threadState?: AgentThreadState | null;
  hasImage: boolean;
}): AgentSpecialistRoute | null {
  const normalized = normalizeCompactText(params.text);
  const workflowHints = getImageWorkflowPlanningHints(params.threadState);
  const hasWorkflow = hasImageEditWorkflow(params.threadState);
  const mentionsFollowupEdit = mentionsImageEditFollowup(params.text);

  if (!hasWorkflow && !mentionsFollowupEdit && workflowHints.length === 0) {
    return null;
  }

  const planningHints = [
    hasWorkflow
      ? "当前回合处于连续修图 workflow，应优先理解为对上一轮结果图或目标图继续编辑。"
      : "",
    mentionsFollowupEdit
      ? "用户当前表达的是继续修图或局部重调，不要误解成重新生图。"
      : "",
    params.hasImage
      ? "如果本轮带了图片，应优先把它理解为当前修图输入或替换后的参考图。"
      : "",
    ...workflowHints,
  ].filter(Boolean);

  return {
    domain: "image_edit_memory",
    specialistId: "image-edit-subagent",
    reason: hasWorkflow ? "当前回合命中连续修图工作流" : "当前回合命中连续修图语义",
    planningHints,
    toolAllowlist: [
      "respondDirect",
      "resolveTargetImage",
      "editTargetImage",
      "editImage",
      "removeBackground",
      "removeWatermark",
      "enhanceImage",
      "extendImage",
      "resizeImage",
      "angleEditImage",
      "localEditImage",
      "boxCutout",
      "useSelectedCanvasImage",
      "useLatestResultImage",
      "useOriginalUploadImage",
      "publishResultToCanvas",
      "publishToCanvas",
    ],
  };
}

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

function upsertTool(
  toolCalls: AgentPlannedToolCall[],
  name: AgentPlannedToolCall["name"],
  reason: string,
  options?: { prepend?: boolean },
) {
  const existing = toolCalls.find((tool) => tool.name === name);
  if (existing) return toolCalls;
  const nextTool = { name, reason } satisfies AgentPlannedToolCall;
  return options?.prepend ? [nextTool, ...toolCalls] : [...toolCalls, nextTool];
}

export function refinePlanWithImageEditSubagent(params: {
  text: string;
  plan: AgentExecutionPlan;
  availableTools: AgentProjectToolDefinition[];
  threadState?: AgentThreadState | null;
  hasImage: boolean;
}): AgentExecutionPlan {
  const hasWorkflow = hasImageEditWorkflow(params.threadState);
  const hasImageWorkflowHints = getImageWorkflowPlanningHints(params.threadState).length > 0;
  const mentionsFollowupEdit = mentionsImageEditFollowup(params.text);

  if (!hasWorkflow && !hasImageWorkflowHints && !mentionsFollowupEdit) {
    return params.plan;
  }

  let toolCalls = [...params.plan.toolCalls];
  if (hasTool(params.availableTools, "resolveTargetImage")) {
    toolCalls = upsertTool(
      toolCalls,
      "resolveTargetImage",
      "image-edit 子代理优先绑定上一轮结果图作为连续修图目标",
      { prepend: true }
    );
  } else if (hasTool(params.availableTools, "useLatestResultImage")) {
    toolCalls = upsertTool(
      toolCalls,
      "useLatestResultImage",
      "image-edit 子代理优先绑定上一轮结果图作为连续修图目标",
      { prepend: true }
    );
  }
  if (hasTool(params.availableTools, "editTargetImage")) {
    toolCalls = upsertTool(
      toolCalls,
      "editTargetImage",
      "image-edit 子代理要求优先走连续修图执行链路"
    );
  } else if (hasTool(params.availableTools, "editImage")) {
    toolCalls = upsertTool(
      toolCalls,
      "editImage",
      "image-edit 子代理要求优先走连续修图执行链路"
    );
  }
  if (hasTool(params.availableTools, "publishResultToCanvas")) {
    toolCalls = upsertTool(
      toolCalls,
      "publishResultToCanvas",
      "image-edit 子代理建议把修图结果同步回画布"
    );
  } else if (hasTool(params.availableTools, "publishToCanvas")) {
    toolCalls = upsertTool(
      toolCalls,
      "publishToCanvas",
      "image-edit 子代理建议把修图结果同步回画布"
    );
  }

  return {
    ...params.plan,
    nextAction: "edit_image",
    intent: "image_edit",
    responseMode: "tool_then_execute",
    preferredCapability:
      params.plan.preferredCapability === "respond"
        ? "generic_edit"
        : params.plan.preferredCapability,
    shouldUseStructuredReply: true,
    toolCalls,
  };
}