跳到内容

编写代理

本教程将重点介绍系统的核心:代理本身,指导您在 Agent-lightning 中定义代理逻辑的各种方法。

任何代理的基本要求是

  1. 它必须接受单个 任务 作为输入。
  2. 它必须接受一组可调整的 资源(例如 PromptTemplateLLM)。
  3. 它必须 发出 跟踪跨度数据,以便算法可以理解其行为并从中学习。最简单的方法是返回最终奖励。

在实践中,请记住任务、资源和跨度具有额外的要求,以便使其在 Agent-lightning 中可训练

  1. 您需要一个训练数据集,其中包含一组任务,与您的代理期望的输入类型相同。
  2. 可调整的资源与算法相关。例如,我们看到的 APO 算法会调整一个 PromptTemplate。其他算法可能会调整模型权重或其他配置。
  3. 算法可以使用的跨度类型各不相同。几乎所有算法都支持在 rollout 结束时使用单个最终奖励跨度。但是,并非所有算法都支持在 rollout 中间发出奖励,更不用说像异常或日志消息之类的其他类型的跨度了。

本教程将向您展示如何编写可以处理各种任务和资源的代理,并发出所有类型的跨度。但是,您应该理解代理和算法通常是协同设计的。在算法中支持新的资源类型或跨度通常比仅仅将其添加到代理中复杂得多。

@rollout 装饰器

创建代理的最简单方法是编写标准的 Python 函数并使用 @rollout 装饰器对其进行标记。这种方法非常适合逻辑简单的代理,不需要复杂的状态管理。

Agent-lightning 会自动检查您的函数签名并注入所需的资源。例如,如果您的函数有一个名为 prompt_template 的参数,Agent-lightning 将找到当前 rollout 的 PromptTemplate 资源并将其传递进去。

让我们重新回顾一下第一个教程中的 room_selector 代理

from typing import TypedDict
from agentlightning import PromptTemplate, rollout

# Define a data structure for the task input
class RoomSelectionTask(TypedDict):
    # ... fields for the task ...
    pass

@rollout
def room_selector(task: RoomSelectionTask, prompt_template: PromptTemplate) -> float:
    # 1. Use the injected prompt_template to format the input for the LLM
    prompt = prompt_template.format(**task)

    # 2. Execute the agent's logic (e.g., call an LLM, use tools)
    # ...

    # 3. Grade the final choice to get a reward
    reward = room_selection_grader(final_message, task["expected_choice"])

    # 4. Return the final reward as a float
    return reward

当您训练此代理时,数据集预计将是 RoomSelectionTask 对象列表

from agentlightning import Dataset, Trainer

dataset: Dataset[RoomSelectionTask] = [
    RoomSelectionTask(date="2025-10-15", time="10:00", duration_min=60, attendees=10),
    RoomSelectionTask(date="2025-10-16", time="10:00", duration_min=60, attendees=10),
]

Trainer().fit(agent=room_selector, train_dataset=dataset)

在幕后,@rollout 装饰器会将您的函数包装在 FunctionalLitAgent 对象中,该对象是 LitAgent 的子类(如下所述),使其与 TrainerRunner 兼容。它支持诸如 taskprompt_templatellmrollout 等参数,为您提供了对执行上下文的灵活访问。

这里还有一个使用 llmrollout 参数进行更高级用法的示例。llm 参数为您提供了一个 OpenAI 兼容的 LLM 端点进行交互,该端点可以在底层由算法进行调整。rollout 参数为您提供完整的 Rollout 对象,其中包含 rollout ID、rollout 模式(训练或验证)等。

from openai import OpenAI
from agentlightning import LLM, Rollout

class FlightBookingTask(TypedDict):
    request: str
    expected_booking: dict

@rollout
def flight_assistant(task: FlightBookingTask, llm: LLM, rollout: Rollout) -> float:
    print(f"Rollout ID: {rollout.rollout_id}")
    print(f"Rollout Mode: {rollout.mode}")

    # Use the tuned LLM resource to create an OpenAI client
    client = OpenAI(
        # This endpoint could be a proxy to a proxy to a proxy ...
        # It could be different every time `flight_assistant` is called
        # But it should be OpenAI-API compatible
        base_url=llm.endpoint,

        # Use a dummy key if not provided
        # Usually this does not matter because the training LLM is often not guarded by an API key
        # But you can use `or os.environ["OPENAI_API_KEY"]` to make the function compatible with 3rd-party LLMs
        api_key=llm.api_key or "dummy-key",
    )

    # Make an API call with the specified model
    response = client.chat.completions.create(
        model=llm.model,
        messages=[{"role": "user", "content": task["request"]}],
    )
    # Whether the API supports features like streaming, tool calls, etc. depends on
    # the endpoint that algorithms are serving to you.
    final_message = response.choices[0].message.content

    # Grade the result and return a reward
    reward = grade_flight_booking(final_message, task["expected_booking"])
    return reward

从代理返回值

您的代理函数返回的值(即由 @rollout 装饰器装饰的函数的返回值)至关重要,因为它是报告 rollout 结果的主要方式。Agent-lightning 支持多种返回类型以适应不同的场景,从简单的奖励到详细的自定义跟踪。

  • float:这是最简单和最常见的返回类型。float 被视为整个 rollout 的最终奖励。Agent-lightning 会根据此值自动创建一个最终奖励跨度。

  • None:返回 None 告诉 runner 跟踪收集完全由 Tracer 通过自动检测(例如,通过 AgentOps)处理。在这种情况下,runner 将简单地检索 tracer 已经捕获的跨度。

发出最终奖励

当返回 None 时,您仍然必须确保记录最终奖励。您可以通过使用 emit_reward 函数(在 使用 Emitters 文档中介绍)来执行此操作。使用 @reward 装饰器包装您的奖励计算函数不再是推荐的方法。

  • list[ReadableSpan]list[SpanCoreFields]list[Span]:对于高级用例,您可以手动构建并返回 rollout 的所有跨度列表。这为您提供了对跟踪数据的完全控制。您可以返回 OpenTelemetry ReadableSpan 对象列表或 Agent-lightning 的本机 Span 对象列表。

对于大多数用户,返回一个简单的代理的 float 或返回 None 并使用 emitter 进行更复杂的代理是推荐的方法。

基于类的代理

对于需要状态、辅助方法或训练与验证不同的逻辑的更复杂的代理,您可以创建一个继承自 LitAgent 的类。这种面向对象的方法为代理的生命周期提供了更多的结构和控制。

要创建基于类的代理,您需要子类化 agentlightning.LitAgent 并实现其 rollout 方法。

以下是如何将 room_selector 实现为一个类。rollout 方法的签名与基于函数的代理略有不同,主要在于它如何处理资源。简单来说,算法不会只将一个 PromptTemplate 发送到代理,而是发送 NamedResources,这是一个从资源键到 Resource 的映射。这种设计是为了允许更高级的功能,例如多资源调整。

使用 @rollout 装饰器时,具有正确匹配类型的资源将自动注入到 rollout 方法中。但是,当您使用基于类的代理时,您需要手动从 resources 字典访问资源。内置算法列出了它们的资源键命名约定 这里

import agentlightning as agl

class RoomSelectorAgent(agl.LitAgent[RoomSelectionTask]):
    def rollout(self, task: RoomSelectionTask, resources: agl.NamedResources, rollout: agl.Rollout) -> float:
        # 1. Access the prompt_template from the resources dictionary
        prompt_template = resources["prompt_template"]

        # 2. Execute the agent's logic
        prompt = prompt_template.format(**task)
        # ...

        # 3. Grade the final choice
        reward = room_selection_grader(final_message, task["expected_choice"])

        # 4. Return the final reward
        return reward

# To use it with the trainer:
# agent = RoomSelectorAgent()
# trainer.fit(agent=agent, ...)

LitAgent 类提供了您可以覆盖的几个方法,以获得更细粒度的控制

注意

无论代理是异步的还是同步的,rollout 始终在异步上下文中执行。如果您的同步代理包含一些 asyncio.run() 调用,可能会引发一个错误,即已经有一个事件循环正在运行。为了避免阻塞事件循环,建议将内部异步操作卸载到单独的线程中。这里有一个示例代码

import asyncio
import queue
import threading

def run_sync_ephemeral(coro) -> Any:
    """
    Run an async coroutine from sync code.
    - If no loop in this thread: use asyncio.run() directly.
    - If already in an event loop: spawn a worker thread that calls asyncio.run()
    (which creates and closes a brand-new event loop per call).
    """
    try:
        asyncio.get_running_loop()
    except RuntimeError:
        # No running loop in this thread; safe to use asyncio.run
        return asyncio.run(coro)

    # Already in a running loop -> execute in a worker thread
    q = queue.Queue[Any]()

    def worker():
        try:
            result = asyncio.run(coro)  # creates & closes its own loop
            q.put((True, result))
        except BaseException as e:
            q.put((False, e))

    t = threading.Thread(target=worker, daemon=True)
    t.start()
    ok, payload = q.get()
    t.join()
    if ok:
        return payload
    raise payload