nl2sql/planner.py

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


class Planner:
    name = "planner"

    def __init__(self, llm: LLMProvider) -> None:
        self.llm = llm

    def run(self, *, user_query: str, schema_preview: str) -> StageResult:
        t0 = time.perf_counter()
        plan_text, t_in, t_out, cost = self.llm.plan(
            user_query=user_query, 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={"len_plan": len(plan_text)},
        )
        return StageResult(ok=True, data={"plan": plan_text}, trace=trace)