kwaak 0.3.0

Run a team of autonomous agents on your code, right from your terminal
Documentation 
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
139
140
141
142
143
144
145
146
147
148
149
150
151

use std::sync::Arc;

use async_trait::async_trait;
use swiftide::chat_completion;
use tokio::sync::mpsc;
use tokio_util::task::AbortOnDropHandle;
use uuid::Uuid;

use crate::{
    chat_message::ChatMessage,
    commands::{CommandResponse, Responder},
};

use super::ui_event::UIEvent;

/// Handles responses from commands application wide
///
/// Basically converts command responses into UI events
/// The responder is send with commands so that the backend has a way to communicate with the
/// frontend, without knowing about the frontend
///
/// Only one is expected to be running at a time
///
/// TODO: If only used in app, singleton is not needed
#[derive(Debug)]
pub struct AppCommandResponder {
    // ui_tx: mpsc::UnboundedSender<UIEvent>,
    tx: mpsc::UnboundedSender<CommandResponse>,
    _handle: AbortOnDropHandle<()>,
}

#[derive(Debug, Clone)]
pub struct AppCommandResponderForChatId {
    inner: mpsc::UnboundedSender<CommandResponse>,
    uuid: uuid::Uuid,
}

impl AppCommandResponder {
    pub fn spawn_for(ui_tx: mpsc::UnboundedSender<UIEvent>) -> AppCommandResponder {
        tracing::info!("Initializing app command responder");
        let (tx, mut rx) = mpsc::unbounded_channel();
        let handle = tokio::spawn(async move {
            while let Some(response) = rx.recv().await {
                tracing::debug!("[RESPONDER] Received response: {:?}", response);
                let ui_event = match response {
                    CommandResponse::Chat(uuid, msg) => UIEvent::ChatMessage(uuid, msg.into()),
                    CommandResponse::Activity(uuid, state) => UIEvent::ActivityUpdate(uuid, state),
                    CommandResponse::RenameChat(uuid, name) => UIEvent::RenameChat(uuid, name),
                    CommandResponse::Completed(uuid) => UIEvent::CommandDone(uuid),
                    CommandResponse::BackendMessage(uuid, msg) => {
                        UIEvent::ChatMessage(uuid, ChatMessage::new_system(&msg))
                    }
                };

                if let Err(err) = ui_tx.send(ui_event) {
                    tracing::error!("Failed to send response to ui: {:#}", err);
                }
            }
            tracing::info!("App command responder shutting down");
        });

        // Create the app responder, spawn a task to handle responses, the once cell returns a
        // clone of the responder without the rx
        AppCommandResponder {
            tx,
            _handle: AbortOnDropHandle::new(handle),
        }
    }

    #[must_use]
    pub fn for_chat_id(&self, uuid: Uuid) -> Arc<dyn Responder> {
        Arc::new(AppCommandResponderForChatId {
            inner: self.tx.clone(),
            uuid,
        }) as Arc<dyn Responder>
    }
}

#[async_trait]
impl Responder for AppCommandResponderForChatId {
    fn send(&self, response: CommandResponse) {
        tracing::debug!("[RESPONDER SENDER] Sending response: {:?}", response);
        let response = response.with_uuid(self.uuid);
        if let Err(err) = self.inner.send(response) {
            tracing::error!("Failed to send response for command: {:?}", err);
        }
    }

    fn system_message(&self, message: &str) {
        self.send(CommandResponse::Chat(
            self.uuid,
            chat_completion::ChatMessage::new_system(message),
        ));
    }

    fn update(&self, state: &str) {
        self.send(CommandResponse::Activity(self.uuid, state.into()));
    }

    fn rename(&self, name: &str) {
        self.send(CommandResponse::RenameChat(self.uuid, name.into()));
    }

    fn agent_message(&self, message: chat_completion::ChatMessage) {
        self.send(CommandResponse::Chat(self.uuid, message));
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use tokio::sync::mpsc;
    use uuid::Uuid;

    const TEST_UUID: Uuid = Uuid::from_u128(0x1234_5678_90ab_cdef_1234_5678_90ab_cdef);

    #[tokio::test]
    async fn test_app_responder() {
        let (ui_tx, mut ui_rx) = mpsc::unbounded_channel();
        let app = AppCommandResponder::spawn_for(ui_tx);

        let responder = app.for_chat_id(TEST_UUID);

        responder.system_message("Test message");

        let Some(ui_event) = ui_rx.recv().await else {
            panic!("No UI event received");
        };

        match ui_event {
            UIEvent::ChatMessage(received_uuid, received_message) => {
                assert_eq!(received_uuid, TEST_UUID);
                assert_eq!(received_message.content(), "Test message");
                assert!(received_message.role().is_system());
            }
            _ => panic!("Unexpected UI event received"),
        }

        responder.send(CommandResponse::Completed(Uuid::new_v4()));

        // Verify the UI event is received
        if let Some(ui_event) = ui_rx.recv().await {
            match ui_event {
                UIEvent::CommandDone(received_uuid) => assert_eq!(received_uuid, TEST_UUID),
                _ => panic!("Unexpected UI event received"),
            }
        } else {
            panic!("No UI event received");
        }
    }
}