Skip to content

stages

__all__ = ['BaseRagStage', 'QueryRagStage', 'ResponseRagStage', 'RetrievalRagStage'] module-attribute

BaseRagStage

Bases: FuturesExecutorMixin, ABC

Source code in griptape/engines/rag/stages/base_rag_stage.py 
11
12
13
14
15
16
17
18
@define(kw_only=True)
class BaseRagStage(FuturesExecutorMixin, ABC):
    @abstractmethod
    def run(self, context: RagContext) -> RagContext: ...

    @property
    @abstractmethod
    def modules(self) -> Sequence[BaseRagModule]: ...

modules abstractmethod property

run(context) abstractmethod

Source code in griptape/engines/rag/stages/base_rag_stage.py 
13
14
@abstractmethod
def run(self, context: RagContext) -> RagContext: ...

QueryRagStage

Bases: BaseRagStage

Source code in griptape/engines/rag/stages/query_rag_stage.py 
17
18
19
20
21
22
23
24
25
26
27
28
29
30
@define(kw_only=True)
class QueryRagStage(BaseRagStage):
    query_modules: list[BaseQueryRagModule] = field()

    @property
    def modules(self) -> Sequence[BaseRagModule]:
        return self.query_modules

    def run(self, context: RagContext) -> RagContext:
        logging.info("QueryRagStage: running %s query generation modules sequentially", len(self.query_modules))

        [qm.run(context) for qm in self.query_modules]

        return context

modules property

query_modules = field() class-attribute instance-attribute

run(context)

Source code in griptape/engines/rag/stages/query_rag_stage.py 
25
26
27
28
29
30
def run(self, context: RagContext) -> RagContext:
    logging.info("QueryRagStage: running %s query generation modules sequentially", len(self.query_modules))

    [qm.run(context) for qm in self.query_modules]

    return context

ResponseRagStage

Bases: BaseRagStage

Source code in griptape/engines/rag/stages/response_rag_stage.py 
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
@define(kw_only=True)
class ResponseRagStage(BaseRagStage):
    response_modules: list[BaseResponseRagModule] = field()

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

        ms.extend(self.response_modules)

        return ms

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

        with self.create_futures_executor() as futures_executor:
            results = utils.execute_futures_list(
                [futures_executor.submit(with_contextvars(r.run), context) for r in self.response_modules]
            )

        context.outputs = results

        return context

modules property

response_modules = field() class-attribute instance-attribute

run(context)

Source code in griptape/engines/rag/stages/response_rag_stage.py 
32
33
34
35
36
37
38
39
40
41
42
def run(self, context: RagContext) -> RagContext:
    logging.info("ResponseRagStage: running %s retrieval modules in parallel", len(self.response_modules))

    with self.create_futures_executor() as futures_executor:
        results = utils.execute_futures_list(
            [futures_executor.submit(with_contextvars(r.run), context) for r in self.response_modules]
        )

    context.outputs = results

    return context

RetrievalRagStage

Bases: BaseRagStage

Source code in griptape/engines/rag/stages/retrieval_rag_stage.py 
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
68
69
@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))

        with self.create_futures_executor() as futures_executor:
            results = utils.execute_futures_list(
                [futures_executor.submit(with_contextvars(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 = field(default=None) class-attribute instance-attribute

modules property

rerank_module = field(default=None) class-attribute instance-attribute

retrieval_modules = field() class-attribute instance-attribute

run(context)

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

    with self.create_futures_executor() as futures_executor:
        results = utils.execute_futures_list(
            [futures_executor.submit(with_contextvars(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