Update src/streamlit_app.py

src/streamlit_app.py

@@ -5,17 +5,13 @@ import requests
 
 import streamlit as st
 import torch
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, TextIteratorStreamer
 
-# Optional LangSmith (trace + feedback)
+# Optional LangSmith
 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
 
 
 # =========================
@@ -27,6 +23,9 @@ MAX_NEW_TOKENS = 192
 TOP_K_SEARCH = 3
 LOGO_URL = "https://teapotai.com/assets/logo.gif"
 
+# How many (user,assistant) pairs to include in the prompt by default
+MAX_TURNS_IN_PROMPT = 6
+
 st.set_page_config(page_title="TeapotAI Chat", page_icon="🫖", layout="centered")
 
 
@@ -50,8 +49,8 @@ tokenizer, model, device = load_model()
 # =========================
 @st.cache_resource
 def get_langsmith():
-    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:
+    key = os.getenv("LANGCHAIN_API_KEY") or os.getenv("LANGSMITH_API_KEY")
+    if key and LangSmithClient:
         return LangSmithClient()
     return None
 
@@ -70,11 +69,10 @@ if "needs_answer" not in st.session_state:
 
 # =========================
 # HEADER (prevent logo flash)
-# Use a fixed pixel width to avoid layout shift / big flash.
 # =========================
 col1, col2 = st.columns([1, 7], vertical_alignment="center")
 with col1:
-    st.image(LOGO_URL, width=56)  # fixed width prevents "flash huge"
+    st.image(LOGO_URL, width=56)
 with col2:
     st.markdown("## TeapotAI Chat")
     st.caption("Grounded answers with web context")
@@ -104,6 +102,14 @@ with st.sidebar:
         placeholder="Extra context appended after web snippets…",
     )
 
+    max_turns = st.slider(
+        "Conversation turns in prompt",
+        min_value=0,
+        max_value=12,
+        value=MAX_TURNS_IN_PROMPT,
+        help="How many recent (user, assistant) pairs to include in the prompt.",
+    )
+
 
 # =========================
 # WEB SEARCH (ALWAYS ON)
@@ -139,93 +145,101 @@ def web_search_snippets(query: str):
 
 
 # =========================
-# CONTEXT TRUNCATION (TAIL)
+# CONTEXT + PROMPT BUILDING
 # =========================
-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:
+def count_tokens(text: str) -> int:
+    return len(tokenizer.encode(text)) if text else 0
+
+
+def build_conversation(messages, turns: int) -> str:
+    """
+    Build a compact transcript from the last `turns` (user,assistant) pairs.
+    """
+    if turns <= 0:
         return ""
-    ctx_tokens = tokenizer.encode(ctx)
-    if len(ctx_tokens) <= budget:
-        return ctx
-    return tokenizer.decode(ctx_tokens[-budget:], skip_special_tokens=True)
 
+    # Collect last 2*turns messages ending at the most recent message
+    # Keep only user/assistant roles.
+    filtered = [m for m in messages if m.get("role") in ("user", "assistant")]
+
+    # Take tail, but ensure we start on a user message if possible
+    tail = filtered[-(2 * turns) :]
+    # If first is assistant, drop it (misaligned pair)
+    if tail and tail[0]["role"] == "assistant":
+        tail = tail[1:]
+
+    lines = []
+    for m in tail:
+        role = "User" if m["role"] == "user" else "Assistant"
+        content = (m.get("content") or "").strip()
+        if content:
+            lines.append(f"{role}: {content}")
+    return "\n".join(lines).strip()
+
+
+def truncate_to_token_budget(full_prompt: str, max_tokens: int) -> str:
+    ids = tokenizer.encode(full_prompt)
+    if len(ids) <= max_tokens:
+        return full_prompt
+    # Tail truncate to keep the most recent instruction + question
+    ids = ids[-max_tokens:]
+    return tokenizer.decode(ids, skip_special_tokens=True)
+
+
+def build_prompt(web_ctx: str, local_ctx: str, system: str, convo: str, question: str) -> str:
+    # Order matters: context first, then system, then convo, then question.
+    parts = []
+    ctx = f"{web_ctx}\n\n{local_ctx}".strip()
+    if ctx:
+        parts.append("Context:\n" + ctx)
 
-def count_tokens(text: str) -> int:
-    return len(tokenizer.encode(text)) if text else 0
+    parts.append("System:\n" + system.strip())
+
+    if convo:
+        parts.append("Conversation:\n" + convo)
+
+    parts.append("User:\n" + question.strip())
+    parts.append("Assistant:\n")  # encourages continuation style
+
+    raw = "\n\n".join(parts).strip() + "\n"
+    # Enforce input budget at token level
+    return truncate_to_token_budget(raw, MAX_INPUT_TOKENS)
 
 
 # =========================
-# LANGSMITH-TRACED ANSWER FUNCTION
-# (signature exactly: context, system_prompt, question -> answer)
+# STREAM GENERATION
 # =========================
-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)
+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,
+        )
 
+    threading.Thread(target=run, daemon=True).start()
 
-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
+    acc = ""
+    for chunk in streamer:
+        acc += chunk
+        yield acc
 
 
 # =========================
-# FEEDBACK HANDLER (attached to trace_id)
+# FEEDBACK
 # =========================
 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")
-
-    # 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
+    # If you later add LangSmith run_ids per message, hook it here.
+    # (Keeping this simple/stable like your previous version.)
+    # if ls_client and st.session_state.messages[idx].get("run_id"): ...
 
 
 # =========================
@@ -233,43 +247,35 @@ 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"])
 
-        # 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,),
-        )
+        if msg["role"] == "assistant":
+            # Inline row: info popover + thumbs + metrics
+            c_info, c_fb, c_metrics = st.columns([1.2, 1.4, 10], vertical_alignment="center")
+
+            with c_info:
+                with st.popover("ℹ️", help="Inspect prompt"):
+                    st.markdown("**Prompt sent to model**")
+                    st.code(msg.get("prompt", ""), language="text")
+
+            with c_fb:
+                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,),
+                )
+
+            with c_metrics:
+                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']}"
+                )
 
 
 # =========================
@@ -296,69 +302,66 @@ if (
     # Web search
     web_ctx, search_time = web_search_snippets(question)
 
-    # Context + truncation
-    context = truncate_context(web_ctx, local_context, system_prompt, question)
-    prompt = f"{context}\n{system_prompt}\n{question}\n"
+    # Conversation transcript (from prior messages, excluding current user msg is fine either way;
+    # keeping it includes the last user msg too, but we also add question explicitly.)
+    convo = build_conversation(st.session_state.messages[:-1], turns=max_turns)
+
+    # Prompt
+    prompt = build_prompt(
+        web_ctx=web_ctx,
+        local_ctx=local_context,
+        system=system_prompt,
+        convo=convo,
+        question=question,
+    )
+
     input_tokens = count_tokens(prompt)
 
-    # Run traced answer (returns answer; trace_id obtained from current run tree)
+    # Stream normally
     with st.chat_message("assistant"):
         placeholder = st.empty()
-
         start = time.perf_counter()
+        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)
+        for partial in stream_generate(prompt):
+            final_text = partial
+            placeholder.markdown(final_text)
 
         gen_time = time.perf_counter() - start
-        output_tokens = count_tokens(answer)
+        output_tokens = count_tokens(final_text)
         tps = output_tokens / gen_time if gen_time > 0 else 0.0
 
-        # 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**")
+        # Inline row under the live message
+        c_info, c_fb, c_metrics = st.columns([1.2, 1.4, 10], vertical_alignment="center")
+
+        with c_info:
+            with st.popover("ℹ️", help="Inspect prompt"):
+                st.markdown("**Prompt sent to model**")
                 st.code(prompt, language="text")
-        with c2:
+
+        # For the live message, we don't have a saved index yet; show disabled thumbs placeholder
+        with c_fb:
+            st.feedback("thumbs", key="fb_live", disabled=True)
+
+        with c_metrics:
             st.caption(
-                f"🔎 {search_time:.2f}s  "
-                f"🧠 {gen_time:.2f}s  "
-                f"⚡ {tps:.1f} tok/s  "
-                f"🧾 in {input_tokens} • out {output_tokens}"
+                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
+    # Persist assistant message (so feedback attaches properly after rerun)
     st.session_state.messages.append(
         {
             "role": "assistant",
-            "content": answer,
-            "context": context,
-            "system_prompt": system_prompt,
-            "question": question,
+            "content": final_text,
             "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,
         }
     )