1use crate::{
2 EditPredictionId, EditPredictionModelInput, cursor_excerpt,
3 open_ai_compatible::{self, load_open_ai_compatible_api_key_if_needed},
4 prediction::EditPredictionResult,
5};
6use anyhow::{Context as _, Result, anyhow};
7use gpui::{App, AppContext as _, Entity, Task};
8use language::{
9 Anchor, Buffer, BufferSnapshot, OffsetRangeExt as _, ToOffset, ToPoint as _,
10 language_settings::all_language_settings,
11};
12use settings::EditPredictionPromptFormat;
13use std::{path::Path, sync::Arc, time::Instant};
14use zeta_prompt::ZetaPromptInput;
15
16const FIM_CONTEXT_TOKENS: usize = 512;
17
18struct FimRequestOutput {
19 request_id: String,
20 edits: Vec<(std::ops::Range<Anchor>, Arc<str>)>,
21 snapshot: BufferSnapshot,
22 response_received_at: Instant,
23 inputs: ZetaPromptInput,
24 buffer: Entity<Buffer>,
25 buffer_snapshotted_at: Instant,
26}
27
28pub fn request_prediction(
29 EditPredictionModelInput {
30 buffer,
31 snapshot,
32 position,
33 events,
34 ..
35 }: EditPredictionModelInput,
36 prompt_format: EditPredictionPromptFormat,
37 cx: &mut App,
38) -> Task<Result<Option<EditPredictionResult>>> {
39 let settings = &all_language_settings(None, cx).edit_predictions;
40 let provider = settings.provider;
41
42 let full_path: Arc<Path> = snapshot
43 .file()
44 .map(|file| file.full_path(cx))
45 .unwrap_or_else(|| "untitled".into())
46 .into();
47
48 let http_client = cx.http_client();
49 let cursor_point = position.to_point(&snapshot);
50 let buffer_snapshotted_at = Instant::now();
51
52 let Some(settings) = (match provider {
53 settings::EditPredictionProvider::Ollama => settings.ollama.clone(),
54 settings::EditPredictionProvider::OpenAiCompatibleApi => {
55 settings.open_ai_compatible_api.clone()
56 }
57 _ => None,
58 }) else {
59 return Task::ready(Err(anyhow!("Unsupported edit prediction provider for FIM")));
60 };
61
62 let api_key = load_open_ai_compatible_api_key_if_needed(provider, cx);
63
64 let result = cx.background_spawn(async move {
65 let (excerpt_range, _) = cursor_excerpt::editable_and_context_ranges_for_cursor_position(
66 cursor_point,
67 &snapshot,
68 FIM_CONTEXT_TOKENS,
69 0,
70 );
71 let excerpt_offset_range = excerpt_range.to_offset(&snapshot);
72 let cursor_offset = cursor_point.to_offset(&snapshot);
73
74 let inputs = ZetaPromptInput {
75 events,
76 related_files: Vec::new(),
77 cursor_offset_in_excerpt: cursor_offset - excerpt_offset_range.start,
78 cursor_path: full_path.clone(),
79 excerpt_start_row: Some(excerpt_range.start.row),
80 cursor_excerpt: snapshot
81 .text_for_range(excerpt_range)
82 .collect::<String>()
83 .into(),
84 excerpt_ranges: Default::default(),
85 experiment: None,
86 in_open_source_repo: false,
87 can_collect_data: false,
88 };
89
90 let prefix = inputs.cursor_excerpt[..inputs.cursor_offset_in_excerpt].to_string();
91 let suffix = inputs.cursor_excerpt[inputs.cursor_offset_in_excerpt..].to_string();
92 let prompt = format_fim_prompt(prompt_format, &prefix, &suffix);
93 let stop_tokens = get_fim_stop_tokens();
94
95 let max_tokens = settings.max_output_tokens;
96
97 let (response_text, request_id) = open_ai_compatible::send_custom_server_request(
98 provider,
99 &settings,
100 prompt,
101 max_tokens,
102 stop_tokens,
103 api_key,
104 &http_client,
105 )
106 .await?;
107
108 let response_received_at = Instant::now();
109
110 log::debug!(
111 "fim: completion received ({:.2}s)",
112 (response_received_at - buffer_snapshotted_at).as_secs_f64()
113 );
114
115 let completion: Arc<str> = clean_fim_completion(&response_text).into();
116 let edits = if completion.is_empty() {
117 vec![]
118 } else {
119 let cursor_offset = cursor_point.to_offset(&snapshot);
120 let anchor = snapshot.anchor_after(cursor_offset);
121 vec![(anchor..anchor, completion)]
122 };
123
124 anyhow::Ok(FimRequestOutput {
125 request_id,
126 edits,
127 snapshot,
128 response_received_at,
129 inputs,
130 buffer,
131 buffer_snapshotted_at,
132 })
133 });
134
135 cx.spawn(async move |cx: &mut gpui::AsyncApp| {
136 let output = result.await.context("fim edit prediction failed")?;
137 anyhow::Ok(Some(
138 EditPredictionResult::new(
139 EditPredictionId(output.request_id.into()),
140 &output.buffer,
141 &output.snapshot,
142 output.edits.into(),
143 None,
144 output.buffer_snapshotted_at,
145 output.response_received_at,
146 output.inputs,
147 None,
148 cx,
149 )
150 .await,
151 ))
152 })
153}
154
155fn format_fim_prompt(
156 prompt_format: EditPredictionPromptFormat,
157 prefix: &str,
158 suffix: &str,
159) -> String {
160 match prompt_format {
161 EditPredictionPromptFormat::CodeLlama => {
162 format!("<PRE> {prefix} <SUF>{suffix} <MID>")
163 }
164 EditPredictionPromptFormat::StarCoder => {
165 format!("<fim_prefix>{prefix}<fim_suffix>{suffix}<fim_middle>")
166 }
167 EditPredictionPromptFormat::DeepseekCoder => {
168 format!("<｜fim▁begin｜>{prefix}<｜fim▁hole｜>{suffix}<｜fim▁end｜>")
169 }
170 EditPredictionPromptFormat::Qwen | EditPredictionPromptFormat::CodeGemma => {
171 format!("<|fim_prefix|>{prefix}<|fim_suffix|>{suffix}<|fim_middle|>")
172 }
173 EditPredictionPromptFormat::Codestral => {
174 format!("[SUFFIX]{suffix}[PREFIX]{prefix}")
175 }
176 EditPredictionPromptFormat::Glm => {
177 format!("<|code_prefix|>{prefix}<|code_suffix|>{suffix}<|code_middle|>")
178 }
179 _ => {
180 format!("<fim_prefix>{prefix}<fim_suffix>{suffix}<fim_middle>")
181 }
182 }
183}
184
185fn get_fim_stop_tokens() -> Vec<String> {
186 vec![
187 "<|endoftext|>".to_string(),
188 "<|file_separator|>".to_string(),
189 "<|fim_pad|>".to_string(),
190 "<|fim_prefix|>".to_string(),
191 "<|fim_middle|>".to_string(),
192 "<|fim_suffix|>".to_string(),
193 "<fim_prefix>".to_string(),
194 "<fim_middle>".to_string(),
195 "<fim_suffix>".to_string(),
196 "<PRE>".to_string(),
197 "<SUF>".to_string(),
198 "<MID>".to_string(),
199 "[PREFIX]".to_string(),
200 "[SUFFIX]".to_string(),
201 ]
202}
203
204fn clean_fim_completion(response: &str) -> String {
205 let mut result = response.to_string();
206
207 let end_tokens = [
208 "<|endoftext|>",
209 "<|file_separator|>",
210 "<|fim_pad|>",
211 "<|fim_prefix|>",
212 "<|fim_middle|>",
213 "<|fim_suffix|>",
214 "<fim_prefix>",
215 "<fim_middle>",
216 "<fim_suffix>",
217 "<PRE>",
218 "<SUF>",
219 "<MID>",
220 "[PREFIX]",
221 "[SUFFIX]",
222 ];
223
224 for token in &end_tokens {
225 if let Some(pos) = result.find(token) {
226 result.truncate(pos);
227 }
228 }
229
230 result
231}