Corrective RAG

Grade and fall back

Implement the Corrective RAG pattern in this codebase, using the LLM client and language we already use.

## Goal

Answer a question from an internal index, but grade the retrieved documents first, fall back to a rewritten web search when none is relevant, and check the answer against its documents.

## How it works

1. Retrieve the top 4 documents for the question from the internal index.
2. Grade each document with the Document grader prompt, parsing its JSON into relevant and reason, and keep the relevant ones.
3. If none is relevant, call the Query rewriter prompt with the question and the rejection reasons, run its query on web search, and use the web results as the documents.
4. Generate an answer from the documents, then run the Answer check prompt on it and parse the list of unsupported claims.
5. If any claim is unsupported, regenerate with that list, telling the model to drop or fix those claims, and check again.

## Prompts

### Document grader prompt

```markdown
You grade whether a retrieved document helps answer a question.

<question>{{question}}</question>
<document>{{doc}}</document>

Respond in JSON: {"relevant": true|false, "reason": "one sentence"}
Mark relevant only if the document contains facts needed for the answer.
A document on the same topic that does not answer the question is not relevant.
```

### Query rewriter prompt

```markdown
The internal search returned nothing useful for this question.

<question>{{question}}</question>
<rejected_reasons>{{reasons}}</rejected_reasons>

Write one web search query that would find the missing facts. Use specific
terms and names, not the full question. Return only the query.
```

### Answer check prompt

```markdown
Check whether every claim in the answer is supported by the documents.

<documents>{{docs}}</documents>
<answer>{{answer}}</answer>

Respond in JSON:
{"supported": true|false, "unsupported_claims": ["claim", "..."]}
```

## Stop rule

Stop when the check finds every claim supported, or after 2 regenerations. At the cap, return the last answer with its unsupported claims removed, and a status saying the check did not fully pass.

## Requirements

- Make the documents retrieved (default 4) and the regeneration cap (default 2) parameters.
- If a grader reply is not valid JSON, treat that document as not relevant and log the raw reply.
- If web search fails or returns nothing, reply that the answer could not be found rather than generating from rejected documents.
- Log every grade with its reason, the rewritten query and each check result.
- Add a test with a stub grader that rejects all 4 documents and a stub web search; assert the rewriter ran once, web search ran once, and the generator received only the web documents.

Self-RAG

Implement the Corrective RAG pattern (Self-RAG variant) in this codebase, using the LLM client and language we already use.

## Goal

Let one model decide on demand whether to retrieve, judge each passage's relevance and its own output's support with reflection tokens, and keep the best-scored continuation.

## How it works

1. Use a model trained, or prompted with the Self-reflective answer prompt, to emit reflection tokens: Retrieve (yes, no, or continue to reuse the current passages), IsRel (relevant or irrelevant), IsSup (fully, partially or not supported) and IsUse (1 to 5).
2. Generate the next segment. If the model emits Retrieve=yes, search the index with the question plus the answer so far and take the top 4 passages.
3. For each passage, generate a candidate segment in parallel and read the IsRel and IsSup tokens the model emits for it.
4. Score each candidate as a weighted sum of its IsRel, IsSup and IsUse tokens, keep the best one and append it to the answer (the paper keeps the top 2 with a beam search).
5. When the model emits Retrieve=no, generate the segment without passages.
6. If a candidate's reflection tokens cannot be parsed, give it the lowest score instead of failing the request.

## Prompts

### Self-reflective answer prompt

```markdown
Answer the question one short segment at a time.

<question>{{question}}</question>
<answer_so_far>{{answer}}</answer_so_far>
{{#if passage}}<passage>{{passage}}</passage>{{/if}}

Before each segment, write [Retrieve=Yes] if you need a new document,
[Retrieve=Continue] to keep using the current passage, else [Retrieve=No].
If a passage is given, write [IsRel=Relevant] or [IsRel=Irrelevant] for it.
If a passage was given, after the segment write [IsSup=Fully], [IsSup=Partially]
or [IsSup=No] for how well it supports the segment.
Finally write [IsUse=1..5] for the segment's usefulness.
```

## Stop rule

Stop when the model ends its answer, or after 6 segments. Return the answer with each segment's IsSup tag, so partly supported text can be shown as such.

## Requirements

- Make the segment cap (default 6), the passages per retrieval (default 4) and the score weights (defaults IsRel 1.0, IsSup 1.0, IsUse 0.5) parameters.
- Expose a retrieval threshold: retrieve when the probability of Retrieve=yes is above it (default 0.2), where the model exposes token probabilities.
- Log every segment's candidates with their reflection tokens and scores.
- Add a test with a stub model that emits Retrieve=no for a rewording task; assert the retriever is never called and one segment is returned.

Groundedness loop

Implement the Corrective RAG pattern (Groundedness loop variant) in this codebase, using the LLM client and language we already use.

## Goal

Answer from retrieved documents, then grade the answer itself: regenerate when it is not grounded in the documents, and rewrite the query and retrieve again when it does not answer the question.

## How it works

1. Retrieve the top 4 documents for the question and generate an answer from them.
2. Grade the answer with the Answer grader prompt, parsing its JSON into grounded, answers_question and unsupported claims.
3. If both checks are true, return the answer.
4. If grounded is false, regenerate from the same documents, passing the unsupported claims to drop or fix.
5. If grounded is true but answers_question is false, call the Query rewriter prompt for a new query, retrieve again and regenerate.
6. If the grader's JSON does not parse, retry once; if it fails again, treat the answer as not grounded.

## Prompts

### Answer grader prompt

```markdown
Grade an answer against the documents it was written from and the question.

<question>{{question}}</question>
<documents>{{docs}}</documents>
<answer>{{answer}}</answer>

Respond in JSON:
{"grounded": true|false, "answers_question": true|false,
 "unsupported_claims": ["claim", "..."]}
grounded is true only if every claim is stated in the documents.
```

### Query rewriter prompt

```markdown
The answer below did not address the question, so the search missed.

<question>{{question}}</question>
<answer>{{answer}}</answer>

Write one search query that would find the missing facts. Use specific
terms and names, not the full question. Return only the query.
```

## Stop rule

Stop when both checks pass, or after 3 loops in total, regenerations and re-retrievals together. At the cap, reply that no supported answer was found and return the last documents.

## Requirements

- Make the total loop cap (default 3) and the documents per retrieval (default 4) parameters.
- Cap re-retrievals separately (default 2), so a poor index cannot use up every loop. If a retrieval fails or returns nothing, reply that no supported answer was found.
- Log each grade with the loop it triggered.
- Add a test with a stub grader that returns answers_question=false once and then passes both checks, and a stub rewriter that returns 'payroll vendor contract'; assert two retrievals, the second with the query 'payroll vendor contract', and that the second answer is returned.