feat: integrate postgres database to save clinical notes
Build Agents Runtime / build (push) Successful in 1m25s Details

This commit is contained in:
Deep Koluguri 2026-05-22 22:04:07 -04:00
parent dbd68fc529
commit a3b214855d
8 changed files with 219 additions and 7 deletions

View File

@ -2,3 +2,4 @@ apiVersion: kustomize.config.k8s.io/v1beta1
kind: Kustomization
resources:
- curaflow-agent/deployment.yaml
- postgres.yaml

39
k8s/postgres.yaml Normal file
View File

@ -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

View File

@ -4,3 +4,6 @@ langchain
langchain-openai
temporalio
pydantic
SQLAlchemy==2.0.29
psycopg2-binary==2.9.9

View File

@ -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()

33
src/curaflow_agent/db.py Normal file
View File

@ -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)

View File

@ -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
],
)

View File

@ -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
}

92
test_agents.py Normal file
View File

@ -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())