Upload 5 files

app.py

@@ -1,37 +1,57 @@
-import math
+import streamlit as st
 import os
-import re
-from pathlib import Path
-from statistics import median
-import json
 import pandas as pd
-import streamlit as st
-from bs4 import BeautifulSoup
+from command_center import CommandCenter
+from process_documents import process_documents
+from embed_documents import create_retriever
+import json
 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.retrievers.multi_query import MultiQueryRetriever
 from langchain_openai import ChatOpenAI
-from ragatouille import RAGPretrainedModel
 
 st.set_page_config(layout="wide")
 os.environ["OPENAI_API_KEY"] = "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"
 
-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
+    ss[:2] for ss in session_state if not ss[0].startswith("/")
 ]
+ai_message_format = lambda message, references: (
+    f"{message}\n\n---\n\n{references}" if references != "" else message
+)
+
+
+def process_documents_wrapper(inputs):
+    snippets = process_documents(inputs)
+    st.session_state.retriever = create_retriever(snippets)
+    st.session_state.source_doc_urls = inputs
+    st.session_state.index = [snip.metadata["header"] for snip in snippets]
+    response = f"Uploaded and processed documents {inputs}"
+    st.session_state.messages.append((f"/upload {inputs}", response, ""))
+    return response
 
 
-def get_conversation_history():
-    return json.dumps(
+def index_documents_wrapper(inputs=None):
+    response = pd.Series(st.session_state.index, name="references").to_markdown()
+    st.session_state.messages.append(("/index", response, ""))
+    return response
+
+
+def calculate_cost_wrapper(inputs=None):
+    try:
+        stats_df = pd.DataFrame(st.session_state.costing)
+        stats_df.loc["total"] = stats_df.sum()
+        response = stats_df.to_markdown()
+    except ValueError:
+        response = "No costing incurred yet"
+    st.session_state.messages.append(("/cost", response, ""))
+    return response
+
+
+def download_conversation_wrapper(inputs=None):
+    conversation_data = json.dumps(
         {
             "document_urls": (
                 st.session_state.source_doc_urls
@@ -39,7 +59,7 @@ def get_conversation_history():
                 else []
             ),
             "document_snippets": (
-                st.session_state.headers.to_list()
+                st.session_state.index.to_list()
                 if "headers" in st.session_state
                 else []
             ),
@@ -60,38 +80,28 @@ def get_conversation_history():
             ),
         }
     )
-
-
-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")
-    colbert.index(
-        collection=[chunk.page_content for chunk in texts],
-        split_documents=False,
-        document_metadatas=[chunk.metadata for chunk in texts],
-        index_name="vector_store",
-    )
-    retriever = colbert.as_langchain_retriever(k=5)
-    retriever = MultiQueryRetriever.from_llm(
-        retriever=retriever, llm=ChatOpenAI(temperature=0)
+    st.sidebar.download_button(
+        "Download Conversation",
+        conversation_data,
+        file_name="conversation_data.json",
+        mime="application/json",
     )
-    return retriever
+    st.session_state.messages.append(("/download", "Conversation data downloaded", ""))
 
 
-def query_llm(retriever, query):
+def query_llm_wrapper(inputs):
+    retriever = st.session_state.retriever
     qa_chain = ConversationalRetrievalChain.from_llm(
         llm=ChatOpenAI(model="gpt-4-0125-preview", temperature=0),
         retriever=retriever,
         return_source_documents=True,
         chain_type="stuff",
     )
-    relevant_docs = retriever.get_relevant_documents(query)
+    relevant_docs = retriever.get_relevant_documents(inputs)
     with get_openai_callback() as cb:
         result = qa_chain(
             {
-                "question": query,
+                "question": inputs,
                 "chat_history": session_state_2_llm_chat_history(
                     st.session_state.messages
                 ),
@@ -100,192 +110,62 @@ def query_llm(retriever, query):
         stats = cb
     result = result["answer"]
     references = get_references(relevant_docs)
-    st.session_state.messages.append((query, result, references))
-    return result, references, 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=50
-        ).split("\n")
-    ]
-
-
-def process_documents():
-    try:
-        snippets = []
-        for url in st.session_state.source_doc_urls:
-            if url.endswith(".pdf"):
-                snippets.extend(process_pdf(url))
-            else:
-                snippets.extend(process_web(url))
-        st.session_state.retriever = embeddings_on_local_vectordb(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}")
-
-
-def process_pdf(url):
-    data = PDFMinerPDFasHTMLLoader(url).load()[0]
-    content = BeautifulSoup(data.page_content, "html.parser").find_all("div")
-    snippets = get_pdf_snippets(content)
-    filtered_snippets = filter_pdf_snippets(snippets, new_line_threshold_ratio=0.4)
-    median_font_size = math.ceil(
-        median([font_size for _, font_size in filtered_snippets])
-    )
-    semantic_snippets = get_pdf_semantic_snippets(filtered_snippets, median_font_size)
-    document_snippets = [
-        Document(
-            page_content=deep_strip(snip[1]["header_text"]) + " " + deep_strip(snip[0]),
-            metadata={
-                "header": " ".join(snip[1]["header_text"].split()[:10]),
-                "source_url": url,
-                "source_type": "pdf",
-                "chunk_id": i,
-            },
-        )
-        for i, snip in enumerate(semantic_snippets)
-    ]
-    return document_snippets
-
-
-def get_pdf_snippets(content):
-    current_font_size = None
-    current_text = ""
-    snippets = []
-    for cntnt in content:
-        span = cntnt.find("span")
-        if not span:
-            continue
-        style = span.get("style")
-        if not style:
-            continue
-        font_size = re.findall("font-size:(\d+)px", style)
-        if not font_size:
-            continue
-        font_size = int(font_size[0])
-
-        if not current_font_size:
-            current_font_size = font_size
-        if font_size == current_font_size:
-            current_text += cntnt.text
-        else:
-            snippets.append((current_text, current_font_size))
-            current_font_size = font_size
-            current_text = cntnt.text
-    snippets.append((current_text, current_font_size))
-    return snippets
-
-
-def filter_pdf_snippets(content_list, new_line_threshold_ratio):
-    filtered_list = []
-    for e, (content, font_size) in enumerate(content_list):
-        newline_count = content.count("\n")
-        total_chars = len(content)
-        ratio = newline_count / total_chars
-        if ratio <= new_line_threshold_ratio:
-            filtered_list.append((content, font_size))
-    return filtered_list
-
-
-def get_pdf_semantic_snippets(filtered_snippets, median_font_size):
-    semantic_snippets = []
-    current_header = None
-    current_content = []
-    header_font_size = None
-    content_font_sizes = []
-
-    for content, font_size in filtered_snippets:
-        if font_size > median_font_size:
-            if current_header is not None:
-                metadata = {
-                    "header_font_size": header_font_size,
-                    "content_font_size": (
-                        median(content_font_sizes) if content_font_sizes else None
-                    ),
-                    "header_text": current_header,
-                }
-                semantic_snippets.append((current_content, metadata))
-                current_content = []
-                content_font_sizes = []
-
-            current_header = content
-            header_font_size = font_size
-        else:
-            content_font_sizes.append(font_size)
-            if current_content:
-                current_content += " " + content
-            else:
-                current_content = content
-
-    if current_header is not None:
-        metadata = {
-            "header_font_size": header_font_size,
-            "content_font_size": (
-                median(content_font_sizes) if content_font_sizes else None
-            ),
-            "header_text": current_header,
+    st.session_state.messages.append((inputs, result, references))
+    st.session_state.costing.append(
+        {
+            "prompt tokens": stats.prompt_tokens,
+            "completion tokens": stats.completion_tokens,
+            "cost": stats.total_cost,
         }
-        semantic_snippets.append((current_content, metadata))
-    return semantic_snippets
-
-
-def process_web(url):
-    data = WebBaseLoader(url).load()[0]
-    document_snippets = [
-        Document(
-            page_content=deep_strip(data.page_content),
-            metadata={
-                "header": data.metadata["title"],
-                "source_url": url,
-                "source_type": "web",
-            },
-        )
-    ]
-    return document_snippets
+    )
+    return result, references
 
 
-def boot():
+def boot(command_center):
     st.title("Agent Xi - An ArXiv Chatbot")
-    st.sidebar.title("Input Documents")
-    input_fields()
-    st.sidebar.button("Submit Documents", on_click=process_documents)
-    if "headers" in st.session_state:
-        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:
         st.chat_message("human").write(message[0])
         st.chat_message("ai").write(ai_message_format(message[1], message[2]))
     if query := st.chat_input():
         st.chat_message("human").write(query)
-        response, references, stats = query_llm(st.session_state.retriever, query)
-        st.chat_message("ai").write(ai_message_format(response, references))
+        response = command_center.execute_command(query)
+        if response is None:
+            pass
+        elif type(response) == tuple:
+            result, references = response
+            st.chat_message("ai").write(ai_message_format(result, references))
+        else:
+            st.chat_message("ai").write(response)
+
 
-        st.session_state.costing.append(
+if __name__ == "__main__":
+    all_commands = [
+        ("/upload", list, process_documents_wrapper, "Upload and process documents"),
+        ("/index", None, index_documents_wrapper, "View index of processed documents"),
+        ("/cost", None, calculate_cost_wrapper, "Calculate cost of conversation"),
+        (
+            "/download",
+            None,
+            download_conversation_wrapper,
+            "Download conversation data",
+        ),
+    ]
+    st.sidebar.title("Commands Menu")
+    st.sidebar.write(
+        pd.DataFrame(
             {
-                "prompt tokens": stats.prompt_tokens,
-                "completion tokens": stats.completion_tokens,
-                "cost": stats.total_cost,
+                "Command": [command[0] for command in all_commands],
+                "Description": [command[3] for command in all_commands],
             }
         )
-        stats_df = pd.DataFrame(st.session_state.costing)
-        stats_df.loc["total"] = stats_df.sum()
-        st.sidebar.write(stats_df)
-    st.sidebar.download_button(
-        "Download Conversation",
-        get_conversation_history(),
-        "conversation.json",
     )
-
-
-if __name__ == "__main__":
-    boot()
+    command_center = CommandCenter(
+        default_input_type=str,
+        default_function=query_llm_wrapper,
+        all_commands=[command[:3] for command in all_commands],
+    )
+    boot(command_center)

command_center.py

@@ -0,0 +1,44 @@
+class CommandCenter:
+    def __init__(self, default_input_type, default_function=None, all_commands=None):
+        self.commands = {}
+        self.add_command("/default", default_input_type, default_function)
+        if all_commands:
+            for command, input_type, function in all_commands:
+                self.add_command(command, input_type, function)
+
+    def add_command(self, command, input_type, function=None):
+        assert input_type in [None, str, int, float, bool, list], "Invalid input type"
+        self.commands[command] = {"input_type": input_type, "function": function}
+
+    def parse_command(self, input_string):
+        # parsing the input string
+        if not input_string.startswith("/"):
+            command = "/default"
+            argument = input_string.split(" ")
+        else:
+            inputs = input_string.split(" ")
+            command = inputs[0]
+            argument = inputs[1:]
+
+        # type casting the arguments
+        if self.commands[command]["input_type"] == str:
+            argument = " ".join(argument)
+        elif self.commands[command]["input_type"] == int:
+            argument = int(" ".join(argument))
+        elif self.commands[command]["input_type"] == float:
+            argument = float(" ".join(argument))
+        elif self.commands[command]["input_type"] == bool:
+            argument = bool(" ".join(argument))
+        elif self.commands[command]["input_type"] == list:
+            argument = argument
+        else:
+            argument = None
+
+        return command, argument
+
+    def execute_command(self, input_string):
+        command, argument = self.parse_command(input_string)
+        if command in self.commands:
+            return self.commands[command]["function"](argument)
+        else:
+            return "Invalid command"

embed_documents.py

@@ -0,0 +1,18 @@
+from langchain.retrievers.multi_query import MultiQueryRetriever
+from langchain_openai import ChatOpenAI
+from ragatouille import RAGPretrainedModel
+
+
+def create_retriever(texts):
+    colbert = RAGPretrainedModel.from_pretrained("colbert-ir/colbertv1.9")
+    colbert.index(
+        collection=[chunk.page_content for chunk in texts],
+        split_documents=False,
+        document_metadatas=[chunk.metadata for chunk in texts],
+        index_name="vector_store",
+    )
+    retriever = colbert.as_langchain_retriever(k=5)
+    retriever = MultiQueryRetriever.from_llm(
+        retriever=retriever, llm=ChatOpenAI(temperature=0)
+    )
+    return retriever

process_documents.py

@@ -0,0 +1,141 @@
+import math
+import re
+from statistics import median
+from bs4 import BeautifulSoup
+from langchain.docstore.document import Document
+from langchain.document_loaders import PDFMinerPDFasHTMLLoader, WebBaseLoader
+
+deep_strip = lambda text: re.sub(r"\s+", " ", text or "").strip()
+
+
+def process_documents(urls):
+    snippets = []
+    for url in urls:
+        if url.endswith(".pdf"):
+            snippets.extend(process_pdf(url))
+        else:
+            snippets.extend(process_web(url))
+    for e, snippet in enumerate(snippets):
+        snippet.metadata["chunk_id"] = e
+    return snippets
+
+
+def process_pdf(url):
+    data = PDFMinerPDFasHTMLLoader(url).load()[0]
+    content = BeautifulSoup(data.page_content, "html.parser").find_all("div")
+    snippets = get_pdf_snippets(content)
+    filtered_snippets = filter_pdf_snippets(snippets, new_line_threshold_ratio=0.4)
+    median_font_size = math.ceil(
+        median([font_size for _, font_size in filtered_snippets])
+    )
+    semantic_snippets = get_pdf_semantic_snippets(filtered_snippets, median_font_size)
+    document_snippets = [
+        Document(
+            page_content=deep_strip(snip[1]["header_text"]) + " " + deep_strip(snip[0]),
+            metadata={
+                "header": " ".join(snip[1]["header_text"].split()[:10]),
+                "source_url": url,
+                "source_type": "pdf",
+                "chunk_id": i,
+            },
+        )
+        for i, snip in enumerate(semantic_snippets)
+    ]
+    return document_snippets
+
+
+def get_pdf_snippets(content):
+    current_font_size = None
+    current_text = ""
+    snippets = []
+    for cntnt in content:
+        span = cntnt.find("span")
+        if not span:
+            continue
+        style = span.get("style")
+        if not style:
+            continue
+        font_size = re.findall("font-size:(\d+)px", style)
+        if not font_size:
+            continue
+        font_size = int(font_size[0])
+
+        if not current_font_size:
+            current_font_size = font_size
+        if font_size == current_font_size:
+            current_text += cntnt.text
+        else:
+            snippets.append((current_text, current_font_size))
+            current_font_size = font_size
+            current_text = cntnt.text
+    snippets.append((current_text, current_font_size))
+    return snippets
+
+
+def filter_pdf_snippets(content_list, new_line_threshold_ratio):
+    filtered_list = []
+    for e, (content, font_size) in enumerate(content_list):
+        newline_count = content.count("\n")
+        total_chars = len(content)
+        ratio = newline_count / total_chars
+        if ratio <= new_line_threshold_ratio:
+            filtered_list.append((content, font_size))
+    return filtered_list
+
+
+def get_pdf_semantic_snippets(filtered_snippets, median_font_size):
+    semantic_snippets = []
+    current_header = None
+    current_content = []
+    header_font_size = None
+    content_font_sizes = []
+
+    for content, font_size in filtered_snippets:
+        if font_size > median_font_size:
+            if current_header is not None:
+                metadata = {
+                    "header_font_size": header_font_size,
+                    "content_font_size": (
+                        median(content_font_sizes) if content_font_sizes else None
+                    ),
+                    "header_text": current_header,
+                }
+                semantic_snippets.append((current_content, metadata))
+                current_content = []
+                content_font_sizes = []
+
+            current_header = content
+            header_font_size = font_size
+        else:
+            content_font_sizes.append(font_size)
+            if current_content:
+                current_content += " " + content
+            else:
+                current_content = content
+
+    if current_header is not None:
+        metadata = {
+            "header_font_size": header_font_size,
+            "content_font_size": (
+                median(content_font_sizes) if content_font_sizes else None
+            ),
+            "header_text": current_header,
+        }
+        semantic_snippets.append((current_content, metadata))
+    return semantic_snippets
+
+
+def process_web(url):
+    data = WebBaseLoader(url).load()[0]
+    document_snippets = [
+        Document(
+            page_content=deep_strip(data.page_content),
+            metadata={
+                "header": data.metadata["title"],
+                "source_url": url,
+                "source_type": "web",
+                "chunk_id": 0,
+            },
+        )
+    ]
+    return document_snippets

requirements.txt

@@ -4,9 +4,9 @@ langchain-community==0.0.24
 langchain-core==0.1.27
 langchain-experimental==0.0.49
 langchain-openai==0.0.8
-chromadb==0.4.22
 tiktoken==0.5.2
 pdfminer.six==20231228
 beautifulsoup4==4.12.3
 RAGatouille==0.0.7.post7
-pandas==2.2.1
+pandas==2.2.1
+tabulate==0.9.0