Alex Hancock
332 lines
13 KiB
Rust
332 lines
13 KiB
Rust
use crate::agents::subagent_task_config::DEFAULT_SUBAGENT_MAX_TURNS;
|
|
use crate::{
|
|
agents::extension::ExtensionConfig,
|
|
agents::{extension_manager::ExtensionManager, Agent, TaskConfig},
|
|
config::ExtensionConfigManager,
|
|
prompt_template::render_global_file,
|
|
providers::errors::ProviderError,
|
|
};
|
|
use anyhow::anyhow;
|
|
use chrono::{DateTime, Utc};
|
|
use rmcp::model::Tool;
|
|
use rmcp::model::{ErrorCode, ErrorData};
|
|
use serde::{Deserialize, Serialize};
|
|
// use serde_json::{self};
|
|
use crate::conversation::message::{Message, MessageContent, ToolRequest};
|
|
use crate::conversation::Conversation;
|
|
use std::{collections::HashMap, sync::Arc};
|
|
use tokio::sync::{Mutex, RwLock};
|
|
use tokio_util::sync::CancellationToken;
|
|
use tracing::{debug, error, instrument};
|
|
|
|
/// Status of a subagent
|
|
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
|
pub enum SubAgentStatus {
|
|
Ready, // Ready to process messages
|
|
Processing, // Currently working on a task
|
|
Completed(String), // Task completed (with optional message for success/error)
|
|
Terminated, // Manually terminated
|
|
}
|
|
|
|
/// Progress information for a subagent
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct SubAgentProgress {
|
|
pub subagent_id: String,
|
|
pub status: SubAgentStatus,
|
|
pub message: String,
|
|
pub turn: usize,
|
|
pub max_turns: Option<usize>,
|
|
pub timestamp: DateTime<Utc>,
|
|
}
|
|
|
|
/// A specialized agent that can handle specific tasks independently
|
|
pub struct SubAgent {
|
|
pub id: String,
|
|
pub conversation: Arc<Mutex<Conversation>>,
|
|
pub status: Arc<RwLock<SubAgentStatus>>,
|
|
pub config: TaskConfig,
|
|
pub turn_count: Arc<Mutex<usize>>,
|
|
pub created_at: DateTime<Utc>,
|
|
pub extension_manager: Arc<RwLock<ExtensionManager>>,
|
|
}
|
|
|
|
impl SubAgent {
|
|
/// Create a new subagent with the given configuration and provider
|
|
#[instrument(skip(task_config))]
|
|
pub async fn new(task_config: TaskConfig) -> Result<Arc<Self>, anyhow::Error> {
|
|
debug!("Creating new subagent with id: {}", task_config.id);
|
|
|
|
// Create a new extension manager for this subagent
|
|
let mut extension_manager = ExtensionManager::new();
|
|
|
|
// Add extensions based on task_type:
|
|
// 1. If executing dynamic task (task_type = 'text_instruction'), default to using all enabled extensions
|
|
// 2. (TODO) If executing a sub-recipe task, only use recipe extensions
|
|
|
|
// Get all enabled extensions from config
|
|
let enabled_extensions = ExtensionConfigManager::get_all()
|
|
.unwrap_or_default()
|
|
.into_iter()
|
|
.filter(|ext| ext.enabled)
|
|
.map(|ext| ext.config)
|
|
.collect::<Vec<ExtensionConfig>>();
|
|
|
|
// Add enabled extensions to the subagent's extension manager
|
|
for extension in enabled_extensions {
|
|
if let Err(e) = extension_manager.add_extension(extension).await {
|
|
debug!("Failed to add extension to subagent: {}", e);
|
|
// Continue with other extensions even if one fails
|
|
}
|
|
}
|
|
|
|
let subagent = Arc::new(SubAgent {
|
|
id: task_config.id.clone(),
|
|
conversation: Arc::new(Mutex::new(Conversation::new_unvalidated(Vec::new()))),
|
|
status: Arc::new(RwLock::new(SubAgentStatus::Ready)),
|
|
config: task_config,
|
|
turn_count: Arc::new(Mutex::new(0)),
|
|
created_at: Utc::now(),
|
|
extension_manager: Arc::new(RwLock::new(extension_manager)),
|
|
});
|
|
|
|
debug!("Subagent {} created successfully", subagent.id);
|
|
Ok(subagent)
|
|
}
|
|
|
|
/// Update the status of the subagent
|
|
async fn set_status(&self, status: SubAgentStatus) {
|
|
// Update the status first, then release the lock
|
|
{
|
|
let mut current_status = self.status.write().await;
|
|
*current_status = status.clone();
|
|
} // Write lock is released here!
|
|
}
|
|
|
|
/// Process a message and generate a response using the subagent's provider
|
|
#[instrument(skip(self, message))]
|
|
pub async fn reply_subagent(
|
|
&self,
|
|
message: String,
|
|
task_config: TaskConfig,
|
|
) -> Result<Conversation, anyhow::Error> {
|
|
debug!("Processing message for subagent {}", self.id);
|
|
|
|
// Get provider from task config
|
|
let provider = self
|
|
.config
|
|
.provider
|
|
.as_ref()
|
|
.ok_or_else(|| anyhow!("No provider configured for subagent"))?;
|
|
|
|
// Set status to processing
|
|
self.set_status(SubAgentStatus::Processing).await;
|
|
|
|
// Add user message to conversation
|
|
let user_message = Message::user().with_text(message.clone());
|
|
{
|
|
let mut conversation = self.conversation.lock().await;
|
|
conversation.push(user_message.clone());
|
|
}
|
|
|
|
// Get the current conversation for context
|
|
let mut messages = {
|
|
let conversation = self.conversation.lock().await;
|
|
conversation.clone()
|
|
};
|
|
|
|
// Get tools from the subagent's own extension manager
|
|
let tools: Vec<Tool> = self
|
|
.extension_manager
|
|
.read()
|
|
.await
|
|
.get_prefixed_tools(None)
|
|
.await
|
|
.unwrap_or_default();
|
|
|
|
let toolshim_tools: Vec<Tool> = vec![];
|
|
|
|
// Build system prompt using the template
|
|
let system_prompt = self.build_system_prompt(&tools).await?;
|
|
|
|
// Generate response from provider with loop for tool processing (max_turns iterations)
|
|
let mut loop_count = 0;
|
|
let max_turns = self.config.max_turns.unwrap_or(DEFAULT_SUBAGENT_MAX_TURNS);
|
|
let mut last_error: Option<anyhow::Error> = None;
|
|
|
|
// Generate response from provider
|
|
loop {
|
|
loop_count += 1;
|
|
|
|
match Agent::generate_response_from_provider(
|
|
Arc::clone(provider),
|
|
&system_prompt,
|
|
messages.messages(),
|
|
&tools,
|
|
&toolshim_tools,
|
|
)
|
|
.await
|
|
{
|
|
Ok((response, _usage)) => {
|
|
// Process any tool calls in the response
|
|
let tool_requests: Vec<ToolRequest> = response
|
|
.content
|
|
.iter()
|
|
.filter_map(|content| {
|
|
if let MessageContent::ToolRequest(req) = content {
|
|
Some(req.clone())
|
|
} else {
|
|
None
|
|
}
|
|
})
|
|
.collect();
|
|
|
|
// If there are no tool requests, we're done
|
|
if tool_requests.is_empty() || loop_count >= max_turns {
|
|
self.add_message(response.clone()).await;
|
|
messages.push(response.clone());
|
|
|
|
// Set status back to ready
|
|
self.set_status(SubAgentStatus::Completed("Completed!".to_string()))
|
|
.await;
|
|
break;
|
|
}
|
|
|
|
// Add the assistant message with tool calls to the conversation
|
|
messages.push(response.clone());
|
|
|
|
// Process each tool request and create user response messages
|
|
for request in &tool_requests {
|
|
if let Ok(tool_call) = &request.tool_call {
|
|
// Handle platform tools or dispatch to extension manager
|
|
let tool_result = match self
|
|
.extension_manager
|
|
.read()
|
|
.await
|
|
.dispatch_tool_call(tool_call.clone(), CancellationToken::default())
|
|
.await
|
|
{
|
|
Ok(result) => result.result.await,
|
|
Err(e) => Err(ErrorData::new(
|
|
ErrorCode::INTERNAL_ERROR,
|
|
e.to_string(),
|
|
None,
|
|
)),
|
|
};
|
|
|
|
match tool_result {
|
|
Ok(result) => {
|
|
// Create a user message with the tool response
|
|
let tool_response_message = Message::user()
|
|
.with_tool_response(request.id.clone(), Ok(result.clone()));
|
|
messages.push(tool_response_message);
|
|
}
|
|
Err(e) => {
|
|
// Create a user message with the tool error
|
|
let tool_error_message = Message::user().with_tool_response(
|
|
request.id.clone(),
|
|
Err(ErrorData::new(
|
|
ErrorCode::INTERNAL_ERROR,
|
|
e.to_string(),
|
|
None,
|
|
)),
|
|
);
|
|
messages.push(tool_error_message);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// Continue the loop to get the next response from the provider
|
|
}
|
|
Err(ProviderError::ContextLengthExceeded(_)) => {
|
|
self.set_status(SubAgentStatus::Completed(
|
|
"Context length exceeded".to_string(),
|
|
))
|
|
.await;
|
|
last_error = Some(anyhow::anyhow!("Context length exceeded"));
|
|
break;
|
|
}
|
|
Err(ProviderError::RateLimitExceeded(_)) => {
|
|
self.set_status(SubAgentStatus::Completed("Rate limit exceeded".to_string()))
|
|
.await;
|
|
last_error = Some(anyhow::anyhow!("Rate limit exceeded"));
|
|
break;
|
|
}
|
|
Err(e) => {
|
|
self.set_status(SubAgentStatus::Completed(format!("Error: {}", e)))
|
|
.await;
|
|
error!("Error: {}", e);
|
|
last_error = Some(anyhow::anyhow!("Provider error: {}", e));
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
|
|
// Handle error cases or return the last message
|
|
if let Some(error) = last_error {
|
|
Err(error)
|
|
} else {
|
|
Ok(messages)
|
|
}
|
|
}
|
|
|
|
/// Add a message to the conversation (for tracking agent responses)
|
|
async fn add_message(&self, message: Message) {
|
|
let mut conversation = self.conversation.lock().await;
|
|
conversation.push(message);
|
|
}
|
|
|
|
/// Build the system prompt for the subagent using the template
|
|
async fn build_system_prompt(&self, available_tools: &[Tool]) -> Result<String, anyhow::Error> {
|
|
let mut context = HashMap::new();
|
|
|
|
// Add basic context
|
|
context.insert(
|
|
"current_date_time",
|
|
serde_json::Value::String(Utc::now().format("%Y-%m-%d %H:%M:%S UTC").to_string()),
|
|
);
|
|
context.insert("subagent_id", serde_json::Value::String(self.id.clone()));
|
|
|
|
// Add max turns if configured
|
|
if let Some(max_turns) = self.config.max_turns {
|
|
context.insert(
|
|
"max_turns",
|
|
serde_json::Value::Number(serde_json::Number::from(max_turns)),
|
|
);
|
|
}
|
|
|
|
// Add available tools with descriptions for better context
|
|
let tools_with_descriptions: Vec<String> = available_tools
|
|
.iter()
|
|
.map(|t| {
|
|
if let Some(description) = &t.description {
|
|
format!("{}: {}", t.name, description)
|
|
} else {
|
|
t.name.to_string()
|
|
}
|
|
})
|
|
.collect();
|
|
|
|
context.insert(
|
|
"available_tools",
|
|
serde_json::Value::String(if tools_with_descriptions.is_empty() {
|
|
"None".to_string()
|
|
} else {
|
|
tools_with_descriptions.join(", ")
|
|
}),
|
|
);
|
|
|
|
// Add tool count for context
|
|
context.insert(
|
|
"tool_count",
|
|
serde_json::Value::Number(serde_json::Number::from(available_tools.len())),
|
|
);
|
|
|
|
// Render the subagent system prompt template
|
|
let system_prompt = render_global_file("subagent_system.md", &context)
|
|
.map_err(|e| anyhow!("Failed to render subagent system prompt: {}", e))?;
|
|
|
|
Ok(system_prompt)
|
|
}
|
|
}
|