Pydantic-AI中Anthropic流式处理与AWS Bedrock集成的问题分析

背景介绍

Pydantic-AI是一个基于Pydantic的AI代理框架，它提供了与多种AI模型集成的能力。在实际应用中，开发者经常需要处理流式响应，特别是在与AWS Bedrock上的Anthropic模型交互时。本文将深入分析在使用Pydantic-AI框架时遇到的流式处理问题及其解决方案。

问题现象

当开发者使用Pydantic-AI的run_stream方法与Anthropic模型交互时，发现流式处理仅返回了代理的第一个消息（通常是解释将要执行工具调用的内容），而不是预期的最终结果。相比之下，同步的run方法则能正确返回工具调用后的最终结果。

问题复现

通过一个简单的示例可以清晰复现这个问题。创建一个带有随机数生成工具的代理，分别使用run_stream和run方法调用：

agent = Agent(
    AnthropicModel(model_name="anthropic.claude-3-5-sonnet-20241022-v2:0", 
                  anthropic_client=AsyncAnthropicBedrock(aws_region="us-west-2")),
    system_prompt="Return a random number using your tools",
    tools=[get_random_number],
)

# 流式调用
async with agent.run_stream(user_prompt="生成随机数") as result:
    async for token in result.stream():
        print(token)

# 同步调用
sync_result = await agent.run(user_prompt="生成随机数")
print(sync_result.data)

流式调用仅输出类似"I'll help you generate a random number using the get_random_number function"的内容，而同步调用则能正确返回生成的随机数及总结。

问题分析

深入分析消息流发现，run_stream方法在处理过程中存在以下问题：

1. 流式处理在第一个工具调用后提前终止
2. 最终模型响应未被正确处理
3. 开发者无法通过常规方式获取完整的消息历史

相比之下，同步的run方法能够正确处理完整的交互流程：初始请求→工具调用→工具返回→最终响应。

解决方案

Pydantic-AI框架提供了iter()方法作为替代方案，能够正确处理完整的交互流程。开发者可以通过以下方式实现只流式传输最终消息的需求：

with agent.iter(**input_args) as agent_run:
    async for node in agent_run:
        if Agent.is_model_request_node(node):
            async with node.stream(agent_run.ctx) as request_stream:
                async for event in request_stream:
                    if isinstance(event, PartStartEvent):
                        if isinstance(event.part, TextPart) and "\n\nFinal" in event.part.content:
                            is_final = True
                            await output_msg.stream_token(event.part.content.replace("\nFinal",""))
                    if isinstance(event, PartDeltaEvent) and is_final and isinstance(event.delta, TextPartDelta):
                        await output_msg.stream_token(event.delta.content_delta)

这种方法的关键点在于：

1. 使用iter()方法遍历整个代理执行过程
2. 识别并只处理标记为最终响应的消息部分
3. 通过事件流处理机制实现精细控制

最佳实践建议

1. 对于需要完整控制流的场景，优先使用iter()方法而非run_stream
2. 可以通过在系统提示中添加特殊标记（如"\n\nFinal"）来帮助识别最终响应
3. 考虑实现自定义的结果过滤器，根据业务需求筛选需要流式传输的内容
4. 对于简单场景，同步run方法可能更为可靠

总结

Pydantic-AI框架在与Anthropic模型和AWS Bedrock集成时，流式处理功能存在一定的局限性。开发者可以通过使用iter()方法替代run_stream来实现更精细的控制。理解框架内部的消息处理机制对于构建可靠的AI应用至关重要。随着框架的持续发展，这些问题有望在未来的版本中得到更好的解决。