Custom Concept Grouping for Cohorts

{
  "include": [{ "concept_id": 4115260, "include_descendants": true }],
  "exclude": [{ "concept_id": 37311178, "include_descendants": true }]
}

PAGE_SIZE = 200   # the endpoint clamps anything larger to 200
MAX_LEVELS = 20   # the endpoint's maximum; it defaults to only 10


def fetch_all(path, key):
    """Read every page of a paginated hierarchy response.

    Three separate caps can silently shrink a hierarchy result. All three have to
    be handled or your concept set is quietly incomplete:

    * `page_size`  - one call returns a single page. Page until has_next is False.
    * `max_levels` - defaults to 10, so descendants deeper than 10 levels are
      dropped. Nothing in the response flags this: `truncated` reports the row
      cap only. Ask for the maximum (20).
    * `max_results` - the row cap. `truncated` does flag this one.
    """
    page, out = 1, []
    while True:
        body = client.get(
            path,
            max_levels=MAX_LEVELS,
            max_results=5000,
            include_invalid=False,
            page=page,
            page_size=PAGE_SIZE,
        )
        out.extend(body["data"][key])       # note: payload is nested under "data"

        summary = body["data"]["hierarchy_summary"]
        if summary.get("truncated"):
            raise RuntimeError(
                f"{path} exceeds {summary['result_limit']} rows and was truncated. "
                "Narrow it with domain_ids / vocabulary_ids, or split it into "
                "several include seeds."
            )

        if not body["meta"]["pagination"]["has_next"]:
            return out
        page += 1


def ancestors(concept_id):
    return fetch_all(f"/v1/concepts/{concept_id}/ancestors", "ancestors")


def descendants(concept_id):
    return fetch_all(f"/v1/concepts/{concept_id}/descendants", "descendants")


def find_groupers(seed_ids):
    """Ancestors shared by ALL seeds, tightest first.

    Each candidate is scored by its distance to the seed it is *furthest* from,
    then by its total distance. Ranking on a single seed's
    `min_levels_of_separation` would make the answer depend on which seed you
    happened to process last - the same candidate sits at different distances
    from different seeds.

    Seeds are de-duplicated first. `distances[candidate]` is keyed by seed, so a
    repeated seed collapses to one key while `len(seeds)` still counts it twice -
    the "shared by every seed" test could then never pass and this would return an
    empty list. Extraction pipelines produce duplicates routinely, whenever two
    phrases resolve to the same concept.
    """
    seeds = sorted(set(seed_ids))
    distances, meta = {}, {}

    for cid in seeds:
        for a in ancestors(cid):
            distances.setdefault(a["concept_id"], {})[cid] = a["min_levels_of_separation"]
            meta[a["concept_id"]] = a

    shared = [c for c, d in distances.items() if len(d) == len(seeds)]

    return sorted(
        (
            {**meta[c], "worst_distance": max(distances[c].values())}
            for c in shared
        ),
        key=lambda a: (
            a["worst_distance"],
            sum(distances[a["concept_id"]].values()),
            a["concept_id"],   # deterministic tie-break: candidates do tie on both
        ),
    )

4115260  Lung finding                      worst_distance=2
4115259  Lower respiratory tract finding   worst_distance=3
4185503  Finding of region of thorax       worst_distance=3
4227253  Viscus structure finding          worst_distance=3
4024567  Respiratory finding               worst_distance=4

expression = {
    "include": [
        {"concept_id": 4115260, "include_descendants": True},   # Lung finding
    ],
    "exclude": [
        {"concept_id": 37311178, "include_descendants": True},  # Normal lung
        {"concept_id": 4064736,  "include_descendants": True},  # Lung function testing normal
        {"concept_id": 40481136, "include_descendants": True},  # Lungs in normal arrangement
        {"concept_id": 4300172,  "include_descendants": True},  # Chest percussion normal
    ],
    "domain_id": "Condition",
    "standard_only": True,
}

def keep(concept, expression):
    """The domain / standard filters, applied to any concept in the set."""
    if expression.get("standard_only") and concept["standard_concept"] != "S":
        return False
    if expression.get("domain_id") and concept["domain_id"] != expression["domain_id"]:
        return False
    return True


def expand(expression):
    def resolve(rules):
        out = set()
        for rule in rules:
            # The rule's own root is subject to the same filters as its
            # descendants. Skipping it here is how a classification ('C') or
            # wrong-domain grouper leaks into a cohort it can never legally
            # appear in.
            root = client.get(f"/v1/concepts/{rule['concept_id']}")["data"]
            if keep(root, expression):
                out.add(root["concept_id"])

            if rule.get("include_descendants"):
                out.update(
                    d["concept_id"]
                    for d in descendants(rule["concept_id"])
                    if keep(d, expression)
                )
        return out

    return resolve(expression["include"]) - resolve(expression.get("exclude", []))

agent proposes terms
      │
      ▼
resolve to seed concept_ids
      │
      ▼
intersect /ancestors  ──────►  candidate groupers (nearest first)
      │
      ▼
author expression  {include, exclude}   ◄── you version this
      │
      ▼
expand  →  filter domain + standard_concept='S'
      │
      ▼
concept_id list  →  WHERE condition_concept_id IN (...)

{
  "queries": [
    { "query_id": "a", "concept_id": 4115260, "operation": "descendants",
      "params": { "max_results": 5000 } }
  ]
}