Skip to main content

llimorse_chat/
lib.rs

1//! Framework for a llimorse-based chat application
2
3#![warn(missing_docs)]
4#![warn(clippy::missing_docs_in_private_items)]
5
6pub mod agent;
7pub mod history;
8pub mod log;
9pub mod tools;
10pub mod ui;
11
12use agent::ChatAgent;
13use anyhow::Result;
14use futures::FutureExt;
15pub use history::ChatHistory;
16use llimorse::Agent;
17use llimorse::line_format::ChatMessage;
18use std::sync::{Arc, Mutex};
19use std::thread::{self, JoinHandle};
20use tokio::sync::mpsc;
21use tokio::time::{self, Duration};
22pub use tools::{SubagentNotifier, UserToolGate};
23pub use ui::UiState;
24
25/// The application state
26pub struct App<I: UiState> {
27    /// Agent thread running concurrently
28    agent_thread: Option<JoinHandle<()>>,
29
30    /// UI state
31    ui: I,
32
33    /// Notifications to the agent
34    agent_notifications: mpsc::UnboundedSender<agent::Notification>,
35
36    /// Notification to the UI
37    ui_notifications: mpsc::UnboundedReceiver<ui::Notification>,
38}
39
40impl<I: UiState> App<I> {
41    /// Create a new application state around `agent`, pre-feeding the chat log with `history`.
42    pub fn new_with_history<F: FnOnce(&Agent, Arc<Mutex<ChatHistory>>) -> Result<I>>(
43        mut agent: Agent,
44        history: &[ChatMessage],
45        ui_notifications: ui::NotificationChannel,
46        create_ui: F,
47    ) -> Result<Self> {
48        let mut chat_history = ChatHistory::default();
49        for message in history {
50            chat_history.push_raw(&agent, message);
51        }
52        chat_history.force_resolve_unresolved_tool_calls(&mut agent);
53
54        let chat_history = Arc::new(Mutex::new(chat_history));
55
56        let ui = create_ui(&agent, Arc::clone(&chat_history))?;
57
58        let (agent_notifications, recv_agent_notifications) = mpsc::unbounded_channel();
59        let ui_notification_sender = ui_notifications.sender();
60
61        let agent_thread = thread::spawn({
62            move || {
63                tokio::runtime::Builder::new_current_thread()
64                    .enable_all()
65                    .build()
66                    .unwrap()
67                    .block_on(async move {
68                        let mut wba = ChatAgent::new(
69                            chat_history,
70                            recv_agent_notifications,
71                            Arc::new(ui_notification_sender),
72                        );
73                        if let Err(err) = wba.run(agent).await {
74                            panic!("Agent error: {err}");
75                        }
76                    })
77            }
78        });
79
80        Ok(App {
81            agent_thread: Some(agent_thread),
82
83            ui,
84
85            agent_notifications,
86            ui_notifications: ui_notifications.into_receiver(),
87        })
88    }
89
90    /// Create a new application state around `agent`.
91    pub fn new<F: FnOnce(&Agent, Arc<Mutex<ChatHistory>>) -> Result<I>>(
92        agent: Agent,
93        ui_notifications: ui::NotificationChannel,
94        create_ui: F,
95    ) -> Result<Self> {
96        Self::new_with_history(agent, &[], ui_notifications, create_ui)
97    }
98
99    /// Run the application until it finds it should exit.
100    pub async fn run(&mut self) -> Result<()> {
101        // Redraw at a steady cadence, independent of how busy the notification
102        // channel is (a hot channel must not starve the idle tick, which e.g.
103        // animates the spinner). The first tick is immediate, so this also
104        // covers the initial draw.
105        let mut redraw_interval = time::interval(Duration::from_millis(500));
106
107        loop {
108            let event_result = futures::select! {
109                result = self.ui.get_event().fuse() => result.map(Some).map_err(Into::into),
110                _ = redraw_interval.tick().fuse() => {
111                    self.ui.notify(ui::Notification::Update).map_err(Into::into)?;
112                    Ok(None)
113                }
114                notification = self.ui_notifications.recv().fuse() => {
115                    if let Some(notification) = notification {
116                        let exit = matches!(notification, ui::Notification::Exit);
117                        self.ui.notify(notification).map_err(Into::into)?;
118                        if exit {
119                            // TODO: Fix this, it's a bit of a hack here
120                            return Ok(());
121                        }
122                    }
123                    Ok(None)
124                }
125            };
126
127            if let Some(event) = event_result? {
128                match event {
129                    ui::Event::Exit => {
130                        let _ = self.agent_notifications.send(agent::Notification::Exit);
131                        return Ok(());
132                    }
133                    ui::Event::Input(message) => {
134                        self.ui
135                            .notify(ui::Notification::PromptQueued(message.clone()))
136                            .map_err(Into::into)?;
137
138                        let _ = self
139                            .agent_notifications
140                            .send(agent::Notification::QueuePrompt(message));
141                    }
142                    ui::Event::ForceSubmitQueued => {
143                        let _ = self
144                            .agent_notifications
145                            .send(agent::Notification::ForceSubmitQueued);
146                    }
147                    ui::Event::Continue => {
148                        let _ = self.agent_notifications.send(agent::Notification::Continue);
149                    }
150                }
151            }
152        }
153    }
154}
155
156impl<I: UiState> Drop for App<I> {
157    fn drop(&mut self) {
158        if let Some(agent_thread) = self.agent_thread.take() {
159            let _ = self.agent_notifications.send(agent::Notification::Exit);
160            let _ = agent_thread.join();
161        }
162    }
163}