Skip to content

retrieval_rag_stage

RetrievalRagStage

Bases: BaseRagStage

Source code in griptape/engines/rag/stages/retrieval_rag_stage.py 
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
63
64
65
66
67
@define(kw_only=True)
class RetrievalRagStage(BaseRagStage):
    retrieval_modules: list[BaseRetrievalRagModule] = field()
    rerank_module: Optional[BaseRerankRagModule] = field(default=None)
    max_chunks: Optional[int] = field(default=None)

    @property
    def modules(self) -> list[BaseRagModule]:
        ms = []

        ms.extend(self.retrieval_modules)

        if self.rerank_module is not None:
            ms.append(self.rerank_module)

        return ms

    def run(self, context: RagContext) -> RagContext:
        logging.info("RetrievalRagStage: running %s retrieval modules in parallel", len(self.retrieval_modules))

        results = utils.execute_futures_list(
            [self.futures_executor.submit(r.run, context) for r in self.retrieval_modules]
        )

        # flatten the list of lists
        results = list(itertools.chain.from_iterable(results))

        # deduplicate the list
        chunks_before_dedup = len(results)
        results = list({str(c.value): c for c in results}.values())
        chunks_after_dedup = len(results)

        logging.info(
            "RetrievalRagStage: deduplicated %s " "chunks (%s - %s)",
            chunks_before_dedup - chunks_after_dedup,
            chunks_before_dedup,
            chunks_after_dedup,
        )

        context.text_chunks = [a for a in results if isinstance(a, TextArtifact)]

        if self.rerank_module:
            logging.info("RetrievalRagStage: running rerank module on %s chunks", chunks_after_dedup)

            context.text_chunks = [a for a in self.rerank_module.run(context) if isinstance(a, TextArtifact)]

        if self.max_chunks:
            context.text_chunks = context.text_chunks[: self.max_chunks]

        return context

max_chunks: Optional[int] = field(default=None) class-attribute instance-attribute

modules: list[BaseRagModule] property

rerank_module: Optional[BaseRerankRagModule] = field(default=None) class-attribute instance-attribute

retrieval_modules: list[BaseRetrievalRagModule] = field() class-attribute instance-attribute

run(context)

Source code in griptape/engines/rag/stages/retrieval_rag_stage.py 
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
def run(self, context: RagContext) -> RagContext:
    logging.info("RetrievalRagStage: running %s retrieval modules in parallel", len(self.retrieval_modules))

    results = utils.execute_futures_list(
        [self.futures_executor.submit(r.run, context) for r in self.retrieval_modules]
    )

    # flatten the list of lists
    results = list(itertools.chain.from_iterable(results))

    # deduplicate the list
    chunks_before_dedup = len(results)
    results = list({str(c.value): c for c in results}.values())
    chunks_after_dedup = len(results)

    logging.info(
        "RetrievalRagStage: deduplicated %s " "chunks (%s - %s)",
        chunks_before_dedup - chunks_after_dedup,
        chunks_before_dedup,
        chunks_after_dedup,
    )

    context.text_chunks = [a for a in results if isinstance(a, TextArtifact)]

    if self.rerank_module:
        logging.info("RetrievalRagStage: running rerank module on %s chunks", chunks_after_dedup)

        context.text_chunks = [a for a in self.rerank_module.run(context) if isinstance(a, TextArtifact)]

    if self.max_chunks:
        context.text_chunks = context.text_chunks[: self.max_chunks]

    return context