Update app.py

app.py

@@ -1,20 +1,10 @@
 import os
 import requests
-import json
-from bs4 import BeautifulSoup
-from textwrap import shorten
-
 import gradio as gr
 
-from langchain_core.documents import Document
-from langchain_text_splitters import RecursiveCharacterTextSplitter
-from langchain_community.vectorstores import Chroma
-from langchain_community.embeddings import HuggingFaceEmbeddings
+# ---------------- RAG DOCUMENT SETUP ---------------- #
 
-# -----------------------
-# 1. SCRAPE K8S DOCS
-# -----------------------
-urls = {
+K8S_DOC_URLS = {
     "pods": "https://kubernetes.io/docs/concepts/workloads/pods/",
     "deployments": "https://kubernetes.io/docs/concepts/workloads/controllers/deployment/",
     "services": "https://kubernetes.io/docs/concepts/services-networking/service/",
@@ -27,121 +17,113 @@ urls = {
     "autoscaling": "https://kubernetes.io/docs/tasks/run-application/horizontal-pod-autoscale/"
 }
 
-def scrape_docs():
-    docs = []
-    for name, url in urls.items():
-        try:
-            r = requests.get(url, timeout=20)
-            soup = BeautifulSoup(r.text, "html.parser")
-            content = soup.find("div", class_="td-content")
-            if not content:
-                continue
-            text = content.get_text(separator="\n").strip()
-            docs.append(Document(page_content=text, metadata={"doc_id": name, "url": url}))
-        except Exception:
-            continue
-    return docs
-
-docs = scrape_docs()
-
-# -----------------------
-# 2. CHUNK + EMBED + VECTOR DB
-# -----------------------
-splitter = RecursiveCharacterTextSplitter(chunk_size=800, chunk_overlap=100)
-chunks = splitter.split_documents(docs)
-
-embedding = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
-vectordb = Chroma.from_documents(chunks, embedding)
-retriever = vectordb.as_retriever(
-    search_type="similarity_score_threshold",
-    search_kwargs={"k": 5, "score_threshold": 0.4}
-)
-
-# -----------------------
-# 3. RAG HELPERS
-# -----------------------
-def build_context_with_citations(query: str):
-    retrieved_docs = retriever.invoke(query)
-    context = ""
-    mapping = []
-
-    for i, d in enumerate(retrieved_docs, start=1):
-        label = f"[{i}]"
-        context += f"{label} {d.page_content[:1000]}\n\nSource: {d.metadata['url']}\n\n"
-        mapping.append({
-            "label": label,
-            "url": d.metadata["url"],
-            "doc": d.metadata["doc_id"],
-            "preview": shorten(d.page_content, width=200)
-        })
-    return context, mapping
-
-def build_prompt(query, context):
-    return f"""
-You are a Kubernetes expert.
-Use ONLY the context below.
-Add citations like [1][2] after each fact.
-If not found, say: 'Not in docs'.
-
-QUESTION:
-{query}
-
-CONTEXT:
-{context}
-""".strip()
-
-# -----------------------
-# 4. OPENROUTER LLM
-# -----------------------
-import requests as req
-
-OPENROUTER_API_KEY = os.environ.get("OPENROUTER_API_KEY", "")
-
-def call_llm(prompt: str) -> str:
-    if not OPENROUTER_API_KEY:
-        return "OpenRouter API key is not set. Please configure OPENROUTER_API_KEY in the Space settings."
-
+def fetch_doc(url):
+    try:
+        response = requests.get(url, timeout=10)
+        if response.status_code == 200:
+            return response.text
+    except:
+        return ""
+    return ""
+
+DOCUMENTS = [
+    {"doc": name, "url": url, "text": fetch_doc(url)}
+    for name, url in K8S_DOC_URLS.items()
+]
+
+def search_docs(query, top_k=3):
+    query = query.lower()
+    matches = []
+    for doc in DOCUMENTS:
+        text = doc["text"].lower()
+        if query in text:
+            snippet_start = text.index(query)
+            snippet_end = snippet_start + 350
+            snippet = doc["text"][snippet_start:snippet_end].replace("\n", " ")
+            matches.append((snippet, doc["url"], doc["doc"]))
+    return matches[:top_k]
+
+
+# --------------- LLM CALL (OpenRouter) ---------------- #
+
+def call_llm(prompt):
     url = "https://openrouter.ai/api/v1/chat/completions"
     headers = {
-        "Authorization": f"Bearer {OPENROUTER_API_KEY}",
-        "Content-Type": "application/json"
+        "Authorization": f"Bearer {os.getenv('OPENROUTER_API_KEY')}",
+        "HTTP-Referer": "https://huggingface.co/",
+        "X-Title": "Kubernetes RAG Assistant"
     }
+
     data = {
         "model": "meta-llama/llama-3.1-8b-instruct",
-        "messages": [
-            {"role": "system", "content": "You are a Kubernetes expert. Only use provided context."},
-            {"role": "user", "content": prompt}
-        ],
-        "temperature": 0.0
+        "messages": [{"role": "user", "content": prompt}],
+        "max_tokens": 350
     }
-    response = req.post(url, headers=headers, data=json.dumps(data))
-    out = response.json()
-    return out.get("choices", [{"message": {"content": "No response"}}])[0]["message"]["content"]
 
-def answer_question(query: str):
-    context, sources = build_context_with_citations(query)
-    prompt = build_prompt(query, context)
+    res = requests.post(url, json=data, headers=headers)
+    out = res.json()
+
+    if "choices" in out:
+        return out["choices"][0]["message"]["content"]
+    print("DEBUG LLM Error:", out)
+    return "⚠ Model error. Try again."
+
+
+# ----------- RAG + Prompt Construction ---------------- #
+
+def build_answer(query):
+    results = search_docs(query)
+    context = ""
+    citations = []
+
+    for i, (snippet, url, doc) in enumerate(results, start=1):
+        label = f"[{i}]"
+        context += f"{label}: {snippet}\n\n"
+        citations.append(f"{label} → {url}")
+
+    prompt = f"""
+    Use the context below to answer the question clearly.
+    Add citations like [1], [2] at the end of sentences.
+
+    Context:
+    {context}
+
+    Question: {query}
+    """
+
     answer = call_llm(prompt)
-    return answer, sources
-
-# -----------------------
-# 5. GRADIO CHAT APP
-# -----------------------
-def chat_fn(message, history):
-    answer, sources = answer_question(message)
-    src_lines = [f"{s['label']} – {s['url']}" for s in sources]
-    sources_text = "\n".join(src_lines) if src_lines else "No sources found."
-    full_answer = f"{answer}\n\n---\nSources:\n{sources_text}"
-    return full_answer
-
-demo = gr.ChatInterface(
-    fn=chat_fn,
-    title="Kubernetes RAG Assistant",
-    description="Ask Kubernetes questions. Answers are grounded in official docs and include citations."
-)
-
-def main():
-    return demo
-
-if __name__ == "__main__":
-    demo.launch()
+    citations_text = "\n".join(citations) or "No sources found."
+
+    return answer, citations_text
+
+
+# ---------------------- UI --------------------------- #
+
+custom_css = """
+.source-box {
+    font-size: 14px;
+    background: #1b2733;
+    padding: 10px;
+    border-radius: 8px;
+    color: #c9e2ff;
+    border: 1px solid #4a90e2;
+}
+"""
+
+with gr.Blocks(css=custom_css, theme="soft") as app:
+
+    gr.HTML("""
+    <h1 style='color:#326ce5; text-align:center;'>☸️ Kubernetes RAG Assistant</h1>
+    <p style='text-align:center; font-size:17px; color:#ddd;'>Ask any Kubernetes question and get answers with docs citations 📌</p>
+    """)
+
+    question = gr.Textbox(label="Ask a Kubernetes Question:", placeholder="e.g., What is RBAC in Kubernetes?")
+
+    answer = gr.Markdown(label="Answer")
+    sources = gr.Markdown(label="Sources", elem_classes=["source-box"])
+
+    submit = gr.Button("Ask ☸️")
+
+    submit.click(build_answer, inputs=question, outputs=[answer, sources])
+
+app.launch()