Fix handling of excerpt regions in EP CLI

crates/edit_prediction/src/capture_example.rs

@@ -1,10 +1,6 @@
 use crate::{
-    StoredEvent,
-    cursor_excerpt::editable_and_context_ranges_for_cursor_position,
-    example_spec::{
-        CapturedEvent, CapturedPromptInput, CapturedRelatedExcerpt, CapturedRelatedFile,
-        ExampleSpec, MAX_CURSOR_FILE_SIZE,
-    },
+    StoredEvent, cursor_excerpt::editable_and_context_ranges_for_cursor_position,
+    example_spec::ExampleSpec,
 };
 use anyhow::Result;
 use buffer_diff::BufferDiffSnapshot;
@@ -13,14 +9,13 @@ use gpui::{App, Entity, Task};
 use language::{Buffer, ToPoint as _};
 use project::{Project, WorktreeId};
 use std::{collections::hash_map, fmt::Write as _, ops::Range, path::Path, sync::Arc};
-use text::{BufferSnapshot as TextBufferSnapshot, Point, ToOffset as _};
+use text::{BufferSnapshot as TextBufferSnapshot, Point};
 
 pub fn capture_example(
     project: Entity<Project>,
     buffer: Entity<Buffer>,
     cursor_anchor: language::Anchor,
     mut events: Vec<StoredEvent>,
-    related_files: Vec<zeta_prompt::RelatedFile>,
     populate_expected_patch: bool,
     cx: &mut App,
 ) -> Option<Task<Result<ExampleSpec>>> {
@@ -60,14 +55,6 @@ pub fn capture_example(
             .map(|s| s.to_string())
             .unwrap_or_default();
 
-        let full_cursor_offset = cursor_anchor.to_offset(&snapshot);
-        let cursor_point = cursor_anchor.to_point(&snapshot);
-        let cursor_file_content = if snapshot.len() <= MAX_CURSOR_FILE_SIZE {
-            Some(snapshot.text())
-        } else {
-            None
-        };
-
         let (cursor_excerpt, cursor_offset_in_excerpt, cursor_excerpt_range) = cx
             .background_executor()
             .spawn(async move { compute_cursor_excerpt(&snapshot, cursor_anchor) })
@@ -109,56 +96,6 @@ pub fn capture_example(
             rejected_patch = Some(empty_patch);
         }
 
-        let prompt_input = cursor_file_content.map(|content| {
-            let captured_events: Vec<CapturedEvent> = events
-                .iter()
-                .map(|stored_event| {
-                    let zeta_prompt::Event::BufferChange {
-                        path,
-                        old_path,
-                        diff,
-                        predicted,
-                        in_open_source_repo,
-                    } = stored_event.event.as_ref();
-                    CapturedEvent {
-                        path: strip_root_name(path, &root_name).into(),
-                        old_path: strip_root_name(old_path, &root_name).into(),
-                        diff: diff.clone(),
-                        predicted: *predicted,
-                        in_open_source_repo: *in_open_source_repo,
-                    }
-                })
-                .collect();
-
-            let captured_related_files: Vec<CapturedRelatedFile> = related_files
-                .iter()
-                .map(|rf| CapturedRelatedFile {
-                    path: strip_root_name(&rf.path, &root_name).into(),
-                    max_row: rf.max_row,
-                    excerpts: rf
-                        .excerpts
-                        .iter()
-                        .map(|e| CapturedRelatedExcerpt {
-                            row_range: e.row_range.clone(),
-                            text: e.text.to_string(),
-                        })
-                        .collect(),
-                })
-                .collect();
-
-            CapturedPromptInput {
-                cursor_file_content: content,
-                cursor_offset: full_cursor_offset,
-                cursor_row: cursor_point.row,
-                cursor_column: cursor_point.column,
-                excerpt_start_row: Some(0),
-                events: captured_events,
-                related_files: captured_related_files,
-                in_open_source_repo: false,
-                zed_version: None,
-            }
-        });
-
         let mut spec = ExampleSpec {
             name: generate_timestamp_name(),
             repository_url,
@@ -171,7 +108,6 @@ pub fn capture_example(
             edit_history,
             expected_patches,
             rejected_patch,
-            captured_prompt_input: prompt_input,
             telemetry: None,
             human_feedback: Vec::new(),
             rating: None,
@@ -466,7 +402,6 @@ mod tests {
                     buffer.clone(),
                     Anchor::MIN,
                     events,
-                    Vec::new(),
                     true,
                     cx,
                 )
@@ -584,38 +519,11 @@ mod tests {
                     "}
                     .to_string()
                 ),
-                captured_prompt_input: example.captured_prompt_input.clone(),
                 telemetry: None,
                 human_feedback: Vec::new(),
                 rating: None,
             }
         );
-
-        let prompt_input = example
-            .captured_prompt_input
-            .expect("should have captured prompt input");
-        assert!(
-            prompt_input.cursor_file_content.contains("fn main()"),
-            "cursor_file_content should contain file content"
-        );
-        assert_eq!(
-            prompt_input.cursor_offset, 0,
-            "cursor at Anchor::MIN should be offset 0"
-        );
-        assert_eq!(
-            prompt_input.cursor_row, 0,
-            "cursor at Anchor::MIN should be row 0"
-        );
-        assert_eq!(
-            prompt_input.cursor_column, 0,
-            "cursor at Anchor::MIN should be column 0"
-        );
-        assert!(prompt_input.events.len() > 0, "should have captured events");
-        assert_eq!(
-            prompt_input.related_files.len(),
-            0,
-            "should have no related files (none passed)"
-        );
     }
 
     fn init_test(cx: &mut TestAppContext) {

crates/edit_prediction/src/example_spec.rs

@@ -1,7 +1,7 @@
 use crate::udiff::DiffLine;
 use anyhow::{Context as _, Result};
 use serde::{Deserialize, Serialize};
-use std::{borrow::Cow, fmt::Write as _, mem, ops::Range, path::Path, sync::Arc};
+use std::{borrow::Cow, fmt::Write as _, mem, path::Path, sync::Arc};
 use telemetry_events::EditPredictionRating;
 
 pub const CURSOR_POSITION_MARKER: &str = "[CURSOR_POSITION]";
@@ -81,8 +81,6 @@ pub struct ExampleSpec {
     #[serde(default, skip_serializing_if = "Option::is_none")]
     pub rejected_patch: Option<String>,
     #[serde(default, skip_serializing_if = "Option::is_none")]
-    pub captured_prompt_input: Option<CapturedPromptInput>,
-    #[serde(default, skip_serializing_if = "Option::is_none")]
     pub telemetry: Option<TelemetrySource>,
     #[serde(default, skip_serializing_if = "Vec::is_empty")]
     pub human_feedback: Vec<HumanFeedback>,
@@ -105,76 +103,6 @@ pub struct TelemetrySource {
     pub was_shown: bool,
 }
 
-/// All data needed to run format_prompt without loading the project.
-#[derive(Clone, Debug, PartialEq, Hash, Serialize, Deserialize)]
-pub struct CapturedPromptInput {
-    pub cursor_file_content: String,
-    pub cursor_offset: usize,
-    pub cursor_row: u32,
-    pub cursor_column: u32,
-    #[serde(default, skip_serializing_if = "Option::is_none")]
-    pub excerpt_start_row: Option<u32>,
-    pub events: Vec<CapturedEvent>,
-    pub related_files: Vec<CapturedRelatedFile>,
-    #[serde(default)]
-    pub in_open_source_repo: bool,
-    #[serde(default, skip_serializing_if = "Option::is_none")]
-    pub zed_version: Option<String>,
-}
-
-#[derive(Clone, Debug, PartialEq, Hash, Serialize, Deserialize)]
-pub struct CapturedEvent {
-    pub path: Arc<Path>,
-    pub old_path: Arc<Path>,
-    pub diff: String,
-    pub predicted: bool,
-    #[serde(default)]
-    pub in_open_source_repo: bool,
-}
-
-impl CapturedEvent {
-    pub fn to_event(&self) -> zeta_prompt::Event {
-        zeta_prompt::Event::BufferChange {
-            path: self.path.clone(),
-            old_path: self.old_path.clone(),
-            diff: self.diff.clone(),
-            predicted: self.predicted,
-            in_open_source_repo: self.in_open_source_repo,
-        }
-    }
-}
-
-#[derive(Clone, Debug, PartialEq, Hash, Serialize, Deserialize)]
-pub struct CapturedRelatedFile {
-    pub path: Arc<Path>,
-    pub max_row: u32,
-    pub excerpts: Vec<CapturedRelatedExcerpt>,
-}
-
-impl CapturedRelatedFile {
-    pub fn to_related_file(&self) -> zeta_prompt::RelatedFile {
-        zeta_prompt::RelatedFile {
-            path: self.path.clone(),
-            max_row: self.max_row,
-            in_open_source_repo: false,
-            excerpts: self
-                .excerpts
-                .iter()
-                .map(|e| zeta_prompt::RelatedExcerpt {
-                    row_range: e.row_range.clone(),
-                    text: e.text.clone().into(),
-                })
-                .collect(),
-        }
-    }
-}
-
-#[derive(Clone, Debug, PartialEq, Hash, Serialize, Deserialize)]
-pub struct CapturedRelatedExcerpt {
-    pub row_range: Range<u32>,
-    pub text: String,
-}
-
 const REASONING_HEADING: &str = "Reasoning";
 const UNCOMMITTED_DIFF_HEADING: &str = "Uncommitted Diff";
 const EDIT_HISTORY_HEADING: &str = "Edit History";
@@ -320,7 +248,6 @@ impl ExampleSpec {
             edit_history: String::new(),
             expected_patches: Vec::new(),
             rejected_patch: None,
-            captured_prompt_input: None,
             telemetry: None,
             human_feedback: Vec::new(),
             rating: None,
@@ -654,7 +581,6 @@ mod tests {
             edit_history: String::new(),
             expected_patches: Vec::new(),
             rejected_patch: None,
-            captured_prompt_input: None,
             telemetry: None,
             human_feedback: Vec::new(),
             rating: None,
@@ -791,7 +717,6 @@ mod tests {
             edit_history: String::new(),
             expected_patches: Vec::new(),
             rejected_patch: None,
-            captured_prompt_input: None,
             telemetry: None,
             human_feedback: Vec::new(),
             rating: None,
@@ -864,7 +789,6 @@ mod tests {
             edit_history: String::new(),
             expected_patches: Vec::new(),
             rejected_patch: None,
-            captured_prompt_input: None,
             telemetry: None,
             human_feedback: Vec::new(),
             rating: None,

crates/edit_prediction_cli/src/example.rs

@@ -15,9 +15,8 @@ use std::{
     collections::VecDeque,
     io::Read,
     path::{Path, PathBuf},
-    sync::Arc,
 };
-use zeta_prompt::RelatedFile;
+use zeta_prompt::ZetaPromptInput;
 
 #[derive(Clone, Debug, Serialize, Deserialize)]
 pub struct Example {
@@ -27,7 +26,7 @@ pub struct Example {
     /// 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 prompt_inputs: Option<ExamplePromptInputs>,
+    pub prompt_inputs: Option<ZetaPromptInput>,
 
     /// The input and expected output from the edit prediction model.
     #[serde(skip_serializing_if = "Option::is_none")]
@@ -46,6 +45,9 @@ pub struct Example {
     #[serde(default, skip_serializing_if = "Vec::is_empty")]
     pub qa: Vec<Option<QaResult>>,
 
+    /// The Zed version used to generate this example.
+    pub zed_version: Option<String>,
+
     /// The application state used to process this example.
     #[serde(skip)]
     pub state: Option<ExampleState>,
@@ -59,18 +61,6 @@ pub struct ExampleState {
     pub _open_buffers: OpenedBuffers,
 }
 
-#[derive(Clone, Debug, Serialize, Deserialize)]
-pub struct ExamplePromptInputs {
-    pub content: String,
-    pub cursor_row: u32,
-    pub cursor_column: u32,
-    pub cursor_offset: usize,
-    #[serde(default, skip_serializing_if = "Option::is_none")]
-    pub excerpt_start_row: Option<u32>,
-    pub edit_history: Vec<Arc<zeta_prompt::Event>>,
-    pub related_files: Option<Vec<RelatedFile>>,
-}
-
 #[derive(Clone, Debug, Serialize, Deserialize)]
 pub struct ExamplePrompt {
     pub input: String,
@@ -340,5 +330,6 @@ fn parse_markdown_example(input: &str) -> Result<Example> {
         score: Vec::new(),
         qa: Vec::new(),
         state: None,
+        zed_version: None,
     })
 }