use std::collections::{HashMap, HashSet};
use std::sync::Arc;
use bevy_ecs::lifecycle::HookContext;
use bevy_ecs::prelude::*;
use bevy_ecs::system::SystemParam;
use bevy_ecs::world::DeferredWorld;
use bevy_reflect::prelude::*;
use rig_core::catalog::ModelSpec;
use rig_core::completion::Message;
use rig_core::conversation::{Aside, AsideId, Author, AuthorKind, Entry, Halt};
use rig_core::effect::{EffectKind, HandlerKey};
use rig_core::error::{ErrorKind, ErrorReport};
use rig_core::message::UserContent;
use serde::{Deserialize, Serialize};
use web_time::SystemTime;
use super::agent::{
ActiveTurn, Agent, CallOf, Conversation, EffectParent, Notice, TurnOf, TurnOutcome,
};
use super::model::Connection;
#[cfg(doc)]
use super::stage::Stage;
use super::stage::{
AgentWork, Decision, Effect, ForAgent, Holds, Outcome, ReadyToStart, Settled, Started,
SupersededBy, moved,
};
use super::tasks::Wake;
use super::turn::{Exiting, TurnCall, spawn_turn_call, take_outcome};
#[derive(Reflect, Clone, Copy, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum DeliveryMode {
#[default]
Steer,
Queue,
#[serde(rename = "note")]
Inform,
}
#[derive(Reflect, Clone, Debug, Serialize, Deserialize)]
#[reflect(opaque, Clone, Debug)]
pub struct Attachment {
pub label: String,
pub content: UserContent,
}
#[derive(Reflect, Clone, Debug, Serialize, Deserialize)]
#[reflect(Clone, Debug)]
#[serde(rename_all = "snake_case")]
pub enum MessagePart {
Text(String),
Attachment(Attachment),
}
#[derive(Component, Reflect, Clone, Debug, Serialize, Deserialize)]
#[component(on_add = sent)]
#[reflect(Component, Clone, Debug)]
#[require(Effect, Name = Name::new("message"))]
pub struct Incoming {
pub parts: Vec<MessagePart>,
#[reflect(ignore)]
pub author: Author,
pub mode: DeliveryMode,
pub literal: bool,
}
impl Incoming {
pub fn new(text: impl Into<String>, mode: DeliveryMode) -> Self {
Self {
parts: vec![MessagePart::Text(text.into())],
author: Author::user(),
mode,
literal: false,
}
}
pub fn text(&self) -> &str {
texts(&self.parts).next().unwrap_or_default()
}
pub fn text_mut(&mut self) -> &mut String {
let first = self
.parts
.iter()
.position(|part| matches!(part, MessagePart::Text(_)));
let (at, part) = match first {
Some(at) => (at, self.parts.remove(at)),
None => (self.parts.len(), MessagePart::Text(String::new())),
};
match self.parts.insert_mut(at, part) {
MessagePart::Text(text) => text,
MessagePart::Attachment(attachment) => &mut attachment.label,
}
}
fn is_empty(&self) -> bool {
self.parts.iter().all(|part| match part {
MessagePart::Text(text) => text.trim().is_empty(),
MessagePart::Attachment(_) => false,
})
}
#[must_use]
pub fn with_author(self, author: Author) -> Self {
Self { author, ..self }
}
#[must_use]
pub fn literal(self) -> Self {
Self {
literal: true,
..self
}
}
pub fn command(&self) -> Option<&str> {
let user = self.author.kind == AuthorKind::User && !self.literal;
user.then(|| self.text().trim_start().strip_prefix('/'))
.flatten()
}
}
#[derive(EntityEvent, Reflect, Clone, Debug)]
#[reflect(Event, Clone, Debug)]
pub struct Deliver {
pub entity: Entity,
pub message: Incoming,
}
impl Deliver {
pub fn new(agent: Entity, text: impl Into<String>, mode: DeliveryMode) -> Self {
Self {
entity: agent,
message: Incoming::new(text, mode),
}
}
#[must_use]
pub fn with_author(self, author: Author) -> Self {
Self {
message: self.message.with_author(author),
..self
}
}
#[must_use]
pub fn literal(self) -> Self {
Self {
message: self.message.literal(),
..self
}
}
}
#[derive(Message, Clone, Debug)]
pub struct Recalled {
pub agent: Entity,
pub text: String,
pub why: Option<String>,
}
#[derive(Component, Reflect, Clone, Copy, Debug)]
#[reflect(Component, Clone, Debug)]
pub struct DeliveredIn(pub Entity);
#[derive(Component, Clone, Debug)]
#[component(on_insert = moved::<Unsent>)]
pub struct Unsent {
pub ended: Option<TurnOutcome>,
}
const INBOX_MARK: &str = "rig.inbox";
#[derive(Component, Clone, Copy, Debug)]
pub(crate) struct InboxMark(AsideId);
fn mark(conversation: &mut Conversation, message: &Incoming) -> InboxMark {
let after = conversation.entries().next_back().map(|entry| entry.id);
let value = serde_json::to_value(message).unwrap_or_default();
let aside = Aside::new(INBOX_MARK, Author::rig(), value).after(after);
InboxMark(conversation.add_aside(aside))
}
pub(crate) fn kept(conversation: &Conversation) -> Vec<(InboxMark, Incoming)> {
let marks = conversation
.asides()
.filter(|(_, aside)| aside.kind == INBOX_MARK);
let read = |aside: &Aside| Incoming::deserialize(&aside.value).ok();
marks
.filter_map(|(id, aside)| Some((InboxMark(id), read(aside)?)))
.collect()
}
pub(crate) fn unmark(
answered: On<Insert<Outcome>>,
marked: Query<(&InboxMark, &ForAgent)>,
mut agents: Query<&mut Conversation>,
) {
if let Ok((&InboxMark(id), &ForAgent(agent))) = marked.get(answered.entity)
&& let Ok(mut conversation) = agents.get_mut(agent)
{
conversation.remove_aside(id);
}
}
pub(crate) fn on_deliver(
deliver: On<Deliver>,
agents: Query<(), With<Agent>>,
mut commands: Commands,
) {
if agents.contains(deliver.entity) {
commands.spawn((deliver.message.clone(), ForAgent(deliver.entity)));
}
}
fn sent(mut world: DeferredWorld, context: HookContext) {
let entity = context.entity;
let parent = world
.get::<ForAgent>(entity)
.and_then(|agent| world.get::<EffectParent>(agent.0))
.map(|parent| parent.0);
if let Some(parent) = parent
&& let Some(mut effect) = world.get_mut::<Effect>(entity)
&& effect.parent.is_none()
{
effect.parent = Some(parent);
}
if let Some(wake) = world.get_resource::<Wake>() {
wake.wake();
}
}
pub(crate) fn accept(
ready: Query<(Entity, &Effect, &Incoming, Option<&ForAgent>), ReadyToStart>,
mut agents: Query<(&mut Conversation, Option<&Connection>, Has<ActiveTurn>), With<Agent>>,
mut commands: Commands,
mut notices: MessageWriter<Notice>,
) {
let mut ready: Vec<_> = ready.iter().collect();
ready.sort_unstable_by_key(|(_, effect, ..)| effect.id());
let at = SystemTime::now();
let mut starting = HashMap::new();
for (entity, effect, message, agent) in ready {
let agent = agent.map(|agent| agent.0);
let Some((agent, (mut conversation, connection, busy))) =
agent.and_then(|agent| Some((agent, agents.get_mut(agent).ok()?)))
else {
let why = ErrorReport::new(ErrorKind::Internal, "a message needs an agent");
commands.entity(entity).try_insert_if_new(Outcome(Err(why)));
continue;
};
if message.is_empty() {
commands
.entity(entity)
.try_insert_if_new(Outcome::custom("empty"));
continue;
}
let joins = starting.get(&agent).copied();
if busy && joins.is_none() {
let mark = mark(&mut conversation, message);
commands.entity(entity).try_insert((Started { at }, mark));
continue;
}
let spec = connection.map(|connection| &*connection.spec);
let turn = match joins {
None if message.mode == DeliveryMode::Inform => None,
None => Some(commands.spawn(TurnOf(agent)).id()),
joins => joins,
};
commands.entity(entity).try_insert(Started { at });
let conversation = &mut *conversation;
let into = (turn, &mut commands);
deliver(
agent,
spec,
[(entity, message)],
conversation,
&mut notices,
into,
);
match (turn, joins) {
(None, _) => {
conversation.halt(Halt::Kept);
}
(Some(turn), None) => {
starting.insert(agent, turn);
commands.spawn((
TurnCall,
CallOf(turn),
ForAgent(agent),
Effect::with_parent_id(Some(effect.id())),
));
}
(Some(_), Some(_)) => {}
}
}
}
pub(crate) fn carry_on_accepted(
accepted: Query<(&Incoming, &ForAgent), (Added<Started>, Without<Outcome>)>,
mut agents: Query<(&mut Conversation, Option<&Connection>, &ActiveTurn)>,
turns: Query<(&TurnOutcome, Option<&Holds>)>,
inboxes: Inboxes,
mut commands: Commands,
mut notices: MessageWriter<Notice>,
) {
let mut carried = HashSet::new();
for (message, &ForAgent(agent)) in &accepted {
if message.mode == DeliveryMode::Inform || !carried.insert(agent) {
continue;
}
let Ok((mut conversation, connection, active)) = agents.get_mut(agent) else {
continue;
};
let turn = active.turn();
let Ok((TurnOutcome::Answered(_), holds)) = turns.get(turn) else {
continue;
};
let Some(next) = carry_on(inboxes.of(agent)) else {
continue;
};
let spec = connection.map(|connection| &*connection.spec);
let conversation = &mut *conversation;
let into = (Some(turn), &mut commands);
deliver(agent, spec, next, conversation, &mut notices, into);
take_outcome(&mut commands, turn, holds);
spawn_turn_call(&mut commands, turn, agent);
}
}
#[derive(SystemParam)]
pub(crate) struct Inboxes<'w, 's> {
work: Query<'w, 's, &'static AgentWork>,
accepted: Query<
'w,
's,
(
Entity,
&'static Effect,
&'static Incoming,
&'static Decision,
),
(With<Started>, Without<Outcome>, Without<Settled>),
>,
on_the_way: Query<
'w,
's,
(&'static Incoming, &'static Decision),
(
Without<Started>,
Without<Outcome>,
Without<Settled>,
Without<SupersededBy>,
),
>,
}
impl Inboxes<'_, '_> {
pub(crate) fn of(&self, agent: Entity) -> Vec<(Entity, &Incoming)> {
let work = self.work.get(agent).into_iter().flat_map(AgentWork::iter);
let accepted = self.accepted.iter_many(work).flatten();
let mut waiting: Vec<_> = accepted
.filter(|(.., decision)| !decision.is_decided())
.collect();
waiting.sort_unstable_by_key(|(_, effect, ..)| effect.id());
waiting
.into_iter()
.map(|(entity, _, message, _)| (entity, message))
.collect()
}
pub(crate) fn coming(&self, agent: Entity) -> bool {
let work = self.work.get(agent).into_iter().flat_map(AgentWork::iter);
self.on_the_way
.iter_many(work)
.flatten()
.any(|(message, decision)| {
message.mode != DeliveryMode::Inform && !decision.is_decided()
})
}
}
pub(crate) fn carry_on(waiting: Vec<(Entity, &Incoming)>) -> Option<Vec<(Entity, &Incoming)>> {
let steering = waiting
.iter()
.any(|(_, message)| message.mode == DeliveryMode::Steer);
let carries = match steering {
true => DeliveryMode::Steer,
false => DeliveryMode::Queue,
};
let next: Vec<_> = waiting
.into_iter()
.filter(|(_, message)| {
matches!(message.mode, DeliveryMode::Inform) || message.mode == carries
})
.collect();
next.iter()
.any(|(_, message)| message.mode == carries)
.then_some(next)
}
pub(crate) fn unsent_on_turn_end(
end: On<Remove<TurnOf>>,
turns: Query<(&TurnOf, Option<&TurnOutcome>)>,
inboxes: Inboxes,
exiting: Option<Res<Exiting>>,
mut commands: Commands,
) {
let Ok((&TurnOf(agent), ended)) = turns.get(end.entity) else {
return;
};
if exiting.is_some() {
return;
}
for (message, _) in inboxes.of(agent) {
let ended = ended.cloned();
commands.entity(message).try_insert(Unsent { ended });
}
}
pub(crate) fn settle_unsent(
unsent: Query<
(Entity, &Effect, &Incoming, &ForAgent, &Decision),
(With<Unsent>, Without<Outcome>, Without<Holds>),
>,
mut agents: Query<(&mut Conversation, Option<&Connection>, Has<ActiveTurn>)>,
mut commands: Commands,
mut recalled: MessageWriter<Recalled>,
mut notices: MessageWriter<Notice>,
) {
let mut unsent: Vec<_> = unsent.iter().collect();
unsent.sort_unstable_by_key(|(_, effect, ..)| effect.id());
let mut typed: HashMap<Entity, Vec<&str>> = HashMap::new();
for (entity, _, message, &ForAgent(agent), decision) in unsent {
let Ok((mut conversation, connection, busy)) = agents.get_mut(agent) else {
continue;
};
if decision.is_decided() {
continue;
}
if busy {
commands.entity(entity).try_remove::<Unsent>();
} else if message.author.kind == AuthorKind::User {
commands
.entity(entity)
.try_insert_if_new(Outcome::custom("recalled"));
typed.entry(agent).or_default().push(message.text().trim());
} else {
let spec = connection.map(|connection| &*connection.spec);
let conversation = &mut *conversation;
let into = (None, &mut commands);
deliver(
agent,
spec,
[(entity, message)],
conversation,
&mut notices,
into,
);
conversation.halt(Halt::Kept);
}
}
for (agent, typed) in typed {
recalled.write(Recalled {
agent,
text: typed.join("\n\n"),
why: None,
});
}
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct MessageRecord {
pub text: String,
pub author: Author,
pub mode: DeliveryMode,
}
impl MessageRecord {
pub const KIND: &str = "message";
pub fn read(kind: &EffectKind) -> Option<Self> {
match kind {
EffectKind::Custom { kind, payload } if &**kind == Self::KIND => {
serde_json::from_value(payload.clone()).ok()
}
_ => None,
}
}
}
pub(crate) fn message_kind(message: &Incoming) -> (HandlerKey, EffectKind) {
let record = MessageRecord {
text: texts(&message.parts).collect::<Vec<_>>().join("\n"),
author: message.author.clone(),
mode: message.mode,
};
let kind = EffectKind::Custom {
kind: Arc::from(MessageRecord::KIND),
payload: serde_json::to_value(record).unwrap_or_default(),
};
(HandlerKey::new(MessageRecord::KIND), kind)
}
pub(crate) fn deliver<'a>(
agent: Entity,
spec: Option<&ModelSpec>,
messages: impl IntoIterator<Item = (Entity, &'a Incoming)>,
conversation: &mut Conversation,
notices: &mut MessageWriter<Notice>,
(into, commands): (Option<Entity>, &mut Commands),
) {
for (entity, message) in messages {
push_messages(agent, spec, [message], conversation, notices);
let mut delivered = commands.entity(entity);
delivered.try_insert_if_new(Outcome::custom("delivered"));
if let Some(turn) = into {
delivered.try_insert(DeliveredIn(turn));
}
}
}
fn texts(parts: &[MessagePart]) -> impl Iterator<Item = &str> {
parts.iter().filter_map(|part| match part {
MessagePart::Text(text) => Some(text.as_str()),
MessagePart::Attachment(_) => None,
})
}
pub(crate) fn push_messages<'a>(
agent: Entity,
spec: Option<&ModelSpec>,
messages: impl IntoIterator<Item = &'a Incoming>,
conversation: &mut Conversation,
notices: &mut MessageWriter<Notice>,
) {
for message in messages {
let mut header = message.author.header();
let mut content = Vec::with_capacity(message.parts.len());
let mut refused = Vec::new();
for part in &message.parts {
let attachment = match part {
MessagePart::Text(text) => {
let text = text.trim();
let text = match header.take() {
Some(header) if text.is_empty() => header,
Some(header) => format!("{header}\n{text}"),
None => text.to_owned(),
};
if !text.is_empty() {
content.push(UserContent::text(text));
}
continue;
}
MessagePart::Attachment(attachment) => attachment,
};
let Attachment {
label,
content: item,
} = attachment;
let why = match (item, spec) {
(UserContent::Image(_), Some(spec)) if !spec.input.image => Some(format!(
"{} does not read images, so {label} is sent as its name only.",
spec.display_name
)),
(UserContent::Image(_), None) => Some(format!(
"No model is connected, so {label} is sent as its name only."
)),
_ => None,
};
match why {
Some(note) => {
notices.write(Notice::info(agent, note));
refused.push(label.as_str());
}
None => content.push(item.clone()),
}
}
if content.is_empty() {
content.extend(refused.into_iter().map(UserContent::text));
}
if let Some(header) = header {
content.insert(0, UserContent::text(header));
}
let entry = Entry::new(Message::User { content });
let entry = match message.author == Author::user() {
true => entry,
false => entry.by(message.author.clone()),
};
conversation.push(entry);
}
}
#[cfg(test)]
mod tests;