nl2sql/repair.py

import time


from nl2sql.types import StageTrace, StageResult
from adapters.llm.base import LLMProvider

GUIDELINES = """
When repairing:
1. Keep query SELECT-only.
2. Explicitly qualify ambiguous columns with table names.
3. Match GROUP BY fields with aggregations.
4. Use known foreign keys for JOIN.
5. Add a reasonable LIMIT if missing.
Return only the corrected SQL.
"""


class Repair:
    name = "repair"

    def __init__(self, llm: LLMProvider):
        self.llm = llm

    def run(self, sql: str, error_msg: str, schema_preview: str) -> StageResult:
        t0 = time.perf_counter()
        fixed_sql, t_in, t_out, cost = self.llm.repair(
            sql=sql,
            error_msg=f"{GUIDELINES}\n\n{error_msg}",
            schema_preview=schema_preview,
        )
        trace = StageTrace(
            stage=self.name,
            duration_ms=(time.perf_counter() - t0) * 1000,
            token_in=t_in,
            token_out=t_out,
            cost_usd=cost,
            notes={"old_sql_len": len(sql), "new_sql_len": len(fixed_sql)},
        )
        return StageResult(ok=True, data={"sql": fixed_sql}, trace=trace)