Update app.py

app.py

@@ -31,75 +31,88 @@ URLS = {
     "autoscaling": "https://kubernetes.io/docs/tasks/run-application/horizontal-pod-autoscale/",
 }
 
+
+# ----------------- SCRAPING + KB ----------------- #
+
 def scrape_page(name, url):
     try:
         r = requests.get(url, timeout=20)
+        r.raise_for_status()
         soup = BeautifulSoup(r.text, "html.parser")
         content = soup.find("div", class_="td-content")
         if not content:
             return None
         text = content.get_text(separator="\n").strip()
         return Document(page_content=text, metadata={"doc_id": name, "url": url})
-    except:
+    except Exception as e:
+        print(f"[ERROR] scraping {url}: {e}")
         return None
 
+
 def build_or_load_kb():
     embedding_model = HuggingFaceEmbeddings(
         model_name="sentence-transformers/all-MiniLM-L6-v2"
     )
-    
+
+    # If DB exists, load it
     if os.path.isdir(PERSIST_DIR):
+        print("[INFO] Loading existing Chroma DB")
         vectordb = Chroma(
             embedding_function=embedding_model,
-            persist_directory=PERSIST_DIR
+            persist_directory=PERSIST_DIR,
         )
-        data = vectordb._collection.get(include=["documents", "metadatas"])
+        raw = vectordb._collection.get(include=["documents", "metadatas"])
         chunks = [
             Document(page_content=doc, metadata=meta)
-            for doc, meta in zip(data["documents"], data["metadatas"])
+            for doc, meta in zip(raw["documents"], raw["metadatas"])
         ]
         return vectordb, chunks
 
+    # Else: scrape + build
+    print("[INFO] No DB found, scraping docs...")
     docs = []
     for name, url in URLS.items():
         d = scrape_page(name, url)
         if d:
             docs.append(d)
+    print(f"[INFO] Scraped {len(docs)} docs")
 
     splitter = RecursiveCharacterTextSplitter(chunk_size=900, chunk_overlap=200)
     chunks = splitter.split_documents(docs)
 
     vectordb = Chroma.from_documents(
-        chunks,
-        embedding_model,
-        persist_directory=PERSIST_DIR
+        chunks, embedding_model, persist_directory=PERSIST_DIR
     )
     return vectordb, chunks
 
+
 vectordb, chunks = build_or_load_kb()
 
+# ----------------- HYBRID SEARCH ----------------- #
+
 bm25_corpus = [doc.page_content.split() for doc in chunks]
 bm25 = BM25Okapi(bm25_corpus)
-
 reranker = CrossEncoder("cross-encoder/ms-marco-MiniLM-L-12-v2")
 
 retriever = vectordb.as_retriever(
     search_type="similarity_score_threshold",
-    search_kwargs={"k": 8, "score_threshold": 0.4},
+    search_kwargs={"k": 8, "score_threshold": 0.35},
 )
 
-def hybrid_search(query):
-    vresults = retriever.invoke(query)
-    tokens = query.lower().split()
-    bm_scores = bm25.get_scores(tokens)
-    bm_ranked = sorted(zip(bm_scores, chunks), key=lambda x: x[0], reverse=True)
-    bmresults = [d for _, d in bm_ranked[:5]]
 
-    combined = vresults + bmresults
+def hybrid_search(query, top_k=5):
+    vector_results = retriever.invoke(query)
+
+    tokenized_query = query.lower().split()
+    bm25_scores = bm25.get_scores(tokenized_query)
+    bm25_ranked = sorted(zip(bm25_scores, chunks), key=lambda x: x[0], reverse=True)
+    bm25_results = [d for _, d in bm25_ranked[:top_k]]
+
+    combined = vector_results + bm25_results
     seen = set()
     unique = []
     for d in combined:
-        key = (d.metadata.get("doc_id"), d.page_content[:50])
+        key = (d.metadata.get("doc_id"), d.page_content[:80])
         if key not in seen:
             seen.add(key)
             unique.append(d)
@@ -107,124 +120,199 @@ def hybrid_search(query):
     if not unique:
         return []
 
-    pairs = [(query, d.page_content) for d in unique]
+    pairs = [(query, doc.page_content) for doc in unique]
     scores = reranker.predict(pairs)
-    ranked = sorted(zip(scores, unique), key=lambda x: x[0], reverse=True)[:5]
+    ranked = sorted(zip(scores, unique), key=lambda x: x[0], reverse=True)[:top_k]
+
+    for s, doc in ranked:
+        doc.metadata["rerank_score"] = float(s)
+    return [doc for _, doc in ranked]
 
-    for s, d in ranked:
-        d.metadata["rerank_score"] = float(s)
 
-    return [d for _, d in ranked]
+# ----------------- LLM CALL ----------------- #
 
-def call_llm(prompt):
+def call_llm(prompt: str) -> str:
     api_key = os.getenv("OPENROUTER_API_KEY")
     if not api_key:
-        return "⚠️ Missing OPENROUTER_API_KEY environment variable.\nGroundedness: 0%"
+        return "⚠️ Missing OPENROUTER_API_KEY in Space secrets.\nGroundedness: 0%"
 
     try:
-        res = requests.post(
+        r = requests.post(
             "https://openrouter.ai/api/v1/chat/completions",
             headers={
                 "Authorization": f"Bearer {api_key}",
                 "HTTP-Referer": "https://huggingface.co/",
-                "X-Title": "Kubernetes RAG Assistant"
+                "X-Title": "Kubernetes RAG Assistant",
             },
             json={
                 "model": "meta-llama/llama-3.1-8b-instruct",
                 "messages": [{"role": "user", "content": prompt}],
-                "max_tokens": 300,
-                "temperature": 0.2
-            }
+                "max_tokens": 400,
+                "temperature": 0.0,
+            },
+            timeout=60,
         )
-        res.raise_for_status()
-        return res.json()["choices"][0]["message"]["content"]
+        r.raise_for_status()
+        data = r.json()
+        return data["choices"][0]["message"]["content"]
     except Exception as e:
-        return f"⚠️ LLM Error: {e}\nGroundedness: 0%"
+        print("[ERROR] LLM:", e)
+        return f"⚠️ LLM error: {e}\nGroundedness: 0%"
+
+
+# ----------------- CONTEXT + METRICS ----------------- #
 
-def build_context(query):
+def build_context(query: str):
     docs = hybrid_search(query)
     if not docs:
         return "", [], []
-    ctx, srcs, scores = "", [], []
+
+    context, sources, scores = "", [], []
     for i, d in enumerate(docs, start=1):
-        ctx += f"[{i}] {d.page_content[:900]}\nSource: {d.metadata['url']}\n\n"
-        srcs.append(f"[{i}] → {d.metadata['url']}")
+        label = f"[{i}]"
+        context += f"{label} {d.page_content[:900]}\nSource: {d.metadata['url']}\n\n"
+        sources.append(f"{label} → {d.metadata['url']}")
         scores.append(d.metadata["rerank_score"])
-    return ctx, srcs, scores
+    return context, sources, scores
+
+
+def classify_query(q: str) -> str:
+    q = q.lower()
+    if "how" in q:
+        return "how-to"
+    if "error" in q or "fail" in q:
+        return "debug"
+    return "general"
+
 
 def init_metrics():
-    return {"q":[], "lat":[], "tok":[], "g":[], "cit":[], "r":[], "type":[]}
+    return {"q": [], "lat": [], "tok": [], "g": [], "r": [], "c": [], "t": []}
+
+
+# global analytics, no gr.State
+METRICS = init_metrics()
+
+
+# ----------------- CHAT HANDLER ----------------- #
 
-def classify(q):
-    q=q.lower()
-    return "how-to" if "how" in q else ("debug" if "error" in q else "general")
+def answer_question(query, history):
+    global METRICS
+    if METRICS is None:
+        METRICS = init_metrics()
 
-def answer(q, history, metrics):
-    if metrics is None: metrics = init_metrics()
-    start  = time.time()
+    start = time.time()
+    ctx, sources, scores = build_context(query)
 
-    ctx, srcs, scores = build_context(q)
     if not ctx:
-        txt = "Not in docs.\nGroundedness: 0%"
-        history.append((q, txt))
-        return history,"",metrics
-        
-    prompt = f"""Use context ONLY. Cite every sentence as [n].
-User question: {q}
+        reply = "Not in docs or insufficient context.\nGroundedness: 0%"
+        history.append((query, reply))
+        return history, ""
+
+    prompt = f"""
+Use ONLY the context below to answer.
+Every sentence MUST end with a citation like [1].
+
+Question: {query}
 
 Context:
 {ctx}
-Groundedness MUST appear as: Groundedness: XX%"""
-    txt = call_llm(prompt)
 
+At the end add a line: Groundedness: XX%
+"""
+    answer = call_llm(prompt)
     latency = time.time() - start
-    grounded = int(re.search(r"Groundedness:\s*(\d+)%", txt).group(1)) if "Groundedness" in txt else 0
-    tokens = len(txt.split())
-    cites = len(set(re.findall(r"\[(\d+)\]", txt)))
-    avg = sum(scores)/len(scores)
-
-    final = txt+"\n\nSources:\n"+"\n".join(srcs)
-    history.append((q, final))
-
-    metrics["q"].append(q)
-    metrics["lat"].append(latency)
-    metrics["tok"].append(tokens)
-    metrics["g"].append(grounded)
-    metrics["cit"].append(cites)
-    metrics["r"].append(avg)
-    metrics["type"].append(classify(q))
-
-    return history,"",metrics
-
-def render(metrics):
-    if len(metrics["q"])==0: return [],0,0,0
-    rows=[[
-        i+1, metrics["q"][i], round(metrics["lat"][i],3),
-        metrics["tok"][i], metrics["g"][i],
-        round(metrics["r"][i],2), metrics["cit"][i], metrics["type"][i]
-    ] for i in range(len(metrics["q"]))]
-    avgL=sum(metrics["g"])/len(metrics["g"])
-    avgG=sum(metrics["lat"])/len(metrics["lat"])
-    avgT=sum(metrics["tok"])/len(metrics["tok"])
-    return rows,avgL,avgG,avgT
-
-metrics_state = gr.State(init_metrics())
+
+    # robust groundedness parsing
+    grounded = 0
+    m = re.search(r"Groundedness:\s*(\d+)%", answer)
+    if m:
+        try:
+            grounded = int(m.group(1))
+        except ValueError:
+            grounded = 0
+
+    cites = len(set(re.findall(r"\[(\d+)\]", answer)))
+    avg_score = sum(scores) / len(scores) if scores else 0.0
+    tokens = len(answer.split()) + len(prompt.split())
+
+    alert = ""
+    if grounded < 70 or cites == 0:
+        alert = "⚠️ Low support from docs; please verify in official Kubernetes docs.\n\n"
+
+    final = alert + answer + "\n\n---\nSources:\n" + "\n".join(sources)
+
+    history.append((query, final))
+
+    METRICS["q"].append(query)
+    METRICS["lat"].append(latency)
+    METRICS["tok"].append(tokens)
+    METRICS["g"].append(grounded)
+    METRICS["r"].append(avg_score)
+    METRICS["c"].append(cites)
+    METRICS["t"].append(classify_query(query))
+
+    return history, ""
+
+
+# ----------------- ANALYTICS HELPERS ----------------- #
+
+def render_metrics():
+    if len(METRICS["q"]) == 0:
+        return [], 0.0, 0.0, 0.0
+
+    rows = []
+    for i, q in enumerate(METRICS["q"]):
+        rows.append([
+            i + 1,
+            q,
+            round(METRICS["lat"][i], 3),
+            METRICS["tok"][i],
+            METRICS["g"][i],
+            round(METRICS["r"][i], 3),
+            METRICS["c"][i],
+            METRICS["t"][i],
+        ])
+
+    avg_ground = sum(METRICS["g"]) / len(METRICS["g"])
+    avg_lat = sum(METRICS["lat"]) / len(METRICS["lat"])
+    avg_tok = sum(METRICS["tok"]) / len(METRICS["tok"])
+
+    return rows, avg_ground, avg_lat, avg_tok
+
+
+# ----------------- GRADIO UI ----------------- #
 
 with gr.Blocks(title="Kubernetes RAG Assistant") as app:
     gr.Markdown("# ☸ Kubernetes RAG Assistant")
+
     with gr.Tab("Chat"):
-        chat = gr.Chatbot()
-        inp  = gr.Textbox(label="Ask anything about Kubernetes")
-        clear= gr.Button("Reset")
-        inp.submit(answer,[inp,chat,metrics_state],[chat,inp,metrics_state])
-        clear.click(lambda: ([], "", init_metrics()), None, [chat,inp,metrics_state])
+        chat = gr.Chatbot(height=450)
+        inp = gr.Textbox(label="Ask anything about Kubernetes")
+        clear_btn = gr.Button("Reset Conversation")
+
+        inp.submit(answer_question, [inp, chat], [chat, inp])
+        clear_btn.click(lambda: ([], ""), None, [chat, inp])
 
     with gr.Tab("Analytics"):
-        table   = gr.DataFrame(headers=["ID","Query","Latency","Tokens","Grounded","Relevance","Citations","Type"])
-        avgL    = gr.Number(label="Avg Groundedness")
-        avgG    = gr.Number(label="Avg Latency")  
-        avgT    = gr.Number(label="Avg Tokens")
+        gr.Markdown("### 📊 Query Analytics (this session)")
+        table = gr.DataFrame(
+            headers=[
+                "ID",
+                "Query",
+                "Latency (s)",
+                "Tokens",
+                "Groundedness (%)",
+                "Avg Rerank Score",
+                "Citations",
+                "Type",
+            ],
+            interactive=False,
+        )
+        avgG = gr.Number(label="Avg Groundedness (%)")
+        avgL = gr.Number(label="Avg Latency (s)")
+        avgT = gr.Number(label="Avg Tokens")
+
         refresh = gr.Button("Update Dashboard")
-        refresh.click(render,[metrics_state],[table,avgL,avgG,avgT])
+        refresh.click(render_metrics, None, [table, avgG, avgL, avgT])
 
 app.launch()