1use std::{cmp, sync::Arc};
2
3use client::{Client, UserStore};
4use cloud_llm_client::EditPredictionRejectReason;
5use edit_prediction_types::{
6 DataCollectionState, EditPredictionDelegate, EditPredictionDismissReason,
7 EditPredictionIconSet, SuggestionDisplayType,
8};
9use gpui::{App, Entity, prelude::*};
10use language::{Buffer, ToPoint as _};
11use project::Project;
12use ui::prelude::*;
13
14use crate::{BufferEditPrediction, EditPredictionModel, EditPredictionStore};
15
16pub struct ZedEditPredictionDelegate {
17 store: Entity<EditPredictionStore>,
18 project: Entity<Project>,
19 singleton_buffer: Option<Entity<Buffer>>,
20}
21
22impl ZedEditPredictionDelegate {
23 pub fn new(
24 project: Entity<Project>,
25 singleton_buffer: Option<Entity<Buffer>>,
26 client: &Arc<Client>,
27 user_store: &Entity<UserStore>,
28 cx: &mut Context<Self>,
29 ) -> Self {
30 let store = EditPredictionStore::global(client, user_store, cx);
31 store.update(cx, |store, cx| {
32 store.register_project(&project, cx);
33 });
34
35 cx.observe(&store, |_this, _ep_store, cx| {
36 cx.notify();
37 })
38 .detach();
39
40 Self {
41 project: project,
42 store: store,
43 singleton_buffer,
44 }
45 }
46}
47
48impl EditPredictionDelegate for ZedEditPredictionDelegate {
49 fn name() -> &'static str {
50 "zed-predict"
51 }
52
53 fn display_name() -> &'static str {
54 "Zed's Edit Predictions"
55 }
56
57 fn show_predictions_in_menu() -> bool {
58 true
59 }
60
61 fn show_tab_accept_marker() -> bool {
62 true
63 }
64
65 fn icons(&self, cx: &App) -> EditPredictionIconSet {
66 match self.store.read(cx).edit_prediction_model {
67 EditPredictionModel::Ollama => EditPredictionIconSet::new(IconName::AiOllama),
68 EditPredictionModel::Sweep => EditPredictionIconSet::new(IconName::SweepAi)
69 .with_disabled(IconName::SweepAiDisabled)
70 .with_up(IconName::SweepAiUp)
71 .with_down(IconName::SweepAiDown)
72 .with_error(IconName::SweepAiError),
73 EditPredictionModel::Mercury => EditPredictionIconSet::new(IconName::Inception),
74 EditPredictionModel::Zeta1 | EditPredictionModel::Zeta2 { .. } => {
75 EditPredictionIconSet::new(IconName::ZedPredict)
76 .with_disabled(IconName::ZedPredictDisabled)
77 .with_up(IconName::ZedPredictUp)
78 .with_down(IconName::ZedPredictDown)
79 .with_error(IconName::ZedPredictError)
80 }
81 }
82 }
83
84 fn data_collection_state(&self, cx: &App) -> DataCollectionState {
85 if let Some(buffer) = &self.singleton_buffer
86 && let Some(file) = buffer.read(cx).file()
87 {
88 let is_project_open_source =
89 self.store
90 .read(cx)
91 .is_file_open_source(&self.project, file, cx);
92 if self.store.read(cx).data_collection_choice.is_enabled(cx) {
93 DataCollectionState::Enabled {
94 is_project_open_source,
95 }
96 } else {
97 DataCollectionState::Disabled {
98 is_project_open_source,
99 }
100 }
101 } else {
102 return DataCollectionState::Disabled {
103 is_project_open_source: false,
104 };
105 }
106 }
107
108 fn toggle_data_collection(&mut self, cx: &mut App) {
109 self.store.update(cx, |store, cx| {
110 store.toggle_data_collection_choice(cx);
111 });
112 }
113
114 fn usage(&self, cx: &App) -> Option<client::EditPredictionUsage> {
115 self.store.read(cx).usage(cx)
116 }
117
118 fn is_enabled(
119 &self,
120 _buffer: &Entity<language::Buffer>,
121 _cursor_position: language::Anchor,
122 cx: &App,
123 ) -> bool {
124 let store = self.store.read(cx);
125 if store.edit_prediction_model == EditPredictionModel::Sweep {
126 store.has_sweep_api_token(cx)
127 } else {
128 true
129 }
130 }
131
132 fn is_refreshing(&self, cx: &App) -> bool {
133 self.store.read(cx).is_refreshing(&self.project)
134 }
135
136 fn refresh(
137 &mut self,
138 buffer: Entity<language::Buffer>,
139 cursor_position: language::Anchor,
140 _debounce: bool,
141 cx: &mut Context<Self>,
142 ) {
143 let store = self.store.read(cx);
144
145 if store.user_store.read_with(cx, |user_store, _cx| {
146 user_store.account_too_young() || user_store.has_overdue_invoices()
147 }) {
148 return;
149 }
150
151 self.store.update(cx, |store, cx| {
152 if let Some(current) =
153 store.prediction_at(&buffer, Some(cursor_position), &self.project, cx)
154 && let BufferEditPrediction::Local { prediction } = current
155 && prediction.interpolate(buffer.read(cx)).is_some()
156 {
157 return;
158 }
159
160 store.refresh_context(&self.project, &buffer, cursor_position, cx);
161 store.refresh_prediction_from_buffer(self.project.clone(), buffer, cursor_position, cx)
162 });
163 }
164
165 fn accept(&mut self, cx: &mut Context<Self>) {
166 self.store.update(cx, |store, cx| {
167 store.accept_current_prediction(&self.project, cx);
168 });
169 }
170
171 fn discard(&mut self, reason: EditPredictionDismissReason, cx: &mut Context<Self>) {
172 self.store.update(cx, |store, cx| {
173 store.reject_current_prediction(
174 EditPredictionRejectReason::Discarded,
175 &self.project,
176 reason,
177 cx,
178 );
179 });
180 }
181
182 fn did_show(&mut self, display_type: SuggestionDisplayType, cx: &mut Context<Self>) {
183 self.store.update(cx, |store, cx| {
184 store.did_show_current_prediction(&self.project, display_type, cx);
185 });
186 }
187
188 fn suggest(
189 &mut self,
190 buffer: &Entity<language::Buffer>,
191 cursor_position: language::Anchor,
192 cx: &mut Context<Self>,
193 ) -> Option<edit_prediction_types::EditPrediction> {
194 self.store.update(cx, |store, cx| {
195 let prediction =
196 store.prediction_at(buffer, Some(cursor_position), &self.project, cx)?;
197
198 let prediction = match prediction {
199 BufferEditPrediction::Local { prediction } => prediction,
200 BufferEditPrediction::Jump { prediction } => {
201 return Some(edit_prediction_types::EditPrediction::Jump {
202 id: Some(prediction.id.to_string().into()),
203 snapshot: prediction.snapshot.clone(),
204 target: prediction.edits.first().unwrap().0.start,
205 });
206 }
207 };
208
209 let buffer = buffer.read(cx);
210 let snapshot = buffer.snapshot();
211
212 let Some(edits) = prediction.interpolate(&snapshot) else {
213 store.reject_current_prediction(
214 EditPredictionRejectReason::InterpolatedEmpty,
215 &self.project,
216 EditPredictionDismissReason::Ignored,
217 cx,
218 );
219 return None;
220 };
221
222 let cursor_row = cursor_position.to_point(&snapshot).row;
223 let (closest_edit_ix, (closest_edit_range, _)) =
224 edits.iter().enumerate().min_by_key(|(_, (range, _))| {
225 let distance_from_start =
226 cursor_row.abs_diff(range.start.to_point(&snapshot).row);
227 let distance_from_end = cursor_row.abs_diff(range.end.to_point(&snapshot).row);
228 cmp::min(distance_from_start, distance_from_end)
229 })?;
230
231 let mut edit_start_ix = closest_edit_ix;
232 for (range, _) in edits[..edit_start_ix].iter().rev() {
233 let distance_from_closest_edit = closest_edit_range.start.to_point(&snapshot).row
234 - range.end.to_point(&snapshot).row;
235 if distance_from_closest_edit <= 1 {
236 edit_start_ix -= 1;
237 } else {
238 break;
239 }
240 }
241
242 let mut edit_end_ix = closest_edit_ix + 1;
243 for (range, _) in &edits[edit_end_ix..] {
244 let distance_from_closest_edit = range.start.to_point(buffer).row
245 - closest_edit_range.end.to_point(&snapshot).row;
246 if distance_from_closest_edit <= 1 {
247 edit_end_ix += 1;
248 } else {
249 break;
250 }
251 }
252
253 Some(edit_prediction_types::EditPrediction::Local {
254 id: Some(prediction.id.to_string().into()),
255 edits: edits[edit_start_ix..edit_end_ix].to_vec(),
256 edit_preview: Some(prediction.edit_preview.clone()),
257 })
258 })
259 }
260}