rag/config.py

"""Constants and static configuration, shared by the app and the offline scripts."""

from dataclasses import dataclass
from pathlib import Path

# --- Paths -----------------------------------------------------------------

PROJECT_ROOT = Path(__file__).resolve().parent.parent
DATA_DIR = PROJECT_ROOT / "data"
SOURCES_DIR = DATA_DIR / "sources"
MANIFEST_PATH = DATA_DIR / "manifest.json"
CHUNKS_PATH = DATA_DIR / "chunks.jsonl"
EVAL_DIR = PROJECT_ROOT / "eval"
EVAL_DATASET_PATH = EVAL_DIR / "dataset.json"


# --- Corpus ----------------------------------------------------------------


@dataclass(frozen=True)
class BookSpec:
    """One of the source books that make up the corpus.

    `repo` and `sha` tell us which repository to clone and at which commit, and
    `base_url` is where the published pages live so that citations can link to
    them. We keep a base URL per book rather than deriving one, because
    async-book is published on rust-lang.github.io while the other four are on
    doc.rust-lang.org.

    `exclude_parts` lists mdBook part headings whose chapters we skip. async-book
    keeps its pre-rewrite chapters under "Old chapters", and since those
    contradict the current text we leave them out of the corpus.

    `exclude_paths` lists individual chapters to skip. The Reference already
    turns off search for three generated index pages in its own `book.toml`, and
    they are machine-generated tables rather than prose, so we follow that
    decision.
    """

    title: str
    repo: str
    sha: str
    base_url: str
    license: str
    exclude_parts: tuple[str, ...] = ()
    exclude_paths: tuple[str, ...] = ()


BOOKS: dict[str, BookSpec] = {
    "book": BookSpec(
        title="The Rust Programming Language",
        repo="rust-lang/book",
        sha="917544888a55",
        base_url="https://doc.rust-lang.org/book/",
        license="MIT OR Apache-2.0",
    ),
    "reference": BookSpec(
        title="The Rust Reference",
        repo="rust-lang/reference",
        sha="bec6b5e6631b",
        base_url="https://doc.rust-lang.org/reference/",
        license="MIT OR Apache-2.0",
        exclude_paths=("grammar.md", "syntax-index.md", "test-summary.md"),
    ),
    "rust-by-example": BookSpec(
        title="Rust by Example",
        repo="rust-lang/rust-by-example",
        sha="15308f3e9518",
        base_url="https://doc.rust-lang.org/rust-by-example/",
        license="MIT OR Apache-2.0",
    ),
    "nomicon": BookSpec(
        title="The Rustonomicon",
        repo="rust-lang/nomicon",
        sha="5012a37c682b",
        base_url="https://doc.rust-lang.org/nomicon/",
        license="MIT OR Apache-2.0",
    ),
    "async-book": BookSpec(
        title="Asynchronous Programming in Rust",
        repo="rust-lang/async-book",
        sha="43891cedf954",
        base_url="https://rust-lang.github.io/async-book/",
        license="MIT",
        exclude_parts=("Old chapters",),
    ),
}


# --- Chunking --------------------------------------------------------------

HEADING_SEPARATOR = " › "

EMBED_MODEL = "Qwen/Qwen3-Embedding-0.6B"
EMBED_DIM = 1024

# Sized to stay cheap to retrieve and read while still holding a useful amount of
# information. Counted with the embedding model's own tokenizer, so changing
# EMBED_MODEL re-chunks the corpus, and the index and the eval dataset have to be
# rebuilt with it.
CHUNK_TARGET_TOKENS = 400
CHUNK_MAX_TOKENS = 512

# Headings at or above this level start a new chunk. We use h3 rather than h1
# because the Book's chapter pages open at h2 with their subsections at h3.
CHUNK_HEADING_LEVEL = 3

# Qwen3-Embedding embeds a query and a document differently: its own
# `config_sentence_transformers.json` sets `prompts.query` to this instruction and
# `prompts.document` to an empty string.
#
# LlamaIndex reports `query_instruction = None` while quietly applying an
# instruction of its own. We pass ours explicitly to override that, because the
# two embed the same query differently and so retrieve different passages.
QUERY_INSTRUCTION = (
    "Instruct: Given a web search query, retrieve relevant passages that answer the query\nQuery: "
)


# --- Index -----------------------------------------------------------------

CHROMA_DIR = DATA_DIR / "chroma"
CHROMA_COLLECTION = "rust_docs"

CHROMA_DISTANCE = "cosine"

HF_INDEX_REPO_ENV = "HF_INDEX_REPO"


# --- Retrieval -------------------------------------------------------------

RETRIEVAL_TOP_K = 8

# How the retriever combines its parts. One of "vector", "bm25", "hybrid" or "rerank" 
RETRIEVAL_MODE = "hybrid"

CANDIDATE_TOP_K = 20

RERANK_MODEL = "cross-encoder/ms-marco-MiniLM-L-6-v2"

MAX_SEARCHES_PER_TURN = 3

MAX_AGENT_ITERATIONS = 6

MAX_ANSWER_TOKENS = 4096


# --- Providers -------------------------------------------------------------

@dataclass(frozen=True)
class ProviderSpec:
    label: str
    key_label: str
    models: tuple[str, ...]
    env_var: str

    @property
    def default_model(self) -> str:
        return self.models[0]

PROVIDERS: dict[str, ProviderSpec] = {
    "openai": ProviderSpec(
        label="OpenAI",
        key_label="OpenAI API key",
        models=("gpt-5-mini", "gpt-5"),
        env_var="OPENAI_API_KEY",
    ),
    "gemini": ProviderSpec(
        label="Google Gemini",
        key_label="Google Gemini API key",
        models=("gemini-2.5-flash", "gemini-2.5-pro"),
        env_var="GOOGLE_API_KEY",
    ),
    "anthropic": ProviderSpec(
        label="Anthropic Claude",
        key_label="Anthropic API key",
        models=("claude-haiku-4-5", "claude-sonnet-4-5"),
        env_var="ANTHROPIC_API_KEY",
    ),
}

DEFAULT_PROVIDER = "openai"