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
- chunk your data: split your documents into smaller, retrieval-sized pieces (skip this if your data already lives in a vector database)
- upload/link your corpus: index those chunks for search, or connect an existing vector database
- generate question-answer pairs: castform auto-generates synthetic questions and answers grounded in your corpus
- define environment and rewards: configure the search tool the model uses and the reward signals that score its answers
- 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
- chunking: customize how documents are split
- corpus: backend options and setup
- qa generation: how
Pipelineworks, full config reference - search environment: how the rl training environment works