Update app.py

app.py

@@ -44,14 +44,19 @@ def scrape_page(name, url):
         return None
 
 def build_or_load_kb():
-    embedding_model = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
-
+    embedding_model = HuggingFaceEmbeddings(
+        model_name="sentence-transformers/all-MiniLM-L6-v2"
+    )
+    
     if os.path.isdir(PERSIST_DIR):
-        vectordb = Chroma(embedding_function=embedding_model, persist_directory=PERSIST_DIR)
-        raw = vectordb._collection.get(include=["documents", "metadatas"])
+        vectordb = Chroma(
+            embedding_function=embedding_model,
+            persist_directory=PERSIST_DIR
+        )
+        data = vectordb._collection.get(include=["documents", "metadatas"])
         chunks = [
             Document(page_content=doc, metadata=meta)
-            for doc, meta in zip(raw["documents"], raw["metadatas"])
+            for doc, meta in zip(data["documents"], data["metadatas"])
         ]
         return vectordb, chunks
 
@@ -64,165 +69,162 @@ def build_or_load_kb():
     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.persist()
-
+    vectordb = Chroma.from_documents(
+        chunks,
+        embedding_model,
+        persist_directory=PERSIST_DIR
+    )
     return vectordb, chunks
 
 vectordb, chunks = build_or_load_kb()
 
 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.35},
+    search_kwargs={"k": 8, "score_threshold": 0.4},
 )
 
-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
+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
     seen = set()
     unique = []
     for d in combined:
-        key = (d.metadata.get("doc_id"), d.page_content[:80])
+        key = (d.metadata.get("doc_id"), d.page_content[:50])
         if key not in seen:
             seen.add(key)
             unique.append(d)
+
     if not unique:
         return []
-    pairs = [(query, doc.page_content) for doc in unique]
+
+    pairs = [(query, d.page_content) for d in unique]
     scores = reranker.predict(pairs)
-    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]
+    ranked = sorted(zip(scores, unique), key=lambda x: x[0], reverse=True)[:5]
+
+    for s, d in ranked:
+        d.metadata["rerank_score"] = float(s)
+
+    return [d for _, d in ranked]
 
 def call_llm(prompt):
-    url = "https://openrouter.ai/api/v1/chat/completions"
     api_key = os.getenv("OPENROUTER_API_KEY")
     if not api_key:
-        return "⚠ Missing API key.\nGroundedness: 0%"
-    res = requests.post(url, headers={
-        "Authorization": f"Bearer {api_key}",
-        "HTTP-Referer": "https://huggingface.co/",
-        "X-Title": "Kubernetes RAG Assistant"
-    }, json={
-        "model": "meta-llama/llama-3.1-8b-instruct",
-        "messages": [{"role": "user", "content": prompt}],
-        "max_tokens": 400,
-        "temperature": 0
-    }).json()
-    return res["choices"][0]["message"]["content"]
-
-def build_context(query, history):
+        return "⚠️ Missing OPENROUTER_API_KEY environment variable.\nGroundedness: 0%"
+
+    try:
+        res = 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"
+            },
+            json={
+                "model": "meta-llama/llama-3.1-8b-instruct",
+                "messages": [{"role": "user", "content": prompt}],
+                "max_tokens": 300,
+                "temperature": 0.2
+            }
+        )
+        res.raise_for_status()
+        return res.json()["choices"][0]["message"]["content"]
+    except Exception as e:
+        return f"⚠️ LLM Error: {e}\nGroundedness: 0%"
+
+def build_context(query):
     docs = hybrid_search(query)
     if not docs:
         return "", [], []
-    context, sources, scores = "", [], []
+    ctx, srcs, 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']}")
+        ctx += f"[{i}] {d.page_content[:900]}\nSource: {d.metadata['url']}\n\n"
+        srcs.append(f"[{i}] → {d.metadata['url']}")
         scores.append(d.metadata["rerank_score"])
-    return context, sources, scores
+    return ctx, srcs, scores
 
-def classify_query(q):
+def init_metrics():
+    return {"q":[], "lat":[], "tok":[], "g":[], "cit":[], "r":[], "type":[]}
+
+def classify(q):
     q=q.lower()
-    if "how" in q: return "how-to"
-    if "error" in q: return "debug"
-    return "general"
+    return "how-to" if "how" in q else ("debug" if "error" in q else "general")
 
-def init_metrics():
-    return {"q":[], "lat":[], "tok":[], "g":[],"r":[],"c":[],"t":[]}
+def answer(q, history, metrics):
+    if metrics is None: metrics = init_metrics()
+    start  = time.time()
 
-def answer_question(query, history, metrics):
-    if metrics is None or metrics == {}: metrics = init_metrics()
-    start = time.time()
-    ctx, sources, scores = build_context(query, history)
+    ctx, srcs, scores = build_context(q)
     if not ctx:
-        reply="Not in docs.\nGroundedness: 0%"
-        history.append((query, reply))
+        txt = "Not in docs.\nGroundedness: 0%"
+        history.append((q, txt))
         return history,"",metrics
-    prompt=f"""
-Use ONLY context. Every sentence must end with citation [n].
-Answer:
-Question: {query}
+        
+    prompt = f"""Use context ONLY. Cite every sentence as [n].
+User question: {q}
+
 Context:
 {ctx}
-Groundedness must be in final line as: Groundedness: XX%
-"""
-    answer=call_llm(prompt)
-    latency=time.time()-start
-    grounded=int(re.search(r"Groundedness:\s*(\d+)%", answer).group(1)) if "Groundedness" in answer else 0
-    cites=len(set(re.findall(r"\[(\d+)\]", answer)))
-    avg_score=sum(scores)/len(scores)
-    tokens=len(answer.split())+len(prompt.split())
-    alert="⚠ Low support.\n\n" if grounded<70 or cites==0 else ""
-    final=alert+answer+"\n\n---\nSources:\n"+"\n".join(sources)
-    history.append((query,final))
-    metrics["q"].append(query)
+Groundedness MUST appear as: Groundedness: XX%"""
+    txt = 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["r"].append(avg_score)
-    metrics["c"].append(cites)
-    metrics["t"].append(classify_query(query))
+    metrics["cit"].append(cites)
+    metrics["r"].append(avg)
+    metrics["type"].append(classify(q))
+
     return history,"",metrics
 
 def render(metrics):
-    rows=[[i+1,metrics["q"][i],round(metrics["lat"][i],3),
-           metrics["tok"][i],metrics["g"][i],
-           round(metrics["r"][i],3),metrics["c"][i],metrics["t"][i]]
-          for i in range(len(metrics["q"]))]
-    avg_lat=sum(metrics["lat"])/len(metrics["lat"])
-    avg_g=sum(metrics["g"])/len(metrics["g"])
-    avg_tok=sum(metrics["tok"])/len(metrics["tok"])
-    return rows,avg_lat,avg_g,avg_tok
-
-def charts(metrics):
-    df=pd.DataFrame({
-        "Latency":metrics["lat"],
-        "Groundedness":metrics["g"],
-        "Tokens":metrics["tok"],
-        "Type":metrics["t"]
-    })
-    fig_l,ax=plt.subplots();ax.plot(df["Latency"]);ax.set_title("Latency");ax.set_xlabel("#");ax.set_ylabel("s")
-    fig_g,ax=plt.subplots();ax.plot(df["Groundedness"]);ax.set_title("Groundedness");ax.set_xlabel("#");ax.set_ylabel("%")
-    fig_t,ax=plt.subplots();ax.plot(df["Tokens"]);ax.set_title("Tokens");ax.set_xlabel("#");ax.set_ylabel("count")
-    fig_p,ax=plt.subplots();df["Type"].value_counts().plot.pie(ax=ax,autopct="%1.1f%");ax.set_ylabel("");ax.set_title("Query Types")
-    return fig_l,fig_g,fig_t,fig_p
-
-def export_csv(metrics):
-    df=pd.DataFrame(metrics)
-    path="analytics.csv";df.to_csv(path,index=False);return path
-
-def clear_all(): return [],"",init_metrics()
-
-metrics_state=gr.State(init_metrics())
-
-with gr.Blocks() as app:
+    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())
+
+with gr.Blocks(title="Kubernetes RAG Assistant") as app:
     gr.Markdown("# ☸ Kubernetes RAG Assistant")
     with gr.Tab("Chat"):
-        chat=gr.Chatbot()
-        user_in=gr.Textbox(label="Ask about Kubernetes")
-        clear=gr.Button("Clear")
-        user_in.submit(answer_question,[user_in,chat,metrics_state],[chat,user_in,metrics_state])
-        clear.click(clear_all,outputs=[chat,user_in,metrics_state])
+        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])
+
     with gr.Tab("Analytics"):
-        table=gr.Dataframe(headers=["ID","Query","Latency","Tokens","Grounded","Rerank","Citations","Type"])
-        avgL=gr.Number(label="Avg Latency");avgG=gr.Number(label="Avg Grounded");avgT=gr.Number(label="Avg Tokens")
-        p1,p2,p3,p4=gr.Plot(),gr.Plot(),gr.Plot(),gr.Plot()
-        refresh=gr.Button("Refresh")
-        export=gr.Button("Export CSV")
-        file=gr.File()
+        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")
+        refresh = gr.Button("Update Dashboard")
         refresh.click(render,[metrics_state],[table,avgL,avgG,avgT])
-        refresh.click(charts,[metrics_state],[p1,p2,p3,p4])
-        export.click(export_csv,[metrics_state],[file])
 
 app.launch()