imagegen-tools/src/lib/canvas/prompt-mentions.ts

import type {
  CanvasConnection,
  CanvasImageNodeMetadata,
  CanvasNode,
  CanvasPromptNodeMetadata,
} from "@/lib/canvas/types";

export type PromptMentionSource = {
  alias: string;
  token: string;
  nodeId: string;
  kind: "image" | "prompt";
  label: string;
  description: string;
  imageUrl?: string;
  prompt?: string;
  index: number;
};

const imageMentionUrlPrefix = "canvas-image://";
const promptMentionUrlPrefix = "canvas-prompt://";

export function getPromptMentionSources(
  nodes: CanvasNode[],
  connections: CanvasConnection[],
  promptNodeId: string,
): PromptMentionSource[] {
  const nodeById = new Map(nodes.map((node) => [node.id, node] as const));
  const sources: PromptMentionSource[] = [];
  let imageIndex = 0;
  let promptIndex = 0;

  for (const connection of connections) {
    if (connection.toNodeId !== promptNodeId) continue;

    const node = nodeById.get(connection.fromNodeId);
    if (!node) continue;

    if (node.type === "prompt") {
      const metadata = node.metadata as CanvasPromptNodeMetadata;
      const prompt = metadata.prompt?.trim() || "";
      promptIndex += 1;
      sources.push({
        alias: `@文本${promptIndex}`,
        token: createPromptMentionToken(`文本${promptIndex}`, node.id, "prompt"),
        nodeId: node.id,
        kind: "prompt",
        label: `文本${promptIndex}`,
        index: promptIndex,
        description: prompt || node.title.trim() || `已连接文本 ${promptIndex}`,
        prompt,
      });
      continue;
    }

    if (node.type === "image") {
      const metadata = node.metadata as CanvasImageNodeMetadata;
      if (!metadata.imageUrl) continue;
      imageIndex += 1;
      sources.push({
        alias: `@图${imageIndex}`,
        token: createPromptMentionToken(`图${imageIndex}`, node.id, "image"),
        nodeId: node.id,
        kind: "image",
        label: `图${imageIndex}`,
        index: imageIndex,
        description:
          metadata.prompt?.trim() || node.title.trim() || `已连接素材 ${imageIndex}`,
        imageUrl: metadata.imageUrl,
      });
    }
  }

  return sources;
}

export function getConnectedPromptSources(
  nodes: CanvasNode[],
  connections: CanvasConnection[],
  targetNodeId: string,
): PromptMentionSource[] {
  const promptById = new Map(
    nodes
      .filter((node) => node.type === "prompt")
      .map((node) => [node.id, node] as const),
  );

  return connections
    .filter((connection) => connection.toNodeId === targetNodeId)
    .map((connection) => promptById.get(connection.fromNodeId))
    .filter((node): node is CanvasNode => Boolean(node))
    .map((node, index) => {
      const metadata = node.metadata as CanvasPromptNodeMetadata;
      const prompt = metadata.prompt?.trim() || "";
      return {
        alias: `@文本${index + 1}`,
        token: createPromptMentionToken(`文本${index + 1}`, node.id, "prompt"),
        nodeId: node.id,
        kind: "prompt",
        label: `文本${index + 1}`,
        index: index + 1,
        description: prompt || node.title.trim() || `已连接文本 ${index + 1}`,
        prompt,
      } satisfies PromptMentionSource;
    });
}

export function createPromptMentionToken(
  label: string,
  nodeId: string,
  kind: PromptMentionSource["kind"] = "image",
) {
  const prefix = kind === "prompt" ? promptMentionUrlPrefix : imageMentionUrlPrefix;
  return `[${label}](${prefix}${nodeId})`;
}

export function parsePromptMentionUrl(value: string) {
  if (value.startsWith(imageMentionUrlPrefix)) {
    return value.slice(imageMentionUrlPrefix.length);
  }
  if (value.startsWith(promptMentionUrlPrefix)) {
    return value.slice(promptMentionUrlPrefix.length);
  }
  return null;
}

export function resolvePromptMentions(
  prompt: string,
  mentions: PromptMentionSource[],
) {
  const explicitlyReferencedNodeIds: string[] = [];
  const usedMentions: PromptMentionSource[] = [];
  let resolvedPrompt = prompt;

  for (const mention of mentions) {
    const replacement =
      mention.kind === "image"
        ? `${mention.label}（第 ${mention.index} 张输入参考图）`
        : `${mention.label}（第 ${mention.index} 段输入文本）`;
    const usedMarkdownToken = resolvedPrompt.includes(mention.token);
    const usedAlias = resolvedPrompt.includes(mention.alias);
    if (!usedMarkdownToken && !usedAlias) continue;
    explicitlyReferencedNodeIds.push(mention.nodeId);
    usedMentions.push(mention);
    resolvedPrompt = resolvedPrompt.split(mention.token).join(replacement);
    resolvedPrompt = resolvedPrompt.split(mention.alias).join(replacement);
  }

  if (mentions.length > 0) {
    const imageGuide = mentions
      .filter((mention) => mention.kind === "image")
      .map(
        (mention) =>
          `${mention.label} = 第 ${mention.index} 张输入参考图：${mention.description}`,
      );
    const promptGuide = mentions
      .filter((mention) => mention.kind === "prompt")
      .map(
        (mention) =>
          `${mention.label} = 第 ${mention.index} 段输入文本：${mention.description}`,
      );
    const guide = [
      imageGuide.length ? `已连接参考图编号：\n${imageGuide.join("\n")}` : "",
      promptGuide.length ? `已连接文本编号：\n${promptGuide.join("\n")}` : "",
    ]
      .filter(Boolean)
      .join("\n\n");
    resolvedPrompt = `${resolvedPrompt.trim()}\n\n${guide}`;
  }

  return {
    prompt: resolvedPrompt.trim(),
    referencedNodeIds: mentions.map((mention) => mention.nodeId),
    explicitlyReferencedNodeIds,
  };
}