使用 Agent-lightning 编写第一个算法¶

在第一个教程“训练第一个 Agent”中，我们介绍了 Trainer，并展示了如何使用预构建的算法，例如 自动提示优化 (APO) 来提高 Agent 的性能。 Trainer 处理了所有复杂的交互，让我们专注于 Agent 的逻辑。

现在，我们将更进一步。如果您有一个独特的训练想法，不符合标准算法，该怎么办？本教程将向您展示如何从头开始编写自己的自定义算法。我们将构建一个简单的算法，系统地测试提示模板列表，并识别出具有最高奖励的模板。

通过本教程，您将了解 Algorithm、Runner 和一个新组件 Store 的核心机制，它们共同创建了 Agent-lightning 核心训练循环的强大功能。

提示

本教程可帮助您构建对 Agent-lightning 核心组件的交互方式的基本理解。建议所有定制算法的用户阅读本教程，即使对于那些不打算进行提示优化的用户也是如此。

训练的核心概念¶

在深入了解 LightningStore 之前，让我们定义两个关键概念，它们是 Agent-lightning 中任何训练过程的核心：资源和 Tracer。

资源：可调用的资产¶

资源是您的算法试图改进的资产。可以将它们视为 Agent 执行任务所使用的“配方”。此配方可以是

指导 LLM 的 提示模板。
机器学习模型的权重。
Agent 所需的任何其他配置或数据。

算法的工作是运行实验并迭代更新这些资源，以找到性能最佳的版本。

Tracer：数据收集器¶

算法如何知道更改是否是改进？它需要数据。这就是 Tracer 的作用。

Tracer 自动 instrument（即修改/修补）Agent 的代码。这意味着它会监视重要事件，例如 LLM 调用、工具的使用或 奖励信号，并记录发生情况的详细日志。每个这样的日志称为 Span（已经在上一个教程中介绍过）。

来自单个任务执行的 Span 集合为算法提供了 Agent 行为的完整、逐步跟踪，这对于学习和改进至关重要。我们的默认 tracer 构建在 AgentOps SDK 之上，以支持 instrumenting 用各种 Agent/非 Agent 框架编写的代码。

中央枢纽：LightningStore¶

现在，所有这些资源、任务和 Span 存储在哪里？它们都由 LightningStore 管理。

LightningStore 作为整个系统的中央数据库和消息队列。它是 Algorithm 与 Runners 之间解耦的唯一事实来源。

注意

在上一个教程中，我们简化了训练循环，说 Algorithm 和 Agent “通过 Trainer” 进行通信。从宏观层面上来说，这是正确的，但实际上使这一切成为可能的是 LightningStore。

Algorithm 连接到 Store 以 enqueue_rollout（任务）和 update_resources（如新的提示模板）。它还查询 Store 以检索已完成 rollout 的结果 Span 和奖励。
Runners 连接到 Store 以 dequeue_rollout（轮询可用任务）。执行任务后，它们使用 Tracer 将结果 Span 和状态更新写回 Store。

这种架构是 Agent-lightning 可扩展性的关键。由于 Algorithm 和 Runners 仅与 Store 通信，因此它们可以在不同的进程甚至不同的机器上运行。

Store Architecture

对 Store 包含内容的心理模型

LightningStore 不仅仅是一个简单的数据库；它是一个用于管理整个训练生命周期的组织系统。以下是它跟踪的内容

任务队列：一个等待 Runner 拾取、可通过 enqueue_rollout 和 dequeue_rollout 交互的待处理 Rollout 队列。
Rollout：单个任务的记录。Rollout 包含有关任务的元数据，并跟踪完成它的所有 Attempt，可通过 query_rollouts 和 wait_for_rollouts 交互。
Attempt：Rollout 的一次执行。如果尝试失败（例如，由于网络错误），如果已配置，Store 可以自动安排重试。每个尝试都链接到其父 Rollout，并包含状态和时间信息。Rollout 状态与子状态同步。 对于初学者，您可以假设每个 Rollout 只有一个 Attempt，除非您已明确配置重试。
Span：由 Tracer 在尝试期间生成的详细结构化日志。每个 Span 链接到其父 Attempt 和 Rollout。
资源：Algorithm 创建的资产（如提示模板）的版本化集合。每个 Rollout 链接到它应使用的资源的特定版本。

构建自定义算法¶

让我们构建一个算法，从预定义的列表中找到最佳系统提示。逻辑很简单

从候选提示模板列表中开始。
对于每个模板，在 Store 中创建一个“资源”包。
Enqueue 一个 rollout（一个任务），告诉 Runner 使用此特定资源。
等待 Runner 拾取任务并完成它。
查询 Store 以获取 rollout 的 Span 中的最终奖励。
测试完所有模板后，比较奖励并宣布最佳模板。

我们可以将此作为与 LightningStore 直接交互的简单 Python 函数来实现。

async def find_best_prompt(store, prompts_to_test, task_input):
    """A simple algorithm to find the best prompt from a list."""
    results = []

    # Iterate through each prompt to test it
    for prompt in prompts_to_test:
        print(f"[Algo] Updating prompt template to: '{prompt}'")

        # 1. Update the resources in the store with the new prompt
        resources_update = await store.add_resources(
            resources={"prompt_template": prompt}
        )

        # 2. Enqueue a rollout task for a runner to execute
        print("[Algo] Queuing task for clients...")
        rollout = await store.enqueue_rollout(
            input=task_input,
            resources_id=resources_update.resources_id,
        )
        print(f"[Algo] Task '{rollout.rollout_id}' is now available for clients.")

        # 3. Wait for the rollout to be completed by a runner
        await store.wait_for_rollouts([rollout.rollout_id])

        # 4. Query the completed rollout and its spans
        completed_rollout = await store.get_rollout_by_id(rollout.rollout_id)
        print(f"[Algo] Received Result: {completed_rollout.model_dump_json(indent=None)}")

        spans = await store.query_spans(rollout.rollout_id)
        # We expect at least two spans: one for the LLM call and one for the final reward
        print(f"[Algo] Queried Spans:\n  - " + "\n  - ".join(str(span) for span in spans))
        # find_final_reward is a helper function to extract the reward span
        final_reward = find_final_reward(spans)
        print(f"[Algo] Final reward: {final_reward}\n")

        results.append((prompt, final_reward))

    # 5. Find and print the best prompt based on the collected rewards
    print(f"[Algo] All prompts and their rewards: {results}")
    best_prompt, best_reward = max(results, key=lambda item: item[1])
    print(f"[Algo] Best prompt found: '{best_prompt}' with reward {best_reward}")

异步操作

您会注意到 async 和 await 关键字。Agent-lightning 构建在 asyncio 之上，以有效地处理并发操作。与 store 的所有交互都是异步网络调用，因此必须等待它们。

Agent 和 Runner¶

我们的算法需要一个 agent 来执行任务和一个 runner 来管理该过程。

runner 是一个长期运行的 worker 进程。它的工作很简单

通过 LightningStoreClient 连接到 LightningStore。
进入循环，不断向 LightningStore 询问新任务（dequeue_rollout）。
当它获得任务时，它会运行 simple_agent 函数。
至关重要的是，runner 使用 Tracer 包装 agent 执行。tracer 会自动捕获所有重要事件（例如 LLM 调用和最终奖励）作为 Span 并将其发送回 LightningStore。

# Connecting to Store
store = agl.LightningStoreClient("https://:4747")  # or some other address
runner = LitAgentRunner[str](tracer=AgentOpsTracer())
with runner.run_context(agent=simple_agent, store=store):  # <-- where the wrapping and instrumentation happens
    await runner.iter()  # polling for new tasks forever

对于本示例，agent 的工作是获取资源中的提示，使用它向 LLM 提问，并返回分数。

def simple_agent(task: str, prompt_template: PromptTemplate) -> float:
    """An agent that answers a question and gets judged by an LLM."""
    client = OpenAI()

    # Generate a response using the provided prompt template
    prompt = prompt_template.format(any_question=task)
    response = client.chat.completions.create(
        model="gpt-4.1-nano", messages=[{"role": "user", "content": prompt}]
    )
    llm_output = response.choices[0].message.content
    print(f"[Rollout] LLM returned: {llm_output}")

    # This llm_output and the final score are automatically logged as spans by the Tracer
    score = random.uniform(0, 1)  # Replace with actual scoring logic if needed
    return score

运行示例¶

要查看所有内容的操作，您需要在三个单独的终端窗口中。

提示

如果您想跟随，可以在 apo_custom_algorithm.py 文件中找到此示例的完整代码。

1. 启动 Store： 在第一个终端中，启动 LightningStore 服务器。此组件将等待来自算法和 runner 的连接。store 将默认在端口 4747 ⚡ 上侦听。

agl store

2. 启动 Runner： 在第二个终端中，启动 runner 进程。它将连接到 store 并等待任务。

启动 runner 的代码如下所示

export OPENAI_API_KEY=sk-... # Your OpenAI API key
python apo_custom_algorithm.py runner

您将看到输出，表明 runner 已启动并正在等待 rollout。

2025-10-14 22:23:41,339 [INFO] ... [Worker 0] Setting up tracer...
2025-10-14 22:23:41,343 [INFO] ... [Worker 0] Instrumentation applied.
2025-10-14 22:23:41,494 [INFO] ... [Worker 0] AgentOps client initialized.
2025-10-14 22:23:41,494 [INFO] ... [Worker 0] Started async rollouts (max: unlimited).

3. 启动 Algorithm： 在第三个终端中，运行算法。这将启动整个过程。

例如，我们使用以下参数运行算法代码

prompts_to_test = [
    "You are a helpful assistant. {any_question}",
    "You are a knowledgeable AI. {any_question}",
    "You are a friendly chatbot. {any_question}",
]
task_input = "Why is the sky blue?"
store = agl.LightningStoreClient("https://:4747")
find_best_prompt(store, prompts_to_test, task_input)

或者，您可以使用我们预先编写的脚本来试用

python apo_custom_algorithm.py algo

理解输出¶

当算法运行时，您将在所有三个终端中看到日志出现，显示组件实时交互。

Algorithm 输出： 算法终端显示主要控制流：更新提示、排队任务和接收最终结果。您还可以看到它从 store 中检索的原始 Span 数据。

[Algo] Updating prompt template to: 'You are a helpful assistant. {any_question}'
[Algo] Queuing task for clients...
[Algo] Task 'ro-1d18988581cd' is now available for clients.
[Algo] Received Result: rollout_id='ro-1d18988581cd' ... status='succeeded' ...
[Algo] Queried Spans:
  - Span(name='openai.chat.completion', attributes={'gen_ai.prompt.0.content': 'You are a helpful assistant...', 'gen_ai.completion.0.content': 'The sky appears blue...'})
  - Span(name='reward', attributes={'value': 0.95})
[Algo] Final reward: 0.95

[Algo] Updating prompt template to: 'You are a knowledgeable AI. {any_question}'
...
[Algo] Final reward: 0.95

[Algo] Updating prompt template to: 'You are a friendly chatbot. {any_question}'
...
[Algo] Final reward: 1.0

[Algo] All prompts and their rewards: [('You are a helpful assistant. {any_question}', 0.95), ('You are a knowledgeable AI. {any_question}', 0.95), ('You are a friendly chatbot. {any_question}', 1.0)]
[Algo] Best prompt found: 'You are a friendly chatbot. {any_question}' with reward 1.0

Runner 输出： runner 终端显示它拾取每个任务、执行 agent 逻辑和报告完成情况。

[Rollout] LLM returned: The sky appears blue due to Rayleigh scattering...
2025-10-14 22:25:50,803 [INFO] ... [Worker 0 | Rollout ro-a9f54ac19af5] Completed in 4.24s. ...

[Rollout] LLM returned: The sky looks blue because of a process called Rayleigh scattering...
2025-10-14 22:25:59,863 [INFO] ... [Worker 0 | Rollout ro-c67eaa9016b6] Completed in 4.06s. ...

Store Server 输出： store 终端显示详细的日志，记录了所有交互，确认了它作为中央枢纽的作用。您可以查看 enqueue 和 dequeue rollout、添加 Span 以及更新状态的请求。

... "POST /enqueue_rollout HTTP/1.1" 200 ...
... "GET /dequeue_rollout HTTP/1.1" 200 ...
... "POST /add_span HTTP/1.1" 200 ...
... "POST /update_attempt HTTP/1.1" 200 ...
... "POST /wait_for_rollouts HTTP/1.1" 200 ...
... "GET /query_spans/ro-c67eaa9016b6 HTTP/1.1" 200 ...

训练器在哪里？

你可能想知道为什么上一个教程关注的是Trainer类，但我们在这里还没有使用它。

将Trainer视为一个方便的包装器，它可以为你管理整个训练过程。当你希望将预构建的算法应用于你的智能体，而无需担心其底层机制时，它非常完美。Trainer会处理启动LightningStore，协调Runners，管理它们的生命周期，并处理错误。

然而，在本教程中，我们的目标是构建一个新的算法。为此，我们需要直接与核心组件交互：Store、Runner以及算法逻辑本身。分别运行它们可以让你获得更多的控制权和更清晰、隔离的日志，这对于开发和调试来说是理想的。

一旦你的自定义算法成熟，你可以将其打包以符合我们的标准接口（@algo或Algorithm）。这样就可以再次使用Trainer，从而获得自动生命周期管理的所有好处，同时使用你自己的自定义逻辑。一个执行此操作的示例代码可在apo_custom_algorithm_trainer.py中找到。