244 lines · 7.8 KB · Rust
RawDownload
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}