| @@ -135,78 +135,129 @@ pub struct RelatedExcerpt { · 134 unmodified lines |
| 135 | } | 135 | } |
| 136 | | 136 | |
| 137 | pub fn format_zeta_prompt(input: &ZetaPromptInput, version: ZetaVersion) -> String { | 137 | pub fn format_zeta_prompt(input: &ZetaPromptInput, version: ZetaVersion) -> String { |
| 138 | let mut prompt = String::new(); | 138 | format_zeta_prompt_with_budget(input, version, MAX_PROMPT_TOKENS) |
| 139 | let mut related_file_ranges = write_related_files(&mut prompt, &input.related_files); | 139 | } |
| 140 | let mut event_ranges = write_edit_history_section(&mut prompt, input); | | |
| 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 | match version { | 147 | match version { |
| 143 | ZetaVersion::V0112MiddleAtEnd => { | 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 | ZetaVersion::V0113Ordered | ZetaVersion::V0114180EditableRegion => { | 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 | ZetaVersion::V0120GitMergeMarkers => { | 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( | 159 | let cursor_tokens = estimate_tokens(cursor_section.len()); |
| 156 | &mut prompt, | 160 | let budget_after_cursor = max_tokens.saturating_sub(cursor_tokens); |
| 157 | &mut related_file_ranges, | 161 | |
| 158 | &mut event_ranges, | 162 | let edit_history_section = |
| 159 | MAX_PROMPT_TOKENS, | 163 | format_edit_history_within_budget(&input.events, budget_after_cursor); |
| 160 | ); | 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 | prompt | 174 | prompt |
| 163 | } | 175 | } |
| 164 | | 176 | |
| 165 | fn truncate_prompt_to_budget( | 177 | fn format_edit_history_within_budget(events: &[Arc<Event>], max_tokens: usize) -> String { |
| 166 | prompt: &mut String, | 178 | let header = "<|file_sep|>edit history\n"; |
| 167 | related_file_ranges: &mut Vec<Range<usize>>, | 179 | let header_tokens = estimate_tokens(header.len()); |
| 168 | event_ranges: &mut Vec<Range<usize>>, | 180 | if header_tokens >= max_tokens { |
| 169 | max_tokens: usize, | 181 | return String::new(); |
| 170 | ) { | 182 | } |
| 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 | }; | | |
| 179 | | 183 | |
| 180 | let Some(range) = range_to_remove else { | 184 | let mut event_strings: Vec<String> = Vec::new(); |
| 181 | if remove_from_related_files && !event_ranges.is_empty() { | 185 | let mut total_tokens = header_tokens; |
| 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 | }; | | |
| 191 | | 186 | |
| 192 | let removed_len = range.end - range.start; | 187 | for event in events.iter().rev() { |
| 193 | prompt.replace_range(range.clone(), ""); | 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() { | 192 | if total_tokens + event_tokens > max_tokens { |
| 196 | if r.start > range.start { | 193 | break; |
| 197 | r.start -= removed_len; | 194 | } |
| 198 | r.end -= removed_len; | 195 | total_tokens += event_tokens; |
| 199 | } | 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() { | 222 | |
| 202 | if r.start > range.start { | 223 | let mut file_tokens = header_tokens; |
| 203 | r.start -= removed_len; | 224 | let mut excerpts_to_include = 0; |
| 204 | r.end -= removed_len; | 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 | pub fn write_related_files( | 263 | pub fn write_related_files( |
| @@ -233,18 +284,6 @@ pub fn write_related_files( · 20 unmodified lines |
| 233 | ranges | 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 | mod v0112_middle_at_end { | 287 | mod v0112_middle_at_end { |
| 249 | use super::*; | 288 | use super::*; |
| 250 | | 289 | |
| @@ -376,3 +415,265 @@ pub mod v0120_git_merge_markers { · 125 unmodified lines |
| 376 | prompt.push_str(SEPARATOR); | 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 | } |