Merge branch 'main' into main

Author: pycoder49

.config/nextest.toml

@@ -81,6 +81,8 @@ slow-timeout = { period = "60s", terminate-after = 2 }
 
 [[profile.default.overrides]]
 filter = 'binary(e2e) and test(providers::sglang::)'
+# The model we run on SGLang often fails to emit valid tool calls, so we need many retries
+retries = { backoff = "fixed", count = 8, delay = "10s", jitter = true }
 slow-timeout = { period = "60s", terminate-after = 2 }
 
 [[profile.default.overrides]]

evaluations/tests/common/mod.rs

@@ -7,7 +7,10 @@ use std::{
 use tensorzero::{
     ChatInferenceDatapoint, Client, ClientBuilder, ClientBuilderMode, JsonInferenceDatapoint,
 };
-use tensorzero_core::clickhouse::test_helpers::{get_clickhouse, CLICKHOUSE_URL};
+use tensorzero_core::clickhouse::{
+    test_helpers::{get_clickhouse, CLICKHOUSE_URL},
+    TableName,
+};
 use uuid::Uuid;
 
 /// Takes a chat fixture as a path to a JSONL file and writes the fixture to the dataset.
@@ -32,7 +35,7 @@ pub async fn write_chat_fixture_to_dataset(
     }
     let clickhouse = get_clickhouse().await;
     clickhouse
-        .write(&datapoints, "ChatInferenceDatapoint")
+        .write(&datapoints, TableName::ChatInferenceDatapoint)
         .await
         .unwrap();
 }
@@ -56,7 +59,7 @@ pub async fn write_json_fixture_to_dataset(
     }
     let clickhouse = get_clickhouse().await;
     clickhouse
-        .write(&datapoints, "JsonInferenceDatapoint")
+        .write(&datapoints, TableName::JsonInferenceDatapoint)
         .await
         .unwrap();
 }

tensorzero-core/src/cache.rs

@@ -1,7 +1,7 @@
 use std::collections::HashMap;
 use std::sync::Arc;
 
-use crate::clickhouse::ClickHouseConnectionInfo;
+use crate::clickhouse::{ClickHouseConnectionInfo, TableName};
 use crate::embeddings::{EmbeddingRequest, EmbeddingResponse};
 use crate::error::{Error, ErrorDetails};
 use crate::inference::types::file::serialize_with_file_data;
@@ -270,7 +270,7 @@ pub fn start_cache_write<T: Serialize + CacheOutput + Send + Sync + 'static>(
                         finish_reason,
                     },
                 }],
-                "ModelInferenceCache",
+                TableName::ModelInferenceCache,
             )
             .await
         {
@@ -337,7 +337,7 @@ pub fn start_cache_write_streaming(
                         finish_reason,
                     },
                 }],
-                "ModelInferenceCache",
+                TableName::ModelInferenceCache,
             )
             .await
     });

tensorzero-core/src/clickhouse/migration_manager/migrations/migration_0033.rs

@@ -1,7 +1,7 @@
 use super::check_table_exists;
 use crate::clickhouse::migration_manager::migration_trait::Migration;
 use crate::clickhouse::migration_manager::migrations::table_is_nonempty;
-use crate::clickhouse::ClickHouseConnectionInfo;
+use crate::clickhouse::{ClickHouseConnectionInfo, TableName};
 use crate::error::Error;
 use async_trait::async_trait;
 use uuid::Uuid;
@@ -59,7 +59,7 @@ impl Migration for Migration0033<'_> {
                 &[serde_json::json!({
                     "deployment_id": deployment_id,
                 })],
-                "DeploymentID",
+                TableName::DeploymentID,
             )
             .await?;
 

tensorzero-core/src/clickhouse/migration_manager/mod.rs

@@ -3,7 +3,7 @@ pub mod migrations;
 
 use std::time::{Duration, Instant};
 
-use crate::clickhouse::ClickHouseConnectionInfo;
+use crate::clickhouse::{ClickHouseConnectionInfo, TableName};
 use crate::endpoints::status::TENSORZERO_VERSION;
 use crate::error::{Error, ErrorDetails};
 use crate::serde_util::deserialize_u64;
@@ -147,7 +147,7 @@ async fn insert_migration_record(
                 execution_time_ms: execution_time.as_millis() as u64,
                 applied_at: None,
             }],
-            "TensorZeroMigration",
+            TableName::TensorZeroMigration,
         )
         .await?;
     Ok(())

tensorzero-core/src/clickhouse/mod.rs

@@ -51,6 +51,47 @@ pub fn make_clickhouse_http_client() -> Result<Client, Error> {
         })
 }
 
+/// Defines all of the ClickHouse tables that we write to from Rust
+/// This will be used to implement per-table ClickHouse write batching.
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
+pub enum TableName {
+    BatchModelInference,
+    BatchRequest,
+    ChatInference,
+    ChatInferenceDatapoint,
+    JsonInference,
+    JsonInferenceDatapoint,
+    ModelInference,
+    ModelInferenceCache,
+    DeploymentID,
+    TensorZeroMigration,
+    BooleanMetricFeedback,
+    FloatMetricFeedback,
+    DemonstrationFeedback,
+    CommentFeedback,
+}
+
+impl TableName {
+    pub fn as_str(self) -> &'static str {
+        match self {
+            TableName::BatchModelInference => "BatchModelInference",
+            TableName::BatchRequest => "BatchRequest",
+            TableName::ChatInference => "ChatInference",
+            TableName::ChatInferenceDatapoint => "ChatInferenceDatapoint",
+            TableName::JsonInference => "JsonInference",
+            TableName::JsonInferenceDatapoint => "JsonInferenceDatapoint",
+            TableName::ModelInference => "ModelInference",
+            TableName::ModelInferenceCache => "ModelInferenceCache",
+            TableName::DeploymentID => "DeploymentID",
+            TableName::TensorZeroMigration => "TensorZeroMigration",
+            TableName::BooleanMetricFeedback => "BooleanMetricFeedback",
+            TableName::FloatMetricFeedback => "FloatMetricFeedback",
+            TableName::DemonstrationFeedback => "DemonstrationFeedback",
+            TableName::CommentFeedback => "CommentFeedback",
+        }
+    }
+}
+
 impl ClickHouseConnectionInfo {
     /// Create a new ClickHouse connection info from a database URL.
     /// You should always use this function in production code or generic integration tests that
@@ -121,18 +162,18 @@ impl ClickHouseConnectionInfo {
     pub async fn write(
         &self,
         rows: &[impl Serialize + Send + Sync],
-        table: &str,
+        table: TableName,
     ) -> Result<(), Error> {
         match self {
             Self::Disabled => Ok(()),
             Self::Mock { mock_data, .. } => {
-                write_mock(rows, table, &mut mock_data.write().await).await
+                write_mock(rows, table.as_str(), &mut mock_data.write().await).await
             }
             Self::Production {
                 database_url,
                 client,
                 ..
-            } => write_production(database_url, client, rows, table).await,
+            } => write_production(database_url, client, rows, table.as_str()).await,
         }
     }
 

tensorzero-core/src/endpoints/batch_inference.rs

@@ -18,7 +18,7 @@ use super::inference::{
     InferenceIds, InferenceModels, InferenceParams, InferenceResponse, JsonInferenceResponse,
 };
 use crate::cache::{CacheEnabledMode, CacheOptions};
-use crate::clickhouse::ClickHouseConnectionInfo;
+use crate::clickhouse::{ClickHouseConnectionInfo, TableName};
 use crate::config_parser::Config;
 use crate::error::{Error, ErrorDetails};
 use crate::function::{sample_variant, FunctionConfig};
@@ -632,7 +632,7 @@ async fn write_start_batch_inference<'a>(
     }
 
     clickhouse_connection_info
-        .write(rows.as_slice(), "BatchModelInference")
+        .write(rows.as_slice(), TableName::BatchModelInference)
         .await?;
 
     let batch_request_insert = BatchRequestRow::new(UnparsedBatchRequestRow {
@@ -663,7 +663,7 @@ pub async fn write_batch_request_row(
     batch_request: &BatchRequestRow<'_>,
 ) -> Result<(), Error> {
     clickhouse_connection_info
-        .write(&[batch_request], "BatchRequest")
+        .write(&[batch_request], TableName::BatchRequest)
         .await
 }
 
@@ -761,7 +761,7 @@ async fn write_batch_request_status_update(
         errors: vec![], // TODO (#503): add better error handling
     });
     clickhouse_connection_info
-        .write(&[batch_request_insert], "BatchRequest")
+        .write(&[batch_request_insert], TableName::BatchRequest)
         .await?;
     Ok(())
 }
@@ -919,18 +919,18 @@ pub async fn write_completed_batch_inference<'a>(
     match &**function {
         FunctionConfig::Chat(_chat_function) => {
             clickhouse_connection_info
-                .write(&inference_rows_to_write, "ChatInference")
+                .write(&inference_rows_to_write, TableName::ChatInference)
                 .await?;
         }
         FunctionConfig::Json(_json_function) => {
             clickhouse_connection_info
-                .write(&inference_rows_to_write, "JsonInference")
+                .write(&inference_rows_to_write, TableName::JsonInference)
                 .await?;
         }
     }
     // Write all the ModelInference rows to the database
     clickhouse_connection_info
-        .write(&model_inference_rows_to_write, "ModelInference")
+        .write(&model_inference_rows_to_write, TableName::ModelInference)
         .await?;
 
     Ok(inferences)

tensorzero-core/src/endpoints/datasets.rs

@@ -1,3 +1,4 @@
+use crate::clickhouse::TableName;
 use crate::function::FunctionConfigType;
 #[cfg(feature = "pyo3")]
 use crate::inference::types::pyo3_helpers::{
@@ -1088,10 +1089,10 @@ pub enum DatapointKind {
 }
 
 impl DatapointKind {
-    pub fn table_name(&self) -> &'static str {
+    pub fn table_name(&self) -> TableName {
         match self {
-            DatapointKind::Chat => "ChatInferenceDatapoint",
-            DatapointKind::Json => "JsonInferenceDatapoint",
+            DatapointKind::Chat => TableName::ChatInferenceDatapoint,
+            DatapointKind::Json => TableName::JsonInferenceDatapoint,
         }
     }
 }

tensorzero-core/src/endpoints/feedback.rs

@@ -12,7 +12,7 @@ use tokio::time::Instant;
 use tracing::instrument;
 use uuid::Uuid;
 
-use crate::clickhouse::ClickHouseConnectionInfo;
+use crate::clickhouse::{ClickHouseConnectionInfo, TableName};
 use crate::config_parser::{Config, MetricConfigLevel, MetricConfigType};
 use crate::error::{Error, ErrorDetails};
 use crate::function::FunctionConfig;
@@ -274,7 +274,9 @@ async fn write_comment(
     });
     if !dryrun {
         tokio::spawn(async move {
-            let _ = connection_info.write(&[payload], "CommentFeedback").await;
+            let _ = connection_info
+                .write(&[payload], TableName::CommentFeedback)
+                .await;
         });
     }
     Ok(())
@@ -314,7 +316,7 @@ async fn write_demonstration(
     if !dryrun {
         tokio::spawn(async move {
             let _ = connection_info
-                .write(&[payload], "DemonstrationFeedback")
+                .write(&[payload], TableName::DemonstrationFeedback)
                 .await;
         });
     }
@@ -353,7 +355,7 @@ async fn write_float(
     if !dryrun {
         tokio::spawn(async move {
             let _ = connection_info
-                .write(&[payload], "FloatMetricFeedback")
+                .write(&[payload], TableName::FloatMetricFeedback)
                 .await;
         });
     }
@@ -390,7 +392,7 @@ async fn write_boolean(
     if !dryrun {
         tokio::spawn(async move {
             let _ = connection_info
-                .write(&[payload], "BooleanMetricFeedback")
+                .write(&[payload], TableName::BooleanMetricFeedback)
                 .await;
         });
     }

tensorzero-core/src/endpoints/inference.rs

@@ -21,7 +21,7 @@ use tracing::instrument;
 use uuid::Uuid;
 
 use crate::cache::{CacheOptions, CacheParamsOptions};
-use crate::clickhouse::ClickHouseConnectionInfo;
+use crate::clickhouse::{ClickHouseConnectionInfo, TableName};
 use crate::config_parser::{Config, ObjectStoreInfo};
 use crate::embeddings::EmbeddingModelTable;
 use crate::error::{Error, ErrorDetails};
@@ -843,7 +843,7 @@ async fn write_inference(
         // Write the model responses to the ModelInference table
         for response in model_responses {
             let _ = clickhouse_connection_info
-                .write(&[response], "ModelInference")
+                .write(&[response], TableName::ModelInference)
                 .await;
         }
         // Write the inference to the Inference table
@@ -852,14 +852,14 @@ async fn write_inference(
                 let chat_inference =
                     ChatInferenceDatabaseInsert::new(result, input.clone(), metadata);
                 let _ = clickhouse_connection_info
-                    .write(&[chat_inference], "ChatInference")
+                    .write(&[chat_inference], TableName::ChatInference)
                     .await;
             }
             InferenceResult::Json(result) => {
                 let json_inference =
                     JsonInferenceDatabaseInsert::new(result, input.clone(), metadata);
                 let _ = clickhouse_connection_info
-                    .write(&[json_inference], "JsonInference")
+                    .write(&[json_inference], TableName::JsonInference)
                     .await;
             }
         }