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):
    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)
    prompt = st.text_input("Enter Your Question:", placeholder="Ask something", key="1")
    if prompt:
        st.info(prompt, icon="🧐")
        result = qa({"question": prompt})
        st.success(result['answer'], icon="🤖")


def summary_strategy(docs):
    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))
    prompt = st.text_input("Enter Your Question:", placeholder="Ask something", key="2")
    if prompt:
        st.info(prompt, icon="🧐")
        result = qa({"question": prompt})
        st.success(result['answer'], icon="🤖")


def hypothetical_questions_strategy(docs):
    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))
    prompt = st.text_input("Enter Your Question:", placeholder="Ask something", key="3")
    if prompt:
        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()