Skip to content

RAG types

Base class for every RAG type.

To add one, inherit from BaseRAG and: 1. describe its steps in FLOW (the docs playground draws this diagram) 2. implement run(), recording every step with a Trace answer() then works for free.

BaseRAG

Common interface for RAG types.

FLOW lists the steps and the arrows between them

{"nodes": [{"id": "retrieve", "label": "Retrieve", "kind": "search"}, ...], "edges": [["retrieve", "rerank"], ...]}

Every step that run() records must be a node, and two steps in a row must be joined by an edge.

Source code in src/rag_titan/rag_types/base_rag.py 
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
class BaseRAG:
    """
    Common interface for RAG types.

    FLOW lists the steps and the arrows between them:
        {"nodes": [{"id": "retrieve", "label": "Retrieve", "kind": "search"}, ...],
         "edges": [["retrieve", "rerank"], ...]}
    Every step that run() records must be a node, and two steps in a row
    must be joined by an edge.
    """

    FLOW = {"nodes": [], "edges": []}

    def run(self, query):
        """
        Answer the question and show how.

        Returns:
            A dict with "answer" (the text), "sources" (the chunks used for
            the answer) and "steps" (one {"step", "input", "output", "ms"}
            dict per step, in order).
        """
        raise NotImplementedError

    def answer(self, query):
        """Return only the answer text."""
        return self.run(query)["answer"]

answer(query)

Return only the answer text.

Source code in src/rag_titan/rag_types/base_rag.py 
38
39
40
def answer(self, query):
    """Return only the answer text."""
    return self.run(query)["answer"]

run(query)

Answer the question and show how.

Returns:

Type Description

A dict with "answer" (the text), "sources" (the chunks used for

the answer) and "steps" (one {"step", "input", "output", "ms"}

dict per step, in order).
Source code in src/rag_titan/rag_types/base_rag.py 
27
28
29
30
31
32
33
34
35
36
def run(self, query):
    """
    Answer the question and show how.

    Returns:
        A dict with "answer" (the text), "sources" (the chunks used for
        the answer) and "steps" (one {"step", "input", "output", "ms"}
        dict per step, in order).
    """
    raise NotImplementedError

Records every step a RAG type takes, so you can see (and draw) how an answer was made.

trace = Trace()
matches = trace.record("retrieve", retriever.retrieve, query="leave?", top_k=4)
print(describe_steps(trace.steps))

Trace

A list of steps. Each step is a dict: {"step": "retrieve", "input": {...}, "output": ..., "ms": 1.25}

Source code in src/rag_titan/rag_types/trace.py 
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
class Trace:
    """
    A list of steps. Each step is a dict:
        {"step": "retrieve", "input": {...}, "output": ..., "ms": 1.25}
    """

    def __init__(self):
        self.steps = []

    def record(self, step, function, **inputs):
        """
        Call function(**inputs), save it as a step, and return its output.

        Args:
            step: the step name, one of the RAG type's FLOW node ids.
            function: the function to call, e.g. retriever.retrieve.
            inputs: the keyword arguments for the function.
        """
        start = time.perf_counter()
        output = function(**inputs)
        ms = (time.perf_counter() - start) * 1000

        self.add(step, inputs, output, ms)
        return output

    def add(self, step, inputs, output, ms=0.0):
        """Save a step that was not made by calling a function (e.g. giving up)."""
        self.steps.append({"step": step, "input": inputs, "output": output, "ms": round(ms, 3)})

    def result(self, answer, sources=None):
        """Return what run() gives back: the answer, its sources and the steps."""
        return {"answer": answer, "sources": list(sources or []), "steps": self.steps}

add(step, inputs, output, ms=0.0)

Save a step that was not made by calling a function (e.g. giving up).

Source code in src/rag_titan/rag_types/trace.py 
40
41
42
def add(self, step, inputs, output, ms=0.0):
    """Save a step that was not made by calling a function (e.g. giving up)."""
    self.steps.append({"step": step, "input": inputs, "output": output, "ms": round(ms, 3)})

record(step, function, **inputs)

Call function(**inputs), save it as a step, and return its output.

Parameters:

Name Type Description Default
step

the step name, one of the RAG type's FLOW node ids.

 required
function

the function to call, e.g. retriever.retrieve.

 required
inputs

the keyword arguments for the function.

 {}
Source code in src/rag_titan/rag_types/trace.py 
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
def record(self, step, function, **inputs):
    """
    Call function(**inputs), save it as a step, and return its output.

    Args:
        step: the step name, one of the RAG type's FLOW node ids.
        function: the function to call, e.g. retriever.retrieve.
        inputs: the keyword arguments for the function.
    """
    start = time.perf_counter()
    output = function(**inputs)
    ms = (time.perf_counter() - start) * 1000

    self.add(step, inputs, output, ms)
    return output

result(answer, sources=None)

Return what run() gives back: the answer, its sources and the steps.

Source code in src/rag_titan/rag_types/trace.py 
44
45
46
def result(self, answer, sources=None):
    """Return what run() gives back: the answer, its sources and the steps."""
    return {"answer": answer, "sources": list(sources or []), "steps": self.steps}

describe_steps(steps)

Return a short, readable log of the steps, one line per step.

Example line
  1. retrieve 0.21 ms query='how many leave days?', top_k=4 -> 3 chunks: chunk_0 hr.pdf (0.894), chunk_2 it.pdf (0.000), ...
Source code in src/rag_titan/rag_types/trace.py 
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
def describe_steps(steps):
    """
    Return a short, readable log of the steps, one line per step.

    Example line:
        2. retrieve   0.21 ms  query='how many leave days?', top_k=4
                               -> 3 chunks: chunk_0 hr.pdf (0.894), chunk_2 it.pdf (0.000), ...
    """
    if not steps:
        return "(no steps: the question was not searched)"

    lines = []
    width = max(len(step["step"]) for step in steps)

    for number, step in enumerate(steps, start=1):
        # chunk lists in the input are the previous step's output, so only count them
        inputs = ", ".join(
            f"{key}={_short(value, brief=True)}" for key, value in step["input"].items()
        )
        head = f"{number}. {step['step']:<{width}}  {step['ms']:7.2f} ms  "
        lines.append(head + inputs)
        lines.append(" " * len(head) + "-> " + _short(step["output"]))

    return "\n".join(lines)
Simple RAG

Question -> Retrieve -> Rerank -> Generate

SimpleRAG

Bases: BaseRAG

The classic RAG: one search, then one answer.

Parameters:

Name Type Description Default
retriever

any BaseRetriever, for example a DenseRetriever.

 required
reranker

a BaseReranker.

 required
generator

a BaseGenerator.

 required
top_k

how many chunks to retrieve.

 5
top_n

how many chunks to keep after reranking.

 2
Source code in src/rag_titan/rag_types/simple_rag.py 
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
class SimpleRAG(BaseRAG):
    """
    The classic RAG: one search, then one answer.

    Args:
        retriever: any BaseRetriever, for example a DenseRetriever.
        reranker: a BaseReranker.
        generator: a BaseGenerator.
        top_k: how many chunks to retrieve.
        top_n: how many chunks to keep after reranking.
    """

    FLOW = {
        "nodes": [
            {"id": "retrieve", "label": "Retrieve", "kind": "search"},
            {"id": "rerank", "label": "Rerank", "kind": "search"},
            {"id": "generate", "label": "Generate answer", "kind": "output"},
        ],
        "edges": [["retrieve", "rerank"], ["rerank", "generate"]],
    }

    def __init__(self, retriever, reranker, generator, top_k=5, top_n=2):
        self.retriever = retriever
        self.reranker = reranker
        self.generator = generator
        self.top_k = top_k
        self.top_n = top_n

    def run(self, query):
        """Answer the question from the indexed documents and record each step."""
        query = query.strip()
        trace = Trace()

        if not query:
            return trace.result("Please enter a question.")

        if not self.retriever.has_documents():
            return trace.result("Please index your documents first.")

        matches = trace.record("retrieve", self.retriever.retrieve, query=query, top_k=self.top_k)
        best_chunks = trace.record(
            "rerank", self.reranker.rerank, query=query, matches=matches, top_n=self.top_n
        )

        for chunk in best_chunks:
            logger.info("Selected %s from %s (rerank %.4f)",
                        chunk["id"], chunk["source"], chunk["rerank_score"])

        answer = trace.record("generate", self.generator.generate, query=query, chunks=best_chunks)
        return trace.result(answer, best_chunks)

run(query)

Answer the question from the indexed documents and record each step.

Source code in src/rag_titan/rag_types/simple_rag.py 
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
def run(self, query):
    """Answer the question from the indexed documents and record each step."""
    query = query.strip()
    trace = Trace()

    if not query:
        return trace.result("Please enter a question.")

    if not self.retriever.has_documents():
        return trace.result("Please index your documents first.")

    matches = trace.record("retrieve", self.retriever.retrieve, query=query, top_k=self.top_k)
    best_chunks = trace.record(
        "rerank", self.reranker.rerank, query=query, matches=matches, top_n=self.top_n
    )

    for chunk in best_chunks:
        logger.info("Selected %s from %s (rerank %.4f)",
                    chunk["id"], chunk["source"], chunk["rerank_score"])

    answer = trace.record("generate", self.generator.generate, query=query, chunks=best_chunks)
    return trace.result(answer, best_chunks)

Adaptive RAG and its helpers (router, evidence checker, query rewriter).

AdaptiveRAG

Bases: BaseRAG

A RAG that changes its effort based on the question.

Greetings get a direct reply, simple questions a small search, complex questions a bigger search. When the evidence is weak, the question is rewritten and the search is tried again.

Parameters:

Name Type Description Default
llm

any BaseLLM, used for routing, checking and rewriting.

 required
retriever

any BaseRetriever, for example a DenseRetriever.

 required
reranker

a BaseReranker.

 required
generator

a BaseGenerator.

 required
max_attempts

how many searches to try before giving up.

 2
Source code in src/rag_titan/rag_types/adaptive/adaptive_rag.py 
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
class AdaptiveRAG(BaseRAG):
    """
    A RAG that changes its effort based on the question.

    Greetings get a direct reply, simple questions a small search,
    complex questions a bigger search. When the evidence is weak,
    the question is rewritten and the search is tried again.

    Args:
        llm: any BaseLLM, used for routing, checking and rewriting.
        retriever: any BaseRetriever, for example a DenseRetriever.
        reranker: a BaseReranker.
        generator: a BaseGenerator.
        max_attempts: how many searches to try before giving up.
    """

    FLOW = {
        "nodes": [
            {"id": "route", "label": "Choose route", "kind": "llm"},
            {"id": "direct_reply", "label": "Direct reply", "kind": "output"},
            {"id": "retrieve", "label": "Retrieve", "kind": "search"},
            {"id": "rerank", "label": "Rerank", "kind": "search"},
            {"id": "check", "label": "Check evidence", "kind": "llm"},
            {"id": "rewrite", "label": "Rewrite query", "kind": "llm"},
            {"id": "generate", "label": "Generate answer", "kind": "output"},
            {"id": "give_up", "label": "No answer", "kind": "output"},
        ],
        "edges": [
            ["route", "direct_reply"], ["route", "retrieve"],
            ["retrieve", "rerank"], ["rerank", "check"],
            ["check", "generate"], ["check", "rewrite"], ["check", "give_up"],
            ["rewrite", "retrieve"],
        ],
    }

    def __init__(self, llm, retriever, reranker, generator, max_attempts=2):
        self.llm = llm
        self.retriever = retriever
        self.reranker = reranker
        self.generator = generator
        self.max_attempts = max_attempts

        self.router = Router(llm)
        self.evidence_checker = EvidenceChecker(llm)
        self.query_rewriter = QueryRewriter(llm)

    def run(self, query):
        """Answer the question from the indexed documents and record each step."""
        query = query.strip()
        trace = Trace()

        if not query:
            return trace.result("Please enter a question.")

        route = trace.record("route", self.router.route, query=query)
        logger.info("Selected route: %s", route)

        if route == "direct":
            reply = trace.record("direct_reply", self.llm.ask,
                                 instructions=DIRECT_REPLY_PROMPT, user_input=query)
            return trace.result(reply)

        if not self.retriever.has_documents():
            return trace.result("Please index your documents first.")

        sizes = SEARCH_SIZES[route]
        return self._search_and_answer(query, sizes["top_k"], sizes["top_n"], trace)

    def _search_and_answer(self, query, top_k, top_n, trace):
        """Search, check the evidence, and retry with a rewritten query."""
        # the original question never changes, only the search query does
        search_query = query

        for attempt in range(1, self.max_attempts + 1):
            logger.info("Search attempt %d: '%s'", attempt, search_query)

            matches = trace.record("retrieve", self.retriever.retrieve,
                                   query=search_query, top_k=top_k)
            chunks = trace.record("rerank", self.reranker.rerank,
                                  query=query, matches=matches, top_n=top_n)
            enough = trace.record("check", self.evidence_checker.has_enough,
                                  query=query, chunks=chunks)

            if enough:
                logger.info("Enough evidence found")
                answer = trace.record("generate", self.generator.generate,
                                      query=query, chunks=chunks)
                return trace.result(answer, chunks)

            if attempt == self.max_attempts:
                break

            logger.info("Not enough evidence, rewriting and widening the search")
            search_query = trace.record("rewrite", self.query_rewriter.rewrite, query=query)
            top_k += 4
            top_n += 2

        trace.add("give_up", {"attempts": self.max_attempts}, NO_ANSWER)
        return trace.result(NO_ANSWER)

run(query)

Answer the question from the indexed documents and record each step.

Source code in src/rag_titan/rag_types/adaptive/adaptive_rag.py 
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
def run(self, query):
    """Answer the question from the indexed documents and record each step."""
    query = query.strip()
    trace = Trace()

    if not query:
        return trace.result("Please enter a question.")

    route = trace.record("route", self.router.route, query=query)
    logger.info("Selected route: %s", route)

    if route == "direct":
        reply = trace.record("direct_reply", self.llm.ask,
                             instructions=DIRECT_REPLY_PROMPT, user_input=query)
        return trace.result(reply)

    if not self.retriever.has_documents():
        return trace.result("Please index your documents first.")

    sizes = SEARCH_SIZES[route]
    return self._search_and_answer(query, sizes["top_k"], sizes["top_n"], trace)

EvidenceChecker

Asks the LLM whether the chunks are enough to answer the question.

Parameters:

Name Type Description Default
llm

any BaseLLM.

 required
Source code in src/rag_titan/rag_types/adaptive/evidence_checker.py 
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
class EvidenceChecker:
    """
    Asks the LLM whether the chunks are enough to answer the question.

    Args:
        llm: any BaseLLM.
    """

    def __init__(self, llm):
        self.llm = llm

    def has_enough(self, query, chunks):
        """Return True when the LLM replies "yes"."""
        if not chunks:
            return False

        user_input = f"QUESTION:\n{query}\n\nCONTEXT:\n{build_context(chunks)}"
        reply = self.llm.ask(EVIDENCE_CHECK_PROMPT, user_input)

        return reply.strip().lower().startswith("yes")

has_enough(query, chunks)

Return True when the LLM replies "yes".

Source code in src/rag_titan/rag_types/adaptive/evidence_checker.py 
18
19
20
21
22
23
24
25
26
def has_enough(self, query, chunks):
    """Return True when the LLM replies "yes"."""
    if not chunks:
        return False

    user_input = f"QUESTION:\n{query}\n\nCONTEXT:\n{build_context(chunks)}"
    reply = self.llm.ask(EVIDENCE_CHECK_PROMPT, user_input)

    return reply.strip().lower().startswith("yes")

QueryRewriter

Rewrites a question into a better search query.

Parameters:

Name Type Description Default
llm

any BaseLLM.

 required
Source code in src/rag_titan/rag_types/adaptive/query_rewriter.py 
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
class QueryRewriter:
    """
    Rewrites a question into a better search query.

    Args:
        llm: any BaseLLM.
    """

    def __init__(self, llm):
        self.llm = llm

    def rewrite(self, query):
        """Return the rewritten query, or the original if the LLM is empty."""
        rewritten = self.llm.ask(QUERY_REWRITE_PROMPT, query)
        return rewritten or query

rewrite(query)

Return the rewritten query, or the original if the LLM is empty.

Source code in src/rag_titan/rag_types/adaptive/query_rewriter.py 
17
18
19
20
def rewrite(self, query):
    """Return the rewritten query, or the original if the LLM is empty."""
    rewritten = self.llm.ask(QUERY_REWRITE_PROMPT, query)
    return rewritten or query

Router

Decides how much work a question needs: direct, simple or complex.

Parameters:

Name Type Description Default
llm

any BaseLLM.

 required
default_route

route used when the LLM reply is not a known label.

 'simple'
Source code in src/rag_titan/rag_types/adaptive/router.py 
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
class Router:
    """
    Decides how much work a question needs: direct, simple or complex.

    Args:
        llm: any BaseLLM.
        default_route: route used when the LLM reply is not a known label.
    """

    def __init__(self, llm, default_route="simple"):
        self.llm = llm
        self.default_route = default_route

    def route(self, query):
        """Return "direct", "simple" or "complex"."""
        reply = self.llm.ask(ROUTER_PROMPT, query)
        label = reply.strip().lower().strip(".!\"' ")

        # small models sometimes add extra words, so fall back safely
        if label not in ROUTES:
            return self.default_route

        return label

route(query)

Return "direct", "simple" or "complex".

Source code in src/rag_titan/rag_types/adaptive/router.py 
21
22
23
24
25
26
27
28
29
30
def route(self, query):
    """Return "direct", "simple" or "complex"."""
    reply = self.llm.ask(ROUTER_PROMPT, query)
    label = reply.strip().lower().strip(".!\"' ")

    # small models sometimes add extra words, so fall back safely
    if label not in ROUTES:
        return self.default_route

    return label