rag training

rag Mar 23, 2026 4 min read

retrieval-augmented generation (rag) is when an llm uses a search tool to find relevant information before answering a question. a rag agent can be finetuned to a specific corpus, learning to formulate better search queries and reason more effectively over your data.

all you’ll need to do is link/upload your corpus, and we’ll help you generate the question-answer pairs as training examples, define the environment and rewards to train your model.

the workflow

  1. chunk your data: split your documents into smaller, retrieval-sized pieces (skip this if your data already lives in a vector database)
  2. upload/link your corpus: index those chunks for search, or connect an existing vector database
  3. generate question-answer pairs: castform auto-generates synthetic questions and answers grounded in your corpus
  4. define environment and rewards: configure the search tool the model uses and the reward signals that score its answers
  5. train: launch the run, and the model learns to search your corpus and answer with citations

you can stand up a rag agent in minutes with your coding agent: run castform setup and describe the task. for a worked example, see this run trained to answer policy questions over the gitlab company handbook — its overview walks through the corpus, dataset, and results.

sdk overview

the sdk exposes each stage of the pipeline. if you already have a populated vector database, skip to step 3 and point a corpus backend at it. if you’re starting from raw documents, begin with chunking.

start a local castform workflow by ensuring you have an active session. this is a no-op after castform login or when an explicit platform credential is already configured:

from castform.platform import ensure_session

ensure_session()

1. chunk your data

split your documents into retrieval-sized pieces. built-in chunkers handle markdown and emails, or bring your own.

from castform.rag.chunkers.markdown import MarkdownChunker

chunker = MarkdownChunker(min_char=1024, max_char=2048)
chunks = chunker.chunk_folder("path/to/docs")

see chunking for configuration options and the email chunker.

2. upload to a corpus backend

index your chunks for search. the simplest option is the castform corpus, exposed through PostgresChunkSource:

from castform.rag.corpus.postgres.source import PostgresChunkSource

source = PostgresChunkSource(
    corpus_name="my-docs",
)
await source.populate_from_chunks(chunks)

already have your data in an external vector db? skip chunking and upload, and use that backend directly. see corpus.

3. generate qa pairs

the pipeline generates synthetic question-answer pairs grounded in your corpus:

from castform.rag.qa_generation.pipeline_config import PipelineConfig, PlatformConfig, CorpusConfig, TargetsConfig
from castform.rag.qa_generation.pipeline import Pipeline

cfg = PipelineConfig(
    platform=PlatformConfig(),
    corpus=CorpusConfig(corpus_name="my-docs", corpus_id=source.corpus_id),
    targets=TargetsConfig(total_samples=200),
)

pipeline = Pipeline(cfg)
result = pipeline.run()

train_data = result["train_dataset"]
eval_data = result["eval_dataset"]

see qa generation for the full config reference.

4. define the environment and launch

SearchEnv is defined in the postgres-search example’s main.py (examples/postgres-search in the benchmax repo). run this from within that example directory:

from main import SearchEnv

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
import dataclasses

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

bundle = dump_bundle(
    SearchEnv,
    constructor_args={
        "search": search,
        "judge_base_url": config.llm_url(),
        "judge_model": "gpt-5.4-mini",
    },
)

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

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

see search environment for reward configuration and custom search backends. see launching for upload_assets() parameters.

next steps