From a3b214855d1ecea9fc032a509b7f583ee9f4d45f Mon Sep 17 00:00:00 2001 From: Deep Koluguri Date: Fri, 22 May 2026 22:04:07 -0400 Subject: [PATCH] feat: integrate postgres database to save clinical notes --- k8s/kustomization.yaml | 1 + k8s/postgres.yaml | 39 ++++++++++++++ requirements.txt | 3 ++ src/curaflow_agent/activities.py | 32 +++++++++++ src/curaflow_agent/db.py | 33 ++++++++++++ src/curaflow_agent/worker.py | 6 ++- src/curaflow_agent/workflows.py | 20 +++++-- test_agents.py | 92 ++++++++++++++++++++++++++++++++ 8 files changed, 219 insertions(+), 7 deletions(-) create mode 100644 k8s/postgres.yaml create mode 100644 src/curaflow_agent/db.py create mode 100644 test_agents.py diff --git a/k8s/kustomization.yaml b/k8s/kustomization.yaml index 72b7a45..ab228ae 100644 --- a/k8s/kustomization.yaml +++ b/k8s/kustomization.yaml @@ -2,3 +2,4 @@ apiVersion: kustomize.config.k8s.io/v1beta1 kind: Kustomization resources: - curaflow-agent/deployment.yaml + - postgres.yaml diff --git a/k8s/postgres.yaml b/k8s/postgres.yaml new file mode 100644 index 0000000..874fd33 --- /dev/null +++ b/k8s/postgres.yaml @@ -0,0 +1,39 @@ +apiVersion: apps/v1 +kind: Deployment +metadata: + name: postgres + namespace: agents-runtime +spec: + replicas: 1 + selector: + matchLabels: + app: postgres + template: + metadata: + labels: + app: postgres + spec: + containers: + - name: postgres + image: postgres:15-alpine + env: + - name: POSTGRES_DB + value: curaflow + - name: POSTGRES_USER + value: admin + - name: POSTGRES_PASSWORD + value: secret123 + ports: + - containerPort: 5432 +--- +apiVersion: v1 +kind: Service +metadata: + name: postgres + namespace: agents-runtime +spec: + selector: + app: postgres + ports: + - port: 5432 + targetPort: 5432 diff --git a/requirements.txt b/requirements.txt index e7cee6a..cd5bf7f 100644 --- a/requirements.txt +++ b/requirements.txt @@ -4,3 +4,6 @@ langchain langchain-openai temporalio pydantic + +SQLAlchemy==2.0.29 +psycopg2-binary==2.9.9 diff --git a/src/curaflow_agent/activities.py b/src/curaflow_agent/activities.py index b54b71a..85148ed 100644 --- a/src/curaflow_agent/activities.py +++ b/src/curaflow_agent/activities.py @@ -102,3 +102,35 @@ async def send_alert_activity(alert_data: dict) -> str: # 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." + +@activity.defn +async def save_clinical_record_activity(payload: dict) -> str: + """ + Saves the raw dictation, extracted clinical note, and billing codes to the database. + payload should contain: + - raw_dictation: str + - clinical_note: dict + - billing_codes: dict + """ + # Import here to avoid circular imports if any, or just import db at the top. + from .db import SessionLocal, ClinicalRecord, init_db + + # Ensure tables exist (in production this would be handled by migrations like Alembic) + init_db() + + db = SessionLocal() + try: + record = ClinicalRecord( + raw_dictation=payload.get("raw_dictation", ""), + clinical_note_json=payload.get("clinical_note", {}), + billing_codes_json=payload.get("billing_codes", {}) + ) + db.add(record) + db.commit() + db.refresh(record) + return f"Successfully saved record with ID: {record.id}" + except Exception as e: + db.rollback() + raise e + finally: + db.close() diff --git a/src/curaflow_agent/db.py b/src/curaflow_agent/db.py new file mode 100644 index 0000000..555c429 --- /dev/null +++ b/src/curaflow_agent/db.py @@ -0,0 +1,33 @@ +import os +from sqlalchemy import create_engine, Column, Integer, String, Text, DateTime, JSON +from sqlalchemy.orm import declarative_base, sessionmaker +from datetime import datetime + +DATABASE_URL = os.environ.get( + "DATABASE_URL", + "postgresql://admin:secret123@postgres.agents-runtime.svc.cluster.local:5432/curaflow" +) + +engine = create_engine(DATABASE_URL) +SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine) +Base = declarative_base() + +class ClinicalRecord(Base): + __tablename__ = "clinical_records" + + id = Column(Integer, primary_key=True, index=True) + patient_id = Column(String(50), index=True, default="unknown") + created_at = Column(DateTime, default=datetime.utcnow) + + # Raw Data + raw_dictation = Column(Text, nullable=False) + + # Structured Note (JSON) + clinical_note_json = Column(JSON, nullable=True) + + # Billing + billing_codes_json = Column(JSON, nullable=True) + +# Create tables if they don't exist +def init_db(): + Base.metadata.create_all(bind=engine) diff --git a/src/curaflow_agent/worker.py b/src/curaflow_agent/worker.py index 5ef18a5..3df020e 100644 --- a/src/curaflow_agent/worker.py +++ b/src/curaflow_agent/worker.py @@ -15,7 +15,8 @@ from curaflow_agent.activities import ( structure_note_activity, generate_billing_codes_activity, analyze_lab_anomaly_activity, - send_alert_activity + send_alert_activity, + save_clinical_record_activity ) async def main(): @@ -40,7 +41,8 @@ async def main(): structure_note_activity, generate_billing_codes_activity, analyze_lab_anomaly_activity, - send_alert_activity + send_alert_activity, + save_clinical_record_activity ], ) diff --git a/src/curaflow_agent/workflows.py b/src/curaflow_agent/workflows.py index a9d3ee0..0bd0164 100644 --- a/src/curaflow_agent/workflows.py +++ b/src/curaflow_agent/workflows.py @@ -6,7 +6,8 @@ with workflow.unsafe.imports_passed_through(): structure_note_activity, generate_billing_codes_activity, analyze_lab_anomaly_activity, - send_alert_activity + send_alert_activity, + save_clinical_record_activity ) @workflow.defn @@ -32,13 +33,22 @@ class ClinicalIntakeWorkflow: start_to_close_timeout=timedelta(minutes=10) ) - # Step 3: Return the final, combined object - # Note: We are returning it directly here so it's visible in the Temporal UI. - # In a real integration, we would trigger a 3rd activity here to POST this - # JSON payload to the client's BYOD (Bring Your Own Database) API endpoint. + # Step 3: Save to the client's database via ORM + save_payload = { + "raw_dictation": raw_dictation, + "clinical_note": structured_note, + "billing_codes": billing_codes + } + + save_status = await workflow.execute_activity( + save_clinical_record_activity, + save_payload, + start_to_close_timeout=timedelta(minutes=1) + ) return { "status": "success", + "db_status": save_status, "clinical_note": structured_note, "billing_codes": billing_codes } diff --git a/test_agents.py b/test_agents.py new file mode 100644 index 0000000..acce328 --- /dev/null +++ b/test_agents.py @@ -0,0 +1,92 @@ +import asyncio +import uuid +import sys +from temporalio.client import Client + +async def test_scribe_and_coder(client: Client): + print("\n=== Testing AI Scribe & Medical Coder ===") + dictation = ( + "Patient is a 45-year-old male presenting with severe, crushing chest pain radiating to his left arm. " + "He reports shortness of breath and sweating. He has a past medical history of hypertension and hyperlipidemia. " + "Current medications include Lisinopril 20mg daily and Atorvastatin 40mg daily. " + "He is allergic to Penicillin. I suspect acute myocardial infarction. I will order an EKG, Troponin levels, " + "and administer Aspirin 324mg and sublingual Nitroglycerin immediately. Admitting to Cardiac ICU." + ) + + print(f"Sending Dictation: {dictation[:100]}...") + + try: + # Start the Clinical Intake Workflow + result = await client.execute_workflow( + "ClinicalIntakeWorkflow", + dictation, + id=f"intake-{uuid.uuid4()}", + task_queue="curaflow-tasks", + ) + print("Success! Workflow completed.") + import json + print(json.dumps(result, indent=2)) + except Exception as e: + print(f"Workflow failed: {e}") + + +async def test_lab_watcher(client: Client): + print("\n=== Testing Lab Result Anomaly Watcher ===") + patient_id = f"patient-{uuid.uuid4()}" + + # 1. Start the long-running workflow with baseline data + print(f"Starting long-running Lab Result Watcher for patient: {patient_id}") + baseline_data = { + "baseline": "Hypertension, CKD Stage 3", + "medications": ["Lisinopril", "Spironolactone"] + } + + # We use start_workflow instead of execute_workflow because it runs infinitely + handle = await client.start_workflow( + "LabResultWatcherWorkflow", + baseline_data, + id=patient_id, + task_queue="curaflow-tasks", + ) + + # 2. Wait a moment + await asyncio.sleep(2) + + # 3. Send a completely normal lab result + normal_lab = "Serum Potassium: 4.2 mEq/L (Normal range 3.6-5.2)" + print(f"Sending normal lab result: {normal_lab}") + await handle.signal("add_lab_result", normal_lab) + + await asyncio.sleep(10) # Give AI time to process + + # 4. Send a dangerous lab result! + # Lisinopril + Spironolactone + High Potassium = Life threatening hyperkalemia + dangerous_lab = "Serum Potassium: 6.8 mEq/L (Critical High)" + print(f"Sending dangerous lab result: {dangerous_lab}") + await handle.signal("add_lab_result", dangerous_lab) + + await asyncio.sleep(10) # Give AI time to process and trigger alert + + print("Discharging patient to close the workflow...") + await handle.signal("discharge_patient") + + result = await handle.result() + print(f"Patient Discharged. Final Workflow Result: {result}") + + +async def main(): + # Use localhost:7233 if port forwarding, or cluster internal if running inside cluster + temporal_url = "temporal-frontend.ai-core.svc.cluster.local:7233" + print(f"Connecting to Temporal at {temporal_url}...") + try: + client = await Client.connect(temporal_url, namespace="default") + print("Connected!") + except Exception as e: + print(f"Failed to connect: {e}") + sys.exit(1) + + await test_scribe_and_coder(client) + await test_lab_watcher(client) + +if __name__ == "__main__": + asyncio.run(main())