src/engine.py

import time
import logging
from typing import List, Dict, Any, Tuple
from src.config import config
from src.retriever import HybridRetriever, Chunk
from src.reranker import Reranker
from src.query_expansion import QueryExpander
from src.hallucination_guard import HallucinationGuard
from src.llm_client import LocalLLMClient

logger = logging.getLogger("VigilantRAG")
logging.basicConfig(level=logging.INFO)

class VigilantRAGEngine:
    def __init__(self):
        self.retriever = HybridRetriever()
        self.reranker = Reranker()
        self.query_expander = QueryExpander()
        self.hallucination_guard = HallucinationGuard()
        self.llm = LocalLLMClient()

    def query(self, user_query: str, config_overrides: Dict[str, Any] = None) -> Dict[str, Any]:
        """
        Runs the self-correcting RAG pipeline on a user query.
        
        Returns:
            Dictionary containing final answer and detailed step-by-step execution telemetry.
        """
        start_time = time.time()
        
        # Apply temporary configuration overrides if any
        active_config = config
        if config_overrides:
            active_config = config.copy(update=config_overrides)
            
        telemetry = {
            "original_query": user_query,
            "final_query": user_query,
            "retrieval_attempts": [],
            "query_expansion_triggered": False,
            "generation_attempts": [],
            "hallucination_blocked_count": 0,
            "success": False,
            "execution_time_sec": 0.0
        }

        current_query = user_query
        retrieval_success = False
        top_chunks: List[Dict[str, Any]] = []

        # --- STAGE 1: Retrieval Loop (with Query Expansion) ---
        for attempt in range(active_config.max_query_expansion_retries + 1):
            attempt_log = {
                "attempt": attempt,
                "query": current_query,
                "dense_hits": 0,
                "sparse_hits": 0,
                "total_candidates": 0,
                "top_score": 0.0,
                "chunks": [],
                "status": "pending"
            }
            
             # 1. Hybrid Search
            candidates = self.retriever.search(
                current_query, 
                top_k_dense=active_config.top_k_dense, 
                top_k_sparse=active_config.top_k_sparse
            )
            
            # [SaaS Limit] If limit_document_id is set (Free plan: single doc search restriction)
            if config_overrides and config_overrides.get("limit_document_id"):
                doc_id_limit = config_overrides["limit_document_id"]
                candidates = [c for c in candidates if c["chunk"].doc_id == doc_id_limit]
            
            # Count source types for telemetry
            dense_count = sum(1 for c in candidates if "dense" in c["sources"])
            sparse_count = sum(1 for c in candidates if "sparse" in c["sources"])
            attempt_log["dense_hits"] = dense_count
            attempt_log["sparse_hits"] = sparse_count
            attempt_log["total_candidates"] = len(candidates)

            # 2. Re-ranking (Bypassed if use_reranking is False in Free plan)
            if config_overrides and config_overrides.get("use_reranking") is False:
                # Bypass Cross-Encoder. Direct sort chunks by dense score
                reranked = [
                    {
                        "chunk": item["chunk"],
                        "cross_score": item.get("dense_score", 0.0), # Fallback mockup score
                        "dense_score": item.get("dense_score", 0.0),
                        "sparse_score": item.get("sparse_score", 0.0),
                        "sources": item.get("sources", ["dense"])
                    }
                    for item in candidates[:active_config.top_n_final]
                ]
            else:
                # Standard Cross-Encoder Reranker execution
                reranked = self.reranker.rerank(current_query, candidates, top_n=active_config.top_n_final)

            # 3. Evaluate Retrieval Relevance
            top_score = reranked[0]["cross_score"] if reranked else -1.0
            attempt_log["top_score"] = top_score
            
            if top_score >= active_config.relevance_threshold:
                attempt_log["status"] = "relevant"
                telemetry["retrieval_attempts"].append(attempt_log)
                top_chunks = reranked
                retrieval_success = True
                break
            else:
                attempt_log["status"] = "irrelevant"
                telemetry["retrieval_attempts"].append(attempt_log)
                
                # If we have retries left, trigger Query Expansion
                if attempt < active_config.max_query_expansion_retries:
                    telemetry["query_expansion_triggered"] = True
                    logger.info(f"Retrieval top score {top_score:.3f} below threshold {active_config.relevance_threshold}. Triggering query expansion...")
                    
                    # Expand query using LLM (if model is initialized) or fall back to thesaurus
                    # To pass the LLM generator dynamically:
                    def llm_gen(p, temp, sys):
                        return self.llm.generate(p, temp, sys)
                    
                    current_query = self.query_expander.expand_llm(current_query, llm_gen)
                    logger.info(f"Query expanded to: '{current_query}'")
                else:
                    logger.warning("Query expansion retry limit reached. Proceeding with best available chunks.")
                    top_chunks = reranked

        telemetry["final_query"] = current_query
        
        # If absolutely no chunks were found (database is empty)
        if not top_chunks:
            telemetry["execution_time_sec"] = time.time() - start_time
            return {
                "answer": "I don't have any documents indexed yet. Please upload or ingest some documents first.",
                "telemetry": telemetry
            }

        # Combine texts from top 5 chunks to serve as context (Premise)
        context_text = "\n\n".join([f"Document [{item['chunk'].doc_title}]: {item['chunk'].text}" for item in top_chunks])

        # --- STAGE 2: Generation & Hallucination Guard Loop ---
        final_answer = ""
        current_temp = active_config.temperature_default
        
        # Default strict system prompt
        system_prompt = (
            "You are a factual assistant. Answer the query ONLY using the provided context. "
            "If the answer cannot be found in the context, say 'I cannot find the answer in the provided context'. "
            "Do not use external knowledge or make up facts."
        )

        for attempt in range(active_config.max_hallucination_retries + 1):
            gen_log = {
                "attempt": attempt,
                "temperature": current_temp,
                "system_prompt": system_prompt,
                "response_draft": "",
                "nli_scores": {},
                "status": "pending"
            }
            
            # Format Prompt
            generation_prompt = (
                f"Context information is below.\n"
                f"---------------------\n"
                f"{context_text}\n"
                f"---------------------\n"
                f"Given the context information above, answer the query.\n"
                f"Query: {telemetry['original_query']}\n"
                f"Answer:"
            )
            
            # 1. Generate Response
            draft = self.llm.generate(generation_prompt, temperature=current_temp, system_prompt=system_prompt)
            gen_log["response_draft"] = draft
            
            # 2. Audit Response via NLI Guard (Bypassed if use_nli_guard is False in Free plan)
            if config_overrides and config_overrides.get("use_nli_guard") is False:
                # Bypass DeBERTa NLI checks. Auto-verify the draft response
                is_hallucination = False
                nli_scores = {"entailment": 1.0, "contradiction": 0.0, "neutral": 0.0}
            else:
                # Standard DeBERTa validation audit execution
                is_hallucination, nli_scores = self.hallucination_guard.evaluate_response(context_text, draft)
            
            gen_log["nli_scores"] = nli_scores
            
            if not is_hallucination:
                final_answer = draft
                telemetry["success"] = True
                gen_log["status"] = "passed"
                telemetry["generation_attempts"].append(gen_log)
                break
            else:
                telemetry["hallucination_blocked_count"] += 1
                gen_log["status"] = "hallucinated"
                telemetry["generation_attempts"].append(gen_log)
                
                # Adjust generation constraints for next try
                if attempt < active_config.max_hallucination_retries:
                    # Switch to a higher temperature to explore alternative wordings,
                    # and tighten the system prompt to enforce strict context alignment.
                    current_temp = active_config.temperature_retry
                    system_prompt = (
                        "CRITICAL: You generated information not present in the context. "
                        "You must ONLY use the provided context. If the query cannot be answered exactly "
                        "using the context, say 'I cannot find the answer in the provided context'. "
                        "DO NOT make up facts. Answer strictly and factually."
                    )
                else:
                    logger.warning("Max hallucination retries reached. Returning the most factual fallback statement.")
                    final_answer = (
                        "I cannot answer this question with absolute certainty based on the provided documents. "
                        "My NLI guard detected a potential hallucination in the drafted response."
                    )

        telemetry["execution_time_sec"] = time.time() - start_time
        
        return {
            "answer": final_answer,
            "telemetry": telemetry
        }