ai-chatbot/lib/ai/image-workflow-memory.ts

import "server-only";

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

export interface ImageWorkflowMemory {
  lastIntent?: "image_gen" | "image_edit" | "vision";
  lastOptimizedPrompt?: string;
  lastUserPrompt?: string;
  lastRequestedImageCount?: number;
  lastAspectRatio?: string;
  lastImageSize?: string;
  lastSourceImageUrls?: string[];
  lastResultImageUrls?: string[];
}

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

function normalizeImageUrls(value: unknown): string[] | undefined {
  if (!Array.isArray(value)) return undefined;
  const urls = Array.from(
    new Set(
      value
        .map((item) => String(item || "").trim())
        .filter((url) => /^https?:\/\//i.test(url))
    )
  );
  return urls.length > 0 ? urls : undefined;
}

function normalizeImageCount(value: unknown): number | undefined {
  const parsed = Number(value);
  if (!Number.isFinite(parsed)) return undefined;
  const rounded = Math.round(parsed);
  if (rounded < 1 || rounded > 12) return undefined;
  return rounded;
}

export function extractImageWorkflowMemory(
  workingState?: {
    lastTurnSummary?: unknown;
    lastResolvedTargetImageUrls?: string[] | null;
    lastGeneratedImageUrls?: string[] | null;
  } | null,
): ImageWorkflowMemory | null {
  const summary = workingState?.lastTurnSummary;
  let memory: ImageWorkflowMemory | null = null;

  if (summary && typeof summary === "object") {
    const summaryRecord = summary as Record<string, unknown>;
    const imageWorkflowMemory = summaryRecord.imageWorkflowMemory;
    if (imageWorkflowMemory && typeof imageWorkflowMemory === "object") {
      const record = imageWorkflowMemory as Record<string, unknown>;
      memory = {
        lastIntent:
          record.lastIntent === "image_gen" ||
          record.lastIntent === "image_edit" ||
          record.lastIntent === "vision"
            ? record.lastIntent
            : undefined,
        lastOptimizedPrompt: normalizeText(record.lastOptimizedPrompt, 4000),
        lastUserPrompt: normalizeText(record.lastUserPrompt, 1000),
        lastRequestedImageCount: normalizeImageCount(record.lastRequestedImageCount),
        lastAspectRatio: normalizeText(record.lastAspectRatio, 32),
        lastImageSize: normalizeText(record.lastImageSize, 32),
        lastSourceImageUrls: normalizeImageUrls(record.lastSourceImageUrls),
        lastResultImageUrls: normalizeImageUrls(record.lastResultImageUrls),
      };
    }
  }

  const fallbackTargetUrls = normalizeImageUrls(workingState?.lastResolvedTargetImageUrls);
  const fallbackGeneratedUrls = normalizeImageUrls(workingState?.lastGeneratedImageUrls);
  if (!memory) {
    if (!fallbackTargetUrls && !fallbackGeneratedUrls) return null;
    memory = {};
  }

  if (!memory.lastSourceImageUrls && fallbackTargetUrls) {
    memory.lastSourceImageUrls = fallbackTargetUrls;
  }
  if (!memory.lastResultImageUrls && fallbackGeneratedUrls) {
    memory.lastResultImageUrls = fallbackGeneratedUrls;
  }

  return memory;
}

export function buildImageWorkflowMemoryUpdate(params: {
  existing?: ImageWorkflowMemory | null;
  lastIntent?: "image_gen" | "image_edit" | "vision";
  lastOptimizedPrompt?: string;
  lastUserPrompt?: string;
  lastRequestedImageCount?: number;
  lastAspectRatio?: string;
  lastImageSize?: string;
  lastSourceImageUrls?: string[] | null;
  lastResultImageUrls?: string[] | null;
}) {
  return {
    lastIntent: params.lastIntent ?? params.existing?.lastIntent,
    lastOptimizedPrompt:
      normalizeText(params.lastOptimizedPrompt, 4000) ??
      params.existing?.lastOptimizedPrompt,
    lastUserPrompt:
      normalizeText(params.lastUserPrompt, 1000) ?? params.existing?.lastUserPrompt,
    lastRequestedImageCount:
      normalizeImageCount(params.lastRequestedImageCount) ??
      params.existing?.lastRequestedImageCount,
    lastAspectRatio:
      normalizeText(params.lastAspectRatio, 32) ?? params.existing?.lastAspectRatio,
    lastImageSize:
      normalizeText(params.lastImageSize, 32) ?? params.existing?.lastImageSize,
    lastSourceImageUrls:
      normalizeImageUrls(params.lastSourceImageUrls) ??
      params.existing?.lastSourceImageUrls,
    lastResultImageUrls:
      normalizeImageUrls(params.lastResultImageUrls) ??
      params.existing?.lastResultImageUrls,
  } satisfies ImageWorkflowMemory;
}

export function getImageWorkflowPlanningHints(
  threadState?: AgentThreadState | null,
): string[] {
  const metadata = threadState?.metadata;
  const memory =
    metadata && typeof metadata === "object" && !Array.isArray(metadata)
      ? ((metadata.imageWorkflowMemory as ImageWorkflowMemory | undefined) || undefined)
      : undefined;
  if (!memory) return [];

  const hints: string[] = [];
  if (memory.lastOptimizedPrompt) {
    hints.push(
      `上一轮可复用提示词摘要：${memory.lastOptimizedPrompt.replace(/\s+/g, " ").slice(0, 180)}`
    );
  }
  if (memory.lastRequestedImageCount) {
    hints.push(`用户上一轮期望出图数量：${memory.lastRequestedImageCount} 张。`);
  }
  if (memory.lastAspectRatio) {
    hints.push(`用户上一轮常用画幅：${memory.lastAspectRatio}。`);
  }
  if (memory.lastImageSize) {
    hints.push(`用户上一轮常用清晰度：${memory.lastImageSize}。`);
  }
  if (memory.lastSourceImageUrls?.length) {
    hints.push(`上一轮存在 ${memory.lastSourceImageUrls.length} 张参考/目标图。`);
  }
  if (memory.lastResultImageUrls?.length) {
    hints.push(`上一轮生成结果数：${memory.lastResultImageUrls.length} 张。`);
  }
  return hints;
}