LangGraph进阶：条件边与工具调用Agent实现

在前两篇文章中,我们讨论了LCEL和AgentExecutor的局限性,以及LangGraph的基础概念。今天,我们将深入探讨LangGraph的高级特性,重点关注条件边的使用和如何实现一个完整的工具调用Agent。

条件边的高级用法

条件边是LangGraph中最强大的特性之一,它允许我们基于状态动态决定执行流程。让我们深入了解一些高级用法。

1. 多条件路由

from typing import List, Dict, Literal
from pydantic import BaseModel
from langgraph.graph import StateGraph, END

class AgentState(BaseModel):
    messages: List[Dict[str, str]] = []
    current_input: str = ""
    tools_output: Dict[str, str] = {}
    status: str = "RUNNING"
    error_count: int = 0

def route_by_status(state: AgentState) -> Literal["process", "retry", "error", "end"]:
    """复杂的路由逻辑"""
    if state.status == "SUCCESS":
        return "end"
    elif state.status == "ERROR":
        if state.error_count >= 3:
            return "error"
        return "retry"
    elif state.status == "NEED_TOOL":
        return "process"
    return "process"

# 构建图结构
workflow = StateGraph(AgentState)

# 添加条件边
workflow.add_conditional_edges(
    "check_status",
    route_by_status,
    {
        "process": "execute_tool",
        "retry": "retry_handler",
        "error": "error_handler",
        "end": END
    }
)

2. 并行执行

LangGraph支持并行执行多个节点,这在处理复杂任务时特别有用:

async def parallel_tools_execution(state: AgentState) -> AgentState:
    """并行执行多个工具"""
    tools = identify_required_tools(state.current_input)
    
    async def execute_tool(tool):
        result = await tool.ainvoke(state.current_input)
        return {tool.name: result}
    
    # 并行执行所有工具
    results = await asyncio.gather(*[execute_tool(tool) for tool in tools])
    
    # 合并结果
    tools_output = {}
    for result in results:
        tools_output.update(result)
    
    return AgentState(
        messages=state.messages,
        current_input=state.current_input,
        tools_output=tools_output,
        status="SUCCESS"
    )

实现完整的工具调用Agent

让我们通过实现一个完整的工具调用Agent来展示LangGraph的强大功能。这个Agent能够:

理解用户输入
选择合适的工具
执行工具调用
生成最终响应

1. 定义状态和工具

from typing import List, Dict, Optional
from pydantic import BaseModel
from langchain.tools import BaseTool
from langchain.tools.calculator import CalculatorTool
from langchain.tools.wikipedia import WikipediaQueryRun
from langchain_core.language_models import ChatOpenAI

class Tool(BaseModel):
    name: str
    description: str
    func: callable

class AgentState(BaseModel):
    messages: List[Dict[str, str]] = []
    current_input: str = ""
    thought: str = ""
    selected_tool: Optional[str] = None
    tool_input: str = ""
    tool_output: str = ""
    final_answer: str = ""
    status: str = "STARTING"

# 定义可用工具
tools = [
    Tool(
        name="calculator",
        description="用于执行数学计算",
        func=CalculatorTool()
    ),
    Tool(
        name="wikipedia",
        description="用于查询维基百科信息",
        func=WikipediaQueryRun()
    )
]

2. 实现核心节点

async def think(state: AgentState) -> AgentState:
    """思考下一步行动"""
    prompt = f"""
    基于用户输入和当前对话历史，思考下一步行动。
    用户输入: {state.current_input}
    可用工具: {[t.name + ': ' + t.description for t in tools]}
    
    请决定:
    1. 是否需要使用工具
    2. 如果需要，选择哪个工具
    3. 使用什么参数调用工具
    
    以JSON格式返回: {{"thought": "思考过程", "need_tool": true/false, "tool": "工具名", "tool_input": "参数"}}
    """
    
    llm = ChatOpenAI(temperature=0)
    response = await llm.ainvoke(prompt)
    result = json.loads(response)
    
    return AgentState(
        **state.dict(),
        thought=result["thought"],
        selected_tool=result.get("tool"),
        tool_input=result.get("tool_input"),
        status="NEED_TOOL" if result["need_tool"] else "GENERATE_RESPONSE"
    )

async def execute_tool(state: AgentState) -> AgentState:
    """执行工具调用"""
    tool = next((t for t in tools if t.name == state.selected_tool), None)
    if not tool:
        return AgentState(
            **state.dict(),
            status="ERROR",
            thought="Selected tool not found"
        )
    
    try:
        result = await tool.func.ainvoke(state.tool_input)
        return AgentState(
            **state.dict(),
            tool_output=str(result),
            status="GENERATE_RESPONSE"
        )
    except Exception as e:
        return AgentState(
            **state.dict(),
            status="ERROR",
            thought=f"Tool execution failed: {str(e)}"
        )

async def generate_response(state: AgentState) -> AgentState:
    """生成最终响应"""
    prompt = f"""
    基于以下信息生成对用户的回复:
    用户输入: {state.current_input}
    思考过程: {state.thought}
    工具输出: {state.tool_output}
    
    请生成一个清晰、有帮助的回复。
    """
    
    llm = ChatOpenAI(temperature=0.7)
    response = await llm.ainvoke(prompt)
    
    return AgentState(
        **state.dict(),
        final_answer=response,
        status="SUCCESS"
    )

3. 构建完整的工作流

# 创建图结构
workflow = StateGraph(AgentState)

# 添加节点
workflow.add_node("think", think)
workflow.add_node("execute_tool", execute_tool)
workflow.add_node("generate_response", generate_response)

# 定义路由函数
def route_next_step(state: AgentState) -> str:
    if state.status == "ERROR":
        return "error"
    elif state.status == "NEED_TOOL":
        return "execute_tool"
    elif state.status == "GENERATE_RESPONSE":
        return "generate_response"
    elif state.status == "SUCCESS":
        return "end"
    return "think"

# 添加条件边
workflow.add_conditional_edges(
    "think",
    route_next_step,
    {
        "execute_tool": "execute_tool",
        "generate_response": "generate_response",
        "error": "error_handler",
        "end": END
    }
)

workflow.add_conditional_edges(
    "execute_tool",
    route_next_step,
    {
        "generate_response": "generate_response",
        "error": "error_handler",
        "end": END
    }
)

workflow.add_edge("generate_response", END)

# 编译图
app = workflow.compile()

4. 使用示例

async def run_agent(user_input: str):
    state = AgentState(current_input=user_input)
    final_state = await app.ainvoke(state)
    return final_state.final_answer

# 使用示例
async def main():
    questions = [
        "计算23乘以45等于多少？",
        "谁发明了相对论？",
        "计算圆周率乘以10等于多少？"
    ]
    
    for question in questions:
        print(f"\n问题: {question}")
        answer = await run_agent(question)
        print(f"回答: {answer}")

# 运行示例
import asyncio
asyncio.run(main())

高级特性和技巧

1. 状态持久化

LangGraph支持将状态持久化到数据库中:

from langgraph.persist import GraphStatePersist

class DBStatePersist(GraphStatePersist):
    async def persist(self, state: AgentState):
        # 实现状态持久化逻辑
        pass
    
    async def load(self) -> AgentState:
        # 实现状态加载逻辑
        pass

# 使用持久化
workflow = StateGraph(AgentState, persist_handler=DBStatePersist())

2. 流式输出

支持实时显示执行过程:

async def stream_handler(state: AgentState):
    """处理流式输出"""
    yield f"思考中: {state.thought}\n"
    if state.selected_tool:
        yield f"使用工具: {state.selected_tool}\n"
    if state.tool_output:
        yield f"工具输出: {state.tool_output}\n"
    yield f"最终答案: {state.final_answer}\n"

# 在图中启用流式输出
workflow.set_stream_handler(stream_handler)

3. 错误处理和重试机制

async def error_handler(state: AgentState) -> AgentState:
    """处理错误情况"""
    if state.error_count < 3:
        return AgentState(
            **state.dict(),
            error_count=state.error_count + 1,
            status="RETRY"
        )
    return AgentState(
        **state.dict(),
        final_answer="抱歉，我无法完成这个任务。",
        status="ERROR"
    )

# 添加错误处理节点
workflow.add_node("error_handler", error_handler)