tenant.openagents/omega

1use std::{ops::Range, sync::Arc};
2
3use cloud_llm_client::{EditPredictionRejectReason, PredictEditsRequestTrigger};
4use edit_prediction_types::{PredictedCursorPosition, interpolate_edits};
5use gpui::{AsyncApp, Entity, SharedString};
6use language::{Anchor, Buffer, BufferSnapshot, EditPreview, TextBufferSnapshot};
7use zeta_prompt::{Zeta2PromptInput, Zeta3PromptInput};
8
9#[derive(Clone, Default, Debug, PartialEq, Eq, Hash)]
10pub struct EditPredictionId(pub SharedString);
11
12impl From<EditPredictionId> for gpui::ElementId {
13    fn from(value: EditPredictionId) -> Self {
14        gpui::ElementId::Name(value.0)
15    }
16}
17
18impl std::fmt::Display for EditPredictionId {
19    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
20        write!(f, "{}", self.0)
21    }
22}
23
24#[derive(Clone, serde::Serialize)]
25#[serde(tag = "version", content = "input")]
26pub enum EditPredictionInputs {
27    V2(Zeta2PromptInput),
28    V3(Zeta3PromptInput),
29}
30
31/// A prediction response that was returned from the provider, whether it was ultimately valid or not.
32pub struct EditPredictionResult {
33    pub prediction: EditPrediction,
34    pub reject_reason: Option<EditPredictionRejectReason>,
35    pub e2e_latency: std::time::Duration,
36}
37
38impl EditPredictionResult {
39    pub async fn new(
40        id: EditPredictionId,
41        edited_buffer: &Entity<Buffer>,
42        edited_buffer_snapshot: &BufferSnapshot,
43        edits: Arc<[(Range<Anchor>, Arc<str>)]>,
44        cursor_position: Option<PredictedCursorPosition>,
45        editable_range: Option<Range<Anchor>>,
46        inputs: EditPredictionInputs,
47        model_version: Option<String>,
48        trigger: PredictEditsRequestTrigger,
49        e2e_latency: std::time::Duration,
50        cx: &mut AsyncApp,
51    ) -> Self {
52        let (edits, reject_reason, new_snapshot): (
53            Arc<[(Range<Anchor>, Arc<str>)]>,
54            Option<EditPredictionRejectReason>,
55            Option<BufferSnapshot>,
56        ) = if edits.is_empty() {
57            (
58                Arc::default(),
59                Some(EditPredictionRejectReason::Empty),
60                None,
61            )
62        } else {
63            edited_buffer.read_with(cx, |buffer, _cx| {
64                let new_snapshot = buffer.snapshot();
65                match interpolate_edits(&edited_buffer_snapshot, &new_snapshot, &edits) {
66                    Some(edits) if edits.is_empty() => (
67                        Arc::default(),
68                        Some(EditPredictionRejectReason::InterpolatedEmpty),
69                        None,
70                    ),
71                    Some(edits) => (Arc::from(edits), None, Some(new_snapshot)),
72                    None => (
73                        Arc::default(),
74                        Some(EditPredictionRejectReason::InterpolateFailed),
75                        None,
76                    ),
77                }
78            })
79        };
80        let snapshot = new_snapshot.unwrap_or_else(|| edited_buffer_snapshot.clone());
81
82        let edit_preview = if !edits.is_empty() {
83            edited_buffer
84                .read_with(cx, |buffer, cx| buffer.preview_edits(edits.clone(), cx))
85                .await
86        } else {
87            EditPreview::unchanged(edited_buffer_snapshot)
88        };
89
90        Self {
91            prediction: EditPrediction {
92                id,
93                edits,
94                cursor_position,
95                editable_range,
96                snapshot,
97                edit_preview,
98                inputs,
99                buffer: edited_buffer.clone(),
100                model_version,
101                trigger,
102            },
103            reject_reason,
104            e2e_latency,
105        }
106    }
107
108    pub fn new_rejected(
109        id: EditPredictionId,
110        edited_buffer: &Entity<Buffer>,
111        edited_buffer_snapshot: &BufferSnapshot,
112        inputs: EditPredictionInputs,
113        model_version: Option<String>,
114        trigger: PredictEditsRequestTrigger,
115        e2e_latency: std::time::Duration,
116        reject_reason: EditPredictionRejectReason,
117    ) -> Self {
118        Self {
119            prediction: EditPrediction {
120                id,
121                edits: Arc::default(),
122                cursor_position: None,
123                editable_range: None,
124                snapshot: edited_buffer_snapshot.clone(),
125                edit_preview: EditPreview::unchanged(edited_buffer_snapshot),
126                inputs,
127                buffer: edited_buffer.clone(),
128                model_version,
129                trigger,
130            },
131            reject_reason: Some(reject_reason),
132            e2e_latency,
133        }
134    }
135}
136
137#[derive(Clone)]
138pub struct EditPrediction {
139    pub id: EditPredictionId,
140    pub edits: Arc<[(Range<Anchor>, Arc<str>)]>,
141    pub cursor_position: Option<PredictedCursorPosition>,
142    pub editable_range: Option<Range<Anchor>>,
143    pub snapshot: BufferSnapshot,
144    pub edit_preview: EditPreview,
145    pub inputs: EditPredictionInputs,
146    pub buffer: Entity<Buffer>,
147    pub model_version: Option<String>,
148    pub trigger: PredictEditsRequestTrigger,
149}
150
151impl EditPrediction {
152    pub fn interpolate(
153        &self,
154        new_snapshot: &TextBufferSnapshot,
155    ) -> Option<Vec<(Range<Anchor>, Arc<str>)>> {
156        interpolate_edits(&self.snapshot, new_snapshot, &self.edits)
157    }
158
159    pub fn targets_buffer(&self, buffer: &Buffer) -> bool {
160        self.snapshot.remote_id() == buffer.remote_id()
161    }
162}
163
164impl std::fmt::Debug for EditPrediction {
165    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
166        f.debug_struct("EditPrediction")
167            .field("id", &self.id)
168            .field("edits", &self.edits)
169            .finish()
170    }
171}
172
173#[cfg(test)]
174mod tests {
175    use std::path::Path;
176
177    use super::*;
178    use gpui::{App, Entity, TestAppContext, prelude::*};
179    use language::{Buffer, ToOffset as _};
180    use zeta_prompt::Zeta2PromptInput;
181
182    #[gpui::test]
183    async fn test_edit_prediction_basic_interpolation(cx: &mut TestAppContext) {
184        let buffer = cx.new(|cx| Buffer::local("Lorem ipsum dolor", cx));
185        let edits: Arc<[(Range<Anchor>, Arc<str>)]> = cx.update(|cx| {
186            to_prediction_edits([(2..5, "REM".into()), (9..11, "".into())], &buffer, cx).into()
187        });
188
189        let edit_preview = cx
190            .read(|cx| buffer.read(cx).preview_edits(edits.clone(), cx))
191            .await;
192
193        let prediction = EditPrediction {
194            id: EditPredictionId("prediction-1".into()),
195            edits,
196            cursor_position: None,
197            editable_range: None,
198            snapshot: cx.read(|cx| buffer.read(cx).snapshot()),
199            buffer: buffer.clone(),
200            edit_preview,
201            model_version: None,
202            trigger: PredictEditsRequestTrigger::Other,
203            inputs: EditPredictionInputs::V2(Zeta2PromptInput {
204                events: vec![],
205                related_files: Some(vec![]),
206                active_buffer_diagnostics: vec![],
207                cursor_path: Path::new("path.txt").into(),
208                cursor_offset_in_excerpt: 0,
209                cursor_excerpt: "".into(),
210                excerpt_start_row: None,
211                excerpt_ranges: Default::default(),
212                syntax_ranges: None,
213                in_open_source_repo: false,
214                can_collect_data: false,
215                repo_url: None,
216            }),
217        };
218
219        cx.update(|cx| {
220            assert_eq!(
221                from_prediction_edits(
222                    &prediction.interpolate(&buffer.read(cx).snapshot()).unwrap(),
223                    &buffer,
224                    cx
225                ),
226                vec![(2..5, "REM".into()), (9..11, "".into())]
227            );
228
229            buffer.update(cx, |buffer, cx| buffer.edit([(2..5, "")], None, cx));
230            assert_eq!(
231                from_prediction_edits(
232                    &prediction.interpolate(&buffer.read(cx).snapshot()).unwrap(),
233                    &buffer,
234                    cx
235                ),
236                vec![(2..2, "REM".into()), (6..8, "".into())]
237            );
238
239            buffer.update(cx, |buffer, cx| buffer.undo(cx));
240            assert_eq!(
241                from_prediction_edits(
242                    &prediction.interpolate(&buffer.read(cx).snapshot()).unwrap(),
243                    &buffer,
244                    cx
245                ),
246                vec![(2..5, "REM".into()), (9..11, "".into())]
247            );
248
249            buffer.update(cx, |buffer, cx| buffer.edit([(2..5, "R")], None, cx));
250            assert_eq!(
251                from_prediction_edits(
252                    &prediction.interpolate(&buffer.read(cx).snapshot()).unwrap(),
253                    &buffer,
254                    cx
255                ),
256                vec![(3..3, "EM".into()), (7..9, "".into())]
257            );
258
259            buffer.update(cx, |buffer, cx| buffer.edit([(3..3, "E")], None, cx));
260            assert_eq!(
261                from_prediction_edits(
262                    &prediction.interpolate(&buffer.read(cx).snapshot()).unwrap(),
263                    &buffer,
264                    cx
265                ),
266                vec![(4..4, "M".into()), (8..10, "".into())]
267            );
268
269            buffer.update(cx, |buffer, cx| buffer.edit([(4..4, "M")], None, cx));
270            assert_eq!(
271                from_prediction_edits(
272                    &prediction.interpolate(&buffer.read(cx).snapshot()).unwrap(),
273                    &buffer,
274                    cx
275                ),
276                vec![(9..11, "".into())]
277            );
278
279            buffer.update(cx, |buffer, cx| buffer.edit([(4..5, "")], None, cx));
280            assert_eq!(
281                from_prediction_edits(
282                    &prediction.interpolate(&buffer.read(cx).snapshot()).unwrap(),
283                    &buffer,
284                    cx
285                ),
286                vec![(4..4, "M".into()), (8..10, "".into())]
287            );
288
289            buffer.update(cx, |buffer, cx| buffer.edit([(8..10, "")], None, cx));
290            assert_eq!(
291                from_prediction_edits(
292                    &prediction.interpolate(&buffer.read(cx).snapshot()).unwrap(),
293                    &buffer,
294                    cx
295                ),
296                vec![(4..4, "M".into())]
297            );
298
299            buffer.update(cx, |buffer, cx| buffer.edit([(4..6, "")], None, cx));
300            assert_eq!(prediction.interpolate(&buffer.read(cx).snapshot()), None);
301        })
302    }
303
304    fn to_prediction_edits(
305        iterator: impl IntoIterator<Item = (Range<usize>, Arc<str>)>,
306        buffer: &Entity<Buffer>,
307        cx: &App,
308    ) -> Vec<(Range<Anchor>, Arc<str>)> {
309        let buffer = buffer.read(cx);
310        iterator
311            .into_iter()
312            .map(|(range, text)| {
313                (
314                    buffer.anchor_after(range.start)..buffer.anchor_before(range.end),
315                    text,
316                )
317            })
318            .collect()
319    }
320
321    fn from_prediction_edits(
322        editor_edits: &[(Range<Anchor>, Arc<str>)],
323        buffer: &Entity<Buffer>,
324        cx: &App,
325    ) -> Vec<(Range<usize>, Arc<str>)> {
326        let buffer = buffer.read(cx);
327        editor_edits
328            .iter()
329            .map(|(range, text)| {
330                (
331                    range.start.to_offset(buffer)..range.end.to_offset(buffer),
332                    text.clone(),
333                )
334            })
335            .collect()
336    }
337}
338