Subgraphs: Composing Large Agent Systems

from langgraph.graph import StateGraph, END
from typing import TypedDict, Annotated
import operator

# --- Subgraph 1: Drug Research ---
class ResearchSubState(TypedDict):
    drug_name: str
    research_findings: Annotated[list[str], operator.add]
    is_complete: bool

def search_literature(state: ResearchSubState) -> dict:
    """Search medical literature for drug information."""
    finding = f"Found mechanism data for {state['drug_name']}: inhibits VKOR enzyme"
    return {"research_findings": [finding]}

def extract_key_facts(state: ResearchSubState) -> dict:
    """Extract and structure key facts from research."""
    finding = f"Extracted: {state['drug_name']} — monitoring: INR 2.0-3.0"
    return {"research_findings": [finding], "is_complete": True}

research_graph = StateGraph(ResearchSubState)
research_graph.add_node("search", search_literature)
research_graph.add_node("extract", extract_key_facts)
research_graph.set_entry_point("search")
research_graph.add_edge("search", "extract")
research_graph.add_edge("extract", END)

research_subgraph = research_graph.compile()  # Compiled subgraph

# --- Subgraph 2: Safety Check ---
class SafetySubState(TypedDict):
    drug_name: str
    interactions: Annotated[list[str], operator.add]
    risk_level: str

def check_interactions(state: SafetySubState) -> dict:
    return {"interactions": [f"{state['drug_name']} interacts with NSAIDs (major)"]}

def assess_risk(state: SafetySubState) -> dict:
    return {"risk_level": "moderate"}

safety_graph = StateGraph(SafetySubState)
safety_graph.add_node("check", check_interactions)
safety_graph.add_node("assess", assess_risk)
safety_graph.set_entry_point("check")
safety_graph.add_edge("check", "assess")
safety_graph.add_edge("assess", END)

safety_subgraph = safety_graph.compile()

class ParentState(TypedDict):
    drug_name: str
    # These fields must overlap with subgraph states
    research_findings: Annotated[list[str], operator.add]
    interactions: Annotated[list[str], operator.add]
    risk_level: str
    is_complete: bool
    final_report: str

def synthesize(state: ParentState) -> dict:
    """Combine all subgraph outputs into a final report."""
    report = f"""Drug Report: {state['drug_name']}

Research findings:
{chr(10).join(f'- {f}' for f in state['research_findings'])}

Safety:
{chr(10).join(f'- {i}' for i in state['interactions'])}
Risk level: {state.get('risk_level', 'unknown')}"""

    return {"final_report": report}

# Build parent graph
parent = StateGraph(ParentState)

# Add subgraphs as nodes
parent.add_node("research", research_subgraph)
parent.add_node("safety", safety_subgraph)
parent.add_node("synthesize", synthesize)

parent.set_entry_point("research")
parent.add_edge("research", "safety")
parent.add_edge("safety", "synthesize")
parent.add_edge("synthesize", END)

parent_app = parent.compile()

# Run the full pipeline
result = parent_app.invoke({
    "drug_name": "warfarin",
    "research_findings": [],
    "interactions": [],
    "risk_level": "",
    "is_complete": False,
    "final_report": "",
})

print(result["final_report"])

# Subgraph has: drug_name (input), internal_temp_data (internal), findings (output)
# Parent must have: drug_name (to pass in), findings (to receive output)
# Parent does NOT need: internal_temp_data (stays inside subgraph)

from langgraph.types import Send

class ParallelAnalysisState(TypedDict):
    drug_list: list[str]
    per_drug_results: Annotated[list[dict], operator.add]
    combined_report: str

def create_parallel_tasks(state: ParallelAnalysisState) -> list[Send]:
    """Fan out: one research task per drug."""
    return [
        Send("research_drug", {"drug_name": drug, "research_findings": [], "is_complete": False})
        for drug in state["drug_list"]
    ]

def collect_results(state: ParallelAnalysisState) -> dict:
    """Fan in: collect all parallel results."""
    report = f"Analyzed {len(state['per_drug_results'])} drugs:\n"
    for result in state["per_drug_results"]:
        report += f"- {result.get('drug_name', 'unknown')}\n"
    return {"combined_report": report}

multi_drug_graph = StateGraph(ParallelAnalysisState)
multi_drug_graph.add_node("research_drug", research_subgraph)
multi_drug_graph.add_node("collect", collect_results)

multi_drug_graph.add_conditional_edges("__start__", create_parallel_tasks)
multi_drug_graph.add_edge("research_drug", "collect")
multi_drug_graph.add_edge("collect", END)

# Test research subgraph independently
research_result = research_subgraph.invoke({
    "drug_name": "metformin",
    "research_findings": [],
    "is_complete": False,
})

assert research_result["is_complete"] is True
assert len(research_result["research_findings"]) > 0
assert "metformin" in research_result["research_findings"][0].lower()
print("Research subgraph test passed")

# Test safety subgraph independently
safety_result = safety_subgraph.invoke({
    "drug_name": "metformin",
    "interactions": [],
    "risk_level": "",
})

assert safety_result["risk_level"] in ("low", "moderate", "high")
print("Safety subgraph test passed")