Develop the streamlit app

.dockerignore

@@ -43,4 +43,7 @@ ENV/
 notebook_version/
 
 # Ignore PDF data files (will be mounted at runtime)
-data/*.pdf 
+data/*.pdf 
+
+# Pre-processed data (will be generated inside the container)
+processed_data/ 

.gitignore

@@ -44,5 +44,9 @@ env/
 # Mac OS
 .DS_Store
 
-# Local data
-/data/*.pdf 
+# Local data - keep PDFs private and never commit them
+/data/*.pdf
+/data/**/*.pdf
+
+# Do NOT ignore processed_data anymore - we want to commit this
+# /processed_data/ 

Dockerfile

@@ -14,18 +14,42 @@ COPY requirements.txt .
 # Install Python dependencies
 RUN pip install --no-cache-dir -r requirements.txt
 
-# Copy application code
+# Copy the rest of the application code
 COPY . .
 
-# Create data directory for PDFs
+# Create data directory for PDFs (if not already created)
 RUN mkdir -p data
 
-# Expose port for Chainlit
-EXPOSE 8000
+# Create directory for pre-processed data
+RUN mkdir -p processed_data
 
-# Set environment variables
-ENV PYTHONPATH=/app
-ENV OPENAI_API_KEY=${OPENAI_API_KEY}
+# Configure Streamlit to run in headless mode (no welcome screen)
+RUN mkdir -p /root/.streamlit && \
+    echo '[general]' > /root/.streamlit/config.toml && \
+    echo 'email = ""' >> /root/.streamlit/config.toml && \
+    echo 'showWarningOnDirectExecution = false' >> /root/.streamlit/config.toml && \
+    echo '' >> /root/.streamlit/config.toml && \
+    echo '[server]' >> /root/.streamlit/config.toml && \
+    echo 'headless = true' >> /root/.streamlit/config.toml
 
-# Command to run the application
-CMD ["chainlit", "run", "app.py", "--port", "8000", "--host", "0.0.0.0"] 
+# Set environment variables
+ENV HOST=0.0.0.0
+ENV PORT=8000
+
+# Allow both Chainlit and Streamlit apps to run (use environment variable to choose)
+ENV APP_MODE=streamlit
+
+# Create the entrypoint script
+RUN echo '#!/bin/bash' > /app/entrypoint.sh && \
+    echo 'if [ "$APP_MODE" = "chainlit" ]; then' >> /app/entrypoint.sh && \
+    echo '    chainlit run app.py --host $HOST --port $PORT' >> /app/entrypoint.sh && \
+    echo 'else' >> /app/entrypoint.sh && \
+    echo '    streamlit run streamlit_app.py --server.address $HOST --server.port $PORT' >> /app/entrypoint.sh && \
+    echo 'fi' >> /app/entrypoint.sh && \
+    chmod +x /app/entrypoint.sh
+
+# Expose the port
+EXPOSE $PORT
+
+# Run the application
+ENTRYPOINT ["/app/entrypoint.sh"] 

README.md

@@ -8,6 +8,8 @@ An application that helps answer questions about AB Testing using a collection o
 - **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
+- **Pre-processed Data**: Faster startup with pre-processed document embeddings
+- **Privacy Preserving**: Deploy with pre-processed data only, keeping source PDFs private
 
 ## Prerequisites
 
@@ -52,12 +54,17 @@ An application that helps answer questions about AB Testing using a collection o
    # Copy your AB Testing PDFs to the data directory
    ```
 
-7. Run the application
+7. Pre-process the PDF files (this step significantly speeds up application startup)
+   ```bash
+   python scripts/preprocess_data.py
+   ```
+
+8. Run the application
    ```bash
    chainlit run app.py
    ```
 
-8. Open your browser to `http://localhost:8000`
+9. Open your browser to `http://localhost:8000`
 
 ### Using Docker
 
@@ -83,28 +90,65 @@ An application that helps answer questions about AB Testing using a collection o
    # Copy your AB Testing PDFs to the data directory
    ```
 
-5. Build and run with Docker Compose
+5. Build and run with Docker Compose (the container will automatically pre-process PDFs at startup)
    ```bash
    docker-compose up -d
    ```
 
 6. Open your browser to `http://localhost:8000`
 
-## Deploying to Hugging Face Spaces
+## Deploying to Hugging Face Spaces (Privacy-Preserving Method)
+
+This method allows you to deploy your app without exposing your source PDF files:
+
+1. Pre-process your PDFs locally:
+   ```bash
+   # Ensure PDFs are in the data directory
+   python scripts/preprocess_data.py
+   ```
+
+2. Commit only the pre-processed data (not the PDFs) to your repository:
+   ```bash
+   # PDFs are excluded in .gitignore already
+   git add processed_data app.py requirements.txt Dockerfile scripts/docker-entrypoint.sh
+   git commit -m "Add pre-processed data for deployment"
+   ```
+
+3. Create a new Hugging Face Space with Docker deployment
 
-1. Create a new Hugging Face Space with Docker deployment
-2. Set the following environment variables in your Space settings:
+4. 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
+
+5. Push your code to the Space's repository:
+   ```bash
+   git push huggingface main
+   ```
+
+The advantage of this approach is that:
+- Your original PDF files remain private and are never uploaded
+- Only the pre-processed embeddings and chunks are deployed
+- The app will work immediately without needing to process PDFs at runtime
+
+## Pre-processing Explained
+
+The pre-processing step:
+1. Loads PDF files from the `data` directory
+2. Splits documents into chunks with correct page references
+3. Generates embeddings for each chunk using OpenAI's embedding model
+4. Saves both document chunks and the vector database to disk
+
+This approach has several advantages:
+- The app starts up much faster since it doesn't need to process PDFs at runtime
+- You don't need to upload large PDF files to Hugging Face
+- The vector database is persisted, providing consistent performance
+- You can deploy without exposing your source materials
 
 ## 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:
+1. The application loads pre-processed document chunks and vector store from disk
+2. When a user asks a question:
    - The query is rephrased to be more specific for better retrieval
-   - Relevant document chunks are retrieved
+   - Relevant document chunks are retrieved from the vector store
    - OpenAI's GPT model generates an answer based on the retrieved context
    - Source references are tracked and displayed alongside the answer
 

app.log

@@ -0,0 +1 @@
+2025-04-29 21:20:25 - Loaded .env file

app.py

@@ -1,149 +1,106 @@
 import os
+import pickle
 import chainlit as cl
-from chainlit.playground.providers.openai import ChatOpenAI
+from langchain_openai.chat_models 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
+from pathlib import Path
 import functools
+from qdrant_client import QdrantClient
+from langchain_core.vectorstores import VectorStoreRetriever
+from langchain_core.documents import Document
 
 # 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
+# Paths to pre-processed data
+PROCESSED_DATA_DIR = Path("processed_data")
+CHUNKS_FILE = PROCESSED_DATA_DIR / "document_chunks.pkl"
+QDRANT_DIR = PROCESSED_DATA_DIR / "qdrant_vectorstore"
+
+# Check if pre-processed data exists
+def check_preprocessed_data():
+    """Check if pre-processed data exists and is readable."""
+    if not PROCESSED_DATA_DIR.exists():
+        raise FileNotFoundError(f"Processed data directory not found: {PROCESSED_DATA_DIR}. Please run scripts/preprocess_data.py first.")
+    
+    if not CHUNKS_FILE.exists():
+        raise FileNotFoundError(f"Document chunks file not found: {CHUNKS_FILE}. Please run scripts/preprocess_data.py first.")
+    
+    if not QDRANT_DIR.exists():
+        raise FileNotFoundError(f"Vector store directory not found: {QDRANT_DIR}. Please run scripts/preprocess_data.py first.")
+
+# Load pre-processed document chunks
 @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()
+def load_document_chunks():
+    """Load pre-processed document chunks from disk."""
+    with open(CHUNKS_FILE, 'rb') as f:
+        return pickle.load(f)
+
+# Load and create a custom retriever for Qdrant
+@cl.cache
+def load_qdrant_retriever():
+    """Load Qdrant client and create a retriever."""
+    embedding_model = OpenAIEmbeddings(model="text-embedding-3-small")
     
-    # Create a mapping of merged document chunks back to original pages
-    page_map = {}
-    current_index = 0
+    # Load the document chunks for metadata access
+    chunks = load_document_chunks()
     
-    # For source tracking, we'll store page information before merging
-    docs_by_source = defaultdict(list)
+    # Create a dictionary mapping IDs to documents for quick lookup
+    docs_by_id = {i: doc for i, doc in enumerate(chunks)}
     
-    # Group documents by their source file
-    for doc in all_docs:
-        source = doc.metadata.get("source", "")
-        docs_by_source[source].append(doc)
+    # Initialize Qdrant client
+    client = QdrantClient(path=str(QDRANT_DIR))
     
-    # 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
+    # Create a custom retriever function
+    def retrieve_docs(query, k=5):
+        # Get query embedding
+        query_embedding = embedding_model.embed_query(query)
         
-        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
-            })
+        # Search Qdrant
+        results = client.search(
+            collection_name="kohavi_ab_testing_pdf_collection",
+            query_vector=query_embedding,
+            limit=k
+        )
         
-        # 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
-        }
+        # Convert results to documents
+        documents = []
+        for result in results:
+            doc_id = result.id
+            if doc_id in docs_by_id:
+                documents.append(docs_by_id[doc_id])
         
-        # 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)
+        return documents
     
-    # 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
+    # Create a VectorStoreRetriever from the custom retriever function
+    retriever = VectorStoreRetriever(
+        vectorstore=None,  # Not using a standard vectorstore
+        search_type="similarity",
+        search_kwargs={"k": 5},
+        retriever_fn=retrieve_docs
     )
     
-    # 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
+    return retriever
 
-# 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})
+# Check for pre-processed data at startup
+try:
+    check_preprocessed_data()
+    # Load document chunks
+    docs = load_document_chunks()
+    # Load retriever
+    retriever = load_qdrant_retriever()
+    print(f"Successfully loaded pre-processed data with {len(docs)} document chunks")
+except FileNotFoundError as e:
+    # If pre-processed data doesn't exist, print a helpful error message
+    print(f"ERROR: {e}")
+    print("\nTo pre-process your data, run:")
+    print("  python scripts/preprocess_data.py")
+    # Raise the error to prevent the app from starting without pre-processed data
+    raise
 
 # Define prompts
 RAG_PROMPT = """
@@ -260,13 +217,14 @@ async def on_message(message: cl.Message):
     
     # Stream the response
     final_answer = await chat_model.astream(
-        [
-            HumanMessage(
-                content=rag_prompt.format(
+        messages=[
+            {
+                "role": "user",
+                "content": rag_prompt.format(
                     context=retrieval_result.get("context", ""),
                     question=query
                 )
-            )
+            }
         ],
         callbacks=[
             cl.AsyncLangchainCallbackHandler(
@@ -281,6 +239,9 @@ async def on_message(message: cl.Message):
         source_elements = []
         for i, source in enumerate(sources):
             title = source.get("title", "Unknown")
+            # Remove the .pdf extension if present
+            if title.lower().endswith('.pdf'):
+                title = title[:-4]  # Remove the last 4 characters (.pdf)
             page = source.get("page", "Unknown")
             
             source_elements.append(

app_debug.py

@@ -0,0 +1,302 @@
+import os
+import pickle
+import sys
+import chainlit as cl
+from langchain_openai.chat_models import ChatOpenAI
+from chainlit.element import Text
+from langchain_core.prompts import ChatPromptTemplate
+from langchain.schema.output_parser import StrOutputParser
+from langchain_openai.embeddings import OpenAIEmbeddings
+from pathlib import Path
+import functools
+from qdrant_client import QdrantClient
+from langchain_core.vectorstores import VectorStoreRetriever
+from langchain_core.documents import Document
+
+print("Starting app_debug.py execution...")
+
+# 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", "")
+print("Configured OpenAI API key")
+
+# Paths to pre-processed data
+PROCESSED_DATA_DIR = Path("processed_data")
+CHUNKS_FILE = PROCESSED_DATA_DIR / "document_chunks.pkl"
+QDRANT_DIR = PROCESSED_DATA_DIR / "qdrant_vectorstore"
+print(f"Configured paths: {PROCESSED_DATA_DIR}, {CHUNKS_FILE}, {QDRANT_DIR}")
+
+# Check if pre-processed data exists
+def check_preprocessed_data():
+    """Check if pre-processed data exists and is readable."""
+    print("Checking preprocessed data...")
+    if not PROCESSED_DATA_DIR.exists():
+        raise FileNotFoundError(f"Processed data directory not found: {PROCESSED_DATA_DIR}. Please run scripts/preprocess_data.py first.")
+    
+    if not CHUNKS_FILE.exists():
+        raise FileNotFoundError(f"Document chunks file not found: {CHUNKS_FILE}. Please run scripts/preprocess_data.py first.")
+    
+    if not QDRANT_DIR.exists():
+        raise FileNotFoundError(f"Vector store directory not found: {QDRANT_DIR}. Please run scripts/preprocess_data.py first.")
+    print("Preprocessed data checks passed")
+
+# Load pre-processed document chunks
+@cl.cache
+def load_document_chunks():
+    """Load pre-processed document chunks from disk."""
+    print("Loading document chunks...")
+    with open(CHUNKS_FILE, 'rb') as f:
+        chunks = pickle.load(f)
+    print(f"Loaded {len(chunks)} document chunks")
+    return chunks
+
+# Load and create a custom retriever for Qdrant
+@cl.cache
+def load_qdrant_retriever():
+    """Load Qdrant client and create a retriever."""
+    print("Creating Qdrant retriever...")
+    embedding_model = OpenAIEmbeddings(model="text-embedding-3-small")
+    print("Created embedding model")
+    
+    # Load the document chunks for metadata access
+    chunks = load_document_chunks()
+    
+    # Create a dictionary mapping IDs to documents for quick lookup
+    docs_by_id = {i: doc for i, doc in enumerate(chunks)}
+    print(f"Created mapping of {len(docs_by_id)} documents")
+    
+    # Initialize Qdrant client
+    print(f"Initializing Qdrant client at {QDRANT_DIR}...")
+    client = QdrantClient(path=str(QDRANT_DIR))
+    print("Initialized Qdrant client")
+    
+    # Create a custom retriever function
+    def retrieve_docs(query, k=5):
+        print(f"Retrieving documents for query: {query[:30]}...")
+        # Get query embedding
+        query_embedding = embedding_model.embed_query(query)
+        
+        # Search Qdrant
+        results = client.search(
+            collection_name="kohavi_ab_testing_pdf_collection",
+            query_vector=query_embedding,
+            limit=k
+        )
+        
+        # Convert results to documents
+        documents = []
+        for result in results:
+            doc_id = result.id
+            if doc_id in docs_by_id:
+                documents.append(docs_by_id[doc_id])
+        print(f"Retrieved {len(documents)} documents")
+        return documents
+    
+    # Create a VectorStoreRetriever from the custom retriever function
+    retriever = VectorStoreRetriever(
+        vectorstore=None,  # Not using a standard vectorstore
+        search_type="similarity",
+        search_kwargs={"k": 5},
+        retriever_fn=retrieve_docs
+    )
+    print("Created retriever")
+    return retriever
+
+# Check for pre-processed data at startup
+try:
+    print("Checking for pre-processed data...")
+    check_preprocessed_data()
+    # Load document chunks
+    print("Loading document chunks...")
+    docs = load_document_chunks()
+    # Load retriever
+    print("Loading retriever...")
+    retriever = load_qdrant_retriever()
+    print(f"Successfully loaded pre-processed data with {len(docs)} document chunks")
+except FileNotFoundError as e:
+    # If pre-processed data doesn't exist, print a helpful error message
+    print(f"ERROR: {e}")
+    print("\nTo pre-process your data, run:")
+    print("  python scripts/preprocess_data.py")
+    # Raise the error to prevent the app from starting without pre-processed data
+    raise
+
+# 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)
+print("Defined prompts")
+
+# Setup chat model
+def get_openai_chat_model():
+    print("Creating OpenAI chat model...")
+    model = ChatOpenAI(
+        model="gpt-4-turbo",
+        temperature=0,
+        streaming=True,
+    )
+    print("Created OpenAI chat model")
+    return model
+
+# Setup rephrased retrieval chain
+def setup_rephrased_retriever_chain():
+    print("Setting up rephrased retriever chain...")
+    chat_model = get_openai_chat_model()
+    
+    def retrieve_and_format_documents(query):
+        print(f"Retrieving documents for: {query[:30]}...")
+        rephrased_query_chain = rephrase_query_prompt | chat_model | StrOutputParser()
+        
+        @cl.step(name="Rephrasing query for better retrieval")
+        async def rephrase_query():
+            print("Rephrasing query...")
+            rephrased_query = await cl.make_async(rephrased_query_chain.invoke)({"question": query})
+            print(f"Rephrased query: {rephrased_query[:30]}...")
+            await cl.Message(content=f"Rephrased query: {rephrased_query}").send()
+            return rephrased_query
+        
+        rephrased_query_result = cl.run_sync(rephrase_query())
+        print(f"Received rephrased query: {rephrased_query_result[:30]}...")
+        
+        # Get relevant documents based on rephrased query
+        docs = retriever.get_relevant_documents(rephrased_query_result)
+        print(f"Retrieved {len(docs)} documents")
+        
+        # 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])
+        print(f"Formatted {len(docs)} documents")
+        
+        return {"context": formatted_docs, "sources": sources, "question": query}
+    
+    print("Created retrieve_and_format_documents function")
+    return retrieve_and_format_documents
+
+print("Defined setup functions")
+
+@cl.on_chat_start
+async def on_chat_start():
+    print("on_chat_start called")
+    # 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)
+    print("Stored chat model and retriever chain in user session")
+    
+    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()
+    print("Sent welcome message")
+
+@cl.on_message
+async def on_message(message: cl.Message):
+    print(f"on_message called with content: {message.content[:30]}...")
+    # 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")
+    print("Retrieved chat model and retriever chain from user session")
+    
+    query = message.content
+    
+    # Step 1: Retrieve documents and extract sources
+    @cl.step(name="Retrieving relevant documents")
+    async def retrieve_docs():
+        print("retrieve_docs step started")
+        result = await cl.make_async(retriever_chain)(query)
+        print("retrieve_docs step completed")
+        return result
+    
+    retrieval_result = await retrieve_docs()
+    print("Retrieved documents and extracted sources")
+    
+    # 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()
+    print("Sent empty message to be streamed to")
+    
+    # Stream the response
+    final_answer = await chat_model.astream(
+        messages=[
+            {
+                "role": "user",
+                "content": rag_prompt.format(
+                    context=retrieval_result.get("context", ""),
+                    question=query
+                )
+            }
+        ],
+        callbacks=[
+            cl.AsyncLangchainCallbackHandler(
+                stream_to_message=msg,
+                append_to_message=True,
+            )
+        ],
+    )
+    print("Streamed response from OpenAI")
+    
+    # Display sources
+    if sources:
+        source_elements = []
+        for i, source in enumerate(sources):
+            title = source.get("title", "Unknown")
+            # Remove the .pdf extension if present
+            if title.lower().endswith('.pdf'):
+                title = title[:-4]  # Remove the last 4 characters (.pdf)
+            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()
+        print("Updated message with source elements")
+
+print("Defined event handlers")
+print("Ready to start Chainlit app") 

debug_app.py

@@ -0,0 +1,53 @@
+import os
+import pickle
+import sys
+from pathlib import Path
+from qdrant_client import QdrantClient
+
+# Paths to pre-processed data
+PROCESSED_DATA_DIR = Path("processed_data")
+CHUNKS_FILE = PROCESSED_DATA_DIR / "document_chunks.pkl"
+QDRANT_DIR = PROCESSED_DATA_DIR / "qdrant_vectorstore"
+
+def main():
+    print("Starting debug process...")
+    
+    # Step 1: Check if directories and files exist
+    print(f"Checking if processed_data directory exists: {PROCESSED_DATA_DIR.exists()}")
+    print(f"Checking if document_chunks.pkl exists: {CHUNKS_FILE.exists()}")
+    print(f"Checking if qdrant_vectorstore directory exists: {QDRANT_DIR.exists()}")
+    
+    # Step 2: Try to load document chunks
+    try:
+        print("Attempting to load document chunks...")
+        with open(CHUNKS_FILE, 'rb') as f:
+            chunks = pickle.load(f)
+        print(f"Successfully loaded {len(chunks)} document chunks")
+    except Exception as e:
+        print(f"Error loading document chunks: {e}")
+        return
+    
+    # Step 3: Try to initialize Qdrant client
+    try:
+        print(f"Attempting to initialize Qdrant client at {QDRANT_DIR}...")
+        client = QdrantClient(path=str(QDRANT_DIR))
+        
+        # Try to list collections
+        collections = client.get_collections()
+        print(f"Available collections: {collections}")
+        
+        # Try to get collection info
+        try:
+            collection_info = client.get_collection("kohavi_ab_testing_pdf_collection")
+            print(f"Collection info: {collection_info}")
+        except Exception as e:
+            print(f"Error getting collection info: {e}")
+        
+    except Exception as e:
+        print(f"Error initializing Qdrant client: {e}")
+        return
+    
+    print("Debug process completed successfully!")
+
+if __name__ == "__main__":
+    main() 

debug_output.log

@@ -0,0 +1 @@
+2025-04-29 21:22:54 - Loaded .env file

debug_start.py

@@ -0,0 +1,88 @@
+import os
+import pickle
+import sys
+from pathlib import Path
+from qdrant_client import QdrantClient
+from langchain_openai.embeddings import OpenAIEmbeddings
+from dotenv import load_dotenv
+
+# Load .env file for OpenAI API key
+load_dotenv()
+print("Loaded .env file")
+
+# Paths to pre-processed data
+PROCESSED_DATA_DIR = Path("processed_data")
+CHUNKS_FILE = PROCESSED_DATA_DIR / "document_chunks.pkl"
+QDRANT_DIR = PROCESSED_DATA_DIR / "qdrant_vectorstore"
+
+def main():
+    print("Starting debug process...")
+    
+    # Step 1: Check if directories and files exist
+    print(f"Checking if processed_data directory exists: {PROCESSED_DATA_DIR.exists()}")
+    print(f"Checking if document_chunks.pkl exists: {CHUNKS_FILE.exists()}")
+    print(f"Checking if qdrant_vectorstore directory exists: {QDRANT_DIR.exists()}")
+    
+    # Step 2: Try to load document chunks
+    try:
+        print("Attempting to load document chunks...")
+        with open(CHUNKS_FILE, 'rb') as f:
+            chunks = pickle.load(f)
+        print(f"Successfully loaded {len(chunks)} document chunks")
+    except Exception as e:
+        print(f"Error loading document chunks: {e}")
+        return
+    
+    # Step 3: Try to initialize OpenAI embeddings
+    try:
+        print("Initializing OpenAI embeddings...")
+        embedding_model = OpenAIEmbeddings(model="text-embedding-3-small")
+        print("Successfully initialized OpenAI embeddings")
+        
+        # Test with a simple query
+        test_embedding = embedding_model.embed_query("This is a test")
+        print(f"Successfully created test embedding with length {len(test_embedding)}")
+    except Exception as e:
+        print(f"Error initializing OpenAI embeddings: {e}")
+        return
+    
+    # Step 4: Try to initialize Qdrant client
+    try:
+        print("Initializing Qdrant client...")
+        client = QdrantClient(path=str(QDRANT_DIR))
+        
+        # List collections
+        collections = client.get_collections()
+        print(f"Available collections: {collections}")
+        
+        # Get collection info
+        collection_name = "kohavi_ab_testing_pdf_collection"
+        collection_info = client.get_collection(collection_name)
+        print(f"Collection info: {collection_info}")
+        
+        # Test a simple search
+        print("Testing search functionality...")
+        query_embedding = embedding_model.embed_query("What is AB testing?")
+        search_results = client.search(
+            collection_name=collection_name,
+            query_vector=query_embedding,
+            limit=2
+        )
+        print(f"Successfully retrieved {len(search_results)} search results")
+        
+        # Display a sample result
+        if search_results:
+            print("First result ID:", search_results[0].id)
+            print("First result score:", search_results[0].score)
+            print("First result payload text (first 100 chars):", search_results[0].payload["text"][:100])
+        
+    except Exception as e:
+        print(f"Error with Qdrant operations: {e}")
+        return
+    
+    print("Debug process completed successfully!")
+
+if __name__ == "__main__":
+    print("Script execution started")
+    main()
+    print("Script execution finished") 

docker-compose.yml

@@ -10,6 +10,7 @@ services:
       - "8000:8000"
     volumes:
       - ./data:/app/data
+      - ./processed_data:/app/processed_data
       - ./.env:/app/.env
     environment:
       - OPENAI_API_KEY=${OPENAI_API_KEY}

huggingface-space.yml

@@ -1,11 +1,22 @@
 ---
-title: AB Testing RAG Agent
+title: AB Testing RAG
 emoji: 📊
-colorFrom: indigo
-colorTo: blue
+colorFrom: blue
+colorTo: indigo
 sdk: docker
 pinned: false
 license: mit
+
+# Environment variables (fill these in on Hugging Face)
+# OPENAI_API_KEY: Your OpenAI API Key
+
+# Configuration for Dockerfile
+dockerfile:
+  # Use environment variables to specify which app to run
+  env:
+    APP_MODE: streamlit  # Options: streamlit or chainlit
+
+# Streamlit app configuration
 app_port: 8000
 ---
 

minimal_app.py

@@ -0,0 +1,19 @@
+import os
+import chainlit as cl
+from dotenv import load_dotenv
+
+# Load environment variables
+load_dotenv()
+print("Loaded .env file")
+
+@cl.on_chat_start
+async def start():
+    print("Chat started")
+    await cl.Message(content="Hello! I am a minimal AB Testing RAG agent.").send()
+    print("Welcome message sent")
+
+@cl.on_message
+async def main(message):
+    print(f"Received message: {message.content[:30]}...")
+    await cl.Message(content=f"You asked: {message.content}\n\nThis is a minimal test response.").send()
+    print("Response sent") 

requirements.txt

@@ -10,4 +10,6 @@ tiktoken==0.5.2
 python-dotenv==1.0.1
 unstructured==0.12.5
 pypdf==3.17.4
-numpy==1.26.3 
+numpy==1.26.3
+streamlit==1.32.0
+requests==2.31.0 

run_log.txt

@@ -0,0 +1 @@
+2025-04-29 21:21:30 - Loaded .env file

scripts/docker-entrypoint.sh

@@ -0,0 +1,34 @@
+#!/bin/bash
+set -e
+
+# Function to check if pre-processed data exists
+check_preprocessed_data() {
+  if [ -d "/app/processed_data/qdrant_vectorstore" ] && [ -f "/app/processed_data/document_chunks.pkl" ]; then
+    return 0  # Data exists
+  else
+    return 1  # Data doesn't exist
+  fi
+}
+
+# Check if pre-processed data already exists
+if check_preprocessed_data; then
+  echo "Pre-processed data already exists and will be used"
+  echo "No need to process PDF files as processed data is already available"
+else
+  # Count PDF files in data directory
+  pdf_count=$(find /app/data -name "*.pdf" | wc -l)
+  echo "Found $pdf_count PDF files in data directory"
+  
+  # If there are PDF files but no pre-processed data, run pre-processing
+  if [ "$pdf_count" -gt 0 ]; then
+    echo "Pre-processing PDF files..."
+    python /app/scripts/preprocess_data.py
+    echo "Pre-processing complete"
+  else
+    echo "WARNING: No PDF files found in data directory and no pre-processed data exists"
+    echo "The application will likely fail to start"
+  fi
+fi
+
+# Execute the provided command (usually starting the Chainlit app)
+exec "$@" 

scripts/prepare_for_deployment.py

@@ -0,0 +1,161 @@
+#!/usr/bin/env python3
+"""
+Prepares the repository for deployment to Hugging Face Spaces.
+
+This script:
+1. Verifies that pre-processed data exists
+2. Checks that PDF files aren't included in the repository
+3. Lists the files that will be included in the deployment
+4. Provides instructions for deploying to Hugging Face
+
+Run this after successfully pre-processing your data with preprocess_data.py
+"""
+
+import os
+import sys
+from pathlib import Path
+import shutil
+
+PROCESSED_DATA_DIR = Path("processed_data")
+CHUNKS_FILE = PROCESSED_DATA_DIR / "document_chunks.pkl"
+QDRANT_DIR = PROCESSED_DATA_DIR / "qdrant_vectorstore"
+DATA_DIR = Path("data")
+
+def check_preprocessed_data():
+    """Check if pre-processed data exists and is ready for deployment."""
+    print("\n=== Checking for pre-processed data ===")
+    
+    if not PROCESSED_DATA_DIR.exists():
+        print(f"❌ ERROR: Processed data directory not found: {PROCESSED_DATA_DIR}")
+        print("   Please run scripts/preprocess_data.py first.")
+        return False
+    
+    if not CHUNKS_FILE.exists():
+        print(f"❌ ERROR: Document chunks file not found: {CHUNKS_FILE}")
+        print("   Please run scripts/preprocess_data.py first.")
+        return False
+    
+    if not QDRANT_DIR.exists():
+        print(f"❌ ERROR: Vector store directory not found: {QDRANT_DIR}")
+        print("   Please run scripts/preprocess_data.py first.")
+        return False
+    
+    print(f"✅ Found processed data directory: {PROCESSED_DATA_DIR}")
+    print(f"✅ Found document chunks file: {CHUNKS_FILE}")
+    print(f"✅ Found vector store directory: {QDRANT_DIR}")
+    
+    # Check if vector store actually has content
+    qdrant_files = list(QDRANT_DIR.glob("**/*"))
+    if len(qdrant_files) < 5:  # Arbitrary threshold for a minimum number of files
+        print(f"⚠️  WARNING: Vector store directory might be empty or incomplete.")
+        print(f"   Only found {len(qdrant_files)} files in {QDRANT_DIR}")
+    else:
+        print(f"✅ Vector store directory contains {len(qdrant_files)} files/directories")
+    
+    return True
+
+def check_for_pdf_files():
+    """Check that PDF files aren't included in the data directory."""
+    print("\n=== Checking for PDF files ===")
+    
+    pdf_files = list(DATA_DIR.glob("**/*.pdf"))
+    
+    if pdf_files:
+        print(f"⚠️  WARNING: Found {len(pdf_files)} PDF files in the data directory.")
+        print("   These files will NOT be committed to the repository if you follow the instructions below.")
+        print("   PDFs in the data directory are excluded in .gitignore.")
+        for pdf in pdf_files[:5]:  # Show first 5 PDFs only
+            print(f"   - {pdf}")
+        if len(pdf_files) > 5:
+            print(f"   - ... and {len(pdf_files) - 5} more")
+    else:
+        print("✅ No PDF files found in the data directory - good!")
+    
+    return True
+
+def list_deployment_files():
+    """List the essential files that will be included in the deployment."""
+    print("\n=== Files to include in deployment ===")
+    
+    essential_files = [
+        "app.py",
+        "requirements.txt",
+        "Dockerfile",
+        "docker-compose.yml",
+        "chainlit.md",
+        "chainlit.json",
+        "README.md",
+        "scripts/docker-entrypoint.sh",
+        ".dockerignore",
+        "processed_data/",
+    ]
+    
+    print("The following files and directories should be included in your deployment:")
+    for file in essential_files:
+        file_path = Path(file)
+        if file_path.exists() or (file.endswith('/') and Path(file.rstrip('/')).exists()):
+            print(f"✅ {file}")
+        else:
+            print(f"❌ {file} - Not found!")
+    
+    return True
+
+def provide_deployment_instructions():
+    """Provide instructions for deploying to Hugging Face."""
+    print("\n=== Deployment Instructions ===")
+    
+    print("""
+To deploy to Hugging Face Spaces:
+
+1. Ensure you have the Hugging Face CLI installed:
+   pip install huggingface_hub
+
+2. Log in to Hugging Face:
+   huggingface-cli login
+
+3. Create a new Hugging Face Space with Docker deployment from the Hugging Face website
+
+4. Add your repository as a remote:
+   git remote add huggingface https://huggingface.co/spaces/YOUR_USERNAME/YOUR_SPACE_NAME
+
+5. Stage the necessary files (do NOT include PDF files):
+   git add app.py requirements.txt Dockerfile docker-compose.yml chainlit.md chainlit.json README.md scripts/docker-entrypoint.sh .dockerignore processed_data/
+
+6. Commit the changes:
+   git commit -m "Prepare for Hugging Face deployment with pre-processed data"
+
+7. Push to Hugging Face:
+   git push huggingface main
+
+8. Set your OpenAI API key in the Hugging Face Space settings
+""")
+    
+    return True
+
+def main():
+    """Main entry point of the script."""
+    print("=" * 80)
+    print("PREPARING FOR HUGGING FACE DEPLOYMENT")
+    print("=" * 80)
+    
+    checks = [
+        check_preprocessed_data,
+        check_for_pdf_files,
+        list_deployment_files,
+        provide_deployment_instructions
+    ]
+    
+    all_passed = True
+    for check in checks:
+        if not check():
+            all_passed = False
+    
+    if all_passed:
+        print("\n✅ All checks passed! Your repository is ready for deployment to Hugging Face Spaces.")
+        print("   Follow the deployment instructions above to deploy your application.")
+    else:
+        print("\n❌ Some checks failed. Please fix the issues before deploying.")
+        sys.exit(1)
+
+if __name__ == "__main__":
+    main() 

scripts/preprocess_data.py

@@ -0,0 +1,247 @@
+#!/usr/bin/env python3
+"""
+Pre-processes PDF files in the data directory, creating and saving:
+1. Document chunks with metadata
+2. Vector embeddings
+
+This allows the app to load pre-processed data directly instead of processing
+PDFs at runtime, making the app start faster and eliminating the need to
+upload PDFs to Hugging Face.
+"""
+
+import os
+import pickle
+import tiktoken
+from pathlib import Path
+from collections import defaultdict
+import json
+
+# Import required LangChain modules
+from langchain_community.document_loaders import DirectoryLoader
+from langchain_community.document_loaders import PyPDFLoader
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain_openai.embeddings import OpenAIEmbeddings
+from langchain_qdrant import Qdrant
+from langchain_core.documents import Document
+from qdrant_client import QdrantClient
+from qdrant_client.models import Distance, VectorParams, PointStruct
+
+# Set OpenAI API key
+from dotenv import load_dotenv
+load_dotenv()
+
+# Ensure OpenAI API key is available
+if not os.environ.get("OPENAI_API_KEY"):
+    raise EnvironmentError("OPENAI_API_KEY environment variable not found. Please set it in your .env file.")
+
+# Create directories to store pre-processed data
+PROCESSED_DATA_DIR = Path("processed_data")
+PROCESSED_DATA_DIR.mkdir(exist_ok=True)
+
+CHUNKS_FILE = PROCESSED_DATA_DIR / "document_chunks.pkl"
+QDRANT_DIR = PROCESSED_DATA_DIR / "qdrant_vectorstore"
+
+def load_and_process_documents():
+    """
+    Load PDF documents, merge them by source, and split into chunks.
+    This is the same process used in the notebook and app.py.
+    """
+    print("Loading PDF documents...")
+    path = "data/"
+    loader = DirectoryLoader(path, glob="*.pdf", loader_cls=PyPDFLoader)
+    all_docs = loader.load()
+    print(f"Loaded {len(all_docs)} pages from PDF files")
+    
+    # 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)
+    
+    print(f"Created {len(merged_docs)} merged documents")
+    
+    # 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
+    print("Splitting documents into chunks...")
+    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)
+    print(f"Created {len(split_chunks)} document chunks")
+    
+    return split_chunks
+
+def create_and_save_vectorstore(chunks):
+    """
+    Create a vector store from document chunks and save it to disk.
+    """
+    print("Creating embeddings and vector store...")
+    embedding_model = OpenAIEmbeddings(model="text-embedding-3-small")
+    
+    # Extract text and metadata for separate processing
+    texts = [doc.page_content for doc in chunks]
+    metadatas = [doc.metadata for doc in chunks]
+    
+    # Ensure the directory exists
+    QDRANT_DIR.mkdir(exist_ok=True, parents=True)
+
+    # Create a local Qdrant client
+    client = QdrantClient(path=str(QDRANT_DIR))
+    
+    # Get the embedding dimension
+    sample_embedding = embedding_model.embed_query("Sample text")
+    
+    # Create the collection if it doesn't exist
+    collection_name = "kohavi_ab_testing_pdf_collection"
+    try:
+        collection_info = client.get_collection(collection_name)
+        print(f"Collection {collection_name} already exists")
+    except Exception:
+        # Collection doesn't exist, create it
+        print(f"Creating collection {collection_name}")
+        client.create_collection(
+            collection_name=collection_name,
+            vectors_config=VectorParams(
+                size=len(sample_embedding),
+                distance=Distance.COSINE
+            )
+        )
+    
+    # Process in batches to avoid memory issues
+    batch_size = 100
+    print(f"Processing {len(texts)} documents in batches of {batch_size}")
+    
+    for i in range(0, len(texts), batch_size):
+        batch_texts = texts[i:i+batch_size]
+        batch_metadatas = metadatas[i:i+batch_size]
+        
+        print(f"Processing batch {i//batch_size + 1}/{(len(texts) + batch_size - 1)//batch_size}")
+        
+        # Get embeddings for this batch
+        embeddings = embedding_model.embed_documents(batch_texts)
+        
+        # Create points for this batch
+        points = []
+        for j, (text, embedding, metadata) in enumerate(zip(batch_texts, embeddings, batch_metadatas)):
+            points.append(PointStruct(
+                id=i + j,
+                vector=embedding,
+                payload={
+                    "text": text,
+                    "metadata": metadata
+                }
+            ))
+        
+        # Upsert points into the collection
+        client.upsert(
+            collection_name=collection_name,
+            points=points
+        )
+    
+    print(f"Vector store created and saved to {QDRANT_DIR}")
+    return True
+
+def main():
+    # Load and process documents
+    print("Starting pre-processing of PDF files...")
+    chunks = load_and_process_documents()
+    
+    # Save chunks to disk
+    print(f"Saving {len(chunks)} document chunks to {CHUNKS_FILE}...")
+    with open(CHUNKS_FILE, 'wb') as f:
+        pickle.dump(chunks, f)
+    print(f"Chunks saved to {CHUNKS_FILE}")
+    
+    # Create and save vector store
+    success = create_and_save_vectorstore(chunks)
+    
+    if success:
+        print("Pre-processing complete! The application can now use these pre-processed files.")
+        print(f"- Document chunks: {CHUNKS_FILE}")
+        print(f"- Vector store: {QDRANT_DIR}")
+    else:
+        print("Error creating vector store. Please check the logs.")
+
+if __name__ == "__main__":
+    main() 

streamlit_app.py

@@ -0,0 +1,476 @@
+import os
+import pickle
+import streamlit as st
+from pathlib import Path
+from dotenv import load_dotenv
+from langchain_openai.chat_models import ChatOpenAI
+from langchain_openai.embeddings import OpenAIEmbeddings
+from langchain_core.prompts import ChatPromptTemplate
+from qdrant_client import QdrantClient
+from langchain_core.documents import Document
+from langchain.agents import AgentExecutor, create_openai_tools_agent
+from langchain_core.tools import tool
+from langchain.agents.format_scratchpad.openai_tools import format_to_openai_tool_messages
+from langchain_core.messages import AIMessage, HumanMessage
+import requests
+import json
+from langchain_core.output_parsers import StrOutputParser
+
+# Global variable to store ArXiv sources
+ARXIV_SOURCES = []
+
+# Load environment variables
+load_dotenv()
+print("Loaded .env file")
+
+# 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", "")
+
+# Paths to pre-processed data
+PROCESSED_DATA_DIR = Path("processed_data")
+CHUNKS_FILE = PROCESSED_DATA_DIR / "document_chunks.pkl"
+QDRANT_DIR = PROCESSED_DATA_DIR / "qdrant_vectorstore"
+
+# Define prompts
+INITIAL_RAG_PROMPT = """
+CONTEXT:
+{context}
+
+QUERY:
+{question}
+
+You are a helpful assistant with expertise in AB Testing. 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, but I can try a different approach if you'd like."
+"""
+
+EVALUATE_RESPONSE_PROMPT = """
+QUERY:
+{question}
+
+RESPONSE:
+{response}
+
+Evaluate if the response sufficiently answers the query based on the following criteria:
+1. Relevance: Does the response directly address the query topic?
+2. Completeness: Does the response fully answer all aspects of the query?
+3. Accuracy: Is the information provided factually correct and helpful?
+
+Return only "SUFFICIENT" if the response meets all criteria, or "INSUFFICIENT" if the response needs improvement.
+"""
+
+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.
+"""
+
+AGENT_PROMPT = """
+You are an expert AB Testing assistant. Your job is to provide helpful, accurate information about AB Testing topics. 
+
+You have access to several tools:
+1. You can search for relevant documents in the database using search_documents - use this for general AB testing questions
+2. You can rephrase a query to get better search results using search_with_rephrased_query - use this when initial searches don't yield good results
+3. You can search ArXiv for academic papers using search_arxiv - use this for:
+   a) Specific academic papers, their authors, or publications
+   b) As a fallback when other tools don't yield satisfactory results
+   c) Technical questions that might be better answered with academic research
+
+When the user asks about specific papers, authors of papers, or academic publications, you should IMMEDIATELY use the search_arxiv tool rather than the document search tools.
+
+For general AB testing questions, follow this process:
+1. First try search_documents
+2. If that doesn't provide good information, try search_with_rephrased_query
+3. If still insufficient, try search_arxiv as a final resource before giving up
+
+Use these tools to provide the best possible answer.
+"""
+
+@st.cache_resource
+def load_document_chunks():
+    """Load pre-processed document chunks from disk."""
+    with open(CHUNKS_FILE, 'rb') as f:
+        return pickle.load(f)
+
+@st.cache_resource
+def get_chat_model():
+    """Get the chat model for initial RAG."""
+    return ChatOpenAI(
+        model="gpt-4.1-mini",
+        temperature=0,
+    )
+
+@st.cache_resource
+def get_agent_model():
+    """Get the more powerful model for agent and evaluation."""
+    return ChatOpenAI(
+        model="gpt-4.1",
+        temperature=0,
+    )
+
+@st.cache_resource
+def get_embedding_model():
+    """Get the embedding model."""
+    return OpenAIEmbeddings(model="text-embedding-3-small")
+
+@st.cache_resource
+def setup_qdrant_client():
+    """Set up the Qdrant client."""
+    return QdrantClient(path=str(QDRANT_DIR))
+
+def retrieve_documents(query, k=5):
+    """Retrieve relevant documents for a query."""
+    # Get models and data
+    embedding_model = get_embedding_model()
+    chunks = load_document_chunks()
+    client = setup_qdrant_client()
+    
+    # Create a mapping of IDs to documents
+    docs_by_id = {i: doc for i, doc in enumerate(chunks)}
+    
+    # Get query embedding
+    query_embedding = embedding_model.embed_query(query)
+    
+    # Search Qdrant
+    results = client.search(
+        collection_name="kohavi_ab_testing_pdf_collection",
+        query_vector=query_embedding,
+        limit=k
+    )
+    
+    # Convert results to documents
+    documents = []
+    sources_dict = {}  # Use a dictionary to track unique sources by file+page
+    
+    print(f"Retrieved {len(results)} search results")
+    
+    for result in results:
+        doc_id = result.id
+        if doc_id in docs_by_id:
+            doc = docs_by_id[doc_id]
+            documents.append(doc)
+            
+            # Debug the metadata
+            print(f"Document metadata: {doc.metadata}")
+            
+            # Extract source info
+            source_path = doc.metadata.get("source", "")
+            filename = source_path.split("/")[-1] if "/" in source_path else source_path
+            
+            # Remove .pdf extension if present
+            if filename.lower().endswith('.pdf'):
+                filename = filename[:-4]
+            
+            # Default to the full filename if we can't extract a title
+            if not filename:
+                filename = "Unknown Source"
+                
+            # Get page number, use a default if not available
+            page = doc.metadata.get("page", "unknown")
+            
+            # All PDF sources in data directory are by Ron Kohavi, so add his name as prefix
+            title = f"Ron Kohavi: {filename}"
+            
+            # Create a unique key for this source based on filename and page
+            source_key = f"{filename}_{page}"
+            
+            # Only add to sources if we haven't seen this exact source (same file, same page) before
+            if source_key not in sources_dict:
+                sources_dict[source_key] = {
+                    "title": title,
+                    "page": page,
+                    "score": float(result.score),
+                    "type": "pdf"
+                }
+                print(f"Added source: {title}, Page: {page}")
+            else:
+                print(f"Skipping duplicate source: {title}, Page: {page}")
+    
+    # Convert the dictionary of unique sources back to a list
+    sources = list(sources_dict.values())
+    
+    print(f"Returning {len(documents)} documents with {len(sources)} unique sources")
+    return documents, sources
+
+def rephrase_query(query):
+    """Rephrase the query to improve retrieval."""
+    chat_model = get_chat_model()
+    prompt = ChatPromptTemplate.from_template(REPHRASE_QUERY_PROMPT)
+    messages = prompt.format_messages(question=query)
+    response = chat_model.invoke(messages)
+    return response.content
+
+def generate_answer(context, question):
+    """Generate an answer using the context and question."""
+    chat_model = get_chat_model()
+    prompt = ChatPromptTemplate.from_template(INITIAL_RAG_PROMPT)
+    messages = prompt.format_messages(context=context, question=question)
+    response = chat_model.invoke(messages)
+    return response.content
+
+def evaluate_response(question, response):
+    """Evaluate if the response is sufficient."""
+    # Check if this is likely a question about academic papers or authors
+    paper_keywords = ["author", "paper", "publication", "published", "journal", "conference", "researchers", "wrote"]
+    
+    is_paper_query = any(keyword in question.lower() for keyword in paper_keywords) and ("'" in question or '"' in question)
+    
+    # If it's a paper query and the response indicates lack of knowledge, mark as insufficient
+    if is_paper_query and ("don't know" in response.lower() or "cannot " in response.lower() or "i don't have" in response.lower()):
+        print("Paper-specific query detected with insufficient answer, sending to agent")
+        return False
+    
+    # Otherwise, use the LLM evaluation
+    agent_model = get_agent_model()
+    prompt = ChatPromptTemplate.from_template(EVALUATE_RESPONSE_PROMPT)
+    messages = prompt.format_messages(question=question, response=response)
+    result = agent_model.invoke(messages)
+    return "SUFFICIENT" in result.content
+
+@tool
+def search_documents(query: str) -> str:
+    """Search for relevant documents in the AB Testing database."""
+    documents, _ = retrieve_documents(query)
+    if not documents:
+        return "No relevant documents found"
+    return "\n\n".join([doc.page_content for doc in documents])
+
+@tool
+def search_with_rephrased_query(query: str) -> str:
+    """Rephrase the query and then search for relevant documents."""
+    rephrased = rephrase_query(query)
+    documents, _ = retrieve_documents(rephrased)
+    if not documents:
+        return "No relevant documents found even with rephrased query"
+    return "\n\n".join([doc.page_content for doc in documents])
+
+@tool
+def search_arxiv(query: str) -> str:
+    """Search ArXiv for academic papers related to the query."""
+    global ARXIV_SOURCES
+    ARXIV_SOURCES = []  # Reset sources for new search
+    
+    try:
+        # Check if the query is looking for a specific paper by title
+        if "paper" in query.lower() and ("title" in query.lower() or "called" in query.lower() or "named" in query.lower() or "'" in query or '"' in query):
+            # Try to extract paper title from quotes if present
+            import re
+            title_match = re.search(r'[\'"]([^\'"]+)[\'"]', query)
+            
+            if title_match:
+                paper_title = title_match.group(1)
+                # Use title-specific search with exact match
+                formatted_query = f'ti:"{paper_title}"'
+            else:
+                # Fall back to general search but optimize for title
+                formatted_query = query.replace(' ', '+')
+                formatted_query = f'all:{formatted_query}'
+        else:
+            # General query
+            formatted_query = query.replace(' ', '+')
+            formatted_query = f'all:{formatted_query}'
+        
+        print(f"Searching ArXiv with query: {formatted_query}")
+        url = f"http://export.arxiv.org/api/query?search_query={formatted_query}&start=0&max_results=5"
+        
+        response = requests.get(url)
+        if response.status_code != 200:
+            return "Error fetching data from ArXiv"
+        
+        # Parse response
+        import xml.etree.ElementTree as ET
+        root = ET.fromstring(response.text)
+        
+        results = []
+        ns = {'atom': 'http://www.w3.org/2005/Atom'}
+        
+        # Count total entries
+        total_entries = len(root.findall('atom:entry', ns))
+        print(f"Found {total_entries} papers on ArXiv")
+        
+        # Clear previous sources and add new ones
+        ARXIV_SOURCES.clear()
+        
+        for entry in root.findall('atom:entry', ns):
+            title = entry.find('atom:title', ns).text
+            authors = [author.find('atom:name', ns).text for author in entry.findall('atom:author', ns)]
+            summary = entry.find('atom:summary', ns).text
+            link = entry.find('atom:id', ns).text
+            
+            # Add to global sources list
+            ARXIV_SOURCES.append({
+                "title": title,
+                "authors": ", ".join(authors),
+                "type": "arxiv"
+            })
+            
+            results.append({
+                "title": title,
+                "authors": ", ".join(authors),
+                "summary": summary,
+                "link": link
+            })
+        
+        if not results:
+            return "No papers found on ArXiv matching the query"
+        
+        # Format results as text
+        text_results = []
+        for i, paper in enumerate(results):
+            text_results.append(f"Paper {i+1}:\nTitle: {paper['title']}\nAuthors: {paper['authors']}\nSummary: {paper['summary'][:300]}...\nLink: {paper['link']}\n")
+        
+        return "\n".join(text_results)
+    except Exception as e:
+        return f"Error searching ArXiv: {str(e)}"
+
+def setup_agent():
+    """Set up the agent with tools."""
+    agent_model = get_agent_model()
+    tools = [search_documents, search_with_rephrased_query, search_arxiv]
+    prompt = ChatPromptTemplate.from_messages([
+        ("system", AGENT_PROMPT),
+        ("human", "{input}"),
+        ("ai", "{agent_scratchpad}")
+    ])
+    
+    agent = create_openai_tools_agent(agent_model, tools, prompt)
+    executor = AgentExecutor(
+        agent=agent, 
+        tools=tools, 
+        verbose=True,
+        handle_parsing_errors=True
+    )
+    
+    return executor
+
+# Streamlit UI
+st.set_page_config(
+    page_title="AB Testing RAG Agent",
+    page_icon="📊",
+    layout="wide"
+)
+
+st.title("📊 AB Testing RAG Agent")
+st.markdown("""
+This specialized agent can answer questions about A/B Testing using a collection of Ron Kohavi's work. If it can't fully answer your A/B Testing questions using this collection, it will then automatically search Arxiv. Let's begin!
+""")
+
+# Initialize chat history
+if "messages" not in st.session_state:
+    st.session_state.messages = []
+
+# Display chat history
+for message in st.session_state.messages:
+    with st.chat_message(message["role"]):
+        st.markdown(message["content"])
+        
+        # Display sources if available
+        if "sources" in message and message["sources"]:
+            st.markdown("#### Sources")
+            for i, source in enumerate(message["sources"]):
+                title = source.get("title", "Unknown")
+                
+                # Display differently based on source type
+                if source.get("type") == "arxiv":
+                    authors = source.get("authors", "Unknown authors")
+                    st.markdown(f"**{i+1}. {title}**\nAuthors: {authors}")
+                else:
+                    # PDF source with page number
+                    page = source.get("page", "Unknown")
+                    st.markdown(f"**{i+1}. {title}** (Page: {page})")
+
+# Input for new question
+query = st.chat_input("Ask a question about A/B Testing")
+
+if query:
+    # Add user message to chat history
+    st.session_state.messages.append({"role": "user", "content": query})
+    
+    # Display user message
+    with st.chat_message("user"):
+        st.markdown(query)
+    
+    # Display assistant response
+    with st.chat_message("assistant"):
+        message_placeholder = st.empty()
+        
+        with st.status("Processing your query...", expanded=True) as status:
+            # Try initial RAG approach first (matching original architecture)
+            st.write("Attempting initial search...")
+            documents, sources = retrieve_documents(query)
+            
+            # If no documents found, try rephrasing right away
+            if not documents:
+                st.write("No relevant documents found, trying with rephrased query...")
+                rephrased_query = rephrase_query(query)
+                st.write(f"Rephrased query: {rephrased_query}")
+                documents, sources = retrieve_documents(rephrased_query)
+            
+            # Format documents into a string for context
+            context = "\n\n".join([doc.page_content for doc in documents])
+            st.write(f"Retrieved {len(documents)} documents")
+            
+            # If we have context, try the initial RAG approach
+            if context:
+                st.write("Generating initial answer...")
+                initial_answer = generate_answer(context, query)
+                
+                # Evaluate if the initial answer is sufficient
+                st.write("Evaluating answer quality...")
+                is_sufficient = evaluate_response(query, initial_answer)
+                
+                if is_sufficient:
+                    # If the initial answer is good, use it
+                    st.write("Initial answer is sufficient")
+                    answer = initial_answer
+                else:
+                    # If not sufficient, use the agent with tools
+                    st.write("Initial answer needs improvement, using advanced tools...")
+                    agent = setup_agent()
+                    agent_response = agent.invoke({"input": query})
+                    answer = agent_response["output"]
+                    
+                    # Check if the agent used ArXiv and has sources
+                    if ARXIV_SOURCES:
+                        st.write(f"Found {len(ARXIV_SOURCES)} ArXiv sources")
+                        # Replace sources with ArXiv sources
+                        sources = ARXIV_SOURCES
+            else:
+                # If no context at all, use the agent directly
+                st.write("No relevant documents found, using advanced tools...")
+                agent = setup_agent()
+                agent_response = agent.invoke({"input": query})
+                answer = agent_response["output"]
+                
+                # Check if ArXiv sources are available
+                if ARXIV_SOURCES:
+                    st.write(f"Found {len(ARXIV_SOURCES)} ArXiv sources")
+                    # Replace sources with ArXiv sources
+                    sources = ARXIV_SOURCES
+            
+            status.update(label="Completed!", state="complete", expanded=False)
+        
+        # Display answer
+        message_placeholder.markdown(answer)
+        
+        # Display sources directly in the message (if available)
+        if sources:
+            st.markdown("#### Sources")
+            for i, source in enumerate(sources):
+                title = source.get("title", "Unknown")
+                
+                # Display differently based on source type
+                if source.get("type") == "arxiv":
+                    authors = source.get("authors", "Unknown authors")
+                    st.markdown(f"**{i+1}. {title}**\nAuthors: {authors}")
+                else:
+                    # PDF source with page number
+                    page = source.get("page", "Unknown")
+                    st.markdown(f"**{i+1}. {title}** (Page: {page})")
+        
+        # Add assistant message to chat history with sources
+        st.session_state.messages.append({
+            "role": "assistant", 
+            "content": answer, 
+            "sources": sources if sources else []
+        })