Skip to repository content987 lines · 39.8 KB · rust
tenant.openagents/omega
No repository description is available.
OpenAgents Git authority 2026-07-28T03:59:56.235Z Public web read
NIP-34 coordinate
30617:7649603503856e5148d571eac2766b288a8ff1e9e35d380337a1d2b0015b4f92:omegaMaintainersHidden in public view
References2 branches · 1 tag
Read-only clone
git clone https://openagents.com/git/tenant.openagents/omega.gitBrowse files
zeta.rs
1use crate::{
2 CloudRequestTimeoutError, CurrentEditPrediction, DebugEvent, EditPredictionFinishedDebugEvent,
3 EditPredictionId, EditPredictionInputs, EditPredictionModelInput,
4 EditPredictionStartedDebugEvent, EditPredictionStore, PromptHistoryBoundary,
5 ZedUpdateRequiredError, buffer_path_with_id_fallback,
6 cursor_excerpt::{self, compute_cursor_excerpt, compute_syntax_ranges},
7 data_collection::CapturedPredictionContext,
8 prediction::EditPredictionResult,
9 udiff::prediction_edits_for_single_file_diff,
10};
11use anyhow::{Context as _, Result};
12use cloud_llm_client::{
13 AcceptEditPredictionBody, EditPredictionRejectReason, predict_edits_v3::RawCompletionRequest,
14};
15use edit_prediction_types::PredictedCursorPosition;
16use gpui::{App, AppContext as _, Entity, Task, TaskExt, WeakEntity, prelude::*};
17use language::{
18 Buffer, BufferSnapshot, DiagnosticSeverity, EditPredictionPromptFormat, OffsetRangeExt as _,
19 ToOffset as _, ZetaVersion, language_settings::all_language_settings, text_diff,
20};
21use project::Project;
22use release_channel::AppVersion;
23use text::{Anchor, Bias, Point};
24use ui::SharedString;
25use workspace::notifications::simple_message_notification::MessageNotification;
26use workspace::notifications::{NotificationId, show_app_notification};
27use workspace::workspace_error::{ErrorAction, ErrorSeverity, WorkspaceError};
28
29use std::{ops::Range, path::Path, sync::Arc};
30use zeta_prompt::{
31 FilePosition, ParsedOutput, Zeta2PromptInput, Zeta3PromptInput, ZetaFormat,
32 excerpt_ranges_for_format, format_zeta_prompt, get_prefill, parse_zeta2_model_output,
33 stop_tokens_for_format,
34 zeta1::{self, EDITABLE_REGION_END_MARKER},
35};
36
37use crate::open_ai_compatible::{
38 load_open_ai_compatible_api_key_if_needed, send_custom_server_request,
39};
40
41pub(crate) fn request_prediction_with_zeta(
42 store: &mut EditPredictionStore,
43 EditPredictionModelInput {
44 buffer,
45 snapshot,
46 position,
47 related_files,
48 editable_context,
49 events,
50 debug_tx,
51 mode,
52 trigger,
53 project,
54 diagnostic_search_range,
55 can_collect_data,
56 is_open_source,
57 allow_jump,
58 ..
59 }: EditPredictionModelInput,
60 context_task: Option<Task<Result<CapturedPredictionContext>>>,
61 prompt_history_boundary: Option<PromptHistoryBoundary>,
62 repo_url: Option<String>,
63 cx: &mut Context<EditPredictionStore>,
64) -> Task<Result<Option<EditPredictionResult>>> {
65 let settings = &all_language_settings(None, cx).edit_predictions;
66 let provider = settings.provider;
67 let custom_server_settings = match provider {
68 settings::EditPredictionProvider::Ollama => settings.ollama.clone(),
69 settings::EditPredictionProvider::OpenAiCompatibleApi => {
70 settings.open_ai_compatible_api.clone()
71 }
72 _ => None,
73 };
74
75 let http_client = cx.http_client();
76 let request_start = cx.background_executor().now();
77 let raw_config = store.zeta2_raw_config().cloned();
78 let preferred_experiment = store.preferred_experiment().map(|s| s.to_owned());
79 let open_ai_compatible_api_key = load_open_ai_compatible_api_key_if_needed(provider, cx);
80
81 let excerpt_path = buffer_path_with_id_fallback(snapshot.file(), &snapshot.text, cx);
82
83 let repo_url = repo_url.filter(|_| can_collect_data);
84 let client = store.client.clone();
85 let llm_token = store.llm_token.clone();
86 let organization_id = store
87 .user_store
88 .read(cx)
89 .current_organization()
90 .map(|organization| organization.id.clone());
91 let app_version = AppVersion::global(cx);
92
93 enum PendingPrediction {
94 Resolved {
95 inputs: EditPredictionInputs,
96 buffer: Entity<Buffer>,
97 snapshot: BufferSnapshot,
98 edits: Vec<(Range<Anchor>, Arc<str>)>,
99 cursor_position: Option<PredictedCursorPosition>,
100 editable_range_in_buffer: Option<Range<usize>>,
101 },
102 Patch {
103 inputs: EditPredictionInputs,
104 project: Entity<Project>,
105 buffer: Entity<Buffer>,
106 snapshot: BufferSnapshot,
107 patch: String,
108 },
109 }
110
111 let project_for_request = project.clone();
112 let request_task = cx.background_spawn({
113 async move {
114 let local_zeta_version = custom_server_settings
115 .as_ref()
116 .and_then(|settings| match settings.prompt_format {
117 EditPredictionPromptFormat::Zeta(version) => Some(version),
118 EditPredictionPromptFormat::Infer => {
119 match settings.model.to_ascii_lowercase().as_str() {
120 "zeta" | "zeta1" => Some(ZetaVersion::Zeta1),
121 "zeta2" => Some(ZetaVersion::Zeta2),
122 "zeta2.1" => Some(ZetaVersion::Zeta2_1),
123 _ => None,
124 }
125 }
126 _ => None,
127 })
128 .unwrap_or_default();
129 let zeta_format = raw_config
130 .as_ref()
131 .map(|config| config.format)
132 .or(match local_zeta_version {
133 ZetaVersion::Zeta1 => None,
134 ZetaVersion::Zeta2 => Some(ZetaFormat::V0211SeedCoder),
135 ZetaVersion::Zeta2_1 => Some(ZetaFormat::V0318SeedMultiRegions),
136 })
137 .unwrap_or_default();
138
139 enum RequestInput {
140 V3 {
141 full_context_offset_range: Range<usize>,
142 prompt_input: Zeta2PromptInput,
143 formatted_prompt: Option<String>,
144 },
145 V4(Zeta3PromptInput),
146 }
147
148 let cursor_offset = position.to_offset(&snapshot);
149 let request_input =
150 if allow_jump && custom_server_settings.is_none() && raw_config.is_none() {
151 let cursor_point = snapshot.offset_to_point(cursor_offset);
152 let editable_context = editable_context
153 .context("missing editable context task")?
154 .await?;
155 let active_buffer_diagnostics = active_buffer_diagnostics(
156 &snapshot,
157 diagnostic_search_range,
158 cursor_point.row,
159 ACTIVE_BUFFER_DIAGNOSTIC_ADDITIONAL_CONTEXT_TOKEN_COUNT,
160 );
161 let syntax_ranges =
162 compute_syntax_ranges(&snapshot, cursor_offset, &(0..snapshot.len()));
163
164 RequestInput::V4(Zeta3PromptInput {
165 cursor_path: excerpt_path,
166 cursor_position: FilePosition {
167 row: cursor_point.row,
168 column: cursor_point.column,
169 },
170 events,
171 editable_context,
172 syntax_ranges,
173 active_buffer_diagnostics,
174 in_open_source_repo: is_open_source,
175 can_collect_data,
176 repo_url,
177 })
178 } else {
179 let (full_context_offset_range, prompt_input) = zeta2_prompt_input(
180 &snapshot,
181 related_files,
182 events,
183 diagnostic_search_range,
184 excerpt_path,
185 cursor_offset,
186 is_open_source,
187 can_collect_data,
188 repo_url,
189 );
190 let formatted_prompt = format_zeta_prompt(&prompt_input, zeta_format);
191
192 RequestInput::V3 {
193 full_context_offset_range,
194 prompt_input,
195 formatted_prompt,
196 }
197 };
198
199 if let Some(debug_tx) = &debug_tx {
200 let prompt = match &request_input {
201 RequestInput::V3 {
202 formatted_prompt, ..
203 } => formatted_prompt.clone(),
204 RequestInput::V4(_) => None,
205 };
206 debug_tx
207 .unbounded_send(DebugEvent::EditPredictionStarted(
208 EditPredictionStartedDebugEvent {
209 buffer: buffer.downgrade(),
210 prompt,
211 position,
212 },
213 ))
214 .ok();
215 }
216
217 log::trace!("Sending edit prediction request");
218
219 let (full_context_offset_range, prompt_input, formatted_prompt) = match request_input {
220 RequestInput::V3 {
221 full_context_offset_range,
222 prompt_input,
223 formatted_prompt,
224 } => (full_context_offset_range, prompt_input, formatted_prompt),
225 RequestInput::V4(v4_prompt_input) => {
226 let inputs = EditPredictionInputs::V3(v4_prompt_input.clone());
227 let (response, usage) = EditPredictionStore::send_v4_request(
228 v4_prompt_input,
229 preferred_experiment.clone(),
230 client,
231 llm_token,
232 organization_id,
233 app_version,
234 trigger,
235 mode,
236 )
237 .await?;
238
239 let request_id = EditPredictionId(response.request_id.into());
240 let model_version = response.model_version;
241 return anyhow::Ok((
242 Some((
243 request_id,
244 Some(PendingPrediction::Patch {
245 inputs,
246 project: project_for_request,
247 buffer,
248 snapshot: snapshot.clone(),
249 patch: response.patch,
250 }),
251 model_version,
252 )),
253 usage,
254 ));
255 }
256 };
257
258 let Some((request_id, output, model_version, usage)) =
259 (if let Some(custom_settings) = &custom_server_settings {
260 let max_tokens = custom_settings.max_output_tokens * 4;
261
262 Some(match local_zeta_version {
263 ZetaVersion::Zeta1 => {
264 let ranges = &prompt_input.excerpt_ranges;
265 let editable_range_in_excerpt = ranges.editable_350.clone();
266 let prompt = zeta1::format_zeta1_from_input(
267 &prompt_input,
268 editable_range_in_excerpt.clone(),
269 ranges.editable_350_context_150.clone(),
270 );
271 let stop_tokens = vec![
272 EDITABLE_REGION_END_MARKER.to_string(),
273 format!("{EDITABLE_REGION_END_MARKER}\n"),
274 format!("{EDITABLE_REGION_END_MARKER}\n\n"),
275 format!("{EDITABLE_REGION_END_MARKER}\n\n\n"),
276 ];
277
278 let (response_text, request_id) = send_custom_server_request(
279 provider,
280 custom_settings,
281 prompt,
282 max_tokens,
283 stop_tokens,
284 open_ai_compatible_api_key.clone(),
285 &http_client,
286 )
287 .await?;
288
289 let request_id = EditPredictionId(request_id.into());
290 let output_text = zeta1::clean_zeta1_model_output(&response_text);
291 let parsed_output = output_text.map(|text| ParsedOutput {
292 new_editable_region: text,
293 range_in_excerpt: editable_range_in_excerpt,
294 cursor_offset_in_new_editable_region: None,
295 });
296
297 (request_id, parsed_output, None, None)
298 }
299 ZetaVersion::Zeta2 | ZetaVersion::Zeta2_1 => {
300 let Some(prompt) = formatted_prompt.clone() else {
301 return Ok((None, None));
302 };
303 let prefill = get_prefill(&prompt_input, zeta_format);
304 let prompt = format!("{prompt}{prefill}");
305
306 let (response_text, request_id) = send_custom_server_request(
307 provider,
308 custom_settings,
309 prompt,
310 max_tokens,
311 stop_tokens_for_format(zeta_format)
312 .iter()
313 .map(|token| token.to_string())
314 .collect(),
315 open_ai_compatible_api_key.clone(),
316 &http_client,
317 )
318 .await?;
319
320 let request_id = EditPredictionId(request_id.into());
321 let output_text = if response_text.is_empty() {
322 None
323 } else {
324 let output = format!("{prefill}{response_text}");
325 Some(parse_zeta2_model_output(
326 &output,
327 zeta_format,
328 &prompt_input,
329 )?)
330 };
331
332 (request_id, output_text, None, None)
333 }
334 })
335 } else if let Some(config) = &raw_config {
336 let Some(prompt) = format_zeta_prompt(&prompt_input, config.format) else {
337 return Ok((None, None));
338 };
339 let prefill = get_prefill(&prompt_input, config.format);
340 let prompt = format!("{prompt}{prefill}");
341 let environment = config
342 .environment
343 .clone()
344 .or_else(|| Some(config.format.to_string().to_lowercase()));
345 let request = RawCompletionRequest {
346 model: config.model_id.clone().unwrap_or_default(),
347 prompt,
348 temperature: None,
349 stop: stop_tokens_for_format(config.format)
350 .iter()
351 .map(|token| std::borrow::Cow::Borrowed(*token))
352 .collect(),
353 max_tokens: Some(2048),
354 environment,
355 };
356
357 let (mut response, usage) = EditPredictionStore::send_raw_llm_request(
358 request,
359 client,
360 None,
361 llm_token,
362 organization_id,
363 app_version,
364 )
365 .await?;
366
367 let request_id = EditPredictionId(response.id.clone().into());
368 let output = if let Some(choice) = response.choices.pop() {
369 let response = &choice.text;
370 let output = format!("{prefill}{response}");
371 Some(parse_zeta2_model_output(
372 &output,
373 config.format,
374 &prompt_input,
375 )?)
376 } else {
377 None
378 };
379
380 Some((request_id, output, None, usage))
381 } else {
382 let (response, usage) = EditPredictionStore::send_v3_request(
383 prompt_input.clone(),
384 preferred_experiment.clone(),
385 client,
386 llm_token,
387 organization_id,
388 app_version,
389 trigger,
390 mode,
391 )
392 .await?;
393
394 let request_id = EditPredictionId(response.request_id.into());
395 let model_version = response.model_version;
396 let parsed_output = ParsedOutput {
397 new_editable_region: response.output,
398 range_in_excerpt: response.editable_range,
399 cursor_offset_in_new_editable_region: response.cursor_offset,
400 };
401
402 Some((request_id, Some(parsed_output), model_version, usage))
403 })
404 else {
405 return Ok((None, None));
406 };
407
408 log::trace!("Got edit prediction response");
409
410 let (new_editable_region, cursor_offset_in_output, editable_range_in_excerpt) =
411 if let Some(output) = output {
412 let ParsedOutput {
413 new_editable_region: output_text,
414 range_in_excerpt: editable_range_in_excerpt,
415 cursor_offset_in_new_editable_region: cursor_offset_in_output,
416 } = output;
417 (
418 Some(output_text),
419 cursor_offset_in_output,
420 editable_range_in_excerpt,
421 )
422 } else {
423 let (editable_range, _) =
424 excerpt_ranges_for_format(zeta_format, &prompt_input.excerpt_ranges);
425 (None, None, editable_range)
426 };
427
428 let editable_range_in_buffer = editable_range_in_excerpt.start
429 + full_context_offset_range.start
430 ..editable_range_in_excerpt.end + full_context_offset_range.start;
431
432 if let Some(debug_tx) = &debug_tx {
433 debug_tx
434 .unbounded_send(DebugEvent::EditPredictionFinished(
435 EditPredictionFinishedDebugEvent {
436 buffer: buffer.downgrade(),
437 position,
438 model_output: new_editable_region.clone(),
439 },
440 ))
441 .ok();
442 }
443 let (edits, cursor_position) = new_editable_region
444 .map(|mut output_text| {
445 let mut old_text = snapshot
446 .text_for_range(editable_range_in_buffer.clone())
447 .collect::<String>();
448
449 if !output_text.is_empty() && !output_text.ends_with('\n') {
450 output_text.push('\n');
451 }
452 if !old_text.is_empty() && !old_text.ends_with('\n') {
453 old_text.push('\n');
454 }
455
456 compute_edits_and_cursor_position(
457 old_text,
458 &output_text,
459 editable_range_in_buffer.start,
460 cursor_offset_in_output,
461 &snapshot,
462 )
463 })
464 .unwrap_or_default();
465
466 let prediction = Some(PendingPrediction::Resolved {
467 inputs: EditPredictionInputs::V2(prompt_input),
468 buffer,
469 snapshot: snapshot.clone(),
470 edits,
471 cursor_position,
472 editable_range_in_buffer: Some(editable_range_in_buffer),
473 });
474
475 anyhow::Ok((Some((request_id, prediction, model_version)), usage))
476 }
477 });
478
479 cx.spawn(async move |this, cx| {
480 let Some((id, prediction, model_version)) =
481 handle_api_response(&this, request_task.await, cx)?
482 else {
483 return Ok(None);
484 };
485 let request_duration = cx.background_executor().now() - request_start;
486
487 let Some(prediction) = prediction else {
488 return Ok(None);
489 };
490
491 let (
492 inputs,
493 edited_buffer,
494 edited_buffer_snapshot,
495 edits,
496 cursor_position,
497 editable_range_in_buffer,
498 ) = match prediction {
499 PendingPrediction::Resolved {
500 inputs,
501 buffer,
502 snapshot,
503 edits,
504 cursor_position,
505 editable_range_in_buffer,
506 } => (
507 inputs,
508 buffer,
509 snapshot,
510 edits,
511 cursor_position,
512 editable_range_in_buffer,
513 ),
514 PendingPrediction::Patch {
515 inputs,
516 project,
517 buffer: fallback_buffer,
518 snapshot: fallback_snapshot,
519 patch,
520 } => {
521 let (buffer, snapshot, edits, cursor_position) =
522 match prediction_edits_for_single_file_diff(&patch, &project, cx).await {
523 Ok(Some(edits)) => edits,
524 Ok(None) => {
525 return Ok(Some(
526 EditPredictionResult::new(
527 id,
528 &fallback_buffer,
529 &fallback_snapshot,
530 Arc::new([]),
531 None,
532 None,
533 inputs,
534 model_version,
535 trigger,
536 request_duration,
537 cx,
538 )
539 .await,
540 ));
541 }
542 Err(error) => {
543 log::error!("failed to apply edit prediction patch: {error:?}");
544 return Ok(Some(EditPredictionResult::new_rejected(
545 id,
546 &fallback_buffer,
547 &fallback_snapshot,
548 inputs,
549 model_version,
550 trigger,
551 request_duration,
552 EditPredictionRejectReason::PatchApplyFailed,
553 )));
554 }
555 };
556 let editable_range_in_buffer =
557 edits
558 .iter()
559 .fold(None, |range: Option<Range<usize>>, (edit_range, _)| {
560 let edit_range = edit_range.start.to_offset(&snapshot)
561 ..edit_range.end.to_offset(&snapshot);
562 Some(match range {
563 Some(range) => {
564 range.start.min(edit_range.start)..range.end.max(edit_range.end)
565 }
566 None => edit_range,
567 })
568 });
569
570 (
571 inputs,
572 buffer,
573 snapshot,
574 edits,
575 cursor_position,
576 editable_range_in_buffer,
577 )
578 }
579 };
580
581 let result = EditPredictionResult::new(
582 id,
583 &edited_buffer,
584 &edited_buffer_snapshot,
585 edits.into(),
586 cursor_position,
587 editable_range_in_buffer
588 .clone()
589 .map(|range| edited_buffer_snapshot.anchor_range_inside(range)),
590 inputs,
591 model_version,
592 trigger,
593 request_duration,
594 cx,
595 )
596 .await;
597
598 if let Some(editable_range_in_buffer) = editable_range_in_buffer.clone() {
599 let prediction = &result.prediction;
600 let weak_this = this.clone();
601 let request_id = prediction.id.clone();
602 let edited_buffer = edited_buffer.clone();
603 let edited_buffer_snapshot = edited_buffer_snapshot.clone();
604 let edit_preview = prediction.edit_preview.clone();
605 let model_version = prediction.model_version.clone();
606 cx.spawn(async move |cx| {
607 weak_this
608 .update(cx, |this, cx| {
609 this.enqueue_settled_prediction(
610 request_id.clone(),
611 &project,
612 &edited_buffer,
613 &edited_buffer_snapshot,
614 editable_range_in_buffer,
615 &edit_preview,
616 context_task,
617 prompt_history_boundary,
618 model_version,
619 request_duration,
620 cx,
621 );
622 })
623 .ok();
624 })
625 .detach();
626 }
627
628 Ok(Some(result))
629 })
630}
631
632fn handle_api_response<T>(
633 this: &WeakEntity<EditPredictionStore>,
634 response: Result<(T, Option<client::EditPredictionUsage>)>,
635 cx: &mut gpui::AsyncApp,
636) -> Result<T> {
637 match response {
638 Ok((data, usage)) => {
639 if let Some(usage) = usage {
640 this.update(cx, |this, cx| {
641 this.user_store.update(cx, |user_store, cx| {
642 user_store.update_edit_prediction_usage(usage, cx);
643 });
644 })
645 .ok();
646 }
647 Ok(data)
648 }
649 Err(err) => {
650 if err.is::<CloudRequestTimeoutError>() {
651 this.update(cx, |this, cx| this.back_off_requests_after_timeout(cx))
652 .ok();
653 }
654
655 if err.is::<ZedUpdateRequiredError>() {
656 cx.update(|cx| {
657 this.update(cx, |this, _cx| {
658 this.update_required = true;
659 })
660 .ok();
661
662 let message: SharedString = err.to_string().into();
663
664 struct UpdateRequiredError {
665 message: SharedString,
666 }
667 impl WorkspaceError for UpdateRequiredError {
668 fn primary_message(&self) -> SharedString {
669 self.message.clone()
670 }
671 fn severity(&self) -> ErrorSeverity {
672 ErrorSeverity::Critical
673 }
674 fn primary_action(&self) -> ErrorAction {
675 ErrorAction::link(
676 "Update Omega",
677 "https://github.com/OpenAgentsInc/omega/releases",
678 )
679 }
680 }
681
682 show_app_notification(
683 NotificationId::unique::<ZedUpdateRequiredError>(),
684 cx,
685 move |cx| {
686 cx.new({
687 let message = message.clone();
688 move |cx| {
689 let error = UpdateRequiredError { message };
690 MessageNotification::from_workspace_error(error, cx)
691 }
692 })
693 },
694 );
695 });
696 }
697 Err(err)
698 }
699 }
700}
701
702const ACTIVE_BUFFER_DIAGNOSTIC_ADDITIONAL_CONTEXT_TOKEN_COUNT: usize = 100;
703const MAX_ACTIVE_BUFFER_DIAGNOSTICS_TO_COLLECT: usize = 20;
704pub(crate) const MAX_ACTIVE_BUFFER_DIAGNOSTIC_MESSAGE_TOKENS_TO_COLLECT: usize = 512;
705pub(crate) const MAX_ACTIVE_BUFFER_DIAGNOSTIC_SNIPPET_TOKENS_TO_COLLECT: usize = 512;
706
707pub(crate) fn active_buffer_diagnostics(
708 snapshot: &language::BufferSnapshot,
709 diagnostic_search_range: Range<Point>,
710 cursor_row: u32,
711 additional_context_token_count: usize,
712) -> Vec<zeta_prompt::ActiveBufferDiagnostic> {
713 let mut diagnostics = snapshot
714 .diagnostics_in_range::<Point, Point>(diagnostic_search_range, false)
715 .collect::<Vec<_>>();
716 diagnostics.sort_by_key(|entry| {
717 cursor_row.abs_diff(entry.range.start.row) + cursor_row.abs_diff(entry.range.end.row)
718 });
719
720 diagnostics
721 .into_iter()
722 .map(|entry| {
723 let diagnostic_point_range = entry.range.clone();
724 let snippet_point_range = cursor_excerpt::expand_context_syntactically_then_linewise(
725 snapshot,
726 diagnostic_point_range.clone(),
727 additional_context_token_count,
728 );
729
730 let severity = match entry.diagnostic.severity {
731 DiagnosticSeverity::ERROR => Some(1),
732 DiagnosticSeverity::WARNING => Some(2),
733 DiagnosticSeverity::INFORMATION => Some(3),
734 DiagnosticSeverity::HINT => Some(4),
735 _ => None,
736 };
737 (
738 severity,
739 zeta_prompt::clamp_text_to_token_count(
740 &entry.diagnostic.message,
741 MAX_ACTIVE_BUFFER_DIAGNOSTIC_MESSAGE_TOKENS_TO_COLLECT,
742 )
743 .to_string(),
744 diagnostic_point_range,
745 snippet_point_range,
746 )
747 })
748 .take(MAX_ACTIVE_BUFFER_DIAGNOSTICS_TO_COLLECT)
749 .map(
750 |(severity, message, diagnostic_point_range, snippet_point_range)| {
751 let (snippet, diagnostic_range_in_snippet) = if snippet_point_range.start
752 == Point::new(0, 0)
753 && snippet_point_range.end == snapshot.max_point()
754 {
755 (String::new(), 0..0)
756 } else {
757 let snippet = snapshot
758 .text_for_range(snippet_point_range.clone())
759 .collect::<String>();
760 let snippet = zeta_prompt::clamp_text_to_token_count(
761 &snippet,
762 MAX_ACTIVE_BUFFER_DIAGNOSTIC_SNIPPET_TOKENS_TO_COLLECT,
763 )
764 .to_string();
765 let snippet_start_offset = snippet_point_range.start.to_offset(snapshot);
766 let diagnostic_offset_range = diagnostic_point_range.to_offset(snapshot);
767 let diagnostic_range_start = diagnostic_offset_range
768 .start
769 .saturating_sub(snippet_start_offset)
770 .min(snippet.len());
771 let diagnostic_range_end = diagnostic_offset_range
772 .end
773 .saturating_sub(snippet_start_offset)
774 .min(snippet.len());
775 (snippet, diagnostic_range_start..diagnostic_range_end)
776 };
777 zeta_prompt::ActiveBufferDiagnostic {
778 severity,
779 message,
780 snippet,
781 snippet_buffer_row_range: diagnostic_point_range.start.row
782 ..diagnostic_point_range.end.row,
783 diagnostic_range_in_snippet,
784 }
785 },
786 )
787 .collect()
788}
789
790pub fn zeta2_prompt_input(
791 snapshot: &language::BufferSnapshot,
792 mut related_files: Vec<zeta_prompt::RelatedFile>,
793 events: Vec<Arc<zeta_prompt::Event>>,
794 diagnostic_search_range: Range<Point>,
795 excerpt_path: Arc<Path>,
796 cursor_offset: usize,
797 is_open_source: bool,
798 can_collect_data: bool,
799 repo_url: Option<String>,
800) -> (Range<usize>, zeta_prompt::Zeta2PromptInput) {
801 let (excerpt_point_range, excerpt_offset_range, cursor_offset_in_excerpt) =
802 compute_cursor_excerpt(snapshot, cursor_offset);
803
804 let cursor_excerpt: Arc<str> = snapshot
805 .text_for_range(excerpt_point_range.clone())
806 .collect::<String>()
807 .into();
808 let syntax_ranges = compute_syntax_ranges(snapshot, cursor_offset, &excerpt_offset_range);
809 let excerpt_ranges = zeta_prompt::compute_legacy_excerpt_ranges(
810 &cursor_excerpt,
811 cursor_offset_in_excerpt,
812 &syntax_ranges,
813 );
814
815 let active_buffer_diagnostics = active_buffer_diagnostics(
816 snapshot,
817 diagnostic_search_range,
818 snapshot.offset_to_point(cursor_offset).row,
819 ACTIVE_BUFFER_DIAGNOSTIC_ADDITIONAL_CONTEXT_TOKEN_COUNT,
820 );
821 for file in &mut related_files {
822 for excerpt in &mut file.excerpts {
823 excerpt.context_source = zeta_prompt::ContextSource::Lsp;
824 }
825 }
826
827 let prompt_input = zeta_prompt::Zeta2PromptInput {
828 cursor_path: excerpt_path,
829 cursor_excerpt,
830 cursor_offset_in_excerpt,
831 excerpt_start_row: Some(excerpt_point_range.start.row),
832 events,
833 related_files: Some(related_files),
834 active_buffer_diagnostics,
835 excerpt_ranges,
836 syntax_ranges: Some(syntax_ranges),
837 in_open_source_repo: is_open_source,
838 can_collect_data,
839 repo_url,
840 };
841 (excerpt_offset_range, prompt_input)
842}
843
844pub(crate) fn edit_prediction_accepted(
845 store: &EditPredictionStore,
846 current_prediction: CurrentEditPrediction,
847 cx: &App,
848) {
849 if store.zeta2_raw_config().is_some() {
850 return;
851 }
852
853 let request_id = current_prediction.prediction.id.to_string();
854 let model_version = current_prediction.prediction.model_version;
855 let e2e_latency = current_prediction.e2e_latency;
856 let client = store.client.clone();
857 let llm_token = store.llm_token.clone();
858 let organization_id = store
859 .user_store
860 .read(cx)
861 .current_organization()
862 .map(|organization| organization.id.clone());
863 let app_version = AppVersion::global(cx);
864
865 cx.background_spawn(async move {
866 let body = serde_json::to_string(&AcceptEditPredictionBody {
867 request_id,
868 model_version,
869 e2e_latency_ms: Some(e2e_latency.as_millis()),
870 })?;
871
872 let url = client
873 .http_client()
874 .build_zed_llm_url("/predict_edits/accept", &[])?;
875 EditPredictionStore::send_api_request::<()>(
876 move |builder| Ok(builder.uri(url.as_ref()).body(body.clone().into())?),
877 client,
878 llm_token,
879 organization_id,
880 app_version,
881 )
882 .await?;
883 anyhow::Ok(())
884 })
885 .detach_and_log_err(cx);
886}
887
888pub fn compute_edits(
889 old_text: String,
890 new_text: &str,
891 offset: usize,
892 snapshot: &BufferSnapshot,
893) -> Vec<(Range<Anchor>, Arc<str>)> {
894 compute_edits_and_cursor_position(old_text, new_text, offset, None, snapshot).0
895}
896
897pub fn compute_edits_and_cursor_position(
898 old_text: String,
899 new_text: &str,
900 offset: usize,
901 cursor_offset_in_new_text: Option<usize>,
902 snapshot: &BufferSnapshot,
903) -> (
904 Vec<(Range<Anchor>, Arc<str>)>,
905 Option<PredictedCursorPosition>,
906) {
907 let diffs = text_diff(&old_text, new_text);
908
909 // Delta represents the cumulative change in byte count from all preceding edits.
910 // new_offset = old_offset + delta, so old_offset = new_offset - delta
911 let mut delta: isize = 0;
912 let mut cursor_position: Option<PredictedCursorPosition> = None;
913 let buffer_len = snapshot.len();
914
915 let edits = diffs
916 .iter()
917 .map(|(raw_old_range, new_text)| {
918 // Compute cursor position if it falls within or before this edit.
919 if let (Some(cursor_offset), None) = (cursor_offset_in_new_text, cursor_position) {
920 let edit_start_in_new = (raw_old_range.start as isize + delta) as usize;
921 let edit_end_in_new = edit_start_in_new + new_text.len();
922
923 if cursor_offset < edit_start_in_new {
924 let cursor_in_old = (cursor_offset as isize - delta) as usize;
925 let buffer_offset = (offset + cursor_in_old).min(buffer_len);
926 cursor_position = Some(PredictedCursorPosition::at_anchor(
927 snapshot.anchor_after(buffer_offset),
928 ));
929 } else if cursor_offset < edit_end_in_new {
930 let buffer_offset = (offset + raw_old_range.start).min(buffer_len);
931 let offset_within_insertion = cursor_offset - edit_start_in_new;
932 cursor_position = Some(PredictedCursorPosition::new(
933 snapshot.anchor_before(buffer_offset),
934 offset_within_insertion,
935 ));
936 }
937
938 delta += new_text.len() as isize - raw_old_range.len() as isize;
939 }
940
941 // Compute the edit with prefix/suffix trimming.
942 let mut old_range = raw_old_range.clone();
943 let old_slice = &old_text[old_range.clone()];
944
945 let prefix_len = common_prefix(old_slice.chars(), new_text.chars());
946 let suffix_len = common_prefix(
947 old_slice[prefix_len..].chars().rev(),
948 new_text[prefix_len..].chars().rev(),
949 );
950
951 old_range.start += offset;
952 old_range.end += offset;
953 old_range.start += prefix_len;
954 old_range.end -= suffix_len;
955
956 old_range.start = old_range.start.min(buffer_len);
957 old_range.end = old_range.end.min(buffer_len);
958
959 let new_text = new_text[prefix_len..new_text.len() - suffix_len].into();
960 let range = if old_range.is_empty() {
961 let anchor = snapshot.anchor_after(old_range.start);
962 anchor..anchor
963 } else {
964 snapshot.anchor_after(old_range.start)..snapshot.anchor_before(old_range.end)
965 };
966 (range, new_text)
967 })
968 .collect();
969
970 if let (Some(cursor_offset), None) = (cursor_offset_in_new_text, cursor_position) {
971 let cursor_in_old = (cursor_offset as isize - delta) as usize;
972 let buffer_offset = snapshot.clip_offset(offset + cursor_in_old, Bias::Right);
973 cursor_position = Some(PredictedCursorPosition::at_anchor(
974 snapshot.anchor_after(buffer_offset),
975 ));
976 }
977
978 (edits, cursor_position)
979}
980
981fn common_prefix<T1: Iterator<Item = char>, T2: Iterator<Item = char>>(a: T1, b: T2) -> usize {
982 a.zip(b)
983 .take_while(|(a, b)| a == b)
984 .map(|(a, _)| a.len_utf8())
985 .sum()
986}
987