src/core/rag_engine.py

"""
Human+ RAG Engine
Loads knowledge base (.md files) → chunks → FAISS vector store.
Cached with st.cache_resource so it only runs once per deployment.
"""

import streamlit as st
from pathlib import Path
from typing import Optional

from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_community.vectorstores import FAISS
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_core.documents import Document


# ── Config ───────────────────────────────────────────────────────────
KNOWLEDGE_DIR   = Path("/app/knowledge")
KNOWLEDGE_DIR_DEV = Path("knowledge")          # fallback saat dev lokal

EMBEDDING_MODEL = "sentence-transformers/all-mpnet-base-v2"

CHUNK_SIZE      = 800
CHUNK_OVERLAP   = 150
SEPARATORS      = ["\n\n## ", "\n\n### ", "\n\n", "\n", " ", "."]

# Retriever config (MMR = Maximal Marginal Relevance — kurangi redundansi)
RETRIEVER_K        = 8
RETRIEVER_FETCH_K  = 20
RETRIEVER_LAMBDA   = 0.7   # 0 = diversity, 1 = relevance


class RAGEngineError(Exception):
    """Raised when RAG engine fails to load or retrieve."""
    pass


@st.cache_resource(show_spinner=False, ttl=7200)
def get_vector_store() -> FAISS:
    """
    Load knowledge base dan build FAISS vector store.
    Cached selama 2 jam — auto-reload jika knowledge base diupdate.

    Returns:
        FAISS vector store yang siap diquery

    Raises:
        RAGEngineError: jika folder atau dokumen tidak ditemukan
    """
    knowledge_dir = _resolve_knowledge_dir()
    documents     = _load_markdown_documents(knowledge_dir)
    chunks        = _split_documents(documents)
    embeddings    = _build_embeddings()
    vector_store  = FAISS.from_documents(chunks, embeddings)
    return vector_store


def retrieve_context(query: str, vector_store: Optional[FAISS] = None) -> str:
    """
    Retrieve relevant knowledge chunks untuk query biomarker/protocol.

    Args:
        query: teks query (biasanya extracted biomarkers)
        vector_store: optional — jika None, akan di-load otomatis

    Returns:
        context string yang siap dimasukkan ke prompt
    """
    if vector_store is None:
        vector_store = get_vector_store()

    retriever = vector_store.as_retriever(
        search_type="mmr",
        search_kwargs={
            "k":           RETRIEVER_K,
            "fetch_k":     RETRIEVER_FETCH_K,
            "lambda_mult": RETRIEVER_LAMBDA,
        },
    )

    try:
        docs = retriever.invoke(query)
    except Exception as e:
        raise RAGEngineError(f"Retrieval gagal: {e}") from e

    context = "\n\n".join(
        f"[Sumber: {doc.metadata.get('source', 'unknown')}]\n{doc.page_content}"
        for doc in docs
    )
    return context


def get_knowledge_stats() -> dict:
    """Return stats tentang knowledge base yang ter-load (untuk debugging)."""
    try:
        knowledge_dir = _resolve_knowledge_dir()
        files = list(knowledge_dir.glob("*.md"))
        return {
            "ok":       True,
            "dir":      str(knowledge_dir),
            "n_files":  len(files),
            "filenames": [f.name for f in files],
        }
    except RAGEngineError as e:
        return {"ok": False, "error": str(e)}


# ── Private helpers ──────────────────────────────────────────────────

def _resolve_knowledge_dir() -> Path:
    """Cari knowledge directory — Docker path atau dev local."""
    for candidate in [KNOWLEDGE_DIR, KNOWLEDGE_DIR_DEV]:
        if candidate.exists() and candidate.is_dir():
            return candidate
    raise RAGEngineError(
        f"Folder knowledge tidak ditemukan. "
        f"Dicari di: {KNOWLEDGE_DIR} dan {KNOWLEDGE_DIR_DEV}"
    )


def _load_markdown_documents(knowledge_dir: Path) -> list[Document]:
    """Load semua file .md dari knowledge_dir sebagai LangChain Documents."""
    documents = []
    for file_path in sorted(knowledge_dir.glob("*.md")):
        try:
            content = file_path.read_text(encoding="utf-8").strip()
            if content:
                documents.append(
                    Document(
                        page_content=content,
                        metadata={
                            "source":   file_path.name,
                            "filepath": str(file_path),
                        },
                    )
                )
        except Exception as e:
            # Log tapi jangan stop — satu file gagal tidak perlu matikan semua
            st.warning(f"⚠️ Gagal membaca {file_path.name}: {e}")

    if not documents:
        raise RAGEngineError(
            f"Tidak ada dokumen .md yang berhasil dibaca dari {knowledge_dir}"
        )

    return documents


def _split_documents(documents: list[Document]) -> list[Document]:
    """Split documents menjadi chunks yang optimal untuk embedding."""
    splitter = RecursiveCharacterTextSplitter(
        chunk_size=CHUNK_SIZE,
        chunk_overlap=CHUNK_OVERLAP,
        separators=SEPARATORS,
    )
    return splitter.split_documents(documents)


def _build_embeddings() -> HuggingFaceEmbeddings:
    """Build HuggingFace embeddings model."""
    return HuggingFaceEmbeddings(
        model_name=EMBEDDING_MODEL,
        model_kwargs={"device": "cpu"},
        encode_kwargs={"normalize_embeddings": True},
    )