Merge pull request #5 from jonfairbanks/develop

UX Improvements
jonfairbanks · Feb 21, 2024 · 1bdd6f4 · 1bdd6f4
2 parents 497908d + b4cc043
commit 1bdd6f4
Show file tree

Hide file tree

Showing 9 changed files with 144 additions and 72 deletions.
diff --git a/.github/workflows/main.yaml b/.github/workflows/main.yaml
@@ -27,7 +27,7 @@ jobs:
           severity: 'CRITICAL,HIGH,MODERATE'
 
       - name: Upload Trivy scan results to GitHub Security tab
-        uses: github/codeql-action/upload-sarif@v2
+        uses: github/codeql-action/upload-sarif@v3
         with:
           sarif_file: 'trivy-results.sarif'
 

diff --git a/README.md b/README.md
@@ -10,8 +10,7 @@ Offline, Open-Source RAG
 
 Ingest files for retrieval augmented generation (RAG) with open-source Large Language Models (LLMs), all without 3rd parties or sensitive data leaving your network.
 
-Features:
-- Supports 100% offline embeddings & LLMs; No OpenAI!
+- Offline Embeddings & LLMs Support (No OpenAI!)
 - Streaming Responses
 - Conversation Memory
 - Chat Export
@@ -55,9 +54,10 @@ Docker:
 - [ ] Implement Log Library
 
 ### Known Issues & Bugs
-- [ ] Refreshing the page loses all state (expected streamlit behavior; need to implement local-storage)
+- [ ] Refreshing the page loses all state (expected Streamlit behavior; need to implement local-storage)
 - [x] Files can be uploaded before Ollama config is set, leading to embedding errors
-- [ ] Assuming localhost is used, Models are not automatically loaded at first page load
+- [ ] Assuming Ollama is hosted on localhost, Models are automatically loaded and selected, but the dropdown does not render the selected option
+- [ ] Upon first sending the first Chat message, the File Processing expander appears to re-run itself
 
 ### Resources
 - [Ollama](https://ollama.com/)

diff --git a/components/header.py b/components/header.py
@@ -2,7 +2,7 @@
 
 
 def set_page_header():
-    st.header("📚 Offline, Open-Source RAG")
+    st.header("📚 Offline, Open-Source RAG", anchor=False)
     st.caption(
         "Ingest files for retrieval augmented generation (RAG) with open-source Large Language Models (LLMs), all without 3rd parties or sensitive data leaving your network."
     )
diff --git a/components/page_state.py b/components/page_state.py
@@ -1,47 +1,74 @@
 import streamlit as st
 
+from utils.ollama import get_models
+
 
 def set_initial_state():
+
+    ###########
+    # General #
+    ###########
+
     if "ollama_endpoint" not in st.session_state:
-        st.session_state["ollama_endpoint"] = "http://localhost:11434"  # Revert to None
+        st.session_state["ollama_endpoint"] = "http://localhost:11434"
 
     if "embedding_model" not in st.session_state:
         st.session_state["embedding_model"] = None
 
     if "ollama_models" not in st.session_state:
-        st.session_state["ollama_models"] = []
+        try:
+            models = get_models()
+            st.session_state["ollama_models"] = models
+        except Exception as err:
+            print(
+                f"Warning: Initial loading of Ollama models failed. You might be hosting Ollama somewhere other than localhost. -- {err}"
+            )
+            st.session_state["ollama_models"] = []
+            pass
 
     if "selected_model" not in st.session_state:
         st.session_state["selected_model"] = None
 
-    if "top_k" not in st.session_state:
-        st.session_state["top_k"] = None
+    if "messages" not in st.session_state:
+        st.session_state["messages"] = [
+            {
+                "role": "assistant",
+                "content": "Hello! Import some files or ingest a GitHub repo and we can get started.",
+            }
+        ]
 
-    if "github_repo" not in st.session_state:
-        st.session_state["github_repo"] = None
+    ######################
+    #  Files & Documents #
+    ######################
 
     if "file_list" not in st.session_state:
         st.session_state["file_list"] = []
 
-    if "embedding_model" not in st.session_state:
-        st.session_state["embedding_model"] = None
-
-    if "chunk_size" not in st.session_state:
-        st.session_state["chunk_size"] = None
+    if "github_repo" not in st.session_state:
+        st.session_state["github_repo"] = None
 
-    if "advanced" not in st.session_state:
-        st.session_state["advanced"] = False
+    ###############
+    # Llama-Index #
+    ###############
 
     if "documents" not in st.session_state:
         st.session_state["documents"] = None
 
     if "query_engine" not in st.session_state:
         st.session_state["query_engine"] = None
 
-    if "messages" not in st.session_state:
-        st.session_state["messages"] = [
-            {
-                "role": "assistant",
-                "content": "Hello! Import some files or ingest a GitHub repo and we can get started.",
-            }
-        ]
+    #####################
+    # Advanced Settings #
+    #####################
+
+    if "advanced" not in st.session_state:
+        st.session_state["advanced"] = False
+
+    if "top_k" not in st.session_state:
+        st.session_state["top_k"] = None
+
+    if "embedding_model" not in st.session_state:
+        st.session_state["embedding_model"] = None
+
+    if "chunk_size" not in st.session_state:
+        st.session_state["chunk_size"] = None
diff --git a/components/sidebar.py b/components/sidebar.py
@@ -8,18 +8,13 @@
 
 def sidebar():
     with st.sidebar:
-        tab1, tab2, tab3, tab4 = st.sidebar.tabs(
-            ["File Upload", "GitHub Repo", "Settings", "About"]
-        )
+        tab1, tab2, tab3 = st.sidebar.tabs(["My Files", "Settings", "About"])
 
         with tab1:
             file_upload()
 
         with tab2:
-            github_repo()
-
-        with tab3:
             settings()
 
-        with tab4:
+        with tab3:
             about()
diff --git a/components/tabs/about.py b/components/tabs/about.py
@@ -8,27 +8,52 @@ def about():
     st.caption(f"Developed by Jon Fairbanks &copy; {datetime.now().year}")
     st.write("")
 
-    st.subheader("Links")
-    st.markdown(
-        """
-        * [GitHub](https://github.com/jonfairbanks/local-rag)
-        * [Docker Hub](#)
+    links_html = """
+    <ul style="list-style-type:none; padding-left:0;">
+        <li>
+            <a href="https://github.com/jonfairbanks/local-rag" style="color: grey;">GitHub</a>
+        </li>
+        <li>
+            <a href="https://hub.docker.com/r/jonfairbanks/local-rag" style="color: grey;">Docker Hub</a>
+        </li>
+    </ul>
     """
-    )
 
-    st.subheader("Resources")
-    st.markdown(
-        """
-        * [What is RAG?](https://blogs.nvidia.com/blog/what-is-retrieval-augmented-generation/)
-        * [What are embeddings?](https://aws.amazon.com/what-is/embeddings-in-machine-learning/)
-        * [Embeddings Leaderboard](https://huggingface.co/spaces/mteb/leaderboard)
+    resources_html = """
+    <ul style="list-style-type:none; padding-left:0;">
+        <li>
+            <a href="https://blogs.nvidia.com/blog/what-is-retrieval-augmented-generation/" style="color: grey;">
+                What is RAG?
+            </a>
+        </li>
+        <li>
+            <a href="https://aws.amazon.com/what-is/embeddings-in-machine-learning/" style="color: grey;">
+                What are embeddings?
+            </a>
+        </li>
+    </ul>
     """
-    )
 
-    st.subheader("Help")
-    st.markdown(
-        """
-        * [Bug Reports](https://github.com/jonfairbanks/local-rag/issues)
-        * [Feature Requests](https://github.com/jonfairbanks/local-rag/discussions/new?category=ideas)
+    help_html = """
+    <ul style="list-style-type:none; padding-left:0;">
+        <li>
+            <a href="https://github.com/jonfairbanks/local-rag/issues" style="color: grey;">
+                Bug Reports
+            </a>
+        </li>
+        <li>
+            <a href="https://github.com/jonfairbanks/local-rag/discussions/new?category=ideas" style="color: grey;">
+                Feature Requests
+            </a>
+        </li>
+    </ul>
     """
-    )
+
+    st.subheader("Links")
+    st.markdown(links_html, unsafe_allow_html=True)
+
+    st.subheader("Resources")
+    st.markdown(resources_html, unsafe_allow_html=True)
+
+    st.subheader("Help")
+    st.markdown(help_html, unsafe_allow_html=True)
diff --git a/components/tabs/file_upload.py b/components/tabs/file_upload.py
@@ -8,41 +8,52 @@
 
 
 def file_upload():
-    st.title("Directly import local files")
-    st.caption("Convert your local files to embeddings for utilization during chat")
+    st.title("Directly import your files")
+    st.caption("Convert your files to embeddings for utilization during chat")
     st.write("")
-    if(st.session_state['selected_model'] is not None):
+
+    if st.session_state["selected_model"] is not None:
         uploaded_files = st.file_uploader(
             "Select Files",
             accept_multiple_files=True,
             type=("csv", "docx", "epub", "ipynb", "json", "md", "pdf", "ppt", "pptx",),
         )
     else:
-        st.info("Please configure Ollama settings before proceeding!")
+        st.warning("Please configure Ollama settings before proceeding!", icon="⚠️")
         uploaded_files = st.file_uploader(
             "Select Files",
             accept_multiple_files=True,
             type=("csv", "docx", "epub", "ipynb", "json", "md", "pdf", "ppt", "pptx",),
-            disabled=True
+            disabled=True,
         )
 
     if len(uploaded_files) > 0:
         st.session_state["file_list"] = uploaded_files
 
         with st.status("Preparing your data...", expanded=True) as status:
             error = None
+
+            ######################
+            # Save Files to Disk #
+            ######################
+
             st.caption("Uploading Files Locally")
-            # Save the files to disk
             for uploaded_file in uploaded_files:
                 with st.spinner(f"Processing {uploaded_file.name}..."):
                     save_dir = os.getcwd() + "/data"
                     func.save_uploaded_file(uploaded_file, save_dir)
 
             st.caption("Loading Embedding Model")
-            # Create llama-index service-context to use local LLMs and embeddings
+
+            ######################################
+            # Create Llama-Index service-context #
+            # to use local LLMs and embeddings   #
+            ######################################
+
             try:
                 llm = ollama.create_ollama_llm(
-                    st.session_state['selected_model'], st.session_state['ollama_endpoint'],
+                    st.session_state["selected_model"],
+                    st.session_state["ollama_endpoint"],
                 )
                 # resp = llm.complete("Hello!")
                 # print(resp)
@@ -51,21 +62,33 @@ def file_upload():
                 print(f"Setting up Service Context failed: {err}")
                 error = err
 
+            #######################################
+            # Load files from the data/ directory #
+            #######################################
+
             st.caption("Processing File Data")
             try:
                 documents = llama_index.load_documents(save_dir)
-                st.session_state.documents = documents
+                st.session_state["documents"] = documents
             except Exception as err:
                 print(f"Document Load Error: {err}")
                 error = err
 
+            ###########################################
+            # Create an index from ingested documents #
+            ###########################################
+
             st.caption("Creating File Index")
             try:
                 llama_index.create_query_engine(documents, service_context)
             except Exception as err:
                 print(f"Index Creation Error: {err}")
                 error = err
 
+            #####################
+            # Show Final Status #
+            #####################
+
             if error is not None:
                 status.update(
                     label="File processing failed.", state="error", expanded=True,
@@ -78,6 +101,5 @@ def file_upload():
                     expanded=False,
                 )
 
-    # st.caption(
-    #     "Although any uploads are supported, you will get the best results with: _csv, docx, epub, ipynb, md, pdf, ppt, pptx_"
-    # )
+    with st.expander("GitHub Repo", expanded=False):
+        st.write(":grey[Coming Soon&trade;]")
diff --git a/components/tabs/settings.py b/components/tabs/settings.py
@@ -4,8 +4,6 @@
 
 import utils.ollama as ollama
 
-# import utils.llama_index as llama_index
-
 from datetime import datetime
 
 
@@ -20,10 +18,11 @@ def settings():
             "Ollama Endpoint",
             key="ollama_endpoint",
             placeholder="http://localhost:11434",
-            value=st.session_state["ollama_endpoint"],
             on_change=ollama.get_models,
         )
-        st.selectbox("Model", st.session_state.ollama_models, key="selected_model")
+        st.selectbox(
+            "Model", st.session_state["ollama_models"], key="selected_model",
+        )
         st.button(
             "Refresh", on_click=ollama.get_models,
         )
@@ -38,12 +37,13 @@ def settings():
             st.text_input(
                 "System Prompt",
                 value="You are a sophisticated virtual assistant designed to assist users in comprehensively understanding and extracting insights from a wide range of documents at their disposal. Your expertise lies in tackling complex inquiries and providing insightful analyses based on the information contained within these documents.",
+                key="system_prompt",
                 disabled=True,
             )
-            option = st.selectbox(
+            st.selectbox(
                 "Chat Mode",
                 ("Best", "Condense Question", "Context", "Condense + Context"),
-                help="Sets the [Llama-Index Chat Mode](https://docs.llamaindex.ai/en/stable/module_guides/deploying/chat_engines/usage_pattern.html#available-chat-modes) to be used.",
+                help="Sets the [Llama-Index Chat Mode](https://docs.llamaindex.ai/en/stable/module_guides/deploying/chat_engines/usage_pattern.html#available-chat-modes) used when creating the Query Engine.",
                 disabled=True,
             )
             st.write("")
@@ -60,11 +60,14 @@ def settings():
             disabled=True,
         )
         if st.session_state["advanced"] == True:
+            st.caption(
+                "View the [Embeddings Leaderboard](https://huggingface.co/spaces/mteb/leaderboard)"
+            )
             st.text_input(
                 "Chunk Size",
                 help="This should not exceed the value provided by your embedding model.",
                 key="chunk_size",
-                placeholder="1024",
+                placeholder="512",
                 value=st.session_state["chunk_size"],
                 disabled=True,
             )

diff --git a/utils/llama_index.py b/utils/llama_index.py
@@ -66,7 +66,7 @@ def load_documents(data_dir: str):
     try:
         files = SimpleDirectoryReader(input_dir=data_dir, recursive=True)
         documents = files.load_data(files)
-        #print(f"Loaded {len(documents):,} documents")
+        # print(f"Loaded {len(documents):,} documents")
         return documents
     except Exception as err:
         print(f"Error creating data index: {err}")