apds/rag_utils.py

import os
import time
from pathlib import Path
from typing import List, Optional

import streamlit as st
import google.generativeai as genai
from langchain_community.document_loaders import PyPDFLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_community.vectorstores import FAISS
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain.chains import RetrievalQA

# ── Configuration ───────────────────────────────────────────────────────────

DB_PATH = Path(__file__).resolve().parent.parent / "models" / "vector_db"
KNOWLEDGE_DIR = Path(__file__).resolve().parent.parent / "docs"

@st.cache_resource
def get_embeddings():
    """Download and cache local embedding model (Zero tokens used)."""
    return HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")

@st.cache_resource
def get_vector_db():
    """Load or create the local FAISS index (Zero tokens used for indexing)."""
    embeddings = get_embeddings()
    
    if DB_PATH.exists():
        return FAISS.load_local(str(DB_PATH), embeddings, allow_dangerous_deserialization=True)
    
    # Create new index if not exists
    st.info("AI is reading engineering manuals for the first time... please wait.")
    all_docs = []
    
    for pdf in KNOWLEDGE_DIR.glob("*.pdf"):
        loader = PyPDFLoader(str(pdf))
        all_docs.extend(loader.load())
    
    text_splitter = RecursiveCharacterTextSplitter(chunk_size=600, chunk_overlap=100)
    splits = text_splitter.split_documents(all_docs)
    
    vector_db = FAISS.from_documents(splits, embeddings)
    vector_db.save_local(str(DB_PATH))
    return vector_db

def get_ai_consultant():
    """Initialize the Gemini LLM for the final expert report."""
    api_key = os.getenv("GOOGLE_API_KEY")
    if not api_key:
        return None
    
    genai.configure(api_key=api_key)
    return ChatGoogleGenerativeAI(model="gemini-flash-latest", google_api_key=api_key, version="v1")

def generate_expert_report(detection_summary: str) -> str:
    """
    RAG Pipeline: Retrieve context from IRC manuals and query Gemini.
    (Tokens only used for the final reasoning, not for searching).
    """
    try:
        db = get_vector_db()
        llm = get_ai_consultant()
        
        if not llm:
            return "Error: Gemini API Key not found. Please check your .env file."

        # Setup QA chain with limited context (Token Optimization)
        qa_chain = RetrievalQA.from_chain_type(
            llm=llm,
            chain_type="stuff",
            retriever=db.as_retriever(search_kwargs={"k": 3}), # Only 3 chunks to save tokens
            return_source_documents=True
        )

        prompt = (
            f"As a Senior Civil Engineer specializing in Indian Road Congress (IRC) standards, "
            f"provide a professional maintenance recommendation for the following detection: {detection_summary}. "
            f"Keep your response concise (max 200 words) to optimize for speed. "
            f"Always cite which IRC clause or manual you are referencing."
        )

        response = qa_chain.invoke({"query": prompt})
        
        result = response["result"]
        sources = set([doc.metadata['source'].split('/')[-1].split('\\')[-1] for doc in response["source_documents"]])
        
        final_output = f"{result}\n\n**Reference Sources:** {', '.join(sources)}"
        return final_output

    except Exception as e:
        return f"Expert logic error: {str(e)}"