Update app.py

app.py

@@ -1,67 +1,128 @@
-import os
-import tempfile
+from langchain.chains import RetrievalQA
+from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
+from langchain.callbacks.manager import CallbackManager
+from langchain_community.llms import Ollama
+from langchain_community.embeddings.ollama import OllamaEmbeddings
+from langchain_community.vectorstores import Chroma
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain_community.document_loaders import PyPDFLoader
+from langchain.prompts import PromptTemplate
+from langchain.memory import ConversationBufferMemory
 import streamlit as st
-from streamlit_chat import message
-
-from rag import ChatPDF
-
-st.set_page_config(page_title="ChatPDF")
-
-
-def display_messages():
-    st.subheader("Chat")
-    for i, (msg, is_user) in enumerate(st.session_state["messages"]):
-        message(msg, is_user=is_user, key=str(i))
-    st.session_state["thinking_spinner"] = st.empty()
-
-
-def process_input():
-    if st.session_state["user_input"] and len(st.session_state["user_input"].strip()) > 0:
-        user_text = st.session_state["user_input"].strip()
-        with st.session_state["thinking_spinner"], st.spinner(f"Thinking"):
-            agent_text = st.session_state["assistant"].ask(user_text)
-
-        st.session_state["messages"].append((user_text, True))
-        st.session_state["messages"].append((agent_text, False))
-
-
-def read_and_save_file():
-    st.session_state["assistant"].clear()
-    st.session_state["messages"] = []
-    st.session_state["user_input"] = ""
-
-    for file in st.session_state["file_uploader"]:
-        with tempfile.NamedTemporaryFile(delete=False) as tf:
-            tf.write(file.getbuffer())
-            file_path = tf.name
+import os
+import time
 
-        with st.session_state["ingestion_spinner"], st.spinner(f"Ingesting {file.name}"):
-            st.session_state["assistant"].ingest(file_path)
-        os.remove(file_path)
+if not os.path.exists('files'):
+    os.mkdir('files')
 
+if not os.path.exists('jj'):
+    os.mkdir('jj')
 
-def page():
-    if len(st.session_state) == 0:
-        st.session_state["messages"] = []
-        st.session_state["assistant"] = ChatPDF()
+if 'template' not in st.session_state:
+    st.session_state.template = """You are a knowledgeable chatbot, here to help with questions of the user. Your tone should be professional and informative.
 
-    st.header("ChatPDF")
+    Context: {context}
+    History: {history}
 
-    st.subheader("Upload a document")
-    st.file_uploader(
-        "Upload document",
-        type=["pdf"],
-        key="file_uploader",
-        on_change=read_and_save_file,
-        label_visibility="collapsed",
-        accept_multiple_files=True,
+    User: {question}
+    Chatbot:"""
+if 'prompt' not in st.session_state:
+    st.session_state.prompt = PromptTemplate(
+        input_variables=["history", "context", "question"],
+        template=st.session_state.template,
     )
-
-    st.session_state["ingestion_spinner"] = st.empty()
-
-    display_messages()
-    st.text_input("Message", key="user_input", on_change=process_input)
-
-
-if __name__ == "__main__":
-    page()
+if 'memory' not in st.session_state:
+    st.session_state.memory = ConversationBufferMemory(
+        memory_key="history",
+        return_messages=True,
+        input_key="question")
+if 'vectorstore' not in st.session_state:
+    st.session_state.vectorstore = Chroma(persist_directory='jj',
+                                          embedding_function=OllamaEmbeddings(base_url='http://localhost:11434',
+                                                                              model="mistral")
+                                          )
+if 'llm' not in st.session_state:
+    st.session_state.llm = Ollama(base_url="http://localhost:11434",
+                                  model="mistral",
+                                  verbose=True,
+                                  callback_manager=CallbackManager(
+                                      [StreamingStdOutCallbackHandler()]),
+                                  )
+
+# Initialize session state
+if 'chat_history' not in st.session_state:
+    st.session_state.chat_history = []
+
+st.title("PDF Chatbot")
+
+# Upload a PDF file
+uploaded_file = st.file_uploader("Upload your PDF", type='pdf')
+
+for message in st.session_state.chat_history:
+    with st.chat_message(message["role"]):
+        st.markdown(message["message"])
+
+if uploaded_file is not None:
+    if not os.path.isfile("files/"+uploaded_file.name+".pdf"):
+        with st.status("Analyzing your document..."):
+            bytes_data = uploaded_file.read()
+            f = open("files/"+uploaded_file.name+".pdf", "wb")
+            f.write(bytes_data)
+            f.close()
+            loader = PyPDFLoader("files/"+uploaded_file.name+".pdf")
+            data = loader.load()
+
+            # Initialize text splitter
+            text_splitter = RecursiveCharacterTextSplitter(
+                chunk_size=1500,
+                chunk_overlap=200,
+                length_function=len
+            )
+            all_splits = text_splitter.split_documents(data)
+
+            # Create and persist the vector store
+            st.session_state.vectorstore = Chroma.from_documents(
+                documents=all_splits,
+                embedding=OllamaEmbeddings(model="mistral")
+            )
+            st.session_state.vectorstore.persist()
+
+    st.session_state.retriever = st.session_state.vectorstore.as_retriever()
+    # Initialize the QA chain
+    if 'qa_chain' not in st.session_state:
+        st.session_state.qa_chain = RetrievalQA.from_chain_type(
+            llm=st.session_state.llm,
+            chain_type='stuff',
+            retriever=st.session_state.retriever,
+            verbose=True,
+            chain_type_kwargs={
+                "verbose": True,
+                "prompt": st.session_state.prompt,
+                "memory": st.session_state.memory,
+            }
+        )
+
+    # Chat input
+    if user_input := st.chat_input("You:", key="user_input"):
+        user_message = {"role": "user", "message": user_input}
+        st.session_state.chat_history.append(user_message)
+        with st.chat_message("user"):
+            st.markdown(user_input)
+        with st.chat_message("assistant"):
+            with st.spinner("Assistant is typing..."):
+                response = st.session_state.qa_chain(user_input)
+            message_placeholder = st.empty()
+            full_response = ""
+            for chunk in response['result'].split():
+                full_response += chunk + " "
+                time.sleep(0.05)
+                # Add a blinking cursor to simulate typing
+                message_placeholder.markdown(full_response + "▌")
+            message_placeholder.markdown(full_response)
+
+        chatbot_message = {"role": "assistant", "message": response['result']}
+        st.session_state.chat_history.append(chatbot_message)
+
+
+else:
+    st.write("Please upload a PDF file.")