Upload 2 files

app.py

@@ -9,11 +9,10 @@ from bs4 import BeautifulSoup
 from langchain.chains import ConversationalRetrievalChain
 from langchain.docstore.document import Document
 from langchain.document_loaders import PDFMinerPDFasHTMLLoader, WebBaseLoader
-from langchain.embeddings.openai import OpenAIEmbeddings
-from langchain_openai import ChatOpenAI, OpenAI
-from langchain.vectorstores import Chroma
+from langchain_openai import ChatOpenAI
 from langchain.retrievers.multi_query import MultiQueryRetriever
 from ragatouille import RAGPretrainedModel
+import pandas as pd
 
 
 st.set_page_config(layout="wide")
@@ -55,7 +54,10 @@ def query_llm(retriever, query):
 
 def input_fields():
     st.session_state.source_doc_urls = [
-        url.strip() for url in st.sidebar.text_input("Source Document URLs").split(",")
+        url.strip()
+        for url in st.sidebar.text_area(
+            "Source Document URLs\n(New line separated)", height=200
+        ).split("\n")
     ]
 
 
@@ -68,9 +70,9 @@ def process_documents():
             else:
                 snippets.extend(process_web(url))
         st.session_state.retriever = embeddings_on_local_vectordb(snippets)
-        st.session_state.headers = [
-            " ".join(snip.metadata["header"].split()[:10]) for snip in snippets
-        ]
+        st.session_state.headers = pd.Series(
+            [snip.metadata["header"] for snip in snippets], name="references"
+        )
     except Exception as e:
         st.error(f"An error occurred: {e}")
 
@@ -88,12 +90,13 @@ def process_pdf(url):
         Document(
             page_content=deep_strip(snip[1]["header_text"]) + " " + deep_strip(snip[0]),
             metadata={
-                "header": deep_strip(snip[1]["header_text"]),
+                "header": " ".join(snip[1]["header_text"].split()[:10]),
                 "source_url": url,
                 "source_type": "pdf",
+                "chunk_id": i,
             },
         )
-        for snip in semantic_snippets
+        for i, snip in enumerate(semantic_snippets)
     ]
     return document_snippets
 
@@ -196,12 +199,13 @@ def process_web(url):
 
 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:
-        for header in st.session_state.headers:
-            col2.info(header)
+        col2.write("### References")
+        col2.write(st.session_state.headers)
     if "messages" not in st.session_state:
         st.session_state.messages = []
     for message in st.session_state.messages:
@@ -210,11 +214,11 @@ def boot():
     if query := col1.chat_input():
         col1.chat_message("human").write(query)
         references, response = query_llm(st.session_state.retriever, query)
-        for snip in references:
-            st.sidebar.success(
-                f'Section {" ".join(snip.metadata["header"].split()[:10])}'
-            )
-        col1.chat_message("ai").write(response)
+        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
+        )
 
 
 if __name__ == "__main__":

requirements.txt

@@ -8,4 +8,5 @@ chromadb==0.4.22
 tiktoken==0.5.2
 pdfminer.six==20231228
 beautifulsoup4==4.12.3
-RAGatouille==0.0.7.post7
+RAGatouille==0.0.7.post7
+pandas==2.2.1