Skip to main content

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
19/// Router agent that directs queries to specialized agents.
20///
21/// Uses an LLM to analyze user queries and determine which
22/// specialized agent is best suited to handle them.
23pub struct RouterAgent {
24    llm: Box<dyn LLMClient>,
25}
26
27impl RouterAgent {
28    /// Creates a new RouterAgent with the given LLM client.
29    pub fn new(llm: Box<dyn LLMClient>) -> Self {
30        Self { llm }
31    }
32
33    /// Parse routing decision from LLM output
34    ///
35    /// This handles various LLM output formats:
36    /// - Clean output: "product"
37    /// - With whitespace: "  product  "
38    /// - With extra text: "I would route this to product"
39    /// - Agent suffix: "product agent"
40    fn parse_routing_decision(output: &str) -> Option<String> {
41        let trimmed = output.trim().to_lowercase();
42
43        // First, try exact match
44        if VALID_AGENTS.contains(&trimmed.as_str()) {
45            return Some(trimmed);
46        }
47
48        // Try to extract valid agent name from output
49        // Split by common delimiters and check each word
50        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        // Check if any valid agent name is contained in the output
58        for agent in VALID_AGENTS {
59            if trimmed.contains(agent) {
60                return Some(agent.to_string());
61            }
62        }
63
64        None
65    }
66
67    /// Routes a query to the appropriate agent type.
68    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        // Parse the response with robust matching
73        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                // Default to orchestrator for complex queries or unrecognized routing
84                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        // Note: RouterAgent.route() is called by the orchestrator/chat handler,
98        // not through the Agent trait execute() method. This is a placeholder.
99        // If called directly, return the router agent name.
100        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}