Update app.py

app.py

@@ -16,6 +16,7 @@ from langchain_community.vectorstores import Chroma
 from rank_bm25 import BM25Okapi
 from sentence_transformers import CrossEncoder
 
+
 PERSIST_DIR = "k8s_chroma_db"
 
 URLS = {
@@ -32,16 +33,18 @@ URLS = {
 }
 
 
-# ----------------- SCRAPING + KB ----------------- #
+# ================= Knowledge Base ================= #
 
 def scrape_page(name, url):
     try:
-        r = requests.get(url, timeout=20)
-        r.raise_for_status()
-        soup = BeautifulSoup(r.text, "html.parser")
+        response = requests.get(url, timeout=20)
+        response.raise_for_status()
+        soup = BeautifulSoup(response.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:
@@ -50,48 +53,48 @@ 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 DB exists, load it
     if os.path.isdir(PERSIST_DIR):
-        print("[INFO] Loading existing Chroma DB")
+        print("[INFO] Loading existing vector DB...")
         vectordb = Chroma(
             embedding_function=embedding_model,
             persist_directory=PERSIST_DIR,
         )
         raw = vectordb._collection.get(include=["documents", "metadatas"])
         chunks = [
-            Document(page_content=doc, metadata=meta)
-            for doc, meta in zip(raw["documents"], raw["metadatas"])
+            Document(page_content=d, metadata=m)
+            for d, m in zip(raw["documents"], raw["metadatas"])
         ]
         return vectordb, chunks
 
-    # Else: scrape + build
-    print("[INFO] No DB found, scraping docs...")
+    print("[INFO] No DB found — scraping docs...")
     docs = []
     for name, url in URLS.items():
-        d = scrape_page(name, url)
-        if d:
-            docs.append(d)
+        doc = scrape_page(name, url)
+        if doc:
+            docs.append(doc)
     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
-    )
+    vectordb = Chroma.from_documents(chunks, embedding_model, persist_directory=PERSIST_DIR)
+    vectordb.persist()
+
+    print("[INFO] Vector DB built & saved.")
     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)
+# ================= Search & Reranker ================= #
+
+bm25 = BM25Okapi([c.page_content.split() for c in chunks])
 reranker = CrossEncoder("cross-encoder/ms-marco-MiniLM-L-12-v2")
 
 retriever = vectordb.as_retriever(
@@ -103,41 +106,40 @@ retriever = vectordb.as_retriever(
 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]]
+    bm_scores = bm25.get_scores(query.lower().split())
+    bm_ranked = sorted(zip(bm_scores, chunks), reverse=True)
+    bm_results = [doc for _, doc in bm_ranked[:top_k]]
 
-    combined = vector_results + bm25_results
+    unique_docs = []
     seen = set()
-    unique = []
-    for d in combined:
-        key = (d.metadata.get("doc_id"), d.page_content[:80])
+    for doc in vector_results + bm_results:
+        key = (doc.metadata.get("doc_id"), doc.page_content[:50])
         if key not in seen:
             seen.add(key)
-            unique.append(d)
+            unique_docs.append(doc)
 
-    if not unique:
+    if not unique_docs:
         return []
 
-    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)[:top_k]
+    rerank_pairs = [(query, doc.page_content) for doc in unique_docs]
+    scores = reranker.predict(rerank_pairs)
+    ranked = sorted(zip(scores, unique_docs), reverse=True)[:top_k]
 
     for s, doc in ranked:
         doc.metadata["rerank_score"] = float(s)
+
     return [doc for _, doc in ranked]
 
 
-# ----------------- LLM CALL ----------------- #
+# ================= LLM ================= #
 
-def call_llm(prompt: str) -> str:
+def call_llm(prompt):
     api_key = os.getenv("OPENROUTER_API_KEY")
     if not api_key:
-        return "⚠️ Missing OPENROUTER_API_KEY in Space secrets.\nGroundedness: 0%"
+        return "⚠️ Missing API key.\nGroundedness: 0%"
 
     try:
-        r = requests.post(
+        res = requests.post(
             "https://openrouter.ai/api/v1/chat/completions",
             headers={
                 "Authorization": f"Bearer {api_key}",
@@ -147,36 +149,20 @@ def call_llm(prompt: str) -> str:
             json={
                 "model": "meta-llama/llama-3.1-8b-instruct",
                 "messages": [{"role": "user", "content": prompt}],
-                "max_tokens": 400,
                 "temperature": 0.0,
+                "max_tokens": 400,
             },
-            timeout=60,
         )
-        r.raise_for_status()
-        data = r.json()
+        res.raise_for_status()
+        data = res.json()
         return data["choices"][0]["message"]["content"]
     except Exception as e:
-        print("[ERROR] LLM:", e)
-        return f"⚠️ LLM error: {e}\nGroundedness: 0%"
-
-
-# ----------------- CONTEXT + METRICS ----------------- #
-
-def build_context(query: str):
-    docs = hybrid_search(query)
-    if not docs:
-        return "", [], []
+        return f"⚠️ LLM Error: {e}\nGroundedness: 0%"
 
-    context, sources, scores = "", [], []
-    for i, d in enumerate(docs, start=1):
-        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 context, sources, scores
 
+# ================= Analytics ================= #
 
-def classify_query(q: str) -> str:
+def classify_query(q):
     q = q.lower()
     if "how" in q:
         return "how-to"
@@ -185,134 +171,110 @@ def classify_query(q: str) -> str:
     return "general"
 
 
-def init_metrics():
-    return {"q": [], "lat": [], "tok": [], "g": [], "r": [], "c": [], "t": []}
-
+METRICS = {"q": [], "lat": [], "tok": [], "g": [], "r": [], "c": [], "t": []}
 
-# global analytics, no gr.State
-METRICS = init_metrics()
 
-
-# ----------------- CHAT HANDLER ----------------- #
+# ================= Chat Handler ================= #
 
 def answer_question(query, history):
-    global METRICS
-    if METRICS is None:
-        METRICS = init_metrics()
-
     start = time.time()
-    ctx, sources, scores = build_context(query)
+    docs = hybrid_search(query)
 
-    if not ctx:
-        reply = "Not in docs or insufficient context.\nGroundedness: 0%"
-        history.append((query, reply))
-        return history, ""
+    if not docs:
+        reply = "Not found in docs.\nGroundedness: 0%"
+        return history + [
+            {"role": "user", "content": query},
+            {"role": "assistant", "content": reply}
+        ], ""
+
+    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']}")
+        scores.append(d.metadata["rerank_score"])
 
     prompt = f"""
-Use ONLY the context below to answer.
-Every sentence MUST end with a citation like [1].
+Answer the question ONLY using the context below.
+Each sentence MUST end with a citation like [1].
 
 Question: {query}
 
 Context:
 {ctx}
 
-At the end add a line: Groundedness: XX%
+End with: Groundedness: XX%
 """
+
     answer = call_llm(prompt)
     latency = time.time() - start
 
-    # 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
+        grounded = int(m.group(1"))
 
     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)
+    avg_score = sum(scores) / len(scores) if scores else 0
 
-    history.append((query, final))
+    final = answer + "\n\n---\nSources:\n" + "\n".join(sources)
 
     METRICS["q"].append(query)
     METRICS["lat"].append(latency)
-    METRICS["tok"].append(tokens)
+    METRICS["tok"].append(len(answer.split()))
     METRICS["g"].append(grounded)
     METRICS["r"].append(avg_score)
     METRICS["c"].append(cites)
     METRICS["t"].append(classify_query(query))
 
+    history.append({"role": "user", "content": query})
+    history.append({"role": "assistant", "content": final})
     return history, ""
 
 
-# ----------------- ANALYTICS HELPERS ----------------- #
+def update_dashboard():
+    rows = list(zip(
+        range(1, len(METRICS["q"])+1),
+        METRICS["q"],
+        METRICS["lat"],
+        METRICS["tok"],
+        METRICS["g"],
+        METRICS["r"],
+        METRICS["c"],
+        METRICS["t"],
+    ))
 
-def render_metrics():
-    if len(METRICS["q"]) == 0:
-        return [], 0.0, 0.0, 0.0
+    avgG = round(sum(METRICS["g"]) / len(METRICS["g"]), 2)
+    avgL = round(sum(METRICS["lat"]) / len(METRICS["lat"]), 2)
+    avgT = round(sum(METRICS["tok"]) / len(METRICS["tok"]), 2)
 
-    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],
-        ])
+    return rows, avgG, avgL, avgT
 
-    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 ----------------- #
+# ================= UI ================= #
 
 with gr.Blocks(title="Kubernetes RAG Assistant") as app:
     gr.Markdown("# ☸ Kubernetes RAG Assistant")
 
     with gr.Tab("Chat"):
         chat = gr.Chatbot(height=450)
-        inp = gr.Textbox(label="Ask anything about Kubernetes")
-        clear_btn = gr.Button("Reset Conversation")
+        user_in = gr.Textbox(label="Ask anything about Kubernetes")
+        reset = gr.Button("Reset")
 
-        inp.submit(answer_question, [inp, chat], [chat, inp])
-        clear_btn.click(lambda: ([], ""), None, [chat, inp])
+        user_in.submit(answer_question, [user_in, chat], [chat, user_in])
+        reset.click(lambda: ([], ""), None, [chat, user_in])
 
     with gr.Tab("Analytics"):
-        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)")
+        gr.Markdown("### 📊 Analytics This Session")
+        table = gr.DataFrame(headers=[
+            "ID","Query","Latency","Tokens","Grounded","Rerank","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("Update Dashboard")
-        refresh.click(render_metrics, None, [table, avgG, avgL, avgT])
+        refresh = gr.Button("Refresh")
+        refresh.click(update_dashboard, None, [table, avgG, avgL, avgT])
 
 app.launch()