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