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.clone())
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 excerpt_start_row: Some(input_excerpt.context_range.start.row),
172 };
173
174 (prompt, stop_tokens, Some(editable_offset_range), inputs)
175 } else {
176 let (excerpt_range, _) =
177 cursor_excerpt::editable_and_context_ranges_for_cursor_position(
178 cursor_point,
179 &snapshot,
180 FIM_CONTEXT_TOKENS,
181 0,
182 );
183 let excerpt_offset_range = excerpt_range.to_offset(&snapshot);
184 let cursor_offset = cursor_point.to_offset(&snapshot);
185
186 let inputs = ZetaPromptInput {
187 events,
188 related_files: Vec::new(),
189 cursor_offset_in_excerpt: cursor_offset - excerpt_offset_range.start,
190 editable_range_in_excerpt: cursor_offset - excerpt_offset_range.start
191 ..cursor_offset - excerpt_offset_range.start,
192 cursor_path: full_path.clone(),
193 excerpt_start_row: Some(excerpt_range.start.row),
194 cursor_excerpt: snapshot
195 .text_for_range(excerpt_range)
196 .collect::<String>()
197 .into(),
198 };
199
200 let prefix = inputs.cursor_excerpt[..inputs.cursor_offset_in_excerpt].to_string();
201 let suffix = inputs.cursor_excerpt[inputs.cursor_offset_in_excerpt..].to_string();
202 let prompt = format_fim_prompt(&model, &prefix, &suffix);
203 let stop_tokens = get_fim_stop_tokens();
204
205 (prompt, stop_tokens, None, inputs)
206 };
207
208 let request = OllamaGenerateRequest {
209 model: model.clone(),
210 prompt,
211 raw: true,
212 stream: false,
213 options: Some(OllamaGenerateOptions {
214 num_predict: Some(max_output_tokens),
215 temperature: Some(0.2),
216 stop: Some(stop_tokens),
217 }),
218 };
219
220 let request_body = serde_json::to_string(&request)?;
221 let http_request = http_client::Request::builder()
222 .method(http_client::Method::POST)
223 .uri(format!("{}/api/generate", api_url))
224 .header("Content-Type", "application/json")
225 .body(http_client::AsyncBody::from(request_body))?;
226
227 let mut response = http_client.send(http_request).await?;
228 let status = response.status();
229
230 log::debug!("Ollama: Response status: {}", status);
231
232 if !status.is_success() {
233 let mut body = String::new();
234 response.body_mut().read_to_string(&mut body).await?;
235 return Err(anyhow::anyhow!("Ollama API error: {} - {}", status, body));
236 }
237
238 let mut body = String::new();
239 response.body_mut().read_to_string(&mut body).await?;
240
241 let ollama_response: OllamaGenerateResponse =
242 serde_json::from_str(&body).context("Failed to parse Ollama response")?;
243
244 let response_received_at = Instant::now();
245
246 log::debug!(
247 "Ollama: Completion received ({:.2}s)",
248 (response_received_at - buffer_snapshotted_at).as_secs_f64()
249 );
250
251 let edits = if is_zeta {
252 let editable_range =
253 editable_range_override.expect("zeta model should have editable range");
254
255 log::trace!("ollama response: {}", ollama_response.response);
256
257 let response = clean_zeta_completion(&ollama_response.response);
258 match zeta1::parse_edits(&response, editable_range, &snapshot) {
259 Ok(edits) => edits,
260 Err(err) => {
261 log::warn!("Ollama zeta: Failed to parse response: {}", err);
262 vec![]
263 }
264 }
265 } else {
266 let completion: Arc<str> = clean_fim_completion(&ollama_response.response).into();
267 if completion.is_empty() {
268 vec![]
269 } else {
270 let cursor_offset = cursor_point.to_offset(&snapshot);
271 let anchor = snapshot.anchor_after(cursor_offset);
272 vec![(anchor..anchor, completion)]
273 }
274 };
275
276 anyhow::Ok(OllamaRequestOutput {
277 created_at: ollama_response.created_at,
278 edits,
279 snapshot,
280 response_received_at,
281 inputs,
282 buffer,
283 buffer_snapshotted_at,
284 })
285 });
286
287 cx.spawn(async move |cx: &mut gpui::AsyncApp| {
288 let output = result.await.context("Ollama edit prediction failed")?;
289 anyhow::Ok(Some(
290 EditPredictionResult::new(
291 EditPredictionId(output.created_at.into()),
292 &output.buffer,
293 &output.snapshot,
294 output.edits.into(),
295 None,
296 output.buffer_snapshotted_at,
297 output.response_received_at,
298 output.inputs,
299 cx,
300 )
301 .await,
302 ))
303 })
304 }
305}
306
307fn is_zeta_model(model: &str) -> bool {
308 model.to_lowercase().contains("zeta")
309}
310
311fn get_zeta_stop_tokens() -> Vec<String> {
312 vec![EDITABLE_REGION_END_MARKER.to_string(), "```".to_string()]
313}
314
315fn format_fim_prompt(model: &str, prefix: &str, suffix: &str) -> String {
316 let model_base = model.split(':').next().unwrap_or(model);
317
318 match model_base {
319 "codellama" | "code-llama" => {
320 format!("<PRE> {prefix} <SUF>{suffix} <MID>")
321 }
322 "starcoder" | "starcoder2" | "starcoderbase" => {
323 format!("<fim_prefix>{prefix}<fim_suffix>{suffix}<fim_middle>")
324 }
325 "deepseek-coder" | "deepseek-coder-v2" => {
326 format!("<｜fim▁begin｜>{prefix}<｜fim▁hole｜>{suffix}<｜fim▁end｜>")
327 }
328 "qwen2.5-coder" | "qwen-coder" | "qwen" => {
329 format!("<|fim_prefix|>{prefix}<|fim_suffix|>{suffix}<|fim_middle|>")
330 }
331 "codegemma" => {
332 format!("<|fim_prefix|>{prefix}<|fim_suffix|>{suffix}<|fim_middle|>")
333 }
334 "codestral" | "mistral" => {
335 format!("[SUFFIX]{suffix}[PREFIX]{prefix}")
336 }
337 "glm" | "glm-4" | "glm-4.5" => {
338 format!("<|code_prefix|>{prefix}<|code_suffix|>{suffix}<|code_middle|>")
339 }
340 _ => {
341 format!("<fim_prefix>{prefix}<fim_suffix>{suffix}<fim_middle>")
342 }
343 }
344}
345
346fn get_fim_stop_tokens() -> Vec<String> {
347 vec![
348 "<|endoftext|>".to_string(),
349 "<|file_separator|>".to_string(),
350 "<|fim_pad|>".to_string(),
351 "<|fim_prefix|>".to_string(),
352 "<|fim_middle|>".to_string(),
353 "<|fim_suffix|>".to_string(),
354 "<fim_prefix>".to_string(),
355 "<fim_middle>".to_string(),
356 "<fim_suffix>".to_string(),
357 "<PRE>".to_string(),
358 "<SUF>".to_string(),
359 "<MID>".to_string(),
360 "[PREFIX]".to_string(),
361 "[SUFFIX]".to_string(),
362 ]
363}
364
365fn clean_zeta_completion(mut response: &str) -> &str {
366 if let Some(last_newline_ix) = response.rfind('\n') {
367 let last_line = &response[last_newline_ix + 1..];
368 if EDITABLE_REGION_END_MARKER.starts_with(&last_line) {
369 response = &response[..last_newline_ix]
370 }
371 }
372 response
373}
374
375fn clean_fim_completion(response: &str) -> String {
376 let mut result = response.to_string();
377
378 let end_tokens = [
379 "<|endoftext|>",
380 "<|file_separator|>",
381 "<|fim_pad|>",
382 "<|fim_prefix|>",
383 "<|fim_middle|>",
384 "<|fim_suffix|>",
385 "<fim_prefix>",
386 "<fim_middle>",
387 "<fim_suffix>",
388 "<PRE>",
389 "<SUF>",
390 "<MID>",
391 "[PREFIX]",
392 "[SUFFIX]",
393 ];
394
395 for token in &end_tokens {
396 if let Some(pos) = result.find(token) {
397 result.truncate(pos);
398 }
399 }
400
401 result
402}