zeta2 cli: Include section ranges in new full output format (#39203)

Agus Zubiaga and Oleksiy Syvokon created 1 month ago

Release Notes:

- N/A

---------

Co-authored-by: Oleksiy Syvokon <[e-mail]>

Change summary

Cargo.lock                                          | 1 +
crates/cloud_zeta2_prompt/Cargo.toml                | 1 +
crates/cloud_zeta2_prompt/src/cloud_zeta2_prompt.rs | 3 ++-
crates/zeta_cli/src/main.rs                         | 8 +++++---
4 files changed, 9 insertions(+), 4 deletions(-)

Detailed changes

Cargo.lock

@@ -3217,6 +3217,7 @@ dependencies = [
  "indoc",
  "ordered-float 2.10.1",
  "rustc-hash 2.1.1",
+ "serde",
  "strum 0.27.1",
  "workspace-hack",
 ]

crates/cloud_zeta2_prompt/Cargo.toml

@@ -17,5 +17,6 @@ cloud_llm_client.workspace = true
 indoc.workspace = true
 ordered-float.workspace = true
 rustc-hash.workspace = true
+serde.workspace = true
 strum.workspace = true
 workspace-hack.workspace = true

crates/cloud_zeta2_prompt/src/cloud_zeta2_prompt.rs

@@ -5,6 +5,7 @@ use cloud_llm_client::predict_edits_v3::{self, Event, PromptFormat, ReferencedDe
 use indoc::indoc;
 use ordered_float::OrderedFloat;
 use rustc_hash::{FxHashMap, FxHashSet};
+use serde::Serialize;
 use std::fmt::Write;
 use std::sync::Arc;
 use std::{cmp::Reverse, collections::BinaryHeap, ops::Range, path::Path};
@@ -75,7 +76,7 @@ pub enum DeclarationStyle {
     Declaration,
 }
 
-#[derive(Clone, Debug)]
+#[derive(Clone, Debug, Serialize)]
 pub struct SectionLabels {
     pub excerpt_index: usize,
     pub section_ranges: Vec<(Arc<Path>, Range<usize>)>,

crates/zeta_cli/src/main.rs

@@ -122,7 +122,7 @@ enum OutputFormat {
     #[default]
     Prompt,
     Request,
-    Both,
+    Full,
 }
 
 #[derive(Debug, Clone)]
@@ -299,15 +299,17 @@ async fn get_context(
                         .await?;
 
                     let planned_prompt = cloud_zeta2_prompt::PlannedPrompt::populate(&request)?;
-                    let prompt_string = planned_prompt.to_prompt_string()?.0;
+                    let (prompt_string, section_labels) = planned_prompt.to_prompt_string()?;
+
                     match zeta2_args.output_format {
                         OutputFormat::Prompt => anyhow::Ok(prompt_string),
                         OutputFormat::Request => {
                             anyhow::Ok(serde_json::to_string_pretty(&request)?)
                         }
-                        OutputFormat::Both => anyhow::Ok(serde_json::to_string_pretty(&json!({
+                        OutputFormat::Full => anyhow::Ok(serde_json::to_string_pretty(&json!({
                             "request": request,
                             "prompt": prompt_string,
+                            "section_labels": section_labels,
                         }))?),
                     }
                 })