Add support for chat

Author: punitarani

app.py

@@ -1,7 +1,9 @@
 """Streamlit App"""
 
+import openai
 import streamlit as st
 from langchain.schema import Document
+from streamlit_chat import message
 
 from rabbithole import summarize_document
 from rabbithole.embedding import embed_document
@@ -14,6 +16,41 @@
 for state_var in ["uploaded_files", "documents", "embeddings", "keywords", "summaries"]:
     if state_var not in st.session_state:
         st.session_state[state_var] = {}
+if "processed" not in st.session_state:
+    st.session_state.processed = False
+if "bot_messages" not in st.session_state:
+    st.session_state.bot_messages = [
+        "Hello, I am here to help you learn more efficiently"
+    ]
+if "user_messages" not in st.session_state:
+    st.session_state.user_messages = []
+
+
+def generate_response(prompt):
+    """Generate a response to a prompt using the GPT-3.5 Turbo model"""
+    st.session_state['user_messages'].append(prompt)
+
+    messages = []
+    # Alternate between bot and the assistant until the conversation is over
+    msg_idx = 0
+    while True:
+        if len(st.session_state['bot_messages']) > msg_idx:
+            messages.append({"role": "assistant", "content": st.session_state['bot_messages'][msg_idx]})
+        else:
+            break
+        if len(st.session_state['user_message']) > msg_idx:
+            messages.append({"role": "user", "content": st.session_state['user_messages'][msg_idx]})
+        else:
+            break
+        msg_idx += 1
+    print(messages)
+    completion = openai.ChatCompletion.create(
+        model="gpt-4",
+        messages=messages
+    )
+    response = completion.choices[0].message.content
+    st.session_state['bot_messages'].append(response)
+    print(response)
 
 
 def load_files_with_spinner(files: list) -> dict[str, list[Document]]:
@@ -87,56 +124,82 @@ def generate_plan_with_spinner() -> dict:
 
 st.title("RabbitHole")
 
-uploaded_files = st.file_uploader("Upload content",
-                                  type=["docx", "pdf", "txt", *SUPPORTED_IMG_FILE_TYPES, *SUPPORTED_AV_FILE_TYPES],
-                                  accept_multiple_files=True)
-
-if st.button("Dive in"):
-    if not uploaded_files:
-        st.warning("Please upload a file first.")
-        st.stop()
-
-    # Check if uploaded files have changed
-    uploaded_files_changed = False
-    if len(uploaded_files) != len(st.session_state.uploaded_files):
-        uploaded_files_changed = True
-    else:
-        for new_file, old_file in zip(uploaded_files, st.session_state.uploaded_files):
-            if new_file != old_file:
-                uploaded_files_changed = True
-                break
-
-    if uploaded_files_changed:
-        st.session_state.uploaded_files = uploaded_files
-
-        # Load the text from the uploaded PDF files
-        st.session_state.documents = load_files_with_spinner(st.session_state.uploaded_files)
-        st.session_state.embeddings = embed_documents_with_spinner(st.session_state.documents)
-        st.session_state.keywords = extract_keywords_with_spinner(st.session_state.embeddings)
-        st.session_state.summaries = generate_summary_with_spinner(st.session_state.documents)
-
-    # Display the keywords and summaries
-    for doc_name, doc_keywords in st.session_state.keywords.items():
-        st.header(doc_name)
-        st.caption("Keywords: " + ", ".join(doc_keywords))
-        st.write(st.session_state.summaries[doc_name])
-        st.divider()
-
-    # Display the plan
-    st.header("Study Plan")
-    plan = generate_plan_with_spinner()
-    for data in plan.get("plan", []):
-        for doc_name, doc_data in data.items():
-            st.subheader(doc_name)
-            st.write(f"**Background Concepts**")
-            for concept in doc_data.get("Background Concepts", []):
-                st.write(f"- {concept}")
-            st.write(f"**Key Concepts**")
-            for concept in doc_data.get("Key Concepts", []):
-                st.write(f"- {concept}")
-            st.write(f"**Further Reading**")
-            for concept in doc_data.get("Further Reading", []):
-                st.write(f"- {concept}")
-        st.write("")
-
-    st.success('Summarization completed.')
+if not st.session_state.processed:
+    uploaded_files = st.file_uploader("Upload content",
+                                      type=["docx", "pdf", "txt", *SUPPORTED_IMG_FILE_TYPES, *SUPPORTED_AV_FILE_TYPES],
+                                      accept_multiple_files=True)
+
+    if st.button("Dive in"):
+        if not uploaded_files:
+            st.warning("Please upload a file first.")
+            st.stop()
+
+        # Check if uploaded files have changed
+        uploaded_files_changed = False
+        if len(uploaded_files) != len(st.session_state.uploaded_files):
+            uploaded_files_changed = True
+        else:
+            for new_file, old_file in zip(uploaded_files, st.session_state.uploaded_files):
+                if new_file != old_file:
+                    uploaded_files_changed = True
+                    break
+
+        if uploaded_files_changed:
+            st.session_state.uploaded_files = uploaded_files
+
+            # Load the text from the uploaded PDF files
+            st.session_state.documents = load_files_with_spinner(st.session_state.uploaded_files)
+            st.session_state.embeddings = embed_documents_with_spinner(st.session_state.documents)
+            st.session_state.keywords = extract_keywords_with_spinner(st.session_state.embeddings)
+            st.session_state.summaries = generate_summary_with_spinner(st.session_state.documents)
+
+        # Display the keywords and summaries
+        for doc_name, doc_keywords in st.session_state.keywords.items():
+            st.header(doc_name)
+            st.caption("Keywords: " + ", ".join(doc_keywords))
+            st.write(st.session_state.summaries[doc_name])
+            st.divider()
+
+        # Display the plan
+        st.header("Study Plan")
+        plan = generate_plan_with_spinner()
+        for data in plan.get("plan", []):
+            for doc_name, doc_data in data.items():
+                st.subheader(doc_name)
+                st.write(f"**Background Concepts**")
+                for concept in doc_data.get("Background Concepts", []):
+                    st.write(f"- {concept}")
+                st.write(f"**Key Concepts**")
+                for concept in doc_data.get("Key Concepts", []):
+                    st.write(f"- {concept}")
+                st.write(f"**Further Reading**")
+                for concept in doc_data.get("Further Reading", []):
+                    st.write(f"- {concept}")
+            st.write("")
+
+        st.session_state.processed = True
+        st.success('Summarization completed.')
+
+if st.session_state.processed:
+    st.header("Loaded Files")
+    for file in st.session_state.uploaded_files:
+        st.write(file.name)
+
+    st.header("Chat")
+    # Iterate through the bot and user message and print them alternatively
+    message_i = 0
+    while True:
+        if len(st.session_state.bot_messages) > message_i:
+            message(st.session_state.bot_messages[message_i])
+        else:
+            break
+        if len(st.session_state.user_messages) > message_i:
+            message(st.session_state.user_messages[message_i], is_user=True)
+        else:
+            break
+        message_i += 1
+
+    user_input = st.text_input("What do you want to learn more about?", key="user_message")
+    if st.button("Send"):
+        generate_response(user_input)
+        st.experimental_rerun()