tenant.openagents/omega

1use crate::{
2    EditPredictionId, EditPredictionInputs, EditPredictionModelInput, cursor_excerpt,
3    open_ai_compatible::{self, load_open_ai_compatible_api_key_if_needed},
4    prediction::EditPredictionResult,
5};
6use anyhow::{Context as _, Result, anyhow};
7use gpui::{App, AppContext as _, Entity, Task};
8use language::{
9    Anchor, Buffer, BufferSnapshot, EditPredictionPromptFormat, ToOffset, ToPoint as _,
10    language_settings::all_language_settings,
11};
12use std::{path::Path, sync::Arc, time::Instant};
13use zeta_prompt::{Zeta2PromptInput, compute_editable_and_context_ranges};
14
15const FIM_CONTEXT_TOKENS: usize = 512;
16
17struct FimRequestOutput {
18    request_id: String,
19    edits: Vec<(std::ops::Range<Anchor>, Arc<str>)>,
20    editable_range: std::ops::Range<Anchor>,
21    snapshot: BufferSnapshot,
22    inputs: Zeta2PromptInput,
23    buffer: Entity<Buffer>,
24}
25
26pub fn request_prediction(
27    EditPredictionModelInput {
28        buffer,
29        snapshot,
30        position,
31        events,
32        trigger,
33        ..
34    }: EditPredictionModelInput,
35    prompt_format: EditPredictionPromptFormat,
36    cx: &mut App,
37) -> Task<Result<Option<EditPredictionResult>>> {
38    let settings = &all_language_settings(None, cx).edit_predictions;
39    let provider = settings.provider;
40
41    let full_path: Arc<Path> = snapshot
42        .file()
43        .map(|file| file.full_path(cx))
44        .unwrap_or_else(|| "untitled".into())
45        .into();
46
47    let http_client = cx.http_client();
48    let cursor_point = position.to_point(&snapshot);
49    let request_start = cx.background_executor().now();
50
51    let Some(settings) = (match provider {
52        settings::EditPredictionProvider::Ollama => settings.ollama.clone(),
53        settings::EditPredictionProvider::OpenAiCompatibleApi => {
54            settings.open_ai_compatible_api.clone()
55        }
56        _ => None,
57    }) else {
58        return Task::ready(Err(anyhow!("Unsupported edit prediction provider for FIM")));
59    };
60
61    let api_key = load_open_ai_compatible_api_key_if_needed(provider, cx);
62
63    let result = cx.background_spawn(async move {
64        let cursor_offset = cursor_point.to_offset(&snapshot);
65        let (excerpt_point_range, excerpt_offset_range, cursor_offset_in_excerpt) =
66            cursor_excerpt::compute_cursor_excerpt(&snapshot, cursor_offset);
67        let cursor_excerpt: Arc<str> = snapshot
68            .text_for_range(excerpt_point_range.clone())
69            .collect::<String>()
70            .into();
71        let syntax_ranges =
72            cursor_excerpt::compute_syntax_ranges(&snapshot, cursor_offset, &excerpt_offset_range);
73        let (editable_range, _) = compute_editable_and_context_ranges(
74            &cursor_excerpt,
75            cursor_offset_in_excerpt,
76            &syntax_ranges,
77            FIM_CONTEXT_TOKENS,
78            0,
79        );
80
81        let inputs = Zeta2PromptInput {
82            events,
83            related_files: Some(Vec::new()),
84            active_buffer_diagnostics: Vec::new(),
85            cursor_offset_in_excerpt: cursor_offset - excerpt_offset_range.start,
86            cursor_path: full_path.clone(),
87            excerpt_start_row: Some(excerpt_point_range.start.row),
88            cursor_excerpt,
89            excerpt_ranges: Default::default(),
90            syntax_ranges: None,
91            in_open_source_repo: false,
92            can_collect_data: false,
93            repo_url: None,
94        };
95
96        let editable_text = &inputs.cursor_excerpt[editable_range.clone()];
97        let cursor_in_editable = cursor_offset_in_excerpt.saturating_sub(editable_range.start);
98        let prefix = editable_text[..cursor_in_editable].to_string();
99        let suffix = editable_text[cursor_in_editable..].to_string();
100        let prompt = format_fim_prompt(prompt_format, &prefix, &suffix);
101        let stop_tokens = get_fim_stop_tokens();
102
103        let max_tokens = settings.max_output_tokens;
104
105        let (response_text, request_id) = open_ai_compatible::send_custom_server_request(
106            provider,
107            &settings,
108            prompt,
109            max_tokens,
110            stop_tokens,
111            api_key,
112            &http_client,
113        )
114        .await?;
115
116        let response_received_at = Instant::now();
117
118        log::debug!(
119            "fim: completion received ({:.2}s)",
120            (response_received_at - request_start).as_secs_f64()
121        );
122
123        let completion: Arc<str> = clean_fim_completion(&response_text).into();
124        let edits = if completion.is_empty() {
125            vec![]
126        } else {
127            let cursor_offset = cursor_point.to_offset(&snapshot);
128            let anchor = snapshot.anchor_after(cursor_offset);
129            vec![(anchor..anchor, completion)]
130        };
131
132        let editable_range = snapshot.anchor_range_inside(
133            (excerpt_offset_range.start + editable_range.start)
134                ..(excerpt_offset_range.start + editable_range.end),
135        );
136
137        anyhow::Ok(FimRequestOutput {
138            request_id,
139            edits,
140            editable_range,
141            snapshot,
142            inputs,
143            buffer,
144        })
145    });
146
147    cx.spawn(async move |cx: &mut gpui::AsyncApp| {
148        let output = result.await.context("fim edit prediction failed")?;
149        anyhow::Ok(Some(
150            EditPredictionResult::new(
151                EditPredictionId(output.request_id.into()),
152                &output.buffer,
153                &output.snapshot,
154                output.edits.into(),
155                None,
156                Some(output.editable_range),
157                EditPredictionInputs::V2(output.inputs),
158                None,
159                trigger,
160                cx.background_executor().now() - request_start,
161                cx,
162            )
163            .await,
164        ))
165    })
166}
167
168fn format_fim_prompt(
169    prompt_format: EditPredictionPromptFormat,
170    prefix: &str,
171    suffix: &str,
172) -> String {
173    match prompt_format {
174        EditPredictionPromptFormat::CodeLlama => {
175            format!("<PRE> {prefix} <SUF>{suffix} <MID>")
176        }
177        EditPredictionPromptFormat::StarCoder => {
178            format!("<fim_prefix>{prefix}<fim_suffix>{suffix}<fim_middle>")
179        }
180        EditPredictionPromptFormat::DeepseekCoder => {
181            format!("<｜fim▁begin｜>{prefix}<｜fim▁hole｜>{suffix}<｜fim▁end｜>")
182        }
183        EditPredictionPromptFormat::Qwen | EditPredictionPromptFormat::CodeGemma => {
184            format!("<|fim_prefix|>{prefix}<|fim_suffix|>{suffix}<|fim_middle|>")
185        }
186        EditPredictionPromptFormat::Codestral => {
187            format!("[SUFFIX]{suffix}[PREFIX]{prefix}")
188        }
189        EditPredictionPromptFormat::Glm => {
190            format!("<|code_prefix|>{prefix}<|code_suffix|>{suffix}<|code_middle|>")
191        }
192        _ => {
193            format!("<fim_prefix>{prefix}<fim_suffix>{suffix}<fim_middle>")
194        }
195    }
196}
197
198fn get_fim_stop_tokens() -> Vec<String> {
199    vec![
200        "<|endoftext|>".to_string(),
201        "<|file_separator|>".to_string(),
202        "<|fim_pad|>".to_string(),
203        "<|fim_prefix|>".to_string(),
204        "<|fim_middle|>".to_string(),
205        "<|fim_suffix|>".to_string(),
206        "<fim_prefix>".to_string(),
207        "<fim_middle>".to_string(),
208        "<fim_suffix>".to_string(),
209        "<PRE>".to_string(),
210        "<SUF>".to_string(),
211        "<MID>".to_string(),
212        "[PREFIX]".to_string(),
213        "[SUFFIX]".to_string(),
214    ]
215}
216
217fn clean_fim_completion(response: &str) -> String {
218    let mut result = response.to_string();
219
220    let end_tokens = [
221        "<|endoftext|>",
222        "<|file_separator|>",
223        "<|fim_pad|>",
224        "<|fim_prefix|>",
225        "<|fim_middle|>",
226        "<|fim_suffix|>",
227        "<fim_prefix>",
228        "<fim_middle>",
229        "<fim_suffix>",
230        "<PRE>",
231        "<SUF>",
232        "<MID>",
233        "[PREFIX]",
234        "[SUFFIX]",
235    ];
236
237    for token in &end_tokens {
238        if let Some(pos) = result.find(token) {
239            result.truncate(pos);
240        }
241    }
242
243    result
244}
245