src/rag.py

import os
from functools import lru_cache
from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnablePassthrough
from langchain_community.vectorstores import FAISS
from dotenv import load_dotenv

load_dotenv()

# Use absolute path relative to this file's location
SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
DB_FAISS_PATH = os.path.join(SCRIPT_DIR, "vectorstore", "db_faiss")

# Global cache for embeddings and database
_embeddings = None
_db = None

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

@lru_cache(maxsize=1)
def get_embeddings():
    """Load embeddings model once and cache it"""
    try:
        from langchain_huggingface import HuggingFaceEmbeddings
        return HuggingFaceEmbeddings(
            model_name='sentence-transformers/all-MiniLM-L6-v2',
            model_kwargs={'device': 'cpu'}
        )
    except Exception as e:
        raise ImportError(f"Failed to load HuggingFaceEmbeddings: {e}")

def get_vector_db():
    """Load vector database once and cache it"""
    global _db
    if _db is None:
        try:
            embeddings = get_embeddings()
            _db = FAISS.load_local(DB_FAISS_PATH, embeddings, allow_dangerous_deserialization=True)
        except Exception as e:
            raise RuntimeError(f"Error loading FAISS index: {e}\n(Did you run ingest.py?)")
    return _db

def get_answer(query, status_placeholder=None):
    """Get answer for query. status_placeholder is an optional st.empty() for status updates."""
    
    def update_status(text):
        if status_placeholder:
            status_placeholder.markdown(text)
    
    # Load cached embeddings and database
    try:
        update_status("🔍 *Searching knowledge base...*")
        db = get_vector_db()
    except Exception as e:
        return f"❌ Database error: {str(e)}"

    # Setup LLM
    try:
        update_status("🤖 *Connecting to AI...*")
        llm = ChatOpenAI(
            model="google/gemini-2.5-flash",
            api_key=os.getenv("OPENROUTER_API_KEY"),
            base_url="https://openrouter.ai/api/v1",
            temperature=0.7,
            timeout=30,
            max_tokens=2000,
        )
    except Exception as e:
        return f"❌ LLM init error: {str(e)}"

    # Define Prompt
    template = """You are Sagar's personal assistant! 🌟 You're a friendly, warm, and always helpful AI assistant.

    Your personality traits:
    - You're enthusiastic and positive! 😊
    - You use emojis naturally in your responses to make conversations more engaging 🎉
    - You're helpful and always try your best to assist
    - You speak in a warm, conversational tone like a good friend would
    - You're supportive and encouraging 💪

    Answer questions based on the following context about Sagar:
    {context}

    Guidelines:
    1. Answer based on the provided context about Sagar.
    2. Be friendly, warm, and use emojis naturally throughout your response! 😄
    3. Keep your responses conversational and engaging.
    4. If the answer is not in the context, say something like: "Hmm, I don't have that info yet! 🤔 But you can always reach out to Sagar directly at sagar_rathi@zohomail.in 📧"
    5. Do not make up facts - only share what's in the context.
    6. Make the person feel welcome and helped! ✨

    Question: {question}
    """
    prompt = ChatPromptTemplate.from_template(template)

    # Create LCEL Chain
    update_status("📚 *Finding relevant info...*")
    retriever = db.as_retriever()
    rag_chain = (
        {"context": retriever | format_docs, "question": RunnablePassthrough()}
        | prompt
        | llm
        | StrOutputParser()
    )

    update_status("� *Thinking...*")
    
    try:
        response = rag_chain.invoke(query)
        # Clear the status when done
        if status_placeholder:
            status_placeholder.empty()
        return response
    except Exception as e:
        return f"❌ API call error: {str(e)}"