1use crate::{
2 EditPredictionId, EditPredictionModelInput, cursor_excerpt,
3 prediction::EditPredictionResult,
4 zeta1::{
5 self, MAX_CONTEXT_TOKENS as ZETA_MAX_CONTEXT_TOKENS,
6 MAX_EVENT_TOKENS as ZETA_MAX_EVENT_TOKENS,
7 },
8};
9use anyhow::{Context as _, Result};
10use futures::AsyncReadExt as _;
11use gpui::{App, AppContext as _, Entity, SharedString, Task, http_client};
12use language::{
13 Anchor, Buffer, BufferSnapshot, OffsetRangeExt as _, ToOffset, ToPoint as _,
14 language_settings::all_language_settings,
15};
16use language_model::{LanguageModelProviderId, LanguageModelRegistry};
17use serde::{Deserialize, Serialize};
18use std::{path::Path, sync::Arc, time::Instant};
19use zeta_prompt::{
20 ZetaPromptInput,
21 zeta1::{EDITABLE_REGION_END_MARKER, format_zeta1_prompt},
22};
23
24const FIM_CONTEXT_TOKENS: usize = 512;
25
26pub struct Ollama;
27
28#[derive(Debug, Serialize)]
29struct OllamaGenerateRequest {
30 model: String,
31 prompt: String,
32 raw: bool,
33 stream: bool,
34 #[serde(skip_serializing_if = "Option::is_none")]
35 options: Option<OllamaGenerateOptions>,
36}
37
38#[derive(Debug, Serialize)]
39struct OllamaGenerateOptions {
40 #[serde(skip_serializing_if = "Option::is_none")]
41 num_predict: Option<u32>,
42 #[serde(skip_serializing_if = "Option::is_none")]
43 temperature: Option<f32>,
44 #[serde(skip_serializing_if = "Option::is_none")]
45 stop: Option<Vec<String>>,
46}
47
48#[derive(Debug, Deserialize)]
49struct OllamaGenerateResponse {
50 created_at: String,
51 response: String,
52}
53
54const PROVIDER_ID: LanguageModelProviderId = LanguageModelProviderId::new("ollama");
55
56pub fn is_available(cx: &App) -> bool {
57 LanguageModelRegistry::read_global(cx)
58 .provider(&PROVIDER_ID)
59 .is_some_and(|provider| provider.is_authenticated(cx))
60}
61
62pub fn ensure_authenticated(cx: &mut App) {
63 if let Some(provider) = LanguageModelRegistry::read_global(cx).provider(&PROVIDER_ID) {
64 provider.authenticate(cx).detach_and_log_err(cx);
65 }
66}
67
68pub fn fetch_models(cx: &mut App) -> Vec<SharedString> {
69 let Some(provider) = LanguageModelRegistry::read_global(cx).provider(&PROVIDER_ID) else {
70 return Vec::new();
71 };
72 provider.authenticate(cx).detach_and_log_err(cx);
73 let mut models: Vec<SharedString> = provider
74 .provided_models(cx)
75 .into_iter()
76 .map(|model| SharedString::from(model.id().0.to_string()))
77 .collect();
78 models.sort();
79 models
80}
81
82/// Output from the Ollama HTTP request, containing all data needed to create the prediction result.
83struct OllamaRequestOutput {
84 created_at: String,
85 edits: Vec<(std::ops::Range<Anchor>, Arc<str>)>,
86 snapshot: BufferSnapshot,
87 response_received_at: Instant,
88 inputs: ZetaPromptInput,
89 buffer: Entity<Buffer>,
90 buffer_snapshotted_at: Instant,
91}
92
93impl Ollama {
94 pub fn new() -> Self {
95 Self
96 }
97
98 pub fn request_prediction(
99 &self,
100 EditPredictionModelInput {
101 buffer,
102 snapshot,
103 position,
104 events,
105 ..
106 }: EditPredictionModelInput,
107 cx: &mut App,
108 ) -> Task<Result<Option<EditPredictionResult>>> {
109 let settings = &all_language_settings(None, cx).edit_predictions.ollama;
110 let Some(model) = settings.model.clone() else {
111 return Task::ready(Ok(None));
112 };
113 let api_url = settings.api_url.clone();
114
115 log::debug!("Ollama: Requesting completion (model: {})", model);
116
117 let full_path: Arc<Path> = snapshot
118 .file()
119 .map(|file| file.full_path(cx))
120 .unwrap_or_else(|| "untitled".into())
121 .into();
122
123 let http_client = cx.http_client();
124 let cursor_point = position.to_point(&snapshot);
125 let buffer_snapshotted_at = Instant::now();
126
127 let is_zeta = is_zeta_model(&model);
128
129 // Zeta generates more tokens than FIM models. Ideally, we'd use MAX_REWRITE_TOKENS,
130 // but this might be too slow for local deployments. So we make it configurable,
131 // but we also have this hardcoded multiplier for now.
132 let max_output_tokens = if is_zeta {
133 settings.max_output_tokens * 4
134 } else {
135 settings.max_output_tokens
136 };
137
138 let result = cx.background_spawn(async move {
139 let zeta_editable_region_tokens = max_output_tokens as usize;
140
141 // For zeta models, use the dedicated zeta1 functions which handle their own
142 // range computation with the correct token limits.
143 let (prompt, stop_tokens, editable_range_override, inputs) = if is_zeta {
144 let path_str = full_path.to_string_lossy();
145 let input_excerpt = zeta1::excerpt_for_cursor_position(
146 cursor_point,
147 &path_str,
148 &snapshot,
149 zeta_editable_region_tokens,
150 ZETA_MAX_CONTEXT_TOKENS,
151 );
152 let input_events = zeta1::prompt_for_events(&events, ZETA_MAX_EVENT_TOKENS);
153 let prompt = format_zeta1_prompt(&input_events, &input_excerpt.prompt);
154 let editable_offset_range = input_excerpt.editable_range.to_offset(&snapshot);
155 let context_offset_range = input_excerpt.context_range.to_offset(&snapshot);
156 let stop_tokens = get_zeta_stop_tokens();
157
158 let inputs = ZetaPromptInput {
159 events,
160 related_files: Vec::new(),
161 cursor_offset_in_excerpt: cursor_point.to_offset(&snapshot)
162 - context_offset_range.start,
163 cursor_path: full_path.clone(),
164 cursor_excerpt: snapshot
165 .text_for_range(input_excerpt.context_range)
166 .collect::<String>()
167 .into(),
168 editable_range_in_excerpt: (editable_offset_range.start
169 - context_offset_range.start)
170 ..(editable_offset_range.end - context_offset_range.start),
171 };
172
173 (prompt, stop_tokens, Some(editable_offset_range), inputs)
174 } else {
175 let (excerpt_range, _) =
176 cursor_excerpt::editable_and_context_ranges_for_cursor_position(
177 cursor_point,
178 &snapshot,
179 FIM_CONTEXT_TOKENS,
180 0,
181 );
182 let excerpt_offset_range = excerpt_range.to_offset(&snapshot);
183 let cursor_offset = cursor_point.to_offset(&snapshot);
184
185 let inputs = ZetaPromptInput {
186 events,
187 related_files: Vec::new(),
188 cursor_offset_in_excerpt: cursor_offset - excerpt_offset_range.start,
189 editable_range_in_excerpt: cursor_offset - excerpt_offset_range.start
190 ..cursor_offset - excerpt_offset_range.start,
191 cursor_path: full_path.clone(),
192 cursor_excerpt: snapshot
193 .text_for_range(excerpt_range)
194 .collect::<String>()
195 .into(),
196 };
197
198 let prefix = inputs.cursor_excerpt[..inputs.cursor_offset_in_excerpt].to_string();
199 let suffix = inputs.cursor_excerpt[inputs.cursor_offset_in_excerpt..].to_string();
200 let prompt = format_fim_prompt(&model, &prefix, &suffix);
201 let stop_tokens = get_fim_stop_tokens();
202
203 (prompt, stop_tokens, None, inputs)
204 };
205
206 let request = OllamaGenerateRequest {
207 model: model.clone(),
208 prompt,
209 raw: true,
210 stream: false,
211 options: Some(OllamaGenerateOptions {
212 num_predict: Some(max_output_tokens),
213 temperature: Some(0.2),
214 stop: Some(stop_tokens),
215 }),
216 };
217
218 let request_body = serde_json::to_string(&request)?;
219 let http_request = http_client::Request::builder()
220 .method(http_client::Method::POST)
221 .uri(format!("{}/api/generate", api_url))
222 .header("Content-Type", "application/json")
223 .body(http_client::AsyncBody::from(request_body))?;
224
225 let mut response = http_client.send(http_request).await?;
226 let status = response.status();
227
228 log::debug!("Ollama: Response status: {}", status);
229
230 if !status.is_success() {
231 let mut body = String::new();
232 response.body_mut().read_to_string(&mut body).await?;
233 return Err(anyhow::anyhow!("Ollama API error: {} - {}", status, body));
234 }
235
236 let mut body = String::new();
237 response.body_mut().read_to_string(&mut body).await?;
238
239 let ollama_response: OllamaGenerateResponse =
240 serde_json::from_str(&body).context("Failed to parse Ollama response")?;
241
242 let response_received_at = Instant::now();
243
244 log::debug!(
245 "Ollama: Completion received ({:.2}s)",
246 (response_received_at - buffer_snapshotted_at).as_secs_f64()
247 );
248
249 let edits = if is_zeta {
250 let editable_range =
251 editable_range_override.expect("zeta model should have editable range");
252
253 log::trace!("ollama response: {}", ollama_response.response);
254
255 let response = clean_zeta_completion(&ollama_response.response);
256 match zeta1::parse_edits(&response, editable_range, &snapshot) {
257 Ok(edits) => edits,
258 Err(err) => {
259 log::warn!("Ollama zeta: Failed to parse response: {}", err);
260 vec![]
261 }
262 }
263 } else {
264 let completion: Arc<str> = clean_fim_completion(&ollama_response.response).into();
265 if completion.is_empty() {
266 vec![]
267 } else {
268 let cursor_offset = cursor_point.to_offset(&snapshot);
269 let anchor = snapshot.anchor_after(cursor_offset);
270 vec![(anchor..anchor, completion)]
271 }
272 };
273
274 anyhow::Ok(OllamaRequestOutput {
275 created_at: ollama_response.created_at,
276 edits,
277 snapshot,
278 response_received_at,
279 inputs,
280 buffer,
281 buffer_snapshotted_at,
282 })
283 });
284
285 cx.spawn(async move |cx: &mut gpui::AsyncApp| {
286 let output = result.await.context("Ollama edit prediction failed")?;
287 anyhow::Ok(Some(
288 EditPredictionResult::new(
289 EditPredictionId(output.created_at.into()),
290 &output.buffer,
291 &output.snapshot,
292 output.edits.into(),
293 output.buffer_snapshotted_at,
294 output.response_received_at,
295 output.inputs,
296 cx,
297 )
298 .await,
299 ))
300 })
301 }
302}
303
304fn is_zeta_model(model: &str) -> bool {
305 model.to_lowercase().contains("zeta")
306}
307
308fn get_zeta_stop_tokens() -> Vec<String> {
309 vec![EDITABLE_REGION_END_MARKER.to_string(), "```".to_string()]
310}
311
312fn format_fim_prompt(model: &str, prefix: &str, suffix: &str) -> String {
313 let model_base = model.split(':').next().unwrap_or(model);
314
315 match model_base {
316 "codellama" | "code-llama" => {
317 format!("<PRE> {prefix} <SUF>{suffix} <MID>")
318 }
319 "starcoder" | "starcoder2" | "starcoderbase" => {
320 format!("<fim_prefix>{prefix}<fim_suffix>{suffix}<fim_middle>")
321 }
322 "deepseek-coder" | "deepseek-coder-v2" => {
323 format!("<｜fim▁begin｜>{prefix}<｜fim▁hole｜>{suffix}<｜fim▁end｜>")
324 }
325 "qwen2.5-coder" | "qwen-coder" | "qwen" => {
326 format!("<|fim_prefix|>{prefix}<|fim_suffix|>{suffix}<|fim_middle|>")
327 }
328 "codegemma" => {
329 format!("<|fim_prefix|>{prefix}<|fim_suffix|>{suffix}<|fim_middle|>")
330 }
331 "codestral" | "mistral" => {
332 format!("[SUFFIX]{suffix}[PREFIX]{prefix}")
333 }
334 "glm" | "glm-4" | "glm-4.5" => {
335 format!("<|code_prefix|>{prefix}<|code_suffix|>{suffix}<|code_middle|>")
336 }
337 _ => {
338 format!("<fim_prefix>{prefix}<fim_suffix>{suffix}<fim_middle>")
339 }
340 }
341}
342
343fn get_fim_stop_tokens() -> Vec<String> {
344 vec![
345 "<|endoftext|>".to_string(),
346 "<|file_separator|>".to_string(),
347 "<|fim_pad|>".to_string(),
348 "<|fim_prefix|>".to_string(),
349 "<|fim_middle|>".to_string(),
350 "<|fim_suffix|>".to_string(),
351 "<fim_prefix>".to_string(),
352 "<fim_middle>".to_string(),
353 "<fim_suffix>".to_string(),
354 "<PRE>".to_string(),
355 "<SUF>".to_string(),
356 "<MID>".to_string(),
357 "[PREFIX]".to_string(),
358 "[SUFFIX]".to_string(),
359 ]
360}
361
362fn clean_zeta_completion(mut response: &str) -> &str {
363 if let Some(last_newline_ix) = response.rfind('\n') {
364 let last_line = &response[last_newline_ix + 1..];
365 if EDITABLE_REGION_END_MARKER.starts_with(&last_line) {
366 response = &response[..last_newline_ix]
367 }
368 }
369 response
370}
371
372fn clean_fim_completion(response: &str) -> String {
373 let mut result = response.to_string();
374
375 let end_tokens = [
376 "<|endoftext|>",
377 "<|file_separator|>",
378 "<|fim_pad|>",
379 "<|fim_prefix|>",
380 "<|fim_middle|>",
381 "<|fim_suffix|>",
382 "<fim_prefix>",
383 "<fim_middle>",
384 "<fim_suffix>",
385 "<PRE>",
386 "<SUF>",
387 "<MID>",
388 "[PREFIX]",
389 "[SUFFIX]",
390 ];
391
392 for token in &end_tokens {
393 if let Some(pos) = result.find(token) {
394 result.truncate(pos);
395 }
396 }
397
398 result
399}