use crate::node::NodeType;
use crate::types::GraphState;
pub trait Router: Send + Sync {
fn next(&self, state: &GraphState, current: NodeType) -> NextNode;
}
#[derive(Debug, Clone, PartialEq)]
pub enum NextNode {
LLM,
Tool,
End,
}
pub struct SimpleRouter;
impl Router for SimpleRouter {
fn next(&self, state: &GraphState, current: NodeType) -> NextNode {
match current {
NodeType::LLM => {
if state.has_pending_tool_calls() {
NextNode::Tool
} else {
NextNode::End
}
}
NodeType::Tool => {
NextNode::LLM
}
}
}
}