Pi Agent源码解析:执行循环、工具注册与上下文管理
先说明一下这篇笔记是系列第二篇。第一篇我把 Pi Agent 的工程结构、模块划分和整体调用链捋了一遍包括agent、tool、memory、context、orchestrator这几个核心包各自负责什么。当时留了个尾巴说要把核心实现细节逐个过一遍这篇就来填坑。如果你是第一次接触这个框架建议先把第一篇扫一眼或者是直接去 GitHub 把仓库 clone 下来边看这篇笔记边对着pi_agent的源码目录翻。没有源码在手光看笔记会丢失很多细节尤其是我下面要讲的这些“源码里藏着但文档里不写”的坑。这篇笔记适合两种人一种是正在读 Pi Agent 源码但有点摸不着头脑的另一种是想自己写一个轻量 Agent 框架、需要参考设计思路的。不管哪种建议你手里开着源码编辑器遇到代码段就写个小 demo 跑一下效果比光看强得多。1. 核心抽象Agent 类与执行循环先看最核心的东西——pi_agent/agent/base.py里的Agent类。这个类承担的职责是整个框架最关键的接收任务、调度工具、对话模型、整理结果。我拉源码时首先盯的就是它因为只要看懂Agent.run()方法整个框架的执行逻辑就通了七八成。1.1 Agent 基类设计Agent类的__init__大概是这样的class Agent: def __init__( self, name: str, system_prompt: str, model: BaseChatModel, tools: Optional[List[ToolEntry]] None, memory: Optional[BaseMemory] None, max_iterations: int 10, temperature: float 0.7, ): self.name name self.system_prompt system_prompt self.model model self.tools {tool.name: tool for tool in (tools or [])} self.memory memory or ConversationBufferMemory() self.max_iterations max_iterations self.temperature temperature这里有几个值得注意的设计决策。第一个是tools用了字典而不是列表key 是工具名。这个看着不起眼但实际是给后文“模型返回工具名 → 查表调用”的打底——模型返回的tool_call.name上来就是字符串字典查询 O(1)而且天然防止同名工具注册。第二个是传入的是BaseChatModel而不是直接把某个具体模型类写死也就是说你可以在不破坏框架逻辑的前提下把模型从 GPT 换成本地部署的 Llama。这是老生常谈的依赖倒置但源码里真正执行得彻底的框架不多Pi Agent 算一个。再看system_prompt。这个参数被包装过了源码里_build_system_prompt()方法会在你传入的 system prompt 后面自动追加一段工具说明你是一个智能助手可以使用以下工具完成任务 - name: 工具名 description: 工具描述 parameters: {参数名: 参数说明}这段自动生成的提示文本不是可选的是每次都拼的。为什么因为 Agent 本身不知道你有没有替它把工具信息写进 system prompt为了保证工具在模型端“可见”框架在这一步做了一次保险。这点我觉得做得聪明调用方只需要声明工具不需要操心模型 API 层和 prompt 模板层的拼接框架把这些脏活包了。1.2 执行循环Run Loop源码逐行拆解run()方法是理解 Agent 的钥匙。我直接把简化后的核心逻辑贴出来然后逐段解释def run(self, task: str) - str: messages self.memory.get_messages() if not messages: messages [{role: system, content: self._build_system_prompt()}] messages.append({role: user, content: task}) for step in range(self.max_iterations): response self.model.chat(messages, toolsself._tool_schemas()) if response.tool_calls: messages.append(self._to_assistant_message(response)) for call in response.tool_calls: tool_result self._execute_tool(call) messages.append( { role: tool, tool_call_id: call.id, content: tool_result, } ) continue final_answer response.content self.memory.add_user_message(task) self.memory.add_ai_message(final_answer) return final_answer raise MaxIterationsError( fAgent {self.name} 超过最大迭代次数 {self.max_iterations} )先说整体流程用户 task 进来先塞进消息队列然后进入循环。循环里先调模型看返回值里有没有工具调用。有工具调用就执行工具把结果回填到 messages然后continue再调模型没有工具调用说明模型认为任务可以总结了于是返回答案并写入 memory。循环上限max_iterations兜底防止模型一直要求调工具不停下来。看到这里你可能会问为什么不把工具执行结果当作“最终答案”直接返回给用户而是回填后再让模型总结一次这个问题我在社区里见过不少人问。答案在于工具结果往往是结构化数据JSON、状态码、日志片段它们不是自然语言答案。比如一个搜索工具返回了 50 条 URL你得让模型对这些结果做二次加工——挑重点、排序、组织描述——才能形成用户能看懂的答案。所以“工具执行 ! 任务结束”必须再走一轮模型调用。另外一个细节是messages的复用方式它没有在每轮循环里重建列表而是直接 append。这是一个性能优化也是一个必须严格注意的点。如果你在循环体内部错误地重建了 messages那么前一轮工具调用的结果就丢了极容易造成循环死转。源码里没有犯这个错但你要是二次开发时改动了这段千万留神。关于MaxIterationsError我看到有使用者吐槽这个异常不够细致。确实只有“超过最大迭代次数”一个维度看不出来到底是工具反复报错还是模型死活不确定。我的排查经验是遇到这个异常先去查日志里最近几轮工具调用的返回内容如果都是异常信息那大概率是工具实现有 bug别急着调大max_iterations。2. 工具注册与调用链从装饰器到 JSON Schema工具系统是 Agent 框架里使用频率最高的子系统。Pi Agent 的工具注册走的是装饰器风格我一开始觉得只是语法糖深入研究后发现这里的“糖”是包了实际逻辑的。2.1 工具协议与注册机制先看一个最简单的工具声明agent.tool(description计算两个整数的和) def add(a: int, b: int) - int: return a bagent.tool()装饰器背后做了什么源码位于pi_agent/tool/registry.py简化后大概是def tool(self, description: str ): def decorator(func): schema generate_schema(func, description) entry ToolEntry( namefunc.__name__, funcfunc, schemaschema, ) self.tools[func.__name__] entry return func return decorator核心动作有两个生成 JSON Schema 和注册。generate_schema这个函数有意思它是靠函数签名和类型注解反推参数结构的。也就是说你在函数里写a: int它就能推断出a的类型是integer写b: str就能推断出是string。源码里用的方式是inspect.signature 类型注解 默认值推断def generate_schema(func, description): sig inspect.signature(func) properties {} required [] for name, param in sig.parameters.items(): properties[name] {type: map_python_type(param.annotation)} if param.default is inspect.Parameter.empty: required.append(name) return { type: function, function: { name: func.__name__, description: description, parameters: { type: object, properties: properties, required: required, }, }, }这套自动推断方案让开发者省去了手写 JSON Schema 的麻烦。但是有个坑必须提如果函数参数没有类型注解map_python_type会落到一个默认的string分支上然后模型端就会“自由发挥”可能传一个和预期完全不符的类型过来。我踩过这个坑建议你的工具函数务必写全类型注解。工具入口的__name__作为注册 key 这件事也要注意。如果同一个 Agent 下挂了两个同名的函数比如不同模块里都写了get_data后面注册的会静默覆盖前面注册的。源码里没有做冲突检查这是个设计上的小缺陷二次开发时可以加个日志。2.2 工具调用的参数解析与结果回填模型端发起工具调用的返回结果里参数是一个 JSON 字符串。_execute_tool方法要做的就是从 JSON 字符串解析出参数再传给真正的函数def _execute_tool(self, call): tool_entry self.tools.get(call.name) if not tool_entry: return f错误: 工具 {call.name} 不存在 try: args json.loads(call.arguments) except json.JSONDecodeError: # 模型偶尔会返回截断的 JSON return 错误: 工具参数不是合法 JSON try: result tool_entry.func(**args) except TypeError as e: return f参数错误: {e} return result这一段里藏着好几个实战要点。第一call.arguments是字符串必须先json.loads但模型返回的 JSON 偶尔是残缺的特别是长参数列表时容易截断所以JSONDecodeError要接住。第二tool_entry.func(**args)这行会把 Python 的TypeError引出来如果模型给的参数名不匹配整个 Agent 循环就会断掉。源码的处理是把错误信息回传给模型让它自己纠正这么做的好处是至少 Agent 不会崩溃。你要注意_execute_tool的返回值不是直接给用户看的。它会被塞进 messages 里当作role: tool的消息然后模型会根据这个内容决定下一步。所以在设计工具返回时尽量返回结构化、可被模型“二次引用”的文本比如 JSON 序列化后的数据而不是一大段渲染好的 HTML。工具可以返回 JSON 字符串但要注意如果返回过大会占用上下文窗口所以能在工具内部做好数据精简就尽量精简比如只返回 top-10 而不是所有搜索结果。3. 上下文管理与记忆机制窗口不够用的日子写 Agent 框架绕不开一个现实大模型的上下文窗口是有限资源。Pi Agent 对这个问题的处理方案分成了两个层面单轮内的上下文组装和跨轮次的记忆持久化。这两个层面的源码分别在context和memory包里。3.1 上下文组装策略继续回到run()方法中的 messages。你可能会想messages 是不是把历史全部倒给模型不是。Pi Agent 在context.py里有个build_context函数它做的最重要一件事是分层系统指令层system prompt 工具描述每轮强制保留任务状态层当前的任务描述、中间变量每轮更新动态记忆层历史对话/过往工具结果根据窗口余量动态取舍每轮循环进入模型调用前build_context会先从 memory 里取历史消息再拼上当前轮的 messages最后做一个 token 估算。源码里 token 估算用的是tiktoken如果你换成了别的模型就得自己接入对应的 tokenizer否则估算会偏差较大。这个分层的思路是源码中最值得学习的设计之一。很多新手在做 Agent 时喜欢把所有内容一股脑塞进去结果系统提示被历史对话挤到下半场模型的“系统人格”就崩了。Pi Agent 用“分词预算”的方式避免了这个问题系统指令永远占第一个位置并且优先级最高。3.2 记忆压缩与裁剪算法跨轮次记忆是 Agent 能不能“记得住事情”的关键。Pi Agent 的默认实现是ConversationBufferMemory就一个列表存储历史消息。但真正有意思的是memory/compressor.py里的compress_messages函数def compress_messages(messages, max_tokens6000, summary_modelNone): total_tokens count_tokens(messages) if total_tokens max_tokens: return messages # 先裁剪大体积的工具结果 target_budget max_tokens * 0.8 for msg in messages: if msg[role] tool and len(msg[content]) 500: msg[content] msg[content][:500] ...(已截断) total_tokens count_tokens(messages) # 再压缩最旧的对话为摘要 while total_tokens target_budget: user_msg, ai_msg messages.pop(0), messages.pop(0) summary summary_model.summarize(f{user_msg}: {ai_msg}) messages.insert(0, {role: system, content: f历史摘要: {summary}}) total_tokens count_tokens(messages) return messages这里有两个反直觉但非常实用的细节。首先是裁剪优先级工具结果优先被压缩而不是优先删除最早的消息。原因是工具结果通常是一条巨长的 JSON而里面大部分内容对于当前决策是冗余的——模型只需要知道“这个工具调用成功并返回了 N 条结果”就行了。相比于把几轮前的用户消息删掉折掉工具结果对信息损失更小。其次是“摘要替换”的操作把最早的一组 user/assistant 消息压缩成一条 system 类型的摘要。这一手是保存长期记忆的常见策略但位置不要搞错——把摘要插到列表头部不要插到尾部否则模型会把它当作当前的指令来理解导致行为异常。compress_messages没在每次run()里自动调用而是暴露出来给 memory 对象在存储的时候自行选择是否调用。默认的ConversationBufferMemory不调用它只有当你显式使用SummaryMemory时才会触发。老实说这个设计我一开始觉得有点折腾但后来想通了压缩是有代价的每一次摘要都意味着细节丢失框架把选择权交给了使用者而不是一刀切。另外Pi Agent 的 memory 还有个window_size参数用来控制单轮 messages 里保留多少条历史。说句实在话window_size按“条数”切而不是按“token 数”切精准度一般但在很多本地部署的场景里够用了。如果你要做生产级记忆管理建议在此基础上再叠加一层更细粒度的 token 预估否则容易被对话轮次很多但每轮都挺短的场景欺骗。4. 多 Agent 协作与任务编排Orchestrator 的实现细节单 Agent 能做简单任务但稍微复杂一点的场景——比如“生成代码然后跑测试再修订”这种闭环——就需要多个角色协同。Pi Agent 在这块提供了Orchestrator我在源码里看它的实现时觉得挺有启发。4.1 角色定义与通信机制多 Agent 不是把一堆 Agent 类塞进一个列表就行。Pi Agent 在orchestrator.py中定义了一个AgentMessage数据结构dataclass class AgentMessage: sender: str receiver: str content: str metadata: dict field(default_factorydict) created_at: float field(default_factorytime.time)为什么要引入一个消息对象而不是直接函数调用两个原因。第一是解耦消息只通过“发件人/收件人”传递每个 Agent 不知道对方内部怎么实现的第二是可观测所有 Agent 之间交换的信息都会经过 Orchestrator 的消息队列你可以把完整通信链路打到日志里哪里出了问题一目了然。我在实际调试中全靠这个日志否则多 Agent 互相调用时你根本不知道谁在给谁传空消息。Orchestrator 的核心调度可以简化成def run(self, task: str) - str: first_agent self.schedule_agent(None, task) result self.agents[first_agent].run(task) for _ in range(self.max_turns): if not self.messages_out: break while self.messages_out: msg self.messages_out.popleft() receiver self.agents[msg.receiver] result receiver.run(msg.content) self.results[msg.receiver] result return result注意schedule_agent的选路逻辑。默认实现是硬编码第一个任务总是发给列表里的第一个 Agent后续消息根据receiver路由。如果你想做更动态的调度比如根据任务类型分发给不同 Agent可以覆写schedule_agent。这个类是源码里少数几个有意“留白”给开发者扩展的地方其他大部分模块都偏“开箱即用”。4.2 任务分配与结果汇总的实现细节让我用一个具体例子来说明多 Agent 编排的执行过程。假设你有一个PlannerAgent、一个CoderAgent、一个TesterAgent任务要求“写一个计算阶乘的函数并验证它”。Orchestrator 的流程是把“写一个计算阶乘的函数并验证它”派发给PlannerAgent。PlannerAgent返回一个拆解结果“步骤一实现阶乘步骤二编写测试用例步骤三运行测试”。PlannerAgent产生一条AgentMessage发给CoderAgent内容是第一步的细化指令。CoderAgent执行生成代码返回结果同时继续把任务推给TesterAgent。TesterAgent执行测试返回通过/失败的结论。如果测试通过Orchestrator 直接把最终结果返回如果失败TesterAgent会把失败信息推回给CoderAgent进入下一轮修订。这个链条里的关键是消息的同步性。Orchestrator 的 while 循环是同步的——每次popleft一条消息就停下来等对应 Agent 的结果结果回来再继续处理队列下一条。这是源码中默认的行为单线程、容易调试但吞吐量低。如果要并行你得对接消息队列系统比如 Redis Stream这在源码里留了接口但没实现。多 Agent 场景最容易出的问题是“流程不终止”。如果max_turns设得过大或者循环条件没写对Agent 们会无限地互相发消息。源码用max_turns兜底默认值 30。但如果两个 Agent 都在“修订→再评审→再修订”的循环里出不来30 轮也会打满。我的经验是把max_turns调小一点比如 8宁可过早终止让用户重新发起也不要让任务卡死在后台烧 token。另外一个经验是给每个 Agent 设置不同的system_prompt时务必要写明“什么时候该结束并把控制权交回 Orchestrator”。这个写清楚比在 Orchestrator 里做各种终止条件判断都有效。框架文档里没提这一点但实际用过一轮的人应该都会有感觉。5. 常见问题与排查技巧实录源码看完了实操中肯定会碰到问题。我把这几周跑 Pi Agent 时踩过的坑集中整理一下每一条都是在真实场景里撞见、并且从源码层面定位原因之后才敢写的。现象根因分析排查/解决方向模型始终不调用工具工具 schema 生成失败模型看不到工具列表检查函数注解打印_tool_schemas()结果看是否符合 API 规范工具返回TypeError: xxx字符串模型传入了错误类型的参数_execute_tool捕获 TypeError 后原样返回工具函数里建议做显式类型校验返回友好错误不要依赖 Python 隐式转换超出max_iterations工具本身不稳定一直在报错导致模型反复重试先看日志里的工具调用链修复工具后再考虑调大阈值多 Agent 任务不结束某个 Agent 的 system prompt 没有说明终止条件重写该 Agent 的系统提示明确“任务完成时输出 DONE”历史对话影响新任务判断window_size过大或者没有正确使用摘要压缩启用SummaryMemory或者手动调用compress_messages模型返回的 JSON 参数残缺上下文被截断或者模型能力较弱缩短工具返回内容长度给工具的 description 写清楚参数边界工具名称冲突覆盖两个函数同名后一个覆盖前一个注册前加if name in self.tools检查并打日志排查技巧上我最常用的是在run()方法里临时加一行print(messages)把每轮循环的 messages 打出来看能直接看到模型返回了什么、工具执行结果是什么、下一轮模型是否拿到了正确的回填。这个方法虽然原始但比任何 debugger 都直观尤其是参数多的时候。另外一个技巧是“最小化测试工具”。当工具出问题时先写一个不带任何工具调用的纯对话任务确认模型本身通畅再挂一个只做return 1的桩工具确认工具调用链通然后再逐步加载真实工具。这能帮你迅速把“模型问题”和“工具问题”区分开。我每次给 Pi Agent 接入新工具都用这套流程基本十分钟内就能定位问题。我个人的体会是读 Pi Agent 源码最大的收获不是学会了某个具体调用而是理解了一个 Agent 框架在“上下文管理、工具调用、多角色协同”这三件事上是如何做权衡的。不同框架的取舍不同但 Pi Agent 这套实现足够轻、足够清晰适合当作研读 Agent 框架的入门蓝本。最后分享一个扩展思路如果你想把 Pi Agent 用在自己的项目里不要只把它当黑盒调试着在Agent.run()的循环体里加一个钩子函数——比如每轮调完模型后把 messages 快照写入本地——这样你就能沉淀出你自己的 Agent 行为数据集后续做 prompt 调优、成本分析都有数据可依。这步钩子接口目前在框架里没有正式暴露但run()方法内部加一行代码就可以实现改动成本很低。