Skip to main content

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;