app-gemini.py

import logging
import time
import os

import pdfplumber
import streamlit as st
from dotenv import load_dotenv

from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_google_genai import (
    GoogleGenerativeAIEmbeddings,
    ChatGoogleGenerativeAI,
)
from langchain_community.vectorstores import FAISS
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnablePassthrough

# -------------------------------------------------
# Load Gemini API Key
# -------------------------------------------------

load_dotenv()

api_key = os.getenv("GOOGLE_API_KEY")

if not api_key:
    st.error("GOOGLE_API_KEY not found!")
    st.stop()

os.environ["GOOGLE_API_KEY"] = api_key

# -------------------------------------------------
# Logging
# -------------------------------------------------

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
    handlers=[
        logging.FileHandler("rag_app.log"),
        logging.StreamHandler()
    ]
)

logger = logging.getLogger("rag_app")


# -------------------------------------------------
# Timer Class
# -------------------------------------------------

class Timer:

    def __init__(self, label, ui_placeholder=None):
        self.label = label
        self.ui_placeholder = ui_placeholder

    def __enter__(self):
        self.start = time.perf_counter()
        logger.info(f"START {self.label}")
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        elapsed = time.perf_counter() - self.start

        logger.info(f"FINISH {self.label} -- {elapsed:.2f}s")

        if self.ui_placeholder:
            self.ui_placeholder.write(f"⏱ {self.label}: {elapsed:.2f} seconds")

        st.session_state.setdefault("timings", {})[self.label] = elapsed


# -------------------------------------------------
# Streamlit UI
# -------------------------------------------------

st.set_page_config(page_title="PDF Chatbot", layout="wide")

st.title("📄 PDF Chatbot using Gemini")

with st.sidebar:
    st.header("Upload PDF")
    file = st.file_uploader(
        "Upload a PDF",
        type=["pdf"]
    )

# -------------------------------------------------
# Process PDF
# -------------------------------------------------

if file is not None:

    st.subheader("Processing")

    timing_box = st.container()

    # ---------------------------------------------
    # Extract Text
    # ---------------------------------------------

    with Timer("1. PDF Text Extraction", timing_box):

        text = ""

        with pdfplumber.open(file) as pdf:

            for page in pdf.pages:

                page_text = page.extract_text()

                if page_text:
                    text += page_text + "\n"

    logger.info(f"Extracted {len(text)} characters")

    # ---------------------------------------------
    # Chunking
    # ---------------------------------------------

    with Timer("2. Chunking", timing_box):

        splitter = RecursiveCharacterTextSplitter(
            chunk_size=1000,
            chunk_overlap=200,
            separators=["\n\n", "\n", ". ", " ", ""]
        )

        chunks = splitter.split_text(text)

    logger.info(f"Created {len(chunks)} chunks")

    # ---------------------------------------------
    # Gemini Embeddings
    # ---------------------------------------------

    embeddings = GoogleGenerativeAIEmbeddings(
        model="models/gemini-embedding-001"
    )

    # ---------------------------------------------
    # FAISS
    # ---------------------------------------------

    with Timer(
        f"3. Embeddings + FAISS ({len(chunks)} chunks)",
        timing_box
    ):
        pass
        # vector_store = FAISS.from_texts(
        #     chunks,
        #     embedding=embeddings
        # )

    if st.session_state.get("vector_store") is None:
        vector_store = FAISS.from_texts(
            chunks,
            embedding=embeddings
        )

        st.session_state.vector_store = vector_store
        st.success("PDF processed successfully!")

    else:
        vector_store = st.session_state.vector_store
        st.success("Using cached vector store!")

    # ---------------------------------------------
    # Retriever
    # ---------------------------------------------

    retriever = vector_store.as_retriever(
        search_type="mmr",
        search_kwargs={"k": 4}
    )

    # ---------------------------------------------
    # Gemini LLM
    # ---------------------------------------------

    llm = ChatGoogleGenerativeAI(
    model="gemini-2.5-flash",
    temperature=0.3,
    max_tokens=2048
)

    # ---------------------------------------------
    # Prompt
    # ---------------------------------------------

    prompt = ChatPromptTemplate.from_messages(
        [
            (
                "system",
                """You are a helpful assistant answering questions about a PDF.

Rules:

1. Only answer from the given context.
2. If the answer is not available, say:
   "The information is not available in the document."
3. Give detailed explanations.
4. Use bullet points whenever helpful.
5. Include numbers and facts from the document.
6. Do not make up information.

Context:

{context}
"""
            ),
            ("human", "{question}")
        ]
    )

    # ---------------------------------------------
    # Helper
    # ---------------------------------------------

    def format_docs(docs):
        return "\n\n".join(doc.page_content for doc in docs)

    # ---------------------------------------------
    # RAG Chain
    # ---------------------------------------------

    chain = (
        {
            "context": retriever | format_docs,
            "question": RunnablePassthrough()
        }
        | prompt
        | llm
        | StrOutputParser()
    )

    # ---------------------------------------------
    # User Question
    # ---------------------------------------------

    question = st.text_input(
        "Ask a question about the PDF"
    )

    if question:

        with Timer("4. Retrieval + Gemini Response", timing_box):

            response = chain.invoke(question)

        st.subheader("Answer")

        st.write(response)

        with st.expander("Timing Summary"):

            timings = st.session_state.get("timings", {})

            for stage, sec in timings.items():
                st.write(f"{stage}: {sec:.2f} seconds")