Fix a missing newline in the zeta prompt (#46677)

Max Brunsfeld created 1 week ago

Release Notes:

- N/A

Change summary

crates/edit_prediction/src/edit_prediction.rs | 11 ++++++++---
crates/edit_prediction/src/zeta2.rs           |  2 ++
crates/edit_prediction_cli/src/predict.rs     |  9 ++++++++-
crates/zeta_prompt/src/zeta_prompt.rs         |  4 +++-
4 files changed, 21 insertions(+), 5 deletions(-)

Detailed changes

crates/edit_prediction/src/edit_prediction.rs

@@ -1888,14 +1888,19 @@ impl EditPredictionStore {
     async fn send_raw_llm_request(
         request: RawCompletionRequest,
         client: Arc<Client>,
+        custom_url: Option<Arc<Url>>,
         llm_token: LlmApiToken,
         app_version: Version,
         #[cfg(feature = "cli-support")] eval_cache: Option<Arc<dyn EvalCache>>,
         #[cfg(feature = "cli-support")] eval_cache_kind: EvalCacheEntryKind,
     ) -> Result<(RawCompletionResponse, Option<EditPredictionUsage>)> {
-        let url = client
-            .http_client()
-            .build_zed_llm_url("/predict_edits/raw", &[])?;
+        let url = if let Some(custom_url) = custom_url {
+            custom_url.as_ref().clone()
+        } else {
+            client
+                .http_client()
+                .build_zed_llm_url("/predict_edits/raw", &[])?
+        };
 
         #[cfg(feature = "cli-support")]
         let cache_key = if let Some(cache) = eval_cache {

crates/edit_prediction/src/zeta2.rs

@@ -35,6 +35,7 @@ pub fn request_prediction_with_zeta2(
     cx: &mut Context<EditPredictionStore>,
 ) -> Task<Result<Option<EditPredictionResult>>> {
     let buffer_snapshotted_at = Instant::now();
+    let url = store.custom_predict_edits_url.clone();
 
     let Some(excerpt_path) = snapshot
         .file()
@@ -88,6 +89,7 @@ pub fn request_prediction_with_zeta2(
             let response = EditPredictionStore::send_raw_llm_request(
                 request,
                 client,
+                url,
                 llm_token,
                 app_version,
                 #[cfg(feature = "cli-support")]

crates/edit_prediction_cli/src/predict.rs

@@ -1,7 +1,7 @@
 use crate::{
     PredictionProvider, PromptFormat,
     anthropic_client::AnthropicClient,
-    example::{Example, ExamplePrediction},
+    example::{Example, ExamplePrediction, ExamplePrompt},
     format_prompt::{TeacherPrompt, run_format_prompt},
     headless::EpAppState,
     load_project::run_load_project,
@@ -123,6 +123,13 @@ pub async fn run_prediction(
 
                         if let Some(prompt) = request.prompt {
                             fs::write(run_dir.join("prediction_prompt.md"), &prompt)?;
+                            if provider == PredictionProvider::Zeta2 {
+                                updated_example.prompt.get_or_insert(ExamplePrompt {
+                                    input: prompt,
+                                    expected_output: String::new(),
+                                    format: PromptFormat::Zeta2,
+                                });
+                            }
                         }
                     }
                     DebugEvent::EditPredictionFinished(request) => {

crates/zeta_prompt/src/zeta_prompt.rs

@@ -109,6 +109,9 @@ fn write_cursor_excerpt_section(prompt: &mut String, input: &ZetaPromptInput) {
 
     prompt.push_str("<|fim_suffix|>\n");
     prompt.push_str(&input.cursor_excerpt[input.editable_range_in_excerpt.end..]);
+    if !prompt.ends_with('\n') {
+        prompt.push('\n');
+    }
 
     prompt.push_str("<|fim_middle|>current\n");
     prompt.push_str(
@@ -119,7 +122,6 @@ fn write_cursor_excerpt_section(prompt: &mut String, input: &ZetaPromptInput) {
     prompt.push_str(
         &input.cursor_excerpt[input.cursor_offset_in_excerpt..input.editable_range_in_excerpt.end],
     );
-
     if !prompt.ends_with('\n') {
         prompt.push('\n');
     }