Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,31 +1,26 @@
 import streamlit as st
 
-import altair as alt # Keep if used, otherwise remove
-import numpy as np   # Keep if used, otherwise remove
+import altair as alt
+import numpy as np
 import pandas as pd
-import rag  # Import your rag module
+import rag
 import os
 
-# --- Environment Variable Setup for Caching (run once at the top) ---
 cache_dir_base = "/app/.cache_app"
-if not os.path.exists(cache_dir_base): # Ensure base exists, useful if Dockerfile step was missed
+if not os.path.exists(cache_dir_base):
     try:
         os.makedirs(cache_dir_base, exist_ok=True)
     except Exception as e:
         st.error(f"Failed to create base cache directory {cache_dir_base}: {e}")
 
-# Set environment variables for Hugging Face libraries if not already set by Dockerfile
-# Note: Setting them here is a fallback; best set in Dockerfile or Space secrets for build-time tools.
 os.environ['HF_HOME'] = os.environ.get('HF_HOME', os.path.join(cache_dir_base, 'huggingface_hub'))
 os.environ['TRANSFORMERS_CACHE'] = os.environ.get('TRANSFORMERS_CACHE', os.path.join(cache_dir_base, 'transformers'))
 os.environ['SENTENCE_TRANSFORMERS_HOME'] = os.environ.get('SENTENCE_TRANSFORMERS_HOME', os.path.join(cache_dir_base, 'sentence_transformers'))
 os.environ['PIP_CACHE_DIR'] = os.environ.get('PIP_CACHE_DIR', os.path.join(cache_dir_base, 'pip'))
 
-# --- Page Configuration (should be the first Streamlit command) ---
 st.set_page_config(layout="wide", page_title="Communication Board")
 
-# --- Session State Initialization ---
-if 'assistant_initialized_flag' not in st.session_state: # Use a flag to track if full init ran
+if 'assistant_initialized_flag' not in st.session_state:
     st.session_state.current_page = "main"
     st.session_state.show_custom_words = False
     st.session_state.custom_words = []
@@ -34,67 +29,31 @@ if 'assistant_initialized_flag' not in st.session_state: # Use a flag to track i
     st.session_state.speech_rate = 1.0
     st.session_state.voice_option = "Default Voice"
     st.session_state.messages = []
-    st.session_state.assistant = None # Will be initialized by initialize_assistant
+    st.session_state.assistant = None
     st.session_state.text_output = ""
-    st.session_state.assistant_initialized_flag = True # Mark that initial session state setup is done
+    st.session_state.assistant_initialized_flag = True
 
-
-# --- Theme Colors ---
 theme_colors = {
-    "Default": {
-        "bg": "#FFFFFF", "text": "#000000",
-        "pronoun": "#FFFF99", "verb": "#CCFFCC", 
-        "question": "#CCCCFF", "common": "#FFCC99",
-        "preposition": "#99CCFF", "descriptive": "#CCFF99",
-        "misc": "#FFB6C1"
-    },
-    "High Contrast": {
-        "bg": "#FFFFFF", "text": "#000000",
-        "pronoun": "#FFFF00", "verb": "#00FF00",
-        "question": "#0000FF", "common": "#FF6600",
-        "preposition": "#00CCFF", "descriptive": "#66FF33",
-        "misc": "#FF3366"
-    },
-    "Pastel": {
-        "bg": "#F8F8FF", "text": "#333333",
-        "pronoun": "#FFEFD5", "verb": "#E0FFFF",
-        "question": "#D8BFD8", "common": "#FFE4B5",
-        "preposition": "#B0E0E6", "descriptive": "#F0FFF0",
-        "misc": "#FFF0F5"
-    },
-    "Dark Mode": {
-        "bg": "#333333", "text": "#FFFFFF",
-        "pronoun": "#8B8B00", "verb": "#006400",
-        "question": "#00008B", "common": "#8B4500",
-        "preposition": "#00688B", "descriptive": "#698B22",
-        "misc": "#8B1A1A"
-    }
+    "Default": {"bg": "#FFFFFF", "text": "#000000", "pronoun": "#FFFF99", "verb": "#CCFFCC", "question": "#CCCCFF", "common": "#FFCC99", "preposition": "#99CCFF", "descriptive": "#CCFF99", "misc": "#FFB6C1"},
+    "High Contrast": {"bg": "#FFFFFF", "text": "#000000", "pronoun": "#FFFF00", "verb": "#00FF00", "question": "#0000FF", "common": "#FF6600", "preposition": "#00CCFF", "descriptive": "#66FF33", "misc": "#FF3366"},
+    "Pastel": {"bg": "#F8F8FF", "text": "#333333", "pronoun": "#FFEFD5", "verb": "#E0FFFF", "question": "#D8BFD8", "common": "#FFE4B5", "preposition": "#B0E0E6", "descriptive": "#F0FFF0", "misc": "#FFF0F5"},
+    "Dark Mode": {"bg": "#333333", "text": "#FFFFFF", "pronoun": "#8B8B00", "verb": "#006400", "question": "#00008B", "common": "#8B4500", "preposition": "#00688B", "descriptive": "#698B22", "misc": "#8B1A1A"}
 }
 colors = theme_colors[st.session_state.theme]
 
-# --- Helper Function to Initialize Assistant ---
-@st.cache_resource # Cache the assistant object
+@st.cache_resource
 def initialize_assistant(doc_path):
-    """Initializes the AACAssistant."""
-    # Create a dummy document if it doesn't exist for demonstration
-    # This part is fine, but ensure doc_path is correct relative to where streamlit_app.py is.
-    # If streamlit_app.py is in src/, and aac_user_experiences.txt is also in src/, then "aac_user_experiences.txt" is enough.
-    # If streamlit_app.py is at /app and text file is /app/src/aac_user_experiences.txt, then "src/aac_user_experiences.txt" is correct.
     if not os.path.exists(doc_path):
         st.sidebar.warning(f"Doc '{os.path.basename(doc_path)}' not found at '{doc_path}'. Creating dummy.")
         try:
-            os.makedirs(os.path.dirname(doc_path), exist_ok=True) # Ensure directory exists
+            parent_dir = os.path.dirname(doc_path)
+            if parent_dir: # Ensure parent_dir is not empty if doc_path is just a filename
+                 os.makedirs(parent_dir, exist_ok=True)
             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.
-                """)
+                f.write("I grew up in Seattle and love the rain.\nMy favorite hobby is playing chess.\nI have a dog named Max.\nI studied computer science.\nI enjoy sci-fi movies.\n")
         except Exception as e:
             st.sidebar.error(f"Failed to create dummy doc at '{doc_path}': {e}")
-            return None 
+            return None
     try:
         st.sidebar.info("Initializing AAC Assistant... This may take some time.")
         assistant = rag.AACAssistant(doc_path)
@@ -103,331 +62,183 @@ def initialize_assistant(doc_path):
     except Exception as e:
         st.sidebar.error(f"Fatal Error Initializing AAC Assistant: {e}")
         st.sidebar.error("The assistant could not be started. Please check the logs for details.")
-        # import traceback # Optional: show full traceback in sidebar for debugging
-        # st.sidebar.text(traceback.format_exc())
         return None
 
-DEFAULT_DOCUMENT_PATH = "src/aac_user_experiences.txt" # Make sure this path is correct
+DEFAULT_DOCUMENT_PATH = "src/aac_user_experiences.txt"
 
-# --- Initialize Assistant (runs only if assistant is None in session_state) ---
 if st.session_state.assistant is None:
     st.session_state.assistant = initialize_assistant(DEFAULT_DOCUMENT_PATH)
 
-# --- CSS Styling ---
-# Your CSS remains largely the same, ensure variables like st.session_state.text_size and colors are defined before this.
 css = f"""
-    <style>
-        /* ... your existing CSS using {st.session_state.text_size} and {colors["..."]} ... */
-        .big-font {{
-            font-size:{st.session_state.text_size}px !important;
-            text-align: center;
-        }}
-        .output-box {{
-            border: 2px solid #ddd;
-            border-radius: 5px;
-            padding: 15px;
-            min-height: 100px;
-            background-color: white;
-            margin-bottom: 15px;
-            font-size: {st.session_state.text_size}px;
-            box-shadow: 0 2px 4px rgba(0,0,0,0.1);
-        }}
-        div[data-testid="stHorizontalBlock"] {{
-            gap: 5px;
-        }}
-        section[data-testid="stSidebar"] {{
-            width: 20rem !important; /* Ensure sidebar width is applied */
-            background-color: {colors["bg"]};
-        }}
-        body {{
-            background-color: {colors["bg"]};
-            color: {colors["text"]};
-        }}
-        .stButton>button {{
-            width: 100%;
-            height: 60px;
-            font-size: {max(16, st.session_state.text_size - 6)}px;
-            font-weight: bold;
-            white-space: normal;
-            word-wrap: break-word; /* Ensure long words in buttons wrap */
-            padding: 0px;
-            transition: transform 0.1s ease;
-        }}
-        .stButton>button:hover {{
-            filter: brightness(95%);
-            transform: scale(1.03);
-            box-shadow: 0 2px 3px rgba(0,0,0,0.1);
-        }}
-        .control-button {{ /* This class was defined but not used on control buttons */
-            height: 80px !important;
-            font-size: {max(18, st.session_state.text_size - 4)}px !important;
-        }}
-        .btn-pronoun {{ background-color: {colors["pronoun"]} !important; border: 1px solid #888 !important; }}
-        .btn-verb {{ background-color: {colors["verb"]} !important; border: 1px solid #888 !important; }}
-        .btn-question {{ background-color: {colors["question"]} !important; border: 1px solid #888 !important; }}
-        .btn-common {{ background-color: {colors["common"]} !important; border: 1px solid #888 !important; }}
-        .btn-preposition {{ background-color: {colors["preposition"]} !important; border: 1px solid #888 !important; }}
-        .btn-descriptive {{ background-color: {colors["descriptive"]} !important; border: 1px solid #888 !important; }}
-        .btn-misc {{ background-color: {colors["misc"]} !important; border: 1px solid #888 !important; }}
-        .sidebar .stChatMessage {{
-            background-color: {colors.get('bg', '#FFFFFF')};
-            border-radius: 8px;
-        }}
-    </style>
-    """
+<style>
+.big-font {{ font-size:{st.session_state.text_size}px !important; text-align: center; }}
+.output-box {{ border: 2px solid #ddd; border-radius: 5px; padding: 15px; min-height: 100px; background-color: white; margin-bottom: 15px; font-size: {st.session_state.text_size}px; box-shadow: 0 2px 4px rgba(0,0,0,0.1); }}
+div[data-testid="stHorizontalBlock"] {{ gap: 5px; }}
+section[data-testid="stSidebar"] {{ width: 20rem !important; background-color: {colors["bg"]}; }}
+body {{ background-color: {colors["bg"]}; color: {colors["text"]}; }}
+.stButton>button {{ width: 100%; height: 60px; font-size: {max(16, st.session_state.text_size - 6)}px; font-weight: bold; white-space: normal; word-wrap: break-word; padding: 0px; transition: transform 0.1s ease; }}
+.stButton>button:hover {{ filter: brightness(95%); transform: scale(1.03); box-shadow: 0 2px 3px rgba(0,0,0,0.1); }}
+.control-button {{ height: 80px !important; font-size: {max(18, st.session_state.text_size - 4)}px !important; }}
+.btn-pronoun {{ background-color: {colors["pronoun"]} !important; border: 1px solid #888 !important; }}
+.btn-verb {{ background-color: {colors["verb"]} !important; border: 1px solid #888 !important; }}
+.btn-question {{ background-color: {colors["question"]} !important; border: 1px solid #888 !important; }}
+.btn-common {{ background-color: {colors["common"]} !important; border: 1px solid #888 !important; }}
+.btn-preposition {{ background-color: {colors["preposition"]} !important; border: 1px solid #888 !important; }}
+.btn-descriptive {{ background-color: {colors["descriptive"]} !important; border: 1px solid #888 !important; }}
+.btn-misc {{ background-color: {colors["misc"]} !important; border: 1px solid #888 !important; }}
+.sidebar .stChatMessage {{ background-color: {colors.get('bg', '#FFFFFF')}; border-radius: 8px; }}
+</style>
+"""
 st.markdown(css, unsafe_allow_html=True)
 
-# --- JS for Button Coloring ---
-# Your JS remains the same
 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) { // Ensure category is present
-                const category = parts[3];
-                // Remove old category classes before adding new one
-                button.className = button.className.replace(/btn-[a-z]+/g, '').trim();
-                button.classList.add('btn-' + category);
-            }
-        });
-    }
-    // Run on initial load
-    document.addEventListener('DOMContentLoaded', colorButtons);
-
-    // Re-run when Streamlit updates the DOM (e.g., after button clicks that cause re-render)
-    // Using a more robust approach for Streamlit that re-evaluates when elements are added
-    const streamlitDoc = window.parent.document;
-    const observer = new MutationObserver((mutationsList, observer) => {
-        for(const mutation of mutationsList) {
-            if (mutation.type === 'childList' && mutation.addedNodes.length > 0) {
-                colorButtons(); // Re-apply colors if new nodes are added
-                // We could be more specific here to only run if relevant buttons are added
-            }
+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.className = button.className.replace(/btn-[a-z]+/g, '').trim();
+            button.classList.add('btn-' + category);
         }
     });
-    observer.observe(streamlitDoc.body, { childList: true, subtree: true });
-
-    // Initial call in case DOM is already ready from a script rerun
-    colorButtons();
+}
+document.addEventListener('DOMContentLoaded', colorButtons);
+const streamlitDoc = window.parent.document;
+const observer = new MutationObserver((mutationsList, observer) => {
+    for(const mutation of mutationsList) {
+        if (mutation.type === 'childList' && mutation.addedNodes.length > 0) {
+            colorButtons();
+        }
+    }
+});
+observer.observe(streamlitDoc.body, { childList: true, subtree: true });
+colorButtons();
 </script>
 """
 st.markdown(js, unsafe_allow_html=True)
 
-
-# --- Main UI Layout ---
 st.title("Communication Board")
 
-# Output Box (moved to top as per your original structure)
 st.session_state.text_output = st.text_area(
-    "Compose Message:", value=st.session_state.text_output, height=100, key="main_text_output_area" # Changed key to avoid conflict
+    "Compose Message:", value=st.session_state.get("text_output", ""), height=100, key="main_text_output_area"
 )
 
-# Keyboard Layout
 layout = [
-    # ... (your existing layout definition - keep as is) ...
-        [
-            {"word": "I", "category": "pronoun"}, {"word": "am", "category": "verb"},
-            {"word": "how", "category": "question"}, {"word": "what", "category": "question"},
-            {"word": "when", "category": "question"}, {"word": "where", "category": "question"},
-            {"word": "who", "category": "question"}, {"word": "why", "category": "question"},
-            {"word": "That", "category": "pronoun"}, {"word": "Please", "category": "common"}
-        ],
-        [
-            {"word": "me", "category": "pronoun"}, {"word": "are", "category": "verb"},
-            {"word": "is", "category": "verb"}, {"word": "was", "category": "verb"},
-            {"word": "will", "category": "verb"}, {"word": "help", "category": "verb"},
-            {"word": "need", "category": "verb"}, {"word": "want", "category": "verb"},
-            {"word": "thank you", "category": "common"}, {"word": "sorry", "category": "common"}
-        ],
-        [
-            {"word": "my", "category": "pronoun"}, {"word": "can", "category": "verb"},
-            {"word": "A", "category": "misc"}, {"word": "B", "category": "misc"},
-            {"word": "C", "category": "misc"}, {"word": "D", "category": "misc"},
-            {"word": "E", "category": "misc"}, {"word": "F", "category": "misc"},
-            {"word": "G", "category": "misc"}, {"word": "H", "category": "misc"}
-        ],
-        [
-            {"word": "it", "category": "pronoun"}, {"word": "did", "category": "verb"},
-            {"word": "letter_I", "category": "misc", "display": "I"}, {"word": "J", "category": "misc"},
-            {"word": "K", "category": "misc"}, {"word": "L", "category": "misc"},
-            {"word": "M", "category": "misc"}, {"word": "N", "category": "misc"},
-            {"word": "O", "category": "misc"}, {"word": "P", "category": "misc"}
-        ],
-        [
-            {"word": "they", "category": "pronoun"}, {"word": "do", "category": "verb"},
-            {"word": "Q", "category": "misc"}, {"word": "R", "category": "misc"},
-            {"word": "S", "category": "misc"}, {"word": "T", "category": "misc"},
-            {"word": "U", "category": "misc"}, {"word": "V", "category": "misc"},
-            {"word": "W", "category": "misc"}, {"word": "X", "category": "misc"}
-        ],
-        [
-            {"word": "we", "category": "pronoun"}, {"word": "Y", "category": "misc"},
-            {"word": "Z", "category": "misc"}, {"word": "1", "category": "misc"},
-            {"word": "2", "category": "misc"}, {"word": "3", "category": "misc"},
-            {"word": "4", "category": "misc"}, {"word": "5", "category": "misc"}, # Corrected, was 4
-            {"word": ".", "category": "misc"}, {"word": "?", "category": "misc"}
-        ],
-        [
-            {"word": "you", "category": "pronoun"}, {"word": "6", "category": "misc"},
-            {"word": "7", "category": "misc"}, {"word": "8", "category": "misc"},
-            {"word": "9", "category": "misc"}, {"word": "0", "category": "misc"},
-            {"word": "-", "category": "misc"}, {"word": "!", "category": "misc"},
-            {"word": ",", "category": "misc"}, {"word": "SPACE", "category": "misc"}
-        ],
-        [
-            {"word": "your", "category": "pronoun"}, {"word": "like", "category": "verb"},
-            {"word": "to", "category": "preposition"}, {"word": "with", "category": "preposition"},
-            {"word": "in", "category": "preposition"}, {"word": "the", "category": "misc"},
-            {"word": "and", "category": "misc"}, {"word": "but", "category": "misc"},
-            {"word": "not", "category": "descriptive"}, {"word": "yes", "category": "common"}
-        ]
+    [{"word": "I", "category": "pronoun"}, {"word": "am", "category": "verb"}, {"word": "how", "category": "question"}, {"word": "what", "category": "question"}, {"word": "when", "category": "question"}, {"word": "where", "category": "question"}, {"word": "who", "category": "question"}, {"word": "why", "category": "question"}, {"word": "That", "category": "pronoun"}, {"word": "Please", "category": "common"}],
+    [{"word": "me", "category": "pronoun"}, {"word": "are", "category": "verb"}, {"word": "is", "category": "verb"}, {"word": "was", "category": "verb"}, {"word": "will", "category": "verb"}, {"word": "help", "category": "verb"}, {"word": "need", "category": "verb"}, {"word": "want", "category": "verb"}, {"word": "thank you", "category": "common"}, {"word": "sorry", "category": "common"}],
+    [{"word": "my", "category": "pronoun"}, {"word": "can", "category": "verb"}, {"word": "A", "category": "misc"}, {"word": "B", "category": "misc"}, {"word": "C", "category": "misc"}, {"word": "D", "category": "misc"}, {"word": "E", "category": "misc"}, {"word": "F", "category": "misc"}, {"word": "G", "category": "misc"}, {"word": "H", "category": "misc"}],
+    [{"word": "it", "category": "pronoun"}, {"word": "did", "category": "verb"}, {"word": "letter_I", "category": "misc", "display": "I"}, {"word": "J", "category": "misc"}, {"word": "K", "category": "misc"}, {"word": "L", "category": "misc"}, {"word": "M", "category": "misc"}, {"word": "N", "category": "misc"}, {"word": "O", "category": "misc"}, {"word": "P", "category": "misc"}],
+    [{"word": "they", "category": "pronoun"}, {"word": "do", "category": "verb"}, {"word": "Q", "category": "misc"}, {"word": "R", "category": "misc"}, {"word": "S", "category": "misc"}, {"word": "T", "category": "misc"}, {"word": "U", "category": "misc"}, {"word": "V", "category": "misc"}, {"word": "W", "category": "misc"}, {"word": "X", "category": "misc"}],
+    [{"word": "we", "category": "pronoun"}, {"word": "Y", "category": "misc"}, {"word": "Z", "category": "misc"}, {"word": "1", "category": "misc"}, {"word": "2", "category": "misc"}, {"word": "3", "category": "misc"}, {"word": "4", "category": "misc"}, {"word": "5", "category": "misc"}, {"word": ".", "category": "misc"}, {"word": "?", "category": "misc"}],
+    [{"word": "you", "category": "pronoun"}, {"word": "6", "category": "misc"}, {"word": "7", "category": "misc"}, {"word": "8", "category": "misc"}, {"word": "9", "category": "misc"}, {"word": "0", "category": "misc"}, {"word": "-", "category": "misc"}, {"word": "!", "category": "misc"}, {"word": ",", "category": "misc"}, {"word": "SPACE", "category": "misc"}],
+    [{"word": "your", "category": "pronoun"}, {"word": "like", "category": "verb"}, {"word": "to", "category": "preposition"}, {"word": "with", "category": "preposition"}, {"word": "in", "category": "preposition"}, {"word": "the", "category": "misc"}, {"word": "and", "category": "misc"}, {"word": "but", "category": "misc"}, {"word": "not", "category": "descriptive"}, {"word": "yes", "category": "common"}]
 ]
 
-# --- Add Custom Words to Layout (if any) ---
+dynamic_layout = list(layout) # Make a mutable copy for adding custom words
+
 if st.session_state.custom_words:
-    # This logic for adding custom words seems okay, but ensure layout is mutable if modified often
-    # Or rebuild the full layout list including custom words each time
     current_custom_row = []
     for idx, word_info in enumerate(st.session_state.custom_words):
         word = word_info["word"]
         category = word_info["category"]
         current_custom_row.append({"word": f"custom_{idx}_{word}", "display": word, "category": category})
-        if len(current_custom_row) == 10: # Assuming 10 columns for keyboard
-            layout.append(list(current_custom_row)) # Add a copy
+        if len(current_custom_row) == 10:
+            dynamic_layout.append(list(current_custom_row))
             current_custom_row = []
-    if current_custom_row: # Add any remaining custom words
+    if current_custom_row:
         while len(current_custom_row) < 10:
-            current_custom_row.append({"word": "", "category": "misc"}) # Pad with empty placeholders
-        layout.append(list(current_custom_row))
+            current_custom_row.append({"word": "", "category": "misc"})
+        dynamic_layout.append(list(current_custom_row))
 
-# --- Keyboard Rendering ---
 def add_to_output(word_to_add):
     current_text = st.session_state.get("text_output", "")
     if word_to_add == "SPACE":
         current_text += " "
     elif word_to_add in [".", "?", "!", ",", "-"]:
         current_text += word_to_add
-    elif word_to_add.isdigit() or (len(word_to_add) == 1 and word_to_add.isalpha()): # For single letters/numbers
+    elif word_to_add.isdigit() or (len(word_to_add) == 1 and word_to_add.isalpha()):
         current_text += word_to_add
-    else: # For full words
+    else:
         if current_text and not current_text.endswith(" "):
             current_text += " "
         current_text += word_to_add
     st.session_state.text_output = current_text
 
 st.markdown("### Communication Keyboard")
-for row_idx, row_items in enumerate(layout):
+for row_idx, row_items in enumerate(dynamic_layout):
     cols = st.columns(len(row_items))
     for col_idx, item_info in enumerate(row_items):
-        # Ensure item_info is a dictionary and has 'word'
         if not isinstance(item_info, dict) or "word" not in item_info or item_info["word"] == "":
-            continue 
-            
+            continue
         word = item_info["word"]
-        category = item_info.get("category", "misc") # Default category if missing
+        category = item_info.get("category", "misc")
         display_text = item_info.get("display", word)
-        
-        button_key = f"key_{row_idx}_{col_idx}_{category}_{word.replace(' ', '_')}" # Make key more robust
+        button_key = f"key_{row_idx}_{col_idx}_{category}_{word.replace(' ', '_').replace('.', 'dot').replace('?', 'qmark')}" # Make key even more robust
 
-        # Use a lambda with default arguments to capture current word for the callback
-        # This avoids the late binding issue with loops if not using functools.partial or similar
         if cols[col_idx].button(
-                display_text if word != "SPACE" else "␣", # Display text
+                display_text if word != "SPACE" else "␣",
                 key=button_key,
                 use_container_width=True
             ):
-            # Action happens here when button is clicked, then Streamlit reruns
-            # For complex actions or to avoid issues with default args in loops for callbacks,
-            # direct action or ensuring unique callbacks is better.
-            # Here, the action is simple enough that direct modification is fine.
-            word_to_add_on_click = display_text # Use display_text for 'letter_I' or custom words
-            if word.startswith("custom_") or word.startswith("letter_"):
-                 add_to_output(display_text)
-            else:
-                 add_to_output(word)
-            st.rerun() # Rerun to update the text_area with the new character/word
-
+            word_to_add_on_click = display_text if word.startswith("custom_") or word.startswith("letter_") else word
+            add_to_output(word_to_add_on_click)
+            st.rerun()
 
-# --- Control Buttons & SEND Action Processing ---
 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="secondary"): # Changed type
+    if st.button("CLEAR", key="clear_btn", help="Clear all text", use_container_width=True, type="secondary"):
         st.session_state.text_output = ""
-        st.rerun() 
+        st.rerun()
 with col2:
-    if st.button("SPEAK", key="speak_btn", help="Speak the current text (feature placeholder)", use_container_width=True, type="secondary"): # Changed type
+    if st.button("SPEAK", key="speak_btn", help="Speak the current text (feature placeholder)", use_container_width=True, type="secondary"):
         if st.session_state.text_output:
             st.toast(f"Speaking (simulated): {st.session_state.text_output}", icon="🔊")
-            # Actual text-to-speech integration would go here
-
-    # Moved DEL button here as SPEAK often stands alone or is larger
-    if st.button("⌫ DEL", key="backspace_btn", help="Delete last character", use_container_width=True): # Changed key
+    
+    if st.button("⌫ DEL", key="backspace_btn", 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]
-            st.rerun() 
-with col4:
-    if st.button("⌫ WORD", key="backspace_word_btn", help="Delete last word", use_container_width=True): # Changed key
-        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: # Add a space if there are still words left
-                    st.session_state.text_output += " "
-            else: # If no words left after split (e.g. only spaces before), clear it
-                st.session_state.text_output = ""
             st.rerun()
+with col4:
+    if st.button("⌫ WORD", key="backspace_word_btn", help="Delete last word", use_container_width=True):
+        current_text = st.session_state.text_output.rstrip()
+        if ' ' in current_text:
+            st.session_state.text_output = current_text.rsplit(' ', 1)[0] + " "
+        else:
+            st.session_state.text_output = ""
+        st.rerun()
 
-with col3: 
+with col3:
     if st.button("SEND 🗣️", key="send_btn", help="Send message to assistant", use_container_width=True, type="primary"):
         if 'assistant' not in st.session_state or st.session_state.assistant is None:
-            st.error("AAC Assistant is not initialized. Please check sidebar for errors.")
+            st.error("AAC Assistant is not initialized. Please check sidebar for errors or wait for initialization to complete.")
         elif not st.session_state.text_output.strip():
             st.toast("Please compose a message first.", icon="✍️")
         else:
-            user_message_to_send = st.session_state.text_output # Capture message
-            
-            # Add user message to chat history
+            user_message_to_send = st.session_state.text_output
             st.session_state.messages.append({"role": "user", "content": user_message_to_send})
+            st.session_state.text_output = ""
             
-            # Clear the input text area immediately
-            st.session_state.text_output = "" 
-            
-            # Process with AACAssistant and show a spinner
             with st.spinner("Elliot is thinking... please wait a moment..."):
                 try:
                     print(f"DEBUG [Streamlit UI]: User typed: '{user_message_to_send}'")
                     print(f"DEBUG [Streamlit UI]: Calling assistant.process_query for '{user_message_to_send}'")
-                    
-                    # CRITICAL: This is where your rag.py's process_query is called
-                    # Ensure it returns ONLY the assistant's clean response
                     assistant_response = st.session_state.assistant.process_query(user_message_to_send)
-                    
                     print(f"DEBUG [Streamlit UI]: Received assistant response: '{assistant_response}'")
                     st.session_state.messages.append({"role": "assistant", "content": assistant_response})
                 except Exception as e:
                     error_message_display = f"An error occurred while processing your message: {e}"
-                    st.error(error_message_display) 
+                    st.error(error_message_display)
                     st.session_state.messages.append({"role": "assistant", "content": f"*Error: Could not get a response. Details: {e}*"})
-                    # For more detailed debugging in the UI itself if needed:
-                    # import traceback
-                    # st.text_area("Error Traceback:", traceback.format_exc(), height=200)
-            
-            st.rerun() # Rerun to update the chat display and clear the (now empty) input field visually.
+            st.rerun()
 
-# --- Settings Sidebar ---
 with st.sidebar:
-    # Your settings code (commented out in your original, kept as is)
-    # st.title("Settings")
-    # ...
     st.divider()
-    # --- Chat History Display in Sidebar ---
     st.subheader("Conversation Log")
-    chat_container = st.container(height=400 if st.session_state.messages else 100) # Adjust height if empty
+    chat_container = st.container(height=400 if st.session_state.messages else 100)
     with chat_container:
         if not st.session_state.messages:
             st.caption("No messages yet...")