Send some traffic to zeta2 for testing (#47710)

Max Brunsfeld created 1 week ago

Release Notes:

- N/A

Change summary

crates/cloud_llm_client/src/cloud_llm_client.rs |  1 
crates/edit_prediction/src/capture_example.rs   |  5 ++++
crates/edit_prediction/src/edit_prediction.rs   | 20 +++++++++++++++++-
3 files changed, 24 insertions(+), 2 deletions(-)

Detailed changes

crates/cloud_llm_client/src/cloud_llm_client.rs

@@ -142,6 +142,7 @@ pub struct PredictEditsBody {
 
 #[derive(Default, Debug, Clone, Copy, Serialize, Deserialize)]
 pub enum PredictEditsRequestTrigger {
+    Testing,
     Diagnostics,
     Cli,
     #[default]

crates/edit_prediction/src/capture_example.rs

@@ -18,6 +18,7 @@ use text::{BufferSnapshot as TextBufferSnapshot, Point, ToOffset as _};
 
 pub(crate) const DEFAULT_EXAMPLE_CAPTURE_RATE_PER_10K_PREDICTIONS: u16 = 10;
 pub(crate) const DEFAULT_STAFF_EXAMPLE_CAPTURE_RATE_PER_10K_PREDICTIONS: u16 = 100;
+pub(crate) const ZETA2_TESTING_RATE_PER_10K_PREDICTION: u16 = 500;
 
 pub fn capture_example(
     project: Entity<Project>,
@@ -317,6 +318,10 @@ pub(crate) fn should_sample_edit_prediction_example_capture(cx: &App) -> bool {
         && rand::random::<u16>() % 10_000 < capture_rate
 }
 
+pub(crate) fn should_send_testing_zeta2_request() -> bool {
+    rand::random::<u16>() % 10_000 < ZETA2_TESTING_RATE_PER_10K_PREDICTION
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

crates/edit_prediction/src/edit_prediction.rs

@@ -71,7 +71,9 @@ pub mod zeta2;
 #[cfg(test)]
 mod edit_prediction_tests;
 
-use crate::capture_example::should_sample_edit_prediction_example_capture;
+use crate::capture_example::{
+    should_sample_edit_prediction_example_capture, should_send_testing_zeta2_request,
+};
 use crate::license_detection::LicenseDetectionWatcher;
 use crate::mercury::Mercury;
 use crate::onboarding_modal::ZedPredictModal;
@@ -170,6 +172,7 @@ pub enum EditPredictionModel {
     Mercury,
 }
 
+#[derive(Clone)]
 pub struct EditPredictionModelInput {
     project: Entity<Project>,
     buffer: Entity<Buffer>,
@@ -1763,7 +1766,20 @@ impl EditPredictionStore {
             }
         }
         let task = match self.edit_prediction_model {
-            EditPredictionModel::Zeta1 => zeta1::request_prediction_with_zeta1(self, inputs, cx),
+            EditPredictionModel::Zeta1 => {
+                if should_send_testing_zeta2_request() {
+                    let mut zeta2_inputs = inputs.clone();
+                    zeta2_inputs.trigger = PredictEditsRequestTrigger::Testing;
+                    zeta2::request_prediction_with_zeta2(
+                        self,
+                        zeta2_inputs,
+                        Default::default(),
+                        cx,
+                    )
+                    .detach();
+                }
+                zeta1::request_prediction_with_zeta1(self, inputs, cx)
+            }
             EditPredictionModel::Zeta2 { version } => {
                 zeta2::request_prediction_with_zeta2(self, inputs, version, cx)
             }