How to build a RAG system

Table of contents



Understanding retrieval-augmented generation (RAG)

The basic RAG process



Step 1: Query

def preprocess_query(query):
    """
    Preprocesses the input query for retrieval.

    Args:
        query (str): The input query.

    Returns:
        str: The preprocessed query.
    """
    # Tokenize the query
    tokens = word_tokenize(query)
    # Lowercase the tokens
    tokens = [token.lower() for token in tokens]
    preprocessed_query = " ".join(tokens)
    return preprocessed_query

Step 2: Retrieval

class TFIDFRetriever(weave.Model):
    """
    A retriever model that uses TF-IDF for indexing and searching documents.

    Attributes:
        vectorizer (TfidfVectorizer): The TF-IDF vectorizer.
        index (list): The indexed data.
        data (list): The data to be indexed.
    """

    vectorizer: TfidfVectorizer = TfidfVectorizer()
    index: list = None
    data: list = None

    def index_data(self, data):
        """
        Indexes the provided data using TF-IDF.

        Args:
            data (list): A list of documents to be indexed. Each document should be a dictionary
                         containing a key 'cleaned_content' with the text to be indexed.
        """
        self.data = data
        docs = [doc["cleaned_content"] for doc in data]
        self.index = self.vectorizer.fit_transform(docs)

    @weave.op
    def search(self, query, k=5):
        """
        Searches the indexed data for the given query using cosine similarity.

        Args:
            query (str): The search query.
            k (int): The number of top results to return. Default is 5.

        Returns:
            list: A list of dictionaries containing the source, text, and score of the top-k results.
        """
        query_vec = self.vectorizer.transform([query])
        cosine_distances = cdist(
            query_vec.todense(), self.index.todense(), metric="cosine"
        )[0]
        top_k_indices = cosine_distances.argsort()[:k]
        output = []
        for idx in top_k_indices:
            output.append(
                {
                    "source": self.data[idx]["metadata"]["source"],
                    "text": self.data[idx]["cleaned_content"],
                    "score": 1 - cosine_distances[idx],
                }
            )
        return output

    @weave.op
    def predict(self, query: str, k: int):
        """
        Predicts the top-k results for the given query.

        Args:
            query (str): The search query.
            k (int): The number of top results to return.

        Returns:
            list: A list of dictionaries containing the source, text, and score of the top-k results.
        """
        return self.search(query, k)

Step 3: Context integration

@weave.op
 def generate_context(self, context: List[Dict[str, any]]) -> List[Dict[str, any]]:
     """
     Generate a list of contexts from the provided context list.

     Args:
         context (List[Dict[str, any]]): A list of dictionaries containing context data.

     Returns:
         List[Dict[str, any]]: A list of dictionaries with 'source' and 'text' keys.
     """
     contexts = [
         {"data": {"source": item["source"], "text": item["text"]}}
         for item in context
     ]
     return contexts

Step 4: Generation and response

class SimpleResponseGenerator(weave.Model):
    """
    A simple response generator model using Cohere's API.

    Attributes:
        model (str): The model name for generating responses.
        prompt (str): The prompt to be used for generating responses.
        client (cohere.ClientV2): The Cohere client for interacting with the Cohere API.
    """

    model: str
    prompt: str
    client: cohere.ClientV2 = None

    def __init__(self, **kwargs):
        """
        Initialize the SimpleResponseGenerator with the provided keyword arguments.
        Sets up the Cohere client using the API key from environment variables.
        """
        super().__init__(**kwargs)
        self.client = cohere.ClientV2(
            api_key=os.environ["COHERE_API_KEY"],
        )

    def create_messages(self, query: str):
        """
        Create a list of messages for the chat model based on the query.

        Args:
            query (str): The user's query.

        Returns:
            List[Dict[str, any]]: A list of messages formatted for the chat model.
        """
        messages = [
            {"role": "system", "content": self.prompt},
            {"role": "user", "content": query},
        ]
        return messages

    @weave.op()
    def generate_response(self, query: str, context: List[Dict[str, any]]) -> str:
        """
        Generate a response from the chat model based on the query and context.

        Args:
            query (str): The user's query.
            context (List[Dict[str, any]]): A list of dictionaries containing context data.

        Returns:
            str: The generated response from the chat model.
        """
        documents = self.generate_context(context)
        messages = self.create_messages(query)
        response = self.client.chat(
            messages=messages,
            model=self.model,
            temperature=0.1,
            max_tokens=2000,
            documents=documents,
        )
        return response.message.content[0].text

Advanced RAG components



1. Query enhancement

Query expansion using synonyms or related terms
Key-phrase recognition to identify key concepts
Query reformulation based on user intent analysis

@weave.op
async def predict(self, query: str) -> Dict[str, Any]:
    """
    Predict the language, generate search queries, and get intent predictions for a given query.

    Args:
        query (str): The input query to process.

    Returns:
        Dict[str, Any]: A dictionary containing the original query, detected language, generated search queries, and intent predictions.
    """
    language = detect_language(query.replace("\n", " "))["lang"]
    search_queries = await self.generate_cohere_queries(query)
    intents = await self.get_intent_prediction(query)
    return {
        "query": query,
        "language": language,
        "search_queries": search_queries,
        "intents": intents["intents"],
    }

2. Reranking

Semantic similarity scoring between query and retrieved documents
Consideration of document freshness and source reliability
Machine learning models trained on historical query-document relevance data

@weave.op
def rerank(self, query, docs, top_n=None):
    """
    Reranks the given documents based on their relevance to the query.

    Args:
        query (str): The query string.
        docs (List[Dict[str, Any]]): A list of documents to be reranked.
        top_n (int, optional): The number of top documents to return. Defaults to None.

    Returns:
        List[Dict[str, Any]]: A list of reranked documents with relevance scores.
    """
    client = cohere.ClientV2(os.environ["COHERE_API_KEY"])
    documents = [doc["text"] for doc in docs]
    response = client. rerank (
        model=self.model, query=query, documents=documents, top_n=top_n or len(docs)
    )

    outputs = []
    for doc in response.results:
        reranked_doc = docs[doc.index]
        reranked_doc["relevance_score"] = doc.relevance_score
        outputs.append(reranked_doc)
    return outputs[:top_n]

3. Response validation

Fact-checking against the retrieved information
Coherence assessment of the generated response
Confidence scoring to determine when to fall back to human intervention

Query enhancement improves the quality of information retrieval.
Reranking ensures the most relevant information is used for generation.
Response validation adds an extra layer of quality control to the output.

Best practices for building RAG systems

Start with a clear purpose:
Define specific goals and success metrics for your RAG system.
Understand the primary use cases and user needs you're addressing.
Ensure high-quality data:
Maintain an accurate and up-to-date knowledge base.
Regularly review and clean your data to ensure relevance and accuracy.
Embrace evaluation and iteration:
Continuously test your system against predefined metrics.
Be prepared to refine and adjust based on performance data and user feedback.
Involve subject matter experts:
Collaborate with domain experts to validate system outputs.
Use their insights to improve data quality and response accuracy.
Prioritize user experience:
Design intuitive interfaces that make it easy for users to interact with the system.
Ensure the system's responses are accurate, helpful, and easy to understand.
Be transparent:
Inform users that they are interacting with an AI system.
Provide citations or sources for the information in responses when possible.
Commit to continuous improvement:
Regularly update the system to adapt to changing user needs and new information.
Stay informed about advancements in RAG technology and incorporate relevant improvements.

Applying the 80/20 rule in RAG systems

Pareto principle in RAG context

Approximately 80% of user queries can be effectively addressed by 20% of your knowledge base.
20% of your development efforts may yield 80% of the system's performance improvements.
80% of user satisfaction might come from getting 20% of critical features right.

Strategies for applying the 80/20 rule

Focus on high-impact data:
Identify and prioritize the most frequently accessed or queried 20% of your knowledge base.
Ensure this core data is of the highest quality, up-to-date, and optimized for retrieval.
Prioritize key features:
Concentrate on developing and refining the features that address the most common user needs.
Implement improvements that greatly benefit system performance and user experience.
Optimize query handling:
Analyze user query patterns to identify the most common types of questions.
Fine-tune your system to excel at handling these frequent queries.
Streamline development cycles:
Focus sprints on tackling the 20% of issues or enhancements that will yield 80% of desired improvements.
Avoid getting bogged down in minor optimizations that offer diminishing returns.

Benefits of applying the 80/20 rule

Efficient resource utilization:
By focusing on high-impact areas, you can achieve significant improvements with less effort.
This approach allows for more effective allocation of development resources.
Faster development cycles:
Prioritizing the most impactful features and improvements leads to quicker iterations and releases.
Users can benefit from meaningful updates more frequently.
Enhanced user satisfaction:
By addressing the most common queries and pain points first, you can quickly improve the overall user experience.
This approach often leads to higher user adoption and satisfaction rates.
Scalable improvement process:
As you tackle the "low-hanging fruit," you create a solid foundation for addressing more complex issues later.
This iterative process allows for continuous, manageable improvements over time.

Wandbot – A working RAG system in production

An introduction to Wandbot

Features and Capabilities

Customer support via chat and email:
Wandbot offers immediate assistance to users by answering questions about Weights & Biases products and services.
It handles various inquiries, from basic troubleshooting to complex technical questions.
Documentation search:
Tech support employees can quickly find relevant information from documents and guides.
This feature reduces the time spent searching for specific details or procedures.
Developer assistance:
Wandbot provides code snippets and API documentation to developers.
It helps streamline development by offering quick access to relevant technical information.

Lessons Learned

Flexibility and adaptability:
Wandbot's development revealed the importance of creating a flexible system that can evolve with changing needs.
We learned to design for scalability, allowing the system to handle an increasing variety of tasks and data sources.
Importance of continuous iteration:
Regular updates and refinements based on user interactions and feedback were crucial for improving Wandbot's performance.
We emphasized the value of an iterative development approach, constantly evaluating and enhancing the system.
Balancing accuracy and speed:
Finding the right balance between response accuracy and latency was an ongoing challenge.
We learned to optimize retrieval and generation processes to maintain high accuracy without sacrificing response time.
Value of tracing interactions:
Incorporating traces of user interactions was essential in identifying areas for improvement and new feature development.
We established processes to regularly trace and analyze user interactions.

Impact and Benefits

Reduced workload on human support teams by handling a significant portion of user inquiries.
Improved response times for customer queries, leading to higher user satisfaction.
Enhanced internal efficiency by providing quick access to documentation and technical information.
Demonstrated the scalability and versatility of RAG systems in handling diverse tasks.

Managing risks and rewards in RAG development

Addressing the challenges of building RAG systems

Handling hallucinations and inaccuracies in LLMs:
Challenge: Large Language Models can sometimes generate plausible-sounding but incorrect information.
Solution: Implement robust fact-checking mechanisms and provide source citations for generated responses.
Example: Wandbot includes source citations in its responses, allowing users to verify information.
Managing technical debt while scaling features:
Challenge: Rapidly adding features can lead to accumulated technical debt, making the system harder to maintain and improve.
Solution: Adopt a balanced approach to development, allocating time for new features and code refactoring.
Strategy: Regularly review and optimize existing code and prioritize modular design for more accessible future enhancements.

Balancing Act

Trade-offs between latency and accuracy:
Challenge: Improving response accuracy often comes at the cost of increased processing time.
Approach: Continuously optimize retrieval and generation processes to balance speed and accuracy.
Example: Wandbot's development involved fine-tuning the number of retrieved documents and the complexity of the generation model to find an optimal balance.
Phased rollouts to mitigate risks:
Strategy: Implement new features or significant changes in stages, starting with a small user group before broader deployment.
Benefit: This approach allows for early detection and correction of issues, minimizing the impact on the overall user base.
Implementation: Wandbot was initially rolled out to free users on Discord before expanding to broader chat and email support.

Realizing Rewards

Enhanced user sxperience with instant, accurate responses:
Benefit: Users receive immediate, relevant answers to their queries, significantly improving their interaction with the product or service.
Impact: This leads to higher user satisfaction and potentially increased user retention.
Reduced workload on support teams:
Advantage: By handling many common queries, RAG systems free up human support teams to focus on more complex issues.
Result: More efficient resource allocation and potentially reduced operational costs.
Global accessibility and availability:
Benefit: RAG systems can provide 24/7 support across different time zones and languages.
Impact: This global reach enhances the user experience for a diverse, international user base.
Continuous learning and improvement:
Advantage: RAG systems can be designed to learn from interactions, continuously improving their performance over time.
Outcome: The system becomes more valuable and effective with increased usage.

Common challenges in RAG system development

Keeping up with rapidly evolving LLMs

Challenge: Large language models are rapidly evolving, with new models and APIs being released frequently.
Impact: Constant changes can complicate system stability and performance.
Strategies:
Implement a modular architecture that allows for easy integration of new models.
Regularly evaluate new LLMs and APIs to assess potential improvements.
Develop a systematic approach for testing and integrating updates without disrupting existing functionality.

Balancing feature development and system refinement

Challenge: There's often pressure to add new features while also needing to refine and optimize existing components.
Impact: Focusing too much on new features can lead to a bloated, unstable system while over-focusing on refinement can slow innovation.
Solution:
Use evaluation frameworks to prioritize developments that offer the most value.
Implement a balanced development cycle that allocates new features and system optimization time.
Regularly reassess priorities based on user feedback and performance metrics.

Creating representative evaluation datasets

Challenge: Developing datasets that accurately represent real-world queries and scenarios is crucial but challenging.
Impact: Inadequate evaluation data can lead to poor system performance in production environments.
Approach:
Combine automated data collection methods with expert analysis of chat logs.
Continuously update evaluation datasets to reflect changing user needs and query patterns.
Involve subject matter experts in the creation and validation of evaluation data.

Navigating trade offs between latency and accuracy

Challenge: Improving response accuracy often comes at the cost of increased latency, and vice versa.
Impact: Poor balance can result in slow responses or inaccurate information, frustrating users.
Techniques:
Implement efficient retrieval algorithms to reduce response times without sacrificing relevance.
Use caching mechanisms for frequently asked queries to improve response speed.
Explore techniques like query optimization and parallel processing to enhance performance.

Ensuring continuous evolution

Challenge: RAG systems must adapt to changing information, user needs, and technological advancements.
Impact: Failure to evolve can lead to outdated responses and decreased system effectiveness.
Solutions:
Regular Dataset Updates:
Implement automated processes to keep the knowledge base current.
Establish protocols for reviewing and incorporating new information.
Granular Evaluations:
Conduct detailed assessments of each component in the RAG pipeline.
Use metrics that reflect real-world performance and user satisfaction.
Incorporating User Feedback:
Develop mechanisms to collect and analyze user interactions and feedback.
Use this information to guide system improvements and feature development.
Fine-tuning the Entire Pipeline:
Regularly optimize each stage of the RAG process, from retrieval to response generation.
Implement A/B testing to evaluate the impact of changes before full deployment.

Key takeaways

Integration of LLMs and information retrieval: RAG effectively combines large language models with traditional information retrieval techniques, ensuring current and contextually relevant responses.
Addressing LLM imitations: Using real-time data retrieval, RAG overcomes the static knowledge and hallucination issues commonly associated with LLMs, leading to more accurate and informed outputs.
Core process and enhancements: The RAG pipeline involves query processing, data retrieval, context integration, and response generation. Advanced components like query enhancement, reranking, and response validation improve the pipeline's accuracy and reliability.
Strategic sevelopment practices: Successful RAG systems benefit from clear goals, high-quality data, continuous iteration, user experience prioritization, and transparency. Applying the 80/20 rule can optimize resource allocation and accelerate development.
Real-world application through Wandbot: The Wandbot case study illustrates practical RAG implementation, highlighting the importance of flexibility, iteration, and balancing accuracy with latency. It demonstrates how RAG can enhance customer support and internal efficiencies.
Challenges and opportunities: Developing RAG systems involves managing technical debt, evolving LLMs, feature balancing, latency-accuracy trade-offs, and continuous adaptation to changing information and user needs.

Conclusion