Spaces:
Runtime error
Runtime error
| import os | |
| import tempfile | |
| from typing import TypedDict, List | |
| import gradio as gr | |
| 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 | |
| EMBEDDING_MODEL = "sentence-transformers/all-MiniLM-L6-v2" | |
| llm = ChatGoogleGenerativeAI( | |
| model="gemini-2.5-flash", | |
| google_api_key=os.getenv("GOOGLE_API_KEY"), | |
| temperature=0.2, | |
| ) | |
| embeddings = HuggingFaceEmbeddings(model_name=EMBEDDING_MODEL) | |
| VECTORSTORE = None | |
| RETRIEVER = None | |
| 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." | |
| 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"} | |
| def router_node(state: AgentState): | |
| return state.get("route", "general") | |
| def rag_answer_node(state: AgentState): | |
| prompt = f""" | |
| 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." | |
| Context: | |
| {state['context']} | |
| Question: | |
| {state['question']} | |
| """ | |
| return {"answer": llm.invoke(prompt).content} | |
| 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. | |
| 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") | |
| 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) | |