use crate::run_support;
use bevy_ecs::prelude::*;
use rig_cassette::ecs::EffectLogResource;
use rig_cassette::effect_log::EffectLogRecorder;
use rig_core::{
effect::EffectKind,
message::{Message, UserContent},
};
use rig_ecs::{
agent::{
Context, DocumentId, DocumentText, Grant, MessageParts, Options, ProviderOptions,
RunResult, Settled, UsesModel, Utterance,
},
bus::{PendingEffect, RigSchedule},
systems::{RigSet, RunCommands},
};
use run_support::*;
const MODEL: &str = "t/model:default";
fn add_before_assemble<M>(
app: &mut bevy_app::App,
system: impl IntoScheduleConfigs<bevy_ecs::system::ScheduleSystem, M>,
) {
app.world_mut()
.resource_mut::<Schedules>()
.add_systems(RigSchedule, system.before(RigSet::Assemble));
}
fn settle(app: &mut bevy_app::App, run: Entity, what: &str) -> String {
tick_until(app, what, |world| world.get::<Settled>(run).is_some());
app.world()
.get::<RunResult>(run)
.expect("settled with an answer")
.0
.clone()
}
#[derive(Resource, Default)]
struct Once(bool);
fn despawn_the_assistant_utterance(
utterances: Query<(Entity, &rig_ecs::agent::Role), With<Utterance>>,
mut once: ResMut<Once>,
mut commands: Commands,
) {
if once.0 {
return;
}
for (entity, role) in &utterances {
if *role == rig_ecs::agent::Role::Assistant {
commands.entity(entity).despawn();
once.0 = true;
}
}
}
#[test]
fn an_utterance_despawned_before_assemble_leaves_the_next_request() {
let mut app = app();
let (agent, requests) = capturing_agent(&mut app, MODEL, MODEL, "ok");
app.init_resource::<Once>();
add_before_assemble(&mut app, despawn_the_assistant_utterance);
let history = vec![
MessageParts::User {
content: vec![UserContent::text("A")],
},
MessageParts::Assistant(rig_core::message::AssistantMessage::new(vec![
rig_core::message::AssistantContent::text("B"),
])),
];
let run = app
.world_mut()
.spawn_run(agent, &history, "C", false, Some(1));
settle(&mut app, run, "answered");
let requests = requests.lock().expect("requests");
assert_eq!(requests.len(), 1);
assert_eq!(
texts(&requests[0]),
vec!["system:You are terse.", "user:A", "user:C"],
"the assistant utterance left the history before the fold"
);
}
#[test]
fn one_document_attached_to_two_runs_appears_in_both_requests() {
let mut app = app();
let (agent, requests) = capturing_agent(&mut app, MODEL, MODEL, "ok");
let document = app
.world_mut()
.spawn((
DocumentId("shared".to_owned()),
DocumentText("one document, many turns".to_owned()),
))
.id();
app.world_mut().spawn((Context(document), ChildOf(agent)));
let first = app
.world_mut()
.spawn_run(agent, &[], "first?", false, Some(1));
let second = app
.world_mut()
.spawn_run(agent, &[], "second?", false, Some(1));
settle(&mut app, first, "first");
settle(&mut app, second, "second");
let requests = requests.lock().expect("requests");
assert_eq!(requests.len(), 2);
for request in requests.iter() {
assert_eq!(request.documents.len(), 1);
assert_eq!(request.documents[0].id, "shared");
assert_eq!(request.documents[0].text, "one document, many turns");
}
let documents = app
.world_mut()
.query::<&DocumentId>()
.iter(app.world())
.count();
assert_eq!(documents, 1, "one entity, two requests");
}
#[test]
fn a_tool_granted_by_a_relationship_is_advertised_and_gone_after_removal() {
let mut app = app();
let (agent, requests) = capturing_agent(&mut app, MODEL, MODEL, "ok");
let tool = register(
&mut app,
"t/tool:add#0",
NeverCalled {
name: "add".to_owned(),
},
);
let before = app.world_mut().spawn_run(agent, &[], "one", false, Some(1));
settle(&mut app, before, "before the grant");
let grant = app.world_mut().spawn((Grant(tool), ChildOf(agent))).id();
let during = app.world_mut().spawn_run(agent, &[], "two", false, Some(1));
settle(&mut app, during, "with the grant");
app.world_mut().despawn(grant);
let after = app
.world_mut()
.spawn_run(agent, &[], "three", false, Some(1));
settle(&mut app, after, "after the grant");
let requests = requests.lock().expect("requests");
let advertised: Vec<Vec<String>> = requests
.iter()
.map(|request| {
request
.tools
.iter()
.map(|tool| tool.name.to_string())
.collect()
})
.collect();
assert_eq!(
advertised,
vec![Vec::<String>::new(), vec!["add".to_owned()], Vec::new()]
);
}
#[test]
fn uses_model_swapped_on_a_run_changes_the_next_requests_key() {
let mut app = app();
let (default_model, default_requests) = Capturing::new(MODEL, "from default");
let (fast_model, fast_requests) = Capturing::new("t/model:fast", "from fast");
let default_model = register(&mut app, MODEL, default_model);
let fast_model = register(&mut app, "t/model:fast", fast_model);
EffectLogResource::install(app.world_mut(), EffectLogRecorder::new());
let agent = spawn_agent(app.world_mut(), "t", default_model);
let first = app.world_mut().spawn_run(agent, &[], "one", false, Some(1));
assert_eq!(settle(&mut app, first, "default"), "from default");
let second = app.world_mut().spawn_run(agent, &[], "two", false, Some(1));
app.world_mut()
.entity_mut(second)
.insert(UsesModel(fast_model));
assert_eq!(settle(&mut app, second, "fast"), "from fast");
assert_eq!(default_requests.lock().expect("requests").len(), 1);
assert_eq!(fast_requests.lock().expect("requests").len(), 1);
let log = app.world().resource::<EffectLogResource>().log();
let keys: Vec<String> = log
.records
.iter()
.map(|record| record.key.to_string())
.collect();
assert_eq!(keys, vec!["t/model:default", "t/model:fast"]);
}
fn patch_temperature(mut effects: Query<&mut PendingEffect, Added<PendingEffect>>) {
for mut effect in &mut effects {
if let EffectKind::Completion { request, .. } = &mut effect.kind {
request.temperature = Some(0.5);
}
}
}
#[test]
fn a_patch_system_rewrites_the_folded_request_and_the_record_holds_it() {
let mut app = app();
let (model, requests) = Capturing::new(MODEL, "ok");
let model = register(&mut app, MODEL, model);
EffectLogResource::install(app.world_mut(), EffectLogRecorder::new());
app.world_mut()
.resource_mut::<Schedules>()
.add_systems(RigSchedule, patch_temperature.in_set(RigSet::Patch));
let agent = spawn_agent(app.world_mut(), "t", model);
let run = app.world_mut().spawn_run(agent, &[], "one", false, Some(1));
settle(&mut app, run, "patched");
assert_eq!(
requests.lock().expect("requests")[0].temperature,
Some(0.5),
"the handler saw the patch"
);
let log = app.world().resource::<EffectLogResource>().log();
let EffectKind::Completion { request, .. } = &log.records[0].kind else {
panic!("a completion");
};
assert_eq!(
request.temperature,
Some(0.5),
"the record is the patched request"
);
assert_eq!(
app.world()
.get::<rig_ecs::agent::Temperature>(agent)
.expect("a temperature")
.0,
None
);
}
fn shout_the_prompt(
utterances: Query<&rig_ecs::agent::Role, With<Utterance>>,
mut parts: Query<(&ChildOf, &mut rig_ecs::agent::content::parts::ContentPart)>,
) {
for (parent, mut part) in &mut parts {
let rig_ecs::agent::content::parts::ContentPart::Text(text) = &mut *part else {
continue;
};
if utterances
.get(parent.parent())
.is_ok_and(|role| *role == rig_ecs::agent::Role::User)
&& !text.text.ends_with('!')
{
text.text.push('!');
}
}
}
#[test]
fn a_system_before_assemble_rewrites_an_utterance() {
let mut app = app();
let (agent, requests) = capturing_agent(&mut app, MODEL, MODEL, "ok");
add_before_assemble(&mut app, shout_the_prompt);
let run = app
.world_mut()
.spawn_run(agent, &[], "quietly", false, Some(1));
settle(&mut app, run, "shouted");
let requests = requests.lock().expect("requests");
assert_eq!(
texts(&requests[0]),
vec!["system:You are terse.", "user:quietly!"]
);
let Message::User { content } = &requests[0].chat_history[1] else {
panic!("the prompt");
};
assert_eq!(content.len(), 1);
}
fn stop_the_run(effects: Query<Entity, Added<PendingEffect>>, mut commands: Commands) {
for effect in &effects {
commands.entity(effect).despawn();
}
}
#[test]
fn despawning_the_effect_in_patch_fails_the_run_cancelled() {
let mut app = app();
let (model, requests) = Capturing::new(MODEL, "ok");
let model = register(&mut app, MODEL, model);
EffectLogResource::install(app.world_mut(), EffectLogRecorder::new());
app.world_mut()
.resource_mut::<Schedules>()
.add_systems(RigSchedule, stop_the_run.in_set(RigSet::Patch));
let agent = spawn_agent(app.world_mut(), "t", model);
let run = app.world_mut().spawn_run(agent, &[], "one", false, Some(1));
tick_until(&mut app, "cancelled", |world| {
world.get::<rig_ecs::agent::Failed>(run).is_some()
});
let failed = app
.world()
.get::<rig_ecs::agent::Failed>(run)
.expect("failed");
assert!(
matches!(&failed.0, rig_ecs::agent::Failure::Cancelled(report) if report.kind == rig_core::error::ErrorKind::Cancelled),
"{failed:?}"
);
assert!(
requests.lock().expect("requests").is_empty(),
"the handler never saw it"
);
let log = app.world().resource::<EffectLogResource>().log();
assert!(
log.records.is_empty(),
"despawned before dispatch: no record"
);
}
#[test]
fn the_runtime_measures_itself_into_bevy_diagnostics() {
use bevy_diagnostic::DiagnosticsStore;
use rig_ecs::systems::diagnostics::RUNS_LIVE;
let mut app = app();
let (agent, _) = capturing_agent(&mut app, "t/model:m", "m", "hello");
let run = app.world_mut().spawn_run(agent, &[], "hi", false, None);
app.update();
ended(&mut app, run, "settled");
let store = app.world().resource::<DiagnosticsStore>();
let live = store.get(&RUNS_LIVE).expect("registered");
assert_eq!(live.value(), Some(0.0));
assert!(live.history_len() >= 1);
}
#[test]
fn a_runs_options_overlay_the_agents() {
use rig_core::completion::{CacheRetention, Effort, GenerationOptions, Reasoning};
let mut app = app();
let (agent, requests) = capturing_agent(&mut app, MODEL, MODEL, "ok");
app.world_mut().entity_mut(agent).insert(Options(
GenerationOptions::default().reasoning(Effort::High).seed(7),
));
let plain = app
.world_mut()
.spawn_run(agent, &[], "first?", false, Some(1));
settle(&mut app, plain, "the agent's options");
let overlaid = app
.world_mut()
.spawn_run(agent, &[], "second?", false, Some(1));
app.world_mut().entity_mut(overlaid).insert(Options(
GenerationOptions::default()
.cache(CacheRetention::Long)
.seed(9),
));
settle(&mut app, overlaid, "the run's options");
let requests = requests.lock().expect("requests");
let [first, second] = requests.as_slice() else {
panic!("two requests: {}", requests.len());
};
assert_eq!(
first.options,
GenerationOptions::default().reasoning(Effort::High).seed(7)
);
assert_eq!(
second.options.reasoning,
Some(Reasoning::Effort(Effort::High))
);
assert_eq!(second.options.cache, Some(CacheRetention::Long));
assert_eq!(second.options.seed, Some(9));
}
#[test]
fn a_runs_provider_options_overlay_the_agents() {
fn entries(value: serde_json::Value) -> rig_core::completion::ProviderOptions {
serde_json::from_value(value).expect("provider options")
}
let mut app = app();
let (agent, requests) = capturing_agent(&mut app, MODEL, MODEL, "ok");
let agents = entries(serde_json::json!({
"alpha": {"*": {"top_k": 4}},
"beta": {"*": {"min_p": 0.1}},
}));
app.world_mut()
.entity_mut(agent)
.insert(ProviderOptions(agents.clone()));
let plain = app
.world_mut()
.spawn_run(agent, &[], "first?", false, Some(1));
settle(&mut app, plain, "the agent's provider options");
let overlaid = app
.world_mut()
.spawn_run(agent, &[], "second?", false, Some(1));
app.world_mut()
.entity_mut(overlaid)
.insert(ProviderOptions(entries(
serde_json::json!({"alpha": {"*": {"top_k": 8}}}),
)));
settle(&mut app, overlaid, "the run's provider options");
let requests = requests.lock().expect("requests");
let [first, second] = requests.as_slice() else {
panic!("two requests: {}", requests.len());
};
assert_eq!(first.provider_options, agents);
assert_eq!(
second.provider_options,
entries(serde_json::json!({
"alpha": {"*": {"top_k": 8}},
"beta": {"*": {"min_p": 0.1}},
}))
);
}