Improve orchestration agent model selection

This commit is contained in:
2026-08-26 02:15:52 -05:00
parent 92ab03be07
commit 2022716f91
18 changed files with 486 additions and 60 deletions
+1
View File
@@ -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;
+6 -3
View File
@@ -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,
+85
View File
@@ -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>,
+47 -2
View File
@@ -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() {
+1
View File
@@ -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();
+1
View File
@@ -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>>()?,
+4 -1
View File
@@ -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]