headless_lms_chatbot/chatbot_tools/
output_limits.rs

1//! Bounds on how much a tool call may hand back, so that no tool can produce an output the
2//! conversation cannot store.
3//!
4//! `chatbot_conversation_message_tool_outputs.output` is `VARCHAR(131072)`, and `record_tool_call`
5//! writes the call and its output in one transaction: an oversized output does not degrade the
6//! answer, it rolls the call back and ends the turn. [truncate_tool_output] is the backstop that
7//! makes that unreachable; [CappedList] is what tools use so the backstop rarely has to fire.
8
9use serde::Serialize;
10use std::borrow::Cow;
11
12use headless_lms_utils::strings::truncate_utf8_at_boundary;
13
14/// The most a single tool output may carry, wrapper and instructions included.
15///
16/// Postgres counts `VARCHAR(131072)` in characters and this is a byte budget, so staying under the
17/// column is guaranteed rather than merely likely. The gap to the column is also the headroom that
18/// keeps one result from eating a turn's whole context window.
19const MAX_TOOL_OUTPUT_BYTES: usize = 100_000;
20
21/// What a truncated tool output leaves out, phrased for the model rather than for a log.
22const TRUNCATION_INSTRUCTION: &str = "This result was too large to return in full and was cut off \
23    mid-way, so it may end in the middle of a value. Treat it as a partial view: say so before \
24    answering from it, do not claim anything about what a complete result would have contained, \
25    and prefer narrowing the call (a single id, fewer facets) over reasoning from the fragment.";
26
27/// Cuts `output` down to what a tool output row can hold, returning the text to send and, when it
28/// had to cut, the instruction telling the model what it is looking at.
29///
30/// The notice is returned separately rather than appended so the caller can put it in the
31/// instructions block: inside the output delimiters it would read as part of the data it is
32/// warning about.
33pub(crate) fn truncate_tool_output(output: &str) -> (Cow<'_, str>, Option<&'static str>) {
34    if output.len() <= MAX_TOOL_OUTPUT_BYTES {
35        return (Cow::Borrowed(output), None);
36    }
37    (
38        Cow::Borrowed(truncate_utf8_at_boundary(output, MAX_TOOL_OUTPUT_BYTES)),
39        Some(TRUNCATION_INSTRUCTION),
40    )
41}
42
43/// How much of a list was left out of a tool output.
44#[derive(Serialize)]
45pub(crate) struct ListTruncation {
46    shown: usize,
47    total: usize,
48}
49
50/// A list a tool caps before serializing it, so one long list cannot crowd out the fields beside
51/// it and so the model is never handed a partial list that looks complete.
52///
53/// Prefer aggregating over capping where the rows are repetitive: a cap answers "what are the
54/// first N" when the question was usually "how many".
55#[derive(Serialize)]
56pub(crate) struct CappedList<T> {
57    items: Vec<T>,
58    /// Absent when everything fit, which is the common case and the one where an extra field
59    /// would only invite the model to comment on it.
60    #[serde(skip_serializing_if = "Option::is_none")]
61    truncated: Option<ListTruncation>,
62}
63
64impl<T> CappedList<T> {
65    /// Keeps at most `max_items` of `items`, recording what that left out.
66    pub(crate) fn new(mut items: Vec<T>, max_items: usize) -> Self {
67        let total = items.len();
68        if total <= max_items {
69            return Self {
70                items,
71                truncated: None,
72            };
73        }
74        items.truncate(max_items);
75        Self {
76            items,
77            truncated: Some(ListTruncation {
78                shown: max_items,
79                total,
80            }),
81        }
82    }
83
84    /// Whether anything was left out, for a tool that wants to say so in its instructions as well
85    /// as in its data.
86    pub(crate) fn is_truncated(&self) -> bool {
87        self.truncated.is_some()
88    }
89}
90
91/// Reads as the rows that survived the cap, so a tool can inspect them to word its instructions
92/// without the truncation bookkeeping getting in the way.
93impl<T> std::ops::Deref for CappedList<T> {
94    type Target = [T];
95
96    fn deref(&self) -> &Self::Target {
97        &self.items
98    }
99}
100
101#[cfg(test)]
102mod tests {
103    use super::*;
104
105    #[test]
106    fn an_output_within_the_budget_is_returned_untouched() {
107        let (output, notice) = truncate_tool_output("small");
108        assert_eq!(output, "small");
109        assert!(notice.is_none());
110    }
111
112    /// The budget is in bytes and the cut has to land on a char boundary, so a multi-byte
113    /// character straddling the limit must not be split into invalid UTF-8.
114    #[test]
115    fn an_oversized_output_is_cut_to_the_budget_on_a_char_boundary() {
116        let output = "ä".repeat(MAX_TOOL_OUTPUT_BYTES);
117        let (truncated, notice) = truncate_tool_output(&output);
118        assert!(truncated.len() <= MAX_TOOL_OUTPUT_BYTES);
119        assert!(truncated.chars().all(|c| c == 'ä'));
120        assert!(notice.is_some());
121    }
122
123    #[test]
124    fn a_capped_list_reports_only_what_it_left_out() {
125        let untruncated = CappedList::new(vec![1, 2, 3], 3);
126        assert!(!untruncated.is_truncated());
127        assert_eq!(
128            serde_json::to_string(&untruncated).unwrap(),
129            r#"{"items":[1,2,3]}"#
130        );
131
132        let truncated = CappedList::new(vec![1, 2, 3, 4], 2);
133        assert!(truncated.is_truncated());
134        assert_eq!(
135            serde_json::to_string(&truncated).unwrap(),
136            r#"{"items":[1,2],"truncated":{"shown":2,"total":4}}"#
137        );
138    }
139}
headless_lms_chatbot/chatbot_tools/output_limits.rs

headless_lms_chatbot/chatbot_tools/
output_limits.rs
