Update src/streamlit_app.py

src/streamlit_app.py

@@ -3,26 +3,21 @@ import numpy as np
 import pandas as pd
 import streamlit as st
 
-"""
-# Welcome to Streamlit!
-
-Edit `/streamlit_app.py` to customize this app to your heart's desire :heart:.
-If you have any questions, checkout our [documentation](https://docs.streamlit.io) and [community
-forums](https://discuss.streamlit.io).
 
-In the meantime, below is an example of what you can do with just a few lines of code:
-"""
+import streamlit as st
+import pandas as pd
+import rag  # Import the rag module
+import os   # Import os for file path handling
 import streamlit as st
 import pandas as pd
 import rag  # Import the rag module
 import os   # Import os for file path handling
 
-def main():
-    st.set_page_config(layout="wide", page_title="Communication Board")
 
-    # --- Session State Initialization ---
-    # Check specifically for 'assistant' to ensure it's initialized
-    if 'assistant' not in st.session_state:
+st.set_page_config(layout="wide", page_title="Communication Board")
+# --- Session State Initialization ---
+# Check specifically for 'assistant' to ensure it's initialized
+if 'assistant' not in st.session_state:
         st.session_state.current_page = "main"
         st.session_state.show_custom_words = False
         st.session_state.custom_words = []
@@ -36,7 +31,7 @@ def main():
         st.session_state.text_output = ""
 
     # --- Theme Colors ---
-    theme_colors = {
+theme_colors = {
         "Default": {
             "bg": "#FFFFFF", "text": "#000000",
             "pronoun": "#FFFF99", "verb": "#CCFFCC", 
@@ -66,40 +61,38 @@ def main():
             "misc": "#8B1A1A"
         }
     }
-    colors = theme_colors[st.session_state.theme]
-
-    # --- Helper Function to Initialize Assistant (Adapted from previous main.py) ---
-    @st.cache_resource # Cache the assistant
-    def initialize_assistant(doc_path):
-        """Initializes the AACAssistant."""
+colors = theme_colors[st.session_state.theme]
+# --- Helper Function to Initialize Assistant (Adapted from previous main.py) ---
+@st.cache_resource # Cache the assistant
+def initialize_assistant(doc_path):
+    """Initializes the AACAssistant."""
         # Create a dummy document if it doesn't exist for demonstration
-        if not os.path.exists(doc_path):
-            st.sidebar.warning(f"Doc '{os.path.basename(doc_path)}' not found. Creating dummy.")
-            try:
-                with open(doc_path, "w") as f:
-                    f.write("""
-                    I grew up in Seattle and love the rain.
-                    My favorite hobby is playing chess.
-                    I have a dog named Max.
-                    I studied computer science.
-                    I enjoy sci-fi movies.
-                    """)
-            except Exception as e:
-                st.sidebar.error(f"Failed to create dummy doc: {e}")
-                return None
+    if not os.path.exists(doc_path):
+        st.sidebar.warning(f"Doc '{os.path.basename(doc_path)}' not found. Creating dummy.")
         try:
-            assistant = rag.AACAssistant(doc_path)
-            st.sidebar.success("AAC Assistant Initialized.")
-            return assistant
+            with open(doc_path, "w") as f:
+                f.write("""
+                I grew up in Seattle and love the rain.
+                My favorite hobby is playing chess.
+                I have a dog named Max.
+                I studied computer science.
+                I enjoy sci-fi movies.
+                """)
         except Exception as e:
+            st.sidebar.error(f"Failed to create dummy doc: {e}")
+            return None
+    try:
+        assistant = rag.AACAssistant(doc_path)
+        st.sidebar.success("AAC Assistant Initialized.")
+        return assistant
+    except Exception as e:
             st.sidebar.error(f"Error initializing AAC Assistant: {e}")
             st.sidebar.error("Ensure Ollama/LM Studio running.")
             return None
 
-    DEFAULT_DOCUMENT_PATH = "aac_user_experiences.txt"
-
-    # --- CSS Styling ---
-    css = f"""
+DEFAULT_DOCUMENT_PATH = "aac_user_experiences.txt"
+# --- CSS Styling ---
+css = f"""
     <style>
         .big-font {{
             font-size:{st.session_state.text_size}px !important;
@@ -181,29 +174,28 @@ def main():
     """
 
     # --- JS for Button Coloring (no delay, no setTimeout) ---
-    js = """
-    <script>
-        function colorButtons() {
-            const buttons = document.querySelectorAll('button[id^="key_"]');
-            buttons.forEach(button => {
-                const id = button.id;
-                const parts = id.split('_');
-                if (parts.length >= 4) {
-                    const category = parts[3];
-                    button.classList.add('btn-' + category);
-                }
-            });
-        }
-        document.addEventListener('DOMContentLoaded', colorButtons);
-        new MutationObserver(colorButtons).observe(document.body, { childList: true, subtree: true });
-    </script>
-    """
-
-    st.markdown(css, unsafe_allow_html=True)
-    st.markdown(js, unsafe_allow_html=True)
+js = """
+<script>
+    function colorButtons() {
+        const buttons = document.querySelectorAll('button[id^="key_"]');
+        buttons.forEach(button => {
+            const id = button.id;
+            const parts = id.split('_');
+            if (parts.length >= 4) {
+                const category = parts[3];
+                button.classList.add('btn-' + category);
+            }
+        });
+    }
+    document.addEventListener('DOMContentLoaded', colorButtons);
+    new MutationObserver(colorButtons).observe(document.body, { childList: true, subtree: true });
+</script>
+"""
+st.markdown(css, unsafe_allow_html=True)
+st.markdown(js, unsafe_allow_html=True)
 
     # --- Keyboard Layout ---
-    layout = [
+layout = [
         [
             {"word": "I", "category": "pronoun"},
             {"word": "am", "category": "verb"},
@@ -303,101 +295,98 @@ def main():
     ]
 
     # --- Add Custom Words ---
-    if st.session_state.custom_words:
-        custom_row = []
-        for idx, word_info in enumerate(st.session_state.custom_words):
-            word = word_info["word"]
-            category = word_info["category"]
-            custom_row.append({"word": f"custom_{idx}_{word}", "display": word, "category": category})
-            if len(custom_row) == 10:
-                layout.append(custom_row)
-                custom_row = []
-        if custom_row:
-            while len(custom_row) < 10:
-                custom_row.append({"word": "", "category": "misc"})
+if st.session_state.custom_words:
+    custom_row = []
+    for idx, word_info in enumerate(st.session_state.custom_words):
+        word = word_info["word"]
+        category = word_info["category"]
+        custom_row.append({"word": f"custom_{idx}_{word}", "display": word, "category": category})
+        if len(custom_row) == 10:
             layout.append(custom_row)
+            custom_row = []
+    if custom_row:
+        while len(custom_row) < 10:
+            custom_row.append({"word": "", "category": "misc"})
+        layout.append(custom_row)
 
     # --- Initialize Assistant ---
     # Attempt initialization only once or if it failed previously
-    if st.session_state.assistant is None: # This check is now safe
-        st.session_state.assistant = initialize_assistant(DEFAULT_DOCUMENT_PATH)
-
-    # --- Output Box (move to top, before keyboard) ---
-    st.title("Communication Board")
+if st.session_state.assistant is None: # This check is now safe
+    st.session_state.assistant = initialize_assistant(DEFAULT_DOCUMENT_PATH)
+# --- Output Box (move to top, before keyboard) ---
+st.title("Communication Board")
     # Use st.text_area directly for input and display, bound to session state
-    st.session_state.text_output = st.text_area(
+st.session_state.text_output = st.text_area(
         "Compose Message:", value=st.session_state.text_output, height=100, key="main_text_output"
     )
 
     # --- Keyboard Rendering (no delay, instant update) ---
-    send_action = False # Flag for SEND button
-    def add_to_output(word):
-        if word == "SPACE":
+send_action = False # Flag for SEND button
+def add_to_output(word):
+    if word == "SPACE":
+        st.session_state.text_output += " "
+    elif word in [".", "?", "!", ",", "-"]:
+        st.session_state.text_output += word
+    elif word.isdigit() or (len(word) == 1 and word.isalpha()):
+        st.session_state.text_output += word
+    else:
+        if st.session_state.text_output and not st.session_state.text_output.endswith(" "):
             st.session_state.text_output += " "
-        elif word in [".", "?", "!", ",", "-"]:
-            st.session_state.text_output += word
-        elif word.isdigit() or (len(word) == 1 and word.isalpha()):
-            st.session_state.text_output += word
-        else:
-            if st.session_state.text_output and not st.session_state.text_output.endswith(" "):
-                st.session_state.text_output += " "
-            st.session_state.text_output += word
+        st.session_state.text_output += word
 
-    st.markdown("### Communication Keyboard")
-    for row_idx, row in enumerate(layout):
-        cols = st.columns(len(row))
-        for col_idx, item in enumerate(cols):
-            word_info = row[col_idx]
-            if "word" not in word_info or word_info["word"] == "":
-                continue
-            word = word_info["word"]
-            category = word_info["category"]
-            display = word_info.get("display", word)
-            key = f"key_{row_idx}_{col_idx}_{category}_{word}"
-            def make_callback(w=word, d=display):
-                def cb():
-                    if w.startswith("custom_") or w.startswith("letter_"):
-                        add_to_output(d)
-                    else:
-                        add_to_output(w)
-                return cb
-            with cols[col_idx]:
-                st.button(
+st.markdown("### Communication Keyboard")
+for row_idx, row in enumerate(layout):
+    cols = st.columns(len(row))
+    for col_idx, item in enumerate(cols):
+        word_info = row[col_idx]
+        if "word" not in word_info or word_info["word"] == "":
+            continue
+        word = word_info["word"]
+        category = word_info["category"]
+        display = word_info.get("display", word)
+        key = f"key_{row_idx}_{col_idx}_{category}_{word}"
+        def make_callback(w=word, d=display):
+            def cb():
+                if w.startswith("custom_") or w.startswith("letter_"):
+                    add_to_output(d)
+                else:
+                    add_to_output(w)
+            return cb
+        with cols[col_idx]:
+            st.button(
                     display if word != "SPACE" else "␣",
                     key=key,
                     on_click=make_callback()
                 )
 
     # --- Control Buttons ---
-    col1, col2, col3, col4 = st.columns(4)
-    with col1:
-        if st.button("CLEAR", key="clear_btn", help="Clear all text", use_container_width=True, type="primary"):
-            st.session_state.text_output = ""
-            st.rerun() # Rerun to reflect change
-    with col2:
-        if st.button("SPEAK", key="speak_btn", help="Speak the current text", use_container_width=True, type="primary"):
-            if st.session_state.text_output:
-                st.toast(f"Speaking: {st.session_state.text_output}", icon="🔊")
-        if st.button("⌫ DEL", key="backspace", help="Delete last character", use_container_width=True):
-            if st.session_state.text_output:
-                st.session_state.text_output = st.session_state.text_output[:-1]
-    with col4:
-        if st.button("⌫ WORD", key="backspace_word", help="Delete last word", use_container_width=True):
-            if st.session_state.text_output:
-                words = st.session_state.text_output.rstrip().split()
+col1, col2, col3, col4 = st.columns(4)
+with col1:
+    if st.button("CLEAR", key="clear_btn", help="Clear all text", use_container_width=True, type="primary"):
+        st.session_state.text_output = ""
+        st.rerun() # Rerun to reflect change
+with col2:
+    if st.button("SPEAK", key="speak_btn", help="Speak the current text", use_container_width=True, type="primary"):
+        if st.session_state.text_output:
+            st.toast(f"Speaking: {st.session_state.text_output}", icon="🔊")
+    if st.button("⌫ DEL", key="backspace", help="Delete last character", use_container_width=True):
+        if st.session_state.text_output:
+            st.session_state.text_output = st.session_state.text_output[:-1]
+with col4:
+    if st.button("⌫ WORD", key="backspace_word", help="Delete last word", use_container_width=True):
+        if st.session_state.text_output:
+            words = st.session_state.text_output.rstrip().split()
+            if words:
+                words.pop()
+                st.session_state.text_output = " ".join(words)
                 if words:
-                    words.pop()
-                    st.session_state.text_output = " ".join(words)
-                    if words:
-                        st.session_state.text_output += " "
-    with col3: # Use the 3rd column for SEND
-        if st.button("SEND", key="send_btn", help="Send message to assistant", use_container_width=True, type="primary"):
-             if st.session_state.text_output:
-                 send_action = True # Set flag to process sending
-
-
-    # --- Settings Sidebar ---
-    with st.sidebar:
+                    st.session_state.text_output += " "
+with col3: # Use the 3rd column for SEND
+    if st.button("SEND", key="send_btn", help="Send message to assistant", use_container_width=True, type="primary"):
+         if st.session_state.text_output:
+             send_action = True # Set flag to process sending
+# --- Settings Sidebar ---
+with st.sidebar:
         # --- Settings Section Commented Out ---
         # st.title("Settings")
         # st.subheader("Interface")
@@ -448,34 +437,31 @@ def main():
         #                 st.rerun()
         # --- End of Settings Section Commented Out ---
 
-        st.divider()
-        # --- Chat History Display in Sidebar ---
-        st.subheader("Conversation")
-        chat_container = st.container(height=400) # Fixed height container
-        with chat_container:
-            for message in st.session_state.messages:
-                with st.chat_message(message["role"]):
-                    st.markdown(message["content"])
+    st.divider()
+    # --- Chat History Display in Sidebar ---
+    st.subheader("Conversation")
+    chat_container = st.container(height=400) # Fixed height container
+    with chat_container:
+        for message in st.session_state.messages:
+            with st.chat_message(message["role"]):
+                st.markdown(message["content"])
 
-    # --- Process SEND Action ---
-    if send_action and st.session_state.assistant:
-        user_message = st.session_state.text_output
-        st.session_state.messages.append({"role": "user", "content": user_message})
+# --- Process SEND Action ---
+if send_action and st.session_state.assistant:
+    user_message = st.session_state.text_output
+    st.session_state.messages.append({"role": "user", "content": user_message})
 
         # Process with AACAssistant
-        try:
-            # Get the response from the AACAssistant
-            response = st.session_state.assistant.process_query(user_message)
-            # Add assistant response to chat history
-            st.session_state.messages.append({"role": "assistant", "content": response})
-        except Exception as e:
-            error_message = f"An error occurred: {e}"
-            st.error(error_message) # Show error in main area
-            st.session_state.messages.append({"role": "assistant", "content": f"*Error processing: {error_message}*"})
-
-        # Clear the board's text area after sending
-        st.session_state.text_output = ""
-        st.rerun() # Rerun to update chat and clear board
+    try:
+        # Get the response from the AACAssistant
+        response = st.session_state.assistant.process_query(user_message)
+        # Add assistant response to chat history
+        st.session_state.messages.append({"role": "assistant", "content": response})
+    except Exception as e:
+        error_message = f"An error occurred: {e}"
+        st.error(error_message) # Show error in main area
+        st.session_state.messages.append({"role": "assistant", "content": f"*Error processing: {error_message}*"})
+    # Clear the board's text area after sending
+    st.session_state.text_output = ""
+    st.rerun() # Rerun to update chat and clear board
 
-if __name__ == "__main__":
-    main()