app/pilot_retrieve.py

"""Paediatrics retrieval over SQLite FTS (HF Spaces / no-OpenSearch pilot)."""
from __future__ import annotations

import logging

from .config import settings
from .paeds_retrieve import (
    enrich_with_relationships,
    related_premium_codes_for_primaries,
    relationship_meta,
    rrf_fuse,
)
from .sob_grounding import filter_current_sob_only, is_current_sob
from .specialty_scope import (
    filter_catalog_by_prefixes,
    is_paediatrics_26_premium,
    is_paediatrics_26_primary,
    paeds_premium_allowlist_codes,
)
from . import sqlite_store

logger = logging.getLogger(__name__)


def get_code(code: str) -> dict | None:
    return sqlite_store.get_code(code)


def fetch_all_paediatrics_26_primaries() -> list[dict]:
    return [
        enrich_with_relationships(d)
        for d in sqlite_store.all_current_docs()
        if is_paediatrics_26_primary(d)
    ]


def hybrid_search_paediatrics(
    query_text: str,
    query_vector: list[float] | None = None,  # unused — BM25-only pilot
    *,
    prefixes: list[str] | None = None,
    top_k: int | None = None,
) -> list[dict]:
    del query_vector  # pilot is lexical-only
    k = top_k or max(settings.hybrid_top_k, 40)
    hits = sqlite_store.fts_search(query_text, top_k=max(k * 3, 80))
    hits = [d for d in hits if is_paediatrics_26_primary(d)]
    if prefixes:
        hits = filter_catalog_by_prefixes(hits, prefixes)
    return [enrich_with_relationships(d) for d in hits[:k]]


def rank_full_paeds_catalog(
    query_text: str,
    query_vector: list[float] | None = None,
    *,
    prefixes: list[str] | None = None,
) -> list[dict]:
    catalog = fetch_all_paediatrics_26_primaries()
    catalog = filter_catalog_by_prefixes(catalog, prefixes or [])
    if not catalog and prefixes:
        catalog = fetch_all_paediatrics_26_primaries()

    top_k = max(len(catalog), settings.hybrid_top_k, 50)
    hybrid_hits = hybrid_search_paediatrics(
        query_text,
        query_vector,
        prefixes=prefixes,
        top_k=top_k,
    )
    by_code = {d["billing_code"]: d for d in hybrid_hits}

    relationship_boost: list[dict] = []
    for hit in hybrid_hits[:12]:
        for mirror in hit.get("relationship_mirrors") or []:
            if prefixes and mirror[:1] not in {p.upper() for p in prefixes}:
                continue
            if mirror in by_code:
                continue
            doc = get_code(mirror)
            if doc and is_paediatrics_26_primary(doc):
                relationship_boost.append(enrich_with_relationships(doc))

    fused = rrf_fuse(
        [hybrid_hits, relationship_boost] if relationship_boost else [hybrid_hits]
    )
    fused_codes = {d["billing_code"] for d in fused}
    for doc in catalog:
        code = doc["billing_code"]
        if code not in fused_codes:
            fused.append(enrich_with_relationships(doc))
            fused_codes.add(code)
    logger.info(
        "SQLite paeds catalog: hybrid=%d catalog=%d fused=%d",
        len(hybrid_hits),
        len(catalog),
        len(fused),
    )
    return fused


def hybrid_premium_search_paediatrics(
    query_text: str,
    query_vector: list[float] | None = None,
    *,
    primary_codes: list[str],
    top_k: int = 24,
) -> list[dict]:
    del query_vector
    premium_query = (
        f"Paediatrics specialty 26. {query_text} "
        "premium add-on immunization injection special visit chronic disease "
        "OTHER PREMIUMS after hours procedure E409 E410"
    )
    hits = sqlite_store.fts_search(premium_query, top_k=top_k * 3)
    hits = filter_current_sob_only(hits)
    hits = [h for h in hits if is_paediatrics_26_premium(h)]

    related = related_premium_codes_for_primaries(primary_codes)
    # Also force-include the static Paediatrics premium allowlist so BM25 misses
    # (e.g. G372) still reach the LLM pool.
    related = list(dict.fromkeys([*related, *sorted(paeds_premium_allowlist_codes())]))

    by_code = {h["billing_code"]: enrich_with_relationships(h) for h in hits}
    for code in related:
        if code in by_code:
            continue
        doc = get_code(code)
        if not is_current_sob(doc):
            continue
        if not is_paediatrics_26_premium(doc):
            continue
        by_code[code] = enrich_with_relationships(doc)

    ordered: list[dict] = []
    seen: set[str] = set()
    for code in related:
        if code in by_code and code not in seen:
            ordered.append(by_code[code])
            seen.add(code)
    for code, doc in by_code.items():
        if code not in seen:
            ordered.append(doc)
            seen.add(code)
    return ordered[: max(top_k, len(related))]


# Re-export helpers used by analyze post-processing.
__all__ = [
    "enrich_with_relationships",
    "get_code",
    "hybrid_premium_search_paediatrics",
    "rank_full_paeds_catalog",
    "relationship_meta",
]