Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,6 +1,8 @@
+# streamlit_app.py
 import os
 import re
 import time
+import warnings
 from typing import List, Dict
 
 import requests
@@ -11,6 +13,13 @@ from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 from teapotai import TeapotAI
 
 
+# -----------------------
+# Silence noisy 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
 # -----------------------
@@ -31,36 +40,30 @@ st.set_page_config(
 
 CUSTOM_CSS = f"""
 <style>
-/* App background */
 .stApp {{
   background: {TEA_BG};
   color: {TEA_TEXT};
 }}
 
-/* Sidebar */
 section[data-testid="stSidebar"] {{
   background: {TEA_PANEL};
   border-right: 1px solid {TEA_BORDER};
 }}
 
-/* Chat bubbles */
 div[data-testid="stChatMessage"] {{
   border-radius: 16px;
   padding: 8px 10px;
 }}
 
-/* Inputs */
 .stTextInput > div > div, .stTextArea > div > div {{
   border-radius: 12px !important;
 }}
 
-/* Buttons */
 .stButton button {{
   border-radius: 12px;
   border: 1px solid {TEA_BORDER};
 }}
 
-/* Accent-ish links */
 a {{
   color: {TEA_ACCENT} !important;
 }}
@@ -79,7 +82,6 @@ TIMEOUT_SECS = 15
 MODEL_TINY = "teapotai/tinyteapot"
 MODEL_LLM = "teapotai/teapotllm"
 
-
 DEFAULT_SYSTEM_PROMPT = (
     "You are Teapot, an open-source AI assistant optimized for low-end devices, "
     "providing short, accurate responses without hallucinating while excelling at "
@@ -96,15 +98,26 @@ DEFAULT_DOCUMENTS = [
 # -----------------------
 # 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.
+    """
+    try:
+        st.image(img_url, use_container_width=True)
+    except TypeError:
+        st.image(img_url, use_column_width=True)
+
+
 def get_brave_key() -> str:
-    # Streamlit Cloud secrets support + local env var support
-    return st.secrets.get("BRAVE_API_KEY") if hasattr(st, "secrets") and "BRAVE_API_KEY" in st.secrets else os.getenv("BRAVE_API_KEY")
+    # HF Spaces typically provides secrets as env vars.
+    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 env var or Streamlit secrets).")
+        raise RuntimeError("Missing BRAVE_API_KEY (set Space secret or env var).")
 
     headers = {"Accept": "application/json", "X-Subscription-Token": brave_api_key}
     params = {"q": query, "count": top_k}
@@ -132,7 +145,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; plus you asked to strip <strong> tags.
+    Stable formatting + strip <strong> tags as requested.
     """
     if not results:
         return ""
@@ -143,7 +156,6 @@ 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()
 
-        # strip <strong> tags specifically, as requested
         title = title.replace("<strong>", "").replace("</strong>", "")
         snippet = snippet.replace("<strong>", "").replace("</strong>", "")
 
@@ -155,13 +167,28 @@ def format_context_from_results(results: List[Dict[str, str]]) -> str:
     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.
+    """
+    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)
+
+
 # -----------------------
 # Model / TeapotAI loader
 # -----------------------
 @st.cache_resource
 def load_teapot_ai(model_name: str) -> TeapotAI:
     """
-    Cached per model_name. TinyTeapot will be loaded on startup (we call it once).
+    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)
@@ -170,37 +197,19 @@ def load_teapot_ai(model_name: str) -> TeapotAI:
     model.to(device)
     model.eval()
 
-    teapot_ai = TeapotAI(
+    return TeapotAI(
         tokenizer=tokenizer,
         model=model,
         documents=DEFAULT_DOCUMENTS,
     )
-    return teapot_ai
-
-
-def typewriter_render(text: str, container, speed_chars_per_sec: float = 250.0):
-    """
-    TeapotAI.query isn't streamed (in this code), so we do a simple typewriter effect.
-    """
-    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)
 
 
 # -----------------------
-# UI
+# UI Header
 # -----------------------
-# Header with logo
 col1, col2 = st.columns([1, 3], vertical_alignment="center")
 with col1:
-    # Streamlit will fetch the gif directly
-    st.image(TEAPOT_LOGO_GIF, use_container_width=True)
+    st_image_full_width(TEAPOT_LOGO_GIF)
 with col2:
     st.markdown("## TeapotAI Chat")
     st.caption("Brave Search (top 3 snippets) → context → TeapotAI.query()")
@@ -215,36 +224,36 @@ with st.sidebar:
         help="TinyTeapot loads by default. Switching loads the other model (cached).",
     )
 
-    system_prompt = st.text_area("System prompt", value=DEFAULT_SYSTEM_PROMPT, height=140)
+    system_prompt = st.text_area("System prompt", value=DEFAULT_SYSTEM_PROMPT, height=150)
     show_sources = st.checkbox("Show sources/context", value=True)
-
-    # Optional: “typing” effect
     typing_effect = st.checkbox("Typing effect", value=True)
 
 
-# Load TinyTeapot on startup, regardless of current selection (your requirement)
+# Requirement: load tiny model on startup regardless of selection
 _ = load_teapot_ai(MODEL_TINY)
 
 # Load selected model (cached after first load)
 teapot_ai = load_teapot_ai(model_choice)
 
+
+# -----------------------
 # Chat state
+# -----------------------
 if "messages" not in st.session_state:
     st.session_state.messages = []  # [{"role": "user"/"assistant", "content": str}]
 
-# Render history
 for m in st.session_state.messages:
     with st.chat_message(m["role"]):
         st.markdown(m["content"])
 
-question = st.chat_input("Ask a question… (@sources are fetched via Brave)")
+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)
 
-    # Brave context
+    # Brave search context
     try:
         results = brave_search_snippets(question, top_k=TOP_K)
         context = format_context_from_results(results)
@@ -252,8 +261,7 @@ if question:
         results = []
         context = ""
 
-    # TeapotAI query (context comes from Brave)
-    # NOTE: you explicitly want context="" param to hold Brave results after stripping strong tags.
+    # TeapotAI query (context includes Brave results)
     answer = teapot_ai.query(
         query=question,
         context=context,