Search Agent Example

cd examples/search
python run_search_agent.py

import asyncio
import os
from dotenv import load_dotenv
from local_retrieval_tool import LocalRetrievalTool
from transformers import AutoTokenizer
from rllm.agents.system_prompts import SEARCH_SYSTEM_PROMPT
from rllm.agents.tool_agent import ToolAgent
from rllm.data.dataset import DatasetRegistry
from rllm.engine.agent_execution_engine import AgentExecutionEngine
from rllm.environments.tools.tool_env import ToolEnvironment
from rllm.rewards.reward_fn import search_reward_fn
from rllm.utils import save_trajectories

os.environ["TOKENIZERS_PARALLELISM"] = "true"
if "RETRIEVAL_SERVER_URL" not in os.environ:
    os.environ["RETRIEVAL_SERVER_URL"] = "http://127.0.0.1:8000"

load_dotenv()

n_parallel_agents = 64
model_name = "Qwen/Qwen3-4B"

tokenizer = AutoTokenizer.from_pretrained(model_name)
sampling_params = {"temperature": 0.6, "top_p": 0.95, "model": model_name}

# Configure custom search tool
tool_map = {"local_search": LocalRetrievalTool}

engine = AgentExecutionEngine(
    agent_class=ToolAgent,
    agent_args={
        "tool_map": tool_map,
        "system_prompt": SEARCH_SYSTEM_PROMPT,
        "parser_name": "qwen"
    },
    env_class=ToolEnvironment,
    env_args={
        "tool_map": tool_map,
        "reward_fn": search_reward_fn
    },
    rollout_engine=None,
    engine_name="openai",
    tokenizer=tokenizer,
    sampling_params=sampling_params,
    rollout_engine_args={
        "base_url": "http://localhost:30000/v1",
        "api_key": "None",
    },
    max_response_length=16384,
    max_prompt_length=4096,
    config=None,
    n_parallel_agents=n_parallel_agents,
)

test_dataset = DatasetRegistry.load_dataset("hotpotqa", "test")
tasks = test_dataset.get_data()

results = asyncio.run(engine.execute_tasks(tasks))
save_trajectories(results, filename="search_trajectories.pt")

from rllm.environments.tools.base_tool import BaseTool
import requests

class LocalRetrievalTool(BaseTool):
    """Tool for searching local Wikipedia corpus."""
    
    def __init__(self, server_url: str = "http://127.0.0.1:8000", max_results: int = 5):
        self.server_url = server_url
        self.max_results = max_results
    
    def __call__(self, query: str) -> str:
        """Search for documents matching the query."""
        response = requests.post(
            f"{self.server_url}/search",
            json={"query": query, "k": self.max_results},
            timeout=30
        )
        results = response.json()["results"]
        
        # Format results for the agent
        formatted = ""
        for i, doc in enumerate(results, 1):
            formatted += f"\n[Document {i}]\n{doc['text']}\n"
        return formatted
    
    @property
    def name(self) -> str:
        return "local_search"
    
    @property
    def description(self) -> str:
        return "Search Wikipedia for information. Input: search query string."

bash examples/search/train_search_agent.sh

SEARCH_SYSTEM_PROMPT = """You are a helpful AI assistant with access to a search tool.

When answering questions:
1. Use the search tool to find relevant information
2. You can search multiple times to gather information
3. Reason step-by-step based on the search results
4. Put your final answer in \\boxed{} format

Example:
Question: What is the capital of France?

You should:
1. Search for "capital of France"
2. Review the results
3. Provide answer: \\boxed{Paris}
"""

Question: "In what year was the director of Pulp Fiction born?"

Agent trajectory:
1. Search: "director of Pulp Fiction"
   → Result: "Quentin Tarantino"
2. Search: "Quentin Tarantino birth year"
   → Result: "1963"
3. Answer: \boxed{1963}