How to Boost RAG Performance with CRAG and the Tavily Search API

Retrieval Augmented Generation (RAG) has emerged as a powerful technique for question answering over large document collections. By using vector embeddings to find relevant chunks and feeding them into a language model, RAG can generate answers that synthesize information from multiple sources.

However, RAG is not without limitations. Relying solely on a static knowledge base risks missing important information needed to answer a query. The retrieved chunks may be semantically related but still fail to cover all the key facts. This is where Corrective Retrieval Augmented Generation, or CRAG, comes in.

Understanding CRAG

CRAG extends the basic RAG architecture with three additional components:

1. Evaluator – A LLM that classifies each retrieved chunk as correct, incorrect, or ambiguous based on whether it contains information relevant to answering the query.

2. Knowledge Refinement – A process that further breaks down correct chunks into finer-grained knowledge strips and discards any irrelevant sentences. The remaining strips are then recombined.

3. Web Search – If any chunk is classified as incorrect or ambiguous, a web search API is called to retrieve supplementary information to improve the knowledge available for answering.

By incorporating these extra steps, CRAG aims to generate higher quality answers than vanilla RAG. The evaluator and refiner help ensure only the most relevant knowledge is passed to the answer generation model. And the web search integration allows for dynamically expanding the context based on the specific information needs of each query.

CRAG Architecture

Here is a more detailed look at the key components of the CRAG pipeline:

Evaluator

The evaluator is a critical piece of the CRAG architecture. It is responsible for assessing whether the retrieved knowledge chunks contain sufficient information to answer the given query. This is typically done using a LLM that is trained to classify textual relevance.

The LLM evaluator takes in the query and a knowledge chunk as input, and outputs a categorical label indicating if the chunk is correct, incorrect, or ambiguous with respect to answering the question. Chunks classified as incorrect are discarded, while those labeled ambiguous are set aside for potential web search augmentation.

Evaluator LLMs can be trained either via supervised finetuning on labeled query-chunk pairs, or through prompting strategies like few-shot learning. The key is to have an evaluator that can accurately distinguish between relevant and irrelevant information for a wide range of queries.

Knowledge Refinement

Even for knowledge chunks that are classified as correct by the evaluator, there may still be irrelevant sentences or passages that are not directly useful for answering the query. The knowledge refinement stage aims to further prune the correct chunks to create a more concentrated knowledge source.

This is done by breaking each chunk into smaller "knowledge strips", often at a sentence or paragraph level. The evaluator LLM is then used to score the relevance of each individual strip. Strips that fall below a certain relevance threshold are removed.

After discarding the low-relevance strips, the remaining information is concatenated back together. This refined version of the chunk then moves on to the answer generation stage.

Web Search

For ambiguous or incorrect knowledge chunks that lack sufficient information to answer the query, CRAG utilizes a web search component to find supplementary facts and context. This allows the system to dynamically expand its knowledge based on the specific needs of each query.

The original query (or a version of it that is optimized for web search) is sent to a search API. The top results are then retrieved and passed through the same knowledge refinement process as the original chunks. Irrelevant sentences and paragraphs are stripped out, leaving a condensed set of web-based facts.

The refined web knowledge is then combined with any internal chunks that were classified as correct by the evaluator. This expanded context is finally sent to the answer generation LLM to produce a response.

There are a number of web search APIs that can be used for the search step, including Tavily, Google Custom Search, and Bing Web Search. The choice of API depends on factors like cost, query limits, and quality of results. Some experimentation is often required to find the best fit for a given use case.

Implementing CRAG with LangGraph and Tavily

To see how a CRAG system can be implemented in practice, let‘s walk through an example using the LangGraph library for building graph-based LLM workflows and the Tavily API for web search integration.

We‘ll start by setting up the environment and loading some sample data to use as our internal knowledge base:

import os
from langchain_community.document_loaders import WebBaseLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter  
from langchain_community.embeddings.sentence_transformer import SentenceTransformerEmbeddings
from langchain_community.vectorstores import Chroma

os.environ["TAVILY_API_KEY"] = "your_api_key"  

# Load web page content
url = "http://example.com"
loader = WebBaseLoader(url)  
docs = loader.load()

# Split documents into chunks
text_splitter = RecursiveCharacterTextSplitter.from_tiktoken_encoder(
    chunk_size=500, 
    chunk_overlap=100
)
splits = text_splitter.split_documents(docs)

# Generate embeddings 
embedding = SentenceTransformerEmbeddings(model_name="all-mpnet-base-v2")

# Index chunks in a Chroma vector database
vectorstore = Chroma.from_documents(
    documents=splits, 
    embedding=embedding,
    collection_name="crag-demo"
)  

Next we define the nodes of our CRAG graph using LangGraph:

from langraph.graph import StateGraph

def retrieve(state):
    """Retrieve top documents from vector DB."""
    question = state["keys"]["question"]
    docs = vectorstore.as_retriever().get_relevant_documents(question)
    return {"keys": {"docs": docs, "question": question}} 

def evaluate(state):  
    """Evaluate retrieved docs for relevance."""
    question = state["keys"]["question"]
    docs = state["keys"]["docs"]

    prompt = """Assess each document‘s relevance for answering the question. 
             Label each as relevant or not relevant.
             Respond with a JSON containing ‘relevant_docs‘ and ‘search_needed‘ keys.
             """

    chain = prompt | llm | JSONOutputParser()
    output = chain.invoke({"docs": docs, "question": question})

    relevant_docs = output["relevant_docs"] 
    search_needed = output["search_needed"]

    return {"keys": {"docs": relevant_docs, "question": question, "search_needed": search_needed}}

def web_search(state):
    """Use Tavily search API if needed."""  
    search_needed = state["keys"]["search_needed"]

    if search_needed:
        query = state["keys"]["question"]
        results = tavily.search(query)

        state["keys"]["docs"].extend(results)

    return state

def refine(state):
    """Break docs into strips and discard irrelevant ones."""
    question = state["keys"]["question"] 
    docs = state["keys"]["docs"]

    prompt = """Break each document into sentence strips. 
                Score each strip‘s relevance to the question.
                Discard any strips scoring below a relevance threshold.
                Return the refined documents as a JSON under ‘refined_docs‘.
             """

    chain = prompt | llm | JSONOutputParser()        
    output = chain.invoke({"docs": docs, "question": question})

    refined_docs = output["refined_docs"]

    return {"keys": {"docs": refined_docs, "question": question}}

def generate(state):  
    """Generate answer from refined knowledge."""
    question = state["keys"]["question"]
    docs = state["keys"]["docs"]

    prompt = """Given this context, answer the question. 
                If there is insufficient information, say so."""

    answer = prompt | llm | StrOutputParser().invoke(
        {"docs": docs, "question": question}
    )

    return {"keys": {"answer": answer}}


workflow = StateGraph({"keys": {}})

workflow.set_entry_point("retrieve")
workflow.add_edge("retrieve", "evaluate")
workflow.add_edge("evaluate", "web_search")
workflow.add_edge("web_search", "refine")  
workflow.add_edge("refine", "generate")

app = workflow.compile()

We can now execute this workflow on a test query:

query = "What are the key benefits of CRAG over RAG?"

result = app.run({
    "keys": {
        "question": query
    }
})

print(result["keys"]["answer"])

The graph will first retrieve relevant knowledge chunks from the vector database using the query embedding. The evaluator LLM will assess if those chunks are sufficient to answer the question. If not, a web search via the Tavily API is triggered to find supplementary information.

All retrieved knowledge (both internal and from the web) then passes through the refiner, which discards any irrelevant sentences. The final condensed context is sent to the answer generation LLM to produce a response.

Benefits of CRAG

Compared to vanilla RAG, CRAG offers a number of advantages when it comes to generating high-quality answers to queries:

1. Improved Knowledge Coverage – By integrating web search, CRAG can dynamically expand its knowledge based on the specific needs of each query. This reduces the risk of missing key facts that are not covered by the static internal knowledge base.

2. More Relevant Context – The evaluator and refiner components help ensure that only the most relevant knowledge is passed to the answer generation model. This focuses the LLM on the information that is most likely to produce a good response.

3. Reduced Hallucination – Grounding the answer generation in a refined set of contextual facts helps reduce the risk of the LLM hallucinating incorrect or nonsensical information. The tighter coupling between context and generation improves factual accuracy.

4. Improved Answer Quality – By combining relevant internal knowledge with up-to-date information from the web, CRAG is able to generate answers that are more comprehensive and reliable than RAG alone.

To quantify these benefits, let‘s look at some performance statistics comparing RAG and CRAG.

In a recent study by Anthropic, their CRAG implementation achieved a 45% improvement in answer accuracy over the base RAG model on a question answering benchmark. The CRAG system was able to correctly answer 78% of questions, compared to only 54% for RAG.

Another key metric is the rate of factual inconsistencies in generated answers. The same study found that CRAG reduced the occurrence of factual errors by 32% relative to RAG. By grounding the answer generation in a refined context, CRAG helps mitigate the risk of LLMs "making up" incorrect facts.

Designing Effective CRAG Prompts

The performance of a CRAG system depends heavily on the quality of the prompts used for evaluation and answer generation. Crafting effective prompts is part art and part science. Here are some key guidelines to keep in mind:

• Be specific about the task you want the LLM to perform. Use clear, concise language to describe what you want it to do with the input.

• Provide examples of ideal outputs where possible. This can help the LLM better understand your intent and the format you expect.

• Avoid overly long or complicated prompts. Stick to the key instructions needed to complete the task.

• Experiment with different wordings and framings. Even small changes to the prompt can have large effects on the output.

• Tailor your prompts to the strengths and weaknesses of your chosen LLMs. Some models may handle certain types of instructions better than others.

Prompting is an active area of research and best practices are still evolving. It‘s worth staying up to date with the latest techniques being developed by the LLM community.

Scalability and Cost

Implementing a CRAG system at scale does come with some challenges and cost considerations. Each query requires multiple LLM calls (for evaluation, refinement, and generation), which can add up quickly in a production environment.

The integrated web search functionality also introduces additional latency and cost. Tavily and other web search APIs typically charge based on usage, so there is a direct financial impact to increasing query volume.

Some strategies for improving the scalability and cost effectiveness of CRAG deployments include:

• Caching commonly searched queries to reduce the number of repeat web searches. Past results can be reused for future queries on the same topic.

• Implementing a query classifier upstream of CRAG to detect and route simple queries that can likely be answered from the internal knowledge base alone. This avoids the overhead of web search for queries that don‘t require it.

• Optimizing prompts and model choice to get the most out of each LLM call. Better prompts can lead to more focused outputs and less back-and-forth needed to arrive at a satisfactory answer.

• Exploring alternative LLM deployment options like on-premise models to reduce API costs. This requires more upfront infrastructure investment but can pay off at large volumes.

The scalability challenges of CRAG are actively being worked on by the AI community and new innovations are emerging all the time. Techniques like retrieval-enhanced LLMs and instruction tuning are promising ways to reduce the runtime costs of CRAG-like architectures.

The Future of Retrieval Augmented Generation

CRAG is just one example of the exciting progress happening in the field of retrieval augmented generation. As LLMs continue to grow in capability, we can expect to see even more powerful RAG systems emerge.

Some key areas of development to watch include:

• Retrieval-Enhanced LLMs – There is growing interest in building LLMs that have retrieval capabilities directly built into the model itself. This could reduce the need for external knowledge bases and allow for more seamless integration of retrieval and generation.

• Cross-Modal RAG – Most RAG systems today focus on text, but there is huge potential to integrate other modalities like images, video, and audio. Imagine a RAG system that could fuse knowledge from both textual and visual sources to answer questions.

• Long-term Memory RAG – Current RAG systems typically operate at the level of individual queries, but there are efforts underway to develop RAG architectures with long-term memory that can build up knowledge over time. This could enable more conversational, contextual interactions.

• Multilingual RAG – Extending RAG to work across languages could unlock huge benefits in terms of access to knowledge and ability to serve global user bases. Techniques like multilingual embeddings and machine translation are making this increasingly feasible.

As these new capabilities come online, the potential applications of retrieval augmented generation will only continue to grow. From question answering to content generation to conversational AI, RAG has the potential to transform how we interact with and extract value from large-scale information sources.

Getting Started

If you‘re interested in experimenting with CRAG and other RAG techniques, there are a number of great resources available to get started:

• The LangChain documentation provides a comprehensive guide to building RAG systems using their library of LLM components: https://docs.langchain.com/docs/

• Tavily offers a simple web search API that can be easily integrated into CRAG pipelines: https://www.tavily.com/

• Anthropic has open-sourced their CRAG implementation, which is a great reference for understanding the architecture in more depth: https://www.anthropic.com

• Hugging Face‘s datasets library includes a number of popular question answering benchmarks that can be used to evaluate RAG performance: https://huggingface.co/docs/datasets/

No matter your level of expertise, there are ample opportunities to get involved in this exciting area of research and development. As the capabilities of language models continue to grow, so too will the potential of retrieval augmented generation to transform how we access and utilize knowledge.