大模型的开发应用（十七）：多智能体项目

1 项目需求

假如我想让AI生成一份关于成都市GDP增长率分析的报告，我们假设生成报告需要通过以下步骤：（1）收集资料；（2）对资料中的数据进行计算分析；（3）撰写报告。

三个步骤可以对应图中的三个节点，并且每个节点，我们都可以让大模型完成相应的功能。

本项目的工作流节点流程图如下所示：

图中，边上有判断条件的，表示条件边。

2 模型配置

我们先把模型推理服务吊起来：

lmdeploy serve api_server /data/coding/models/Qwen/Qwen3-4B --cache-max-entry-count 0.4 --reasoning-parser deepseek-r1

我用的是 RTX 3060，显存只有12GB，4B的模型要想能放下，需要指定–cache-max-entry-count，这个参数表示 KV Cache 占用的显存比例，默认是0.8，如果使用默认值，则会报错，说显存不够，我这里设定其为 0.4。

因为 Qwen3 是带有思维链的模型，像 deepseek 一样，但我不想打印思考过程，所以我这里需要设定 --reasoning-parser 参数为 deepseek-r1，这个参数表示将模型的输出按照 deepseek-r1 的格式来解析。

我们先把要用到的库导进来：

from typing import Literal, Annotated
from langchain_openai import ChatOpenAI
from langgraph.graph import StateGraph, END
from langchain_core.messages import HumanMessage, SystemMessage, AIMessage
from pydantic import BaseModel, Field
import re

然后我们进行模型配置，这里我们使用 Qwen3-4B：

# ================== 模型配置 ==================
model = ChatOpenAI(
    openai_api_base="http://localhost:23333/v1",
    model_name="/data/coding/models/Qwen/Qwen3-4B",
    openai_api_key="XXXXX",
    max_tokens=400,
    temperature=0.2,
)

3 状态类

这里我们不但要定义属性，还要定义一些方法，主要就是把字典转为状态类，或者把状态类中的属性转成字典，代码如下：

# ================== 状态类 ==================
class AgentState(BaseModel):
    # messages 用于接收上一个节点或者起始节点的信息
    messages: Annotated[list, Field(default_factory=list)]  
    # next_agent 用于指定下一个执行节点，下一个节点只能是 researcher、calculator、writer 和 done 其中之一
    next_agent: Literal["researcher", "calculator", "writer", "done"] = "researcher"   
    # step_count 用于记录步数
    step_count: int = 0
    # is_finalized 用于记录是否终止
    is_finalized: bool = False
    search_data: dict = Field(default_factory=dict)  # 新增搜索数据存储

    @classmethod
    def from_raw(cls, raw):
        # 这个函数用于判断 raw 是否为 cls 类（AgentState）对象
        # 如果 raw 属于 AgentState 类，则直接把 raw 返回，否则转成 AgentState 类对象
        """统一状态转换方法"""
        if isinstance(raw, cls):
            return raw
        if isinstance(raw, dict):
            return cls(
                messages=raw.get('messages', []),
                next_agent=raw.get('next_agent', 'researcher'),
                step_count=raw.get('step_count', 0),
                is_finalized=raw.get('is_finalized', False),
                search_data=raw.get('search_data', {})
            )
        raise ValueError(f"无效状态类型: {type(raw)}")  

    def to_safe_dict(self):
        """安全转换为字典"""
        return {
            "messages": [msg.dict() for msg in self.messages],
            "next_agent": self.next_agent,
            "step_count": self.step_count,
            "is_finalized": self.is_finalized,
            "search_data": self.search_data
        }

4 入口节点

这里我们不使用 langgraph 预定义的 SATRT，而是自己创建一个入口节点，方便进行逻辑判断。

# ================== 创建Supervisor节点（相当于 START 节点）函数 ==================
def supervisor(raw_state):
    print("================执行 supervisor 节点================")
    state = AgentState.from_raw(raw_state)
    if state.is_finalized:
        return state

    try:
        # 判断步数是否达到循环上限，若达到上限，则强行走 writer 节点
        # 因为 writer 节点执行完之后会添加 [END] 标识，相当于强制结束
        if state.step_count >= 8:   
            print("[强制终止] 达到最大步数限制")
            return AgentState(
                messages=state.messages,
                next_agent="writer",
                step_count=state.step_count + 1,
                is_finalized=False
            )

        # 消息转换，转成 langchain 能处理的格式
        messages = convert_messages(state.messages)

        # 在消息列表中添加流程控制提示词
        control_prompt = '''请严格选择下一步：\n1.需要更多数据 → researcher\n2.需要计算 → calculator \n3.生成报告 → writer\n4. 完成 → done'''
        messages = messages + [SystemMessage(content=control_prompt)]

        # 模型推理
        response = model.invoke(messages)

        # 查看模型的回复中，是否包含 researcher|calculator|writer，若不包含，则 new_agent 置为 done
        decision = re.search(r'\b(researcher|calculator|writer)\b', response.content.lower())
        new_agent = decision.group(1).lower() if decision else "done"

        # 若 new_agent 为 done，但是还没有经过 writer 节点，那只能强制调用 writer 生成报告
        # 如果经过 writer，则模型回复中会带有 [END] 字符串，因此可以据此判断是否经过 writer 节点
        if new_agent == "done" and not any("[END]" in msg.content for msg in messages):
            new_agent = "writer"

        return AgentState(
            messages=state.messages,
            next_agent=new_agent,
            step_count=state.step_count + 1,
            is_finalized=False
        )
    except Exception as e:
        # 如果发生错误，则强制结束
        print(f"[Supervisor错误] {str(e)}")
        return AgentState(
            messages=state.messages,
            next_agent="writer",
            step_count=state.step_count + 1,
            is_finalized=False
        )

这里出现了一个消息转换函数，它的功能是将消息中的字典转换为 langchain 中的消息格式：

# ================== 消息转换 ==================
def convert_messages(raw_messages):
    converted = []
    for msg in raw_messages:
        if isinstance(msg, dict):
            try:
                # 获取当前消息的角色，如果不存在，则置为 assistant
                role = msg.get("role", "assistant")
                # 获取当前消息的内容
                content = msg.get("content", "")
                # 根据角色获取类名
                role_dict = {
                    "user": HumanMessage,
                    "system": SystemMessage,
                    "assistant": AIMessage
                }
                msg_type = role_dict[role]
                # 将内容封装成 langchain 中对应角色的标准格式
                std_msg = msg_type(content=content)
                # 添加到消息列表
                converted.append(std_msg)
            except KeyError:
                converted.append(AIMessage(content=str(msg)))
        elif hasattr(msg, 'content'):
            converted.append(msg)
        else:
            converted.append(AIMessage(content=str(msg)))
    return converted

5 创建节点

因为在向图中添加节点的时候，使用的是函数名，所以如果我们使用一个函数来创建节点，那么需要使用内置函数，外层函数返回的是内层函数的函数名。

本项目中，节点与节点之间的区别在于，使用了不同的提示词，我们把提示词封装成系统消息，然后将其作为最后一条消息加入到消息列表中，随后将消息列表输入到模型中，如果当前节点是 writer，那么要给模型输出的末尾加上 [END]。下一个节点具体是什么，需要从模型的回复中解析出来，并赋值给状态类的 next_agent 变量，因此，提示词中需要让模型给出下一个节点的名称。

代码如下：

# ================== 创建普通节点函数 ==================
def create_agent(role: str, prompt: str):
    # 这里 role 不是 AI、HUMAN、SYSTEM 等角色，而是指具体哪个节点
    def agent(raw_state):
        print(f"================执行 {role} 节点================")
        state = AgentState.from_raw(raw_state)
        if state.is_finalized:
            return state

        try:
            # 消息转换
            messages = convert_messages(state.messages)

            # 往messages中添加系统信息，让节点处的模型来选择下一步该走哪个节点
            # 如果当前节点是 writer，那么 prompt 是“生成包含[END]标识的最终报告”
            messages.append(SystemMessage(
                content=f"{prompt}\n最后必须用'建议下一步：选项'格式结尾（选项只能是researcher/calculator/writer/done）"
            ))

            # 获取模型响应
            response = model.invoke(messages)
            print(f"\n[{role}输出] {response.content[:500]}...")  # 限制日志长度

            # 如果当前节点是 writer，且模型的回复中不含 [END] 标识，那就强制添加
            # 虽然在提示词中要求 writer 节点的输出以 [END]，但我们使用的是小模型，小模型的指令遵循能力有限
            if role == "writer" and "[END]" not in response.content:
                response.content += "\n[END]"
                print(f"[{role}修正] 已添加[END]标识")

            # 解析下一步建议，模型未必会按照提示词来输出，所以这里可能需要强制匹配
            next_agent = "done"     # 这里先预定义下一个节点为 End
            
            # 这里使用正则匹配，:= 是海象运算符
            if match := re.search(r'建议下一步\s*[:：]\s*(\w+)', response.content):
                # 从match中解析出suggestion
                suggestion = match.group(1).lower()

                # 如果suggestion在集合是 researcher、calculator、writer 其中之一，
                # 那下一个节点就按照 suggestion，否则就用预定义的节点，即 END
                if suggestion in {"researcher", "calculator", "writer", "done"}:
                    next_agent = suggestion

            return AgentState(
                messages=messages + [response],
                next_agent=next_agent,
                step_count=state.step_count + 1,
                is_finalized="[END]" in response.content    # 如果 [END] 在模型答复中，则 is_finalized 为True
            )
        except Exception as e:
            print(f"[{role}错误] {str(e)}")
            return AgentState(
                messages=state.messages,
                next_agent="done",
                step_count=state.step_count + 1,
                is_finalized=True
            )
    return agent

6 执行入口

执行入口首先要建工作流，即建图，然后添加节点，添加边，构建好后进行编译，编译结束后是流式输出。代码如下：

# ================== 执行入口 ==================
def main():
    # ================== 工作流配置（建图） ==================
    builder = StateGraph(AgentState)
    builder.add_node("supervisor", supervisor)  # supervisor 相当于开始节点，
    builder.add_node("researcher", create_agent("researcher", "收集经济数据并分析趋势"))
    builder.add_node("calculator", create_agent("calculator", "执行GDP增长率计算"))
    builder.add_node("writer", create_agent("writer", "生成包含[END]标识的最终报告"))

    # 添加条件边
    # 因为逻辑不是串行的，这里只有发起者是固定的，但边的另一端是哪个节点，需要通过状态来判断
    # 多智能体建图的时候一般都是按照下面这样的方式，即构建一条从一个节点到多个节点的条件边
    builder.add_conditional_edges(
        "supervisor",
        lambda state: state.next_agent,
        {
            "researcher": "researcher",
            "calculator": "calculator", 
            "writer": "writer",
            "done": END
        }
    )

    # 为每一个节点构建一条指向 supervisor 的边
    for agent in ["researcher", "calculator", "writer"]:
        builder.add_edge(agent, "supervisor")

    # 设定入口
    builder.set_entry_point("supervisor")

    # 编译图
    workflow = builder.compile()

    try:
        # 初始化状态时确保类型正确
        initial_state = AgentState(messages=[HumanMessage(content="成都市GDP增长率分析")])
        
        # 流式输出
        for step in workflow.stream(initial_state):
            node_name, raw_state = step.popitem()       # 这里 raw_state 是字典
            current_state = AgentState.from_raw(raw_state)
            
            # print(f"\n[系统状态] 当前节点: {node_name}")
            print(f"下一步: {current_state.next_agent}")
            print(f"步数: {current_state.step_count}")
            print(f"完成状态: {current_state.is_finalized}")

            # 检查终止条件
            if current_state.is_finalized or node_name == "__end__":
                print("\n================流程完成================")
                if any("[END]" in msg.content for msg in current_state.messages):
                    print("最终报告内容:")
                    print(next(msg.content for msg in reversed(current_state.messages) if "[END]" in msg.content))
                else:
                    print("警告：未检测到完整报告")
                break
                
    except Exception as e:
        print(f"\n!!! 流程异常终止: {str(e)}")
        if 'current_state' in locals():
            print("最后状态:", current_state.to_safe_dict())


if __name__ == "__main__":
    main()

从建图的程序可以看到，三个功能节点（researcher、writer、calculator）都有一条边指向了 supervisor，并且不是条件边，也就是说三个功能节点执行后，都会无条件执行supervisor，这使得三个功能节点执行后状态中的 next_agent 属性并不可靠。

其实也没必要让三个功能节点执行时给出下一步，因为下一步一定是 supervisor，而 next_agent 会在 supervisor 节点执行时更新，next_agent 只在 supervisor 执行完选择下一个节点时有用，因为是条件边。

当然，上面的程序，执行完 writer 节点后，就直接跳出了，没有执行 supervisor，因为已经把报告写出来了，没必要返回 supervisor。

输出：

================执行 supervisor 节点================
下一步: researcher
步数: 1
完成状态: False
================执行 researcher 节点================

[researcher输出] 

成都市GDP增长率分析（2020-2023年）：  
1. **数据来源**：国家统计局、成都市统计局、《成都统计年鉴》。  
2. **趋势分析**：  
   - 2020年：受疫情影响，GDP增速降至3.2%（全国平均为2.3%），但成都通过数字经济、服务业等韧性产业快速恢复。  
   - 2021年：增速回升至6.1%，高于全国平均水平（8.1%），受益于成渝双城经济圈政策和电子信息产业扩张。  
   - 2022年：增速5.8%，略低于全国（8.1%），但成都作为西部经济高地，增速仍居全国前列。  
   - 2023年：预计增速5.5%-6%，受消费复苏和科技创新驱动，但面临房地产调整和外部环境压力。  
3. **关键驱动因素**：  
   - **产业结构升级**：电子信息、生物医药、新能源等新兴产业占比提升。  
   - **政策支持**：成渝双城经济圈建设、西部科学城规划。  
   - **消费市场**：成都作为“天府之国”，文旅、餐饮等服务业持续增长。  
4. **挑战**：房地产行业调整、外部经济不确定性、人口老龄化。  

建议下一步：**re...
下一步: done
步数: 2
完成状态: False
================执行 supervisor 节点================
下一步: writer
步数: 3
完成状态: False
================执行 writer 节点================

[writer输出] 

**成都市GDP增长率分析报告（2020-2023年）**  
[END]  

**核心结论**：  
- 成都市GDP年均增速稳定在5.5%-6.1%，高于全国平均水平（2020-2023年全国平均增速约4.5%）。  
- 2023年增速预计5.5%-6%，受消费复苏、科技创新和成渝双城经济圈政策推动，但面临房地产调整和外部环境压力。  
- 产业结构持续优化，电子信息、生物医药、新能源等新兴产业贡献率显著提升。  

**建议下一步：**  
done...
下一步: done
步数: 4
完成状态: True

================流程完成================
最终报告内容:


**成都市GDP增长率分析报告（2020-2023年）**  
[END]  

**核心结论**：  
- 成都市GDP年均增速稳定在5.5%-6.1%，高于全国平均水平（2020-2023年全国平均增速约4.5%）。  
- 2023年增速预计5.5%-6%，受消费复苏、科技创新和成渝双城经济圈政策推动，但面临房地产调整和外部环境压力。  
- 产业结构持续优化，电子信息、生物医药、新能源等新兴产业贡献率显著提升。  

**建议下一步：**  
done

上面的 “最终报告内容” 之所以在 成都市GDP增长率分析报告（2020-2023年 后面会有个 [END] 标识符，是因为在执行 writer 节点的时候，我们的提示词为 生成包含[END]标识的最终报告，结果生成的 [END] 标识符并不是在末尾，而是在标题的中间，从打印结果也能看到：

7 调用商业模型实现

这里我们调用智谱商业大模型，模型配置的代码修改如下：

# ================== 模型配置 ==================
model = ChatOpenAI(
    openai_api_base="https://open.bigmodel.cn/api/paas/v4",
    model_name="glm-4-plus",
    openai_api_key="XXXXX",		# 去智谱官网申请 API Key
    max_tokens=400,
    temperature=0.2
)

输出：

================执行 supervisor 节点================
下一步: writer
步数: 1
完成状态: False
================执行 writer 节点================

[writer输出] 成都市作为中国西南地区的重要经济中心，其GDP增长率一直是衡量当地经济发展状况的重要指标。以下是对成都市GDP增长率的分析，主要从历史数据、影响因素和未来趋势三个方面进行探讨。

### 一、历史数据分析

1. **近年GDP增长率**：
   - **2015-2020年**：这一时期，成都市的GDP增长率基本保持在7%-8%之间，显示出较为稳健的增长态势。其中，2018年和2019年受全球经济形势和国内经济结构调整的影响，增长率略有放缓。
   - **2020年**：受新冠疫情影响，成都市GDP增长率有所下降，但仍然保持在正增长区间，显示出较强的经济韧性。
   - **2021-2022年**：随着疫情逐步得到控制和经济恢复，成都市的GDP增长率有所回升，重回7%以上的增长水平。

2. **产业结构变化**：
   - **第一产业**：占比逐年下降，但对农业现代化的投入增加，提升了农业产值。
   - **第二产业**：工业和服务业成为拉动GDP增长的主要动力，特别是电子信息、汽车制造等高新技术产业的快速发展。
   - **第三产业**：服务业占比持续提升，尤其是金融、...
[writer修正] 已添加[END]标识
下一步: done
步数: 2
完成状态: True

================流程完成================
最终报告内容:
成都市作为中国西南地区的重要经济中心，其GDP增长率一直是衡量当地经济发展状况的重要指标。以下是对成都市GDP增长率的分析，主要从历史数据、影响因素和未来趋势三个方面进行探讨。

### 一、历史数据分析

1. **近年GDP增长率**：
   - **2015-2020年**：这一时期，成都市的GDP增长率基本保持在7%-8%之间，显示出较为稳健的增长态势。其中，2018年和2019年受全球经济形势和国内经济结构调整的影响，增长率略有放缓。
   - **2020年**：受新冠疫情影响，成都市GDP增长率有所下降，但仍然保持在正增长区间，显示出较强的经济韧性。
   - **2021-2022年**：随着疫情逐步得到控制和经济恢复，成都市的GDP增长率有所回升，重回7%以上的增长水平。

2. **产业结构变化**：
   - **第一产业**：占比逐年下降，但对农业现代化的投入增加，提升了农业产值。
   - **第二产业**：工业和服务业成为拉动GDP增长的主要动力，特别是电子信息、汽车制造等高新技术产业的快速发展。
   - **第三产业**：服务业占比持续提升，尤其是金融、物流、文化旅游等现代服务业的增长显著。

### 二、影响因素分析

1. **政策支持**：
   - 国家对西部地区的扶持政策，如西部大开发战略，为成都市经济发展提供了政策红利。
   - 成都市政府出台的一系列招商引资、科技创新和产业升级政策，有效推动了经济增长。

2. **区位优势**：
   - 成都市地处西南地区中心，交通便利，是连接西南、西北和华中地区的重要枢纽，区位优势明显。
   - 成都天府国际机场的建成投运，进一步提升了成都的国际化水平和经济辐射能力。

3. **产业结构优化**：
   - 高新技术产业的快速发展，特别是电子信息、生物医药等领域的突破，为经济增长注入新动力。
   - 传统产业的转型升级，提升了产业链的整体竞争力。

4. **人口红利**：
   - 成都市作为人口净流入城市，劳动力资源丰富，消费市场潜力巨大。

### 三、未来趋势预测

1. **持续增长**：
   - 预计未来几年，成都市的GDP增长率将保持在6%-8%之间，继续保持稳健增长态势。

2. **产业结构进一步优化**：
   - 高新技术产业和服务业将继续成为经济增长的主要驱动力，传统产业将进一步转型升级。

3. **区域协同发展**：
   - 成渝地区双城经济圈建设的推进，将为成都市带来更多发展机遇，促进区域经济协同发展。

4. **绿色发展**：
   - 随着环保意识的增强和绿色政策的实施，成都市将更加注重可持续发展，绿色经济将成为新的增长点。

### 结论

总体来看，成都市GDP增长率在过去几年中表现出较强的稳定性和韧性，未来在政策支持、区位优势和产业结构优化的共同作用下，预计将继续保持稳健增长。同时，成都市应继续加强科技创新、优化营商环境、推动绿色发展，以实现更高质量的经济增长。
[END]

可以看到，调用商业模型后，程序只用了两步就完成了报告，并且跳过了收集数据的步骤，而是过了 supervisor 后直接走 writer 节点，然后结束。

当然，虽然glm-4-plus模型的能力确实比较强，但在 writer 节点中，也并未严格按照提示词生成 [END] 标识，另外，glm-4-plus 模型预训练用的数据比较老了，生成的报告只有2022年的。

8 总结

本项目还是比较复杂的，主要是 supervisor 和 create_agent 这两个函数比较复杂，需要比较长的时间才能把逻辑理顺，但只要把这两个函数搞明白了，基本上这个项目就梳理清楚了。

一般情况下，如果存在条件边，让模型去做判断下一个节点走哪里，就需要选择适当的提示词（比如在 supervisor 节点中的提示词，让模型选择下一个节点走哪里），把所有可能的情况都写到提示词里，最后根据模型的回复来确定走哪个节点。因为要保证模型能尽量按指令执行，并且具有较强的选择能力，那么必须在条件边的出点（本项目为 supervisor 节点）使用推理能力比较强的模型，同时在程序中把规则写死，防止模型给出的选择结果不在我们提供的选项范围内，依次做到流程控制。

附录：完整代码

from typing import Literal, Annotated
from langchain_openai import ChatOpenAI
from langgraph.graph import StateGraph, START, END
from langchain_core.messages import HumanMessage, SystemMessage, AIMessage
from pydantic import BaseModel, Field
import re

# ================== 模型配置 ==================
model = ChatOpenAI(
    openai_api_base="http://localhost:23333/v1",
    model_name="/data/coding/models/Qwen/Qwen3-4B",
    openai_api_key="XXXXX",
    max_tokens=400,
    temperature=0.2,
)


# ================== 状态类 ==================
class AgentState(BaseModel):
    # messages 用于接收上一个节点或者起始节点的信息
    messages: Annotated[list, Field(default_factory=list)]  
    # next_agent 用于指定下一个执行节点，下一个节点只能是 researcher、calculator、writer 和 done 其中之一
    next_agent: Literal["researcher", "calculator", "writer", "done"] = "researcher"   
    # step_count 用于记录步数
    step_count: int = 0
    # is_finalized 用于记录是否终止
    is_finalized: bool = False
    search_data: dict = Field(default_factory=dict)  # 新增搜索数据存储

    @classmethod
    def from_raw(cls, raw):
        # 这个函数用于判断 raw 是否为 cls 类对象，这里 cls 是类名，比如 AgentState
        # 如果 raw 属于 AgentState 类，则直接把 raw 返回，否则转成 AgentState 类对象
        """统一状态转换方法"""
        if isinstance(raw, cls):
            return raw
        if isinstance(raw, dict):
            return cls(
                messages=raw.get('messages', []),
                next_agent=raw.get('next_agent', 'researcher'),
                step_count=raw.get('step_count', 0),
                is_finalized=raw.get('is_finalized', False),
                search_data=raw.get('search_data', {})
            )
        raise ValueError(f"无效状态类型: {type(raw)}")  

    def to_safe_dict(self):
        """安全转换为字典"""
        return {
            "messages": [msg.dict() for msg in self.messages],
            "next_agent": self.next_agent,
            "step_count": self.step_count,
            "is_finalized": self.is_finalized,
            "search_data": self.search_data
        }


# ================== 创建Supervisor节点（相当于 START 节点）函数 ==================
def supervisor(raw_state):
    print("================执行 supervisor 节点================")
    state = AgentState.from_raw(raw_state)
    if state.is_finalized:
        return state

    try:
        # 判断步数是否达到循环上限，若达到上限，则强行走 writer 节点
        # 因为 writer 节点执行完之后会添加 [END] 标识，相当于强制结束
        if state.step_count >= 8:   
            print("[强制终止] 达到最大步数限制")
            return AgentState(
                messages=state.messages,
                next_agent="writer",
                step_count=state.step_count + 1,
                is_finalized=False
            )

        # 消息转换，转成 langchain 能处理的格式
        messages = convert_messages(state.messages)

        # 在消息列表中添加流程控制提示词
        control_prompt = '''请严格选择下一步：\n1.需要更多数据 → researcher\n2.需要计算 → calculator \n3.生成报告 → writer\n4. 完成 → done'''
        messages = messages + [SystemMessage(content=control_prompt)]

        # 模型推理
        response = model.invoke(messages)

        # 查看模型的回复中，是否包含 researcher|calculator|writer，若不包含，则 new_agent 置为 done
        decision = re.search(r'\b(researcher|calculator|writer)\b', response.content.lower())
        new_agent = decision.group(1).lower() if decision else "done"

        # 若 new_agent 为 done，但是还没有经过 writer 节点，那只能强制调用 writer 生成报告
        # 如果经过 writer，则模型回复中会带有 [END] 字符串，因此可以据此判断是否经过 writer 节点
        if new_agent == "done" and not any("[END]" in msg.content for msg in messages):
            new_agent = "writer"

        return AgentState(
            messages=state.messages,
            next_agent=new_agent,
            step_count=state.step_count + 1,
            is_finalized=False
        )
    except Exception as e:
        # 如果发生错误，则强制结束
        print(f"[Supervisor错误] {str(e)}")
        return AgentState(
            messages=state.messages,
            next_agent="writer",
            step_count=state.step_count + 1,
            is_finalized=False
        )


# ================== 消息转换 ==================
def convert_messages(raw_messages):
    converted = []
    for msg in raw_messages:
        if isinstance(msg, dict):
            try:
                # 获取当前消息的角色，如果不存在，则置为 assistant
                role = msg.get("role", "assistant")
                # 获取当前消息的内容
                content = msg.get("content", "")
                # 根据角色获取类名
                role_dict = {
                    "user": HumanMessage,
                    "system": SystemMessage,
                    "assistant": AIMessage
                }
                msg_type = role_dict[role]
                # 将内容封装成 langchain 中对应角色的标准格式
                std_msg = msg_type(content=content)
                # 添加到消息列表
                converted.append(std_msg)
            except KeyError:
                converted.append(AIMessage(content=str(msg)))
        elif hasattr(msg, 'content'):
            converted.append(msg)
        else:
            converted.append(AIMessage(content=str(msg)))
    return converted


# ================== 创建普通节点函数 ==================
def create_agent(role: str, prompt: str):
    # 这里 role 不是 AI、HUMAN、SYSTEM 等角色，而是指具体哪个节点
    def agent(raw_state):
        print(f"================执行 {role} 节点================")
        state = AgentState.from_raw(raw_state)
        if state.is_finalized:
            return state

        try:
            # 消息转换
            messages = convert_messages(state.messages)

            # 往messages中添加系统信息，让节点处的模型来选择下一步该走哪个节点
            # 如果当前节点是 writer，那么 prompt 是 生成包含[END]标识的最终报告
            messages.append(SystemMessage(
                content=f"{prompt}\n最后必须用'建议下一步：选项'格式结尾（选项只能是researcher/calculator/writer/done）"
            ))

            # 获取模型响应
            response = model.invoke(messages)
            print(f"\n[{role}输出] {response.content[:500]}...")  # 限制日志长度

            # 如果当前节点是 writer，且模型的回复中不含 [END] 标识，那就强制添加
            # 虽然在提示词中要求 writer 节点的输出以 [END]，但我们使用的是小模型，小模型的指令遵循能力有限
            if role == "writer" and "[END]" not in response.content:
                response.content += "\n[END]"
                print(f"[{role}修正] 已添加[END]标识")

            # 解析下一步建议，模型未必会按照提示词来输出，所以这里可能需要强制匹配
            next_agent = "done"     # 这里先预定义下一个节点为 End
            
            # 这里使用正则匹配，:= 是海象运算符
            if match := re.search(r'建议下一步\s*[:：]\s*(\w+)', response.content):
                # 从match中解析出suggestion
                suggestion = match.group(1).lower()

                # 如果suggestion在集合是 researcher、calculator、writer 其中之一，
                # 那下一个节点就按照 suggestion，否则就用预定义的节点，即 END
                if suggestion in {"researcher", "calculator", "writer", "done"}:
                    next_agent = suggestion

            return AgentState(
                messages=messages + [response],
                next_agent=next_agent,
                step_count=state.step_count + 1,
                is_finalized="[END]" in response.content    # 如果 [END] 在模型答复中，则 is_finalized 为True
            )
        except Exception as e:
            print(f"[{role}错误] {str(e)}")
            return AgentState(
                messages=state.messages,
                next_agent="done",
                step_count=state.step_count + 1,
                is_finalized=True
            )
    return agent


# ================== 执行入口 ==================
def main():
    # ================== 工作流配置（建图） ==================
    builder = StateGraph(AgentState)
    builder.add_node("supervisor", supervisor)  # supervisor 相当于开始节点，
    builder.add_node("researcher", create_agent("researcher", "收集经济数据并分析趋势"))
    builder.add_node("calculator", create_agent("calculator", "执行GDP增长率计算"))
    builder.add_node("writer", create_agent("writer", "生成包含[END]标识的最终报告"))

    # 添加条件边
    # 因为逻辑不是串行的，这里只有发起者是固定的，但边的另一端是哪个节点，需要通过状态来判断
    # 多智能体建图的时候一般都是按照下面这样的方式，即构建一条从一个节点到多个节点的条件边
    builder.add_conditional_edges(
        "supervisor",
        lambda state: state.next_agent,
        {
            "researcher": "researcher",
            "calculator": "calculator", 
            "writer": "writer",
            "done": END
        }
    )

    # 为每一个节点构建一条指向 supervisor 的边
    for agent in ["researcher", "calculator", "writer"]:
        builder.add_edge(agent, "supervisor")

    # 设定入口
    builder.set_entry_point("supervisor")

    # 编译图
    workflow = builder.compile()

    try:
        # 初始化状态时确保类型正确
        initial_state = AgentState(messages=[HumanMessage(content="成都市GDP增长率分析")])
        
        # 流式输出
        for step in workflow.stream(initial_state):
            node_name, raw_state = step.popitem()       # 这里 raw_state 是字典
            current_state = AgentState.from_raw(raw_state)
            
            # print(f"\n[系统状态] 当前节点: {node_name}")
            print(f"下一步: {current_state.next_agent}")
            print(f"步数: {current_state.step_count}")
            print(f"完成状态: {current_state.is_finalized}")

            # 检查终止条件
            if current_state.is_finalized or node_name == "__end__":
                print("\n================流程完成================")
                if any("[END]" in msg.content for msg in current_state.messages):
                    print("最终报告内容:")
                    print(next(msg.content for msg in reversed(current_state.messages) if "[END]" in msg.content))
                else:
                    print("警告：未检测到完整报告")
                break
                
    except Exception as e:
        print(f"\n!!! 流程异常终止: {str(e)}")
        if 'current_state' in locals():
            print("最后状态:", current_state.to_safe_dict())


if __name__ == "__main__":
    main()