backend.py

# backend.py
"""
Backend logic for EduRAG: model loading, search, and preprocessing functions.
"""
import os
from dotenv import load_dotenv
from langchain_community.vectorstores import Chroma
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain.prompts import PromptTemplate
from langchain.chains import LLMChain
from langchain_google_genai import ChatGoogleGenerativeAI
from PIL import Image
import io
import base64

# --- CONFIGURATION ---
VECTOR_STORE_PATH = "chroma_db"
MODEL_NAME = "all-MiniLM-L6-v2"
GEMINI_MODEL_NAME = "gemini-2.5-flash"

load_dotenv()

def load_models():
    """
    Loads the embedding model, LLM, and the Chroma vector store.
    """
    from langchain_community.embeddings import HuggingFaceEmbeddings
    embeddings = HuggingFaceEmbeddings(model_name=MODEL_NAME)
    if not os.path.exists(VECTOR_STORE_PATH):
        raise FileNotFoundError(f"Chroma database not found at '{VECTOR_STORE_PATH}'. Please run 'python ingest.py' first.")
    vectorstore = Chroma(
        persist_directory=VECTOR_STORE_PATH,
        embedding_function=embeddings
    )
    llm = ChatGoogleGenerativeAI(
        model=GEMINI_MODEL_NAME,
        temperature=0.3,
        google_api_key=os.getenv("GOOGLE_API_KEY")
    )
    prompt_template = """
You are an expert educational assistant helping students with NCERT textbooks. You are having a conversation with a student and should respond in a helpful, conversational manner.

Instructions:
1. Use the provided context from textbooks to answer questions
2. Reference previous parts of the conversation when relevant
3. Be conversational and encouraging
4. If no direct answer exists in the context, look for related information
5. Mention chapter titles, sections, or topics when they help explain concepts
6. Only if absolutely no relevant information exists, suggest how the student might rephrase their question
7. Keep your responses engaging and educational

Previous conversation:
{conversation_history}

Context from textbook(s):
{context}

Student's Question: {question}

Your Response:
"""
    PROMPT = PromptTemplate(template=prompt_template, input_variables=["conversation_history", "context", "question"])
    chain = LLMChain(llm=llm, prompt=PROMPT)
    return vectorstore, chain

def hybrid_search(vectorstore, query, k=5):
    """
    Perform hybrid search combining semantic and keyword-based retrieval.
    """
    semantic_docs = vectorstore.similarity_search_with_score(query, k=k)
    all_docs = vectorstore.get()
    keyword_matches = []
    query_words = set(query.lower().split())
    if 'documents' in all_docs and 'metadatas' in all_docs:
        for i, (doc_text, metadata) in enumerate(zip(all_docs['documents'], all_docs['metadatas'])):
            doc_words = set(doc_text.lower().split())
            overlap = len(query_words.intersection(doc_words))
            if overlap > 0:
                from langchain.schema import Document
                doc = Document(page_content=doc_text, metadata=metadata)
                score = 1.0 / (overlap + 1)
                keyword_matches.append((doc, score))
    all_results = semantic_docs + keyword_matches
    seen_content = set()
    unique_results = []
    for doc, score in all_results:
        content_hash = hash(doc.page_content[:100])
        if content_hash not in seen_content:
            seen_content.add(content_hash)
            unique_results.append((doc, score))
    return sorted(unique_results, key=lambda x: x[1])[:k]

def preprocess_query(query):
    processed_query = query.lower().strip()
    physics_synonyms = {
        "chapters": ["topics", "sections", "units"],
        "physics": ["physical science", "mechanics", "motion"],
        "energy": ["power", "force", "work"],
        "conservation": ["preservation", "constant"],
        "law": ["principle", "rule", "theorem"],
        "motion": ["movement", "kinematics"],
        "electricity": ["electric", "electrical", "current"],
        "magnetism": ["magnetic", "magnet"],
        "light": ["optics", "optical", "rays"],
        "waves": ["wave", "vibration", "oscillation"]
    }
    query_terms = processed_query.split()
    expanded_terms = []
    for term in query_terms:
        expanded_terms.append(term)
        for key, synonyms in physics_synonyms.items():
            if term in key or key in term:
                expanded_terms.extend(synonyms)
    expanded_query = " ".join(expanded_terms)
    return query, expanded_query

def query_documents(query, vectorstore, chain, conversation_history="", k=3):
    """
    Query the RAG system and generate a response.
    
    Args:
        query: User's question
        vectorstore: Chroma vectorstore instance
        chain: LLMChain instance
        conversation_history: Previous conversation context (optional)
        k: Number of documents to retrieve
        
    Returns:
        Generated response string
    """
    # Retrieve relevant documents
    docs = vectorstore.similarity_search(query, k=k)
    
    # Prepare context
    context = "\n\n".join([doc.page_content for doc in docs])
    
    # Generate response
    try:
        # Check if chain expects conversation_history
        if "conversation_history" in chain.prompt.input_variables:
            result = chain.run(
                conversation_history=conversation_history,
                context=context,
                question=query
            )
        else:
            # Fallback for simpler prompts
            result = chain.run(context=context, question=query)
        return result
    except Exception as e:
        # If there's an issue with the chain, return a basic response
        return f"Error generating response: {str(e)}"