ai-chatbot/lib/ai/image-generation-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";
import {
  getPromptMemoryPlanningHints,
  isPromptCarryForwardRequest,
} from "@/lib/ai/prompt-memory";

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

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

  return /继续生成|继续出图|继续做图|再来(一|两|三|四|五|六|七|八|九|\d+)?张|再出(一|两|三|四|五|六|七|八|九|\d+)?张|再做(一|两|三|四|五|六|七|八|九|\d+)?组|再给我(一|两|三|四|五|六|七|八|九|\d+)?张|补(一|两|三|四|五|六|七|八|九|\d+)?张|多来(一|两|三|四|五|六|七|八|九|\d+)?张|换一版|再换一版|同风格再来|同样风格再来|保持.*(不变|一致).*再来|同提示词再来|沿用上一个效果再来/.test(
    normalized,
  );
}

export function buildImageGenerationSpecialistRoute(params: {
  text: string;
  threadState?: AgentThreadState | null;
  hasImage: boolean;
}): AgentSpecialistRoute | null {
  const normalized = normalizeCompactText(params.text);
  const promptHints = getPromptMemoryPlanningHints(params.threadState);
  const workflowHints = getImageWorkflowPlanningHints(params.threadState);
  const hasPromptMemory = promptHints.length > 0;
  const hasImageWorkflowMemory = workflowHints.length > 0;
  const carryForward = isPromptCarryForwardRequest({
    text: params.text,
    hasImage: params.hasImage,
  });
  const mentionsGenerationFollowup = mentionsImageGenerationFollowup(params.text);
  const mentionsPromptExecution =
    /提示词|prompt|关键词|咒语/.test(normalized) &&
    /生成|生图|出图|继续做|继续生成|执行/.test(normalized);

  if (
    !hasPromptMemory &&
    !hasImageWorkflowMemory &&
    !carryForward &&
    !mentionsPromptExecution &&
    !mentionsGenerationFollowup
  ) {
    return null;
  }

  const planningHints = [
    carryForward
      ? "用户当前要求沿用上一轮优化提示词继续执行，应优先理解为连续生图执行。"
      : "",
    mentionsGenerationFollowup
      ? "如果用户说“再来几张 / 换一版 / 同风格再来”，优先理解为沿用上一轮生图上下文继续执行，而不是开启全新任务。"
      : "",
    hasPromptMemory
      ? "当前会话里存在上一轮可复用提示词，如果用户在说“上一个提示词/刚才那个 prompt”，优先复用这段提示词。"
      : "",
    hasImageWorkflowMemory
      ? "当前会话里存在图片工作流记忆，优先参考上一轮参考图、结果图、数量和画幅偏好。"
      : "",
    params.hasImage
      ? "如果当前回合上传了新图片，并要求沿用上一轮提示词，应理解为“上一轮提示词 + 当前图片条件”的联合执行。"
      : "",
    ...promptHints,
    ...workflowHints,
  ].filter(Boolean);

  return {
    domain: "image_generation_memory",
    specialistId: "image-generation-subagent",
    reason: carryForward
      ? "当前回合命中提示词接力生图"
      : mentionsGenerationFollowup
        ? "当前回合命中连续生图跟进语义"
        : "当前回合命中提示词相关的连续生图任务",
    planningHints,
    toolAllowlist: [
      "respondDirect",
      "resolveTargetImage",
      "generateNewImage",
      "generateImage",
      "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 refinePlanWithImageGenerationSubagent(params: {
  text: string;
  plan: AgentExecutionPlan;
  availableTools: AgentProjectToolDefinition[];
  threadState?: AgentThreadState | null;
  hasImage: boolean;
}): AgentExecutionPlan {
  const carryForward = isPromptCarryForwardRequest({
    text: params.text,
    hasImage: params.hasImage,
  });
  const mentionsGenerationFollowup = mentionsImageGenerationFollowup(params.text);
  const hasPromptMemory = getPromptMemoryPlanningHints(params.threadState).length > 0;
  const hasImageWorkflowMemory = getImageWorkflowPlanningHints(params.threadState).length > 0;
  if (
    !carryForward &&
    !mentionsGenerationFollowup &&
    !hasPromptMemory &&
    !hasImageWorkflowMemory
  ) {
    return params.plan;
  }

  let toolCalls = [...params.plan.toolCalls];
  const targetResolverTools = toolCalls.filter((tool) =>
    ["resolveTargetImage", "useSelectedCanvasImage", "useLatestResultImage", "useOriginalUploadImage"].includes(
      tool.name,
    ),
  );
  if (hasTool(params.availableTools, "generateNewImage")) {
    toolCalls = [
      ...targetResolverTools,
      { name: "generateNewImage", reason: "image-generation 子代理要求优先走连续生图执行链路" },
    ];
  } else if (hasTool(params.availableTools, "generateImage")) {
    toolCalls = [
      ...targetResolverTools,
      { name: "generateImage", reason: "image-generation 子代理要求优先走连续生图执行链路" },
    ];
  }
  if (hasTool(params.availableTools, "publishResultToCanvas")) {
    toolCalls = upsertTool(
      toolCalls,
      "publishResultToCanvas",
      "image-generation 子代理建议把连续生图结果同步到画布"
    );
  } else if (hasTool(params.availableTools, "publishToCanvas")) {
    toolCalls = upsertTool(
      toolCalls,
      "publishToCanvas",
      "image-generation 子代理建议把连续生图结果同步到画布"
    );
  }

  return {
    ...params.plan,
    nextAction: "generate_image",
    intent: "image_gen",
    responseMode: "tool_then_execute",
    preferredCapability: "generate_image",
    shouldUseStructuredReply: true,
    toolCalls,
  };
}