// plan-and-execute.ts
import {
  StateGraph,
  START,
  END,
  Annotation,
  MemorySaver,
} from "@langchain/langgraph";
 
const PlanExecuteState = Annotation.Root({
  // 用户原始任务
  input: Annotation<string>,
 
  // 当前执行计划（剩余步骤）
  plan: Annotation<string[]>({
    reducer: (_, update) => update, // 整体替换
    default: () => [],
  }),
 
  // 已完成步骤及其结果
  pastSteps: Annotation<Array<{ step: string; result: string }>>({
    reducer: (current, update) => [...current, ...update],
    default: () => [],
  }),
 
  // 最终回答（replanner 决定收工时写入）
  response: Annotation<string>({
    reducer: (_, update) => update,
    default: () => "",
  }),
});

import { ChatAnthropic } from "@langchain/anthropic";
import { z } from "zod";
 
const planSchema = z.object({
  steps: z
    .array(z.string())
    .describe("按顺序排列的执行步骤列表，每步一句话，3 到 7 步"),
});
 
const plannerModel = new ChatAnthropic({
  model: "claude-opus-4-7",
  temperature: 0,
}).withStructuredOutput(planSchema);
 
async function plannerNode(state: typeof PlanExecuteState.State) {
  const result = await plannerModel.invoke([
    {
      role: "system",
      content: `你是一个任务规划专家。把用户任务拆成 3 到 7 个有序步骤。
每个步骤必须：
1. 足够具体，可以独立交给执行者去做
2. 前后有清晰的依赖关系
3. 一句话讲清楚
 
最后一步必须是"基于前面所有信息给出最终答案"，不要写"汇总"这种空话。`,
    },
    { role: "user", content: state.input },
  ]);
 
  return { plan: result.steps };
}

import { createAgent } from "langchain";
import { tool } from "@langchain/core/tools";
 
const webSearch = tool(
  async ({ query }) => {
    // 真实场景接 Tavily / SerpAPI
    return `关于"${query}"的检索结果：行业 2025 年增长 28%，主要玩家包括 A、B、C 三家...`;
  },
  {
    name: "web_search",
    description: "从公开互联网检索事实信息",
    schema: z.object({ query: z.string() }),
  }
);
 
const analyze = tool(
  async ({ data, question }) => {
    return `针对"${question}"对数据"${data}"的分析：呈现集中度提升趋势，CR3 达到 65%...`;
  },
  {
    name: "analyze_data",
    description: "对一段数据进行结构化分析",
    schema: z.object({
      data: z.string(),
      question: z.string(),
    }),
  }
);
 
const executorAgent = createAgent({
  model: new ChatAnthropic({ model: "claude-sonnet-4-6", temperature: 0 }),
  tools: [webSearch, analyze],
  systemPrompt: "你是任务执行者。专注完成当前这一步，使用工具获取数据，给出简洁的步骤结果。",
});
 
async function executorNode(state: typeof PlanExecuteState.State) {
  // 当前要执行的是 plan 的第一项
  const currentStep = state.plan[0];
 
  // 给 executor 提供上下文：原始任务 + 已完成步骤 + 当前步骤
  const context = state.pastSteps
    .map((s, i) => `${i + 1}. ${s.step}\n   结果: ${s.result}`)
    .join("\n");
 
  const prompt = `原始任务：${state.input}
 
${context ? `已完成的步骤：\n${context}\n` : ""}
现在请执行这一步：${currentStep}
 
执行完后给出简洁的步骤结果（不超过 200 字）。`;
 
  const result = await executorAgent.invoke({
    messages: [{ role: "user", content: prompt }],
  });
 
  const lastMessage = result.messages.at(-1);
  const stepResult = extractText(lastMessage);
 
  return {
    pastSteps: [{ step: currentStep, result: stepResult }],
  };
}
 
function extractText(msg: unknown): string {
  if (!msg || typeof msg !== "object") return "";
  const m = msg as { content?: unknown; contentBlocks?: unknown };
  if (Array.isArray(m.contentBlocks)) {
    return (m.contentBlocks as Array<{ type?: string; text?: string }>)
      .filter((b) => b.type === "text")
      .map((b) => b.text ?? "")
      .join("");
  }
  return typeof m.content === "string" ? m.content : "";
}

const replanSchema = z.object({
  action: z.enum(["continue", "replan", "finish"]),
  updatedPlan: z
    .array(z.string())
    .optional()
    .describe("如果 action 是 replan，给出更新后的剩余步骤"),
  finalAnswer: z
    .string()
    .optional()
    .describe("如果 action 是 finish，给出最终回答"),
});
 
const replannerModel = new ChatAnthropic({
  model: "claude-sonnet-4-6",
  temperature: 0,
}).withStructuredOutput(replanSchema);
 
async function replannerNode(state: typeof PlanExecuteState.State) {
  const completed = state.pastSteps
    .map((s, i) => `${i + 1}. ${s.step}\n   结果: ${s.result}`)
    .join("\n");
 
  const remaining = state.plan.slice(1); // 把刚执行完的那一步从计划里去掉
 
  const result = await replannerModel.invoke([
    {
      role: "system",
      content: `你是任务规划评估者。根据已完成步骤的结果，决定下一步：
- continue: 按原剩余计划继续
- replan: 根据新信息调整剩余计划
- finish: 已经收集到足够信息，直接给最终答案
 
如果剩余步骤为空，必须 finish。`,
    },
    {
      role: "user",
      content: `原始任务：${state.input}
 
已完成：
${completed}
 
剩余计划：
${remaining.map((s, i) => `${i + 1}. ${s}`).join("\n") || "（空）"}`,
    },
  ]);
 
  if (result.action === "finish") {
    return { response: result.finalAnswer ?? "（无回答）", plan: [] };
  }
 
  if (result.action === "replan" && result.updatedPlan) {
    return { plan: result.updatedPlan };
  }
 
  // continue：把刚执行完的步骤从计划里弹掉
  return { plan: remaining };
}

function shouldEnd(state: typeof PlanExecuteState.State) {
  if (state.response) return END;
  if (state.plan.length === 0) return END;
  return "executor";
}
 
const graph = new StateGraph(PlanExecuteState)
  .addNode("planner", plannerNode)
  .addNode("executor", executorNode)
  .addNode("replanner", replannerNode)
  .addEdge(START, "planner")
  .addEdge("planner", "executor")
  .addEdge("executor", "replanner")
  .addConditionalEdges("replanner", shouldEnd, {
    executor: "executor",
    [END]: END,
  });
 
const app = graph.compile({ checkpointer: new MemorySaver() });

const result = await app.invoke(
  {
    input:
      "分析 2025 年中国新能源汽车市场的竞争格局，包括头部玩家市场份额和主要技术路线对比",
  },
  { configurable: { thread_id: "research-001" } }
);
 
console.log("===== 最终计划路径 =====");
result.pastSteps.forEach((s, i) => {
  console.log(`步骤 ${i + 1}: ${s.step}`);
  console.log(`  → ${s.result.slice(0, 150)}...`);
});
 
console.log("\n===== 最终回答 =====");
console.log(result.response);

for await (const update of app.stream(
  { input: "..." },
  { streamMode: "updates", configurable: { thread_id: "research-002" } }
)) {
  for (const [node, payload] of Object.entries(update)) {
    console.log(`[${node}]`, JSON.stringify(payload).slice(0, 200));
  }
}