use super::*;
#[derive(Debug)]
pub struct ResolvedStage {
pub provider_name: String,
pub model: String,
pub tools: Vec<Tool>,
pub fallbacks: Vec<leviath_core::blueprint::ModelEntry>,
pub output: Option<leviath_core::output::OutputSpec>,
}
pub(crate) const DEFAULT_CONTEXT_WINDOW_TOKENS: usize = 8192;
pub(crate) fn context_window_tokens(world: &World, provider_name: &str, model: &str) -> usize {
match world
.get_resource::<Providers>()
.and_then(|p| p.0.get(provider_name))
{
Some(provider) => provider.max_context_tokens(model),
None => {
tracing::warn!(
provider = provider_name,
model,
"provider not registered; using default context window for percentage budgets"
);
DEFAULT_CONTEXT_WINDOW_TOKENS
}
}
}
pub(crate) fn stage_setup_from(
stage: &leviath_core::Stage,
global_hints: leviath_core::config::PromptHints,
agent_hints: leviath_core::config::PromptHintOverrides,
output: Option<leviath_core::output::OutputSpec>,
) -> StageSetup {
let temperature = stage
.model
.parameters
.get("temperature")
.and_then(|v| v.as_f64())
.map(|t| t as f32);
let extra_params: serde_json::Map<String, serde_json::Value> = stage
.model
.parameters
.iter()
.filter(|(k, _)| k.as_str() != "temperature" && k.as_str() != "max_output_tokens")
.map(|(k, v)| (k.clone(), v.clone()))
.collect();
let max_output_tokens = stage
.model
.parameters
.get("max_output_tokens")
.and_then(|v| v.as_u64())
.map(|t| t as usize);
let base_prompt = stage
.config
.get("system_prompt")
.and_then(|v| v.as_str())
.map(String::from);
let system_prompt = match &stage.mode {
leviath_core::blueprint::StageMode::FanOut { config }
if !config.split_prompt.trim().is_empty() =>
{
Some(match base_prompt {
Some(base) => format!("{base}\n\n{}", config.split_prompt),
None => config.split_prompt.clone(),
})
}
_ => base_prompt,
};
let system_prompt = match (&output, stage.require_output) {
(Some(spec), true) => {
let described = leviath_core::describe_spec(spec);
let demand = match described.is_empty() {
true => format!(
"Before this stage ends you must call `{tool}` with your final answer. It is \
the only thing the caller receives.",
tool = leviath_core::blueprint::SUBMIT_OUTPUT_TOOL
),
false => format!(
"Before this stage ends you must call `{tool}` with your final answer. It is \
the only thing the caller receives.\n\n{described}",
tool = leviath_core::blueprint::SUBMIT_OUTPUT_TOOL
),
};
Some(match system_prompt {
Some(base) => format!("{base}\n\n{demand}"),
None => demand,
})
}
_ => system_prompt,
};
let batch_tool_hint = leviath_core::taint::resolve_batch_tool_hint(
global_hints.batch_tool,
agent_hints.batch_tool,
stage.batch_tool_hint,
);
let shell_hint = leviath_core::taint::resolve_shell_hint(
global_hints.shell,
agent_hints.shell,
stage.shell_hint,
);
StageSetup {
inference_config: InferenceConfig {
temperature,
max_output_tokens,
extra_params,
batch_tool_hint,
shell_hint,
request_timeout_secs: stage.model.request_timeout_secs,
},
routing: stage.tool_result_routing.clone(),
accepts_messages: stage.accepts_messages,
context_layout: stage.context_layout.clone(),
system_prompt,
output,
}
}
pub fn spawn_agent(
world: &mut World,
agent_id: String,
blueprint: leviath_core::Blueprint,
task: &str,
stages: Vec<ResolvedStage>,
global_hints: leviath_core::config::PromptHints,
) -> Result<Entity, String> {
let seeds = std::collections::HashMap::from([("task".to_string(), task.to_string())]);
spawn_agent_seeded(
world,
SeededSpawn {
agent_id,
blueprint,
seeds,
stages,
global_hints,
global_nudge: leviath_core::NudgeConfig::default(),
region_scripts: std::collections::HashMap::new(),
},
)
}
pub struct SeededSpawn {
pub agent_id: String,
pub blueprint: leviath_core::Blueprint,
pub seeds: std::collections::HashMap<String, String>,
pub stages: Vec<ResolvedStage>,
pub global_hints: leviath_core::config::PromptHints,
pub global_nudge: leviath_core::NudgeConfig,
pub region_scripts: std::collections::HashMap<
String,
std::sync::Arc<leviath_scripting::region_hook::RegionScript>,
>,
}
pub fn spawn_agent_seeded(world: &mut World, spawn: SeededSpawn) -> Result<Entity, String> {
let SeededSpawn {
agent_id,
mut blueprint,
seeds,
stages,
global_hints,
global_nudge,
region_scripts,
} = spawn;
let seeds = &seeds;
let stage_windows: Vec<usize> = stages
.iter()
.map(|rs| context_window_tokens(world, &rs.provider_name, &rs.model))
.collect();
blueprint.context_layout = blueprint.context_layout.resolved(stage_windows[0]);
for (i, stage) in blueprint.stages.iter_mut().enumerate() {
if let Some(layout) = &stage.context_layout {
stage.context_layout = Some(layout.resolved(stage_windows[i]));
}
}
blueprint
.context_layout
.validate()
.map_err(|e| e.to_string())?;
for stage in &blueprint.stages {
if let Some(layout) = &stage.context_layout {
layout.validate().map_err(|e| e.to_string())?;
}
}
let stage_outputs: Vec<Option<leviath_core::output::OutputSpec>> =
stages.iter().map(|rs| rs.output.clone()).collect();
let stage_infs: Vec<StageInference> = stages
.into_iter()
.map(|rs| StageInference {
provider_name: rs.provider_name,
model: rs.model,
tools: rs.tools,
tool_filter: None, fallbacks: rs.fallbacks,
output: rs.output,
})
.collect();
let agent_hints = leviath_core::config::PromptHintOverrides {
batch_tool: blueprint.batch_tool_hint,
shell: blueprint.shell_hint,
};
let setups: Vec<StageSetup> = blueprint
.stages
.iter()
.zip(stage_outputs)
.map(|(s, output)| stage_setup_from(s, global_hints, agent_hints, output))
.collect();
let mut window = ContextWindow::new(blueprint.context_layout.total_budget_tokens);
window.region_scripts = region_scripts;
crate::context_setup::init_window_seeded(&mut window, &blueprint, seeds);
let prompts: Vec<Option<String>> = setups.iter().map(|s| s.system_prompt.clone()).collect();
crate::context_setup::ensure_stage_instructions_region(&mut window, &prompts);
apply_stage_context(&setups[0], &mut window)?;
let stage0_name = blueprint.stages[0].name.clone();
let stage0_inf = stage_infs[0].clone();
let setup0 = &setups[0];
let stage0_cfg = setup0.inference_config.clone();
let stage0_routing = setup0.routing.clone();
let accepts_messages = setup0.accepts_messages;
let mut visits = VisitCounts::default();
*visits.0.entry(stage0_name.clone()).or_insert(0) += 1;
let ledger = StageLedger(
blueprint
.stages
.iter()
.enumerate()
.map(|(i, s)| leviath_core::run_meta::StageRecord::new(s.name.clone(), i))
.collect(),
);
let repetition = blueprint
.repetition_detection
.as_ref()
.map(crate::repetition::RepetitionDetector::from_detection_config);
let entity = world
.spawn((
AgentBlueprint(blueprint),
AgentState {
agent_id,
current_stage: stage0_name,
iteration: 0,
status: AgentStatus::Active,
spawned_children_ids: vec![],
pending_wait: None,
accepts_messages,
},
MessageInbox::default(),
StageCursor { index: 0 },
StageProgress::default(),
StageInferences(stage_infs),
StageSetups(setups),
visits,
window,
stage0_inf,
stage0_cfg,
ReadyToInfer,
))
.id();
world.entity_mut(entity).insert((
ledger,
StageIoBuffer::default(),
crate::pipeline::response::GlobalNudge(global_nudge),
));
if let Some(detector) = repetition {
world.entity_mut(entity).insert(detector);
}
if let Some(routing) = stage0_routing {
world
.entity_mut(entity)
.insert(crate::components::ToolResultRoutingComponent { routing });
}
Ok(entity)
}
#[cfg(test)]
mod stage_instructions_fit_tests {
fn layout(window: usize) -> leviath_core::layout::ContextLayout {
use leviath_core::layout::{BudgetSpec, ContextLayout, RegionDefinition};
let pct = |p: f64| BudgetSpec::Percent {
percent: p,
min: None,
max: None,
};
let mut task =
RegionDefinition::new("task".to_string(), leviath_core::RegionKind::Pinned, 0);
task.budget = pct(0.02);
let mut instr = RegionDefinition::new(
leviath_core::layout::STAGE_INSTRUCTIONS_REGION.to_string(),
leviath_core::RegionKind::Pinned,
0,
);
instr.budget = pct(0.03);
ContextLayout::new(vec![task, instr], window).resolved(window)
}
fn big_prompt() -> String {
"word ".repeat(2_600)
}
#[test]
fn a_stage_prompt_measured_at_spawn_uses_the_declared_region() {
let window_tokens = 128_000;
let layout = layout(window_tokens);
let task_max = layout
.regions
.iter()
.find(|r| r.name == "task")
.expect("task")
.max_tokens;
let instr_max = layout
.regions
.iter()
.find(|r| r.name == leviath_core::layout::STAGE_INSTRUCTIONS_REGION)
.expect("stage_instructions")
.max_tokens;
let prompt = big_prompt();
let tokens = leviath_core::estimate_tokens(&format!("[Stage instructions: {prompt}]"));
assert!(
tokens > task_max && tokens < instr_max,
"the fixture must reproduce the reported shape: {tokens} vs task {task_max} / \
stage_instructions {instr_max}"
);
let bp = leviath_core::Blueprint::new(
"t".to_string(),
"d".to_string(),
vec![leviath_core::Stage::new(
"work".to_string(),
leviath_core::blueprint::ModelConfig::new("p".to_string(), "m".to_string()),
)],
layout,
);
let mut window = crate::components::ContextWindow::new(window_tokens);
crate::context_setup::init_window_seeded(
&mut window,
&bp,
&std::collections::HashMap::new(),
);
let setup = crate::pipeline::transition::StageSetup {
inference_config: crate::components::InferenceConfig {
temperature: None,
max_output_tokens: None,
extra_params: Default::default(),
batch_tool_hint: false,
shell_hint: false,
request_timeout_secs: None,
},
routing: None,
accepts_messages: true,
context_layout: None,
system_prompt: Some(prompt),
output: None,
};
crate::pipeline::transition::apply_stage_context(&setup, &mut window)
.expect("the prompt fits the region declared for it");
let instr = window
.get_region(leviath_core::layout::STAGE_INSTRUCTIONS_REGION)
.expect("region exists");
assert!(
instr.content.iter().any(|e| e.content.contains("word")),
"the prompt landed in stage_instructions"
);
}
#[test]
fn a_blueprint_that_declares_no_region_still_gets_one() {
use leviath_core::layout::{BudgetSpec, ContextLayout, RegionDefinition};
let window_tokens = 128_000;
let prompt = big_prompt();
let mut task =
RegionDefinition::new("task".to_string(), leviath_core::RegionKind::Pinned, 0);
task.budget = BudgetSpec::Percent {
percent: 0.02,
min: None,
max: None,
};
let only_task = ContextLayout::new(vec![task], window_tokens).resolved(window_tokens);
let bp = leviath_core::Blueprint::new(
"t".to_string(),
"d".to_string(),
vec![leviath_core::Stage::new(
"work".to_string(),
leviath_core::blueprint::ModelConfig::new("p".to_string(), "m".to_string()),
)],
only_task,
);
let mut window = crate::components::ContextWindow::new(window_tokens);
crate::context_setup::init_window_seeded(
&mut window,
&bp,
&std::collections::HashMap::new(),
);
let prompts = vec![Some(prompt.clone())];
crate::context_setup::ensure_stage_instructions_region(&mut window, &prompts);
let setup = crate::pipeline::transition::StageSetup {
inference_config: crate::components::InferenceConfig {
temperature: None,
max_output_tokens: None,
extra_params: Default::default(),
batch_tool_hint: false,
shell_hint: false,
request_timeout_secs: None,
},
routing: None,
accepts_messages: true,
context_layout: None,
system_prompt: Some(prompt),
output: None,
};
crate::pipeline::transition::apply_stage_context(&setup, &mut window)
.expect("the prompt no longer has to fit the caller's task region");
let task_region = window.get_region("task").expect("task");
assert!(
task_region.content.is_empty(),
"the task region is left for the caller's task"
);
let instr = window
.get_region(leviath_core::layout::STAGE_INSTRUCTIONS_REGION)
.expect("the runtime made one");
assert!(instr.content.iter().any(|e| e.content.contains("word")));
}
#[test]
fn no_region_is_made_when_no_stage_has_a_prompt() {
let mut window = crate::components::ContextWindow::new(1_000);
crate::context_setup::ensure_stage_instructions_region(&mut window, &[None, None]);
assert!(
window
.get_region(leviath_core::layout::STAGE_INSTRUCTIONS_REGION)
.is_none()
);
}
#[test]
fn a_declared_region_is_not_resized() {
let mut window = crate::components::ContextWindow::new(100_000);
window.add_region(leviath_core::Region::new(
leviath_core::layout::STAGE_INSTRUCTIONS_REGION.to_string(),
leviath_core::RegionKind::Pinned,
4_242,
));
crate::context_setup::ensure_stage_instructions_region(&mut window, &[Some(big_prompt())]);
assert_eq!(
window
.get_region(leviath_core::layout::STAGE_INSTRUCTIONS_REGION)
.expect("declared")
.max_tokens,
4_242
);
}
#[test]
fn an_impossible_prompt_is_still_refused_and_names_the_right_region() {
let mut window = crate::components::ContextWindow::new(1_000);
window.add_region(leviath_core::Region::new(
"task".to_string(),
leviath_core::RegionKind::Pinned,
40,
));
let prompt = "z".repeat(100_000);
crate::context_setup::ensure_stage_instructions_region(
&mut window,
&[Some(prompt.clone())],
);
assert_eq!(
window
.get_region(leviath_core::layout::STAGE_INSTRUCTIONS_REGION)
.expect("made")
.max_tokens,
250
);
let setup = crate::pipeline::transition::StageSetup {
inference_config: crate::components::InferenceConfig {
temperature: None,
max_output_tokens: None,
extra_params: Default::default(),
batch_tool_hint: false,
shell_hint: false,
request_timeout_secs: None,
},
routing: None,
accepts_messages: true,
context_layout: None,
system_prompt: Some(prompt),
output: None,
};
let err = crate::pipeline::transition::apply_stage_context(&setup, &mut window)
.expect_err("a prompt larger than the window cannot be housed");
assert!(
err.contains(leviath_core::layout::STAGE_INSTRUCTIONS_REGION),
"{err}"
);
}
#[test]
fn the_region_is_sized_for_the_widest_prompt() {
let mut window = crate::components::ContextWindow::new(100_000);
let small = "word ".repeat(10);
let large = big_prompt();
let expected = leviath_core::estimate_tokens(&format!("[Stage instructions: {large}]"));
crate::context_setup::ensure_stage_instructions_region(
&mut window,
&[Some(small), Some(large)],
);
assert_eq!(
window
.get_region(leviath_core::layout::STAGE_INSTRUCTIONS_REGION)
.expect("made")
.max_tokens,
expected
);
}
#[test]
fn a_scoped_stage_layout_still_routes_to_the_declared_region() {
use leviath_core::layout::{BudgetSpec, ContextLayout, RegionDefinition};
let window_tokens = 128_000;
let prompt = big_prompt();
let mut scoped_task =
RegionDefinition::new("task".to_string(), leviath_core::RegionKind::Pinned, 0);
scoped_task.budget = BudgetSpec::Percent {
percent: 0.02,
min: None,
max: None,
};
let scoped = ContextLayout::new(vec![scoped_task], window_tokens).resolved(window_tokens);
let bp = leviath_core::Blueprint::new(
"t".to_string(),
"d".to_string(),
vec![leviath_core::Stage::new(
"work".to_string(),
leviath_core::blueprint::ModelConfig::new("p".to_string(), "m".to_string()),
)],
layout(window_tokens),
);
let mut window = crate::components::ContextWindow::new(window_tokens);
crate::context_setup::init_window_seeded(
&mut window,
&bp,
&std::collections::HashMap::new(),
);
let setup = crate::pipeline::transition::StageSetup {
inference_config: crate::components::InferenceConfig {
temperature: None,
max_output_tokens: None,
extra_params: Default::default(),
batch_tool_hint: false,
shell_hint: false,
request_timeout_secs: None,
},
routing: None,
accepts_messages: true,
context_layout: Some(scoped),
system_prompt: Some(prompt),
output: None,
};
crate::pipeline::transition::apply_stage_context(&setup, &mut window)
.expect("the prompt fits the region declared for it");
let instr = window
.get_region(leviath_core::layout::STAGE_INSTRUCTIONS_REGION)
.expect("carried through the scoped layout");
assert!(
instr.content.iter().any(|e| e.content.contains("word")),
"the prompt landed in stage_instructions, not in the scoped task region"
);
}
}