Update src/streamlit_app.py

src/streamlit_app.py

@@ -5,13 +5,17 @@ import requests
 
 import streamlit as st
 import torch
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, TextIteratorStreamer
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 
-# Optional LangSmith
+# Optional LangSmith (trace + feedback)
 try:
     from langsmith import Client as LangSmithClient
+    from langsmith import traceable
+    from langsmith.run_helpers import get_current_run_tree
 except Exception:
     LangSmithClient = None
+    traceable = None
+    get_current_run_tree = None
 
 
 # =========================
@@ -46,8 +50,8 @@ tokenizer, model, device = load_model()
 # =========================
 @st.cache_resource
 def get_langsmith():
-    key = os.getenv("LANGCHAIN_API_KEY") or os.getenv("LANGSMITH_API_KEY")
-    if key and LangSmithClient:
+    key = os.getenv("LANGCHAIN_API_KEY") or os.getenv("LANGSMITH_API_KEY") or os.getenv("LANGCHAIN_TRACING_V2")
+    if (os.getenv("LANGCHAIN_API_KEY") or os.getenv("LANGSMITH_API_KEY")) and LangSmithClient:
         return LangSmithClient()
     return None
 
@@ -60,14 +64,17 @@ ls_client = get_langsmith()
 # =========================
 if "messages" not in st.session_state:
     st.session_state.messages = []
+if "needs_answer" not in st.session_state:
+    st.session_state.needs_answer = False
 
 
 # =========================
-# HEADER
+# HEADER (prevent logo flash)
+# Use a fixed pixel width to avoid layout shift / big flash.
 # =========================
-col1, col2 = st.columns([1, 6])
+col1, col2 = st.columns([1, 7], vertical_alignment="center")
 with col1:
-    st.image(LOGO_URL, use_column_width=True)
+    st.image(LOGO_URL, width=56)  # fixed width prevents "flash huge"
 with col2:
     st.markdown("## TeapotAI Chat")
     st.caption("Grounded answers with web context")
@@ -124,8 +131,7 @@ def web_search_snippets(query: str):
 
     snippets = []
     for item in data.get("web", {}).get("results", [])[:TOP_K_SEARCH]:
-        desc = (item.get("description") or "")
-        desc = desc.replace("<strong>", "").replace("</strong>", "").strip()
+        desc = (item.get("description") or "").replace("<strong>", "").replace("</strong>", "").strip()
         if desc:
             snippets.append(desc)
 
@@ -137,18 +143,16 @@ def web_search_snippets(query: str):
 # =========================
 def truncate_context(web_ctx: str, local_ctx: str, system: str, question: str) -> str:
     ctx = f"{web_ctx}\n\n{local_ctx}".strip()
-
     base = f"\n{system}\n{question}\n"
     base_tokens = tokenizer.encode(base)
     budget = MAX_INPUT_TOKENS - len(base_tokens)
-
     if budget <= 0:
         return ""
-
-    ctx_tokens = tokenizer.encode(ctx) if ctx else []
+    if not ctx:
+        return ""
+    ctx_tokens = tokenizer.encode(ctx)
     if len(ctx_tokens) <= budget:
         return ctx
-
     return tokenizer.decode(ctx_tokens[-budget:], skip_special_tokens=True)
 
 
@@ -157,44 +161,68 @@ def count_tokens(text: str) -> int:
 
 
 # =========================
-# STREAM GENERATION
+# LANGSMITH-TRACED ANSWER FUNCTION
+# (signature exactly: context, system_prompt, question -> answer)
 # =========================
-def stream_generate(prompt: str):
-    inputs = tokenizer(prompt, return_tensors="pt").to(device)
-    streamer = TextIteratorStreamer(tokenizer, skip_special_tokens=True)
-
-    def run():
-        model.generate(
-            **inputs,
-            max_new_tokens=MAX_NEW_TOKENS,
-            do_sample=False,
-            num_beams=1,
-            streamer=streamer,
-        )
+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,
+            )
+        text = tokenizer.decode(out[0], skip_special_tokens=True)
+        return text
+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)
 
-    threading.Thread(target=run, daemon=True).start()
 
-    acc = ""
-    for chunk in streamer:
-        acc += chunk
-        yield acc
+def get_trace_id_if_available() -> str | None:
+    # Works when running inside a @traceable function call
+    if not get_current_run_tree:
+        return None
+    try:
+        run = get_current_run_tree()
+        return str(run.id) if run and getattr(run, "id", None) else None
+    except Exception:
+        return None
 
 
 # =========================
-# FEEDBACK HANDLER
+# 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]
-    if ls_client and msg.get("run_id"):
+    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:
+            # LangSmith SDK supports trace_id= for feedback association
             ls_client.create_feedback(
-                run_id=msg["run_id"],
+                trace_id=trace_id,
                 key="thumb_rating",
                 score=score,
+                comment="thumbs_up" if score else "thumbs_down",
             )
         except Exception:
             pass
@@ -205,26 +233,43 @@ def handle_feedback(idx: int):
 # =========================
 for i, msg in enumerate(st.session_state.messages):
     with st.chat_message(msg["role"]):
+        if msg["role"] == "user":
+            st.markdown(msg["content"])
+            continue
+
+        # Assistant
         st.markdown(msg["content"])
 
-        if msg["role"] == "assistant":
+        # Info icon popover with full prompt/context
+        # (st.popover is stable in your Streamlit range; no rerun on open/close)
+        c1, c2 = st.columns([1, 12], vertical_alignment="center")
+        with c1:
+            with st.popover("ℹ️", help="Inspect"):
+                st.markdown("**Context**")
+                st.code(msg.get("context", ""), language="text")
+                st.markdown("**System**")
+                st.code(msg.get("system_prompt", ""), language="text")
+                st.markdown("**Question**")
+                st.code(msg.get("question", ""), language="text")
+                st.markdown("**Prompt**")
+                st.code(msg.get("prompt", ""), language="text")
+        with c2:
             st.caption(
-                f"{msg['search_time']:.2f}s • {msg['gen_time']:.2f}s • "
-                f"{msg['tps']:.1f} tok/s • in {msg['input_tokens']} • out {msg['output_tokens']}"
+                f"🔎 {msg['search_time']:.2f}s  "
+                f"🧠 {msg['gen_time']:.2f}s  "
+                f"⚡ {msg['tps']:.1f} tok/s  "
+                f"🧾 in {msg['input_tokens']} • out {msg['output_tokens']}"
             )
 
-            with st.expander("Inspect context"):
-                st.code(msg.get("prompt", ""), language="text")
-
-            key = f"fb_{i}"
-            st.session_state.setdefault(key, msg.get("feedback"))
-            st.feedback(
-                "thumbs",
-                key=key,
-                disabled=msg.get("feedback") is not None,
-                on_change=handle_feedback,
-                args=(i,),
-            )
+        key = f"fb_{i}"
+        st.session_state.setdefault(key, msg.get("feedback"))
+        st.feedback(
+            "thumbs",
+            key=key,
+            disabled=msg.get("feedback") is not None,
+            on_change=handle_feedback,
+            args=(i,),
+        )
 
 
 # =========================
@@ -234,79 +279,89 @@ query = st.chat_input("Ask a question...")
 
 if query:
     st.session_state.messages.append({"role": "user", "content": query})
+    st.session_state.needs_answer = True
     st.rerun()
 
 
 # =========================
-# GENERATE
+# GENERATE (once per user message)
 # =========================
-if st.session_state.messages and st.session_state.messages[-1]["role"] == "user":
+if (
+    st.session_state.needs_answer
+    and st.session_state.messages
+    and st.session_state.messages[-1]["role"] == "user"
+):
     question = st.session_state.messages[-1]["content"]
 
+    # Web search
     web_ctx, search_time = web_search_snippets(question)
 
-    final_context = truncate_context(
-        web_ctx,
-        local_context,
-        system_prompt,
-        question,
-    )
-
-    prompt = f"{final_context}\n{system_prompt}\n{question}\n"
+    # Context + truncation
+    context = truncate_context(web_ctx, local_context, system_prompt, question)
+    prompt = f"{context}\n{system_prompt}\n{question}\n"
     input_tokens = count_tokens(prompt)
 
-    run_id = None
-    if ls_client:
-        try:
-            run = ls_client.create_run(
-                name="teapot_chat",
-                run_type="llm",
-                inputs={"prompt": prompt, "question": question},
-            )
-            run_id = run.id
-        except Exception:
-            pass
-
+    # Run traced answer (returns answer; trace_id obtained from current run tree)
     with st.chat_message("assistant"):
         placeholder = st.empty()
+
         start = time.perf_counter()
-        final_text = ""
 
-        for partial in stream_generate(prompt):
-            final_text = partial
-            placeholder.markdown(final_text)
+        # Generate full answer first (traced), then "stream" it to UI quickly.
+        # This keeps LangSmith tracing simple/reliable while still giving a streaming UX.
+        answer = traced_answer(context, system_prompt, question)
+        trace_id = get_trace_id_if_available()
+
+        # Typewriter-ish stream (fast, looks normal)
+        buf = ""
+        for ch in answer:
+            buf += ch
+            placeholder.markdown(buf)
+            # small delay; tune if you want faster/slower
+            time.sleep(0.002)
 
         gen_time = time.perf_counter() - start
-        output_tokens = count_tokens(final_text)
+        output_tokens = count_tokens(answer)
         tps = output_tokens / gen_time if gen_time > 0 else 0.0
 
-        st.caption(
-            f"{search_time:.2f}s • {gen_time:.2f}s • "
-            f"{tps:.1f} tok/s • in {input_tokens} • out {output_tokens}"
-        )
-
-        with st.expander("Inspect context"):
-            st.code(prompt, language="text")
-
-    if ls_client and run_id:
-        try:
-            ls_client.update_run(run_id, outputs={"answer": final_text})
-        except Exception:
-            pass
+        # Metrics + info popover for this live message
+        c1, c2 = st.columns([1, 12], vertical_alignment="center")
+        with c1:
+            with st.popover("ℹ️", help="Inspect"):
+                st.markdown("**Context**")
+                st.code(context, language="text")
+                st.markdown("**System**")
+                st.code(system_prompt, language="text")
+                st.markdown("**Question**")
+                st.code(question, language="text")
+                st.markdown("**Prompt**")
+                st.code(prompt, language="text")
+        with c2:
+            st.caption(
+                f"🔎 {search_time:.2f}s  "
+                f"🧠 {gen_time:.2f}s  "
+                f"⚡ {tps:.1f} tok/s  "
+                f"🧾 in {input_tokens} • out {output_tokens}"
+            )
 
+    # Persist assistant message
     st.session_state.messages.append(
         {
             "role": "assistant",
-            "content": final_text,
+            "content": answer,
+            "context": context,
+            "system_prompt": system_prompt,
+            "question": question,
             "prompt": prompt,
             "search_time": search_time,
             "gen_time": gen_time,
             "input_tokens": input_tokens,
             "output_tokens": output_tokens,
             "tps": tps,
-            "run_id": run_id,
+            "trace_id": trace_id,
             "feedback": None,
         }
     )
 
+    st.session_state.needs_answer = False
     st.rerun()