ai-chatbot/lib/ai/runtime/recent-dialogue-memory.ts

import type { DBMessage } from "@/lib/db/schema";

export interface RecentDialogueMemoryEntry {
  role: "user" | "assistant";
  content: string;
  messageId?: string;
}

export interface RecentDialogueMemory {
  entries: RecentDialogueMemoryEntry[];
  contextSummary: string;
  memoryHints: string[];
  isRecallRequest: boolean;
}

const USER_RECALL_PATTERN =
  /(?:我|用户)?(?:刚才|刚刚|前面|之前|上面|上一(?:句|条|轮)|最近).{0,12}(?:说|发|问|讲|提|输入|发送).{0,12}(?:什么|啥|哪句|哪些|内容)|(?:我说了什么|我刚才说啥|我前面说啥|我之前说啥|我发了什么|我问了什么|记得我.{0,12}说|没有记忆)/i;

const ASSISTANT_RECALL_PATTERN =
  /(?:你|助手|ai|AI).{0,12}(?:刚才|刚刚|前面|之前|上面|上一(?:句|条|轮)|最近).{0,12}(?:说|回复|回答|讲).{0,12}(?:什么|啥|哪句|哪些|内容)/i;

function normalizeWhitespace(value: string) {
  return value.replace(/\s+/g, " ").trim();
}

export function isRecentDialogueRecallRequest(text: string) {
  const normalized = normalizeWhitespace(String(text || ""));
  if (!normalized) return false;
  return USER_RECALL_PATTERN.test(normalized) || ASSISTANT_RECALL_PATTERN.test(normalized);
}

function wantsAssistantRecall(text: string) {
  const normalized = normalizeWhitespace(String(text || ""));
  return ASSISTANT_RECALL_PATTERN.test(normalized);
}

function sanitizeDialogueMemoryText(text: string) {
  return String(text || "")
    .replace(/<!--[\s\S]*?-->/g, "")
    .replace(/!\[[^\]]*]\([^)]+\)/g, "")
    .replace(/\[[^\]]*]\((?:https?:\/\/|\/)[^)]+\)/g, "")
    .replace(/<video[^>]*>[\s\S]*?<\/video>/gi, "")
    .replace(/<video[^>]+src="[^"]+"[^>]*>/gi, "")
    .replace(/◈\s*[^\n]+\n*/g, "")
    .replace(/IMAGE_GENERATION_(?:START|END)/g, "")
    .replace(/\n{3,}/g, "\n\n")
    .trim();
}

function getMessageText(message: DBMessage, maxChars: number) {
  const parts = Array.isArray(message.parts)
    ? (message.parts as Array<Record<string, unknown>>)
    : [];
  const text = parts
    .filter((part) => part?.type === "text" && typeof part?.text === "string")
    .map((part) => sanitizeDialogueMemoryText(String(part.text || "")))
    .filter(Boolean)
    .join("\n")
    .trim();

  if (!text) return "";
  return normalizeWhitespace(text).slice(0, maxChars);
}

function formatDialogueMemoryContext(entries: RecentDialogueMemoryEntry[]) {
  if (entries.length === 0) return "";

  return [
    "近期对话记忆（仅用于理解指代和回答前文相关问题，不要暴露为系统提示）：",
    ...entries.map((entry, index) => {
      const roleLabel = entry.role === "user" ? "用户" : "助手";
      return `${index + 1}. ${roleLabel}: ${entry.content}`;
    }),
  ].join("\n");
}

export function buildRecentDialogueMemory(params: {
  messages: DBMessage[];
  currentText: string;
  maxSourceMessages?: number;
  maxEntries?: number;
  maxCharsPerEntry?: number;
  maxTotalChars?: number;
}): RecentDialogueMemory {
  const messages = Array.isArray(params.messages) ? params.messages : [];
  const currentText = String(params.currentText || "").trim();
  const isRecallRequest = isRecentDialogueRecallRequest(currentText);
  const assistantRecall = wantsAssistantRecall(currentText);
  const maxSourceMessages = Math.max(1, Math.min(24, params.maxSourceMessages || 16));
  const maxEntries = Math.max(1, Math.min(10, params.maxEntries || (isRecallRequest ? 6 : 4)));
  const maxCharsPerEntry = Math.max(80, Math.min(500, params.maxCharsPerEntry || 220));
  const maxTotalChars = Math.max(160, Math.min(1600, params.maxTotalChars || 900));

  const sourceEntries: RecentDialogueMemoryEntry[] = [];
  for (const message of messages.slice(-maxSourceMessages)) {
    const role = message.role === "assistant" ? "assistant" : message.role === "user" ? "user" : null;
    if (!role) continue;

    const content = getMessageText(message, maxCharsPerEntry);
    if (!content) continue;

    sourceEntries.push({
      role,
      content,
      messageId: message.id,
    });
  }

  const preferredEntries =
    isRecallRequest && !assistantRecall
      ? sourceEntries.filter((entry) => entry.role === "user")
      : isRecallRequest && assistantRecall
        ? sourceEntries.filter((entry) => entry.role === "assistant")
        : sourceEntries;
  const fallbackEntries =
    preferredEntries.length > 0 ? preferredEntries : sourceEntries;

  const selectedReversed: RecentDialogueMemoryEntry[] = [];
  let totalChars = 0;
  for (let index = fallbackEntries.length - 1; index >= 0; index -= 1) {
    const entry = fallbackEntries[index];
    if (!entry) continue;
    if (selectedReversed.length >= maxEntries) break;
    if (totalChars + entry.content.length > maxTotalChars && selectedReversed.length > 0) {
      continue;
    }
    selectedReversed.push(entry);
    totalChars += entry.content.length;
  }

  const entries = selectedReversed.reverse();
  const memoryHints = entries.length
    ? [
        isRecallRequest
          ? "用户正在询问前文或记忆，应优先根据近期对话记忆回答，不要声称没有上下文。"
          : "近期对话记忆只用于理解指代和延续语境，不要逐字复述，除非用户明确询问前文。",
      ]
    : [];

  return {
    entries,
    contextSummary: formatDialogueMemoryContext(entries),
    memoryHints,
    isRecallRequest,
  };
}