LangGraph之结合结构化输出构建Router Agent（数据库）

在构建Router Agent的过程中，我们旨在实现一个智能系统，它能够根据用户的输入，精准判断是执行数据库操作还是直接生成模型回复。这一过程涉及多个关键步骤，下面将详细阐述。

一、定义Pydantic模型类

使用Pydantic库定义模型类，用于结构化输出和数据验证。

1. UserInfo模型类对应数据库中用户信息表的结构，包括姓名、年龄、邮箱和电话等字段，方便对用户相关信息进行结构化处理和存储。
2. ConversationalResponse模型类用于处理模型生成的自然语言回复，它只有一个response字段，用于存储模型针对用户问题生成的自然语言回答内容。
3. FinalResponse模型类则通过Union类型将UserInfo和ConversationalResponse包含进来，这样在使用大模型实例的with_structured_output方法时，就可以灵活选择输出类型，实现根据不同业务逻辑返回不同格式数据的功能。

from typing import Union, Optional
from pydantic import BaseModel, Field

# 定义数据库插入的用户信息模型
class UserInfo(BaseModel):
    name: str = Field(description="The name of the user")
    age: Optional[int] = Field(description="The age of the user")
    email: str = Field(description="The email address of the user")
    phone: Optional[str] = Field(description="The phone number of the user")

# 定义正常生成模型回复的模型
class ConversationalResponse(BaseModel):
    response: str = Field(description="A conversational response to the user's query")

# 定义最终响应模型，包含UserInfo或ConversationalResponse
class FinalResponse(BaseModel):
    final_output: Union[UserInfo, ConversationalResponse]

二、配置数据库连接并建表

借助sqlalchemy库来配置与MySQL数据库的连接并创建相应的表结构。

1. 创建一个基类Base，它是所有数据库模型类的基类，为定义具体的表结构提供基础。
2. 定义User类，继承自Base，对应数据库中的users表，明确表中的字段及其数据类型，如id为主键，name、age、email和phone分别对应相应的用户信息字段。
3. 设置数据库连接URI，包含数据库用户名、密码、服务器IP地址、数据库名称和字符集等信息，根据实际情况进行替换。通过create_engine创建数据库引擎，使用Base.metadata.create_all(engine)检查并创建表结构，如果表不存在则创建。
4. 通过sessionmaker创建会话类，并实例化一个会话对象，用于后续与数据库的交互操作。

from sqlalchemy import create_engine, Table, Column, Integer, String, MetaData
from sqlalchemy.orm import declarative_base, sessionmaker

# 创建基类
Base = declarative_base()

# 定义UserInfo模型对应的表结构
class User(Base):
    __tablename__ = 'users'
    id = Column(Integer, primary_key=True)
    name = Column(String(50))
    age = Column(Integer)
    email = Column(String(100))
    phone = Column(String(15))

# 数据库连接URI，需替换为实际信息
DATABASE_URI ='mysql+pymysql://root:snowba11950123@192.168.110.131/langgraph_agent?charset=utf8mb4'
engine = create_engine(DATABASE_URI, echo=True)

# 如果表不存在，则创建表
Base.metadata.create_all(engine)

# 创建会话
Session = sessionmaker(bind=engine)
session = Session()

三、定义节点函数

节点函数是Router Agent的核心执行单元，负责具体的任务处理。

1. chat_with_model函数作为路由节点函数，接收全局共享状态state，从state中获取用户输入的消息列表messages。使用大模型实例11m的with_structured_output方法，并传入FinalResponse模型类，将用户输入的自然语言文本转化为结构化输出。最后将包含结构化输出结果的消息列表返回，以便后续处理。
2. final_answer函数用于处理直接生成自然语言回复的情况。它从共享状态state中获取最新的消息（即最后一条消息），通过结构化解析获取消息中的final_output中的response字段，这个字段就是模型生成的自然语言回复内容。将该回复内容作为消息列表返回，作为最终传递给用户的响应。
3. insert_db函数用于执行数据库插入操作。同样从共享状态state中获取最新消息，提取消息中的final_output，根据UserInfo模型类的结构，通过访问output的各个字段（如name、age、email和phone）获取具体的用户信息值。使用这些值创建User表的实例，并将其添加到数据库会话中，提交事务以完成数据插入操作。如果插入过程中出现异常，则回滚事务并返回错误信息；如果成功，则返回数据已成功存储的消息。最后关闭会话，确保资源正确释放。

from langchain_openai import ChatOpenAI
from langchain_core.messages import AnyMessage, SystemMessage, HumanMessage, ToolMessage

# 生成模型实例
11m = ChatOpenAI(model="gpt-40-mini")

def chat_with_model(state):
    """将用户输入文本转化为结构化输出"""
    messages = state['messages']
    structured_11m = 11m.with_structured_output(FinalResponse)
    response = structured_11m.invoke(messages)
    return {"messages": [response]}

def final_answer(state):
    """生成自然语言回复"""
    messages = state['messages'][-1]
    response = messages.final_output.response
    return {"messages": [response]}

def insert_db(state):
    """执行数据库插入操作"""
    session = Session()
    try:
        result = state['messages'][-1]
        output = result.final_output
        user = User(name=output.name, age=output.age, email=output.email, phone=output.phone)
        session.add(user)
        session.commit()
        return {"messages": ["数据已成功存储至Mysql数据库。"]}
    except Exception as e:
        session.rollback()
        return {"messages": [f"数据存储失败,错误原因:{e}"]}
    finally:
        session.close()

四、定义路由函数并构建图

路由函数generate_branch用于根据输出类型决定路由分支走向。

1. 它接收包含消息列表的状态state，从列表中获取最后一条消息（即最新消息），提取消息中的final_output。通过isinstance函数判断output的类型，如果是UserInfo类型，则返回True，表示应执行数据库插入操作；如果是ConversationalResponse类型，则返回False，表示应直接生成自然语言回复。
2. 构建状态图时，首先导入相关的库和类，定义AgentState类型，它是一个TypedDict，用于描述状态中消息的结构，其中messages是一个列表，通过annotated和operator.add确保消息列表能够正确追加新消息。
3. 使用StateGraph类创建状态图对象，添加三个关键节点：chat_with_model作为路由节点，负责处理用户输入并转化为结构化输出；final_answer用于生成自然语言回复；insert_db用于执行数据库插入操作。
4. 设置chat_with_model为启动节点，接收用户输入。通过add_conditional_edges方法设置条件边，根据generate_branch函数的返回值决定路由走向。如果返回True，则连接到insert_db节点；如果返回False，则连接到final_answer节点。
5. 最后将final_answer和insert_db设置为终止节点，并编译状态图，使其可以运行。

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

class AgentState(TypedDict):
    messages: Annotated[list[AnyMessage], operator.add]

def generate_branch(state: AgentState):
    """根据输出类型决定路由分支"""
    result = state['messages'][-1]
    output = result.final_output
    if isinstance(output, UserInfo):
        return True
    elif isinstance(output, ConversationalResponse):
        return False

graph = StateGraph(AgentState)
graph.add_node("chat_with_model", chat_with_model)
graph.add_node("final_answer", final_answer)
graph.add_node("insert_db", insert_db)
graph.set_entry_point("chat_with_model")
graph.add_conditional_edges(
    "chat_with_model",
    generate_branch,
    {True: "insert_db", False: "final_answer"}
)
graph.set_finish_point("final_answer")
graph.set_finish_point("insert_db")
graph = graph.compile()

五、测试Router Agent

为了验证Router Agent的功能，进行了两组测试。

1. 第一组测试插入数据库分支，构造一个包含用户详细信息的查询query1，将其包装成HumanMessage并放入input_message1字典的messages列表中。使用状态图的invoke方法执行查询，查看输出结果。通过观察输出，确认是否正确执行了数据库插入操作，包括检查数据库中是否成功插入数据，以及返回的消息是否表明插入成功。
2. 第二组测试直接生成回复分支，构造一个询问介绍的查询query2，同样包装成HumanMessage放入input_message2中。调用invoke方法执行查询，检查输出结果是否为模型生成的自然语言回复，同时确认数据库中没有插入数据，以此验证直接生成回复功能的正确性。

from langchain_core.messages import HumanMessage

# 测试插入数据库分支
query1 = "我叫木羽,今年28岁,邮箱地址是snow@gmial.com,电话是1323521313"
input_message1 = {"messages": [HumanMessage(content=query1)]}
result1 = graph.invoke(input_message1)
print(result1)

# 测试直接生成回复分支
query2 = "你好,请你介绍一下你自己"
input_message2 = {"messages": [HumanMessage(content=query2)]}
result2 = graph.invoke(input_message2)
print(result2)

最终结果：

通过上述步骤，我们成功构建并测试了Router Agent，它能够根据用户输入准确地选择执行数据库操作或直接生成模型回复，展示了其在处理特定业务逻辑时的高效性和精准性。

六、Router Agent 总结