src/nl_sql/api/main.py

"""FastAPI surface for the NL→SQL Assistant.

Endpoints:
    GET  /healthz        — liveness probe + provider configuration snapshot.
    GET  /readyz         — readiness probe (Chroma + DB registry reachable).
    GET  /databases      — list registered DBs + table counts.
    POST /ask            — translate question to SQL, execute, return result.
    GET  /eval/latest    — metadata of the latest committed eval report.

Auth:
    Set ``NL_SQL_API_KEY`` (env, .env, or settings). When set, every request
    to /ask and /databases must include ``X-API-Key`` matching it. /healthz
    and /readyz are always open for orchestrator probes.

Rate limit:
    In-process token bucket per API key (60 req/min default). No external
    Redis — this is a single-replica portfolio demo, not a fleet service.
"""

from __future__ import annotations

import logging
import os
import secrets
import time
import uuid
from collections import defaultdict, deque
from functools import lru_cache
from typing import Any, NamedTuple

from fastapi import Depends, FastAPI, Header, HTTPException, Request, status
from pydantic import BaseModel, Field

from nl_sql import __version__
from nl_sql.agent.graph import (
    PipelineConfig,
    PipelineRunResult,
    build_pipeline,
    run_pipeline,
)
from nl_sql.config import Settings, get_settings
from nl_sql.db.registry import DatabaseRegistry, get_default_registry
from nl_sql.llm.cache import CachingEmbeddingProvider, CachingLLMProvider
from nl_sql.llm.providers import build_provider
from nl_sql.llm.providers.base import EmbeddingProvider, LLMProvider
from nl_sql.llm.providers.mistral import MistralProvider
from nl_sql.paths import under_root
from nl_sql.schema_index.indexer import SchemaIndex

logger = logging.getLogger("nl_sql.api")

# ---------------------------------------------------------- response models


class HealthResponse(BaseModel):
    status: str
    version: str
    providers_configured: list[str]


class ReadyResponse(BaseModel):
    status: str
    chroma_ok: bool
    registry_ok: bool
    registered_dbs: int
    schema_chunks: int


class DatabaseInfo(BaseModel):
    db_id: str
    dialect: str
    description: str = ""
    table_count: int


class DatabasesResponse(BaseModel):
    databases: list[DatabaseInfo]


class AskRequest(BaseModel):
    question: str = Field(min_length=1, max_length=2000)
    db_id: str = Field(min_length=1)


class TraceStep(BaseModel):
    node: str
    model: str | None = None
    tokens_in: int | None = None
    tokens_out: int | None = None
    confidence: float | None = None


class AskResponse(BaseModel):
    trace_id: str
    db_id: str
    sql: str
    rationale: str
    confidence: float
    confidence_label: str
    rows: list[list[Any]] | None
    columns: list[str] | None
    row_count: int
    truncated: bool
    caption: str
    output_format: str | None
    error_kind: str | None
    error_message: str
    repair_attempted: bool
    latency_ms: float
    trace: list[TraceStep]


class EvalLatestResponse(BaseModel):
    configuration: str
    sql_model: str
    overall_ea: float | None
    n: int
    report_path: str


# ---------------------------------------------------------- helpers


def _confidence_label(value: float) -> str:
    if value >= 0.8:
        return "High"
    if value >= 0.5:
        return "Medium"
    if value > 0.0:
        return "Low"
    return "Unknown"


def _result_to_response(result: PipelineRunResult, *, latency_ms: float) -> AskResponse:
    rows: list[list[Any]] | None = None
    columns: list[str] | None = None
    row_count = 0
    truncated = False
    if result.outcome is not None and result.outcome.result is not None:
        rows = [list(r) for r in result.outcome.result.rows]
        columns = list(result.outcome.result.columns)
        row_count = result.outcome.result.row_count
        truncated = result.outcome.result.truncated

    trace_steps: list[TraceStep] = []
    for step in result.trace:
        trace_steps.append(
            TraceStep(
                node=str(step.get("node", "?")),
                model=step.get("model"),  # type: ignore[arg-type]
                tokens_in=step.get("input_tokens"),  # type: ignore[arg-type]
                tokens_out=step.get("output_tokens"),  # type: ignore[arg-type]
                confidence=step.get("confidence"),  # type: ignore[arg-type]
            )
        )

    fmt_name = None if result.output_format is None else type(result.output_format).__name__

    return AskResponse(
        trace_id=str(uuid.uuid4()),
        db_id=result.db_id,
        sql=result.sql,
        rationale=result.rationale,
        confidence=result.confidence,
        confidence_label=_confidence_label(result.confidence),
        rows=rows,
        columns=columns,
        row_count=row_count,
        truncated=truncated,
        caption=result.caption,
        output_format=fmt_name,
        error_kind=result.error_kind.value if result.error_kind else None,
        error_message=result.error_message,
        repair_attempted=result.repair_attempted,
        latency_ms=latency_ms,
        trace=trace_steps,
    )


# ---------------------------------------------------------- rate limit


class _TokenBucket:
    """Sliding-window token bucket per key.

    Default: 60 requests per 60 seconds. Single-process state — fine for the
    portfolio demo. Move to Redis if/when running multiple replicas.

    Buckets for keys that stopped calling are swept once per window: without
    that, `_hits` is an unbounded dict keyed by remote input and grows for the
    life of the process.
    """

    def __init__(self, *, max_req: int = 60, window_s: int = 60) -> None:
        self.max_req = max_req
        self.window_s = window_s
        self._hits: dict[str, deque[float]] = defaultdict(deque)
        self._last_sweep = time.time()

    def check(self, key: str) -> tuple[bool, int]:
        now = time.time()
        cutoff = now - self.window_s
        if now - self._last_sweep >= self.window_s:
            self._sweep(cutoff)
        bucket = self._hits[key]
        while bucket and bucket[0] < cutoff:
            bucket.popleft()
        if len(bucket) >= self.max_req:
            retry_after = int(self.window_s - (now - bucket[0]))
            return False, max(retry_after, 1)
        bucket.append(now)
        return True, 0

    def _sweep(self, cutoff: float) -> None:
        stale = [k for k, hits in self._hits.items() if not hits or hits[-1] < cutoff]
        for k in stale:
            del self._hits[k]
        self._last_sweep = time.time()


def _key_matches(supplied: str | None, expected: str) -> bool:
    """Constant-time API-key compare that can't be crashed by the header.

    `secrets.compare_digest` raises TypeError on `str` holding any codepoint
    above 127, and Starlette decodes headers as latin-1 — so a single accented
    byte in X-API-Key turned a 401 into an unhandled 500. Compare bytes.
    """
    if supplied is None:
        return False
    return secrets.compare_digest(
        supplied.encode("utf-8", "surrogateescape"),
        expected.encode("utf-8", "surrogateescape"),
    )


# ---------------------------------------------------------- bootstrap


def _build_pipeline_components(
    settings: Settings,
) -> tuple[DatabaseRegistry, SchemaIndex, LLMProvider, LLMProvider]:
    if not settings.mistral_api_key:
        raise RuntimeError("MISTRAL_API_KEY is not set — API can't bootstrap embeddings.")
    raw_sql = build_provider(settings.default_provider, settings=settings)
    sql_provider: LLMProvider = CachingLLMProvider(raw_sql, cache_dir=settings.llm_cache_dir)
    explain_provider: LLMProvider = sql_provider
    raw_embed: EmbeddingProvider = MistralProvider(
        api_key=settings.mistral_api_key,
        gen_model=settings.mistral_gen_model,
        embed_model=settings.mistral_embed_model,
        base_url=settings.mistral_base_url,
    )
    embedder: EmbeddingProvider = CachingEmbeddingProvider(
        raw_embed, cache_dir=settings.llm_cache_dir
    )
    schema_index = SchemaIndex(persist_dir=str(under_root("chroma_data")), embedder=embedder)
    registry = get_default_registry(
        pg_dsn=settings.pg_dsn,
        pg_db_id=settings.pg_db_id,
        pg_description=settings.pg_description,
    )
    return registry, schema_index, sql_provider, explain_provider


class Singletons(NamedTuple):
    """The four runtime objects the API routes share.

    Exposed as a public type so tests can construct mock instances and feed
    them through ``app.dependency_overrides[get_singletons]``.
    """

    pipeline: Any
    registry: DatabaseRegistry
    schema_index: SchemaIndex
    sql_provider: LLMProvider


@lru_cache(maxsize=1)
def _make_singletons() -> Singletons:
    """Lazy: build the pipeline only when the first /ask hits — keeps /healthz
    fast and avoids touching Chroma when the API is used for status probes."""
    settings = get_settings()
    registry, schema_index, sql_provider, explain_provider = _build_pipeline_components(settings)
    # Eval-script env toggles bootstrap into PipelineConfig once at boot;
    # individual nodes never read os.environ at runtime (see graph.py docstrings).
    config = PipelineConfig(
        sql_provider=sql_provider,
        explain_provider=explain_provider,
        schema_index=schema_index,
        registry=registry,
        fewshot_top_k=3,
        sort_schema_block=True,
        cross_db_fewshot=True,
        verify_retry_on_empty=True,
        use_m_schema=os.environ.get("NLSQL_M_SCHEMA") == "1",
        use_dac_prompt=os.environ.get("NLSQL_DAC") == "1",
    )
    pipeline = build_pipeline(config)
    return Singletons(pipeline, registry, schema_index, sql_provider)


def get_singletons() -> Singletons:
    """FastAPI Depends-able factory; tests override via ``app.dependency_overrides``."""
    return _make_singletons()


def create_app() -> FastAPI:
    app = FastAPI(
        title="NL→SQL Assistant",
        version=__version__,
        description=(
            "Portfolio API: natural-language questions → SQL → executed rows. "
            "BIRD Mini-Dev 61.5% reproducible single-run (codestral, $0); the "
            "hint-assisted headline is eval-only. Chinook 100%. AST safety guards."
        ),
    )
    settings = get_settings()
    rate_limiter = _TokenBucket(max_req=60, window_s=60)

    # Read from Settings (honours .env), not os.environ directly.
    api_key_env = settings.api_key

    async def require_api_key(
        request: Request,
        x_api_key: str | None = Header(default=None, alias="X-API-Key"),
    ) -> str:
        # 1. Auth — only enforced when a key is configured. Constant-time compare
        #    so a timing side-channel can't leak the key byte by byte.
        if api_key_env and not _key_matches(x_api_key, api_key_env):
            raise HTTPException(
                status_code=status.HTTP_401_UNAUTHORIZED,
                detail="missing or invalid X-API-Key header",
            )

        # 2. Rate limit — ALWAYS on, even with no key. A keyless public deploy is
        #    otherwise unthrottled (the README promises 60 req/min).
        #    Bucket by API key ONLY when a key is configured — in that branch the
        #    header has just been verified, so it is trusted input. With no key
        #    configured the header is unauthenticated attacker input: bucketing on
        #    it let a caller mint a fresh bucket per request with a random header
        #    and never hit the limit at all. Fall back to the client IP.
        client_host = request.client.host if request.client else "unknown"
        bucket_key = x_api_key if (api_key_env and x_api_key) else client_host
        ok, retry_after = rate_limiter.check(bucket_key)
        if not ok:
            raise HTTPException(
                status_code=status.HTTP_429_TOO_MANY_REQUESTS,
                detail=f"rate limit exceeded; retry in {retry_after}s",
                headers={"Retry-After": str(retry_after)},
            )
        return x_api_key or "anonymous"

    # --------------------------------------------------------- health / ready

    @app.get("/healthz", response_model=HealthResponse, tags=["status"])
    def healthz() -> HealthResponse:
        configured: list[str] = []
        if settings.mistral_api_key:
            configured.append("mistral")
        if settings.github_token:
            configured.append("github_models")
        if settings.groq_api_key:
            configured.append("groq")
        configured.append("ollama")
        return HealthResponse(
            status="ok",
            version=__version__,
            providers_configured=sorted(configured),
        )

    @app.get("/readyz", response_model=ReadyResponse, tags=["status"])
    def readyz() -> ReadyResponse:
        chroma_ok = False
        registry_ok = False
        registered = 0
        schema_chunks = 0
        try:
            factory: Any = app.dependency_overrides.get(get_singletons, get_singletons)
            singletons: Singletons = factory()
            registered = len(singletons.registry.ids())
            registry_ok = registered > 0
            schema_chunks = singletons.schema_index.schema_collection.count()
            chroma_ok = schema_chunks > 0
        except Exception:
            pass
        all_ok = chroma_ok and registry_ok
        return ReadyResponse(
            status="ok" if all_ok else "not_ready",
            chroma_ok=chroma_ok,
            registry_ok=registry_ok,
            registered_dbs=registered,
            schema_chunks=schema_chunks,
        )

    # --------------------------------------------------------- product API

    @app.get("/databases", response_model=DatabasesResponse, tags=["catalog"])
    def databases(
        _auth: str = Depends(require_api_key),
        singletons: Singletons = Depends(get_singletons),  # noqa: B008
    ) -> DatabasesResponse:
        infos: list[DatabaseInfo] = []
        for db_id in singletons.registry.ids():
            spec = singletons.registry.get(db_id)
            try:
                records = singletons.schema_index.schema_collection.get(
                    where={"db_id": db_id}, include=["metadatas"]
                )
                table_count = len(records.get("metadatas") or [])
            except Exception:
                table_count = 0
            infos.append(
                DatabaseInfo(
                    db_id=db_id,
                    dialect=str(spec.dialect),
                    description=str(getattr(spec, "description", "") or ""),
                    table_count=table_count,
                )
            )
        return DatabasesResponse(databases=infos)

    @app.post("/ask", response_model=AskResponse, tags=["nl-sql"])
    def ask(
        req: AskRequest,
        _auth: str = Depends(require_api_key),
        singletons: Singletons = Depends(get_singletons),  # noqa: B008
    ) -> AskResponse:
        if req.db_id not in singletons.registry.ids():
            raise HTTPException(
                status_code=status.HTTP_404_NOT_FOUND,
                detail=f"unknown db_id: {req.db_id!r}; see /databases for the list",
            )
        spec = singletons.registry.get(req.db_id)
        t0 = time.perf_counter()
        try:
            result = run_pipeline(
                singletons.pipeline,
                question=req.question,
                db_id=req.db_id,
                dialect=spec.dialect,
                verify_retry_on_empty=True,
            )
        except Exception as exc:
            # Never leak exception text (paths, driver internals, provider
            # payloads) to the client. Log the detail server-side under a
            # trace_id the caller can quote in a bug report.
            trace_id = str(uuid.uuid4())
            logger.exception("pipeline crashed [trace_id=%s]", trace_id)
            raise HTTPException(
                status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
                detail=f"internal pipeline error (trace_id={trace_id})",
            ) from exc
        latency_ms = (time.perf_counter() - t0) * 1000.0
        return _result_to_response(result, latency_ms=latency_ms)

    @app.get("/eval/latest", response_model=EvalLatestResponse, tags=["transparency"])
    def eval_latest() -> EvalLatestResponse:
        """Metadata of the reproducible single-run baseline (codestral, no voting,
        no BIRD rescue hints) — the number this API's own pipeline achieves. The
        hint-assisted headline is an eval-only layer; see README for the
        three-tier breakdown."""
        import json

        baseline = under_root("eval", "baselines", "reproducible_n200.json")
        if not baseline.exists():
            raise HTTPException(
                status_code=status.HTTP_404_NOT_FOUND,
                detail="no committed baseline yet — run scripts/eval_baseline.py",
            )
        data = json.loads(baseline.read_text(encoding="utf-8"))
        overall = data.get("overall") or {}
        return EvalLatestResponse(
            configuration=str(data.get("configuration", "unknown")),
            sql_model=str(data.get("sql_model", "unknown")),
            overall_ea=overall.get("ea"),
            n=int(overall.get("n") or 0),
            report_path=str(baseline),
        )

    return app


app = create_app()