use bevy_ecs::prelude::*;
use bevy_reflect::Reflect;
use rig_core::message::{self, AssistantContent, AssistantMessage, ToolResultContent, UserContent};
use serde::{Deserialize, Serialize};
use std::collections::BTreeMap;
use super::binary::{BinaryAssets, BinaryError, PartSource};
use crate::agent::{MessageParts, Role, Utterance};
#[derive(Component, Debug, Clone, PartialEq, Serialize, Deserialize, Reflect)]
#[reflect(Component)]
pub enum ContentPart {
Text(#[reflect(remote = super::reflect::TextPartReflect)] message::Text),
Image(ImagePart),
Audio(AudioPart),
Video(VideoPart),
Document(DocumentPart),
ToolCall(#[reflect(remote = super::reflect::ToolCallPartReflect)] message::ToolCall),
Reasoning(#[reflect(remote = super::reflect::ReasoningPartReflect)] message::Reasoning),
Opaque(#[reflect(remote = super::reflect::OpaquePartReflect)] message::Opaque),
ToolResult {
#[reflect(remote = crate::agent::reflect::ToolCallIdReflect)]
call: message::CallId,
#[reflect(remote = crate::agent::reflect::ToolNameReflect)]
name: message::ToolName,
#[serde(default, skip_serializing_if = "std::ops::Not::not")]
is_error: bool,
},
Json(#[reflect(remote = super::reflect::JsonPartReflect)] serde_json::Value),
}
#[derive(Component, Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Reflect)]
#[reflect(Component)]
pub enum RequestPartEdit {
Text(String),
Remove,
}
#[derive(Component, Debug, Clone, Copy, PartialEq, Eq, Reflect)]
#[relationship(relationship_target = EditedBy)]
#[reflect(Component)]
pub struct EditTarget(pub Entity);
#[derive(Component, Debug, Default, Reflect)]
#[relationship_target(relationship = EditTarget)]
#[reflect(Component)]
pub struct EditedBy(Vec<Entity>);
#[derive(Component, Debug, Clone, Default, PartialEq, Serialize, Deserialize, Reflect)]
#[reflect(Component)]
pub struct AssistantHead {
#[reflect(remote = super::reflect::OriginReflect)]
pub origin: Option<message::Origin>,
#[reflect(remote = super::reflect::StopReflect)]
pub stop: Option<message::StopReason>,
}
impl AssistantHead {
pub fn of(message: &AssistantMessage) -> Self {
Self {
origin: message.origin.clone(),
stop: message.stop.clone(),
}
}
pub fn message(&self, content: Vec<AssistantContent>) -> AssistantMessage {
AssistantMessage::default()
.with_origin(self.origin.clone())
.with_stop(self.stop.clone())
.with_content(content)
}
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Reflect)]
pub struct ImagePart {
pub source: PartSource,
#[reflect(remote = super::reflect::ImageMediaReflect)]
pub media_type: Option<message::ImageMediaType>,
#[reflect(remote = super::reflect::NativeReflect)]
pub native: Option<message::Native>,
#[reflect(remote = super::reflect::ImageDetailReflect)]
pub detail: Option<message::ImageDetail>,
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Reflect)]
pub struct AudioPart {
pub source: PartSource,
#[reflect(remote = super::reflect::AudioMediaReflect)]
pub media_type: Option<message::AudioMediaType>,
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Reflect)]
pub struct VideoPart {
pub source: PartSource,
#[reflect(remote = super::reflect::VideoMediaReflect)]
pub media_type: Option<message::VideoMediaType>,
#[reflect(remote = super::reflect::PartParamsReflect)]
pub additional_params: Option<serde_json::Value>,
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Reflect)]
pub struct DocumentPart {
pub source: PartSource,
#[reflect(remote = super::reflect::DocumentMediaReflect)]
pub media_type: Option<message::DocumentMediaType>,
#[reflect(remote = super::reflect::PartParamsReflect)]
pub additional_params: Option<serde_json::Value>,
}
#[derive(Component, Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Reflect)]
#[reflect(Component)]
#[serde(rename_all = "snake_case")]
pub enum ToolResultStatus {
Ok,
Error,
Refused,
Skipped,
Denied,
WrongFamily,
}
pub const TOOL_RESULT_LIMIT_MARKER: &str = "\n[โฆ {omitted} bytes omitted โฆ]\n";
#[derive(Component, Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Reflect)]
#[reflect(Component)]
pub struct ToolResultLimit {
pub max_bytes: usize,
pub marker: String,
}
impl ToolResultLimit {
#[must_use]
pub fn new(max_bytes: usize) -> Self {
Self {
max_bytes,
marker: TOOL_RESULT_LIMIT_MARKER.to_owned(),
}
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, thiserror::Error)]
pub enum ContentError {
#[error(transparent)]
Binary(#[from] BinaryError),
#[error("content graph is missing an entity or required component")]
Missing,
#[error("content part has an invalid type or role")]
Shape,
}
type PreparedParts = Vec<(ContentPart, Vec<ContentPart>)>;
fn image(assets: &mut BinaryAssets, value: message::Image) -> Result<ImagePart, ContentError> {
Ok(ImagePart {
source: assets.intern(value.data)?,
media_type: value.media_type,
detail: value.detail,
native: value.native,
})
}
fn user(
assets: &mut BinaryAssets,
value: UserContent,
) -> Result<(ContentPart, Vec<ContentPart>), ContentError> {
let mut children = Vec::new();
let part = match value {
UserContent::Text(value) => ContentPart::Text(value),
UserContent::Image(value) => ContentPart::Image(image(assets, value)?),
UserContent::Audio(value) => ContentPart::Audio(AudioPart {
source: assets.intern(value.data)?,
media_type: value.media_type,
}),
UserContent::Video(value) => ContentPart::Video(VideoPart {
source: assets.intern(value.data)?,
media_type: value.media_type,
additional_params: value.additional_params,
}),
UserContent::Document(value) => ContentPart::Document(DocumentPart {
source: assets.intern(value.data)?,
media_type: value.media_type,
additional_params: value.additional_params,
}),
UserContent::ToolResult(value) => {
children = value
.content
.into_iter()
.map(|item| {
Ok(match item {
ToolResultContent::Text(text) => ContentPart::Text(text),
ToolResultContent::Image(value) => {
ContentPart::Image(image(assets, value)?)
}
ToolResultContent::Json { value } => ContentPart::Json(value),
})
})
.collect::<Result<_, ContentError>>()?;
ContentPart::ToolResult {
call: value.call,
name: value.name,
is_error: value.is_error,
}
}
};
Ok((part, children))
}
fn prepare(
assets: &mut BinaryAssets,
parts: MessageParts,
) -> Result<(Role, Option<AssistantHead>, PreparedParts), ContentError> {
Ok(match parts {
MessageParts::User { content } => (
Role::User,
None,
content
.into_iter()
.map(|part| user(assets, part))
.collect::<Result<_, _>>()?,
),
MessageParts::Assistant(message) => (
Role::Assistant,
Some(AssistantHead::of(&message)),
message
.content
.into_iter()
.map(|part| {
let part = match part {
AssistantContent::Text(value) => ContentPart::Text(value),
AssistantContent::Image(value) => ContentPart::Image(image(assets, value)?),
AssistantContent::ToolCall(value) => ContentPart::ToolCall(value),
AssistantContent::Reasoning(value) => ContentPart::Reasoning(value),
AssistantContent::Opaque(value) => ContentPart::Opaque(value),
_ => return Err(ContentError::Shape),
};
Ok((part, Vec::new()))
})
.collect::<Result<_, ContentError>>()?,
),
})
}
fn spawn_parts(world: &mut World, parent: Entity, parts: PreparedParts) {
for (part, children) in parts {
let entity = world.spawn(ChildOf(parent)).insert(part).id();
for child in children {
world.spawn(ChildOf(entity)).insert(child);
}
}
}
pub fn write_message(
world: &mut World,
utterance: Entity,
parts: MessageParts,
) -> Result<(), ContentError> {
if world.get::<Utterance>(utterance).is_none() {
return Err(ContentError::Missing);
}
world.init_resource::<BinaryAssets>();
let (role, head, parts) = prepare(&mut world.resource_mut::<BinaryAssets>(), parts)?;
replace_parts(world, utterance, role, head, parts);
Ok(())
}
fn replace_parts(
world: &mut World,
utterance: Entity,
role: Role,
head: Option<AssistantHead>,
parts: PreparedParts,
) {
let old: Vec<_> = world
.get::<Children>(utterance)
.map(|children| children.iter().collect())
.unwrap_or_default();
for child in old {
world.despawn(child);
}
let mut entity = world.entity_mut(utterance);
entity.insert(role).remove::<AssistantHead>();
if let Some(head) = head {
entity.insert(head);
}
spawn_parts(world, utterance, parts);
}
fn ordered<'a>(
get: &impl Fn(Entity) -> Option<EntityRef<'a>>,
parent: Entity,
) -> Result<Vec<Entity>, ContentError> {
let parent = get(parent).ok_or(ContentError::Missing)?;
let children = parent.get::<Children>();
let mut ordered = Vec::new();
for child in children.into_iter().flat_map(|c| c.iter()) {
let entity = get(child).ok_or(ContentError::Missing)?;
if entity.get::<ContentPart>().is_none() {
return Err(ContentError::Shape);
}
ordered.push(child);
}
Ok(ordered)
}
fn read_image(assets: &BinaryAssets, value: &ImagePart) -> Result<message::Image, ContentError> {
Ok(message::Image {
data: assets.resolve(&value.source)?,
media_type: value.media_type.clone(),
detail: value.detail.clone(),
native: value.native.clone(),
})
}
fn read_edited_part<'a>(
get: &impl Fn(Entity) -> Option<EntityRef<'a>>,
entity: Entity,
nested: bool,
edits: &BTreeMap<Entity, RequestPartEdit>,
) -> Result<Option<ContentPart>, ContentError> {
let entity_ref = get(entity).ok_or(ContentError::Missing)?;
let mut value = entity_ref
.get::<ContentPart>()
.ok_or(ContentError::Shape)?
.clone();
if matches!(value, ContentPart::ToolResult { .. }) {
if nested {
return Err(ContentError::Shape);
}
} else if entity_ref
.get::<Children>()
.is_some_and(|children| !children.is_empty())
{
return Err(ContentError::Shape);
}
match edits.get(&entity) {
Some(RequestPartEdit::Remove) => return Ok(None),
Some(RequestPartEdit::Text(text)) => match &mut value {
ContentPart::Text(part) => part.text.clone_from(text),
_ => return Err(ContentError::Shape),
},
None => {}
}
Ok(Some(value))
}
fn to_user<'a>(
get: &impl Fn(Entity) -> Option<EntityRef<'a>>,
assets: &BinaryAssets,
entity: Entity,
value: ContentPart,
edits: &BTreeMap<Entity, RequestPartEdit>,
) -> Result<UserContent, ContentError> {
Ok(match value {
ContentPart::Text(value) => UserContent::Text(value),
ContentPart::Image(value) => UserContent::Image(read_image(assets, &value)?),
ContentPart::Audio(value) => UserContent::Audio(message::Audio {
data: assets.resolve(&value.source)?,
media_type: value.media_type,
}),
ContentPart::Video(value) => UserContent::Video(message::Video {
data: assets.resolve(&value.source)?,
media_type: value.media_type,
additional_params: value.additional_params,
}),
ContentPart::Document(value) => UserContent::Document(message::Document {
data: assets.resolve(&value.source)?,
media_type: value.media_type,
additional_params: value.additional_params,
}),
ContentPart::ToolResult {
call,
name,
is_error,
} => UserContent::ToolResult(message::ToolResult {
call,
name,
is_error,
content: ordered(get, entity)?
.into_iter()
.map(|child| read_edited_part(get, child, true, edits))
.collect::<Result<Vec<_>, _>>()?
.into_iter()
.flatten()
.map(|part| {
Ok(match part {
ContentPart::Text(value) => ToolResultContent::Text(value),
ContentPart::Image(value) => {
ToolResultContent::Image(read_image(assets, &value)?)
}
ContentPart::Json(value) => ToolResultContent::Json { value },
_ => return Err(ContentError::Shape),
})
})
.collect::<Result<Vec<_>, ContentError>>()
.and_then(non_empty)?,
}),
_ => return Err(ContentError::Shape),
})
}
pub fn read_message(world: &World, utterance: Entity) -> Result<MessageParts, ContentError> {
let empty = BinaryAssets::default();
let assets = world.get_resource::<BinaryAssets>().unwrap_or(&empty);
read_message_from(
&|entity| world.get_entity(entity).ok(),
assets,
utterance,
&BTreeMap::new(),
)
}
#[derive(bevy_ecs::system::SystemParam)]
pub struct ContentGraph<'w, 's> {
entities: Query<'w, 's, EntityRef<'static>, Without<bevy_ecs::resource::IsResource>>,
assets: Option<Res<'w, BinaryAssets>>,
}
impl ContentGraph<'_, '_> {
pub fn message_with(
&self,
utterance: Entity,
edits: &BTreeMap<Entity, RequestPartEdit>,
) -> Result<MessageParts, ContentError> {
if edits.is_empty() {
return self.message(utterance);
}
self.message(utterance)?;
let empty = BinaryAssets::default();
read_message_from(
&|entity| self.entities.get(entity).ok(),
self.assets.as_deref().unwrap_or(&empty),
utterance,
edits,
)
}
pub fn target_utterance(&self, target: Entity) -> Result<Entity, ContentError> {
let part = self
.entities
.get(target)
.map_err(|_| ContentError::Missing)?;
if !part.contains::<ContentPart>() {
return Err(ContentError::Shape);
}
let parent = part.get::<ChildOf>().ok_or(ContentError::Missing)?.parent();
let owner = self
.entities
.get(parent)
.map_err(|_| ContentError::Missing)?;
if owner.contains::<Utterance>() {
return Ok(parent);
}
if !matches!(
owner.get::<ContentPart>(),
Some(ContentPart::ToolResult { .. })
) {
return Err(ContentError::Shape);
}
let utterance = owner
.get::<ChildOf>()
.ok_or(ContentError::Missing)?
.parent();
if !self
.entities
.get(utterance)
.map_err(|_| ContentError::Missing)?
.contains::<Utterance>()
{
return Err(ContentError::Shape);
}
Ok(utterance)
}
pub fn message(&self, utterance: Entity) -> Result<MessageParts, ContentError> {
let empty = BinaryAssets::default();
let assets = self.assets.as_deref().unwrap_or(&empty);
read_message_from(
&|entity| self.entities.get(entity).ok(),
assets,
utterance,
&BTreeMap::new(),
)
}
}
fn read_message_from<'a>(
get: &impl Fn(Entity) -> Option<EntityRef<'a>>,
assets: &BinaryAssets,
utterance: Entity,
edits: &BTreeMap<Entity, RequestPartEdit>,
) -> Result<MessageParts, ContentError> {
let entity = get(utterance).ok_or(ContentError::Missing)?;
if entity.get::<Utterance>().is_none() {
return Err(ContentError::Missing);
}
let role = entity.get::<Role>().ok_or(ContentError::Missing)?;
let values = ordered(get, utterance)?
.into_iter()
.map(|entity| {
read_edited_part(get, entity, false, edits).map(|part| part.map(|part| (entity, part)))
})
.collect::<Result<Vec<_>, _>>()?
.into_iter()
.flatten()
.collect::<Vec<_>>();
match role {
Role::User => {
if entity.get::<AssistantHead>().is_some() {
return Err(ContentError::Shape);
}
MessageParts::user(
values
.into_iter()
.map(|(entity, value)| to_user(get, assets, entity, value, edits))
.collect::<Result<_, _>>()?,
)
}
Role::Assistant => MessageParts::assistant(
entity
.get::<AssistantHead>()
.ok_or(ContentError::Missing)?
.message(
values
.into_iter()
.map(|(_, value)| {
Ok(match value {
ContentPart::Text(value) => AssistantContent::Text(value),
ContentPart::Image(value) => {
AssistantContent::Image(read_image(assets, &value)?)
}
ContentPart::ToolCall(value) => AssistantContent::ToolCall(value),
ContentPart::Reasoning(value) => AssistantContent::Reasoning(value),
ContentPart::Opaque(value) => AssistantContent::Opaque(value),
_ => return Err(ContentError::Shape),
})
})
.collect::<Result<_, ContentError>>()?,
),
),
}
}
fn non_empty<T>(items: Vec<T>) -> Result<Vec<T>, ContentError> {
if items.is_empty() {
return Err(ContentError::Shape);
}
Ok(items)
}
pub fn collect_binary_assets(
world: &mut World,
pins: impl IntoIterator<Item = super::binary::BinaryId>,
) -> Result<(), BinaryError> {
let mut roots: Vec<_> = pins.into_iter().collect();
for part in world.query::<&ContentPart>().iter(world) {
let source = match part {
ContentPart::Image(part) => &part.source,
ContentPart::Audio(part) => &part.source,
ContentPart::Video(part) => &part.source,
ContentPart::Document(part) => &part.source,
_ => continue,
};
if let PartSource::Binary { id, .. } = source {
roots.push(*id);
}
}
match world.get_resource_mut::<BinaryAssets>() {
Some(mut assets) => assets.retain(roots),
None if roots.is_empty() => Ok(()),
None => Err(BinaryError::Missing),
}
}
pub(crate) fn spawn_deferred(
commands: &mut Commands,
assets: &mut BinaryAssets,
parent: Entity,
message: MessageParts,
) -> Result<Entity, ContentError> {
spawn_deferred_with(commands, assets, parent, message, Vec::new())
}
pub(crate) fn spawn_deferred_with(
commands: &mut Commands,
assets: &mut BinaryAssets,
parent: Entity,
message: MessageParts,
statuses: Vec<ToolResultStatus>,
) -> Result<Entity, ContentError> {
let (role, head, parts) = prepare(assets, message)?;
let mut utterance = commands.spawn((Utterance, role, ChildOf(parent)));
if let Some(head) = head {
utterance.insert(head);
}
let entity = utterance.id();
commands.queue(move |world: &mut World| {
spawn_parts(world, entity, parts);
stamp_statuses(world, entity, &statuses);
});
Ok(entity)
}
fn stamp_statuses(world: &mut World, utterance: Entity, statuses: &[ToolResultStatus]) {
if statuses.is_empty() {
return;
}
let results: Vec<Entity> = world
.get::<Children>(utterance)
.into_iter()
.flat_map(|children| children.iter())
.filter(|child| {
matches!(
world.get::<ContentPart>(*child),
Some(ContentPart::ToolResult { .. })
)
})
.collect();
for (entity, status) in results.into_iter().zip(statuses) {
world.entity_mut(entity).insert(*status);
}
}
pub(crate) fn replace_deferred(
commands: &mut Commands,
assets: &mut BinaryAssets,
utterance: Entity,
message: MessageParts,
) -> Result<(), ContentError> {
let (role, head, parts) = prepare(assets, message)?;
commands.queue(move |world: &mut World| {
if world.get::<Utterance>(utterance).is_none() {
return;
}
replace_parts(world, utterance, role, head, parts);
});
Ok(())
}