app/app.py

import os
import pickle
import streamlit as st
from pathlib import Path
from typing import Annotated, List, TypedDict, Dict, Any, Union
import operator
import functools
import numpy as np
from scipy.spatial.distance import cosine

from dotenv import load_dotenv
from langchain_core.messages import AIMessage, BaseMessage, HumanMessage
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.tools import tool
from langchain_openai import ChatOpenAI
from langchain_community.tools.arxiv.tool import ArxivQueryRun
from langchain.schema.output_parser import StrOutputParser
from langchain_openai.embeddings import OpenAIEmbeddings
from langchain_core.vectorstores import VectorStore
from langchain_core.documents import Document
from langgraph.graph import StateGraph, END
from langgraph.prebuilt import ToolNode

# Load environment variables
load_dotenv()

# Check for API key
if not os.environ.get("OPENAI_API_KEY"):
    st.error("OpenAI API key not found. Please set the OPENAI_API_KEY environment variable.")
    st.stop()

# Custom vector store implementation
class CustomVectorStore(VectorStore):
    def __init__(self, embedded_docs, embedding_model):
        self.embedded_docs = embedded_docs
        self.embedding_model = embedding_model
    
    def similarity_search_with_score(self, query, k=5):
        # Get the query embedding
        query_embedding = self.embedding_model.embed_query(query)
        
        # Calculate similarity scores
        results = []
        for doc in self.embedded_docs:
            # Calculate cosine similarity (1 - cosine distance)
            similarity = 1 - cosine(query_embedding, doc["embedding"])
            results.append((doc, similarity))
        
        # Sort by similarity score (highest first)
        results.sort(key=lambda x: x[1], reverse=True)
        
        # Convert to Document objects and return top k
        documents_with_scores = []
        for doc, score in results[:k]:
            document = Document(
                page_content=doc["text"],
                metadata=doc["metadata"]
            )
            documents_with_scores.append((document, score))
        
        return documents_with_scores
    
    def similarity_search(self, query, k=5):
        docs_with_scores = self.similarity_search_with_score(query, k)
        return [doc for doc, _ in docs_with_scores]
    
    def as_retriever(self, search_kwargs=None):
        if search_kwargs is None:
            search_kwargs = {"k": 5}
        
        # Create a simple retriever function
        def retriever(query):
            return self.similarity_search(query, k=search_kwargs.get("k", 5))
        
        # Add get_relevant_documents method to make it compatible with langchain
        retriever.get_relevant_documents = retriever
        
        return retriever
    
    @classmethod
    def from_texts(cls, texts, embedding, metadatas=None, **kwargs):
        """Implement required abstract method from VectorStore base class.
        
        This is a stub implementation that won't actually be used in our case.
        """
        # Create embeddings for the texts
        embeddings = embedding.embed_documents(texts)
        
        # Create embedded docs format
        embedded_docs = []
        for i, (text, embedding_vector) in enumerate(zip(texts, embeddings)):
            metadata = metadatas[i] if metadatas else {}
            embedded_docs.append({
                "text": text,
                "embedding": embedding_vector,
                "metadata": metadata
            })
        
        # Return an instance with the embedded docs
        return cls(embedded_docs, embedding)

# Function to find the processed_data directory
def find_processed_data():
    """Find the processed_data directory path"""
    # Check common locations
    possible_paths = [
        "data/processed_data",
        "app/data/processed_data",
        "/data/processed_data"
    ]
    
    for path in possible_paths:
        if os.path.exists(path):
            return path
    
    # Check relative to the current file
    current_dir = Path(__file__).parent
    for path in [current_dir / "../data/processed_data", current_dir / "data/processed_data"]:
        if path.exists():
            return str(path.resolve())
    
    raise FileNotFoundError("Could not find processed_data directory")

# Initialize the vectorstore
@st.cache_resource
def initialize_vectorstore():
    """Initialize the vectorstore from processed data"""
    try:
        processed_data_path = find_processed_data()
        
        # Load chunks for reference
        chunks_path = os.path.join(processed_data_path, "chunks.pkl")
        try:
            with open(chunks_path, "rb") as f:
                chunks = pickle.load(f)
        except Exception as e:
            chunks = []
            raise RuntimeError(f"Error loading chunks.pkl: {str(e)}")
        
        # Load embedded docs
        embedded_docs_path = os.path.join(processed_data_path, "embedded_docs.pkl")
        try:
            with open(embedded_docs_path, "rb") as f:
                embedded_docs = pickle.load(f)
        except Exception as e:
            embedded_docs = []
            raise RuntimeError(f"Error loading embedded_docs.pkl: {str(e)}")
        
        if not chunks or not embedded_docs:
            # Return empty vectorstore as fallback
            embedding_model = OpenAIEmbeddings(model="text-embedding-3-small")
            vectorstore = CustomVectorStore([], embedding_model)
            return vectorstore, []
            
        # Initialize embedding model
        try:
            embedding_model = OpenAIEmbeddings(model="text-embedding-3-small")
        except Exception as e:
            raise RuntimeError(f"Error initializing OpenAI embeddings model: {str(e)}")
        
        # Create custom vectorstore
        vectorstore = CustomVectorStore(embedded_docs, embedding_model)
        
        return vectorstore, chunks
    except Exception as e:
        raise RuntimeError(f"Error in vectorstore initialization: {str(e)}")

# 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.
"""

# Define the GraphState for the LangGraph
class GraphState(TypedDict):
    messages: Annotated[List[BaseMessage], operator.add]
    sources: Annotated[List[Dict[str, Any]], operator.add]  # Track all sources

# Initialize the AB Testing QA system
@st.cache_resource
def initialize_qa_system(_vectorstore):
    """Initialize the AB Testing QA system"""
    # Create a retriever
    retriever = _vectorstore.as_retriever(search_kwargs={"k": 5})
    
    # Initialize prompt templates
    rag_prompt = ChatPromptTemplate.from_template(RAG_PROMPT)
    rephrase_query_prompt = ChatPromptTemplate.from_template(REPHRASE_QUERY_PROMPT)
    
    # Initialize models (with streaming enabled)
    openai_chat_model = ChatOpenAI(model="gpt-4.1-mini", temperature=0, streaming=True)
    
    # Define the RAG chain node
    def rag_chain_node(state: GraphState) -> GraphState:
        query = state["messages"][-1].content
        
        # 1. Retrieve documents once
        docs = retriever(query)
        
        # 2. Extract sources from the documents
        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"),
            })
        
        # 3. Use a simplified RAG chain without retrieval
        # Create context from documents
        context = "\n\n".join([doc.page_content for doc in docs])
        
        # Format the prompt with context and query
        formatted_prompt = rag_prompt.format(context=context, question=query)
        
        # Get a placeholder for streaming
        if "streaming_container" in state:
            streaming_container = state["streaming_container"]
            full_response = []
            
            # Stream the response
            for chunk in openai_chat_model.stream(formatted_prompt):
                content = chunk.content
                full_response.append(content)
                streaming_container.markdown("".join(full_response))
            
            response_text = "".join(full_response)
        else:
            # If no streaming container provided, fall back to non-streaming
            response = openai_chat_model.invoke(formatted_prompt)
            response_text = StrOutputParser().invoke(response)
        
        return {
            "messages": [AIMessage(content=response_text)],
            "sources": sources
        }
    
    # Define the tools
    @tool
    def retrieve_information(
        query: Annotated[str, "query to ask the retrieve information tool"]
        ):
        """Use Retrieval Augmented Generation to retrieve information about AB Testing."""
        # 1. Retrieve documents
        docs = retriever(query)
        
        # 2. Extract and store sources
        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"),
            })
        
        # Store sources for later access
        retrieve_information.last_sources = sources
        
        # 3. Return just the formatted document contents
        formatted_content = "\n\n".join([f"Retrieved Information: {i+1}\n{doc.page_content}" 
                                        for i, doc in enumerate(docs)])
        return formatted_content
    
    @tool
    def retrieve_information_with_rephrased_query(
        query: Annotated[str, "query to be rephrased before asking the retrieve information tool"]
        ):
        """This tool will intelligently rephrase your AB testing query and then will use Retrieval Augmented Generation to retrieve information about the rephrased query."""
        
        # 1. Rephrase the query first
        rephrased_query = rephrase_query_prompt.format(question=query)
        rephrased_query = openai_chat_model.invoke(rephrased_query)
        rephrased_query = StrOutputParser().invoke(rephrased_query)
        
        # 2. Retrieve documents using the rephrased query
        docs = retriever(rephrased_query)
        
        # 3. Extract and store sources
        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"),
            })
        
        # Store sources for later access
        retrieve_information_with_rephrased_query.last_sources = sources
        
        # 4. Return formatted content with rephrased query
        formatted_content = f"Rephrased query: {rephrased_query}\n\n" + "\n\n".join(
            [f"Retrieved Information: {i+1}\n{doc.page_content}" for i, doc in enumerate(docs)]
        )
        return formatted_content
    
    # Create tool belt
    tool_belt = [
        retrieve_information,
        retrieve_information_with_rephrased_query,
        ArxivQueryRun(),
    ]
    
    # Create tool node
    tool_node = ToolNode(tool_belt)
    
    # Setup agent model (with streaming)
    model = ChatOpenAI(model="gpt-4.1", temperature=0, streaming=True)
    model = model.bind_tools(tool_belt)
    
    # Define model calling function
    def call_model(state):
        messages = state["messages"]
        
        # Check if we have a streaming container
        streaming_container = state.get("streaming_container", None)
        
        # For streaming response
        if streaming_container:
            full_response = []
            
            # Stream the response
            for chunk in model.stream(messages):
                if hasattr(chunk, "content") and chunk.content is not None:
                    content = chunk.content
                    full_response.append(content)
                    streaming_container.markdown("".join(full_response))
            
            # Get the final response
            if full_response:
                response = AIMessage(content="".join(full_response))
            else:
                # Fall back to non-streaming if needed
                response = model.invoke(messages)
        else:
            # Non-streaming fallback
            response = model.invoke(messages)
        
        # Extract sources if available by examining the last message
        sources = []
        if len(messages) > 0:
            last_message = messages[-1]
            if hasattr(last_message, 'content'):
                content = last_message.content
                
                # Check for specific patterns in the content
                if isinstance(content, str):
                    if "Rephrased query:" in content and hasattr(retrieve_information_with_rephrased_query, "last_sources"):
                        sources = retrieve_information_with_rephrased_query.last_sources
                    elif "Retrieved Information:" in content and hasattr(retrieve_information, "last_sources"):
                        sources = retrieve_information.last_sources
                    elif "Title:" in content and "Authors:" in content:  # ArxivQueryRun pattern
                        # Extract paper titles and IDs from ArXiv results
                        import re
                        titles = re.findall(r"Title: (.*?)$", content, re.MULTILINE)
                        # Try to extract the arxiv IDs - match both old and new format IDs
                        arxiv_ids = re.findall(r"URL: https://arxiv\.org/abs/([0-9v\.]+)", content) 
                        
                        sources = []
                        for i, title in enumerate(titles):
                            source = {"title": title, "type": "arxiv_paper"}
                            # Add arxiv_id if available
                            if i < len(arxiv_ids):
                                source["arxiv_id"] = arxiv_ids[i]
                            sources.append(source)
        
        # Return both the response and sources
        return {
            "messages": [response],
            "sources": sources
        }
    
    # Define continuation condition
    def should_continue(state):
        last_message = state["messages"][-1]
        
        if last_message.tool_calls:
            return "action"
        
        return "end"
    
    # Define helpfulness check
    def NonAB_Testing_or_helpful_RAG_or_continue(state):
        initial_query = state["messages"][0]
        final_response = state["messages"][-1]
        
        prompt_template = """\
        
        Given an initial query, determine if the initial query is related to AB Testing (even vaguely e.g. statistics, A/B testing, etc.) or not. If not related to AB Testing, return 'Y'. If related to AB Testing, then given the initial query and a final response, determine if the final response is extremely helpful or not. If extremely helpful, return 'Y'. If not extremely helpful, return 'N'.

        Initial Query:
        {initial_query}

        Final Response:
        {final_response}"""
        
        prompt_template = ChatPromptTemplate.from_template(prompt_template)
        
        helpfulness_check_model = ChatOpenAI(model="gpt-4.1", temperature=0)
        
        helpfulness_chain = prompt_template | helpfulness_check_model | StrOutputParser()
        
        helpfulness_response = helpfulness_chain.invoke({
            "initial_query": initial_query.content, 
            "final_response": final_response.content
        })
        
        if "Y" in helpfulness_response:
            return "end"
        else:
            return "agent"
    
    # Create graph
    graph = StateGraph(GraphState)
    
    # Add nodes
    graph.add_node("Initial_RAG", rag_chain_node)
    graph.add_node("agent", call_model)
    graph.add_node("action", tool_node)
    
    # Set entry point
    graph.set_entry_point("Initial_RAG")
    
    # Add edges
    graph.add_conditional_edges(
        "Initial_RAG",
        NonAB_Testing_or_helpful_RAG_or_continue,
        {
            "agent": "agent",
            "end": END
        }
    )
    
    graph.add_conditional_edges(
        "agent",
        should_continue,
        {
            "action": "action",
            "end": END
        }
    )
    
    graph.add_edge("action", "agent")
    
    # Compile graph
    return graph.compile()

# Streamlit interface
st.title("📊 A/B Testing RAG Agent")
st.markdown("""
This specialized agent answers your A/B Testing questions using a thorough collection of Ron Kohavi's work. If it can't answer your questions using this collection, it will then search Arxiv. It has been trained to only answer based on the sources it retrieves. Let's begin!
""")

# Initialize the system
try:
    # Show loading indicator
    loading_placeholder = st.empty()
    with loading_placeholder.container():
        import time
        for dots in [".", "..", "..."]:
            loading_placeholder.text(f"Loading{dots}")
            time.sleep(0.2)
    
    # Initialize components (but hide the details)
    vectorstore, chunks = initialize_vectorstore()
    qa_system = initialize_qa_system(vectorstore)
    
    # Clear loading indicator
    loading_placeholder.empty()
except Exception as e:
    st.error(f"Error initializing the system: {str(e)}")
    st.stop()

# Initialize session state for chat history
if "messages" not in st.session_state:
    st.session_state.messages = []

# Display chat history
for i, message in enumerate(st.session_state.messages):
    if message["role"] == "user":
        st.chat_message("user").write(message["content"])
    else:
        with st.chat_message("assistant"):
            st.write(message["content"])
            

# Chat input
query = st.chat_input("Ask me anything about A/B Testing...")

if query:
    # Display user message
    st.chat_message("user").write(query)
    st.session_state.messages.append({"role": "user", "content": query})
    
    # Process query
    with st.spinner("Thinking..."):
        # Create a placeholder for streaming output
        with st.chat_message("assistant"):
            streaming_container = st.empty()
        
            # Create input state for the graph with streaming container
            input_state = {
                "messages": [HumanMessage(content=query)],
                "sources": [],
                "streaming_container": streaming_container  # Pass the container for streaming
            }
        
            # Execute graph
            result = qa_system.invoke(input_state)
        
            # Extract result
            answer = result["messages"][-1].content
            sources = result["sources"]


            # Process sources to remove duplicates and format properly
            unique_sources = set()
            sources_text = ""
               
            for source in sources:
                if "type" in source and source["type"] == "arxiv_paper":
                    # Extract arXiv ID from Entry ID metadata
                    entry_id = source.get('Entry ID', '')  # This is the key field containing the ID
                    if entry_id:
                        # Extract arXiv ID from format like "http://arxiv.org/abs/2404.19647v1"
                        arxiv_id = entry_id.split('/abs/')[-1].split('v')[0]  # Removes version suffix
                        sources_text += f"- ArXiv paper: [{source['title']}](https://arxiv.org/abs/{arxiv_id})\n"
                    else:
                        sources_text += f"- ArXiv paper: {source['title']}\n"

                else:
                    # Handle retrieval sources (Ron Kohavi's work)
                    # Remove .pdf extension if present
                    title = source['title'].replace('.pdf', '')
                        
                    # Create a unique identifier for this source
                    source_id = f"{title}|{source['page']}"
                        
                    # Only add if not already added
                    if source_id not in unique_sources:
                        unique_sources.add(source_id)
                        sources_text += f"- Ron Kohavi: {title}, page {source['page']}\n"
                

            # Final display with the complete answer and sources
            if "I don't know" in answer:
                answers_and_sources = answer
            else:
                answers_and_sources = answer + "\n\n" + "**Sources:**" + "\n\n" + sources_text
            streaming_container.markdown(answers_and_sources)
        
            
    
    # Save to chat history (still save sources for internal use, even if not displayed)
    st.session_state.messages.append({
        "role": "assistant", 
        "content": answer,
        "sources": sources
    })