app/services/retrieval.py

"""RAG benchmark store: ChromaDB (in-memory) + BGE-small embeddings.

Loads `app/benchmarks/seed_clauses.json` at warmup, embeds with sentence-transformers
BGE-small, and exposes `fetch_benchmark_matches()` returning BenchmarkMatch objects
that carry semantic similarity scores (0-1).
"""

from __future__ import annotations

import json
import logging
import threading
from pathlib import Path

from app.schemas import BenchmarkMatch, BenchmarkRef, ClauseType

logger = logging.getLogger(__name__)

_BENCHMARK_PATH = Path(__file__).parents[1] / "benchmarks" / "seed_clauses.json"
_BGE_MODEL = "BAAI/bge-small-en-v1.5"
_COLLECTION_NAME = "benchmark_clauses"

_warmup_lock = threading.Lock()
_warmed = False
_collection = None
_embedder = None
_by_type: dict[str, list[BenchmarkRef]] = {}


def _load_seed() -> list[BenchmarkRef]:
    if not _BENCHMARK_PATH.exists():
        logger.warning("seed clauses file not found at %s", _BENCHMARK_PATH)
        return []
    raw = json.loads(_BENCHMARK_PATH.read_text(encoding="utf-8"))
    return [BenchmarkRef(**item) for item in raw]


def warmup() -> None:
    global _warmed, _collection, _embedder, _by_type

    with _warmup_lock:
        if _warmed:
            return

        seed = _load_seed()
        for ref in seed:
            _by_type.setdefault(ref.clause_type, []).append(ref)

        logger.info("Loaded %d benchmark clauses into type index", len(seed))

        try:
            import chromadb
            from sentence_transformers import SentenceTransformer
        except Exception as exc:
            logger.warning(
                "Chroma or sentence-transformers unavailable (%s); type-only retrieval active",
                exc,
            )
            _warmed = True
            return

        try:
            _embedder = SentenceTransformer(_BGE_MODEL)
            client = chromadb.EphemeralClient()
            try:
                client.delete_collection(_COLLECTION_NAME)
            except Exception:
                pass
            _collection = client.create_collection(_COLLECTION_NAME)

            if seed:
                ids = [ref.id for ref in seed]
                docs = [ref.text for ref in seed]
                metas = [
                    {"clause_type": ref.clause_type, "doc_type": ref.doc_type}
                    for ref in seed
                ]
                embeddings = _embedder.encode(docs, normalize_embeddings=True).tolist()
                _collection.add(ids=ids, documents=docs, embeddings=embeddings, metadatas=metas)
                logger.info(
                    "ChromaDB vector store ready — %d clauses indexed with BGE-small embeddings",
                    len(seed),
                )
        except Exception as exc:
            logger.exception("Chroma initialization failed (%s); falling back to type-only", exc)
            _collection = None
            _embedder = None

        _warmed = True


def fetch_benchmark_matches(
    clause_text: str, clause_type: ClauseType, k: int = 2
) -> list[BenchmarkMatch]:
    """Return up to k BenchmarkMatch objects ranked by semantic similarity."""
    same_type = _by_type.get(clause_type, [])
    if not same_type:
        return []

    if _collection is None or _embedder is None:
        # Fallback: return type-only matches with zero similarity
        return [BenchmarkMatch(ref=ref, similarity=0.0) for ref in same_type[:k]]

    try:
        query_vec = _embedder.encode([clause_text], normalize_embeddings=True).tolist()
        n = min(max(k * 3, 6), len(same_type))
        result = _collection.query(
            query_embeddings=query_vec,
            n_results=n,
            where={"clause_type": clause_type},
        )
        hit_ids = result.get("ids", [[]])[0]
        distances = result.get("distances", [[]])[0]
        by_id: dict[str, BenchmarkRef] = {ref.id: ref for ref in same_type}

        matches: list[BenchmarkMatch] = []
        for hit_id, dist in zip(hit_ids, distances):
            if hit_id in by_id:
                # ChromaDB cosine distance: 0=identical, 2=opposite → similarity = 1 - dist/2
                sim = max(0.0, min(1.0, 1.0 - dist / 2.0))
                matches.append(BenchmarkMatch(ref=by_id[hit_id], similarity=sim))

        return matches[:k] if matches else [BenchmarkMatch(ref=r, similarity=0.0) for r in same_type[:k]]
    except Exception as exc:
        logger.warning("semantic retrieval failed (%s); returning type-only", exc)
        return [BenchmarkMatch(ref=ref, similarity=0.0) for ref in same_type[:k]]


# Legacy shim for any callers that still use the old signature
def fetch_benchmarks_by_text(clause_text: str, clause_type: ClauseType, k: int = 2):
    return [m.ref for m in fetch_benchmark_matches(clause_text, clause_type, k)]


def fetch_benchmarks(clause_type: ClauseType, k: int = 2):
    refs = _by_type.get(clause_type, [])
    return refs[:k]