rig_memory/
recent.rs

1//! Recent messages: which newest messages of a conversation a request
2//! sends. A [`Recent`] rule says where the messages it keeps start; [`view`] moves that
3//! start past results whose calls it cut off, so every result it keeps
4//! still follows its call.
5//!
6//! ```
7//! use std::sync::Arc;
8//!
9//! use rig_core::completion::Message;
10//! use rig_memory::recent::{LastMessages, view};
11//!
12//! let messages = [Message::user("1"), Message::assistant("2"), Message::user("3")].map(Arc::new);
13//! assert_eq!(view(&messages, &LastMessages(2)), &messages[1..]);
14//! ```
15
16use std::sync::Arc;
17
18use rig_core::completion::Message;
19use rig_core::conversation::pairing::safe_start;
20use rig_core::wasm_compat::{WasmCompatSend, WasmCompatSync};
21
22use crate::{HeuristicTokenCounter, TokenCounter};
23
24/// Where the newest messages a request sends start.
25pub trait Recent: WasmCompatSend + WasmCompatSync {
26    /// The index of the first message of `messages` to keep, at most
27    /// `messages.len()`.
28    fn start(&self, messages: &[Arc<Message>]) -> usize;
29}
30
31/// The newest so many messages.
32#[derive(Clone, Copy, Debug, PartialEq, Eq)]
33pub struct LastMessages(pub usize);
34
35impl Recent for LastMessages {
36    fn start(&self, messages: &[Arc<Message>]) -> usize {
37        messages.len().saturating_sub(self.0)
38    }
39}
40
41/// The newest messages within `max_tokens`, as `counter` counts them:
42/// walking back from the newest, it stops before the first message that
43/// does not fit.
44#[derive(Clone, Copy, Debug)]
45pub struct TokenBudget<C: TokenCounter = HeuristicTokenCounter> {
46    /// The budget.
47    pub max_tokens: usize,
48    /// How a message is counted.
49    pub counter: C,
50}
51
52impl TokenBudget {
53    /// A budget of `max_tokens`, counted by the default
54    /// [`HeuristicTokenCounter`].
55    pub fn new(max_tokens: usize) -> Self {
56        Self {
57            max_tokens,
58            counter: HeuristicTokenCounter::default(),
59        }
60    }
61}
62
63impl<C: TokenCounter> Recent for TokenBudget<C> {
64    fn start(&self, messages: &[Arc<Message>]) -> usize {
65        let mut budget = self.max_tokens;
66        let mut start = messages.len();
67        for (index, message) in messages.iter().enumerate().rev() {
68            let cost = self.counter.count(message);
69            if cost > budget {
70                break;
71            }
72            budget -= cost;
73            start = index;
74        }
75        start
76    }
77}
78
79/// The messages of `messages` that `recent` keeps, from its start moved
80/// forward by [`safe_start`] so no result loses its call.
81pub fn view<'m>(messages: &'m [Arc<Message>], recent: &dyn Recent) -> &'m [Arc<Message>] {
82    let start = safe_start(messages, recent.start(messages));
83    messages.get(start..).unwrap_or_default()
84}
85
86#[cfg(test)]
87mod tests;
rig_memory/recent.rs

rig_memory/
recent.rs
