use super::default::types::{ProjectRequirement, Subtask};
use super::llm::{ChatMessage, LLMProvider, parse_llm_json};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::RwLock;
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AgentInfo {
pub id: String,
pub name: String,
pub description: String,
pub capabilities: Vec<String>,
pub supported_task_types: Vec<String>,
pub prompt_template: Option<String>,
pub current_load: u32,
pub available: bool,
pub performance_score: f32,
pub metadata: HashMap<String, String>,
}
impl AgentInfo {
pub fn new(id: impl Into<String>, name: impl Into<String>) -> Self {
Self {
id: id.into(),
name: name.into(),
description: String::new(),
capabilities: Vec::new(),
supported_task_types: Vec::new(),
prompt_template: None,
current_load: 0,
available: true,
performance_score: 0.8,
metadata: HashMap::new(),
}
}
pub fn with_description(mut self, desc: impl Into<String>) -> Self {
self.description = desc.into();
self
}
pub fn with_capability(mut self, cap: impl Into<String>) -> Self {
self.capabilities.push(cap.into());
self
}
pub fn with_capabilities(mut self, caps: Vec<String>) -> Self {
self.capabilities.extend(caps);
self
}
pub fn with_task_type(mut self, task_type: impl Into<String>) -> Self {
self.supported_task_types.push(task_type.into());
self
}
pub fn with_prompt_template(mut self, template: impl Into<String>) -> Self {
self.prompt_template = Some(template.into());
self
}
pub fn with_performance_score(mut self, score: f32) -> Self {
self.performance_score = score;
self
}
}
#[async_trait::async_trait]
pub trait AgentProvider: Send + Sync {
async fn list_agents(&self) -> Vec<AgentInfo>;
async fn get_agent(&self, agent_id: &str) -> Option<AgentInfo>;
async fn filter_by_capabilities(&self, capabilities: &[String]) -> Vec<AgentInfo> {
let agents = self.list_agents().await;
agents
.into_iter()
.filter(|agent| {
capabilities.is_empty()
|| capabilities
.iter()
.any(|cap| agent.capabilities.contains(cap))
})
.collect()
}
async fn filter_by_task_type(&self, task_type: &str) -> Vec<AgentInfo> {
let agents = self.list_agents().await;
agents
.into_iter()
.filter(|agent| {
agent.supported_task_types.is_empty()
|| agent.supported_task_types.contains(&task_type.to_string())
})
.collect()
}
async fn update_agent_status(&self, agent_id: &str, load: u32, available: bool);
async fn register_agent(&self, _agent: AgentInfo) -> anyhow::Result<()> {
Ok(())
}
async fn unregister_agent(&self, _agent_id: &str) -> anyhow::Result<()> {
Ok(())
}
}
#[derive(Clone)]
pub struct InMemoryAgentProvider {
agents: Arc<RwLock<HashMap<String, AgentInfo>>>,
}
impl InMemoryAgentProvider {
pub fn new() -> Self {
Self {
agents: Arc::new(RwLock::new(HashMap::new())),
}
}
pub async fn add_agent(&self, agent: AgentInfo) {
let mut agents = self.agents.write().await;
agents.insert(agent.id.clone(), agent);
}
pub async fn remove_agent(&self, agent_id: &str) {
let mut agents = self.agents.write().await;
agents.remove(agent_id);
}
}
impl Default for InMemoryAgentProvider {
fn default() -> Self {
Self::new()
}
}
#[async_trait::async_trait]
impl AgentProvider for InMemoryAgentProvider {
async fn list_agents(&self) -> Vec<AgentInfo> {
let agents = self.agents.read().await;
agents.values().filter(|a| a.available).cloned().collect()
}
async fn get_agent(&self, agent_id: &str) -> Option<AgentInfo> {
let agents = self.agents.read().await;
agents.get(agent_id).cloned()
}
async fn update_agent_status(&self, agent_id: &str, load: u32, available: bool) {
let mut agents = self.agents.write().await;
if let Some(agent) = agents.get_mut(agent_id) {
agent.current_load = load;
agent.available = available;
}
}
async fn register_agent(&self, agent: AgentInfo) -> anyhow::Result<()> {
self.add_agent(agent).await;
Ok(())
}
async fn unregister_agent(&self, agent_id: &str) -> anyhow::Result<()> {
self.remove_agent(agent_id).await;
Ok(())
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RoutingContext {
pub subtask: Subtask,
pub todo_id: String,
pub raw_idea: String,
pub requirement: Option<ProjectRequirement>,
pub metadata: HashMap<String, String>,
pub conversation_history: Vec<String>,
}
impl RoutingContext {
pub fn new(subtask: Subtask, todo_id: &str) -> Self {
Self {
subtask,
todo_id: todo_id.to_string(),
raw_idea: String::new(),
requirement: None,
metadata: HashMap::new(),
conversation_history: Vec::new(),
}
}
pub fn with_raw_idea(mut self, idea: &str) -> Self {
self.raw_idea = idea.to_string();
self
}
pub fn with_requirement(mut self, req: ProjectRequirement) -> Self {
self.requirement = Some(req);
self
}
pub fn with_metadata(mut self, key: &str, value: &str) -> Self {
self.metadata.insert(key.to_string(), value.to_string());
self
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RoutingDecision {
pub agent_id: String,
pub reason: String,
pub confidence: f32,
pub alternatives: Vec<String>,
pub decision_type: RoutingDecisionType,
pub needs_human_confirmation: bool,
pub execution_params: HashMap<String, String>,
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub enum RoutingDecisionType {
ExactMatch,
CapabilityMatch,
LLMInference,
RuleMatch,
Default,
HumanAssigned,
}
#[async_trait::async_trait]
pub trait AgentRouter: Send + Sync {
fn name(&self) -> &str;
async fn route(
&self,
context: &RoutingContext,
available_agents: &[AgentInfo],
) -> anyhow::Result<RoutingDecision>;
async fn route_batch(
&self,
contexts: &[RoutingContext],
available_agents: &[AgentInfo],
) -> anyhow::Result<Vec<RoutingDecision>> {
let mut results = Vec::new();
for ctx in contexts {
let decision = self.route(ctx, available_agents).await?;
results.push(decision);
}
Ok(results)
}
async fn validate_decision(
&self,
decision: &RoutingDecision,
available_agents: &[AgentInfo],
) -> bool {
available_agents.iter().any(|a| a.id == decision.agent_id)
}
}
pub struct LLMAgentRouter {
llm: Arc<dyn LLMProvider>,
system_prompt: Option<String>,
routing_prompt_template: Option<String>,
explain_decisions: bool,
confidence_threshold: f32,
}
impl LLMAgentRouter {
pub fn new(llm: Arc<dyn LLMProvider>) -> Self {
Self {
llm,
system_prompt: None,
routing_prompt_template: None,
explain_decisions: true,
confidence_threshold: 0.7,
}
}
pub fn with_system_prompt(mut self, prompt: impl Into<String>) -> Self {
self.system_prompt = Some(prompt.into());
self
}
pub fn with_routing_prompt_template(mut self, template: impl Into<String>) -> Self {
self.routing_prompt_template = Some(template.into());
self
}
pub fn with_confidence_threshold(mut self, threshold: f32) -> Self {
self.confidence_threshold = threshold;
self
}
pub fn with_explain_decisions(mut self, explain: bool) -> Self {
self.explain_decisions = explain;
self
}
fn generate_routing_prompt(&self, context: &RoutingContext, agents: &[AgentInfo]) -> String {
if let Some(ref template) = self.routing_prompt_template {
template
.replace("{subtask_description}", &context.subtask.description)
.replace(
"{required_capabilities}",
&context.subtask.required_capabilities.join(", "),
)
.replace("{raw_idea}", &context.raw_idea)
.replace(
"{agents}",
&agents
.iter()
.map(|a| {
format!(
"- {}: {} (能力: {})",
a.id,
a.description,
a.capabilities.join(", ")
)
})
.collect::<Vec<_>>()
.join("\n"),
)
} else {
format!(
r#"请为以下任务选择最合适的执行Agent。
## 任务信息
- 任务描述: {}
- 所需能力: {}
- 原始需求: {}
## 可用Agent列表
{}
## 要求
请以JSON格式返回路由决策:
```json
{{
"agent_id": "选中的Agent ID",
"reason": "选择理由",
"confidence": 0.0-1.0的置信度,
"alternatives": ["备选Agent ID列表"]
}}
```
请基于任务需求和Agent能力做出最佳匹配。"#,
context.subtask.description,
context.subtask.required_capabilities.join(", "),
context.raw_idea,
agents
.iter()
.map(|a| format!(
"- ID: {}\n 名称: {}\n 描述: {}\n 能力: {}\n 任务类型: {}\n 性能评分: {:.2}",
a.id,
a.name,
a.description,
a.capabilities.join(", "),
a.supported_task_types.join(", "),
a.performance_score
))
.collect::<Vec<_>>()
.join("\n\n")
)
}
}
fn default_system_prompt(&self) -> &'static str {
r#"你是一个任务路由专家,负责将任务分配给最合适的执行Agent。
你需要考虑以下因素:
1. 任务所需的能力与Agent的能力是否匹配
2. Agent的历史性能评分
3. 任务类型与Agent支持的任务类型是否匹配
4. Agent的当前负载情况
做出决策时,请确保:
- 选择能力最匹配的Agent
- 如果有多个匹配的Agent,优先选择性能评分高的
- 如果没有完全匹配的,选择最接近的并降低置信度
- 始终提供清晰的决策理由"#
}
}
#[async_trait::async_trait]
impl AgentRouter for LLMAgentRouter {
fn name(&self) -> &str {
"llm_router"
}
async fn route(
&self,
context: &RoutingContext,
available_agents: &[AgentInfo],
) -> anyhow::Result<RoutingDecision> {
if available_agents.is_empty() {
return Err(anyhow::anyhow!("No available agents"));
}
let messages = vec![
ChatMessage::system(
self.system_prompt
.as_deref()
.unwrap_or_else(|| self.default_system_prompt()),
),
ChatMessage::user(self.generate_routing_prompt(context, available_agents)),
];
let response = self.llm.chat(messages).await?;
#[derive(Deserialize)]
struct LLMRoutingResponse {
agent_id: String,
reason: String,
confidence: Option<f32>,
alternatives: Option<Vec<String>>,
}
match parse_llm_json::<LLMRoutingResponse>(&response) {
Ok(parsed) => {
let confidence = parsed.confidence.unwrap_or(0.8);
Ok(RoutingDecision {
agent_id: parsed.agent_id,
reason: parsed.reason,
confidence,
alternatives: parsed.alternatives.unwrap_or_default(),
decision_type: RoutingDecisionType::LLMInference,
needs_human_confirmation: confidence < self.confidence_threshold,
execution_params: HashMap::new(),
})
}
Err(_) => {
let fallback_agent = available_agents
.first()
.ok_or_else(|| anyhow::anyhow!("No available agents"))?;
Ok(RoutingDecision {
agent_id: fallback_agent.id.clone(),
reason: "LLM响应解析失败,使用默认分配".to_string(),
confidence: 0.5,
alternatives: available_agents
.iter()
.skip(1)
.map(|a| a.id.clone())
.collect(),
decision_type: RoutingDecisionType::Default,
needs_human_confirmation: true,
execution_params: HashMap::new(),
})
}
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RoutingRule {
pub id: String,
pub name: String,
pub priority: i32,
pub conditions: Vec<RuleCondition>,
pub condition_logic: ConditionLogic,
pub target_agent_id: String,
pub enabled: bool,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RuleCondition {
pub field: RuleField,
pub operator: RuleOperator,
pub value: String,
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub enum RuleField {
SubtaskDescription,
RequiredCapability,
TaskType,
Priority,
Metadata(String),
RawIdea,
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub enum RuleOperator {
Equals,
NotEquals,
Contains,
NotContains,
StartsWith,
EndsWith,
Regex,
In,
NotIn,
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub enum ConditionLogic {
All,
Any,
}
pub struct RuleBasedRouter {
rules: Arc<RwLock<Vec<RoutingRule>>>,
default_agent_id: Option<String>,
confirm_on_match: bool,
}
impl RuleBasedRouter {
pub fn new() -> Self {
Self {
rules: Arc::new(RwLock::new(Vec::new())),
default_agent_id: None,
confirm_on_match: false,
}
}
pub fn with_default_agent(mut self, agent_id: impl Into<String>) -> Self {
self.default_agent_id = Some(agent_id.into());
self
}
pub fn with_confirm_on_match(mut self, confirm: bool) -> Self {
self.confirm_on_match = confirm;
self
}
pub async fn add_rule(&self, rule: RoutingRule) {
let mut rules = self.rules.write().await;
rules.push(rule);
rules.sort_by(|a, b| b.priority.cmp(&a.priority));
}
pub async fn remove_rule(&self, rule_id: &str) {
let mut rules = self.rules.write().await;
rules.retain(|r| r.id != rule_id);
}
fn check_condition(&self, condition: &RuleCondition, context: &RoutingContext) -> bool {
let field_value = match &condition.field {
RuleField::SubtaskDescription => context.subtask.description.clone(),
RuleField::RequiredCapability => context.subtask.required_capabilities.join(","),
RuleField::TaskType => context
.metadata
.get("task_type")
.cloned()
.unwrap_or_default(),
RuleField::Priority => context
.metadata
.get("priority")
.cloned()
.unwrap_or_default(),
RuleField::Metadata(key) => context.metadata.get(key).cloned().unwrap_or_default(),
RuleField::RawIdea => context.raw_idea.clone(),
};
match &condition.operator {
RuleOperator::Equals => field_value == condition.value,
RuleOperator::NotEquals => field_value != condition.value,
RuleOperator::Contains => field_value.contains(&condition.value),
RuleOperator::NotContains => !field_value.contains(&condition.value),
RuleOperator::StartsWith => field_value.starts_with(&condition.value),
RuleOperator::EndsWith => field_value.ends_with(&condition.value),
RuleOperator::Regex => regex::Regex::new(&condition.value)
.map(|re| re.is_match(&field_value))
.unwrap_or(false),
RuleOperator::In => condition.value.split(',').any(|v| v.trim() == field_value),
RuleOperator::NotIn => !condition.value.split(',').any(|v| v.trim() == field_value),
}
}
fn check_rule(&self, rule: &RoutingRule, context: &RoutingContext) -> bool {
if !rule.enabled {
return false;
}
match rule.condition_logic {
ConditionLogic::All => rule
.conditions
.iter()
.all(|c| self.check_condition(c, context)),
ConditionLogic::Any => rule
.conditions
.iter()
.any(|c| self.check_condition(c, context)),
}
}
}
impl Default for RuleBasedRouter {
fn default() -> Self {
Self::new()
}
}
#[async_trait::async_trait]
impl AgentRouter for RuleBasedRouter {
fn name(&self) -> &str {
"rule_based_router"
}
async fn route(
&self,
context: &RoutingContext,
available_agents: &[AgentInfo],
) -> anyhow::Result<RoutingDecision> {
let rules = self.rules.read().await;
for rule in rules.iter() {
if self.check_rule(rule, context) {
if available_agents
.iter()
.any(|a| a.id == rule.target_agent_id)
{
return Ok(RoutingDecision {
agent_id: rule.target_agent_id.clone(),
reason: format!("匹配规则: {} ({})", rule.name, rule.id),
confidence: 1.0,
alternatives: Vec::new(),
decision_type: RoutingDecisionType::RuleMatch,
needs_human_confirmation: self.confirm_on_match,
execution_params: HashMap::new(),
});
}
}
}
let agent_id = self
.default_agent_id
.clone()
.or_else(|| available_agents.first().map(|a| a.id.clone()))
.ok_or_else(|| anyhow::anyhow!("No available agents"))?;
Ok(RoutingDecision {
agent_id,
reason: "无匹配规则,使用默认分配".to_string(),
confidence: 0.5,
alternatives: available_agents.iter().map(|a| a.id.clone()).collect(),
decision_type: RoutingDecisionType::Default,
needs_human_confirmation: true,
execution_params: HashMap::new(),
})
}
}
pub struct CapabilityRouter {
capability_weights: HashMap<String, f32>,
load_balancing: bool,
performance_weight: f32,
}
impl CapabilityRouter {
pub fn new() -> Self {
Self {
capability_weights: HashMap::new(),
load_balancing: true,
performance_weight: 0.3,
}
}
pub fn with_capability_weight(mut self, capability: impl Into<String>, weight: f32) -> Self {
self.capability_weights.insert(capability.into(), weight);
self
}
pub fn with_load_balancing(mut self, enabled: bool) -> Self {
self.load_balancing = enabled;
self
}
pub fn with_performance_weight(mut self, weight: f32) -> Self {
self.performance_weight = weight;
self
}
fn calculate_match_score(&self, agent: &AgentInfo, required_caps: &[String]) -> f32 {
if required_caps.is_empty() {
return 1.0;
}
let mut score = 0.0;
let mut total_weight = 0.0;
for cap in required_caps {
let weight = self.capability_weights.get(cap).copied().unwrap_or(1.0);
total_weight += weight;
if agent.capabilities.contains(cap) {
score += weight;
}
}
let capability_score = if total_weight > 0.0 {
score / total_weight
} else {
1.0
};
let load_score = if self.load_balancing {
1.0 - (agent.current_load as f32 / 100.0)
} else {
1.0
};
let performance_score = agent.performance_score;
capability_score * 0.5 + load_score * 0.2 + performance_score * self.performance_weight
}
}
impl Default for CapabilityRouter {
fn default() -> Self {
Self::new()
}
}
#[async_trait::async_trait]
impl AgentRouter for CapabilityRouter {
fn name(&self) -> &str {
"capability_router"
}
async fn route(
&self,
context: &RoutingContext,
available_agents: &[AgentInfo],
) -> anyhow::Result<RoutingDecision> {
if available_agents.is_empty() {
return Err(anyhow::anyhow!("No available agents"));
}
let required_caps = &context.subtask.required_capabilities;
let mut scored_agents: Vec<(&AgentInfo, f32)> = available_agents
.iter()
.map(|a| (a, self.calculate_match_score(a, required_caps)))
.collect();
scored_agents.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
let (best_agent, best_score) = scored_agents[0];
Ok(RoutingDecision {
agent_id: best_agent.id.clone(),
reason: format!(
"能力匹配得分最高 ({:.2}),匹配能力: {:?}",
best_score,
required_caps
.iter()
.filter(|c| best_agent.capabilities.contains(c))
.collect::<Vec<_>>()
),
confidence: best_score,
alternatives: scored_agents
.iter()
.skip(1)
.map(|(a, _)| a.id.clone())
.collect(),
decision_type: RoutingDecisionType::CapabilityMatch,
needs_human_confirmation: best_score < 0.5,
execution_params: HashMap::new(),
})
}
}
pub struct CompositeRouter {
routers: Vec<Arc<dyn AgentRouter>>,
fallback_router: Option<Arc<dyn AgentRouter>>,
}
impl CompositeRouter {
pub fn new() -> Self {
Self {
routers: Vec::new(),
fallback_router: None,
}
}
pub fn add_router(mut self, router: Arc<dyn AgentRouter>) -> Self {
self.routers.push(router);
self
}
pub fn with_fallback(mut self, router: Arc<dyn AgentRouter>) -> Self {
self.fallback_router = Some(router);
self
}
}
impl Default for CompositeRouter {
fn default() -> Self {
Self::new()
}
}
#[async_trait::async_trait]
impl AgentRouter for CompositeRouter {
fn name(&self) -> &str {
"composite_router"
}
async fn route(
&self,
context: &RoutingContext,
available_agents: &[AgentInfo],
) -> anyhow::Result<RoutingDecision> {
for router in &self.routers {
match router.route(context, available_agents).await {
Ok(decision) if decision.confidence >= 0.5 => {
return Ok(decision);
}
_ => continue,
}
}
if let Some(ref fallback) = self.fallback_router {
return fallback.route(context, available_agents).await;
}
let agent = available_agents
.first()
.ok_or_else(|| anyhow::anyhow!("No available agents"))?;
Ok(RoutingDecision {
agent_id: agent.id.clone(),
reason: "所有路由器均无高置信度匹配,使用默认分配".to_string(),
confidence: 0.3,
alternatives: available_agents
.iter()
.skip(1)
.map(|a| a.id.clone())
.collect(),
decision_type: RoutingDecisionType::Default,
needs_human_confirmation: true,
execution_params: HashMap::new(),
})
}
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn test_in_memory_agent_provider() {
let provider = InMemoryAgentProvider::new();
provider
.add_agent(
AgentInfo::new("agent_1", "Test Agent 1")
.with_capability("backend")
.with_performance_score(0.9),
)
.await;
provider
.add_agent(
AgentInfo::new("agent_2", "Test Agent 2")
.with_capability("frontend")
.with_performance_score(0.85),
)
.await;
let agents = provider.list_agents().await;
assert_eq!(agents.len(), 2);
let backend_agents = provider
.filter_by_capabilities(&["backend".to_string()])
.await;
assert_eq!(backend_agents.len(), 1);
assert_eq!(backend_agents[0].id, "agent_1");
}
#[tokio::test]
async fn test_capability_router() {
let router = CapabilityRouter::new()
.with_load_balancing(true)
.with_performance_weight(0.3);
let agents = vec![
AgentInfo::new("agent_1", "Backend Agent")
.with_capability("backend")
.with_capability("database")
.with_performance_score(0.9),
AgentInfo::new("agent_2", "Frontend Agent")
.with_capability("frontend")
.with_capability("ui")
.with_performance_score(0.85),
];
let context = RoutingContext::new(
Subtask {
id: "task_1".to_string(),
description: "Build API".to_string(),
required_capabilities: vec!["backend".to_string()],
order: 1,
depends_on: Vec::new(),
},
"todo_1",
);
let decision = router.route(&context, &agents).await.unwrap();
assert_eq!(decision.agent_id, "agent_1");
assert_eq!(decision.decision_type, RoutingDecisionType::CapabilityMatch);
}
#[tokio::test]
async fn test_rule_based_router() {
let router = RuleBasedRouter::new();
router
.add_rule(RoutingRule {
id: "rule_1".to_string(),
name: "Backend Rule".to_string(),
priority: 10,
conditions: vec![RuleCondition {
field: RuleField::RequiredCapability,
operator: RuleOperator::Contains,
value: "backend".to_string(),
}],
condition_logic: ConditionLogic::All,
target_agent_id: "backend_agent".to_string(),
enabled: true,
})
.await;
let agents = vec![
AgentInfo::new("backend_agent", "Backend Agent"),
AgentInfo::new("frontend_agent", "Frontend Agent"),
];
let context = RoutingContext::new(
Subtask {
id: "task_1".to_string(),
description: "Build API".to_string(),
required_capabilities: vec!["backend".to_string()],
order: 1,
depends_on: Vec::new(),
},
"todo_1",
);
let decision = router.route(&context, &agents).await.unwrap();
assert_eq!(decision.agent_id, "backend_agent");
assert_eq!(decision.decision_type, RoutingDecisionType::RuleMatch);
}
#[tokio::test]
async fn test_composite_router() {
let rule_router = Arc::new(RuleBasedRouter::new());
let capability_router = Arc::new(CapabilityRouter::new());
let composite = CompositeRouter::new()
.add_router(rule_router)
.with_fallback(capability_router);
let agents = vec![
AgentInfo::new("agent_1", "Agent 1")
.with_capability("general")
.with_performance_score(0.8),
];
let context = RoutingContext::new(
Subtask {
id: "task_1".to_string(),
description: "General task".to_string(),
required_capabilities: vec![],
order: 1,
depends_on: Vec::new(),
},
"todo_1",
);
let decision = composite.route(&context, &agents).await.unwrap();
assert_eq!(decision.agent_id, "agent_1");
}
}