RAG Engine

RAG Engine#

RAG Engines#

RAG Engine is an abstraction for implementing modular retrieval augmented generation (RAG) pipelines.

RAG Stages#

RagEngine s consist of three stages: QueryRagStage , RetrievalRagStage , and ResponseRagStage . These stages are always executed sequentially. Each stage comprises multiple modules, which are executed in a customized manner. Due to this unique structure, RagEngines are not intended to replace Workflows or HuggingFace Pipelines .

RAG Modules#

RAG modules are used to implement actions in the different stages of the RAG pipeline. RagEngine enables developers to easily add new modules to experiment with novel RAG strategies.

The three stages of the pipeline implemented in RAG Engines, together with their purposes and associated modules, are as follows:

Query Stage#

This stage is used for modifying input queries before they are submitted.

Query Stage Modules#

  • TranslateQueryRagModule is for translating the query into another language.

Retrieval Stage#

Results are retrieved in this stage, either from a vector store in the form of chunks, or with a text loader. You may optionally use a rerank module in this stage to rerank results in order of their relevance to the original query.

Retrieval Stage Modules#

  • TextChunksRerankRagModule is for re-ranking retrieved results.

  • TextLoaderRetrievalRagModule is for retrieving data with text loaders in real time.

  • VectorStoreRetrievalRagModule is for retrieving text chunks from a vector store.

Response Stage#

Responses are generated in this final stage.

Response Stage Modules#

  • PromptResponseRagModule is for generating responses based on retrieved text chunks.

  • TextChunksResponseRagModule is for responding with retrieved text chunks.

  • FootnotePromptResponseRagModule is for responding with automatic footnotes from text chunk references.

RAG Context#

RagContext is a container object for passing around queries, text chunks, module configs, and other metadata. RagContext is modified by modules when appropriate. Some modules support runtime config overrides through RagContext.module_configs .

Simple Example#

The following example shows a simple RAG pipeline that retrieves data from a local vector store and generates a response:

from griptape.chunkers import TextChunker
from griptape.drivers.embedding.openai import OpenAiEmbeddingDriver
from griptape.drivers.vector.local import LocalVectorStoreDriver
from griptape.engines.rag import RagContext, RagEngine
from griptape.engines.rag.modules import (
    PromptResponseRagModule,
    VectorStoreRetrievalRagModule,
)
from griptape.engines.rag.stages import (
    ResponseRagStage,
    RetrievalRagStage,
)
from griptape.loaders import WebLoader

vector_store = LocalVectorStoreDriver(embedding_driver=OpenAiEmbeddingDriver())

#  Load some data from a couple sources.

web_artifact = WebLoader().load("https://griptape.ai")
chunks = TextChunker(max_tokens=250).chunk(web_artifact)
vector_store.upsert_collection({"griptape": chunks})

rag_engine = RagEngine(
    # This stage is responsible for retrieving the relevant chunks.
    retrieval_stage=RetrievalRagStage(
        retrieval_modules=[
            VectorStoreRetrievalRagModule(
                name="WebRetriever",
                vector_store_driver=vector_store,
                query_params={"namespace": "griptape"},
            ),
        ],
    ),
    # This stage is responsible for generating the final response.
    response_stage=ResponseRagStage(
        response_modules=[
            PromptResponseRagModule(),
        ]
    ),
)

rag_context = RagContext(query="What is Griptape?")
rag_context = rag_engine.process(rag_context)
print(rag_context.outputs[0].to_text())

Advanced Example#

The following example shows an advanced RAG pipeline that does the following:

  • Translates incoming queries into English.

  • Retrieves data from a local vector store and a local file.

  • Reranks the results using the local rerank driver.

  • Generates multiple types of response.

from rich.pretty import pprint

from griptape.chunkers import TextChunker
from griptape.common.reference import Reference
from griptape.drivers.embedding.openai import OpenAiEmbeddingDriver
from griptape.drivers.prompt.openai import OpenAiChatPromptDriver
from griptape.drivers.rerank.local import LocalRerankDriver
from griptape.drivers.vector.local import LocalVectorStoreDriver
from griptape.engines.rag import RagContext, RagEngine
from griptape.engines.rag.modules import (
    FootnotePromptResponseRagModule,
    PromptResponseRagModule,
    TextChunksRerankRagModule,
    TextChunksResponseRagModule,
    TextLoaderRetrievalRagModule,
    TranslateQueryRagModule,
    VectorStoreRetrievalRagModule,
)
from griptape.engines.rag.stages import (
    QueryRagStage,
    ResponseRagStage,
    RetrievalRagStage,
)
from griptape.loaders import TextLoader, WebLoader
from griptape.rules import Rule, Ruleset

prompt_driver = OpenAiChatPromptDriver(model="gpt-4.1")
vector_store = LocalVectorStoreDriver(embedding_driver=OpenAiEmbeddingDriver())
rerank_driver = LocalRerankDriver()
web_loader = WebLoader()
text_chunker = TextChunker(max_tokens=250)

#  Load some data from a couple sources.

sites = [
    {
        "title": "Griptape Site",
        "url": "https://www.griptape.ai",
    },
    {"title": "Griptape GitHub", "url": "https://github.com/griptape-ai/griptape"},
]
site_artifacts = list(web_loader.load_collection([site["url"] for site in sites]).values())

#  Set a reference on each artifact so that the FootnotePromptResponseRagModule can generate footnotes.

for site_arifact, site in zip(site_artifacts, sites):
    site_arifact.reference = Reference(title=site["title"], url=site["url"])

#  Chunk each site artifact.

site_artifacts_chunks = [text_chunker.chunk(artifact) for artifact in site_artifacts]

#  Flatten the list of chunks before inserting them into the vector store.

site_artifacts_chunks = [
    site_artifact_chunk for site_artifact_chunk in site_artifacts_chunks for site_artifact_chunk in site_artifact_chunk
]
vector_store.upsert_collection({"griptape": site_artifacts_chunks})

rag_engine = RagEngine(
    # This stage is responsible for producing the query. It can include things like translation, rewriting, etc.
    query_stage=QueryRagStage(query_modules=[TranslateQueryRagModule(prompt_driver=prompt_driver, language="english")]),
    # This stage is responsible for retrieving the relevant chunks.
    retrieval_stage=RetrievalRagStage(
        max_chunks=5,
        retrieval_modules=[  # Modules can pull from different sources.
            # Such as a vector store
            VectorStoreRetrievalRagModule(
                name="WebRetriever",
                vector_store_driver=vector_store,
                query_params={"top_n": 20, "namespace": "griptape"},
            ),
            # Or a text source.
            TextLoaderRetrievalRagModule(
                loader=TextLoader(),
                vector_store_driver=vector_store,
                source="README.md",
            ),
        ],
        # We can rerank the chunks before passing them to the response stage to ensure the best ones are used.
        rerank_module=TextChunksRerankRagModule(rerank_driver=LocalRerankDriver()),
    ),
    # This stage is responsible for generating the final response.
    response_stage=ResponseRagStage(
        response_modules=[
            # You can have multiple response modules to generate different types of responses.
            PromptResponseRagModule(
                prompt_driver=prompt_driver,
                rulesets=[Ruleset(name="Concise", rules=[Rule("Answer concisely")])],
            ),
            # This one generates a response with footnotes.
            FootnotePromptResponseRagModule(prompt_driver=prompt_driver),
            # This one just returns the text chunks directly with no LLM generation.
            TextChunksResponseRagModule(),
        ]
    ),
)

rag_context = RagContext(
    query="¿Qué ofrecen los servicios en la nube de Griptape?",  # What do Griptapes cloud services offer?
)
rag_context = rag_engine.process(rag_context)

#  Lets print out the interesting responses

for output in rag_context.outputs[:2]:
    print(output.to_text() + "\n")

#  We can also see the references that were configured up-front

pprint(rag_context.get_references())

RAG Tool#

See RagTool for an example of how to integrate the RAG Tool with an Agent.

Contents