From 65250fe08d29c27d7414cdea5201550f720a7307 Mon Sep 17 00:00:00 2001 From: Michael Sloan Date: Tue, 29 Jul 2025 11:28:18 -0600 Subject: [PATCH] cloud provider: Use `CompletionEvent` type from `zed_llm_client` (#35285) Release Notes: - N/A --- crates/language_models/src/provider/cloud.rs | 31 ++++++++------------ 1 file changed, 12 insertions(+), 19 deletions(-) diff --git a/crates/language_models/src/provider/cloud.rs b/crates/language_models/src/provider/cloud.rs index 09a2ac6e0a..1e6e7b96d0 100644 --- a/crates/language_models/src/provider/cloud.rs +++ b/crates/language_models/src/provider/cloud.rs @@ -35,8 +35,8 @@ use ui::{TintColor, prelude::*}; use util::{ResultExt as _, maybe}; use zed_llm_client::{ CLIENT_SUPPORTS_STATUS_MESSAGES_HEADER_NAME, CURRENT_PLAN_HEADER_NAME, CompletionBody, - CompletionRequestStatus, CountTokensBody, CountTokensResponse, EXPIRED_LLM_TOKEN_HEADER_NAME, - ListModelsResponse, MODEL_REQUESTS_RESOURCE_HEADER_VALUE, + CompletionEvent, CompletionRequestStatus, CountTokensBody, CountTokensResponse, + EXPIRED_LLM_TOKEN_HEADER_NAME, ListModelsResponse, MODEL_REQUESTS_RESOURCE_HEADER_VALUE, SERVER_SUPPORTS_STATUS_MESSAGES_HEADER_NAME, SUBSCRIPTION_LIMIT_RESOURCE_HEADER_NAME, TOOL_USE_LIMIT_REACHED_HEADER_NAME, ZED_VERSION_HEADER_NAME, }; @@ -1040,15 +1040,8 @@ impl LanguageModel for CloudLanguageModel { } } -#[derive(Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum CloudCompletionEvent { - Status(CompletionRequestStatus), - Event(T), -} - fn map_cloud_completion_events( - stream: Pin>> + Send>>, + stream: Pin>> + Send>>, mut map_callback: F, ) -> BoxStream<'static, Result> where @@ -1063,10 +1056,10 @@ where Err(error) => { vec![Err(LanguageModelCompletionError::from(error))] } - Ok(CloudCompletionEvent::Status(event)) => { + Ok(CompletionEvent::Status(event)) => { vec![Ok(LanguageModelCompletionEvent::StatusUpdate(event))] } - Ok(CloudCompletionEvent::Event(event)) => map_callback(event), + Ok(CompletionEvent::Event(event)) => map_callback(event), }) }) .boxed() @@ -1074,9 +1067,9 @@ where fn usage_updated_event( usage: Option, -) -> impl Stream>> { +) -> impl Stream>> { futures::stream::iter(usage.map(|usage| { - Ok(CloudCompletionEvent::Status( + Ok(CompletionEvent::Status( CompletionRequestStatus::UsageUpdated { amount: usage.amount as usize, limit: usage.limit, @@ -1087,9 +1080,9 @@ fn usage_updated_event( fn tool_use_limit_reached_event( tool_use_limit_reached: bool, -) -> impl Stream>> { +) -> impl Stream>> { futures::stream::iter(tool_use_limit_reached.then(|| { - Ok(CloudCompletionEvent::Status( + Ok(CompletionEvent::Status( CompletionRequestStatus::ToolUseLimitReached, )) })) @@ -1098,7 +1091,7 @@ fn tool_use_limit_reached_event( fn response_lines( response: Response, includes_status_messages: bool, -) -> impl Stream>> { +) -> impl Stream>> { futures::stream::try_unfold( (String::new(), BufReader::new(response.into_body())), move |(mut line, mut body)| async move { @@ -1106,9 +1099,9 @@ fn response_lines( Ok(0) => Ok(None), Ok(_) => { let event = if includes_status_messages { - serde_json::from_str::>(&line)? + serde_json::from_str::>(&line)? } else { - CloudCompletionEvent::Event(serde_json::from_str::(&line)?) + CompletionEvent::Event(serde_json::from_str::(&line)?) }; line.clear();