add config.gateway.unstable_disable_feedback_target_validation (tensorzero#2944)

* added a config.gateway.unstable_disable_feedback_target_validation flag, implementation, and tests

* slowed down polling for feedback get_function name

* changed margin for jitter, added comment

* fixed bindings

* fixed type in Feedback table story

Author: virajmehta

clients/rust/src/lib.rs

@@ -410,8 +410,7 @@ impl Client {
                         .await
                         .map_err(err_to_http)
                 })
-                .await?
-                .0)
+                .await?)
             }
         }
     }

internal/tensorzero-node/lib/bindings/GatewayConfig.ts

@@ -11,4 +11,5 @@ export type GatewayConfig = {
   export: ExportConfig;
   base_path: string | null;
   unstable_error_json: boolean;
+  unstable_disable_feedback_target_validation: boolean;
 };

tensorzero-core/src/config_parser/gateway.rs

@@ -25,6 +25,9 @@ pub struct UninitializedGatewayConfig {
     // If set, all of the HTTP endpoints will have this path prepended.
     // E.g. a base path of `/custom/prefix` will cause the inference endpoint to become `/custom/prefix/inference`.
     pub base_path: Option<String>,
+    // If set to `true`, disables validation on feedback queries (read from ClickHouse to check that the target is valid)
+    #[serde(default)]
+    pub unstable_disable_feedback_target_validation: bool,
     /// If enabled, adds an 'error_json' field alongside the human-readable 'error' field
     /// in HTTP error responses. This contains a JSON-serialized version of the error.
     /// While 'error_json' will always be valid JSON when present, the exact contents is unstable,
@@ -63,6 +66,8 @@ impl UninitializedGatewayConfig {
             export: self.export,
             base_path: self.base_path,
             unstable_error_json: self.unstable_error_json,
+            unstable_disable_feedback_target_validation: self
+                .unstable_disable_feedback_target_validation,
         })
     }
 }
@@ -80,6 +85,7 @@ pub struct GatewayConfig {
     // E.g. a base path of `/custom/prefix` will cause the inference endpoint to become `/custom/prefix/inference`.
     pub base_path: Option<String>,
     pub unstable_error_json: bool,
+    pub unstable_disable_feedback_target_validation: bool,
 }
 
 fn serialize_optional_socket_addr<S>(

tensorzero-core/src/endpoints/feedback.rs

@@ -32,11 +32,15 @@ use super::validate_tags;
 ///
 /// This is the amount of time we want to wait after the target was supposed to have been written
 /// before we decide that the target was actually not written because we can't find it in the database.
-const FEEDBACK_COOLDOWN_PERIOD: Duration = Duration::from_secs(5);
+/// This should really be read at 5000ms but since there might be some jitter we want to make sure there's
+/// a read at ~5s
+const FEEDBACK_COOLDOWN_PERIOD: Duration = Duration::from_millis(6000);
 /// Since we can't be sure that an inference actually completed when the ID says it was
 /// (the ID is generated at the start of the inference), we wait a minimum amount of time
 /// before we decide that the target was actually not written because we can't find it in the database.
-const FEEDBACK_MINIMUM_WAIT_TIME: Duration = Duration::from_millis(1200);
+const FEEDBACK_MINIMUM_WAIT_TIME: Duration = Duration::from_millis(1000);
+/// We also poll in the intermediate time so that we can return as soon as we find a target entry.
+const FEEDBACK_TARGET_POLL_INTERVAL: Duration = Duration::from_millis(2000);
 
 /// The expected payload is a JSON object with the following fields:
 #[derive(Debug, Default, Serialize, Deserialize)]
@@ -94,7 +98,7 @@ pub async fn feedback_handler(
     State(app_state): AppState,
     StructuredJson(params): StructuredJson<Params>,
 ) -> Result<Json<FeedbackResponse>, Error> {
-    feedback(app_state, params).await
+    Ok(Json(feedback(app_state, params).await?))
 }
 
 // Helper function to avoid requiring axum types in the client
@@ -105,7 +109,7 @@ pub async fn feedback(
         ..
     }: AppStateData,
     params: Params,
-) -> Result<Json<FeedbackResponse>, Error> {
+) -> Result<FeedbackResponse, Error> {
     validate_tags(&params.tags, params.internal)?;
     validate_feedback_specific_tags(&params.tags)?;
     // Get the metric config or return an error if it doesn't exist
@@ -139,6 +143,7 @@ pub async fn feedback(
                 feedback_metadata.level,
                 feedback_id,
                 dryrun,
+                config.gateway.unstable_disable_feedback_target_validation,
             )
             .await?;
         }
@@ -161,6 +166,7 @@ pub async fn feedback(
                 feedback_metadata.target_id,
                 feedback_id,
                 dryrun,
+                config.gateway.unstable_disable_feedback_target_validation,
             )
             .await?;
         }
@@ -172,12 +178,13 @@ pub async fn feedback(
                 feedback_metadata.target_id,
                 feedback_id,
                 dryrun,
+                config.gateway.unstable_disable_feedback_target_validation,
             )
             .await?;
         }
     }
 
-    Ok(Json(FeedbackResponse { feedback_id }))
+    Ok(FeedbackResponse { feedback_id })
 }
 
 #[derive(Debug)]
@@ -248,10 +255,13 @@ async fn write_comment(
     level: &MetricConfigLevel,
     feedback_id: Uuid,
     dryrun: bool,
+    disable_validation: bool,
 ) -> Result<(), Error> {
     let Params { value, tags, .. } = params;
     // Verify that the function name exists.
-    let _ = throttled_get_function_name(&connection_info, level, &target_id).await?;
+    if !disable_validation {
+        let _ = throttled_get_function_name(&connection_info, level, &target_id).await?;
+    }
     let value = value.as_str().ok_or_else(|| ErrorDetails::InvalidRequest {
         message: "Feedback value for a comment must be a string".to_string(),
     })?;
@@ -318,6 +328,7 @@ async fn write_float(
     target_id: Uuid,
     feedback_id: Uuid,
     dryrun: bool,
+    disable_validation: bool,
 ) -> Result<(), Error> {
     let Params {
         metric_name,
@@ -327,8 +338,11 @@ async fn write_float(
     } = params;
     let metric_config: &crate::config_parser::MetricConfig =
         config.get_metric_or_err(metric_name)?;
-    // Verify that the function name exists.
-    let _ = throttled_get_function_name(&connection_info, &metric_config.level, &target_id).await?;
+    if !disable_validation {
+        // Verify that the function name exists.
+        let _ =
+            throttled_get_function_name(&connection_info, &metric_config.level, &target_id).await?;
+    }
 
     let value = value.as_f64().ok_or_else(|| {
         Error::new(ErrorDetails::InvalidRequest {
@@ -353,6 +367,7 @@ async fn write_boolean(
     target_id: Uuid,
     feedback_id: Uuid,
     dryrun: bool,
+    disable_validation: bool,
 ) -> Result<(), Error> {
     let Params {
         metric_name,
@@ -361,8 +376,11 @@ async fn write_boolean(
         ..
     } = params;
     let metric_config = config.get_metric_or_err(metric_name)?;
-    // Verify that the function name exists.
-    let _ = throttled_get_function_name(&connection_info, &metric_config.level, &target_id).await?;
+    if !disable_validation {
+        // Verify that the function name exists.
+        let _ =
+            throttled_get_function_name(&connection_info, &metric_config.level, &target_id).await?;
+    }
     let value = value.as_bool().ok_or_else(|| {
         Error::new(ErrorDetails::InvalidRequest {
             message: format!("Feedback value for metric `{metric_name}` must be a boolean"),
@@ -420,7 +438,7 @@ async fn throttled_get_function_name(
                 }
             }
         }
-        tokio::time::sleep(Duration::from_millis(500)).await;
+        tokio::time::sleep(FEEDBACK_TARGET_POLL_INTERVAL).await;
     }
 }
 

tensorzero-core/src/gateway_util.rs

@@ -278,6 +278,7 @@ mod tests {
             export: Default::default(),
             base_path: None,
             unstable_error_json: false,
+            unstable_disable_feedback_target_validation: false,
         };
 
         let config = Box::leak(Box::new(Config {
@@ -332,6 +333,7 @@ mod tests {
             export: Default::default(),
             base_path: None,
             unstable_error_json: false,
+            unstable_disable_feedback_target_validation: false,
         };
 
         let config = Box::leak(Box::new(Config {
@@ -356,6 +358,7 @@ mod tests {
             export: Default::default(),
             base_path: None,
             unstable_error_json: false,
+            unstable_disable_feedback_target_validation: false,
         };
         let config = Box::leak(Box::new(Config {
             gateway: gateway_config,
@@ -382,6 +385,7 @@ mod tests {
             export: Default::default(),
             base_path: None,
             unstable_error_json: false,
+            unstable_disable_feedback_target_validation: false,
         };
         let config = Config {
             gateway: gateway_config,

tensorzero-core/tests/e2e/feedback.rs

@@ -2,6 +2,11 @@ use reqwest::{Client, StatusCode};
 use serde_json::{json, Value};
 use tensorzero_core::{
     clickhouse::test_helpers::{select_feedback_clickhouse, select_feedback_tags_clickhouse},
+    config_parser::{
+        Config, MetricConfig, MetricConfigLevel, MetricConfigOptimize, MetricConfigType,
+    },
+    endpoints::feedback::{feedback, Params},
+    gateway_util::AppStateData,
     inference::types::{ContentBlockChatOutput, JsonInferenceOutput, Role, Text, TextKind},
 };
 use tokio::time::{sleep, Duration};
@@ -170,6 +175,39 @@ async fn e2e_test_comment_feedback_with_payload(inference_payload: serde_json::V
     assert_eq!(retrieved_value, "bad job!");
 }
 
+#[tokio::test(flavor = "multi_thread")]
+async fn e2e_test_comment_feedback_validation_disabled() {
+    let mut config = Config::default();
+    let clickhouse = get_clickhouse().await;
+    config.gateway.unstable_disable_feedback_target_validation = true;
+    let state = AppStateData::new_with_clickhouse_and_http_client(
+        config.into(),
+        clickhouse.clone(),
+        reqwest::Client::new(),
+    );
+    let inference_id = Uuid::now_v7();
+    let params = Params {
+        inference_id: Some(inference_id),
+        metric_name: "comment".to_string(),
+        value: json!("foo bar"),
+        ..Default::default()
+    };
+    let val = feedback(state, params).await.unwrap();
+    tokio::time::sleep(Duration::from_millis(500)).await;
+
+    // Check that this was correctly written to ClickHouse
+    let query = format!(
+        "SELECT * FROM CommentFeedback WHERE target_id='{inference_id}' FORMAT JsonEachRow"
+    );
+    let response = clickhouse
+        .run_query_synchronous_no_params(query)
+        .await
+        .unwrap();
+    let result: Value = serde_json::from_str(&response.response).unwrap();
+    let clickhouse_feedback_id = Uuid::parse_str(result["id"].as_str().unwrap()).unwrap();
+    assert_eq!(val.feedback_id, clickhouse_feedback_id);
+}
+
 #[tokio::test]
 async fn e2e_test_demonstration_feedback_normal_function() {
     e2e_test_demonstration_feedback_with_payload(serde_json::json!({
@@ -1160,6 +1198,47 @@ async fn e2e_test_float_feedback_with_payload(inference_payload: serde_json::Val
     assert_eq!(metric_name, "brevity_score");
 }
 
+#[tokio::test(flavor = "multi_thread")]
+async fn e2e_test_float_feedback_validation_disabled() {
+    let mut config = Config::default();
+    let metric_config = MetricConfig {
+        r#type: MetricConfigType::Float,
+        optimize: MetricConfigOptimize::Max,
+        level: MetricConfigLevel::Inference,
+    };
+    config
+        .metrics
+        .insert("user_score".to_string(), metric_config);
+    let clickhouse = get_clickhouse().await;
+    config.gateway.unstable_disable_feedback_target_validation = true;
+    let state = AppStateData::new_with_clickhouse_and_http_client(
+        config.into(),
+        clickhouse.clone(),
+        reqwest::Client::new(),
+    );
+    let inference_id = Uuid::now_v7();
+    let params = Params {
+        inference_id: Some(inference_id),
+        metric_name: "user_score".to_string(),
+        value: json!(3.1),
+        ..Default::default()
+    };
+    let val = feedback(state, params).await.unwrap();
+    tokio::time::sleep(Duration::from_millis(500)).await;
+
+    // Check that this was correctly written to ClickHouse
+    let query = format!(
+        "SELECT * FROM FloatMetricFeedback WHERE target_id='{inference_id}' FORMAT JsonEachRow"
+    );
+    let response = clickhouse
+        .run_query_synchronous_no_params(query)
+        .await
+        .unwrap();
+    let result: Value = serde_json::from_str(&response.response).unwrap();
+    let clickhouse_feedback_id = Uuid::parse_str(result["id"].as_str().unwrap()).unwrap();
+    assert_eq!(val.feedback_id, clickhouse_feedback_id);
+}
+
 #[tokio::test]
 async fn e2e_test_boolean_feedback_normal_function() {
     e2e_test_boolean_feedback_with_payload(serde_json::json!({
@@ -1353,6 +1432,47 @@ async fn e2e_test_boolean_feedback_with_payload(inference_payload: serde_json::V
     assert_eq!(metric_name, "goal_achieved");
 }
 
+#[tokio::test(flavor = "multi_thread")]
+async fn e2e_test_boolean_feedback_validation_disabled() {
+    let mut config = Config::default();
+    let metric_config = MetricConfig {
+        r#type: MetricConfigType::Boolean,
+        optimize: MetricConfigOptimize::Max,
+        level: MetricConfigLevel::Inference,
+    };
+    config
+        .metrics
+        .insert("task_success".to_string(), metric_config);
+    let clickhouse = get_clickhouse().await;
+    config.gateway.unstable_disable_feedback_target_validation = true;
+    let state = AppStateData::new_with_clickhouse_and_http_client(
+        config.into(),
+        clickhouse.clone(),
+        reqwest::Client::new(),
+    );
+    let inference_id = Uuid::now_v7();
+    let params = Params {
+        inference_id: Some(inference_id),
+        metric_name: "task_success".to_string(),
+        value: json!(true),
+        ..Default::default()
+    };
+    let val = feedback(state, params).await.unwrap();
+    tokio::time::sleep(Duration::from_millis(500)).await;
+
+    // Check that this was correctly written to ClickHouse
+    let query = format!(
+        "SELECT * FROM BooleanMetricFeedback WHERE target_id='{inference_id}' FORMAT JsonEachRow"
+    );
+    let response = clickhouse
+        .run_query_synchronous_no_params(query)
+        .await
+        .unwrap();
+    let result: Value = serde_json::from_str(&response.response).unwrap();
+    let clickhouse_feedback_id = Uuid::parse_str(result["id"].as_str().unwrap()).unwrap();
+    assert_eq!(val.feedback_id, clickhouse_feedback_id);
+}
+
 #[tokio::test(flavor = "multi_thread")]
 #[traced_test]
 async fn test_fast_inference_then_feedback() {

ui/app/components/feedback/FeedbackTable.stories.tsx

@@ -31,6 +31,7 @@ const config: Config = {
     bind_address: "localhost:8080",
     base_path: "/",
     unstable_error_json: false,
+    unstable_disable_feedback_target_validation: false,
   },
   object_store_info: { kind: { type: "disabled" } },
   provider_types: {