kwaak 0.1.6

Run a team of autonomous agents on your code, right from your terminal
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138

use ratatui::prelude::*;

use crate::{
    chat::Chat,
    chat_message::{ChatMessage, ChatRole},
};

mod message_styles {
    use super::{Color, Modifier, Style};

    pub const USER: Style = Style::new().fg(Color::Cyan).add_modifier(Modifier::ITALIC);

    pub const ASSISTANT: Style = Style::new()
        .fg(Color::Rgb(200, 160, 255))
        .add_modifier(Modifier::BOLD);

    pub const SYSTEM: Style = Style::new().fg(Color::DarkGray).add_modifier(Modifier::DIM);

    pub const TOOL_DONE: Style = Style::new().fg(Color::Green).add_modifier(Modifier::DIM);

    pub const TOOL_CALLED: Style = Style::new().fg(Color::DarkGray).add_modifier(Modifier::DIM);

    pub const COMMAND: Style = Style::new()
        .fg(Color::LightMagenta)
        .add_modifier(Modifier::BOLD);
}
#[allow(clippy::trivially_copy_pass_by_ref)]
pub fn get_style_and_prefix(role: &ChatRole) -> (&'static str, Style) {
    match role {
        ChatRole::User => ("▶ ", message_styles::USER),
        ChatRole::Assistant => ("✦ ", message_styles::ASSISTANT),
        ChatRole::System => ("ℹ ", message_styles::SYSTEM),
        ChatRole::Tool => ("⚙ ", message_styles::TOOL_DONE), // unused
        ChatRole::Command => ("» ", message_styles::COMMAND),
    }
}

// TODO: Maybe have tool state just on the message?
pub fn format_chat_message<'a>(current_chat: &Chat, message: &'a ChatMessage) -> Text<'a> {
    let (prefix, style) = get_style_and_prefix(message.role());

    // Render markdown first
    let mut rendered_text = tui_markdown::from_str(message.content());

    // Prepend the styled prefix to the first line
    if let Some(first_line) = rendered_text.lines.first_mut() {
        first_line.spans.insert(0, Span::styled(prefix, style));
    }

    // Remove background color from all lines
    // tui_markdown adds background colors and underlines to some headers
    rendered_text.lines.iter_mut().for_each(|line| {
        line.spans.iter_mut().for_each(|span| {
            span.style = span
                .style
                .bg(Color::Reset)
                .remove_modifier(Modifier::UNDERLINED);
        });
    });

    if let Some(swiftide::chat_completion::ChatMessage::Assistant(.., Some(tool_calls))) =
        message.original()
    {
        if !message.content().is_empty() {
            rendered_text.push_line(Line::from("\n\n"));
        }
        for tool_call in tool_calls {
            if tool_call.name() == "stop" {
                continue;
            }
            let is_done = current_chat.is_tool_call_completed(tool_call.id());
            let tool_call_text = format_tool_call(tool_call);
            let tool_prefix = "⚙ ";

            if is_done {
                // add a suffix checkmark on the end
                let checkmark = " ✓";
                rendered_text.lines.push(Line::styled(
                    [tool_prefix, &tool_call_text, checkmark].join(" "),
                    message_styles::TOOL_DONE,
                ));
            } else {
                rendered_text.lines.push(Line::styled(
                    [tool_prefix, &tool_call_text].join(" "),
                    message_styles::TOOL_CALLED,
                ));
            }
        }
    }

    if !message.role().is_assistant() {
        for line in &mut rendered_text.lines {
            for span in &mut line.spans {
                span.style = style;
            }
        }
    }

    rendered_text
}

fn format_tool_call(tool_call: &swiftide::chat_completion::ToolCall) -> String {
    // If args, parse them as a json value, then if its just one, render only the value, otherwise
    // limit the output to 20 characters
    let formatted_args = tool_call.args().and_then(|args| {
        if let Ok(parsed) = serde_json::from_str::<serde_json::Value>(args) {
            if let Some(obj) = parsed.as_object() {
                if obj.is_empty() {
                    return None;
                }
                if obj.keys().count() == 1 {
                    let key = obj.keys().next().unwrap();
                    let val = obj[key].as_str().unwrap_or_default();

                    if val.len() > 20 {
                        return Some(format!("{} ...", &val[..20]));
                    }

                    return Some(val.to_string());
                }
                if args.len() > 20 {
                    return Some(format!("{} ...", &args[..20]));
                }

                return Some(args.to_string());
            }
            None
        } else {
            None
        }
    });

    if let Some(args) = formatted_args {
        format!("calling tool `{}` with `{}`", tool_call.name(), args)
    } else {
        format!("calling tool `{}`", tool_call.name())
    }
}