use bevy_app::prelude::*;
use bevy_ecs::prelude::*;
use bevy_reflect::prelude::*;
use rig_core::conversation::{Conversation, MessageId};
use rig_core::effect::Decider;
use rig_ecs::stage::{Decision, Outcome, Settled};
use crate::ask::{Answer, Ask};
use crate::commands::RunCommand;
use crate::main_agent::{ShownAgent, keep_shown};
#[derive(Reflect, Clone, Debug, Default, PartialEq, Eq)]
#[reflect(Clone, Debug, Default, PartialEq)]
pub struct PickItem {
pub label: String,
pub command: String,
pub detail: Option<String>,
pub checked: bool,
}
#[derive(Component, Reflect, Clone, Debug, Default, PartialEq, Eq)]
#[reflect(Component, Clone, Debug, Default, PartialEq)]
pub struct Picker {
pub agent: Option<Entity>,
pub title: String,
pub body: String,
pub items: Vec<PickItem>,
pub selected: usize,
pub free_text: bool,
pub multiple: bool,
}
impl Picker {
pub fn new(title: impl Into<String>) -> Self {
Self {
title: title.into(),
..Self::default()
}
}
pub fn for_agent(mut self, agent: Entity) -> Self {
self.agent = Some(agent);
self
}
pub fn body(mut self, body: impl Into<String>) -> Self {
self.body = body.into();
self
}
pub fn item(mut self, label: impl Into<String>, command: impl Into<String>) -> Self {
self.items.push(PickItem {
label: label.into(),
command: command.into(),
detail: None,
checked: false,
});
self
}
pub fn checked(mut self, checked: bool) -> Self {
if let Some(item) = self.items.last_mut() {
item.checked = checked;
}
self
}
pub fn multiple(mut self) -> Self {
self.multiple = true;
self
}
pub fn free_text(mut self) -> Self {
self.free_text = true;
self
}
}
#[derive(Reflect, Clone, Debug, Default, PartialEq, Eq)]
#[reflect(Clone, Debug, Default, PartialEq)]
pub struct Chosen {
pub items: Vec<usize>,
pub text: Option<String>,
}
impl Chosen {
pub fn item(index: usize) -> Self {
Self::items([index])
}
pub fn items(indices: impl IntoIterator<Item = usize>) -> Self {
Self {
items: indices.into_iter().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
}
fn positions(&self) -> Vec<usize> {
let mut positions = self.items.clone();
positions.sort_unstable();
positions.dedup();
positions
}
}
#[derive(Component, Reflect, Clone, Copy, Debug, Default, PartialEq, Eq)]
#[reflect(Component, Clone, Debug, Default, PartialEq)]
pub struct Highlighted(pub usize);
#[derive(Component, Clone, Debug, PartialEq)]
pub enum PickPreview {
Text(String),
Conversation {
conversation: Conversation,
at: Option<MessageId>,
},
}
#[derive(EntityEvent, Reflect, Clone, Debug)]
#[reflect(Event, Clone, Debug)]
pub struct Pick {
pub entity: Entity,
pub chosen: Option<Chosen>,
#[reflect(ignore, default = "person")]
pub by: Decider,
}
fn person() -> Decider {
Decider::Person
}
impl Pick {
pub fn person(entity: Entity, chosen: Option<Chosen>) -> Self {
Self {
entity,
chosen,
by: Decider::Person,
}
}
}
#[derive(EntityEvent, Reflect, Clone, Debug)]
#[reflect(Event, Clone, Debug)]
pub struct Picked {
pub entity: Entity,
pub chosen: Option<Chosen>,
}
pub struct PickPlugin;
impl Plugin for PickPlugin {
fn build(&self, app: &mut App) {
app.init_resource::<ShownAgent>()
.add_systems(First, keep_shown)
.add_observer(pick);
}
}
fn pick(
pick: On<Pick>,
pickers: Query<&Picker>,
mut asks: Query<(&Ask, &mut Decision), (Without<Outcome>, Without<Settled>)>,
shown: Res<ShownAgent>,
mut commands: Commands,
) {
let (entity, chosen) = (pick.entity, pick.chosen.clone());
if let Ok(picker) = pickers.get(entity) {
let items = chosen.as_ref().map(Chosen::positions).unwrap_or_default();
let items = items.iter().filter_map(|index| picker.items.get(*index));
if let Some(agent) = picker.agent.or(shown.0) {
for item in items.filter(|item| !item.command.is_empty()) {
commands.trigger(RunCommand::new(agent, item.command.clone()));
}
}
commands.trigger(Picked { entity, chosen });
commands.entity(entity).try_despawn();
return;
}
let Ok((ask, mut decision)) = asks.get_mut(entity) else {
return;
};
let answer = match chosen {
Some(chosen) => {
let ids = chosen
.positions()
.into_iter()
.map(|index| ask.choices.get(index));
let ids: Option<Vec<String>> = ids.map(|c| c.map(|c| c.id.clone())).collect();
let typed = chosen.text.is_some();
ids.filter(|ids| ask.multiple || ids.len() == 1 || (ids.is_empty() && typed))
.map(|choices| {
let answer = Answer {
choices,
text: chosen.text,
};
answer.outcome()
})
}
None => Some(Outcome::refused("The user closed the question.")),
};
if let Some(answer) = answer {
decision.decide(pick.by.clone(), answer);
}
}
#[cfg(test)]
mod tests;