Upload app.py

app.py

@@ -22,6 +22,48 @@ LOCAL_VECTOR_STORE_DIR = Path(__file__).resolve().parent.joinpath("vector_store"
 
 deep_strip = lambda text: re.sub(r"\s+", " ", text or "").strip()
 
+get_references = lambda relevant_docs: " ".join(
+    [f"[{ref}]" for ref in sorted([ref.metadata["chunk_id"] for ref in relevant_docs])]
+)
+session_state_2_llm_chat_history = lambda session_state: [
+    ss[:2] for ss in session_state
+]
+
+
+def get_conversation_history():
+    return json.dumps(
+        {
+            "document_urls": (
+                st.session_state.source_doc_urls
+                if "source_doc_urls" in st.session_state
+                else []
+            ),
+            "document_snippets": (
+                st.session_state.headers.to_list()
+                if "headers" in st.session_state
+                else []
+            ),
+            "conversation": [
+                {"human": message[0], "ai": message[1], "references": message[2]}
+                for message in st.session_state.messages
+            ],
+            "costing": (
+                st.session_state.costing if "costing" in st.session_state else []
+            ),
+            "total_cost": (
+                {
+                    k: sum(d[k] for d in st.session_state.costing)
+                    for k in st.session_state.costing[0]
+                }
+                if "costing" in st.session_state and len(st.session_state.costing) > 0
+                else {}
+            ),
+        }
+    )
+
+
+ai_message_format = lambda message, references: f"{message}\n\n---\n\n{references}"
+
 
 def embeddings_on_local_vectordb(texts):
     colbert = RAGPretrainedModel.from_pretrained("colbert-ir/colbertv1.9")
@@ -48,12 +90,18 @@ def query_llm(retriever, query):
     relevant_docs = retriever.get_relevant_documents(query)
     with get_openai_callback() as cb:
         result = qa_chain(
-            {"question": query, "chat_history": st.session_state.messages}
+            {
+                "question": query,
+                "chat_history": session_state_2_llm_chat_history(
+                    st.session_state.messages
+                ),
+            }
         )
         stats = cb
     result = result["answer"]
-    st.session_state.messages.append((query, result))
-    return relevant_docs, result, stats
+    references = get_references(relevant_docs)
+    st.session_state.messages.append((query, result, references))
+    return result, references, stats
 
 
 def input_fields():
@@ -202,7 +250,7 @@ def process_web(url):
 
 
 def boot():
-    st.title("Xi Chatbot")
+    st.title("Agent Xi - An ArXiv Chatbot")
     st.sidebar.title("Input Documents")
     input_fields()
     st.sidebar.button("Submit Documents", on_click=process_documents)
@@ -216,19 +264,17 @@ def boot():
 
     for message in st.session_state.messages:
         st.chat_message("human").write(message[0])
-        st.chat_message("ai").write(message[1])
+        st.chat_message("ai").write(ai_message_format(message[1], message[2]))
     if query := st.chat_input():
         st.chat_message("human").write(query)
-        references, response, stats = query_llm(st.session_state.retriever, query)
-        sorted_references = sorted([ref.metadata["chunk_id"] for ref in references])
-        references_str = " ".join([f"[{ref}]" for ref in sorted_references])
-        st.chat_message("ai").write(response + "\n\n---\nReferences:" + references_str)
+        response, references, stats = query_llm(st.session_state.retriever, query)
+        st.chat_message("ai").write(ai_message_format(response, references))
 
         st.session_state.costing.append(
             {
                 "prompt tokens": stats.prompt_tokens,
                 "completion tokens": stats.completion_tokens,
-                "total cost": stats.total_cost,
+                "cost": stats.total_cost,
             }
         )
         stats_df = pd.DataFrame(st.session_state.costing)
@@ -236,15 +282,10 @@ def boot():
         st.sidebar.write(stats_df)
     st.sidebar.download_button(
         "Download Conversation",
-        json.dumps(
-            [
-                {"human": message[0], "ai": message[1]}
-                for message in st.session_state.messages
-            ]
-        ),
+        get_conversation_history(),
         "conversation.json",
     )
 
 
 if __name__ == "__main__":
-    boot()
+    boot()