diff --git a/.gitignore b/.gitignore index 28e2df9..0a27682 100644 --- a/.gitignore +++ b/.gitignore @@ -18,3 +18,8 @@ ui/framework # and the BIRD golden eval set. Never committed. ctrl/seed/data api/evals/data + +# Recon build artefacts — generated by `make recon` from the per-dataset +# YAML + live warehouse introspection. Re-buildable, not source. +api/datasets/*/recon.json +api/datasets/*/extracted_schema.json diff --git a/Makefile b/Makefile index 70be936..202ae52 100644 --- a/Makefile +++ b/Makefile @@ -1,5 +1,5 @@ -.PHONY: kind tilt-up tilt-down seed seed-fetch seed-push evals test \ - compose-up compose-down compose-clean compose-seed compose-evals \ +.PHONY: kind tilt-up tilt-down seed seed-fetch seed-push recon evals test \ + compose-up compose-down compose-clean compose-seed compose-recon compose-evals \ docs-graphs COMPOSE := docker compose -f ctrl/docker-compose.yml @@ -32,6 +32,22 @@ seed-push: seed: seed-fetch seed-push kubectl $(KCTX) $(KNS) exec deploy/api -- uv run python -m seed.load_bird +# Rebuild api/datasets//{extracted_schema,recon}.json from the YAML + +# live warehouse introspection. Runs inside the api pod (needs Postgres +# access), then copies the artefacts back to the host so they appear next +# to the source YAMLs. +recon: + @POD=$$(kubectl $(KCTX) $(KNS) get pod -l app=api -o jsonpath='{.items[0].metadata.name}'); \ + kubectl $(KCTX) $(KNS) exec $$POD -- uv run python -m api.recon.build; \ + for ds in $$(ls api/datasets); do \ + [ -d "api/datasets/$$ds" ] || continue; \ + case "$$ds" in __*|.*) continue ;; esac; \ + for f in extracted_schema.json recon.json; do \ + kubectl $(KCTX) $(KNS) cp $$POD:/app/api/datasets/$$ds/$$f api/datasets/$$ds/$$f 2>/dev/null \ + && echo " → api/datasets/$$ds/$$f" || true; \ + done; \ + done + evals: kubectl $(KCTX) $(KNS) exec deploy/api -- uv run python -m api.evals.run_evals @@ -55,6 +71,9 @@ compose-clean: compose-seed: seed-fetch $(COMPOSE) exec api uv run python -m seed.load_bird +compose-recon: + $(COMPOSE) exec api uv run python -m api.recon.build + compose-evals: $(COMPOSE) exec api uv run python -m api.evals.run_evals diff --git a/api/analyses/compare_periods.py b/api/analyses/compare_periods.py index 14f9769..b1cd401 100644 --- a/api/analyses/compare_periods.py +++ b/api/analyses/compare_periods.py @@ -18,8 +18,8 @@ from api.analyses.types import Finding from api.llm import chat from api.prompts import load, render from api.runtime import events +from api.recon import load_recon from api.tools.execute_sql import execute_sql -from api.tools.schema import load_schema logger = logging.getLogger("nvi.analyses.compare_periods") @@ -110,7 +110,7 @@ class ComparePeriods(Analysis): def _generate_pair(question: str, period_a: str, period_b: str) -> dict[str, dict[str, str]]: - schema = load_schema() + schema = load_recon() text = chat( system=load("compare_periods.pair"), user=render( diff --git a/api/analyses/drill_down/__init__.py b/api/analyses/drill_down/__init__.py new file mode 100644 index 0000000..ecd588c --- /dev/null +++ b/api/analyses/drill_down/__init__.py @@ -0,0 +1,3 @@ +from api.analyses.drill_down.analysis import DrillDown + +__all__ = ["DrillDown"] diff --git a/api/analyses/drill_down/analysis.py b/api/analyses/drill_down/analysis.py new file mode 100644 index 0000000..2b0a98d --- /dev/null +++ b/api/analyses/drill_down/analysis.py @@ -0,0 +1,89 @@ +"""drill_down Analysis — ReAct loop. + +Pattern: bounded loop that picks a dimension to slice by, generates SQL, +executes it, and decides whether to continue. From the runtime's view this +is just `(args, question) → Finding`; internally it's a small ReAct agent. + +This file owns the orchestration. The decision step, the slice execution, +the interpretation, and the prompt-context formatters all live in +`helpers.py`. Types and constants live in `types.py`. +""" + +from __future__ import annotations + +import logging +from typing import Any + +from api import langfuse_client as lf +from api.analyses.base import Analysis +from api.analyses.drill_down.helpers import decide_next, execute_slice, interpret +from api.analyses.drill_down.types import ARGS_SCHEMA, DrillDownArgs, Slice +from api.analyses.types import Finding + +logger = logging.getLogger("nvi.analyses.drill_down") + + +class DrillDown(Analysis): + name = "drill_down" + description = ( + "Iteratively slice a metric by candidate dimensions to find which " + "ones explain the most variance. ReAct loop: pick a dimension, " + "query, decide whether to continue. Best for open-ended 'which " + "factors explain X' or 'why are some segments different' questions." + ) + args_schema = ARGS_SCHEMA + + async def run(self, args: dict[str, Any], question: str) -> Finding: + a = DrillDownArgs.from_raw(args, default_question=question) + + with lf.span("analysis.drill_down", input={"question": a.question, "metric": a.metric}) as span: + try: + slices = await self._loop(a) + if not slices: + return Finding( + analysis=self.name, summary="", + error="drill_down stopped before producing any slice", + ) + summary = await interpret(a.question, a.metric, slices) + except Exception as e: + logger.exception("drill_down failed") + span.update(output={"error": str(e)}) + return Finding(analysis=self.name, summary="", error=str(e)) + + span.update(output={"summary": summary, "iterations": len(slices)}) + return self._finalise(a, slices, summary) + + async def _loop(self, a: DrillDownArgs) -> list[Slice]: + """Bounded ReAct loop. Each iteration: decide → slice. Stops when the + LLM says so, when the budget runs out, or when a repeat is detected.""" + slices: list[Slice] = [] + tried: set[str] = set() + for it in range(a.max_iters): + remaining = a.max_iters - it + decision = await decide_next(a.question, a.metric, a.dimensions, slices, remaining) + if decision.get("action") == "stop": + break + dim = decision.get("dimension") + if not dim or dim in tried: + break + tried.add(dim) + slices.append(await execute_slice( + a.question, a.metric, dim, decision.get("reason", ""), + )) + return slices + + def _finalise(self, a: DrillDownArgs, slices: list[Slice], summary: str) -> Finding: + rows_combined = [ + {"dimension": s.dimension, **r} for s in slices for r in s.rows + ] + return Finding( + analysis=self.name, + summary=summary, + rows=rows_combined, + sql=[s.sql for s in slices], + metadata={ + "metric": a.metric, + "iterations": len(slices), + "dimensions_tried": [s.dimension for s in slices], + }, + ) diff --git a/api/analyses/drill_down/helpers.py b/api/analyses/drill_down/helpers.py new file mode 100644 index 0000000..984a3de --- /dev/null +++ b/api/analyses/drill_down/helpers.py @@ -0,0 +1,173 @@ +"""drill_down helpers — one function per concern. + +Kept separate from `analysis.py` so the DrillDown class stays readable as +high-level orchestration: decide → slice → loop → interpret. +""" + +from __future__ import annotations + +import json +import logging +import re +from typing import Any + +from api import langfuse_client as lf +from api.analyses.drill_down.types import Slice +from api.llm import chat +from api.prompts import load, render +from api.recon import load_recon +from api.runtime import events +from api.tools.execute_sql import execute_sql +from api.tools.text_to_sql import text_to_sql + +logger = logging.getLogger("nvi.analyses.drill_down.helpers") + + +# ── Decision step ── + +async def decide_next(question: str, metric: str, dimensions: list[str], + slices: list[Slice], budget: int) -> dict[str, Any]: + """One LLM call to pick the next dimension or stop. + + If the LLM picks something not in the candidate list, drop the choice + and stop — the planner's fallback (if any) takes over. We don't try to + coerce the LLM into a valid dimension because it might be hallucinating + a name (e.g. a metric name) that has no obvious mapping. + """ + system = load("drill_down.next.system") + user = render( + "drill_down.next.user", + question=question, + metric=metric, + dimensions=", ".join(dimensions), + history=format_history(slices), + budget=budget, + ) + with lf.span("drill_down.next", as_type="generation", input={"budget": budget}) as span: + await events.publish_current(events.llm_call( + "drill_down.next", system_len=len(system), user_len=len(user), + )) + decision = _parse_json(chat(system=system, user=user, max_tokens=256)) + + # Hard guard: the chosen dimension MUST be in the candidate list. + if decision.get("action") == "drill": + dim = decision.get("dimension") + if dim not in dimensions: + logger.warning( + "drill_down.next picked %r (not in candidates %s); stopping", + dim, dimensions, + ) + decision = { + "action": "stop", + "reason": f"LLM picked {dim!r}, which isn't in the candidate dimensions", + } + span.update(output=decision) + return decision + + +# ── Slice execution ── + +def build_slice_question(question: str, metric: str, dim: str) -> str: + """Construct a slice question with explicit table/join hints from the + recon graph. Stops the LLM from inventing FROM clauses that omit the + table that owns the dimension column. + """ + recon = load_recon() + dim_owners = recon.owning_tables(dim) + metric_def = recon.metrics.get(metric) + metric_table = metric_def.from_table if metric_def else None + + hints: list[str] = [] + if dim_owners: + hints.append(f"- The dimension column `{dim}` lives in table: {', '.join(dim_owners)}.") + if metric_table: + hints.append(f"- The metric `{metric}` is defined over table `{metric_table}`.") + if metric_def and metric_def.filter: + hints.append(f" Apply this filter for the metric: {metric_def.filter}.") + if metric_def: + hints.append(f" Compute the metric as: {metric_def.sql} (use this expression verbatim).") + if dim_owners and metric_table and dim_owners[0] != metric_table: + path = recon.join_path(metric_table, dim_owners[0]) + if path: + hints.append(f"- Required JOIN path: {' → '.join(path)}.") + + hint_block = ("\n".join(hints) + "\n\n") if hints else "" + return ( + f"{question}\n\n" + f"Slice the metric `{metric}` by `{dim}` and return the top rows by metric value.\n\n" + f"{hint_block}" + f"GROUP BY the dimension. ORDER BY the metric DESC. Limit to top 10." + ) + + +async def execute_slice(question: str, metric: str, dim: str, reason: str) -> Slice: + """Generate SQL for one slice, execute it, emit tool-call events + around both, and return the Slice.""" + slice_q = build_slice_question(question, metric, dim) + + await events.publish_current(events.tool_call_start("text_to_sql", input={"question": slice_q})) + t2s = text_to_sql(slice_q) + await events.publish_current(events.tool_call_end( + "text_to_sql", output={"sql": t2s.sql, "tables": t2s.used_tables}, + )) + + await events.publish_current(events.tool_call_start( + "execute_sql", input={"sql": t2s.sql, "dimension": dim}, + )) + result = execute_sql(t2s.sql) + await events.publish_current(events.tool_call_end( + "execute_sql", + output={"dimension": dim, "row_count": result.row_count, + "preview": result.as_dicts()[:5]}, + )) + + return Slice( + dimension=dim, + sql=t2s.sql, + rows=result.as_dicts()[:20], + reason=reason, + ) + + +# ── Final interpretation ── + +async def interpret(question: str, metric: str, slices: list[Slice]) -> str: + system = load("drill_down.interpret.system") + user = render( + "drill_down.interpret.user", + question=question, + metric=metric, + slices_block=format_slices_block(slices), + ) + await events.publish_current(events.llm_call( + "drill_down.interpret", system_len=len(system), user_len=len(user), + )) + return chat(system=system, user=user, max_tokens=512) + + +# ── Prompt-context formatters ── + +def format_history(slices: list[Slice]) -> str: + """Render the 'already tried' block fed to drill_down.next.""" + if not slices: + return "(none yet)" + return "\n".join(f"- {s.dimension}: {repr(s.rows[:5])}" for s in slices) + + +def format_slices_block(slices: list[Slice]) -> str: + """Render every slice's rows for drill_down.interpret.""" + return "\n\n".join( + f"dimension: {s.dimension}\nrows: {repr(s.rows)}" + for s in slices + ) + + +# ── JSON extraction ── + +def _parse_json(text: str) -> dict[str, Any]: + m = re.search(r"```(?:json)?\s*(\{.*\})\s*```", text, re.DOTALL) + raw = m.group(1) if m else text + start, end = raw.find("{"), raw.rfind("}") + if start < 0 or end <= start: + raise ValueError(f"drill_down.next returned no JSON: {text[:200]!r}") + return json.loads(raw[start:end + 1]) diff --git a/api/analyses/drill_down/types.py b/api/analyses/drill_down/types.py new file mode 100644 index 0000000..ae03c98 --- /dev/null +++ b/api/analyses/drill_down/types.py @@ -0,0 +1,61 @@ +"""Public data shapes + constants for the drill_down Analysis.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any + +DEFAULT_DIMENSIONS: list[str] = ["A2", "A3", "A11", "A12", "A13", "A14"] +DEFAULT_METRIC: str = "loan_default_rate" +DEFAULT_MAX_ITERS: int = 3 + +# JSON-schema-ish surface the planner reads to know which args this +# Analysis accepts. Kept as a plain dict so `registry.catalog()` can dump +# it straight to JSON. +ARGS_SCHEMA: dict[str, dict[str, Any]] = { + "question": { + "type": "string", + "description": "Refined question this Analysis should answer.", + }, + "metric": { + "type": "string", + "description": "Metric or column to slice (e.g. 'loan_default_rate', 'amount').", + }, + "dimensions": { + "type": "array", + "items": "string", + "description": "Candidate dimensions to drill into. Defaults to district demographics.", + }, + "max_iters": { + "type": "integer", + "description": f"Cap on slice iterations. Default {DEFAULT_MAX_ITERS}.", + }, +} + + +@dataclass +class DrillDownArgs: + """Parsed, defaulted args for one drill_down run.""" + question: str + metric: str = DEFAULT_METRIC + dimensions: list[str] = field(default_factory=lambda: list(DEFAULT_DIMENSIONS)) + max_iters: int = DEFAULT_MAX_ITERS + + @classmethod + def from_raw(cls, raw: dict[str, Any], default_question: str) -> "DrillDownArgs": + return cls( + question=raw.get("question") or default_question, + metric=raw.get("metric") or DEFAULT_METRIC, + dimensions=raw.get("dimensions") or list(DEFAULT_DIMENSIONS), + max_iters=int(raw.get("max_iters") or DEFAULT_MAX_ITERS), + ) + + +@dataclass +class Slice: + """One iteration's slice: the chosen dimension, the SQL it produced, + its rows, and the LLM's stated reason for picking the dimension.""" + dimension: str + sql: str + rows: list[dict[str, Any]] + reason: str = "" diff --git a/api/analyses/registry.py b/api/analyses/registry.py index 4a2f272..6b2c17f 100644 --- a/api/analyses/registry.py +++ b/api/analyses/registry.py @@ -7,10 +7,12 @@ from typing import Any from api.analyses.base import Analysis from api.analyses.compare_periods import ComparePeriods from api.analyses.direct_answer import DirectAnswer +from api.analyses.drill_down import DrillDown REGISTRY: dict[str, Analysis] = { DirectAnswer.name: DirectAnswer(), ComparePeriods.name: ComparePeriods(), + DrillDown.name: DrillDown(), } diff --git a/api/datasets/financial/schema_docs.yaml b/api/datasets/financial/schema_docs.yaml index 88b2dd8..c38940f 100644 --- a/api/datasets/financial/schema_docs.yaml +++ b/api/datasets/financial/schema_docs.yaml @@ -1,93 +1,99 @@ -# Per-table / per-column human descriptions for the BIRD `financial` schema -# (PKDD'99 Czech bank). The text-to-SQL tool prepends the relevant subset to -# the LLM prompt so it can disambiguate cryptic column names like A2, A3, etc. +# Human-authored augmentation for the BIRD `financial` dataset. +# +# The warehouse STRUCTURE (tables/columns/types/nullability) is extracted +# from Postgres into extracted_schema.json by `make recon`. This file only +# carries what humans add on top: table/column descriptions and relationships. +# (Metrics live in metrics.yaml.) +# +# Format: +# tables.: one-line description. +# columns..: one-line description. Sparse — only the ones +# worth describing. Undescribed columns just have +# no description in the recon. +# relationships: declared FKs (BIRD ships SQLite without them). schema: financial tables: - account: - description: One row per bank account. The pivot table — most facts hang off account_id. - columns: - account_id: Surrogate primary key for the account. - district_id: Geographic district the account belongs to (joins to district). - frequency: Statement issuance frequency. "POPLATEK MESICNE" = monthly, "POPLATEK TYDNE" = weekly, "POPLATEK PO OBRATU" = on transaction. - date: Date the account was opened (YYMMDD as integer). + account: One row per bank account. The pivot table — most facts hang off account_id. + client: One row per bank customer. + disp: Disposition — many-to-many link between client and account. Type tells you whether the client is the OWNER or just a USER of the account. + district: Demographic and economic data for each Czech district. The A-columns are infamous — most analysis questions referencing demographics use these. + loan: One row per loan granted to an account. + card: One row per credit/debit card issued to a disposition. + trans: Transaction history. Largest table by far (~1M rows). Every credit, debit, transfer. + order: Standing payment orders (recurring debits). Note the table name collides with a SQL keyword and must be quoted. - client: - description: One row per bank customer. - columns: - client_id: Surrogate primary key for the client. - gender: '"M" or "F". Derived from birth_number.' - birth_date: Date of birth, normalised. Original birth_number encoded gender by adding 50 to the month field for women. - district_id: District the client lives in (joins to district). +columns: + account.district_id: Geographic district the account belongs to (joins to district). + account.frequency: Statement issuance frequency. "POPLATEK MESICNE" = monthly, "POPLATEK TYDNE" = weekly, "POPLATEK PO OBRATU" = on transaction. + account.date: Date the account was opened (YYMMDD as integer). - disp: - description: Disposition — many-to-many link between client and account. Type tells you whether the client is the OWNER or just a USER of the account. - columns: - disp_id: Surrogate key. - client_id: Joins to client. - account_id: Joins to account. - type: '"OWNER" or "DISPONENT" (USER). Only owners can request loans / cards.' + client.gender: '"M" or "F". Derived from birth_number.' + client.birth_date: Date of birth, normalised. Original birth_number encoded gender by adding 50 to the month field for women. + client.district_id: District the client lives in (joins to district). - district: - description: Demographic and economic data for each Czech district. The A-columns are infamous — most analysis questions referencing demographics use these. - columns: - district_id: Surrogate key. - A2: District name. - A3: Region (a grouping of districts). - A4: Number of inhabitants. - A5: Number of municipalities with population < 499. - A6: Number of municipalities with population 500-1999. - A7: Number of municipalities with population 2000-9999. - A8: Number of municipalities with population > 10000. - A9: Number of cities. - A10: Ratio of urban inhabitants (percent). - A11: Average salary (in CZK). - A12: Unemployment rate in 1995 (percent). NULL means missing data. - A13: Unemployment rate in 1996 (percent). - A14: Number of entrepreneurs per 1000 inhabitants. - A15: Number of crimes committed in 1995. NULL means missing data. - A16: Number of crimes committed in 1996. + disp.client_id: Joins to client. + disp.account_id: Joins to account. + disp.type: '"OWNER" or "DISPONENT" (USER). Only owners can request loans / cards.' - loan: - description: One row per loan granted to an account. - columns: - loan_id: Surrogate key. - account_id: The account the loan was granted to. - date: Loan grant date (YYMMDD as integer). - amount: Total amount of the loan (CZK). - duration: Loan duration in months. - payments: Monthly payment (CZK). - status: Loan status code. "A" = contract finished, no problems. "B" = contract finished, loan not paid (DEFAULT). "C" = running contract, payments on time. "D" = running contract, client in debt (AT RISK). Defaults are A,B finished or B,D paid-flag attached. + district.A2: District name. + district.A3: Region (a grouping of districts). + district.A4: Number of inhabitants. + district.A5: Number of municipalities with population < 499. + district.A6: Number of municipalities with population 500-1999. + district.A7: Number of municipalities with population 2000-9999. + district.A8: Number of municipalities with population > 10000. + district.A9: Number of cities. + district.A10: Ratio of urban inhabitants (percent). + district.A11: Average salary (in CZK). + district.A12: Unemployment rate in 1995 (percent). NULL means missing data. + district.A13: Unemployment rate in 1996 (percent). + district.A14: Number of entrepreneurs per 1000 inhabitants. + district.A15: Number of crimes committed in 1995. NULL means missing data. + district.A16: Number of crimes committed in 1996. - card: - description: One row per credit/debit card issued to a disposition. - columns: - card_id: Surrogate key. - disp_id: The disposition (and via it, the account) the card belongs to. - type: '"junior", "classic", or "gold".' - issued: Date the card was issued. + loan.account_id: The account the loan was granted to. + loan.date: Loan grant date (YYMMDD as integer). + loan.amount: Total amount of the loan (CZK). + loan.duration: Loan duration in months. + loan.payments: Monthly payment (CZK). + loan.status: Loan status code. "A" = contract finished, no problems. "B" = contract finished, loan not paid (DEFAULT). "C" = running contract, payments on time. "D" = running contract, client in debt (AT RISK). - trans: - description: Transaction history. Largest table by far (~1M rows). Every credit, debit, transfer. - columns: - trans_id: Surrogate key. - account_id: The account the transaction belongs to. - date: Transaction date (YYMMDD as integer). - type: '"PRIJEM" (credit) or "VYDAJ" (debit).' - operation: Mode of operation, e.g. "VKLAD" (cash credit), "VYBER" (cash withdrawal), "PREVOD Z UCTU" (transfer from another bank), etc. - amount: Amount (CZK). - balance: Account balance after the transaction. - k_symbol: Characterisation of the payment. "POJISTNE" = insurance, "SLUZBY" = household, "UROK" = interest credited, "SANKC. UROK" = penalty interest, "SIPO" = household payment, "DUCHOD" = old-age pension, "UVER" = loan payment. - bank: Counterparty bank code (only set for inter-bank transfers). - account: Counterparty account number (only set for inter-bank transfers). + card.disp_id: The disposition (and via it, the account) the card belongs to. + card.type: '"junior", "classic", or "gold".' + card.issued: Date the card was issued. - "order": - description: Standing payment orders (recurring debits). Note the table name collides with a SQL keyword and must be quoted. - columns: - order_id: Surrogate key. - account_id: The originating account. - bank_to: Recipient bank. - account_to: Recipient account. - amount: Order amount (CZK). - k_symbol: Purpose, same vocabulary as trans.k_symbol. + trans.account_id: The account the transaction belongs to. + trans.date: Transaction date (YYMMDD as integer). + trans.type: '"PRIJEM" (credit) or "VYDAJ" (debit).' + trans.operation: Mode of operation, e.g. "VKLAD" (cash credit), "VYBER" (cash withdrawal), "PREVOD Z UCTU" (transfer from another bank). + trans.amount: Amount (CZK). + trans.balance: Account balance after the transaction. + trans.k_symbol: Characterisation of the payment. "POJISTNE" = insurance, "SLUZBY" = household, "UROK" = interest credited, "SANKC. UROK" = penalty interest, "SIPO" = household payment, "DUCHOD" = old-age pension, "UVER" = loan payment. + trans.bank: Counterparty bank code (only set for inter-bank transfers). + trans.account: Counterparty account number (only set for inter-bank transfers). + + order.account_id: The originating account. + order.bank_to: Recipient bank. + order.account_to: Recipient account. + order.amount: Order amount (CZK). + order.k_symbol: Purpose, same vocabulary as trans.k_symbol. + +relationships: + - from: account.district_id + to: district.district_id + - from: client.district_id + to: district.district_id + - from: disp.client_id + to: client.client_id + - from: disp.account_id + to: account.account_id + - from: loan.account_id + to: account.account_id + - from: card.disp_id + to: disp.disp_id + - from: trans.account_id + to: account.account_id + - from: order.account_id + to: account.account_id diff --git a/api/plan/planner.py b/api/plan/planner.py index 43864a0..03dbb33 100644 --- a/api/plan/planner.py +++ b/api/plan/planner.py @@ -15,13 +15,13 @@ from api.analyses.registry import catalog from api.llm import chat from api.plan.types import Plan from api.prompts import load, render -from api.tools.schema import load_schema +from api.recon import load_recon logger = logging.getLogger("nvi.plan.planner") def plan(question: str) -> Plan: - schema = load_schema() + schema = load_recon() user = render( "planner.user", question=question, diff --git a/api/plan/types.py b/api/plan/types.py index f082844..c48274c 100644 --- a/api/plan/types.py +++ b/api/plan/types.py @@ -3,7 +3,7 @@ from __future__ import annotations from dataclasses import dataclass, field -from typing import Any +from typing import Any, Optional @dataclass @@ -11,17 +11,30 @@ class PlanStep: analysis: str args: dict[str, Any] = field(default_factory=dict) why: str = "" + # Optional sibling step to run when the primary errors or yields no + # finding. Lets the planner hedge between a speculative Analysis (e.g. + # `drill_down`) and a safer one (e.g. `direct_answer`). + fallback: Optional["PlanStep"] = None @classmethod def from_dict(cls, raw: dict[str, Any]) -> "PlanStep": + fb_raw = raw.get("fallback") return cls( analysis=raw["analysis"], args=raw.get("args") or {}, why=raw.get("why", ""), + fallback=cls.from_dict(fb_raw) if fb_raw else None, ) def to_dict(self) -> dict[str, Any]: - return {"analysis": self.analysis, "args": self.args, "why": self.why} + out: dict[str, Any] = { + "analysis": self.analysis, + "args": self.args, + "why": self.why, + } + if self.fallback is not None: + out["fallback"] = self.fallback.to_dict() + return out @dataclass diff --git a/api/prompts/drill_down.interpret.system.txt b/api/prompts/drill_down.interpret.system.txt new file mode 100644 index 0000000..4f60d36 --- /dev/null +++ b/api/prompts/drill_down.interpret.system.txt @@ -0,0 +1,7 @@ +Task: write a short interpretation of an iterative drill-down across a metric. + +You're given the analyst's question, the metric, and the slices tried (each: a dimension and its top rows). Identify the dimension(s) that show the strongest signal — biggest variance, most concentrated outliers, clearest pattern — and write a short, concrete summary. + +Output: two to four sentences. Lead with the dimension(s) that mattered most. Cite specific values from the slices (e.g. "districts with unemployment above 5% accounted for 73% of defaults"). No preamble. + +If no slice produced a meaningful signal, say so plainly. diff --git a/api/prompts/drill_down.interpret.user.txt b/api/prompts/drill_down.interpret.user.txt new file mode 100644 index 0000000..f0fe8e9 --- /dev/null +++ b/api/prompts/drill_down.interpret.user.txt @@ -0,0 +1,7 @@ +Question: +{question} + +Metric: {metric} + +Slices (in order tried): +{slices_block} diff --git a/api/prompts/drill_down.next.system.txt b/api/prompts/drill_down.next.system.txt new file mode 100644 index 0000000..2320e6a --- /dev/null +++ b/api/prompts/drill_down.next.system.txt @@ -0,0 +1,16 @@ +Task: pick the next dimension to slice by, in an iterative drill-down investigation. + +You're given the analyst's question, the candidate dimensions, the dimensions already tried (with their slice rows), and a remaining iteration budget. Decide which dimension to try next, or stop if the investigation has converged or budget is exhausted. + +Output: a JSON object in one fenced ```json block: +{ + "action": "drill" | "stop", + "dimension": "", // required when action is "drill" + "reason": "" +} + +Rules: +- Pick a dimension that hasn't been tried yet and that the previous slices suggest might explain variance in the metric. +- If every candidate has been tried, or the previous slice already shows a clear story, return action="stop". +- Don't repeat dimensions. +- Never invent a dimension that isn't in the candidate list. diff --git a/api/prompts/drill_down.next.user.txt b/api/prompts/drill_down.next.user.txt new file mode 100644 index 0000000..2e97227 --- /dev/null +++ b/api/prompts/drill_down.next.user.txt @@ -0,0 +1,13 @@ +Question: +{question} + +Metric we're slicing: +{metric} + +Candidate dimensions: +{dimensions} + +Already tried (with row previews): +{history} + +Remaining iterations: {budget} diff --git a/api/prompts/planner.system.txt b/api/prompts/planner.system.txt index 0a9443d..7aa734b 100644 --- a/api/prompts/planner.system.txt +++ b/api/prompts/planner.system.txt @@ -6,7 +6,16 @@ Output: a JSON object with exactly this shape, in one fenced ```json block: { "rationale": "", "steps": [ - {"analysis": "", "args": { ... }, "why": ""} + { + "analysis": "", + "args": { ... }, + "why": "", + "fallback": { // OPTIONAL — see rules + "analysis": "", + "args": { ... }, + "why": "" + } + } ] } @@ -15,3 +24,4 @@ Rules: - The `args` keys must match the analysis's args_schema; additional keys are ignored, missing optional ones are fine. - Most questions are single-step. Only chain analyses when the second genuinely needs the first's output. - Comparing two periods is ONE `compare_periods` step, not two `direct_answer` steps. +- A step's `fallback` is optional. Include one when the primary analysis is speculative (e.g. an exploratory `drill_down` that might find nothing useful) and there's a safer Analysis that can still produce an answer. The fallback runs only if the primary errors or yields no finding. Fallbacks must not themselves have fallbacks. diff --git a/api/prompts/text_to_sql.system.txt b/api/prompts/text_to_sql.system.txt index f9ba9dc..c019c59 100644 --- a/api/prompts/text_to_sql.system.txt +++ b/api/prompts/text_to_sql.system.txt @@ -4,8 +4,9 @@ Output: exactly one fenced ```sql block, nothing else — no prose, no second bl Constraints: - search_path is set to `financial`; don't qualify table names with the schema. +- **Postgres folds unquoted identifiers to lowercase.** Always double-quote any column or table name that contains an uppercase letter — including all the `district` columns: `"A2"`, `"A3"`, `"A4"`, `"A5"`, `"A6"`, `"A7"`, `"A8"`, `"A9"`, `"A10"`, `"A11"`, `"A12"`, `"A13"`, `"A14"`, `"A15"`, `"A16"`. Mixing quoted and unquoted forms of the same column (`"A3"` and `A13`) will fail at runtime — be consistent. - Quote identifiers that collide with SQL keywords (notably `"order"`). -- Dates in the source are YYMMDD integers (e.g. 950315 → 1995-03-15). Convert with TO_DATE(LPAD(date::text, 6, '0'), 'YYMMDD') whenever you need a real date for filtering, comparison, or display. +- Dates in the source are YYMMDD integers (e.g. 950315 → 1995-03-15). Convert with `TO_DATE(LPAD(date::text, 6, '0'), 'YYMMDD')` whenever you need a real date for filtering, comparison, or display. - Loan `status` values: 'A' finished OK, 'B' finished defaulted, 'C' running OK, 'D' running in debt. - Transaction `type` values: 'PRIJEM' credit, 'VYDAJ' debit. - Read-only: never emit DDL or DML. diff --git a/api/recon/__init__.py b/api/recon/__init__.py new file mode 100644 index 0000000..df2af55 --- /dev/null +++ b/api/recon/__init__.py @@ -0,0 +1,50 @@ +"""Recon package — load the dataset's knowledge graph. + +Typical use: + from api.recon import load_recon + recon = load_recon() + tables = recon.owning_tables("A2") # ["district"] + path = recon.join_path("loan", "district") # ["loan", "account", "district"] + +The Recon is read from `api/datasets//recon.json`. If the +file is missing on first call, it's built on demand (with a log line) so +local dev doesn't require remembering `make recon`. For deploys, run +`python -m api.recon.build` (or `make recon`) to pre-bake. +""" + +from __future__ import annotations + +import json +import logging +from functools import lru_cache +from pathlib import Path + +from api.config import get_settings +from api.datasets import DATASETS_DIR +from api.recon.types import ( + Column, Metric, Recon, Relationship, Table, +) + +logger = logging.getLogger("nvi.recon") + +__all__ = [ + "load_recon", "recon_path", + "Recon", "Column", "Table", "Metric", "Relationship", +] + + +def recon_path(dataset: str | None = None) -> Path: + dataset = dataset or get_settings().dataset + return DATASETS_DIR / dataset / "recon.json" + + +@lru_cache(maxsize=1) +def load_recon() -> Recon: + dataset = get_settings().dataset + path = recon_path(dataset) + if not path.exists(): + logger.info("recon.json missing for %s; building on demand", dataset) + from api.recon.build import build_recon, write_recon + write_recon(dataset, build_recon(dataset)) + data = json.loads(path.read_text()) + return Recon.from_dict(data) diff --git a/api/recon/build.py b/api/recon/build.py new file mode 100644 index 0000000..907a5bb --- /dev/null +++ b/api/recon/build.py @@ -0,0 +1,192 @@ +"""Recon build — two-stage: + +1. **DDL extraction** (auto, from Postgres): reads every table's columns + + types + nullability via SQLAlchemy Inspector and writes + `api/datasets//extracted_schema.json`. This is the structural + source of truth. Re-run whenever the warehouse changes. + +2. **Augmentation merge** (human YAML): reads `schema_docs.yaml` (sparse — + table descriptions, column descriptions, relationships) and `metrics.yaml`, + merges them onto the extracted schema, and writes + `api/datasets//recon.json` — the artefact the runtime reads. + +Run via `make recon` or `python -m api.recon.build`. Auto-triggered by +load_recon() on first call if recon.json is missing. +""" + +from __future__ import annotations + +import argparse +import json +import logging +import sys +from pathlib import Path +from typing import Any + +import yaml +from sqlalchemy import inspect + +from api.datasets import DATASETS_DIR +from api.recon.types import Column, Metric, Recon, Relationship, Table +from api.tools.db import get_engine + +logger = logging.getLogger("nvi.recon.build") + + +# ── Stage 1: DDL extraction ── + +def extract_ddl(dataset: str) -> dict[str, Any]: + """Introspect Postgres for `` schema → structural dict. + + Shape: + { + "schema": "", + "tables": { + "
": { + "columns": [ + {"name": ..., "sql_type": ..., "nullable": ...}, + ... + ] + }, + ... + } + } + """ + insp = inspect(get_engine()) + tables: dict[str, dict[str, Any]] = {} + for name in insp.get_table_names(schema=dataset): + cols = [] + for c in insp.get_columns(name, schema=dataset): + cols.append({ + "name": c["name"], + "sql_type": str(c["type"]).upper(), + "nullable": bool(c["nullable"]), + }) + tables[name] = {"columns": cols} + return {"schema": dataset, "tables": tables} + + +def write_extracted_schema(dataset: str, extracted: dict[str, Any]) -> Path: + out = DATASETS_DIR / dataset / "extracted_schema.json" + out.write_text(json.dumps(extracted, indent=2) + "\n", encoding="utf-8") + return out + + +# ── Stage 2: augmentation merge ── + +def _read_yaml(path: Path) -> dict[str, Any]: + if not path.exists(): + return {} + return yaml.safe_load(path.read_text()) or {} + + +def _parse_column_descs(raw: dict[str, Any]) -> dict[str, str]: + """Accept either the flat form `{table.col: desc}` or the nested form + `{table: {col: desc}}`. Returns flat form.""" + out: dict[str, str] = {} + for k, v in raw.items(): + if isinstance(v, dict): + for col, desc in v.items(): + out[f"{k}.{col}"] = desc + else: + out[k] = v + return out + + +def merge_into_recon(dataset: str, extracted: dict[str, Any]) -> Recon: + """Combine extracted DDL with human-authored augmentation YAML.""" + schema_name = extracted.get("schema", dataset) + + aug = _read_yaml(DATASETS_DIR / dataset / "schema_docs.yaml") + metrics_yaml = _read_yaml(DATASETS_DIR / dataset / "metrics.yaml") + + table_descs: dict[str, str] = aug.get("tables", {}) or {} + col_descs = _parse_column_descs(aug.get("columns", {}) or {}) + + tables: dict[str, Table] = {} + for tname, t_data in extracted["tables"].items(): + cols = [ + Column( + name=c["name"], + sql_type=c["sql_type"], + nullable=c["nullable"], + description=col_descs.get(f"{tname}.{c['name']}"), + ) + for c in t_data["columns"] + ] + tables[tname] = Table( + name=tname, + description=table_descs.get(tname), + columns=cols, + ) + + metrics: dict[str, Metric] = { + name: Metric.from_spec(name, spec) + for name, spec in (metrics_yaml.get("metrics", {}) or {}).items() + } + + relationships = [Relationship.parse(r) for r in (aug.get("relationships", []) or [])] + + column_to_tables: dict[str, list[str]] = {} + for t in tables.values(): + for c in t.columns: + column_to_tables.setdefault(c.name, []).append(t.name) + for k in column_to_tables: + column_to_tables[k] = sorted(column_to_tables[k]) + + return Recon( + schema=schema_name, + tables=tables, + metrics=metrics, + column_to_tables=column_to_tables, + relationships=relationships, + ) + + +def write_recon(dataset: str, recon: Recon) -> Path: + out = DATASETS_DIR / dataset / "recon.json" + out.write_text(json.dumps(recon.to_dict(), indent=2) + "\n", encoding="utf-8") + return out + + +# ── Public entry: full build ── + +def build_recon(dataset: str) -> Recon: + """End-to-end: extract DDL, write the extracted artefact, merge with + YAML, write the recon artefact. Returns the in-memory Recon.""" + extracted = extract_ddl(dataset) + write_extracted_schema(dataset, extracted) + recon = merge_into_recon(dataset, extracted) + write_recon(dataset, recon) + return recon + + +def main() -> None: + logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(name)s %(message)s") + ap = argparse.ArgumentParser(description="Build the recon artefacts for one or all datasets.") + ap.add_argument("--dataset", help="Dataset name (default: all subdirs of api/datasets/).") + args = ap.parse_args() + + if args.dataset: + targets = [args.dataset] + else: + targets = [ + p.name for p in DATASETS_DIR.iterdir() + if p.is_dir() and not p.name.startswith("_") and not p.name.startswith(".") + ] + + if not targets: + print("no datasets to build", file=sys.stderr) + sys.exit(1) + + for name in targets: + recon = build_recon(name) + print( + f"recon[{name}]: {len(recon.tables)} tables, " + f"{len(recon.metrics)} metrics, {len(recon.relationships)} relationships " + f"→ api/datasets/{name}/{{extracted_schema,recon}}.json" + ) + + +if __name__ == "__main__": + main() diff --git a/api/recon/types.py b/api/recon/types.py new file mode 100644 index 0000000..9377759 --- /dev/null +++ b/api/recon/types.py @@ -0,0 +1,250 @@ +"""Recon types — the dataset's knowledge graph as seen by the rest of the api. + +The dataset-model types (Column, Table, Metric, SchemaContext) used to live +in `api/tools/types.py`. They're authored here now because they describe the +dataset, not the tools that consume it. + +`Recon` is a SchemaContext extended with derived indexes that let Analyses +ask things like "which table owns column A2?" or "how do I join loan to +district?" without re-deriving from raw YAML each time. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any + + +# ── Schema model ── + +@dataclass +class Column: + name: str + sql_type: str + nullable: bool + description: str | None = None + + @classmethod + def from_inspector(cls, info: dict[str, Any], description: str | None = None) -> "Column": + """Build from a SQLAlchemy Inspector column dict.""" + return cls( + name=info["name"], + sql_type=str(info["type"]).upper(), + nullable=bool(info["nullable"]), + description=description, + ) + + def to_dict(self) -> dict[str, Any]: + return { + "name": self.name, + "sql_type": self.sql_type, + "nullable": self.nullable, + "description": self.description, + } + + @classmethod + def from_dict(cls, d: dict[str, Any]) -> "Column": + return cls(name=d["name"], sql_type=d["sql_type"], + nullable=d["nullable"], description=d.get("description")) + + +@dataclass +class Table: + name: str + description: str | None + columns: list[Column] + + def to_dict(self) -> dict[str, Any]: + return { + "name": self.name, + "description": self.description, + "columns": [c.to_dict() for c in self.columns], + } + + @classmethod + def from_dict(cls, d: dict[str, Any]) -> "Table": + return cls( + name=d["name"], + description=d.get("description"), + columns=[Column.from_dict(c) for c in d.get("columns", [])], + ) + + +@dataclass +class Metric: + name: str + description: str + sql: str + from_table: str + filter: str | None + unit: str | None + + @classmethod + def from_spec(cls, name: str, spec: dict[str, Any]) -> "Metric": + return cls( + name=name, + description=spec.get("description", ""), + sql=spec["sql"], + from_table=spec["from_table"], + filter=spec.get("filter"), + unit=spec.get("unit"), + ) + + def to_dict(self) -> dict[str, Any]: + return { + "name": self.name, + "description": self.description, + "sql": self.sql, + "from_table": self.from_table, + "filter": self.filter, + "unit": self.unit, + } + + @classmethod + def from_dict(cls, d: dict[str, Any]) -> "Metric": + return cls( + name=d["name"], description=d.get("description", ""), + sql=d["sql"], from_table=d["from_table"], + filter=d.get("filter"), unit=d.get("unit"), + ) + + +# ── Relationships ── + +@dataclass +class Relationship: + from_table: str + from_column: str + to_table: str + to_column: str + + @classmethod + def parse(cls, raw: dict[str, str]) -> "Relationship": + """Accept either `{from: 'a.b', to: 'c.d'}` or `{from_table, from_column, ...}`.""" + if "from_table" in raw: + return cls(raw["from_table"], raw["from_column"], raw["to_table"], raw["to_column"]) + ft, fc = raw["from"].split(".", 1) + tt, tc = raw["to"].split(".", 1) + return cls(ft, fc, tt, tc) + + def to_dict(self) -> dict[str, str]: + return { + "from_table": self.from_table, "from_column": self.from_column, + "to_table": self.to_table, "to_column": self.to_column, + } + + +# ── Recon — top-level dataset knowledge bundle ── + +@dataclass +class Recon: + schema: str + tables: dict[str, Table] + metrics: dict[str, Metric] + column_to_tables: dict[str, list[str]] = field(default_factory=dict) + relationships: list[Relationship] = field(default_factory=list) + + # ── Read-side helpers used by Analyses + text_to_sql ── + + def table_names(self) -> list[str]: + return sorted(self.tables) + + def metric_names(self) -> list[str]: + return sorted(self.metrics) + + def owning_tables(self, column: str) -> list[str]: + """Tables that have a column with this name (case-sensitive).""" + return self.column_to_tables.get(column, []) + + def join_path(self, src: str, dst: str) -> list[str] | None: + """Shortest sequence of tables from src to dst via declared relationships. + Returns None if no path exists. The path includes both endpoints.""" + if src == dst: + return [src] + if src not in self.tables or dst not in self.tables: + return None + # Build undirected adjacency once. + adj: dict[str, set[str]] = {t: set() for t in self.tables} + for r in self.relationships: + adj.setdefault(r.from_table, set()).add(r.to_table) + adj.setdefault(r.to_table, set()).add(r.from_table) + # BFS. + from collections import deque + prev: dict[str, str | None] = {src: None} + q: deque[str] = deque([src]) + while q: + cur = q.popleft() + if cur == dst: + # reconstruct + path: list[str] = [] + node: str | None = cur + while node is not None: + path.append(node) + node = prev[node] + return list(reversed(path)) + for nb in adj.get(cur, ()): + if nb not in prev: + prev[nb] = cur + q.append(nb) + return None + + # ── Prompt-rendering helpers ── + + def render_tables(self, names: list[str] | None = None) -> str: + """CREATE TABLE-like rendering for prompt context.""" + sel = [self.tables[n] for n in (names or self.table_names()) if n in self.tables] + out: list[str] = [] + for t in sel: + header = f"-- {t.description}\n" if t.description else "" + cols: list[str] = [] + for c in t.columns: + line = f' "{c.name}" {c.sql_type}' + if not c.nullable: + line += " NOT NULL" + if c.description: + line += f" -- {c.description}" + cols.append(line) + out.append(header + f'CREATE TABLE "{t.name}" (\n' + ",\n".join(cols) + "\n);") + return "\n\n".join(out) + + def render_metrics(self) -> str: + if not self.metrics: + return "(no metrics defined)" + lines: list[str] = [] + for m in self.metrics.values(): + lines.append( + f"- {m.name} ({m.unit or 'unitless'}): {m.description}\n" + f" sql: {m.sql}\n" + f" from: {m.from_table}" + + (f"\n filter: {m.filter}" if m.filter else "") + ) + return "\n".join(lines) + + def render_relationships(self) -> str: + if not self.relationships: + return "(no relationships declared)" + return "\n".join( + f" {r.from_table}.{r.from_column} → {r.to_table}.{r.to_column}" + for r in self.relationships + ) + + # ── Cache serialisation ── + + def to_dict(self) -> dict[str, Any]: + return { + "schema": self.schema, + "tables": {n: t.to_dict() for n, t in self.tables.items()}, + "metrics": {n: m.to_dict() for n, m in self.metrics.items()}, + "column_to_tables": self.column_to_tables, + "relationships": [r.to_dict() for r in self.relationships], + } + + @classmethod + def from_dict(cls, d: dict[str, Any]) -> "Recon": + return cls( + schema=d["schema"], + tables={n: Table.from_dict(t) for n, t in d.get("tables", {}).items()}, + metrics={n: Metric.from_dict(m) for n, m in d.get("metrics", {}).items()}, + column_to_tables=d.get("column_to_tables", {}), + relationships=[Relationship(**r) for r in d.get("relationships", [])], + ) diff --git a/api/recon/validate.py b/api/recon/validate.py new file mode 100644 index 0000000..df46d0d --- /dev/null +++ b/api/recon/validate.py @@ -0,0 +1,76 @@ +"""SQL validation against the recon schema. + +`validate_sql` parses the SQL, runs sqlglot's `qualify` optimizer pass with +the dataset's column→type schema, and raises with a clear message when the +LLM references a column that doesn't exist on the table it's bound to. + +This is the deterministic complement to the LLM prompt: prompt hints help +the LLM get it right; the validator GUARANTEES we don't ship a schema- +wrong query to Postgres. If validation fails, the caller can either bubble +the error up (no silent retry) or do one *guided* re-prompt that includes +the validator's message — which is a correction with concrete facts, not +a blind retry. +""" + +from __future__ import annotations + +import sqlglot +from sqlglot.errors import OptimizeError +from sqlglot.optimizer.qualify import qualify + +from api.recon import load_recon +from api.recon.types import Recon + + +class ReconValidationError(ValueError): + """The SQL references a table/column combination that doesn't exist + in the recon. The original sqlglot error is in `__cause__`.""" + + +def _build_sqlglot_schema(recon: Recon) -> dict[str, dict[str, str]]: + """Convert recon → sqlglot's expected schema shape: {table: {col: type}}.""" + return { + name: {col.name: col.sql_type for col in t.columns} + for name, t in recon.tables.items() + } + + +def validate_sql(sql: str, recon: Recon | None = None) -> None: + """Raise `ReconValidationError` if any column reference in `sql` doesn't + exist on the table it's bound to. Returns None on success.""" + recon = recon or load_recon() + schema = _build_sqlglot_schema(recon) + try: + parsed = sqlglot.parse_one(sql, dialect="postgres") + qualify(parsed, schema=schema, dialect="postgres") + except OptimizeError as e: + # Try to enrich the message with hints from the recon. + msg = str(e) + hint = _column_hint(msg, recon) + full = f"{msg}{hint}" if hint else msg + raise ReconValidationError(full) from e + + +_COL_PATTERNS = [ + r"Column '([^']+)'", # "Column 'X' could not be resolved." + r"Unknown column:\s*\"?([^\"\s,]+)\"?", # "Unknown column: X" + r"column \"([^\"]+)\"", # 'column "X" does not exist' +] + + +def _column_hint(msg: str, recon: Recon) -> str: + """If the error names a specific column, append a hint about which + table(s) actually own it. Tries the wordings sqlglot and psycopg use.""" + import re + col_name: str | None = None + for pat in _COL_PATTERNS: + m = re.search(pat, msg) + if m: + col_name = m.group(1) + break + if col_name is None: + return "" + owners = recon.owning_tables(col_name) + if owners: + return f" (Column '{col_name}' actually lives in: {', '.join(owners)}.)" + return f" (No table in the {recon.schema!r} schema has a column named '{col_name}'.)" diff --git a/api/runtime/context.py b/api/runtime/context.py index 37dd23e..6b9dc24 100644 --- a/api/runtime/context.py +++ b/api/runtime/context.py @@ -1,8 +1,8 @@ -"""Per-run context — exposes the active run_id to code deep in the call stack -without threading it through every signature. +"""Per-run context — exposes the active run_id (and active Analysis) to code +deep in the call stack without threading them through every signature. -Set by the runtime at run entry; read by Analyses and helper functions so they -can publish trace events without needing the run_id passed in explicitly. +Set by the runtime; read by Analyses, Tools, and helper functions so they +can publish trace events with the right associations. """ from __future__ import annotations @@ -10,3 +10,8 @@ from __future__ import annotations from contextvars import ContextVar current_run_id: ContextVar[str | None] = ContextVar("nvi.current_run_id", default=None) + +# Name of the Analysis currently executing (e.g. "direct_answer"). Set by the +# runtime in `_execute_node` before each Analysis runs, reset after. Used by +# event factories to annotate tool/llm calls so the UI can group them. +current_analysis: ContextVar[str | None] = ContextVar("nvi.current_analysis", default=None) diff --git a/api/runtime/events.py b/api/runtime/events.py index 9ce1fd7..c679e50 100644 --- a/api/runtime/events.py +++ b/api/runtime/events.py @@ -5,6 +5,10 @@ transition; the FastAPI SSE endpoint subscribes and forwards. Event payloads are constructed via the module-level factory functions below so the dict shape lives in one place and is grep-able. + +Many event types carry an `analysis` field that names the active Analysis +(read from `current_analysis` ContextVar). The UI uses it to group +sub-events under the right Analysis node without inferring from time order. """ from __future__ import annotations @@ -14,7 +18,7 @@ import json from typing import Any from api.analyses.types import Finding -from api.runtime.context import current_run_id +from api.runtime.context import current_analysis, current_run_id # run_id -> queue of events. Events are plain dicts. _queues: dict[str, asyncio.Queue[dict[str, Any] | None]] = {} @@ -80,6 +84,8 @@ async def stream(run_id: str): # ── Event factories ── # Each returns the dict payload to pass into `publish()` / # `publish_current()`. Names mirror the `type` field for grep-ability. +# Factories that fire from inside an Analysis pick up the analysis name from +# the `current_analysis` ContextVar so callers don't have to pass it. def run_start(question: str) -> dict[str, Any]: return {"type": "run_start", "question": question} @@ -101,33 +107,66 @@ def node_update(node: str, update: Any) -> dict[str, Any]: } -def analysis_start(name: str, args: dict[str, Any], why: str) -> dict[str, Any]: - return {"type": "analysis_start", "analysis": name, "args": args, "why": why} +def analysis_start(name: str, args: dict[str, Any], why: str, *, + step_id: str | None = None) -> dict[str, Any]: + return {"type": "analysis_start", "analysis": name, "step_id": step_id, + "args": args, "why": why} -def analysis_end(name: str, finding: Finding) -> dict[str, Any]: +def analysis_end(name: str, finding: Finding, *, + step_id: str | None = None) -> dict[str, Any]: return { "type": "analysis_end", "analysis": name, + "step_id": step_id, "summary": finding.summary, "error": finding.error, "row_count": len(finding.rows), } +def analysis_fallback(*, primary: str, fallback: str, reason: str, + step_id: str | None = None) -> dict[str, Any]: + return { + "type": "analysis_fallback", + "primary": primary, + "fallback": fallback, + "reason": reason, + "step_id": step_id, + } + + def synth_start() -> dict[str, Any]: return {"type": "synth_start"} def tool_call_start(tool: str, *, input: Any = None) -> dict[str, Any]: - return {"type": "tool_call_start", "tool": tool, "input": input} + return { + "type": "tool_call_start", + "tool": tool, + "analysis": current_analysis.get(), + "input": input, + } -def tool_call_end(tool: str, *, output: Any = None, error: str | None = None) -> dict[str, Any]: - return {"type": "tool_call_end", "tool": tool, "output": output, "error": error} +def tool_call_end(tool: str, *, output: Any = None, + error: str | None = None) -> dict[str, Any]: + return { + "type": "tool_call_end", + "tool": tool, + "analysis": current_analysis.get(), + "output": output, + "error": error, + } def llm_call(label: str, *, system_len: int, user_len: int) -> dict[str, Any]: """Lightweight breadcrumb for an LLM call. Doesn't carry the prompt contents (those go through Langfuse) — just the fact and rough size.""" - return {"type": "llm_call", "label": label, "system_chars": system_len, "user_chars": user_len} + return { + "type": "llm_call", + "label": label, + "analysis": current_analysis.get(), + "system_chars": system_len, + "user_chars": user_len, + } diff --git a/api/runtime/runner.py b/api/runtime/runner.py index 23875ba..2ce8feb 100644 --- a/api/runtime/runner.py +++ b/api/runtime/runner.py @@ -2,8 +2,9 @@ Graph: START → plan → execute → synthesize → END -`execute` iterates Plan steps and dispatches each to its Analysis. Sequential -for now; swap to a fan-out with parallel nodes once we want to parallelise. +`execute` iterates Plan steps and dispatches each to its Analysis. A step +can declare an optional `fallback` Analysis that runs when the primary +errors or returns an empty finding. """ from __future__ import annotations @@ -21,7 +22,7 @@ from api.llm import chat from api.plan.planner import plan as run_planner from api.prompts import load, render from api.runtime import events -from api.runtime.context import current_run_id +from api.runtime.context import current_analysis, current_run_id from api.runtime.state import RunState logger = logging.getLogger("nvi.runtime") @@ -35,25 +36,74 @@ def _plan_node(state: RunState) -> dict[str, Any]: } +def _step_id(idx: int, name: str, suffix: str | None = None) -> str: + base = f"{name}#{idx}" + return f"{base}.{suffix}" if suffix else base + + +async def _invoke_analysis(name: str, args: dict[str, Any], question: str, + *, step_id: str, run_id: str, why: str) -> Finding: + """Run a single Analysis with contextvar set, events around it, and + error→Finding catch. Returns the Finding (success or failure).""" + analysis = ANALYSES.get(name) + await events.publish(run_id, events.analysis_start(name, args, why, step_id=step_id)) + if analysis is None: + f = Finding(analysis=name, summary="", error=f"unknown analysis {name!r}") + else: + token = current_analysis.set(step_id) + try: + f = await analysis.run(args, question) + except Exception as e: + logger.exception("analysis %s raised", name) + f = Finding(analysis=name, summary="", error=str(e)) + finally: + current_analysis.reset(token) + await events.publish(run_id, events.analysis_end(name, f, step_id=step_id)) + return f + + +def _finding_is_empty_or_errored(f: Finding) -> bool: + if f.error: + return True + if not f.summary and not f.rows: + return True + return False + + async def _execute_node(state: RunState) -> dict[str, Any]: findings: list[dict[str, Any]] = [] - for step in state.get("plan_steps", []): - name = step["analysis"] - analysis = ANALYSES.get(name) - await events.publish( - state["run_id"], - events.analysis_start(name, step.get("args", {}), step.get("why", "")), + run_id = state["run_id"] + for idx, step in enumerate(state.get("plan_steps", [])): + primary_name = step["analysis"] + primary_args = step.get("args", {}) + primary_why = step.get("why", "") + primary_step_id = _step_id(idx, primary_name) + + primary = await _invoke_analysis( + primary_name, primary_args, state["question"], + step_id=primary_step_id, run_id=run_id, why=primary_why, ) - if analysis is None: - f = Finding(analysis=name, summary="", error=f"unknown analysis {name!r}") - else: - try: - f = await analysis.run(step.get("args", {}), state["question"]) - except Exception as e: - logger.exception("analysis %s raised", name) - f = Finding(analysis=name, summary="", error=str(e)) - findings.append(dataclasses.asdict(f)) - await events.publish(state["run_id"], events.analysis_end(name, f)) + primary_dict = dataclasses.asdict(primary) + primary_dict["step_id"] = primary_step_id + findings.append(primary_dict) + + fb = step.get("fallback") + if fb and _finding_is_empty_or_errored(primary): + fb_name = fb["analysis"] + fb_step_id = _step_id(idx, fb_name, suffix="fallback") + reason = primary.error or "primary returned empty finding" + await events.publish(run_id, events.analysis_fallback( + primary=primary_name, fallback=fb_name, reason=reason, + step_id=fb_step_id, + )) + fb_finding = await _invoke_analysis( + fb_name, fb.get("args", {}), state["question"], + step_id=fb_step_id, run_id=run_id, why=fb.get("why", ""), + ) + fb_dict = dataclasses.asdict(fb_finding) + fb_dict["step_id"] = fb_step_id + fb_dict["is_fallback_for"] = primary_step_id + findings.append(fb_dict) return {"findings": findings} @@ -61,14 +111,17 @@ async def _synthesize_node(state: RunState) -> dict[str, Any]: await events.publish(state["run_id"], events.synth_start()) findings = state.get("findings", []) - if not findings or all(f.get("error") for f in findings): + # Effective findings = primary OR its fallback if the primary was empty. + # Pass everything through; the LLM weighs which one to cite. + usable = [f for f in findings if not _finding_is_empty_or_errored(_dict_to_finding(f))] + if not usable: return { "answer": "I couldn't answer this question — see the per-analysis errors above.", "error": "all_analyses_failed", } findings_block = "\n\n".join( - f"[{f['analysis']}]\n" + f"[{f['analysis']}{ ' (fallback)' if f.get('is_fallback_for') else '' }]\n" f"summary: {f.get('summary') or '(none)'}\n" f"error: {f.get('error') or '(none)'}\n" f"sql: {f.get('sql')}" @@ -89,6 +142,17 @@ async def _synthesize_node(state: RunState) -> dict[str, Any]: return {"answer": answer} +def _dict_to_finding(d: dict[str, Any]) -> Finding: + return Finding( + analysis=d.get("analysis", ""), + summary=d.get("summary", ""), + rows=d.get("rows") or [], + sql=d.get("sql") or [], + error=d.get("error"), + metadata=d.get("metadata") or {}, + ) + + def build_graph(): g: StateGraph = StateGraph(RunState) g.add_node("plan", _plan_node) diff --git a/api/tools/schema.py b/api/tools/schema.py deleted file mode 100644 index 84df945..0000000 --- a/api/tools/schema.py +++ /dev/null @@ -1,57 +0,0 @@ -"""Warehouse schema retrieval. - -Physical layer via SQLAlchemy's Inspector against the active dataset's -Postgres schema; semantic layer (table/column descriptions, metric catalog) -from YAML files under `api/datasets//`. - -Cached on first call — the schema is stable for a process's lifetime. -Restart the api to pick up a dataset switch or YAML edits. -""" - -from __future__ import annotations - -from functools import lru_cache -from typing import Any - -from sqlalchemy import inspect - -from api.config import get_settings -from api.datasets import read_yaml -from api.tools.db import get_engine -from api.tools.types import Column, Metric, SchemaContext, Table - - -@lru_cache(maxsize=1) -def load_schema() -> SchemaContext: - dataset = get_settings().dataset - docs = read_yaml(dataset, "schema_docs.yaml") - metrics_yaml = read_yaml(dataset, "metrics.yaml") - - # Postgres schema name comes from the dataset id; the schema_docs file - # can override it, but normally they match. - schema_name: str = docs.get("schema", dataset) - table_descs: dict[str, dict[str, Any]] = docs.get("tables", {}) or {} - - insp = inspect(get_engine()) - tables: dict[str, Table] = { - name: Table( - name=name, - description=(desc := table_descs.get(name, {})).get("description"), - columns=[ - Column.from_inspector(c, (desc.get("columns") or {}).get(c["name"])) - for c in insp.get_columns(name, schema=schema_name) - ], - ) - for name in insp.get_table_names(schema=schema_name) - } - - metrics: dict[str, Metric] = { - name: Metric.from_spec(name, spec) - for name, spec in (metrics_yaml.get("metrics", {}) or {}).items() - } - - return SchemaContext(schema=schema_name, tables=tables, metrics=metrics) - - -def get_metric(name: str) -> Metric | None: - return load_schema().metrics.get(name) diff --git a/api/tools/text_to_sql.py b/api/tools/text_to_sql.py index ad01ac4..a5b6591 100644 --- a/api/tools/text_to_sql.py +++ b/api/tools/text_to_sql.py @@ -10,23 +10,24 @@ import sqlglot from api import langfuse_client as lf from api.llm import chat from api.prompts import load, render -from api.tools.schema import load_schema +from api.recon import load_recon +from api.recon.validate import ReconValidationError, validate_sql from api.tools.types import T2SResult logger = logging.getLogger("nvi.tools.text_to_sql") def text_to_sql(question: str, *, hint_tables: list[str] | None = None) -> T2SResult: - schema = load_schema() - tables = hint_tables or schema.table_names() + recon = load_recon() + tables = hint_tables or recon.table_names() system = load("text_to_sql.system") def _user(retry_hint: str = "") -> str: return render( "text_to_sql.user", question=question, - schema_block=schema.render_tables(tables), - metrics_block=schema.render_metrics(), + schema_block=recon.render_tables(tables), + metrics_block=recon.render_metrics(), retry_hint=retry_hint, ) @@ -35,14 +36,9 @@ def text_to_sql(question: str, *, hint_tables: list[str] | None = None) -> T2SRe as_type="generation", input={"question": question, "hint_tables": hint_tables}, ) as span: - sql = _extract_sql(chat(system=system, user=_user())) - try: - _validate(sql) - except Exception as e: - logger.info("t2s first attempt invalid (%s); retrying once", e) - hint = f"\n\nPrevious attempt failed parsing: {e}. Fix and return only the SQL." - sql = _extract_sql(chat(system=system, user=_user(hint))) - _validate(sql) + raw = _extract_sql(chat(system=system, user=_user())) + sql = _normalize(raw) # raises on parse error — fail fast + validate_sql(sql, recon) # raises ReconValidationError — fail fast result = T2SResult(sql=sql, used_tables=_extract_tables(sql)) span.update(output={"sql": sql, "used_tables": result.used_tables}) @@ -56,10 +52,20 @@ def _extract_sql(text: str) -> str: return text.strip().rstrip(";").strip() -def _validate(sql: str) -> None: +def _normalize(sql: str) -> str: + """Parse the LLM's SQL, then re-render with every identifier quoted. + + Postgres folds unquoted identifiers to lowercase, so `d.A13` becomes + `d.a13` and breaks against case-preserving columns like district."A13". + Re-rendering with `identify=True` puts double quotes around every + identifier — case is preserved, and Postgres treats a quoted lowercase + identifier the same as the unquoted form, so this is safe in both + directions. + """ parsed = sqlglot.parse_one(sql, dialect="postgres") if parsed is None: raise ValueError("sqlglot returned no parse tree") + return parsed.sql(dialect="postgres", identify=True) def _extract_tables(sql: str) -> list[str]: diff --git a/api/tools/types.py b/api/tools/types.py index d76790f..f913b84 100644 --- a/api/tools/types.py +++ b/api/tools/types.py @@ -1,11 +1,7 @@ """Public data shapes for the Tools layer. -Behavior modules in api/tools/ import from here. Keeping types in one file -makes the layer's surface visible at a glance. - -Secondary constructors (`from_inspector`, `from_spec`, …) live as -classmethods so call sites read as one line instead of multi-line kwarg -blocks. +Tool-output types only. The dataset model (Column, Table, Metric, Recon) +lives in `api/recon/types.py` — Tools consume it but don't own it. """ from __future__ import annotations @@ -14,100 +10,6 @@ from dataclasses import dataclass from typing import Any -# ── Schema model ── - -@dataclass -class Column: - name: str - sql_type: str - nullable: bool - description: str | None = None - - @classmethod - def from_inspector(cls, info: dict[str, Any], description: str | None = None) -> "Column": - """Build from a SQLAlchemy Inspector column dict.""" - return cls( - name=info["name"], - sql_type=str(info["type"]).upper(), - nullable=bool(info["nullable"]), - description=description, - ) - - -@dataclass -class Table: - name: str - description: str | None - columns: list[Column] - - -@dataclass -class Metric: - name: str - description: str - sql: str - from_table: str - filter: str | None - unit: str | None - - @classmethod - def from_spec(cls, name: str, spec: dict[str, Any]) -> "Metric": - """Build from a metrics.yaml entry.""" - return cls( - name=name, - description=spec.get("description", ""), - sql=spec["sql"], - from_table=spec["from_table"], - filter=spec.get("filter"), - unit=spec.get("unit"), - ) - - -@dataclass -class SchemaContext: - schema: str - tables: dict[str, Table] - metrics: dict[str, Metric] - - def table_names(self) -> list[str]: - return sorted(self.tables) - - def metric_names(self) -> list[str]: - return sorted(self.metrics) - - def render_tables(self, names: list[str] | None = None) -> str: - """CREATE TABLE-like rendering for prompt context.""" - sel = [self.tables[n] for n in (names or self.table_names()) if n in self.tables] - out: list[str] = [] - for t in sel: - header = f"-- {t.description}\n" if t.description else "" - cols: list[str] = [] - for c in t.columns: - line = f' "{c.name}" {c.sql_type}' - if not c.nullable: - line += " NOT NULL" - if c.description: - line += f" -- {c.description}" - cols.append(line) - out.append(header + f'CREATE TABLE "{t.name}" (\n' + ",\n".join(cols) + "\n);") - return "\n\n".join(out) - - def render_metrics(self) -> str: - if not self.metrics: - return "(no metrics defined)" - lines: list[str] = [] - for m in self.metrics.values(): - lines.append( - f"- {m.name} ({m.unit or 'unitless'}): {m.description}\n" - f" sql: {m.sql}\n" - f" from: {m.from_table}" - + (f"\n filter: {m.filter}" if m.filter else "") - ) - return "\n".join(lines) - - -# ── Query execution ── - @dataclass class QueryResult: columns: list[str] @@ -119,8 +21,6 @@ class QueryResult: return [dict(zip(self.columns, r)) for r in self.rows] -# ── Text-to-SQL output ── - @dataclass class T2SResult: sql: str diff --git a/ctrl/k8s/base/api.yaml b/ctrl/k8s/base/api.yaml index d3dc661..0736283 100644 --- a/ctrl/k8s/base/api.yaml +++ b/ctrl/k8s/base/api.yaml @@ -13,6 +13,14 @@ spec: labels: app: api spec: + # lng.local.ar resolves to ::1 via the host's dnsmasq → coredns chain, + # which is the pod's own loopback (wrong). Override to the kind bridge + # gateway so traffic from this pod reaches the host's system Caddy and + # gets reverse-proxied to Langfuse on :3000. + hostAliases: + - ip: "172.19.0.1" + hostnames: + - "lng.local.ar" containers: - name: api image: nvi-api diff --git a/ctrl/recon.sh b/ctrl/recon.sh new file mode 100755 index 0000000..e7feb46 --- /dev/null +++ b/ctrl/recon.sh @@ -0,0 +1,46 @@ +#!/usr/bin/env bash +# Rebuild the per-dataset recon artefacts inside the api pod and copy them +# back to the host so they're visible next to the source YAMLs. +# +# Outputs per dataset under api/datasets//: +# - extracted_schema.json (auto, from Postgres introspection) +# - recon.json (merged: extracted + schema_docs.yaml + metrics.yaml) +# +# Usage: +# bash ctrl/recon.sh # all datasets +# bash ctrl/recon.sh financial # one dataset + +set -euo pipefail +cd "$(dirname "$0")/.." + +CTX="${KCTX:---context kind-nvi}" +NS="${KNS:--n nvi}" + +POD=$(kubectl $CTX $NS get pod -l app=api -o jsonpath='{.items[0].metadata.name}') +if [ -z "$POD" ]; then + echo "no api pod found; is \`make tilt-up\` running?" >&2 + exit 1 +fi + +# Run the build inside the pod (needs live Postgres access). +build_args=() +if [ "${1-}" != "" ]; then + build_args+=(--dataset "$1") +fi +kubectl $CTX $NS exec "$POD" -- uv run python -m api.recon.build "${build_args[@]}" + +# Pull each artefact back to the host. +for ds_path in api/datasets/*/; do + ds=$(basename "$ds_path") + case "$ds" in + __*|.*) continue ;; + esac + if [ "${1-}" != "" ] && [ "$1" != "$ds" ]; then + continue + fi + for f in extracted_schema.json recon.json; do + if kubectl $CTX $NS cp "$POD:/app/api/datasets/$ds/$f" "api/datasets/$ds/$f" 2>/dev/null; then + echo " → api/datasets/$ds/$f" + fi + done +done diff --git a/tests/test_analyses_registry.py b/tests/test_analyses_registry.py new file mode 100644 index 0000000..e17f2a3 --- /dev/null +++ b/tests/test_analyses_registry.py @@ -0,0 +1,23 @@ +"""Regression: Analysis registry exposes the expected Analyses.""" + +from api.analyses.registry import REGISTRY, catalog, get + + +def test_registry_has_expected_analyses(): + assert set(REGISTRY) >= {"direct_answer", "compare_periods", "drill_down"} + + +def test_registry_get_returns_instance_or_none(): + a = get("direct_answer") + assert a is not None + assert a.name == "direct_answer" + assert get("does_not_exist") is None + + +def test_catalog_shape_for_planner(): + cat = catalog() + names = {entry["name"] for entry in cat} + assert names >= {"direct_answer", "compare_periods", "drill_down"} + for entry in cat: + assert isinstance(entry["description"], str) and entry["description"] + assert isinstance(entry["args_schema"], dict) diff --git a/tests/test_events.py b/tests/test_events.py index 6158417..9120ff0 100644 --- a/tests/test_events.py +++ b/tests/test_events.py @@ -2,6 +2,7 @@ from api.analyses.types import Finding from api.runtime import events +from api.runtime.context import current_analysis def test_run_start_shape(): @@ -22,10 +23,11 @@ def test_plan_ready_shape(): def test_analysis_start_shape(): - e = events.analysis_start("direct_answer", {"q": "?"}, "fits") + e = events.analysis_start("direct_answer", {"q": "?"}, "fits", step_id="direct_answer#0") assert e == { "type": "analysis_start", "analysis": "direct_answer", + "step_id": "direct_answer#0", "args": {"q": "?"}, "why": "fits", } @@ -33,19 +35,65 @@ def test_analysis_start_shape(): def test_analysis_end_uses_finding(): f = Finding(analysis="direct_answer", summary="ok", rows=[{"a": 1}], sql=["SELECT 1"]) - e = events.analysis_end("direct_answer", f) + e = events.analysis_end("direct_answer", f, step_id="direct_answer#0") assert e["type"] == "analysis_end" assert e["analysis"] == "direct_answer" + assert e["step_id"] == "direct_answer#0" assert e["summary"] == "ok" assert e["error"] is None assert e["row_count"] == 1 -def test_tool_call_shapes(): +def test_analysis_fallback_shape(): + e = events.analysis_fallback( + primary="drill_down", fallback="direct_answer", + reason="empty result", step_id="direct_answer#0.fallback", + ) + assert e == { + "type": "analysis_fallback", + "primary": "drill_down", + "fallback": "direct_answer", + "reason": "empty result", + "step_id": "direct_answer#0.fallback", + } + + +def test_tool_call_shapes_carry_analysis_field(): + """Tool/LLM events must include the active Analysis name from the + contextvar — UI groups sub-events by it. None when outside an Analysis.""" + # Outside any Analysis: analysis is None. s = events.tool_call_start("text_to_sql", input={"q": "?"}) - assert s == {"type": "tool_call_start", "tool": "text_to_sql", "input": {"q": "?"}} + assert s["type"] == "tool_call_start" + assert s["tool"] == "text_to_sql" + assert s["analysis"] is None + assert s["input"] == {"q": "?"} + e = events.tool_call_end("text_to_sql", output={"sql": "..."}) - assert e == {"type": "tool_call_end", "tool": "text_to_sql", "output": {"sql": "..."}, "error": None} + assert e["analysis"] is None + assert e["output"] == {"sql": "..."} + assert e["error"] is None + + # Inside an Analysis context: analysis field is populated. + token = current_analysis.set("direct_answer#0") + try: + s2 = events.tool_call_start("execute_sql", input={"sql": "SELECT 1"}) + e2 = events.tool_call_end("execute_sql", output={"row_count": 1}) + assert s2["analysis"] == "direct_answer#0" + assert e2["analysis"] == "direct_answer#0" + finally: + current_analysis.reset(token) + + +def test_llm_call_carries_analysis_field(): + token = current_analysis.set("compare_periods#0") + try: + e = events.llm_call("compare_periods.interpret", system_len=100, user_len=200) + assert e["type"] == "llm_call" + assert e["analysis"] == "compare_periods#0" + assert e["system_chars"] == 100 + assert e["user_chars"] == 200 + finally: + current_analysis.reset(token) def test_synth_start_shape(): diff --git a/tests/test_imports.py b/tests/test_imports.py index 46eca73..976779a 100644 --- a/tests/test_imports.py +++ b/tests/test_imports.py @@ -17,8 +17,10 @@ MODULES = [ "api.llm", "api.datasets", "api.prompts", + "api.recon", + "api.recon.types", + "api.recon.build", "api.tools.db", - "api.tools.schema", "api.tools.text_to_sql", "api.tools.execute_sql", "api.tools.types", diff --git a/tests/test_plan_types.py b/tests/test_plan_types.py index 49542cc..ef213c0 100644 --- a/tests/test_plan_types.py +++ b/tests/test_plan_types.py @@ -29,3 +29,28 @@ def test_plan_roundtrip(): assert plan.rationale == "single step suffices" assert len(plan.steps) == 1 assert plan.steps[0].analysis == "direct_answer" + + +def test_planstep_fallback_roundtrip(): + raw = { + "analysis": "drill_down", + "args": {"metric": "amount"}, + "why": "speculative", + "fallback": { + "analysis": "direct_answer", + "args": {"question": "fallback q"}, + "why": "safe option", + }, + } + step = PlanStep.from_dict(raw) + assert step.fallback is not None + assert step.fallback.analysis == "direct_answer" + assert step.fallback.args == {"question": "fallback q"} + assert step.fallback.why == "safe option" + assert step.to_dict() == raw + + +def test_planstep_no_fallback_omits_key(): + step = PlanStep.from_dict({"analysis": "direct_answer"}) + assert step.fallback is None + assert "fallback" not in step.to_dict() diff --git a/tests/test_recon.py b/tests/test_recon.py new file mode 100644 index 0000000..3dd6452 --- /dev/null +++ b/tests/test_recon.py @@ -0,0 +1,106 @@ +"""Regression tests for the recon package — type round-trips, join_path, +and column_to_tables index.""" + +from api.recon.types import Recon, Relationship, Table, Column, Metric + + +def _sample_recon() -> Recon: + cols_account = [ + Column("account_id", "BIGINT", False), + Column("district_id", "BIGINT", True), + ] + cols_district = [ + Column("district_id", "BIGINT", False), + Column("A2", "TEXT", True), + Column("A11", "BIGINT", True), + ] + cols_loan = [ + Column("loan_id", "BIGINT", False), + Column("account_id", "BIGINT", True), + Column("amount", "BIGINT", True), + Column("status", "TEXT", True), + ] + rels = [ + Relationship("account", "district_id", "district", "district_id"), + Relationship("loan", "account_id", "account", "account_id"), + ] + c2t = { + "account_id": sorted(["account", "loan"]), + "district_id": sorted(["account", "district"]), + "A2": ["district"], + "A11": ["district"], + "loan_id": ["loan"], + "amount": ["loan"], + "status": ["loan"], + } + return Recon( + schema="financial", + tables={ + "account": Table("account", "Accounts", cols_account), + "district": Table("district", "Districts", cols_district), + "loan": Table("loan", "Loans", cols_loan), + }, + metrics={ + "loan_volume": Metric("loan_volume", "Total CZK lent", + "SUM(amount)", "loan", None, "CZK"), + }, + column_to_tables=c2t, + relationships=rels, + ) + + +def test_owning_tables(): + r = _sample_recon() + assert r.owning_tables("A2") == ["district"] + assert r.owning_tables("amount") == ["loan"] + assert r.owning_tables("account_id") == ["account", "loan"] + assert r.owning_tables("nonsense") == [] + + +def test_join_path_direct(): + r = _sample_recon() + assert r.join_path("loan", "account") == ["loan", "account"] + assert r.join_path("account", "district") == ["account", "district"] + + +def test_join_path_multi_hop(): + r = _sample_recon() + assert r.join_path("loan", "district") == ["loan", "account", "district"] + + +def test_join_path_same(): + r = _sample_recon() + assert r.join_path("loan", "loan") == ["loan"] + + +def test_join_path_unknown_returns_none(): + r = _sample_recon() + assert r.join_path("loan", "does_not_exist") is None + + +def test_recon_to_dict_from_dict_roundtrip(): + r = _sample_recon() + d = r.to_dict() + r2 = Recon.from_dict(d) + assert r2.schema == r.schema + assert sorted(r2.tables) == sorted(r.tables) + assert sorted(r2.metrics) == sorted(r.metrics) + assert r2.column_to_tables == r.column_to_tables + assert len(r2.relationships) == len(r.relationships) + assert r2.join_path("loan", "district") == ["loan", "account", "district"] + + +def test_relationship_parse_dotted(): + r = Relationship.parse({"from": "loan.account_id", "to": "account.account_id"}) + assert r.from_table == "loan" + assert r.from_column == "account_id" + assert r.to_table == "account" + assert r.to_column == "account_id" + + +def test_relationship_parse_explicit(): + r = Relationship.parse({ + "from_table": "loan", "from_column": "account_id", + "to_table": "account", "to_column": "account_id", + }) + assert r.from_table == "loan" diff --git a/tests/test_recon_validate.py b/tests/test_recon_validate.py new file mode 100644 index 0000000..d5fb2c5 --- /dev/null +++ b/tests/test_recon_validate.py @@ -0,0 +1,94 @@ +"""Regression: recon SQL validator catches LLM-generated schema mismatches. + +Bug captured: drill_down generated `SELECT district_id FROM loan` — +syntactically valid, parses, but `loan` doesn't have a `district_id` column. +We want to catch this BEFORE Postgres does, with a clear message naming +where the column actually lives. +""" + +import pytest + +from api.recon.types import Column, Metric, Recon, Relationship, Table +from api.recon.validate import ReconValidationError, validate_sql + + +def _recon() -> Recon: + return Recon( + schema="financial", + tables={ + "loan": Table("loan", None, [ + Column("loan_id", "BIGINT", False), + Column("account_id", "BIGINT", True), + Column("amount", "BIGINT", True), + Column("status", "TEXT", True), + ]), + "account": Table("account", None, [ + Column("account_id", "BIGINT", False), + Column("district_id", "BIGINT", True), + ]), + "district": Table("district", None, [ + Column("district_id", "BIGINT", False), + Column("A2", "TEXT", True), + Column("A13", "BIGINT", True), + ]), + }, + metrics={}, + column_to_tables={ + "loan_id": ["loan"], + "account_id": ["account", "loan"], + "amount": ["loan"], + "status": ["loan"], + "district_id": ["account", "district"], + "A2": ["district"], + "A13": ["district"], + }, + relationships=[ + Relationship("account", "district_id", "district", "district_id"), + Relationship("loan", "account_id", "account", "account_id"), + ], + ) + + +def test_valid_sql_passes(): + sql = "SELECT account_id, amount FROM loan WHERE status = 'B'" + validate_sql(sql, _recon()) + + +def test_valid_join_passes(): + sql = """ + SELECT d.A2, AVG(l.amount) + FROM loan l + JOIN account a ON l.account_id = a.account_id + JOIN district d ON a.district_id = d.district_id + GROUP BY d.A2 + """ + validate_sql(sql, _recon()) + + +def test_unknown_column_on_table_raises(): + # The exact bug we hit in production: district_id picked from loan + # without joining through account. + sql = """ + WITH x AS (SELECT district_id, AVG(amount) FROM loan GROUP BY district_id) + SELECT * FROM x + """ + with pytest.raises(ReconValidationError) as exc: + validate_sql(sql, _recon()) + msg = str(exc.value) + assert "district_id" in msg + # The hint should name where district_id actually lives. + assert "account" in msg and "district" in msg + + +def test_misspelled_column_raises(): + sql = "SELECT amnt FROM loan" + with pytest.raises(ReconValidationError) as exc: + validate_sql(sql, _recon()) + assert "amnt" in str(exc.value) + + +def test_qualified_column_on_wrong_table_raises(): + # district.amount doesn't exist (amount is on loan). + sql = "SELECT d.amount FROM district d" + with pytest.raises(ReconValidationError): + validate_sql(sql, _recon()) diff --git a/tests/test_yaml_datasets.py b/tests/test_yaml_datasets.py index 8d1d5d2..33702b5 100644 --- a/tests/test_yaml_datasets.py +++ b/tests/test_yaml_datasets.py @@ -3,6 +3,11 @@ Bug captured: schema_docs.yaml had values that began with a quote character (e.g. `gender: 'M' or 'F'.`), which YAML treats as a quoted scalar — anything after the closing quote is a parse error. + +The YAML structure is now flat — tables. is a description string, +columns are keyed `table.column`, and relationships are an explicit list. +The DDL structure (column types, nullability) lives in extracted_schema.json +generated by `make recon`. """ from pathlib import Path @@ -20,6 +25,9 @@ def test_schema_docs_parses(dataset): assert docs is not None assert "tables" in docs assert isinstance(docs["tables"], dict) and docs["tables"] + # New shape: each table value is a description string, not a dict. + for name, val in docs["tables"].items(): + assert isinstance(val, str), f"tables.{name} should be a string description, got {type(val)}" @pytest.mark.parametrize("dataset", DATASETS) @@ -31,16 +39,24 @@ def test_metrics_parses(dataset): def test_financial_columns_with_quotes_survive_parse(): - """Specific regression: values that contain quoted codes (M, F, OWNER, - PRIJEM, etc.) must parse without truncation.""" + """Values containing quoted codes (M, F, OWNER, PRIJEM, ...) must parse + without truncation. Columns are flat `table.col` keys now.""" docs = yaml.safe_load( (DATASETS_DIR / "financial" / "schema_docs.yaml").read_text() ) - cols = docs["tables"]["client"]["columns"] - assert '"M"' in cols["gender"] and '"F"' in cols["gender"] + cols = docs["columns"] + assert '"M"' in cols["client.gender"] and '"F"' in cols["client.gender"] + assert '"OWNER"' in cols["disp.type"] and '"DISPONENT"' in cols["disp.type"] + assert '"PRIJEM"' in cols["trans.type"] and '"VYDAJ"' in cols["trans.type"] - disp = docs["tables"]["disp"]["columns"] - assert '"OWNER"' in disp["type"] and '"DISPONENT"' in disp["type"] - trans = docs["tables"]["trans"]["columns"] - assert '"PRIJEM"' in trans["type"] and '"VYDAJ"' in trans["type"] +def test_financial_has_relationships(): + docs = yaml.safe_load( + (DATASETS_DIR / "financial" / "schema_docs.yaml").read_text() + ) + rels = docs.get("relationships") or [] + assert len(rels) >= 5 # we have 8; assert ≥5 to allow some pruning + # all entries have from/to + for r in rels: + assert "from" in r and "to" in r + assert "." in r["from"] and "." in r["to"] diff --git a/ui/app/src/composables/useRunStream.ts b/ui/app/src/composables/useRunStream.ts index 7170b54..81fb18a 100644 --- a/ui/app/src/composables/useRunStream.ts +++ b/ui/app/src/composables/useRunStream.ts @@ -1,5 +1,4 @@ -import { ref, onUnmounted, computed } from 'vue' -import type { LogEntry } from 'soleprint-ui' +import { computed, onUnmounted, ref } from 'vue' import { apiUrl } from '../config' /* Mirror of the backend factories in api/runtime/events.py */ @@ -9,26 +8,41 @@ export interface PlanStep { analysis: string args: Record why: string + fallback?: PlanStep | null } +export type NodeKind = 'plan' | 'analysis' | 'fallback' | 'synthesize' + export interface GraphNode { - id: string // e.g. "plan", "direct_answer", "synthesize" - kind: 'plan' | 'analysis' | 'synthesize' + /** Stable id used for scroll-into-view + bidirectional links */ + id: string + kind: NodeKind label: string - status: 'pending' | 'running' | 'done' | 'error' + status: 'pending' | 'running' | 'done' | 'error' | 'skipped' detail?: string + /** For fallback nodes: id of the primary they back. */ + fallbackFor?: string } export interface ToolCall { + /** owning Analysis node id (e.g. "direct_answer#0"); null for run-level calls */ + analysisId: string | null tool: string input?: any output?: any error?: string | null - /** ms since the run started */ t_started: number t_ended?: number } +export interface LogEntry { + level: string + stage: string + msg: string + ts: string +} + + export function useRunStream() { const status = ref('idle') const runId = ref(null) @@ -80,6 +94,19 @@ export function useRunStream() { } } + /** Sub-events (tool calls + LLM calls) grouped by their owning Analysis + * step_id. UI uses this to render the nested sub-graph in each Analysis + * expansion. */ + const subEventsByAnalysis = computed(() => { + const out: Record = {} + for (const c of toolCalls.value) { + const key = c.analysisId ?? '__root__' + if (!out[key]) out[key] = [] + out[key]!.push(c) + } + return out + }) + async function ask(q: string): Promise { reset() status.value = 'connecting' @@ -140,36 +167,61 @@ export function useRunStream() { rationale.value = d.rationale planSteps.value = d.steps - // Replace placeholder analysis nodes with real ones from the plan. const planNode = nodeBy('plan')! planNode.status = 'done' planNode.detail = d.rationale - const synth = graphNodes.value.find(n => n.id === 'synthesize') - graphNodes.value = [ - planNode, - ...d.steps.map((s, i): GraphNode => ({ - id: `${s.analysis}#${i}`, + + // Build the chain: plan → step[0] (and its fallback) → step[1] → ... → synth. + const synth: GraphNode = nodeBy('synthesize') + || { id: 'synthesize', kind: 'synthesize', label: 'Synthesize', status: 'pending' } + + const nodes: GraphNode[] = [planNode] + d.steps.forEach((s, i) => { + const primaryId = `${s.analysis}#${i}` + nodes.push({ + id: primaryId, kind: 'analysis', label: s.analysis, status: 'pending', detail: s.why, - })), - synth || { id: 'synthesize', kind: 'synthesize', label: 'Synthesize', status: 'pending' }, - ] - push('info', 'plan', `${d.steps.length} step(s): ${d.steps.map(s => s.analysis).join(' → ')}`) + }) + if (s.fallback) { + nodes.push({ + id: `${s.fallback.analysis}#${i}.fallback`, + kind: 'fallback', + label: s.fallback.analysis, + status: 'skipped', // dimmed until proven needed + detail: `fallback for ${s.analysis}: ${s.fallback.why}`, + fallbackFor: primaryId, + }) + } + }) + nodes.push(synth) + graphNodes.value = nodes + + push('info', 'plan', `${d.steps.length} step(s): ${d.steps.map(s => + s.fallback ? `${s.analysis} (fallback: ${s.fallback.analysis})` : s.analysis, + ).join(' → ')}`) }) es.addEventListener('analysis_start', (e: MessageEvent) => { - const d = JSON.parse(e.data) as { analysis: string; args: any; why: string } - // Mark the first matching pending analysis node as running. - const n = graphNodes.value.find(g => g.kind === 'analysis' && g.label === d.analysis && g.status === 'pending') - if (n) n.status = 'running' + const d = JSON.parse(e.data) as { analysis: string; step_id: string; args: any; why: string } + // Prefer the step_id-based match (works for both primary and fallback). + const n = nodeBy(d.step_id) || graphNodes.value.find( + g => (g.kind === 'analysis' || g.kind === 'fallback') && g.label === d.analysis && g.status === 'pending' + ) + if (n) { + n.id = d.step_id // ensure id matches step_id for bidirectional linking + n.status = 'running' + } push('info', d.analysis, `start: ${d.why || JSON.stringify(d.args)}`) }) es.addEventListener('analysis_end', (e: MessageEvent) => { - const d = JSON.parse(e.data) as { analysis: string; summary: string; error: string | null; row_count: number } - const n = graphNodes.value.find(g => g.kind === 'analysis' && g.label === d.analysis && g.status === 'running') + const d = JSON.parse(e.data) as { + analysis: string; step_id: string; summary: string; error: string | null; row_count: number + } + const n = nodeBy(d.step_id) if (n) { n.status = d.error ? 'error' : 'done' n.detail = d.error || d.summary @@ -179,9 +231,26 @@ export function useRunStream() { push(level, d.analysis, msg) }) + es.addEventListener('analysis_fallback', (e: MessageEvent) => { + const d = JSON.parse(e.data) as { + primary: string; fallback: string; reason: string; step_id: string + } + const fbNode = nodeBy(d.step_id) + if (fbNode) { + fbNode.status = 'running' // promoted from 'skipped' + fbNode.detail = `fallback fired — ${d.reason}` + } + push('warn', 'fallback', `${d.primary} → ${d.fallback}: ${d.reason}`) + }) + es.addEventListener('tool_call_start', (e: MessageEvent) => { - const d = JSON.parse(e.data) as { tool: string; input: any } - toolCalls.value.push({ tool: d.tool, input: d.input, t_started: elapsed() }) + const d = JSON.parse(e.data) as { tool: string; analysis: string | null; input: any } + toolCalls.value.push({ + tool: d.tool, + analysisId: d.analysis, + input: d.input, + t_started: elapsed(), + }) const detail = d.tool === 'execute_sql' ? (d.input?.sql || '') : d.tool === 'text_to_sql' ? (d.input?.question || '') : JSON.stringify(d.input) @@ -189,8 +258,10 @@ export function useRunStream() { }) es.addEventListener('tool_call_end', (e: MessageEvent) => { - const d = JSON.parse(e.data) as { tool: string; output: any; error: string | null } - const call = [...toolCalls.value].reverse().find(c => c.tool === d.tool && c.t_ended === undefined) + const d = JSON.parse(e.data) as { tool: string; analysis: string | null; output: any; error: string | null } + const call = [...toolCalls.value].reverse().find( + c => c.tool === d.tool && c.analysisId === d.analysis && c.t_ended === undefined, + ) if (call) { call.output = d.output call.error = d.error @@ -203,7 +274,8 @@ export function useRunStream() { if (d.tool === 'text_to_sql' && d.output?.sql) { push('info', d.tool, `✓ ${d.output.sql}`) } else if (d.tool === 'execute_sql') { - const label = d.output?.label ? `${d.output.label}: ` : '' + const label = d.output?.dimension ? `${d.output.dimension}: ` : + d.output?.label ? `${d.output.label}: ` : '' push('info', d.tool, `✓ ${label}${d.output?.row_count} rows`) } else if (d.tool === 'generate_pair') { push('info', d.tool, `✓ a=${d.output?.a?.label} · b=${d.output?.b?.label}`) @@ -213,7 +285,7 @@ export function useRunStream() { }) es.addEventListener('llm_call', (e: MessageEvent) => { - const d = JSON.parse(e.data) as { label: string; system_chars: number; user_chars: number } + const d = JSON.parse(e.data) as { label: string; analysis: string | null; system_chars: number; user_chars: number } push('debug', d.label, `llm ← sys:${d.system_chars} usr:${d.user_chars}`) }) @@ -241,8 +313,6 @@ export function useRunStream() { }) es.addEventListener('error', (e: MessageEvent) => { - // The browser also fires a generic `error` event on disconnect; only - // surface it when the run hasn't already ended. if (status.value === 'streaming') { status.value = 'error' const msg = e?.data ? JSON.parse(e.data)?.message : 'stream error' @@ -261,7 +331,7 @@ export function useRunStream() { return { status, runId, question, rationale, planSteps, graphNodes, - log, toolCalls, answer, error, isRunning, + log, toolCalls, subEventsByAnalysis, answer, error, isRunning, ask, disconnect, } } diff --git a/ui/app/src/pages/Ask.vue b/ui/app/src/pages/Ask.vue index a2d149b..b7de85a 100644 --- a/ui/app/src/pages/Ask.vue +++ b/ui/app/src/pages/Ask.vue @@ -1,19 +1,19 @@