"""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"]