Improve orchestration agent model selection
This commit is contained in:
@@ -15,3 +15,4 @@ pub(crate) use provider_run_coordinator::{
|
||||
pub(crate) use rig::{
|
||||
prepare_provider_run, provider_runtime_for_request, PreparedProviderRun, ProviderActionContext,
|
||||
};
|
||||
pub(crate) use rig_request::OrchestrationModelOption;
|
||||
|
||||
@@ -13,8 +13,9 @@ use uuid::Uuid;
|
||||
use warp_multi_agent_api::ToolType;
|
||||
|
||||
use super::rig_request::{
|
||||
prepare_bedrock_rig_turn_for_mode, prepare_rig_turn, prepare_rig_turn_for_mode, MCPToolTarget,
|
||||
PreparedRigTurn, RigRequestMode,
|
||||
add_orchestration_model_options, prepare_bedrock_rig_turn_for_mode, prepare_rig_turn,
|
||||
prepare_rig_turn_for_mode, MCPToolTarget, OrchestrationModelOption, PreparedRigTurn,
|
||||
RigRequestMode,
|
||||
};
|
||||
use super::rig_tool::action_from_tool_call;
|
||||
use super::ProviderRunProfile;
|
||||
@@ -76,6 +77,7 @@ pub(crate) async fn prepare_provider_run(
|
||||
base_provider_config: crate::ai::provider::ProviderConfig,
|
||||
cli_provider_config: crate::ai::provider::ProviderConfig,
|
||||
mut params: RequestParams,
|
||||
orchestration_models: Vec<OrchestrationModelOption>,
|
||||
) -> anyhow::Result<PreparedProviderRun> {
|
||||
let (supported_tools, supported_cli_agent_tools) =
|
||||
crate::ai::agent::api::prepare_direct_provider_params(&mut params);
|
||||
@@ -112,7 +114,7 @@ pub(crate) async fn prepare_provider_run(
|
||||
}
|
||||
};
|
||||
|
||||
let (base_runtime, prepared) = prepare_provider_profile(
|
||||
let (base_runtime, mut prepared) = prepare_provider_profile(
|
||||
base_provider_config,
|
||||
params,
|
||||
supported_tools.clone(),
|
||||
@@ -120,6 +122,7 @@ pub(crate) async fn prepare_provider_run(
|
||||
None,
|
||||
)
|
||||
.await?;
|
||||
add_orchestration_model_options(&mut prepared.request.tools, &orchestration_models);
|
||||
let (cli_runtime, cli_prepared) = prepare_provider_profile(
|
||||
cli_provider_config,
|
||||
cli_params,
|
||||
|
||||
@@ -33,6 +33,91 @@ pub(crate) struct PreparedRigTurn {
|
||||
pub mcp_tool_aliases: HashMap<String, MCPToolTarget>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
pub(crate) struct OrchestrationModelOption {
|
||||
pub id: String,
|
||||
pub display_name: String,
|
||||
pub provider: String,
|
||||
pub quality: Option<f32>,
|
||||
pub cost: Option<f32>,
|
||||
pub credit_multiplier: Option<f32>,
|
||||
}
|
||||
|
||||
pub(crate) fn add_orchestration_model_options(
|
||||
tools: &mut [ToolDefinition],
|
||||
models: &[OrchestrationModelOption],
|
||||
) {
|
||||
if models.is_empty() {
|
||||
return;
|
||||
}
|
||||
let Some(tool) = tools.iter_mut().find(|tool| tool.name == "run_agents") else {
|
||||
return;
|
||||
};
|
||||
|
||||
let catalog = models
|
||||
.iter()
|
||||
.map(|model| {
|
||||
let mut details = vec![
|
||||
format!("id={:?}", model.id),
|
||||
format!("name={:?}", model.display_name),
|
||||
format!("provider={:?}", model.provider),
|
||||
];
|
||||
if let Some(quality) = model.quality {
|
||||
details.push(format!("quality_score={quality:.2}"));
|
||||
}
|
||||
if let Some(cost) = model.cost {
|
||||
details.push(format!("cost_score={cost:.2}"));
|
||||
}
|
||||
if let Some(multiplier) = model.credit_multiplier {
|
||||
details.push(format!("credit_multiplier={multiplier:.2}x"));
|
||||
}
|
||||
format!("- {}", details.join(", "))
|
||||
})
|
||||
.collect::<Vec<_>>()
|
||||
.join("\n");
|
||||
let description = format!(
|
||||
"Required model for this child. Select exactly one available model ID. Prioritize the model best suited to the child's task and most likely to succeed. Among similarly capable models, prefer the lower-cost option; do not sacrifice material capability merely to choose the cheapest model. Cost scores represent relative consumption, with higher values costing more.\nAvailable models:\n{catalog}"
|
||||
);
|
||||
let model_ids = models
|
||||
.iter()
|
||||
.map(|model| serde_json::Value::String(model.id.clone()))
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
let Some(agent_items) = tool
|
||||
.input_schema
|
||||
.get_mut("properties")
|
||||
.and_then(|properties| properties.get_mut("agent_run_configs"))
|
||||
.and_then(|configs| configs.get_mut("items"))
|
||||
else {
|
||||
return;
|
||||
};
|
||||
let Some(properties) = agent_items
|
||||
.get_mut("properties")
|
||||
.and_then(serde_json::Value::as_object_mut)
|
||||
else {
|
||||
return;
|
||||
};
|
||||
properties.insert(
|
||||
"model_id".to_string(),
|
||||
serde_json::json!({
|
||||
"type": "string",
|
||||
"enum": model_ids,
|
||||
"description": description,
|
||||
}),
|
||||
);
|
||||
let Some(required) = agent_items
|
||||
.get_mut("required")
|
||||
.and_then(serde_json::Value::as_array_mut)
|
||||
else {
|
||||
return;
|
||||
};
|
||||
if !required.iter().any(|field| field == "model_id") {
|
||||
required.push(serde_json::Value::String("model_id".to_string()));
|
||||
}
|
||||
|
||||
tool.description = "Start one or more child agents. Assign each child the best-fit, cost-effective model from its required model_id choices, prioritizing capability and likelihood of success over price. Use independent, uniquely named tasks and do not launch duplicate agents for follow-up work. After launch, call wait_for_events when you need child-agent results instead of repeating their work yourself.".to_string();
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Eq, PartialEq, serde::Serialize, serde::Deserialize)]
|
||||
pub(crate) struct MCPToolTarget {
|
||||
pub server_id: Option<Uuid>,
|
||||
|
||||
@@ -8,8 +8,8 @@ use galaxy_util::local_or_remote_path::LocalOrRemotePath;
|
||||
use warp_multi_agent_api::ToolType;
|
||||
|
||||
use super::{
|
||||
input_messages, prepare_bedrock_rig_turn, prepare_rig_turn, prepare_rig_turn_for_mode,
|
||||
tool_definitions, RigRequestMode,
|
||||
add_orchestration_model_options, input_messages, prepare_bedrock_rig_turn, prepare_rig_turn,
|
||||
prepare_rig_turn_for_mode, tool_definitions, OrchestrationModelOption, RigRequestMode,
|
||||
};
|
||||
use crate::ai::agent::api::RequestParams;
|
||||
use crate::ai::agent::task::TaskId;
|
||||
@@ -631,6 +631,51 @@ fn modern_and_legacy_orchestration_tools_follow_subagent_capabilities() {
|
||||
.any(|tool| matches!(tool.name.as_str(), "run_agents" | "start_agent")));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn run_agents_tool_requires_a_best_fit_model_from_the_available_catalog() {
|
||||
let (mut tools, _) = tool_definitions(&[ToolType::Subagent], None);
|
||||
add_orchestration_model_options(
|
||||
&mut tools,
|
||||
&[
|
||||
OrchestrationModelOption {
|
||||
id: "strong-model".to_string(),
|
||||
display_name: "Strong Model".to_string(),
|
||||
provider: "Provider A".to_string(),
|
||||
quality: Some(0.95),
|
||||
cost: Some(0.8),
|
||||
credit_multiplier: Some(2.0),
|
||||
},
|
||||
OrchestrationModelOption {
|
||||
id: "efficient-model".to_string(),
|
||||
display_name: "Efficient Model".to_string(),
|
||||
provider: "Provider B".to_string(),
|
||||
quality: Some(0.9),
|
||||
cost: Some(0.3),
|
||||
credit_multiplier: Some(0.5),
|
||||
},
|
||||
],
|
||||
);
|
||||
|
||||
let run_agents = tools
|
||||
.iter()
|
||||
.find(|tool| tool.name == "run_agents")
|
||||
.expect("run_agents tool");
|
||||
let model = &run_agents.input_schema["properties"]["agent_run_configs"]["items"]["properties"]
|
||||
["model_id"];
|
||||
assert_eq!(
|
||||
model["enum"],
|
||||
serde_json::json!(["strong-model", "efficient-model"])
|
||||
);
|
||||
assert!(model["description"]
|
||||
.as_str()
|
||||
.is_some_and(|description| description.contains("Prioritize the model best suited")));
|
||||
assert!(
|
||||
run_agents.input_schema["properties"]["agent_run_configs"]["items"]["required"]
|
||||
.as_array()
|
||||
.is_some_and(|required| required.iter().any(|field| field == "model_id"))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[allow(deprecated)]
|
||||
fn grouped_mcp_tool_names_use_the_installation_id_not_the_display_name() {
|
||||
|
||||
@@ -31,6 +31,7 @@ async fn missing_cli_provider_route_falls_back_to_base_provider_profile() {
|
||||
ProviderConfig::OpenAI(openai_config("base-provider-model")),
|
||||
ProviderConfig::None,
|
||||
params,
|
||||
Vec::new(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
@@ -181,6 +181,7 @@ pub(crate) fn action_from_tool_call(
|
||||
name: required_nonempty_string(config, "name")?,
|
||||
prompt: required_nonempty_string(config, "prompt")?,
|
||||
title: optional_string(config, "title")?.unwrap_or_default(),
|
||||
model_id: optional_string(config, "model_id")?.unwrap_or_default(),
|
||||
})
|
||||
})
|
||||
.collect::<Result<_, String>>()?,
|
||||
|
||||
@@ -444,7 +444,8 @@ fn run_agents_calls_decode_to_local_domain_requests_with_safe_defaults() {
|
||||
{
|
||||
"name": "runtime",
|
||||
"prompt": "Inspect runtime behavior",
|
||||
"title": "Runtime investigator"
|
||||
"title": "Runtime investigator",
|
||||
"model_id": "strong-model"
|
||||
},
|
||||
{
|
||||
"name": "tests",
|
||||
@@ -476,9 +477,11 @@ fn run_agents_calls_decode_to_local_domain_requests_with_safe_defaults() {
|
||||
"Inspect runtime behavior"
|
||||
);
|
||||
assert_eq!(request.agent_run_configs[0].title, "Runtime investigator");
|
||||
assert_eq!(request.agent_run_configs[0].model_id, "strong-model");
|
||||
assert_eq!(request.agent_run_configs[1].name, "tests");
|
||||
assert_eq!(request.agent_run_configs[1].prompt, "Design focused tests");
|
||||
assert!(request.agent_run_configs[1].title.is_empty());
|
||||
assert!(request.agent_run_configs[1].model_id.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
||||
Reference in New Issue
Block a user