verbose live UI + tool-level SSE events + Groq default + regression tests

api/plan/planner.py

@@ -0,0 +1,52 @@
+"""Plan layer — one LLM call that selects Analyses for a question.
+
+Public surface is framework-free. Returns a `Plan` that the runtime iterates.
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+import re
+from typing import Any
+
+from api import langfuse_client as lf
+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
+
+logger = logging.getLogger("nvi.plan.planner")
+
+
+def plan(question: str) -> Plan:
+    schema = load_schema()
+    user = render(
+        "planner.user",
+        question=question,
+        table_names=", ".join(schema.table_names()),
+        metrics_block=schema.render_metrics(),
+        catalog=json.dumps(catalog(), indent=2),
+    )
+
+    with lf.span("plan", as_type="generation", input={"question": question}) as span:
+        result = Plan.from_dict(
+            _parse(chat(system=load("planner.system"), user=user, max_tokens=1024))
+        )
+        span.update(output={
+            "rationale": result.rationale,
+            "steps": [s.to_dict() for s in result.steps],
+        })
+        if not result.steps:
+            raise ValueError("planner returned empty plan")
+        return result
+
+
+def _parse(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"planner returned no JSON: {text[:200]!r}")
+    return json.loads(raw[start:end + 1])

api/plan/types.py

@@ -0,0 +1,37 @@
+"""Public data shapes for the Plan layer."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any
+
+
+@dataclass
+class PlanStep:
+    analysis: str
+    args: dict[str, Any] = field(default_factory=dict)
+    why: str = ""
+
+    @classmethod
+    def from_dict(cls, raw: dict[str, Any]) -> "PlanStep":
+        return cls(
+            analysis=raw["analysis"],
+            args=raw.get("args") or {},
+            why=raw.get("why", ""),
+        )
+
+    def to_dict(self) -> dict[str, Any]:
+        return {"analysis": self.analysis, "args": self.args, "why": self.why}
+
+
+@dataclass
+class Plan:
+    rationale: str
+    steps: list[PlanStep]
+
+    @classmethod
+    def from_dict(cls, raw: dict[str, Any]) -> "Plan":
+        return cls(
+            rationale=raw.get("rationale", ""),
+            steps=[PlanStep.from_dict(s) for s in raw.get("steps", [])],
+        )