{"id":26515873,"url":"https://github.com/fareedkhan-dev/rag-with-rl","last_synced_at":"2026-04-04T20:32:55.263Z","repository":{"id":283488980,"uuid":"951932591","full_name":"FareedKhan-dev/rag-with-rl","owner":"FareedKhan-dev","description":"Maximizing the Performance of a Simple RAG using RL","archived":false,"fork":false,"pushed_at":"2025-03-20T14:09:41.000Z","size":103,"stargazers_count":63,"open_issues_count":0,"forks_count":16,"subscribers_count":2,"default_branch":"main","last_synced_at":"2025-07-04T03:05:38.378Z","etag":null,"topics":["llm","openai","python","rag","reinforcement-learning","rl"],"latest_commit_sha":null,"homepage":"https://medium.com/@fareedkhandev/maximizing-simple-rag-performance-using-rl-rewards-in-python-d4c14cbadf59","language":"Jupyter Notebook","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"mit","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/FareedKhan-dev.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":"LICENSE","code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null,"dei":null,"publiccode":null,"codemeta":null}},"created_at":"2025-03-20T13:21:54.000Z","updated_at":"2025-06-27T06:26:41.000Z","dependencies_parsed_at":"2025-03-21T14:49:20.582Z","dependency_job_id":null,"html_url":"https://github.com/FareedKhan-dev/rag-with-rl","commit_stats":null,"previous_names":["fareedkhan-dev/rag-with-rl"],"tags_count":0,"template":false,"template_full_name":null,"purl":"pkg:github/FareedKhan-dev/rag-with-rl","repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/FareedKhan-dev%2Frag-with-rl","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/FareedKhan-dev%2Frag-with-rl/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/FareedKhan-dev%2Frag-with-rl/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/FareedKhan-dev%2Frag-with-rl/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/FareedKhan-dev","download_url":"https://codeload.github.com/FareedKhan-dev/rag-with-rl/tar.gz/refs/heads/main","sbom_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/FareedKhan-dev%2Frag-with-rl/sbom","scorecard":null,"host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":286080680,"owners_count":31413269,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2026-04-04T20:09:54.854Z","status":"ssl_error","status_checked_at":"2026-04-04T20:09:44.350Z","response_time":60,"last_error":"SSL_read: unexpected eof while reading","robots_txt_status":"success","robots_txt_updated_at":"2025-07-24T06:49:26.215Z","robots_txt_url":"https://github.com/robots.txt","online":false,"can_crawl_api":true,"host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":["llm","openai","python","rag","reinforcement-learning","rl"],"created_at":"2025-03-21T06:18:38.610Z","updated_at":"2026-04-04T20:32:55.235Z","avatar_url":"https://github.com/FareedKhan-dev.png","language":"Jupyter Notebook","funding_links":[],"categories":[],"sub_categories":[],"readme":"\u003c!-- omit in toc --\u003e\n# Simple RAG with RL\n\n[![Python 3.7+](https://img.shields.io/badge/python-3.7+-blue.svg)](https://www.python.org/downloads/release/python-370/) [![Nebius AI](https://img.shields.io/badge/Nebius%20AI-LLM-brightgreen)](https://cloud.nebius.ai/services/llm-embedding) [![OpenAI](https://img.shields.io/badge/OpenAI-API-lightgrey)](https://openai.com/) [![Medium](https://img.shields.io/badge/Medium-Blog-black?logo=medium)](https://medium.com/@fareedkhandev/maximizing-simple-rag-performance-using-rl-rewards-in-python-d4c14cbadf59)\n\nA simple RAG works in three simple steps:\n\n1. **Indexing**: Break documents into chunks and convert to vector embeddings.\n\n2. **Retrieval**: When a question is asked, find the most relevant chunks.\n\n3. **Generation**: Combine the question with retrieved chunks and let the AI generate an answer using this information.\n\nThe actual problem is to generate an answer to a given question using the provided documents. Simple RAG often fails to generate accurate answers due to the lack of context in the retrieved chunks. In this notebook, we will use the `RL RAG` approach to generate answers to the given questions using the provided documents.\n\n\u003c!-- omit in toc --\u003e\n# Table of Contents\n\n- [Setting Up the Environment](../main.md#setting-up-the-environment)\n- [Data Preprocessing](../main.md#data-preprocessing)\n- [Document Embedding Generation](../main.md#document-embedding-generation)\n- [Vector Store Implementation](../main.md#vector-store-implementation)\n- [Simple Retrieval Implementation](../main.md#simple-retrieval-implementation)\n- [LLM Response Generation](../main.md#llm-response-generation)\n- [Basic RAG Pipeline](../main.md#basic-rag-pipeline)\n- [Evaluate the basic RAG pipeline](../main.md#evaluate-the-basic-rag-pipeline)\n- [Reinforcement Learning for RAG](../main.md#reinforcement-learning-for-rag)\n- [State, Action Space, and Reward Methodology](../main.md#state-action-space-and-reward-methodology)\n- [Action Function Logic](../main.md#action-function-logic)\n- [Policy Network](../main.md#policy-network)\n- [Single RL Step](../main.md#single-rl-step)\n- [Training Parameters and Policy Update](../main.md#training-parameters-and-policy-update)\n- [Training Loop](../main.md#training-loop)\n- [Performance Comparison Logic](../main.md#performance-comparison-logic)\n- [Evaluation Framework (**Optional**)](../main.md#evaluation-framework-optional)\n- [Evaluating (RL vs Simple) RAG](../main.md#evaluating-rl-vs-simple-rag)\n- [Saving the Comparison Results](../main.md#saving-the-comparison-results)\n- [What can we conclude?](../main.md#what-can-we-conclude)\n\n## Setting Up the Environment\n\nClone the repository and install the required libraries.\n\n```bash\n!git clone https://github.com/FareedKhan-dev/rag-with-rl.git\n```\n\nInstall the required libraries.\n\n```bash\n!pip install -r rag-with-rl/requirements.txt\n```\n\nFirst, we need to import the necessary libraries and set up the environment. We will be using HuggingFace Models hosted under **Nebius** platform. Obviously, you can use your own models as long as they are compatible with OpenAI's API.\n\n\n```python\n# Importing the os module for interacting with the operating system\nimport os\n\n# Importing the OpenAI module for working with OpenAI's API\nfrom openai import OpenAI\n\n# Importing numpy for numerical operations\nimport numpy as np\n\n# Importing json for working with JSON data\nimport json\n\n# Typing module for type hints\nfrom typing import Dict, List, Tuple, Optional, Union\n```\n\nNext, we need to initialize the client responsible for response and embedding generation.\n\n\n```python\n# Set up the API connection using the OpenAI client\n# Replace the base_url and api_key with your own values\n\nclient = OpenAI(\n    base_url=\"https://api.studio.nebius.com/v1/\",  # Base URL for (eg. ollama api, anyother llm api provider)\n    api_key= os.environ[\"OPENAI_API_KEY\"]  # API key for authentication \n)\n```\n\n## Data Preprocessing\nNow that we have moved onto the data preprocessing stage, we need to load the data and preprocess it. Let's create a function that will load all the `.txt` files from a directory and return a list of documents.\n\n\n```python\n# Function to load documents from a directory\ndef load_documents(directory_path: str) -\u003e List[str]:\n    \"\"\"\n    Load all text documents from the specified directory.\n\n    Args:\n        directory_path (str): Path to the directory containing text files.\n\n    Returns:\n        List[str]: A list of strings, where each string is the content of a text file.\n    \"\"\"\n    documents = []  # Initialize an empty list to store document contents\n    for filename in os.listdir(directory_path):  # Iterate through all files in the directory\n        if filename.endswith(\".txt\"):  # Check if the file has a .txt extension\n            # Open the file in read mode with UTF-8 encoding and append its content to the list\n            with open(os.path.join(directory_path, filename), 'r', encoding='utf-8') as file:\n                documents.append(file.read())\n    return documents  # Return the list of document contents\n```\n\nWe need to create a function that performs chunking of the documents once they are loaded. We are using a `chunk_size` of `100` characters, but you can adjust it as per your requirements.\n\n\n```python\n# Function to split documents into chunks\ndef split_into_chunks(documents: List[str], chunk_size: int = 30) -\u003e List[str]:\n    \"\"\"\n    Split documents into smaller chunks of specified size.\n\n    Args:\n        documents (List[str]): A list of document strings to be split into chunks.\n        chunk_size (int): The maximum number of words in each chunk. Default is 100.\n\n    Returns:\n        List[str]: A list of chunks, where each chunk is a string containing up to `chunk_size` words.\n    \"\"\"\n    chunks = []  # Initialize an empty list to store the chunks\n    for doc in documents:  # Iterate through each document\n        words = doc.split()  # Split the document into words\n        # Create chunks of the specified size\n        for i in range(0, len(words), chunk_size):\n            chunk = \" \".join(words[i:i + chunk_size])  # Join words to form a chunk\n            chunks.append(chunk)  # Add the chunk to the list\n    return chunks  # Return the list of chunks\n```\n\nThis step is **optional**, where we preprocess each chunk by removing special characters, converting to lowercase, etc.\n\n\n```python\n# Function to preprocess text (e.g., lowercasing, removing special characters)\ndef preprocess_text(text: str) -\u003e str:\n    \"\"\"\n    Preprocess the input text by converting it to lowercase and removing special characters.\n\n    Args:\n        text (str): The input text to preprocess.\n\n    Returns:\n        str: The preprocessed text with only alphanumeric characters and spaces.\n    \"\"\"\n    # Convert the text to lowercase\n    text = text.lower()\n    # Remove special characters, keeping only alphanumeric characters and spaces\n    text = ''.join(char for char in text if char.isalnum() or char.isspace())\n    return text\n```\n\nHowever, if you are using the previous preprocessing step, you can simply create a function to preprocess the entire document.\n\n\n```python\n# Function to preprocess all chunks\ndef preprocess_chunks(chunks: List[str]) -\u003e List[str]:\n    \"\"\"\n    Apply preprocessing to all text chunks.\n\n    Args:\n        chunks (List[str]): A list of text chunks to preprocess.\n\n    Returns:\n        List[str]: A list of preprocessed text chunks.\n    \"\"\"\n    # Apply the preprocess_text function to each chunk in the list\n    return [preprocess_text(chunk) for chunk in chunks]\n```\n\nNow that we have implemented all the functions for data preprocessing, we can load the documents from the directory, split them into chunks, and preprocess the chunks.\n\n\n```python\n# Specify the directory path containing the text files\ndirectory_path = \"data\"\n\n# Load all text documents from the specified directory\ndocuments = load_documents(directory_path)\n\n# Split the loaded documents into smaller chunks of text\nchunks = split_into_chunks(documents)\n\n# Preprocess the chunks (e.g., lowercasing, removing special characters)\npreprocessed_chunks = preprocess_chunks(chunks)\n```\n\nPrint the first 200 characters of the first two chunks\n\n\n```python\n# Print the first 2 preprocessed chunks, displaying only the first 200 characters of each chunk\nfor i in range(2):\n    # Use slicing to limit the output to the first 200\n    print(f\"Chunk {i+1}: {preprocessed_chunks[i][:50]} ... \")\n    print(\"-\" * 50)  # Print a separator line\n```\n\n    Chunk 1: quantum computing principles progress and possibil ... \n    --------------------------------------------------\n    Chunk 2: process information in binary digits bits quantum  ... \n    --------------------------------------------------\n    \n\n## Document Embedding Generation\n\nIn the previous step, we chunked our document. Now it's time to generate embeddings for the chunk dataset. When working with RAG, our knowledge base is typically quite large. Therefore, we need to perform embedding generation in batches. Let's create a core function to generate embeddings for the chunks in batches.\n\nThe embedding model we are using is `BAAI/bge-en-icl`.\n\n\n```python\n# Function to generate embeddings for a single batch of text chunks\ndef generate_embeddings_batch(chunks_batch: List[str], model: str = \"BAAI/bge-en-icl\") -\u003e List[List[float]]:\n    \"\"\"\n    Generate embeddings for a batch of text chunks using the OpenAI client.\n\n    Args:\n        chunks_batch (List[str]): A batch of text chunks to generate embeddings for.\n        model (str): The model to use for embedding generation. Default is \"BAAI/bge-en-icl\".\n\n    Returns:\n        List[List[float]]: A list of embeddings, where each embedding is a list of floats.\n    \"\"\"\n    # Use the OpenAI client to create embeddings for the input batch\n    response = client.embeddings.create(\n        model=model,  # Specify the model to use for embedding generation\n        input=chunks_batch  # Provide the batch of text chunks as input\n    )\n    # Extract embeddings from the response and return them\n    embeddings = [item.embedding for item in response.data]\n    return embeddings\n```\n\nNext, we will define a function to generate embeddings for all text chunks in batches. This function will take a list of text chunks as input and generate embeddings for each batch of chunks using the OpenAI client. The function will return a list of embeddings corresponding to all the text chunks.\n\n\n```python\n# Function to generate embeddings for all chunks with batching\ndef generate_embeddings(chunks: List[str], batch_size: int = 10) -\u003e np.ndarray:\n    \"\"\"\n    Generate embeddings for all text chunks in batches.\n\n    Args:\n        chunks (List[str]): A list of text chunks to generate embeddings for.\n        batch_size (int): The number of chunks to process in each batch. Default is 10.\n\n    Returns:\n        np.ndarray: A NumPy array containing embeddings for all chunks.\n    \"\"\"\n    all_embeddings = []  # Initialize an empty list to store all embeddings\n\n    # Iterate through the chunks in batches\n    for i in range(0, len(chunks), batch_size):\n        # Extract the current batch of chunks\n        batch = chunks[i:i + batch_size]\n        # Generate embeddings for the current batch\n        embeddings = generate_embeddings_batch(batch)\n        # Extend the list of all embeddings with the embeddings from the current batch\n        all_embeddings.extend(embeddings)\n\n    # Convert the list of embeddings to a NumPy array and return it\n    return np.array(all_embeddings)\n```\n\nLet's create another function to save the embeddings to a file in JSON format.\n\n\n```python\n# Function to save embeddings to a file\ndef save_embeddings(embeddings: np.ndarray, output_file: str) -\u003e None:\n    \"\"\"\n    Save embeddings to a JSON file.\n\n    Args:\n        embeddings (np.ndarray): A NumPy array containing the embeddings to save.\n        output_file (str): The path to the output JSON file where embeddings will be saved.\n\n    Returns:\n        None\n    \"\"\"\n    # Open the specified file in write mode with UTF-8 encoding\n    with open(output_file, 'w', encoding='utf-8') as file:\n        # Convert the NumPy array to a list and save it as JSON\n        json.dump(embeddings.tolist(), file)\n```\n\nNow that we have implemented all the functions for embedding generation, we can proceed to generate embeddings for the preprocessed text chunks and save them to a JSON file.\n\n\n```python\n# Ensure the chunks are preprocessed before generating embeddings\npreprocessed_chunks = preprocess_chunks(chunks)\n\n# Generate embeddings for the preprocessed chunks\nembeddings = generate_embeddings(preprocessed_chunks)\n\n# Save the generated embeddings to a JSON file named \"embeddings.json\"\nsave_embeddings(embeddings, \"embeddings.json\")\n```\n\n## Vector Store Implementation\nSince we are not using any python libraries for vector storage, we will implement a simple vector store using a dictionary.\n\n\n```python\n# Initialize an in-memory vector store as a dictionary\n# The keys will be unique identifiers (integers), and the values will be dictionaries containing embeddings and corresponding text chunks\nvector_store: dict[int, dict[str, object]] = {}\n\n# Function to add embeddings and corresponding text chunks to the vector store\ndef add_to_vector_store(embeddings: np.ndarray, chunks: List[str]) -\u003e None:\n    \"\"\"\n    Add embeddings and their corresponding text chunks to the vector store.\n\n    Args:\n        embeddings (np.ndarray): A NumPy array containing the embeddings to add.\n        chunks (List[str]): A list of text chunks corresponding to the embeddings.\n\n    Returns:\n        None\n    \"\"\"\n    # Iterate over embeddings and chunks simultaneously\n    for embedding, chunk in zip(embeddings, chunks):\n        # Add each embedding and its corresponding chunk to the vector store\n        # Use the current length of the vector store as the unique key\n        vector_store[len(vector_store)] = {\"embedding\": embedding, \"chunk\": chunk}\n```\n\n## Simple Retrieval Implementation\n\nWe do know for retrieving the most similar text chunks to a given query, we can use the cosine similarity between the query embedding and the embeddings of all text chunks. The higher the cosine similarity, the more similar the text chunks are. We can then sort the chunks based on their similarity scores and return the top-k most similar chunks.\n    \nSo, let's implement a simple cosine similarity-based retrieval function.\n\nThe cosine similarity between two vectors $A$ and $B$ is calculated as:\n\n$$\\text{cosine similarity} = \\frac{A \\cdot B}{||A|| \\times ||B||} = \\frac{\\sum_{i=1}^{n} A_i B_i}{\\sqrt{\\sum_{i=1}^{n} A_i^2} \\times \\sqrt{\\sum_{i=1}^{n} B_i^2}}$$\n\nWhere:\n- $A \\cdot B$ is the dot product of vectors $A$ and $B$\n- $||A||$ and $||B||$ are the Euclidean norms (magnitudes) of the vectors\n- $n$ is the dimension of the vectors\n\n\n```python\n# Function to compute cosine similarity between two vectors\ndef cosine_similarity(vec1: np.ndarray, vec2: np.ndarray) -\u003e float:\n    \"\"\"\n    Compute the cosine similarity between two vectors.\n\n    Args:\n        vec1 (np.ndarray): The first vector.\n        vec2 (np.ndarray): The second vector.\n\n    Returns:\n        float: The cosine similarity between the two vectors, ranging from -1 to 1.\n    \"\"\"\n    # Compute the dot product of the two vectors\n    dot_product = np.dot(vec1, vec2)\n    # Compute the magnitude (norm) of the first vector\n    norm_vec1 = np.linalg.norm(vec1)\n    # Compute the magnitude (norm) of the second vector\n    norm_vec2 = np.linalg.norm(vec2)\n    # Return the cosine similarity as the ratio of the dot product to the product of the norms\n    return dot_product / (norm_vec1 * norm_vec2)\n```\n\nWhen we calculate the cosine similarity between a query and all the chunks, we can perform a similarity search. Based on the `top_k` parameter, we retrieve the top k most similar chunks.\n\n\n```python\n# Function to perform similarity search in the vector store\ndef similarity_search(query_embedding: np.ndarray, top_k: int = 5) -\u003e List[str]:\n    \"\"\"\n    Perform similarity search in the vector store and return the top_k most similar chunks.\n\n    Args:\n        query_embedding (np.ndarray): The embedding vector of the query.\n        top_k (int): The number of most similar chunks to retrieve. Default is 5.\n\n    Returns:\n        List[str]: A list of the top_k most similar text chunks.\n    \"\"\"\n    similarities = []  # Initialize a list to store similarity scores and corresponding keys\n\n    # Iterate through all items in the vector store\n    for key, value in vector_store.items():\n        # Compute the cosine similarity between the query embedding and the stored embedding\n        similarity = cosine_similarity(query_embedding, value[\"embedding\"])\n        # Append the key and similarity score as a tuple to the list\n        similarities.append((key, similarity))\n\n    # Sort the list of similarities in descending order based on the similarity score\n    similarities = sorted(similarities, key=lambda x: x[1], reverse=True)\n\n    # Retrieve the top_k most similar chunks based on their keys\n    return [vector_store[key][\"chunk\"] for key, _ in similarities[:top_k]]\n```\n\nOnce we have the similarity search function ready, we can simply code a retrieval function on top of it that will provide the relevant chunks based on the query.\n\n\n```python\n# Function to retrieve relevant document chunks for a query\ndef retrieve_relevant_chunks(query_text: str, top_k: int = 5) -\u003e List[str]:\n    \"\"\"\n    Retrieve the most relevant document chunks for a given query text.\n\n    Args:\n        query_text (str): The query text for which relevant chunks are to be retrieved.\n        top_k (int): The number of most relevant chunks to retrieve. Default is 5.\n\n    Returns:\n        List[str]: A list of the top_k most relevant text chunks.\n    \"\"\"\n    # Generate embedding for the query text using the embedding model\n    query_embedding = generate_embeddings([query_text])[0]\n    \n    # Perform similarity search to find the most relevant chunks\n    relevant_chunks = similarity_search(query_embedding, top_k=top_k)\n    \n    # Return the list of relevant chunks\n    return relevant_chunks\n```\n\nNow that we have implemented all the functions for retrieval, we can proceed to test the retrieval system with a sample query.\n\n\n```python\n# Add the generated embeddings and their corresponding preprocessed chunks to the vector store\nadd_to_vector_store(embeddings, preprocessed_chunks)\n\n# Define a query text for which we want to retrieve relevant document chunks\nquery_text = \"What is Quantum Computing?\"\n\n# Retrieve the most relevant chunks from the vector store based on the query text\nrelevant_chunks = retrieve_relevant_chunks(query_text)\n\n# Print the first 50 characters of each retrieved relevant chunk\nfor idx, chunk in enumerate(relevant_chunks):\n    print(f\"Chunk {idx + 1}: {chunk[:50]} ... \")\n    print(\"-\" * 50)  # Print a separator line\n```\n\n    Chunk 1: quantum computing principles progress and possibil ... \n    --------------------------------------------------\n    Chunk 2: through distinct stages 1 nisq era current 2 error ... \n    --------------------------------------------------\n    Chunk 3: quantum advantage and practical applications quant ... \n    --------------------------------------------------\n    Chunk 4: process information in binary digits bits quantum  ... \n    --------------------------------------------------\n    Chunk 5: measuring the correct answer quantum gates and cir ... \n    --------------------------------------------------\n    \n\n## LLM Response Generation\n\nWhen we have a query and a set of relevant document chunks, we can use a large language model (LLM) to generate a response based on the query and the retrieved information. In this section, we will use the OpenAI API to generate a response to a query by providing the query text and the relevant document chunks as context to the LLM.\n\nFirst we need a function to construct the input prompt for the LLM, which includes the query text and the relevant document chunks as context.\n\n\n```python\n# Function to construct a prompt with context\ndef construct_prompt(query: str, context_chunks: List[str]) -\u003e str:\n    \"\"\"\n    Construct a prompt by combining the query with the retrieved context chunks.\n\n    Args:\n        query (str): The query text for which the prompt is being constructed.\n        context_chunks (List[str]): A list of relevant context chunks to include in the prompt.\n\n    Returns:\n        str: The constructed prompt to be used as input for the LLM.\n    \"\"\"\n    # Combine all context chunks into a single string, separated by newlines\n    context = \"\\n\".join(context_chunks)\n    \n    # Define the system message to guide the LLM's behavior\n    system_message = (\n        \"You are a helpful assistant. Only use the provided context to answer the question. \"\n        \"If the context doesn't contain the information needed, say 'I don't have enough information to answer this question.'\"\n    )\n    \n    # Construct the final prompt by combining the system message, context, and query\n    prompt = f\"System: {system_message}\\n\\nContext:\\n{context}\\n\\nQuestion:\\n{query}\\n\\nAnswer:\"\n    \n    return prompt\n```\n\nTo generate an LLM response, we need to implement a function that takes the constructed input prompt and sends it to the OpenAI API for response generation.\n\n\n```python\n# Function to generate a response using the OpenAI chat model\ndef generate_response(\n    prompt: str,\n    model: str = \"google/gemma-2-2b-it\",\n    max_tokens: int = 512,\n    temperature: float = 1,\n    top_p: float = 0.9,\n    top_k: int = 50\n) -\u003e str:\n    \"\"\"\n    Generate a response from the OpenAI chat model based on the constructed prompt.\n\n    Args:\n        prompt (str): The input prompt to provide to the chat model.\n        model (str): The model to use for generating the response. Default is \"google/gemma-2-2b-it\".\n        max_tokens (int): Maximum number of tokens in the response. Default is 512.\n        temperature (float): Sampling temperature for response diversity. Default is 0.5.\n        top_p (float): Probability mass for nucleus sampling. Default is 0.9.\n        top_k (int): Number of highest probability tokens to consider. Default is 50.\n\n    Returns:\n        str: The generated response from the chat model.\n    \"\"\"\n    # Use the OpenAI client to create a chat completion\n    response = client.chat.completions.create(\n        model=model,  # Specify the model to use for generating the response\n        max_tokens=max_tokens,  # Maximum number of tokens in the response\n        temperature=temperature,  # Sampling temperature for response diversity\n        top_p=top_p,  # Probability mass for nucleus sampling\n        extra_body={  # Additional parameters for the request\n            \"top_k\": top_k  # Number of highest probability tokens to consider\n        },\n        messages=[  # List of messages to provide context for the chat model\n            {\n                \"role\": \"user\",  # Role of the message sender (user in this case)\n                \"content\": [  # Content of the message\n                    {\n                        \"type\": \"text\",  # Type of content (text in this case)\n                        \"text\": prompt  # The actual prompt text\n                    }\n                ]\n            }\n        ]\n    )\n    # Return the content of the first choice in the response\n    return response.choices[0].message.content\n```\n\n## Basic RAG Pipeline\n\nWe cannot run small pieces of code repeatedly. Therefore, we need to create a simple RAG pipeline that takes only one parameter, which is our query, and returns the LLM response.\n\n\n```python\n# Function to implement the basic Retrieval-Augmented Generation (RAG) pipeline\ndef basic_rag_pipeline(query: str) -\u003e str:\n    \"\"\"\n    Implement the basic Retrieval-Augmented Generation (RAG) pipeline:\n    retrieve relevant chunks, construct a prompt, and generate a response.\n\n    Args:\n        query (str): The input query for which a response is to be generated.\n\n    Returns:\n        str: The generated response from the LLM based on the query and retrieved context.\n    \"\"\"\n    # Step 1: Retrieve the most relevant chunks for the given query\n    relevant_chunks: List[str] = retrieve_relevant_chunks(query)\n    \n    # Step 2: Construct a prompt using the query and the retrieved chunks\n    prompt: str = construct_prompt(query, relevant_chunks)\n    \n    # Step 3: Generate a response from the LLM using the constructed prompt\n    response: str = generate_response(prompt)\n    \n    # Return the generated response\n    return response\n```\n\n## Evaluate the basic RAG pipeline\n\nNow that we have coded the basic RAG pipeline, we can use it for evaluation. Our evaluation queries contain different targeted segments, such as `factual_queries` and `complex_nature`. We are going to test the factual knowledge of our RAG pipeline.\n\nLet's load our evaluation queries and their expected answers.\n\n\n```python\n# Open the validation data file in read mode and load its content as a dictionary\nwith open('data/val.json', 'r') as file:\n    validation_data = json.load(file)\n\n# Test the basic RAG pipeline with a sample query\nsample_query = validation_data['basic_factual_questions'][0]['question']  # Extract the query text\nexpected_answer = validation_data['basic_factual_questions'][0]['answer']  # Extract the ground truth answer\n\n# print the sample query and expected answer\nprint(f\"Sample Query: {sample_query}\\n\")\nprint(f\"Expected Answer: {expected_answer}\\n\")\n```\n\n    Sample Query: What is the mathematical representation of a qubit in superposition?\n    \n    Expected Answer: |ψ⟩ = α|0⟩ + β|1⟩, where α and β are complex numbers satisfying |α|² + |β|² = 1, representing the probability amplitudes for measuring the qubit in state |0⟩ or |1⟩ respectively.\n    \n    \n\nLet's test the basic RAG pipeline with this eval query and see how well it performs.\n\n\n```python\n# Print a message to indicate the start of the RAG pipeline\nprint(\"🔍 Running the Retrieval-Augmented Generation (RAG) pipeline...\")\nprint(f\"📥 Query: {sample_query}\\n\")\n\n# Run the RAG pipeline and get the response\nresponse = basic_rag_pipeline(sample_query)\n\n# Print the response with better formatting\nprint(\"🤖 AI Response:\")\nprint(\"-\" * 50)\nprint(response.strip())\nprint(\"-\" * 50)\n\n# Print the ground truth answer for comparison\nprint(\"✅ Ground Truth Answer:\")\nprint(\"-\" * 50)\nprint(expected_answer)\nprint(\"-\" * 50)\n```\n\n    🔍 Running the Retrieval-Augmented Generation (RAG) pipeline...\n    📥 Query: What is the mathematical representation of a qubit in superposition?\n    \n    🤖 AI Response:\n    --------------------------------------------------\n    ψ  α0  β1\n    --------------------------------------------------\n    ✅ Ground Truth Answer:\n    --------------------------------------------------\n    |ψ⟩ = α|0⟩ + β|1⟩, where α and β are complex numbers satisfying |α|² + |β|² = 1, representing the probability amplitudes for measuring the qubit in state |0⟩ or |1⟩ respectively.\n    --------------------------------------------------\n    \n\nThe simple RAG pipeline doesn't seem to perform well in its current state. The generated response is not only irrelevant to the ground truth but also misses critical information.\n\nBut don't worry! In the upcoming steps, we will implement a Reinforcement Learning-based RAG pipeline to address these shortcomings. This will help us improve the retrieval and generation process, making the responses more accurate and contextually relevant.\n\nStay tuned as we take our RAG pipeline to the next level! 🚀\n\n## Reinforcement Learning for RAG\n\nReinforcement Learning (RL) is a type of machine learning where an agent learns to make decisions by taking actions in an environment to maximize some notion of cumulative reward. Unlike supervised learning, the agent is not explicitly told which actions to take, but instead must discover which actions yield the most reward through trial and error.\n\nFollow are the main components of a reinforcement learning system:\n\n1. **Agent**: The learner or decision-maker\n2. **Environment**: The world with which the agent interacts\n3. **State (S)**: The current situation of the agent in the environment\n4. **Action (A)**: A set of possible moves the agent can make\n5. **Reward (R)**: Feedback from the environment after each action\n6. **Policy (π)**: Strategy that the agent follows to determine the next action\n\nThe goal in reinforcement learning is to learn a policy π that maximizes the expected cumulative reward:\n\n$$\\pi^* = \\arg\\max_\\pi \\mathbb{E}\\left[ \\sum_{t=0}^{T} \\gamma^t R_t \\right]$$\n\nWhere:\n- $\\pi^*$ is the optimal policy\n- $\\gamma$ is the discount factor (0 ≤ γ ≤ 1)\n- $R_t$ is the reward at time step t\n- $T$ is the final time step\n\nIn the context of RAG systems, reinforcement learning can be used to:\n- Improve retrieval by learning which documents are most helpful\n- Refine prompt construction based on user feedback\n- Optimize the generation process by learning from successful responses\n\n## State, Action Space, and Reward Methodology\n\nThe very first step when coding an RL algorithm is to define three things:\n\n- **State**: It is the current situation of the environment. In our case, the initial state is our simple RAG pipeline (query, context, response).\n- **Action Space**: It is the decision that the agent takes based on the state. In our case, the actions can include changing the model, modifying the context, altering the query, etc.\n- **Reward**: It is the feedback that the agent receives after taking an action. In our case, the reward can be the similarity between the generated response and the ground truth answer.\n\nOur state will be changing constantly as we perform training. For that, we need to save the state after each `training episode` so that our RL agent can learn from it and avoid making the same mistakes again.\n\n\n```python\n# Function to define the state representation for reinforcement learning\ndef define_state(\n    query: str, \n    context_chunks: List[str], \n    rewritten_query: str = None, \n    previous_responses: List[str] = None, \n    previous_rewards: List[float] = None\n) -\u003e dict:\n    \"\"\"\n    Define the state representation for the reinforcement learning agent.\n    \n    Args:\n        query (str): The original user query.\n        context_chunks (List[str]): Retrieved context chunks from the knowledge base.\n        rewritten_query (str, optional): A reformulated version of the original query.\n        previous_responses (List[str], optional): List of previously generated responses.\n        previous_rewards (List[float], optional): List of rewards received for previous actions.\n    \n    Returns:\n        dict: A dictionary representing the current state with all relevant information.\n    \"\"\"\n    state = {\n        \"original_query\": query,                                    # The initial query from the user\n        \"current_query\": rewritten_query if rewritten_query else query,  # Current version of the query (may be rewritten)\n        \"context\": context_chunks,                                 # Retrieved context chunks from the knowledge base\n        \"previous_responses\": previous_responses if previous_responses else [],  # History of generated responses\n        \"previous_rewards\": previous_rewards if previous_rewards else []         # History of received rewards\n    }\n    return state\n```\n\nWe have defined the state representation for the RL agent, including the user query, retrieved context chunks, rewritten query (if any), and histories of responses and rewards. This state will guide the agent in generating better responses. \n\nNext we need to define the action space for the reinforcement learning agent. The action space consists of the set of possible actions that the agent can take at each step. In this case, we define four actions:\n- `rewrite_query`: Reformulate the original query to improve retrieval\n- `expand_context`: Retrieve additional context chunks\n- `filter_context`: Remove irrelevant context chunks\n- `generate_response`: Generate a response based on the current query and context\n\n\n```python\n# Function to define the action space for reinforcement learning\ndef define_action_space() -\u003e List[str]:\n    \"\"\"\n    Define the set of possible actions the reinforcement learning agent can take.\n    \n    Actions include:\n    - rewrite_query: Reformulate the original query to improve retrieval\n    - expand_context: Retrieve additional context chunks\n    - filter_context: Remove irrelevant context chunks\n    - generate_response: Generate a response based on current query and context\n    \n    Returns:\n        List[str]: A list of available actions.\n    \"\"\"\n\n    # Define the set of actions the agent can take\n    actions = [\"rewrite_query\", \"expand_context\", \"filter_context\", \"generate_response\"]\n    return actions\n```\n\nObviously, when our RL agent takes an action, it will be based on the current state and the action space. It will be rewarded based on the quality of the response generated by the RAG pipeline. The reward function will be based on the cosine similarity between the generated response and the ground truth answer.\n\n\n```python\n# Function to calculate the reward based on response quality\ndef calculate_reward(response: str, ground_truth: str) -\u003e float:\n    \"\"\"\n    Calculate a reward value by comparing the generated response to the ground truth.\n    \n    Uses cosine similarity between the embeddings of the response and ground truth\n    to determine how close the response is to the expected answer.\n    \n    Args:\n        response (str): The generated response from the RAG pipeline.\n        ground_truth (str): The expected correct answer.\n    \n    Returns:\n        float: A reward value between -1 and 1, where higher values indicate \n               greater similarity to the ground truth.\n    \"\"\"\n    # Generate embeddings for both the response and ground truth\n    response_embedding = generate_embeddings([response])[0]\n    ground_truth_embedding = generate_embeddings([ground_truth])[0]\n    \n    # Calculate cosine similarity between the embeddings as the reward\n    similarity = cosine_similarity(response_embedding, ground_truth_embedding)\n    return similarity\n```\n\nOur goal is to maximize the reward by generating responses that are similar to the ground truth answer. Higher reward values indicate that the generated response is more aligned with the expected answer.\n\n## Action Function Logic\n\nNow that we have defined the action space, we need to implement the logic for each action. This logic will determine how the RAG pipeline should be modified based on the action taken by the RL agent.\n\nJust to revisit, the four actions are:\n- `rewrite_query`: Reformulate the original query to improve retrieval\n- `expand_context`: Retrieve additional context chunks\n- `filter_context`: Remove irrelevant context chunks\n- `generate_response`: Generate a response based on the current query and context\n\nLet's create our first action logic for the agent. The first action we will implement is the `rewrite_query` action, which involves reformulating the original user query to improve retrieval performance. This action is crucial for enhancing the relevance of the retrieved context and generating more accurate responses.\n\n\n```python\n# Function to rewrite the query for better document retrieval\ndef rewrite_query(\n    query: str, \n    context_chunks: List[str], \n    model: str = \"google/gemma-2-2b-it\", \n    max_tokens: int = 100, \n    temperature: float = 0.3\n) -\u003e str:\n    \"\"\"\n    Use the LLM to rewrite the query for better document retrieval.\n\n    Args:\n        query (str): The original query text.\n        context_chunks (List[str]): A list of context chunks retrieved so far.\n        model (str): The model to use for generating the rewritten query. Default is \"google/gemma-2-2b-it\".\n        max_tokens (int): Maximum number of tokens in the rewritten query. Default is 100.\n        temperature (float): Sampling temperature for response diversity. Default is 0.3.\n\n    Returns:\n        str: The rewritten query optimized for document retrieval.\n    \"\"\"\n    # Construct a prompt for the LLM to rewrite the query\n    rewrite_prompt = f\"\"\"\n    You are a query optimization assistant. Your task is to rewrite the given query to make it more effective \n    for retrieving relevant information. The query will be used for document retrieval.\n    \n    Original query: {query}\n    \n    Based on the context retrieved so far:\n    {' '.join(context_chunks[:2]) if context_chunks else 'No context available yet'}\n    \n    Rewrite the query to be more specific and targeted to retrieve better information.\n    Rewritten query:\n    \"\"\"\n    \n    # Use the LLM to generate a rewritten query\n    response = client.chat.completions.create(\n        model=model, # Specify the model to use for generating the response\n        max_tokens=max_tokens, # Maximum number of tokens in the response\n        temperature=temperature, # Sampling temperature for response diversity\n        messages=[\n            {\n                \"role\": \"user\",\n                \"content\": rewrite_prompt\n            }\n        ]\n    )\n    \n    # Extract and return the rewritten query from the response\n    rewritten_query = response.choices[0].message.content.strip()\n    return rewritten_query\n```\n\nThis action is crucial for enhancing the relevance of the retrieved context and generating more accurate responses.\n\nLet's code our next action logic, which is to expand the context by retrieving additional chunks. We will use the existing function `retrieve_relevant_chunks` to get more context chunks and then filter out any duplicates from the current context. We will limit the number of new chunks to be added to the context to a specified top_k value.\n\n\n```python\n# Function to expand the context by retrieving additional chunks\ndef expand_context(query: str, current_chunks: List[str], top_k: int = 3) -\u003e List[str]:\n    \"\"\"\n    Expand the context by retrieving additional chunks.\n\n    Args:\n        query (str): The query text for which additional context is needed.\n        current_chunks (List[str]): The current list of context chunks.\n        top_k (int): The number of additional chunks to retrieve. Default is 3.\n\n    Returns:\n        List[str]: The expanded list of context chunks including new unique chunks.\n    \"\"\"\n    # Retrieve more chunks than currently available\n    additional_chunks = retrieve_relevant_chunks(query, top_k=top_k + len(current_chunks))\n    \n    # Filter out chunks that are already in the current context\n    new_chunks = []\n    for chunk in additional_chunks:\n        if chunk not in current_chunks:\n            new_chunks.append(chunk)\n    \n    # Add new unique chunks to the current context, limited to top_k\n    expanded_context = current_chunks + new_chunks[:top_k]\n    return expanded_context\n```\n\nWe need to filter the context to keep only the most relevant chunks for the query. This filtering step is crucial to ensure that the context provided to the language model is concise and focused on the most relevant information.\n\n\n```python\n# Function to filter the context to keep only the most relevant chunks\ndef filter_context(query: str, context_chunks: List[str]) -\u003e List[str]:\n    \"\"\"\n    Filter the context to keep only the most relevant chunks.\n\n    Args:\n        query (str): The query text for which relevance is calculated.\n        context_chunks (List[str]): The list of context chunks to filter.\n\n    Returns:\n        List[str]: A filtered list of the most relevant context chunks.\n    \"\"\"\n    if not context_chunks:\n        return []\n        \n    # Generate embeddings for the query and each chunk\n    query_embedding = generate_embeddings([query])[0]\n    chunk_embeddings = [generate_embeddings([chunk])[0] for chunk in context_chunks]\n    \n    # Calculate relevance scores for each chunk\n    relevance_scores = []\n    for chunk_embedding in chunk_embeddings:\n        score = cosine_similarity(query_embedding, chunk_embedding)\n        relevance_scores.append(score)\n    \n    # Sort chunks by relevance scores in descending order\n    sorted_chunks = [x for _, x in sorted(zip(relevance_scores, context_chunks), reverse=True)]\n    \n    # Keep the top 5 most relevant chunks or fewer if less than 5 are available\n    filtered_chunks = sorted_chunks[:min(5, len(sorted_chunks))]\n    \n    return filtered_chunks\n```\n\nThis action will help the agent explore more information relevant to the query.\n\n## Policy Network\n\nPreviously, we defined our state, actions, and reward logic. Next, we need to create a policy network that will select an action based on the current state.\n\nA policy network is a function that takes the current state and the action space as input and returns the selected action based on the state.\n\nThe policy network can use a simple heuristic to select an action based on the current state. For example, if there are no previous responses, the policy network can prioritize rewriting the query. If the context has too many chunks, the policy network can choose to filter the context.\n\n\n```python\n# Function to define a policy network to select an action based on the state\ndef policy_network(\n    state: dict, \n    action_space: List[str], \n    epsilon: float = 0.2\n) -\u003e str:\n    \"\"\"\n    Define a policy network to select an action based on the current state using an epsilon-greedy strategy.\n\n    Args:\n        state (dict): The current state of the environment, including query, context, responses, and rewards.\n        action_space (List[str]): The list of possible actions the agent can take.\n        epsilon (float): The probability of choosing a random action for exploration. Default is 0.2.\n\n    Returns:\n        str: The selected action from the action space.\n    \"\"\"\n    # Use epsilon-greedy strategy: random exploration vs. exploitation\n    if np.random.random() \u003c epsilon:\n        # Exploration: randomly select an action from the action space\n        action = np.random.choice(action_space)\n    else:\n        # Exploitation: select the best action based on the current state using a simple heuristic\n\n        # If there are no previous responses, prioritize rewriting the query\n        if len(state[\"previous_responses\"]) == 0:\n            action = \"rewrite_query\"\n        # If there are previous responses but the rewards are low, try expanding the context\n        elif state[\"previous_rewards\"] and max(state[\"previous_rewards\"]) \u003c 0.7:\n            action = \"expand_context\"\n        # If the context has too many chunks, try filtering the context\n        elif len(state[\"context\"]) \u003e 5:\n            action = \"filter_context\"\n        # Otherwise, generate a response\n        else:\n            action = \"generate_response\"\n    \n    return action\n```\n\nSo our policy network works like this:\n- If there are no previous responses, prioritize rewriting the query.\n- If there are previous responses but the rewards are low, try expanding the context.\n- If the context has too many chunks, try filtering the context.\n- Otherwise, generate a response.\n\n## Single RL Step\n\nWe have coded an important component of the RL pipeline. For any developer who has done any kind of training, there exists a training loop where each iteration is a single step in which the RL agent takes an action, rewards are calculated, states are updated, and so on. So, we need to code a single step of our training loop. Let's do that.\n\n\n```python\n# Function to perform a single RL step\ndef rl_step(\n    state: dict, \n    action_space: List[str], \n    ground_truth: str\n) -\u003e tuple[dict, str, float, str]:\n    \"\"\"\n    Perform a single RL step: select an action, execute it, and calculate the reward.\n\n    Args:\n        state (dict): The current state of the environment, including query, context, responses, and rewards.\n        action_space (List[str]): The list of possible actions the agent can take.\n        ground_truth (str): The expected correct answer to calculate the reward.\n\n    Returns:\n        tuple: A tuple containing:\n            - state (dict): The updated state after executing the action.\n            - action (str): The action selected by the policy network.\n            - reward (float): The reward received for the action.\n            - response (str): The response generated (if applicable).\n    \"\"\"\n    # Select an action using the policy network\n    action: str = policy_network(state, action_space)\n    response: str = None  # Initialize response as None\n    reward: float = 0  # Initialize reward as 0\n\n    # Execute the selected action\n    if action == \"rewrite_query\":\n        # Rewrite the query to improve retrieval\n        rewritten_query: str = rewrite_query(state[\"original_query\"], state[\"context\"])\n        state[\"current_query\"] = rewritten_query  # Update the current query in the state\n        # Retrieve new context based on the rewritten query\n        new_context: List[str] = retrieve_relevant_chunks(rewritten_query)\n        state[\"context\"] = new_context  # Update the context in the state\n\n    elif action == \"expand_context\":\n        # Expand the context by retrieving additional chunks\n        expanded_context: List[str] = expand_context(state[\"current_query\"], state[\"context\"])\n        state[\"context\"] = expanded_context  # Update the context in the state\n\n    elif action == \"filter_context\":\n        # Filter the context to keep only the most relevant chunks\n        filtered_context: List[str] = filter_context(state[\"current_query\"], state[\"context\"])\n        state[\"context\"] = filtered_context  # Update the context in the state\n\n    elif action == \"generate_response\":\n        # Construct a prompt using the current query and context\n        prompt: str = construct_prompt(state[\"current_query\"], state[\"context\"])\n        # Generate a response using the LLM\n        response: str = generate_response(prompt)\n        # Calculate the reward based on the similarity between the response and the ground truth\n        reward: float = calculate_reward(response, ground_truth)\n        # Update the state with the new response and reward\n        state[\"previous_responses\"].append(response)\n        state[\"previous_rewards\"].append(reward)\n\n    # Return the updated state, selected action, reward, and response\n    return state, action, reward, response\n```\n\nIn our single step function, we first select an action using the policy network. The policy network uses an epsilon-greedy strategy to balance exploration and exploitation. If the random number is less than epsilon, we choose a random action from the action space for exploration. Otherwise, we select the best action based on the current state using a simple heuristic.\n\n## Training Parameters and Policy Update\n\nWe need to define some training parameters for our training loop and also define a function to update the policy based on the rewards received.\n\nThough the training parameters function is **optional**, it can be used for advanced implementations of the RL pipeline.\n\n\n```python\n# Function to initialize training parameters\ndef initialize_training_params() -\u003e Dict[str, Union[float, int]]:\n    \"\"\"\n    Initialize training parameters such as learning rate, number of episodes, and discount factor.\n\n    Returns:\n        Dict[str, Union[float, int]]: A dictionary containing the initialized training parameters.\n    \"\"\"\n    params = {\n        \"learning_rate\": 0.01,  # Learning rate for policy updates\n        \"num_episodes\": 100,   # Total number of training episodes\n        \"discount_factor\": 0.99  # Discount factor for future rewards\n    }\n    return params\n```\n\nSimilar to how our state changes after each step in the RL process, the policy also needs to be updated based on the rewards received. The update_policy function takes the current policy, state, action, reward, and learning rate as input and returns the updated policy.\n\n\n```python\n# Function to update policy based on reward\ndef update_policy(\n    policy: Dict[str, Dict[str, Union[float, str]]], \n    state: Dict[str, object], \n    action: str, \n    reward: float, \n    learning_rate: float\n) -\u003e Dict[str, Dict[str, Union[float, str]]]:\n    \"\"\"\n    Update the policy based on the reward received.\n\n    Args:\n        policy (Dict[str, Dict[str, Union[float, str]]]): The current policy to be updated.\n        state (Dict[str, object]): The current state of the environment.\n        action (str): The action taken by the agent.\n        reward (float): The reward received for the action.\n        learning_rate (float): The learning rate for updating the policy.\n\n    Returns:\n        Dict[str, Dict[str, Union[float, str]]]: The updated policy.\n    \"\"\"\n    # Example: Simple policy update (to be replaced with a proper RL algorithm)\n    policy[state[\"query\"]] = {\n        \"action\": action,  # Store the action taken\n        \"reward\": reward   # Store the reward received\n    }\n    return policy\n```\n\nIn the above `update_policy` logic, we store the action taken and the reward received for each query in the policy dictionary. In a more advanced RL algorithm, the policy update would involve more sophisticated methods such as policy gradients or Q-learning.\n\nFinally, we need to implement progress tracking logic to monitor the training process. This will help us understand how the model is learning and improving over time.\n\n\n```python\n# Function to track training progress\ndef track_progress(\n    episode: int, \n    reward: float, \n    rewards_history: List[float]\n) -\u003e List[float]:\n    \"\"\"\n    Track the training progress by storing rewards for each episode.\n\n    Args:\n        episode (int): The current episode number.\n        reward (float): The reward received in the current episode.\n        rewards_history (List[float]): A list to store the rewards for all episodes.\n\n    Returns:\n        List[float]: The updated rewards history.\n    \"\"\"\n    # Append the current reward to the rewards history\n    rewards_history.append(reward)\n    \n    # Print progress every 10 episodes\n    print(f\"Episode {episode}: Reward = {reward}\")\n    \n    return rewards_history\n```\n\n## Training Loop\n\nNow that we have coded every part of the training loop, we can put it all together in a single function that implements the training loop for the RL-enhanced RAG system.\n\n\n```python\n# Function to implement the training loop\ndef training_loop(\n    query_text: str, \n    ground_truth: str, \n    params: Optional[Dict[str, Union[float, int]]] = None\n) -\u003e Tuple[Dict[str, Dict[str, Union[float, str]]], List[float], List[List[str]], Optional[str]]:\n    \"\"\"\n    Implement the training loop for RL-enhanced RAG.\n\n    Args:\n        query_text (str): The input query text for the RAG pipeline.\n        ground_truth (str): The expected correct answer for the query.\n        params (Optional[Dict[str, Union[float, int]]]): Training parameters such as learning rate, \n            number of episodes, and discount factor. If None, default parameters are initialized.\n\n    Returns:\n        Tuple: A tuple containing:\n            - policy (Dict[str, Dict[str, Union[float, str]]]): The updated policy after training.\n            - rewards_history (List[float]): A list of rewards received in each episode.\n            - actions_history (List[List[str]]): A list of actions taken in each episode.\n            - best_response (Optional[str]): The best response generated during training.\n    \"\"\"\n    # Initialize training parameters if not provided\n    if params is None:\n        params = initialize_training_params()\n    \n    # Initialize variables to track progress\n    rewards_history: List[float] = []  # List to store rewards for each episode\n    actions_history: List[List[str]] = []  # List to store actions taken in each episode\n    policy: Dict[str, Dict[str, Union[float, str]]] = {}  # Policy dictionary to store actions and rewards\n    action_space: List[str] = define_action_space()  # Define the action space\n    best_response: Optional[str] = None  # Variable to store the best response\n    best_reward: float = -1  # Initialize the best reward to a very low value\n    \n    # Get initial performance from the simple RAG pipeline for comparison\n    simple_response: str = basic_rag_pipeline(query_text)\n    simple_reward: float = calculate_reward(simple_response, ground_truth)\n    print(f\"Simple RAG reward: {simple_reward:.4f}\")\n\n    # Start the training loop\n    for episode in range(params[\"num_episodes\"]):\n        # Reset the environment with the same query\n        context_chunks: List[str] = retrieve_relevant_chunks(query_text)\n        state: Dict[str, object] = define_state(query_text, context_chunks)\n        episode_reward: float = 0  # Initialize the reward for the current episode\n        episode_actions: List[str] = []  # Initialize the list of actions for the current episode\n        \n        # Maximum number of steps per episode to prevent infinite loops\n        for step in range(10):\n            # Perform a single RL step\n            state, action, reward, response = rl_step(state, action_space, ground_truth)\n            episode_actions.append(action)  # Record the action taken\n            \n            # If a response is generated, end the episode\n            if response:\n                episode_reward = reward  # Update the episode reward\n                \n                # Track the best response and reward\n                if reward \u003e best_reward:\n                    best_reward = reward\n                    best_response = response\n                \n                break  # Exit the loop as the episode ends\n        \n        # Update rewards and actions history\n        rewards_history.append(episode_reward)\n        actions_history.append(episode_actions)\n        \n        # Print progress every 5 episodes\n        if episode % 5 == 0:\n            print(f\"Episode {episode}: Reward = {episode_reward:.4f}, Actions = {episode_actions}\")\n    \n    # Compare the best RL-enhanced RAG reward with the simple RAG reward\n    improvement: float = best_reward - simple_reward\n    print(f\"\\nTraining completed:\")\n    print(f\"Simple RAG reward: {simple_reward:.4f}\")\n    print(f\"Best RL-enhanced RAG reward: {best_reward:.4f}\")\n    print(f\"Improvement: {improvement:.4f} ({improvement * 100:.2f}%)\")\n\n    return policy, rewards_history, actions_history, best_response\n```\n\nThis function will take the input query text, the expected ground truth answer, and optionally some training parameters. It will return the updated policy, a list of rewards received in each episode, a list of actions taken in each episode, and the best response generated during training.\n\nIn more detail, the `training_loop` function will:\n- Initialize training parameters if not provided.\n- Get the initial performance from the simple RAG pipeline for comparison.\n- Start the training loop for the specified number of episodes.\n- Perform a single RL step in each episode.\n- Update rewards and actions history for each episode.\n- Print progress every 5 episodes.\n- Compare the best RL-enhanced RAG reward with the simple RAG reward.\n- Return the updated policy, rewards history, actions history, and the best response generated during training.\n\n## Performance Comparison Logic\n\nAlthough we can manually compare the simple RAG pipeline with the RL-based RAG pipeline, a function can definitely help us in this regard. So, let's define a function to compare the performance of the simple RAG pipeline with the RL-enhanced RAG pipeline.\n\n\n```python\n# Function to compare Simple RAG vs RL-Enhanced RAG\ndef compare_rag_approaches(query_text: str, ground_truth: str) -\u003e Tuple[str, str, float, float]:\n    \"\"\"\n    Compare the outputs of simple RAG versus RL-enhanced RAG.\n\n    Args:\n        query_text (str): The input query text for the RAG pipeline.\n        ground_truth (str): The expected correct answer for the query.\n\n    Returns:\n        Tuple[str, str, float, float]: A tuple containing:\n            - simple_response (str): The response generated by the simple RAG pipeline.\n            - best_rl_response (str): The best response generated by the RL-enhanced RAG pipeline.\n            - simple_similarity (float): The similarity score of the simple RAG response to the ground truth.\n            - rl_similarity (float): The similarity score of the RL-enhanced RAG response to the ground truth.\n    \"\"\"\n    print(\"=\" * 80)\n    print(f\"Query: {query_text}\")\n    print(\"=\" * 80)\n    \n    # Step 1: Generate a response using the simple RAG pipeline\n    # The basic RAG pipeline retrieves relevant chunks and generates a response without reinforcement learning.\n    simple_response: str = basic_rag_pipeline(query_text)\n    # Calculate the similarity score between the simple RAG response and the ground truth.\n    simple_similarity: float = calculate_reward(simple_response, ground_truth)\n    \n    print(\"\\nSimple RAG Output:\")\n    print(\"-\" * 40)\n    print(simple_response)\n    print(f\"Similarity to ground truth: {simple_similarity:.4f}\")\n    \n    # Step 2: Train the RL-enhanced RAG model\n    print(\"\\nTraining RL-enhanced RAG model...\")\n    # Initialize training parameters (e.g., learning rate, number of episodes, discount factor).\n    params: Dict[str, float | int] = initialize_training_params()\n    # Set the number of episodes to a smaller value for demonstration purposes.\n    params[\"num_episodes\"] = 5\n    \n    # Run the training loop for the RL-enhanced RAG model.\n    # This loop trains the model to optimize its responses using reinforcement learning.\n    _, rewards_history, actions_history, best_rl_response = training_loop(\n        query_text, ground_truth, params\n    )\n    \n    # If no response was generated during training, generate one using the current query and context.\n    if best_rl_response is None:\n        # Retrieve relevant chunks for the query.\n        context_chunks: List[str] = retrieve_relevant_chunks(query_text)\n        # Construct a prompt using the query and retrieved context.\n        prompt: str = construct_prompt(query_text, context_chunks)\n        # Generate a response using the language model.\n        best_rl_response: str = generate_response(prompt)\n    \n    # Calculate the similarity score between the RL-enhanced RAG response and the ground truth.\n    rl_similarity: float = calculate_reward(best_rl_response, ground_truth)\n    \n    print(\"\\nRL-enhanced RAG Output:\")\n    print(\"-\" * 40)\n    print(best_rl_response)\n    print(f\"Similarity to ground truth: {rl_similarity:.4f}\")\n    \n    # Step 3: Evaluate and compare the results\n    # Calculate the improvement in similarity score achieved by the RL-enhanced RAG model.\n    improvement: float = rl_similarity - simple_similarity\n    \n    print(\"\\nEvaluation Results:\")\n    print(\"-\" * 40)\n    print(f\"Simple RAG similarity to ground truth: {simple_similarity:.4f}\")\n    print(f\"RL-enhanced RAG similarity to ground truth: {rl_similarity:.4f}\")\n    print(f\"Improvement: {improvement * 100:.2f}%\")\n    \n    # Step 4: Plot the reward history (if there are enough episodes and matplotlib is available)\n    if len(rewards_history) \u003e 1:\n        try:\n            import matplotlib.pyplot as plt\n            # Create a plot to visualize the reward history during RL training.\n            plt.figure(figsize=(10, 6))\n            plt.plot(rewards_history)\n            plt.title('Reward History During RL Training')\n            plt.xlabel('Episode')\n            plt.ylabel('Reward')\n            plt.grid(True)\n            plt.show()\n        except ImportError:\n            # If matplotlib is not available, print a message instead of plotting.\n            print(\"Matplotlib not available for plotting rewards\")\n    \n    # Return the results: responses and similarity scores for both approaches.\n    return simple_response, best_rl_response, simple_similarity, rl_similarity\n```\n\nSo our performance comparison logic is not very complicated but is based on 4 steps:\n1. Generate a response using the simple RAG pipeline.\n2. Train the RL-enhanced RAG model using the training loop.\n3. Evaluate and compare the results.\n4. Plot the reward history (if available).\n\n## Evaluation Framework (**Optional**)\n\nThis step is optional but in case you want to evaluate all the eval queries in the validation data, you can use the following code.\n\nFirst, to check the relevance of the retrieved chunks and the ground truth, we need to have a function that evaluates the relevance of the retrieved chunks.\n\n\n```python\n# Function to evaluate relevance of retrieved chunks\ndef evaluate_relevance(retrieved_chunks: List[str], ground_truth_chunks: List[str]) -\u003e float:\n    \"\"\"\n    Evaluate the relevance of retrieved chunks by comparing them to ground truth chunks.\n\n    Args:\n        retrieved_chunks (List[str]): A list of text chunks retrieved by the system.\n        ground_truth_chunks (List[str]): A list of ground truth text chunks for comparison.\n\n    Returns:\n        float: The average relevance score between the retrieved chunks and the ground truth chunks.\n    \"\"\"\n    relevance_scores: List[float] = []  # Initialize a list to store relevance scores\n\n    # Iterate through pairs of retrieved and ground truth chunks\n    for retrieved, ground_truth in zip(retrieved_chunks, ground_truth_chunks):\n        # Calculate the cosine similarity between the embeddings of the retrieved and ground truth chunks\n        relevance: float = cosine_similarity(\n            generate_embeddings([retrieved])[0],\n            generate_embeddings([ground_truth])[0]\n        )\n        # Append the relevance score to the list\n        relevance_scores.append(relevance)\n\n    # Return the average relevance score\n    return np.mean(relevance_scores)\n```\n\nTo evaluate the accuracy of the generated responses, we can use the cosine similarity between the embeddings of the generated responses and the ground truth. So let's define a function to evaluate the accuracy of the responses based on this similarity metric.\n\n\n```python\n# Function to evaluate the accuracy of generated responses\ndef evaluate_accuracy(responses: List[str], ground_truth_responses: List[str]) -\u003e float:\n    \"\"\"\n    Evaluate the accuracy of generated responses by comparing them to ground truth responses.\n\n    Args:\n        responses (List[str]): A list of generated responses to evaluate.\n        ground_truth_responses (List[str]): A list of ground truth responses to compare against.\n\n    Returns:\n        float: The average accuracy score, calculated as the mean cosine similarity \n               between the embeddings of the generated responses and the ground truth responses.\n    \"\"\"\n    accuracy_scores: List[float] = []  # Initialize a list to store accuracy scores\n\n    # Iterate through each pair of generated response and ground truth response\n    for response, ground_truth in zip(responses, ground_truth_responses):\n        # Calculate the cosine similarity between the embeddings of the response and ground truth\n        accuracy: float = cosine_similarity(\n            generate_embeddings([response])[0],\n            generate_embeddings([ground_truth])[0]\n        )\n        # Append the accuracy score to the list\n        accuracy_scores.append(accuracy)\n\n    # Return the mean of the accuracy scores\n    return np.mean(accuracy_scores)\n```\n\nWe also need to measure the response quality and assign a relevant score for it to be used in the reinforcement learning process.\n\n\n```python\n# Function to evaluate response quality\ndef evaluate_response_quality(responses: List[str]) -\u003e float:\n    \"\"\"\n    Evaluate the quality of responses using a heuristic or external model.\n\n    Args:\n        responses (List[str]): A list of generated responses to evaluate.\n\n    Returns:\n        float: The average quality score of the responses, ranging from 0 to 1.\n    \"\"\"\n    quality_scores: List[float] = []  # Initialize a list to store quality scores for each response\n\n    for response in responses:\n        # Example heuristic: Calculate a quality score based on response length\n        # Normalize the length by a maximum of 100 words and cap the score at 1.0\n        quality: float = len(response.split()) / 100\n        quality_scores.append(min(quality, 1.0))  # Append the capped quality score to the list\n\n    # Return the average quality score across all responses\n    return np.mean(quality_scores)\n```\n\nThen we can evaluate the performance of the RL-enhanced RAG model on the validation dataset:\n\n\n```python\n# Function to evaluate RAG performance\ndef evaluate_rag_performance(\n    queries: List[str], \n    ground_truth_chunks: List[str], \n    ground_truth_responses: List[str]\n) -\u003e Dict[str, float]:\n    \"\"\"\n    Evaluate the performance of the RAG pipeline using relevance, accuracy, and response quality metrics.\n\n    Args:\n        queries (List[str]): A list of query strings to evaluate.\n        ground_truth_chunks (List[str]): A list of ground truth text chunks corresponding to the queries.\n        ground_truth_responses (List[str]): A list of ground truth responses corresponding to the queries.\n\n    Returns:\n        Dict[str, float]: A dictionary containing the average relevance, accuracy, and quality scores.\n    \"\"\"\n    # Initialize lists to store scores for each metric\n    relevance_scores: List[float] = []\n    accuracy_scores: List[float] = []\n    quality_scores: List[float] = []\n\n    # Iterate through each query and its corresponding ground truth data\n    for query, ground_truth_chunk, ground_truth_response in zip(queries, ground_truth_chunks, ground_truth_responses):\n        # Retrieve relevant chunks for the query\n        retrieved_chunks: List[str] = retrieve_relevant_chunks(query)\n        \n        # Evaluate the relevance of the retrieved chunks compared to the ground truth chunk\n        relevance: float = evaluate_relevance(retrieved_chunks, [ground_truth_chunk])\n        relevance_scores.append(relevance)\n\n        # Generate a response using the basic RAG pipeline\n        response: str = basic_rag_pipeline(query)\n        \n        # Evaluate the accuracy of the generated response compared to the ground truth response\n        accuracy: float = evaluate_accuracy([response], [ground_truth_response])\n        accuracy_scores.append(accuracy)\n\n        # Evaluate the quality of the generated response\n        quality: float = evaluate_response_quality([response])\n        quality_scores.append(quality)\n\n    # Calculate the average scores for each metric\n    avg_relevance: float = np.mean(relevance_scores)\n    avg_accuracy: float = np.mean(accuracy_scores)\n    avg_quality: float = np.mean(quality_scores)\n\n    # Return the average scores as a dictionary\n    return {\n        \"average_relevance\": avg_relevance,\n        \"average_accuracy\": avg_accuracy,\n        \"average_quality\": avg_quality\n    }\n```\n\n## Evaluating (RL vs Simple) RAG\n\nAh, the moment of truth! Let's evaluate the performance of the simple RAG pipeline against the RL-enhanced RAG pipeline on our factual query, where the simple RAG previously failed to provide the correct answer. Let's see if the RL-enhanced RAG pipeline can perform better.\n\nLet's revisit our evaluation query and see what the simple RAG pipeline generates for it.\n\n\n```python\n# Print a message to indicate the start of the RAG pipeline\nprint(\"🔍 Running the Retrieval-Augmented Generation (RAG) pipeline...\")\nprint(f\"📥 Query: {sample_query}\\n\")\n\n# Run the RAG pipeline and get the response\nresponse = basic_rag_pipeline(sample_query)\n\n# Print the response with better formatting\nprint(\"🤖 AI Response:\")\nprint(\"-\" * 50)\nprint(response.strip())\nprint(\"-\" * 50)\n\n# Print the ground truth answer for comparison\nprint(\"✅ Ground Truth Answer:\")\nprint(\"-\" * 50)\nprint(expected_answer)\nprint(\"-\" * 50)\n```\n\n    🔍 Running the Retrieval-Augmented Generation (RAG) pipeline...\n    📥 Query: What is the mathematical representation of a qubit in superposition?\n    \n    🤖 AI Response:\n    --------------------------------------------------\n    ψ  α0  β1\n    --------------------------------------------------\n    ✅ Ground Truth Answer:\n    --------------------------------------------------\n    |ψ⟩ = α|0⟩ + β|1⟩, where α and β are complex numbers satisfying |α|² + |β|² = 1, representing the probability amplitudes for measuring the qubit in state |0⟩ or |1⟩ respectively.\n    --------------------------------------------------\n    \n\n\n```python\n# Compare the performance of the simple RAG pipeline and the RL-enhanced RAG pipeline\n# using the sample query and its expected answer.\n# The function returns:\n# - simple_response: The response generated by the simple RAG pipeline.\n# - rl_response: The best response generated by the RL-enhanced RAG pipeline.\n# - simple_sim: The similarity score of the simple RAG response to the ground truth.\n# - rl_sim: The similarity score of the RL-enhanced RAG response to the ground truth.\nsimple_response, rl_response, simple_sim, rl_sim = compare_rag_approaches(sample_query, expected_answer)\n```\n\n    ================================================================================\n    Query: What is the mathematical representation of a qubit in superposition?\n    ================================================================================\n    \n    Simple RAG Output:\n    ----------------------------------------\n    ψ  α0  β1 \n    \n    Similarity to ground truth: 0.6726\n    \n    Training RL-enhanced RAG model...\n    Simple RAG reward: 0.6772\n    Episode 0: Reward = 0.0000, Actions = ['rewrite_query', 'rewrite_query', np.str_('rewrite_query'), 'rewrite_query', np.str_('rewrite_query'), 'rewrite_query', 'rewrite_query', 'rewrite_query', np.str_('expand_context'), 'rewrite_query']\n    \n    Training completed:\n    Simple RAG reward: 0.6772\n    Best RL-enhanced RAG reward: 0.8652\n    Improvement: 0.1879 (18.79%)\n    \n    RL-enhanced RAG Output:\n    ----------------------------------------\n    The mathematical representation of a qubit in superposition is: \n    ψ = α0 + β1 \n    \n    Where:\n    \n    * α and β are complex numbers.\n    * α² + β² = 1  \n    \n    \n    Let me know if you would like a deeper explanation of any of these terms! \n    \n    Similarity to ground truth: 0.8652\n    \n    Evaluation Results:\n    ----------------------------------------\n    Simple RAG similarity to ground truth: 0.6726\n    RL-enhanced RAG similarity to ground truth: 0.8652\n    Improvement: 19.26%\n    \n![png](data/reward_plot.png)\n\nYou can clearly see that the response generated by the RL-enhanced RAG model is more accurate and relevant compared to the simple RAG pipeline. The improvement in similarity to the ground truth is evident, indicating that the RL-enhanced model has learned to generate better responses through training.\n\n## Saving the Comparison Results\n\nAfter implementing the RL algorithm, we can save the comparison results to check the performance of the RL implementation later.\n\n\n```python\n# Save the results for later comparison\nresults = {\n    \"query\": query_text,  # The input query text\n    \"ground_truth\": expected_answer,  # The expected correct answer for the query\n    \"simple_rag\": {\n        \"response\": simple_response,  # The response generated by the simple RAG pipeline\n        \"similarity\": float(simple_sim)  # The similarity score of the simple RAG response to the ground truth\n    },\n    \"rl_rag\": {\n        \"response\": rl_response,  # The response generated by the RL-enhanced RAG pipeline\n        \"similarity\": float(rl_sim)  # The similarity score of the RL-enhanced RAG response to the ground truth\n    },\n    \"improvement\": float(rl_sim - simple_sim)  # The improvement in similarity score achieved by RL-enhanced RAG\n}\n\n# Save the results to a JSON file for future reference\nwith open('rl_rag_results.json', 'w') as f:\n    json.dump(results, f, indent=2)  # Write the results dictionary to the file with indentation for readability\n\n# Print a confirmation message to indicate that the results have been saved\nprint(\"\\nResults saved to rl_rag_results.json\")\n```\n\n    \n    Results saved to rl_rag_results.json\n    \n\n## What can we conclude?\n\n- The performance of the simple RAG is lower compared to the RL-enhanced RAG on factual queries.\n- The RL-enhanced RAG achieved a 19.5% improvement in the similarity score within 5 episodes.\n- Further improvements can be achieved by:\n    - Training for more episodes.\n    - Tuning hyperparameters.\n- Time is a key constraint for training.\n- Parallel implementation of the RL algorithm can help reduce training time.\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Ffareedkhan-dev%2Frag-with-rl","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Ffareedkhan-dev%2Frag-with-rl","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Ffareedkhan-dev%2Frag-with-rl/lists"}