| @@ -135,78 +135,129 @@ pub struct RelatedExcerpt { · 134 unmodified lines |
| 135 | 135 | } |
| 136 | 136 | |
| 137 | 137 | pub fn format_zeta_prompt(input: &ZetaPromptInput, version: ZetaVersion) -> String { |
| 138 | | − let mut prompt = String::new(); |
| 139 | | − let mut related_file_ranges = write_related_files(&mut prompt, &input.related_files); |
| 140 | | − let mut event_ranges = write_edit_history_section(&mut prompt, input); |
| 138 | + format_zeta_prompt_with_budget(input, version, MAX_PROMPT_TOKENS) |
| 139 | +} |
| 141 | 140 | |
| 141 | +fn format_zeta_prompt_with_budget( |
| 142 | + input: &ZetaPromptInput, |
| 143 | + version: ZetaVersion, |
| 144 | + max_tokens: usize, |
| 145 | +) -> String { |
| 146 | + let mut cursor_section = String::new(); |
| 142 | 147 | match version { |
| 143 | 148 | ZetaVersion::V0112MiddleAtEnd => { |
| 144 | | − v0112_middle_at_end::write_cursor_excerpt_section(&mut prompt, input); |
| 149 | + v0112_middle_at_end::write_cursor_excerpt_section(&mut cursor_section, input); |
| 145 | 150 | } |
| 146 | 151 | ZetaVersion::V0113Ordered | ZetaVersion::V0114180EditableRegion => { |
| 147 | | − v0113_ordered::write_cursor_excerpt_section(&mut prompt, input) |
| 152 | + v0113_ordered::write_cursor_excerpt_section(&mut cursor_section, input) |
| 148 | 153 | } |
| 149 | | − |
| 150 | 154 | ZetaVersion::V0120GitMergeMarkers => { |
| 151 | | − v0120_git_merge_markers::write_cursor_excerpt_section(&mut prompt, input) |
| 155 | + v0120_git_merge_markers::write_cursor_excerpt_section(&mut cursor_section, input) |
| 152 | 156 | } |
| 153 | 157 | } |
| 154 | 158 | |
| 155 | | − truncate_prompt_to_budget( |
| 156 | | − &mut prompt, |
| 157 | | − &mut related_file_ranges, |
| 158 | | − &mut event_ranges, |
| 159 | | − MAX_PROMPT_TOKENS, |
| 160 | | − ); |
| 159 | + let cursor_tokens = estimate_tokens(cursor_section.len()); |
| 160 | + let budget_after_cursor = max_tokens.saturating_sub(cursor_tokens); |
| 161 | + |
| 162 | + let edit_history_section = |
| 163 | + format_edit_history_within_budget(&input.events, budget_after_cursor); |
| 164 | + let edit_history_tokens = estimate_tokens(edit_history_section.len()); |
| 165 | + let budget_after_edit_history = budget_after_cursor.saturating_sub(edit_history_tokens); |
| 166 | + |
| 167 | + let related_files_section = |
| 168 | + format_related_files_within_budget(&input.related_files, budget_after_edit_history); |
| 161 | 169 | |
| 170 | + let mut prompt = String::new(); |
| 171 | + prompt.push_str(&related_files_section); |
| 172 | + prompt.push_str(&edit_history_section); |
| 173 | + prompt.push_str(&cursor_section); |
| 162 | 174 | prompt |
| 163 | 175 | } |
| 164 | 176 | |
| 165 | | −fn truncate_prompt_to_budget( |
| 166 | | − prompt: &mut String, |
| 167 | | − related_file_ranges: &mut Vec<Range<usize>>, |
| 168 | | − event_ranges: &mut Vec<Range<usize>>, |
| 169 | | − max_tokens: usize, |
| 170 | | −) { |
| 171 | | − let mut remove_from_related_files = true; |
| 172 | | − |
| 173 | | − while estimate_tokens(prompt.len()) > max_tokens { |
| 174 | | − let range_to_remove = if remove_from_related_files { |
| 175 | | − related_file_ranges.pop() |
| 176 | | − } else { |
| 177 | | − event_ranges.pop() |
| 178 | | − }; |
| 177 | +fn format_edit_history_within_budget(events: &[Arc<Event>], max_tokens: usize) -> String { |
| 178 | + let header = "<|file_sep|>edit history\n"; |
| 179 | + let header_tokens = estimate_tokens(header.len()); |
| 180 | + if header_tokens >= max_tokens { |
| 181 | + return String::new(); |
| 182 | + } |
| 179 | 183 | |
| 180 | | − let Some(range) = range_to_remove else { |
| 181 | | − if remove_from_related_files && !event_ranges.is_empty() { |
| 182 | | − remove_from_related_files = false; |
| 183 | | − continue; |
| 184 | | − } else if !remove_from_related_files && !related_file_ranges.is_empty() { |
| 185 | | − remove_from_related_files = true; |
| 186 | | − continue; |
| 187 | | − } else { |
| 188 | | − break; |
| 189 | | − } |
| 190 | | − }; |
| 184 | + let mut event_strings: Vec<String> = Vec::new(); |
| 185 | + let mut total_tokens = header_tokens; |
| 191 | 186 | |
| 192 | | − let removed_len = range.end - range.start; |
| 193 | | − prompt.replace_range(range.clone(), ""); |
| 187 | + for event in events.iter().rev() { |
| 188 | + let mut event_str = String::new(); |
| 189 | + write_event(&mut event_str, event); |
| 190 | + let event_tokens = estimate_tokens(event_str.len()); |
| 194 | 191 | |
| 195 | | − for r in related_file_ranges.iter_mut() { |
| 196 | | − if r.start > range.start { |
| 197 | | − r.start -= removed_len; |
| 198 | | − r.end -= removed_len; |
| 199 | | − } |
| 192 | + if total_tokens + event_tokens > max_tokens { |
| 193 | + break; |
| 194 | + } |
| 195 | + total_tokens += event_tokens; |
| 196 | + event_strings.push(event_str); |
| 197 | + } |
| 198 | + |
| 199 | + if event_strings.is_empty() { |
| 200 | + return String::new(); |
| 201 | + } |
| 202 | + |
| 203 | + let mut result = String::from(header); |
| 204 | + for event_str in event_strings.iter().rev() { |
| 205 | + result.push_str(&event_str); |
| 206 | + } |
| 207 | + result |
| 208 | +} |
| 209 | + |
| 210 | +fn format_related_files_within_budget(related_files: &[RelatedFile], max_tokens: usize) -> String { |
| 211 | + let mut result = String::new(); |
| 212 | + let mut total_tokens = 0; |
| 213 | + |
| 214 | + for file in related_files { |
| 215 | + let path_str = file.path.to_string_lossy(); |
| 216 | + let header_len = "<|file_sep|>".len() + path_str.len() + 1; |
| 217 | + let header_tokens = estimate_tokens(header_len); |
| 218 | + |
| 219 | + if total_tokens + header_tokens > max_tokens { |
| 220 | + break; |
| 200 | 221 | } |
| 201 | | − for r in event_ranges.iter_mut() { |
| 202 | | − if r.start > range.start { |
| 203 | | − r.start -= removed_len; |
| 204 | | − r.end -= removed_len; |
| 222 | + |
| 223 | + let mut file_tokens = header_tokens; |
| 224 | + let mut excerpts_to_include = 0; |
| 225 | + |
| 226 | + for excerpt in &file.excerpts { |
| 227 | + let needs_newline = !excerpt.text.ends_with('\n'); |
| 228 | + let needs_ellipsis = excerpt.row_range.end < file.max_row; |
| 229 | + let excerpt_len = excerpt.text.len() |
| 230 | + + if needs_newline { "\n".len() } else { "".len() } |
| 231 | + + if needs_ellipsis { |
| 232 | + "...\n".len() |
| 233 | + } else { |
| 234 | + "".len() |
| 235 | + }; |
| 236 | + |
| 237 | + let excerpt_tokens = estimate_tokens(excerpt_len); |
| 238 | + if total_tokens + file_tokens + excerpt_tokens > max_tokens { |
| 239 | + break; |
| 205 | 240 | } |
| 241 | + file_tokens += excerpt_tokens; |
| 242 | + excerpts_to_include += 1; |
| 206 | 243 | } |
| 207 | 244 | |
| 208 | | − remove_from_related_files = !remove_from_related_files; |
| 245 | + if excerpts_to_include > 0 { |
| 246 | + total_tokens += file_tokens; |
| 247 | + write!(result, "<|file_sep|>{}\n", path_str).ok(); |
| 248 | + for excerpt in file.excerpts.iter().take(excerpts_to_include) { |
| 249 | + result.push_str(&excerpt.text); |
| 250 | + if !result.ends_with('\n') { |
| 251 | + result.push('\n'); |
| 252 | + } |
| 253 | + if excerpt.row_range.end < file.max_row { |
| 254 | + result.push_str("...\n"); |
| 255 | + } |
| 256 | + } |
| 257 | + } |
| 209 | 258 | } |
| 259 | + |
| 260 | + result |
| 210 | 261 | } |
| 211 | 262 | |
| 212 | 263 | pub fn write_related_files( |
| @@ -233,18 +284,6 @@ pub fn write_related_files( · 20 unmodified lines |
| 233 | 284 | ranges |
| 234 | 285 | } |
| 235 | 286 | |
| 236 | | −fn write_edit_history_section(prompt: &mut String, input: &ZetaPromptInput) -> Vec<Range<usize>> { |
| 237 | | − let mut ranges = Vec::new(); |
| 238 | | − prompt.push_str("<|file_sep|>edit history\n"); |
| 239 | | − for event in &input.events { |
| 240 | | − let start = prompt.len(); |
| 241 | | − write_event(prompt, event); |
| 242 | | − let end = prompt.len(); |
| 243 | | − ranges.push(start..end); |
| 244 | | − } |
| 245 | | − ranges |
| 246 | | −} |
| 247 | | − |
| 248 | 287 | mod v0112_middle_at_end { |
| 249 | 288 | use super::*; |
| 250 | 289 | |
| @@ -376,3 +415,265 @@ pub mod v0120_git_merge_markers { · 125 unmodified lines |
| 376 | 415 | prompt.push_str(SEPARATOR); |
| 377 | 416 | } |
| 378 | 417 | } |
| 418 | + |
| 419 | +#[cfg(test)] |
| 420 | +mod tests { |
| 421 | + use super::*; |
| 422 | + use indoc::indoc; |
| 423 | + |
| 424 | + fn make_input( |
| 425 | + cursor_excerpt: &str, |
| 426 | + editable_range: Range<usize>, |
| 427 | + cursor_offset: usize, |
| 428 | + events: Vec<Event>, |
| 429 | + related_files: Vec<RelatedFile>, |
| 430 | + ) -> ZetaPromptInput { |
| 431 | + ZetaPromptInput { |
| 432 | + cursor_path: Path::new("test.rs").into(), |
| 433 | + cursor_excerpt: cursor_excerpt.into(), |
| 434 | + editable_range_in_excerpt: editable_range, |
| 435 | + cursor_offset_in_excerpt: cursor_offset, |
| 436 | + events: events.into_iter().map(Arc::new).collect(), |
| 437 | + related_files, |
| 438 | + } |
| 439 | + } |
| 440 | + |
| 441 | + fn make_event(path: &str, diff: &str) -> Event { |
| 442 | + Event::BufferChange { |
| 443 | + path: Path::new(path).into(), |
| 444 | + old_path: Path::new(path).into(), |
| 445 | + diff: diff.to_string(), |
| 446 | + predicted: false, |
| 447 | + in_open_source_repo: false, |
| 448 | + } |
| 449 | + } |
| 450 | + |
| 451 | + fn make_related_file(path: &str, content: &str) -> RelatedFile { |
| 452 | + RelatedFile { |
| 453 | + path: Path::new(path).into(), |
| 454 | + max_row: content.lines().count() as u32, |
| 455 | + excerpts: vec![RelatedExcerpt { |
| 456 | + row_range: 0..content.lines().count() as u32, |
| 457 | + text: content.into(), |
| 458 | + }], |
| 459 | + } |
| 460 | + } |
| 461 | + |
| 462 | + fn format_with_budget(input: &ZetaPromptInput, max_tokens: usize) -> String { |
| 463 | + format_zeta_prompt_with_budget(input, ZetaVersion::V0114180EditableRegion, max_tokens) |
| 464 | + } |
| 465 | + |
| 466 | + #[test] |
| 467 | + fn test_no_truncation_when_within_budget() { |
| 468 | + let input = make_input( |
| 469 | + "prefix\neditable\nsuffix", |
| 470 | + 7..15, |
| 471 | + 10, |
| 472 | + vec![make_event("a.rs", "-old\n+new\n")], |
| 473 | + vec![make_related_file("related.rs", "fn helper() {}\n")], |
| 474 | + ); |
| 475 | + |
| 476 | + assert_eq!( |
| 477 | + format_with_budget(&input, 10000), |
| 478 | + indoc! {r#" |
| 479 | + <|file_sep|>related.rs |
| 480 | + fn helper() {} |
| 481 | + <|file_sep|>edit history |
| 482 | + --- a/a.rs |
| 483 | + +++ b/a.rs |
| 484 | + -old |
| 485 | + +new |
| 486 | + <|file_sep|>test.rs |
| 487 | + <|fim_prefix|> |
| 488 | + prefix |
| 489 | + <|fim_middle|>current |
| 490 | + edi<|user_cursor|>table |
| 491 | + <|fim_suffix|> |
| 492 | + |
| 493 | + suffix |
| 494 | + <|fim_middle|>updated |
| 495 | + "#} |
| 496 | + ); |
| 497 | + } |
| 498 | + |
| 499 | + #[test] |
| 500 | + fn test_truncation_drops_edit_history_when_budget_tight() { |
| 501 | + let input = make_input( |
| 502 | + "code", |
| 503 | + 0..4, |
| 504 | + 2, |
| 505 | + vec![make_event("a.rs", "-x\n+y\n")], |
| 506 | + vec![ |
| 507 | + make_related_file("r1.rs", "a\n"), |
| 508 | + make_related_file("r2.rs", "b\n"), |
| 509 | + ], |
| 510 | + ); |
| 511 | + |
| 512 | + assert_eq!( |
| 513 | + format_with_budget(&input, 10000), |
| 514 | + indoc! {r#" |
| 515 | + <|file_sep|>r1.rs |
| 516 | + a |
| 517 | + <|file_sep|>r2.rs |
| 518 | + b |
| 519 | + <|file_sep|>edit history |
| 520 | + --- a/a.rs |
| 521 | + +++ b/a.rs |
| 522 | + -x |
| 523 | + +y |
| 524 | + <|file_sep|>test.rs |
| 525 | + <|fim_prefix|> |
| 526 | + <|fim_middle|>current |
| 527 | + co<|user_cursor|>de |
| 528 | + <|fim_suffix|> |
| 529 | + <|fim_middle|>updated |
| 530 | + "#} |
| 531 | + ); |
| 532 | + |
| 533 | + assert_eq!( |
| 534 | + format_with_budget(&input, 50), |
| 535 | + indoc! {r#" |
| 536 | + <|file_sep|>r1.rs |
| 537 | + a |
| 538 | + <|file_sep|>r2.rs |
| 539 | + b |
| 540 | + <|file_sep|>test.rs |
| 541 | + <|fim_prefix|> |
| 542 | + <|fim_middle|>current |
| 543 | + co<|user_cursor|>de |
| 544 | + <|fim_suffix|> |
| 545 | + <|fim_middle|>updated |
| 546 | + "#} |
| 547 | + ); |
| 548 | + } |
| 549 | + |
| 550 | + #[test] |
| 551 | + fn test_truncation_includes_partial_excerpts() { |
| 552 | + let input = make_input( |
| 553 | + "x", |
| 554 | + 0..1, |
| 555 | + 0, |
| 556 | + vec![], |
| 557 | + vec![RelatedFile { |
| 558 | + path: Path::new("big.rs").into(), |
| 559 | + max_row: 30, |
| 560 | + excerpts: vec![ |
| 561 | + RelatedExcerpt { |
| 562 | + row_range: 0..10, |
| 563 | + text: "first excerpt\n".into(), |
| 564 | + }, |
| 565 | + RelatedExcerpt { |
| 566 | + row_range: 10..20, |
| 567 | + text: "second excerpt\n".into(), |
| 568 | + }, |
| 569 | + RelatedExcerpt { |
| 570 | + row_range: 20..30, |
| 571 | + text: "third excerpt\n".into(), |
| 572 | + }, |
| 573 | + ], |
| 574 | + }], |
| 575 | + ); |
| 576 | + |
| 577 | + assert_eq!( |
| 578 | + format_with_budget(&input, 10000), |
| 579 | + indoc! {r#" |
| 580 | + <|file_sep|>big.rs |
| 581 | + first excerpt |
| 582 | + ... |
| 583 | + second excerpt |
| 584 | + ... |
| 585 | + third excerpt |
| 586 | + <|file_sep|>test.rs |
| 587 | + <|fim_prefix|> |
| 588 | + <|fim_middle|>current |
| 589 | + <|user_cursor|>x |
| 590 | + <|fim_suffix|> |
| 591 | + <|fim_middle|>updated |
| 592 | + "#} |
| 593 | + ); |
| 594 | + |
| 595 | + assert_eq!( |
| 596 | + format_with_budget(&input, 50), |
| 597 | + indoc! {r#" |
| 598 | + <|file_sep|>big.rs |
| 599 | + first excerpt |
| 600 | + ... |
| 601 | + <|file_sep|>test.rs |
| 602 | + <|fim_prefix|> |
| 603 | + <|fim_middle|>current |
| 604 | + <|user_cursor|>x |
| 605 | + <|fim_suffix|> |
| 606 | + <|fim_middle|>updated |
| 607 | + "#} |
| 608 | + ); |
| 609 | + } |
| 610 | + |
| 611 | + #[test] |
| 612 | + fn test_truncation_drops_older_events_first() { |
| 613 | + let input = make_input( |
| 614 | + "x", |
| 615 | + 0..1, |
| 616 | + 0, |
| 617 | + vec![make_event("old.rs", "-1\n"), make_event("new.rs", "-2\n")], |
| 618 | + vec![], |
| 619 | + ); |
| 620 | + |
| 621 | + assert_eq!( |
| 622 | + format_with_budget(&input, 10000), |
| 623 | + indoc! {r#" |
| 624 | + <|file_sep|>edit history |
| 625 | + --- a/old.rs |
| 626 | + +++ b/old.rs |
| 627 | + -1 |
| 628 | + --- a/new.rs |
| 629 | + +++ b/new.rs |
| 630 | + -2 |
| 631 | + <|file_sep|>test.rs |
| 632 | + <|fim_prefix|> |
| 633 | + <|fim_middle|>current |
| 634 | + <|user_cursor|>x |
| 635 | + <|fim_suffix|> |
| 636 | + <|fim_middle|>updated |
| 637 | + "#} |
| 638 | + ); |
| 639 | + |
| 640 | + assert_eq!( |
| 641 | + format_with_budget(&input, 55), |
| 642 | + indoc! {r#" |
| 643 | + <|file_sep|>edit history |
| 644 | + --- a/new.rs |
| 645 | + +++ b/new.rs |
| 646 | + -2 |
| 647 | + <|file_sep|>test.rs |
| 648 | + <|fim_prefix|> |
| 649 | + <|fim_middle|>current |
| 650 | + <|user_cursor|>x |
| 651 | + <|fim_suffix|> |
| 652 | + <|fim_middle|>updated |
| 653 | + "#} |
| 654 | + ); |
| 655 | + } |
| 656 | + |
| 657 | + #[test] |
| 658 | + fn test_cursor_excerpt_always_included_with_minimal_budget() { |
| 659 | + let input = make_input( |
| 660 | + "fn main() {}", |
| 661 | + 0..12, |
| 662 | + 3, |
| 663 | + vec![make_event("a.rs", "-old\n+new\n")], |
| 664 | + vec![make_related_file("related.rs", "helper\n")], |
| 665 | + ); |
| 666 | + |
| 667 | + assert_eq!( |
| 668 | + format_with_budget(&input, 30), |
| 669 | + indoc! {r#" |
| 670 | + <|file_sep|>test.rs |
| 671 | + <|fim_prefix|> |
| 672 | + <|fim_middle|>current |
| 673 | + fn <|user_cursor|>main() {} |
| 674 | + <|fim_suffix|> |
| 675 | + <|fim_middle|>updated |
| 676 | + "#} |
| 677 | + ); |
| 678 | + } |
| 679 | +} |