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()