Update src/streamlit_app.py

src/streamlit_app.py

@@ -12,29 +12,470 @@ 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
+
+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.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 = ""
+
+    # --- 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"
+        }
+    }
+    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
+        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"""
+    <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;
+            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;
+            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 chat styling */
+        .sidebar .stChatMessage {{
+            background-color: {colors.get('bg', '#FFFFFF')}; /* Use theme background */
+            border-radius: 8px;
+        }}
+    </style>
+    """
+
+    # --- 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)
+
+    # --- Keyboard Layout ---
+    layout = [
+        [
+            {"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 ---
+    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")
+    # 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
+
+    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()
+                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:
+        # --- 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 ---
+
+        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 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
 
-num_points = st.slider("Number of points in spiral", 1, 10000, 1100)
-num_turns = st.slider("Number of turns in spiral", 1, 300, 31)
-
-indices = np.linspace(0, 1, num_points)
-theta = 2 * np.pi * num_turns * indices
-radius = indices
-
-x = radius * np.cos(theta)
-y = radius * np.sin(theta)
-
-df = pd.DataFrame({
-    "x": x,
-    "y": y,
-    "idx": indices,
-    "rand": np.random.randn(num_points),
-})
-
-st.altair_chart(alt.Chart(df, height=700, width=700)
-    .mark_point(filled=True)
-    .encode(
-        x=alt.X("x", axis=None),
-        y=alt.Y("y", axis=None),
-        color=alt.Color("idx", legend=None, scale=alt.Scale()),
-        size=alt.Size("rand", legend=None, scale=alt.Scale(range=[1, 150])),
-    ))
+if __name__ == "__main__":
+    main()