| """Node: combine retrieve_schema + retrieve_examples into one ContextBundle. |
| |
| Thin wrapper over `nl_sql.schema_index.retrieve_context`. Per arch v2 §3, |
| this node also owns dialect-adapter hints (Postgres vs SQLite). For v1 we |
| just pass dialect through state — the prompt assembler picks dialect-specific |
| phrasing once we observe model failure modes during eval. |
| """ |
|
|
| from __future__ import annotations |
|
|
| from collections.abc import Callable |
|
|
| from sqlalchemy.engine import Engine |
|
|
| from nl_sql.agent.state import PipelineState |
| from nl_sql.db.registry import DatabaseRegistry |
| from nl_sql.schema_index.indexer import SchemaIndex |
| from nl_sql.schema_index.retriever import retrieve_context |
|
|
|
|
| def make_context_builder_node( |
| index: SchemaIndex, |
| *, |
| schema_top_k: int = 5, |
| fewshot_top_k: int = 3, |
| fk_hops: int = 1, |
| table_budget: int = 12, |
| registry: DatabaseRegistry | None = None, |
| primary_sample_size: int = 3, |
| extended_sample_size: int = 0, |
| cross_db_fewshot: bool = False, |
| ) -> Callable[[PipelineState], PipelineState]: |
| """Construct the context-builder node. |
| |
| Sample mixture wiring: when `registry` is provided AND |
| `extended_sample_size > primary_sample_size`, the node opens the |
| db's read-only engine for the current question and asks |
| `retrieve_context` to attach an "extended samples" appendix to the |
| bundle. `render_schema_block` then formats it as a supplementary |
| block. No-op when either flag is missing — the production default. |
| """ |
|
|
| mixture_enabled = registry is not None and extended_sample_size > primary_sample_size |
|
|
| def node(state: PipelineState) -> PipelineState: |
| question = state.get("question", "") |
| db_id = state.get("db_id", "") |
| if not question or not db_id: |
| return { |
| "context": None, |
| "trace": _append_trace(state, "context_builder", note="missing question or db_id"), |
| } |
| engine: Engine | None = None |
| if mixture_enabled: |
| assert registry is not None |
| engine = registry.get(db_id).make_engine() |
| try: |
| bundle = retrieve_context( |
| index, |
| question, |
| db_id=db_id, |
| schema_top_k=schema_top_k, |
| fewshot_top_k=fewshot_top_k, |
| fk_hops=fk_hops, |
| table_budget=table_budget, |
| engine=engine, |
| primary_sample_size=primary_sample_size, |
| extended_sample_size=extended_sample_size, |
| cross_db_fewshot=cross_db_fewshot, |
| ) |
| finally: |
| if engine is not None: |
| engine.dispose() |
| return { |
| "context": bundle, |
| "trace": _append_trace( |
| state, |
| "context_builder", |
| tables=bundle.all_tables, |
| fewshots=len(bundle.fewshots), |
| truncated=bundle.truncated, |
| extended_sample_tables=( |
| sorted(bundle.extended_samples) if bundle.extended_samples else [] |
| ), |
| ), |
| } |
|
|
| return node |
|
|
|
|
| def _append_trace(state: PipelineState, node: str, **details: object) -> list[dict[str, object]]: |
| trace = list(state.get("trace") or []) |
| trace.append({"node": node, **details}) |
| return trace |
|
|
