# Corrective RAG (CRAG) Pipeline Documentation

This document outlines the architecture, execution workflow, node contracts, and core configuration settings for the Corrective RAG (`CRAG`) engine in Study Buddy.

---

## 1. Pipeline Graph

<div align="center">
  <a href="crag_graph.png" target="_blank" title="Click to view in full resolution">
    <img width="2717" height="4096" alt="crag_graph" src="https://github.com/user-attachments/assets/b938e601-bba1-4618-9b10-248157df44d2" />
  </a>
  <br>
  <sub><em>(Click graph to open full resolution)</em></sub>
</div>

<details>
<summary>View raw Mermaid specification</summary>

```mermaid
graph TD;
	__start__([__start__])
	triage(triage)
	reply_chitchat(reply_chitchat)
	reply_refusal(reply_refusal)
	init_state(init_state)
	router(router)
	retrieve(retrieve)
	normalize(normalize)
	merge_pool(merge_pool)
	rerank(rerank)
	evaluator(evaluator)
	retry_check(retry_check)
	generate(generate)
	fallback(fallback)
	__end__([__end__])

	__start__ --> triage;
	triage -.->|chitchat| reply_chitchat;
	triage -.->|unsafe| reply_refusal;
	triage -.->|study query| init_state;
	reply_chitchat --> __end__;
	reply_refusal --> __end__;
	init_state --> router;
	router --> retrieve;
	retrieve --> normalize;
	normalize --> merge_pool;
	merge_pool --> rerank;
	rerank -.->|top_score >= 0.65| generate;
	rerank -.->|0.30 < score < 0.65| evaluator;
	rerank -.->|top_score <= 0.30| retry_check;
	evaluator -.->|correct| generate;
	evaluator -.->|ambiguous / incorrect| retry_check;
	retry_check -.->|retries <= 2| router;
	retry_check -.->|retries > 2| fallback;
	generate --> __end__;
	fallback --> __end__;
```

</details>

---

## 2. End-to-End Workflow

The pipeline processes user queries through three distinct pathways:

1. **Zero-Retrieval Fast Paths (Bypass)**:
   - Evaluated at `triage` via Tier 0 (regex/rules, `<1ms`) and Tier 1 (fast LLM classification, line-based output).
   - If `chitchat`: routes directly to `reply_chitchat` with localized canned greetings/thanks and exits.
   - If `unsafe`: routes directly to `reply_refusal` with policy refusals (gibberish, NSFW, politics, religion) and exits.

2. **Standard RAG Pipeline (Happy Path)**:
   - Initialized at `init_state` (passes `intent` and sub-queries, sets `retries=0`).
   - `router` selects search sources and formulates query plans (defaults to curated Supabase materials on pass 0).
   - `retrieve` executes external tool queries (Supabase vector store, Wikipedia, ArXiv, Web).
   - `normalize` cleans text into structured `Chunk` objects and computes SHA-256 content IDs.
   - `merge_pool` deduplicates passages and merges them into the cumulative pool.
   - If top reranker score $\ge 0.65$ (`HIGH_THRESHOLD`) on simple queries, `score_gate` short-circuits directly to `generate` (skipping the evaluator LLM call).
   - `generate` drafts an educational answer strictly cited with chunk numbers `[n]`.

3. **Evaluation & Corrective Feedback Loop**:
   - For simple queries with score between $0.30$ and $0.65$, or complex queries with score $> 0.30$, `evaluator` inspects relevance and coverage (`correct`, `ambiguous`, `incorrect`).
   - If rerank score is $\le 0.30$ (`LOW_THRESHOLD`), `score_gate` routes directly to `retry_check`.
   - If evaluated as `correct`, proceeds to `generate`.
   - If `ambiguous` or `incorrect`, `retry_check` tracks attempt count.
   - If `retries <= MAX_RETRIES (2)`, loops back to `router` to re-query (`bad_query` rewrites query targeting `missing`; sources remain fixed with no tool escalation).
   - If retries are exhausted, routes to `fallback` (synthesizes hybrid response or general-knowledge answer with `low_confidence=True`).

---

## 3. Node Specifications & Outputs

### 1. `triage`
- **Purpose**: Classifies user input into intent (`chitchat`, `unsafe`, `simple`, `complex`) and decomposes complex queries into 1–3 sub-queries. Runs Tier 0 regex heuristics first, falling back to a structured line-based LLM call.
- **Output fields**:
  - `intent`: `"chitchat" | "unsafe" | "simple" | "complex"`
  - `sub_questions`: `list[str]` (non-empty only for complex)
  - `reply`: canned string for chitchat/unsafe, otherwise `None`

### 2. `reply_chitchat`
- **Purpose**: Emits immediate conversational response in English or Arabic without touching vector stores or tools.
- **Output fields**:
  - `answer`: localized greeting/thanks message
  - `citations`: `[]`
  - `low_confidence`: `False`

### 3. `reply_refusal`
- **Purpose**: Emits standardized safety refusal for out-of-scope queries (gibberish, adult, politics, religion).
- **Output fields**:
  - `answer`: standardized refusal string
  - `citations`: `[]`
  - `low_confidence`: `False`

### 4. `init_state`
- **Purpose**: Initializes graph state variables for legitimate study queries.
- **Output fields**:
  - `original_q`: validated question string
  - `sources_to_use`: `["supabase"]` (default first pass)
  - `retries`: `0`, `pool`: `[]`, `ranked`: `[]`, `verdict`: `None`

### 5. `router`
- **Purpose**: Selects target retrieval sources and constructs retrieval queries.
  - **Pass 0**: If triage `intent == "complex"`, queries **all tools concurrently** (`supabase`, `arxiv`, `wikipedia`, `web`) with zero LLM delay. If `intent == "simple"`, runs **automatic tool calling** via LangChain agentic tool selection.
  - **Retries**: Keeps existing sources fixed (no adding tools / no tool escalation); strictly performs **requerying** by rewriting queries targeting evaluator `missing` feedback.
- **Output fields**:
  - `sources_to_use`: list of sources (`["supabase", "arxiv", "wikipedia", "web"]`)
  - `search_query`: optimized search keywords
  - `tool_queries`: mapping of source to specific search query

### 6. `retrieve`
- **Purpose**: Calls registered tool functions concurrently across designated sources with intent- and retry-aware dynamic timeouts:
  - **Base timeout**: `7.0s` (`TOOL_TIMEOUT_SIMPLE_S`) for simple queries, `13.0s` (`TOOL_TIMEOUT_COMPLEX_S`) for complex queries.
  - **Progressive timeout escalation for empty context (`context == []`)**:
    - `retries == 1`: tool timeout escalates to `16.0s` (`TOOL_TIMEOUT_RETRY1_EMPTY_S`).
    - `retries == 2`: tool timeout escalates to `25.0s` (`TOOL_TIMEOUT_RETRY2_EMPTY_S`).
- **Output fields**:
  - `raw_results`: raw documents/payloads returned by tool APIs
  - `all_retrievers_failed`: boolean indicating whether all tool retrievers failed/timed out

### 7. `normalize`
- **Purpose**: Parses varied tool response formats, splits text into target chunk lengths, and computes deterministic content hashes.
- **Output fields**:
  - `new_chunks`: list of validated `Chunk` instances
  - `tried_sources`: updated cumulative list of queried sources

### 8. `merge_pool`
- **Purpose**: Accumulates incoming chunks into the existing pool, preserves scores, and ensures top-tier document slots are retained without unbounded growth.
- **Output fields**:
  - `pool`: deduped, bounded chunk list (capped at `POOL_MAX`)

### 9. `rerank`
- **Purpose**: Scores chunks against `original_q` using a cross-encoder model. Retains top Supabase chunks to preserve course material grounding.
- **Output fields**:
  - `ranked`: top candidate chunks sorted by descending score (length `FINAL_CONTEXT_N`)

### 10. `evaluator`
- **Purpose**: Fast LLM judge checking context relevance and completeness against the question. Emits line-based key-values (`VERDICT`, `CAUSE`, `MISSING`). The only cause for retrying is `bad_query` (requerying).
- **Output fields**:
  - `verdict`: `"correct" | "ambiguous" | "incorrect"`
  - `likely_cause`: `"none" | "bad_query"`
  - `missing`: description of missing factual concepts

### 11. `retry_check`
- **Purpose**: Increments retry counter and records diagnostic logs.
- **Output fields**:
  - `retries`: updated integer counter (`retries + 1`)

### 12. `generate`
- **Purpose**: Synthesizes the final educational answer strictly grounded in ranked context chunks cleanly without bracketed numbers (`[1]`, `[2]`).
- **Output fields**:
  - `answer`: synthesized text
  - `citations`: list of cited source titles/URLs
  - `low_confidence`: `False`

### 13. `fallback`
- **Purpose**: Failsafe generator triggered when retries are exhausted or retrieval fails. Generates hybrid context-assisted or knowledge-based explanations.
- **Output fields**:
  - `answer`: fallback educational text
  - `citations`: list of citations (if partial context was used)
  - `low_confidence`: `True`

---

## 4. Conditional Edge Gates

- **`triage_gate`**:
  - `intent == "chitchat"` $\rightarrow$ `reply_chitchat`
  - `intent == "unsafe"` $\rightarrow$ `reply_refusal`
  - default $\rightarrow$ `init_state`

- **`score_gate`**:
  - **For Simple Queries (`intent == "simple"`)**:
    - `top_score >= HIGH_THRESHOLD (0.65)` $\rightarrow$ `generate` (high confidence single-concept shortcut, skips evaluator).
    - `LOW_THRESHOLD (0.30) < top_score < HIGH_THRESHOLD (0.65)` $\rightarrow$ `evaluator` (ambiguous zone, LLM verifies relevance).
    - `top_score <= LOW_THRESHOLD (0.30)` $\rightarrow$ `retry_check` $\rightarrow$ retries with **all tools** (`supabase`, `arxiv`, `wikipedia`, `web`) using keyword-style search queries.
  - **For Complex Queries (`intent == "complex"`)**:
    - `top_score <= LOW_THRESHOLD (0.30)` $\rightarrow$ `retry_check` (immediate retry, context too weak).
    - `top_score > LOW_THRESHOLD (0.30)` $\rightarrow$ `evaluator` (ALWAYS verifies that ALL sub-parts are satisfied).

- **`evaluator_gate`**:
  - `verdict == "correct"` $\rightarrow$ `generate`
  - `verdict in {"ambiguous", "incorrect"}` $\rightarrow$ `retry_check`

- **`retry_check_gate`**:
  - `retries <= MAX_RETRIES (2)` $\rightarrow$ `router` (retry requerying missing facts)
  - `retries > MAX_RETRIES (2)` $\rightarrow$ `fallback` (exhausted retries)

---

## 5. Model Architecture & Tier Assignment

CRAG employs a high-efficiency 8B / 14B split tailored to the specific latency and reasoning requirements of each graph node, powered directly by Mistral AI:

| Node / Role | Model Tier | Why |
| :--- | :--- | :--- |
| **triage, reply_chitchat, reply_refusal, router** (re-query formatting) | `ministral-8b-latest` | Fast, lightweight classification calls where low latency (<1s) matters more than heavy parameter capacity. |
| **tool_selector** (agentic tool routing in router/retrieve), **evaluator**, **generate**, **fallback** | `ministral-14b-latest` | 14-billion parameter reasoning model with native tool-calling, Pydantic structured output, and high-fidelity cited generation. |

---

## 6. Key Configuration & Constants

All runtime thresholds and parameters are defined in `src/rag/constants.py` and validated via Pydantic in `src/rag/config.py`:

| Parameter | Default | Purpose / Behavior |
| :--- | :--- | :--- |
| `HIGH_THRESHOLD` | `0.65` | Quality cutoff. If top rerank score $\ge 0.65$ on simple queries, skips evaluator and answers immediately. |
| `LOW_THRESHOLD` | `0.30` | Minimum viability floor. If top score $\le 0.30$, skips evaluator and triggers corrective search. |
| `KEEP_THRESHOLD` | `0.5` | Threshold for carrying over candidate chunks into subsequent retry pools. |
| `MAX_RETRIES` | `2` | Maximum corrective search iterations before falling back to graceful synthesis. |
| `TOP_K` | `5` | Base number of chunks retrieved per tool call. |
| `RETRIEVE_MULTIPLIER`| `3` | Over-retrieval factor (retrieves `TOP_K * MULTIPLIER` candidates for reranker pool). |
| `FINAL_CONTEXT_N` | `8` | Maximum number of top reranked chunks passed into generation prompt context. |
| `POOL_MAX` | `50` | Hard cap on total deduped candidate chunks retained in pool across passes. |
| `CARRY_MAX` | `4` | Maximum high-scoring chunks carried forward across retry iterations. |
| `TRIAGE_TIMEOUT_S` | `15.0` | Max timeout allowed for Tier 1 triage classification before failing to complex path. |
| `TOOL_TIMEOUT_SIMPLE_S` | `7.0` | Default timeout for tool retrieval calls on simple queries. |
| `TOOL_TIMEOUT_COMPLEX_S` | `13.0` | Default timeout for tool retrieval calls on complex queries. |
| `TOOL_TIMEOUT_RETRY1_EMPTY_S` | `16.0` | Escalated tool timeout on retry 1 when context pool is empty. |
| `TOOL_TIMEOUT_RETRY2_EMPTY_S` | `25.0` | Escalated tool timeout on retry 2 when context pool remains empty. |
| `MAX_SUBQUERIES` | `3` | Upper bound on sub-queries generated for complex multi-part questions. |
| `TRIAGE_MODEL` | `ministral-8b-latest` | Fast classifier for intent triage and canned chitchat/refusal. |
| `ROUTER_MODEL` | `ministral-8b-latest` | Fast model for router re-query keyword formulation. |
| `TOOL_SELECTOR_MODEL`| `ministral-14b-latest` | Mistral 14B for high-fidelity agentic tool selection. |
| `EVALUATOR_MODEL` | `ministral-14b-latest` | Mistral 14B for structured context evaluation. |
| `GENERATOR_MODEL` | `ministral-14b-latest` | Mistral 14B for rigorous educational answer generation. |
| `FALLBACK_MODEL` | `ministral-14b-latest` | Mistral 14B for graceful fallback answer generation. |
| `RERANKER_MODEL` | `jina-reranker-v2-base-multilingual` | Via Jina API Key (Tier 1 cascade). |