diff --git a/crates/zeta_prompt/src/zeta_prompt.rs b/crates/zeta_prompt/src/zeta_prompt.rs index f2bf994ae551b930aba873597dcdd210a6ad7c92..2cc5322db0ce5b2e11002f06c036832357199d97 100644 --- a/crates/zeta_prompt/src/zeta_prompt.rs +++ b/crates/zeta_prompt/src/zeta_prompt.rs @@ -25,6 +25,11 @@ fn estimate_tokens(bytes: usize) -> usize { bytes / 3 } +/// Leave some slack to avoid overflow. +fn apply_prompt_budget_margin(max_tokens: usize) -> usize { + (max_tokens as f64 * 0.9).floor() as usize +} + #[derive(Clone, Debug, PartialEq, Hash, Serialize, Deserialize)] pub struct ZetaPromptInput { pub cursor_path: Arc, @@ -557,13 +562,14 @@ pub fn format_prompt_with_budget_for_format( cursor_offset, ); + let budget_with_margin = apply_prompt_budget_margin(max_tokens); seed_coder::assemble_fim_prompt( context, &editable_range, &cursor_section, &input.events, related_files, - max_tokens, + budget_with_margin, ) } _ => { @@ -577,24 +583,25 @@ pub fn format_prompt_with_budget_for_format( cursor_offset, ); + let mut remaining_budget = apply_prompt_budget_margin(max_tokens); let cursor_tokens = estimate_tokens(cursor_section.len()); - let budget_after_cursor = max_tokens.saturating_sub(cursor_tokens); + remaining_budget = remaining_budget.saturating_sub(cursor_tokens); let edit_history_section = format_edit_history_within_budget( &input.events, "<|file_sep|>", "edit history", - budget_after_cursor, + remaining_budget, max_edit_event_count_for_format(&format), ); let edit_history_tokens = estimate_tokens(edit_history_section.len()); - let budget_after_edit_history = budget_after_cursor.saturating_sub(edit_history_tokens); + remaining_budget = remaining_budget.saturating_sub(edit_history_tokens); let related_files_section = format_related_files_within_budget( &related_files, "<|file_sep|>", "", - budget_after_edit_history, + remaining_budget, ); let mut prompt = String::new(); @@ -4251,6 +4258,10 @@ mod tests { format_prompt_with_budget_for_format(input, ZetaFormat::V0114180EditableRegion, max_tokens) } + fn budget_with_margin(requested_tokens: usize) -> usize { + ((requested_tokens as f64) / 0.9).ceil() as usize + } + #[test] fn test_no_truncation_when_within_budget() { let input = make_input( @@ -4321,7 +4332,7 @@ mod tests { ); assert_eq!( - format_with_budget(&input, 55), + format_with_budget(&input, budget_with_margin(55)), Some( indoc! {r#" <|file_sep|>edit history @@ -4392,7 +4403,7 @@ mod tests { ); assert_eq!( - format_with_budget(&input, 50).unwrap(), + format_with_budget(&input, budget_with_margin(50)).unwrap(), indoc! {r#" <|file_sep|>big.rs first excerpt @@ -4464,7 +4475,7 @@ mod tests { // file_b header (7) + excerpt (7) = 14 tokens, which fits. // file_a would need another 14 tokens, which doesn't fit. assert_eq!( - format_with_budget(&input, 52).unwrap(), + format_with_budget(&input, budget_with_margin(52)).unwrap(), indoc! {r#" <|file_sep|>file_b.rs high priority content @@ -4537,7 +4548,7 @@ mod tests { // With tight budget, only order<=1 excerpts included (header + important fn). assert_eq!( - format_with_budget(&input, 55).unwrap(), + format_with_budget(&input, budget_with_margin(55)).unwrap(), indoc! {r#" <|file_sep|>mod.rs mod header @@ -4661,6 +4672,34 @@ mod tests { ); } + #[test] + fn test_v0317_formats_prompt_with_many_related_files() { + let related_files = (0..900) + .map(|index| { + make_related_file( + &format!("related_{index}.rs"), + "fn helper() {\n let value = 1;\n}\n", + ) + }) + .collect(); + + let input = make_input( + "code", + 0..4, + 2, + vec![make_event("a.rs", "-x\n+y\n")], + related_files, + ); + + let prompt = + format_prompt_with_budget_for_format(&input, ZetaFormat::V0317SeedMultiRegions, 4096); + + assert!(prompt.is_some()); + let prompt = prompt.expect("v0317 should produce a prompt under high related-file count"); + assert!(prompt.contains("test.rs")); + assert!(prompt.contains(CURSOR_MARKER)); + } + #[test] fn test_seed_coder_no_context() { let input = make_input("before\nmiddle\nafter", 7..13, 10, vec![], vec![]);