AG-UI: 状态管理

对于一个智能体来说，如果要从简单的聊天机器人，提升为真正解决实际问题的智能助手，状态管理是很重要的一环。

在聊天机器人式的交互中，用户消息和 LLM 响应交替出现。每次 LLM 基于全部的历史消息生成响应。这是一种线性的状态叠加方式。全部状态保存在历史消息中。对于文本生成的场景来说，这样的交互模式是适用的。用户关注的总是最近一次 LLM 的响应。

对于一些复杂的场景，这种聊天机器人的交互模式变得不再适用。原因在于，用户的每次消息只提供部分信息，而智能体的响应需要根据当前已累积的全部信息来做出判断。如果没有状态管理，LLM 只能从历史消息中尝试抽取全部信息，很可能是不准确的。

状态是用户和智能体所共用的工作区。用户和智能体都可以对状态进行修改。用户的主要角色是评估和判断，智能体的主要角色是提供智能建议。

举一个例子，行程规划。要得到一个完整的行程规划，需要用到多方面的数据。

• 行程的基本信息：人数，起始和结束日期，起始和结束地点
• 动态的环境信息：天气，景点的开放时间
• 用户的通用偏好：喜欢去的地方 / 不喜欢去的地方
• 用户的当前偏好：之前步骤中添加或删除的地点
• 当前的行程计划：上一次智能体的规划结果

这些数据构成了行程规划智能体的状态。在多轮与智能体的交互中，用户最终形成了一个完整的行程规划。

可以把状态看成一个大的对象。用户和智能体都会对该对象中的某些属性进行修改。以刚才的行程规划智能体为例：

• 用户以表单填写或自然语言的方式，提供行程的基本信息。
• 智能体通过工具调用，获取天气信息。
• 智能体展示规划结果，用户去掉其中的某个景点。
• 用户手动输入想要去的景点。

菜谱规划

这里用菜谱规划作为实现状态管理的示例。沿用之前的做菜助手智能体的例子。基本的场景是用户和智能体之间通过交互来确定所使用的食材，再基于这些食材推荐菜谱。

具体的场景经过了简化。智能体会有一份初始的食材列表。用户通过前端界面直接对食材列表进行修改。智能体也会根据用户的输入，对食材列表进行修改。

在真实的应用场景中，智能体可以通过智能冰箱获取实际可用的食材。如果食材的购买也通过该智能体完成，智能体可以自动追踪食材的购买和消耗。

具体实现

智能体的状态中包含了两个属性，available_ingredients 和 selected_ingredients，分别表示可用的食材列表和已选择的食材列表。

智能体的状态

class AgentState(MessagesState):
    available_ingredients: list[str]
    selected_ingredients: list[str]

LLM 工具

获取可用的食材列表

可用的食材列表由工具 get_available_ingredients 获取，这里只是返回一个固定的列表。当调用该工具时，如果 selected_ingredients 不存在，说明用户没有明确的选择，默认把全部食材更新到 selected_ingredients.

获取可用食材的工具

@tool
def get_available_ingredients(
    tool_call_id: Annotated[str, InjectedToolCallId],
    state: Annotated[dict, InjectedState],
) -> Command:
    """Return available ingredients and initialize the selected list if needed."""
    available_ingredients = [
        "chicken breast",
        "eggs",
        "tomatoes",
        "onions",
        "garlic",
        "potatoes",
        "carrots",
    ]
    update = {"available_ingredients": available_ingredients}
    if "selected_ingredients" not in state:
        update["selected_ingredients"] = available_ingredients
    return Command(update={
        **update,
        "messages": [ToolMessage(
            "Available ingredients state updated.",
            tool_call_id=tool_call_id,
        )],
    })

更新所选择的食材列表

工具 update_ingredients 的作用是更新状态。当用户要求添加食材时，LLM 会调用该工具来更新状态。

更新食材的工具

@tool
def update_ingredients(
    selected_ingredients: list[str],
    tool_call_id: Annotated[str, InjectedToolCallId],
) -> Command:
    """Update the selected ingredients list in the agent state.

    Args:
        selected_ingredients: The list of ingredients selected for cooking.

    Returns:
        A Command that updates the agent state with the provided ingredients.
    """
    return Command(update={
        "selected_ingredients": selected_ingredients,
        "messages": [ToolMessage("Selected ingredients state updated.", tool_call_id=tool_call_id)],
    })

LangGraph Middleware

当发送请求给 LLM 时，需要提供当前选择的食材列表。这是通过 LangGraph 的 middleware 来实现的。该 middleware 会修改发送给 LLM 的 system prompt。这样就确保了发送给 LLM 的食材列表总是来自当前的状态。

提供当前选择的食材列表的 middleware

def _inject_ingredients(request):
    ingredients = (request.state or {}).get("selected_ingredients") or []
    current_ingredients = ", ".join(ingredients) if ingredients else "none"
    content = (
        BASE_PROMPT
        + f"\n\nThe user currently has these selected ingredients: {current_ingredients}."
        + "\nThis selected ingredients list is authoritative. Ignore any ingredient lists "
        + "from earlier messages or tool results."
    )
    return request.override(system_message=SystemMessage(content=content))

class IngredientsMiddleware(AgentMiddleware):
    def wrap_model_call(self, request, handler):
        return handler(_inject_ingredients(request))

    async def awrap_model_call(self, request, handler):
        return await handler(_inject_ingredients(request))

前端实现

在前端界面上会展示当前已经选择的食材的列表。智能体的状态会通过 AG-UI 事件来传递。CopilotKit 内置提供了对状态同步的支持。这个同步是双向的。Web 应用可以订阅 onStateSnapshotEvent 事件把智能体后端的状态值同步给界面组件。也可以使用 setState 设置智能体的状态。当下一次运行智能体时，所设置的状态会被发送给智能体，更新后端的状态。

AG-UI LangGraph 集成库会自动使用请求中的状态更新 graph 内部的状态。

在下面的代码中，syncState 负责从智能体的后端同步状态。updateIngredients 调用 agent.setState 设置智能体的状态。

客户端状态同步

const { agent } = useAgent();
const [ingredients, setIngredients] = useState<string[]>([]);
const [availableIngredients, setAvailableIngredients] = useState<string[]>([]);

useEffect(() => {
    const syncState = (state: AgentState) => {
        setAvailableIngredients(state.available_ingredients ?? []);
        setIngredients(state.selected_ingredients ?? []);
    };

    syncState(agent.state as AgentState);

    const subscriber = agent.subscribe({
        onStateSnapshotEvent: ({ state }) => syncState(state as AgentState),
        onRunFinalized: ({ state }) => syncState(state as AgentState),
    });
    return () => subscriber.unsubscribe();
}, [agent]);

const updateIngredients = (next: string[]) => {
    setIngredients(next);
    agent.setState({ ...(agent.state as object), selected_ingredients: next });
};

演示

以下是智能体运行之后的演示界面。

获取可用食材

LLM 调用工具 get_available_ingredients 获取可用食材，并更新已选择食材的列表。界面右侧的食材列表同步更新。

修改所选择的食材

用户通过界面右侧的列表添加或删除食材，也可以在输入中说明。LLM 调用工具 update_ingredients 进行更新。

推荐菜谱

当需要推荐菜谱时，LLM 会根据状态中已选择的食材列表进行推荐。

源代码