recon-driven sql composer; pick → compose → execute; llm out of structural sql

api/analyses/_pick.py

@@ -0,0 +1,144 @@
+"""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