From 0227e8384946b5d2331a8d7706a81b15fc28ea8b Mon Sep 17 00:00:00 2001 From: Ben Kunkle Date: Wed, 24 Jun 2026 21:32:40 -0400 Subject: [PATCH] ep: Batch settled edit prediction telemetry (#59857) # Objective Closes EP-78 - Describe the objective or issue this PR addresses. - If you're fixing a specific issue, use "Fixes #X" for each issue as [described in the GitHub docs](https://docs.github.com/en/issues/tracking-your-work-with-issues/using-issues/linking-a-pull-request-to-an-issue#linking-a-pull-request-to-an-issue-using-a-keyword). ## Solution - Describe the solution used to achieve the objective above. ## Testing - Did you test these changes? If so, how? - Are there any parts that need more testing? - How can other people (reviewers) test your changes? Is there anything specific they need to know? - If relevant, what platforms did you test these changes on, and are there any important ones you can't test? ## Self-Review Checklist: - [x] I've reviewed my own diff for quality, security, and reliability - [x] Unsafe blocks (if any) have justifying comments - [x] The content adheres to Zed's UI standards ([UX/UI](https://github.com/zed-industries/zed/blob/main/CONTRIBUTING.md#uiux-checklist) and [icon](https://github.com/zed-industries/zed/blob/main/crates/icons/README.md) guidelines) - [x] Tests cover the new/changed behavior - [x] Performance impact has been considered and is acceptable ## Showcase > This section is optional. If this PR does not include a visual change or does not add a new user-facing feature, you can delete this section. - Help others understand the result of this PR by showcasing your awesome work! - If this PR includes a visual change, consider adding a screenshot, GIF, or video - A before/after comparison is very useful for changes to existing features! While a showcase should aim to be brief and digestible, you can use a toggleable section to save space on longer showcases:
Click to view showcase My super cool demos here
--- Release Notes: - N/A or Added/Fixed/Improved ... --- .../src/capture_prediction_context.rs | 95 ------ crates/edit_prediction/src/data_collection.rs | 93 +++++- crates/edit_prediction/src/edit_prediction.rs | 283 ++++++++++-------- .../src/edit_prediction_tests.rs | 84 +++++- crates/edit_prediction/src/zeta.rs | 60 ++-- 5 files changed, 357 insertions(+), 258 deletions(-) delete mode 100644 crates/edit_prediction/src/capture_prediction_context.rs diff --git a/crates/edit_prediction/src/capture_prediction_context.rs b/crates/edit_prediction/src/capture_prediction_context.rs deleted file mode 100644 index 39f05532e7..0000000000 --- a/crates/edit_prediction/src/capture_prediction_context.rs +++ /dev/null @@ -1,95 +0,0 @@ -use crate::{ - EditPredictionStore, StoredEvent, - data_collection::{compute_uncommitted_diff, uncommitted_diffs_for_events}, - zeta, -}; -use anyhow::Result; -use gpui::{Context, Entity, Task}; -use language::{Buffer, Point, ToPoint as _}; -use project::Project; -use text::OffsetRangeExt as _; - -const MAX_UNCOMMITTED_DIFF_SIZE: usize = 64 * 1024; - -pub(crate) struct CapturedPredictionContext { - pub(crate) repository_url: Option, - pub(crate) revision: Option, - pub(crate) uncommitted_diff: Option, - pub(crate) buffer_diagnostics: Vec, - pub(crate) editable_context: Vec, -} - -pub(crate) fn capture_prediction_context( - project: Entity, - buffer: Entity, - cursor_anchor: language::Anchor, - stored_events: Vec, - repository_url: Option, - revision: Option, - editable_context_task: Task>>, - cx: &mut Context, -) -> Option>> { - let snapshot = buffer.read(cx).snapshot(); - let worktree_id = snapshot.file()?.worktree_id(cx); - let uncommitted_diff_task = - uncommitted_diffs_for_events(project, worktree_id, stored_events, cx); - - Some(cx.spawn(async move |_this, cx| { - let uncommitted_diff_snapshot = match uncommitted_diff_task.await { - Ok(snapshot) => Some(snapshot), - Err(error) => { - log::debug!("failed to capture uncommitted diff: {error:?}"); - None - } - }; - - let uncommitted_diff = if let Some(uncommitted_diff_snapshot) = uncommitted_diff_snapshot { - let estimated_uncommitted_diff_size = uncommitted_diff_snapshot - .iter() - .map(|(_, buffer_snapshot, diff_snapshot)| { - diff_snapshot - .hunks(buffer_snapshot) - .map(|hunk| { - hunk.diff_base_byte_range.len() - + hunk.range.to_offset(buffer_snapshot).len() - }) - .sum::() - }) - .sum::(); - - if estimated_uncommitted_diff_size <= MAX_UNCOMMITTED_DIFF_SIZE { - let uncommitted_diff = cx - .background_executor() - .spawn(async move { compute_uncommitted_diff(uncommitted_diff_snapshot) }) - .await; - (uncommitted_diff.len() <= MAX_UNCOMMITTED_DIFF_SIZE).then_some(uncommitted_diff) - } else { - None - } - } else { - None - }; - - let buffer_diagnostics = zeta::active_buffer_diagnostics( - &snapshot, - Point::new(0, 0)..snapshot.max_point(), - cursor_anchor.to_point(&snapshot).row, - 100, - ); - let editable_context = match editable_context_task.await { - Ok(editable_context) => editable_context, - Err(error) => { - log::debug!("failed to capture editable context: {error:?}"); - Vec::new() - } - }; - - Ok(CapturedPredictionContext { - repository_url, - revision, - uncommitted_diff, - buffer_diagnostics, - editable_context, - }) - })) -} diff --git a/crates/edit_prediction/src/data_collection.rs b/crates/edit_prediction/src/data_collection.rs index 4be37c03b8..217124157a 100644 --- a/crates/edit_prediction/src/data_collection.rs +++ b/crates/edit_prediction/src/data_collection.rs @@ -1,20 +1,105 @@ -use crate::{EditPredictionStore, StoredEvent}; +use crate::{EditPredictionStore, StoredEvent, zeta}; -use anyhow::Context as _; +use anyhow::{Context as _, Result}; use buffer_diff::BufferDiffSnapshot; use collections::HashMap; use gpui::{Context, Entity, Task}; -use language::BufferSnapshot; +use language::{Buffer, BufferSnapshot, ToPoint as _}; use project::{Project, ProjectPath, WorktreeId}; use std::{fmt::Write as _, ops::Range, path::Path, sync::Arc}; -use text::Point; +use text::{OffsetRangeExt as _, Point}; use util::rel_path::RelPath; +const MAX_UNCOMMITTED_DIFF_SIZE: usize = 64 * 1024; + pub type UncommittedDiffSnapshot = Vec<(Arc, BufferSnapshot, BufferDiffSnapshot)>; pub type UncommittedDiffResult = std::result::Result>; pub use zeta_prompt::udiff::CURSOR_POSITION_MARKER; +pub(crate) struct CapturedPredictionContext { + pub(crate) repository_url: Option, + pub(crate) revision: Option, + pub(crate) uncommitted_diff: Option, + pub(crate) buffer_diagnostics: Vec, + pub(crate) editable_context: Vec, +} + +pub(crate) fn capture_prediction_context( + project: Entity, + buffer: Entity, + cursor_anchor: language::Anchor, + stored_events: Vec, + repository_url: Option, + revision: Option, + editable_context_task: Task>>, + cx: &mut Context, +) -> Option>> { + let snapshot = buffer.read(cx).snapshot(); + let worktree_id = snapshot.file()?.worktree_id(cx); + let uncommitted_diff_task = + uncommitted_diffs_for_events(project, worktree_id, stored_events, cx); + + Some(cx.spawn(async move |_this, cx| { + let uncommitted_diff_snapshot = match uncommitted_diff_task.await { + Ok(snapshot) => Some(snapshot), + Err(error) => { + log::debug!("failed to capture uncommitted diff: {error:?}"); + None + } + }; + + let uncommitted_diff = if let Some(uncommitted_diff_snapshot) = uncommitted_diff_snapshot { + let estimated_uncommitted_diff_size = uncommitted_diff_snapshot + .iter() + .map(|(_, buffer_snapshot, diff_snapshot)| { + diff_snapshot + .hunks(buffer_snapshot) + .map(|hunk| { + hunk.diff_base_byte_range.len() + + hunk.range.to_offset(buffer_snapshot).len() + }) + .sum::() + }) + .sum::(); + + if estimated_uncommitted_diff_size <= MAX_UNCOMMITTED_DIFF_SIZE { + let uncommitted_diff = cx + .background_executor() + .spawn(async move { compute_uncommitted_diff(uncommitted_diff_snapshot) }) + .await; + (uncommitted_diff.len() <= MAX_UNCOMMITTED_DIFF_SIZE).then_some(uncommitted_diff) + } else { + None + } + } else { + None + }; + + let buffer_diagnostics = zeta::active_buffer_diagnostics( + &snapshot, + Point::new(0, 0)..snapshot.max_point(), + cursor_anchor.to_point(&snapshot).row, + 100, + ); + let editable_context = match editable_context_task.await { + Ok(editable_context) => editable_context, + Err(error) => { + log::debug!("failed to capture editable context: {error:?}"); + Vec::new() + } + }; + + Ok(CapturedPredictionContext { + repository_url, + revision, + uncommitted_diff, + buffer_diagnostics, + editable_context, + }) + })) +} + pub fn uncommitted_diffs_for_events( project: Entity, worktree_id: WorktreeId, diff --git a/crates/edit_prediction/src/edit_prediction.rs b/crates/edit_prediction/src/edit_prediction.rs index f9b07daa4f..5e4cdf6da7 100644 --- a/crates/edit_prediction/src/edit_prediction.rs +++ b/crates/edit_prediction/src/edit_prediction.rs @@ -3,8 +3,10 @@ use buffer_diff::BufferDiff; use client::{Client, EditPredictionUsage, UserStore, global_llm_token}; use cloud_api_client::LlmApiToken; use cloud_api_types::{ - EditPredictionRecentFile, EditPredictionSettledKeptChars, OrganizationId, - SettledEditPrediction, SettledEditPredictionSampleData, SubmitEditPredictionFeedbackBody, + EditPredictionRecentFile, EditPredictionSettledKeptChars, + MAX_EDIT_PREDICTION_SETTLED_PER_REQUEST, OrganizationId, SettledEditPrediction, + SettledEditPredictionSampleData, SubmitEditPredictionFeedbackBody, + SubmitEditPredictionSettledBatchBody, SubmitEditPredictionSettledResponse, }; use cloud_llm_client::predict_edits_v3::{ PREDICT_EDITS_MODE_HEADER_NAME, PREDICT_EDITS_REQUEST_ID_HEADER_NAME, @@ -83,7 +85,6 @@ mod prediction; pub mod udiff; -mod capture_prediction_context; pub mod open_ai_compatible; mod zed_edit_prediction_delegate; pub mod zeta; @@ -91,8 +92,8 @@ pub mod zeta; #[cfg(test)] mod edit_prediction_tests; -use crate::capture_prediction_context::{CapturedPredictionContext, capture_prediction_context}; use crate::cursor_excerpt::expand_context_syntactically_then_linewise; +use crate::data_collection::{CapturedPredictionContext, capture_prediction_context}; use crate::example_spec::RecentFile; use crate::license_detection::LicenseDetectionWatcher; use crate::mercury::Mercury; @@ -130,8 +131,6 @@ const REQUEST_TIMEOUT_BACKOFF: Duration = Duration::from_secs(10); const EDIT_PREDICTION_SETTLED_TTL: Duration = Duration::from_secs(60 * 5); const EDIT_PREDICTION_SETTLED_QUIESCENCE: Duration = Duration::from_secs(10); const EDIT_PREDICTION_CAPTURE_MAX_FUTURE_EVENTS: usize = 4; -/// The server rejects settled bodies larger than 64 KiB (compressed). -const EDIT_PREDICTION_SETTLED_MAX_BODY_BYTES: usize = 63 * 1024; const EDIT_PREDICTION_SETTLED_MAX_EDITABLE_REGION_BYTES: usize = 4 * 1024; pub struct EditPredictionJumpsFeatureFlag; @@ -2026,25 +2025,11 @@ impl EditPredictionStore { } }); + let mut ready_predictions_by_organization_id: HashMap<_, Vec<_>> = HashMap::default(); for (pending_capture, settled_editable_region) in ready_predictions { - let PendingPredictionCapture { - request_id, - editable_region_before_prediction, - predicted_editable_region, - ts_error_count_before_prediction, - ts_error_count_after_prediction, - organization_id, - can_collect_data, - is_in_open_source_repo, - sample_data, - model_version, - e2e_latency, - .. - } = pending_capture; - let settled_editable_region_for_metrics = settled_editable_region.clone(); #[cfg(test)] { - let request_id = request_id.clone(); + let request_id = pending_capture.request_id.clone(); let settled_editable_region = settled_editable_region.clone(); this.update(cx, |this, _| { if let Some(callback) = &this.settled_event_callback { @@ -2052,111 +2037,28 @@ impl EditPredictionStore { } }); } - cx.background_spawn({ - let client = client.clone(); - let llm_token = llm_token.clone(); - let app_version = app_version.clone(); - async move { - let kept_rate_result = compute_kept_rate( - &editable_region_before_prediction, - &predicted_editable_region, - &settled_editable_region_for_metrics, - ); - - let result: anyhow::Result<()> = async { - let sample_data = if can_collect_data - && let Some(sample_data) = sample_data - && let Ok(context) = sample_data.context_task.await - { - Some(SettledEditPredictionSampleData { - repository_url: context.repository_url, - revision: context.revision, - uncommitted_diff: context.uncommitted_diff, - editable_path: sample_data.editable_path, - editable_offset_range: sample_data.editable_offset_range, - buffer_diagnostics: context.buffer_diagnostics, - editable_context: context.editable_context, - future_edit_history_events: sample_data - .future_edit_history_events, - navigation_history: sample_data - .navigation_history - .into_iter() - .map(|file| EditPredictionRecentFile { - path: file.path, - cursor_position: file.cursor_position, - }) - .collect(), - edit_events_before_quiescence: sample_data - .edit_events_before_quiescence, - next_edit_cursor_offset: sample_data.next_edit_cursor_offset, - }) - } else { - None - }; - let settled_editable_region = - can_collect_data.then_some(settled_editable_region); - - let mut body = SettledEditPrediction { - request_id: request_id.0.to_string(), - settled_editable_region, - ts_error_count_before_prediction, - ts_error_count_after_prediction, - can_collect_data, - is_in_open_source_repo, - sample_data, - kept_chars: EditPredictionSettledKeptChars { - candidate_new: kept_rate_result.candidate_new_chars, - reference_new: kept_rate_result.reference_new_chars, - candidate_deleted: kept_rate_result.candidate_deleted_chars, - reference_deleted: kept_rate_result.reference_deleted_chars, - kept: kept_rate_result.kept_chars, - correctly_deleted: kept_rate_result.correctly_deleted_chars, - discarded: kept_rate_result.discarded_chars, - context: kept_rate_result.context_chars, - kept_rate: kept_rate_result.kept_rate, - recall_rate: kept_rate_result.recall_rate, - }, - example: None, - model_version, - e2e_latency_ms: e2e_latency.as_millis().min(u128::from(u64::MAX)) - as u64, - }; - - let mut compressed = - zstd::encode_all(&serde_json::to_vec(&body)?[..], 3)?; - if compressed.len() > EDIT_PREDICTION_SETTLED_MAX_BODY_BYTES { - body.sample_data = None; - compressed = zstd::encode_all(&serde_json::to_vec(&body)?[..], 3)?; - } - - let url = client - .http_client() - .build_zed_llm_url("/predict_edits/settled", &[])?; - Self::send_api_request::( - |builder| { - Ok(builder - .uri(url.as_ref()) - .header("Content-Encoding", "zstd") - .body(compressed.clone().into())?) - }, - client, - llm_token, - organization_id, - app_version, - ) - .await?; - Ok(()) - } - .await; - - if let Err(error) = result { - log::error!("failed to submit edit prediction settled: {error:?}"); - } - } - }) - .detach(); + ready_predictions_by_organization_id + .entry(pending_capture.organization_id.clone()) + .or_default() + .push((pending_capture, settled_editable_region)); } + cx.background_spawn({ + let client = client.clone(); + let llm_token = llm_token.clone(); + let app_version = app_version.clone(); + async move { + send_settled_batches( + client, + llm_token, + app_version, + ready_predictions_by_organization_id, + ) + .await; + } + }) + .detach(); + next_wake_time = oldest_edited_at.map(|time| time + EDIT_PREDICTION_SETTLED_QUIESCENCE); } } @@ -2416,6 +2318,137 @@ impl EditPredictionStore { pub const THROTTLE_TIMEOUT: Duration = Duration::from_millis(300); } +async fn send_settled_batches( + client: Arc, + llm_token: LlmApiToken, + app_version: Version, + ready_predictions_by_organization_id: hash_map::HashMap< + Option, + Vec<(PendingPredictionCapture, String)>, + collections::FxBuildHasher, + >, +) { + let Some(url) = client + .http_client() + .build_zed_llm_url("/predict_edits/settled", &[]) + .context("failed to build edit predictions settled url") + .log_err() + else { + return; + }; + + for (organization_id, ready_predictions) in ready_predictions_by_organization_id { + let mut ready_predictions = ready_predictions.into_iter(); + loop { + let done_batch = ready_predictions + .by_ref() + .take(MAX_EDIT_PREDICTION_SETTLED_PER_REQUEST); + let mut batch = Vec::with_capacity(MAX_EDIT_PREDICTION_SETTLED_PER_REQUEST); + for (pending_capture, settled_editable_region) in done_batch { + let PendingPredictionCapture { + request_id, + editable_region_before_prediction, + predicted_editable_region, + ts_error_count_before_prediction, + ts_error_count_after_prediction, + can_collect_data, + is_in_open_source_repo, + sample_data, + model_version, + e2e_latency, + .. + } = pending_capture; + let kept_rate_result = compute_kept_rate( + &editable_region_before_prediction, + &predicted_editable_region, + &settled_editable_region, + ); + + let sample_data = if can_collect_data + && let Some(sample_data) = sample_data + && let Ok(context) = sample_data.context_task.await + { + Some(SettledEditPredictionSampleData { + repository_url: context.repository_url, + revision: context.revision, + uncommitted_diff: context.uncommitted_diff, + editable_path: sample_data.editable_path, + editable_offset_range: sample_data.editable_offset_range, + buffer_diagnostics: context.buffer_diagnostics, + editable_context: context.editable_context, + future_edit_history_events: sample_data.future_edit_history_events, + navigation_history: sample_data + .navigation_history + .into_iter() + .map(|file| EditPredictionRecentFile { + path: file.path, + cursor_position: file.cursor_position, + }) + .collect(), + edit_events_before_quiescence: sample_data.edit_events_before_quiescence, + next_edit_cursor_offset: sample_data.next_edit_cursor_offset, + }) + } else { + None + }; + + batch.push(SettledEditPrediction { + request_id: request_id.0.to_string(), + settled_editable_region: can_collect_data.then_some(settled_editable_region), + ts_error_count_before_prediction, + ts_error_count_after_prediction, + can_collect_data, + is_in_open_source_repo, + sample_data, + kept_chars: EditPredictionSettledKeptChars { + candidate_new: kept_rate_result.candidate_new_chars, + reference_new: kept_rate_result.reference_new_chars, + candidate_deleted: kept_rate_result.candidate_deleted_chars, + reference_deleted: kept_rate_result.reference_deleted_chars, + kept: kept_rate_result.kept_chars, + correctly_deleted: kept_rate_result.correctly_deleted_chars, + discarded: kept_rate_result.discarded_chars, + context: kept_rate_result.context_chars, + kept_rate: kept_rate_result.kept_rate, + recall_rate: kept_rate_result.recall_rate, + }, + example: None, + model_version, + e2e_latency_ms: e2e_latency.as_millis().min(u128::from(u64::MAX)) as u64, + }); + } + + if batch.is_empty() { + break; + } + + let result = async { + let body = SubmitEditPredictionSettledBatchBody { predictions: batch }; + let compressed = zstd::encode_all(&serde_json::to_vec(&body)?[..], 3)?; + EditPredictionStore::send_api_request::( + |builder| { + Ok(builder + .uri(url.as_ref()) + .header("Content-Encoding", "zstd") + .body(compressed.clone().into())?) + }, + client.clone(), + llm_token.clone(), + organization_id.clone(), + app_version.clone(), + ) + .await?; + anyhow::Ok(()) + } + .await; + + if let Err(error) = result { + log::error!("failed to submit edit predictions settled: {error:?}"); + } + } + } +} + fn currently_following(project: &Entity, cx: &App) -> bool { let Some(app_state) = AppState::try_global(cx) else { return false; diff --git a/crates/edit_prediction/src/edit_prediction_tests.rs b/crates/edit_prediction/src/edit_prediction_tests.rs index 6c13139019..7e4ec8abef 100644 --- a/crates/edit_prediction/src/edit_prediction_tests.rs +++ b/crates/edit_prediction/src/edit_prediction_tests.rs @@ -4,7 +4,7 @@ use clock::ReplicaId; use cloud_api_types::{ CreateLlmTokenResponse, LlmToken, Organization, OrganizationConfiguration, OrganizationEditPredictionConfiguration, OrganizationId, SettledEditPrediction, - SubmitEditPredictionSettledResponse, + SubmitEditPredictionSettledBatchBody, SubmitEditPredictionSettledResponse, }; use cloud_llm_client::{ EditPredictionRejectReason, EditPredictionRejection, PredictEditsRequestTrigger, @@ -259,6 +259,68 @@ async fn test_simple_request(cx: &mut TestAppContext) { assert_eq!(prediction.edits[0].1.as_ref(), " are you?"); } +#[gpui::test] +async fn test_zeta_request_sends_settled_body_when_data_collection_is_disabled( + cx: &mut TestAppContext, +) { + let (ep_store, mut requests) = init_test_with_fake_client(cx); + let fs = FakeFs::new(cx.executor()); + fs.insert_tree( + "/root", + json!({ + "foo.md": "Hello!\nHow\nBye\n" + }), + ) + .await; + let project = Project::test(fs, vec![path!("/root").as_ref()], cx).await; + + let buffer = project + .update(cx, |project, cx| { + let path = project.find_project_path(path!("root/foo.md"), cx).unwrap(); + project.open_buffer(path, cx) + }) + .await + .unwrap(); + let snapshot = buffer.read_with(cx, |buffer, _cx| buffer.snapshot()); + let position = snapshot.anchor_before(language::Point::new(1, 3)); + + ep_store.update(cx, |ep_store, cx| { + ep_store.register_buffer(&buffer, &project, cx); + }); + + let prediction_task = ep_store.update(cx, |ep_store, cx| { + ep_store.request_prediction( + &project, + &buffer, + position, + PredictEditsRequestTrigger::Other, + cx, + ) + }); + + let (request, respond_tx) = requests.predict.next().await.unwrap(); + assert!(!request.input.can_collect_data); + respond_tx + .send(model_response(&request, SIMPLE_DIFF)) + .unwrap(); + + prediction_task.await.unwrap().unwrap(); + cx.run_until_parked(); + cx.executor() + .advance_clock(EDIT_PREDICTION_SETTLED_QUIESCENCE); + cx.run_until_parked(); + + let settled_request = requests + .settled + .next() + .now_or_never() + .flatten() + .expect("settled request should be sent"); + assert!(!settled_request.can_collect_data); + assert_eq!(settled_request.settled_editable_region, None); + assert_eq!(settled_request.sample_data, None); +} + #[gpui::test] async fn test_request_events(cx: &mut TestAppContext) { let (ep_store, mut requests) = init_test_with_fake_client(cx); @@ -1270,6 +1332,19 @@ async fn test_empty_prediction(cx: &mut TestAppContext) { e2e_latency_ms: Some(0), }] ); + cx.executor() + .advance_clock(EDIT_PREDICTION_SETTLED_QUIESCENCE); + cx.run_until_parked(); + + let settled_request = requests + .settled + .next() + .now_or_never() + .flatten() + .expect("empty prediction should still send settled request"); + assert_eq!(settled_request.request_id, id); + assert_eq!(settled_request.settled_editable_region, None); + assert_eq!(settled_request.sample_data, None); } #[gpui::test] @@ -2533,8 +2608,11 @@ fn init_test_with_fake_client_and_legacy_data_collection( } else { buf }; - let req = serde_json::from_slice(&body).unwrap(); - settled_req_tx.unbounded_send(req).unwrap(); + let req: SubmitEditPredictionSettledBatchBody = + serde_json::from_slice(&body).unwrap(); + for prediction in req.predictions { + settled_req_tx.unbounded_send(prediction).unwrap(); + } serde_json::to_string(&SubmitEditPredictionSettledResponse {}).unwrap() } _ => { diff --git a/crates/edit_prediction/src/zeta.rs b/crates/edit_prediction/src/zeta.rs index 8b388ef180..321ab27c02 100644 --- a/crates/edit_prediction/src/zeta.rs +++ b/crates/edit_prediction/src/zeta.rs @@ -3,8 +3,8 @@ use crate::{ EditPredictionId, EditPredictionModelInput, EditPredictionStartedDebugEvent, EditPredictionStore, PromptHistoryBoundary, ZedUpdateRequiredError, buffer_path_with_id_fallback, - capture_prediction_context::CapturedPredictionContext, cursor_excerpt::{self, compute_cursor_excerpt, compute_syntax_ranges}, + data_collection::CapturedPredictionContext, prediction::EditPredictionResult, }; use anyhow::Result; @@ -404,36 +404,34 @@ pub(crate) fn request_prediction_with_zeta( ) .await; - if can_collect_data { - let prediction = &result.prediction; - let weak_this = this.clone(); - let request_id = prediction.id.clone(); - let edited_buffer = edited_buffer.clone(); - let edited_buffer_snapshot = edited_buffer_snapshot.clone(); - let editable_range_in_buffer = editable_range_in_buffer.clone(); - let edit_preview = prediction.edit_preview.clone(); - let model_version = prediction.model_version.clone(); - cx.spawn(async move |cx| { - weak_this - .update(cx, |this, cx| { - this.enqueue_settled_prediction( - request_id.clone(), - &project, - &edited_buffer, - &edited_buffer_snapshot, - editable_range_in_buffer, - &edit_preview, - context_task, - prompt_history_boundary, - model_version, - request_duration, - cx, - ); - }) - .ok(); - }) - .detach(); - } + let prediction = &result.prediction; + let weak_this = this.clone(); + let request_id = prediction.id.clone(); + let edited_buffer = edited_buffer.clone(); + let edited_buffer_snapshot = edited_buffer_snapshot.clone(); + let editable_range_in_buffer = editable_range_in_buffer.clone(); + let edit_preview = prediction.edit_preview.clone(); + let model_version = prediction.model_version.clone(); + cx.spawn(async move |cx| { + weak_this + .update(cx, |this, cx| { + this.enqueue_settled_prediction( + request_id.clone(), + &project, + &edited_buffer, + &edited_buffer_snapshot, + editable_range_in_buffer, + &edit_preview, + context_task, + prompt_history_boundary, + model_version, + request_duration, + cx, + ); + }) + .ok(); + }) + .detach(); Ok(Some(result)) })