tenant.openagents/omega

1use anyhow::Result;
2use edit_prediction::cursor_excerpt;
3use edit_prediction_types::{
4    EditPrediction, EditPredictionDelegate, EditPredictionDiscardReason, EditPredictionIconSet,
5    EditPredictionRequestTrigger,
6};
7use futures::AsyncReadExt;
8use gpui::{App, AppContext as _, Context, Entity, Global, SharedString, Task};
9use http_client::HttpClient;
10use icons::IconName;
11use language::{
12    Anchor, Buffer, BufferSnapshot, EditPreview, language_settings::all_language_settings,
13};
14use language_model::{ApiKeyState, AuthenticateError, EnvVar, env_var};
15use serde::{Deserialize, Serialize};
16
17use std::{
18    ops::Range,
19    sync::Arc,
20    time::{Duration, Instant},
21};
22use text::ToOffset;
23
24pub const CODESTRAL_API_URL: &str = "https://codestral.mistral.ai";
25pub const DEBOUNCE_TIMEOUT: Duration = Duration::from_millis(150);
26
27static CODESTRAL_API_KEY_ENV_VAR: std::sync::LazyLock<EnvVar> = env_var!("CODESTRAL_API_KEY");
28
29struct GlobalCodestralApiKey(Entity<ApiKeyState>);
30
31impl Global for GlobalCodestralApiKey {}
32
33pub fn codestral_api_key_state(cx: &mut App) -> Entity<ApiKeyState> {
34    if let Some(global) = cx.try_global::<GlobalCodestralApiKey>() {
35        return global.0.clone();
36    }
37    let entity =
38        cx.new(|cx| ApiKeyState::new(codestral_api_url(cx), CODESTRAL_API_KEY_ENV_VAR.clone()));
39    cx.set_global(GlobalCodestralApiKey(entity.clone()));
40    entity
41}
42
43pub fn codestral_api_key(cx: &App) -> Option<Arc<str>> {
44    let url = codestral_api_url(cx);
45    cx.try_global::<GlobalCodestralApiKey>()?
46        .0
47        .read(cx)
48        .key(&url)
49}
50
51pub fn load_codestral_api_key(cx: &mut App) -> Task<Result<(), AuthenticateError>> {
52    let credentials_provider = zed_credentials_provider::global(cx);
53    let api_url = codestral_api_url(cx);
54    codestral_api_key_state(cx).update(cx, |key_state, cx| {
55        key_state.load_if_needed(api_url, |s| s, credentials_provider, cx)
56    })
57}
58
59pub fn codestral_api_url(cx: &App) -> SharedString {
60    all_language_settings(None, cx)
61        .edit_predictions
62        .codestral
63        .api_url
64        .clone()
65        .unwrap_or_else(|| CODESTRAL_API_URL.to_string())
66        .into()
67}
68
69/// Represents a completion that has been received and processed from Codestral.
70/// This struct maintains the state needed to interpolate the completion as the user types.
71#[derive(Clone)]
72struct CurrentCompletion {
73    /// The buffer snapshot at the time the completion was generated.
74    /// Used to detect changes and interpolate edits.
75    snapshot: BufferSnapshot,
76    /// The edits that should be applied to transform the original text into the predicted text.
77    /// Each edit is a range in the buffer and the text to replace it with.
78    edits: Arc<[(Range<Anchor>, Arc<str>)]>,
79    /// Preview of how the buffer will look after applying the edits.
80    edit_preview: EditPreview,
81}
82
83impl CurrentCompletion {
84    /// Attempts to adjust the edits based on changes made to the buffer since the completion was generated.
85    /// Returns None if no predicted edits remain or the user's edits conflict with the predicted edits.
86    fn interpolate(&self, new_snapshot: &BufferSnapshot) -> Option<Vec<(Range<Anchor>, Arc<str>)>> {
87        edit_prediction_types::interpolate_edits(&self.snapshot, new_snapshot, &self.edits)
88            .filter(|edits| !edits.is_empty())
89    }
90}
91
92pub struct CodestralEditPredictionDelegate {
93    http_client: Arc<dyn HttpClient>,
94    pending_request: Option<Task<Result<()>>>,
95    current_completion: Option<CurrentCompletion>,
96}
97
98impl CodestralEditPredictionDelegate {
99    pub fn new(http_client: Arc<dyn HttpClient>) -> Self {
100        Self {
101            http_client,
102            pending_request: None,
103            current_completion: None,
104        }
105    }
106
107    pub fn ensure_api_key_loaded(cx: &mut App) {
108        load_codestral_api_key(cx).detach();
109    }
110
111    /// Uses Codestral's Fill-in-the-Middle API for code completion.
112    async fn fetch_completion(
113        http_client: Arc<dyn HttpClient>,
114        api_key: &str,
115        prompt: String,
116        suffix: String,
117        model: String,
118        max_tokens: Option<u32>,
119        api_url: String,
120    ) -> Result<String> {
121        let start_time = Instant::now();
122
123        log::debug!(
124            "Codestral: Requesting completion (model: {}, max_tokens: {:?})",
125            model,
126            max_tokens
127        );
128
129        let request = CodestralRequest {
130            model,
131            prompt,
132            suffix: if suffix.is_empty() {
133                None
134            } else {
135                Some(suffix)
136            },
137            max_tokens: max_tokens.or(Some(350)),
138            temperature: Some(0.2),
139            top_p: Some(1.0),
140            stream: Some(false),
141            stop: None,
142            random_seed: None,
143            min_tokens: None,
144        };
145
146        let request_body = serde_json::to_string(&request)?;
147
148        log::debug!("Codestral: Sending FIM request");
149
150        let http_request = http_client::Request::builder()
151            .method(http_client::Method::POST)
152            .uri(format!("{}/v1/fim/completions", api_url))
153            .header("Content-Type", "application/json")
154            .header("Authorization", format!("Bearer {}", api_key))
155            .body(http_client::AsyncBody::from(request_body))?;
156
157        let mut response = http_client.send(http_request).await?;
158        let status = response.status();
159
160        log::debug!("Codestral: Response status: {}", status);
161
162        if !status.is_success() {
163            let mut body = String::new();
164            response.body_mut().read_to_string(&mut body).await?;
165            return Err(anyhow::anyhow!(
166                "Codestral API error: {} - {}",
167                status,
168                body
169            ));
170        }
171
172        let mut body = String::new();
173        response.body_mut().read_to_string(&mut body).await?;
174
175        let codestral_response: CodestralResponse = serde_json::from_str(&body)?;
176
177        let elapsed = start_time.elapsed();
178
179        if let Some(choice) = codestral_response.choices.first() {
180            let completion = &choice.message.content;
181
182            log::debug!(
183                "Codestral: Completion received ({} tokens, {:.2}s)",
184                codestral_response.usage.completion_tokens,
185                elapsed.as_secs_f64()
186            );
187
188            // Return just the completion text for insertion at cursor
189            Ok(completion.clone())
190        } else {
191            log::error!("Codestral: No completion returned in response");
192            Err(anyhow::anyhow!("No completion returned from Codestral"))
193        }
194    }
195}
196
197impl EditPredictionDelegate for CodestralEditPredictionDelegate {
198    fn name() -> &'static str {
199        "codestral"
200    }
201
202    fn display_name() -> &'static str {
203        "Codestral"
204    }
205
206    fn show_predictions_in_menu() -> bool {
207        true
208    }
209
210    fn icons(&self, _cx: &App) -> EditPredictionIconSet {
211        EditPredictionIconSet::new(IconName::AiMistral)
212    }
213
214    fn is_enabled(&self, _buffer: &Entity<Buffer>, _cursor_position: Anchor, cx: &App) -> bool {
215        codestral_api_key(cx).is_some()
216    }
217
218    fn is_refreshing(&self, _cx: &App) -> bool {
219        self.pending_request.is_some()
220    }
221
222    fn refresh(
223        &mut self,
224        buffer: Entity<Buffer>,
225        cursor_position: language::Anchor,
226        debounce: bool,
227        _trigger: EditPredictionRequestTrigger,
228        cx: &mut Context<Self>,
229    ) {
230        log::debug!("Codestral: Refresh called (debounce: {})", debounce);
231
232        let Some(api_key) = codestral_api_key(cx) else {
233            log::warn!("Codestral: No API key configured, skipping refresh");
234            return;
235        };
236
237        let snapshot = buffer.read(cx).snapshot();
238
239        // Check if current completion is still valid
240        if let Some(current_completion) = self.current_completion.as_ref() {
241            if current_completion.interpolate(&snapshot).is_some() {
242                return;
243            }
244        }
245
246        let http_client = self.http_client.clone();
247
248        // Get settings
249        let settings = all_language_settings(None, cx);
250        let model = settings
251            .edit_predictions
252            .codestral
253            .model
254            .clone()
255            .unwrap_or_else(|| "codestral-latest".to_string());
256        let max_tokens = settings.edit_predictions.codestral.max_tokens;
257        let api_url = codestral_api_url(cx).to_string();
258
259        self.pending_request = Some(cx.spawn(async move |this, cx| {
260            if debounce {
261                log::debug!("Codestral: Debouncing for {:?}", DEBOUNCE_TIMEOUT);
262                cx.background_executor().timer(DEBOUNCE_TIMEOUT).await;
263            }
264
265            let cursor_offset = cursor_position.to_offset(&snapshot);
266
267            const MAX_EDITABLE_TOKENS: usize = 350;
268            const MAX_CONTEXT_TOKENS: usize = 150;
269
270            let (excerpt_point_range, excerpt_offset_range, cursor_offset_in_excerpt) =
271                cursor_excerpt::compute_cursor_excerpt(&snapshot, cursor_offset);
272            let syntax_ranges = cursor_excerpt::compute_syntax_ranges(
273                &snapshot,
274                cursor_offset,
275                &excerpt_offset_range,
276            );
277            let excerpt_text: String = snapshot.text_for_range(excerpt_point_range).collect();
278            let (_, context_range) = zeta_prompt::compute_editable_and_context_ranges(
279                &excerpt_text,
280                cursor_offset_in_excerpt,
281                &syntax_ranges,
282                MAX_EDITABLE_TOKENS,
283                MAX_CONTEXT_TOKENS,
284            );
285            let context_text = &excerpt_text[context_range.clone()];
286            let cursor_within_excerpt = cursor_offset_in_excerpt
287                .saturating_sub(context_range.start)
288                .min(context_text.len());
289            let prompt = context_text[..cursor_within_excerpt].to_string();
290            let suffix = context_text[cursor_within_excerpt..].to_string();
291
292            let completion_text = match Self::fetch_completion(
293                http_client,
294                &api_key,
295                prompt,
296                suffix,
297                model,
298                max_tokens,
299                api_url,
300            )
301            .await
302            {
303                Ok(completion) => completion,
304                Err(e) => {
305                    log::error!("Codestral: Failed to fetch completion: {}", e);
306                    this.update(cx, |this, cx| {
307                        this.pending_request = None;
308                        cx.notify();
309                    })?;
310                    return Err(e);
311                }
312            };
313
314            if completion_text.trim().is_empty() {
315                log::debug!("Codestral: Completion was empty after trimming; ignoring");
316                this.update(cx, |this, cx| {
317                    this.pending_request = None;
318                    cx.notify();
319                })?;
320                return Ok(());
321            }
322
323            let edits: Arc<[(Range<Anchor>, Arc<str>)]> =
324                vec![(cursor_position..cursor_position, completion_text.into())].into();
325            let edit_preview = buffer
326                .read_with(cx, |buffer, cx| buffer.preview_edits(edits.clone(), cx))
327                .await;
328
329            this.update(cx, |this, cx| {
330                this.current_completion = Some(CurrentCompletion {
331                    snapshot,
332                    edits,
333                    edit_preview,
334                });
335                this.pending_request = None;
336                cx.notify();
337            })?;
338
339            Ok(())
340        }));
341    }
342
343    fn accept(&mut self, _cx: &mut Context<Self>) {
344        log::debug!("Codestral: Completion accepted");
345        self.pending_request = None;
346        self.current_completion = None;
347    }
348
349    fn discard(&mut self, _reason: EditPredictionDiscardReason, _cx: &mut Context<Self>) {
350        log::debug!("Codestral: Completion discarded");
351        self.pending_request = None;
352        self.current_completion = None;
353    }
354
355    /// Returns the completion suggestion, adjusted or invalidated based on user edits
356    fn suggest(
357        &mut self,
358        buffer: &Entity<Buffer>,
359        _cursor_position: Anchor,
360        cx: &mut Context<Self>,
361    ) -> Option<EditPrediction> {
362        let current_completion = self.current_completion.as_ref()?;
363        let buffer = buffer.read(cx);
364        let edits = current_completion.interpolate(&buffer.snapshot())?;
365        if edits.is_empty() {
366            return None;
367        }
368        Some(EditPrediction::Local {
369            id: None,
370            edits,
371            cursor_position: None,
372            edit_preview: Some(current_completion.edit_preview.clone()),
373        })
374    }
375}
376
377#[derive(Debug, Serialize, Deserialize)]
378pub struct CodestralRequest {
379    pub model: String,
380    pub prompt: String,
381    #[serde(skip_serializing_if = "Option::is_none")]
382    pub suffix: Option<String>,
383    #[serde(skip_serializing_if = "Option::is_none")]
384    pub max_tokens: Option<u32>,
385    #[serde(skip_serializing_if = "Option::is_none")]
386    pub temperature: Option<f32>,
387    #[serde(skip_serializing_if = "Option::is_none")]
388    pub top_p: Option<f32>,
389    #[serde(skip_serializing_if = "Option::is_none")]
390    pub stream: Option<bool>,
391    #[serde(skip_serializing_if = "Option::is_none")]
392    pub stop: Option<Vec<String>>,
393    #[serde(skip_serializing_if = "Option::is_none")]
394    pub random_seed: Option<u32>,
395    #[serde(skip_serializing_if = "Option::is_none")]
396    pub min_tokens: Option<u32>,
397}
398
399#[derive(Debug, Deserialize)]
400pub struct CodestralResponse {
401    pub id: String,
402    pub object: String,
403    pub model: String,
404    pub usage: Usage,
405    pub created: u64,
406    pub choices: Vec<Choice>,
407}
408
409#[derive(Debug, Deserialize)]
410pub struct Usage {
411    pub prompt_tokens: u32,
412    pub completion_tokens: u32,
413    pub total_tokens: u32,
414}
415
416#[derive(Debug, Deserialize)]
417pub struct Choice {
418    pub index: u32,
419    pub message: Message,
420    pub finish_reason: String,
421}
422
423#[derive(Debug, Deserialize)]
424pub struct Message {
425    pub content: String,
426    pub role: String,
427}
428