graph.py

# -----------------------------------------------
# graph.py
# Wires all 6 agents together using LangGraph
# Defines the pipeline flow
# -----------------------------------------------

from langgraph.graph import StateGraph, END
from typing import TypedDict, List, Dict, Any
from agents.document_agent    import run_document_agent
from agents.abnormality_agent import run_abnormality_agent
from agents.rag_agent         import run_rag_agent
from agents.explanation_agent import run_explanation_agent
from agents.specialist_agent  import run_specialist_agent
from agents.report_agent      import run_report_agent

# -----------------------------------------------
# Shared State — passed between all agents
# -----------------------------------------------
class GraphState(TypedDict):
    # Input
    file_path:             str
    language:              str
    is_imaging:            bool

    # Document Agent output
    document_type:         str
    raw_text:              str
    findings:              str
    abnormalities:         str
    patient_info:          str
    doctor_notes:          str
    gemini_raw:            str
    error:                 str

    # Abnormality Agent output
    abnormality_output:    Dict
    overall_severity:      str
    urgent_flags:          List
    abnormal_values:       List
    normal_values:         List
    severity_reason:       str

    # RAG Agent output
    rag_results:           str
    matched_conditions:    List

    # Explanation Agent output
    explanation:           str
    explanation_parsed:    Dict
    summary:               str
    what_to_do:            str

    # Specialist Agent output
    primary_specialist:    str
    secondary_specialists: List
    urgency:               str
    urgency_timing:        str
    urgency_reason:        str
    where_to_go:           str
    questions_for_doctor:  List
    what_to_bring:         List
    quick_specialists:     List
    specialist_raw:        str

    # Report Agent output
    report_path:           str
    report_success:        bool

# -----------------------------------------------
# Conditional edge — stop if error
# -----------------------------------------------
def should_continue(state: GraphState) -> str:
    if state.get("error") and state.get("document_type") == "unknown":
        print(f"[Graph] Stopping pipeline — error: {state['error']}")
        return "end"
    return "continue"

# -----------------------------------------------
# Build LangGraph pipeline
# -----------------------------------------------
def build_graph():
    graph = StateGraph(GraphState)

    # Add all agents as nodes
    graph.add_node("document_agent",    run_document_agent)
    graph.add_node("abnormality_agent", run_abnormality_agent)
    graph.add_node("rag_agent",         run_rag_agent)
    graph.add_node("explanation_agent", run_explanation_agent)
    graph.add_node("specialist_agent",  run_specialist_agent)
    graph.add_node("report_agent",      run_report_agent)

    # Set entry point
    graph.set_entry_point("document_agent")

    # Conditional edge after document agent
    # Stop if document can't be read
    graph.add_conditional_edges(
        "document_agent",
        should_continue,
        {
            "continue": "abnormality_agent",
            "end":      END
        }
    )

    # Linear flow for remaining agents
    graph.add_edge("abnormality_agent", "rag_agent")
    graph.add_edge("rag_agent",         "explanation_agent")
    graph.add_edge("explanation_agent", "specialist_agent")
    graph.add_edge("specialist_agent",  "report_agent")
    graph.add_edge("report_agent",      END)

    return graph.compile()

# -----------------------------------------------
# Run the complete pipeline
# -----------------------------------------------
def run_pipeline(file_path: str, language: str = "English") -> dict:
    print("\n" + "=" * 60)
    print("MEDICAL REPORT ANALYZER — PIPELINE STARTED")
    print("=" * 60)

    # Build initial state
    initial_state = GraphState(
        # Input
        file_path             = file_path,
        language              = language,
        is_imaging = False,

        # Document Agent
        document_type         = "",
        raw_text              = "",
        findings              = "",
        abnormalities         = "",
        patient_info          = "",
        doctor_notes          = "",
        gemini_raw            = "",
        error                 = "",

        # Abnormality Agent
        abnormality_output    = {},
        overall_severity      = "",
        urgent_flags          = [],
        abnormal_values       = [],
        normal_values         = [],
        severity_reason       = "",

        # RAG Agent
        rag_results           = "",
        matched_conditions    = [],

        # Explanation Agent
        explanation           = "",
        explanation_parsed    = {},
        summary               = "",
        what_to_do            = "",

        # Specialist Agent
        primary_specialist    = "",
        secondary_specialists = [],
        urgency               = "",
        urgency_timing        = "",
        urgency_reason        = "",
        where_to_go           = "",
        questions_for_doctor  = [],
        what_to_bring         = [],
        quick_specialists     = [],
        specialist_raw        = "",

        # Report Agent
        report_path           = "",
        report_success        = False,
    )

    # Build and run graph
    pipeline = build_graph()
    final_state = pipeline.invoke(initial_state)

    print("\n" + "=" * 60)
    print("PIPELINE COMPLETE")
    print("=" * 60)
    print(f"Severity    : {final_state.get('overall_severity')}")
    print(f"Conditions  : {final_state.get('matched_conditions')}")
    print(f"Specialist  : {final_state.get('primary_specialist')}")
    print(f"Report      : {final_state.get('report_path')}")

    return final_state