rag_engine.py

"""
rag_engine.py — Retrieval-Augmented Generation pipeline.

Embeds PDF chunks with sentence-transformers, stores them in a FAISS index,
and answers caregiver questions using a local llama.cpp model (Qwen 2.5 14B).
"""

from __future__ import annotations

import json
import os
from pathlib import Path
from typing import List, Tuple, Optional

from config import (
    EMBEDDING_MODEL,
    TOP_K_RETRIEVAL,
    TEMPERATURE,
    MAX_TOKENS,
    CONTEXT_SIZE,
    FAISS_INDEX_PATH,
    CHUNKS_JSON_PATH,
    LM_STUDIO_URL,  # kept for backwards-compat with local LM Studio setups
    get_model_path,
)
from pdf_loader import (
    ingest_pdf_file,
    load_chunks_from_cache,
    get_indexed_sources,
)


# ── System prompt ─────────────────────────────────────────────────────────────

SYSTEM_PROMPT = (
    "You are AidAiLine, a medical caregiver assistant.\n"
    "Each user message includes up to two context blocks for the ACTIVE care profile:\n"
    "1. Care profile & tracker data (demographics, insurance, PCP, medications, "
    "allergies, food preferences, appointments)\n"
    "2. Excerpts retrieved from uploaded medical documents for that profile\n\n"
    "Answer using ONLY the provided context. Use tracker/profile data for meds, "
    "allergies, appointments, and demographics. Use document excerpts for clinical "
    "details found in uploaded PDFs.\n"
    "If the answer is not in either source, say clearly that you could not find it "
    "in the profile or documents.\n"
    "When asked about totals, amounts, or lists: find ALL matching items from "
    "the provided context, not just the first or most recent one. List each item "
    "with its details. If multiple items match, show them all.\n"
    "Be concise and specific — cite dates, names, and numbers when available."
)

_MODEL_DOWNLOAD_MSG = """\
⚠️  Local AI model not found at: {model_path}

The Documents chat uses a hosted AI model (Hugging Face Inference API)
when no local GGUF is available. The hosted model is Qwen 2.5 7B Instruct
running on Hugging Face infrastructure.

To use a fully local model instead:
1. Place a GGUF at the path above
2. Install llama-cpp-python
3. Restart this app
"""


# Hosted model used when no local GGUF is available.
# Small enough for the free inference tier, good enough for short answers.
_INFERENCE_MODEL = "Qwen/Qwen2.5-7B-Instruct"


# ── Lazy-loaded globals ───────────────────────────────────────────────────────

_embedder = None
_faiss_index = None
_faiss_chunks: List[dict] = []
_llama_model = None


def _get_embedder():
    global _embedder
    if _embedder is None:
        from sentence_transformers import SentenceTransformer
        _embedder = SentenceTransformer(EMBEDDING_MODEL)
    return _embedder


def _get_faiss_index():
    """Load or create the FAISS index."""
    global _faiss_index, _faiss_chunks
    import numpy as np

    try:
        import faiss
    except ImportError:
        raise ImportError("faiss-cpu is not installed. Run: pip install faiss-cpu")

    if _faiss_index is not None:
        return _faiss_index, _faiss_chunks

    if FAISS_INDEX_PATH.exists() and CHUNKS_JSON_PATH.exists():
        _faiss_index = faiss.read_index(str(FAISS_INDEX_PATH))
        with open(CHUNKS_JSON_PATH, "r", encoding="utf-8") as f:
            _faiss_chunks = json.load(f)
    else:
        # Start with an empty flat index (384-dim for all-MiniLM-L6-v2)
        _faiss_index = faiss.IndexFlatL2(384)
        _faiss_chunks = []

    return _faiss_index, _faiss_chunks


def _save_faiss_index():
    """Flush FAISS index to disk."""
    try:
        import faiss
        FAISS_INDEX_PATH.parent.mkdir(parents=True, exist_ok=True)
        faiss.write_index(_faiss_index, str(FAISS_INDEX_PATH))
    except Exception as e:
        print(f"[rag_engine] Warning: could not save FAISS index: {e}")


def _get_inference_client():
    """
    Lazy-create the HF Inference API client for the hosted Qwen model.
    Returns None if huggingface_hub isn't installed (caller handles gracefully).
    """
    try:
        from huggingface_hub import InferenceClient
    except ImportError:
        return None
    token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGINGFACE_TOKEN")
    return InferenceClient(model=_INFERENCE_MODEL, token=token)


def _check_lm_studio() -> bool:
    """Check if LM Studio is reachable at the configured URL."""
    model_path = get_model_path()
    if not Path(model_path).exists():
        return False
    try:
        import urllib.request
        import json
        base = LM_STUDIO_URL.rsplit("/v1/", 1)[0] + "/v1/models"
        req = urllib.request.Request(base, method="GET")
        with urllib.request.urlopen(req, timeout=3) as resp:
            return resp.status == 200
    except Exception:
        return False


# ── Profile context (for RAG prompt) ─────────────────────────────────────────

def _profile_has_indexed_docs(profile_id: str) -> bool:
    """True if the FAISS index contains at least one chunk for profile_id."""
    if not profile_id:
        return False
    try:
        _, chunks = _get_faiss_index()
        return any(c.get("profile_id") == profile_id for c in chunks)
    except Exception:
        return False


def build_profile_context(profile_id: str) -> str:
    """
    Serialize the active care profile plus scoped tracker data for the LLM prompt.
    Returns empty string when profile_id is missing or not found.
    """
    if not profile_id:
        return ""

    import profiles
    import med_tracker
    import food_tracker
    import appointment_tracker
    from doc_forms import _resolve_emergency_contacts

    p = profiles.get_profile(profile_id)
    if not p:
        return ""

    def _line(label: str, value: str) -> None:
        v = (value or "").strip()
        if v:
            lines.append(f"{label}: {v}")

    lines: List[str] = ["=== ACTIVE CARE PROFILE ==="]
    _line("Patient", profiles.profile_display_name(p))
    if (p.get("label") or "").strip():
        _line("Profile label", p.get("label", ""))
    _line("DOB", p.get("dob", ""))
    _line("Phone", p.get("phone_number", ""))
    _line("Email", p.get("email", ""))

    addr = profiles.compose_address(*profiles.address_parts(p))
    if addr:
        _line("Address", addr.replace("\n", ", "))

    _line("Care mode", p.get("care_mode", ""))
    if (p.get("care_mode") or "").strip() == "Managing someone else's care":
        _line("Caregiver", p.get("caregiver_name", ""))
        _line("Caregiver phone", p.get("caregiver_phone", ""))
        _line("Caregiver email", p.get("caregiver_email", ""))

    _line("Insurance provider", p.get("insurance_provider", ""))
    _line("Policy ID", p.get("policy_id", ""))
    _line("Group ID", p.get("group_id", ""))

    _line("Primary care doctor", p.get("pcp_name", ""))
    _line("PCP clinic", p.get("pcp_clinic", ""))
    _line("PCP phone", p.get("pcp_phone", ""))
    pcp_addr = profiles.compose_address(*profiles.address_parts(p, pcp=True))
    if pcp_addr:
        _line("PCP address", pcp_addr.replace("\n", ", "))

    tracked = profiles.get_tracked_symptoms(profile_id)
    if tracked:
        lines.append("")
        lines.append("Tracked symptoms (ongoing):")
        for s in tracked:
            lines.append(f"  • {s}")

    contacts = _resolve_emergency_contacts(p)
    if contacts:
        lines.append("")
        lines.append("Emergency contacts:")
        for i, c in enumerate(contacts, start=1):
            parts = [c.get("name") or "_(name not set)_"]
            if c.get("relationship"):
                parts.append(f"({c['relationship']})")
            if c.get("phone"):
                parts.append(f"phone {c['phone']}")
            lines.append(f"  {i}. {' '.join(parts)}")

    lines.append("")
    lines.append("=== CURRENT MEDICATIONS ===")
    current_meds = med_tracker.get_medications(filter="current", profile_id=profile_id)
    if current_meds:
        for m in current_meds:
            bits = [m.get("name") or "Unnamed"]
            if m.get("dosage"):
                bits.append(m["dosage"])
            if m.get("frequency"):
                bits.append(m["frequency"])
            if m.get("category"):
                bits.append(f"[{m['category']}]")
            line = " • " + " — ".join(bits)
            if m.get("side_effects"):
                line += f" (side effects: {m['side_effects']})"
            if m.get("personal_notes"):
                line += f" (notes: {m['personal_notes']})"
            lines.append(line)
    else:
        lines.append("(None recorded)")

    lines.append("")
    lines.append("=== ALLERGIES & FOOD ===")
    food = food_tracker.get_food(profile_id=profile_id)
    for heading, key in (
        ("Allergies", "allergies"),
        ("Preferred foods", "liked_foods"),
        ("Food aversions", "disliked_foods"),
    ):
        entries = food.get(key) or []
        if entries:
            names = ", ".join(e.get("name", "") for e in entries if e.get("name"))
            lines.append(f"{heading}: {names}")
        else:
            lines.append(f"{heading}: (none recorded)")

    lines.append("")
    lines.append("=== UPCOMING APPOINTMENTS ===")
    appts = appointment_tracker.get_all_appointments(
        include_past=False, profile_id=profile_id,
    )
    if appts:
        for a in appts[:12]:
            title = a.get("title") or a.get("provider") or "Appointment"
            when = " ".join(
                x for x in (a.get("date", ""), a.get("time", "")) if x
            )
            where = (a.get("location") or "").strip()
            row = f" • {when} — {title}"
            if where:
                row += f" @ {where}"
            lines.append(row)
        if len(appts) > 12:
            lines.append(f"  … and {len(appts) - 12} more")
    else:
        lines.append("(None scheduled)")

    return "\n".join(lines)


# ── Public API ────────────────────────────────────────────────────────────────

def reset_engine():
    """Force reload of index (call after settings change)."""
    global _faiss_index, _faiss_chunks
    _faiss_index = None
    _faiss_chunks = []


def index_document(pdf_path: str | Path, profile_id: str = "") -> str:
    """
    Parse, chunk, embed, and index a PDF file.
    Returns a status string suitable for display.
    """
    global _faiss_index, _faiss_chunks

    try:
        import numpy as np
        import faiss

        # 1. Parse & chunk
        all_chunks = ingest_pdf_file(pdf_path, profile_id=profile_id)

        # 2. Re-embed all chunks from scratch (simplest correctness strategy)
        embedder = _get_embedder()
        texts = [c["text"] for c in all_chunks]

        if not texts:
            return "⚠️ No text could be extracted from that PDF."

        embeddings = embedder.encode(texts, show_progress_bar=False, normalize_embeddings=True)
        embeddings = embeddings.astype("float32")

        # 3. Rebuild index
        dim = embeddings.shape[1]
        _faiss_index = faiss.IndexFlatIP(dim)  # Inner-product (cosine after normalize)
        _faiss_index.add(embeddings)
        _faiss_chunks = all_chunks

        # 4. Persist
        _save_faiss_index()
        with open(CHUNKS_JSON_PATH, "w", encoding="utf-8") as f:
            json.dump(all_chunks, f, indent=2, ensure_ascii=False)

        sources = sorted({c["source"] for c in all_chunks})
        return (
            f"✅ Indexed {len(all_chunks)} chunks from {len(sources)} document(s):\n"
            + "\n".join(f"  • {s}" for s in sources)
        )

    except Exception as e:
        return f"❌ Error indexing document: {e}"


def retrieve(query: str, top_k: int = TOP_K_RETRIEVAL, profile_id: str = "") -> List[dict]:
    """Return top-k most relevant chunks for the query, optionally scoped to profile."""
    import numpy as np

    index, chunks = _get_faiss_index()
    if index.ntotal == 0:
        return []

    embedder = _get_embedder()
    q_vec = embedder.encode([query], normalize_embeddings=True).astype("float32")
    search_k = min(index.ntotal, max(top_k * 10, top_k)) if profile_id else min(top_k, index.ntotal)
    distances, indices = index.search(q_vec, search_k)

    results = []
    for dist, idx in zip(distances[0], indices[0]):
        if idx < 0 or idx >= len(chunks):
            continue
        chunk = chunks[idx]
        if profile_id and chunk.get("profile_id") != profile_id:
            continue
        results.append({**chunk, "score": float(dist)})
        if len(results) >= top_k:
            break
    return results


def answer_question(question: str, profile_id: str = "") -> Tuple[str, List[dict]]:
    """
    Full RAG pipeline: profile context + retrieve → build prompt → generate answer.

    Returns (answer_text, source_chunks).
    """
    profile_text = build_profile_context(profile_id) if profile_id else ""
    has_profile_docs = _profile_has_indexed_docs(profile_id)

    if not profile_text and not has_profile_docs:
        return (
            "📂 No care profile data or indexed documents are available yet.\n\n"
            "Sign in with a care profile in Settings & Profiles, add profile details, "
            "and/or upload PDFs in the Documents tab.",
            [],
        )

    # Retrieve document excerpts when this profile has indexed docs
    context_chunks: List[dict] = []
    if has_profile_docs:
        context_chunks = retrieve(question, profile_id=profile_id)

    if not profile_text and not context_chunks:
        return (
            "I couldn't find that in your uploaded documents, and no active profile "
            "data is loaded.",
            [],
        )

    user_blocks: List[str] = []
    if profile_text:
        user_blocks.append(
            "Care profile & tracker data (Settings, Medications, Food, Appointments):\n"
            + profile_text
        )
    if context_chunks:
        context_text = "\n\n---\n\n".join(
            f"[Source: {c['source']}, Page {c['page']}]\n{c['text']}"
            for c in context_chunks
        )
        user_blocks.append(f"Retrieved document excerpts:\n{context_text}")
    elif has_profile_docs:
        user_blocks.append(
            "Retrieved document excerpts:\n"
            "(No closely matching passages in uploaded documents for this question.)"
        )
    else:
        user_blocks.append(
            "Retrieved document excerpts:\n"
            "(No documents indexed for this profile yet.)"
        )

    user_blocks.append(f"Question: {question}")
    user_message = "\n\n".join(user_blocks)

    client = _get_inference_client()
    if client is None:
        return (
            "⚠️ huggingface-hub is not installed.\n\n"
            "Install it with: `pip install huggingface-hub`",
            context_chunks,
        )

    try:
        response = client.chat_completion(
            messages=[
                {"role": "system", "content": SYSTEM_PROMPT},
                {"role": "user", "content": user_message},
            ],
            max_tokens=MAX_TOKENS,
            temperature=TEMPERATURE,
        )
        answer = response.choices[0].message.content.strip()
        return answer, context_chunks

    except Exception as e:
        return f"❌ Error calling hosted model: {e}", context_chunks


def delete_document(source_name: str) -> str:
    """Remove a document and all its chunks from the index.
    
    Returns a status string suitable for display.
    """
    global _faiss_index, _faiss_chunks
    import numpy as np
    import faiss

    try:
        index, chunks = _get_faiss_index()
        before = len(chunks)
        
        # Filter out chunks from this source
        _faiss_chunks = [c for c in chunks if c.get("source") != source_name]
        removed = before - len(_faiss_chunks)
        
        if removed == 0:
            return f"⚠️ No chunks found for: {source_name}"
        
        # Rebuild index from scratch
        if _faiss_chunks:
            embedder = _get_embedder()
            texts = [c["text"] for c in _faiss_chunks]
            embeddings = embedder.encode(texts, show_progress_bar=False, normalize_embeddings=True)
            embeddings = embeddings.astype("float32")
            dim = embeddings.shape[1]
            _faiss_index = faiss.IndexFlatIP(dim)
            _faiss_index.add(embeddings)
        else:
            _faiss_index = faiss.IndexFlatIP(384)  # Fresh empty index
        
        # Persist
        _save_faiss_index()
        with open(CHUNKS_JSON_PATH, "w", encoding="utf-8") as f:
            json.dump(_faiss_chunks, f, indent=2, ensure_ascii=False)
        
        # Also delete the source file
        doc_path = Path(source_name)
        if doc_path.is_absolute() and doc_path.exists():
            doc_path.unlink()
        
        return f"✅ Deleted {removed} chunks from: {Path(source_name).name}"
    
    except Exception as e:
        return f"❌ Error deleting document: {e}"


def get_document_status(profile_id: str = "") -> dict:
    """Return info about the current index and model state."""
    model_path = get_model_path()
    model_exists = Path(model_path).exists()

    try:
        index, chunks = _get_faiss_index()
        if profile_id:
            chunks = [c for c in chunks if c.get("profile_id") == profile_id]
        sources = sorted({c["source"] for c in chunks}) if chunks else []
        total_chunks = len(chunks) if profile_id else index.ntotal
    except Exception:
        sources = []
        total_chunks = 0

    return {
        "model_found": model_exists,
        "model_path": str(model_path),
        "total_chunks": total_chunks,
        "indexed_sources": sources,
        "num_documents": len(sources),
    }


def delete_for_profile(profile_id: str) -> int:
    """Remove all indexed chunks belonging to profile_id. Returns count removed."""
    global _faiss_index, _faiss_chunks
    import faiss

    try:
        index, chunks = _get_faiss_index()
        before = len(chunks)
        _faiss_chunks = [c for c in chunks if c.get("profile_id") != profile_id]
        removed = before - len(_faiss_chunks)
        if removed == 0:
            return 0

        if _faiss_chunks:
            embedder = _get_embedder()
            texts = [c["text"] for c in _faiss_chunks]
            embeddings = embedder.encode(texts, show_progress_bar=False, normalize_embeddings=True)
            embeddings = embeddings.astype("float32")
            dim = embeddings.shape[1]
            _faiss_index = faiss.IndexFlatIP(dim)
            _faiss_index.add(embeddings)
        else:
            _faiss_index = faiss.IndexFlatIP(384)

        _save_faiss_index()
        with open(CHUNKS_JSON_PATH, "w", encoding="utf-8") as f:
            json.dump(_faiss_chunks, f, indent=2, ensure_ascii=False)
        return removed
    except Exception:
        return 0