ai-chatbot/lib/ai/conversation-action-resolver.ts

import "server-only";

import type { ImageWorkflowMemory } from "@/lib/ai/image-workflow-memory";
import { mentionsImageEditFollowup } from "@/lib/ai/image-edit-subagent";
import { mentionsImageGenerationFollowup } from "@/lib/ai/image-generation-subagent";
import { isPromptCarryForwardRequest } from "@/lib/ai/prompt-memory";

export interface PendingGridConfirmation {
  gridType: "grid_2x2" | "grid_3x2" | "grid_3x3";
  shotCount: 4 | 6 | 9;
  label: string;
  originalPrompt: string;
}

export type ConversationActionKind =
  | "default"
  | "grid_confirmation_followup"
  | "active_storyboard_followup"
  | "image_edit_followup"
  | "image_generation_followup"
  | "prompt_carry_forward";

export interface ConversationActionResolution {
  kind: ConversationActionKind;
  forceAgentConversation: boolean;
  shouldKeepSelectedModel: boolean;
  effectiveInputText?: string;
  usePromptCarryForward: boolean;
}

export function isConfirmedPromptGenerationRequest(text: string): boolean {
  const normalized = String(text || "").toLowerCase().replace(/\s+/g, "");
  if (!normalized) return false;

  const refersPrompt =
    /根据这个提示词|按这个提示词|就按这个提示词|按刚才的提示词|按上面的提示词|用这个提示词|用刚才的提示词|按该提示词/.test(
      normalized,
    );
  const wantsImageGeneration =
    /生图|生成图片|出图|画出来|生成一张|来一张|开始生成|按此生成|按这个生成/.test(normalized);

  return refersPrompt && wantsImageGeneration;
}

export function detectRequestedStoryboardGrid(text: string): PendingGridConfirmation | null {
  const normalized = String(text || "").toLowerCase().replace(/\s+/g, "");
  if (!normalized) return null;

  if (/4宫格|四宫格|2x2|2×2/.test(normalized)) {
    return {
      gridType: "grid_2x2",
      shotCount: 4,
      label: "4宫格",
      originalPrompt: String(text || "").trim(),
    };
  }

  if (/6宫格|六宫格|3x2|3×2|2x3|2×3/.test(normalized)) {
    return {
      gridType: "grid_3x2",
      shotCount: 6,
      label: "6宫格",
      originalPrompt: String(text || "").trim(),
    };
  }

  if (/9宫格|九宫格|3x3|3×3/.test(normalized)) {
    return {
      gridType: "grid_3x3",
      shotCount: 9,
      label: "9宫格",
      originalPrompt: String(text || "").trim(),
    };
  }

  return null;
}

export function isGridConfirmationReply(text: string): boolean {
  const normalized = String(text || "").toLowerCase().replace(/\s+/g, "");
  if (!normalized) return false;
  return /确认|就按这个|开始生成|可以开始了|没问题|对的|就这样/.test(normalized);
}

export function extractPendingGridConfirmationFromWorkingState(
  workingState?: {
    lastTurnSummary?: unknown;
  } | null,
): PendingGridConfirmation | null {
  const summary = workingState?.lastTurnSummary;
  if (!summary || typeof summary !== "object") return null;
  const summaryRecord = summary as Record<string, unknown>;
  const pending = summaryRecord.pendingGridConfirmation;
  if (!pending || typeof pending !== "object") return null;
  const record = pending as Record<string, unknown>;
  const gridType = String(record.gridType || "").trim();
  const shotCount = Number(record.shotCount || 0);
  const label = String(record.label || "").trim();
  const originalPrompt = String(record.originalPrompt || "").trim();

  if (
    !originalPrompt ||
    !label ||
    !["grid_2x2", "grid_3x2", "grid_3x3"].includes(gridType) ||
    ![4, 6, 9].includes(shotCount)
  ) {
    return null;
  }

  return {
    gridType: gridType as PendingGridConfirmation["gridType"],
    shotCount: shotCount as PendingGridConfirmation["shotCount"],
    label,
    originalPrompt,
  };
}

export function buildGridConfirmationMessage(grid: PendingGridConfirmation): string {
  return [
    `我看到你希望按 ${grid.label} 来生成分镜图。`,
    "为了确保后续自动切片稳定，我会按标准网格、统一格线、固定顺序来出图。",
    `请确认是否按 ${grid.label} 标准版式开始生成。`,
    "",
    "如果确认，请回复：",
    `- 确认按${grid.label}生成`,
    "",
    "如果要修改，也可以直接回复：",
    "- 改成 9 宫格",
    "- 改成 6 宫格",
    "- 改成 4 宫格",
  ].join("\n");
}

export function resolveConversationAction(params: {
  text: string;
  pendingGridConfirmation?: PendingGridConfirmation | null;
  hasActiveStoryboardTask: boolean;
  imageWorkflowMemory?: ImageWorkflowMemory | null;
  confirmedPrompt?: string;
  hasImageInputForPromptCarryForward: boolean;
}) {
  const requestedGrid = detectRequestedStoryboardGrid(params.text);
  const shouldContinuePendingGrid =
    Boolean(params.pendingGridConfirmation) &&
    (isGridConfirmationReply(params.text) || requestedGrid !== null);

  if (shouldContinuePendingGrid) {
    return {
      kind: "grid_confirmation_followup",
      forceAgentConversation: true,
      shouldKeepSelectedModel: true,
      effectiveInputText: params.pendingGridConfirmation?.originalPrompt,
      usePromptCarryForward: false,
    } satisfies ConversationActionResolution;
  }

  if (params.hasActiveStoryboardTask) {
    return {
      kind: "active_storyboard_followup",
      forceAgentConversation: true,
      shouldKeepSelectedModel: true,
      usePromptCarryForward: false,
    } satisfies ConversationActionResolution;
  }

  const hasImageEditFollowup =
    params.imageWorkflowMemory?.lastIntent === "image_edit" &&
    (mentionsImageEditFollowup(params.text) ||
      Boolean(params.imageWorkflowMemory?.lastResultImageUrls?.length));

  if (hasImageEditFollowup) {
    return {
      kind: "image_edit_followup",
      forceAgentConversation: true,
      shouldKeepSelectedModel: true,
      usePromptCarryForward: false,
    } satisfies ConversationActionResolution;
  }

  const promptCarryForwardRequested =
    Boolean(params.confirmedPrompt) &&
    (isConfirmedPromptGenerationRequest(params.text) ||
      isPromptCarryForwardRequest({
        text: params.text,
        hasImage: params.hasImageInputForPromptCarryForward,
      }));

  const hasImageGenerationFollowup =
    params.imageWorkflowMemory?.lastIntent === "image_gen" &&
    (mentionsImageGenerationFollowup(params.text) || promptCarryForwardRequested);

  if (hasImageGenerationFollowup) {
    return {
      kind: "image_generation_followup",
      forceAgentConversation: true,
      shouldKeepSelectedModel: true,
      usePromptCarryForward: promptCarryForwardRequested,
    } satisfies ConversationActionResolution;
  }

  const usePromptCarryForward = promptCarryForwardRequested;

  if (usePromptCarryForward) {
    return {
      kind: "prompt_carry_forward",
      forceAgentConversation: false,
      shouldKeepSelectedModel: false,
      usePromptCarryForward: true,
    } satisfies ConversationActionResolution;
  }

  return {
    kind: "default",
    forceAgentConversation: false,
    shouldKeepSelectedModel: false,
    usePromptCarryForward: false,
  } satisfies ConversationActionResolution;
}