Update src/streamlit_app.py

src/streamlit_app.py

@@ -8,7 +8,7 @@ import streamlit as st
 import torch
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, TextIteratorStreamer
 
-# Optional parsing libs (best-effort)
+# Optional parsing libs (safe fallbacks)
 try:
     from pypdf import PdfReader  # pip install pypdf
 except Exception:
@@ -25,6 +25,7 @@ try:
 except Exception:
     LangSmithClient = None
 
+
 # =========================
 # CONFIG
 # =========================
@@ -76,11 +77,12 @@ if "pending_response" not in st.session_state:
 
 
 # =========================
-# HEADER (LOGO)
+# HEADER (SAFE IMAGE CALL)
 # =========================
 col1, col2 = st.columns([1, 6])
 with col1:
-    st.image(LOGO_URL, use_container_width=True)
+    # IMPORTANT: use_column_width=True (works on your Streamlit version)
+    st.image(LOGO_URL, use_column_width=True)
 with col2:
     st.markdown("## TeapotAI Chat")
     st.caption("Fast grounded answers with clean web context")
@@ -104,25 +106,25 @@ with st.sidebar:
         height=180,
     )
 
-    st.markdown("### Local Context")
+    st.markdown("### Local Context (Text)")
     local_context_text = st.text_area(
         "Paste additional context (optional)",
         height=140,
         placeholder="This will be appended after web content...",
     )
 
+    st.markdown("### Local Context (File Upload)")
     uploaded_files = st.file_uploader(
-        "Upload files to add to Local Context (pdf, txt, csv, md, json, etc.)",
-        type=None,
+        "Upload files (pdf, txt, csv, md, json, etc.)",
         accept_multiple_files=True,
+        type=None,
     )
 
 
 # =========================
-# FILE PARSING -> STRING
+# FILE PARSING
 # =========================
-def _safe_decode(b: bytes) -> str:
-    # best effort decode without throwing
+def safe_decode(b: bytes) -> str:
     for enc in ("utf-8", "utf-16", "latin-1"):
         try:
             return b.decode(enc)
@@ -131,58 +133,46 @@ def _safe_decode(b: bytes) -> str:
     return b.decode("utf-8", errors="ignore")
 
 
-def parse_uploaded_file_to_text(file) -> str:
+def parse_file_to_text(file) -> str:
     name = (file.name or "").lower()
     raw = file.getvalue()
 
     # PDF
-    if name.endswith(".pdf"):
-        if not PdfReader:
-            return (
-                f"[{file.name}] PDF parsing not available (install pypdf). "
-                f"Raw bytes={len(raw)}"
-            )
+    if name.endswith(".pdf") and PdfReader:
         try:
             reader = PdfReader(io.BytesIO(raw))
-            parts = []
-            for i, page in enumerate(reader.pages):
+            pages = []
+            for page in reader.pages:
                 txt = page.extract_text() or ""
-                txt = txt.strip()
-                if txt:
-                    parts.append(txt)
-            return "\n\n".join(parts).strip()
+                if txt.strip():
+                    pages.append(txt.strip())
+            return "\n\n".join(pages)
         except Exception as e:
-            return f"[{file.name}] PDF parse error: {e}"
+            return f"[PDF parse error: {e}]"
 
     # CSV
-    if name.endswith(".csv"):
-        if not pd:
-            return (
-                f"[{file.name}] CSV parsing not available (install pandas). "
-                f"Raw bytes={len(raw)}"
-            )
+    if name.endswith(".csv") and pd:
         try:
             df = pd.read_csv(io.BytesIO(raw))
-            # Keep it compact but readable
             return df.to_csv(index=False)
-        except Exception as e:
-            # fallback: raw text
-            return f"[{file.name}] CSV parse error ({e}). Raw:\n{_safe_decode(raw)}"
+        except Exception:
+            return safe_decode(raw)
 
-    # JSON / TXT / MD / others -> decode
-    return _safe_decode(raw).strip()
+    # TXT / MD / JSON / fallback
+    return safe_decode(raw)
 
 
-def build_local_context(text_area: str, files) -> str:
+def build_local_context(text_block: str, files) -> str:
     chunks = []
-    if text_area.strip():
-        chunks.append(text_area.strip())
+
+    if text_block and text_block.strip():
+        chunks.append(text_block.strip())
 
     if files:
         for f in files:
-            parsed = parse_uploaded_file_to_text(f).strip()
-            if parsed:
-                chunks.append(f"\n\n--- FILE: {f.name} ---\n{parsed}")
+            parsed = parse_file_to_text(f)
+            if parsed and parsed.strip():
+                chunks.append(f"\n\n--- FILE: {f.name} ---\n{parsed.strip()}")
 
     return "\n\n".join(chunks).strip()
 
@@ -191,7 +181,7 @@ local_context = build_local_context(local_context_text, uploaded_files)
 
 
 # =========================
-# WEB SEARCH (SNIPPETS ONLY) - ALWAYS ON
+# WEB SEARCH (ALWAYS ON)
 # =========================
 def web_search_snippets(query: str):
     api_key = os.getenv("BRAVE_API_KEY") or st.secrets.get("BRAVE_API_KEY", None)
@@ -221,14 +211,13 @@ def web_search_snippets(query: str):
         if desc:
             snippets.append(desc)
 
-    clean_context = "\n\n".join(snippets)  # paragraph-separated only
-    return clean_context, (t1 - t0)
+    return "\n\n".join(snippets), (t1 - t0)
 
 
 # =========================
-# TRUNCATE TO LAST 512 TOKENS (TAIL)
+# CONTEXT TRUNCATION (TAIL)
 # =========================
-def truncate_context(web_ctx: str, local_ctx: str, system: str, question: str) -> str:
+def truncate_context(web_ctx, local_ctx, system, question):
     ordered_context = f"{web_ctx}\n\n{local_ctx}".strip()
 
     base = f"\n{system}\n{question}\n"
@@ -236,13 +225,13 @@ def truncate_context(web_ctx: str, local_ctx: str, system: str, question: str) -
     budget = MAX_INPUT_TOKENS - len(base_tokens)
 
     if budget <= 0:
-        return ""  # system+question already consume budget
+        return ""
 
     ctx_tokens = tokenizer.encode(ordered_context) if ordered_context else []
     if len(ctx_tokens) <= budget:
         return ordered_context
 
-    truncated = ctx_tokens[-budget:]  # keep MOST RECENT tokens
+    truncated = ctx_tokens[-budget:]
     return tokenizer.decode(truncated, skip_special_tokens=True)
 
 
@@ -262,7 +251,7 @@ def stream_generate(prompt: str):
             streamer=streamer,
         )
 
-    thread = threading.Thread(target=run, daemon=True)
+    thread = threading.Thread(target=run)
     thread.start()
 
     text = ""
@@ -272,7 +261,7 @@ def stream_generate(prompt: str):
 
 
 # =========================
-# FEEDBACK HANDLER (Native st.feedback)
+# FEEDBACK HANDLER
 # =========================
 def handle_feedback(idx: int):
     val = st.session_state[f"feedback_{idx}"]
@@ -301,7 +290,7 @@ for i, msg in enumerate(st.session_state.messages):
             st.markdown(msg["content"])
             continue
 
-        # Assistant messages: collapsed-by-default expander = "whole message response be the dropdown"
+        # Entire response as collapsed dropdown (less visible inspector)
         with st.expander("🫖 Assistant response (click to expand)", expanded=False):
             st.markdown(msg["content"])
 
@@ -312,17 +301,9 @@ for i, msg in enumerate(st.session_state.messages):
                 f"🧾 in={msg['input_tokens']} • out={msg['output_tokens']}"
             )
 
-            # Show EXACT prompt passed into the model (and the parts)
-            st.markdown("---")
-            st.markdown("#### Prompt & Inputs (exactly what was passed to the model)")
-            st.markdown("**System prompt:**")
-            st.code(msg.get("system_prompt", ""), language="text")
-            st.markdown("**Question:**")
-            st.code(msg.get("question", ""), language="text")
-            st.markdown("**Full model input (prompt):**")
-            st.code(msg.get("prompt", ""), language="text")
-
-        # Native thumbs feedback (outside expander so it's still reachable)
+            st.markdown("### Exact Model Input (Prompt)")
+            st.code(msg["prompt"], language="text")
+
         key = f"feedback_{i}"
         st.session_state.setdefault(key, msg.get("feedback"))
         st.feedback(
@@ -340,7 +321,6 @@ for i, msg in enumerate(st.session_state.messages):
 query = st.chat_input("Ask a question...")
 
 if query:
-    # show user message first
     st.session_state.messages.append({"role": "user", "content": query})
     st.rerun()
 
@@ -355,24 +335,22 @@ if (
 ):
     question = st.session_state.messages[-1]["content"]
 
-    # --- Web Search (always on) ---
+    # Always do web search
     web_ctx, search_time = web_search_snippets(question)
 
-    # --- Strict Order Context ---
     final_context = truncate_context(
-        web_ctx=web_ctx,
-        local_ctx=local_context,
-        system=system_prompt,
-        question=question,
+        web_ctx,
+        local_context,
+        system_prompt,
+        question,
     )
 
-    # IMPORTANT: prompt is EXACTLY what we pass to the model
-    prompt = f"{final_context}\n{system_prompt}\n{question}\n".strip() + "\n"
+    # EXACT prompt passed to model
+    prompt = f"{final_context}\n{system_prompt}\n{question}\n"
 
-    # Token accounting (split input vs output)
     input_tokens = len(tokenizer.encode(prompt))
 
-    # LangSmith run
+    # LangSmith run (optional)
     run_id = None
     if ls_client:
         try:
@@ -380,19 +358,14 @@ if (
                 name="teapot_chat",
                 run_type="llm",
                 inputs={
-                    "web_content": web_ctx,
-                    "local_context": local_context,
-                    "system_prompt": system_prompt,
-                    "question": question,
-                    "final_context": final_context,
                     "prompt": prompt,
+                    "question": question,
                 },
             )
             run_id = run.id
         except Exception:
             pass
 
-    # --- Stream UI: assistant response itself is a dropdown ---
     with st.chat_message("assistant"):
         with st.expander("🫖 Assistant response (click to expand)", expanded=False):
             placeholder = st.empty()
@@ -414,13 +387,7 @@ if (
                 f"🧾 in={input_tokens} • out={output_tokens}"
             )
 
-            st.markdown("---")
-            st.markdown("#### Prompt & Inputs (exactly what was passed to the model)")
-            st.markdown("**System prompt:**")
-            st.code(system_prompt, language="text")
-            st.markdown("**Question:**")
-            st.code(question, language="text")
-            st.markdown("**Full model input (prompt):**")
+            st.markdown("### Exact Model Input (Prompt)")
             st.code(prompt, language="text")
 
     if ls_client and run_id:
@@ -433,11 +400,6 @@ if (
         {
             "role": "assistant",
             "content": final_text,
-            "system_prompt": system_prompt,
-            "question": question,
-            "web_context": web_ctx,
-            "local_context": local_context,
-            "final_context": final_context,
             "prompt": prompt,
             "search_time": search_time,
             "gen_time": gen_time,