diff --git a/src/curaflow_agent/activities.py b/src/curaflow_agent/activities.py index 7ef4648..c254e81 100644 --- a/src/curaflow_agent/activities.py +++ b/src/curaflow_agent/activities.py @@ -4,6 +4,7 @@ from temporalio import activity from langchain_openai import ChatOpenAI from langchain_core.prompts import PromptTemplate from .models import ClinicalNote, BillingCodesOutput +from .lab_models import AnomalyReport # In the cluster, Ollama provides an OpenAI-compatible endpoint. # We fetch model/URL from env vars so it's easy to override for testing. @@ -58,3 +59,43 @@ async def generate_billing_codes_activity(clinical_note_dict: dict) -> dict: result: BillingCodesOutput = await chain.ainvoke({"note": note_str}) return result.model_dump() + +@activity.defn +async def analyze_lab_anomaly_activity(input_data: dict) -> dict: + """ + Analyzes a new lab result against a patient's baseline and medications. + input_data should contain: + - baseline: string describing baseline conditions + - medications: list of strings + - new_lab_result: string describing the new lab result + """ + llm = get_llm() + structured_llm = llm.with_structured_output(AnomalyReport) + + prompt = PromptTemplate.from_template( + "You are an expert AI Medical Diagnostic Assistant.\n" + "Review the following patient data and the new lab result. Determine if there is a dangerous anomaly or correlation (e.g., a lab result that is dangerous given the patient's current medications or baseline).\n\n" + "Baseline Conditions: {baseline}\n" + "Current Medications: {medications}\n" + "New Lab Result: {new_lab_result}\n\n" + "Provide a structured assessment." + ) + + chain = prompt | structured_llm + result: AnomalyReport = await chain.ainvoke({ + "baseline": input_data.get("baseline", "None"), + "medications": ", ".join(input_data.get("medications", [])), + "new_lab_result": input_data.get("new_lab_result", "Unknown") + }) + + return result.model_dump() + +@activity.defn +async def send_alert_activity(alert_data: dict) -> str: + """ + Mocks sending an alert to the physician. + """ + # In a real app, this might send an SMS via Twilio or a push notification. + # We will just print it out. + print(f"!!! CRITICAL ALERT Triggered !!!\n{alert_data['danger_explanation']}\nRecommended Action: {alert_data['recommended_action']}") + return "Alert successfully dispatched to physician." diff --git a/src/curaflow_agent/lab_models.py b/src/curaflow_agent/lab_models.py new file mode 100644 index 0000000..e92e94e --- /dev/null +++ b/src/curaflow_agent/lab_models.py @@ -0,0 +1,6 @@ +from pydantic import BaseModel, Field + +class AnomalyReport(BaseModel): + is_dangerous: bool = Field(description="True if the new lab result presents a dangerous correlation or anomaly compared to baseline and medications.") + danger_explanation: str = Field(description="A concise explanation of why the result is dangerous, or 'Normal' if safe.") + recommended_action: str = Field(description="The recommended next step for the physician (e.g., 'Immediately stop Lisinopril', 'Monitor closely').") diff --git a/src/curaflow_agent/worker.py b/src/curaflow_agent/worker.py index 6a91956..5ef18a5 100644 --- a/src/curaflow_agent/worker.py +++ b/src/curaflow_agent/worker.py @@ -10,8 +10,13 @@ from temporalio.worker import Worker logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s") logger = logging.getLogger(__name__) -from curaflow_agent.workflows import ClinicalIntakeWorkflow -from curaflow_agent.activities import structure_note_activity, generate_billing_codes_activity +from curaflow_agent.workflows import ClinicalIntakeWorkflow, LabResultWatcherWorkflow +from curaflow_agent.activities import ( + structure_note_activity, + generate_billing_codes_activity, + analyze_lab_anomaly_activity, + send_alert_activity +) async def main(): temporal_url = os.environ.get("TEMPORAL_URL", "temporal-frontend.ai-core.svc.cluster.local:7233") @@ -30,8 +35,13 @@ async def main(): worker = Worker( client, task_queue=task_queue, - workflows=[ClinicalIntakeWorkflow], - activities=[structure_note_activity, generate_billing_codes_activity], + workflows=[ClinicalIntakeWorkflow, LabResultWatcherWorkflow], + activities=[ + structure_note_activity, + generate_billing_codes_activity, + analyze_lab_anomaly_activity, + send_alert_activity + ], ) logger.info(f"Starting CuraFlow agent worker on queue '{task_queue}'...") diff --git a/src/curaflow_agent/workflows.py b/src/curaflow_agent/workflows.py index 18e52ab..a9d3ee0 100644 --- a/src/curaflow_agent/workflows.py +++ b/src/curaflow_agent/workflows.py @@ -2,7 +2,12 @@ from datetime import timedelta from temporalio import workflow with workflow.unsafe.imports_passed_through(): - from .activities import structure_note_activity, generate_billing_codes_activity + from .activities import ( + structure_note_activity, + generate_billing_codes_activity, + analyze_lab_anomaly_activity, + send_alert_activity + ) @workflow.defn class ClinicalIntakeWorkflow: @@ -37,3 +42,60 @@ class ClinicalIntakeWorkflow: "clinical_note": structured_note, "billing_codes": billing_codes } + +@workflow.defn +class LabResultWatcherWorkflow: + def __init__(self) -> None: + self._pending_lab_results = [] + self._is_discharged = False + + @workflow.signal + def add_lab_result(self, result: str) -> None: + self._pending_lab_results.append(result) + + @workflow.signal + def discharge_patient(self) -> None: + self._is_discharged = True + + @workflow.run + async def run(self, baseline_data: dict) -> dict: + """ + Runs continuously, processing new lab results via signals. + baseline_data: {"baseline": "...", "medications": ["..."]} + """ + alerts_triggered = 0 + + while not self._is_discharged: + # Wait until a new lab result arrives OR the patient is discharged + await workflow.wait_condition( + lambda: len(self._pending_lab_results) > 0 or self._is_discharged + ) + + # Process all pending results + while self._pending_lab_results: + new_result = self._pending_lab_results.pop(0) + + input_data = { + "baseline": baseline_data.get("baseline", "None"), + "medications": baseline_data.get("medications", []), + "new_lab_result": new_result + } + + anomaly_report = await workflow.execute_activity( + analyze_lab_anomaly_activity, + input_data, + start_to_close_timeout=timedelta(minutes=5) + ) + + if anomaly_report.get("is_dangerous", False): + await workflow.execute_activity( + send_alert_activity, + anomaly_report, + start_to_close_timeout=timedelta(minutes=1) + ) + alerts_triggered += 1 + + return { + "status": "discharged", + "total_alerts_triggered": alerts_triggered + }