1use crate::{
2 agents::Agent,
3 llm::LLMClient,
4 types::{AgentContext, AgentType, Result},
5};
6use async_trait::async_trait;
7
8const VALID_AGENTS: &[&str] = &[
10 "product",
11 "invoice",
12 "sales",
13 "finance",
14 "hr",
15 "orchestrator",
16 "research",
17];
18
19pub struct RouterAgent {
24 llm: Box<dyn LLMClient>,
25}
26
27impl RouterAgent {
28 pub fn new(llm: Box<dyn LLMClient>) -> Self {
30 Self { llm }
31 }
32
33 fn parse_routing_decision(output: &str) -> Option<String> {
41 let trimmed = output.trim().to_lowercase();
42
43 if VALID_AGENTS.contains(&trimmed.as_str()) {
45 return Some(trimmed);
46 }
47
48 for word in trimmed.split(|c: char| c.is_whitespace() || c == ':' || c == ',' || c == '.') {
51 let word = word.trim();
52 if VALID_AGENTS.contains(&word) {
53 return Some(word.to_string());
54 }
55 }
56
57 for agent in VALID_AGENTS {
59 if trimmed.contains(agent) {
60 return Some(agent.to_string());
61 }
62 }
63
64 None
65 }
66
67 pub async fn route(&self, query: &str, _context: &AgentContext) -> Result<AgentType> {
69 let system_prompt = self.system_prompt();
70 let response = self.llm.generate_with_system(&system_prompt, query).await?;
71
72 let agent_name = Self::parse_routing_decision(&response);
74
75 match agent_name.as_deref() {
76 Some("product") => Ok(AgentType::Product),
77 Some("invoice") => Ok(AgentType::Invoice),
78 Some("sales") => Ok(AgentType::Sales),
79 Some("finance") => Ok(AgentType::Finance),
80 Some("hr") => Ok(AgentType::HR),
81 Some("orchestrator") | Some("research") => Ok(AgentType::Orchestrator),
82 _ => {
83 tracing::debug!(
85 "Router could not parse output '{}', defaulting to orchestrator",
86 response
87 );
88 Ok(AgentType::Orchestrator)
89 }
90 }
91 }
92}
93
94#[async_trait]
95impl Agent for RouterAgent {
96 async fn execute(&self, _input: &str, _context: &AgentContext) -> Result<String> {
97 Ok("router".to_string())
101 }
102
103 fn system_prompt(&self) -> String {
104 r#"You are a routing agent that classifies user queries and routes them to the appropriate specialized agent.
105
106Available agents:
107- product: Product information, recommendations, catalog queries
108- invoice: Invoice processing, billing questions, payment status
109- sales: Sales data, analytics, performance metrics
110- finance: Financial reports, budgets, expense analysis
111- hr: Human resources, employee information, policies
112- orchestrator: Complex queries requiring multiple agents or research
113
114Analyze the user's query and respond with ONLY the agent name (lowercase, one word).
115Examples:
116- "What products do we have?" → product
117- "Show me last quarter's sales" → sales
118- "What's our hiring policy?" → hr
119- "Create a comprehensive market analysis" → orchestrator
120
121Respond with ONLY the agent name, nothing else."#.to_string()
122 }
123
124 fn agent_type(&self) -> AgentType {
125 AgentType::Router
126 }
127}