File size: 5,451 Bytes
8f55ee4 d3b134d b5ce09d 8f55ee4 | 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 161 162 163 164 165 166 167 168 169 170 171 172 173 | # =======================================
# 📘 RAG App Pro – Gemini + Smart Embeddings (Multi-User Safe)
# =======================================
import os, re, shutil, textwrap, requests, uuid
from bs4 import BeautifulSoup
import google.generativeai as genai
from sentence_transformers import SentenceTransformer
import chromadb
import gradio as gr
from langchain_community.document_loaders import UnstructuredPDFLoader
import camelot
# ======================
# 🔹 Gemini + Local Setup
# ======================
genai.configure(api_key="AQ.Ab8RN6I4za-ifPTOLhC78wesRWQxGaWeQQZZa44Cv7dBtaxX6A") # 🔑 apni Gemini key
MODEL = "gemini-3.5-flash"
embedder = SentenceTransformer("all-MiniLM-L6-v2")
chroma_client = chromadb.Client()
# ======================
# 🔹 Utils
# ======================
def clean_text(text):
return re.sub(r"\s+", " ", text).strip()
def adaptive_chunk_text(text):
length = len(text)
if length < 3000:
size = 500
elif length < 10000:
size = 1000
else:
size = 1500
chunks = []
for i in range(0, len(text), size - 150):
chunks.append(text[i:i + size])
return chunks
def extract_pdf_text(pdf_path):
"""Smart PDF extractor (tables + text)"""
full_text = ""
try:
tables = camelot.read_pdf(pdf_path, pages="all")
for i, table in enumerate(tables):
full_text += f"\n\n[Table {i+1}]\n" + table.df.to_string(index=False)
except Exception:
pass
try:
loader = UnstructuredPDFLoader(pdf_path)
docs = loader.load()
full_text += "\n\n".join([doc.page_content for doc in docs])
except Exception as e:
full_text += f"\n\n[Error extracting text: {e}]"
return clean_text(full_text)
# ======================
# 🔹 Session Handling
# ======================
def create_user_collection():
"""Each user/session gets unique collection"""
session_id = f"user_{str(uuid.uuid4())[:8]}"
collection = chroma_client.create_collection(name=session_id)
return session_id, collection
def reset_collection(collection_name):
"""Delete previous data for same user"""
try:
chroma_client.delete_collection(name=collection_name)
except Exception:
pass
return chroma_client.create_collection(name=collection_name)
# ======================
# 🔹 Ingestion Logic
# ======================
def ingest_source(source, from_url, collection_name):
# Delete previous user data
collection = reset_collection(collection_name)
if from_url:
html = requests.get(source, timeout=15).text
soup = BeautifulSoup(html, "html.parser")
text = clean_text(soup.get_text())
else:
text = extract_pdf_text(source)
if not text.strip():
return "⚠️ No readable text found (maybe image-only PDF)."
chunks = adaptive_chunk_text(text)
embeddings = embedder.encode(chunks).tolist()
for i, emb in enumerate(embeddings):
collection.add(ids=[f"{collection_name}_{i}"], embeddings=[emb], documents=[chunks[i]])
return f"✅ [{collection_name}] Ingested {len(chunks)} chunks successfully!"
# ======================
# 🔹 Query Logic
# ======================
def rag_query(query, collection_name):
try:
collection = chroma_client.get_collection(name=collection_name)
q_emb = embedder.encode([query]).tolist()
results = collection.query(query_embeddings=q_emb, n_results=4)
if not results["documents"]:
return "⚠️ No context found. Try ingesting data first."
context = "\n\n".join(results["documents"][0])
prompt = f"""
You are a knowledgeable AI assistant.
Use the context below to answer clearly and in multiple lines.
Context:
{context}
Question: {query}
Answer:
"""
response = genai.GenerativeModel(MODEL).generate_content(prompt)
ans = response.text.replace(". ", ".\n")
return ans
except Exception as e:
return f"⚠️ Error: {e}"
# ======================
# 🔹 Gradio UI
# ======================
def start_new_session():
session_id, _ = create_user_collection()
return session_id
session_id = start_new_session()
def ingest_website(url):
return ingest_source(url, True, session_id)
def ingest_pdf(file):
return ingest_source(file.name, False, session_id)
def query_ask(q):
return rag_query(q, session_id)
with gr.Blocks(theme=gr.themes.Soft(primary_hue="emerald")) as demo:
gr.Markdown("# 🤖 Smart RAG App Pro (Gemini + Adaptive PDF + Multi-User Mode)")
gr.Markdown(f"🆕 **Private Session ID:** `{session_id}` – Your data is isolated and auto-clears on refresh.")
with gr.Tab("🌐 Ingest Website"):
url_in = gr.Textbox(label="Enter Website URL")
url_btn = gr.Button("Ingest Website")
url_out = gr.Textbox(label="Status")
url_btn.click(fn=ingest_website, inputs=url_in, outputs=url_out)
with gr.Tab("📄 Ingest PDF"):
pdf_in = gr.File(label="Upload PDF")
pdf_btn = gr.Button("Ingest PDF")
pdf_out = gr.Textbox(label="Status")
pdf_btn.click(fn=ingest_pdf, inputs=pdf_in, outputs=pdf_out)
with gr.Tab("💬 Ask Questions"):
q_in = gr.Textbox(label="Ask anything from ingested data")
q_btn = gr.Button("Ask Gemini")
q_out = gr.Markdown(label="Answer")
q_btn.click(fn=query_ask, inputs=q_in, outputs=q_out)
demo.launch()
|