rig_tui/history.rs
1//! The prompt history.
2
3use std::sync::Arc;
4
5use bevy_ecs::prelude::*;
6use bevy_log::warn;
7use rig_harness::prelude::{Files, PluginData};
8
9use super::jsonl::{self, Kept};
10
11/// The file of the history in its [`PromptHistory::file`].
12const HISTORY_FILE: &str = "tui/history.jsonl";
13
14/// The prompts Up and Down browse, oldest first, kept across sessions in
15/// the user's [`PluginData`] as `tui/history.jsonl` (`file: None` saves
16/// nothing). Prompts are kept as typed, slash commands included.
17///
18/// The view loads it in `PreStartup`, appends each prompt [`Sent`] in
19/// [`HistorySystems::Record`], and in [`HistorySystems::Save`] adds to the
20/// file those recorded since that are still in `entries`, keeping the
21/// newest [`Self::KEPT`]. A plugin filters `entries` between the two, such
22/// as to keep secrets out.
23///
24/// ```no_run
25/// use rig_harness::prelude::*;
26/// use rig_tui::{HistorySystems, PromptHistory};
27///
28/// /// Prompts that look like they hold a key are never kept.
29/// fn scrub(mut history: ResMut<PromptHistory>) {
30/// history.entries.retain(|prompt| !prompt.contains("sk-"));
31/// }
32///
33/// fn build(app: &mut App) {
34/// app.add_systems(Update, scrub.after(HistorySystems::Record).before(HistorySystems::Save));
35/// }
36/// ```
37#[derive(Resource, Default)]
38pub struct PromptHistory {
39 /// The prompts, oldest first.
40 pub entries: Vec<String>,
41 /// The files prompts are saved in, as `tui/history.jsonl`, one JSON
42 /// string per line.
43 pub file: Option<Arc<dyn Files>>,
44 /// The prompts recorded since the last save.
45 unsaved: Vec<String>,
46}
47
48impl PromptHistory {
49 /// How many prompts are kept.
50 pub const KEPT: usize = 500;
51
52 /// The newest prompts kept in `files`, which new ones are saved to.
53 pub fn load(files: Arc<dyn Files>) -> Self {
54 Self {
55 entries: jsonl::tail(&Kept(files.clone(), HISTORY_FILE), Self::KEPT),
56 file: Some(files),
57 unsaved: Vec::new(),
58 }
59 }
60}
61
62/// Loads the history kept in the user's [`PluginData`], before any
63/// plugin's `Startup` system edits it.
64pub(crate) fn load(data: Option<Res<PluginData>>, mut history: ResMut<PromptHistory>) {
65 if let Some(data) = data {
66 *history = PromptHistory::load(data.user.clone());
67 }
68}
69
70/// Where the terminal view records the prompts sent and saves them, in
71/// `Update`, in this order.
72#[derive(SystemSet, Debug, Clone, Copy, PartialEq, Eq, Hash)]
73pub enum HistorySystems {
74 /// Appends the prompts sent this frame to [`PromptHistory::entries`].
75 Record,
76 /// Saves the prompts recorded that are still there.
77 Save,
78}
79
80/// A prompt sent, which [`HistorySystems::Record`] adds to the
81/// [`PromptHistory`] unless it repeats the last one: the input writes one
82/// for each prompt it sends, and a plugin writes one for a prompt it sent
83/// for the user, such as from a panel of its own.
84#[derive(Message, Clone, Debug, PartialEq, Eq)]
85pub struct Sent(pub String);
86
87/// Adds each prompt sent to the history, unless it repeats the last one.
88pub(crate) fn record(mut sent: MessageReader<Sent>, mut history: ResMut<PromptHistory>) {
89 for Sent(text) in sent.read() {
90 let text = text.trim();
91 if text.is_empty() || history.entries.last().is_some_and(|last| last == text) {
92 continue;
93 }
94 history.entries.push(text.to_owned());
95 history.unsaved.push(text.to_owned());
96 }
97}
98
99/// Adds the prompts recorded since the last save that are still kept to
100/// the file, then keeps the newest.
101pub(crate) fn save(mut history: ResMut<PromptHistory>) {
102 if history.unsaved.is_empty() && history.entries.len() <= PromptHistory::KEPT {
103 return;
104 }
105 let unsaved = std::mem::take(&mut history.unsaved);
106 let kept = unsaved.iter().filter(|text| history.entries.contains(text));
107 if let Some(file) = &history.file
108 && let Err(failure) = jsonl::append(&Kept(file.clone(), HISTORY_FILE), kept)
109 {
110 warn!("could not save the prompt history: {failure}");
111 }
112 let excess = history.entries.len().saturating_sub(PromptHistory::KEPT);
113 history.entries.drain(..excess);
114}
115
116#[cfg(test)]
117mod tests;