app/services/rag.py

"""
ByteAstra — ChromaDB RAG service.

Handles:
  - Initialising the ChromaDB client and per-domain collections
  - Embedding queries using sentence-transformers
  - Retrieving the top-K most relevant chunks for a query
  - Indexing new documents (used by ingest.py)
"""
from __future__ import annotations
import logging
from dataclasses import dataclass
from functools import lru_cache

import chromadb
from chromadb import Settings as ChromaSettings
from chromadb.api.types import Documents, EmbeddingFunction, Embeddings
import torch
from transformers import AutoTokenizer, AutoModel
from app.config import get_settings
from app.schemas import Citation

logger = logging.getLogger(__name__)
settings = get_settings()


class TransformersEmbeddingFunction(EmbeddingFunction[Documents]):
    def __init__(self, model_name: str = "sentence-transformers/all-MiniLM-L6-v2", device: str = "cpu"):
        self.device = device
        if "/" not in model_name:
            model_name = f"sentence-transformers/{model_name}"
        
        # Pin embedding CPU threads to 1 to save the second core for the LLM
        torch.set_num_threads(1)
        
        try:
            import os
            os.environ["HF_HUB_OFFLINE"] = "1"
            os.environ["TRANSFORMERS_OFFLINE"] = "1"
            self.tokenizer = AutoTokenizer.from_pretrained(model_name, local_files_only=True)
            self.model = AutoModel.from_pretrained(model_name, local_files_only=True).to(self.device)
            logger.info("Loaded embedding model %s from local files.", model_name)
        except Exception:
            logger.info("Local files not found for %s. Enabling online download...", model_name)
            import os
            os.environ["HF_HUB_OFFLINE"] = "0"
            os.environ["TRANSFORMERS_OFFLINE"] = "0"
            self.tokenizer = AutoTokenizer.from_pretrained(model_name, local_files_only=False)
            self.model = AutoModel.from_pretrained(model_name, local_files_only=False).to(self.device)
            logger.info("Successfully downloaded and loaded embedding model %s.", model_name)

    def __call__(self, input: Documents) -> Embeddings:
        encoded_input = self.tokenizer(input, padding=True, truncation=True, return_tensors='pt').to(self.device)
        with torch.no_grad():
            model_output = self.model(**encoded_input)
        
        token_embeddings = model_output[0]
        attention_mask = encoded_input['attention_mask']
        input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
        sum_embeddings = torch.sum(token_embeddings * input_mask_expanded, 1)
        sum_mask = torch.clamp(input_mask_expanded.sum(1), min=1e-9)
        embeddings = sum_embeddings / sum_mask
        
        import torch.nn.functional as F
        embeddings = F.normalize(embeddings, p=2, dim=1)
        
        return embeddings.cpu().numpy().tolist()


@dataclass
class RetrievedChunk:
    chunk_id: str
    content: str
    source: str
    chapter: str | None
    section: str | None
    relevance_score: float  # cosine similarity (0–1)

    def to_citation(self) -> Citation:
        return Citation(
            chunk_id=self.chunk_id,
            source=self.source,
            chapter=self.chapter,
            section=self.section,
            relevance_score=self.relevance_score,
            content=self.content,
        )


@lru_cache(maxsize=1)
def _get_embedding_fn() -> TransformersEmbeddingFunction:
    logger.info("Loading embedding model: %s", settings.embedding_model)
    return TransformersEmbeddingFunction(model_name=settings.embedding_model, device="cpu")


@lru_cache(maxsize=1)
def _get_chroma_client() -> chromadb.ClientAPI:
    logger.info("Initialising ChromaDB at: %s", settings.chroma_persist_path)
    return chromadb.PersistentClient(
        path=settings.chroma_persist_path,
        settings=ChromaSettings(anonymized_telemetry=False),
    )


def get_or_create_collection(collection_name: str) -> chromadb.Collection:
    """Get or create a ChromaDB collection for a specific domain."""
    client = _get_chroma_client()
    return client.get_or_create_collection(
        name=collection_name,
        metadata={"hnsw:space": "cosine"},
    )


def retrieve(
    query: str,
    collection_name: str,
    top_k: int = 5,
    relevance_threshold: float = 0.45,
) -> list[RetrievedChunk]:
    """
    Retrieve the most relevant chunks from the domain's collection.
    Only returns chunks with similarity score >= relevance_threshold.

    Returns an empty list if nothing meets the threshold — the engine
    will treat this as an out-of-syllabus query.
    """
    import re
    collection = get_or_create_collection(collection_name)
    
    # Compute query embedding using the custom function
    embedding_fn = _get_embedding_fn()
    query_embeddings = embedding_fn([query])

    # Fetch candidates for lexical re-ranking (reduced for speed on CPU)
    candidate_k = max(top_k * 2, 6)
    results = collection.query(
        query_embeddings=query_embeddings,
        n_results=candidate_k,
        include=["documents", "metadatas", "distances"],
    )

    candidates: list[tuple[RetrievedChunk, float]] = []
    if not results["ids"] or not results["ids"][0]:
        return []

    # Extract significant query words for lexical boosting
    q_words = re.findall(r"\b[a-zA-Z]{3,}\b", query.lower())
    stop_words = {"how", "does", "that", "relate", "to", "what", "is", "are", "the", "and", "in", "of", "a", "an", "with", "about", "for", "its", "why", "who"}
    sig_words = [w for w in q_words if w not in stop_words]

    for i, chunk_id in enumerate(results["ids"][0]):
        # ChromaDB cosine distance → similarity: similarity = 1 - distance
        distance = results["distances"][0][i]
        similarity = 1.0 - distance
        if similarity < relevance_threshold:
            logger.debug("Chunk %s below threshold (%.3f < %.3f) — skipped", chunk_id, similarity, relevance_threshold)
            continue

        content = results["documents"][0][i]
        meta = results["metadatas"][0][i] or {}
        
        # Calculate lexical boost based on query term frequency/overlap
        boost = 0.0
        if sig_words:
            content_lower = content.lower()
            for w in sig_words:
                if re.search(r"\b" + re.escape(w) + r"\b", content_lower):
                    boost += 0.03
        
        boosted_score = similarity + boost
        
        chunk = RetrievedChunk(
            chunk_id=chunk_id,
            content=content,
            source=meta.get("source", "Unknown Source"),
            chapter=meta.get("chapter"),
            section=meta.get("section"),
            relevance_score=round(similarity, 4),
        )
        candidates.append((chunk, boosted_score))

    # Sort candidates by boosted score descending and take top_k
    candidates.sort(key=lambda x: x[1], reverse=True)
    selected_chunks = [x[0] for x in candidates[:top_k]]

    logger.info("Retrieved %d relevant chunks for query (candidates=%d, threshold=%.2f)", len(selected_chunks), len(candidates), relevance_threshold)
    return selected_chunks


def index_chunks(
    collection_name: str,
    chunks: list[dict],
) -> None:
    """
    Index a batch of text chunks into the domain's ChromaDB collection.

    Each chunk dict must have:
      - id: str  (unique, e.g. "ayurveda_charaka_ch1_000")
      - text: str
      - source: str
      - chapter: str | None
      - section: str | None
    """
    collection = get_or_create_collection(collection_name)

    ids = [c["id"] for c in chunks]
    documents = [c["text"] for c in chunks]
    metadatas = [
        {
            "source": c.get("source", ""),
            "chapter": c.get("chapter") or "",
            "section": c.get("section") or "",
        }
        for c in chunks
    ]

    # Compute embeddings manually to bypass ChromaDB internal conflicts
    embedding_fn = _get_embedding_fn()
    embeddings = embedding_fn(documents)

    collection.upsert(ids=ids, embeddings=embeddings, documents=documents, metadatas=metadatas)
    logger.info("Indexed %d chunks into collection '%s'", len(chunks), collection_name)