Add initial support for edit predictions via Ollama (#48233)
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>
This commit is contained in:
23 files changed
+1025
-48
No files matched your search
Generated
+2
@@ -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 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 newline_ix = content.find('\n').context("could not find newline")?;
|
||||
let content = &content[newline_ix + 1..];
|
||||
|
||||
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::*;
|
||||
|
||||
Reference in new issue
Block a user