Router Node

class BaseRouter(Node):
    async def process(self, task_context: TaskContext) -> TaskContext:
        pass

    def route(self, task_context: TaskContext) -> Node:
        for route_node in self.routes:
            route_node.task_context = task_context
            next_node = route_node.determine_next_node(task_context)
            if next_node:
                return next_node
        return self.fallback if self.fallback else None

class RouterNode(ABC):
    @abstractmethod
    def determine_next_node(self, task_context: TaskContext) -> Optional[Node]:
        pass

    @property
    def node_name(self):
        return self.__class__.__name__

class TicketRouterNode(BaseRouter):
    def __init__(self):
        self.routes = [
            CloseTicketRouter(),
            EscalationRouter(),
            InvoiceRouter(),
        ]
        self.fallback = GenerateResponseNode()

class CloseTicketRouter(RouterNode):
    def determine_next_node(self, task_context: TaskContext) -> Optional[Node]:
        output: FilterSpamNode.OutputType = task_context.nodes["FilterSpamNode"][
            "result"
        ].output
        if not output.is_human and output.confidence > 0.8:
            return CloseTicketNode()
        return None

class EscalationRouter(RouterNode):
    def determine_next_node(self, task_context: TaskContext) -> Optional[Node]:
        analysis = task_context.nodes["DetermineTicketIntentNode"]["result"].output
        if analysis.intent.escalate or analysis.escalate:
            return EscalateTicketNode()
        return None

NodeConfig(
    node=TicketRouterNode,
    connections=[CloseTicketNode, EscalateTicketNode, ProcessInvoiceNode, GenerateResponseNode],
    is_router=True,
    description="Route tickets based on analysis results",
)