graph.py

from langgraph.graph import StateGraph, START, END
from state import AgentState
from nodes import clarity_agent, research_agent, validator_agent, synthesis_agent

# Routing Functions

def route_clarity(state: AgentState):
    if state["clarity_status"] == "needs_clarification":
        return END # End run, will display the clarification question to user
    return "research_agent"

def route_validator(state: AgentState):
    # Logic: If insufficient AND attempts < 3 -> Retry Research
    # Else -> Synthesis
    if state["validation_result"] == "insufficient" and state.get("attempt_count", 0) < 3:
        return "research_agent"
    return "synthesis_agent"

def route_research(state: AgentState):
    # Research always goes to Validator
    # But we can check confidence here if we wanted to skip Validator as per requirements?
    # Req: "ROUTES TO: Validator Agent (if confidence < 6) OR Synthesis Agent (if confidence >= 6)"
    # Wait, requirements say:
    # "ROUTES TO: Validator Agent (if confidence < 6) OR Synthesis Agent (if confidence >= 6)"
    # Correction: The requirements actually describe a specific flow.
    # Let's align exactly with requirements text.
    
    # Requirement 1b says:
    # ROUTES TO: Validator Agent (if confidence < 6) OR Synthesis Agent (if confidence >= 6)
    
    score = state.get("confidence_score", 0)
    if score >= 6.0:
        return "synthesis_agent"
    return "validator_agent"


# Build Graph
builder = StateGraph(AgentState)

# Add Nodes
builder.add_node("clarity_agent", clarity_agent)
builder.add_node("research_agent", research_agent)
builder.add_node("validator_agent", validator_agent)
builder.add_node("synthesis_agent", synthesis_agent)

# Add Edges
builder.add_edge(START, "clarity_agent")

builder.add_conditional_edges(
    "clarity_agent",
    route_clarity,
    {END: END, "research_agent": "research_agent"}
)

builder.add_conditional_edges(
    "research_agent",
    route_research,
    {"validator_agent": "validator_agent", "synthesis_agent": "synthesis_agent"}
)

builder.add_conditional_edges(
    "validator_agent",
    route_validator,
    {"research_agent": "research_agent", "synthesis_agent": "synthesis_agent"}
)

builder.add_edge("synthesis_agent", END)

graph = builder.compile()