ep: Fix raw request shape (#46711)

Agus Zubiaga created 1 week ago

Release Notes:

- N/A

Change summary

crates/cloud_llm_client/src/predict_edits_v3.rs | 4 +++-
crates/edit_prediction/src/zeta2.rs             | 2 +-
2 files changed, 4 insertions(+), 2 deletions(-)

Detailed changes

crates/cloud_llm_client/src/predict_edits_v3.rs

@@ -219,7 +219,9 @@ impl Sub for Line {
 pub struct RawCompletionRequest {
     pub model: String,
     pub prompt: String,
-    pub max_tokens: u32,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub max_tokens: Option<u32>,
+    #[serde(skip_serializing_if = "Option::is_none")]
     pub temperature: Option<f32>,
     pub stop: Vec<Cow<'static, str>>,
 }

crates/edit_prediction/src/zeta2.rs

@@ -81,7 +81,7 @@ pub fn request_prediction_with_zeta2(
                 prompt,
                 temperature: None,
                 stop: vec![],
-                max_tokens: 1024,
+                max_tokens: Some(2048),
             };
 
             log::trace!("Sending edit prediction request");