File size: 4,031 Bytes
e888b38
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
"""
LeukemiaScope Agentic Workflow
LangGraph-based multi-agent system for blood cell analysis
"""

from typing import TypedDict, Annotated
from langgraph.graph import StateGraph, END
from PIL import Image

# Import agent nodes
import sys
import os
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from agents.image_analyzer import image_analyzer_node
from agents.clinical_advisor import clinical_advisor_node
from agents.report_generator import report_generator_node


class WorkflowState(TypedDict):
    """State shared across all agents in the workflow"""
    # Input
    image: Image.Image
    patient_id: str
    patient_context: str
    
    # Image Analysis Results
    classification: str
    confidence: float
    is_leukemia: bool
    raw_response: str
    analysis_complete: bool
    
    # Clinical Advice
    clinical_advice: str
    next_steps: list
    severity: str
    requires_urgent_action: bool
    clinical_complete: bool
    
    # Report
    report: str
    report_complete: bool
    
    # Errors
    error: str


def should_consult_clinical_advisor(state: WorkflowState) -> str:
    """
    Conditional edge: Route to Clinical Advisor only if leukemia is detected
    
    Args:
        state: Current workflow state
        
    Returns:
        "clinical_advisor" if leukemia detected, else "report_generator"
    """
    if state.get("is_leukemia", False):
        return "clinical_advisor"
    return "report_generator"


def create_workflow() -> StateGraph:
    """
    Create the LangGraph workflow for blood cell analysis
    
    Workflow:
    1. Image Analyzer (MedGemma) - Always runs first
    2. Clinical Advisor - Only if leukemia detected
    3. Report Generator - Always runs, generates final report
    """
    # Create the graph
    workflow = StateGraph(WorkflowState)
    
    # Add nodes
    workflow.add_node("image_analyzer", image_analyzer_node)
    workflow.add_node("clinical_advisor", clinical_advisor_node)
    workflow.add_node("report_generator", report_generator_node)
    
    # Set entry point
    workflow.set_entry_point("image_analyzer")
    
    # Add conditional edge from image_analyzer
    workflow.add_conditional_edges(
        "image_analyzer",
        should_consult_clinical_advisor,
        {
            "clinical_advisor": "clinical_advisor",
            "report_generator": "report_generator"
        }
    )
    
    # Clinical advisor always leads to report generator
    workflow.add_edge("clinical_advisor", "report_generator")
    
    # Report generator ends the workflow
    workflow.add_edge("report_generator", END)
    
    return workflow


def compile_workflow():
    """Compile the workflow into an executable graph"""
    workflow = create_workflow()
    return workflow.compile()


# Global compiled workflow
_app = None


def get_app():
    """Get or create the compiled workflow app"""
    global _app
    if _app is None:
        _app = compile_workflow()
    return _app


def run_analysis(image: Image.Image, patient_id: str = "Anonymous", patient_context: str = "") -> dict:
    """
    Run the complete analysis workflow
    
    Args:
        image: PIL Image of blood cell
        patient_id: Optional patient identifier
        patient_context: Optional context about patient
        
    Returns:
        Final state with all results
    """
    app = get_app()
    
    # Initial state
    initial_state = {
        "image": image,
        "patient_id": patient_id,
        "patient_context": patient_context,
        "classification": "",
        "confidence": 0.0,
        "is_leukemia": False,
        "raw_response": "",
        "analysis_complete": False,
        "clinical_advice": "",
        "next_steps": [],
        "severity": "Low",
        "requires_urgent_action": False,
        "clinical_complete": False,
        "report": "",
        "report_complete": False,
        "error": ""
    }
    
    # Run workflow
    final_state = app.invoke(initial_state)
    
    return final_state