Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,7 +1,7 @@
 import os
 import time
-import requests
 import threading
+import requests
 
 import streamlit as st
 import torch
@@ -72,11 +72,7 @@ DEFAULT_SYSTEM_PROMPT = (
 )
 
 SAMPLE_SYSTEM_PROMPT = DEFAULT_SYSTEM_PROMPT
-
-SAMPLE_CONTEXT = (
-    "Teapot is an open-source AI assistant optimized for running on low-end cpu devices."
-)
-
+SAMPLE_CONTEXT = "Teapot is an open-source AI assistant optimized for running on low-end cpu devices."
 SAMPLE_ANSWER = "I am Teapot, an open-source AI assistant optimized for running on low-end cpu devices."
 SAMPLE_PROMPT = f"{SAMPLE_CONTEXT}\n{SAMPLE_SYSTEM_PROMPT}\n{SAMPLE_QUESTION}\n"
 
@@ -106,6 +102,7 @@ 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
@@ -215,13 +212,41 @@ def get_trace_id_if_available() -> str | None:
 
 
 # =========================
-# TRACED STREAMING GENERATION
+# FEEDBACK HANDLER (attached to trace_id)
+# =========================
+def handle_feedback(idx: int):
+    val = st.session_state.get(f"fb_{idx}")
+    st.session_state.messages[idx]["feedback"] = val
+
+    msg = st.session_state.messages[idx]
+    trace_id = msg.get("trace_id")
+
+    if ls_client and trace_id:
+        score = 1 if val == "👍" else 0
+        try:
+            ls_client.create_feedback(
+                trace_id=trace_id,
+                key="thumb_rating",
+                score=score,
+                comment="thumbs_up" if score else "thumbs_down",
+            )
+        except Exception:
+            pass
+
+
+# =========================
+# STREAMING + LANGSMITH FIX
+# - We do NOT trace a generator.
+# - We stream to UI while returning a SINGLE final string.
 # =========================
-def stream_answer_tokens(context: str, system_prompt: str, question: str):
+_UI_STREAM = {"placeholder": None}  # set per-request
+
+
+def _generate_with_streamer(prompt: str) -> str:
     """
-    Yields decoded text increments as they are generated (real streaming).
+    Runs model.generate with a TextIteratorStreamer and updates a Streamlit placeholder
+    as chunks arrive. Returns the final full text.
     """
-    prompt = f"{context}\n{system_prompt}\n{question}\n"
     inputs = tokenizer(prompt, return_tensors="pt").to(device)
 
     streamer = TextIteratorStreamer(
@@ -238,52 +263,38 @@ def stream_answer_tokens(context: str, system_prompt: str, question: str):
         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
+    buf = ""
+    ph = _UI_STREAM.get("placeholder")
+    if ph is not None:
+        ph.markdown("")  # ensure element exists before first chunk
 
+    for chunk in streamer:
+        buf += chunk
+        if ph is not None:
+            ph.markdown(buf)
 
-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
+    return buf
 
 
-# =========================
-# FEEDBACK HANDLER (attached to trace_id)
-# =========================
-def handle_feedback(idx: int):
-    val = st.session_state.get(f"fb_{idx}")
-    st.session_state.messages[idx]["feedback"] = val
+if traceable:
 
-    msg = st.session_state.messages[idx]
-    trace_id = msg.get("trace_id")
+    @traceable(name="teapot_answer")
+    def traced_answer_streaming(context: str, system_prompt: str, question: str) -> str:
+        prompt = f"{context}\n{system_prompt}\n{question}\n"
+        return _generate_with_streamer(prompt)
 
-    if ls_client and trace_id:
-        score = 1 if val == "👍" else 0
-        try:
-            ls_client.create_feedback(
-                trace_id=trace_id,
-                key="thumb_rating",
-                score=score,
-                comment="thumbs_up" if score else "thumbs_down",
-            )
-        except Exception:
-            pass
+else:
+
+    def traced_answer_streaming(context: str, system_prompt: str, question: str) -> str:
+        prompt = f"{context}\n{system_prompt}\n{question}\n"
+        return _generate_with_streamer(prompt)
 
 
 # =========================
-# INPUT FIRST (so new user msg renders immediately)
+# INPUT FIRST (so latest user msg renders immediately)
 # =========================
 query = st.chat_input("Ask a question...")
 
@@ -292,7 +303,9 @@ if query:
 
 
 # =========================
-# RENDER HISTORY (now includes latest user msg)
+# RENDER HISTORY
+# Row 1: message + feedback
+# Row 2: inspect + debug metrics
 # =========================
 for i, msg in enumerate(st.session_state.messages):
     with st.chat_message(msg["role"]):
@@ -336,7 +349,6 @@ for i, msg in enumerate(st.session_state.messages):
 
 # =========================
 # GENERATE ONLY IF THIS RUN RECEIVED A NEW QUERY
-# (We detect by: query is not None)
 # =========================
 if query:
     question = query
@@ -349,25 +361,25 @@ if query:
     prompt = f"{context}\n{system_prompt}\n{question}\n"
     input_tokens = count_tokens(prompt)
 
-    # Stream assistant response (real streaming)
+    # Assistant response (stream to UI, return full string for LangSmith)
     with st.chat_message("assistant"):
+        # Row 1: message + feedback (disabled live)
         msg_col, fb_col = st.columns([14, 1], vertical_alignment="center")
         with msg_col:
             placeholder = st.empty()
         with fb_col:
             st.feedback("thumbs", key="live_fb", disabled=True)
 
-        start = time.perf_counter()
-
-        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)
+        _UI_STREAM["placeholder"] = placeholder
 
+        start = time.perf_counter()
+        answer = traced_answer_streaming(context, system_prompt, question)
         trace_id = get_trace_id_if_available()
         gen_time = time.perf_counter() - start
-        output_tokens = count_tokens(buf)
+
+        _UI_STREAM["placeholder"] = None  # cleanup
+
+        output_tokens = count_tokens(answer)
         tps = output_tokens / gen_time if gen_time > 0 else 0.0
 
         # Row 2: inspect + metrics
@@ -394,7 +406,7 @@ if query:
     st.session_state.messages.append(
         {
             "role": "assistant",
-            "content": buf,
+            "content": answer,
             "context": context,
             "system_prompt": system_prompt,
             "question": question,