diff --git a/crates/agent/src/active_thread.rs b/crates/agent/src/active_thread.rs index 9a5a181cd6..242ceff376 100644 --- a/crates/agent/src/active_thread.rs +++ b/crates/agent/src/active_thread.rs @@ -1293,6 +1293,7 @@ impl ActiveThread { let request = language_model::LanguageModelRequest { thread_id: None, prompt_id: None, + mode: None, messages: vec![request_message], tools: vec![], stop: vec![], diff --git a/crates/agent/src/buffer_codegen.rs b/crates/agent/src/buffer_codegen.rs index b13270c0b4..4e0ee6a9c3 100644 --- a/crates/agent/src/buffer_codegen.rs +++ b/crates/agent/src/buffer_codegen.rs @@ -460,6 +460,7 @@ impl CodegenAlternative { LanguageModelRequest { thread_id: None, prompt_id: None, + mode: None, tools: Vec::new(), stop: Vec::new(), temperature: None, diff --git a/crates/agent/src/message_editor.rs b/crates/agent/src/message_editor.rs index 25c1ea3f8f..cc4f439f68 100644 --- a/crates/agent/src/message_editor.rs +++ b/crates/agent/src/message_editor.rs @@ -11,6 +11,7 @@ use editor::{ ContextMenuOptions, ContextMenuPlacement, Editor, EditorElement, EditorEvent, EditorMode, EditorStyle, MultiBuffer, }; +use feature_flags::{FeatureFlagAppExt, NewBillingFeatureFlag}; use file_icons::FileIcons; use fs::Fs; use futures::future::Shared; @@ -33,6 +34,7 @@ use theme::ThemeSettings; use ui::{Disclosure, KeyBinding, PopoverMenuHandle, Tooltip, prelude::*}; use util::ResultExt as _; use workspace::Workspace; +use zed_llm_client::CompletionMode; use crate::assistant_model_selector::AssistantModelSelector; use crate::context_picker::{ContextPicker, ContextPickerCompletionProvider}; @@ -50,7 +52,6 @@ pub struct MessageEditor { thread: Entity, incompatible_tools_state: Entity, editor: Entity, - #[allow(dead_code)] workspace: WeakEntity, project: Entity, context_store: Entity, @@ -419,6 +420,37 @@ impl MessageEditor { } } + fn render_max_mode_toggle(&self, cx: &mut Context) -> Option { + if !cx.has_flag::() { + return None; + } + + let model = LanguageModelRegistry::read_global(cx) + .default_model() + .map(|default| default.model.clone())?; + if !model.supports_max_mode() { + return None; + } + + let active_completion_mode = self.thread.read(cx).completion_mode(); + + Some( + IconButton::new("max-mode", IconName::SquarePlus) + .icon_size(IconSize::Small) + .toggle_state(active_completion_mode == Some(CompletionMode::Max)) + .on_click(cx.listener(move |this, _event, _window, cx| { + this.thread.update(cx, |thread, _cx| { + thread.set_completion_mode(match active_completion_mode { + Some(CompletionMode::Max) => Some(CompletionMode::Normal), + Some(CompletionMode::Normal) | None => Some(CompletionMode::Max), + }); + }); + })) + .tooltip(Tooltip::text("Max Mode")) + .into_any_element(), + ) + } + fn render_editor( &self, font_size: Rems, @@ -579,6 +611,7 @@ impl MessageEditor { }), ) }) + .children(self.render_max_mode_toggle(cx)) .child(self.model_selector.clone()) .map({ let focus_handle = focus_handle.clone(); @@ -1100,6 +1133,7 @@ impl MessageEditor { let request = language_model::LanguageModelRequest { thread_id: None, prompt_id: None, + mode: None, messages: vec![request_message], tools: vec![], stop: vec![], diff --git a/crates/agent/src/terminal_inline_assistant.rs b/crates/agent/src/terminal_inline_assistant.rs index b7690fa580..1056feb3fb 100644 --- a/crates/agent/src/terminal_inline_assistant.rs +++ b/crates/agent/src/terminal_inline_assistant.rs @@ -276,6 +276,7 @@ impl TerminalInlineAssistant { LanguageModelRequest { thread_id: None, prompt_id: None, + mode: None, messages: vec![request_message], tools: Vec::new(), stop: Vec::new(), diff --git a/crates/agent/src/thread.rs b/crates/agent/src/thread.rs index 7d7c6fefa9..f439f1f8fc 100644 --- a/crates/agent/src/thread.rs +++ b/crates/agent/src/thread.rs @@ -34,6 +34,7 @@ use settings::Settings; use thiserror::Error; use util::{ResultExt as _, TryFutureExt as _, post_inc}; use uuid::Uuid; +use zed_llm_client::CompletionMode; use crate::context::{AgentContext, ContextLoadResult, LoadedContext}; use crate::thread_store::{ @@ -290,6 +291,7 @@ pub struct Thread { summary: Option, pending_summary: Task>, detailed_summary_state: DetailedSummaryState, + completion_mode: Option, messages: Vec, next_message_id: MessageId, last_prompt_id: PromptId, @@ -339,6 +341,7 @@ impl Thread { summary: None, pending_summary: Task::ready(None), detailed_summary_state: DetailedSummaryState::NotGenerated, + completion_mode: None, messages: Vec::new(), next_message_id: MessageId(0), last_prompt_id: PromptId::new(), @@ -394,6 +397,7 @@ impl Thread { summary: Some(serialized.summary), pending_summary: Task::ready(None), detailed_summary_state: serialized.detailed_summary_state, + completion_mode: None, messages: serialized .messages .into_iter() @@ -518,6 +522,14 @@ impl Thread { } } + pub fn completion_mode(&self) -> Option { + self.completion_mode + } + + pub fn set_completion_mode(&mut self, mode: Option) { + self.completion_mode = mode; + } + pub fn message(&self, id: MessageId) -> Option<&Message> { let index = self .messages @@ -930,6 +942,12 @@ impl Thread { self.remaining_turns -= 1; let mut request = self.to_completion_request(cx); + request.mode = if model.supports_max_mode() { + self.completion_mode + } else { + None + }; + if model.supports_tools() { request.tools = self .tools() @@ -967,6 +985,7 @@ impl Thread { let mut request = LanguageModelRequest { thread_id: Some(self.id.to_string()), prompt_id: Some(self.last_prompt_id.to_string()), + mode: None, messages: vec![], tools: Vec::new(), stop: Vec::new(), @@ -1063,6 +1082,7 @@ impl Thread { let mut request = LanguageModelRequest { thread_id: None, prompt_id: None, + mode: None, messages: vec![], tools: Vec::new(), stop: Vec::new(), diff --git a/crates/assistant/src/inline_assistant.rs b/crates/assistant/src/inline_assistant.rs index 11ef6bd53e..6495bea21d 100644 --- a/crates/assistant/src/inline_assistant.rs +++ b/crates/assistant/src/inline_assistant.rs @@ -2981,6 +2981,7 @@ impl CodegenAlternative { Ok(LanguageModelRequest { thread_id: None, prompt_id: None, + mode: None, messages, tools: Vec::new(), stop: Vec::new(), diff --git a/crates/assistant/src/terminal_inline_assistant.rs b/crates/assistant/src/terminal_inline_assistant.rs index b7ac8fb43d..fb47173acd 100644 --- a/crates/assistant/src/terminal_inline_assistant.rs +++ b/crates/assistant/src/terminal_inline_assistant.rs @@ -294,6 +294,7 @@ impl TerminalInlineAssistant { Ok(LanguageModelRequest { thread_id: None, prompt_id: None, + mode: None, messages, tools: Vec::new(), stop: Vec::new(), diff --git a/crates/assistant_context_editor/src/context.rs b/crates/assistant_context_editor/src/context.rs index 37a0327d29..5f5a6e8ee8 100644 --- a/crates/assistant_context_editor/src/context.rs +++ b/crates/assistant_context_editor/src/context.rs @@ -2559,6 +2559,7 @@ impl AssistantContext { let mut completion_request = LanguageModelRequest { thread_id: None, prompt_id: None, + mode: None, messages: Vec::new(), tools: Vec::new(), stop: Vec::new(), diff --git a/crates/eval/src/instance.rs b/crates/eval/src/instance.rs index b92bd1bd0f..dcbd9998df 100644 --- a/crates/eval/src/instance.rs +++ b/crates/eval/src/instance.rs @@ -573,6 +573,7 @@ impl ExampleInstance { let request = LanguageModelRequest { thread_id: None, prompt_id: None, + mode: None, messages: vec![LanguageModelRequestMessage { role: Role::User, content: vec![MessageContent::Text(to_prompt(assertion.description))], diff --git a/crates/git_ui/src/git_panel.rs b/crates/git_ui/src/git_panel.rs index 6f1d57c7a2..d2d2811c08 100644 --- a/crates/git_ui/src/git_panel.rs +++ b/crates/git_ui/src/git_panel.rs @@ -1746,6 +1746,7 @@ impl GitPanel { let request = LanguageModelRequest { thread_id: None, prompt_id: None, + mode: None, messages: vec![LanguageModelRequestMessage { role: Role::User, content: vec![content.into()], diff --git a/crates/language_model/src/language_model.rs b/crates/language_model/src/language_model.rs index 8c9860ca75..81a03c3305 100644 --- a/crates/language_model/src/language_model.rs +++ b/crates/language_model/src/language_model.rs @@ -240,6 +240,26 @@ pub trait LanguageModel: Send + Sync { /// Whether this model supports tools. fn supports_tools(&self) -> bool; + /// Returns whether this model supports "max mode"; + fn supports_max_mode(&self) -> bool { + if self.provider_id().0 != ZED_CLOUD_PROVIDER_ID { + return false; + } + + const MAX_MODE_CAPABLE_MODELS: &[CloudModel] = &[ + CloudModel::Anthropic(anthropic::Model::Claude3_7Sonnet), + CloudModel::Anthropic(anthropic::Model::Claude3_7SonnetThinking), + ]; + + for model in MAX_MODE_CAPABLE_MODELS { + if self.id().0 == model.id() { + return true; + } + } + + false + } + fn tool_input_format(&self) -> LanguageModelToolSchemaFormat { LanguageModelToolSchemaFormat::JsonSchema } diff --git a/crates/language_model/src/request.rs b/crates/language_model/src/request.rs index 9816bc73c0..19ce8f8b7d 100644 --- a/crates/language_model/src/request.rs +++ b/crates/language_model/src/request.rs @@ -11,6 +11,7 @@ use gpui::{ use image::{DynamicImage, ImageDecoder, codecs::png::PngEncoder, imageops::resize}; use serde::{Deserialize, Serialize}; use util::ResultExt; +use zed_llm_client::CompletionMode; #[derive(Clone, PartialEq, Eq, Serialize, Deserialize, Hash)] pub struct LanguageModelImage { @@ -252,6 +253,7 @@ pub struct LanguageModelRequestTool { pub struct LanguageModelRequest { pub thread_id: Option, pub prompt_id: Option, + pub mode: Option, pub messages: Vec, pub tools: Vec, pub stop: Vec, diff --git a/crates/rules_library/src/rules_library.rs b/crates/rules_library/src/rules_library.rs index 74a5d4b948..8f5fb4f953 100644 --- a/crates/rules_library/src/rules_library.rs +++ b/crates/rules_library/src/rules_library.rs @@ -922,6 +922,7 @@ impl RulesLibrary { LanguageModelRequest { thread_id: None, prompt_id: None, + mode: None, messages: vec![LanguageModelRequestMessage { role: Role::System, content: vec![body.to_string().into()], diff --git a/crates/semantic_index/src/summary_index.rs b/crates/semantic_index/src/summary_index.rs index f1b829f9b5..c09ae4beba 100644 --- a/crates/semantic_index/src/summary_index.rs +++ b/crates/semantic_index/src/summary_index.rs @@ -559,6 +559,7 @@ impl SummaryIndex { let request = LanguageModelRequest { thread_id: None, prompt_id: None, + mode: None, messages: vec![LanguageModelRequestMessage { role: Role::User, content: vec![prompt.into()],