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
@@ -0,0 +1,641 @@
|
||||
use crate::metrics::{self, Scores};
|
||||
use std::{
|
||||
collections::HashMap,
|
||||
io::{IsTerminal, Write},
|
||||
sync::Arc,
|
||||
};
|
||||
|
||||
use anyhow::Result;
|
||||
use edit_prediction::{EditPredictionStore, udiff::DiffLine};
|
||||
use gpui::{AsyncApp, Entity};
|
||||
use project::Project;
|
||||
use util::ResultExt as _;
|
||||
|
||||
use crate::{
|
||||
EvaluateArguments, PredictionOptions,
|
||||
example::{Example, NamedExample},
|
||||
headless::ZetaCliAppState,
|
||||
paths::print_run_data_dir,
|
||||
predict::{PredictionDetails, perform_predict, setup_store},
|
||||
};
|
||||
|
||||
#[derive(Debug)]
|
||||
pub(crate) struct ExecutionData {
|
||||
execution_id: String,
|
||||
diff: String,
|
||||
reasoning: String,
|
||||
}
|
||||
|
||||
pub async fn run_evaluate(
|
||||
args: EvaluateArguments,
|
||||
app_state: &Arc<ZetaCliAppState>,
|
||||
cx: &mut AsyncApp,
|
||||
) {
|
||||
if args.example_paths.is_empty() {
|
||||
eprintln!("No examples provided");
|
||||
return;
|
||||
}
|
||||
|
||||
let all_tasks = args.example_paths.into_iter().map(|path| {
|
||||
let options = args.options.clone();
|
||||
let app_state = app_state.clone();
|
||||
let example = NamedExample::load(&path).expect("Failed to load example");
|
||||
|
||||
cx.spawn(async move |cx| {
|
||||
let project = example.setup_project(&app_state, cx).await.unwrap();
|
||||
|
||||
let providers = (0..args.repetitions)
|
||||
.map(|_| setup_store(args.options.provider, &project, &app_state, cx).unwrap())
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
let _edited_buffers = example.apply_edit_history(&project, cx).await.unwrap();
|
||||
|
||||
let tasks = providers
|
||||
.into_iter()
|
||||
.enumerate()
|
||||
.map(move |(repetition_ix, store)| {
|
||||
let repetition_ix = (args.repetitions > 1).then(|| repetition_ix as u16);
|
||||
let example = example.clone();
|
||||
let project = project.clone();
|
||||
let options = options.clone();
|
||||
|
||||
cx.spawn(async move |cx| {
|
||||
let name = example.name.clone();
|
||||
run_evaluate_one(
|
||||
example,
|
||||
repetition_ix,
|
||||
project,
|
||||
store,
|
||||
options,
|
||||
!args.skip_prediction,
|
||||
cx,
|
||||
)
|
||||
.await
|
||||
.map_err(|err| (err, name, repetition_ix))
|
||||
})
|
||||
});
|
||||
futures::future::join_all(tasks).await
|
||||
})
|
||||
});
|
||||
let all_results = futures::future::join_all(all_tasks).await;
|
||||
|
||||
write_aggregated_scores(&mut std::io::stdout(), &all_results).unwrap();
|
||||
if let Some(mut output_file) =
|
||||
std::fs::File::create(crate::paths::RUN_DIR.join("aggregated_results.md")).log_err()
|
||||
{
|
||||
write_aggregated_scores(&mut output_file, &all_results).log_err();
|
||||
};
|
||||
|
||||
if args.repetitions > 1 {
|
||||
if let Err(e) = write_bucketed_analysis(&all_results) {
|
||||
eprintln!("Failed to write bucketed analysis: {:?}", e);
|
||||
}
|
||||
}
|
||||
|
||||
print_run_data_dir(args.repetitions == 1, std::io::stdout().is_terminal());
|
||||
}
|
||||
|
||||
fn write_aggregated_scores(
|
||||
w: &mut impl std::io::Write,
|
||||
all_results: &Vec<
|
||||
Vec<Result<(EvaluationResult, ExecutionData), (anyhow::Error, String, Option<u16>)>>,
|
||||
>,
|
||||
) -> Result<()> {
|
||||
let mut successful = Vec::new();
|
||||
let mut failed_count = 0;
|
||||
|
||||
for result in all_results.iter().flatten() {
|
||||
match result {
|
||||
Ok((eval_result, _execution_data)) => successful.push(eval_result),
|
||||
Err((err, name, repetition_ix)) => {
|
||||
if failed_count == 0 {
|
||||
writeln!(w, "## Errors\n")?;
|
||||
}
|
||||
|
||||
failed_count += 1;
|
||||
writeln!(w, "{}", fmt_evaluation_error(err, name, repetition_ix))?;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if successful.len() > 1 {
|
||||
let edit_scores = successful
|
||||
.iter()
|
||||
.filter_map(|r| r.edit_scores.clone())
|
||||
.collect::<Vec<_>>();
|
||||
let has_edit_predictions = edit_scores.len() > 0;
|
||||
let aggregated_result = EvaluationResult {
|
||||
context_scores: Scores::aggregate(successful.iter().map(|r| &r.context_scores)),
|
||||
edit_scores: has_edit_predictions.then(|| EditScores::aggregate(&edit_scores)),
|
||||
prompt_len: successful.iter().map(|r| r.prompt_len).sum::<usize>() / successful.len(),
|
||||
generated_len: successful.iter().map(|r| r.generated_len).sum::<usize>()
|
||||
/ successful.len(),
|
||||
};
|
||||
|
||||
writeln!(w, "\n{}", "-".repeat(80))?;
|
||||
writeln!(w, "\n## TOTAL SCORES")?;
|
||||
writeln!(w, "{:#}", aggregated_result)?;
|
||||
}
|
||||
|
||||
if successful.len() + failed_count > 1 {
|
||||
writeln!(
|
||||
w,
|
||||
"\nCongratulations! {}/{} ({:.2}%) of runs weren't outright failures 🎉",
|
||||
successful.len(),
|
||||
successful.len() + failed_count,
|
||||
(successful.len() as f64 / (successful.len() + failed_count) as f64) * 100.0
|
||||
)?;
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn run_evaluate_one(
|
||||
example: NamedExample,
|
||||
repetition_ix: Option<u16>,
|
||||
project: Entity<Project>,
|
||||
store: Entity<EditPredictionStore>,
|
||||
prediction_options: PredictionOptions,
|
||||
predict: bool,
|
||||
cx: &mut AsyncApp,
|
||||
) -> Result<(EvaluationResult, ExecutionData)> {
|
||||
let predict_result = perform_predict(
|
||||
example.clone(),
|
||||
project,
|
||||
store,
|
||||
repetition_ix,
|
||||
prediction_options,
|
||||
cx,
|
||||
)
|
||||
.await?;
|
||||
|
||||
let evaluation_result = evaluate(&example.example, &predict_result, predict);
|
||||
|
||||
if repetition_ix.is_none() {
|
||||
write_eval_result(
|
||||
&example,
|
||||
&predict_result,
|
||||
&evaluation_result,
|
||||
&mut std::io::stdout(),
|
||||
std::io::stdout().is_terminal(),
|
||||
predict,
|
||||
)?;
|
||||
}
|
||||
|
||||
if let Some(mut results_file) =
|
||||
std::fs::File::create(predict_result.run_example_dir.join("results.md")).log_err()
|
||||
{
|
||||
write_eval_result(
|
||||
&example,
|
||||
&predict_result,
|
||||
&evaluation_result,
|
||||
&mut results_file,
|
||||
false,
|
||||
predict,
|
||||
)
|
||||
.log_err();
|
||||
}
|
||||
|
||||
let execution_data = ExecutionData {
|
||||
execution_id: if let Some(rep_ix) = repetition_ix {
|
||||
format!("{:03}", rep_ix)
|
||||
} else {
|
||||
example.name.clone()
|
||||
},
|
||||
diff: predict_result.diff.clone(),
|
||||
reasoning: std::fs::read_to_string(
|
||||
predict_result
|
||||
.run_example_dir
|
||||
.join("prediction_response.md"),
|
||||
)
|
||||
.unwrap_or_default(),
|
||||
};
|
||||
|
||||
anyhow::Ok((evaluation_result, execution_data))
|
||||
}
|
||||
|
||||
fn write_eval_result(
|
||||
example: &NamedExample,
|
||||
predictions: &PredictionDetails,
|
||||
evaluation_result: &EvaluationResult,
|
||||
out: &mut impl Write,
|
||||
use_color: bool,
|
||||
predict: bool,
|
||||
) -> Result<()> {
|
||||
if predict {
|
||||
writeln!(
|
||||
out,
|
||||
"## Expected edit prediction:\n\n```diff\n{}\n```\n",
|
||||
compare_diffs(
|
||||
&example.example.expected_patch,
|
||||
&predictions.diff,
|
||||
use_color
|
||||
)
|
||||
)?;
|
||||
writeln!(
|
||||
out,
|
||||
"## Actual edit prediction:\n\n```diff\n{}\n```\n",
|
||||
compare_diffs(
|
||||
&predictions.diff,
|
||||
&example.example.expected_patch,
|
||||
use_color
|
||||
)
|
||||
)?;
|
||||
}
|
||||
|
||||
writeln!(out, "{:#}", evaluation_result)?;
|
||||
|
||||
anyhow::Ok(())
|
||||
}
|
||||
|
||||
#[derive(Debug, Default, Clone)]
|
||||
pub struct EditScores {
|
||||
pub line_match: Scores,
|
||||
pub chr_f: f64,
|
||||
}
|
||||
|
||||
impl EditScores {
|
||||
pub fn aggregate(scores: &[EditScores]) -> EditScores {
|
||||
let line_match = Scores::aggregate(scores.iter().map(|s| &s.line_match));
|
||||
let chr_f = scores.iter().map(|s| s.chr_f).sum::<f64>() / scores.len() as f64;
|
||||
|
||||
EditScores { line_match, chr_f }
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Default)]
|
||||
pub struct EvaluationResult {
|
||||
pub edit_scores: Option<EditScores>,
|
||||
pub context_scores: Scores,
|
||||
pub prompt_len: usize,
|
||||
pub generated_len: usize,
|
||||
}
|
||||
|
||||
impl std::fmt::Display for EvaluationResult {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
if f.alternate() {
|
||||
self.fmt_table(f)
|
||||
} else {
|
||||
self.fmt_markdown(f)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl EvaluationResult {
|
||||
fn fmt_markdown(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
write!(
|
||||
f,
|
||||
r#"
|
||||
### Context Scores
|
||||
{}
|
||||
"#,
|
||||
self.context_scores.to_markdown(),
|
||||
)?;
|
||||
if let Some(scores) = &self.edit_scores {
|
||||
write!(
|
||||
f,
|
||||
r#"
|
||||
### Edit Prediction Scores
|
||||
{}"#,
|
||||
scores.line_match.to_markdown()
|
||||
)?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn fmt_table(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
writeln!(f, "#### Prompt Statistics")?;
|
||||
writeln!(f, "─────────────────────────")?;
|
||||
writeln!(f, "Prompt_len Generated_len")?;
|
||||
writeln!(f, "─────────────────────────")?;
|
||||
writeln!(f, "{:<11} {:<14}", self.prompt_len, self.generated_len,)?;
|
||||
writeln!(f)?;
|
||||
writeln!(f)?;
|
||||
writeln!(f, "#### Performance Scores")?;
|
||||
writeln!(
|
||||
f,
|
||||
"──────────────────────────────────────────────────────────────────"
|
||||
)?;
|
||||
writeln!(
|
||||
f,
|
||||
" TP FP FN Precision Recall F1"
|
||||
)?;
|
||||
writeln!(
|
||||
f,
|
||||
"──────────────────────────────────────────────────────────────────"
|
||||
)?;
|
||||
writeln!(
|
||||
f,
|
||||
"Context Retrieval {:<6} {:<6} {:<6} {:>8.2} {:>7.2} {:>6.2}",
|
||||
self.context_scores.true_positives,
|
||||
self.context_scores.false_positives,
|
||||
self.context_scores.false_negatives,
|
||||
self.context_scores.precision() * 100.0,
|
||||
self.context_scores.recall() * 100.0,
|
||||
self.context_scores.f1_score() * 100.0
|
||||
)?;
|
||||
if let Some(edit_scores) = &self.edit_scores {
|
||||
let line_match = &edit_scores.line_match;
|
||||
writeln!(f, "Edit Prediction")?;
|
||||
writeln!(
|
||||
f,
|
||||
" ├─ exact lines {:<6} {:<6} {:<6} {:>8.2} {:>7.2} {:>6.2}",
|
||||
line_match.true_positives,
|
||||
line_match.false_positives,
|
||||
line_match.false_negatives,
|
||||
line_match.precision() * 100.0,
|
||||
line_match.recall() * 100.0,
|
||||
line_match.f1_score() * 100.0
|
||||
)?;
|
||||
writeln!(
|
||||
f,
|
||||
" └─ diff chrF {:<6} {:<6} {:<6} {:>8} {:>8} {:>6.2}",
|
||||
"-", "-", "-", "-", "-", edit_scores.chr_f
|
||||
)?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
fn evaluate(example: &Example, preds: &PredictionDetails, predict: bool) -> EvaluationResult {
|
||||
let mut eval_result = EvaluationResult {
|
||||
prompt_len: preds.prompt_len,
|
||||
generated_len: preds.generated_len,
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
if predict {
|
||||
// todo: alternatives for patches
|
||||
let expected_patch = example
|
||||
.expected_patch
|
||||
.lines()
|
||||
.map(DiffLine::parse)
|
||||
.collect::<Vec<_>>();
|
||||
let actual_patch = preds.diff.lines().map(DiffLine::parse).collect::<Vec<_>>();
|
||||
|
||||
let line_match = metrics::line_match_score(&expected_patch, &actual_patch);
|
||||
let chr_f = metrics::delta_chr_f(&expected_patch, &actual_patch);
|
||||
|
||||
eval_result.edit_scores = Some(EditScores { line_match, chr_f });
|
||||
}
|
||||
|
||||
eval_result
|
||||
}
|
||||
|
||||
/// Return annotated `patch_a` so that:
|
||||
/// Additions and deletions that are not present in `patch_b` will be highlighted in red.
|
||||
/// Additions and deletions that are present in `patch_b` will be highlighted in green.
|
||||
pub fn compare_diffs(patch_a: &str, patch_b: &str, use_color: bool) -> String {
|
||||
let green = if use_color { "\x1b[32m✓ " } else { "" };
|
||||
let red = if use_color { "\x1b[31m✗ " } else { "" };
|
||||
let neutral = if use_color { " " } else { "" };
|
||||
let reset = if use_color { "\x1b[0m" } else { "" };
|
||||
let lines_a = patch_a.lines().map(DiffLine::parse);
|
||||
let lines_b: Vec<_> = patch_b.lines().map(DiffLine::parse).collect();
|
||||
|
||||
let annotated = lines_a
|
||||
.map(|line| match line {
|
||||
DiffLine::Addition(_) | DiffLine::Deletion(_) => {
|
||||
if lines_b.contains(&line) {
|
||||
format!("{green}{line}{reset}")
|
||||
} else {
|
||||
format!("{red}{line}{reset}")
|
||||
}
|
||||
}
|
||||
_ => format!("{neutral}{line}{reset}"),
|
||||
})
|
||||
.collect::<Vec<String>>();
|
||||
|
||||
annotated.join("\n")
|
||||
}
|
||||
|
||||
fn write_bucketed_analysis(
|
||||
all_results: &Vec<
|
||||
Vec<Result<(EvaluationResult, ExecutionData), (anyhow::Error, String, Option<u16>)>>,
|
||||
>,
|
||||
) -> Result<()> {
|
||||
#[derive(Debug)]
|
||||
struct EditBucket {
|
||||
diff: String,
|
||||
is_correct: bool,
|
||||
execution_indices: Vec<String>,
|
||||
reasoning_samples: Vec<String>,
|
||||
}
|
||||
|
||||
let mut total_executions = 0;
|
||||
let mut empty_predictions = Vec::new();
|
||||
let mut errors = Vec::new();
|
||||
|
||||
let mut buckets: HashMap<String, EditBucket> = HashMap::new();
|
||||
|
||||
for result in all_results.iter().flatten() {
|
||||
total_executions += 1;
|
||||
|
||||
let (evaluation_result, execution_data) = match result {
|
||||
Ok((eval_result, execution_data)) => {
|
||||
if execution_data.diff.is_empty() {
|
||||
empty_predictions.push(execution_data);
|
||||
continue;
|
||||
}
|
||||
(eval_result, execution_data)
|
||||
}
|
||||
Err(err) => {
|
||||
errors.push(err);
|
||||
continue;
|
||||
}
|
||||
};
|
||||
|
||||
buckets
|
||||
.entry(execution_data.diff.clone())
|
||||
.and_modify(|bucket| {
|
||||
bucket
|
||||
.execution_indices
|
||||
.push(execution_data.execution_id.clone());
|
||||
bucket
|
||||
.reasoning_samples
|
||||
.push(execution_data.reasoning.clone());
|
||||
})
|
||||
.or_insert_with(|| EditBucket {
|
||||
diff: execution_data.diff.clone(),
|
||||
is_correct: {
|
||||
evaluation_result
|
||||
.edit_scores
|
||||
.as_ref()
|
||||
.map_or(false, |edit_scores| {
|
||||
edit_scores.line_match.false_positives == 0
|
||||
&& edit_scores.line_match.false_negatives == 0
|
||||
&& edit_scores.line_match.true_positives > 0
|
||||
})
|
||||
},
|
||||
execution_indices: vec![execution_data.execution_id.clone()],
|
||||
reasoning_samples: vec![execution_data.reasoning.clone()],
|
||||
});
|
||||
}
|
||||
|
||||
let mut sorted_buckets = buckets.into_values().collect::<Vec<_>>();
|
||||
sorted_buckets.sort_by(|a, b| match (a.is_correct, b.is_correct) {
|
||||
(true, false) => std::cmp::Ordering::Less,
|
||||
(false, true) => std::cmp::Ordering::Greater,
|
||||
_ => b.execution_indices.len().cmp(&a.execution_indices.len()),
|
||||
});
|
||||
|
||||
let output_path = crate::paths::RUN_DIR.join("bucketed_analysis.md");
|
||||
let mut output = std::fs::File::create(&output_path)?;
|
||||
|
||||
writeln!(output, "# Bucketed Edit Analysis\n")?;
|
||||
|
||||
writeln!(output, "## Summary\n")?;
|
||||
writeln!(output, "- **Total executions**: {}", total_executions)?;
|
||||
|
||||
let correct_count: usize = sorted_buckets
|
||||
.iter()
|
||||
.filter(|b| b.is_correct)
|
||||
.map(|b| b.execution_indices.len())
|
||||
.sum();
|
||||
|
||||
let incorrect_count: usize = sorted_buckets
|
||||
.iter()
|
||||
.filter(|b| !b.is_correct)
|
||||
.map(|b| b.execution_indices.len())
|
||||
.sum();
|
||||
|
||||
writeln!(
|
||||
output,
|
||||
"- **Correct predictions**: {} ({:.1}%)",
|
||||
correct_count,
|
||||
(correct_count as f64 / total_executions as f64) * 100.0
|
||||
)?;
|
||||
|
||||
writeln!(
|
||||
output,
|
||||
"- **Incorrect predictions**: {} ({:.1}%)",
|
||||
incorrect_count,
|
||||
(incorrect_count as f64 / total_executions as f64) * 100.0
|
||||
)?;
|
||||
|
||||
writeln!(
|
||||
output,
|
||||
"- **No Predictions**: {} ({:.1}%)",
|
||||
empty_predictions.len(),
|
||||
(empty_predictions.len() as f64 / total_executions as f64) * 100.0
|
||||
)?;
|
||||
|
||||
let unique_incorrect = sorted_buckets.iter().filter(|b| !b.is_correct).count();
|
||||
writeln!(
|
||||
output,
|
||||
"- **Unique incorrect edit patterns**: {}\n",
|
||||
unique_incorrect
|
||||
)?;
|
||||
|
||||
writeln!(output, "---\n")?;
|
||||
|
||||
for (idx, bucket) in sorted_buckets.iter().filter(|b| b.is_correct).enumerate() {
|
||||
if idx == 0 {
|
||||
writeln!(
|
||||
output,
|
||||
"## Correct Predictions ({} occurrences)\n",
|
||||
bucket.execution_indices.len()
|
||||
)?;
|
||||
}
|
||||
|
||||
writeln!(output, "**Predicted Edit:**\n")?;
|
||||
writeln!(output, "```diff")?;
|
||||
writeln!(output, "{}", bucket.diff)?;
|
||||
writeln!(output, "```\n")?;
|
||||
|
||||
writeln!(
|
||||
output,
|
||||
"**Executions:** {}\n",
|
||||
bucket.execution_indices.join(", ")
|
||||
)?;
|
||||
writeln!(output, "---\n")?;
|
||||
}
|
||||
|
||||
for (idx, bucket) in sorted_buckets.iter().filter(|b| !b.is_correct).enumerate() {
|
||||
writeln!(
|
||||
output,
|
||||
"## Incorrect Prediction #{} ({} occurrences)\n",
|
||||
idx + 1,
|
||||
bucket.execution_indices.len()
|
||||
)?;
|
||||
|
||||
writeln!(output, "**Predicted Edit:**\n")?;
|
||||
writeln!(output, "```diff")?;
|
||||
writeln!(output, "{}", bucket.diff)?;
|
||||
writeln!(output, "```\n")?;
|
||||
|
||||
writeln!(
|
||||
output,
|
||||
"**Executions:** {}\n",
|
||||
bucket.execution_indices.join(", ")
|
||||
)?;
|
||||
|
||||
for (exec_id, reasoning) in bucket
|
||||
.execution_indices
|
||||
.iter()
|
||||
.zip(bucket.reasoning_samples.iter())
|
||||
{
|
||||
writeln!(output, "{}", fmt_execution(exec_id, reasoning))?;
|
||||
}
|
||||
|
||||
writeln!(output, "\n---\n")?;
|
||||
}
|
||||
|
||||
if !empty_predictions.is_empty() {
|
||||
writeln!(
|
||||
output,
|
||||
"## No Predictions ({} occurrences)\n",
|
||||
empty_predictions.len()
|
||||
)?;
|
||||
|
||||
for execution_data in &empty_predictions {
|
||||
writeln!(
|
||||
output,
|
||||
"{}",
|
||||
fmt_execution(&execution_data.execution_id, &execution_data.reasoning)
|
||||
)?;
|
||||
}
|
||||
writeln!(output, "\n---\n")?;
|
||||
}
|
||||
|
||||
if !errors.is_empty() {
|
||||
writeln!(output, "## Errors ({} occurrences)\n", errors.len())?;
|
||||
|
||||
for (err, name, repetition_ix) in &errors {
|
||||
writeln!(output, "{}", fmt_evaluation_error(err, name, repetition_ix))?;
|
||||
}
|
||||
writeln!(output, "\n---\n")?;
|
||||
}
|
||||
|
||||
fn fmt_execution(exec_id: &str, reasoning: &str) -> String {
|
||||
let exec_content = format!(
|
||||
"\n### Execution {} `{}/{}/prediction_response.md`{}",
|
||||
exec_id,
|
||||
crate::paths::RUN_DIR.display(),
|
||||
exec_id,
|
||||
indent_text(&format!("\n\n```\n{}\n```\n", reasoning,), 2)
|
||||
);
|
||||
indent_text(&exec_content, 2)
|
||||
}
|
||||
|
||||
fn indent_text(text: &str, spaces: usize) -> String {
|
||||
let indent = " ".repeat(spaces);
|
||||
text.lines()
|
||||
.collect::<Vec<_>>()
|
||||
.join(&format!("\n{}", indent))
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn fmt_evaluation_error(err: &anyhow::Error, name: &str, repetition_ix: &Option<u16>) -> String {
|
||||
let err = format!("{err:?}")
|
||||
.replace("<edits", "```xml\n<edits")
|
||||
.replace("</edits>", "</edits>\n```");
|
||||
format!(
|
||||
"### ERROR {name}{}\n\n{err}\n",
|
||||
repetition_ix
|
||||
.map(|ix| format!(" [RUN {ix:03}]"))
|
||||
.unwrap_or_default()
|
||||
)
|
||||
}
|
||||
@@ -0,0 +1,561 @@
|
||||
use std::{
|
||||
borrow::Cow,
|
||||
cell::RefCell,
|
||||
fmt::{self, Display},
|
||||
fs,
|
||||
io::Write,
|
||||
mem,
|
||||
path::{Path, PathBuf},
|
||||
sync::{Arc, OnceLock},
|
||||
};
|
||||
|
||||
use crate::headless::ZetaCliAppState;
|
||||
use anyhow::{Context as _, Result, anyhow};
|
||||
use clap::ValueEnum;
|
||||
use cloud_zeta2_prompt::CURSOR_MARKER;
|
||||
use collections::HashMap;
|
||||
use edit_prediction::udiff::OpenedBuffers;
|
||||
use futures::{
|
||||
AsyncWriteExt as _,
|
||||
lock::{Mutex, OwnedMutexGuard},
|
||||
};
|
||||
use futures::{FutureExt as _, future::Shared};
|
||||
use gpui::{AsyncApp, Entity, Task, http_client::Url};
|
||||
use language::{Anchor, Buffer};
|
||||
use project::{Project, ProjectPath};
|
||||
use pulldown_cmark::CowStr;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use util::{paths::PathStyle, rel_path::RelPath};
|
||||
|
||||
use crate::paths::{REPOS_DIR, WORKTREES_DIR};
|
||||
|
||||
const UNCOMMITTED_DIFF_HEADING: &str = "Uncommitted Diff";
|
||||
const EDIT_HISTORY_HEADING: &str = "Edit History";
|
||||
const CURSOR_POSITION_HEADING: &str = "Cursor Position";
|
||||
const EXPECTED_PATCH_HEADING: &str = "Expected Patch";
|
||||
const EXPECTED_CONTEXT_HEADING: &str = "Expected Context";
|
||||
const REPOSITORY_URL_FIELD: &str = "repository_url";
|
||||
const REVISION_FIELD: &str = "revision";
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct NamedExample {
|
||||
pub name: String,
|
||||
pub example: Example,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Serialize, Deserialize)]
|
||||
pub struct Example {
|
||||
pub repository_url: String,
|
||||
pub revision: String,
|
||||
pub uncommitted_diff: String,
|
||||
pub cursor_path: PathBuf,
|
||||
pub cursor_position: String,
|
||||
pub edit_history: String,
|
||||
pub expected_patch: String,
|
||||
}
|
||||
|
||||
pub type ActualExcerpt = Excerpt;
|
||||
|
||||
#[derive(Clone, Debug, Serialize, Deserialize)]
|
||||
pub struct Excerpt {
|
||||
pub path: PathBuf,
|
||||
pub text: String,
|
||||
}
|
||||
|
||||
#[derive(ValueEnum, Debug, Clone)]
|
||||
pub enum ExampleFormat {
|
||||
Json,
|
||||
Toml,
|
||||
Md,
|
||||
}
|
||||
|
||||
impl NamedExample {
|
||||
pub fn load(path: impl AsRef<Path>) -> Result<Self> {
|
||||
let path = path.as_ref();
|
||||
let content = std::fs::read_to_string(path)?;
|
||||
let ext = path.extension();
|
||||
|
||||
match ext.and_then(|s| s.to_str()) {
|
||||
Some("json") => Ok(Self {
|
||||
name: path.file_stem().unwrap_or_default().display().to_string(),
|
||||
example: serde_json::from_str(&content)?,
|
||||
}),
|
||||
Some("toml") => Ok(Self {
|
||||
name: path.file_stem().unwrap_or_default().display().to_string(),
|
||||
example: toml::from_str(&content)?,
|
||||
}),
|
||||
Some("md") => Self::parse_md(&content),
|
||||
Some(_) => {
|
||||
anyhow::bail!("Unrecognized example extension: {}", ext.unwrap().display());
|
||||
}
|
||||
None => {
|
||||
anyhow::bail!(
|
||||
"Failed to determine example type since the file does not have an extension."
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn parse_md(input: &str) -> Result<Self> {
|
||||
use pulldown_cmark::{CodeBlockKind, Event, HeadingLevel, Parser, Tag, TagEnd};
|
||||
|
||||
let parser = Parser::new(input);
|
||||
|
||||
let mut named = NamedExample {
|
||||
name: String::new(),
|
||||
example: Example {
|
||||
repository_url: String::new(),
|
||||
revision: String::new(),
|
||||
uncommitted_diff: String::new(),
|
||||
cursor_path: PathBuf::new(),
|
||||
cursor_position: String::new(),
|
||||
edit_history: String::new(),
|
||||
expected_patch: String::new(),
|
||||
},
|
||||
};
|
||||
|
||||
let mut text = String::new();
|
||||
let mut block_info: CowStr = "".into();
|
||||
|
||||
#[derive(PartialEq)]
|
||||
enum Section {
|
||||
UncommittedDiff,
|
||||
EditHistory,
|
||||
CursorPosition,
|
||||
ExpectedExcerpts,
|
||||
ExpectedPatch,
|
||||
Other,
|
||||
}
|
||||
|
||||
let mut current_section = Section::Other;
|
||||
|
||||
for event in parser {
|
||||
match event {
|
||||
Event::Text(line) => {
|
||||
text.push_str(&line);
|
||||
|
||||
if !named.name.is_empty()
|
||||
&& current_section == Section::Other
|
||||
// in h1 section
|
||||
&& let Some((field, value)) = line.split_once('=')
|
||||
{
|
||||
match field.trim() {
|
||||
REPOSITORY_URL_FIELD => {
|
||||
named.example.repository_url = value.trim().to_string();
|
||||
}
|
||||
REVISION_FIELD => {
|
||||
named.example.revision = value.trim().to_string();
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
Event::End(TagEnd::Heading(HeadingLevel::H1)) => {
|
||||
if !named.name.is_empty() {
|
||||
anyhow::bail!(
|
||||
"Found multiple H1 headings. There should only be one with the name of the example."
|
||||
);
|
||||
}
|
||||
named.name = mem::take(&mut text);
|
||||
}
|
||||
Event::End(TagEnd::Heading(HeadingLevel::H2)) => {
|
||||
let title = mem::take(&mut text);
|
||||
current_section = if title.eq_ignore_ascii_case(UNCOMMITTED_DIFF_HEADING) {
|
||||
Section::UncommittedDiff
|
||||
} else if title.eq_ignore_ascii_case(EDIT_HISTORY_HEADING) {
|
||||
Section::EditHistory
|
||||
} else if title.eq_ignore_ascii_case(CURSOR_POSITION_HEADING) {
|
||||
Section::CursorPosition
|
||||
} else if title.eq_ignore_ascii_case(EXPECTED_PATCH_HEADING) {
|
||||
Section::ExpectedPatch
|
||||
} else if title.eq_ignore_ascii_case(EXPECTED_CONTEXT_HEADING) {
|
||||
Section::ExpectedExcerpts
|
||||
} else {
|
||||
Section::Other
|
||||
};
|
||||
}
|
||||
Event::End(TagEnd::Heading(HeadingLevel::H3)) => {
|
||||
mem::take(&mut text);
|
||||
}
|
||||
Event::End(TagEnd::Heading(HeadingLevel::H4)) => {
|
||||
mem::take(&mut text);
|
||||
}
|
||||
Event::End(TagEnd::Heading(level)) => {
|
||||
anyhow::bail!("Unexpected heading level: {level}");
|
||||
}
|
||||
Event::Start(Tag::CodeBlock(kind)) => {
|
||||
match kind {
|
||||
CodeBlockKind::Fenced(info) => {
|
||||
block_info = info;
|
||||
}
|
||||
CodeBlockKind::Indented => {
|
||||
anyhow::bail!("Unexpected indented codeblock");
|
||||
}
|
||||
};
|
||||
}
|
||||
Event::Start(_) => {
|
||||
text.clear();
|
||||
block_info = "".into();
|
||||
}
|
||||
Event::End(TagEnd::CodeBlock) => {
|
||||
let block_info = block_info.trim();
|
||||
match current_section {
|
||||
Section::UncommittedDiff => {
|
||||
named.example.uncommitted_diff = mem::take(&mut text);
|
||||
}
|
||||
Section::EditHistory => {
|
||||
named.example.edit_history.push_str(&mem::take(&mut text));
|
||||
}
|
||||
Section::CursorPosition => {
|
||||
named.example.cursor_path = block_info.into();
|
||||
named.example.cursor_position = mem::take(&mut text);
|
||||
}
|
||||
Section::ExpectedExcerpts => {
|
||||
mem::take(&mut text);
|
||||
}
|
||||
Section::ExpectedPatch => {
|
||||
named.example.expected_patch = mem::take(&mut text);
|
||||
}
|
||||
Section::Other => {}
|
||||
}
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
|
||||
if named.example.cursor_path.as_path() == Path::new("")
|
||||
|| named.example.cursor_position.is_empty()
|
||||
{
|
||||
anyhow::bail!("Missing cursor position codeblock");
|
||||
}
|
||||
|
||||
Ok(named)
|
||||
}
|
||||
|
||||
pub fn write(&self, format: ExampleFormat, mut out: impl Write) -> Result<()> {
|
||||
match format {
|
||||
ExampleFormat::Json => Ok(serde_json::to_writer(out, &self.example)?),
|
||||
ExampleFormat::Toml => {
|
||||
Ok(out.write_all(toml::to_string_pretty(&self.example)?.as_bytes())?)
|
||||
}
|
||||
ExampleFormat::Md => Ok(write!(out, "{}", self)?),
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn setup_project(
|
||||
&self,
|
||||
app_state: &Arc<ZetaCliAppState>,
|
||||
cx: &mut AsyncApp,
|
||||
) -> Result<Entity<Project>> {
|
||||
let worktree_path = self.setup_worktree().await?;
|
||||
|
||||
static AUTHENTICATED: OnceLock<Shared<Task<()>>> = OnceLock::new();
|
||||
|
||||
AUTHENTICATED
|
||||
.get_or_init(|| {
|
||||
let client = app_state.client.clone();
|
||||
cx.spawn(async move |cx| {
|
||||
client
|
||||
.sign_in_with_optional_connect(true, cx)
|
||||
.await
|
||||
.unwrap();
|
||||
})
|
||||
.shared()
|
||||
})
|
||||
.clone()
|
||||
.await;
|
||||
|
||||
let project = cx.update(|cx| {
|
||||
Project::local(
|
||||
app_state.client.clone(),
|
||||
app_state.node_runtime.clone(),
|
||||
app_state.user_store.clone(),
|
||||
app_state.languages.clone(),
|
||||
app_state.fs.clone(),
|
||||
None,
|
||||
cx,
|
||||
)
|
||||
})?;
|
||||
|
||||
let worktree = project
|
||||
.update(cx, |project, cx| {
|
||||
project.create_worktree(&worktree_path, true, cx)
|
||||
})?
|
||||
.await?;
|
||||
worktree
|
||||
.read_with(cx, |worktree, _cx| {
|
||||
worktree.as_local().unwrap().scan_complete()
|
||||
})?
|
||||
.await;
|
||||
|
||||
anyhow::Ok(project)
|
||||
}
|
||||
|
||||
pub async fn setup_worktree(&self) -> Result<PathBuf> {
|
||||
let (repo_owner, repo_name) = self.repo_name()?;
|
||||
let file_name = self.file_name();
|
||||
|
||||
let repo_dir = REPOS_DIR.join(repo_owner.as_ref()).join(repo_name.as_ref());
|
||||
let repo_lock = lock_repo(&repo_dir).await;
|
||||
|
||||
if !repo_dir.is_dir() {
|
||||
fs::create_dir_all(&repo_dir)?;
|
||||
run_git(&repo_dir, &["init"]).await?;
|
||||
run_git(
|
||||
&repo_dir,
|
||||
&["remote", "add", "origin", &self.example.repository_url],
|
||||
)
|
||||
.await?;
|
||||
}
|
||||
|
||||
// Resolve the example to a revision, fetching it if needed.
|
||||
let revision = run_git(
|
||||
&repo_dir,
|
||||
&[
|
||||
"rev-parse",
|
||||
&format!("{}^{{commit}}", self.example.revision),
|
||||
],
|
||||
)
|
||||
.await;
|
||||
let revision = if let Ok(revision) = revision {
|
||||
revision
|
||||
} else {
|
||||
run_git(
|
||||
&repo_dir,
|
||||
&["fetch", "--depth", "1", "origin", &self.example.revision],
|
||||
)
|
||||
.await?;
|
||||
let revision = run_git(&repo_dir, &["rev-parse", "FETCH_HEAD"]).await?;
|
||||
if revision != self.example.revision {
|
||||
run_git(&repo_dir, &["tag", &self.example.revision, &revision]).await?;
|
||||
}
|
||||
revision
|
||||
};
|
||||
|
||||
// Create the worktree for this example if needed.
|
||||
let worktree_path = WORKTREES_DIR.join(&file_name).join(repo_name.as_ref());
|
||||
if worktree_path.is_dir() {
|
||||
run_git(&worktree_path, &["clean", "--force", "-d"]).await?;
|
||||
run_git(&worktree_path, &["reset", "--hard", "HEAD"]).await?;
|
||||
run_git(&worktree_path, &["checkout", revision.as_str()]).await?;
|
||||
} else {
|
||||
let worktree_path_string = worktree_path.to_string_lossy();
|
||||
run_git(&repo_dir, &["branch", "-f", &file_name, revision.as_str()]).await?;
|
||||
run_git(
|
||||
&repo_dir,
|
||||
&["worktree", "add", "-f", &worktree_path_string, &file_name],
|
||||
)
|
||||
.await?;
|
||||
}
|
||||
drop(repo_lock);
|
||||
|
||||
// Apply the uncommitted diff for this example.
|
||||
if !self.example.uncommitted_diff.is_empty() {
|
||||
let mut apply_process = smol::process::Command::new("git")
|
||||
.current_dir(&worktree_path)
|
||||
.args(&["apply", "-"])
|
||||
.stdin(std::process::Stdio::piped())
|
||||
.spawn()?;
|
||||
|
||||
let mut stdin = apply_process.stdin.take().unwrap();
|
||||
stdin
|
||||
.write_all(self.example.uncommitted_diff.as_bytes())
|
||||
.await?;
|
||||
stdin.close().await?;
|
||||
drop(stdin);
|
||||
|
||||
let apply_result = apply_process.output().await?;
|
||||
if !apply_result.status.success() {
|
||||
anyhow::bail!(
|
||||
"Failed to apply uncommitted diff patch with status: {}\nstderr:\n{}\nstdout:\n{}",
|
||||
apply_result.status,
|
||||
String::from_utf8_lossy(&apply_result.stderr),
|
||||
String::from_utf8_lossy(&apply_result.stdout),
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
Ok(worktree_path)
|
||||
}
|
||||
|
||||
pub fn file_name(&self) -> String {
|
||||
self.name
|
||||
.chars()
|
||||
.map(|c| {
|
||||
if c.is_whitespace() {
|
||||
'-'
|
||||
} else {
|
||||
c.to_ascii_lowercase()
|
||||
}
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn repo_name(&self) -> Result<(Cow<'_, str>, Cow<'_, str>)> {
|
||||
// git@github.com:owner/repo.git
|
||||
if self.example.repository_url.contains('@') {
|
||||
let (owner, repo) = self
|
||||
.example
|
||||
.repository_url
|
||||
.split_once(':')
|
||||
.context("expected : in git url")?
|
||||
.1
|
||||
.split_once('/')
|
||||
.context("expected / in git url")?;
|
||||
Ok((
|
||||
Cow::Borrowed(owner),
|
||||
Cow::Borrowed(repo.trim_end_matches(".git")),
|
||||
))
|
||||
// http://github.com/owner/repo.git
|
||||
} else {
|
||||
let url = Url::parse(&self.example.repository_url)?;
|
||||
let mut segments = url.path_segments().context("empty http url")?;
|
||||
let owner = segments
|
||||
.next()
|
||||
.context("expected owner path segment")?
|
||||
.to_string();
|
||||
let repo = segments
|
||||
.next()
|
||||
.context("expected repo path segment")?
|
||||
.trim_end_matches(".git")
|
||||
.to_string();
|
||||
assert!(segments.next().is_none());
|
||||
|
||||
Ok((owner.into(), repo.into()))
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn cursor_position(
|
||||
&self,
|
||||
project: &Entity<Project>,
|
||||
cx: &mut AsyncApp,
|
||||
) -> Result<(Entity<Buffer>, Anchor)> {
|
||||
let worktree = project.read_with(cx, |project, cx| {
|
||||
project.visible_worktrees(cx).next().unwrap()
|
||||
})?;
|
||||
let cursor_path = RelPath::new(&self.example.cursor_path, PathStyle::Posix)?.into_arc();
|
||||
let cursor_buffer = project
|
||||
.update(cx, |project, cx| {
|
||||
project.open_buffer(
|
||||
ProjectPath {
|
||||
worktree_id: worktree.read(cx).id(),
|
||||
path: cursor_path,
|
||||
},
|
||||
cx,
|
||||
)
|
||||
})?
|
||||
.await?;
|
||||
let cursor_offset_within_excerpt = self
|
||||
.example
|
||||
.cursor_position
|
||||
.find(CURSOR_MARKER)
|
||||
.ok_or_else(|| anyhow!("missing cursor marker"))?;
|
||||
let mut cursor_excerpt = self.example.cursor_position.clone();
|
||||
cursor_excerpt.replace_range(
|
||||
cursor_offset_within_excerpt..(cursor_offset_within_excerpt + CURSOR_MARKER.len()),
|
||||
"",
|
||||
);
|
||||
let excerpt_offset = cursor_buffer.read_with(cx, |buffer, _cx| {
|
||||
let text = buffer.text();
|
||||
|
||||
let mut matches = text.match_indices(&cursor_excerpt);
|
||||
let Some((excerpt_offset, _)) = matches.next() else {
|
||||
anyhow::bail!(
|
||||
"\nExcerpt:\n\n{cursor_excerpt}\nBuffer text:\n{text}\n.Cursor excerpt did not exist in buffer."
|
||||
);
|
||||
};
|
||||
assert!(matches.next().is_none());
|
||||
|
||||
Ok(excerpt_offset)
|
||||
})??;
|
||||
|
||||
let cursor_offset = excerpt_offset + cursor_offset_within_excerpt;
|
||||
let cursor_anchor =
|
||||
cursor_buffer.read_with(cx, |buffer, _| buffer.anchor_after(cursor_offset))?;
|
||||
Ok((cursor_buffer, cursor_anchor))
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub async fn apply_edit_history(
|
||||
&self,
|
||||
project: &Entity<Project>,
|
||||
cx: &mut AsyncApp,
|
||||
) -> Result<OpenedBuffers<'_>> {
|
||||
edit_prediction::udiff::apply_diff(&self.example.edit_history, project, cx).await
|
||||
}
|
||||
}
|
||||
|
||||
async fn run_git(repo_path: &Path, args: &[&str]) -> Result<String> {
|
||||
let output = smol::process::Command::new("git")
|
||||
.current_dir(repo_path)
|
||||
.args(args)
|
||||
.output()
|
||||
.await?;
|
||||
|
||||
anyhow::ensure!(
|
||||
output.status.success(),
|
||||
"`git {}` within `{}` failed with status: {}\nstderr:\n{}\nstdout:\n{}",
|
||||
args.join(" "),
|
||||
repo_path.display(),
|
||||
output.status,
|
||||
String::from_utf8_lossy(&output.stderr),
|
||||
String::from_utf8_lossy(&output.stdout),
|
||||
);
|
||||
Ok(String::from_utf8(output.stdout)?.trim().to_string())
|
||||
}
|
||||
|
||||
impl Display for NamedExample {
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
write!(f, "# {}\n\n", self.name)?;
|
||||
write!(
|
||||
f,
|
||||
"{REPOSITORY_URL_FIELD} = {}\n",
|
||||
self.example.repository_url
|
||||
)?;
|
||||
write!(f, "{REVISION_FIELD} = {}\n\n", self.example.revision)?;
|
||||
|
||||
write!(f, "## {UNCOMMITTED_DIFF_HEADING}\n\n")?;
|
||||
write!(f, "`````diff\n")?;
|
||||
write!(f, "{}", self.example.uncommitted_diff)?;
|
||||
write!(f, "`````\n")?;
|
||||
|
||||
if !self.example.edit_history.is_empty() {
|
||||
write!(f, "`````diff\n{}`````\n", self.example.edit_history)?;
|
||||
}
|
||||
|
||||
write!(
|
||||
f,
|
||||
"## {CURSOR_POSITION_HEADING}\n\n`````{}\n{}`````\n",
|
||||
self.example.cursor_path.display(),
|
||||
self.example.cursor_position
|
||||
)?;
|
||||
write!(f, "## {EDIT_HISTORY_HEADING}\n\n")?;
|
||||
|
||||
if !self.example.expected_patch.is_empty() {
|
||||
write!(
|
||||
f,
|
||||
"\n## {EXPECTED_PATCH_HEADING}\n\n`````diff\n{}`````\n",
|
||||
self.example.expected_patch
|
||||
)?;
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
thread_local! {
|
||||
static REPO_LOCKS: RefCell<HashMap<PathBuf, Arc<Mutex<()>>>> = RefCell::new(HashMap::default());
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub async fn lock_repo(path: impl AsRef<Path>) -> OwnedMutexGuard<()> {
|
||||
REPO_LOCKS
|
||||
.with(|cell| {
|
||||
cell.borrow_mut()
|
||||
.entry(path.as_ref().to_path_buf())
|
||||
.or_default()
|
||||
.clone()
|
||||
})
|
||||
.lock_owned()
|
||||
.await
|
||||
}
|
||||
@@ -0,0 +1,122 @@
|
||||
use client::{Client, ProxySettings, UserStore};
|
||||
use extension::ExtensionHostProxy;
|
||||
use fs::RealFs;
|
||||
use gpui::http_client::read_proxy_from_env;
|
||||
use gpui::{App, AppContext, Entity};
|
||||
use gpui_tokio::Tokio;
|
||||
use language::LanguageRegistry;
|
||||
use language_extension::LspAccess;
|
||||
use node_runtime::{NodeBinaryOptions, NodeRuntime};
|
||||
use project::project_settings::ProjectSettings;
|
||||
use release_channel::{AppCommitSha, AppVersion};
|
||||
use reqwest_client::ReqwestClient;
|
||||
use settings::{Settings, SettingsStore};
|
||||
use std::path::PathBuf;
|
||||
use std::sync::Arc;
|
||||
use util::ResultExt as _;
|
||||
|
||||
/// Headless subset of `workspace::AppState`.
|
||||
pub struct ZetaCliAppState {
|
||||
pub languages: Arc<LanguageRegistry>,
|
||||
pub client: Arc<Client>,
|
||||
pub user_store: Entity<UserStore>,
|
||||
pub fs: Arc<dyn fs::Fs>,
|
||||
pub node_runtime: NodeRuntime,
|
||||
}
|
||||
|
||||
// TODO: dedupe with crates/eval/src/eval.rs
|
||||
pub fn init(cx: &mut App) -> ZetaCliAppState {
|
||||
let app_commit_sha = option_env!("ZED_COMMIT_SHA").map(|s| AppCommitSha::new(s.to_owned()));
|
||||
|
||||
let app_version = AppVersion::load(
|
||||
env!("ZED_PKG_VERSION"),
|
||||
option_env!("ZED_BUILD_ID"),
|
||||
app_commit_sha,
|
||||
);
|
||||
release_channel::init(app_version.clone(), cx);
|
||||
gpui_tokio::init(cx);
|
||||
|
||||
let settings_store = SettingsStore::new(cx, &settings::default_settings());
|
||||
cx.set_global(settings_store);
|
||||
|
||||
// Set User-Agent so we can download language servers from GitHub
|
||||
let user_agent = format!(
|
||||
"Zeta CLI/{} ({}; {})",
|
||||
app_version,
|
||||
std::env::consts::OS,
|
||||
std::env::consts::ARCH
|
||||
);
|
||||
let proxy_str = ProxySettings::get_global(cx).proxy.to_owned();
|
||||
let proxy_url = proxy_str
|
||||
.as_ref()
|
||||
.and_then(|input| input.parse().ok())
|
||||
.or_else(read_proxy_from_env);
|
||||
let http = {
|
||||
let _guard = Tokio::handle(cx).enter();
|
||||
|
||||
ReqwestClient::proxy_and_user_agent(proxy_url, &user_agent)
|
||||
.expect("could not start HTTP client")
|
||||
};
|
||||
cx.set_http_client(Arc::new(http));
|
||||
|
||||
let client = Client::production(cx);
|
||||
cx.set_http_client(client.http_client());
|
||||
|
||||
let git_binary_path = None;
|
||||
let fs = Arc::new(RealFs::new(
|
||||
git_binary_path,
|
||||
cx.background_executor().clone(),
|
||||
));
|
||||
|
||||
let mut languages = LanguageRegistry::new(cx.background_executor().clone());
|
||||
languages.set_language_server_download_dir(paths::languages_dir().clone());
|
||||
let languages = Arc::new(languages);
|
||||
|
||||
let user_store = cx.new(|cx| UserStore::new(client.clone(), cx));
|
||||
|
||||
extension::init(cx);
|
||||
|
||||
let (mut tx, rx) = watch::channel(None);
|
||||
cx.observe_global::<SettingsStore>(move |cx| {
|
||||
let settings = &ProjectSettings::get_global(cx).node;
|
||||
let options = NodeBinaryOptions {
|
||||
allow_path_lookup: !settings.ignore_system_version,
|
||||
allow_binary_download: true,
|
||||
use_paths: settings.path.as_ref().map(|node_path| {
|
||||
let node_path = PathBuf::from(shellexpand::tilde(node_path).as_ref());
|
||||
let npm_path = settings
|
||||
.npm_path
|
||||
.as_ref()
|
||||
.map(|path| PathBuf::from(shellexpand::tilde(&path).as_ref()));
|
||||
(
|
||||
node_path.clone(),
|
||||
npm_path.unwrap_or_else(|| {
|
||||
let base_path = PathBuf::new();
|
||||
node_path.parent().unwrap_or(&base_path).join("npm")
|
||||
}),
|
||||
)
|
||||
}),
|
||||
};
|
||||
tx.send(Some(options)).log_err();
|
||||
})
|
||||
.detach();
|
||||
let node_runtime = NodeRuntime::new(client.http_client(), None, rx);
|
||||
|
||||
let extension_host_proxy = ExtensionHostProxy::global(cx);
|
||||
|
||||
debug_adapter_extension::init(extension_host_proxy.clone(), cx);
|
||||
language_extension::init(LspAccess::Noop, extension_host_proxy, languages.clone());
|
||||
language_model::init(client.clone(), cx);
|
||||
language_models::init(user_store.clone(), client.clone(), cx);
|
||||
languages::init(languages.clone(), fs.clone(), node_runtime.clone(), cx);
|
||||
prompt_store::init(cx);
|
||||
terminal_view::init(cx);
|
||||
|
||||
ZetaCliAppState {
|
||||
languages,
|
||||
client,
|
||||
user_store,
|
||||
fs,
|
||||
node_runtime,
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,507 @@
|
||||
mod evaluate;
|
||||
mod example;
|
||||
mod headless;
|
||||
mod metrics;
|
||||
mod paths;
|
||||
mod predict;
|
||||
mod source_location;
|
||||
mod util;
|
||||
|
||||
use crate::{
|
||||
evaluate::run_evaluate,
|
||||
example::{ExampleFormat, NamedExample},
|
||||
headless::ZetaCliAppState,
|
||||
predict::run_predict,
|
||||
source_location::SourceLocation,
|
||||
util::{open_buffer, open_buffer_with_language_server},
|
||||
};
|
||||
use ::util::paths::PathStyle;
|
||||
use anyhow::{Result, anyhow};
|
||||
use clap::{Args, Parser, Subcommand, ValueEnum};
|
||||
use cloud_llm_client::predict_edits_v3;
|
||||
use edit_prediction::udiff::DiffLine;
|
||||
use edit_prediction_context::EditPredictionExcerptOptions;
|
||||
use gpui::{Application, AsyncApp, Entity, prelude::*};
|
||||
use language::{Bias, Buffer, BufferSnapshot, Point};
|
||||
use metrics::delta_chr_f;
|
||||
use project::{Project, Worktree, lsp_store::OpenLspBufferHandle};
|
||||
use reqwest_client::ReqwestClient;
|
||||
use std::io::{self};
|
||||
use std::{collections::HashSet, path::PathBuf, str::FromStr, sync::Arc};
|
||||
|
||||
#[derive(Parser, Debug)]
|
||||
#[command(name = "zeta")]
|
||||
struct ZetaCliArgs {
|
||||
#[arg(long, default_value_t = false)]
|
||||
printenv: bool,
|
||||
#[command(subcommand)]
|
||||
command: Option<Command>,
|
||||
}
|
||||
|
||||
#[derive(Subcommand, Debug)]
|
||||
enum Command {
|
||||
Context(ContextArgs),
|
||||
Predict(PredictArguments),
|
||||
Eval(EvaluateArguments),
|
||||
ConvertExample {
|
||||
path: PathBuf,
|
||||
#[arg(long, value_enum, default_value_t = ExampleFormat::Md)]
|
||||
output_format: ExampleFormat,
|
||||
},
|
||||
Score {
|
||||
golden_patch: PathBuf,
|
||||
actual_patch: PathBuf,
|
||||
},
|
||||
Clean,
|
||||
}
|
||||
|
||||
#[derive(Debug, Args)]
|
||||
struct ContextArgs {
|
||||
#[arg(long)]
|
||||
provider: ContextProvider,
|
||||
#[arg(long)]
|
||||
worktree: PathBuf,
|
||||
#[arg(long)]
|
||||
cursor: SourceLocation,
|
||||
#[arg(long)]
|
||||
use_language_server: bool,
|
||||
#[arg(long)]
|
||||
edit_history: Option<FileOrStdin>,
|
||||
#[clap(flatten)]
|
||||
zeta2_args: Zeta2Args,
|
||||
}
|
||||
|
||||
#[derive(clap::ValueEnum, Default, Debug, Clone, Copy)]
|
||||
enum ContextProvider {
|
||||
Zeta1,
|
||||
#[default]
|
||||
Zeta2,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Args)]
|
||||
struct Zeta2Args {
|
||||
#[arg(long, default_value_t = 8192)]
|
||||
max_prompt_bytes: usize,
|
||||
#[arg(long, default_value_t = 2048)]
|
||||
max_excerpt_bytes: usize,
|
||||
#[arg(long, default_value_t = 1024)]
|
||||
min_excerpt_bytes: usize,
|
||||
#[arg(long, default_value_t = 0.66)]
|
||||
target_before_cursor_over_total_bytes: f32,
|
||||
#[arg(long, default_value_t = 1024)]
|
||||
max_diagnostic_bytes: usize,
|
||||
#[arg(long, value_enum, default_value_t = PromptFormat::default())]
|
||||
prompt_format: PromptFormat,
|
||||
#[arg(long, value_enum, default_value_t = Default::default())]
|
||||
output_format: OutputFormat,
|
||||
#[arg(long, default_value_t = 42)]
|
||||
file_indexing_parallelism: usize,
|
||||
#[arg(long, default_value_t = false)]
|
||||
disable_imports_gathering: bool,
|
||||
#[arg(long, default_value_t = u8::MAX)]
|
||||
max_retrieved_definitions: u8,
|
||||
}
|
||||
|
||||
#[derive(Debug, Args)]
|
||||
pub struct PredictArguments {
|
||||
#[clap(long, short, value_enum, default_value_t = PredictionsOutputFormat::Md)]
|
||||
format: PredictionsOutputFormat,
|
||||
example_path: PathBuf,
|
||||
#[clap(flatten)]
|
||||
options: PredictionOptions,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Args)]
|
||||
pub struct PredictionOptions {
|
||||
#[clap(flatten)]
|
||||
zeta2: Zeta2Args,
|
||||
#[clap(long)]
|
||||
provider: PredictionProvider,
|
||||
#[clap(long, value_enum, default_value_t = CacheMode::default())]
|
||||
cache: CacheMode,
|
||||
}
|
||||
|
||||
#[derive(Debug, ValueEnum, Default, Clone, Copy, PartialEq)]
|
||||
pub enum CacheMode {
|
||||
/// Use cached LLM requests and responses, except when multiple repetitions are requested
|
||||
#[default]
|
||||
Auto,
|
||||
/// Use cached LLM requests and responses, based on the hash of the prompt and the endpoint.
|
||||
#[value(alias = "request")]
|
||||
Requests,
|
||||
/// Ignore existing cache entries for both LLM and search.
|
||||
Skip,
|
||||
/// Use cached LLM responses AND search results for full determinism. Fails if they haven't been cached yet.
|
||||
/// Useful for reproducing results and fixing bugs outside of search queries
|
||||
Force,
|
||||
}
|
||||
|
||||
impl CacheMode {
|
||||
fn use_cached_llm_responses(&self) -> bool {
|
||||
self.assert_not_auto();
|
||||
matches!(self, CacheMode::Requests | CacheMode::Force)
|
||||
}
|
||||
|
||||
fn use_cached_search_results(&self) -> bool {
|
||||
self.assert_not_auto();
|
||||
matches!(self, CacheMode::Force)
|
||||
}
|
||||
|
||||
fn assert_not_auto(&self) {
|
||||
assert_ne!(
|
||||
*self,
|
||||
CacheMode::Auto,
|
||||
"Cache mode should not be auto at this point!"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(clap::ValueEnum, Debug, Clone)]
|
||||
pub enum PredictionsOutputFormat {
|
||||
Json,
|
||||
Md,
|
||||
Diff,
|
||||
}
|
||||
|
||||
#[derive(Debug, Args)]
|
||||
pub struct EvaluateArguments {
|
||||
example_paths: Vec<PathBuf>,
|
||||
#[clap(flatten)]
|
||||
options: PredictionOptions,
|
||||
#[clap(short, long, default_value_t = 1, alias = "repeat")]
|
||||
repetitions: u16,
|
||||
#[arg(long)]
|
||||
skip_prediction: bool,
|
||||
}
|
||||
|
||||
#[derive(clap::ValueEnum, Default, Debug, Clone, Copy, PartialEq)]
|
||||
enum PredictionProvider {
|
||||
Zeta1,
|
||||
#[default]
|
||||
Zeta2,
|
||||
Sweep,
|
||||
}
|
||||
|
||||
fn zeta2_args_to_options(args: &Zeta2Args) -> edit_prediction::ZetaOptions {
|
||||
edit_prediction::ZetaOptions {
|
||||
context: EditPredictionExcerptOptions {
|
||||
max_bytes: args.max_excerpt_bytes,
|
||||
min_bytes: args.min_excerpt_bytes,
|
||||
target_before_cursor_over_total_bytes: args.target_before_cursor_over_total_bytes,
|
||||
},
|
||||
max_prompt_bytes: args.max_prompt_bytes,
|
||||
prompt_format: args.prompt_format.into(),
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(clap::ValueEnum, Default, Debug, Clone, Copy)]
|
||||
enum PromptFormat {
|
||||
OnlySnippets,
|
||||
#[default]
|
||||
OldTextNewText,
|
||||
Minimal,
|
||||
MinimalQwen,
|
||||
SeedCoder1120,
|
||||
}
|
||||
|
||||
impl Into<predict_edits_v3::PromptFormat> for PromptFormat {
|
||||
fn into(self) -> predict_edits_v3::PromptFormat {
|
||||
match self {
|
||||
Self::OnlySnippets => predict_edits_v3::PromptFormat::OnlySnippets,
|
||||
Self::OldTextNewText => predict_edits_v3::PromptFormat::OldTextNewText,
|
||||
Self::Minimal => predict_edits_v3::PromptFormat::Minimal,
|
||||
Self::MinimalQwen => predict_edits_v3::PromptFormat::MinimalQwen,
|
||||
Self::SeedCoder1120 => predict_edits_v3::PromptFormat::SeedCoder1120,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(clap::ValueEnum, Default, Debug, Clone)]
|
||||
enum OutputFormat {
|
||||
#[default]
|
||||
Prompt,
|
||||
Request,
|
||||
Full,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
enum FileOrStdin {
|
||||
File(PathBuf),
|
||||
Stdin,
|
||||
}
|
||||
|
||||
impl FileOrStdin {
|
||||
async fn read_to_string(&self) -> Result<String, std::io::Error> {
|
||||
match self {
|
||||
FileOrStdin::File(path) => smol::fs::read_to_string(path).await,
|
||||
FileOrStdin::Stdin => smol::unblock(|| std::io::read_to_string(std::io::stdin())).await,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl FromStr for FileOrStdin {
|
||||
type Err = <PathBuf as FromStr>::Err;
|
||||
|
||||
fn from_str(s: &str) -> Result<Self, Self::Err> {
|
||||
match s {
|
||||
"-" => Ok(Self::Stdin),
|
||||
_ => Ok(Self::File(PathBuf::from_str(s)?)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
struct LoadedContext {
|
||||
full_path_str: String,
|
||||
snapshot: BufferSnapshot,
|
||||
clipped_cursor: Point,
|
||||
worktree: Entity<Worktree>,
|
||||
project: Entity<Project>,
|
||||
buffer: Entity<Buffer>,
|
||||
lsp_open_handle: Option<OpenLspBufferHandle>,
|
||||
}
|
||||
|
||||
async fn load_context(
|
||||
args: &ContextArgs,
|
||||
app_state: &Arc<ZetaCliAppState>,
|
||||
cx: &mut AsyncApp,
|
||||
) -> Result<LoadedContext> {
|
||||
let ContextArgs {
|
||||
worktree: worktree_path,
|
||||
cursor,
|
||||
use_language_server,
|
||||
..
|
||||
} = args;
|
||||
|
||||
let worktree_path = worktree_path.canonicalize()?;
|
||||
|
||||
let project = cx.update(|cx| {
|
||||
Project::local(
|
||||
app_state.client.clone(),
|
||||
app_state.node_runtime.clone(),
|
||||
app_state.user_store.clone(),
|
||||
app_state.languages.clone(),
|
||||
app_state.fs.clone(),
|
||||
None,
|
||||
cx,
|
||||
)
|
||||
})?;
|
||||
|
||||
let worktree = project
|
||||
.update(cx, |project, cx| {
|
||||
project.create_worktree(&worktree_path, true, cx)
|
||||
})?
|
||||
.await?;
|
||||
|
||||
let mut ready_languages = HashSet::default();
|
||||
let (lsp_open_handle, buffer) = if *use_language_server {
|
||||
let (lsp_open_handle, _, buffer) = open_buffer_with_language_server(
|
||||
project.clone(),
|
||||
worktree.clone(),
|
||||
cursor.path.clone(),
|
||||
&mut ready_languages,
|
||||
cx,
|
||||
)
|
||||
.await?;
|
||||
(Some(lsp_open_handle), buffer)
|
||||
} else {
|
||||
let buffer =
|
||||
open_buffer(project.clone(), worktree.clone(), cursor.path.clone(), cx).await?;
|
||||
(None, buffer)
|
||||
};
|
||||
|
||||
let full_path_str = worktree
|
||||
.read_with(cx, |worktree, _| worktree.root_name().join(&cursor.path))?
|
||||
.display(PathStyle::local())
|
||||
.to_string();
|
||||
|
||||
let snapshot = cx.update(|cx| buffer.read(cx).snapshot())?;
|
||||
let clipped_cursor = snapshot.clip_point(cursor.point, Bias::Left);
|
||||
if clipped_cursor != cursor.point {
|
||||
let max_row = snapshot.max_point().row;
|
||||
if cursor.point.row < max_row {
|
||||
return Err(anyhow!(
|
||||
"Cursor position {:?} is out of bounds (line length is {})",
|
||||
cursor.point,
|
||||
snapshot.line_len(cursor.point.row)
|
||||
));
|
||||
} else {
|
||||
return Err(anyhow!(
|
||||
"Cursor position {:?} is out of bounds (max row is {})",
|
||||
cursor.point,
|
||||
max_row
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
Ok(LoadedContext {
|
||||
full_path_str,
|
||||
snapshot,
|
||||
clipped_cursor,
|
||||
worktree,
|
||||
project,
|
||||
buffer,
|
||||
lsp_open_handle,
|
||||
})
|
||||
}
|
||||
|
||||
async fn zeta2_context(
|
||||
args: ContextArgs,
|
||||
app_state: &Arc<ZetaCliAppState>,
|
||||
cx: &mut AsyncApp,
|
||||
) -> Result<String> {
|
||||
let LoadedContext {
|
||||
worktree,
|
||||
project,
|
||||
buffer,
|
||||
clipped_cursor,
|
||||
lsp_open_handle: _handle,
|
||||
..
|
||||
} = load_context(&args, app_state, cx).await?;
|
||||
|
||||
// wait for worktree scan before starting zeta2 so that wait_for_initial_indexing waits for
|
||||
// the whole worktree.
|
||||
worktree
|
||||
.read_with(cx, |worktree, _cx| {
|
||||
worktree.as_local().unwrap().scan_complete()
|
||||
})?
|
||||
.await;
|
||||
let output = cx
|
||||
.update(|cx| {
|
||||
let store = cx.new(|cx| {
|
||||
edit_prediction::EditPredictionStore::new(
|
||||
app_state.client.clone(),
|
||||
app_state.user_store.clone(),
|
||||
cx,
|
||||
)
|
||||
});
|
||||
store.update(cx, |store, cx| {
|
||||
store.set_options(zeta2_args_to_options(&args.zeta2_args));
|
||||
store.register_buffer(&buffer, &project, cx);
|
||||
});
|
||||
cx.spawn(async move |cx| {
|
||||
let updates_rx = store.update(cx, |store, cx| {
|
||||
let cursor = buffer.read(cx).snapshot().anchor_before(clipped_cursor);
|
||||
store.set_use_context(true);
|
||||
store.refresh_context(&project, &buffer, cursor, cx);
|
||||
store.project_context_updates(&project).unwrap()
|
||||
})?;
|
||||
|
||||
updates_rx.recv().await.ok();
|
||||
|
||||
let context = store.update(cx, |store, cx| {
|
||||
store.context_for_project(&project, cx).to_vec()
|
||||
})?;
|
||||
|
||||
anyhow::Ok(serde_json::to_string_pretty(&context).unwrap())
|
||||
})
|
||||
})?
|
||||
.await?;
|
||||
|
||||
Ok(output)
|
||||
}
|
||||
|
||||
async fn zeta1_context(
|
||||
args: ContextArgs,
|
||||
app_state: &Arc<ZetaCliAppState>,
|
||||
cx: &mut AsyncApp,
|
||||
) -> Result<edit_prediction::zeta1::GatherContextOutput> {
|
||||
let LoadedContext {
|
||||
full_path_str,
|
||||
snapshot,
|
||||
clipped_cursor,
|
||||
..
|
||||
} = load_context(&args, app_state, cx).await?;
|
||||
|
||||
let events = match args.edit_history {
|
||||
Some(events) => events.read_to_string().await?,
|
||||
None => String::new(),
|
||||
};
|
||||
|
||||
let prompt_for_events = move || (events, 0);
|
||||
cx.update(|cx| {
|
||||
edit_prediction::zeta1::gather_context(
|
||||
full_path_str,
|
||||
&snapshot,
|
||||
clipped_cursor,
|
||||
prompt_for_events,
|
||||
cloud_llm_client::PredictEditsRequestTrigger::Cli,
|
||||
cx,
|
||||
)
|
||||
})?
|
||||
.await
|
||||
}
|
||||
|
||||
fn main() {
|
||||
zlog::init();
|
||||
zlog::init_output_stderr();
|
||||
let args = ZetaCliArgs::parse();
|
||||
let http_client = Arc::new(ReqwestClient::new());
|
||||
let app = Application::headless().with_http_client(http_client);
|
||||
|
||||
app.run(move |cx| {
|
||||
let app_state = Arc::new(headless::init(cx));
|
||||
cx.spawn(async move |cx| {
|
||||
match args.command {
|
||||
None => {
|
||||
if args.printenv {
|
||||
::util::shell_env::print_env();
|
||||
} else {
|
||||
panic!("Expected a command");
|
||||
}
|
||||
}
|
||||
Some(Command::Context(context_args)) => {
|
||||
let result = match context_args.provider {
|
||||
ContextProvider::Zeta1 => {
|
||||
let context =
|
||||
zeta1_context(context_args, &app_state, cx).await.unwrap();
|
||||
serde_json::to_string_pretty(&context.body).unwrap()
|
||||
}
|
||||
ContextProvider::Zeta2 => {
|
||||
zeta2_context(context_args, &app_state, cx).await.unwrap()
|
||||
}
|
||||
};
|
||||
println!("{}", result);
|
||||
}
|
||||
Some(Command::Predict(arguments)) => {
|
||||
run_predict(arguments, &app_state, cx).await;
|
||||
}
|
||||
Some(Command::Eval(arguments)) => {
|
||||
run_evaluate(arguments, &app_state, cx).await;
|
||||
}
|
||||
Some(Command::ConvertExample {
|
||||
path,
|
||||
output_format,
|
||||
}) => {
|
||||
let example = NamedExample::load(path).unwrap();
|
||||
example.write(output_format, io::stdout()).unwrap();
|
||||
}
|
||||
Some(Command::Score {
|
||||
golden_patch,
|
||||
actual_patch,
|
||||
}) => {
|
||||
let golden_content = std::fs::read_to_string(golden_patch).unwrap();
|
||||
let actual_content = std::fs::read_to_string(actual_patch).unwrap();
|
||||
|
||||
let golden_diff: Vec<DiffLine> = golden_content
|
||||
.lines()
|
||||
.map(|line| DiffLine::parse(line))
|
||||
.collect();
|
||||
|
||||
let actual_diff: Vec<DiffLine> = actual_content
|
||||
.lines()
|
||||
.map(|line| DiffLine::parse(line))
|
||||
.collect();
|
||||
|
||||
let score = delta_chr_f(&golden_diff, &actual_diff);
|
||||
println!("{:.2}", score);
|
||||
}
|
||||
Some(Command::Clean) => {
|
||||
std::fs::remove_dir_all(&*crate::paths::TARGET_ZETA_DIR).unwrap()
|
||||
}
|
||||
};
|
||||
|
||||
let _ = cx.update(|cx| cx.quit());
|
||||
})
|
||||
.detach();
|
||||
});
|
||||
}
|
||||
@@ -0,0 +1,380 @@
|
||||
use collections::{HashMap, HashSet};
|
||||
use edit_prediction::udiff::DiffLine;
|
||||
|
||||
type Counts = HashMap<String, usize>;
|
||||
type CountsDelta = HashMap<String, isize>;
|
||||
|
||||
#[derive(Default, Debug, Clone)]
|
||||
pub struct Scores {
|
||||
pub true_positives: usize,
|
||||
pub false_positives: usize,
|
||||
pub false_negatives: usize,
|
||||
}
|
||||
|
||||
impl Scores {
|
||||
pub fn from_sets(expected: &HashSet<String>, actual: &HashSet<String>) -> Scores {
|
||||
let true_positives = expected.intersection(actual).count();
|
||||
let false_positives = actual.difference(expected).count();
|
||||
let false_negatives = expected.difference(actual).count();
|
||||
|
||||
Scores {
|
||||
true_positives,
|
||||
false_positives,
|
||||
false_negatives,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn from_counts(expected: &Counts, actual: &Counts) -> Scores {
|
||||
let mut true_positives = 0;
|
||||
let mut false_positives = 0;
|
||||
let mut false_negatives = 0;
|
||||
|
||||
for (ngram, &expected_count) in expected {
|
||||
let actual_count = *actual.get(ngram).unwrap_or(&0);
|
||||
if actual_count > expected_count {
|
||||
false_positives += actual_count - expected_count;
|
||||
} else {
|
||||
false_negatives += expected_count - actual_count;
|
||||
}
|
||||
true_positives += expected_count.min(actual_count);
|
||||
}
|
||||
|
||||
for (ngram, &actual_count) in actual {
|
||||
if !expected.contains_key(ngram) {
|
||||
false_positives += actual_count;
|
||||
}
|
||||
}
|
||||
|
||||
Scores {
|
||||
true_positives,
|
||||
false_positives,
|
||||
false_negatives,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn to_markdown(&self) -> String {
|
||||
format!(
|
||||
"
|
||||
Precision : {:.4}
|
||||
Recall : {:.4}
|
||||
F1 Score : {:.4}
|
||||
True Positives : {}
|
||||
False Positives : {}
|
||||
False Negatives : {}",
|
||||
self.precision(),
|
||||
self.recall(),
|
||||
self.f1_score(),
|
||||
self.true_positives,
|
||||
self.false_positives,
|
||||
self.false_negatives
|
||||
)
|
||||
}
|
||||
|
||||
pub fn aggregate<'a>(scores: impl Iterator<Item = &'a Scores>) -> Scores {
|
||||
let mut true_positives = 0;
|
||||
let mut false_positives = 0;
|
||||
let mut false_negatives = 0;
|
||||
|
||||
for score in scores {
|
||||
true_positives += score.true_positives;
|
||||
false_positives += score.false_positives;
|
||||
false_negatives += score.false_negatives;
|
||||
}
|
||||
|
||||
Scores {
|
||||
true_positives,
|
||||
false_positives,
|
||||
false_negatives,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn precision(&self) -> f64 {
|
||||
if self.true_positives + self.false_positives == 0 {
|
||||
0.0
|
||||
} else {
|
||||
self.true_positives as f64 / (self.true_positives + self.false_positives) as f64
|
||||
}
|
||||
}
|
||||
|
||||
pub fn recall(&self) -> f64 {
|
||||
if self.true_positives + self.false_negatives == 0 {
|
||||
0.0
|
||||
} else {
|
||||
self.true_positives as f64 / (self.true_positives + self.false_negatives) as f64
|
||||
}
|
||||
}
|
||||
|
||||
pub fn f1_score(&self) -> f64 {
|
||||
let recall = self.recall();
|
||||
let precision = self.precision();
|
||||
if precision + recall == 0.0 {
|
||||
0.0
|
||||
} else {
|
||||
2.0 * precision * recall / (precision + recall)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn line_match_score(expected_patch: &[DiffLine], actual_patch: &[DiffLine]) -> Scores {
|
||||
let expected_change_lines = expected_patch
|
||||
.iter()
|
||||
.filter(|line| matches!(line, DiffLine::Addition(_) | DiffLine::Deletion(_)))
|
||||
.map(|line| line.to_string())
|
||||
.collect();
|
||||
|
||||
let actual_change_lines = actual_patch
|
||||
.iter()
|
||||
.filter(|line| matches!(line, DiffLine::Addition(_) | DiffLine::Deletion(_)))
|
||||
.map(|line| line.to_string())
|
||||
.collect();
|
||||
|
||||
Scores::from_sets(&expected_change_lines, &actual_change_lines)
|
||||
}
|
||||
|
||||
enum ChrfWhitespace {
|
||||
#[allow(unused)]
|
||||
Unchanged,
|
||||
Ignore,
|
||||
}
|
||||
|
||||
const CHR_F_CHAR_ORDER: usize = 6;
|
||||
const CHR_F_BETA: f64 = 2.0;
|
||||
const CHR_F_WHITESPACE: ChrfWhitespace = ChrfWhitespace::Ignore;
|
||||
|
||||
/// Computes a delta-chrF score that compares two sets of edits.
|
||||
///
|
||||
/// This metric works by:
|
||||
/// 1. Reconstructing original, golden (expected result), and actual texts from diffs
|
||||
/// 2. Computing n-gram count differences (deltas) between original→golden and original→actual
|
||||
/// 3. Comparing these deltas to measure how well actual edits match expected edits
|
||||
pub fn delta_chr_f(expected: &[DiffLine], actual: &[DiffLine]) -> f64 {
|
||||
// Reconstruct texts from diffs
|
||||
let mut original_text = String::new(); // state of the text before any edits
|
||||
let mut golden_text = String::new(); // text after applying golden edits
|
||||
let mut actual_text = String::new(); // text after applying actual edits
|
||||
|
||||
for line in expected {
|
||||
match line {
|
||||
DiffLine::Context(s) => {
|
||||
original_text.push_str(s);
|
||||
golden_text.push_str(s);
|
||||
}
|
||||
DiffLine::Deletion(s) => {
|
||||
original_text.push_str(s);
|
||||
}
|
||||
DiffLine::Addition(s) => {
|
||||
golden_text.push_str(s);
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
|
||||
for line in actual {
|
||||
match line {
|
||||
DiffLine::Context(s) | DiffLine::Addition(s) => {
|
||||
actual_text.push_str(s);
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
|
||||
// Edge case
|
||||
if original_text == golden_text && golden_text == actual_text {
|
||||
return 100.0;
|
||||
}
|
||||
|
||||
// Compute the metric
|
||||
let original_ngrams = chr_f_ngram_counts(&original_text);
|
||||
let golden_ngrams = chr_f_ngram_counts(&golden_text);
|
||||
let actual_ngrams = chr_f_ngram_counts(&actual_text);
|
||||
|
||||
let mut total_precision = 0.0;
|
||||
let mut total_recall = 0.0;
|
||||
|
||||
for order in 0..CHR_F_CHAR_ORDER {
|
||||
let expected_delta = compute_ngram_delta(&golden_ngrams[order], &original_ngrams[order]);
|
||||
let actual_delta = compute_ngram_delta(&actual_ngrams[order], &original_ngrams[order]);
|
||||
|
||||
if expected_delta.is_empty() && actual_delta.is_empty() {
|
||||
total_precision += 1.0;
|
||||
total_recall += 1.0;
|
||||
continue;
|
||||
}
|
||||
|
||||
let expected_counts = ngram_delta_to_counts(&expected_delta);
|
||||
let actual_counts = ngram_delta_to_counts(&actual_delta);
|
||||
|
||||
let score = Scores::from_counts(&expected_counts, &actual_counts);
|
||||
total_precision += score.precision();
|
||||
total_recall += score.recall();
|
||||
}
|
||||
|
||||
let prec = total_precision / CHR_F_CHAR_ORDER as f64;
|
||||
let recall = total_recall / CHR_F_CHAR_ORDER as f64;
|
||||
let f_score = if prec + recall == 0.0 {
|
||||
0.0
|
||||
} else {
|
||||
(1.0 + CHR_F_BETA * CHR_F_BETA) * prec * recall / (CHR_F_BETA * CHR_F_BETA * prec + recall)
|
||||
};
|
||||
|
||||
f_score * 100.0
|
||||
}
|
||||
|
||||
fn chr_f_ngram_counts(text: &str) -> Vec<Counts> {
|
||||
// Ignore whitespace. The original chrF implementation skips all
|
||||
// whitespace. We should consider compressing multiple consecutive
|
||||
// spaces into one -- this may reflect our task more closely.
|
||||
let text = match CHR_F_WHITESPACE {
|
||||
ChrfWhitespace::Unchanged => text.to_string(),
|
||||
ChrfWhitespace::Ignore => text
|
||||
.chars()
|
||||
.filter(|c| !c.is_whitespace())
|
||||
.collect::<String>(),
|
||||
};
|
||||
|
||||
(1..=CHR_F_CHAR_ORDER)
|
||||
.map(|order| count_ngrams(&text, order))
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn compute_ngram_delta(after: &Counts, before: &Counts) -> CountsDelta {
|
||||
let mut delta = CountsDelta::default();
|
||||
|
||||
for (ngram, &before_count) in before {
|
||||
let after_count = *after.get(ngram).unwrap_or(&0);
|
||||
delta.insert(ngram.clone(), after_count as isize - before_count as isize);
|
||||
}
|
||||
|
||||
for (ngram, &after_count) in after {
|
||||
if !before.contains_key(ngram) {
|
||||
delta.insert(ngram.clone(), after_count as isize);
|
||||
}
|
||||
}
|
||||
|
||||
delta
|
||||
}
|
||||
|
||||
/// Convert negative counts to special deletion tokens.
|
||||
/// For example, if expected delta is {"foo": -1} and actual delta is {"bar": -1},
|
||||
/// we convert it to {"¬foo": +1} and {"¬bar": +1}. This way _not_ deleting "foo"
|
||||
/// will result in a false negative, and mistakenly deleting "bar" will result in a false positive.
|
||||
fn ngram_delta_to_counts(delta: &CountsDelta) -> Counts {
|
||||
let mut counts = Counts::default();
|
||||
|
||||
for (ngram, &delta) in delta {
|
||||
if delta > 0 {
|
||||
counts.insert(ngram.clone(), delta as usize);
|
||||
} else {
|
||||
counts.insert(format!("¬{ngram}"), delta.unsigned_abs());
|
||||
}
|
||||
}
|
||||
|
||||
counts
|
||||
}
|
||||
|
||||
fn count_ngrams(text: &str, n: usize) -> Counts {
|
||||
let chars: Vec<char> = text.chars().collect();
|
||||
let mut counts = Counts::default();
|
||||
|
||||
for window in chars.windows(n) {
|
||||
let ngram: String = window.iter().collect();
|
||||
*counts.entry(ngram).or_insert(0) += 1;
|
||||
}
|
||||
|
||||
counts
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod test {
|
||||
use super::*;
|
||||
use edit_prediction::udiff::DiffLine;
|
||||
|
||||
#[test]
|
||||
fn test_delta_chr_f_perfect_match() {
|
||||
let diff = vec![
|
||||
DiffLine::Context("fn main() {"),
|
||||
DiffLine::Deletion(" println!(\"Hello\");"),
|
||||
DiffLine::Addition(" println!(\"Hello, World!\");"),
|
||||
DiffLine::Context("}"),
|
||||
];
|
||||
|
||||
let score = delta_chr_f(&diff, &diff);
|
||||
assert!((score - 100.0).abs() < 1e-2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_delta_chr_f_wrong_edit() {
|
||||
// When the edit is wrong
|
||||
let expected = vec![
|
||||
DiffLine::Context("one "),
|
||||
DiffLine::Deletion("two "),
|
||||
DiffLine::Context("three"),
|
||||
];
|
||||
|
||||
let actual = vec![
|
||||
DiffLine::Context("one "),
|
||||
DiffLine::Context("two "),
|
||||
DiffLine::Deletion("three"),
|
||||
DiffLine::Addition("four"),
|
||||
];
|
||||
|
||||
// Then the score should be low
|
||||
let score = delta_chr_f(&expected, &actual);
|
||||
assert!(score > 20.0 && score < 40.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_delta_chr_f_partial_match() {
|
||||
let expected = vec![
|
||||
DiffLine::Deletion("let x = 42;"),
|
||||
DiffLine::Addition("let x = 100;"),
|
||||
];
|
||||
|
||||
let actual = vec![
|
||||
DiffLine::Deletion("let x = 42;"),
|
||||
DiffLine::Addition("let x = 99;"),
|
||||
];
|
||||
|
||||
// We got the edit location right, but the replacement text is wrong.
|
||||
// Deleted ngrams will match, bringing the score somewhere in the middle.
|
||||
let score = delta_chr_f(&expected, &actual);
|
||||
assert!(score > 40.0 && score < 60.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_delta_chr_f_missed_edit() {
|
||||
// When predictions makes no changes
|
||||
let expected = vec![
|
||||
DiffLine::Context("prefix "),
|
||||
DiffLine::Deletion("old"),
|
||||
DiffLine::Addition("new"),
|
||||
DiffLine::Context(" suffix"),
|
||||
];
|
||||
|
||||
let actual = vec![
|
||||
DiffLine::Context("prefix "),
|
||||
DiffLine::Context("old"),
|
||||
DiffLine::Context(" suffix"),
|
||||
];
|
||||
|
||||
// Then the score should be low (all expected changes are false negatives)
|
||||
let score = delta_chr_f(&expected, &actual);
|
||||
assert!(score < 20.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_delta_chr_f_extra_edit() {
|
||||
// When adding unexpected content
|
||||
let expected = vec![DiffLine::Context("hello"), DiffLine::Context("world")];
|
||||
|
||||
let actual = vec![
|
||||
DiffLine::Context("hello"),
|
||||
DiffLine::Addition("extra"),
|
||||
DiffLine::Context("world"),
|
||||
];
|
||||
|
||||
// Then the score should be low (all actual changes are false positives)
|
||||
let score = delta_chr_f(&expected, &actual);
|
||||
assert!(score < 20.0);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,57 @@
|
||||
use std::{env, path::PathBuf, sync::LazyLock};
|
||||
|
||||
pub static TARGET_ZETA_DIR: LazyLock<PathBuf> =
|
||||
LazyLock::new(|| env::current_dir().unwrap().join("target/zeta"));
|
||||
pub static CACHE_DIR: LazyLock<PathBuf> = LazyLock::new(|| TARGET_ZETA_DIR.join("cache"));
|
||||
pub static REPOS_DIR: LazyLock<PathBuf> = LazyLock::new(|| TARGET_ZETA_DIR.join("repos"));
|
||||
pub static WORKTREES_DIR: LazyLock<PathBuf> = LazyLock::new(|| TARGET_ZETA_DIR.join("worktrees"));
|
||||
pub static RUN_DIR: LazyLock<PathBuf> = LazyLock::new(|| {
|
||||
TARGET_ZETA_DIR
|
||||
.join("runs")
|
||||
.join(chrono::Local::now().format("%d-%m-%y-%H_%M_%S").to_string())
|
||||
});
|
||||
pub static LATEST_EXAMPLE_RUN_DIR: LazyLock<PathBuf> =
|
||||
LazyLock::new(|| TARGET_ZETA_DIR.join("latest"));
|
||||
|
||||
pub fn print_run_data_dir(deep: bool, use_color: bool) {
|
||||
println!("\n## Run Data\n");
|
||||
let mut files = Vec::new();
|
||||
|
||||
let current_dir = std::env::current_dir().unwrap();
|
||||
for file in std::fs::read_dir(&*RUN_DIR).unwrap() {
|
||||
let file = file.unwrap();
|
||||
if file.file_type().unwrap().is_dir() && deep {
|
||||
for file in std::fs::read_dir(file.path()).unwrap() {
|
||||
let path = file.unwrap().path();
|
||||
let path = path.strip_prefix(¤t_dir).unwrap_or(&path);
|
||||
files.push(format!(
|
||||
"- {}/{}{}{}",
|
||||
path.parent().unwrap().display(),
|
||||
if use_color { "\x1b[34m" } else { "" },
|
||||
path.file_name().unwrap().display(),
|
||||
if use_color { "\x1b[0m" } else { "" },
|
||||
));
|
||||
}
|
||||
} else {
|
||||
let path = file.path();
|
||||
let path = path.strip_prefix(¤t_dir).unwrap_or(&path);
|
||||
files.push(format!(
|
||||
"- {}/{}{}{}",
|
||||
path.parent().unwrap().display(),
|
||||
if use_color { "\x1b[34m" } else { "" },
|
||||
path.file_name().unwrap().display(),
|
||||
if use_color { "\x1b[0m" } else { "" }
|
||||
));
|
||||
}
|
||||
}
|
||||
files.sort();
|
||||
|
||||
for file in files {
|
||||
println!("{}", file);
|
||||
}
|
||||
|
||||
println!(
|
||||
"\n💡 Tip of the day: {} always points to the latest run\n",
|
||||
LATEST_EXAMPLE_RUN_DIR.display()
|
||||
);
|
||||
}
|
||||
@@ -0,0 +1,373 @@
|
||||
use crate::example::{ActualExcerpt, NamedExample};
|
||||
use crate::headless::ZetaCliAppState;
|
||||
use crate::paths::{CACHE_DIR, LATEST_EXAMPLE_RUN_DIR, RUN_DIR, print_run_data_dir};
|
||||
use crate::{
|
||||
CacheMode, PredictArguments, PredictionOptions, PredictionProvider, PredictionsOutputFormat,
|
||||
};
|
||||
use ::serde::Serialize;
|
||||
use anyhow::{Context, Result, anyhow};
|
||||
use cloud_zeta2_prompt::{CURSOR_MARKER, write_codeblock};
|
||||
use edit_prediction::{EditPredictionStore, EvalCache, EvalCacheEntryKind, EvalCacheKey};
|
||||
use futures::StreamExt as _;
|
||||
use gpui::{AppContext, AsyncApp, Entity};
|
||||
use project::Project;
|
||||
use project::buffer_store::BufferStoreEvent;
|
||||
use serde::Deserialize;
|
||||
use std::fs;
|
||||
use std::io::{IsTerminal, Write};
|
||||
use std::path::PathBuf;
|
||||
use std::sync::Arc;
|
||||
use std::sync::Mutex;
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
pub async fn run_predict(
|
||||
args: PredictArguments,
|
||||
app_state: &Arc<ZetaCliAppState>,
|
||||
cx: &mut AsyncApp,
|
||||
) {
|
||||
let example = NamedExample::load(args.example_path).unwrap();
|
||||
let project = example.setup_project(app_state, cx).await.unwrap();
|
||||
let store = setup_store(args.options.provider, &project, app_state, cx).unwrap();
|
||||
let _edited_buffers = example.apply_edit_history(&project, cx).await.unwrap();
|
||||
let result = perform_predict(example, project, store, None, args.options, cx)
|
||||
.await
|
||||
.unwrap();
|
||||
result.write(args.format, std::io::stdout()).unwrap();
|
||||
|
||||
print_run_data_dir(true, std::io::stdout().is_terminal());
|
||||
}
|
||||
|
||||
pub fn setup_store(
|
||||
provider: PredictionProvider,
|
||||
project: &Entity<Project>,
|
||||
app_state: &Arc<ZetaCliAppState>,
|
||||
cx: &mut AsyncApp,
|
||||
) -> Result<Entity<EditPredictionStore>> {
|
||||
let store = cx.new(|cx| {
|
||||
edit_prediction::EditPredictionStore::new(
|
||||
app_state.client.clone(),
|
||||
app_state.user_store.clone(),
|
||||
cx,
|
||||
)
|
||||
})?;
|
||||
|
||||
store.update(cx, |store, _cx| {
|
||||
let model = match provider {
|
||||
PredictionProvider::Zeta1 => edit_prediction::EditPredictionModel::Zeta1,
|
||||
PredictionProvider::Zeta2 => edit_prediction::EditPredictionModel::Zeta2,
|
||||
PredictionProvider::Sweep => edit_prediction::EditPredictionModel::Sweep,
|
||||
};
|
||||
store.set_edit_prediction_model(model);
|
||||
})?;
|
||||
|
||||
let buffer_store = project.read_with(cx, |project, _| project.buffer_store().clone())?;
|
||||
|
||||
cx.subscribe(&buffer_store, {
|
||||
let project = project.clone();
|
||||
let store = store.clone();
|
||||
move |_, event, cx| match event {
|
||||
BufferStoreEvent::BufferAdded(buffer) => {
|
||||
store.update(cx, |store, cx| store.register_buffer(&buffer, &project, cx));
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
})?
|
||||
.detach();
|
||||
|
||||
anyhow::Ok(store)
|
||||
}
|
||||
|
||||
pub async fn perform_predict(
|
||||
example: NamedExample,
|
||||
project: Entity<Project>,
|
||||
store: Entity<EditPredictionStore>,
|
||||
repetition_ix: Option<u16>,
|
||||
options: PredictionOptions,
|
||||
cx: &mut AsyncApp,
|
||||
) -> Result<PredictionDetails> {
|
||||
let mut cache_mode = options.cache;
|
||||
if repetition_ix.is_some() {
|
||||
if cache_mode != CacheMode::Auto && cache_mode != CacheMode::Skip {
|
||||
panic!("Repetitions are not supported in Auto cache mode");
|
||||
} else {
|
||||
cache_mode = CacheMode::Skip;
|
||||
}
|
||||
} else if cache_mode == CacheMode::Auto {
|
||||
cache_mode = CacheMode::Requests;
|
||||
}
|
||||
|
||||
let mut example_run_dir = RUN_DIR.join(&example.file_name());
|
||||
if let Some(repetition_ix) = repetition_ix {
|
||||
example_run_dir = example_run_dir.join(format!("{:03}", repetition_ix));
|
||||
}
|
||||
fs::create_dir_all(&example_run_dir)?;
|
||||
if LATEST_EXAMPLE_RUN_DIR.is_symlink() {
|
||||
fs::remove_file(&*LATEST_EXAMPLE_RUN_DIR)?;
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
std::os::unix::fs::symlink(&example_run_dir, &*LATEST_EXAMPLE_RUN_DIR)
|
||||
.context("creating latest link")?;
|
||||
|
||||
#[cfg(windows)]
|
||||
std::os::windows::fs::symlink_dir(&example_run_dir, &*LATEST_EXAMPLE_RUN_DIR)
|
||||
.context("creating latest link")?;
|
||||
|
||||
store.update(cx, |store, _cx| {
|
||||
store.with_eval_cache(Arc::new(RunCache {
|
||||
example_run_dir: example_run_dir.clone(),
|
||||
cache_mode,
|
||||
}));
|
||||
})?;
|
||||
|
||||
let (cursor_buffer, cursor_anchor) = example.cursor_position(&project, cx).await?;
|
||||
|
||||
let result = Arc::new(Mutex::new(PredictionDetails::new(example_run_dir.clone())));
|
||||
|
||||
let prompt_format = options.zeta2.prompt_format;
|
||||
|
||||
store.update(cx, |store, _cx| {
|
||||
let mut options = store.options().clone();
|
||||
options.prompt_format = prompt_format.into();
|
||||
store.set_options(options);
|
||||
})?;
|
||||
|
||||
let mut debug_task = gpui::Task::ready(Ok(()));
|
||||
|
||||
if options.provider == crate::PredictionProvider::Zeta2 {
|
||||
let mut debug_rx = store.update(cx, |store, _| store.debug_info())?;
|
||||
|
||||
debug_task = cx.background_spawn({
|
||||
let result = result.clone();
|
||||
async move {
|
||||
let mut start_time = None;
|
||||
let mut retrieval_finished_at = None;
|
||||
while let Some(event) = debug_rx.next().await {
|
||||
match event {
|
||||
edit_prediction::DebugEvent::ContextRetrievalStarted(info) => {
|
||||
start_time = Some(info.timestamp);
|
||||
fs::write(
|
||||
example_run_dir.join("search_prompt.md"),
|
||||
&info.search_prompt,
|
||||
)?;
|
||||
}
|
||||
edit_prediction::DebugEvent::ContextRetrievalFinished(info) => {
|
||||
retrieval_finished_at = Some(info.timestamp);
|
||||
for (key, value) in &info.metadata {
|
||||
if *key == "search_queries" {
|
||||
fs::write(
|
||||
example_run_dir.join("search_queries.json"),
|
||||
value.as_bytes(),
|
||||
)?;
|
||||
}
|
||||
}
|
||||
}
|
||||
edit_prediction::DebugEvent::EditPredictionRequested(request) => {
|
||||
let prediction_started_at = Instant::now();
|
||||
start_time.get_or_insert(prediction_started_at);
|
||||
let prompt = request.local_prompt.unwrap_or_default();
|
||||
fs::write(example_run_dir.join("prediction_prompt.md"), &prompt)?;
|
||||
|
||||
{
|
||||
let mut result = result.lock().unwrap();
|
||||
result.prompt_len = prompt.chars().count();
|
||||
|
||||
for included_file in request.inputs.included_files {
|
||||
let insertions =
|
||||
vec![(request.inputs.cursor_point, CURSOR_MARKER)];
|
||||
result.excerpts.extend(included_file.excerpts.iter().map(
|
||||
|excerpt| ActualExcerpt {
|
||||
path: included_file.path.components().skip(1).collect(),
|
||||
text: String::from(excerpt.text.as_ref()),
|
||||
},
|
||||
));
|
||||
write_codeblock(
|
||||
&included_file.path,
|
||||
included_file.excerpts.iter(),
|
||||
if included_file.path == request.inputs.cursor_path {
|
||||
&insertions
|
||||
} else {
|
||||
&[]
|
||||
},
|
||||
included_file.max_row,
|
||||
false,
|
||||
&mut result.excerpts_text,
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
let response =
|
||||
request.response_rx.await?.0.map_err(|err| anyhow!(err))?;
|
||||
let response = edit_prediction::zeta2::text_from_response(response)
|
||||
.unwrap_or_default();
|
||||
let prediction_finished_at = Instant::now();
|
||||
fs::write(example_run_dir.join("prediction_response.md"), &response)?;
|
||||
|
||||
let mut result = result.lock().unwrap();
|
||||
result.generated_len = response.chars().count();
|
||||
result.retrieval_time =
|
||||
retrieval_finished_at.unwrap() - start_time.unwrap();
|
||||
result.prediction_time = prediction_finished_at - prediction_started_at;
|
||||
result.total_time = prediction_finished_at - start_time.unwrap();
|
||||
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
anyhow::Ok(())
|
||||
}
|
||||
});
|
||||
|
||||
store.update(cx, |store, cx| {
|
||||
store.refresh_context(&project, &cursor_buffer, cursor_anchor, cx)
|
||||
})?;
|
||||
}
|
||||
|
||||
let prediction = store
|
||||
.update(cx, |store, cx| {
|
||||
store.request_prediction(
|
||||
&project,
|
||||
&cursor_buffer,
|
||||
cursor_anchor,
|
||||
cloud_llm_client::PredictEditsRequestTrigger::Cli,
|
||||
cx,
|
||||
)
|
||||
})?
|
||||
.await?;
|
||||
|
||||
debug_task.await?;
|
||||
|
||||
let mut result = Arc::into_inner(result).unwrap().into_inner().unwrap();
|
||||
|
||||
result.diff = prediction
|
||||
.and_then(|prediction| {
|
||||
let prediction = prediction.prediction.ok()?;
|
||||
prediction.edit_preview.as_unified_diff(&prediction.edits)
|
||||
})
|
||||
.unwrap_or_default();
|
||||
|
||||
anyhow::Ok(result)
|
||||
}
|
||||
|
||||
struct RunCache {
|
||||
cache_mode: CacheMode,
|
||||
example_run_dir: PathBuf,
|
||||
}
|
||||
|
||||
impl RunCache {
|
||||
fn output_cache_path((kind, key): &EvalCacheKey) -> PathBuf {
|
||||
CACHE_DIR.join(format!("{kind}_out_{key:x}.json",))
|
||||
}
|
||||
|
||||
fn input_cache_path((kind, key): &EvalCacheKey) -> PathBuf {
|
||||
CACHE_DIR.join(format!("{kind}_in_{key:x}.json",))
|
||||
}
|
||||
|
||||
fn link_to_run(&self, key: &EvalCacheKey) {
|
||||
let output_link_path = self.example_run_dir.join(format!("{}_out.json", key.0));
|
||||
fs::hard_link(Self::output_cache_path(key), &output_link_path).unwrap();
|
||||
|
||||
let input_link_path = self.example_run_dir.join(format!("{}_in.json", key.0));
|
||||
fs::hard_link(Self::input_cache_path(key), &input_link_path).unwrap();
|
||||
}
|
||||
}
|
||||
|
||||
impl EvalCache for RunCache {
|
||||
fn read(&self, key: EvalCacheKey) -> Option<String> {
|
||||
let path = RunCache::output_cache_path(&key);
|
||||
|
||||
if path.exists() {
|
||||
let use_cache = match key.0 {
|
||||
EvalCacheEntryKind::Search => self.cache_mode.use_cached_search_results(),
|
||||
EvalCacheEntryKind::Context | EvalCacheEntryKind::Prediction => {
|
||||
self.cache_mode.use_cached_llm_responses()
|
||||
}
|
||||
};
|
||||
if use_cache {
|
||||
log::info!("Using cache entry: {}", path.display());
|
||||
self.link_to_run(&key);
|
||||
Some(fs::read_to_string(path).unwrap())
|
||||
} else {
|
||||
log::trace!("Skipping cached entry: {}", path.display());
|
||||
None
|
||||
}
|
||||
} else if matches!(self.cache_mode, CacheMode::Force) {
|
||||
panic!(
|
||||
"No cached entry found for {:?}. Run without `--cache force` at least once.",
|
||||
key.0
|
||||
);
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
fn write(&self, key: EvalCacheKey, input: &str, output: &str) {
|
||||
fs::create_dir_all(&*CACHE_DIR).unwrap();
|
||||
|
||||
let input_path = RunCache::input_cache_path(&key);
|
||||
fs::write(&input_path, input).unwrap();
|
||||
|
||||
let output_path = RunCache::output_cache_path(&key);
|
||||
log::trace!("Writing cache entry: {}", output_path.display());
|
||||
fs::write(&output_path, output).unwrap();
|
||||
|
||||
self.link_to_run(&key);
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Serialize, Deserialize)]
|
||||
pub struct PredictionDetails {
|
||||
pub diff: String,
|
||||
pub excerpts: Vec<ActualExcerpt>,
|
||||
pub excerpts_text: String, // TODO: contains the worktree root path. Drop this field and compute it on the fly
|
||||
pub retrieval_time: Duration,
|
||||
pub prediction_time: Duration,
|
||||
pub total_time: Duration,
|
||||
pub run_example_dir: PathBuf,
|
||||
pub prompt_len: usize,
|
||||
pub generated_len: usize,
|
||||
}
|
||||
|
||||
impl PredictionDetails {
|
||||
pub fn new(run_example_dir: PathBuf) -> Self {
|
||||
Self {
|
||||
diff: Default::default(),
|
||||
excerpts: Default::default(),
|
||||
excerpts_text: Default::default(),
|
||||
retrieval_time: Default::default(),
|
||||
prediction_time: Default::default(),
|
||||
total_time: Default::default(),
|
||||
run_example_dir,
|
||||
prompt_len: 0,
|
||||
generated_len: 0,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn write(&self, format: PredictionsOutputFormat, mut out: impl Write) -> Result<()> {
|
||||
let formatted = match format {
|
||||
PredictionsOutputFormat::Md => self.to_markdown(),
|
||||
PredictionsOutputFormat::Json => serde_json::to_string_pretty(self)?,
|
||||
PredictionsOutputFormat::Diff => self.diff.clone(),
|
||||
};
|
||||
|
||||
Ok(out.write_all(formatted.as_bytes())?)
|
||||
}
|
||||
|
||||
pub fn to_markdown(&self) -> String {
|
||||
format!(
|
||||
"## Excerpts\n\n\
|
||||
{}\n\n\
|
||||
## Prediction\n\n\
|
||||
{}\n\n\
|
||||
## Time\n\n\
|
||||
Retrieval: {}ms\n\
|
||||
Prediction: {}ms\n\n\
|
||||
Total: {}ms\n",
|
||||
self.excerpts_text,
|
||||
self.diff,
|
||||
self.retrieval_time.as_millis(),
|
||||
self.prediction_time.as_millis(),
|
||||
self.total_time.as_millis(),
|
||||
)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,70 @@
|
||||
use std::{fmt, fmt::Display, path::Path, str::FromStr, sync::Arc};
|
||||
|
||||
use ::util::{paths::PathStyle, rel_path::RelPath};
|
||||
use anyhow::{Result, anyhow};
|
||||
use language::Point;
|
||||
use serde::{Deserialize, Deserializer, Serialize, Serializer};
|
||||
|
||||
#[derive(Debug, Clone, Hash, Eq, PartialEq)]
|
||||
pub struct SourceLocation {
|
||||
pub path: Arc<RelPath>,
|
||||
pub point: Point,
|
||||
}
|
||||
|
||||
impl Serialize for SourceLocation {
|
||||
fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
|
||||
where
|
||||
S: Serializer,
|
||||
{
|
||||
serializer.serialize_str(&self.to_string())
|
||||
}
|
||||
}
|
||||
|
||||
impl<'de> Deserialize<'de> for SourceLocation {
|
||||
fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
|
||||
where
|
||||
D: Deserializer<'de>,
|
||||
{
|
||||
let s = String::deserialize(deserializer)?;
|
||||
s.parse().map_err(serde::de::Error::custom)
|
||||
}
|
||||
}
|
||||
|
||||
impl Display for SourceLocation {
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
write!(
|
||||
f,
|
||||
"{}:{}:{}",
|
||||
self.path.display(PathStyle::Posix),
|
||||
self.point.row + 1,
|
||||
self.point.column + 1
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
impl FromStr for SourceLocation {
|
||||
type Err = anyhow::Error;
|
||||
|
||||
fn from_str(s: &str) -> Result<Self> {
|
||||
let parts: Vec<&str> = s.split(':').collect();
|
||||
if parts.len() != 3 {
|
||||
return Err(anyhow!(
|
||||
"Invalid source location. Expected 'file.rs:line:column', got '{}'",
|
||||
s
|
||||
));
|
||||
}
|
||||
|
||||
let path = RelPath::new(Path::new(&parts[0]), PathStyle::local())?.into_arc();
|
||||
let line: u32 = parts[1]
|
||||
.parse()
|
||||
.map_err(|_| anyhow!("Invalid line number: '{}'", parts[1]))?;
|
||||
let column: u32 = parts[2]
|
||||
.parse()
|
||||
.map_err(|_| anyhow!("Invalid column number: '{}'", parts[2]))?;
|
||||
|
||||
// Convert from 1-based to 0-based indexing
|
||||
let point = Point::new(line.saturating_sub(1), column.saturating_sub(1));
|
||||
|
||||
Ok(SourceLocation { path, point })
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,198 @@
|
||||
use anyhow::{Result, anyhow};
|
||||
use futures::channel::mpsc;
|
||||
use futures::{FutureExt as _, StreamExt as _};
|
||||
use gpui::{AsyncApp, Entity, Task};
|
||||
use language::{Buffer, LanguageId, LanguageNotFound, LanguageServerId, ParseStatus};
|
||||
use project::lsp_store::OpenLspBufferHandle;
|
||||
use project::{Project, ProjectPath, Worktree};
|
||||
use std::collections::HashSet;
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
use util::rel_path::RelPath;
|
||||
|
||||
pub fn open_buffer(
|
||||
project: Entity<Project>,
|
||||
worktree: Entity<Worktree>,
|
||||
path: Arc<RelPath>,
|
||||
cx: &AsyncApp,
|
||||
) -> Task<Result<Entity<Buffer>>> {
|
||||
cx.spawn(async move |cx| {
|
||||
let project_path = worktree.read_with(cx, |worktree, _cx| ProjectPath {
|
||||
worktree_id: worktree.id(),
|
||||
path,
|
||||
})?;
|
||||
|
||||
let buffer = project
|
||||
.update(cx, |project, cx| project.open_buffer(project_path, cx))?
|
||||
.await?;
|
||||
|
||||
let mut parse_status = buffer.read_with(cx, |buffer, _cx| buffer.parse_status())?;
|
||||
while *parse_status.borrow() != ParseStatus::Idle {
|
||||
parse_status.changed().await?;
|
||||
}
|
||||
|
||||
Ok(buffer)
|
||||
})
|
||||
}
|
||||
|
||||
pub async fn open_buffer_with_language_server(
|
||||
project: Entity<Project>,
|
||||
worktree: Entity<Worktree>,
|
||||
path: Arc<RelPath>,
|
||||
ready_languages: &mut HashSet<LanguageId>,
|
||||
cx: &mut AsyncApp,
|
||||
) -> Result<(OpenLspBufferHandle, LanguageServerId, Entity<Buffer>)> {
|
||||
let buffer = open_buffer(project.clone(), worktree, path.clone(), cx).await?;
|
||||
|
||||
let (lsp_open_handle, path_style) = project.update(cx, |project, cx| {
|
||||
(
|
||||
project.register_buffer_with_language_servers(&buffer, cx),
|
||||
project.path_style(cx),
|
||||
)
|
||||
})?;
|
||||
|
||||
let language_registry = project.read_with(cx, |project, _| project.languages().clone())?;
|
||||
let result = language_registry
|
||||
.load_language_for_file_path(path.as_std_path())
|
||||
.await;
|
||||
|
||||
if let Err(error) = result
|
||||
&& !error.is::<LanguageNotFound>()
|
||||
{
|
||||
anyhow::bail!(error);
|
||||
}
|
||||
|
||||
let Some(language_id) = buffer.read_with(cx, |buffer, _cx| {
|
||||
buffer.language().map(|language| language.id())
|
||||
})?
|
||||
else {
|
||||
return Err(anyhow!("No language for {}", path.display(path_style)));
|
||||
};
|
||||
|
||||
let log_prefix = format!("{} | ", path.display(path_style));
|
||||
if !ready_languages.contains(&language_id) {
|
||||
wait_for_lang_server(&project, &buffer, log_prefix, cx).await?;
|
||||
ready_languages.insert(language_id);
|
||||
}
|
||||
|
||||
let lsp_store = project.read_with(cx, |project, _cx| project.lsp_store())?;
|
||||
|
||||
// hacky wait for buffer to be registered with the language server
|
||||
for _ in 0..100 {
|
||||
let Some(language_server_id) = lsp_store.update(cx, |lsp_store, cx| {
|
||||
buffer.update(cx, |buffer, cx| {
|
||||
lsp_store
|
||||
.language_servers_for_local_buffer(&buffer, cx)
|
||||
.next()
|
||||
.map(|(_, language_server)| language_server.server_id())
|
||||
})
|
||||
})?
|
||||
else {
|
||||
cx.background_executor()
|
||||
.timer(Duration::from_millis(10))
|
||||
.await;
|
||||
continue;
|
||||
};
|
||||
|
||||
return Ok((lsp_open_handle, language_server_id, buffer));
|
||||
}
|
||||
|
||||
return Err(anyhow!("No language server found for buffer"));
|
||||
}
|
||||
|
||||
// TODO: Dedupe with similar function in crates/eval/src/instance.rs
|
||||
pub fn wait_for_lang_server(
|
||||
project: &Entity<Project>,
|
||||
buffer: &Entity<Buffer>,
|
||||
log_prefix: String,
|
||||
cx: &mut AsyncApp,
|
||||
) -> Task<Result<()>> {
|
||||
eprintln!("{}⏵ Waiting for language server", log_prefix);
|
||||
|
||||
let (mut tx, mut rx) = mpsc::channel(1);
|
||||
|
||||
let lsp_store = project
|
||||
.read_with(cx, |project, _| project.lsp_store())
|
||||
.unwrap();
|
||||
|
||||
let has_lang_server = buffer
|
||||
.update(cx, |buffer, cx| {
|
||||
lsp_store.update(cx, |lsp_store, cx| {
|
||||
lsp_store
|
||||
.language_servers_for_local_buffer(buffer, cx)
|
||||
.next()
|
||||
.is_some()
|
||||
})
|
||||
})
|
||||
.unwrap_or(false);
|
||||
|
||||
if has_lang_server {
|
||||
project
|
||||
.update(cx, |project, cx| project.save_buffer(buffer.clone(), cx))
|
||||
.unwrap()
|
||||
.detach();
|
||||
}
|
||||
let (mut added_tx, mut added_rx) = mpsc::channel(1);
|
||||
|
||||
let subscriptions = [
|
||||
cx.subscribe(&lsp_store, {
|
||||
let log_prefix = log_prefix.clone();
|
||||
move |_, event, _| {
|
||||
if let project::LspStoreEvent::LanguageServerUpdate {
|
||||
message:
|
||||
client::proto::update_language_server::Variant::WorkProgress(
|
||||
client::proto::LspWorkProgress {
|
||||
message: Some(message),
|
||||
..
|
||||
},
|
||||
),
|
||||
..
|
||||
} = event
|
||||
{
|
||||
eprintln!("{}⟲ {message}", log_prefix)
|
||||
}
|
||||
}
|
||||
}),
|
||||
cx.subscribe(project, {
|
||||
let buffer = buffer.clone();
|
||||
move |project, event, cx| match event {
|
||||
project::Event::LanguageServerAdded(_, _, _) => {
|
||||
let buffer = buffer.clone();
|
||||
project
|
||||
.update(cx, |project, cx| project.save_buffer(buffer, cx))
|
||||
.detach();
|
||||
added_tx.try_send(()).ok();
|
||||
}
|
||||
project::Event::DiskBasedDiagnosticsFinished { .. } => {
|
||||
tx.try_send(()).ok();
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}),
|
||||
];
|
||||
|
||||
cx.spawn(async move |cx| {
|
||||
if !has_lang_server {
|
||||
// some buffers never have a language server, so this aborts quickly in that case.
|
||||
let timeout = cx.background_executor().timer(Duration::from_secs(500));
|
||||
futures::select! {
|
||||
_ = added_rx.next() => {},
|
||||
_ = timeout.fuse() => {
|
||||
anyhow::bail!("Waiting for language server add timed out after 5 seconds");
|
||||
}
|
||||
};
|
||||
}
|
||||
let timeout = cx.background_executor().timer(Duration::from_secs(60 * 5));
|
||||
let result = futures::select! {
|
||||
_ = rx.next() => {
|
||||
eprintln!("{}⚑ Language server idle", log_prefix);
|
||||
anyhow::Ok(())
|
||||
},
|
||||
_ = timeout.fuse() => {
|
||||
anyhow::bail!("LSP wait timed out after 5 minutes");
|
||||
}
|
||||
};
|
||||
drop(subscriptions);
|
||||
result
|
||||
})
|
||||
}
|
||||
Reference in New Issue
Block a user