Spaces:
Runtime error
Runtime error
File size: 2,522 Bytes
3f713ce 2948416 3f713ce 2948416 3f713ce 2948416 3f713ce 2948416 3f713ce 2948416 3f713ce | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 | 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())
|