tenant.openagents/omega

1use anyhow::Result;
2use collections::BTreeMap;
3use credentials_provider::CredentialsProvider;
4use futures::{FutureExt, StreamExt, future::BoxFuture};
5use gpui::{App, AppContext, AsyncApp, Context, Entity, SharedString, Task};
6use http_client::{CustomHeaders, HttpClient};
7use language_model::{
8    ApiKeyConfiguration, ApiKeyState, AuthenticateError, CompactionResult, EnvVar,
9    FastModeConfirmation, IconOrSvg, LanguageModel, LanguageModelCompletionError,
10    LanguageModelCompletionEvent, LanguageModelEffortLevel, LanguageModelId, LanguageModelName,
11    LanguageModelProvider, LanguageModelProviderId, LanguageModelProviderName,
12    LanguageModelProviderState, LanguageModelRequest, LanguageModelToolChoice, OPEN_AI_PROVIDER_ID,
13    OPEN_AI_PROVIDER_NAME, ProviderSettingsView, RateLimiter, env_var,
14};
15use open_ai::{
16    ResponseStreamEvent,
17    responses::{
18        CompactRequest, CompactedResponse, Request as ResponseRequest,
19        StreamEvent as ResponsesStreamEvent, compact_response, stream_response,
20    },
21    stream_completion,
22};
23use settings::{OpenAiAvailableModel as AvailableModel, Settings, SettingsStore};
24use std::sync::{Arc, LazyLock};
25use strum::IntoEnumIterator;
26use ui::IconName;
27
28use open_ai::completion::token_usage_from_response_usage;
29pub use open_ai::completion::{
30    ChatCompletionMaxTokensParameter, OpenAiEventMapper, OpenAiResponseEventMapper, into_open_ai,
31    into_open_ai_response,
32};
33
34const PROVIDER_ID: LanguageModelProviderId = OPEN_AI_PROVIDER_ID;
35const PROVIDER_NAME: LanguageModelProviderName = OPEN_AI_PROVIDER_NAME;
36
37const API_KEY_ENV_VAR_NAME: &str = "OPENAI_API_KEY";
38static API_KEY_ENV_VAR: LazyLock<EnvVar> = env_var!(API_KEY_ENV_VAR_NAME);
39
40#[derive(Default, Clone, Debug, PartialEq)]
41pub struct OpenAiSettings {
42    pub api_url: String,
43    pub available_models: Vec<AvailableModel>,
44    pub custom_headers: CustomHeaders,
45}
46
47pub struct OpenAiLanguageModelProvider {
48    http_client: Arc<dyn HttpClient>,
49    state: Entity<State>,
50}
51
52pub struct State {
53    api_key_state: ApiKeyState,
54    credentials_provider: Arc<dyn CredentialsProvider>,
55}
56
57impl State {
58    fn is_authenticated(&self) -> bool {
59        self.api_key_state.has_key()
60    }
61
62    fn set_api_key(&mut self, api_key: Option<String>, cx: &mut Context<Self>) -> Task<Result<()>> {
63        let credentials_provider = self.credentials_provider.clone();
64        let api_url = OpenAiLanguageModelProvider::api_url(cx);
65        self.api_key_state.store(
66            api_url,
67            api_key,
68            |this| &mut this.api_key_state,
69            credentials_provider,
70            cx,
71        )
72    }
73
74    fn authenticate(&mut self, cx: &mut Context<Self>) -> Task<Result<(), AuthenticateError>> {
75        let credentials_provider = self.credentials_provider.clone();
76        let api_url = OpenAiLanguageModelProvider::api_url(cx);
77        self.api_key_state.load_if_needed(
78            api_url,
79            |this| &mut this.api_key_state,
80            credentials_provider,
81            cx,
82        )
83    }
84}
85
86impl OpenAiLanguageModelProvider {
87    pub fn new(
88        http_client: Arc<dyn HttpClient>,
89        credentials_provider: Arc<dyn CredentialsProvider>,
90        cx: &mut App,
91    ) -> Self {
92        let state = cx.new(|cx| {
93            cx.observe_global::<SettingsStore>(|this: &mut State, cx| {
94                let credentials_provider = this.credentials_provider.clone();
95                let api_url = Self::api_url(cx);
96                this.api_key_state.handle_url_change(
97                    api_url,
98                    |this| &mut this.api_key_state,
99                    credentials_provider,
100                    cx,
101                );
102                cx.notify();
103            })
104            .detach();
105            State {
106                api_key_state: ApiKeyState::new(Self::api_url(cx), (*API_KEY_ENV_VAR).clone()),
107                credentials_provider,
108            }
109        });
110
111        Self { http_client, state }
112    }
113
114    fn create_language_model(&self, model: open_ai::Model) -> Arc<dyn LanguageModel> {
115        Arc::new(OpenAiLanguageModel {
116            id: LanguageModelId::from(model.id().to_string()),
117            model,
118            state: self.state.clone(),
119            http_client: self.http_client.clone(),
120            request_limiter: RateLimiter::new(4),
121        })
122    }
123
124    fn settings(cx: &App) -> &OpenAiSettings {
125        &crate::AllLanguageModelSettings::get_global(cx).openai
126    }
127
128    fn api_url(cx: &App) -> SharedString {
129        let api_url = &Self::settings(cx).api_url;
130        if api_url.is_empty() {
131            open_ai::OPEN_AI_API_URL.into()
132        } else {
133            SharedString::new(api_url.as_str())
134        }
135    }
136}
137
138impl LanguageModelProviderState for OpenAiLanguageModelProvider {
139    type ObservableEntity = State;
140
141    fn observable_entity(&self) -> Option<Entity<Self::ObservableEntity>> {
142        Some(self.state.clone())
143    }
144}
145
146impl LanguageModelProvider for OpenAiLanguageModelProvider {
147    fn id(&self) -> LanguageModelProviderId {
148        PROVIDER_ID
149    }
150
151    fn name(&self) -> LanguageModelProviderName {
152        PROVIDER_NAME
153    }
154
155    fn icon(&self) -> IconOrSvg {
156        IconOrSvg::Icon(IconName::AiOpenAi)
157    }
158
159    fn default_model(&self, _cx: &App) -> Option<Arc<dyn LanguageModel>> {
160        Some(self.create_language_model(open_ai::Model::default()))
161    }
162
163    fn default_fast_model(&self, _cx: &App) -> Option<Arc<dyn LanguageModel>> {
164        Some(self.create_language_model(open_ai::Model::default_fast()))
165    }
166
167    fn provided_models(&self, cx: &App) -> Vec<Arc<dyn LanguageModel>> {
168        let mut models = BTreeMap::default();
169
170        // Add base models from open_ai::Model::iter()
171        for model in open_ai::Model::iter() {
172            if !matches!(model, open_ai::Model::Custom { .. }) {
173                models.insert(model.id().to_string(), model);
174            }
175        }
176
177        // Override with available models from settings
178        for model in &OpenAiLanguageModelProvider::settings(cx).available_models {
179            models.insert(
180                model.name.clone(),
181                open_ai::Model::Custom {
182                    name: model.name.clone(),
183                    display_name: model.display_name.clone(),
184                    max_tokens: model.max_tokens,
185                    max_output_tokens: model.max_output_tokens,
186                    max_completion_tokens: model.max_completion_tokens,
187                    reasoning_effort: model.reasoning_effort,
188                    supports_chat_completions: model.capabilities.chat_completions,
189                    supports_images: model.capabilities.images,
190                },
191            );
192        }
193
194        models
195            .into_values()
196            .map(|model| self.create_language_model(model))
197            .collect()
198    }
199
200    fn is_authenticated(&self, cx: &App) -> bool {
201        self.state.read(cx).is_authenticated()
202    }
203
204    fn authenticate(&self, cx: &mut App) -> Task<Result<(), AuthenticateError>> {
205        self.state.update(cx, |state, cx| state.authenticate(cx))
206    }
207
208    fn settings_view(&self, cx: &mut App) -> Option<ProviderSettingsView> {
209        let state = self.state.read(cx);
210        Some(ProviderSettingsView::ApiKey(ApiKeyConfiguration::new(
211            state.api_key_state.has_key(),
212            state.api_key_state.is_from_env_var(),
213            state.api_key_state.env_var_name().clone(),
214            "https://platform.openai.com/api-keys".into(),
215        )))
216    }
217
218    fn set_api_key(&self, api_key: Option<String>, cx: &mut App) -> Task<Result<()>> {
219        self.state
220            .update(cx, |state, cx| state.set_api_key(api_key, cx))
221    }
222
223    fn fast_mode_confirmation(&self, _cx: &App) -> Option<FastModeConfirmation> {
224        Some(FastModeConfirmation {
225            title: "Enable Fast Mode for OpenAI?".into(),
226            message: "Fast mode sends requests using OpenAI's Priority processing tier, which \
227                targets significantly lower latency than the standard tier and is billed at a \
228                premium per-token rate."
229                .into(),
230        })
231    }
232}
233
234fn default_thinking_reasoning_effort(model: &open_ai::Model) -> Option<open_ai::ReasoningEffort> {
235    use open_ai::ReasoningEffort;
236
237    model
238        .reasoning_effort()
239        .filter(|effort| open_ai_reasoning_effort_is_supported(*effort))
240        .or_else(|| {
241            let supported_efforts = model.supported_reasoning_efforts();
242            if supported_efforts.contains(&ReasoningEffort::Medium) {
243                Some(ReasoningEffort::Medium)
244            } else {
245                supported_efforts
246                    .iter()
247                    .copied()
248                    .find(|effort| open_ai_reasoning_effort_is_supported(*effort))
249            }
250        })
251}
252
253fn open_ai_reasoning_effort_is_supported(effort: open_ai::ReasoningEffort) -> bool {
254    effort != open_ai::ReasoningEffort::None
255}
256
257fn normalize_open_ai_response_thinking_effort(
258    request: &mut LanguageModelRequest,
259    model: &open_ai::Model,
260) {
261    let selected_effort_is_supported = request
262        .thinking_effort
263        .as_deref()
264        .and_then(|effort| effort.parse::<open_ai::ReasoningEffort>().ok())
265        .is_some_and(|effort| {
266            open_ai_reasoning_effort_is_supported(effort)
267                && model.supported_reasoning_efforts().contains(&effort)
268        });
269
270    if !selected_effort_is_supported {
271        request.thinking_effort = None;
272    }
273}
274
275fn supports_selectable_thinking_effort(model: &open_ai::Model) -> bool {
276    model.uses_responses_api()
277        && model
278            .supported_reasoning_efforts()
279            .iter()
280            .any(|effort| open_ai_reasoning_effort_is_supported(*effort))
281}
282
283fn supported_thinking_effort_levels(model: &open_ai::Model) -> Vec<LanguageModelEffortLevel> {
284    if !supports_selectable_thinking_effort(model) {
285        return Vec::new();
286    }
287
288    let default_effort = default_thinking_reasoning_effort(model);
289    model
290        .supported_reasoning_efforts()
291        .iter()
292        .copied()
293        .filter_map(|effort| {
294            if !open_ai_reasoning_effort_is_supported(effort) {
295                return None;
296            }
297
298            Some(LanguageModelEffortLevel {
299                name: effort.label().into(),
300                value: effort.value().into(),
301                is_default: Some(effort) == default_effort,
302            })
303        })
304        .collect()
305}
306
307#[cfg(test)]
308mod tests {
309    use super::*;
310
311    #[test]
312    fn supported_thinking_effort_levels_hide_none() {
313        let effort_levels = supported_thinking_effort_levels(&open_ai::Model::FivePointTwo);
314        let values = effort_levels
315            .iter()
316            .map(|level| level.value.as_ref())
317            .collect::<Vec<_>>();
318
319        assert_eq!(values, ["low", "medium", "high", "xhigh"]);
320        assert_eq!(
321            effort_levels
322                .iter()
323                .find(|level| level.is_default)
324                .map(|level| level.value.as_ref()),
325            Some("medium")
326        );
327    }
328
329    #[test]
330    fn models_supporting_only_none_have_no_selectable_thinking_effort() {
331        let model = open_ai::Model::Custom {
332            name: "custom-model".to_string(),
333            display_name: None,
334            max_tokens: 128_000,
335            max_output_tokens: None,
336            max_completion_tokens: None,
337            reasoning_effort: Some(open_ai::ReasoningEffort::None),
338            supports_chat_completions: false,
339            supports_images: true,
340        };
341
342        assert!(!supports_selectable_thinking_effort(&model));
343        assert!(supported_thinking_effort_levels(&model).is_empty());
344        assert!(
345            model
346                .supported_reasoning_efforts()
347                .contains(&open_ai::ReasoningEffort::None)
348        );
349    }
350}
351
352pub struct OpenAiLanguageModel {
353    id: LanguageModelId,
354    model: open_ai::Model,
355    state: Entity<State>,
356    http_client: Arc<dyn HttpClient>,
357    request_limiter: RateLimiter,
358}
359
360impl OpenAiLanguageModel {
361    fn stream_completion(
362        &self,
363        request: open_ai::Request,
364        cx: &AsyncApp,
365    ) -> BoxFuture<'static, Result<futures::stream::BoxStream<'static, Result<ResponseStreamEvent>>>>
366    {
367        let http_client = self.http_client.clone();
368
369        let (api_key, api_url, extra_headers) = self.state.read_with(cx, |state, cx| {
370            let api_url = OpenAiLanguageModelProvider::api_url(cx);
371            let extra_headers = OpenAiLanguageModelProvider::settings(cx)
372                .custom_headers
373                .clone();
374            (state.api_key_state.key(&api_url), api_url, extra_headers)
375        });
376
377        let future = self.request_limiter.stream(async move {
378            let provider = PROVIDER_NAME;
379            let Some(api_key) = api_key else {
380                return Err(LanguageModelCompletionError::NoApiKey { provider });
381            };
382            let request = stream_completion(
383                http_client.as_ref(),
384                provider.0.as_str(),
385                &api_url,
386                &api_key,
387                request,
388                &extra_headers,
389            );
390            let response = request.await?;
391            Ok(response)
392        });
393
394        async move { Ok(future.await?.boxed()) }.boxed()
395    }
396
397    fn stream_response(
398        &self,
399        request: ResponseRequest,
400        cx: &AsyncApp,
401    ) -> BoxFuture<'static, Result<futures::stream::BoxStream<'static, Result<ResponsesStreamEvent>>>>
402    {
403        let http_client = self.http_client.clone();
404
405        let (api_key, api_url, extra_headers) = self.state.read_with(cx, |state, cx| {
406            let api_url = OpenAiLanguageModelProvider::api_url(cx);
407            let extra_headers = OpenAiLanguageModelProvider::settings(cx)
408                .custom_headers
409                .clone();
410            (state.api_key_state.key(&api_url), api_url, extra_headers)
411        });
412
413        let provider = PROVIDER_NAME;
414        let future = self.request_limiter.stream(async move {
415            let Some(api_key) = api_key else {
416                return Err(LanguageModelCompletionError::NoApiKey { provider });
417            };
418            let request = stream_response(
419                http_client.as_ref(),
420                provider.0.as_str(),
421                &api_url,
422                &api_key,
423                request,
424                &extra_headers,
425            );
426            let response = request.await?;
427            Ok(response)
428        });
429
430        async move { Ok(future.await?.boxed()) }.boxed()
431    }
432
433    fn compact_response(
434        &self,
435        request: CompactRequest,
436        cx: &AsyncApp,
437    ) -> BoxFuture<'static, Result<CompactedResponse, LanguageModelCompletionError>> {
438        let http_client = self.http_client.clone();
439
440        let (api_key, api_url, extra_headers) = self.state.read_with(cx, |state, cx| {
441            let api_url = OpenAiLanguageModelProvider::api_url(cx);
442            let extra_headers = OpenAiLanguageModelProvider::settings(cx)
443                .custom_headers
444                .clone();
445            (state.api_key_state.key(&api_url), api_url, extra_headers)
446        });
447
448        let provider = PROVIDER_NAME;
449        let future = self.request_limiter.run(async move {
450            let Some(api_key) = api_key else {
451                return Err(LanguageModelCompletionError::NoApiKey { provider });
452            };
453            Ok(compact_response(
454                http_client.as_ref(),
455                provider.0.as_str(),
456                &api_url,
457                &api_key,
458                request,
459                &extra_headers,
460            )
461            .await?)
462        });
463
464        future.boxed()
465    }
466}
467
468impl LanguageModel for OpenAiLanguageModel {
469    fn id(&self) -> LanguageModelId {
470        self.id.clone()
471    }
472
473    fn name(&self) -> LanguageModelName {
474        LanguageModelName::from(self.model.display_name().to_string())
475    }
476
477    fn provider_id(&self) -> LanguageModelProviderId {
478        PROVIDER_ID
479    }
480
481    fn provider_name(&self) -> LanguageModelProviderName {
482        PROVIDER_NAME
483    }
484
485    fn supports_tools(&self) -> bool {
486        true
487    }
488
489    fn supports_images(&self) -> bool {
490        use open_ai::Model;
491        match &self.model {
492            Model::FourOmniMini
493            | Model::Five
494            | Model::FiveMini
495            | Model::FiveNano
496            | Model::FivePointOne
497            | Model::FivePointTwo
498            | Model::FivePointThreeCodex
499            | Model::FivePointFour
500            | Model::FivePointFourMini
501            | Model::FivePointFourNano
502            | Model::FivePointFourPro
503            | Model::FivePointFive
504            | Model::FivePointFivePro
505            | Model::FivePointSixSol
506            | Model::FivePointSixTerra
507            | Model::FivePointSixLuna
508            | Model::O3 => true,
509            Model::Four => false,
510            Model::Custom {
511                supports_images, ..
512            } => *supports_images,
513        }
514    }
515
516    fn supports_tool_choice(&self, choice: LanguageModelToolChoice) -> bool {
517        match choice {
518            LanguageModelToolChoice::Auto => true,
519            LanguageModelToolChoice::Any => true,
520            LanguageModelToolChoice::None => true,
521        }
522    }
523
524    fn supports_streaming_tools(&self) -> bool {
525        true
526    }
527
528    fn supports_thinking(&self) -> bool {
529        supports_selectable_thinking_effort(&self.model)
530    }
531
532    fn supports_fast_mode(&self) -> bool {
533        self.model.supports_priority()
534    }
535
536    fn supports_server_side_compaction(&self) -> bool {
537        self.model.supports_compaction()
538    }
539
540    fn supports_explicit_compaction(&self) -> bool {
541        self.model.supports_compaction()
542    }
543
544    fn compact(
545        &self,
546        mut request: LanguageModelRequest,
547        cx: &AsyncApp,
548    ) -> BoxFuture<'static, Result<CompactionResult, LanguageModelCompletionError>> {
549        if !self.supports_explicit_compaction() {
550            return async {
551                Err(LanguageModelCompletionError::Other(anyhow::anyhow!(
552                    "this OpenAI model does not support explicit compaction"
553                )))
554            }
555            .boxed();
556        }
557
558        normalize_open_ai_response_thinking_effort(&mut request, &self.model);
559        let request = match into_open_ai_response(
560            request,
561            self.model.id(),
562            self.model.supports_parallel_tool_calls(),
563            self.model.supports_prompt_cache_key(),
564            self.max_output_tokens(),
565            default_thinking_reasoning_effort(&self.model),
566            self.model
567                .supported_reasoning_efforts()
568                .contains(&open_ai::ReasoningEffort::None),
569            &OPEN_AI_PROVIDER_ID,
570        ) {
571            Ok(request) => request,
572            Err(error) => return async move { Err(error.into()) }.boxed(),
573        };
574        let request = request.into_compact_request();
575        let response = self.compact_response(request, cx);
576        async move {
577            let response = response.await?;
578            let usage = token_usage_from_response_usage(&response.usage);
579            let context = response
580                .into_compacted_context(OPEN_AI_PROVIDER_ID)
581                .map_err(LanguageModelCompletionError::Other)?;
582            Ok(CompactionResult { context, usage })
583        }
584        .boxed()
585    }
586
587    fn supported_effort_levels(&self) -> Vec<LanguageModelEffortLevel> {
588        supported_thinking_effort_levels(&self.model)
589    }
590
591    fn supports_split_token_display(&self) -> bool {
592        true
593    }
594
595    fn telemetry_id(&self) -> String {
596        format!("openai/{}", self.model.id())
597    }
598
599    fn max_token_count(&self) -> u64 {
600        self.model.max_token_count()
601    }
602
603    fn max_output_tokens(&self) -> Option<u64> {
604        self.model.max_output_tokens()
605    }
606
607    fn stream_completion(
608        &self,
609        mut request: LanguageModelRequest,
610        cx: &AsyncApp,
611    ) -> BoxFuture<
612        'static,
613        Result<
614            futures::stream::BoxStream<
615                'static,
616                Result<LanguageModelCompletionEvent, LanguageModelCompletionError>,
617            >,
618            LanguageModelCompletionError,
619        >,
620    > {
621        if !self.model.supports_priority() {
622            request.speed = None;
623        }
624        if self.model.uses_responses_api() {
625            normalize_open_ai_response_thinking_effort(&mut request, &self.model);
626            let request = match into_open_ai_response(
627                request,
628                self.model.id(),
629                self.model.supports_parallel_tool_calls(),
630                self.model.supports_prompt_cache_key(),
631                self.max_output_tokens(),
632                default_thinking_reasoning_effort(&self.model),
633                self.model
634                    .supported_reasoning_efforts()
635                    .contains(&open_ai::ReasoningEffort::None),
636                &OPEN_AI_PROVIDER_ID,
637            ) {
638                Ok(request) => request,
639                Err(error) => return async move { Err(error.into()) }.boxed(),
640            };
641            let completions = self.stream_response(request, cx);
642            async move {
643                let mapper = OpenAiResponseEventMapper::new(OPEN_AI_PROVIDER_ID);
644                Ok(mapper.map_stream(completions.await?).boxed())
645            }
646            .boxed()
647        } else {
648            let request = match into_open_ai(
649                request,
650                self.model.id(),
651                self.model.supports_parallel_tool_calls(),
652                self.model.supports_prompt_cache_key(),
653                self.max_output_tokens(),
654                ChatCompletionMaxTokensParameter::MaxCompletionTokens,
655                None,
656                false,
657            ) {
658                Ok(request) => request,
659                Err(error) => return async move { Err(error.into()) }.boxed(),
660            };
661            let completions = self.stream_completion(request, cx);
662            async move {
663                let mapper = OpenAiEventMapper::new();
664                Ok(mapper.map_stream(completions.await?).boxed())
665            }
666            .boxed()
667        }
668    }
669}
670