Spaces:
Runtime error
Runtime error
| import logging | |
| from langgraph.graph import StateGraph, END | |
| from langgraph.checkpoint.memory import MemorySaver | |
| from src.workflow.state import GovBridgeState | |
| from src.agents.input_agent import input_agent | |
| from src.agents.ideas_agent import ideas_agent | |
| from src.agents.complaint_agent import complaint_agent | |
| from src.agents.email_writer_agent import email_writer_agent | |
| _log = logging.getLogger(__name__) | |
| def _route_after_input(state: GovBridgeState) -> str: | |
| category = state.get("category") | |
| if category == "idea": | |
| _log.info("routing → ideas_agent") | |
| return "ideas_agent" | |
| if category == "complaint": | |
| _log.info("routing → complaint_agent") | |
| return "complaint_agent" | |
| _log.info("routing → input_agent (no category yet)") | |
| return "input_agent" | |
| def _route_after_email_writer(state: GovBridgeState) -> str: | |
| if state.get("draft_approved"): | |
| _log.info("draft approved — routing to END") | |
| return END | |
| _log.info("draft not yet approved — looping back to email_writer") | |
| return "email_writer" | |
| def _route_after_conversation(node_name: str): | |
| def _route(state: GovBridgeState) -> str: | |
| if state.get("conversation_complete"): | |
| _log.info("%s conversation complete — routing to email_writer", node_name) | |
| return "email_writer" | |
| _log.info("%s conversation continuing — looping back", node_name) | |
| return node_name | |
| return _route | |
| def build_graph(): | |
| graph = StateGraph(GovBridgeState) | |
| graph.add_node("input_agent", input_agent) | |
| graph.add_node("ideas_agent", ideas_agent) | |
| graph.add_node("complaint_agent", complaint_agent) | |
| graph.add_node("email_writer", email_writer_agent) | |
| graph.set_entry_point("input_agent") | |
| graph.add_conditional_edges("input_agent", _route_after_input, { | |
| "ideas_agent": "ideas_agent", | |
| "complaint_agent": "complaint_agent", | |
| "input_agent": "input_agent", | |
| }) | |
| graph.add_conditional_edges("ideas_agent", _route_after_conversation("ideas_agent"), { | |
| "email_writer": "email_writer", | |
| "ideas_agent": "ideas_agent", | |
| }) | |
| graph.add_conditional_edges("complaint_agent", _route_after_conversation("complaint_agent"), { | |
| "email_writer": "email_writer", | |
| "complaint_agent": "complaint_agent", | |
| }) | |
| graph.add_conditional_edges("email_writer", _route_after_email_writer, { | |
| END: END, | |
| "email_writer": "email_writer", | |
| }) | |
| return graph.compile(checkpointer=MemorySaver()) | |