Spaces:
Sleeping
Sleeping
| 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 | |
| # ----------------------- | |
| # 1. SCRAPE K8S DOCS | |
| # ----------------------- | |
| 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/", | |
| "namespaces": "https://kubernetes.io/docs/concepts/overview/working-with-objects/namespaces/", | |
| "nodes": "https://kubernetes.io/docs/concepts/architecture/nodes/", | |
| "statefulsets": "https://kubernetes.io/docs/concepts/workloads/controllers/statefulset/", | |
| "rbac": "https://kubernetes.io/docs/reference/access-authn-authz/rbac/", | |
| "persistent-volumes": "https://kubernetes.io/docs/concepts/storage/persistent-volumes/", | |
| "ingress": "https://kubernetes.io/docs/concepts/services-networking/ingress/", | |
| "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." | |
| url = "https://openrouter.ai/api/v1/chat/completions" | |
| headers = { | |
| "Authorization": f"Bearer {OPENROUTER_API_KEY}", | |
| "Content-Type": "application/json" | |
| } | |
| 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 | |
| } | |
| 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) | |
| 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() | |