app.py

import os
import fitz 
import zipfile
import requests
import gradio as gr
from dotenv import load_dotenv

# LangChain & Vector DB imports
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_chroma import Chroma
from langchain_groq import ChatGroq
from langchain_core.documents import Document
from langchain_core.prompts import (
    PromptTemplate,
    ChatPromptTemplate,
    SystemMessagePromptTemplate,
    HumanMessagePromptTemplate,
    MessagesPlaceholder,
)
from langchain_core.runnables.history import RunnableWithMessageHistory
from langchain_core.runnables import RunnablePassthrough
from langchain_core.output_parsers import StrOutputParser
from langchain_community.chat_message_histories import ChatMessageHistory
from langchain_core.messages import HumanMessage, AIMessage

load_dotenv()

# --- Configurations & Secrets ---
HF_TOKEN = os.getenv('HF_TOKEN')
GROQ_API_KEY = os.getenv("GROQ_API_KEY")
SOURCE_URL = os.getenv('URL')
PERSIST_DIR = './chroma_db/'
GROQ_MODEL = "llama-3.3-70b-versatile"
DESTINATION_FOLDER = "Model_TS"

# --- Initialization Downloading ---
def download_and_extract_zip(url, destination_folder):
    """Downloads a zip file from a URL and extracts its contents to the specified destination folder."""
    zip_file_path = "temp.zip"
    
    try:
        # Send an HTTP GET request to the OneDrive link to download the file
        response = requests.get(url)
        # Check if the request was successful (status code 200)
        response.raise_for_status()
        # Save the zip file to a temporary location
        with open(zip_file_path, "wb") as f:
            f.write(response.content)
        # Create the destination folder if it doesn't exist
        os.makedirs(destination_folder, exist_ok=True)
        # Extract the contents of the zip file
        with zipfile.ZipFile(zip_file_path, 'r') as zip_ref:
            zip_ref.extractall(destination_folder)
        print(f"Zip file downloaded and extracted to: {destination_folder}")

    except requests.exceptions.RequestException as e:
        print(f"Error downloading file: {e}")

    finally:
        # Remove the temporary zip file
        if os.path.exists(zip_file_path):
            os.remove(zip_file_path)


# Splitting, Initialize Embeddings and VectorDB Storage
download_and_extract_zip(SOURCE_URL, os.getcwd())

def gen_splits(folder_name):
    file_paths = os.listdir(folder_name)
    new_file_paths = [os.path.join(os.getcwd(), folder_name, file) for file in file_paths]
    
    splits = []
    empty_pages = 0

    for file_path in new_file_paths:
        if not file_path.lower().endswith(".pdf"):
            continue
            
        doc = fitz.open(file_path)
        file_name = os.path.basename(file_path)
        
        for page_num in range(len(doc)):
            page = doc.load_page(page_num)
            text = page.get_text("text").strip()   # ← strip whitespace

            # ── Skip empty/image-only pages ────────────────────────────────
            if not text or len(text) < 20:         # ← 20 chars minimum threshold
                empty_pages += 1
                continue

            page_doc = Document(
                page_content=text,
                metadata={
                    "source": file_name,
                    "page": page_num + 1,
                    "total_pages": len(doc),
                    "format": "PDF",
                    "extraction_method": "PyMuPDF"
                }
            )
            splits.append(page_doc)
            
        doc.close()

    print(f"✓ Loaded {len(splits)} pages | Skipped {empty_pages} empty/image-only pages")
    return splits

splits = gen_splits(DESTINATION_FOLDER)
embedding_func = HuggingFaceEmbeddings(model_name='all-MiniLM-L6-v2')

def vectordb_from_splits(splits):
    # ── Reuse existing ChromaDB if persist dir already populated ──────────────
    if os.path.exists(PERSIST_DIR) and os.listdir(PERSIST_DIR):
        print("✓ Loading existing ChromaDB from disk — skipping re-embedding.")
        return Chroma(persist_directory=PERSIST_DIR, embedding_function=embedding_func)

    if not splits:
        raise ValueError("No text content extracted. Check if PDFs are scanned images.")

    print(f"Building ChromaDB from {len(splits)} chunks...")
    vectordb = Chroma.from_documents(
        documents=splits,
        persist_directory=PERSIST_DIR,
        embedding=embedding_func
    )
    print(f"✓ ChromaDB built successfully.")
    return vectordb


vectordb = vectordb_from_splits(splits)


# RAG Chain
GROQ_API_KEY = os.getenv("GROQ_API_KEY")  # set in HF Spaces → Settings → Secrets

# ── Model options on Groq free tier (swap as needed) ──────────────────────────
# "llama-3.3-70b-versatile"   ← RECOMMENDED: best reasoning, table fidelity
# "llama3-8b-8192"            ← fallback if hitting TPM limits
# "qwen-qwq-32b"              ← strong reasoning, good for clause referencing
# "deepseek-r1-distill-llama-70b" ← chain-of-thought style; verbose but thorough
GROQ_MODEL = "llama-3.3-70b-versatile"

# ── Session store ──────────────────────────────────────────────────────────────
session_store: dict = {}

def get_session_history(session_id: str) -> ChatMessageHistory:
    if session_id not in session_store:
        session_store[session_id] = ChatMessageHistory()
    return session_store[session_id]


def get_file(source_documents):
    references, files_in_order = [], []
    seen_refs, seen_files = set(), set()
    for doc in source_documents:
        source = os.path.basename(doc.metadata.get("source", "unknown"))
        page   = doc.metadata.get("page", 0) + 1
        ref    = f"Page-{page} of {source}"
        if ref not in seen_refs:
            references.append(ref)
            seen_refs.add(ref)
        if source not in seen_files:
            files_in_order.append(source)
            seen_files.add(source)
    return references, files_in_order


def build_chain(vectordb: Chroma):
    system_instruction = (
        "You are an expert **Electrical Engineer AI Assistant**, specialized in power systems "
        "and substation design (AIS/GIS up to 765kV), providing insights strictly from the provided context.\n\n"
        "**Formatting Guidelines:**\n"
        "1. Be Precise and Organize using **bullet points or numbered lists** where appropriate.\n"
        "2. **Bold** key technical terms, parameters, and essential facts.\n"
        "3. Use **technical language** consistent with IEC/IEEE/POWERGRID standards.\n"
        "4. For multi-step explanations, use **sub-headings** (e.g., `## Sub-section`).\n"
        "5. **Always include clause references (e.g., Clause XX.XX) for every piece of information.**\n"
        "6. **CRITICAL: If context contains a table, reproduce it EXACTLY — preserve all rows, "
        "columns, headers, and alignment. Never paraphrase table data.**\n\n"
        "**Context Prioritization:**\n"
        "1. Prioritize documents directly related to the queried equipment type.\n"
        "2. 'Specific Requirements' clauses **supersede** all other documents — reflect modified clauses first.\n"
        "3. If context is insufficient: 'The available documents do not contain information regarding [detail].'\n"
        "4. **Do not invent information** outside the provided context."
    )

    prompt = ChatPromptTemplate.from_messages([
        SystemMessagePromptTemplate.from_template(system_instruction),
        MessagesPlaceholder(variable_name="chat_history"),
        HumanMessagePromptTemplate.from_template(
            "Context:\n{context}\n\nQuestion:\n{question}"
        ),
    ])

    # ── Groq LLM ───────────────────────────────────────────────────────────────
    llm = ChatGroq(
        model=GROQ_MODEL,
        temperature=0.1,
        max_tokens=1024,
        api_key=GROQ_API_KEY,
    )

    # ── Retriever ──────────────────────────────────────────────────────────────
    retriever = vectordb.as_retriever(
        search_type="mmr",
        search_kwargs={"k": 3, "lambda_mult": 0.5, "fetch_k": 15},
    )

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

    rag_core = (
        RunnablePassthrough.assign(
            context=lambda x: format_docs(retriever.invoke(x["question"]))
        )
        | prompt
        | llm
        | StrOutputParser()
    )

    chain_with_history = RunnableWithMessageHistory(
        rag_core,
        get_session_history,
        input_messages_key="question",
        history_messages_key="chat_history",
    )

    return chain_with_history, retriever


# ── Build once at startup (not per Gradio call) ───────────────────────────────
chain, retriever = build_chain(vectordb)   # vectordb initialised elsewhere

# retriever = vectordb.as_retriever(
#         search_type="mmr",
#         search_kwargs={"k": 3, "lambda_mult": 0.5, "fetch_k": 15},
# )
    

# Query Re-write
def rewrite_query(question: str, llm) -> str:
    """
    Rewrites the user query to improve retrieval against POWERGRID
    technical specification documents (IEC/IEEE standards, GIS/AIS
    substation specs, protection & control documents).
    """
    rewrite_prompt = PromptTemplate.from_template("""
You are an expert query rewriter for a POWERGRID technical document retrieval system.
The document corpus contains:
- Model Technical Specifications for various equipments used in GIS/AIS substations (132kV /220kV / 400kV / 765kV)
- IEC and IEEE standards referenced in POWERGRID specs
- Equipment-specific specs: Circuit Breakers, Isolators, Surge Arresters, CTs, VTs, Gas Insulated Switchgears,
  Power Transformers, Reactors, Protection Relays, Control & Relay Panels, Visual Monitoring Systems (VMS), Switchyard Erection 
- Specific Requirements Document (which supersede other docs)

Your task:
1. Expand abbreviations (e.g., CB → Circuit Breaker, SA → Surge Arrester, CT → Current Transformer)
2. Add relevant technical keywords likely present in the documents
3. Include clause/section indicators if the query implies a specific requirement
4. If the query is vague, make it specific to power system Substation context
5. Preserve the original intent — do NOT change what is being asked
6. Output ONLY the rewritten query, nothing else

Original Query: {question}

Rewritten Query:""")

    chain = rewrite_prompt | llm | StrOutputParser()
    rewritten = chain.invoke({"question": question})
    return rewritten.strip()


## RAG_PDF without re-writing query
def rag_pdf(question: str, chat_history: list, session_id: str = "default"):
    response_text = chain.invoke(
        {"question": question},
        config={"configurable": {"session_id": session_id}},
    )

    source_docs = retriever.invoke(question)
    source_docs = source_docs[:3]
    references, unique_sources = get_file(source_docs)

    if references:
        response_text += "\n\n**References:**\n"
        for i, ref in enumerate(references, 1):      # numbered list, all 3
            response_text += f"{i}. {ref}\n"

    file_paths = [
        os.path.realpath(os.path.join(DESTINATION_FOLDER, src))
        for src in unique_sources
    ]
    return response_text, file_paths


# After query rewriter
def rag_pdf_query_rewrite(question: str, history: list):

    # ── LLM (same instance used for rewriting + generation) ───────────────────
    llm = ChatGroq(
        model=GROQ_MODEL,
        temperature=0.1,
        max_tokens=1024,
        api_key=GROQ_API_KEY,
    )

    # ── Step 1: Rewrite the query ──────────────────────────────────────────────
    rewritten_question = rewrite_query(question, llm)
    print(f"\n[Query Rewriter]\n  Original : {question}\n  Rewritten: {rewritten_question}\n")

    # ── Step 2: Retrieve using rewritten query ─────────────────────────────────
    source_docs = retriever.invoke(rewritten_question)
    source_docs = source_docs[:3]

    # ── Step 3: Build context from retrieved docs ──────────────────────────────
    context = "\n\n---\n\n".join(doc.page_content for doc in source_docs)

    # ── Step 4: Build prompt with original + rewritten query ───────────────────
    system_instruction = (
        "You are an expert **Electrical Engineer AI Assistant**, specialized in power systems "
        "and substation design (AIS/GIS up to 765kV), providing insights strictly from the provided context.\n\n"
        "**Formatting Guidelines:**\n"
        "1. Be Precise and Organize using **bullet points or numbered lists** where appropriate.\n"
        "2. **Bold** key technical terms, parameters, and essential facts.\n"
        "3. Use technical language consistent with IEC/IEEE/POWERGRID standards.\n"
        "4. For multi-step explanations use **sub-headings**.\n"
        "5. **Always include clause references (e.g., Clause XX.XX) for every fact.**\n"
        "6. **CRITICAL: Reproduce tables EXACTLY — preserve all rows, columns, headers. Never paraphrase table data.**\n\n"
        "**Context Prioritization:**\n"
        "1. Prioritize documents directly related to the queried equipment.\n"
        "2. 'Specific Requirements' clauses supersede all other documents.\n"
        "3. If context is insufficient: state 'The available documents do not contain information regarding [detail].'\n"
        "4. Do not invent information outside the provided context."
    )

    prompt = ChatPromptTemplate.from_messages([
        SystemMessagePromptTemplate.from_template(system_instruction),
        MessagesPlaceholder(variable_name="chat_history"),
        HumanMessagePromptTemplate.from_template(
            "Context:\n{context}\n\n"
            "Original Question: {original_question}\n"
            "Rewritten Question: {rewritten_question}"
        ),
    ])

    # ── Step 5: Convert history to LangChain messages ─────────────────────────
    # Gradio 6 passes history as list of dicts: {"role": .., "content": ..}
    from langchain_core.messages import HumanMessage, AIMessage
    lc_history = []
    for msg in history:
        if msg["role"] == "user":
            lc_history.append(HumanMessage(content=msg["content"]))
        elif msg["role"] == "assistant":
            lc_history.append(AIMessage(content=msg["content"]))

    # ── Step 6: Generate response ──────────────────────────────────────────────
    chain = prompt | llm | StrOutputParser()
    response_text = chain.invoke({
        "context": context,
        "original_question": question,
        "rewritten_question": rewritten_question,
        "chat_history": lc_history,
    })

    # ── Step 7: Attach references ──────────────────────────────────────────────
    references, unique_sources = get_file(source_docs)
    if references:
        response_text += "\n\n**References:**\n"
        for i, ref in enumerate(references, 1):
            response_text += f"{i}. {ref}\n"

    file_paths = [
        os.path.realpath(os.path.join(DESTINATION_FOLDER, src))
        for src in unique_sources
    ]

    return response_text, file_paths


## BACKEND Interface
# ── Pre-define components with render=False ────────────────────────────────────
file_output = gr.File(
    render=False,
    label="Reference Documents",
    file_count="multiple",
    interactive=False,
)

chatbot = gr.Chatbot(
    render=False,
    height=500,
    show_label=False,
    placeholder="Ask a question about POWERGRID Technical Specifications...",
    layout="bubble",
)

# ── Wrapper function ───────────────────────────────────────────────────────────
def rag_pdf_ui(question: str, history: list) -> tuple:
    response_text, file_paths = rag_pdf_query_rewrite(question, history)
    return response_text, file_paths

# ── Text ───────────────────────────────────────────────────────────────────────
Title = "# PG-ATLAS : POWERGRID - AI Technical Library & Assistance System"

Description = """
## Welcome to the AI-Powered Search Engine
### Model Technical Specifications | Engineering — Substation Department
---
This intelligent assistant leverages a **Large Language Model (LLM)** to provide precise, context-aware answers directly from POWERGRID's official documentation.

**Document Coverage:**
* 📑 Model Technical Specifications — Engineering (Substation) Department, POWERGRID
* ⚡ Applicable to AIS/GIS Substations up to **765kV**
* 🔄 Updated to the latest revision as on 03-May-2026.
---
**Guidelines for Effective Use:**
* Frame queries with **specific equipment names** and **voltage class** (e.g., *765kV GIS Circuit Breaker*, *400kV Surge Arrester*, *220kV Isolator*)
* Include **clause keywords** for targeted retrieval (e.g., *type test*, *insulation level*, *earthing*, *interpole cabling*)
* For comparative queries, specify the parameter of interest (e.g., *rated current*, *BIL*, *SF₆ gas pressure*)
* Clear the chat session if responses appear off-context or drift from the query intent
> ⚠️ *Responses are strictly based on the provided documentation. Always verify critical design parameters against the original Model TS before application.*
"""

# ── Layout ─────────────────────────────────────────────────────────────────────
with gr.Blocks(fill_height=True) as demo:

    with gr.Column():

        with gr.Row():
            with gr.Column(scale=1):
                gr.Image(
                    value="Images/PG Logo.png",
                    width=200,
                    show_label=False,
                    interactive=False,
                    elem_id="Logo",
                    buttons=[],
                )
            with gr.Column(scale=3, elem_classes=["center-title"]):
                gr.Markdown(Title)

        with gr.Row():
            with gr.Column():
                gr.Markdown(Description)

        with gr.Row():
            with gr.Column(elem_classes=["chat_container"]):

                with gr.Tab("Model TS"):
                    gr.ChatInterface(
                        fn=rag_pdf_ui,
                        chatbot=chatbot,
                        title=None,
                        concurrency_limit=5,
                        fill_height=True,
                        delete_cache=(300, 360),
                        examples=[
                            "Type Tests for HV Switchgears.",
                            "What should be the height of GIB outside GIS hall for any type of crossings?",
                            "What is the resistivity of stone for ground spreading in switchyard?",
                            "Specify the details of Earthing System.",
                            "Specify details for Interpole cabling in CB.",
                        ],
                        additional_outputs=[file_output],
                        editable=True,
                        flagging_mode="never",
                        cache_examples=False,
                    )
                    file_output.render()

                # with gr.Tab("Pre-Bid Schemes"):
                #     gr.Markdown(
                #         "### Pre-Bid Scheme Query\n"
                #         "_Interface under development — coming soon._"
                #     )

demo.launch(
    css="CSS/style.css",
    theme=gr.themes.Base()
)