Spaces:
Sleeping
Sleeping
File size: 3,518 Bytes
6dc9d46 aefac4f 6dc9d46 aefac4f 6dc9d46 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 | """
MediGuard AI RAG-Helper
State definitions for LangGraph workflow
"""
from typing import Dict, List, Any, Optional, Annotated
from typing_extensions import TypedDict
from pydantic import BaseModel, ConfigDict
from src.config import ExplanationSOP
import operator
class AgentOutput(BaseModel):
"""Structured output from each specialist agent"""
agent_name: str
findings: Any
metadata: Optional[Dict[str, Any]] = None
class BiomarkerFlag(BaseModel):
"""Structure for flagged biomarker values"""
name: str
value: float
unit: str
status: str # "NORMAL", "HIGH", "LOW", "CRITICAL_HIGH", "CRITICAL_LOW"
reference_range: str
warning: Optional[str] = None
class SafetyAlert(BaseModel):
"""Structure for safety warnings"""
severity: str # "LOW", "MEDIUM", "HIGH", "CRITICAL"
biomarker: Optional[str] = None
message: str
action: str
class KeyDriver(BaseModel):
"""Biomarker contribution to prediction"""
biomarker: str
value: Any
contribution: Optional[str] = None
explanation: str
evidence: Optional[str] = None
class GuildState(TypedDict):
"""
The shared state/workspace for the Clinical Insight Guild.
Passed between all agent nodes in the LangGraph workflow.
"""
# === Input Data ===
patient_biomarkers: Dict[str, float] # Raw biomarker values
model_prediction: Dict[str, Any] # Disease prediction from ML model
patient_context: Optional[Dict[str, Any]] # Age, gender, BMI, etc.
# === Workflow Control ===
plan: Optional[Dict[str, Any]] # Execution plan from Planner
sop: ExplanationSOP # Current operating procedures
# === Agent Outputs (Accumulated) - Use Annotated with operator.add for parallel updates ===
agent_outputs: Annotated[List[AgentOutput], operator.add]
biomarker_flags: Annotated[List[BiomarkerFlag], operator.add]
safety_alerts: Annotated[List[SafetyAlert], operator.add]
biomarker_analysis: Optional[Dict[str, Any]]
# === Final Structured Output ===
final_response: Optional[Dict[str, Any]]
# === Metadata ===
processing_timestamp: Optional[str]
sop_version: Optional[str]
# === Input Schema for Patient Data ===
class PatientInput(BaseModel):
"""Standard input format for patient assessment"""
biomarkers: Dict[str, float]
model_prediction: Dict[str, Any] # Contains: disease, confidence, probabilities
patient_context: Optional[Dict[str, Any]] = None
def model_post_init(self, __context: Any) -> None:
if self.patient_context is None:
self.patient_context = {"age": None, "gender": None, "bmi": None}
model_config = ConfigDict(json_schema_extra={
"example": {
"biomarkers": {
"Glucose": 185,
"HbA1c": 8.2,
"Hemoglobin": 13.5,
"Platelets": 220000,
"Cholesterol": 210
},
"model_prediction": {
"disease": "Diabetes",
"confidence": 0.89,
"probabilities": {
"Diabetes": 0.89,
"Heart Disease": 0.06,
"Anemia": 0.03,
"Thalassemia": 0.01,
"Thrombocytopenia": 0.01
}
},
"patient_context": {
"age": 52,
"gender": "male",
"bmi": 31.2
}
}
})
|