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, min_coverage=None):
    """Shared ancestors of the seeds, best candidate first.

    Candidates are ranked by how many seeds they cover, then by their distance to
    the seed they are *furthest* from, then by 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.

    `min_coverage` defaults to "every seed", but partial coverage is deliberately
    reported rather than discarded. One outlier seed otherwise drags the only
    surviving candidate all the way up to "Clinical finding"; a 4-of-5 candidate
    with 2,400 descendants is a far better grouper than a 5-of-5 candidate with
    194,000, and you handle the fifth seed as its own include rule.

    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 coverage 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))
    floor = min_coverage or len(seeds)
    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

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


def with_cost(candidates, top_n=8):
    """Annotate the top candidates with what they would drag in.

    `/level` returns the count without fetching the set, so this is one cheap call
    per candidate. Note `total_descendants` comes back as a **string** - compare it
    as a number or the size cliff below will not sort the way you expect.
    """
    for c in candidates[:top_n]:
        level = client.get(f"/v1/concepts/{c['concept_id']}/level")["data"]
        c["total_descendants"] = int(level["total_descendants"])
    return candidates[:top_n]

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", []))

255573    Condition   Chronic obstructive pulmonary disease   Disorder
36308250  Meas Value  Chronic obstructive pulmonary disease   Answer     ← rank 2

GET /v1/search/semantic?query=chronic+obstructive+pulmonary+disease
      &standard_concept=S&domain_ids=Condition

4183159  Pulmonary blastoma (Observation / Morph Abnormality)
   --Asso morph of-->  40391740  Pulmonary blastoma                      (Condition / Disorder)
                       37167738  Primary pulmonary blastoma              (Condition / Disorder)
                       36561198  Pulmonary blastoma of lower lobe, lung  (ICDO Condition)
                       ... 10+ more

agent proposes terms
      │
      ▼
resolve to seed concept_ids   ◄── pass domain_ids; drop isolated concepts
      │
      ▼
intersect /ancestors  ──────►  candidate groupers (coverage, then size)
      │
      ▼
author expression  {include, exclude}   ◄── you version this
      │
      ▼
expand  →  filter domain + standard_concept='S'
      │
      ▼
concept_id list  →  WHERE condition_concept_id IN (...)
      │
      ▼
/concepts/recommended  ──────►  anything outside the expansion?
      │                                   │
      │  no - done                        │  yes - add as include seed
      ▼                                   └──────────► (back to expression)

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