145 lines
5.0 KiB
Python
145 lines
5.0 KiB
Python
"""pick_for_question — one constrained LLM call that returns a typed Pick.
|
|
|
|
Used by L2 Analyses that need to translate a free-form question into the
|
|
composer's input. The model is shown:
|
|
- the narrowed metric catalog (name + description + unit)
|
|
- the narrowed candidate column refs (bare or `table.column`)
|
|
- the filter grammar (via the system prompt)
|
|
|
|
It returns JSON. We validate the JSON into a `Pick` and fail fast on any
|
|
shape error — no local retry. Per the project's retry feedback, recovery
|
|
belongs at a future global layer that emits a visible event when it fires.
|
|
|
|
If the model returns `{"error": "..."}` (the prompt's escape hatch for
|
|
"this question can't be expressed"), we raise PickValidationError with
|
|
the model's reason — the caller surfaces it cleanly.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import re
|
|
from dataclasses import dataclass
|
|
|
|
from api import langfuse_client as lf
|
|
from api.analyses._narrow import Candidates, narrow_candidates
|
|
from api.composer.types import Pick, PickValidationError
|
|
from api.llm import chat
|
|
from api.prompts import load, render
|
|
from api.recon import load_recon
|
|
from api.recon.types import Recon
|
|
|
|
|
|
@dataclass
|
|
class PickOutcome:
|
|
pick: Pick
|
|
candidates: Candidates # the narrowed set the model saw (kept for tracing)
|
|
|
|
|
|
def pick_for_question(
|
|
question: str,
|
|
*,
|
|
recon: Recon | None = None,
|
|
allowed_metrics: list[str] | None = None,
|
|
max_tokens: int = 512,
|
|
span_name: str = "pick.gen",
|
|
) -> PickOutcome:
|
|
"""Run one LLM call to translate `question` into a Pick.
|
|
|
|
`allowed_metrics`: optional list to restrict the candidate metrics
|
|
(used by Analyses that already know which metric the planner chose,
|
|
so we don't waste tokens listing every metric in the catalog).
|
|
"""
|
|
recon = recon or load_recon()
|
|
candidates = narrow_candidates(recon, allowed_metrics=allowed_metrics)
|
|
|
|
metrics_block = _render_metrics_block(candidates)
|
|
columns_block = _render_columns_block(candidates)
|
|
|
|
system = load("pick.system")
|
|
user = render(
|
|
"pick.user",
|
|
question=question,
|
|
metrics_block=metrics_block,
|
|
columns_block=columns_block,
|
|
)
|
|
|
|
with lf.span(
|
|
"pick_for_question",
|
|
input={
|
|
"question": question,
|
|
"metric_candidates": [m.name for m in candidates.metrics],
|
|
"column_candidate_count": len(candidates.column_refs),
|
|
},
|
|
) as span:
|
|
raw = chat(system=system, user=user, max_tokens=max_tokens, span_name=span_name)
|
|
payload = _parse_json(raw)
|
|
if "error" in payload:
|
|
raise PickValidationError(
|
|
f"LLM declined to pick: {payload['error']!s}"
|
|
)
|
|
pick = Pick.from_dict(payload)
|
|
# Bind-check the pick against recon up-front — the composer would
|
|
# raise the same errors later, but raising here makes the failure
|
|
# event happen at the pick step instead of compose.
|
|
_validate_against_recon(pick, recon)
|
|
span.update(output=pick.to_dict())
|
|
return PickOutcome(pick=pick, candidates=candidates)
|
|
|
|
|
|
def _render_metrics_block(candidates: Candidates) -> str:
|
|
if not candidates.metrics:
|
|
return "(no candidate metrics)"
|
|
lines: list[str] = []
|
|
for m in candidates.metrics:
|
|
unit = f" [{m.unit}]" if m.unit else ""
|
|
lines.append(f"- {m.name}{unit}: {m.description}")
|
|
return "\n".join(lines)
|
|
|
|
|
|
def _render_columns_block(candidates: Candidates) -> str:
|
|
if not candidates.column_refs:
|
|
return "(no candidate columns)"
|
|
return "\n".join(f"- {c}" for c in candidates.column_refs)
|
|
|
|
|
|
def _parse_json(text: str) -> dict:
|
|
"""Best-effort: prefer a ```json``` fence, otherwise extract the first
|
|
{...} block. Matches the existing pattern in drill_down._parse_json."""
|
|
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 PickValidationError(
|
|
f"pick_for_question returned no JSON object: {text[:200]!r}"
|
|
)
|
|
try:
|
|
return json.loads(raw[start:end + 1])
|
|
except json.JSONDecodeError as e:
|
|
raise PickValidationError(
|
|
f"pick_for_question returned invalid JSON: {e!s}; raw={raw[:200]!r}"
|
|
) from e
|
|
|
|
|
|
def _validate_against_recon(pick: Pick, recon: Recon) -> None:
|
|
"""Surface metric/column resolution errors as PickValidationError up
|
|
front so the trace marks the pick step (not compose) as the failure."""
|
|
if pick.metric not in recon.metrics:
|
|
raise PickValidationError(
|
|
f"picked metric {pick.metric!r} is not in recon"
|
|
)
|
|
for ref in pick.group_by:
|
|
try:
|
|
recon.resolve_column(ref)
|
|
except ValueError as e:
|
|
raise PickValidationError(
|
|
f"group_by column {ref!r}: {e}"
|
|
) from e
|
|
for f in pick.where:
|
|
try:
|
|
recon.resolve_column(f.column)
|
|
except ValueError as e:
|
|
raise PickValidationError(
|
|
f"where filter on {f.column!r}: {e}"
|
|
) from e
|