File size: 4,978 Bytes
649fda5
a3e86ec
 
649fda5
 
 
a3e86ec
 
 
 
 
649fda5
a3e86ec
 
 
649fda5
a3e86ec
649fda5
a3e86ec
 
 
 
 
 
649fda5
a3e86ec
 
649fda5
a3e86ec
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
649fda5
 
a3e86ec
 
 
 
 
 
 
 
 
649fda5
 
a3e86ec
 
649fda5
 
a3e86ec
649fda5
a3e86ec
 
649fda5
a3e86ec
 
649fda5
a3e86ec
 
649fda5
a3e86ec
649fda5
 
a3e86ec
 
 
 
649fda5
a3e86ec
 
 
 
 
 
 
 
 
 
 
 
 
 
649fda5
a3e86ec
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
649fda5
a3e86ec
 
 
 
 
 
 
649fda5
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
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)