AI Agent Budget Planning - Graph Builder

graph_builder.py

def route_after_parse(state: AgentState) -> str:
    return "error" if "error" in state else "run"

def route_after_run(state: AgentState) -> str:
    return "error" if "error" in state else "summarize"

def build_budagent_graph():

    graph_builder = StateGraph(AgentState)
    graph_builder.add_node("parse", parse_email)
    graph_builder.add_node("run", run_budget)
    graph_builder.add_node("summarize", summarize)
    graph_builder.add_node("error", handle_error)
    graph_builder.add_edge(START, "parse")
    graph_builder.add_conditional_edges("parse", route_after_parse, {
        "run": "run",
        "error": "error",
    })
    graph_builder.add_conditional_edges("run", route_after_run, {
        "summarize": "summarize",
        "error": "error",
    })
    graph_builder.add_edge("summarize", END)
    graph_builder.add_edge("error", END)

    return graph_builder.compile()