53 lines
1.6 KiB
Python
53 lines
1.6 KiB
Python
"""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])
|