Update src/streamlit_app.py

src/streamlit_app.py

@@ -3,7 +3,7 @@ import os
 import re
 import time
 import warnings
-from typing import List, Dict
+from typing import List, Dict, Any, Optional
 
 import requests
 import streamlit as st
@@ -14,73 +14,30 @@ from teapotai import TeapotAI
 
 
 # -----------------------
-# Silence noisy warnings (optional)
+# Warnings (optional)
 # -----------------------
 warnings.filterwarnings("ignore", message="pkg_resources is deprecated as an API.*")
 warnings.filterwarnings("ignore", message='Field name "schema" in "TeapotTool" shadows.*')
 
 
 # -----------------------
-# Branding / Theme
+# Brand / Assets
 # -----------------------
 TEAPOT_LOGO_GIF = "https://teapotai.com/assets/logo.gif"
 
-TEA_BG = "#fbf7ef"       # warm off-white
-TEA_PANEL = "#fffaf2"    # slightly brighter
-TEA_TEXT = "#1f2937"     # slate-ish
-TEA_MUTED = "#6b7280"    # gray
-TEA_ACCENT = "#c0841d"   # warm amber
-TEA_BORDER = "rgba(31, 41, 55, 0.10)"
-
-st.set_page_config(
-    page_title="TeapotAI Chat",
-    page_icon="🫖",
-    layout="centered",
-)
-
-CUSTOM_CSS = f"""
-<style>
-.stApp {{
-  background: {TEA_BG};
-  color: {TEA_TEXT};
-}}
-
-section[data-testid="stSidebar"] {{
-  background: {TEA_PANEL};
-  border-right: 1px solid {TEA_BORDER};
-}}
-
-div[data-testid="stChatMessage"] {{
-  border-radius: 16px;
-  padding: 8px 10px;
-}}
-
-.stTextInput > div > div, .stTextArea > div > div {{
-  border-radius: 12px !important;
-}}
-
-.stButton button {{
-  border-radius: 12px;
-  border: 1px solid {TEA_BORDER};
-}}
-
-a {{
-  color: {TEA_ACCENT} !important;
-}}
-</style>
-"""
-st.markdown(CUSTOM_CSS, unsafe_allow_html=True)
+SUGGESTED_QUERIES = [
+    "Who are you?",
+    "Tell me about teapotllm",
+    "What is the weather like in NYC today?",
+]
 
 
 # -----------------------
-# Config
+# Models
 # -----------------------
-BRAVE_ENDPOINT = "https://api.search.brave.com/res/v1/web/search"
-TOP_K = 3
-TIMEOUT_SECS = 15
-
 MODEL_TINY = "teapotai/tinyteapot"
-MODEL_LLM = "teapotai/teapotllm"
+MODEL_LLM = "teapotai/teapotllm"  # if you keep the toggle elsewhere later
+DEFAULT_MODEL = MODEL_TINY
 
 DEFAULT_SYSTEM_PROMPT = (
     "You are Teapot, an open-source AI assistant optimized for low-end devices, "
@@ -95,39 +52,173 @@ DEFAULT_DOCUMENTS = [
 ]
 
 
+# -----------------------
+# Web search (Brave) config
+# -----------------------
+BRAVE_ENDPOINT = "https://api.search.brave.com/res/v1/web/search"
+TOP_K = 3
+TIMEOUT_SECS = 15
+
+
+# -----------------------
+# Streamlit page config
+# -----------------------
+st.set_page_config(page_title="TeapotAI Chat", page_icon="🫖", layout="centered")
+
+
+# -----------------------
+# Theme (light/dark) via CSS variables
+# -----------------------
+LIGHT_THEME = dict(
+    bg="#fbf7ef",
+    panel="#fffaf2",
+    text="#111827",
+    muted="#6b7280",
+    accent="#c0841d",
+    border="rgba(17, 24, 39, 0.12)",
+    bubble_user="#eef2ff",
+    bubble_asst="#ffffff",
+    code_bg="#0b1220",
+)
+
+DARK_THEME = dict(
+    bg="#0b0f19",
+    panel="#0f1626",
+    text="#e5e7eb",
+    muted="#9ca3af",
+    accent="#f59e0b",
+    border="rgba(229, 231, 235, 0.12)",
+    bubble_user="#111827",
+    bubble_asst="#0f172a",
+    code_bg="#0b1220",
+)
+
+
+def inject_css(theme: dict):
+    css = f"""
+    <style>
+      :root {{
+        --bg: {theme["bg"]};
+        --panel: {theme["panel"]};
+        --text: {theme["text"]};
+        --muted: {theme["muted"]};
+        --accent: {theme["accent"]};
+        --border: {theme["border"]};
+        --bubble_user: {theme["bubble_user"]};
+        --bubble_asst: {theme["bubble_asst"]};
+        --code_bg: {theme["code_bg"]};
+      }}
+
+      .stApp {{
+        background: var(--bg);
+        color: var(--text);
+      }}
+
+      section[data-testid="stSidebar"] {{
+        background: var(--panel);
+        border-right: 1px solid var(--border);
+      }}
+
+      /* Header title */
+      h1, h2, h3, p, span, label {{
+        color: var(--text) !important;
+      }}
+
+      a {{
+        color: var(--accent) !important;
+        text-decoration: none;
+      }}
+      a:hover {{ text-decoration: underline; }}
+
+      /* Chat message containers */
+      div[data-testid="stChatMessage"] {{
+        border-radius: 18px;
+        padding: 10px 12px;
+      }}
+
+      /* We color bubbles using a wrapper class we add via markdown */
+      .bubble-user {{
+        background: var(--bubble_user);
+        border: 1px solid var(--border);
+        border-radius: 16px;
+        padding: 10px 12px;
+      }}
+      .bubble-asst {{
+        background: var(--bubble_asst);
+        border: 1px solid var(--border);
+        border-radius: 16px;
+        padding: 10px 12px;
+      }}
+
+      /* Inputs and buttons */
+      .stTextArea textarea {{
+        border-radius: 12px !important;
+      }}
+      .stButton button {{
+        border-radius: 999px !important;
+        border: 1px solid var(--border) !important;
+      }}
+
+      /* Make popover trigger button look like an icon */
+      button[kind="secondary"] {{
+        border-radius: 999px !important;
+      }}
+
+      /* Code blocks */
+      code, pre {{
+        background: var(--code_bg) !important;
+      }}
+
+      /* Suggested chips */
+      .suggest-row {{
+        display: flex;
+        gap: 10px;
+        flex-wrap: wrap;
+        margin-top: 6px;
+      }}
+      .suggest-chip {{
+        display: inline-block;
+        padding: 10px 12px;
+        border: 1px solid var(--border);
+        border-radius: 999px;
+        background: var(--panel);
+        color: var(--text);
+        cursor: pointer;
+        user-select: none;
+      }}
+      .suggest-chip:hover {{
+        border-color: var(--accent);
+      }}
+    </style>
+    """
+    st.markdown(css, unsafe_allow_html=True)
+
+
 # -----------------------
 # Helpers
 # -----------------------
 def st_image_full_width(img_url: str):
-    """
-    HF Spaces sometimes pins an older streamlit build where st.image doesn't accept
-    use_container_width. Fall back to use_column_width.
-    """
+    # Streamlit API varies across HF images
     try:
         st.image(img_url, use_container_width=True)
     except TypeError:
         st.image(img_url, use_column_width=True)
 
 
-def get_brave_key() -> str:
-    # HF Spaces typically provides secrets as env vars.
+def get_brave_key() -> Optional[str]:
+    # HF Spaces secrets are usually env vars; also allow Streamlit secrets
     return os.getenv("BRAVE_API_KEY") or (st.secrets.get("BRAVE_API_KEY") if hasattr(st, "secrets") else None)
 
 
 def brave_search_snippets(query: str, top_k: int = 3) -> List[Dict[str, str]]:
     brave_api_key = get_brave_key()
     if not brave_api_key:
-        raise RuntimeError("Missing BRAVE_API_KEY (set Space secret or env var).")
+        raise RuntimeError("Missing BRAVE_API_KEY (set Space secret / env var).")
 
     headers = {"Accept": "application/json", "X-Subscription-Token": brave_api_key}
     params = {"q": query, "count": top_k}
 
-    resp = requests.get(
-        BRAVE_ENDPOINT,
-        headers=headers,
-        params=params,
-        timeout=TIMEOUT_SECS,
-    )
+    resp = requests.get(BRAVE_ENDPOINT, headers=headers, params=params, timeout=TIMEOUT_SECS)
     resp.raise_for_status()
     data = resp.json()
 
@@ -144,9 +235,7 @@ def brave_search_snippets(query: str, top_k: int = 3) -> List[Dict[str, str]]:
 
 
 def format_context_from_results(results: List[Dict[str, str]]) -> str:
-    """
-    Stable formatting + strip <strong> tags as requested.
-    """
+    # Stable formatting + strip <strong> tags exactly as requested
     if not results:
         return ""
 
@@ -159,18 +248,11 @@ def format_context_from_results(results: List[Dict[str, str]]) -> str:
         title = title.replace("<strong>", "").replace("</strong>", "")
         snippet = snippet.replace("<strong>", "").replace("</strong>", "")
 
-        blocks.append(
-            f"[{i}] {title}\n"
-            f"URL: {url}\n"
-            f"Snippet: {snippet}"
-        )
+        blocks.append(f"[{i}] {title}\nURL: {url}\nSnippet: {snippet}")
     return "\n\n".join(blocks)
 
 
-def typewriter_render(text: str, container, speed_chars_per_sec: float = 350.0):
-    """
-    TeapotAI.query returns a full string; mimic streaming with a typewriter effect.
-    """
+def typewriter_render(text: str, container, speed_chars_per_sec: float = 450.0):
     if not text:
         container.markdown("")
         return
@@ -183,13 +265,10 @@ def typewriter_render(text: str, container, speed_chars_per_sec: float = 350.0):
 
 
 # -----------------------
-# Model / TeapotAI loader
+# Model / TeapotAI loader (cached)
 # -----------------------
 @st.cache_resource
 def load_teapot_ai(model_name: str) -> TeapotAI:
-    """
-    Cached per model_name. TinyTeapot is loaded on startup via an explicit call.
-    """
     tokenizer = AutoTokenizer.from_pretrained(model_name)
     model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
 
@@ -205,81 +284,166 @@ def load_teapot_ai(model_name: str) -> TeapotAI:
 
 
 # -----------------------
-# UI Header
+# Session state init
 # -----------------------
-col1, col2 = st.columns([1, 3], vertical_alignment="center")
-with col1:
-    st_image_full_width(TEAPOT_LOGO_GIF)
-with col2:
-    st.markdown("## TeapotAI Chat")
-    st.caption("Brave Search (top 3 snippets) → context → TeapotAI.query()")
+if "theme" not in st.session_state:
+    st.session_state.theme = "Light"
+if "messages" not in st.session_state:
+    # Each message can optionally include sources/context metadata
+    # {"role": "assistant", "content": "...", "sources": [...], "context": "..."}
+    st.session_state.messages = []
+if "pending_query" not in st.session_state:
+    st.session_state.pending_query = None
 
-with st.sidebar:
-    st.header("Settings")
+# Inject theme CSS
+inject_css(LIGHT_THEME if st.session_state.theme == "Light" else DARK_THEME)
 
-    model_choice = st.radio(
-        "Model",
-        options=[MODEL_TINY, MODEL_LLM],
-        index=0,
-        help="TinyTeapot loads by default. Switching loads the other model (cached).",
-    )
 
-    system_prompt = st.text_area("System prompt", value=DEFAULT_SYSTEM_PROMPT, height=150)
-    show_sources = st.checkbox("Show sources/context", value=True)
-    typing_effect = st.checkbox("Typing effect", value=True)
+# -----------------------
+# Sidebar (ONLY configurable: theme toggle, system prompt, web search checkbox)
+# -----------------------
+with st.sidebar:
+    st.markdown("### Settings")
+    st.session_state.theme = st.radio("Theme", ["Light", "Dark"], horizontal=True, index=0 if st.session_state.theme == "Light" else 1)
+
+    system_prompt = st.text_area("System prompt", value=DEFAULT_SYSTEM_PROMPT, height=160)
+    use_web_search = st.checkbox("Use web search", value=True)
 
+    # Re-inject after theme changes (Streamlit reruns)
+    inject_css(LIGHT_THEME if st.session_state.theme == "Light" else DARK_THEME)
 
-# Requirement: load tiny model on startup regardless of selection
-_ = load_teapot_ai(MODEL_TINY)
+# -----------------------
+# Header
+# -----------------------
+c1, c2 = st.columns([1, 4], vertical_alignment="center")
+with c1:
+    st_image_full_width(TEAPOT_LOGO_GIF)
+with c2:
+    st.markdown("## TeapotAI Chat")
+    st.caption("A lightweight, grounded chat experience.")
 
-# Load selected model (cached after first load)
-teapot_ai = load_teapot_ai(model_choice)
+# Load tiny model on startup (your earlier requirement) and use it (no model toggle in UI now per your request)
+_ = load_teapot_ai(DEFAULT_MODEL)
+teapot_ai = load_teapot_ai(DEFAULT_MODEL)
 
 
 # -----------------------
-# Chat state
+# Suggested searches (shown only when chat is empty)
 # -----------------------
-if "messages" not in st.session_state:
-    st.session_state.messages = []  # [{"role": "user"/"assistant", "content": str}]
+if len(st.session_state.messages) == 0 and st.session_state.pending_query is None:
+    st.markdown("#### Try one of these")
+    cols = st.columns(3)
+    for i, q in enumerate(SUGGESTED_QUERIES):
+        with cols[i]:
+            if st.button(q, key=f"suggest_{i}"):
+                st.session_state.pending_query = q
+                st.rerun()
 
-for m in st.session_state.messages:
-    with st.chat_message(m["role"]):
-        st.markdown(m["content"])
 
-question = st.chat_input("Ask a question… (Brave will fetch top 3 snippets)")
-
-if question:
-    st.session_state.messages.append({"role": "user", "content": question})
-    with st.chat_message("user"):
-        st.markdown(question)
+# -----------------------
+# Render chat history
+# -----------------------
+for idx, m in enumerate(st.session_state.messages):
+    if m["role"] == "user":
+        with st.chat_message("user"):
+            st.markdown(f'<div class="bubble-user">{m["content"]}</div>', unsafe_allow_html=True)
+    else:
+        with st.chat_message("assistant"):
+            st.markdown(f'<div class="bubble-asst">{m["content"]}</div>', unsafe_allow_html=True)
+
+            # Info icon → expands sources/context + links
+            sources = m.get("sources") or []
+            context = m.get("context") or ""
+
+            if sources or context:
+                # Prefer popover if available, else expander
+                try:
+                    with st.popover("ℹ️", use_container_width=False):
+                        st.markdown("**Sources**")
+                        if sources:
+                            for j, s in enumerate(sources, start=1):
+                                title = s.get("title", "").strip() or f"Result {j}"
+                                url = s.get("url", "").strip()
+                                snippet = s.get("snippet", "").strip()
+                                if url:
+                                    st.markdown(f"- [{title}]({url})")
+                                else:
+                                    st.markdown(f"- {title}")
+                                if snippet:
+                                    st.caption(snippet)
+                        else:
+                            st.caption("(No sources returned.)")
+
+                        st.markdown("**Full context**")
+                        if context.strip():
+                            st.code(context)
+                        else:
+                            st.caption("(Empty context.)")
+                except Exception:
+                    with st.expander("ℹ️ Sources / Context"):
+                        st.markdown("**Sources**")
+                        if sources:
+                            for j, s in enumerate(sources, start=1):
+                                title = s.get("title", "").strip() or f"Result {j}"
+                                url = s.get("url", "").strip()
+                                snippet = s.get("snippet", "").strip()
+                                if url:
+                                    st.markdown(f"- [{title}]({url})")
+                                else:
+                                    st.markdown(f"- {title}")
+                                if snippet:
+                                    st.caption(snippet)
+                        else:
+                            st.caption("(No sources returned.)")
+
+                        st.markdown("**Full context**")
+                        if context.strip():
+                            st.code(context)
+                        else:
+                            st.caption("(Empty context.)")
 
-    # Brave search context
-    try:
-        results = brave_search_snippets(question, top_k=TOP_K)
-        context = format_context_from_results(results)
-    except Exception:
-        results = []
-        context = ""
 
-    # TeapotAI query (context includes Brave results)
+# -----------------------
+# Input handling
+# -----------------------
+user_input = st.chat_input("Ask a question…")
+
+if st.session_state.pending_query and not user_input:
+    user_input = st.session_state.pending_query
+    st.session_state.pending_query = None
+
+if user_input:
+    # Add user message
+    st.session_state.messages.append({"role": "user", "content": user_input})
+
+    # Build context (optional web search)
+    results: List[Dict[str, str]] = []
+    context = ""
+
+    if use_web_search:
+        try:
+            results = brave_search_snippets(user_input, top_k=TOP_K)
+            context = format_context_from_results(results)
+        except Exception:
+            # If web search fails, keep context empty so system prompt can enforce refusal behavior
+            results = []
+            context = ""
+
+    # Query TeapotAI
     answer = teapot_ai.query(
-        query=question,
+        query=user_input,
         context=context,
         system_prompt=system_prompt,
     )
 
-    with st.chat_message("assistant"):
-        if show_sources:
-            with st.expander("Sources / Context used", expanded=False):
-                if context.strip():
-                    st.code(context)
-                else:
-                    st.write("(No search context returned.)")
-
-        placeholder = st.empty()
-        if typing_effect:
-            typewriter_render(answer, placeholder, speed_chars_per_sec=350.0)
-        else:
-            placeholder.markdown(answer)
-
-    st.session_state.messages.append({"role": "assistant", "content": answer})
+    # Append assistant message with metadata for info popover
+    st.session_state.messages.append(
+        {
+            "role": "assistant",
+            "content": answer,
+            "sources": results,
+            "context": context,
+        }
+    )
+
+    st.rerun()