跳到主要内容

Agent 主循环:一问一答怎么跑完

30 秒导读: 你问 PandasAI 一句自然语言(比如"哪个国家 GDP 最高?"),它内部不是直接答你, 而是先让 LLM 写一段 Python 代码,再执行这段代码,把代码的运行结果当答案返回。这一章讲的就是 从"收到问题"到"吐出结果"这条端到端主线的控制流:两个阶段(生成、执行)、每个阶段各自的重试、 以及贯穿全程的上下文对象 AgentState。

本章只讲主线控制流。代码怎么被校验/清洗/安全执行,见 代码即 SQL; 喂给 LLM 的 prompt 长什么样,见 Prompt 工程;数据集/视图/SQL 编译,见 语义层。

1. 先建立直觉:它把"回答问题"变成"写代码 + 跑代码"​

传统聊天机器人:你问一句,模型直接生成一段文字答案。PandasAI 不一样——它多绕了一道:

  • 它不让模型直接回答,而是让模型写一段操作数据的 Python 代码。
  • 然后在本地执行这段代码,代码算出来的东西(一个数字、一张表、一张图)才是最终答案。

为什么要绕这道?因为模型擅长写代码、不擅长算数。让它对着几百万行数据"心算"必然出错; 但让它写一句 SELECT country ORDER BY gdp DESC LIMIT 1,交给真正的引擎去跑,结果就精确可信。

一句类比:你不是问会计"总和是多少",而是让会计写好公式、按下计算器。模型负责"写公式", 引擎负责"按计算器"。

最小使用示例(帮你感受输入输出形态):

import pandasai as pai

df = pai.DataFrame("data.csv")
# 一次 chat:内部会生成代码 + 执行代码,返回的是"执行结果"而不是"模型原话"
result = df.chat("哪个国家的 GDP 最高?")

# 追问:复用上一轮的记忆和上一段代码,做增量修改
result = df.follow_up("那前三名呢?")
  • chat = 开一段新对话(先清空记忆),问第一句。
  • follow_up = 在已有对话上追问,带着上下文继续。

两者的区别只有一处:chat 会先清记忆,follow_up 不会。真正干活的是它们共同调用的 _process_query。

真实入口在 pandasai/agent/base.py:chat(:92)、follow_up(:105)。 chat 里先做了个前置检查——没有 LLM credits 直接报错(:96),因为"没有模型就没法写代码"。

2. 顶层全景:两阶段流水线​

整条主线可以浓缩成一句话:生成代码(带重试)→ 执行代码(带重试)。

              ┌─────────────────────── _process_query (base.py:271) ───────────────────────┐
              │                                                                             │
 用户问题 ──▶  │  ① generate_code_with_retries        ② execute_with_retries                │ ──▶ 结果
              │     "让 LLM 写代码"                     "跑代码 + 解析结果"                   │
              │        │  失败?                            │  失败?                        │
              │        ▼  把 traceback 喂回 LLM            ▼  把 traceback 喂回 LLM          │
              │     _regenerate_code_after_error ◀────────┘  (两阶段共用这个改错函数)         │
              └─────────────────────────────────────────────────────────────────────────────┘
                         贯穿全程:AgentState(记忆 memory + last_code_generated + config/llm)

怎么读这张图: 从左到右是主线顺序;两个阶段都可能失败,失败时都掉进同一个"把报错喂回 LLM 让它改"的 函数 _regenerate_code_after_error;底部那行 AgentState 是从头贯穿到尾的上下文与记忆。

各部件一句话职责:

部件干什么在哪里
chat / follow_up对外入口;chat 先清记忆再问,follow_up 直接续问agent/base.py:92 / :105
_process_query主线编排:调生成、调执行、兜底异常agent/base.py:271
generate_code_with_retries生成阶段:调 LLM 写代码,失败按重试改agent/base.py:171
execute_with_retries执行阶段:跑代码 + 解析结果,失败按重试改agent/base.py:197
_regenerate_code_after_error两阶段共用的"改错":把 traceback 喂回 LLMagent/base.py:296
_handle_exception兜底:重试用光后返回 ErrorResponseagent/base.py:310
CodeGenerator.generate_code"生成→校验→清洗"三步,产出可执行代码core/code_generation/base.py:16
AgentState贯穿上下文:记忆、上一段代码、中间值、config/llmagent/state.py:24

主线走一遍(高层,不进代码): 用户问题进来 → _process_query 先给这轮分配一个 prompt_id → 调 ① 生成代码(内部让 LLM 写、校验、清洗)→ 调 ② 执行代码(跑出结果、解析成响应)→ 返回结果。 中途任何阶段崩了,先在本阶段内重试改错;重试用光了才向上抛,由 _handle_exception 兜成一个错误响应。

3. AgentState:贯穿全程的"上下文与记忆"​

在讲两个阶段之前,得先认识那个从头贯穿到尾的对象——AgentState(agent/state.py:24)。 它是一个 @dataclass,注释直说自己是"在流水线各步之间传递属性的上下文类"。你可以把它理解为 这一整段对话的共享白板:谁都能往上写、往上读。

白板上最关键的几格:

字段装什么谁在用
memory多轮对话历史(问/答交替)生成 prompt 时取上下文;chat 时被清空
last_code_generated上一次生成并清洗后的代码决定"给初始模板还是给上一段代码";改错时作为基底
intermediate_values中间产物的键值缓存执行期存取(add/get,:98/:106)
output_type本轮期望的返回类型生成 prompt 与校验时约束
last_prompt_id本轮的唯一 id日志追踪(assign_prompt_id,:87)

config 与 llm 的解析是 AgentState 一个容易被忽略的巧处。config 是一个 property(:110):

  • 如果本地设了 _config,就用本地的;否则回退到全局 pai.config.get()(:118)。
  • 初始化时 _get_config(:73)负责把 dict 转成 Config 对象,或在没传时取全局默认。

LLM 就挂在 config.llm 上——所以 _process_query 一进来就能 self._state.config.llm.type(base.py:276) 拿到当前模型;真正写代码那一步是 config.llm.generate_code(...)(code_generation/base.py:33)。 这就是"config/llm 解析"如何把主循环和背后的模型接起来。

Agent 在构造时初始化这块白板:self._state.initialize(dfs, config, memory_size, vectorstore, description) (base.py:83,展开见 state.py:45)。

4. 阶段①:生成代码(带重试)​

这一节讲 generate_code_with_retries(base.py:171)——"让 LLM 把代码写出来,写不出来就按报错重写"。

4.1 一次干净的生成:memory 先记账,再拼 prompt​

单次生成走 generate_code(base.py:111),三个动作:

# 真实逻辑(base.py:111),精简展示
self._state.memory.add(str(query), is_user=True)   # ① 把用户问题记进记忆
prompt = get_chat_prompt_for_sql(self._state)      # ② 用当前 state 拼出 prompt
code = self._code_generator.generate_code(prompt)  # ③ 交给 CodeGenerator 产出代码

注意顺序:先把 query 写进 memory,再拼 prompt。这一步顺序很关键,下面 §4.3 会解释它如何影响 "给初始模板还是给上一段代码"。

4.2 CodeGenerator 的"生成→校验→清洗"三步​

generate_code(prompt) 拿到 prompt 后,真正干活的是 CodeGenerator.generate_code (core/code_generation/base.py:16)。它把"从 prompt 到可用代码"拆成三步:

   prompt
     │
     ▼
 ① 生成   config.llm.generate_code(prompt, ctx)   → 原始代码,存进 last_code_generated
     │
     ▼
 ② 校验   CodeRequirementValidator.validate(code)  → 不满足要求就 raise ValueError
     │
     ▼
 ③ 清洗   CodeCleaner.clean_code(code)             → 清洗后的代码,再次覆盖 last_code_generated
     │
     ▼
 返回清洗后的代码

对应源码:生成在 :33,校验在 validate_and_clean_code(:54)里先 validate(:57), 清洗 clean_code(:63)。一个细节:last_code_generated 被写了两次—— 先存原始生成(:35,便于日志),清洗后再覆盖成最终版(:41,便于后续多轮复用)。 校验和清洗的内部规则不在本章,见 代码即 SQL。

4.3 memory 如何驱动"初始模板 vs 上一段代码"​

这是主循环里最精妙的一处:同一个 prompt 模板,第一轮和后续轮喂给 LLM 的"起点"不一样。

判断逻辑写在模板里(core/prompts/templates/generate_python_code_with_sql.tmpl:9):

{% if last_code_generated and context.memory.count() > 0 %}
Last code generated:
{{ last_code_generated }}
{% else %}
Update this initial code:
```python
# TODO: import the required dependencies
import pandas as pd
# Write code here
# Declare result var: ...
```
{% endif %}

拆开看这个条件:

  • context.memory.count() > 0:几乎总为真。因为 §4.1 里 generate_code 会先 memory.add 再拼 prompt,所以拼 prompt 时记忆里至少已有当前这条 query。
  • 真正的开关是 last_code_generated 有没有值:
轮次last_code_generatedLLM 收到的起点
首轮(chat 后第一问)None(初始)"Update this initial code" + 一段空白代码模板
后续(follow_up / 重试)上一轮清洗后的代码"Last code generated:" + 上一段真实代码

直觉:第一次没有历史代码可改,就给一张"填空模板"让 LLM 从零写起;之后有了上一段代码,就把它交回去, 让 LLM 在既有代码上做增量修改——这也是 follow_up 能"接着上文改"的底层原因。

这条 prompt = get_chat_prompt_for_sql(self._state)(base.py:117)里, get_chat_prompt_for_sql 正是把 context.last_code_generated 传进模板的(core/prompts/__init__.py:19)。

4.4 生成阶段的重试语义​

回到带重试的外层 generate_code_with_retries(base.py:171):

# 精简自 base.py:171
max_retries = self._state.config.max_retries        # 默认 3(config.py:13)
try:
    return self.generate_code(query)                # 先正常生成一次
except Exception as e:
    exception = e
    while attempts <= max_retries:
        try:
            return self._regenerate_code_after_error(   # 带着上段代码 + 报错重新生成
                self._state.last_code_generated, exception)
        except Exception as e:
            exception = e
            attempts += 1
            if attempts > max_retries:
                raise                                    # 用光重试,向上抛

要点:

  • 首次生成不算重试;失败后最多再试 max_retries 次(默认 3,config.py:13)。
  • 每次重试都调 _regenerate_code_after_error,把 last_code_generated + 异常一起喂回 LLM。
  • 用光仍失败就 raise,冒泡到 _process_query。

5. 阶段②:执行代码(带重试)​

生成阶段产出一段代码后,_process_query 把它交给 execute_with_retries(base.py:197)。 这一阶段的结构和生成阶段对称,但重试触发点不同——这里是代码跑崩了而不是写不出来。

# 精简自 base.py:197
while attempts <= max_retries:
    try:
        result = self.execute_code(code)                 # 跑代码
        return self._response_parser.parse(result, code) # 解析成最终响应
    except Exception as e:
        attempts += 1
        if attempts > max_retries:
            raise
        code = self._regenerate_code_after_error(code, e)  # 报错喂回 LLM,拿到新代码,下轮再跑

关键差异对照:

生成阶段 generate_code_with_retries执行阶段 execute_with_retries
触发重试LLM 生成/校验/清洗失败代码执行或结果解析抛异常
上限max_retries(默认 3)max_retries(默认 3),独立计数
改错方式_regenerate_code_after_error_regenerate_code_after_error(同一函数)
重试产物新的代码字符串新代码 → 回到 while 顶部再执行一次

注意:两个阶段各自最多 max_retries 次,互相独立。一次问答最坏情况下可能触发多达 (生成侧几次 + 执行侧几次)的 LLM 调用。execute_code 本身(跑代码、沙箱、解析)属于第 2 章范畴。

6. 两阶段共用的"改错":把 traceback 喂回 LLM​

两个阶段失败时都掉进同一个函数 _regenerate_code_after_error(base.py:296)。 它的核心思想是 agent 式自我纠错:不是简单重跑,而是把出错的代码和完整报错栈交给 LLM,让它读懂错误再改。

# 精简自 base.py:296
error_trace = traceback.format_exc()                  # 抓完整报错栈
if isinstance(error, InvalidLLMOutputType):           # 错的是"返回类型不对"
    prompt = get_correct_output_type_error_prompt(self._state, code, error_trace)
else:                                                  # 其它执行错误(如 SQL 用法错)
    prompt = get_correct_error_prompt_for_sql(self._state, code, error_trace)
return self._code_generator.generate_code(prompt)     # 再走一遍"生成→校验→清洗"

两点值得记:

  • 按错误类型选不同的纠错 prompt:类型错误(InvalidLLMOutputType)用一套 prompt, 普通执行错误用另一套(core/prompts/__init__.py:27/:35)。prompt 细节见 第 3 章。
  • 改错本身也走 CodeGenerator 三步:纠错生成的新代码同样要过校验+清洗,不是拿到就用。

7. 兜底:重试用光之后​

_process_query(base.py:271)是主线编排者,也是最后的安全网:

# 精简自 base.py:271
try:
    self._state.assign_prompt_id()               # 给本轮分配 id(日志追踪)
    code = self.generate_code_with_retries(str(query))  # 阶段①
    result = self.execute_with_retries(code)            # 阶段②
    return result
except CodeExecutionError:
    return self._handle_exception(code)          # 兜底:不抛,返回错误响应
  • 只有 CodeExecutionError 会被这里捕获,交给 _handle_exception(base.py:310)—— 它把 traceback 包成一个 ErrorResponse(带 last_code_executed 和 error)返回, 不让异常冒泡到用户。
  • 生成阶段用光重试后 raise 的异常,若不是 CodeExecutionError,则会直接向上传播(不被这里兜)。

至此,一次"一问一答"的完整生命周期结束:入口 → 分配 id → 生成(带重试)→ 执行(带重试)→ 返回结果或错误响应。全程 AgentState 这块白板记录着记忆和上一段代码,让下一轮 follow_up 能接上。

8. 巧妙之处(可带走的精华)​

  • "生成 + 执行"两阶段各自重试,却共用一个改错函数。 生成崩(写不出)和执行崩(跑不通)是两类 不同的失败,但纠错手段统一为"把 traceback 喂回 LLM"(_regenerate_code_after_error,base.py:296), 减少了重复逻辑。
  • 用 last_code_generated 是否为空,一个字段区分首轮与续轮。 不需要额外的"轮次计数器", 模板里 {% if last_code_generated ... %}(generate_python_code_with_sql.tmpl:9)就实现了 "首轮给填空模板 / 续轮给上段代码增量改"。
  • memory 先记账再拼 prompt。 generate_code 里 memory.add 在拼 prompt 之前(base.py:114), 保证了模板条件里 memory.count() > 0 恒真,把开关权完全交给 last_code_generated。
  • 兜底只吞 CodeExecutionError。 _process_query 只捕获这一类异常(base.py:293), 让用户永远拿到结构化的 ErrorResponse 而非裸异常,同时不掩盖其它编程错误。

9. 边界与局限(诚实)​

  • 重试次数硬上限。 两阶段各 max_retries(默认 3,config.py:13)。若 LLM 反复改不对, 最终仍会失败——生成侧 raise,执行侧被兜成 ErrorResponse。
  • chat 需要 LLM credits。 chat 一进门就检查 config.llm,为空直接 ValueError(base.py:96)。
  • 本章不覆盖的黑箱: 代码怎么被校验/清洗/在沙箱里执行(第 2 章)、prompt 模板每个片段 (第 3 章)、数据集/视图/SQL 编译(第 4 章)。本章只保证你看清"主线控制流"这一层。

10. 代码地图(导航索引)​

主题文件路径符号名
对外入口:新对话pandasai/agent/base.pyAgent.chat
对外入口:追问pandasai/agent/base.pyAgent.follow_up
主线编排 + 异常兜底pandasai/agent/base.pyAgent._process_query
生成阶段(带重试)pandasai/agent/base.pyAgent.generate_code_with_retries
单次生成(记忆+拼prompt)pandasai/agent/base.pyAgent.generate_code
执行阶段(带重试)pandasai/agent/base.pyAgent.execute_with_retries
改错:traceback 喂回 LLMpandasai/agent/base.pyAgent._regenerate_code_after_error
兜底:包成 ErrorResponsepandasai/agent/base.pyAgent._handle_exception
生成→校验→清洗三步pandasai/core/code_generation/base.pyCodeGenerator.generate_code
贯穿上下文与记忆pandasai/agent/state.pyAgentState
config/llm 解析pandasai/agent/state.pyAgentState.config / _get_config
首轮/续轮模板开关pandasai/core/prompts/templates/generate_python_code_with_sql.tmpl{% if last_code_generated ... %}
拼 chat promptpandasai/core/prompts/__init__.pyget_chat_prompt_for_sql
记忆存储pandasai/helpers/memory.pyMemory.add / get_last_message