Skip to content

Commit 947e08d

Browse files
authored
Merge pull request #436 from smart-on-fhir/mikix/irae-prompt2
feat(irae): update the IRAE system & user prompts
2 parents 03929c6 + 13266d9 commit 947e08d

File tree

4 files changed

+44
-28
lines changed

4 files changed

+44
-28
lines changed

cumulus_etl/etl/studies/irae/irae_tasks.py

Lines changed: 21 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
"""Define tasks for the irae study"""
22

33
import enum
4+
import json
45
from typing import ClassVar
56

67
import pydantic
@@ -42,23 +43,31 @@ class DSAMention(pydantic.BaseModel):
4243

4344

4445
class BaseIraeTask(tasks.BaseOpenAiTaskWithSpans):
45-
task_version = 0
46+
task_version = 1
4647
# Task Version History:
47-
# ** 0 (2025-08): Initial work, still in flux **
48+
# ** 1 (2025-08): Updated prompt **
49+
# ** 0 (2025-08): Initial version **
4850

4951
system_prompt = (
50-
"You are a helpful assistant reviewing kidney-transplant notes for donor-specific "
51-
"antibody (DSA) information. Return *only* valid JSON."
52+
"You are a clinical chart reviewer for a kidney transplant outcomes study.\n"
53+
"Your task is to extract patient-specific information from an unstructured clinical "
54+
"document and map it into a predefined Pydantic schema.\n"
55+
"\n"
56+
"Core Rules:\n"
57+
"1. Base all assertions ONLY on patient-specific information in the clinical document.\n"
58+
" - Never negate or exclude information just because it is not mentioned.\n"
59+
" - Never conflate family history or population-level risk with patient findings.\n"
60+
"2. Do not invent or infer facts beyond what is documented.\n"
61+
"3. Maintain high fidelity to the clinical document language when citing spans.\n"
62+
"4. Always produce structured JSON that conforms to the Pydantic schema provided below.\n"
63+
"\n"
64+
"Pydantic Schema:\n" + json.dumps(DSAMention.model_json_schema())
5265
)
5366
user_prompt = (
54-
"Evaluate the following chart for donor-specific antibody (DSA) information.\n"
55-
"Here is the chart for you to analyze:\n"
56-
"%CLINICAL-NOTE%\n"
57-
"Keep the pydantic structure previously provided in mind when structuring your output.\n"
58-
"Finally, ensure that your final assertions are based on Patient-specific information "
59-
"only.\n"
60-
"For example, we should never deny the presence of an observation because of a lack of "
61-
"family history."
67+
"Evaluate the following clinical document for kidney transplant variables and outcomes.\n"
68+
"Here is the clinical document for you to analyze:\n"
69+
"\n"
70+
"%CLINICAL-NOTE%"
6271
)
6372
response_format = DSAMention
6473

tests/data/irae/output.ndjson

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1 +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]]}}
1+
{"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

Lines changed: 19 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,7 @@
11
"""Tests for etl/studies/irae/"""
22

3+
import json
4+
35
import ddt
46

57
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):
4345
"messages": [
4446
{
4547
"role": "system",
46-
"content": "You are a helpful assistant reviewing kidney-transplant notes "
47-
"for donor-specific antibody (DSA) information. Return *only* valid JSON.",
48+
"content": "You are a clinical chart reviewer for a kidney transplant outcomes study.\n"
49+
"Your task is to extract patient-specific information from an unstructured clinical "
50+
"document and map it into a predefined Pydantic schema.\n"
51+
"\n"
52+
"Core Rules:\n"
53+
"1. Base all assertions ONLY on patient-specific information in the clinical document.\n"
54+
" - Never negate or exclude information just because it is not mentioned.\n"
55+
" - Never conflate family history or population-level risk with patient findings.\n"
56+
"2. Do not invent or infer facts beyond what is documented.\n"
57+
"3. Maintain high fidelity to the clinical document language when citing spans.\n"
58+
"4. Always produce structured JSON that conforms to the Pydantic schema provided below.\n"
59+
"\n"
60+
"Pydantic Schema:\n" + json.dumps(DSAMention.model_json_schema()),
4861
},
4962
{
5063
"role": "user",
51-
"content": "Evaluate the following chart for donor-specific antibody (DSA) "
52-
"information.\n"
53-
"Here is the chart for you to analyze:\n"
54-
"Test note 1\n"
55-
"Keep the pydantic structure previously provided in mind when structuring "
56-
"your output.\n"
57-
"Finally, ensure that your final assertions are based on Patient-specific "
58-
"information only.\n"
59-
"For example, we should never deny the presence of an observation because "
60-
"of a lack of family history.",
64+
"content": "Evaluate the following clinical document for kidney "
65+
"transplant variables and outcomes.\n"
66+
"Here is the clinical document for you to analyze:\n\n"
67+
"Test note 1",
6168
},
6269
],
6370
"model": model_id,

tests/nlp/test_openai.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -90,7 +90,7 @@ async def test_caching(self):
9090
await irae.IraeGptOss120bTask(self.job_config, self.scrubber).run()
9191

9292
self.assertEqual(self.mock_create.call_count, 1)
93-
cache_dir = f"{self.phi_dir}/nlp-cache/irae__nlp_gpt_oss_120b_v0/06ee/"
93+
cache_dir = f"{self.phi_dir}/nlp-cache/irae__nlp_gpt_oss_120b_v1/06ee"
9494
cache_file = f"{cache_dir}/sha256-06ee538c626fbf4bdcec2199b7225c8034f26e2b46a7b5cb7ab385c8e8c00efa.cache"
9595
self.assertEqual(
9696
common.read_json(cache_file),
@@ -171,7 +171,7 @@ async def test_output_fields(self):
171171
"6beb306dc5b91513f353ecdb6aaedee8a9864b3a2f20d91f0d5b27510152acf2",
172172
"generated_on": "2021-09-14T21:23:45+00:00",
173173
"system_fingerprint": "test-fp",
174-
"task_version": 0,
174+
"task_version": 1,
175175
"result": {
176176
"spans": [(1, 3)],
177177
"dsa_history": False,
@@ -194,7 +194,7 @@ async def test_trailing_whitespace_removed(self):
194194
kwargs = self.mock_create.call_args.kwargs
195195
for message in kwargs["messages"]:
196196
if message["role"] == "user":
197-
self.assertIn("\nTest\n lines\n", message["content"])
197+
self.assertIn("\nTest\n lines", message["content"])
198198
break
199199
else:
200200
assert False, "No user message found"

0 commit comments

Comments
 (0)