search environment

an environment defines the tools your model can use and the reward signals for training. for rag, this means a search tool over your corpus and a reward function that checks if the model retrieved the right information.

how it works

we provide a default environment, SearchEnv, that comes with a default search tool implementation, system prompt, and rewards. within each rollout,

• the model is given a prompt from the generated qa dataset, formatted to also include a system prompt that instructs them on how to use the search tool to answer the question and cite their sources.
• the model can use the search tool up to a certain number of times to query your corpus. it works across any of our supported search backends.
• then, the environment scores the answer on four components:

component	what it measures
answer_correctness	an llm judge compares the answer to the ground truth. the gate: the components below are multiplied by it
retrieval_hit	fraction of gold sources cited in the answer. ungated, so citing gold counts even when the answer is wrong
citation_precision	of the sources it cited, how many are gold (gated on correctness)
answer_length	deterministic brevity score (gated on correctness)

only answer_correctness calls the llm judge (which is why SearchEnv requires a judge_model); the rest are deterministic. each component is weighted; you can retune the weights below. the earlier conciseness-judge and search-efficiency scorers remain as opt-in helpers for subclasses that override compute_reward.

quickstart

the only code you write is a small subclass that sets the system prompt for your corpus. SearchEnv is defined in the postgres-search example’s main.py (examples/postgres-search in the benchmax repo), so run this from within that example directory. then pass your corpus and judge through constructor_args and launch:

from main import SearchEnv

from benchmax.auth import InjectedAuth
from benchmax.bundle import dump_bundle
from castform import config
from castform.rag.corpus.postgres.search import PostgresSearch
from castform.platform.environment_assets import upload_assets
from castform.platform.client import TrainerClient
from castform.platform import ensure_session
import dataclasses

MAX_SEARCH_CALLS = 10
ensure_session()

class MySearchEnv(SearchEnv):
    # the system prompt tells the model how to search, answer, and cite
    system_prompt = SearchEnv.render_system_prompt(
        corpus_description="acme's internal support docs",
        max_search_calls=MAX_SEARCH_CALLS,
    )

search = PostgresSearch(
    corpus_name="my-docs",
    base_url=config.platform_url(),
)

bundle = dump_bundle(
    MySearchEnv,
    constructor_args={
        "search": search,
        "judge_base_url": config.llm_url(),
        "judge_model": "gpt-5.4-mini",
        "judge_auth": InjectedAuth("judge"),
        "max_search_calls": MAX_SEARCH_CALLS,
    },
)

uploaded = upload_assets(
    bundle=bundle,
    train_dataset=train_data,
    eval_dataset=eval_data,
    run_name="my-search-model",
)

trainer = TrainerClient()
run_id = trainer.launch_training_run(
    **dataclasses.asdict(uploaded),
)

constructor_args is the dict castform passes to SearchEnv.__init__ when the run starts on the trainer; that’s how the environment receives your search client, judge, and settings.

PostgresSearch resolves its corpus credential at request time: the local session during validation and the trainer’s rotating act-as token during training. InjectedAuth("judge") works the same way for the managed judge, keeping both credentials out of the serialized bundle. for your own external judge endpoint, pass StaticBearerAuth(provider_api_key) explicitly; that static key is then part of the bundle.

SearchEnv expects your dataset rows to have question, answer, and reference_chunks columns, matching exactly what qa generation produces.

basic customization

you configure SearchEnv through constructor_args and the system prompt on your subclass. no rewriting the environment.

the search tool

the search tool runs against whatever search client you pass as search. swap the client to change backend or search modes; the prompt and rewards stay the same.

client	import	modes
PostgresSearch	castform.rag.corpus.postgres.search	lexical
TpufSearch	castform.rag.corpus.turbopuffer.search	lexical, vector, hybrid
PineconeSearch	castform.rag.corpus.pinecone.search	vector
ChromaSearch	castform.rag.corpus.chroma.search	vector, lexical, hybrid

when a backend supports more than one mode, the model chooses per query; with "auto" (the default) SearchEnv picks hybrid > lexical > vector. see corpus backends for setup.

reward weights

every component has a weight you can raise, lower, or zero out. set any of these in constructor_args:

arg	default	component
w_correctness	1.0	answer correctness
w_retrieval_hit	0.3	retrieval hit
w_citation_precision	0.3	citation precision
w_length	0.2	answer length

set a weight to 0 to drop that component. for example, if citations don’t matter for your task, set w_retrieval_hit and w_citation_precision to 0.

search budget

max_search_calls (default 10) caps how many times the model may call search; a rollout that exceeds it scores 0. set it in both render_system_prompt (so the prompt states the right budget) and constructor_args (so the environment enforces it), as shown in the quickstart.

the system prompt

the system prompt is what makes the rest work: it instructs the model to reason in <think>, call search, and put its final answer in <answer>...</answer> with [Source: <id>] citations, both of which the reward parses. render_system_prompt fills your corpus_description and max_search_calls into a default template.

to change the wording, override SYSTEM_PROMPT_TEMPLATE on your subclass. keep the {corpus_description} and {max_search_calls} placeholders, and the <answer> / [Source: ...] conventions the reward depends on:

class MySearchEnv(SearchEnv):
    SYSTEM_PROMPT_TEMPLATE = """\
Answer questions about {corpus_description} using the search tool.
You may search up to {max_search_calls} times.
Put your final answer in <answer>...</answer> and cite sources as [Source: <id>].
"""
    system_prompt = SearchEnv.render_system_prompt(
        corpus_description="acme's support docs",
        max_search_calls=10,
    )

advanced customization

custom search backend

the built-in clients implement a small SearchClient protocol. implement it yourself to search any store; you pass your client as search with no subclassing of SearchEnv needed.

from castform.rag.corpus.search_client import SearchClient

class SearchClient(Protocol):
    async def search(self, query: str, mode: str = "auto", top_k: int = 10) -> list[dict[str, Any]]
    async def embed(self, text: str) -> list[float] | None

    @property
    def available_modes(self) -> list[str]

    def get_params(self) -> dict[str, Any]

• search() returns dicts with content, source, metadata, and score
• available_modes reports which of lexical / vector / hybrid the backend supports
• get_params() returns serializable connection parameters

the client must be pickle-safe: store connection parameters and reconstruct sdk clients lazily after unpickling, since it’s shipped to remote workers.

from castform.rag.corpus.search_client import SearchClient

class MySearch:
    def __init__(self, endpoint: str, api_key: str):
        self._endpoint = endpoint
        self._api_key = api_key

    async def search(self, query: str, mode: str = "auto", top_k: int = 10) -> list[dict[str, Any]]:
        # return dicts with content, source, metadata, score
        ...

    async def embed(self, text: str) -> list[float] | None:
        return None  # optional

    @property
    def available_modes(self) -> list[str]:
        return ["lexical"]

    def get_params(self) -> dict[str, Any]:
        return {"endpoint": self._endpoint, "api_key": self._api_key}

# pass it as `search` — no SearchEnv subclass needed for the backend
constructor_args = {
    "search": MySearch("https://...", "key"),
    "judge_base_url": "https://api.openai.com/v1",
    "judge_model": "gpt-5.4-mini",
}

corpus-specific citations

citations are matched by document id. by default SearchEnv reads the file (or file_path) metadata key from your reference chunks and from each [Source: <id>] tag. if your sources use a different id, override _extract_reference_ids or _canonicalize_id on your subclass.

fully custom rewards

to change scoring beyond the weights, override compute_reward on your subclass, or build an environment from scratch. see writing your own environment.

next steps

• see qa generation to generate the question / answer / reference_chunks dataset over your corpus
• see launching a training run to launch a training job using your environment and dataset
• see corpus backends for setting up turbopuffer, pinecone, or chroma