ares/agents/
router.rs

1use crate::{
2    agents::Agent,
3    llm::LLMClient,
4    types::{AgentContext, AgentType, Result},
5};
6use async_trait::async_trait;
7
8/// Valid agent names for routing
9const VALID_AGENTS: &[&str] = &[
10    "product",
11    "invoice",
12    "sales",
13    "finance",
14    "hr",
15    "orchestrator",
16    "research",
17];
18
19pub struct RouterAgent {
20    llm: Box<dyn LLMClient>,
21}
22
23impl RouterAgent {
24    pub fn new(llm: Box<dyn LLMClient>) -> Self {
25        Self { llm }
26    }
27
28    /// Parse routing decision from LLM output
29    ///
30    /// This handles various LLM output formats:
31    /// - Clean output: "product"
32    /// - With whitespace: "  product  "
33    /// - With extra text: "I would route this to product"
34    /// - Agent suffix: "product agent"
35    fn parse_routing_decision(output: &str) -> Option<String> {
36        let trimmed = output.trim().to_lowercase();
37
38        // First, try exact match
39        if VALID_AGENTS.contains(&trimmed.as_str()) {
40            return Some(trimmed);
41        }
42
43        // Try to extract valid agent name from output
44        // Split by common delimiters and check each word
45        for word in trimmed.split(|c: char| c.is_whitespace() || c == ':' || c == ',' || c == '.') {
46            let word = word.trim();
47            if VALID_AGENTS.contains(&word) {
48                return Some(word.to_string());
49            }
50        }
51
52        // Check if any valid agent name is contained in the output
53        for agent in VALID_AGENTS {
54            if trimmed.contains(agent) {
55                return Some(agent.to_string());
56            }
57        }
58
59        None
60    }
61
62    // Route the query to the appropriate agent
63    pub async fn route(&self, query: &str, _context: &AgentContext) -> Result<AgentType> {
64        let system_prompt = self.system_prompt();
65        let response = self.llm.generate_with_system(&system_prompt, query).await?;
66
67        // Parse the response with robust matching
68        let agent_name = Self::parse_routing_decision(&response);
69
70        match agent_name.as_deref() {
71            Some("product") => Ok(AgentType::Product),
72            Some("invoice") => Ok(AgentType::Invoice),
73            Some("sales") => Ok(AgentType::Sales),
74            Some("finance") => Ok(AgentType::Finance),
75            Some("hr") => Ok(AgentType::HR),
76            Some("orchestrator") | Some("research") => Ok(AgentType::Orchestrator),
77            _ => {
78                // Default to orchestrator for complex queries or unrecognized routing
79                tracing::debug!(
80                    "Router could not parse output '{}', defaulting to orchestrator",
81                    response
82                );
83                Ok(AgentType::Orchestrator)
84            }
85        }
86    }
87}
88
89#[async_trait]
90impl Agent for RouterAgent {
91    async fn execute(&self, input: &str, context: &AgentContext) -> Result<String> {
92        let agent_type = self.route(input, context).await?;
93        // Return lowercase agent name for workflow engine compatibility
94        let agent_name = match agent_type {
95            AgentType::Router => "router",
96            AgentType::Orchestrator => "orchestrator",
97            AgentType::Product => "product",
98            AgentType::Invoice => "invoice",
99            AgentType::Sales => "sales",
100            AgentType::Finance => "finance",
101            AgentType::HR => "hr",
102        };
103        Ok(agent_name.to_string())
104    }
105
106    fn system_prompt(&self) -> String {
107        r#"You are a routing agent that classifies user queries and routes them to the appropriate specialized agent.
108
109Available agents:
110- product: Product information, recommendations, catalog queries
111- invoice: Invoice processing, billing questions, payment status
112- sales: Sales data, analytics, performance metrics
113- finance: Financial reports, budgets, expense analysis
114- hr: Human resources, employee information, policies
115- orchestrator: Complex queries requiring multiple agents or research
116
117Analyze the user's query and respond with ONLY the agent name (lowercase, one word).
118Examples:
119- "What products do we have?" → product
120- "Show me last quarter's sales" → sales
121- "What's our hiring policy?" → hr
122- "Create a comprehensive market analysis" → orchestrator
123
124Respond with ONLY the agent name, nothing else."#.to_string()
125    }
126
127    fn agent_type(&self) -> AgentType {
128        AgentType::Router
129    }
130}
ares/agents/router.rs

ares/agents/
router.rs
