Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,10 +1,11 @@
 import os
 import time
 import requests
+import threading
 
 import streamlit as st
 import torch
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, TextIteratorStreamer
 
 # Optional LangSmith (trace + feedback)
 try:
@@ -105,7 +106,6 @@ if "messages" not in st.session_state:
 if "seeded" not in st.session_state:
     st.session_state.seeded = False
 
-# Seed exactly once on first load
 if (not st.session_state.seeded) and (len(st.session_state.messages) == 0):
     st.session_state.messages = [SAMPLE_USER_MSG, SAMPLE_ASSISTANT_MSG]
     st.session_state.seeded = True
@@ -204,39 +204,6 @@ def count_tokens(text: str) -> int:
     return len(tokenizer.encode(text)) if text else 0
 
 
-# =========================
-# LANGSMITH-TRACED ANSWER FUNCTION
-# =========================
-if traceable:
-
-    @traceable(name="teapot_answer")
-    def traced_answer(context: str, system_prompt: str, question: str) -> str:
-        prompt = f"{context}\n{system_prompt}\n{question}\n"
-        inputs = tokenizer(prompt, return_tensors="pt").to(device)
-        with torch.no_grad():
-            out = model.generate(
-                **inputs,
-                max_new_tokens=MAX_NEW_TOKENS,
-                do_sample=False,
-                num_beams=1,
-            )
-        return tokenizer.decode(out[0], skip_special_tokens=True)
-
-else:
-
-    def traced_answer(context: str, system_prompt: str, question: str) -> str:
-        prompt = f"{context}\n{system_prompt}\n{question}\n"
-        inputs = tokenizer(prompt, return_tensors="pt").to(device)
-        with torch.no_grad():
-            out = model.generate(
-                **inputs,
-                max_new_tokens=MAX_NEW_TOKENS,
-                do_sample=False,
-                num_beams=1,
-            )
-        return tokenizer.decode(out[0], skip_special_tokens=True)
-
-
 def get_trace_id_if_available() -> str | None:
     if not get_current_run_tree:
         return None
@@ -247,6 +214,51 @@ def get_trace_id_if_available() -> str | None:
         return None
 
 
+# =========================
+# TRACED STREAMING GENERATION
+# =========================
+def stream_answer_tokens(context: str, system_prompt: str, question: str):
+    """
+    Yields decoded text increments as they are generated (real streaming).
+    """
+    prompt = f"{context}\n{system_prompt}\n{question}\n"
+    inputs = tokenizer(prompt, return_tensors="pt").to(device)
+
+    streamer = TextIteratorStreamer(
+        tokenizer,
+        skip_prompt=True,
+        skip_special_tokens=True,
+    )
+
+    gen_kwargs = dict(
+        **inputs,
+        max_new_tokens=MAX_NEW_TOKENS,
+        do_sample=False,
+        num_beams=1,
+        streamer=streamer,
+    )
+
+    # Run generate in a background thread so we can iterate streamer in the main thread.
+    t = threading.Thread(target=model.generate, kwargs=gen_kwargs, daemon=True)
+    t.start()
+
+    for text in streamer:
+        # text can be tiny chunks; yield as-is
+        yield text
+
+
+if traceable:
+    # Wrap a traced function around the streaming loop (LangSmith will see a single run)
+    @traceable(name="teapot_answer_stream")
+    def traced_stream_answer(context: str, system_prompt: str, question: str):
+        for chunk in stream_answer_tokens(context, system_prompt, question):
+            yield chunk
+else:
+    def traced_stream_answer(context: str, system_prompt: str, question: str):
+        for chunk in stream_answer_tokens(context, system_prompt, question):
+            yield chunk
+
+
 # =========================
 # FEEDBACK HANDLER (attached to trace_id)
 # =========================
@@ -257,7 +269,6 @@ def handle_feedback(idx: int):
     msg = st.session_state.messages[idx]
     trace_id = msg.get("trace_id")
 
-    # Attach feedback to this traced run
     if ls_client and trace_id:
         score = 1 if val == "👍" else 0
         try:
@@ -272,9 +283,16 @@ def handle_feedback(idx: int):
 
 
 # =========================
-# RENDER HISTORY
-# Row 1: message + feedback
-# Row 2: inspect + debug metrics
+# INPUT FIRST (so new user msg renders immediately)
+# =========================
+query = st.chat_input("Ask a question...")
+
+if query:
+    st.session_state.messages.append({"role": "user", "content": query})
+
+
+# =========================
+# RENDER HISTORY (now includes latest user msg)
 # =========================
 for i, msg in enumerate(st.session_state.messages):
     with st.chat_message(msg["role"]):
@@ -299,7 +317,6 @@ for i, msg in enumerate(st.session_state.messages):
 
         # Row 2
         inspect_col, metrics_col = st.columns([12, 1], vertical_alignment="center")
-
         with inspect_col:
             st.caption(
                 f"🔎 {msg.get('search_time', 0.0):.2f}s (search) "
@@ -307,7 +324,6 @@ for i, msg in enumerate(st.session_state.messages):
                 f"⚡ {msg.get('tps', 0.0):.1f} tok/s  "
                 f"🧾 {msg.get('input_tokens', 0)} input tokens • {msg.get('output_tokens', 0)} output tokens"
             )
-
         with metrics_col:
             with st.popover("ℹ️", help="Inspect"):
                 st.markdown("**Context**")
@@ -319,15 +335,10 @@ for i, msg in enumerate(st.session_state.messages):
 
 
 # =========================
-# INPUT + GENERATE (NO RERUN / NO FLASH)
+# GENERATE ONLY IF THIS RUN RECEIVED A NEW QUERY
+# (We detect by: query is not None)
 # =========================
-query = st.chat_input("Ask a question...")
-
 if query:
-    # Persist user message
-    st.session_state.messages.append({"role": "user", "content": query})
-
-    # Generate immediately in the same run (no st.rerun)
     question = query
 
     # Web search
@@ -338,9 +349,8 @@ if query:
     prompt = f"{context}\n{system_prompt}\n{question}\n"
     input_tokens = count_tokens(prompt)
 
-    # Assistant response
+    # Stream assistant response (real streaming)
     with st.chat_message("assistant"):
-        # Row 1: message + feedback (feedback disabled until persisted)
         msg_col, fb_col = st.columns([14, 1], vertical_alignment="center")
         with msg_col:
             placeholder = st.empty()
@@ -348,20 +358,16 @@ if query:
             st.feedback("thumbs", key="live_fb", disabled=True)
 
         start = time.perf_counter()
-        answer = traced_answer(context, system_prompt, question)
-        trace_id = get_trace_id_if_available()
 
-        # Typewriter render (reduce updates to avoid jitter)
         buf = ""
-        for j, ch in enumerate(answer, 1):
-            buf += ch
-            if j % 6 == 0:  # update every 6 chars
-                placeholder.markdown(buf)
-                time.sleep(0.001)
-        placeholder.markdown(buf)
+        placeholder.markdown("")  # ensures first token updates a visible element
+        for chunk in traced_stream_answer(context, system_prompt, question):
+            buf += chunk
+            placeholder.markdown(buf)
 
+        trace_id = get_trace_id_if_available()
         gen_time = time.perf_counter() - start
-        output_tokens = count_tokens(answer)
+        output_tokens = count_tokens(buf)
         tps = output_tokens / gen_time if gen_time > 0 else 0.0
 
         # Row 2: inspect + metrics
@@ -384,11 +390,11 @@ if query:
                 st.markdown("**Prompt**")
                 st.code(prompt, language="text")
 
-    # Persist assistant message (so it shows on subsequent runs)
+    # Persist assistant message
     st.session_state.messages.append(
         {
             "role": "assistant",
-            "content": answer,
+            "content": buf,
             "context": context,
             "system_prompt": system_prompt,
             "question": question,