Reorganize edit prediction code and remove old experiments (#44187)
Release Notes: - N/A --------- Co-authored-by: Agus Zubiaga <agus@zed.dev> Co-authored-by: Ben Kunkle <ben@zed.dev>
This commit is contained in:
Max Brunsfeld
co-authored by
Agus Zubiaga
Ben Kunkle
Agus Zubiaga
Ben Kunkle
parent
76167109db
commit
42583c1141
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,843 @@
|
||||
use std::{
|
||||
collections::BTreeSet,
|
||||
fmt::{Display, Formatter},
|
||||
ops::Range,
|
||||
path::PathBuf,
|
||||
sync::{Arc, LazyLock},
|
||||
};
|
||||
|
||||
use anyhow::{Result, anyhow};
|
||||
use fs::Fs;
|
||||
use futures::StreamExt as _;
|
||||
use gpui::{App, AppContext as _, Entity, Subscription, Task};
|
||||
use itertools::Itertools;
|
||||
use postage::watch;
|
||||
use project::Worktree;
|
||||
use strum::VariantArray;
|
||||
use util::{ResultExt as _, maybe, rel_path::RelPath};
|
||||
use worktree::ChildEntriesOptions;
|
||||
|
||||
/// Matches the most common license locations, with US and UK English spelling.
|
||||
static LICENSE_FILE_NAME_REGEX: LazyLock<regex::bytes::Regex> = LazyLock::new(|| {
|
||||
regex::bytes::RegexBuilder::new(
|
||||
"^ \
|
||||
(?: license | licence)? \
|
||||
(?: [\\-._]? \
|
||||
(?: apache (?: [\\-._] (?: 2.0 | 2 ))? | \
|
||||
0? bsd (?: [\\-._] [0123])? (?: [\\-._] clause)? | \
|
||||
isc | \
|
||||
mit | \
|
||||
upl | \
|
||||
zlib))? \
|
||||
(?: [\\-._]? (?: license | licence))? \
|
||||
(?: \\.txt | \\.md)? \
|
||||
$",
|
||||
)
|
||||
.ignore_whitespace(true)
|
||||
.case_insensitive(true)
|
||||
.build()
|
||||
.unwrap()
|
||||
});
|
||||
|
||||
#[derive(Debug, Clone, Copy, Eq, Ord, PartialOrd, PartialEq, VariantArray)]
|
||||
pub enum OpenSourceLicense {
|
||||
Apache2_0,
|
||||
BSDZero,
|
||||
BSD,
|
||||
ISC,
|
||||
MIT,
|
||||
UPL1_0,
|
||||
Zlib,
|
||||
}
|
||||
|
||||
impl Display for OpenSourceLicense {
|
||||
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
|
||||
write!(f, "{}", self.spdx_identifier())
|
||||
}
|
||||
}
|
||||
|
||||
impl OpenSourceLicense {
|
||||
/// These are SPDX identifiers for the licenses, except for BSD, where the variants are not
|
||||
/// distinguished.
|
||||
pub fn spdx_identifier(&self) -> &'static str {
|
||||
match self {
|
||||
OpenSourceLicense::Apache2_0 => "apache-2.0",
|
||||
OpenSourceLicense::BSDZero => "0bsd",
|
||||
OpenSourceLicense::BSD => "bsd",
|
||||
OpenSourceLicense::ISC => "isc",
|
||||
OpenSourceLicense::MIT => "mit",
|
||||
OpenSourceLicense::UPL1_0 => "upl-1.0",
|
||||
OpenSourceLicense::Zlib => "zlib",
|
||||
}
|
||||
}
|
||||
|
||||
pub fn patterns(&self) -> &'static [&'static str] {
|
||||
match self {
|
||||
OpenSourceLicense::Apache2_0 => &[
|
||||
include_str!("../license_patterns/apache-2.0-pattern"),
|
||||
include_str!("../license_patterns/apache-2.0-reference-pattern"),
|
||||
],
|
||||
OpenSourceLicense::BSDZero => &[include_str!("../license_patterns/0bsd-pattern")],
|
||||
OpenSourceLicense::BSD => &[include_str!("../license_patterns/bsd-pattern")],
|
||||
OpenSourceLicense::ISC => &[include_str!("../license_patterns/isc-pattern")],
|
||||
OpenSourceLicense::MIT => &[include_str!("../license_patterns/mit-pattern")],
|
||||
OpenSourceLicense::UPL1_0 => &[include_str!("../license_patterns/upl-1.0-pattern")],
|
||||
OpenSourceLicense::Zlib => &[include_str!("../license_patterns/zlib-pattern")],
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TODO: Consider using databake or similar to not parse at runtime.
|
||||
static LICENSE_PATTERNS: LazyLock<LicensePatterns> = LazyLock::new(|| {
|
||||
let mut approximate_max_length = 0;
|
||||
let mut patterns = Vec::new();
|
||||
for license in OpenSourceLicense::VARIANTS {
|
||||
for pattern in license.patterns() {
|
||||
let (pattern, length) = parse_pattern(pattern).unwrap();
|
||||
patterns.push((*license, pattern));
|
||||
approximate_max_length = approximate_max_length.max(length);
|
||||
}
|
||||
}
|
||||
LicensePatterns {
|
||||
patterns,
|
||||
approximate_max_length,
|
||||
}
|
||||
});
|
||||
|
||||
fn detect_license(text: &str) -> Option<OpenSourceLicense> {
|
||||
let text = canonicalize_license_text(text);
|
||||
for (license, pattern) in LICENSE_PATTERNS.patterns.iter() {
|
||||
log::trace!("Checking if license is {}", license);
|
||||
if check_pattern(&pattern, &text) {
|
||||
return Some(*license);
|
||||
}
|
||||
}
|
||||
|
||||
None
|
||||
}
|
||||
|
||||
struct LicensePatterns {
|
||||
patterns: Vec<(OpenSourceLicense, Vec<PatternPart>)>,
|
||||
approximate_max_length: usize,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default, PartialEq, Eq)]
|
||||
struct PatternPart {
|
||||
/// Indicates that matching `text` is optional. Skipping `match_any_chars` is conditional on
|
||||
/// matching `text`.
|
||||
optional: bool,
|
||||
/// Indicates the number of characters that can be skipped before matching `text`.
|
||||
match_any_chars: Range<usize>,
|
||||
/// The text to match, may be empty.
|
||||
text: String,
|
||||
}
|
||||
|
||||
/// Lines that start with "-- " begin a `PatternPart`. `-- 1..10` specifies `match_any_chars:
|
||||
/// 1..10`. `-- 1..10 optional:` additionally specifies `optional: true`. It's a parse error for a
|
||||
/// line to start with `--` without matching this format.
|
||||
///
|
||||
/// Text that does not have `--` prefixes participate in the `text` field and are canonicalized by
|
||||
/// lowercasing, replacing all runs of whitespace with a single space, and otherwise only keeping
|
||||
/// ascii alphanumeric characters.
|
||||
fn parse_pattern(pattern_source: &str) -> Result<(Vec<PatternPart>, usize)> {
|
||||
let mut pattern = Vec::new();
|
||||
let mut part = PatternPart::default();
|
||||
let mut approximate_max_length = 0;
|
||||
for line in pattern_source.lines() {
|
||||
if let Some(directive) = line.trim().strip_prefix("--") {
|
||||
if part != PatternPart::default() {
|
||||
pattern.push(part);
|
||||
part = PatternPart::default();
|
||||
}
|
||||
let valid = maybe!({
|
||||
let directive_chunks = directive.split_whitespace().collect::<Vec<_>>();
|
||||
if !(1..=2).contains(&directive_chunks.len()) {
|
||||
return None;
|
||||
}
|
||||
if directive_chunks.len() == 2 {
|
||||
part.optional = true;
|
||||
}
|
||||
let range_chunks = directive_chunks[0].split("..").collect::<Vec<_>>();
|
||||
if range_chunks.len() != 2 {
|
||||
return None;
|
||||
}
|
||||
part.match_any_chars.start = range_chunks[0].parse::<usize>().ok()?;
|
||||
part.match_any_chars.end = range_chunks[1].parse::<usize>().ok()?;
|
||||
if part.match_any_chars.start > part.match_any_chars.end {
|
||||
return None;
|
||||
}
|
||||
approximate_max_length += part.match_any_chars.end;
|
||||
Some(())
|
||||
});
|
||||
if valid.is_none() {
|
||||
return Err(anyhow!("Invalid pattern directive: {}", line));
|
||||
}
|
||||
continue;
|
||||
}
|
||||
approximate_max_length += line.len() + 1;
|
||||
let line = canonicalize_license_text(line);
|
||||
if line.is_empty() {
|
||||
continue;
|
||||
}
|
||||
if !part.text.is_empty() {
|
||||
part.text.push(' ');
|
||||
}
|
||||
part.text.push_str(&line);
|
||||
}
|
||||
if part != PatternPart::default() {
|
||||
pattern.push(part);
|
||||
}
|
||||
Ok((pattern, approximate_max_length))
|
||||
}
|
||||
|
||||
/// Checks a pattern against text by iterating over the pattern parts in reverse order, and checking
|
||||
/// matches with the end of a prefix of the input. Assumes that `canonicalize_license_text` has
|
||||
/// already been applied to the input.
|
||||
fn check_pattern(pattern: &[PatternPart], input: &str) -> bool {
|
||||
let mut input_ix = input.len();
|
||||
let mut match_any_chars = 0..0;
|
||||
for part in pattern.iter().rev() {
|
||||
if part.text.is_empty() {
|
||||
match_any_chars.start += part.match_any_chars.start;
|
||||
match_any_chars.end += part.match_any_chars.end;
|
||||
continue;
|
||||
}
|
||||
|
||||
let search_range_end = n_chars_before_offset(match_any_chars.start, input_ix, input);
|
||||
let search_range_start = n_chars_before_offset(
|
||||
match_any_chars.len() + part.text.len(),
|
||||
search_range_end,
|
||||
input,
|
||||
);
|
||||
let found_ix = input[search_range_start..search_range_end].rfind(&part.text);
|
||||
|
||||
if let Some(found_ix) = found_ix {
|
||||
input_ix = search_range_start + found_ix;
|
||||
match_any_chars = part.match_any_chars.clone();
|
||||
} else if !part.optional {
|
||||
log::trace!(
|
||||
"Failed to match pattern\n`...{}`\nagainst input\n`...{}`",
|
||||
&part.text[n_chars_before_offset(128, part.text.len(), &part.text)..],
|
||||
&input[n_chars_before_offset(128, search_range_end, input)..search_range_end],
|
||||
);
|
||||
return false;
|
||||
}
|
||||
}
|
||||
is_char_count_within_range(&input[..input_ix], match_any_chars)
|
||||
}
|
||||
|
||||
fn n_chars_before_offset(char_count: usize, offset: usize, string: &str) -> usize {
|
||||
if char_count == 0 {
|
||||
return offset;
|
||||
}
|
||||
string[..offset]
|
||||
.char_indices()
|
||||
.nth_back(char_count.saturating_sub(1))
|
||||
.map_or(0, |(byte_ix, _)| byte_ix)
|
||||
}
|
||||
|
||||
fn is_char_count_within_range(string: &str, char_count_range: Range<usize>) -> bool {
|
||||
if string.len() >= char_count_range.start * 4 && string.len() < char_count_range.end {
|
||||
return true;
|
||||
}
|
||||
if string.len() < char_count_range.start || string.len() >= char_count_range.end * 4 {
|
||||
return false;
|
||||
}
|
||||
char_count_range.contains(&string.chars().count())
|
||||
}
|
||||
|
||||
/// Canonicalizes license text by removing all non-alphanumeric characters, lowercasing, and turning
|
||||
/// runs of whitespace into a single space. Unicode alphanumeric characters are intentionally
|
||||
/// preserved since these should cause license mismatch when not within a portion of the license
|
||||
/// where arbitrary text is allowed.
|
||||
fn canonicalize_license_text(license: &str) -> String {
|
||||
license
|
||||
.chars()
|
||||
.filter(|c| c.is_ascii_whitespace() || c.is_alphanumeric())
|
||||
.map(|c| c.to_ascii_lowercase())
|
||||
.collect::<String>()
|
||||
.split_ascii_whitespace()
|
||||
.join(" ")
|
||||
}
|
||||
|
||||
pub enum LicenseDetectionWatcher {
|
||||
Local {
|
||||
is_open_source_rx: watch::Receiver<bool>,
|
||||
_is_open_source_task: Task<()>,
|
||||
_worktree_subscription: Subscription,
|
||||
},
|
||||
SingleFile,
|
||||
Remote,
|
||||
}
|
||||
|
||||
impl LicenseDetectionWatcher {
|
||||
pub fn new(worktree: &Entity<Worktree>, cx: &mut App) -> Self {
|
||||
let worktree_ref = worktree.read(cx);
|
||||
if worktree_ref.is_single_file() {
|
||||
return Self::SingleFile;
|
||||
}
|
||||
|
||||
let (files_to_check_tx, mut files_to_check_rx) = futures::channel::mpsc::unbounded();
|
||||
|
||||
let Worktree::Local(local_worktree) = worktree_ref else {
|
||||
return Self::Remote;
|
||||
};
|
||||
let fs = local_worktree.fs().clone();
|
||||
|
||||
let options = ChildEntriesOptions {
|
||||
include_files: true,
|
||||
include_dirs: false,
|
||||
include_ignored: true,
|
||||
};
|
||||
for top_file in local_worktree.child_entries_with_options(RelPath::empty(), options) {
|
||||
let path_bytes = top_file.path.as_unix_str().as_bytes();
|
||||
if top_file.is_created() && LICENSE_FILE_NAME_REGEX.is_match(path_bytes) {
|
||||
let rel_path = top_file.path.clone();
|
||||
files_to_check_tx.unbounded_send(rel_path).ok();
|
||||
}
|
||||
}
|
||||
|
||||
let _worktree_subscription =
|
||||
cx.subscribe(worktree, move |_worktree, event, _cx| match event {
|
||||
worktree::Event::UpdatedEntries(updated_entries) => {
|
||||
for updated_entry in updated_entries.iter() {
|
||||
let rel_path = &updated_entry.0;
|
||||
let path_bytes = rel_path.as_unix_str().as_bytes();
|
||||
if LICENSE_FILE_NAME_REGEX.is_match(path_bytes) {
|
||||
files_to_check_tx.unbounded_send(rel_path.clone()).ok();
|
||||
}
|
||||
}
|
||||
}
|
||||
worktree::Event::DeletedEntry(_) | worktree::Event::UpdatedGitRepositories(_) => {}
|
||||
});
|
||||
|
||||
let worktree_snapshot = worktree.read(cx).snapshot();
|
||||
let (mut is_open_source_tx, is_open_source_rx) = watch::channel_with::<bool>(false);
|
||||
|
||||
let _is_open_source_task = cx.background_spawn(async move {
|
||||
let mut eligible_licenses = BTreeSet::new();
|
||||
while let Some(rel_path) = files_to_check_rx.next().await {
|
||||
let abs_path = worktree_snapshot.absolutize(&rel_path);
|
||||
let was_open_source = !eligible_licenses.is_empty();
|
||||
if Self::is_path_eligible(&fs, abs_path).await.unwrap_or(false) {
|
||||
eligible_licenses.insert(rel_path);
|
||||
} else {
|
||||
eligible_licenses.remove(&rel_path);
|
||||
}
|
||||
let is_open_source = !eligible_licenses.is_empty();
|
||||
if is_open_source != was_open_source {
|
||||
*is_open_source_tx.borrow_mut() = is_open_source;
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
Self::Local {
|
||||
is_open_source_rx,
|
||||
_is_open_source_task,
|
||||
_worktree_subscription,
|
||||
}
|
||||
}
|
||||
|
||||
async fn is_path_eligible(fs: &Arc<dyn Fs>, abs_path: PathBuf) -> Option<bool> {
|
||||
log::debug!("checking if `{abs_path:?}` is an open source license");
|
||||
// resolve symlinks so that the file size from metadata is correct
|
||||
let Some(abs_path) = fs.canonicalize(&abs_path).await.ok() else {
|
||||
log::debug!(
|
||||
"`{abs_path:?}` license file probably deleted (error canonicalizing the path)"
|
||||
);
|
||||
return None;
|
||||
};
|
||||
let metadata = fs.metadata(&abs_path).await.log_err()??;
|
||||
if metadata.len > LICENSE_PATTERNS.approximate_max_length as u64 {
|
||||
log::debug!(
|
||||
"`{abs_path:?}` license file was skipped \
|
||||
because its size of {} bytes was larger than the max size of {} bytes",
|
||||
metadata.len,
|
||||
LICENSE_PATTERNS.approximate_max_length
|
||||
);
|
||||
return None;
|
||||
}
|
||||
let text = fs.load(&abs_path).await.log_err()?;
|
||||
let is_eligible = detect_license(&text).is_some();
|
||||
if is_eligible {
|
||||
log::debug!(
|
||||
"`{abs_path:?}` matches a license that is eligible for data collection (if enabled)"
|
||||
);
|
||||
} else {
|
||||
log::debug!(
|
||||
"`{abs_path:?}` does not match a license that is eligible for data collection"
|
||||
);
|
||||
}
|
||||
Some(is_eligible)
|
||||
}
|
||||
|
||||
/// Answers false until we find out it's open source
|
||||
pub fn is_project_open_source(&self) -> bool {
|
||||
match self {
|
||||
Self::Local {
|
||||
is_open_source_rx, ..
|
||||
} => *is_open_source_rx.borrow(),
|
||||
Self::SingleFile | Self::Remote => false,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::path::Path;
|
||||
|
||||
use fs::FakeFs;
|
||||
use gpui::TestAppContext;
|
||||
use rand::Rng as _;
|
||||
use serde_json::json;
|
||||
use settings::SettingsStore;
|
||||
|
||||
use super::*;
|
||||
|
||||
const APACHE_2_0_TXT: &str = include_str!("../license_examples/apache-2.0-ex0.txt");
|
||||
const ISC_TXT: &str = include_str!("../license_examples/isc.txt");
|
||||
const MIT_TXT: &str = include_str!("../license_examples/mit-ex0.txt");
|
||||
const UPL_1_0_TXT: &str = include_str!("../license_examples/upl-1.0.txt");
|
||||
const BSD_0_TXT: &str = include_str!("../license_examples/0bsd.txt");
|
||||
|
||||
#[track_caller]
|
||||
fn assert_matches_license(text: &str, license: OpenSourceLicense) {
|
||||
assert_eq!(detect_license(text), Some(license));
|
||||
assert!(text.len() < LICENSE_PATTERNS.approximate_max_length);
|
||||
}
|
||||
|
||||
/*
|
||||
// Uncomment this and run with `cargo test -p zeta -- --no-capture &> licenses-output` to
|
||||
// traverse your entire home directory and run license detection on every file that has a
|
||||
// license-like name.
|
||||
#[test]
|
||||
fn test_check_all_licenses_in_home_dir() {
|
||||
let mut detected = Vec::new();
|
||||
let mut unrecognized = Vec::new();
|
||||
let mut walked_entries = 0;
|
||||
let homedir = std::env::home_dir().unwrap();
|
||||
for entry in walkdir::WalkDir::new(&homedir) {
|
||||
walked_entries += 1;
|
||||
if walked_entries % 10000 == 0 {
|
||||
println!(
|
||||
"So far visited {} files in {}",
|
||||
walked_entries,
|
||||
homedir.display()
|
||||
);
|
||||
}
|
||||
let Ok(entry) = entry else {
|
||||
continue;
|
||||
};
|
||||
if !LICENSE_FILE_NAME_REGEX.is_match(entry.file_name().as_encoded_bytes()) {
|
||||
continue;
|
||||
}
|
||||
let Ok(contents) = std::fs::read_to_string(entry.path()) else {
|
||||
continue;
|
||||
};
|
||||
let path_string = entry.path().to_string_lossy().into_owned();
|
||||
let license = detect_license(&contents);
|
||||
match license {
|
||||
Some(license) => detected.push((license, path_string)),
|
||||
None => unrecognized.push(path_string),
|
||||
}
|
||||
}
|
||||
println!("\nDetected licenses:\n");
|
||||
detected.sort();
|
||||
for (license, path) in &detected {
|
||||
println!("{}: {}", license.spdx_identifier(), path);
|
||||
}
|
||||
println!("\nUnrecognized licenses:\n");
|
||||
for path in &unrecognized {
|
||||
println!("{}", path);
|
||||
}
|
||||
panic!(
|
||||
"{} licenses detected, {} unrecognized",
|
||||
detected.len(),
|
||||
unrecognized.len()
|
||||
);
|
||||
println!("This line has a warning to make sure this test is always commented out");
|
||||
}
|
||||
*/
|
||||
|
||||
#[test]
|
||||
fn test_apache_positive_detection() {
|
||||
assert_matches_license(APACHE_2_0_TXT, OpenSourceLicense::Apache2_0);
|
||||
assert_matches_license(
|
||||
include_str!("../license_examples/apache-2.0-ex1.txt"),
|
||||
OpenSourceLicense::Apache2_0,
|
||||
);
|
||||
assert_matches_license(
|
||||
include_str!("../license_examples/apache-2.0-ex2.txt"),
|
||||
OpenSourceLicense::Apache2_0,
|
||||
);
|
||||
assert_matches_license(
|
||||
include_str!("../license_examples/apache-2.0-ex3.txt"),
|
||||
OpenSourceLicense::Apache2_0,
|
||||
);
|
||||
assert_matches_license(
|
||||
include_str!("../license_examples/apache-2.0-ex4.txt"),
|
||||
OpenSourceLicense::Apache2_0,
|
||||
);
|
||||
assert_matches_license(
|
||||
include_str!("../../../LICENSE-APACHE"),
|
||||
OpenSourceLicense::Apache2_0,
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_apache_negative_detection() {
|
||||
assert_eq!(
|
||||
detect_license(&format!(
|
||||
"{APACHE_2_0_TXT}\n\nThe terms in this license are void if P=NP."
|
||||
)),
|
||||
None
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_bsd_1_clause_positive_detection() {
|
||||
assert_matches_license(
|
||||
include_str!("../license_examples/bsd-1-clause.txt"),
|
||||
OpenSourceLicense::BSD,
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_bsd_2_clause_positive_detection() {
|
||||
assert_matches_license(
|
||||
include_str!("../license_examples/bsd-2-clause-ex0.txt"),
|
||||
OpenSourceLicense::BSD,
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_bsd_3_clause_positive_detection() {
|
||||
assert_matches_license(
|
||||
include_str!("../license_examples/bsd-3-clause-ex0.txt"),
|
||||
OpenSourceLicense::BSD,
|
||||
);
|
||||
assert_matches_license(
|
||||
include_str!("../license_examples/bsd-3-clause-ex1.txt"),
|
||||
OpenSourceLicense::BSD,
|
||||
);
|
||||
assert_matches_license(
|
||||
include_str!("../license_examples/bsd-3-clause-ex2.txt"),
|
||||
OpenSourceLicense::BSD,
|
||||
);
|
||||
assert_matches_license(
|
||||
include_str!("../license_examples/bsd-3-clause-ex3.txt"),
|
||||
OpenSourceLicense::BSD,
|
||||
);
|
||||
assert_matches_license(
|
||||
include_str!("../license_examples/bsd-3-clause-ex4.txt"),
|
||||
OpenSourceLicense::BSD,
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_bsd_0_positive_detection() {
|
||||
assert_matches_license(BSD_0_TXT, OpenSourceLicense::BSDZero);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_isc_positive_detection() {
|
||||
assert_matches_license(ISC_TXT, OpenSourceLicense::ISC);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_isc_negative_detection() {
|
||||
let license_text = format!(
|
||||
r#"{ISC_TXT}
|
||||
|
||||
This project is dual licensed under the ISC License and the MIT License."#
|
||||
);
|
||||
|
||||
assert_eq!(detect_license(&license_text), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_mit_positive_detection() {
|
||||
assert_matches_license(MIT_TXT, OpenSourceLicense::MIT);
|
||||
assert_matches_license(
|
||||
include_str!("../license_examples/mit-ex1.txt"),
|
||||
OpenSourceLicense::MIT,
|
||||
);
|
||||
assert_matches_license(
|
||||
include_str!("../license_examples/mit-ex2.txt"),
|
||||
OpenSourceLicense::MIT,
|
||||
);
|
||||
assert_matches_license(
|
||||
include_str!("../license_examples/mit-ex3.txt"),
|
||||
OpenSourceLicense::MIT,
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_mit_negative_detection() {
|
||||
let license_text = format!(
|
||||
r#"{MIT_TXT}
|
||||
|
||||
This project is dual licensed under the MIT License and the Apache License, Version 2.0."#
|
||||
);
|
||||
assert_eq!(detect_license(&license_text), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_upl_positive_detection() {
|
||||
assert_matches_license(UPL_1_0_TXT, OpenSourceLicense::UPL1_0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_upl_negative_detection() {
|
||||
let license_text = format!(
|
||||
r#"{UPL_1_0_TXT}
|
||||
|
||||
This project is dual licensed under the UPL License and the MIT License."#
|
||||
);
|
||||
|
||||
assert_eq!(detect_license(&license_text), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_zlib_positive_detection() {
|
||||
assert_matches_license(
|
||||
include_str!("../license_examples/zlib-ex0.txt"),
|
||||
OpenSourceLicense::Zlib,
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn random_strings_negative_detection() {
|
||||
for _i in 0..20 {
|
||||
let random_string = rand::rng()
|
||||
.sample_iter::<char, _>(rand::distr::StandardUniform)
|
||||
.take(512)
|
||||
.collect::<String>();
|
||||
assert_eq!(detect_license(&random_string), None);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_n_chars_before_offset() {
|
||||
assert_eq!(n_chars_before_offset(2, 4, "hello"), 2);
|
||||
|
||||
let input = "ㄒ乇丂ㄒ";
|
||||
assert_eq!(n_chars_before_offset(2, input.len(), input), "ㄒ乇".len());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_is_char_count_within_range() {
|
||||
// TODO: make this into a proper property test.
|
||||
for _i in 0..20 {
|
||||
let mut rng = rand::rng();
|
||||
let random_char_count = rng.random_range(0..64);
|
||||
let random_string = rand::rng()
|
||||
.sample_iter::<char, _>(rand::distr::StandardUniform)
|
||||
.take(random_char_count)
|
||||
.collect::<String>();
|
||||
let min_chars = rng.random_range(0..10);
|
||||
let max_chars = rng.random_range(min_chars..32);
|
||||
let char_count_range = min_chars..max_chars;
|
||||
assert_eq!(
|
||||
is_char_count_within_range(&random_string, char_count_range.clone()),
|
||||
char_count_range.contains(&random_char_count),
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_license_file_name_regex() {
|
||||
// Test basic license file names
|
||||
assert!(LICENSE_FILE_NAME_REGEX.is_match(b"LICENSE"));
|
||||
assert!(LICENSE_FILE_NAME_REGEX.is_match(b"LICENCE"));
|
||||
assert!(LICENSE_FILE_NAME_REGEX.is_match(b"license"));
|
||||
assert!(LICENSE_FILE_NAME_REGEX.is_match(b"licence"));
|
||||
|
||||
// Test with extensions
|
||||
assert!(LICENSE_FILE_NAME_REGEX.is_match(b"LICENSE.txt"));
|
||||
assert!(LICENSE_FILE_NAME_REGEX.is_match(b"LICENSE.md"));
|
||||
assert!(LICENSE_FILE_NAME_REGEX.is_match(b"LICENCE.txt"));
|
||||
assert!(LICENSE_FILE_NAME_REGEX.is_match(b"LICENCE.md"));
|
||||
|
||||
// Test with specific license types
|
||||
assert!(LICENSE_FILE_NAME_REGEX.is_match(b"LICENSE-APACHE"));
|
||||
assert!(LICENSE_FILE_NAME_REGEX.is_match(b"LICENSE-MIT"));
|
||||
assert!(LICENSE_FILE_NAME_REGEX.is_match(b"LICENSE.MIT"));
|
||||
assert!(LICENSE_FILE_NAME_REGEX.is_match(b"LICENSE_MIT"));
|
||||
assert!(LICENSE_FILE_NAME_REGEX.is_match(b"LICENSE-ISC"));
|
||||
assert!(LICENSE_FILE_NAME_REGEX.is_match(b"LICENSE-UPL"));
|
||||
|
||||
// Test with "license" coming after
|
||||
assert!(LICENSE_FILE_NAME_REGEX.is_match(b"APACHE-LICENSE"));
|
||||
|
||||
// Test version numbers
|
||||
assert!(LICENSE_FILE_NAME_REGEX.is_match(b"APACHE-2"));
|
||||
assert!(LICENSE_FILE_NAME_REGEX.is_match(b"APACHE-2.0"));
|
||||
assert!(LICENSE_FILE_NAME_REGEX.is_match(b"BSD-1"));
|
||||
assert!(LICENSE_FILE_NAME_REGEX.is_match(b"BSD-2"));
|
||||
assert!(LICENSE_FILE_NAME_REGEX.is_match(b"BSD-3"));
|
||||
assert!(LICENSE_FILE_NAME_REGEX.is_match(b"BSD-3-CLAUSE"));
|
||||
|
||||
// Test combinations
|
||||
assert!(LICENSE_FILE_NAME_REGEX.is_match(b"LICENSE-MIT.txt"));
|
||||
assert!(LICENSE_FILE_NAME_REGEX.is_match(b"LICENCE.ISC.md"));
|
||||
assert!(LICENSE_FILE_NAME_REGEX.is_match(b"license_upl"));
|
||||
assert!(LICENSE_FILE_NAME_REGEX.is_match(b"LICENSE.APACHE.2.0"));
|
||||
|
||||
// Test case insensitive
|
||||
assert!(LICENSE_FILE_NAME_REGEX.is_match(b"License"));
|
||||
assert!(LICENSE_FILE_NAME_REGEX.is_match(b"license-mit.TXT"));
|
||||
assert!(LICENSE_FILE_NAME_REGEX.is_match(b"LICENCE_isc.MD"));
|
||||
|
||||
// Test edge cases that should match
|
||||
assert!(LICENSE_FILE_NAME_REGEX.is_match(b"license.mit"));
|
||||
assert!(LICENSE_FILE_NAME_REGEX.is_match(b"licence-upl.txt"));
|
||||
|
||||
// Test non-matching patterns
|
||||
assert!(!LICENSE_FILE_NAME_REGEX.is_match(b"COPYING"));
|
||||
assert!(!LICENSE_FILE_NAME_REGEX.is_match(b"LICENSE.html"));
|
||||
assert!(!LICENSE_FILE_NAME_REGEX.is_match(b"MYLICENSE"));
|
||||
assert!(!LICENSE_FILE_NAME_REGEX.is_match(b"src/LICENSE"));
|
||||
assert!(!LICENSE_FILE_NAME_REGEX.is_match(b"LICENSE.old"));
|
||||
assert!(!LICENSE_FILE_NAME_REGEX.is_match(b"LICENSE-GPL"));
|
||||
assert!(!LICENSE_FILE_NAME_REGEX.is_match(b"LICENSEABC"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_canonicalize_license_text() {
|
||||
let input = " Paragraph 1\nwith multiple lines\n\n\n\nParagraph 2\nwith more lines\n ";
|
||||
let expected = "paragraph 1 with multiple lines paragraph 2 with more lines";
|
||||
assert_eq!(canonicalize_license_text(input), expected);
|
||||
|
||||
// Test tabs and mixed whitespace
|
||||
let input = "Word1\t\tWord2\n\n Word3\r\n\r\n\r\nWord4 ";
|
||||
let expected = "word1 word2 word3 word4";
|
||||
assert_eq!(canonicalize_license_text(input), expected);
|
||||
}
|
||||
|
||||
fn init_test(cx: &mut TestAppContext) {
|
||||
cx.update(|cx| {
|
||||
let settings_store = SettingsStore::test(cx);
|
||||
cx.set_global(settings_store);
|
||||
});
|
||||
}
|
||||
|
||||
#[gpui::test]
|
||||
async fn test_watcher_single_file(cx: &mut TestAppContext) {
|
||||
init_test(cx);
|
||||
|
||||
let fs = FakeFs::new(cx.background_executor.clone());
|
||||
fs.insert_tree("/root", json!({ "main.rs": "fn main() {}" }))
|
||||
.await;
|
||||
|
||||
let worktree = Worktree::local(
|
||||
Path::new("/root/main.rs"),
|
||||
true,
|
||||
fs.clone(),
|
||||
Default::default(),
|
||||
&mut cx.to_async(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let watcher = cx.update(|cx| LicenseDetectionWatcher::new(&worktree, cx));
|
||||
assert!(matches!(watcher, LicenseDetectionWatcher::SingleFile));
|
||||
assert!(!watcher.is_project_open_source());
|
||||
}
|
||||
|
||||
#[gpui::test]
|
||||
async fn test_watcher_updates_on_changes(cx: &mut TestAppContext) {
|
||||
init_test(cx);
|
||||
|
||||
let fs = FakeFs::new(cx.background_executor.clone());
|
||||
fs.insert_tree("/root", json!({ "main.rs": "fn main() {}" }))
|
||||
.await;
|
||||
|
||||
let worktree = Worktree::local(
|
||||
Path::new("/root"),
|
||||
true,
|
||||
fs.clone(),
|
||||
Default::default(),
|
||||
&mut cx.to_async(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let watcher = cx.update(|cx| LicenseDetectionWatcher::new(&worktree, cx));
|
||||
assert!(matches!(watcher, LicenseDetectionWatcher::Local { .. }));
|
||||
assert!(!watcher.is_project_open_source());
|
||||
|
||||
fs.write(Path::new("/root/LICENSE-MIT"), MIT_TXT.as_bytes())
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
cx.background_executor.run_until_parked();
|
||||
assert!(watcher.is_project_open_source());
|
||||
|
||||
fs.write(Path::new("/root/LICENSE-APACHE"), APACHE_2_0_TXT.as_bytes())
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
cx.background_executor.run_until_parked();
|
||||
assert!(watcher.is_project_open_source());
|
||||
|
||||
fs.write(Path::new("/root/LICENSE-MIT"), "Nevermind".as_bytes())
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// Still considered open source as LICENSE-APACHE is present
|
||||
cx.background_executor.run_until_parked();
|
||||
assert!(watcher.is_project_open_source());
|
||||
|
||||
fs.write(
|
||||
Path::new("/root/LICENSE-APACHE"),
|
||||
"Also nevermind".as_bytes(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
cx.background_executor.run_until_parked();
|
||||
assert!(!watcher.is_project_open_source());
|
||||
}
|
||||
|
||||
#[gpui::test]
|
||||
async fn test_watcher_initially_opensource_and_then_deleted(cx: &mut TestAppContext) {
|
||||
init_test(cx);
|
||||
|
||||
let fs = FakeFs::new(cx.background_executor.clone());
|
||||
fs.insert_tree(
|
||||
"/root",
|
||||
json!({ "main.rs": "fn main() {}", "LICENSE-MIT": MIT_TXT }),
|
||||
)
|
||||
.await;
|
||||
|
||||
let worktree = Worktree::local(
|
||||
Path::new("/root"),
|
||||
true,
|
||||
fs.clone(),
|
||||
Default::default(),
|
||||
&mut cx.to_async(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let watcher = cx.update(|cx| LicenseDetectionWatcher::new(&worktree, cx));
|
||||
assert!(matches!(watcher, LicenseDetectionWatcher::Local { .. }));
|
||||
|
||||
cx.background_executor.run_until_parked();
|
||||
assert!(watcher.is_project_open_source());
|
||||
|
||||
fs.remove_file(
|
||||
Path::new("/root/LICENSE-MIT"),
|
||||
fs::RemoveOptions {
|
||||
recursive: false,
|
||||
ignore_if_not_exists: false,
|
||||
},
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
cx.background_executor.run_until_parked();
|
||||
assert!(!watcher.is_project_open_source());
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,174 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
use crate::ZedPredictUpsell;
|
||||
use ai_onboarding::EditPredictionOnboarding;
|
||||
use client::{Client, UserStore};
|
||||
use db::kvp::Dismissable;
|
||||
use fs::Fs;
|
||||
use gpui::{
|
||||
ClickEvent, DismissEvent, Entity, EventEmitter, FocusHandle, Focusable, MouseDownEvent, Render,
|
||||
linear_color_stop, linear_gradient,
|
||||
};
|
||||
use language::language_settings::EditPredictionProvider;
|
||||
use settings::update_settings_file;
|
||||
use ui::{Vector, VectorName, prelude::*};
|
||||
use workspace::{ModalView, Workspace};
|
||||
|
||||
#[macro_export]
|
||||
macro_rules! onboarding_event {
|
||||
($name:expr) => {
|
||||
telemetry::event!($name, source = "Edit Prediction Onboarding");
|
||||
};
|
||||
($name:expr, $($key:ident $(= $value:expr)?),+ $(,)?) => {
|
||||
telemetry::event!($name, source = "Edit Prediction Onboarding", $($key $(= $value)?),+);
|
||||
};
|
||||
}
|
||||
|
||||
/// Introduces user to Zed's Edit Prediction feature
|
||||
pub struct ZedPredictModal {
|
||||
onboarding: Entity<EditPredictionOnboarding>,
|
||||
focus_handle: FocusHandle,
|
||||
}
|
||||
|
||||
pub(crate) fn set_edit_prediction_provider(provider: EditPredictionProvider, cx: &mut App) {
|
||||
let fs = <dyn Fs>::global(cx);
|
||||
update_settings_file(fs, cx, move |settings, _| {
|
||||
settings
|
||||
.project
|
||||
.all_languages
|
||||
.features
|
||||
.get_or_insert(Default::default())
|
||||
.edit_prediction_provider = Some(provider);
|
||||
});
|
||||
}
|
||||
|
||||
impl ZedPredictModal {
|
||||
pub fn toggle(
|
||||
workspace: &mut Workspace,
|
||||
user_store: Entity<UserStore>,
|
||||
client: Arc<Client>,
|
||||
window: &mut Window,
|
||||
cx: &mut Context<Workspace>,
|
||||
) {
|
||||
workspace.toggle_modal(window, cx, |_window, cx| {
|
||||
let weak_entity = cx.weak_entity();
|
||||
Self {
|
||||
onboarding: cx.new(|cx| {
|
||||
EditPredictionOnboarding::new(
|
||||
user_store.clone(),
|
||||
client.clone(),
|
||||
copilot::Copilot::global(cx)
|
||||
.is_some_and(|copilot| copilot.read(cx).status().is_configured()),
|
||||
Arc::new({
|
||||
let this = weak_entity.clone();
|
||||
move |_window, cx| {
|
||||
ZedPredictUpsell::set_dismissed(true, cx);
|
||||
set_edit_prediction_provider(EditPredictionProvider::Zed, cx);
|
||||
this.update(cx, |_, cx| cx.emit(DismissEvent)).ok();
|
||||
}
|
||||
}),
|
||||
Arc::new({
|
||||
let this = weak_entity.clone();
|
||||
move |window, cx| {
|
||||
ZedPredictUpsell::set_dismissed(true, cx);
|
||||
set_edit_prediction_provider(EditPredictionProvider::Copilot, cx);
|
||||
this.update(cx, |_, cx| cx.emit(DismissEvent)).ok();
|
||||
copilot::initiate_sign_in(window, cx);
|
||||
}
|
||||
}),
|
||||
cx,
|
||||
)
|
||||
}),
|
||||
focus_handle: cx.focus_handle(),
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
fn cancel(&mut self, _: &menu::Cancel, _: &mut Window, cx: &mut Context<Self>) {
|
||||
ZedPredictUpsell::set_dismissed(true, cx);
|
||||
cx.emit(DismissEvent);
|
||||
}
|
||||
}
|
||||
|
||||
impl EventEmitter<DismissEvent> for ZedPredictModal {}
|
||||
|
||||
impl Focusable for ZedPredictModal {
|
||||
fn focus_handle(&self, _cx: &App) -> FocusHandle {
|
||||
self.focus_handle.clone()
|
||||
}
|
||||
}
|
||||
|
||||
impl ModalView for ZedPredictModal {
|
||||
fn on_before_dismiss(
|
||||
&mut self,
|
||||
_window: &mut Window,
|
||||
cx: &mut Context<Self>,
|
||||
) -> workspace::DismissDecision {
|
||||
ZedPredictUpsell::set_dismissed(true, cx);
|
||||
workspace::DismissDecision::Dismiss(true)
|
||||
}
|
||||
}
|
||||
|
||||
impl Render for ZedPredictModal {
|
||||
fn render(&mut self, window: &mut Window, cx: &mut Context<Self>) -> impl IntoElement {
|
||||
let window_height = window.viewport_size().height;
|
||||
let max_height = window_height - px(200.);
|
||||
|
||||
v_flex()
|
||||
.id("edit-prediction-onboarding")
|
||||
.key_context("ZedPredictModal")
|
||||
.relative()
|
||||
.w(px(550.))
|
||||
.h_full()
|
||||
.max_h(max_height)
|
||||
.p_4()
|
||||
.gap_2()
|
||||
.elevation_3(cx)
|
||||
.track_focus(&self.focus_handle(cx))
|
||||
.overflow_hidden()
|
||||
.on_action(cx.listener(Self::cancel))
|
||||
.on_action(cx.listener(|_, _: &menu::Cancel, _window, cx| {
|
||||
onboarding_event!("Cancelled", trigger = "Action");
|
||||
cx.emit(DismissEvent);
|
||||
}))
|
||||
.on_any_mouse_down(cx.listener(|this, _: &MouseDownEvent, window, _cx| {
|
||||
this.focus_handle.focus(window);
|
||||
}))
|
||||
.child(
|
||||
div()
|
||||
.opacity(0.5)
|
||||
.absolute()
|
||||
.top(px(-8.0))
|
||||
.right_0()
|
||||
.w(px(400.))
|
||||
.h(px(92.))
|
||||
.child(
|
||||
Vector::new(VectorName::AiGrid, rems_from_px(400.), rems_from_px(92.))
|
||||
.color(Color::Custom(cx.theme().colors().text.alpha(0.32))),
|
||||
),
|
||||
)
|
||||
.child(
|
||||
div()
|
||||
.absolute()
|
||||
.top_0()
|
||||
.right_0()
|
||||
.w(px(660.))
|
||||
.h(px(401.))
|
||||
.overflow_hidden()
|
||||
.bg(linear_gradient(
|
||||
75.,
|
||||
linear_color_stop(cx.theme().colors().panel_background.alpha(0.01), 1.0),
|
||||
linear_color_stop(cx.theme().colors().panel_background, 0.45),
|
||||
)),
|
||||
)
|
||||
.child(h_flex().absolute().top_2().right_2().child(
|
||||
IconButton::new("cancel", IconName::Close).on_click(cx.listener(
|
||||
|_, _: &ClickEvent, _window, cx| {
|
||||
onboarding_event!("Cancelled", trigger = "X click");
|
||||
cx.emit(DismissEvent);
|
||||
},
|
||||
)),
|
||||
))
|
||||
.child(self.onboarding.clone())
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,333 @@
|
||||
use std::{
|
||||
ops::Range,
|
||||
path::Path,
|
||||
sync::Arc,
|
||||
time::{Duration, Instant},
|
||||
};
|
||||
|
||||
use cloud_llm_client::EditPredictionRejectReason;
|
||||
use gpui::{AsyncApp, Entity, SharedString};
|
||||
use language::{Anchor, Buffer, BufferSnapshot, EditPreview, OffsetRangeExt, TextBufferSnapshot};
|
||||
use serde::Serialize;
|
||||
|
||||
#[derive(Clone, Default, Debug, PartialEq, Eq, Hash)]
|
||||
pub struct EditPredictionId(pub SharedString);
|
||||
|
||||
impl From<EditPredictionId> for gpui::ElementId {
|
||||
fn from(value: EditPredictionId) -> Self {
|
||||
gpui::ElementId::Name(value.0)
|
||||
}
|
||||
}
|
||||
|
||||
impl std::fmt::Display for EditPredictionId {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
write!(f, "{}", self.0)
|
||||
}
|
||||
}
|
||||
|
||||
/// A prediction response that was returned from the provider, whether it was ultimately valid or not.
|
||||
pub struct EditPredictionResult {
|
||||
pub id: EditPredictionId,
|
||||
pub prediction: Result<EditPrediction, EditPredictionRejectReason>,
|
||||
}
|
||||
|
||||
impl EditPredictionResult {
|
||||
pub async fn new(
|
||||
id: EditPredictionId,
|
||||
edited_buffer: &Entity<Buffer>,
|
||||
edited_buffer_snapshot: &BufferSnapshot,
|
||||
edits: Arc<[(Range<Anchor>, Arc<str>)]>,
|
||||
buffer_snapshotted_at: Instant,
|
||||
response_received_at: Instant,
|
||||
inputs: EditPredictionInputs,
|
||||
cx: &mut AsyncApp,
|
||||
) -> Self {
|
||||
if edits.is_empty() {
|
||||
return Self {
|
||||
id,
|
||||
prediction: Err(EditPredictionRejectReason::Empty),
|
||||
};
|
||||
}
|
||||
|
||||
let Some((edits, snapshot, edit_preview_task)) = edited_buffer
|
||||
.read_with(cx, |buffer, cx| {
|
||||
let new_snapshot = buffer.snapshot();
|
||||
let edits: Arc<[_]> =
|
||||
interpolate_edits(&edited_buffer_snapshot, &new_snapshot, edits)?.into();
|
||||
|
||||
Some((edits.clone(), new_snapshot, buffer.preview_edits(edits, cx)))
|
||||
})
|
||||
.ok()
|
||||
.flatten()
|
||||
else {
|
||||
return Self {
|
||||
id,
|
||||
prediction: Err(EditPredictionRejectReason::InterpolatedEmpty),
|
||||
};
|
||||
};
|
||||
|
||||
let edit_preview = edit_preview_task.await;
|
||||
|
||||
Self {
|
||||
id: id.clone(),
|
||||
prediction: Ok(EditPrediction {
|
||||
id,
|
||||
edits,
|
||||
snapshot,
|
||||
edit_preview,
|
||||
inputs,
|
||||
buffer: edited_buffer.clone(),
|
||||
buffer_snapshotted_at,
|
||||
response_received_at,
|
||||
}),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct EditPrediction {
|
||||
pub id: EditPredictionId,
|
||||
pub edits: Arc<[(Range<Anchor>, Arc<str>)]>,
|
||||
pub snapshot: BufferSnapshot,
|
||||
pub edit_preview: EditPreview,
|
||||
pub buffer: Entity<Buffer>,
|
||||
pub buffer_snapshotted_at: Instant,
|
||||
pub response_received_at: Instant,
|
||||
pub inputs: EditPredictionInputs,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
pub struct EditPredictionInputs {
|
||||
pub events: Vec<Arc<cloud_llm_client::predict_edits_v3::Event>>,
|
||||
pub included_files: Vec<cloud_llm_client::predict_edits_v3::RelatedFile>,
|
||||
pub cursor_point: cloud_llm_client::predict_edits_v3::Point,
|
||||
pub cursor_path: Arc<Path>,
|
||||
}
|
||||
|
||||
impl EditPrediction {
|
||||
pub fn interpolate(
|
||||
&self,
|
||||
new_snapshot: &TextBufferSnapshot,
|
||||
) -> Option<Vec<(Range<Anchor>, Arc<str>)>> {
|
||||
interpolate_edits(&self.snapshot, new_snapshot, self.edits.clone())
|
||||
}
|
||||
|
||||
pub fn targets_buffer(&self, buffer: &Buffer) -> bool {
|
||||
self.snapshot.remote_id() == buffer.remote_id()
|
||||
}
|
||||
|
||||
pub fn latency(&self) -> Duration {
|
||||
self.response_received_at - self.buffer_snapshotted_at
|
||||
}
|
||||
}
|
||||
|
||||
impl std::fmt::Debug for EditPrediction {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
f.debug_struct("EditPrediction")
|
||||
.field("id", &self.id)
|
||||
.field("edits", &self.edits)
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
pub fn interpolate_edits(
|
||||
old_snapshot: &TextBufferSnapshot,
|
||||
new_snapshot: &TextBufferSnapshot,
|
||||
current_edits: Arc<[(Range<Anchor>, Arc<str>)]>,
|
||||
) -> Option<Vec<(Range<Anchor>, Arc<str>)>> {
|
||||
let mut edits = Vec::new();
|
||||
|
||||
let mut model_edits = current_edits.iter().peekable();
|
||||
for user_edit in new_snapshot.edits_since::<usize>(&old_snapshot.version) {
|
||||
while let Some((model_old_range, _)) = model_edits.peek() {
|
||||
let model_old_range = model_old_range.to_offset(old_snapshot);
|
||||
if model_old_range.end < user_edit.old.start {
|
||||
let (model_old_range, model_new_text) = model_edits.next().unwrap();
|
||||
edits.push((model_old_range.clone(), model_new_text.clone()));
|
||||
} else {
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
if let Some((model_old_range, model_new_text)) = model_edits.peek() {
|
||||
let model_old_offset_range = model_old_range.to_offset(old_snapshot);
|
||||
if user_edit.old == model_old_offset_range {
|
||||
let user_new_text = new_snapshot
|
||||
.text_for_range(user_edit.new.clone())
|
||||
.collect::<String>();
|
||||
|
||||
if let Some(model_suffix) = model_new_text.strip_prefix(&user_new_text) {
|
||||
if !model_suffix.is_empty() {
|
||||
let anchor = old_snapshot.anchor_after(user_edit.old.end);
|
||||
edits.push((anchor..anchor, model_suffix.into()));
|
||||
}
|
||||
|
||||
model_edits.next();
|
||||
continue;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return None;
|
||||
}
|
||||
|
||||
edits.extend(model_edits.cloned());
|
||||
|
||||
if edits.is_empty() { None } else { Some(edits) }
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use gpui::{App, Entity, TestAppContext, prelude::*};
|
||||
use language::{Buffer, ToOffset as _};
|
||||
|
||||
#[gpui::test]
|
||||
async fn test_edit_prediction_basic_interpolation(cx: &mut TestAppContext) {
|
||||
let buffer = cx.new(|cx| Buffer::local("Lorem ipsum dolor", cx));
|
||||
let edits: Arc<[(Range<Anchor>, Arc<str>)]> = cx.update(|cx| {
|
||||
to_prediction_edits([(2..5, "REM".into()), (9..11, "".into())], &buffer, cx).into()
|
||||
});
|
||||
|
||||
let edit_preview = cx
|
||||
.read(|cx| buffer.read(cx).preview_edits(edits.clone(), cx))
|
||||
.await;
|
||||
|
||||
let prediction = EditPrediction {
|
||||
id: EditPredictionId("prediction-1".into()),
|
||||
edits,
|
||||
snapshot: cx.read(|cx| buffer.read(cx).snapshot()),
|
||||
buffer: buffer.clone(),
|
||||
edit_preview,
|
||||
inputs: EditPredictionInputs {
|
||||
events: vec![],
|
||||
included_files: vec![],
|
||||
cursor_point: cloud_llm_client::predict_edits_v3::Point {
|
||||
line: cloud_llm_client::predict_edits_v3::Line(0),
|
||||
column: 0,
|
||||
},
|
||||
cursor_path: Path::new("path.txt").into(),
|
||||
},
|
||||
buffer_snapshotted_at: Instant::now(),
|
||||
response_received_at: Instant::now(),
|
||||
};
|
||||
|
||||
cx.update(|cx| {
|
||||
assert_eq!(
|
||||
from_prediction_edits(
|
||||
&prediction.interpolate(&buffer.read(cx).snapshot()).unwrap(),
|
||||
&buffer,
|
||||
cx
|
||||
),
|
||||
vec![(2..5, "REM".into()), (9..11, "".into())]
|
||||
);
|
||||
|
||||
buffer.update(cx, |buffer, cx| buffer.edit([(2..5, "")], None, cx));
|
||||
assert_eq!(
|
||||
from_prediction_edits(
|
||||
&prediction.interpolate(&buffer.read(cx).snapshot()).unwrap(),
|
||||
&buffer,
|
||||
cx
|
||||
),
|
||||
vec![(2..2, "REM".into()), (6..8, "".into())]
|
||||
);
|
||||
|
||||
buffer.update(cx, |buffer, cx| buffer.undo(cx));
|
||||
assert_eq!(
|
||||
from_prediction_edits(
|
||||
&prediction.interpolate(&buffer.read(cx).snapshot()).unwrap(),
|
||||
&buffer,
|
||||
cx
|
||||
),
|
||||
vec![(2..5, "REM".into()), (9..11, "".into())]
|
||||
);
|
||||
|
||||
buffer.update(cx, |buffer, cx| buffer.edit([(2..5, "R")], None, cx));
|
||||
assert_eq!(
|
||||
from_prediction_edits(
|
||||
&prediction.interpolate(&buffer.read(cx).snapshot()).unwrap(),
|
||||
&buffer,
|
||||
cx
|
||||
),
|
||||
vec![(3..3, "EM".into()), (7..9, "".into())]
|
||||
);
|
||||
|
||||
buffer.update(cx, |buffer, cx| buffer.edit([(3..3, "E")], None, cx));
|
||||
assert_eq!(
|
||||
from_prediction_edits(
|
||||
&prediction.interpolate(&buffer.read(cx).snapshot()).unwrap(),
|
||||
&buffer,
|
||||
cx
|
||||
),
|
||||
vec![(4..4, "M".into()), (8..10, "".into())]
|
||||
);
|
||||
|
||||
buffer.update(cx, |buffer, cx| buffer.edit([(4..4, "M")], None, cx));
|
||||
assert_eq!(
|
||||
from_prediction_edits(
|
||||
&prediction.interpolate(&buffer.read(cx).snapshot()).unwrap(),
|
||||
&buffer,
|
||||
cx
|
||||
),
|
||||
vec![(9..11, "".into())]
|
||||
);
|
||||
|
||||
buffer.update(cx, |buffer, cx| buffer.edit([(4..5, "")], None, cx));
|
||||
assert_eq!(
|
||||
from_prediction_edits(
|
||||
&prediction.interpolate(&buffer.read(cx).snapshot()).unwrap(),
|
||||
&buffer,
|
||||
cx
|
||||
),
|
||||
vec![(4..4, "M".into()), (8..10, "".into())]
|
||||
);
|
||||
|
||||
buffer.update(cx, |buffer, cx| buffer.edit([(8..10, "")], None, cx));
|
||||
assert_eq!(
|
||||
from_prediction_edits(
|
||||
&prediction.interpolate(&buffer.read(cx).snapshot()).unwrap(),
|
||||
&buffer,
|
||||
cx
|
||||
),
|
||||
vec![(4..4, "M".into())]
|
||||
);
|
||||
|
||||
buffer.update(cx, |buffer, cx| buffer.edit([(4..6, "")], None, cx));
|
||||
assert_eq!(prediction.interpolate(&buffer.read(cx).snapshot()), None);
|
||||
})
|
||||
}
|
||||
|
||||
fn to_prediction_edits(
|
||||
iterator: impl IntoIterator<Item = (Range<usize>, Arc<str>)>,
|
||||
buffer: &Entity<Buffer>,
|
||||
cx: &App,
|
||||
) -> Vec<(Range<Anchor>, Arc<str>)> {
|
||||
let buffer = buffer.read(cx);
|
||||
iterator
|
||||
.into_iter()
|
||||
.map(|(range, text)| {
|
||||
(
|
||||
buffer.anchor_after(range.start)..buffer.anchor_before(range.end),
|
||||
text,
|
||||
)
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn from_prediction_edits(
|
||||
editor_edits: &[(Range<Anchor>, Arc<str>)],
|
||||
buffer: &Entity<Buffer>,
|
||||
cx: &App,
|
||||
) -> Vec<(Range<usize>, Arc<str>)> {
|
||||
let buffer = buffer.read(cx);
|
||||
editor_edits
|
||||
.iter()
|
||||
.map(|(range, text)| {
|
||||
(
|
||||
range.start.to_offset(buffer)..range.end.to_offset(buffer),
|
||||
text.clone(),
|
||||
)
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,440 @@
|
||||
use anyhow::{Context as _, Result};
|
||||
use cloud_llm_client::predict_edits_v3::Event;
|
||||
use credentials_provider::CredentialsProvider;
|
||||
use edit_prediction_context::RelatedFile;
|
||||
use futures::{AsyncReadExt as _, FutureExt, future::Shared};
|
||||
use gpui::{
|
||||
App, AppContext as _, Entity, Task,
|
||||
http_client::{self, AsyncBody, Method},
|
||||
};
|
||||
use language::{Buffer, BufferSnapshot, Point, ToOffset as _, ToPoint as _};
|
||||
use lsp::DiagnosticSeverity;
|
||||
use project::{Project, ProjectPath};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::{
|
||||
collections::VecDeque,
|
||||
fmt::{self, Write as _},
|
||||
ops::Range,
|
||||
path::Path,
|
||||
sync::Arc,
|
||||
time::Instant,
|
||||
};
|
||||
|
||||
use crate::{EditPredictionId, EditPredictionInputs, prediction::EditPredictionResult};
|
||||
|
||||
const SWEEP_API_URL: &str = "https://autocomplete.sweep.dev/backend/next_edit_autocomplete";
|
||||
|
||||
pub struct SweepAi {
|
||||
pub api_token: Shared<Task<Option<String>>>,
|
||||
pub debug_info: Arc<str>,
|
||||
}
|
||||
|
||||
impl SweepAi {
|
||||
pub fn new(cx: &App) -> Self {
|
||||
SweepAi {
|
||||
api_token: load_api_token(cx).shared(),
|
||||
debug_info: debug_info(cx),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn set_api_token(&mut self, api_token: Option<String>, cx: &mut App) -> Task<Result<()>> {
|
||||
self.api_token = Task::ready(api_token.clone()).shared();
|
||||
store_api_token_in_keychain(api_token, cx)
|
||||
}
|
||||
|
||||
pub fn request_prediction_with_sweep(
|
||||
&self,
|
||||
project: &Entity<Project>,
|
||||
active_buffer: &Entity<Buffer>,
|
||||
snapshot: BufferSnapshot,
|
||||
position: language::Anchor,
|
||||
events: Vec<Arc<Event>>,
|
||||
recent_paths: &VecDeque<ProjectPath>,
|
||||
related_files: Vec<RelatedFile>,
|
||||
diagnostic_search_range: Range<Point>,
|
||||
cx: &mut App,
|
||||
) -> Task<Result<Option<EditPredictionResult>>> {
|
||||
let debug_info = self.debug_info.clone();
|
||||
let Some(api_token) = self.api_token.clone().now_or_never().flatten() else {
|
||||
return Task::ready(Ok(None));
|
||||
};
|
||||
let full_path: Arc<Path> = snapshot
|
||||
.file()
|
||||
.map(|file| file.full_path(cx))
|
||||
.unwrap_or_else(|| "untitled".into())
|
||||
.into();
|
||||
|
||||
let project_file = project::File::from_dyn(snapshot.file());
|
||||
let repo_name = project_file
|
||||
.map(|file| file.worktree.read(cx).root_name_str())
|
||||
.unwrap_or("untitled")
|
||||
.into();
|
||||
let offset = position.to_offset(&snapshot);
|
||||
|
||||
let recent_buffers = recent_paths.iter().cloned();
|
||||
let http_client = cx.http_client();
|
||||
|
||||
let recent_buffer_snapshots = recent_buffers
|
||||
.filter_map(|project_path| {
|
||||
let buffer = project.read(cx).get_open_buffer(&project_path, cx)?;
|
||||
if active_buffer == &buffer {
|
||||
None
|
||||
} else {
|
||||
Some(buffer.read(cx).snapshot())
|
||||
}
|
||||
})
|
||||
.take(3)
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
let cursor_point = position.to_point(&snapshot);
|
||||
let buffer_snapshotted_at = Instant::now();
|
||||
|
||||
let result = cx.background_spawn(async move {
|
||||
let text = snapshot.text();
|
||||
|
||||
let mut recent_changes = String::new();
|
||||
for event in &events {
|
||||
write_event(event.as_ref(), &mut recent_changes).unwrap();
|
||||
}
|
||||
|
||||
let mut file_chunks = recent_buffer_snapshots
|
||||
.into_iter()
|
||||
.map(|snapshot| {
|
||||
let end_point = Point::new(30, 0).min(snapshot.max_point());
|
||||
FileChunk {
|
||||
content: snapshot.text_for_range(Point::zero()..end_point).collect(),
|
||||
file_path: snapshot
|
||||
.file()
|
||||
.map(|f| f.path().as_unix_str())
|
||||
.unwrap_or("untitled")
|
||||
.to_string(),
|
||||
start_line: 0,
|
||||
end_line: end_point.row as usize,
|
||||
timestamp: snapshot.file().and_then(|file| {
|
||||
Some(
|
||||
file.disk_state()
|
||||
.mtime()?
|
||||
.to_seconds_and_nanos_for_persistence()?
|
||||
.0,
|
||||
)
|
||||
}),
|
||||
}
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
let retrieval_chunks = related_files
|
||||
.iter()
|
||||
.flat_map(|related_file| {
|
||||
related_file.excerpts.iter().map(|excerpt| FileChunk {
|
||||
file_path: related_file.path.path.as_unix_str().to_string(),
|
||||
start_line: excerpt.point_range.start.row as usize,
|
||||
end_line: excerpt.point_range.end.row as usize,
|
||||
content: excerpt.text.to_string(),
|
||||
timestamp: None,
|
||||
})
|
||||
})
|
||||
.collect();
|
||||
|
||||
let diagnostic_entries = snapshot.diagnostics_in_range(diagnostic_search_range, false);
|
||||
let mut diagnostic_content = String::new();
|
||||
let mut diagnostic_count = 0;
|
||||
|
||||
for entry in diagnostic_entries {
|
||||
let start_point: Point = entry.range.start;
|
||||
|
||||
let severity = match entry.diagnostic.severity {
|
||||
DiagnosticSeverity::ERROR => "error",
|
||||
DiagnosticSeverity::WARNING => "warning",
|
||||
DiagnosticSeverity::INFORMATION => "info",
|
||||
DiagnosticSeverity::HINT => "hint",
|
||||
_ => continue,
|
||||
};
|
||||
|
||||
diagnostic_count += 1;
|
||||
|
||||
writeln!(
|
||||
&mut diagnostic_content,
|
||||
"{} at line {}: {}",
|
||||
severity,
|
||||
start_point.row + 1,
|
||||
entry.diagnostic.message
|
||||
)?;
|
||||
}
|
||||
|
||||
if !diagnostic_content.is_empty() {
|
||||
file_chunks.push(FileChunk {
|
||||
file_path: format!("Diagnostics for {}", full_path.display()),
|
||||
start_line: 0,
|
||||
end_line: diagnostic_count,
|
||||
content: diagnostic_content,
|
||||
timestamp: None,
|
||||
});
|
||||
}
|
||||
|
||||
let request_body = AutocompleteRequest {
|
||||
debug_info,
|
||||
repo_name,
|
||||
file_path: full_path.clone(),
|
||||
file_contents: text.clone(),
|
||||
original_file_contents: text,
|
||||
cursor_position: offset,
|
||||
recent_changes: recent_changes.clone(),
|
||||
changes_above_cursor: true,
|
||||
multiple_suggestions: false,
|
||||
branch: None,
|
||||
file_chunks,
|
||||
retrieval_chunks,
|
||||
recent_user_actions: vec![],
|
||||
use_bytes: true,
|
||||
// TODO
|
||||
privacy_mode_enabled: false,
|
||||
};
|
||||
|
||||
let mut buf: Vec<u8> = Vec::new();
|
||||
let writer = brotli::CompressorWriter::new(&mut buf, 4096, 11, 22);
|
||||
serde_json::to_writer(writer, &request_body)?;
|
||||
let body: AsyncBody = buf.into();
|
||||
|
||||
let inputs = EditPredictionInputs {
|
||||
events,
|
||||
included_files: vec![cloud_llm_client::predict_edits_v3::RelatedFile {
|
||||
path: full_path.clone(),
|
||||
max_row: cloud_llm_client::predict_edits_v3::Line(snapshot.max_point().row),
|
||||
excerpts: vec![cloud_llm_client::predict_edits_v3::Excerpt {
|
||||
start_line: cloud_llm_client::predict_edits_v3::Line(0),
|
||||
text: request_body.file_contents.into(),
|
||||
}],
|
||||
}],
|
||||
cursor_point: cloud_llm_client::predict_edits_v3::Point {
|
||||
column: cursor_point.column,
|
||||
line: cloud_llm_client::predict_edits_v3::Line(cursor_point.row),
|
||||
},
|
||||
cursor_path: full_path.clone(),
|
||||
};
|
||||
|
||||
let request = http_client::Request::builder()
|
||||
.uri(SWEEP_API_URL)
|
||||
.header("Content-Type", "application/json")
|
||||
.header("Authorization", format!("Bearer {}", api_token))
|
||||
.header("Connection", "keep-alive")
|
||||
.header("Content-Encoding", "br")
|
||||
.method(Method::POST)
|
||||
.body(body)?;
|
||||
|
||||
let mut response = http_client.send(request).await?;
|
||||
|
||||
let mut body: Vec<u8> = Vec::new();
|
||||
response.body_mut().read_to_end(&mut body).await?;
|
||||
|
||||
let response_received_at = Instant::now();
|
||||
if !response.status().is_success() {
|
||||
anyhow::bail!(
|
||||
"Request failed with status: {:?}\nBody: {}",
|
||||
response.status(),
|
||||
String::from_utf8_lossy(&body),
|
||||
);
|
||||
};
|
||||
|
||||
let response: AutocompleteResponse = serde_json::from_slice(&body)?;
|
||||
|
||||
let old_text = snapshot
|
||||
.text_for_range(response.start_index..response.end_index)
|
||||
.collect::<String>();
|
||||
let edits = language::text_diff(&old_text, &response.completion)
|
||||
.into_iter()
|
||||
.map(|(range, text)| {
|
||||
(
|
||||
snapshot.anchor_after(response.start_index + range.start)
|
||||
..snapshot.anchor_before(response.start_index + range.end),
|
||||
text,
|
||||
)
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
anyhow::Ok((
|
||||
response.autocomplete_id,
|
||||
edits,
|
||||
snapshot,
|
||||
response_received_at,
|
||||
inputs,
|
||||
))
|
||||
});
|
||||
|
||||
let buffer = active_buffer.clone();
|
||||
|
||||
cx.spawn(async move |cx| {
|
||||
let (id, edits, old_snapshot, response_received_at, inputs) = result.await?;
|
||||
anyhow::Ok(Some(
|
||||
EditPredictionResult::new(
|
||||
EditPredictionId(id.into()),
|
||||
&buffer,
|
||||
&old_snapshot,
|
||||
edits.into(),
|
||||
buffer_snapshotted_at,
|
||||
response_received_at,
|
||||
inputs,
|
||||
cx,
|
||||
)
|
||||
.await,
|
||||
))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
pub const SWEEP_CREDENTIALS_URL: &str = "https://autocomplete.sweep.dev";
|
||||
pub const SWEEP_CREDENTIALS_USERNAME: &str = "sweep-api-token";
|
||||
|
||||
pub fn load_api_token(cx: &App) -> Task<Option<String>> {
|
||||
if let Some(api_token) = std::env::var("SWEEP_AI_TOKEN")
|
||||
.ok()
|
||||
.filter(|value| !value.is_empty())
|
||||
{
|
||||
return Task::ready(Some(api_token));
|
||||
}
|
||||
let credentials_provider = <dyn CredentialsProvider>::global(cx);
|
||||
cx.spawn(async move |cx| {
|
||||
let (_, credentials) = credentials_provider
|
||||
.read_credentials(SWEEP_CREDENTIALS_URL, &cx)
|
||||
.await
|
||||
.ok()??;
|
||||
String::from_utf8(credentials).ok()
|
||||
})
|
||||
}
|
||||
|
||||
fn store_api_token_in_keychain(api_token: Option<String>, cx: &App) -> Task<Result<()>> {
|
||||
let credentials_provider = <dyn CredentialsProvider>::global(cx);
|
||||
|
||||
cx.spawn(async move |cx| {
|
||||
if let Some(api_token) = api_token {
|
||||
credentials_provider
|
||||
.write_credentials(
|
||||
SWEEP_CREDENTIALS_URL,
|
||||
SWEEP_CREDENTIALS_USERNAME,
|
||||
api_token.as_bytes(),
|
||||
cx,
|
||||
)
|
||||
.await
|
||||
.context("Failed to save Sweep API token to system keychain")
|
||||
} else {
|
||||
credentials_provider
|
||||
.delete_credentials(SWEEP_CREDENTIALS_URL, cx)
|
||||
.await
|
||||
.context("Failed to delete Sweep API token from system keychain")
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
struct AutocompleteRequest {
|
||||
pub debug_info: Arc<str>,
|
||||
pub repo_name: String,
|
||||
pub branch: Option<String>,
|
||||
pub file_path: Arc<Path>,
|
||||
pub file_contents: String,
|
||||
pub recent_changes: String,
|
||||
pub cursor_position: usize,
|
||||
pub original_file_contents: String,
|
||||
pub file_chunks: Vec<FileChunk>,
|
||||
pub retrieval_chunks: Vec<FileChunk>,
|
||||
pub recent_user_actions: Vec<UserAction>,
|
||||
pub multiple_suggestions: bool,
|
||||
pub privacy_mode_enabled: bool,
|
||||
pub changes_above_cursor: bool,
|
||||
pub use_bytes: bool,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
struct FileChunk {
|
||||
pub file_path: String,
|
||||
pub start_line: usize,
|
||||
pub end_line: usize,
|
||||
pub content: String,
|
||||
pub timestamp: Option<u64>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
struct UserAction {
|
||||
pub action_type: ActionType,
|
||||
pub line_number: usize,
|
||||
pub offset: usize,
|
||||
pub file_path: String,
|
||||
pub timestamp: u64,
|
||||
}
|
||||
|
||||
#[allow(dead_code)]
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
|
||||
#[serde(rename_all = "SCREAMING_SNAKE_CASE")]
|
||||
enum ActionType {
|
||||
CursorMovement,
|
||||
InsertChar,
|
||||
DeleteChar,
|
||||
InsertSelection,
|
||||
DeleteSelection,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Deserialize)]
|
||||
struct AutocompleteResponse {
|
||||
pub autocomplete_id: String,
|
||||
pub start_index: usize,
|
||||
pub end_index: usize,
|
||||
pub completion: String,
|
||||
#[allow(dead_code)]
|
||||
pub confidence: f64,
|
||||
#[allow(dead_code)]
|
||||
pub logprobs: Option<serde_json::Value>,
|
||||
#[allow(dead_code)]
|
||||
pub finish_reason: Option<String>,
|
||||
#[allow(dead_code)]
|
||||
pub elapsed_time_ms: u64,
|
||||
#[allow(dead_code)]
|
||||
#[serde(default, rename = "completions")]
|
||||
pub additional_completions: Vec<AdditionalCompletion>,
|
||||
}
|
||||
|
||||
#[allow(dead_code)]
|
||||
#[derive(Debug, Clone, Deserialize)]
|
||||
struct AdditionalCompletion {
|
||||
pub start_index: usize,
|
||||
pub end_index: usize,
|
||||
pub completion: String,
|
||||
pub confidence: f64,
|
||||
pub autocomplete_id: String,
|
||||
pub logprobs: Option<serde_json::Value>,
|
||||
pub finish_reason: Option<String>,
|
||||
}
|
||||
|
||||
fn write_event(
|
||||
event: &cloud_llm_client::predict_edits_v3::Event,
|
||||
f: &mut impl fmt::Write,
|
||||
) -> fmt::Result {
|
||||
match event {
|
||||
cloud_llm_client::predict_edits_v3::Event::BufferChange {
|
||||
old_path,
|
||||
path,
|
||||
diff,
|
||||
..
|
||||
} => {
|
||||
if old_path != path {
|
||||
// TODO confirm how to do this for sweep
|
||||
// writeln!(f, "User renamed {:?} to {:?}\n", old_path, new_path)?;
|
||||
}
|
||||
|
||||
if !diff.is_empty() {
|
||||
write!(f, "File: {}:\n{}\n", path.display(), diff)?
|
||||
}
|
||||
|
||||
fmt::Result::Ok(())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn debug_info(cx: &gpui::App) -> Arc<str> {
|
||||
format!(
|
||||
"Zed v{version} ({sha}) - OS: {os} - Zed v{version}",
|
||||
version = release_channel::AppVersion::global(cx),
|
||||
sha = release_channel::AppCommitSha::try_global(cx)
|
||||
.map_or("unknown".to_string(), |sha| sha.full()),
|
||||
os = client::telemetry::os_name(),
|
||||
)
|
||||
.into()
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,637 @@
|
||||
use anyhow::{Context as _, Result};
|
||||
use language::{Anchor, BufferSnapshot, OffsetRangeExt as _, Point};
|
||||
use std::{cmp, ops::Range, path::Path, sync::Arc};
|
||||
|
||||
const EDITS_TAG_NAME: &'static str = "edits";
|
||||
const OLD_TEXT_TAG_NAME: &'static str = "old_text";
|
||||
const NEW_TEXT_TAG_NAME: &'static str = "new_text";
|
||||
const XML_TAGS: &[&str] = &[EDITS_TAG_NAME, OLD_TEXT_TAG_NAME, NEW_TEXT_TAG_NAME];
|
||||
|
||||
pub async fn parse_xml_edits<'a>(
|
||||
input: &'a str,
|
||||
get_buffer: impl Fn(&Path) -> Option<(&'a BufferSnapshot, &'a [Range<Anchor>])> + Send,
|
||||
) -> Result<(&'a BufferSnapshot, Vec<(Range<Anchor>, Arc<str>)>)> {
|
||||
parse_xml_edits_inner(input, get_buffer)
|
||||
.await
|
||||
.with_context(|| format!("Failed to parse XML edits:\n{input}"))
|
||||
}
|
||||
|
||||
async fn parse_xml_edits_inner<'a>(
|
||||
input: &'a str,
|
||||
get_buffer: impl Fn(&Path) -> Option<(&'a BufferSnapshot, &'a [Range<Anchor>])> + Send,
|
||||
) -> Result<(&'a BufferSnapshot, Vec<(Range<Anchor>, Arc<str>)>)> {
|
||||
let xml_edits = extract_xml_replacements(input)?;
|
||||
|
||||
let (buffer, context_ranges) = get_buffer(xml_edits.file_path.as_ref())
|
||||
.with_context(|| format!("no buffer for file {}", xml_edits.file_path))?;
|
||||
|
||||
let mut all_edits = vec![];
|
||||
for (old_text, new_text) in xml_edits.replacements {
|
||||
let match_range = fuzzy_match_in_ranges(old_text, buffer, context_ranges)?;
|
||||
let matched_old_text = buffer
|
||||
.text_for_range(match_range.clone())
|
||||
.collect::<String>();
|
||||
let edits_within_hunk = language::text_diff(&matched_old_text, new_text);
|
||||
all_edits.extend(
|
||||
edits_within_hunk
|
||||
.into_iter()
|
||||
.map(move |(inner_range, inner_text)| {
|
||||
(
|
||||
buffer.anchor_after(match_range.start + inner_range.start)
|
||||
..buffer.anchor_before(match_range.start + inner_range.end),
|
||||
inner_text,
|
||||
)
|
||||
}),
|
||||
);
|
||||
}
|
||||
|
||||
Ok((buffer, all_edits))
|
||||
}
|
||||
|
||||
fn fuzzy_match_in_ranges(
|
||||
old_text: &str,
|
||||
buffer: &BufferSnapshot,
|
||||
context_ranges: &[Range<Anchor>],
|
||||
) -> Result<Range<usize>> {
|
||||
let mut state = FuzzyMatcher::new(buffer, old_text);
|
||||
let mut best_match = None;
|
||||
let mut tie_match_range = None;
|
||||
|
||||
for range in context_ranges {
|
||||
let best_match_cost = best_match.as_ref().map(|(score, _)| *score);
|
||||
match (best_match_cost, state.match_range(range.to_offset(buffer))) {
|
||||
(Some(lowest_cost), Some((new_cost, new_range))) => {
|
||||
if new_cost == lowest_cost {
|
||||
tie_match_range = Some(new_range);
|
||||
} else if new_cost < lowest_cost {
|
||||
tie_match_range.take();
|
||||
best_match = Some((new_cost, new_range));
|
||||
}
|
||||
}
|
||||
(None, Some(new_match)) => {
|
||||
best_match = Some(new_match);
|
||||
}
|
||||
(None, None) | (Some(_), None) => {}
|
||||
};
|
||||
}
|
||||
|
||||
if let Some((_, best_match_range)) = best_match {
|
||||
if let Some(tie_match_range) = tie_match_range {
|
||||
anyhow::bail!(
|
||||
"Multiple ambiguous matches:\n{:?}:\n{}\n\n{:?}:\n{}",
|
||||
best_match_range.clone(),
|
||||
buffer.text_for_range(best_match_range).collect::<String>(),
|
||||
tie_match_range.clone(),
|
||||
buffer.text_for_range(tie_match_range).collect::<String>()
|
||||
);
|
||||
}
|
||||
return Ok(best_match_range);
|
||||
}
|
||||
|
||||
anyhow::bail!(
|
||||
"Failed to fuzzy match `old_text`:\n{}\nin:\n```\n{}\n```",
|
||||
old_text,
|
||||
context_ranges
|
||||
.iter()
|
||||
.map(|range| buffer.text_for_range(range.clone()).collect::<String>())
|
||||
.collect::<Vec<String>>()
|
||||
.join("```\n```")
|
||||
);
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
struct XmlEdits<'a> {
|
||||
file_path: &'a str,
|
||||
/// Vec of (old_text, new_text) pairs
|
||||
replacements: Vec<(&'a str, &'a str)>,
|
||||
}
|
||||
|
||||
fn extract_xml_replacements(input: &str) -> Result<XmlEdits<'_>> {
|
||||
let mut cursor = 0;
|
||||
|
||||
let (edits_body_start, edits_attrs) =
|
||||
find_tag_open(input, &mut cursor, EDITS_TAG_NAME)?.context("No edits tag found")?;
|
||||
|
||||
let file_path = edits_attrs
|
||||
.trim_start()
|
||||
.strip_prefix("path")
|
||||
.context("no path attribute on edits tag")?
|
||||
.trim_end()
|
||||
.strip_prefix('=')
|
||||
.context("no value for path attribute")?
|
||||
.trim()
|
||||
.trim_start_matches('"')
|
||||
.trim_end_matches('"');
|
||||
|
||||
cursor = edits_body_start;
|
||||
let mut edits_list = Vec::new();
|
||||
|
||||
while let Some((old_body_start, _)) = find_tag_open(input, &mut cursor, OLD_TEXT_TAG_NAME)? {
|
||||
let old_body_end = find_tag_close(input, &mut cursor)?;
|
||||
let old_text = trim_surrounding_newlines(&input[old_body_start..old_body_end]);
|
||||
|
||||
let (new_body_start, _) = find_tag_open(input, &mut cursor, NEW_TEXT_TAG_NAME)?
|
||||
.context("no new_text tag following old_text")?;
|
||||
let new_body_end = find_tag_close(input, &mut cursor)?;
|
||||
let new_text = trim_surrounding_newlines(&input[new_body_start..new_body_end]);
|
||||
|
||||
edits_list.push((old_text, new_text));
|
||||
}
|
||||
|
||||
Ok(XmlEdits {
|
||||
file_path,
|
||||
replacements: edits_list,
|
||||
})
|
||||
}
|
||||
|
||||
/// Trims a single leading and trailing newline
|
||||
fn trim_surrounding_newlines(input: &str) -> &str {
|
||||
let start = input.strip_prefix('\n').unwrap_or(input);
|
||||
let end = start.strip_suffix('\n').unwrap_or(start);
|
||||
end
|
||||
}
|
||||
|
||||
fn find_tag_open<'a>(
|
||||
input: &'a str,
|
||||
cursor: &mut usize,
|
||||
expected_tag: &str,
|
||||
) -> Result<Option<(usize, &'a str)>> {
|
||||
let mut search_pos = *cursor;
|
||||
|
||||
while search_pos < input.len() {
|
||||
let Some(tag_start) = input[search_pos..].find("<") else {
|
||||
break;
|
||||
};
|
||||
let tag_start = search_pos + tag_start;
|
||||
if !input[tag_start + 1..].starts_with(expected_tag) {
|
||||
search_pos = search_pos + tag_start + 1;
|
||||
continue;
|
||||
};
|
||||
|
||||
let after_tag_name = tag_start + expected_tag.len() + 1;
|
||||
let close_bracket = input[after_tag_name..]
|
||||
.find('>')
|
||||
.with_context(|| format!("missing > after <{}", expected_tag))?;
|
||||
let attrs_end = after_tag_name + close_bracket;
|
||||
let body_start = attrs_end + 1;
|
||||
|
||||
let attributes = input[after_tag_name..attrs_end].trim();
|
||||
*cursor = body_start;
|
||||
|
||||
return Ok(Some((body_start, attributes)));
|
||||
}
|
||||
|
||||
Ok(None)
|
||||
}
|
||||
|
||||
fn find_tag_close(input: &str, cursor: &mut usize) -> Result<usize> {
|
||||
let mut depth = 1;
|
||||
let mut search_pos = *cursor;
|
||||
|
||||
while search_pos < input.len() && depth > 0 {
|
||||
let Some(bracket_offset) = input[search_pos..].find('<') else {
|
||||
break;
|
||||
};
|
||||
let bracket_pos = search_pos + bracket_offset;
|
||||
|
||||
if input[bracket_pos..].starts_with("</")
|
||||
&& let Some(close_end) = input[bracket_pos + 2..].find('>')
|
||||
{
|
||||
let close_start = bracket_pos + 2;
|
||||
let tag_name = input[close_start..close_start + close_end].trim();
|
||||
|
||||
if XML_TAGS.contains(&tag_name) {
|
||||
depth -= 1;
|
||||
if depth == 0 {
|
||||
*cursor = close_start + close_end + 1;
|
||||
return Ok(bracket_pos);
|
||||
}
|
||||
}
|
||||
search_pos = close_start + close_end + 1;
|
||||
continue;
|
||||
} else if let Some(close_bracket_offset) = input[bracket_pos..].find('>') {
|
||||
let close_bracket_pos = bracket_pos + close_bracket_offset;
|
||||
let tag_name = &input[bracket_pos + 1..close_bracket_pos].trim();
|
||||
if XML_TAGS.contains(&tag_name) {
|
||||
depth += 1;
|
||||
}
|
||||
}
|
||||
|
||||
search_pos = bracket_pos + 1;
|
||||
}
|
||||
|
||||
anyhow::bail!("no closing tag found")
|
||||
}
|
||||
|
||||
const REPLACEMENT_COST: u32 = 1;
|
||||
const INSERTION_COST: u32 = 3;
|
||||
const DELETION_COST: u32 = 10;
|
||||
|
||||
/// A fuzzy matcher that can process text chunks incrementally
|
||||
/// and return the best match found so far at each step.
|
||||
struct FuzzyMatcher<'a> {
|
||||
snapshot: &'a BufferSnapshot,
|
||||
query_lines: Vec<&'a str>,
|
||||
matrix: SearchMatrix,
|
||||
}
|
||||
|
||||
impl<'a> FuzzyMatcher<'a> {
|
||||
fn new(snapshot: &'a BufferSnapshot, old_text: &'a str) -> Self {
|
||||
let query_lines = old_text.lines().collect();
|
||||
Self {
|
||||
snapshot,
|
||||
query_lines,
|
||||
matrix: SearchMatrix::new(0),
|
||||
}
|
||||
}
|
||||
|
||||
fn match_range(&mut self, range: Range<usize>) -> Option<(u32, Range<usize>)> {
|
||||
let point_range = range.to_point(&self.snapshot);
|
||||
let buffer_line_count = (point_range.end.row - point_range.start.row + 1) as usize;
|
||||
|
||||
self.matrix
|
||||
.reset(self.query_lines.len() + 1, buffer_line_count + 1);
|
||||
let query_line_count = self.query_lines.len();
|
||||
|
||||
for row in 0..query_line_count {
|
||||
let query_line = self.query_lines[row].trim();
|
||||
let leading_deletion_cost = (row + 1) as u32 * DELETION_COST;
|
||||
|
||||
self.matrix.set(
|
||||
row + 1,
|
||||
0,
|
||||
SearchState::new(leading_deletion_cost, SearchDirection::Up),
|
||||
);
|
||||
|
||||
let mut buffer_lines = self.snapshot.text_for_range(range.clone()).lines();
|
||||
|
||||
let mut col = 0;
|
||||
while let Some(buffer_line) = buffer_lines.next() {
|
||||
let buffer_line = buffer_line.trim();
|
||||
let up = SearchState::new(
|
||||
self.matrix
|
||||
.get(row, col + 1)
|
||||
.cost
|
||||
.saturating_add(DELETION_COST),
|
||||
SearchDirection::Up,
|
||||
);
|
||||
let left = SearchState::new(
|
||||
self.matrix
|
||||
.get(row + 1, col)
|
||||
.cost
|
||||
.saturating_add(INSERTION_COST),
|
||||
SearchDirection::Left,
|
||||
);
|
||||
let diagonal = SearchState::new(
|
||||
if query_line == buffer_line {
|
||||
self.matrix.get(row, col).cost
|
||||
} else if fuzzy_eq(query_line, buffer_line) {
|
||||
self.matrix.get(row, col).cost + REPLACEMENT_COST
|
||||
} else {
|
||||
self.matrix
|
||||
.get(row, col)
|
||||
.cost
|
||||
.saturating_add(DELETION_COST + INSERTION_COST)
|
||||
},
|
||||
SearchDirection::Diagonal,
|
||||
);
|
||||
self.matrix
|
||||
.set(row + 1, col + 1, up.min(left).min(diagonal));
|
||||
col += 1;
|
||||
}
|
||||
}
|
||||
|
||||
// Find all matches with the best cost
|
||||
let mut best_cost = u32::MAX;
|
||||
let mut matches_with_best_cost = Vec::new();
|
||||
|
||||
for col in 1..=buffer_line_count {
|
||||
let cost = self.matrix.get(query_line_count, col).cost;
|
||||
if cost < best_cost {
|
||||
best_cost = cost;
|
||||
matches_with_best_cost.clear();
|
||||
matches_with_best_cost.push(col as u32);
|
||||
} else if cost == best_cost {
|
||||
matches_with_best_cost.push(col as u32);
|
||||
}
|
||||
}
|
||||
|
||||
// Find ranges for the matches
|
||||
for &match_end_col in &matches_with_best_cost {
|
||||
let mut matched_lines = 0;
|
||||
let mut query_row = query_line_count;
|
||||
let mut match_start_col = match_end_col;
|
||||
while query_row > 0 && match_start_col > 0 {
|
||||
let current = self.matrix.get(query_row, match_start_col as usize);
|
||||
match current.direction {
|
||||
SearchDirection::Diagonal => {
|
||||
query_row -= 1;
|
||||
match_start_col -= 1;
|
||||
matched_lines += 1;
|
||||
}
|
||||
SearchDirection::Up => {
|
||||
query_row -= 1;
|
||||
}
|
||||
SearchDirection::Left => {
|
||||
match_start_col -= 1;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let buffer_row_start = match_start_col + point_range.start.row;
|
||||
let buffer_row_end = match_end_col + point_range.start.row;
|
||||
|
||||
let matched_buffer_row_count = buffer_row_end - buffer_row_start;
|
||||
let matched_ratio = matched_lines as f32
|
||||
/ (matched_buffer_row_count as f32).max(query_line_count as f32);
|
||||
if matched_ratio >= 0.8 {
|
||||
let buffer_start_ix = self
|
||||
.snapshot
|
||||
.point_to_offset(Point::new(buffer_row_start, 0));
|
||||
let buffer_end_ix = self.snapshot.point_to_offset(Point::new(
|
||||
buffer_row_end - 1,
|
||||
self.snapshot.line_len(buffer_row_end - 1),
|
||||
));
|
||||
return Some((best_cost, buffer_start_ix..buffer_end_ix));
|
||||
}
|
||||
}
|
||||
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
fn fuzzy_eq(left: &str, right: &str) -> bool {
|
||||
const THRESHOLD: f64 = 0.8;
|
||||
|
||||
let min_levenshtein = left.len().abs_diff(right.len());
|
||||
let min_normalized_levenshtein =
|
||||
1. - (min_levenshtein as f64 / cmp::max(left.len(), right.len()) as f64);
|
||||
if min_normalized_levenshtein < THRESHOLD {
|
||||
return false;
|
||||
}
|
||||
|
||||
strsim::normalized_levenshtein(left, right) >= THRESHOLD
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord)]
|
||||
enum SearchDirection {
|
||||
Up,
|
||||
Left,
|
||||
Diagonal,
|
||||
}
|
||||
|
||||
#[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord)]
|
||||
struct SearchState {
|
||||
cost: u32,
|
||||
direction: SearchDirection,
|
||||
}
|
||||
|
||||
impl SearchState {
|
||||
fn new(cost: u32, direction: SearchDirection) -> Self {
|
||||
Self { cost, direction }
|
||||
}
|
||||
}
|
||||
|
||||
struct SearchMatrix {
|
||||
cols: usize,
|
||||
rows: usize,
|
||||
data: Vec<SearchState>,
|
||||
}
|
||||
|
||||
impl SearchMatrix {
|
||||
fn new(cols: usize) -> Self {
|
||||
SearchMatrix {
|
||||
cols,
|
||||
rows: 0,
|
||||
data: Vec::new(),
|
||||
}
|
||||
}
|
||||
|
||||
fn reset(&mut self, rows: usize, cols: usize) {
|
||||
self.rows = rows;
|
||||
self.cols = cols;
|
||||
self.data
|
||||
.fill(SearchState::new(0, SearchDirection::Diagonal));
|
||||
self.data.resize(
|
||||
self.rows * self.cols,
|
||||
SearchState::new(0, SearchDirection::Diagonal),
|
||||
);
|
||||
}
|
||||
|
||||
fn get(&self, row: usize, col: usize) -> SearchState {
|
||||
debug_assert!(row < self.rows);
|
||||
debug_assert!(col < self.cols);
|
||||
self.data[row * self.cols + col]
|
||||
}
|
||||
|
||||
fn set(&mut self, row: usize, col: usize, state: SearchState) {
|
||||
debug_assert!(row < self.rows && col < self.cols);
|
||||
self.data[row * self.cols + col] = state;
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use gpui::TestAppContext;
|
||||
use indoc::indoc;
|
||||
use language::Point;
|
||||
use project::{FakeFs, Project};
|
||||
use serde_json::json;
|
||||
use settings::SettingsStore;
|
||||
use util::path;
|
||||
|
||||
#[test]
|
||||
fn test_extract_xml_edits() {
|
||||
let input = indoc! {r#"
|
||||
<edits path="test.rs">
|
||||
<old_text>
|
||||
old content
|
||||
</old_text>
|
||||
<new_text>
|
||||
new content
|
||||
</new_text>
|
||||
</edits>
|
||||
"#};
|
||||
|
||||
let result = extract_xml_replacements(input).unwrap();
|
||||
assert_eq!(result.file_path, "test.rs");
|
||||
assert_eq!(result.replacements.len(), 1);
|
||||
assert_eq!(result.replacements[0].0, "old content");
|
||||
assert_eq!(result.replacements[0].1, "new content");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_extract_xml_edits_with_wrong_closing_tags() {
|
||||
let input = indoc! {r#"
|
||||
<edits path="test.rs">
|
||||
<old_text>
|
||||
old content
|
||||
</new_text>
|
||||
<new_text>
|
||||
new content
|
||||
</old_text>
|
||||
</ edits >
|
||||
"#};
|
||||
|
||||
let result = extract_xml_replacements(input).unwrap();
|
||||
assert_eq!(result.file_path, "test.rs");
|
||||
assert_eq!(result.replacements.len(), 1);
|
||||
assert_eq!(result.replacements[0].0, "old content");
|
||||
assert_eq!(result.replacements[0].1, "new content");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_extract_xml_edits_with_xml_like_content() {
|
||||
let input = indoc! {r#"
|
||||
<edits path="component.tsx">
|
||||
<old_text>
|
||||
<foo><bar></bar></foo>
|
||||
</old_text>
|
||||
<new_text>
|
||||
<foo><bar><baz></baz></bar></foo>
|
||||
</new_text>
|
||||
</edits>
|
||||
"#};
|
||||
|
||||
let result = extract_xml_replacements(input).unwrap();
|
||||
assert_eq!(result.file_path, "component.tsx");
|
||||
assert_eq!(result.replacements.len(), 1);
|
||||
assert_eq!(result.replacements[0].0, "<foo><bar></bar></foo>");
|
||||
assert_eq!(
|
||||
result.replacements[0].1,
|
||||
"<foo><bar><baz></baz></bar></foo>"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_extract_xml_edits_with_conflicting_content() {
|
||||
let input = indoc! {r#"
|
||||
<edits path="component.tsx">
|
||||
<old_text>
|
||||
<new_text></new_text>
|
||||
</old_text>
|
||||
<new_text>
|
||||
<old_text></old_text>
|
||||
</new_text>
|
||||
</edits>
|
||||
"#};
|
||||
|
||||
let result = extract_xml_replacements(input).unwrap();
|
||||
assert_eq!(result.file_path, "component.tsx");
|
||||
assert_eq!(result.replacements.len(), 1);
|
||||
assert_eq!(result.replacements[0].0, "<new_text></new_text>");
|
||||
assert_eq!(result.replacements[0].1, "<old_text></old_text>");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_extract_xml_edits_multiple_pairs() {
|
||||
let input = indoc! {r#"
|
||||
Some reasoning before edits. Lots of thinking going on here
|
||||
|
||||
<edits path="test.rs">
|
||||
<old_text>
|
||||
first old
|
||||
</old_text>
|
||||
<new_text>
|
||||
first new
|
||||
</new_text>
|
||||
<old_text>
|
||||
second old
|
||||
</edits>
|
||||
<new_text>
|
||||
second new
|
||||
</old_text>
|
||||
</edits>
|
||||
"#};
|
||||
|
||||
let result = extract_xml_replacements(input).unwrap();
|
||||
assert_eq!(result.file_path, "test.rs");
|
||||
assert_eq!(result.replacements.len(), 2);
|
||||
assert_eq!(result.replacements[0].0, "first old");
|
||||
assert_eq!(result.replacements[0].1, "first new");
|
||||
assert_eq!(result.replacements[1].0, "second old");
|
||||
assert_eq!(result.replacements[1].1, "second new");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_extract_xml_edits_unexpected_eof() {
|
||||
let input = indoc! {r#"
|
||||
<edits path="test.rs">
|
||||
<old_text>
|
||||
first old
|
||||
</
|
||||
"#};
|
||||
|
||||
extract_xml_replacements(input).expect_err("Unexpected end of file");
|
||||
}
|
||||
|
||||
#[gpui::test]
|
||||
async fn test_parse_xml_edits(cx: &mut TestAppContext) {
|
||||
let fs = init_test(cx);
|
||||
|
||||
let buffer_1_text = indoc! {r#"
|
||||
one two three four
|
||||
five six seven eight
|
||||
nine ten eleven twelve
|
||||
thirteen fourteen fifteen
|
||||
sixteen seventeen eighteen
|
||||
"#};
|
||||
|
||||
fs.insert_tree(
|
||||
path!("/root"),
|
||||
json!({
|
||||
"file1": buffer_1_text,
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
|
||||
let project = Project::test(fs, [path!("/root").as_ref()], cx).await;
|
||||
let buffer = project
|
||||
.update(cx, |project, cx| {
|
||||
project.open_local_buffer(path!("/root/file1"), cx)
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
let buffer_snapshot = buffer.read_with(cx, |buffer, _| buffer.snapshot());
|
||||
|
||||
let edits = indoc! {r#"
|
||||
<edits path="root/file1">
|
||||
<old_text>
|
||||
nine ten eleven twelve
|
||||
</old_text>
|
||||
<new_text>
|
||||
nine TEN eleven twelve!
|
||||
</new_text>
|
||||
</edits>
|
||||
"#};
|
||||
|
||||
let included_ranges = [(buffer_snapshot.anchor_before(Point::new(1, 0))..Anchor::MAX)];
|
||||
let (buffer, edits) = parse_xml_edits(edits, |_path| {
|
||||
Some((&buffer_snapshot, included_ranges.as_slice()))
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let edits = edits
|
||||
.into_iter()
|
||||
.map(|(range, text)| (range.to_point(&buffer), text))
|
||||
.collect::<Vec<_>>();
|
||||
assert_eq!(
|
||||
edits,
|
||||
&[
|
||||
(Point::new(2, 5)..Point::new(2, 8), "TEN".into()),
|
||||
(Point::new(2, 22)..Point::new(2, 22), "!".into())
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
fn init_test(cx: &mut TestAppContext) -> Arc<FakeFs> {
|
||||
cx.update(|cx| {
|
||||
let settings_store = SettingsStore::test(cx);
|
||||
cx.set_global(settings_store);
|
||||
});
|
||||
|
||||
FakeFs::new(cx.background_executor.clone())
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,239 @@
|
||||
use std::{cmp, sync::Arc};
|
||||
|
||||
use client::{Client, UserStore};
|
||||
use cloud_llm_client::EditPredictionRejectReason;
|
||||
use edit_prediction_types::{DataCollectionState, Direction, EditPredictionDelegate};
|
||||
use gpui::{App, Entity, prelude::*};
|
||||
use language::{Buffer, ToPoint as _};
|
||||
use project::Project;
|
||||
|
||||
use crate::{BufferEditPrediction, EditPredictionModel, EditPredictionStore};
|
||||
|
||||
pub struct ZedEditPredictionDelegate {
|
||||
store: Entity<EditPredictionStore>,
|
||||
project: Entity<Project>,
|
||||
singleton_buffer: Option<Entity<Buffer>>,
|
||||
}
|
||||
|
||||
impl ZedEditPredictionDelegate {
|
||||
pub fn new(
|
||||
project: Entity<Project>,
|
||||
singleton_buffer: Option<Entity<Buffer>>,
|
||||
client: &Arc<Client>,
|
||||
user_store: &Entity<UserStore>,
|
||||
cx: &mut Context<Self>,
|
||||
) -> Self {
|
||||
let store = EditPredictionStore::global(client, user_store, cx);
|
||||
store.update(cx, |store, cx| {
|
||||
store.register_project(&project, cx);
|
||||
});
|
||||
|
||||
cx.observe(&store, |_this, _ep_store, cx| {
|
||||
cx.notify();
|
||||
})
|
||||
.detach();
|
||||
|
||||
Self {
|
||||
project: project,
|
||||
store: store,
|
||||
singleton_buffer,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl EditPredictionDelegate for ZedEditPredictionDelegate {
|
||||
fn name() -> &'static str {
|
||||
"zed-predict"
|
||||
}
|
||||
|
||||
fn display_name() -> &'static str {
|
||||
"Zed's Edit Predictions"
|
||||
}
|
||||
|
||||
fn show_predictions_in_menu() -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
fn show_tab_accept_marker() -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
fn data_collection_state(&self, cx: &App) -> DataCollectionState {
|
||||
if let Some(buffer) = &self.singleton_buffer
|
||||
&& let Some(file) = buffer.read(cx).file()
|
||||
{
|
||||
let is_project_open_source =
|
||||
self.store
|
||||
.read(cx)
|
||||
.is_file_open_source(&self.project, file, cx);
|
||||
if self.store.read(cx).data_collection_choice.is_enabled() {
|
||||
DataCollectionState::Enabled {
|
||||
is_project_open_source,
|
||||
}
|
||||
} else {
|
||||
DataCollectionState::Disabled {
|
||||
is_project_open_source,
|
||||
}
|
||||
}
|
||||
} else {
|
||||
return DataCollectionState::Disabled {
|
||||
is_project_open_source: false,
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
fn toggle_data_collection(&mut self, cx: &mut App) {
|
||||
self.store.update(cx, |store, cx| {
|
||||
store.toggle_data_collection_choice(cx);
|
||||
});
|
||||
}
|
||||
|
||||
fn usage(&self, cx: &App) -> Option<client::EditPredictionUsage> {
|
||||
self.store.read(cx).usage(cx)
|
||||
}
|
||||
|
||||
fn is_enabled(
|
||||
&self,
|
||||
_buffer: &Entity<language::Buffer>,
|
||||
_cursor_position: language::Anchor,
|
||||
cx: &App,
|
||||
) -> bool {
|
||||
let store = self.store.read(cx);
|
||||
if store.edit_prediction_model == EditPredictionModel::Sweep {
|
||||
store.has_sweep_api_token()
|
||||
} else {
|
||||
true
|
||||
}
|
||||
}
|
||||
|
||||
fn is_refreshing(&self, cx: &App) -> bool {
|
||||
self.store.read(cx).is_refreshing(&self.project)
|
||||
}
|
||||
|
||||
fn refresh(
|
||||
&mut self,
|
||||
buffer: Entity<language::Buffer>,
|
||||
cursor_position: language::Anchor,
|
||||
_debounce: bool,
|
||||
cx: &mut Context<Self>,
|
||||
) {
|
||||
let store = self.store.read(cx);
|
||||
|
||||
if store.user_store.read_with(cx, |user_store, _cx| {
|
||||
user_store.account_too_young() || user_store.has_overdue_invoices()
|
||||
}) {
|
||||
return;
|
||||
}
|
||||
|
||||
if let Some(current) = store.current_prediction_for_buffer(&buffer, &self.project, cx)
|
||||
&& let BufferEditPrediction::Local { prediction } = current
|
||||
&& prediction.interpolate(buffer.read(cx)).is_some()
|
||||
{
|
||||
return;
|
||||
}
|
||||
|
||||
self.store.update(cx, |store, cx| {
|
||||
store.refresh_context(&self.project, &buffer, cursor_position, cx);
|
||||
store.refresh_prediction_from_buffer(self.project.clone(), buffer, cursor_position, cx)
|
||||
});
|
||||
}
|
||||
|
||||
fn cycle(
|
||||
&mut self,
|
||||
_buffer: Entity<language::Buffer>,
|
||||
_cursor_position: language::Anchor,
|
||||
_direction: Direction,
|
||||
_cx: &mut Context<Self>,
|
||||
) {
|
||||
}
|
||||
|
||||
fn accept(&mut self, cx: &mut Context<Self>) {
|
||||
self.store.update(cx, |store, cx| {
|
||||
store.accept_current_prediction(&self.project, cx);
|
||||
});
|
||||
}
|
||||
|
||||
fn discard(&mut self, cx: &mut Context<Self>) {
|
||||
self.store.update(cx, |store, _cx| {
|
||||
store.reject_current_prediction(EditPredictionRejectReason::Discarded, &self.project);
|
||||
});
|
||||
}
|
||||
|
||||
fn did_show(&mut self, cx: &mut Context<Self>) {
|
||||
self.store.update(cx, |store, cx| {
|
||||
store.did_show_current_prediction(&self.project, cx);
|
||||
});
|
||||
}
|
||||
|
||||
fn suggest(
|
||||
&mut self,
|
||||
buffer: &Entity<language::Buffer>,
|
||||
cursor_position: language::Anchor,
|
||||
cx: &mut Context<Self>,
|
||||
) -> Option<edit_prediction_types::EditPrediction> {
|
||||
let prediction =
|
||||
self.store
|
||||
.read(cx)
|
||||
.current_prediction_for_buffer(buffer, &self.project, cx)?;
|
||||
|
||||
let prediction = match prediction {
|
||||
BufferEditPrediction::Local { prediction } => prediction,
|
||||
BufferEditPrediction::Jump { prediction } => {
|
||||
return Some(edit_prediction_types::EditPrediction::Jump {
|
||||
id: Some(prediction.id.to_string().into()),
|
||||
snapshot: prediction.snapshot.clone(),
|
||||
target: prediction.edits.first().unwrap().0.start,
|
||||
});
|
||||
}
|
||||
};
|
||||
|
||||
let buffer = buffer.read(cx);
|
||||
let snapshot = buffer.snapshot();
|
||||
|
||||
let Some(edits) = prediction.interpolate(&snapshot) else {
|
||||
self.store.update(cx, |store, _cx| {
|
||||
store.reject_current_prediction(
|
||||
EditPredictionRejectReason::InterpolatedEmpty,
|
||||
&self.project,
|
||||
);
|
||||
});
|
||||
return None;
|
||||
};
|
||||
|
||||
let cursor_row = cursor_position.to_point(&snapshot).row;
|
||||
let (closest_edit_ix, (closest_edit_range, _)) =
|
||||
edits.iter().enumerate().min_by_key(|(_, (range, _))| {
|
||||
let distance_from_start = cursor_row.abs_diff(range.start.to_point(&snapshot).row);
|
||||
let distance_from_end = cursor_row.abs_diff(range.end.to_point(&snapshot).row);
|
||||
cmp::min(distance_from_start, distance_from_end)
|
||||
})?;
|
||||
|
||||
let mut edit_start_ix = closest_edit_ix;
|
||||
for (range, _) in edits[..edit_start_ix].iter().rev() {
|
||||
let distance_from_closest_edit = closest_edit_range.start.to_point(&snapshot).row
|
||||
- range.end.to_point(&snapshot).row;
|
||||
if distance_from_closest_edit <= 1 {
|
||||
edit_start_ix -= 1;
|
||||
} else {
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
let mut edit_end_ix = closest_edit_ix + 1;
|
||||
for (range, _) in &edits[edit_end_ix..] {
|
||||
let distance_from_closest_edit =
|
||||
range.start.to_point(buffer).row - closest_edit_range.end.to_point(&snapshot).row;
|
||||
if distance_from_closest_edit <= 1 {
|
||||
edit_end_ix += 1;
|
||||
} else {
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
Some(edit_prediction_types::EditPrediction::Local {
|
||||
id: Some(prediction.id.to_string().into()),
|
||||
edits: edits[edit_start_ix..edit_end_ix].to_vec(),
|
||||
edit_preview: Some(prediction.edit_preview.clone()),
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,504 @@
|
||||
mod input_excerpt;
|
||||
|
||||
use std::{fmt::Write, ops::Range, path::Path, sync::Arc, time::Instant};
|
||||
|
||||
use crate::{
|
||||
EditPredictionId, EditPredictionStore, ZedUpdateRequiredError,
|
||||
prediction::{EditPredictionInputs, EditPredictionResult},
|
||||
};
|
||||
use anyhow::{Context as _, Result};
|
||||
use cloud_llm_client::{
|
||||
PredictEditsBody, PredictEditsGitInfo, PredictEditsRequestTrigger, PredictEditsResponse,
|
||||
predict_edits_v3::Event,
|
||||
};
|
||||
use gpui::{App, AppContext as _, AsyncApp, Context, Entity, SharedString, Task};
|
||||
use input_excerpt::excerpt_for_cursor_position;
|
||||
use language::{
|
||||
Anchor, Buffer, BufferSnapshot, OffsetRangeExt as _, Point, ToPoint as _, text_diff,
|
||||
};
|
||||
use project::{Project, ProjectPath};
|
||||
use release_channel::AppVersion;
|
||||
use workspace::notifications::{ErrorMessagePrompt, NotificationId, show_app_notification};
|
||||
|
||||
const CURSOR_MARKER: &str = "<|user_cursor_is_here|>";
|
||||
const START_OF_FILE_MARKER: &str = "<|start_of_file|>";
|
||||
const EDITABLE_REGION_START_MARKER: &str = "<|editable_region_start|>";
|
||||
const EDITABLE_REGION_END_MARKER: &str = "<|editable_region_end|>";
|
||||
|
||||
pub(crate) const MAX_CONTEXT_TOKENS: usize = 150;
|
||||
pub(crate) const MAX_REWRITE_TOKENS: usize = 350;
|
||||
pub(crate) const MAX_EVENT_TOKENS: usize = 500;
|
||||
|
||||
pub(crate) fn request_prediction_with_zeta1(
|
||||
store: &mut EditPredictionStore,
|
||||
project: &Entity<Project>,
|
||||
buffer: &Entity<Buffer>,
|
||||
snapshot: BufferSnapshot,
|
||||
position: language::Anchor,
|
||||
events: Vec<Arc<Event>>,
|
||||
trigger: PredictEditsRequestTrigger,
|
||||
cx: &mut Context<EditPredictionStore>,
|
||||
) -> Task<Result<Option<EditPredictionResult>>> {
|
||||
let buffer = buffer.clone();
|
||||
let buffer_snapshotted_at = Instant::now();
|
||||
let client = store.client.clone();
|
||||
let llm_token = store.llm_token.clone();
|
||||
let app_version = AppVersion::global(cx);
|
||||
|
||||
let (git_info, can_collect_file) = if let Some(file) = snapshot.file() {
|
||||
let can_collect_file = store.can_collect_file(project, file, cx);
|
||||
let git_info = if can_collect_file {
|
||||
git_info_for_file(project, &ProjectPath::from_file(file.as_ref(), cx), cx)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
(git_info, can_collect_file)
|
||||
} else {
|
||||
(None, false)
|
||||
};
|
||||
|
||||
let full_path: Arc<Path> = snapshot
|
||||
.file()
|
||||
.map(|f| Arc::from(f.full_path(cx).as_path()))
|
||||
.unwrap_or_else(|| Arc::from(Path::new("untitled")));
|
||||
let full_path_str = full_path.to_string_lossy().into_owned();
|
||||
let cursor_point = position.to_point(&snapshot);
|
||||
let prompt_for_events = {
|
||||
let events = events.clone();
|
||||
move || prompt_for_events_impl(&events, MAX_EVENT_TOKENS)
|
||||
};
|
||||
let gather_task = gather_context(
|
||||
full_path_str,
|
||||
&snapshot,
|
||||
cursor_point,
|
||||
prompt_for_events,
|
||||
trigger,
|
||||
cx,
|
||||
);
|
||||
|
||||
cx.spawn(async move |this, cx| {
|
||||
let GatherContextOutput {
|
||||
mut body,
|
||||
context_range,
|
||||
editable_range,
|
||||
included_events_count,
|
||||
} = gather_task.await?;
|
||||
let done_gathering_context_at = Instant::now();
|
||||
|
||||
let included_events = &events[events.len() - included_events_count..events.len()];
|
||||
body.can_collect_data = can_collect_file
|
||||
&& this
|
||||
.read_with(cx, |this, _| this.can_collect_events(included_events))
|
||||
.unwrap_or(false);
|
||||
if body.can_collect_data {
|
||||
body.git_info = git_info;
|
||||
}
|
||||
|
||||
log::debug!(
|
||||
"Events:\n{}\nExcerpt:\n{:?}",
|
||||
body.input_events,
|
||||
body.input_excerpt
|
||||
);
|
||||
|
||||
let http_client = client.http_client();
|
||||
|
||||
let response = EditPredictionStore::send_api_request::<PredictEditsResponse>(
|
||||
|request| {
|
||||
let uri = if let Ok(predict_edits_url) = std::env::var("ZED_PREDICT_EDITS_URL") {
|
||||
predict_edits_url
|
||||
} else {
|
||||
http_client
|
||||
.build_zed_llm_url("/predict_edits/v2", &[])?
|
||||
.as_str()
|
||||
.into()
|
||||
};
|
||||
Ok(request
|
||||
.uri(uri)
|
||||
.body(serde_json::to_string(&body)?.into())?)
|
||||
},
|
||||
client,
|
||||
llm_token,
|
||||
app_version,
|
||||
)
|
||||
.await;
|
||||
|
||||
let inputs = EditPredictionInputs {
|
||||
events: included_events.into(),
|
||||
included_files: vec![cloud_llm_client::predict_edits_v3::RelatedFile {
|
||||
path: full_path.clone(),
|
||||
max_row: cloud_llm_client::predict_edits_v3::Line(snapshot.max_point().row),
|
||||
excerpts: vec![cloud_llm_client::predict_edits_v3::Excerpt {
|
||||
start_line: cloud_llm_client::predict_edits_v3::Line(context_range.start.row),
|
||||
text: snapshot
|
||||
.text_for_range(context_range)
|
||||
.collect::<String>()
|
||||
.into(),
|
||||
}],
|
||||
}],
|
||||
cursor_point: cloud_llm_client::predict_edits_v3::Point {
|
||||
column: cursor_point.column,
|
||||
line: cloud_llm_client::predict_edits_v3::Line(cursor_point.row),
|
||||
},
|
||||
cursor_path: full_path,
|
||||
};
|
||||
|
||||
// let response = perform_predict_edits(PerformPredictEditsParams {
|
||||
// client,
|
||||
// llm_token,
|
||||
// app_version,
|
||||
// body,
|
||||
// })
|
||||
// .await;
|
||||
|
||||
let (response, usage) = match response {
|
||||
Ok(response) => response,
|
||||
Err(err) => {
|
||||
if err.is::<ZedUpdateRequiredError>() {
|
||||
cx.update(|cx| {
|
||||
this.update(cx, |ep_store, _cx| {
|
||||
ep_store.update_required = true;
|
||||
})
|
||||
.ok();
|
||||
|
||||
let error_message: SharedString = err.to_string().into();
|
||||
show_app_notification(
|
||||
NotificationId::unique::<ZedUpdateRequiredError>(),
|
||||
cx,
|
||||
move |cx| {
|
||||
cx.new(|cx| {
|
||||
ErrorMessagePrompt::new(error_message.clone(), cx)
|
||||
.with_link_button("Update Zed", "https://zed.dev/releases")
|
||||
})
|
||||
},
|
||||
);
|
||||
})
|
||||
.ok();
|
||||
}
|
||||
|
||||
return Err(err);
|
||||
}
|
||||
};
|
||||
|
||||
let received_response_at = Instant::now();
|
||||
log::debug!("completion response: {}", &response.output_excerpt);
|
||||
|
||||
if let Some(usage) = usage {
|
||||
this.update(cx, |this, cx| {
|
||||
this.user_store.update(cx, |user_store, cx| {
|
||||
user_store.update_edit_prediction_usage(usage, cx);
|
||||
});
|
||||
})
|
||||
.ok();
|
||||
}
|
||||
|
||||
let edit_prediction = process_completion_response(
|
||||
response,
|
||||
buffer,
|
||||
&snapshot,
|
||||
editable_range,
|
||||
inputs,
|
||||
buffer_snapshotted_at,
|
||||
received_response_at,
|
||||
cx,
|
||||
)
|
||||
.await;
|
||||
|
||||
let finished_at = Instant::now();
|
||||
|
||||
// record latency for ~1% of requests
|
||||
if rand::random::<u8>() <= 2 {
|
||||
telemetry::event!(
|
||||
"Edit Prediction Request",
|
||||
context_latency = done_gathering_context_at
|
||||
.duration_since(buffer_snapshotted_at)
|
||||
.as_millis(),
|
||||
request_latency = received_response_at
|
||||
.duration_since(done_gathering_context_at)
|
||||
.as_millis(),
|
||||
process_latency = finished_at.duration_since(received_response_at).as_millis()
|
||||
);
|
||||
}
|
||||
|
||||
edit_prediction.map(Some)
|
||||
})
|
||||
}
|
||||
|
||||
fn process_completion_response(
|
||||
prediction_response: PredictEditsResponse,
|
||||
buffer: Entity<Buffer>,
|
||||
snapshot: &BufferSnapshot,
|
||||
editable_range: Range<usize>,
|
||||
inputs: EditPredictionInputs,
|
||||
buffer_snapshotted_at: Instant,
|
||||
received_response_at: Instant,
|
||||
cx: &AsyncApp,
|
||||
) -> Task<Result<EditPredictionResult>> {
|
||||
let snapshot = snapshot.clone();
|
||||
let request_id = prediction_response.request_id;
|
||||
let output_excerpt = prediction_response.output_excerpt;
|
||||
cx.spawn(async move |cx| {
|
||||
let output_excerpt: Arc<str> = output_excerpt.into();
|
||||
|
||||
let edits: Arc<[(Range<Anchor>, Arc<str>)]> = cx
|
||||
.background_spawn({
|
||||
let output_excerpt = output_excerpt.clone();
|
||||
let editable_range = editable_range.clone();
|
||||
let snapshot = snapshot.clone();
|
||||
async move { parse_edits(output_excerpt, editable_range, &snapshot) }
|
||||
})
|
||||
.await?
|
||||
.into();
|
||||
|
||||
let id = EditPredictionId(request_id.into());
|
||||
Ok(EditPredictionResult::new(
|
||||
id,
|
||||
&buffer,
|
||||
&snapshot,
|
||||
edits,
|
||||
buffer_snapshotted_at,
|
||||
received_response_at,
|
||||
inputs,
|
||||
cx,
|
||||
)
|
||||
.await)
|
||||
})
|
||||
}
|
||||
|
||||
fn parse_edits(
|
||||
output_excerpt: Arc<str>,
|
||||
editable_range: Range<usize>,
|
||||
snapshot: &BufferSnapshot,
|
||||
) -> Result<Vec<(Range<Anchor>, Arc<str>)>> {
|
||||
let content = output_excerpt.replace(CURSOR_MARKER, "");
|
||||
|
||||
let start_markers = content
|
||||
.match_indices(EDITABLE_REGION_START_MARKER)
|
||||
.collect::<Vec<_>>();
|
||||
anyhow::ensure!(
|
||||
start_markers.len() == 1,
|
||||
"expected exactly one start marker, found {}",
|
||||
start_markers.len()
|
||||
);
|
||||
|
||||
let end_markers = content
|
||||
.match_indices(EDITABLE_REGION_END_MARKER)
|
||||
.collect::<Vec<_>>();
|
||||
anyhow::ensure!(
|
||||
end_markers.len() == 1,
|
||||
"expected exactly one end marker, found {}",
|
||||
end_markers.len()
|
||||
);
|
||||
|
||||
let sof_markers = content
|
||||
.match_indices(START_OF_FILE_MARKER)
|
||||
.collect::<Vec<_>>();
|
||||
anyhow::ensure!(
|
||||
sof_markers.len() <= 1,
|
||||
"expected at most one start-of-file marker, found {}",
|
||||
sof_markers.len()
|
||||
);
|
||||
|
||||
let codefence_start = start_markers[0].0;
|
||||
let content = &content[codefence_start..];
|
||||
|
||||
let newline_ix = content.find('\n').context("could not find newline")?;
|
||||
let content = &content[newline_ix + 1..];
|
||||
|
||||
let codefence_end = content
|
||||
.rfind(&format!("\n{EDITABLE_REGION_END_MARKER}"))
|
||||
.context("could not find end marker")?;
|
||||
let new_text = &content[..codefence_end];
|
||||
|
||||
let old_text = snapshot
|
||||
.text_for_range(editable_range.clone())
|
||||
.collect::<String>();
|
||||
|
||||
Ok(compute_edits(
|
||||
old_text,
|
||||
new_text,
|
||||
editable_range.start,
|
||||
snapshot,
|
||||
))
|
||||
}
|
||||
|
||||
pub fn compute_edits(
|
||||
old_text: String,
|
||||
new_text: &str,
|
||||
offset: usize,
|
||||
snapshot: &BufferSnapshot,
|
||||
) -> Vec<(Range<Anchor>, Arc<str>)> {
|
||||
text_diff(&old_text, new_text)
|
||||
.into_iter()
|
||||
.map(|(mut old_range, new_text)| {
|
||||
old_range.start += offset;
|
||||
old_range.end += offset;
|
||||
|
||||
let prefix_len = common_prefix(
|
||||
snapshot.chars_for_range(old_range.clone()),
|
||||
new_text.chars(),
|
||||
);
|
||||
old_range.start += prefix_len;
|
||||
|
||||
let suffix_len = common_prefix(
|
||||
snapshot.reversed_chars_for_range(old_range.clone()),
|
||||
new_text[prefix_len..].chars().rev(),
|
||||
);
|
||||
old_range.end = old_range.end.saturating_sub(suffix_len);
|
||||
|
||||
let new_text = new_text[prefix_len..new_text.len() - suffix_len].into();
|
||||
let range = if old_range.is_empty() {
|
||||
let anchor = snapshot.anchor_after(old_range.start);
|
||||
anchor..anchor
|
||||
} else {
|
||||
snapshot.anchor_after(old_range.start)..snapshot.anchor_before(old_range.end)
|
||||
};
|
||||
(range, new_text)
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn common_prefix<T1: Iterator<Item = char>, T2: Iterator<Item = char>>(a: T1, b: T2) -> usize {
|
||||
a.zip(b)
|
||||
.take_while(|(a, b)| a == b)
|
||||
.map(|(a, _)| a.len_utf8())
|
||||
.sum()
|
||||
}
|
||||
|
||||
fn git_info_for_file(
|
||||
project: &Entity<Project>,
|
||||
project_path: &ProjectPath,
|
||||
cx: &App,
|
||||
) -> Option<PredictEditsGitInfo> {
|
||||
let git_store = project.read(cx).git_store().read(cx);
|
||||
if let Some((repository, _repo_path)) =
|
||||
git_store.repository_and_path_for_project_path(project_path, cx)
|
||||
{
|
||||
let repository = repository.read(cx);
|
||||
let head_sha = repository
|
||||
.head_commit
|
||||
.as_ref()
|
||||
.map(|head_commit| head_commit.sha.to_string());
|
||||
let remote_origin_url = repository.remote_origin_url.clone();
|
||||
let remote_upstream_url = repository.remote_upstream_url.clone();
|
||||
if head_sha.is_none() && remote_origin_url.is_none() && remote_upstream_url.is_none() {
|
||||
return None;
|
||||
}
|
||||
Some(PredictEditsGitInfo {
|
||||
head_sha,
|
||||
remote_origin_url,
|
||||
remote_upstream_url,
|
||||
})
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
pub struct GatherContextOutput {
|
||||
pub body: PredictEditsBody,
|
||||
pub context_range: Range<Point>,
|
||||
pub editable_range: Range<usize>,
|
||||
pub included_events_count: usize,
|
||||
}
|
||||
|
||||
pub fn gather_context(
|
||||
full_path_str: String,
|
||||
snapshot: &BufferSnapshot,
|
||||
cursor_point: language::Point,
|
||||
prompt_for_events: impl FnOnce() -> (String, usize) + Send + 'static,
|
||||
trigger: PredictEditsRequestTrigger,
|
||||
cx: &App,
|
||||
) -> Task<Result<GatherContextOutput>> {
|
||||
cx.background_spawn({
|
||||
let snapshot = snapshot.clone();
|
||||
async move {
|
||||
let input_excerpt = excerpt_for_cursor_position(
|
||||
cursor_point,
|
||||
&full_path_str,
|
||||
&snapshot,
|
||||
MAX_REWRITE_TOKENS,
|
||||
MAX_CONTEXT_TOKENS,
|
||||
);
|
||||
let (input_events, included_events_count) = prompt_for_events();
|
||||
let editable_range = input_excerpt.editable_range.to_offset(&snapshot);
|
||||
|
||||
let body = PredictEditsBody {
|
||||
input_events,
|
||||
input_excerpt: input_excerpt.prompt,
|
||||
can_collect_data: false,
|
||||
diagnostic_groups: None,
|
||||
git_info: None,
|
||||
outline: None,
|
||||
speculated_output: None,
|
||||
trigger,
|
||||
};
|
||||
|
||||
Ok(GatherContextOutput {
|
||||
body,
|
||||
context_range: input_excerpt.context_range,
|
||||
editable_range,
|
||||
included_events_count,
|
||||
})
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
fn prompt_for_events_impl(events: &[Arc<Event>], mut remaining_tokens: usize) -> (String, usize) {
|
||||
let mut result = String::new();
|
||||
for (ix, event) in events.iter().rev().enumerate() {
|
||||
let event_string = format_event(event.as_ref());
|
||||
let event_tokens = guess_token_count(event_string.len());
|
||||
if event_tokens > remaining_tokens {
|
||||
return (result, ix);
|
||||
}
|
||||
|
||||
if !result.is_empty() {
|
||||
result.insert_str(0, "\n\n");
|
||||
}
|
||||
result.insert_str(0, &event_string);
|
||||
remaining_tokens -= event_tokens;
|
||||
}
|
||||
return (result, events.len());
|
||||
}
|
||||
|
||||
pub fn format_event(event: &Event) -> String {
|
||||
match event {
|
||||
Event::BufferChange {
|
||||
path,
|
||||
old_path,
|
||||
diff,
|
||||
..
|
||||
} => {
|
||||
let mut prompt = String::new();
|
||||
|
||||
if old_path != path {
|
||||
writeln!(
|
||||
prompt,
|
||||
"User renamed {} to {}\n",
|
||||
old_path.display(),
|
||||
path.display()
|
||||
)
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
if !diff.is_empty() {
|
||||
write!(
|
||||
prompt,
|
||||
"User edited {}:\n```diff\n{}\n```",
|
||||
path.display(),
|
||||
diff
|
||||
)
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
prompt
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Typical number of string bytes per token for the purposes of limiting model input. This is
|
||||
/// intentionally low to err on the side of underestimating limits.
|
||||
pub(crate) const BYTES_PER_TOKEN_GUESS: usize = 3;
|
||||
|
||||
fn guess_token_count(bytes: usize) -> usize {
|
||||
bytes / BYTES_PER_TOKEN_GUESS
|
||||
}
|
||||
@@ -0,0 +1,231 @@
|
||||
use super::{
|
||||
CURSOR_MARKER, EDITABLE_REGION_END_MARKER, EDITABLE_REGION_START_MARKER, START_OF_FILE_MARKER,
|
||||
guess_token_count,
|
||||
};
|
||||
use language::{BufferSnapshot, Point};
|
||||
use std::{fmt::Write, ops::Range};
|
||||
|
||||
#[derive(Debug)]
|
||||
pub struct InputExcerpt {
|
||||
pub context_range: Range<Point>,
|
||||
pub editable_range: Range<Point>,
|
||||
pub prompt: String,
|
||||
}
|
||||
|
||||
pub fn excerpt_for_cursor_position(
|
||||
position: Point,
|
||||
path: &str,
|
||||
snapshot: &BufferSnapshot,
|
||||
editable_region_token_limit: usize,
|
||||
context_token_limit: usize,
|
||||
) -> InputExcerpt {
|
||||
let mut scope_range = position..position;
|
||||
let mut remaining_edit_tokens = editable_region_token_limit;
|
||||
|
||||
while let Some(parent) = snapshot.syntax_ancestor(scope_range.clone()) {
|
||||
let parent_tokens = guess_token_count(parent.byte_range().len());
|
||||
let parent_point_range = Point::new(
|
||||
parent.start_position().row as u32,
|
||||
parent.start_position().column as u32,
|
||||
)
|
||||
..Point::new(
|
||||
parent.end_position().row as u32,
|
||||
parent.end_position().column as u32,
|
||||
);
|
||||
if parent_point_range == scope_range {
|
||||
break;
|
||||
} else if parent_tokens <= editable_region_token_limit {
|
||||
scope_range = parent_point_range;
|
||||
remaining_edit_tokens = editable_region_token_limit - parent_tokens;
|
||||
} else {
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
let editable_range = expand_range(snapshot, scope_range, remaining_edit_tokens);
|
||||
let context_range = expand_range(snapshot, editable_range.clone(), context_token_limit);
|
||||
|
||||
let mut prompt = String::new();
|
||||
|
||||
writeln!(&mut prompt, "```{path}").unwrap();
|
||||
if context_range.start == Point::zero() {
|
||||
writeln!(&mut prompt, "{START_OF_FILE_MARKER}").unwrap();
|
||||
}
|
||||
|
||||
for chunk in snapshot.chunks(context_range.start..editable_range.start, false) {
|
||||
prompt.push_str(chunk.text);
|
||||
}
|
||||
|
||||
push_editable_range(position, snapshot, editable_range.clone(), &mut prompt);
|
||||
|
||||
for chunk in snapshot.chunks(editable_range.end..context_range.end, false) {
|
||||
prompt.push_str(chunk.text);
|
||||
}
|
||||
write!(prompt, "\n```").unwrap();
|
||||
|
||||
InputExcerpt {
|
||||
context_range,
|
||||
editable_range,
|
||||
prompt,
|
||||
}
|
||||
}
|
||||
|
||||
fn push_editable_range(
|
||||
cursor_position: Point,
|
||||
snapshot: &BufferSnapshot,
|
||||
editable_range: Range<Point>,
|
||||
prompt: &mut String,
|
||||
) {
|
||||
writeln!(prompt, "{EDITABLE_REGION_START_MARKER}").unwrap();
|
||||
for chunk in snapshot.chunks(editable_range.start..cursor_position, false) {
|
||||
prompt.push_str(chunk.text);
|
||||
}
|
||||
prompt.push_str(CURSOR_MARKER);
|
||||
for chunk in snapshot.chunks(cursor_position..editable_range.end, false) {
|
||||
prompt.push_str(chunk.text);
|
||||
}
|
||||
write!(prompt, "\n{EDITABLE_REGION_END_MARKER}").unwrap();
|
||||
}
|
||||
|
||||
fn expand_range(
|
||||
snapshot: &BufferSnapshot,
|
||||
range: Range<Point>,
|
||||
mut remaining_tokens: usize,
|
||||
) -> Range<Point> {
|
||||
let mut expanded_range = range;
|
||||
expanded_range.start.column = 0;
|
||||
expanded_range.end.column = snapshot.line_len(expanded_range.end.row);
|
||||
loop {
|
||||
let mut expanded = false;
|
||||
|
||||
if remaining_tokens > 0 && expanded_range.start.row > 0 {
|
||||
expanded_range.start.row -= 1;
|
||||
let line_tokens =
|
||||
guess_token_count(snapshot.line_len(expanded_range.start.row) as usize);
|
||||
remaining_tokens = remaining_tokens.saturating_sub(line_tokens);
|
||||
expanded = true;
|
||||
}
|
||||
|
||||
if remaining_tokens > 0 && expanded_range.end.row < snapshot.max_point().row {
|
||||
expanded_range.end.row += 1;
|
||||
expanded_range.end.column = snapshot.line_len(expanded_range.end.row);
|
||||
let line_tokens = guess_token_count(expanded_range.end.column as usize);
|
||||
remaining_tokens = remaining_tokens.saturating_sub(line_tokens);
|
||||
expanded = true;
|
||||
}
|
||||
|
||||
if !expanded {
|
||||
break;
|
||||
}
|
||||
}
|
||||
expanded_range
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use gpui::{App, AppContext};
|
||||
use indoc::indoc;
|
||||
use language::{Buffer, Language, LanguageConfig, LanguageMatcher, tree_sitter_rust};
|
||||
use std::sync::Arc;
|
||||
|
||||
#[gpui::test]
|
||||
fn test_excerpt_for_cursor_position(cx: &mut App) {
|
||||
let text = indoc! {r#"
|
||||
fn foo() {
|
||||
let x = 42;
|
||||
println!("Hello, world!");
|
||||
}
|
||||
|
||||
fn bar() {
|
||||
let x = 42;
|
||||
let mut sum = 0;
|
||||
for i in 0..x {
|
||||
sum += i;
|
||||
}
|
||||
println!("Sum: {}", sum);
|
||||
return sum;
|
||||
}
|
||||
|
||||
fn generate_random_numbers() -> Vec<i32> {
|
||||
let mut rng = rand::thread_rng();
|
||||
let mut numbers = Vec::new();
|
||||
for _ in 0..5 {
|
||||
numbers.push(rng.random_range(1..101));
|
||||
}
|
||||
numbers
|
||||
}
|
||||
"#};
|
||||
let buffer = cx.new(|cx| Buffer::local(text, cx).with_language(Arc::new(rust_lang()), cx));
|
||||
let snapshot = buffer.read(cx).snapshot();
|
||||
|
||||
// Ensure we try to fit the largest possible syntax scope, resorting to line-based expansion
|
||||
// when a larger scope doesn't fit the editable region.
|
||||
let excerpt = excerpt_for_cursor_position(Point::new(12, 5), "main.rs", &snapshot, 50, 32);
|
||||
assert_eq!(
|
||||
excerpt.prompt,
|
||||
indoc! {r#"
|
||||
```main.rs
|
||||
let x = 42;
|
||||
println!("Hello, world!");
|
||||
<|editable_region_start|>
|
||||
}
|
||||
|
||||
fn bar() {
|
||||
let x = 42;
|
||||
let mut sum = 0;
|
||||
for i in 0..x {
|
||||
sum += i;
|
||||
}
|
||||
println!("Sum: {}", sum);
|
||||
r<|user_cursor_is_here|>eturn sum;
|
||||
}
|
||||
|
||||
fn generate_random_numbers() -> Vec<i32> {
|
||||
<|editable_region_end|>
|
||||
let mut rng = rand::thread_rng();
|
||||
let mut numbers = Vec::new();
|
||||
```"#}
|
||||
);
|
||||
|
||||
// The `bar` function won't fit within the editable region, so we resort to line-based expansion.
|
||||
let excerpt = excerpt_for_cursor_position(Point::new(12, 5), "main.rs", &snapshot, 40, 32);
|
||||
assert_eq!(
|
||||
excerpt.prompt,
|
||||
indoc! {r#"
|
||||
```main.rs
|
||||
fn bar() {
|
||||
let x = 42;
|
||||
let mut sum = 0;
|
||||
<|editable_region_start|>
|
||||
for i in 0..x {
|
||||
sum += i;
|
||||
}
|
||||
println!("Sum: {}", sum);
|
||||
r<|user_cursor_is_here|>eturn sum;
|
||||
}
|
||||
|
||||
fn generate_random_numbers() -> Vec<i32> {
|
||||
let mut rng = rand::thread_rng();
|
||||
<|editable_region_end|>
|
||||
let mut numbers = Vec::new();
|
||||
for _ in 0..5 {
|
||||
numbers.push(rng.random_range(1..101));
|
||||
```"#}
|
||||
);
|
||||
}
|
||||
|
||||
fn rust_lang() -> Language {
|
||||
Language::new(
|
||||
LanguageConfig {
|
||||
name: "Rust".into(),
|
||||
matcher: LanguageMatcher {
|
||||
path_suffixes: vec!["rs".to_string()],
|
||||
..Default::default()
|
||||
},
|
||||
..Default::default()
|
||||
},
|
||||
Some(tree_sitter_rust::LANGUAGE.into()),
|
||||
)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,358 @@
|
||||
#[cfg(feature = "eval-support")]
|
||||
use crate::EvalCacheEntryKind;
|
||||
use crate::prediction::EditPredictionResult;
|
||||
use crate::{
|
||||
DebugEvent, EDIT_PREDICTIONS_MODEL_ID, EditPredictionId, EditPredictionInputs,
|
||||
EditPredictionRequestedDebugEvent, EditPredictionStore,
|
||||
};
|
||||
use anyhow::{Result, anyhow, bail};
|
||||
use cloud_llm_client::predict_edits_v3::{self, Event, PromptFormat};
|
||||
use cloud_llm_client::{EditPredictionRejectReason, PredictEditsRequestTrigger};
|
||||
use cloud_zeta2_prompt::CURSOR_MARKER;
|
||||
use edit_prediction_context::{EditPredictionExcerpt, Line};
|
||||
use edit_prediction_context::{RelatedExcerpt, RelatedFile};
|
||||
use futures::channel::oneshot;
|
||||
use gpui::{Entity, Task, prelude::*};
|
||||
use language::{Anchor, BufferSnapshot};
|
||||
use language::{Buffer, Point, ToOffset as _, ToPoint};
|
||||
use project::{Project, ProjectItem as _};
|
||||
use release_channel::AppVersion;
|
||||
use std::{
|
||||
env,
|
||||
path::Path,
|
||||
sync::Arc,
|
||||
time::{Duration, Instant},
|
||||
};
|
||||
|
||||
pub fn request_prediction_with_zeta2(
|
||||
store: &mut EditPredictionStore,
|
||||
project: &Entity<Project>,
|
||||
active_buffer: &Entity<Buffer>,
|
||||
active_snapshot: BufferSnapshot,
|
||||
position: Anchor,
|
||||
events: Vec<Arc<Event>>,
|
||||
mut included_files: Vec<RelatedFile>,
|
||||
trigger: PredictEditsRequestTrigger,
|
||||
cx: &mut Context<EditPredictionStore>,
|
||||
) -> Task<Result<Option<EditPredictionResult>>> {
|
||||
let options = store.options.clone();
|
||||
let buffer_snapshotted_at = Instant::now();
|
||||
|
||||
let Some((excerpt_path, active_project_path)) = active_snapshot
|
||||
.file()
|
||||
.map(|file| -> Arc<Path> { file.full_path(cx).into() })
|
||||
.zip(active_buffer.read(cx).project_path(cx))
|
||||
else {
|
||||
return Task::ready(Err(anyhow!("No file path for excerpt")));
|
||||
};
|
||||
|
||||
let client = store.client.clone();
|
||||
let llm_token = store.llm_token.clone();
|
||||
let app_version = AppVersion::global(cx);
|
||||
let debug_tx = store.debug_tx.clone();
|
||||
|
||||
let file = active_buffer.read(cx).file();
|
||||
|
||||
let active_file_full_path = file.as_ref().map(|f| f.full_path(cx));
|
||||
|
||||
// TODO data collection
|
||||
let can_collect_data = file
|
||||
.as_ref()
|
||||
.map_or(false, |file| store.can_collect_file(project, file, cx));
|
||||
|
||||
#[cfg(feature = "eval-support")]
|
||||
let eval_cache = store.eval_cache.clone();
|
||||
|
||||
let request_task = cx.background_spawn({
|
||||
let active_buffer = active_buffer.clone();
|
||||
async move {
|
||||
let cursor_offset = position.to_offset(&active_snapshot);
|
||||
let cursor_point = cursor_offset.to_point(&active_snapshot);
|
||||
|
||||
let before_retrieval = Instant::now();
|
||||
|
||||
let excerpt_options = options.context;
|
||||
|
||||
let Some(excerpt) = EditPredictionExcerpt::select_from_buffer(
|
||||
cursor_point,
|
||||
&active_snapshot,
|
||||
&excerpt_options,
|
||||
) else {
|
||||
return Ok((None, None));
|
||||
};
|
||||
|
||||
let excerpt_anchor_range = active_snapshot.anchor_after(excerpt.range.start)
|
||||
..active_snapshot.anchor_before(excerpt.range.end);
|
||||
let related_excerpt = RelatedExcerpt {
|
||||
anchor_range: excerpt_anchor_range.clone(),
|
||||
point_range: Point::new(excerpt.line_range.start.0, 0)
|
||||
..Point::new(excerpt.line_range.end.0, 0),
|
||||
text: active_snapshot.as_rope().slice(excerpt.range),
|
||||
};
|
||||
|
||||
if let Some(buffer_ix) = included_files
|
||||
.iter()
|
||||
.position(|file| file.buffer.entity_id() == active_buffer.entity_id())
|
||||
{
|
||||
let file = &mut included_files[buffer_ix];
|
||||
file.excerpts.push(related_excerpt);
|
||||
file.merge_excerpts();
|
||||
let last_ix = included_files.len() - 1;
|
||||
included_files.swap(buffer_ix, last_ix);
|
||||
} else {
|
||||
let active_file = RelatedFile {
|
||||
path: active_project_path,
|
||||
buffer: active_buffer.downgrade(),
|
||||
excerpts: vec![related_excerpt],
|
||||
max_row: active_snapshot.max_point().row,
|
||||
};
|
||||
included_files.push(active_file);
|
||||
}
|
||||
|
||||
let included_files = included_files
|
||||
.iter()
|
||||
.map(|related_file| predict_edits_v3::RelatedFile {
|
||||
path: Arc::from(related_file.path.path.as_std_path()),
|
||||
max_row: Line(related_file.max_row),
|
||||
excerpts: related_file
|
||||
.excerpts
|
||||
.iter()
|
||||
.map(|excerpt| predict_edits_v3::Excerpt {
|
||||
start_line: Line(excerpt.point_range.start.row),
|
||||
text: excerpt.text.to_string().into(),
|
||||
})
|
||||
.collect(),
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
let cloud_request = predict_edits_v3::PredictEditsRequest {
|
||||
excerpt_path,
|
||||
excerpt: String::new(),
|
||||
excerpt_line_range: Line(0)..Line(0),
|
||||
excerpt_range: 0..0,
|
||||
cursor_point: predict_edits_v3::Point {
|
||||
line: predict_edits_v3::Line(cursor_point.row),
|
||||
column: cursor_point.column,
|
||||
},
|
||||
related_files: included_files,
|
||||
events,
|
||||
can_collect_data,
|
||||
debug_info: debug_tx.is_some(),
|
||||
prompt_max_bytes: Some(options.max_prompt_bytes),
|
||||
prompt_format: options.prompt_format,
|
||||
excerpt_parent: None,
|
||||
git_info: None,
|
||||
trigger,
|
||||
};
|
||||
|
||||
let prompt_result = cloud_zeta2_prompt::build_prompt(&cloud_request);
|
||||
|
||||
let inputs = EditPredictionInputs {
|
||||
included_files: cloud_request.related_files,
|
||||
events: cloud_request.events,
|
||||
cursor_point: cloud_request.cursor_point,
|
||||
cursor_path: cloud_request.excerpt_path,
|
||||
};
|
||||
|
||||
let retrieval_time = Instant::now() - before_retrieval;
|
||||
|
||||
let debug_response_tx = if let Some(debug_tx) = &debug_tx {
|
||||
let (response_tx, response_rx) = oneshot::channel();
|
||||
|
||||
debug_tx
|
||||
.unbounded_send(DebugEvent::EditPredictionRequested(
|
||||
EditPredictionRequestedDebugEvent {
|
||||
inputs: inputs.clone(),
|
||||
retrieval_time,
|
||||
buffer: active_buffer.downgrade(),
|
||||
local_prompt: match prompt_result.as_ref() {
|
||||
Ok(prompt) => Ok(prompt.clone()),
|
||||
Err(err) => Err(err.to_string()),
|
||||
},
|
||||
position,
|
||||
response_rx,
|
||||
},
|
||||
))
|
||||
.ok();
|
||||
Some(response_tx)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
if cfg!(debug_assertions) && env::var("ZED_ZETA2_SKIP_REQUEST").is_ok() {
|
||||
if let Some(debug_response_tx) = debug_response_tx {
|
||||
debug_response_tx
|
||||
.send((Err("Request skipped".to_string()), Duration::ZERO))
|
||||
.ok();
|
||||
}
|
||||
anyhow::bail!("Skipping request because ZED_ZETA2_SKIP_REQUEST is set")
|
||||
}
|
||||
|
||||
let prompt = prompt_result?;
|
||||
let generation_params =
|
||||
cloud_zeta2_prompt::generation_params(cloud_request.prompt_format);
|
||||
let request = open_ai::Request {
|
||||
model: EDIT_PREDICTIONS_MODEL_ID.clone(),
|
||||
messages: vec![open_ai::RequestMessage::User {
|
||||
content: open_ai::MessageContent::Plain(prompt),
|
||||
}],
|
||||
stream: false,
|
||||
max_completion_tokens: None,
|
||||
stop: generation_params.stop.unwrap_or_default(),
|
||||
temperature: generation_params.temperature.unwrap_or(0.7),
|
||||
tool_choice: None,
|
||||
parallel_tool_calls: None,
|
||||
tools: vec![],
|
||||
prompt_cache_key: None,
|
||||
reasoning_effort: None,
|
||||
};
|
||||
|
||||
log::trace!("Sending edit prediction request");
|
||||
|
||||
let before_request = Instant::now();
|
||||
let response = EditPredictionStore::send_raw_llm_request(
|
||||
request,
|
||||
client,
|
||||
llm_token,
|
||||
app_version,
|
||||
#[cfg(feature = "eval-support")]
|
||||
eval_cache,
|
||||
#[cfg(feature = "eval-support")]
|
||||
EvalCacheEntryKind::Prediction,
|
||||
)
|
||||
.await;
|
||||
let received_response_at = Instant::now();
|
||||
let request_time = received_response_at - before_request;
|
||||
|
||||
log::trace!("Got edit prediction response");
|
||||
|
||||
if let Some(debug_response_tx) = debug_response_tx {
|
||||
debug_response_tx
|
||||
.send((
|
||||
response
|
||||
.as_ref()
|
||||
.map_err(|err| err.to_string())
|
||||
.map(|response| response.0.clone()),
|
||||
request_time,
|
||||
))
|
||||
.ok();
|
||||
}
|
||||
|
||||
let (res, usage) = response?;
|
||||
let request_id = EditPredictionId(res.id.clone().into());
|
||||
let Some(mut output_text) = text_from_response(res) else {
|
||||
return Ok((Some((request_id, None)), usage));
|
||||
};
|
||||
|
||||
if output_text.contains(CURSOR_MARKER) {
|
||||
log::trace!("Stripping out {CURSOR_MARKER} from response");
|
||||
output_text = output_text.replace(CURSOR_MARKER, "");
|
||||
}
|
||||
|
||||
let get_buffer_from_context = |path: &Path| {
|
||||
if Some(path) == active_file_full_path.as_deref() {
|
||||
Some((
|
||||
&active_snapshot,
|
||||
std::slice::from_ref(&excerpt_anchor_range),
|
||||
))
|
||||
} else {
|
||||
None
|
||||
}
|
||||
};
|
||||
|
||||
let (_, edits) = match options.prompt_format {
|
||||
PromptFormat::Minimal | PromptFormat::MinimalQwen | PromptFormat::SeedCoder1120 => {
|
||||
if output_text.contains("--- a/\n+++ b/\nNo edits") {
|
||||
let edits = vec![];
|
||||
(&active_snapshot, edits)
|
||||
} else {
|
||||
crate::udiff::parse_diff(&output_text, get_buffer_from_context).await?
|
||||
}
|
||||
}
|
||||
PromptFormat::OldTextNewText => {
|
||||
crate::xml_edits::parse_xml_edits(&output_text, get_buffer_from_context).await?
|
||||
}
|
||||
_ => {
|
||||
bail!("unsupported prompt format {}", options.prompt_format)
|
||||
}
|
||||
};
|
||||
|
||||
anyhow::Ok((
|
||||
Some((
|
||||
request_id,
|
||||
Some((
|
||||
inputs,
|
||||
active_buffer,
|
||||
active_snapshot.clone(),
|
||||
edits,
|
||||
received_response_at,
|
||||
)),
|
||||
)),
|
||||
usage,
|
||||
))
|
||||
}
|
||||
});
|
||||
|
||||
cx.spawn(async move |this, cx| {
|
||||
let Some((id, prediction)) =
|
||||
EditPredictionStore::handle_api_response(&this, request_task.await, cx)?
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
let Some((inputs, edited_buffer, edited_buffer_snapshot, edits, received_response_at)) =
|
||||
prediction
|
||||
else {
|
||||
return Ok(Some(EditPredictionResult {
|
||||
id,
|
||||
prediction: Err(EditPredictionRejectReason::Empty),
|
||||
}));
|
||||
};
|
||||
|
||||
Ok(Some(
|
||||
EditPredictionResult::new(
|
||||
id,
|
||||
&edited_buffer,
|
||||
&edited_buffer_snapshot,
|
||||
edits.into(),
|
||||
buffer_snapshotted_at,
|
||||
received_response_at,
|
||||
inputs,
|
||||
cx,
|
||||
)
|
||||
.await,
|
||||
))
|
||||
})
|
||||
}
|
||||
|
||||
pub fn text_from_response(mut res: open_ai::Response) -> Option<String> {
|
||||
let choice = res.choices.pop()?;
|
||||
let output_text = match choice.message {
|
||||
open_ai::RequestMessage::Assistant {
|
||||
content: Some(open_ai::MessageContent::Plain(content)),
|
||||
..
|
||||
} => content,
|
||||
open_ai::RequestMessage::Assistant {
|
||||
content: Some(open_ai::MessageContent::Multipart(mut content)),
|
||||
..
|
||||
} => {
|
||||
if content.is_empty() {
|
||||
log::error!("No output from Baseten completion response");
|
||||
return None;
|
||||
}
|
||||
|
||||
match content.remove(0) {
|
||||
open_ai::MessagePart::Text { text } => text,
|
||||
open_ai::MessagePart::Image { .. } => {
|
||||
log::error!("Expected text, got an image");
|
||||
return None;
|
||||
}
|
||||
}
|
||||
}
|
||||
_ => {
|
||||
log::error!("Invalid response message: {:?}", choice.message);
|
||||
return None;
|
||||
}
|
||||
};
|
||||
Some(output_text)
|
||||
}
|
||||
Reference in New Issue
Block a user