# multi_agent_demo.py
# 演示 Supervisor 模式的 Multi-Agent 协作
# 包含：Supervisor Agent、Researcher Agent、Writer Agent

import operator
from typing import TypedDict, List, Annotated, Literal
from langgraph.graph import StateGraph, START, END

# ============================================================
# 1. 定义共享 State
# 所有 Agent 读写同一个 State，通过 State 实现通信
# ============================================================

class MultiAgentState(TypedDict):
    topic: str                          # 研究主题
    search_results: List[str]           # 研究员收集的资料
    draft: str                          # 写作员的草稿
    final_output: str                   # 最终输出
    # operator.add 保证日志是追加而非覆盖，避免节点间相互覆盖历史记录
    messages: Annotated[List[str], operator.add]
    # 下一步执行哪个节点，由 Supervisor 设置
    next_worker: str
    # 记录迭代次数，防止死循环
    iteration_count: int


# ============================================================
# 2. Researcher Agent
# 职责：收集信息。实际场景这里会调用搜索工具或爬虫。
# ============================================================

def researcher_agent(state: MultiAgentState) -> dict:
    """
    研究员 Agent：负责信息收集。
    这里用模拟数据代替真实搜索，实际使用时接入 Tavily、Bing 等搜索 API。
    """
    topic = state["topic"]

    # 模拟搜索结果（实际场景替换为真实搜索调用）
    simulated_results = [
        f"【基础知识】{topic} 的定义：指通过...方式实现...的技术方法。",
        f"【应用案例】{topic} 在工业界的主要应用包括：自动化流程、智能决策支持等。",
        f"【最新进展】截至 2026 年，{topic} 领域的主要突破在于...研究方向取得了显著进展。",
        f"【常见问题】实践中 {topic} 面临的挑战主要有：数据质量、模型可解释性、部署成本。",
    ]

    return {
        "search_results": simulated_results,
        # messages 是 Annotated[List, operator.add]，返回列表会被追加，不会覆盖
        "messages": [f"[Researcher] 完成资料收集，找到 {len(simulated_results)} 条相关信息"],
    }


# ============================================================
# 3. Writer Agent
# 职责：基于研究结果生成文章草稿。
# ============================================================

def writer_agent(state: MultiAgentState) -> dict:
    """
    写作员 Agent：基于 Researcher 收集的资料撰写文章。
    实际场景这里会调用 LLM，传入 search_results 作为上下文。
    """
    topic = state["topic"]
    search_results = state.get("search_results", [])

    if not search_results:
        # 防御性处理：如果没有搜索结果，不应该被调用到，但加上保护更稳健
        return {
            "draft": f"[错误] 没有可用的研究资料，无法生成关于「{topic}」的文章。",
            "messages": ["[Writer] 错误：没有研究资料，无法写作"],
        }

    # 模拟基于资料生成文章（实际替换为 LLM 调用）
    results_text = "\n".join(f"- {r}" for r in search_results)
    draft = f"""# {topic} 综述

## 概述

{topic} 是当前技术领域的重要研究方向。

## 核心内容

基于最新研究资料整理如下：

{results_text}

## 总结

综合以上资料，{topic} 在实际应用中具有重要价值，未来发展空间广阔。
"""

    return {
        "draft": draft,
        "messages": [f"[Writer] 完成文章草稿，共 {len(draft)} 字符"],
    }


# ============================================================
# 4. Finalizer Agent
# 职责：对草稿做最终润色和输出
# ============================================================

def finalizer_agent(state: MultiAgentState) -> dict:
    """
    最终整合节点：检查草稿质量，生成最终输出。
    实际场景这里可以做格式检查、事实核对、排版优化等。
    """
    draft = state.get("draft", "")

    if not draft or draft.startswith("[错误]"):
        final_output = "任务未能完成，请检查错误信息后重试。"
    else:
        # 模拟最终润色（添加元信息）
        final_output = draft + f"\n\n---\n*本文由 Multi-Agent 系统自动生成，共经历 {state['iteration_count']} 次协作迭代。*"

    return {
        "final_output": final_output,
        "messages": ["[Finalizer] 最终报告生成完成"],
    }


# ============================================================
# 5. Supervisor Agent
# 这是整个系统的核心：决定下一步调用哪个 Worker
# ============================================================

def supervisor_agent(state: MultiAgentState) -> dict:
    """
    Supervisor：任务调度中枢。
    根据当前 State 决定下一步该做什么。

    决策逻辑：
    1. 没有搜索结果 → 先去搜索
    2. 有搜索结果但没有草稿 → 去写作
    3. 有草稿 → 进入最终化
    4. 迭代次数超限 → 强制结束，防止死循环
    """
    iteration = state.get("iteration_count", 0) + 1

    # 安全阀：超过最大迭代次数强制结束
    # 这是 Multi-Agent 系统中防止死循环的关键机制
    MAX_ITERATIONS = 5
    if iteration > MAX_ITERATIONS:
        return {
            "next_worker": "finalizer",
            "iteration_count": iteration,
            "messages": [f"[Supervisor] 达到最大迭代次数 {MAX_ITERATIONS}，强制结束"],
        }

    has_research = bool(state.get("search_results"))
    has_draft = bool(state.get("draft"))

    if not has_research:
        next_step = "researcher"
        reason = "尚未收集资料，分配给研究员"
    elif not has_draft:
        next_step = "writer"
        reason = "资料已收集，分配给写作员"
    else:
        next_step = "finalizer"
        reason = "草稿已完成，进入最终化"

    return {
        "next_worker": next_step,
        "iteration_count": iteration,
        "messages": [f"[Supervisor] 第 {iteration} 次决策：{reason} → {next_step}"],
    }


# ============================================================
# 6. 条件路由函数
# 根据 Supervisor 设置的 next_worker 字段决定跳转到哪个节点
# ============================================================

def route_by_supervisor(state: MultiAgentState) -> Literal["researcher", "writer", "finalizer"]:
    """
    条件边的路由函数。
    读取 Supervisor 写入的 next_worker 字段，返回下一个节点名称。
    这是 Supervisor 模式的关键连接点。
    """
    return state["next_worker"]


# ============================================================
# 7. 构建图
# ============================================================

def build_multi_agent_graph():
    graph = StateGraph(MultiAgentState)

    # 注册所有节点
    graph.add_node("supervisor", supervisor_agent)
    graph.add_node("researcher", researcher_agent)
    graph.add_node("writer", writer_agent)
    graph.add_node("finalizer", finalizer_agent)

    # 入口：从 START 进入 Supervisor
    graph.add_edge(START, "supervisor")

    # Supervisor 出发的条件边：根据 next_worker 字段决定去哪里
    graph.add_conditional_edges(
        "supervisor",
        route_by_supervisor,
        {
            "researcher": "researcher",
            "writer": "writer",
            "finalizer": "finalizer",
        }
    )

    # Worker 完成后，回到 Supervisor 进行下一轮决策
    # 这形成了 Worker → Supervisor → Worker 的循环结构
    graph.add_edge("researcher", "supervisor")
    graph.add_edge("writer", "supervisor")

    # Finalizer 完成后直接结束
    graph.add_edge("finalizer", END)

    return graph.compile()


# ============================================================
# 8. 运行演示
# ============================================================

def run_demo():
    app = build_multi_agent_graph()

    initial_state = {
        "topic": "大型语言模型的 Agent 应用",
        "search_results": [],
        "draft": "",
        "final_output": "",
        "messages": [],
        "next_worker": "",
        "iteration_count": 0,
    }

    print("=" * 60)
    print(f"研究主题：{initial_state['topic']}")
    print("=" * 60)

    # 使用 stream 模式观察每个节点的执行过程
    for event in app.stream(initial_state, stream_mode="updates"):
        for node_name, updates in event.items():
            if node_name.startswith("__"):
                continue
            # 打印每个节点产生的日志消息
            new_messages = updates.get("messages", [])
            for msg in new_messages:
                print(msg)

    # 获取最终结果
    final_state = app.invoke(initial_state)
    print("\n" + "=" * 60)
    print("最终输出：")
    print("=" * 60)
    print(final_state["final_output"])

    print("\n执行日志：")
    for msg in final_state["messages"]:
        print(f"  {msg}")


if __name__ == "__main__":
    run_demo()

from langgraph.graph import StateGraph, START, END

# 子图：独立的研究 Agent
def build_research_subgraph():
    sub = StateGraph(ResearchState)
    sub.add_node("fetch", fetch_node)
    sub.add_node("summarize", summarize_node)
    sub.add_edge(START, "fetch")
    sub.add_edge("fetch", "summarize")
    sub.add_edge("summarize", END)
    # 编译子图
    return sub.compile()

# 父图中把子图作为一个节点注册
research_subgraph = build_research_subgraph()

parent_graph = StateGraph(MainState)
# 子图可以直接作为节点传入，LangGraph 会自动处理 State 的映射
parent_graph.add_node("research", research_subgraph)
# ...

MAX_ITERATIONS = 10

def supervisor(state):
    count = state.get("iteration_count", 0) + 1
    if count >= MAX_ITERATIONS:
        return {"next_worker": "end", "iteration_count": count}
    # 正常决策逻辑...

# 在 invoke 时设置，超过后抛出 GraphRecursionError
app.invoke(initial_state, config={"recursion_limit": 25})

class State(TypedDict):
    is_complete: bool  # 明确的完成标志

def supervisor(state):
    if state.get("is_complete"):
        return {"next_worker": "end"}
    # ...