src/nl_sql/schema_index/retriever.py

"""Build a `ContextBundle` for the LangGraph context_builder node.

Pipeline (per docs/02_architecture_v2.md §3 + §4):
1. Dense retrieve top-k schema chunks for `db_id` (Mistral embeddings).
2. FK-graph traversal up to `fk_hops` to add neighbour tables, capped by
   `table_budget`.
3. Dense retrieve top-k few-shot Q→SQL pairs for the same db_id.

Returns flat structures the prompt assembler can render directly. No prompt
text is built here — that lives in `agent/prompts/` once stage 4 lands.
"""

from __future__ import annotations

from collections import deque
from dataclasses import dataclass, field
from typing import Any, cast

from sqlalchemy.engine import Engine

from nl_sql.schema_index.indexer import FewShotHit, SchemaIndex, SchemaQueryHit
from nl_sql.schema_index.introspector import fetch_extended_samples


@dataclass(frozen=True, slots=True)
class ContextBundle:
    db_id: str
    question: str
    schema_hits: list[SchemaQueryHit]
    fk_neighbours: list[SchemaQueryHit]
    fewshots: list[FewShotHit]
    truncated: bool = False
    notes: list[str] = field(default_factory=list)
    extended_samples: dict[str, dict[str, tuple[Any, ...]]] | None = None
    """Per-difficulty sample mixture. Maps `table_name → col_name → tail
    sample values` (i.e. samples 4..N when chunks were built at
    sample_size=3). Set by `retrieve_context` when called with both
    `engine` and `extended_sample_size > primary_sample_size`. Rendered
    as an "additional sample values" appendix by `render_schema_block`.
    """

    @property
    def all_tables(self) -> list[str]:
        seen: list[str] = []
        for hit in (*self.schema_hits, *self.fk_neighbours):
            if hit.table_name and hit.table_name not in seen:
                seen.append(hit.table_name)
        return seen


def retrieve_context(
    index: SchemaIndex,
    question: str,
    *,
    db_id: str,
    schema_top_k: int = 5,
    fewshot_top_k: int = 3,
    fk_hops: int = 1,
    table_budget: int = 12,
    engine: Engine | None = None,
    primary_sample_size: int = 3,
    extended_sample_size: int = 0,
    cross_db_fewshot: bool = False,
) -> ContextBundle:
    """One call → schema cards + FK neighbours + fewshots, db_id-scoped.

    `fk_hops` of 1 (default) adds direct neighbours of every retrieved table;
    2 adds neighbours-of-neighbours, etc. `table_budget` caps the total table
    count (top-k hits + neighbours together) so the downstream prompt stays
    within the context window.

    Sample mixture (optional): when `engine` is provided and
    `extended_sample_size > primary_sample_size`, the bundle's
    `extended_samples` field is populated with the tail (rows
    primary..extended) of each retrieved table's column samples. Wired
    by `make_context_builder_node` via the registry's read-only engine.
    """
    schema_hits = (
        index.query_schema(question, db_id=db_id, top_k=schema_top_k) if schema_top_k > 0 else []
    )
    fewshots = (
        index.query_fewshots(
            question,
            db_id=db_id,
            top_k=fewshot_top_k,
            cross_db=cross_db_fewshot,
        )
        if fewshot_top_k > 0
        else []
    )

    notes: list[str] = []
    if not schema_hits:
        notes.append(f"no schema chunks indexed for db_id={db_id!r}")

    seed_tables = [h.table_name for h in schema_hits if h.table_name]
    fk_extra: list[SchemaQueryHit] = []
    truncated = False

    if seed_tables and fk_hops > 0:
        graph = index.fk_graph(db_id)
        seen = set(seed_tables)
        frontier: deque[tuple[str, int]] = deque((t, 0) for t in seed_tables)
        neighbour_order: list[str] = []
        while frontier:
            table, depth = frontier.popleft()
            if depth >= fk_hops:
                continue
            for nb in sorted(graph.get(table, ())):
                if nb in seen:
                    continue
                seen.add(nb)
                neighbour_order.append(nb)
                frontier.append((nb, depth + 1))

        # Total budget = seed tables already in schema_hits, plus neighbours.
        slots_left = max(0, table_budget - len(seed_tables))
        chosen_neighbours = neighbour_order[:slots_left]
        if len(neighbour_order) > slots_left:
            truncated = True
            notes.append(
                f"FK traversal yielded {len(neighbour_order)} neighbours, "
                f"capped to {slots_left} by table_budget={table_budget}"
            )

        if chosen_neighbours:
            fk_extra = _materialise_neighbours(index, db_id=db_id, names=chosen_neighbours)

    extended_samples: dict[str, dict[str, tuple[Any, ...]]] | None = None
    if engine is not None and extended_sample_size > primary_sample_size:
        all_tables = list(seed_tables)
        for hit in fk_extra:
            if hit.table_name and hit.table_name not in all_tables:
                all_tables.append(hit.table_name)
        if all_tables:
            extended_samples = fetch_extended_samples(
                engine,
                all_tables,
                primary_size=primary_sample_size,
                extended_size=extended_sample_size,
            )

    return ContextBundle(
        db_id=db_id,
        question=question,
        schema_hits=schema_hits,
        fk_neighbours=fk_extra,
        fewshots=fewshots,
        truncated=truncated,
        notes=notes,
        extended_samples=extended_samples,
    )


def _materialise_neighbours(
    index: SchemaIndex,
    *,
    db_id: str,
    names: list[str],
) -> list[SchemaQueryHit]:
    """Pull schema chunks for FK-graph neighbours by exact table_name match.

    Distance is set to ``inf`` to mark these as graph-derived (not dense
    retrieval). Order matches `names` (caller already prioritised).
    """
    if not names:
        return []
    where_clause = cast(
        Any,
        {"$and": [{"db_id": db_id}, {"table_name": {"$in": names}}]},
    )
    records = index.schema_collection.get(
        where=where_clause,
        include=["documents", "metadatas"],
    )
    ids = records.get("ids") or []
    docs = records.get("documents") or []
    metas = records.get("metadatas") or []

    by_name: dict[str, SchemaQueryHit] = {}
    for i, _id in enumerate(ids):
        meta = metas[i] if i < len(metas) and metas[i] else {}
        doc = docs[i] if i < len(docs) else ""
        table_name = str(meta.get("table_name") or "")
        if not table_name:
            continue
        by_name[table_name] = SchemaQueryHit(
            chunk_id=str(_id),
            table_name=table_name,
            db_id=str(meta.get("db_id") or ""),
            text=str(doc),
            distance=float("inf"),
            metadata=dict(meta),
        )
    return [by_name[n] for n in names if n in by_name]