lovart3/lib/ai/auto-agent-graph.ts

import { Annotation, StateGraph, START, END } from "@langchain/langgraph";
import type { AnalysisResult } from "@/lib/ai/smart-analyzer";
import type { RouteResult } from "@/lib/ai/model-router";

interface HistoryMessage {
  role: "user" | "assistant";
  content: string;
}

export interface AutoAgentInput {
  text: string;
  hasImage: boolean;
  imageUrls?: string[];
  imageFromHistory?: boolean;
  historyMessages?: HistoryMessage[];
  context?: string;
}

export interface AutoAgentResult {
  analysis?: AnalysisResult;
  friendlyMessage?: string;
  optimizedPrompt?: string;
  route?: RouteResult;
}

const AutoAgentState = Annotation.Root({
  text: Annotation<string>,
  hasImage: Annotation<boolean>,
  imageUrls: Annotation<string[]>,
  imageFromHistory: Annotation<boolean>,
  historyMessages: Annotation<HistoryMessage[]>,
  context: Annotation<string | null>,
  analysis: Annotation<AnalysisResult | null>,
  friendlyMessage: Annotation<string | null>,
  optimizedPrompt: Annotation<string | null>,
  route: Annotation<RouteResult | null>,
});

type AutoAgentStateType = typeof AutoAgentState.State;

function compactText(text: string, max = 48) {
  const normalized = text.replace(/\s+/g, " ").trim();
  if (normalized.length <= max) return normalized;
  return `${normalized.slice(0, max)}...`;
}

function buildStrategicFriendlyMessage(
  analysis: AnalysisResult,
  state: AutoAgentStateType,
): string | null {
  const fallbackQuestion = analysis.clarifyingQuestions?.[0];

  if (analysis.needsClarification) {
    return analysis.conversationalReply || fallbackQuestion || "我先确认一下，你希望我优先调整哪一部分？";
  }

  if (analysis.intent === "image_gen") {
    const demand = compactText(state.text);
    return [
      "收到，我来帮你走自动模式生成方案。",
      `- 需求理解：${demand}`,
      "- 执行策略：先锁定主体与场景，再控制风格/光影/构图，最后输出可继续迭代的版本。",
      "- 继续优化：你可以直接说“保留A，改B”，我会按最新要求快速迭代。",
    ].join("\n");
  }

  if (analysis.intent === "image_edit") {
    const sourceHint = state.hasImage
      ? state.imageFromHistory
        ? "我会基于你上一轮图片继续编辑。"
        : "我会基于你刚上传的图片编辑。"
      : "我会先按当前上下文进行编辑。";
    return [
      "收到，我来帮你精准编辑这张图。",
      `- 执行基准：${sourceHint}`,
      "- 处理方式：只改你提到的目标区域，其他主体、构图与风格尽量保持不变。",
      "- 交付方式：先出一版结果，你再告诉我细调方向（如颜色、材质、氛围、细节）。",
    ].join("\n");
  }

  return analysis.conversationalReply || null;
}

const analyzeNode = async (state: AutoAgentStateType) => {
  const { analyzeIntent } = await import("./smart-analyzer");
  const analysis = await analyzeIntent({
    text: state.text,
    hasImage: state.hasImage,
    imageUrls: state.imageUrls,
    imageFromHistory: state.imageFromHistory,
    historyMessages: state.historyMessages,
    context: state.context ?? undefined,
  });

  const fallbackQuestion = analysis.needsClarification
    ? analysis.clarifyingQuestions?.[0]
    : undefined;

  const strategicMessage = buildStrategicFriendlyMessage(analysis, state);

  return {
    analysis,
    friendlyMessage: strategicMessage || analysis.conversationalReply || fallbackQuestion || null,
  };
};

const shouldOptimize = (state: AutoAgentStateType) => {
  if (!state.analysis) return "route";
  if (state.analysis.needsClarification) return "route";
  if (state.analysis.intent === "image_gen") {
    return "optimize";
  }
  if (state.analysis.intent === "image_edit" && state.analysis.subIntent) {
    return "optimize";
  }
  return "route";
};

const optimizeNode = async (state: AutoAgentStateType) => {
  if (!state.analysis) {
    return {};
  }
  if (state.analysis.intent === "image_edit" && !state.analysis.subIntent) {
    return {};
  }
  if (state.analysis.intent !== "image_edit" && state.analysis.intent !== "image_gen") {
    return {};
  }
  const { optimizePromptWithLLM } = await import("./prompt-optimizer");
  const optimized = await optimizePromptWithLLM({
    originalPrompt: state.text,
    intent: state.analysis.intent,
    subIntent: state.analysis.subIntent,
    historyMessages: state.historyMessages,
    context: state.context ?? undefined,
  });
  return { optimizedPrompt: optimized.optimizedPrompt };
};

const routeNode = async (state: AutoAgentStateType) => {
  const { routeToModel } = await import("./model-router");
  const route = await routeToModel({
    intent: state.analysis?.intent || "chat",
    subIntent: state.analysis?.subIntent,
    hasImage: state.hasImage,
  });
  return { route };
};

let compiledGraph: ReturnType<typeof buildGraph> | null = null;

function buildGraph() {
  return new StateGraph(AutoAgentState)
    .addNode("analyze", analyzeNode)
    .addNode("optimize", optimizeNode)
    .addNode("route", routeNode)
    .addEdge(START, "analyze")
    .addConditionalEdges("analyze", shouldOptimize, {
      optimize: "optimize",
      route: "route",
    })
    .addEdge("optimize", "route")
    .addEdge("route", END)
    .compile();
}

function getGraph() {
  if (!compiledGraph) {
    compiledGraph = buildGraph();
  }
  return compiledGraph;
}

export async function runAutoAgentGraph(params: AutoAgentInput): Promise<AutoAgentResult> {
  const graph = getGraph();
  const result = await graph.invoke({
    text: params.text,
    hasImage: params.hasImage,
    imageUrls: params.imageUrls ?? [],
    imageFromHistory: params.imageFromHistory ?? false,
    historyMessages: params.historyMessages ?? [],
    context: params.context ?? null,
    analysis: null,
    friendlyMessage: null,
    optimizedPrompt: null,
    route: null,
  });

  return {
    analysis: result.analysis ?? undefined,
    friendlyMessage: result.friendlyMessage ?? undefined,
    optimizedPrompt: result.optimizedPrompt ?? undefined,
    route: result.route ?? undefined,
  };
}