utils/rag_engine.py

"""
RAG engine: retrieves relevant context from vector store,
builds a strict prompt, and queries the LLM.

Backend priority:
1. GROQ_API_KEY set → Groq API (fast, 100K tokens/day free tier)
2. USE_HF_LLM=1 + HF_TOKEN set → HuggingFace Inference API
3. Otherwise → Ollama (local)
"""
import os
import re
import logging
from typing import List, Dict, Any, Generator

from utils.vector_store import VectorStoreManager
from utils.memory import ConversationMemory

logger = logging.getLogger(__name__)

DEFAULT_OLLAMA_MODEL  = "llama3.2"
DEFAULT_GROQ_MODEL    = "llama-3.3-70b-versatile"
GROQ_FALLBACK_MODEL   = "llama-3.1-8b-instant"   # separate daily quota, higher limits
DEFAULT_HF_MODEL      = "meta-llama/Llama-3.1-8B-Instruct"

HF_TOKEN     = os.environ.get("HF_TOKEN") or os.environ.get("AgenticMultiModalRag_Token", "")
GROQ_API_KEY = os.environ.get("GROQ_API_KEY")
USE_HF_LLM   = os.environ.get("USE_HF_LLM", "").lower() in ("1", "true", "yes")

# Pick backend: Groq first (fast), then HF Inference (only if USE_HF_LLM=1), then Ollama
if GROQ_API_KEY:
    BACKEND = "groq"
elif USE_HF_LLM and HF_TOKEN:
    BACKEND = "hf"
else:
    BACKEND = "ollama"

OLLAMA_HOST = os.environ.get("OLLAMA_HOST", "http://localhost:11434")

SYSTEM_PROMPT = """You are a document assistant. Answer questions using ONLY the [CONTEXT] provided.
Rules:
1. If the context contains information relevant to the question, answer from it — even if only partially relevant.
2. Combine information from multiple context chunks if needed.
3. Only say "I DON'T KNOW" if the context truly contains NO relevant information at all.
4. Be concise. Cite source and page when available.
5. Do NOT make up information that is not in the context.
6. When answering questions about tables or structured data, apply ALL filter conditions from the question. Only include rows that match every condition — do not display or reference rows that do not match.
7. Give ONLY the final answer. Do NOT show reasoning steps, intermediate calculations, excluded rows, or any explanation of how you arrived at the answer.
8. Format key values, measurements, thresholds, and important requirements in **bold** (markdown). Use plain text for surrounding prose.
9. OCR-scanned documents may contain garbled characters in numbers and ranges (e.g., "50H,000" means "501-1,000"; "O" may mean "0"; "l" may mean "1"). Use numeric context and table structure to interpret such artifacts — do NOT use OCR garbling as a reason to say "I DON'T KNOW".
10. When the question asks about a specific person, only answer from context that explicitly mentions that person by name. Do NOT apply one person's details (experience, role, background) to another person. If the context does not contain information about the requested person, say "I DON'T KNOW".
"""

GENERAL_PROMPT = """You are a helpful AI assistant. Answer directly and concisely — final answer only, no reasoning steps.
If you don't know the answer, say so honestly.
"""

# Cosine distance threshold (0=identical, 2=opposite).
RELEVANCE_THRESHOLD = 1.2


def _make_hf_client():
    from huggingface_hub import InferenceClient
    return InferenceClient(token=HF_TOKEN)


def _make_groq_client():
    from groq import Groq
    return Groq(api_key=GROQ_API_KEY)


def _make_ollama_client():
    import ollama
    return ollama.Client(host=OLLAMA_HOST)


def _is_rate_limit(exc: Exception) -> bool:
    msg = str(exc).lower()
    return "429" in msg or "rate_limit" in msg or "rate limit" in msg


class RAGEngine:
    def __init__(self, vector_store: VectorStoreManager, model: str = None):
        self.vs = vector_store
        if BACKEND == "hf":
            self.model = os.environ.get("HF_MODEL", DEFAULT_HF_MODEL)
            self._client = _make_hf_client()
            logger.info(f"LLM backend: HuggingFace Inference ({self.model})")
        elif BACKEND == "groq":
            self.model = os.environ.get("GROQ_MODEL", DEFAULT_GROQ_MODEL)
            self._client = _make_groq_client()
            logger.info(f"LLM backend: Groq ({self.model})")
        else:
            self.model = model or os.environ.get("OLLAMA_MODEL", DEFAULT_OLLAMA_MODEL)
            self._client = _make_ollama_client()
            logger.info(f"LLM backend: Ollama ({self.model})")

    @staticmethod
    def _fix_ocr(text: str) -> str:
        """Fix common OCR garbling in scanned building-code PDFs before sending to LLM."""
        # "50H,000" → "501-1,000"  (range rows in IBC Table 1006.3.3 and similar tables)
        text = re.sub(r'(\d+)H,(\d{3})', lambda m: f"{m.group(1)}1-1,{m.group(2)}", text)
        # "81/z" / "91/z" / "81/i" / "91/i" → "8½" / "9½"  (PDF fraction "½" OCR'd as "1/z" or "1/i")
        text = re.sub(r'(\d+)1/[zi]', lambda m: f"{m.group(1)}½", text)
        return text

    @staticmethod
    def _is_question_only(text: str) -> bool:
        """True when a chunk is mostly a list of questions with no answer content."""
        lines = [l.strip() for l in text.splitlines() if l.strip()]
        if len(lines) < 2:
            return False
        questions = sum(1 for l in lines if l.endswith("?"))
        return questions / len(lines) >= 0.6

    @staticmethod
    def _is_non_content_chunk(text: str) -> bool:
        """True when a chunk is noise: URL lists, error logs, or markdown headers with no prose."""
        lines = [l.strip() for l in text.splitlines() if l.strip()]
        if not lines:
            return True
        url_or_header = sum(
            1 for l in lines
            if l.startswith("http") or l.startswith("chrome-extension://")
            or (l.startswith("#") and len(l.split()) <= 6)
        )
        if url_or_header / len(lines) >= 0.5:
            return True
        # Error/log lines (common patterns from Building_code samples.txt noise)
        noise_keywords = ("402 client error", "payment required", "depleted", "rate limit", "token limit")
        lowered = text.lower()
        if any(kw in lowered for kw in noise_keywords):
            return True
        return False

    def _build_context(self, results: List[Dict[str, Any]]) -> str:
        if not results:
            return "No relevant documents found."
        parts = []
        idx = 1
        for r in results:
            if self._is_question_only(r["text"]):
                continue
            if self._is_non_content_chunk(r["text"]):
                continue
            meta = r["metadata"]
            source   = meta.get("source", "unknown")
            page     = meta.get("page", "")
            doc_type = meta.get("type", "text")
            page_str = f", Page {page}" if page else ""
            type_str = f" [{doc_type}]" if doc_type != "text" else ""
            parts.append(f"[Doc {idx} — {source}{page_str}{type_str}]\n{self._fix_ocr(r['text'])}")
            idx += 1
        return "\n\n---\n\n".join(parts) if parts else "No relevant documents found."

    def _build_messages(self, question: str, context: str, memory: ConversationMemory, extra_context: str = ""):
        if extra_context:
            extra_section = (
                f"[STRUCTURED DATA from tables]\n{extra_context}\n\n"
                "Merge the structured data above with the document context below into one clean, readable answer. "
                "Bold key values and measurements.\n\n"
            )
        else:
            extra_section = ""
        user_message = (
            f"{extra_section}[CONTEXT from documents]\n{context}\n\n"
            f"[QUESTION]\n{question}\n\n"
            "Remember: Answer from the structured data and context above. If not found, say \"I DON'T KNOW\"."
        )
        return [
            {"role": "system", "content": SYSTEM_PROMPT},
            *memory.get_history_for_prompt(),
            {"role": "user", "content": user_message},
        ]

    def _is_off_topic(self, results: List[Dict[str, Any]]) -> bool:
        if not results:
            return True
        return all(r.get("distance", 1.0) > RELEVANCE_THRESHOLD for r in results)

    def _build_general_messages(self, question: str, memory: ConversationMemory):
        return [
            {"role": "system", "content": GENERAL_PROMPT},
            *memory.get_history_for_prompt(),
            {"role": "user", "content": question},
        ]

    def query(
        self,
        question: str,
        memory: ConversationMemory,
        n_results: int = 8,
        temperature: float = 0.0,
        stream: bool = False,
        source_filter: list = None,
        pre_fetched_results: list = None,
        extra_context: str = "",
    ) -> Generator[str, None, None]:
        results = pre_fetched_results if pre_fetched_results is not None else self.vs.query(question, n_results=n_results, source_filter=source_filter)

        if self._is_off_topic(results) and not extra_context:
            logger.info(f"Off-topic query (no relevant chunks): '{question[:60]}'")
            messages = self._build_general_messages(question, memory)
        else:
            context  = self._build_context(results)
            messages = self._build_messages(question, context, memory, extra_context)

        try:
            if BACKEND == "hf":
                yield from self._query_hf(messages, memory, question, temperature, stream)
            elif BACKEND == "groq":
                yield from self._query_groq(messages, memory, question, temperature, stream)
            else:
                yield from self._query_ollama(messages, memory, question, temperature, stream)
        except Exception as e:
            error_msg = f"Error: {str(e)}"
            logger.error(error_msg, exc_info=True)
            yield error_msg

    def _query_hf(self, messages, memory, question, temperature, stream):
        try:
            resp = self._client.chat_completion(
                model=self.model,
                messages=messages,
                temperature=max(temperature, 0.01),
                max_tokens=2048,
            )
            answer = resp.choices[0].message.content
            memory.add("user", question)
            memory.add("assistant", answer)
            yield answer
        except Exception as e:
            if "402" in str(e) or "payment" in str(e).lower() or "depleted" in str(e).lower():
                logger.warning("HF Inference credits depleted — falling back to Ollama")
                try:
                    yield from self._ollama_fallback(messages, memory, question, temperature)
                except Exception:
                    yield "⚠️ HF Inference credits depleted and Ollama is not available. Please set a GROQ_API_KEY."
            else:
                raise

    def _ollama_fallback(self, messages, memory, question, temperature):
        import ollama
        model = os.environ.get("OLLAMA_MODEL", DEFAULT_OLLAMA_MODEL)
        client = ollama.Client(host=OLLAMA_HOST)
        response = client.chat(
            model=model,
            messages=messages,
            options={"temperature": temperature},
        )
        answer = response["message"]["content"]
        memory.add("user", question)
        memory.add("assistant", answer)
        yield answer

    def _query_groq(self, messages, memory, question, temperature, stream):
        try:
            if stream:
                response_text = ""
                with self._client.chat.completions.create(
                    model=self.model,
                    messages=messages,
                    temperature=temperature,
                    stream=True,
                ) as stream_resp:
                    for chunk in stream_resp:
                        token = chunk.choices[0].delta.content or ""
                        response_text += token
                        yield token
                memory.add("user", question)
                memory.add("assistant", response_text)
            else:
                resp = self._client.chat.completions.create(
                    model=self.model,
                    messages=messages,
                    temperature=temperature,
                )
                answer = resp.choices[0].message.content
                memory.add("user", question)
                memory.add("assistant", answer)
                yield answer
        except Exception as e:
            if _is_rate_limit(e):
                # 1. Try a smaller Groq model — it has a separate, larger daily quota
                if self.model != GROQ_FALLBACK_MODEL:
                    try:
                        logger.warning(f"Groq rate limit on {self.model} — trying {GROQ_FALLBACK_MODEL}")
                        resp = self._client.chat.completions.create(
                            model=GROQ_FALLBACK_MODEL,
                            messages=messages,
                            temperature=temperature,
                        )
                        answer = resp.choices[0].message.content
                        memory.add("user", question)
                        memory.add("assistant", answer)
                        yield answer
                        return
                    except Exception as fb_exc:
                        logger.warning(f"Groq {GROQ_FALLBACK_MODEL} also failed: {fb_exc}")
                # 2. HF Inference
                if HF_TOKEN:
                    try:
                        logger.warning("Trying HF Inference fallback")
                        hf_client = _make_hf_client()
                        hf_model = os.environ.get("HF_MODEL", DEFAULT_HF_MODEL)
                        resp = hf_client.chat_completion(
                            model=hf_model,
                            messages=messages,
                            temperature=max(temperature, 0.01),
                            max_tokens=2048,
                        )
                        answer = resp.choices[0].message.content
                        memory.add("user", question)
                        memory.add("assistant", answer)
                        yield answer
                        return
                    except Exception as hf_exc:
                        logger.warning(f"HF Inference fallback failed: {hf_exc}")
                # 3. Ollama
                try:
                    yield from self._ollama_fallback(messages, memory, question, temperature)
                except Exception:
                    yield "⚠️ All LLM backends are currently unavailable (Groq daily limit reached, Ollama not running). Please try again later or upgrade at https://console.groq.com/settings/billing"
            else:
                raise

    def _query_ollama(self, messages, memory, question, temperature, stream):
        if stream:
            response_text = ""
            stream_resp = self._client.chat(
                model=self.model,
                messages=messages,
                stream=True,
                options={"temperature": temperature},
            )
            for chunk in stream_resp:
                token = chunk["message"]["content"]
                response_text += token
                yield token
            memory.add("user", question)
            memory.add("assistant", response_text)
        else:
            response = self._client.chat(
                model=self.model,
                messages=messages,
                options={"temperature": temperature},
            )
            answer = response["message"]["content"]
            memory.add("user", question)
            memory.add("assistant", answer)
            yield answer

    def list_available_models(self) -> List[str]:
        return [self.model]