src/backend/planner.py

"""Query planner: one natural-language request -> a list of search criteria.

This is a planner, not a chatbot. It has no memory, no dialogue, and never
answers a clinical question. It decomposes a cohort description into concepts
the existing search surfaces can already serve, and hands each one back as a
chip the user clicks. Nothing is searched automatically.

Two boundaries are load-bearing:

* Cohort logic is not ours. Age, care setting, index date, lookback -- the
  research plan assigns those to SAGE. The planner extracts them so the user
  can see they were understood, and returns them in `constraints`, which the
  UI must render as explicitly NOT searched. Silently dropping them would let
  a reader believe the result set honoured "after 2020".

* Every criterion carries a `quote`: the span of the user's own query it came
  from, verified server-side to be a substring of that query. `concept` may be
  a normalization ("HF" -> "heart failure") so it searches well, but the quote
  keeps the decomposition auditable and kills the failure mode where a model
  invents a criterion nobody asked for. This mirrors the rule already applied
  to phenotype facets: every evidence span must appear in its source.
"""

from __future__ import annotations

from .llm import (KINDS, LlmError, builtin_catalog, builtin_configured,
                  builtin_model, builtin_provider, complete_json,
                  normalize_provider, parse_json_object, stream_json)

# Must stay in sync with the frontend CATS table and the served code
# categories. An unknown category is dropped, not coerced -- a chip that
# opens the wrong search surface is worse than a missing chip.
CATEGORIES = {"phenotype", "diagnosis", "medication", "lab", "procedure"}
CONSTRAINT_KINDS = {"age", "setting", "temporal", "demographic", "other"}

MAX_CRITERIA = 8
MAX_CONSTRAINTS = 8
MAX_QUERY_CHARS = 1000
MAX_FIELD_CHARS = 120

# Short, but not as short as it can be. Measured against the 40 SAGE CIPHER
# cohort specs: cutting the original 2,010-char prompt to 808 chars cost real
# quality -- phenotype routing collapsed from 4/8 to 1/8 on the affected
# queries, because the rule pairing a named condition with both phenotype and
# diagnosis had been trimmed away. Restoring that one line (932 chars) matches
# the long prompt's routing.
#
# Prompt length is not the latency lever it looks like: 808 vs 2,010 chars
# moved the median only ~0.7s. The cost is reasoning tokens, which scale with
# how complex the *request* is, not how long the instructions are. So keep
# this prompt tight for clarity and token cost, but do not trim rules hoping
# to buy speed -- you will pay in routing quality and get almost nothing back.
# Anything added here must earn its place against the 40-spec set.
SYSTEM = """\
Split a cohort request into ENCODE search criteria. Plan only: never answer \
clinical questions or produce codes.

Categories: phenotype (the condition naming the cohort), diagnosis (ICD \
codes for a condition), medication, lab, procedure.

Rules:
- One criterion per distinct clinical concept.
- The condition naming the cohort is a phenotype; also emit it as diagnosis \
when the request wants its ICD codes.
- concept: short searchable term, abbreviations expanded. Not a sentence.
- quote: the exact substring of the request it came from, verbatim.
- Cohort logic (age, setting, dates, sex, site, counts) goes in constraints, \
never criteria.
- Invent nothing; no concepts means empty criteria.

Return only this JSON:
{"criteria":[{"concept":"heart failure","category":"phenotype",\
"quote":"heart failure","rationale":"names the cohort"}],
 "constraints":[{"text":"adults","kind":"age"}]}
kind: age|setting|temporal|demographic|other"""


def available() -> bool:
    """Whether this deployment ships a planner model of its own. A user can
    still bring their own model when this is False."""
    return builtin_configured()


def _clean(value: object, limit: int = MAX_FIELD_CHARS) -> str:
    if not isinstance(value, str):
        return ""
    return " ".join(value.split())[:limit].strip()


def _valid_criteria(raw: object, query_lower: str) -> list[dict]:
    """Keep only well-formed, category-known, query-grounded criteria."""
    out: list[dict] = []
    seen: set[tuple[str, str]] = set()
    if not isinstance(raw, list):
        return out
    for item in raw:
        if not isinstance(item, dict):
            continue
        concept = _clean(item.get("concept"))
        category = _clean(item.get("category"), 32).lower()
        quote = _clean(item.get("quote"))
        if not concept or category not in CATEGORIES:
            continue
        # Grounding: the cited span must really be in the user's query.
        # A model that cannot point at the text it used does not get a chip.
        if not quote or quote.lower() not in query_lower:
            continue
        key = (category, concept.lower())
        if key in seen:
            continue
        seen.add(key)
        out.append({"concept": concept, "category": category, "quote": quote,
                    "rationale": _clean(item.get("rationale"))})
        if len(out) >= MAX_CRITERIA:
            break
    return out


def _valid_constraints(raw: object) -> list[dict]:
    out: list[dict] = []
    seen: set[str] = set()
    if not isinstance(raw, list):
        return out
    for item in raw:
        if not isinstance(item, dict):
            continue
        text = _clean(item.get("text"))
        if not text or text.lower() in seen:
            continue
        kind = _clean(item.get("kind"), 32).lower()
        seen.add(text.lower())
        out.append({"text": text, "kind": kind if kind in CONSTRAINT_KINDS else "other"})
        if len(out) >= MAX_CONSTRAINTS:
            break
    return out


MAX_MODEL_TEXT = 100_000


def _assemble(q: str, raw: dict, label: str) -> dict:
    criteria = _valid_criteria(raw.get("criteria"), q.lower())
    constraints = _valid_constraints(raw.get("constraints"))

    # A well-formed response that grounds nothing is a real outcome, not an
    # error: the UI says so plainly rather than showing an empty panel.
    note = None
    if not criteria:
        note = ("No searchable clinical concept was found in that description. "
                "Try naming a condition, drug, lab, or procedure.")

    return {"query": q, "criteria": criteria, "constraints": constraints,
            "note": note, "model": label}


def plan(query: str, model: object = None, builtin: str | None = None) -> dict:
    """Decompose `query` by calling a model from here. `model` optionally
    overrides the deployment's model with a user-supplied provider; `builtin`
    names which of the deployment's own models to use."""
    q = " ".join((query or "").split())[:MAX_QUERY_CHARS]
    if not q:
        raise LlmError("Empty query")

    provider = normalize_provider(model) if model else builtin_provider(builtin)
    raw = complete_json(SYSTEM, q, provider=provider)
    return _assemble(q, raw, provider["label"])


def plan_streaming(query: str, builtin: str | None = None):
    """Yield ("thinking", delta) while the model reasons, then ("plan", dict).

    Same validation as plan(); only the transport differs, so the built-in
    model can show its reasoning instead of a blank box for ~25 seconds.
    """
    q = " ".join((query or "").split())[:MAX_QUERY_CHARS]
    if not q:
        raise LlmError("Empty query")
    provider = builtin_provider(builtin)
    text = ""
    for kind, payload in stream_json(SYSTEM, q, provider=provider):
        if kind == "thinking":
            yield ("thinking", payload)
        elif kind == "usage":
            yield ("usage", payload)
        else:
            text = payload
    yield ("plan", _assemble(q, parse_json_object(text), provider["label"]))


def plan_from_text(query: str, text: object, label: object = None) -> dict:
    """Validate model output the *caller* obtained, without ever seeing their
    credentials.

    The browser calls its own model directly and posts the reply here. Parsing
    and every schema rule -- grounding, the category allowlist, the caps --
    then run in exactly one place, instead of being reimplemented in JS and
    drifting from this file.
    """
    q = " ".join((query or "").split())[:MAX_QUERY_CHARS]
    if not q:
        raise LlmError("Empty query")
    if not isinstance(text, str) or not text.strip():
        raise LlmError("The model returned an empty response")
    if len(text) > MAX_MODEL_TEXT:
        raise LlmError("The model response was too large to read")
    shown = _clean(label, 60) if isinstance(label, str) else ""
    return _assemble(q, parse_json_object(text), shown or "your model")


__all__ = ["plan", "plan_streaming", "plan_from_text", "available", "builtin_model", "SYSTEM",
           "LlmError", "CATEGORIES", "KINDS"]