Update app.py

app.py

@@ -33,18 +33,16 @@ URLS = {
 }
 
 
-# ================= Knowledge Base ================= #
+# ------------------ Knowledge Base ------------------ #
 
 def scrape_page(name, url):
     try:
-        response = requests.get(url, timeout=20)
-        response.raise_for_status()
-        soup = BeautifulSoup(response.text, "html.parser")
+        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 Exception as e:
@@ -53,13 +51,12 @@ def scrape_page(name, url):
 
 
 def build_or_load_kb():
-    print("[INFO] Loading embedding model...")
     embedding_model = HuggingFaceEmbeddings(
         model_name="sentence-transformers/all-MiniLM-L6-v2"
     )
 
     if os.path.isdir(PERSIST_DIR):
-        print("[INFO] Loading existing vector DB...")
+        print("[INFO] Loading existing DB...")
         vectordb = Chroma(
             embedding_function=embedding_model,
             persist_directory=PERSIST_DIR,
@@ -74,9 +71,9 @@ def build_or_load_kb():
     print("[INFO] No DB found — scraping docs...")
     docs = []
     for name, url in URLS.items():
-        doc = scrape_page(name, url)
-        if doc:
-            docs.append(doc)
+        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)
@@ -84,16 +81,12 @@ def build_or_load_kb():
 
     vectordb = Chroma.from_documents(chunks, embedding_model, persist_directory=PERSIST_DIR)
     vectordb.persist()
-
-    print("[INFO] Vector DB built & saved.")
+    print("[INFO] DB created.")
     return vectordb, chunks
 
 
 vectordb, chunks = build_or_load_kb()
 
-
-# ================= Search & Reranker ================= #
-
 bm25 = BM25Okapi([c.page_content.split() for c in chunks])
 reranker = CrossEncoder("cross-encoder/ms-marco-MiniLM-L-12-v2")
 
@@ -113,7 +106,7 @@ def hybrid_search(query, top_k=5):
     unique_docs = []
     seen = set()
     for doc in vector_results + bm_results:
-        key = (doc.metadata.get("doc_id"), doc.page_content[:50])
+        key = (doc.metadata.get("doc_id"), doc.page_content[:60])
         if key not in seen:
             seen.add(key)
             unique_docs.append(doc)
@@ -121,22 +114,22 @@ def hybrid_search(query, top_k=5):
     if not unique_docs:
         return []
 
-    rerank_pairs = [(query, doc.page_content) for doc in unique_docs]
-    scores = reranker.predict(rerank_pairs)
+    pairs = [(query, doc.page_content) for doc in unique_docs]
+    scores = reranker.predict(pairs)
     ranked = sorted(zip(scores, unique_docs), reverse=True)[:top_k]
 
-    for s, doc in ranked:
-        doc.metadata["rerank_score"] = float(s)
+    for score, doc in ranked:
+        doc.metadata["rerank_score"] = float(score)
 
-    return [doc for _, doc in ranked]
+    return [doc for score, doc in ranked]
 
 
-# ================= LLM ================= #
+# ------------------ LLM ------------------ #
 
 def call_llm(prompt):
     api_key = os.getenv("OPENROUTER_API_KEY")
     if not api_key:
-        return "⚠️ Missing API key.\nGroundedness: 0%"
+        return "⚠ Missing OPENROUTER_API_KEY\nGroundedness: 0%"
 
     try:
         res = requests.post(
@@ -149,33 +142,29 @@ def call_llm(prompt):
             json={
                 "model": "meta-llama/llama-3.1-8b-instruct",
                 "messages": [{"role": "user", "content": prompt}],
-                "temperature": 0.0,
                 "max_tokens": 400,
+                "temperature": 0.0,
             },
+            timeout=60
         )
         res.raise_for_status()
-        data = res.json()
-        return data["choices"][0]["message"]["content"]
+        return res.json()["choices"][0]["message"]["content"]
     except Exception as e:
-        return f"⚠️ LLM Error: {e}\nGroundedness: 0%"
+        return f"⚠ LLM error: {e}\nGroundedness: 0%"
 
 
-# ================= Analytics ================= #
+# ------------------ Chat + Metrics ------------------ #
+
+METRICS = {"q": [], "lat": [], "tok": [], "g": [], "r": [], "c": [], "t": []}
+
 
 def classify_query(q):
     q = q.lower()
-    if "how" in q:
-        return "how-to"
-    if "error" in q or "fail" in q:
-        return "debug"
+    if "how" in q: return "how-to"
+    if "error" in q or "fail" in q: return "debug"
     return "general"
 
 
-METRICS = {"q": [], "lat": [], "tok": [], "g": [], "r": [], "c": [], "t": []}
-
-
-# ================= Chat Handler ================= #
-
 def answer_question(query, history):
     start = time.time()
     docs = hybrid_search(query)
@@ -184,21 +173,20 @@ def answer_question(query, history):
         reply = "Not found in docs.\nGroundedness: 0%"
         return history + [
             {"role": "user", "content": query},
-            {"role": "assistant", "content": reply}
+            {"role": "assistant", "content": reply},
         ], ""
 
+    scores = []
     ctx = ""
     sources = []
-    scores = []
     for i, d in enumerate(docs, 1):
-        label = f"[{i}]"
-        ctx += f"{label} {d.page_content[:900]}\nSource: {d.metadata['url']}\n\n"
-        sources.append(f"{label} → {d.metadata['url']}")
+        ctx += f"[{i}] {d.page_content[:900]}\nSource: {d.metadata['url']}\n\n"
+        sources.append(f"[{i}] → {d.metadata['url']}")
         scores.append(d.metadata["rerank_score"])
 
     prompt = f"""
-Answer the question ONLY using the context below.
-Each sentence MUST end with a citation like [1].
+Answer using ONLY the context below.
+Each sentence MUST include citation like [1].
 
 Question: {query}
 
@@ -214,13 +202,14 @@ End with: Groundedness: XX%
     grounded = 0
     m = re.search(r"Groundedness:\s*(\d+)%", answer)
     if m:
-        grounded = int(m.group(1"))
+        grounded = int(m.group(1))
 
     cites = len(set(re.findall(r"\[(\d+)\]", answer)))
-    avg_score = sum(scores) / len(scores) if scores else 0
+    avg_score = sum(scores) / len(scores)
 
     final = answer + "\n\n---\nSources:\n" + "\n".join(sources)
 
+    # Log metrics correctly
     METRICS["q"].append(query)
     METRICS["lat"].append(latency)
     METRICS["tok"].append(len(answer.split()))
@@ -236,7 +225,7 @@ End with: Groundedness: XX%
 
 def update_dashboard():
     rows = list(zip(
-        range(1, len(METRICS["q"])+1),
+        range(1, len(METRICS["q"]) + 1),
         METRICS["q"],
         METRICS["lat"],
         METRICS["tok"],
@@ -253,28 +242,32 @@ def update_dashboard():
     return rows, avgG, avgL, avgT
 
 
-# ================= UI ================= #
+# ------------------ UI ------------------ #
 
 with gr.Blocks(title="Kubernetes RAG Assistant") as app:
     gr.Markdown("# ☸ Kubernetes RAG Assistant")
 
     with gr.Tab("Chat"):
         chat = gr.Chatbot(height=450)
-        user_in = gr.Textbox(label="Ask anything about Kubernetes")
-        reset = gr.Button("Reset")
+        user_in = gr.Textbox(label="Ask about Kubernetes")
+        clear = gr.Button("Clear")
 
         user_in.submit(answer_question, [user_in, chat], [chat, user_in])
-        reset.click(lambda: ([], ""), None, [chat, user_in])
+        clear.click(lambda: ([], ""), None, [chat, user_in])
 
     with gr.Tab("Analytics"):
-        gr.Markdown("### 📊 Analytics This Session")
-        table = gr.DataFrame(headers=[
-            "ID","Query","Latency","Tokens","Grounded","Rerank","Citations","Type"
-        ], interactive=False)
+        gr.Markdown("### 📊 Query Analytics")
+        table = gr.DataFrame(
+            headers=[
+                "ID", "Query", "Latency", "Tokens",
+                "Groundedness", "Rerank Score", "Citations", "Type",
+            ],
+            interactive=False
+        )
         avgG = gr.Number(label="Avg Groundedness")
         avgL = gr.Number(label="Avg Latency")
         avgT = gr.Number(label="Avg Tokens")
-        refresh = gr.Button("Refresh")
-        refresh.click(update_dashboard, None, [table, avgG, avgL, avgT])
+        update = gr.Button("Refresh Dashboard")
+        update.click(update_dashboard, None, [table, avgG, avgL, avgT])
 
 app.launch()