use crate::{ EditPredictionId, EditPredictionInputs, EditPredictionModelInput, cursor_excerpt, open_ai_compatible::{self, load_open_ai_compatible_api_key_if_needed}, prediction::EditPredictionResult,};use anyhow::{Context as _, Result, anyhow};use gpui::{App, AppContext as _, Entity, Task};use language::{ Anchor, Buffer, BufferSnapshot, EditPredictionPromptFormat, ToOffset, ToPoint as _, language_settings::all_language_settings,};use std::{path::Path, sync::Arc, time::Instant};use zeta_prompt::{Zeta2PromptInput, compute_editable_and_context_ranges}; const FIM_CONTEXT_TOKENS: usize = 512; struct FimRequestOutput { request_id: String, edits: Vec<(std::ops::Range<Anchor>, Arc<str>)>, editable_range: std::ops::Range<Anchor>, snapshot: BufferSnapshot, inputs: Zeta2PromptInput, buffer: Entity<Buffer>,} pub fn request_prediction( EditPredictionModelInput { buffer, snapshot, position, events, trigger, .. }: EditPredictionModelInput, prompt_format: EditPredictionPromptFormat, cx: &mut App,) -> Task<Result<Option<EditPredictionResult>>> { let settings = &all_language_settings(None, cx).edit_predictions; let provider = settings.provider; 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 request_start = cx.background_executor().now(); let Some(settings) = (match provider { settings::EditPredictionProvider::Ollama => settings.ollama.clone(), settings::EditPredictionProvider::OpenAiCompatibleApi => { settings.open_ai_compatible_api.clone() } _ => None, }) else { return Task::ready(Err(anyhow!("Unsupported edit prediction provider for FIM"))); }; let api_key = load_open_ai_compatible_api_key_if_needed(provider, cx); let result = cx.background_spawn(async move { let cursor_offset = cursor_point.to_offset(&snapshot); let (excerpt_point_range, excerpt_offset_range, cursor_offset_in_excerpt) = cursor_excerpt::compute_cursor_excerpt(&snapshot, cursor_offset); let cursor_excerpt: Arc<str> = snapshot .text_for_range(excerpt_point_range.clone()) .collect::<String>() .into(); let syntax_ranges = cursor_excerpt::compute_syntax_ranges(&snapshot, cursor_offset, &excerpt_offset_range); let (editable_range, _) = compute_editable_and_context_ranges( &cursor_excerpt, cursor_offset_in_excerpt, &syntax_ranges, FIM_CONTEXT_TOKENS, 0, ); let inputs = Zeta2PromptInput { events, related_files: Some(Vec::new()), active_buffer_diagnostics: Vec::new(), cursor_offset_in_excerpt: cursor_offset - excerpt_offset_range.start, cursor_path: full_path.clone(), excerpt_start_row: Some(excerpt_point_range.start.row), cursor_excerpt, excerpt_ranges: Default::default(), syntax_ranges: None, in_open_source_repo: false, can_collect_data: false, repo_url: None, }; let editable_text = &inputs.cursor_excerpt[editable_range.clone()]; let cursor_in_editable = cursor_offset_in_excerpt.saturating_sub(editable_range.start); let prefix = editable_text[..cursor_in_editable].to_string(); let suffix = editable_text[cursor_in_editable..].to_string(); let prompt = format_fim_prompt(prompt_format, &prefix, &suffix); let stop_tokens = get_fim_stop_tokens(); let max_tokens = settings.max_output_tokens; let (response_text, request_id) = open_ai_compatible::send_custom_server_request( provider, &settings, prompt, max_tokens, stop_tokens, api_key, &http_client, ) .await?; let response_received_at = Instant::now(); log::debug!( "fim: completion received ({:.2}s)", (response_received_at - request_start).as_secs_f64() ); let completion: Arc<str> = clean_fim_completion(&response_text).into(); let edits = 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)] }; let editable_range = snapshot.anchor_range_inside( (excerpt_offset_range.start + editable_range.start) ..(excerpt_offset_range.start + editable_range.end), ); anyhow::Ok(FimRequestOutput { request_id, edits, editable_range, snapshot, inputs, buffer, }) }); cx.spawn(async move |cx: &mut gpui::AsyncApp| { let output = result.await.context("fim edit prediction failed")?; anyhow::Ok(Some( EditPredictionResult::new( EditPredictionId(output.request_id.into()), &output.buffer, &output.snapshot, output.edits.into(), None, Some(output.editable_range), EditPredictionInputs::V2(output.inputs), None, trigger, cx.background_executor().now() - request_start, cx, ) .await, )) })} fn format_fim_prompt( prompt_format: EditPredictionPromptFormat, prefix: &str, suffix: &str,) -> String { match prompt_format { EditPredictionPromptFormat::CodeLlama => { format!("<PRE> {prefix} <SUF>{suffix} <MID>") } EditPredictionPromptFormat::StarCoder => { format!("<fim_prefix>{prefix}<fim_suffix>{suffix}<fim_middle>") } EditPredictionPromptFormat::DeepseekCoder => { format!("<|fim▁begin|>{prefix}<|fim▁hole|>{suffix}<|fim▁end|>") } EditPredictionPromptFormat::Qwen | EditPredictionPromptFormat::CodeGemma => { format!("<|fim_prefix|>{prefix}<|fim_suffix|>{suffix}<|fim_middle|>") } EditPredictionPromptFormat::Codestral => { format!("[SUFFIX]{suffix}[PREFIX]{prefix}") } EditPredictionPromptFormat::Glm => { 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_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}
use crate::{
EditPredictionId, EditPredictionInputs, EditPredictionModelInput, cursor_excerpt,
open_ai_compatible::{self, load_open_ai_compatible_api_key_if_needed},
prediction::EditPredictionResult,
};
use anyhow::{Context as _, Result, anyhow};
use gpui::{App, AppContext as _, Entity, Task};
use language::{
Anchor, Buffer, BufferSnapshot, EditPredictionPromptFormat, ToOffset, ToPoint as _,
language_settings::all_language_settings,
};
use std::{path::Path, sync::Arc, time::Instant};
use zeta_prompt::{Zeta2PromptInput, compute_editable_and_context_ranges};
const FIM_CONTEXT_TOKENS: usize = 512;
struct FimRequestOutput {
request_id: String,
edits: Vec<(std::ops::Range<Anchor>, Arc<str>)>,
editable_range: std::ops::Range<Anchor>,
snapshot: BufferSnapshot,
inputs: Zeta2PromptInput,
buffer: Entity<Buffer>,
}
pub fn request_prediction(
EditPredictionModelInput {
buffer,
snapshot,
position,
events,
trigger,
..
}: EditPredictionModelInput,
prompt_format: EditPredictionPromptFormat,
cx: &mut App,
) -> Task<Result<Option<EditPredictionResult>>> {
let settings = &all_language_settings(None, cx).edit_predictions;
let provider = settings.provider;
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 request_start = cx.background_executor().now();
let Some(settings) = (match provider {
settings::EditPredictionProvider::Ollama => settings.ollama.clone(),
settings::EditPredictionProvider::OpenAiCompatibleApi => {
settings.open_ai_compatible_api.clone()
}
_ => None,
}) else {
return Task::ready(Err(anyhow!("Unsupported edit prediction provider for FIM")));
};
let api_key = load_open_ai_compatible_api_key_if_needed(provider, cx);
let result = cx.background_spawn(async move {
let cursor_offset = cursor_point.to_offset(&snapshot);
let (excerpt_point_range, excerpt_offset_range, cursor_offset_in_excerpt) =
cursor_excerpt::compute_cursor_excerpt(&snapshot, cursor_offset);
let cursor_excerpt: Arc<str> = snapshot
.text_for_range(excerpt_point_range.clone())
.collect::<String>()
.into();
let syntax_ranges =
cursor_excerpt::compute_syntax_ranges(&snapshot, cursor_offset, &excerpt_offset_range);
let (editable_range, _) = compute_editable_and_context_ranges(
&cursor_excerpt,
cursor_offset_in_excerpt,
&syntax_ranges,
FIM_CONTEXT_TOKENS,
0,
);
let inputs = Zeta2PromptInput {
events,
related_files: Some(Vec::new()),
active_buffer_diagnostics: Vec::new(),
cursor_offset_in_excerpt: cursor_offset - excerpt_offset_range.start,
cursor_path: full_path.clone(),
excerpt_start_row: Some(excerpt_point_range.start.row),
cursor_excerpt,
excerpt_ranges: Default::default(),
syntax_ranges: None,
in_open_source_repo: false,
can_collect_data: false,
repo_url: None,
};
let editable_text = &inputs.cursor_excerpt[editable_range.clone()];
let cursor_in_editable = cursor_offset_in_excerpt.saturating_sub(editable_range.start);
let prefix = editable_text[..cursor_in_editable].to_string();
let suffix = editable_text[cursor_in_editable..].to_string();
let prompt = format_fim_prompt(prompt_format, &prefix, &suffix);
let stop_tokens = get_fim_stop_tokens();
let max_tokens = settings.max_output_tokens;
let (response_text, request_id) = open_ai_compatible::send_custom_server_request(
provider,
&settings,
prompt,
max_tokens,
stop_tokens,
api_key,
&http_client,
)
.await?;
let response_received_at = Instant::now();
log::debug!(
"fim: completion received ({:.2}s)",
(response_received_at - request_start).as_secs_f64()
);
let completion: Arc<str> = clean_fim_completion(&response_text).into();
let edits = 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)]
};
let editable_range = snapshot.anchor_range_inside(
(excerpt_offset_range.start + editable_range.start)
..(excerpt_offset_range.start + editable_range.end),
);
anyhow::Ok(FimRequestOutput {
request_id,
edits,
editable_range,
snapshot,
inputs,
buffer,
})
});
cx.spawn(async move |cx: &mut gpui::AsyncApp| {
let output = result.await.context("fim edit prediction failed")?;
anyhow::Ok(Some(
EditPredictionResult::new(
EditPredictionId(output.request_id.into()),
&output.buffer,
&output.snapshot,
output.edits.into(),
None,
Some(output.editable_range),
EditPredictionInputs::V2(output.inputs),
None,
trigger,
cx.background_executor().now() - request_start,
cx,
)
.await,
))
})
}
fn format_fim_prompt(
prompt_format: EditPredictionPromptFormat,
prefix: &str,
suffix: &str,
) -> String {
match prompt_format {
EditPredictionPromptFormat::CodeLlama => {
format!("<PRE> {prefix} <SUF>{suffix} <MID>")
}
EditPredictionPromptFormat::StarCoder => {
format!("<fim_prefix>{prefix}<fim_suffix>{suffix}<fim_middle>")
}
EditPredictionPromptFormat::DeepseekCoder => {
format!("<|fim▁begin|>{prefix}<|fim▁hole|>{suffix}<|fim▁end|>")
}
EditPredictionPromptFormat::Qwen | EditPredictionPromptFormat::CodeGemma => {
format!("<|fim_prefix|>{prefix}<|fim_suffix|>{suffix}<|fim_middle|>")
}
EditPredictionPromptFormat::Codestral => {
format!("[SUFFIX]{suffix}[PREFIX]{prefix}")
}
EditPredictionPromptFormat::Glm => {
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_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
}