search backends

if you already have your data in an external database, or wish to use vector/hybrid search, use one of our supported backends instead of the castform corpus.

all backends implement the same ChunkSource protocol, so switching backends doesn’t change the rest of your pipeline. the qa generation pipeline, default search environment, and training work identically regardless of which backend you choose.

backend	Class	supported search modes
castform	PostgresChunkSource	lexical
turbopuffer	TpufChunkSource	lexical, vector, hybrid
pinecone	PineconeChunkSource	vector
chroma	ChromaChunkSource	lexical, vector, hybrid

---

turbopuffer

turbopuffer supports lexical, vector, and hybrid search natively. you’ll need your own turbopuffer API key.

corpus setup

from castform.rag.corpus.turbopuffer.source import TpufChunkSource

# lexical-only (no embeddings needed)
source = TpufChunkSource(
    api_key="tpuf_...",
    namespace="my-docs",
)
source.populate_from_folder("./docs/")

to enable vector and hybrid search, index with embeddings. chunk sources take a synchronous embed_fn:

from openai import OpenAI

client = OpenAI()

def embed(texts: list[str]) -> list[list[float]]:
    response = client.embeddings.create(model="text-embedding-3-large", input=texts)
    return [item.embedding for item in response.data]

source = TpufChunkSource(
    api_key="tpuf_...",
    namespace="my-docs",
    embed_fn=embed,
)

search client

search clients take an async embed_fn; OpenAIEmbedder is built for that side (it is pickle-safe, so it can ship inside the environment bundle):

import os

from benchmax.auth import StaticBearerAuth
from castform.rag.corpus.embed import OpenAIEmbedder
from castform.rag.corpus.turbopuffer.search import TpufSearch

search = TpufSearch(
    namespace="my-docs",
    embed_fn=OpenAIEmbedder(
        model="text-embedding-3-large",
        base_url="https://api.openai.com/v1",
        auth=StaticBearerAuth(os.environ["OPENAI_API_KEY"]),
    ),
)

TpufSearch reads TPUF_API_KEY when it runs. Pass token_provider= only when the runtime supplies the key through another per-call source.

parameter	default	description
namespace	required	turbopuffer namespace
region	"aws-us-east-1"	turbopuffer region
embed_fn	None	embedding function; required for vector/hybrid
content_attr	None	metadata fields to concatenate as content
token_provider	TPUF_API_KEY	optional per-call key provider

---

pinecone

pinecone provides managed vector search. supports vector search only. you’ll need your own pinecone API key.

corpus setup

from castform.rag.corpus.pinecone.source import PineconeChunkSource

source = PineconeChunkSource(
    api_key="pc_...",
    index_name="my-docs",
)
source.populate_from_folder("./docs/")

by default, pinecone uses its hosted inference API (multilingual-e5-large) for embeddings. to use a custom embedding function:

source = PineconeChunkSource(
    api_key="pc_...",
    index_name="my-docs",
    embed_fn=embedder,
)

for existing indexes with custom metadata field names, use field_mapping:

source = PineconeChunkSource(
    api_key="pc_...",
    index_name="existing-index",
    field_mapping={"content": "body_text", "file_path": "source_file"},
)

search client

from castform.rag.corpus.pinecone.search import PineconeSearch

search = PineconeSearch(
    index_name="my-docs",
)

PineconeSearch reads PINECONE_API_KEY when it runs. Pass token_provider= only when the runtime supplies the key through another per-call source.

parameter	default	description
index_name	required	pinecone index name
index_host	None	direct host URL (skips index lookup)
namespace	""	pinecone namespace
embed_fn	None	custom embedding function
embed_model	"multilingual-e5-large"	hosted inference model (used when no embed_fn)
field_mapping	None	maps custom metadata field names for existing indexes
token_provider	PINECONE_API_KEY	optional per-call key provider

---

chroma

chroma is an open-source embedding database you can self-host. supports vector, lexical (BM25), and hybrid search.

corpus setup

client-server mode is required for training since the model needs network access to the server during remote training.

from castform.rag.corpus.chroma.source import ChromaChunkSource

source = ChromaChunkSource(
    collection_name="my-docs",
    host="chroma.example.com",
    port=8000,
)
source.populate_from_folder("./docs/")

chroma auto-detects which search modes are available:

mode	requires	description
vector	always available	embedding-based similarity search
lexical	BM25 via Chroma Search API	keyword matching
hybrid	both vector + BM25	reciprocal rank fusion

search client

from castform.rag.corpus.chroma.search import ChromaSearch

search = ChromaSearch(
    collection_name="my-docs",
    host="chroma.example.com",
    port=8000,
)

parameter	default	description
collection_name	required	chroma collection name
host	None	self-hosted chroma server hostname (omit for chroma cloud)
port	8000	chroma server port
embed_fn	None	custom embedding function
enable_bm25	True	attempt to use BM25 if available on server
content_attr	None	metadata fields to concatenate as content