Update app.py

app.py

@@ -0,0 +1,451 @@
+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, DESTINATION_FOLDER)
+
+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 = []
+    for file_path in new_file_paths:
+        if not file_path.lower().endswith(".pdf"):
+            continue
+            
+        # Open document using fitz
+        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") # "text" maintains logical flow; "blocks" is better for tables
+            
+            # Creating a LangChain Document object for each page
+            # This replaces the need for RecursiveCharacterTextSplitter
+            page_doc = Document(
+                page_content=text,
+                metadata={
+                    "source": file_name,
+                    "page": page_num + 1,  # 1-indexed for user readability
+                    "total_pages": len(doc),
+                    "format": "PDF",
+                    "extraction_method": "PyMuPDF"
+                }
+            )
+            splits.append(page_doc)
+            
+        doc.close()
+        
+    return splits
+
+splits = gen_splits(DESTINATION_FOLDER)
+embedding_func = HuggingFaceEmbeddings(model_name='all-MiniLM-L6-v2')
+
+def vectordb_from_splits(splits):
+    vectordb = Chroma.from_documents(documents=splits, persist_directory=PERSIST_DIR, embedding=embedding_func)
+    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. 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=2048,
+        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
+
+
+# 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 GIS/AIS substations (220kV / 400kV / 765kV)
+- IEC and IEEE standards referenced in POWERGRID specs
+- Equipment-specific specs: Circuit Breakers, Isolators, Surge Arresters, CTs, VTs, 
+  Power Transformers, Reactors, Protection Relays, Control & Relay Panels
+- Specific Requirements documents (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(folder_name, 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=2048,
+        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. 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(folder_name, 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 = "# Engineering SS — Model Technical Specifications"
+
+Description = """
+## Welcome to the Search Engine for POWERGRID Model Technical Specifications! 👋
+
+This tool leverages the power of AI to answer queries based on:
+
+* **Model Technical Specifications** — POWERGRID Engineering Department 📑
+* Includes latest Model TS as on date.
+
+**Tips for Effective Use:**
+* Use elaborate, keyword-rich questions for precise responses. 🎯
+* Include equipment names (e.g., *GIS*, *Circuit Breaker*, *Surge Arrester*). 🔌
+* Clear chat if responses seem off-context. 🔄
+"""
+
+# ── Layout ─────────────────────────────────────────────────────────────────────
+with gr.Blocks(
+    css="CSS/style.css",
+    fill_height=True,
+    theme=gr.themes.Base(),
+) 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",
+                    )
+                    file_output.render()
+
+                # with gr.Tab("Pre-Bid Schemes"):
+                #     gr.Markdown(
+                #         "### Pre-Bid Scheme Query\n"
+                #         "_Interface under development — coming soon._"
+                #     )
+
+    demo.queue()