Review/Advices of an autonomous multi-agent design

[vincent-herlemont]

I've been experimenting an autonomous multi-agent system with rig. I'd love to get your thoughts on whether this approach is appropriate or if there's a more idiomatic way to design such a system with rig.

Experiment: The multi-agent system successfully coordinated between agents Aurelie, Paul, and Julie to locate the keys, ultimately finding them under the entry of the house after recursive task delegation and handling an error in task extraction.

Result:

Question: Hi Aurelie, where are the keys?
 Coordinated prompt: Hi Aurelie, can you tell me where the keys are?
 new tasks:
  agent: "Aurelie", task: "Provide information on the location of the keys."
 Delegating to agent "Aurelie": "Provide information on the location of the keys."
 Response from agent Aurelie: I'm not sure where the keys are located, but I recommend asking Paul first, as he might have the information you need. If Paul is unable to help, then Julie will definitely know where they are.
 Generated new task from response: Hi Paul, do you know where the keys are located? Hi Julie, if Paul doesn't know, do you know where the keys are located?
 new tasks:
  agent: "Paul", task: "Inform if you know the location of the keys."
  agent: "Julie", task: "Inform if you know the location of the keys, in case Paul doesn't know."
 Delegating to agent "Paul": "Inform if you know the location of the keys."
 Response from agent Paul: I don't have the information about the location of the keys. Perhaps Julie might know where they are. Would you like me to check with her?
 Generated new task from response: Hi Julie, do you know where the keys are located?
 new tasks:
  agent: "Julie", task: "Find out and provide the location of the keys."
 Delegating to agent "Julie": "Find out and provide the location of the keys."
 Response from agent Julie: The keys are located under the entry of the house.
Multi-Agent System: Hey there! It looks like the keys are under the entry of the house. Hope that helps!

Key Design Points

1. Structured Data Flow: Using extractors to parse agent responses into structured data (AgentResponse, AgentTasks)
2. Recursive Processing: Allowing agents to generate new tasks that are processed recursively
3. Clear Agent Roles: Each agent has a specific role and knowledge domain
4. Coordination Layer: Using coordinator and task generator agents to manage the flow

// run.rs
use anyhow::Result;
use rig::completion::Chat;
use rig::providers::{self};

use super::MultiAgentSystem;

pub async fn run() -> Result<()> {
    // Create OpenAI client
    let openai_client = providers::openai::Client::from_env();
    let completion_model = openai_client.completion_model(providers::openai::GPT_4O);

    let multi_agent_system = MultiAgentSystem::new(completion_model, openai_client);

    let question = "Hi Aurelie, where are the keys?";
    println!("Question: {}", question);
    println!(
        "Multi-Agent System: {}",
        multi_agent_system.chat(question, vec![]).await?
    );

    Ok(())
}

// model.rs
use schemars::JsonSchema;
use serde::{Deserialize, Serialize};

#[derive(Debug, Deserialize, JsonSchema, Serialize, PartialEq, Eq, Hash)]
pub enum AgentRole {
    Paul,
    Julie,
    Aurelie,
    William,
}

#[derive(Debug, Deserialize, JsonSchema, Serialize)]
/// # AgentTask
/// A task assigned to an agent in the multi-agent system
pub struct AgentTask {
    /// The agent assigned to this task
    pub agent: AgentRole,
    /// The task instructions for the agent
    pub task: String,
}

#[derive(Debug, Deserialize, JsonSchema, Serialize)]
/// # AgentTasks
///
/// Tasks extracted from the coordinator's prompt
pub struct AgentTasks {
    /// The tasks to be distributed among agents
    pub tasks: Vec<AgentTask>,
}

#[derive(Debug, Deserialize, JsonSchema, Serialize)]
/// # AgentResponse
///
/// The result of the multi-agent investigation
pub struct AgentResponse {
    /// # The location of the keys.
    /// If the keys are not found leave it empty.
    pub keys_location: Option<String>,
}

// investigation.rs
use anyhow::Result;
use rig::{
    agent::Agent,
    completion::{Chat, Message, PromptError},
    extractor::Extractor,
    providers::{self, openai},
};
use std::collections::HashMap;

use super::{
    agent_config::{
        create_coordinator_agent, create_response_formatter, create_specialized_agents,
        create_task_generator,
    },
    models::{AgentResponse, AgentRole, AgentTasks},
};

pub struct MultiAgentSystem {
    coordinator: Agent<openai::CompletionModel>,
    task_generator: Agent<openai::CompletionModel>,
    specialized_agents: HashMap<AgentRole, Agent<openai::CompletionModel>>,
    task_extractor: Extractor<openai::CompletionModel, AgentTasks>,
    response_extractor: Extractor<openai::CompletionModel, AgentResponse>,
    response_formatter: Agent<openai::CompletionModel>,
}

impl MultiAgentSystem {
    pub fn new(model: openai::CompletionModel, client: openai::Client) -> Self {
        let specialized_agents = create_specialized_agents(model.clone())
            .into_iter()
            .map(|(k, v)| (k, v.build()))
            .collect();

        Self {
            coordinator: create_coordinator_agent(model.clone()).build(),
            task_generator: create_task_generator(model.clone()).build(),
            specialized_agents,
            task_extractor: client
                .extractor::<AgentTasks>(providers::openai::GPT_4O)
                .build(),
            response_extractor: client
                .extractor::<AgentResponse>(providers::openai::GPT_4O)
                .build(),
            response_formatter: create_response_formatter(model).build(),
        }
    }

    async fn process_tasks(
        &self,
        tasks: AgentTasks,
        chat_history: Vec<Message>,
    ) -> Result<Option<String>, PromptError> {
        println!(" new tasks:");
        for task in &tasks.tasks {
            println!("  agent: \"{:?}\", task: \"{}\"", task.agent, task.task);
        }
        for task in tasks.tasks {
            println!(
                " Delegating to agent \"{:?}\": \"{}\"",
                task.agent, task.task
            );
            let response = self.specialized_agents[&task.agent]
                .chat(&task.task, chat_history.clone())
                .await?;
            println!(" Response from agent {:?}: {}", task.agent, response);
            let result = match self.response_extractor.extract(&response).await {
                Ok(result) => result,
                Err(e) => {
                    println!(" Error extracting response: {}", e);
                    continue;
                }
            };
            if let Some(location) = result.keys_location {
                return Ok(Some(location));
            }
            let new_task = self
                .task_generator
                .chat(&response, chat_history.clone())
                .await?;
            println!(" Generated new task from response: {}", new_task);
            let new_tasks = match self.task_extractor.extract(&new_task).await {
                Ok(tasks) => tasks,
                Err(e) => {
                    println!(" Error extracting new tasks: {}", e);
                    continue;
                }
            };
            // Use Box::pin for the recursive async call
            if let Ok(Some(location)) =
                Box::pin(self.process_tasks(new_tasks, chat_history.clone())).await
            {
                return Ok(Some(location));
            }
        }
        Ok(None)
    }
}

impl Chat for MultiAgentSystem {
    async fn chat(&self, prompt: &str, chat_history: Vec<Message>) -> Result<String, PromptError> {
        let coordinated_prompt = self.coordinator.chat(prompt, chat_history.clone()).await?;
        println!(" Coordinated prompt: {}", coordinated_prompt);

        let initial_tasks = match self.task_extractor.extract(&coordinated_prompt).await {
            Ok(tasks) => tasks,
            Err(e) => return Ok(format!("Error extracting tasks: {}", e)),
        };

        let keys_location = match self
            .process_tasks(initial_tasks, chat_history.clone())
            .await
        {
            Ok(location) => location,
            Err(e) => return Ok(format!("Error processing tasks: {}", e)),
        };

        // Format the final response using the response_formatter
        let format_context = format!(
            "Initial question: {}\nCoordinated question: {}\nKeys location: {}",
            prompt,
            coordinated_prompt,
            keys_location.unwrap_or_else(|| "not found".to_string())
        );
        self.response_formatter
            .chat(&format_context, chat_history)
            .await
    }
}

// agent_config.rs
use rig::{agent::AgentBuilder, providers::openai};
use std::collections::HashMap;

use crate::experiments::e4_multi_agents::models::AgentRole;

pub fn create_coordinator_agent(
    model: openai::CompletionModel,
) -> AgentBuilder<openai::CompletionModel> {
    AgentBuilder::new(model)
        .preamble(
            "You are a coordinator agent who rephrases the user's question or request in a most precise way without adding any additional information.
            Explicitly Repeat Actions: Always state each action explicitly for every subject. Avoid combining multiple subjects under a single verb unless they share the action simultaneously in the same context.
            Preserve Context: Retain all relevant contextual elements, including salutations, names, or specific references, unless explicitly instructed to remove them.",
        )
}

pub fn create_task_generator(
    model: openai::CompletionModel,
) -> AgentBuilder<openai::CompletionModel> {
    AgentBuilder::new(model)
        .preamble(
            "You are a task generator agent who converts input into a direct question addressed to the relevant agent(s).
Explicitly Repeat Actions: Always state each action explicitly for every agent. Avoid combining multiple agents under a single verb unless they share the action simultaneously in the same context.
Direct Questions: Transform input statements into direct, first-person questions directed at the appropriate agent. The output should take the form of: 'Hi [Name], [question].' Ensure the question is concise, direct, and conversational.",
        )
}

pub fn create_response_formatter(
    model: openai::CompletionModel,
) -> AgentBuilder<openai::CompletionModel> {
    AgentBuilder::new(model).preamble(
        "You are a response formatter agent who formats the final response based on the multi-agent investigation results.
When the keys are found, respond with a clear and concise message stating their location.
When the keys are not found, respond with a message indicating that the search was unsuccessful.
Always maintain a natural, conversational tone while being direct and informative.",
    )
}

pub fn create_specialized_agents(
    model: openai::CompletionModel,
) -> HashMap<AgentRole, AgentBuilder<openai::CompletionModel>> {
    let mut agents = HashMap::new();

    agents.insert(
        AgentRole::Paul,
        AgentBuilder::new(model.clone()).preamble(
            "You are Paul, a specialized agent in the multi-agent system.
You don't know where the keys are, but maybe Julie knows where they are.",
        ),
    );

    agents.insert(
        AgentRole::Julie,
        AgentBuilder::new(model.clone()).preamble(
            "You are Julie, a specialized agent in the multi-agent system. 
You know where the keys are: they are located under the entry of the house.",
        ),
    );

    agents.insert(
        AgentRole::Aurelie,
        AgentBuilder::new(model.clone())
            .preamble(
                "You are Aurelie, a specialized agent in the multi-agent system.
You don't know where the keys are, but you know that Paul might know, and if he doesn't know then Julie will definitely know.",
            ),
    );

    agents.insert(
        AgentRole::William,
        AgentBuilder::new(model.clone())
            .preamble(
                "You are William, a specialized agent in the multi-agent system.
You don't know where the keys are, but you know that Julie might know, and if she doesn't know then Paul will definitely know.",
            ),
    );

    agents
}

Questions:

1. Is using extractors for structured handoffs between agents a good approach? Are there better alternatives?
2. How could this design be improved to handle more complex scenarios or different types of agent interactions?
3. Are there any potential issues with the current implementation that could cause problems at scale?

[mateobelanger]

@cvauclair @0xMochan