Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,91 +1,90 @@
-import altair as alt
-import numpy as np
-import pandas as pd
-import streamlit as st
-
-
 import streamlit as st
-# import streamlit as st
-# Disable all filesystem watching
-# st.set_option("server.fileWatcherType", "none")
+st.set_option("server.fileWatcherType", "none") # Disable filesystem watcher - good for deployed apps
 
+import altair as alt # Keep if used, otherwise remove
+import numpy as np   # Keep if used, otherwise remove
 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
-cache_dir_base = "/app/.cache_app"
-
-# Create the base directory if it doesn't exist
-os.makedirs(cache_dir_base, exist_ok=True)
+import rag  # Import your rag module
+import os
 
-# Set environment variables for Hugging Face libraries
-os.environ['HF_HOME'] = os.path.join(cache_dir_base, 'huggingface_hub') # For hub downloads
-os.environ['TRANSFORMERS_CACHE'] = os.path.join(cache_dir_base, 'transformers') # For transformers models
-os.environ['SENTENCE_TRANSFORMERS_HOME'] = os.path.join(cache_dir_base, 'sentence_transformers') # For sentence-transformers
-
-# Set environment variable for pip cache (optional but good practice)
-os.environ['PIP_CACHE_DIR'] = os.path.join(cache_dir_base, 'pip')
+# --- 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
+    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 ---
-# 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 = []
-        st.session_state.text_size = 22
-        st.session_state.theme = "Default"
-        st.session_state.speech_rate = 1.0
-        st.session_state.voice_option = "Default Voice"
-        st.session_state.messages = []
-        st.session_state.assistant = None
-        # Initialize text_output only if assistant is not initialized, assuming they go together
-        st.session_state.text_output = ""
+if 'assistant_initialized_flag' not in st.session_state: # Use a flag to track if full init ran
+    st.session_state.current_page = "main"
+    st.session_state.show_custom_words = False
+    st.session_state.custom_words = []
+    st.session_state.text_size = 22
+    st.session_state.theme = "Default"
+    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.text_output = ""
+    st.session_state.assistant_initialized_flag = True # Mark that initial session state setup is done
 
-    # --- Theme Colors ---
+
+# --- 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 (Adapted from previous main.py) ---
-@st.cache_resource # Cache the assistant
+
+# --- Helper Function to Initialize Assistant ---
+@st.cache_resource # Cache the assistant object
 def initialize_assistant(doc_path):
     """Initializes the AACAssistant."""
-        # Create a dummy document if it doesn't exist for demonstration
+    # 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. Creating dummy.")
+        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
             with open(doc_path, "w") as f:
                 f.write("""
                 I grew up in Seattle and love the rain.
@@ -95,21 +94,31 @@ def initialize_assistant(doc_path):
                 I enjoy sci-fi movies.
                 """)
         except Exception as e:
-            st.sidebar.error(f"Failed to create dummy doc: {e}")
-            return None
+            st.sidebar.error(f"Failed to create dummy doc at '{doc_path}': {e}")
+            return None 
     try:
+        st.sidebar.info("Initializing AAC Assistant... This may take some time.")
         assistant = rag.AACAssistant(doc_path)
-        st.sidebar.success("AAC Assistant Initialized.")
+        st.sidebar.success("AAC Assistant Initialized and ready!")
         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
+        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
+
+# --- 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)
 
-DEFAULT_DOCUMENT_PATH = "src/aac_user_experiences.txt"
 # --- 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;
@@ -128,7 +137,7 @@ css = f"""
             gap: 5px;
         }}
         section[data-testid="stSidebar"] {{
-            width: 20rem;
+            width: 20rem !important; /* Ensure sidebar width is applied */
             background-color: {colors["bg"]};
         }}
         body {{
@@ -141,6 +150,7 @@ css = f"""
             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;
         }}
@@ -149,47 +159,27 @@ css = f"""
             transform: scale(1.03);
             box-shadow: 0 2px 3px rgba(0,0,0,0.1);
         }}
-        .control-button {{
+        .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 chat styling */
+        .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')}; /* Use theme background */
+            background-color: {colors.get('bg', '#FFFFFF')};
             border-radius: 8px;
         }}
     </style>
     """
+st.markdown(css, unsafe_allow_html=True)
 
-    # --- JS for Button Coloring (no delay, no setTimeout) ---
+# --- JS for Button Coloring ---
+# Your JS remains the same
 js = """
 <script>
     function colorButtons() {
@@ -197,287 +187,252 @@ js = """
         buttons.forEach(button => {
             const id = button.id;
             const parts = id.split('_');
-            if (parts.length >= 4) {
+            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);
-    new MutationObserver(colorButtons).observe(document.body, { childList: true, subtree: true });
+
+    // 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
+            }
+        }
+    });
+    observer.observe(streamlitDoc.body, { childList: true, subtree: true });
+
+    // Initial call in case DOM is already ready from a script rerun
+    colorButtons();
 </script>
 """
-st.markdown(css, unsafe_allow_html=True)
 st.markdown(js, unsafe_allow_html=True)
 
-    # --- Keyboard Layout ---
+
+# --- 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
+)
+
+# 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": "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": "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": "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": "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": "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": "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": "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": "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 ---
+# --- Add Custom Words to Layout (if any) ---
 if st.session_state.custom_words:
-    custom_row = []
+    # 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"]
-        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)
+        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
+            current_custom_row = []
+    if current_custom_row: # Add any remaining custom words
+        while len(current_custom_row) < 10:
+            current_custom_row.append({"word": "", "category": "misc"}) # Pad with empty placeholders
+        layout.append(list(current_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")
-    # Use st.text_area directly for input and display, bound to session state
-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":
-        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
+# --- 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
+        current_text += word_to_add
+    else: # For full words
+        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 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()
-                )
+for row_idx, row_items in enumerate(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 
+            
+        word = item_info["word"]
+        category = item_info.get("category", "misc") # Default category if missing
+        display_text = item_info.get("display", word)
+        
+        button_key = f"key_{row_idx}_{col_idx}_{category}_{word.replace(' ', '_')}" # Make key 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
+                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
 
-    # --- Control Buttons ---
+
+# --- 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="primary"):
+    if st.button("CLEAR", key="clear_btn", help="Clear all text", use_container_width=True, type="secondary"): # Changed type
         st.session_state.text_output = ""
-        st.rerun() # Rerun to reflect change
+        st.rerun() 
 with col2:
-    if st.button("SPEAK", key="speak_btn", help="Speak the current text", use_container_width=True, type="primary"):
+    if st.button("SPEAK", key="speak_btn", help="Speak the current text (feature placeholder)", use_container_width=True, type="secondary"): # Changed type
         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):
+            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.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", help="Delete last word", use_container_width=True):
+    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:
+                if words: # Add a space if there are still words left
                     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
+            else: # If no words left after split (e.g. only spaces before), clear it
+                st.session_state.text_output = ""
+            st.rerun()
+
+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.")
+        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
+            st.session_state.messages.append({"role": "user", "content": user_message_to_send})
+            
+            # 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.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.
+
 # --- Settings Sidebar ---
 with st.sidebar:
-        # --- Settings Section Commented Out ---
-        # st.title("Settings")
-        # st.subheader("Interface")
-        # theme_options = list(theme_colors.keys())
-        # theme_index = theme_options.index(st.session_state.theme)
-        # new_theme = st.selectbox("Theme", theme_options, index=theme_index)
-        # new_text_size = st.slider("Text Size", 12, 36, st.session_state.text_size)
-        # if st.button("Apply Settings Changes", type="primary"):
-        #     changed = False
-        #     if new_theme != st.session_state.theme:
-        #         st.session_state.theme = new_theme
-        #         changed = True
-        #     if new_text_size != st.session_state.text_size:
-        #         st.session_state.text_size = new_text_size
-        #         changed = True
-        #     if changed:
-        #         st.rerun()
-        # with st.expander("Speech Settings"):
-        #     speech_rate = st.slider("Rate", 0.5, 2.0, st.session_state.speech_rate, 0.1)
-        #     voice_options = ["Default Voice", "Female Voice", "Male Voice", "Child Voice"]
-        #     voice_index = voice_options.index(st.session_state.voice_option)
-        #     voice = st.selectbox("Voice", voice_options, index=voice_index)
-        #     if st.button("Apply Speech Settings"):
-        #         st.session_state.speech_rate = speech_rate
-        #         st.session_state.voice_option = voice
-        # st.subheader("Custom Words")
-        # with st.expander("Add New Word"):
-        #     word = st.text_input("Word")
-        #     cat = st.selectbox("Category", list(colors.keys()))
-        #     col1, col2 = st.columns(2)
-        #     with col1:
-        #         if st.button("Add", key="add_word"):
-        #             if word and cat:
-        #                 st.session_state.custom_words.append({"word": word, "category": cat})
-        #                 st.success(f"Added '{word}'")
-        #                 st.rerun()
-        #     if st.session_state.custom_words:
-        #         st.write("Current custom words:")
-        #         words_df = pd.DataFrame([
-        #             {"Word": w["word"], "Category": w["category"]}
-        #             for w in st.session_state.custom_words
-        #         ])
-        #         st.dataframe(words_df, hide_index=True)
-        #         with col2:
-        #             if st.button("Clear All", key="clear_words"):
-        #                 st.session_state.custom_words = []
-        #                 st.success("Words cleared")
-        #                 st.rerun()
-        # --- End of Settings Section Commented Out ---
-
+    # Your settings code (commented out in your original, kept as is)
+    # st.title("Settings")
+    # ...
     st.divider()
     # --- Chat History Display in Sidebar ---
-    st.subheader("Conversation")
-    chat_container = st.container(height=400) # Fixed height container
+    st.subheader("Conversation Log")
+    chat_container = st.container(height=400 if st.session_state.messages else 100) # Adjust height if empty
     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 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
-
+        if not st.session_state.messages:
+            st.caption("No messages yet...")
+        else:
+            for message in st.session_state.messages:
+                with st.chat_message(message["role"]):
+                    st.markdown(message["content"])