feat: add example__* tasks

These are bare bone prompts, which can run on sample shipped data.
They are meant to be used when following some docs on how to step
through the NLP workflow.

Author: mikix

cumulus_etl/etl/pipeline.py

@@ -174,6 +174,9 @@ async def run_pipeline(
     # record filesystem options like --s3-region before creating Roots
     store.set_user_fs_options(vars(args))
 
+    if args.dir_input == "%EXAMPLE%" and not os.path.exists(args.dir_input):
+        args.dir_input = os.path.join(os.path.dirname(__file__), "studies/example/ndjson")
+
     root_input = store.Root(args.dir_input)
     root_output = store.Root(args.dir_output)
     root_phi = store.Root(args.dir_phi, create=True)

cumulus_etl/etl/studies/example/__init__.py

@@ -0,0 +1,7 @@
+"""The example study"""
+
+from .example_tasks import ExampleGpt4oTask as ExampleGpt4oTask
+from .example_tasks import ExampleGpt4Task as ExampleGpt4Task
+from .example_tasks import ExampleGpt5Task as ExampleGpt5Task
+from .example_tasks import ExampleGptOss120bTask as ExampleGptOss120bTask
+from .example_tasks import ExampleLlama4ScoutTask as ExampleLlama4ScoutTask

cumulus_etl/etl/studies/example/example_tasks.py

@@ -0,0 +1,65 @@
+"""Define tasks for the example/sample study"""
+
+import json
+
+import pydantic
+
+from cumulus_etl import nlp
+from cumulus_etl.etl import tasks
+
+
+class AgeMention(pydantic.BaseModel):
+    has_mention: bool | None = pydantic.Field(None)
+    spans: list[str] = pydantic.Field(default_factory=list, description="Supporting text spans")
+    age: int | None = pydantic.Field(None, description="The age of the patient")
+
+
+class BaseExampleTask(tasks.BaseOpenAiTaskWithSpans):
+    task_version = 0
+    # Task Version History:
+    # ** 0 (2025-08): Initial work, still in flux **
+
+    system_prompt = (
+        "You are a clinical chart reviewer.\n"
+        "Your task is to extract patient-specific information from an unstructured clinical "
+        "document and map it into a predefined Pydantic schema.\n"
+        "\n"
+        "Core Rules:\n"
+        "1. Base all assertions ONLY on patient-specific information in the clinical document.\n"
+        "   - Never negate or exclude information just because it is not mentioned.\n"
+        "   - Never conflate family history or population-level risk with patient findings.\n"
+        "2. Do not invent or infer facts beyond what is documented.\n"
+        "3. Maintain high fidelity to the clinical document language when citing spans.\n"
+        "4. Always produce structured JSON that conforms to the Pydantic schema provided below.\n"
+        "\n"
+        "Pydantic Schema:\n" + json.dumps(AgeMention.model_json_schema())
+    )
+    response_format = AgeMention
+
+
+# Have a task for every ETL-supported model, to allow sites to choose whatever model works for them.
+
+
+class ExampleGpt4Task(BaseExampleTask):
+    name = "example__nlp_gpt4"
+    client_class = nlp.Gpt4Model
+
+
+class ExampleGpt4oTask(BaseExampleTask):
+    name = "example__nlp_gpt4o"
+    client_class = nlp.Gpt4oModel
+
+
+class ExampleGpt5Task(BaseExampleTask):
+    name = "example__nlp_gpt5"
+    client_class = nlp.Gpt5Model
+
+
+class ExampleGptOss120bTask(BaseExampleTask):
+    name = "example__nlp_gpt_oss_120b"
+    client_class = nlp.GptOss120bModel
+
+
+class ExampleLlama4ScoutTask(BaseExampleTask):
+    name = "example__nlp_llama4_scout"
+    client_class = nlp.Llama4ScoutModel

cumulus_etl/etl/studies/example/ndjson/docref.ndjson.gz

@@ -6,6 +6,7 @@
 import re
 import string
 import sys
+import types
 import typing
 from collections.abc import AsyncIterator, Callable
 from typing import ClassVar
@@ -223,7 +224,7 @@ def convert_pydantic_fields_to_pyarrow(
     ) -> pyarrow.DataType:
         return pyarrow.struct(
             [
-                pyarrow.field(name, cls._convert_type_to_pyarrow(info.annotation))
+                pyarrow.field(name, cls._convert_type_to_pyarrow(info.annotation), nullable=True)
                 for name, info in fields.items()
             ]
         )
@@ -232,16 +233,23 @@ def convert_pydantic_fields_to_pyarrow(
     def _convert_type_to_pyarrow(cls, annotation) -> pyarrow.DataType:
         # Since we only need to handle a small amount of possible types, we just do this ourselves
         # rather than relying on an external library.
-        if issubclass(annotation, str):
+        if origin := typing.get_origin(annotation):  # e.g. "UnionType" or "list"
+            sub_type = typing.get_args(annotation)[0]
+            if issubclass(origin, types.UnionType):
+                # This is gonna be something like "str | None" so just grab first arg.
+                # We mark all our fields are nullable at the pyarrow layer.
+                return cls._convert_type_to_pyarrow(sub_type)
+            elif issubclass(origin, list):
+                # Note: does not handle struct types underneath yet
+                return pyarrow.list_(cls._convert_type_to_pyarrow(sub_type))
+        elif issubclass(annotation, str):
             return pyarrow.string()
         elif issubclass(annotation, bool):
             return pyarrow.bool_()
-        elif issubclass(typing.get_origin(annotation), list):
-            sub_type = typing.get_args(annotation)[0]
-            # Note: does not handle struct types underneath yet
-            return pyarrow.list_(cls._convert_type_to_pyarrow(sub_type))
-        else:
-            raise ValueError(f"Unsupported type {annotation}")  # pragma: no cover
+        elif issubclass(annotation, int):
+            return pyarrow.int32()
+
+        raise ValueError(f"Unsupported type {annotation}")  # pragma: no cover
 
 
 class BaseOpenAiTaskWithSpans(BaseOpenAiTask):

cumulus_etl/etl/studies/example/ndjson/dxreport.ndjson.gz

@@ -5,7 +5,7 @@
 from typing import TypeVar
 
 from cumulus_etl import cli_utils, errors
-from cumulus_etl.etl.studies import covid_symptom, irae
+from cumulus_etl.etl.studies import covid_symptom, example, irae
 from cumulus_etl.etl.tasks import basic_tasks
 
 AnyTask = TypeVar("AnyTask", bound="EtlTask")  # noqa: F821
@@ -31,6 +31,11 @@ def get_nlp_tasks() -> list[type[AnyTask]]:
         covid_symptom.CovidSymptomNlpResultsGpt4Task,
         covid_symptom.CovidSymptomNlpResultsTask,
         covid_symptom.CovidSymptomNlpResultsTermExistsTask,
+        example.ExampleGpt4Task,
+        example.ExampleGpt4oTask,
+        example.ExampleGpt5Task,
+        example.ExampleGptOss120bTask,
+        example.ExampleLlama4ScoutTask,
         irae.IraeGptOss120bTask,
         irae.IraeGpt4oTask,
         irae.IraeGpt5Task,

cumulus_etl/etl/tasks/nlp_task.py

@@ -31,9 +31,9 @@ async def post_init_check(self) -> None:
         try:
             models = self.client.models.list()
             names = {model.id async for model in models}
-        except openai.APIError:
+        except openai.APIError as exc:
             errors.fatal(
-                f"NLP server '{self.USER_ID}' is unreachable.\n"
+                f"NLP server '{self.USER_ID}' is unreachable: {exc}.\n"
                 f"If it's a local server, try running 'docker compose up {self.USER_ID} --wait'.",
                 errors.SERVICE_MISSING,
             )
@@ -45,6 +45,9 @@ async def post_init_check(self) -> None:
             )
 
     async def prompt(self, system: str, user: str, schema: BaseModel) -> chat.ParsedChatCompletion:
+        return await self._parse_prompt(system, user, schema)
+
+    async def _parse_prompt(self, system: str, user: str, schema) -> chat.ParsedChatCompletion:
         return await self.client.chat.completions.parse(
             model=self.MODEL_NAME,
             messages=[
@@ -75,12 +78,20 @@ async def pre_init_check(cls) -> None:
             errors.fatal("\n".join(messages), errors.ARGS_INVALID)
 
     def make_client(self) -> openai.AsyncOpenAI:
-        return openai.AsyncAzureOpenAI(api_version="2024-06-01")
+        return openai.AsyncAzureOpenAI(api_version="2024-10-21")
 
 
-class Gpt35Model(AzureModel):
+class Gpt35Model(AzureModel):  # deprecated, do not use in new code (doesn't support JSON schemas)
     MODEL_NAME = "gpt-35-turbo-0125"
 
+    # 3.5 doesn't support a pydantic JSON schema, so we do some work to keep it using the same API
+    # as the rest of our code.
+    async def prompt(self, system: str, user: str, schema: BaseModel) -> chat.ParsedChatCompletion:
+        response = await self._parse_prompt(system, user, {"type": "json_object"})
+        parsed = schema.model_validate_json(response.choices[0].message.content)
+        response.choices[0].message.parsed = parsed
+        return response
+
 
 class Gpt4Model(AzureModel):
     MODEL_NAME = "gpt-4"

cumulus_etl/etl/tasks/task_factory.py

@@ -0,0 +1,116 @@
+---
+title: Example Workflow
+parent: NLP
+grand_parent: ETL
+nav_order: 1
+# audience: engineer familiar with the project
+# type: tutorial
+---
+
+# An Example NLP Workflow
+
+Let's work through an end-to-end NLP workflow, as if you were doing a real study.
+But we'll use an example study shipped with the ETL for testing purposes instead.
+
+This will take us from the initial actual NLP, then to chart review,
+then finally to analyzing accuracy.
+
+You don't need to prepare your own clinical notes for this run-through.
+We'll use synthetic notes shipped with Cumulus ETL for this very purpose.
+
+This example study we'll use is just a very simple age range study.
+The NLP will only be tasked with extracting an age from a clinical note.
+
+## The NLP Itself
+
+Before we start, you'll need to have Cumulus ETL and your AWS infrastructure ready.
+Follow the [setup instructions](../setup) if you haven't done so already, then come back here.
+
+### Model Setup
+
+You have a choice of model for this.
+Real studies might require one specific model or another.
+But this example task is fairly liberal.
+
+Here are the options, along with the task name to use.
+
+#### Azure Cloud Options
+For these, you'll want to set a couple variables first:
+```sh
+export AZURE_OPENAI_API_KEY=xxx
+export AZURE_OPENAI_ENDPOINT=https://xxx.openai.azure.com/
+```
+
+Task names:
+- GPT4: `example__nlp_gpt4`
+- GPT4o: `example__nlp_gpt4o`
+- GPT5: `example__nlp_gpt5`
+
+#### Local (On-Prem) Options
+For these, you'll need to start up the appropriate model on your machine:
+```sh
+docker compose up --wait gpt-oss-120b
+```
+
+Task names:
+- GPT-OSS 120B (needs 80GB of GPU memory): `example__nlp_gpt_oss_120b`
+
+### Running the ETL
+
+Now that your model is ready, let's run the ETL on some notes!
+
+Below is the command line to use.
+You'll need to change the bucket names and paths to wherever you set up your AWS infrastructure.
+And you'll want to change the task name as appropriate for your model.
+Leave the odd looking `%EXAMPLE%` bit in place;
+that just tells Cumulus ETL to use its built-in example documents as the input.
+
+The output and PHI bucket locations should be the same as your normal ETL runs on raw FHIR data.
+There's no actual PHI in this example run because of the synthetic data,
+but normally there is, and that PHI bucket is where Cumulus ETL keeps caches of NLP results.
+
+```sh
+docker compose run --rm \
+  cumulus-etl nlp \
+  %EXAMPLE% \
+  s3://my-output-bucket/ \
+  s3://my-phi-bucket/ \
+  --task example__nlp_gpt4
+```
+
+(If this were a real study, you'd probably do this a bit differently.
+You'd point at your real DocumentReference resources for example.
+And you'd probably restrict the set of documents you run NLP on with an argument
+like `--cohort-athena-table study__my_cohort`.
+But for this run-through, we're going to hand-wave all the document selection pieces.)
+
+### Running the Crawler
+
+Whenever you write a new table to S3, you'll want to run your AWS Glue crawler again,
+so that the table's schema gets set correctly in Athena.
+
+First, confirm that your AWS Cloud Formation templates have the `example__nlp_*` tables
+configured in them. If not, try copying the Glue crawler definition from
+[the sample template we provide](../setup/aws.md).
+
+Then go to your AWS console, in the AWS Glue service, in the sidebar under Data Catalog, and
+choose Crawlers.
+You should see your crawler listed there. Select it, click Run, and wait for it to finish.
+
+### Confirm the Data in Athena
+
+While you're in the AWS console, switch to the Athena service and select the appropriate
+Cumulus workgroup and database.
+
+Then if you make a query like below (assuming you used the GPT4 model),
+you should see eight results with extracted ages.
+```sql
+select * from example__nlp_gpt4
+```
+
+**Congratulations!**
+You've now run NLP on some synthetic clinical notes and uploaded the results to Athena.
+Those extracted ages could now be post-processed by the `example` study to calculate age ranges,
+and then confirmed with chart review by humans.
+
+At least that's the flow you'd use for a real study.

cumulus_etl/nlp/openai.py

@@ -2,6 +2,7 @@
 title: NLP
 parent: ETL
 nav_order: 5
+has_children: true
 # audience: non-programmers, conversational tone, selling a bit
 # type: explanation
 ---