Update app.py

app.py

@@ -8,20 +8,17 @@ import io
 from langchain.text_splitter import CharacterTextSplitter
 from langchain.embeddings import HuggingFaceEmbeddings
 from langchain.vectorstores import FAISS
-from langchain.memory import ConversationBufferMemory
-from langchain.chains import ConversationalRetrievalChain
-from langchain.llms import HuggingFaceHub
 
 load_dotenv()
 
 # Initialize session state variables
-if "conversation" not in st.session_state:
-    st.session_state.conversation = None
+if "vectorstore" not in st.session_state:
+    st.session_state.vectorstore = None
 if "chat_history" not in st.session_state:
     st.session_state.chat_history = []
 
 def reset_conversation():
-    st.session_state.conversation = None
+    st.session_state.vectorstore = None
     st.session_state.chat_history = []
 
 def get_pdf_text(pdf_docs):
@@ -47,38 +44,59 @@ def get_vectorstore(text_chunks):
     vectorstore = FAISS.from_texts(texts=text_chunks, embedding=embeddings)
     return vectorstore
 
-def get_conversation_chain(vectorstore):
-    llm = HuggingFaceHub(repo_id="google/flan-t5-xxl", model_kwargs={"temperature":0.5, "max_length":512})
-
-    memory = ConversationBufferMemory(
-        memory_key='chat_history', return_messages=True)
-    conversation_chain = ConversationalRetrievalChain.from_llm(
-        llm=llm,
-        retriever=vectorstore.as_retriever(),
-        memory=memory
-    )
-    return conversation_chain
+def get_together_response(prompt, history):
+    url = "https://api.together.xyz/v1/chat/completions"
+    model_link = "NousResearch/Nous-Hermes-2-Yi-34B"
+
+    messages = [{"role": "system", "content": "You are an AI assistant that helps users understand the content of their PDFs. Provide concise and relevant answers based on the information in the documents."}]
+    
+    for human, ai in history:
+        messages.append({"role": "user", "content": human})
+        messages.append({"role": "assistant", "content": ai})
+    
+    messages.append({"role": "user", "content": prompt})
+
+    payload = {
+        "model": model_link,
+        "messages": messages,
+        "temperature": 0.7,
+        "top_p": 0.95,
+        "top_k": 50,
+        "repetition_penalty": 1,
+        "max_tokens": 1024
+    }
+
+    headers = {
+        "accept": "application/json",
+        "content-type": "application/json",
+        "Authorization": f"Bearer {os.getenv('TOGETHER_API_KEY')}"
+    }
+
+    try:
+        response = requests.post(url, json=payload, headers=headers)
+        response.raise_for_status()
+        return response.json()['choices'][0]['message']['content']
+    except requests.exceptions.RequestException as e:
+        return f"Error: {str(e)}"
 
 def handle_userinput(user_question):
-    response = st.session_state.conversation({'question': user_question})
-    st.session_state.chat_history = response['chat_history']
-
-    for i, message in enumerate(st.session_state.chat_history):
-        if i % 2 == 0:
-            st.write(user_template.replace(
-                "{{MSG}}", message.content), unsafe_allow_html=True)
-        else:
-            st.write(bot_template.replace(
-                "{{MSG}}", message.content), unsafe_allow_html=True)
+    if st.session_state.vectorstore:
+        docs = st.session_state.vectorstore.similarity_search(user_question)
+        context = "\n".join([doc.page_content for doc in docs])
+        prompt = f"Context from PDFs:\n{context}\n\nQuestion: {user_question}\nAnswer:"
+        
+        response = get_together_response(prompt, st.session_state.chat_history)
+        st.session_state.chat_history.append((user_question, response))
+
+        return response
+    else:
+        return "Please upload and process PDF documents first."
 
 # Streamlit application
 st.set_page_config(page_title="Chat with your PDFs", page_icon=":books:")
 
 st.header("Chat with your PDFs :books:")
 
-user_template = '<div style="background-color: #e6f3ff; padding: 10px; border-radius: 5px; margin-bottom: 10px;"><strong>Human:</strong> {{MSG}}</div>'
-bot_template = '<div style="background-color: #f0f0f0; padding: 10px; border-radius: 5px; margin-bottom: 10px;"><strong>AI:</strong> {{MSG}}</div>'
-
 # Sidebar
 with st.sidebar:
     st.subheader("Your documents")
@@ -92,17 +110,26 @@ with st.sidebar:
             text_chunks = get_text_chunks(raw_text)
 
             # Create vector store
-            vectorstore = get_vectorstore(text_chunks)
-
-            # Create conversation chain
-            st.session_state.conversation = get_conversation_chain(vectorstore)
+            st.session_state.vectorstore = get_vectorstore(text_chunks)
+            
+            st.success("PDFs processed successfully!")
 
     st.button('Reset Chat', on_click=reset_conversation)
 
 # Main chat interface
-if st.session_state.conversation is None:
+if st.session_state.vectorstore is None:
     st.write("Please upload PDF documents and click 'Process' to start chatting.")
 else:
     user_question = st.text_input("Ask a question about your documents:")
     if user_question:
-        handle_userinput(user_question)
+        response = handle_userinput(user_question)
+        
+        st.write("Human: " + user_question)
+        st.write("AI: " + response)
+
+# Display chat history
+st.subheader("Chat History")
+for human, ai in st.session_state.chat_history:
+    st.write("Human: " + human)
+    st.write("AI: " + ai)
+    st.write("---")