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, Any, Optional
+from typing import List, Dict, Optional
 
 import requests
 import streamlit as st
@@ -14,14 +14,14 @@ from teapotai import TeapotAI
 
 
 # -----------------------
-# Warnings (optional)
+# Optional: quiet noisy warnings from deps
 # -----------------------
 warnings.filterwarnings("ignore", message="pkg_resources is deprecated as an API.*")
 warnings.filterwarnings("ignore", message='Field name "schema" in "TeapotTool" shadows.*')
 
 
 # -----------------------
-# Brand / Assets
+# Config
 # -----------------------
 TEAPOT_LOGO_GIF = "https://teapotai.com/assets/logo.gif"
 
@@ -31,13 +31,7 @@ SUGGESTED_QUERIES = [
     "What is the weather like in NYC today?",
 ]
 
-
-# -----------------------
-# Models
-# -----------------------
-MODEL_TINY = "teapotai/tinyteapot"
-MODEL_LLM = "teapotai/teapotllm"  # if you keep the toggle elsewhere later
-DEFAULT_MODEL = MODEL_TINY
+MODEL_NAME = "teapotai/tinyteapot"
 
 DEFAULT_SYSTEM_PROMPT = (
     "You are Teapot, an open-source AI assistant optimized for low-end devices, "
@@ -51,154 +45,23 @@ DEFAULT_DOCUMENTS = [
     """Teapot (Tiny Teapot) is an open-source small language model (~77 million parameters) fine-tuned on synthetic data and optimized to run locally on resource-constrained devices such as smartphones and CPUs. Teapot is trained to only answer using context from documents, reducing hallucinations. Teapot can perform a variety of tasks, including hallucination-resistant Question Answering (QnA), Retrieval-Augmented Generation (RAG), and JSON extraction. TeapotLLM is a fine tune of flan-t5-large that was trained on synthetic data generated by Deepseek v3 TeapotLLM can be hosted on low-power devices with as little as 2GB of CPU RAM such as a Raspberry Pi. Teapot is a model built by and for the community."""
 ]
 
-
-# -----------------------
-# Web search (Brave) config
-# -----------------------
+# Brave Search
 BRAVE_ENDPOINT = "https://api.search.brave.com/res/v1/web/search"
 TOP_K = 3
 TIMEOUT_SECS = 15
 
 
 # -----------------------
-# Streamlit page config
+# Streamlit setup (no custom theming)
 # -----------------------
 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):
-    # Streamlit API varies across HF images
+    # Streamlit API varies across builds
     try:
         st.image(img_url, use_container_width=True)
     except TypeError:
@@ -206,23 +69,22 @@ def st_image_full_width(img_url: str):
 
 
 def get_brave_key() -> Optional[str]:
-    # HF Spaces secrets are usually env vars; also allow Streamlit secrets
+    # HF Spaces secrets are commonly env vars; support st.secrets too
     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 / env var).")
+    key = get_brave_key()
+    if not key:
+        raise RuntimeError("Missing BRAVE_API_KEY (set as a Space secret / env var).")
 
-    headers = {"Accept": "application/json", "X-Subscription-Token": brave_api_key}
+    headers = {"Accept": "application/json", "X-Subscription-Token": key}
     params = {"q": query, "count": top_k}
+    r = requests.get(BRAVE_ENDPOINT, headers=headers, params=params, timeout=TIMEOUT_SECS)
+    r.raise_for_status()
+    data = r.json()
 
-    resp = requests.get(BRAVE_ENDPOINT, headers=headers, params=params, timeout=TIMEOUT_SECS)
-    resp.raise_for_status()
-    data = resp.json()
-
-    results = []
+    results: List[Dict[str, str]] = []
     web = data.get("web") or {}
     items = web.get("results") or []
     for item in items[:top_k]:
@@ -235,7 +97,6 @@ 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 exactly as requested
     if not results:
         return ""
 
@@ -245,97 +106,121 @@ def format_context_from_results(results: List[Dict[str, str]]) -> str:
         url = re.sub(r"\s+", " ", r.get("url", "")).strip()
         snippet = re.sub(r"\s+", " ", r.get("snippet", "")).strip()
 
+        # per your requirement: strip <strong> tags
         title = title.replace("<strong>", "").replace("</strong>", "")
         snippet = snippet.replace("<strong>", "").replace("</strong>", "")
 
         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 = 450.0):
+def count_tokens(tokenizer: AutoTokenizer, text: str) -> int:
     if not text:
-        container.markdown("")
-        return
-    delay = 1.0 / max(speed_chars_per_sec, 1.0)
-    out = ""
-    for ch in text:
-        out += ch
-        container.markdown(out)
-        time.sleep(delay)
+        return 0
+    try:
+        return len(tokenizer.encode(text))
+    except Exception:
+        return 0
+
+
+def render_sources_popover(sources: List[Dict[str, str]], context: str):
+    """
+    Renders ℹ️ popover if available; otherwise uses expander.
+    """
+    def _body():
+        st.markdown("**Sources**")
+        if sources:
+            for j, s in enumerate(sources, start=1):
+                title = (s.get("title") or "").strip() or f"Result {j}"
+                url = (s.get("url") or "").strip()
+                snippet = (s.get("snippet") or "").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.)")
+
+    try:
+        with st.popover("ℹ️"):
+            _body()
+    except Exception:
+        with st.expander("ℹ️ Sources / Context"):
+            _body()
 
 
 # -----------------------
-# Model / TeapotAI loader (cached)
+# Load model + TeapotAI (cached)
 # -----------------------
 @st.cache_resource
-def load_teapot_ai(model_name: str) -> TeapotAI:
-    tokenizer = AutoTokenizer.from_pretrained(model_name)
-    model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
+def load_teapot_ai_and_tokenizer():
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+    model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME)
 
     device = "cuda" if torch.cuda.is_available() else "cpu"
     model.to(device)
     model.eval()
 
-    return TeapotAI(
+    teapot_ai = TeapotAI(
         tokenizer=tokenizer,
         model=model,
         documents=DEFAULT_DOCUMENTS,
     )
+    return teapot_ai, tokenizer
 
 
 # -----------------------
-# Session state init
+# Session state
 # -----------------------
-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": "..."}
+    # Each assistant message includes: sources/context + timing/tokens
     st.session_state.messages = []
 if "pending_query" not in st.session_state:
     st.session_state.pending_query = None
 
-# Inject theme CSS
-inject_css(LIGHT_THEME if st.session_state.theme == "Light" else DARK_THEME)
+
+# -----------------------
+# Header
+# -----------------------
+c1, c2 = st.columns([1, 5], vertical_alignment="center")
+with c1:
+    st_image_full_width(TEAPOT_LOGO_GIF)
+with c2:
+    st.markdown("## TeapotAI Chat")
+    st.caption("Fast, grounded answers — with optional web context.")
 
 
 # -----------------------
-# Sidebar (ONLY configurable: theme toggle, system prompt, web search checkbox)
+# Sidebar (ONLY: system prompt + web search toggle)
 # -----------------------
 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)
-
-# -----------------------
-# 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 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)
+# Load tiny model on startup
+teapot_ai, hf_tokenizer = load_teapot_ai_and_tokenizer()
 
 
 # -----------------------
-# Suggested searches (shown only when chat is empty)
+# Suggested queries on empty chat
 # -----------------------
 if len(st.session_state.messages) == 0 and st.session_state.pending_query is None:
-    st.markdown("#### Try one of these")
+    st.markdown("#### Suggested")
     cols = st.columns(3)
     for i, q in enumerate(SUGGESTED_QUERIES):
         with cols[i]:
-            if st.button(q, key=f"suggest_{i}"):
+            if st.button(q, key=f"suggest_{i}", use_container_width=True):
                 st.session_state.pending_query = q
                 st.rerun()
 
@@ -343,64 +228,33 @@ if len(st.session_state.messages) == 0 and st.session_state.pending_query is Non
 # -----------------------
 # Render chat history
 # -----------------------
-for idx, m in enumerate(st.session_state.messages):
+for m in 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)
+            st.markdown(m["content"])
     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.)")
+            st.markdown(m["content"])
+
+            # metadata row
+            meta_cols = st.columns([1, 3, 3, 5])
+            with meta_cols[0]:
+                render_sources_popover(m.get("sources", []), m.get("context", ""))
+
+            # tokens/sec, and token counts
+            tps = m.get("tps", None)
+            out_toks = m.get("output_tokens", None)
+            secs = m.get("seconds", None)
+
+            with meta_cols[1]:
+                if tps is not None:
+                    st.caption(f"⚡ {tps:.1f} tokens/s")
+            with meta_cols[2]:
+                if out_toks is not None:
+                    st.caption(f"🧮 {out_toks} output tokens")
+            with meta_cols[3]:
+                if secs is not None:
+                    st.caption(f"⏱️ {secs:.2f}s")
 
 
 # -----------------------
@@ -413,37 +267,41 @@ if st.session_state.pending_query and not user_input:
     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]] = []
+    sources: 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)
+            sources = brave_search_snippets(user_input, top_k=TOP_K)
+            context = format_context_from_results(sources)
         except Exception:
-            # If web search fails, keep context empty so system prompt can enforce refusal behavior
-            results = []
+            sources = []
             context = ""
 
-    # Query TeapotAI
+    # Teapot inference + timing
+    t0 = time.perf_counter()
     answer = teapot_ai.query(
         query=user_input,
         context=context,
         system_prompt=system_prompt,
     )
+    t1 = time.perf_counter()
+
+    elapsed = max(t1 - t0, 1e-6)
+    output_tokens = count_tokens(hf_tokenizer, answer)
+    tps = output_tokens / elapsed if output_tokens > 0 else 0.0
 
-    # Append assistant message with metadata for info popover
     st.session_state.messages.append(
         {
             "role": "assistant",
             "content": answer,
-            "sources": results,
+            "sources": sources,
             "context": context,
+            "seconds": elapsed,
+            "output_tokens": output_tokens,
+            "tps": tps,
         }
     )
-
     st.rerun()