Merge pull request #436 from smart-on-fhir/mikix/irae-prompt2

feat(irae): update the IRAE system & user prompts

Author: mikix

cumulus_etl/etl/studies/irae/irae_tasks.py

@@ -1,6 +1,7 @@
 """Define tasks for the irae study"""
 
 import enum
+import json
 from typing import ClassVar
 
 import pydantic
@@ -42,23 +43,31 @@ class DSAMention(pydantic.BaseModel):
 
 
 class BaseIraeTask(tasks.BaseOpenAiTaskWithSpans):
-    task_version = 0
+    task_version = 1
     # Task Version History:
-    # ** 0 (2025-08): Initial work, still in flux **
+    # ** 1 (2025-08): Updated prompt **
+    # ** 0 (2025-08): Initial version **
 
     system_prompt = (
-        "You are a helpful assistant reviewing kidney-transplant notes for donor-specific "
-        "antibody (DSA) information. Return *only* valid JSON."
+        "You are a clinical chart reviewer for a kidney transplant outcomes study.\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(DSAMention.model_json_schema())
     )
     user_prompt = (
-        "Evaluate the following chart for donor-specific antibody (DSA) information.\n"
-        "Here is the chart for you to analyze:\n"
-        "%CLINICAL-NOTE%\n"
-        "Keep the pydantic structure previously provided in mind when structuring your output.\n"
-        "Finally, ensure that your final assertions are based on Patient-specific information "
-        "only.\n"
-        "For example, we should never deny the presence of an observation because of a lack of "
-        "family history."
+        "Evaluate the following clinical document for kidney transplant variables and outcomes.\n"
+        "Here is the clinical document for you to analyze:\n"
+        "\n"
+        "%CLINICAL-NOTE%"
     )
     response_format = DSAMention
 

tests/data/irae/output.ndjson

@@ -1 +1 @@
-{"note_ref": "DocumentReference/c31a3dbf188ed241b2c06b2475cd56159017fa1df1ea882d3fc4beab860fc24d", "subject_ref": "Patient/00680c7c0e2e1712e9c4a01eb5c6dfb8949871faef6337c5db204d19e1d9ca58", "encounter_ref": "Encounter/b3d0707624491d8b71a808bd20b63625981af48f526b95214146de2a15f7dd43", "generated_on": "2021-09-14T21:23:45+00:00", "task_version": 0, "system_fingerprint": "test-fp", "result": {"dsa_history": false, "dsa_mentioned": true, "dsa_present": "DSA Treatment prescribed or DSA treatment administered", "spans": [[5, 9]]}}
+{"note_ref": "DocumentReference/c31a3dbf188ed241b2c06b2475cd56159017fa1df1ea882d3fc4beab860fc24d", "subject_ref": "Patient/00680c7c0e2e1712e9c4a01eb5c6dfb8949871faef6337c5db204d19e1d9ca58", "encounter_ref": "Encounter/b3d0707624491d8b71a808bd20b63625981af48f526b95214146de2a15f7dd43", "generated_on": "2021-09-14T21:23:45+00:00", "task_version": 1, "system_fingerprint": "test-fp", "result": {"dsa_history": false, "dsa_mentioned": true, "dsa_present": "DSA Treatment prescribed or DSA treatment administered", "spans": [[5, 9]]}}

tests/nlp/test_irae.py

@@ -1,5 +1,7 @@
 """Tests for etl/studies/irae/"""
 
+import json
+
 import ddt
 
 from cumulus_etl.etl.studies.irae.irae_tasks import DSAMention, DSAPresent
@@ -43,21 +45,26 @@ async def test_basic_etl(self, task_name, model_id):
                 "messages": [
                     {
                         "role": "system",
-                        "content": "You are a helpful assistant reviewing kidney-transplant notes "
-                        "for donor-specific antibody (DSA) information. Return *only* valid JSON.",
+                        "content": "You are a clinical chart reviewer for a kidney transplant outcomes study.\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(DSAMention.model_json_schema()),
                     },
                     {
                         "role": "user",
-                        "content": "Evaluate the following chart for donor-specific antibody (DSA) "
-                        "information.\n"
-                        "Here is the chart for you to analyze:\n"
-                        "Test note 1\n"
-                        "Keep the pydantic structure previously provided in mind when structuring "
-                        "your output.\n"
-                        "Finally, ensure that your final assertions are based on Patient-specific "
-                        "information only.\n"
-                        "For example, we should never deny the presence of an observation because "
-                        "of a lack of family history.",
+                        "content": "Evaluate the following clinical document for kidney "
+                        "transplant variables and outcomes.\n"
+                        "Here is the clinical document for you to analyze:\n\n"
+                        "Test note 1",
                     },
                 ],
                 "model": model_id,

tests/nlp/test_openai.py

@@ -90,7 +90,7 @@ async def test_caching(self):
         await irae.IraeGptOss120bTask(self.job_config, self.scrubber).run()
 
         self.assertEqual(self.mock_create.call_count, 1)
-        cache_dir = f"{self.phi_dir}/nlp-cache/irae__nlp_gpt_oss_120b_v0/06ee/"
+        cache_dir = f"{self.phi_dir}/nlp-cache/irae__nlp_gpt_oss_120b_v1/06ee"
         cache_file = f"{cache_dir}/sha256-06ee538c626fbf4bdcec2199b7225c8034f26e2b46a7b5cb7ab385c8e8c00efa.cache"
         self.assertEqual(
             common.read_json(cache_file),
@@ -171,7 +171,7 @@ async def test_output_fields(self):
                 "6beb306dc5b91513f353ecdb6aaedee8a9864b3a2f20d91f0d5b27510152acf2",
                 "generated_on": "2021-09-14T21:23:45+00:00",
                 "system_fingerprint": "test-fp",
-                "task_version": 0,
+                "task_version": 1,
                 "result": {
                     "spans": [(1, 3)],
                     "dsa_history": False,
@@ -194,7 +194,7 @@ async def test_trailing_whitespace_removed(self):
         kwargs = self.mock_create.call_args.kwargs
         for message in kwargs["messages"]:
             if message["role"] == "user":
-                self.assertIn("\nTest\n  lines\n", message["content"])
+                self.assertIn("\nTest\n  lines", message["content"])
                 break
         else:
             assert False, "No user message found"