Update app.py

app.py

@@ -7,10 +7,10 @@ from textwrap import shorten
 
 from langchain_core.documents import Document
 from langchain_text_splitters import RecursiveCharacterTextSplitter
-from langchain_community.embeddings import HuggingFaceEmbeddings
+from langchain_huggingface import HuggingFaceEmbeddings
 from langchain_community.vectorstores import Chroma
 
-# ------------------ 1. SCRAPE K8S DOCS ------------------ #
+# ------------------ SCRAPE KUBERNETES DOCS ------------------ #
 
 URLS = {
     "pods": "https://kubernetes.io/docs/concepts/workloads/pods/",
@@ -37,8 +37,7 @@ def scrape_page(name, url):
             page_content=text,
             metadata={"doc_id": name, "url": url}
         )
-    except Exception as e:
-        print(f"Error scraping {name}: {e}")
+    except:
         return None
 
 docs = []
@@ -47,17 +46,12 @@ for name, url in URLS.items():
     if d:
         docs.append(d)
 
-# ------------------ 2. CHUNK + EMBED + CHROMA ------------------ #
+# ------------------ CHUNK + EMBEDDINGS + VECTOR DB ------------------ #
 
-splitter = RecursiveCharacterTextSplitter(
-    chunk_size=800,
-    chunk_overlap=120
-)
+splitter = RecursiveCharacterTextSplitter(chunk_size=900, chunk_overlap=200)
 chunks = splitter.split_documents(docs)
 
-embedding_model = HuggingFaceEmbeddings(
-    model_name="sentence-transformers/all-MiniLM-L6-v2"
-)
+embedding_model = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
 
 vectordb = Chroma.from_documents(chunks, embedding_model)
 retriever = vectordb.as_retriever(
@@ -65,119 +59,80 @@ retriever = vectordb.as_retriever(
     search_kwargs={"k": 5, "score_threshold": 0.4}
 )
 
-# ------------------ 3. RAG HELPERS ------------------ #
-
-def build_context_with_citations(query: str):
-    retrieved = retriever.invoke(query)
-    context = ""
-    mapping = []
-    for i, d in enumerate(retrieved, start=1):
-        label = f"[{i}]"
-        context += (
-            f"{label} {d.page_content[:900]}\n"
-            f"Source: {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, history_str: str):
-    return f"""
-You are a Kubernetes expert assistant.
-
-Follow these rules:
-1. Use ONLY the context below.
-2. Every factual statement MUST have citations like [1], [2].
-3. If the answer is not in the context, say: "Not in docs."
-
-Conversation so far:
-{history_str}
-
-User question: {query}
-
-Context:
-{context}
-""".strip()
-
-# ------------------ 4. OPENROUTER LLM ------------------ #
-
-def call_llm(prompt: str) -> str:
-    api_key = os.getenv("OPENROUTER_API_KEY", "")
-    if not api_key:
-        return "⚠ OPENROUTER_API_KEY is not set in this Space."
+# ------------------ LLM CALL (OPENROUTER) ------------------ #
 
+def call_llm(prompt):
     url = "https://openrouter.ai/api/v1/chat/completions"
     headers = {
-        "Authorization": f"Bearer {api_key}",
+        "Authorization": f"Bearer {os.getenv('OPENROUTER_API_KEY')}",
         "HTTP-Referer": "https://huggingface.co/",
         "X-Title": "Kubernetes RAG Assistant"
     }
-    payload = {
+    data = {
         "model": "meta-llama/llama-3.1-8b-instruct",
-        "messages": [
-            {"role": "system", "content": "You answer only from provided context."},
-            {"role": "user", "content": prompt}
-        ],
+        "messages": [{"role": "user", "content": prompt}],
+        "max_tokens": 400,
         "temperature": 0.0,
-        "max_tokens": 500
     }
-    resp = requests.post(url, headers=headers, json=payload, timeout=60)
-    data = resp.json()
-    if "choices" in data:
-        return data["choices"][0]["message"]["content"]
-    print("LLM error:", json.dumps(data, indent=2))
-    return "⚠ LLM error. Please try again."
-
-def answer_question(query: str, history):
-    # history is list of [user, bot]
-    history_str = ""
-    for u, b in history[-4:]:  # last 4 turns
-        history_str += f"User: {u}\nAssistant: {b}\n"
-    ctx, sources = build_context_with_citations(query)
-    prompt = build_prompt(query, ctx, history_str)
+    r = requests.post(url, headers=headers, json=data)
+    res = r.json()
+    if "choices" in res:
+        return res["choices"][0]["message"]["content"]
+    print("🚨 LLM ERROR:", res)
+    return "⚠️ Error: No response from model"
+
+# ------------------ BUILD ANSWER WITH CITATIONS ------------------ #
+
+def build_context_with_citations(query):
+    docs = retriever.invoke(query)
+    context = ""
+    sources = []
+    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']}")
+    return context, sources
+
+def answer_question(query, history):
+    context, sources = build_context_with_citations(query)
+    prompt = f"""
+Answer the question strictly using the context below.
+Every sentence must include citation like [1], [2].
+If missing info → reply: "Not in docs."
+
+Question: {query}
+
+Context:
+{context}
+"""
     answer = call_llm(prompt)
-    return answer, sources
+    src = "\n".join(sources) if sources else "No sources available."
+    history.append((query, answer + "\n\n---\nSources:\n" + src))
+    return history, ""
 
-# ------------------ 5. GRADIO CHAT UI ------------------ #
+# ------------------ GRADIO UI ------------------ #
 
 custom_css = """
 .source-box {
-    font-size: 14px;
-    background: #111827;
+    background: #1e293b;
     padding: 10px;
     border-radius: 8px;
-    color: #d1e4ff;
-    border: 1px solid #2563eb;
+    color: #dbeafe;
+    border: 1px solid #3b82f6;
 }
 """
 
-def chat_fn(message, history):
-    answer, refs = answer_question(message, history)
-    src_lines = [f"{s['label']} – {s['url']}" for s in refs]
-    sources_text = "\n".join(src_lines) if src_lines else "No sources found."
-    full_answer = f"{answer}\n\n---\n**Sources**:\n{sources_text}"
-    history.append((message, answer))
-    return history, ""
+with gr.Blocks(theme="soft") as app:
+    gr.HTML(f"<style>{custom_css}</style>")
 
-with gr.Blocks(css=custom_css, theme="soft") as demo:
-    gr.HTML(
-        "<h1 style='text-align:center;color:#3b82f6;'>☸ Kubernetes RAG Assistant</h1>"
-        "<p style='text-align:center;color:#e5e7eb;'>Ask Kubernetes questions. "
-        "Answers are grounded in official docs and include citations.</p>"
-    )
+    gr.HTML("<h1 style='text-align:center;color:#3b82f6'>☸ Kubernetes RAG Assistant</h1>"
+            "<p style='text-align:center;color:#cbd5e1'>Ask Kubernetes questions — answers include official docs citations 📌</p>")
 
     chat = gr.Chatbot(label="Conversation", height=450)
-    msg = gr.Textbox(label="Your question", placeholder="e.g. What is a StatefulSet?")
+    msg = gr.Textbox(label="Ask a question...", placeholder="What is a pod?")
     clear = gr.Button("Clear Chat")
 
-    def respond(message, history):
-        return chat_fn(message, history)
-
-    msg.submit(respond, [msg, chat], [chat, msg])
+    msg.submit(answer_question, [msg, chat], [chat, msg])
     clear.click(lambda: ([], ""), None, [chat, msg])
 
-demo.launch()
+app.launch()