api/rag.py

import os
from typing import Optional
from langchain_chroma import Chroma
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_community.document_loaders import PyPDFLoader
from langchain_groq import ChatGroq
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_classic.chains import create_history_aware_retriever
from langchain_classic.chains.combine_documents import create_stuff_documents_chain
from langchain_community.chat_message_histories import ChatMessageHistory
from langchain_core.chat_history import BaseChatMessageHistory
from langchain_community.retrievers import BM25Retriever
from langchain_core.documents import Document
from typing import List, Tuple
from .core import config


def get_embeddings():
    if config.HUGGINGFACE_TOKEN:
        os.environ["HUGGINGFACE_TOKEN"] = config.HUGGINGFACE_TOKEN
    # Use all-MiniLM-L6-v2: smaller model (~90MB) that works well on free tier
    # all-mpnet-base-v2 (~420MB) is too large for Render free tier (512MB RAM)
    return HuggingFaceEmbeddings(
        model_name="all-MiniLM-L6-v2",
        encode_kwargs={"normalize_embeddings": True},
    )


def get_user_chroma_dir(user_id: str, session_id: str | None = None) -> str:
    # Use /tmp for ChromaDB to avoid permission issues in HF Spaces
    base = "/tmp/chroma_db"
    if session_id:
        return os.path.join(base, f"user_{user_id}", f"session_{session_id}")
    return os.path.join(base, f"user_{user_id}")


def get_vectorstore_for_user(user_id: str, session_id: str | None = None) -> Chroma:
    if not session_id:
        # Enforce per-session isolation; caller must provide session_id
        raise ValueError("session_id is required for vectorstore access")
    
    # Try to use persistent directory, fall back to in-memory if it fails
    try:
        persist_dir = get_user_chroma_dir(user_id, session_id)
        os.makedirs(persist_dir, exist_ok=True)
        embeddings = get_embeddings()
        return Chroma(persist_directory=persist_dir, embedding_function=embeddings)
    except Exception as e:
        print(f"⚠️ Persistent ChromaDB failed ({e}), using in-memory mode")
        # Fallback to in-memory ChromaDB (no persistence)
        embeddings = get_embeddings()
        return Chroma(embedding_function=embeddings)


from .core.ocr import extract_text_from_pdf_with_ocr

def index_pdf_for_user(user_id: str, temp_pdf_path: str, session_id: str | None = None):
    if not session_id:
        raise ValueError("session_id is required for indexing")
    
    loader = PyPDFLoader(temp_pdf_path)
    # Load all docs first
    raw_docs = loader.load()
    
    final_docs = []
    pages_needs_ocr = []
    
    # Analyze each page
    # If a page has very little text, it might be an image/scan -> mark for OCR
    for i, doc in enumerate(raw_docs):
        content = doc.page_content or ""
        # Simple heuristic: if less than 50 chars of meaningful text, try OCR
        # This covers empty pages or pages with just "Scanned by CamScanner" etc.
        if len(content.strip()) < 50:
            pages_needs_ocr.append(i)
        else:
            final_docs.append(doc)
            
    # Run OCR on identified pages
    if pages_needs_ocr:
        print(f"OCR needed for {len(pages_needs_ocr)} pages: {pages_needs_ocr}")
        try:
            ocr_docs = extract_text_from_pdf_with_ocr(temp_pdf_path, pages_needs_ocr)
            # Re-insert OCR docs in correct order (though order matters less for bag-of-chunks, it helps context)
            # Since final_docs is already populated with non-OCR pages, we can just append and sort or just append specific ones.
            # Simpler: just extend final_docs with whatever we got. 
            # Note: The OCR docs metadata 'page' corresponds to 0-indexed page num.
            
            # Filter out OCR failures (empty text)
            ocr_docs = [d for d in ocr_docs if d.page_content.strip()]
            final_docs.extend(ocr_docs)
            
        except Exception as e:
            print(f"Warning: OCR failed ({e}). Proceeding with what we have.")
    
    # If we still have absolutely no text after everything
    if not final_docs:
         # One last desperate attempt: Force OCR on ALL pages if we gathered nothing so far
         # (Only if we haven't already tried OCR on all pages)
         if len(pages_needs_ocr) != len(raw_docs):
             print("No text found in initial pass. Attempting OCR on ALL pages...")
             try:
                 final_docs = extract_text_from_pdf_with_ocr(temp_pdf_path)
                 final_docs = [d for d in final_docs if d.page_content.strip()]
             except Exception as e:
                print(f"Fallback OCR failed: {e}")

    if not final_docs:
        raise ValueError("No extractable text found in the PDF, even after OCR attempt.")

    # Sort solely for debugging sanity (optional)
    final_docs.sort(key=lambda x: x.metadata.get("page", 0))

    # Slightly smaller chunks generally improve recall; keep modest overlap for continuity
    splitter = RecursiveCharacterTextSplitter(chunk_size=900, chunk_overlap=150)
    splits = splitter.split_documents(final_docs)
    
    if not splits:
        raise ValueError("No text chunks generated from the PDF.")
        
    vs = get_vectorstore_for_user(user_id, session_id)
    vs.add_documents(splits)


def get_llm() -> ChatGroq:
    # Deterministic answers; we rely on retrieved context only
    # Using Groq's free open-source model: openai/gpt-oss-120b
    return ChatGroq(api_key=config.GROQ_API_KEY, model="openai/gpt-oss-120b", temperature=0)


def build_conversational_chain(user_id: str, history: Optional[BaseChatMessageHistory], session_id: str | None = None):
    if not session_id:
        raise ValueError("session_id is required for chat")
    vs = get_vectorstore_for_user(user_id, session_id)
    # Embedding retriever (primary). Avoid score_threshold here due to Chroma compatibility.
    embedding_retriever = vs.as_retriever(search_kwargs={"k": 20})

    # Build a lightweight BM25 retriever over all docs in the session for hybrid search
    bm25 = None
    try:
        # Try to get all documents from the collection
        collection = vs._collection
        all_data = collection.get(include=["documents", "metadatas"])
        texts = all_data.get("documents", []) or []
        metas = all_data.get("metadatas", []) or []
        
        print(f"Chroma collection has {len(texts)} documents")
        
        if texts and len(texts) > 0:
            bm25_docs: List[Document] = [Document(page_content=t, metadata=m or {}) for t, m in zip(texts, metas)]
            bm25 = BM25Retriever.from_documents(bm25_docs)
            bm25.k = 20
            print(f"BM25 initialized with {len(bm25_docs)} documents")
        else:
            print("WARNING: No documents found in Chroma collection - did you upload a PDF?")
    except Exception as e:
        print(f"BM25 initialization failed: {e}")
        import traceback
        traceback.print_exc()
        bm25 = None
    llm = get_llm()

    contextualize_q_system_prompt = (
        "Given a chat history and the latest user question"
        " which might reference context in the chat history, "
        "formulate a standalone question which can be understood "
        "without the chat history. Do NOT answer the question, "
        "just reformulate it if needed and otherwise return it as is."
    )
    contextualize_q_prompt = ChatPromptTemplate.from_messages(
        [
            ("system", contextualize_q_system_prompt),
            MessagesPlaceholder("chat_history"),
            ("human", "{input}"),
        ]
    )

    history_aware_retriever = create_history_aware_retriever(llm, embedding_retriever, contextualize_q_prompt)

    system_prompt = (
        "You are a grounded RAG assistant.\n"
        "Use ONLY the information in the retrieved context to answer.\n"
        "Do NOT use prior knowledge or invent facts.\n\n"
        "When answering from context, be clear and structured (headings, bullet points, numbered lists as needed).\n\n"
        "Retrieved context follows.\n{context}"
    )
    qa_prompt = ChatPromptTemplate.from_messages(
        [
            ("system", system_prompt),
            MessagesPlaceholder("chat_history"),
            ("human", "{input}"),
        ]
    )

    question_answer_chain = create_stuff_documents_chain(llm, qa_prompt)

    # Compose a custom retrieval function that performs multi-query expansion and RRF fusion
    def retrieve(query: str, chat_history) -> List[Document]:
        # Multi-query expansion: generate several paraphrases of the user query
        # Simplified to avoid breaking on subsequent queries
        queries = [query]
        try:
            mq_prompt = ChatPromptTemplate.from_messages([
                ("system", "Generate 2 alternative search queries to find relevant information. Return ONLY a JSON array of strings, nothing else. Example: [\"query 1\", \"query 2\"]"),
                ("human", "{q}")
            ])
            mq = llm.invoke(mq_prompt.format_messages(q=query)).content.strip()
            import json
            # Try to extract JSON array if wrapped in markdown code blocks
            if "```" in mq:
                # Extract content between ```json and ``` or ``` and ```
                start = mq.find("[")
                end = mq.rfind("]") + 1
                if start != -1 and end > start:
                    mq = mq[start:end]
            parsed = json.loads(mq)
            if isinstance(parsed, list):
                for alt in parsed:
                    if isinstance(alt, str) and alt.strip() and alt not in queries:
                        queries.append(alt.strip())
                print(f"Multi-query expansion: Generated {len(queries)-1} additional queries")
        except Exception as e:
            # Log for debugging but don't fail - single query still works fine
            print(f"Multi-query expansion skipped ({e}). Continuing with original query.")
            pass

        def dedup_by_text(docs: List[Document]) -> List[Document]:
            seen = set()
            unique = []
            for d in docs:
                key = (d.page_content.strip(), str(d.metadata))
                if key in seen:
                    continue
                seen.add(key)
                unique.append(d)
            return unique

        # Collect candidates per retriever
        candidates: List[Tuple[Document, int]] = []  # (doc, rank)
        print(f"Retrieve: Processing {len(queries)} queries: {[q[:50] for q in queries]}")
        for i, q in enumerate(queries):
            # Embedding hits - always retrieve, don't filter by threshold at this stage
            try:
                docs = embedding_retriever.invoke(q)
                print(f"  Query {i+1}: Embedding retriever returned {len(docs)} docs for: '{q[:50]}...'")
            except Exception as e:
                print(f"  Query {i+1}: Embedding invoke failed: {e}, trying get_relevant_documents")
                try:
                    docs = embedding_retriever.get_relevant_documents(q)
                    print(f"  Query {i+1}: get_relevant_documents returned {len(docs)} docs")
                except Exception as e2:
                    print(f"  Query {i+1}: get_relevant_documents also failed: {e2}")
                    docs = []
            for rank, d in enumerate(docs):
                candidates.append((d, rank))
            # BM25 hits
            if bm25 is not None:
                try:
                    # Try invoke first (newer LangChain), fall back to get_relevant_documents
                    try:
                        bm25_docs = bm25.invoke(q)
                    except AttributeError:
                        bm25_docs = bm25.get_relevant_documents(q)
                    print(f"  BM25 returned {len(bm25_docs)} docs for query: {q[:50]}")
                    for rank, d in enumerate(bm25_docs):
                        candidates.append((d, rank))
                except Exception as e:
                    print(f"  BM25 retrieval failed: {e}")
                    pass

        # Reciprocal Rank Fusion
        scores = {}
        for d, r in candidates:
            key = (d.page_content, tuple(sorted(d.metadata.items()))) if isinstance(d.metadata, dict) else (d.page_content, str(d.metadata))
            scores[key] = scores.get(key, 0) + 1.0 / (60 + r)  # 60 for stability

        # Rebuild documents with aggregated scores
        scored_docs = []
        for d, r in candidates:
            key = (d.page_content, tuple(sorted(d.metadata.items()))) if isinstance(d.metadata, dict) else (d.page_content, str(d.metadata))
            if key in scores:
                d.metadata = dict(d.metadata or {})
                d.metadata["rrf_score"] = scores[key]
                scored_docs.append(d)
        # Sort by fused score desc, then truncate
        scored_docs.sort(key=lambda x: x.metadata.get("rrf_score", 0), reverse=True)
        out = dedup_by_text(scored_docs)[:15]
        print(f"Retrieve: Final result: {len(out)} documents after deduplication and ranking")
        
        return out

    # Return a simple invokable object that mirrors the output shape of create_retrieval_chain
    class SimpleRAG:
        def invoke(self, inputs):
            q = inputs.get("input", "")
            chat_history = inputs.get("chat_history", [])
            print(f"SimpleRAG: Processing query: '{q[:100]}...'")
            docs = retrieve(q, chat_history)
            print(f"SimpleRAG: Retrieved {len(docs)} documents")
            if not docs:
                print("SimpleRAG: No documents retrieved, returning 'I don't know' response")
                return {"answer": "I don't know based on the uploaded documents. Please make sure you have uploaded PDF documents to this session.", "context": []}
            answer = question_answer_chain.invoke({
                "input": q,
                "chat_history": chat_history,
                "context": docs,
            })
            print(f"SimpleRAG: Generated answer: '{answer[:100]}...'")
            # create_stuff_documents_chain returns a string by default
            return {"answer": answer, "context": docs}

    return SimpleRAG()