玩转LlamaIndex:状态管理

21 阅读4分钟

这篇我们来说说共享状态、共享数据,以供后续步骤使用,如:用户信息、累计值

img

一、状态管理:让步骤之间"共享数据"

1.1 为什么要状态管理?

有些数据是整个 Workflow 共享的,比如:

  • 用户信息(从第一步输入,后面每步都要用)
  • 累积的统计数据(每步加一点)
  • 缓存的中间结果

这种"全局共享数据"就放在 Context 里。

1.2 Context 基础 API

from llama_index.core.workflow import Context

# 写
await ctx.set("user_name", "张三")
await ctx.set("counter", 0)

# 读
user_name = await ctx.get("user_name")  # "张三"
counter = await ctx.get("counter", default=0)  # 带默认值

# 自增
counter = await ctx.get("counter", default=0)
await ctx.set("counter", counter + 1)

1.3 完整示例:多步累积计算

class AccumulateWorkflow(Workflow):

    @step
    async def init(self, ctx: Context, ev: StartEvent) -> NumberEvent:
       """初始化 state,接收输入"""
       await ctx.store.set("total", 0)
       await ctx.store.set("count", 0)
       return NumberEvent(numbers = ev.numbers)
    @step
    async def accumulate(self, ctx: Context, ev: NumberEvent) -> StopEvent:
       """每收到一个数字就累加,所有数字走完就停"""
       total = await ctx.store.get("total")
       count = await ctx.store.get("count")

       for n in ev.numbers:
          total += n
          count += 1
          print(f"当前累加: {total},已处理 {count} 个数")

       await ctx.store.set("total", total)
       await ctx.store.set("count", count)

       return StopEvent(result = f"总和: {total},数量: {count}")

运行结果:

img

1.4 关键:step 怎么拿到 ctx?

async def accumulate(self, ctx: Context, ev: NumberEvent) -> StopEvent:
#                    ^^^                        ^^^
#                  第一个参数                    第二个参数

记住这个顺序:ctx 永远是第一个参数(放在 self 后面),ev 是事件参数(紧跟其后)。

框架通过类型注解 Context 来识别这个参数,自动注入

不用手动 Context(workflow)

1.5 状态序列化:跨进程持久化

Context 状态可以序列化到磁盘,实现"程序重启后状态还在":

# 保存
state = await ctx.to_dict()
import json
with open("workflow_state.json", "w") as f:
    json.dump(state, f)

# 恢复(在新的 Workflow 实例中)
with open("workflow_state.json", "r") as f:
    state = json.load(f)
new_ctx = Context(WorkflowClass)
await new_ctx.from_dict(state)

这个特性在做长任务(比如几个小时的数据处理)时特别有用——中途崩了可以从断点继续。

二、状态 vs 事件:啥时候用哪个?

这是新手最容易混的地方:

img

简单判断:

  • 数据只在一个 step 里用 → 用局部变量
  • 数据在相邻两步之间传 → 用事件
  • 数据多个 step 共享 → 用状态

三、组合实战:分支 + 循环 + 状态

来个真实点的例子:一个"猜数字"Workflow,数字不对就一直猜,猜对就停。

class GuessNumberWorkflow(Workflow):

    @step
    async def init(self, ctx: Context, ev: StartEvent) -> GuessEvent:
       """初初始化:把目标数字存到 state"""
       target = ev.target or random.randint(1, 100)

       # 设定初始搜索区间
       low = 1
       high = 100

       # 如果用户提供了第一次猜测,直接使用;否则取中点
       if ev.first_guess is not None:
          first_guess = ev.first_guess
       else:
          first_guess = (low + high) // 2

       # 存储到上下文
       await ctx.store.set("target", target)
       await ctx.store.set("low", low)
       await ctx.store.set("high", high)
       await ctx.store.set("attempts", 0)

       return GuessEvent(guess = first_guess, feedback = "开始猜数")

    @step
    async def accumulate(self, ctx: Context, ev: GuessEvent) -> StopEvent|GuessEvent:
       """每收到一个数字就累加,所有数字走完就停"""
       target = await ctx.store.get("target")
       attempts = await ctx.store.get("attempts", default = 0)
       low = await ctx.store.get("low")
       high = await ctx.store.get("high")

       attempts += 1

       await ctx.store.set("attempts", attempts)

       if ev.guess == target:
          return StopEvent(result = f"🎉 猜对了!用了 {attempts} 次")
       elif ev.guess < target:
          feedback = "太小了"
          low = max(low, ev.guess + 1)  # 目标在 guess+1 到 high 之间
       else:
          feedback = "太大了"
          high = min(high, ev.guess - 1)  # 目标在 low 到 guess-1 之间

       if attempts >= 10:
          return StopEvent(result = f"💀 超时失败,目标数字是 {target}")

       # 继续猜:返回 GuessEvent 又回到自己
       next_guess = (low + high) // 2  # 二分法
       print(f"第 {attempts} 次:猜 {ev.guess}({feedback}),区间 [{low}, {high}],下次猜 {next_guess}")

       # 更新存储的区间
       await ctx.store.set("low", low)
       await ctx.store.set("high", high)

       return GuessEvent(guess = next_guess, feedback = feedback)

运行会看到:

img

四、新手最常踩的坑

坑 1:状态忘了初始化

# ❌ 错误:没初始化就 get
@step
async def step1(self, ctx: Context, ev: StartEvent):
    total = await ctx.get("total")  # KeyError!
    ...

# ✅ 正确:get 时带默认值,或者先 set
total = await ctx.get("total", default=0)

坑 2:Context 的坑——并发写会覆盖

# ❌ 错误:多个 step 同时写同一个 key
@step
async def step_a(self, ctx: Context, ev: Ev) -> Ev:
    await ctx.set("counter", 10)
    ...

@step
async def step_b(self, ctx: Context, ev: Ev) -> Ev:
    await ctx.set("counter", 20)  # 可能覆盖 step_a 的写入
    ...

坑 3:Context赋值问题

# 此方式已过时
await ctx.set("counter", 10)
    
# 应该使用
await ctx.store.set("counter", 10)

DeepSeek和豆包都没找到问题所在,minmax直接找到问题点

五、本篇小结

这一篇我们把 Workflow 从"动态业务引擎"升级成了"可共享数据的动态业务引擎":

  • 循环:让 step 返回"上一步的事件类型",形成事件回路
  • 状态管理:ctx.get() / ctx.set() 共享全局数据
  • 状态序列化:ctx.to_dict() / ctx.from_dict() 跨进程持久化
  • 俩大坑:状态要初始化、并发要避免覆盖

核心心法:事件是接力棒,状态是记录本——事件负责"动",状态负责"存"。

下一篇我们讲流式输出与并发执行——让 Workflow 边跑边吐结果、多个分支同时跑,大幅提升性能和用户体验。