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