ai-chatbot/lib/ai/prompt-memory.ts

import "server-only";

import type { AgentThreadState } from "@/lib/ai/runtime/agent-thread-state";

export interface PromptMemoryEntry {
  optimizedPrompt: string;
  sourceIntent?: string;
  sourceText?: string;
  reusable?: boolean;
}

function normalizeText(value: unknown, max = 2000): string | undefined {
  const normalized = String(value || "").trim();
  if (!normalized) return undefined;
  return normalized.slice(0, max);
}

export function extractPromptMemoryFromWorkingState(
  workingState?: {
    lastTurnSummary?: unknown;
  } | null,
): PromptMemoryEntry | null {
  const summary = workingState?.lastTurnSummary;
  if (!summary || typeof summary !== "object") return null;
  const summaryRecord = summary as Record<string, unknown>;
  const memory = summaryRecord.promptMemory;
  if (memory && typeof memory === "object") {
    const memoryRecord = memory as Record<string, unknown>;
    const optimizedPrompt = normalizeText(memoryRecord.optimizedPrompt, 4000);
    if (!optimizedPrompt) return null;
    return {
      optimizedPrompt,
      sourceIntent: normalizeText(memoryRecord.sourceIntent, 64),
      sourceText: normalizeText(memoryRecord.sourceText, 500),
      reusable:
        typeof memoryRecord.reusable === "boolean" ? memoryRecord.reusable : true,
    };
  }

  const resolutionMode = normalizeText(summaryRecord.resolutionMode, 64);
  const optimizedPrompt = normalizeText(summaryRecord.optimizedPrompt, 4000);
  if (resolutionMode === "prompt_only" && optimizedPrompt) {
    return {
      optimizedPrompt,
      sourceIntent: "image_gen",
      sourceText: normalizeText(summaryRecord.text, 500),
      reusable: true,
    };
  }
  return null;
}

export function extractConfirmedPromptFromWorkingState(
  workingState?: {
    lastTurnSummary?: unknown;
  } | null,
): string | undefined {
  const summary = workingState?.lastTurnSummary;
  if (!summary || typeof summary !== "object") return undefined;
  const summaryRecord = summary as Record<string, unknown>;

  const resolutionMode =
    typeof summaryRecord.resolutionMode === "string" ? summaryRecord.resolutionMode.trim() : "";
  const optimizedPrompt =
    typeof summaryRecord.optimizedPrompt === "string" ? summaryRecord.optimizedPrompt.trim() : "";

  if (resolutionMode !== "prompt_only" || !optimizedPrompt) return undefined;
  return optimizedPrompt;
}

export function buildPromptMemoryUpdate(params: {
  optimizedPrompt: string;
  sourceIntent?: string;
  sourceText?: string;
  reusable?: boolean;
}): PromptMemoryEntry {
  return {
    optimizedPrompt: normalizeText(params.optimizedPrompt, 4000) || "",
    sourceIntent: normalizeText(params.sourceIntent, 64),
    sourceText: normalizeText(params.sourceText, 500),
    reusable: params.reusable !== false,
  };
}

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

  const refersPreviousPrompt =
    /根据(上一个|上一句|刚才|上面|之前|那个)?提示词|按(上一个|上一句|刚才|上面|之前|那个)?提示词|用(上一个|上一句|刚才|上面|之前|那个)?提示词|用刚才输出的提示词|按刚才输出的提示词|按上面的提示词|沿用上个提示词|继续用那个提示词/.test(
      normalized,
    );

  if (!refersPreviousPrompt) return false;

  const wantsExecution =
    /生图|生成图片|出图|画出来|做图|继续生成|继续做|重新做|重绘|优化这张图|优化这个图|处理这张图|处理这个图|按这个做|照这个做|按此生成/.test(
      normalized,
    );

  if (wantsExecution) return true;

  if (params.hasImage) {
    return /优化|处理|重绘|参考这张图|结合这张图|基于这张图/.test(normalized);
  }

  return false;
}

export function getPromptMemoryPlanningHints(
  threadState?: AgentThreadState | null,
): string[] {
  const metadata = threadState?.metadata;
  const memory =
    metadata && typeof metadata === "object" && !Array.isArray(metadata)
      ? ((metadata.promptMemory as PromptMemoryEntry | undefined) || undefined)
      : undefined;

  if (!memory?.optimizedPrompt) return [];

  const preview = memory.optimizedPrompt.replace(/\s+/g, " ").trim().slice(0, 180);
  const hints = [
    "当前会话中存在上一轮可复用的优化提示词；如果用户明确要求沿用、根据上一个提示词继续做图，应优先复用这段提示词。",
    `上一轮优化提示词摘要：${preview}`,
  ];
  if (memory.sourceText) {
    hints.push(`该提示词来源于用户上一轮需求：${memory.sourceText}`);
  }
  return hints;
}