oak-team/oak-gpui
3
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Rework edit prediction CLI (#44562)

Browse Source 
This PR restructures the commands of the Edit Prediction CLI (now called
`ep`), to support some flows that are important for the training
process:
* generating zeta2 prompt and expected output, without running
predictions
* scoring outputs that are generated by a system other than the
production code (to evaluate the model during training)

To achieve this, we've restructured the CLI commands so that they all
take as input, and produce as output, a consistent, uniform data format:
a set of one or more `Example` structs, expressible either as the
original markdown format, or as a JSON lines. The `Example` struct
starts with the basic fields that are in human-readable eval format, but
contain a number of optional fields that are filled in by different
steps in the processing pipeline (`context`, `predict`, `format-prompt`,
and `score`).

### To do

* [x] Adjust the teacher model output parsing to use the full buffer
contents
* [x] Move udiff to cli
* [x] Align `format-prompt` with Zeta2's production code
* [x] Change score output to assume same provider
* [x] Move pretty reporting to `eval` command
* [x] Store cursor point in addition to cursor offset
* [x] Rename `edit_prediction_cli2` -> `edit_prediction_cli` (nuke the
old one)

Release Notes:

- N/A

---------

Co-authored-by: Oleksiy Syvokon <oleksiy@zed.dev>
Co-authored-by: Agus Zubiaga <agus@zed.dev>
Co-authored-by: Ben Kunkle <ben@zed.dev>
This commit is contained in:
Max Brunsfeld
2025-12-10 17:36:51 -08:00
committed by GitHub
co-authored by Oleksiy Syvokon Agus Zubiaga Ben Kunkle
parent 37077a8ebb
commit 25d74480aa
44 changed files with 2502 additions and 4801 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/edit_prediction_cli/src/training/llm_client.rs → crates/edit_prediction_cli/src/anthropic_client.rs
+26 -25
View File
@@ -5,11 +5,13 @@ use anthropic::{
use anyhow::Result;
use http_client::HttpClient;
use indoc::indoc;
use reqwest_client::ReqwestClient;
use sqlez::bindable::Bind;
use sqlez::bindable::StaticColumnCount;
use sqlez_macros::sql;
use std::hash::Hash;
use std::hash::Hasher;
use std::path::Path;
use std::sync::Arc;
pub struct PlainLlmClient {
@@ -18,7 +20,8 @@ pub struct PlainLlmClient {
}
impl PlainLlmClient {
fn new(http_client: Arc<dyn HttpClient>) -> Result<Self> {
fn new() -> Result<Self> {
let http_client: Arc<dyn http_client::HttpClient> = Arc::new(ReqwestClient::new());
let api_key = std::env::var("ANTHROPIC_API_KEY")
.map_err(|_| anyhow::anyhow!("ANTHROPIC_API_KEY environment variable not set"))?;
Ok(Self {
@@ -29,12 +32,12 @@ impl PlainLlmClient {
async fn generate(
&self,
model: String,
model: &str,
max_tokens: u64,
messages: Vec<Message>,
) -> Result<AnthropicResponse> {
let request = AnthropicRequest {
model,
model: model.to_string(),
max_tokens,
messages,
tools: Vec::new(),
@@ -105,11 +108,12 @@ struct SerializableMessage {
}
impl BatchingLlmClient {
fn new(cache_path: &str, http_client: Arc<dyn HttpClient>) -> Result<Self> {
fn new(cache_path: &Path) -> Result<Self> {
let http_client: Arc<dyn http_client::HttpClient> = Arc::new(ReqwestClient::new());
let api_key = std::env::var("ANTHROPIC_API_KEY")
.map_err(|_| anyhow::anyhow!("ANTHROPIC_API_KEY environment variable not set"))?;
let connection = sqlez::connection::Connection::open_file(&cache_path);
let connection = sqlez::connection::Connection::open_file(&cache_path.to_str().unwrap());
let mut statement = sqlez::statement::Statement::prepare(
&connection,
indoc! {"
@@ -182,16 +186,16 @@ impl BatchingLlmClient {
async fn generate(
&self,
model: String,
model: &str,
max_tokens: u64,
messages: Vec<Message>,
) -> Result<Option<AnthropicResponse>> {
let response = self.lookup(&model, max_tokens, &messages)?;
let response = self.lookup(model, max_tokens, &messages)?;
if let Some(response) = response {
return Ok(Some(response));
}
self.mark_for_batch(&model, max_tokens, &messages)?;
self.mark_for_batch(model, max_tokens, &messages)?;
Ok(None)
}
@@ -258,7 +262,7 @@ impl BatchingLlmClient {
}
}
}
log::info!("Uploaded {} successful requests", success_count);
log::info!("Downloaded {} successful requests", success_count);
}
}
@@ -363,23 +367,20 @@ fn message_content_to_string(content: &[RequestContent]) -> String {
.join("\n")
}
pub enum LlmClient {
pub enum AnthropicClient {
// No batching
Plain(PlainLlmClient),
Batch(BatchingLlmClient),
Dummy,
}
impl LlmClient {
pub fn plain(http_client: Arc<dyn HttpClient>) -> Result<Self> {
Ok(Self::Plain(PlainLlmClient::new(http_client)?))
impl AnthropicClient {
pub fn plain() -> Result<Self> {
Ok(Self::Plain(PlainLlmClient::new()?))
}
pub fn batch(cache_path: &str, http_client: Arc<dyn HttpClient>) -> Result<Self> {
Ok(Self::Batch(BatchingLlmClient::new(
cache_path,
http_client,
)?))
pub fn batch(cache_path: &Path) -> Result<Self> {
Ok(Self::Batch(BatchingLlmClient::new(cache_path)?))
}
#[allow(dead_code)]
@@ -389,29 +390,29 @@ impl LlmClient {
pub async fn generate(
&self,
model: String,
model: &str,
max_tokens: u64,
messages: Vec<Message>,
) -> Result<Option<AnthropicResponse>> {
match self {
LlmClient::Plain(plain_llm_client) => plain_llm_client
AnthropicClient::Plain(plain_llm_client) => plain_llm_client
.generate(model, max_tokens, messages)
.await
.map(Some),
LlmClient::Batch(batching_llm_client) => {
AnthropicClient::Batch(batching_llm_client) => {
batching_llm_client
.generate(model, max_tokens, messages)
.await
}
LlmClient::Dummy => panic!("Dummy LLM client is not expected to be used"),
AnthropicClient::Dummy => panic!("Dummy LLM client is not expected to be used"),
}
}
pub async fn sync_batches(&self) -> Result<()> {
match self {
LlmClient::Plain(_) => Ok(()),
LlmClient::Batch(batching_llm_client) => batching_llm_client.sync_batches().await,
LlmClient::Dummy => panic!("Dummy LLM client is not expected to be used"),
AnthropicClient::Plain(_) => Ok(()),
AnthropicClient::Batch(batching_llm_client) => batching_llm_client.sync_batches().await,
AnthropicClient::Dummy => panic!("Dummy LLM client is not expected to be used"),
}
}
}
crates/edit_prediction_cli/src/evaluate.rs
-641
View File
@@ -1,641 +0,0 @@
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()
)
}
crates/edit_prediction_cli/src/example.rs
+296 -493
View File
@@ -1,59 +1,103 @@
use crate::{
PredictionProvider, PromptFormat,
metrics::ClassificationMetrics,
paths::{REPOS_DIR, WORKTREES_DIR},
};
use anyhow::{Context as _, Result};
use edit_prediction::udiff::OpenedBuffers;
use gpui::Entity;
use http_client::Url;
use language::{Anchor, Buffer};
use project::Project;
use serde::{Deserialize, Serialize};
use std::sync::Arc;
use std::{
borrow::Cow,
cell::RefCell,
fmt::{self, Display},
fs,
hash::Hash,
hash::Hasher,
io::Write,
io::{Read, Write},
mem,
path::{Path, PathBuf},
sync::{Arc, OnceLock},
};
use zeta_prompt::RelatedFile;
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, Hash, Serialize, Deserialize)]
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct Example {
#[serde(default)]
pub name: String,
pub repository_url: String,
pub revision: String,
pub uncommitted_diff: String,
pub cursor_path: PathBuf,
pub cursor_path: Arc<Path>,
pub cursor_position: String,
pub edit_history: String,
pub expected_patch: String,
/// The full content of the file where an edit is being predicted, and the
/// actual cursor offset.
#[serde(skip_serializing_if = "Option::is_none")]
pub buffer: Option<ExampleBuffer>,
/// The context retrieved for the prediction. This requires the worktree to
/// be loaded and the language server to be started.
#[serde(skip_serializing_if = "Option::is_none")]
pub context: Option<ExampleContext>,
/// The input and expected output from the edit prediction model.
#[serde(skip_serializing_if = "Option::is_none")]
pub prompt: Option<ExamplePrompt>,
/// The actual predictions from the model.
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub predictions: Vec<ExamplePrediction>,
/// The scores, for how well the actual predictions match the expected
/// predictions.
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub score: Vec<ExampleScore>,
/// The application state used to process this example.
#[serde(skip)]
pub state: Option<ExampleState>,
}
#[derive(Clone, Debug)]
pub struct ExampleState {
pub project: Entity<Project>,
pub buffer: Entity<Buffer>,
pub cursor_position: Anchor,
pub _open_buffers: OpenedBuffers,
}
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct ExampleContext {
pub files: Arc<[RelatedFile]>,
}
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct ExampleBuffer {
pub content: String,
pub cursor_row: u32,
pub cursor_column: u32,
pub cursor_offset: usize,
}
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct ExamplePrompt {
pub input: String,
pub expected_output: String,
pub format: PromptFormat,
}
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct ExamplePrediction {
pub actual_patch: String,
pub actual_output: String,
pub provider: PredictionProvider,
}
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct ExampleScore {
pub delta_chr_f: f32,
pub line_match: ClassificationMetrics,
}
impl Example {
@@ -90,485 +134,244 @@ impl Example {
}
}
pub async fn setup_worktree(&self, file_name: String) -> Result<PathBuf> {
let (repo_owner, repo_name) = self.repo_name()?;
pub fn worktree_path(&self) -> PathBuf {
WORKTREES_DIR
.join(&self.name)
.join(self.repo_name().unwrap().1.as_ref())
}
let repo_dir = REPOS_DIR.join(repo_owner.as_ref()).join(repo_name.as_ref());
let repo_lock = lock_repo(&repo_dir).await;
pub fn repo_path(&self) -> PathBuf {
let (repo_owner, repo_name) = self.repo_name().expect("failed to get repo name");
REPOS_DIR.join(repo_owner.as_ref()).join(repo_name.as_ref())
}
}
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.repository_url],
)
.await?;
}
pub fn read_examples(inputs: &[PathBuf]) -> Vec<Example> {
let mut examples = Vec::new();
// Resolve the example to a revision, fetching it if needed.
let revision = run_git(
&repo_dir,
&["rev-parse", &format!("{}^{{commit}}", self.revision)],
)
.await;
let revision = if let Ok(revision) = revision {
revision
let stdin_path: PathBuf = PathBuf::from("-");
let inputs = if inputs.is_empty() {
&[stdin_path]
} else {
inputs
};
for path in inputs {
let is_stdin = path.as_path() == Path::new("-");
let content = if is_stdin {
let mut buffer = String::new();
std::io::stdin()
.read_to_string(&mut buffer)
.expect("Failed to read from stdin");
buffer
} else {
if run_git(
&repo_dir,
&["fetch", "--depth", "1", "origin", &self.revision],
)
.await
.is_err()
{
run_git(&repo_dir, &["fetch", "origin"]).await?;
}
let revision = run_git(&repo_dir, &["rev-parse", "FETCH_HEAD"]).await?;
if revision != self.revision {
run_git(&repo_dir, &["tag", &self.revision, &revision]).await?;
}
revision
std::fs::read_to_string(path)
.unwrap_or_else(|_| panic!("Failed to read path: {:?}", &path))
};
let filename = path.file_stem().unwrap().to_string_lossy().to_string();
let ext = if !is_stdin {
path.extension()
.map(|ext| ext.to_string_lossy().to_string())
.unwrap_or_else(|| panic!("{} should have an extension", path.display()))
} else {
"jsonl".to_string()
};
// 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.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.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),
);
match ext.as_ref() {
"json" => {
let mut example =
serde_json::from_str::<Example>(&content).unwrap_or_else(|error| {
panic!("Failed to parse example file: {}\n{error}", path.display())
});
if example.name.is_empty() {
example.name = filename;
}
examples.push(example);
}
"jsonl" => examples.extend(
content
.lines()
.enumerate()
.map(|(line_ix, line)| {
let mut example =
serde_json::from_str::<Example>(line).unwrap_or_else(|_| {
panic!(
"Failed to parse example on {}:{}",
path.display(),
line_ix + 1
)
});
if example.name.is_empty() {
example.name = format!("{filename}-{line_ix}")
}
example
})
.collect::<Vec<Example>>(),
),
"md" => {
examples.push(parse_markdown_example(filename, &content).unwrap());
}
ext => {
panic!("{} has invalid example extension `{ext}`", path.display())
}
}
Ok(worktree_path)
}
examples
}
pub fn unique_name(&self) -> String {
let mut hasher = std::hash::DefaultHasher::new();
self.hash(&mut hasher);
let disambiguator = hasher.finish();
let hash = format!("{:04x}", disambiguator);
format!("{}_{}", &self.revision[..8], &hash[..4])
pub fn write_examples(examples: &[Example], output_path: Option<&PathBuf>) {
let mut content = String::new();
for example in examples {
let line = serde_json::to_string(example).unwrap();
content.push_str(&line);
content.push('\n');
}
if let Some(output_path) = output_path {
std::fs::write(output_path, content).expect("Failed to write examples");
} else {
std::io::stdout().write_all(&content.as_bytes()).unwrap();
}
}
pub type ActualExcerpt = Excerpt;
fn parse_markdown_example(id: String, input: &str) -> Result<Example> {
use pulldown_cmark::{CodeBlockKind, CowStr, Event, HeadingLevel, Parser, Tag, TagEnd};
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct Excerpt {
pub path: PathBuf,
pub text: String,
}
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(ValueEnum, Debug, Clone)]
pub enum ExampleFormat {
Json,
Toml,
Md,
}
let parser = Parser::new(input);
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();
let mut example = Example {
name: id,
repository_url: String::new(),
revision: String::new(),
uncommitted_diff: String::new(),
cursor_path: PathBuf::new().into(),
cursor_position: String::new(),
edit_history: String::new(),
expected_patch: String::new(),
buffer: None,
context: None,
prompt: None,
predictions: Vec::new(),
score: Vec::new(),
state: None,
};
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."
);
}
}
let mut name = String::new();
let mut text = String::new();
let mut block_info: CowStr = "".into();
#[derive(PartialEq)]
enum Section {
UncommittedDiff,
EditHistory,
CursorPosition,
ExpectedExcerpts,
ExpectedPatch,
Other,
}
pub fn parse_md(input: &str) -> Result<Self> {
use pulldown_cmark::{CodeBlockKind, Event, HeadingLevel, Parser, Tag, TagEnd};
let mut current_section = Section::Other;
let parser = Parser::new(input);
for event in parser {
match event {
Event::Text(line) => {
text.push_str(&line);
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();
}
_ => {}
if let Some((field, value)) = line.split_once('=') {
match field.trim() {
REPOSITORY_URL_FIELD => {
example.repository_url = value.trim().to_string();
}
REVISION_FIELD => {
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())?)
Event::End(TagEnd::Heading(HeadingLevel::H1)) => {
if !name.is_empty() {
anyhow::bail!(
"Found multiple H1 headings. There should only be one with the name of the example."
);
}
name = mem::take(&mut text);
}
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> {
self.example.setup_worktree(self.file_name()).await
}
pub fn file_name(&self) -> String {
self.name
.chars()
.map(|c| {
if c.is_whitespace() {
'-'
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 {
c.to_ascii_lowercase()
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 => {
example.uncommitted_diff = mem::take(&mut text);
}
Section::EditHistory => {
example.edit_history.push_str(&mem::take(&mut text));
}
Section::CursorPosition => {
example.cursor_path = Path::new(block_info).into();
example.cursor_position = mem::take(&mut text);
}
Section::ExpectedExcerpts => {
mem::take(&mut text);
}
Section::ExpectedPatch => {
example.expected_patch = mem::take(&mut text);
}
Section::Other => {}
}
})
.collect()
}
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(())
}
}
if example.cursor_path.as_ref() == Path::new("") || example.cursor_position.is_empty() {
anyhow::bail!("Missing cursor position codeblock");
}
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
Ok(example)
}
crates/edit_prediction_cli/src/format_prompt.rs
+280
View File
@@ -0,0 +1,280 @@
use crate::{
PromptFormat,
example::{Example, ExamplePrompt},
headless::EpAppState,
retrieve_context::run_context_retrieval,
};
use edit_prediction::{EditPredictionStore, zeta2::zeta2_prompt_input};
use gpui::AsyncApp;
use std::sync::Arc;
use zeta_prompt::format_zeta_prompt;
pub async fn run_format_prompt(
example: &mut Example,
prompt_format: PromptFormat,
app_state: Arc<EpAppState>,
mut cx: AsyncApp,
) {
run_context_retrieval(example, app_state, cx.clone()).await;
let prompt = match prompt_format {
PromptFormat::Teacher => TeacherPrompt::format(example),
PromptFormat::Zeta2 => {
let ep_store = cx
.update(|cx| EditPredictionStore::try_global(cx).unwrap())
.unwrap();
let state = example.state.as_ref().unwrap();
let snapshot = state
.buffer
.read_with(&cx, |buffer, _| buffer.snapshot())
.unwrap();
let project = state.project.clone();
let (_, input) = ep_store
.update(&mut cx, |ep_store, _cx| {
zeta2_prompt_input(
&snapshot,
example.context.as_ref().unwrap().files.clone(),
ep_store.edit_history_for_project(&project),
example.cursor_path.clone(),
example.buffer.as_ref().unwrap().cursor_offset,
)
})
.unwrap();
format_zeta_prompt(&input)
}
};
example.prompt = Some(ExamplePrompt {
input: prompt,
expected_output: example.expected_patch.clone(), // TODO
format: prompt_format,
});
}
pub trait PromptFormatter {
fn format(example: &Example) -> String;
}
pub trait PromptParser {
/// Return unified diff patch of prediction given raw LLM response
fn parse(example: &Example, response: &str) -> String;
}
pub struct TeacherPrompt;
impl PromptFormatter for TeacherPrompt {
fn format(example: &Example) -> String {
let edit_history = Self::format_edit_history(&example.edit_history);
let context = Self::format_context(example);
let editable_region = Self::format_editable_region(example);
let prompt = Self::PROMPT
.replace("{{context}}", &context)
.replace("{{edit_history}}", &edit_history)
.replace("{{editable_region}}", &editable_region);
prompt
}
}
impl TeacherPrompt {
const PROMPT: &str = include_str!("teacher.prompt.md");
pub(crate) const EDITABLE_REGION_START: &str = "<|editable_region_start|>\n";
pub(crate) const EDITABLE_REGION_END: &str = "<|editable_region_end|>";
/// Truncate edit history to this number of last lines
const MAX_HISTORY_LINES: usize = 128;
fn format_edit_history(edit_history: &str) -> String {
// Strip comments ("garbage lines") from edit history
let lines = edit_history
.lines()
.filter(|&s| Self::is_udiff_content_line(s))
.collect::<Vec<_>>();
let history_lines = if lines.len() > Self::MAX_HISTORY_LINES {
&lines[lines.len() - Self::MAX_HISTORY_LINES..]
} else {
&lines
};
if history_lines.is_empty() {
return "(No edit history)".to_string();
}
history_lines.join("\n")
}
fn format_context(example: &Example) -> String {
if example.context.is_none() {
panic!("Missing context retriever step");
}
let mut prompt = String::new();
zeta_prompt::write_related_files(&mut prompt, &example.context.as_ref().unwrap().files);
prompt
}
fn format_editable_region(example: &Example) -> String {
let mut result = String::new();
let path_str = example.cursor_path.to_string_lossy();
result.push_str(&format!("`````path=\"{path_str}\"\n"));
result.push_str(Self::EDITABLE_REGION_START);
// TODO: control number of lines around cursor
result.push_str(&example.cursor_position);
if !example.cursor_position.ends_with('\n') {
result.push('\n');
}
result.push_str(&format!("{}\n", Self::EDITABLE_REGION_END));
result.push_str("`````");
result
}
fn extract_editable_region(text: &str) -> String {
let start = text
.find(Self::EDITABLE_REGION_START)
.map_or(0, |pos| pos + Self::EDITABLE_REGION_START.len());
let end = text.find(Self::EDITABLE_REGION_END).unwrap_or(text.len());
let region = &text[start..end];
region.replace("<|user_cursor|>", "")
}
fn is_udiff_content_line(s: &str) -> bool {
s.starts_with("-")
|| s.starts_with("+")
|| s.starts_with(" ")
|| s.starts_with("---")
|| s.starts_with("+++")
|| s.starts_with("@@")
}
}
impl PromptParser for TeacherPrompt {
fn parse(example: &Example, response: &str) -> String {
// Ideally, we should always be able to find cursor position in the retrieved context.
// In reality, sometimes we don't find it for these reasons:
// 1. `example.cursor_position` contains _more_ context than included in the retrieved context
// (can be fixed by getting cursor coordinates at the load_example stage)
// 2. Context retriever just didn't include cursor line.
//
// In that case, fallback to using `cursor_position` as excerpt.
let cursor_file = &example
.buffer
.as_ref()
.expect("`buffer` should be filled in in the context collection step")
.content;
// Extract updated (new) editable region from the model response
let new_editable_region = extract_last_codeblock(response);
// Reconstruct old editable region we sent to the model
let old_editable_region = Self::format_editable_region(example);
let old_editable_region = Self::extract_editable_region(&old_editable_region);
if !cursor_file.contains(&old_editable_region) {
panic!("Something's wrong: editable_region is not found in the cursor file")
}
// Apply editable region to a larger context and compute diff.
// This is needed to get a better context lines around the editable region
let edited_file = cursor_file.replace(&old_editable_region, &new_editable_region);
let diff = language::unified_diff(&cursor_file, &edited_file);
let diff = indoc::formatdoc! {"
--- a/{path}
+++ b/{path}
{diff}
",
path = example.cursor_path.to_string_lossy(),
diff = diff,
};
diff
}
}
fn extract_last_codeblock(text: &str) -> String {
let mut last_block = None;
let mut search_start = 0;
while let Some(start) = text[search_start..].find("```") {
let start = start + search_start;
let bytes = text.as_bytes();
let mut backtick_end = start;
while backtick_end < bytes.len() && bytes[backtick_end] == b'`' {
backtick_end += 1;
}
let backtick_count = backtick_end - start;
let closing_backticks = "`".repeat(backtick_count);
while backtick_end < bytes.len() && bytes[backtick_end] != b'\n' {
backtick_end += 1;
}
if let Some(end_pos) = text[backtick_end..].find(&closing_backticks) {
let code_block = &text[backtick_end + 1..backtick_end + end_pos - 1];
last_block = Some(code_block.to_string());
search_start = backtick_end + end_pos + backtick_count;
} else {
break;
}
}
last_block.unwrap_or_else(|| text.to_string())
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_extract_last_code_block() {
let text = indoc::indoc! {"
Some thinking
```
first block
```
`````path='something' lines=1:2
last block
`````
"};
let last_block = extract_last_codeblock(text);
assert_eq!(last_block, "last block");
}
#[test]
fn test_extract_editable_region() {
let text = indoc::indoc! {"
some lines
are
here
<|editable_region_start|>
one
two three
<|editable_region_end|>
more
lines here
"};
let parsed = TeacherPrompt::extract_editable_region(text);
assert_eq!(
parsed,
indoc::indoc! {"
one
two three
"}
);
}
}
crates/edit_prediction_cli/src/headless.rs
+3 -3
View File
@@ -16,7 +16,7 @@ use std::sync::Arc;
use util::ResultExt as _;
/// Headless subset of `workspace::AppState`.
pub struct ZetaCliAppState {
pub struct EpAppState {
pub languages: Arc<LanguageRegistry>,
pub client: Arc<Client>,
pub user_store: Entity<UserStore>,
@@ -25,7 +25,7 @@ pub struct ZetaCliAppState {
}
// TODO: dedupe with crates/eval/src/eval.rs
pub fn init(cx: &mut App) -> ZetaCliAppState {
pub fn init(cx: &mut App) -> EpAppState {
let app_commit_sha = option_env!("ZED_COMMIT_SHA").map(|s| AppCommitSha::new(s.to_owned()));
let app_version = AppVersion::load(
@@ -112,7 +112,7 @@ pub fn init(cx: &mut App) -> ZetaCliAppState {
prompt_store::init(cx);
terminal_view::init(cx);
ZetaCliAppState {
EpAppState {
languages,
client,
user_store,
crates/edit_prediction_cli/src/load_project.rs
+320
View File
@@ -0,0 +1,320 @@
use crate::{
example::{Example, ExampleBuffer, ExampleState},
headless::EpAppState,
};
use anyhow::{Result, anyhow};
use collections::HashMap;
use edit_prediction::EditPredictionStore;
use edit_prediction::udiff::OpenedBuffers;
use futures::{
AsyncWriteExt as _,
lock::{Mutex, OwnedMutexGuard},
};
use gpui::{AsyncApp, Entity};
use language::{Anchor, Buffer, ToOffset, ToPoint};
use project::buffer_store::BufferStoreEvent;
use project::{Project, ProjectPath};
use std::{
cell::RefCell,
fs,
path::{Path, PathBuf},
sync::Arc,
};
use util::{paths::PathStyle, rel_path::RelPath};
use zeta_prompt::CURSOR_MARKER;
pub async fn run_load_project(example: &mut Example, app_state: Arc<EpAppState>, mut cx: AsyncApp) {
if example.state.is_some() {
return;
}
let project = setup_project(example, &app_state, &mut cx).await;
let buffer_store = project
.read_with(&cx, |project, _| project.buffer_store().clone())
.unwrap();
let ep_store = cx
.update(|cx| EditPredictionStore::try_global(cx).unwrap())
.unwrap();
cx.subscribe(&buffer_store, {
let project = project.clone();
move |_, event, cx| match event {
BufferStoreEvent::BufferAdded(buffer) => {
ep_store.update(cx, |store, cx| store.register_buffer(&buffer, &project, cx));
}
_ => {}
}
})
.unwrap()
.detach();
let _open_buffers = apply_edit_history(example, &project, &mut cx)
.await
.unwrap();
let (buffer, cursor_position) = cursor_position(example, &project, &mut cx).await;
example.buffer = buffer
.read_with(&cx, |buffer, _cx| {
let cursor_point = cursor_position.to_point(&buffer);
Some(ExampleBuffer {
content: buffer.text(),
cursor_row: cursor_point.row,
cursor_column: cursor_point.column,
cursor_offset: cursor_position.to_offset(&buffer),
})
})
.unwrap();
example.state = Some(ExampleState {
buffer,
project,
cursor_position,
_open_buffers,
});
}
async fn cursor_position(
example: &Example,
project: &Entity<Project>,
cx: &mut AsyncApp,
) -> (Entity<Buffer>, Anchor) {
let worktree = project
.read_with(cx, |project, cx| {
project.visible_worktrees(cx).next().unwrap()
})
.unwrap();
let cursor_path = RelPath::new(&example.cursor_path, PathStyle::Posix)
.unwrap()
.into_arc();
let cursor_buffer = project
.update(cx, |project, cx| {
project.open_buffer(
ProjectPath {
worktree_id: worktree.read(cx).id(),
path: cursor_path,
},
cx,
)
})
.unwrap()
.await
.unwrap();
let cursor_offset_within_excerpt = example
.cursor_position
.find(CURSOR_MARKER)
.ok_or_else(|| anyhow!("missing cursor marker"))
.unwrap();
let mut cursor_excerpt = 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 (excerpt_offset, _) = matches.next().unwrap_or_else(|| {
panic!(
"\nExcerpt:\n\n{cursor_excerpt}\nBuffer text:\n{text}\n.Cursor excerpt did not exist in buffer."
);
});
assert!(matches.next().is_none(), "More than one cursor position match found for {}", &example.name);
excerpt_offset
}).unwrap();
let cursor_offset = excerpt_offset + cursor_offset_within_excerpt;
let cursor_anchor = cursor_buffer
.read_with(cx, |buffer, _| buffer.anchor_after(cursor_offset))
.unwrap();
(cursor_buffer, cursor_anchor)
}
async fn setup_project(
example: &mut Example,
app_state: &Arc<EpAppState>,
cx: &mut AsyncApp,
) -> Entity<Project> {
setup_worktree(example).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,
)
})
.unwrap();
let worktree = project
.update(cx, |project, cx| {
project.create_worktree(&example.worktree_path(), true, cx)
})
.unwrap()
.await
.unwrap();
worktree
.read_with(cx, |worktree, _cx| {
worktree.as_local().unwrap().scan_complete()
})
.unwrap()
.await;
project
}
pub async fn setup_worktree(example: &Example) {
let repo_dir = example.repo_path();
let repo_lock = lock_repo(&repo_dir).await;
if !repo_dir.is_dir() {
fs::create_dir_all(&repo_dir).unwrap();
run_git(&repo_dir, &["init"]).await.unwrap();
run_git(
&repo_dir,
&["remote", "add", "origin", &example.repository_url],
)
.await
.unwrap();
}
// Resolve the example to a revision, fetching it if needed.
let revision = run_git(
&repo_dir,
&["rev-parse", &format!("{}^{{commit}}", example.revision)],
)
.await;
let revision = if let Ok(revision) = revision {
revision
} else {
if run_git(
&repo_dir,
&["fetch", "--depth", "1", "origin", &example.revision],
)
.await
.is_err()
{
run_git(&repo_dir, &["fetch", "origin"]).await.unwrap();
}
let revision = run_git(&repo_dir, &["rev-parse", "FETCH_HEAD"])
.await
.unwrap();
if revision != example.revision {
run_git(&repo_dir, &["tag", &example.revision, &revision])
.await
.unwrap();
}
revision
};
// Create the worktree for this example if needed.
let worktree_path = example.worktree_path();
if worktree_path.is_dir() {
run_git(&worktree_path, &["clean", "--force", "-d"])
.await
.unwrap();
run_git(&worktree_path, &["reset", "--hard", "HEAD"])
.await
.unwrap();
run_git(&worktree_path, &["checkout", revision.as_str()])
.await
.unwrap();
} else {
let worktree_path_string = worktree_path.to_string_lossy();
run_git(
&repo_dir,
&["branch", "-f", &example.name, revision.as_str()],
)
.await
.unwrap();
run_git(
&repo_dir,
&[
"worktree",
"add",
"-f",
&worktree_path_string,
&example.name,
],
)
.await
.unwrap();
}
drop(repo_lock);
// Apply the uncommitted diff for this example.
if !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()
.unwrap();
let mut stdin = apply_process.stdin.take().unwrap();
stdin
.write_all(example.uncommitted_diff.as_bytes())
.await
.unwrap();
stdin.close().await.unwrap();
drop(stdin);
let apply_result = apply_process.output().await.unwrap();
if !apply_result.status.success() {
panic!(
"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),
);
}
}
}
async fn apply_edit_history(
example: &Example,
project: &Entity<Project>,
cx: &mut AsyncApp,
) -> Result<OpenedBuffers> {
edit_prediction::udiff::apply_diff(&example.edit_history, project, cx).await
}
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
}
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())
}
crates/edit_prediction_cli/src/main.rs
+140 -466
View File
@@ -1,522 +1,196 @@
mod evaluate;
mod anthropic_client;
mod example;
mod format_prompt;
mod headless;
mod load_project;
mod metrics;
mod paths;
mod predict;
mod source_location;
mod training;
mod util;
mod retrieve_context;
mod score;
use crate::{
evaluate::run_evaluate,
example::{ExampleFormat, NamedExample},
headless::ZetaCliAppState,
predict::run_predict,
source_location::SourceLocation,
training::{context::ContextType, distill::run_distill},
util::{open_buffer, open_buffer_with_language_server},
};
use ::util::{ResultExt, 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 clap::{Args, CommandFactory, Parser, Subcommand, ValueEnum};
use edit_prediction::EditPredictionStore;
use gpui::Application;
use reqwest_client::ReqwestClient;
use std::io::{self};
use std::{collections::HashSet, path::PathBuf, str::FromStr, sync::Arc};
use serde::{Deserialize, Serialize};
use std::{path::PathBuf, sync::Arc};
use crate::example::{read_examples, write_examples};
use crate::format_prompt::run_format_prompt;
use crate::load_project::run_load_project;
use crate::predict::run_prediction;
use crate::retrieve_context::run_context_retrieval;
use crate::score::run_scoring;
#[derive(Parser, Debug)]
#[command(name = "zeta")]
struct ZetaCliArgs {
#[command(name = "ep")]
struct EpArgs {
#[arg(long, default_value_t = false)]
printenv: bool,
#[clap(long, default_value_t = 10)]
max_parallelism: usize,
#[command(subcommand)]
command: Option<Command>,
#[clap(global = true)]
inputs: Vec<PathBuf>,
#[arg(long, short, global = true)]
output: Option<PathBuf>,
#[arg(long, short, global = true)]
in_place: bool,
}
#[derive(Subcommand, Debug)]
enum Command {
Context(ContextArgs),
Predict(PredictArguments),
Eval(EvaluateArguments),
Distill(DistillArguments),
ConvertExample {
path: PathBuf,
#[arg(long, value_enum, default_value_t = ExampleFormat::Md)]
output_format: ExampleFormat,
},
Score {
golden_patch: PathBuf,
actual_patch: PathBuf,
},
/// Parse markdown examples and output a combined .jsonl file
ParseExample,
/// Create git worktrees for each example and load file contents
LoadBuffer,
/// Retrieve context for input examples.
Context,
/// Generate a prompt string for a specific model
FormatPrompt(FormatPromptArgs),
/// Runs edit prediction
Predict(PredictArgs),
/// Computes a score based on actual and expected patches
Score(PredictArgs),
/// Print aggregated scores
Eval(PredictArgs),
/// Remove git repositories and worktrees
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,
struct FormatPromptArgs {
#[clap(long)]
prompt_format: PromptFormat,
}
#[derive(clap::ValueEnum, Default, Debug, Clone, Copy)]
enum ContextProvider {
Zeta1,
#[default]
#[derive(Clone, Copy, Debug, ValueEnum, Serialize, Deserialize)]
enum PromptFormat {
Teacher,
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(Debug, Args)]
pub struct DistillArguments {
split_commit_dataset: PathBuf,
#[clap(long, value_enum, default_value_t = ContextType::CurrentFile)]
context_type: ContextType,
#[clap(long)]
batch: Option<String>,
}
#[derive(Clone, Debug, Args)]
pub struct PredictionOptions {
#[clap(flatten)]
zeta2: Zeta2Args,
struct PredictArgs {
#[clap(long)]
provider: PredictionProvider,
#[clap(long, value_enum, default_value_t = CacheMode::default())]
cache: CacheMode,
#[clap(long, default_value_t = 1)]
repetitions: usize,
}
#[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)]
#[derive(Clone, Copy, Debug, ValueEnum, Serialize, Deserialize)]
enum PredictionProvider {
Zeta1,
#[default]
Zeta2,
Sweep,
Mercury,
Zeta1,
Zeta2,
Teacher,
}
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)
));
impl EpArgs {
fn output_path(&self) -> Option<PathBuf> {
if self.in_place {
if self.inputs.len() == 1 {
self.inputs.first().cloned()
} else {
panic!("--in-place requires exactly one input file")
}
} else {
return Err(anyhow!(
"Cursor position {:?} is out of bounds (max row is {})",
cursor.point,
max_row
));
self.output.clone()
}
}
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 args = EpArgs::parse();
if args.printenv {
::util::shell_env::print_env();
return;
}
let output = args.output_path();
let command = match args.command {
Some(cmd) => cmd,
None => {
EpArgs::command().print_help().unwrap();
return;
}
};
match &command {
Command::Clean => {
std::fs::remove_dir_all(&*paths::DATA_DIR).unwrap();
return;
}
_ => {}
}
let mut examples = read_examples(&args.inputs);
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));
EditPredictionStore::global(&app_state.client, &app_state.user_store, 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()
match &command {
Command::Predict(args) => predict::sync_batches(&args.provider).await,
_ => (),
};
for data in examples.chunks_mut(args.max_parallelism) {
let mut futures = Vec::new();
for example in data.iter_mut() {
let cx = cx.clone();
let app_state = app_state.clone();
futures.push(async {
match &command {
Command::ParseExample => {}
Command::LoadBuffer => {
run_load_project(example, app_state.clone(), cx).await;
}
Command::Context => {
run_context_retrieval(example, app_state, cx).await;
}
Command::FormatPrompt(args) => {
run_format_prompt(example, args.prompt_format, app_state, cx).await;
}
Command::Predict(args) => {
run_prediction(
example,
Some(args.provider),
args.repetitions,
app_state.clone(),
cx,
)
.await;
}
Command::Score(args) | Command::Eval(args) => {
run_scoring(example, &args, app_state, cx).await;
}
Command::Clean => {
unreachable!()
}
}
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::Distill(arguments)) => {
let _guard = cx
.update(|cx| gpui_tokio::Tokio::handle(cx))
.unwrap()
.enter();
run_distill(arguments).await.log_err();
}
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();
futures::future::join_all(futures).await;
}
let golden_diff: Vec<DiffLine> = golden_content
.lines()
.map(|line| DiffLine::parse(line))
.collect();
if args.output.is_some() || !matches!(command, Command::Eval(_)) {
write_examples(&examples, output.as_ref());
}
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()
}
match &command {
Command::Predict(args) => predict::sync_batches(&args.provider).await,
Command::Eval(_) => score::print_report(&examples),
_ => (),
};
let _ = cx.update(|cx| cx.quit());
crates/edit_prediction_cli/src/metrics.rs
+21 -30
View File
@@ -1,30 +1,34 @@
use collections::{HashMap, HashSet};
use edit_prediction::udiff::DiffLine;
use serde::{Deserialize, Serialize};
type Counts = HashMap<String, usize>;
type CountsDelta = HashMap<String, isize>;
#[derive(Default, Debug, Clone)]
pub struct Scores {
#[derive(Default, Debug, Clone, Serialize, Deserialize)]
pub struct ClassificationMetrics {
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 {
impl ClassificationMetrics {
pub fn from_sets(
expected: &HashSet<String>,
actual: &HashSet<String>,
) -> ClassificationMetrics {
let true_positives = expected.intersection(actual).count();
let false_positives = actual.difference(expected).count();
let false_negatives = expected.difference(actual).count();
Scores {
ClassificationMetrics {
true_positives,
false_positives,
false_negatives,
}
}
pub fn from_counts(expected: &Counts, actual: &Counts) -> Scores {
pub fn from_counts(expected: &Counts, actual: &Counts) -> ClassificationMetrics {
let mut true_positives = 0;
let mut false_positives = 0;
let mut false_negatives = 0;
@@ -45,32 +49,16 @@ impl Scores {
}
}
Scores {
ClassificationMetrics {
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 {
pub fn aggregate<'a>(
scores: impl Iterator<Item = &'a ClassificationMetrics>,
) -> ClassificationMetrics {
let mut true_positives = 0;
let mut false_positives = 0;
let mut false_negatives = 0;
@@ -81,7 +69,7 @@ False Negatives : {}",
false_negatives += score.false_negatives;
}
Scores {
ClassificationMetrics {
true_positives,
false_positives,
false_negatives,
@@ -115,7 +103,10 @@ False Negatives : {}",
}
}
pub fn line_match_score(expected_patch: &[DiffLine], actual_patch: &[DiffLine]) -> Scores {
pub fn line_match_score(
expected_patch: &[DiffLine],
actual_patch: &[DiffLine],
) -> ClassificationMetrics {
let expected_change_lines = expected_patch
.iter()
.filter(|line| matches!(line, DiffLine::Addition(_) | DiffLine::Deletion(_)))
@@ -128,7 +119,7 @@ pub fn line_match_score(expected_patch: &[DiffLine], actual_patch: &[DiffLine])
.map(|line| line.to_string())
.collect();
Scores::from_sets(&expected_change_lines, &actual_change_lines)
ClassificationMetrics::from_sets(&expected_change_lines, &actual_change_lines)
}
enum ChrfWhitespace {
@@ -204,7 +195,7 @@ pub fn delta_chr_f(expected: &[DiffLine], actual: &[DiffLine]) -> f64 {
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);
let score = ClassificationMetrics::from_counts(&expected_counts, &actual_counts);
total_precision += score.precision();
total_recall += score.recall();
}
crates/edit_prediction_cli/src/paths.rs
+18 -50
View File
@@ -1,57 +1,25 @@
use std::{env, path::PathBuf, sync::LazyLock};
use std::{
path::{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 DATA_DIR: LazyLock<PathBuf> = LazyLock::new(|| {
let dir = dirs::home_dir().unwrap().join(".zed_ep");
ensure_dir(&dir)
});
pub static CACHE_DIR: LazyLock<PathBuf> = LazyLock::new(|| ensure_dir(&DATA_DIR.join("cache")));
pub static REPOS_DIR: LazyLock<PathBuf> = LazyLock::new(|| ensure_dir(&DATA_DIR.join("repos")));
pub static WORKTREES_DIR: LazyLock<PathBuf> =
LazyLock::new(|| ensure_dir(&DATA_DIR.join("worktrees")));
pub static RUN_DIR: LazyLock<PathBuf> = LazyLock::new(|| {
TARGET_ZETA_DIR
DATA_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 static LATEST_EXAMPLE_RUN_DIR: LazyLock<PathBuf> = LazyLock::new(|| DATA_DIR.join("latest"));
pub static LLM_CACHE_DB: LazyLock<PathBuf> = LazyLock::new(|| CACHE_DIR.join("llm_cache.sqlite"));
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(&current_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(&current_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()
);
fn ensure_dir(path: &Path) -> PathBuf {
std::fs::create_dir_all(path).expect("Failed to create directory");
path.to_path_buf()
}
crates/edit_prediction_cli/src/predict.rs
+245 -348
View File
@@ -1,374 +1,271 @@
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,
PredictionProvider, PromptFormat,
anthropic_client::AnthropicClient,
example::{Example, ExamplePrediction},
format_prompt::{PromptParser, TeacherPrompt, run_format_prompt},
headless::EpAppState,
load_project::run_load_project,
paths::{LATEST_EXAMPLE_RUN_DIR, RUN_DIR},
retrieve_context::run_context_retrieval,
};
use edit_prediction::{DebugEvent, EditPredictionStore};
use futures::{FutureExt as _, StreamExt as _, future::Shared};
use gpui::{AppContext as _, AsyncApp, Task};
use std::{
fs,
sync::{
Arc, Mutex, OnceLock,
atomic::{AtomicUsize, Ordering::SeqCst},
},
};
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,
pub async fn run_prediction(
example: &mut Example,
provider: Option<PredictionProvider>,
repetition_count: usize,
app_state: Arc<EpAppState>,
mut cx: 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
if !example.predictions.is_empty() {
return;
}
run_load_project(example, app_state.clone(), cx.clone()).await;
run_context_retrieval(example, app_state.clone(), cx.clone()).await;
let provider = provider.unwrap();
if matches!(provider, PredictionProvider::Teacher) {
if example.prompt.is_none() {
run_format_prompt(example, PromptFormat::Teacher, app_state.clone(), cx).await;
}
let batched = true;
return predict_anthropic(example, repetition_count, batched).await;
}
if matches!(
provider,
PredictionProvider::Zeta1 | PredictionProvider::Zeta2
) {
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 ep_store = cx
.update(|cx| EditPredictionStore::try_global(cx).unwrap())
.unwrap();
result.write(args.format, std::io::stdout()).unwrap();
print_run_data_dir(true, std::io::stdout().is_terminal());
}
ep_store
.update(&mut 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,
PredictionProvider::Mercury => edit_prediction::EditPredictionModel::Mercury,
PredictionProvider::Teacher => unreachable!(),
};
store.set_edit_prediction_model(model);
})
.unwrap();
let state = example.state.as_ref().unwrap();
let run_dir = RUN_DIR.join(&example.name);
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,
)
})?;
let updated_example = Arc::new(Mutex::new(example.clone()));
let current_run_ix = Arc::new(AtomicUsize::new(0));
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 mut debug_rx = ep_store
.update(&mut cx, |store, cx| store.debug_info(&state.project, cx))
.unwrap();
let debug_task = cx.background_spawn({
let updated_example = updated_example.clone();
let current_run_ix = current_run_ix.clone();
let run_dir = run_dir.clone();
async move {
while let Some(event) = debug_rx.next().await {
let run_ix = current_run_ix.load(SeqCst);
let mut updated_example = updated_example.lock().unwrap();
let buffer_store = project.read_with(cx, |project, _| project.buffer_store().clone())?;
let run_dir = if repetition_count > 1 {
run_dir.join(format!("{:03}", run_ix))
} else {
run_dir.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();
match event {
DebugEvent::EditPredictionStarted(request) => {
assert_eq!(updated_example.predictions.len(), run_ix + 1);
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,
)?;
if let Some(prompt) = request.prompt {
fs::write(run_dir.join("prediction_prompt.md"), &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(),
)?;
}
}
}
DebugEvent::EditPredictionFinished(request) => {
assert_eq!(updated_example.predictions.len(), run_ix + 1);
if let Some(output) = request.model_output {
fs::write(run_dir.join("prediction_response.md"), &output)?;
updated_example
.predictions
.last_mut()
.unwrap()
.actual_output = output;
}
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::open_ai_response::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();
if run_ix >= repetition_count {
break;
}
}
_ => {}
}
anyhow::Ok(())
}
});
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
);
for ix in 0..repetition_count {
current_run_ix.store(ix, SeqCst);
let run_dir = if repetition_count > 1 {
run_dir.join(format!("{:03}", ix))
} 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(),
run_dir.clone()
};
Ok(out.write_all(formatted.as_bytes())?)
fs::create_dir_all(&run_dir).unwrap();
if LATEST_EXAMPLE_RUN_DIR.is_symlink() {
fs::remove_file(&*LATEST_EXAMPLE_RUN_DIR).unwrap();
}
#[cfg(unix)]
std::os::unix::fs::symlink(&run_dir, &*LATEST_EXAMPLE_RUN_DIR).unwrap();
#[cfg(windows)]
std::os::windows::fs::symlink_dir(&run_dir, &*LATEST_EXAMPLE_RUN_DIR).unwrap();
updated_example
.lock()
.unwrap()
.predictions
.push(ExamplePrediction {
actual_patch: String::new(),
actual_output: String::new(),
provider,
});
let prediction = ep_store
.update(&mut cx, |store, cx| {
store.request_prediction(
&state.project,
&state.buffer,
state.cursor_position,
cloud_llm_client::PredictEditsRequestTrigger::Cli,
cx,
)
})
.unwrap()
.await
.unwrap();
updated_example
.lock()
.unwrap()
.predictions
.last_mut()
.unwrap()
.actual_patch = prediction
.and_then(|prediction| {
let prediction = prediction.prediction.ok()?;
prediction.edit_preview.as_unified_diff(&prediction.edits)
})
.unwrap_or_default();
}
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(),
)
ep_store
.update(&mut cx, |store, _| {
store.remove_project(&state.project);
})
.unwrap();
debug_task.await.unwrap();
*example = Arc::into_inner(updated_example)
.unwrap()
.into_inner()
.unwrap();
}
async fn predict_anthropic(example: &mut Example, _repetition_count: usize, batched: bool) {
let llm_model_name = "claude-sonnet-4-5";
let max_tokens = 16384;
let llm_client = if batched {
AnthropicClient::batch(&crate::paths::LLM_CACHE_DB.as_ref())
} else {
AnthropicClient::plain()
};
let llm_client = llm_client.expect("Failed to create LLM client");
let prompt = example
.prompt
.as_ref()
.unwrap_or_else(|| panic!("Prompt is required for an example {}", &example.name));
let messages = vec![anthropic::Message {
role: anthropic::Role::User,
content: vec![anthropic::RequestContent::Text {
text: prompt.input.clone(),
cache_control: None,
}],
}];
let Some(response) = llm_client
.generate(llm_model_name, max_tokens, messages)
.await
.unwrap()
else {
// Request stashed for batched processing
return;
};
let actual_output = response
.content
.into_iter()
.filter_map(|content| match content {
anthropic::ResponseContent::Text { text } => Some(text),
_ => None,
})
.collect::<Vec<String>>()
.join("\n");
let actual_patch = TeacherPrompt::parse(example, &actual_output);
let prediction = ExamplePrediction {
actual_patch,
actual_output,
provider: PredictionProvider::Teacher,
};
example.predictions.push(prediction);
}
pub async fn sync_batches(provider: &PredictionProvider) {
match provider {
PredictionProvider::Teacher => {
let cache_path = crate::paths::LLM_CACHE_DB.as_ref();
let llm_client =
AnthropicClient::batch(cache_path).expect("Failed to create LLM client");
llm_client
.sync_batches()
.await
.expect("Failed to sync batches");
}
_ => (),
}
}
crates/edit_prediction_cli/src/util.rs → crates/edit_prediction_cli/src/retrieve_context.rs
+98 -68
View File
@@ -1,106 +1,136 @@
use anyhow::{Result, anyhow};
use futures::channel::mpsc;
use futures::{FutureExt as _, StreamExt as _};
use crate::{
example::{Example, ExampleContext},
headless::EpAppState,
load_project::run_load_project,
};
use anyhow::Result;
use collections::HashSet;
use edit_prediction::{DebugEvent, EditPredictionStore};
use futures::{FutureExt as _, StreamExt as _, channel::mpsc};
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;
use language::{Buffer, LanguageNotFound};
use project::Project;
use std::{sync::Arc, time::Duration};
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,
})?;
pub async fn run_context_retrieval(
example: &mut Example,
app_state: Arc<EpAppState>,
mut cx: AsyncApp,
) {
if example.context.is_some() {
return;
}
let buffer = project
.update(cx, |project, cx| project.open_buffer(project_path, cx))?
.await?;
run_load_project(example, app_state.clone(), cx.clone()).await;
let mut parse_status = buffer.read_with(cx, |buffer, _cx| buffer.parse_status())?;
while *parse_status.borrow() != ParseStatus::Idle {
parse_status.changed().await?;
let state = example.state.as_ref().unwrap();
let project = state.project.clone();
let _lsp_handle = project
.update(&mut cx, |project, cx| {
project.register_buffer_with_language_servers(&state.buffer, cx)
})
.unwrap();
wait_for_language_server_to_start(example, &project, &state.buffer, &mut cx).await;
let ep_store = cx
.update(|cx| EditPredictionStore::try_global(cx).unwrap())
.unwrap();
let mut events = ep_store
.update(&mut cx, |store, cx| {
store.register_buffer(&state.buffer, &project, cx);
store.set_use_context(true);
store.refresh_context(&project, &state.buffer, state.cursor_position, cx);
store.debug_info(&project, cx)
})
.unwrap();
while let Some(event) = events.next().await {
match event {
DebugEvent::ContextRetrievalFinished(_) => {
break;
}
_ => {}
}
}
Ok(buffer)
})
let context_files = ep_store
.update(&mut cx, |store, cx| store.context_for_project(&project, cx))
.unwrap();
example.context = Some(ExampleContext {
files: context_files,
});
}
pub async fn open_buffer_with_language_server(
project: Entity<Project>,
worktree: Entity<Worktree>,
path: Arc<RelPath>,
ready_languages: &mut HashSet<LanguageId>,
async fn wait_for_language_server_to_start(
example: &Example,
project: &Entity<Project>,
buffer: &Entity<Buffer>,
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 language_registry = project
.read_with(cx, |project, _| project.languages().clone())
.unwrap();
let result = language_registry
.load_language_for_file_path(path.as_std_path())
.load_language_for_file_path(&example.cursor_path)
.await;
if let Err(error) = result
&& !error.is::<LanguageNotFound>()
{
anyhow::bail!(error);
panic!("Failed to load language for file path: {}", error);
}
let Some(language_id) = buffer.read_with(cx, |buffer, _cx| {
buffer.language().map(|language| language.id())
})?
let Some(language_id) = buffer
.read_with(cx, |buffer, _cx| {
buffer.language().map(|language| language.id())
})
.unwrap()
else {
return Err(anyhow!("No language for {}", path.display(path_style)));
panic!("No language for {:?}", example.cursor_path);
};
let log_prefix = format!("{} | ", path.display(path_style));
let mut ready_languages = HashSet::default();
let log_prefix = format!("{} | ", example.name);
if !ready_languages.contains(&language_id) {
wait_for_lang_server(&project, &buffer, log_prefix, cx).await?;
wait_for_lang_server(&project, &buffer, log_prefix, cx)
.await
.unwrap();
ready_languages.insert(language_id);
}
let lsp_store = project.read_with(cx, |project, _cx| project.lsp_store())?;
let lsp_store = project
.read_with(cx, |project, _cx| project.lsp_store())
.unwrap();
// 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())
if 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 {
.unwrap()
.is_some()
{
return;
} 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"));
panic!("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>,
crates/edit_prediction_cli/src/score.rs
+119
View File
@@ -0,0 +1,119 @@
use crate::{
PredictArgs,
example::{Example, ExampleScore},
headless::EpAppState,
metrics::{self, ClassificationMetrics},
predict::run_prediction,
};
use edit_prediction::udiff::DiffLine;
use gpui::AsyncApp;
use std::sync::Arc;
pub async fn run_scoring(
example: &mut Example,
args: &PredictArgs,
app_state: Arc<EpAppState>,
cx: AsyncApp,
) {
run_prediction(
example,
Some(args.provider),
args.repetitions,
app_state,
cx,
)
.await;
let expected_patch = parse_patch(&example.expected_patch);
let mut scores = vec![];
for pred in &example.predictions {
let actual_patch = parse_patch(&pred.actual_patch);
let line_match = metrics::line_match_score(&expected_patch, &actual_patch);
let delta_chr_f = metrics::delta_chr_f(&expected_patch, &actual_patch) as f32;
scores.push(ExampleScore {
delta_chr_f,
line_match,
});
}
example.score = scores;
}
fn parse_patch(patch: &str) -> Vec<DiffLine<'_>> {
patch.lines().map(DiffLine::parse).collect()
}
pub fn print_report(examples: &[Example]) {
eprintln!(
"──────────────────────────────────────────────────────────────────────────────────────"
);
eprintln!(
"{:<30} {:>4} {:>4} {:>4} {:>10} {:>8} {:>8} {:>10}",
"Example name", "TP", "FP", "FN", "Precision", "Recall", "F1", "DeltaChrF"
);
eprintln!(
"──────────────────────────────────────────────────────────────────────────────────────"
);
let mut all_line_match_scores = Vec::new();
let mut all_delta_chr_f_scores = Vec::new();
for example in examples {
for score in example.score.iter() {
let line_match = &score.line_match;
eprintln!(
"{:<30} {:>4} {:>4} {:>4} {:>9.2}% {:>7.2}% {:>7.2}% {:>9.2}",
truncate_name(&example.name, 30),
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,
score.delta_chr_f
);
all_line_match_scores.push(line_match.clone());
all_delta_chr_f_scores.push(score.delta_chr_f);
}
}
eprintln!(
"──────────────────────────────────────────────────────────────────────────────────────"
);
if !all_line_match_scores.is_empty() {
let total_line_match = ClassificationMetrics::aggregate(all_line_match_scores.iter());
let avg_delta_chr_f: f32 =
all_delta_chr_f_scores.iter().sum::<f32>() / all_delta_chr_f_scores.len() as f32;
eprintln!(
"{:<30} {:>4} {:>4} {:>4} {:>9.2}% {:>7.2}% {:>7.2}% {:>9.2}",
"TOTAL",
total_line_match.true_positives,
total_line_match.false_positives,
total_line_match.false_negatives,
total_line_match.precision() * 100.0,
total_line_match.recall() * 100.0,
total_line_match.f1_score() * 100.0,
avg_delta_chr_f
);
eprintln!(
"──────────────────────────────────────────────────────────────────────────────────────"
);
}
eprintln!("\n");
}
fn truncate_name(name: &str, max_len: usize) -> String {
if name.len() <= max_len {
name.to_string()
} else {
format!("{}...", &name[..max_len - 3])
}
}
crates/edit_prediction_cli/src/source_location.rs
-70
View File
@@ -1,70 +0,0 @@
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 })
}
}
crates/edit_prediction_cli/src/training/teacher.prompt.md → crates/edit_prediction_cli/src/teacher.prompt.md
+4
View File
@@ -46,3 +46,7 @@ Output example:
## Code Context
{{context}}
## Editable region
{{editable_region}}
crates/edit_prediction_cli/src/training/context.rs
-89
View File
@@ -1,89 +0,0 @@
use std::path::Path;
use crate::{source_location::SourceLocation, training::teacher::TeacherModel};
#[derive(Debug, Clone, Default, clap::ValueEnum)]
pub enum ContextType {
#[default]
CurrentFile,
}
const MAX_CONTEXT_SIZE: usize = 32768;
pub fn collect_context(
context_type: &ContextType,
worktree_dir: &Path,
cursor: SourceLocation,
) -> String {
let context = match context_type {
ContextType::CurrentFile => {
let file_path = worktree_dir.join(cursor.path.as_std_path());
let context = std::fs::read_to_string(&file_path).unwrap_or_default();
let context = add_special_tags(&context, worktree_dir, cursor);
context
}
};
let region_end_offset = context.find(TeacherModel::REGION_END);
if context.len() <= MAX_CONTEXT_SIZE {
return context;
}
if let Some(region_end_offset) = region_end_offset
&& region_end_offset + TeacherModel::REGION_END.len() > MAX_CONTEXT_SIZE
{
let to_truncate = context.len() - MAX_CONTEXT_SIZE;
format!(
"[...{} bytes truncated]\n{}\n",
to_truncate,
&context[to_truncate..]
)
} else {
format!(
"{}\n[...{} bytes truncated]\n",
&context[..MAX_CONTEXT_SIZE],
context.len() - MAX_CONTEXT_SIZE
)
}
}
/// Add <|editable_region_start/end|> tags
fn add_special_tags(context: &str, worktree_dir: &Path, cursor: SourceLocation) -> String {
let path = worktree_dir.join(cursor.path.as_std_path());
let file = std::fs::read_to_string(&path).unwrap_or_default();
let lines = file.lines().collect::<Vec<_>>();
let cursor_row = cursor.point.row as usize;
let start_line = cursor_row.saturating_sub(TeacherModel::LEFT_CONTEXT_SIZE);
let end_line = (cursor_row + TeacherModel::RIGHT_CONTEXT_SIZE).min(lines.len());
let snippet = lines[start_line..end_line].join("\n");
if context.contains(&snippet) {
let mut cursor_line = lines[cursor_row].to_string();
cursor_line.insert_str(cursor.point.column as usize, TeacherModel::USER_CURSOR);
let mut snippet_with_tags_lines = vec![];
snippet_with_tags_lines.push(TeacherModel::REGION_START);
snippet_with_tags_lines.extend(&lines[start_line..cursor_row]);
snippet_with_tags_lines.push(&cursor_line);
snippet_with_tags_lines.extend(&lines[cursor_row + 1..end_line]);
snippet_with_tags_lines.push(TeacherModel::REGION_END);
let snippet_with_tags = snippet_with_tags_lines.join("\n");
context.replace(&snippet, &snippet_with_tags)
} else {
log::warn!(
"Can't find area around the cursor in the context; proceeding without special tags"
);
context.to_string()
}
}
pub fn strip_special_tags(context: &str) -> String {
context
.replace(TeacherModel::REGION_START, "")
.replace(TeacherModel::REGION_END, "")
.replace(TeacherModel::USER_CURSOR, "")
}
crates/edit_prediction_cli/src/training/distill.rs
-94
View File
@@ -1,94 +0,0 @@
use serde::Deserialize;
use std::sync::Arc;
use crate::{
DistillArguments,
example::Example,
source_location::SourceLocation,
training::{
context::ContextType,
llm_client::LlmClient,
teacher::{TeacherModel, TeacherOutput},
},
};
use anyhow::Result;
use reqwest_client::ReqwestClient;
#[derive(Debug, Deserialize)]
pub struct SplitCommit {
repo_url: String,
commit_sha: String,
edit_history: String,
expected_patch: String,
cursor_position: String,
}
pub async fn run_distill(arguments: DistillArguments) -> Result<()> {
let split_commits: Vec<SplitCommit> = std::fs::read_to_string(&arguments.split_commit_dataset)
.expect("Failed to read split commit dataset")
.lines()
.map(|line| serde_json::from_str(line).expect("Failed to parse JSON line"))
.collect();
let http_client: Arc<dyn http_client::HttpClient> = Arc::new(ReqwestClient::new());
let llm_client = if let Some(cache_path) = arguments.batch {
LlmClient::batch(&cache_path, http_client)?
} else {
LlmClient::plain(http_client)?
};
let mut teacher = TeacherModel::new(
"claude-sonnet-4-5".to_string(),
ContextType::CurrentFile,
llm_client,
);
let mut num_marked_for_batching = 0;
for commit in split_commits {
if let Some(distilled) = distill_one(&mut teacher, commit).await? {
println!("{}", serde_json::to_string(&distilled)?);
} else {
if num_marked_for_batching == 0 {
log::warn!("Marked for batching");
}
num_marked_for_batching += 1;
}
}
eprintln!(
"{} requests are marked for batching",
num_marked_for_batching
);
let llm_client = teacher.client;
llm_client.sync_batches().await?;
Ok(())
}
pub async fn distill_one(
teacher: &mut TeacherModel,
commit: SplitCommit,
) -> Result<Option<TeacherOutput>> {
let cursor: SourceLocation = commit
.cursor_position
.parse()
.expect("Failed to parse cursor position");
let path = cursor.path.to_rel_path_buf();
let example = Example {
repository_url: commit.repo_url,
revision: commit.commit_sha,
uncommitted_diff: commit.edit_history.clone(),
cursor_path: path.as_std_path().to_path_buf(),
cursor_position: commit.cursor_position,
edit_history: commit.edit_history, // todo: trim
expected_patch: commit.expected_patch,
};
let prediction = teacher.predict(example).await;
prediction
}
crates/edit_prediction_cli/src/training/mod.rs
-4
View File
@@ -1,4 +0,0 @@
pub mod context;
pub mod distill;
pub mod llm_client;
pub mod teacher;
crates/edit_prediction_cli/src/training/teacher.rs
-266
View File
@@ -1,266 +0,0 @@
use crate::{
example::Example,
source_location::SourceLocation,
training::{
context::{ContextType, collect_context, strip_special_tags},
llm_client::LlmClient,
},
};
use anthropic::{Message, RequestContent, ResponseContent, Role};
use anyhow::Result;
pub struct TeacherModel {
pub llm_name: String,
pub context: ContextType,
pub client: LlmClient,
}
#[derive(Debug, serde::Serialize)]
pub struct TeacherOutput {
parsed_output: String,
prompt: String,
raw_llm_response: String,
context: String,
diff: String,
}
impl TeacherModel {
const PROMPT: &str = include_str!("teacher.prompt.md");
pub(crate) const REGION_START: &str = "<|editable_region_start|>\n";
pub(crate) const REGION_END: &str = "<|editable_region_end|>";
pub(crate) const USER_CURSOR: &str = "<|user_cursor|>";
/// Number of lines to include before the cursor position
pub(crate) const LEFT_CONTEXT_SIZE: usize = 5;
/// Number of lines to include after the cursor position
pub(crate) const RIGHT_CONTEXT_SIZE: usize = 5;
/// Truncate edit history to this number of last lines
const MAX_HISTORY_LINES: usize = 128;
pub fn new(llm_name: String, context: ContextType, client: LlmClient) -> Self {
TeacherModel {
llm_name,
context,
client,
}
}
pub async fn predict(&self, input: Example) -> Result<Option<TeacherOutput>> {
let name = input.unique_name();
let worktree_dir = input.setup_worktree(name).await?;
let cursor: SourceLocation = input
.cursor_position
.parse()
.expect("Failed to parse cursor position");
let context = collect_context(&self.context, &worktree_dir, cursor.clone());
let edit_history = Self::format_edit_history(&input.edit_history);
let prompt = Self::PROMPT
.replace("{{context}}", &context)
.replace("{{edit_history}}", &edit_history);
let messages = vec![Message {
role: Role::User,
content: vec![RequestContent::Text {
text: prompt.clone(),
cache_control: None,
}],
}];
let Some(response) = self
.client
.generate(self.llm_name.clone(), 16384, messages)
.await?
else {
return Ok(None);
};
let response_text = response
.content
.into_iter()
.filter_map(|content| match content {
ResponseContent::Text { text } => Some(text),
_ => None,
})
.collect::<Vec<String>>()
.join("\n");
let parsed_output = self.parse_response(&response_text);
let original_editable_region = Self::extract_editable_region(&context);
let context_after_edit = context.replace(&original_editable_region, &parsed_output);
let context_after_edit = strip_special_tags(&context_after_edit);
let context_before_edit = strip_special_tags(&context);
let diff = language::unified_diff(&context_before_edit, &context_after_edit);
// zeta distill --batch batch_results.txt
// zeta distill
// 1. Run `zeta distill <2000 examples <- all examples>` for the first time
// - store LLM requests in a batch, don't actual send the request
// - send the batch (2000 requests) after all inputs are processed
// 2. `zeta send-batches`
// - upload the batch to Anthropic
// https://platform.claude.com/docs/en/build-with-claude/batch-processing
// https://crates.io/crates/anthropic-sdk-rust
// - poll for results
// - when ready, store results in cache (a database)
// 3. `zeta distill` again
// - use the cached results this time
Ok(Some(TeacherOutput {
parsed_output,
prompt,
raw_llm_response: response_text,
context,
diff,
}))
}
fn parse_response(&self, content: &str) -> String {
let codeblock = Self::extract_last_codeblock(content);
let editable_region = Self::extract_editable_region(&codeblock);
editable_region
}
/// Extract content from the last code-fenced block if any, or else return content as is
fn extract_last_codeblock(text: &str) -> String {
let mut last_block = None;
let mut search_start = 0;
while let Some(start) = text[search_start..].find("```") {
let start = start + search_start;
let bytes = text.as_bytes();
let mut backtick_end = start;
while backtick_end < bytes.len() && bytes[backtick_end] == b'`' {
backtick_end += 1;
}
let backtick_count = backtick_end - start;
let closing_backticks = "`".repeat(backtick_count);
if let Some(end_pos) = text[backtick_end..].find(&closing_backticks) {
let code_block = &text[backtick_end + 1..backtick_end + end_pos - 1];
last_block = Some(code_block.to_string());
search_start = backtick_end + end_pos + backtick_count;
} else {
break;
}
}
last_block.unwrap_or_else(|| text.to_string())
}
fn extract_editable_region(text: &str) -> String {
let start = text
.find(Self::REGION_START)
.map_or(0, |pos| pos + Self::REGION_START.len());
let end = text.find(Self::REGION_END).unwrap_or(text.len());
text[start..end].to_string()
}
/// Truncates edit history to a maximum length and removes comments (unified diff garbage lines)
fn format_edit_history(edit_history: &str) -> String {
let lines = edit_history
.lines()
.filter(|&s| Self::is_content_line(s))
.collect::<Vec<_>>();
let history_lines = if lines.len() > Self::MAX_HISTORY_LINES {
&lines[lines.len() - Self::MAX_HISTORY_LINES..]
} else {
&lines
};
history_lines.join("\n")
}
fn is_content_line(s: &str) -> bool {
s.starts_with("-")
|| s.starts_with("+")
|| s.starts_with(" ")
|| s.starts_with("---")
|| s.starts_with("+++")
|| s.starts_with("@@")
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_parse_response() {
let teacher = TeacherModel::new(
"test".to_string(),
ContextType::CurrentFile,
LlmClient::dummy(),
);
let response = "This is a test response.";
let parsed = teacher.parse_response(response);
assert_eq!(parsed, response.to_string());
let response = indoc::indoc! {"
Some thinking
`````
actual response
`````
"};
let parsed = teacher.parse_response(response);
assert_eq!(parsed, "actual response");
}
#[test]
fn test_extract_last_code_block() {
let text = indoc::indoc! {"
Some thinking
```
first block
```
`````
last block
`````
"};
let last_block = TeacherModel::extract_last_codeblock(text);
assert_eq!(last_block, "last block");
}
#[test]
fn test_extract_editable_region() {
let teacher = TeacherModel::new(
"test".to_string(),
ContextType::CurrentFile,
LlmClient::dummy(),
);
let response = indoc::indoc! {"
some lines
are
here
<|editable_region_start|>
one
two three
<|editable_region_end|>
more
lines here
"};
let parsed = teacher.parse_response(response);
assert_eq!(
parsed,
indoc::indoc! {"
one
two three
"}
);
}
}