Upload 3 files

README.md

@@ -1,26 +1,20 @@
-# LangGraph Multi-Agent: Gemini + Salary Prediction
+---
+title: Gemini LangGraph RAG
+emoji: 🤖
+colorFrom: blue
+colorTo: green
+sdk: gradio
+sdk_version: "5.0.0"
+python_version: "3.10"
+app_file: app.py
+pinned: false
+---
 
-This Hugging Face Space replaces DuckDuckGo with Gemini.
+# Gemini LangGraph RAG
 
-## Files
-- `app.py` - Gradio + LangGraph application
-- `salary_model.joblib` - uploaded trained salary model
-- `requirements.txt` - dependencies
+Upload PDF, DOCX, TXT, MD, or CSV files, build a temporary Chroma knowledge base, and ask RAG questions using Gemini and LangGraph.
 
-## Hugging Face Setup
-1. Create a new Hugging Face Space.
-2. Select **Gradio** as the SDK.
-3. Upload all files in this folder.
-4. Go to **Settings > Repository secrets**.
-5. Add:
+## Secret Required
+Add this Hugging Face Space secret:
 
-```text
-GOOGLE_API_KEY = your_gemini_api_key
-```
-
-## Run locally
-
-```bash
-pip install -r requirements.txt
-python app.py
-```
+`GOOGLE_API_KEY = your Gemini API key`

app.py

@@ -1,291 +1,159 @@
 import os
-import re
-import datetime
-from typing import TypedDict
+import tempfile
+from typing import TypedDict, List
 
 import gradio as gr
-import joblib
-import pandas as pd
-
-from langgraph.graph import StateGraph, END
 from langchain_google_genai import ChatGoogleGenerativeAI
+from langchain_huggingface import HuggingFaceEmbeddings
+from langchain_chroma import Chroma
+from langchain_text_splitters import RecursiveCharacterTextSplitter
+from langchain_core.documents import Document
+from langgraph.graph import StateGraph, END
 
+# Optional readers
+from pypdf import PdfReader
+import docx
 
-# =========================
-# 1. Load Gemini
-# =========================
-GOOGLE_API_KEY = os.environ.get("GOOGLE_API_KEY")
-
-if GOOGLE_API_KEY:
-    llm = ChatGoogleGenerativeAI(
-        model="gemini-2.5-flash",
-        google_api_key=GOOGLE_API_KEY,
-        temperature=0.2
-    )
-else:
-    llm = None
-
-
-# =========================
-# 2. Load salary model
-# =========================
-MODEL_PATH = "salary_model.joblib"
-
-salary_model_artifact = joblib.load(MODEL_PATH)
-
-if isinstance(salary_model_artifact, dict):
-    salary_model = salary_model_artifact.get("model")
-    salary_feature = salary_model_artifact.get("feature", "Age")
-    salary_target = salary_model_artifact.get("target", "Salary")
-else:
-    salary_model = salary_model_artifact
-    salary_feature = "Age"
-    salary_target = "Salary"
-
-
-# =========================
-# 3. Simple internal knowledge base
-#    Replace this with your real documents later.
-# =========================
-PRODUCT_KB = """
-Apple Store Product Notes:
-- iPhone 16 is available in 128GB, 256GB, and 512GB versions.
-- MacBook Air is designed for light and portable productivity.
-- AirPods Pro supports active noise cancellation.
-- Promotional bundles may be offered during festive periods.
-"""
-
-POLICY_KB = """
-Store Policy Notes:
-- Standard return period is 14 days if the item is unused and in original packaging.
-- Warranty coverage depends on the product category.
-- Shipping time depends on location and stock availability.
-- Refunds are processed after inspection.
-"""
-
-TECH_KB = """
-Technical Support Notes:
-- If a phone overheats, close background apps, update software, and avoid charging during heavy use.
-- If Wi-Fi is unstable, restart the router and device, then forget and reconnect to the network.
-- If battery drains quickly, check battery health and disable unnecessary background refresh.
-"""
-
-
-# =========================
-# 4. Helper functions
-# =========================
-def call_gemini(prompt: str) -> str:
-    if llm is None:
-        return (
-            "Gemini API key is missing. In Hugging Face Spaces, add GOOGLE_API_KEY "
-            "under Settings > Repository secrets."
-        )
-
-    try:
-        return llm.invoke(prompt).content
-    except Exception as e:
-        return f"Gemini error: {e}"
-
+EMBEDDING_MODEL = "sentence-transformers/all-MiniLM-L6-v2"
 
-def predict_salary_from_text(query: str) -> str:
-    try:
-        dob_match = re.search(r"\b(\d{4}-\d{2}-\d{2})\b", query)
+llm = ChatGoogleGenerativeAI(
+    model="gemini-2.5-flash",
+    google_api_key=os.getenv("GOOGLE_API_KEY"),
+    temperature=0.2,
+)
+embeddings = HuggingFaceEmbeddings(model_name=EMBEDDING_MODEL)
 
-        if dob_match:
-            dob = pd.to_datetime(dob_match.group(1), errors="coerce")
-            if pd.isna(dob):
-                return "Could not parse DOB. Please use YYYY-MM-DD format."
+VECTORSTORE = None
+RETRIEVER = None
 
-            today = pd.Timestamp.today()
-            age = int((today - dob).days // 365)
-        else:
-            age_match = re.search(r"\b(\d{1,3})\b", query)
-            if not age_match:
-                return "Please provide an age or DOB, for example: 'Predict salary for age 30'."
-            age = int(age_match.group(1))
+class AgentState(TypedDict):
+    question: str
+    context: str
+    answer: str
+    route: str
+
+
+def read_file(path: str) -> str:
+    ext = os.path.splitext(path)[1].lower()
+    if ext == ".pdf":
+        reader = PdfReader(path)
+        return "\n".join(page.extract_text() or "" for page in reader.pages)
+    if ext == ".docx":
+        d = docx.Document(path)
+        return "\n".join(p.text for p in d.paragraphs)
+    if ext in [".txt", ".md", ".csv"]:
+        with open(path, "r", encoding="utf-8", errors="ignore") as f:
+            return f.read()
+    return ""
+
+
+def build_rag(files):
+    global VECTORSTORE, RETRIEVER
+    if not files:
+        return "Please upload at least one PDF, DOCX, TXT, MD, or CSV file."
+
+    docs: List[Document] = []
+    for f in files:
+        text = read_file(f.name)
+        if text.strip():
+            docs.append(Document(page_content=text, metadata={"source": os.path.basename(f.name)}))
+
+    if not docs:
+        return "No readable text found. Please upload PDF, DOCX, TXT, MD, or CSV files."
+
+    splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
+    chunks = splitter.split_documents(docs)
+
+    persist_dir = tempfile.mkdtemp(prefix="chroma_")
+    VECTORSTORE = Chroma.from_documents(
+        documents=chunks,
+        embedding=embeddings,
+        persist_directory=persist_dir,
+        collection_name="rag_collection",
+    )
+    RETRIEVER = VECTORSTORE.as_retriever(search_kwargs={"k": 4})
+    return f"Knowledge base ready. Loaded {len(docs)} file(s) and created {len(chunks)} chunks."
 
-        prediction = salary_model.predict([[age]])
 
-        # Handles both scalar and nested output
-        predicted_salary = prediction[0]
-        if hasattr(predicted_salary, "__len__"):
-            predicted_salary = predicted_salary[0]
+def retrieve_node(state: AgentState):
+    global RETRIEVER
+    if RETRIEVER is None:
+        return {"context": "", "route": "general"}
+    docs = RETRIEVER.invoke(state["question"])
+    context = "\n\n".join(
+        [f"Source: {d.metadata.get('source', 'uploaded file')}\n{d.page_content}" for d in docs]
+    )
+    return {"context": context, "route": "rag" if context.strip() else "general"}
 
-        return (
-            f"Estimated {salary_target}: {float(predicted_salary):,.2f}\n\n"
-            f"Input used: {salary_feature} = {age}\n"
-            f"Note: This is a model estimate, not a guaranteed salary."
-        )
 
-    except Exception as e:
-        return f"Salary prediction error: {e}"
+def router_node(state: AgentState):
+    return state.get("route", "general")
 
 
-def simple_retrieve(query: str, kb: str) -> str:
-    # For teaching/demo: pass the relevant internal notes directly to Gemini.
+def rag_answer_node(state: AgentState):
     prompt = f"""
-Use the internal knowledge below to answer the user question.
-If the answer is not available in the internal knowledge, say that clearly.
+You are a helpful RAG assistant. Answer using ONLY the context below.
+If the answer is not in the context, say: "I cannot find this in the uploaded documents."
 
-Internal knowledge:
-{kb}
+Context:
+{state['context']}
 
-User question:
-{query}
-
-Answer concisely:
+Question:
+{state['question']}
 """
-    return call_gemini(prompt)
+    return {"answer": llm.invoke(prompt).content}
 
 
-# =========================
-# 5. LangGraph State
-# =========================
-class AgentState(TypedDict):
-    query: str
-    response: str
-    next_node: str
-    debug_log: str
-
+def general_answer_node(state: AgentState):
+    prompt = f"""
+The user has not built a usable knowledge base yet, or no relevant context was found.
+Answer briefly and remind the user to upload documents for RAG-based answers.
 
-# =========================
-# 6. Router Node
-# =========================
-def router_node(state: AgentState):
-    query = state["query"].lower()
-
-    if any(word in query for word in ["salary", "income", "pay", "wage", "dob", "birth", "age"]):
-        next_node = "salary_agent"
-    elif any(word in query for word in ["return", "refund", "warranty", "shipping", "policy"]):
-        next_node = "policy_agent"
-    elif any(word in query for word in ["fix", "troubleshoot", "overheat", "wifi", "battery", "technical"]):
-        next_node = "tech_agent"
-    elif any(word in query for word in ["product", "iphone", "macbook", "airpods", "available", "buy"]):
-        next_node = "product_agent"
-    else:
-        next_node = "general_agent"
-
-    return {
-        "next_node": next_node,
-        "debug_log": f"Router selected: {next_node}"
-    }
-
-
-# =========================
-# 7. Specialist Nodes
-# =========================
-def product_agent_node(state: AgentState):
-    answer = simple_retrieve(state["query"], PRODUCT_KB)
-    return {
-        "response": answer,
-        "debug_log": state["debug_log"] + "\nProduct Agent used internal product notes + Gemini."
-    }
-
-
-def policy_agent_node(state: AgentState):
-    answer = simple_retrieve(state["query"], POLICY_KB)
-    return {
-        "response": answer,
-        "debug_log": state["debug_log"] + "\nPolicy Agent used internal policy notes + Gemini."
-    }
-
-
-def tech_agent_node(state: AgentState):
-    answer = simple_retrieve(state["query"], TECH_KB)
-    return {
-        "response": answer,
-        "debug_log": state["debug_log"] + "\nTech Agent used internal technical notes + Gemini."
-    }
-
-
-def salary_agent_node(state: AgentState):
-    answer = predict_salary_from_text(state["query"])
-    return {
-        "response": answer,
-        "debug_log": state["debug_log"] + "\nSalary Agent used salary_model.joblib."
-    }
-
-
-def general_agent_node(state: AgentState):
-    answer = call_gemini(state["query"])
-    return {
-        "response": answer,
-        "debug_log": state["debug_log"] + "\nGeneral Agent used Gemini."
-    }
-
-
-# =========================
-# 8. Build LangGraph
-# =========================
-workflow = StateGraph(AgentState)
-
-workflow.add_node("router", router_node)
-workflow.add_node("product_agent", product_agent_node)
-workflow.add_node("policy_agent", policy_agent_node)
-workflow.add_node("tech_agent", tech_agent_node)
-workflow.add_node("salary_agent", salary_agent_node)
-workflow.add_node("general_agent", general_agent_node)
-
-workflow.set_entry_point("router")
-
-workflow.add_conditional_edges(
-    "router",
-    lambda state: state["next_node"],
-    {
-        "product_agent": "product_agent",
-        "policy_agent": "policy_agent",
-        "tech_agent": "tech_agent",
-        "salary_agent": "salary_agent",
-        "general_agent": "general_agent",
-    }
+Question: {state['question']}
+"""
+    return {"answer": llm.invoke(prompt).content}
+
+
+graph = StateGraph(AgentState)
+graph.add_node("retrieve", retrieve_node)
+graph.add_node("rag_answer", rag_answer_node)
+graph.add_node("general_answer", general_answer_node)
+graph.set_entry_point("retrieve")
+graph.add_conditional_edges(
+    "retrieve",
+    router_node,
+    {"rag": "rag_answer", "general": "general_answer"},
 )
+graph.add_edge("rag_answer", END)
+graph.add_edge("general_answer", END)
+app_graph = graph.compile()
+
+
+def ask_rag(question):
+    if not question or not question.strip():
+        return "Please enter a question."
+    result = app_graph.invoke({"question": question, "context": "", "answer": "", "route": ""})
+    return result["answer"]
+
+with gr.Blocks(title="Gemini LangGraph RAG") as demo:
+    gr.Markdown("# Gemini LangGraph RAG on Hugging Face")
+    gr.Markdown("Upload documents, build a temporary Chroma knowledge base, then ask questions.")
+
+    with gr.Row():
+        files = gr.File(
+            label="Upload documents",
+            file_count="multiple",
+            file_types=[".pdf", ".docx", ".txt", ".md", ".csv"],
+        )
+        status = gr.Textbox(label="Status")
 
-workflow.add_edge("product_agent", END)
-workflow.add_edge("policy_agent", END)
-workflow.add_edge("tech_agent", END)
-workflow.add_edge("salary_agent", END)
-workflow.add_edge("general_agent", END)
-
-graph_app = workflow.compile()
-
-
-# =========================
-# 9. Gradio UI
-# =========================
-def chat_fn(message, history):
-    try:
-        result = graph_app.invoke({
-            "query": message,
-            "response": "",
-            "next_node": "",
-            "debug_log": ""
-        })
-
-        answer = result.get("response", "No answer returned.")
-        log = result.get("debug_log", "No debug log.")
-
-        return f"{answer}\n\n---\nRouting trace:\n{log}"
-
-    except Exception as e:
-        return f"App error: {e}"
-
-
-demo = gr.ChatInterface(
-    fn=chat_fn,
-    title="LangGraph Multi-Agent: Gemini + Salary Prediction",
-    description=(
-        "Ask about products, policies, technical support, general questions, "
-        "or salary prediction using age/DOB. DuckDuckGo has been replaced with Gemini."
-    ),
-    examples=[
-        "What products are available?",
-        "What is the return policy?",
-        "How do I fix an overheating iPhone?",
-        "Predict salary for age 30",
-        "Predict salary for DOB 1995-05-20",
-        "Explain agentic AI simply."
-    ]
-)
+    build_btn = gr.Button("Build Knowledge Base")
+    build_btn.click(build_rag, inputs=files, outputs=status)
+
+    question = gr.Textbox(label="Question", placeholder="Ask about your uploaded documents")
+    answer = gr.Textbox(label="Answer", lines=10)
+    ask_btn = gr.Button("Ask RAG")
+    ask_btn.click(ask_rag, inputs=question, outputs=answer)
 
 if __name__ == "__main__":
     demo.launch(server_name="0.0.0.0", server_port=7860)

requirements.txt

@@ -2,7 +2,9 @@ gradio
 langgraph
 langchain
 langchain-google-genai
-google-generativeai
-scikit-learn==1.6.1
-joblib
-pandas
+langchain-huggingface
+langchain-chroma
+chromadb
+sentence-transformers
+pypdf
+python-docx