use std::collections::HashMap;
use std::sync::Arc;
use std::time::Duration;
use bevy_app::App;
use bevy_ecs::prelude::*;
use bevy_ecs::system::SystemParam;
use bevy_log::warn;
use ratatui::style::{Style, Stylize};
use ratatui::text::{Line, Span};
use rig_core::completion::{AssistantContent, Message};
use rig_core::message::{Audio, Document, Image, ToolCall, ToolResult, UserContent, Video};
use rig_core::tool::{ExitStatus, KeptEnd, OutputCut, ToolResultContext};
use rig_ecs::prelude::{Aside, AsideId, Author, AuthorKind, MessageId};
use rig_harness::rig_core::conversation::Note;
use super::{diff, markdown};
use crate::Theme;
pub struct ToolCallView<'a> {
pub call: &'a ToolCall,
pub result: Option<&'a ToolResult>,
pub state: CallState,
pub theme: Theme,
pub facts: Option<&'a ToolResultContext>,
pub took: Option<Duration>,
pub output: Option<&'a str>,
pub children: &'a [ChildCall],
}
#[derive(Clone, Debug, PartialEq)]
pub struct ChildCall {
pub depth: usize,
pub call: Option<ToolCall>,
pub name: String,
pub agent: Option<String>,
pub state: CallState,
pub failed: bool,
pub answer: Option<String>,
}
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
pub enum CallState {
Streaming,
Waiting,
Held(Vec<String>),
Running(Option<String>),
Done,
}
pub(crate) const UNNAMED_HOLD: &str = "unnamed hold";
impl<'a> ToolCallView<'a> {
pub fn new(call: &'a ToolCall, result: Option<&'a ToolResult>, state: CallState) -> Self {
Self {
call,
result,
state,
theme: Theme::default(),
facts: None,
took: None,
output: None,
children: &[],
}
}
pub fn name(&self) -> &str {
&self.call.function.name
}
pub fn argument(&self, key: &str) -> Option<&str> {
self.call.function.arguments.get(key)?.as_str()
}
pub fn result_text(&self) -> Option<String> {
self.result.map(result_text)
}
pub fn failed(&self) -> bool {
self.result.is_some_and(|result| result.is_error)
|| self.exit().is_some_and(|exit| !exit.success())
}
fn exit(&self) -> Option<ExitStatus> {
self.facts?.get::<ExitStatus>().ok().flatten()
}
pub fn header(&self, title: impl Into<String>, detail: impl Into<String>) -> Line<'static> {
let detail = detail.into();
let detail = (!detail.is_empty()).then(|| Span::from(detail).dim());
self.header_spans(title, detail.into_iter().collect())
}
pub fn header_spans(
&self,
title: impl Into<String>,
detail: Vec<Span<'static>>,
) -> Line<'static> {
let failure = || match self.exit() {
Some(ExitStatus::Exited(code)) if code != 0 => format!("exit {code}"),
Some(ExitStatus::Signaled) => "killed by a signal".to_owned(),
Some(ExitStatus::TimedOut) => "timed out".to_owned(),
Some(ExitStatus::Killed) => "killed".to_owned(),
_ => "failed".to_owned(),
};
let (band, state) = match (self.result, &self.state) {
(Some(_), _) if self.failed() => (self.theme.call_failed, failure()),
(_, CallState::Held(holds)) => (
self.theme.call_waiting,
format!("held: {}", holds.join(", ")),
),
(Some(_), _) => (self.theme.call_done, String::new()),
(None, CallState::Running(_)) => (self.theme.call_running, "running".to_owned()),
(None, _) => (self.theme.call_waiting, String::new()),
};
let mut spans = vec![Span::from(" "), Span::from(title.into()).bold()];
if !detail.is_empty() {
spans.push(Span::from(" "));
spans.extend(detail);
}
if !state.is_empty() {
spans.push(Span::from(format!(" · {state}")).italic());
}
Line::from(spans).bg(band)
}
pub fn result_lines(&self) -> Vec<Line<'static>> {
let Some(text) = self.result_text() else {
return self.progress_lines();
};
if !self.failed()
&& let Some(lines) = diff::unified_lines(&text, &self.theme)
{
return lines;
}
let (output, marker) = self.cut_marker(&text);
result_rows(output.into_iter(), marker, result_style(self.failed()))
}
fn cut_marker<'t>(&self, text: &'t str) -> (Vec<&'t str>, Option<(&'t str, KeptEnd)>) {
let mut lines: Vec<&str> = text.lines().collect();
let cut = self
.facts
.and_then(|facts| facts.get::<OutputCut>().ok().flatten());
let marker = cut.and_then(|cut| {
let line = match cut.kept {
KeptEnd::Ends if lines.is_empty() => None,
KeptEnd::Ends => Some(lines.remove(0)),
KeptEnd::Start => lines.pop(),
};
Some((line?, cut.kept))
});
(lines, marker)
}
pub fn progress_lines(&self) -> Vec<Line<'static>> {
let CallState::Running(Some(progress)) = &self.state else {
return Vec::new();
};
result_rows(progress.lines(), None, Style::new().yellow())
}
pub fn default_lines(&self) -> Vec<Line<'static>> {
let arguments = serde_json::Value::Object(self.call.function.arguments.clone()).to_string();
let mut lines = vec![self.header(self.name().to_owned(), arguments)];
lines.extend(self.result_lines());
lines
}
}
pub struct EntryView<'a> {
pub plugin: &'a str,
pub title: &'a str,
pub text: &'a str,
}
impl EntryView<'_> {
pub fn default_lines(&self) -> Vec<Line<'static>> {
let mut lines = vec![Line::styled(
format!("⤶ {}", self.title),
Style::new().magenta().bold(),
)];
lines.extend(text_rows(self.text, Style::new().dim()));
lines
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub enum ContentKind {
Image,
Audio,
Video,
Document,
}
impl ContentKind {
pub fn name(self) -> &'static str {
match self {
Self::Image => "image",
Self::Audio => "audio",
Self::Video => "video",
Self::Document => "document",
}
}
}
#[derive(Clone, Copy, Debug)]
pub enum ContentBlock<'a> {
Image(&'a Image),
Audio(&'a Audio),
Video(&'a Video),
Document(&'a Document),
}
impl<'a> ContentBlock<'a> {
pub(crate) fn of_user(content: &'a UserContent) -> Option<Self> {
match content {
UserContent::Image(image) => Some(Self::Image(image)),
UserContent::Audio(audio) => Some(Self::Audio(audio)),
UserContent::Video(video) => Some(Self::Video(video)),
UserContent::Document(document) => Some(Self::Document(document)),
UserContent::Text(_) | UserContent::ToolResult(_) => None,
}
}
pub(crate) fn of_assistant(content: &'a AssistantContent) -> Option<Self> {
match content {
AssistantContent::Image(image) => Some(Self::Image(image)),
_ => None,
}
}
pub fn kind(self) -> ContentKind {
match self {
Self::Image(_) => ContentKind::Image,
Self::Audio(_) => ContentKind::Audio,
Self::Video(_) => ContentKind::Video,
Self::Document(_) => ContentKind::Document,
}
}
}
pub struct ContentView<'a> {
pub message: Option<MessageId>,
pub index: usize,
pub block: ContentBlock<'a>,
pub assistant: bool,
}
impl ContentView<'_> {
pub fn default_lines(&self) -> Vec<Line<'static>> {
let name = self.block.kind().name();
vec![Line::from(format!(" [{name}]")).dim()]
}
}
pub(crate) fn plugin_name(author: &Author) -> &str {
match &author.kind {
AuthorKind::Plugin(name) => name,
_ => "",
}
}
pub(crate) fn result_text(result: &ToolResult) -> String {
result
.content
.iter()
.filter_map(|content| content.as_text())
.collect::<Vec<_>>()
.join("\n")
}
pub(crate) fn result_style(failed: bool) -> Style {
if failed {
Style::new().red()
} else {
Style::new().dim()
}
}
pub const OUTPUT_INDENT: &str = " ";
pub fn text_rows(text: &str, style: Style) -> Vec<Line<'static>> {
result_rows(text.lines(), None, style)
}
fn result_rows<'t>(
lines: impl Iterator<Item = &'t str>,
marker: Option<(&'t str, KeptEnd)>,
style: Style,
) -> Vec<Line<'static>> {
let (before, after) = match marker {
Some((marker, KeptEnd::Ends)) => (Some(marker), None),
Some((marker, KeptEnd::Start)) => (None, Some(marker)),
None => (None, None),
};
let rows = before.into_iter().chain(lines).chain(after);
rows.map(|line| {
let line = line.replace('\t', " ");
Line::styled(format!("{OUTPUT_INDENT}{line}"), style)
})
.collect()
}
pub type RenderToolCall = dyn Fn(&ToolCallView<'_>) -> Vec<Line<'static>> + Send + Sync;
#[derive(Component, Clone)]
pub struct ToolRenderer {
pub tool: String,
pub render: Arc<RenderToolCall>,
}
impl ToolRenderer {
pub fn new(
tool: &str,
render: impl Fn(&ToolCallView<'_>) -> Vec<Line<'static>> + Send + Sync + 'static,
) -> Self {
let (tool, render) = (tool.to_owned(), Arc::new(render));
Self { tool, render }
}
}
pub type RenderEntry = dyn Fn(&EntryView<'_>) -> Vec<Line<'static>> + Send + Sync;
#[derive(Component, Clone)]
pub struct EntryRenderer {
pub plugin: String,
pub render: Arc<RenderEntry>,
}
impl EntryRenderer {
pub fn new(
plugin: &str,
render: impl Fn(&EntryView<'_>) -> Vec<Line<'static>> + Send + Sync + 'static,
) -> Self {
let (plugin, render) = (plugin.to_owned(), Arc::new(render));
Self { plugin, render }
}
}
pub type RenderContent = dyn Fn(&ContentView<'_>) -> Vec<Line<'static>> + Send + Sync;
#[derive(Component, Clone)]
pub struct ContentRenderer {
pub kind: ContentKind,
pub render: Arc<RenderContent>,
}
impl ContentRenderer {
pub fn new(
kind: ContentKind,
render: impl Fn(&ContentView<'_>) -> Vec<Line<'static>> + Send + Sync + 'static,
) -> Self {
let render = Arc::new(render);
Self { kind, render }
}
}
pub type RenderAside = dyn Fn(&AsideView<'_>) -> Vec<Line<'static>> + Send + Sync;
#[derive(Component, Clone)]
pub struct AsideRenderer {
pub kind: String,
pub render: Arc<RenderAside>,
}
impl AsideRenderer {
pub fn new(
kind: &str,
render: impl Fn(&AsideView<'_>) -> Vec<Line<'static>> + Send + Sync + 'static,
) -> Self {
let (kind, render) = (kind.to_owned(), Arc::new(render));
Self { kind, render }
}
}
pub(crate) type Decorated = Vec<(String, std::ops::Range<usize>)>;
#[derive(Component, Clone)]
pub struct CallDecorator {
pub name: String,
pub render: Arc<RenderToolCall>,
}
impl CallDecorator {
pub fn new(
name: &str,
render: impl Fn(&ToolCallView<'_>) -> Vec<Line<'static>> + Send + Sync + 'static,
) -> Self {
let (name, render) = (name.to_owned(), Arc::new(render));
Self { name, render }
}
}
pub type DecorateMessage = dyn Fn(&DecoratedMessage<'_>) -> Vec<Line<'static>> + Send + Sync;
#[derive(Component, Clone)]
pub struct MessageDecorator {
pub name: String,
pub render: Arc<DecorateMessage>,
}
impl MessageDecorator {
pub fn new(
name: &str,
render: impl Fn(&DecoratedMessage<'_>) -> Vec<Line<'static>> + Send + Sync + 'static,
) -> Self {
let (name, render) = (name.to_owned(), Arc::new(render));
Self { name, render }
}
}
pub struct AsideView<'a> {
pub id: AsideId,
pub aside: &'a Aside,
pub theme: Theme,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub enum Role {
User,
Assistant,
Reasoning,
}
impl Role {
pub fn name(self) -> &'static str {
match self {
Self::User => "user",
Self::Assistant => "assistant",
Self::Reasoning => "reasoning",
}
}
}
pub struct MessageView<'a> {
pub role: Role,
pub message: Option<MessageId>,
pub block: usize,
pub text: &'a str,
pub marks: &'a [(usize, Span<'static>)],
pub width: usize,
pub theme: Theme,
}
impl MessageView<'_> {
pub fn default_lines(&self) -> Vec<Line<'static>> {
match self.role {
Role::User => {
let lines = self.text.lines().enumerate();
let rows = lines.map(|(index, line)| {
let prefix = if index == 0 { "› " } else { " " };
let text = line.replace('\t', " ");
Line::from(vec![Span::from(prefix).dim(), Span::from(text).bold()])
});
rows.collect()
}
Role::Assistant => {
markdown::render_marked(self.text, self.marks, &self.theme, Some(self.width))
}
Role::Reasoning => {
let lines = self.text.lines();
let style = Style::new().dim().italic();
let rows = lines.map(|line| Line::styled(line.replace('\t', " "), style));
rows.collect()
}
}
}
}
pub struct DecoratedMessage<'a> {
pub message: &'a Message,
pub id: Option<MessageId>,
pub author: Option<&'a Author>,
}
pub trait AppToolRenderersExt {
fn add_tool_renderer(
&mut self,
tool: &str,
render: impl Fn(&ToolCallView<'_>) -> Vec<Line<'static>> + Send + Sync + 'static,
) -> &mut Self;
fn add_entry_renderer(
&mut self,
plugin: &str,
render: impl Fn(&EntryView<'_>) -> Vec<Line<'static>> + Send + Sync + 'static,
) -> &mut Self;
fn add_content_renderer(
&mut self,
kind: ContentKind,
render: impl Fn(&ContentView<'_>) -> Vec<Line<'static>> + Send + Sync + 'static,
) -> &mut Self;
fn add_aside_renderer(
&mut self,
kind: &str,
render: impl Fn(&AsideView<'_>) -> Vec<Line<'static>> + Send + Sync + 'static,
) -> &mut Self;
fn add_call_decorator(
&mut self,
name: &str,
render: impl Fn(&ToolCallView<'_>) -> Vec<Line<'static>> + Send + Sync + 'static,
) -> &mut Self;
}
fn add_look<T: Component>(app: &mut App, name: String, look: T, same: impl Fn(&T) -> bool) {
let world = app.world_mut();
let mut looks = world.query_filtered::<&T, Without<ChildOf>>();
if looks.iter(world).any(same) {
warn!("look not added: `{name}` has one; insert `Disabled` on it to replace it");
return;
}
world.spawn((Name::new(name), look));
}
impl AppToolRenderersExt for App {
fn add_tool_renderer(
&mut self,
tool: &str,
render: impl Fn(&ToolCallView<'_>) -> Vec<Line<'static>> + Send + Sync + 'static,
) -> &mut Self {
let look = ToolRenderer::new(tool, render);
add_look(self, format!("renderer:{tool}"), look, |other| {
other.tool == tool
});
self
}
fn add_entry_renderer(
&mut self,
plugin: &str,
render: impl Fn(&EntryView<'_>) -> Vec<Line<'static>> + Send + Sync + 'static,
) -> &mut Self {
let look = EntryRenderer::new(plugin, render);
add_look(self, format!("entry:{plugin}"), look, |other| {
other.plugin == plugin
});
self
}
fn add_content_renderer(
&mut self,
kind: ContentKind,
render: impl Fn(&ContentView<'_>) -> Vec<Line<'static>> + Send + Sync + 'static,
) -> &mut Self {
let look = ContentRenderer::new(kind, render);
let name = format!("content:{}", kind.name());
add_look(self, name, look, |other| other.kind == kind);
self
}
fn add_aside_renderer(
&mut self,
kind: &str,
render: impl Fn(&AsideView<'_>) -> Vec<Line<'static>> + Send + Sync + 'static,
) -> &mut Self {
let look = AsideRenderer::new(kind, render);
add_look(self, format!("aside:{kind}"), look, |other| {
other.kind == kind
});
self
}
fn add_call_decorator(
&mut self,
name: &str,
render: impl Fn(&ToolCallView<'_>) -> Vec<Line<'static>> + Send + Sync + 'static,
) -> &mut Self {
let look = CallDecorator::new(name, render);
add_look(self, format!("decorator:{name}"), look, |other| {
other.name == name
});
self
}
}
type Scoped<T> = (Entity, Ref<'static, T>, Option<&'static ChildOf>);
#[derive(SystemParam)]
pub(crate) struct Looks<'w, 's> {
tools: Query<'w, 's, Scoped<ToolRenderer>>,
entries: Query<'w, 's, Scoped<EntryRenderer>>,
contents: Query<'w, 's, Scoped<ContentRenderer>>,
asides: Query<'w, 's, Scoped<AsideRenderer>>,
decorators: Query<'w, 's, Scoped<CallDecorator>>,
message_decorators: Query<'w, 's, Scoped<MessageDecorator>>,
}
fn scoped<'a, T: Component, K, V>(
looks: impl Iterator<Item = (Entity, Ref<'a, T>, Option<&'a ChildOf>)>,
pane: Option<Entity>,
key: impl Fn(&'a T) -> (K, V),
) -> Vec<(K, V)> {
let mut found: Vec<(bool, Entity, &'a T)> = looks
.filter_map(|(entity, look, parent)| match parent.map(ChildOf::parent) {
None => Some((true, entity, look.into_inner())),
Some(of) if Some(of) == pane => Some((false, entity, look.into_inner())),
Some(_) => None,
})
.collect();
found.sort_by_key(|(global, entity, _)| (*global, *entity));
found.into_iter().map(|(.., look)| key(look)).collect()
}
fn first_of<K: Eq + std::hash::Hash, V>(looks: Vec<(K, V)>) -> HashMap<K, V> {
let mut map = HashMap::new();
for (key, look) in looks {
map.entry(key).or_insert(look);
}
map
}
impl Looks<'_, '_> {
pub(crate) fn by_kind(&self, pane: Option<Entity>) -> Renderers<'_> {
let tools = scoped(self.tools.iter(), pane, |look| {
(look.tool.as_str(), &look.render)
});
let entries = scoped(self.entries.iter(), pane, |look| {
(look.plugin.as_str(), &look.render)
});
let contents = scoped(self.contents.iter(), pane, |look| (look.kind, &look.render));
let asides = scoped(self.asides.iter(), pane, |look| {
(look.kind.as_str(), &look.render)
});
let decorators = scoped(self.decorators.iter(), pane, |look| {
(look.name.as_str(), &look.render)
});
let message_decorators = scoped(self.message_decorators.iter(), pane, |look| {
(look.name.as_str(), &look.render)
});
Renderers {
tools: first_of(tools),
entries: first_of(entries),
contents: first_of(contents),
asides: first_of(asides),
decorators,
message_decorators,
}
}
pub(crate) fn changed(&self) -> bool {
self.tools.iter().any(|(_, look, _)| look.is_changed())
|| self.entries.iter().any(|(_, look, _)| look.is_changed())
|| self.contents.iter().any(|(_, look, _)| look.is_changed())
|| self.asides.iter().any(|(_, look, _)| look.is_changed())
|| self.decorators.iter().any(|(_, look, _)| look.is_changed())
|| (self.message_decorators.iter()).any(|(_, look, _)| look.is_changed())
}
}
#[derive(Default)]
pub(crate) struct Renderers<'a> {
pub(crate) tools: HashMap<&'a str, &'a Arc<RenderToolCall>>,
pub(crate) entries: HashMap<&'a str, &'a Arc<RenderEntry>>,
pub(crate) contents: HashMap<ContentKind, &'a Arc<RenderContent>>,
pub(crate) asides: HashMap<&'a str, &'a Arc<RenderAside>>,
pub(crate) decorators: Vec<(&'a str, &'a Arc<RenderToolCall>)>,
pub(crate) message_decorators: Vec<(&'a str, &'a Arc<DecorateMessage>)>,
}
impl Renderers<'_> {
pub(crate) fn call(&self, view: &ToolCallView<'_>) -> Vec<Line<'static>> {
self.call_decorated(view).0
}
pub(crate) fn call_decorated(
&self,
view: &ToolCallView<'_>,
) -> (Vec<Line<'static>>, Decorated) {
let mut lines = match self.tools.get(view.name()) {
Some(render) => render(view),
None => view.default_lines(),
};
let mut decorated = Vec::new();
for (name, decorate) in &self.decorators {
let start = lines.len();
lines.extend(decorate(view));
if lines.len() > start {
decorated.push(((*name).to_owned(), start..lines.len()));
}
}
(lines, decorated)
}
pub(crate) fn entry(&self, view: &EntryView<'_>) -> Vec<Line<'static>> {
match self.entries.get(view.plugin) {
Some(render) => render(view),
None => view.default_lines(),
}
}
pub(crate) fn content(&self, view: &ContentView<'_>) -> Vec<Line<'static>> {
match self.contents.get(&view.block.kind()) {
Some(render) => render(view),
None => view.default_lines(),
}
}
pub(crate) fn aside(&self, view: &AsideView<'_>) -> Option<Vec<Line<'static>>> {
if let Some(note) = Note::of(view.aside) {
let entry = EntryView {
plugin: plugin_name(&view.aside.author),
title: ¬e.title,
text: ¬e.text,
};
return Some(self.entry(&entry));
}
let render = self.asides.get(view.aside.kind.as_str())?;
Some(render(view))
}
pub(crate) fn draws(&self, aside: &Aside) -> bool {
aside.kind == Note::KIND || self.asides.contains_key(aside.kind.as_str())
}
pub(crate) fn message_decorations(
&self,
message: &DecoratedMessage<'_>,
) -> (Vec<Line<'static>>, Decorated) {
let (mut lines, mut decorated) = (Vec::new(), Vec::new());
for (name, decorate) in &self.message_decorators {
let start = lines.len();
lines.extend(decorate(message));
if lines.len() > start {
decorated.push(((*name).to_owned(), start..lines.len()));
}
}
(lines, decorated)
}
}
#[cfg(test)]
mod tests;