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