diff --git a/crates/zeta/src/zeta.rs b/crates/zeta/src/zeta.rs index 2cf39588a02663..0eefc78f6600e9 100644 --- a/crates/zeta/src/zeta.rs +++ b/crates/zeta/src/zeta.rs @@ -53,6 +53,18 @@ const BYTES_PER_TOKEN_GUESS: usize = 3; /// for the model to specify insertions. const BUFFER_EXCERPT_BYTE_LIMIT: usize = (2048 * 2 / 3) * BYTES_PER_TOKEN_GUESS; +/// Note that this is not the limit for the overall prompt, just for the inputs to the template +/// instantiated in `crates/collab/src/llm.rs`. +const TOTAL_BYTE_LIMIT: usize = BUFFER_EXCERPT_BYTE_LIMIT * 2; + +/// Maximum number of events to include in the prompt. +const MAX_EVENT_COUNT: usize = 16; + +/// Maximum number of string bytes in a single event. Arbitrarily choosing this to be 4x the size of +/// equally splitting up the the remaining bytes after the largest possible buffer excerpt. +const PER_EVENT_BYTE_LIMIT: usize = + (TOTAL_BYTE_LIMIT - BUFFER_EXCERPT_BYTE_LIMIT) / MAX_EVENT_COUNT * 4; + actions!(edit_prediction, [ClearHistory]); #[derive(Copy, Clone, Default, Debug, PartialEq, Eq, Hash)] @@ -235,8 +247,6 @@ impl Zeta { } fn push_event(&mut self, event: Event) { - const MAX_EVENT_COUNT: usize = 16; - if let Some(Event::BufferChange { new_snapshot: last_new_snapshot, timestamp: last_timestamp, @@ -338,7 +348,8 @@ impl Zeta { let (input_excerpt, excerpt_range) = prompt_for_excerpt(&snapshot, cursor_point, cursor_offset)?; - let input_events = prompt_for_events(events.iter()); + let chars_remaining = TOTAL_BYTE_LIMIT.saturating_sub(input_excerpt.len()); + let input_events = prompt_for_events(events.iter(), chars_remaining); // Note that input_outline is not currently used in prompt generation and so // is not counted towards TOTAL_BYTE_LIMIT. @@ -1085,14 +1096,26 @@ fn prompt_for_excerpt( anyhow::Ok((builder.to_string(), editable_range)) } -fn prompt_for_events<'a>(events: impl Iterator) -> String { +fn prompt_for_events<'a>( + events: impl Iterator, + mut bytes_remaining: usize, +) -> String { let mut result = String::new(); for event in events { if !result.is_empty() { result.push('\n'); result.push('\n'); } - result.push_str(&event.to_prompt()); + let event_string = event.to_prompt(); + let len = event_string.len(); + if len > PER_EVENT_BYTE_LIMIT { + continue; + } + if len > bytes_remaining { + break; + } + bytes_remaining -= len; + result.push_str(&event_string); } result }