Single Agent with Tool Registration

This is the first and simplest example in the Agent framework series. It demonstrates how to stand up a single Dragon agent with an LLM backend, register a sync tool, and execute a task. This example is the foundation for all subsequent examples — read it first to understand the basic lifecycle and configuration.

What you’ll learn:

• How to configure an AgentConfig with minimal required parameters

• How to register a sync tool using registry.register(fn)

• How to set up an inference pipeline and queue

• How to create a Pipeline with a single agent node

• How to launch and run a task through the DAGOrchestrator

Architecture:

The agent process receives a task via Dragon Queue, calls the LLM, uses tools, and returns results via the Scoreboard DDict.

Main Code

Below is the complete example:

Listing 28 01_single_agent.py: Minimal single-agent example

"""01 — Single Agent with Tool Registration.

**Prerequisites:** None — this is the starting point.

**What you'll learn:**

* How to stand up a single Dragon agent with an LLM backend
* How to register a sync tool using ``registry.register(fn)``
* Minimal configuration: ``AgentConfig``, ``ToolRegistry``, ``InferenceConfig``
* The agent lifecycle: create → listen → process → shutdown

Architecture::

    Main Process
      └── Inference.initialize()         ← vLLM engine on 1 GPU
            └── CPUWorker
                  └── InferenceWorker

    planner_agent (Dragon Process)
      └── receives task via Dragon Queue
      └── calls propose_experiment tool
      └── returns result

This is the simplest possible Dragon agent setup.  One agent, one tool,
one inference pipeline.

Usage::

    dragon 01_single_agent.py
"""

import asyncio

import dragon
import multiprocessing as mp

from dragon.ai.agent.core import create_sub_agent
from dragon.ai.agent.config import (
    AgentConfig,
    OrchestratorConfig,
    Pipeline,
    PipelineNode,
)
from dragon.ai.agent.tools import ToolRegistry
from dragon.ai.agent.orchestrator import DAGOrchestrator
from dragon.native.event import Event
from dragon.native.process import Process
from dragon.native.queue import Queue
from dragon.workflows.batch import Batch

from dragon.ai.inference import (
    BatchingConfig,
    HardwareConfig,
    InferenceConfig,
    ModelConfig,
    Inference,
)

# --- Tool implementation ---------------------------------------------------
from tools import propose_experiment


# ===========================================================================
# User-configurable constants
# ===========================================================================

MODEL_NAME = "/path/to/your/model"       # any vLLM-compatible checkpoint
HF_TOKEN = ""                            # set if the model is gated                              # set if the model is gated


# ===========================================================================
# Inference Pipeline Configuration
#
# Minimal setup: 1 node, 1 GPU, 1 inference worker.
# ===========================================================================

INFERENCE_CONFIG = InferenceConfig(
    model=ModelConfig(
        model_name=MODEL_NAME,
        hf_token=HF_TOKEN,
        tp_size=1,                    # single GPU
        max_tokens=8192,
        max_model_len=32768,
    ),
    hardware=HardwareConfig(
        num_nodes=1,
        num_gpus=1,
        num_inf_workers_per_cpu=1,
    ),
    batching=BatchingConfig(
        batch_wait_seconds=0.1,
        max_batch_size=32,
    ),
)


# ===========================================================================
# Tool registry
#
# register() accepts any callable — sync or async.  It auto-wraps it in
# a FunctionTool, deriving name, description, and parameter schemas from
# the function's __name__, __doc__, and type annotations.
# ===========================================================================

registry = ToolRegistry()
registry.register(propose_experiment)


# ===========================================================================
# Pipeline — single node, single agent
# ===========================================================================

pipeline = Pipeline(nodes=[
    PipelineNode(
        agent_id="planner_agent",
        task_description=(
            "You are a scientific experiment planner.  The user wants to "
            "study Monte Carlo convergence for estimating π.\n\n"
            "Propose an experiment plan by calling propose_experiment with:\n"
            "  - description, sample_sizes, convergence_target, methodology\n\n"
            "Report the approved plan verbatim as your final answer."
        ),
        depends_on=[],
    ),
])


# ===========================================================================
# Main
# ===========================================================================

async def main():
    input_queue = Queue()

    print("[startup] Initializing inference pipeline...", flush=True)

    inference_pipeline = None
    try:
        inference_pipeline = Inference(INFERENCE_CONFIG, input_queue)
        inference_pipeline.initialize()
    except Exception as exc:
        import traceback
        print(f"\n[FATAL] Inference pipeline failed to initialize: {exc}", flush=True)
        traceback.print_exc()
        if inference_pipeline is not None:
            inference_pipeline.destroy()
        return
    print("[startup] Inference pipeline ready.\n", flush=True)

    procs, agent_specs = [], []
    try:
        # Create the single agent
        agent_spec = {
            "config": AgentConfig(
                agent_id="planner_agent",
                name="Experiment Planner",
                role=(
                    "You are an experiment planner for Monte Carlo "
                    "convergence studies.  Propose plans via "
                    "propose_experiment."
                ),
                inference_queue=input_queue,
            ),
            "tool_registry": registry,
            "shutdown_event": Event(),
            "reply_queue": Queue(),
        }
        agent_specs = [agent_spec]

        # Launch agent as a Dragon Process
        p = Process(target=create_sub_agent, kwargs=agent_spec)
        p.start()
        procs.append(p)

        # Wait for agent to publish its input queue
        agent_input_queue = agent_spec["reply_queue"].get()
        agent_spec["config"].input_queue = agent_input_queue
        print("[startup] Agent 'planner_agent' ready.", flush=True)

        # Create orchestrator (even for a single agent, it manages
        # the DDict, Batch dispatch, and result collection)
        orchestrator = DAGOrchestrator(
            config=OrchestratorConfig(
                agents=[agent_spec["config"]],
                poll_interval=0.5,
                poll_timeout=120.0,
            ),
            pipeline=pipeline,
        )

        user_input = (
            "Propose a Monte Carlo experiment to estimate π using sample "
            "sizes of 1000, 10000, and 100000."
        )

        batch = Batch()
        try:
            print("=" * 60, flush=True)
            print("Dragon AI — 01 Single Agent", flush=True)
            print("=" * 60, flush=True)
            print(f"Request: {user_input}\n", flush=True)

            result = orchestrator.run(
                user_input=user_input,
                batch=batch,
            )

            print("\n" + "=" * 60, flush=True)
            print("FINAL RESULT", flush=True)
            print("=" * 60, flush=True)
            print(result, flush=True)

        except Exception as exc:
            import traceback
            print(f"\n[error] Pipeline failed: {exc}", flush=True)
            traceback.print_exc()
        finally:
            orchestrator.destroy()
            batch.join()

    except Exception as exc:
        import traceback
        print(f"\n[error] Fatal: {exc}", flush=True)
        traceback.print_exc()
    finally:
        for spec in agent_specs:
            try:
                spec["shutdown_event"].set()
            except Exception:
                pass
        for p in procs:
            try:
                p.join()
            except Exception:
                pass
        print("\n[teardown] Agent stopped.", flush=True)
        try:
            inference_pipeline.destroy()
        except Exception:
            pass
        print("[teardown] Inference pipeline stopped.", flush=True)


if __name__ == "__main__":
    mp.set_start_method("dragon")
    asyncio.run(main())

Key Concepts

Agent Lifecycle:

1. AgentConfig defines the agent’s identity, tools, and inference backend

2. PipelineNode registers the agent in the workflow DAG

3. DAGOrchestrator spawns the agent as a Dragon process

4. Agent listens on its input queue for DispatchHeader messages

5. For each task, agent processes via the LLM + tool loop

6. Results are written to the Scoreboard DDict

7. Agent signals completion via an Event

Tool Registration:

Use registry.register(fn) to register a simple callable. The framework extracts parameter annotations and generates a JSON schema for the LLM.

Inference Setup:

Pass an inference queue to the agent so it can send requests to the vLLM backend. Multiple agents in the same Dragon runtime share this queue, allowing them to submit inference requests to the same GPU worker(s). The backend can run on the same machine or a different node.

Installation

Install Dragon with the ai optional dependency set before running the agent examples that use the shared inference backend:

pip install "dragonhpc[ai]"

For a source checkout, use pip install -e "src[ai]" from the repository root. The inference backend also requires the Dragon vLLM compatibility plugin; see Dragon Inference Service - User Guide for the current plugin installation and vLLM patching steps.

System Description

• For a minimal run: 1 node, 1 GPU (for vLLM), any CPU available

How to Run

Step 1: Edit the model path

Open 01_single_agent.py and set MODEL_NAME to your vLLM-compatible checkpoint (e.g., meta-llama/Llama-2-7b-hf).

Step 2: Set HuggingFace token (if using gated models)

export HF_TOKEN="hf_your_token_here"

Step 3: Run

dragon 01_single_agent.py

Example output:

$ dragon 01_single_agent.py
Agent 'planner' started
Task 'Estimate the convergence rate' received
Calling LLM with tool: estimate_convergence_rate
Result: {'convergence_rate': 0.95, 'confidence': 0.87}
Agent 'planner' completed

Next Steps

Once this example works, proceed to:

• 02 — Multi-Agent DAG (multi-agent orchestration, function nodes, registration styles)

• 03 — Human-in-the-Loop (approval gates before tool execution)

• 04 — Memory Management (history strategies, dedicated summarizer LLM)

• 05 — MCP Tools (integrate remote MCP servers)

• 06 — Full Pipeline (all features combined with tracing)