Skip to content

Build & Finetune Agents

How It Fits Together

NexAU defines and runs agents, NexRL turns agent trajectories into RL training data, and Weaver performs remote model training and weight sync. You can use all three together or replace any layer.

What is NexAU?

NexAU is a framework for tool-using LLM agents. It supports YAML agent definitions, flexible tools, LLM backend configuration, tracing, and training-framework integration.

Common use cases:

  • Web research, news search, SQL agents, and other tool-use tasks.
  • Multi-turn reasoning and long-context workflows.
  • Turning agent traces into RL signals.

Build a Web Research Agent

A simplified deep research agent:

yaml
type: agent
name: deep_research_agent
max_context_tokens: 100000

system_prompt: |
  Date: {{date}}
  You are a research agent. Search for information using web_search and web_read.

system_prompt_type: jinja
tool_call_mode: openai

llm_config:
  temperature: 0.7
  max_tokens: 4096
  api_type: openai_chat_completion

tools:
  - name: web_search
    yaml_path: ./tools/WebSearch.yaml
    binding: nexau.archs.tool.builtin.web_tool:web_search
  - name: web_read
    yaml_path: ./tools/WebRead.yaml
    binding: nexau.archs.tool.builtin.web_tool:web_read

Run it:

python
from datetime import datetime
import os

from nexau import Agent, AgentConfig, LLMConfig


agent_config = AgentConfig.from_yaml("deep_research_agent.yaml")
agent_config.llm_config = LLMConfig(
    model=os.getenv("LLM_MODEL"),
    base_url=os.getenv("LLM_BASE_URL"),
    api_key=os.getenv("LLM_API_KEY"),
)

agent = Agent(agent_config)
response = agent.run(
    "What are the latest developments in quantum computing?",
    context={"date": datetime.now().strftime("%Y-%m-%d")},
)
print(response)

Connect to NexRL Training

To let NexRL read the agent config, run the agent, collect trajectories, and compute rewards, organize task files in a workspace:

text
recipe/my_task/agent_workspace/
├── agent_config.yaml
├── evaluator.py
├── custom_worker.py
└── tools/
FileRole
agent_config.yamlNexAU agent definition.
evaluator.pyReward computation and answer checking.
custom_worker.pyOptional formatting, prompting, or trajectory conversion.
tools/Custom tool definitions and implementations.

Enable Tracing for RL

NexRL needs tokens, tool calls, and response traces from each agent run. Agent configs usually enable tracing:

yaml
llm_config:
  temperature: 0.7
  max_tokens: 8192
  api_type: openai_chat_completion

  logprobs: true
  extra_body:
    train_mode: true
    include_stop_str_in_output: true
    skip_special_tokens: false
    return_tokens_as_token_ids: true

tracers:
  - import: nexau.archs.tracer.adapters.in_memory:InMemoryTracer

These settings make rollout results convertible into tokens, old logprobs, masks, and rewards/advantages for Weaver training.

Reference the Agent in a NexRL Recipe

yaml
rollout_worker:
  type: "nexau"
  nexau_agent_config_path: "recipe/my_task/agent_workspace/agent_config.yaml"
  evaluator_module_path: "recipe/my_task/agent_workspace/evaluator.py:MyEvaluator"
  nexau_agent_workspace: "recipe/my_task/agent_workspace"
  task_name: "my_task"

service:
  train_service:
    backend: weaver

Existing Weaver + NexAU examples:

  • NexRL/recipe/nexau_deepsearch/weaver.yaml
  • NexRL/recipe/nexau_news/weaver.yaml

Evaluator Example

An evaluator converts an agent run into a reward. The exact interface depends on the current NexRL version; a common shape looks like this:

python
from typing import Any

from nexrl.rollout_worker import EvaluationRunResult, Evaluator, NexAUEvaluationTarget


class MyEvaluator(Evaluator):
    def evaluate(
        self,
        data: Any,
        evaluation_target: NexAUEvaluationTarget,
    ) -> EvaluationRunResult:
        response = evaluation_target.response
        expected = data.get("answer", "")
        reward = 1.0 if expected and expected in response else 0.0

        return EvaluationRunResult(
            reward=reward,
            extra_info={
                "expected": expected,
                "response": response,
            },
        )

Training Flow

bash
nexrl -m training-service \
  -c recipe/nexau_deepsearch/weaver.yaml \
  --run-nexrl \
  --tag deepsearch-weaver-v1

During training:

  1. NexAU runs tasks and generates tool-use trajectories.
  2. The evaluator assigns rewards.
  3. NexRL aggregates trajectories and computes advantages.
  4. The NexRL trainer calls Weaver forward_backward() and optim_step().
  5. Weaver exports sampler weights, and NexRL syncs them back to rollout services.

Practical Advice

  • Validate agent config, tool schemas, and evaluator logic on a small task set first.
  • Start rewards from verifiable signals such as exact match, unit tests, SQL execution, or retrieval hits.
  • Cap max response length for long-context agents to control rollout cost.
  • Keep debug traces early in training and inspect masks, rewards, and old logprobs.
  • Use Weaver checkpoints for long experiments instead of relying only on short-lived sampler exports.

Next Steps

Weaver API Documentation