Simplify UI layout

app.py

@@ -1,151 +1,80 @@
-import os
-import sys
-import traceback
 from typing import cast
 
-# Print to container logs for debugging
-print("=" * 50, file=sys.stderr)
-print("Starting app.py...", file=sys.stderr)
-print("=" * 50, file=sys.stderr)
-
 import streamlit as st
 
-print("Streamlit imported successfully", file=sys.stderr)
-
-# Configure Streamlit for Hugging Face Spaces
-# Port is set via command line argument in Dockerfile
-# Don't set STREAMLIT_SERVER_PORT here to avoid conflicts
-os.environ["STREAMLIT_SERVER_ADDRESS"] = "0.0.0.0"
-print("Set STREAMLIT_SERVER_ADDRESS", file=sys.stderr)
-
-# Try to import modules with error handling
-print("Attempting to import modules...", file=sys.stderr)
-try:
-    from gateway_client import delete_profile, ingest_and_rewrite
-    print("✓ gateway_client imported", file=sys.stderr)
-    from llm import chat, set_model
-    print("✓ llm imported", file=sys.stderr)
-    from model_config import MODEL_CHOICES, MODEL_TO_PROVIDER
-    print("✓ model_config imported", file=sys.stderr)
-    IMPORTS_SUCCESSFUL = True
-    print("All imports successful!", file=sys.stderr)
-except Exception as e:
-    IMPORTS_SUCCESSFUL = False
-    IMPORT_ERROR = str(e)
-    IMPORT_TRACEBACK = traceback.format_exc()
-    print(f"✗ Import failed: {IMPORT_ERROR}", file=sys.stderr)
-    print(IMPORT_TRACEBACK, file=sys.stderr)
-    # Create dummy functions to prevent NameError
-    def delete_profile(user_id: str) -> bool:
-        return False
-    def ingest_and_rewrite(user_id: str, query: str, model_type: str = "openai") -> str:
-        return query
-    def chat(messages, persona):
-        return "Error: Could not initialize chat function", 0, 0, 0
-    def set_model(model_id: str) -> None:
-        pass
-    MODEL_CHOICES = ["gpt-4.1-mini"]
-    MODEL_TO_PROVIDER = {"gpt-4.1-mini": "openai"}
+from gateway_client import delete_profile, ingest_and_rewrite
+from llm import chat, set_model
+from model_config import MODEL_CHOICES, MODEL_TO_PROVIDER
 
 
 def rewrite_message(
-    msg: str, persona_name: str, show_rationale: bool, skip_rewrite: bool, provider: str = "openai"
+    msg: str,
+    persona_name: str,
+    show_rationale: bool,
+    skip_rewrite: bool,
+    provider: str,
 ) -> str:
+    """Rewrite the user message via MemMachine unless skip is requested."""
     if skip_rewrite:
         rewritten_msg = msg
         if show_rationale:
             rewritten_msg += " At the beginning of your response, please say the following in ITALIC: 'Persona Rationale: No personalization applied.'. Begin your answer on the next line."
-    else:
-        try:
-            rewritten_msg = ingest_and_rewrite(
-                user_id=persona_name, query=msg, model_type=provider
-            )
-            if show_rationale:
-                rewritten_msg += " At the beginning of your response, please say the following in ITALIC: 'Persona Rationale: ' followed by 1 sentence about how your reasoning for how the persona traits influenced this response, also in italics. Begin your answer on the next line."
+        return rewritten_msg
 
-        except Exception as e:
-            # If backend is unavailable, use original message without rewriting
-            st.warning(f"Backend memory server unavailable. Using message without personalization: {e}")
-            rewritten_msg = msg
-            if show_rationale:
-                rewritten_msg += " At the beginning of your response, please say the following in ITALIC: 'Persona Rationale: No personalization applied (backend unavailable).'. Begin your answer on the next line."
-    return rewritten_msg
+    try:
+        rewritten_msg = ingest_and_rewrite(
+            user_id=persona_name, query=msg, model_type=provider
+        )
+        if show_rationale:
+            rewritten_msg += " At the beginning of your response, please say the following in ITALIC: 'Persona Rationale: ' followed by 1 sentence about how your reasoning for how the persona traits influenced this response, also in italics. Begin your answer on the next line."
+        return rewritten_msg
+    except Exception as e:
+        st.warning(
+            f"Backend memory server unavailable. Using message without personalization: {e}"
+        )
+        rewritten_msg = msg
+        if show_rationale:
+            rewritten_msg += " At the beginning of your response, please say the following in ITALIC: 'Persona Rationale: No personalization applied (backend unavailable).'. Begin your answer on the next line."
+        return rewritten_msg
 
 
 # ──────────────────────────────────────────────────────────────
 # Page setup & CSS
 # ──────────────────────────────────────────────────────────────
-print("Setting up page config...", file=sys.stderr)
-try:
-    st.set_page_config(page_title="MemMachine Chatbot", layout="wide")
-    print("✓ Page config set", file=sys.stderr)
-except Exception as e:
-    print(f"✗ Failed to set page config: {e}", file=sys.stderr)
-    st.error(f"Failed to set page config: {e}")
-
-# Show import errors if any
-if not IMPORTS_SUCCESSFUL:
-    print("Showing import error to user...", file=sys.stderr)
-    st.error("⚠️ Application Error: Failed to import required modules")
-    st.error(f"Error: {IMPORT_ERROR}")
-    with st.expander("Show full error details"):
-        st.code(IMPORT_TRACEBACK)
-    st.stop()
-
-# Debug: Show that app is loading
-print("Rendering initial content...", file=sys.stderr)
-sys.stderr.flush()  # Force flush to ensure logs appear
-st.markdown(
-    '<div class="pill on">🚀 Playground Ready</div>',
-    unsafe_allow_html=True,
-)
-print("✓ Initial content rendered", file=sys.stderr)
-sys.stderr.flush()
+st.set_page_config(page_title="MemMachine Chatbot", layout="wide")
 
 try:
     with open("./styles.css") as f:
         st.markdown(f"<style>{f.read()}</style>", unsafe_allow_html=True)
 except FileNotFoundError:
-    st.markdown("<style></style>", unsafe_allow_html=True)
+    pass
 
 
 # ──────────────────────────────────────────────────────────────
 # Sidebar
 # ──────────────────────────────────────────────────────────────
-try:
-    with st.sidebar:
-        try:
-            st.image("./assets/memmachine_logo.png", use_container_width=True)
-        except (FileNotFoundError, Exception) as e:
-            st.markdown("### MemMachine")
-
-        st.markdown("#### Choose Model")
-
-        try:
-            model_id = st.selectbox(
-                "Choose Model", MODEL_CHOICES, index=0, label_visibility="collapsed"
-            )
-            provider = MODEL_TO_PROVIDER[model_id]
-            set_model(model_id)
-        except Exception as e:
-            st.error(f"Error in model selection: {e}")
-            model_id = MODEL_CHOICES[0] if MODEL_CHOICES else "gpt-4.1-mini"
-            provider = MODEL_TO_PROVIDER.get(model_id, "openai")
-            set_model(model_id)
-except Exception as e:
-    st.error(f"Error initializing sidebar: {e}")
-    import traceback
-    st.code(traceback.format_exc())
-    # Set defaults
-    model_id = MODEL_CHOICES[0] if MODEL_CHOICES else "gpt-4.1-mini"
+default_model = MODEL_CHOICES[0] if MODEL_CHOICES else "gpt-4.1-mini"
+model_id = default_model
+provider = MODEL_TO_PROVIDER.get(default_model, "openai")
+selected_persona = "Charlie"
+persona_name = "Charlie"
+skip_rewrite = False
+compare_personas = False
+show_rationale = False
+
+with st.sidebar:
+    try:
+        st.image("./assets/memmachine_logo.png", use_container_width=True)
+    except (FileNotFoundError, Exception):
+        st.markdown("### MemMachine")
+
+    st.markdown("#### Choose Model")
+    model_id = st.selectbox(
+        "Choose Model", MODEL_CHOICES, index=0, label_visibility="collapsed"
+    )
     provider = MODEL_TO_PROVIDER.get(model_id, "openai")
-    selected_persona = "Charlie"
-    persona_name = "Charlie"
-    skip_rewrite = False
-    compare_personas = False
-    show_rationale = False
-else:
-    # Continue with sidebar if no exception
+    set_model(model_id)
+
     st.markdown("#### Choose user persona")
     selected_persona = st.selectbox(
         "Choose user persona",
@@ -174,80 +103,7 @@ else:
             st.error(f"Failed to delete profile for '{persona_name}'.")
     st.divider()
 
-# ──────────────────────────────────────────────────────────────
-# Hero / Context panels
-# ──────────────────────────────────────────────────────────────
-persona_badge = persona_name if persona_name == selected_persona else f"{persona_name} • custom"
-
-hero_html = f"""
-<div class="hero-card">
-  <div>
-    <p class="eyebrow">Live Persona Playground</p>
-    <h2>MemMachine Conversational Memory</h2>
-    <p>Use rapid persona rewrites, compare against control personas, and stream
-    LLM output from your preferred provider without leaving this panel.</p>
-  </div>
-  <div class="status-chip">
-    <span>Active model</span>
-    <strong>{model_id}</strong>
-    <small>{provider.title()}</small>
-  </div>
-</div>
-"""
-st.markdown(hero_html, unsafe_allow_html=True)
-
-stat_cols = st.columns(3)
-stat_cols[0].markdown(
-    f"""
-    <div class="stat-card">
-        <span>Persona</span>
-        <h4>{persona_badge}</h4>
-    </div>
-    """,
-    unsafe_allow_html=True,
-)
-stat_cols[1].markdown(
-    f"""
-    <div class="stat-card">
-        <span>Comparison mode</span>
-        <h4>{"Control vs Persona" if compare_personas else "Single persona"}</h4>
-    </div>
-    """,
-    unsafe_allow_html=True,
-)
-stat_cols[2].markdown(
-    f"""
-    <div class="stat-card">
-        <span>Rewrite</span>
-        <h4>{"Enabled" if not skip_rewrite else "Skipped"}</h4>
-    </div>
-    """,
-    unsafe_allow_html=True,
-)
-
-st.markdown(
-    f"""
-    <div class="pill-group">
-        <span class="pill {'on' if not skip_rewrite else ''}">Rewrite</span>
-        <span class="pill {'on' if compare_personas else ''}">Compare personas</span>
-        <span class="pill {'on' if show_rationale else ''}">Persona rationale</span>
-        <span class="pill on">Backend: {os.getenv("MEMORY_SERVER_URL", "default")}</span>
-    </div>
-    """,
-    unsafe_allow_html=True,
-)
-
-with st.expander("How to use this playground", expanded=False):
-    st.markdown(
-        """
-        1. Pick a model + persona in the sidebar.  
-        2. Toggle rewrite / comparison to control how memories are applied.  
-        3. Type a message below — MemMachine will rewrite user input if enabled, call the model, and stream the response.  
-        4. Use *Clear chat* or *Delete profile* to reset memories.
-        """
-    )
 
-# ──────────────────────────────────────────────────────────────
 # ──────────────────────────────────────────────────────────────
 # Session state
 # ──────────────────────────────────────────────────────────────
@@ -285,12 +141,8 @@ def append_user_turn(msgs: list[dict], new_user_msg: str) -> list[dict]:
 # ──────────────────────────────────────────────────────────────
 # Title
 # ──────────────────────────────────────────────────────────────
-try:
-    st.title("MemMachine Chatbot")
-except Exception as e:
-    st.error(f"Error rendering title: {e}")
-    import traceback
-    st.code(traceback.format_exc())
+st.title("MemMachine Chatbot")
+
 
 # ──────────────────────────────────────────────────────────────
 # Chat logic
@@ -298,35 +150,39 @@ except Exception as e:
 msg = st.chat_input("Type your message…")
 if msg:
     st.session_state.history.append({"role": "user", "content": msg})
-    # rewritten_msg = "Use the persona profile to personalize your naswer only when applicable.\n"
     if compare_personas:
         all_answers = {}
-        rewritten_msg = rewrite_message(msg, persona_name, show_rationale, False, provider)
+        rewritten_msg = rewrite_message(
+            msg, persona_name, show_rationale, False, provider
+        )
         msgs = clean_history(st.session_state.history, persona_name)
         msgs = append_user_turn(msgs, rewritten_msg)
-        txt, lat, tok, tps = chat(msgs, persona_name)
+        txt, *_ = chat(msgs, persona_name)
         all_answers[persona_name] = txt
 
-        rewritten_msg_control = rewrite_message(msg, "Control", show_rationale, True, provider)
+        rewritten_msg_control = rewrite_message(
+            msg, "Control", show_rationale, True, provider
+        )
         msgs_control = clean_history(st.session_state.history, "Control")
         msgs_control = append_user_turn(msgs_control, rewritten_msg_control)
-        txt_control, lat, tok, tps = chat(msgs_control, "Arnold")
+        txt_control, *_ = chat(msgs_control, "Arnold")
         all_answers["Control"] = txt_control
 
         st.session_state.history.append(
             {"role": "assistant_all", "axis": "role", "content": all_answers}
         )
     else:
-        rewritten_msg = rewrite_message(msg, persona_name, show_rationale, skip_rewrite, provider)
+        rewritten_msg = rewrite_message(
+            msg, persona_name, show_rationale, skip_rewrite, provider
+        )
         msgs = clean_history(st.session_state.history, persona_name)
         msgs = append_user_turn(msgs, rewritten_msg)
-        txt, lat, tok, tps = chat(
-            msgs, "Arnold" if persona_name == "Control" else persona_name
-        )
+        txt, *_ = chat(msgs, "Arnold" if persona_name == "Control" else persona_name)
         st.session_state.history.append(
             {"role": "assistant", "persona": persona_name, "content": txt}
         )
 
+
 # ──────────────────────────────────────────────────────────────
 # Chat history display
 # ───────────────────────────────────────────��──────────────────