Update app.py

app.py

@@ -4,16 +4,16 @@ import re
 from pathlib import Path
 from statistics import median
 
+import pandas as pd
 import streamlit as st
 from bs4 import BeautifulSoup
+from langchain.callbacks import get_openai_callback
 from langchain.chains import ConversationalRetrievalChain
 from langchain.docstore.document import Document
 from langchain.document_loaders import PDFMinerPDFasHTMLLoader, WebBaseLoader
-from langchain_openai import ChatOpenAI
 from langchain.retrievers.multi_query import MultiQueryRetriever
+from langchain_openai import ChatOpenAI
 from ragatouille import RAGPretrainedModel
-import pandas as pd
-
 
 st.set_page_config(layout="wide")
 os.environ["OPENAI_API_KEY"] = "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"
@@ -46,17 +46,21 @@ def query_llm(retriever, query):
         chain_type="stuff",
     )
     relevant_docs = retriever.get_relevant_documents(query)
-    result = qa_chain({"question": query, "chat_history": st.session_state.messages})
+    with get_openai_callback() as cb:
+        result = qa_chain(
+            {"question": query, "chat_history": st.session_state.messages}
+        )
+        stats = cb
     result = result["answer"]
     st.session_state.messages.append((query, result))
-    return relevant_docs, result
+    return relevant_docs, result, stats
 
 
 def input_fields():
     st.session_state.source_doc_urls = [
         url.strip()
         for url in st.sidebar.text_area(
-            "Source Document URLs\n(New line separated)", height=200
+            "Source Document URLs\n(New line separated)", height=50
         ).split("\n")
     ]
 
@@ -201,25 +205,34 @@ def boot():
     st.title("Xi Chatbot")
     st.sidebar.title("Input Documents")
     input_fields()
-    col1, col2 = st.columns([4, 1])
     st.sidebar.button("Submit Documents", on_click=process_documents)
     if "headers" in st.session_state:
-        col2.write("### References")
-        col2.write(st.session_state.headers)
+        st.sidebar.write("### References")
+        st.sidebar.write(st.session_state.headers)
+    if "costing" not in st.session_state:
+        st.session_state.costing = []
     if "messages" not in st.session_state:
         st.session_state.messages = []
     for message in st.session_state.messages:
-        col1.chat_message("human").write(message[0])
-        col1.chat_message("ai").write(message[1])
-    if query := col1.chat_input():
-        col1.chat_message("human").write(query)
-        references, response = query_llm(st.session_state.retriever, query)
+        st.chat_message("human").write(message[0])
+        st.chat_message("ai").write(message[1])
+    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])
-        col1.chat_message("ai").write(
-            response + "\n\n---\nReferences:" + references_str
+        st.chat_message("ai").write(response + "\n\n---\nReferences:" + references_str)
+        st.session_state.costing.append(
+            {
+                "prompt tokens": stats.prompt_tokens,
+                "completion tokens": stats.completion_tokens,
+                "total cost": stats.total_cost,
+            }
         )
+        stats_df = pd.DataFrame(st.session_state.costing)
+        stats_df.loc["total"] = stats_df.sum()
+        st.sidebar.write(stats_df)
 
 
 if __name__ == "__main__":
-    boot()
+    boot()