1use crate::{
2 DebugEvent, EditPredictionFinishedDebugEvent, EditPredictionId, EditPredictionModelInput,
3 EditPredictionStartedDebugEvent, open_ai_response::text_from_response,
4 prediction::EditPredictionResult, zeta1::compute_edits,
5};
6use anyhow::{Context as _, Result};
7use edit_prediction_types::EditPredictionDismissReason;
8use futures::AsyncReadExt as _;
9use gpui::{
10 App, AppContext as _, Entity, Global, SharedString, Task,
11 http_client::{self, AsyncBody, HttpClient, Method},
12};
13use language::{OffsetRangeExt as _, ToOffset, ToPoint as _};
14use language_model::{ApiKeyState, EnvVar, env_var};
15use release_channel::AppVersion;
16use serde::Serialize;
17use std::{mem, ops::Range, path::Path, sync::Arc, time::Instant};
18
19use zeta_prompt::ZetaPromptInput;
20
21const MERCURY_API_URL: &str = "https://api.inceptionlabs.ai/v1/edit/completions";
22const MAX_REWRITE_TOKENS: usize = 150;
23const MAX_CONTEXT_TOKENS: usize = 350;
24
25pub struct Mercury {
26 pub api_token: Entity<ApiKeyState>,
27}
28
29impl Mercury {
30 pub fn new(cx: &mut App) -> Self {
31 Mercury {
32 api_token: mercury_api_token(cx),
33 }
34 }
35
36 pub(crate) fn request_prediction(
37 &self,
38 EditPredictionModelInput {
39 buffer,
40 snapshot,
41 position,
42 events,
43 related_files,
44 debug_tx,
45 ..
46 }: EditPredictionModelInput,
47 cx: &mut App,
48 ) -> Task<Result<Option<EditPredictionResult>>> {
49 self.api_token.update(cx, |key_state, cx| {
50 _ = key_state.load_if_needed(MERCURY_CREDENTIALS_URL, |s| s, cx);
51 });
52 let Some(api_token) = self.api_token.read(cx).key(&MERCURY_CREDENTIALS_URL) else {
53 return Task::ready(Ok(None));
54 };
55 let full_path: Arc<Path> = snapshot
56 .file()
57 .map(|file| file.full_path(cx))
58 .unwrap_or_else(|| "untitled".into())
59 .into();
60
61 let http_client = cx.http_client();
62 let cursor_point = position.to_point(&snapshot);
63 let buffer_snapshotted_at = Instant::now();
64 let active_buffer = buffer.clone();
65
66 let result = cx.background_spawn(async move {
67 let (editable_range, context_range) =
68 crate::cursor_excerpt::editable_and_context_ranges_for_cursor_position(
69 cursor_point,
70 &snapshot,
71 MAX_CONTEXT_TOKENS,
72 MAX_REWRITE_TOKENS,
73 );
74
75 let related_files = crate::filter_redundant_excerpts(
76 related_files,
77 full_path.as_ref(),
78 context_range.start.row..context_range.end.row,
79 );
80
81 let context_offset_range = context_range.to_offset(&snapshot);
82
83 let editable_offset_range = editable_range.to_offset(&snapshot);
84
85 let inputs = zeta_prompt::ZetaPromptInput {
86 events,
87 related_files,
88 cursor_offset_in_excerpt: cursor_point.to_offset(&snapshot)
89 - context_range.start.to_offset(&snapshot),
90 cursor_path: full_path.clone(),
91 cursor_excerpt: snapshot
92 .text_for_range(context_range)
93 .collect::<String>()
94 .into(),
95 editable_range_in_excerpt: (editable_offset_range.start
96 - context_offset_range.start)
97 ..(editable_offset_range.end - context_offset_range.start),
98 };
99
100 let prompt = build_prompt(&inputs);
101
102 if let Some(debug_tx) = &debug_tx {
103 debug_tx
104 .unbounded_send(DebugEvent::EditPredictionStarted(
105 EditPredictionStartedDebugEvent {
106 buffer: active_buffer.downgrade(),
107 prompt: Some(prompt.clone()),
108 position,
109 },
110 ))
111 .ok();
112 }
113
114 let request_body = open_ai::Request {
115 model: "mercury-coder".into(),
116 messages: vec![open_ai::RequestMessage::User {
117 content: open_ai::MessageContent::Plain(prompt),
118 }],
119 stream: false,
120 max_completion_tokens: None,
121 stop: vec![],
122 temperature: None,
123 tool_choice: None,
124 parallel_tool_calls: None,
125 tools: vec![],
126 prompt_cache_key: None,
127 reasoning_effort: None,
128 };
129
130 let buf = serde_json::to_vec(&request_body)?;
131 let body: AsyncBody = buf.into();
132
133 let request = http_client::Request::builder()
134 .uri(MERCURY_API_URL)
135 .header("Content-Type", "application/json")
136 .header("Authorization", format!("Bearer {}", api_token))
137 .header("Connection", "keep-alive")
138 .method(Method::POST)
139 .body(body)
140 .context("Failed to create request")?;
141
142 let mut response = http_client
143 .send(request)
144 .await
145 .context("Failed to send request")?;
146
147 let mut body: Vec<u8> = Vec::new();
148 response
149 .body_mut()
150 .read_to_end(&mut body)
151 .await
152 .context("Failed to read response body")?;
153
154 let response_received_at = Instant::now();
155 if !response.status().is_success() {
156 anyhow::bail!(
157 "Request failed with status: {:?}\nBody: {}",
158 response.status(),
159 String::from_utf8_lossy(&body),
160 );
161 };
162
163 let mut response: open_ai::Response =
164 serde_json::from_slice(&body).context("Failed to parse response")?;
165
166 let id = mem::take(&mut response.id);
167 let response_str = text_from_response(response).unwrap_or_default();
168
169 if let Some(debug_tx) = &debug_tx {
170 debug_tx
171 .unbounded_send(DebugEvent::EditPredictionFinished(
172 EditPredictionFinishedDebugEvent {
173 buffer: active_buffer.downgrade(),
174 model_output: Some(response_str.clone()),
175 position,
176 },
177 ))
178 .ok();
179 }
180
181 let response_str = response_str.strip_prefix("```\n").unwrap_or(&response_str);
182 let response_str = response_str.strip_suffix("\n```").unwrap_or(&response_str);
183
184 let mut edits = Vec::new();
185 const NO_PREDICTION_OUTPUT: &str = "None";
186
187 if response_str != NO_PREDICTION_OUTPUT {
188 let old_text = snapshot
189 .text_for_range(editable_offset_range.clone())
190 .collect::<String>();
191 edits = compute_edits(
192 old_text,
193 &response_str,
194 editable_offset_range.start,
195 &snapshot,
196 );
197 }
198
199 anyhow::Ok((id, edits, snapshot, response_received_at, inputs))
200 });
201
202 cx.spawn(async move |cx| {
203 let (id, edits, old_snapshot, response_received_at, inputs) =
204 result.await.context("Mercury edit prediction failed")?;
205 anyhow::Ok(Some(
206 EditPredictionResult::new(
207 EditPredictionId(id.into()),
208 &buffer,
209 &old_snapshot,
210 edits.into(),
211 buffer_snapshotted_at,
212 response_received_at,
213 inputs,
214 cx,
215 )
216 .await,
217 ))
218 })
219 }
220}
221
222fn build_prompt(inputs: &ZetaPromptInput) -> String {
223 const RECENTLY_VIEWED_SNIPPETS_START: &str = "<|recently_viewed_code_snippets|>\n";
224 const RECENTLY_VIEWED_SNIPPETS_END: &str = "<|/recently_viewed_code_snippets|>\n";
225 const RECENTLY_VIEWED_SNIPPET_START: &str = "<|recently_viewed_code_snippet|>\n";
226 const RECENTLY_VIEWED_SNIPPET_END: &str = "<|/recently_viewed_code_snippet|>\n";
227 const CURRENT_FILE_CONTENT_START: &str = "<|current_file_content|>\n";
228 const CURRENT_FILE_CONTENT_END: &str = "<|/current_file_content|>\n";
229 const CODE_TO_EDIT_START: &str = "<|code_to_edit|>\n";
230 const CODE_TO_EDIT_END: &str = "<|/code_to_edit|>\n";
231 const EDIT_DIFF_HISTORY_START: &str = "<|edit_diff_history|>\n";
232 const EDIT_DIFF_HISTORY_END: &str = "<|/edit_diff_history|>\n";
233 const CURSOR_TAG: &str = "<|cursor|>";
234 const CODE_SNIPPET_FILE_PATH_PREFIX: &str = "code_snippet_file_path: ";
235 const CURRENT_FILE_PATH_PREFIX: &str = "current_file_path: ";
236
237 let mut prompt = String::new();
238
239 push_delimited(
240 &mut prompt,
241 RECENTLY_VIEWED_SNIPPETS_START..RECENTLY_VIEWED_SNIPPETS_END,
242 |prompt| {
243 for related_file in inputs.related_files.iter() {
244 for related_excerpt in &related_file.excerpts {
245 push_delimited(
246 prompt,
247 RECENTLY_VIEWED_SNIPPET_START..RECENTLY_VIEWED_SNIPPET_END,
248 |prompt| {
249 prompt.push_str(CODE_SNIPPET_FILE_PATH_PREFIX);
250 prompt.push_str(related_file.path.to_string_lossy().as_ref());
251 prompt.push('\n');
252 prompt.push_str(related_excerpt.text.as_ref());
253 },
254 );
255 }
256 }
257 },
258 );
259
260 push_delimited(
261 &mut prompt,
262 CURRENT_FILE_CONTENT_START..CURRENT_FILE_CONTENT_END,
263 |prompt| {
264 prompt.push_str(CURRENT_FILE_PATH_PREFIX);
265 prompt.push_str(inputs.cursor_path.as_os_str().to_string_lossy().as_ref());
266 prompt.push('\n');
267
268 prompt.push_str(&inputs.cursor_excerpt[0..inputs.editable_range_in_excerpt.start]);
269 push_delimited(prompt, CODE_TO_EDIT_START..CODE_TO_EDIT_END, |prompt| {
270 prompt.push_str(
271 &inputs.cursor_excerpt
272 [inputs.editable_range_in_excerpt.start..inputs.cursor_offset_in_excerpt],
273 );
274 prompt.push_str(CURSOR_TAG);
275 prompt.push_str(
276 &inputs.cursor_excerpt
277 [inputs.cursor_offset_in_excerpt..inputs.editable_range_in_excerpt.end],
278 );
279 });
280 prompt.push_str(&inputs.cursor_excerpt[inputs.editable_range_in_excerpt.end..]);
281 },
282 );
283
284 push_delimited(
285 &mut prompt,
286 EDIT_DIFF_HISTORY_START..EDIT_DIFF_HISTORY_END,
287 |prompt| {
288 for event in inputs.events.iter() {
289 zeta_prompt::write_event(prompt, &event);
290 }
291 },
292 );
293
294 prompt
295}
296
297fn push_delimited(prompt: &mut String, delimiters: Range<&str>, cb: impl FnOnce(&mut String)) {
298 prompt.push_str(delimiters.start);
299 cb(prompt);
300 prompt.push('\n');
301 prompt.push_str(delimiters.end);
302}
303
304pub const MERCURY_CREDENTIALS_URL: SharedString =
305 SharedString::new_static("https://api.inceptionlabs.ai/v1/edit/completions");
306pub const MERCURY_CREDENTIALS_USERNAME: &str = "mercury-api-token";
307pub static MERCURY_TOKEN_ENV_VAR: std::sync::LazyLock<EnvVar> = env_var!("MERCURY_AI_TOKEN");
308
309struct GlobalMercuryApiKey(Entity<ApiKeyState>);
310
311impl Global for GlobalMercuryApiKey {}
312
313pub fn mercury_api_token(cx: &mut App) -> Entity<ApiKeyState> {
314 if let Some(global) = cx.try_global::<GlobalMercuryApiKey>() {
315 return global.0.clone();
316 }
317 let entity =
318 cx.new(|_| ApiKeyState::new(MERCURY_CREDENTIALS_URL, MERCURY_TOKEN_ENV_VAR.clone()));
319 cx.set_global(GlobalMercuryApiKey(entity.clone()));
320 entity
321}
322
323pub fn load_mercury_api_token(cx: &mut App) -> Task<Result<(), language_model::AuthenticateError>> {
324 mercury_api_token(cx).update(cx, |key_state, cx| {
325 key_state.load_if_needed(MERCURY_CREDENTIALS_URL, |s| s, cx)
326 })
327}
328
329const FEEDBACK_API_URL: &str = "https://api-feedback.inceptionlabs.ai/feedback";
330
331#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
332#[serde(rename_all = "snake_case")]
333enum MercuryUserAction {
334 Accept,
335 Reject,
336 Ignore,
337}
338
339#[derive(Serialize)]
340struct FeedbackRequest {
341 request_id: SharedString,
342 provider_name: &'static str,
343 user_action: MercuryUserAction,
344 provider_version: String,
345}
346
347pub(crate) fn edit_prediction_accepted(
348 prediction_id: EditPredictionId,
349 http_client: Arc<dyn HttpClient>,
350 cx: &App,
351) {
352 send_feedback(prediction_id, MercuryUserAction::Accept, http_client, cx);
353}
354
355pub(crate) fn edit_prediction_rejected(
356 prediction_id: EditPredictionId,
357 was_shown: bool,
358 dismiss_reason: EditPredictionDismissReason,
359 http_client: Arc<dyn HttpClient>,
360 cx: &App,
361) {
362 if !was_shown {
363 return;
364 }
365 let action = match dismiss_reason {
366 EditPredictionDismissReason::Rejected => MercuryUserAction::Reject,
367 EditPredictionDismissReason::Ignored => MercuryUserAction::Ignore,
368 };
369 send_feedback(prediction_id, action, http_client, cx);
370}
371
372fn send_feedback(
373 prediction_id: EditPredictionId,
374 action: MercuryUserAction,
375 http_client: Arc<dyn HttpClient>,
376 cx: &App,
377) {
378 let request_id = prediction_id.0;
379 let app_version = AppVersion::global(cx);
380 cx.background_spawn(async move {
381 let body = FeedbackRequest {
382 request_id,
383 provider_name: "zed",
384 user_action: action,
385 provider_version: app_version.to_string(),
386 };
387
388 let request = http_client::Request::builder()
389 .uri(FEEDBACK_API_URL)
390 .method(Method::POST)
391 .header("Content-Type", "application/json")
392 .body(AsyncBody::from(serde_json::to_vec(&body)?))?;
393
394 let response = http_client.send(request).await?;
395 if !response.status().is_success() {
396 anyhow::bail!("Feedback API returned status: {}", response.status());
397 }
398
399 log::debug!(
400 "Mercury feedback sent: request_id={}, action={:?}",
401 body.request_id,
402 body.user_action
403 );
404
405 anyhow::Ok(())
406 })
407 .detach_and_log_err(cx);
408}