Max Brunsfeld , Oleksiy Syvokon , and Ben Kunkle created
Closes https://github.com/zed-industries/zed/issues/15968
Release Notes:
- Added the ability to use Ollama as an edit prediction provider
---------
Co-authored-by: Oleksiy Syvokon <oleksiy@zed.dev>
Co-authored-by: Ben Kunkle <ben@zed.dev>Cargo.lock | 2
assets/settings/default.json | 5
crates/agent_ui/src/agent_ui.rs | 1
crates/edit_prediction/src/edit_prediction.rs | 53
crates/edit_prediction/src/ollama.rs | 371 ++++
crates/edit_prediction/src/udiff.rs | 4
crates/edit_prediction/src/zed_edit_prediction_delegate.rs | 1
crates/edit_prediction/src/zeta1.rs | 51
crates/edit_prediction_ui/src/edit_prediction_button.rs | 62
crates/language/src/language.rs | 5
crates/language/src/language_settings.rs | 19
crates/language/src/text_diff.rs | 17
crates/ollama/src/ollama.rs | 3
crates/reqwest_client/Cargo.toml | 1
crates/reqwest_client/src/reqwest_client.rs | 11
crates/settings_content/src/language.rs | 49
crates/settings_ui/Cargo.toml | 1
crates/settings_ui/src/components.rs | 2
crates/settings_ui/src/components/ollama_model_picker.rs | 236 ++
crates/settings_ui/src/pages/edit_prediction_provider_setup.rs | 133 +
crates/settings_ui/src/settings_ui.rs | 4
crates/zed/src/zed/edit_prediction_registry.rs | 7
crates/zeta_prompt/src/zeta_prompt.rs | 35
23 files changed, 1,025 insertions(+), 48 deletions(-)
@@ -13957,6 +13957,7 @@ dependencies = [
"regex",
"serde",
"tokio",
+ "util",
"zed-reqwest",
]
@@ -15214,6 +15215,7 @@ dependencies = [
"heck 0.5.0",
"itertools 0.14.0",
"language",
+ "language_model",
"language_models",
"log",
"menu",
@@ -1500,6 +1500,11 @@
// to improve the service.
"privacy_mode": false,
},
+ "ollama": {
+ "api_url": "http://localhost:11434",
+ "model": "qwen2.5-coder:7b-base",
+ "max_output_tokens": 256,
+ },
// Whether edit predictions are enabled when editing text threads in the agent panel.
// This setting has no effect if globally disabled.
"enabled_in_text_threads": true,
@@ -401,6 +401,7 @@ fn update_command_palette_filter(cx: &mut App) {
}
EditPredictionProvider::Zed
| EditPredictionProvider::Codestral
+ | EditPredictionProvider::Ollama
| EditPredictionProvider::Sweep
| EditPredictionProvider::Mercury
| EditPredictionProvider::Experimental(_) => {
@@ -56,6 +56,7 @@ pub mod cursor_excerpt;
pub mod example_spec;
mod license_detection;
pub mod mercury;
+pub mod ollama;
mod onboarding_modal;
pub mod open_ai_response;
mod prediction;
@@ -76,6 +77,7 @@ use crate::capture_example::{
};
use crate::license_detection::LicenseDetectionWatcher;
use crate::mercury::Mercury;
+use crate::ollama::Ollama;
use crate::onboarding_modal::ZedPredictModal;
pub use crate::prediction::EditPrediction;
pub use crate::prediction::EditPredictionId;
@@ -153,6 +155,7 @@ pub struct EditPredictionStore {
edit_prediction_model: EditPredictionModel,
pub sweep_ai: SweepAi,
pub mercury: Mercury,
+ pub ollama: Ollama,
data_collection_choice: DataCollectionChoice,
reject_predictions_tx: mpsc::UnboundedSender<EditPredictionRejection>,
shown_predictions: VecDeque<EditPrediction>,
@@ -169,6 +172,7 @@ pub enum EditPredictionModel {
},
Sweep,
Mercury,
+ Ollama,
}
#[derive(Clone)]
@@ -307,17 +311,25 @@ impl ProjectState {
) {
self.cancelled_predictions.insert(pending_prediction.id);
- cx.spawn(async move |this, cx| {
- let Some(prediction_id) = pending_prediction.task.await else {
- return;
- };
+ if pending_prediction.drop_on_cancel {
+ drop(pending_prediction.task);
+ } else {
+ cx.spawn(async move |this, cx| {
+ let Some(prediction_id) = pending_prediction.task.await else {
+ return;
+ };
- this.update(cx, |this, _cx| {
- this.reject_prediction(prediction_id, EditPredictionRejectReason::Canceled, false);
+ this.update(cx, |this, _cx| {
+ this.reject_prediction(
+ prediction_id,
+ EditPredictionRejectReason::Canceled,
+ false,
+ );
+ })
+ .ok();
})
- .ok();
- })
- .detach()
+ .detach()
+ }
}
fn active_buffer(
@@ -399,6 +411,9 @@ impl PredictionRequestedBy {
struct PendingPrediction {
id: usize,
task: Task<Option<EditPredictionId>>,
+ /// If true, the task is dropped immediately on cancel (cancelling the HTTP request).
+ /// If false, the task is awaited to completion so rejection can be reported.
+ drop_on_cancel: bool,
}
/// A prediction from the perspective of a buffer.
@@ -632,6 +647,7 @@ impl EditPredictionStore {
},
sweep_ai: SweepAi::new(cx),
mercury: Mercury::new(cx),
+ ollama: Ollama::new(),
data_collection_choice,
reject_predictions_tx: reject_tx,
@@ -675,6 +691,9 @@ impl EditPredictionStore {
.with_down(IconName::ZedPredictDown)
.with_error(IconName::ZedPredictError)
}
+ EditPredictionModel::Ollama => {
+ edit_prediction_types::EditPredictionIconSet::new(IconName::AiOllama)
+ }
}
}
@@ -1207,7 +1226,7 @@ impl EditPredictionStore {
EditPredictionModel::Sweep => {
sweep_ai::edit_prediction_accepted(self, current_prediction, cx)
}
- EditPredictionModel::Mercury => {}
+ EditPredictionModel::Mercury | EditPredictionModel::Ollama => {}
EditPredictionModel::Zeta1 | EditPredictionModel::Zeta2 { .. } => {
zeta2::edit_prediction_accepted(self, current_prediction, cx)
}
@@ -1347,7 +1366,9 @@ impl EditPredictionStore {
return;
}
}
- EditPredictionModel::Sweep | EditPredictionModel::Mercury => return,
+ EditPredictionModel::Sweep
+ | EditPredictionModel::Mercury
+ | EditPredictionModel::Ollama => return,
}
self.reject_predictions_tx
@@ -1536,6 +1557,9 @@ impl EditPredictionStore {
-> Task<Result<Option<(EditPredictionResult, PredictionRequestedBy)>>>
+ 'static,
) {
+ let is_ollama = self.edit_prediction_model == EditPredictionModel::Ollama;
+ let drop_on_cancel = is_ollama;
+ let max_pending_predictions = if is_ollama { 1 } else { 2 };
let project_state = self.get_or_init_project(&project, cx);
let pending_prediction_id = project_state.next_pending_prediction_id;
project_state.next_pending_prediction_id += 1;
@@ -1650,16 +1674,18 @@ impl EditPredictionStore {
new_prediction_id
});
- if project_state.pending_predictions.len() <= 1 {
+ if project_state.pending_predictions.len() < max_pending_predictions {
project_state.pending_predictions.push(PendingPrediction {
id: pending_prediction_id,
task,
+ drop_on_cancel,
});
- } else if project_state.pending_predictions.len() == 2 {
+ } else {
let pending_prediction = project_state.pending_predictions.pop().unwrap();
project_state.pending_predictions.push(PendingPrediction {
id: pending_prediction_id,
task,
+ drop_on_cancel,
});
project_state.cancel_pending_prediction(pending_prediction, cx);
}
@@ -1794,6 +1820,7 @@ impl EditPredictionStore {
}
EditPredictionModel::Sweep => self.sweep_ai.request_prediction_with_sweep(inputs, cx),
EditPredictionModel::Mercury => self.mercury.request_prediction(inputs, cx),
+ EditPredictionModel::Ollama => self.ollama.request_prediction(inputs, cx),
};
cx.spawn(async move |this, cx| {
@@ -0,0 +1,371 @@
+use crate::{
+ EditPredictionId, EditPredictionModelInput, cursor_excerpt,
+ prediction::EditPredictionResult,
+ zeta1::{
+ self, MAX_CONTEXT_TOKENS as ZETA_MAX_CONTEXT_TOKENS,
+ MAX_EVENT_TOKENS as ZETA_MAX_EVENT_TOKENS,
+ },
+};
+use anyhow::{Context as _, Result};
+use futures::AsyncReadExt as _;
+use gpui::{App, AppContext as _, Entity, Task, http_client};
+use language::{
+ Anchor, Buffer, BufferSnapshot, OffsetRangeExt as _, ToOffset, ToPoint as _,
+ language_settings::all_language_settings,
+};
+use language_model::{LanguageModelProviderId, LanguageModelRegistry};
+use serde::{Deserialize, Serialize};
+use std::{path::Path, sync::Arc, time::Instant};
+use zeta_prompt::{
+ ZetaPromptInput,
+ zeta1::{EDITABLE_REGION_END_MARKER, format_zeta1_prompt},
+};
+
+const FIM_CONTEXT_TOKENS: usize = 512;
+
+pub struct Ollama;
+
+#[derive(Debug, Serialize)]
+struct OllamaGenerateRequest {
+ model: String,
+ prompt: String,
+ raw: bool,
+ stream: bool,
+ #[serde(skip_serializing_if = "Option::is_none")]
+ options: Option<OllamaGenerateOptions>,
+}
+
+#[derive(Debug, Serialize)]
+struct OllamaGenerateOptions {
+ #[serde(skip_serializing_if = "Option::is_none")]
+ num_predict: Option<u32>,
+ #[serde(skip_serializing_if = "Option::is_none")]
+ temperature: Option<f32>,
+ #[serde(skip_serializing_if = "Option::is_none")]
+ stop: Option<Vec<String>>,
+}
+
+#[derive(Debug, Deserialize)]
+struct OllamaGenerateResponse {
+ created_at: String,
+ response: String,
+}
+
+pub fn is_available(cx: &App) -> bool {
+ let ollama_provider_id = LanguageModelProviderId::new("ollama");
+ LanguageModelRegistry::read_global(cx)
+ .provider(&ollama_provider_id)
+ .is_some_and(|provider| provider.is_authenticated(cx))
+}
+
+/// Output from the Ollama HTTP request, containing all data needed to create the prediction result.
+struct OllamaRequestOutput {
+ created_at: String,
+ edits: Vec<(std::ops::Range<Anchor>, Arc<str>)>,
+ snapshot: BufferSnapshot,
+ response_received_at: Instant,
+ inputs: ZetaPromptInput,
+ buffer: Entity<Buffer>,
+ buffer_snapshotted_at: Instant,
+}
+
+impl Ollama {
+ pub fn new() -> Self {
+ Self
+ }
+
+ pub fn request_prediction(
+ &self,
+ EditPredictionModelInput {
+ buffer,
+ snapshot,
+ position,
+ events,
+ ..
+ }: EditPredictionModelInput,
+ cx: &mut App,
+ ) -> Task<Result<Option<EditPredictionResult>>> {
+ let settings = &all_language_settings(None, cx).edit_predictions.ollama;
+ let Some(model) = settings.model.clone() else {
+ return Task::ready(Ok(None));
+ };
+ let max_output_tokens = settings.max_output_tokens;
+ let api_url = settings.api_url.clone();
+
+ log::debug!("Ollama: Requesting completion (model: {})", model);
+
+ let full_path: Arc<Path> = snapshot
+ .file()
+ .map(|file| file.full_path(cx))
+ .unwrap_or_else(|| "untitled".into())
+ .into();
+
+ let http_client = cx.http_client();
+ let cursor_point = position.to_point(&snapshot);
+ let buffer_snapshotted_at = Instant::now();
+
+ let is_zeta = is_zeta_model(&model);
+
+ let result = cx.background_spawn(async move {
+ // For zeta models, use the dedicated zeta1 functions which handle their own
+ // range computation with the correct token limits.
+ let (prompt, stop_tokens, editable_range_override, inputs) = if is_zeta {
+ let path_str = full_path.to_string_lossy();
+ let input_excerpt = zeta1::excerpt_for_cursor_position(
+ cursor_point,
+ &path_str,
+ &snapshot,
+ max_output_tokens as usize,
+ ZETA_MAX_CONTEXT_TOKENS,
+ );
+ let input_events = zeta1::prompt_for_events(&events, ZETA_MAX_EVENT_TOKENS);
+ let prompt = format_zeta1_prompt(&input_events, &input_excerpt.prompt);
+ let editable_offset_range = input_excerpt.editable_range.to_offset(&snapshot);
+ let context_offset_range = input_excerpt.context_range.to_offset(&snapshot);
+ let stop_tokens = get_zeta_stop_tokens();
+
+ let inputs = ZetaPromptInput {
+ events,
+ related_files: Vec::new(),
+ cursor_offset_in_excerpt: cursor_point.to_offset(&snapshot)
+ - context_offset_range.start,
+ cursor_path: full_path.clone(),
+ cursor_excerpt: snapshot
+ .text_for_range(input_excerpt.context_range.clone())
+ .collect::<String>()
+ .into(),
+ editable_range_in_excerpt: (editable_offset_range.start
+ - context_offset_range.start)
+ ..(editable_offset_range.end - context_offset_range.start),
+ excerpt_start_row: Some(input_excerpt.context_range.start.row),
+ };
+
+ (prompt, stop_tokens, Some(editable_offset_range), inputs)
+ } else {
+ let (excerpt_range, _) =
+ cursor_excerpt::editable_and_context_ranges_for_cursor_position(
+ cursor_point,
+ &snapshot,
+ FIM_CONTEXT_TOKENS,
+ 0,
+ );
+ let excerpt_offset_range = excerpt_range.to_offset(&snapshot);
+ let cursor_offset = cursor_point.to_offset(&snapshot);
+
+ let inputs = ZetaPromptInput {
+ events,
+ related_files: Vec::new(),
+ cursor_offset_in_excerpt: cursor_offset - excerpt_offset_range.start,
+ editable_range_in_excerpt: cursor_offset - excerpt_offset_range.start
+ ..cursor_offset - excerpt_offset_range.start,
+ cursor_path: full_path.clone(),
+ excerpt_start_row: Some(excerpt_range.start.row),
+ cursor_excerpt: snapshot
+ .text_for_range(excerpt_range)
+ .collect::<String>()
+ .into(),
+ };
+
+ let prefix = inputs.cursor_excerpt[..inputs.cursor_offset_in_excerpt].to_string();
+ let suffix = inputs.cursor_excerpt[inputs.cursor_offset_in_excerpt..].to_string();
+ let prompt = format_fim_prompt(&model, &prefix, &suffix);
+ let stop_tokens = get_fim_stop_tokens();
+
+ (prompt, stop_tokens, None, inputs)
+ };
+
+ let request = OllamaGenerateRequest {
+ model: model.clone(),
+ prompt,
+ raw: true,
+ stream: false,
+ options: Some(OllamaGenerateOptions {
+ num_predict: Some(max_output_tokens),
+ temperature: Some(0.2),
+ stop: Some(stop_tokens),
+ }),
+ };
+
+ let request_body = serde_json::to_string(&request)?;
+ let http_request = http_client::Request::builder()
+ .method(http_client::Method::POST)
+ .uri(format!("{}/api/generate", api_url))
+ .header("Content-Type", "application/json")
+ .body(http_client::AsyncBody::from(request_body))?;
+
+ let mut response = http_client.send(http_request).await?;
+ let status = response.status();
+
+ log::debug!("Ollama: Response status: {}", status);
+
+ if !status.is_success() {
+ let mut body = String::new();
+ response.body_mut().read_to_string(&mut body).await?;
+ return Err(anyhow::anyhow!("Ollama API error: {} - {}", status, body));
+ }
+
+ let mut body = String::new();
+ response.body_mut().read_to_string(&mut body).await?;
+
+ let ollama_response: OllamaGenerateResponse =
+ serde_json::from_str(&body).context("Failed to parse Ollama response")?;
+
+ let response_received_at = Instant::now();
+
+ log::debug!(
+ "Ollama: Completion received ({:.2}s)",
+ (response_received_at - buffer_snapshotted_at).as_secs_f64()
+ );
+
+ let edits = if is_zeta {
+ let editable_range =
+ editable_range_override.expect("zeta model should have editable range");
+
+ log::trace!("ollama response: {}", ollama_response.response);
+
+ let response = clean_zeta_completion(&ollama_response.response);
+ match zeta1::parse_edits(&response, editable_range, &snapshot) {
+ Ok(edits) => edits,
+ Err(err) => {
+ log::warn!("Ollama zeta: Failed to parse response: {}", err);
+ vec![]
+ }
+ }
+ } else {
+ let completion: Arc<str> = clean_fim_completion(&ollama_response.response).into();
+ if completion.is_empty() {
+ vec![]
+ } else {
+ let cursor_offset = cursor_point.to_offset(&snapshot);
+ let anchor = snapshot.anchor_after(cursor_offset);
+ vec![(anchor..anchor, completion)]
+ }
+ };
+
+ anyhow::Ok(OllamaRequestOutput {
+ created_at: ollama_response.created_at,
+ edits,
+ snapshot,
+ response_received_at,
+ inputs,
+ buffer,
+ buffer_snapshotted_at,
+ })
+ });
+
+ cx.spawn(async move |cx: &mut gpui::AsyncApp| {
+ let output = result.await.context("Ollama edit prediction failed")?;
+ anyhow::Ok(Some(
+ EditPredictionResult::new(
+ EditPredictionId(output.created_at.into()),
+ &output.buffer,
+ &output.snapshot,
+ output.edits.into(),
+ None,
+ output.buffer_snapshotted_at,
+ output.response_received_at,
+ output.inputs,
+ cx,
+ )
+ .await,
+ ))
+ })
+ }
+}
+
+fn is_zeta_model(model: &str) -> bool {
+ model.to_lowercase().contains("zeta")
+}
+
+fn get_zeta_stop_tokens() -> Vec<String> {
+ vec![EDITABLE_REGION_END_MARKER.to_string(), "```".to_string()]
+}
+
+fn format_fim_prompt(model: &str, prefix: &str, suffix: &str) -> String {
+ let model_base = model.split(':').next().unwrap_or(model);
+
+ match model_base {
+ "codellama" | "code-llama" => {
+ format!("<PRE> {prefix} <SUF>{suffix} <MID>")
+ }
+ "starcoder" | "starcoder2" | "starcoderbase" => {
+ format!("<fim_prefix>{prefix}<fim_suffix>{suffix}<fim_middle>")
+ }
+ "deepseek-coder" | "deepseek-coder-v2" => {
+ format!("<ο½fimβbeginο½>{prefix}<ο½fimβholeο½>{suffix}<ο½fimβendο½>")
+ }
+ "qwen2.5-coder" | "qwen-coder" | "qwen" => {
+ format!("<|fim_prefix|>{prefix}<|fim_suffix|>{suffix}<|fim_middle|>")
+ }
+ "codegemma" => {
+ format!("<|fim_prefix|>{prefix}<|fim_suffix|>{suffix}<|fim_middle|>")
+ }
+ "codestral" | "mistral" => {
+ format!("[SUFFIX]{suffix}[PREFIX]{prefix}")
+ }
+ "glm" | "glm-4" | "glm-4.5" => {
+ format!("<|code_prefix|>{prefix}<|code_suffix|>{suffix}<|code_middle|>")
+ }
+ _ => {
+ format!("<fim_prefix>{prefix}<fim_suffix>{suffix}<fim_middle>")
+ }
+ }
+}
+
+fn get_fim_stop_tokens() -> Vec<String> {
+ vec![
+ "<|endoftext|>".to_string(),
+ "<|file_separator|>".to_string(),
+ "<|fim_pad|>".to_string(),
+ "<|fim_prefix|>".to_string(),
+ "<|fim_middle|>".to_string(),
+ "<|fim_suffix|>".to_string(),
+ "<fim_prefix>".to_string(),
+ "<fim_middle>".to_string(),
+ "<fim_suffix>".to_string(),
+ "<PRE>".to_string(),
+ "<SUF>".to_string(),
+ "<MID>".to_string(),
+ "[PREFIX]".to_string(),
+ "[SUFFIX]".to_string(),
+ ]
+}
+
+fn clean_zeta_completion(mut response: &str) -> &str {
+ if let Some(last_newline_ix) = response.rfind('\n') {
+ let last_line = &response[last_newline_ix + 1..];
+ if EDITABLE_REGION_END_MARKER.starts_with(&last_line) {
+ response = &response[..last_newline_ix]
+ }
+ }
+ response
+}
+
+fn clean_fim_completion(response: &str) -> String {
+ let mut result = response.to_string();
+
+ let end_tokens = [
+ "<|endoftext|>",
+ "<|file_separator|>",
+ "<|fim_pad|>",
+ "<|fim_prefix|>",
+ "<|fim_middle|>",
+ "<|fim_suffix|>",
+ "<fim_prefix>",
+ "<fim_middle>",
+ "<fim_suffix>",
+ "<PRE>",
+ "<SUF>",
+ "<MID>",
+ "[PREFIX]",
+ "[SUFFIX]",
+ ];
+
+ for token in &end_tokens {
+ if let Some(pos) = result.find(token) {
+ result.truncate(pos);
+ }
+ }
+
+ result
+}
@@ -676,9 +676,9 @@ pub enum DiffLine<'a> {
#[derive(Debug, PartialEq)]
pub struct HunkLocation {
- start_line_old: u32,
+ pub start_line_old: u32,
count_old: u32,
- start_line_new: u32,
+ pub start_line_new: u32,
count_new: u32,
}
@@ -76,6 +76,7 @@ impl EditPredictionDelegate for ZedEditPredictionDelegate {
.with_down(IconName::ZedPredictDown)
.with_error(IconName::ZedPredictError)
}
+ EditPredictionModel::Ollama => EditPredictionIconSet::new(IconName::AiOllama),
}
}
@@ -6,7 +6,7 @@ use crate::{
cursor_excerpt::{editable_and_context_ranges_for_cursor_position, guess_token_count},
prediction::EditPredictionResult,
};
-use anyhow::{Context as _, Result};
+use anyhow::Result;
use cloud_llm_client::{
PredictEditsBody, PredictEditsGitInfo, PredictEditsRequestTrigger, PredictEditsResponse,
};
@@ -19,12 +19,13 @@ use project::{Project, ProjectPath};
use release_channel::AppVersion;
use text::Bias;
use workspace::notifications::{ErrorMessagePrompt, NotificationId, show_app_notification};
-use zeta_prompt::{Event, ZetaPromptInput};
-
-const CURSOR_MARKER: &str = "<|user_cursor_is_here|>";
-const START_OF_FILE_MARKER: &str = "<|start_of_file|>";
-const EDITABLE_REGION_START_MARKER: &str = "<|editable_region_start|>";
-const EDITABLE_REGION_END_MARKER: &str = "<|editable_region_end|>";
+use zeta_prompt::{
+ Event, ZetaPromptInput,
+ zeta1::{
+ CURSOR_MARKER, EDITABLE_REGION_END_MARKER, EDITABLE_REGION_START_MARKER,
+ START_OF_FILE_MARKER,
+ },
+};
pub(crate) const MAX_CONTEXT_TOKENS: usize = 150;
pub(crate) const MAX_REWRITE_TOKENS: usize = 350;
@@ -265,7 +266,7 @@ fn process_completion_response(
let output_excerpt = output_excerpt.clone();
let editable_range = editable_range.clone();
let snapshot = snapshot.clone();
- async move { parse_edits(output_excerpt, editable_range, &snapshot) }
+ async move { parse_edits(output_excerpt.as_ref(), editable_range, &snapshot) }
})
.await?
.into();
@@ -286,8 +287,8 @@ fn process_completion_response(
})
}
-fn parse_edits(
- output_excerpt: Arc<str>,
+pub(crate) fn parse_edits(
+ output_excerpt: &str,
editable_range: Range<usize>,
snapshot: &BufferSnapshot,
) -> Result<Vec<(Range<Anchor>, Arc<str>)>> {
@@ -297,8 +298,8 @@ fn parse_edits(
.match_indices(EDITABLE_REGION_START_MARKER)
.collect::<Vec<_>>();
anyhow::ensure!(
- start_markers.len() == 1,
- "expected exactly one start marker, found {}",
+ start_markers.len() <= 1,
+ "expected at most one start marker, found {}",
start_markers.len()
);
@@ -306,8 +307,8 @@ fn parse_edits(
.match_indices(EDITABLE_REGION_END_MARKER)
.collect::<Vec<_>>();
anyhow::ensure!(
- end_markers.len() == 1,
- "expected exactly one end marker, found {}",
+ end_markers.len() <= 1,
+ "expected at most one end marker, found {}",
end_markers.len()
);
@@ -320,16 +321,16 @@ fn parse_edits(
sof_markers.len()
);
- let codefence_start = start_markers[0].0;
- let content = &content[codefence_start..];
-
- let newline_ix = content.find('\n').context("could not find newline")?;
- let content = &content[newline_ix + 1..];
+ let content_start = start_markers
+ .first()
+ .map(|e| e.0 + EDITABLE_REGION_START_MARKER.len() + 1) // +1 to skip \n after marker
+ .unwrap_or(0);
+ let content_end = end_markers
+ .first()
+ .map(|e| e.0.saturating_sub(1)) // -1 to exclude \n before marker
+ .unwrap_or(content.strip_suffix("\n").unwrap_or(&content).len());
- let codefence_end = content
- .rfind(&format!("\n{EDITABLE_REGION_END_MARKER}"))
- .context("could not find end marker")?;
- let new_text = &content[..codefence_end];
+ let new_text = &content[content_start..content_end];
let old_text = snapshot
.text_for_range(editable_range.clone())
@@ -515,6 +516,10 @@ pub fn gather_context(
})
}
+pub(crate) fn prompt_for_events(events: &[Arc<Event>], max_tokens: usize) -> String {
+ prompt_for_events_impl(events, max_tokens).0
+}
+
fn prompt_for_events_impl(events: &[Arc<Event>], mut remaining_tokens: usize) -> (String, usize) {
let mut result = String::new();
for (ix, event) in events.iter().rev().enumerate() {
@@ -89,9 +89,9 @@ impl Render for EditPredictionButton {
return div().hidden();
}
- let all_language_settings = all_language_settings(None, cx);
+ let language_settings = all_language_settings(None, cx);
- match all_language_settings.edit_predictions.provider {
+ match language_settings.edit_predictions.provider {
EditPredictionProvider::Copilot => {
let Some(copilot) = EditPredictionStore::try_global(cx)
.and_then(|store| store.read(cx).copilot_for_project(&self.project.upgrade()?))
@@ -303,6 +303,60 @@ impl Render for EditPredictionButton {
.with_handle(self.popover_menu_handle.clone()),
)
}
+ EditPredictionProvider::Ollama => {
+ let enabled = self.editor_enabled.unwrap_or(true);
+ let this = cx.weak_entity();
+
+ div().child(
+ PopoverMenu::new("ollama")
+ .menu(move |window, cx| {
+ this.update(cx, |this, cx| {
+ this.build_edit_prediction_context_menu(
+ EditPredictionProvider::Ollama,
+ window,
+ cx,
+ )
+ })
+ .ok()
+ })
+ .anchor(Corner::BottomRight)
+ .trigger_with_tooltip(
+ IconButton::new("ollama-icon", IconName::AiOllama)
+ .shape(IconButtonShape::Square)
+ .when(!enabled, |this| {
+ this.indicator(Indicator::dot().color(Color::Ignored))
+ .indicator_border_color(Some(
+ cx.theme().colors().status_bar_background,
+ ))
+ }),
+ move |_window, cx| {
+ let settings = all_language_settings(None, cx);
+ let tooltip_meta = match settings
+ .edit_predictions
+ .ollama
+ .model
+ .as_deref()
+ {
+ Some(model) if !model.trim().is_empty() => {
+ format!("Powered by Ollama ({model})")
+ }
+ _ => {
+ "Ollama model not configured β configure a model before use"
+ .to_string()
+ }
+ };
+
+ Tooltip::with_meta(
+ "Edit Prediction",
+ Some(&ToggleMenu),
+ tooltip_meta,
+ cx,
+ )
+ },
+ )
+ .with_handle(self.popover_menu_handle.clone()),
+ )
+ }
provider @ (EditPredictionProvider::Experimental(_)
| EditPredictionProvider::Zed
| EditPredictionProvider::Sweep
@@ -1318,6 +1372,10 @@ pub fn get_available_providers(cx: &mut App) -> Vec<EditPredictionProvider> {
providers.push(EditPredictionProvider::Codestral);
}
+ if edit_prediction::ollama::is_available(cx) {
+ providers.push(EditPredictionProvider::Ollama);
+ }
+
if cx.has_flag::<SweepFeatureFlag>()
&& edit_prediction::sweep_ai::sweep_api_token(cx)
.read(cx)
@@ -66,8 +66,9 @@ use syntax_map::{QueryCursorHandle, SyntaxSnapshot};
use task::RunnableTag;
pub use task_context::{ContextLocation, ContextProvider, RunnableRange};
pub use text_diff::{
- DiffOptions, apply_diff_patch, line_diff, text_diff, text_diff_with_options, unified_diff,
- unified_diff_with_context, unified_diff_with_offsets, word_diff_ranges,
+ DiffOptions, apply_diff_patch, apply_reversed_diff_patch, line_diff, text_diff,
+ text_diff_with_options, unified_diff, unified_diff_with_context, unified_diff_with_offsets,
+ word_diff_ranges,
};
use theme::SyntaxTheme;
pub use toolchain::{
@@ -390,6 +390,8 @@ pub struct EditPredictionSettings {
pub codestral: CodestralSettings,
/// Settings specific to Sweep.
pub sweep: SweepSettings,
+ /// Settings specific to Ollama.
+ pub ollama: OllamaSettings,
/// Whether edit predictions are enabled in the assistant panel.
/// This setting has no effect if globally disabled.
pub enabled_in_text_threads: bool,
@@ -448,6 +450,16 @@ pub struct SweepSettings {
pub privacy_mode: bool,
}
+#[derive(Clone, Debug, Default)]
+pub struct OllamaSettings {
+ /// Model to use for completions.
+ pub model: Option<String>,
+ /// Maximum tokens to generate.
+ pub max_output_tokens: u32,
+ /// Custom API URL to use for Ollama.
+ pub api_url: Arc<str>,
+}
+
impl AllLanguageSettings {
/// Returns the [`LanguageSettings`] for the language with the specified name.
pub fn language<'a>(
@@ -678,6 +690,12 @@ impl settings::Settings for AllLanguageSettings {
let sweep_settings = SweepSettings {
privacy_mode: sweep.privacy_mode.unwrap(),
};
+ let ollama = edit_predictions.ollama.unwrap();
+ let ollama_settings = OllamaSettings {
+ model: ollama.model.map(|m| m.0),
+ max_output_tokens: ollama.max_output_tokens.unwrap(),
+ api_url: ollama.api_url.unwrap().into(),
+ };
let enabled_in_text_threads = edit_predictions.enabled_in_text_threads.unwrap();
@@ -717,6 +735,7 @@ impl settings::Settings for AllLanguageSettings {
copilot: copilot_settings,
codestral: codestral_settings,
sweep: sweep_settings,
+ ollama: ollama_settings,
enabled_in_text_threads,
examples_dir: edit_predictions.examples_dir,
example_capture_rate: edit_predictions.example_capture_rate,
@@ -301,6 +301,12 @@ pub fn apply_diff_patch(base_text: &str, patch: &str) -> Result<String, anyhow::
result.map_err(|err| anyhow!(err))
}
+pub fn apply_reversed_diff_patch(base_text: &str, patch: &str) -> Result<String, anyhow::Error> {
+ let patch = diffy::Patch::from_str(patch).context("Failed to parse patch")?;
+ let reversed = patch.reverse();
+ diffy::apply(base_text, &reversed).map_err(|err| anyhow!(err))
+}
+
fn should_perform_word_diff_within_hunk(
old_row_range: &Range<u32>,
old_byte_range: &Range<usize>,
@@ -462,6 +468,17 @@ mod tests {
assert_eq!(apply_diff_patch(old_text, &patch).unwrap(), new_text);
}
+ #[test]
+ fn test_apply_reversed_diff_patch() {
+ let old_text = "one two\nthree four five\nsix seven eight nine\nten\n";
+ let new_text = "one two\nthree FOUR five\nsix SEVEN eight nine\nten\nELEVEN\n";
+ let patch = unified_diff(old_text, new_text);
+ assert_eq!(
+ apply_reversed_diff_patch(new_text, &patch).unwrap(),
+ old_text
+ );
+ }
+
#[test]
fn test_unified_diff_with_offsets() {
let old_text = "foo\nbar\nbaz\n";
@@ -1,4 +1,5 @@
-use anyhow::{Context as _, Result};
+use anyhow::{Context, Result};
+
use futures::{AsyncBufReadExt, AsyncReadExt, StreamExt, io::BufReader, stream::BoxStream};
use http_client::{AsyncBody, HttpClient, HttpRequestExt, Method, Request as HttpRequest};
use serde::{Deserialize, Serialize};
@@ -26,6 +26,7 @@ log.workspace = true
tokio = { workspace = true, features = ["rt", "rt-multi-thread"] }
regex.workspace = true
reqwest.workspace = true
+util.workspace = true
[dev-dependencies]
gpui.workspace = true
@@ -2,6 +2,8 @@ use std::error::Error;
use std::sync::{LazyLock, OnceLock};
use std::{borrow::Cow, mem, pin::Pin, task::Poll, time::Duration};
+use util::defer;
+
use anyhow::anyhow;
use bytes::{BufMut, Bytes, BytesMut};
use futures::{AsyncRead, FutureExt as _, TryStreamExt as _};
@@ -249,10 +251,11 @@ impl http_client::HttpClient for ReqwestClient {
let handle = self.handle.clone();
async move {
- let mut response = handle
- .spawn(async { request.send().await })
- .await?
- .map_err(redact_error)?;
+ let join_handle = handle.spawn(async { request.send().await });
+ let abort_handle = join_handle.abort_handle();
+ let _abort_on_drop = defer(move || abort_handle.abort());
+
+ let mut response = join_handle.await?.map_err(redact_error)?;
let headers = mem::take(response.headers_mut());
let mut builder = http::Response::builder()
@@ -84,6 +84,7 @@ pub enum EditPredictionProvider {
Supermaven,
Zed,
Codestral,
+ Ollama,
Sweep,
Mercury,
Experimental(&'static str),
@@ -104,6 +105,7 @@ impl<'de> Deserialize<'de> for EditPredictionProvider {
Supermaven,
Zed,
Codestral,
+ Ollama,
Sweep,
Mercury,
Experimental(String),
@@ -115,6 +117,7 @@ impl<'de> Deserialize<'de> for EditPredictionProvider {
Content::Supermaven => EditPredictionProvider::Supermaven,
Content::Zed => EditPredictionProvider::Zed,
Content::Codestral => EditPredictionProvider::Codestral,
+ Content::Ollama => EditPredictionProvider::Ollama,
Content::Sweep => EditPredictionProvider::Sweep,
Content::Mercury => EditPredictionProvider::Mercury,
Content::Experimental(name)
@@ -142,6 +145,7 @@ impl EditPredictionProvider {
| EditPredictionProvider::Copilot
| EditPredictionProvider::Supermaven
| EditPredictionProvider::Codestral
+ | EditPredictionProvider::Ollama
| EditPredictionProvider::Sweep
| EditPredictionProvider::Mercury
| EditPredictionProvider::Experimental(_) => false,
@@ -160,6 +164,7 @@ impl EditPredictionProvider {
EXPERIMENTAL_ZETA2_EDIT_PREDICTION_PROVIDER_NAME,
) => Some("Zeta2"),
EditPredictionProvider::None | EditPredictionProvider::Experimental(_) => None,
+ EditPredictionProvider::Ollama => Some("Ollama"),
}
}
}
@@ -183,6 +188,8 @@ pub struct EditPredictionSettingsContent {
pub codestral: Option<CodestralSettingsContent>,
/// Settings specific to Sweep.
pub sweep: Option<SweepSettingsContent>,
+ /// Settings specific to Ollama.
+ pub ollama: Option<OllamaEditPredictionSettingsContent>,
/// Whether edit predictions are enabled in the assistant prompt editor.
/// This has no effect if globally disabled.
pub enabled_in_text_threads: Option<bool>,
@@ -242,6 +249,48 @@ pub struct SweepSettingsContent {
pub privacy_mode: Option<bool>,
}
+/// Ollama model name for edit predictions.
+#[with_fallible_options]
+#[derive(Clone, Debug, Serialize, Deserialize, JsonSchema, MergeFrom, PartialEq, Eq)]
+#[serde(transparent)]
+pub struct OllamaModelName(pub String);
+
+impl AsRef<str> for OllamaModelName {
+ fn as_ref(&self) -> &str {
+ &self.0
+ }
+}
+
+impl From<String> for OllamaModelName {
+ fn from(value: String) -> Self {
+ Self(value)
+ }
+}
+
+impl From<OllamaModelName> for String {
+ fn from(value: OllamaModelName) -> Self {
+ value.0
+ }
+}
+
+#[with_fallible_options]
+#[derive(Clone, Debug, Default, Serialize, Deserialize, JsonSchema, MergeFrom, PartialEq)]
+pub struct OllamaEditPredictionSettingsContent {
+ /// Model to use for completions.
+ ///
+ /// Default: none
+ pub model: Option<OllamaModelName>,
+ /// Maximum tokens to generate for FIM models.
+ /// This setting does not apply to sweep models.
+ ///
+ /// Default: 256
+ pub max_output_tokens: Option<u32>,
+ /// Api URL to use for completions.
+ ///
+ /// Default: "http://localhost:11434"
+ pub api_url: Option<String>,
+}
+
/// The mode in which edit predictions should be displayed.
#[derive(
Copy,
@@ -31,6 +31,7 @@ fuzzy.workspace = true
gpui.workspace = true
heck.workspace = true
itertools.workspace = true
+language_model.workspace = true
language_models.workspace = true
language.workspace = true
log.workspace = true
@@ -3,6 +3,7 @@ mod font_picker;
mod icon_theme_picker;
mod input_field;
mod number_field;
+mod ollama_model_picker;
mod section_items;
mod theme_picker;
@@ -11,5 +12,6 @@ pub use font_picker::font_picker;
pub use icon_theme_picker::icon_theme_picker;
pub use input_field::*;
pub use number_field::*;
+pub use ollama_model_picker::render_ollama_model_picker;
pub use section_items::*;
pub use theme_picker::theme_picker;
@@ -0,0 +1,236 @@
+use std::sync::Arc;
+
+use fuzzy::StringMatch;
+use gpui::{AnyElement, App, Context, DismissEvent, ReadGlobal, SharedString, Task, Window, px};
+use language_model::{LanguageModelProviderId, LanguageModelRegistry};
+use picker::{Picker, PickerDelegate};
+use settings::SettingsStore;
+use ui::{ListItem, ListItemSpacing, PopoverMenu, prelude::*};
+use util::ResultExt;
+
+use crate::{
+ SettingField, SettingsFieldMetadata, SettingsUiFile, render_picker_trigger_button,
+ update_settings_file,
+};
+
+type OllamaModelPicker = Picker<OllamaModelPickerDelegate>;
+
+struct OllamaModelPickerDelegate {
+ models: Vec<SharedString>,
+ filtered_models: Vec<StringMatch>,
+ selected_index: usize,
+ on_model_changed: Arc<dyn Fn(SharedString, &mut Window, &mut App) + 'static>,
+}
+
+impl OllamaModelPickerDelegate {
+ fn new(
+ current_model: SharedString,
+ on_model_changed: impl Fn(SharedString, &mut Window, &mut App) + 'static,
+ cx: &mut Context<OllamaModelPicker>,
+ ) -> Self {
+ let mut models = Self::fetch_ollama_models(cx);
+
+ let current_in_list = models.contains(¤t_model);
+ if !current_model.is_empty() && !current_in_list {
+ models.insert(0, current_model.clone());
+ }
+
+ let selected_index = models
+ .iter()
+ .position(|model| *model == current_model)
+ .unwrap_or(0);
+
+ let filtered_models = models
+ .iter()
+ .enumerate()
+ .map(|(index, model)| StringMatch {
+ candidate_id: index,
+ string: model.to_string(),
+ positions: Vec::new(),
+ score: 0.0,
+ })
+ .collect();
+
+ Self {
+ models,
+ filtered_models,
+ selected_index,
+ on_model_changed: Arc::new(on_model_changed),
+ }
+ }
+
+ fn fetch_ollama_models(cx: &mut App) -> Vec<SharedString> {
+ let ollama_provider_id = LanguageModelProviderId::new("ollama");
+
+ let Some(provider) = LanguageModelRegistry::read_global(cx).provider(&ollama_provider_id)
+ else {
+ return Vec::new();
+ };
+
+ // Re-fetch models in case ollama has been started or updated since
+ // Zed was launched.
+ provider.authenticate(cx).detach_and_log_err(cx);
+
+ let mut models: Vec<SharedString> = provider
+ .provided_models(cx)
+ .into_iter()
+ .map(|model| SharedString::from(model.id().0.to_string()))
+ .collect();
+
+ models.sort();
+ models
+ }
+}
+
+impl PickerDelegate for OllamaModelPickerDelegate {
+ type ListItem = AnyElement;
+
+ fn match_count(&self) -> usize {
+ self.filtered_models.len()
+ }
+
+ fn selected_index(&self) -> usize {
+ self.selected_index
+ }
+
+ fn set_selected_index(
+ &mut self,
+ ix: usize,
+ _: &mut Window,
+ cx: &mut Context<OllamaModelPicker>,
+ ) {
+ self.selected_index = ix.min(self.filtered_models.len().saturating_sub(1));
+ cx.notify();
+ }
+
+ fn placeholder_text(&self, _window: &mut Window, _cx: &mut App) -> Arc<str> {
+ "Search modelsβ¦".into()
+ }
+
+ fn update_matches(
+ &mut self,
+ query: String,
+ _window: &mut Window,
+ cx: &mut Context<OllamaModelPicker>,
+ ) -> Task<()> {
+ let query_lower = query.to_lowercase();
+
+ self.filtered_models = self
+ .models
+ .iter()
+ .enumerate()
+ .filter(|(_, model)| query.is_empty() || model.to_lowercase().contains(&query_lower))
+ .map(|(index, model)| StringMatch {
+ candidate_id: index,
+ string: model.to_string(),
+ positions: Vec::new(),
+ score: 0.0,
+ })
+ .collect();
+
+ self.selected_index = 0;
+ cx.notify();
+
+ Task::ready(())
+ }
+
+ fn confirm(
+ &mut self,
+ _secondary: bool,
+ window: &mut Window,
+ cx: &mut Context<OllamaModelPicker>,
+ ) {
+ let Some(model_match) = self.filtered_models.get(self.selected_index) else {
+ return;
+ };
+
+ (self.on_model_changed)(model_match.string.clone().into(), window, cx);
+ cx.emit(DismissEvent);
+ }
+
+ fn dismissed(&mut self, window: &mut Window, cx: &mut Context<OllamaModelPicker>) {
+ cx.defer_in(window, |picker, window, cx| {
+ picker.set_query("", window, cx);
+ });
+ cx.emit(DismissEvent);
+ }
+
+ fn render_match(
+ &self,
+ ix: usize,
+ selected: bool,
+ _window: &mut Window,
+ _cx: &mut Context<OllamaModelPicker>,
+ ) -> Option<Self::ListItem> {
+ let model_match = self.filtered_models.get(ix)?;
+
+ Some(
+ ListItem::new(ix)
+ .inset(true)
+ .spacing(ListItemSpacing::Sparse)
+ .toggle_state(selected)
+ .child(Label::new(model_match.string.clone()))
+ .into_any_element(),
+ )
+ }
+}
+
+pub fn render_ollama_model_picker(
+ field: SettingField<settings::OllamaModelName>,
+ file: SettingsUiFile,
+ _metadata: Option<&SettingsFieldMetadata>,
+ _window: &mut Window,
+ cx: &mut App,
+) -> AnyElement {
+ let (_, value) = SettingsStore::global(cx).get_value_from_file(file.to_settings(), field.pick);
+ let current_value: SharedString = value
+ .map(|m| m.0.clone().into())
+ .unwrap_or_else(|| "".into());
+
+ PopoverMenu::new("ollama-model-picker")
+ .trigger(render_picker_trigger_button(
+ "ollama_model_picker_trigger".into(),
+ if current_value.is_empty() {
+ "Select a modelβ¦".into()
+ } else {
+ current_value.clone()
+ },
+ ))
+ .menu(move |window, cx| {
+ Some(cx.new(|cx| {
+ let file = file.clone();
+ let current_value = current_value.clone();
+ let delegate = OllamaModelPickerDelegate::new(
+ current_value,
+ move |model_name, window, cx| {
+ update_settings_file(
+ file.clone(),
+ field.json_path,
+ window,
+ cx,
+ move |settings, _cx| {
+ (field.write)(
+ settings,
+ Some(settings::OllamaModelName(model_name.to_string())),
+ );
+ },
+ )
+ .log_err();
+ },
+ cx,
+ );
+
+ Picker::uniform_list(delegate, window, cx)
+ .show_scrollbar(true)
+ .width(rems_from_px(210.))
+ .max_height(Some(rems(18.).into()))
+ }))
+ })
+ .anchor(gpui::Corner::TopLeft)
+ .offset(gpui::Point {
+ x: px(0.0),
+ y: px(2.0),
+ })
+ .with_handle(ui::PopoverMenuHandle::default())
+ .into_any_element()
+}
@@ -12,6 +12,9 @@ use settings::Settings as _;
use ui::{ButtonLink, ConfiguredApiCard, ContextMenu, DropdownMenu, DropdownStyle, prelude::*};
use workspace::AppState;
+const OLLAMA_API_URL_PLACEHOLDER: &str = "http://localhost:11434";
+const OLLAMA_MODEL_PLACEHOLDER: &str = "qwen2.5-coder:3b-base";
+
use crate::{
SettingField, SettingItem, SettingsFieldMetadata, SettingsPageItem, SettingsWindow, USER,
components::{SettingsInputField, SettingsSectionHeader},
@@ -61,6 +64,7 @@ pub(crate) fn render_edit_prediction_setup_page(
)
.into_any_element()
}),
+ Some(render_ollama_provider(settings_window, window, cx).into_any_element()),
Some(
render_api_key_provider(
IconName::AiMistral,
@@ -328,6 +332,135 @@ fn sweep_settings() -> Box<[SettingsPageItem]> {
})])
}
+fn render_ollama_provider(
+ settings_window: &SettingsWindow,
+ window: &mut Window,
+ cx: &mut Context<SettingsWindow>,
+) -> impl IntoElement {
+ let ollama_settings = ollama_settings();
+ let additional_fields = settings_window
+ .render_sub_page_items_section(ollama_settings.iter().enumerate(), true, window, cx)
+ .into_any_element();
+
+ v_flex()
+ .id("ollama")
+ .min_w_0()
+ .pt_8()
+ .gap_1p5()
+ .child(
+ SettingsSectionHeader::new("Ollama")
+ .icon(IconName::ZedPredict)
+ .no_padding(true),
+ )
+ .child(
+ Label::new("Configure the local Ollama server and model used for edit predictions.")
+ .size(LabelSize::Small)
+ .color(Color::Muted),
+ )
+ .child(div().px_neg_8().child(additional_fields))
+}
+
+fn ollama_settings() -> Box<[SettingsPageItem]> {
+ Box::new([
+ SettingsPageItem::SettingItem(SettingItem {
+ title: "API URL",
+ description: "The base URL of your Ollama server.",
+ field: Box::new(SettingField {
+ pick: |settings| {
+ settings
+ .project
+ .all_languages
+ .edit_predictions
+ .as_ref()?
+ .ollama
+ .as_ref()?
+ .api_url
+ .as_ref()
+ },
+ write: |settings, value| {
+ settings
+ .project
+ .all_languages
+ .edit_predictions
+ .get_or_insert_default()
+ .ollama
+ .get_or_insert_default()
+ .api_url = value;
+ },
+ json_path: Some("edit_predictions.ollama.api_url"),
+ }),
+ metadata: Some(Box::new(SettingsFieldMetadata {
+ placeholder: Some(OLLAMA_API_URL_PLACEHOLDER),
+ ..Default::default()
+ })),
+ files: USER,
+ }),
+ SettingsPageItem::SettingItem(SettingItem {
+ title: "Model",
+ description: "The Ollama model to use for edit predictions.",
+ field: Box::new(SettingField {
+ pick: |settings| {
+ settings
+ .project
+ .all_languages
+ .edit_predictions
+ .as_ref()?
+ .ollama
+ .as_ref()?
+ .model
+ .as_ref()
+ },
+ write: |settings, value| {
+ settings
+ .project
+ .all_languages
+ .edit_predictions
+ .get_or_insert_default()
+ .ollama
+ .get_or_insert_default()
+ .model = value;
+ },
+ json_path: Some("edit_predictions.ollama.model"),
+ }),
+ metadata: Some(Box::new(SettingsFieldMetadata {
+ placeholder: Some(OLLAMA_MODEL_PLACEHOLDER),
+ ..Default::default()
+ })),
+ files: USER,
+ }),
+ SettingsPageItem::SettingItem(SettingItem {
+ title: "Max Output Tokens",
+ description: "The maximum number of tokens to generate.",
+ field: Box::new(SettingField {
+ pick: |settings| {
+ settings
+ .project
+ .all_languages
+ .edit_predictions
+ .as_ref()?
+ .ollama
+ .as_ref()?
+ .max_output_tokens
+ .as_ref()
+ },
+ write: |settings, value| {
+ settings
+ .project
+ .all_languages
+ .edit_predictions
+ .get_or_insert_default()
+ .ollama
+ .get_or_insert_default()
+ .max_output_tokens = value;
+ },
+ json_path: Some("edit_predictions.ollama.max_output_tokens"),
+ }),
+ metadata: None,
+ files: USER,
+ }),
+ ])
+}
+
fn codestral_settings() -> Box<[SettingsPageItem]> {
Box::new([
SettingsPageItem::SettingItem(SettingItem {
@@ -45,7 +45,8 @@ use zed_actions::{OpenProjectSettings, OpenSettings, OpenSettingsAt};
use crate::components::{
EnumVariantDropdown, NumberField, NumberFieldMode, NumberFieldType, SettingsInputField,
- SettingsSectionHeader, font_picker, icon_theme_picker, theme_picker,
+ SettingsSectionHeader, font_picker, icon_theme_picker, render_ollama_model_picker,
+ theme_picker,
};
const NAVBAR_CONTAINER_TAB_INDEX: isize = 0;
@@ -535,6 +536,7 @@ fn init_renderers(cx: &mut App) {
.add_basic_renderer::<settings::RelativeLineNumbers>(render_dropdown)
.add_basic_renderer::<settings::WindowDecorations>(render_dropdown)
.add_basic_renderer::<settings::FontSize>(render_editable_number_field)
+ .add_basic_renderer::<settings::OllamaModelName>(render_ollama_model_picker)
// please semicolon stay on next line
;
}
@@ -193,6 +193,7 @@ fn assign_edit_prediction_provider(
}
value @ (EditPredictionProvider::Experimental(_)
| EditPredictionProvider::Zed
+ | EditPredictionProvider::Ollama
| EditPredictionProvider::Sweep
| EditPredictionProvider::Mercury) => {
let ep_store = edit_prediction::EditPredictionStore::global(client, &user_store, cx);
@@ -209,6 +210,12 @@ fn assign_edit_prediction_provider(
EditPredictionProvider::Mercury if cx.has_flag::<MercuryFeatureFlag>() => {
edit_prediction::EditPredictionModel::Mercury
}
+ EditPredictionProvider::Ollama => {
+ if !edit_prediction::ollama::is_available(cx) {
+ return false;
+ }
+ edit_prediction::EditPredictionModel::Ollama
+ }
EditPredictionProvider::Experimental(name)
if name == EXPERIMENTAL_ZETA2_EDIT_PREDICTION_PROVIDER_NAME
&& cx.has_flag::<Zeta2FeatureFlag>() =>
@@ -483,6 +483,41 @@ pub mod v0131_git_merge_markers_prefix {
}
}
+/// The zeta1 prompt format
+pub mod zeta1 {
+ pub const CURSOR_MARKER: &str = "<|user_cursor_is_here|>";
+ pub const START_OF_FILE_MARKER: &str = "<|start_of_file|>";
+ pub const EDITABLE_REGION_START_MARKER: &str = "<|editable_region_start|>";
+ pub const EDITABLE_REGION_END_MARKER: &str = "<|editable_region_end|>";
+
+ const INSTRUCTION_HEADER: &str = concat!(
+ "### Instruction:\n",
+ "You are a code completion assistant and your task is to analyze user edits and then rewrite an ",
+ "excerpt that the user provides, suggesting the appropriate edits within the excerpt, taking ",
+ "into account the cursor location.\n\n",
+ "### User Edits:\n\n"
+ );
+ const EXCERPT_HEADER: &str = "\n\n### User Excerpt:\n\n";
+ const RESPONSE_HEADER: &str = "\n\n### Response:\n";
+
+ /// Formats a complete zeta1 prompt from the input events and excerpt.
+ pub fn format_zeta1_prompt(input_events: &str, input_excerpt: &str) -> String {
+ let mut prompt = String::with_capacity(
+ INSTRUCTION_HEADER.len()
+ + input_events.len()
+ + EXCERPT_HEADER.len()
+ + input_excerpt.len()
+ + RESPONSE_HEADER.len(),
+ );
+ prompt.push_str(INSTRUCTION_HEADER);
+ prompt.push_str(input_events);
+ prompt.push_str(EXCERPT_HEADER);
+ prompt.push_str(input_excerpt);
+ prompt.push_str(RESPONSE_HEADER);
+ prompt
+ }
+}
+
#[cfg(test)]
mod tests {
use super::*;