tenant.openagents/omega

1mod api_key;
2mod registry;
3mod request;
4
5#[cfg(any(test, feature = "test-support"))]
6pub mod fake_provider;
7
8pub use language_model_core::*;
9
10use anyhow::Result;
11use futures::FutureExt;
12use futures::{StreamExt, future::BoxFuture, stream::BoxStream};
13use gpui::{AnyView, App, AsyncApp, Task, Window};
14use icons::IconName;
15use parking_lot::Mutex;
16use std::sync::Arc;
17
18pub type CreateProviderSettingsView = Arc<dyn Fn(&mut Window, &mut App) -> AnyView + 'static>;
19
20pub use crate::api_key::{ApiKey, ApiKeyState};
21pub use crate::registry::*;
22pub use crate::request::{LanguageModelImageExt, gpui_size_to_image_size, image_size_to_gpui};
23pub use env_var::{EnvVar, env_var};
24
25pub fn init(cx: &mut App) {
26    registry::init(cx);
27}
28
29#[derive(Debug, Clone, PartialEq, Eq)]
30pub struct DisabledReason(pub SharedString);
31
32impl DisabledReason {
33    pub fn new(reason: impl Into<SharedString>) -> Self {
34        Self(reason.into())
35    }
36}
37
38/// The outcome of an explicit [`LanguageModel::compact`] request.
39#[derive(Debug, Clone, PartialEq)]
40pub struct CompactionResult {
41    /// The replacement context to persist and use in subsequent requests.
42    pub context: CompactedContext,
43    /// Token usage of the compaction request itself, as reported by the
44    /// provider.
45    pub usage: TokenUsage,
46}
47
48pub struct LanguageModelTextStream {
49    pub message_id: Option<String>,
50    pub stream: BoxStream<'static, Result<String, LanguageModelCompletionError>>,
51    // Has complete token usage after the stream has finished
52    pub last_token_usage: Arc<Mutex<TokenUsage>>,
53}
54
55impl Default for LanguageModelTextStream {
56    fn default() -> Self {
57        Self {
58            message_id: None,
59            stream: Box::pin(futures::stream::empty()),
60            last_token_usage: Arc::new(Mutex::new(TokenUsage::default())),
61        }
62    }
63}
64
65pub trait LanguageModel: Send + Sync {
66    fn id(&self) -> LanguageModelId;
67    fn name(&self) -> LanguageModelName;
68    fn provider_id(&self) -> LanguageModelProviderId;
69    fn provider_name(&self) -> LanguageModelProviderName;
70    fn upstream_provider_id(&self) -> LanguageModelProviderId {
71        self.provider_id()
72    }
73    fn upstream_provider_name(&self) -> LanguageModelProviderName {
74        self.provider_name()
75    }
76
77    /// Returns whether this model is the "latest", so we can highlight it in the UI.
78    fn is_latest(&self) -> bool {
79        false
80    }
81
82    /// Whether the model is currently disabled and, if so, why this is the case.
83    fn is_disabled(&self) -> Option<DisabledReason> {
84        None
85    }
86
87    /// Whether requests to this model require the user to consent to the
88    /// upstream provider retaining inference logs (i.e. the model cannot be
89    /// offered with Zero Data Retention).
90    fn requires_data_retention(&self) -> bool {
91        false
92    }
93
94    /// When this model refuses a request, the model ID to fall back to (same provider).
95    fn refusal_fallback_model_id(&self) -> Option<&'static str> {
96        None
97    }
98
99    fn telemetry_id(&self) -> String;
100
101    fn api_key(&self, _cx: &App) -> Option<String> {
102        None
103    }
104
105    /// Information about the cost of using this model, if available.
106    fn model_cost_info(&self) -> Option<LanguageModelCostInfo> {
107        None
108    }
109
110    /// Whether this model supports thinking.
111    fn supports_thinking(&self) -> bool {
112        false
113    }
114
115    /// Whether thinking can be turned off entirely for this model. Some
116    /// models (e.g. Claude Fable 5) always think and cannot honor an "off"
117    /// request. Only meaningful when `supports_thinking` returns `true`.
118    fn supports_disabling_thinking(&self) -> bool {
119        true
120    }
121
122    fn supports_fast_mode(&self) -> bool {
123        false
124    }
125
126    /// Returns the list of supported effort levels that can be used when thinking.
127    fn supported_effort_levels(&self) -> Vec<LanguageModelEffortLevel> {
128        Vec::new()
129    }
130
131    /// Returns the default effort level to use when thinking.
132    fn default_effort_level(&self) -> Option<LanguageModelEffortLevel> {
133        self.supported_effort_levels()
134            .into_iter()
135            .find(|effort_level| effort_level.is_default)
136    }
137
138    /// Whether this model supports provider-side automatic context
139    /// compaction (requested via `LanguageModelRequest::compact_at_tokens`).
140    fn supports_server_side_compaction(&self) -> bool {
141        false
142    }
143
144    fn supports_explicit_compaction(&self) -> bool {
145        false
146    }
147
148    fn compact(
149        &self,
150        _request: LanguageModelRequest,
151        _cx: &AsyncApp,
152    ) -> BoxFuture<'static, Result<CompactionResult, LanguageModelCompletionError>> {
153        let provider = self.provider_name();
154        async move {
155            Err(LanguageModelCompletionError::Other(anyhow::anyhow!(
156                "{provider} does not support explicit compaction"
157            )))
158        }
159        .boxed()
160    }
161
162    /// Whether this model supports images
163    fn supports_images(&self) -> bool;
164
165    /// Whether this model supports tools.
166    fn supports_tools(&self) -> bool;
167
168    /// Whether this model supports choosing which tool to use.
169    fn supports_tool_choice(&self, choice: LanguageModelToolChoice) -> bool;
170
171    /// Returns whether this model or provider supports streaming tool calls;
172    fn supports_streaming_tools(&self) -> bool {
173        false
174    }
175
176    /// Returns whether this model/provider reports accurate split input/output token counts.
177    /// When true, the UI may show separate input/output token indicators.
178    fn supports_split_token_display(&self) -> bool {
179        false
180    }
181
182    fn tool_input_format(&self) -> LanguageModelToolSchemaFormat {
183        LanguageModelToolSchemaFormat::JsonSchema
184    }
185
186    fn max_token_count(&self) -> u64;
187    fn max_output_tokens(&self) -> Option<u64> {
188        None
189    }
190
191    fn stream_completion(
192        &self,
193        request: LanguageModelRequest,
194        cx: &AsyncApp,
195    ) -> BoxFuture<
196        'static,
197        Result<
198            BoxStream<'static, Result<LanguageModelCompletionEvent, LanguageModelCompletionError>>,
199            LanguageModelCompletionError,
200        >,
201    >;
202
203    fn stream_completion_text(
204        &self,
205        request: LanguageModelRequest,
206        cx: &AsyncApp,
207    ) -> BoxFuture<'static, Result<LanguageModelTextStream, LanguageModelCompletionError>> {
208        let future = self.stream_completion(request, cx);
209
210        async move {
211            let events = future.await?;
212            let mut events = events.fuse();
213            let mut message_id = None;
214            let mut first_item_text = None;
215            let last_token_usage = Arc::new(Mutex::new(TokenUsage::default()));
216
217            if let Some(first_event) = events.next().await {
218                match first_event {
219                    Ok(LanguageModelCompletionEvent::StartMessage { message_id: id }) => {
220                        message_id = Some(id);
221                    }
222                    Ok(LanguageModelCompletionEvent::Text(text)) => {
223                        first_item_text = Some(text);
224                    }
225                    _ => (),
226                }
227            }
228
229            let stream = futures::stream::iter(first_item_text.map(Ok))
230                .chain(events.filter_map({
231                    let last_token_usage = last_token_usage.clone();
232                    move |result| {
233                        let last_token_usage = last_token_usage.clone();
234                        async move {
235                            match result {
236                                Ok(LanguageModelCompletionEvent::Queued { .. }) => None,
237                                Ok(LanguageModelCompletionEvent::Started) => None,
238                                Ok(LanguageModelCompletionEvent::StartMessage { .. }) => None,
239                                Ok(LanguageModelCompletionEvent::Text(text)) => Some(Ok(text)),
240                                Ok(LanguageModelCompletionEvent::Thinking { .. }) => None,
241                                Ok(LanguageModelCompletionEvent::RedactedThinking { .. }) => None,
242                                Ok(LanguageModelCompletionEvent::ReasoningDetails(_)) => None,
243                                Ok(LanguageModelCompletionEvent::Stop(_)) => None,
244                                Ok(LanguageModelCompletionEvent::ToolUse(_)) => None,
245                                Ok(LanguageModelCompletionEvent::ToolUseJsonParseError {
246                                    ..
247                                }) => None,
248                                Ok(LanguageModelCompletionEvent::Compaction(_)) => None,
249                                Ok(LanguageModelCompletionEvent::UsageUpdate(token_usage)) => {
250                                    *last_token_usage.lock() = token_usage;
251                                    None
252                                }
253                                Err(err) => Some(Err(err)),
254                            }
255                        }
256                    }
257                }))
258                .boxed();
259
260            Ok(LanguageModelTextStream {
261                message_id,
262                stream,
263                last_token_usage,
264            })
265        }
266        .boxed()
267    }
268
269    fn stream_completion_tool(
270        &self,
271        request: LanguageModelRequest,
272        cx: &AsyncApp,
273    ) -> BoxFuture<'static, Result<LanguageModelToolUse, LanguageModelCompletionError>> {
274        let future = self.stream_completion(request, cx);
275
276        async move {
277            let events = future.await?;
278            let mut events = events.fuse();
279
280            // Iterate through events until we find a complete ToolUse
281            while let Some(event) = events.next().await {
282                match event {
283                    Ok(LanguageModelCompletionEvent::ToolUse(tool_use))
284                        if tool_use.is_input_complete =>
285                    {
286                        return Ok(tool_use);
287                    }
288                    Err(err) => {
289                        return Err(err);
290                    }
291                    _ => {}
292                }
293            }
294
295            // Stream ended without a complete tool use
296            Err(LanguageModelCompletionError::Other(anyhow::anyhow!(
297                "Stream ended without receiving a complete tool use"
298            )))
299        }
300        .boxed()
301    }
302
303    #[cfg(any(test, feature = "test-support"))]
304    fn as_fake(&self) -> &fake_provider::FakeLanguageModel {
305        unimplemented!()
306    }
307}
308
309impl std::fmt::Debug for dyn LanguageModel {
310    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
311        f.debug_struct("<dyn LanguageModel>")
312            .field("id", &self.id())
313            .field("name", &self.name())
314            .field("provider_id", &self.provider_id())
315            .field("provider_name", &self.provider_name())
316            .field("upstream_provider_name", &self.upstream_provider_name())
317            .field("upstream_provider_id", &self.upstream_provider_id())
318            .field("upstream_provider_id", &self.upstream_provider_id())
319            .field("supports_streaming_tools", &self.supports_streaming_tools())
320            .finish()
321    }
322}
323
324/// Either a built-in icon name or a path to an external SVG.
325#[derive(Debug, Clone, PartialEq, Eq)]
326pub enum IconOrSvg {
327    /// A built-in icon from Zed's icon set.
328    Icon(IconName),
329    /// Path to a custom SVG icon file.
330    Svg(SharedString),
331}
332
333impl Default for IconOrSvg {
334    fn default() -> Self {
335        Self::Icon(IconName::OmegaAssistant)
336    }
337}
338
339pub trait LanguageModelProvider: 'static {
340    fn id(&self) -> LanguageModelProviderId;
341    fn name(&self) -> LanguageModelProviderName;
342    fn icon(&self) -> IconOrSvg {
343        IconOrSvg::default()
344    }
345    fn default_model(&self, cx: &App) -> Option<Arc<dyn LanguageModel>>;
346    fn default_fast_model(&self, cx: &App) -> Option<Arc<dyn LanguageModel>>;
347    fn provided_models(&self, cx: &App) -> Vec<Arc<dyn LanguageModel>>;
348    fn recommended_models(&self, _cx: &App) -> Vec<Arc<dyn LanguageModel>> {
349        Vec::new()
350    }
351    fn is_authenticated(&self, cx: &App) -> bool;
352    fn authenticate(&self, cx: &mut App) -> Task<Result<(), AuthenticateError>>;
353    fn settings_view(&self, cx: &mut App) -> Option<ProviderSettingsView>;
354
355    fn set_api_key(&self, _key: Option<String>, _cx: &mut App) -> Task<Result<()>> {
356        Task::ready(Ok(()))
357    }
358
359    /// Copy shown when this provider rejects a request as unauthenticated
360    /// (HTTP 401). The default assumes API-key authentication; providers using
361    /// other mechanisms (account or subscription based auth) should override
362    /// this so users aren't told to check an API key they don't have.
363    fn authentication_error_message(&self) -> SharedString {
364        format!(
365            "The API key for {} is invalid or has expired. \
366            Update your key in Settings > AI > LLM Providers to continue.",
367            self.name().0
368        )
369        .into()
370    }
371
372    /// Copy shown when a request fails because no credentials are configured
373    /// for this provider. The default assumes API-key authentication;
374    /// providers using other mechanisms (account or subscription based auth)
375    /// should override this.
376    fn missing_credentials_error_message(&self) -> SharedString {
377        format!(
378            "No API key is configured for {}. \
379            Add your key in Settings > AI > LLM Providers to continue.",
380            self.name().0
381        )
382        .into()
383    }
384
385    /// Copy shown the first time a user enables fast mode for a model from
386    /// this provider. Returning `None` skips the confirmation prompt and lets
387    /// the toggle apply silently.
388    fn fast_mode_confirmation(&self, _cx: &App) -> Option<FastModeConfirmation> {
389        None
390    }
391}
392
393/// A provider's settings UI, modeled as mutually exclusive presentation modes.
394#[derive(Clone)]
395pub enum ProviderSettingsView {
396    ApiKey(ApiKeyConfiguration),
397    Inline(InlineProviderSettings),
398    SubPage(SubPageProviderSettings),
399}
400
401#[derive(Clone)]
402pub struct InlineProviderSettings {
403    pub title: Option<SharedString>,
404    pub description: Option<InlineDescription>,
405    pub create_view: CreateProviderSettingsView,
406}
407
408#[derive(Clone)]
409pub struct SubPageProviderSettings {
410    pub description: Option<InlineDescription>,
411    pub create_view: CreateProviderSettingsView,
412}
413
414impl SubPageProviderSettings {
415    pub fn new(create_view: impl Fn(&mut Window, &mut App) -> AnyView + 'static) -> Self {
416        Self {
417            description: None,
418            create_view: Arc::new(create_view),
419        }
420    }
421
422    pub fn description(mut self, description: InlineDescription) -> Self {
423        self.description = Some(description);
424        self
425    }
426}
427
428impl ApiKeyConfiguration {
429    pub fn new(
430        has_key: bool,
431        is_from_env_var: bool,
432        env_var_name: SharedString,
433        api_key_url: SharedString,
434    ) -> Self {
435        Self {
436            has_key,
437            is_from_env_var,
438            env_var_name,
439            api_key_url,
440        }
441    }
442}
443
444/// A live snapshot of a single-API-key provider's credential state, used by the
445/// settings UI to render the provider's "API Key" section.
446#[derive(Clone)]
447pub struct ApiKeyConfiguration {
448    pub has_key: bool,
449    pub is_from_env_var: bool,
450    pub env_var_name: SharedString,
451    pub api_key_url: SharedString,
452}
453
454/// The subtitle rendered beneath a provider's name when its configuration is
455/// shown inline.
456#[derive(Clone)]
457pub enum InlineDescription {
458    /// A clickable "Where to find key" link pointing at the given URL, for
459    /// API-key based providers.
460    ApiKeyUrl(SharedString),
461    /// Plain descriptive text, e.g. explaining a sign-in based provider.
462    Text(SharedString),
463}
464
465/// Provider-specific copy shown the first time a user enables fast mode.
466#[derive(Debug, Clone)]
467pub struct FastModeConfirmation {
468    pub title: SharedString,
469    pub message: SharedString,
470}
471
472pub trait LanguageModelProviderState: 'static {
473    type ObservableEntity;
474
475    fn observable_entity(&self) -> Option<gpui::Entity<Self::ObservableEntity>>;
476
477    fn subscribe<T: 'static>(
478        &self,
479        cx: &mut gpui::Context<T>,
480        callback: impl Fn(&mut T, &mut gpui::Context<T>) + 'static,
481    ) -> Option<gpui::Subscription> {
482        let entity = self.observable_entity()?;
483        Some(cx.observe(&entity, move |this, _, cx| {
484            callback(this, cx);
485        }))
486    }
487}
488
489#[derive(Clone, Debug, PartialEq)]
490pub enum LanguageModelCostInfo {
491    /// Cost per 1,000 input and output tokens
492    TokenCost {
493        input_token_cost_per_1m: f64,
494        output_token_cost_per_1m: f64,
495    },
496    /// Cost per request
497    RequestCost { cost_per_request: f64 },
498}
499
500impl LanguageModelCostInfo {
501    pub fn to_shared_string(&self) -> SharedString {
502        match self {
503            LanguageModelCostInfo::RequestCost { cost_per_request } => {
504                let cost_str = format!("{}×", Self::cost_value_to_string(cost_per_request));
505                SharedString::from(cost_str)
506            }
507            LanguageModelCostInfo::TokenCost {
508                input_token_cost_per_1m,
509                output_token_cost_per_1m,
510            } => {
511                let input_cost = Self::cost_value_to_string(input_token_cost_per_1m);
512                let output_cost = Self::cost_value_to_string(output_token_cost_per_1m);
513                SharedString::from(format!("{}$/{}$", input_cost, output_cost))
514            }
515        }
516    }
517
518    fn cost_value_to_string(cost: &f64) -> SharedString {
519        if (cost.fract() - 0.0).abs() < std::f64::EPSILON {
520            SharedString::from(format!("{:.0}", cost))
521        } else {
522            SharedString::from(format!("{:.2}", cost))
523        }
524    }
525}
526