Minor Fixes

app.py

@@ -21,6 +21,7 @@ from insights_bullet_chain import insights_bullet_chain
 from synopsis_chain import synopsis_chain
 from custom_exceptions import InvalidArgumentError, InvalidCommandError
 from openai_configuration import openai_parser
+import pickle
 
 st.set_page_config(layout="wide")
 
@@ -36,14 +37,14 @@ Here's a quick guide to getting started with me:
 | `/configure --key <api key> --model <model>` | Configure the OpenAI API key and model for our conversation. |
 | `/add-papers <list of urls>` | Upload and process documents for our conversation. |
 | `/library` | View an index of processed documents to easily navigate your research. |
-| `/view-doc <document id>` | View the content of a specific document. |
+| `/view-snip <snippet id>` | View the content of a specific snnippet. |
 | `/session-expense` | Calculate the cost of our conversation, ensuring transparency in resource usage. |
 | `/export` | Download conversation data for your records or further analysis. |
-| `/auto-insight <document id>` | Automatically generate questions and answers for the paper. |
-| `/condense-summary <document id>` | Generate increasingly concise, entity-dense summaries of the paper. |
-| `/insight-bullets <list of document ids>` | Extract and summarize key insights, methods, results, and conclusions. |
-| `/paper-synopsis <document id>` | Generate a synopsis of the paper. |
-| `/deep-dive [<list of document ids>] <query>` | Query me with a specific context. |
+| `/auto-insight <list of snippet ids>` | Automatically generate questions and answers for the paper. |
+| `/condense-summary <list of snippet ids>` | Generate increasingly concise, entity-dense summaries of the paper. |
+| `/insight-bullets <list of snippet ids>` | Extract and summarize key insights, methods, results, and conclusions. |
+| `/paper-synopsis <list of snippet ids>` | Generate a synopsis of the paper. |
+| `/deep-dive [<list of snippet ids>] <query>` | Query me with a specific context. |
 
 
 <br>
@@ -75,7 +76,7 @@ def index_documents_wrapper(inputs=None):
 
 def view_document_wrapper(inputs):
     response = st.session_state.documents[inputs].page_content
-    st.session_state.messages.append((f"/view-doc {inputs}", response, "identity"))
+    st.session_state.messages.append((f"/view-snip {inputs}", response, "identity"))
     return (response, "identity")
 
 
@@ -91,6 +92,7 @@ def calculate_cost_wrapper(inputs=None):
 
 
 def download_conversation_wrapper(inputs=None):
+    pickle.dump(st.session_state.messages, open("conversation_data.pkl", "wb"))
     conversation_data = json.dumps(
         {
             "document_urls": (
@@ -180,9 +182,9 @@ def query_llm_wrapper(inputs):
 
 
 def chain_of_density_wrapper(inputs):
-    if inputs == "":
-        raise InvalidArgumentError("Please provide a document id")
-    document = st.session_state.documents[inputs].page_content
+    if inputs == []:
+        raise InvalidArgumentError("Please provide snippet ids")
+    document = "\n\n".join([st.session_state.documents[c].page_content for c in inputs])
     llm = ChatOpenAI(model=st.session_state.model, temperature=0)
     with get_openai_callback() as cb:
         summary = chain_of_density_chain(llm).invoke({"paper": document})
@@ -199,9 +201,9 @@ def chain_of_density_wrapper(inputs):
 
 
 def synopsis_wrapper(inputs):
-    if inputs == "":
-        raise InvalidArgumentError("Please provide a document id")
-    document = st.session_state.documents[inputs].page_content
+    if inputs == []:
+        raise InvalidArgumentError("Please provide snippet ids")
+    document = "\n\n".join([st.session_state.documents[c].page_content for c in inputs])
     llm = ChatOpenAI(model=st.session_state.model, temperature=0)
     with get_openai_callback() as cb:
         summary = synopsis_chain(llm).invoke({"paper": document})
@@ -218,8 +220,8 @@ def synopsis_wrapper(inputs):
 
 
 def insights_bullet_wrapper(inputs):
-    if inputs == "":
-        raise InvalidArgumentError("Please provide a document id")
+    if inputs == []:
+        raise InvalidArgumentError("Please provide snippet ids")
     document = "\n\n".join([st.session_state.documents[c].page_content for c in inputs])
     llm = ChatOpenAI(model=st.session_state.model, temperature=0)
     with get_openai_callback() as cb:
@@ -237,9 +239,9 @@ def insights_bullet_wrapper(inputs):
 
 
 def auto_qa_chain_wrapper(inputs):
-    if inputs == "":
-        raise InvalidArgumentError("Please provide a document id")
-    document = st.session_state.documents[inputs].page_content
+    if inputs == []:
+        raise InvalidArgumentError("Please provide snippet ids")
+    document = "\n\n".join([st.session_state.documents[c].page_content for c in inputs])
     llm = ChatOpenAI(model=st.session_state.model, temperature=0)
     auto_qa_conversation = []
     with get_openai_callback() as cb:
@@ -310,15 +312,15 @@ if __name__ == "__main__":
         ("/configure", str, configure_openai_wrapper),
         ("/add-papers", list, process_documents_wrapper),
         ("/library", None, index_documents_wrapper),
-        ("/view-doc", str, view_document_wrapper),
+        ("/view-snip", str, view_document_wrapper),
         ("/session-expense", None, calculate_cost_wrapper),
         ("/export", None, download_conversation_wrapper),
         ("/help-me", None, lambda x: (welcome_message, "identity")),
-        ("/auto-insight", str, auto_qa_chain_wrapper),
+        ("/auto-insight", list, auto_qa_chain_wrapper),
         ("/deep-dive", str, query_llm_wrapper),
-        ("/condense-summary", str, chain_of_density_wrapper),
+        ("/condense-summary", list, chain_of_density_wrapper),
         ("/insight-bullets", list, insights_bullet_wrapper),
-        ("/paper-synopsis", str, synopsis_wrapper),
+        ("/paper-synopsis", list, synopsis_wrapper),
     ]
     command_center = CommandCenter(
         default_input_type=str,
@@ -334,7 +336,9 @@ if __name__ == "__main__":
     }
     jsonify_functions = {
         "identity": lambda x: x,
-        "dataframe": lambda x: x.to_dict(orient="records"),
+        "dataframe": lambda x: (
+            x.to_dict(orient="records") if isinstance(x, pd.DataFrame) else x
+        ),
         "reponse_with_citations": lambda x: x,
     }
     boot(command_center, formating_functions)