{
 "nbformat": 4,
 "nbformat_minor": 0,
 "metadata": {
  "colab": {
   "provenance": [],
   "gpuType": "T4",
   "name": "StudyOcean_RAG_Colab.ipynb"
  },
  "kernelspec": {
   "name": "python3",
   "display_name": "Python 3"
  },
  "language_info": {
   "name": "python"
  },
  "accelerator": "GPU"
 },
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 🎓 Study Ocean AI Mentor — Google Colab Pipeline\n",
    "\n",
    "> **Run this notebook on Colab with T4 GPU**  \n",
    "> Runtime → Change runtime type → **T4 GPU**\n",
    "\n",
    "### What this notebook does:\n",
    "1. ✅ Mounts your Google Drive\n",
    "2. ✅ Installs all dependencies (with `faiss-gpu` for T4 acceleration)\n",
    "3. ✅ Extracts PDFs from `StudyOceanPDFs.zip`\n",
    "4. ✅ Runs the ingestion pipeline (GPU-accelerated embeddings)\n",
    "5. ✅ Saves the FAISS vector store back to Google Drive\n",
    "6. ✅ Launches the Gradio chat UI with a public shareable link\n",
    "\n",
    "---\n",
    "### 📋 Before you start:\n",
    "- Upload `StudyOceanPDFs.zip` to your **Google Drive** (e.g., `MyDrive/StudyOcean/StudyOceanPDFs.zip`)\n",
    "- Get your **Google Gemini API key** from https://aistudio.google.com/apikey\n",
    "- Set runtime to **T4 GPU**: Runtime → Change runtime type → T4 GPU\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 🔧 Step 1 — Verify GPU & Mount Google Drive"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ── Verify T4 GPU is available ───────────────────────────────\n",
    "import subprocess\n",
    "result = subprocess.run(['nvidia-smi'], capture_output=True, text=True)\n",
    "print(result.stdout if result.returncode == 0 else \"⚠️ No GPU found! Go to Runtime → Change runtime type → T4 GPU\")\n",
    "\n",
    "import torch\n",
    "print(f\"\\n🔥 PyTorch CUDA available: {torch.cuda.is_available()}\")\n",
    "if torch.cuda.is_available():\n",
    "    print(f\"   GPU: {torch.cuda.get_device_name(0)}\")\n",
    "    print(f\"   VRAM: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ── Mount Google Drive ────────────────────────────────────────\n",
    "from google.colab import drive\n",
    "drive.mount('/content/drive')\n",
    "print(\"✅ Google Drive mounted at /content/drive\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 📦 Step 2 — Install Dependencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ── Install all packages ──────────────────────────────────────\n",
    "# Using faiss-gpu for T4 GPU acceleration during similarity search\n",
    "!pip install -q \\\n",
    "    langchain \\\n",
    "    langchain-community \\\n",
    "    langchain-core \\\n",
    "    langchain-google-genai \\\n",
    "    langchain-huggingface \\\n",
    "    pymupdf \\\n",
    "    faiss-gpu \\\n",
    "    sentence-transformers \\\n",
    "    google-generativeai \\\n",
    "    python-dotenv \\\n",
    "    gradio \\\n",
    "    tiktoken \\\n",
    "    accelerate\n",
    "\n",
    "print(\"\\n✅ All packages installed!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## ⚙️ Step 3 — Configure Settings\n",
    "\n",
    "> **Edit the variables in this cell before running!**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "# ════════════════════════════════════════════════════════\n",
    "#              ✏️  EDIT THESE SETTINGS\n",
    "# ════════════════════════════════════════════════════════\n",
    "\n",
    "# Your Google Gemini API key — get it from https://aistudio.google.com/apikey\n",
    "GOOGLE_API_KEY = \"\"  # ← PASTE YOUR KEY HERE\n",
    "\n",
    "# Path to StudyOceanPDFs.zip in your Google Drive\n",
    "ZIP_DRIVE_PATH = \"/content/drive/MyDrive/StudyOcean/StudyOceanPDFs.zip\"\n",
    "\n",
    "# Where to save the vector store in Google Drive (so it persists!)\n",
    "VECTOR_STORE_DRIVE_PATH = \"/content/drive/MyDrive/StudyOcean/vector_store\"\n",
    "\n",
    "# Embedding model — 'cuda' will use the T4 GPU!\n",
    "EMBEDDING_MODEL = \"all-MiniLM-L6-v2\"   # Fast & accurate\n",
    "EMBEDDING_DEVICE = \"cuda\"               # ← T4 GPU acceleration\n",
    "\n",
    "# Gemini LLM model\n",
    "LLM_MODEL = \"gemini-2.0-flash\"\n",
    "\n",
    "# Text splitting config\n",
    "CHUNK_SIZE    = 1000\n",
    "CHUNK_OVERLAP = 200\n",
    "\n",
    "# RAG retrieval\n",
    "TOP_K = 5\n",
    "\n",
    "# ════════════════════════════════════════════════════════\n",
    "\n",
    "# Local working paths (inside Colab VM)\n",
    "PDF_DIR          = \"/content/studyocean_pdfs\"\n",
    "VECTOR_STORE_LOCAL = \"/content/vector_store\"\n",
    "\n",
    "os.makedirs(PDF_DIR, exist_ok=True)\n",
    "os.makedirs(VECTOR_STORE_LOCAL, exist_ok=True)\n",
    "\n",
    "# Set the API key as env var for LangChain\n",
    "os.environ[\"GOOGLE_API_KEY\"] = GOOGLE_API_KEY\n",
    "\n",
    "print(\"✅ Configuration set!\")\n",
    "print(f\"   ZIP source:    {ZIP_DRIVE_PATH}\")\n",
    "print(f\"   PDF extract:   {PDF_DIR}\")\n",
    "print(f\"   Vector store:  {VECTOR_STORE_LOCAL}\")\n",
    "print(f\"   Drive backup:  {VECTOR_STORE_DRIVE_PATH}\")\n",
    "print(f\"   Embedding GPU: {EMBEDDING_DEVICE}\")\n",
    "print(f\"   LLM:           {LLM_MODEL}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 📂 Step 4 — Extract PDFs from ZIP"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import zipfile\n",
    "import time\n",
    "\n",
    "# Check if zip exists\n",
    "if not os.path.exists(ZIP_DRIVE_PATH):\n",
    "    print(f\"❌ ZIP not found at: {ZIP_DRIVE_PATH}\")\n",
    "    print(\"   Please upload StudyOceanPDFs.zip to your Google Drive\")\n",
    "    print(\"   Then update ZIP_DRIVE_PATH in Step 3.\")\n",
    "    raise FileNotFoundError(f\"ZIP not found: {ZIP_DRIVE_PATH}\")\n",
    "\n",
    "zip_size_gb = os.path.getsize(ZIP_DRIVE_PATH) / 1e9\n",
    "print(f\"📦 Found ZIP: {zip_size_gb:.2f} GB\")\n",
    "print(f\"   Extracting to {PDF_DIR} ...\")\n",
    "print(\"   (This may take a few minutes for large archives)\")\n",
    "\n",
    "start = time.time()\n",
    "with zipfile.ZipFile(ZIP_DRIVE_PATH, 'r') as zf:\n",
    "    zf.extractall(PDF_DIR)\n",
    "elapsed = time.time() - start\n",
    "\n",
    "# Count extracted PDFs\n",
    "pdf_count = sum(\n",
    "    1 for root, _, files in os.walk(PDF_DIR)\n",
    "    for f in files if f.lower().endswith('.pdf')\n",
    ")\n",
    "\n",
    "print(f\"\\n✅ Extraction complete in {elapsed:.1f}s\")\n",
    "print(f\"   Found {pdf_count} PDF files\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 🔍 Step 5 — PDF Ingestion Pipeline (GPU-Accelerated)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import fitz  # PyMuPDF\n",
    "from langchain.text_splitter import RecursiveCharacterTextSplitter\n",
    "from langchain_community.vectorstores import FAISS\n",
    "from langchain_huggingface import HuggingFaceEmbeddings\n",
    "from langchain.docstore.document import Document\n",
    "\n",
    "\n",
    "def collect_pdf_paths(pdf_dir):\n",
    "    \"\"\"Recursively collect all PDF paths.\"\"\"\n",
    "    paths = []\n",
    "    for root, _, files in os.walk(pdf_dir):\n",
    "        for f in files:\n",
    "            if f.lower().endswith('.pdf'):\n",
    "                paths.append(os.path.join(root, f))\n",
    "    return sorted(paths)\n",
    "\n",
    "\n",
    "def extract_text(pdf_path):\n",
    "    \"\"\"Extract text from a PDF using PyMuPDF.\"\"\"\n",
    "    text = \"\"\n",
    "    try:\n",
    "        doc = fitz.open(pdf_path)\n",
    "        for page in doc:\n",
    "            text += page.get_text()\n",
    "        doc.close()\n",
    "    except Exception as e:\n",
    "        print(f\"  ⚠ Error reading {os.path.basename(pdf_path)}: {e}\")\n",
    "    return text\n",
    "\n",
    "\n",
    "def build_documents(pdf_paths):\n",
    "    \"\"\"Extract and chunk all PDFs into LangChain Documents.\"\"\"\n",
    "    splitter = RecursiveCharacterTextSplitter(\n",
    "        chunk_size=CHUNK_SIZE,\n",
    "        chunk_overlap=CHUNK_OVERLAP,\n",
    "        separators=[\"\\n\\n\", \"\\n\", \". \", \" \", \"\"],\n",
    "    )\n",
    "    all_docs = []\n",
    "    skipped = 0\n",
    "\n",
    "    for i, path in enumerate(pdf_paths, 1):\n",
    "        rel = os.path.relpath(path, PDF_DIR)\n",
    "        parts = rel.replace(\"\\\\\", \"/\").split(\"/\")\n",
    "        subject = parts[0] if len(parts) > 1 else \"General\"\n",
    "\n",
    "        if i % 50 == 0 or i == 1:\n",
    "            print(f\"  [{i}/{len(pdf_paths)}] {rel}\")\n",
    "\n",
    "        text = extract_text(path)\n",
    "        if not text.strip():\n",
    "            skipped += 1\n",
    "            continue\n",
    "\n",
    "        chunks = splitter.create_documents(\n",
    "            texts=[text],\n",
    "            metadatas=[{\"source\": rel, \"filename\": os.path.basename(path), \"subject\": subject}]\n",
    "        )\n",
    "        all_docs.extend(chunks)\n",
    "\n",
    "    print(f\"\\n  Skipped (image/scanned PDFs): {skipped}\")\n",
    "    return all_docs\n",
    "\n",
    "\n",
    "# ── Step 5a: Collect PDFs ────────────────────────────────────\n",
    "print(\"=\" * 60)\n",
    "print(\"📄 PHASE 1: Collecting & chunking PDFs\")\n",
    "print(\"=\" * 60)\n",
    "pdf_paths = collect_pdf_paths(PDF_DIR)\n",
    "print(f\"   {len(pdf_paths)} PDFs found\")\n",
    "\n",
    "t0 = time.time()\n",
    "documents = build_documents(pdf_paths)\n",
    "t_extract = time.time() - t0\n",
    "\n",
    "print(f\"\\n✅ Created {len(documents)} chunks in {t_extract:.1f}s\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ── Step 5b: Load Embedding Model onto T4 GPU ────────────────\n",
    "print(\"=\" * 60)\n",
    "print(f\"🧠 PHASE 2: Loading embedding model on {EMBEDDING_DEVICE.upper()}\")\n",
    "print(\"=\" * 60)\n",
    "\n",
    "import torch\n",
    "device = EMBEDDING_DEVICE if torch.cuda.is_available() else \"cpu\"\n",
    "if device == \"cpu\" and EMBEDDING_DEVICE == \"cuda\":\n",
    "    print(\"⚠️  CUDA not available, falling back to CPU\")\n",
    "\n",
    "embeddings = HuggingFaceEmbeddings(\n",
    "    model_name=EMBEDDING_MODEL,\n",
    "    model_kwargs={\"device\": device},\n",
    "    encode_kwargs={\"normalize_embeddings\": True, \"batch_size\": 256},  # Larger batch on GPU!\n",
    ")\n",
    "\n",
    "print(f\"✅ Embedding model loaded on {device.upper()}\")\n",
    "print(f\"   Model: {EMBEDDING_MODEL}\")\n",
    "print(f\"   Batch size: 256 (optimized for T4)\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ── Step 5c: Build FAISS Vector Store ───────────────────────\n",
    "print(\"=\" * 60)\n",
    "print(f\"📊 PHASE 3: Building FAISS vector store ({len(documents)} chunks)\")\n",
    "print(\"=\" * 60)\n",
    "\n",
    "BATCH_SIZE = 1000  # Larger batches since we have GPU\n",
    "vector_store = None\n",
    "t0 = time.time()\n",
    "\n",
    "total_batches = (len(documents) + BATCH_SIZE - 1) // BATCH_SIZE\n",
    "for i in range(0, len(documents), BATCH_SIZE):\n",
    "    batch = documents[i : i + BATCH_SIZE]\n",
    "    batch_num = (i // BATCH_SIZE) + 1\n",
    "\n",
    "    elapsed = time.time() - t0\n",
    "    rate = (i + 1) / elapsed if elapsed > 0 else 0\n",
    "    eta = (len(documents) - i) / rate if rate > 0 else 0\n",
    "    print(f\"  Batch {batch_num:3d}/{total_batches} | {len(batch)} chunks | \"\n",
    "          f\"ETA: {eta/60:.1f} min\")\n",
    "\n",
    "    if vector_store is None:\n",
    "        vector_store = FAISS.from_documents(batch, embeddings)\n",
    "    else:\n",
    "        batch_store = FAISS.from_documents(batch, embeddings)\n",
    "        vector_store.merge_from(batch_store)\n",
    "\n",
    "t_embed = time.time() - t0\n",
    "print(f\"\\n✅ Vector store built in {t_embed/60:.1f} minutes\")\n",
    "print(f\"   Total vectors: {vector_store.index.ntotal}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 💾 Step 6 — Save Vector Store to Google Drive"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import shutil\n",
    "\n",
    "# Save locally first\n",
    "print(\"💾 Saving vector store locally...\")\n",
    "vector_store.save_local(VECTOR_STORE_LOCAL)\n",
    "print(f\"   ✅ Saved to {VECTOR_STORE_LOCAL}\")\n",
    "\n",
    "# Copy to Google Drive for persistence\n",
    "print(\"\\n☁️  Copying to Google Drive...\")\n",
    "os.makedirs(VECTOR_STORE_DRIVE_PATH, exist_ok=True)\n",
    "\n",
    "for fname in os.listdir(VECTOR_STORE_LOCAL):\n",
    "    src = os.path.join(VECTOR_STORE_LOCAL, fname)\n",
    "    dst = os.path.join(VECTOR_STORE_DRIVE_PATH, fname)\n",
    "    shutil.copy2(src, dst)\n",
    "    size = os.path.getsize(dst) / 1e6\n",
    "    print(f\"   ✅ {fname}  ({size:.1f} MB)\")\n",
    "\n",
    "print(f\"\\n🎉 Vector store saved to Google Drive!\")\n",
    "print(f\"   Path: {VECTOR_STORE_DRIVE_PATH}\")\n",
    "print(f\"   Total vectors: {vector_store.index.ntotal}\")\n",
    "print(\"\\n📌 You can now download this folder to your laptop\")\n",
    "print(\"   and place it at: Study_Ocean Rag/data/vector_store/\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 🧪 Step 7 — Test the Vector Store (Quick Sanity Check)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ── Quick retrieval test ─────────────────────────────────────\n",
    "test_question = \"What is Newton's second law of motion?\"\n",
    "print(f\"🔍 Test query: '{test_question}'\\n\")\n",
    "\n",
    "results = vector_store.similarity_search(test_question, k=3)\n",
    "for i, doc in enumerate(results, 1):\n",
    "    print(f\"--- Result {i} ---\")\n",
    "    print(f\"Subject : {doc.metadata.get('subject', 'N/A')}\")\n",
    "    print(f\"File    : {doc.metadata.get('filename', 'N/A')}\")\n",
    "    print(f\"Preview : {doc.page_content[:200]}...\")\n",
    "    print()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 🤖 Step 8 — Build the RAG Chain (Gemini + FAISS)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_google_genai import ChatGoogleGenerativeAI\n",
    "from langchain.chains import RetrievalQA\n",
    "from langchain.prompts import PromptTemplate\n",
    "\n",
    "# Validate API key\n",
    "if not GOOGLE_API_KEY or GOOGLE_API_KEY == \"\":\n",
    "    print(\"❌ GOOGLE_API_KEY is not set! Go back to Step 3 and paste your key.\")\n",
    "    raise ValueError(\"Missing GOOGLE_API_KEY\")\n",
    "\n",
    "# ── Student Mentor Prompt ────────────────────────────────────\n",
    "MENTOR_PROMPT = PromptTemplate(\n",
    "    input_variables=[\"context\", \"question\"],\n",
    "    template=\"\"\"You are **Study Ocean AI Mentor** — a friendly, patient, and knowledgeable tutor who helps students understand their study material deeply.\n",
    "\n",
    "Your role:\n",
    "- Solve student doubts with clear, step-by-step explanations\n",
    "- Use simple language that students can easily understand\n",
    "- Provide examples and analogies when helpful\n",
    "- If the answer is in the context, use it; if not, say so honestly\n",
    "- Encourage students and make learning enjoyable\n",
    "\n",
    "**Guidelines:**\n",
    "1. Start with a brief, direct answer, then explain in detail\n",
    "2. Use bullet points, numbered steps, or tables for clarity\n",
    "3. For math/science problems, show the solution step by step\n",
    "4. Mention which topic/chapter the information relates to when possible\n",
    "5. If the topic isn't covered, say: \"This isn't in your study materials, but here's what I know...\"\n",
    "\n",
    "---\n",
    "📚 **Relevant Study Material:**\n",
    "{context}\n",
    "\n",
    "---\n",
    "🙋 **Student's Question:** {question}\n",
    "\n",
    "---\n",
    "🎓 **Mentor's Response:**\n",
    "\"\"\"\n",
    ")\n",
    "\n",
    "# ── Initialize LLM ───────────────────────────────────────────\n",
    "llm = ChatGoogleGenerativeAI(\n",
    "    model=LLM_MODEL,\n",
    "    google_api_key=GOOGLE_API_KEY,\n",
    "    temperature=0.3,\n",
    "    max_output_tokens=2048,\n",
    "    convert_system_message_to_human=True,\n",
    ")\n",
    "\n",
    "# ── Build RAG Chain ──────────────────────────────────────────\n",
    "retriever = vector_store.as_retriever(\n",
    "    search_type=\"similarity\",\n",
    "    search_kwargs={\"k\": TOP_K},\n",
    ")\n",
    "\n",
    "rag_chain = RetrievalQA.from_chain_type(\n",
    "    llm=llm,\n",
    "    chain_type=\"stuff\",\n",
    "    retriever=retriever,\n",
    "    return_source_documents=True,\n",
    "    chain_type_kwargs={\"prompt\": MENTOR_PROMPT},\n",
    ")\n",
    "\n",
    "print(\"✅ RAG Chain ready!\")\n",
    "print(f\"   LLM: {LLM_MODEL}\")\n",
    "print(f\"   Retriever: top-{TOP_K} chunks\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 🧪 Step 9 — Test the Full RAG Pipeline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def ask(question):\n",
    "    \"\"\"Ask the mentor a question.\"\"\"\n",
    "    print(f\"\\n🙋 Question: {question}\")\n",
    "    print(\"🔍 Searching study material...\\n\")\n",
    "    result = rag_chain.invoke({\"query\": question})\n",
    "    print(\"🎓 Mentor's Answer:\")\n",
    "    print(\"─\" * 60)\n",
    "    print(result[\"result\"])\n",
    "    print(\"─\" * 60)\n",
    "    print(\"\\n📖 Sources:\")\n",
    "    seen = set()\n",
    "    for doc in result.get(\"source_documents\", []):\n",
    "        key = doc.metadata.get(\"source\", \"?\")\n",
    "        if key not in seen:\n",
    "            seen.add(key)\n",
    "            print(f\"  [{doc.metadata.get('subject','?')}] {doc.metadata.get('filename','?')}\")\n",
    "\n",
    "# Test it!\n",
    "ask(\"Explain Newton's laws of motion with examples\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 🌐 Step 10 — Launch Gradio Chat UI (Public Link)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import gradio as gr\n",
    "\n",
    "CUSTOM_CSS = \"\"\"\n",
    ".gradio-container {\n",
    "    max-width: 900px !important;\n",
    "    margin: auto !important;\n",
    "    font-family: 'Inter', 'Segoe UI', sans-serif !important;\n",
    "}\n",
    ".header-section {\n",
    "    text-align: center;\n",
    "    padding: 20px 0 10px;\n",
    "    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);\n",
    "    border-radius: 16px;\n",
    "    margin-bottom: 20px;\n",
    "    color: white;\n",
    "}\n",
    ".header-section h1 { font-size: 2em; margin: 0; color: white !important; }\n",
    ".header-section p  { opacity: 0.9; margin: 5px 0 0; font-size: 1.05em; color: #e8e8ff !important; }\n",
    ".chatbot { border-radius: 12px !important; min-height: 450px !important; }\n",
    ".primary-btn {\n",
    "    background: linear-gradient(135deg, #667eea, #764ba2) !important;\n",
    "    border: none !important; color: white !important;\n",
    "    border-radius: 10px !important; font-weight: 600 !important;\n",
    "}\n",
    ".primary-btn:hover { opacity: 0.9 !important; transform: translateY(-1px); }\n",
    "\"\"\"\n",
    "\n",
    "def format_sources(docs):\n",
    "    if not docs: return \"\"\n",
    "    seen, lines = set(), [\"\\n\\n---\\n📖 **Sources Referenced:**\\n\"]\n",
    "    for doc in docs:\n",
    "        key = f\"{doc.metadata.get('subject','?')}/{doc.metadata.get('source','?')}\"\n",
    "        if key not in seen:\n",
    "            seen.add(key)\n",
    "            lines.append(f\"**•** `{doc.metadata.get('subject','?')}` → _{doc.metadata.get('filename','?')}_\")\n",
    "    return \"\\n\".join(lines)\n",
    "\n",
    "def respond(message, history):\n",
    "    if not message.strip():\n",
    "        return \"\", history\n",
    "    result = rag_chain.invoke({\"query\": message})\n",
    "    answer = result[\"result\"] + format_sources(result.get(\"source_documents\", []))\n",
    "    history.append((message, answer))\n",
    "    return \"\", history\n",
    "\n",
    "def clear_chat():\n",
    "    return [], \"\"\n",
    "\n",
    "with gr.Blocks(css=CUSTOM_CSS, title=\"Study Ocean AI Mentor\", theme=gr.themes.Soft()) as demo:\n",
    "    gr.HTML(\"\"\"\n",
    "    <div class=\"header-section\">\n",
    "        <h1>🎓 Study Ocean AI Mentor</h1>\n",
    "        <p>Ask any doubt from your study material — I'll explain it step by step!</p>\n",
    "    </div>\n",
    "    \"\"\")\n",
    "    chatbot = gr.Chatbot(\n",
    "        label=\"💬 Chat with Your Mentor\",\n",
    "        elem_classes=[\"chatbot\"],\n",
    "        height=480,\n",
    "        show_copy_button=True,\n",
    "        bubble_full_width=False,\n",
    "        avatar_images=(None, \"https://api.dicebear.com/7.x/bottts/svg?seed=mentor&backgroundColor=6366f1\"),\n",
    "    )\n",
    "    with gr.Row():\n",
    "        msg = gr.Textbox(\n",
    "            placeholder=\"Type your doubt here... e.g., 'Explain Newton's third law'\",\n",
    "            scale=5, lines=1, max_lines=3, show_label=False, container=False,\n",
    "        )\n",
    "        submit_btn = gr.Button(\"Ask 🚀\", elem_classes=[\"primary-btn\"], scale=1)\n",
    "    with gr.Row():\n",
    "        clear_btn = gr.Button(\"🗑️ Clear Chat\", size=\"sm\")\n",
    "\n",
    "    gr.Examples(\n",
    "        examples=[\n",
    "            \"What are Newton's laws of motion? Explain with examples.\",\n",
    "            \"Explain the difference between mitosis and meiosis.\",\n",
    "            \"What is the Pythagorean theorem and how do I apply it?\",\n",
    "            \"Explain the concept of supply and demand in economics.\",\n",
    "            \"What are the key features of the Indian Constitution?\",\n",
    "        ],\n",
    "        inputs=msg, label=\"💡 Try these example questions:\",\n",
    "    )\n",
    "\n",
    "    msg.submit(respond, [msg, chatbot], [msg, chatbot])\n",
    "    submit_btn.click(respond, [msg, chatbot], [msg, chatbot])\n",
    "    clear_btn.click(clear_chat, outputs=[chatbot, msg])\n",
    "\n",
    "# Launch with share=True to get a public Gradio URL!\n",
    "demo.launch(share=True, debug=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "## 💻 Step 11 — Use Pre-built Vector Store (Skip Re-ingestion)\n",
    "\n",
    "> **Run this cell INSTEAD of Steps 4–6 if you already have a saved vector store on Drive**\n",
    "> This lets you skip the ~1 hour ingestion on future sessions!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ── Load existing vector store from Drive ────────────────────\n",
    "# ⚠️ Only run this cell if you've already done ingestion before!\n",
    "# Skip this if running ingestion fresh (Steps 4-6).\n",
    "\n",
    "from langchain_community.vectorstores import FAISS\n",
    "from langchain_huggingface import HuggingFaceEmbeddings\n",
    "import torch, os\n",
    "\n",
    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "print(f\"🔧 Loading embeddings on {device.upper()}...\")\n",
    "\n",
    "embeddings = HuggingFaceEmbeddings(\n",
    "    model_name=EMBEDDING_MODEL,\n",
    "    model_kwargs={\"device\": device},\n",
    "    encode_kwargs={\"normalize_embeddings\": True},\n",
    ")\n",
    "\n",
    "# Try Drive path first, fall back to local\n",
    "store_path = VECTOR_STORE_DRIVE_PATH if os.path.exists(VECTOR_STORE_DRIVE_PATH) else VECTOR_STORE_LOCAL\n",
    "print(f\"📂 Loading vector store from: {store_path}\")\n",
    "\n",
    "vector_store = FAISS.load_local(\n",
    "    store_path,\n",
    "    embeddings,\n",
    "    allow_dangerous_deserialization=True,\n",
    ")\n",
    "\n",
    "print(f\"✅ Loaded! Total vectors: {vector_store.index.ntotal}\")\n",
    "print(\"   Now run Steps 8 → 10 to start chatting!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "## 📥 Step 12 — Download Vector Store to Your Laptop\n",
    "\n",
    "> Run this to zip and download the vector store directly to your PC (no Drive needed)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import shutil\n",
    "from google.colab import files\n",
    "\n",
    "# Zip the vector store\n",
    "print(\"📦 Zipping vector store...\")\n",
    "zip_path = \"/content/vector_store_export.zip\"\n",
    "shutil.make_archive(\"/content/vector_store_export\", 'zip', VECTOR_STORE_LOCAL)\n",
    "\n",
    "size_mb = os.path.getsize(zip_path) / 1e6\n",
    "print(f\"   ZIP size: {size_mb:.1f} MB\")\n",
    "\n",
    "# Download to local machine\n",
    "print(\"⬇️  Downloading to your computer...\")\n",
    "files.download(zip_path)\n",
    "print(\"\\n✅ Download started!\")\n",
    "print(\"   Extract to: Study_Ocean Rag/data/vector_store/\")\n",
    "print(\"   Then run locally: python main.py chat\")"
   ]
  }
 ]
}