From 1038e1c2ef7c09c27e013592fd5f1b45eb230f2c Mon Sep 17 00:00:00 2001 From: Richard Feldman Date: Mon, 8 Dec 2025 16:59:49 -0500 Subject: [PATCH] Clean up some duplicated code --- Cargo.lock | 1 - crates/language_model/src/registry.rs | 5 - crates/language_models/Cargo.toml | 1 - crates/language_models/src/api_key.rs | 14 - crates/language_models/src/extension.rs | 42 +- crates/language_models/src/language_models.rs | 2 +- crates/language_models/src/provider/cloud.rs | 442 +----------------- crates/language_models/src/provider/google.rs | 18 - 8 files changed, 7 insertions(+), 518 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index bfe1ab0dd14709a423735c35842d0b14cac7628d..bf910847a0c0fd8dc636ab86b6f426ada881753a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -8858,7 +8858,6 @@ dependencies = [ "credentials_provider", "deepseek", "editor", - "extension", "extension_host", "fs", "futures 0.3.31", diff --git a/crates/language_model/src/registry.rs b/crates/language_model/src/registry.rs index e14ebb0d3b82569a5a30be0b5601a478294490e8..bb563b14d70d4b938c739bdc29bdd46f38fa05b4 100644 --- a/crates/language_model/src/registry.rs +++ b/crates/language_model/src/registry.rs @@ -237,11 +237,6 @@ impl LanguageModelRegistry { } } - /// Returns the set of installed LLM extension IDs. - pub fn installed_llm_extension_ids(&self) -> &HashSet> { - &self.installed_llm_extension_ids - } - /// Returns true if a provider should be hidden from the UI. /// Built-in providers are hidden when their corresponding extension is installed. pub fn should_hide_provider(&self, provider_id: &LanguageModelProviderId) -> bool { diff --git a/crates/language_models/Cargo.toml b/crates/language_models/Cargo.toml index 4aaf9dcec5d33a8625297ffba98e1dbbc1c57fa8..0bde1d6666d5446ceb6d6caad5d5a17d14b84413 100644 --- a/crates/language_models/Cargo.toml +++ b/crates/language_models/Cargo.toml @@ -28,7 +28,6 @@ convert_case.workspace = true copilot.workspace = true credentials_provider.workspace = true deepseek = { workspace = true, features = ["schemars"] } -extension.workspace = true extension_host.workspace = true fs.workspace = true futures.workspace = true diff --git a/crates/language_models/src/api_key.rs b/crates/language_models/src/api_key.rs index 122234b6ced6d0bf1b7a0d684683c841824ccd2d..1f4271d04a24dc9bd1a22a1bb7682527086d7639 100644 --- a/crates/language_models/src/api_key.rs +++ b/crates/language_models/src/api_key.rs @@ -223,10 +223,6 @@ impl ApiKeyState { } impl ApiKey { - pub fn key(&self) -> &str { - &self.key - } - pub fn from_env(env_var_name: SharedString, key: &str) -> Self { Self { source: ApiKeySource::EnvVar(env_var_name), @@ -234,16 +230,6 @@ impl ApiKey { } } - pub async fn load_from_system_keychain( - url: &str, - credentials_provider: &dyn CredentialsProvider, - cx: &AsyncApp, - ) -> Result { - Self::load_from_system_keychain_impl(url, credentials_provider, cx) - .await - .into_authenticate_result() - } - async fn load_from_system_keychain_impl( url: &str, credentials_provider: &dyn CredentialsProvider, diff --git a/crates/language_models/src/extension.rs b/crates/language_models/src/extension.rs index 59dc98f211c19520c92124655ddace799d189158..ed26bfb511684fb0e338408e8cb6b81dd786c314 100644 --- a/crates/language_models/src/extension.rs +++ b/crates/language_models/src/extension.rs @@ -1,12 +1,9 @@ use collections::HashMap; -use extension::{ExtensionLanguageModelProviderProxy, LanguageModelProviderRegistration}; -use gpui::{App, Entity}; -use language_model::{LanguageModelProviderId, LanguageModelRegistry}; -use std::sync::{Arc, LazyLock}; +use std::sync::LazyLock; /// Maps built-in provider IDs to their corresponding extension IDs. /// When an extension with this ID is installed, the built-in provider should be hidden. -pub static BUILTIN_TO_EXTENSION_MAP: LazyLock> = +static BUILTIN_TO_EXTENSION_MAP: LazyLock> = LazyLock::new(|| { let mut map = HashMap::default(); map.insert("anthropic", "anthropic"); @@ -21,38 +18,3 @@ pub static BUILTIN_TO_EXTENSION_MAP: LazyLock Option<&'static str> { BUILTIN_TO_EXTENSION_MAP.get(provider_id).copied() } - -/// Returns true if the given provider ID is a built-in provider that can be hidden by an extension. -pub fn is_hideable_builtin_provider(provider_id: &str) -> bool { - BUILTIN_TO_EXTENSION_MAP.contains_key(provider_id) -} - -/// Proxy implementation that registers extension-based language model providers -/// with the LanguageModelRegistry. -pub struct ExtensionLanguageModelProxy { - registry: Entity, -} - -impl ExtensionLanguageModelProxy { - pub fn new(registry: Entity) -> Self { - Self { registry } - } -} - -impl ExtensionLanguageModelProviderProxy for ExtensionLanguageModelProxy { - fn register_language_model_provider( - &self, - provider_id: Arc, - register_fn: LanguageModelProviderRegistration, - cx: &mut App, - ) { - let _ = provider_id; - register_fn(cx); - } - - fn unregister_language_model_provider(&self, provider_id: Arc, cx: &mut App) { - self.registry.update(cx, |registry, cx| { - registry.unregister_provider(LanguageModelProviderId::from(provider_id), cx); - }); - } -} diff --git a/crates/language_models/src/language_models.rs b/crates/language_models/src/language_models.rs index 2e72539768c7f08192ef677ad19a36a416f2c3ed..2444a5e68c381556e32041ff9adc934f2ef7a5cf 100644 --- a/crates/language_models/src/language_models.rs +++ b/crates/language_models/src/language_models.rs @@ -14,7 +14,7 @@ pub mod provider; mod settings; pub mod ui; -pub use crate::extension::{extension_for_builtin_provider, is_hideable_builtin_provider}; +pub use crate::extension::extension_for_builtin_provider; pub use crate::google_ai_api_key::api_key_for_gemini_cli; use crate::provider::anthropic::AnthropicLanguageModelProvider; use crate::provider::bedrock::BedrockLanguageModelProvider; diff --git a/crates/language_models/src/provider/cloud.rs b/crates/language_models/src/provider/cloud.rs index 3730db2b42654bd2f9b01693d60683167731aa51..a19a427dbacb32883b1877888ec04899a2b8d427 100644 --- a/crates/language_models/src/provider/cloud.rs +++ b/crates/language_models/src/provider/cloud.rs @@ -1,8 +1,5 @@ use ai_onboarding::YoungAccountBanner; -use anthropic::{ - AnthropicModelMode, ContentDelta, Event, ResponseContent, ToolResultContent, ToolResultPart, - Usage, -}; +use anthropic::AnthropicModelMode; use anyhow::{Context as _, Result, anyhow}; use chrono::{DateTime, Utc}; use client::{Client, ModelRequestUsage, UserStore, zed_urls}; @@ -26,9 +23,8 @@ use language_model::{ LanguageModelCompletionError, LanguageModelCompletionEvent, LanguageModelId, LanguageModelName, LanguageModelProvider, LanguageModelProviderId, LanguageModelProviderName, LanguageModelProviderState, LanguageModelRequest, LanguageModelToolChoice, - LanguageModelToolResultContent, LanguageModelToolSchemaFormat, LanguageModelToolUse, - LanguageModelToolUseId, LlmApiToken, MessageContent, ModelRequestLimitReachedError, - PaymentRequiredError, RateLimiter, RefreshLlmTokenListener, Role, StopReason, + LanguageModelToolSchemaFormat, LlmApiToken, ModelRequestLimitReachedError, + PaymentRequiredError, RateLimiter, RefreshLlmTokenListener, }; use release_channel::AppVersion; use schemars::JsonSchema; @@ -46,6 +42,7 @@ use thiserror::Error; use ui::{TintColor, prelude::*}; use util::{ResultExt as _, maybe}; +use crate::provider::anthropic::{AnthropicEventMapper, count_anthropic_tokens, into_anthropic}; use crate::provider::google::{GoogleEventMapper, into_google}; use crate::provider::open_ai::{OpenAiEventMapper, count_open_ai_tokens, into_open_ai}; use crate::provider::x_ai::count_xai_tokens; @@ -1397,434 +1394,3 @@ mod tests { } } } - -fn count_anthropic_tokens( - request: LanguageModelRequest, - cx: &App, -) -> BoxFuture<'static, Result> { - use gpui::AppContext as _; - cx.background_spawn(async move { - let messages = request.messages; - let mut tokens_from_images = 0; - let mut string_messages = Vec::with_capacity(messages.len()); - - for message in messages { - let mut string_contents = String::new(); - - for content in message.content { - match content { - MessageContent::Text(text) => { - string_contents.push_str(&text); - } - MessageContent::Thinking { .. } => {} - MessageContent::RedactedThinking(_) => {} - MessageContent::Image(image) => { - tokens_from_images += image.estimate_tokens(); - } - MessageContent::ToolUse(_tool_use) => {} - MessageContent::ToolResult(tool_result) => match &tool_result.content { - LanguageModelToolResultContent::Text(text) => { - string_contents.push_str(text); - } - LanguageModelToolResultContent::Image(image) => { - tokens_from_images += image.estimate_tokens(); - } - }, - } - } - - if !string_contents.is_empty() { - string_messages.push(tiktoken_rs::ChatCompletionRequestMessage { - role: match message.role { - Role::User => "user".into(), - Role::Assistant => "assistant".into(), - Role::System => "system".into(), - }, - content: Some(string_contents), - name: None, - function_call: None, - }); - } - } - - tiktoken_rs::num_tokens_from_messages("gpt-4", &string_messages) - .map(|tokens| (tokens + tokens_from_images) as u64) - }) - .boxed() -} - -fn into_anthropic( - request: LanguageModelRequest, - model: String, - default_temperature: f32, - max_output_tokens: u64, - mode: AnthropicModelMode, -) -> anthropic::Request { - let mut new_messages: Vec = Vec::new(); - let mut system_message = String::new(); - - for message in request.messages { - if message.contents_empty() { - continue; - } - - match message.role { - Role::User | Role::Assistant => { - let mut anthropic_message_content: Vec = message - .content - .into_iter() - .filter_map(|content| match content { - MessageContent::Text(text) => { - let text = if text.chars().last().is_some_and(|c| c.is_whitespace()) { - text.trim_end().to_string() - } else { - text - }; - if !text.is_empty() { - Some(anthropic::RequestContent::Text { - text, - cache_control: None, - }) - } else { - None - } - } - MessageContent::Thinking { - text: thinking, - signature, - } => { - if !thinking.is_empty() { - Some(anthropic::RequestContent::Thinking { - thinking, - signature: signature.unwrap_or_default(), - cache_control: None, - }) - } else { - None - } - } - MessageContent::RedactedThinking(data) => { - if !data.is_empty() { - Some(anthropic::RequestContent::RedactedThinking { data }) - } else { - None - } - } - MessageContent::Image(image) => Some(anthropic::RequestContent::Image { - source: anthropic::ImageSource { - source_type: "base64".to_string(), - media_type: "image/png".to_string(), - data: image.source.to_string(), - }, - cache_control: None, - }), - MessageContent::ToolUse(tool_use) => { - Some(anthropic::RequestContent::ToolUse { - id: tool_use.id.to_string(), - name: tool_use.name.to_string(), - input: tool_use.input, - cache_control: None, - }) - } - MessageContent::ToolResult(tool_result) => { - Some(anthropic::RequestContent::ToolResult { - tool_use_id: tool_result.tool_use_id.to_string(), - is_error: tool_result.is_error, - content: match tool_result.content { - LanguageModelToolResultContent::Text(text) => { - ToolResultContent::Plain(text.to_string()) - } - LanguageModelToolResultContent::Image(image) => { - ToolResultContent::Multipart(vec![ToolResultPart::Image { - source: anthropic::ImageSource { - source_type: "base64".to_string(), - media_type: "image/png".to_string(), - data: image.source.to_string(), - }, - }]) - } - }, - cache_control: None, - }) - } - }) - .collect(); - let anthropic_role = match message.role { - Role::User => anthropic::Role::User, - Role::Assistant => anthropic::Role::Assistant, - Role::System => unreachable!("System role should never occur here"), - }; - if let Some(last_message) = new_messages.last_mut() - && last_message.role == anthropic_role - { - last_message.content.extend(anthropic_message_content); - continue; - } - - if message.cache { - let cache_control_value = Some(anthropic::CacheControl { - cache_type: anthropic::CacheControlType::Ephemeral, - }); - for message_content in anthropic_message_content.iter_mut().rev() { - match message_content { - anthropic::RequestContent::RedactedThinking { .. } => {} - anthropic::RequestContent::Text { cache_control, .. } - | anthropic::RequestContent::Thinking { cache_control, .. } - | anthropic::RequestContent::Image { cache_control, .. } - | anthropic::RequestContent::ToolUse { cache_control, .. } - | anthropic::RequestContent::ToolResult { cache_control, .. } => { - *cache_control = cache_control_value; - break; - } - } - } - } - - new_messages.push(anthropic::Message { - role: anthropic_role, - content: anthropic_message_content, - }); - } - Role::System => { - if !system_message.is_empty() { - system_message.push_str("\n\n"); - } - system_message.push_str(&message.string_contents()); - } - } - } - - anthropic::Request { - model, - messages: new_messages, - max_tokens: max_output_tokens, - system: if system_message.is_empty() { - None - } else { - Some(anthropic::StringOrContents::String(system_message)) - }, - thinking: if request.thinking_allowed - && let AnthropicModelMode::Thinking { budget_tokens } = mode - { - Some(anthropic::Thinking::Enabled { budget_tokens }) - } else { - None - }, - tools: request - .tools - .into_iter() - .map(|tool| anthropic::Tool { - name: tool.name, - description: tool.description, - input_schema: tool.input_schema, - }) - .collect(), - tool_choice: request.tool_choice.map(|choice| match choice { - LanguageModelToolChoice::Auto => anthropic::ToolChoice::Auto, - LanguageModelToolChoice::Any => anthropic::ToolChoice::Any, - LanguageModelToolChoice::None => anthropic::ToolChoice::None, - }), - metadata: None, - stop_sequences: Vec::new(), - temperature: request.temperature.or(Some(default_temperature)), - top_k: None, - top_p: None, - } -} - -struct AnthropicEventMapper { - tool_uses_by_index: collections::HashMap, - usage: Usage, - stop_reason: StopReason, -} - -impl AnthropicEventMapper { - fn new() -> Self { - Self { - tool_uses_by_index: collections::HashMap::default(), - usage: Usage::default(), - stop_reason: StopReason::EndTurn, - } - } - - fn map_event( - &mut self, - event: Event, - ) -> Vec> { - match event { - Event::ContentBlockStart { - index, - content_block, - } => match content_block { - ResponseContent::Text { text } => { - vec![Ok(LanguageModelCompletionEvent::Text(text))] - } - ResponseContent::Thinking { thinking } => { - vec![Ok(LanguageModelCompletionEvent::Thinking { - text: thinking, - signature: None, - })] - } - ResponseContent::RedactedThinking { data } => { - vec![Ok(LanguageModelCompletionEvent::RedactedThinking { data })] - } - ResponseContent::ToolUse { id, name, .. } => { - self.tool_uses_by_index.insert( - index, - RawToolUse { - id, - name, - input_json: String::new(), - }, - ); - Vec::new() - } - }, - Event::ContentBlockDelta { index, delta } => match delta { - ContentDelta::TextDelta { text } => { - vec![Ok(LanguageModelCompletionEvent::Text(text))] - } - ContentDelta::ThinkingDelta { thinking } => { - vec![Ok(LanguageModelCompletionEvent::Thinking { - text: thinking, - signature: None, - })] - } - ContentDelta::SignatureDelta { signature } => { - vec![Ok(LanguageModelCompletionEvent::Thinking { - text: "".to_string(), - signature: Some(signature), - })] - } - ContentDelta::InputJsonDelta { partial_json } => { - if let Some(tool_use) = self.tool_uses_by_index.get_mut(&index) { - tool_use.input_json.push_str(&partial_json); - - let event = serde_json::from_str::(&tool_use.input_json) - .ok() - .and_then(|input| { - let input_json_roundtripped = serde_json::to_string(&input).ok()?; - - if !tool_use.input_json.starts_with(&input_json_roundtripped) { - return None; - } - - Some(LanguageModelCompletionEvent::ToolUse( - LanguageModelToolUse { - id: LanguageModelToolUseId::from(tool_use.id.clone()), - name: tool_use.name.clone().into(), - raw_input: tool_use.input_json.clone(), - input, - is_input_complete: false, - thought_signature: None, - }, - )) - }); - - if let Some(event) = event { - vec![Ok(event)] - } else { - Vec::new() - } - } else { - Vec::new() - } - } - }, - Event::ContentBlockStop { index } => { - if let Some(tool_use) = self.tool_uses_by_index.remove(&index) { - let event_result = match serde_json::from_str(&tool_use.input_json) { - Ok(input) => Ok(LanguageModelCompletionEvent::ToolUse( - LanguageModelToolUse { - id: LanguageModelToolUseId::from(tool_use.id), - name: tool_use.name.into(), - raw_input: tool_use.input_json, - input, - is_input_complete: true, - thought_signature: None, - }, - )), - Err(json_parse_err) => { - Ok(LanguageModelCompletionEvent::ToolUseJsonParseError { - id: LanguageModelToolUseId::from(tool_use.id), - tool_name: tool_use.name.into(), - raw_input: tool_use.input_json.into(), - json_parse_error: json_parse_err.to_string(), - }) - } - }; - - vec![event_result] - } else { - Vec::new() - } - } - Event::MessageStart { message } => { - update_anthropic_usage(&mut self.usage, &message.usage); - vec![ - Ok(LanguageModelCompletionEvent::UsageUpdate( - convert_anthropic_usage(&self.usage), - )), - Ok(LanguageModelCompletionEvent::StartMessage { - message_id: message.id, - }), - ] - } - Event::MessageDelta { delta, usage } => { - update_anthropic_usage(&mut self.usage, &usage); - if let Some(stop_reason) = delta.stop_reason.as_deref() { - self.stop_reason = match stop_reason { - "end_turn" => StopReason::EndTurn, - "max_tokens" => StopReason::MaxTokens, - "tool_use" => StopReason::ToolUse, - "refusal" => StopReason::Refusal, - _ => { - log::error!("Unexpected anthropic stop_reason: {stop_reason}"); - StopReason::EndTurn - } - }; - } - vec![Ok(LanguageModelCompletionEvent::UsageUpdate( - convert_anthropic_usage(&self.usage), - ))] - } - Event::MessageStop => { - vec![Ok(LanguageModelCompletionEvent::Stop(self.stop_reason))] - } - Event::Error { error } => { - vec![Err(error.into())] - } - _ => Vec::new(), - } - } -} - -struct RawToolUse { - id: String, - name: String, - input_json: String, -} - -fn update_anthropic_usage(usage: &mut Usage, new: &Usage) { - if let Some(input_tokens) = new.input_tokens { - usage.input_tokens = Some(input_tokens); - } - if let Some(output_tokens) = new.output_tokens { - usage.output_tokens = Some(output_tokens); - } - if let Some(cache_creation_input_tokens) = new.cache_creation_input_tokens { - usage.cache_creation_input_tokens = Some(cache_creation_input_tokens); - } - if let Some(cache_read_input_tokens) = new.cache_read_input_tokens { - usage.cache_read_input_tokens = Some(cache_read_input_tokens); - } -} - -fn convert_anthropic_usage(usage: &Usage) -> language_model::TokenUsage { - language_model::TokenUsage { - input_tokens: usage.input_tokens.unwrap_or(0), - output_tokens: usage.output_tokens.unwrap_or(0), - cache_creation_input_tokens: usage.cache_creation_input_tokens.unwrap_or(0), - cache_read_input_tokens: usage.cache_read_input_tokens.unwrap_or(0), - } -} diff --git a/crates/language_models/src/provider/google.rs b/crates/language_models/src/provider/google.rs index c5a5affcd3d9e8c34f6306f86cb5348f86397892..a80d36c414e126cc08c641c44ce5ea2ee732a5c6 100644 --- a/crates/language_models/src/provider/google.rs +++ b/crates/language_models/src/provider/google.rs @@ -1,6 +1,5 @@ use anyhow::{Context as _, Result, anyhow}; use collections::BTreeMap; -use credentials_provider::CredentialsProvider; use futures::{FutureExt, Stream, StreamExt, future, future::BoxFuture}; use google_ai::{ FunctionDeclaration, GenerateContentResponse, GoogleModelMode, Part, SystemInstruction, @@ -33,7 +32,6 @@ use ui_input::InputField; use util::ResultExt; use zed_env_vars::EnvVar; -use crate::api_key::ApiKey; use crate::api_key::ApiKeyState; use crate::ui::{ConfiguredApiCard, InstructionListItem}; @@ -128,22 +126,6 @@ impl GoogleLanguageModelProvider { }) } - pub fn api_key_for_gemini_cli(cx: &mut App) -> Task> { - if let Some(key) = API_KEY_ENV_VAR.value.clone() { - return Task::ready(Ok(key)); - } - let credentials_provider = ::global(cx); - let api_url = Self::api_url(cx).to_string(); - cx.spawn(async move |cx| { - Ok( - ApiKey::load_from_system_keychain(&api_url, credentials_provider.as_ref(), cx) - .await? - .key() - .to_string(), - ) - }) - } - fn settings(cx: &App) -> &GoogleSettings { &crate::AllLanguageModelSettings::get_global(cx).google }