improve operator handling and deep link cohort building

Author: maggie-li-yd

plugins/functions/code-suggestion/src/agent/cohortAgent.ts

@@ -30,6 +30,7 @@ import { getCohortAgentPrompt } from "./cohortAgentPrompt";
 // degrades tool choice and adds latency to every step.
 const COHORT_AGENT_SERVER_TOOLS = new Set([
   "list_cohort_filters",
+  "list_cohort_filter_values",
   "build_d2e_cohort_deeplink",
   "search_concepts",
   "check_concept_coverage_in_dataset",

plugins/functions/code-suggestion/src/agent/cohortAgentPrompt.ts

@@ -143,19 +143,123 @@ export function buildCohortCatalog(config: any): CohortCatalog {
   return { cards };
 }
 
+/** Case/whitespace-insensitive comparison key for a card/attribute name. */
+export function normName(s: string): string {
+  return String(s ?? "").trim().toLowerCase();
+}
+
+/** Find a card by display name (exact ci, then substring). */
+export function findCardByName(
+  catalog: CohortCatalog,
+  name: string,
+): CatalogCard | undefined {
+  const n = normName(name);
+  return (
+    catalog.cards.find((c) => normName(c.name) === n) ??
+    catalog.cards.find(
+      (c) => normName(c.name).includes(n) || n.includes(normName(c.name)),
+    )
+  );
+}
+
+/** Find an attribute within a card by display name (exact ci, then substring). */
+export function findAttributeByName(
+  card: CatalogCard,
+  name: string,
+): CatalogAttribute | undefined {
+  const n = normName(name);
+  return (
+    card.attributes.find((a) => normName(a.name) === n) ??
+    card.attributes.find(
+      (a) => normName(a.name).includes(n) || n.includes(normName(a.name)),
+    )
+  );
+}
+
+/**
+ * Every card that exposes an attribute by this name.
+ *
+ * Used to turn "card X has no attribute Y" into "…but card Z does". The model
+ * puts demographics on the event card it is thinking about ("ER visits under
+ * 80" → Age on the Visit card), and without this hint it reads the rejection as
+ * "this dataset cannot filter on age" and gives up — the attribute is on Basic
+ * Data, one clause away.
+ */
+export function findAttributeAcrossCards(
+  catalog: CohortCatalog,
+  name: string,
+): { card: CatalogCard; attribute: CatalogAttribute }[] {
+  const out: { card: CatalogCard; attribute: CatalogAttribute }[] = [];
+  for (const card of catalog.cards) {
+    const attribute = findAttributeByName(card, name);
+    if (attribute) out.push({ card, attribute });
+  }
+  return out;
+}
+
+/** The card's primary concept-set attribute (cohortDefinitionKey "CodesetId"). */
+export function primaryConceptAttribute(
+  card: CatalogCard,
+): CatalogAttribute | undefined {
+  return card.attributes.find(
+    (a) => a.kind === "conceptSet" && a.cohortDefinitionKey === "CodesetId",
+  );
+}
+
 /**
  * Compact human/LLM-readable summary of the catalog: one line per card listing
  * each attribute as `Name[kind]`, so the model can ground its filter choices on
  * the real cards/attributes before calling the build tool.
+ *
+ * The trailing notes are not decoration. Two failure modes come straight from a
+ * bare card→attribute listing: the model looks for Age on the event card it is
+ * building and concludes the dataset cannot filter by age (demographics are
+ * patient-card attributes), and it invents a `category` token instead of looking
+ * the column up. Both are answered here, next to the data they apply to.
  */
 export function summarizeCatalog(catalog: CohortCatalog): string {
   const lines = catalog.cards.map((c) => {
-    const attrs = c.attributes.map((a) => `${a.name}[${a.kind}]`).join(", ");
+    const attrs = c.attributes
+      .map(
+        (a) =>
+          `${a.name}[${a.kind}${
+            a.kind === "conceptSet" && a.cohortDefinitionKey === "CodesetId"
+              ? "*"
+              : ""
+          }]`,
+      )
+      .join(", ");
     return `- ${c.name}: ${attrs || "(no filterable attributes)"}`;
   });
-  return (
-    "Filter cards available on this dataset (use only these cards and " +
-    "attributes when composing cohort filters):\n" +
-    lines.join("\n")
-  );
+  const patient = catalog.cards.find((c) => c.key === "patient");
+  const demographics = (patient?.attributes ?? [])
+    .filter((a) => a.kind !== "conceptSet")
+    .map((a) => a.name)
+    .join(", ");
+
+  return [
+    "Filter cards available on this dataset. Use these exact card and attribute " +
+    "NAMES in build_d2e_cohort_deeplink clauses (one clause per card occurrence, " +
+    "clauses are AND-ed):",
+    lines.join("\n"),
+    "",
+    "Attribute kinds — how to supply the value:",
+    "- num       → { attribute, op: '>='|'<='|'<'|'>'|'='|'!=', value: <number> }, " +
+    "or op 'range' with value [low, high]",
+    "- category  → the EXACT stored token. Look it up with " +
+    "list_cohort_filter_values({ card, attribute, query? }) and pass a value it " +
+    "returned; never hardcode or guess one. Use op 'in' with a list to match any " +
+    "of several tokens.",
+    "- conceptSet→ a persisted concept-set id (from list_concept_sets / " +
+    "create_concept_set), NOT an OMOP concept id. '*' marks the card's primary " +
+    "concept set, which is what a clause's `conceptSetId` attaches to.",
+    "- datetime  → not supported by the deep-link builder yet.",
+    "",
+    `"${patient?.name ?? "Basic Data"}" is the PATIENT card: ${
+      demographics || "the patient demographics"
+    } live only there. ` +
+    "A demographic filter always goes in its own clause for that card — event " +
+    "cards (encounters, conditions, drugs, …) have no age or gender attribute, " +
+    "which is how the data is modelled, not a limit on what you can filter.",
+  ].join("\n");
 }

plugins/functions/mcp-server/src/lib/cohortCatalog.ts

@@ -4,6 +4,12 @@ import type {
   CatalogCard,
   CatalogAttribute,
 } from "./cohortCatalog";
+import {
+  findAttributeAcrossCards,
+  findAttributeByName,
+  findCardByName,
+  primaryConceptAttribute,
+} from "./cohortCatalog";
 import type { CohortConstraint, CohortExpression } from "./cohortBookmarkTree";
 
 /**
@@ -38,42 +44,37 @@ export interface ResolverDeps {
 
 const NUM_OPS = new Set([">=", "<=", "<", ">", "=", "!="]);
 
-function norm(s: string): string {
-  return s.trim().toLowerCase();
-}
+/** Ops a category/text attribute accepts. `in` OR-s a list of stored tokens. */
+const CATEGORY_OPS = new Set(["=", "!=", "in", "not in"]);
 
-/** Find a card by display name (exact ci, then substring). */
-function findCard(
+/**
+ * "Card X has no attribute Y" — plus, when Y exists elsewhere, where to put it.
+ *
+ * The bare version of this error is what ended a real session with "the Visit
+ * card doesn't have an age attribute" and no cohort: age is a patient
+ * attribute, one clause away, and the model had no way to know that from the
+ * rejection alone.
+ */
+function unknownAttributeError(
   catalog: CohortCatalog,
-  name: string,
-): CatalogCard | undefined {
-  const n = norm(name);
-  return (
-    catalog.cards.find((c) => norm(c.name) === n) ??
-    catalog.cards.find(
-      (c) => norm(c.name).includes(n) || n.includes(norm(c.name)),
-    )
-  );
-}
-
-/** Find an attribute within a card by display name (exact ci, then substring). */
-function findAttr(
   card: CatalogCard,
-  name: string,
-): CatalogAttribute | undefined {
-  const n = norm(name);
-  return (
-    card.attributes.find((a) => norm(a.name) === n) ??
-    card.attributes.find(
-      (a) => norm(a.name).includes(n) || n.includes(norm(a.name)),
-    )
+  requested: string,
+): Error {
+  const available = card.attributes.map((a) => a.name).join(", ") || "(none)";
+  const elsewhere = findAttributeAcrossCards(catalog, requested).filter(
+    (hit) => hit.card.key !== card.key,
   );
-}
-
-/** The card's primary concept-set attribute (cohortDefinitionKey "CodesetId"). */
-function primaryConceptAttr(card: CatalogCard): CatalogAttribute | undefined {
-  return card.attributes.find(
-    (a) => a.kind === "conceptSet" && a.cohortDefinitionKey === "CodesetId",
+  const hint = elsewhere.length
+    ? ` "${elsewhere[0].attribute.name}" IS available on card "${elsewhere[0].card.name}"` +
+      `${
+        elsewhere.length > 1
+          ? ` (also: ${elsewhere.slice(1).map((h) => `"${h.card.name}"`).join(", ")})`
+          : ""
+      } — move that constraint into its own clause for that card instead of ` +
+      `dropping it. Demographics (age, gender, race) always live on the patient card.`
+    : "";
+  return new Error(
+    `Card "${card.name}" has no attribute "${requested}". Available: ${available}.${hint}`,
   );
 }
 
@@ -150,6 +151,56 @@ function numExpressions(c: ClauseConstraint): {
   return { expressions: [{ operator: c.op, value: v }], combine: "OR" };
 }
 
+/**
+ * Build category/text expressions, resolving each raw term to the token the
+ * dataset actually stores.
+ *
+ * A list (`op:"in"`, or an array value) becomes several OR-ed expressions on the
+ * one attribute. That matters for real questions: an encounter-type column
+ * splits "an ER visit" across several tokens ("Emergency Room Visit",
+ * "Emergency Room and Inpatient Visit"), and forcing one token per constraint
+ * would quietly answer a narrower question than the user asked. Negation is
+ * AND-ed instead — "neither A nor B" is not "not A or not B".
+ */
+async function categoryExpressions(
+  c: ClauseConstraint,
+  card: CatalogCard,
+  attr: CatalogAttribute,
+  deps: ResolverDeps,
+): Promise<{ expressions: CohortExpression[]; combine: "AND" | "OR" }> {
+  const op = String(c.op ?? "=").trim().toLowerCase();
+  if (!CATEGORY_OPS.has(op)) {
+    throw new Error(
+      `Unsupported operator "${c.op}" for text attribute "${attr.name}". Use ` +
+        `"=" (or "in" with a list of values to match any of them), or "!=" / ` +
+        `"not in" to exclude.`,
+    );
+  }
+  const raws = (Array.isArray(c.value) ? c.value : [c.value])
+    .map((v) => String(v ?? "").trim())
+    .filter(Boolean);
+  if (raws.length === 0) {
+    throw new Error(`Constraint on "${attr.name}" has no value.`);
+  }
+
+  const negate = op === "!=" || op === "not in";
+  const resolved: string[] = [];
+  // Sequential on purpose: the resolver's fetches share a per-attribute cache,
+  // so the second token reuses the first one's domain read instead of racing it.
+  for (const raw of raws) {
+    const value = await deps.resolveValue(card, attr, raw);
+    if (!resolved.includes(value)) resolved.push(value);
+  }
+
+  return {
+    expressions: resolved.map((value) => ({
+      operator: negate ? "!=" : "=",
+      value,
+    })),
+    combine: negate ? "AND" : "OR",
+  };
+}
+
 /**
  * Resolve clauses to constraints. Throws an actionable Error on any clause that
  * can't be resolved (unknown card/attribute, missing concept attribute,
@@ -164,7 +215,7 @@ export async function resolveClausesToConstraints(
 
   for (let i = 0; i < clauses.length; i++) {
     const clause = clauses[i];
-    const card = findCard(catalog, clause.card);
+    const card = findCardByName(catalog, clause.card);
     if (!card) {
       throw new Error(
         `Unknown filter card "${clause.card}". Available: ${catalog.cards
@@ -196,7 +247,7 @@ export async function resolveClausesToConstraints(
     //    the agent already resolved the id via the concept-set tools).
     if (hasConcept) {
       await assertValidConceptSetId(clause.conceptSetId, card.name, deps);
-      const attr = primaryConceptAttr(card);
+      const attr = primaryConceptAttribute(card);
       if (!attr) {
         throw new Error(
           `Card "${card.name}" has no concept set to attach concept set ${clause.conceptSetId} to.`,
@@ -212,25 +263,22 @@ export async function resolveClausesToConstraints(
 
     // 2. explicit attribute constraints.
     for (const cc of clause.constraints ?? []) {
-      const attr = findAttr(card, cc.attribute);
+      const attr = findAttributeByName(card, cc.attribute);
       if (!attr) {
-        throw new Error(
-          `Card "${card.name}" has no attribute "${cc.attribute}". Available: ${card.attributes
-            .map((a) => a.name)
-            .join(", ")}.`,
-        );
+        throw unknownAttributeError(catalog, card, cc.attribute);
       }
 
       let expressions: CohortExpression[];
       let combine: "AND" | "OR";
       if (attr.kind === "num") {
         ({ expressions, combine } = numExpressions(cc));
       } else if (attr.kind === "category") {
-        const resolved = await deps.resolveValue(card, attr, String(cc.value));
-        expressions = [
-          { operator: cc.op === "!=" ? "!=" : "=", value: resolved },
-        ];
-        combine = "OR";
+        ({ expressions, combine } = await categoryExpressions(
+          cc,
+          card,
+          attr,
+          deps,
+        ));
       } else if (attr.kind === "conceptSet") {
         // value IS the concept-set id (agent-resolved), e.g. a unit set.
         await assertValidConceptSetId(cc.value, card.name, deps);