Save current image generation workspace

src/lib/canvas/prompt-mentions.ts

@@ -0,0 +1,92 @@
+import type {
+  CanvasConnection,
+  CanvasImageNodeMetadata,
+  CanvasNode,
+} from "@/lib/canvas/types";
+
+export type PromptMentionSource = {
+  alias: string;
+  token: string;
+  nodeId: string;
+  label: string;
+  description: string;
+  imageUrl: string;
+  index: number;
+};
+
+const mentionUrlPrefix = "canvas-image://";
+
+export function getPromptMentionSources(
+  nodes: CanvasNode[],
+  connections: CanvasConnection[],
+  promptNodeId: string,
+) {
+  const imageById = new Map(
+    nodes
+      .filter((node) => node.type === "image")
+      .map((node) => [node.id, node] as const),
+  );
+
+  return connections
+    .filter((connection) => connection.toNodeId === promptNodeId)
+    .map((connection) => imageById.get(connection.fromNodeId))
+    .filter((node): node is CanvasNode => Boolean(node))
+    .map((node, index) => {
+      const metadata = node.metadata as CanvasImageNodeMetadata;
+      return {
+        alias: `@图${index + 1}`,
+        token: createPromptMentionToken(`图${index + 1}`, node.id),
+        nodeId: node.id,
+        label: `图${index + 1}`,
+        index: index + 1,
+        description:
+          metadata.prompt?.trim() || node.title.trim() || `已连接素材 ${index + 1}`,
+        imageUrl: metadata.imageUrl,
+      } satisfies PromptMentionSource;
+    })
+    .filter((item) => Boolean(item.imageUrl));
+}
+
+export function createPromptMentionToken(label: string, nodeId: string) {
+  return `[${label}](${mentionUrlPrefix}${nodeId})`;
+}
+
+export function parsePromptMentionUrl(value: string) {
+  return value.startsWith(mentionUrlPrefix)
+    ? value.slice(mentionUrlPrefix.length)
+    : null;
+}
+
+export function resolvePromptMentions(
+  prompt: string,
+  mentions: PromptMentionSource[],
+) {
+  const explicitlyReferencedNodeIds: string[] = [];
+  let resolvedPrompt = prompt;
+
+  for (const mention of mentions) {
+    const replacement = `${mention.label}（第 ${mention.index} 张输入参考图）`;
+    const usedMarkdownToken = resolvedPrompt.includes(mention.token);
+    const usedAlias = resolvedPrompt.includes(mention.alias);
+    if (!usedMarkdownToken && !usedAlias) continue;
+    explicitlyReferencedNodeIds.push(mention.nodeId);
+    resolvedPrompt = resolvedPrompt.split(mention.token).join(replacement);
+    resolvedPrompt = resolvedPrompt.split(mention.alias).join(replacement);
+  }
+
+  if (mentions.length > 0) {
+    const referenceGuide = mentions
+      .map(
+        (mention) =>
+          `${mention.label} = 第 ${mention.index} 张输入参考图：${mention.description}`,
+      )
+      .join("\n");
+    resolvedPrompt = `${resolvedPrompt.trim()}\n\n已连接参考图编号：\n${referenceGuide}`;
+  }
+
+  return {
+    prompt: resolvedPrompt.trim(),
+    referencedNodeIds: mentions.map((mention) => mention.nodeId),
+    explicitlyReferencedNodeIds,
+  };
+}