"""Build the Chroma schema index for one (or all) registered databases.

Live tool — calls Mistral `mistral-embed` for vectors. Idempotent: re-runs

upsert chunks under the same `chunk_id` (db::table), so vectors get refreshed

in place and stale chunks for renamed tables are NOT auto-pruned (run with

``--reset`` to clear the collection first if you have schema deletions).

The default ``--sample-size`` is imported from ``PipelineConfig.primary_sample_size``

so the index is built with the same density runtime expects. Pass an explicit

value only if you want to rebuild for a non-default runtime configuration.

Usage:

uv run python scripts/build_index.py --db chinook

uv run python scripts/build_index.py --db all --persist chroma_data

uv run python scripts/build_index.py --db chinook --reset

"""

from __future__ import annotations

import argparse

import sys

from dataclasses import fields

from pathlib import Path

import chromadb

from nl_sql.agent.graph import PipelineConfig

from nl_sql.config import get_settings

from nl_sql.db.registry import get_default_registry

from nl_sql.llm.cache import CachingEmbeddingProvider

from nl_sql.llm.providers.base import EmbeddingProvider

from nl_sql.llm.providers.mistral import MistralProvider

from nl_sql.schema_index.chunker import to_chunks

from nl_sql.schema_index.indexer import SCHEMA_COLLECTION, SchemaIndex

from nl_sql.schema_index.introspector import introspect

def _runtime_sample_size_default() -> int:

"""Read `PipelineConfig.primary_sample_size` default without constructing

the dataclass (it requires live providers/registry we don't have here)."""

for field_ in fields(PipelineConfig):

if field_.name == "primary_sample_size":

default = field_.default

if isinstance(default, int):

return default

raise RuntimeError("PipelineConfig.primary_sample_size default missing")

DEFAULT_SAMPLE_SIZE: int = _runtime_sample_size_default()

"""Source of truth for the sample density baked into Chroma chunks.

Runtime expects this to equal `PipelineConfig.primary_sample_size`; the

mixture appendix breaks if the index is built with more samples than

runtime advertises."""

def build_for_db(idx: SchemaIndex, db_id: str, *, sample_size: int = DEFAULT_SAMPLE_SIZE) -> int:

registry = get_default_registry()

spec = registry.get(db_id)

print(f"[introspect] {db_id} ({spec.url})")

tables = introspect(spec.make_engine(), sample_size=sample_size)

print(f"[chunk] {len(tables)} tables → chunks")

chunks = to_chunks(tables, db_id=db_id)

print(f"[index] embedding + upserting {len(chunks)} chunks")

n = idx.index_schema(chunks)

print(f"[done] {db_id}: {n} chunks indexed")

return n

def build_parser() -> argparse.ArgumentParser:

parser = argparse.ArgumentParser(description=__doc__)

parser.add_argument(

"--db",

required=True,

help="Database id (e.g. 'chinook', 'bird_california_schools') or 'all'.",

)

parser.add_argument(

"--persist",

default="chroma_data",

help="Chroma persist directory (default: chroma_data/)",

)

parser.add_argument(

"--sample-size",

type=int,

default=DEFAULT_SAMPLE_SIZE,

help=(

"Top-K sample values per column to bake into each chunk "

f"(default: {DEFAULT_SAMPLE_SIZE} = PipelineConfig.primary_sample_size). "

"Keep aligned with runtime or the sample-mixture appendix breaks."

),

)

parser.add_argument(

"--reset",

action="store_true",

help="Drop the schema_chunks collection before indexing.",

)

parser.add_argument(

"--no-cache",

action="store_true",

help="Disable diskcache wrapper around the embedding provider.",

)

return parser

def main(argv: list[str] | None = None) -> int:

args = build_parser().parse_args(argv)

settings = get_settings()

persist = Path(args.persist)

persist.mkdir(parents=True, exist_ok=True)

client = chromadb.PersistentClient(path=str(persist))

if args.reset:

try:

client.delete_collection(SCHEMA_COLLECTION)

print(f"[reset] dropped {SCHEMA_COLLECTION}")

except Exception as exc:

print(f"[reset] no existing {SCHEMA_COLLECTION} to drop ({exc})")

raw_embedder = 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 = (

raw_embedder

if args.no_cache

else CachingEmbeddingProvider(

raw_embedder,

cache_dir=settings.llm_cache_dir,

size_limit_gb=settings.llm_cache_size_limit_gb,

)

)

idx = SchemaIndex(persist_dir=persist, embedder=embedder, client=client)

registry = get_default_registry()

targets = registry.ids() if args.db == "all" else [args.db]

total = 0

for db_id in targets:

total += build_for_db(idx, db_id, sample_size=args.sample_size)

print(f"[summary] indexed {total} chunks across {len(targets)} db(s)")

return 0

if __name__ == "__main__":

sys.exit(main())