Pydantic-AI项目中Agent工具调用与流式响应的结合实践

概述

在使用Pydantic-AI框架构建智能代理应用时，开发者经常会遇到需要将Agent的工具调用功能与流式响应结合的场景。本文将通过一个实际的骰子游戏案例，深入分析如何正确处理Agent执行过程中的工具调用与流式输出的时序问题。

核心问题分析

在Pydantic-AI框架中，Agent执行通常包含以下几个阶段：

1. 接收用户输入
2. 解析意图并决定是否需要调用工具
3. 执行工具函数
4. 生成最终响应

当开发者尝试将这些步骤与FastAPI的流式响应(StreamingResponse)结合时，可能会遇到工具调用内容无法正确返回的问题。这是因为工具调用发生在流式响应开始之后，导致最终响应内容为空。

解决方案详解

1. 使用iter方法替代stream

Pydantic-AI框架提供了更底层的iter方法，可以更精细地控制Agent执行流程。通过这种方法，开发者可以明确区分不同节点的输出，包括工具调用节点(CallToolsNode)和最终响应节点。

async def query(prompt: str, uname: str) -> StreamingResponse:
    async def generate_content():
        async with agent.run_stream(prompt, deps=uname) as result:
            async for node in result.iter():
                if node.type == "CallToolsNode":
                    # 处理工具调用
                    pass
                elif node.type == "ResponseNode":
                    # 处理最终响应
                    async for text in node.stream():
                        yield text
    
    return StreamingResponse(generate_content())

2. 优化工具调用与响应时序

在实际应用中，我们需要确保工具调用的结果能够被正确捕获并包含在流式响应中。可以通过以下方式实现：

1. 预先收集工具调用结果
2. 将工具结果与Agent响应合并
3. 通过统一的流式接口输出

3. 错误处理与日志记录

在流式响应中加入适当的错误处理机制和日志记录，可以帮助开发者更好地调试和优化Agent行为：

async def generate_content():
    try:
        async with agent.run_stream(prompt, deps=uname) as result:
            tool_results = []
            async for node in result.iter():
                if node.type == "CallToolsNode":
                    tool_results.append(await node.execute())
                elif node.type == "ResponseNode":
                    # 先输出工具结果
                    for res in tool_results:
                        yield f"工具结果: {res}\n"
                    # 再输出Agent响应
                    async for text in node.stream():
                        yield text
    except Exception as e:
        yield f"处理过程中发生错误: {str(e)}"

实际应用案例

以骰子游戏为例，完整的解决方案应该包含以下要素：

1. 系统提示定义游戏规则
2. 工具函数实现骰子投掷和玩家名称获取
3. 流式响应处理所有中间步骤和最终结果

system_prompt = """\
你是一个骰子游戏，需要掷骰子并判断结果是否与用户猜测一致。
如果一致，告诉玩家他们赢了。在响应中使用玩家名称。
"""

@app.post("/query")
async def query(prompt: str, uname: str) -> StreamingResponse:
    async def game_stream():
        async with agent.run_stream(prompt, deps=uname) as result:
            # 先处理工具调用
            dice_result = None
            async for node in result.iter():
                if node.type == "CallToolsNode":
                    if node.func.__name__ == "roll_die":
                        dice_result = await node.execute()
                        yield f"骰子结果: {dice_result}\n"
            
            # 再处理最终响应
            async for text in result.stream():
                yield text
    
    return StreamingResponse(game_stream())

性能优化建议

1. 去抖动优化：适当设置debounce_by参数，平衡响应速度和性能
2. 上下文管理：确保RunContext正确传递依赖项
3. 并发处理：对于多个工具调用，考虑并行执行提高效率
4. 缓存机制：对频繁使用的工具结果实施缓存

总结

在Pydantic-AI框架中实现Agent工具调用与流式响应的完美结合，关键在于理解框架的执行流程和节点类型。通过使用iter方法替代简单的stream方法，开发者可以获得更细粒度的控制能力，确保工具调用结果能够被正确捕获并包含在最终响应中。本文提供的解决方案和优化建议，可以帮助开发者构建更健壮、响应更及时的AI应用。