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
 
@@ -65,17 +69,15 @@ if "needs_answer" not in st.session_state:
 
 
 # =========================
-# HEADER (reduce logo flash)
-# - fixed width prevents "giant image then shrink"
-# - container keeps layout stable
+# HEADER (prevent logo flash)
+# Use a fixed pixel width to avoid layout shift / big flash.
 # =========================
-with st.container():
-    col1, col2 = st.columns([1, 7], vertical_alignment="center")
-    with col1:
-        st.image(LOGO_URL, width=56)  # fixed width = stable, no huge flash
-    with col2:
-        st.markdown("## TeapotAI Chat")
-        st.caption("Grounded answers with web context")
+col1, col2 = st.columns([1, 7], vertical_alignment="center")
+with col1:
+    st.image(LOGO_URL, width=56)  # fixed width prevents "flash huge"
+with col2:
+    st.markdown("## TeapotAI Chat")
+    st.caption("Grounded answers with web context")
 
 
 # =========================
@@ -137,94 +139,93 @@ def web_search_snippets(query: str):
 
 
 # =========================
-# PROMPT BUILDING (NO CONVERSATION HISTORY)
+# CONTEXT TRUNCATION (TAIL)
 # =========================
-def count_tokens(text: str) -> int:
-    return len(tokenizer.encode(text)) if text else 0
-
-
-def truncate_to_token_budget(text: str, max_tokens: int) -> str:
-    ids = tokenizer.encode(text)
-    if len(ids) <= max_tokens:
-        return text
-    ids = ids[-max_tokens:]  # tail truncate
-    return tokenizer.decode(ids, skip_special_tokens=True)
+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 ""
+    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)
 
 
-def build_prompt(web_ctx: str, local_ctx: str, system: str, question: str) -> str:
-    ctx = f"{web_ctx}\n\n{local_ctx}".strip()
-    parts = []
-    if ctx:
-        parts.append(ctx)
-    parts.append(system.strip())
-    parts.append(question.strip())
-    raw = "\n\n".join(parts).strip() + "\n"
-    return truncate_to_token_budget(raw, MAX_INPUT_TOKENS)
+def count_tokens(text: str) -> int:
+    return len(tokenizer.encode(text)) if text else 0
 
 
 # =========================
-# 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
+# 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 you want to wire LangSmith feedback to a run later, store run_id per message and use it here.
-    # For now we keep it stable and local like the earlier version.
-
-
-def render_inline_controls(msg: dict, feedback_key: str, feedback_disabled: bool, feedback_idx: int | None):
-    """
-    Inline row under assistant message:
-      ℹ️ popover (full prompt), thumbs, metrics
-    """
-    c_info, c_fb, c_metrics = st.columns([1.1, 1.7, 10], vertical_alignment="center")
-
-    with c_info:
-        with st.popover("ℹ️", help="Inspect"):
-            st.markdown("**Prompt (sent to model)**")
-            st.code(msg.get("prompt", ""), language="text")
-
-    with c_fb:
-        st.feedback(
-            "thumbs",
-            key=feedback_key,
-            disabled=feedback_disabled,
-            on_change=(handle_feedback if (not feedback_disabled and feedback_idx is not None) else None),
-            args=((feedback_idx,) if (not feedback_disabled and feedback_idx is not None) else None),
-        )
-
-    with c_metrics:
-        st.caption(
-            f"🔎 {msg.get('search_time', 0.0):.2f}s "
-            f"• 🧠 {msg.get('gen_time', 0.0):.2f}s "
-            f"• ⚡ {msg.get('tps', 0.0):.1f} tok/s "
-            f"• 🧾 in {msg.get('input_tokens', 0)} • out {msg.get('output_tokens', 0)}"
-        )
+    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:
+            # LangSmith SDK supports trace_id= for feedback association
+            ls_client.create_feedback(
+                trace_id=trace_id,
+                key="thumb_rating",
+                score=score,
+                comment="thumbs_up" if score else "thumbs_down",
+            )
+        except Exception:
+            pass
 
 
 # =========================
@@ -232,20 +233,44 @@ def render_inline_controls(msg: dict, feedback_key: str, feedback_disabled: bool
 # =========================
 for i, msg in enumerate(st.session_state.messages):
     with st.chat_message(msg["role"]):
-        st.markdown(msg["content"])
+        if msg["role"] == "user":
+            st.markdown(msg["content"])
+            continue
 
-        if msg["role"] == "assistant":
-            # Ensure feedback state exists
-            k = f"fb_{i}"
-            st.session_state.setdefault(k, msg.get("feedback"))
+        # Assistant
+        st.markdown(msg["content"])
 
-            render_inline_controls(
-                msg=msg,
-                feedback_key=k,
-                feedback_disabled=(msg.get("feedback") is not None),
-                feedback_idx=i,
+        # 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  "
+                f"🧠 {msg['gen_time']:.2f}s  "
+                f"⚡ {msg['tps']:.1f} tok/s  "
+                f"🧾 in {msg['input_tokens']} • out {msg['output_tokens']}"
             )
 
+        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,),
+        )
+
 
 # =========================
 # INPUT
@@ -268,57 +293,72 @@ if (
 ):
     question = st.session_state.messages[-1]["content"]
 
+    # Web search
     web_ctx, search_time = web_search_snippets(question)
 
-    prompt = build_prompt(
-        web_ctx=web_ctx,
-        local_ctx=local_context,
-        system=system_prompt,
-        question=question,
-    )
+    # 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 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
 
-        live_msg = {
-            "prompt": prompt,
-            "search_time": search_time,
-            "gen_time": gen_time,
-            "input_tokens": input_tokens,
-            "output_tokens": output_tokens,
-            "tps": tps,
-        }
-
-        # Inline controls for the live message (thumbs disabled until it’s saved)
-        render_inline_controls(
-            msg=live_msg,
-            feedback_key="fb_live",
-            feedback_disabled=True,
-            feedback_idx=None,
-        )
+        # 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 (search) "
+                f"🧠 {gen_time:.2f}s (generation) "
+                f"⚡ {tps:.1f} tok/s  "
+                f"🧾 {input_tokens} input tokens • {output_tokens} output tokens"
+            )
 
-    # Save assistant message so thumbs attach after rerun
+    # 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,
+            "trace_id": trace_id,
             "feedback": None,
         }
     )