app.py

import os
import uuid
import tempfile
import streamlit as st
import openai
from langchain.retrievers.multi_vector import MultiVectorRetriever
from langchain.vectorstores import Chroma
from langchain.embeddings import OpenAIEmbeddings
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.storage import InMemoryStore
from langchain.memory import ConversationBufferMemory
from langchain.llms import OpenAI
from langchain.chains import ConversationalRetrievalChain
from langchain.chat_models import ChatOpenAI
from langchain.prompts import ChatPromptTemplate
from langchain.schema.output_parser import StrOutputParser
import uuid
from langchain.schema.document import Document
from langchain.output_parsers.openai_functions import JsonKeyOutputFunctionsParser
from langchain.document_loaders import PyPDFLoader

# Set OpenAI API key
OPENAI_API_KEY = st.secrets["OPENAI_API_KEY"]
if not OPENAI_API_KEY:
    st.error("OPENAI_API_KEY not set in environment variables!")
    raise SystemExit
openai.api_key = OPENAI_API_KEY


def process_pdf(uploaded_file):
    with st.spinner("Processing PDF..."):
        with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as tmp:
            tmp.write(uploaded_file.getvalue())
            tmp_path = tmp.name
            loaders = [PyPDFLoader(tmp_path)]
            docs = []
            for l in loaders:
                docs.extend(l.load())
            text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000)
            docs = text_splitter.split_documents(docs)
    return docs


def smaller_chunks_strategy(docs):
    prompt = st.text_input("Enter Your Question:", placeholder="Ask something", key="1")
    if prompt:
        with st.spinner('Processing with smaller_chunks_strategy'):
            vectorstore = Chroma(
                collection_name="full_documents",
                embedding_function=OpenAIEmbeddings()
            )
            store = InMemoryStore()
            id_key = "doc_id"
            retriever = MultiVectorRetriever(
                vectorstore=vectorstore,
                docstore=store,
                id_key=id_key,
            )
            doc_ids = [str(uuid.uuid4()) for _ in docs]
            child_text_splitter = RecursiveCharacterTextSplitter(chunk_size=400)
            sub_docs = []
            for i, doc in enumerate(docs):
                _id = doc_ids[i]
                _sub_docs = child_text_splitter.split_documents([doc])
                for _doc in _sub_docs:
                    _doc.metadata[id_key] = _id
                sub_docs.extend(_sub_docs)
    
            retriever.vectorstore.add_documents(sub_docs)
            retriever.docstore.mset(list(zip(doc_ids, docs)))
            memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True)
            qa = ConversationalRetrievalChain.from_llm(OpenAI(temperature=0), retriever, memory=memory)
            st.info(prompt, icon="🧐")
            result = qa({"question": prompt})
            st.success(result['answer'], icon="🤖")


def summary_strategy(docs):
    prompt = st.text_input("Enter Your Question:", placeholder="Ask something", key="2")
    if prompt:
        with st.spinner('Processing with summary_strategy'):
            chain = (
                {"doc": lambda x: x.page_content}
                | ChatPromptTemplate.from_template("Summarize the following document:\n\n{doc}")
                | ChatOpenAI(max_retries=0)
                | StrOutputParser()
            )
            summaries = chain.batch(docs, {"max_concurrency": 5})
            vectorstore = Chroma(
                collection_name="summaries",
                embedding_function= OpenAIEmbeddings()
            )
            store = InMemoryStore()
            id_key = "doc_id"
            retriever = MultiVectorRetriever(
                vectorstore=vectorstore,
                docstore=store,
                id_key=id_key,
            )
            doc_ids = [str(uuid.uuid4()) for _ in docs]
            summary_docs = [Document(page_content=s, metadata={id_key: doc_ids[i]}) for i, s in enumerate(summaries)]
            retriever.vectorstore.add_documents(summary_docs)
            retriever.docstore.mset(list(zip(doc_ids, docs)))
            qa = ConversationalRetrievalChain.from_llm(OpenAI(temperature=0), retriever, memory=ConversationBufferMemory(memory_key="chat_history", return_messages=True))
            st.info(prompt, icon="🧐")
            result = qa({"question": prompt})
            st.success(result['answer'], icon="🤖")


def hypothetical_questions_strategy(docs):
    prompt = st.text_input("Enter Your Question:", placeholder="Ask something", key="3")
    if prompt:
        with st.spinner('Processing with hypothetical_questions_strategy'):
            functions = [
                {
                    "name": "hypothetical_questions",
                    "description": "Generate hypothetical questions",
                    "parameters": {
                        "type": "object",
                        "properties": {
                            "questions": {
                                "type": "array",
                                "items": {
                                    "type": "string"
                                },
                            },
                        },
                        "required": ["questions"]
                    }
                }
            ]
            chain = (
                {"doc": lambda x: x.page_content}
                | ChatPromptTemplate.from_template("Generate a list of 3 hypothetical questions that the below document could be used to answer:\n\n{doc}")
                | ChatOpenAI(max_retries=0, model="gpt-4").bind(functions=functions, function_call={"name": "hypothetical_questions"})
                | JsonKeyOutputFunctionsParser(key_name="questions")
            )
            hypothetical_questions = chain.batch(docs, {"max_concurrency": 5})
            vectorstore = Chroma(
                collection_name="hypo-questions",
                embedding_function=OpenAIEmbeddings()
            )
            store = InMemoryStore()
            id_key = "doc_id"
            retriever = MultiVectorRetriever(
                vectorstore=vectorstore,
                docstore=store,
                id_key=id_key,
            )
            doc_ids = [str(uuid.uuid4()) for _ in docs]
            question_docs = []
            for i, question_list in enumerate(hypothetical_questions):
                question_docs.extend([Document(page_content=s, metadata={id_key: doc_ids[i]}) for s in question_list])
            retriever.vectorstore.add_documents(question_docs)
            retriever.docstore.mset(list(zip(doc_ids, docs)))
            qa = ConversationalRetrievalChain.from_llm(OpenAI(temperature=0), retriever, memory=ConversationBufferMemory(memory_key="chat_history", return_messages=True))
            st.info(prompt, icon="🧐")
            result = qa({"question": prompt})
            st.success(result['answer'], icon="🤖")



def app():
    image_path = "icon.png"
    st.sidebar.image(image_path, caption="icon", use_column_width=True)
    st.title("VecDBCompare 0.0.1")
    st.sidebar.markdown("""
        # 🚀 **VecDBCompare: Your Vector DB Strategy Tester**
        ## 📌 **What is it?**
        VecDBCompare lets you evaluate and compare three vector database retrieval strategies in a snap!
        ## 📤 **How to Use?**
        1. **Upload a PDF** 📄
        2. Get **Three QABots** 🤖🤖🤖, each with a different strategy.
        3. **Ask questions** ❓ and see how each bot responds differently.
        4. **Decide** ✅ which strategy works best for you!
        ## 🌟 **Why VecDBCompare?**
        - **Simple & Fast** ⚡: Upload, ask, and compare!
        - **Real-time Comparison** 🔍: See strategies in action side-by-side.
        - **Empower Your Choice** 💡: Pick the best strategy for your needs.
        Dive in and discover with VecDBCompare! 🌐
    """)
    uploaded_file = st.file_uploader("Choose a PDF file", type=["pdf"])
    if uploaded_file:
        docs = process_pdf(uploaded_file)
        option = st.selectbox(
            "Which retrieval strategy would you like to use?",
            ("Smaller Chunks", "Summary", "Hypothetical Questions")
        )
        if option == 'Smaller Chunks':
            smaller_chunks_strategy(docs)
        elif option == 'Summary':
            summary_strategy(docs)
        elif option == 'Hypothetical Questions':
            hypothetical_questions_strategy(docs)


if __name__ == "__main__":
    st.set_page_config(layout="wide")
    app()