nl2sql/verifier.py

import sqlglot
from sqlglot import expressions as exp
from nl2sql.types import StageResult, StageTrace

class Verifier:
    name = "verifier"

    def run(self, sql: str, exec_result: StageResult) -> StageResult:
        if not exec_result.ok:
            return StageResult(ok=False, data=None,
                               trace=StageTrace(stage=self.name, duration_ms=0,
                               notes={"reason": "execution_error"}),
                               error=exec_result.errors)

        # Rule 1: check SELECT / GROUP consistency
        issues = []
        try:
            tree = sqlglot.parse_one(sql)
            if isinstance(tree, exp.Select):
                group = tree.args.get("group")
                aggs = [a for a in tree.find_all(exp.AggFunc)]
                if aggs and not group:
                    issues.append("Aggregation without GROUP BY.")
        except Exception as e:
            issues.append(f"Parse error during verification: {e}")

        if issues:
            return StageResult(ok=False, data=None,
                               trace=StageTrace(stage=self.name, duration_ms=0,
                               notes={"issues": issues}),
                               error=issues)
        return StageResult(ok=True, data={"verified": True},
                           trace=StageTrace(stage=self.name, duration_ms=0))