Initial app setup

.dockerignore

@@ -0,0 +1,46 @@
+# Git
+.git
+.gitignore
+
+# Python
+__pycache__/
+*.py[cod]
+*$py.class
+*.so
+.Python
+env/
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+
+# Virtual Environment
+venv/
+ENV/
+.env
+
+# IDE specific files
+.idea/
+.vscode/
+*.swp
+*.swo
+
+# OS specific
+.DS_Store
+
+# Notebook version (keep in repo but not in Docker)
+notebook_version/
+
+# Ignore PDF data files (will be mounted at runtime)
+data/*.pdf 

.env.example

@@ -0,0 +1,10 @@
+# OpenAI API key
+OPENAI_API_KEY=your_openai_api_key
+
+# Optional backup key if primary is not set
+OPENAI_API_KEY_BACKUP=
+
+# Chainlit configuration
+CHAINLIT_MAX_STEPS_HISTORY=10
+CHAINLIT_AUTH_SECRET=
+CHAINLIT_SERVER_URL= 

.gitignore

@@ -0,0 +1,48 @@
+# Python
+__pycache__/
+*.py[cod]
+*$py.class
+*.so
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+
+# Environment and secrets
+.env
+.venv
+venv/
+ENV/
+env/
+.env.local
+.env.development.local
+.env.test.local
+.env.production.local
+
+# Chainlit
+.chainlit/
+.chainlit/cache/
+
+# IDE specific files
+.idea/
+.vscode/
+*.swp
+*.swo
+
+# Mac OS
+.DS_Store
+
+# Local data
+/data/*.pdf 

Dockerfile

@@ -0,0 +1,31 @@
+FROM python:3.11-slim
+
+WORKDIR /app
+
+# Install system dependencies
+RUN apt-get update && apt-get install -y \
+    build-essential \
+    python3-dev \
+    && rm -rf /var/lib/apt/lists/*
+
+# Copy requirements file
+COPY requirements.txt .
+
+# Install Python dependencies
+RUN pip install --no-cache-dir -r requirements.txt
+
+# Copy application code
+COPY . .
+
+# Create data directory for PDFs
+RUN mkdir -p data
+
+# Expose port for Chainlit
+EXPOSE 8000
+
+# Set environment variables
+ENV PYTHONPATH=/app
+ENV OPENAI_API_KEY=${OPENAI_API_KEY}
+
+# Command to run the application
+CMD ["chainlit", "run", "app.py", "--port", "8000", "--host", "0.0.0.0"] 

Makefile

@@ -0,0 +1,23 @@
+.PHONY: setup run docker-build docker-run clean
+
+setup:
+	python -m pip install -r requirements.txt
+
+run:
+	chainlit run app.py
+
+docker-build:
+	docker build -t ab-testing-rag-agent .
+
+docker-run:
+	docker run -p 8000:8000 --env-file .env -v $(PWD)/data:/app/data ab-testing-rag-agent
+
+docker-compose-up:
+	docker-compose up -d
+
+docker-compose-down:
+	docker-compose down
+
+clean:
+	rm -rf __pycache__
+	rm -rf .chainlit 

README.md

@@ -1,12 +1,119 @@
----
-title: AB Testing RAG Agent
-emoji: 🚀
-colorFrom: purple
-colorTo: gray
-sdk: docker
-pinned: false
-license: mit
-short_description: AB Testing RAG app using Ron Kohavi's work
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# AB Testing RAG Agent
+
+An application that helps answer questions about AB Testing using a collection of PDF documents and a RAG (Retrieval Augmented Generation) approach. This agent was built using LangChain, OpenAI, Qdrant, and Chainlit.
+
+## Features
+
+- **Document Retrieval**: Retrieves relevant information from AB Testing documents
+- **Query Rephrasing**: Improves retrieval by rephrasing your query for better results
+- **Source References**: Shows the exact document sources used to answer your question
+- **Streaming Interface**: See the response as it's being generated
+
+## Prerequisites
+
+- Python 3.9+
+- Docker and Docker Compose (optional)
+- OpenAI API key
+
+## Setup
+
+### Local Development
+
+1. Clone the repository
+   ```bash
+   git clone https://github.com/yourusername/AB_Testing_RAG_Agent.git
+   cd AB_Testing_RAG_Agent
+   ```
+
+2. Create a virtual environment
+   ```bash
+   python -m venv venv
+   source venv/bin/activate  # On Windows: venv\Scripts\activate
+   ```
+
+3. Install dependencies
+   ```bash
+   pip install -r requirements.txt
+   ```
+
+4. Create a `.env` file from the example
+   ```bash
+   cp .env.example .env
+   ```
+
+5. Add your OpenAI API key to the `.env` file
+   ```bash
+   OPENAI_API_KEY=your_openai_api_key
+   ```
+
+6. Add PDF files to the `data` directory
+   ```bash
+   mkdir -p data
+   # Copy your AB Testing PDFs to the data directory
+   ```
+
+7. Run the application
+   ```bash
+   chainlit run app.py
+   ```
+
+8. Open your browser to `http://localhost:8000`
+
+### Using Docker
+
+1. Clone the repository
+   ```bash
+   git clone https://github.com/yourusername/AB_Testing_RAG_Agent.git
+   cd AB_Testing_RAG_Agent
+   ```
+
+2. Create a `.env` file from the example
+   ```bash
+   cp .env.example .env
+   ```
+
+3. Add your OpenAI API key to the `.env` file
+   ```bash
+   OPENAI_API_KEY=your_openai_api_key
+   ```
+
+4. Add PDF files to the `data` directory
+   ```bash
+   mkdir -p data
+   # Copy your AB Testing PDFs to the data directory
+   ```
+
+5. Build and run with Docker Compose
+   ```bash
+   docker-compose up -d
+   ```
+
+6. Open your browser to `http://localhost:8000`
+
+## Deploying to Hugging Face Spaces
+
+1. Create a new Hugging Face Space with Docker deployment
+2. Set the following environment variables in your Space settings:
+   - `OPENAI_API_KEY`: Your OpenAI API key
+3. Push your code to the Space's repository
+
+## How It Works
+
+1. The application loads and processes PDF documents using LangChain's document loaders
+2. Documents are chunked and embedded using OpenAI's embedding model
+3. Qdrant vector store is used for semantic search
+4. When a user asks a question:
+   - The query is rephrased to be more specific for better retrieval
+   - Relevant document chunks are retrieved
+   - OpenAI's GPT model generates an answer based on the retrieved context
+   - Source references are tracked and displayed alongside the answer
+
+## License
+
+MIT
+
+## Acknowledgements
+
+- Based on the notebook version in `notebook_version/AB_Testing_RAG_Agent.ipynb`
+- Thanks to the creators of LangChain, Chainlit, and Qdrant for their excellent tools
+- Special thanks to all the AB Testing papers and documents used as the knowledge base

aimakerspace/openai_utils/chatmodel.py

@@ -1,45 +0,0 @@
-from openai import OpenAI, AsyncOpenAI
-from dotenv import load_dotenv
-import os
-
-load_dotenv()
-
-
-class ChatOpenAI:
-    def __init__(self, model_name: str = "gpt-4o-mini"):
-        self.model_name = model_name
-        self.openai_api_key = os.getenv("OPENAI_API_KEY")
-        if self.openai_api_key is None:
-            raise ValueError("OPENAI_API_KEY is not set")
-
-    def run(self, messages, text_only: bool = True, **kwargs):
-        if not isinstance(messages, list):
-            raise ValueError("messages must be a list")
-
-        client = OpenAI()
-        response = client.chat.completions.create(
-            model=self.model_name, messages=messages, **kwargs
-        )
-
-        if text_only:
-            return response.choices[0].message.content
-
-        return response
-    
-    async def astream(self, messages, **kwargs):
-        if not isinstance(messages, list):
-            raise ValueError("messages must be a list")
-        
-        client = AsyncOpenAI()
-
-        stream = await client.chat.completions.create(
-            model=self.model_name,
-            messages=messages,
-            stream=True,
-            **kwargs
-        )
-
-        async for chunk in stream:
-            content = chunk.choices[0].delta.content
-            if content is not None:
-                yield content

aimakerspace/openai_utils/embedding.py

@@ -1,59 +0,0 @@
-from dotenv import load_dotenv
-from openai import AsyncOpenAI, OpenAI
-import openai
-from typing import List
-import os
-import asyncio
-
-
-class EmbeddingModel:
-    def __init__(self, embeddings_model_name: str = "text-embedding-3-small"):
-        load_dotenv()
-        self.openai_api_key = os.getenv("OPENAI_API_KEY")
-        self.async_client = AsyncOpenAI()
-        self.client = OpenAI()
-
-        if self.openai_api_key is None:
-            raise ValueError(
-                "OPENAI_API_KEY environment variable is not set. Please set it to your OpenAI API key."
-            )
-        openai.api_key = self.openai_api_key
-        self.embeddings_model_name = embeddings_model_name
-
-    async def async_get_embeddings(self, list_of_text: List[str]) -> List[List[float]]:
-        embedding_response = await self.async_client.embeddings.create(
-            input=list_of_text, model=self.embeddings_model_name
-        )
-
-        return [embeddings.embedding for embeddings in embedding_response.data]
-
-    async def async_get_embedding(self, text: str) -> List[float]:
-        embedding = await self.async_client.embeddings.create(
-            input=text, model=self.embeddings_model_name
-        )
-
-        return embedding.data[0].embedding
-
-    def get_embeddings(self, list_of_text: List[str]) -> List[List[float]]:
-        embedding_response = self.client.embeddings.create(
-            input=list_of_text, model=self.embeddings_model_name
-        )
-
-        return [embeddings.embedding for embeddings in embedding_response.data]
-
-    def get_embedding(self, text: str) -> List[float]:
-        embedding = self.client.embeddings.create(
-            input=text, model=self.embeddings_model_name
-        )
-
-        return embedding.data[0].embedding
-
-
-if __name__ == "__main__":
-    embedding_model = EmbeddingModel()
-    print(asyncio.run(embedding_model.async_get_embedding("Hello, world!")))
-    print(
-        asyncio.run(
-            embedding_model.async_get_embeddings(["Hello, world!", "Goodbye, world!"])
-        )
-    )

aimakerspace/openai_utils/prompts.py

@@ -1,78 +0,0 @@
-import re
-
-
-class BasePrompt:
-    def __init__(self, prompt):
-        """
-        Initializes the BasePrompt object with a prompt template.
-
-        :param prompt: A string that can contain placeholders within curly braces
-        """
-        self.prompt = prompt
-        self._pattern = re.compile(r"\{([^}]+)\}")
-
-    def format_prompt(self, **kwargs):
-        """
-        Formats the prompt string using the keyword arguments provided.
-
-        :param kwargs: The values to substitute into the prompt string
-        :return: The formatted prompt string
-        """
-        matches = self._pattern.findall(self.prompt)
-        return self.prompt.format(**{match: kwargs.get(match, "") for match in matches})
-
-    def get_input_variables(self):
-        """
-        Gets the list of input variable names from the prompt string.
-
-        :return: List of input variable names
-        """
-        return self._pattern.findall(self.prompt)
-
-
-class RolePrompt(BasePrompt):
-    def __init__(self, prompt, role: str):
-        """
-        Initializes the RolePrompt object with a prompt template and a role.
-
-        :param prompt: A string that can contain placeholders within curly braces
-        :param role: The role for the message ('system', 'user', or 'assistant')
-        """
-        super().__init__(prompt)
-        self.role = role
-
-    def create_message(self, format=True, **kwargs):
-        """
-        Creates a message dictionary with a role and a formatted message.
-
-        :param kwargs: The values to substitute into the prompt string
-        :return: Dictionary containing the role and the formatted message
-        """
-        if format:
-            return {"role": self.role, "content": self.format_prompt(**kwargs)}
-        
-        return {"role": self.role, "content": self.prompt}
-
-
-class SystemRolePrompt(RolePrompt):
-    def __init__(self, prompt: str):
-        super().__init__(prompt, "system")
-
-
-class UserRolePrompt(RolePrompt):
-    def __init__(self, prompt: str):
-        super().__init__(prompt, "user")
-
-
-class AssistantRolePrompt(RolePrompt):
-    def __init__(self, prompt: str):
-        super().__init__(prompt, "assistant")
-
-
-if __name__ == "__main__":
-    prompt = BasePrompt("Hello {name}, you are {age} years old")
-    print(prompt.format_prompt(name="John", age=30))
-
-    prompt = SystemRolePrompt("Hello {name}, you are {age} years old")
-    print(prompt.create_message(name="John", age=30))
-    print(prompt.get_input_variables())

aimakerspace/text_utils.py

@@ -1,146 +0,0 @@
-import os
-import re
-from typing import List
-import PyPDF2
-
-
-class TextFileLoader:
-    def __init__(self, path: str, encoding: str = "utf-8"):
-        self.documents = []
-        self.path = path
-        self.encoding = encoding
-
-    def load(self):
-        if os.path.isdir(self.path):
-            self.load_directory()
-        else:
-            self.load_file()
-
-    def load_file(self):
-        with open(self.path, "r", encoding=self.encoding) as f:
-            self.documents.append(f.read())
-
-    def load_directory(self):
-        for root, _, files in os.walk(self.path):
-            for file in files:
-                if file.endswith(".txt"):
-                    with open(
-                        os.path.join(root, file), "r", encoding=self.encoding
-                    ) as f:
-                        self.documents.append(f.read())
-
-    def load_documents(self):
-        self.load()
-        return self.documents
-
-
-class CharacterTextSplitter:
-    def __init__(
-        self,
-        chunk_size: int = 1000,
-        chunk_overlap: int = 200,
-    ):
-        assert (
-            chunk_size > chunk_overlap
-        ), "Chunk size must be greater than chunk overlap"
-
-        self.chunk_size = chunk_size
-        self.chunk_overlap = chunk_overlap
-
-    def split(self, text: str) -> List[str]:
-        chunks = []
-        for i in range(0, len(text), self.chunk_size - self.chunk_overlap):
-            chunks.append(text[i : i + self.chunk_size])
-        return chunks
-
-    def split_texts(self, texts: List[str]) -> List[str]:
-        chunks = []
-        for text in texts:
-            chunks.extend(self.split(text))
-        return chunks
-
-
-class PDFLoader:
-    def __init__(self, path: str):
-        self.documents = []
-        self.path = path
-        print(f"PDFLoader initialized with path: {self.path}")
-
-    def load(self):
-        print(f"Loading PDF from path: {self.path}")
-        print(f"Path exists: {os.path.exists(self.path)}")
-        print(f"Is file: {os.path.isfile(self.path)}")
-        print(f"Is directory: {os.path.isdir(self.path)}")
-        print(f"File permissions: {oct(os.stat(self.path).st_mode)[-3:]}")
-        
-        try:
-            # Try to open the file first to verify access
-            with open(self.path, 'rb') as test_file:
-                pass
-            
-            # If we can open it, proceed with loading
-            if os.path.isdir(self.path):
-                self.load_directory()
-            else:
-                self.load_file()
-            
-        except IOError as e:
-            raise ValueError(f"Cannot access file at '{self.path}': {str(e)}")
-        except Exception as e:
-            raise ValueError(f"Error processing file at '{self.path}': {str(e)}")
-
-    def load_file(self):
-        with open(self.path, 'rb') as file:
-            # Create PDF reader object
-            pdf_reader = PyPDF2.PdfReader(file)
-            
-            # Extract text from each page separately
-            self.documents = []  # Clear existing documents
-            for page_num, page in enumerate(pdf_reader.pages):
-                page_text = page.extract_text() or ""  # Handle None returns
-                if page_text.strip():  # Only add non-empty pages
-                    self.documents.append(page_text)
-                else:
-                    print(f"Warning: Page {page_num + 1} is empty in {os.path.basename(self.path)}")
-
-            print(f"Loaded {len(self.documents)} pages from {os.path.basename(self.path)}")
-
-    def load_directory(self):
-        self.documents = []  # Clear existing documents
-        for root, _, files in os.walk(self.path):
-            for file in files:
-                if file.lower().endswith('.pdf'):
-                    file_path = os.path.join(root, file)
-                    try:
-                        with open(file_path, 'rb') as f:
-                            pdf_reader = PyPDF2.PdfReader(f)
-                            
-                            # Extract text from each page separately
-                            for page_num, page in enumerate(pdf_reader.pages):
-                                page_text = page.extract_text() or ""
-                                if page_text.strip():
-                                    self.documents.append(page_text)
-                                else:
-                                    print(f"Warning: Page {page_num + 1} is empty in {file}")
-                    except Exception as e:
-                        print(f"Error processing {file}: {str(e)}")
-
-    def load_documents(self):
-        if not self.documents:  # Only load if not already loaded
-            self.load()
-        return self.documents
-
-
-if __name__ == "__main__":
-    loader = TextFileLoader("data/KingLear.txt")
-    loader.load()
-    splitter = CharacterTextSplitter()
-    chunks = splitter.split_texts(loader.documents)
-    print(len(chunks))
-    print(chunks[0])
-    print("--------")
-    print(chunks[1])
-    print("--------")
-    print(chunks[-2])
-    print("--------")
-    print(chunks[-1])

aimakerspace/vectordatabase.py

@@ -1,81 +0,0 @@
-import numpy as np
-from collections import defaultdict
-from typing import List, Tuple, Callable
-from aimakerspace.openai_utils.embedding import EmbeddingModel
-import asyncio
-
-
-def cosine_similarity(vector_a: np.array, vector_b: np.array) -> float:
-    """Computes the cosine similarity between two vectors."""
-    dot_product = np.dot(vector_a, vector_b)
-    norm_a = np.linalg.norm(vector_a)
-    norm_b = np.linalg.norm(vector_b)
-    return dot_product / (norm_a * norm_b)
-
-
-class VectorDatabase:
-    def __init__(self, embedding_model: EmbeddingModel = None):
-        self.vectors = defaultdict(np.array)
-        self.embedding_model = embedding_model or EmbeddingModel()
-
-    def insert(self, key: str, vector: np.array) -> None:
-        self.vectors[key] = vector
-
-    def search(
-        self,
-        query_vector: np.array,
-        k: int,
-        distance_measure: Callable = cosine_similarity,
-    ) -> List[Tuple[str, float]]:
-        scores = [
-            (key, distance_measure(query_vector, vector))
-            for key, vector in self.vectors.items()
-        ]
-        return sorted(scores, key=lambda x: x[1], reverse=True)[:k]
-
-    def search_by_text(
-        self,
-        query_text: str,
-        k: int,
-        distance_measure: Callable = cosine_similarity,
-        return_as_text: bool = False,
-    ) -> List[Tuple[str, float]]:
-        query_vector = self.embedding_model.get_embedding(query_text)
-        results = self.search(query_vector, k, distance_measure)
-        return [result[0] for result in results] if return_as_text else results
-
-    def retrieve_from_key(self, key: str) -> np.array:
-        return self.vectors.get(key, None)
-
-    async def abuild_from_list(self, list_of_text: List[str]) -> "VectorDatabase":
-        embeddings = await self.embedding_model.async_get_embeddings(list_of_text)
-        for text, embedding in zip(list_of_text, embeddings):
-            self.insert(text, np.array(embedding))
-        return self
-
-
-if __name__ == "__main__":
-    list_of_text = [
-        "I like to eat broccoli and bananas.",
-        "I ate a banana and spinach smoothie for breakfast.",
-        "Chinchillas and kittens are cute.",
-        "My sister adopted a kitten yesterday.",
-        "Look at this cute hamster munching on a piece of broccoli.",
-    ]
-
-    vector_db = VectorDatabase()
-    vector_db = asyncio.run(vector_db.abuild_from_list(list_of_text))
-    k = 2
-
-    searched_vector = vector_db.search_by_text("I think fruit is awesome!", k=k)
-    print(f"Closest {k} vector(s):", searched_vector)
-
-    retrieved_vector = vector_db.retrieve_from_key(
-        "I like to eat broccoli and bananas."
-    )
-    print("Retrieved vector:", retrieved_vector)
-
-    relevant_texts = vector_db.search_by_text(
-        "I think fruit is awesome!", k=k, return_as_text=True
-    )
-    print(f"Closest {k} text(s):", relevant_texts)

app.py

@@ -0,0 +1,295 @@
+import os
+import chainlit as cl
+from chainlit.playground.providers.openai import ChatOpenAI
+from chainlit.element import Text
+from langchain_core.prompts import ChatPromptTemplate
+from langchain.schema.output_parser import StrOutputParser
+from langchain_community.document_loaders import DirectoryLoader
+from langchain_community.document_loaders import PyPDFLoader
+from langchain_openai.embeddings import OpenAIEmbeddings
+from langchain_community.vectorstores import Qdrant
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain_core.runnables import RunnablePassthrough
+from langchain_core.messages import AIMessage, HumanMessage
+from operator import itemgetter
+from collections import defaultdict
+from langchain_core.documents import Document
+import tiktoken
+import re
+import functools
+
+# Configure OpenAI API key from environment variable
+if not os.environ.get("OPENAI_API_KEY"):
+    os.environ["OPENAI_API_KEY"] = os.environ.get("OPENAI_API_KEY_BACKUP", "")
+
+# Document loading and processing
+@cl.cache
+def load_and_process_documents():
+    # Load all PDF documents (each page as a separate document)
+    path = "data/"
+    loader = DirectoryLoader(path, glob="*.pdf", loader_cls=PyPDFLoader)
+    all_docs = loader.load()
+    
+    # Create a mapping of merged document chunks back to original pages
+    page_map = {}
+    current_index = 0
+    
+    # For source tracking, we'll store page information before merging
+    docs_by_source = defaultdict(list)
+    
+    # Group documents by their source file
+    for doc in all_docs:
+        source = doc.metadata.get("source", "")
+        docs_by_source[source].append(doc)
+    
+    # Merge pages from the same PDF but track page ranges
+    merged_docs = []
+    for source, source_docs in docs_by_source.items():
+        # Sort by page number if available
+        source_docs.sort(key=lambda x: x.metadata.get("page", 0))
+        
+        # Merge the content
+        merged_content = ""
+        page_ranges = []
+        current_pos = 0
+        
+        for doc in source_docs:
+            # Get the page number (1-indexed for human readability)
+            page_num = doc.metadata.get("page", 0) + 1
+            
+            # Add a separator between pages for clarity
+            if merged_content:
+                merged_content += "\n\n"
+            
+            # Record where this page's content starts in the merged document
+            start_pos = len(merged_content)
+            merged_content += doc.page_content
+            end_pos = len(merged_content)
+            
+            # Store the mapping of character ranges to original page numbers
+            page_ranges.append({
+                "start": start_pos,
+                "end": end_pos,
+                "page": page_num,
+                "source": source
+            })
+        
+        # Create merged metadata that includes page mapping information
+        merged_metadata = {
+            "source": source,
+            "title": source.split("/")[-1],
+            "page_count": len(source_docs),
+            "merged": True,
+            "page_ranges": page_ranges  # Store the page ranges for later reference
+        }
+        
+        # Create a new document with the merged content
+        merged_doc = Document(page_content=merged_content, metadata=merged_metadata)
+        merged_docs.append(merged_doc)
+    
+    # tiktoken_len counts tokens (not characters) using the gpt-4o-mini tokenizer
+    def tiktoken_len(text):
+        tokens = tiktoken.encoding_for_model("gpt-4o-mini").encode(
+            text,
+        )
+        return len(tokens)
+    
+    # Process splits to add page info based on character position
+    def add_page_info_to_splits(splits):
+        for split in splits:
+            # Get the start position of this chunk
+            start_pos = split.metadata.get("start_index", 0)
+            end_pos = start_pos + len(split.page_content)
+            
+            # Find which page this chunk belongs to
+            if "page_ranges" in split.metadata:
+                for page_range in split.metadata["page_ranges"]:
+                    # If chunk significantly overlaps with this page range
+                    if (start_pos <= page_range["end"] and 
+                        end_pos >= page_range["start"]):
+                        # Use this page number
+                        split.metadata["page"] = page_range["page"]
+                        break
+        return splits
+    
+    # Split the text with start index tracking
+    text_splitter = RecursiveCharacterTextSplitter(
+        chunk_size=300,
+        chunk_overlap=50,
+        length_function=tiktoken_len,
+        add_start_index=True
+    )
+    
+    # Split and then process to add page information
+    raw_splits = text_splitter.split_documents(merged_docs)
+    split_chunks = add_page_info_to_splits(raw_splits)
+    
+    return split_chunks
+
+# Load and process documents at startup
+docs = load_and_process_documents()
+
+# Set up embeddings and vector store
+@cl.cache
+def setup_vector_store():
+    embedding_model = OpenAIEmbeddings(model="text-embedding-3-small")
+    vectorstore = Qdrant.from_documents(
+        docs,
+        embedding_model,
+        location=":memory:",
+        collection_name="kohavi_ab_testing_pdf_collection",
+    )
+    return vectorstore
+
+# Setup vector store at startup
+vectorstore = setup_vector_store()
+retriever = vectorstore.as_retriever(search_kwargs={"k": 5})
+
+# Define prompts
+RAG_PROMPT = """
+CONTEXT:
+{context}
+
+QUERY:
+{question}
+
+You are a helpful assistant. Use the available context to answer the question. Do not use your own knowledge! If you cannot answer the question based on the context, you must say "I don't know".
+"""
+
+REPHRASE_QUERY_PROMPT = """
+QUERY:
+{question}
+
+You are a helpful assistant. Rephrase the provided query to be more specific and to the point in order to improve retrieval in our RAG pipeline about AB Testing.
+"""
+
+rag_prompt = ChatPromptTemplate.from_template(RAG_PROMPT)
+rephrase_query_prompt = ChatPromptTemplate.from_template(REPHRASE_QUERY_PROMPT)
+
+# Setup chat model
+def get_openai_chat_model():
+    model = ChatOpenAI(
+        model="gpt-4-turbo",
+        temperature=0,
+        streaming=True,
+    )
+    return model
+
+# Setup rephrased retrieval chain
+def setup_rephrased_retriever_chain():
+    chat_model = get_openai_chat_model()
+    
+    def retrieve_and_format_documents(query):
+        rephrased_query_chain = rephrase_query_prompt | chat_model | StrOutputParser()
+        
+        @cl.step(name="Rephrasing query for better retrieval")
+        async def rephrase_query():
+            rephrased_query = await cl.make_async(rephrased_query_chain.invoke)({"question": query})
+            await cl.Message(content=f"Rephrased query: {rephrased_query}").send()
+            return rephrased_query
+        
+        rephrased_query_result = cl.run_sync(rephrase_query())
+        
+        # Get relevant documents based on rephrased query
+        docs = retriever.get_relevant_documents(rephrased_query_result)
+        
+        # Extract sources from the documents for later display
+        sources = []
+        for doc in docs:
+            source_path = doc.metadata.get("source", "")
+            filename = source_path.split("/")[-1] if "/" in source_path else source_path
+            
+            sources.append({
+                "title": filename,
+                "page": doc.metadata.get("page", "unknown"),
+            })
+        
+        # Format documents into a string for the context
+        formatted_docs = "\n\n".join([doc.page_content for doc in docs])
+        
+        return {"context": formatted_docs, "sources": sources, "question": query}
+    
+    return retrieve_and_format_documents
+
+@cl.on_chat_start
+async def on_chat_start():
+    # Initialize the chat model and chain in the user session
+    chat_model = get_openai_chat_model()
+    retriever_chain = setup_rephrased_retriever_chain()
+    
+    # Store in user session
+    cl.user_session.set("chat_model", chat_model)
+    cl.user_session.set("retriever_chain", retriever_chain)
+    
+    welcome_message = """
+# 📊 AB Testing RAG Agent
+
+This agent can answer questions about AB Testing using a collection of PDF documents.
+
+**Examples of questions you can ask:**
+- What is AB Testing?
+- How do I interpret p-values in experiments?
+- What are the best practices for running experiments with low traffic?
+- How should I choose metrics for my experiments?
+    """
+    
+    await cl.Message(content=welcome_message).send()
+
+@cl.on_message
+async def on_message(message: cl.Message):
+    # Get the chat model and retriever chain from user session
+    chat_model = cl.user_session.get("chat_model")
+    retriever_chain = cl.user_session.get("retriever_chain")
+    
+    query = message.content
+    
+    # Step 1: Retrieve documents and extract sources
+    @cl.step(name="Retrieving relevant documents")
+    async def retrieve_docs():
+        result = await cl.make_async(retriever_chain)(query)
+        return result
+    
+    retrieval_result = await retrieve_docs()
+    
+    # Store sources for later display
+    sources = retrieval_result.get("sources", [])
+    
+    # Step 2: Generate response with the retrieved context
+    msg = cl.Message(content="")
+    await msg.send()
+    
+    # Stream the response
+    final_answer = await chat_model.astream(
+        [
+            HumanMessage(
+                content=rag_prompt.format(
+                    context=retrieval_result.get("context", ""),
+                    question=query
+                )
+            )
+        ],
+        callbacks=[
+            cl.AsyncLangchainCallbackHandler(
+                stream_to_message=msg,
+                append_to_message=True,
+            )
+        ],
+    )
+    
+    # Display sources
+    if sources:
+        source_elements = []
+        for i, source in enumerate(sources):
+            title = source.get("title", "Unknown")
+            page = source.get("page", "Unknown")
+            
+            source_elements.append(
+                Text(
+                    name=f"Source {i+1}",
+                    content=f"{title}, Page: {page}",
+                    display="inline",
+                )
+            )
+        
+        await msg.elements.extend(source_elements)
+        await msg.update() 

chainlit.json

@@ -0,0 +1,18 @@
+{
+  "name": "AB Testing RAG Agent",
+  "description": "Ask questions about AB Testing and get answers with sources",
+  "theme": {
+    "color": "#4F6AFF",
+    "darkMode": true
+  },
+  "features": {
+    "multi_modal": false
+  },
+  "ui": {
+    "show_header": true,
+    "show_footer": true,
+    "message_size": "large",
+    "default_expand_messages": true
+  },
+  "markdown": "chainlit.md"
+} 

chainlit.md

@@ -0,0 +1,28 @@
+# 📊 AB Testing RAG Agent
+
+This application helps answer questions about AB Testing using a collection of PDF documents.
+
+## Features
+
+- **Document Retrieval**: Retrieves relevant information from AB Testing documents
+- **Query Rephrasing**: Improves retrieval by rephrasing your query for better results
+- **Source References**: Shows the exact document sources used to answer your question
+- **Streaming Interface**: See the response as it's being generated
+
+## Example Questions
+
+- What is AB Testing?
+- How do I interpret p-values in experiments?
+- What are the best practices for running experiments with low traffic?
+- What is CUPED and when should I use it?
+- How should I choose metrics for my experiments?
+
+## About
+
+This application was built using:
+- LangChain
+- OpenAI models
+- Qdrant vector database
+- Chainlit frontend
+
+The information is sourced from various AB Testing documents including research papers and articles by experts in the field. 

docker-compose.yml

@@ -0,0 +1,16 @@
+version: '3'
+
+services:
+  ab-testing-rag:
+    build:
+      context: .
+      dockerfile: Dockerfile
+    container_name: ab-testing-rag
+    ports:
+      - "8000:8000"
+    volumes:
+      - ./data:/app/data
+      - ./.env:/app/.env
+    environment:
+      - OPENAI_API_KEY=${OPENAI_API_KEY}
+    restart: unless-stopped 

huggingface-space.yml

@@ -0,0 +1,20 @@
+---
+title: AB Testing RAG Agent
+emoji: 📊
+colorFrom: indigo
+colorTo: blue
+sdk: docker
+pinned: false
+license: mit
+app_port: 8000
+---
+
+# AB Testing RAG Agent
+
+This is a Hugging Face Spaces deployment of the AB Testing RAG Agent. The agent helps answer questions about AB Testing using Ron Kohavi's work and articles from Arxiv.
+
+## Environment Variables
+
+Environment variables needed:
+
+- `OPENAI_API_KEY`: Your OpenAI API key 

requirements.txt

@@ -0,0 +1,13 @@
+chainlit==1.0.301
+langchain==0.1.12
+langchain-community==0.0.30
+langchain-openai==0.0.6
+langchain-core==0.1.30
+langchain-qdrant==0.0.1
+qdrant-client==1.7.3
+openai==1.12.0
+tiktoken==0.5.2
+python-dotenv==1.0.1
+unstructured==0.12.5
+pypdf==3.17.4
+numpy==1.26.3 

scripts/__init__.py

@@ -0,0 +1,3 @@
+"""
+Utility scripts for the AB Testing RAG Agent.
+""" 

scripts/prepare_data.py

@@ -0,0 +1,40 @@
+#!/usr/bin/env python3
+"""
+Utility script to copy PDF files from notebook_version/data to the data directory.
+This makes it easier to set up the application with the existing data.
+"""
+
+import os
+import shutil
+import glob
+from pathlib import Path
+
+def main():
+    # Create data directory if it doesn't exist
+    data_dir = Path("data")
+    data_dir.mkdir(exist_ok=True)
+    
+    # Source directory with PDF files
+    source_dir = Path("notebook_version/data")
+    if not source_dir.exists():
+        print(f"Error: Source directory {source_dir} does not exist.")
+        return
+    
+    # Find all PDF files in the source directory
+    pdf_files = glob.glob(str(source_dir / "*.pdf"))
+    if not pdf_files:
+        print(f"No PDF files found in {source_dir}")
+        return
+    
+    # Copy each PDF file to the data directory
+    for pdf_file in pdf_files:
+        filename = os.path.basename(pdf_file)
+        destination = data_dir / filename
+        
+        print(f"Copying {filename}...")
+        shutil.copy2(pdf_file, destination)
+    
+    print(f"Copied {len(pdf_files)} PDF files to {data_dir}")
+
+if __name__ == "__main__":
+    main()