From f89bb2f0d223d029ea6180faf8257843c362af3e Mon Sep 17 00:00:00 2001 From: Max Brunsfeld Date: Thu, 6 Nov 2025 23:34:05 -0800 Subject: [PATCH] Small zeta cli fixes (#42170) * Fix a panic that happened because we lost the `ContextRetrievalStarted` debug message, so we didn't assign `t0`. * Write the edit prediction response log file as a markdown file containing the text, not a JSON file. We mostly always want the text content. Release Notes: - N/A --- crates/zeta2/src/zeta2.rs | 87 +++++++++++++++++++--------------- crates/zeta_cli/src/example.rs | 5 +- crates/zeta_cli/src/predict.rs | 14 +++--- 3 files changed, 57 insertions(+), 49 deletions(-) diff --git a/crates/zeta2/src/zeta2.rs b/crates/zeta2/src/zeta2.rs index d7a645794fd8a21f8b70b3786b435a23a9babe63..7dd5f8d27dcef5455a14854faa8426d4badac4b8 100644 --- a/crates/zeta2/src/zeta2.rs +++ b/crates/zeta2/src/zeta2.rs @@ -919,40 +919,10 @@ impl Zeta { .ok(); } - let (mut res, usage) = response?; - + let (res, usage) = response?; let request_id = EditPredictionId(Uuid::from_str(&res.id)?); - - let Some(choice) = res.choices.pop() else { - return Ok((None, usage)); - }; - - let output_text = match choice.message { - open_ai::RequestMessage::Assistant { - content: Some(open_ai::MessageContent::Plain(content)), - .. - } => content, - open_ai::RequestMessage::Assistant { - content: Some(open_ai::MessageContent::Multipart(mut content)), - .. - } => { - if content.is_empty() { - log::error!("No output from Baseten completion response"); - return Ok((None, usage)); - } - - match content.remove(0) { - open_ai::MessagePart::Text { text } => text, - open_ai::MessagePart::Image { .. } => { - log::error!("Expected text, got an image"); - return Ok((None, usage)); - } - } - } - _ => { - log::error!("Invalid response message: {:?}", choice.message); - return Ok((None, usage)); - } + let Some(output_text) = text_from_response(res) else { + return Ok((None, usage)) }; let (edited_buffer_snapshot, edits) = @@ -1216,6 +1186,10 @@ impl Zeta { return Task::ready(Ok(())); }; + let app_version = AppVersion::global(cx); + let client = self.client.clone(); + let llm_token = self.llm_token.clone(); + let debug_tx = self.debug_tx.clone(); let current_file_path: Arc = snapshot .file() .map(|f| f.full_path(cx).into()) @@ -1240,10 +1214,17 @@ impl Zeta { } }; - let app_version = AppVersion::global(cx); - let client = self.client.clone(); - let llm_token = self.llm_token.clone(); - let debug_tx = self.debug_tx.clone(); + if let Some(debug_tx) = &debug_tx { + debug_tx + .unbounded_send(ZetaDebugInfo::ContextRetrievalStarted( + ZetaContextRetrievalStartedDebugInfo { + project: project.clone(), + timestamp: Instant::now(), + search_prompt: prompt.clone(), + }, + )) + .ok(); + } let (tool_schema, tool_description) = &*cloud_zeta2_prompt::retrieval_prompt::TOOL_SCHEMA; @@ -1503,6 +1484,38 @@ impl Zeta { } } +pub fn text_from_response(mut res: open_ai::Response) -> Option { + let choice = res.choices.pop()?; + let output_text = match choice.message { + open_ai::RequestMessage::Assistant { + content: Some(open_ai::MessageContent::Plain(content)), + .. + } => content, + open_ai::RequestMessage::Assistant { + content: Some(open_ai::MessageContent::Multipart(mut content)), + .. + } => { + if content.is_empty() { + log::error!("No output from Baseten completion response"); + return None; + } + + match content.remove(0) { + open_ai::MessagePart::Text { text } => text, + open_ai::MessagePart::Image { .. } => { + log::error!("Expected text, got an image"); + return None; + } + } + } + _ => { + log::error!("Invalid response message: {:?}", choice.message); + return None; + } + }; + Some(output_text) +} + #[derive(Error, Debug)] #[error( "You must update to Zed version {minimum_version} or higher to continue using edit predictions." diff --git a/crates/zeta_cli/src/example.rs b/crates/zeta_cli/src/example.rs index e3c988dd4b0046f117999ad61050c3d710f618a1..a470effa575f5e8ece3c59781dc09d9d1c5e822e 100644 --- a/crates/zeta_cli/src/example.rs +++ b/crates/zeta_cli/src/example.rs @@ -166,9 +166,7 @@ impl NamedExample { REVISION_FIELD => { named.example.revision = value.trim().to_string(); } - _ => { - eprintln!("Warning: Unrecognized field `{field}`"); - } + _ => {} } } } @@ -193,7 +191,6 @@ impl NamedExample { } else if title.eq_ignore_ascii_case(EXPECTED_CONTEXT_HEADING) { Section::ExpectedExcerpts } else { - eprintln!("Warning: Unrecognized section `{title:?}`"); Section::Other }; } diff --git a/crates/zeta_cli/src/predict.rs b/crates/zeta_cli/src/predict.rs index c94353d5b1dd02648265b70b4b584705f3769f96..1bc2411c825a2fa7147ff0c657804908b687d9ff 100644 --- a/crates/zeta_cli/src/predict.rs +++ b/crates/zeta_cli/src/predict.rs @@ -148,14 +148,6 @@ pub async fn zeta2_predict( &request.local_prompt.unwrap_or_default(), )?; - let response = request.response_rx.await?.0.map_err(|err| anyhow!(err))?; - prediction_finished_at = Some(Instant::now()); - - fs::write( - LOGS_DIR.join("prediction_response.json"), - &serde_json::to_string_pretty(&response).unwrap(), - )?; - for included_file in request.request.included_files { let insertions = vec![(request.request.cursor_point, CURSOR_MARKER)]; result @@ -177,6 +169,12 @@ pub async fn zeta2_predict( &mut excerpts_text, ); } + + let response = request.response_rx.await?.0.map_err(|err| anyhow!(err))?; + let response = zeta2::text_from_response(response).unwrap_or_default(); + prediction_finished_at = Some(Instant::now()); + fs::write(LOGS_DIR.join("prediction_response.md"), &response)?; + break; } }