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. 理解用户输入
2. 选择合适的工具
3. 执行工具调用
4. 生成最终响应

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)

结语

通过这个完整的工具调用Agent实现,我们可以看到LangGraph在处理复杂LLM应用时的强大能力:

1. 灵活的流程控制: 通过条件边实现复杂的决策逻辑
2. 状态管理: 清晰的状态定义和转换
3. 错误处理: 完善的错误处理和重试机制
4. 可扩展性: 易于添加新的工具和功能

LangGraph不仅解决了传统方法的局限性,还提供了一个更加强大和灵活的框架来构建复杂的LLM应用。通过合理使用这些高级特性,我们可以构建出更加智能和可靠的AI应用。