use bevy_app::prelude::*;
use bevy_ecs::lifecycle::HookContext;
use bevy_ecs::prelude::*;
use bevy_ecs::system::SystemParam;
use bevy_ecs::world::DeferredWorld;
use rig_core::conversation::AuthorKind;
use rig_core::effect::{Answered as EffectOutcome, Decider, EffectKind};
use rig_ecs::inbox::Incoming;
use rig_ecs::stage::{
AgentPass, Decision, Effect, EffectRequest, ForAgent, IncomingMessages, Outcome, PluginEffects,
ReadyToStart, Settled, Stage,
};
use crate::commands::CommandSystems;
use serde::{Deserialize, Serialize};
#[derive(Component, Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[component(on_add = request)]
#[require(Name = Name::new("ask"))]
pub struct Ask {
pub title: String,
pub body: String,
pub choices: Vec<Choice>,
pub free_text: bool,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub initial: Option<String>,
#[serde(default, skip_serializing_if = "std::ops::Not::not")]
pub multiline: bool,
#[serde(default, skip_serializing_if = "std::ops::Not::not")]
pub inline: bool,
#[serde(default, skip_serializing_if = "std::ops::Not::not")]
pub multiple: bool,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct Choice {
pub id: String,
pub label: String,
#[serde(default, skip_serializing_if = "std::ops::Not::not")]
pub checked: bool,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub more: Option<String>,
}
impl Choice {
pub fn new(id: impl Into<String>, label: impl Into<String>) -> Self {
Self {
id: id.into(),
label: label.into(),
checked: false,
more: None,
}
}
#[must_use]
pub fn checked(mut self) -> Self {
self.checked = true;
self
}
#[must_use]
pub fn more(mut self, text: impl Into<String>) -> Self {
self.more = Some(text.into());
self
}
}
impl Ask {
pub fn new(title: impl Into<String>) -> Self {
Self {
title: title.into(),
body: String::new(),
choices: Vec::new(),
free_text: false,
initial: None,
multiline: false,
inline: false,
multiple: false,
}
}
pub fn yes_no(title: impl Into<String>) -> Self {
Self::new(title).choice("yes", "Yes").choice("no", "No")
}
pub fn body(mut self, body: impl Into<String>) -> Self {
self.body = body.into();
self
}
pub fn choice(self, id: impl Into<String>, label: impl Into<String>) -> Self {
self.with(Choice::new(id, label))
}
pub fn with(mut self, choice: Choice) -> Self {
self.choices.push(choice);
self
}
pub fn free_text(mut self) -> Self {
self.free_text = true;
self
}
pub fn initial(mut self, text: impl Into<String>) -> Self {
self.initial = Some(text.into());
self.free_text = true;
self
}
pub fn multiline(mut self) -> Self {
self.multiline = true;
self.free_text = true;
self
}
pub fn inline(mut self) -> Self {
self.inline = true;
self
}
pub fn multiple(mut self) -> Self {
self.multiple = true;
self
}
}
#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct Answer {
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub choices: Vec<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub text: Option<String>,
}
impl Answer {
pub fn choice(id: impl Into<String>) -> Self {
Self::choices([id])
}
pub fn choices(ids: impl IntoIterator<Item = impl Into<String>>) -> Self {
Self {
choices: ids.into_iter().map(Into::into).collect(),
text: None,
}
}
pub fn text(text: impl Into<String>) -> Self {
Self::default().with_text(text)
}
pub fn with_text(mut self, text: impl Into<String>) -> Self {
self.text = Some(text.into());
self
}
pub fn outcome(&self) -> Outcome {
Outcome::custom(serde_json::to_value(self).unwrap_or_default())
}
pub fn of(outcome: &Outcome) -> Option<Self> {
match &outcome.0 {
Ok(EffectOutcome::Custom { payload }) => Self::deserialize(payload).ok(),
_ => None,
}
}
pub fn is(&self, id: &str) -> bool {
self.choices.iter().any(|choice| choice == id)
}
}
#[derive(Component, Clone, Copy, Debug, Default, PartialEq, Eq)]
pub struct Answerer {
pub agent: Option<Entity>,
}
#[derive(SystemParam)]
pub struct Answerers<'w, 's> {
answerers: Query<'w, 's, &'static Answerer>,
}
impl Answerers<'_, '_> {
pub fn can_answer(&self, agent: Option<Entity>) -> bool {
self.answerers
.iter()
.any(|answerer| answerer.agent.is_none() || answerer.agent == agent)
}
}
pub struct AskPlugin;
#[derive(SystemSet, Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub struct AskSystems;
impl Plugin for AskPlugin {
fn build(&self, app: &mut App) {
app.configure_sets(
AgentPass,
AskSystems.in_set(Stage::Gate).after(CommandSystems),
)
.add_systems(AgentPass, answer_with_text.in_set(AskSystems));
}
}
fn request(mut world: DeferredWorld, context: HookContext) {
let Some(ask) = world.get::<Ask>(context.entity) else {
return;
};
let kind = EffectKind::Custom {
kind: "ask".into(),
payload: serde_json::to_value(ask).unwrap_or_default(),
};
world
.commands()
.entity(context.entity)
.insert_if_new(EffectRequest {
kind,
handler: None,
});
}
fn answer_with_text(
mut typed: Query<(&Incoming, &ForAgent, &mut Decision), (ReadyToStart, IncomingMessages)>,
mut asks: Query<
(Entity, &Ask, &ForAgent, &Effect, &mut Decision),
(PluginEffects, Without<Outcome>, Without<Settled>),
>,
) {
for (incoming, &ForAgent(agent), mut taken) in &mut typed {
let typed_by_user = incoming.author.kind == AuthorKind::User;
if !typed_by_user || incoming.command().is_some() || taken.is_decided() {
continue;
}
let open = asks
.iter()
.filter(|(_, question, to, _, decision)| {
question.free_text && !question.inline && to.0 == agent && !decision.is_decided()
})
.min_by_key(|(_, _, _, effect, _)| effect.id())
.map(|(ask, ..)| ask);
let Some(Ok((.., mut decision))) = open.map(|ask| asks.get_mut(ask)) else {
continue;
};
let answer = Answer::text(incoming.text()).outcome();
decision.decide(Decider::Person, answer);
taken.decide(Decider::plugin("rig-harness"), Outcome::custom("answer"));
}
}
#[cfg(test)]
mod tests;