Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,180 +1,167 @@
 import os
-import re
 import time
 import threading
-from typing import List, Dict, Optional
-
+import requests
 import streamlit as st
 import torch
-import requests
-
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, TextIteratorStreamer
 
-# Fast RAG
-from sklearn.feature_extraction.text import TfidfVectorizer
-import numpy as np
-
-# File parsing
-import io
-try:
-    import docx
-except:
-    docx = None
-try:
-    import PyPDF2
-except:
-    PyPDF2 = None
-
-# LangSmith
+# Optional LangSmith
 try:
     from langsmith import Client as LangSmithClient
 except:
     LangSmithClient = None
 
 
-# -----------------------
+# =========================
 # CONFIG
-# -----------------------
+# =========================
 MODEL_NAME = "teapotai/tinyteapot"
 MAX_INPUT_TOKENS = 512
 MAX_NEW_TOKENS = 192
 TOP_K_SEARCH = 3
-TOP_K_RAG = 3
-
-st.set_page_config(page_title="TeapotAI Chat", page_icon="🫖", layout="centered")
+LOGO_URL = "https://teapotai.com/assets/logo.gif"
 
+st.set_page_config(
+    page_title="TeapotAI Chat",
+    page_icon="🫖",
+    layout="centered"
+)
 
-# -----------------------
-# MODEL LOADING (CACHED)
-# -----------------------
+# =========================
+# LOAD MODEL (CACHED)
+# =========================
 @st.cache_resource
 def load_model():
     tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
     model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME)
     device = "cuda" if torch.cuda.is_available() else "cpu"
-    model.to(device).eval()
+    model.to(device)
+    model.eval()
     return tokenizer, model, device
 
-
 tokenizer, model, device = load_model()
 
 
-# -----------------------
+# =========================
 # LANGSMITH
-# -----------------------
+# =========================
 @st.cache_resource
 def get_langsmith():
-    if LangSmithClient and (os.getenv("LANGCHAIN_API_KEY") or os.getenv("LANGSMITH_API_KEY")):
+    api_key = os.getenv("LANGCHAIN_API_KEY") or os.getenv("LANGSMITH_API_KEY")
+    if api_key and LangSmithClient:
         return LangSmithClient()
     return None
 
 ls_client = get_langsmith()
 
 
-# -----------------------
-# FAST FILE → TEXT
-# -----------------------
-def extract_text(file) -> str:
-    name = file.name.lower()
+# =========================
+# SESSION STATE
+# =========================
+if "messages" not in st.session_state:
+    st.session_state.messages = []
+
+# =========================
+# HEADER (LOGO RESTORED)
+# =========================
+col1, col2 = st.columns([1, 6])
+with col1:
+    st.image(LOGO_URL, use_column_width=True)
+with col2:
+    st.markdown("## TeapotAI Chat")
+    st.caption("Fast, grounded answers with web context")
+
 
-    if name.endswith(".txt") or name.endswith(".md"):
-        return file.read().decode("utf-8", errors="ignore")
+# =========================
+# SIDEBAR SETTINGS
+# =========================
+with st.sidebar:
+    st.markdown("### Settings")
 
-    if name.endswith(".pdf") and PyPDF2:
-        reader = PyPDF2.PdfReader(file)
-        return "\n".join(page.extract_text() or "" for page in reader.pages)
+    system_prompt = st.text_area(
+        "System Prompt",
+        value=(
+            "You are Teapot, an open-source AI assistant optimized for low-end devices, "
+            "providing short, accurate responses without hallucinating while excelling at "
+            "information extraction and text summarization. "
+            "If the context does not answer the question, reply exactly: "
+            "'I am sorry but I don't have any information on that'."
+        ),
+        height=180
+    )
 
-    if name.endswith(".docx") and docx:
-        doc = docx.Document(file)
-        return "\n".join(p.text for p in doc.paragraphs)
+    st.markdown("### Extra Context (Optional)")
+    user_context = st.text_area(
+        "Paste context to append to web results",
+        height=150,
+        placeholder="Add any custom context here..."
+    )
 
-    return ""
+    use_web = st.checkbox("Use web search", value=True)
 
 
-def chunk_by_paragraph(text: str) -> List[str]:
-    chunks = [p.strip() for p in re.split(r"\n{2,}|\n", text) if len(p.strip()) > 30]
-    return chunks[:2000]  # safety cap
+# =========================
+# WEB SEARCH (FAST)
+# =========================
+def web_search(query: str):
+    api_key = os.getenv("BRAVE_API_KEY") or st.secrets.get("BRAVE_API_KEY", None)
+    if not api_key:
+        return "", 0.0
 
+    headers = {
+        "X-Subscription-Token": api_key,
+        "Accept": "application/json"
+    }
 
-# -----------------------
-# FAST TFIDF RAG (CACHED)
-# -----------------------
-@st.cache_resource
-def build_tfidf(chunks: List[str]):
-    if not chunks:
-        return None, None
-    vectorizer = TfidfVectorizer(stop_words="english", max_features=20000)
-    matrix = vectorizer.fit_transform(chunks)
-    return vectorizer, matrix
-
-
-def retrieve_top_chunks(query: str, chunks: List[str], vectorizer, matrix, k=3):
-    if not chunks or vectorizer is None:
-        return []
-    q_vec = vectorizer.transform([query])
-    scores = (matrix @ q_vec.T).toarray().ravel()
-    top_idx = np.argsort(scores)[-k:][::-1]
-    return [chunks[i] for i in top_idx if scores[i] > 0]
-
-
-# -----------------------
-# SEARCH (FAST)
-# -----------------------
-def web_search(query):
-    key = os.getenv("BRAVE_API_KEY") or st.secrets.get("BRAVE_API_KEY", None)
-    if not key:
-        return [], ""
-
-    headers = {"X-Subscription-Token": key, "Accept": "application/json"}
     params = {"q": query, "count": TOP_K_SEARCH}
 
     t0 = time.perf_counter()
-    r = requests.get(
-        "https://api.search.brave.com/res/v1/web/search",
-        headers=headers,
-        params=params,
-        timeout=8,
-    )
-    data = r.json()
+    try:
+        r = requests.get(
+            "https://api.search.brave.com/res/v1/web/search",
+            headers=headers,
+            params=params,
+            timeout=6,
+        )
+        data = r.json()
+    except:
+        return "", 0.0
     t1 = time.perf_counter()
 
-    results = []
-    ctx_blocks = []
-
+    blocks = []
     for i, item in enumerate(data.get("web", {}).get("results", [])[:TOP_K_SEARCH], 1):
         title = item.get("title", "")
         url = item.get("url", "")
-        desc = item.get("description", "").replace("<strong>", "").replace("</strong>", "")
-
-        results.append({"title": title, "url": url, "snippet": desc})
-        ctx_blocks.append(f"[{i}] {title}\nURL: {url}\nSnippet: {desc}")
-
-    return results, "\n\n".join(ctx_blocks), (t1 - t0)
+        desc = item.get("description", "")
+        desc = desc.replace("<strong>", "").replace("</strong>", "")
+        blocks.append(f"[{i}] {title}\nURL: {url}\nSnippet: {desc}")
 
+    context = "\n\n".join(blocks)
+    return context, (t1 - t0)
 
-# -----------------------
-# PROMPT + TRUNCATION
-# -----------------------
-def build_prompt(context, system, question):
-    return f"{context}\n{system}\n{question}\n"
 
-
-def truncate_context(context, system, question):
-    base = build_prompt("", system, question)
-    base_tokens = tokenizer.encode(base)
+# =========================
+# TRUNCATE TO LAST 512 TOKENS
+# =========================
+def truncate_to_512(context: str, system: str, question: str):
+    base_prompt = f"\n{system}\n{question}\n"
+    base_tokens = tokenizer.encode(base_prompt)
     budget = MAX_INPUT_TOKENS - len(base_tokens)
 
     ctx_tokens = tokenizer.encode(context)
     if len(ctx_tokens) <= budget:
         return context
 
-    return tokenizer.decode(ctx_tokens[-budget:], skip_special_tokens=True)
+    # Keep MOST RECENT tokens (tail truncation)
+    truncated = ctx_tokens[-budget:]
+    return tokenizer.decode(truncated, skip_special_tokens=True)
 
 
-# -----------------------
+# =========================
 # STREAM GENERATION
-# -----------------------
-def stream_generate(prompt):
+# =========================
+def stream_generate(prompt: str):
     inputs = tokenizer(prompt, return_tensors="pt").to(device)
     streamer = TextIteratorStreamer(tokenizer, skip_special_tokens=True)
 
@@ -183,6 +170,7 @@ def stream_generate(prompt):
             **inputs,
             max_new_tokens=MAX_NEW_TOKENS,
             do_sample=False,
+            num_beams=1,
             streamer=streamer,
         )
 
@@ -195,70 +183,34 @@ def stream_generate(prompt):
         yield text
 
 
-# -----------------------
-# SESSION STATE
-# -----------------------
-if "messages" not in st.session_state:
-    st.session_state.messages = []
-
-if "rag_chunks" not in st.session_state:
-    st.session_state.rag_chunks = []
+# =========================
+# LANGSMITH FEEDBACK HANDLER
+# =========================
+def handle_feedback(idx: int):
+    val = st.session_state[f"feedback_{idx}"]
+    msg = st.session_state.messages[idx]
 
-if "vectorizer" not in st.session_state:
-    st.session_state.vectorizer = None
+    if val is None:
+        return
 
-if "matrix" not in st.session_state:
-    st.session_state.matrix = None
+    msg["feedback"] = val
 
-
-# -----------------------
-# HEADER
-# -----------------------
-st.markdown("## 🫖 TeapotAI Chat")
-st.caption("Fast, grounded answers with hybrid RAG")
-
-
-# -----------------------
-# SETTINGS SIDEBAR (UPGRADED)
-# -----------------------
-with st.sidebar:
-    st.markdown("### Settings")
-
-    system_prompt = st.text_area(
-        "System Prompt",
-        value=(
-            "You are Teapot, an open-source AI assistant optimized for low-end devices. "
-            "Answer using the provided context only. "
-            "If the context does not answer the question, reply exactly: "
-            "'I am sorry but I don't have any information on that'."
-        ),
-        height=150,
-    )
-
-    st.markdown("### Custom Context (RAG)")
-    pasted = st.text_area("Paste context text (optional)", height=150)
-    uploaded = st.file_uploader("Or upload file (.txt, .pdf, .docx, .md)", type=["txt", "pdf", "docx", "md"])
-
-
-    # Build RAG index (VERY FAST)
-    combined_text = ""
-    if pasted:
-        combined_text += pasted + "\n"
-    if uploaded:
-        combined_text += extract_text(uploaded)
-
-    if combined_text:
-        chunks = chunk_by_paragraph(combined_text)
-        vec, mat = build_tfidf(chunks)
-        st.session_state.rag_chunks = chunks
-        st.session_state.vectorizer = vec
-        st.session_state.matrix = mat
-        st.success(f"Indexed {len(chunks)} context chunks")
+    if ls_client and msg.get("run_id"):
+        score = 1 if val == "👍" else 0
+        try:
+            ls_client.create_feedback(
+                run_id=msg["run_id"],
+                key="thumb_rating",
+                score=score,
+                comment="thumbs_up" if score else "thumbs_down",
+            )
+        except Exception as e:
+            print("LangSmith feedback error:", e)
 
 
-# -----------------------
-# CHAT HISTORY
-# -----------------------
+# =========================
+# RENDER CHAT
+# =========================
 for i, msg in enumerate(st.session_state.messages):
     with st.chat_message(msg["role"]):
         st.markdown(msg["content"])
@@ -271,67 +223,70 @@ for i, msg in enumerate(st.session_state.messages):
                 f"🧮 {msg['tokens']} tokens"
             )
 
-            # Inline clean thumbs (better UX)
-            col1, col2, _ = st.columns([1, 1, 6])
-            if col1.button("👍", key=f"up_{i}", disabled=msg.get("feedback") is not None):
-                msg["feedback"] = 1
-                if ls_client and msg.get("run_id"):
-                    ls_client.create_feedback(msg["run_id"], key="feedback", score=1.0)
-            if col2.button("👎", key=f"down_{i}", disabled=msg.get("feedback") is not None):
-                msg["feedback"] = -1
-                if ls_client and msg.get("run_id"):
-                    ls_client.create_feedback(msg["run_id"], key="feedback", score=-1.0)
-
-
-# -----------------------
-# INPUT
-# -----------------------
+            feedback_key = f"feedback_{i}"
+            st.session_state.setdefault(feedback_key, msg.get("feedback"))
+
+            st.feedback(
+                "thumbs",
+                key=feedback_key,
+                disabled=msg.get("feedback") is not None,
+                on_change=handle_feedback,
+                args=(i,),
+            )
+
+
+# =========================
+# CHAT INPUT
+# =========================
 query = st.chat_input("Ask a question...")
 
 if query:
     st.session_state.messages.append({"role": "user", "content": query})
 
-    # 1️⃣ Web Search (timed)
-    results, web_context, search_time = web_search(query)
-
-    # 2️⃣ TFIDF RAG (VERY FAST ~1-3ms)
-    rag_context = ""
-    if st.session_state.vectorizer:
-        top_chunks = retrieve_top_chunks(
-            query,
-            st.session_state.rag_chunks,
-            st.session_state.vectorizer,
-            st.session_state.matrix,
-            k=TOP_K_RAG,
-        )
-        rag_context = "\n\n".join(top_chunks)
+    # ---- WEB SEARCH ----
+    web_context = ""
+    search_time = 0.0
+    if use_web:
+        web_context, search_time = web_search(query)
+
+    # ---- COMBINED CONTEXT (WEB + USER BOX) ----
+    combined_context = ""
+    if user_context:
+        combined_context += user_context.strip() + "\n\n"
+    if web_context:
+        combined_context += web_context
+
+    truncated_context = truncate_to_512(
+        combined_context,
+        system_prompt,
+        query
+    )
 
-    # 3️⃣ Hybrid Context
-    full_context = (rag_context + "\n\n" + web_context).strip()
-    truncated_context = truncate_context(full_context, system_prompt, query)
-    prompt = build_prompt(truncated_context, system_prompt, query)
+    prompt = f"{truncated_context}\n{system_prompt}\n{query}\n"
 
-    # LangSmith Run
+    # ---- LANGSMITH RUN ----
     run_id = None
     if ls_client:
-        run = ls_client.create_run(
-            name="teapot_chat",
-            run_type="llm",
-            inputs={
-                "context": truncated_context,
-                "system_prompt": system_prompt,
-                "question": query,
-            },
-        )
-        run_id = run.id
+        try:
+            run = ls_client.create_run(
+                name="teapot_chat",
+                run_type="llm",
+                inputs={
+                    "context": truncated_context,
+                    "system_prompt": system_prompt,
+                    "question": query,
+                },
+            )
+            run_id = run.id
+        except:
+            pass
 
-    # 4️⃣ Stream Generation (timed)
+    # ---- STREAM OUTPUT ----
     with st.chat_message("assistant"):
         placeholder = st.empty()
 
         gen_start = time.perf_counter()
         final_text = ""
-
         for partial in stream_generate(prompt):
             final_text = partial
             placeholder.markdown(final_text)
@@ -348,7 +303,10 @@ if query:
         )
 
     if ls_client and run_id:
-        ls_client.update_run(run_id, outputs={"answer": final_text})
+        try:
+            ls_client.update_run(run_id, outputs={"answer": final_text})
+        except:
+            pass
 
     st.session_state.messages.append(
         {