from openai import OpenAI import streamlit as st from dotenv import load_dotenv import os import shelve from langchain_openai import OpenAIEmbeddings, ChatOpenAI from langchain_core.prompts import PromptTemplate from langchain.chains import RetrievalQA from ragatouille import RAGPretrainedModel from langchain.document_loaders import PyPDFLoader RAG = RAGPretrainedModel.from_pretrained("colbert-ir/colbertv2.0") load_dotenv() file_path = "policy_data.pdf" loader = PyPDFLoader(file_path) pages = loader.load() full_document = "" for page in pages: full_document += page.page_content index_dir = ".ragatouille/colbert/indexes/policy_data/" ivf_path = os.path.join(index_dir, "ivf.pid.pt") index_exists = os.path.exists(ivf_path) if not index_exists: st.write("Indexing document...") loader = PyPDFLoader(file_path) pages = loader.load() full_document = "".join(page.page_content for page in pages) RAG.index( collection=[full_document], index_name="policy_data", max_document_length=512, split_documents=True, ) retriever = RAG.as_langchain_retriever(k=3) st.write("Indexing complete!") else: st.write("Loading existing index...") RAG = RAGPretrainedModel.from_index(".ragatouille/colbert/indexes/policy_data/") retriever = RAG.as_langchain_retriever(k=3) template = """Use the context below to answer the question. Keep the answer concise and to the point. If you are unsure about the answer, just say i do not know the answer to the question do not create your own answer and make sure the answer is concise and to the point. Summarize the information such that main points are covered and if you think that there needs to be some more information added to the answer then you can add that information as well. {context} Question: {question} Helpful Answer:""" prompt = PromptTemplate.from_template(template) chain_type_kwargs = {"prompt": prompt} st.title("Streamlit Chatbot Interface") USER_AVATAR = "👤" BOT_AVATAR = "🤖" llm = ChatOpenAI(model="gpt-3.5-turbo", temperature=0.4, max_tokens=500, streaming=True) chain = RetrievalQA.from_chain_type( llm = llm, chain_type="stuff", retriever=retriever, chain_type_kwargs=chain_type_kwargs, ) # Ensure openai_model is initialized in session state if "openai_model" not in st.session_state: st.session_state["openai_model"] = "gpt-3.5-turbo" # Load chat history from shelve file def load_chat_history(): with shelve.open("chat_history") as db: return db.get("messages", []) # Save chat history to shelve file def save_chat_history(messages): with shelve.open("chat_history") as db: db["messages"] = messages # Initialize or load chat history if "messages" not in st.session_state: st.session_state.messages = load_chat_history() # Sidebar with a button to delete chat history with st.sidebar: if st.button("Delete Chat History"): st.session_state.messages = [] save_chat_history([]) # Display chat messages for message in st.session_state.messages: avatar = USER_AVATAR if message["role"] == "user" else BOT_AVATAR with st.chat_message(message["role"], avatar=avatar): st.markdown(message["content"]) # Main chat interface if prompt := st.chat_input("How can I help?"): st.session_state.messages.append({"role": "user", "content": prompt}) with st.chat_message("user", avatar=USER_AVATAR): st.markdown(prompt) with st.chat_message("assistant", avatar=BOT_AVATAR): message_placeholder = st.empty() full_response = chain.invoke(prompt)['result'] message_placeholder.markdown(chain.invoke(prompt)['result']) st.session_state.messages.append({"role": "assistant", "content": full_response}) save_chat_history(st.session_state.messages)