From 6301dbd512b857f610683a216ad14b09ef90950b Mon Sep 17 00:00:00 2001 From: Shuyang Li Date: Sun, 25 Jan 2026 20:48:50 -0500 Subject: [PATCH 01/46] Move feedback-related inference queries to InferenceQueries trait (#5830) Add FeedbackTargetInfo struct and three new trait methods to InferenceQueries: - get_feedback_target_info: retrieves function info by target_id (inference or episode) - get_chat_inference_tool_params: gets tool params for demonstration validation - get_json_inference_output_schema: gets output schema for demonstration validation This removes inline SQL queries from feedback/mod.rs in favor of trait methods, enabling future Postgres support for these operations. Remove serde alias from FunctionInfo.function_name --- .../src/db/clickhouse/inference_queries.rs | 146 +++++++++++++++++- .../mock_clickhouse_connection_info.rs | 37 ++++- tensorzero-core/src/db/inferences.rs | 48 +++++- .../src/endpoints/feedback/human_feedback.rs | 7 +- tensorzero-core/src/endpoints/feedback/mod.rs | 142 +++-------------- 5 files changed, 251 insertions(+), 129 deletions(-) diff --git a/tensorzero-core/src/db/clickhouse/inference_queries.rs b/tensorzero-core/src/db/clickhouse/inference_queries.rs index 9191807444a..d14dc3ad31d 100644 --- a/tensorzero-core/src/db/clickhouse/inference_queries.rs +++ b/tensorzero-core/src/db/clickhouse/inference_queries.rs @@ -1,6 +1,8 @@ use async_trait::async_trait; use itertools::Itertools; +use uuid::Uuid; +use crate::config::MetricConfigLevel; use crate::db::clickhouse::ClickHouseConnectionInfo; use crate::db::clickhouse::query_builder::parameters::add_parameter; use crate::db::clickhouse::query_builder::{ @@ -9,15 +11,20 @@ use crate::db::clickhouse::query_builder::{ }; use crate::db::inferences::{ ClickHouseStoredInferenceWithDispreferredOutputs, CountInferencesParams, - DEFAULT_INFERENCE_QUERY_LIMIT, InferenceMetadata, InferenceOutputSource, InferenceQueries, - ListInferenceMetadataParams, ListInferencesParams, PaginationParams, + DEFAULT_INFERENCE_QUERY_LIMIT, FunctionInfo, InferenceMetadata, InferenceOutputSource, + InferenceQueries, ListInferenceMetadataParams, ListInferencesParams, PaginationParams, }; use crate::function::FunctionConfigType; +use crate::tool::ToolCallConfigDatabaseInsert; +use crate::tool::deserialize_optional_tool_info; use crate::{ config::Config, error::{Error, ErrorDetails}, stored_inference::StoredInferenceDatabase, }; +use serde::Deserialize; +use serde_json::Value; +use std::collections::HashMap; /// Represents the structured parts of a single-table query /// This allows the caller to insert JOINs between the SELECT/FROM and WHERE clauses @@ -126,6 +133,141 @@ impl InferenceQueries for ClickHouseConnectionInfo { Ok(count) } + + async fn get_function_info( + &self, + target_id: &Uuid, + level: MetricConfigLevel, + ) -> Result, Error> { + let query = match level { + MetricConfigLevel::Inference => { + r" + SELECT function_name, function_type, variant_name, episode_id + FROM InferenceById + WHERE id_uint = toUInt128({target_id:UUID}) + LIMIT 1 + FORMAT JSONEachRow + SETTINGS max_threads=1 + " + } + MetricConfigLevel::Episode => { + r" + SELECT function_name, function_type, variant_name, uint_to_uuid(episode_id_uint) as episode_id + FROM InferenceByEpisodeId + WHERE episode_id_uint = toUInt128({target_id:UUID}) + LIMIT 1 + FORMAT JSONEachRow + SETTINGS max_threads=1 + " + } + }; + + let target_id_str = target_id.to_string(); + let params = HashMap::from([("target_id", target_id_str.as_str())]); + let response = self + .inner + .run_query_synchronous(query.to_string(), ¶ms) + .await?; + + if response.response.is_empty() { + return Ok(None); + } + + let info: FunctionInfo = serde_json::from_str(&response.response).map_err(|e| { + Error::new(ErrorDetails::ClickHouseDeserialization { + message: e.to_string(), + }) + })?; + + Ok(Some(info)) + } + + async fn get_chat_inference_tool_params( + &self, + function_name: &str, + inference_id: Uuid, + ) -> Result, Error> { + let query = r" + SELECT tool_params, dynamic_tools, dynamic_provider_tools, allowed_tools, tool_choice + FROM ChatInference + WHERE function_name = {function_name:String} AND id = {inference_id:String} + FORMAT JSONEachRow + "; + + let inference_id_str = inference_id.to_string(); + let params = HashMap::from([ + ("function_name", function_name), + ("inference_id", inference_id_str.as_str()), + ]); + let response = self + .inner + .run_query_synchronous(query.to_string(), ¶ms) + .await?; + + if response.response.is_empty() { + return Ok(None); + } + + #[derive(Debug, Deserialize)] + struct ToolParamsResult { + #[serde(flatten, deserialize_with = "deserialize_optional_tool_info")] + tool_params: Option, + } + + let result: ToolParamsResult = serde_json::from_str(&response.response).map_err(|e| { + Error::new(ErrorDetails::ClickHouseQuery { + message: format!("Failed to parse tool params result: {e}"), + }) + })?; + + Ok(result.tool_params) + } + + async fn get_json_inference_output_schema( + &self, + function_name: &str, + inference_id: Uuid, + ) -> Result, Error> { + let query = r" + SELECT output_schema + FROM JsonInference + WHERE function_name = {function_name:String} AND id = {inference_id:String} + FORMAT JSONEachRow + "; + + let inference_id_str = inference_id.to_string(); + let params = HashMap::from([ + ("function_name", function_name), + ("inference_id", inference_id_str.as_str()), + ]); + let response = self + .inner + .run_query_synchronous(query.to_string(), ¶ms) + .await?; + + if response.response.is_empty() { + return Ok(None); + } + + #[derive(Debug, Deserialize)] + struct OutputSchemaResult { + output_schema: String, + } + + let result: OutputSchemaResult = serde_json::from_str(&response.response).map_err(|e| { + Error::new(ErrorDetails::ClickHouseQuery { + message: format!("Failed to parse output schema result: {e}"), + }) + })?; + + let output_schema: Value = serde_json::from_str(&result.output_schema).map_err(|e| { + Error::new(ErrorDetails::ClickHouseQuery { + message: format!("Failed to parse output schema: {e}"), + }) + })?; + + Ok(Some(output_schema)) + } } /// Generates the SQL query and parameters for counting inferences. diff --git a/tensorzero-core/src/db/clickhouse/mock_clickhouse_connection_info.rs b/tensorzero-core/src/db/clickhouse/mock_clickhouse_connection_info.rs index 6423e48d5c7..4b3879eb8c4 100644 --- a/tensorzero-core/src/db/clickhouse/mock_clickhouse_connection_info.rs +++ b/tensorzero-core/src/db/clickhouse/mock_clickhouse_connection_info.rs @@ -3,6 +3,7 @@ use async_trait::async_trait; use uuid::Uuid; use crate::config::Config; +use crate::config::MetricConfigLevel; use crate::config::snapshot::{ConfigSnapshot, SnapshotHash}; use crate::db::datasets::{ DatasetMetadata, DatasetQueries, GetDatapointParams, GetDatapointsParams, @@ -15,8 +16,8 @@ use crate::db::inference_count::{ VariantThroughput, }; use crate::db::inferences::{ - CountInferencesParams as ListInferencesCountParams, InferenceMetadata, InferenceQueries, - ListInferenceMetadataParams, ListInferencesParams, MockInferenceQueries, + CountInferencesParams as ListInferencesCountParams, FunctionInfo, InferenceMetadata, + InferenceQueries, ListInferenceMetadataParams, ListInferencesParams, MockInferenceQueries, }; use crate::db::model_inferences::{MockModelInferenceQueries, ModelInferenceQueries}; use crate::db::stored_datapoint::StoredDatapoint; @@ -24,6 +25,8 @@ use crate::db::{ConfigQueries, MockConfigQueries}; use crate::error::Error; use crate::inference::types::StoredModelInference; use crate::stored_inference::StoredInferenceDatabase; +use crate::tool::ToolCallConfigDatabaseInsert; +use serde_json::Value; /// Mock struct that implements all traits on ClickHouseConnectionInfo. /// Usage: in tests, create a new mutable instance of this struct, and use the appropriate expect_ methods on the fields inside to mock @@ -81,6 +84,36 @@ impl InferenceQueries for MockClickHouseConnectionInfo { .count_inferences(config, params) .await } + + async fn get_function_info( + &self, + target_id: &Uuid, + level: MetricConfigLevel, + ) -> Result, Error> { + self.inference_queries + .get_function_info(target_id, level) + .await + } + + async fn get_chat_inference_tool_params( + &self, + function_name: &str, + inference_id: Uuid, + ) -> Result, Error> { + self.inference_queries + .get_chat_inference_tool_params(function_name, inference_id) + .await + } + + async fn get_json_inference_output_schema( + &self, + function_name: &str, + inference_id: Uuid, + ) -> Result, Error> { + self.inference_queries + .get_json_inference_output_schema(function_name, inference_id) + .await + } } #[async_trait] diff --git a/tensorzero-core/src/db/inferences.rs b/tensorzero-core/src/db/inferences.rs index 981eddc9294..46935c9d3f7 100644 --- a/tensorzero-core/src/db/inferences.rs +++ b/tensorzero-core/src/db/inferences.rs @@ -11,7 +11,7 @@ use uuid::Uuid; #[cfg(test)] use mockall::automock; -use crate::config::Config; +use crate::config::{Config, MetricConfigLevel}; use crate::db::clickhouse::query_builder::{InferenceFilter, OrderBy}; use crate::endpoints::inference::InferenceParams; use crate::error::{Error, ErrorDetails}; @@ -313,6 +313,16 @@ pub struct CountInferencesParams<'a> { pub search_query_experimental: Option<&'a str>, } +/// Function information retrieved for feedback validation. +/// Contains the function name, type, variant, and episode ID associated with an inference or episode. +#[derive(Debug, Deserialize, PartialEq)] +pub struct FunctionInfo { + pub function_name: String, + pub function_type: FunctionType, + pub variant_name: String, + pub episode_id: Uuid, +} + #[async_trait] #[cfg_attr(test, automock)] pub trait InferenceQueries { @@ -335,4 +345,40 @@ pub trait InferenceQueries { config: &Config, params: &CountInferencesParams<'_>, ) -> Result; + + /// Get function information for feedback validation by target_id. + /// + /// When `level` is `Inference`, queries by inference_id from `InferenceById`. + /// When `level` is `Episode`, queries by episode_id from `InferenceByEpisodeId`. + /// + /// Returns `None` if the target doesn't exist. + async fn get_function_info( + &self, + target_id: &Uuid, + level: MetricConfigLevel, + ) -> Result, Error>; + + /// Get tool parameters from a chat inference for demonstration validation. + /// + /// Returns the tool configuration that was used at inference time, which is needed + /// to validate demonstration feedback against the actual tools available. + /// + /// Returns `None` if the inference doesn't exist. + async fn get_chat_inference_tool_params( + &self, + function_name: &str, + inference_id: Uuid, + ) -> Result, Error>; + + /// Get output schema from a json inference for demonstration validation. + /// + /// Returns the output schema that was used at inference time, which is needed + /// to validate demonstration feedback against the actual schema. + /// + /// Returns `None` if the inference doesn't exist. + async fn get_json_inference_output_schema( + &self, + function_name: &str, + inference_id: Uuid, + ) -> Result, Error>; } diff --git a/tensorzero-core/src/endpoints/feedback/human_feedback.rs b/tensorzero-core/src/endpoints/feedback/human_feedback.rs index f3faef525d8..b888d32a30b 100644 --- a/tensorzero-core/src/endpoints/feedback/human_feedback.rs +++ b/tensorzero-core/src/endpoints/feedback/human_feedback.rs @@ -1,4 +1,5 @@ -use super::{FunctionInfo, throttled_get_function_info}; +use super::throttled_get_function_info; +use crate::db::inferences::FunctionInfo; use crate::{ config::MetricConfigLevel, db::clickhouse::{ClickHouseConnectionInfo, TableName}, @@ -108,7 +109,7 @@ async fn get_output( let FunctionInfo { function_type, episode_id, - name, + function_name, variant_name, } = function_info; let table_name = function_type.inference_table_name(); @@ -119,7 +120,7 @@ async fn get_output( WHERE id = '{inference_id}' AND episode_id = '{episode_id}' AND - function_name = '{name}' AND + function_name = '{function_name}' AND variant_name = '{variant_name}' LIMIT 1 FORMAT JSONEachRow diff --git a/tensorzero-core/src/endpoints/feedback/mod.rs b/tensorzero-core/src/endpoints/feedback/mod.rs index de964f8b810..6863504dea9 100644 --- a/tensorzero-core/src/endpoints/feedback/mod.rs +++ b/tensorzero-core/src/endpoints/feedback/mod.rs @@ -22,16 +22,14 @@ use crate::db::feedback::{ BooleanMetricFeedbackInsert, CommentFeedbackInsert, CommentTargetType, DemonstrationFeedbackInsert, FeedbackQueries, FloatMetricFeedbackInsert, }; +use crate::db::inferences::{FunctionInfo, InferenceQueries}; use crate::error::{Error, ErrorDetails}; use crate::function::FunctionConfig; use crate::inference::types::{ - ContentBlockChatOutput, ContentBlockOutput, FunctionType, Text, parse_chat_output, + ContentBlockChatOutput, ContentBlockOutput, Text, parse_chat_output, }; use crate::jsonschema_util::JSONSchema; -use crate::tool::{ - StaticToolConfig, ToolCall, ToolCallConfig, ToolCallConfigDatabaseInsert, - deserialize_optional_tool_info, -}; +use crate::tool::{StaticToolConfig, ToolCall, ToolCallConfig}; use crate::utils::gateway::{AppState, AppStateData, StructuredJson}; use crate::utils::uuid::uuid_elapsed; use tensorzero_auth::middleware::RequestApiKeyExtension; @@ -355,11 +353,11 @@ async fn write_demonstration( &inference_id, ) .await?; - let function_config = config.get_function(&function_info.name)?; + let function_config = config.get_function(&function_info.function_name)?; let dynamic_demonstration_info = get_dynamic_demonstration_info( &connection_info, inference_id, - &function_info.name, + &function_info.function_name, &function_config, &config.tools, ) @@ -539,8 +537,11 @@ async fn throttled_get_function_info( // Poll every 500ms until the deadline is reached. loop { // If an error occurs during lookup (distinct from the target_id not existing), we bail out immediately. - match get_function_info(connection_info, metric_config_level, target_id).await? { - Some(identifier) => return Ok(identifier), + let feedback_target_info = connection_info + .get_function_info(target_id, metric_config_level.clone()) + .await?; + match feedback_target_info { + Some(feedback_target_info) => return Ok(feedback_target_info), None => { if Instant::now() >= deadline { let identifier_type = match metric_config_level { @@ -563,67 +564,6 @@ async fn throttled_get_function_info( } } -/// Retrieves the function name associated with a given `target_id` of the inference or episode. -/// -/// # Arguments -/// -/// * `connection_info` - Connection details for the ClickHouse database. -/// * `metric_config_level` - The level of metric configuration, either `Inference` or `Episode`. -/// * `target_id` - The UUID of the target to be validated and retrieved. -/// -/// # Returns -/// -/// * On success: -/// - Returns a `FunctionInfo` containing the function name and type. -/// - Returns `None` if the `target_id` does not exist. -/// * On failure: -/// - Returns an `Error` if the `target_id` exists, but is invalid -async fn get_function_info( - connection_info: &ClickHouseConnectionInfo, - metric_config_level: &MetricConfigLevel, - target_id: &Uuid, -) -> Result, Error> { - let query = match metric_config_level { - MetricConfigLevel::Inference => format!( - "SELECT function_name as name, function_type, variant_name, episode_id - FROM InferenceById - WHERE id_uint = toUInt128(toUUID('{target_id}')) - LIMIT 1 - FORMAT JSONEachRow - SETTINGS max_threads=1" - ), - MetricConfigLevel::Episode => format!( - "SELECT function_name as name, function_type, variant_name, uint_to_uuid(episode_id_uint) as episode_id - FROM InferenceByEpisodeId - WHERE episode_id_uint = toUInt128(toUUID('{target_id}')) - LIMIT 1 - FORMAT JSONEachRow - SETTINGS max_threads=1" - ), - }; - let response = connection_info - .run_query_synchronous_no_params(query) - .await?; - if response.response.is_empty() { - return Ok(None); - }; - Ok(Some(serde_json::from_str(&response.response).map_err( - |e| { - Error::new(ErrorDetails::ClickHouseDeserialization { - message: e.to_string(), - }) - }, - )?)) -} - -#[derive(Debug, Deserialize, PartialEq)] -struct FunctionInfo { - name: String, - function_type: FunctionType, - variant_name: String, - episode_id: Uuid, -} - #[derive(Debug, Deserialize, PartialEq)] struct DemonstrationToolCall { name: String, @@ -767,12 +707,6 @@ pub enum DynamicDemonstrationInfo { Json(Value), } -#[derive(Debug, Deserialize)] -struct ToolParamsResult { - #[serde(flatten, deserialize_with = "deserialize_optional_tool_info")] - tool_params: Option, -} - /// In order to properly validate demonstration data we need to fetch the information that was /// passed to the inference at runtime. /// If we don't do this then we might allow some e.g. tool calls or output schemas that would not @@ -789,52 +723,22 @@ async fn get_dynamic_demonstration_info( ) -> Result { match function_config { FunctionConfig::Chat(..) => { - let parameterized_query = "SELECT tool_params, dynamic_tools, dynamic_provider_tools, allowed_tools, tool_choice FROM ChatInference WHERE function_name={function_name:String} and id={inference_id:String} FORMAT JSONEachRow".to_string(); - let result = clickhouse_client - .run_query_synchronous( - parameterized_query, - &HashMap::from([ - ("function_name", function_name), - ("inference_id", &inference_id.to_string()), - ]), - ) + let tool_params = clickhouse_client + .get_chat_inference_tool_params(function_name, inference_id) .await?; - let tool_params_result = serde_json::from_str::(&result.response) - .map_err(|e| { - Error::new(ErrorDetails::ClickHouseQuery { - message: format!("Failed to parse demonstration result: {e}"), - }) - })?; - Ok(DynamicDemonstrationInfo::Chat( // If the tool params are not present in the database, we use the default tool params (empty tools). // This is consistent with how they are serialized at inference time. - tool_params_result - .tool_params + tool_params .unwrap_or_default() .into_tool_call_config(function_config, static_tools)?, )) } FunctionConfig::Json(..) => { - let parameterized_query = "SELECT output_schema FROM JsonInference WHERE function_name={function_name:String} and id={inference_id:String} FORMAT JSONEachRow".to_string(); - let result = clickhouse_client - .run_query_synchronous( - parameterized_query, - &HashMap::from([ - ("function_name", function_name), - ("inference_id", &inference_id.to_string()), - ]), - ) - .await?; - let result_value = serde_json::from_str::(&result.response).map_err(|e| { - Error::new(ErrorDetails::ClickHouseQuery { - message: format!("Failed to parse demonstration result: {e}"), - }) - })?; - let output_schema_str = result_value - .get("output_schema") - .and_then(|v| v.as_str()) + let output_schema = clickhouse_client + .get_json_inference_output_schema(function_name, inference_id) + .await? .ok_or_else(|| { Error::new(ErrorDetails::ClickHouseQuery { message: "Failed to get output schema from demonstration result" @@ -842,15 +746,11 @@ async fn get_dynamic_demonstration_info( }) })?; - let mut output_schema = - serde_json::from_str::(output_schema_str).map_err(|e| { - Error::new(ErrorDetails::ClickHouseQuery { - message: format!("Failed to parse output schema: {e}"), - }) - })?; - if function_name.starts_with("tensorzero::llm_judge") { - output_schema = handle_llm_judge_output_schema(output_schema); - } + let output_schema = if function_name.starts_with("tensorzero::llm_judge") { + handle_llm_judge_output_schema(output_schema) + } else { + output_schema + }; Ok(DynamicDemonstrationInfo::Json(output_schema)) } } From b44f1be8534b0d672f09e0b546901838af487e82 Mon Sep 17 00:00:00 2001 From: Gabriel Bianconi <1275491+GabrielBianconi@users.noreply.github.com> Date: Sun, 25 Jan 2026 21:36:15 -0500 Subject: [PATCH 02/46] Update docs (#5845) --- docs/deployment/valkey-redis.mdx | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/docs/deployment/valkey-redis.mdx b/docs/deployment/valkey-redis.mdx index dd5ce058bf7..ac063205b9b 100644 --- a/docs/deployment/valkey-redis.mdx +++ b/docs/deployment/valkey-redis.mdx @@ -85,6 +85,5 @@ If both `TENSORZERO_VALKEY_URL` and `TENSORZERO_POSTGRES_URL` are set, the gatew ### Durability -A critical failure of Valkey (e.g. crash or power outage) could mean losing rate limiting data since the last backup. -When the rate limiting window is granular (e.g. minutes), this can be tolerated. -If your rate limiting config includes larger windows, we recommend setting up [recurrent RDB backups](https://valkey.io/topics/persistence/) (point-in-time snapshots) for better durability guarantees. +A critical failure of Valkey (e.g. server crash) may result in loss of rate limiting data since the last backup. +This is generally tolerable if your rate limiting windows are short (e.g. minutes), but if you require precise limits or longer time windows, we recommend configuring [recurring RDB (point-in-time) snapshots](https://valkey.io/topics/persistence/) for improved durability. From 1ec7fe4a301818a0cd2e0e0755aacc2f6e3556c6 Mon Sep 17 00:00:00 2001 From: Gabriel Bianconi <1275491+GabrielBianconi@users.noreply.github.com> Date: Sun, 25 Jan 2026 22:34:57 -0500 Subject: [PATCH 03/46] Add Cursor to CLA whitelist (#5853) --- .github/workflows/cla.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/cla.yml b/.github/workflows/cla.yml index c3520f0effb..1e1cef5b4d6 100644 --- a/.github/workflows/cla.yml +++ b/.github/workflows/cla.yml @@ -29,6 +29,6 @@ jobs: path-to-signatures: "ci/cla-signatures.json" path-to-document: "https://github.com/tensorzero/tensorzero/blob/main/CLA.md" branch: "cla-signatures" - allowlist: claude,dependabot[bot],TensorZero-Experimental-CI-Bot[bot] + allowlist: claude,cursoragent,dependabot[bot],TensorZero-Experimental-CI-Bot[bot] custom-pr-sign-comment: "I have read the Contributor License Agreement (CLA) and hereby sign the CLA." lock-pullrequest-aftermerge: false From 1b3e27257154f8327a96c41c612a0957f2ab236c Mon Sep 17 00:00:00 2001 From: Shuyang Li Date: Sun, 25 Jan 2026 22:41:28 -0500 Subject: [PATCH 04/46] Implement chat json inference tables and count queries for Postgres (#5692) * Add migration * Step 1 (partial): Add chat/json inference tables and InferenceCountQueries - Add migration 20260116194140_inference_tables.sql creating partitioned chat_inferences and json_inferences tables with indexes - Add migration 20260116195939_inference_default_partitions.sql creating default partitions for backfilling historical data - Implement InferenceCountQueries trait for PostgresConnectionInfo - Uses static SQL queries to satisfy sqlx 0.9 SqlSafeStr requirements - Feedback count methods stubbed until step-2 (feedback tables) - Add scripts/reset-dev-postgres.sh for dev database reset Address code review feedback for inference_count.rs - Log warning instead of silently falling back to epoch on RFC3339 parse error - Remove misleading test that tested ClickHouse table naming in Postgres module Refactor Postgres inference_count to use sqlx QueryBuilder - Replace format! SQL construction with QueryBuilder for type safety - Use push() for trusted SQL fragments, push_bind() for parameters - Remove local time_window_units helper, use TimeWindow::to_postgres_time_unit() - Extract reusable helper functions for count queries Use ENABLE_POSTGRES_READ flag to dispatch inference count queries - Update handlers to check ENABLE_POSTGRES_READ flag and use either Postgres or ClickHouse connection - Affected handlers: - get_inference_count_handler - get_inference_with_feedback_count_handler - get_function_throughput_by_variant_handler - list_functions_with_inference_count_handler - get_episode_inference_count_handler - Add tests verifying Postgres path returns error when disabled - Business logic functions are database-agnostic (take &impl InferenceCountQueries) Fix nested dollar quoting in pg_cron setup Use $cron_setup$ tag for outer DO block to avoid conflict with inner $$ used in cron.schedule command strings. Make pg_trgm extension and trigram indexes conditional - Only create pg_trgm extension if available (check pg_available_extensions) - Only create trigram indexes if pg_trgm is enabled - Graceful degradation: substring search falls back to sequential scans Add pg_cron partition management and inference retention config - Update tensorzero_drop_old_partitions to accept retention key parameter - Add pg_cron jobs to inference_tables migration (conditional on pg_cron availability): - Daily partition creation at 00:05 UTC - Daily old partition cleanup at 00:30 UTC (reads inference_retention_days) - Add inference_retention_days config option to PostgresConfig - Gateway writes retention config to tensorzero_retention_config table on startup - If retention not configured, partitions are retained indefinitely Add TODO for inference_retention_days documentation (#5764) Warn that enabling retention will immediately start dropping old partitions, and clients must manage their own backups. Add ErrorDetails::NotImplemented for stubbed Postgres features - Add NotImplemented variant with message field - Returns HTTP 501 Not Implemented status code - Replace PostgresConnection errors in inference_count.rs stubs - Reference TODO(#5691) for implementing feedback table queries Add load_fixtures_postgres.sh for loading small fixtures Loads chat_inferences and json_inferences from JSONL fixtures. Only small fixtures for now (Step 1 tables). Usage: ./ui/fixtures/load_fixtures_postgres.sh Env vars: TENSORZERO_POSTGRES_URL, TENSORZERO_SKIP_TRUNCATE Add load_fixtures_postgres.sh for loading small fixtures Loads chat_inferences and json_inferences from JSONL fixtures. Derives created_at from UUIDv7 timestamp (first 48 bits = ms since epoch). Usage: ./ui/fixtures/load_fixtures_postgres.sh Env vars: TENSORZERO_POSTGRES_URL, TENSORZERO_SKIP_TRUNCATE * e2e tests * PR comments created_at should not have a default since we compute it Test postgres functions Update fixtures * I don't think we should sleep --- .github/workflows/general.yml | 24 +- AGENTS.md | 8 + tensorzero-core/src/config/mod.rs | 11 + tensorzero-core/src/db/mod.rs | 12 + tensorzero-core/src/db/postgres/README.md | 84 ++++ .../src/db/postgres/inference_count.rs | 300 +++++++++++++ .../20260116161927_partitioning_functions.sql | 72 +++ .../20260116194140_inference_tables.sql | 104 +++++ tensorzero-core/src/db/postgres/mod.rs | 56 +++ .../src/db/postgres/rate_limiting.rs | 15 +- .../src/db/postgres/test_helpers.rs | 23 + .../internal/get_episode_inference_count.rs | 8 +- .../src/endpoints/internal/inference_count.rs | 181 +++++--- tensorzero-core/src/error/mod.rs | 24 +- .../src/function/function_config.rs | 10 + tensorzero-core/src/utils/gateway.rs | 13 + .../tests/e2e/db/inference_count_queries.rs | 361 +++++++-------- tensorzero-core/tests/e2e/db/mod.rs | 23 + tensorzero-core/tests/e2e/db/postgres/mod.rs | 1 + .../db/postgres/postgres_function_tests.rs | 416 ++++++++++++++++++ .../tests/e2e/docker-compose.live.yml | 28 ++ ui/fixtures/Dockerfile | 4 +- ui/fixtures/download-small-fixtures.py | 2 +- ui/fixtures/load_fixtures_postgres.sh | 144 +++++- ui/fixtures/reset-dev-postgres.sh | 37 ++ ui/fixtures/upload-small-fixtures.sh | 2 +- 26 files changed, 1677 insertions(+), 286 deletions(-) create mode 100644 tensorzero-core/src/db/postgres/README.md create mode 100644 tensorzero-core/src/db/postgres/inference_count.rs create mode 100644 tensorzero-core/src/db/postgres/migrations/20260116161927_partitioning_functions.sql create mode 100644 tensorzero-core/src/db/postgres/migrations/20260116194140_inference_tables.sql create mode 100644 tensorzero-core/src/db/postgres/test_helpers.rs create mode 100644 tensorzero-core/tests/e2e/db/postgres/postgres_function_tests.rs create mode 100755 ui/fixtures/reset-dev-postgres.sh diff --git a/.github/workflows/general.yml b/.github/workflows/general.yml index c8ef0e0c554..024e9874397 100644 --- a/.github/workflows/general.yml +++ b/.github/workflows/general.yml @@ -918,8 +918,12 @@ jobs: name: "Postgres tests" permissions: contents: read + # Permission to download artifacts and for rust-cache actions: read + # Permission to fetch GitHub OIDC token authentication for Namespace login + id-token: write runs-on: ubuntu-latest + needs: [build-gateway-container] if: ${{ github.event_name == 'merge_group' || (github.event.pull_request.head.repo.full_name == github.repository) }} steps: @@ -950,7 +954,25 @@ jobs: with: tool: cargo-nextest - - name: Set environment variables for Postgres and Valkey + - name: Install Namespace CLI + uses: namespacelabs/nscloud-setup@d1c625762f7c926a54bd39252efff0705fd11c64 + continue-on-error: ${{ github.event.pull_request.head.repo.full_name != github.repository || github.actor == 'dependabot[bot]' }} + + - name: Login to Namespace registry + run: nsc docker login + continue-on-error: ${{ github.event.pull_request.head.repo.full_name != github.repository || github.actor == 'dependabot[bot]' }} + + - name: Download gateway container image + run: | + docker pull nscr.io/igvf4asmf8kri/gateway:sha-${{ github.sha }} + docker tag nscr.io/igvf4asmf8kri/gateway:sha-${{ github.sha }} tensorzero/gateway:sha-${{ github.sha }} + continue-on-error: ${{ github.event.pull_request.head.repo.full_name != github.repository || github.actor == 'dependabot[bot]' }} + + - name: Set TENSORZERO_GATEWAY_TAG + run: | + echo "TENSORZERO_GATEWAY_TAG=sha-${{ github.sha }}" >> $GITHUB_ENV + + - name: Set environment variables for Postgres run: | echo "DATABASE_URL=postgres://postgres:postgres@localhost:5432/tensorzero-e2e-tests" >> $GITHUB_ENV echo "TENSORZERO_POSTGRES_URL=postgres://postgres:postgres@localhost:5432/tensorzero-e2e-tests" >> $GITHUB_ENV diff --git a/AGENTS.md b/AGENTS.md index 91056d26ee9..986a11d4927 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -26,6 +26,14 @@ - Business logic layer will generate all data that TensorZero is responsible for (e.g. UUIDs for new datapoints, `staled_at` timestamps). - Database layer (ClickHouse and/or Postgres) will insert data as-is into the backing database, with the only exception of `updated_at` timestamps which we insert by calling native functions in the database. +## For Postgres (sqlx) + +- **Do not use `format!` for SQL queries.** Use `sqlx::QueryBuilder` for dynamic queries. + - Use `.push()` for trusted SQL fragments (table names, SQL keywords). + - Use `.push_bind()` for user-provided values (prevents SQL injection, handles types). + - Use `.build_query_scalar()` for scalar results, `.build()` for row results. +- For static queries, use `sqlx::query!` directly. + # Python Dependencies We use `uv` to manage Python dependencies. diff --git a/tensorzero-core/src/config/mod.rs b/tensorzero-core/src/config/mod.rs index e0b93526bc9..5106486c2ca 100644 --- a/tensorzero-core/src/config/mod.rs +++ b/tensorzero-core/src/config/mod.rs @@ -2266,6 +2266,16 @@ pub struct PostgresConfig { pub enabled: Option, #[serde(default = "default_connection_pool_size")] pub connection_pool_size: u32, + /// Retention period in days for inference tables (chat_inferences, json_inferences). + /// If set, old partitions beyond this age will be dropped by pg_cron. + /// If not set, partitions are retained indefinitely. + /// + /// TODO(#5764): Document clearly in user-facing docs: + /// - WARNING: When first set (or lowered), pg_cron will immediately start dropping + /// partitions older than this value on its next daily run. This is irreversible. + /// - Clients are responsible for managing their own data backups before enabling retention. + /// - Recommend testing in non-production first and starting with a longer period. + pub inference_retention_days: Option, } fn default_connection_pool_size() -> u32 { @@ -2277,6 +2287,7 @@ impl Default for PostgresConfig { Self { enabled: None, connection_pool_size: 20, + inference_retention_days: None, } } } diff --git a/tensorzero-core/src/db/mod.rs b/tensorzero-core/src/db/mod.rs index 95987d430f3..3af9fb27143 100644 --- a/tensorzero-core/src/db/mod.rs +++ b/tensorzero-core/src/db/mod.rs @@ -97,6 +97,18 @@ impl TimeWindow { TimeWindow::Cumulative => "year", // Cumulative uses a full year as fallback } } + + /// Converts the time window to the PostgreSQL date_trunc time unit. + pub fn to_postgres_time_unit(&self) -> &'static str { + match self { + TimeWindow::Minute => "minute", + TimeWindow::Hour => "hour", + TimeWindow::Day => "day", + TimeWindow::Week => "week", + TimeWindow::Month => "month", + TimeWindow::Cumulative => "year", // Not used, but uses a full year as fallback + } + } } #[cfg_attr(feature = "ts-bindings", derive(ts_rs::TS))] diff --git a/tensorzero-core/src/db/postgres/README.md b/tensorzero-core/src/db/postgres/README.md new file mode 100644 index 00000000000..6c151206986 --- /dev/null +++ b/tensorzero-core/src/db/postgres/README.md @@ -0,0 +1,84 @@ +# Postgres Development + +## Local Development + +For developing Postgres features (writing queries, testing locally): + +```bash +# 1. Set environment variables +export TENSORZERO_POSTGRES_URL="postgres://postgres:postgres@localhost:5432/postgres_migration_dev" +export TENSORZERO_INTERNAL_FLAG_ENABLE_POSTGRES_READ=1 + +# 2. Reset Postgres and load fixture data +./ui/fixtures/reset-dev-postgres.sh +``` + +### Rebuilding SQLx Cache + +When you edit SQL queries, rebuild the sqlx cache: + +```bash +DATABASE_URL="postgres://postgres:postgres@localhost:5432/postgres_migration_dev" cargo sqlx prepare +``` + +We build sqlx cache within each crate (see https://github.com/tensorzero/tensorzero/pull/5622). + +## Running E2E Tests + +### With New Migrations (Fast Workflow) + +When you add new migration files, run migrations via `cargo` instead of rebuilding the Docker container: + +```bash +# 1. (Optional) Remove existing Postgres container if iterating on migrations +docker compose -f tensorzero-core/tests/e2e/docker-compose.yml down -v postgres + +# 2. Start Postgres and ClickHouse +docker compose -f tensorzero-core/tests/e2e/docker-compose.yml up -d postgres clickhouse + +# 3. Wait for Postgres to be ready +until docker compose -f tensorzero-core/tests/e2e/docker-compose.yml exec postgres pg_isready -U postgres; do sleep 1; done + +# 4. Run migrations via cargo (uses your local code) +TENSORZERO_POSTGRES_URL=postgres://postgres:postgres@localhost:5432/tensorzero-e2e-tests cargo migrate-postgres + +# 5. Load Postgres fixtures +TENSORZERO_POSTGRES_URL=postgres://postgres:postgres@localhost:5432/tensorzero-e2e-tests ./ui/fixtures/load_fixtures_postgres.sh + +# 6. Run Postgres E2E tests +TENSORZERO_POSTGRES_URL=postgres://postgres:postgres@localhost:5432/tensorzero-e2e-tests \ + cargo test --package tensorzero-core --features e2e_tests postgres::inference_count_queries +``` + +### Without New Migrations (Full Docker Compose) + +When using existing migrations (no local changes), use the standard docker-compose workflow: + +```bash +# 1. Start all services (uses published gateway image for migrations) +docker compose -f tensorzero-core/tests/e2e/docker-compose.yml up -d + +# 2. Wait for fixtures to load +docker compose -f tensorzero-core/tests/e2e/docker-compose.yml ps # Check health status + +# 3. Run tests +TENSORZERO_POSTGRES_URL=postgres://postgres:postgres@localhost:5432/tensorzero-e2e-tests \ +TENSORZERO_CLICKHOUSE_URL=http://chuser:chpassword@localhost:8123/tensorzero_e2e_tests \ + cargo test --package tensorzero-core --features e2e_tests +``` + +### With New Migrations (Slow Alternative) + +If you need the full docker-compose orchestration with new migrations: + +```bash +# 1. Rebuild gateway image with new migrations (~5 min) +docker compose -f tensorzero-core/tests/e2e/docker-compose.yml build gateway-postgres-migrations + +# 2. Start all services +docker compose -f tensorzero-core/tests/e2e/docker-compose.yml up -d +``` + +## CI Behavior + +CI sets `TENSORZERO_GATEWAY_TAG=sha-${{ github.sha }}`, which uses the freshly-built image from the PR's commit. The container rebuild is only slow for local development. diff --git a/tensorzero-core/src/db/postgres/inference_count.rs b/tensorzero-core/src/db/postgres/inference_count.rs new file mode 100644 index 00000000000..d535026e941 --- /dev/null +++ b/tensorzero-core/src/db/postgres/inference_count.rs @@ -0,0 +1,300 @@ +//! Postgres queries for inference count. + +use async_trait::async_trait; +use chrono::{DateTime, Utc}; +use sqlx::{PgPool, QueryBuilder, Row}; + +use super::PostgresConnectionInfo; +use crate::db::TimeWindow; +use crate::db::inference_count::{ + CountByVariant, CountInferencesParams, CountInferencesWithDemonstrationFeedbacksParams, + CountInferencesWithFeedbackParams, FunctionInferenceCount, + GetFunctionThroughputByVariantParams, InferenceCountQueries, VariantThroughput, +}; +use crate::error::{Error, ErrorDetails}; +use crate::function::FunctionConfigType; + +/// Builds and executes a count query for inferences. +async fn count_inferences( + pool: &PgPool, + function_type: FunctionConfigType, + function_name: &str, + variant_name: Option<&str>, +) -> Result { + let table = function_type.postgres_table_name(); + + let mut qb = QueryBuilder::new("SELECT COUNT(*) FROM "); + qb.push(table); + qb.push(" WHERE function_name = "); + qb.push_bind(function_name); + + if let Some(variant) = variant_name { + qb.push(" AND variant_name = "); + qb.push_bind(variant); + } + + qb.build_query_scalar().fetch_one(pool).await +} + +/// Builds and executes a count-by-variant query for inferences. +async fn count_by_variant( + pool: &PgPool, + function_type: FunctionConfigType, + function_name: &str, + variant_name: Option<&str>, +) -> Result, sqlx::Error> { + let table = function_type.postgres_table_name(); + + let mut qb = QueryBuilder::new( + r#"SELECT + variant_name, + COUNT(*) AS inference_count, + to_char(MAX(created_at), 'YYYY-MM-DD"T"HH24:MI:SS.MS"Z"') AS last_used_at + FROM "#, + ); + qb.push(table); + qb.push(" WHERE function_name = "); + qb.push_bind(function_name); + + if let Some(variant) = variant_name { + qb.push(" AND variant_name = "); + qb.push_bind(variant); + } + + qb.push(" GROUP BY variant_name ORDER BY inference_count DESC"); + + let rows = qb.build().fetch_all(pool).await?; + + Ok(rows + .into_iter() + .map(|row| { + let variant_name: String = row.get("variant_name"); + let inference_count: i64 = row.get("inference_count"); + let last_used_at: String = row.get("last_used_at"); + CountByVariant { + variant_name, + inference_count: inference_count as u64, + last_used_at, + } + }) + .collect()) +} + +/// Builds and executes a throughput-by-variant query. +async fn throughput_by_variant( + pool: &PgPool, + function_name: &str, + time_window: TimeWindow, + max_periods: u32, +) -> Result, Error> { + let rows = if time_window == TimeWindow::Cumulative { + // For cumulative, return all-time data grouped by variant with fixed epoch start + let mut qb = QueryBuilder::new( + r"SELECT + '1970-01-01T00:00:00.000Z'::text AS period_start, + variant_name, + COUNT(*)::INT AS count + FROM ( + SELECT variant_name FROM tensorzero.chat_inferences WHERE function_name = ", + ); + qb.push_bind(function_name); + qb.push( + " UNION ALL SELECT variant_name FROM tensorzero.json_inferences WHERE function_name = ", + ); + qb.push_bind(function_name); + qb.push(") AS combined GROUP BY variant_name ORDER BY variant_name DESC"); + + qb.build().fetch_all(pool).await? + } else { + let unit = time_window.to_postgres_time_unit(); + + let mut qb = QueryBuilder::new( + "WITH combined AS ( + SELECT variant_name, created_at FROM tensorzero.chat_inferences WHERE function_name = ", + ); + qb.push_bind(function_name); + qb.push(" UNION ALL SELECT variant_name, created_at FROM tensorzero.json_inferences WHERE function_name = "); + qb.push_bind(function_name); + qb.push( + "), + max_time AS ( + SELECT MAX(created_at) AS max_ts FROM combined + ) + SELECT + to_char(date_trunc('", + ); + qb.push(unit); + qb.push( + r#"', c.created_at), 'YYYY-MM-DD"T"HH24:MI:SS.000"Z"') AS period_start, + c.variant_name, + COUNT(*)::INT AS count + FROM combined c, max_time m + WHERE c.created_at >= m.max_ts - INTERVAL '1 "#, + ); + qb.push(unit); + qb.push("' * ("); + qb.push_bind(max_periods as i32); + qb.push(" + 1) GROUP BY date_trunc('"); + qb.push(unit); + qb.push("', c.created_at), c.variant_name ORDER BY period_start DESC, variant_name DESC"); + + qb.build().fetch_all(pool).await? + }; + + let variant_throughputs = rows + .into_iter() + .map(|row| { + let period_start_str: String = row.get("period_start"); + let period_start = DateTime::parse_from_rfc3339(&period_start_str) + .map(|dt| dt.with_timezone(&Utc)) + .map_err(|err| { + Error::new(ErrorDetails::PostgresResult { + result_type: "variant_throughput", + message: format!( + "Failed to parse `period_start` value `{period_start_str}`: {err}" + ), + }) + })?; + let variant_name: String = row.get("variant_name"); + let count: i32 = row.get("count"); + Ok(VariantThroughput { + period_start, + variant_name, + count: count as u32, + }) + }) + .collect::, Error>>()?; + Ok(variant_throughputs) +} + +#[async_trait] +impl InferenceCountQueries for PostgresConnectionInfo { + async fn count_inferences_for_function( + &self, + params: CountInferencesParams<'_>, + ) -> Result { + let pool = self.get_pool_result()?; + let count = count_inferences( + pool, + params.function_type, + params.function_name, + params.variant_name, + ) + .await + .map_err(Error::from)?; + Ok(count as u64) + } + + async fn count_inferences_by_variant( + &self, + params: CountInferencesParams<'_>, + ) -> Result, Error> { + let pool = self.get_pool_result()?; + count_by_variant( + pool, + params.function_type, + params.function_name, + params.variant_name, + ) + .await + .map_err(Error::from) + } + + async fn count_inferences_with_feedback( + &self, + _params: CountInferencesWithFeedbackParams<'_>, + ) -> Result { + // TODO(#5691): Implement when feedback tables are added in step-2 + Err(Error::new(ErrorDetails::NotImplemented { + message: "count_inferences_with_feedback not yet implemented for Postgres".to_string(), + })) + } + + async fn count_inferences_with_demonstration_feedback( + &self, + _params: CountInferencesWithDemonstrationFeedbacksParams<'_>, + ) -> Result { + // TODO(#5691): Implement when feedback tables are added in step-2 + Err(Error::new(ErrorDetails::NotImplemented { + message: + "count_inferences_with_demonstration_feedback not yet implemented for Postgres" + .to_string(), + })) + } + + async fn count_inferences_for_episode(&self, episode_id: uuid::Uuid) -> Result { + let pool = self.get_pool_result()?; + + let mut qb = QueryBuilder::new( + r"SELECT COUNT(*) FROM ( + SELECT id FROM tensorzero.chat_inferences WHERE episode_id = ", + ); + qb.push_bind(episode_id); + qb.push(" UNION ALL SELECT id FROM tensorzero.json_inferences WHERE episode_id = "); + qb.push_bind(episode_id); + qb.push(") AS combined"); + + let count: i64 = qb + .build_query_scalar() + .fetch_one(pool) + .await + .map_err(Error::from)?; + + Ok(count as u64) + } + + async fn get_function_throughput_by_variant( + &self, + params: GetFunctionThroughputByVariantParams<'_>, + ) -> Result, Error> { + let pool = self.get_pool_result()?; + throughput_by_variant( + pool, + params.function_name, + params.time_window, + params.max_periods, + ) + .await + } + + async fn list_functions_with_inference_count( + &self, + ) -> Result, Error> { + let pool = self.get_pool_result()?; + + let rows = sqlx::query( + r" + SELECT + function_name, + MAX(created_at) AS last_inference_timestamp, + COUNT(*)::INT AS inference_count + FROM ( + SELECT function_name, created_at FROM tensorzero.chat_inferences + UNION ALL + SELECT function_name, created_at FROM tensorzero.json_inferences + ) AS combined + GROUP BY function_name + ORDER BY last_inference_timestamp DESC + ", + ) + .fetch_all(pool) + .await + .map_err(Error::from)?; + + let results = rows + .into_iter() + .map(|row| { + let function_name: String = row.get("function_name"); + let last_inference_timestamp: DateTime = row.get("last_inference_timestamp"); + let inference_count: i32 = row.get("inference_count"); + FunctionInferenceCount { + function_name, + last_inference_timestamp, + inference_count: inference_count as u32, + } + }) + .collect(); + + Ok(results) + } +} diff --git a/tensorzero-core/src/db/postgres/migrations/20260116161927_partitioning_functions.sql b/tensorzero-core/src/db/postgres/migrations/20260116161927_partitioning_functions.sql new file mode 100644 index 00000000000..db8357160e6 --- /dev/null +++ b/tensorzero-core/src/db/postgres/migrations/20260116161927_partitioning_functions.sql @@ -0,0 +1,72 @@ +-- Schema wrapping all tensorzero constructs +CREATE SCHEMA IF NOT EXISTS tensorzero; + +-- Retention configuration table +-- Gateway writes retention_days on startup; pg_cron jobs read it +CREATE TABLE IF NOT EXISTS tensorzero.retention_config ( + key TEXT PRIMARY KEY, + value TEXT NOT NULL, + updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW() +); + +-- Partition management: create future partitions for a given table +-- Called by pg_cron daily or manually, or when creating new partitioned tables +-- p_table_name should be unqualified (e.g., 'chat_inferences'), tables are created in tensorzero schema +CREATE OR REPLACE FUNCTION tensorzero.create_partitions(p_table_name TEXT) +RETURNS void AS $$ +DECLARE + partition_date DATE; + partition_name TEXT; +BEGIN + FOR i IN 0..7 LOOP + partition_date := CURRENT_DATE + i; + partition_name := p_table_name || '_' || to_char(partition_date, 'YYYY_MM_DD'); + IF NOT EXISTS ( + SELECT 1 FROM pg_tables WHERE schemaname = 'tensorzero' AND tablename = partition_name + ) THEN + EXECUTE format( + 'CREATE TABLE IF NOT EXISTS tensorzero.%I PARTITION OF tensorzero.%I FOR VALUES FROM (%L) TO (%L)', + partition_name, + p_table_name, + partition_date, + partition_date + 1 + ); + END IF; + END LOOP; +END; +$$ LANGUAGE plpgsql; + +-- Partition management: drop old partitions for a given table +-- Called by pg_cron daily or manually +-- p_retention_key: key in tensorzero.retention_config (e.g., 'inference_retention_days') +CREATE OR REPLACE FUNCTION tensorzero.drop_old_partitions(p_table_name TEXT, p_retention_key TEXT) +RETURNS void AS $$ +DECLARE + retention_days INT; + cutoff_date DATE; + partition_record RECORD; + pattern TEXT; +BEGIN + SELECT value::INT INTO retention_days + FROM tensorzero.retention_config + WHERE key = p_retention_key; + + IF retention_days IS NULL THEN + RAISE NOTICE '% not configured, skipping partition cleanup for %', p_retention_key, p_table_name; + RETURN; + END IF; + + cutoff_date := CURRENT_DATE - retention_days; + pattern := '^' || p_table_name || '_\d{4}_\d{2}_\d{2}$'; + + FOR partition_record IN + SELECT tablename + FROM pg_tables + WHERE schemaname = 'tensorzero' AND tablename ~ pattern + LOOP + IF to_date(substring(partition_record.tablename from '\d{4}_\d{2}_\d{2}$'), 'YYYY_MM_DD') < cutoff_date THEN + EXECUTE format('DROP TABLE tensorzero.%I', partition_record.tablename); + END IF; + END LOOP; +END; +$$ LANGUAGE plpgsql; diff --git a/tensorzero-core/src/db/postgres/migrations/20260116194140_inference_tables.sql b/tensorzero-core/src/db/postgres/migrations/20260116194140_inference_tables.sql new file mode 100644 index 00000000000..9bca6fe5d20 --- /dev/null +++ b/tensorzero-core/src/db/postgres/migrations/20260116194140_inference_tables.sql @@ -0,0 +1,104 @@ +-- TODO(#5691): Add indexes for substring search on input/output JSONB columns + +-- Extract timestamp from UUIDv7 +-- UUIDv7 stores milliseconds since Unix epoch in the first 48 bits +CREATE OR REPLACE FUNCTION tensorzero.uuid_v7_to_timestamp(p_uuid UUID) +RETURNS TIMESTAMPTZ AS $$ +BEGIN + RETURN to_timestamp( + ('x' || substring(replace(p_uuid::text, '-', ''), 1, 12))::bit(48)::bigint / 1000.0 + ) AT TIME ZONE 'UTC'; +END; +$$ LANGUAGE plpgsql IMMUTABLE STRICT; + +-- chat_inferences (partitioned by day) +CREATE TABLE tensorzero.chat_inferences ( + id UUID NOT NULL, + function_name TEXT NOT NULL, + variant_name TEXT NOT NULL, + episode_id UUID NOT NULL, + input JSONB NOT NULL, + output JSONB NOT NULL, + tool_params JSONB NOT NULL, + inference_params JSONB NOT NULL, + processing_time_ms INTEGER, + ttft_ms INTEGER, + tags JSONB NOT NULL DEFAULT '{}', + extra_body JSONB NOT NULL DEFAULT '[]', + dynamic_tools JSONB NOT NULL DEFAULT '[]', + dynamic_provider_tools JSONB NOT NULL DEFAULT '[]', + allowed_tools JSONB, + tool_choice TEXT, + parallel_tool_calls BOOLEAN, + snapshot_hash BYTEA, + created_at TIMESTAMPTZ NOT NULL, + PRIMARY KEY (id, created_at) +) PARTITION BY RANGE (created_at); + +-- Indexes for chat_inferences +CREATE INDEX idx_chat_inferences_function_variant_episode + ON tensorzero.chat_inferences(function_name, variant_name, episode_id); +CREATE INDEX idx_chat_inferences_episode_id ON tensorzero.chat_inferences(episode_id); +CREATE INDEX idx_chat_inferences_function_id ON tensorzero.chat_inferences(function_name, id); +CREATE INDEX idx_chat_inferences_tags ON tensorzero.chat_inferences USING GIN (tags); + +-- json_inferences (partitioned by day) +CREATE TABLE tensorzero.json_inferences ( + id UUID NOT NULL, + function_name TEXT NOT NULL, + variant_name TEXT NOT NULL, + episode_id UUID NOT NULL, + input JSONB NOT NULL, + output JSONB NOT NULL, + output_schema JSONB NOT NULL, + inference_params JSONB NOT NULL, + processing_time_ms INTEGER, + ttft_ms INTEGER, + tags JSONB NOT NULL DEFAULT '{}', + extra_body JSONB NOT NULL DEFAULT '[]', + auxiliary_content JSONB NOT NULL, + snapshot_hash BYTEA, + created_at TIMESTAMPTZ NOT NULL, + PRIMARY KEY (id, created_at) +) PARTITION BY RANGE (created_at); + +-- Indexes for json_inferences +CREATE INDEX idx_json_inferences_function_variant_episode + ON tensorzero.json_inferences(function_name, variant_name, episode_id); +CREATE INDEX idx_json_inferences_episode_id ON tensorzero.json_inferences(episode_id); +CREATE INDEX idx_json_inferences_function_id ON tensorzero.json_inferences(function_name, id); +CREATE INDEX idx_json_inferences_tags ON tensorzero.json_inferences USING GIN (tags); + +-- Create initial partitions for the next 7 days +SELECT tensorzero.create_partitions('chat_inferences'); +SELECT tensorzero.create_partitions('json_inferences'); + +-- Schedule pg_cron jobs for partition management (if pg_cron is available) +DO $cron_setup$ +BEGIN + -- Check if pg_cron extension is available + IF EXISTS (SELECT 1 FROM pg_extension WHERE extname = 'pg_cron') THEN + -- Create future partitions daily at 00:05 UTC + PERFORM cron.schedule( + 'tensorzero_create_inference_partitions', + '5 0 * * *', + $$SELECT tensorzero.create_partitions('chat_inferences'); SELECT tensorzero.create_partitions('json_inferences');$$ + ); + + -- Drop old partitions daily at 00:30 UTC (only acts if retention is configured) + PERFORM cron.schedule( + 'tensorzero_drop_old_inference_partitions', + '30 0 * * *', + $$SELECT tensorzero.drop_old_partitions('chat_inferences', 'inference_retention_days'); SELECT tensorzero.drop_old_partitions('json_inferences', 'inference_retention_days');$$ + ); + + RAISE NOTICE 'pg_cron jobs scheduled for inference partition management'; + ELSE + RAISE NOTICE 'pg_cron extension not available - partition management must be scheduled externally'; + END IF; +END $cron_setup$; + +-- Default partitions for backfilling historical data +-- These catch any rows that don't fit into the date-based partitions +CREATE TABLE tensorzero.chat_inferences_default PARTITION OF tensorzero.chat_inferences DEFAULT; +CREATE TABLE tensorzero.json_inferences_default PARTITION OF tensorzero.json_inferences DEFAULT; diff --git a/tensorzero-core/src/db/postgres/mod.rs b/tensorzero-core/src/db/postgres/mod.rs index 5ac8ea3de9e..9ccfd270550 100644 --- a/tensorzero-core/src/db/postgres/mod.rs +++ b/tensorzero-core/src/db/postgres/mod.rs @@ -11,8 +11,12 @@ use crate::error::{Error, ErrorDetails}; use super::HealthCheckable; pub mod experimentation; +pub mod inference_count; pub mod rate_limiting; +#[cfg(any(test, feature = "e2e_tests"))] +pub mod test_helpers; + const RUN_MIGRATIONS_COMMAND: &str = "Please see our documentation to learn more about deploying Postgres: https://www.tensorzero.com/docs/deployment/postgres"; #[derive(Debug, Clone)] @@ -120,6 +124,58 @@ impl PostgresConnectionInfo { } Ok(()) } + + /// Writes retention configuration to the `tensorzero.retention_config` table. + /// This is called on gateway startup to sync config from tensorzero.toml to Postgres. + pub async fn write_retention_config( + &self, + inference_retention_days: Option, + ) -> Result<(), Error> { + let Some(pool) = self.get_pool() else { + return Ok(()); + }; + + match inference_retention_days { + Some(days) => { + sqlx::query( + r" + INSERT INTO tensorzero.retention_config (key, value, updated_at) + VALUES ('inference_retention_days', $1, NOW()) + ON CONFLICT (key) DO UPDATE SET value = $1, updated_at = NOW() + ", + ) + .bind(days.to_string()) + .execute(pool) + .await + .map_err(|e| { + Error::new(ErrorDetails::PostgresQuery { + message: format!("Failed to write inference_retention_days config: {e}"), + }) + })?; + tracing::info!( + inference_retention_days = days, + "Configured inference retention policy" + ); + } + None => { + sqlx::query( + "DELETE FROM tensorzero.retention_config WHERE key = 'inference_retention_days'", + ) + .execute(pool) + .await + .map_err(|e| { + Error::new(ErrorDetails::PostgresQuery { + message: format!("Failed to clear inference_retention_days config: {e}"), + }) + })?; + tracing::debug!( + "Inference retention policy not configured (partitions retained indefinitely)" + ); + } + } + + Ok(()) + } } #[async_trait] diff --git a/tensorzero-core/src/db/postgres/rate_limiting.rs b/tensorzero-core/src/db/postgres/rate_limiting.rs index 60ce46f1138..8da458d3a81 100644 --- a/tensorzero-core/src/db/postgres/rate_limiting.rs +++ b/tensorzero-core/src/db/postgres/rate_limiting.rs @@ -33,7 +33,6 @@ impl RateLimitQueries for PostgresConnectionInfo { let pool = self.get_pool().ok_or_else(|| { Error::new(ErrorDetails::PostgresQuery { message: "Failed to consume tickets for rate limiting: PostgreSQL connection is disabled.".to_string(), - function_name: None, }) })?; @@ -76,9 +75,8 @@ impl RateLimitQueries for PostgresConnectionInfo { } let pool = self.get_pool().ok_or_else(|| { - Error::new(ErrorDetails::PostgresQuery { + Error::new(ErrorDetails::PostgresConnection { message: "PostgreSQL connection is disabled".to_string(), - function_name: None, }) })?; @@ -106,8 +104,9 @@ impl RateLimitQueries for PostgresConnectionInfo { .await .map_err(|e| { Error::new(ErrorDetails::PostgresQuery { - message: format!("Database query failed: {e}"), - function_name: Some("return_multiple_resource_tickets".to_string()), + message: format!( + "Database query failed in function return_multiple_resource_tickets: {e}" + ), }) })?; @@ -129,7 +128,6 @@ impl RateLimitQueries for PostgresConnectionInfo { let pool = self.get_pool().ok_or_else(|| { Error::new(ErrorDetails::PostgresQuery { message: "PostgreSQL connection is disabled".to_string(), - function_name: None, }) })?; @@ -144,8 +142,9 @@ impl RateLimitQueries for PostgresConnectionInfo { .await .map_err(|e| { Error::new(ErrorDetails::PostgresQuery { - message: format!("Database query failed: {e}"), - function_name: Some("get_resource_bucket_balance".to_string()), + message: format!( + "Database query failed in function get_resource_bucket_balance: {e}" + ), }) })?; diff --git a/tensorzero-core/src/db/postgres/test_helpers.rs b/tensorzero-core/src/db/postgres/test_helpers.rs new file mode 100644 index 00000000000..7276db18169 --- /dev/null +++ b/tensorzero-core/src/db/postgres/test_helpers.rs @@ -0,0 +1,23 @@ +#![expect(clippy::expect_used, clippy::missing_panics_doc, clippy::print_stdout)] +use std::sync::LazyLock; + +use sqlx::postgres::PgPoolOptions; + +use super::PostgresConnectionInfo; + +pub static POSTGRES_URL: LazyLock = LazyLock::new(|| { + std::env::var("TENSORZERO_POSTGRES_URL") + .expect("Environment variable TENSORZERO_POSTGRES_URL must be set") +}); + +pub async fn get_postgres() -> PostgresConnectionInfo { + let postgres_url = &*POSTGRES_URL; + let start = std::time::Instant::now(); + println!("Connecting to PostgreSQL"); + let pool = PgPoolOptions::new() + .connect(postgres_url) + .await + .expect("Failed to connect to PostgreSQL"); + println!("Connected to PostgreSQL in {:?}", start.elapsed()); + PostgresConnectionInfo::new_with_pool(pool) +} diff --git a/tensorzero-core/src/endpoints/episodes/internal/get_episode_inference_count.rs b/tensorzero-core/src/endpoints/episodes/internal/get_episode_inference_count.rs index 412fb63e83a..40346812000 100644 --- a/tensorzero-core/src/endpoints/episodes/internal/get_episode_inference_count.rs +++ b/tensorzero-core/src/endpoints/episodes/internal/get_episode_inference_count.rs @@ -1,6 +1,7 @@ use crate::{ db::inference_count::InferenceCountQueries, error::Error, + feature_flags::ENABLE_POSTGRES_READ, utils::gateway::{AppState, AppStateData}, }; use axum::{ @@ -31,8 +32,11 @@ pub async fn get_episode_inference_count_handler( State(app_state): AppState, Path(episode_id): Path, ) -> Result, Error> { - let stats = - get_episode_inference_count(&app_state.clickhouse_connection_info, episode_id).await?; + let stats = if ENABLE_POSTGRES_READ.get() { + get_episode_inference_count(&app_state.postgres_connection_info, episode_id).await? + } else { + get_episode_inference_count(&app_state.clickhouse_connection_info, episode_id).await? + }; Ok(Json(stats)) } diff --git a/tensorzero-core/src/endpoints/internal/inference_count.rs b/tensorzero-core/src/endpoints/internal/inference_count.rs index 2533a7da45f..ff7923e3e3f 100644 --- a/tensorzero-core/src/endpoints/internal/inference_count.rs +++ b/tensorzero-core/src/endpoints/internal/inference_count.rs @@ -14,6 +14,7 @@ use crate::db::inference_count::{ GetFunctionThroughputByVariantParams, InferenceCountQueries, VariantThroughput, }; use crate::error::{Error, ErrorDetails}; +use crate::feature_flags::ENABLE_POSTGRES_READ; use crate::function::DEFAULT_FUNCTION_NAME; use crate::utils::gateway::{AppState, AppStateData}; @@ -136,15 +137,14 @@ pub async fn get_inference_count_handler( Path(function_name): Path, Query(params): Query, ) -> Result, Error> { - Ok(Json( - get_inference_count( - &app_state.config, - &app_state.clickhouse_connection_info, - &function_name, - params, - ) - .await?, - )) + let database: &(dyn InferenceCountQueries + Sync) = if ENABLE_POSTGRES_READ.get() { + &app_state.postgres_connection_info + } else { + &app_state.clickhouse_connection_info + }; + + let response = get_inference_count(&app_state.config, database, &function_name, params).await?; + Ok(Json(response)) } /// HTTP handler for the feedback stats endpoint @@ -162,22 +162,27 @@ pub async fn get_inference_with_feedback_count_handler( Path((function_name, metric_name)): Path<(String, String)>, Query(params): Query, ) -> Result, Error> { - Ok(Json( - get_inference_with_feedback_count( - &app_state.config, - &app_state.clickhouse_connection_info, - function_name, - metric_name, - params, - ) - .await?, - )) + let database: &(dyn InferenceCountQueries + Sync) = if ENABLE_POSTGRES_READ.get() { + &app_state.postgres_connection_info + } else { + &app_state.clickhouse_connection_info + }; + + let response = get_inference_with_feedback_count( + &app_state.config, + database, + function_name, + metric_name, + params, + ) + .await?; + Ok(Json(response)) } /// Core business logic for getting inference count async fn get_inference_count( config: &Config, - clickhouse: &impl InferenceCountQueries, + database: &(dyn InferenceCountQueries + Sync), function_name: &str, params: InferenceCountQueryParams, ) -> Result { @@ -205,7 +210,7 @@ async fn get_inference_count( // Handle group_by=variant case if let Some(InferenceCountGroupBy::Variant) = params.group_by { - let variant_rows = clickhouse.count_inferences_by_variant(count_params).await?; + let variant_rows = database.count_inferences_by_variant(count_params).await?; let inference_count = variant_rows.iter().map(|r| r.inference_count).sum(); let count_by_variant = variant_rows.into_iter().map(Into::into).collect(); @@ -216,9 +221,7 @@ async fn get_inference_count( }); } - let inference_count = clickhouse - .count_inferences_for_function(count_params) - .await?; + let inference_count = database.count_inferences_for_function(count_params).await?; Ok(InferenceCountResponse { inference_count, @@ -229,7 +232,7 @@ async fn get_inference_count( /// Core business logic for getting feedback count async fn get_inference_with_feedback_count( config: &Config, - clickhouse_connection_info: &impl InferenceCountQueries, + database: &(dyn InferenceCountQueries + Sync), function_name: String, metric_name: String, params: InferenceWithFeedbackCountQueryParams, @@ -242,7 +245,7 @@ async fn get_inference_with_feedback_count( // TODO(shuyangli): it's probably wrong that we're not distinguishing between the feedback type // ("demonstration") and the metric name, but we can fix it later. if metric_name == "demonstration" { - let feedback_count = clickhouse_connection_info + let feedback_count = database .count_inferences_with_demonstration_feedback( CountInferencesWithDemonstrationFeedbacksParams { function_name: &function_name, @@ -263,24 +266,20 @@ async fn get_inference_with_feedback_count( // Get feedback and matching inference counts based on metric type let (feedback_count, inference_count) = try_join( - clickhouse_connection_info.count_inferences_with_feedback( - CountInferencesWithFeedbackParams { - function_name: &function_name, - function_type, - metric_name: &metric_name, - metric_config, - metric_threshold: None, - }, - ), - clickhouse_connection_info.count_inferences_with_feedback( - CountInferencesWithFeedbackParams { - function_name: &function_name, - function_type, - metric_name: &metric_name, - metric_config, - metric_threshold: Some(params.threshold), - }, - ), + database.count_inferences_with_feedback(CountInferencesWithFeedbackParams { + function_name: &function_name, + function_type, + metric_name: &metric_name, + metric_config, + metric_threshold: None, + }), + database.count_inferences_with_feedback(CountInferencesWithFeedbackParams { + function_name: &function_name, + function_type, + metric_name: &metric_name, + metric_config, + metric_threshold: Some(params.threshold), + }), ) .await?; @@ -302,13 +301,15 @@ pub async fn get_function_throughput_by_variant_handler( Path(function_name): Path, Query(params): Query, ) -> Result, Error> { - let response = get_function_throughput_by_variant( - &state.config, - &state.clickhouse_connection_info, - &function_name, - params, - ) - .await?; + let database: &(dyn InferenceCountQueries + Sync) = if ENABLE_POSTGRES_READ.get() { + &state.postgres_connection_info + } else { + &state.clickhouse_connection_info + }; + + let response = + get_function_throughput_by_variant(&state.config, database, &function_name, params).await?; + Ok(Json(response)) } @@ -316,14 +317,14 @@ pub async fn get_function_throughput_by_variant_handler( /// Validates the function exists and returns throughput data grouped by variant and time period. pub async fn get_function_throughput_by_variant( config: &Config, - clickhouse_connection_info: &impl InferenceCountQueries, + database: &(dyn InferenceCountQueries + Sync), function_name: &str, params: FunctionThroughputByVariantQueryParams, ) -> Result { // Validate function exists config.get_function(function_name)?; - let throughput = clickhouse_connection_info + let throughput = database .get_function_throughput_by_variant(GetFunctionThroughputByVariantParams { function_name, time_window: params.time_window, @@ -340,17 +341,21 @@ pub async fn get_function_throughput_by_variant( pub async fn list_functions_with_inference_count_handler( State(state): State, ) -> Result, Error> { - let response = list_functions_with_inference_count(&state.clickhouse_connection_info).await?; + let database: &(dyn InferenceCountQueries + Sync) = if ENABLE_POSTGRES_READ.get() { + &state.postgres_connection_info + } else { + &state.clickhouse_connection_info + }; + + let response = list_functions_with_inference_count(database).await?; Ok(Json(response)) } /// Core business logic for listing all functions with their inference counts async fn list_functions_with_inference_count( - clickhouse_connection_info: &impl InferenceCountQueries, + database: &(dyn InferenceCountQueries + Sync), ) -> Result { - let functions = clickhouse_connection_info - .list_functions_with_inference_count() - .await?; + let functions = database.list_functions_with_inference_count().await?; Ok(ListFunctionsWithInferenceCountResponse { functions }) } @@ -360,6 +365,7 @@ mod tests { use super::*; use crate::config::{Config, ConfigFileGlob}; use crate::db::clickhouse::MockClickHouseConnectionInfo; + use crate::db::postgres::PostgresConnectionInfo; use crate::function::FunctionConfigType; use crate::testing::get_unit_test_gateway_handle; use std::io::Write; @@ -630,4 +636,63 @@ mod tests { assert_eq!(result.inference_count, 10); } + + // Tests for ENABLE_POSTGRES_READ flag dispatch + // Note: The business logic functions take `&impl InferenceCountQueries` and are database-agnostic. + // The flag only affects which connection the handlers pass to the business logic. + // These tests verify the Postgres implementation can be invoked through the same interface. + + #[tokio::test] + async fn test_get_inference_count_with_postgres_disabled_returns_error() { + // When Postgres is disabled, calling the Postgres implementation should return an error + + let config_str = r#" + [functions.test_function] + type = "chat" + + [functions.test_function.variants.test_variant] + type = "chat_completion" + model = "openai::gpt-4" + "#; + + let mut temp_file = NamedTempFile::new().unwrap(); + temp_file.write_all(config_str.as_bytes()).unwrap(); + + let config = Config::load_from_path_optional_verify_credentials( + &ConfigFileGlob::new_from_path(temp_file.path()).unwrap(), + false, + ) + .await + .unwrap() + .into_config_without_writing_for_tests(); + + let postgres = PostgresConnectionInfo::new_disabled(); + let params = InferenceCountQueryParams { + variant_name: None, + group_by: None, + }; + + let result = get_inference_count(&config, &postgres, "test_function", params).await; + + assert!( + result.is_err(), + "Should return error when Postgres is disabled" + ); + let err = result.unwrap_err(); + assert!( + err.to_string().contains("disabled"), + "Error should indicate database is disabled: {err}" + ); + } + + #[tokio::test] + async fn test_list_functions_with_postgres_disabled_returns_error() { + let postgres = PostgresConnectionInfo::new_disabled(); + let result = list_functions_with_inference_count(&postgres).await; + + assert!( + result.is_err(), + "Should return error when Postgres is disabled" + ); + } } diff --git a/tensorzero-core/src/error/mod.rs b/tensorzero-core/src/error/mod.rs index 10178ceb17f..7aa2b5a6773 100644 --- a/tensorzero-core/src/error/mod.rs +++ b/tensorzero-core/src/error/mod.rs @@ -493,6 +493,10 @@ pub enum ErrorDetails { message: String, }, NoFallbackVariantsRemaining, + /// Feature not yet implemented. Used for stubbed functionality during incremental migration. + NotImplemented { + message: String, + }, Observability { message: String, }, @@ -518,7 +522,6 @@ pub enum ErrorDetails { message: String, }, PostgresQuery { - function_name: Option, message: String, }, PostgresResult { @@ -716,6 +719,7 @@ impl ErrorDetails { ErrorDetails::ModelNotFound { .. } => tracing::Level::WARN, ErrorDetails::ModelValidation { .. } => tracing::Level::ERROR, ErrorDetails::NoFallbackVariantsRemaining => tracing::Level::WARN, + ErrorDetails::NotImplemented { .. } => tracing::Level::ERROR, ErrorDetails::Observability { .. } => tracing::Level::WARN, ErrorDetails::OutputParsing { .. } => tracing::Level::WARN, ErrorDetails::OutputValidation { .. } => tracing::Level::WARN, @@ -872,6 +876,7 @@ impl ErrorDetails { ErrorDetails::ModelProvidersExhausted { .. } => StatusCode::INTERNAL_SERVER_ERROR, ErrorDetails::ModelValidation { .. } => StatusCode::INTERNAL_SERVER_ERROR, ErrorDetails::NoFallbackVariantsRemaining => StatusCode::BAD_GATEWAY, + ErrorDetails::NotImplemented { .. } => StatusCode::NOT_IMPLEMENTED, ErrorDetails::Observability { .. } => StatusCode::INTERNAL_SERVER_ERROR, ErrorDetails::OptimizationResponse { .. } => StatusCode::INTERNAL_SERVER_ERROR, ErrorDetails::OutputParsing { .. } => StatusCode::INTERNAL_SERVER_ERROR, @@ -1474,6 +1479,9 @@ impl std::fmt::Display for ErrorDetails { ErrorDetails::NoFallbackVariantsRemaining => { write!(f, "No fallback variants remaining.") } + ErrorDetails::NotImplemented { message } => { + write!(f, "Not implemented: {message}") + } ErrorDetails::ModelValidation { message } => { write!(f, "Failed to validate model: {message}") } @@ -1522,16 +1530,9 @@ impl std::fmt::Display for ErrorDetails { "Unexpected Postgres result of type {result_type}: {message}" ) } - ErrorDetails::PostgresQuery { - function_name, - message, - } => match function_name { - Some(function_name) => write!( - f, - "Postgres query failed in function {function_name} with message: {message}" - ), - None => write!(f, "Postgres query failed: {message}"), - }, + ErrorDetails::PostgresQuery { message } => { + write!(f, "Postgres query failed: {message}") + } ErrorDetails::ValkeyConnection { message } => { write!(f, "Error connecting to Valkey: {message}") } @@ -1708,7 +1709,6 @@ impl From for Error { fn from(err: sqlx::Error) -> Self { Self::new(ErrorDetails::PostgresQuery { message: err.to_string(), - function_name: None, }) } } diff --git a/tensorzero-core/src/function/function_config.rs b/tensorzero-core/src/function/function_config.rs index a369aecfb97..4fb148cf48b 100644 --- a/tensorzero-core/src/function/function_config.rs +++ b/tensorzero-core/src/function/function_config.rs @@ -87,6 +87,7 @@ pub enum FunctionConfigType { } impl FunctionConfigType { + /// Returns the ClickHouse table name for the given function type. pub fn table_name(&self) -> &'static str { match self { FunctionConfigType::Chat => "ChatInference", @@ -94,6 +95,15 @@ impl FunctionConfigType { } } + /// Returns the Postgres table name for the given function type. + pub fn postgres_table_name(&self) -> &'static str { + match self { + FunctionConfigType::Chat => "tensorzero.chat_inferences", + FunctionConfigType::Json => "tensorzero.json_inferences", + } + } + + /// Returns the ClickHouse datapoint table name for the given function type. pub fn datapoint_table_name(&self) -> &'static str { match self { FunctionConfigType::Chat => "ChatInferenceDatapoint", diff --git a/tensorzero-core/src/utils/gateway.rs b/tensorzero-core/src/utils/gateway.rs index dfd16d5c32c..c08f9ae8734 100644 --- a/tensorzero-core/src/utils/gateway.rs +++ b/tensorzero-core/src/utils/gateway.rs @@ -464,6 +464,8 @@ async fn create_postgres_connection( Ok(connection_info) } +// TODO(#5764): We should test that on startup we issue the correct SQL for write_retention_config, +// but this is currently structured that's difficult to swap in a Mock. pub async fn setup_postgres( config: &Config, postgres_url: Option, @@ -498,6 +500,11 @@ pub async fn setup_postgres( } }; + // Write retention config to Postgres (syncs tensorzero.toml -> database) + postgres_connection_info + .write_retention_config(config.postgres.inference_retention_days) + .await?; + Ok(postgres_connection_info) } @@ -923,6 +930,7 @@ mod tests { postgres: PostgresConfig { enabled: Some(false), connection_pool_size: 20, + inference_retention_days: None, }, ..Default::default() })); @@ -941,6 +949,7 @@ mod tests { postgres: PostgresConfig { enabled: Some(false), connection_pool_size: 20, + inference_retention_days: None, }, ..Default::default() })); @@ -964,6 +973,7 @@ mod tests { postgres: PostgresConfig { enabled: None, connection_pool_size: 20, + inference_retention_days: None, }, ..Default::default() })); @@ -982,6 +992,7 @@ mod tests { postgres: PostgresConfig { enabled: Some(true), connection_pool_size: 20, + inference_retention_days: None, }, ..Default::default() })); @@ -1000,6 +1011,7 @@ mod tests { postgres: PostgresConfig { enabled: Some(true), connection_pool_size: 20, + inference_retention_days: None, }, ..Default::default() })); @@ -1016,6 +1028,7 @@ mod tests { postgres: PostgresConfig { enabled: Some(false), connection_pool_size: 20, + inference_retention_days: None, }, rate_limiting: Default::default(), ..Default::default() diff --git a/tensorzero-core/tests/e2e/db/inference_count_queries.rs b/tensorzero-core/tests/e2e/db/inference_count_queries.rs index 3288be73edc..0cb7c58dc9e 100644 --- a/tensorzero-core/tests/e2e/db/inference_count_queries.rs +++ b/tensorzero-core/tests/e2e/db/inference_count_queries.rs @@ -1,5 +1,9 @@ -//! E2E tests for inference count ClickHouse queries. +//! Shared test logic for InferenceCountQueries implementations (ClickHouse and Postgres). +//! +//! Each test function accepts a connection implementing `InferenceCountQueries`. +//! Tests use `>=` assertions since ClickHouse may accumulate data from other tests. +use crate::db::get_test_postgres; use tensorzero_core::db::clickhouse::test_helpers::get_clickhouse; use tensorzero_core::db::inference_count::{ CountInferencesParams, CountInferencesWithDemonstrationFeedbacksParams, @@ -10,148 +14,145 @@ use tensorzero_core::{ function::FunctionConfigType, }; -#[tokio::test] -async fn test_count_inferences_for_chat_function() { - let clickhouse = get_clickhouse().await; +/// Generates test functions for both ClickHouse and Postgres backends. +macro_rules! make_db_test { + ($test_name:ident) => { + paste::paste! { + #[tokio::test] + async fn [<$test_name _clickhouse>]() { + let conn = get_clickhouse().await; + $test_name(conn).await; + } + + #[tokio::test] + async fn [<$test_name _postgres>]() { + let conn = get_test_postgres().await; + $test_name(conn).await; + } + } + }; +} + +/// Generates test functions for ClickHouse only. +macro_rules! make_clickhouse_only_test { + ($test_name:ident) => { + paste::paste! { + #[tokio::test] + async fn [<$test_name _clickhouse>]() { + let conn = get_clickhouse().await; + $test_name(conn).await; + } + } + }; +} +// ===== SHARED TEST IMPLEMENTATIONS ===== + +async fn test_count_inferences_for_chat_function(conn: impl InferenceCountQueries) { let params = CountInferencesParams { function_name: "write_haiku", function_type: FunctionConfigType::Chat, variant_name: None, }; - let count = clickhouse - .count_inferences_for_function(params) - .await - .unwrap(); + let count = conn.count_inferences_for_function(params).await.unwrap(); - // The test data has at least 804 inferences for write_haiku (base fixture data) - // The count may be higher due to other tests adding more inferences + // The fixture data has 804 inferences for write_haiku + // ClickHouse may have more due to other tests assert!( count >= 804, "Expected at least 804 inferences for write_haiku, got {count}" ); } +make_db_test!(test_count_inferences_for_chat_function); -#[tokio::test] -async fn test_count_inferences_for_json_function() { - let clickhouse = get_clickhouse().await; - +async fn test_count_inferences_for_json_function(conn: impl InferenceCountQueries) { let params = CountInferencesParams { function_name: "extract_entities", function_type: FunctionConfigType::Json, variant_name: None, }; - let count = clickhouse - .count_inferences_for_function(params) - .await - .unwrap(); + let count = conn.count_inferences_for_function(params).await.unwrap(); - // The test data has at least 604 inferences for extract_entities (base fixture data) - // The count may be higher due to other tests adding more inferences + // The fixture data has 604 inferences for extract_entities assert!( count >= 604, "Expected at least 604 inferences for extract_entities, got {count}" ); } +make_db_test!(test_count_inferences_for_json_function); -#[tokio::test] -async fn test_count_inferences_for_chat_function_with_variant() { - let clickhouse = get_clickhouse().await; - +async fn test_count_inferences_for_chat_function_with_variant(conn: impl InferenceCountQueries) { let params = CountInferencesParams { function_name: "write_haiku", function_type: FunctionConfigType::Chat, variant_name: Some("initial_prompt_gpt4o_mini"), }; - let count = clickhouse - .count_inferences_for_function(params) - .await - .unwrap(); + let count = conn.count_inferences_for_function(params).await.unwrap(); - // The test data has at least 649 inferences for write_haiku with variant initial_prompt_gpt4o_mini + // The fixture data has 649 inferences for write_haiku with variant initial_prompt_gpt4o_mini assert!( count >= 649, "Expected at least 649 inferences for write_haiku/initial_prompt_gpt4o_mini, got {count}" ); } +make_db_test!(test_count_inferences_for_chat_function_with_variant); -#[tokio::test] -async fn test_count_inferences_for_json_function_with_variant() { - let clickhouse = get_clickhouse().await; - +async fn test_count_inferences_for_json_function_with_variant(conn: impl InferenceCountQueries) { let params = CountInferencesParams { function_name: "extract_entities", function_type: FunctionConfigType::Json, variant_name: Some("gpt4o_initial_prompt"), }; - let count = clickhouse - .count_inferences_for_function(params) - .await - .unwrap(); + let count = conn.count_inferences_for_function(params).await.unwrap(); - // The test data has at least 132 inferences for extract_entities with variant gpt4o_initial_prompt + // The fixture data has 132 inferences for extract_entities with variant gpt4o_initial_prompt assert!( count >= 132, "Expected at least 132 inferences for extract_entities/gpt4o_initial_prompt, got {count}" ); } +make_db_test!(test_count_inferences_for_json_function_with_variant); -#[tokio::test] -async fn test_count_inferences_for_nonexistent_function() { - let clickhouse = get_clickhouse().await; - +async fn test_count_inferences_for_nonexistent_function(conn: impl InferenceCountQueries) { let params = CountInferencesParams { function_name: "nonexistent_function", function_type: FunctionConfigType::Chat, variant_name: None, }; - let count = clickhouse - .count_inferences_for_function(params) - .await - .unwrap(); + let count = conn.count_inferences_for_function(params).await.unwrap(); - // Should return 0 for nonexistent function - assert_eq!(count, 0); + assert_eq!(count, 0, "Expected 0 for nonexistent function"); } +make_db_test!(test_count_inferences_for_nonexistent_function); -#[tokio::test] -async fn test_count_inferences_for_nonexistent_variant() { - let clickhouse = get_clickhouse().await; - +async fn test_count_inferences_for_nonexistent_variant(conn: impl InferenceCountQueries) { let params = CountInferencesParams { function_name: "write_haiku", function_type: FunctionConfigType::Chat, variant_name: Some("nonexistent_variant"), }; - let count = clickhouse - .count_inferences_for_function(params) - .await - .unwrap(); + let count = conn.count_inferences_for_function(params).await.unwrap(); - // Should return 0 for nonexistent variant - assert_eq!(count, 0); + assert_eq!(count, 0, "Expected 0 for nonexistent variant"); } +make_db_test!(test_count_inferences_for_nonexistent_variant); -#[tokio::test] -async fn test_count_inferences_by_variant_for_chat_function() { - let clickhouse = get_clickhouse().await; - +async fn test_count_inferences_by_variant_for_chat_function( + conn: impl InferenceCountQueries + Clone, +) { let params = CountInferencesParams { function_name: "write_haiku", function_type: FunctionConfigType::Chat, variant_name: None, }; - let rows = clickhouse - .count_inferences_by_variant(params) - .await - .unwrap(); + let rows = conn.count_inferences_by_variant(params).await.unwrap(); // There should be at least 2 variants for write_haiku assert!( @@ -168,7 +169,7 @@ async fn test_count_inferences_by_variant_for_chat_function() { function_type: FunctionConfigType::Chat, variant_name: None, }; - let total_count = clickhouse + let total_count = conn .count_inferences_for_function(total_params) .await .unwrap(); @@ -186,7 +187,7 @@ async fn test_count_inferences_by_variant_for_chat_function() { ); } - // Each row should have a valid last_used timestamp + // Each row should have a valid last_used timestamp in RFC 3339 format let rfc3339_millis_regex = regex::Regex::new(r"^\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}\.\d{3}Z$").unwrap(); for row in &rows { @@ -198,21 +199,16 @@ async fn test_count_inferences_by_variant_for_chat_function() { ); } } +make_db_test!(test_count_inferences_by_variant_for_chat_function); -#[tokio::test] -async fn test_count_inferences_by_variant_for_json_function() { - let clickhouse = get_clickhouse().await; - +async fn test_count_inferences_by_variant_for_json_function(conn: impl InferenceCountQueries) { let params = CountInferencesParams { function_name: "extract_entities", function_type: FunctionConfigType::Json, variant_name: None, }; - let rows = clickhouse - .count_inferences_by_variant(params) - .await - .unwrap(); + let rows = conn.count_inferences_by_variant(params).await.unwrap(); // There should be at least 2 variants for extract_entities assert!( @@ -228,31 +224,80 @@ async fn test_count_inferences_by_variant_for_json_function() { "Expected gpt4o_initial_prompt variant to be present" ); } +make_db_test!(test_count_inferences_by_variant_for_json_function); -#[tokio::test] -async fn test_count_inferences_by_variant_for_nonexistent_function() { - let clickhouse = get_clickhouse().await; - +async fn test_count_inferences_by_variant_for_nonexistent_function( + conn: impl InferenceCountQueries, +) { let params = CountInferencesParams { function_name: "nonexistent_function", function_type: FunctionConfigType::Chat, variant_name: None, }; - let rows = clickhouse - .count_inferences_by_variant(params) - .await - .unwrap(); + let rows = conn.count_inferences_by_variant(params).await.unwrap(); + + assert!( + rows.is_empty(), + "Expected empty result for nonexistent function" + ); +} +make_db_test!(test_count_inferences_by_variant_for_nonexistent_function); + +async fn test_list_functions_with_inference_count(conn: impl InferenceCountQueries) { + let rows = conn.list_functions_with_inference_count().await.unwrap(); + + // Should return multiple functions + assert!( + rows.len() >= 2, + "Expected at least 2 functions, got {}", + rows.len() + ); + + // write_haiku is a chat function, extract_entities is a json function + // Both should be present in the results + let write_haiku = rows.iter().find(|r| r.function_name == "write_haiku"); + let extract_entities = rows.iter().find(|r| r.function_name == "extract_entities"); + + // Verify inference_counts are reasonable based on test fixtures + let write_haiku = write_haiku.expect("Chat function write_haiku should be present"); + let extract_entities = + extract_entities.expect("Json function extract_entities should be present"); + assert!( + write_haiku.inference_count >= 804, + "Expected at least 804 inferences for write_haiku, got {}", + write_haiku.inference_count + ); + assert!( + extract_entities.inference_count >= 604, + "Expected at least 604 inferences for extract_entities, got {}", + extract_entities.inference_count + ); + + // Results should be ordered by last_inference_timestamp DESC + for i in 1..rows.len() { + assert!( + rows[i - 1].last_inference_timestamp >= rows[i].last_inference_timestamp, + "Results should be ordered by last_inference_timestamp DESC" + ); + } - // Should return empty for nonexistent function - assert!(rows.is_empty()); + // Each row should have a positive inference_count + for row in &rows { + assert!( + row.inference_count > 0, + "Each function should have at least one inference, {} has inference_count {}", + row.function_name, + row.inference_count + ); + } } +make_db_test!(test_list_functions_with_inference_count); -/// Test counting feedbacks for a float metric -#[tokio::test] -async fn test_count_feedbacks_for_float_metric() { - let clickhouse = get_clickhouse().await; +// ===== CLICKHOUSE-ONLY TEST IMPLEMENTATIONS ===== +// These tests require feedback tables which are not yet implemented for Postgres. +async fn test_count_feedbacks_for_float_metric(conn: impl InferenceCountQueries) { let metric_config = MetricConfig { r#type: MetricConfigType::Float, optimize: MetricConfigOptimize::Max, @@ -268,23 +313,17 @@ async fn test_count_feedbacks_for_float_metric() { metric_threshold: None, }; - let count = clickhouse - .count_inferences_with_feedback(params) - .await - .unwrap(); + let count = conn.count_inferences_with_feedback(params).await.unwrap(); - // The test database should have some haiku_rating feedbacks assert!( count > 0, "Should have feedbacks for haiku_rating metric on write_haiku" ); } +// TODO(#5691): Implement feedback tables for Postgres and switch to make_db_test +make_clickhouse_only_test!(test_count_feedbacks_for_float_metric); -/// Test counting feedbacks for a boolean metric -#[tokio::test] -async fn test_count_feedbacks_for_boolean_metric() { - let clickhouse = get_clickhouse().await; - +async fn test_count_feedbacks_for_boolean_metric(conn: impl InferenceCountQueries) { let metric_config = MetricConfig { r#type: MetricConfigType::Boolean, optimize: MetricConfigOptimize::Max, @@ -300,23 +339,17 @@ async fn test_count_feedbacks_for_boolean_metric() { metric_threshold: None, }; - // Query executes successfully, count may be 0 if no feedbacks exist - let count = clickhouse - .count_inferences_with_feedback(params) - .await - .unwrap(); + let count = conn.count_inferences_with_feedback(params).await.unwrap(); assert!( count > 0, "Should have feedbacks for exact_match metric on extract_entities" ); } +// TODO(#5691): Implement feedback tables for Postgres and switch to make_db_test +make_clickhouse_only_test!(test_count_feedbacks_for_boolean_metric); -/// Test counting inferences with feedback meeting threshold for a float metric -#[tokio::test] -async fn test_count_inferences_with_threshold_float_metric() { - let clickhouse = get_clickhouse().await; - +async fn test_count_inferences_with_threshold_float_metric(conn: impl InferenceCountQueries) { let metric_config = MetricConfig { r#type: MetricConfigType::Float, optimize: MetricConfigOptimize::Max, @@ -333,7 +366,7 @@ async fn test_count_inferences_with_threshold_float_metric() { metric_threshold: None, }; - let total_feedbacks = clickhouse + let total_feedbacks = conn .count_inferences_with_feedback(total_params) .await .unwrap(); @@ -347,7 +380,7 @@ async fn test_count_inferences_with_threshold_float_metric() { metric_threshold: Some(0.5), }; - let threshold_count = clickhouse + let threshold_count = conn .count_inferences_with_feedback(threshold_params) .await .unwrap(); @@ -358,12 +391,10 @@ async fn test_count_inferences_with_threshold_float_metric() { "Threshold count should be < total feedbacks" ); } +// TODO(#5691): Implement feedback tables for Postgres and switch to make_db_test +make_clickhouse_only_test!(test_count_inferences_with_threshold_float_metric); -/// Test counting inferences with feedback meeting threshold for a boolean metric (optimize max) -#[tokio::test] -async fn test_count_inferences_with_threshold_boolean_metric_max() { - let clickhouse = get_clickhouse().await; - +async fn test_count_inferences_with_threshold_boolean_metric_max(conn: impl InferenceCountQueries) { let metric_config = MetricConfig { r#type: MetricConfigType::Boolean, optimize: MetricConfigOptimize::Max, @@ -380,7 +411,7 @@ async fn test_count_inferences_with_threshold_boolean_metric_max() { metric_threshold: None, }; - let total_feedbacks = clickhouse + let total_feedbacks = conn .count_inferences_with_feedback(total_params) .await .unwrap(); @@ -394,7 +425,7 @@ async fn test_count_inferences_with_threshold_boolean_metric_max() { metric_threshold: Some(0.0), // Not used for boolean, but required }; - let threshold_count = clickhouse + let threshold_count = conn .count_inferences_with_feedback(threshold_params) .await .unwrap(); @@ -405,30 +436,26 @@ async fn test_count_inferences_with_threshold_boolean_metric_max() { "Threshold count should be < total feedbacks" ); } +// TODO(#5691): Implement feedback tables for Postgres and switch to make_db_test +make_clickhouse_only_test!(test_count_inferences_with_threshold_boolean_metric_max); -/// Test counting demonstration feedbacks for a function -#[tokio::test] -async fn test_count_demonstration_feedbacks() { - let clickhouse = get_clickhouse().await; - +async fn test_count_demonstration_feedbacks(conn: impl InferenceCountQueries) { let params = CountInferencesWithDemonstrationFeedbacksParams { function_name: "extract_entities", function_type: FunctionConfigType::Json, }; - let count = clickhouse + let count = conn .count_inferences_with_demonstration_feedback(params) .await .unwrap(); assert!(count > 0, "Should have demonstrations for extract_entities"); } +// TODO(#5691): Implement feedback tables for Postgres and switch to make_db_test +make_clickhouse_only_test!(test_count_demonstration_feedbacks); -/// Test with episode-level metric -#[tokio::test] -async fn test_count_feedbacks_for_episode_level_boolean_metric() { - let clickhouse = get_clickhouse().await; - +async fn test_count_feedbacks_for_episode_level_boolean_metric(conn: impl InferenceCountQueries) { let metric_config = MetricConfig { r#type: MetricConfigType::Boolean, optimize: MetricConfigOptimize::Max, @@ -444,10 +471,7 @@ async fn test_count_feedbacks_for_episode_level_boolean_metric() { metric_threshold: None, }; - let count = clickhouse - .count_inferences_with_feedback(params) - .await - .unwrap(); + let count = conn.count_inferences_with_feedback(params).await.unwrap(); // We're verifying the query executes successfully with episode-level join assert!( @@ -455,11 +479,10 @@ async fn test_count_feedbacks_for_episode_level_boolean_metric() { "Should have feedbacks for boolean metric exact_match_episode on extract_entities" ); } +// TODO(#5691): Implement feedback tables for Postgres and switch to make_db_test +make_clickhouse_only_test!(test_count_feedbacks_for_episode_level_boolean_metric); -#[tokio::test] -async fn test_count_feedbacks_for_episode_level_float_metric() { - let clickhouse = get_clickhouse().await; - +async fn test_count_feedbacks_for_episode_level_float_metric(conn: impl InferenceCountQueries) { let metric_config = MetricConfig { r#type: MetricConfigType::Float, optimize: MetricConfigOptimize::Max, @@ -475,10 +498,7 @@ async fn test_count_feedbacks_for_episode_level_float_metric() { metric_threshold: None, }; - let count = clickhouse - .count_inferences_with_feedback(params) - .await - .unwrap(); + let count = conn.count_inferences_with_feedback(params).await.unwrap(); // We're verifying the query executes successfully with episode-level join assert!( @@ -486,58 +506,5 @@ async fn test_count_feedbacks_for_episode_level_float_metric() { "Should have feedbacks for float metric jaccard_similarity_episode on extract_entities" ); } - -/// Test list_functions_with_inference_count returns expected functions -#[tokio::test] -async fn test_list_functions_with_inference_count() { - let clickhouse = get_clickhouse().await; - - let rows = clickhouse - .list_functions_with_inference_count() - .await - .unwrap(); - - // Should return multiple functions - assert!( - rows.len() >= 2, - "Expected at least 2 functions, got {}", - rows.len() - ); - - // write_haiku is a chat function, extract_entities is a json function - // Both should be present in the results - let write_haiku = rows.iter().find(|r| r.function_name == "write_haiku"); - let extract_entities = rows.iter().find(|r| r.function_name == "extract_entities"); - - // Verify inference_counts are reasonable based on test fixtures - let write_haiku = write_haiku.expect("Chat function write_haiku should be present"); - let extract_entities = - extract_entities.expect("Json function extract_entities should be present"); - assert!( - write_haiku.inference_count >= 804, - "Expected at least 804 inferences for write_haiku, got {}", - write_haiku.inference_count - ); - assert!( - extract_entities.inference_count >= 604, - "Expected at least 604 inferences for extract_entities, got {}", - extract_entities.inference_count - ); - - // Results should be ordered by last_inference_timestamp DESC - for i in 1..rows.len() { - assert!( - rows[i - 1].last_inference_timestamp >= rows[i].last_inference_timestamp, - "Results should be ordered by last_inference_timestamp DESC" - ); - } - // Each row should have a positive inference_count - for row in &rows { - assert!( - row.inference_count > 0, - "Each function should have at least one inference, {} has inference_count {}", - row.function_name, - row.inference_count - ); - } -} +// TODO(#5691): Implement feedback tables for Postgres and switch to make_db_test +make_clickhouse_only_test!(test_count_feedbacks_for_episode_level_float_metric); diff --git a/tensorzero-core/tests/e2e/db/mod.rs b/tensorzero-core/tests/e2e/db/mod.rs index 4005b98b8a0..554f778be1b 100644 --- a/tensorzero-core/tests/e2e/db/mod.rs +++ b/tensorzero-core/tests/e2e/db/mod.rs @@ -1,3 +1,10 @@ +//! Database E2E tests and helpers. +//! +//! This module provides connection helpers for e2e tests against ClickHouse and Postgres. + +use sqlx::postgres::PgPoolOptions; +use tensorzero_core::db::postgres::PostgresConnectionInfo; + mod bandit_queries; mod batch_inference_queries; mod dataset_queries; @@ -9,3 +16,19 @@ mod postgres; mod rate_limit_queries; mod select_queries; mod workflow_evaluation_queries; + +// ===== CONNECTION HELPERS ===== + +pub async fn get_test_postgres() -> PostgresConnectionInfo { + let postgres_url = std::env::var("TENSORZERO_POSTGRES_URL") + .expect("Environment variable TENSORZERO_POSTGRES_URL must be set"); + + let start = std::time::Instant::now(); + println!("Connecting to PostgreSQL"); + let pool = PgPoolOptions::new() + .connect(&postgres_url) + .await + .expect("Failed to connect to PostgreSQL"); + println!("Connected to PostgreSQL in {:?}", start.elapsed()); + PostgresConnectionInfo::new_with_pool(pool) +} diff --git a/tensorzero-core/tests/e2e/db/postgres/mod.rs b/tensorzero-core/tests/e2e/db/postgres/mod.rs index 9b31281e367..db61b1d0490 100644 --- a/tensorzero-core/tests/e2e/db/postgres/mod.rs +++ b/tensorzero-core/tests/e2e/db/postgres/mod.rs @@ -1 +1,2 @@ mod experimentation_queries; +mod postgres_function_tests; diff --git a/tensorzero-core/tests/e2e/db/postgres/postgres_function_tests.rs b/tensorzero-core/tests/e2e/db/postgres/postgres_function_tests.rs new file mode 100644 index 00000000000..984566f2eca --- /dev/null +++ b/tensorzero-core/tests/e2e/db/postgres/postgres_function_tests.rs @@ -0,0 +1,416 @@ +//! Tests for Postgres functions in the tensorzero schema. + +use crate::db::get_test_postgres; +use chrono::{Days, TimeZone, Utc}; +use sqlx::QueryBuilder; +use uuid::{NoContext, Timestamp, Uuid}; + +/// Tests that the Postgres `uuid_v7_to_timestamp` function extracts the same timestamp +/// that Rust's `uuid` crate encodes into a UUIDv7. +#[tokio::test] +async fn test_uuid_v7_to_timestamp() { + let conn = get_test_postgres().await; + let pool = conn.get_pool().expect("Pool should be available"); + + // Generate a UUIDv7 in Rust + let rust_uuid = Uuid::now_v7(); + + // Extract timestamp using Postgres function + let (postgres_timestamp,): (chrono::DateTime,) = + sqlx::query_as("SELECT tensorzero.uuid_v7_to_timestamp($1::uuid)") + .bind(rust_uuid) + .fetch_one(pool) + .await + .expect("Query should succeed"); + + // Generate a new UUIDv7 using the timestamp extracted by Postgres + let postgres_millis = postgres_timestamp.timestamp_millis() as u64; + let reconstructed_uuid = Uuid::new_v7(Timestamp::from_unix( + NoContext, + postgres_millis / 1000, + (postgres_millis % 1000) as u32 * 1_000_000, + )); + + // The first 48 bits (timestamp portion) should match + let rust_timestamp_bits = &rust_uuid.as_bytes()[..6]; + let reconstructed_timestamp_bits = &reconstructed_uuid.as_bytes()[..6]; + + assert_eq!( + rust_timestamp_bits, reconstructed_timestamp_bits, + "Postgres-extracted timestamp should match the original UUIDv7 timestamp" + ); +} + +/// Tests the function with a known timestamp value. +#[tokio::test] +async fn test_uuid_v7_to_timestamp_known_value() { + let conn = get_test_postgres().await; + let pool = conn.get_pool().expect("Pool should be available"); + + // Create a UUIDv7 with a known timestamp: Jan 1, 2024 00:00:00.123 UTC + let known_time = + Utc.with_ymd_and_hms(2024, 1, 1, 0, 0, 0).unwrap() + chrono::Duration::milliseconds(123); + let secs = known_time.timestamp() as u64; + let nanos = known_time.timestamp_subsec_nanos(); + let rust_uuid = Uuid::new_v7(Timestamp::from_unix(NoContext, secs, nanos)); + + // Extract timestamp using Postgres function + let (postgres_timestamp,): (chrono::DateTime,) = + sqlx::query_as("SELECT tensorzero.uuid_v7_to_timestamp($1::uuid)") + .bind(rust_uuid) + .fetch_one(pool) + .await + .expect("Query should succeed"); + + // Postgres should return the same millisecond-precision timestamp + assert_eq!( + postgres_timestamp.timestamp_millis(), + known_time.timestamp_millis(), + "Postgres should extract the correct timestamp from UUIDv7" + ); +} + +const TEST_PARTITIONED_TABLE: &str = "test_partitioned_table"; +const TEST_RETENTION_TABLE: &str = "test_retention_table"; +const TEST_RETENTION_KEY: &str = "test_retention_days"; + +/// Tests that `create_partitions` creates partitions for the next 8 days. +#[tokio::test] +async fn test_create_partitions() { + let conn = get_test_postgres().await; + let pool = conn.get_pool().expect("Pool should be available"); + + // Clean up any existing test table and its partitions + let mut qb: QueryBuilder = + QueryBuilder::new("DROP TABLE IF EXISTS tensorzero."); + qb.push(TEST_PARTITIONED_TABLE); + qb.push(" CASCADE"); + qb.build() + .execute(pool) + .await + .expect("Cleanup should succeed"); + + // Create a partitioned table + let mut qb: QueryBuilder = QueryBuilder::new("CREATE TABLE tensorzero."); + qb.push(TEST_PARTITIONED_TABLE); + qb.push(" (id UUID NOT NULL, created_at DATE NOT NULL, PRIMARY KEY (id, created_at)) PARTITION BY RANGE (created_at)"); + qb.build() + .execute(pool) + .await + .expect("Table creation should succeed"); + + // Call create_partitions + sqlx::query("SELECT tensorzero.create_partitions($1)") + .bind(TEST_PARTITIONED_TABLE) + .execute(pool) + .await + .expect("create_partitions should succeed"); + + // Verify partitions were created for days 0..7 (8 partitions) + let (partition_count,): (i64,) = sqlx::query_as( + "SELECT COUNT(*) FROM pg_tables WHERE schemaname = 'tensorzero' AND tablename LIKE $1", + ) + .bind(format!("{TEST_PARTITIONED_TABLE}_%")) + .fetch_one(pool) + .await + .expect("Count query should succeed"); + + assert_eq!( + partition_count, 8, + "create_partitions should create 8 partitions (today + 7 future days)" + ); + + // Verify the partition names match expected dates + let today = Utc::now().date_naive(); + for i in 0..8 { + let partition_date = today + Days::new(i); + let expected_partition = format!( + "{}_{}", + TEST_PARTITIONED_TABLE, + partition_date.format("%Y_%m_%d") + ); + + let (exists,): (bool,) = sqlx::query_as( + "SELECT EXISTS(SELECT 1 FROM pg_tables WHERE schemaname = 'tensorzero' AND tablename = $1)", + ) + .bind(&expected_partition) + .fetch_one(pool) + .await + .expect("Existence check should succeed"); + + assert!( + exists, + "Partition {expected_partition} should exist for day offset {i}" + ); + } + + // Clean up + let mut qb: QueryBuilder = + QueryBuilder::new("DROP TABLE IF EXISTS tensorzero."); + qb.push(TEST_PARTITIONED_TABLE); + qb.push(" CASCADE"); + qb.build() + .execute(pool) + .await + .expect("Cleanup should succeed"); +} + +/// Tests that `drop_old_partitions` drops partitions older than the retention period. +#[tokio::test] +async fn test_drop_old_partitions() { + let conn = get_test_postgres().await; + let pool = conn.get_pool().expect("Pool should be available"); + + // Clean up any existing test table and retention config + let mut qb: QueryBuilder = + QueryBuilder::new("DROP TABLE IF EXISTS tensorzero."); + qb.push(TEST_RETENTION_TABLE); + qb.push(" CASCADE"); + qb.build() + .execute(pool) + .await + .expect("Cleanup should succeed"); + + sqlx::query("DELETE FROM tensorzero.retention_config WHERE key = $1") + .bind(TEST_RETENTION_KEY) + .execute(pool) + .await + .expect("Retention config cleanup should succeed"); + + // Create a partitioned table + let mut qb: QueryBuilder = QueryBuilder::new("CREATE TABLE tensorzero."); + qb.push(TEST_RETENTION_TABLE); + qb.push(" (id UUID NOT NULL, created_at DATE NOT NULL, PRIMARY KEY (id, created_at)) PARTITION BY RANGE (created_at)"); + qb.build() + .execute(pool) + .await + .expect("Table creation should succeed"); + + // Manually create partitions: some old (30+ days ago), some recent + let today = Utc::now().date_naive(); + let old_dates = [ + today - Days::new(60), + today - Days::new(45), + today - Days::new(31), + ]; + let recent_dates = [today - Days::new(5), today - Days::new(1), today]; + + for date in old_dates.iter().chain(recent_dates.iter()) { + let partition_name = format!("{}_{}", TEST_RETENTION_TABLE, date.format("%Y_%m_%d")); + let next_date = *date + Days::new(1); + + let mut qb: QueryBuilder = + QueryBuilder::new("CREATE TABLE IF NOT EXISTS tensorzero."); + qb.push(&partition_name); + qb.push(" PARTITION OF tensorzero."); + qb.push(TEST_RETENTION_TABLE); + qb.push(" FOR VALUES FROM ('"); + qb.push(date.to_string()); + qb.push("') TO ('"); + qb.push(next_date.to_string()); + qb.push("')"); + qb.build() + .execute(pool) + .await + .expect("Partition creation should succeed"); + } + + // Verify all 6 partitions exist + let (initial_count,): (i64,) = sqlx::query_as( + "SELECT COUNT(*) FROM pg_tables WHERE schemaname = 'tensorzero' AND tablename LIKE $1", + ) + .bind(format!("{TEST_RETENTION_TABLE}_%")) + .fetch_one(pool) + .await + .expect("Count query should succeed"); + + assert_eq!(initial_count, 6, "Should have 6 partitions initially"); + + // Set retention to 30 days + sqlx::query( + "INSERT INTO tensorzero.retention_config (key, value) VALUES ($1, '30') ON CONFLICT (key) DO UPDATE SET value = '30'", + ) + .bind(TEST_RETENTION_KEY) + .execute(pool) + .await + .expect("Retention config insert should succeed"); + + // Call drop_old_partitions + sqlx::query("SELECT tensorzero.drop_old_partitions($1, $2)") + .bind(TEST_RETENTION_TABLE) + .bind(TEST_RETENTION_KEY) + .execute(pool) + .await + .expect("drop_old_partitions should succeed"); + + // Verify old partitions were dropped (only 3 recent should remain) + let (final_count,): (i64,) = sqlx::query_as( + "SELECT COUNT(*) FROM pg_tables WHERE schemaname = 'tensorzero' AND tablename LIKE $1", + ) + .bind(format!("{TEST_RETENTION_TABLE}_%")) + .fetch_one(pool) + .await + .expect("Count query should succeed"); + + assert_eq!( + final_count, 3, + "Should have 3 partitions after dropping old ones (30+ days old)" + ); + + // Verify the old partitions are gone + for date in &old_dates { + let partition_name = format!("{}_{}", TEST_RETENTION_TABLE, date.format("%Y_%m_%d")); + let (exists,): (bool,) = sqlx::query_as( + "SELECT EXISTS(SELECT 1 FROM pg_tables WHERE schemaname = 'tensorzero' AND tablename = $1)", + ) + .bind(&partition_name) + .fetch_one(pool) + .await + .expect("Existence check should succeed"); + + assert!( + !exists, + "Old partition {partition_name} should have been dropped" + ); + } + + // Verify recent partitions still exist + for date in &recent_dates { + let partition_name = format!("{}_{}", TEST_RETENTION_TABLE, date.format("%Y_%m_%d")); + let (exists,): (bool,) = sqlx::query_as( + "SELECT EXISTS(SELECT 1 FROM pg_tables WHERE schemaname = 'tensorzero' AND tablename = $1)", + ) + .bind(&partition_name) + .fetch_one(pool) + .await + .expect("Existence check should succeed"); + + assert!( + exists, + "Recent partition {partition_name} should still exist" + ); + } + + // Clean up + let mut qb: QueryBuilder = + QueryBuilder::new("DROP TABLE IF EXISTS tensorzero."); + qb.push(TEST_RETENTION_TABLE); + qb.push(" CASCADE"); + qb.build() + .execute(pool) + .await + .expect("Cleanup should succeed"); + + sqlx::query("DELETE FROM tensorzero.retention_config WHERE key = $1") + .bind(TEST_RETENTION_KEY) + .execute(pool) + .await + .expect("Retention config cleanup should succeed"); +} + +const TEST_NO_RETENTION_TABLE: &str = "test_no_retention_table"; +const TEST_MISSING_RETENTION_KEY: &str = "test_missing_retention_key"; + +/// Tests that `drop_old_partitions` does nothing when the retention key is not configured. +#[tokio::test] +async fn test_drop_old_partitions_skips_when_retention_not_configured() { + let conn = get_test_postgres().await; + let pool = conn.get_pool().expect("Pool should be available"); + + // Ensure the retention key does not exist + sqlx::query("DELETE FROM tensorzero.retention_config WHERE key = $1") + .bind(TEST_MISSING_RETENTION_KEY) + .execute(pool) + .await + .expect("Retention config cleanup should succeed"); + + // Clean up any existing test table + let mut qb: QueryBuilder = + QueryBuilder::new("DROP TABLE IF EXISTS tensorzero."); + qb.push(TEST_NO_RETENTION_TABLE); + qb.push(" CASCADE"); + qb.build() + .execute(pool) + .await + .expect("Cleanup should succeed"); + + // Create a partitioned table + let mut qb: QueryBuilder = QueryBuilder::new("CREATE TABLE tensorzero."); + qb.push(TEST_NO_RETENTION_TABLE); + qb.push(" (id UUID NOT NULL, created_at DATE NOT NULL, PRIMARY KEY (id, created_at)) PARTITION BY RANGE (created_at)"); + qb.build() + .execute(pool) + .await + .expect("Table creation should succeed"); + + // Create some old partitions that would be dropped if retention were configured + let today = Utc::now().date_naive(); + let old_dates = [ + today - Days::new(60), + today - Days::new(45), + today - Days::new(31), + ]; + + for date in &old_dates { + let partition_name = format!("{}_{}", TEST_NO_RETENTION_TABLE, date.format("%Y_%m_%d")); + let next_date = *date + Days::new(1); + + let mut qb: QueryBuilder = + QueryBuilder::new("CREATE TABLE IF NOT EXISTS tensorzero."); + qb.push(&partition_name); + qb.push(" PARTITION OF tensorzero."); + qb.push(TEST_NO_RETENTION_TABLE); + qb.push(" FOR VALUES FROM ('"); + qb.push(date.to_string()); + qb.push("') TO ('"); + qb.push(next_date.to_string()); + qb.push("')"); + qb.build() + .execute(pool) + .await + .expect("Partition creation should succeed"); + } + + // Verify all 3 partitions exist before calling drop_old_partitions + let (initial_count,): (i64,) = sqlx::query_as( + "SELECT COUNT(*) FROM pg_tables WHERE schemaname = 'tensorzero' AND tablename LIKE $1", + ) + .bind(format!("{TEST_NO_RETENTION_TABLE}_%")) + .fetch_one(pool) + .await + .expect("Count query should succeed"); + + assert_eq!(initial_count, 3, "Should have 3 partitions initially"); + + // Call drop_old_partitions with a non-existent retention key + sqlx::query("SELECT tensorzero.drop_old_partitions($1, $2)") + .bind(TEST_NO_RETENTION_TABLE) + .bind(TEST_MISSING_RETENTION_KEY) + .execute(pool) + .await + .expect("drop_old_partitions should succeed even without retention config"); + + // Verify all partitions still exist (nothing was dropped) + let (final_count,): (i64,) = sqlx::query_as( + "SELECT COUNT(*) FROM pg_tables WHERE schemaname = 'tensorzero' AND tablename LIKE $1", + ) + .bind(format!("{TEST_NO_RETENTION_TABLE}_%")) + .fetch_one(pool) + .await + .expect("Count query should succeed"); + + assert_eq!( + final_count, 3, + "All partitions should remain when retention key is not configured" + ); + + // Clean up + let mut qb: QueryBuilder = + QueryBuilder::new("DROP TABLE IF EXISTS tensorzero."); + qb.push(TEST_NO_RETENTION_TABLE); + qb.push(" CASCADE"); + qb.build() + .execute(pool) + .await + .expect("Cleanup should succeed"); +} diff --git a/tensorzero-core/tests/e2e/docker-compose.live.yml b/tensorzero-core/tests/e2e/docker-compose.live.yml index 1dd060184f9..e0e9a23c694 100644 --- a/tensorzero-core/tests/e2e/docker-compose.live.yml +++ b/tensorzero-core/tests/e2e/docker-compose.live.yml @@ -143,6 +143,32 @@ services: retries: 72 # Retry for up to 6 minutes start_period: 5s # Give the script time to start + # fixtures-postgres and fixtures have exactly the same build context, so the fixture is downloaded only once and shared. + fixtures-postgres: + image: tensorzero/fixtures:${TENSORZERO_COMMIT_TAG:-latest} + build: + dockerfile: ui/fixtures/Dockerfile + context: ../../.. + volumes: + - ../../../ui/fixtures:/fixtures + - ~/.aws:/root/.aws + environment: + - TENSORZERO_POSTGRES_URL=postgres://postgres:postgres@postgres:5432/${TENSORZERO_E2E_TESTS_DATABASE:-tensorzero-e2e-tests} + - R2_ACCESS_KEY_ID + - R2_SECRET_ACCESS_KEY + depends_on: + postgres: + condition: service_healthy + gateway-postgres-migrations: + condition: service_healthy + command: [ "bash", "-c", "cd /fixtures && ./load_fixtures_postgres.sh && touch /load_complete_postgres.marker" ] + healthcheck: + test: [ "CMD", "test", "-f", "/load_complete_postgres.marker" ] + interval: 5s + timeout: 1s + retries: 72 # Retry for up to 6 minutes + start_period: 5s # Give the script time to start + live-tests: image: tensorzero/live-tests:${TENSORZERO_COMMIT_TAG} build: @@ -210,6 +236,8 @@ services: depends_on: fixtures: condition: service_completed_successfully + fixtures-postgres: + condition: service_completed_successfully clickhouse: condition: service_healthy postgres: diff --git a/ui/fixtures/Dockerfile b/ui/fixtures/Dockerfile index 437011e8804..99adc968e13 100644 --- a/ui/fixtures/Dockerfile +++ b/ui/fixtures/Dockerfile @@ -1,7 +1,7 @@ -# This Dockerfile is used to load fixtures into a (separate) ClickHouse server +# This Dockerfile is used to load fixtures into (separate) ClickHouse and Postgres servers FROM clickhouse:lts -RUN apt-get update && apt-get install -y cmake curl +RUN apt-get update && apt-get install -y cmake curl postgresql-client RUN ARCH=$(uname -m) && \ if [ "$ARCH" = "x86_64" ]; then \ AWSCLI_URL="https://awscli.amazonaws.com/awscli-exe-linux-x86_64.zip"; \ diff --git a/ui/fixtures/download-small-fixtures.py b/ui/fixtures/download-small-fixtures.py index 811a3738015..a6a42e87608 100644 --- a/ui/fixtures/download-small-fixtures.py +++ b/ui/fixtures/download-small-fixtures.py @@ -24,7 +24,7 @@ # and keep the local name unchanged (e.g., "file.jsonl") FIXTURES = { "model_inference_examples.jsonl": "model_inference_examples.jsonl", - "chat_inference_examples_20260120_201507.jsonl": "chat_inference_examples.jsonl", + "chat_inference_examples_20260123.jsonl": "chat_inference_examples.jsonl", "json_inference_examples.jsonl": "json_inference_examples.jsonl", "boolean_metric_feedback_examples.jsonl": "boolean_metric_feedback_examples.jsonl", "float_metric_feedback_examples.jsonl": "float_metric_feedback_examples.jsonl", diff --git a/ui/fixtures/load_fixtures_postgres.sh b/ui/fixtures/load_fixtures_postgres.sh index 0087c186852..f42e7a2cfd6 100755 --- a/ui/fixtures/load_fixtures_postgres.sh +++ b/ui/fixtures/load_fixtures_postgres.sh @@ -1,6 +1,142 @@ #!/bin/bash -set -euxo pipefail +set -euo pipefail -# Placeholder for actual postgres fixture loading script. -echo "Placeholder script" -exit 0 +# Load small fixtures into Postgres inference tables. +# Only loads chat_inferences and json_inferences (Step 1 tables). +# +# Usage: +# ./ui/fixtures/load_fixtures_postgres.sh +# +# Environment variables: +# TENSORZERO_POSTGRES_URL - Postgres connection URL (default: postgres://postgres:postgres@localhost:5432/tensorzero_ui_fixtures) +# TENSORZERO_SKIP_TRUNCATE - Set to 1 to skip truncating tables before loading + +POSTGRES_URL="${TENSORZERO_POSTGRES_URL:-postgres://postgres:postgres@localhost:5432/tensorzero_ui_fixtures}" +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" + +cd "$SCRIPT_DIR" + +# Helper function to load JSONL into a table via temp TEXT table +load_jsonl() { + local file="$1" + local table="$2" + local insert_sql="$3" + + if [ ! -f "$file" ]; then + echo "Warning: $file not found, skipping" + return + fi + + echo "Loading $file into $table..." + + psql -q "$POSTGRES_URL" <>'id')::uuid, + j->>'function_name', + j->>'variant_name', + (j->>'episode_id')::uuid, + COALESCE(j->'input', '{}')::jsonb, + COALESCE(j->'output', '{}')::jsonb, + COALESCE(j->'tool_params', '{}')::jsonb, + COALESCE(j->'inference_params', '{}')::jsonb, + (j->>'processing_time_ms')::integer, + (j->>'ttft_ms')::integer, + COALESCE(j->'tags', '{}')::jsonb, + COALESCE(j->'extra_body', '[]')::jsonb, + COALESCE(j->'dynamic_tools', '[]')::jsonb, + COALESCE(j->'dynamic_provider_tools', '[]')::jsonb, + j->'allowed_tools', + j->>'tool_choice', + (j->>'parallel_tool_calls')::boolean, + tensorzero.uuid_v7_to_timestamp((j->>'id')::uuid) +FROM tmp_jsonl, LATERAL (SELECT data::jsonb AS j) AS parsed +ON CONFLICT (id, created_at) DO NOTHING; +" + +# JSON Inferences +# Note: input, output, output_schema, inference_params, auxiliary_content are JSONB in our schema +# created_at is derived from the UUIDv7 id using tensorzero.uuid_v7_to_timestamp() +load_jsonl "json_inference_examples.jsonl" "tensorzero.json_inferences" " +INSERT INTO tensorzero.json_inferences ( + id, function_name, variant_name, episode_id, + input, output, output_schema, inference_params, + processing_time_ms, ttft_ms, tags, extra_body, auxiliary_content, created_at +) +SELECT + (j->>'id')::uuid, + j->>'function_name', + j->>'variant_name', + (j->>'episode_id')::uuid, + COALESCE(j->'input', '{}')::jsonb, + COALESCE(j->'output', '{}')::jsonb, + COALESCE(j->'output_schema', '{}')::jsonb, + COALESCE(j->'inference_params', '{}')::jsonb, + (j->>'processing_time_ms')::integer, + (j->>'ttft_ms')::integer, + COALESCE(j->'tags', '{}')::jsonb, + COALESCE(j->'extra_body', '[]')::jsonb, + COALESCE(j->'auxiliary_content', '{}')::jsonb, + tensorzero.uuid_v7_to_timestamp((j->>'id')::uuid) +FROM tmp_jsonl, LATERAL (SELECT data::jsonb AS j) AS parsed +ON CONFLICT (id, created_at) DO NOTHING; +" + +echo "" +echo "All fixtures loaded successfully!" + +# Print row counts +echo "" +echo "Row counts:" +psql -q "$POSTGRES_URL" < Resetting dev database: $DB_NAME" + +# Terminate existing connections and drop/recreate the database +PGPASSWORD="$DB_PASSWORD" psql -h "$DB_HOST" -p "$DB_PORT" -U "$DB_USER" -d postgres < pg_backend_pid(); + +DROP DATABASE IF EXISTS "$DB_NAME"; +CREATE DATABASE "$DB_NAME"; +EOF + +echo "==> Database recreated" + +# Run migrations +export TENSORZERO_POSTGRES_URL="postgres://$DB_USER:$DB_PASSWORD@$DB_HOST:$DB_PORT/$DB_NAME" +echo "==> Running migrations..." + +SQLX_OFFLINE=1 cargo run --package gateway --bin gateway -- --run-postgres-migrations + +./ui/fixtures/load_fixtures_postgres.sh + +echo "==> Done! Database $DB_NAME is ready." +echo " Connection: $TENSORZERO_POSTGRES_URL" diff --git a/ui/fixtures/upload-small-fixtures.sh b/ui/fixtures/upload-small-fixtures.sh index b1c93112aae..d2878543f1d 100755 --- a/ui/fixtures/upload-small-fixtures.sh +++ b/ui/fixtures/upload-small-fixtures.sh @@ -8,7 +8,7 @@ cd "$(dirname "$0")" # and update download-small-fixtures.py to use the new filename JSONL_FILES=( "model_inference_examples.jsonl" - "chat_inference_examples_20260120_201507.jsonl" + "chat_inference_examples_20260123.jsonl" "json_inference_examples.jsonl" "boolean_metric_feedback_examples.jsonl" "float_metric_feedback_examples.jsonl" From 516de1608fa5fce2a7757e5cfa3f17b77410ec26 Mon Sep 17 00:00:00 2001 From: Gabriel Bianconi <1275491+GabrielBianconi@users.noreply.github.com> Date: Sun, 25 Jan 2026 23:29:16 -0500 Subject: [PATCH 05/46] Address CI flakiness in test_rate_limiting_time_intervals (#5855) --- tensorzero-core/tests/e2e/rate_limiting.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tensorzero-core/tests/e2e/rate_limiting.rs b/tensorzero-core/tests/e2e/rate_limiting.rs index 88c62249879..8184d10f1f4 100644 --- a/tensorzero-core/tests/e2e/rate_limiting.rs +++ b/tensorzero-core/tests/e2e/rate_limiting.rs @@ -948,12 +948,12 @@ make_rate_limit_tests!(test_rate_limiting_time_intervals); async fn test_rate_limiting_time_intervals(backend: &str, stream: bool) { // Test different time intervals work correctly - // We use concurrent requests to avoid flakiness from token refill if the requests are not cached + // We use `refill_rate = 0` to avoid flakiness from tokens refilling if requests serialize due to database locking delays let id = Uuid::now_v7(); let config = generate_rate_limit_config( &[&format!( r#"[[rate_limiting.rules]] -model_inferences_per_second = {{ capacity = 2, refill_rate = 2 }} +model_inferences_per_second = {{ capacity = 2, refill_rate = 0 }} always = true scope = [ {{ tag_key = "test12_user_id_{id}", tag_value = "123" }} From 66efa2e05f0d60c72a310e5c653e2c3bf27a00ae Mon Sep 17 00:00:00 2001 From: Gabriel Bianconi <1275491+GabrielBianconi@users.noreply.github.com> Date: Mon, 26 Jan 2026 08:41:21 -0500 Subject: [PATCH 06/46] Bump migration timeout (#5857) --- ci/buildkite/test-clickhouse-cloud.sh | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ci/buildkite/test-clickhouse-cloud.sh b/ci/buildkite/test-clickhouse-cloud.sh index e564cdab544..403a220ccb3 100755 --- a/ci/buildkite/test-clickhouse-cloud.sh +++ b/ci/buildkite/test-clickhouse-cloud.sh @@ -95,7 +95,7 @@ cargo run --bin gateway --features e2e_tests -- --run-postgres-migrations cargo run-e2e > e2e_logs.txt 2>&1 & count=0 - max_attempts=90 + max_attempts=180 while ! curl -s -f http://localhost:3000/health >/dev/null 2>&1; do echo "Waiting for gateway to be healthy..." sleep 1 From 261945c37b987b90958c39e98723efd9f716eb86 Mon Sep 17 00:00:00 2001 From: Gabriel Bianconi <1275491+GabrielBianconi@users.noreply.github.com> Date: Mon, 26 Jan 2026 08:41:38 -0500 Subject: [PATCH 07/46] Add timeouts (so we get retries) to pytest (#5856) * Add timeouts (so we get retries) to pytest * Fix --- clients/python/pyproject.toml | 1 + clients/python/pytest.ini | 2 ++ clients/python/requirements.txt | 4 ++++ clients/python/uv.lock | 14 ++++++++++++++ 4 files changed, 21 insertions(+) diff --git a/clients/python/pyproject.toml b/clients/python/pyproject.toml index 05574627624..940187f8e2d 100644 --- a/clients/python/pyproject.toml +++ b/clients/python/pyproject.toml @@ -59,6 +59,7 @@ dev = [ "pytest-rerunfailures", "pytest", "pytest-httpserver", + "pytest-timeout", "openai", "clickhouse-connect", "pandas", diff --git a/clients/python/pytest.ini b/clients/python/pytest.ini index 0102b0a974e..bde480ba0b7 100644 --- a/clients/python/pytest.ini +++ b/clients/python/pytest.ini @@ -1,2 +1,4 @@ [pytest] asyncio_default_fixture_loop_scope = function +timeout = 60 +timeout_method = thread diff --git a/clients/python/requirements.txt b/clients/python/requirements.txt index 52fd035c3b0..0ac2fb6f249 100644 --- a/clients/python/requirements.txt +++ b/clients/python/requirements.txt @@ -720,6 +720,7 @@ pytest==8.3.5 \ # via # pytest-asyncio # pytest-rerunfailures + # pytest-timeout # pytest-xdist pytest-asyncio==0.26.0 \ --hash=sha256:7b51ed894f4fbea1340262bdae5135797ebbe21d8638978e35d31c6d19f72fb0 \ @@ -730,6 +731,9 @@ pytest-httpserver==1.1.3 \ pytest-rerunfailures==15.1 \ --hash=sha256:c6040368abd7b8138c5b67288be17d6e5611b7368755ce0465dda0362c8ece80 \ --hash=sha256:f674c3594845aba8b23c78e99b1ff8068556cc6a8b277f728071fdc4f4b0b355 +pytest-timeout==2.4.0 \ + --hash=sha256:7e68e90b01f9eff71332b25001f85c75495fc4e3a836701876183c4bcfd0540a \ + --hash=sha256:c42667e5cdadb151aeb5b26d114aff6bdf5a907f176a007a30b940d3d865b5c2 pytest-xdist==3.6.1 \ --hash=sha256:9ed4adfb68a016610848639bb7e02c9352d5d9f03d04809919e2dafc3be4cca7 \ --hash=sha256:ead156a4db231eec769737f57668ef58a2084a34b2e55c4a8fa20d861107300d diff --git a/clients/python/uv.lock b/clients/python/uv.lock index fb639224b77..c598f18c257 100644 --- a/clients/python/uv.lock +++ b/clients/python/uv.lock @@ -1002,6 +1002,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/f3/30/11d836ff01c938969efa319b4ebe2374ed79d28043a12bfc908577aab9f3/pytest_rerunfailures-15.1-py3-none-any.whl", hash = "sha256:f674c3594845aba8b23c78e99b1ff8068556cc6a8b277f728071fdc4f4b0b355", size = 13274, upload_time = "2025-05-08T06:36:32.029Z" }, ] +[[package]] +name = "pytest-timeout" +version = "2.4.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "pytest" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/ac/82/4c9ecabab13363e72d880f2fb504c5f750433b2b6f16e99f4ec21ada284c/pytest_timeout-2.4.0.tar.gz", hash = "sha256:7e68e90b01f9eff71332b25001f85c75495fc4e3a836701876183c4bcfd0540a", size = 17973, upload_time = "2025-05-05T19:44:34.99Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/fa/b6/3127540ecdf1464a00e5a01ee60a1b09175f6913f0644ac748494d9c4b21/pytest_timeout-2.4.0-py3-none-any.whl", hash = "sha256:c42667e5cdadb151aeb5b26d114aff6bdf5a907f176a007a30b940d3d865b5c2", size = 14382, upload_time = "2025-05-05T19:44:33.502Z" }, +] + [[package]] name = "pytest-xdist" version = "3.6.1" @@ -1091,6 +1103,7 @@ dev = [ { name = "pytest-asyncio" }, { name = "pytest-httpserver" }, { name = "pytest-rerunfailures" }, + { name = "pytest-timeout" }, { name = "pytest-xdist" }, { name = "requests" }, ] @@ -1115,6 +1128,7 @@ dev = [ { name = "pytest-asyncio" }, { name = "pytest-httpserver" }, { name = "pytest-rerunfailures" }, + { name = "pytest-timeout" }, { name = "pytest-xdist" }, { name = "requests" }, ] From 59d298e2958cb6dc0441fdb5ea189657a53a5af7 Mon Sep 17 00:00:00 2001 From: Gabriel Bianconi <1275491+GabrielBianconi@users.noreply.github.com> Date: Mon, 26 Jan 2026 08:44:54 -0500 Subject: [PATCH 08/46] Improve `download-large-fixtures.py` (#5850) * Improve download-large-fixtures.py * Improve download-large-fixtures.py * Retry ETag verification during R2 sync Co-authored-by: gabriel * Allow cursoragent in CLA workflow Co-authored-by: gabriel * Fix * Fix --- ui/fixtures/download-large-fixtures.py | 179 +++++++++++++++++++++---- 1 file changed, 153 insertions(+), 26 deletions(-) diff --git a/ui/fixtures/download-large-fixtures.py b/ui/fixtures/download-large-fixtures.py index 20f8112526d..f62edd25d44 100644 --- a/ui/fixtures/download-large-fixtures.py +++ b/ui/fixtures/download-large-fixtures.py @@ -17,7 +17,11 @@ # cd to directory of this file os.chdir(os.path.dirname(os.path.abspath(__file__))) + +# ============================================================================= # Constants +# ============================================================================= + PART_SIZE = 8388608 FIXTURES = [ "large_chat_inference_v2.parquet", @@ -33,10 +37,16 @@ "large_json_comment_feedback.parquet", "large_json_demonstration_feedback.parquet", ] -R2_BUCKET = "https://pub-147e9850a60643208c411e70b636e956.r2.dev" +R2_PUBLIC_BUCKET_URL = "https://pub-147e9850a60643208c411e70b636e956.r2.dev" +R2_S3_ENDPOINT_URL = "https://19918a216783f0ac9e052233569aef60.r2.cloudflarestorage.com/" S3_FIXTURES_DIR = Path("./s3-fixtures") +# ============================================================================= +# Shared utilities +# ============================================================================= + + def calculate_etag(file_path): """Calculate S3/R2 style ETag for a file.""" file_size = os.path.getsize(file_path) @@ -61,19 +71,130 @@ def calculate_etag(file_path): return f"{combined_md5}-{num_parts}" -def get_remote_etag(filename): - """Get ETag from R2 bucket.""" - response = requests.head(f"{R2_BUCKET}/{filename}") - return response.headers.get("ETag", "").strip('"') +def get_remote_etag(filename, retries: int = 3): + """Get ETag from R2 bucket via public URL with retry logic.""" + last_error = None + for attempt in range(retries): + try: + response = requests.head(f"{R2_PUBLIC_BUCKET_URL}/{filename}", timeout=30) + response.raise_for_status() + etag = response.headers.get("ETag") + if not etag: + raise Exception(f"Missing ETag header for {filename}") + return etag.strip('"') + except Exception as exc: + last_error = exc + if attempt < retries - 1: + sleep_time = 3**attempt + print( + f"Error fetching ETag for {filename} (attempt {attempt + 1} of {retries}): {exc}", + flush=True, + ) + print(f"Retrying in {sleep_time} seconds...", flush=True) + time.sleep(sleep_time) + + raise Exception(f"Failed to fetch ETag for {filename} after {retries} attempts") from last_error + + +def verify_etags(): + """Verify ETags of all downloaded fixtures match remote.""" + mismatches = [] + for fixture in FIXTURES: + local_file = S3_FIXTURES_DIR / fixture + if not local_file.exists(): + raise Exception(f"Fixture {fixture} not found after sync") + + local_etag = calculate_etag(local_file) + remote_etag = get_remote_etag(fixture) + + if local_etag != remote_etag: + mismatches.append(f"{fixture}: local={local_etag}, remote={remote_etag}") + else: + print(f"ETag OK: {fixture}", flush=True) + + if mismatches: + raise Exception("ETag mismatches after sync:\n" + "\n".join(mismatches)) + + +# ============================================================================= +# Authenticated path: S3 sync (used in CI with R2 credentials) +# ============================================================================= + + +def sync_fixtures_from_r2(retries: int = 3) -> None: + """Sync fixtures from R2 using aws s3 sync with retry logic.""" + cmd = [ + "aws", + "s3", + "--region", + "auto", + "--endpoint-url", + R2_S3_ENDPOINT_URL, + "--no-progress", + "--cli-connect-timeout", + "30", + "--cli-read-timeout", + "180", + "sync", + "s3://tensorzero-fixtures/", + str(S3_FIXTURES_DIR), + # Only download the files in `FIXTURES` + "--exclude", + "*", + *[arg for f in FIXTURES for arg in ("--include", f)], + ] + + env = { + **os.environ, # Preserve PATH and other environment variables + "AWS_ACCESS_KEY_ID": os.environ["R2_ACCESS_KEY_ID"], + "AWS_SECRET_ACCESS_KEY": os.environ["R2_SECRET_ACCESS_KEY"], + "AWS_MAX_ATTEMPTS": "15", + "AWS_RETRY_MODE": "adaptive", + } + last_error = None + for attempt in range(retries): + print(f"Running aws s3 sync (attempt {attempt + 1} of {retries})...", flush=True) + result = subprocess.run(cmd, env=env) -def download_file(filename, remote_etag): - """Download file from R2 bucket.""" + if result.returncode == 0: + print("Sync completed successfully. Verifying ETags...", flush=True) + try: + verify_etags() + return + except Exception as exc: + last_error = exc + print( + f"ETag verification failed (attempt {attempt + 1} of {retries}): {exc}", + flush=True, + ) + else: + last_error = Exception(f"aws s3 sync failed with exit code {result.returncode}") + print( + f"aws s3 sync failed with exit code {result.returncode} (attempt {attempt + 1} of {retries})", + flush=True, + ) + + if attempt < retries - 1: + sleep_time = 3**attempt # Exponential backoff: 1, 3, 9 seconds + print(f"Retrying in {sleep_time} seconds...", flush=True) + time.sleep(sleep_time) + + raise Exception(f"Fixture sync failed after {retries} attempts") from last_error + + +# ============================================================================= +# Public HTTP fallback (used locally without R2 credentials) +# ============================================================================= + + +def download_file_http(filename, remote_etag): + """Download a single file from R2 via public HTTP URL.""" RETRIES = 3 for i in range(RETRIES): try: - url = f"{R2_BUCKET}/{filename}" - response = requests.get(url, stream=True) + url = f"{R2_PUBLIC_BUCKET_URL}/{filename}" + response = requests.get(url, stream=True, timeout=300) response.raise_for_status() local_file = S3_FIXTURES_DIR / filename @@ -95,21 +216,8 @@ def download_file(filename, remote_etag): raise Exception(f"Failed to download `{filename}` after {RETRIES} attempts") -def main(): - # Create s3-fixtures directory if it doesn't exist - S3_FIXTURES_DIR.mkdir(exist_ok=True) - - if os.environ.get("R2_ACCESS_KEY_ID") is not None and os.environ.get("R2_SECRET_ACCESS_KEY") != "": - print("R2_ACCESS_KEY_ID set, downloading fixtures using 'aws s3 sync'") - subprocess.check_call( - f"aws s3 --region auto --endpoint-url https://19918a216783f0ac9e052233569aef60.r2.cloudflarestorage.com/ sync s3://tensorzero-fixtures/ {S3_FIXTURES_DIR}", - env={ - "AWS_ACCESS_KEY_ID": os.environ["R2_ACCESS_KEY_ID"], - "AWS_SECRET_ACCESS_KEY": os.environ["R2_SECRET_ACCESS_KEY"], - }, - shell=True, - ) - return +def download_fixtures_http(): + """Download all fixtures via public HTTP (fallback when no R2 credentials).""" def process_fixture(fixture): local_file = S3_FIXTURES_DIR / fixture @@ -117,7 +225,7 @@ def process_fixture(fixture): if not local_file.exists(): print(f"Downloading {fixture} (file doesn't exist locally)", flush=True) - download_file(fixture, remote_etag) + download_file_http(fixture, remote_etag) return local_etag = calculate_etag(local_file) @@ -126,7 +234,7 @@ def process_fixture(fixture): print(f"Downloading {fixture} (ETag mismatch)", flush=True) print(f"Local ETag: {local_etag}", flush=True) print(f"Remote ETag: {remote_etag}", flush=True) - download_file(fixture, remote_etag) + download_file_http(fixture, remote_etag) else: print(f"Skipping {fixture} (up to date)", flush=True) @@ -136,6 +244,25 @@ def process_fixture(fixture): for result in executor.map(process_fixture, FIXTURES): assert result is None + +# ============================================================================= +# Main +# ============================================================================= + + +def main(): + S3_FIXTURES_DIR.mkdir(exist_ok=True) + + if os.environ.get("R2_ACCESS_KEY_ID") and os.environ.get("R2_SECRET_ACCESS_KEY"): + print("R2 credentials found, downloading fixtures using `aws s3 sync`", flush=True) + sync_fixtures_from_r2() + else: + print( + "WARNING: `R2_ACCESS_KEY_ID` or `R2_SECRET_ACCESS_KEY` not set. Falling back to slow public HTTP downloads.", + flush=True, + ) + download_fixtures_http() + for fixture in FIXTURES: print(f"Fixture {fixture}:", flush=True) subprocess.run( From 2681771cd7f4f0289ed9257a2590ad2045cf424b Mon Sep 17 00:00:00 2001 From: Viraj Mehta Date: Mon, 26 Jan 2026 14:58:38 -0500 Subject: [PATCH 09/46] temporarily disable python client build job (#5874) * temporarily disable python client build job * remove the build-windows job too * fixed general job stuff --- .github/workflows/general.yml | 9 ++--- .github/workflows/python-client-build.yml | 11 ++++-- .../tests/e2e/docker-compose.live.yml | 38 ++++++++++++++----- tensorzero-core/tests/e2e/docker-compose.yml | 18 +++++++-- 4 files changed, 54 insertions(+), 22 deletions(-) diff --git a/.github/workflows/general.yml b/.github/workflows/general.yml index 024e9874397..83e590837c5 100644 --- a/.github/workflows/general.yml +++ b/.github/workflows/general.yml @@ -956,11 +956,11 @@ jobs: - name: Install Namespace CLI uses: namespacelabs/nscloud-setup@d1c625762f7c926a54bd39252efff0705fd11c64 - continue-on-error: ${{ github.event.pull_request.head.repo.full_name != github.repository || github.actor == 'dependabot[bot]' }} + continue-on-error: ${{ github.event.pull_request.head.repo.full_name != github.repository || github.actor == 'dependabot[bot]' }} - name: Login to Namespace registry run: nsc docker login - continue-on-error: ${{ github.event.pull_request.head.repo.full_name != github.repository || github.actor == 'dependabot[bot]' }} + continue-on-error: ${{ github.event.pull_request.head.repo.full_name != github.repository || github.actor == 'dependabot[bot]' }} - name: Download gateway container image run: | @@ -1264,8 +1264,7 @@ jobs: check-all-general-jobs-passed: permissions: {} if: always() && github.repository == 'tensorzero/tensorzero' - needs: - [ + needs: [ check-version-consistency, check-production-docker-container, check-if-edited-then-edit, @@ -1274,7 +1273,7 @@ jobs: check-python-client-build, check-node-bindings, check-python-schemas, - build-windows, + # TODO (#5873): Add build-windows back when re-enabled build-ui-container, build-gateway-container, build-gateway-e2e-container, diff --git a/.github/workflows/python-client-build.yml b/.github/workflows/python-client-build.yml index 604a83eb8e9..eba2b958742 100644 --- a/.github/workflows/python-client-build.yml +++ b/.github/workflows/python-client-build.yml @@ -120,7 +120,8 @@ jobs: windows: runs-on: ${{ matrix.platform.runner }} - if: github.repository == 'tensorzero/tensorzero' + # TODO (#5873): Temporarily disabled - re-enable when ready + if: false && github.repository == 'tensorzero/tensorzero' strategy: matrix: platform: @@ -254,7 +255,8 @@ jobs: uv run twine check ./clients/python/dist/tensorzero-*.tar.gz check-artifacts: - needs: [linux, musllinux, windows, macos, sdist] + # TODO (#5873): Add windows back when re-enabled + needs: [linux, musllinux, macos, sdist] runs-on: ubuntu-latest if: github.repository == 'tensorzero/tensorzero' steps: @@ -264,6 +266,7 @@ jobs: - name: Count wheels run: | ls -lh ./wheels-*/* - # Count the number of files named 'tensorzero-*.whl', and verify that it equals 6 + # Count the number of files named 'tensorzero-*.whl', and verify that it equals 5 # If we add new platforms / python versions, update this number - ls -lh ./wheels-*/* | grep -c 'tensorzero-.*\.whl' | grep -q 6 + # TODO (#5873): Update to 6 when windows is re-enabled + ls -lh ./wheels-*/* | grep -c 'tensorzero-.*\.whl' | grep -q 5 diff --git a/tensorzero-core/tests/e2e/docker-compose.live.yml b/tensorzero-core/tests/e2e/docker-compose.live.yml index e0e9a23c694..231fd1f37dd 100644 --- a/tensorzero-core/tests/e2e/docker-compose.live.yml +++ b/tensorzero-core/tests/e2e/docker-compose.live.yml @@ -6,7 +6,6 @@ volumes: # We need this for e2e tests that write to a tmpdir from a test, and then read it from the gateway shared-tmpdir: - services: provider-proxy: image: tensorzero/provider-proxy:${TENSORZERO_COMMIT_TAG} @@ -23,9 +22,16 @@ services: # Cache mode can be overridden via PROVIDER_PROXY_CACHE_MODE environment variable. # Options: read-old-write-new (default), read-only, read-write # See: https://github.com/tensorzero/tensorzero/issues/5380 - command: [ "--cache-path", "/app/ci/provider-proxy-cache", "--mode", "${PROVIDER_PROXY_CACHE_MODE:-read-old-write-new}" ] + command: + [ + "--cache-path", + "/app/ci/provider-proxy-cache", + "--mode", + "${PROVIDER_PROXY_CACHE_MODE:-read-old-write-new}", + ] healthcheck: - test: [ "CMD", "wget", "--spider", "--tries=1", "http://localhost:3004/health" ] + test: + ["CMD", "wget", "--spider", "--tries=1", "http://localhost:3004/health"] start_period: 10s start_interval: 1s timeout: 1s @@ -83,7 +89,7 @@ services: VLLM_API_KEY: ${VLLM_API_KEY:-} VOYAGE_API_KEY: ${VOYAGE_API_KEY:-} XAI_API_KEY: ${XAI_API_KEY:-} - command: [ --config-file, tensorzero-core/tests/e2e/config/tensorzero.*.toml ] + command: [--config-file, tensorzero-core/tests/e2e/config/tensorzero.*.toml] volumes: # Mount the e2e directory so the config path exists inside the container - ./:/app/tensorzero-core/tests/e2e:ro @@ -111,7 +117,15 @@ services: extra_hosts: - "howdy.tensorzero.com:127.0.0.1" healthcheck: - test: [ "CMD", "wget", "--no-verbose", "--tries=1", "--spider", "http://localhost:3000/health" ] + test: + [ + "CMD", + "wget", + "--no-verbose", + "--tries=1", + "--spider", + "http://localhost:3000/health", + ] start_period: 1s start_interval: 1s timeout: 1s @@ -135,9 +149,10 @@ services: gateway: condition: service_healthy # Keep this running to make `check-docker-compose.sh` detect that all containers are healthy - command: [ "bash", "-c", "cd /fixtures && ./load_fixtures.sh tensorzero_e2e_tests" ] + command: + ["bash", "-c", "cd /fixtures && ./load_fixtures.sh tensorzero_e2e_tests"] healthcheck: - test: [ "CMD", "test", "-f", "/load_complete.marker" ] + test: ["CMD", "test", "-f", "/load_complete.marker"] interval: 5s timeout: 1s retries: 72 # Retry for up to 6 minutes @@ -161,9 +176,14 @@ services: condition: service_healthy gateway-postgres-migrations: condition: service_healthy - command: [ "bash", "-c", "cd /fixtures && ./load_fixtures_postgres.sh && touch /load_complete_postgres.marker" ] + command: + [ + "bash", + "-c", + "cd /fixtures && ./load_fixtures_postgres.sh && touch /load_complete_postgres.marker", + ] healthcheck: - test: [ "CMD", "test", "-f", "/load_complete_postgres.marker" ] + test: ["CMD", "test", "-f", "/load_complete_postgres.marker"] interval: 5s timeout: 1s retries: 72 # Retry for up to 6 minutes diff --git a/tensorzero-core/tests/e2e/docker-compose.yml b/tensorzero-core/tests/e2e/docker-compose.yml index 98a56420c02..916c17178a1 100644 --- a/tensorzero-core/tests/e2e/docker-compose.yml +++ b/tensorzero-core/tests/e2e/docker-compose.yml @@ -38,9 +38,14 @@ services: gateway-clickhouse-migrations: condition: service_healthy # Keep this running to make `check-docker-compose.sh` detect that all containers are healthy - command: [ "bash", "-c", "cd /fixtures && ./load_fixtures.sh ${TENSORZERO_E2E_TESTS_DATABASE:-tensorzero_e2e_tests} && sleep infinity" ] + command: + [ + "bash", + "-c", + "cd /fixtures && ./load_fixtures.sh ${TENSORZERO_E2E_TESTS_DATABASE:-tensorzero_e2e_tests} && sleep infinity", + ] healthcheck: - test: [ "CMD", "test", "-f", "/load_complete.marker" ] + test: ["CMD", "test", "-f", "/load_complete.marker"] interval: 5s timeout: 1s retries: 72 # Retry for up to 6 minutes @@ -64,9 +69,14 @@ services: condition: service_healthy gateway-postgres-migrations: condition: service_healthy - command: [ "bash", "-c", "cd /fixtures && ./load_fixtures_postgres.sh && touch /load_complete_postgres.marker && sleep infinity" ] + command: + [ + "bash", + "-c", + "cd /fixtures && ./load_fixtures_postgres.sh && touch /load_complete_postgres.marker && sleep infinity", + ] healthcheck: - test: [ "CMD", "test", "-f", "/load_complete_postgres.marker" ] + test: ["CMD", "test", "-f", "/load_complete_postgres.marker"] interval: 5s timeout: 1s retries: 72 # Retry for up to 6 minutes From 9040b186217e9a83bea656736f8cbdd5bf4d2fd4 Mon Sep 17 00:00:00 2001 From: Aaron Hill Date: Mon, 26 Jan 2026 16:40:33 -0500 Subject: [PATCH 10/46] Remove spurious inference request in check_test_streaming_cache_with_err (#5878) * Remove spurious inference request in check_test_streaming_cache_with_err We were performing two streaming requests, and ignoring the first one. The first one could have populated the cache in the background, causing an unexpected cache hit for the real request * Fix lint --- tensorzero-core/tests/e2e/cache.rs | 10 ---------- 1 file changed, 10 deletions(-) diff --git a/tensorzero-core/tests/e2e/cache.rs b/tensorzero-core/tests/e2e/cache.rs index 7fc7e5d909f..284dc8fd52b 100644 --- a/tensorzero-core/tests/e2e/cache.rs +++ b/tensorzero-core/tests/e2e/cache.rs @@ -3,7 +3,6 @@ use futures::StreamExt; use rand::Rng; use reqwest::Client; -use reqwest::StatusCode; use reqwest_eventsource::Event; use reqwest_eventsource::RequestBuilderExt; use serde_json::Value; @@ -526,15 +525,6 @@ pub async fn check_test_streaming_cache_with_err( "cache_options": {"enabled": "on", "lookback_s": 10} }); - let response = Client::new() - .post(get_gateway_endpoint("/inference")) - .json(&payload) - .send() - .await - .unwrap(); - // Check Response is OK, then fields in order - assert_eq!(response.status(), StatusCode::OK); - let mut event_source = Client::new() .post(get_gateway_endpoint("/inference")) .json(&payload) From 5a8fb5ba81ab3a6fa1dbd82797bcc7dab3aefee1 Mon Sep 17 00:00:00 2001 From: Viraj Mehta Date: Mon, 26 Jan 2026 22:28:14 -0500 Subject: [PATCH 11/46] Combined fixes to CI (#5886) * filter * Retry CH checks * update system template * Use R2 sync in CI * Faster cleanup * Faster cleanup * Fix * made dev images use debug builds * Speed up CI * Make CH calls more resilient in CI * Bump image tests timeout This should fix some of the flakes we're seeing on CI (as some of these tests can take over 30 seconds, even when running locally) * Fix * disabled clickhouse cloud fast channel tests --------- Co-authored-by: Andrew Jesson Co-authored-by: Gabriel Bianconi <1275491+GabrielBianconi@users.noreply.github.com> Co-authored-by: Aaron Hill --- .buildkite/pipeline.yml | 9 +- .config/nextest.toml | 4 +- .github/workflows/docker-hub-publish.yml | 4 + .github/workflows/general.yml | 44 ++-- .../ui-tests-e2e-model-inference-cache.yml | 3 + .github/workflows/ui-tests.yml | 8 + ci/buildkite/test-clickhouse-cloud.sh | 25 +- ci/delete-clickhouse-dbs.sh | 8 +- ci/free-disk-space.sh | 69 ++++-- evaluations/Dockerfile | 12 +- tensorzero-optimizers/src/gepa/analyze.rs | 229 +++++++++++++++++- .../analyze/system_template.minijinja | 2 +- 12 files changed, 354 insertions(+), 63 deletions(-) diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index 88b15fadaf5..7ce32147acb 100644 --- a/.buildkite/pipeline.yml +++ b/.buildkite/pipeline.yml @@ -8,10 +8,11 @@ steps: concurrency: 1 concurrency_group: "clickhouse-cloud-normal-${CLICKHOUSE_ID}" command: CLICKHOUSE_PREFIX="dev-tensorzero-e2e-tests-instance-" ./ci/buildkite/test-clickhouse-cloud.sh - - label: "Clickhouse Cloud tests - fast release channel" - concurrency: 1 - concurrency_group: "clickhouse-cloud-fast-${CLICKHOUSE_ID}" - command: CLICKHOUSE_PREFIX="dev-tensorzero-e2e-tests-fast-instance-" TENSORZERO_CLICKHOUSE_FAST_CHANNEL=1 ./ci/buildkite/test-clickhouse-cloud.sh + # Temporarily disabled - see https://github.com/tensorzero/tensorzero/issues/5887 + # - label: "Clickhouse Cloud tests - fast release channel" + # concurrency: 1 + # concurrency_group: "clickhouse-cloud-fast-${CLICKHOUSE_ID}" + # command: CLICKHOUSE_PREFIX="dev-tensorzero-e2e-tests-fast-instance-" TENSORZERO_CLICKHOUSE_FAST_CHANNEL=1 ./ci/buildkite/test-clickhouse-cloud.sh notify: - github_commit_status: diff --git a/.config/nextest.toml b/.config/nextest.toml index 2de3c9cc02e..1228b08fe8e 100644 --- a/.config/nextest.toml +++ b/.config/nextest.toml @@ -155,9 +155,9 @@ slow-timeout = { period = "30s", terminate-after = 3 } filter = 'test(websearch)' slow-timeout = { period = "120s", terminate-after = 2 } -# Image inference seems to be slow on GCP Vertex Gemini, so we give it a longer timeout +# Image inference seems to be slow on many providers, so we give it a longer timeout [[profile.default.overrides]] -filter = 'test(providers::gcp_vertex_gemini::test_image_inference)' +filter = 'test(test_image_inference)' slow-timeout = { period = "60s", terminate-after = 2 } [[profile.default.overrides]] diff --git a/.github/workflows/docker-hub-publish.yml b/.github/workflows/docker-hub-publish.yml index 846705c2d24..ef9aef84d3c 100644 --- a/.github/workflows/docker-hub-publish.yml +++ b/.github/workflows/docker-hub-publish.yml @@ -97,6 +97,10 @@ jobs: tags: ${{ env.DOCKERHUB_USER }}/${{ matrix.container.name }} labels: ${{ steps.meta.outputs.labels }} sbom: true + # Use debug builds for dev images + build-args: | + PROFILE=${{ inputs.is_dev && 'dev' || 'performance' }} + DEBUG_BUILD=${{ inputs.is_dev && '1' || '0' }} - name: Export digest run: | diff --git a/.github/workflows/general.yml b/.github/workflows/general.yml index 83e590837c5..23d7c25235a 100644 --- a/.github/workflows/general.yml +++ b/.github/workflows/general.yml @@ -286,6 +286,8 @@ jobs: lint-rust: permissions: contents: read + # Permission for rust-cache + actions: read runs-on: ubuntu-latest strategy: matrix: @@ -305,6 +307,11 @@ jobs: fi done shell: bash + - uses: Swatinem/rust-cache@9d47c6ad4b02e050fd481d890b2ea34778fd09d6 + with: + cache-provider: "buildjet" + shared-key: "lint-rust-cache" + save-if: ${{ github.event_name == 'merge_group' }} - name: Install cargo-hack uses: taiki-e/install-action@60581cd7025e0e855cebd745379013e286d9c787 @@ -347,23 +354,11 @@ jobs: - name: uv-lock run: | - bash -c ' - git ls-files "**/pyproject.toml" \ - | while read f; do - dir=$(dirname "$f") - (cd "$dir" && uv lock --project="pyproject.toml") - done - ' + git ls-files "**/pyproject.toml" | xargs -P4 -I{} bash -c 'cd "$(dirname "{}")" && uv lock --project="pyproject.toml"' - name: uv-export run: | - bash -c ' - git ls-files "**/pyproject.toml" \ - | while read f; do - dir=$(dirname "$f") - (cd "$dir" && uv export --project="pyproject.toml" --output-file=requirements.txt --quiet) - done - ' + git ls-files "**/pyproject.toml" | xargs -P4 -I{} bash -c 'cd "$(dirname "{}")" && uv export --project="pyproject.toml" --output-file=requirements.txt --quiet' - name: verify uv generated files run: git diff --exit-code @@ -1114,6 +1109,9 @@ jobs: uses: ./.github/workflows/ui-tests.yml with: is_merge_group: ${{ github.event_name == 'merge_group' }} + secrets: + R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }} + R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }} needs: [build-gateway-container, build-mock-provider-api-container] ui-tests-e2e: @@ -1139,16 +1137,26 @@ jobs: S3_SECRET_ACCESS_KEY: ${{ secrets.AWS_SECRET_ACCESS_KEY }} GCP_JWT_KEY: ${{ secrets.GCP_JWT_KEY }} - check-production-docker-container: + check-production-deployment-docker-compose: + needs: [build-gateway-container] permissions: contents: read + # Permission to download artifacts + actions: read runs-on: ubuntu-latest if: github.repository == 'tensorzero/tensorzero' && github.event_name == 'merge_group' steps: - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 - - name: Build Docker container for production deployment tests - run: docker build -t tensorzero/gateway -f gateway/Dockerfile . + - name: Download gateway container image + uses: actions/download-artifact@018cc2cf5baa6db3ef3c5f8a56943fffe632ef53 + with: + name: build-gateway-container + + - name: Load gateway container + run: | + docker load < gateway-container.tar + docker tag tensorzero/gateway:sha-${{ github.sha }} tensorzero/gateway:latest - name: Launch ClickHouse container for E2E tests run: | @@ -1266,7 +1274,7 @@ jobs: if: always() && github.repository == 'tensorzero/tensorzero' needs: [ check-version-consistency, - check-production-docker-container, + check-production-deployment-docker-compose, check-if-edited-then-edit, check-docker-compose-released, check-docker-compose-commit, diff --git a/.github/workflows/ui-tests-e2e-model-inference-cache.yml b/.github/workflows/ui-tests-e2e-model-inference-cache.yml index 500c859c74c..89454df20ca 100644 --- a/.github/workflows/ui-tests-e2e-model-inference-cache.yml +++ b/.github/workflows/ui-tests-e2e-model-inference-cache.yml @@ -256,6 +256,9 @@ jobs: - name: Start dependency Docker containers and apply fixtures working-directory: ui + env: + R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }} + R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }} run: | # Environment variables only used by the gateway container # We deliberately leave these unset when starting the UI container, to ensure diff --git a/.github/workflows/ui-tests.yml b/.github/workflows/ui-tests.yml index c3b706b4835..abc56f1272a 100644 --- a/.github/workflows/ui-tests.yml +++ b/.github/workflows/ui-tests.yml @@ -2,6 +2,11 @@ name: UI Tests on: workflow_call: + secrets: + R2_ACCESS_KEY_ID: + required: true + R2_SECRET_ACCESS_KEY: + required: true inputs: is_merge_group: required: true @@ -82,6 +87,9 @@ jobs: - name: Start Docker containers and apply fixtures working-directory: ui + env: + R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }} + R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }} run: | echo "FIREWORKS_ACCOUNT_ID=not_used" >> fixtures/.env echo "TENSORZERO_GATEWAY_TAG=sha-${{ github.sha }}" >> fixtures/.env diff --git a/ci/buildkite/test-clickhouse-cloud.sh b/ci/buildkite/test-clickhouse-cloud.sh index 403a220ccb3..1932fcfe0a7 100755 --- a/ci/buildkite/test-clickhouse-cloud.sh +++ b/ci/buildkite/test-clickhouse-cloud.sh @@ -15,7 +15,9 @@ export R2_SECRET_ACCESS_KEY=$(buildkite-agent secret get R2_SECRET_ACCESS_KEY) # We concatenate our clickhouse instance prefix, along with our chosen clickhouse id (e.g. 'dev-tensorzero-e2e-tests-instance-' and '0'), to form the instance name # Then, we look up the instance url for this name, and add basic-auth credentials to the url to get our full TENSORZERO_CLICKHOUSE_URL # The 'export' statements go on separate lines to prevent the return code from the $() command from being ignored -CURL_OUTPUT=$(curl --retry 3 --user "$CLICKHOUSE_API_KEY:$CLICKHOUSE_KEY_SECRET" https://api.clickhouse.cloud/v1/organizations/b55f1935-803f-4931-90b3-4d26089004d4/services) +CURL_OUTPUT=$(curl \ + --retry 10 --retry-delay 5 --retry-max-time 300 --retry-all-errors --max-time 30 \ + --user "$CLICKHOUSE_API_KEY:$CLICKHOUSE_KEY_SECRET" https://api.clickhouse.cloud/v1/organizations/b55f1935-803f-4931-90b3-4d26089004d4/services) echo "ClickHouse API response: $CURL_OUTPUT" TENSORZERO_CLICKHOUSE_URL=$(echo "$CURL_OUTPUT" | jq -r ".result[] | select(.name == \"${CLICKHOUSE_PREFIX}${CLICKHOUSE_ID}\") | .endpoints[] | select(.protocol == \"https\") | \"https://$CLICKHOUSE_USERNAME:$CLICKHOUSE_PASSWORD@\" + .host + \":\" + (.port | tostring)") export TENSORZERO_CLICKHOUSE_URL @@ -42,7 +44,9 @@ cleanup_database() { echo "Cleaning up database: $DB_NAME" # Remove the database path from URL for the cleanup call BASE_URL=$(echo "$TENSORZERO_CLICKHOUSE_URL" | sed 's|/[^/]*$||') - curl -X POST "${BASE_URL%/}/?param_target=$DB_NAME" \ + curl -X POST \ + --retry 10 --retry-delay 5 --retry-max-time 300 --retry-all-errors --max-time 30 \ + "${BASE_URL%/}/?param_target=$DB_NAME" \ --data-binary "DROP DATABASE IF EXISTS {target:Identifier}" || true echo "Cleanup completed for database: $DB_NAME" fi @@ -74,7 +78,9 @@ echo "deb https://packages.clickhouse.com/deb stable main" | sudo tee /etc/apt/s sudo apt-get update sudo apt-get install -y clickhouse-client -curl "$TENSORZERO_CLICKHOUSE_URL" --data-binary 'SHOW DATABASES' +curl \ + --retry 10 --retry-delay 5 --retry-max-time 300 --retry-all-errors --max-time 30 \ + "$TENSORZERO_CLICKHOUSE_URL" --data-binary 'SHOW DATABASES' curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y . "$HOME/.cargo/env" curl -LsSf https://astral.sh/uv/0.6.17/install.sh | sh @@ -112,7 +118,18 @@ export CLICKHOUSE_USER="$CLICKHOUSE_USERNAME" export CLICKHOUSE_PASSWORD="$CLICKHOUSE_PASSWORD" export SQLX_OFFLINE=1 cd ui/fixtures -./load_fixtures.sh $DATABASE_NAME +max_retries=3 +for attempt in $(seq 1 $max_retries); do + if ./load_fixtures.sh $DATABASE_NAME; then + break + fi + if [ $attempt -eq $max_retries ]; then + echo "load_fixtures.sh failed after $max_retries attempts" + exit 1 + fi + echo "load_fixtures.sh failed (attempt $attempt/$max_retries), retrying..." + sleep 5 +done cd ../.. sleep 2 diff --git a/ci/delete-clickhouse-dbs.sh b/ci/delete-clickhouse-dbs.sh index 828ac3a0f2a..1049f0092a2 100755 --- a/ci/delete-clickhouse-dbs.sh +++ b/ci/delete-clickhouse-dbs.sh @@ -7,7 +7,9 @@ BASE_URL=$(echo "$TENSORZERO_CLICKHOUSE_URL" | sed 's|/[^/]*$||') echo "Fetching databases" # Get all 'tensorzero_e2e_tests_migration_manager' databases older than 1 hour (to avoid deleting dbs used by currently running tests) # We use the 'metadata_modification_time' of a known table ('ChatInference') to get a lower bound on the database age. -databases=$(curl -X POST "$BASE_URL" \ +databases=$(curl -X POST \ + --retry 10 --retry-delay 5 --retry-max-time 300 --retry-all-errors --max-time 30 \ + "$BASE_URL" \ --data-binary "select database from system.tables where name = 'ChatInference' and startsWith(database, 'tensorzero_e2e_tests_migration_manager') and metadata_modification_time < subtractHours(now(), 1);") echo "The following databases will be deleted:" @@ -17,7 +19,9 @@ echo "$databases" echo "$databases" | while read db; do if [ -n "$db" ]; then echo "Dropping database: $db" - curl -X POST "${BASE_URL%/}/?param_target=$db" \ + curl -X POST \ + --retry 10 --retry-delay 5 --retry-max-time 300 --retry-all-errors --max-time 30 \ + "${BASE_URL%/}/?param_target=$db" \ --data-binary "DROP DATABASE IF EXISTS {target:Identifier}" echo "Database $db dropped." fi diff --git a/ci/free-disk-space.sh b/ci/free-disk-space.sh index c796a13e9bb..2610546d608 100755 --- a/ci/free-disk-space.sh +++ b/ci/free-disk-space.sh @@ -5,13 +5,31 @@ set -euo pipefail # This script EXTREMELY aggressively removes large directories and packages not needed for CI tasks # Optimized for ClickHouse tests - targeting 20GB+ space savings +# Minimum free space required (in GB) - exit early once achieved +REQUIRED_FREE_GB=25 + +# Check if we have enough free space and exit early if so +check_space_and_maybe_exit() { + df -h + local avail_kb + avail_kb=$(df -k / | tail -1 | awk '{print $4}') + local avail_gb=$((avail_kb / 1024 / 1024)) + if [ "$avail_gb" -ge "$REQUIRED_FREE_GB" ]; then + echo "=== Sufficient free space achieved: ${avail_gb}GB >= ${REQUIRED_FREE_GB}GB ===" + df -h + echo "=== Early exit - disk cleanup completed successfully ===" + exit 0 + fi + echo "Current free space: ${avail_gb}GB (need ${REQUIRED_FREE_GB}GB)" +} + echo "=== Disk usage before cleanup ===" -df -h -echo "=== Detailed disk analysis ===" -du -sh /usr/share/* 2>/dev/null | sort -hr | head -30 || true -du -sh /opt/* 2>/dev/null | sort -hr | head -30 || true -du -sh /usr/local/* 2>/dev/null | sort -hr | head -20 || true -du -sh /var/* 2>/dev/null | sort -hr | head -20 || true +check_space_and_maybe_exit +# echo "=== Detailed disk analysis ===" +# du -sh /usr/share/* 2>/dev/null | sort -hr | head -30 || true +# du -sh /opt/* 2>/dev/null | sort -hr | head -30 || true +# du -sh /usr/local/* 2>/dev/null | sort -hr | head -20 || true +# du -sh /var/* 2>/dev/null | sort -hr | head -20 || true echo "=== Starting EXTREME cleanup - targeting 20GB+ savings ===" @@ -23,6 +41,7 @@ echo "Phase 1: Removing massive directories in parallel..." (sudo rm -rf /opt/ghc &) # ~1.6GB GHC (sudo rm -rf /usr/share/swift &) # ~2.8GB Swift wait +check_space_and_maybe_exit # Phase 2: Remove ALL hosted toolchains (parallel) echo "Phase 2: Removing ALL hosted toolchains..." @@ -34,6 +53,7 @@ echo "Phase 2: Removing ALL hosted toolchains..." (sudo rm -rf /opt/hostedtoolcache/Java_Temurin-Hotspot &) (sudo rm -rf /opt/hostedtoolcache/* &) # Remove any remaining toolchains wait +check_space_and_maybe_exit # Phase 3: Remove cloud tools and runtimes (parallel) echo "Phase 3: Removing cloud tools and additional runtimes..." @@ -44,6 +64,7 @@ echo "Phase 3: Removing cloud tools and additional runtimes..." (sudo rm -rf /usr/local/share/powershell &) # PowerShell (sudo rm -rf /opt/microsoft &) # ~772MB Microsoft tools wait +check_space_and_maybe_exit # Phase 4: EXTREME package removal - CI doesn't need most of these echo "Phase 4: Removing unnecessary packages and development tools..." @@ -58,6 +79,7 @@ echo "Phase 4: Removing unnecessary packages and development tools..." (sudo rm -rf /usr/share/icons &) # ~47MB Icons (sudo rm -rf /usr/share/python-babel-localedata &) # ~31MB Python locale wait +check_space_and_maybe_exit # Phase 5: Ultra-aggressive system cleanup echo "Phase 5: EXTREME Docker and system cleanup..." @@ -70,12 +92,14 @@ sudo systemctl stop docker || true sudo rm -rf /var/lib/docker/* || true sudo rm -rf /var/lib/containerd/* || true sudo systemctl start docker || true +check_space_and_maybe_exit # Phase 6: Snap packages removal (they take significant space) echo "Phase 6: Removing snap packages..." sudo snap list 2>/dev/null | awk 'NR>1 {print $1}' | xargs -r sudo snap remove || true sudo rm -rf /var/lib/snapd/snaps/* || true sudo rm -rf /snap/* || true +check_space_and_maybe_exit # Phase 7: Remove large system components not needed for CI echo "Phase 7: Removing large system components..." @@ -87,6 +111,7 @@ echo "Phase 7: Removing large system components..." (sudo rm -rf /usr/share/pixmaps/* &) # Pixmaps (sudo rm -rf /usr/share/applications/* &) # Desktop applications wait +check_space_and_maybe_exit # Phase 8: Aggressive APT cleanup with package removal echo "Phase 8: Extreme APT cleanup and package removal..." @@ -114,6 +139,7 @@ sudo apt-get clean || true sudo rm -rf /var/lib/apt/lists/* || true sudo rm -rf /var/cache/apt/* || true sudo rm -rf /var/lib/dpkg/info/*.list || true +check_space_and_maybe_exit # Phase 9: Remove all documentation, manuals, and locales echo "Phase 9: Removing ALL documentation and locales..." @@ -123,6 +149,7 @@ sudo rm -rf /usr/share/info/* || true sudo rm -rf /usr/share/locale/* || true sudo rm -rf /usr/share/i18n/locales/* || true sudo rm -rf /usr/lib/locale/* || true +check_space_and_maybe_exit # Phase 10: Clean logs, cache, and temporary files extremely aggressively echo "Phase 10: Extreme cleanup of logs, cache, and temp files..." @@ -134,29 +161,27 @@ sudo rm -rf /root/.cache/* || true sudo rm -rf /home/*/.cache/* || true sudo rm -rf /var/cache/* || true sudo rm -rf /usr/share/mime/* || true +check_space_and_maybe_exit # Phase 11: Remove kernel modules and headers we don't need echo "Phase 11: Removing unnecessary kernel components..." -# Keep only the current kernel, remove others -current_kernel=$(uname -r) sudo rm -rf /lib/modules/*/kernel/sound/* || true sudo rm -rf /lib/modules/*/kernel/drivers/media/* || true sudo rm -rf /lib/modules/*/kernel/drivers/staging/* || true sudo rm -rf /usr/src/linux-headers-* || true -echo "=== EXTREME cleanup completed - checking results ===" -df -h -echo "=== Verifying critical tools still work ===" -which cargo && echo "✓ Cargo available" || echo "âš  Cargo not found" -which docker && echo "✓ Docker available" || echo "âš  Docker not found" -which curl && echo "✓ Curl available" || echo "âš  Curl not found" -which git && echo "✓ Git available" || echo "âš  Git not found" -which rustc && echo "✓ Rust available" || echo "âš  Rust not found" - -echo "=== Final disk analysis ===" -df -h -echo "=== Remaining large directories (>100MB) ===" -du -sh /usr/* 2>/dev/null | awk '$1 ~ /[0-9]+[GM]/ && $1+0 >= 100 {print}' | sort -hr | head -20 || true -du -sh /opt/* 2>/dev/null | awk '$1 ~ /[0-9]+[GM]/ && $1+0 >= 100 {print}' | sort -hr | head -10 || true +# echo "=== EXTREME cleanup completed - checking results ===" +# df -h +# echo "=== Verifying critical tools still work ===" +# which cargo && echo "✓ Cargo available" || echo "âš  Cargo not found" +# which docker && echo "✓ Docker available" || echo "âš  Docker not found" +# which curl && echo "✓ Curl available" || echo "âš  Curl not found" +# which git && echo "✓ Git available" || echo "âš  Git not found" +# which rustc && echo "✓ Rust available" || echo "âš  Rust not found" + +# echo "=== Remaining large directories (>100MB) ===" +# du -sh /usr/* 2>/dev/null | awk '$1 ~ /[0-9]+[GM]/ && $1+0 >= 100 {print}' | sort -hr | head -20 || true +# du -sh /opt/* 2>/dev/null | awk '$1 ~ /[0-9]+[GM]/ && $1+0 >= 100 {print}' | sort -hr | head -10 || true echo "=== EXTREME disk cleanup completed successfully ===" +df -h diff --git a/evaluations/Dockerfile b/evaluations/Dockerfile index 3012c882df1..781819ff8cb 100644 --- a/evaluations/Dockerfile +++ b/evaluations/Dockerfile @@ -1,6 +1,12 @@ +# Global build args +ARG CARGO_BUILD_FLAGS="" +ARG PROFILE="release" + # ========== builder ========== FROM rust:1.88.0 AS builder +ARG CARGO_BUILD_FLAGS +ARG PROFILE WORKDIR /src @@ -8,10 +14,8 @@ RUN apt-get update && apt-get install -y clang libc++-dev && rm -rf /var/lib/apt COPY . . -ARG CARGO_BUILD_FLAGS="" - -RUN cargo build --release -p evaluations $CARGO_BUILD_FLAGS && \ - cp -r /src/target/release /release +RUN cargo build --profile ${PROFILE} -p evaluations $CARGO_BUILD_FLAGS +RUN if [ "${PROFILE}" = "dev" ]; then cp -r /src/target/debug /release; else cp -r /src/target/${PROFILE} /release; fi # ========== base ========== diff --git a/tensorzero-optimizers/src/gepa/analyze.rs b/tensorzero-optimizers/src/gepa/analyze.rs index 44438f676e2..03452d03045 100644 --- a/tensorzero-optimizers/src/gepa/analyze.rs +++ b/tensorzero-optimizers/src/gepa/analyze.rs @@ -33,6 +33,37 @@ use evaluations::stats::EvaluationInfo; use crate::gepa::validate::FunctionContext; +/// Fields to exclude when serializing datapoints for GEPA functions. +/// These metadata fields are not relevant for prompt optimization and waste tokens. +const DATAPOINT_FIELDS_TO_DROP: &[&str] = &[ + "dataset_name", + "id", + "episode_id", + "auxiliary", + "is_deleted", + "is_custom", + "source_inference_id", + "staled_at", + "updated_at", + "name", +]; + +/// Serialize a datapoint for GEPA functions, excluding superfluous metadata fields. +/// +/// This reduces token usage by removing fields like IDs, timestamps, and flags +/// that are not relevant for prompt optimization analysis. +fn serialize_filtered_datapoint(datapoint: &Datapoint) -> Result { + let mut value = to_value(datapoint)?; + + if let Some(obj) = value.as_object_mut() { + for field in DATAPOINT_FIELDS_TO_DROP { + obj.remove(*field); + } + } + + Ok(value) +} + // ============================================================================ // Type-safe thought signature stripping // ============================================================================ @@ -202,13 +233,16 @@ pub fn build_analyze_input( // Strip thought signatures at the type level before serialization. // Only Chat types can contain Thought blocks; Json outputs are left intact. - let datapoint_value = match eval_info.datapoint.clone() { - Datapoint::Chat(mut chat_datapoint) => { - strip_signatures_from_chat_datapoint(&mut chat_datapoint); - to_value(Datapoint::Chat(chat_datapoint))? + let mut datapoint_for_serialization = eval_info.datapoint.clone(); + match &mut datapoint_for_serialization { + Datapoint::Chat(chat_datapoint) => { + strip_signatures_from_chat_datapoint(chat_datapoint); } - json_datapoint @ Datapoint::Json(_) => to_value(&json_datapoint)?, - }; + Datapoint::Json(_) => {} // No signatures to strip in JSON + } + + // Serialize with filtered fields + let datapoint_value = serialize_filtered_datapoint(&datapoint_for_serialization)?; map.insert("datapoint".to_string(), datapoint_value); let output_value = match &eval_info.response { @@ -1248,4 +1282,187 @@ mod tests { "text should be preserved" ); } + + // ============================================================================ + // Unit Tests for datapoint field filtering + // ============================================================================ + + #[test] + fn test_serialize_filtered_datapoint_chat() { + let eval_info = create_test_evaluation_info(); + let datapoint = match &eval_info.datapoint { + Datapoint::Chat(dp) => Datapoint::Chat(dp.clone()), + Datapoint::Json(_) => panic!("Expected Chat datapoint"), + }; + + let filtered = serialize_filtered_datapoint(&datapoint) + .expect("Serialization should succeed, expected Chat datapoint with valid fields"); + let obj = filtered + .as_object() + .expect("Filtered result should be a JSON object"); + + // Verify kept fields are present + assert!( + obj.contains_key("function_name"), + "Expected function_name to be kept for context" + ); + assert!( + obj.contains_key("input"), + "Expected input to be kept as core data for analysis" + ); + // Note: output might be None in test data, so we check it's either present or absent as an Option + // tool_params is flattened so its individual fields would be in the object + // tags is optional but if present should be kept + + // Verify dropped fields are absent + assert!( + !obj.contains_key("dataset_name"), + "Expected dataset_name to be dropped as not relevant for optimization" + ); + assert!( + !obj.contains_key("id"), + "Expected id to be dropped as internal identifier" + ); + assert!( + !obj.contains_key("episode_id"), + "Expected episode_id to be dropped as internal identifier" + ); + assert!( + !obj.contains_key("auxiliary"), + "Expected auxiliary to be dropped as extra metadata" + ); + assert!( + !obj.contains_key("is_deleted"), + "Expected is_deleted to be dropped as internal state flag" + ); + assert!( + !obj.contains_key("is_custom"), + "Expected is_custom to be dropped as internal flag" + ); + assert!( + !obj.contains_key("source_inference_id"), + "Expected source_inference_id to be dropped as internal identifier" + ); + assert!( + !obj.contains_key("staled_at"), + "Expected staled_at to be dropped as timestamp not relevant" + ); + assert!( + !obj.contains_key("updated_at"), + "Expected updated_at to be dropped as timestamp not relevant" + ); + assert!( + !obj.contains_key("name"), + "Expected name to be dropped as optional name field typically not relevant" + ); + } + + #[test] + fn test_serialize_filtered_datapoint_json() { + use tensorzero_core::{ + db::stored_datapoint::StoredJsonInferenceDatapoint, + inference::types::JsonInferenceOutput, + }; + + // Create a JSON datapoint + let input = Input { + messages: vec![], + system: None, + }; + + let stored_input = serde_json::from_value(to_value(&input).unwrap()).unwrap(); + + let output_schema = json!({ + "type": "object", + "properties": { + "result": {"type": "string"} + } + }); + + let stored_datapoint = StoredJsonInferenceDatapoint { + dataset_name: "test_dataset".to_string(), + function_name: "test_function".to_string(), + id: Uuid::now_v7(), + episode_id: Some(Uuid::now_v7()), + input: stored_input, + output: Some(JsonInferenceOutput { + raw: Some(r#"{"result": "test"}"#.to_string()), + parsed: Some(json!({"result": "test"})), + }), + output_schema: output_schema.clone(), + tags: Some(HashMap::new()), + auxiliary: String::new(), + is_deleted: false, + is_custom: false, + source_inference_id: None, + staled_at: None, + updated_at: "2025-01-01T00:00:00Z".to_string(), + name: Some("test_name".to_string()), + snapshot_hash: None, + }; + + let datapoint = Datapoint::Json(stored_datapoint.into_datapoint()); + + let filtered = serialize_filtered_datapoint(&datapoint) + .expect("Serialization should succeed, expected JSON datapoint with valid fields"); + let obj = filtered + .as_object() + .expect("Filtered result should be a JSON object"); + + // Verify kept fields are present + assert!( + obj.contains_key("function_name"), + "Expected function_name to be kept for context" + ); + assert!( + obj.contains_key("input"), + "Expected input to be kept as core data for analysis" + ); + assert!( + obj.contains_key("output_schema"), + "Expected output_schema to be kept for understanding output format" + ); + + // Verify dropped fields are absent (same as Chat datapoint) + assert!( + !obj.contains_key("dataset_name"), + "Expected dataset_name to be dropped as not relevant for optimization" + ); + assert!( + !obj.contains_key("id"), + "Expected id to be dropped as internal identifier" + ); + assert!( + !obj.contains_key("episode_id"), + "Expected episode_id to be dropped as internal identifier" + ); + assert!( + !obj.contains_key("auxiliary"), + "Expected auxiliary to be dropped as extra metadata" + ); + assert!( + !obj.contains_key("is_deleted"), + "Expected is_deleted to be dropped as internal state flag" + ); + assert!( + !obj.contains_key("is_custom"), + "Expected is_custom to be dropped as internal flag" + ); + assert!( + !obj.contains_key("source_inference_id"), + "Expected source_inference_id to be dropped as internal identifier" + ); + assert!( + !obj.contains_key("staled_at"), + "Expected staled_at to be dropped as timestamp not relevant" + ); + assert!( + !obj.contains_key("updated_at"), + "Expected updated_at to be dropped as timestamp not relevant" + ); + assert!( + !obj.contains_key("name"), + "Expected name to be dropped as optional name field typically not relevant" + ); + } } diff --git a/tensorzero-optimizers/src/gepa/functions/analyze/system_template.minijinja b/tensorzero-optimizers/src/gepa/functions/analyze/system_template.minijinja index f063e346129..235323929f5 100644 --- a/tensorzero-optimizers/src/gepa/functions/analyze/system_template.minijinja +++ b/tensorzero-optimizers/src/gepa/functions/analyze/system_template.minijinja @@ -11,7 +11,7 @@ The **function_context** contains: The **inference_context** contains: - **message_templates:** The system, user, or assistant templates/instructions used for the LLM -- **datapoint:** The complete datapoint including the LLM input, tags, episode_id, function_name, and other metadata. +- **datapoint:** The datapoint including the function_name and LLM input along with the LLM output, tags, and other metadata (if provided). - The datapoint may include an output (e.g., human annotation or synthetic generation). Use this for reference with caution as it may be suboptimal or even erroneous. The **inference_output** is the LLM response given the **inference_context.** From 138bddc1e824df46daf7f320d914ea4eaf721ac0 Mon Sep 17 00:00:00 2001 From: Shuyang Li Date: Mon, 26 Jan 2026 22:48:58 -0500 Subject: [PATCH 12/46] Namespace valkey keys to be tensorzero_... (#5870) * Namespace valkey keys to be tensorzero_... * Version functions * Graceful migration --- .../db/valkey/lua/tensorzero_ratelimit.lua | 72 +++++++-- tensorzero-core/src/db/valkey/mod.rs | 23 +++ .../src/db/valkey/rate_limiting.rs | 28 ++-- .../tests/e2e/rate_limiting_startup.rs | 152 ++++++++++++++++++ 4 files changed, 250 insertions(+), 25 deletions(-) diff --git a/tensorzero-core/src/db/valkey/lua/tensorzero_ratelimit.lua b/tensorzero-core/src/db/valkey/lua/tensorzero_ratelimit.lua index 4b2e94e15f3..b41a1187ad2 100644 --- a/tensorzero-core/src/db/valkey/lua/tensorzero_ratelimit.lua +++ b/tensorzero-core/src/db/valkey/lua/tensorzero_ratelimit.lua @@ -1,4 +1,4 @@ -#!lua name=tensorzero_ratelimit_v1 +#!lua name=tensorzero_ratelimit_v2 -- IMPORTANT: If you change the implementation of any function in a way that affects -- its behavior or return format, update the function name version suffix (e.g., v1 -> v2) @@ -11,14 +11,18 @@ -- use 1-based indexing. See https://github.com/valkey-io/valkey-doc/blob/main/topics/functions-intro.md -- -- Rate limiting state in Valkey is stored as a Hash: --- key = 'ratelimit:', +-- key = 'tensorzero_ratelimit:', -- value = { 'balance': balance, 'last_refilled': timestamp at microseconds precision } --- LINT.IfEdited() +-- Lint.IfEdited() -- Minimum TTL in seconds to avoid very short expirations local MIN_TTL_SECONDS = 3600 -- 1 hour +-- Rate limiting key prefixes +local RATE_LIMITING_KEY_PREFIX = 'tensorzero_ratelimit:' +local OLD_RATE_LIMITING_KEY_PREFIX_FOR_MIGRATION = 'ratelimit:' + -- Get current server time in microseconds local function get_server_time_micros() local time = server.call('TIME') @@ -97,7 +101,7 @@ local function consume_tickets(keys, args) -- HMGET reads multiple fields from the hash in one call, and returns a table with two values: -- [balance, last_refilled] -- Returns [nil, nil] if the key does not exist. - local data = server.call('HMGET', 'ratelimit:' .. key, 'balance', 'last_refilled') + local data = server.call('HMGET', RATE_LIMITING_KEY_PREFIX .. key, 'balance', 'last_refilled') local balance = tonumber(data[1]) or capacity -- New bucket starts at capacity local last_refilled = tonumber(data[2]) or now @@ -122,7 +126,7 @@ local function consume_tickets(keys, args) -- All succeed: deduct from all and persist for i = 1, num_keys do local key = keys[i] - local redis_key = 'ratelimit:' .. key + local redis_key = RATE_LIMITING_KEY_PREFIX .. key local new_balance = state[i].balance - params[i].requested -- Store aligned last_refilled, NOT current time @@ -179,7 +183,7 @@ local function return_tickets(keys, args) for i = 1, num_keys do local key = keys[i] - local redis_key = 'ratelimit:' .. key + local redis_key = RATE_LIMITING_KEY_PREFIX .. key local base_idx = (i - 1) * 4 local returned = tonumber(args[base_idx + 1]) local capacity = tonumber(args[base_idx + 2]) @@ -226,7 +230,7 @@ local function get_balance(keys, args) local refill_interval = tonumber(args[3]) local now = get_server_time_micros() - local data = server.call('HMGET', 'ratelimit:' .. key, 'balance', 'last_refilled') + local data = server.call('HMGET', RATE_LIMITING_KEY_PREFIX .. key, 'balance', 'last_refilled') local balance = tonumber(data[1]) or capacity local last_refilled = tonumber(data[2]) or now @@ -235,13 +239,59 @@ local function get_balance(keys, args) return cjson.encode({ balance = new_balance }) end +-- Function 4: migrate_old_keys +-- Migrates old rate limit keys ('ratelimit:*') to new prefixed keys ('tensorzero_ratelimit:*'). +-- This preserves existing rate limit state during upgrades from older versions. +-- Keys are only copied if the new key doesn't already exist. +-- KEYS: (none) +-- ARGV: (none) +-- Returns: JSON object {migrated_count} +local function migrate_old_keys(keys, args) + local migrated_count = 0 + local cursor = '0' + + repeat + -- SCAN for old keys + local result = server.call('SCAN', cursor, 'MATCH', OLD_RATE_LIMITING_KEY_PREFIX_FOR_MIGRATION .. '*', 'COUNT', 100) + cursor = result[1] + local found_keys = result[2] + + for _, old_key in ipairs(found_keys) do + -- Extract the suffix after the old prefix + local suffix = string.sub(old_key, #OLD_RATE_LIMITING_KEY_PREFIX_FOR_MIGRATION + 1) + local new_key = RATE_LIMITING_KEY_PREFIX .. suffix + + -- Only migrate if new key doesn't exist + local exists = server.call('EXISTS', new_key) + if exists == 0 then + -- Copy hash data + local data = server.call('HGETALL', old_key) + if #data > 0 then + server.call('HSET', new_key, unpack(data)) + + -- Copy TTL if present + local ttl = server.call('TTL', old_key) + if ttl > 0 then + server.call('EXPIRE', new_key, ttl) + end + + migrated_count = migrated_count + 1 + end + end + end + until cursor == '0' + + return cjson.encode({ migrated_count = migrated_count }) +end + -- Register all functions with version suffix for safe rolling deploys -server.register_function('t0_consume_tickets_v1', consume_tickets) -server.register_function('t0_return_tickets_v1', return_tickets) +server.register_function('tensorzero_consume_tickets_v2', consume_tickets) +server.register_function('tensorzero_return_tickets_v2', return_tickets) server.register_function{ - function_name = 't0_get_balance_v1', + function_name = 'tensorzero_get_balance_v2', callback = get_balance, flags = { 'no-writes' } -- Enables FCALL_RO on read replicas } +server.register_function('tensorzero_migrate_old_keys_v1', migrate_old_keys) --- LINT.ThenEdit(tensorzero-core/src/db/valkey/rate_limiting.rs) +-- Lint.ThenEdit(tensorzero-core/src/db/valkey/rate_limiting.rs) diff --git a/tensorzero-core/src/db/valkey/mod.rs b/tensorzero-core/src/db/valkey/mod.rs index 6a648885eea..80ab3391a10 100644 --- a/tensorzero-core/src/db/valkey/mod.rs +++ b/tensorzero-core/src/db/valkey/mod.rs @@ -39,6 +39,9 @@ impl ValkeyConnectionInfo { // When creating the connection, load the function library into Valkey. Self::load_function_library(&mut connection).await?; + // Migrate old rate limit keys to new prefixed keys for backwards compatibility. + Self::migrate_old_ratelimit_keys(&mut connection).await?; + Ok(Self::Enabled { connection: Box::new(connection), }) @@ -73,6 +76,26 @@ impl ValkeyConnectionInfo { }) }) } + + /// Migrate old rate limit keys (`ratelimit:*`) to new prefixed keys (`tensorzero_ratelimit:*`). + /// This preserves existing rate limit state during upgrades from older versions. + /// Keys are only copied if the new key doesn't already exist. + /// The migration runs entirely in Lua for efficiency (single round-trip). + async fn migrate_old_ratelimit_keys(connection: &mut ConnectionManager) -> Result<(), Error> { + // Call the Lua function to perform the migration atomically on the server + let _result: String = redis::cmd("FCALL") + .arg("tensorzero_migrate_old_keys_v1") + .arg(0) // No keys passed + .query_async(connection) + .await + .map_err(|e| { + Error::new(ErrorDetails::ValkeyQuery { + message: format!("Failed to migrate old rate limit keys: {e}"), + }) + })?; + + Ok(()) + } } const HEALTH_CHECK_TIMEOUT_MS: u64 = 1000; diff --git a/tensorzero-core/src/db/valkey/rate_limiting.rs b/tensorzero-core/src/db/valkey/rate_limiting.rs index 5ec464d089a..07fb49375ee 100644 --- a/tensorzero-core/src/db/valkey/rate_limiting.rs +++ b/tensorzero-core/src/db/valkey/rate_limiting.rs @@ -16,7 +16,7 @@ use super::ValkeyConnectionInfo; // Important: these types must match the response types for the Lua functions in tensorzero-core/src/db/valkey/lua/tensorzero_ratelimit.lua. // Lint.IfEdited() -/// Response from t0_consume_tickets_v1 Lua function +/// Response from tensorzero_consume_tickets_v2 Lua function #[derive(Debug, Deserialize)] struct ConsumeTicketsResponse { key: String, @@ -25,14 +25,14 @@ struct ConsumeTicketsResponse { consumed: i64, } -/// Response from t0_return_tickets_v1 Lua function +/// Response from tensorzero_return_tickets_v2 Lua function #[derive(Debug, Deserialize)] struct ReturnTicketsResponse { key: String, balance: i64, } -/// Response from t0_get_balance_v1 Lua function +/// Response from tensorzero_get_balance_v2 Lua function #[derive(Debug, Deserialize)] struct GetBalanceResponse { balance: i64, @@ -59,7 +59,7 @@ async fn execute_consume_tickets( // Call the versioned Valkey function let mut cmd = redis::cmd("FCALL"); - cmd.arg("t0_consume_tickets_v1").arg(keys.len()); + cmd.arg("tensorzero_consume_tickets_v2").arg(keys.len()); for key in &keys { cmd.arg(*key); } @@ -109,7 +109,7 @@ async fn execute_return_tickets( // Call the versioned Valkey function let mut cmd = redis::cmd("FCALL"); - cmd.arg("t0_return_tickets_v1").arg(keys.len()); + cmd.arg("tensorzero_return_tickets_v2").arg(keys.len()); for key in &keys { cmd.arg(*key); } @@ -148,7 +148,7 @@ async fn execute_get_balance( ) -> Result { // Use FCALL_RO for read-only operations (works with replicas) let result: String = redis::cmd("FCALL_RO") - .arg("t0_get_balance_v1") + .arg("tensorzero_get_balance_v2") .arg(1) // number of keys .arg(key) .arg(capacity as i64) @@ -312,12 +312,12 @@ mod tests { #[tokio::test] async fn test_consume_tickets_sends_correct_command() { - // Expected command: FCALL t0_consume_tickets_v1 1 test_key 10 100 5 1000000 + // Expected command: FCALL tensorzero_consume_tickets_v2 1 test_key 10 100 5 1000000 let expected_response = r#"[{"key":"test_key","success":true,"remaining":90,"consumed":10}]"#; let mut mock = MockRedisConnection::new(vec![MockCmd::new( redis::cmd("FCALL") - .arg("t0_consume_tickets_v1") + .arg("tensorzero_consume_tickets_v2") .arg(1) // numkeys .arg("test_key") .arg(10i64) // requested @@ -348,11 +348,11 @@ mod tests { #[tokio::test] async fn test_consume_tickets_multiple_keys_sends_correct_command() { - // Expected: FCALL t0_consume_tickets_v1 2 key1 key2 10 100 5 1000000 20 200 10 60000000 + // Expected: FCALL tensorzero_consume_tickets_v2 2 key1 key2 10 100 5 1000000 20 200 10 60000000 let expected_response = r#"[{"key":"key1","success":true,"remaining":90,"consumed":10},{"key":"key2","success":true,"remaining":180,"consumed":20}]"#; let mut mock = MockRedisConnection::new(vec![MockCmd::new( redis::cmd("FCALL") - .arg("t0_consume_tickets_v1") + .arg("tensorzero_consume_tickets_v2") .arg(2) // numkeys .arg("key1") .arg("key2") @@ -399,7 +399,7 @@ mod tests { r#"[{"key":"test_key","success":false,"remaining":5,"consumed":0}]"#; let mut mock = MockRedisConnection::new(vec![MockCmd::new( redis::cmd("FCALL") - .arg("t0_consume_tickets_v1") + .arg("tensorzero_consume_tickets_v2") .arg(1) .arg("test_key") .arg(100i64) // requesting more than available @@ -431,7 +431,7 @@ mod tests { let expected_response = r#"[{"key":"test_key","balance":60}]"#; let mut mock = MockRedisConnection::new(vec![MockCmd::new( redis::cmd("FCALL") - .arg("t0_return_tickets_v1") + .arg("tensorzero_return_tickets_v2") .arg(1) // numkeys .arg("test_key") .arg(10i64) // returned @@ -460,7 +460,7 @@ mod tests { let expected_response = r#"{"balance":75}"#; let mut mock = MockRedisConnection::new(vec![MockCmd::new( redis::cmd("FCALL_RO") - .arg("t0_get_balance_v1") + .arg("tensorzero_get_balance_v2") .arg(1) // numkeys .arg("test_key") .arg(100i64) // capacity @@ -490,7 +490,7 @@ mod tests { let expected_response = r#"{"balance":100}"#; let mut mock = MockRedisConnection::new(vec![MockCmd::new( redis::cmd("FCALL_RO") - .arg("t0_get_balance_v1") + .arg("tensorzero_get_balance_v2") .arg(1) .arg("test_key") .arg(100i64) diff --git a/tensorzero-core/tests/e2e/rate_limiting_startup.rs b/tensorzero-core/tests/e2e/rate_limiting_startup.rs index 7bc63af5207..8bf3d22caed 100644 --- a/tensorzero-core/tests/e2e/rate_limiting_startup.rs +++ b/tensorzero-core/tests/e2e/rate_limiting_startup.rs @@ -7,11 +7,14 @@ //! - Starting with both backends available //! - Backend selection via config +use redis::AsyncCommands; use tempfile::NamedTempFile; use tensorzero::ClientBuilder; use tensorzero::ClientBuilderMode; use tensorzero::PostgresConfig; use tensorzero_core::db::clickhouse::test_helpers::CLICKHOUSE_URL; +use tensorzero_core::db::valkey::ValkeyConnectionInfo; +use uuid::Uuid; /// Helper to get the test Postgres URL from environment fn postgres_url() -> Option { @@ -425,3 +428,152 @@ model = "dummy" "Gateway should start successfully with disabled rate limiting and without backends, got: {result:?}", ); } + +/// Test that old rate limit keys (`ratelimit:*`) are migrated to new keys (`tensorzero_ratelimit:*`) +/// when creating a Valkey connection. +#[tokio::test] +async fn test_valkey_migration_old_ratelimit_keys() { + let valkey_url = + valkey_url().expect("Valkey tests should have a TENSORZERO_VALKEY_URL env variable"); + + // Create a raw Redis client to set up the old key + let client = redis::Client::open(valkey_url.as_str()).expect("Failed to create Redis client"); + let mut raw_conn = client + .get_multiplexed_async_connection() + .await + .expect("Failed to connect to Valkey"); + + // Create a unique test key to avoid interference with other tests + let test_id = Uuid::now_v7().to_string(); + let old_key = format!("ratelimit:{test_id}"); + let new_key = format!("tensorzero_ratelimit:{test_id}"); + + // Set up the old key with some rate limit state + let _: () = raw_conn + .hset_multiple( + &old_key, + &[("balance", "42"), ("last_refilled", "1234567890")], + ) + .await + .expect("Failed to set old key"); + + // Set a TTL on the old key + let _: () = raw_conn + .expire(&old_key, 3600) + .await + .expect("Failed to set TTL on old key"); + + // Verify the old key exists and new key doesn't + let old_exists: bool = raw_conn + .exists(&old_key) + .await + .expect("Failed to check old key existence"); + assert!(old_exists, "Old key should exist before migration"); + + let new_exists_before: bool = raw_conn + .exists(&new_key) + .await + .expect("Failed to check new key existence"); + assert!( + !new_exists_before, + "New key should not exist before migration" + ); + + // Create the ValkeyConnectionInfo, which triggers the migration + let _valkey_conn = ValkeyConnectionInfo::new(&valkey_url) + .await + .expect("Failed to create ValkeyConnectionInfo"); + + // Verify the new key now exists with the same data + let new_exists_after: bool = raw_conn + .exists(&new_key) + .await + .expect("Failed to check new key existence after migration"); + assert!(new_exists_after, "New key should exist after migration"); + + // Verify the data was copied correctly + let new_data: Vec<(String, String)> = raw_conn + .hgetall(&new_key) + .await + .expect("Failed to get new key data"); + + let balance = new_data + .iter() + .find(|(k, _)| k == "balance") + .map(|(_, v)| v.as_str()); + let last_refilled = new_data + .iter() + .find(|(k, _)| k == "last_refilled") + .map(|(_, v)| v.as_str()); + + assert_eq!( + balance, + Some("42"), + "Balance should be copied from old key to new key" + ); + assert_eq!( + last_refilled, + Some("1234567890"), + "last_refilled should be copied from old key to new key" + ); + + // Verify TTL was copied (should be close to 3600, accounting for test execution time) + let new_ttl: i64 = raw_conn + .ttl(&new_key) + .await + .expect("Failed to get new key TTL"); + assert!( + new_ttl > 3500 && new_ttl <= 3600, + "TTL should be copied from old key, got {new_ttl}" + ); +} + +/// Test that migration is idempotent - existing new keys are not overwritten +#[tokio::test] +async fn test_valkey_migration_does_not_overwrite_existing_new_keys() { + let valkey_url = + valkey_url().expect("Valkey tests should have a TENSORZERO_VALKEY_URL env variable"); + + let client = redis::Client::open(valkey_url.as_str()).expect("Failed to create Redis client"); + let mut raw_conn = client + .get_multiplexed_async_connection() + .await + .expect("Failed to connect to Valkey"); + + let test_id = Uuid::now_v7().to_string(); + let old_key = format!("ratelimit:{test_id}"); + let new_key = format!("tensorzero_ratelimit:{test_id}"); + + // Set up both old and new keys with different values + let _: () = raw_conn + .hset_multiple(&old_key, &[("balance", "100"), ("last_refilled", "111")]) + .await + .expect("Failed to set old key"); + + let _: () = raw_conn + .hset_multiple(&new_key, &[("balance", "50"), ("last_refilled", "222")]) + .await + .expect("Failed to set new key"); + + // Create the ValkeyConnectionInfo, which triggers the migration + let _valkey_conn = ValkeyConnectionInfo::new(&valkey_url) + .await + .expect("Failed to create ValkeyConnectionInfo"); + + // Verify the new key still has its original data (not overwritten) + let new_data: Vec<(String, String)> = raw_conn + .hgetall(&new_key) + .await + .expect("Failed to get new key data"); + + let balance = new_data + .iter() + .find(|(k, _)| k == "balance") + .map(|(_, v)| v.as_str()); + + assert_eq!( + balance, + Some("50"), + "New key should not be overwritten by migration" + ); +} From 0f03ff2dc13a203db4ade6e26692e2f899ee7d83 Mon Sep 17 00:00:00 2001 From: Shuyang Li Date: Mon, 26 Jan 2026 22:49:10 -0500 Subject: [PATCH 13/46] Fix an erroneous image for postgres fixture (#5865) --- tensorzero-core/tests/e2e/docker-compose.yml | 1 - 1 file changed, 1 deletion(-) diff --git a/tensorzero-core/tests/e2e/docker-compose.yml b/tensorzero-core/tests/e2e/docker-compose.yml index 916c17178a1..2afaa63018a 100644 --- a/tensorzero-core/tests/e2e/docker-compose.yml +++ b/tensorzero-core/tests/e2e/docker-compose.yml @@ -53,7 +53,6 @@ services: # fixtures-postgres and fixtures have exactly the same build context, so the fixture is downloaded only once and shared. fixtures-postgres: - image: tensorzero/fixtures:${TENSORZERO_COMMIT_TAG:-latest} build: dockerfile: ui/fixtures/Dockerfile context: ../../.. From 2db386b31e19b0ed1eff4f2ca172538245799af1 Mon Sep 17 00:00:00 2001 From: Viraj Mehta Date: Mon, 26 Jan 2026 23:10:12 -0500 Subject: [PATCH 14/46] Handle unknown autopilot tool calls (#5836) * added a durable-tool for rejecting tool calls * cleanup * cleanup * consolidated the tool list to one place * made filtering consistent for all unavailable tools * set up types that enforce that the UI shouldnt see missing tool calls * added conversions between event and input message types * directly register internal publishing task * make autoreject a tasktool * fix bug in filtering logic * cleaned up PR comments * added a check to avoid duplicate rejections --- Cargo.lock | 1 + gateway/Cargo.toml | 3 +- gateway/src/main.rs | 9 +- internal/autopilot-client/src/client.rs | 200 +++++++++++++++- internal/autopilot-client/src/error.rs | 8 + internal/autopilot-client/src/lib.rs | 17 +- .../src/reject_missing_tool.rs | 98 ++++++++ internal/autopilot-client/src/types.rs | 222 +++++++++++++++++- internal/autopilot-tools/src/lib.rs | 41 +++- .../src/tools/prod/auto_reject_tool_call.rs | 96 ++++++++ .../autopilot-tools/src/tools/prod/mod.rs | 2 + internal/autopilot-tools/src/visitor.rs | 71 +++++- internal/autopilot-tools/tests/common/mod.rs | 2 +- internal/autopilot-worker/src/worker.rs | 15 +- internal/autopilot-worker/src/wrapper.rs | 2 +- internal/durable-tools/src/lib.rs | 3 +- .../src/tensorzero_client/client_ext.rs | 9 +- .../src/tensorzero_client/embedded.rs | 9 +- .../src/tensorzero_client/mod.rs | 9 +- .../tensorzero-node/lib/bindings/Event.ts | 12 - .../lib/bindings/EventPayload.ts | 4 +- .../lib/bindings/GatewayEvent.ts | 14 ++ .../lib/bindings/GatewayEventPayload.ts | 23 ++ ...atewayEventPayloadToolCallAuthorization.ts | 14 ++ ...sponse.ts => GatewayListEventsResponse.ts} | 12 +- .../lib/bindings/GatewayStreamUpdate.ts | 13 + .../GatewayToolCallAuthorizationStatus.ts | 11 + .../lib/bindings/StreamUpdate.ts | 5 - .../bindings/ToolCallAuthorizationStatus.ts | 3 +- .../lib/bindings/ToolCallDecisionSource.ts | 2 +- .../tensorzero-node/lib/bindings/index.ts | 9 +- tensorzero-core/src/client/mod.rs | 3 + .../src/endpoints/internal/autopilot.rs | 12 +- tensorzero-core/src/error/mod.rs | 10 + tensorzero-core/src/utils/gateway.rs | 28 ++- tensorzero-core/tests/e2e/feedback.rs | 8 +- tensorzero-core/tests/e2e/howdy.rs | 2 + .../autopilot/EventStream.stories.tsx | 12 +- ui/app/components/autopilot/EventStream.tsx | 18 +- .../autopilot/PendingToolCallCard.tsx | 4 +- ui/app/hooks/useAutopilotEventStream.ts | 24 +- .../$session_id/events/authorize.route.ts | 4 +- .../autopilot/sessions/$session_id/route.tsx | 14 +- ui/app/utils/tensorzero/autopilot-client.ts | 14 +- 44 files changed, 951 insertions(+), 141 deletions(-) create mode 100644 internal/autopilot-client/src/reject_missing_tool.rs create mode 100644 internal/autopilot-tools/src/tools/prod/auto_reject_tool_call.rs delete mode 100644 internal/tensorzero-node/lib/bindings/Event.ts create mode 100644 internal/tensorzero-node/lib/bindings/GatewayEvent.ts create mode 100644 internal/tensorzero-node/lib/bindings/GatewayEventPayload.ts create mode 100644 internal/tensorzero-node/lib/bindings/GatewayEventPayloadToolCallAuthorization.ts rename internal/tensorzero-node/lib/bindings/{ListEventsResponse.ts => GatewayListEventsResponse.ts} (67%) create mode 100644 internal/tensorzero-node/lib/bindings/GatewayStreamUpdate.ts create mode 100644 internal/tensorzero-node/lib/bindings/GatewayToolCallAuthorizationStatus.ts delete mode 100644 internal/tensorzero-node/lib/bindings/StreamUpdate.ts diff --git a/Cargo.lock b/Cargo.lock index c92b91d91a9..b271494aebe 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2340,6 +2340,7 @@ name = "gateway" version = "2026.1.5" dependencies = [ "async-stream", + "autopilot-tools", "autopilot-worker", "axum", "brotli", diff --git a/gateway/Cargo.toml b/gateway/Cargo.toml index 318145dcc3b..85ccf40f8cf 100644 --- a/gateway/Cargo.toml +++ b/gateway/Cargo.toml @@ -21,6 +21,7 @@ futures = { workspace = true } tokio-stream = { workspace = true } tensorzero-auth = { path = "../internal/tensorzero-auth" } autopilot-worker = { path = "../internal/autopilot-worker" } +autopilot-tools = { path = "../internal/autopilot-tools" } durable-tools = { path = "../internal/durable-tools" } sqlx = { version = "0.9.0-alpha.1", features = ["sqlx-toml", "postgres", "runtime-tokio", "chrono"] } secrecy = { workspace = true } @@ -34,7 +35,7 @@ async-stream = { workspace = true } workspace = true [features] -e2e_tests = ["tensorzero-core/e2e_tests", "autopilot-worker/e2e_tests"] +e2e_tests = ["tensorzero-core/e2e_tests", "autopilot-worker/e2e_tests", "autopilot-tools/e2e_tests"] default = [] ts-bindings = ["dep:ts-rs", "tensorzero-core/ts-bindings", "tensorzero-optimizers/ts-bindings"] diff --git a/gateway/src/main.rs b/gateway/src/main.rs index bd4b7fe3d18..d92415bf1da 100644 --- a/gateway/src/main.rs +++ b/gateway/src/main.rs @@ -1,3 +1,5 @@ +#![recursion_limit = "256"] + use clap::Parser; use futures::{FutureExt, StreamExt}; use mimalloc::MiMalloc; @@ -236,8 +238,13 @@ async fn run() -> Result<(), ExitCode> { ); } + // Collect available tool names for autopilot (single source of truth) + let available_tools = autopilot_tools::collect_tool_names() + .await + .log_err_pretty("Failed to collect autopilot tool names")?; + // Initialize GatewayHandle - let gateway_handle = gateway::GatewayHandle::new(unwritten_config) + let gateway_handle = gateway::GatewayHandle::new(unwritten_config, available_tools) .await .log_err_pretty("Failed to initialize AppState")?; diff --git a/internal/autopilot-client/src/client.rs b/internal/autopilot-client/src/client.rs index 153dab57b63..fdac07646d0 100644 --- a/internal/autopilot-client/src/client.rs +++ b/internal/autopilot-client/src/client.rs @@ -1,6 +1,8 @@ //! Autopilot API client implementation. +use std::collections::HashSet; use std::fmt; +use std::sync::Arc; use std::time::Duration; use durable_tools_spawn::{SpawnClient, SpawnOptions}; @@ -15,11 +17,12 @@ use uuid::Uuid; use crate::StreamUpdate; use crate::error::AutopilotError; +use crate::reject_missing_tool::reject_missing_tool; use crate::types::{ ApproveAllToolCallsRequest, ApproveAllToolCallsResponse, CreateEventRequest, CreateEventResponse, ErrorResponse, Event, EventPayload, EventPayloadToolCall, - ListEventsParams, ListEventsResponse, ListSessionsParams, ListSessionsResponse, - StreamEventsParams, ToolCallAuthorizationStatus, + GatewayListEventsResponse, GatewayStreamUpdate, ListEventsParams, ListEventsResponse, + ListSessionsParams, ListSessionsResponse, StreamEventsParams, ToolCallAuthorizationStatus, }; /// Default base URL for the Autopilot API. @@ -55,6 +58,7 @@ pub struct AutopilotClientBuilder { spawn_queue_name: String, tool_call_cache_capacity: u64, tool_call_cache_ttl: Duration, + available_tools: Option>, } impl Default for AutopilotClientBuilder { @@ -69,6 +73,7 @@ impl Default for AutopilotClientBuilder { spawn_queue_name: DEFAULT_SPAWN_QUEUE_NAME.to_string(), tool_call_cache_capacity: DEFAULT_TOOL_CALL_CACHE_CAPACITY, tool_call_cache_ttl: DEFAULT_TOOL_CALL_CACHE_TTL, + available_tools: None, } } } @@ -139,18 +144,31 @@ impl AutopilotClientBuilder { self } + /// Sets the set of available tool names for filtering unknown tool calls. + /// + /// This is a required field. When tool calls are received for tools not in + /// this set, they will be automatically rejected with a `NotAvailable` status. + pub fn available_tools(mut self, tools: HashSet) -> Self { + self.available_tools = Some(tools); + self + } + /// Builds the [`AutopilotClient`]. /// /// # Errors /// /// Returns an error if: /// - The API key is not set + /// - The available tools are not set /// - The HTTP client cannot be built /// - Durable spawn configuration is missing or invalid pub async fn build(self) -> Result { let api_key = self .api_key .ok_or(AutopilotError::MissingConfig("api_key"))?; + let available_tools = self + .available_tools + .ok_or(AutopilotError::MissingConfig("available_tools"))?; let base_url = match self.base_url { Some(url) => url, @@ -185,7 +203,7 @@ impl AutopilotClientBuilder { )); }; - let spawn_client = spawn_builder.build().await?; + let spawn_client = Arc::new(spawn_builder.build().await?); let tool_call_cache = Cache::builder() .max_capacity(self.tool_call_cache_capacity) .time_to_live(self.tool_call_cache_ttl) @@ -198,6 +216,7 @@ impl AutopilotClientBuilder { api_key, spawn_client, tool_call_cache, + available_tools: Arc::new(available_tools), }) } } @@ -212,8 +231,10 @@ pub struct AutopilotClient { sse_http_client: reqwest::Client, base_url: Url, api_key: SecretString, - spawn_client: SpawnClient, + spawn_client: Arc, tool_call_cache: Cache, + /// Set of available tool names for filtering unknown tool calls. + available_tools: Arc>, } impl fmt::Debug for AutopilotClient { @@ -247,6 +268,28 @@ impl AutopilotClient { } } + /// Check if a tool call is for an unknown tool. + fn is_unknown_tool(&self, tool_call: &EventPayloadToolCall) -> bool { + !self.available_tools.contains(&tool_call.name) + } + + /// Check if an event should be filtered from responses. + /// + /// Returns `true` for: + /// - `ToolCallAuthorization` events with `NotAvailable` status + /// - `ToolCall` events for unknown tools (not in `available_tools`) + fn should_filter_event(event: &Event, available_tools: &HashSet) -> bool { + match &event.payload { + EventPayload::ToolCallAuthorization(auth) + if matches!(auth.status, ToolCallAuthorizationStatus::NotAvailable) => + { + true + } + EventPayload::ToolCall(tool_call) if !available_tools.contains(&tool_call.name) => true, + _ => false, + } + } + /// Spawns a durable task to execute an approved tool call. /// /// This is called after a `ToolCallAuthorization` event with `Approved` status @@ -359,11 +402,18 @@ impl AutopilotClient { // ------------------------------------------------------------------------- /// Lists events for a session. + /// + /// Unknown tool calls (those not in `available_tools`) are filtered from + /// both `events` and `pending_tool_calls`, and automatically rejected via a durable task. + /// `ToolCallAuthorization::NotAvailable` events are also filtered from results. + /// + /// Returns `GatewayListEventsResponse` which uses `GatewayEvent` - a narrower type + /// that excludes `NotAvailable` authorization status. pub async fn list_events( &self, session_id: Uuid, params: ListEventsParams, - ) -> Result { + ) -> Result { let url = self .base_url .join(&format!("/v1/sessions/{session_id}/events"))?; @@ -378,7 +428,46 @@ impl AutopilotClient { let body: ListEventsResponse = response.json().await?; self.cache_tool_call_events(&body.events); self.cache_tool_call_events(&body.pending_tool_calls); - Ok(body) + + // Filter out unknown tool calls and NotAvailable authorizations, then convert to gateway types. + // We don't reject unknown tool calls from events - only from pending_tool_calls below. + let mut filtered_events = Vec::new(); + for event in body.events { + if Self::should_filter_event(&event, &self.available_tools) { + continue; + } + // Conversion should succeed since we filtered out NotAvailable events + let gateway_event = event.try_into().map_err(|e| { + AutopilotError::Internal(format!("Event conversion failed after filtering: {e}")) + })?; + filtered_events.push(gateway_event); + } + + // Filter pending tool calls and reject unknown ones + let mut filtered_pending_tool_calls = Vec::new(); + for event in body.pending_tool_calls { + // Reject unknown tool calls (but not other filtered events like NotAvailable authorizations) + if let EventPayload::ToolCall(tool_call) = &event.payload + && self.is_unknown_tool(tool_call) + { + reject_missing_tool(&self.spawn_client, tool_call).await?; + } + if Self::should_filter_event(&event, &self.available_tools) { + continue; + } + // Conversion should succeed since we filtered out NotAvailable events + let gateway_event = event.try_into().map_err(|e| { + AutopilotError::Internal(format!("Event conversion failed after filtering: {e}")) + })?; + filtered_pending_tool_calls.push(gateway_event); + } + + Ok(GatewayListEventsResponse { + events: filtered_events, + previous_user_message_event_id: body.previous_user_message_event_id, + status: body.status, + pending_tool_calls: filtered_pending_tool_calls, + }) } /// Gets a single event by ID. @@ -413,8 +502,9 @@ impl AutopilotClient { let tool_call_event_id = match &request.payload { EventPayload::ToolCallAuthorization(auth) => match auth.status { ToolCallAuthorizationStatus::Approved => Some(auth.tool_call_event_id), - // Don't start the tool if rejected - ToolCallAuthorizationStatus::Rejected { .. } => None, + // Don't start the tool if rejected or not available + ToolCallAuthorizationStatus::Rejected { .. } + | ToolCallAuthorizationStatus::NotAvailable => None, }, _ => None, }; @@ -527,12 +617,16 @@ impl AutopilotClient { /// Returns `AutopilotError::Http` if the server returns an error status code /// (e.g., 401 Unauthorized, 404 Not Found). This is checked before returning /// the stream, so connection errors are caught immediately. + /// Returns a stream of `GatewayStreamUpdate` - a narrower type that excludes + /// `NotAvailable` authorization status. pub async fn stream_events( &self, session_id: Uuid, params: StreamEventsParams, - ) -> Result> + use<>, AutopilotError> - { + ) -> Result< + impl Stream> + use<>, + AutopilotError, + > { let mut url = self .base_url .join(&format!("/v1/sessions/{session_id}/events/stream"))?; @@ -570,8 +664,13 @@ impl AutopilotClient { // Connection is good, return the stream let cache = self.tool_call_cache.clone(); + let available_tools = self.available_tools.clone(); + let spawn_client = self.spawn_client.clone(); + let stream = event_source.filter_map(move |result| { let cache = cache.clone(); + let available_tools = available_tools.clone(); + let spawn_client = spawn_client.clone(); async move { match result { Ok(SseEvent::Open) => None, @@ -579,11 +678,32 @@ impl AutopilotClient { if message.event == "event" { match serde_json::from_str::(&message.data) { Ok(update) => { + // Cache tool calls for later lookup if let EventPayload::ToolCall(tool_call) = &update.event.payload { cache.insert(update.event.id, tool_call.clone()); + + // Reject unknown tool calls + if !available_tools.contains(&tool_call.name) + && let Err(e) = + reject_missing_tool(&spawn_client, tool_call).await + { + return Some(Err(e)); + } + } + + // Filter out NotAvailable authorizations and unknown tool calls + if Self::should_filter_event(&update.event, &available_tools) { + return None; + } + + // Convert to gateway type - should succeed since we filtered NotAvailable above + match update.try_into() { + Ok(gateway_update) => Some(Ok(gateway_update)), + Err(e) => Some(Err(AutopilotError::Internal(format!( + "StreamUpdate conversion failed after filtering: {e}" + )))), } - Some(Ok(update)) } Err(e) => Some(Err(AutopilotError::Json(e))), } @@ -654,3 +774,61 @@ impl AutopilotClient { }) } } + +#[cfg(test)] +mod tests { + use super::*; + #[test] + fn test_is_unknown_tool_with_known_tool() { + let mut tools = HashSet::new(); + tools.insert("inference".to_string()); + tools.insert("feedback".to_string()); + let available_tools: Arc> = Arc::new(tools); + + let tool_name = "inference".to_string(); + + let is_unknown = !available_tools.contains(&tool_name); + assert!(!is_unknown, "Known tool should not be marked as unknown"); + } + + #[test] + fn test_is_unknown_tool_with_unknown_tool() { + let mut tools = HashSet::new(); + tools.insert("inference".to_string()); + tools.insert("feedback".to_string()); + let available_tools: Arc> = Arc::new(tools); + + let tool_name = "fake_nonexistent_tool".to_string(); + + let is_unknown = !available_tools.contains(&tool_name); + assert!(is_unknown, "Unknown tool should be marked as unknown"); + } + + #[test] + fn test_default_base_url_is_valid() { + // Test that DEFAULT_BASE_URL can be parsed + let url: Url = DEFAULT_BASE_URL + .parse() + .expect("DEFAULT_BASE_URL should be a valid URL"); + assert_eq!(url.scheme(), "https", "Should use HTTPS"); + assert!(url.host_str().is_some(), "Should have a host"); + } + + #[test] + fn test_build_fails_without_available_tools() { + let result = futures::executor::block_on( + AutopilotClient::builder() + .api_key("test_key") + .spawn_database_url("postgres://localhost/test") + .build(), + ); + + assert!( + matches!( + result, + Err(AutopilotError::MissingConfig("available_tools")) + ), + "Building without available_tools should fail with MissingConfig error" + ); + } +} diff --git a/internal/autopilot-client/src/error.rs b/internal/autopilot-client/src/error.rs index 26e8780e1ec..9ed4bad7882 100644 --- a/internal/autopilot-client/src/error.rs +++ b/internal/autopilot-client/src/error.rs @@ -32,6 +32,10 @@ pub enum AutopilotError { #[error("Spawn error: {0}")] Spawn(#[from] durable_tools_spawn::SpawnError), + /// Database error. + #[error("Database error: {0}")] + Database(#[from] sqlx::Error), + /// Missing required configuration. #[error("Missing required configuration: {0}")] MissingConfig(&'static str), @@ -47,4 +51,8 @@ pub enum AutopilotError { response: ApproveAllToolCallsResponse, errors: Vec<(Uuid, Box)>, }, + + /// Internal error - indicates a bug in the client. + #[error("Internal error: {0}")] + Internal(String), } diff --git a/internal/autopilot-client/src/lib.rs b/internal/autopilot-client/src/lib.rs index 5d66bed23e2..9da71128dea 100644 --- a/internal/autopilot-client/src/lib.rs +++ b/internal/autopilot-client/src/lib.rs @@ -23,8 +23,9 @@ //! let response = client.create_event( //! Uuid::nil(), //! CreateEventRequest { -//! deployment_id: Uuid::now_v7(), +//! deployment_id: Uuid::now_v7().to_string(), //! tensorzero_version: "2025.1.0".to_string(), +//! config_snapshot_hash: Some("abc123".to_string()), //! payload: EventPayload::Message(EventPayloadMessage { //! role: Role::User, //! content: vec![EventPayloadMessageContent::Text(Text { @@ -42,20 +43,24 @@ mod client; mod error; +mod reject_missing_tool; mod types; pub use client::{ AutopilotClient, AutopilotClientBuilder, DEFAULT_BASE_URL, DEFAULT_SPAWN_QUEUE_NAME, }; pub use error::AutopilotError; +pub use reject_missing_tool::reject_missing_tool; pub use types::{ ApproveAllToolCallsRequest, ApproveAllToolCallsResponse, AutopilotSideInfo, AutopilotStatus, AutopilotToolResult, Base64File, CreateEventRequest, CreateEventResponse, ErrorDetail, ErrorResponse, Event, EventPayload, EventPayloadError, EventPayloadMessage, EventPayloadMessageContent, EventPayloadStatusUpdate, EventPayloadToolCall, - EventPayloadToolCallAuthorization, EventPayloadToolResult, File, ListEventsParams, - ListEventsResponse, ListSessionsParams, ListSessionsResponse, ObjectStoragePointer, - OptimizationWorkflowSideInfo, RawText, Role, Session, StatusUpdate, StreamEventsParams, - StreamUpdate, Template, Text, Thought, ToolCallAuthorizationStatus, ToolCallDecisionSource, - ToolCallWrapper, ToolOutcome, Unknown, UrlFile, + EventPayloadToolCallAuthorization, EventPayloadToolResult, File, GatewayEvent, + GatewayEventPayload, GatewayEventPayloadToolCallAuthorization, GatewayListEventsResponse, + GatewayStreamUpdate, GatewayToolCallAuthorizationStatus, ListEventsParams, ListEventsResponse, + ListSessionsParams, ListSessionsResponse, ObjectStoragePointer, OptimizationWorkflowSideInfo, + RawText, Role, Session, StatusUpdate, StreamEventsParams, StreamUpdate, Template, Text, + Thought, ToolCallAuthorizationStatus, ToolCallDecisionSource, ToolCallWrapper, ToolOutcome, + Unknown, UrlFile, }; diff --git a/internal/autopilot-client/src/reject_missing_tool.rs b/internal/autopilot-client/src/reject_missing_tool.rs new file mode 100644 index 00000000000..2cf1a685e27 --- /dev/null +++ b/internal/autopilot-client/src/reject_missing_tool.rs @@ -0,0 +1,98 @@ +//! Functionality for rejecting tool calls for tools that are not available. + +use std::sync::Arc; + +use durable_tools_spawn::{SpawnClient, SpawnOptions}; +use uuid::Uuid; + +use crate::error::AutopilotError; +use crate::types::{AutopilotSideInfo, EventPayloadToolCall}; + +/// Spawn a durable task to reject a tool call for a missing/unavailable tool. +/// +/// This enqueues a task that will send a `NotAvailable` authorization event +/// to the autopilot API so the session can continue without hanging on an +/// unknown tool. +/// +/// # Duplicate Prevention +/// +/// This function checks if a rejection task already exists for the tool call +/// before spawning a new one to prevent duplicate rejections. +/// +/// # Errors +/// +/// Returns an error if checking for existing tasks or spawning the durable task fails. +pub async fn reject_missing_tool( + spawn_client: &Arc, + tool_call: &EventPayloadToolCall, +) -> Result<(), AutopilotError> { + let tool_call_event_id = tool_call.side_info.tool_call_event_id; + + // Check if we've already spawned a rejection for this tool call + if check_tool_rejection_exists(spawn_client, tool_call_event_id).await? { + tracing::debug!( + tool_name = %tool_call.name, + tool_call_event_id = %tool_call_event_id, + "Rejection task already exists, skipping" + ); + return Ok(()); + } + + let side_info = AutopilotSideInfo { + session_id: tool_call.side_info.session_id, + tool_call_event_id, + config_snapshot_hash: tool_call.side_info.config_snapshot_hash.clone(), + optimization: tool_call.side_info.optimization.clone(), + }; + + let episode_id = Uuid::now_v7(); + spawn_client + .spawn_tool_by_name( + "__auto_reject_tool_call__", + serde_json::json!({}), + serde_json::to_value(&side_info)?, + episode_id, + SpawnOptions::default(), + ) + .await?; + + tracing::info!( + tool_name = %tool_call.name, + session_id = %side_info.session_id, + tool_call_event_id = %side_info.tool_call_event_id, + "Rejecting tool call for missing tool" + ); + + Ok(()) +} + +/// Check if a rejection task already exists for the given tool call event. +/// +/// This queries the durable tasks table to see if we've already spawned +/// a `__auto_reject_tool_call__` task for this specific tool call. +pub async fn check_tool_rejection_exists( + spawn_client: &Arc, + tool_call_event_id: Uuid, +) -> Result { + let queue_name = spawn_client.queue_name(); + + // Query for an existing rejection task with matching tool_call_event_id in side_info. + // The table name is derived from the queue name, which is defaulted or comes from an env var. + let query = format!( + r" + SELECT EXISTS( + SELECT 1 + FROM durable.t_{queue_name} + WHERE task_name = '__auto_reject_tool_call__' + AND params->'side_info'->>'tool_call_event_id' = $1 + ) + " + ); + + let exists: bool = sqlx::query_scalar(sqlx::AssertSqlSafe(query)) + .bind(tool_call_event_id.to_string()) + .fetch_one(spawn_client.pool()) + .await?; + + Ok(exists) +} diff --git a/internal/autopilot-client/src/types.rs b/internal/autopilot-client/src/types.rs index a530e5d7077..06d542a4756 100644 --- a/internal/autopilot-client/src/types.rs +++ b/internal/autopilot-client/src/types.rs @@ -12,6 +12,7 @@ pub use tensorzero_types::{ Base64File, File, ObjectStoragePointer, RawText, Role, Template, Text, Thought, ToolCallWrapper, Unknown, UrlFile, }; +use tensorzero_types::{InputMessage, InputMessageContent}; use uuid::Uuid; // ============================================================================= @@ -41,6 +42,41 @@ pub struct EventPayloadMessage { pub content: Vec, } +impl TryFrom for EventPayloadMessage { + type Error = &'static str; + + fn try_from(msg: InputMessage) -> Result { + let content = msg + .content + .into_iter() + .map(|c| match c { + InputMessageContent::Text(text) => Ok(EventPayloadMessageContent::Text(text)), + _ => Err("EventPayloadMessage only supports Text content blocks"), + }) + .collect::, _>>()?; + + Ok(EventPayloadMessage { + role: msg.role, + content, + }) + } +} + +impl From for InputMessage { + fn from(msg: EventPayloadMessage) -> Self { + InputMessage { + role: msg.role, + content: msg + .content + .into_iter() + .map(|c| match c { + EventPayloadMessageContent::Text(text) => InputMessageContent::Text(text), + }) + .collect(), + } + } +} + /// A session representing an autopilot conversation. #[cfg_attr(feature = "ts-bindings", derive(ts_rs::TS))] #[derive(Debug, Clone, Serialize, Deserialize)] @@ -54,10 +90,11 @@ pub struct Session { pub created_at: DateTime, } -/// An event within a session. +/// Internal event type - consumers should use `GatewayEvent` instead. +/// +/// Note: TS derive is needed for types that reference this, but we don't export it. #[cfg_attr(feature = "ts-bindings", derive(ts_rs::TS))] #[derive(Debug, Clone, Serialize, Deserialize)] -#[cfg_attr(feature = "ts-bindings", ts(export))] pub struct Event { pub id: Uuid, pub payload: EventPayload, @@ -65,6 +102,32 @@ pub struct Event { pub created_at: DateTime, } +/// An event as seen by gateway consumers. +/// +/// Uses `GatewayEventPayload` which excludes `NotAvailable` authorization status. +#[cfg_attr(feature = "ts-bindings", derive(ts_rs::TS))] +#[derive(Debug, Clone, Serialize, Deserialize)] +#[cfg_attr(feature = "ts-bindings", ts(export))] +pub struct GatewayEvent { + pub id: Uuid, + pub payload: GatewayEventPayload, + pub session_id: Uuid, + pub created_at: DateTime, +} + +impl TryFrom for GatewayEvent { + type Error = &'static str; + + fn try_from(event: Event) -> Result { + Ok(GatewayEvent { + id: event.id, + payload: event.payload.try_into()?, + session_id: event.session_id, + created_at: event.created_at, + }) + } +} + /// The UX-relevant status of the Autopilot. #[cfg_attr(feature = "ts-bindings", derive(ts_rs::TS))] #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] @@ -79,14 +142,38 @@ pub enum AutopilotStatus { Failed, } +/// Internal stream update type - consumers should use `GatewayStreamUpdate` instead. +/// +/// Note: TS derive is needed for types that reference this, but we don't export it. #[cfg_attr(feature = "ts-bindings", derive(ts_rs::TS))] #[derive(Debug, Clone, Serialize, Deserialize)] -#[cfg_attr(feature = "ts-bindings", ts(export))] pub struct StreamUpdate { pub event: Event, pub status: AutopilotStatus, } +/// Stream update as seen by gateway consumers. +/// +/// Uses `GatewayEvent` which excludes `NotAvailable` authorization status. +#[cfg_attr(feature = "ts-bindings", derive(ts_rs::TS))] +#[derive(Debug, Clone, Serialize, Deserialize)] +#[cfg_attr(feature = "ts-bindings", ts(export))] +pub struct GatewayStreamUpdate { + pub event: GatewayEvent, + pub status: AutopilotStatus, +} + +impl TryFrom for GatewayStreamUpdate { + type Error = &'static str; + + fn try_from(update: StreamUpdate) -> Result { + Ok(GatewayStreamUpdate { + event: update.event.try_into()?, + status: update.status, + }) + } +} + /// Error payload for an event. #[cfg_attr(feature = "ts-bindings", derive(ts_rs::TS))] #[derive(Debug, Clone, Serialize, Deserialize)] @@ -109,14 +196,13 @@ pub struct EventPayloadToolResult { pub outcome: ToolOutcome, } -/// The payload of an event. +/// Internal event payload type - consumers should use `GatewayEventPayload` instead. +/// +/// Note: TS derive is needed for types that reference this, but we don't export it. #[cfg_attr(feature = "ts-bindings", derive(ts_rs::TS))] #[derive(Debug, Clone, Serialize, Deserialize)] #[serde(tag = "type", rename_all = "snake_case")] -#[cfg_attr( - feature = "ts-bindings", - ts(export, tag = "type", rename_all = "snake_case") -)] +#[cfg_attr(feature = "ts-bindings", ts(tag = "type", rename_all = "snake_case"))] pub enum EventPayload { Message(EventPayloadMessage), Error(EventPayloadError), @@ -141,6 +227,46 @@ impl EventPayload { } } +/// Event payload as seen by gateway consumers. +/// +/// Uses `GatewayEventPayloadToolCallAuthorization` which excludes `NotAvailable` status. +#[cfg_attr(feature = "ts-bindings", derive(ts_rs::TS))] +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(tag = "type", rename_all = "snake_case")] +#[cfg_attr( + feature = "ts-bindings", + ts(export, tag = "type", rename_all = "snake_case") +)] +pub enum GatewayEventPayload { + Message(EventPayloadMessage), + Error(EventPayloadError), + StatusUpdate(EventPayloadStatusUpdate), + ToolCall(EventPayloadToolCall), + ToolCallAuthorization(GatewayEventPayloadToolCallAuthorization), + ToolResult(EventPayloadToolResult), + #[serde(other)] + #[serde(alias = "other")] // legacy name + Unknown, +} + +impl TryFrom for GatewayEventPayload { + type Error = &'static str; + + fn try_from(payload: EventPayload) -> Result>::Error> { + match payload { + EventPayload::Message(m) => Ok(GatewayEventPayload::Message(m)), + EventPayload::Error(e) => Ok(GatewayEventPayload::Error(e)), + EventPayload::StatusUpdate(s) => Ok(GatewayEventPayload::StatusUpdate(s)), + EventPayload::ToolCall(t) => Ok(GatewayEventPayload::ToolCall(t)), + EventPayload::ToolCallAuthorization(auth) => { + Ok(GatewayEventPayload::ToolCallAuthorization(auth.try_into()?)) + } + EventPayload::ToolResult(r) => Ok(GatewayEventPayload::ToolResult(r)), + EventPayload::Unknown => Ok(GatewayEventPayload::Unknown), + } + } +} + /// A status update within a session. #[cfg_attr(feature = "ts-bindings", derive(ts_rs::TS))] #[derive(Debug, Clone, Serialize, Deserialize)] @@ -269,6 +395,7 @@ pub struct AutopilotToolResult { #[serde(tag = "type", rename_all = "snake_case")] pub enum ToolCallDecisionSource { Ui, + Automatic, } #[cfg_attr(feature = "ts-bindings", derive(ts_rs::TS))] @@ -279,12 +406,66 @@ pub struct EventPayloadToolCallAuthorization { pub status: ToolCallAuthorizationStatus, } +/// Tool call authorization payload as seen by gateway consumers. +/// +/// Uses `GatewayToolCallAuthorizationStatus` which excludes `NotAvailable`. +#[cfg_attr(feature = "ts-bindings", derive(ts_rs::TS))] +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct GatewayEventPayloadToolCallAuthorization { + pub source: ToolCallDecisionSource, + pub tool_call_event_id: Uuid, + pub status: GatewayToolCallAuthorizationStatus, +} + +impl TryFrom for GatewayEventPayloadToolCallAuthorization { + type Error = &'static str; + + fn try_from(auth: EventPayloadToolCallAuthorization) -> Result { + Ok(GatewayEventPayloadToolCallAuthorization { + source: auth.source, + tool_call_event_id: auth.tool_call_event_id, + status: auth.status.try_into()?, + }) + } +} + #[cfg_attr(feature = "ts-bindings", derive(ts_rs::TS))] #[derive(Debug, Clone, Serialize, Deserialize)] #[serde(tag = "type", rename_all = "snake_case")] pub enum ToolCallAuthorizationStatus { Approved, Rejected { reason: String }, + NotAvailable, +} + +/// Authorization status for tool calls as seen by gateway consumers. +/// +/// This is a narrower type than `ToolCallAuthorizationStatus` that excludes +/// `NotAvailable` since that status is filtered out before reaching consumers. +#[cfg_attr(feature = "ts-bindings", derive(ts_rs::TS))] +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(tag = "type", rename_all = "snake_case")] +pub enum GatewayToolCallAuthorizationStatus { + Approved, + Rejected { reason: String }, +} + +impl TryFrom for GatewayToolCallAuthorizationStatus { + type Error = &'static str; + + fn try_from(status: ToolCallAuthorizationStatus) -> Result { + match status { + ToolCallAuthorizationStatus::Approved => { + Ok(GatewayToolCallAuthorizationStatus::Approved) + } + ToolCallAuthorizationStatus::Rejected { reason } => { + Ok(GatewayToolCallAuthorizationStatus::Rejected { reason }) + } + ToolCallAuthorizationStatus::NotAvailable => { + Err("NotAvailable status should be filtered before conversion") + } + } + } } #[cfg_attr(feature = "ts-bindings", derive(ts_rs::TS))] @@ -398,10 +579,11 @@ pub struct CreateEventResponse { pub session_id: Uuid, } -/// Response from listing events. +/// Internal response type - consumers should use `GatewayListEventsResponse` instead. +/// +/// Note: TS derive is needed for types that reference this, but we don't export it. #[cfg_attr(feature = "ts-bindings", derive(ts_rs::TS))] #[derive(Debug, Clone, Serialize, Deserialize)] -#[cfg_attr(feature = "ts-bindings", ts(export))] pub struct ListEventsResponse { pub events: Vec, /// The most recent `message` event with role `user` in this session. @@ -416,6 +598,26 @@ pub struct ListEventsResponse { pub pending_tool_calls: Vec, } +/// Response from listing events as seen by gateway consumers. +/// +/// Uses `GatewayEvent` which excludes `NotAvailable` authorization status. +#[cfg_attr(feature = "ts-bindings", derive(ts_rs::TS))] +#[derive(Debug, Clone, Serialize, Deserialize)] +#[cfg_attr(feature = "ts-bindings", ts(export))] +pub struct GatewayListEventsResponse { + pub events: Vec, + /// The most recent `message` event with role `user` in this session. + pub previous_user_message_event_id: Uuid, + /// The current status of the Autopilot in this session. + /// Ignores pagination parameters. + pub status: AutopilotStatus, + /// All tool calls in Event history that do not have responses. + /// These may be duplicates of some of the values in events. + /// All EventPayloads in these Events should be of type ToolCall. + #[serde(default)] + pub pending_tool_calls: Vec, +} + /// Response from listing sessions. #[cfg_attr(feature = "ts-bindings", derive(ts_rs::TS))] #[derive(Debug, Clone, Serialize, Deserialize)] diff --git a/internal/autopilot-tools/src/lib.rs b/internal/autopilot-tools/src/lib.rs index 0c10abb5154..b6cbad192d0 100644 --- a/internal/autopilot-tools/src/lib.rs +++ b/internal/autopilot-tools/src/lib.rs @@ -40,13 +40,32 @@ //! - `ErrorSimpleTool` - Always returns an error //! - `SlowSimpleTool` - Sleeps for configurable duration +use std::collections::HashSet; + pub mod error; pub mod tools; mod visitor; pub use error::{AutopilotToolError, AutopilotToolResult}; -pub use visitor::ToolVisitor; +pub use visitor::{ToolNameCollector, ToolVisitor}; + +/// Collect all available tool names. +/// +/// This uses the visitor pattern with `for_each_tool` to derive the set of +/// tool names from the single source of truth. +/// +/// This is used by `AutopilotClient` to know which tools are available +/// for filtering unknown tool calls. +/// +/// # Errors +/// +/// Returns an error if visiting any tool fails (e.g., lock poisoning). +pub async fn collect_tool_names() -> Result, String> { + let collector = ToolNameCollector::new(); + for_each_tool(&collector).await?; + Ok(collector.into_names()) +} /// Iterate over all tools with a visitor. /// @@ -67,8 +86,8 @@ pub use visitor::ToolVisitor; /// impl ToolVisitor for LocalVisitor<'_> { /// type Error = ToolError; /// -/// async fn visit_task_tool(&self) -> Result<(), ToolError> { -/// self.0.register_task_tool::().await?; +/// async fn visit_task_tool(&self, tool: T) -> Result<(), ToolError> { +/// self.0.register_task_tool_instance(tool).await?; /// Ok(()) /// } /// @@ -92,8 +111,8 @@ pub use visitor::ToolVisitor; /// impl ToolVisitor for RemoteVisitor<'_> { /// type Error = ToolError; /// -/// async fn visit_task_tool(&self) -> Result<(), ToolError> { -/// self.0.register_client_tool::().await?; +/// async fn visit_task_tool(&self, tool: T) -> Result<(), ToolError> { +/// self.0.register_client_tool_instance(tool).await?; /// Ok(()) /// } /// @@ -139,7 +158,7 @@ pub async fn for_each_tool(visitor: &V) -> Result<(), V::Error> .visit_simple_tool::() .await?; visitor - .visit_task_tool::() + .visit_task_tool(tools::LaunchOptimizationWorkflowTool) .await?; visitor .visit_simple_tool::() @@ -172,11 +191,11 @@ pub async fn for_each_tool(visitor: &V) -> Result<(), V::Error> #[cfg(feature = "e2e_tests")] { // TaskTools - visitor.visit_task_tool::().await?; - visitor.visit_task_tool::().await?; - visitor.visit_task_tool::().await?; - visitor.visit_task_tool::().await?; - visitor.visit_task_tool::().await?; + visitor.visit_task_tool(tools::EchoTool).await?; + visitor.visit_task_tool(tools::SlowTool).await?; + visitor.visit_task_tool(tools::FailingTool).await?; + visitor.visit_task_tool(tools::FlakyTool).await?; + visitor.visit_task_tool(tools::PanicTool).await?; // SimpleTools visitor.visit_simple_tool::().await?; diff --git a/internal/autopilot-tools/src/tools/prod/auto_reject_tool_call.rs b/internal/autopilot-tools/src/tools/prod/auto_reject_tool_call.rs new file mode 100644 index 00000000000..25b9425fdd3 --- /dev/null +++ b/internal/autopilot-tools/src/tools/prod/auto_reject_tool_call.rs @@ -0,0 +1,96 @@ +//! Auto-reject tool for handling unknown tool calls. +//! +//! This tool is called when the autopilot client receives a tool call for a tool +//! that doesn't exist in the TensorZero deployment. It sends a `NotAvailable` +//! authorization event to the autopilot API. + +use std::borrow::Cow; + +use async_trait::async_trait; +use durable_tools::{NonControlToolError, TaskTool, ToolContext, ToolMetadata, ToolResult}; +use schemars::Schema; +use serde::{Deserialize, Serialize}; + +use crate::error::AutopilotToolError; +use autopilot_client::{ + AutopilotSideInfo, EventPayload, EventPayloadToolCallAuthorization, + ToolCallAuthorizationStatus, ToolCallDecisionSource, +}; +use schemars::JsonSchema; +use tensorzero_core::endpoints::internal::autopilot::CreateEventGatewayRequest; + +/// Parameters for the auto-reject tool (not visible to LLM - internal use only). +#[derive(Debug, Clone, Default, Serialize, Deserialize, JsonSchema)] +pub struct AutoRejectToolCallParams {} + +/// Built-in tool to send `NotAvailable` authorization for unknown tool calls. +/// +/// This tool is spawned by the autopilot client when it receives a tool call +/// for a tool that doesn't exist in the deployment. It sends a `NotAvailable` +/// event to the autopilot API so the session can continue without hanging. +/// +/// Reserved name: `__auto_reject_tool_call__` +#[derive(Default)] +pub struct AutoRejectToolCallTool; + +impl ToolMetadata for AutoRejectToolCallTool { + type SideInfo = AutopilotSideInfo; + type Output = (); + type LlmParams = AutoRejectToolCallParams; + + fn name(&self) -> Cow<'static, str> { + Cow::Borrowed("__auto_reject_tool_call__") + } + + fn description(&self) -> Cow<'static, str> { + Cow::Borrowed( + "Internal tool for rejecting unknown tool calls. Not intended for direct use.", + ) + } + + fn parameters_schema(&self) -> ToolResult { + let schema = serde_json::json!({ + "type": "object", + "description": "Internal tool for auto-rejecting unknown tool calls.", + "properties": {}, + "required": [] + }); + + serde_json::from_value(schema).map_err(|e| { + NonControlToolError::SchemaGeneration { + message: e.to_string(), + } + .into() + }) + } +} + +#[async_trait] +impl TaskTool for AutoRejectToolCallTool { + async fn execute( + &self, + _llm_params: ::LlmParams, + side_info: ::SideInfo, + ctx: &mut ToolContext<'_>, + ) -> ToolResult<::Output> { + // Send ToolCallAuthorization::NotAvailable to autopilot API + ctx.client() + .create_autopilot_event( + side_info.session_id, + CreateEventGatewayRequest { + payload: EventPayload::ToolCallAuthorization( + EventPayloadToolCallAuthorization { + source: ToolCallDecisionSource::Automatic, + tool_call_event_id: side_info.tool_call_event_id, + status: ToolCallAuthorizationStatus::NotAvailable, + }, + ), + previous_user_message_event_id: None, + }, + ) + .await + .map_err(|e| AutopilotToolError::client_error("auto_reject_tool_call", e))?; + + Ok(()) + } +} diff --git a/internal/autopilot-tools/src/tools/prod/mod.rs b/internal/autopilot-tools/src/tools/prod/mod.rs index 05c27a6ea46..572843547df 100644 --- a/internal/autopilot-tools/src/tools/prod/mod.rs +++ b/internal/autopilot-tools/src/tools/prod/mod.rs @@ -3,6 +3,7 @@ //! This module contains production-ready tools that can be used by autopilot //! to perform actions like inference, feedback, and other operations. +mod auto_reject_tool_call; mod create_datapoints; mod create_datapoints_from_inferences; mod delete_datapoints; @@ -20,6 +21,7 @@ mod run_evaluation; mod update_datapoints; mod write_config; +pub use auto_reject_tool_call::AutoRejectToolCallTool; pub use create_datapoints::{CreateDatapointsTool, CreateDatapointsToolParams}; pub use create_datapoints_from_inferences::{ CreateDatapointsFromInferencesTool, CreateDatapointsFromInferencesToolParams, diff --git a/internal/autopilot-tools/src/visitor.rs b/internal/autopilot-tools/src/visitor.rs index 4e8407759d0..c81f84af09a 100644 --- a/internal/autopilot-tools/src/visitor.rs +++ b/internal/autopilot-tools/src/visitor.rs @@ -4,6 +4,9 @@ //! different registration strategies while ensuring the same set of tools is //! processed regardless of the visitor implementation. +use std::collections::HashSet; +use std::sync::Mutex; + use async_trait::async_trait; use durable_tools::{SimpleTool, TaskTool}; use serde::Serialize; @@ -39,9 +42,9 @@ pub trait ToolVisitor { /// /// For local execution, this typically calls `register_task_tool`. /// For remote execution, this wraps the tool in an adapter. - async fn visit_task_tool(&self) -> Result<(), Self::Error> + async fn visit_task_tool(&self, tool: T) -> Result<(), Self::Error> where - T: TaskTool + Default, + T: TaskTool, T::SideInfo: TryFrom + Serialize, >::Error: std::fmt::Display; @@ -55,3 +58,67 @@ pub trait ToolVisitor { T::SideInfo: TryFrom + Serialize, >::Error: std::fmt::Display; } + +/// A visitor that collects tool names from `for_each_tool`. +/// +/// This is used by `collect_tool_names` to derive the set of available +/// tool names from the single source of truth in `for_each_tool`. +pub struct ToolNameCollector { + names: Mutex>, +} + +impl ToolNameCollector { + pub fn new() -> Self { + Self { + names: Mutex::new(HashSet::new()), + } + } + + /// Consume the collector and return the collected tool names. + /// + /// If the mutex was poisoned (which would only happen if a panic occurred + /// while holding the lock), this returns the data anyway since we still + /// want to use it. + pub fn into_names(self) -> HashSet { + self.names.into_inner().unwrap_or_else(|e| e.into_inner()) + } +} + +impl Default for ToolNameCollector { + fn default() -> Self { + Self::new() + } +} + +#[async_trait] +impl ToolVisitor for ToolNameCollector { + type Error = String; + + async fn visit_task_tool(&self, tool: T) -> Result<(), Self::Error> + where + T: TaskTool, + T::SideInfo: TryFrom + Serialize, + >::Error: std::fmt::Display, + { + let mut names = self + .names + .lock() + .map_err(|e| format!("Failed to acquire lock: {e}"))?; + names.insert(tool.name().to_string()); + Ok(()) + } + + async fn visit_simple_tool(&self) -> Result<(), Self::Error> + where + T: SimpleTool + Default, + T::SideInfo: TryFrom + Serialize, + >::Error: std::fmt::Display, + { + let mut names = self + .names + .lock() + .map_err(|e| format!("Failed to acquire lock: {e}"))?; + names.insert(T::default().name().to_string()); + Ok(()) + } +} diff --git a/internal/autopilot-tools/tests/common/mod.rs b/internal/autopilot-tools/tests/common/mod.rs index 5c28b24aef1..3e39fa4a2d6 100644 --- a/internal/autopilot-tools/tests/common/mod.rs +++ b/internal/autopilot-tools/tests/common/mod.rs @@ -56,7 +56,7 @@ mock! { &self, session_id: Uuid, params: durable_tools::ListEventsParams, - ) -> Result; + ) -> Result; async fn list_autopilot_sessions( &self, diff --git a/internal/autopilot-worker/src/worker.rs b/internal/autopilot-worker/src/worker.rs index 5ed9a893935..56a1ca96ff9 100644 --- a/internal/autopilot-worker/src/worker.rs +++ b/internal/autopilot-worker/src/worker.rs @@ -6,6 +6,7 @@ use anyhow::Result; use async_trait::async_trait; use autopilot_client::AutopilotSideInfo; use autopilot_tools::ToolVisitor; +use autopilot_tools::tools::AutoRejectToolCallTool; use durable_tools::{ SimpleTool, TaskTool, TensorZeroClient, ToolError, ToolExecutor, Worker, WorkerOptions, }; @@ -92,6 +93,14 @@ impl AutopilotWorker { executor: &self.executor, }; autopilot_tools::for_each_tool(&visitor).await?; + + // Register internal tools directly without the ClientTaskToolWrapper. + // AutoRejectToolCallTool only writes a NotAvailable authorization - + // it doesn't need the wrapper to publish a tool_result. + self.executor + .register_task_tool_instance(AutoRejectToolCallTool) + .await?; + Ok(()) } @@ -146,14 +155,14 @@ struct LocalToolVisitor<'a> { impl ToolVisitor for LocalToolVisitor<'_> { type Error = ToolError; - async fn visit_task_tool(&self) -> Result<(), ToolError> + async fn visit_task_tool(&self, tool: T) -> Result<(), ToolError> where - T: TaskTool + Default, + T: TaskTool, T::SideInfo: TryFrom + Serialize, >::Error: std::fmt::Display, { self.executor - .register_task_tool_instance(ClientTaskToolWrapper::::default()) + .register_task_tool_instance(ClientTaskToolWrapper::new(tool)) .await?; Ok(()) } diff --git a/internal/autopilot-worker/src/wrapper.rs b/internal/autopilot-worker/src/wrapper.rs index 6c4cb5b81f1..13759996d4c 100644 --- a/internal/autopilot-worker/src/wrapper.rs +++ b/internal/autopilot-worker/src/wrapper.rs @@ -437,7 +437,7 @@ mod tests { &self, session_id: Uuid, params: durable_tools::ListEventsParams, - ) -> Result; + ) -> Result; async fn list_autopilot_sessions( &self, diff --git a/internal/durable-tools/src/lib.rs b/internal/durable-tools/src/lib.rs index 290a4b748f9..cbf7b09ab93 100644 --- a/internal/durable-tools/src/lib.rs +++ b/internal/durable-tools/src/lib.rs @@ -211,7 +211,8 @@ pub use tensorzero_client::MockTensorZeroClient; // Re-export autopilot types for use by tools pub use tensorzero_client::{ CreateEventGatewayRequest, CreateEventResponse, EventPayload, EventPayloadToolResult, - ListEventsParams, ListEventsResponse, ListSessionsParams, ListSessionsResponse, ToolOutcome, + GatewayListEventsResponse, ListEventsParams, ListSessionsParams, ListSessionsResponse, + ToolOutcome, }; // Re-export datapoint types for CRUD operations diff --git a/internal/durable-tools/src/tensorzero_client/client_ext.rs b/internal/durable-tools/src/tensorzero_client/client_ext.rs index ef7d4322ff7..d5bf53ee95e 100644 --- a/internal/durable-tools/src/tensorzero_client/client_ext.rs +++ b/internal/durable-tools/src/tensorzero_client/client_ext.rs @@ -6,6 +6,7 @@ use async_trait::async_trait; use autopilot_client::AutopilotError; +use autopilot_client::GatewayListEventsResponse; use tensorzero::{ Client, ClientExt, ClientInferenceParams, ClientMode, CreateDatapointRequest, CreateDatapointsFromInferenceRequestParams, CreateDatapointsResponse, DeleteDatapointsResponse, @@ -30,9 +31,9 @@ use uuid::Uuid; use crate::action::{ActionInput, ActionInputInfo, ActionResponse}; use super::{ - CreateEventGatewayRequest, CreateEventResponse, ListEventsParams, ListEventsResponse, - ListSessionsParams, ListSessionsResponse, RunEvaluationParams, RunEvaluationResponse, - TensorZeroClient, TensorZeroClientError, + CreateEventGatewayRequest, CreateEventResponse, ListEventsParams, ListSessionsParams, + ListSessionsResponse, RunEvaluationParams, RunEvaluationResponse, TensorZeroClient, + TensorZeroClientError, }; /// Implementation of `TensorZeroClient` for the TensorZero SDK `Client`. @@ -150,7 +151,7 @@ impl TensorZeroClient for Client { &self, session_id: Uuid, params: ListEventsParams, - ) -> Result { + ) -> Result { match self.mode() { ClientMode::HTTPGateway(http) => { let mut url = http diff --git a/internal/durable-tools/src/tensorzero_client/embedded.rs b/internal/durable-tools/src/tensorzero_client/embedded.rs index 146015a3858..53f52df3979 100644 --- a/internal/durable-tools/src/tensorzero_client/embedded.rs +++ b/internal/durable-tools/src/tensorzero_client/embedded.rs @@ -4,6 +4,7 @@ //! and wants to call inference and autopilot endpoints without HTTP overhead. use async_trait::async_trait; +use autopilot_client::GatewayListEventsResponse; use tensorzero::{ ClientInferenceParams, CreateDatapointRequest, CreateDatapointsFromInferenceRequestParams, CreateDatapointsResponse, DeleteDatapointsResponse, FeedbackParams, FeedbackResponse, @@ -32,9 +33,9 @@ use uuid::Uuid; use crate::action::{ActionInput, ActionInputInfo, ActionResponse}; use super::{ - CreateEventGatewayRequest, CreateEventResponse, ListEventsParams, ListEventsResponse, - ListSessionsParams, ListSessionsResponse, RunEvaluationParams, RunEvaluationResponse, - TensorZeroClient, TensorZeroClientError, + CreateEventGatewayRequest, CreateEventResponse, ListEventsParams, ListSessionsParams, + ListSessionsResponse, RunEvaluationParams, RunEvaluationResponse, TensorZeroClient, + TensorZeroClientError, }; /// TensorZero client that uses an existing gateway's state directly. @@ -140,7 +141,7 @@ impl TensorZeroClient for EmbeddedClient { &self, session_id: Uuid, params: ListEventsParams, - ) -> Result { + ) -> Result { let autopilot_client = self .app_state .autopilot_client diff --git a/internal/durable-tools/src/tensorzero_client/mod.rs b/internal/durable-tools/src/tensorzero_client/mod.rs index 6e6969dec34..5d449164952 100644 --- a/internal/durable-tools/src/tensorzero_client/mod.rs +++ b/internal/durable-tools/src/tensorzero_client/mod.rs @@ -35,8 +35,8 @@ pub use embedded::EmbeddedClient; // Re-export autopilot types for use by tools pub use autopilot_client::{ - CreateEventResponse, EventPayload, EventPayloadToolResult, ListEventsParams, - ListEventsResponse, ListSessionsParams, ListSessionsResponse, ToolOutcome, + CreateEventResponse, EventPayload, EventPayloadToolResult, GatewayListEventsResponse, + ListEventsParams, ListSessionsParams, ListSessionsResponse, ToolOutcome, }; pub use tensorzero_core::endpoints::internal::autopilot::CreateEventGatewayRequest; @@ -137,11 +137,14 @@ pub trait TensorZeroClient: Send + Sync + 'static { ) -> Result; /// List events in an autopilot session. + /// + /// Returns `GatewayListEventsResponse` which uses narrower types that exclude + /// `NotAvailable` authorization status. async fn list_autopilot_events( &self, session_id: Uuid, params: ListEventsParams, - ) -> Result; + ) -> Result; /// List autopilot sessions. async fn list_autopilot_sessions( diff --git a/internal/tensorzero-node/lib/bindings/Event.ts b/internal/tensorzero-node/lib/bindings/Event.ts deleted file mode 100644 index cf90d064afb..00000000000 --- a/internal/tensorzero-node/lib/bindings/Event.ts +++ /dev/null @@ -1,12 +0,0 @@ -// This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually. -import type { EventPayload } from "./EventPayload"; - -/** - * An event within a session. - */ -export type Event = { - id: string; - payload: EventPayload; - session_id: string; - created_at: string; -}; diff --git a/internal/tensorzero-node/lib/bindings/EventPayload.ts b/internal/tensorzero-node/lib/bindings/EventPayload.ts index 30945af24fb..fbcc635ebf1 100644 --- a/internal/tensorzero-node/lib/bindings/EventPayload.ts +++ b/internal/tensorzero-node/lib/bindings/EventPayload.ts @@ -7,7 +7,9 @@ import type { EventPayloadToolCallAuthorization } from "./EventPayloadToolCallAu import type { EventPayloadToolResult } from "./EventPayloadToolResult"; /** - * The payload of an event. + * Internal event payload type - consumers should use `GatewayEventPayload` instead. + * + * Note: TS derive is needed for types that reference this, but we don't export it. */ export type EventPayload = | ({ type: "message" } & EventPayloadMessage) diff --git a/internal/tensorzero-node/lib/bindings/GatewayEvent.ts b/internal/tensorzero-node/lib/bindings/GatewayEvent.ts new file mode 100644 index 00000000000..cb35d4a7816 --- /dev/null +++ b/internal/tensorzero-node/lib/bindings/GatewayEvent.ts @@ -0,0 +1,14 @@ +// This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually. +import type { GatewayEventPayload } from "./GatewayEventPayload"; + +/** + * An event as seen by gateway consumers. + * + * Uses `GatewayEventPayload` which excludes `NotAvailable` authorization status. + */ +export type GatewayEvent = { + id: string; + payload: GatewayEventPayload; + session_id: string; + created_at: string; +}; diff --git a/internal/tensorzero-node/lib/bindings/GatewayEventPayload.ts b/internal/tensorzero-node/lib/bindings/GatewayEventPayload.ts new file mode 100644 index 00000000000..0db11356820 --- /dev/null +++ b/internal/tensorzero-node/lib/bindings/GatewayEventPayload.ts @@ -0,0 +1,23 @@ +// This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually. +import type { EventPayloadError } from "./EventPayloadError"; +import type { EventPayloadMessage } from "./EventPayloadMessage"; +import type { EventPayloadStatusUpdate } from "./EventPayloadStatusUpdate"; +import type { EventPayloadToolCall } from "./EventPayloadToolCall"; +import type { EventPayloadToolResult } from "./EventPayloadToolResult"; +import type { GatewayEventPayloadToolCallAuthorization } from "./GatewayEventPayloadToolCallAuthorization"; + +/** + * Event payload as seen by gateway consumers. + * + * Uses `GatewayEventPayloadToolCallAuthorization` which excludes `NotAvailable` status. + */ +export type GatewayEventPayload = + | ({ type: "message" } & EventPayloadMessage) + | ({ type: "error" } & EventPayloadError) + | ({ type: "status_update" } & EventPayloadStatusUpdate) + | ({ type: "tool_call" } & EventPayloadToolCall) + | ({ + type: "tool_call_authorization"; + } & GatewayEventPayloadToolCallAuthorization) + | ({ type: "tool_result" } & EventPayloadToolResult) + | { type: "unknown" }; diff --git a/internal/tensorzero-node/lib/bindings/GatewayEventPayloadToolCallAuthorization.ts b/internal/tensorzero-node/lib/bindings/GatewayEventPayloadToolCallAuthorization.ts new file mode 100644 index 00000000000..3a434fa707c --- /dev/null +++ b/internal/tensorzero-node/lib/bindings/GatewayEventPayloadToolCallAuthorization.ts @@ -0,0 +1,14 @@ +// This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually. +import type { GatewayToolCallAuthorizationStatus } from "./GatewayToolCallAuthorizationStatus"; +import type { ToolCallDecisionSource } from "./ToolCallDecisionSource"; + +/** + * Tool call authorization payload as seen by gateway consumers. + * + * Uses `GatewayToolCallAuthorizationStatus` which excludes `NotAvailable`. + */ +export type GatewayEventPayloadToolCallAuthorization = { + source: ToolCallDecisionSource; + tool_call_event_id: string; + status: GatewayToolCallAuthorizationStatus; +}; diff --git a/internal/tensorzero-node/lib/bindings/ListEventsResponse.ts b/internal/tensorzero-node/lib/bindings/GatewayListEventsResponse.ts similarity index 67% rename from internal/tensorzero-node/lib/bindings/ListEventsResponse.ts rename to internal/tensorzero-node/lib/bindings/GatewayListEventsResponse.ts index 428d90c27a3..9f7b61ae35f 100644 --- a/internal/tensorzero-node/lib/bindings/ListEventsResponse.ts +++ b/internal/tensorzero-node/lib/bindings/GatewayListEventsResponse.ts @@ -1,12 +1,14 @@ // This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually. import type { AutopilotStatus } from "./AutopilotStatus"; -import type { Event } from "./Event"; +import type { GatewayEvent } from "./GatewayEvent"; /** - * Response from listing events. + * Response from listing events as seen by gateway consumers. + * + * Uses `GatewayEvent` which excludes `NotAvailable` authorization status. */ -export type ListEventsResponse = { - events: Array; +export type GatewayListEventsResponse = { + events: Array; /** * The most recent `message` event with role `user` in this session. */ @@ -21,5 +23,5 @@ export type ListEventsResponse = { * These may be duplicates of some of the values in events. * All EventPayloads in these Events should be of type ToolCall. */ - pending_tool_calls: Array; + pending_tool_calls: Array; }; diff --git a/internal/tensorzero-node/lib/bindings/GatewayStreamUpdate.ts b/internal/tensorzero-node/lib/bindings/GatewayStreamUpdate.ts new file mode 100644 index 00000000000..318e6474e17 --- /dev/null +++ b/internal/tensorzero-node/lib/bindings/GatewayStreamUpdate.ts @@ -0,0 +1,13 @@ +// This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually. +import type { AutopilotStatus } from "./AutopilotStatus"; +import type { GatewayEvent } from "./GatewayEvent"; + +/** + * Stream update as seen by gateway consumers. + * + * Uses `GatewayEvent` which excludes `NotAvailable` authorization status. + */ +export type GatewayStreamUpdate = { + event: GatewayEvent; + status: AutopilotStatus; +}; diff --git a/internal/tensorzero-node/lib/bindings/GatewayToolCallAuthorizationStatus.ts b/internal/tensorzero-node/lib/bindings/GatewayToolCallAuthorizationStatus.ts new file mode 100644 index 00000000000..b09f64de1eb --- /dev/null +++ b/internal/tensorzero-node/lib/bindings/GatewayToolCallAuthorizationStatus.ts @@ -0,0 +1,11 @@ +// This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually. + +/** + * Authorization status for tool calls as seen by gateway consumers. + * + * This is a narrower type than `ToolCallAuthorizationStatus` that excludes + * `NotAvailable` since that status is filtered out before reaching consumers. + */ +export type GatewayToolCallAuthorizationStatus = + | { type: "approved" } + | { type: "rejected"; reason: string }; diff --git a/internal/tensorzero-node/lib/bindings/StreamUpdate.ts b/internal/tensorzero-node/lib/bindings/StreamUpdate.ts deleted file mode 100644 index 6065c156843..00000000000 --- a/internal/tensorzero-node/lib/bindings/StreamUpdate.ts +++ /dev/null @@ -1,5 +0,0 @@ -// This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually. -import type { AutopilotStatus } from "./AutopilotStatus"; -import type { Event } from "./Event"; - -export type StreamUpdate = { event: Event; status: AutopilotStatus }; diff --git a/internal/tensorzero-node/lib/bindings/ToolCallAuthorizationStatus.ts b/internal/tensorzero-node/lib/bindings/ToolCallAuthorizationStatus.ts index 572ba6b7eff..7870cb35694 100644 --- a/internal/tensorzero-node/lib/bindings/ToolCallAuthorizationStatus.ts +++ b/internal/tensorzero-node/lib/bindings/ToolCallAuthorizationStatus.ts @@ -2,4 +2,5 @@ export type ToolCallAuthorizationStatus = | { type: "approved" } - | { type: "rejected"; reason: string }; + | { type: "rejected"; reason: string } + | { type: "not_available" }; diff --git a/internal/tensorzero-node/lib/bindings/ToolCallDecisionSource.ts b/internal/tensorzero-node/lib/bindings/ToolCallDecisionSource.ts index 80f89dabfbc..f22c2405ad3 100644 --- a/internal/tensorzero-node/lib/bindings/ToolCallDecisionSource.ts +++ b/internal/tensorzero-node/lib/bindings/ToolCallDecisionSource.ts @@ -1,3 +1,3 @@ // This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually. -export type ToolCallDecisionSource = { type: "ui" }; +export type ToolCallDecisionSource = { type: "ui" } | { type: "automatic" }; diff --git a/internal/tensorzero-node/lib/bindings/index.ts b/internal/tensorzero-node/lib/bindings/index.ts index 32e47ee517c..f812eea9c84 100644 --- a/internal/tensorzero-node/lib/bindings/index.ts +++ b/internal/tensorzero-node/lib/bindings/index.ts @@ -98,7 +98,6 @@ export * from "./EvaluationRunStatsResponse"; export * from "./EvaluationRunSuccessEvent"; export * from "./EvaluationStatistics"; export * from "./EvaluatorConfig"; -export * from "./Event"; export * from "./EventPayload"; export * from "./EventPayloadError"; export * from "./EventPayloadMessage"; @@ -139,6 +138,12 @@ export * from "./GCPVertexGeminiSFTConfig"; export * from "./GCPVertexGeminiSFTJobHandle"; export * from "./GEPAConfig"; export * from "./GEPAJobHandle"; +export * from "./GatewayEvent"; +export * from "./GatewayEventPayload"; +export * from "./GatewayEventPayloadToolCallAuthorization"; +export * from "./GatewayListEventsResponse"; +export * from "./GatewayStreamUpdate"; +export * from "./GatewayToolCallAuthorizationStatus"; export * from "./GetCumulativeFeedbackTimeseriesResponse"; export * from "./GetDatapointCountResponse"; export * from "./GetDatapointParams"; @@ -207,7 +212,6 @@ export * from "./ListDatasetsResponse"; export * from "./ListEpisodesResponse"; export * from "./ListEvaluationRunsResponse"; export * from "./ListEventsParams"; -export * from "./ListEventsResponse"; export * from "./ListFunctionsWithInferenceCountResponse"; export * from "./ListInferenceMetadataResponse"; export * from "./ListInferencesRequest"; @@ -315,7 +319,6 @@ export * from "./StoredJsonInference"; export * from "./StoredOutput"; export * from "./StoredRequestMessage"; export * from "./StreamEventsParams"; -export * from "./StreamUpdate"; export * from "./StreamingTimeouts"; export * from "./System"; export * from "./TGIProvider"; diff --git a/tensorzero-core/src/client/mod.rs b/tensorzero-core/src/client/mod.rs index 0006b76f404..a28c52cbfde 100644 --- a/tensorzero-core/src/client/mod.rs +++ b/tensorzero-core/src/client/mod.rs @@ -1,3 +1,4 @@ +use std::collections::HashSet; use std::{env, fmt::Display, future::Future, path::PathBuf, sync::Arc, time::Duration}; use crate::config::ConfigFileGlob; @@ -632,6 +633,7 @@ impl ClientBuilder { valkey_connection_info, http_client, self.drop_wrapper, + HashSet::new(), // available_tools not needed for embedded client ) .await .map_err(|e| { @@ -677,6 +679,7 @@ impl ClientBuilder { valkey_connection_info.clone(), http_client.clone(), self.drop_wrapper, + HashSet::new(), // available_tools not needed for embedded client ) .await .map_err(|e| { diff --git a/tensorzero-core/src/endpoints/internal/autopilot.rs b/tensorzero-core/src/endpoints/internal/autopilot.rs index f11ae34a9d6..61ff8ed24fc 100644 --- a/tensorzero-core/src/endpoints/internal/autopilot.rs +++ b/tensorzero-core/src/endpoints/internal/autopilot.rs @@ -18,8 +18,8 @@ use uuid::Uuid; use autopilot_client::{ ApproveAllToolCallsRequest, ApproveAllToolCallsResponse, AutopilotClient, CreateEventRequest, - CreateEventResponse, EventPayload, ListEventsParams, ListEventsResponse, ListSessionsParams, - ListSessionsResponse, StreamEventsParams, StreamUpdate, + CreateEventResponse, EventPayload, GatewayListEventsResponse, GatewayStreamUpdate, + ListEventsParams, ListSessionsParams, ListSessionsResponse, StreamEventsParams, }; use crate::endpoints::status::TENSORZERO_VERSION; @@ -95,11 +95,13 @@ pub async fn list_sessions( /// List events for a session from the Autopilot API. /// /// This is the core function called by both the HTTP handler and embedded client. +/// Returns `GatewayListEventsResponse` which uses narrower types that exclude +/// `NotAvailable` authorization status. pub async fn list_events( autopilot_client: &AutopilotClient, session_id: Uuid, params: ListEventsParams, -) -> Result { +) -> Result { autopilot_client .list_events(session_id, params) .await @@ -161,7 +163,7 @@ pub async fn list_events_handler( State(app_state): AppState, Path(session_id): Path, Query(params): Query, -) -> Result, Error> { +) -> Result, Error> { let client = get_autopilot_client(&app_state)?; let response = list_events(&client, session_id, params).await?; Ok(Json(response)) @@ -260,7 +262,7 @@ pub async fn stream_events_handler( // Convert the autopilot event stream to SSE events let sse_stream = stream.map( - |result: Result| match result { + |result: Result| match result { Ok(event) => match serde_json::to_string(&event) { Ok(data) => Ok(SseEvent::default().event("event").data(data)), Err(e) => { diff --git a/tensorzero-core/src/error/mod.rs b/tensorzero-core/src/error/mod.rs index 7aa2b5a6773..09e77807b05 100644 --- a/tensorzero-core/src/error/mod.rs +++ b/tensorzero-core/src/error/mod.rs @@ -1757,6 +1757,10 @@ impl From for Error { message: format!("Spawn error: {e}"), status_code: None, }), + autopilot_client::AutopilotError::Database(e) => Self::new(ErrorDetails::Autopilot { + message: format!("Database error: {e}"), + status_code: None, + }), autopilot_client::AutopilotError::MissingConfig(field) => { Self::new(ErrorDetails::Autopilot { message: format!("Missing config: {field}"), @@ -1780,6 +1784,12 @@ impl From for Error { status_code: None, }) } + autopilot_client::AutopilotError::Internal(message) => { + Self::new(ErrorDetails::Autopilot { + message: format!("Internal error: {message}"), + status_code: None, + }) + } } } } diff --git a/tensorzero-core/src/utils/gateway.rs b/tensorzero-core/src/utils/gateway.rs index c08f9ae8734..3b4a4445e02 100644 --- a/tensorzero-core/src/utils/gateway.rs +++ b/tensorzero-core/src/utils/gateway.rs @@ -1,3 +1,4 @@ +use std::collections::HashSet; use std::future::IntoFuture; use std::net::SocketAddr; use std::sync::Arc; @@ -189,7 +190,10 @@ fn create_auth_cache_from_config(config: &Config) -> Option Result { + pub async fn new( + config: UnwrittenConfig, + available_tools: HashSet, + ) -> Result { let clickhouse_url = std::env::var("TENSORZERO_CLICKHOUSE_URL").ok(); let postgres_url = std::env::var("TENSORZERO_POSTGRES_URL").ok(); let valkey_url = std::env::var("TENSORZERO_VALKEY_URL").ok(); @@ -198,6 +202,7 @@ impl GatewayHandle { clickhouse_url, postgres_url, valkey_url, + available_tools, )) .await } @@ -207,6 +212,7 @@ impl GatewayHandle { clickhouse_url: Option, postgres_url: Option, valkey_url: Option, + available_tools: HashSet, ) -> Result { let clickhouse_connection_info = setup_clickhouse(&config, clickhouse_url, false).await?; let config = Arc::new(Box::pin(config.into_config(&clickhouse_connection_info)).await?); @@ -220,6 +226,7 @@ impl GatewayHandle { valkey_connection_info, http_client, None, + available_tools, ) .await } @@ -272,6 +279,7 @@ impl GatewayHandle { valkey_connection_info: ValkeyConnectionInfo, http_client: TensorzeroHttpClient, drop_wrapper: Option, + available_tools: HashSet, ) -> Result { let rate_limiting_manager = Arc::new(RateLimitingManager::new_from_connections( Arc::new(config.rate_limiting.clone()), @@ -313,8 +321,12 @@ impl GatewayHandle { .build(), ); - let autopilot_client = - setup_autopilot_client(&postgres_connection_info, deployment_id.as_ref()).await?; + let autopilot_client = setup_autopilot_client( + &postgres_connection_info, + deployment_id.as_ref(), + available_tools, + ) + .await?; Ok(Self { app_state: AppStateData { @@ -536,6 +548,7 @@ pub async fn setup_valkey(valkey_url: Option<&str>) -> Result, + available_tools: HashSet, ) -> Result>, Error> { match std::env::var("TENSORZERO_AUTOPILOT_API_KEY") { Ok(api_key) => { @@ -557,10 +570,11 @@ async fn setup_autopilot_client( let queue_name = std::env::var("TENSORZERO_AUTOPILOT_QUEUE_NAME") .unwrap_or_else(|_| "autopilot".to_string()); - let mut builder = AutopilotClient::builder().api_key(api_key); - builder = builder + let mut builder = AutopilotClient::builder() + .api_key(api_key) .spawn_pool(pool.clone()) - .spawn_queue_name(queue_name); + .spawn_queue_name(queue_name) + .available_tools(available_tools); // Allow custom base URL for testing if let Ok(base_url) = std::env::var("TENSORZERO_AUTOPILOT_BASE_URL") { @@ -684,6 +698,7 @@ pub async fn start_openai_compatible_gateway( clickhouse_url, postgres_url, valkey_url, + HashSet::new(), // available_tools )) .await?; @@ -1044,6 +1059,7 @@ mod tests { ValkeyConnectionInfo::Disabled, http_client, None, + HashSet::new(), // available_tools ) .await .expect("Gateway setup should succeed when rate limiting has no rules"); diff --git a/tensorzero-core/tests/e2e/feedback.rs b/tensorzero-core/tests/e2e/feedback.rs index 5443c9eb4f9..e87b2e68d01 100644 --- a/tensorzero-core/tests/e2e/feedback.rs +++ b/tensorzero-core/tests/e2e/feedback.rs @@ -1,6 +1,9 @@ use reqwest::{Client, StatusCode}; use serde_json::{Value, json}; -use std::{collections::HashMap, sync::Arc}; +use std::{ + collections::{HashMap, HashSet}, + sync::Arc, +}; use tensorzero_core::{ config::{Config, MetricConfig, MetricConfigLevel, MetricConfigOptimize, MetricConfigType}, db::{ @@ -256,6 +259,7 @@ async fn e2e_test_comment_feedback_validation_disabled() { ValkeyConnectionInfo::Disabled, TensorzeroHttpClient::new_testing().unwrap(), None, + HashSet::new(), // available_tools ) .await .unwrap(); @@ -1562,6 +1566,7 @@ async fn e2e_test_float_feedback_validation_disabled() { ValkeyConnectionInfo::Disabled, TensorzeroHttpClient::new_testing().unwrap(), None, + HashSet::new(), // available_tools ) .await .unwrap(); @@ -1903,6 +1908,7 @@ async fn e2e_test_boolean_feedback_validation_disabled() { ValkeyConnectionInfo::Disabled, TensorzeroHttpClient::new_testing().unwrap(), None, + HashSet::new(), // available_tools ) .await .unwrap(); diff --git a/tensorzero-core/tests/e2e/howdy.rs b/tensorzero-core/tests/e2e/howdy.rs index 5a85c4e115f..f3e2b638cc9 100644 --- a/tensorzero-core/tests/e2e/howdy.rs +++ b/tensorzero-core/tests/e2e/howdy.rs @@ -1,4 +1,5 @@ #![expect(clippy::print_stdout)] +use std::collections::HashSet; use std::sync::Arc; use crate::clickhouse::get_clean_clickhouse; @@ -57,6 +58,7 @@ async fn get_embedded_client(clickhouse: ClickHouseConnectionInfo) -> tensorzero ValkeyConnectionInfo::Disabled, TensorzeroHttpClient::new_testing().unwrap(), None, + HashSet::new(), // available_tools ) .await .unwrap(); diff --git a/ui/app/components/autopilot/EventStream.stories.tsx b/ui/app/components/autopilot/EventStream.stories.tsx index 89ebb6926fb..493c7db2710 100644 --- a/ui/app/components/autopilot/EventStream.stories.tsx +++ b/ui/app/components/autopilot/EventStream.stories.tsx @@ -1,18 +1,18 @@ import type { Meta, StoryObj } from "@storybook/react-vite"; import EventStream from "./EventStream"; -import type { Event } from "~/types/tensorzero"; +import type { GatewayEvent } from "~/types/tensorzero"; const baseTime = new Date("2026-04-12T10:00:00Z").getTime(); const sessionId = "d1a0b0c0-0000-0000-0000-000000000001"; -function buildEvent(event: Event, index: number): Event { +function buildEvent(event: GatewayEvent, index: number): GatewayEvent { return { ...event, created_at: new Date(baseTime + index * 60 * 1000).toISOString(), }; } -const conversationEvents: Event[] = [ +const conversationEvents: GatewayEvent[] = [ buildEvent( { id: "e2a3f5d6-7b8c-4d9e-8f01-1234567890a1", @@ -46,7 +46,7 @@ const conversationEvents: Event[] = [ ), ]; -const toolingEvents: Event[] = [ +const toolingEvents: GatewayEvent[] = [ buildEvent( { id: "0a1b2c3d-4e5f-4a6b-8c7d-3456789012c3", @@ -101,7 +101,7 @@ const toolingEvents: Event[] = [ ), ]; -const mixedEvents: Event[] = [ +const mixedEvents: GatewayEvent[] = [ buildEvent( { id: "a1b2c3d4-5e6f-4a7b-8c9d-0e1f2a3b4c5d", @@ -297,7 +297,7 @@ const longToolResult = [ longText, ].join("\n\n"); -const longFormEvents: Event[] = [ +const longFormEvents: GatewayEvent[] = [ buildEvent( { id: "7b8c9d0e-1f2a-4b3c-8d4e-0123456789d0", diff --git a/ui/app/components/autopilot/EventStream.tsx b/ui/app/components/autopilot/EventStream.tsx index 2582a7c952d..02ecc29a432 100644 --- a/ui/app/components/autopilot/EventStream.tsx +++ b/ui/app/components/autopilot/EventStream.tsx @@ -9,9 +9,9 @@ import { } from "~/components/ui/tooltip"; import type { AutopilotStatus, - Event, - EventPayload, EventPayloadMessageContent, + GatewayEvent, + GatewayEventPayload, } from "~/types/tensorzero"; import { cn } from "~/utils/common"; @@ -39,7 +39,7 @@ type EventSummary = { }; type EventStreamProps = { - events: Event[]; + events: GatewayEvent[]; className?: string; isLoadingOlder?: boolean; hasReachedStart?: boolean; @@ -76,19 +76,19 @@ export function ToolEventId({ id }: { id: string }) { * Tool event payload types - events related to tool calls. */ export type ToolEventPayload = Extract< - EventPayload, + GatewayEventPayload, { type: "tool_call" | "tool_call_authorization" | "tool_result" } >; /** * An event with a tool-related payload. */ -export type ToolEvent = Event & { payload: ToolEventPayload }; +export type ToolEvent = GatewayEvent & { payload: ToolEventPayload }; /** * Type guard to check if an event is a tool event. */ -export function isToolEvent(event: Event): event is ToolEvent { +export function isToolEvent(event: GatewayEvent): event is ToolEvent { return ( event.payload.type === "tool_call" || event.payload.type === "tool_call_authorization" || @@ -129,7 +129,7 @@ function formatToolError(error: unknown): string { return JSON.stringify(error); } -function summarizeEvent(event: Event): EventSummary { +function summarizeEvent(event: GatewayEvent): EventSummary { const { payload } = event; switch (payload.type) { @@ -174,7 +174,7 @@ function summarizeEvent(event: Event): EventSummary { } } -function renderEventTitle(event: Event) { +function renderEventTitle(event: GatewayEvent) { const { payload } = event; switch (payload.type) { @@ -321,7 +321,7 @@ function EventItem({ event, isPending = false, }: { - event: Event; + event: GatewayEvent; isPending?: boolean; }) { const summary = summarizeEvent(event); diff --git a/ui/app/components/autopilot/PendingToolCallCard.tsx b/ui/app/components/autopilot/PendingToolCallCard.tsx index c4c7e7e76c8..831ad49c57a 100644 --- a/ui/app/components/autopilot/PendingToolCallCard.tsx +++ b/ui/app/components/autopilot/PendingToolCallCard.tsx @@ -1,12 +1,12 @@ import { ChevronRight } from "lucide-react"; import { useState } from "react"; import { TableItemTime } from "~/components/ui/TableItems"; -import type { Event } from "~/types/tensorzero"; +import type { GatewayEvent } from "~/types/tensorzero"; import { cn } from "~/utils/common"; import { getToolCallEventId, isToolEvent, ToolEventId } from "./EventStream"; type PendingToolCallCardProps = { - event: Event; + event: GatewayEvent; isLoading: boolean; loadingAction?: "approving" | "rejecting"; onAuthorize: (approved: boolean) => void; diff --git a/ui/app/hooks/useAutopilotEventStream.ts b/ui/app/hooks/useAutopilotEventStream.ts index c4681c0967b..684940ddfd5 100644 --- a/ui/app/hooks/useAutopilotEventStream.ts +++ b/ui/app/hooks/useAutopilotEventStream.ts @@ -1,22 +1,26 @@ import { useCallback, useEffect, useRef, useState } from "react"; -import type { AutopilotStatus, Event, StreamUpdate } from "~/types/tensorzero"; +import type { + AutopilotStatus, + GatewayEvent, + GatewayStreamUpdate, +} from "~/types/tensorzero"; interface UseAutopilotEventStreamOptions { sessionId: string; - initialEvents: Event[]; - initialPendingToolCalls: Event[]; + initialEvents: GatewayEvent[]; + initialPendingToolCalls: GatewayEvent[]; initialStatus: AutopilotStatus; enabled?: boolean; } interface UseAutopilotEventStreamResult { - events: Event[]; - pendingToolCalls: Event[]; + events: GatewayEvent[]; + pendingToolCalls: GatewayEvent[]; status: AutopilotStatus; isConnected: boolean; error: string | null; isRetrying: boolean; - prependEvents: (newEvents: Event[]) => void; + prependEvents: (newEvents: GatewayEvent[]) => void; } const RETRY_DELAY_MS = 5000; @@ -32,8 +36,8 @@ export function useAutopilotEventStream({ initialStatus, enabled = true, }: UseAutopilotEventStreamOptions): UseAutopilotEventStreamResult { - const [events, setEvents] = useState(initialEvents); - const [pendingToolCalls, setPendingToolCalls] = useState( + const [events, setEvents] = useState(initialEvents); + const [pendingToolCalls, setPendingToolCalls] = useState( initialPendingToolCalls, ); const [status, setStatus] = useState(initialStatus); @@ -122,7 +126,7 @@ export function useAutopilotEventStream({ const data = line.slice(6).trim(); if (data) { try { - const streamUpdate = JSON.parse(data) as StreamUpdate; + const streamUpdate = JSON.parse(data) as GatewayStreamUpdate; const event = streamUpdate.event; lastEventIdRef.current = event.id; @@ -243,7 +247,7 @@ export function useAutopilotEventStream({ }, [initialEvents, initialPendingToolCalls, initialStatus]); // Allow prepending older events (for reverse infinite scroll) - const prependEvents = useCallback((newEvents: Event[]) => { + const prependEvents = useCallback((newEvents: GatewayEvent[]) => { setEvents((prev) => { // Create a set of existing event IDs for deduplication const existingIds = new Set(prev.map((e) => e.id)); diff --git a/ui/app/routes/api/autopilot/sessions/$session_id/events/authorize.route.ts b/ui/app/routes/api/autopilot/sessions/$session_id/events/authorize.route.ts index 6e05c5864d7..b3b8dab5633 100644 --- a/ui/app/routes/api/autopilot/sessions/$session_id/events/authorize.route.ts +++ b/ui/app/routes/api/autopilot/sessions/$session_id/events/authorize.route.ts @@ -1,11 +1,11 @@ import type { ActionFunctionArgs } from "react-router"; import { getAutopilotClient } from "~/utils/tensorzero.server"; -import type { ToolCallAuthorizationStatus } from "~/types/tensorzero"; +import type { GatewayToolCallAuthorizationStatus } from "~/types/tensorzero"; import { logger } from "~/utils/logger"; type AuthorizeRequest = { tool_call_event_id: string; - status: ToolCallAuthorizationStatus; + status: GatewayToolCallAuthorizationStatus; }; /** diff --git a/ui/app/routes/autopilot/sessions/$session_id/route.tsx b/ui/app/routes/autopilot/sessions/$session_id/route.tsx index 617118c8c52..f93c3c9c7c7 100644 --- a/ui/app/routes/autopilot/sessions/$session_id/route.tsx +++ b/ui/app/routes/autopilot/sessions/$session_id/route.tsx @@ -25,7 +25,7 @@ import { ChatInput } from "~/components/autopilot/ChatInput"; import { logger } from "~/utils/logger"; import { getAutopilotClient } from "~/utils/tensorzero.server"; import { useAutopilotEventStream } from "~/hooks/useAutopilotEventStream"; -import type { AutopilotStatus, Event } from "~/types/tensorzero"; +import type { AutopilotStatus, GatewayEvent } from "~/types/tensorzero"; import { useToast } from "~/hooks/use-toast"; // Nil UUID for creating new sessions @@ -42,9 +42,9 @@ export const handle: RouteHandle = { const EVENTS_PER_PAGE = 20; export type EventsData = { - events: Event[]; + events: GatewayEvent[]; hasMoreEvents: boolean; - pendingToolCalls: Event[]; + pendingToolCalls: GatewayEvent[]; status: AutopilotStatus; }; @@ -59,9 +59,9 @@ export async function loader({ params }: Route.LoaderArgs) { return { sessionId: "new", eventsData: { - events: [] as Event[], + events: [] as GatewayEvent[], hasMoreEvents: false, - pendingToolCalls: [] as Event[], + pendingToolCalls: [] as GatewayEvent[], status: { status: "idle" } as AutopilotStatus, }, isNewSession: true, @@ -367,7 +367,9 @@ function EventStreamContent({ return; } - const responseData = (await response.json()) as { events: Event[] }; + const responseData = (await response.json()) as { + events: GatewayEvent[]; + }; logger.debug( `Loaded ${responseData.events.length} older events (requested ${EVENTS_PER_PAGE + 1})`, diff --git a/ui/app/utils/tensorzero/autopilot-client.ts b/ui/app/utils/tensorzero/autopilot-client.ts index 33aea2b08b3..56652166ed9 100644 --- a/ui/app/utils/tensorzero/autopilot-client.ts +++ b/ui/app/utils/tensorzero/autopilot-client.ts @@ -4,12 +4,12 @@ import type { ApproveAllToolCallsResponse, CreateEventGatewayRequest, CreateEventResponse, + GatewayListEventsResponse, + GatewayStreamUpdate, ListEventsParams, - ListEventsResponse, ListSessionsParams, ListSessionsResponse, StreamEventsParams, - StreamUpdate, } from "~/types/tensorzero"; /** @@ -42,7 +42,7 @@ export class AutopilotClient extends BaseTensorZeroClient { async listAutopilotEvents( sessionId: string, params?: ListEventsParams, - ): Promise { + ): Promise { const searchParams = new URLSearchParams(); if (params?.limit) searchParams.set("limit", params.limit.toString()); if (params?.before) searchParams.set("before", params.before); @@ -54,7 +54,7 @@ export class AutopilotClient extends BaseTensorZeroClient { const message = await this.getErrorText(response); this.handleHttpError({ message, response }); } - return (await response.json()) as ListEventsResponse; + return (await response.json()) as GatewayListEventsResponse; } /** @@ -99,12 +99,12 @@ export class AutopilotClient extends BaseTensorZeroClient { /** * Streams events for an autopilot session using Server-Sent Events. - * Returns an async generator that yields Event objects. + * Returns an async generator that yields GatewayStreamUpdate objects. */ async *streamAutopilotEvents( sessionId: string, params?: StreamEventsParams, - ): AsyncGenerator { + ): AsyncGenerator { const searchParams = new URLSearchParams(); if (params?.last_event_id) searchParams.set("last_event_id", params.last_event_id); @@ -145,7 +145,7 @@ export class AutopilotClient extends BaseTensorZeroClient { if (line.startsWith("data: ")) { const data = line.slice(6); if (data) { - const event = JSON.parse(data) as StreamUpdate; + const event = JSON.parse(data) as GatewayStreamUpdate; yield event; } } From 6c032e572cf193a1bd29c78638000b3831244a57 Mon Sep 17 00:00:00 2001 From: simeonlee <15081451+simeonlee@users.noreply.github.com> Date: Mon, 26 Jan 2026 23:10:20 -0500 Subject: [PATCH 15/46] Use hyphens for workflow-evaluations UI route (#5659) * Use hyphens for workflow-evaluations UI route Changes `/workflow_evaluations` to `/workflow-evaluations` in the UI to follow web URL conventions (hyphens for user-facing URLs). - Rename routes directory and update route definitions - Update URL helper functions and tests - Update sidebar and dashboard navigation links - Update e2e test URLs and filenames - Document URL conventions in ui/AGENTS.md Backend API routes (e.g., `/internal/workflow_evaluations`) retain underscores for consistency with gateway API conventions (like Stripe/OpenAI). Closes #5381 * Fix import path for WorkflowEvalRunSelector * Remove unnecessary Stripe/OpenAI reference from docs * Address PR feedback: fix breadcrumb hrefs and move URL docs - Fix breadcrumb hrefs to use hyphenated URL (/workflow-evaluations) - Move URL conventions documentation from AGENTS.md to ui/app/ROUTES.md for progressive discovery (sibling to routes.ts) - Remove bold formatting from documentation per reviewer feedback * Clarify that both RR7 and gateway API routes use underscores * Move URL conventions to ui/app/routes/README.md --- ui/app/components/layout/app.sidebar.tsx | 2 +- ui/app/routes.ts | 8 ++++---- ui/app/routes/README.md | 12 ++++++++++++ ui/app/routes/index.tsx | 2 +- .../WorkflowEvaluationProjectsTable.tsx | 0 .../WorkflowEvaluationRunsTable.tsx | 0 .../layout.tsx | 0 .../$project_name/WorkflowEvalRunSelector.tsx | 0 .../WorkflowEvaluationProjectResultsTable.tsx | 0 .../$project_name/WorkflowEvaluationRunBadge.tsx | 0 .../projects/$project_name/route.tsx | 4 ++-- .../route.tsx | 0 .../runs/$run_id/WorkflowEvaluationRunBasicInfo.tsx | 0 .../$run_id/WorkflowEvaluationRunEpisodesTable.tsx | 0 .../runs/$run_id/route.tsx | 2 +- ui/app/utils/urls.test.ts | 8 ++++---- ui/app/utils/urls.ts | 4 ++-- ...rkflow-evaluations.projects.project_name.spec.ts} | 2 +- ...luations.spec.ts => workflow-evaluations.spec.ts} | 4 ++-- 19 files changed, 30 insertions(+), 18 deletions(-) create mode 100644 ui/app/routes/README.md rename ui/app/routes/{workflow_evaluations => workflow-evaluations}/WorkflowEvaluationProjectsTable.tsx (100%) rename ui/app/routes/{workflow_evaluations => workflow-evaluations}/WorkflowEvaluationRunsTable.tsx (100%) rename ui/app/routes/{workflow_evaluations => workflow-evaluations}/layout.tsx (100%) rename ui/app/routes/{workflow_evaluations => workflow-evaluations}/projects/$project_name/WorkflowEvalRunSelector.tsx (100%) rename ui/app/routes/{workflow_evaluations => workflow-evaluations}/projects/$project_name/WorkflowEvaluationProjectResultsTable.tsx (100%) rename ui/app/routes/{workflow_evaluations => workflow-evaluations}/projects/$project_name/WorkflowEvaluationRunBadge.tsx (100%) rename ui/app/routes/{workflow_evaluations => workflow-evaluations}/projects/$project_name/route.tsx (97%) rename ui/app/routes/{workflow_evaluations => workflow-evaluations}/route.tsx (100%) rename ui/app/routes/{workflow_evaluations => workflow-evaluations}/runs/$run_id/WorkflowEvaluationRunBasicInfo.tsx (100%) rename ui/app/routes/{workflow_evaluations => workflow-evaluations}/runs/$run_id/WorkflowEvaluationRunEpisodesTable.tsx (100%) rename ui/app/routes/{workflow_evaluations => workflow-evaluations}/runs/$run_id/route.tsx (98%) rename ui/e2e_tests/{workflow_evaluations.projects.project_name.spec.ts => workflow-evaluations.projects.project_name.spec.ts} (97%) rename ui/e2e_tests/{workflow_evaluations.spec.ts => workflow-evaluations.spec.ts} (94%) diff --git a/ui/app/components/layout/app.sidebar.tsx b/ui/app/components/layout/app.sidebar.tsx index f8a244883a7..cd187b9cbee 100644 --- a/ui/app/components/layout/app.sidebar.tsx +++ b/ui/app/components/layout/app.sidebar.tsx @@ -116,7 +116,7 @@ const navigation: NavigationSection[] = [ }, { title: "Workflow Evaluations", - url: "/workflow_evaluations", + url: "/workflow-evaluations", icon: SequenceChecks, }, ], diff --git a/ui/app/routes.ts b/ui/app/routes.ts index 635f774e42b..964971c06a6 100644 --- a/ui/app/routes.ts +++ b/ui/app/routes.ts @@ -106,12 +106,12 @@ export default [ ]), // Workflow Evaluations (formerly Dynamic Evaluations) - route("workflow_evaluations", "routes/workflow_evaluations/layout.tsx", [ - index("routes/workflow_evaluations/route.tsx"), - route("runs/:run_id", "routes/workflow_evaluations/runs/$run_id/route.tsx"), + route("workflow-evaluations", "routes/workflow-evaluations/layout.tsx", [ + index("routes/workflow-evaluations/route.tsx"), + route("runs/:run_id", "routes/workflow-evaluations/runs/$run_id/route.tsx"), route( "projects/:project_name", - "routes/workflow_evaluations/projects/$project_name/route.tsx", + "routes/workflow-evaluations/projects/$project_name/route.tsx", ), ]), diff --git a/ui/app/routes/README.md b/ui/app/routes/README.md new file mode 100644 index 00000000000..bd1dc0d1ae0 --- /dev/null +++ b/ui/app/routes/README.md @@ -0,0 +1,12 @@ +# UI Routes + +Routes are defined in `routes.ts`. + +## URL Conventions + +- Frontend routes (user-facing URLs): Use hyphens (e.g., `/api-keys`, `/workflow-evaluations`, `/supervised-fine-tuning`) +- Backend API routes: Use underscores + - RR7 API routes (React Router resource routes): `/api/workflow_evaluations/search_runs` + - Gateway API routes: `/internal/workflow_evaluations`, `/api/curated_inferences` + +This convention aligns frontend URLs with web standards (hyphens) while backend routes use underscores for consistency with the TensorZero gateway API. diff --git a/ui/app/routes/index.tsx b/ui/app/routes/index.tsx index f48b8a4a2e8..e0b83a841d4 100644 --- a/ui/app/routes/index.tsx +++ b/ui/app/routes/index.tsx @@ -294,7 +294,7 @@ export default function Home({ loaderData }: Route.ComponentProps) { description={inferenceEvaluationsDesc} /> diff --git a/ui/app/utils/urls.test.ts b/ui/app/utils/urls.test.ts index d9836d4778d..5a7e0d3170e 100644 --- a/ui/app/utils/urls.test.ts +++ b/ui/app/utils/urls.test.ts @@ -110,10 +110,10 @@ describe("URL helper functions", () => { describe("toWorkflowEvaluationRunUrl", () => { it("should encode run IDs", () => { expect(toWorkflowEvaluationRunUrl("run_1")).toBe( - "/workflow_evaluations/runs/run_1", + "/workflow-evaluations/runs/run_1", ); expect(toWorkflowEvaluationRunUrl("run/test")).toBe( - "/workflow_evaluations/runs/run%2Ftest", + "/workflow-evaluations/runs/run%2Ftest", ); }); }); @@ -121,10 +121,10 @@ describe("URL helper functions", () => { describe("toWorkflowEvaluationProjectUrl", () => { it("should encode project names", () => { expect(toWorkflowEvaluationProjectUrl("project_1")).toBe( - "/workflow_evaluations/projects/project_1", + "/workflow-evaluations/projects/project_1", ); expect(toWorkflowEvaluationProjectUrl("project/test")).toBe( - "/workflow_evaluations/projects/project%2Ftest", + "/workflow-evaluations/projects/project%2Ftest", ); }); }); diff --git a/ui/app/utils/urls.ts b/ui/app/utils/urls.ts index 406d087235a..9e4299d2f3b 100644 --- a/ui/app/utils/urls.ts +++ b/ui/app/utils/urls.ts @@ -81,11 +81,11 @@ export function toEvaluationDatapointUrl( // ============================================================================ export function toWorkflowEvaluationRunUrl(runId: string): string { - return `/workflow_evaluations/runs/${encodeURIComponent(runId)}`; + return `/workflow-evaluations/runs/${encodeURIComponent(runId)}`; } export function toWorkflowEvaluationProjectUrl(projectName: string): string { - return `/workflow_evaluations/projects/${encodeURIComponent(projectName)}`; + return `/workflow-evaluations/projects/${encodeURIComponent(projectName)}`; } // ============================================================================ diff --git a/ui/e2e_tests/workflow_evaluations.projects.project_name.spec.ts b/ui/e2e_tests/workflow-evaluations.projects.project_name.spec.ts similarity index 97% rename from ui/e2e_tests/workflow_evaluations.projects.project_name.spec.ts rename to ui/e2e_tests/workflow-evaluations.projects.project_name.spec.ts index a542bf7f09c..8464907eaa3 100644 --- a/ui/e2e_tests/workflow_evaluations.projects.project_name.spec.ts +++ b/ui/e2e_tests/workflow-evaluations.projects.project_name.spec.ts @@ -3,7 +3,7 @@ import { test, expect } from "@playwright/test"; test("workflow evaluation project page should render and show correct information", async ({ page, }) => { - await page.goto("/workflow_evaluations/projects/beerqa-agentic-rag"); + await page.goto("/workflow-evaluations/projects/beerqa-agentic-rag"); await expect( page .getByRole("navigation", { name: "breadcrumb" }) diff --git a/ui/e2e_tests/workflow_evaluations.spec.ts b/ui/e2e_tests/workflow-evaluations.spec.ts similarity index 94% rename from ui/e2e_tests/workflow_evaluations.spec.ts rename to ui/e2e_tests/workflow-evaluations.spec.ts index 31a3b54f2b1..9e3317e5d90 100644 --- a/ui/e2e_tests/workflow_evaluations.spec.ts +++ b/ui/e2e_tests/workflow-evaluations.spec.ts @@ -3,7 +3,7 @@ import { test, expect } from "@playwright/test"; test("should show the workflow evaluations page and navigate to the workflow evaluation run page", async ({ page, }) => { - await page.goto("/workflow_evaluations"); + await page.goto("/workflow-evaluations"); await expect(page.getByText("Evaluation Runs")).toBeVisible(); await expect( page.getByText("01968d04-142c-7e53-8ea7-3a3255b518dc"), @@ -22,7 +22,7 @@ test("should render comment in workflow evaluation run and open modal when click page, }) => { await page.goto( - "/workflow_evaluations/runs/01968d05-d734-7751-ab33-75dd8b3fb4a3", + "/workflow-evaluations/runs/01968d05-d734-7751-ab33-75dd8b3fb4a3", ); // Wait for the page to load From 3f9e15b118ef3e104ad4a119ef80dadebde7148d Mon Sep 17 00:00:00 2001 From: simeonlee <15081451+simeonlee@users.noreply.github.com> Date: Mon, 26 Jan 2026 23:10:24 -0500 Subject: [PATCH 16/46] Add React Router streaming to api-keys page (#5806) * Add React Router streaming to api-keys page Implements the Suspense/Await pattern for non-blocking page loads: - ApiKeysPageHeader - shared header component - ApiKeysContentSkeleton - skeleton fallback for loading state - ApiKeysErrorState - error state using useAsyncError() - ApiKeysContent - main content component - ApiKeysPage - orchestrates Suspense/Await pattern Naming follows the pattern established in #5651. * Move modal state outside Suspense boundary and add location.key --- ui/app/routes/api-keys/route.tsx | 204 ++++++++++++++++++++++--------- 1 file changed, 146 insertions(+), 58 deletions(-) diff --git a/ui/app/routes/api-keys/route.tsx b/ui/app/routes/api-keys/route.tsx index 6fa0746662c..5324ed74fe2 100644 --- a/ui/app/routes/api-keys/route.tsx +++ b/ui/app/routes/api-keys/route.tsx @@ -1,7 +1,10 @@ import type { Route } from "./+types/route"; +import { Suspense, useState } from "react"; import { + Await, data, - isRouteErrorResponse, + useAsyncError, + useLocation, useNavigate, type RouteHandle, } from "react-router"; @@ -11,8 +14,8 @@ import { PageLayout, SectionLayout, } from "~/components/layout/PageLayout"; -import { useState } from "react"; import { logger } from "~/utils/logger"; +import { PageErrorContent } from "~/components/ui/error"; import { getPostgresClient, isPostgresAvailable, @@ -21,11 +24,102 @@ import AuthTable from "./AuthTable"; import { AuthActions } from "./AuthActions"; import { GenerateApiKeyModal } from "./GenerateApiKeyModal"; import { PostgresRequiredState } from "~/components/ui/PostgresRequiredState"; +import { Skeleton } from "~/components/ui/skeleton"; +import { + Table, + TableBody, + TableCell, + TableHead, + TableHeader, + TableRow, +} from "~/components/ui/table"; +import type { KeyInfo } from "~/types/tensorzero"; export const handle: RouteHandle = { crumb: () => ["TensorZero API Keys"], }; +export type ApiKeysData = { + apiKeys: KeyInfo[]; + offset: number; + limit: number; +}; + +function ApiKeysPageHeader() { + return ; +} + +function ApiKeysContentSkeleton() { + return ( + <> + + +
+ +
+ + + + + Public ID + Description + Created + + + + + {[1, 2, 3].map((i) => ( + + + + + + + + + + + +
+ + +
+ + + ))} + +
+ +
+ + +
+ + + ); +} + +function ApiKeysErrorState() { + const error = useAsyncError(); + return ( + <> + + + + + + ); +} + +async function fetchApiKeys( + limit: number, + offset: number, +): Promise { + const postgresClient = await getPostgresClient(); + const apiKeys = await postgresClient.listApiKeys(limit, offset); + return { apiKeys, offset, limit }; +} + export async function loader({ request }: Route.LoaderArgs) { const url = new URL(request.url); const searchParams = new URLSearchParams(url.search); @@ -38,21 +132,14 @@ export async function loader({ request }: Route.LoaderArgs) { if (!isPostgresAvailable()) { return { - postgresAvailable: false, - apiKeys: [], - offset: 0, - limit: 0, + postgresAvailable: false as const, + apiKeysData: null, }; } - const postgresClient = await getPostgresClient(); - const apiKeys = await postgresClient.listApiKeys(limit, offset); - return { - postgresAvailable: true, - apiKeys, - offset, - limit, + postgresAvailable: true as const, + apiKeysData: fetchApiKeys(limit, offset), }; } @@ -142,15 +229,15 @@ export async function action({ request }: Route.ActionArgs) { }; } -export default function AuthPage({ loaderData }: Route.ComponentProps) { +function ApiKeysContent({ + data, + onOpenModal, +}: { + data: ApiKeysData; + onOpenModal: () => void; +}) { + const { apiKeys, offset, limit } = data; const navigate = useNavigate(); - const { postgresAvailable, apiKeys, offset, limit } = loaderData; - const [generateModalIsOpen, setGenerateModalIsOpen] = useState(false); - const [modalKey, setModalKey] = useState(0); - - if (!postgresAvailable) { - return ; - } const handleNextPage = () => { const searchParams = new URLSearchParams(window.location.search); @@ -164,6 +251,34 @@ export default function AuthPage({ loaderData }: Route.ComponentProps) { navigate(`?${searchParams.toString()}`, { preventScrollReset: true }); }; + return ( + <> + + + + + + + + ); +} + +export default function ApiKeysPage({ loaderData }: Route.ComponentProps) { + const { postgresAvailable, apiKeysData } = loaderData; + const location = useLocation(); + const navigate = useNavigate(); + const [generateModalIsOpen, setGenerateModalIsOpen] = useState(false); + const [modalKey, setModalKey] = useState(0); + + if (!postgresAvailable) { + return ; + } + const handleOpenModal = () => { setModalKey((prev) => prev + 1); setGenerateModalIsOpen(true); @@ -172,27 +287,23 @@ export default function AuthPage({ loaderData }: Route.ComponentProps) { const handleCloseModal = () => { setGenerateModalIsOpen(false); // Reset to first page after creating API key + const searchParams = new URLSearchParams(window.location.search); + const offset = parseInt(searchParams.get("offset") || "0"); if (offset > 0) { - const searchParams = new URLSearchParams(window.location.search); searchParams.set("offset", "0"); - searchParams.set("limit", String(limit)); navigate(`?${searchParams.toString()}`, { preventScrollReset: true }); } }; return ( - - - - - - + }> + }> + {(resolvedData) => ( + + )} + + -

- {error.status} {error.statusText} -

-

{error.data}

- - ); - } else if (error instanceof Error) { - return ( -
-

Error

-

{error.message}

-
- ); - } else { - return ( -
-

Unknown Error

-
- ); - } + return ; } From b5eceff8ea68ba28ed7d90f6c6d8978201cab50d Mon Sep 17 00:00:00 2001 From: Viraj Mehta Date: Mon, 26 Jan 2026 23:10:53 -0500 Subject: [PATCH 17/46] bumped napi to v3, node.js to 24 (#5848) * bumped napi to v3 * also bumped node version * reverted changes to docker compose files --- Cargo.lock | 69 +- internal/tensorzero-node/Cargo.toml | 6 +- internal/tensorzero-node/package.json | 25 +- pnpm-lock.yaml | 1105 ++++++++++++++++++++++++- ui/Dockerfile | 5 +- 5 files changed, 1165 insertions(+), 45 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index b271494aebe..ac3edef1c5f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1357,9 +1357,9 @@ checksum = "3d52eff69cd5e647efe296129160853a42795992097e8af39800e1060caeea9b" [[package]] name = "convert_case" -version = "0.6.0" +version = "0.10.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ec182b0ca2f35d8fc196cf3404988fd8b8c739a4d270ff118a398feb0cbec1ca" +checksum = "633458d4ef8c78b72454de2d54fd6ab2e60f9e02be22f3c6104cdc8a4e0fceb9" dependencies = [ "unicode-segmentation", ] @@ -1518,14 +1518,20 @@ dependencies = [ [[package]] name = "ctor" -version = "0.2.9" +version = "0.6.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "32a2785755761f3ddc1492979ce1e48d2c00d09311c39e4466429188f3dd6501" +checksum = "424e0138278faeb2b401f174ad17e715c829512d74f3d1e81eb43365c2e0590e" dependencies = [ - "quote", - "syn 2.0.114", + "ctor-proc-macro", + "dtor", ] +[[package]] +name = "ctor-proc-macro" +version = "0.0.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "52560adf09603e58c9a7ee1fe1dcb95a16927b17c127f0ac02d6e768a0e25bc1" + [[package]] name = "darling" version = "0.14.4" @@ -1796,6 +1802,21 @@ version = "0.11.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1435fa1053d8b2fbbe9be7e97eca7f33d37b28409959813daefc1446a14247f1" +[[package]] +name = "dtor" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "404d02eeb088a82cfd873006cb713fe411306c7d182c344905e101fb1167d301" +dependencies = [ + "dtor-proc-macro", +] + +[[package]] +name = "dtor-proc-macro" +version = "0.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f678cf4a922c215c63e0de95eb1ff08a958a81d47e485cf9da1e27bf6305cfa5" + [[package]] name = "dunce" version = "1.0.5" @@ -3302,9 +3323,9 @@ dependencies = [ [[package]] name = "libloading" -version = "0.8.9" +version = "0.9.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d7c4b02199fee7c5d21a5ae7d8cfa79a6ef5bb2fc834d6e9058e89c825efdc55" +checksum = "754ca22de805bb5744484a5b151a9e1a8e837d5dc232c2d7d8c2e3492edc8b60" dependencies = [ "cfg-if", "windows-link", @@ -3711,15 +3732,17 @@ dependencies = [ [[package]] name = "napi" -version = "2.16.17" +version = "3.8.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "55740c4ae1d8696773c78fdafd5d0e5fe9bc9f1b071c7ba493ba5c413a9184f3" +checksum = "909805cbad4d569e69b80e101290fe72e92b9742ba9e333b0c1e83b22fb7447b" dependencies = [ "bitflags", "ctor", - "napi-derive", + "futures", + "napi-build", "napi-sys", - "once_cell", + "nohash-hasher", + "rustc-hash", "tokio", ] @@ -3731,12 +3754,12 @@ checksum = "d376940fd5b723c6893cd1ee3f33abbfd86acb1cd1ec079f3ab04a2a3bc4d3b1" [[package]] name = "napi-derive" -version = "2.16.13" +version = "3.5.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7cbe2585d8ac223f7d34f13701434b9d5f4eb9c332cccce8dee57ea18ab8ab0c" +checksum = "04ba21bbdf40b33496b4ee6eadfc64d17a6a6cde57cd31549117b0882d1fef86" dependencies = [ - "cfg-if", "convert_case", + "ctor", "napi-derive-backend", "proc-macro2", "quote", @@ -3745,24 +3768,22 @@ dependencies = [ [[package]] name = "napi-derive-backend" -version = "1.0.75" +version = "5.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1639aaa9eeb76e91c6ae66da8ce3e89e921cd3885e99ec85f4abacae72fc91bf" +checksum = "e9a63791e230572c3218a7acd86ca0a0529fc64294bcbea567cf906d7b04e077" dependencies = [ "convert_case", - "once_cell", "proc-macro2", "quote", - "regex", "semver", "syn 2.0.114", ] [[package]] name = "napi-sys" -version = "2.4.0" +version = "3.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "427802e8ec3a734331fec1035594a210ce1ff4dc5bc1950530920ab717964ea3" +checksum = "8eb602b84d7c1edae45e50bbf1374696548f36ae179dfa667f577e384bb90c2b" dependencies = [ "libloading", ] @@ -3779,6 +3800,12 @@ version = "0.4.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b93853da6d84c2e3c7d730d6473e8817692dd89be387eb01b94d7f108ecb5b8c" +[[package]] +name = "nohash-hasher" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2bf50223579dc7cdcfb3bfcacf7069ff68243f8c363f62ffa99cf000a6b9c451" + [[package]] name = "nom" version = "7.1.3" diff --git a/internal/tensorzero-node/Cargo.toml b/internal/tensorzero-node/Cargo.toml index 50127c26a4b..e34dd46f3ae 100644 --- a/internal/tensorzero-node/Cargo.toml +++ b/internal/tensorzero-node/Cargo.toml @@ -10,15 +10,15 @@ crate-type = ["cdylib"] [dependencies] # Default enable napi4 feature, see https://nodejs.org/api/n-api.html#node-api-version-matrix -napi = { version = "2.12.2", default-features = false, features = [ +napi = { version = "3.8.2", default-features = false, features = [ "napi4", "tokio_rt", ] } -napi-derive = "2.12.2" +napi-derive = "3.5.1" serde_json.workspace = true secrecy.workspace = true tensorzero-core = { path = "../../tensorzero-core" } tensorzero-auth = { path = "../../internal/tensorzero-auth" } [build-dependencies] -napi-build = "2.3.0" +napi-build = "2.3.1" diff --git a/internal/tensorzero-node/package.json b/internal/tensorzero-node/package.json index ae35c2be3a5..1fd4be9ef57 100644 --- a/internal/tensorzero-node/package.json +++ b/internal/tensorzero-node/package.json @@ -11,23 +11,20 @@ "main": "dist/index.js", "types": "dist/index.d.ts", "napi": { - "name": "tensorzero-node", - "triples": { - "defaults": false, - "additional": [ - "aarch64-apple-darwin", - "aarch64-unknown-linux-gnu", - "aarch64-unknown-linux-musl", - "x86_64-apple-darwin", - "x86_64-unknown-linux-gnu", - "x86_64-unknown-linux-musl" - ] - } + "binaryName": "tensorzero-node", + "targets": [ + "aarch64-apple-darwin", + "aarch64-unknown-linux-gnu", + "aarch64-unknown-linux-musl", + "x86_64-apple-darwin", + "x86_64-unknown-linux-gnu", + "x86_64-unknown-linux-musl" + ] }, "license": "Apache-2.0", "devDependencies": { "@eslint/js": "^9.37.0", - "@napi-rs/cli": "^2.18.4", + "@napi-rs/cli": "^3.0.0-alpha.72", "@types/node": "^20.19.4", "eslint": "^9.37.0", "prettier": "^3.6.2", @@ -54,7 +51,7 @@ "format:check": "prettier --check \"**/*.{js,jsx,ts,tsx,css,scss,html,json,yaml,md}\" --ignore-path .prettierignore", "lint": "eslint . --fix --max-warnings=0 --config eslint.config.js --cache --ignore-pattern \"dist/\" --ignore-pattern \"index.d.ts\"", "lint:check": "eslint . --max-warnings=0 --config eslint.config.js --cache --ignore-pattern \"dist/\" --ignore-pattern \"index.d.ts\"", - "universal": "napi universal", + "universal": "napi universalize", "version": "napi version" }, "dependencies": { diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index a80caf42d27..7a0be7cf5a1 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -140,8 +140,8 @@ importers: specifier: ^9.37.0 version: 9.37.0 '@napi-rs/cli': - specifier: ^2.18.4 - version: 2.18.4 + specifier: ^3.0.0-alpha.72 + version: 3.5.1(@emnapi/runtime@1.8.1)(@types/node@20.19.4) '@types/node': specifier: ^20.19.4 version: 20.19.4 @@ -631,6 +631,15 @@ packages: '@codemirror/view@6.38.6': resolution: {integrity: sha512-qiS0z1bKs5WOvHIAC0Cybmv4AJSkAXgX5aD6Mqd2epSLlVJsQl8NG23jCVouIgkh4All/mrbdsf2UOLFnJw0tw==} + '@emnapi/core@1.8.1': + resolution: {integrity: sha512-AvT9QFpxK0Zd8J0jopedNm+w/2fIzvtPKPjqyw9jwvBaReTTqPBk9Hixaz7KbjimP+QNz605/XnjFcDAL2pqBg==} + + '@emnapi/runtime@1.8.1': + resolution: {integrity: sha512-mehfKSMWjjNol8659Z8KxEMrdSJDDot5SXMq00dM8BN4o+CLNXQ0xH2V7EchNHV4RmbZLmmPdEaXZc5H2FXmDg==} + + '@emnapi/wasi-threads@1.1.0': + resolution: {integrity: sha512-WI0DdZ8xFSbgMjR1sFsKABJ/C5OnRrjT06JXbZKexJGrDuPTzZdDYfFlsgcCXCyf+suG5QU2e/y1Wo2V/OapLQ==} + '@esbuild/aix-ppc64@0.25.10': resolution: {integrity: sha512-0NFWnA+7l41irNuaSVlLfgNT12caWJVLzp5eAVhZ0z1qpxbockccEt3s+149rE64VUI3Ml2zt8Nv5JVc4QXTsw==} engines: {node: '>=18'} @@ -1015,6 +1024,140 @@ packages: resolution: {integrity: sha512-bV0Tgo9K4hfPCek+aMAn81RppFKv2ySDQeMoSZuvTASywNTnVJCArCZE2FWqpvIatKu7VMRLWlR1EazvVhDyhQ==} engines: {node: '>=18.18'} + '@inquirer/ansi@2.0.3': + resolution: {integrity: sha512-g44zhR3NIKVs0zUesa4iMzExmZpLUdTLRMCStqX3GE5NT6VkPcxQGJ+uC8tDgBUC/vB1rUhUd55cOf++4NZcmw==} + engines: {node: '>=23.5.0 || ^22.13.0 || ^21.7.0 || ^20.12.0'} + + '@inquirer/checkbox@5.0.4': + resolution: {integrity: sha512-DrAMU3YBGMUAp6ArwTIp/25CNDtDbxk7UjIrrtM25JVVrlVYlVzHh5HR1BDFu9JMyUoZ4ZanzeaHqNDttf3gVg==} + engines: {node: '>=23.5.0 || ^22.13.0 || ^21.7.0 || ^20.12.0'} + peerDependencies: + '@types/node': '>=18' + peerDependenciesMeta: + '@types/node': + optional: true + + '@inquirer/confirm@6.0.4': + resolution: {integrity: sha512-WdaPe7foUnoGYvXzH4jp4wH/3l+dBhZ3uwhKjXjwdrq5tEIFaANxj6zrGHxLdsIA0yKM0kFPVcEalOZXBB5ISA==} + engines: {node: '>=23.5.0 || ^22.13.0 || ^21.7.0 || ^20.12.0'} + peerDependencies: + '@types/node': '>=18' + peerDependenciesMeta: + '@types/node': + optional: true + + '@inquirer/core@11.1.1': + resolution: {integrity: sha512-hV9o15UxX46OyQAtaoMqAOxGR8RVl1aZtDx1jHbCtSJy1tBdTfKxLPKf7utsE4cRy4tcmCQ4+vdV+ca+oNxqNA==} + engines: {node: '>=23.5.0 || ^22.13.0 || ^21.7.0 || ^20.12.0'} + peerDependencies: + '@types/node': '>=18' + peerDependenciesMeta: + '@types/node': + optional: true + + '@inquirer/editor@5.0.4': + resolution: {integrity: sha512-QI3Jfqcv6UO2/VJaEFONH8Im1ll++Xn/AJTBn9Xf+qx2M+H8KZAdQ5sAe2vtYlo+mLW+d7JaMJB4qWtK4BG3pw==} + engines: {node: '>=23.5.0 || ^22.13.0 || ^21.7.0 || ^20.12.0'} + peerDependencies: + '@types/node': '>=18' + peerDependenciesMeta: + '@types/node': + optional: true + + '@inquirer/expand@5.0.4': + resolution: {integrity: sha512-0I/16YwPPP0Co7a5MsomlZLpch48NzYfToyqYAOWtBmaXSB80RiNQ1J+0xx2eG+Wfxt0nHtpEWSRr6CzNVnOGg==} + engines: {node: '>=23.5.0 || ^22.13.0 || ^21.7.0 || ^20.12.0'} + peerDependencies: + '@types/node': '>=18' + peerDependenciesMeta: + '@types/node': + optional: true + + '@inquirer/external-editor@2.0.3': + resolution: {integrity: sha512-LgyI7Agbda74/cL5MvA88iDpvdXI2KuMBCGRkbCl2Dg1vzHeOgs+s0SDcXV7b+WZJrv2+ERpWSM65Fpi9VfY3w==} + engines: {node: '>=23.5.0 || ^22.13.0 || ^21.7.0 || ^20.12.0'} + peerDependencies: + '@types/node': '>=18' + peerDependenciesMeta: + '@types/node': + optional: true + + '@inquirer/figures@2.0.3': + resolution: {integrity: sha512-y09iGt3JKoOCBQ3w4YrSJdokcD8ciSlMIWsD+auPu+OZpfxLuyz+gICAQ6GCBOmJJt4KEQGHuZSVff2jiNOy7g==} + engines: {node: '>=23.5.0 || ^22.13.0 || ^21.7.0 || ^20.12.0'} + + '@inquirer/input@5.0.4': + resolution: {integrity: sha512-4B3s3jvTREDFvXWit92Yc6jF1RJMDy2VpSqKtm4We2oVU65YOh2szY5/G14h4fHlyQdpUmazU5MPCFZPRJ0AOw==} + engines: {node: '>=23.5.0 || ^22.13.0 || ^21.7.0 || ^20.12.0'} + peerDependencies: + '@types/node': '>=18' + peerDependenciesMeta: + '@types/node': + optional: true + + '@inquirer/number@4.0.4': + resolution: {integrity: sha512-CmMp9LF5HwE+G/xWsC333TlCzYYbXMkcADkKzcawh49fg2a1ryLc7JL1NJYYt1lJ+8f4slikNjJM9TEL/AljYQ==} + engines: {node: '>=23.5.0 || ^22.13.0 || ^21.7.0 || ^20.12.0'} + peerDependencies: + '@types/node': '>=18' + peerDependenciesMeta: + '@types/node': + optional: true + + '@inquirer/password@5.0.4': + resolution: {integrity: sha512-ZCEPyVYvHK4W4p2Gy6sTp9nqsdHQCfiPXIP9LbJVW4yCinnxL/dDDmPaEZVysGrj8vxVReRnpfS2fOeODe9zjg==} + engines: {node: '>=23.5.0 || ^22.13.0 || ^21.7.0 || ^20.12.0'} + peerDependencies: + '@types/node': '>=18' + peerDependenciesMeta: + '@types/node': + optional: true + + '@inquirer/prompts@8.2.0': + resolution: {integrity: sha512-rqTzOprAj55a27jctS3vhvDDJzYXsr33WXTjODgVOru21NvBo9yIgLIAf7SBdSV0WERVly3dR6TWyp7ZHkvKFA==} + engines: {node: '>=23.5.0 || ^22.13.0 || ^21.7.0 || ^20.12.0'} + peerDependencies: + '@types/node': '>=18' + peerDependenciesMeta: + '@types/node': + optional: true + + '@inquirer/rawlist@5.2.0': + resolution: {integrity: sha512-CciqGoOUMrFo6HxvOtU5uL8fkjCmzyeB6fG7O1vdVAZVSopUBYECOwevDBlqNLyyYmzpm2Gsn/7nLrpruy9RFg==} + engines: {node: '>=23.5.0 || ^22.13.0 || ^21.7.0 || ^20.12.0'} + peerDependencies: + '@types/node': '>=18' + peerDependenciesMeta: + '@types/node': + optional: true + + '@inquirer/search@4.1.0': + resolution: {integrity: sha512-EAzemfiP4IFvIuWnrHpgZs9lAhWDA0GM3l9F4t4mTQ22IFtzfrk8xbkMLcAN7gmVML9O/i+Hzu8yOUyAaL6BKA==} + engines: {node: '>=23.5.0 || ^22.13.0 || ^21.7.0 || ^20.12.0'} + peerDependencies: + '@types/node': '>=18' + peerDependenciesMeta: + '@types/node': + optional: true + + '@inquirer/select@5.0.4': + resolution: {integrity: sha512-s8KoGpPYMEQ6WXc0dT9blX2NtIulMdLOO3LA1UKOiv7KFWzlJ6eLkEYTDBIi+JkyKXyn8t/CD6TinxGjyLt57g==} + engines: {node: '>=23.5.0 || ^22.13.0 || ^21.7.0 || ^20.12.0'} + peerDependencies: + '@types/node': '>=18' + peerDependenciesMeta: + '@types/node': + optional: true + + '@inquirer/type@4.0.3': + resolution: {integrity: sha512-cKZN7qcXOpj1h+1eTTcGDVLaBIHNMT1Rz9JqJP5MnEJ0JhgVWllx7H/tahUp5YEK1qaByH2Itb8wLG/iScD5kw==} + engines: {node: '>=23.5.0 || ^22.13.0 || ^21.7.0 || ^20.12.0'} + peerDependencies: + '@types/node': '>=18' + peerDependenciesMeta: + '@types/node': + optional: true + '@isaacs/balanced-match@4.0.1': resolution: {integrity: sha512-yzMTt9lEb8Gv7zRioUilSglI0c0smZ9k5D65677DLWLtWJaXIS3CqcGyUFByYKlnUj6TkjLVs54fBl6+TiGQDQ==} engines: {node: 20 || >=22} @@ -1093,10 +1236,338 @@ packages: '@mjackson/node-fetch-server@0.2.0': resolution: {integrity: sha512-EMlH1e30yzmTpGLQjlFmaDAjyOeZhng1/XCd7DExR8PNAnG/G1tyruZxEoUe11ClnwGhGrtsdnyyUx1frSzjng==} - '@napi-rs/cli@2.18.4': - resolution: {integrity: sha512-SgJeA4df9DE2iAEpr3M2H0OKl/yjtg1BnRI5/JyowS71tUWhrfSu2LT0V3vlHET+g1hBVlrO60PmEXwUEKp8Mg==} - engines: {node: '>= 10'} + '@napi-rs/cli@3.5.1': + resolution: {integrity: sha512-XBfLQRDcB3qhu6bazdMJsecWW55kR85l5/k0af9BIBELXQSsCFU0fzug7PX8eQp6vVdm7W/U3z6uP5WmITB2Gw==} + engines: {node: '>= 16'} hasBin: true + peerDependencies: + '@emnapi/runtime': ^1.7.1 + peerDependenciesMeta: + '@emnapi/runtime': + optional: true + + '@napi-rs/cross-toolchain@1.0.3': + resolution: {integrity: sha512-ENPfLe4937bsKVTDA6zdABx4pq9w0tHqRrJHyaGxgaPq03a2Bd1unD5XSKjXJjebsABJ+MjAv1A2OvCgK9yehg==} + peerDependencies: + '@napi-rs/cross-toolchain-arm64-target-aarch64': ^1.0.3 + '@napi-rs/cross-toolchain-arm64-target-armv7': ^1.0.3 + '@napi-rs/cross-toolchain-arm64-target-ppc64le': ^1.0.3 + '@napi-rs/cross-toolchain-arm64-target-s390x': ^1.0.3 + '@napi-rs/cross-toolchain-arm64-target-x86_64': ^1.0.3 + '@napi-rs/cross-toolchain-x64-target-aarch64': ^1.0.3 + '@napi-rs/cross-toolchain-x64-target-armv7': ^1.0.3 + '@napi-rs/cross-toolchain-x64-target-ppc64le': ^1.0.3 + '@napi-rs/cross-toolchain-x64-target-s390x': ^1.0.3 + '@napi-rs/cross-toolchain-x64-target-x86_64': ^1.0.3 + peerDependenciesMeta: + '@napi-rs/cross-toolchain-arm64-target-aarch64': + optional: true + '@napi-rs/cross-toolchain-arm64-target-armv7': + optional: true + '@napi-rs/cross-toolchain-arm64-target-ppc64le': + optional: true + '@napi-rs/cross-toolchain-arm64-target-s390x': + optional: true + '@napi-rs/cross-toolchain-arm64-target-x86_64': + optional: true + '@napi-rs/cross-toolchain-x64-target-aarch64': + optional: true + '@napi-rs/cross-toolchain-x64-target-armv7': + optional: true + '@napi-rs/cross-toolchain-x64-target-ppc64le': + optional: true + '@napi-rs/cross-toolchain-x64-target-s390x': + optional: true + '@napi-rs/cross-toolchain-x64-target-x86_64': + optional: true + + '@napi-rs/lzma-android-arm-eabi@1.4.5': + resolution: {integrity: sha512-Up4gpyw2SacmyKWWEib06GhiDdF+H+CCU0LAV8pnM4aJIDqKKd5LHSlBht83Jut6frkB0vwEPmAkv4NjQ5u//Q==} + engines: {node: '>= 10'} + cpu: [arm] + os: [android] + + '@napi-rs/lzma-android-arm64@1.4.5': + resolution: {integrity: sha512-uwa8sLlWEzkAM0MWyoZJg0JTD3BkPknvejAFG2acUA1raXM8jLrqujWCdOStisXhqQjZ2nDMp3FV6cs//zjfuQ==} + engines: {node: '>= 10'} + cpu: [arm64] + os: [android] + + '@napi-rs/lzma-darwin-arm64@1.4.5': + resolution: {integrity: sha512-0Y0TQLQ2xAjVabrMDem1NhIssOZzF/y/dqetc6OT8mD3xMTDtF8u5BqZoX3MyPc9FzpsZw4ksol+w7DsxHrpMA==} + engines: {node: '>= 10'} + cpu: [arm64] + os: [darwin] + + '@napi-rs/lzma-darwin-x64@1.4.5': + resolution: {integrity: sha512-vR2IUyJY3En+V1wJkwmbGWcYiT8pHloTAWdW4pG24+51GIq+intst6Uf6D/r46citObGZrlX0QvMarOkQeHWpw==} + engines: {node: '>= 10'} + cpu: [x64] + os: [darwin] + + '@napi-rs/lzma-freebsd-x64@1.4.5': + resolution: {integrity: sha512-XpnYQC5SVovO35tF0xGkbHYjsS6kqyNCjuaLQ2dbEblFRr5cAZVvsJ/9h7zj/5FluJPJRDojVNxGyRhTp4z2lw==} + engines: {node: '>= 10'} + cpu: [x64] + os: [freebsd] + + '@napi-rs/lzma-linux-arm-gnueabihf@1.4.5': + resolution: {integrity: sha512-ic1ZZMoRfRMwtSwxkyw4zIlbDZGC6davC9r+2oX6x9QiF247BRqqT94qGeL5ZP4Vtz0Hyy7TEViWhx5j6Bpzvw==} + engines: {node: '>= 10'} + cpu: [arm] + os: [linux] + + '@napi-rs/lzma-linux-arm64-gnu@1.4.5': + resolution: {integrity: sha512-asEp7FPd7C1Yi6DQb45a3KPHKOFBSfGuJWXcAd4/bL2Fjetb2n/KK2z14yfW8YC/Fv6x3rBM0VAZKmJuz4tysg==} + engines: {node: '>= 10'} + cpu: [arm64] + os: [linux] + + '@napi-rs/lzma-linux-arm64-musl@1.4.5': + resolution: {integrity: sha512-yWjcPDgJ2nIL3KNvi4536dlT/CcCWO0DUyEOlBs/SacG7BeD6IjGh6yYzd3/X1Y3JItCbZoDoLUH8iB1lTXo3w==} + engines: {node: '>= 10'} + cpu: [arm64] + os: [linux] + + '@napi-rs/lzma-linux-ppc64-gnu@1.4.5': + resolution: {integrity: sha512-0XRhKuIU/9ZjT4WDIG/qnX7Xz7mSQHYZo9Gb3MP2gcvBgr6BA4zywQ9k3gmQaPn9ECE+CZg2V7DV7kT+x2pUMQ==} + engines: {node: '>= 10'} + cpu: [ppc64] + os: [linux] + + '@napi-rs/lzma-linux-riscv64-gnu@1.4.5': + resolution: {integrity: sha512-QrqDIPEUUB23GCpyQj/QFyMlr8SGxxyExeZz9OWFnHfb70kXdTLWrHS/hEI1Ru+lSbQ/6xRqeoGyQ4Aqdg+/RA==} + engines: {node: '>= 10'} + cpu: [riscv64] + os: [linux] + + '@napi-rs/lzma-linux-s390x-gnu@1.4.5': + resolution: {integrity: sha512-k8RVM5aMhW86E9H0QXdquwojew4H3SwPxbRVbl49/COJQWCUjGi79X6mYruMnMPEznZinUiT1jgKbFo2A00NdA==} + engines: {node: '>= 10'} + cpu: [s390x] + os: [linux] + + '@napi-rs/lzma-linux-x64-gnu@1.4.5': + resolution: {integrity: sha512-6rMtBgnIq2Wcl1rQdZsnM+rtCcVCbws1nF8S2NzaUsVaZv8bjrPiAa0lwg4Eqnn1d9lgwqT+cZgm5m+//K08Kw==} + engines: {node: '>= 10'} + cpu: [x64] + os: [linux] + + '@napi-rs/lzma-linux-x64-musl@1.4.5': + resolution: {integrity: sha512-eiadGBKi7Vd0bCArBUOO/qqRYPHt/VQVvGyYvDFt6C2ZSIjlD+HuOl+2oS1sjf4CFjK4eDIog6EdXnL0NE6iyQ==} + engines: {node: '>= 10'} + cpu: [x64] + os: [linux] + + '@napi-rs/lzma-wasm32-wasi@1.4.5': + resolution: {integrity: sha512-+VyHHlr68dvey6fXc2hehw9gHVFIW3TtGF1XkcbAu65qVXsA9D/T+uuoRVqhE+JCyFHFrO0ixRbZDRK1XJt1sA==} + engines: {node: '>=14.0.0'} + cpu: [wasm32] + + '@napi-rs/lzma-win32-arm64-msvc@1.4.5': + resolution: {integrity: sha512-eewnqvIyyhHi3KaZtBOJXohLvwwN27gfS2G/YDWdfHlbz1jrmfeHAmzMsP5qv8vGB+T80TMHNkro4kYjeh6Deg==} + engines: {node: '>= 10'} + cpu: [arm64] + os: [win32] + + '@napi-rs/lzma-win32-ia32-msvc@1.4.5': + resolution: {integrity: sha512-OeacFVRCJOKNU/a0ephUfYZ2Yt+NvaHze/4TgOwJ0J0P4P7X1mHzN+ig9Iyd74aQDXYqc7kaCXA2dpAOcH87Cg==} + engines: {node: '>= 10'} + cpu: [ia32] + os: [win32] + + '@napi-rs/lzma-win32-x64-msvc@1.4.5': + resolution: {integrity: sha512-T4I1SamdSmtyZgDXGAGP+y5LEK5vxHUFwe8mz6D4R7Sa5/WCxTcCIgPJ9BD7RkpO17lzhlaM2vmVvMy96Lvk9Q==} + engines: {node: '>= 10'} + cpu: [x64] + os: [win32] + + '@napi-rs/lzma@1.4.5': + resolution: {integrity: sha512-zS5LuN1OBPAyZpda2ZZgYOEDC+xecUdAGnrvbYzjnLXkrq/OBC3B9qcRvlxbDR3k5H/gVfvef1/jyUqPknqjbg==} + engines: {node: '>= 10'} + + '@napi-rs/tar-android-arm-eabi@1.1.0': + resolution: {integrity: sha512-h2Ryndraj/YiKgMV/r5by1cDusluYIRT0CaE0/PekQ4u+Wpy2iUVqvzVU98ZPnhXaNeYxEvVJHNGafpOfaD0TA==} + engines: {node: '>= 10'} + cpu: [arm] + os: [android] + + '@napi-rs/tar-android-arm64@1.1.0': + resolution: {integrity: sha512-DJFyQHr1ZxNZorm/gzc1qBNLF/FcKzcH0V0Vwan5P+o0aE2keQIGEjJ09FudkF9v6uOuJjHCVDdK6S6uHtShAw==} + engines: {node: '>= 10'} + cpu: [arm64] + os: [android] + + '@napi-rs/tar-darwin-arm64@1.1.0': + resolution: {integrity: sha512-Zz2sXRzjIX4e532zD6xm2SjXEym6MkvfCvL2RMpG2+UwNVDVscHNcz3d47Pf3sysP2e2af7fBB3TIoK2f6trPw==} + engines: {node: '>= 10'} + cpu: [arm64] + os: [darwin] + + '@napi-rs/tar-darwin-x64@1.1.0': + resolution: {integrity: sha512-EI+CptIMNweT0ms9S3mkP/q+J6FNZ1Q6pvpJOEcWglRfyfQpLqjlC0O+dptruTPE8VamKYuqdjxfqD8hifZDOA==} + engines: {node: '>= 10'} + cpu: [x64] + os: [darwin] + + '@napi-rs/tar-freebsd-x64@1.1.0': + resolution: {integrity: sha512-J0PIqX+pl6lBIAckL/c87gpodLbjZB1OtIK+RDscKC9NLdpVv6VGOxzUV/fYev/hctcE8EfkLbgFOfpmVQPg2g==} + engines: {node: '>= 10'} + cpu: [x64] + os: [freebsd] + + '@napi-rs/tar-linux-arm-gnueabihf@1.1.0': + resolution: {integrity: sha512-SLgIQo3f3EjkZ82ZwvrEgFvMdDAhsxCYjyoSuWfHCz0U16qx3SuGCp8+FYOPYCECHN3ZlGjXnoAIt9ERd0dEUg==} + engines: {node: '>= 10'} + cpu: [arm] + os: [linux] + + '@napi-rs/tar-linux-arm64-gnu@1.1.0': + resolution: {integrity: sha512-d014cdle52EGaH6GpYTQOP9Py7glMO1zz/+ynJPjjzYFSxvdYx0byrjumZk2UQdIyGZiJO2MEFpCkEEKFSgPYA==} + engines: {node: '>= 10'} + cpu: [arm64] + os: [linux] + + '@napi-rs/tar-linux-arm64-musl@1.1.0': + resolution: {integrity: sha512-L/y1/26q9L/uBqiW/JdOb/Dc94egFvNALUZV2WCGKQXc6UByPBMgdiEyW2dtoYxYYYYc+AKD+jr+wQPcvX2vrQ==} + engines: {node: '>= 10'} + cpu: [arm64] + os: [linux] + + '@napi-rs/tar-linux-ppc64-gnu@1.1.0': + resolution: {integrity: sha512-EPE1K/80RQvPbLRJDJs1QmCIcH+7WRi0F73+oTe1582y9RtfGRuzAkzeBuAGRXAQEjRQw/RjtNqr6UTJ+8UuWQ==} + engines: {node: '>= 10'} + cpu: [ppc64] + os: [linux] + + '@napi-rs/tar-linux-s390x-gnu@1.1.0': + resolution: {integrity: sha512-B2jhWiB1ffw1nQBqLUP1h4+J1ovAxBOoe5N2IqDMOc63fsPZKNqF1PvO/dIem8z7LL4U4bsfmhy3gBfu547oNQ==} + engines: {node: '>= 10'} + cpu: [s390x] + os: [linux] + + '@napi-rs/tar-linux-x64-gnu@1.1.0': + resolution: {integrity: sha512-tbZDHnb9617lTnsDMGo/eAMZxnsQFnaRe+MszRqHguKfMwkisc9CCJnks/r1o84u5fECI+J/HOrKXgczq/3Oww==} + engines: {node: '>= 10'} + cpu: [x64] + os: [linux] + + '@napi-rs/tar-linux-x64-musl@1.1.0': + resolution: {integrity: sha512-dV6cODlzbO8u6Anmv2N/ilQHq/AWz0xyltuXoLU3yUyXbZcnWYZuB2rL8OBGPmqNcD+x9NdScBNXh7vWN0naSQ==} + engines: {node: '>= 10'} + cpu: [x64] + os: [linux] + + '@napi-rs/tar-wasm32-wasi@1.1.0': + resolution: {integrity: sha512-jIa9nb2HzOrfH0F8QQ9g3WE4aMH5vSI5/1NYVNm9ysCmNjCCtMXCAhlI3WKCdm/DwHf0zLqdrrtDFXODcNaqMw==} + engines: {node: '>=14.0.0'} + cpu: [wasm32] + + '@napi-rs/tar-win32-arm64-msvc@1.1.0': + resolution: {integrity: sha512-vfpG71OB0ijtjemp3WTdmBKJm9R70KM8vsSExMsIQtV0lVzP07oM1CW6JbNRPXNLhRoue9ofYLiUDk8bE0Hckg==} + engines: {node: '>= 10'} + cpu: [arm64] + os: [win32] + + '@napi-rs/tar-win32-ia32-msvc@1.1.0': + resolution: {integrity: sha512-hGPyPW60YSpOSgzfy68DLBHgi6HxkAM+L59ZZZPMQ0TOXjQg+p2EW87+TjZfJOkSpbYiEkULwa/f4a2hcVjsqQ==} + engines: {node: '>= 10'} + cpu: [ia32] + os: [win32] + + '@napi-rs/tar-win32-x64-msvc@1.1.0': + resolution: {integrity: sha512-L6Ed1DxXK9YSCMyvpR8MiNAyKNkQLjsHsHK9E0qnHa8NzLFqzDKhvs5LfnWxM2kJ+F7m/e5n9zPm24kHb3LsVw==} + engines: {node: '>= 10'} + cpu: [x64] + os: [win32] + + '@napi-rs/tar@1.1.0': + resolution: {integrity: sha512-7cmzIu+Vbupriudo7UudoMRH2OA3cTw67vva8MxeoAe5S7vPFI7z0vp0pMXiA25S8IUJefImQ90FeJjl8fjEaQ==} + engines: {node: '>= 10'} + + '@napi-rs/wasm-runtime@1.1.1': + resolution: {integrity: sha512-p64ah1M1ld8xjWv3qbvFwHiFVWrq1yFvV4f7w+mzaqiR4IlSgkqhcRdHwsGgomwzBH51sRY4NEowLxnaBjcW/A==} + + '@napi-rs/wasm-tools-android-arm-eabi@1.0.1': + resolution: {integrity: sha512-lr07E/l571Gft5v4aA1dI8koJEmF1F0UigBbsqg9OWNzg80H3lDPO+auv85y3T/NHE3GirDk7x/D3sLO57vayw==} + engines: {node: '>= 10'} + cpu: [arm] + os: [android] + + '@napi-rs/wasm-tools-android-arm64@1.0.1': + resolution: {integrity: sha512-WDR7S+aRLV6LtBJAg5fmjKkTZIdrEnnQxgdsb7Cf8pYiMWBHLU+LC49OUVppQ2YSPY0+GeYm9yuZWW3kLjJ7Bg==} + engines: {node: '>= 10'} + cpu: [arm64] + os: [android] + + '@napi-rs/wasm-tools-darwin-arm64@1.0.1': + resolution: {integrity: sha512-qWTI+EEkiN0oIn/N2gQo7+TVYil+AJ20jjuzD2vATS6uIjVz+Updeqmszi7zq7rdFTLp6Ea3/z4kDKIfZwmR9g==} + engines: {node: '>= 10'} + cpu: [arm64] + os: [darwin] + + '@napi-rs/wasm-tools-darwin-x64@1.0.1': + resolution: {integrity: sha512-bA6hubqtHROR5UI3tToAF/c6TDmaAgF0SWgo4rADHtQ4wdn0JeogvOk50gs2TYVhKPE2ZD2+qqt7oBKB+sxW3A==} + engines: {node: '>= 10'} + cpu: [x64] + os: [darwin] + + '@napi-rs/wasm-tools-freebsd-x64@1.0.1': + resolution: {integrity: sha512-90+KLBkD9hZEjPQW1MDfwSt5J1L46EUKacpCZWyRuL6iIEO5CgWU0V/JnEgFsDOGyyYtiTvHc5bUdUTWd4I9Vg==} + engines: {node: '>= 10'} + cpu: [x64] + os: [freebsd] + + '@napi-rs/wasm-tools-linux-arm64-gnu@1.0.1': + resolution: {integrity: sha512-rG0QlS65x9K/u3HrKafDf8cFKj5wV2JHGfl8abWgKew0GVPyp6vfsDweOwHbWAjcHtp2LHi6JHoW80/MTHm52Q==} + engines: {node: '>= 10'} + cpu: [arm64] + os: [linux] + + '@napi-rs/wasm-tools-linux-arm64-musl@1.0.1': + resolution: {integrity: sha512-jAasbIvjZXCgX0TCuEFQr+4D6Lla/3AAVx2LmDuMjgG4xoIXzjKWl7c4chuaD+TI+prWT0X6LJcdzFT+ROKGHQ==} + engines: {node: '>= 10'} + cpu: [arm64] + os: [linux] + + '@napi-rs/wasm-tools-linux-x64-gnu@1.0.1': + resolution: {integrity: sha512-Plgk5rPqqK2nocBGajkMVbGm010Z7dnUgq0wtnYRZbzWWxwWcXfZMPa8EYxrK4eE8SzpI7VlZP1tdVsdjgGwMw==} + engines: {node: '>= 10'} + cpu: [x64] + os: [linux] + + '@napi-rs/wasm-tools-linux-x64-musl@1.0.1': + resolution: {integrity: sha512-GW7AzGuWxtQkyHknHWYFdR0CHmW6is8rG2Rf4V6GNmMpmwtXt/ItWYWtBe4zqJWycMNazpfZKSw/BpT7/MVCXQ==} + engines: {node: '>= 10'} + cpu: [x64] + os: [linux] + + '@napi-rs/wasm-tools-wasm32-wasi@1.0.1': + resolution: {integrity: sha512-/nQVSTrqSsn7YdAc2R7Ips/tnw5SPUcl3D7QrXCNGPqjbatIspnaexvaOYNyKMU6xPu+pc0BTnKVmqhlJJCPLA==} + engines: {node: '>=14.0.0'} + cpu: [wasm32] + + '@napi-rs/wasm-tools-win32-arm64-msvc@1.0.1': + resolution: {integrity: sha512-PFi7oJIBu5w7Qzh3dwFea3sHRO3pojMsaEnUIy22QvsW+UJfNQwJCryVrpoUt8m4QyZXI+saEq/0r4GwdoHYFQ==} + engines: {node: '>= 10'} + cpu: [arm64] + os: [win32] + + '@napi-rs/wasm-tools-win32-ia32-msvc@1.0.1': + resolution: {integrity: sha512-gXkuYzxQsgkj05Zaq+KQTkHIN83dFAwMcTKa2aQcpYPRImFm2AQzEyLtpXmyCWzJ0F9ZYAOmbSyrNew8/us6bw==} + engines: {node: '>= 10'} + cpu: [ia32] + os: [win32] + + '@napi-rs/wasm-tools-win32-x64-msvc@1.0.1': + resolution: {integrity: sha512-rEAf05nol3e3eei2sRButmgXP+6ATgm0/38MKhz9Isne82T4rPIMYsCIFj0kOisaGeVwoi2fnm7O9oWp5YVnYQ==} + engines: {node: '>= 10'} + cpu: [x64] + os: [win32] + + '@napi-rs/wasm-tools@1.0.1': + resolution: {integrity: sha512-enkZYyuCdo+9jneCPE/0fjIta4wWnvVN9hBo2HuiMpRF0q3lzv1J6b/cl7i0mxZUKhBrV3aCKDBQnCOhwKbPmQ==} + engines: {node: '>= 10'} '@neoconfetti/react@1.0.0': resolution: {integrity: sha512-klcSooChXXOzIm+SE5IISIAn3bYzYfPjbX7D7HoqZL84oAfgREeSg5vSIaSFH+DaGzzvImTyWe1OyrJ67vik4A==} @@ -1113,6 +1584,58 @@ packages: resolution: {integrity: sha512-oGB+UxlgWcgQkgwo8GcEGwemoTFt3FIO9ababBmaGwXIoBKZ+GTy0pP185beGg7Llih/NSHSV2XAs1lnznocSg==} engines: {node: '>= 8'} + '@octokit/auth-token@6.0.0': + resolution: {integrity: sha512-P4YJBPdPSpWTQ1NU4XYdvHvXJJDxM6YwpS0FZHRgP7YFkdVxsWcpWGy/NVqlAA7PcPCnMacXlRm1y2PFZRWL/w==} + engines: {node: '>= 20'} + + '@octokit/core@7.0.6': + resolution: {integrity: sha512-DhGl4xMVFGVIyMwswXeyzdL4uXD5OGILGX5N8Y+f6W7LhC1Ze2poSNrkF/fedpVDHEEZ+PHFW0vL14I+mm8K3Q==} + engines: {node: '>= 20'} + + '@octokit/endpoint@11.0.2': + resolution: {integrity: sha512-4zCpzP1fWc7QlqunZ5bSEjxc6yLAlRTnDwKtgXfcI/FxxGoqedDG8V2+xJ60bV2kODqcGB+nATdtap/XYq2NZQ==} + engines: {node: '>= 20'} + + '@octokit/graphql@9.0.3': + resolution: {integrity: sha512-grAEuupr/C1rALFnXTv6ZQhFuL1D8G5y8CN04RgrO4FIPMrtm+mcZzFG7dcBm+nq+1ppNixu+Jd78aeJOYxlGA==} + engines: {node: '>= 20'} + + '@octokit/openapi-types@27.0.0': + resolution: {integrity: sha512-whrdktVs1h6gtR+09+QsNk2+FO+49j6ga1c55YZudfEG+oKJVvJLQi3zkOm5JjiUXAagWK2tI2kTGKJ2Ys7MGA==} + + '@octokit/plugin-paginate-rest@14.0.0': + resolution: {integrity: sha512-fNVRE7ufJiAA3XUrha2omTA39M6IXIc6GIZLvlbsm8QOQCYvpq/LkMNGyFlB1d8hTDzsAXa3OKtybdMAYsV/fw==} + engines: {node: '>= 20'} + peerDependencies: + '@octokit/core': '>=6' + + '@octokit/plugin-request-log@6.0.0': + resolution: {integrity: sha512-UkOzeEN3W91/eBq9sPZNQ7sUBvYCqYbrrD8gTbBuGtHEuycE4/awMXcYvx6sVYo7LypPhmQwwpUe4Yyu4QZN5Q==} + engines: {node: '>= 20'} + peerDependencies: + '@octokit/core': '>=6' + + '@octokit/plugin-rest-endpoint-methods@17.0.0': + resolution: {integrity: sha512-B5yCyIlOJFPqUUeiD0cnBJwWJO8lkJs5d8+ze9QDP6SvfiXSz1BF+91+0MeI1d2yxgOhU/O+CvtiZ9jSkHhFAw==} + engines: {node: '>= 20'} + peerDependencies: + '@octokit/core': '>=6' + + '@octokit/request-error@7.1.0': + resolution: {integrity: sha512-KMQIfq5sOPpkQYajXHwnhjCC0slzCNScLHs9JafXc4RAJI+9f+jNDlBNaIMTvazOPLgb4BnlhGJOTbnN0wIjPw==} + engines: {node: '>= 20'} + + '@octokit/request@10.0.7': + resolution: {integrity: sha512-v93h0i1yu4idj8qFPZwjehoJx4j3Ntn+JhXsdJrG9pYaX6j/XRz2RmasMUHtNgQD39nrv/VwTWSqK0RNXR8upA==} + engines: {node: '>= 20'} + + '@octokit/rest@22.0.1': + resolution: {integrity: sha512-Jzbhzl3CEexhnivb1iQ0KJ7s5vvjMWcmRtq5aUsKmKDrRW6z3r84ngmiFKFvpZjpiU/9/S6ITPFRpn5s/3uQJw==} + engines: {node: '>= 20'} + + '@octokit/types@16.0.0': + resolution: {integrity: sha512-sKq+9r1Mm4efXW1FCk7hFSeJo4QKreL/tTbR0rz/qx/r1Oa2VV83LTA/H/MuCOX7uCIJmQVRKBcbmWoySjAnSg==} + '@playwright/test@1.56.1': resolution: {integrity: sha512-vSMYtL/zOcFpvJCW71Q/OEGQb7KYBPAdKh35WNSkaZA75JlAO8ED8UN6GUNTm3drWomcbcqRPFqQbLae8yBTdg==} engines: {node: '>=18'} @@ -2255,6 +2778,9 @@ packages: peerDependencies: '@testing-library/dom': '>=7.21.4' + '@tybys/wasm-util@0.10.1': + resolution: {integrity: sha512-9tTaPJLSiejZKx+Bmog4uSubteqTvFrVrURwkmHixBo0G4seD0zUxp98E1DzUBJxLQ3NPwXrGKDiVjwx/DpPsg==} + '@types/aria-query@5.0.4': resolution: {integrity: sha512-rfT93uj5s0PRL7EzccGMs3brplhcrghnDoV26NqKhCAS1hVo+WdNsPvE/yb6ilfr5hi2MEk6d5EWJTKdxg8jVw==} @@ -2670,6 +3196,10 @@ packages: resolution: {integrity: sha512-Cxwpt2SfTzTtXcfOlzGEee8O+c+MmUgGrNiBcXnuWxuFJHe6a5Hz7qwhwe5OgaSYI0IJvkLqWX1ASG+cJOkEiA==} engines: {node: '>=10'} + ansi-styles@6.2.3: + resolution: {integrity: sha512-4Dj6M28JB+oAH8kFkTLUo+a2jwOFkuqb3yucU0CANcRRUbxS0cP0nZYCGjcc3BNXwRIsUVmDGgzawme7zvJHvg==} + engines: {node: '>=12'} + arg@5.0.2: resolution: {integrity: sha512-PYjyFOLKQ9y57JvQ6QLo8dAgNqswh8M1RMJYdQduT6xbWSgK36P/Z/v+p888pM69jMMfS8Xd8F6I1kQ/I9HUGg==} @@ -2757,6 +3287,9 @@ packages: resolution: {integrity: sha512-NF+epuEdnUYVlGuhaxbbq+dvJttwLnGY+YixlXlME5KpQ5W3CnXA5cVTneY3SPbPDRkcjMbifrwmFYcClgOZeg==} engines: {node: '>= 0.8'} + before-after-hook@4.0.0: + resolution: {integrity: sha512-q6tR3RPqIB1pMiTRMFcZwuG5T8vwp+vUvEG0vuI6B+Rikh5BfPp2fQ82c925FOs+b0lcFQ8CFrL+KbilfZFhOQ==} + body-parser@1.20.3: resolution: {integrity: sha512-7rAxByjUMqQ3/bHJy7D6OGXvx/MMc4IqBn/X0fcM1QUcAItpZrBEYhWGem+tzXH90c+G01ypMcYJBO9Y30203g==} engines: {node: '>= 0.8', npm: 1.2.8000 || >= 1.4.16} @@ -2834,6 +3367,9 @@ packages: resolution: {integrity: sha512-oKnbhFyRIXpUuez8iBMmyEa4nbj4IOQyuhc/wy9kY7/WVPcwIO9VA668Pu8RkO7+0G76SLROeyw9CpQ061i4mA==} engines: {node: '>=10'} + chardet@2.1.1: + resolution: {integrity: sha512-PsezH1rqdV9VvyNhxxOW32/d75r01NY7TQCmOqomRo15ZSOKbpTFVsfjghxo6JloQUCGnH4k1LGu0R4yCLlWQQ==} + check-error@2.1.1: resolution: {integrity: sha512-OAlb+T7V4Op9OwdkjmguYRqncdlx5JiofwOAUkmTF+jNdHwzTaTs4sRAGpzLF3oOz5xAyDGrPgeIDFQmDOTiJw==} engines: {node: '>= 16'} @@ -2857,6 +3393,15 @@ packages: class-variance-authority@0.7.1: resolution: {integrity: sha512-Ka+9Trutv7G8M6WT6SeiRWz792K5qEqIGEGzXKhAE6xOWAY6pPH8U+9IY3oCMv6kqTmLsv7Xh/2w2RigkePMsg==} + cli-width@4.1.0: + resolution: {integrity: sha512-ouuZd4/dm2Sw5Gmqy6bGyNNNe1qt9RpmxveLSO7KcgsTnU7RXfsw+/bukWGo1abgBiMAic068rclZsO4IWmmxQ==} + engines: {node: '>= 12'} + + clipanion@4.0.0-rc.4: + resolution: {integrity: sha512-CXkMQxU6s9GklO/1f714dkKBMu1lopS1WFF0B8o4AxPykR1hpozxSiUZ5ZUeBjfPgCWqbcNOtZVFhB8Lkfp1+Q==} + peerDependencies: + typanion: '*' + clsx@2.1.1: resolution: {integrity: sha512-eYm0QWBtUrBWZWG0d386OGAw16Z995PiOVo2B7bjWSbHedGl5e0ZWaq65kOGgUSNesEIDkB9ISbTg/JK9dhCZA==} engines: {node: '>=6'} @@ -2877,6 +3422,9 @@ packages: color-name@1.1.4: resolution: {integrity: sha512-dOy+3AuW3a2wNbZHIuMZpTcgjGuLU/uBL/ubcZF9OXbDo8ff4O8yVp5Bf0efS8uEoYo5q4Fx7dY9OgQGXgAsQA==} + colorette@2.0.20: + resolution: {integrity: sha512-IfEDxwoWIjkeXL1eXcDiow4UbKjhLdq6/EuSVR9GMN7KVH3r9gQ83e73hsz1Nd1T3ijd5xv1wcWRYO+D6kCI2w==} + combined-stream@1.0.8: resolution: {integrity: sha512-FQN4MRfuJeHf7cBbBMJFXhKSDq+2kAArBlmRBvcvFE5BB1HZKXtSFASDhdlz9zOYwxh8lDdnvmMOe/+5cdoEdg==} engines: {node: '>= 0.8'} @@ -3109,6 +3657,17 @@ packages: electron-to-chromium@1.5.157: resolution: {integrity: sha512-/0ybgsQd1muo8QlnuTpKwtl0oX5YMlUGbm8xyqgDU00motRkKFFbUJySAQBWcY79rVqNLWIWa87BGVGClwAB2w==} + emnapi@1.8.1: + resolution: {integrity: sha512-34i2BbgHx1LnEO4JCGQYo6h6s4e4KrdWtdTHfllBNLbXSHPmdIHplxKejfabsRK+ukNciqVdalB+fxMibqHdaQ==} + peerDependencies: + node-addon-api: '>= 6.1.0' + peerDependenciesMeta: + node-addon-api: + optional: true + + emoji-regex@10.6.0: + resolution: {integrity: sha512-toUI84YS5YmxW219erniWD0CIVOo46xGKColeNQRgOzDorgBi1v4D71/OFzgD9GO2UGKIv1C3Sp8DAn0+j5w7A==} + empathic@2.0.0: resolution: {integrity: sha512-i6UzDscO/XfAcNYD75CfICkmfLedpyPDdozrLMmQc5ORaQcdMoc21OnlEylMIqI7U8eniKrPMxxtj8k0vhmJhA==} engines: {node: '>=14'} @@ -3163,6 +3722,9 @@ packages: resolution: {integrity: sha512-w+5mJ3GuFL+NjVtJlvydShqE1eN3h3PbI7/5LAsYJP/2qtuMXjfL2LpHSRqo4b4eSF5K/DH1JXKUAHSB2UW50g==} engines: {node: '>= 0.4'} + es-toolkit@1.44.0: + resolution: {integrity: sha512-6penXeZalaV88MM3cGkFZZfOoLGWshWWfdy0tWw/RlVVyhvMaWSBTOvXNeiW3e5FwdS5ePW0LGEu17zT139ktg==} + esbuild@0.25.10: resolution: {integrity: sha512-9RiGKvCwaqxO2owP61uQ4BgNborAQskMR6QusfWzQqv7AZOg5oGehdY2pRJMTKuwxd1IDBP4rSbI5lHzU7SMsQ==} engines: {node: '>=18'} @@ -3277,6 +3839,9 @@ packages: exsolve@1.0.8: resolution: {integrity: sha512-LmDxfWXwcTArk8fUEnOfSZpHOJ6zOMUJKOtFLFqJLoKJetuQG874Uc7/Kki7zFLzYybmZhp1M7+98pfMqeX8yA==} + fast-content-type-parse@3.0.0: + resolution: {integrity: sha512-ZvLdcY8P+N8mGQJahJV5G4U88CSvT1rP8ApL6uETe88MBXrBHAkZlSEySdUlyztF7ccb+Znos3TFqaepHxdhBg==} + fast-deep-equal@3.1.3: resolution: {integrity: sha512-f3qQ9oQy9j2AhBe/H9VC91wLmKBCCU/gDOnKNAYG5hswO7BLKj09Hc5HYNz9cGI++xlpDCIgDaitVs03ATR84Q==} @@ -3389,6 +3954,10 @@ packages: resolution: {integrity: sha512-3hN7NaskYvMDLQY55gnW3NQ+mesEAepTqlg+VEbj7zzqEMBVNhzcGYYeqFo/TlYz6eQiFcp1HcsCZO+nGgS8zg==} engines: {node: '>=6.9.0'} + get-east-asian-width@1.4.0: + resolution: {integrity: sha512-QZjmEOC+IT1uk6Rx0sX22V6uHWVwbdbxf1faPqJ1QhLdGgsRGCZoyaQBm/piRdJy/D2um6hM1UP7ZEeQ4EkP+Q==} + engines: {node: '>=18'} + get-intrinsic@1.3.0: resolution: {integrity: sha512-9fSjSaos/fRIVIp+xSJlE6lfwhES7LNtKaCBIamHsjr2na1BiABJPo0mOjjz8GJDURarmCPGqaiVg5mfjb98CQ==} engines: {node: '>= 0.4'} @@ -3486,6 +4055,10 @@ packages: resolution: {integrity: sha512-v3MXnZAcvnywkTUEZomIActle7RXXeedOR31wwl7VlyoXO4Qi9arvSenNQWne1TcRwhCL1HwLI21bEqdpj8/rA==} engines: {node: '>=0.10.0'} + iconv-lite@0.7.2: + resolution: {integrity: sha512-im9DjEDQ55s9fL4EYzOAv0yMqmMBSZp6G0VvFyTMPKWxiSBHUj9NW/qqLmXUwXrrM7AvqSlTCfvqRb0cM8yYqw==} + engines: {node: '>=0.10.0'} + ignore@5.3.2: resolution: {integrity: sha512-hsBTNUqQTDwkWtcdYI2i06Y/nUBEsNEDJKjWdigLvegy8kDuJAS8uRlpkkcQpyEXL0Z/pjDy5HBmMjRCJ2gq+g==} engines: {node: '>= 4'} @@ -3927,6 +4500,10 @@ packages: ms@2.1.3: resolution: {integrity: sha512-6FlzubTLZG3J2a/NVCAleEhjzq5oxgHyaCU9yYXvcLsvoVaHJq/s5xXI6/XXP6tz7R9xAOtHnSO/tXtF3WRTlA==} + mute-stream@3.0.0: + resolution: {integrity: sha512-dkEJPVvun4FryqBmZ5KhDo0K9iDXAwn08tMLDinNdRBNPcYEDiWYysLcc6k3mjTMlbP9KyylvRpd4wFtwrT9rw==} + engines: {node: ^20.17.0 || >=22.9.0} + nanoid@3.3.11: resolution: {integrity: sha512-N8SpfPUnUp1bK+PMYW8qSWdl9U+wwNWI4QKxOYDy9JAro3WMX7p2OeVRF9v+347pnakNevPmiHhNmZ2HbFA76w==} engines: {node: ^10 || ^12 || ^13.7 || ^14 || >=15.0.1} @@ -3974,6 +4551,9 @@ packages: resolution: {integrity: sha512-gXah6aZrcUxjWg2zR2MwouP2eHlCBzdV4pygudehaKXSGW4v2AsRQUK+lwwXhii6KFZcunEnmSUoYp5CXibxtA==} engines: {node: '>= 0.4'} + obug@2.1.1: + resolution: {integrity: sha512-uTqF9MuPraAQ+IsnPf366RG4cP9RtUi7MLO1N3KEc+wb0a6yKpeL0lmk2IB1jY5KHPAlTc6T/JRdC/YqxHNwkQ==} + on-finished@2.3.0: resolution: {integrity: sha512-ikqdkGAAyf/X/gPhXGvfgAytDZtDbr+bkNUJ0N9h5MI/dmdgCs3l6hoHrcUv41sRKew3jIwrp4qQDXiK99Utww==} engines: {node: '>= 0.8'} @@ -4485,6 +5065,10 @@ packages: siginfo@2.0.0: resolution: {integrity: sha512-ybx0WO1/8bSBLEWXZvEd7gMW3Sn3JFlW3TvX1nREbDLRNQNaeNN8WK0meBwPdAaOI7TtRRRJn/Es1zhrrCHu7g==} + signal-exit@4.1.0: + resolution: {integrity: sha512-bzyZ1e88w9O1iNJbKnOlvYTrWPDl46O1bG0D3XInv+9tkPrxrN8jUUTiFlDkkmKWgn1M6CfIA13SuGqOa9Korw==} + engines: {node: '>=14'} + sirv@3.0.2: resolution: {integrity: sha512-2wcC/oGxHis/BoHkkPwldgiPSYcpZK3JU28WoMVv55yHJgcZ8rlXvuG9iZggz+sU1d4bRgIGASwyWqjxu3FM0g==} engines: {node: '>=18'} @@ -4540,6 +5124,10 @@ packages: prettier: optional: true + string-width@7.2.0: + resolution: {integrity: sha512-tsaTIkKW9b4N+AEj+SVA+WhJzV7/zMhcSu78mLKWSk7cXMOSHsBKFWUs0fWwq8QyK3MgJBQRX6Gbi4kYbdvGkQ==} + engines: {node: '>=18'} + string.prototype.matchall@4.0.12: resolution: {integrity: sha512-6CC9uyBL+/48dYizRf7H7VAYCMCNTBeM78x/VTUe9bFEaxBepPJDa1Ow99LqI/1yF7kuy7Q3cQsYMrcjGUcskA==} engines: {node: '>= 0.4'} @@ -4681,6 +5269,9 @@ packages: engines: {node: '>=18.0.0'} hasBin: true + typanion@3.14.0: + resolution: {integrity: sha512-ZW/lVMRabETuYCd9O9ZvMhAh8GslSqaUjxmK/JLPCh6l73CvLBiuXswj/+7LdnWOgYsQ130FqLzFz5aGT4I3Ug==} + type-check@0.4.0: resolution: {integrity: sha512-XleUoc9uwGXqjWwXaUTZAmzMcFZ5858QA2vvx1Ur5xIcixXIP+8LnFDgRplU30us6teqdlskFfu+ae4K79Ooew==} engines: {node: '>= 0.8.0'} @@ -4745,6 +5336,9 @@ packages: undici-types@7.16.0: resolution: {integrity: sha512-Zz+aZWSj8LE6zoxD+xrjh4VfkIG8Ya6LvYkZqtUQGJPZjYl53ypCaUwWqo7eI0x66KBGeRo+mlBEkMSeSZ38Nw==} + universal-user-agent@7.0.3: + resolution: {integrity: sha512-TmnEAEAsBJVZM/AADELsK76llnwcf9vMKuPz8JflO1frO8Lchitr0fNaN9d+Ap0BjKtqWqd/J17qeDnXh8CL2A==} + universalify@2.0.1: resolution: {integrity: sha512-gptHNQghINnc/vTGIk0SOFGFNXw7JVrlRUtConJRlvaw6DuX0wO5Jeko9sWrMBhh+PsYAZ7oXAiOnf/UKogyiw==} engines: {node: '>= 10.0.0'} @@ -4951,6 +5545,10 @@ packages: resolution: {integrity: sha512-BN22B5eaMMI9UMtjrGd5g5eCYPpCPDUy0FJXbYsaT5zYxjFOckS53SQDE3pWkVoWpHXVb3BrYcEN4Twa55B5cA==} engines: {node: '>=0.10.0'} + wrap-ansi@9.0.2: + resolution: {integrity: sha512-42AtmgqjV+X1VpdOfyTGOYRi0/zsoLqtXQckTmqTeybT+BDIbM/Guxo7x3pE2vtpr1ok6xRqM9OpBe+Jyoqyww==} + engines: {node: '>=18'} + ws@8.18.3: resolution: {integrity: sha512-PEIGCY5tSlUt50cqyMXfCzX+oOPqN0vuGqWzbcJ2xvnkzkq46oOpz7dQaTDBdfICb4N14+GARUDw2XV2N4tvzg==} engines: {node: '>=10.0.0'} @@ -5357,6 +5955,22 @@ snapshots: style-mod: 4.1.2 w3c-keyname: 2.2.8 + '@emnapi/core@1.8.1': + dependencies: + '@emnapi/wasi-threads': 1.1.0 + tslib: 2.8.1 + optional: true + + '@emnapi/runtime@1.8.1': + dependencies: + tslib: 2.8.1 + optional: true + + '@emnapi/wasi-threads@1.1.0': + dependencies: + tslib: 2.8.1 + optional: true + '@esbuild/aix-ppc64@0.25.10': optional: true @@ -5591,6 +6205,125 @@ snapshots: '@humanwhocodes/retry@0.4.3': {} + '@inquirer/ansi@2.0.3': {} + + '@inquirer/checkbox@5.0.4(@types/node@20.19.4)': + dependencies: + '@inquirer/ansi': 2.0.3 + '@inquirer/core': 11.1.1(@types/node@20.19.4) + '@inquirer/figures': 2.0.3 + '@inquirer/type': 4.0.3(@types/node@20.19.4) + optionalDependencies: + '@types/node': 20.19.4 + + '@inquirer/confirm@6.0.4(@types/node@20.19.4)': + dependencies: + '@inquirer/core': 11.1.1(@types/node@20.19.4) + '@inquirer/type': 4.0.3(@types/node@20.19.4) + optionalDependencies: + '@types/node': 20.19.4 + + '@inquirer/core@11.1.1(@types/node@20.19.4)': + dependencies: + '@inquirer/ansi': 2.0.3 + '@inquirer/figures': 2.0.3 + '@inquirer/type': 4.0.3(@types/node@20.19.4) + cli-width: 4.1.0 + mute-stream: 3.0.0 + signal-exit: 4.1.0 + wrap-ansi: 9.0.2 + optionalDependencies: + '@types/node': 20.19.4 + + '@inquirer/editor@5.0.4(@types/node@20.19.4)': + dependencies: + '@inquirer/core': 11.1.1(@types/node@20.19.4) + '@inquirer/external-editor': 2.0.3(@types/node@20.19.4) + '@inquirer/type': 4.0.3(@types/node@20.19.4) + optionalDependencies: + '@types/node': 20.19.4 + + '@inquirer/expand@5.0.4(@types/node@20.19.4)': + dependencies: + '@inquirer/core': 11.1.1(@types/node@20.19.4) + '@inquirer/type': 4.0.3(@types/node@20.19.4) + optionalDependencies: + '@types/node': 20.19.4 + + '@inquirer/external-editor@2.0.3(@types/node@20.19.4)': + dependencies: + chardet: 2.1.1 + iconv-lite: 0.7.2 + optionalDependencies: + '@types/node': 20.19.4 + + '@inquirer/figures@2.0.3': {} + + '@inquirer/input@5.0.4(@types/node@20.19.4)': + dependencies: + '@inquirer/core': 11.1.1(@types/node@20.19.4) + '@inquirer/type': 4.0.3(@types/node@20.19.4) + optionalDependencies: + '@types/node': 20.19.4 + + '@inquirer/number@4.0.4(@types/node@20.19.4)': + dependencies: + '@inquirer/core': 11.1.1(@types/node@20.19.4) + '@inquirer/type': 4.0.3(@types/node@20.19.4) + optionalDependencies: + '@types/node': 20.19.4 + + '@inquirer/password@5.0.4(@types/node@20.19.4)': + dependencies: + '@inquirer/ansi': 2.0.3 + '@inquirer/core': 11.1.1(@types/node@20.19.4) + '@inquirer/type': 4.0.3(@types/node@20.19.4) + optionalDependencies: + '@types/node': 20.19.4 + + '@inquirer/prompts@8.2.0(@types/node@20.19.4)': + dependencies: + '@inquirer/checkbox': 5.0.4(@types/node@20.19.4) + '@inquirer/confirm': 6.0.4(@types/node@20.19.4) + '@inquirer/editor': 5.0.4(@types/node@20.19.4) + '@inquirer/expand': 5.0.4(@types/node@20.19.4) + '@inquirer/input': 5.0.4(@types/node@20.19.4) + '@inquirer/number': 4.0.4(@types/node@20.19.4) + '@inquirer/password': 5.0.4(@types/node@20.19.4) + '@inquirer/rawlist': 5.2.0(@types/node@20.19.4) + '@inquirer/search': 4.1.0(@types/node@20.19.4) + '@inquirer/select': 5.0.4(@types/node@20.19.4) + optionalDependencies: + '@types/node': 20.19.4 + + '@inquirer/rawlist@5.2.0(@types/node@20.19.4)': + dependencies: + '@inquirer/core': 11.1.1(@types/node@20.19.4) + '@inquirer/type': 4.0.3(@types/node@20.19.4) + optionalDependencies: + '@types/node': 20.19.4 + + '@inquirer/search@4.1.0(@types/node@20.19.4)': + dependencies: + '@inquirer/core': 11.1.1(@types/node@20.19.4) + '@inquirer/figures': 2.0.3 + '@inquirer/type': 4.0.3(@types/node@20.19.4) + optionalDependencies: + '@types/node': 20.19.4 + + '@inquirer/select@5.0.4(@types/node@20.19.4)': + dependencies: + '@inquirer/ansi': 2.0.3 + '@inquirer/core': 11.1.1(@types/node@20.19.4) + '@inquirer/figures': 2.0.3 + '@inquirer/type': 4.0.3(@types/node@20.19.4) + optionalDependencies: + '@types/node': 20.19.4 + + '@inquirer/type@4.0.3(@types/node@20.19.4)': + optionalDependencies: + '@types/node': 20.19.4 + '@isaacs/balanced-match@4.0.1': {} '@isaacs/brace-expansion@5.0.0': @@ -5683,7 +6416,250 @@ snapshots: '@mjackson/node-fetch-server@0.2.0': {} - '@napi-rs/cli@2.18.4': {} + '@napi-rs/cli@3.5.1(@emnapi/runtime@1.8.1)(@types/node@20.19.4)': + dependencies: + '@inquirer/prompts': 8.2.0(@types/node@20.19.4) + '@napi-rs/cross-toolchain': 1.0.3 + '@napi-rs/wasm-tools': 1.0.1 + '@octokit/rest': 22.0.1 + clipanion: 4.0.0-rc.4(typanion@3.14.0) + colorette: 2.0.20 + emnapi: 1.8.1 + es-toolkit: 1.44.0 + js-yaml: 4.1.1 + obug: 2.1.1 + semver: 7.7.3 + typanion: 3.14.0 + optionalDependencies: + '@emnapi/runtime': 1.8.1 + transitivePeerDependencies: + - '@napi-rs/cross-toolchain-arm64-target-aarch64' + - '@napi-rs/cross-toolchain-arm64-target-armv7' + - '@napi-rs/cross-toolchain-arm64-target-ppc64le' + - '@napi-rs/cross-toolchain-arm64-target-s390x' + - '@napi-rs/cross-toolchain-arm64-target-x86_64' + - '@napi-rs/cross-toolchain-x64-target-aarch64' + - '@napi-rs/cross-toolchain-x64-target-armv7' + - '@napi-rs/cross-toolchain-x64-target-ppc64le' + - '@napi-rs/cross-toolchain-x64-target-s390x' + - '@napi-rs/cross-toolchain-x64-target-x86_64' + - '@types/node' + - node-addon-api + - supports-color + + '@napi-rs/cross-toolchain@1.0.3': + dependencies: + '@napi-rs/lzma': 1.4.5 + '@napi-rs/tar': 1.1.0 + debug: 4.4.3 + transitivePeerDependencies: + - supports-color + + '@napi-rs/lzma-android-arm-eabi@1.4.5': + optional: true + + '@napi-rs/lzma-android-arm64@1.4.5': + optional: true + + '@napi-rs/lzma-darwin-arm64@1.4.5': + optional: true + + '@napi-rs/lzma-darwin-x64@1.4.5': + optional: true + + '@napi-rs/lzma-freebsd-x64@1.4.5': + optional: true + + '@napi-rs/lzma-linux-arm-gnueabihf@1.4.5': + optional: true + + '@napi-rs/lzma-linux-arm64-gnu@1.4.5': + optional: true + + '@napi-rs/lzma-linux-arm64-musl@1.4.5': + optional: true + + '@napi-rs/lzma-linux-ppc64-gnu@1.4.5': + optional: true + + '@napi-rs/lzma-linux-riscv64-gnu@1.4.5': + optional: true + + '@napi-rs/lzma-linux-s390x-gnu@1.4.5': + optional: true + + '@napi-rs/lzma-linux-x64-gnu@1.4.5': + optional: true + + '@napi-rs/lzma-linux-x64-musl@1.4.5': + optional: true + + '@napi-rs/lzma-wasm32-wasi@1.4.5': + dependencies: + '@napi-rs/wasm-runtime': 1.1.1 + optional: true + + '@napi-rs/lzma-win32-arm64-msvc@1.4.5': + optional: true + + '@napi-rs/lzma-win32-ia32-msvc@1.4.5': + optional: true + + '@napi-rs/lzma-win32-x64-msvc@1.4.5': + optional: true + + '@napi-rs/lzma@1.4.5': + optionalDependencies: + '@napi-rs/lzma-android-arm-eabi': 1.4.5 + '@napi-rs/lzma-android-arm64': 1.4.5 + '@napi-rs/lzma-darwin-arm64': 1.4.5 + '@napi-rs/lzma-darwin-x64': 1.4.5 + '@napi-rs/lzma-freebsd-x64': 1.4.5 + '@napi-rs/lzma-linux-arm-gnueabihf': 1.4.5 + '@napi-rs/lzma-linux-arm64-gnu': 1.4.5 + '@napi-rs/lzma-linux-arm64-musl': 1.4.5 + '@napi-rs/lzma-linux-ppc64-gnu': 1.4.5 + '@napi-rs/lzma-linux-riscv64-gnu': 1.4.5 + '@napi-rs/lzma-linux-s390x-gnu': 1.4.5 + '@napi-rs/lzma-linux-x64-gnu': 1.4.5 + '@napi-rs/lzma-linux-x64-musl': 1.4.5 + '@napi-rs/lzma-wasm32-wasi': 1.4.5 + '@napi-rs/lzma-win32-arm64-msvc': 1.4.5 + '@napi-rs/lzma-win32-ia32-msvc': 1.4.5 + '@napi-rs/lzma-win32-x64-msvc': 1.4.5 + + '@napi-rs/tar-android-arm-eabi@1.1.0': + optional: true + + '@napi-rs/tar-android-arm64@1.1.0': + optional: true + + '@napi-rs/tar-darwin-arm64@1.1.0': + optional: true + + '@napi-rs/tar-darwin-x64@1.1.0': + optional: true + + '@napi-rs/tar-freebsd-x64@1.1.0': + optional: true + + '@napi-rs/tar-linux-arm-gnueabihf@1.1.0': + optional: true + + '@napi-rs/tar-linux-arm64-gnu@1.1.0': + optional: true + + '@napi-rs/tar-linux-arm64-musl@1.1.0': + optional: true + + '@napi-rs/tar-linux-ppc64-gnu@1.1.0': + optional: true + + '@napi-rs/tar-linux-s390x-gnu@1.1.0': + optional: true + + '@napi-rs/tar-linux-x64-gnu@1.1.0': + optional: true + + '@napi-rs/tar-linux-x64-musl@1.1.0': + optional: true + + '@napi-rs/tar-wasm32-wasi@1.1.0': + dependencies: + '@napi-rs/wasm-runtime': 1.1.1 + optional: true + + '@napi-rs/tar-win32-arm64-msvc@1.1.0': + optional: true + + '@napi-rs/tar-win32-ia32-msvc@1.1.0': + optional: true + + '@napi-rs/tar-win32-x64-msvc@1.1.0': + optional: true + + '@napi-rs/tar@1.1.0': + optionalDependencies: + '@napi-rs/tar-android-arm-eabi': 1.1.0 + '@napi-rs/tar-android-arm64': 1.1.0 + '@napi-rs/tar-darwin-arm64': 1.1.0 + '@napi-rs/tar-darwin-x64': 1.1.0 + '@napi-rs/tar-freebsd-x64': 1.1.0 + '@napi-rs/tar-linux-arm-gnueabihf': 1.1.0 + '@napi-rs/tar-linux-arm64-gnu': 1.1.0 + '@napi-rs/tar-linux-arm64-musl': 1.1.0 + '@napi-rs/tar-linux-ppc64-gnu': 1.1.0 + '@napi-rs/tar-linux-s390x-gnu': 1.1.0 + '@napi-rs/tar-linux-x64-gnu': 1.1.0 + '@napi-rs/tar-linux-x64-musl': 1.1.0 + '@napi-rs/tar-wasm32-wasi': 1.1.0 + '@napi-rs/tar-win32-arm64-msvc': 1.1.0 + '@napi-rs/tar-win32-ia32-msvc': 1.1.0 + '@napi-rs/tar-win32-x64-msvc': 1.1.0 + + '@napi-rs/wasm-runtime@1.1.1': + dependencies: + '@emnapi/core': 1.8.1 + '@emnapi/runtime': 1.8.1 + '@tybys/wasm-util': 0.10.1 + optional: true + + '@napi-rs/wasm-tools-android-arm-eabi@1.0.1': + optional: true + + '@napi-rs/wasm-tools-android-arm64@1.0.1': + optional: true + + '@napi-rs/wasm-tools-darwin-arm64@1.0.1': + optional: true + + '@napi-rs/wasm-tools-darwin-x64@1.0.1': + optional: true + + '@napi-rs/wasm-tools-freebsd-x64@1.0.1': + optional: true + + '@napi-rs/wasm-tools-linux-arm64-gnu@1.0.1': + optional: true + + '@napi-rs/wasm-tools-linux-arm64-musl@1.0.1': + optional: true + + '@napi-rs/wasm-tools-linux-x64-gnu@1.0.1': + optional: true + + '@napi-rs/wasm-tools-linux-x64-musl@1.0.1': + optional: true + + '@napi-rs/wasm-tools-wasm32-wasi@1.0.1': + dependencies: + '@napi-rs/wasm-runtime': 1.1.1 + optional: true + + '@napi-rs/wasm-tools-win32-arm64-msvc@1.0.1': + optional: true + + '@napi-rs/wasm-tools-win32-ia32-msvc@1.0.1': + optional: true + + '@napi-rs/wasm-tools-win32-x64-msvc@1.0.1': + optional: true + + '@napi-rs/wasm-tools@1.0.1': + optionalDependencies: + '@napi-rs/wasm-tools-android-arm-eabi': 1.0.1 + '@napi-rs/wasm-tools-android-arm64': 1.0.1 + '@napi-rs/wasm-tools-darwin-arm64': 1.0.1 + '@napi-rs/wasm-tools-darwin-x64': 1.0.1 + '@napi-rs/wasm-tools-freebsd-x64': 1.0.1 + '@napi-rs/wasm-tools-linux-arm64-gnu': 1.0.1 + '@napi-rs/wasm-tools-linux-arm64-musl': 1.0.1 + '@napi-rs/wasm-tools-linux-x64-gnu': 1.0.1 + '@napi-rs/wasm-tools-linux-x64-musl': 1.0.1 + '@napi-rs/wasm-tools-wasm32-wasi': 1.0.1 + '@napi-rs/wasm-tools-win32-arm64-msvc': 1.0.1 + '@napi-rs/wasm-tools-win32-ia32-msvc': 1.0.1 + '@napi-rs/wasm-tools-win32-x64-msvc': 1.0.1 '@neoconfetti/react@1.0.0': {} @@ -5699,6 +6675,68 @@ snapshots: '@nodelib/fs.scandir': 2.1.5 fastq: 1.19.1 + '@octokit/auth-token@6.0.0': {} + + '@octokit/core@7.0.6': + dependencies: + '@octokit/auth-token': 6.0.0 + '@octokit/graphql': 9.0.3 + '@octokit/request': 10.0.7 + '@octokit/request-error': 7.1.0 + '@octokit/types': 16.0.0 + before-after-hook: 4.0.0 + universal-user-agent: 7.0.3 + + '@octokit/endpoint@11.0.2': + dependencies: + '@octokit/types': 16.0.0 + universal-user-agent: 7.0.3 + + '@octokit/graphql@9.0.3': + dependencies: + '@octokit/request': 10.0.7 + '@octokit/types': 16.0.0 + universal-user-agent: 7.0.3 + + '@octokit/openapi-types@27.0.0': {} + + '@octokit/plugin-paginate-rest@14.0.0(@octokit/core@7.0.6)': + dependencies: + '@octokit/core': 7.0.6 + '@octokit/types': 16.0.0 + + '@octokit/plugin-request-log@6.0.0(@octokit/core@7.0.6)': + dependencies: + '@octokit/core': 7.0.6 + + '@octokit/plugin-rest-endpoint-methods@17.0.0(@octokit/core@7.0.6)': + dependencies: + '@octokit/core': 7.0.6 + '@octokit/types': 16.0.0 + + '@octokit/request-error@7.1.0': + dependencies: + '@octokit/types': 16.0.0 + + '@octokit/request@10.0.7': + dependencies: + '@octokit/endpoint': 11.0.2 + '@octokit/request-error': 7.1.0 + '@octokit/types': 16.0.0 + fast-content-type-parse: 3.0.0 + universal-user-agent: 7.0.3 + + '@octokit/rest@22.0.1': + dependencies: + '@octokit/core': 7.0.6 + '@octokit/plugin-paginate-rest': 14.0.0(@octokit/core@7.0.6) + '@octokit/plugin-request-log': 6.0.0(@octokit/core@7.0.6) + '@octokit/plugin-rest-endpoint-methods': 17.0.0(@octokit/core@7.0.6) + + '@octokit/types@16.0.0': + dependencies: + '@octokit/openapi-types': 27.0.0 + '@playwright/test@1.56.1': dependencies: playwright: 1.56.1 @@ -6880,6 +7918,11 @@ snapshots: dependencies: '@testing-library/dom': 10.4.0 + '@tybys/wasm-util@0.10.1': + dependencies: + tslib: 2.8.1 + optional: true + '@types/aria-query@5.0.4': {} '@types/babel__core@7.20.5': @@ -7461,6 +8504,8 @@ snapshots: ansi-styles@5.2.0: {} + ansi-styles@6.2.3: {} + arg@5.0.2: {} argparse@2.0.1: {} @@ -7584,6 +8629,8 @@ snapshots: dependencies: safe-buffer: 5.1.2 + before-after-hook@4.0.0: {} + body-parser@1.20.3: dependencies: bytes: 3.1.2 @@ -7692,6 +8739,8 @@ snapshots: ansi-styles: 4.3.0 supports-color: 7.2.0 + chardet@2.1.1: {} + check-error@2.1.1: {} chokidar@4.0.3: @@ -7704,6 +8753,12 @@ snapshots: dependencies: clsx: 2.1.1 + cli-width@4.1.0: {} + + clipanion@4.0.0-rc.4(typanion@3.14.0): + dependencies: + typanion: 3.14.0 + clsx@2.1.1: {} cmdk@1.1.1(@types/react-dom@19.2.3(@types/react@19.2.7))(@types/react@19.2.7)(react-dom@19.2.1(react@19.2.1))(react@19.2.1): @@ -7734,6 +8789,8 @@ snapshots: color-name@1.1.4: {} + colorette@2.0.20: {} + combined-stream@1.0.8: dependencies: delayed-stream: 1.0.0 @@ -7935,6 +8992,10 @@ snapshots: electron-to-chromium@1.5.157: {} + emnapi@1.8.1: {} + + emoji-regex@10.6.0: {} + empathic@2.0.0: {} encodeurl@1.0.2: {} @@ -8054,6 +9115,8 @@ snapshots: is-date-object: 1.1.0 is-symbol: 1.1.1 + es-toolkit@1.44.0: {} + esbuild@0.25.10: optionalDependencies: '@esbuild/aix-ppc64': 0.25.10 @@ -8275,6 +9338,8 @@ snapshots: exsolve@1.0.8: {} + fast-content-type-parse@3.0.0: {} + fast-deep-equal@3.1.3: {} fast-equals@5.2.2: {} @@ -8376,6 +9441,8 @@ snapshots: gensync@1.0.0-beta.2: {} + get-east-asian-width@1.4.0: {} + get-intrinsic@1.3.0: dependencies: call-bind-apply-helpers: 1.0.2 @@ -8475,6 +9542,10 @@ snapshots: dependencies: safer-buffer: 2.1.2 + iconv-lite@0.7.2: + dependencies: + safer-buffer: 2.1.2 + ignore@5.3.2: {} ignore@7.0.5: {} @@ -8866,6 +9937,8 @@ snapshots: ms@2.1.3: {} + mute-stream@3.0.0: {} + nanoid@3.3.11: {} natural-compare@1.4.0: {} @@ -8912,6 +9985,8 @@ snapshots: define-properties: 1.2.1 es-object-atoms: 1.1.1 + obug@2.1.1: {} + on-finished@2.3.0: dependencies: ee-first: 1.1.1 @@ -9483,6 +10558,8 @@ snapshots: siginfo@2.0.0: {} + signal-exit@4.1.0: {} + sirv@3.0.2: dependencies: '@polka/url': 1.0.0-next.29 @@ -9545,6 +10622,12 @@ snapshots: - react-dom - utf-8-validate + string-width@7.2.0: + dependencies: + emoji-regex: 10.6.0 + get-east-asian-width: 1.4.0 + strip-ansi: 7.1.0 + string.prototype.matchall@4.0.12: dependencies: call-bind: 1.0.8 @@ -9679,6 +10762,8 @@ snapshots: optionalDependencies: fsevents: 2.3.3 + typanion@3.14.0: {} + type-check@0.4.0: dependencies: prelude-ls: 1.2.1 @@ -9768,6 +10853,8 @@ snapshots: undici-types@7.16.0: {} + universal-user-agent@7.0.3: {} + universalify@2.0.1: {} unpipe@1.0.0: {} @@ -10042,6 +11129,12 @@ snapshots: word-wrap@1.2.5: {} + wrap-ansi@9.0.2: + dependencies: + ansi-styles: 6.2.3 + string-width: 7.2.0 + strip-ansi: 7.1.0 + ws@8.18.3: {} wsl-utils@0.1.0: diff --git a/ui/Dockerfile b/ui/Dockerfile index b741987483b..edb04e68380 100644 --- a/ui/Dockerfile +++ b/ui/Dockerfile @@ -46,7 +46,10 @@ RUN pnpm install --frozen-lockfile --prod FROM rust:1.88.0 AS tensorzero-node-build-env ARG DEBUG_BUILD -RUN apt-get update && apt-get install -y nodejs npm && rm -rf /var/lib/apt/lists/* +# Install Node.js 24 (required by @napi-rs/cli v3) +RUN curl -fsSL https://deb.nodesource.com/setup_24.x | bash - && \ + apt-get install -y nodejs && \ + rm -rf /var/lib/apt/lists/* RUN npm install -g pnpm COPY . /build From d95942f2b96fcbc32145cd17662788cbee354b70 Mon Sep 17 00:00:00 2001 From: Andrew Jesson <56548831+anndvision@users.noreply.github.com> Date: Mon, 26 Jan 2026 23:10:58 -0500 Subject: [PATCH 18/46] Add round-trip tests for Python SDK content block types (#5839) * add roundtrip tests for get datapoints / inferences for tool calls * add datapoint tool result validation * Rename test_type_roundtrip.py to test_roundtrip_tool_types.py * Add text content type round-trip tests following inference pattern * Add template content type round-trip tests * Add thought content type round-trip tests using dummy reasoner model * Add None assertions to fix pyright type errors - Add assertions for result.content, input_messages, message.content - Apply pattern across all 4 test files (tool, text, template, thought) - Reduces pyright errors from 81 to 40 * Add type narrowing to fix remaining pyright errors - Add isinstance checks for ChatInferenceResponse/JsonInferenceResponse before accessing inference_id - Add output is not None assertions for stored inferences - Import JsonInferenceResponse in template types - Move isinstance checks before inference_id access for proper type narrowing - Reduces pyright errors from 40 to 0 * Add sleep delays for batch writes in roundtrip tests When batch_writes is enabled, ClickHouse writes are async and not immediately available for retrieval. Add 1-second sleep delays after inference calls to wait for writes to complete before querying. Fixes CI failures where get_inferences returned empty results. * align sync tests * wait after follow-up * wrap all asserts in try --- .../tests/test_roundtrip_template_types.py | 433 +++++++++ .../python/tests/test_roundtrip_text_types.py | 550 +++++++++++ .../tests/test_roundtrip_thought_types.py | 497 ++++++++++ .../python/tests/test_roundtrip_tool_types.py | 895 ++++++++++++++++++ 4 files changed, 2375 insertions(+) create mode 100644 clients/python/tests/test_roundtrip_template_types.py create mode 100644 clients/python/tests/test_roundtrip_text_types.py create mode 100644 clients/python/tests/test_roundtrip_thought_types.py create mode 100644 clients/python/tests/test_roundtrip_tool_types.py diff --git a/clients/python/tests/test_roundtrip_template_types.py b/clients/python/tests/test_roundtrip_template_types.py new file mode 100644 index 00000000000..62b7c3a3c3f --- /dev/null +++ b/clients/python/tests/test_roundtrip_template_types.py @@ -0,0 +1,433 @@ +""" +Test round-trip serialization and reuse of template content types. + +Verifies that template content preserves all fields through: +- Datapoint creation with templates +- Storage types (StoredInputMessageContentTemplate) +- Serialization (asdict, JSON) +- Reuse in new datapoints + +Template content only appears in inputs (not inference responses), so this test +focuses on the input/storage/datapoint lifecycle. + +This catches type generation issues for template content types. +""" + +import asyncio +import json +import time +from dataclasses import asdict + +import pytest +from tensorzero import ( + AsyncTensorZeroGateway, + CreateDatapointsFromInferenceRequestParamsInferenceIds, + JsonInferenceResponse, + TensorZeroGateway, +) +from tensorzero.generated_types import ( + DatapointJson, + InputMessageContentTemplate, + JsonInferenceOutput, + StoredInferenceJson, + StoredInputMessageContentTemplate, +) +from uuid_utils import uuid7 + + +@pytest.mark.asyncio +async def test_async_template_content_roundtrip_complete_flow( + async_client: AsyncTensorZeroGateway, +): + """ + Comprehensive test verifying template content preserves all fields through: + 1. Inference with template input + 2. Storage retrieval (StoredInputMessageContentTemplate) + 3. Serialization (asdict + JSON) + 4. Reuse in follow-up inference + 5. Datapoint creation from inference + 6. Datapoint retrieval and validation (InputMessageContentTemplate) + + This test validates 2 types across the storage/datapoint lifecycle. + Templates only appear in inputs, not in inference responses. + """ + + # ============================================================================ + # Step 1: Create inference with template input + # ============================================================================ + + # json_success function has user_template configured + result = await async_client.inference( + function_name="json_success", + input={ + "system": {"assistant_name": "Test Assistant"}, + "messages": [ + { + "role": "user", + "content": [{"type": "template", "name": "user", "arguments": {"country": "France"}}], + } + ], + }, + stream=False, + ) + + # For JSON functions, we just need the inference_id + assert isinstance(result, JsonInferenceResponse), "Result must be JsonInferenceResponse instance" + assert result.inference_id is not None, "Inference must return ID" + inference_id = str(result.inference_id) + + # Wait for results to be written to ClickHouse (required for batch writes) + await asyncio.sleep(1) + + # ============================================================================ + # Step 2: Query inference back via get_inferences + # ============================================================================ + + get_response = await async_client.get_inferences( + ids=[inference_id], + function_name="json_success", + output_source="inference", + ) + + assert get_response.inferences is not None, "get_inferences must return inferences" + assert len(get_response.inferences) == 1, "Should retrieve exactly 1 inference" + + stored_inference = get_response.inferences[0] + assert isinstance(stored_inference, StoredInferenceJson), "Must be StoredInferenceJson instance" + assert str(stored_inference.inference_id) == inference_id, "Retrieved inference ID must match original" + + # ============================================================================ + # Step 3: Verify stored input Template (StoredInputMessageContentTemplate) + # ============================================================================ + + input_messages = stored_inference.input.messages + assert input_messages is not None, "Input messages must not be None" + assert len(input_messages) >= 1, "Should have at least 1 input message" + + # Find user message with template + user_msg = None + for msg in input_messages: + if msg.role == "user": + user_msg = msg + break + + assert user_msg is not None, "Should have user message in stored input" + assert user_msg.content is not None, "User message content must not be None" + assert len(user_msg.content) > 0, "User message should have content" + + # Verify the template was stored correctly (StoredInputMessageContentTemplate) + template_content = None + for content in user_msg.content: + if content.type == "template": + template_content = content + break + + assert template_content is not None, "Should have template in user message" + assert isinstance(template_content, StoredInputMessageContentTemplate), ( + "Must be StoredInputMessageContentTemplate instance" + ) + + # Field existence + assert hasattr(template_content, "name"), "StoredInputMessageContentTemplate must have 'name' field" + assert hasattr(template_content, "arguments"), "StoredInputMessageContentTemplate must have 'arguments' field" + + # Field values + assert template_content.name == "user", "Template name must be 'user'" + assert template_content.arguments == {"country": "France"}, "Template arguments must be preserved" + + # ============================================================================ + # Step 4: Serialize template content + # ============================================================================ + + template_dict = asdict(template_content) + assert "type" in template_dict and template_dict["type"] == "template" + assert "name" in template_dict and template_dict["name"] == "user" + assert "arguments" in template_dict and template_dict["arguments"] == {"country": "France"} + + # Verify JSON round-trip + template_json = json.dumps(template_dict) + template_reloaded = json.loads(template_json) + assert template_reloaded["name"] == "user" + assert template_reloaded["arguments"] == {"country": "France"} + + # ============================================================================ + # Step 5: Reuse serialized template in follow-up inference + # ============================================================================ + + follow_up_result = await async_client.inference( + function_name="json_success", + input={ + "system": {"assistant_name": "Test Assistant"}, + "messages": [ + { + "role": "user", + "content": [template_dict], # Reuse serialized template + } + ], + }, + stream=False, + ) + + assert isinstance(follow_up_result, JsonInferenceResponse), "Follow-up must return JsonInferenceResponse" + assert follow_up_result.inference_id is not None, ( + "Follow-up inference must succeed when reusing serialized template data" + ) + + follow_up_id = str(follow_up_result.inference_id) + + # Wait for follow-up results to be written to ClickHouse (required for batch writes) + await asyncio.sleep(1) + + # ============================================================================ + # Step 6: Verify follow-up stored data + # ============================================================================ + + follow_up_stored = await async_client.get_inferences( + ids=[follow_up_id], + function_name="json_success", + output_source="inference", + ) + + assert len(follow_up_stored.inferences) == 1, "Should retrieve exactly 1 follow-up inference" + follow_up_inference = follow_up_stored.inferences[0] + + # Verify template in follow-up input + follow_up_messages = follow_up_inference.input.messages + assert follow_up_messages is not None, "Follow-up messages must not be None" + + follow_up_user_msg = None + for msg in follow_up_messages: + if msg.role == "user": + follow_up_user_msg = msg + break + + assert follow_up_user_msg is not None, "Should have user message in follow-up" + assert follow_up_user_msg.content is not None, "Follow-up user message content must not be None" + + follow_up_template = None + for content in follow_up_user_msg.content: + if content.type == "template": + follow_up_template = content + break + + assert follow_up_template is not None, "Should have template in follow-up" + assert isinstance(follow_up_template, StoredInputMessageContentTemplate) + assert follow_up_template.name == "user", "Template name must be preserved" + assert follow_up_template.arguments == {"country": "France"}, "Template arguments must match reused serialized data" + + # ============================================================================ + # Step 7: Create datapoint from follow-up inference + # ============================================================================ + + dataset_name = f"test_template_roundtrip_{uuid7()}" + + try: + datapoint_response = await async_client.create_datapoints_from_inferences( + dataset_name=dataset_name, + params=CreateDatapointsFromInferenceRequestParamsInferenceIds(inference_ids=[follow_up_id]), + output_source="inference", + ) + + assert datapoint_response.ids is not None, "create_datapoints_from_inferences must return IDs" + assert len(datapoint_response.ids) == 1, "Should create exactly 1 datapoint" + + datapoint_id = datapoint_response.ids[0] + + # ============================================================================ + # Step 8: Retrieve datapoint and verify InputMessageContentTemplate + # ============================================================================ + datapoint_get_response = await async_client.get_datapoints(dataset_name=dataset_name, ids=[datapoint_id]) + + assert datapoint_get_response.datapoints is not None, "get_datapoints must return datapoints" + assert len(datapoint_get_response.datapoints) == 1, "Should retrieve exactly 1 datapoint" + + datapoint = datapoint_get_response.datapoints[0] + assert isinstance(datapoint, DatapointJson), "Must be DatapointJson instance" + assert datapoint.id == datapoint_id, "Datapoint ID must match" + + # Verify template in datapoint input + datapoint_messages = datapoint.input.messages + assert datapoint_messages is not None, "Datapoint messages must not be None" + assert len(datapoint_messages) >= 1, "Datapoint should have at least 1 message" + + datapoint_user_msg = None + for msg in datapoint_messages: + if msg.role == "user": + datapoint_user_msg = msg + break + + assert datapoint_user_msg is not None, "Datapoint should have user message" + assert datapoint_user_msg.content is not None, "Datapoint user message content must not be None" + + datapoint_template = None + for content in datapoint_user_msg.content: + if content.type == "template": + datapoint_template = content + break + + assert datapoint_template is not None, "Datapoint should have template in user message" + assert isinstance(datapoint_template, InputMessageContentTemplate), ( + "Datapoint template must be InputMessageContentTemplate instance" + ) + + # Verify fields preserved in datapoint + assert hasattr(datapoint_template, "name"), "Datapoint template must have 'name' field" + assert hasattr(datapoint_template, "arguments"), "Datapoint template must have 'arguments' field" + assert datapoint_template.name == "user", "Datapoint template name must be preserved" + assert datapoint_template.arguments == {"country": "France"}, ( + "Datapoint template arguments must match reused serialized data" + ) + + # Verify output is JsonInferenceOutput + assert isinstance(datapoint.output, JsonInferenceOutput), "Datapoint output must be JsonInferenceOutput" + + finally: + await async_client.delete_dataset(dataset_name=dataset_name) + + +def test_sync_template_content_roundtrip_complete_flow(sync_client: TensorZeroGateway): + """ + Sync version of test_async_template_content_roundtrip_complete_flow. + Tests the same round-trip but with synchronous client. + """ + + # ============================================================================ + # Step 1: Create inference with template + # ============================================================================ + + result = sync_client.inference( + function_name="json_success", + input={ + "system": {"assistant_name": "Test Assistant"}, + "messages": [ + { + "role": "user", + "content": [{"type": "template", "name": "user", "arguments": {"country": "Germany"}}], + } + ], + }, + stream=False, + ) + + assert isinstance(result, JsonInferenceResponse), "Result must be JsonInferenceResponse" + assert result.inference_id is not None, "Result must have inference_id" + + inference_id = str(result.inference_id) + + # Wait for results to be written to ClickHouse (required for batch writes) + time.sleep(1) + + # ============================================================================ + # Step 2: Query via get_inferences + # ============================================================================ + + get_response = sync_client.get_inferences( + ids=[inference_id], + function_name="json_success", + output_source="inference", + ) + + stored_inference = get_response.inferences[0] + assert isinstance(stored_inference, StoredInferenceJson) + + # ============================================================================ + # Step 3: Verify StoredInputMessageContentTemplate + # ============================================================================ + + input_messages_sync = stored_inference.input.messages + assert input_messages_sync is not None, "Input messages must not be None" + + user_msg = None + for msg in input_messages_sync: + if msg.role == "user": + user_msg = msg + break + + assert user_msg is not None, "User message must not be None" + assert user_msg.content is not None, "User message content must not be None" + + template_content = None + for content in user_msg.content: + if content.type == "template": + template_content = content + break + + assert template_content is not None + assert isinstance(template_content, StoredInputMessageContentTemplate) + assert template_content.name == "user" + assert template_content.arguments == {"country": "Germany"} + + # ============================================================================ + # Step 4: Serialize and create datapoint + # ============================================================================ + + template_dict = asdict(template_content) + + # Reuse in another inference + follow_up_result = sync_client.inference( + function_name="json_success", + input={ + "system": {"assistant_name": "Test Assistant"}, + "messages": [ + { + "role": "user", + "content": [template_dict], + } + ], + }, + stream=False, + ) + + assert isinstance(follow_up_result, JsonInferenceResponse), "Follow-up must return JsonInferenceResponse" + assert follow_up_result.inference_id is not None, "Follow-up must have inference_id" + + follow_up_id = str(follow_up_result.inference_id) + + # Wait for follow-up results to be written to ClickHouse (required for batch writes) + time.sleep(1) + + dataset_name = f"test_template_roundtrip_{uuid7()}" + + try: + datapoint_response = sync_client.create_datapoints_from_inferences( + dataset_name=dataset_name, + params=CreateDatapointsFromInferenceRequestParamsInferenceIds(inference_ids=[follow_up_id]), + output_source="inference", + ) + + datapoint_id = datapoint_response.ids[0] + + # ============================================================================ + # Step 5: Verify InputMessageContentTemplate in datapoint + # ============================================================================ + datapoint_get_response = sync_client.get_datapoints(dataset_name=dataset_name, ids=[datapoint_id]) + + datapoint = datapoint_get_response.datapoints[0] + assert isinstance(datapoint, DatapointJson) + + # Find template in datapoint + datapoint_messages_sync = datapoint.input.messages + assert datapoint_messages_sync is not None, "Datapoint messages must not be None" + + datapoint_user_msg = None + for msg in datapoint_messages_sync: + if msg.role == "user": + datapoint_user_msg = msg + break + + assert datapoint_user_msg is not None, "Datapoint user message must not be None" + assert datapoint_user_msg.content is not None, "Datapoint user message content must not be None" + + datapoint_template = None + for content in datapoint_user_msg.content: + if content.type == "template": + datapoint_template = content + break + + assert datapoint_template is not None + assert isinstance(datapoint_template, InputMessageContentTemplate) + assert datapoint_template.name == "user" + assert datapoint_template.arguments == {"country": "Germany"} + + finally: + sync_client.delete_dataset(dataset_name=dataset_name) diff --git a/clients/python/tests/test_roundtrip_text_types.py b/clients/python/tests/test_roundtrip_text_types.py new file mode 100644 index 00000000000..10377603255 --- /dev/null +++ b/clients/python/tests/test_roundtrip_text_types.py @@ -0,0 +1,550 @@ +""" +Test round-trip serialization and reuse of text content types. + +Verifies that text and raw_text content preserves all fields through: +- Inference response types (Text, RawText) +- Storage types (ContentBlockChatOutputText, StoredInputMessageContentText, StoredInputMessageContentRawText) +- Serialization (asdict, JSON) +- Reuse in follow-up inferences +- Datapoint creation and retrieval + +This catches type generation issues for text content types. +""" + +import asyncio +import json +import time +from dataclasses import asdict + +import pytest +from tensorzero import ( + AsyncTensorZeroGateway, + ChatInferenceResponse, + CreateDatapointsFromInferenceRequestParamsInferenceIds, + TensorZeroGateway, + Text, +) +from tensorzero.generated_types import ( + ContentBlockChatOutputText, + DatapointChat, + InputMessageContentRawText, + InputMessageContentText, + StoredInferenceChat, + StoredInputMessageContentRawText, + StoredInputMessageContentText, +) +from uuid_utils import uuid7 + + +@pytest.mark.asyncio +async def test_async_text_content_roundtrip_complete_flow( + async_client: AsyncTensorZeroGateway, +): + """ + Comprehensive test verifying text content preserves all fields through: + 1. Inference response (Text type) + 2. Storage retrieval (ContentBlockChatOutputText, StoredInputMessageContentText types) + 3. Serialization (asdict + JSON) + 4. Reuse in follow-up inference + 5. Datapoint creation from inference + 6. Datapoint retrieval and validation + + This test validates 7 types across the complete lifecycle. + Tests 12 steps total (8 inference + 4 datapoint). + """ + + # ============================================================================ + # Step 1: Create initial inference with text content + # ============================================================================ + + result = await async_client.inference( + function_name="basic_test", + input={ + "system": {"assistant_name": "Test Assistant"}, + "messages": [{"role": "user", "content": "What's your name?"}], + }, + stream=False, + ) + + # Basic result assertions + assert isinstance(result, ChatInferenceResponse), "Result must be ChatInferenceResponse instance" + assert result.content is not None, "Result content must not be None" + assert len(result.content) >= 1, "Result should have at least 1 content block" + assert result.inference_id is not None, "Result must have inference_id" + + # ============================================================================ + # Step 2: Verify response Text fields (types.py Text) + # ============================================================================ + + # Find text content in response + text_response = None + for content in result.content: + if content.type == "text": + text_response = content + break + + assert text_response is not None, "Response must contain text content" + assert isinstance(text_response, Text), "Response must be Text instance" + assert hasattr(text_response, "text"), "Text must have 'text' field" + assert text_response.text is not None, "Text field must not be None" + + original_text = text_response.text + + # Store inference_id for retrieval + inference_id = str(result.inference_id) + + # Wait for results to be written to ClickHouse (required for batch writes) + await asyncio.sleep(1) + + # ============================================================================ + # Step 3: Query inference back via get_inferences + # ============================================================================ + + get_response = await async_client.get_inferences( + ids=[inference_id], + function_name="basic_test", + output_source="inference", + ) + + assert get_response.inferences is not None, "get_inferences must return inferences" + assert len(get_response.inferences) == 1, "Should retrieve exactly 1 inference" + + stored_inference = get_response.inferences[0] + assert isinstance(stored_inference, StoredInferenceChat), "Must be StoredInferenceChat instance" + assert stored_inference.output is not None, "Stored inference output must not be None" + assert str(stored_inference.inference_id) == inference_id, "Retrieved inference ID must match original" + + # ============================================================================ + # Step 4: Verify stored output Text (ContentBlockChatOutputText) + # ============================================================================ + + stored_output = stored_inference.output + assert len(stored_output) >= 1, "Output should have at least 1 content block" + + # Find text in output + stored_text = None + for content in stored_output: + if content.type == "text": + stored_text = content + break + + assert stored_text is not None, "Stored output must contain text" + assert stored_text.type == "text", "Stored output must have type='text'" + assert isinstance(stored_text, ContentBlockChatOutputText), ( + "Stored output must be ContentBlockChatOutputText instance" + ) + assert hasattr(stored_text, "text"), "Stored text must have 'text' field" + assert stored_text.text == original_text, "Stored text must match original" + + # ============================================================================ + # Step 5: Serialize text content + # ============================================================================ + + text_dict = asdict(stored_text) + assert "type" in text_dict and text_dict["type"] == "text" + assert "text" in text_dict and text_dict["text"] == original_text + + # Verify JSON round-trip + text_json = json.dumps(text_dict) + text_reloaded = json.loads(text_json) + assert text_reloaded["text"] == original_text + + # ============================================================================ + # Step 6: Create raw_text content for follow-up + # ============================================================================ + + raw_text_content = {"type": "raw_text", "value": "This is raw text content"} + + # ============================================================================ + # Step 7: Reuse in follow-up inference + # ============================================================================ + + follow_up_result = await async_client.inference( + function_name="basic_test", + input={ + "system": {"assistant_name": "Test Assistant"}, + "messages": [ + {"role": "user", "content": "What's your name?"}, + {"role": "assistant", "content": [text_dict]}, # Reuse serialized text + {"role": "user", "content": [raw_text_content]}, # Add raw_text + ], + }, + stream=False, + ) + + assert isinstance(follow_up_result, ChatInferenceResponse), "Follow-up must return ChatInferenceResponse instance" + assert follow_up_result.inference_id is not None, ( + "Follow-up inference must succeed when reusing serialized text data" + ) + + # Wait for follow-up results to be written to ClickHouse (required for batch writes) + await asyncio.sleep(1) + + # ============================================================================ + # Step 8: Verify follow-up stored data + # ============================================================================ + + follow_up_id = str(follow_up_result.inference_id) + follow_up_stored = await async_client.get_inferences( + ids=[follow_up_id], + function_name="basic_test", + output_source="inference", + ) + + assert len(follow_up_stored.inferences) == 1, "Should retrieve exactly 1 follow-up inference" + follow_up_inference = follow_up_stored.inferences[0] + + # Verify input contains our text and raw_text + input_messages = follow_up_inference.input.messages + assert input_messages is not None, "Input messages must not be None" + assert len(input_messages) >= 3, "Should have user, assistant, and follow-up user messages" + + # Find assistant message with text + assistant_msg = None + for msg in input_messages: + if msg.role == "assistant": + assistant_msg = msg + break + + assert assistant_msg is not None, "Should have assistant message in stored input" + assert assistant_msg.content is not None, "Assistant message content must not be None" + assert len(assistant_msg.content) > 0, "Assistant message should have content" + + # Verify the text was stored correctly (StoredInputMessageContentText) + text_content = None + for content in assistant_msg.content: + if content.type == "text": + text_content = content + break + + assert text_content is not None, "Should have text in assistant message" + assert isinstance(text_content, StoredInputMessageContentText), "Must be StoredInputMessageContentText instance" + assert hasattr(text_content, "text"), "StoredInputMessageContentText must have 'text' field" + assert text_content.text == original_text, "Text must be preserved through round-trip" + + # Verify raw_text was also stored correctly + user_msg_with_raw_text = None + for msg in input_messages: + if msg.role == "user": + for content in msg.content: + if content.type == "raw_text": + user_msg_with_raw_text = content + break + if user_msg_with_raw_text: + break + + assert user_msg_with_raw_text is not None, "Should have raw_text in stored input" + assert isinstance(user_msg_with_raw_text, StoredInputMessageContentRawText), ( + "Must be StoredInputMessageContentRawText instance" + ) + assert user_msg_with_raw_text.value == "This is raw text content", "Raw text must be preserved" + + # ============================================================================ + # Step 9: Create datapoint from follow-up inference + # ============================================================================ + + dataset_name = f"test_text_roundtrip_{uuid7()}" + + try: + datapoint_response = await async_client.create_datapoints_from_inferences( + dataset_name=dataset_name, + params=CreateDatapointsFromInferenceRequestParamsInferenceIds(inference_ids=[follow_up_id]), + output_source="inference", + ) + + assert datapoint_response.ids is not None, "create_datapoints_from_inferences must return IDs" + assert len(datapoint_response.ids) == 1, "Should create exactly 1 datapoint" + + datapoint_id = datapoint_response.ids[0] + + # ============================================================================ + # Step 10: Retrieve datapoint via get_datapoints + # ============================================================================ + datapoint_get_response = await async_client.get_datapoints(dataset_name=dataset_name, ids=[datapoint_id]) + + assert datapoint_get_response.datapoints is not None, "get_datapoints must return datapoints" + assert len(datapoint_get_response.datapoints) == 1, "Should retrieve exactly 1 datapoint" + + datapoint = datapoint_get_response.datapoints[0] + assert isinstance(datapoint, DatapointChat), "Must be DatapointChat instance" + assert datapoint.id == datapoint_id, "Datapoint ID must match" + + # ============================================================================ + # Step 11: Verify text content in datapoint input + # ============================================================================ + + datapoint_input_messages = datapoint.input.messages + assert datapoint_input_messages is not None, "Datapoint input messages must not be None" + assert len(datapoint_input_messages) >= 3, "Datapoint should have user, assistant, and follow-up messages" + + # Find assistant message with text in datapoint input + datapoint_assistant_msg = None + for msg in datapoint_input_messages: + if msg.role == "assistant": + datapoint_assistant_msg = msg + break + + assert datapoint_assistant_msg is not None, "Datapoint should have assistant message" + assert datapoint_assistant_msg.content is not None, "Datapoint assistant message content must not be None" + assert len(datapoint_assistant_msg.content) > 0, "Datapoint assistant message should have content" + + # Verify text in datapoint input + datapoint_text = None + for content in datapoint_assistant_msg.content: + if content.type == "text": + datapoint_text = content + break + + assert datapoint_text is not None, "Datapoint should have text in assistant message" + assert isinstance(datapoint_text, InputMessageContentText), ( + "Datapoint text must be InputMessageContentText instance" + ) + assert hasattr(datapoint_text, "text"), "Datapoint text must have 'text' field" + assert datapoint_text.text == original_text, "Datapoint text must be preserved" + + # Verify raw_text in datapoint input + datapoint_raw_text = None + for msg in datapoint_input_messages: + if msg.role == "user": + for content in msg.content: + if content.type == "raw_text": + datapoint_raw_text = content + break + if datapoint_raw_text: + break + + assert datapoint_raw_text is not None, "Datapoint should have raw_text in input" + assert isinstance(datapoint_raw_text, InputMessageContentRawText), ( + "Datapoint raw_text must be InputMessageContentRawText instance" + ) + assert hasattr(datapoint_raw_text, "value"), "Datapoint raw_text must have 'value' field" + assert datapoint_raw_text.value == "This is raw text content", "Datapoint raw_text must be preserved" + + # ============================================================================ + # Step 12: Verify text in datapoint output + # ============================================================================ + + datapoint_output = datapoint.output + assert datapoint_output is not None, "Datapoint must have output" + assert len(datapoint_output) >= 1, "Datapoint output should have at least 1 content block" + + # Find text in output + datapoint_output_text = None + for content in datapoint_output: + if content.type == "text": + datapoint_output_text = content + break + + assert datapoint_output_text is not None, "Datapoint output must contain text" + assert datapoint_output_text.type == "text", "Datapoint output must have type='text'" + assert isinstance(datapoint_output_text, ContentBlockChatOutputText), ( + "Datapoint output text must be ContentBlockChatOutputText instance" + ) + assert hasattr(datapoint_output_text, "text"), "Datapoint output text must have 'text' field" + + finally: + await async_client.delete_dataset(dataset_name=dataset_name) + + +def test_sync_text_content_roundtrip_complete_flow(sync_client: TensorZeroGateway): + """ + Sync version of test_async_text_content_roundtrip_complete_flow. + Tests the same round-trip but with synchronous client. + """ + + # ============================================================================ + # Step 1: Create initial inference + # ============================================================================ + + result = sync_client.inference( + function_name="basic_test", + input={ + "system": {"assistant_name": "Test Assistant"}, + "messages": [{"role": "user", "content": "What's your name?"}], + }, + stream=False, + ) + + assert isinstance(result, ChatInferenceResponse) + assert result.content is not None, "Result content must not be None" + assert result.inference_id is not None, "Result must have inference_id" + + # Find text in response + text_response = None + for content in result.content: + if content.type == "text": + text_response = content + break + + assert text_response is not None + assert isinstance(text_response, Text) + original_text = text_response.text + + inference_id = str(result.inference_id) + + # Wait for results to be written to ClickHouse (required for batch writes) + time.sleep(1) + + # ============================================================================ + # Step 3: Query via get_inferences + # ============================================================================ + + get_response = sync_client.get_inferences( + ids=[inference_id], + function_name="basic_test", + output_source="inference", + ) + + stored_inference = get_response.inferences[0] + assert isinstance(stored_inference, StoredInferenceChat), "Must be StoredInferenceChat instance" + assert stored_inference.output is not None, "Stored inference output must not be None" + + # Find text in stored output + stored_text = None + for content in stored_inference.output: + if content.type == "text": + stored_text = content + break + + assert stored_text is not None + assert isinstance(stored_text, ContentBlockChatOutputText) + + # ============================================================================ + # Step 5: Serialize and reuse + # ============================================================================ + + text_dict = asdict(stored_text) + raw_text_content = {"type": "raw_text", "value": "This is raw text"} + + follow_up_result = sync_client.inference( + function_name="basic_test", + input={ + "system": {"assistant_name": "Test Assistant"}, + "messages": [ + {"role": "user", "content": "Question"}, + {"role": "assistant", "content": [text_dict]}, + {"role": "user", "content": [raw_text_content]}, + ], + }, + stream=False, + ) + + assert isinstance(follow_up_result, ChatInferenceResponse), "Follow-up must return ChatInferenceResponse" + assert follow_up_result.inference_id is not None, "Follow-up must have inference_id" + + follow_up_id = str(follow_up_result.inference_id) + + # Wait for follow-up results to be written to ClickHouse (required for batch writes) + time.sleep(1) + + follow_up_stored = sync_client.get_inferences( + ids=[follow_up_id], + function_name="basic_test", + output_source="inference", + ) + + follow_up_inference = follow_up_stored.inferences[0] + + # Verify StoredInputMessageContentText + input_messages_sync = follow_up_inference.input.messages + assert input_messages_sync is not None, "Input messages must not be None" + + assistant_msg = None + for msg in input_messages_sync: + if msg.role == "assistant": + assistant_msg = msg + break + + assert assistant_msg is not None, "Assistant message must not be None" + assert assistant_msg.content is not None, "Assistant message content must not be None" + + text_content = None + for content in assistant_msg.content: + if content.type == "text": + text_content = content + break + + assert isinstance(text_content, StoredInputMessageContentText) + assert text_content.text == original_text + + # Verify StoredInputMessageContentRawText + raw_text_found = None + for msg in input_messages_sync: + if msg.role == "user": + assert msg.content is not None, "User message content must not be None" + for content in msg.content: + if content.type == "raw_text": + raw_text_found = content + break + if raw_text_found: + break + + assert isinstance(raw_text_found, StoredInputMessageContentRawText) + + # ============================================================================ + # Step 9: Create datapoint + # ============================================================================ + + dataset_name = f"test_text_roundtrip_{uuid7()}" + + try: + datapoint_response = sync_client.create_datapoints_from_inferences( + dataset_name=dataset_name, + params=CreateDatapointsFromInferenceRequestParamsInferenceIds(inference_ids=[follow_up_id]), + output_source="inference", + ) + + datapoint_id = datapoint_response.ids[0] + datapoint_get_response = sync_client.get_datapoints(dataset_name=dataset_name, ids=[datapoint_id]) + + datapoint = datapoint_get_response.datapoints[0] + assert isinstance(datapoint, DatapointChat) + + # Verify InputMessageContentText + datapoint_input_messages_sync = datapoint.input.messages + assert datapoint_input_messages_sync is not None, "Datapoint input messages must not be None" + + datapoint_assistant_msg = None + for msg in datapoint_input_messages_sync: + if msg.role == "assistant": + datapoint_assistant_msg = msg + break + + assert datapoint_assistant_msg is not None, "Datapoint assistant message must not be None" + assert datapoint_assistant_msg.content is not None, "Datapoint assistant message content must not be None" + + datapoint_text = None + for content in datapoint_assistant_msg.content: + if content.type == "text": + datapoint_text = content + break + + assert isinstance(datapoint_text, InputMessageContentText) + assert datapoint_text.text == original_text + + # Verify InputMessageContentRawText + datapoint_raw_text = None + for msg in datapoint_input_messages_sync: + if msg.role == "user": + assert msg.content is not None, "User message content must not be None" + for content in msg.content: + if content.type == "raw_text": + datapoint_raw_text = content + break + if datapoint_raw_text: + break + + assert isinstance(datapoint_raw_text, InputMessageContentRawText) + + # Verify ContentBlockChatOutputText in output + assert datapoint.output is not None, "Datapoint output must not be None" + output_text = None + for content in datapoint.output: + if content.type == "text": + output_text = content + break + + assert isinstance(output_text, ContentBlockChatOutputText) + + finally: + sync_client.delete_dataset(dataset_name=dataset_name) diff --git a/clients/python/tests/test_roundtrip_thought_types.py b/clients/python/tests/test_roundtrip_thought_types.py new file mode 100644 index 00000000000..75dfa72d6b2 --- /dev/null +++ b/clients/python/tests/test_roundtrip_thought_types.py @@ -0,0 +1,497 @@ +""" +Test round-trip serialization and reuse of thought content types. + +Verifies that thought content preserves all fields through: +- Inference response types (Thought) +- Storage types (ContentBlockChatOutputThought, StoredInputMessageContentThought) +- Serialization (asdict, JSON) +- Reuse in follow-up inferences +- Datapoint creation and retrieval + +This catches type generation issues for thought content types. + +Uses the dummy "reasoner" model which returns thought content blocks. +""" + +import asyncio +import json +import time +from dataclasses import asdict + +import pytest +from tensorzero import ( + AsyncTensorZeroGateway, + ChatInferenceResponse, + CreateDatapointsFromInferenceRequestParamsInferenceIds, + TensorZeroGateway, + Thought, +) +from tensorzero.generated_types import ( + ContentBlockChatOutputThought, + DatapointChat, + InputMessageContentThought, + StoredInferenceChat, + StoredInputMessageContentThought, +) +from uuid_utils import uuid7 + + +@pytest.mark.asyncio +async def test_async_thought_content_roundtrip_complete_flow( + async_client: AsyncTensorZeroGateway, +): + """ + Comprehensive test verifying thought content preserves all fields through: + 1. Inference response (Thought type) + 2. Storage retrieval (ContentBlockChatOutputThought, StoredInputMessageContentThought types) + 3. Serialization (asdict + JSON) + 4. Reuse in follow-up inference + 5. Datapoint creation from inference + 6. Datapoint retrieval and validation + + This test validates 4 types across the complete lifecycle. + Tests 12 steps total (8 inference + 4 datapoint). + + Uses the dummy "reasoner" model which returns thought content. + """ + + # ============================================================================ + # Step 1: Create initial inference with reasoner model + # ============================================================================ + + result = await async_client.inference( + function_name="basic_test", + variant_name="reasoner", # Dummy model that returns thought content + input={ + "system": {"assistant_name": "Test Assistant"}, + "messages": [{"role": "user", "content": "What is 2+2?"}], + }, + stream=False, + ) + + # Basic result assertions + assert isinstance(result, ChatInferenceResponse), "Result must be ChatInferenceResponse instance" + assert result.content is not None, "Result content must not be None" + assert result.inference_id is not None, "Result must have inference_id" + assert len(result.content) >= 1, "Result should have at least 1 content block" + + # ============================================================================ + # Step 2: Verify response Thought fields (types.py Thought) + # ============================================================================ + + # Find thought content in response + thought_response = None + for content in result.content: + if content.type == "thought": + thought_response = content + break + + assert thought_response is not None, "Response must contain thought content" + assert isinstance(thought_response, Thought), "Response must be Thought instance" + assert hasattr(thought_response, "text"), "Thought must have 'text' field" + # Note: text can be None for thoughts with only signature + + original_thought_text = thought_response.text + + # Store inference_id for retrieval + inference_id = str(result.inference_id) + + # Wait for results to be written to ClickHouse (required for batch writes) + await asyncio.sleep(1) + + # ============================================================================ + # Step 3: Query inference back via get_inferences + # ============================================================================ + + get_response = await async_client.get_inferences( + ids=[inference_id], + function_name="basic_test", + output_source="inference", + ) + + assert get_response.inferences is not None, "get_inferences must return inferences" + assert len(get_response.inferences) == 1, "Should retrieve exactly 1 inference" + + stored_inference = get_response.inferences[0] + assert isinstance(stored_inference, StoredInferenceChat), "Must be StoredInferenceChat instance" + assert stored_inference.output is not None, "Stored inference output must not be None" + assert str(stored_inference.inference_id) == inference_id, "Retrieved inference ID must match original" + + # ============================================================================ + # Step 4: Verify stored output Thought (ContentBlockChatOutputThought) + # ============================================================================ + + stored_output = stored_inference.output + assert len(stored_output) >= 1, "Output should have at least 1 content block" + + # Find thought in output + stored_thought = None + for content in stored_output: + if content.type == "thought": + stored_thought = content + break + + assert stored_thought is not None, "Stored output must contain thought" + assert stored_thought.type == "thought", "Stored output must have type='thought'" + assert isinstance(stored_thought, ContentBlockChatOutputThought), ( + "Stored output must be ContentBlockChatOutputThought instance" + ) + assert hasattr(stored_thought, "text"), "Stored thought must have 'text' field" + if original_thought_text is not None: + assert stored_thought.text == original_thought_text, "Stored thought text must match original" + + # ============================================================================ + # Step 5: Serialize thought content + # ============================================================================ + + thought_dict = asdict(stored_thought) + assert "type" in thought_dict and thought_dict["type"] == "thought" + assert "text" in thought_dict + + # Verify JSON round-trip + thought_json = json.dumps(thought_dict) + thought_reloaded = json.loads(thought_json) + assert thought_reloaded["type"] == "thought" + + # ============================================================================ + # Step 6: Reuse in follow-up inference + # ============================================================================ + + follow_up_result = await async_client.inference( + function_name="basic_test", + variant_name="reasoner", + input={ + "system": {"assistant_name": "Test Assistant"}, + "messages": [ + {"role": "user", "content": "What is 2+2?"}, + { + "role": "assistant", + "content": [thought_dict, {"type": "text", "text": "The answer is 4"}], + }, # Reuse serialized thought + text + {"role": "user", "content": "Now what is 3+3?"}, + ], + }, + stream=False, + ) + + assert isinstance(follow_up_result, ChatInferenceResponse), "Follow-up must return ChatInferenceResponse instance" + assert follow_up_result.inference_id is not None, ( + "Follow-up inference must succeed when reusing serialized thought data" + ) + + follow_up_id = str(follow_up_result.inference_id) + + # Wait for follow-up results to be written to ClickHouse (required for batch writes) + await asyncio.sleep(1) + + # ============================================================================ + # Step 7: Verify follow-up stored data + # ============================================================================ + + follow_up_stored = await async_client.get_inferences( + ids=[follow_up_id], + function_name="basic_test", + output_source="inference", + ) + + assert len(follow_up_stored.inferences) == 1, "Should retrieve exactly 1 follow-up inference" + follow_up_inference = follow_up_stored.inferences[0] + + # Verify input contains our thought + input_messages = follow_up_inference.input.messages + assert input_messages is not None, "Input messages must not be None" + assert len(input_messages) >= 3, "Should have user, assistant, and follow-up user messages" + + # Find assistant message with thought + assistant_msg = None + for msg in input_messages: + if msg.role == "assistant": + assistant_msg = msg + break + + assert assistant_msg is not None, "Should have assistant message in stored input" + assert assistant_msg.content is not None, "Assistant message content must not be None" + assert len(assistant_msg.content) > 0, "Assistant message should have content" + + # Verify the thought was stored correctly (StoredInputMessageContentThought) + thought_content = None + for content in assistant_msg.content: + if content.type == "thought": + thought_content = content + break + + assert thought_content is not None, "Should have thought in assistant message" + assert isinstance(thought_content, StoredInputMessageContentThought), ( + "Must be StoredInputMessageContentThought instance" + ) + assert hasattr(thought_content, "text"), "StoredInputMessageContentThought must have 'text' field" + if original_thought_text is not None: + assert thought_content.text == original_thought_text, "Thought must be preserved through round-trip" + + # ============================================================================ + # Step 8: Create datapoint from follow-up inference + # ============================================================================ + + dataset_name = f"test_thought_roundtrip_{uuid7()}" + + try: + datapoint_response = await async_client.create_datapoints_from_inferences( + dataset_name=dataset_name, + params=CreateDatapointsFromInferenceRequestParamsInferenceIds(inference_ids=[follow_up_id]), + output_source="inference", + ) + + assert datapoint_response.ids is not None, "create_datapoints_from_inferences must return IDs" + assert len(datapoint_response.ids) == 1, "Should create exactly 1 datapoint" + + datapoint_id = datapoint_response.ids[0] + + # ============================================================================ + # Step 9: Retrieve datapoint via get_datapoints + # ============================================================================ + datapoint_get_response = await async_client.get_datapoints(dataset_name=dataset_name, ids=[datapoint_id]) + + assert datapoint_get_response.datapoints is not None, "get_datapoints must return datapoints" + assert len(datapoint_get_response.datapoints) == 1, "Should retrieve exactly 1 datapoint" + + datapoint = datapoint_get_response.datapoints[0] + assert isinstance(datapoint, DatapointChat), "Must be DatapointChat instance" + assert datapoint.id == datapoint_id, "Datapoint ID must match" + + # ============================================================================ + # Step 10: Verify thought content in datapoint input + # ============================================================================ + + datapoint_input_messages = datapoint.input.messages + assert datapoint_input_messages is not None, "Datapoint input messages must not be None" + assert len(datapoint_input_messages) >= 3, "Datapoint should have user, assistant, and follow-up messages" + + # Find assistant message with thought in datapoint input + datapoint_assistant_msg = None + for msg in datapoint_input_messages: + if msg.role == "assistant": + datapoint_assistant_msg = msg + break + + assert datapoint_assistant_msg is not None, "Datapoint should have assistant message" + assert datapoint_assistant_msg.content is not None, "Datapoint assistant message content must not be None" + assert len(datapoint_assistant_msg.content) > 0, "Datapoint assistant message should have content" + + # Verify thought in datapoint input + datapoint_thought = None + for content in datapoint_assistant_msg.content: + if content.type == "thought": + datapoint_thought = content + break + + assert datapoint_thought is not None, "Datapoint should have thought in assistant message" + assert isinstance(datapoint_thought, InputMessageContentThought), ( + "Datapoint thought must be InputMessageContentThought instance" + ) + assert hasattr(datapoint_thought, "text"), "Datapoint thought must have 'text' field" + if original_thought_text is not None: + assert datapoint_thought.text == original_thought_text, "Datapoint thought must be preserved" + + # ============================================================================ + # Step 11: Verify thought in datapoint output + # ============================================================================ + + datapoint_output = datapoint.output + assert datapoint_output is not None, "Datapoint must have output" + assert len(datapoint_output) >= 1, "Datapoint output should have at least 1 content block" + + # Find thought in output + datapoint_output_thought = None + for content in datapoint_output: + if content.type == "thought": + datapoint_output_thought = content + break + + assert datapoint_output_thought is not None, "Datapoint output must contain thought" + assert datapoint_output_thought.type == "thought", "Datapoint output must have type='thought'" + assert isinstance(datapoint_output_thought, ContentBlockChatOutputThought), ( + "Datapoint output thought must be ContentBlockChatOutputThought instance" + ) + assert hasattr(datapoint_output_thought, "text"), "Datapoint output thought must have 'text' field" + + finally: + await async_client.delete_dataset(dataset_name=dataset_name) + + +def test_sync_thought_content_roundtrip_complete_flow(sync_client: TensorZeroGateway): + """ + Sync version of test_async_thought_content_roundtrip_complete_flow. + Tests the same round-trip but with synchronous client. + """ + + # ============================================================================ + # Step 1: Create initial inference + # ============================================================================ + + result = sync_client.inference( + function_name="basic_test", + variant_name="reasoner", + input={ + "system": {"assistant_name": "Test Assistant"}, + "messages": [{"role": "user", "content": "What is 2+2?"}], + }, + stream=False, + ) + + assert isinstance(result, ChatInferenceResponse) + assert result.content is not None, "Result content must not be None" + assert result.inference_id is not None, "Result must have inference_id" + + # Find thought in response + thought_response = None + for content in result.content: + if content.type == "thought": + thought_response = content + break + + assert thought_response is not None + assert isinstance(thought_response, Thought) + original_thought_text = thought_response.text + + inference_id = str(result.inference_id) + + # Wait for results to be written to ClickHouse (required for batch writes) + time.sleep(1) + + # ============================================================================ + # Step 3: Query via get_inferences + # ============================================================================ + + get_response = sync_client.get_inferences( + ids=[inference_id], + function_name="basic_test", + output_source="inference", + ) + + stored_inference = get_response.inferences[0] + assert isinstance(stored_inference, StoredInferenceChat), "Must be StoredInferenceChat" + assert stored_inference.output is not None, "Stored inference output must not be None" + + # Find thought in stored output + stored_thought = None + for content in stored_inference.output: + if content.type == "thought": + stored_thought = content + break + + assert stored_thought is not None + assert isinstance(stored_thought, ContentBlockChatOutputThought) + + # ============================================================================ + # Step 5: Serialize and reuse + # ============================================================================ + + thought_dict = asdict(stored_thought) + + follow_up_result = sync_client.inference( + function_name="basic_test", + variant_name="reasoner", + input={ + "system": {"assistant_name": "Test Assistant"}, + "messages": [ + {"role": "user", "content": "Question"}, + {"role": "assistant", "content": [thought_dict, {"type": "text", "text": "Answer"}]}, + {"role": "user", "content": "Follow-up"}, + ], + }, + stream=False, + ) + + assert isinstance(follow_up_result, ChatInferenceResponse), "Follow-up must return ChatInferenceResponse" + assert follow_up_result.inference_id is not None, "Follow-up must have inference_id" + + follow_up_id = str(follow_up_result.inference_id) + + # Wait for follow-up results to be written to ClickHouse (required for batch writes) + time.sleep(1) + + follow_up_stored = sync_client.get_inferences( + ids=[follow_up_id], + function_name="basic_test", + output_source="inference", + ) + + follow_up_inference = follow_up_stored.inferences[0] + + # Verify StoredInputMessageContentThought + input_messages_sync = follow_up_inference.input.messages + assert input_messages_sync is not None, "Input messages must not be None" + + assistant_msg = None + for msg in input_messages_sync: + if msg.role == "assistant": + assistant_msg = msg + break + + assert assistant_msg is not None, "Assistant message must not be None" + assert assistant_msg.content is not None, "Assistant message content must not be None" + + thought_content = None + for content in assistant_msg.content: + if content.type == "thought": + thought_content = content + break + + assert isinstance(thought_content, StoredInputMessageContentThought) + if original_thought_text is not None: + assert thought_content.text == original_thought_text + + # ============================================================================ + # Step 9: Create datapoint + # ============================================================================ + + dataset_name = f"test_thought_roundtrip_{uuid7()}" + + try: + datapoint_response = sync_client.create_datapoints_from_inferences( + dataset_name=dataset_name, + params=CreateDatapointsFromInferenceRequestParamsInferenceIds(inference_ids=[follow_up_id]), + output_source="inference", + ) + + datapoint_id = datapoint_response.ids[0] + datapoint_get_response = sync_client.get_datapoints(dataset_name=dataset_name, ids=[datapoint_id]) + + datapoint = datapoint_get_response.datapoints[0] + assert isinstance(datapoint, DatapointChat) + + # Verify InputMessageContentThought + datapoint_input_messages_sync = datapoint.input.messages + assert datapoint_input_messages_sync is not None, "Datapoint input messages must not be None" + + datapoint_assistant_msg = None + for msg in datapoint_input_messages_sync: + if msg.role == "assistant": + datapoint_assistant_msg = msg + break + + assert datapoint_assistant_msg is not None, "Datapoint assistant message must not be None" + assert datapoint_assistant_msg.content is not None, "Datapoint assistant message content must not be None" + + datapoint_thought = None + for content in datapoint_assistant_msg.content: + if content.type == "thought": + datapoint_thought = content + break + + assert isinstance(datapoint_thought, InputMessageContentThought) + if original_thought_text is not None: + assert datapoint_thought.text == original_thought_text + + # Verify ContentBlockChatOutputThought in output + assert datapoint.output is not None, "Datapoint output must not be None" + output_thought = None + for content in datapoint.output: + if content.type == "thought": + output_thought = content + break + + assert isinstance(output_thought, ContentBlockChatOutputThought) + + finally: + sync_client.delete_dataset(dataset_name=dataset_name) diff --git a/clients/python/tests/test_roundtrip_tool_types.py b/clients/python/tests/test_roundtrip_tool_types.py new file mode 100644 index 00000000000..ac5c89c1220 --- /dev/null +++ b/clients/python/tests/test_roundtrip_tool_types.py @@ -0,0 +1,895 @@ +""" +Test round-trip serialization and reuse of tool calls. + +Verifies that tool call data preserves all fields through: +- Inference response types (ToolCall) +- Storage types (ContentBlockChatOutputToolCall, StoredInputMessageContentToolCall) +- Serialization (asdict, JSON) +- Reuse in follow-up inferences +- Datapoint creation and retrieval (DatapointChat) + +This catches type generation issues like PR #5803. +""" + +import asyncio +import json +import time +from dataclasses import asdict + +import pytest +from tensorzero import ( + AsyncTensorZeroGateway, + ChatInferenceResponse, + CreateDatapointsFromInferenceRequestParamsInferenceIds, + TensorZeroGateway, +) +from tensorzero.generated_types import ( + ContentBlockChatOutputToolCall, + DatapointChat, + InputMessageContentToolCall, + InputMessageContentToolResult, + StoredInferenceChat, + StoredInputMessageContentToolCall, + StoredInputMessageContentToolResult, +) +from tensorzero.types import ToolCall +from uuid_utils import uuid7 + + +@pytest.mark.asyncio +async def test_async_tool_call_roundtrip_complete_flow( + async_client: AsyncTensorZeroGateway, +): + """ + Comprehensive test verifying tool calls preserve all fields through: + 1. Inference response (ToolCall type) + 2. Storage retrieval (StoredInputMessageContentToolCall type) + 3. Serialization (asdict + JSON) + 4. Reuse in follow-up inference (complete tool use flow) + 5. Datapoint creation from inference + 6. Datapoint retrieval and validation (DatapointChat type) + + This test catches type generation issues like PR #5803 where fields were dropped. + Tests 12 steps total (8 inference + 4 datapoint). + """ + + # ============================================================================ + # Step 1: Create initial inference with tool call + # ============================================================================ + + result = await async_client.inference( + function_name="weather_helper", + input={ + "system": {"assistant_name": "Test Assistant"}, + "messages": [{"role": "user", "content": "What's the weather in Brooklyn?"}], + }, + stream=False, + ) + + # Basic result assertions + assert isinstance(result, ChatInferenceResponse), "Result must be ChatInferenceResponse instance" + assert result.content is not None, "Result content must not be None" + assert result.inference_id is not None, "Result must have inference_id" + assert len(result.content) == 1, "Result should have exactly 1 content block" + assert result.content[0].type == "tool_call", "Content block must have type='tool_call'" + assert isinstance(result.content[0], ToolCall), "Content block must be ToolCall instance" + + # ============================================================================ + # Step 2: Verify response ToolCall fields (types.py ToolCall) + # ============================================================================ + + tool_call_response = result.content[0] + + # Type discriminator + assert tool_call_response.type == "tool_call", "Tool call must have type='tool_call'" + + # isinstance check + assert isinstance(tool_call_response, ToolCall), "Response must be ToolCall instance" + + # Field existence (hasattr) + assert hasattr(tool_call_response, "id"), "ToolCall must have 'id' field" + assert hasattr(tool_call_response, "name"), "ToolCall must have 'name' field" + assert hasattr(tool_call_response, "raw_name"), "ToolCall must have 'raw_name' field" + assert hasattr(tool_call_response, "arguments"), "ToolCall must have 'arguments' field" + assert hasattr(tool_call_response, "raw_arguments"), "ToolCall must have 'raw_arguments' field" + + # Field values not None + assert tool_call_response.id is not None, "Tool call id must not be None" + assert tool_call_response.name is not None, "Tool call name must not be None" + assert tool_call_response.raw_name is not None, "Tool call raw_name must not be None" + assert tool_call_response.arguments is not None, "Tool call arguments must not be None" + assert tool_call_response.raw_arguments is not None, "Tool call raw_arguments must not be None" + + # Expected values (from weather_helper deterministic behavior) + assert tool_call_response.id == "0", "Tool call id must be '0'" + assert tool_call_response.name == "get_temperature", "Tool call name must be 'get_temperature'" + assert tool_call_response.raw_name == "get_temperature", "Tool call raw_name must be 'get_temperature'" + assert tool_call_response.arguments == { + "location": "Brooklyn", + "units": "celsius", + }, "Tool call arguments must match expected dict" + assert tool_call_response.raw_arguments == '{"location":"Brooklyn","units":"celsius"}', ( + "Tool call raw_arguments must match expected JSON string" + ) + + # Store inference_id for retrieval + inference_id = str(result.inference_id) + + # Wait for results to be written to ClickHouse (required for batch writes) + await asyncio.sleep(1) + + # ============================================================================ + # Step 3: Query inference back via get_inferences + # ============================================================================ + + get_response = await async_client.get_inferences( + ids=[inference_id], + function_name="weather_helper", # Improves query performance + output_source="inference", + ) + + assert get_response.inferences is not None, "get_inferences must return inferences" + assert len(get_response.inferences) == 1, "Should retrieve exactly 1 inference" + + stored_inference = get_response.inferences[0] + assert isinstance(stored_inference, StoredInferenceChat), "Must be StoredInferenceChat instance" + assert stored_inference.output is not None, "Stored inference output must not be None" + assert str(stored_inference.inference_id) == inference_id, "Retrieved inference ID must match original" + + # ============================================================================ + # Step 4: Verify stored output ToolCall (ContentBlockChatOutputToolCall) + # ============================================================================ + + stored_output = stored_inference.output + assert len(stored_output) == 1, "Output should have exactly 1 content block" + + stored_tool_call = stored_output[0] + + # Type discriminator + assert stored_tool_call.type == "tool_call", "Stored output must have type='tool_call'" + + # isinstance check + assert isinstance(stored_tool_call, ContentBlockChatOutputToolCall), ( + "Stored output must be ContentBlockChatOutputToolCall instance" + ) + + # Field existence (hasattr) + assert hasattr(stored_tool_call, "id"), "Stored tool call must have 'id' field" + assert hasattr(stored_tool_call, "name"), "Stored tool call must have 'name' field" + assert hasattr(stored_tool_call, "raw_name"), "Stored tool call must have 'raw_name' field" + assert hasattr(stored_tool_call, "arguments"), "Stored tool call must have 'arguments' field" + assert hasattr(stored_tool_call, "raw_arguments"), "Stored tool call must have 'raw_arguments' field" + + # Field values not None + assert stored_tool_call.id is not None, "Stored tool call id must not be None" + assert stored_tool_call.name is not None, "Stored tool call name must not be None" + assert stored_tool_call.raw_name is not None, "Stored tool call raw_name must not be None" + assert stored_tool_call.arguments is not None, "Stored tool call arguments must not be None" + assert stored_tool_call.raw_arguments is not None, "Stored tool call raw_arguments must not be None" + + # Values match original + assert stored_tool_call.id == "0", "Stored tool call id must be '0'" + assert stored_tool_call.name == "get_temperature", "Stored tool call name must be 'get_temperature'" + assert stored_tool_call.raw_name == "get_temperature", "Stored tool call raw_name must be 'get_temperature'" + assert stored_tool_call.arguments == { + "location": "Brooklyn", + "units": "celsius", + }, "Stored tool call arguments must match expected dict" + assert stored_tool_call.raw_arguments == '{"location":"Brooklyn","units":"celsius"}', ( + "Stored tool call raw_arguments must match expected JSON string" + ) + + # ============================================================================ + # Step 5: Serialize tool call with asdict() and JSON + # ============================================================================ + + # Serialize to dict + tool_call_dict = asdict(stored_tool_call) + + # Verify serialized dict has all required fields + assert "type" in tool_call_dict and tool_call_dict["type"] == "tool_call", ( + "Serialized dict must include type='tool_call'" + ) + assert "id" in tool_call_dict, "Serialized dict must include 'id' field" + assert "name" in tool_call_dict, "Serialized dict must include 'name' field" + assert "arguments" in tool_call_dict, "Serialized dict must include 'arguments' field" + assert "raw_name" in tool_call_dict, "Serialized dict must include 'raw_name' field" + assert "raw_arguments" in tool_call_dict, "Serialized dict must include 'raw_arguments' field" + + # Verify JSON serialization works (no encoding issues) + tool_call_json = json.dumps(tool_call_dict) + assert isinstance(tool_call_json, str), "Must serialize to JSON string" + + # Verify JSON deserialization works + tool_call_from_json = json.loads(tool_call_json) + assert tool_call_from_json["id"] == "0", "Deserialized JSON must preserve id field" + assert tool_call_from_json["name"] == "get_temperature", "Deserialized JSON must preserve name field" + assert tool_call_from_json["arguments"] == { + "location": "Brooklyn", + "units": "celsius", + }, "Deserialized JSON must preserve arguments field" + + # ============================================================================ + # Step 6: Create mock tool result + # ============================================================================ + + tool_result = { + "type": "tool_result", + "id": "0", + "name": "get_temperature", + "result": '{"temperature": 72, "conditions": "sunny"}', + } + + # ============================================================================ + # Step 7: Reuse in follow-up inference (complete tool use flow) + # ============================================================================ + + follow_up_result = await async_client.inference( + function_name="weather_helper", + input={ + "system": {"assistant_name": "Test Assistant"}, + "messages": [ + # Original user message + {"role": "user", "content": "What's the weather in Brooklyn?"}, + # Assistant message with tool call (using serialized data) + { + "role": "assistant", + "content": [tool_call_dict], # Reuse serialized tool call + }, + # Tool result + {"role": "user", "content": [tool_result]}, + # Follow-up user message + {"role": "user", "content": "How about tomorrow?"}, + ], + }, + stream=False, + ) + + # The critical assertion: the inference succeeded + assert isinstance(follow_up_result, ChatInferenceResponse), "Follow-up must return ChatInferenceResponse instance" + assert follow_up_result.inference_id is not None, ( + "Follow-up inference must succeed when reusing serialized tool call data" + ) + + # Verify the inference ran (not just that it didn't crash) + assert follow_up_result.content is not None, "Follow-up must have content" + assert len(follow_up_result.content) > 0, "Follow-up must generate content" + + # Wait for follow-up results to be written to ClickHouse (required for batch writes) + await asyncio.sleep(1) + + # ============================================================================ + # Step 8: Verify follow-up stored data + # ============================================================================ + + follow_up_id = str(follow_up_result.inference_id) + follow_up_stored = await async_client.get_inferences( + ids=[follow_up_id], + function_name="weather_helper", + output_source="inference", + ) + + assert len(follow_up_stored.inferences) == 1, "Should retrieve exactly 1 follow-up inference" + follow_up_inference = follow_up_stored.inferences[0] + + # Verify input contains our tool call and tool result + input_messages = follow_up_inference.input.messages + assert input_messages is not None, "Input messages must not be None" + assert len(input_messages) >= 3, "Should have user, assistant, and follow-up user messages" + + # Find assistant message with tool call + assistant_msg = None + for msg in input_messages: + if msg.role == "assistant": + assistant_msg = msg + break + + assert assistant_msg is not None, "Should have assistant message in stored input" + assert assistant_msg.content is not None, "Assistant message content must not be None" + assert len(assistant_msg.content) > 0, "Assistant message should have content" + + # Verify the tool call was stored correctly (StoredInputMessageContentToolCall) + tool_call_content = None + for content in assistant_msg.content: + if content.type == "tool_call": + tool_call_content = content + break + + assert tool_call_content is not None, "Should have tool call in assistant message" + assert isinstance(tool_call_content, StoredInputMessageContentToolCall), ( + "Must be StoredInputMessageContentToolCall instance" + ) + + # Critical assertions: verify fields preserved through full round-trip + assert hasattr(tool_call_content, "id"), "StoredInputMessageContentToolCall must have 'id' field" + assert hasattr(tool_call_content, "name"), "StoredInputMessageContentToolCall must have 'name' field" + assert hasattr(tool_call_content, "arguments"), "StoredInputMessageContentToolCall must have 'arguments' field" + + assert tool_call_content.id == "0", "Tool call id must be preserved through round-trip" + assert tool_call_content.name == "get_temperature", "Tool call name must be preserved through round-trip" + assert tool_call_content.arguments == '{"location":"Brooklyn","units":"celsius"}', ( + "Tool call arguments must be preserved through round-trip" + ) + + # Verify tool result was also stored correctly + user_msg_with_result = None + for msg in input_messages: + if msg.role == "user": + for content in msg.content: + if content.type == "tool_result": + user_msg_with_result = content + break + if user_msg_with_result: + break + + assert user_msg_with_result is not None, "Should have tool result in stored input" + assert isinstance(user_msg_with_result, StoredInputMessageContentToolResult), ( + "Must be StoredInputMessageContentToolResult instance" + ) + assert user_msg_with_result.id == "0", "Tool result id must match tool call id" + assert user_msg_with_result.name == "get_temperature", "Tool result name must match tool call name" + assert user_msg_with_result.result == '{"temperature": 72, "conditions": "sunny"}', "Tool result must be preserved" + + # ============================================================================ + # Step 9: Create datapoint from follow-up inference (has complete tool flow) + # ============================================================================ + + dataset_name = f"test_tool_call_roundtrip_{uuid7()}" + + try: + # Use follow_up_id instead of inference_id to get the complete conversation + datapoint_response = await async_client.create_datapoints_from_inferences( + dataset_name=dataset_name, + params=CreateDatapointsFromInferenceRequestParamsInferenceIds( + inference_ids=[follow_up_id] # Use follow-up inference with full conversation + ), + output_source="inference", + ) + + assert datapoint_response.ids is not None, "create_datapoints_from_inferences must return IDs" + assert len(datapoint_response.ids) == 1, "Should create exactly 1 datapoint" + + datapoint_id = datapoint_response.ids[0] + + # ============================================================================ + # Step 10: Retrieve datapoint via get_datapoints + # ============================================================================ + datapoint_get_response = await async_client.get_datapoints(dataset_name=dataset_name, ids=[datapoint_id]) + + assert datapoint_get_response.datapoints is not None, "get_datapoints must return datapoints" + assert len(datapoint_get_response.datapoints) == 1, "Should retrieve exactly 1 datapoint" + + datapoint = datapoint_get_response.datapoints[0] + assert isinstance(datapoint, DatapointChat), "Must be DatapointChat instance" + assert datapoint.id == datapoint_id, "Datapoint ID must match" + + # ============================================================================ + # Step 11: Verify tool calls in datapoint input + # ============================================================================ + + datapoint_input_messages = datapoint.input.messages + assert datapoint_input_messages is not None, "Datapoint input messages must not be None" + assert len(datapoint_input_messages) >= 3, "Datapoint should have user, assistant, and follow-up messages" + + # Find assistant message with tool call in datapoint input + datapoint_assistant_msg = None + for msg in datapoint_input_messages: + if msg.role == "assistant": + datapoint_assistant_msg = msg + break + + assert datapoint_assistant_msg is not None, "Datapoint should have assistant message" + assert datapoint_assistant_msg.content is not None, "Datapoint assistant message content must not be None" + assert len(datapoint_assistant_msg.content) > 0, "Datapoint assistant message should have content" + + # Verify tool call in datapoint input + datapoint_tool_call = None + for content in datapoint_assistant_msg.content: + if content.type == "tool_call": + datapoint_tool_call = content + break + + assert datapoint_tool_call is not None, "Datapoint should have tool call in assistant message" + assert isinstance(datapoint_tool_call, InputMessageContentToolCall), ( + "Datapoint tool call must be InputMessageContentToolCall instance" + ) + + # Verify tool call fields preserved in datapoint + assert hasattr(datapoint_tool_call, "id"), "Datapoint tool call must have 'id' field" + assert hasattr(datapoint_tool_call, "name"), "Datapoint tool call must have 'name' field" + assert hasattr(datapoint_tool_call, "arguments"), "Datapoint tool call must have 'arguments' field" + + assert datapoint_tool_call.id == "0", "Datapoint tool call id must be preserved" + assert datapoint_tool_call.name == "get_temperature", "Datapoint tool call name must be preserved" + assert datapoint_tool_call.arguments == '{"location":"Brooklyn","units":"celsius"}', ( + "Datapoint tool call arguments must be preserved" + ) + + # Verify tool result in datapoint input + datapoint_tool_result = None + for msg in datapoint_input_messages: + if msg.role == "user": + for content in msg.content: + if content.type == "tool_result": + datapoint_tool_result = content + break + if datapoint_tool_result: + break + + assert datapoint_tool_result is not None, "Datapoint should have tool result in input" + assert isinstance(datapoint_tool_result, InputMessageContentToolResult), ( + "Datapoint tool result must be InputMessageContentToolResult instance" + ) + + # Verify tool result fields preserved in datapoint + assert hasattr(datapoint_tool_result, "id"), "Datapoint tool result must have 'id' field" + assert hasattr(datapoint_tool_result, "name"), "Datapoint tool result must have 'name' field" + assert hasattr(datapoint_tool_result, "result"), "Datapoint tool result must have 'result' field" + + assert datapoint_tool_result.id == "0", "Datapoint tool result id must match tool call id" + assert datapoint_tool_result.name == "get_temperature", "Datapoint tool result name must match tool call name" + assert datapoint_tool_result.result == '{"temperature": 72, "conditions": "sunny"}', ( + "Datapoint tool result must be preserved" + ) + + # ============================================================================ + # Step 12: Verify tool calls in datapoint output + # ============================================================================ + + datapoint_output = datapoint.output + assert datapoint_output is not None, "Datapoint must have output" + assert len(datapoint_output) == 1, "Datapoint output should have exactly 1 content block" + + datapoint_output_tool_call = datapoint_output[0] + assert datapoint_output_tool_call.type == "tool_call", "Datapoint output must have type='tool_call'" + assert isinstance(datapoint_output_tool_call, ContentBlockChatOutputToolCall), ( + "Datapoint output tool call must be ContentBlockChatOutputToolCall instance" + ) + + # Verify output tool call fields + assert hasattr(datapoint_output_tool_call, "id"), "Datapoint output tool call must have 'id' field" + assert hasattr(datapoint_output_tool_call, "name"), "Datapoint output tool call must have 'name' field" + assert hasattr(datapoint_output_tool_call, "arguments"), ( + "Datapoint output tool call must have 'arguments' field" + ) + + assert datapoint_output_tool_call.id == "0", "Datapoint output tool call id must be '0'" + assert datapoint_output_tool_call.name == "get_temperature", ( + "Datapoint output tool call name must be 'get_temperature'" + ) + assert datapoint_output_tool_call.arguments == { + "location": "Brooklyn", + "units": "celsius", + }, "Datapoint output tool call arguments must match expected dict" + + finally: + # Clean up: delete the dataset + await async_client.delete_dataset(dataset_name=dataset_name) + + +def test_sync_tool_call_roundtrip_complete_flow(sync_client: TensorZeroGateway): + """ + Comprehensive test verifying tool calls preserve all fields through: + 1. Inference response (ToolCall type) + 2. Storage retrieval (StoredInputMessageContentToolCall type) + 3. Serialization (asdict + JSON) + 4. Reuse in follow-up inference (complete tool use flow) + 5. Datapoint creation from inference + 6. Datapoint retrieval and validation (DatapointChat type) + + This test catches type generation issues like PR #5803 where fields were dropped. + Tests 12 steps total (8 inference + 4 datapoint). + """ + + # ============================================================================ + # Step 1: Create initial inference with tool call + # ============================================================================ + + result = sync_client.inference( + function_name="weather_helper", + input={ + "system": {"assistant_name": "Test Assistant"}, + "messages": [{"role": "user", "content": "What's the weather in Brooklyn?"}], + }, + stream=False, + ) + + # Basic result assertions + assert isinstance(result, ChatInferenceResponse), "Result must be ChatInferenceResponse instance" + assert result.content is not None, "Result content must not be None" + assert result.inference_id is not None, "Result must have inference_id" + assert len(result.content) == 1, "Result should have exactly 1 content block" + assert result.content[0].type == "tool_call", "Content block must have type='tool_call'" + assert isinstance(result.content[0], ToolCall), "Content block must be ToolCall instance" + + # ============================================================================ + # Step 2: Verify response ToolCall fields (types.py ToolCall) + # ============================================================================ + + tool_call_response = result.content[0] + + # Type discriminator + assert tool_call_response.type == "tool_call", "Tool call must have type='tool_call'" + + # isinstance check + assert isinstance(tool_call_response, ToolCall), "Response must be ToolCall instance" + + # Field existence (hasattr) + assert hasattr(tool_call_response, "id"), "ToolCall must have 'id' field" + assert hasattr(tool_call_response, "name"), "ToolCall must have 'name' field" + assert hasattr(tool_call_response, "raw_name"), "ToolCall must have 'raw_name' field" + assert hasattr(tool_call_response, "arguments"), "ToolCall must have 'arguments' field" + assert hasattr(tool_call_response, "raw_arguments"), "ToolCall must have 'raw_arguments' field" + + # Field values not None + assert tool_call_response.id is not None, "Tool call id must not be None" + assert tool_call_response.name is not None, "Tool call name must not be None" + assert tool_call_response.raw_name is not None, "Tool call raw_name must not be None" + assert tool_call_response.arguments is not None, "Tool call arguments must not be None" + assert tool_call_response.raw_arguments is not None, "Tool call raw_arguments must not be None" + + # Expected values (from weather_helper deterministic behavior) + assert tool_call_response.id == "0", "Tool call id must be '0'" + assert tool_call_response.name == "get_temperature", "Tool call name must be 'get_temperature'" + assert tool_call_response.raw_name == "get_temperature", "Tool call raw_name must be 'get_temperature'" + assert tool_call_response.arguments == { + "location": "Brooklyn", + "units": "celsius", + }, "Tool call arguments must match expected dict" + assert tool_call_response.raw_arguments == '{"location":"Brooklyn","units":"celsius"}', ( + "Tool call raw_arguments must match expected JSON string" + ) + + # Store inference_id for retrieval + inference_id = str(result.inference_id) + + # Wait for results to be written to ClickHouse (required for batch writes) + time.sleep(1) + + # ============================================================================ + # Step 3: Query inference back via get_inferences + # ============================================================================ + + get_response = sync_client.get_inferences( + ids=[inference_id], + function_name="weather_helper", # Improves query performance + output_source="inference", + ) + + assert get_response.inferences is not None, "get_inferences must return inferences" + assert len(get_response.inferences) == 1, "Should retrieve exactly 1 inference" + + stored_inference = get_response.inferences[0] + assert isinstance(stored_inference, StoredInferenceChat), "Must be StoredInferenceChat instance" + assert stored_inference.output is not None, "Stored inference output must not be None" + assert str(stored_inference.inference_id) == inference_id, "Retrieved inference ID must match original" + + # ============================================================================ + # Step 4: Verify stored output ToolCall (ContentBlockChatOutputToolCall) + # ============================================================================ + + stored_output = stored_inference.output + assert len(stored_output) == 1, "Output should have exactly 1 content block" + + stored_tool_call = stored_output[0] + + # Type discriminator + assert stored_tool_call.type == "tool_call", "Stored output must have type='tool_call'" + + # isinstance check + assert isinstance(stored_tool_call, ContentBlockChatOutputToolCall), ( + "Stored output must be ContentBlockChatOutputToolCall instance" + ) + + # Field existence (hasattr) + assert hasattr(stored_tool_call, "id"), "Stored tool call must have 'id' field" + assert hasattr(stored_tool_call, "name"), "Stored tool call must have 'name' field" + assert hasattr(stored_tool_call, "raw_name"), "Stored tool call must have 'raw_name' field" + assert hasattr(stored_tool_call, "arguments"), "Stored tool call must have 'arguments' field" + assert hasattr(stored_tool_call, "raw_arguments"), "Stored tool call must have 'raw_arguments' field" + + # Field values not None + assert stored_tool_call.id is not None, "Stored tool call id must not be None" + assert stored_tool_call.name is not None, "Stored tool call name must not be None" + assert stored_tool_call.raw_name is not None, "Stored tool call raw_name must not be None" + assert stored_tool_call.arguments is not None, "Stored tool call arguments must not be None" + assert stored_tool_call.raw_arguments is not None, "Stored tool call raw_arguments must not be None" + + # Values match original + assert stored_tool_call.id == "0", "Stored tool call id must be '0'" + assert stored_tool_call.name == "get_temperature", "Stored tool call name must be 'get_temperature'" + assert stored_tool_call.raw_name == "get_temperature", "Stored tool call raw_name must be 'get_temperature'" + assert stored_tool_call.arguments == { + "location": "Brooklyn", + "units": "celsius", + }, "Stored tool call arguments must match expected dict" + assert stored_tool_call.raw_arguments == '{"location":"Brooklyn","units":"celsius"}', ( + "Stored tool call raw_arguments must match expected JSON string" + ) + + # ============================================================================ + # Step 5: Serialize tool call with asdict() and JSON + # ============================================================================ + + # Serialize to dict + tool_call_dict = asdict(stored_tool_call) + + # Verify serialized dict has all required fields + assert "type" in tool_call_dict and tool_call_dict["type"] == "tool_call", ( + "Serialized dict must include type='tool_call'" + ) + assert "id" in tool_call_dict, "Serialized dict must include 'id' field" + assert "name" in tool_call_dict, "Serialized dict must include 'name' field" + assert "arguments" in tool_call_dict, "Serialized dict must include 'arguments' field" + assert "raw_name" in tool_call_dict, "Serialized dict must include 'raw_name' field" + assert "raw_arguments" in tool_call_dict, "Serialized dict must include 'raw_arguments' field" + + # Verify JSON serialization works (no encoding issues) + tool_call_json = json.dumps(tool_call_dict) + assert isinstance(tool_call_json, str), "Must serialize to JSON string" + + # Verify JSON deserialization works + tool_call_from_json = json.loads(tool_call_json) + assert tool_call_from_json["id"] == "0", "Deserialized JSON must preserve id field" + assert tool_call_from_json["name"] == "get_temperature", "Deserialized JSON must preserve name field" + assert tool_call_from_json["arguments"] == { + "location": "Brooklyn", + "units": "celsius", + }, "Deserialized JSON must preserve arguments field" + + # ============================================================================ + # Step 6: Create mock tool result + # ============================================================================ + + tool_result = { + "type": "tool_result", + "id": "0", + "name": "get_temperature", + "result": '{"temperature": 72, "conditions": "sunny"}', + } + + # ============================================================================ + # Step 7: Reuse in follow-up inference (complete tool use flow) + # ============================================================================ + + follow_up_result = sync_client.inference( + function_name="weather_helper", + input={ + "system": {"assistant_name": "Test Assistant"}, + "messages": [ + # Original user message + {"role": "user", "content": "What's the weather in Brooklyn?"}, + # Assistant message with tool call (using serialized data) + { + "role": "assistant", + "content": [tool_call_dict], # Reuse serialized tool call + }, + # Tool result + {"role": "user", "content": [tool_result]}, + # Follow-up user message + {"role": "user", "content": "How about tomorrow?"}, + ], + }, + stream=False, + ) + + # The critical assertion: the inference succeeded + assert isinstance(follow_up_result, ChatInferenceResponse), "Follow-up must return ChatInferenceResponse instance" + assert follow_up_result.inference_id is not None, ( + "Follow-up inference must succeed when reusing serialized tool call data" + ) + + # Verify the inference ran (not just that it didn't crash) + assert follow_up_result.content is not None, "Follow-up must have content" + assert len(follow_up_result.content) > 0, "Follow-up must generate content" + + # Wait for follow-up results to be written to ClickHouse (required for batch writes) + time.sleep(1) + + # ============================================================================ + # Step 8: Verify follow-up stored data + # ============================================================================ + + follow_up_id = str(follow_up_result.inference_id) + follow_up_stored = sync_client.get_inferences( + ids=[follow_up_id], + function_name="weather_helper", + output_source="inference", + ) + + assert len(follow_up_stored.inferences) == 1, "Should retrieve exactly 1 follow-up inference" + follow_up_inference = follow_up_stored.inferences[0] + + # Verify input contains our tool call and tool result + input_messages = follow_up_inference.input.messages + assert input_messages is not None, "Input messages must not be None" + assert len(input_messages) >= 3, "Should have user, assistant, and follow-up user messages" + + # Find assistant message with tool call + assistant_msg = None + for msg in input_messages: + if msg.role == "assistant": + assistant_msg = msg + break + + assert assistant_msg is not None, "Should have assistant message in stored input" + assert assistant_msg.content is not None, "Assistant message content must not be None" + assert len(assistant_msg.content) > 0, "Assistant message should have content" + + # Verify the tool call was stored correctly (StoredInputMessageContentToolCall) + tool_call_content = None + for content in assistant_msg.content: + if content.type == "tool_call": + tool_call_content = content + break + + assert tool_call_content is not None, "Should have tool call in assistant message" + assert isinstance(tool_call_content, StoredInputMessageContentToolCall), ( + "Must be StoredInputMessageContentToolCall instance" + ) + + # Critical assertions: verify fields preserved through full round-trip + assert hasattr(tool_call_content, "id"), "StoredInputMessageContentToolCall must have 'id' field" + assert hasattr(tool_call_content, "name"), "StoredInputMessageContentToolCall must have 'name' field" + assert hasattr(tool_call_content, "arguments"), "StoredInputMessageContentToolCall must have 'arguments' field" + + assert tool_call_content.id == "0", "Tool call id must be preserved through round-trip" + assert tool_call_content.name == "get_temperature", "Tool call name must be preserved through round-trip" + assert tool_call_content.arguments == '{"location":"Brooklyn","units":"celsius"}', ( + "Tool call arguments must be preserved through round-trip" + ) + + # Verify tool result was also stored correctly + user_msg_with_result = None + for msg in input_messages: + if msg.role == "user": + for content in msg.content: + if content.type == "tool_result": + user_msg_with_result = content + break + if user_msg_with_result: + break + + assert user_msg_with_result is not None, "Should have tool result in stored input" + assert isinstance(user_msg_with_result, StoredInputMessageContentToolResult), ( + "Must be StoredInputMessageContentToolResult instance" + ) + assert user_msg_with_result.id == "0", "Tool result id must match tool call id" + assert user_msg_with_result.name == "get_temperature", "Tool result name must match tool call name" + assert user_msg_with_result.result == '{"temperature": 72, "conditions": "sunny"}', "Tool result must be preserved" + + # ============================================================================ + # Step 9: Create datapoint from follow-up inference (has complete tool flow) + # ============================================================================ + + dataset_name = f"test_tool_call_roundtrip_{uuid7()}" + + try: + # Use follow_up_id instead of inference_id to get the complete conversation + datapoint_response = sync_client.create_datapoints_from_inferences( + dataset_name=dataset_name, + params=CreateDatapointsFromInferenceRequestParamsInferenceIds( + inference_ids=[follow_up_id] # Use follow-up inference with full conversation + ), + output_source="inference", + ) + + assert datapoint_response.ids is not None, "create_datapoints_from_inferences must return IDs" + assert len(datapoint_response.ids) == 1, "Should create exactly 1 datapoint" + + datapoint_id = datapoint_response.ids[0] + + # ============================================================================ + # Step 10: Retrieve datapoint via get_datapoints + # ============================================================================ + datapoint_get_response = sync_client.get_datapoints(dataset_name=dataset_name, ids=[datapoint_id]) + + assert datapoint_get_response.datapoints is not None, "get_datapoints must return datapoints" + assert len(datapoint_get_response.datapoints) == 1, "Should retrieve exactly 1 datapoint" + + datapoint = datapoint_get_response.datapoints[0] + assert isinstance(datapoint, DatapointChat), "Must be DatapointChat instance" + assert datapoint.id == datapoint_id, "Datapoint ID must match" + + # ============================================================================ + # Step 11: Verify tool calls in datapoint input + # ============================================================================ + + datapoint_input_messages = datapoint.input.messages + assert datapoint_input_messages is not None, "Datapoint input messages must not be None" + assert len(datapoint_input_messages) >= 3, "Datapoint should have user, assistant, and follow-up messages" + + # Find assistant message with tool call in datapoint input + datapoint_assistant_msg = None + for msg in datapoint_input_messages: + if msg.role == "assistant": + datapoint_assistant_msg = msg + break + + assert datapoint_assistant_msg is not None, "Datapoint should have assistant message" + assert datapoint_assistant_msg.content is not None, "Datapoint assistant message content must not be None" + assert len(datapoint_assistant_msg.content) > 0, "Datapoint assistant message should have content" + + # Verify tool call in datapoint input + datapoint_tool_call = None + for content in datapoint_assistant_msg.content: + if content.type == "tool_call": + datapoint_tool_call = content + break + + assert datapoint_tool_call is not None, "Datapoint should have tool call in assistant message" + assert isinstance(datapoint_tool_call, InputMessageContentToolCall), ( + "Datapoint tool call must be InputMessageContentToolCall instance" + ) + + # Verify tool call fields preserved in datapoint + assert hasattr(datapoint_tool_call, "id"), "Datapoint tool call must have 'id' field" + assert hasattr(datapoint_tool_call, "name"), "Datapoint tool call must have 'name' field" + assert hasattr(datapoint_tool_call, "arguments"), "Datapoint tool call must have 'arguments' field" + + assert datapoint_tool_call.id == "0", "Datapoint tool call id must be preserved" + assert datapoint_tool_call.name == "get_temperature", "Datapoint tool call name must be preserved" + assert datapoint_tool_call.arguments == '{"location":"Brooklyn","units":"celsius"}', ( + "Datapoint tool call arguments must be preserved" + ) + + # Verify tool result in datapoint input + datapoint_tool_result = None + for msg in datapoint_input_messages: + if msg.role == "user": + for content in msg.content: + if content.type == "tool_result": + datapoint_tool_result = content + break + if datapoint_tool_result: + break + + assert datapoint_tool_result is not None, "Datapoint should have tool result in input" + assert isinstance(datapoint_tool_result, InputMessageContentToolResult), ( + "Datapoint tool result must be InputMessageContentToolResult instance" + ) + + # Verify tool result fields preserved in datapoint + assert hasattr(datapoint_tool_result, "id"), "Datapoint tool result must have 'id' field" + assert hasattr(datapoint_tool_result, "name"), "Datapoint tool result must have 'name' field" + assert hasattr(datapoint_tool_result, "result"), "Datapoint tool result must have 'result' field" + + assert datapoint_tool_result.id == "0", "Datapoint tool result id must match tool call id" + assert datapoint_tool_result.name == "get_temperature", "Datapoint tool result name must match tool call name" + assert datapoint_tool_result.result == '{"temperature": 72, "conditions": "sunny"}', ( + "Datapoint tool result must be preserved" + ) + + # ============================================================================ + # Step 12: Verify tool calls in datapoint output + # ============================================================================ + + datapoint_output = datapoint.output + assert datapoint_output is not None, "Datapoint must have output" + assert len(datapoint_output) == 1, "Datapoint output should have exactly 1 content block" + + datapoint_output_tool_call = datapoint_output[0] + assert datapoint_output_tool_call.type == "tool_call", "Datapoint output must have type='tool_call'" + assert isinstance(datapoint_output_tool_call, ContentBlockChatOutputToolCall), ( + "Datapoint output tool call must be ContentBlockChatOutputToolCall instance" + ) + + # Verify output tool call fields + assert hasattr(datapoint_output_tool_call, "id"), "Datapoint output tool call must have 'id' field" + assert hasattr(datapoint_output_tool_call, "name"), "Datapoint output tool call must have 'name' field" + assert hasattr(datapoint_output_tool_call, "arguments"), ( + "Datapoint output tool call must have 'arguments' field" + ) + + assert datapoint_output_tool_call.id == "0", "Datapoint output tool call id must be '0'" + assert datapoint_output_tool_call.name == "get_temperature", ( + "Datapoint output tool call name must be 'get_temperature'" + ) + assert datapoint_output_tool_call.arguments == { + "location": "Brooklyn", + "units": "celsius", + }, "Datapoint output tool call arguments must match expected dict" + + finally: + # Clean up: delete the dataset + sync_client.delete_dataset(dataset_name=dataset_name) From cbaae0a2f838c916d7db7d007e3c758e3ebe6391 Mon Sep 17 00:00:00 2001 From: Shuyang Li Date: Mon, 26 Jan 2026 23:11:07 -0500 Subject: [PATCH 19/46] Change backend connection errors from 500 to 503 (#5736) (#5833) * Change rate limiting backend connection errors from 500 to 503 (#5736) Change PostgresConnection and ValkeyConnection errors to return HTTP 503 (Service Unavailable) instead of 500 (Internal Server Error). This is more semantically correct for transient backend failures, signals to clients that retrying might succeed, and distinguishes infrastructure issues from application bugs. * Change ClickHouseConnection error from 500 to 503 For consistency with PostgresConnection and ValkeyConnection, change ClickHouseConnection errors to return HTTP 503 (Service Unavailable) instead of 500 (Internal Server Error). This is more semantically correct for transient backend failures and signals to clients that retrying might succeed. --------- Co-authored-by: Claude --- tensorzero-core/src/error/mod.rs | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tensorzero-core/src/error/mod.rs b/tensorzero-core/src/error/mod.rs index 09e77807b05..94a25716b22 100644 --- a/tensorzero-core/src/error/mod.rs +++ b/tensorzero-core/src/error/mod.rs @@ -797,7 +797,7 @@ impl ErrorDetails { ErrorDetails::Cache { .. } => StatusCode::INTERNAL_SERVER_ERROR, ErrorDetails::ChannelWrite { .. } => StatusCode::INTERNAL_SERVER_ERROR, ErrorDetails::ClickHouseConfiguration { .. } => StatusCode::INTERNAL_SERVER_ERROR, - ErrorDetails::ClickHouseConnection { .. } => StatusCode::INTERNAL_SERVER_ERROR, + ErrorDetails::ClickHouseConnection { .. } => StatusCode::SERVICE_UNAVAILABLE, ErrorDetails::ClickHouseDeserialization { .. } => StatusCode::INTERNAL_SERVER_ERROR, ErrorDetails::ClickHouseMigration { .. } => StatusCode::INTERNAL_SERVER_ERROR, ErrorDetails::ClickHouseMigrationsDisabled => StatusCode::INTERNAL_SERVER_ERROR, @@ -885,11 +885,11 @@ impl ErrorDetails { ErrorDetails::PostgresConnectionInitialization { .. } => { StatusCode::INTERNAL_SERVER_ERROR } - ErrorDetails::PostgresConnection { .. } => StatusCode::INTERNAL_SERVER_ERROR, + ErrorDetails::PostgresConnection { .. } => StatusCode::SERVICE_UNAVAILABLE, ErrorDetails::PostgresQuery { .. } => StatusCode::INTERNAL_SERVER_ERROR, ErrorDetails::PostgresResult { .. } => StatusCode::INTERNAL_SERVER_ERROR, ErrorDetails::PostgresMigration { .. } => StatusCode::INTERNAL_SERVER_ERROR, - ErrorDetails::ValkeyConnection { .. } => StatusCode::INTERNAL_SERVER_ERROR, + ErrorDetails::ValkeyConnection { .. } => StatusCode::SERVICE_UNAVAILABLE, ErrorDetails::ValkeyQuery { .. } => StatusCode::INTERNAL_SERVER_ERROR, ErrorDetails::RateLimitExceeded { .. } => StatusCode::TOO_MANY_REQUESTS, ErrorDetails::RateLimitMissingMaxTokens => StatusCode::BAD_REQUEST, From 5a52ae17a0239a8cc123e23c4fb367658606dfbc Mon Sep 17 00:00:00 2001 From: Shuyang Li Date: Mon, 26 Jan 2026 23:11:07 -0500 Subject: [PATCH 20/46] Add clone_datapoints to DatasetQueries trait (#5834) Move the clone datapoints SQL queries from the endpoint handler to the DatasetQueries trait. This centralizes the database logic and enables future Postgres support. The clone operation: 1. Generates new UUIDs for each source datapoint 2. Inserts cloned rows into Chat and Json datapoint tables in parallel 3. Verifies which clones succeeded (source might not exist) 4. Returns mapping of source IDs to new IDs (None if source missing) --- .../src/db/clickhouse/dataset_queries.rs | 133 ++++++++++++++++ .../mock_clickhouse_connection_info.rs | 10 ++ tensorzero-core/src/db/datasets.rs | 14 ++ .../datasets/internal/clone_datapoints.rs | 143 +----------------- 4 files changed, 162 insertions(+), 138 deletions(-) diff --git a/tensorzero-core/src/db/clickhouse/dataset_queries.rs b/tensorzero-core/src/db/clickhouse/dataset_queries.rs index 45ae403f7a3..76a6d74cc75 100644 --- a/tensorzero-core/src/db/clickhouse/dataset_queries.rs +++ b/tensorzero-core/src/db/clickhouse/dataset_queries.rs @@ -502,6 +502,139 @@ impl DatasetQueries for ClickHouseConnectionInfo { written_rows += json_written_rows?; Ok(written_rows) } + + async fn clone_datapoints( + &self, + target_dataset_name: &str, + source_datapoint_ids: &[Uuid], + ) -> Result>, Error> { + if source_datapoint_ids.is_empty() { + return Ok(vec![]); + } + + // Generate all mappings from source to target IDs + let mappings: Vec<(Uuid, Uuid)> = source_datapoint_ids + .iter() + .map(|id| (*id, Uuid::now_v7())) + .collect(); + + // Build the mappings array string for ClickHouse + let mappings_str = format!( + "[{}]", + mappings + .iter() + .map(|(old, new)| format!("('{old}', '{new}')")) + .collect::>() + .join(",") + ); + + // Clone queries using CTE + EXCEPT pattern + let chat_clone_query = r" + INSERT INTO ChatInferenceDatapoint + WITH source AS ( + SELECT ChatInferenceDatapoint.*, mapping.new_id + FROM ChatInferenceDatapoint FINAL + INNER JOIN ( + SELECT + tupleElement(pair, 1) as old_id, + tupleElement(pair, 2) as new_id + FROM ( + SELECT arrayJoin({mappings: Array(Tuple(UUID, UUID))}) as pair + ) + ) AS mapping ON ChatInferenceDatapoint.id = mapping.old_id + WHERE ChatInferenceDatapoint.staled_at IS NULL + ) + SELECT * EXCEPT(new_id) REPLACE( + new_id AS id, + {target_dataset_name: String} AS dataset_name, + now64() AS updated_at + ) + FROM source + "; + + let json_clone_query = r" + INSERT INTO JsonInferenceDatapoint + WITH source AS ( + SELECT JsonInferenceDatapoint.*, mapping.new_id + FROM JsonInferenceDatapoint FINAL + INNER JOIN ( + SELECT + tupleElement(pair, 1) as old_id, + tupleElement(pair, 2) as new_id + FROM ( + SELECT arrayJoin({mappings: Array(Tuple(UUID, UUID))}) as pair + ) + ) AS mapping ON JsonInferenceDatapoint.id = mapping.old_id + WHERE JsonInferenceDatapoint.staled_at IS NULL + ) + SELECT * EXCEPT(new_id) REPLACE( + new_id AS id, + {target_dataset_name: String} AS dataset_name, + now64() AS updated_at + ) + FROM source + "; + + let insert_params = HashMap::from([ + ("target_dataset_name", target_dataset_name), + ("mappings", mappings_str.as_str()), + ]); + + // Execute both inserts in parallel + let (chat_result, json_result) = try_join!( + self.run_query_synchronous(chat_clone_query.to_string(), &insert_params), + self.run_query_synchronous(json_clone_query.to_string(), &insert_params) + )?; + drop(chat_result); + drop(json_result); + + // Verify which new_ids were actually created + let new_ids_str = format!( + "[{}]", + mappings + .iter() + .map(|(_, new)| format!("'{new}'")) + .collect::>() + .join(",") + ); + + let verify_query = r" + SELECT id FROM ( + SELECT id FROM ChatInferenceDatapoint FINAL + WHERE id IN ({new_ids: Array(UUID)}) AND staled_at IS NULL + UNION ALL + SELECT id FROM JsonInferenceDatapoint FINAL + WHERE id IN ({new_ids: Array(UUID)}) AND staled_at IS NULL + ) + "; + let verify_params = HashMap::from([("new_ids", new_ids_str.as_str())]); + let verify_result = self + .run_query_synchronous(verify_query.to_string(), &verify_params) + .await?; + + let created_ids: std::collections::HashSet = verify_result + .response + .lines() + .filter_map(|line| Uuid::parse_str(line.trim()).ok()) + .collect(); + + // Map results based on which new_ids were created + let results: Vec> = mappings + .iter() + .map(|(source_id, new_id)| { + if created_ids.contains(new_id) { + Some(*new_id) + } else { + tracing::warn!( + "Failed to clone datapoint (likely does not exist): {source_id}" + ); + None + } + }) + .collect(); + + Ok(results) + } } /// Converts a vec of OrderBy terms to the correct ClickHouse ORDER BY clauses. diff --git a/tensorzero-core/src/db/clickhouse/mock_clickhouse_connection_info.rs b/tensorzero-core/src/db/clickhouse/mock_clickhouse_connection_info.rs index 4b3879eb8c4..1d92bcb2234 100644 --- a/tensorzero-core/src/db/clickhouse/mock_clickhouse_connection_info.rs +++ b/tensorzero-core/src/db/clickhouse/mock_clickhouse_connection_info.rs @@ -159,6 +159,16 @@ impl DatasetQueries for MockClickHouseConnectionInfo { .delete_datapoints(dataset_name, datapoint_ids) .await } + + async fn clone_datapoints( + &self, + target_dataset_name: &str, + source_datapoint_ids: &[Uuid], + ) -> Result>, Error> { + self.dataset_queries + .clone_datapoints(target_dataset_name, source_datapoint_ids) + .await + } } impl ConfigQueries for MockClickHouseConnectionInfo { diff --git a/tensorzero-core/src/db/datasets.rs b/tensorzero-core/src/db/datasets.rs index 3c0da6c8c5b..3a82d6145fe 100644 --- a/tensorzero-core/src/db/datasets.rs +++ b/tensorzero-core/src/db/datasets.rs @@ -149,4 +149,18 @@ pub trait DatasetQueries { dataset_name: &str, datapoint_ids: Option<&[Uuid]>, ) -> Result; + + /// Clones datapoints to a target dataset, preserving all fields except id and dataset_name. + /// + /// For each source datapoint ID, generates a new UUID and attempts to clone the datapoint + /// to the target dataset. The operation handles both Chat and Json datapoints. + /// + /// Returns a Vec with the same length as `source_datapoint_ids`, where each element is: + /// - `Some(new_id)` if the source datapoint was found and cloned successfully + /// - `None` if the source datapoint doesn't exist + async fn clone_datapoints( + &self, + target_dataset_name: &str, + source_datapoint_ids: &[Uuid], + ) -> Result>, Error>; } diff --git a/tensorzero-core/src/endpoints/datasets/internal/clone_datapoints.rs b/tensorzero-core/src/endpoints/datasets/internal/clone_datapoints.rs index 4af8798ad70..73b4d5c1fae 100644 --- a/tensorzero-core/src/endpoints/datasets/internal/clone_datapoints.rs +++ b/tensorzero-core/src/endpoints/datasets/internal/clone_datapoints.rs @@ -1,10 +1,9 @@ use axum::Json; use axum::extract::{Path, State}; use serde::{Deserialize, Serialize}; -use std::collections::HashMap; use uuid::Uuid; -use crate::db::clickhouse::ClickHouseConnectionInfo; +use crate::db::datasets::DatasetQueries; use crate::error::Error; use crate::utils::gateway::{AppState, StructuredJson}; @@ -37,144 +36,12 @@ pub async fn clone_datapoints_handler( ) -> Result, Error> { validate_dataset_name(&path_params.dataset_name)?; - let new_ids = clone_datapoints( - &path_params.dataset_name, - &request.datapoint_ids, - &app_state.clickhouse_connection_info, - ) - .await?; + let new_ids = app_state + .clickhouse_connection_info + .clone_datapoints(&path_params.dataset_name, &request.datapoint_ids) + .await?; Ok(Json(CloneDatapointsResponse { datapoint_ids: new_ids, })) } - -pub async fn clone_datapoints( - target_dataset_name: &str, - datapoint_ids: &[Uuid], - clickhouse: &ClickHouseConnectionInfo, -) -> Result>, Error> { - if datapoint_ids.is_empty() { - return Ok(vec![]); - } - - // Generate all mappings from source to target IDs - let mappings: Vec<(Uuid, Uuid)> = datapoint_ids - .iter() - .map(|id| (*id, Uuid::now_v7())) - .collect(); - - // Round trip 1: Parallel INSERTs with CTE + EXCEPT pattern - let chat_clone_query = r" - INSERT INTO ChatInferenceDatapoint - WITH source AS ( - SELECT ChatInferenceDatapoint.*, mapping.new_id - FROM ChatInferenceDatapoint FINAL - INNER JOIN ( - SELECT - tupleElement(pair, 1) as old_id, - tupleElement(pair, 2) as new_id - FROM ( - SELECT arrayJoin({mappings: Array(Tuple(UUID, UUID))}) as pair - ) - ) AS mapping ON ChatInferenceDatapoint.id = mapping.old_id - WHERE ChatInferenceDatapoint.staled_at IS NULL - ) - SELECT * EXCEPT(new_id) REPLACE( - new_id AS id, - {target_dataset_name: String} AS dataset_name, - now64() AS updated_at - ) - FROM source - "; - - let json_clone_query = r" - INSERT INTO JsonInferenceDatapoint - WITH source AS ( - SELECT JsonInferenceDatapoint.*, mapping.new_id - FROM JsonInferenceDatapoint FINAL - INNER JOIN ( - SELECT - tupleElement(pair, 1) as old_id, - tupleElement(pair, 2) as new_id - FROM ( - SELECT arrayJoin({mappings: Array(Tuple(UUID, UUID))}) as pair - ) - ) AS mapping ON JsonInferenceDatapoint.id = mapping.old_id - WHERE JsonInferenceDatapoint.staled_at IS NULL - ) - SELECT * EXCEPT(new_id) REPLACE( - new_id AS id, - {target_dataset_name: String} AS dataset_name, - now64() AS updated_at - ) - FROM source - "; - - let mappings_str = format!( - "[{}]", - mappings - .iter() - .map(|(old, new)| format!("('{old}', '{new}')")) - .collect::>() - .join(",") - ); - let insert_params = HashMap::from([ - ("target_dataset_name", target_dataset_name), - ("mappings", mappings_str.as_str()), - ]); - - let chat_future = - clickhouse.run_query_synchronous(chat_clone_query.to_string(), &insert_params); - let json_future = - clickhouse.run_query_synchronous(json_clone_query.to_string(), &insert_params); - - let (chat_result, json_result) = tokio::join!(chat_future, json_future); - chat_result?; - json_result?; - - // Round trip 2: Verify which `new_ids` were actually created (in case some source datapoints don't exist) - let new_ids_str = format!( - "[{}]", - mappings - .iter() - .map(|(_, new)| format!("'{new}'")) - .collect::>() - .join(",") - ); - - let verify_query = r" - SELECT id FROM ( - SELECT id FROM ChatInferenceDatapoint FINAL - WHERE id IN ({new_ids: Array(UUID)}) AND staled_at IS NULL - UNION ALL - SELECT id FROM JsonInferenceDatapoint FINAL - WHERE id IN ({new_ids: Array(UUID)}) AND staled_at IS NULL - ) - "; - let verify_params = HashMap::from([("new_ids", new_ids_str.as_str())]); - let verify_result = clickhouse - .run_query_synchronous(verify_query.to_string(), &verify_params) - .await?; - - let created_ids: std::collections::HashSet = verify_result - .response - .lines() - .filter_map(|line| Uuid::parse_str(line.trim()).ok()) - .collect(); - - // Map results based on which `new_ids` were created - let results: Vec> = mappings - .iter() - .map(|(source_id, new_id)| { - if created_ids.contains(new_id) { - Some(*new_id) - } else { - tracing::warn!("Failed to clone datapoint (likely does not exist): {source_id}"); - None - } - }) - .collect(); - - Ok(results) -} From 6a7f17ee138a6e36416f02f9acd3c9e866639983 Mon Sep 17 00:00:00 2001 From: Gabriel Bianconi <1275491+GabrielBianconi@users.noreply.github.com> Date: Mon, 26 Jan 2026 23:11:10 -0500 Subject: [PATCH 21/46] Run pre-commit check in CI (#5854) * Run pre-commit check in CI * Fix --- .github/workflows/general.yml | 3 +++ 1 file changed, 3 insertions(+) diff --git a/.github/workflows/general.yml b/.github/workflows/general.yml index 23d7c25235a..f8238ea8370 100644 --- a/.github/workflows/general.yml +++ b/.github/workflows/general.yml @@ -641,6 +641,9 @@ jobs: - name: Check helm schema sync run: ./ci/check-helm-schema-sync.sh + - name: prettier-root + run: "uv run --with pre-commit pre-commit run 'prettier: root' --all-files" + # We don't run these two because we want to allow template files to have trailing whitespace # TODO: how do we exclude minijinja files using pre-commit in GHA? # - name: end‑of‑file‑fixer From 38b8ae295fe0d8fe2cb1e7518b8266719adf53c6 Mon Sep 17 00:00:00 2001 From: Aaron Hill Date: Mon, 26 Jan 2026 23:11:25 -0500 Subject: [PATCH 22/46] Wrap inference ClickHouse write in a new 'write_inference' span (#5782) * Wrap inference ClickHouse write in a new 'write_inference' span This will let us track how clickhouse write times contribute to request latency * Adjust test to account for write_inference span --- tensorzero-core/src/endpoints/inference.rs | 16 +++++-- tensorzero-core/tests/e2e/inference/mod.rs | 17 +++++-- tensorzero-core/tests/e2e/otel.rs | 55 +++++++++++++++++++--- 3 files changed, 75 insertions(+), 13 deletions(-) diff --git a/tensorzero-core/src/endpoints/inference.rs b/tensorzero-core/src/endpoints/inference.rs index 60380ed6aa5..60edb5954f2 100644 --- a/tensorzero-core/src/endpoints/inference.rs +++ b/tensorzero-core/src/endpoints/inference.rs @@ -21,6 +21,7 @@ use tokio::time::Instant; use tokio_stream::StreamExt; use tokio_util::task::TaskTracker; use tracing::instrument; +use tracing_futures::Instrument; use tracing_opentelemetry::OpenTelemetrySpanExt; use uuid::Uuid; @@ -766,6 +767,9 @@ async fn infer_variant(args: InferVariantArgs<'_>) -> Result) -> Result(()) - }); + }.instrument(tracing::debug_span!(parent: &parent_span, "write_inference", otel.name = "write_inference", stream = false, inference_id = %inference_id, async_writes = async_writes))); if !async_writes { write_future.await.map_err(|e| { Error::new(ErrorDetails::InternalError { @@ -1000,6 +1004,10 @@ fn create_stream( clickhouse_connection_info: ClickHouseConnectionInfo, deferred_tasks: TaskTracker, ) -> impl FusedStream> + Send { + // Capture the parent span (function_inference) so we can use it as the parent + // for write_inference later, even after function_inference has completed. + let parent_span = tracing::Span::current(); + async_stream::stream! { let mut buffer = vec![]; @@ -1124,7 +1132,7 @@ fn create_stream( } = metadata; let config = config.clone(); - let async_write = config.gateway.observability.async_writes; + let async_writes = config.gateway.observability.async_writes; let write_future = async move { let templates = Arc::clone(&config.templates); let collect_chunks_args = CollectChunksArgs { @@ -1195,8 +1203,8 @@ fn create_stream( } drop(clickhouse_connection_info); - }; - if async_write { + }.instrument(tracing::debug_span!(parent: &parent_span, "write_inference", otel.name = "write_inference", stream = true, inference_id = %inference_id, async_writes = async_writes)); + if async_writes { deferred_tasks.spawn(write_future); } else { write_future.await; diff --git a/tensorzero-core/tests/e2e/inference/mod.rs b/tensorzero-core/tests/e2e/inference/mod.rs index 4f80d878e0e..b79aef7a1bb 100644 --- a/tensorzero-core/tests/e2e/inference/mod.rs +++ b/tensorzero-core/tests/e2e/inference/mod.rs @@ -2108,11 +2108,19 @@ fn check_good_mixture_response(exporter: CapturingOtelExporter, output: ChatInfe assert_eq!(root_attr_map.get("variant_name"), None); let root_children = &spans.span_children[&root_span.span_context.span_id()]; - let [variant_span] = root_children.as_slice() else { - panic!("Expected one child span: {root_children:#?}"); + let (variant_span, write_inference_span) = { + let mut children: Vec<_> = root_children.iter().collect(); + children.sort_by_key(|s| s.name.as_ref()); + match children.as_slice() { + [variant, write] => (*variant, *write), + _ => panic!( + "Expected two child spans (variant_inference, write_inference): {root_children:#?}" + ), + } }; assert_eq!(variant_span.name, "variant_inference"); + assert_eq!(write_inference_span.name, "write_inference"); let variant_attr_map = attrs_to_map(&variant_span.attributes); assert_eq!(variant_attr_map["function_name"], "mixture_of_n".into()); assert_eq!( @@ -2162,7 +2170,10 @@ fn check_good_mixture_response(exporter: CapturingOtelExporter, output: ChatInfe }; check_dummy_model_span(variant1_model_span, &spans, "dummy::alternate", "alternate"); - assert_eq!(num_spans, 16); + assert_eq!( + num_spans, 17, + "Expected 17 spans (added write_inference span)" + ); } fn check_dummy_model_span( diff --git a/tensorzero-core/tests/e2e/otel.rs b/tensorzero-core/tests/e2e/otel.rs index b780da8361e..02f55246b01 100644 --- a/tensorzero-core/tests/e2e/otel.rs +++ b/tensorzero-core/tests/e2e/otel.rs @@ -211,11 +211,17 @@ pub async fn test_reproduce_tracing_bug() { .await .unwrap(); - // We should now see 'variant_inference' as a root span, and have the 'function_inference' span missing + // We should now see 'variant_inference' and 'write_inference' as root spans, + // and have the 'function_inference' span missing let all_spans = exporter.take_spans(); let spans = build_span_map(all_spans); - assert_eq!(spans.root_spans[0].name, "variant_inference"); - assert_eq!(spans.root_spans.len(), 1); + let mut root_span_names: Vec<_> = spans.root_spans.iter().map(|s| s.name.as_ref()).collect(); + root_span_names.sort_unstable(); + assert_eq!( + root_span_names, + vec!["variant_inference", "write_inference"], + "Expected variant_inference and write_inference as root spans (function_inference should be missing due to tracing bug)" + ); } #[tokio::test] @@ -568,8 +574,15 @@ fn check_spans( assert_eq!(tag_count, 3); let root_children = &spans.span_children[&root_span.span_context.span_id()]; - let [variant_span] = root_children.as_slice() else { - panic!("Expected one child span: {root_children:#?}"); + let (variant_span, write_inference_span) = { + let mut children: Vec<_> = root_children.iter().collect(); + children.sort_by_key(|s| s.name.as_ref()); + match children.as_slice() { + [variant, write] => (*variant, *write), + _ => panic!( + "Expected two child spans (variant_inference, write_inference): {root_children:#?}" + ), + } }; assert_eq!(variant_span.name, "variant_inference"); @@ -781,7 +794,37 @@ fn check_spans( ]) ); - assert_eq!(num_spans, 6); + // Check write_inference span + assert_eq!( + write_inference_span.name, "write_inference", + "Expected write_inference span" + ); + assert_eq!( + write_inference_span.status, + Status::Ok, + "write_inference span should have Ok status" + ); + let write_inference_attr_map = attrs_to_map(&write_inference_span.attributes); + assert_eq!( + write_inference_attr_map["stream"], + streaming.into(), + "write_inference span should have correct stream attribute" + ); + assert_eq!( + write_inference_attr_map["inference_id"], + inference_id.to_string().into(), + "write_inference span should have correct inference_id attribute" + ); + assert_eq!( + write_inference_attr_map["async_writes"], + false.into(), + "write_inference span should have async_writes=false in synchronous mode" + ); + + assert_eq!( + num_spans, 7, + "Expected 7 spans total (function_inference, variant_inference, model_inference, model_provider_inference, consume_tickets, return_tickets, write_inference)" + ); } fn remove_unstable_attrs(attrs: &mut HashMap) { From 6b1bc23f540899764a9fcbc32ad85b4e24dbae47 Mon Sep 17 00:00:00 2001 From: Gabriel Bianconi <1275491+GabrielBianconi@users.noreply.github.com> Date: Mon, 26 Jan 2026 23:11:38 -0500 Subject: [PATCH 23/46] Speed up download provider proxy cache (#5871) --- .github/workflows/merge-queue.yml | 16 ++++++++++++++++ ci/download-provider-proxy-cache.sh | 2 +- 2 files changed, 17 insertions(+), 1 deletion(-) diff --git a/.github/workflows/merge-queue.yml b/.github/workflows/merge-queue.yml index b55141eaa91..1b1f2550fb3 100644 --- a/.github/workflows/merge-queue.yml +++ b/.github/workflows/merge-queue.yml @@ -89,6 +89,14 @@ jobs: - name: Cleanup disk space run: ./ci/free-disk-space.sh + - name: Restore provider-proxy cache + if: github.event_name != 'schedule' + uses: actions/cache@5a3ec84eff668545956fd18022155c47e93e2684 + with: + path: ./ci/provider-proxy-cache/ + key: provider-proxy-cache-${{ github.run_id }} + restore-keys: provider-proxy-cache- + - name: Download provider-proxy cache # When running as a cron job, don't use the provider-proxy cache. # The cron job is used to gather information about provider flakiness. @@ -181,6 +189,14 @@ jobs: - name: Cleanup disk space run: ./ci/free-disk-space.sh + - name: Restore client-tests provider-proxy cache + if: github.event_name != 'schedule' + uses: actions/cache@5a3ec84eff668545956fd18022155c47e93e2684 + with: + path: ./ci/provider-proxy-cache/ + key: provider-proxy-cache-client-tests-${{ github.run_id }} + restore-keys: provider-proxy-cache-client-tests- + - name: Download client-tests provider-proxy cache # When running as a cron job, don't use the provider-proxy cache. # The cron job is used to gather information about provider flakiness. diff --git a/ci/download-provider-proxy-cache.sh b/ci/download-provider-proxy-cache.sh index 4c99d3e5f11..03ce9f96e87 100755 --- a/ci/download-provider-proxy-cache.sh +++ b/ci/download-provider-proxy-cache.sh @@ -2,7 +2,7 @@ set -euxo pipefail cd $(dirname $0)/provider-proxy-cache for i in {1..5}; do - if aws s3 --endpoint-url https://19918a216783f0ac9e052233569aef60.r2.cloudflarestorage.com/ sync s3://${PROVIDER_PROXY_CACHE_BUCKET} .; then + if aws s3 --endpoint-url https://19918a216783f0ac9e052233569aef60.r2.cloudflarestorage.com/ sync --only-show-errors s3://${PROVIDER_PROXY_CACHE_BUCKET} .; then break else echo "Attempt $i failed. Retrying..." From b07837268e3876eb81f4bc25b085611b4322ea23 Mon Sep 17 00:00:00 2001 From: Aaron Hill Date: Tue, 27 Jan 2026 09:43:42 -0500 Subject: [PATCH 24/46] Only cancel general.yml early when a required job has failed (#5880) --- .../cancel-merge-queue-on-job-failure.yml | 39 +++++++++++++++++-- 1 file changed, 35 insertions(+), 4 deletions(-) diff --git a/.github/workflows/cancel-merge-queue-on-job-failure.yml b/.github/workflows/cancel-merge-queue-on-job-failure.yml index 235c079e641..0f2d1790c71 100644 --- a/.github/workflows/cancel-merge-queue-on-job-failure.yml +++ b/.github/workflows/cancel-merge-queue-on-job-failure.yml @@ -15,10 +15,32 @@ jobs: runs-on: ubuntu-latest if: github.repository == 'tensorzero/tensorzero' steps: + - name: Checkout repository + uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 + with: + sparse-checkout: .github/workflows/general.yml + sparse-checkout-cone-mode: false + - name: Check running workflows on merge queue branches run: | echo "Checking for running workflows on merge queue branches..." + # Dynamically extract the 'needs' list from 'check-all-general-jobs-passed' job in general.yml + # This ensures we only cancel on failures of jobs that actually block merging + REQUIRED_JOBS=$(yq -r '.jobs["check-all-general-jobs-passed"].needs[]' .github/workflows/general.yml) + + if [ -z "$REQUIRED_JOBS" ]; then + echo "ERROR: Could not extract required jobs from general.yml" + exit 1 + fi + + echo "Required jobs (from check-all-general-jobs-passed needs):" + echo "$REQUIRED_JOBS" + echo "" + + # Build jq filter for required jobs as a JSON array + JQ_FILTER=$(echo "$REQUIRED_JOBS" | jq -R -s 'split("\n") | map(select(length > 0))') + # Get all in-progress workflow runs for general.yml workflow_runs=$(curl -s -H "Authorization: Bearer ${{ github.token }}" \ "https://api.github.com/repos/${{ github.repository }}/actions/workflows/general.yml/runs?status=in_progress&per_page=100") @@ -51,11 +73,20 @@ jobs: jobs=$(curl -s -H "Authorization: Bearer ${{ github.token }}" \ "https://api.github.com/repos/${{ github.repository }}/actions/runs/$run_id/jobs") - # Check if any job has failed - failed_jobs=$(echo "$jobs" | jq -r '.jobs[] | select(.conclusion == "failure") | "\(.name) (status: \(.status), conclusion: \(.conclusion))"') + # Check if any required job has failed + # Job names may include matrix suffixes (e.g., "lint-rust (1)"), so we check if the job name starts with any required job name + failed_jobs=$(echo "$jobs" | jq -r --argjson required "$JQ_FILTER" ' + .jobs[] | + select(.conclusion == "failure") | + select( + .name as $job_name | + any($required[]; . as $req | $job_name | startswith($req)) + ) | + "\(.name) (status: \(.status), conclusion: \(.conclusion))" + ') if [ -n "$failed_jobs" ]; then - echo "Found failed jobs in workflow run $run_id:" + echo "Found failed required jobs in workflow run $run_id:" echo "$failed_jobs" echo "" echo "Cancelling workflow run $run_id on branch $branch..." @@ -75,7 +106,7 @@ jobs: fi echo "" else - echo "No failed jobs found in workflow run $run_id" + echo "No failed required jobs found in workflow run $run_id" fi done <<< "$merge_queue_runs" From b155b73e4bbbc6c3b73adc9e6d69c3e8adefc790 Mon Sep 17 00:00:00 2001 From: simeonlee <15081451+simeonlee@users.noreply.github.com> Date: Tue, 27 Jan 2026 09:56:01 -0500 Subject: [PATCH 25/46] Add React Router streaming to dataset datapoint detail page (#5651) * Add streaming primitives for detail pages - BasicInfoLayoutSkeleton: reusable skeleton for BasicInfoLayout sections, accepts configurable rows prop (default 5) - SectionErrorContainer, SectionErrorNotice: containers for section-level error states - SectionAsyncErrorState: error component for React Router boundaries, extracts error message from async errors * Fix SectionAsyncErrorState to handle non-string error data * Add React Router streaming to dataset datapoint detail page Implements streaming pattern for instant navigation: - Loader returns promise for deferred loading - Suspense/Await wrapper for skeleton and error states - DatapointPageHeader shared across skeleton, error, and content - Uses BasicInfoLayoutSkeleton and SectionAsyncErrorState primitives --- .../$dataset_name/datapoint/$id/route.tsx | 186 ++++++++++++++---- 1 file changed, 150 insertions(+), 36 deletions(-) diff --git a/ui/app/routes/datasets/$dataset_name/datapoint/$id/route.tsx b/ui/app/routes/datasets/$dataset_name/datapoint/$id/route.tsx index d4b343d889c..18463aec944 100644 --- a/ui/app/routes/datasets/$dataset_name/datapoint/$id/route.tsx +++ b/ui/app/routes/datasets/$dataset_name/datapoint/$id/route.tsx @@ -1,6 +1,6 @@ -import { useEffect, useMemo, useState } from "react"; +import { Suspense, useEffect, useMemo, useState, type ReactNode } from "react"; import type { ActionFunctionArgs, RouteHandle } from "react-router"; -import { data, redirect, useFetcher } from "react-router"; +import { Await, data, redirect, useFetcher, useLocation } from "react-router"; import { toDatapointUrl, toDatasetUrl } from "~/utils/urls"; import { InputElement } from "~/components/input_output/InputElement"; import { ChatOutputElement } from "~/components/input_output/ChatOutputElement"; @@ -42,8 +42,100 @@ import type { Datapoint, JsonValue, } from "~/types/tensorzero"; +import { BasicInfoLayoutSkeleton } from "~/components/layout/BasicInfoLayout"; +import { Skeleton } from "~/components/ui/skeleton"; +import { SectionAsyncErrorState } from "~/components/ui/error/ErrorContentPrimitives"; + +interface DatapointPageHeaderProps { + id?: string; + datasetName?: string; + tag?: ReactNode; +} + +function DatapointPageHeader({ + id, + datasetName, + tag, +}: DatapointPageHeaderProps) { + return ( + + } + name={id} + tag={tag} + /> + ); +} + +function DatapointContentSkeleton({ + id, + datasetName, +}: { + id?: string; + datasetName?: string; +}) { + return ( + <> + + + + + + +
+ + + + +
+ + + + + + + + + + + + + + + + ); +} + +function DatapointErrorState({ + id, + datasetName, +}: { + id?: string; + datasetName?: string; +}) { + return ( + <> + + + + ); +} -// Discriminated union for type-safe output state type OutputState = | { type: "chat"; value?: ContentBlockChatOutput[] } | { type: "json"; value?: JsonInferenceOutput; outputSchema: JsonValue }; @@ -128,13 +220,11 @@ export function hasDatapointChanged(params: { originalTags, } = params; - // Check if system has changed (added, removed, or modified) const hasSystemChanged = "system" in currentInput !== "system" in originalInput || JSON.stringify(currentInput.system) !== JSON.stringify(originalInput.system); - // Check if messages changed const hasMessagesChanged = JSON.stringify(currentInput.messages) !== JSON.stringify(originalInput.messages); @@ -350,6 +440,11 @@ export const handle: RouteHandle = { ], }; +export type DatapointData = { + datapoint: Datapoint; + resolvedInput: Input; +}; + export async function loader({ params, }: { @@ -361,24 +456,24 @@ export async function loader({ status: 404, }); } - const datapoint = await getTensorZeroClient().getDatapoint(id, dataset_name); - if (!datapoint) { - throw data(`No datapoint found for ID \`${id}\`.`, { - status: 404, - }); - } - // Load file data for InputElement component - const resolvedInput = await loadFileDataForInput(datapoint.input); + const datapointDataPromise: Promise = getTensorZeroClient() + .getDatapoint(id, dataset_name) + .then(async (datapoint) => { + if (!datapoint) { + throw data(`No datapoint found for ID \`${id}\`.`, { status: 404 }); + } + const resolvedInput = await loadFileDataForInput(datapoint.input); + return { datapoint, resolvedInput }; + }); return { - datapoint, - resolvedInput, + datapointData: datapointDataPromise, }; } -export default function DatapointPage({ loaderData }: Route.ComponentProps) { - const { datapoint, resolvedInput } = loaderData; +function DatapointContent({ data }: { data: DatapointData }) { + const { datapoint, resolvedInput } = data; const [isModalOpen, setIsModalOpen] = useState(false); const [selectedVariant, setSelectedVariant] = useState(null); @@ -492,7 +587,6 @@ export default function DatapointPage({ loaderData }: Route.ComponentProps) { return; } - // Validate schema for JSON output if (output.type === "json") { const schemaValidation = validateJsonSchema(output.outputSchema); if (!schemaValidation.valid) { @@ -557,7 +651,6 @@ export default function DatapointPage({ loaderData }: Route.ComponentProps) { title: "Request Error", description: "Failed to prepare the request. Please try again.", }); - // Reset state on error setLastRequestArgs(null); setIsModalOpen(false); setSelectedVariant(null); @@ -565,7 +658,6 @@ export default function DatapointPage({ loaderData }: Route.ComponentProps) { } catch (error) { logger.error("Failed to prepare inference request:", error); - // Show user-friendly error message based on the error type let errorMessage = "Failed to prepare the request. Please try again."; if (error instanceof Error) { if (error.message.includes("Extra body is not supported")) { @@ -630,22 +722,10 @@ export default function DatapointPage({ loaderData }: Route.ComponentProps) { }; return ( - - - } - name={datapoint.id} + <> + {datapoint.is_custom && ( @@ -785,6 +865,40 @@ export default function DatapointPage({ loaderData }: Route.ComponentProps) { onRefresh={lastRequestArgs ? handleRefresh : null} /> )} + + ); +} + +export default function DatapointPage({ + loaderData, + params, +}: Route.ComponentProps) { + const { datapointData } = loaderData; + const location = useLocation(); + + return ( + + + } + > + + } + > + {(resolvedData) => } + + ); } From 29665a8c69de64ea333a1ea2d924f0e242500fa4 Mon Sep 17 00:00:00 2001 From: Shuyang Li Date: Tue, 27 Jan 2026 10:57:53 -0500 Subject: [PATCH 26/46] Migrate workflow evaluation run queries into WorkflowEvaluationQueries trait (#5832) --- .../src/db/clickhouse/test_helpers.rs | 29 +- .../clickhouse/workflow_evaluation_queries.rs | 433 +++++++++++++++++- .../src/db/workflow_evaluation_queries.rs | 46 ++ .../src/endpoints/workflow_evaluation_run.rs | 214 ++------- 4 files changed, 528 insertions(+), 194 deletions(-) diff --git a/tensorzero-core/src/db/clickhouse/test_helpers.rs b/tensorzero-core/src/db/clickhouse/test_helpers.rs index b4a857f721c..0f7097cfe70 100644 --- a/tensorzero-core/src/db/clickhouse/test_helpers.rs +++ b/tensorzero-core/src/db/clickhouse/test_helpers.rs @@ -4,21 +4,40 @@ clippy::print_stdout, clippy::unwrap_used )] +use std::collections::HashMap; + +use serde::{Deserialize, Serialize}; + use crate::config::BatchWritesConfig; use crate::db::stored_datapoint::StoredChatInferenceDatapoint; use crate::endpoints::datasets::JsonInferenceDatapoint; -use crate::endpoints::workflow_evaluation_run::{ - WorkflowEvaluationRunEpisodeRow, WorkflowEvaluationRunRow, -}; use super::ClickHouseConnectionInfo; + +/// Database row type for workflow evaluation run (used in test helpers). +#[derive(Debug, Deserialize, Serialize)] +pub struct WorkflowEvaluationRunRow { + pub run_id: Uuid, + pub variant_pins: HashMap, + pub tags: HashMap, + pub project_name: Option, + pub run_display_name: Option, +} + +/// Database row type for workflow evaluation run episode (used in test helpers). +#[derive(Debug, Deserialize, Serialize)] +pub struct WorkflowEvaluationRunEpisodeRow { + pub run_id: Uuid, + pub episode_id: Uuid, + pub variant_pins: HashMap, + pub task_name: Option, + pub tags: HashMap, +} #[cfg(feature = "e2e_tests")] use super::escape_string_for_clickhouse_literal; #[cfg(feature = "e2e_tests")] use crate::endpoints::feedback::human_feedback::StaticEvaluationHumanFeedback; use serde_json::Value; -#[cfg(feature = "e2e_tests")] -use std::collections::HashMap; use std::sync::LazyLock; use uuid::Uuid; diff --git a/tensorzero-core/src/db/clickhouse/workflow_evaluation_queries.rs b/tensorzero-core/src/db/clickhouse/workflow_evaluation_queries.rs index b403e966b44..05c00757047 100644 --- a/tensorzero-core/src/db/clickhouse/workflow_evaluation_queries.rs +++ b/tensorzero-core/src/db/clickhouse/workflow_evaluation_queries.rs @@ -5,12 +5,13 @@ use std::collections::HashMap; use async_trait::async_trait; use uuid::Uuid; -use super::ClickHouseConnectionInfo; use super::select_queries::{parse_count, parse_json_rows}; +use super::{ClickHouseConnectionInfo, escape_string_for_clickhouse_literal}; +use crate::config::snapshot::SnapshotHash; use crate::db::workflow_evaluation_queries::{ GroupedWorkflowEvaluationRunEpisodeWithFeedbackRow, WorkflowEvaluationProjectRow, WorkflowEvaluationQueries, WorkflowEvaluationRunEpisodeWithFeedbackRow, - WorkflowEvaluationRunRow, WorkflowEvaluationRunStatisticsRaw, + WorkflowEvaluationRunInfo, WorkflowEvaluationRunRow, WorkflowEvaluationRunStatisticsRaw, WorkflowEvaluationRunStatisticsRow, WorkflowEvaluationRunWithEpisodeCountRow, }; use crate::error::{Error, ErrorDetails}; @@ -681,14 +682,156 @@ impl WorkflowEvaluationQueries for ClickHouseConnectionInfo { }) }) } + + async fn insert_workflow_evaluation_run( + &self, + run_id: Uuid, + variant_pins: &HashMap, + tags: &HashMap, + project_name: Option<&str>, + run_display_name: Option<&str>, + snapshot_hash: &SnapshotHash, + ) -> Result<(), Error> { + let query = r" + INSERT INTO DynamicEvaluationRun ( + run_id_uint, + variant_pins, + tags, + project_name, + run_display_name, + snapshot_hash + ) + VALUES ( + toUInt128({run_id:UUID}), + {variant_pins:Map(String, String)}, + {tags:Map(String, String)}, + {project_name:Nullable(String)}, + {run_display_name:Nullable(String)}, + toUInt256OrNull({snapshot_hash:Nullable(String)}) + ) + "; + + let run_id_str = run_id.to_string(); + let variant_pins_str = to_map_literal(variant_pins); + let tags_str = to_map_literal(tags); + + let mut params = HashMap::new(); + params.insert("run_id", run_id_str.as_str()); + params.insert("variant_pins", variant_pins_str.as_str()); + params.insert("tags", tags_str.as_str()); + // Use \\N to indicate NULL + params.insert("project_name", project_name.unwrap_or("\\N")); + params.insert("run_display_name", run_display_name.unwrap_or("\\N")); + params.insert("snapshot_hash", &**snapshot_hash); + + self.run_query_synchronous(query.to_string(), ¶ms) + .await?; + Ok(()) + } + + async fn insert_workflow_evaluation_run_episode( + &self, + run_id: Uuid, + episode_id: Uuid, + task_name: Option<&str>, + tags: &HashMap, + snapshot_hash: &SnapshotHash, + ) -> Result<(), Error> { + let query = r" + INSERT INTO DynamicEvaluationRunEpisode + ( + run_id, + episode_id_uint, + variant_pins, + datapoint_name, -- for legacy reasons, `task_name` is stored as `datapoint_name` in the database + tags, + snapshot_hash + ) + SELECT + {run_id:UUID} AS run_id, + toUInt128({episode_id:UUID}) AS episode_id_uint, + variant_pins, + {datapoint_name:Nullable(String)} AS datapoint_name, -- for legacy reasons, `task_name` is stored as `datapoint_name` in the database + mapUpdate(tags, {tags:Map(String, String)}) AS tags, -- merge the tags in the params on top of tags in the workflow evaluation run + toUInt256OrNull({snapshot_hash:Nullable(String)}) AS snapshot_hash + FROM DynamicEvaluationRun + WHERE run_id_uint = toUInt128({run_id:UUID}) + "; + + let run_id_str = run_id.to_string(); + let episode_id_str = episode_id.to_string(); + let tags_str = to_map_literal(tags); + + let mut params = HashMap::new(); + params.insert("run_id", run_id_str.as_str()); + params.insert("episode_id", episode_id_str.as_str()); + // Use \\N to indicate NULL; for legacy reasons, stored as `datapoint_name` in the database + params.insert("datapoint_name", task_name.unwrap_or("\\N")); + params.insert("tags", tags_str.as_str()); + params.insert("snapshot_hash", &**snapshot_hash); + + self.run_query_synchronous(query.to_string(), ¶ms) + .await?; + Ok(()) + } + + async fn get_workflow_evaluation_run_by_episode_id( + &self, + episode_id: Uuid, + ) -> Result, Error> { + let query = r" + SELECT variant_pins, tags + FROM DynamicEvaluationRunEpisode + WHERE episode_id_uint = toUInt128({episode_id:UUID}) + FORMAT JSONEachRow + "; + + let episode_id_str = episode_id.to_string(); + let params = HashMap::from([("episode_id", episode_id_str.as_str())]); + + let response = self + .run_query_synchronous(query.to_string(), ¶ms) + .await?; + + if response.response.is_empty() { + return Ok(None); + } + + let info: WorkflowEvaluationRunInfo = + serde_json::from_str(&response.response).map_err(|e| { + Error::new(ErrorDetails::Serialization { + message: format!("Failed to deserialize workflow evaluation run info: {e}"), + }) + })?; + + Ok(Some(info)) + } +} + +/// Converts a HashMap to a ClickHouse map literal string. +/// Example: {"key1": "value1", "key2": "value2"} -> "{'key1':'value1','key2':'value2'}" +fn to_map_literal(map: &HashMap) -> String { + let items: Vec = map + .iter() + .map(|(k, v)| { + format!( + "'{}':'{}'", + escape_string_for_clickhouse_literal(k), + escape_string_for_clickhouse_literal(v) + ) + }) + .collect(); + format!("{{{}}}", items.join(",")) } #[cfg(test)] mod tests { + use std::collections::HashMap; use std::sync::Arc; use uuid::Uuid; + use crate::config::snapshot::SnapshotHash; use crate::db::{ clickhouse::{ ClickHouseConnectionInfo, ClickHouseResponse, ClickHouseResponseMetadata, @@ -1388,4 +1531,290 @@ mod tests { assert_eq!(count, 0); } + + #[tokio::test] + async fn test_insert_workflow_evaluation_run() { + let mut mock_clickhouse_client = MockClickHouseClient::new(); + + mock_clickhouse_client + .expect_run_query_synchronous() + .withf(|query, params| { + assert_query_contains(query, "INSERT INTO DynamicEvaluationRun"); + assert_query_contains(query, "run_id_uint"); + assert_query_contains(query, "variant_pins"); + assert_query_contains(query, "tags"); + assert_query_contains(query, "project_name"); + assert_query_contains(query, "run_display_name"); + assert_query_contains(query, "snapshot_hash"); + assert_query_contains(query, "toUInt128({run_id:UUID})"); + assert_query_contains(query, "{variant_pins:Map(String, String)}"); + assert_query_contains(query, "{tags:Map(String, String)}"); + + assert!(params.contains_key("run_id")); + assert!(params.contains_key("variant_pins")); + assert!(params.contains_key("tags")); + assert!(params.contains_key("project_name")); + assert!(params.contains_key("run_display_name")); + assert!(params.contains_key("snapshot_hash")); + true + }) + .returning(|_, _| { + Ok(ClickHouseResponse { + response: String::new(), + metadata: ClickHouseResponseMetadata { + read_rows: 0, + written_rows: 1, + }, + }) + }); + + let conn = ClickHouseConnectionInfo::new_mock(Arc::new(mock_clickhouse_client)); + + let run_id = Uuid::now_v7(); + let variant_pins = HashMap::from([("fn1".to_string(), "var1".to_string())]); + let tags = HashMap::from([("key".to_string(), "value".to_string())]); + let snapshot_hash = SnapshotHash::default(); + + let result = conn + .insert_workflow_evaluation_run( + run_id, + &variant_pins, + &tags, + Some("my_project"), + Some("My Run"), + &snapshot_hash, + ) + .await; + + assert!( + result.is_ok(), + "insert_workflow_evaluation_run should succeed" + ); + } + + #[tokio::test] + async fn test_insert_workflow_evaluation_run_with_nulls() { + let mut mock_clickhouse_client = MockClickHouseClient::new(); + + mock_clickhouse_client + .expect_run_query_synchronous() + .withf(|_query, params| { + // When project_name and display_name are None, they should be "\\N" + assert_eq!(params.get("project_name"), Some(&"\\N")); + assert_eq!(params.get("run_display_name"), Some(&"\\N")); + true + }) + .returning(|_, _| { + Ok(ClickHouseResponse { + response: String::new(), + metadata: ClickHouseResponseMetadata { + read_rows: 0, + written_rows: 1, + }, + }) + }); + + let conn = ClickHouseConnectionInfo::new_mock(Arc::new(mock_clickhouse_client)); + + let run_id = Uuid::now_v7(); + let variant_pins = HashMap::new(); + let tags = HashMap::new(); + let snapshot_hash = SnapshotHash::default(); + + let result = conn + .insert_workflow_evaluation_run( + run_id, + &variant_pins, + &tags, + None, + None, + &snapshot_hash, + ) + .await; + + assert!( + result.is_ok(), + "insert_workflow_evaluation_run with nulls should succeed" + ); + } + + #[tokio::test] + async fn test_insert_workflow_evaluation_run_episode() { + let mut mock_clickhouse_client = MockClickHouseClient::new(); + + mock_clickhouse_client + .expect_run_query_synchronous() + .withf(|query, params| { + assert_query_contains(query, "INSERT INTO DynamicEvaluationRunEpisode"); + assert_query_contains(query, "run_id"); + assert_query_contains(query, "episode_id_uint"); + assert_query_contains(query, "variant_pins"); + assert_query_contains(query, "datapoint_name"); + assert_query_contains(query, "tags"); + assert_query_contains(query, "snapshot_hash"); + assert_query_contains(query, "SELECT"); + assert_query_contains(query, "FROM DynamicEvaluationRun"); + assert_query_contains(query, "mapUpdate(tags, {tags:Map(String, String)})"); + + assert!(params.contains_key("run_id")); + assert!(params.contains_key("episode_id")); + assert!(params.contains_key("datapoint_name")); + assert!(params.contains_key("tags")); + assert!(params.contains_key("snapshot_hash")); + true + }) + .returning(|_, _| { + Ok(ClickHouseResponse { + response: String::new(), + metadata: ClickHouseResponseMetadata { + read_rows: 1, + written_rows: 1, + }, + }) + }); + + let conn = ClickHouseConnectionInfo::new_mock(Arc::new(mock_clickhouse_client)); + + let run_id = Uuid::now_v7(); + let episode_id = Uuid::now_v7(); + let tags = HashMap::from([("episode_tag".to_string(), "value".to_string())]); + let snapshot_hash = SnapshotHash::default(); + + let result = conn + .insert_workflow_evaluation_run_episode( + run_id, + episode_id, + Some("my_task"), + &tags, + &snapshot_hash, + ) + .await; + + assert!( + result.is_ok(), + "insert_workflow_evaluation_run_episode should succeed" + ); + } + + #[tokio::test] + async fn test_insert_workflow_evaluation_run_episode_with_null_task_name() { + let mut mock_clickhouse_client = MockClickHouseClient::new(); + + mock_clickhouse_client + .expect_run_query_synchronous() + .withf(|_query, params| { + // When task_name is None, datapoint_name should be "\\N" + assert_eq!(params.get("datapoint_name"), Some(&"\\N")); + true + }) + .returning(|_, _| { + Ok(ClickHouseResponse { + response: String::new(), + metadata: ClickHouseResponseMetadata { + read_rows: 1, + written_rows: 1, + }, + }) + }); + + let conn = ClickHouseConnectionInfo::new_mock(Arc::new(mock_clickhouse_client)); + + let run_id = Uuid::now_v7(); + let episode_id = Uuid::now_v7(); + let tags = HashMap::new(); + let snapshot_hash = SnapshotHash::default(); + + let result = conn + .insert_workflow_evaluation_run_episode(run_id, episode_id, None, &tags, &snapshot_hash) + .await; + + assert!( + result.is_ok(), + "insert_workflow_evaluation_run_episode with null task_name should succeed" + ); + } + + #[tokio::test] + async fn test_get_workflow_evaluation_run_by_episode_id() { + let mut mock_clickhouse_client = MockClickHouseClient::new(); + let episode_id = Uuid::parse_str("01968d04-142c-7e53-8ea7-3a3255b518dc").unwrap(); + + mock_clickhouse_client + .expect_run_query_synchronous() + .withf(move |query, params| { + assert_query_contains(query, "SELECT variant_pins, tags"); + assert_query_contains(query, "FROM DynamicEvaluationRunEpisode"); + assert_query_contains( + query, + "WHERE episode_id_uint = toUInt128({episode_id:UUID})", + ); + assert_query_contains(query, "FORMAT JSONEachRow"); + + assert_eq!( + params.get("episode_id"), + Some(&"01968d04-142c-7e53-8ea7-3a3255b518dc") + ); + true + }) + .returning(|_, _| { + Ok(ClickHouseResponse { + response: r#"{"variant_pins":{"fn1":"var1"},"tags":{"key":"value"}}"# + .to_string(), + metadata: ClickHouseResponseMetadata { + read_rows: 1, + written_rows: 0, + }, + }) + }); + + let conn = ClickHouseConnectionInfo::new_mock(Arc::new(mock_clickhouse_client)); + + let result = conn + .get_workflow_evaluation_run_by_episode_id(episode_id) + .await + .unwrap(); + + assert!(result.is_some(), "should return Some for existing episode"); + let info = result.unwrap(); + assert_eq!( + info.variant_pins.get("fn1"), + Some(&"var1".to_string()), + "variant_pins should contain fn1 -> var1" + ); + assert_eq!( + info.tags.get("key"), + Some(&"value".to_string()), + "tags should contain key -> value" + ); + } + + #[tokio::test] + async fn test_get_workflow_evaluation_run_by_episode_id_not_found() { + let mut mock_clickhouse_client = MockClickHouseClient::new(); + + mock_clickhouse_client + .expect_run_query_synchronous() + .returning(|_, _| { + Ok(ClickHouseResponse { + response: String::new(), + metadata: ClickHouseResponseMetadata { + read_rows: 0, + written_rows: 0, + }, + }) + }); + + let conn = ClickHouseConnectionInfo::new_mock(Arc::new(mock_clickhouse_client)); + + let episode_id = Uuid::now_v7(); + let result = conn + .get_workflow_evaluation_run_by_episode_id(episode_id) + .await + .unwrap(); + + assert!( + result.is_none(), + "should return None for non-existent episode" + ); + } } diff --git a/tensorzero-core/src/db/workflow_evaluation_queries.rs b/tensorzero-core/src/db/workflow_evaluation_queries.rs index e971c6b050a..0c1eae529c1 100644 --- a/tensorzero-core/src/db/workflow_evaluation_queries.rs +++ b/tensorzero-core/src/db/workflow_evaluation_queries.rs @@ -13,8 +13,16 @@ use uuid::Uuid; #[cfg(feature = "ts-bindings")] use ts_rs::TS; +use crate::config::snapshot::SnapshotHash; use crate::error::Error; +/// Info returned when looking up a workflow evaluation run by episode ID. +#[derive(Debug, Clone, Deserialize, Serialize)] +pub struct WorkflowEvaluationRunInfo { + pub variant_pins: HashMap, + pub tags: HashMap, +} + /// Database struct for deserializing workflow evaluation project info from ClickHouse. #[derive(Debug, Deserialize)] pub struct WorkflowEvaluationProjectRow { @@ -158,6 +166,44 @@ pub trait WorkflowEvaluationQueries { /// Counts the total number of episodes for a specific workflow evaluation run. async fn count_workflow_evaluation_run_episodes(&self, run_id: Uuid) -> Result; + + /// Inserts a new workflow evaluation run. + /// + /// Note: The table is named `DynamicEvaluationRun` for historical reasons, + /// but this feature is now called "Workflow Evaluations". + async fn insert_workflow_evaluation_run( + &self, + run_id: Uuid, + variant_pins: &HashMap, + tags: &HashMap, + project_name: Option<&str>, + run_display_name: Option<&str>, + snapshot_hash: &SnapshotHash, + ) -> Result<(), Error>; + + /// Inserts a new workflow evaluation run episode. + /// + /// This copies variant_pins from the parent run and merges the provided tags + /// on top of the run's tags. + /// + /// Note: The table is named `DynamicEvaluationRunEpisode` for historical reasons, + /// but this feature is now called "Workflow Evaluations". + async fn insert_workflow_evaluation_run_episode( + &self, + run_id: Uuid, + episode_id: Uuid, + task_name: Option<&str>, + tags: &HashMap, + snapshot_hash: &SnapshotHash, + ) -> Result<(), Error>; + + /// Looks up a workflow evaluation run by episode ID. + /// + /// Returns the variant_pins and tags for the run associated with the given episode. + async fn get_workflow_evaluation_run_by_episode_id( + &self, + episode_id: Uuid, + ) -> Result, Error>; } /// A single workflow evaluation run episode with its associated feedback. diff --git a/tensorzero-core/src/endpoints/workflow_evaluation_run.rs b/tensorzero-core/src/endpoints/workflow_evaluation_run.rs index 249e481fd78..b8e41e707b0 100644 --- a/tensorzero-core/src/endpoints/workflow_evaluation_run.rs +++ b/tensorzero-core/src/endpoints/workflow_evaluation_run.rs @@ -8,8 +8,11 @@ use serde::{Deserialize, Serialize}; use uuid::Uuid; use crate::{ - config::{Config, snapshot::SnapshotHash}, - db::clickhouse::{ClickHouseConnectionInfo, escape_string_for_clickhouse_literal}, + config::Config, + db::{ + clickhouse::ClickHouseConnectionInfo, + workflow_evaluation_queries::WorkflowEvaluationQueries, + }, endpoints::validate_tags, error::{Error, ErrorDetails}, utils::{ @@ -21,12 +24,6 @@ use crate::{ }, }; -#[derive(Debug, Deserialize, Serialize)] -pub struct WorkflowEvaluationRunInfo { - pub variant_pins: HashMap, - pub tags: HashMap, -} - #[derive(Debug, Serialize, Deserialize)] pub struct WorkflowEvaluationRunParams { pub variants: HashMap, @@ -65,16 +62,16 @@ pub async fn workflow_evaluation_run( validate_tags(¶ms.tags, params.internal)?; validate_variant_pins(¶ms.variants, &config)?; let run_id = Uuid::now_v7(); - write_workflow_evaluation_run( - clickhouse_connection_info, - run_id, - params.variants, - params.tags, - params.project_name, - params.display_name, - &config.hash, - ) - .await?; + clickhouse_connection_info + .insert_workflow_evaluation_run( + run_id, + ¶ms.variants, + ¶ms.tags, + params.project_name.as_deref(), + params.display_name.as_deref(), + &config.hash, + ) + .await?; Ok(WorkflowEvaluationRunResponse { run_id }) } @@ -140,15 +137,15 @@ pub async fn workflow_evaluation_run_episode( "tensorzero::workflow_evaluation_run_id".to_string(), run_id_str, ); - write_workflow_evaluation_run_episode( - &clickhouse_connection_info, - params.task_name.as_deref(), - tags, - run_id, - episode_id, - &config.hash, - ) - .await?; + clickhouse_connection_info + .insert_workflow_evaluation_run_episode( + run_id, + episode_id, + params.task_name.as_deref(), + &tags, + &config.hash, + ) + .await?; Ok(WorkflowEvaluationRunEpisodeResponse { episode_id }) } @@ -172,138 +169,6 @@ pub fn validate_variant_pins( Ok(()) } -fn to_map_literal(map: &HashMap) -> String { - let items: Vec = map - .iter() - .map(|(k, v)| { - format!( - "'{}':'{}'", - escape_string_for_clickhouse_literal(k), - escape_string_for_clickhouse_literal(v) - ) - }) - .collect(); - format!("{{{}}}", items.join(",")) -} - -#[derive(Debug, Deserialize, Serialize)] -pub struct WorkflowEvaluationRunRow { - pub run_id: Uuid, - pub variant_pins: HashMap, - pub tags: HashMap, - pub project_name: Option, - pub run_display_name: Option, -} - -/// Writes a workflow evaluation run to the database. -/// -/// Note: The table is named `DynamicEvaluationRun` for historical reasons, -/// but this feature is now called "Workflow Evaluations". -/// -/// Similarly, we write both `tensorzero::dynamic_evaluation_run_id` (old) and -/// `tensorzero::workflow_evaluation_run_id` (new) tags to episode inferences -/// to support backward compatibility during the migration period. -async fn write_workflow_evaluation_run( - clickhouse: ClickHouseConnectionInfo, - run_id: Uuid, - variant_pins: HashMap, - tags: HashMap, - project_name: Option, - run_display_name: Option, - snapshot_hash: &SnapshotHash, -) -> Result<(), Error> { - let query = r" - INSERT INTO DynamicEvaluationRun ( - run_id_uint, - variant_pins, - tags, - project_name, - run_display_name, - snapshot_hash - ) - VALUES ( - toUInt128({run_id:UUID}), - {variant_pins:Map(String, String)}, - {tags:Map(String, String)}, - {project_name:Nullable(String)}, - {run_display_name:Nullable(String)}, - toUInt256OrNull({snapshot_hash:Nullable(String)}) - ) - "; - let mut params = HashMap::new(); - let variant_pins_str = to_map_literal(&variant_pins); - let tags_str = to_map_literal(&tags); - let run_id_str = run_id.to_string(); - params.insert("run_id", run_id_str.as_str()); - params.insert("variant_pins", variant_pins_str.as_str()); - params.insert("tags", tags_str.as_str()); - params.insert("project_name", project_name.as_deref().unwrap_or("\\N")); // Use \\N to indicate NULL - params.insert( - "run_display_name", - run_display_name.as_deref().unwrap_or("\\N"), - ); // Use \\N to indicate NULL - params.insert("snapshot_hash", snapshot_hash); - clickhouse - .run_query_synchronous(query.to_string(), ¶ms) - .await?; - Ok(()) -} - -#[derive(Debug, Deserialize, Serialize)] -pub struct WorkflowEvaluationRunEpisodeRow { - pub run_id: Uuid, - pub episode_id: Uuid, - pub variant_pins: HashMap, - pub task_name: Option, // For legacy reasons, stored as `task_name` in the database - pub tags: HashMap, -} - -/// Writes a workflow evaluation run episode to the database. -/// Note: The table is named `DynamicEvaluationRunEpisode` for historical reasons, -/// but this feature is now called "Workflow Evaluations". -async fn write_workflow_evaluation_run_episode( - clickhouse: &ClickHouseConnectionInfo, - task_name: Option<&str>, - tags: HashMap, - run_id: Uuid, - episode_id: Uuid, - snapshot_hash: &SnapshotHash, -) -> Result<(), Error> { - let query = r" - INSERT INTO DynamicEvaluationRunEpisode - ( - run_id, - episode_id_uint, - variant_pins, - datapoint_name, -- for legacy reasons, `task_name` is stored as `datapoint_name` in the database - tags, - snapshot_hash - ) - SELECT - {run_id:UUID} AS run_id, - toUInt128({episode_id:UUID}) AS episode_id_uint, - variant_pins, - {datapoint_name:Nullable(String)} AS datapoint_name, -- for legacy reasons, `task_name` is stored as `datapoint_name` in the database - mapUpdate(tags, {tags:Map(String, String)}) AS tags, -- merge the tags in the params on top of tags in the workflow evaluation run - toUInt256OrNull({snapshot_hash:Nullable(String)}) AS snapshot_hash - FROM DynamicEvaluationRun - WHERE run_id_uint = toUInt128({run_id:UUID}) - "; - let mut query_params = HashMap::new(); - let run_id_str = run_id.to_string(); - let episode_id_str = episode_id.to_string(); - query_params.insert("run_id", run_id_str.as_str()); - query_params.insert("episode_id", episode_id_str.as_str()); - query_params.insert("datapoint_name", task_name.unwrap_or("\\N")); // Use \\N to indicate NULL; for legacy reasons, stored as `datapoint_name` in the database - let tags_str = to_map_literal(&tags); - query_params.insert("tags", tags_str.as_str()); - query_params.insert("snapshot_hash", &**snapshot_hash); - clickhouse - .run_query_synchronous(query.to_string(), &query_params) - .await?; - Ok(()) -} - /// For workflow evaluation runs, we generate episode IDs that are WORKFLOW_EVALUATION_OFFSET in the future. /// If we come across an episode ID that is at least WORKFLOW_EVALUATION_THRESHOLD, we need to look up the /// appropriate workflow evaluation run and then apply the `variant_name` if unset and the tags if unset. @@ -326,7 +191,9 @@ pub async fn validate_inference_episode_id_and_apply_workflow_evaluation_run( if compare_timestamps(episode_id_timestamp, WORKFLOW_EVALUATION_THRESHOLD) { return validate_tensorzero_uuid(episode_id, "Episode"); } - let workflow_evaluation_run = lookup_workflow_evaluation_run(clickhouse, episode_id).await?; + let workflow_evaluation_run = clickhouse + .get_workflow_evaluation_run_by_episode_id(episode_id) + .await?; let Some(workflow_evaluation_run) = workflow_evaluation_run else { return Err(Error::new(ErrorDetails::InvalidWorkflowEvaluationRun { episode_id, @@ -360,33 +227,6 @@ pub async fn validate_inference_episode_id_and_apply_workflow_evaluation_run( Ok(()) } -/// Looks up a workflow evaluation run from the database. -/// Note: The table is named `DynamicEvaluationRunEpisode` for historical reasons, -/// but this feature is now called "Workflow Evaluations". -async fn lookup_workflow_evaluation_run( - clickhouse: &ClickHouseConnectionInfo, - episode_id: Uuid, -) -> Result, Error> { - let query = r" - SELECT variant_pins, tags FROM DynamicEvaluationRunEpisode WHERE episode_id_uint = toUInt128({episode_id:UUID}) FORMAT JSONEachRow - "; - let episode_id_str = episode_id.to_string(); - let params = HashMap::from([("episode_id", episode_id_str.as_str())]); - let result = clickhouse - .run_query_synchronous(query.to_string(), ¶ms) - .await?; - if result.response.is_empty() { - return Ok(None); - } - let workflow_evaluation_run: WorkflowEvaluationRunInfo = serde_json::from_str(&result.response) - .map_err(|_| { - Error::new(ErrorDetails::Serialization { - message: "Failed to deserialize workflow evaluation run".to_string(), - }) - })?; - Ok(Some(workflow_evaluation_run)) -} - // ============================================================================ // DEPRECATED HANDLERS - For backward compatibility // ============================================================================ From be595a9c0ad7d8a4a40fe44fb76b07ee7aea7b8a Mon Sep 17 00:00:00 2001 From: Viraj Mehta Date: Tue, 27 Jan 2026 11:05:58 -0500 Subject: [PATCH 27/46] bumped version to 2026.1.6 (#5885) --- Cargo.lock | 36 +++++++++---------- Cargo.toml | 2 +- .../production-deployment-k8s-helm/Chart.yaml | 2 +- ui/package.json | 2 +- 4 files changed, 21 insertions(+), 21 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index ac3edef1c5f..4db1c9649c8 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -315,7 +315,7 @@ checksum = "c08606f8c3cbf4ce6ec8e28fb0014a2c086708fe954eaa885384a6165172e7e8" [[package]] name = "autopilot-client" -version = "2026.1.5" +version = "2026.1.6" dependencies = [ "chrono", "durable-tools-spawn", @@ -340,7 +340,7 @@ dependencies = [ [[package]] name = "autopilot-tools" -version = "2026.1.5" +version = "2026.1.6" dependencies = [ "anyhow", "async-trait", @@ -364,7 +364,7 @@ dependencies = [ [[package]] name = "autopilot-worker" -version = "2026.1.5" +version = "2026.1.6" dependencies = [ "anyhow", "async-trait", @@ -1847,7 +1847,7 @@ dependencies = [ [[package]] name = "durable-tools" -version = "2026.1.5" +version = "2026.1.6" dependencies = [ "anyhow", "async-trait", @@ -1877,7 +1877,7 @@ dependencies = [ [[package]] name = "durable-tools-spawn" -version = "2026.1.5" +version = "2026.1.6" dependencies = [ "async-trait", "durable", @@ -2009,7 +2009,7 @@ dependencies = [ [[package]] name = "evaluations" -version = "2026.1.5" +version = "2026.1.6" dependencies = [ "anyhow", "async-trait", @@ -2112,7 +2112,7 @@ dependencies = [ [[package]] name = "feedback-load-test" -version = "2026.1.5" +version = "2026.1.6" dependencies = [ "anyhow", "async-trait", @@ -2358,7 +2358,7 @@ dependencies = [ [[package]] name = "gateway" -version = "2026.1.5" +version = "2026.1.6" dependencies = [ "async-stream", "autopilot-tools", @@ -3587,7 +3587,7 @@ dependencies = [ [[package]] name = "minijinja-utils" -version = "2026.1.5" +version = "2026.1.6" dependencies = [ "minijinja", "serde", @@ -4779,7 +4779,7 @@ dependencies = [ [[package]] name = "rate-limit-load-test" -version = "2026.1.5" +version = "2026.1.6" dependencies = [ "anyhow", "async-trait", @@ -6234,7 +6234,7 @@ dependencies = [ [[package]] name = "tensorzero-auth" -version = "2026.1.5" +version = "2026.1.6" dependencies = [ "axum", "chrono", @@ -6254,7 +6254,7 @@ dependencies = [ [[package]] name = "tensorzero-core" -version = "2026.1.5" +version = "2026.1.6" dependencies = [ "arc-swap", "async-stream", @@ -6375,7 +6375,7 @@ dependencies = [ [[package]] name = "tensorzero-node" -version = "2026.1.5" +version = "2026.1.6" dependencies = [ "napi", "napi-build", @@ -6388,7 +6388,7 @@ dependencies = [ [[package]] name = "tensorzero-optimizers" -version = "2026.1.5" +version = "2026.1.6" dependencies = [ "axum", "base64 0.22.1", @@ -6422,7 +6422,7 @@ dependencies = [ [[package]] name = "tensorzero-python" -version = "2026.1.5" +version = "2026.1.6" dependencies = [ "evaluations", "futures", @@ -6440,7 +6440,7 @@ dependencies = [ [[package]] name = "tensorzero-types" -version = "2026.1.5" +version = "2026.1.6" dependencies = [ "aws-smithy-types", "infer", @@ -6460,7 +6460,7 @@ dependencies = [ [[package]] name = "tensorzero-types-providers" -version = "2026.1.5" +version = "2026.1.6" dependencies = [ "serde", "serde_json", @@ -6468,7 +6468,7 @@ dependencies = [ [[package]] name = "tensorzero-unsafe-helpers" -version = "2026.1.5" +version = "2026.1.6" [[package]] name = "termcolor" diff --git a/Cargo.toml b/Cargo.toml index 7e3930c6052..3038e72245e 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -26,7 +26,7 @@ members = [ resolver = "2" [workspace.package] -version = "2026.1.5" +version = "2026.1.6" rust-version = "1.88.0" license = "Apache-2.0" edition = "2024" diff --git a/examples/production-deployment-k8s-helm/Chart.yaml b/examples/production-deployment-k8s-helm/Chart.yaml index 4dc952329fd..055058fbba6 100644 --- a/examples/production-deployment-k8s-helm/Chart.yaml +++ b/examples/production-deployment-k8s-helm/Chart.yaml @@ -3,4 +3,4 @@ name: tensorzero description: A Helm chart for Kubernetes with deployment, secret, configmap, service, and ingress type: application version: 0.0.0 # updated by CI -appVersion: "2026.1.5" +appVersion: "2026.1.6" diff --git a/ui/package.json b/ui/package.json index 5ac5fc57129..dfb2523e17c 100644 --- a/ui/package.json +++ b/ui/package.json @@ -2,7 +2,7 @@ "name": "tensorzero-ui", "private": true, "type": "module", - "version": "2026.1.5", + "version": "2026.1.6", "scripts": { "build": "NODE_ENV=production react-router build", "dev": "react-router dev", From ec65984f36626171d3fa69ffaecf59c22e4cbad9 Mon Sep 17 00:00:00 2001 From: simeonlee <15081451+simeonlee@users.noreply.github.com> Date: Tue, 27 Jan 2026 11:07:11 -0500 Subject: [PATCH 28/46] Reorganize UI sidebar and overview navigation (#5791) * Reorganize UI sidebar navigation - Rename "Dashboard" to "Overview" with LayoutGrid icon - Promote "Evaluations" to top-level category (Inference Evaluations, Workflow Evaluations) - Remove confusing "Workflows" category - Move API Keys and Docs to sidebar footer - Rebuild sidebar collapse button using SidebarMenuButton for consistent hover states - Add subheading prop to PageHeader component - Add subheading to API Keys page explaining their purpose - Rename "Documentation" to "Docs" in sidebar * Remove subheading prop from PageHeader * Update Dashboard references to Overview * Improve sidebar footer: narrower collapse button, show gateway status when collapsed * Add text fade transitions, fix accessibility and remove dead code * Improve sidebar transitions: height on labels, no truncate, fade on status text * Add fade transition to TensorZero logo text * Update Overview page structure to match sidebar * Add Autopilot tile to Overview page when available * Restructure Overview: Autopilot/Playground as top row, 4-column grid below * Restore 3-column grid for main sections * Finalize sidebar and Overview page reorganization - Move Playground to Tools section in sidebar - Update Overview page layout: 3-col grid for main sections, separate row for Data/Tools - Update tile hover styles: orange-200 border, orange-600 title - Clean up unused imports * Move external link arrow closer to Docs text * Polish Overview page footer links: smaller text, orange hover, rename Docs * Show version mismatch warning in collapsed sidebar tooltip * Restore version mismatch tooltip with explanation * Fix nested tooltips when collapsed with version mismatch * Fix version mismatch layout: show below status text, simplify tooltip * Reorganize sidebar and overview grid - Merge Data and Tools sections into Resources (Playground, Datasets, API Keys) - Remove Docs and API Keys from sidebar footer - Stack Evaluations + Optimization in column 2, Resources as column 3 - Add API Keys to Resources in overview grid - Align Optimization to bottom of column 2 - Use consistent gap spacing on mobile * Remove placeholder API Keys description * Add hover highlight to overview directory cards * Use card-highlight-icon-bg for hover border * Add card-highlight-border semantic token * Clean up unused numEvaluationRunsDesc and comments * Add orange-500 hover color to directory card subtitle --- ui/app/components/layout/ContentLayout.tsx | 2 +- .../layout/TensorZeroStatusIndicator.tsx | 91 +++++++---- ui/app/components/layout/app.sidebar.tsx | 154 ++++++++---------- ui/app/routes/index.tsx | 88 ++++++---- ui/app/tailwind.css | 6 + ui/e2e_tests/error-boundaries.spec.ts | 6 +- 6 files changed, 193 insertions(+), 154 deletions(-) diff --git a/ui/app/components/layout/ContentLayout.tsx b/ui/app/components/layout/ContentLayout.tsx index 0415728ccbc..65c4c9bbed0 100644 --- a/ui/app/components/layout/ContentLayout.tsx +++ b/ui/app/components/layout/ContentLayout.tsx @@ -5,7 +5,7 @@ export function ContentLayout({ children }: React.PropsWithChildren) { const pageTitle = segments.length > 0 ? [...segments.map((b) => b.label), "TensorZero"].join(" • ") - : "Dashboard • TensorZero"; + : "Overview • TensorZero"; return ( <> diff --git a/ui/app/components/layout/TensorZeroStatusIndicator.tsx b/ui/app/components/layout/TensorZeroStatusIndicator.tsx index df5528989a8..c6fb05ff0f7 100644 --- a/ui/app/components/layout/TensorZeroStatusIndicator.tsx +++ b/ui/app/components/layout/TensorZeroStatusIndicator.tsx @@ -2,10 +2,16 @@ import { useTensorZeroStatusFetcher } from "~/routes/api/tensorzero/status"; import { Tooltip, TooltipContent, TooltipTrigger } from "../ui/tooltip"; import { useMemo } from "react"; +interface TensorZeroStatusIndicatorProps { + collapsed?: boolean; +} + /** * A component that displays the status of the TensorZero Gateway. */ -export default function TensorZeroStatusIndicator() { +export default function TensorZeroStatusIndicator({ + collapsed = false, +}: TensorZeroStatusIndicatorProps) { const { status, isLoading } = useTensorZeroStatusFetcher(); const uiVersion = __APP_VERSION__; @@ -26,37 +32,60 @@ export default function TensorZeroStatusIndicator() { return "bg-green-500"; // Everything is good }, [isLoading, status, versionsMatch]); - return ( -
-
-
-
- {isLoading - ? "Checking status..." - : status === undefined - ? "Connecting to Gateway..." - : status - ? `TensorZero Gateway ${serverVersion}` - : "Gateway Unavailable"} -
- {status && !versionsMatch && ( -
- - - - Version mismatch: UI {uiVersion} - - - - Please make sure your UI has the same version as the gateway. - Otherwise you might have compatibility issues. - - -
- )} + const statusText = isLoading + ? "Checking status..." + : status === undefined + ? "Connecting to Gateway..." + : status + ? `TensorZero Gateway ${serverVersion}` + : "Gateway Unavailable"; + + const statusDot = ( +
+
+
+ ); + + const content = ( +
+
+ {statusDot} + + {statusText} +
+ {status && !versionsMatch && !collapsed && ( +
+ + + + Version mismatch: UI {uiVersion} + + + + You may experience compatibility issues + + +
+ )}
); + + if (collapsed) { + return ( + + {content} + + {statusText} + {status && !versionsMatch && ( +
+ Version mismatch: UI {uiVersion} +
+ )} + + + ); + } + + return content; } diff --git a/ui/app/components/layout/app.sidebar.tsx b/ui/app/components/layout/app.sidebar.tsx index cd187b9cbee..2b86f55e1ee 100644 --- a/ui/app/components/layout/app.sidebar.tsx +++ b/ui/app/components/layout/app.sidebar.tsx @@ -1,27 +1,19 @@ import * as React from "react"; -import { useMemo } from "react"; import { - Home, Inferences, Episodes, Functions, SupervisedFineTuning, - Documentation, Dataset, GridCheck, SequenceChecks, Playground, Model, Chat, + SidebarCollapse, + SidebarExpand, } from "~/components/icons/Icons"; -import { KeyRound } from "lucide-react"; -import { useSidebar } from "~/components/ui/sidebar"; -import { useActivePath } from "~/hooks/use-active-path"; -import { useAutopilotAvailable } from "~/context/autopilot-available"; -import { TensorZeroLogo } from "~/components/icons/Icons"; -import { Link } from "react-router"; -import type { IconProps } from "~/components/icons/Icons"; - +import { KeyRound, LayoutGrid } from "lucide-react"; import { Sidebar, SidebarContent, @@ -30,12 +22,17 @@ import { SidebarMenuItem, SidebarMenuButton, SidebarRail, - SidebarTrigger, SidebarFooter, SidebarGroup, SidebarGroupLabel, SidebarGroupContent, + useSidebar, } from "~/components/ui/sidebar"; +import { useActivePath } from "~/hooks/use-active-path"; +import { useAutopilotAvailable } from "~/context/autopilot-available"; +import { TensorZeroLogo } from "~/components/icons/Icons"; +import { Link } from "react-router"; +import type { IconProps } from "~/components/icons/Icons"; import TensorZeroStatusIndicator from "./TensorZeroStatusIndicator"; import { ReadOnlyBadge } from "./ReadOnlyBadge"; @@ -51,16 +48,6 @@ interface NavigationSection { } const navigation: NavigationSection[] = [ - { - title: "Autopilot", - items: [ - { - title: "Sessions", - url: "/autopilot", - icon: Chat, - }, - ], - }, { title: "Observability", items: [ @@ -86,6 +73,21 @@ const navigation: NavigationSection[] = [ }, ], }, + { + title: "Evaluations", + items: [ + { + title: "Inference Evaluations", + url: "/evaluations", + icon: GridCheck, + }, + { + title: "Workflow Evaluations", + url: "/workflow-evaluations", + icon: SequenceChecks, + }, + ], + }, { title: "Optimization", items: [ @@ -97,7 +99,7 @@ const navigation: NavigationSection[] = [ ], }, { - title: "Workflows", + title: "Resources", items: [ { title: "Playground", @@ -110,22 +112,7 @@ const navigation: NavigationSection[] = [ icon: Dataset, }, { - title: "Inference Evaluations", - url: "/evaluations", - icon: GridCheck, - }, - { - title: "Workflow Evaluations", - url: "/workflow-evaluations", - icon: SequenceChecks, - }, - ], - }, - { - title: "Operations", - items: [ - { - title: "TensorZero API Keys", + title: "API Keys", url: "/api-keys", icon: KeyRound, }, @@ -134,18 +121,10 @@ const navigation: NavigationSection[] = [ ]; export function AppSidebar({ ...props }: React.ComponentProps) { - const { state } = useSidebar(); + const { state, toggleSidebar } = useSidebar(); const activePathUtils = useActivePath(); const autopilotAvailable = useAutopilotAvailable(); - // Filter out Autopilot section if not available - const filteredNavigation = useMemo(() => { - if (autopilotAvailable) { - return navigation; - } - return navigation.filter((section) => section.title !== "Autopilot"); - }, [autopilotAvailable]); - return ( @@ -157,9 +136,9 @@ export function AppSidebar({ ...props }: React.ComponentProps) { > - {state === "expanded" && ( - TensorZero - )} + + TensorZero + @@ -171,22 +150,38 @@ export function AppSidebar({ ...props }: React.ComponentProps) { - - {state === "expanded" && Dashboard} + + + Overview + + {autopilotAvailable && ( + + + + + + Autopilot + + + + + )} - {filteredNavigation.map((section) => ( + {navigation.map((section) => ( - {state === "expanded" && ( - {section.title} - )} + {section.title} {section.items?.map((item) => ( @@ -201,7 +196,9 @@ export function AppSidebar({ ...props }: React.ComponentProps) { onClick={(e) => e.stopPropagation()} > - {state === "expanded" && {item.title}} + + {item.title} + @@ -209,33 +206,24 @@ export function AppSidebar({ ...props }: React.ComponentProps) { ))} - - {state === "expanded" && Other} - - - - e.stopPropagation()} - > - - {state === "expanded" && Documentation ↗} - - - - - - {state === "expanded" && } - + + + + {state === "expanded" ? ( + + ) : ( + + )} + + diff --git a/ui/app/routes/index.tsx b/ui/app/routes/index.tsx index e0b83a841d4..e5209dd754c 100644 --- a/ui/app/routes/index.tsx +++ b/ui/app/routes/index.tsx @@ -20,6 +20,7 @@ import { Playground, Model, } from "~/components/icons/Icons"; +import { KeyRound } from "lucide-react"; import { Tooltip, TooltipContent, @@ -49,18 +50,18 @@ function DirectoryCard({ }: DirectoryCardProps) { return ( - -
+ +
-

+

{title}

-

+

{typeof description === "string" ? ( description ) : ( @@ -117,8 +118,8 @@ function FooterLink({ source, icon: Icon, children }: FooterLinkProps) { rel="noopener noreferrer" target="_blank" > - - + + {children} @@ -175,10 +176,6 @@ export async function loader() { (datasets) => `${datasets.datasets.length} datasets`, ); - const numEvaluationRunsDesc = numEvaluationRunsPromise.then( - (runs) => `evaluations, ${runs} runs`, - ); - const inferenceEvaluationsDesc = Promise.all([ configPromise, numEvaluationRunsPromise, @@ -202,7 +199,6 @@ export async function loader() { numVariantsDesc, numEpisodesDesc, numDatasetsDesc, - numEvaluationRunsDesc, inferenceEvaluationsDesc, dynamicEvaluationsDesc, numModelsUsedDesc, @@ -224,7 +220,8 @@ export default function Home({ loaderData }: Route.ComponentProps) { return (

-

Dashboard

+

Overview

+

@@ -258,22 +255,47 @@ export default function Home({ loaderData }: Route.ComponentProps) {

-
-

- Optimization -

-
- +
+
+

+ Evaluations +

+
+ + +
+
+ +
+

+ Optimization +

+
+ +
-
-

Workflows

+
+

Resources

-
@@ -315,7 +331,7 @@ export default function Home({ loaderData }: Route.ComponentProps) { source="https://www.tensorzero.com/docs" icon={Documentation} > - Documentation + Docs { ).toBeVisible(); // Sidebar should still be visible and functional - await expect(page.getByText("Dashboard")).toBeVisible(); + await expect(page.getByText("Overview")).toBeVisible(); await expect(page.getByText("Inferences")).toBeVisible(); }); @@ -24,8 +24,8 @@ test.describe("Error Boundaries", () => { // Verify we're on the error page await expect(page.getByText("Page Not Found")).toBeVisible(); - // Click on Dashboard in sidebar - await page.getByRole("link", { name: "Dashboard" }).click(); + // Click on Overview in sidebar + await page.getByRole("link", { name: "Overview" }).click(); // Should navigate to home page successfully await expect(page.getByText("Ask a question")).toBeVisible(); From 6deaa2460f7cf5cc712cb8546546b52547981580 Mon Sep 17 00:00:00 2001 From: Gabriel Bianconi <1275491+GabrielBianconi@users.noreply.github.com> Date: Tue, 27 Jan 2026 11:20:55 -0500 Subject: [PATCH 29/46] Make Playwright downloads more resilient in CI (#5890) --- .../slash-command-regen-fixtures.yml | 13 ++++++- .../ui-tests-e2e-model-inference-cache.yml | 13 ++++++- .github/workflows/ui-tests-e2e.yml | 39 +++++++++++++++++-- 3 files changed, 60 insertions(+), 5 deletions(-) diff --git a/.github/workflows/slash-command-regen-fixtures.yml b/.github/workflows/slash-command-regen-fixtures.yml index 6162f84c52c..e906fbbb642 100644 --- a/.github/workflows/slash-command-regen-fixtures.yml +++ b/.github/workflows/slash-command-regen-fixtures.yml @@ -47,7 +47,18 @@ jobs: - name: Install `pnpm` dependencies run: pnpm install --frozen-lockfile - name: Setup Playwright - run: pnpm --filter=tensorzero-ui exec playwright install --with-deps chromium + run: | + for attempt in 1 2 3; do + if pnpm --filter=tensorzero-ui exec playwright install --with-deps chromium; then + break + fi + if [ $attempt -eq 3 ]; then + echo "Failed to install Playwright after 3 attempts" + exit 1 + fi + sleep $((10 * attempt)) + done + shell: bash - name: Regenerate fixtures id: e2e_tests run: | diff --git a/.github/workflows/ui-tests-e2e-model-inference-cache.yml b/.github/workflows/ui-tests-e2e-model-inference-cache.yml index 89454df20ca..f0b4295f3f0 100644 --- a/.github/workflows/ui-tests-e2e-model-inference-cache.yml +++ b/.github/workflows/ui-tests-e2e-model-inference-cache.yml @@ -83,7 +83,18 @@ jobs: run: pnpm install --frozen-lockfile - name: Setup Playwright - run: pnpm --filter=tensorzero-ui exec playwright install --with-deps chromium + run: | + for attempt in 1 2 3; do + if pnpm --filter=tensorzero-ui exec playwright install --with-deps chromium; then + break + fi + if [ $attempt -eq 3 ]; then + echo "Failed to install Playwright after 3 attempts" + exit 1 + fi + sleep $((10 * attempt)) + done + shell: bash - name: Download container images uses: actions/download-artifact@018cc2cf5baa6db3ef3c5f8a56943fffe632ef53 diff --git a/.github/workflows/ui-tests-e2e.yml b/.github/workflows/ui-tests-e2e.yml index 57e7deee10e..8e5c6c0a5c5 100644 --- a/.github/workflows/ui-tests-e2e.yml +++ b/.github/workflows/ui-tests-e2e.yml @@ -111,7 +111,18 @@ jobs: run: pnpm install --frozen-lockfile - name: Setup Playwright - run: pnpm --filter=tensorzero-ui exec playwright install --with-deps chromium + run: | + for attempt in 1 2 3; do + if pnpm --filter=tensorzero-ui exec playwright install --with-deps chromium; then + break + fi + if [ $attempt -eq 3 ]; then + echo "Failed to install Playwright after 3 attempts" + exit 1 + fi + sleep $((10 * attempt)) + done + shell: bash - name: Download container images uses: actions/download-artifact@018cc2cf5baa6db3ef3c5f8a56943fffe632ef53 @@ -226,7 +237,18 @@ jobs: run: pnpm install --frozen-lockfile - name: Setup Playwright - run: pnpm --filter=tensorzero-ui exec playwright install --with-deps chromium + run: | + for attempt in 1 2 3; do + if pnpm --filter=tensorzero-ui exec playwright install --with-deps chromium; then + break + fi + if [ $attempt -eq 3 ]; then + echo "Failed to install Playwright after 3 attempts" + exit 1 + fi + sleep $((10 * attempt)) + done + shell: bash - name: Download container images uses: actions/download-artifact@018cc2cf5baa6db3ef3c5f8a56943fffe632ef53 @@ -314,7 +336,18 @@ jobs: run: pnpm install --frozen-lockfile - name: Setup Playwright - run: pnpm --filter=tensorzero-ui exec playwright install --with-deps chromium + run: | + for attempt in 1 2 3; do + if pnpm --filter=tensorzero-ui exec playwright install --with-deps chromium; then + break + fi + if [ $attempt -eq 3 ]; then + echo "Failed to install Playwright after 3 attempts" + exit 1 + fi + sleep $((10 * attempt)) + done + shell: bash - name: Download container images uses: actions/download-artifact@018cc2cf5baa6db3ef3c5f8a56943fffe632ef53 From 673c3232b4a8f1d31797ffc9bdc680f561e0142f Mon Sep 17 00:00:00 2001 From: Gabriel Bianconi <1275491+GabrielBianconi@users.noreply.github.com> Date: Tue, 27 Jan 2026 11:23:04 -0500 Subject: [PATCH 30/46] Fix nightly clippy (#5895) --- tensorzero-core/tests/e2e/db/bandit_queries.rs | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/tensorzero-core/tests/e2e/db/bandit_queries.rs b/tensorzero-core/tests/e2e/db/bandit_queries.rs index 05556ff28b0..5811c09e0fa 100644 --- a/tensorzero-core/tests/e2e/db/bandit_queries.rs +++ b/tensorzero-core/tests/e2e/db/bandit_queries.rs @@ -52,7 +52,7 @@ async fn test_clickhouse_metrics_by_variant_many_results() { assert_eq!(metrics_by_variant.len(), 3); // Sort by count in descending order for deterministic results let mut metrics_by_variant = metrics_by_variant; - metrics_by_variant.sort_by(|a, b| b.count.cmp(&a.count)); + metrics_by_variant.sort_by_key(|b| std::cmp::Reverse(b.count)); let metric = metrics_by_variant.first().unwrap(); assert_eq!(metric.variant_name, "dicl"); assert_eq!(metric.count, 39); @@ -81,7 +81,7 @@ async fn test_clickhouse_metrics_by_variant_episode_boolean() { .unwrap(); // Sort by count in descending order for deterministic results let mut metrics_by_variant = metrics_by_variant; - metrics_by_variant.sort_by(|a, b| b.count.cmp(&a.count)); + metrics_by_variant.sort_by_key(|b| std::cmp::Reverse(b.count)); println!("metrics_by_variant: {metrics_by_variant:?}"); assert_eq!(metrics_by_variant.len(), 3); let metric = metrics_by_variant.first().unwrap(); @@ -112,7 +112,7 @@ async fn test_clickhouse_metrics_by_variant_episode_float() { .unwrap(); // Sort by count in descending order for deterministic results let mut metrics_by_variant = metrics_by_variant; - metrics_by_variant.sort_by(|a, b| b.count.cmp(&a.count)); + metrics_by_variant.sort_by_key(|b| std::cmp::Reverse(b.count)); println!("metrics_by_variant: {metrics_by_variant:?}"); assert_eq!(metrics_by_variant.len(), 3); let metric = metrics_by_variant.first().unwrap(); @@ -893,7 +893,7 @@ async fn test_clickhouse_get_cumulative_feedback_timeseries_baseline_continuity( .collect(); // Sort by period - variant_points.sort_by(|a, b| a.period_end.cmp(&b.period_end)); + variant_points.sort_by_key(|a| a.period_end); if variant_points.len() > 1 { // Verify counts are non-decreasing (cumulative) From 92255dd5a1b19bdbb809580e1d0f2203eca12ddb Mon Sep 17 00:00:00 2001 From: Gabriel Bianconi <1275491+GabrielBianconi@users.noreply.github.com> Date: Tue, 27 Jan 2026 12:09:27 -0500 Subject: [PATCH 31/46] Fix OpenAI error message format (#5818) * Fix OpenAI error message format * Fix OpenAI error message format * Fix * Fix OpenAI auth format with base path Co-authored-by: gabriel * Fix tests * Fix * Fix --------- Co-authored-by: Cursor Agent --- .../tests/openai_compatibility_test.go | 6 +- .../tests/openai-compatibility.test.ts | 2 +- .../python/tests/test_openai_compatibility.py | 6 +- gateway/src/router.rs | 1 + gateway/tests/auth.rs | 126 ++++++ gateway/tests/error_json.rs | 374 +++++++++++++++++- internal/tensorzero-auth/src/middleware.rs | 51 ++- .../openai_compatible/chat_completions.rs | 13 +- .../endpoints/openai_compatible/embeddings.rs | 10 +- .../src/endpoints/openai_compatible/error.rs | 68 ++++ .../src/endpoints/openai_compatible/mod.rs | 3 + tensorzero-core/src/error/mod.rs | 35 +- tensorzero-core/src/utils/gateway.rs | 62 +-- .../tests/e2e/openai_compatible.rs | 35 +- .../tests/e2e/providers/openai/mod.rs | 5 +- .../tests/e2e/raw_usage/openai_compatible.rs | 2 +- 16 files changed, 723 insertions(+), 76 deletions(-) create mode 100644 tensorzero-core/src/endpoints/openai_compatible/error.rs diff --git a/clients/openai-go/tests/openai_compatibility_test.go b/clients/openai-go/tests/openai_compatibility_test.go index cc6c44fb1aa..27cd0584a64 100644 --- a/clients/openai-go/tests/openai_compatibility_test.go +++ b/clients/openai-go/tests/openai_compatibility_test.go @@ -1025,7 +1025,8 @@ func TestStreamingInference(t *testing.T) { var apiErr *openai.Error assert.ErrorAs(t, err, &apiErr, "Expected error to be of type APIError") // ErrorAs assign err to apiErr assert.Equal(t, 404, apiErr.StatusCode, "Expected status code 404") - assert.Contains(t, err.Error(), "404 Not Found \"Unknown function: does_not_exist\"", "Error should indicate 404 Not Found") + assert.Contains(t, err.Error(), "404 Not Found", "Error should indicate 404 Not Found") + assert.Contains(t, err.Error(), "Unknown function: does_not_exist", "Error should indicate unknown function") }) t.Run("it should handle streaming inference with a missing function", func(t *testing.T) { @@ -1075,7 +1076,8 @@ func TestStreamingInference(t *testing.T) { var apiErr *openai.Error assert.ErrorAs(t, err, &apiErr, "Expected error to be of type APIError") assert.Equal(t, 400, apiErr.StatusCode, "Expected status code 404") - assert.Contains(t, apiErr.Error(), "400 Bad Request \"Invalid request to OpenAI-compatible endpoint", "Error should indicate invalid request to OpenAI compartible endpoint") + assert.Contains(t, apiErr.Error(), "400 Bad Request", "Error should indicate 400 Bad Request") + assert.Contains(t, apiErr.Error(), "Invalid request to OpenAI-compatible endpoint", "Error should indicate invalid request to OpenAI compatible endpoint") }) t.Run("it should handle streaming inference with a missing model", func(t *testing.T) { diff --git a/clients/openai-node/tests/openai-compatibility.test.ts b/clients/openai-node/tests/openai-compatibility.test.ts index 44158cf32fb..34e919b2bd8 100644 --- a/clients/openai-node/tests/openai-compatibility.test.ts +++ b/clients/openai-node/tests/openai-compatibility.test.ts @@ -1305,7 +1305,7 @@ it("should reject string input for function with input schema", async () => { // @ts-expect-error - custom TensorZero property "tensorzero::episode_id": episodeId, }) - ).rejects.toThrow(/400 "JSON Schema validation failed/); + ).rejects.toThrow(/400 JSON Schema validation failed/); }); it("should handle multi-turn parallel tool calls", async () => { diff --git a/clients/python/tests/test_openai_compatibility.py b/clients/python/tests/test_openai_compatibility.py index d894473854e..7caff20990d 100644 --- a/clients/python/tests/test_openai_compatibility.py +++ b/clients/python/tests/test_openai_compatibility.py @@ -364,7 +364,7 @@ async def test_async_inference_streaming_nonexistent_function(async_openai_clien # TODO(#3192): handle json errors in patched client assert ( str(exc_info.value) - == "Error code: 404 - {'error': 'Unknown function: does_not_exist', 'error_json': {'UnknownFunction': {'name': 'does_not_exist'}}}" + == "Error code: 404 - {'error': {'message': 'Unknown function: does_not_exist', 'error_json': {'UnknownFunction': {'name': 'does_not_exist'}}, 'tensorzero_error_json': {'UnknownFunction': {'name': 'does_not_exist'}}}}" ) @@ -396,7 +396,7 @@ async def test_async_inference_streaming_missing_function(async_openai_client): # TODO(#3192): handle json errors in patched client assert ( str(exc_info.value) - == """Error code: 400 - {'error': 'Invalid request to OpenAI-compatible endpoint: function_name (passed in model field after "tensorzero::function_name::") cannot be empty', 'error_json': {'InvalidOpenAICompatibleRequest': {'message': 'function_name (passed in model field after "tensorzero::function_name::") cannot be empty'}}}""" + == """Error code: 400 - {'error': {'message': 'Invalid request to OpenAI-compatible endpoint: function_name (passed in model field after "tensorzero::function_name::") cannot be empty', 'error_json': {'InvalidOpenAICompatibleRequest': {'message': 'function_name (passed in model field after "tensorzero::function_name::") cannot be empty'}}, 'tensorzero_error_json': {'InvalidOpenAICompatibleRequest': {'message': 'function_name (passed in model field after "tensorzero::function_name::") cannot be empty'}}}}""" ) @@ -428,7 +428,7 @@ async def test_async_inference_streaming_malformed_function(async_openai_client) # TODO(#3192): handle json errors in patched client assert ( str(exc_info.value) - == """Error code: 400 - {'error': 'Invalid request to OpenAI-compatible endpoint: `model` field must start with `tensorzero::function_name::` or `tensorzero::model_name::`. For example, `tensorzero::function_name::my_function` for a function `my_function` defined in your config, `tensorzero::model_name::my_model` for a model `my_model` defined in your config, or default functions like `tensorzero::model_name::openai::gpt-4o-mini`.', 'error_json': {'InvalidOpenAICompatibleRequest': {'message': '`model` field must start with `tensorzero::function_name::` or `tensorzero::model_name::`. For example, `tensorzero::function_name::my_function` for a function `my_function` defined in your config, `tensorzero::model_name::my_model` for a model `my_model` defined in your config, or default functions like `tensorzero::model_name::openai::gpt-4o-mini`.'}}}""" + == """Error code: 400 - {'error': {'message': 'Invalid request to OpenAI-compatible endpoint: `model` field must start with `tensorzero::function_name::` or `tensorzero::model_name::`. For example, `tensorzero::function_name::my_function` for a function `my_function` defined in your config, `tensorzero::model_name::my_model` for a model `my_model` defined in your config, or default functions like `tensorzero::model_name::openai::gpt-4o-mini`.', 'error_json': {'InvalidOpenAICompatibleRequest': {'message': '`model` field must start with `tensorzero::function_name::` or `tensorzero::model_name::`. For example, `tensorzero::function_name::my_function` for a function `my_function` defined in your config, `tensorzero::model_name::my_model` for a model `my_model` defined in your config, or default functions like `tensorzero::model_name::openai::gpt-4o-mini`.'}}, 'tensorzero_error_json': {'InvalidOpenAICompatibleRequest': {'message': '`model` field must start with `tensorzero::function_name::` or `tensorzero::model_name::`. For example, `tensorzero::function_name::my_function` for a function `my_function` defined in your config, `tensorzero::model_name::my_model` for a model `my_model` defined in your config, or default functions like `tensorzero::model_name::openai::gpt-4o-mini`.'}}}}""" ) diff --git a/gateway/src/router.rs b/gateway/src/router.rs index 2483fdd6b69..c9b094f6e6d 100644 --- a/gateway/src/router.rs +++ b/gateway/src/router.rs @@ -43,6 +43,7 @@ pub fn build_axum_router( auth_cache: app_state.auth_cache.clone(), pool: app_state.postgres_connection_info.get_pool().cloned(), error_json: app_state.config.gateway.unstable_error_json, + base_path: (!base_path.is_empty()).then(|| base_path.to_string()), }); router = router.layer(middleware::from_fn_with_state( state, diff --git a/gateway/tests/auth.rs b/gateway/tests/auth.rs index 457ded6e989..c89210c8e9b 100644 --- a/gateway/tests/auth.rs +++ b/gateway/tests/auth.rs @@ -945,6 +945,132 @@ async fn test_rate_limit_auth_single_key() { .unwrap(); } +/// Tests that auth error responses use the correct format (OpenAI vs TensorZero) based on the route, +/// without a custom base_path. +#[tokio::test] +async fn test_auth_error_format_without_base_path() { + let child_data = start_gateway_on_random_port( + r" + [gateway.auth] + enabled = true + ", + None, + ) + .await; + + let client = reqwest::Client::new(); + + // Test OpenAI-compatible route - should return OpenAI error format + let openai_response = client + .post(format!( + "http://{}/openai/v1/chat/completions", + child_data.addr + )) + .send() + .await + .unwrap(); + + assert_eq!(openai_response.status(), StatusCode::UNAUTHORIZED); + let openai_error: Value = openai_response.json().await.unwrap(); + assert!( + openai_error.get("error").unwrap().get("message").is_some(), + "OpenAI route should return OpenAI error format with nested message, got: {openai_error}" + ); + + // Test non-OpenAI route - should return TensorZero error format + let tensorzero_response = client + .post(format!("http://{}/inference", child_data.addr)) + .send() + .await + .unwrap(); + + assert_eq!(tensorzero_response.status(), StatusCode::UNAUTHORIZED); + let tensorzero_error: Value = tensorzero_response.json().await.unwrap(); + assert!( + tensorzero_error.get("error").unwrap().is_string(), + "TensorZero route should return flat error format, got: {tensorzero_error}" + ); +} + +/// Tests that auth error responses use the correct format (OpenAI vs TensorZero) based on the route, +/// and that this works correctly with a custom base_path. +/// +/// This test verifies: +/// 1. The `MatchedPath` from Axum includes the base path (e.g., `/my/prefix/openai/v1/...`) +/// 2. The `is_openai_compatible_route` function correctly identifies OpenAI routes with base paths +/// 3. OpenAI-compatible routes return errors in OpenAI format: `{"error": {"message": "..."}}` +/// 4. Non-OpenAI routes return errors in TensorZero format: `{"error": "..."}` +#[tokio::test] +async fn test_auth_error_format_with_base_path() { + let child_data = start_gateway_on_random_port( + r#" + base_path = "/my/prefix" + + [gateway.auth] + enabled = true + "#, + None, + ) + .await; + + let client = reqwest::Client::new(); + + // Test OpenAI-compatible route with base path - should return OpenAI error format + let openai_response = client + .post(format!( + "http://{}/my/prefix/openai/v1/chat/completions", + child_data.addr + )) + .send() + .await + .unwrap(); + + assert_eq!(openai_response.status(), StatusCode::UNAUTHORIZED); + let openai_error: Value = openai_response.json().await.unwrap(); + assert!( + openai_error.get("error").unwrap().get("message").is_some(), + "OpenAI route should return OpenAI error format with nested message, got: {openai_error}" + ); + + // Test non-OpenAI route with base path - should return TensorZero error format + let tensorzero_response = client + .post(format!("http://{}/my/prefix/inference", child_data.addr)) + .send() + .await + .unwrap(); + + assert_eq!(tensorzero_response.status(), StatusCode::UNAUTHORIZED); + let tensorzero_error: Value = tensorzero_response.json().await.unwrap(); + assert!( + tensorzero_error.get("error").unwrap().is_string(), + "TensorZero route should return flat error format, got: {tensorzero_error}" + ); + + // When hitting a route without the base path, the MatchedPath extension is None + // (because the route doesn't match any registered route), so the error format + // should fall back to TensorZero format (not OpenAI format) + let no_prefix_response = client + .post(format!( + "http://{}/openai/v1/chat/completions", + child_data.addr + )) + .send() + .await + .unwrap(); + + assert_eq!( + no_prefix_response.status(), + StatusCode::UNAUTHORIZED, + "Route without base path should still require auth" + ); + let no_prefix_error: Value = no_prefix_response.json().await.unwrap(); + // Without a MatchedPath, the function returns false, so we get TensorZero format + assert!( + no_prefix_error.get("error").unwrap().is_string(), + "Route without MatchedPath should return TensorZero error format, got: {no_prefix_error}" + ); +} + #[tokio::test] async fn test_rate_limit_auth_each_key() { let pool = get_postgres_pool_for_testing().await; diff --git a/gateway/tests/error_json.rs b/gateway/tests/error_json.rs index 4558dc0a3d1..7225cea2ca8 100644 --- a/gateway/tests/error_json.rs +++ b/gateway/tests/error_json.rs @@ -1,3 +1,6 @@ +use http::StatusCode; +use serde_json::Value; + use crate::common::start_gateway_on_random_port; mod common; @@ -30,8 +33,375 @@ async fn test_error_json() { .text() .await .unwrap(); + + let json: Value = + serde_json::from_str(&inference_response).expect("Response should be valid JSON"); + + // Verify TensorZero error format with both `error_json` and `unstable_error_json` assert_eq!( - inference_response, - r#"{"error":"Failed to parse the request body as JSON: EOF while parsing a value at line 1 column 0 (400 Bad Request)","error_json":{"JsonRequest":{"message":"Failed to parse the request body as JSON: EOF while parsing a value at line 1 column 0 (400 Bad Request)"}}}"# + json.get("error").and_then(|v| v.as_str()), + Some( + "Failed to parse the request body as JSON: EOF while parsing a value at line 1 column 0 (400 Bad Request)" + ), + "Response should have correct error message" + ); + assert!( + json.get("error_json").is_some(), + "Response should have `error_json` field" + ); +} + +/// Test that OpenAI-compatible chat completions endpoint returns errors in OpenAI format +/// OpenAI format: {"error": {"message": "..."}} +#[tokio::test] +async fn test_openai_compatible_error_format() { + let child_data = start_gateway_on_random_port("", None).await; + let response = reqwest::Client::new() + .post(format!( + "http://{}/openai/v1/chat/completions", + child_data.addr + )) + .send() + .await + .unwrap() + .text() + .await + .unwrap(); + + let json: Value = serde_json::from_str(&response).expect("Response should be valid JSON"); + + // Verify OpenAI error format: {"error": {"message": "..."}} + assert!( + json.get("error").is_some(), + "Response should have an `error` field" + ); + let error_obj = json.get("error").unwrap(); + assert!( + error_obj.is_object(), + "The `error` field should be an object, not a string" + ); + assert!( + error_obj.get("message").is_some(), + "The `error` object should have a `message` field" + ); + assert!( + error_obj.get("message").unwrap().is_string(), + "The `message` field should be a string" + ); +} + +/// Test that OpenAI-compatible embeddings endpoint returns errors in OpenAI format +#[tokio::test] +async fn test_openai_compatible_embeddings_error_format() { + let child_data = start_gateway_on_random_port("", None).await; + let response = reqwest::Client::new() + .post(format!("http://{}/openai/v1/embeddings", child_data.addr)) + .send() + .await + .unwrap() + .text() + .await + .unwrap(); + + let json: Value = serde_json::from_str(&response).expect("Response should be valid JSON"); + + // Verify OpenAI error format: {"error": {"message": "..."}} + assert!( + json.get("error").is_some(), + "Response should have an `error` field" + ); + let error_obj = json.get("error").unwrap(); + assert!( + error_obj.is_object(), + "The `error` field should be an object, not a string" + ); + assert!( + error_obj.get("message").is_some(), + "The `error` object should have a `message` field" + ); + assert!( + error_obj.get("message").unwrap().is_string(), + "The `message` field should be a string" + ); +} + +/// Test that OpenAI-compatible endpoint includes `error_json` when `UNSTABLE_ERROR_JSON` is enabled +#[tokio::test] +async fn test_openai_compatible_error_json() { + let child_data = start_gateway_on_random_port("unstable_error_json = true", None).await; + let response = reqwest::Client::new() + .post(format!( + "http://{}/openai/v1/chat/completions", + child_data.addr + )) + .send() + .await + .unwrap() + .text() + .await + .unwrap(); + + let json: Value = serde_json::from_str(&response).expect("Response should be valid JSON"); + + // Verify OpenAI error format with `error_json` and `tensorzero_error_json`: + // {"error": {"message": "...", "error_json": {...}, "tensorzero_error_json": {...}}} + let error_obj = json + .get("error") + .expect("Response should have an `error` field"); + assert!( + error_obj.is_object(), + "The `error` field should be an object" + ); + assert!( + error_obj.get("message").is_some(), + "The `error` object should have a `message` field" + ); + assert!( + error_obj.get("error_json").is_some(), + "The `error` object should have an `error_json` field when `UNSTABLE_ERROR_JSON` is enabled" + ); + assert!( + error_obj.get("error_json").unwrap().is_object(), + "The `error_json` field should be an object" + ); + assert!( + error_obj.get("tensorzero_error_json").is_some(), + "The `error` object should have a `tensorzero_error_json` field when `UNSTABLE_ERROR_JSON` is enabled" + ); + assert!( + error_obj.get("tensorzero_error_json").unwrap().is_object(), + "The `tensorzero_error_json` field should be an object" + ); + // Both fields should have identical content + assert_eq!( + error_obj.get("error_json"), + error_obj.get("tensorzero_error_json"), + "`error_json` and `tensorzero_error_json` should be identical" + ); +} + +/// Test that auth errors on TensorZero endpoints return TensorZero format +#[tokio::test] +async fn test_auth_error_tensorzero_format() { + let child_data = start_gateway_on_random_port( + " + [gateway.auth] + enabled = true + ", + None, + ) + .await; + + let response = reqwest::Client::new() + .post(format!("http://{}/inference", child_data.addr)) + .send() + .await + .unwrap(); + + assert_eq!(response.status(), StatusCode::UNAUTHORIZED); + let text = response.text().await.unwrap(); + let json: Value = serde_json::from_str(&text).expect("Response should be valid JSON"); + + // Verify TensorZero error format: {"error": "..."} + assert!( + json.get("error").is_some(), + "Response should have an `error` field" + ); + assert!( + json.get("error").unwrap().is_string(), + "The `error` field should be a string for TensorZero format" + ); + assert!( + json.get("error") + .unwrap() + .as_str() + .unwrap() + .contains("TensorZero authentication error"), + "Error message should mention authentication" + ); +} + +/// Test that auth errors on OpenAI endpoints return OpenAI format +#[tokio::test] +async fn test_auth_error_openai_format() { + let child_data = start_gateway_on_random_port( + " + [gateway.auth] + enabled = true + ", + None, + ) + .await; + + let response = reqwest::Client::new() + .post(format!( + "http://{}/openai/v1/chat/completions", + child_data.addr + )) + .send() + .await + .unwrap(); + + assert_eq!(response.status(), StatusCode::UNAUTHORIZED); + let text = response.text().await.unwrap(); + let json: Value = serde_json::from_str(&text).expect("Response should be valid JSON"); + + // Verify OpenAI error format: {"error": {"message": "..."}} + assert!( + json.get("error").is_some(), + "Response should have an `error` field" + ); + let error_obj = json.get("error").unwrap(); + assert!( + error_obj.is_object(), + "The `error` field should be an object for OpenAI format" + ); + assert!( + error_obj.get("message").is_some(), + "The `error` object should have a `message` field" + ); + assert!( + error_obj + .get("message") + .unwrap() + .as_str() + .unwrap() + .contains("TensorZero authentication error"), + "Error message should mention authentication" + ); +} + +/// Test that auth errors on OpenAI endpoints use OpenAI format with a base path +#[tokio::test] +async fn test_auth_error_openai_format_with_base_path() { + let child_data = start_gateway_on_random_port( + r#" + base_path = "/my/prefix" + + [gateway.auth] + enabled = true + "#, + None, + ) + .await; + + let response = reqwest::Client::new() + .post(format!( + "http://{}/my/prefix/openai/v1/chat/completions", + child_data.addr + )) + .send() + .await + .unwrap(); + + assert_eq!( + response.status(), + StatusCode::UNAUTHORIZED, + "Auth errors should return 401 for OpenAI endpoints with a base path" + ); + let text = response.text().await.unwrap(); + let json: Value = serde_json::from_str(&text).expect("Response should be valid JSON"); + + assert!( + json.get("error").is_some(), + "Response should have an `error` field" + ); + let error_obj = json.get("error").unwrap(); + assert!( + error_obj.is_object(), + "The `error` field should be an object for OpenAI format with a base path" + ); + assert!( + error_obj.get("message").and_then(Value::as_str).is_some(), + "The `error` object should have a string `message` field" + ); +} + +/// Test that auth errors on TensorZero endpoints include `error_json` when enabled +#[tokio::test] +async fn test_auth_error_tensorzero_format_with_error_json() { + let child_data = start_gateway_on_random_port( + " + unstable_error_json = true + + [gateway.auth] + enabled = true + ", + None, + ) + .await; + + let response = reqwest::Client::new() + .post(format!("http://{}/inference", child_data.addr)) + .send() + .await + .unwrap(); + + assert_eq!(response.status(), StatusCode::UNAUTHORIZED); + let text = response.text().await.unwrap(); + let json: Value = serde_json::from_str(&text).expect("Response should be valid JSON"); + + // Verify TensorZero error format with `error_json` + assert!( + json.get("error").is_some() && json.get("error").unwrap().is_string(), + "Response should have a string `error` field" + ); + assert!( + json.get("error_json").is_some(), + "Response should have `error_json` field when `unstable_error_json` is enabled" + ); +} + +/// Test that auth errors on OpenAI endpoints include `error_json` when enabled +#[tokio::test] +async fn test_auth_error_openai_format_with_error_json() { + let child_data = start_gateway_on_random_port( + " + unstable_error_json = true + + [gateway.auth] + enabled = true + ", + None, + ) + .await; + + let response = reqwest::Client::new() + .post(format!( + "http://{}/openai/v1/chat/completions", + child_data.addr + )) + .send() + .await + .unwrap(); + + assert_eq!(response.status(), StatusCode::UNAUTHORIZED); + let text = response.text().await.unwrap(); + let json: Value = serde_json::from_str(&text).expect("Response should be valid JSON"); + + // Verify OpenAI error format with `error_json` and `tensorzero_error_json` + let error_obj = json + .get("error") + .expect("Response should have an `error` field"); + assert!( + error_obj.is_object(), + "The `error` field should be an object" + ); + assert!( + error_obj.get("message").is_some(), + "The `error` object should have a `message` field" + ); + assert!( + error_obj.get("error_json").is_some(), + "The `error` object should have an `error_json` field when `unstable_error_json` is enabled" + ); + assert!( + error_obj.get("tensorzero_error_json").is_some(), + "The `error` object should have a `tensorzero_error_json` field when `unstable_error_json` is enabled" + ); + assert_eq!( + error_obj.get("error_json"), + error_obj.get("tensorzero_error_json"), + "`error_json` and `tensorzero_error_json` should be identical" ); } diff --git a/internal/tensorzero-auth/src/middleware.rs b/internal/tensorzero-auth/src/middleware.rs index 4e026e68bae..b284aff5d37 100644 --- a/internal/tensorzero-auth/src/middleware.rs +++ b/internal/tensorzero-auth/src/middleware.rs @@ -8,7 +8,7 @@ use axum::{ response::{IntoResponse, Response}, }; use moka::sync::Cache; -use serde_json::json; +use serde_json::{Value, json}; use tracing::{Instrument, field::Empty}; use crate::{ @@ -16,6 +16,40 @@ use crate::{ postgres::{AuthResult, KeyInfo}, }; +/// Builds an error response body in either TensorZero or OpenAI format. +/// +/// When `openai_format` is true, returns `{"error": {"message": "..."}}`. +/// When `openai_format` is false, returns `{"error": "..."}`. +/// +/// If `error_json` is provided, includes structured error details. +fn build_error_response_body( + message: &str, + openai_format: bool, + error_json: Option, +) -> Value { + let mut body = if openai_format { + json!({"error": {"message": message}}) + } else { + json!({"error": message}) + }; + if let Some(error_json) = error_json { + if openai_format { + body["error"]["error_json"] = error_json.clone(); // DEPRECATED (#5821 / 2026.4+) + body["error"]["tensorzero_error_json"] = error_json; + } else { + body["error_json"] = error_json; + } + } + body +} + +fn is_openai_compatible_route(matched_path: Option<&MatchedPath>, base_path: Option<&str>) -> bool { + matched_path.is_some_and(|p| match base_path { + Some(base) => p.as_str().starts_with(&format!("{base}/openai/v1")), + None => p.as_str().starts_with("/openai/v1"), + }) +} + #[derive(Clone)] pub struct TensorzeroAuthMiddlewareState(Arc); @@ -34,6 +68,8 @@ pub struct TensorzeroAuthMiddlewareStateInner { pub auth_cache: Option>, pub pool: Option, pub error_json: bool, + /// Optional base path prefix for all routes (e.g., "/custom/prefix") + pub base_path: Option, } #[axum::debug_middleware] @@ -160,13 +196,12 @@ pub async fn tensorzero_auth_middleware( auth_span.record("key.organization", &key_info.organization); auth_span.record("key.workspace", &key_info.workspace); } - let message = e.to_string(); - let mut body = json!({ - "error": format!("TensorZero authentication error: {message}"), - }); - if state.error_json { - body["error_json"] = json!(e.to_string()); - } + let message = format!("TensorZero authentication error: {e}"); + let matched_path = request.extensions().get::(); + let is_openai_format = + is_openai_compatible_route(matched_path, state.base_path.as_deref()); + let error_json = state.error_json.then(|| json!(e.to_string())); + let body = build_error_response_body(&message, is_openai_format, error_json); let mut response = (StatusCode::UNAUTHORIZED, Json(body)).into_response(); // Attach the error to the response, so that we can set a nice message in our // `apply_otel_http_trace_layer` middleware diff --git a/tensorzero-core/src/endpoints/openai_compatible/chat_completions.rs b/tensorzero-core/src/endpoints/openai_compatible/chat_completions.rs index e954617248a..7b903ecd101 100644 --- a/tensorzero-core/src/endpoints/openai_compatible/chat_completions.rs +++ b/tensorzero-core/src/endpoints/openai_compatible/chat_completions.rs @@ -14,11 +14,12 @@ use axum::{Extension, debug_handler}; use crate::endpoints::inference::{InferenceOutput, Params, inference}; use crate::error::{Error, ErrorDetails}; -use crate::utils::gateway::{AppState, AppStateData, StructuredJson}; +use crate::utils::gateway::{AppState, AppStateData}; use tensorzero_auth::middleware::RequestApiKeyExtension; use super::types::chat_completions::{OpenAICompatibleParams, OpenAICompatibleResponse}; use super::types::streaming::prepare_serialized_openai_compatible_events; +use super::{OpenAICompatibleError, OpenAIStructuredJson}; /// A handler for the OpenAI-compatible inference endpoint #[debug_handler(state = AppStateData)] @@ -33,15 +34,15 @@ pub async fn chat_completions_handler( .. }): AppState, api_key_ext: Option>, - StructuredJson(openai_compatible_params): StructuredJson, -) -> Result, Error> { + OpenAIStructuredJson(openai_compatible_params): OpenAIStructuredJson, +) -> Result, OpenAICompatibleError> { // Validate `n` parameter if let Some(n) = openai_compatible_params.n && n != 1 { return Err(Error::new(ErrorDetails::InvalidOpenAICompatibleRequest { message: "TensorZero does not support `n` other than 1. Please omit this parameter or set it to 1.".to_string(), - })); + }).into()); } if !openai_compatible_params.unknown_fields.is_empty() { @@ -59,7 +60,7 @@ pub async fn chat_completions_handler( message: format!( "`tensorzero::deny_unknown_fields` is set to true, but found unknown fields in the request: [{unknown_field_names}]" ), - })); + }).into()); } tracing::warn!( "Ignoring unknown fields in OpenAI-compatible request: {:?}", @@ -92,7 +93,7 @@ pub async fn chat_completions_handler( { return Err(Error::new(ErrorDetails::InvalidOpenAICompatibleRequest { message: "`tensorzero::include_raw_usage` requires `stream_options.include_usage` to be true (or omitted) for streaming requests".to_string(), - })); + }).into()); } // OpenAI default: no usage when stream_options is omitted diff --git a/tensorzero-core/src/endpoints/openai_compatible/embeddings.rs b/tensorzero-core/src/endpoints/openai_compatible/embeddings.rs index f05592b3a0c..423f6dc4f8c 100644 --- a/tensorzero-core/src/endpoints/openai_compatible/embeddings.rs +++ b/tensorzero-core/src/endpoints/openai_compatible/embeddings.rs @@ -7,11 +7,11 @@ use axum::{Extension, Json, extract::State}; use crate::endpoints::embeddings::embeddings; -use crate::error::Error; -use crate::utils::gateway::{AppState, AppStateData, StructuredJson}; +use crate::utils::gateway::{AppState, AppStateData}; use tensorzero_auth::middleware::RequestApiKeyExtension; use super::types::embeddings::{OpenAICompatibleEmbeddingParams, OpenAIEmbeddingResponse}; +use super::{OpenAICompatibleError, OpenAIStructuredJson}; pub async fn embeddings_handler( State(AppStateData { @@ -24,8 +24,10 @@ pub async fn embeddings_handler( .. }): AppState, api_key_ext: Option>, - StructuredJson(openai_compatible_params): StructuredJson, -) -> Result, Error> { + OpenAIStructuredJson(openai_compatible_params): OpenAIStructuredJson< + OpenAICompatibleEmbeddingParams, + >, +) -> Result, OpenAICompatibleError> { let embedding_params = openai_compatible_params.try_into()?; let response = embeddings( config, diff --git a/tensorzero-core/src/endpoints/openai_compatible/error.rs b/tensorzero-core/src/endpoints/openai_compatible/error.rs new file mode 100644 index 00000000000..75964e3deb5 --- /dev/null +++ b/tensorzero-core/src/endpoints/openai_compatible/error.rs @@ -0,0 +1,68 @@ +//! OpenAI-compatible error handling. +//! +//! This module provides error types and extractors that format responses according to OpenAI's API specification. + +use std::fmt; + +use axum::Json; +use axum::extract::rejection::JsonRejection; +use axum::extract::{FromRequest, Request}; +use axum::response::{IntoResponse, Response}; +use serde::de::DeserializeOwned; +use tracing::instrument; + +use crate::error::Error; +use crate::utils::gateway::deserialize_json_request; + +/// A wrapper around `Error` that implements `IntoResponse` with OpenAI-compatible error format. +/// +/// OpenAI returns errors as `{"error": {"message": "..."}}` while TensorZero's default +/// format is `{"error": "..."}`. This wrapper ensures OpenAI-compatible endpoints +/// return errors in the expected format. +#[derive(Debug)] +pub struct OpenAICompatibleError(pub Error); + +impl From for OpenAICompatibleError { + fn from(error: Error) -> Self { + OpenAICompatibleError(error) + } +} + +impl fmt::Display for OpenAICompatibleError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + self.0.fmt(f) + } +} + +impl IntoResponse for OpenAICompatibleError { + fn into_response(self) -> Response { + let body = self.0.build_response_body(true); + let mut response = (self.0.status_code(), Json(body)).into_response(); + response.extensions_mut().insert(self.0); + response + } +} + +/// A JSON extractor for OpenAI-compatible endpoints that returns errors in OpenAI format. +/// +/// This is similar to `StructuredJson` but uses `OpenAICompatibleError` as its rejection type +/// so that JSON parsing errors are returned in OpenAI's `{"error": {"message": "..."}}` format. +#[derive(Debug, Clone, Copy, Default)] +pub struct OpenAIStructuredJson(pub T); + +impl FromRequest for OpenAIStructuredJson +where + Json: FromRequest, + S: Send + Sync, + T: Send + Sync + DeserializeOwned, +{ + type Rejection = OpenAICompatibleError; + + #[instrument(skip_all, level = "trace", name = "OpenAIStructuredJson::from_request")] + async fn from_request(req: Request, state: &S) -> Result { + deserialize_json_request(req, state) + .await + .map(OpenAIStructuredJson) + .map_err(OpenAICompatibleError) + } +} diff --git a/tensorzero-core/src/endpoints/openai_compatible/mod.rs b/tensorzero-core/src/endpoints/openai_compatible/mod.rs index 75f9bb77e70..1707d2fd130 100644 --- a/tensorzero-core/src/endpoints/openai_compatible/mod.rs +++ b/tensorzero-core/src/endpoints/openai_compatible/mod.rs @@ -6,8 +6,11 @@ pub mod chat_completions; pub mod embeddings; +pub mod error; pub mod types; +pub use error::{OpenAICompatibleError, OpenAIStructuredJson}; + use chat_completions::chat_completions_handler; use embeddings::embeddings_handler; diff --git a/tensorzero-core/src/error/mod.rs b/tensorzero-core/src/error/mod.rs index 94a25716b22..fe6a6cf3a30 100644 --- a/tensorzero-core/src/error/mod.rs +++ b/tensorzero-core/src/error/mod.rs @@ -168,6 +168,32 @@ impl Error { pub fn is_retryable(&self) -> bool { self.0.is_retryable() } + + /// Builds the JSON response body for this error. + /// + /// When `openai_format` is true, returns `{"error": {"message": "..."}}` (OpenAI-compatible). + /// When `openai_format` is false, returns `{"error": "..."}` (TensorZero default). + /// + /// If `unstable_error_json` is enabled, includes structured error details as `error_json` and `tensorzero_error_json`. + pub fn build_response_body(&self, openai_format: bool) -> Value { + let message = self.to_string(); + let mut body = if openai_format { + json!({"error": {"message": message}}) + } else { + json!({"error": message}) + }; + if *UNSTABLE_ERROR_JSON.get().unwrap_or(&false) { + let error_json = + serde_json::to_value(self.get_details()).unwrap_or_else(|e| json!(e.to_string())); + if openai_format { + body["error"]["error_json"] = error_json.clone(); // DEPRECATED (#5821 / 2026.4+) + body["error"]["tensorzero_error_json"] = error_json; + } else { + body["error_json"] = error_json; + } + } + body + } } // Expect for derive Serialize @@ -1681,14 +1707,7 @@ impl std::fmt::Display for ErrorDetails { impl IntoResponse for Error { /// Log the error and convert it into an Axum response fn into_response(self) -> Response { - let message = self.to_string(); - let mut body = json!({ - "error": message, - }); - if *UNSTABLE_ERROR_JSON.get().unwrap_or(&false) { - body["error_json"] = - serde_json::to_value(self.get_details()).unwrap_or_else(|e| json!(e.to_string())); - } + let body = self.build_response_body(false); let mut response = (self.status_code(), Json(body)).into_response(); // Attach the error to the response, so that we can set a nice message in our // `apply_otel_http_trace_layer` middleware diff --git a/tensorzero-core/src/utils/gateway.rs b/tensorzero-core/src/utils/gateway.rs index 3b4a4445e02..8d7cbb5011f 100644 --- a/tensorzero-core/src/utils/gateway.rs +++ b/tensorzero-core/src/utils/gateway.rs @@ -610,40 +610,54 @@ async fn setup_autopilot_client( /// and instead simply assume that the request body is a JSON object. pub struct StructuredJson(pub T); -impl FromRequest for StructuredJson +/// Shared JSON deserialization logic used by both `StructuredJson` and `OpenAIStructuredJson`. +/// +/// Parses the request body as JSON and deserializes it into the target type `T`, +/// using `serde_path_to_error` for detailed error messages. +pub(crate) async fn deserialize_json_request(req: Request, state: &S) -> Result where - Json: FromRequest, S: Send + Sync, - T: Send + Sync + DeserializeOwned, + T: DeserializeOwned, { - type Rejection = Error; + // Retrieve the request body as Bytes before deserializing it + let bytes = bytes::Bytes::from_request(req, state).await.map_err(|e| { + Error::new(ErrorDetails::JsonRequest { + message: format!("{} ({})", e, e.status()), + }) + })?; - #[instrument(skip_all, level = "trace", name = "StructuredJson::from_request")] - async fn from_request(req: Request, state: &S) -> Result { - // Retrieve the request body as Bytes before deserializing it - let bytes = bytes::Bytes::from_request(req, state).await.map_err(|e| { + // Convert the entire body into `serde_json::Value` + let value = Json::::from_bytes(&bytes) + .map_err(|e| { Error::new(ErrorDetails::JsonRequest { message: format!("{} ({})", e, e.status()), }) - })?; + })? + .0; - // Convert the entire body into `serde_json::Value` - let value = Json::::from_bytes(&bytes) - .map_err(|e| { - Error::new(ErrorDetails::JsonRequest { - message: format!("{} ({})", e, e.status()), - }) - })? - .0; + // Now use `serde_path_to_error::deserialize` to attempt deserialization into `T` + let deserialized: T = serde_path_to_error::deserialize(&value).map_err(|e| { + Error::new(ErrorDetails::JsonRequest { + message: e.to_string(), + }) + })?; - // Now use `serde_path_to_error::deserialize` to attempt deserialization into `T` - let deserialized: T = serde_path_to_error::deserialize(&value).map_err(|e| { - Error::new(ErrorDetails::JsonRequest { - message: e.to_string(), - }) - })?; + Ok(deserialized) +} + +impl FromRequest for StructuredJson +where + Json: FromRequest, + S: Send + Sync, + T: Send + Sync + DeserializeOwned, +{ + type Rejection = Error; - Ok(StructuredJson(deserialized)) + #[instrument(skip_all, level = "trace", name = "StructuredJson::from_request")] + async fn from_request(req: Request, state: &S) -> Result { + deserialize_json_request(req, state) + .await + .map(StructuredJson) } } diff --git a/tensorzero-core/tests/e2e/openai_compatible.rs b/tensorzero-core/tests/e2e/openai_compatible.rs index 4afbdc6165a..76097eea669 100644 --- a/tensorzero-core/tests/e2e/openai_compatible.rs +++ b/tensorzero-core/tests/e2e/openai_compatible.rs @@ -11,14 +11,12 @@ use uuid::Uuid; use crate::common::get_gateway_endpoint; -use tensorzero_core::endpoints::openai_compatible::chat_completions::chat_completions_handler; -use tensorzero_core::{ - db::clickhouse::test_helpers::{ - get_clickhouse, select_chat_inference_clickhouse, select_json_inference_clickhouse, - select_model_inference_clickhouse, - }, - utils::gateway::StructuredJson, +use tensorzero_core::db::clickhouse::test_helpers::{ + get_clickhouse, select_chat_inference_clickhouse, select_json_inference_clickhouse, + select_model_inference_clickhouse, }; +use tensorzero_core::endpoints::openai_compatible::OpenAIStructuredJson; +use tensorzero_core::endpoints::openai_compatible::chat_completions::chat_completions_handler; #[tokio::test(flavor = "multi_thread")] async fn test_openai_compatible_route_new_format() { @@ -36,7 +34,7 @@ async fn test_openai_compatible_route_with_function_name_as_model(model: &str) { let response = chat_completions_handler( State(state), None, - StructuredJson( + OpenAIStructuredJson( serde_json::from_value(serde_json::json!({ "model": model, "messages": [ @@ -471,10 +469,17 @@ async fn test_openai_compatible_route_bad_model_name() { assert_eq!( response_json, json!({ - "error": "Invalid inference target: Invalid model name: Model name 'my_missing_model' not found in model table", - "error_json": { - "InvalidInferenceTarget": { - "message": "Invalid model name: Model name 'my_missing_model' not found in model table" + "error": { + "message": "Invalid inference target: Invalid model name: Model name 'my_missing_model' not found in model table", + "error_json": { + "InvalidInferenceTarget": { + "message": "Invalid model name: Model name 'my_missing_model' not found in model table" + } + }, + "tensorzero_error_json": { + "InvalidInferenceTarget": { + "message": "Invalid model name: Model name 'my_missing_model' not found in model table" + } } } }) @@ -883,7 +888,7 @@ async fn test_openai_compatible_warn_unknown_fields() { chat_completions_handler( State(state), None, - StructuredJson( + OpenAIStructuredJson( serde_json::from_value(serde_json::json!({ "messages": [], "model": "tensorzero::model_name::dummy::good", @@ -907,7 +912,7 @@ async fn test_openai_compatible_deny_unknown_fields() { let err = chat_completions_handler( State(state), None, - StructuredJson( + OpenAIStructuredJson( serde_json::from_value(serde_json::json!({ "messages": [], "model": "tensorzero::model_name::dummy::good", @@ -1061,7 +1066,7 @@ async fn test_openai_compatible_file_with_custom_filename() { let response = chat_completions_handler( State(state), None, - StructuredJson( + OpenAIStructuredJson( serde_json::from_value(serde_json::json!({ "model": "tensorzero::function_name::basic_test_no_system_schema", "messages": [ diff --git a/tensorzero-core/tests/e2e/providers/openai/mod.rs b/tensorzero-core/tests/e2e/providers/openai/mod.rs index 1dd20ad12af..8fb346ab304 100644 --- a/tensorzero-core/tests/e2e/providers/openai/mod.rs +++ b/tensorzero-core/tests/e2e/providers/openai/mod.rs @@ -2002,9 +2002,10 @@ pub async fn test_embedding_extra_headers() { let response_json = response.json::().await.unwrap(); println!("API response: {response_json:?}"); // OpenAI returns 401 for invalid auth, which TensorZero wraps as an error - let error_message = response_json["error"] + // OpenAI-compatible endpoints return {"error": {"message": "..."}} format + let error_message = response_json["error"]["message"] .as_str() - .expect("Expected a string error response due to invalid auth header"); + .expect("Expected an OpenAI-format error response due to invalid auth header"); assert!( !error_message.is_empty(), "Expected an error message due to invalid auth header" diff --git a/tensorzero-core/tests/e2e/raw_usage/openai_compatible.rs b/tensorzero-core/tests/e2e/raw_usage/openai_compatible.rs index 9f197495d96..087e64bc29c 100644 --- a/tensorzero-core/tests/e2e/raw_usage/openai_compatible.rs +++ b/tensorzero-core/tests/e2e/raw_usage/openai_compatible.rs @@ -221,7 +221,7 @@ async fn test_openai_compatible_raw_usage_streaming_error_without_include_usage( ); let error_body: Value = response.json().await.unwrap(); - let error_message = error_body["error"].as_str().unwrap_or(""); + let error_message = error_body["error"]["message"].as_str().unwrap_or(""); assert!( error_message.contains("include_usage"), "Error message should mention include_usage requirement. Got: {error_message}" From 6e21d632715964788c9ce1b0fa098dc20d806c04 Mon Sep 17 00:00:00 2001 From: Gabriel Bianconi <1275491+GabrielBianconi@users.noreply.github.com> Date: Tue, 27 Jan 2026 12:57:54 -0500 Subject: [PATCH 32/46] Clean up load_fixtures (#5882) --- ui/fixtures/load_fixtures.sh | 21 ++------------------- 1 file changed, 2 insertions(+), 19 deletions(-) diff --git a/ui/fixtures/load_fixtures.sh b/ui/fixtures/load_fixtures.sh index 25d19dab31c..889a8c4300d 100755 --- a/ui/fixtures/load_fixtures.sh +++ b/ui/fixtures/load_fixtures.sh @@ -21,25 +21,7 @@ else CLICKHOUSE_SECURE_FLAG="" fi -# Truncate all tables before inserting new data -clickhouse-client --host $CLICKHOUSE_HOST_VAR --user $CLICKHOUSE_USER_VAR --password $CLICKHOUSE_PASSWORD_VAR $CLICKHOUSE_SECURE_FLAG --database "$DATABASE_NAME" --query "TRUNCATE TABLE JsonInference" -clickhouse-client --host $CLICKHOUSE_HOST_VAR --user $CLICKHOUSE_USER_VAR --password $CLICKHOUSE_PASSWORD_VAR $CLICKHOUSE_SECURE_FLAG --database "$DATABASE_NAME" --query "TRUNCATE TABLE BooleanMetricFeedback" -clickhouse-client --host $CLICKHOUSE_HOST_VAR --user $CLICKHOUSE_USER_VAR --password $CLICKHOUSE_PASSWORD_VAR $CLICKHOUSE_SECURE_FLAG --database "$DATABASE_NAME" --query "TRUNCATE TABLE FloatMetricFeedback" -clickhouse-client --host $CLICKHOUSE_HOST_VAR --user $CLICKHOUSE_USER_VAR --password $CLICKHOUSE_PASSWORD_VAR $CLICKHOUSE_SECURE_FLAG --database "$DATABASE_NAME" --query "TRUNCATE TABLE CommentFeedback" -clickhouse-client --host $CLICKHOUSE_HOST_VAR --user $CLICKHOUSE_USER_VAR --password $CLICKHOUSE_PASSWORD_VAR $CLICKHOUSE_SECURE_FLAG --database "$DATABASE_NAME" --query "TRUNCATE TABLE DemonstrationFeedback" -clickhouse-client --host $CLICKHOUSE_HOST_VAR --user $CLICKHOUSE_USER_VAR --password $CLICKHOUSE_PASSWORD_VAR $CLICKHOUSE_SECURE_FLAG --database "$DATABASE_NAME" --query "TRUNCATE TABLE ChatInference" -clickhouse-client --host $CLICKHOUSE_HOST_VAR --user $CLICKHOUSE_USER_VAR --password $CLICKHOUSE_PASSWORD_VAR $CLICKHOUSE_SECURE_FLAG --database "$DATABASE_NAME" --query "TRUNCATE TABLE ModelInference" -clickhouse-client --host $CLICKHOUSE_HOST_VAR --user $CLICKHOUSE_USER_VAR --password $CLICKHOUSE_PASSWORD_VAR $CLICKHOUSE_SECURE_FLAG --database "$DATABASE_NAME" --query "TRUNCATE TABLE ChatInferenceDatapoint" -clickhouse-client --host $CLICKHOUSE_HOST_VAR --user $CLICKHOUSE_USER_VAR --password $CLICKHOUSE_PASSWORD_VAR $CLICKHOUSE_SECURE_FLAG --database "$DATABASE_NAME" --query "TRUNCATE TABLE JsonInferenceDatapoint" -clickhouse-client --host $CLICKHOUSE_HOST_VAR --user $CLICKHOUSE_USER_VAR --password $CLICKHOUSE_PASSWORD_VAR $CLICKHOUSE_SECURE_FLAG --database "$DATABASE_NAME" --query "TRUNCATE TABLE DynamicEvaluationRun" -clickhouse-client --host $CLICKHOUSE_HOST_VAR --user $CLICKHOUSE_USER_VAR --password $CLICKHOUSE_PASSWORD_VAR $CLICKHOUSE_SECURE_FLAG --database "$DATABASE_NAME" --query "TRUNCATE TABLE DynamicEvaluationRunEpisode" -clickhouse-client --host $CLICKHOUSE_HOST_VAR --user $CLICKHOUSE_USER_VAR --password $CLICKHOUSE_PASSWORD_VAR $CLICKHOUSE_SECURE_FLAG --database "$DATABASE_NAME" --query "TRUNCATE TABLE ModelInferenceCache" -clickhouse-client --host $CLICKHOUSE_HOST_VAR --user $CLICKHOUSE_USER_VAR --password $CLICKHOUSE_PASSWORD_VAR $CLICKHOUSE_SECURE_FLAG --database "$DATABASE_NAME" --query "TRUNCATE TABLE DeploymentID" - - - - -# Truncate all tables first to ensure clean loading +# Truncate all tables before loading fixtures echo "Truncating all tables before loading fixtures..." clickhouse-client --host $CLICKHOUSE_HOST_VAR --user $CLICKHOUSE_USER_VAR --password $CLICKHOUSE_PASSWORD_VAR $CLICKHOUSE_SECURE_FLAG --database "$DATABASE_NAME" --query "TRUNCATE TABLE JsonInference" clickhouse-client --host $CLICKHOUSE_HOST_VAR --user $CLICKHOUSE_USER_VAR --password $CLICKHOUSE_PASSWORD_VAR $CLICKHOUSE_SECURE_FLAG --database "$DATABASE_NAME" --query "TRUNCATE TABLE BooleanMetricFeedback" @@ -53,6 +35,7 @@ clickhouse-client --host $CLICKHOUSE_HOST_VAR --user $CLICKHOUSE_USER_VAR --pass clickhouse-client --host $CLICKHOUSE_HOST_VAR --user $CLICKHOUSE_USER_VAR --password $CLICKHOUSE_PASSWORD_VAR $CLICKHOUSE_SECURE_FLAG --database "$DATABASE_NAME" --query "TRUNCATE TABLE DynamicEvaluationRun" clickhouse-client --host $CLICKHOUSE_HOST_VAR --user $CLICKHOUSE_USER_VAR --password $CLICKHOUSE_PASSWORD_VAR $CLICKHOUSE_SECURE_FLAG --database "$DATABASE_NAME" --query "TRUNCATE TABLE DynamicEvaluationRunEpisode" clickhouse-client --host $CLICKHOUSE_HOST_VAR --user $CLICKHOUSE_USER_VAR --password $CLICKHOUSE_PASSWORD_VAR $CLICKHOUSE_SECURE_FLAG --database "$DATABASE_NAME" --query "TRUNCATE TABLE ModelInferenceCache" +clickhouse-client --host $CLICKHOUSE_HOST_VAR --user $CLICKHOUSE_USER_VAR --password $CLICKHOUSE_PASSWORD_VAR $CLICKHOUSE_SECURE_FLAG --database "$DATABASE_NAME" --query "TRUNCATE TABLE DeploymentID" # Download JSONL fixtures from R2 uv run ./download-small-fixtures.py From 2a746df80b78542e5408f504b66908297e3a0429 Mon Sep 17 00:00:00 2001 From: Aaron Hill Date: Tue, 27 Jan 2026 15:10:45 -0500 Subject: [PATCH 33/46] Close autopilot stream when gateway shuts down (#5910) The ui will automatically reconnect. This stops the gateway shutdown from hanging indefinitely when browser tabs are open --- clients/rust/src/lib.rs | 14 ++++++++------ gateway/src/main.rs | 15 +++++++++++---- internal/durable-tools/src/action.rs | 2 ++ .../src/endpoints/internal/autopilot.rs | 11 ++++++++++- tensorzero-core/src/utils/gateway.rs | 10 ++++++---- 5 files changed, 37 insertions(+), 15 deletions(-) diff --git a/clients/rust/src/lib.rs b/clients/rust/src/lib.rs index e4af20551ac..8e57b943c01 100644 --- a/clients/rust/src/lib.rs +++ b/clients/rust/src/lib.rs @@ -598,12 +598,14 @@ impl ClientExt for Client { }), ClientMode::EmbeddedGateway { gateway, timeout } => { Ok(with_embedded_timeout(*timeout, async { - tensorzero_core::endpoints::batch_inference::start_batch_inference( - gateway.handle.app_state.clone(), - params, - // We currently ban auth-enabled configs in embedded gateway mode, - // so we don't have an API key here - None, + Box::pin( + tensorzero_core::endpoints::batch_inference::start_batch_inference( + gateway.handle.app_state.clone(), + params, + // We currently ban auth-enabled configs in embedded gateway mode, + // so we don't have an API key here + None, + ), ) .await .map_err(err_to_http) diff --git a/gateway/src/main.rs b/gateway/src/main.rs index d92415bf1da..e03e18f20d3 100644 --- a/gateway/src/main.rs +++ b/gateway/src/main.rs @@ -379,10 +379,17 @@ async fn run() -> Result<(), ExitCode> { tracing::info!("â”” OpenTelemetry: disabled"); } - let shutdown_signal = shutdown_signal().shared(); + let shutdown_token = gateway_handle.app_state.shutdown_token.clone(); + let shutdown_token_clone = shutdown_token.clone(); + // This is responsible for starting the shutdown + #[expect(clippy::disallowed_methods)] + tokio::spawn(async move { + shutdown_signal().await; + shutdown_token_clone.cancel(); + }); let server_fut = axum::serve(listener, router) - .with_graceful_shutdown(shutdown_signal.clone()) + .with_graceful_shutdown(shutdown_token.clone().cancelled_owned()) .into_future() .map(|r| { let _ = r.log_err_pretty("Failed to start server"); @@ -392,7 +399,7 @@ async fn run() -> Result<(), ExitCode> { // This is a purely informational logging task, so we don't need to wait for it to finish. #[expect(clippy::disallowed_methods)] tokio::spawn(monitor_server_shutdown( - shutdown_signal, + shutdown_token.clone().cancelled_owned(), server_fut.clone(), in_flight_requests_data, )); @@ -562,7 +569,7 @@ async fn spawn_autopilot_worker_if_configured( Ok(Some( spawn_autopilot_worker( &gateway_handle.app_state.deferred_tasks, - gateway_handle.cancel_token.clone(), + gateway_handle.app_state.shutdown_token.clone(), config, ) .await diff --git a/internal/durable-tools/src/action.rs b/internal/durable-tools/src/action.rs index fc885244f20..190d067d152 100644 --- a/internal/durable-tools/src/action.rs +++ b/internal/durable-tools/src/action.rs @@ -154,6 +154,7 @@ pub async fn action( app_state.postgres_connection_info.clone(), app_state.valkey_connection_info.clone(), app_state.deferred_tasks.clone(), + app_state.shutdown_token.clone(), )?; let response = feedback(snapshot_app_state, *feedback_params, None).await?; @@ -168,6 +169,7 @@ pub async fn action( app_state.postgres_connection_info.clone(), app_state.valkey_connection_info.clone(), app_state.deferred_tasks.clone(), + app_state.shutdown_token.clone(), )?; // Run the evaluation using the shared helper diff --git a/tensorzero-core/src/endpoints/internal/autopilot.rs b/tensorzero-core/src/endpoints/internal/autopilot.rs index 61ff8ed24fc..016e2abe879 100644 --- a/tensorzero-core/src/endpoints/internal/autopilot.rs +++ b/tensorzero-core/src/endpoints/internal/autopilot.rs @@ -282,7 +282,14 @@ pub async fn stream_events_handler( }, ); - Ok(Sse::new(sse_stream).keep_alive(KeepAlive::new())) + // Close the stream when the server shuts down + // We do *not* want to wait for the stream to finish when the gateway shuts down, + // as it may stay open indefinitely (on either end - a browser ui ta might be holding open a connection) + // Clients using the endpoint should be able to auto-reconnect, so this is fine. + Ok( + Sse::new(sse_stream.take_until(app_state.shutdown_token.clone().cancelled_owned())) + .keep_alive(KeepAlive::new()), + ) } #[cfg(test)] @@ -293,6 +300,7 @@ mod tests { use crate::db::postgres::PostgresConnectionInfo; use crate::db::valkey::ValkeyConnectionInfo; use crate::http::TensorzeroHttpClient; + use tokio_util::sync::CancellationToken; use tokio_util::task::TaskTracker; fn make_test_app_state_without_autopilot() -> AppStateData { @@ -308,6 +316,7 @@ mod tests { postgres_connection_info, ValkeyConnectionInfo::Disabled, TaskTracker::new(), + CancellationToken::new(), ) .unwrap() } diff --git a/tensorzero-core/src/utils/gateway.rs b/tensorzero-core/src/utils/gateway.rs index 8d7cbb5011f..aa648a557f9 100644 --- a/tensorzero-core/src/utils/gateway.rs +++ b/tensorzero-core/src/utils/gateway.rs @@ -60,7 +60,6 @@ pub type DropWrapper = fn(Box); /// so that it's easy for us to tell where it gets dropped. pub struct GatewayHandle { pub app_state: AppStateData, - pub cancel_token: CancellationToken, drop_wrapper: Option, _private: (), } @@ -69,7 +68,7 @@ impl Drop for GatewayHandle { fn drop(&mut self) { let drop_wrapper = self.drop_wrapper.take(); let mut drop_self = || { - self.cancel_token.cancel(); + self.app_state.shutdown_token.cancel(); let handle = self .app_state .clickhouse_connection_info @@ -156,6 +155,7 @@ pub struct AppStateData { pub deployment_id: Option, /// Token pool manager for rate limiting pre-borrowing pub rate_limiting_manager: Arc, + pub shutdown_token: CancellationToken, // Prevent `AppStateData` from being directly constructed outside of this module // This ensures that `AppStateData` is only ever constructed via explicit `new` methods, // which can ensure that we update global state. @@ -264,9 +264,9 @@ impl GatewayHandle { autopilot_client: None, deployment_id: None, rate_limiting_manager, + shutdown_token: cancel_token, _private: (), }, - cancel_token, drop_wrapper: None, _private: (), } @@ -341,9 +341,9 @@ impl GatewayHandle { autopilot_client, deployment_id, rate_limiting_manager, + shutdown_token: cancel_token, _private: (), }, - cancel_token, drop_wrapper, _private: (), }) @@ -361,6 +361,7 @@ impl AppStateData { postgres_connection_info: PostgresConnectionInfo, valkey_connection_info: ValkeyConnectionInfo, deferred_tasks: TaskTracker, + shutdown_token: CancellationToken, ) -> Result { let rate_limiting_manager = Arc::new(RateLimitingManager::new_from_connections( Arc::new(config.rate_limiting.clone()), @@ -379,6 +380,7 @@ impl AppStateData { autopilot_client: None, deployment_id: None, rate_limiting_manager, + shutdown_token, _private: (), }) } From 0354925f4b47db5c93bb0a1ec90d6fcb6966dcbe Mon Sep 17 00:00:00 2001 From: Gabriel Bianconi <1275491+GabrielBianconi@users.noreply.github.com> Date: Tue, 27 Jan 2026 15:35:04 -0500 Subject: [PATCH 34/46] Improve error details in UI (#5904) --- ui/app/components/autopilot/EventStream.tsx | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/ui/app/components/autopilot/EventStream.tsx b/ui/app/components/autopilot/EventStream.tsx index 02ecc29a432..0e1f13505d0 100644 --- a/ui/app/components/autopilot/EventStream.tsx +++ b/ui/app/components/autopilot/EventStream.tsx @@ -165,8 +165,9 @@ function summarizeEvent(event: GatewayEvent): EventSummary { } return {}; case "error": - // TODO: handle errors - return {}; + return { + description: payload.message, + }; case "unknown": return {}; default: @@ -329,6 +330,7 @@ function EventItem({ const eventIsToolEvent = isToolEvent(event); const isExpandable = event.payload.type === "tool_call" || + event.payload.type === "error" || (event.payload.type === "tool_call_authorization" && event.payload.status.type === "rejected") || (event.payload.type === "tool_result" && @@ -384,7 +386,8 @@ function EventItem({ className={cn( "text-fg-secondary whitespace-pre-wrap", event.payload.type === "tool_call" || - event.payload.type === "tool_result" + event.payload.type === "tool_result" || + event.payload.type === "error" ? "font-mono text-sm" : "text-sm", )} From 337658c66704a08776e820de0a1ab8af59255621 Mon Sep 17 00:00:00 2001 From: Gabriel Bianconi <1275491+GabrielBianconi@users.noreply.github.com> Date: Tue, 27 Jan 2026 15:45:42 -0500 Subject: [PATCH 35/46] Print docker logs in mock-optimization-tests (#5913) --- .github/workflows/general.yml | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/.github/workflows/general.yml b/.github/workflows/general.yml index f8238ea8370..de73ca5ffc0 100644 --- a/.github/workflows/general.yml +++ b/.github/workflows/general.yml @@ -1065,6 +1065,10 @@ jobs: - name: Test (Rust) run: cargo test-optimization-mock + - name: Print docker compose logs + if: always() + run: docker compose -f tensorzero-core/tests/e2e/docker-compose.yml logs -t + - name: Print e2e logs if: always() run: cat e2e_logs.txt From 795488d31d40fc6d1545cfb0c7a70785e55803b0 Mon Sep 17 00:00:00 2001 From: Andrew Jesson <56548831+anndvision@users.noreply.github.com> Date: Tue, 27 Jan 2026 16:45:04 -0500 Subject: [PATCH 36/46] Add `list_datasets` Tool for Autopilot (#5914) * add list datasets tool * generate python schemas * no defaults, optional * robust param parsing * comment --- clients/python/tensorzero/generated_types.py | 54 +++++ clients/rust/src/lib.rs | 77 ++++++- internal/autopilot-tools/src/lib.rs | 4 + .../src/tools/prod/list_datasets.rs | 102 ++++++++ .../autopilot-tools/src/tools/prod/mod.rs | 2 + internal/autopilot-tools/tests/common/mod.rs | 32 ++- .../autopilot-tools/tests/datapoint_tools.rs | 217 +++++++++++++++++- internal/autopilot-worker/src/wrapper.rs | 10 +- .../src/tensorzero_client/client_ext.rs | 14 +- .../src/tensorzero_client/embedded.rs | 18 +- .../src/tensorzero_client/mod.rs | 12 +- .../lib/bindings/ListDatasetsRequest.ts | 21 ++ .../tensorzero-node/lib/bindings/index.ts | 1 + .../endpoints/datasets/v1/list_datasets.rs | 31 +-- .../src/endpoints/datasets/v1/mod.rs | 2 +- .../src/endpoints/datasets/v1/types.rs | 24 +- 16 files changed, 582 insertions(+), 39 deletions(-) create mode 100644 internal/autopilot-tools/src/tools/prod/list_datasets.rs create mode 100644 internal/tensorzero-node/lib/bindings/ListDatasetsRequest.ts diff --git a/clients/python/tensorzero/generated_types.py b/clients/python/tensorzero/generated_types.py index 7bb30146747..ef78cbd9e3a 100644 --- a/clients/python/tensorzero/generated_types.py +++ b/clients/python/tensorzero/generated_types.py @@ -70,6 +70,26 @@ class DatapointMetadataUpdate: """ +@dataclass(kw_only=True) +class DatasetMetadata: + """ + Metadata for a single dataset. + """ + + datapoint_count: int + """ + The total number of datapoints in the dataset. + """ + dataset_name: str + """ + The name of the dataset. + """ + last_updated: str + """ + The timestamp of the last update (ISO 8601 format). + """ + + @dataclass(kw_only=True) class DeleteDatapointsRequest: """ @@ -586,6 +606,40 @@ class JsonInferenceOutput: JsonMode = Literal["off", "on", "strict", "tool"] +@dataclass(kw_only=True) +class ListDatasetsRequest: + """ + Request to list datasets with optional filtering and pagination. + Used by the `GET /internal/datasets` endpoint. + """ + + function_name: str | None = None + """ + Optional function name to filter datasets by. + If provided, only datasets with datapoints for this function will be returned. + """ + limit: int | None = None + """ + The maximum number of datasets to return. + """ + offset: int | None = None + """ + The number of datasets to skip before starting to return results. + """ + + +@dataclass(kw_only=True) +class ListDatasetsResponse: + """ + Response containing a list of datasets. + """ + + datasets: list[DatasetMetadata] + """ + List of dataset metadata. + """ + + @dataclass(kw_only=True) class OpenAICustomToolFormatText: type: Literal["text"] = "text" diff --git a/clients/rust/src/lib.rs b/clients/rust/src/lib.rs index 8e57b943c01..1fcb62e87a2 100644 --- a/clients/rust/src/lib.rs +++ b/clients/rust/src/lib.rs @@ -59,11 +59,11 @@ pub use tensorzero_core::db::{ClickHouseConnection, ModelUsageTimePoint, TimeWin pub use tensorzero_core::endpoints::datasets::v1::types::{ CreateChatDatapointRequest, CreateDatapointRequest, CreateDatapointsFromInferenceRequest, CreateDatapointsFromInferenceRequestParams, CreateDatapointsRequest, CreateDatapointsResponse, - CreateJsonDatapointRequest, DeleteDatapointsRequest, DeleteDatapointsResponse, + CreateJsonDatapointRequest, DatasetMetadata, DeleteDatapointsRequest, DeleteDatapointsResponse, GetDatapointsRequest, GetDatapointsResponse, JsonDatapointOutputUpdate, ListDatapointsRequest, - UpdateChatDatapointRequest, UpdateDatapointMetadataRequest, UpdateDatapointRequest, - UpdateDatapointsMetadataRequest, UpdateDatapointsRequest, UpdateDatapointsResponse, - UpdateJsonDatapointRequest, + ListDatasetsRequest, ListDatasetsResponse, UpdateChatDatapointRequest, + UpdateDatapointMetadataRequest, UpdateDatapointRequest, UpdateDatapointsMetadataRequest, + UpdateDatapointsRequest, UpdateDatapointsResponse, UpdateJsonDatapointRequest, }; pub use tensorzero_core::endpoints::datasets::{ ChatInferenceDatapoint, Datapoint, DatapointKind, JsonInferenceDatapoint, @@ -361,6 +361,24 @@ pub trait ClientExt { params: CreateDatapointsFromInferenceRequestParams, ) -> Result; + /// Lists all datasets with optional filtering and pagination. + /// + /// # Arguments + /// + /// * `request` - The request parameters for listing datasets. + /// + /// # Returns + /// + /// A `ListDatasetsResponse` containing metadata for matching datasets. + /// + /// # Errors + /// + /// Returns a `TensorZeroError` if the request fails. + async fn list_datasets( + &self, + request: ListDatasetsRequest, + ) -> Result; + // ================================================================ // Workflow evaluation operations // ================================================================ @@ -1100,6 +1118,57 @@ impl ClientExt for Client { } } + async fn list_datasets( + &self, + request: ListDatasetsRequest, + ) -> Result { + match self.mode() { + ClientMode::HTTPGateway(client) => { + let ListDatasetsRequest { + function_name, + limit, + offset, + } = &request; + let mut url = client.base_url.join("internal/datasets").map_err(|e| { + TensorZeroError::Other { + source: Error::new(ErrorDetails::InvalidBaseUrl { + message: format!( + "Failed to join base URL with /internal/datasets endpoint: {e}" + ), + }) + .into(), + } + })?; + // Add query params + if let Some(function_name) = function_name { + url.query_pairs_mut() + .append_pair("function_name", function_name); + } + if let Some(limit) = limit { + url.query_pairs_mut() + .append_pair("limit", &limit.to_string()); + } + if let Some(offset) = offset { + url.query_pairs_mut() + .append_pair("offset", &offset.to_string()); + } + let builder = client.http_client.get(url); + Ok(client.send_and_parse_http_response(builder).await?.0) + } + ClientMode::EmbeddedGateway { gateway, timeout } => { + with_embedded_timeout(*timeout, async { + tensorzero_core::endpoints::datasets::v1::list_datasets( + &gateway.handle.app_state.clickhouse_connection_info, + request, + ) + .await + .map_err(err_to_http) + }) + .await + } + } + } + /// Query the Clickhouse database for inferences. /// /// This function is only available in EmbeddedGateway mode. diff --git a/internal/autopilot-tools/src/lib.rs b/internal/autopilot-tools/src/lib.rs index b6cbad192d0..fe414034fe3 100644 --- a/internal/autopilot-tools/src/lib.rs +++ b/internal/autopilot-tools/src/lib.rs @@ -14,6 +14,7 @@ //! - `FeedbackTool` - Submits feedback for inferences or episodes (comments, demonstrations, metrics) //! - `CreateDatapointsTool` - Creates datapoints in a dataset //! - `CreateDatapointsFromInferencesTool` - Creates datapoints from existing inferences +//! - `ListDatasetsTool` - Lists available datasets with metadata //! - `ListDatapointsTool` - Lists datapoints with filtering and pagination //! - `GetDatapointsTool` - Gets specific datapoints by ID //! - `UpdateDatapointsTool` - Updates existing datapoints @@ -145,6 +146,9 @@ pub async fn for_each_tool(visitor: &V) -> Result<(), V::Error> visitor .visit_simple_tool::() .await?; + visitor + .visit_simple_tool::() + .await?; visitor .visit_simple_tool::() .await?; diff --git a/internal/autopilot-tools/src/tools/prod/list_datasets.rs b/internal/autopilot-tools/src/tools/prod/list_datasets.rs new file mode 100644 index 00000000000..68772f61798 --- /dev/null +++ b/internal/autopilot-tools/src/tools/prod/list_datasets.rs @@ -0,0 +1,102 @@ +//! Tool for listing available datasets. + +use std::borrow::Cow; + +use async_trait::async_trait; +use durable_tools::{NonControlToolError, SimpleTool, SimpleToolContext, ToolMetadata, ToolResult}; + +use crate::error::AutopilotToolError; +use schemars::{JsonSchema, Schema}; +use serde::{Deserialize, Serialize}; +use tensorzero::{ListDatasetsRequest, ListDatasetsResponse}; + +use autopilot_client::AutopilotSideInfo; + +/// Parameters for the list_datasets tool (visible to LLM). +#[derive(Debug, Serialize, Deserialize, JsonSchema)] +pub struct ListDatasetsToolParams { + /// Request parameters for listing datasets (filtering, pagination). + #[serde(flatten)] + pub request: ListDatasetsRequest, +} + +/// Tool for listing available datasets. +/// +/// Returns metadata about datasets including name, datapoint count, and last updated timestamp. +/// Use this to discover what datasets are available before using list_datapoints to explore +/// the datapoints within a specific dataset. +#[derive(Default)] +pub struct ListDatasetsTool; + +impl ToolMetadata for ListDatasetsTool { + type SideInfo = AutopilotSideInfo; + type Output = ListDatasetsResponse; + type LlmParams = ListDatasetsToolParams; + + fn name(&self) -> Cow<'static, str> { + Cow::Borrowed("list_datasets") + } + + fn description(&self) -> Cow<'static, str> { + Cow::Borrowed( + "List available datasets with metadata (name, datapoint count, last updated). \ + Use this to discover datasets before exploring their datapoints with list_datapoints.", + ) + } + + // NOTE: This schema must be kept in sync with `ListDatasetsRequest` fields. + // We manually construct it for OpenAI structured outputs compatibility (using anyOf for optional fields). + fn parameters_schema(&self) -> ToolResult { + let schema = serde_json::json!({ + "type": "object", + "description": "List available datasets with optional filtering and pagination.", + "properties": { + "function_name": { + "anyOf": [ + { "type": "string" }, + { "type": "null" } + ], + "description": "Filter by function name. Only datasets with datapoints for this function will be returned." + }, + "limit": { + "anyOf": [ + { "type": "integer" }, + { "type": "null" } + ], + "description": "Maximum number of datasets to return." + }, + "offset": { + "anyOf": [ + { "type": "integer" }, + { "type": "null" } + ], + "description": "Number of datasets to skip for pagination." + } + }, + "required": [], + "additionalProperties": false + }); + + serde_json::from_value(schema).map_err(|e| { + NonControlToolError::SchemaGeneration { + message: e.to_string(), + } + .into() + }) + } +} + +#[async_trait] +impl SimpleTool for ListDatasetsTool { + async fn execute( + llm_params: ::LlmParams, + _side_info: ::SideInfo, + ctx: SimpleToolContext<'_>, + _idempotency_key: &str, + ) -> ToolResult<::Output> { + ctx.client() + .list_datasets(llm_params.request) + .await + .map_err(|e| AutopilotToolError::client_error("list_datasets", e).into()) + } +} diff --git a/internal/autopilot-tools/src/tools/prod/mod.rs b/internal/autopilot-tools/src/tools/prod/mod.rs index 572843547df..a1eec9f2733 100644 --- a/internal/autopilot-tools/src/tools/prod/mod.rs +++ b/internal/autopilot-tools/src/tools/prod/mod.rs @@ -16,6 +16,7 @@ mod get_latest_feedback_by_metric; mod inference; mod launch_optimization_workflow; mod list_datapoints; +mod list_datasets; mod list_inferences; mod run_evaluation; mod update_datapoints; @@ -41,6 +42,7 @@ pub use launch_optimization_workflow::{ LaunchOptimizationWorkflowToolParams, }; pub use list_datapoints::{ListDatapointsTool, ListDatapointsToolParams}; +pub use list_datasets::{ListDatasetsTool, ListDatasetsToolParams}; pub use list_inferences::{ListInferencesTool, ListInferencesToolParams}; pub use run_evaluation::{RunEvaluationTool, RunEvaluationToolParams}; pub use update_datapoints::{UpdateDatapointsTool, UpdateDatapointsToolParams}; diff --git a/internal/autopilot-tools/tests/common/mod.rs b/internal/autopilot-tools/tests/common/mod.rs index 3e39fa4a2d6..4521e945063 100644 --- a/internal/autopilot-tools/tests/common/mod.rs +++ b/internal/autopilot-tools/tests/common/mod.rs @@ -10,10 +10,11 @@ use mockall::mock; use sqlx::types::chrono::Utc; use tensorzero::{ ClientInferenceParams, CreateDatapointRequest, CreateDatapointsFromInferenceRequestParams, - CreateDatapointsResponse, DeleteDatapointsResponse, FeedbackParams, FeedbackResponse, - GetConfigResponse, GetDatapointsResponse, GetInferencesRequest, GetInferencesResponse, - InferenceResponse, ListDatapointsRequest, ListInferencesRequest, Role, StoredChatInference, - StoredInference, UpdateDatapointRequest, UpdateDatapointsResponse, Usage, WriteConfigRequest, + CreateDatapointsResponse, DatasetMetadata, DeleteDatapointsResponse, FeedbackParams, + FeedbackResponse, GetConfigResponse, GetDatapointsResponse, GetInferencesRequest, + GetInferencesResponse, InferenceResponse, ListDatapointsRequest, ListDatasetsRequest, + ListDatasetsResponse, ListInferencesRequest, Role, StoredChatInference, StoredInference, + UpdateDatapointRequest, UpdateDatapointsResponse, Usage, WriteConfigRequest, WriteConfigResponse, }; use tensorzero_core::config::snapshot::SnapshotHash; @@ -91,6 +92,11 @@ mock! { params: CreateDatapointsFromInferenceRequestParams, ) -> Result; + async fn list_datasets( + &self, + request: ListDatasetsRequest, + ) -> Result; + async fn list_datapoints( &self, dataset_name: String, @@ -298,3 +304,21 @@ pub fn create_mock_feedback_by_variant( count, } } + +/// Create a mock ListDatasetsResponse with the given datasets. +pub fn create_mock_list_datasets_response(datasets: Vec) -> ListDatasetsResponse { + ListDatasetsResponse { datasets } +} + +/// Create a mock DatasetMetadata for testing. +pub fn create_mock_dataset_metadata( + dataset_name: &str, + datapoint_count: u32, + last_updated: &str, +) -> DatasetMetadata { + DatasetMetadata { + dataset_name: dataset_name.to_string(), + datapoint_count, + last_updated: last_updated.to_string(), + } +} diff --git a/internal/autopilot-tools/tests/datapoint_tools.rs b/internal/autopilot-tools/tests/datapoint_tools.rs index 38ef296001b..a6e0c0c6eec 100644 --- a/internal/autopilot-tools/tests/datapoint_tools.rs +++ b/internal/autopilot-tools/tests/datapoint_tools.rs @@ -11,7 +11,7 @@ use durable_tools::{ErasedSimpleTool, SimpleToolContext, TensorZeroClientError}; use sqlx::PgPool; use tensorzero::{ CreateChatDatapointRequest, CreateDatapointRequest, CreateDatapointsFromInferenceRequestParams, - ListDatapointsRequest, UpdateChatDatapointRequest, UpdateDatapointRequest, + ListDatapointsRequest, ListDatasetsRequest, UpdateChatDatapointRequest, UpdateDatapointRequest, }; use uuid::Uuid; @@ -19,11 +19,13 @@ use autopilot_tools::tools::{ CreateDatapointsFromInferencesTool, CreateDatapointsFromInferencesToolParams, CreateDatapointsTool, CreateDatapointsToolParams, DeleteDatapointsTool, DeleteDatapointsToolParams, GetDatapointsTool, GetDatapointsToolParams, ListDatapointsTool, - ListDatapointsToolParams, UpdateDatapointsTool, UpdateDatapointsToolParams, + ListDatapointsToolParams, ListDatasetsTool, ListDatasetsToolParams, UpdateDatapointsTool, + UpdateDatapointsToolParams, }; use common::{ MockTensorZeroClient, create_mock_chat_datapoint, create_mock_create_datapoints_response, - create_mock_delete_datapoints_response, create_mock_get_datapoints_response, + create_mock_dataset_metadata, create_mock_delete_datapoints_response, + create_mock_get_datapoints_response, create_mock_list_datasets_response, create_mock_update_datapoints_response, create_test_input, }; @@ -771,3 +773,212 @@ async fn test_create_datapoints_from_inferences_tool_error(pool: PgPool) { assert!(result.is_err(), "Should return error when client fails"); } + +// ===== ListDatasetsTool Tests ===== + +#[sqlx::test(migrator = "MIGRATOR")] +async fn test_list_datasets_tool_basic(pool: PgPool) { + let dataset1 = create_mock_dataset_metadata("dataset_1", 100, "2024-01-01T00:00:00Z"); + let dataset2 = create_mock_dataset_metadata("dataset_2", 50, "2024-01-02T00:00:00Z"); + let mock_response = create_mock_list_datasets_response(vec![dataset1, dataset2]); + + let llm_params = ListDatasetsToolParams { + request: ListDatasetsRequest::default(), + }; + + let side_info = AutopilotSideInfo { + tool_call_event_id: Uuid::now_v7(), + session_id: Uuid::now_v7(), + config_snapshot_hash: "test_hash".to_string(), + optimization: OptimizationWorkflowSideInfo::default(), + }; + + let mut mock_client = MockTensorZeroClient::new(); + mock_client + .expect_list_datasets() + .return_once(move |_| Ok(mock_response)); + + let tool = ListDatasetsTool; + let t0_client: Arc = Arc::new(mock_client); + let ctx = SimpleToolContext::new(&pool, &t0_client); + + let result = tool + .execute_erased( + serde_json::to_value(&llm_params).expect("Failed to serialize llm_params"), + serde_json::to_value(&side_info).expect("Failed to serialize side_info"), + ctx, + "test-idempotency-key", + ) + .await + .expect("ListDatasetsTool execution should succeed"); + + assert!(result.is_object(), "Result should be a JSON object"); + let datasets = result.get("datasets").expect("Should have datasets field"); + assert!(datasets.is_array(), "datasets should be an array"); + assert_eq!( + datasets.as_array().unwrap().len(), + 2, + "Should have 2 datasets" + ); +} + +#[sqlx::test(migrator = "MIGRATOR")] +async fn test_list_datasets_tool_with_function_filter(pool: PgPool) { + let dataset = create_mock_dataset_metadata("filtered_dataset", 25, "2024-01-03T00:00:00Z"); + let mock_response = create_mock_list_datasets_response(vec![dataset]); + + let llm_params = ListDatasetsToolParams { + request: ListDatasetsRequest { + function_name: Some("specific_function".to_string()), + limit: None, + offset: None, + }, + }; + + let side_info = AutopilotSideInfo { + tool_call_event_id: Uuid::now_v7(), + session_id: Uuid::now_v7(), + config_snapshot_hash: "test_hash".to_string(), + optimization: OptimizationWorkflowSideInfo::default(), + }; + + let mut mock_client = MockTensorZeroClient::new(); + mock_client + .expect_list_datasets() + .withf(|request| request.function_name == Some("specific_function".to_string())) + .return_once(move |_| Ok(mock_response)); + + let tool = ListDatasetsTool; + let t0_client: Arc = Arc::new(mock_client); + let ctx = SimpleToolContext::new(&pool, &t0_client); + + let result = tool + .execute_erased( + serde_json::to_value(&llm_params).expect("Failed to serialize llm_params"), + serde_json::to_value(&side_info).expect("Failed to serialize side_info"), + ctx, + "test-idempotency-key", + ) + .await + .expect("ListDatasetsTool execution should succeed"); + + assert!(result.is_object(), "Result should be a JSON object"); +} + +#[sqlx::test(migrator = "MIGRATOR")] +async fn test_list_datasets_tool_with_pagination(pool: PgPool) { + let mock_response = create_mock_list_datasets_response(vec![]); + + let llm_params = ListDatasetsToolParams { + request: ListDatasetsRequest { + function_name: None, + limit: Some(10), + offset: Some(20), + }, + }; + + let side_info = AutopilotSideInfo { + tool_call_event_id: Uuid::now_v7(), + session_id: Uuid::now_v7(), + config_snapshot_hash: "test_hash".to_string(), + optimization: OptimizationWorkflowSideInfo::default(), + }; + + let mut mock_client = MockTensorZeroClient::new(); + mock_client + .expect_list_datasets() + .withf(|request| request.limit == Some(10) && request.offset == Some(20)) + .return_once(move |_| Ok(mock_response)); + + let tool = ListDatasetsTool; + let t0_client: Arc = Arc::new(mock_client); + let ctx = SimpleToolContext::new(&pool, &t0_client); + + let result = tool + .execute_erased( + serde_json::to_value(&llm_params).expect("Failed to serialize llm_params"), + serde_json::to_value(&side_info).expect("Failed to serialize side_info"), + ctx, + "test-idempotency-key", + ) + .await + .expect("ListDatasetsTool execution should succeed"); + + assert!(result.is_object(), "Result should be a JSON object"); +} + +#[sqlx::test(migrator = "MIGRATOR")] +async fn test_list_datasets_tool_empty_response(pool: PgPool) { + let mock_response = create_mock_list_datasets_response(vec![]); + + let llm_params = ListDatasetsToolParams { + request: ListDatasetsRequest::default(), + }; + + let side_info = AutopilotSideInfo { + tool_call_event_id: Uuid::now_v7(), + session_id: Uuid::now_v7(), + config_snapshot_hash: "test_hash".to_string(), + optimization: OptimizationWorkflowSideInfo::default(), + }; + + let mut mock_client = MockTensorZeroClient::new(); + mock_client + .expect_list_datasets() + .return_once(move |_| Ok(mock_response)); + + let tool = ListDatasetsTool; + let t0_client: Arc = Arc::new(mock_client); + let ctx = SimpleToolContext::new(&pool, &t0_client); + + let result = tool + .execute_erased( + serde_json::to_value(&llm_params).expect("Failed to serialize llm_params"), + serde_json::to_value(&side_info).expect("Failed to serialize side_info"), + ctx, + "test-idempotency-key", + ) + .await + .expect("ListDatasetsTool execution should succeed"); + + assert!(result.is_object(), "Result should be a JSON object"); + let datasets = result.get("datasets").expect("Should have datasets field"); + assert!( + datasets.as_array().unwrap().is_empty(), + "Datasets should be empty" + ); +} + +#[sqlx::test(migrator = "MIGRATOR")] +async fn test_list_datasets_tool_error(pool: PgPool) { + let llm_params = ListDatasetsToolParams { + request: ListDatasetsRequest::default(), + }; + + let side_info = AutopilotSideInfo { + tool_call_event_id: Uuid::now_v7(), + session_id: Uuid::now_v7(), + config_snapshot_hash: "test_hash".to_string(), + optimization: OptimizationWorkflowSideInfo::default(), + }; + + let mut mock_client = MockTensorZeroClient::new(); + mock_client + .expect_list_datasets() + .returning(|_| Err(TensorZeroClientError::AutopilotUnavailable)); + + let tool = ListDatasetsTool; + let t0_client: Arc = Arc::new(mock_client); + let ctx = SimpleToolContext::new(&pool, &t0_client); + + let result = tool + .execute_erased( + serde_json::to_value(&llm_params).expect("Failed to serialize llm_params"), + serde_json::to_value(&side_info).expect("Failed to serialize side_info"), + ctx, + "test-idempotency-key", + ) + .await; + + assert!(result.is_err(), "Should return error when client fails"); +} diff --git a/internal/autopilot-worker/src/wrapper.rs b/internal/autopilot-worker/src/wrapper.rs index 13759996d4c..f729e90d118 100644 --- a/internal/autopilot-worker/src/wrapper.rs +++ b/internal/autopilot-worker/src/wrapper.rs @@ -401,8 +401,9 @@ mod tests { ClientInferenceParams, CreateDatapointRequest, CreateDatapointsFromInferenceRequestParams, CreateDatapointsResponse, DeleteDatapointsResponse, FeedbackParams, FeedbackResponse, GetConfigResponse, GetDatapointsResponse, GetInferencesRequest, GetInferencesResponse, - InferenceResponse, ListDatapointsRequest, ListInferencesRequest, UpdateDatapointRequest, - UpdateDatapointsResponse, WriteConfigRequest, WriteConfigResponse, + InferenceResponse, ListDatapointsRequest, ListDatasetsRequest, ListDatasetsResponse, + ListInferencesRequest, UpdateDatapointRequest, UpdateDatapointsResponse, + WriteConfigRequest, WriteConfigResponse, }; use tensorzero_core::config::snapshot::SnapshotHash; use tensorzero_core::db::feedback::FeedbackByVariant; @@ -472,6 +473,11 @@ mod tests { params: CreateDatapointsFromInferenceRequestParams, ) -> Result; + async fn list_datasets( + &self, + request: ListDatasetsRequest, + ) -> Result; + async fn list_datapoints( &self, dataset_name: String, diff --git a/internal/durable-tools/src/tensorzero_client/client_ext.rs b/internal/durable-tools/src/tensorzero_client/client_ext.rs index d5bf53ee95e..a12d2a03473 100644 --- a/internal/durable-tools/src/tensorzero_client/client_ext.rs +++ b/internal/durable-tools/src/tensorzero_client/client_ext.rs @@ -12,8 +12,9 @@ use tensorzero::{ CreateDatapointsFromInferenceRequestParams, CreateDatapointsResponse, DeleteDatapointsResponse, FeedbackParams, FeedbackResponse, GetConfigResponse, GetDatapointsResponse, GetInferencesRequest, GetInferencesResponse, InferenceOutput, InferenceResponse, - ListDatapointsRequest, ListInferencesRequest, TensorZeroError, UpdateDatapointRequest, - UpdateDatapointsResponse, WriteConfigRequest, WriteConfigResponse, + ListDatapointsRequest, ListDatasetsRequest, ListDatasetsResponse, ListInferencesRequest, + TensorZeroError, UpdateDatapointRequest, UpdateDatapointsResponse, WriteConfigRequest, + WriteConfigResponse, }; use tensorzero_core::config::snapshot::SnapshotHash; use tensorzero_core::db::feedback::FeedbackByVariant; @@ -368,6 +369,15 @@ impl TensorZeroClient for Client { .map_err(TensorZeroClientError::TensorZero) } + async fn list_datasets( + &self, + request: ListDatasetsRequest, + ) -> Result { + ClientExt::list_datasets(self, request) + .await + .map_err(TensorZeroClientError::TensorZero) + } + async fn list_datapoints( &self, dataset_name: String, diff --git a/internal/durable-tools/src/tensorzero_client/embedded.rs b/internal/durable-tools/src/tensorzero_client/embedded.rs index 53f52df3979..d7f27bdf29e 100644 --- a/internal/durable-tools/src/tensorzero_client/embedded.rs +++ b/internal/durable-tools/src/tensorzero_client/embedded.rs @@ -9,9 +9,9 @@ use tensorzero::{ ClientInferenceParams, CreateDatapointRequest, CreateDatapointsFromInferenceRequestParams, CreateDatapointsResponse, DeleteDatapointsResponse, FeedbackParams, FeedbackResponse, GetConfigResponse, GetDatapointsResponse, GetInferencesRequest, GetInferencesResponse, - InferenceOutput, InferenceResponse, ListDatapointsRequest, ListInferencesRequest, - TensorZeroError, UpdateDatapointRequest, UpdateDatapointsResponse, WriteConfigRequest, - WriteConfigResponse, + InferenceOutput, InferenceResponse, ListDatapointsRequest, ListDatasetsRequest, + ListDatasetsResponse, ListInferencesRequest, TensorZeroError, UpdateDatapointRequest, + UpdateDatapointsResponse, WriteConfigRequest, WriteConfigResponse, }; use tensorzero_core::config::snapshot::{ConfigSnapshot, SnapshotHash}; use tensorzero_core::config::write_config_snapshot; @@ -278,6 +278,18 @@ impl TensorZeroClient for EmbeddedClient { .map_err(|e| TensorZeroClientError::TensorZero(TensorZeroError::Other { source: e.into() })) } + async fn list_datasets( + &self, + request: ListDatasetsRequest, + ) -> Result { + tensorzero_core::endpoints::datasets::v1::list_datasets( + &self.app_state.clickhouse_connection_info, + request, + ) + .await + .map_err(|e| TensorZeroClientError::TensorZero(TensorZeroError::Other { source: e.into() })) + } + async fn list_datapoints( &self, dataset_name: String, diff --git a/internal/durable-tools/src/tensorzero_client/mod.rs b/internal/durable-tools/src/tensorzero_client/mod.rs index 5d449164952..6fa652cb3d6 100644 --- a/internal/durable-tools/src/tensorzero_client/mod.rs +++ b/internal/durable-tools/src/tensorzero_client/mod.rs @@ -15,9 +15,9 @@ pub use tensorzero::{ Client, ClientBuilder, ClientBuilderError, ClientBuilderMode, ClientInferenceParams, CreateDatapointRequest, CreateDatapointsFromInferenceRequestParams, CreateDatapointsResponse, DeleteDatapointsResponse, FeedbackParams, FeedbackResponse, GetConfigResponse, - GetDatapointsResponse, InferenceResponse, ListDatapointsRequest, PostgresConfig, - TensorZeroError, UpdateDatapointRequest, UpdateDatapointsResponse, WriteConfigRequest, - WriteConfigResponse, + GetDatapointsResponse, InferenceResponse, ListDatapointsRequest, ListDatasetsRequest, + ListDatasetsResponse, PostgresConfig, TensorZeroError, UpdateDatapointRequest, + UpdateDatapointsResponse, WriteConfigRequest, WriteConfigResponse, }; use tensorzero::{GetInferencesRequest, GetInferencesResponse, ListInferencesRequest}; pub use tensorzero_core::cache::CacheEnabledMode; @@ -194,6 +194,12 @@ pub trait TensorZeroClient: Send + Sync + 'static { params: CreateDatapointsFromInferenceRequestParams, ) -> Result; + /// List all datasets with optional filtering and pagination. + async fn list_datasets( + &self, + request: ListDatasetsRequest, + ) -> Result; + /// List datapoints in a dataset with filtering and pagination. async fn list_datapoints( &self, diff --git a/internal/tensorzero-node/lib/bindings/ListDatasetsRequest.ts b/internal/tensorzero-node/lib/bindings/ListDatasetsRequest.ts new file mode 100644 index 00000000000..e31b364ad77 --- /dev/null +++ b/internal/tensorzero-node/lib/bindings/ListDatasetsRequest.ts @@ -0,0 +1,21 @@ +// This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually. + +/** + * Request to list datasets with optional filtering and pagination. + * Used by the `GET /internal/datasets` endpoint. + */ +export type ListDatasetsRequest = { + /** + * Optional function name to filter datasets by. + * If provided, only datasets with datapoints for this function will be returned. + */ + function_name?: string; + /** + * The maximum number of datasets to return. + */ + limit?: number; + /** + * The number of datasets to skip before starting to return results. + */ + offset?: number; +}; diff --git a/internal/tensorzero-node/lib/bindings/index.ts b/internal/tensorzero-node/lib/bindings/index.ts index f812eea9c84..c297584fc64 100644 --- a/internal/tensorzero-node/lib/bindings/index.ts +++ b/internal/tensorzero-node/lib/bindings/index.ts @@ -208,6 +208,7 @@ export * from "./LatestFeedbackIdByMetricResponse"; export * from "./LaunchOptimizationParams"; export * from "./LaunchOptimizationWorkflowParams"; export * from "./ListDatapointsRequest"; +export * from "./ListDatasetsRequest"; export * from "./ListDatasetsResponse"; export * from "./ListEpisodesResponse"; export * from "./ListEvaluationRunsResponse"; diff --git a/tensorzero-core/src/endpoints/datasets/v1/list_datasets.rs b/tensorzero-core/src/endpoints/datasets/v1/list_datasets.rs index ac58fa1afe1..fdceab01661 100644 --- a/tensorzero-core/src/endpoints/datasets/v1/list_datasets.rs +++ b/tensorzero-core/src/endpoints/datasets/v1/list_datasets.rs @@ -1,20 +1,13 @@ use axum::Json; use axum::extract::{Query, State}; -use serde::Deserialize; use tracing::instrument; +use crate::db::clickhouse::ClickHouseConnectionInfo; use crate::db::datasets::{DatasetQueries, GetDatasetMetadataParams}; use crate::error::Error; use crate::utils::gateway::{AppState, AppStateData}; -use super::types::{DatasetMetadata, ListDatasetsResponse}; - -#[derive(Debug, Deserialize)] -pub struct ListDatasetsQueryParams { - function_name: Option, - limit: Option, - offset: Option, -} +use super::types::{DatasetMetadata, ListDatasetsRequest, ListDatasetsResponse}; /// Handler for the GET `/internal/datasets` endpoint. /// Returns metadata for all datasets with optional filtering and pagination. @@ -22,17 +15,25 @@ pub struct ListDatasetsQueryParams { #[instrument(name = "datasets.v1.list_datasets", skip(app_state, params))] pub async fn list_datasets_handler( State(app_state): AppState, - Query(params): Query, + Query(params): Query, ) -> Result, Error> { + let response = list_datasets(&app_state.clickhouse_connection_info, params).await?; + Ok(Json(response)) +} + +/// List datasets with optional filtering and pagination. +/// +/// This is the non-handler function for use by the embedded client. +pub async fn list_datasets( + clickhouse: &ClickHouseConnectionInfo, + params: ListDatasetsRequest, +) -> Result { let db_params = GetDatasetMetadataParams { function_name: params.function_name, limit: params.limit, offset: params.offset, }; - let db_datasets = app_state - .clickhouse_connection_info - .get_dataset_metadata(&db_params) - .await?; + let db_datasets = clickhouse.get_dataset_metadata(&db_params).await?; // Convert from DB type to API type let datasets = db_datasets @@ -44,5 +45,5 @@ pub async fn list_datasets_handler( }) .collect(); - Ok(Json(ListDatasetsResponse { datasets })) + Ok(ListDatasetsResponse { datasets }) } diff --git a/tensorzero-core/src/endpoints/datasets/v1/mod.rs b/tensorzero-core/src/endpoints/datasets/v1/mod.rs index 87aaad93da9..a51dc903da0 100644 --- a/tensorzero-core/src/endpoints/datasets/v1/mod.rs +++ b/tensorzero-core/src/endpoints/datasets/v1/mod.rs @@ -17,7 +17,7 @@ pub use get_datapoints::{ get_datapoints, get_datapoints_by_dataset_handler, get_datapoints_handler, list_datapoints, list_datapoints_handler, }; -pub use list_datasets::list_datasets_handler; +pub use list_datasets::{list_datasets, list_datasets_handler}; pub use update_datapoints::{ update_datapoints, update_datapoints_handler, update_datapoints_metadata, update_datapoints_metadata_handler, diff --git a/tensorzero-core/src/endpoints/datasets/v1/types.rs b/tensorzero-core/src/endpoints/datasets/v1/types.rs index c33c3865e58..e0acc000eb9 100644 --- a/tensorzero-core/src/endpoints/datasets/v1/types.rs +++ b/tensorzero-core/src/endpoints/datasets/v1/types.rs @@ -695,10 +695,29 @@ pub struct DeleteDatapointsResponse { pub num_deleted_datapoints: u64, } +/// Request to list datasets with optional filtering and pagination. +/// Used by the `GET /internal/datasets` endpoint. +#[cfg_attr(feature = "ts-bindings", derive(ts_rs::TS))] +#[derive(Debug, Default, Serialize, Deserialize, JsonSchema)] +#[cfg_attr(feature = "ts-bindings", ts(export, optional_fields))] +#[export_schema] +pub struct ListDatasetsRequest { + /// Optional function name to filter datasets by. + /// If provided, only datasets with datapoints for this function will be returned. + pub function_name: Option, + + /// The maximum number of datasets to return. + pub limit: Option, + + /// The number of datasets to skip before starting to return results. + pub offset: Option, +} + /// Metadata for a single dataset. #[cfg_attr(feature = "ts-bindings", derive(ts_rs::TS))] -#[derive(Debug, Serialize, Deserialize)] +#[derive(Debug, Serialize, Deserialize, JsonSchema)] #[cfg_attr(feature = "ts-bindings", ts(export))] +#[export_schema] pub struct DatasetMetadata { /// The name of the dataset. pub dataset_name: String, @@ -710,8 +729,9 @@ pub struct DatasetMetadata { /// Response containing a list of datasets. #[cfg_attr(feature = "ts-bindings", derive(ts_rs::TS))] -#[derive(Debug, Serialize, Deserialize)] +#[derive(Debug, Serialize, Deserialize, JsonSchema)] #[cfg_attr(feature = "ts-bindings", ts(export))] +#[export_schema] pub struct ListDatasetsResponse { /// List of dataset metadata. pub datasets: Vec, From 90313b982eb1283f71c4a11f805e80b15ac87392 Mon Sep 17 00:00:00 2001 From: Shuyang Li Date: Tue, 27 Jan 2026 16:59:40 -0500 Subject: [PATCH 37/46] Fix SQL injection in human_feedback.rs by adding get_inference_output to InferenceQueries (#5837) The get_output() function in human_feedback.rs used format!() with direct string interpolation for inference_id, episode_id, function_name, variant_name, and table_name. This is a potential SQL injection vulnerability. Added InferenceQueries::get_inference_output() trait method with a proper parameterized query implementation in ClickHouse. Added e2e tests: - test_get_inference_output_chat_inference - test_get_inference_output_json_inference - test_get_inference_output_not_found --- .../src/db/clickhouse/inference_queries.rs | 347 ++++++++++++++++++ .../mock_clickhouse_connection_info.rs | 10 + tensorzero-core/src/db/inferences.rs | 12 + .../src/endpoints/feedback/human_feedback.rs | 46 +-- .../tests/e2e/db/inference_queries.rs | 145 +++++++- 5 files changed, 521 insertions(+), 39 deletions(-) diff --git a/tensorzero-core/src/db/clickhouse/inference_queries.rs b/tensorzero-core/src/db/clickhouse/inference_queries.rs index d14dc3ad31d..86270069531 100644 --- a/tensorzero-core/src/db/clickhouse/inference_queries.rs +++ b/tensorzero-core/src/db/clickhouse/inference_queries.rs @@ -268,6 +268,58 @@ impl InferenceQueries for ClickHouseConnectionInfo { Ok(Some(output_schema)) } + + async fn get_inference_output( + &self, + function_info: &FunctionInfo, + inference_id: Uuid, + ) -> Result, Error> { + let table_name = function_info.function_type.inference_table_name(); + + // Build the query with parameterized values + let query = format!( + r" + SELECT output + FROM {table_name} + WHERE + id = {{inference_id:String}} AND + episode_id = {{episode_id:UUID}} AND + function_name = {{function_name:String}} AND + variant_name = {{variant_name:String}} + LIMIT 1 + FORMAT JSONEachRow + SETTINGS max_threads=1 + " + ); + + let inference_id_str = inference_id.to_string(); + let episode_id_str = function_info.episode_id.to_string(); + let params = HashMap::from([ + ("inference_id", inference_id_str.as_str()), + ("episode_id", episode_id_str.as_str()), + ("function_name", function_info.function_name.as_str()), + ("variant_name", function_info.variant_name.as_str()), + ]); + + let response = self.inner.run_query_synchronous(query, ¶ms).await?; + + if response.response.is_empty() { + return Ok(None); + } + + #[derive(Debug, Deserialize)] + struct OutputResult { + output: String, + } + + let result: OutputResult = serde_json::from_str(&response.response).map_err(|e| { + Error::new(ErrorDetails::ClickHouseDeserialization { + message: format!("Failed to parse output result: {e}"), + }) + })?; + + Ok(Some(result.output)) + } } /// Generates the SQL query and parameters for counting inferences. @@ -1825,6 +1877,301 @@ mod tests { } } + mod get_inference_output_tests { + use crate::db::clickhouse::clickhouse_client::MockClickHouseClient; + use crate::db::clickhouse::query_builder::test_util::assert_query_contains; + use crate::db::clickhouse::{ + ClickHouseConnectionInfo, ClickHouseResponse, ClickHouseResponseMetadata, + }; + use crate::db::inferences::{FunctionInfo, InferenceQueries}; + use crate::inference::types::FunctionType; + use std::sync::Arc; + use uuid::Uuid; + + #[tokio::test] + async fn test_get_inference_output_chat_inference_success() { + let inference_id = Uuid::now_v7(); + let episode_id = Uuid::now_v7(); + let function_info = FunctionInfo { + function_name: "test_chat_function".to_string(), + function_type: FunctionType::Chat, + variant_name: "test_variant".to_string(), + episode_id, + }; + + let mut mock = MockClickHouseClient::new(); + mock.expect_run_query_synchronous() + .withf(move |query, params| { + // Verify query targets ChatInference table + assert_query_contains(query, "FROM ChatInference"); + // Verify parameterized WHERE clause + assert_query_contains(query, "id = {inference_id:String}"); + assert_query_contains(query, "episode_id = {episode_id:UUID}"); + assert_query_contains(query, "function_name = {function_name:String}"); + assert_query_contains(query, "variant_name = {variant_name:String}"); + // Verify parameters are set + assert_eq!( + params.get("function_name"), + Some(&"test_chat_function"), + "function_name parameter should be set" + ); + assert_eq!( + params.get("variant_name"), + Some(&"test_variant"), + "variant_name parameter should be set" + ); + true + }) + .returning(|_, _| { + Ok(ClickHouseResponse { + response: r#"{"output":"[{\"type\":\"text\",\"text\":\"Hello!\"}]"}"# + .to_string(), + metadata: ClickHouseResponseMetadata { + read_rows: 1, + written_rows: 0, + }, + }) + }); + + let conn = ClickHouseConnectionInfo::new_mock(Arc::new(mock)); + let result = conn + .get_inference_output(&function_info, inference_id) + .await + .expect("Should succeed"); + + assert!( + result.is_some(), + "Should return Some for existing inference" + ); + assert_eq!( + result.unwrap(), + r#"[{"type":"text","text":"Hello!"}]"#, + "Should return the output string" + ); + } + + #[tokio::test] + async fn test_get_inference_output_json_inference_success() { + let inference_id = Uuid::now_v7(); + let episode_id = Uuid::now_v7(); + let function_info = FunctionInfo { + function_name: "test_json_function".to_string(), + function_type: FunctionType::Json, + variant_name: "json_variant".to_string(), + episode_id, + }; + + let mut mock = MockClickHouseClient::new(); + mock.expect_run_query_synchronous() + .withf(move |query, params| { + // Verify query targets JsonInference table + assert_query_contains(query, "FROM JsonInference"); + // Verify parameters + assert_eq!( + params.get("function_name"), + Some(&"test_json_function"), + "function_name parameter should be set" + ); + assert_eq!( + params.get("variant_name"), + Some(&"json_variant"), + "variant_name parameter should be set" + ); + true + }) + .returning(|_, _| { + Ok(ClickHouseResponse { + response: r#"{"output":"{\"raw\":\"{\\\"score\\\":0.95}\",\"parsed\":{\"score\":0.95}}"}"# + .to_string(), + metadata: ClickHouseResponseMetadata { + read_rows: 1, + written_rows: 0, + }, + }) + }); + + let conn = ClickHouseConnectionInfo::new_mock(Arc::new(mock)); + let result = conn + .get_inference_output(&function_info, inference_id) + .await + .expect("Should succeed"); + + assert!( + result.is_some(), + "Should return Some for existing JSON inference" + ); + assert!( + result.unwrap().contains("score"), + "Should return the JSON output" + ); + } + + #[tokio::test] + async fn test_get_inference_output_not_found() { + let inference_id = Uuid::now_v7(); + let episode_id = Uuid::now_v7(); + let function_info = FunctionInfo { + function_name: "nonexistent_function".to_string(), + function_type: FunctionType::Chat, + variant_name: "nonexistent_variant".to_string(), + episode_id, + }; + + let mut mock = MockClickHouseClient::new(); + mock.expect_run_query_synchronous().returning(|_, _| { + Ok(ClickHouseResponse { + response: String::new(), // Empty response indicates not found + metadata: ClickHouseResponseMetadata { + read_rows: 0, + written_rows: 0, + }, + }) + }); + + let conn = ClickHouseConnectionInfo::new_mock(Arc::new(mock)); + let result = conn + .get_inference_output(&function_info, inference_id) + .await + .expect("Should succeed even when not found"); + + assert!( + result.is_none(), + "Should return None for non-existent inference" + ); + } + + #[tokio::test] + async fn test_get_inference_output_uses_all_parameters_for_security() { + let inference_id = Uuid::now_v7(); + let episode_id = Uuid::now_v7(); + let function_info = FunctionInfo { + function_name: "secure_function".to_string(), + function_type: FunctionType::Chat, + variant_name: "secure_variant".to_string(), + episode_id, + }; + + let mut mock = MockClickHouseClient::new(); + mock.expect_run_query_synchronous() + .withf(move |query, params| { + // Verify ALL parameters are used in WHERE clause for security + // This prevents unauthorized access by requiring all identifiers to match + assert_query_contains(query, "id = {inference_id:String}"); + assert_query_contains(query, "episode_id = {episode_id:UUID}"); + assert_query_contains(query, "function_name = {function_name:String}"); + assert_query_contains(query, "variant_name = {variant_name:String}"); + + // Verify all parameters are bound (not interpolated) + assert!( + params.contains_key("inference_id"), + "inference_id should be a bound parameter" + ); + assert!( + params.contains_key("episode_id"), + "episode_id should be a bound parameter" + ); + assert!( + params.contains_key("function_name"), + "function_name should be a bound parameter" + ); + assert!( + params.contains_key("variant_name"), + "variant_name should be a bound parameter" + ); + + // Verify no string interpolation in the query (SQL injection prevention) + assert!( + !query.contains("secure_function"), + "function_name should NOT be interpolated into query" + ); + assert!( + !query.contains("secure_variant"), + "variant_name should NOT be interpolated into query" + ); + true + }) + .returning(|_, _| { + Ok(ClickHouseResponse { + response: r#"{"output":"test output"}"#.to_string(), + metadata: ClickHouseResponseMetadata { + read_rows: 1, + written_rows: 0, + }, + }) + }); + + let conn = ClickHouseConnectionInfo::new_mock(Arc::new(mock)); + let result = conn + .get_inference_output(&function_info, inference_id) + .await; + + assert!(result.is_ok(), "Query should execute successfully"); + } + + #[tokio::test] + async fn test_get_inference_output_table_selection_by_function_type() { + // Test Chat function type uses ChatInference table + let chat_function_info = FunctionInfo { + function_name: "chat_func".to_string(), + function_type: FunctionType::Chat, + variant_name: "v1".to_string(), + episode_id: Uuid::now_v7(), + }; + + let mut chat_mock = MockClickHouseClient::new(); + chat_mock + .expect_run_query_synchronous() + .withf(|query, _| { + assert_query_contains(query, "FROM ChatInference"); + true + }) + .returning(|_, _| { + Ok(ClickHouseResponse { + response: String::new(), + metadata: ClickHouseResponseMetadata { + read_rows: 0, + written_rows: 0, + }, + }) + }); + + let conn = ClickHouseConnectionInfo::new_mock(Arc::new(chat_mock)); + let _ = conn + .get_inference_output(&chat_function_info, Uuid::now_v7()) + .await; + + // Test Json function type uses JsonInference table + let json_function_info = FunctionInfo { + function_name: "json_func".to_string(), + function_type: FunctionType::Json, + variant_name: "v1".to_string(), + episode_id: Uuid::now_v7(), + }; + + let mut json_mock = MockClickHouseClient::new(); + json_mock + .expect_run_query_synchronous() + .withf(|query, _| { + assert_query_contains(query, "FROM JsonInference"); + true + }) + .returning(|_, _| { + Ok(ClickHouseResponse { + response: String::new(), + metadata: ClickHouseResponseMetadata { + read_rows: 0, + written_rows: 0, + }, + }) + }); + + let conn = ClickHouseConnectionInfo::new_mock(Arc::new(json_mock)); + let _ = conn + .get_inference_output(&json_function_info, Uuid::now_v7()) + .await; + } + } + mod list_inference_metadata_tests { use crate::db::clickhouse::clickhouse_client::MockClickHouseClient; use crate::db::clickhouse::query_builder::test_util::{ diff --git a/tensorzero-core/src/db/clickhouse/mock_clickhouse_connection_info.rs b/tensorzero-core/src/db/clickhouse/mock_clickhouse_connection_info.rs index 1d92bcb2234..d6b8c82b56e 100644 --- a/tensorzero-core/src/db/clickhouse/mock_clickhouse_connection_info.rs +++ b/tensorzero-core/src/db/clickhouse/mock_clickhouse_connection_info.rs @@ -114,6 +114,16 @@ impl InferenceQueries for MockClickHouseConnectionInfo { .get_json_inference_output_schema(function_name, inference_id) .await } + + async fn get_inference_output( + &self, + function_info: &FunctionInfo, + inference_id: Uuid, + ) -> Result, Error> { + self.inference_queries + .get_inference_output(function_info, inference_id) + .await + } } #[async_trait] diff --git a/tensorzero-core/src/db/inferences.rs b/tensorzero-core/src/db/inferences.rs index 46935c9d3f7..b1e80e073af 100644 --- a/tensorzero-core/src/db/inferences.rs +++ b/tensorzero-core/src/db/inferences.rs @@ -381,4 +381,16 @@ pub trait InferenceQueries { function_name: &str, inference_id: Uuid, ) -> Result, Error>; + + /// Get the output string from an inference for human feedback context. + /// + /// Returns the serialized output of the inference, which is needed when + /// writing static evaluation human feedback records. + /// + /// Returns `None` if the inference doesn't exist. + async fn get_inference_output( + &self, + function_info: &FunctionInfo, + inference_id: Uuid, + ) -> Result, Error>; } diff --git a/tensorzero-core/src/endpoints/feedback/human_feedback.rs b/tensorzero-core/src/endpoints/feedback/human_feedback.rs index b888d32a30b..172a99efaa5 100644 --- a/tensorzero-core/src/endpoints/feedback/human_feedback.rs +++ b/tensorzero-core/src/endpoints/feedback/human_feedback.rs @@ -1,5 +1,5 @@ use super::throttled_get_function_info; -use crate::db::inferences::FunctionInfo; +use crate::db::inferences::{FunctionInfo, InferenceQueries}; use crate::{ config::MetricConfigLevel, db::clickhouse::{ClickHouseConnectionInfo, TableName}, @@ -40,7 +40,14 @@ pub(super) async fn write_static_evaluation_human_feedback_if_necessary( .await? } }; - let output = get_output(clickhouse, &function_info, target_id).await?; + let output = clickhouse + .get_inference_output(&function_info, target_id) + .await? + .ok_or_else(|| { + Error::new(ErrorDetails::InferenceNotFound { + inference_id: target_id, + }) + })?; let row = StaticEvaluationHumanFeedback { output, feedback_id, @@ -101,41 +108,6 @@ struct InferenceEvaluationInfo { evaluator_inference_id: Uuid, } -async fn get_output( - clickhouse: &ClickHouseConnectionInfo, - function_info: &FunctionInfo, - inference_id: Uuid, -) -> Result { - let FunctionInfo { - function_type, - episode_id, - function_name, - variant_name, - } = function_info; - let table_name = function_type.inference_table_name(); - let output: OutputResponse = clickhouse - .run_query_synchronous_no_params_de(format!( - r" - SELECT output FROM {table_name} - WHERE - id = '{inference_id}' AND - episode_id = '{episode_id}' AND - function_name = '{function_name}' AND - variant_name = '{variant_name}' - LIMIT 1 - FORMAT JSONEachRow - SETTINGS max_threads=1" - )) - .await?; - Ok(output.output) -} - -/// This is so we're absolutely sure things are escaped properly. -#[derive(Debug, Deserialize)] -struct OutputResponse { - output: String, -} - /// Represents a row in the StaticEvaluationHumanFeedback database table. /// /// Note: The "Static" prefix is retained for backward compatibility with existing diff --git a/tensorzero-core/tests/e2e/db/inference_queries.rs b/tensorzero-core/tests/e2e/db/inference_queries.rs index bc8723492e1..cb3a44348e1 100644 --- a/tensorzero-core/tests/e2e/db/inference_queries.rs +++ b/tensorzero-core/tests/e2e/db/inference_queries.rs @@ -1,14 +1,16 @@ use std::path::Path; use tensorzero_core::{ - config::{Config, ConfigFileGlob}, + config::{Config, ConfigFileGlob, MetricConfigLevel}, db::{ clickhouse::{query_builder::InferenceFilter, test_helpers::get_clickhouse}, feedback::{FeedbackQueries, FeedbackRow}, - inferences::{InferenceOutputSource, InferenceQueries, ListInferencesParams}, + inferences::{FunctionInfo, InferenceOutputSource, InferenceQueries, ListInferencesParams}, }, endpoints::stored_inferences::v1::types::DemonstrationFeedbackFilter, + inference::types::FunctionType, }; +use uuid::Uuid; async fn get_e2e_config() -> Config { Config::load_from_path_optional_verify_credentials( @@ -115,3 +117,142 @@ async fn test_list_inferences_filtered_to_no_demonstrations() { ); } } + +#[tokio::test(flavor = "multi_thread")] +async fn test_get_inference_output_chat_inference() { + let config = get_e2e_config().await; + let clickhouse = get_clickhouse().await; + + // First, list some chat inferences to get a valid inference_id + let inferences = clickhouse + .list_inferences( + &config, + &ListInferencesParams { + function_name: Some("write_haiku"), + output_source: InferenceOutputSource::Inference, + limit: 1, + ..Default::default() + }, + ) + .await + .unwrap(); + + assert!( + !inferences.is_empty(), + "Should have at least one chat inference" + ); + + let inference = &inferences[0]; + let inference_id = inference.id(); + + // Get function info for this inference + let function_info = clickhouse + .get_function_info(&inference_id, MetricConfigLevel::Inference) + .await + .unwrap() + .expect("Should find function info for existing inference"); + + assert_eq!( + function_info.function_type, + FunctionType::Chat, + "write_haiku should be a Chat function" + ); + + // Now test get_inference_output + let output = clickhouse + .get_inference_output(&function_info, inference_id) + .await + .unwrap(); + + assert!( + output.is_some(), + "Should return output for existing inference" + ); + let output_str = output.unwrap(); + assert!( + !output_str.is_empty(), + "Output should not be empty for existing inference" + ); +} + +#[tokio::test(flavor = "multi_thread")] +async fn test_get_inference_output_json_inference() { + let config = get_e2e_config().await; + let clickhouse = get_clickhouse().await; + + // First, list some json inferences to get a valid inference_id + let inferences = clickhouse + .list_inferences( + &config, + &ListInferencesParams { + function_name: Some("extract_entities"), + output_source: InferenceOutputSource::Inference, + limit: 1, + ..Default::default() + }, + ) + .await + .unwrap(); + + assert!( + !inferences.is_empty(), + "Should have at least one json inference" + ); + + let inference = &inferences[0]; + let inference_id = inference.id(); + + // Get function info for this inference + let function_info = clickhouse + .get_function_info(&inference_id, MetricConfigLevel::Inference) + .await + .unwrap() + .expect("Should find function info for existing inference"); + + assert_eq!( + function_info.function_type, + FunctionType::Json, + "extract_entities should be a Json function" + ); + + // Now test get_inference_output + let output = clickhouse + .get_inference_output(&function_info, inference_id) + .await + .unwrap(); + + assert!( + output.is_some(), + "Should return output for existing json inference" + ); + let output_str = output.unwrap(); + assert!( + !output_str.is_empty(), + "Output should not be empty for existing json inference" + ); +} + +#[tokio::test(flavor = "multi_thread")] +async fn test_get_inference_output_not_found() { + let clickhouse = get_clickhouse().await; + + // Create a fake function_info with a non-existent inference_id + let fake_function_info = FunctionInfo { + function_name: "write_haiku".to_string(), + function_type: FunctionType::Chat, + variant_name: "test_variant".to_string(), + episode_id: Uuid::now_v7(), + }; + + let non_existent_inference_id = Uuid::now_v7(); + + let output = clickhouse + .get_inference_output(&fake_function_info, non_existent_inference_id) + .await + .unwrap(); + + assert!( + output.is_none(), + "Should return None for non-existent inference" + ); +} From 3ba2675c28c69c7740ef5bed5f3b93173fc6b717 Mon Sep 17 00:00:00 2001 From: Aaron Hill Date: Tue, 27 Jan 2026 17:14:14 -0500 Subject: [PATCH 38/46] Try a different disk space cleanup action in flaky job (#5919) --- .github/workflows/ui-tests-e2e-model-inference-cache.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/ui-tests-e2e-model-inference-cache.yml b/.github/workflows/ui-tests-e2e-model-inference-cache.yml index f0b4295f3f0..345c7187aaa 100644 --- a/.github/workflows/ui-tests-e2e-model-inference-cache.yml +++ b/.github/workflows/ui-tests-e2e-model-inference-cache.yml @@ -58,7 +58,7 @@ jobs: uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 - name: Cleanup disk space - run: ./ci/free-disk-space.sh + uses: jlumbroso/free-disk-space@54081f138730dfa15788a46383842cd2f914a1be - name: Setup Node uses: actions/setup-node@49933ea5288caeca8642d1e84afbd3f7d6820020 From fd0b9cfe5101c0d3b0a957d3b3e9b33d7d78b4ab Mon Sep 17 00:00:00 2001 From: Shuyang Li Date: Tue, 27 Jan 2026 17:16:59 -0500 Subject: [PATCH 39/46] E2e tests for endpoints setup for postgres (#5800) * E2e tests for endpoints setup for postgres * Fix tests --- .cargo/config.toml | 8 ++ .config/nextest.toml | 7 ++ .github/workflows/general.yml | 118 ++++++++++++++++++ tensorzero-core/tests/e2e/common.rs | 25 ++++ .../e2e/config/tensorzero.functions.toml | 15 +++ .../endpoints/datasets/clone_datapoints.rs | 5 + .../endpoints/datasets/create_datapoints.rs | 12 ++ .../datasets/create_from_inferences.rs | 5 + .../endpoints/datasets/delete_datapoints.rs | 7 ++ .../e2e/endpoints/datasets/delete_dataset.rs | 6 + .../e2e/endpoints/datasets/get_datapoints.rs | 17 +++ .../e2e/endpoints/datasets/list_datasets.rs | 4 + .../endpoints/datasets/update_datapoints.rs | 56 +++++++-- .../tests/e2e/endpoints/internal/action.rs | 3 + .../tests/e2e/endpoints/internal/config.rs | 5 + .../e2e/endpoints/internal/datapoint_count.rs | 2 + .../tests/e2e/endpoints/internal/episodes.rs | 7 ++ .../e2e/endpoints/internal/evaluations.rs | 34 +++++ .../tests/e2e/endpoints/internal/feedback.rs | 27 ++++ .../tests/e2e/endpoints/internal/functions.rs | 10 ++ .../e2e/endpoints/internal/inference_count.rs | 17 ++- .../endpoints/internal/inference_metadata.rs | 9 ++ .../endpoints/internal/model_inferences.rs | 4 + .../tests/e2e/endpoints/internal/models.rs | 5 + .../internal/workflow_evaluations.rs | 29 +++++ tensorzero-core/tests/e2e/endpoints/mod.rs | 3 + .../stored_inferences/get_inferences.rs | 72 ++++++++--- tensorzero-core/tests/e2e/tests.rs | 1 + 28 files changed, 477 insertions(+), 36 deletions(-) diff --git a/.cargo/config.toml b/.cargo/config.toml index f8221b2ad47..802faf4ea48 100644 --- a/.cargo/config.toml +++ b/.cargo/config.toml @@ -87,6 +87,14 @@ test-clickhouse-fast = [ "0", "--no-fail-fast", ] +test-endpoints = [ + "nextest", + "run", + "--features", + "e2e_tests", + "--profile", + "endpoints", +] test-rate-limit-load = [ "run", "--release", diff --git a/.config/nextest.toml b/.config/nextest.toml index 1228b08fe8e..5bde8e94a34 100644 --- a/.config/nextest.toml +++ b/.config/nextest.toml @@ -62,6 +62,13 @@ default-filter = 'test(batch) and test(mock)' default-filter = 'test(test_dummy_only) | test(clickhouse::) | (test(db::) & !test(postgres) & !test(valkey))' junit.path = "junit.xml" +# Run only endpoint e2e tests (tests under tests/e2e/endpoints/) +# These run against both ClickHouse and Postgres-backed TensorZero gateway. +[profile.endpoints] +default-filter = 'binary(e2e) and test(endpoints::)' +junit.path = "junit.xml" +retries = { backoff = "exponential", count = 4, delay = "5s", jitter = true, max-delay = "60s" } + [profile.optimization] default-filter = 'binary(optimization-live)' diff --git a/.github/workflows/general.yml b/.github/workflows/general.yml index de73ca5ffc0..c3ec01c5a37 100644 --- a/.github/workflows/general.yml +++ b/.github/workflows/general.yml @@ -985,6 +985,123 @@ jobs: if: always() run: docker compose -f tensorzero-core/tests/e2e/docker-compose.yml logs -t + # Run endpoint tests against both ClickHouse and Postgres backends + endpoint-tests: + name: "Endpoint tests (database: ${{ matrix.database }})" + permissions: + contents: read + actions: read + # Permission to fetch GitHub OIDC token authentication for Namespace login + id-token: write + needs: [build-gateway-container, build-mock-provider-api-container] + runs-on: ubuntu-latest + if: ${{ github.event_name == 'merge_group' || (github.event.pull_request.head.repo.full_name == github.repository) }} + strategy: + fail-fast: false + matrix: + database: [clickhouse, postgres] + steps: + - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 + + - name: Cleanup disk space + run: ./ci/free-disk-space.sh + + - name: Install Rust toolchain + run: | + for attempt in 1 2 3; do + if rustup toolchain install stable && rustup default stable; then + break + fi + if [ $attempt -eq 3 ]; then + echo "Failed to install Rust toolchain after 3 attempts" + exit 1 + fi + sleep $((10 * attempt)) + done + shell: bash + + - uses: Swatinem/rust-cache@9d47c6ad4b02e050fd481d890b2ea34778fd09d6 + with: + cache-provider: "buildjet" + shared-key: "build-gateway-cache" + save-if: false + + - name: Install cargo-nextest + uses: taiki-e/install-action@60581cd7025e0e855cebd745379013e286d9c787 + with: + tool: cargo-nextest + + - name: Install Namespace CLI + uses: namespacelabs/nscloud-setup@d1c625762f7c926a54bd39252efff0705fd11c64 + continue-on-error: ${{ github.event.pull_request.head.repo.full_name != github.repository || github.actor == 'dependabot[bot]' }} + + - name: Install uv + uses: astral-sh/setup-uv@ed21f2f24f8dd64503750218de024bcf64c7250a + with: + version: "0.6.17" + + - name: Download ClickHouse fixtures + run: | + uv run ./ui/fixtures/download-small-fixtures.py + + - name: Login to Namespace registry + run: nsc docker login + continue-on-error: ${{ github.event.pull_request.head.repo.full_name != github.repository || github.actor == 'dependabot[bot]' }} + + - name: Download container images + run: | + docker pull nscr.io/igvf4asmf8kri/gateway:sha-${{ github.sha }} + docker pull nscr.io/igvf4asmf8kri/mock-provider-api:sha-${{ github.sha }} + docker tag nscr.io/igvf4asmf8kri/gateway:sha-${{ github.sha }} tensorzero/gateway:sha-${{ github.sha }} + docker tag nscr.io/igvf4asmf8kri/mock-provider-api:sha-${{ github.sha }} tensorzero/mock-provider-api:sha-${{ github.sha }} + continue-on-error: ${{ github.event.pull_request.head.repo.full_name != github.repository || github.actor == 'dependabot[bot]' }} + + - name: Set environment variables + run: | + echo "TENSORZERO_CLICKHOUSE_URL=http://chuser:chpassword@localhost:8123/tensorzero_e2e_tests" >> $GITHUB_ENV + echo "TENSORZERO_POSTGRES_URL=postgres://postgres:postgres@localhost:5432/tensorzero-e2e-tests" >> $GITHUB_ENV + echo "TENSORZERO_GATEWAY_TAG=sha-${{ github.sha }}" >> $GITHUB_ENV + echo "TENSORZERO_MOCK_PROVIDER_API_TAG=sha-${{ github.sha }}" >> $GITHUB_ENV + echo "TENSORZERO_SKIP_LARGE_FIXTURES=1" >> $GITHUB_ENV + + - name: Set Postgres read and write feature flags + if: matrix.database == 'postgres' + run: | + echo "TENSORZERO_INTERNAL_FLAG_ENABLE_POSTGRES_READ=1" >> $GITHUB_ENV + echo "TENSORZERO_INTERNAL_FLAG_ENABLE_POSTGRES_WRITE=1" >> $GITHUB_ENV + + - name: Launch dependency services + run: docker compose -f tensorzero-core/tests/e2e/docker-compose.yml up --wait + + - name: Build the gateway for E2E tests + run: cargo build-e2e + + - name: Launch the gateway for E2E tests + run: | + cargo run-e2e > e2e_logs.txt 2>&1 & + count=0 + max_attempts=20 + while ! curl http://localhost:3000/health; do + echo "Waiting for gateway to be healthy..." + sleep 1 + count=$((count + 1)) + if [ $count -ge $max_attempts ]; then + echo "Gateway failed to become healthy after $max_attempts attempts" + exit 1 + fi + done + + - name: Run endpoint tests + run: cargo test-endpoints + + - name: Print docker compose logs + if: always() + run: docker compose -f tensorzero-core/tests/e2e/docker-compose.yml logs -t + + - name: Print e2e logs + if: always() + run: cat e2e_logs.txt + # Run 'cargo test-optimization' against mock-provider-api mock-optimization-tests: permissions: @@ -1304,6 +1421,7 @@ jobs: lint-rust, clickhouse-tests, postgres-tests, + endpoint-tests, ui-tests, ui-tests-e2e, ui-tests-e2e-regen-model-inference-cache, diff --git a/tensorzero-core/tests/e2e/common.rs b/tensorzero-core/tests/e2e/common.rs index 13235aa10de..4518fec9ac7 100644 --- a/tensorzero-core/tests/e2e/common.rs +++ b/tensorzero-core/tests/e2e/common.rs @@ -4,6 +4,7 @@ use reqwest::Url; use tensorzero_core::{ db::clickhouse::ClickHouseConnectionInfo, endpoints::datasets::{CLICKHOUSE_DATETIME_FORMAT, DatapointKind}, + feature_flags, }; use uuid::Uuid; @@ -12,6 +13,30 @@ lazy_static::lazy_static! { .unwrap_or_else(|_| "http://localhost:3000".to_string()); } +/// Returns true if we're testing against Postgres. +/// +/// This is not really perfect because we rely on the tests running in the same context as when we +/// launch the gateway container, but it's true for our CI setup and is good enough for today. +#[expect(clippy::expect_used)] +pub fn is_postgres_test() -> bool { + feature_flags::init_flags().expect("Failed to initialize feature flags"); + feature_flags::ENABLE_POSTGRES_READ.get() || feature_flags::ENABLE_POSTGRES_WRITE.get() +} + +/// Skips the current test if running against Postgres. +/// Use this for tests that don't have Postgres implementations yet. +/// +/// TODO(#5691): Remove this once we have Postgres implementations for all tests. +#[macro_export] +macro_rules! skip_for_postgres { + () => { + if $crate::common::is_postgres_test() { + eprintln!("Skipping test: Postgres implementation not yet available"); + return; + } + }; +} + pub fn get_gateway_endpoint(endpoint: &str) -> Url { let base_url: Url = GATEWAY_URL .parse() diff --git a/tensorzero-core/tests/e2e/config/tensorzero.functions.toml b/tensorzero-core/tests/e2e/config/tensorzero.functions.toml index 9c382842f3f..96a822ffd59 100644 --- a/tensorzero-core/tests/e2e/config/tensorzero.functions.toml +++ b/tensorzero-core/tests/e2e/config/tensorzero.functions.toml @@ -155,6 +155,21 @@ system_template = "../../../fixtures/config/functions/json_success/prompt/system max_tokens = 100 json_mode = "strict" +[functions.judge_answer] +type = "json" +system_schema = "../../../../ui/fixtures/config/functions/judge_answer/system_schema.json" +user_schema = "../../../../ui/fixtures/config/functions/judge_answer/user_schema.json" +output_schema = "../../../../ui/fixtures/config/functions/judge_answer/output_schema.json" + +[functions.judge_answer.variants."gpt-4.1-mini"] +type = "chat_completion" +model = "openai::gpt-4.1-mini" +system_template = "../../../../ui/fixtures/config/functions/judge_answer/gpt-4.1-mini/system_template.minijinja" +user_template = "../../../../ui/fixtures/config/functions/judge_answer/gpt-4.1-mini/user_template.minijinja" +retries = { num_retries = 4, max_delay_s = 10 } +json_mode = "strict" + + [functions.variant_failover] type = "chat" system_schema = "../../../fixtures/config/functions/basic_test/system_schema.json" diff --git a/tensorzero-core/tests/e2e/endpoints/datasets/clone_datapoints.rs b/tensorzero-core/tests/e2e/endpoints/datasets/clone_datapoints.rs index 8ff30b3c86a..e77e2bc3e3c 100644 --- a/tensorzero-core/tests/e2e/endpoints/datasets/clone_datapoints.rs +++ b/tensorzero-core/tests/e2e/endpoints/datasets/clone_datapoints.rs @@ -44,6 +44,7 @@ async fn get_test_setup() -> &'static (ClickHouseConnectionInfo, Arc) { #[tokio::test(flavor = "multi_thread")] async fn test_clone_datapoint_preserves_source_inference_id() { + skip_for_postgres!(); let client = Client::new(); let (clickhouse, config) = get_test_setup().await; @@ -146,6 +147,7 @@ async fn test_clone_datapoint_preserves_source_inference_id() { #[tokio::test(flavor = "multi_thread")] async fn test_clone_chat_datapoint_success() { + skip_for_postgres!(); let client = Client::new(); let (clickhouse, _config) = get_test_setup().await; @@ -226,6 +228,7 @@ async fn test_clone_chat_datapoint_success() { #[tokio::test(flavor = "multi_thread")] async fn test_clone_to_same_dataset() { + skip_for_postgres!(); let client = Client::new(); let (clickhouse, _config) = get_test_setup().await; @@ -307,6 +310,7 @@ async fn test_clone_to_same_dataset() { #[tokio::test(flavor = "multi_thread")] async fn test_clone_multiple_datapoints() { + skip_for_postgres!(); let client = Client::new(); let (clickhouse, _config) = get_test_setup().await; @@ -405,6 +409,7 @@ async fn test_clone_multiple_datapoints() { #[tokio::test(flavor = "multi_thread")] async fn test_clone_nonexistent_datapoint_returns_null() { + skip_for_postgres!(); let client = Client::new(); // Try to clone a nonexistent datapoint diff --git a/tensorzero-core/tests/e2e/endpoints/datasets/create_datapoints.rs b/tensorzero-core/tests/e2e/endpoints/datasets/create_datapoints.rs index 058cf84f444..cc101cbbece 100644 --- a/tensorzero-core/tests/e2e/endpoints/datasets/create_datapoints.rs +++ b/tensorzero-core/tests/e2e/endpoints/datasets/create_datapoints.rs @@ -34,6 +34,7 @@ async fn get_test_setup() -> &'static (ClickHouseConnectionInfo, Arc) { #[tokio::test(flavor = "multi_thread")] async fn test_create_chat_datapoint_basic() { + skip_for_postgres!(); let client = Client::new(); let (clickhouse, _config) = get_test_setup().await; @@ -146,6 +147,7 @@ async fn test_create_chat_datapoint_basic() { #[tokio::test(flavor = "multi_thread")] async fn test_create_json_datapoint_basic() { + skip_for_postgres!(); let client = Client::new(); let (clickhouse, _config) = get_test_setup().await; @@ -249,6 +251,7 @@ async fn test_create_json_datapoint_basic() { #[tokio::test(flavor = "multi_thread")] async fn test_create_multiple_datapoints() { + skip_for_postgres!(); let client = Client::new(); let (_clickhouse, _config) = get_test_setup().await; @@ -322,6 +325,7 @@ async fn test_create_multiple_datapoints() { #[tokio::test(flavor = "multi_thread")] async fn test_create_chat_datapoint_with_tools() { + skip_for_postgres!(); let client = Client::new(); let (_clickhouse, _config) = get_test_setup().await; @@ -368,6 +372,7 @@ async fn test_create_chat_datapoint_with_tools() { #[tokio::test(flavor = "multi_thread")] async fn test_create_datapoint_with_tags() { + skip_for_postgres!(); let client = Client::new(); let (_clickhouse, _config) = get_test_setup().await; @@ -407,6 +412,7 @@ async fn test_create_datapoint_with_tags() { #[tokio::test(flavor = "multi_thread")] async fn test_create_datapoint_invalid_function() { + skip_for_postgres!(); let client = Client::new(); let (_clickhouse, _config) = get_test_setup().await; @@ -435,6 +441,7 @@ async fn test_create_datapoint_invalid_function() { #[tokio::test(flavor = "multi_thread")] async fn test_create_datapoint_wrong_function_type() { + skip_for_postgres!(); let client = Client::new(); let (_clickhouse, _config) = get_test_setup().await; @@ -470,6 +477,7 @@ async fn test_create_datapoint_wrong_function_type() { #[tokio::test(flavor = "multi_thread")] async fn test_create_datapoint_empty_list() { + skip_for_postgres!(); let client = Client::new(); let (_clickhouse, _config) = get_test_setup().await; @@ -489,6 +497,7 @@ async fn test_create_datapoint_empty_list() { #[tokio::test(flavor = "multi_thread")] async fn test_create_json_datapoint_invalid_schema() { + skip_for_postgres!(); let client = Client::new(); let (clickhouse, _config) = get_test_setup().await; @@ -567,6 +576,7 @@ async fn test_create_json_datapoint_invalid_schema() { #[tokio::test(flavor = "multi_thread")] async fn test_create_datapoint_with_episode_id() { + skip_for_postgres!(); let client = Client::new(); let (_clickhouse, _config) = get_test_setup().await; @@ -609,6 +619,7 @@ async fn test_create_datapoint_with_episode_id() { #[tokio::test(flavor = "multi_thread")] async fn test_create_datapoint_without_output() { + skip_for_postgres!(); let client = Client::new(); let (_clickhouse, _config) = get_test_setup().await; @@ -641,6 +652,7 @@ async fn test_create_datapoint_without_output() { #[tokio::test(flavor = "multi_thread")] async fn test_create_json_datapoint_default_schema() { + skip_for_postgres!(); let client = Client::new(); let (_clickhouse, _config) = get_test_setup().await; diff --git a/tensorzero-core/tests/e2e/endpoints/datasets/create_from_inferences.rs b/tensorzero-core/tests/e2e/endpoints/datasets/create_from_inferences.rs index a83049c4022..961fc8a2c83 100644 --- a/tensorzero-core/tests/e2e/endpoints/datasets/create_from_inferences.rs +++ b/tensorzero-core/tests/e2e/endpoints/datasets/create_from_inferences.rs @@ -38,6 +38,7 @@ async fn get_test_setup() -> &'static (ClickHouseConnectionInfo, Arc) { #[tokio::test(flavor = "multi_thread")] async fn test_create_from_inference_ids_success() { + skip_for_postgres!(); let client = Client::new(); let (clickhouse, config) = get_test_setup().await; @@ -78,6 +79,7 @@ async fn test_create_from_inference_ids_success() { #[tokio::test] async fn test_create_from_inference_query_success() { + skip_for_postgres!(); let client = Client::new(); // Create datapoints using a query (no filters, just function name) @@ -110,6 +112,7 @@ async fn test_create_from_inference_query_success() { #[tokio::test(flavor = "multi_thread")] async fn test_create_from_same_inference_multiple_times_succeeds() { + skip_for_postgres!(); let client = Client::new(); let (clickhouse, config) = get_test_setup().await; @@ -163,6 +166,7 @@ async fn test_create_from_same_inference_multiple_times_succeeds() { #[tokio::test(flavor = "multi_thread")] async fn test_create_from_inference_missing_ids_error() { + skip_for_postgres!(); let client = Client::new(); let (clickhouse, config) = get_test_setup().await; @@ -209,6 +213,7 @@ async fn test_create_from_inference_missing_ids_error() { #[tokio::test] async fn test_create_from_inference_with_filters() { + skip_for_postgres!(); let client = Client::new(); // Create datapoints using a tag filter that exists in the test data diff --git a/tensorzero-core/tests/e2e/endpoints/datasets/delete_datapoints.rs b/tensorzero-core/tests/e2e/endpoints/datasets/delete_datapoints.rs index 1f53a373aa2..ea513185bf5 100644 --- a/tensorzero-core/tests/e2e/endpoints/datasets/delete_datapoints.rs +++ b/tensorzero-core/tests/e2e/endpoints/datasets/delete_datapoints.rs @@ -20,6 +20,7 @@ use crate::common::get_gateway_endpoint; #[tokio::test] async fn test_delete_datapoints_single_chat() { + skip_for_postgres!(); let http_client = Client::new(); let clickhouse = get_clickhouse().await; let dataset_name = format!("test-delete-dp-single-chat-{}", Uuid::now_v7()); @@ -142,6 +143,7 @@ async fn test_delete_datapoints_single_chat() { #[tokio::test] async fn test_delete_datapoints_multiple_mixed() { + skip_for_postgres!(); let http_client = Client::new(); let clickhouse = get_clickhouse().await; let dataset_name = format!("test-delete-dp-multiple-{}", Uuid::now_v7()); @@ -292,6 +294,7 @@ async fn test_delete_datapoints_multiple_mixed() { #[tokio::test] async fn test_delete_datapoints_empty_ids_list() { + skip_for_postgres!(); let http_client = Client::new(); let dataset_name = format!("test-delete-dp-empty-{}", Uuid::now_v7()); @@ -310,6 +313,7 @@ async fn test_delete_datapoints_empty_ids_list() { #[tokio::test] async fn test_delete_datapoints_non_existent_id() { + skip_for_postgres!(); let http_client = Client::new(); let clickhouse = get_clickhouse().await; let dataset_name = format!("test-delete-dp-non-existent-{}", Uuid::now_v7()); @@ -392,6 +396,7 @@ async fn test_delete_datapoints_non_existent_id() { #[tokio::test] async fn test_delete_datapoints_invalid_dataset_name() { + skip_for_postgres!(); let http_client = Client::new(); let datapoint_id = Uuid::now_v7(); @@ -413,6 +418,7 @@ async fn test_delete_datapoints_invalid_dataset_name() { #[tokio::test] async fn test_delete_datapoints_from_empty_dataset() { + skip_for_postgres!(); let http_client = Client::new(); let dataset_name = format!("test-delete-dp-empty-dataset-{}", Uuid::now_v7()); let non_existent_id = Uuid::now_v7(); @@ -437,6 +443,7 @@ async fn test_delete_datapoints_from_empty_dataset() { #[tokio::test] async fn test_delete_datapoints_already_stale() { + skip_for_postgres!(); let http_client = Client::new(); let clickhouse = get_clickhouse().await; let dataset_name = format!("test-delete-dp-already-stale-{}", Uuid::now_v7()); diff --git a/tensorzero-core/tests/e2e/endpoints/datasets/delete_dataset.rs b/tensorzero-core/tests/e2e/endpoints/datasets/delete_dataset.rs index 2b1de671f1a..b48f012bbda 100644 --- a/tensorzero-core/tests/e2e/endpoints/datasets/delete_dataset.rs +++ b/tensorzero-core/tests/e2e/endpoints/datasets/delete_dataset.rs @@ -20,6 +20,7 @@ use crate::common::get_gateway_endpoint; #[tokio::test] async fn test_delete_dataset_with_single_datapoint() { + skip_for_postgres!(); let http_client = Client::new(); let clickhouse = get_clickhouse().await; let dataset_name = format!("test-delete-dataset-single-{}", Uuid::now_v7()); @@ -144,6 +145,7 @@ async fn test_delete_dataset_with_single_datapoint() { #[tokio::test] async fn test_delete_dataset_with_multiple_mixed_datapoints() { + skip_for_postgres!(); let http_client = Client::new(); let clickhouse = get_clickhouse().await; let dataset_name = format!("test-delete-dataset-multiple-{}", Uuid::now_v7()); @@ -297,6 +299,7 @@ async fn test_delete_dataset_with_multiple_mixed_datapoints() { #[tokio::test] async fn test_delete_empty_dataset() { + skip_for_postgres!(); let http_client = Client::new(); let dataset_name = format!("test-delete-dataset-empty-{}", Uuid::now_v7()); @@ -317,6 +320,7 @@ async fn test_delete_empty_dataset() { #[tokio::test] async fn test_delete_dataset_invalid_name() { + skip_for_postgres!(); let http_client = Client::new(); // Try to delete a dataset with invalid characters @@ -332,6 +336,7 @@ async fn test_delete_dataset_invalid_name() { #[tokio::test] async fn test_delete_dataset_twice() { + skip_for_postgres!(); let http_client = Client::new(); let clickhouse = get_clickhouse().await; let dataset_name = format!("test-delete-dataset-twice-{}", Uuid::now_v7()); @@ -405,6 +410,7 @@ async fn test_delete_dataset_twice() { #[tokio::test] async fn test_delete_dataset_with_different_function_names() { + skip_for_postgres!(); let http_client = Client::new(); let clickhouse = get_clickhouse().await; let dataset_name = format!("test-delete-dataset-functions-{}", Uuid::now_v7()); diff --git a/tensorzero-core/tests/e2e/endpoints/datasets/get_datapoints.rs b/tensorzero-core/tests/e2e/endpoints/datasets/get_datapoints.rs index c4a3d4df7e4..445872f3c2c 100644 --- a/tensorzero-core/tests/e2e/endpoints/datasets/get_datapoints.rs +++ b/tensorzero-core/tests/e2e/endpoints/datasets/get_datapoints.rs @@ -26,6 +26,7 @@ mod get_datapoints_tests { #[tokio::test] async fn test_get_datapoints_single_chat_datapoint_without_dataset_name() { + skip_for_postgres!(); let http_client = Client::new(); let clickhouse = get_clickhouse().await; let dataset_name = format!("test-get-dp-single-chat-{}", Uuid::now_v7()); @@ -111,6 +112,7 @@ mod get_datapoints_tests { #[tokio::test] async fn test_get_datapoints_single_chat_datapoint() { + skip_for_postgres!(); let http_client = Client::new(); let clickhouse = get_clickhouse().await; let dataset_name = format!("test-get-dp-single-chat-{}", Uuid::now_v7()); @@ -198,6 +200,7 @@ mod get_datapoints_tests { #[tokio::test] async fn test_get_datapoints_single_json_datapoint() { + skip_for_postgres!(); let http_client = Client::new(); let clickhouse = get_clickhouse().await; let dataset_name = format!("test-get-dp-single-json-{}", Uuid::now_v7()); @@ -281,6 +284,7 @@ mod get_datapoints_tests { #[tokio::test] async fn test_get_datapoints_multiple_mixed_datapoints() { + skip_for_postgres!(); let http_client = Client::new(); let clickhouse = get_clickhouse().await; let dataset_name = format!("test-get-dp-multiple-{}", Uuid::now_v7()); @@ -425,6 +429,7 @@ mod get_datapoints_tests { #[tokio::test] async fn test_get_datapoints_with_non_existent_ids() { + skip_for_postgres!(); let http_client = Client::new(); let clickhouse = get_clickhouse().await; let dataset_name = format!("test-get-dp-non-existent-{}", Uuid::now_v7()); @@ -496,6 +501,7 @@ mod get_datapoints_tests { #[tokio::test] async fn test_get_datapoints_returns_stale_datapoints() { + skip_for_postgres!(); let http_client = Client::new(); let clickhouse = get_clickhouse().await; let dataset_name = format!("test-get-dp-stale-{}", Uuid::now_v7()); @@ -572,6 +578,7 @@ mod get_datapoints_tests { #[tokio::test] async fn test_get_datapoints_empty_ids_list() { + skip_for_postgres!(); let http_client = Client::new(); let clickhouse = get_clickhouse().await; let dataset_name = format!("test-get-dp-empty-ids-list-{}", Uuid::now_v7()); @@ -634,6 +641,7 @@ mod get_datapoints_tests { #[tokio::test] async fn test_get_datapoints_invalid_uuid() { + skip_for_postgres!(); // Create a valid dataset name so we have a valid dataset name. let dataset_name = format!("test-get-dp-invalid-uuid-{}", Uuid::now_v7()); let http_client = Client::new(); @@ -701,6 +709,7 @@ mod list_datapoints_tests { #[tokio::test] async fn test_list_datapoints_basic_pagination() { + skip_for_postgres!(); let http_client = Client::new(); let clickhouse = get_clickhouse().await; let dataset_name = format!("test-list-dp-pagination-{}", Uuid::now_v7()); @@ -815,6 +824,7 @@ mod list_datapoints_tests { #[tokio::test] async fn test_list_datapoints_filter_by_function_name() { + skip_for_postgres!(); let http_client = Client::new(); let clickhouse = get_clickhouse().await; let dataset_name = format!("test-list-dp-function-{}", Uuid::now_v7()); @@ -928,6 +938,7 @@ mod list_datapoints_tests { #[tokio::test] async fn test_list_datapoints_filter_by_tags() { + skip_for_postgres!(); let http_client = Client::new(); let clickhouse = get_clickhouse().await; let dataset_name = format!("test-list-dp-tags-{}", Uuid::now_v7()); @@ -1060,6 +1071,7 @@ mod list_datapoints_tests { #[tokio::test] async fn test_list_datapoints_filter_by_time() { + skip_for_postgres!(); let http_client = Client::new(); let clickhouse = get_clickhouse().await; let dataset_name = format!("test-list-dp-time-{}", Uuid::now_v7()); @@ -1150,6 +1162,7 @@ mod list_datapoints_tests { #[tokio::test] async fn test_list_datapoints_complex_filters() { + skip_for_postgres!(); let http_client = Client::new(); let clickhouse = get_clickhouse().await; let dataset_name = format!("test-list-dp-complex-{}", Uuid::now_v7()); @@ -1341,6 +1354,7 @@ mod list_datapoints_tests { #[tokio::test] async fn test_list_datapoints_empty_dataset() { + skip_for_postgres!(); let http_client = Client::new(); let dataset_name = format!("test-list-dp-empty-{}", Uuid::now_v7()); @@ -1362,6 +1376,7 @@ mod list_datapoints_tests { #[tokio::test] async fn test_list_datapoints_does_not_return_stale() { + skip_for_postgres!(); let http_client = Client::new(); let clickhouse = get_clickhouse().await; let dataset_name = format!("test-list-dp-no-stale-{}", Uuid::now_v7()); @@ -1444,6 +1459,7 @@ mod list_datapoints_tests { #[tokio::test] async fn test_list_datapoints_mixed_chat_and_json() { + skip_for_postgres!(); let http_client = Client::new(); let clickhouse = get_clickhouse().await; let dataset_name = format!("test-list-dp-mixed-{}", Uuid::now_v7()); @@ -1551,6 +1567,7 @@ mod list_datapoints_tests { #[tokio::test] async fn test_list_datapoints_with_large_limit() { + skip_for_postgres!(); let http_client = Client::new(); let clickhouse = get_clickhouse().await; let dataset_name = format!("test-list-dp-large-page-{}", Uuid::now_v7()); diff --git a/tensorzero-core/tests/e2e/endpoints/datasets/list_datasets.rs b/tensorzero-core/tests/e2e/endpoints/datasets/list_datasets.rs index c635cd4473d..5eee8953df1 100644 --- a/tensorzero-core/tests/e2e/endpoints/datasets/list_datasets.rs +++ b/tensorzero-core/tests/e2e/endpoints/datasets/list_datasets.rs @@ -16,6 +16,7 @@ use crate::common::get_gateway_endpoint; #[tokio::test] async fn test_list_datasets_no_params() { + skip_for_postgres!(); let http_client = Client::new(); let clickhouse = get_clickhouse().await; @@ -106,6 +107,7 @@ async fn test_list_datasets_no_params() { #[tokio::test] async fn test_list_datasets_with_function_filter() { + skip_for_postgres!(); let http_client = Client::new(); let clickhouse = get_clickhouse().await; @@ -238,6 +240,7 @@ async fn test_list_datasets_with_function_filter() { #[tokio::test] async fn test_list_datasets_with_pagination() { + skip_for_postgres!(); let http_client = Client::new(); let clickhouse = get_clickhouse().await; @@ -358,6 +361,7 @@ async fn test_list_datasets_with_pagination() { #[tokio::test] async fn test_list_datasets_empty_result() { + skip_for_postgres!(); let http_client = Client::new(); // Filter by a function that doesn't exist diff --git a/tensorzero-core/tests/e2e/endpoints/datasets/update_datapoints.rs b/tensorzero-core/tests/e2e/endpoints/datasets/update_datapoints.rs index e898202ad0d..d199ff1e9fc 100644 --- a/tensorzero-core/tests/e2e/endpoints/datasets/update_datapoints.rs +++ b/tensorzero-core/tests/e2e/endpoints/datasets/update_datapoints.rs @@ -15,6 +15,9 @@ use tensorzero_core::db::datasets::DatasetQueries; use tensorzero_core::db::stored_datapoint::{ StoredChatInferenceDatapoint, StoredDatapoint, StoredJsonInferenceDatapoint, }; +use tensorzero_core::endpoints::datasets::v1::types::{ + DatapointMetadataUpdate, UpdateDatapointMetadataRequest, UpdateDatapointsMetadataRequest, +}; use tensorzero_core::inference::types::{ Arguments, ContentBlockChatOutput, JsonInferenceOutput, Role, StoredInput, StoredInputMessage, StoredInputMessageContent, System, Template, Text, @@ -28,6 +31,7 @@ use crate::common::get_gateway_endpoint; #[tokio::test] async fn test_update_chat_datapoint_output() { + skip_for_postgres!(); let http_client = Client::new(); let clickhouse = get_clickhouse().await; let dataset_name = format!("test-update-chat-{}", Uuid::now_v7()); @@ -161,6 +165,7 @@ async fn test_update_chat_datapoint_output() { #[tokio::test] async fn test_update_json_datapoint_output() { + skip_for_postgres!(); let http_client = Client::new(); let clickhouse = get_clickhouse().await; let dataset_name = format!("test-update-json-{}", Uuid::now_v7()); @@ -267,6 +272,7 @@ async fn test_update_json_datapoint_output() { #[tokio::test] async fn test_update_multiple_datapoints() { + skip_for_postgres!(); let http_client = Client::new(); let clickhouse = get_clickhouse().await; let dataset_name = format!("test-update-batch-{}", Uuid::now_v7()); @@ -421,6 +427,7 @@ async fn test_update_multiple_datapoints() { #[tokio::test] async fn test_update_datapoint_not_found() { + skip_for_postgres!(); let http_client = Client::new(); let dataset_name = format!("test-update-not-found-{}", Uuid::now_v7()); let nonexistent_id = Uuid::now_v7(); @@ -445,6 +452,7 @@ async fn test_update_datapoint_not_found() { #[tokio::test] async fn test_update_datapoint_type_mismatch() { + skip_for_postgres!(); let http_client = Client::new(); let clickhouse = get_clickhouse().await; let dataset_name = format!("test-update-type-mismatch-{}", Uuid::now_v7()); @@ -517,6 +525,7 @@ async fn test_update_datapoint_type_mismatch() { #[tokio::test] async fn test_update_datapoint_with_metadata() { + skip_for_postgres!(); let http_client = Client::new(); let clickhouse = get_clickhouse().await; let dataset_name = format!("test-update-metadata-{}", Uuid::now_v7()); @@ -606,6 +615,7 @@ async fn test_update_datapoint_with_metadata() { #[tokio::test] async fn test_update_datapoint_empty_request() { + skip_for_postgres!(); let http_client = Client::new(); let dataset_name = format!("test-update-empty-{}", Uuid::now_v7()); @@ -625,6 +635,7 @@ async fn test_update_datapoint_empty_request() { #[tokio::test] async fn test_update_datapoint_duplicate_ids() { + skip_for_postgres!(); let http_client = Client::new(); let dataset_name = format!("test-update-duplicate-{}", Uuid::now_v7()); let datapoint_id = Uuid::now_v7(); @@ -656,6 +667,7 @@ async fn test_update_datapoint_duplicate_ids() { #[tokio::test] async fn test_update_chat_datapoint_set_output_to_null() { + skip_for_postgres!(); let http_client = Client::new(); let clickhouse = get_clickhouse().await; let dataset_name = format!("test-update-chat-output-null-{}", Uuid::now_v7()); @@ -747,6 +759,7 @@ async fn test_update_chat_datapoint_set_output_to_null() { #[tokio::test] async fn test_update_chat_datapoint_set_tool_params_to_null() { + skip_for_postgres!(); let http_client = Client::new(); let clickhouse = get_clickhouse().await; let dataset_name = format!("test-update-chat-tool-params-null-{}", Uuid::now_v7()); @@ -860,6 +873,7 @@ async fn test_update_chat_datapoint_set_tool_params_to_null() { #[tokio::test] async fn test_update_chat_datapoint_set_tags_to_empty() { + skip_for_postgres!(); let http_client = Client::new(); let clickhouse = get_clickhouse().await; let dataset_name = format!("test-update-chat-tags-null-{}", Uuid::now_v7()); @@ -954,6 +968,7 @@ async fn test_update_chat_datapoint_set_tags_to_empty() { #[tokio::test] async fn test_update_chat_datapoint_set_name_to_null() { + skip_for_postgres!(); let http_client = Client::new(); let clickhouse = get_clickhouse().await; let dataset_name = format!("test-update-chat-name-null-{}", Uuid::now_v7()); @@ -1045,6 +1060,7 @@ async fn test_update_chat_datapoint_set_name_to_null() { #[tokio::test] async fn test_update_json_datapoint_set_output_to_null() { + skip_for_postgres!(); let http_client = Client::new(); let clickhouse = get_clickhouse().await; let dataset_name = format!("test-update-json-output-null-{}", Uuid::now_v7()); @@ -1143,6 +1159,7 @@ async fn test_update_json_datapoint_set_output_to_null() { #[tokio::test] async fn test_update_json_datapoint_set_tags_to_empty() { + skip_for_postgres!(); let http_client = Client::new(); let clickhouse = get_clickhouse().await; let dataset_name = format!("test-update-json-tags-null-{}", Uuid::now_v7()); @@ -1249,6 +1266,7 @@ async fn test_update_json_datapoint_set_tags_to_empty() { #[tokio::test] async fn test_update_metadata_chat_datapoint() { + skip_for_postgres!(); let http_client = Client::new(); let clickhouse = get_clickhouse().await; let dataset_name = format!("test-update-metadata-chat-{}", Uuid::now_v7()); @@ -1344,6 +1362,7 @@ async fn test_update_metadata_chat_datapoint() { #[tokio::test] async fn test_update_metadata_json_datapoint() { + skip_for_postgres!(); let http_client = Client::new(); let clickhouse = get_clickhouse().await; let dataset_name = format!("test-update-metadata-json-{}", Uuid::now_v7()); @@ -1442,6 +1461,7 @@ async fn test_update_metadata_json_datapoint() { #[tokio::test] async fn test_update_metadata_set_name_to_null() { + skip_for_postgres!(); let http_client = Client::new(); let clickhouse = get_clickhouse().await; let dataset_name = format!("test-update-metadata-null-{}", Uuid::now_v7()); @@ -1522,7 +1542,8 @@ async fn test_update_metadata_set_name_to_null() { } #[tokio::test] -async fn test_update_metadata_batch() { +async fn test_update_metadata_multiple_datapoints() { + skip_for_postgres!(); let http_client = Client::new(); let clickhouse = get_clickhouse().await; let dataset_name = format!("test-update-metadata-batch-{}", Uuid::now_v7()); @@ -1597,22 +1618,27 @@ async fn test_update_metadata_batch() { tokio::time::sleep(Duration::from_millis(500)).await; // Update both datapoints' metadata + let update_request = UpdateDatapointsMetadataRequest { + datapoints: vec![ + UpdateDatapointMetadataRequest { + id: datapoint_id1, + metadata: DatapointMetadataUpdate { + name: Some(Some("updated_name1".to_string())), + }, + }, + UpdateDatapointMetadataRequest { + id: datapoint_id2, + metadata: DatapointMetadataUpdate { + name: Some(Some("updated_name2".to_string())), + }, + }, + ], + }; let resp = http_client .patch(get_gateway_endpoint(&format!( "/v1/datasets/{dataset_name}/datapoints/metadata", ))) - .json(&json!({ - "datapoints": [ - { - "id": datapoint_id1.to_string(), - "metadata": {"name": "updated_name1"} - }, - { - "id": datapoint_id2.to_string(), - "metadata": {"name": "updated_name2"} - } - ] - })) + .json(&update_request) .send() .await .unwrap(); @@ -1658,6 +1684,7 @@ async fn test_update_metadata_batch() { #[tokio::test] async fn test_update_metadata_datapoint_not_found() { + skip_for_postgres!(); let http_client = Client::new(); let dataset_name = format!("test-update-metadata-notfound-{}", Uuid::now_v7()); let non_existent_id = Uuid::now_v7(); @@ -1681,6 +1708,7 @@ async fn test_update_metadata_datapoint_not_found() { #[tokio::test] async fn test_update_metadata_duplicate_ids() { + skip_for_postgres!(); let http_client = Client::new(); let dataset_name = format!("test-update-metadata-duplicate-{}", Uuid::now_v7()); let duplicate_id = Uuid::now_v7(); @@ -1710,6 +1738,7 @@ async fn test_update_metadata_duplicate_ids() { #[tokio::test] async fn test_get_chat_datapoint_modify_and_update_roundtrip() { + skip_for_postgres!(); let http_client = Client::new(); let clickhouse = get_clickhouse().await; let dataset_name = format!("test-roundtrip-{}", Uuid::now_v7()); @@ -1853,6 +1882,7 @@ async fn test_get_chat_datapoint_modify_and_update_roundtrip() { #[tokio::test] async fn test_get_json_datapoint_modify_and_update_roundtrip() { + skip_for_postgres!(); let http_client = Client::new(); let clickhouse = get_clickhouse().await; let dataset_name = format!("test-json-roundtrip-{}", Uuid::now_v7()); diff --git a/tensorzero-core/tests/e2e/endpoints/internal/action.rs b/tensorzero-core/tests/e2e/endpoints/internal/action.rs index 559a852afa9..4fac09512ba 100644 --- a/tensorzero-core/tests/e2e/endpoints/internal/action.rs +++ b/tensorzero-core/tests/e2e/endpoints/internal/action.rs @@ -69,6 +69,7 @@ async fn call_action( /// Test that the action endpoint can execute inference using a historical config /// that has a function not present in the running gateway config. async fn test_action_with_historical_config_impl(client: Client) { + skip_for_postgres!(); let clickhouse = get_clickhouse().await; let id = Uuid::now_v7(); @@ -152,6 +153,7 @@ async fn test_action_with_historical_config_impl_http_gateway() { /// Test that the action endpoint returns an error for a non-existent snapshot hash. async fn test_action_nonexistent_snapshot_hash_impl(client: Client) { + skip_for_postgres!(); // Use a properly formatted hash that simply doesn't exist in the database let nonexistent_hash = SnapshotHash::new_test(); @@ -200,6 +202,7 @@ async fn test_action_nonexistent_snapshot_hash_impl_http_gateway() { /// Test that the action endpoint rejects streaming requests. async fn test_action_streaming_rejected_impl(client: Client) { + skip_for_postgres!(); let clickhouse = get_clickhouse().await; let id = Uuid::now_v7(); diff --git a/tensorzero-core/tests/e2e/endpoints/internal/config.rs b/tensorzero-core/tests/e2e/endpoints/internal/config.rs index 75d95193853..385c09309c6 100644 --- a/tensorzero-core/tests/e2e/endpoints/internal/config.rs +++ b/tensorzero-core/tests/e2e/endpoints/internal/config.rs @@ -12,6 +12,7 @@ use crate::common::get_gateway_endpoint; #[tokio::test(flavor = "multi_thread")] async fn test_get_live_config() { + skip_for_postgres!(); let http_client = Client::new(); let url = get_gateway_endpoint("/internal/config"); @@ -47,6 +48,7 @@ async fn test_get_live_config() { #[tokio::test(flavor = "multi_thread")] async fn test_get_config_by_hash() { + skip_for_postgres!(); let http_client = Client::new(); // First get the live config to obtain the current hash @@ -83,6 +85,7 @@ async fn test_get_config_by_hash() { #[tokio::test(flavor = "multi_thread")] async fn test_get_config_by_nonexistent_hash() { + skip_for_postgres!(); let http_client = Client::new(); // Use a hash that definitely doesn't exist @@ -104,6 +107,7 @@ async fn test_get_config_by_nonexistent_hash() { /// Test writing a config via the Rust client async fn test_write_config_impl(client: TensorZeroClient) { + skip_for_postgres!(); let id = Uuid::now_v7(); // Create a minimal config with a unique metric @@ -170,6 +174,7 @@ tensorzero::make_gateway_test_functions!(test_write_config_impl); /// Test tag merging when writing the same config twice via the Rust client async fn test_write_config_tag_merging_impl(client: TensorZeroClient) { + skip_for_postgres!(); let id = Uuid::now_v7(); // Create a config diff --git a/tensorzero-core/tests/e2e/endpoints/internal/datapoint_count.rs b/tensorzero-core/tests/e2e/endpoints/internal/datapoint_count.rs index c35fa64d9e5..fda54fcb1b3 100644 --- a/tensorzero-core/tests/e2e/endpoints/internal/datapoint_count.rs +++ b/tensorzero-core/tests/e2e/endpoints/internal/datapoint_count.rs @@ -100,6 +100,7 @@ async fn fetch_datapoint_count( #[tokio::test(flavor = "multi_thread")] async fn test_get_datapoint_count_counts_new_datapoints() { + skip_for_postgres!(); let client = Client::new(); let dataset_name = format!("datapoints-count-{}", Uuid::now_v7()); @@ -136,6 +137,7 @@ async fn test_get_datapoint_count_counts_new_datapoints() { #[tokio::test(flavor = "multi_thread")] async fn test_get_datapoint_count_filters_by_function_name() { + skip_for_postgres!(); let client = Client::new(); let dataset_name = format!("datapoints-count-filter-{}", Uuid::now_v7()); diff --git a/tensorzero-core/tests/e2e/endpoints/internal/episodes.rs b/tensorzero-core/tests/e2e/endpoints/internal/episodes.rs index 397d1d58feb..85b3bd30fe9 100644 --- a/tensorzero-core/tests/e2e/endpoints/internal/episodes.rs +++ b/tensorzero-core/tests/e2e/endpoints/internal/episodes.rs @@ -10,6 +10,7 @@ use crate::common::get_gateway_endpoint; #[tokio::test(flavor = "multi_thread")] async fn test_query_episode_table_bounds() { + skip_for_postgres!(); let http_client = Client::new(); let url = get_gateway_endpoint("/internal/episodes/bounds"); @@ -33,6 +34,7 @@ async fn test_query_episode_table_bounds() { #[tokio::test(flavor = "multi_thread")] async fn test_query_episode_table() { + skip_for_postgres!(); let http_client = Client::new(); let url = get_gateway_endpoint("/internal/episodes?limit=10"); @@ -67,6 +69,7 @@ async fn test_query_episode_table() { #[tokio::test(flavor = "multi_thread")] async fn test_query_episode_table_with_pagination() { + skip_for_postgres!(); let http_client = Client::new(); // First, get the bounds to know what episodes exist @@ -110,6 +113,7 @@ async fn test_query_episode_table_with_pagination() { #[tokio::test(flavor = "multi_thread")] async fn test_query_episode_table_limit_zero() { + skip_for_postgres!(); let http_client = Client::new(); let url = get_gateway_endpoint("/internal/episodes?limit=0"); @@ -129,6 +133,7 @@ async fn test_query_episode_table_limit_zero() { #[tokio::test(flavor = "multi_thread")] async fn test_query_episode_table_rejects_limit_over_100() { + skip_for_postgres!(); let http_client = Client::new(); let url = get_gateway_endpoint("/internal/episodes?limit=101"); @@ -142,6 +147,7 @@ async fn test_query_episode_table_rejects_limit_over_100() { #[tokio::test(flavor = "multi_thread")] async fn test_get_episode_inference_count() { + skip_for_postgres!(); let http_client = Client::new(); // First get an episode that exists @@ -169,6 +175,7 @@ async fn test_get_episode_inference_count() { #[tokio::test(flavor = "multi_thread")] async fn test_get_episode_inference_count_nonexistent_episode() { + skip_for_postgres!(); let http_client = Client::new(); // Use a UUID that likely doesn't exist diff --git a/tensorzero-core/tests/e2e/endpoints/internal/evaluations.rs b/tensorzero-core/tests/e2e/endpoints/internal/evaluations.rs index 2e763fee021..66157ece23a 100644 --- a/tensorzero-core/tests/e2e/endpoints/internal/evaluations.rs +++ b/tensorzero-core/tests/e2e/endpoints/internal/evaluations.rs @@ -19,6 +19,7 @@ use crate::common::get_gateway_endpoint; #[tokio::test(flavor = "multi_thread")] async fn test_get_evaluation_run_infos_endpoint() { + skip_for_postgres!(); let http_client = Client::new(); // Use evaluation run IDs from the test fixture data @@ -64,6 +65,7 @@ async fn test_get_evaluation_run_infos_endpoint() { #[tokio::test(flavor = "multi_thread")] async fn test_get_evaluation_run_infos_single_run() { + skip_for_postgres!(); let http_client = Client::new(); let evaluation_run_id = "0196368f-19bd-7082-a677-1c0bf346ff24"; @@ -97,6 +99,7 @@ async fn test_get_evaluation_run_infos_single_run() { #[tokio::test(flavor = "multi_thread")] async fn test_get_evaluation_run_infos_nonexistent_run() { + skip_for_postgres!(); let http_client = Client::new(); let url = get_gateway_endpoint("/internal/evaluations/run_infos").to_string() @@ -120,6 +123,7 @@ async fn test_get_evaluation_run_infos_nonexistent_run() { #[tokio::test(flavor = "multi_thread")] async fn test_get_evaluation_run_infos_wrong_function() { + skip_for_postgres!(); let http_client = Client::new(); // Use a valid evaluation run ID but with wrong function name @@ -150,6 +154,7 @@ async fn test_get_evaluation_run_infos_wrong_function() { #[tokio::test(flavor = "multi_thread")] async fn test_get_evaluation_run_infos_for_datapoint_json_function() { + skip_for_postgres!(); let http_client = Client::new(); // Use datapoint ID from the test fixture data for extract_entities function @@ -187,6 +192,7 @@ async fn test_get_evaluation_run_infos_for_datapoint_json_function() { #[tokio::test(flavor = "multi_thread")] async fn test_get_evaluation_run_infos_for_datapoint_chat_function() { + skip_for_postgres!(); let http_client = Client::new(); // Use datapoint ID from the test fixture data for write_haiku function @@ -223,6 +229,7 @@ async fn test_get_evaluation_run_infos_for_datapoint_chat_function() { #[tokio::test(flavor = "multi_thread")] async fn test_get_evaluation_run_infos_for_datapoint_nonexistent() { + skip_for_postgres!(); let http_client = Client::new(); let datapoint_id = "00000000-0000-0000-0000-000000000000"; @@ -251,6 +258,7 @@ async fn test_get_evaluation_run_infos_for_datapoint_nonexistent() { #[tokio::test(flavor = "multi_thread")] async fn test_get_evaluation_run_infos_for_datapoint_wrong_function() { + skip_for_postgres!(); let http_client = Client::new(); // Use a valid datapoint ID but with wrong function name - this will return an error since @@ -278,6 +286,7 @@ async fn test_get_evaluation_run_infos_for_datapoint_wrong_function() { #[tokio::test(flavor = "multi_thread")] async fn test_get_evaluation_statistics_endpoint() { + skip_for_postgres!(); let http_client = Client::new(); let evaluation_run_id = "0196368f-19bd-7082-a677-1c0bf346ff24"; @@ -315,6 +324,7 @@ async fn test_get_evaluation_statistics_endpoint() { #[tokio::test(flavor = "multi_thread")] async fn test_get_evaluation_statistics_multiple_runs() { + skip_for_postgres!(); let http_client = Client::new(); let evaluation_run_id1 = "0196368f-19bd-7082-a677-1c0bf346ff24"; @@ -345,6 +355,7 @@ async fn test_get_evaluation_statistics_multiple_runs() { #[tokio::test(flavor = "multi_thread")] async fn test_get_evaluation_statistics_empty_run_ids() { + skip_for_postgres!(); let http_client = Client::new(); let url = get_gateway_endpoint("/internal/evaluations/statistics").to_string() @@ -368,6 +379,7 @@ async fn test_get_evaluation_statistics_empty_run_ids() { #[tokio::test(flavor = "multi_thread")] async fn test_get_evaluation_statistics_nonexistent_run() { + skip_for_postgres!(); let http_client = Client::new(); let url = get_gateway_endpoint("/internal/evaluations/statistics").to_string() @@ -391,6 +403,7 @@ async fn test_get_evaluation_statistics_nonexistent_run() { #[tokio::test(flavor = "multi_thread")] async fn test_get_evaluation_statistics_invalid_function_type() { + skip_for_postgres!(); let http_client = Client::new(); let url = get_gateway_endpoint("/internal/evaluations/statistics").to_string() @@ -406,6 +419,7 @@ async fn test_get_evaluation_statistics_invalid_function_type() { #[tokio::test(flavor = "multi_thread")] async fn test_get_evaluation_statistics_invalid_uuid() { + skip_for_postgres!(); let http_client = Client::new(); let url = get_gateway_endpoint("/internal/evaluations/statistics").to_string() @@ -423,6 +437,7 @@ async fn test_get_evaluation_statistics_invalid_uuid() { #[tokio::test(flavor = "multi_thread")] async fn test_get_evaluation_results_haiku() { + skip_for_postgres!(); let http_client = Client::new(); // Use evaluation run ID from the test fixture data for haiku evaluation @@ -462,6 +477,7 @@ async fn test_get_evaluation_results_haiku() { #[tokio::test(flavor = "multi_thread")] async fn test_get_evaluation_results_entity_extraction() { + skip_for_postgres!(); let http_client = Client::new(); // Use evaluation run ID from the test fixture data for entity_extraction (JSON function) @@ -498,6 +514,7 @@ async fn test_get_evaluation_results_entity_extraction() { #[tokio::test(flavor = "multi_thread")] async fn test_get_evaluation_results_multiple_runs() { + skip_for_postgres!(); let http_client = Client::new(); // Use two evaluation run IDs from the test fixture data @@ -540,6 +557,7 @@ async fn test_get_evaluation_results_multiple_runs() { #[tokio::test(flavor = "multi_thread")] async fn test_get_evaluation_results_pagination() { + skip_for_postgres!(); let http_client = Client::new(); let evaluation_run_id = "01963691-9d3c-7793-a8be-3937ebb849c1"; @@ -599,6 +617,7 @@ async fn test_get_evaluation_results_pagination() { #[tokio::test(flavor = "multi_thread")] async fn test_get_evaluation_results_evaluation_not_found() { + skip_for_postgres!(); let http_client = Client::new(); let evaluation_run_id = "01963691-9d3c-7793-a8be-3937ebb849c1"; @@ -618,6 +637,7 @@ async fn test_get_evaluation_results_evaluation_not_found() { #[tokio::test(flavor = "multi_thread")] async fn test_get_evaluation_results_invalid_uuid() { + skip_for_postgres!(); let http_client = Client::new(); let url = get_gateway_endpoint("/internal/evaluations/results").to_string() @@ -633,6 +653,7 @@ async fn test_get_evaluation_results_invalid_uuid() { #[tokio::test(flavor = "multi_thread")] async fn test_get_evaluation_results_nonexistent_run() { + skip_for_postgres!(); let http_client = Client::new(); let url = get_gateway_endpoint("/internal/evaluations/results").to_string() @@ -655,6 +676,7 @@ async fn test_get_evaluation_results_nonexistent_run() { #[tokio::test(flavor = "multi_thread")] async fn test_get_evaluation_results_default_pagination() { + skip_for_postgres!(); let http_client = Client::new(); let evaluation_run_id = "01963691-9d3c-7793-a8be-3937ebb849c1"; @@ -720,6 +742,7 @@ async fn create_test_chat_datapoint( #[tokio::test(flavor = "multi_thread")] async fn test_run_evaluation_streaming_success() { + skip_for_postgres!(); let http_client = Client::new(); let _clickhouse = get_clickhouse().await; @@ -841,6 +864,7 @@ async fn test_run_evaluation_streaming_success() { #[tokio::test(flavor = "multi_thread")] async fn test_run_evaluation_streaming_missing_variant() { + skip_for_postgres!(); let http_client = Client::new(); // Request without variant_name or internal_dynamic_variant_config @@ -880,6 +904,7 @@ async fn test_run_evaluation_streaming_missing_variant() { #[tokio::test(flavor = "multi_thread")] async fn test_run_evaluation_streaming_nonexistent_dataset() { + skip_for_postgres!(); let http_client = Client::new(); let payload = json!({ @@ -950,6 +975,7 @@ async fn test_run_evaluation_streaming_nonexistent_dataset() { #[tokio::test(flavor = "multi_thread")] async fn test_run_evaluation_streaming_with_specific_datapoint_ids() { + skip_for_postgres!(); let http_client = Client::new(); let _clickhouse = get_clickhouse().await; @@ -1034,6 +1060,7 @@ async fn test_run_evaluation_streaming_with_specific_datapoint_ids() { #[tokio::test(flavor = "multi_thread")] async fn test_run_evaluation_streaming_conflicting_variant_config() { + skip_for_postgres!(); let http_client = Client::new(); // Provide both variant_name AND internal_dynamic_variant_config (should fail) @@ -1078,6 +1105,7 @@ async fn test_run_evaluation_streaming_conflicting_variant_config() { #[tokio::test(flavor = "multi_thread")] async fn test_run_evaluation_streaming_invalid_inference_cache() { + skip_for_postgres!(); let http_client = Client::new(); let payload = json!({ @@ -1123,6 +1151,7 @@ async fn test_run_evaluation_streaming_invalid_inference_cache() { /// and then verifies the endpoint returns the correct feedback. #[tokio::test(flavor = "multi_thread")] async fn test_get_human_feedback_returns_feedback_when_exists() { + skip_for_postgres!(); let http_client = Client::new(); // First, run an inference to get an inference_id @@ -1227,6 +1256,7 @@ async fn test_get_human_feedback_returns_feedback_when_exists() { /// Test that get_human_feedback returns None when no feedback exists. #[tokio::test(flavor = "multi_thread")] async fn test_get_human_feedback_returns_none_when_not_exists() { + skip_for_postgres!(); let http_client = Client::new(); let nonexistent_datapoint_id = Uuid::now_v7(); @@ -1260,6 +1290,7 @@ async fn test_get_human_feedback_returns_none_when_not_exists() { /// Test that get_human_feedback works with boolean feedback values. #[tokio::test(flavor = "multi_thread")] async fn test_get_human_feedback_with_boolean_value() { + skip_for_postgres!(); let http_client = Client::new(); // First, run an inference to get an inference_id @@ -1348,6 +1379,7 @@ async fn test_get_human_feedback_with_boolean_value() { /// Test that get_human_feedback returns the correct feedback when output doesn't match. #[tokio::test(flavor = "multi_thread")] async fn test_get_human_feedback_output_mismatch() { + skip_for_postgres!(); let http_client = Client::new(); // First, run an inference to get an inference_id @@ -1444,6 +1476,7 @@ async fn test_get_human_feedback_output_mismatch() { /// Test that get_human_feedback handles invalid UUID in datapoint_id. #[tokio::test(flavor = "multi_thread")] async fn test_get_human_feedback_invalid_uuid() { + skip_for_postgres!(); let http_client = Client::new(); let resp = http_client @@ -1467,6 +1500,7 @@ async fn test_get_human_feedback_invalid_uuid() { /// Test that get_human_feedback requires all parameters. #[tokio::test(flavor = "multi_thread")] async fn test_get_human_feedback_missing_parameters() { + skip_for_postgres!(); let http_client = Client::new(); let datapoint_id = Uuid::now_v7(); diff --git a/tensorzero-core/tests/e2e/endpoints/internal/feedback.rs b/tensorzero-core/tests/e2e/endpoints/internal/feedback.rs index fd276d4eac9..3e8e5d90adf 100644 --- a/tensorzero-core/tests/e2e/endpoints/internal/feedback.rs +++ b/tensorzero-core/tests/e2e/endpoints/internal/feedback.rs @@ -72,6 +72,7 @@ async fn submit_inference_feedback( #[tokio::test(flavor = "multi_thread")] async fn test_get_latest_feedback_id_by_metric_with_data() { + skip_for_postgres!(); let http_client = Client::new(); // Create an inference @@ -113,6 +114,7 @@ async fn test_get_latest_feedback_id_by_metric_with_data() { #[tokio::test(flavor = "multi_thread")] async fn test_get_latest_feedback_id_by_metric_multiple_feedback_same_metric() { + skip_for_postgres!(); let http_client = Client::new(); // Create an inference @@ -147,6 +149,7 @@ async fn test_get_latest_feedback_id_by_metric_multiple_feedback_same_metric() { #[tokio::test(flavor = "multi_thread")] async fn test_get_latest_feedback_id_by_metric_nonexistent_target() { + skip_for_postgres!(); let http_client = Client::new(); // Use a UUID that likely doesn't exist @@ -174,6 +177,7 @@ async fn test_get_latest_feedback_id_by_metric_nonexistent_target() { #[tokio::test(flavor = "multi_thread")] async fn test_get_latest_feedback_id_by_metric_invalid_uuid() { + skip_for_postgres!(); let http_client = Client::new(); // Use an invalid UUID @@ -192,6 +196,7 @@ async fn test_get_latest_feedback_id_by_metric_invalid_uuid() { #[tokio::test(flavor = "multi_thread")] async fn test_get_feedback_bounds_by_target_id_with_feedback() { + skip_for_postgres!(); let http_client = Client::new(); let inference_id = create_inference(&http_client, "basic_test").await; @@ -232,6 +237,7 @@ async fn test_get_feedback_bounds_by_target_id_with_feedback() { #[tokio::test(flavor = "multi_thread")] async fn test_get_feedback_bounds_by_target_id_nonexistent_target() { + skip_for_postgres!(); let http_client = Client::new(); let nonexistent_id = Uuid::now_v7(); @@ -256,6 +262,7 @@ async fn test_get_feedback_bounds_by_target_id_nonexistent_target() { #[tokio::test(flavor = "multi_thread")] async fn test_get_feedback_bounds_by_target_id_invalid_uuid() { + skip_for_postgres!(); let http_client = Client::new(); let url = get_gateway_endpoint("/internal/feedback/not-a-valid-uuid/bounds"); @@ -272,6 +279,7 @@ async fn test_get_feedback_bounds_by_target_id_invalid_uuid() { #[tokio::test(flavor = "multi_thread")] async fn test_get_feedback_by_target_id_with_data() { + skip_for_postgres!(); let http_client = Client::new(); // Create an inference @@ -305,6 +313,7 @@ async fn test_get_feedback_by_target_id_with_data() { #[tokio::test(flavor = "multi_thread")] async fn test_get_feedback_by_target_id_with_pagination() { + skip_for_postgres!(); let http_client = Client::new(); // Create an inference @@ -336,6 +345,7 @@ async fn test_get_feedback_by_target_id_with_pagination() { #[tokio::test(flavor = "multi_thread")] async fn test_get_feedback_by_target_id_nonexistent_target() { + skip_for_postgres!(); let http_client = Client::new(); // Use a UUID that likely doesn't exist @@ -360,6 +370,7 @@ async fn test_get_feedback_by_target_id_nonexistent_target() { #[tokio::test(flavor = "multi_thread")] async fn test_get_feedback_by_target_id_invalid_uuid() { + skip_for_postgres!(); let http_client = Client::new(); // Use an invalid UUID @@ -376,6 +387,7 @@ async fn test_get_feedback_by_target_id_invalid_uuid() { #[tokio::test(flavor = "multi_thread")] async fn test_get_feedback_by_target_id_rejects_both_before_and_after() { + skip_for_postgres!(); let http_client = Client::new(); let target_id = Uuid::now_v7(); @@ -400,6 +412,7 @@ async fn test_get_feedback_by_target_id_rejects_both_before_and_after() { #[tokio::test(flavor = "multi_thread")] async fn test_count_feedback_by_target_id_with_feedback() { + skip_for_postgres!(); let http_client = Client::new(); let inference_id = create_inference(&http_client, "basic_test").await; @@ -425,6 +438,7 @@ async fn test_count_feedback_by_target_id_with_feedback() { #[tokio::test(flavor = "multi_thread")] async fn test_count_feedback_by_target_id_multiple_feedback() { + skip_for_postgres!(); let http_client = Client::new(); let inference_id = create_inference(&http_client, "basic_test").await; @@ -450,6 +464,7 @@ async fn test_count_feedback_by_target_id_multiple_feedback() { #[tokio::test(flavor = "multi_thread")] async fn test_count_feedback_by_target_id_nonexistent_target() { + skip_for_postgres!(); let http_client = Client::new(); let nonexistent_id = Uuid::now_v7(); @@ -470,6 +485,7 @@ async fn test_count_feedback_by_target_id_nonexistent_target() { #[tokio::test(flavor = "multi_thread")] async fn test_count_feedback_by_target_id_invalid_uuid() { + skip_for_postgres!(); let http_client = Client::new(); let url = get_gateway_endpoint("/internal/feedback/not-a-valid-uuid/count"); @@ -486,6 +502,7 @@ async fn test_count_feedback_by_target_id_invalid_uuid() { #[tokio::test(flavor = "multi_thread")] async fn test_get_cumulative_feedback_timeseries_with_data() { + skip_for_postgres!(); let http_client = Client::new(); // Create an inference to generate feedback data @@ -518,6 +535,7 @@ async fn test_get_cumulative_feedback_timeseries_with_data() { #[tokio::test(flavor = "multi_thread")] async fn test_get_cumulative_feedback_timeseries_with_variant_filter() { + skip_for_postgres!(); let http_client = Client::new(); // Query cumulative feedback timeseries with variant_names filter @@ -537,6 +555,7 @@ async fn test_get_cumulative_feedback_timeseries_with_variant_filter() { #[tokio::test(flavor = "multi_thread")] async fn test_get_cumulative_feedback_timeseries_different_time_windows() { + skip_for_postgres!(); let http_client = Client::new(); // Test different time windows (excluding cumulative which is not supported) @@ -559,6 +578,7 @@ async fn test_get_cumulative_feedback_timeseries_different_time_windows() { #[tokio::test(flavor = "multi_thread")] async fn test_get_cumulative_feedback_timeseries_cumulative_window_returns_error() { + skip_for_postgres!(); let http_client = Client::new(); // Cumulative time window is not supported for feedback timeseries @@ -576,6 +596,7 @@ async fn test_get_cumulative_feedback_timeseries_cumulative_window_returns_error #[tokio::test(flavor = "multi_thread")] async fn test_get_cumulative_feedback_timeseries_missing_params() { + skip_for_postgres!(); let http_client = Client::new(); // Missing function_name @@ -625,6 +646,7 @@ async fn test_get_cumulative_feedback_timeseries_missing_params() { #[tokio::test(flavor = "multi_thread")] async fn test_get_cumulative_feedback_timeseries_invalid_time_window() { + skip_for_postgres!(); let http_client = Client::new(); let url = get_gateway_endpoint( @@ -672,6 +694,7 @@ async fn submit_demonstration_feedback( #[tokio::test(flavor = "multi_thread")] async fn test_get_demonstration_feedback_with_data() { + skip_for_postgres!(); let http_client = Client::new(); // Create an inference @@ -712,6 +735,7 @@ async fn test_get_demonstration_feedback_with_data() { #[tokio::test(flavor = "multi_thread")] async fn test_get_demonstration_feedback_with_pagination() { + skip_for_postgres!(); let http_client = Client::new(); // Create an inference @@ -743,6 +767,7 @@ async fn test_get_demonstration_feedback_with_pagination() { #[tokio::test(flavor = "multi_thread")] async fn test_get_demonstration_feedback_nonexistent_target() { + skip_for_postgres!(); let http_client = Client::new(); // Use a UUID that likely doesn't exist @@ -769,6 +794,7 @@ async fn test_get_demonstration_feedback_nonexistent_target() { #[tokio::test(flavor = "multi_thread")] async fn test_get_demonstration_feedback_invalid_uuid() { + skip_for_postgres!(); let http_client = Client::new(); // Use an invalid UUID @@ -785,6 +811,7 @@ async fn test_get_demonstration_feedback_invalid_uuid() { #[tokio::test(flavor = "multi_thread")] async fn test_get_demonstration_feedback_rejects_both_before_and_after() { + skip_for_postgres!(); let http_client = Client::new(); let inference_id = Uuid::now_v7(); diff --git a/tensorzero-core/tests/e2e/endpoints/internal/functions.rs b/tensorzero-core/tests/e2e/endpoints/internal/functions.rs index 95ff543bf97..e80b28b42ac 100644 --- a/tensorzero-core/tests/e2e/endpoints/internal/functions.rs +++ b/tensorzero-core/tests/e2e/endpoints/internal/functions.rs @@ -66,6 +66,7 @@ async fn submit_inference_feedback( #[tokio::test(flavor = "multi_thread")] async fn test_get_function_metrics_with_feedback() { + skip_for_postgres!(); let http_client = Client::new(); // Create inferences with different types of feedback @@ -110,6 +111,7 @@ async fn test_get_function_metrics_with_feedback() { #[tokio::test(flavor = "multi_thread")] async fn test_get_function_metrics_with_variant_filter() { + skip_for_postgres!(); let http_client = Client::new(); // Create inference (which will use a specific variant based on the config) @@ -142,6 +144,7 @@ async fn test_get_function_metrics_with_variant_filter() { #[tokio::test(flavor = "multi_thread")] async fn test_get_function_metrics_nonexistent_function() { + skip_for_postgres!(); let http_client = Client::new(); // Try to query metrics for a function that doesn't exist @@ -161,6 +164,7 @@ async fn test_get_function_metrics_nonexistent_function() { #[tokio::test(flavor = "multi_thread")] async fn test_get_variant_performances_cumulative() { + skip_for_postgres!(); let http_client = Client::new(); // Create an inference and submit feedback @@ -186,6 +190,7 @@ async fn test_get_variant_performances_cumulative() { #[tokio::test(flavor = "multi_thread")] async fn test_get_variant_performances_with_time_window() { + skip_for_postgres!(); let http_client = Client::new(); // Create an inference and submit feedback @@ -220,6 +225,7 @@ async fn test_get_variant_performances_with_time_window() { #[tokio::test(flavor = "multi_thread")] async fn test_get_variant_performances_with_variant_filter() { + skip_for_postgres!(); let http_client = Client::new(); // Create an inference and submit feedback @@ -251,6 +257,7 @@ async fn test_get_variant_performances_with_variant_filter() { #[tokio::test(flavor = "multi_thread")] async fn test_get_variant_performances_nonexistent_function() { + skip_for_postgres!(); let http_client = Client::new(); let url = get_gateway_endpoint( @@ -267,6 +274,7 @@ async fn test_get_variant_performances_nonexistent_function() { #[tokio::test(flavor = "multi_thread")] async fn test_get_variant_performances_nonexistent_metric() { + skip_for_postgres!(); let http_client = Client::new(); let url = get_gateway_endpoint( @@ -283,6 +291,7 @@ async fn test_get_variant_performances_nonexistent_metric() { #[tokio::test(flavor = "multi_thread")] async fn test_get_variant_performances_nonexistent_variant() { + skip_for_postgres!(); let http_client = Client::new(); let url = get_gateway_endpoint( @@ -299,6 +308,7 @@ async fn test_get_variant_performances_nonexistent_variant() { #[tokio::test(flavor = "multi_thread")] async fn test_get_variant_performances_missing_required_params() { + skip_for_postgres!(); let http_client = Client::new(); // Missing metric_name diff --git a/tensorzero-core/tests/e2e/endpoints/internal/inference_count.rs b/tensorzero-core/tests/e2e/endpoints/internal/inference_count.rs index f8b9a2b459d..388142ac866 100644 --- a/tensorzero-core/tests/e2e/endpoints/internal/inference_count.rs +++ b/tensorzero-core/tests/e2e/endpoints/internal/inference_count.rs @@ -111,6 +111,8 @@ async fn submit_episode_feedback( #[tokio::test(flavor = "multi_thread")] async fn test_get_inference_count_chat_function() { + // create_inference doesn't write to Postgres yet + skip_for_postgres!(); let client = Client::new(); // First get the current count @@ -144,6 +146,9 @@ async fn test_get_inference_count_chat_function() { #[tokio::test(flavor = "multi_thread")] async fn test_get_inference_count_json_function() { + // create_inference doesn't write to Postgres yet + skip_for_postgres!(); + let client = Client::new(); // First get the current count @@ -258,6 +263,7 @@ async fn test_get_inference_count_unknown_variant() { #[tokio::test(flavor = "multi_thread")] async fn test_get_feedback_stats_float_metric() { + skip_for_postgres!(); let client = Client::new(); // Create an inference @@ -295,6 +301,7 @@ async fn test_get_feedback_stats_float_metric() { #[tokio::test(flavor = "multi_thread")] async fn test_get_feedback_stats_boolean_metric() { + skip_for_postgres!(); let client = Client::new(); // Create an inference @@ -332,6 +339,7 @@ async fn test_get_feedback_stats_boolean_metric() { #[tokio::test(flavor = "multi_thread")] async fn test_get_feedback_stats_with_threshold() { + skip_for_postgres!(); let client = Client::new(); // Create an inference and submit feedback with a specific value @@ -403,6 +411,7 @@ async fn test_get_feedback_stats_with_threshold() { #[tokio::test(flavor = "multi_thread")] async fn test_get_feedback_stats_demonstration() { + skip_for_postgres!(); let client = Client::new(); // Use json_success which should be able to have demonstrations let url = @@ -450,6 +459,7 @@ async fn test_get_feedback_stats_unknown_metric() { #[tokio::test(flavor = "multi_thread")] async fn test_get_feedback_stats_episode_level_boolean_metric() { + skip_for_postgres!(); let client = Client::new(); // Create an inference to get an episode_id @@ -488,6 +498,7 @@ async fn test_get_feedback_stats_episode_level_boolean_metric() { #[tokio::test(flavor = "multi_thread")] async fn test_get_feedback_stats_episode_level_float_metric() { + skip_for_postgres!(); let client = Client::new(); // Create an inference to get an episode_id @@ -574,7 +585,7 @@ pub async fn test_get_inference_count_basic() { #[tokio::test(flavor = "multi_thread")] pub async fn test_get_inference_count_with_variant_filter() { - let res = get_inference_count_fixture("basic_test", "variant_name=test") + let res = get_inference_count_fixture("write_haiku", "variant_name=initial_prompt_gpt4o_mini") .await .unwrap(); @@ -587,14 +598,14 @@ pub async fn test_get_inference_count_with_variant_filter() { #[tokio::test(flavor = "multi_thread")] pub async fn test_get_inference_count_group_by_variant() { - let res = get_inference_count_fixture("basic_test", "group_by=variant") + let res = get_inference_count_fixture("write_haiku", "group_by=variant") .await .unwrap(); let total_count = res.inference_count; assert!( total_count >= 1, - "Expected at least 1 inference for basic_test, got {total_count}" + "Expected at least 1 inference for write_haiku, got {total_count}" ); let count_by_variant = res diff --git a/tensorzero-core/tests/e2e/endpoints/internal/inference_metadata.rs b/tensorzero-core/tests/e2e/endpoints/internal/inference_metadata.rs index 220d7741e8e..2afa25d4f83 100644 --- a/tensorzero-core/tests/e2e/endpoints/internal/inference_metadata.rs +++ b/tensorzero-core/tests/e2e/endpoints/internal/inference_metadata.rs @@ -8,6 +8,7 @@ use crate::common::get_gateway_endpoint; #[tokio::test(flavor = "multi_thread")] async fn test_list_inference_metadata_no_params() { + skip_for_postgres!(); let http_client = Client::new(); let url = get_gateway_endpoint("/internal/inference_metadata"); @@ -36,6 +37,7 @@ async fn test_list_inference_metadata_no_params() { #[tokio::test(flavor = "multi_thread")] async fn test_list_inference_metadata_with_limit() { + skip_for_postgres!(); let http_client = Client::new(); let url = get_gateway_endpoint("/internal/inference_metadata?limit=5"); @@ -64,6 +66,7 @@ async fn test_list_inference_metadata_with_limit() { #[tokio::test(flavor = "multi_thread")] async fn test_list_inference_metadata_before_and_after_mutually_exclusive() { + skip_for_postgres!(); let http_client = Client::new(); let id = Uuid::now_v7(); let url = get_gateway_endpoint(&format!( @@ -79,6 +82,7 @@ async fn test_list_inference_metadata_before_and_after_mutually_exclusive() { #[tokio::test(flavor = "multi_thread")] async fn test_list_inference_metadata_with_before() { + skip_for_postgres!(); let http_client = Client::new(); // Use a UUID that is likely after any existing data let cursor = Uuid::now_v7(); @@ -109,6 +113,7 @@ async fn test_list_inference_metadata_with_before() { #[tokio::test(flavor = "multi_thread")] async fn test_list_inference_metadata_with_after() { + skip_for_postgres!(); let http_client = Client::new(); // Use a UUID that is likely before any existing data (nil UUID) let cursor = Uuid::nil(); @@ -139,6 +144,7 @@ async fn test_list_inference_metadata_with_after() { #[tokio::test(flavor = "multi_thread")] async fn test_list_inference_metadata_with_function_name() { + skip_for_postgres!(); let http_client = Client::new(); let url = get_gateway_endpoint("/internal/inference_metadata?function_name=basic_test&limit=10"); @@ -162,6 +168,7 @@ async fn test_list_inference_metadata_with_function_name() { #[tokio::test(flavor = "multi_thread")] async fn test_list_inference_metadata_with_variant_name() { + skip_for_postgres!(); let http_client = Client::new(); let url = get_gateway_endpoint("/internal/inference_metadata?variant_name=test&limit=10"); @@ -184,6 +191,7 @@ async fn test_list_inference_metadata_with_variant_name() { #[tokio::test(flavor = "multi_thread")] async fn test_list_inference_metadata_with_episode_id() { + skip_for_postgres!(); let http_client = Client::new(); // First, get some inference metadata to find a valid episode_id let url = get_gateway_endpoint("/internal/inference_metadata?limit=1"); @@ -221,6 +229,7 @@ async fn test_list_inference_metadata_with_episode_id() { #[tokio::test(flavor = "multi_thread")] async fn test_list_inference_metadata_with_multiple_filters() { + skip_for_postgres!(); let http_client = Client::new(); let url = get_gateway_endpoint( "/internal/inference_metadata?function_name=basic_test&variant_name=test&limit=10", diff --git a/tensorzero-core/tests/e2e/endpoints/internal/model_inferences.rs b/tensorzero-core/tests/e2e/endpoints/internal/model_inferences.rs index 2b82ffbdce1..191a200d00e 100644 --- a/tensorzero-core/tests/e2e/endpoints/internal/model_inferences.rs +++ b/tensorzero-core/tests/e2e/endpoints/internal/model_inferences.rs @@ -35,6 +35,7 @@ async fn get_available_inference_id(request: ListInferencesRequest) -> Uuid { #[tokio::test(flavor = "multi_thread")] async fn test_get_model_inferences_for_chat_inference() { + skip_for_postgres!(); // First, get an inference_id from the list_inferences endpoint let inference_id = get_available_inference_id(ListInferencesRequest { function_name: Some("write_haiku".to_string()), @@ -76,6 +77,7 @@ async fn test_get_model_inferences_for_chat_inference() { #[tokio::test(flavor = "multi_thread")] async fn test_get_model_inferences_for_json_inference() { + skip_for_postgres!(); // Get an inference_id for a JSON function let inference_id = get_available_inference_id(ListInferencesRequest { function_name: Some("extract_entities".to_string()), @@ -109,6 +111,7 @@ async fn test_get_model_inferences_for_json_inference() { #[tokio::test(flavor = "multi_thread")] async fn test_get_model_inferences_nonexistent_id() { + skip_for_postgres!(); // Use a random UUID that doesn't exist let nonexistent_id = Uuid::now_v7(); @@ -129,6 +132,7 @@ async fn test_get_model_inferences_nonexistent_id() { #[tokio::test(flavor = "multi_thread")] async fn test_get_model_inferences_invalid_uuid() { + skip_for_postgres!(); let http_client = Client::new(); let url = get_gateway_endpoint("/internal/model_inferences/not-a-valid-uuid"); diff --git a/tensorzero-core/tests/e2e/endpoints/internal/models.rs b/tensorzero-core/tests/e2e/endpoints/internal/models.rs index a848d2fc855..af9c5d75440 100644 --- a/tensorzero-core/tests/e2e/endpoints/internal/models.rs +++ b/tensorzero-core/tests/e2e/endpoints/internal/models.rs @@ -9,6 +9,7 @@ use crate::common::get_gateway_endpoint; #[tokio::test(flavor = "multi_thread")] async fn test_count_models_endpoint() { + skip_for_postgres!(); let http_client = Client::new(); let url = get_gateway_endpoint("/internal/models/count"); @@ -30,6 +31,7 @@ async fn test_count_models_endpoint() { #[tokio::test(flavor = "multi_thread")] async fn test_model_usage_endpoint() { + skip_for_postgres!(); let http_client = Client::new(); let url = get_gateway_endpoint("/internal/models/usage?time_window=week&max_periods=10"); @@ -49,6 +51,7 @@ async fn test_model_usage_endpoint() { #[tokio::test(flavor = "multi_thread")] async fn test_model_latency_endpoint() { + skip_for_postgres!(); let http_client = Client::new(); let url = get_gateway_endpoint("/internal/models/latency?time_window=week"); @@ -68,6 +71,7 @@ async fn test_model_latency_endpoint() { #[tokio::test(flavor = "multi_thread")] async fn test_model_usage_endpoint_missing_params() { + skip_for_postgres!(); let http_client = Client::new(); // Missing required parameters let url = get_gateway_endpoint("/internal/models/usage"); @@ -82,6 +86,7 @@ async fn test_model_usage_endpoint_missing_params() { #[tokio::test(flavor = "multi_thread")] async fn test_model_latency_endpoint_missing_params() { + skip_for_postgres!(); let http_client = Client::new(); // Missing required parameters let url = get_gateway_endpoint("/internal/models/latency"); diff --git a/tensorzero-core/tests/e2e/endpoints/internal/workflow_evaluations.rs b/tensorzero-core/tests/e2e/endpoints/internal/workflow_evaluations.rs index be5c09a1944..0044c40d070 100644 --- a/tensorzero-core/tests/e2e/endpoints/internal/workflow_evaluations.rs +++ b/tensorzero-core/tests/e2e/endpoints/internal/workflow_evaluations.rs @@ -14,6 +14,7 @@ use crate::common::get_gateway_endpoint; #[tokio::test(flavor = "multi_thread")] async fn test_get_workflow_evaluation_projects_endpoint() { + skip_for_postgres!(); let http_client = Client::new(); let url = get_gateway_endpoint("/internal/workflow_evaluations/projects"); @@ -46,6 +47,7 @@ async fn test_get_workflow_evaluation_projects_endpoint() { #[tokio::test(flavor = "multi_thread")] async fn test_get_workflow_evaluation_projects_with_pagination() { + skip_for_postgres!(); let http_client = Client::new(); let url = get_gateway_endpoint("/internal/workflow_evaluations/projects?limit=1&offset=0"); @@ -68,6 +70,7 @@ async fn test_get_workflow_evaluation_projects_with_pagination() { #[tokio::test(flavor = "multi_thread")] async fn test_get_workflow_evaluation_project_count_endpoint() { + skip_for_postgres!(); let http_client = Client::new(); let url = get_gateway_endpoint("/internal/workflow_evaluations/projects/count"); @@ -88,6 +91,7 @@ async fn test_get_workflow_evaluation_project_count_endpoint() { #[tokio::test(flavor = "multi_thread")] async fn test_list_workflow_evaluation_runs_endpoint() { + skip_for_postgres!(); let http_client = Client::new(); let url = get_gateway_endpoint("/internal/workflow_evaluations/list_runs"); @@ -113,6 +117,7 @@ async fn test_list_workflow_evaluation_runs_endpoint() { #[tokio::test(flavor = "multi_thread")] async fn test_list_workflow_evaluation_runs_with_pagination() { + skip_for_postgres!(); let http_client = Client::new(); let url = get_gateway_endpoint("/internal/workflow_evaluations/list_runs?limit=1&offset=0"); @@ -135,6 +140,7 @@ async fn test_list_workflow_evaluation_runs_with_pagination() { #[tokio::test(flavor = "multi_thread")] async fn test_list_workflow_evaluation_runs_with_project_filter() { + skip_for_postgres!(); let http_client = Client::new(); let url = get_gateway_endpoint("/internal/workflow_evaluations/list_runs?project_name=21_questions"); @@ -160,6 +166,7 @@ async fn test_list_workflow_evaluation_runs_with_project_filter() { #[tokio::test(flavor = "multi_thread")] async fn test_count_workflow_evaluation_runs_endpoint() { + skip_for_postgres!(); let http_client = Client::new(); let url = get_gateway_endpoint("/internal/workflow_evaluations/runs/count"); @@ -180,6 +187,7 @@ async fn test_count_workflow_evaluation_runs_endpoint() { #[tokio::test(flavor = "multi_thread")] async fn test_search_workflow_evaluation_runs_endpoint() { + skip_for_postgres!(); let http_client = Client::new(); let url = get_gateway_endpoint("/internal/workflow_evaluations/runs/search"); @@ -201,6 +209,7 @@ async fn test_search_workflow_evaluation_runs_endpoint() { #[tokio::test(flavor = "multi_thread")] async fn test_search_workflow_evaluation_runs_with_project_filter() { + skip_for_postgres!(); let http_client = Client::new(); let url = get_gateway_endpoint( "/internal/workflow_evaluations/runs/search?project_name=21_questions", @@ -227,6 +236,7 @@ async fn test_search_workflow_evaluation_runs_with_project_filter() { #[tokio::test(flavor = "multi_thread")] async fn test_search_workflow_evaluation_runs_with_search_query() { + skip_for_postgres!(); let http_client = Client::new(); let url = get_gateway_endpoint("/internal/workflow_evaluations/runs/search?q=baseline"); @@ -254,6 +264,7 @@ async fn test_search_workflow_evaluation_runs_with_search_query() { #[tokio::test(flavor = "multi_thread")] async fn test_search_workflow_evaluation_runs_with_pagination() { + skip_for_postgres!(); let http_client = Client::new(); let url = get_gateway_endpoint("/internal/workflow_evaluations/runs/search?limit=1&offset=0"); @@ -276,6 +287,7 @@ async fn test_search_workflow_evaluation_runs_with_pagination() { #[tokio::test(flavor = "multi_thread")] async fn test_get_workflow_evaluation_runs_endpoint() { + skip_for_postgres!(); let http_client = Client::new(); // Use a known run ID from the fixture data let run_id = "01968d04-142c-7e53-8ea7-3a3255b518dc"; @@ -308,6 +320,7 @@ async fn test_get_workflow_evaluation_runs_endpoint() { #[tokio::test(flavor = "multi_thread")] async fn test_get_workflow_evaluation_runs_multiple_ids() { + skip_for_postgres!(); let http_client = Client::new(); // Use known run IDs from the fixture data let run_id1 = "01968d04-142c-7e53-8ea7-3a3255b518dc"; @@ -336,6 +349,7 @@ async fn test_get_workflow_evaluation_runs_multiple_ids() { #[tokio::test(flavor = "multi_thread")] async fn test_get_workflow_evaluation_runs_with_project_filter() { + skip_for_postgres!(); let http_client = Client::new(); let run_id = "01968d04-142c-7e53-8ea7-3a3255b518dc"; let url = get_gateway_endpoint(&format!( @@ -366,6 +380,7 @@ async fn test_get_workflow_evaluation_runs_with_project_filter() { #[tokio::test(flavor = "multi_thread")] async fn test_get_workflow_evaluation_runs_empty_ids() { + skip_for_postgres!(); let http_client = Client::new(); let url = get_gateway_endpoint("/internal/workflow_evaluations/get_runs?run_ids="); @@ -393,6 +408,7 @@ async fn test_get_workflow_evaluation_runs_empty_ids() { #[tokio::test(flavor = "multi_thread")] async fn test_list_episodes_by_task_name_endpoint() { + skip_for_postgres!(); let http_client = Client::new(); // Use a known run_id from the fixture data let run_id = "0196a0e5-9600-7c83-ab3b-da81097b66cd"; @@ -434,6 +450,7 @@ async fn test_list_episodes_by_task_name_endpoint() { #[tokio::test(flavor = "multi_thread")] async fn test_list_episodes_by_task_name_with_multiple_run_ids() { + skip_for_postgres!(); let http_client = Client::new(); // Use multiple run_ids from the fixture data let run_ids = "0196a0e5-9600-7c83-ab3b-da81097b66cd,0196a0e5-9600-7c83-ab3b-dabb145a9dbe"; @@ -459,6 +476,7 @@ async fn test_list_episodes_by_task_name_with_multiple_run_ids() { #[tokio::test(flavor = "multi_thread")] async fn test_list_episodes_by_task_name_with_pagination() { + skip_for_postgres!(); let http_client = Client::new(); let run_id = "0196a0e5-9600-7c83-ab3b-da81097b66cd"; let url = get_gateway_endpoint(&format!( @@ -484,6 +502,7 @@ async fn test_list_episodes_by_task_name_with_pagination() { #[tokio::test(flavor = "multi_thread")] async fn test_list_episodes_by_task_name_empty_run_ids() { + skip_for_postgres!(); let http_client = Client::new(); let url = get_gateway_endpoint("/internal/workflow_evaluations/episodes_by_task_name"); @@ -505,6 +524,7 @@ async fn test_list_episodes_by_task_name_empty_run_ids() { #[tokio::test(flavor = "multi_thread")] async fn test_count_episode_groups_endpoint() { + skip_for_postgres!(); let http_client = Client::new(); // Use a known run_id from the fixture data let run_id = "0196a0e5-9600-7c83-ab3b-da81097b66cd"; @@ -531,6 +551,7 @@ async fn test_count_episode_groups_endpoint() { #[tokio::test(flavor = "multi_thread")] async fn test_count_episode_groups_with_multiple_run_ids() { + skip_for_postgres!(); let http_client = Client::new(); // Use multiple run_ids from the fixture data let run_ids = "0196a0e5-9600-7c83-ab3b-da81097b66cd,0196a0e5-9600-7c83-ab3b-dabb145a9dbe"; @@ -557,6 +578,7 @@ async fn test_count_episode_groups_with_multiple_run_ids() { #[tokio::test(flavor = "multi_thread")] async fn test_count_episode_groups_empty_run_ids() { + skip_for_postgres!(); let http_client = Client::new(); let url = get_gateway_endpoint("/internal/workflow_evaluations/episodes_by_task_name/count"); @@ -575,6 +597,7 @@ async fn test_count_episode_groups_empty_run_ids() { #[tokio::test(flavor = "multi_thread")] async fn test_count_matches_list_length_endpoint() { + skip_for_postgres!(); let http_client = Client::new(); let run_id = "0196a0e5-9600-7c83-ab3b-da81097b66cd"; @@ -606,6 +629,7 @@ async fn test_count_matches_list_length_endpoint() { #[tokio::test(flavor = "multi_thread")] async fn test_get_workflow_evaluation_run_episodes_endpoint() { + skip_for_postgres!(); let http_client = Client::new(); // Use a known run ID from the fixture data let run_id = "01968d04-142c-7e53-8ea7-3a3255b518dc"; @@ -643,6 +667,7 @@ async fn test_get_workflow_evaluation_run_episodes_endpoint() { #[tokio::test(flavor = "multi_thread")] async fn test_get_workflow_evaluation_run_episodes_with_pagination() { + skip_for_postgres!(); let http_client = Client::new(); let run_id = "01968d04-142c-7e53-8ea7-3a3255b518dc"; let url = get_gateway_endpoint(&format!( @@ -668,6 +693,7 @@ async fn test_get_workflow_evaluation_run_episodes_with_pagination() { #[tokio::test(flavor = "multi_thread")] async fn test_get_workflow_evaluation_run_episodes_beyond_offset() { + skip_for_postgres!(); let http_client = Client::new(); let run_id = "01968d04-142c-7e53-8ea7-3a3255b518dc"; let url = get_gateway_endpoint(&format!( @@ -698,6 +724,7 @@ async fn test_get_workflow_evaluation_run_episodes_beyond_offset() { #[tokio::test(flavor = "multi_thread")] async fn test_count_workflow_evaluation_run_episodes_endpoint() { + skip_for_postgres!(); let http_client = Client::new(); let run_id = "01968d04-142c-7e53-8ea7-3a3255b518dc"; let url = get_gateway_endpoint(&format!( @@ -723,6 +750,7 @@ async fn test_count_workflow_evaluation_run_episodes_endpoint() { #[tokio::test(flavor = "multi_thread")] async fn test_count_workflow_evaluation_run_episodes_nonexistent_run() { + skip_for_postgres!(); let http_client = Client::new(); // Use a valid but non-existent UUIDv7 let run_id = "01942e26-4693-7e80-8591-47b98e25d999"; @@ -749,6 +777,7 @@ async fn test_count_workflow_evaluation_run_episodes_nonexistent_run() { #[tokio::test(flavor = "multi_thread")] async fn test_count_workflow_evaluation_run_episodes_missing_run_id() { + skip_for_postgres!(); let http_client = Client::new(); let url = get_gateway_endpoint("/internal/workflow_evaluations/run_episodes/count"); diff --git a/tensorzero-core/tests/e2e/endpoints/mod.rs b/tensorzero-core/tests/e2e/endpoints/mod.rs index 4411d1d4c9e..00381f83b63 100644 --- a/tensorzero-core/tests/e2e/endpoints/mod.rs +++ b/tensorzero-core/tests/e2e/endpoints/mod.rs @@ -1,3 +1,6 @@ +#[macro_use] mod datasets; +#[macro_use] mod internal; +#[macro_use] mod stored_inferences; diff --git a/tensorzero-core/tests/e2e/endpoints/stored_inferences/get_inferences.rs b/tensorzero-core/tests/e2e/endpoints/stored_inferences/get_inferences.rs index 1c237baf1b3..80377fe8955 100644 --- a/tensorzero-core/tests/e2e/endpoints/stored_inferences/get_inferences.rs +++ b/tensorzero-core/tests/e2e/endpoints/stored_inferences/get_inferences.rs @@ -39,6 +39,7 @@ async fn list_inferences(request: Value) -> Result, Box Date: Tue, 27 Jan 2026 17:57:19 -0500 Subject: [PATCH 40/46] Bump rlt to 0.4.1 (#5891) This unblocks dependabot (it required a manual change) --- Cargo.lock | 529 ++++++++++-------- Cargo.toml | 2 +- .../tests/load/feedback/src/main.rs | 2 +- .../load/rate-limit-load-test/src/main.rs | 2 +- 4 files changed, 293 insertions(+), 242 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 4db1c9649c8..fb5ba9f06ef 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -99,6 +99,25 @@ dependencies = [ "libc", ] +[[package]] +name = "ansi-str" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "060de1453b69f46304b28274f382132f4e72c55637cf362920926a70d090890d" +dependencies = [ + "ansitok", +] + +[[package]] +name = "ansitok" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c0a8acea8c2f1c60f0a92a8cd26bf96ca97db56f10bbcab238bbe0cceba659ee" +dependencies = [ + "nom 7.1.3", + "vte", +] + [[package]] name = "anstream" version = "0.6.21" @@ -1123,12 +1142,6 @@ dependencies = [ "either", ] -[[package]] -name = "cassowary" -version = "0.3.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "df8670b8c7b9dae1793364eafadf7239c40d669904660c5960d74cfd80b46a53" - [[package]] name = "castaway" version = "0.2.4" @@ -1221,7 +1234,7 @@ version = "4.5.49" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2a0b5487afeab2deb2ff4e03a807ad1a03ac532ff5a2cee5d86884440c7f7671" dependencies = [ - "heck 0.5.0", + "heck", "proc-macro2", "quote", "syn 2.0.114", @@ -1290,13 +1303,14 @@ dependencies = [ [[package]] name = "compact_str" -version = "0.7.1" +version = "0.9.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f86b9c4c00838774a6d902ef931eff7470720c51d90c2e32cfe15dc304737b3f" +checksum = "3fdb1325a1cece981e8a296ab8f0f9b63ae357bd0784a9faaf548cc7b480707a" dependencies = [ "castaway", "cfg-if", "itoa", + "rustversion", "ryu", "static_assertions", ] @@ -1339,7 +1353,7 @@ dependencies = [ "encode_unicode", "libc", "once_cell", - "unicode-width 0.2.2", + "unicode-width", "windows-sys 0.61.2", ] @@ -1477,15 +1491,17 @@ checksum = "d0a5c400df2834b80a4c3327b3aad3a4c4cd4de0629063962b03235697506a28" [[package]] name = "crossterm" -version = "0.27.0" +version = "0.29.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f476fe445d41c9e991fd07515a6f463074b782242ccf4a5b7b1d1012e70824df" +checksum = "d8b9f2e4c67f833b660cdb0a3523065869fb35570177239812ed4c905aeff87b" dependencies = [ "bitflags", "crossterm_winapi", - "libc", - "mio 0.8.11", + "derive_more", + "document-features", + "mio", "parking_lot", + "rustix", "signal-hook", "signal-hook-mio", "winapi", @@ -1562,6 +1578,16 @@ dependencies = [ "darling_macro 0.21.3", ] +[[package]] +name = "darling" +version = "0.23.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "25ae13da2f202d56bd7f91c25fba009e7717a1e4a1cc98a76d844b65ae912e9d" +dependencies = [ + "darling_core 0.23.0", + "darling_macro 0.23.0", +] + [[package]] name = "darling_core" version = "0.14.4" @@ -1604,6 +1630,19 @@ dependencies = [ "syn 2.0.114", ] +[[package]] +name = "darling_core" +version = "0.23.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9865a50f7c335f53564bb694ef660825eb8610e0a53d3e11bf1b0d3df31e03b0" +dependencies = [ + "ident_case", + "proc-macro2", + "quote", + "strsim 0.11.1", + "syn 2.0.114", +] + [[package]] name = "darling_macro" version = "0.14.4" @@ -1638,16 +1677,14 @@ dependencies = [ ] [[package]] -name = "dashmap" -version = "5.5.3" +name = "darling_macro" +version = "0.23.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "978747c1d849a7d2ee5e8adc0159961c48fb7e5db2f06af6723b80123bb53856" +checksum = "ac3984ec7bd6cfa798e62b4a642426a5be0e68f9401cfc2a01e3fa9ea2fcdb8d" dependencies = [ - "cfg-if", - "hashbrown 0.14.5", - "lock_api", - "once_cell", - "parking_lot_core", + "darling_core 0.23.0", + "quote", + "syn 2.0.114", ] [[package]] @@ -1767,6 +1804,28 @@ dependencies = [ "syn 2.0.114", ] +[[package]] +name = "derive_more" +version = "2.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d751e9e49156b02b44f9c1815bcb94b984cdcc4396ecc32521c739452808b134" +dependencies = [ + "derive_more-impl", +] + +[[package]] +name = "derive_more-impl" +version = "2.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "799a97264921d8623a957f6c3b9011f3b5492f557bbb7a5a19b7fa6d06ba8dcb" +dependencies = [ + "convert_case", + "proc-macro2", + "quote", + "rustc_version", + "syn 2.0.114", +] + [[package]] name = "digest" version = "0.10.7" @@ -1790,6 +1849,15 @@ dependencies = [ "syn 2.0.114", ] +[[package]] +name = "document-features" +version = "0.2.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d4b8a88685455ed29a21542a33abd9cb6510b6b129abadabdcef0f4c55bc8f61" +dependencies = [ + "litrs", +] + [[package]] name = "dotenvy" version = "0.15.7" @@ -1960,6 +2028,16 @@ dependencies = [ "syn 2.0.114", ] +[[package]] +name = "env_filter" +version = "0.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1bf3c259d255ca70051b30e2e95b5446cdb8949ac4cd22c0d7fd634d89f568e2" +dependencies = [ + "log", + "regex", +] + [[package]] name = "equator" version = "0.4.2" @@ -1993,7 +2071,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" dependencies = [ "libc", - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -2337,15 +2415,6 @@ dependencies = [ "slab", ] -[[package]] -name = "fxhash" -version = "0.2.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c31b6d751ae2c7f11320402d34e41349dd1016f8d5d45e48c4312bc8625af50c" -dependencies = [ - "byteorder", -] - [[package]] name = "gag" version = "1.0.0" @@ -2558,22 +2627,25 @@ dependencies = [ [[package]] name = "governor" -version = "0.6.3" +version = "0.10.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "68a7f542ee6b35af73b06abc0dad1c1bae89964e4e253bc4b587b91c9637867b" +checksum = "9efcab3c1958580ff1f25a2a41be1668f7603d849bb63af523b208a3cc1223b8" dependencies = [ "cfg-if", - "dashmap 5.5.3", - "futures", + "dashmap", + "futures-sink", "futures-timer", - "no-std-compat", + "futures-util", + "getrandom 0.3.4", + "hashbrown 0.16.1", "nonzero_ext", "parking_lot", "portable-atomic", "quanta", - "rand 0.8.5", + "rand 0.9.2", "smallvec", "spinning_top", + "web-time 1.1.0", ] [[package]] @@ -2627,8 +2699,6 @@ version = "0.15.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9229cfe53dfd69f0609a49f65461bd93001ea1ef889cd5529dd176593f5338a1" dependencies = [ - "allocator-api2", - "equivalent", "foldhash 0.1.5", ] @@ -2666,12 +2736,6 @@ dependencies = [ "num-traits", ] -[[package]] -name = "heck" -version = "0.4.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "95505c38b4572b2d910cecb0281560f54b440a19336cbbcb27bf6ce6adc6f5a8" - [[package]] name = "heck" version = "0.5.0" @@ -3088,7 +3152,7 @@ checksum = "9375e112e4b463ec1b1c6c011953545c65a30164fbab5b581df32b3abf0dcb88" dependencies = [ "console", "portable-atomic", - "unicode-width 0.2.2", + "unicode-width", "unit-prefix", "web-time 1.1.0", ] @@ -3125,6 +3189,19 @@ dependencies = [ "tracing-opentelemetry", ] +[[package]] +name = "instability" +version = "0.3.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "357b7205c6cd18dd2c86ed312d1e70add149aea98e7ef72b9fdf0270e555c11d" +dependencies = [ + "darling 0.23.0", + "indoc", + "proc-macro2", + "quote", + "syn 2.0.114", +] + [[package]] name = "interpolate_name" version = "0.2.4" @@ -3167,15 +3244,6 @@ dependencies = [ "either", ] -[[package]] -name = "itertools" -version = "0.13.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "413ee7dfc52ee1a4949ceeb7dbc8a33f2d6c088194d9f922fb8318faf1f01186" -dependencies = [ - "either", -] - [[package]] name = "itertools" version = "0.14.0" @@ -3278,6 +3346,17 @@ dependencies = [ "simple_asn1", ] +[[package]] +name = "kasuari" +version = "0.4.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8fe90c1150662e858c7d5f945089b7517b0a80d8bf7ba4b1b5ffc984e7230a5b" +dependencies = [ + "hashbrown 0.16.1", + "portable-atomic", + "thiserror 2.0.18", +] + [[package]] name = "lazy_static" version = "1.5.0" @@ -3380,6 +3459,15 @@ dependencies = [ "vcpkg", ] +[[package]] +name = "line-clipping" +version = "0.3.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5f4de44e98ddbf09375cbf4d17714d18f39195f4f4894e8524501726fd9a8a4a" +dependencies = [ + "bitflags", +] + [[package]] name = "linux-raw-sys" version = "0.11.0" @@ -3392,6 +3480,12 @@ version = "0.8.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6373607a59f0be73a39b6fe456b8192fcc3585f602af20751600e974dd455e77" +[[package]] +name = "litrs" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "11d3d7f243d5c5a8b9bb5d6dd2b1602c0cb0b9db1621bafc7ed66e35ff9fe092" + [[package]] name = "lock_api" version = "0.4.14" @@ -3418,11 +3512,11 @@ dependencies = [ [[package]] name = "lru" -version = "0.12.5" +version = "0.16.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "234cf4f4a04dc1f57e24b96cc0cd600cf2af460d4161ac5ecdd0af8e1f3b2a38" +checksum = "a1dc47f592c06f33f8e3aea9591776ec7c9f9e4124778ff8a3c3b87159f7e593" dependencies = [ - "hashbrown 0.15.5", + "hashbrown 0.16.1", ] [[package]] @@ -3610,18 +3704,6 @@ dependencies = [ "simd-adler32", ] -[[package]] -name = "mio" -version = "0.8.11" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a4a650543ca06a924e8b371db273b2756685faae30f8487da1b56505a8f78b0c" -dependencies = [ - "libc", - "log", - "wasi", - "windows-sys 0.48.0", -] - [[package]] name = "mio" version = "1.1.1" @@ -3629,6 +3711,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a69bcab0ad47271a0234d9422b131806bf3968021e5dc9328caf2d4cd58557fc" dependencies = [ "libc", + "log", "wasi", "windows-sys 0.61.2", ] @@ -3794,12 +3877,6 @@ version = "1.0.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "650eef8c711430f1a879fdd01d4745a7deea475becfb90269c06775983bbf086" -[[package]] -name = "no-std-compat" -version = "0.4.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b93853da6d84c2e3c7d730d6473e8817692dd89be387eb01b94d7f108ecb5b8c" - [[package]] name = "nohash-hasher" version = "0.2.0" @@ -3843,7 +3920,7 @@ version = "0.50.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7957b9740744892f114936ab4a57b3f487491bbeafaf8083688b16841a4240e5" dependencies = [ - "windows-sys 0.61.2", + "windows-sys 0.59.0", ] [[package]] @@ -3959,6 +4036,15 @@ dependencies = [ "libm", ] +[[package]] +name = "num_threads" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5c7398b9c8b70908f6371f47ed36737907c87c52af34c268fed0bf0ceb92ead9" +dependencies = [ + "libc", +] + [[package]] name = "object_store" version = "0.13.1" @@ -4123,13 +4209,15 @@ checksum = "1a80800c0488c3a21695ea981a54918fbb37abf04f4d0720c453632255e2ff0e" [[package]] name = "papergrid" -version = "0.11.0" +version = "0.17.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9ad43c07024ef767f9160710b3a6773976194758c7919b17e63b863db0bdf7fb" +checksum = "6978128c8b51d8f4080631ceb2302ab51e32cc6e8615f735ee2f83fd269ae3f1" dependencies = [ + "ansi-str", + "ansitok", "bytecount", "fnv", - "unicode-width 0.1.14", + "unicode-width", ] [[package]] @@ -4336,34 +4424,32 @@ dependencies = [ ] [[package]] -name = "proc-macro-error" -version = "1.0.4" +name = "proc-macro-error-attr2" +version = "2.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "da25490ff9892aab3fcf7c36f08cfb902dd3e71ca0f9f9517bea02a73a5ce38c" +checksum = "96de42df36bb9bba5542fe9f1a054b8cc87e172759a1868aa05c1f3acc89dfc5" dependencies = [ - "proc-macro-error-attr", "proc-macro2", "quote", - "syn 1.0.109", - "version_check", ] [[package]] -name = "proc-macro-error-attr" -version = "1.0.4" +name = "proc-macro-error2" +version = "2.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a1be40180e52ecc98ad80b184934baf3d0d29f979574e439af5a55274b35f869" +checksum = "11ec05c52be0a07b08061f7dd003e7d7092e0472bc731b4af7bb1ef876109802" dependencies = [ + "proc-macro-error-attr2", "proc-macro2", "quote", - "version_check", + "syn 2.0.114", ] [[package]] name = "proc-macro2" -version = "1.0.105" +version = "1.0.106" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "535d180e0ecab6268a3e718bb9fd44db66bbbc256257165fc699dadf70d16fe7" +checksum = "8fd00f0bb2e90d81d1044c2b32617f68fcb9fa3bb7640c23e9c748e53fb30934" dependencies = [ "unicode-ident", ] @@ -4404,7 +4490,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "27c6023962132f4b30eb4c172c91ce92d933da334c59c23cddee82358ddafb0b" dependencies = [ "anyhow", - "itertools 0.14.0", + "itertools 0.11.0", "proc-macro2", "quote", "syn 2.0.114", @@ -4539,7 +4625,7 @@ version = "0.26.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "100246c0ecf400b475341b8455a9213344569af29a3c841d29270e53102e0fcf" dependencies = [ - "heck 0.5.0", + "heck", "proc-macro2", "pyo3-build-config", "quote", @@ -4645,7 +4731,7 @@ dependencies = [ "once_cell", "socket2", "tracing", - "windows-sys 0.60.2", + "windows-sys 0.52.0", ] [[package]] @@ -4758,23 +4844,65 @@ dependencies = [ [[package]] name = "ratatui" -version = "0.27.0" +version = "0.30.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d1ce67fb8ba4446454d1c8dbaeda0557ff5e94d39d5e5ed7f10a65eb4c8266bc" +dependencies = [ + "instability", + "ratatui-core", + "ratatui-crossterm", + "ratatui-widgets", +] + +[[package]] +name = "ratatui-core" +version = "0.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d16546c5b5962abf8ce6e2881e722b4e0ae3b6f1a08a26ae3573c55853ca68d3" +checksum = "5ef8dea09a92caaf73bff7adb70b76162e5937524058a7e5bff37869cbbec293" dependencies = [ "bitflags", - "cassowary", "compact_str", - "crossterm", - "itertools 0.13.0", + "hashbrown 0.16.1", + "indoc", + "itertools 0.14.0", + "kasuari", "lru", - "paste", - "stability", - "strum 0.26.3", - "strum_macros 0.26.4", + "strum", + "thiserror 2.0.18", "unicode-segmentation", "unicode-truncate", - "unicode-width 0.1.14", + "unicode-width", +] + +[[package]] +name = "ratatui-crossterm" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "577c9b9f652b4c121fb25c6a391dd06406d3b092ba68827e6d2f09550edc54b3" +dependencies = [ + "cfg-if", + "crossterm", + "instability", + "ratatui-core", +] + +[[package]] +name = "ratatui-widgets" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d7dbfa023cd4e604c2553483820c5fe8aa9d71a42eea5aa77c6e7f35756612db" +dependencies = [ + "bitflags", + "hashbrown 0.16.1", + "indoc", + "instability", + "itertools 0.14.0", + "line-clipping", + "ratatui-core", + "strum", + "time", + "unicode-segmentation", + "unicode-width", ] [[package]] @@ -5177,27 +5305,29 @@ dependencies = [ [[package]] name = "rlt" -version = "0.2.1" +version = "0.4.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3a5ae299598d2facaedddff60742e225a348c90007b418c6e37de2b326cad979" +checksum = "77b50fa381c338d9e9ae2dafa26637a73fb57fe56e173246f4da21434c6ee6fe" dependencies = [ "anyhow", "async-trait", "byte-unit", "cfg-if", + "chrono", "clap", "crossterm", "governor", "hdrhistogram", "http 1.4.0", "humantime", - "itertools 0.13.0", + "itertools 0.14.0", "log", "nonzero_ext", "parking_lot", "ratatui", "serde", "serde_json", + "strum", "tabled", "tokio", "tokio-util", @@ -5275,7 +5405,7 @@ dependencies = [ "errno", "libc", "linux-raw-sys", - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -5334,7 +5464,7 @@ dependencies = [ "security-framework", "security-framework-sys", "webpki-root-certs", - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -5676,7 +5806,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b75a19a7a740b25bc7944bdee6172368f988763b744e3d4dfe753f6b4ece40cc" dependencies = [ "libc", - "mio 0.8.11", + "mio", "signal-hook", ] @@ -5866,7 +5996,7 @@ dependencies = [ "cfg-if", "dotenvy", "either", - "heck 0.5.0", + "heck", "hex", "proc-macro2", "quote", @@ -5990,16 +6120,6 @@ dependencies = [ "uuid", ] -[[package]] -name = "stability" -version = "0.2.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d904e7009df136af5297832a3ace3370cd14ff1546a232f4f185036c2736fcac" -dependencies = [ - "quote", - "syn 2.0.114", -] - [[package]] name = "stable_deref_trait" version = "1.2.1" @@ -6035,35 +6155,13 @@ version = "0.11.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7da8b5736845d9f2fcb837ea5d9e2628564b3b043a70948a3f0b778838c5fb4f" -[[package]] -name = "strum" -version = "0.26.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8fec0f0aef304996cf250b31b5a10dee7980c85da9d759361292b8bca5a18f06" -dependencies = [ - "strum_macros 0.26.4", -] - [[package]] name = "strum" version = "0.27.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "af23d6f6c1a224baef9d3f61e287d2761385a5b88fdab4eb4c6f11aeb54c4bcf" dependencies = [ - "strum_macros 0.27.2", -] - -[[package]] -name = "strum_macros" -version = "0.26.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4c6bee85a5a24955dc440386795aa378cd9cf82acd5f764469152d2270e581be" -dependencies = [ - "heck 0.5.0", - "proc-macro2", - "quote", - "rustversion", - "syn 2.0.114", + "strum_macros", ] [[package]] @@ -6072,7 +6170,7 @@ version = "0.27.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7695ce3845ea4b33927c055a39dc438a45b059f7c1b3d91d38d10355fb8cbca7" dependencies = [ - "heck 0.5.0", + "heck", "proc-macro2", "quote", "syn 2.0.114", @@ -6149,26 +6247,28 @@ dependencies = [ [[package]] name = "tabled" -version = "0.15.0" +version = "0.20.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4c998b0c8b921495196a48aabaf1901ff28be0760136e31604f7967b0792050e" +checksum = "e39a2ee1fbcd360805a771e1b300f78cc88fec7b8d3e2f71cd37bbf23e725c7d" dependencies = [ + "ansi-str", + "ansitok", "papergrid", "tabled_derive", - "unicode-width 0.1.14", + "testing_table", ] [[package]] name = "tabled_derive" -version = "0.7.0" +version = "0.11.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4c138f99377e5d653a371cdad263615634cfc8467685dfe8e73e2b8e98f44b17" +checksum = "0ea5d1b13ca6cff1f9231ffd62f15eefd72543dab5e468735f1a456728a02846" dependencies = [ - "heck 0.4.1", - "proc-macro-error", + "heck", + "proc-macro-error2", "proc-macro2", "quote", - "syn 1.0.109", + "syn 2.0.114", ] [[package]] @@ -6205,7 +6305,7 @@ dependencies = [ "getrandom 0.3.4", "once_cell", "rustix", - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -6276,7 +6376,7 @@ dependencies = [ "chrono", "clap", "clarabel", - "dashmap 6.1.0", + "dashmap", "derive_builder 0.20.2", "durable", "durable-tools", @@ -6334,7 +6434,7 @@ dependencies = [ "serde_path_to_error", "sha2", "sqlx", - "strum 0.27.2", + "strum", "tempfile", "tensorzero", "tensorzero-auth", @@ -6485,6 +6585,16 @@ version = "0.5.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8f50febec83f5ee1df3015341d8bd429f2d1cc62bcba7ea2076759d315084683" +[[package]] +name = "testing_table" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0f8daae29995a24f65619e19d8d31dea5b389f3d853d8bf297bbf607cd0014cc" +dependencies = [ + "ansitok", + "unicode-width", +] + [[package]] name = "thiserror" version = "1.0.69" @@ -6556,7 +6666,9 @@ checksum = "f9e442fc33d7fdb45aa9bfeb312c095964abdf596f7567261062b2a7107aaabd" dependencies = [ "deranged", "itoa", + "libc", "num-conv", + "num_threads", "powerfmt", "serde_core", "time-core", @@ -6612,7 +6724,7 @@ checksum = "72a2903cd7736441aac9df9d7688bd0ce48edccaadf181c3b90be801e81d3d86" dependencies = [ "bytes", "libc", - "mio 1.1.1", + "mio", "parking_lot", "pin-project-lite", "signal-hook-registry", @@ -7018,18 +7130,19 @@ dependencies = [ [[package]] name = "tui-logger" -version = "0.11.2" +version = "0.18.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "10cd1a0f217c2180e736bc9f3282fea4af182483532c6e719081b6b1c6d6be90" +checksum = "9384df20a5244a6ab204bc4b6959b41f37f0ee7b5e0f2feb7a8a78f58e684d06" dependencies = [ "chrono", - "fxhash", + "env_filter", "lazy_static", "log", "parking_lot", "ratatui", "tracing", "tracing-subscriber", + "unicode-segmentation", ] [[package]] @@ -7085,21 +7198,15 @@ checksum = "f6ccf251212114b54433ec949fd6a7841275f9ada20dddd2f29e9ceea4501493" [[package]] name = "unicode-truncate" -version = "1.1.0" +version = "2.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b3644627a5af5fa321c95b9b235a72fd24cd29c648c2c379431e6628655627bf" +checksum = "16b380a1238663e5f8a691f9039c73e1cdae598a30e9855f541d29b08b53e9a5" dependencies = [ - "itertools 0.13.0", + "itertools 0.14.0", "unicode-segmentation", - "unicode-width 0.1.14", + "unicode-width", ] -[[package]] -name = "unicode-width" -version = "0.1.14" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7dd6e30e90baa6f72411720665d41d89b9a3d039dc45b8faea1ddd07f617f6af" - [[package]] name = "unicode-width" version = "0.2.2" @@ -7224,6 +7331,16 @@ version = "0.8.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5c3082ca00d5a5ef149bb8b555a72ae84c9c59f7250f013ac822ac2e49b19c64" +[[package]] +name = "vte" +version = "0.14.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "231fdcd7ef3037e8330d8e17e61011a2c244126acc0a982f4040ac3f9f0bc077" +dependencies = [ + "arrayvec", + "memchr", +] + [[package]] name = "walkdir" version = "2.5.0" @@ -7431,7 +7548,7 @@ version = "0.1.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" dependencies = [ - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -7519,15 +7636,6 @@ dependencies = [ "windows-targets 0.42.2", ] -[[package]] -name = "windows-sys" -version = "0.48.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "677d2418bec65e3338edb076e806bc1ec15693c5d0104683f2efe857f61056a9" -dependencies = [ - "windows-targets 0.48.5", -] - [[package]] name = "windows-sys" version = "0.52.0" @@ -7579,21 +7687,6 @@ dependencies = [ "windows_x86_64_msvc 0.42.2", ] -[[package]] -name = "windows-targets" -version = "0.48.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9a2fa6e2155d7247be68c096456083145c183cbbbc2764150dda45a87197940c" -dependencies = [ - "windows_aarch64_gnullvm 0.48.5", - "windows_aarch64_msvc 0.48.5", - "windows_i686_gnu 0.48.5", - "windows_i686_msvc 0.48.5", - "windows_x86_64_gnu 0.48.5", - "windows_x86_64_gnullvm 0.48.5", - "windows_x86_64_msvc 0.48.5", -] - [[package]] name = "windows-targets" version = "0.52.6" @@ -7633,12 +7726,6 @@ version = "0.42.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "597a5118570b68bc08d8d59125332c54f1ba9d9adeedeef5b99b02ba2b0698f8" -[[package]] -name = "windows_aarch64_gnullvm" -version = "0.48.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2b38e32f0abccf9987a4e3079dfb67dcd799fb61361e53e2882c3cbaf0d905d8" - [[package]] name = "windows_aarch64_gnullvm" version = "0.52.6" @@ -7657,12 +7744,6 @@ version = "0.42.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e08e8864a60f06ef0d0ff4ba04124db8b0fb3be5776a5cd47641e942e58c4d43" -[[package]] -name = "windows_aarch64_msvc" -version = "0.48.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "dc35310971f3b2dbbf3f0690a219f40e2d9afcf64f9ab7cc1be722937c26b4bc" - [[package]] name = "windows_aarch64_msvc" version = "0.52.6" @@ -7681,12 +7762,6 @@ version = "0.42.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c61d927d8da41da96a81f029489353e68739737d3beca43145c8afec9a31a84f" -[[package]] -name = "windows_i686_gnu" -version = "0.48.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a75915e7def60c94dcef72200b9a8e58e5091744960da64ec734a6c6e9b3743e" - [[package]] name = "windows_i686_gnu" version = "0.52.6" @@ -7717,12 +7792,6 @@ version = "0.42.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "44d840b6ec649f480a41c8d80f9c65108b92d89345dd94027bfe06ac444d1060" -[[package]] -name = "windows_i686_msvc" -version = "0.48.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8f55c233f70c4b27f66c523580f78f1004e8b5a8b659e05a4eb49d4166cca406" - [[package]] name = "windows_i686_msvc" version = "0.52.6" @@ -7741,12 +7810,6 @@ version = "0.42.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8de912b8b8feb55c064867cf047dda097f92d51efad5b491dfb98f6bbb70cb36" -[[package]] -name = "windows_x86_64_gnu" -version = "0.48.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "53d40abd2583d23e4718fddf1ebec84dbff8381c07cae67ff7768bbf19c6718e" - [[package]] name = "windows_x86_64_gnu" version = "0.52.6" @@ -7765,12 +7828,6 @@ version = "0.42.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "26d41b46a36d453748aedef1486d5c7a85db22e56aff34643984ea85514e94a3" -[[package]] -name = "windows_x86_64_gnullvm" -version = "0.48.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0b7b52767868a23d5bab768e390dc5f5c55825b6d30b86c844ff2dc7414044cc" - [[package]] name = "windows_x86_64_gnullvm" version = "0.52.6" @@ -7789,12 +7846,6 @@ version = "0.42.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9aec5da331524158c6d1a4ac0ab1541149c0b9505fde06423b02f5ef0106b9f0" -[[package]] -name = "windows_x86_64_msvc" -version = "0.48.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ed94fce61571a4006852b7389a063ab983c02eb1bb37b47f8272ce92d06d9538" - [[package]] name = "windows_x86_64_msvc" version = "0.52.6" diff --git a/Cargo.toml b/Cargo.toml index 3038e72245e..da92195eb44 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -99,7 +99,7 @@ sqlx = { version = "0.9.0-alpha.1", features = [ "uuid", "sqlx-toml", ] } -rlt = "0.2.1" +rlt = "0.4.1" metrics-exporter-prometheus = { version = "0.18.1", features = [ "http-listener", ], default-features = false } diff --git a/tensorzero-core/tests/load/feedback/src/main.rs b/tensorzero-core/tests/load/feedback/src/main.rs index bce37d1a5b4..af74e5917a0 100644 --- a/tensorzero-core/tests/load/feedback/src/main.rs +++ b/tensorzero-core/tests/load/feedback/src/main.rs @@ -54,7 +54,7 @@ async fn main() -> Result<()> { rlt::cli::Collector::Tui => { tracing_subscriber::registry() .with(EnvFilter::from_default_env()) - .with(rlt::tui_tracing_subscriber_layer()) + .with(rlt::TuiTracingSubscriberLayer) .init(); } rlt::cli::Collector::Silent => { diff --git a/tensorzero-core/tests/load/rate-limit-load-test/src/main.rs b/tensorzero-core/tests/load/rate-limit-load-test/src/main.rs index 113f39d9294..d1a24f6f60d 100644 --- a/tensorzero-core/tests/load/rate-limit-load-test/src/main.rs +++ b/tensorzero-core/tests/load/rate-limit-load-test/src/main.rs @@ -63,7 +63,7 @@ async fn main() -> Result<()> { rlt::cli::Collector::Tui => { tracing_subscriber::registry() .with(EnvFilter::from_default_env()) - .with(rlt::tui_tracing_subscriber_layer()) + .with(rlt::TuiTracingSubscriberLayer) .init(); } rlt::cli::Collector::Silent => { From 2727d04039a33c495c3bc68a9cbe72d18202e64d Mon Sep 17 00:00:00 2001 From: Gabriel Bianconi <1275491+GabrielBianconi@users.noreply.github.com> Date: Tue, 27 Jan 2026 18:12:24 -0500 Subject: [PATCH 41/46] Use Docker in merge queue CI (#5875) * Use Docker in merge queue CI * Fix * Fix * Fix * Fix * Fix * Fix * Fix * Fix * Fix * Fix * Fix --- .../workflows/build-fixtures-container.yml | 55 ++++++ .github/workflows/general.yml | 43 +++-- .github/workflows/merge-queue.yml | 156 +++++++++++++++--- .../tests/e2e/docker-compose.replicated.yml | 1 + tensorzero-core/tests/e2e/docker-compose.yml | 2 + 5 files changed, 219 insertions(+), 38 deletions(-) create mode 100644 .github/workflows/build-fixtures-container.yml diff --git a/.github/workflows/build-fixtures-container.yml b/.github/workflows/build-fixtures-container.yml new file mode 100644 index 00000000000..29eaa07f89b --- /dev/null +++ b/.github/workflows/build-fixtures-container.yml @@ -0,0 +1,55 @@ +name: Build Fixtures Container + +on: + workflow_call: + secrets: + DOCKERHUB_LIMITED_TOKEN: + required: true + +jobs: + build-fixtures-container: + runs-on: ubuntu-latest + if: github.repository == 'tensorzero/tensorzero' + permissions: + # Permission to checkout the repository + contents: read + # Permission to fetch GitHub OIDC token authentication + id-token: write + + steps: + - name: Check out the repo + uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 + + - name: Login to DockerHub + uses: docker/login-action@74a5d142397b4f367a81961eba4e8cd7edddf772 + with: + username: tensorzero + password: ${{ secrets.DOCKERHUB_LIMITED_TOKEN }} + + # We allow the namespace builder setup to fail on Dependabot PRs and PRs from forks + # (where the oidc token is not available) + + - name: Install Namespace CLI + uses: namespacelabs/nscloud-setup@d1c625762f7c926a54bd39252efff0705fd11c64 + continue-on-error: ${{ github.event.pull_request.head.repo.full_name != github.repository || github.actor == 'dependabot[bot]' }} + + - name: Configure Namespace-powered Buildx + uses: namespacelabs/nscloud-setup-buildx-action@91c2e6537780e3b092cb8476406be99a8f91bd5e + with: + wait-for-builder: true + continue-on-error: ${{ github.event.pull_request.head.repo.full_name != github.repository || github.actor == 'dependabot[bot]' }} + + - name: Build `fixtures` container + run: | + docker buildx build -f ui/fixtures/Dockerfile . -t tensorzero/fixtures:sha-${{ github.sha }} -t nscr.io/igvf4asmf8kri/fixtures:sha-${{ github.sha }} + + - name: Push `fixtures` container to Docker Hub + run: docker push tensorzero/fixtures:sha-${{ github.sha }} + + - name: Login to Namespace registry + run: nsc docker login + continue-on-error: ${{ github.event.pull_request.head.repo.full_name != github.repository || github.actor == 'dependabot[bot]' }} + + - name: Push `fixtures` container to Namespace registry + run: docker push nscr.io/igvf4asmf8kri/fixtures:sha-${{ github.sha }} + continue-on-error: ${{ github.event.pull_request.head.repo.full_name != github.repository || github.actor == 'dependabot[bot]' }} diff --git a/.github/workflows/general.yml b/.github/workflows/general.yml index c3ec01c5a37..3ceacb028f5 100644 --- a/.github/workflows/general.yml +++ b/.github/workflows/general.yml @@ -704,7 +704,14 @@ jobs: contents: read # Permission to download artifacts and for rust-cache actions: read - needs: [build-gateway-container, build-mock-provider-api-container] + # Permission to fetch GitHub OIDC token authentication for Namespace login + id-token: write + needs: + [ + build-gateway-container, + build-mock-provider-api-container, + build-fixtures-container, + ] # We don't run many tests here, so use a normal runner with Github Actions caching # to avoid unnecessarily using Namespace credits (it should still always finish before @@ -774,9 +781,11 @@ jobs: run: | docker pull nscr.io/igvf4asmf8kri/gateway:sha-${{ github.sha }} docker pull nscr.io/igvf4asmf8kri/mock-provider-api:sha-${{ github.sha }} + docker pull nscr.io/igvf4asmf8kri/fixtures:sha-${{ github.sha }} # Retag the images to what we expect the names to be docker tag nscr.io/igvf4asmf8kri/gateway:sha-${{ github.sha }} tensorzero/gateway:sha-${{ github.sha }} docker tag nscr.io/igvf4asmf8kri/mock-provider-api:sha-${{ github.sha }} tensorzero/mock-provider-api:sha-${{ github.sha }} + docker tag nscr.io/igvf4asmf8kri/fixtures:sha-${{ github.sha }} tensorzero/fixtures:sha-${{ github.sha }} continue-on-error: ${{ github.event.pull_request.head.repo.full_name != github.repository || github.actor == 'dependabot[bot]' }} - name: Download ClickHouse fixtures @@ -795,21 +804,19 @@ jobs: # Note: In replicated mode, we use the HTTP port (8124) of the second replica so we can ensure this works echo "TENSORZERO_CLICKHOUSE_URL=http://chuser:chpassword@localhost:8124/tensorzero_e2e_tests" >> $GITHUB_ENV - - name: Set TENSORZERO_GATEWAY_TAG + - name: Set container image tags run: | echo "TENSORZERO_GATEWAY_TAG=sha-${{ github.sha }}" >> $GITHUB_ENV - - - name: Set TENSORZERO_MOCK_PROVIDER_API_TAG - run: | echo "TENSORZERO_MOCK_PROVIDER_API_TAG=sha-${{ github.sha }}" >> $GITHUB_ENV + echo "TENSORZERO_COMMIT_TAG=sha-${{ github.sha }}" >> $GITHUB_ENV - name: Launch dependency services with non-replicated ClickHouse container for E2E tests if: matrix.replicated == false - run: TENSORZERO_CLICKHOUSE_VERSION=${{ matrix.clickhouse_version.tag }} docker compose -f tensorzero-core/tests/e2e/docker-compose.yml up --wait + run: TENSORZERO_CLICKHOUSE_VERSION=${{ matrix.clickhouse_version.tag }} docker compose -f tensorzero-core/tests/e2e/docker-compose.yml up --no-build --wait - name: Launch replicated ClickHouse container for E2E tests if: matrix.replicated == true - run: TENSORZERO_CLICKHOUSE_VERSION=${{ matrix.clickhouse_version.tag }} docker compose -f tensorzero-core/tests/e2e/docker-compose.replicated.yml up --wait + run: TENSORZERO_CLICKHOUSE_VERSION=${{ matrix.clickhouse_version.tag }} docker compose -f tensorzero-core/tests/e2e/docker-compose.replicated.yml up --no-build --wait # Make an HTTP request to ClickHouse and check that the version matches '${{ matrix.clickhouse_version }}' - name: Check ClickHouse version @@ -921,7 +928,7 @@ jobs: # Permission to fetch GitHub OIDC token authentication for Namespace login id-token: write runs-on: ubuntu-latest - needs: [build-gateway-container] + needs: [build-gateway-container, build-fixtures-container] if: ${{ github.event_name == 'merge_group' || (github.event.pull_request.head.repo.full_name == github.repository) }} steps: @@ -960,15 +967,18 @@ jobs: run: nsc docker login continue-on-error: ${{ github.event.pull_request.head.repo.full_name != github.repository || github.actor == 'dependabot[bot]' }} - - name: Download gateway container image + - name: Download container images run: | docker pull nscr.io/igvf4asmf8kri/gateway:sha-${{ github.sha }} + docker pull nscr.io/igvf4asmf8kri/fixtures:sha-${{ github.sha }} docker tag nscr.io/igvf4asmf8kri/gateway:sha-${{ github.sha }} tensorzero/gateway:sha-${{ github.sha }} + docker tag nscr.io/igvf4asmf8kri/fixtures:sha-${{ github.sha }} tensorzero/fixtures:sha-${{ github.sha }} continue-on-error: ${{ github.event.pull_request.head.repo.full_name != github.repository || github.actor == 'dependabot[bot]' }} - - name: Set TENSORZERO_GATEWAY_TAG + - name: Set container image tags run: | echo "TENSORZERO_GATEWAY_TAG=sha-${{ github.sha }}" >> $GITHUB_ENV + echo "TENSORZERO_COMMIT_TAG=sha-${{ github.sha }}" >> $GITHUB_ENV - name: Set environment variables for Postgres run: | @@ -976,7 +986,7 @@ jobs: echo "TENSORZERO_POSTGRES_URL=postgres://postgres:postgres@localhost:5432/tensorzero-e2e-tests" >> $GITHUB_ENV - name: Launch Postgres - run: docker compose -f tensorzero-core/tests/e2e/docker-compose.yml up postgres gateway-postgres-migrations fixtures-postgres -d --wait + run: docker compose -f tensorzero-core/tests/e2e/docker-compose.yml up --no-build postgres gateway-postgres-migrations fixtures-postgres -d --wait - name: Test (Rust) run: cargo test-e2e --filterset 'test(db::) & test(postgres)' @@ -1226,6 +1236,16 @@ jobs: # Permission to fetch GitHub OIDC token authentication id-token: write + build-fixtures-container: + uses: ./.github/workflows/build-fixtures-container.yml + permissions: + # Permission to checkout the repository + contents: read + # Permission to fetch GitHub OIDC token authentication + id-token: write + secrets: + DOCKERHUB_LIMITED_TOKEN: ${{ secrets.DOCKERHUB_LIMITED_TOKEN }} + ui-tests: permissions: contents: read @@ -1341,6 +1361,7 @@ jobs: build-gateway-container, build-gateway-e2e-container, build-mock-provider-api-container, + build-fixtures-container, ] if: (github.repository == 'tensorzero/tensorzero' && github.event_name == 'merge_group') permissions: diff --git a/.github/workflows/merge-queue.yml b/.github/workflows/merge-queue.yml index 1b1f2550fb3..c5eade469a8 100644 --- a/.github/workflows/merge-queue.yml +++ b/.github/workflows/merge-queue.yml @@ -56,15 +56,53 @@ env: TENSORZERO_E2E_PROXY: http://localhost:3003 TENSORZERO_COMMIT_TAG: sha-${{ github.sha }} TENSORZERO_GATEWAY_TAG: sha-${{ github.sha }} + TENSORZERO_MOCK_PROVIDER_API_TAG: sha-${{ github.sha }} TENSORZERO_CI: 1 # Nextest filter expression for flaky tests to skip (e.g. "test(hyperbolic)|test(tgi)") CARGO_NEXTEST_FLAKY_TESTS: ${{ vars.CARGO_NEXTEST_FLAKY_TESTS }} jobs: + # Build containers when triggered directly (not from general.yml which already builds them) + # When called from general.yml, github.event_name is 'merge_group' so this is skipped + build-gateway-e2e-container: + if: github.repository == 'tensorzero/tensorzero' && (github.event_name == 'workflow_dispatch' || github.event_name == 'schedule') + uses: ./.github/workflows/build-gateway-e2e-container.yml + permissions: + contents: read + id-token: write + secrets: + DOCKERHUB_USERNAME: ${{ secrets.DOCKERHUB_USERNAME }} + DOCKERHUB_LIMITED_TOKEN: ${{ secrets.DOCKERHUB_LIMITED_TOKEN }} + + build-gateway-container: + if: github.repository == 'tensorzero/tensorzero' && (github.event_name == 'workflow_dispatch' || github.event_name == 'schedule') + uses: ./.github/workflows/build-gateway-container.yml + permissions: + contents: read + id-token: write + + build-mock-provider-api-container: + if: github.repository == 'tensorzero/tensorzero' && (github.event_name == 'workflow_dispatch' || github.event_name == 'schedule') + uses: ./.github/workflows/build-mock-provider-api-container.yml + permissions: + contents: read + id-token: write + + build-fixtures-container: + if: github.repository == 'tensorzero/tensorzero' && (github.event_name == 'workflow_dispatch' || github.event_name == 'schedule') + uses: ./.github/workflows/build-fixtures-container.yml + permissions: + contents: read + id-token: write + secrets: + DOCKERHUB_LIMITED_TOKEN: ${{ secrets.DOCKERHUB_LIMITED_TOKEN }} + live-tests: + needs: [build-gateway-e2e-container] + # Run even when build-gateway-e2e-container is skipped (merge_group case) + if: always() && !failure() && !cancelled() && github.repository == 'tensorzero/tensorzero' name: "live-tests (batch_writes: ${{ matrix.batch_writes }})" runs-on: ubuntu-latest - if: github.repository == 'tensorzero/tensorzero' permissions: # Permission to checkout the repository contents: read @@ -156,9 +194,17 @@ jobs: AWS_ACCESS_KEY_ID=$R2_ACCESS_KEY_ID AWS_SECRET_ACCESS_KEY=$R2_SECRET_ACCESS_KEY PROVIDER_PROXY_CACHE_BUCKET=provider-proxy-cache ./ci/upload-provider-proxy-cache.sh client-tests: + needs: + [ + build-gateway-e2e-container, + build-gateway-container, + build-mock-provider-api-container, + build-fixtures-container, + ] + # Run even when build jobs are skipped (merge_group case where general.yml already built them) + if: always() && !failure() && !cancelled() && github.repository == 'tensorzero/tensorzero' name: "client-tests (batch_writes: ${{ matrix.batch_writes }})" runs-on: ubuntu-latest - if: github.repository == 'tensorzero/tensorzero' permissions: # Permission to checkout the repository contents: read @@ -221,6 +267,18 @@ jobs: - name: Install Namespace CLI uses: namespacelabs/nscloud-setup@d1c625762f7c926a54bd39252efff0705fd11c64 + - name: Login to Namespace registry + run: nsc docker login + + - name: Pull pre-built containers from Namespace registry + run: | + docker pull nscr.io/igvf4asmf8kri/gateway:sha-${{ github.sha }} + docker tag nscr.io/igvf4asmf8kri/gateway:sha-${{ github.sha }} tensorzero/gateway:sha-${{ github.sha }} + docker pull nscr.io/igvf4asmf8kri/mock-provider-api:sha-${{ github.sha }} + docker tag nscr.io/igvf4asmf8kri/mock-provider-api:sha-${{ github.sha }} tensorzero/mock-provider-api:sha-${{ github.sha }} + docker pull nscr.io/igvf4asmf8kri/fixtures:sha-${{ github.sha }} + docker tag nscr.io/igvf4asmf8kri/fixtures:sha-${{ github.sha }} tensorzero/fixtures:sha-${{ github.sha }} + - name: Configure Namespace-powered Buildx uses: namespacelabs/nscloud-setup-buildx-action@91c2e6537780e3b092cb8476406be99a8f91bd5e with: @@ -306,7 +364,7 @@ jobs: - name: Launch dependency services for E2E tests run: | - docker compose -f tensorzero-core/tests/e2e/docker-compose.yml up --build -d --wait + docker compose -f tensorzero-core/tests/e2e/docker-compose.yml up --no-build -d --wait - name: Print ClickHouse container logs if: always() @@ -317,19 +375,70 @@ jobs: run: | ./ci/run-provider-proxy.sh ci - # TODO - get rid of this when the merge queue has a freshly-build gateway image available - - name: Manually run the latest postgres migrations - run: cargo run-e2e --run-postgres-migrations + - name: Pull gateway-e2e container + run: docker pull tensorzero/gateway-e2e:sha-${{ github.sha }} + + - name: Set up gateway environment file + run: | + # TensorZero config + echo "TENSORZERO_CLICKHOUSE_URL=${TENSORZERO_CLICKHOUSE_URL}" >> gateway.env + echo "TENSORZERO_CLICKHOUSE_BATCH_WRITES=${TENSORZERO_CLICKHOUSE_BATCH_WRITES}" >> gateway.env + echo "TENSORZERO_POSTGRES_URL=${TENSORZERO_POSTGRES_URL}" >> gateway.env + echo "TENSORZERO_VALKEY_URL=${TENSORZERO_VALKEY_URL}" >> gateway.env + echo "TENSORZERO_E2E_PROXY=${TENSORZERO_E2E_PROXY}" >> gateway.env + echo "OTEL_EXPORTER_OTLP_TRACES_ENDPOINT=${OTEL_EXPORTER_OTLP_TRACES_ENDPOINT}" >> gateway.env + echo "GCP_VERTEX_CREDENTIALS_PATH=/app/gcp_jwt_key.json" >> gateway.env + echo "GOOGLE_APPLICATION_CREDENTIALS=/app/gcp_jwt_key.json" >> gateway.env + # Provider API keys + echo "ANTHROPIC_API_KEY=${ANTHROPIC_API_KEY}" >> gateway.env + echo "AWS_ACCESS_KEY_ID=${AWS_ACCESS_KEY_ID}" >> gateway.env + echo "AWS_SECRET_ACCESS_KEY=${AWS_SECRET_ACCESS_KEY}" >> gateway.env + echo "AWS_REGION=${AWS_REGION}" >> gateway.env + echo "AZURE_API_KEY=${AZURE_API_KEY}" >> gateway.env + echo "AZURE_OPENAI_EASTUS2_API_KEY=${AZURE_OPENAI_EASTUS2_API_KEY}" >> gateway.env + echo "AZURE_OPENAI_DEPLOYMENT_ID=${AZURE_OPENAI_DEPLOYMENT_ID}" >> gateway.env + echo "AZURE_AI_FOUNDRY_API_KEY=${AZURE_AI_FOUNDRY_API_KEY}" >> gateway.env + echo "DEEPSEEK_API_KEY=${DEEPSEEK_API_KEY}" >> gateway.env + echo "FIREWORKS_API_KEY=${FIREWORKS_API_KEY}" >> gateway.env + echo "FIREWORKS_ACCOUNT_ID=${FIREWORKS_ACCOUNT_ID}" >> gateway.env + echo "GCP_STORAGE_ACCESS_KEY_ID=${GCP_STORAGE_ACCESS_KEY_ID}" >> gateway.env + echo "GCP_STORAGE_SECRET_ACCESS_KEY=${GCP_STORAGE_SECRET_ACCESS_KEY}" >> gateway.env + echo "GOOGLE_AI_STUDIO_API_KEY=${GOOGLE_AI_STUDIO_API_KEY}" >> gateway.env + echo "GROQ_API_KEY=${GROQ_API_KEY}" >> gateway.env + echo "HYPERBOLIC_API_KEY=${HYPERBOLIC_API_KEY}" >> gateway.env + echo "MISTRAL_API_KEY=${MISTRAL_API_KEY}" >> gateway.env + echo "MODAL_KEY=${MODAL_KEY}" >> gateway.env + echo "MODAL_SECRET=${MODAL_SECRET}" >> gateway.env + echo "OPENAI_API_KEY=${OPENAI_API_KEY}" >> gateway.env + echo "OPENROUTER_API_KEY=${OPENROUTER_API_KEY}" >> gateway.env + echo "SGLANG_API_KEY=${SGLANG_API_KEY}" >> gateway.env + echo "TGI_API_KEY=${TGI_API_KEY}" >> gateway.env + echo "TOGETHER_API_KEY=${TOGETHER_API_KEY}" >> gateway.env + echo "VLLM_API_KEY=${VLLM_API_KEY}" >> gateway.env + echo "VOYAGE_API_KEY=${VOYAGE_API_KEY}" >> gateway.env + echo "XAI_API_KEY=${XAI_API_KEY}" >> gateway.env + + - name: Run postgres migrations + run: | + docker run --rm --network host \ + --env-file gateway.env \ + tensorzero/gateway-e2e:sha-${{ github.sha }} \ + --run-postgres-migrations - name: Launch the gateway for E2E tests timeout-minutes: 2 run: | - cargo run-e2e > e2e_logs.txt 2>&1 & + docker run -d --name gateway-e2e --network host \ + --env-file gateway.env \ + -v ${{ github.workspace }}/gcp_jwt_key.json:/app/gcp_jwt_key.json:ro \ + -v ${{ github.workspace }}/tensorzero-core:/app/tensorzero-core:ro \ + -v ${{ github.workspace }}/ui/fixtures:/app/ui/fixtures:ro \ + tensorzero/gateway-e2e:sha-${{ github.sha }} \ + --config-file '/app/tensorzero-core/tests/e2e/config/tensorzero.*.toml' while ! curl -s -f http://localhost:3000/health >/dev/null 2>&1; do echo "Waiting for gateway to be healthy..." sleep 1 done - echo "GATEWAY_PID=$!" >> $GITHUB_ENV - name: Install Python for python async client tests run: uv python install 3.9 @@ -403,27 +512,24 @@ jobs: run: | uv run pytest - - name: Terminate the gateway and wait for it to exit + - name: Print e2e logs if: always() + run: docker logs gateway-e2e 2>&1 || echo "No gateway container logs available" + + - name: Check e2e logs for impossible error messages run: | - echo "Killing gateway with pid $GATEWAY_PID" - kill $GATEWAY_PID - # Wait for at most 30 seconds for the gateway to exit - for i in {1..30}; do - if ! kill -0 $GATEWAY_PID 2>/dev/null; then - echo "Gateway exited" - break - fi - sleep 1 - done - if kill -0 $GATEWAY_PID 2>/dev/null; then - echo "Gateway did not exit after 30 seconds!" + LOGS=$(docker logs gateway-e2e 2>&1) + if [ -z "$LOGS" ]; then + echo "ERROR: Gateway logs are empty" exit 1 fi + grep --invert-match -i "please file a bug report" <<< "$LOGS" - - name: Print e2e logs + - name: Terminate the gateway container if: always() - run: cat e2e_logs.txt + run: | + docker stop gateway-e2e || true + docker rm gateway-e2e || true - name: Print provider-proxy logs if: always() @@ -442,10 +548,6 @@ jobs: continue-on-error: true run: cat vllm_gpt_oss_modal_logs.txt - - name: Check e2e logs for impossible error messages - run: | - grep --invert-match -i "please file a bug report" <<< "$LOGS" - - name: Upload client-tests provider-proxy cache # Only upload the cache when we're running from a 'good' run # (running from the merge queue via 'workflow_call' from general.yml, or a cron job) diff --git a/tensorzero-core/tests/e2e/docker-compose.replicated.yml b/tensorzero-core/tests/e2e/docker-compose.replicated.yml index 6da90593c06..985955b224f 100644 --- a/tensorzero-core/tests/e2e/docker-compose.replicated.yml +++ b/tensorzero-core/tests/e2e/docker-compose.replicated.yml @@ -181,6 +181,7 @@ services: start_interval: 1s timeout: 1s fixtures: + image: tensorzero/fixtures:${TENSORZERO_COMMIT_TAG:-latest} build: dockerfile: ui/fixtures/Dockerfile context: ../../.. diff --git a/tensorzero-core/tests/e2e/docker-compose.yml b/tensorzero-core/tests/e2e/docker-compose.yml index 2afaa63018a..68bd239bacc 100644 --- a/tensorzero-core/tests/e2e/docker-compose.yml +++ b/tensorzero-core/tests/e2e/docker-compose.yml @@ -18,6 +18,7 @@ services: - "3030:3030" fixtures: + image: tensorzero/fixtures:${TENSORZERO_COMMIT_TAG:-latest} build: dockerfile: ui/fixtures/Dockerfile context: ../../.. @@ -53,6 +54,7 @@ services: # fixtures-postgres and fixtures have exactly the same build context, so the fixture is downloaded only once and shared. fixtures-postgres: + image: tensorzero/fixtures:${TENSORZERO_COMMIT_TAG:-latest} build: dockerfile: ui/fixtures/Dockerfile context: ../../.. From dfa25ae044a7168dfd28f84e8f58ff5ff3fe07c5 Mon Sep 17 00:00:00 2001 From: Shuyang Li Date: Tue, 27 Jan 2026 18:20:37 -0500 Subject: [PATCH 42/46] Stop setting up pnpm in a subshell (#5921) --- .github/workflows/slash-command-regen-fixtures.yml | 2 -- .github/workflows/ui-tests-e2e-model-inference-cache.yml | 2 -- .github/workflows/ui-tests-e2e.yml | 6 ------ 3 files changed, 10 deletions(-) diff --git a/.github/workflows/slash-command-regen-fixtures.yml b/.github/workflows/slash-command-regen-fixtures.yml index e906fbbb642..bb0c8e212a0 100644 --- a/.github/workflows/slash-command-regen-fixtures.yml +++ b/.github/workflows/slash-command-regen-fixtures.yml @@ -43,7 +43,6 @@ jobs: fi sleep $((10 * attempt)) done - shell: bash - name: Install `pnpm` dependencies run: pnpm install --frozen-lockfile - name: Setup Playwright @@ -58,7 +57,6 @@ jobs: fi sleep $((10 * attempt)) done - shell: bash - name: Regenerate fixtures id: e2e_tests run: | diff --git a/.github/workflows/ui-tests-e2e-model-inference-cache.yml b/.github/workflows/ui-tests-e2e-model-inference-cache.yml index 345c7187aaa..e7459d46d2f 100644 --- a/.github/workflows/ui-tests-e2e-model-inference-cache.yml +++ b/.github/workflows/ui-tests-e2e-model-inference-cache.yml @@ -77,7 +77,6 @@ jobs: fi sleep $((10 * attempt)) done - shell: bash - name: Install `pnpm` dependencies run: pnpm install --frozen-lockfile @@ -94,7 +93,6 @@ jobs: fi sleep $((10 * attempt)) done - shell: bash - name: Download container images uses: actions/download-artifact@018cc2cf5baa6db3ef3c5f8a56943fffe632ef53 diff --git a/.github/workflows/ui-tests-e2e.yml b/.github/workflows/ui-tests-e2e.yml index 8e5c6c0a5c5..2e59a63d2c8 100644 --- a/.github/workflows/ui-tests-e2e.yml +++ b/.github/workflows/ui-tests-e2e.yml @@ -105,7 +105,6 @@ jobs: fi sleep $((10 * attempt)) done - shell: bash - name: Install `pnpm` dependencies run: pnpm install --frozen-lockfile @@ -122,7 +121,6 @@ jobs: fi sleep $((10 * attempt)) done - shell: bash - name: Download container images uses: actions/download-artifact@018cc2cf5baa6db3ef3c5f8a56943fffe632ef53 @@ -231,7 +229,6 @@ jobs: fi sleep $((10 * attempt)) done - shell: bash - name: Install `pnpm` dependencies run: pnpm install --frozen-lockfile @@ -248,7 +245,6 @@ jobs: fi sleep $((10 * attempt)) done - shell: bash - name: Download container images uses: actions/download-artifact@018cc2cf5baa6db3ef3c5f8a56943fffe632ef53 @@ -330,7 +326,6 @@ jobs: fi sleep $((10 * attempt)) done - shell: bash - name: Install `pnpm` dependencies run: pnpm install --frozen-lockfile @@ -347,7 +342,6 @@ jobs: fi sleep $((10 * attempt)) done - shell: bash - name: Download container images uses: actions/download-artifact@018cc2cf5baa6db3ef3c5f8a56943fffe632ef53 From dbffb93908e26c919b44f2df0861c5d319bb395c Mon Sep 17 00:00:00 2001 From: Gabriel Bianconi <1275491+GabrielBianconi@users.noreply.github.com> Date: Tue, 27 Jan 2026 18:34:29 -0500 Subject: [PATCH 43/46] Update small fixtures and organize downloads (#5852) * Improve download-large-fixtures.py * Improve download-large-fixtures.py * Update small fixtures and organize downloads * Fix * Fix * Fix * Fix * Use R2 sync in CI * Fix * Add creds * Add creds * Fix * Fix * Fix * Fix * Fix * Fix * Fix * Fix * Fix --- .dockerignore | 4 +- .github/workflows/general.yml | 4 +- .../slash-command-regen-fixtures.yml | 7 +- .../ui-tests-e2e-model-inference-cache.yml | 14 +- .github/workflows/ui-tests-e2e.yml | 20 ++ .gitignore | 3 +- CONTRIBUTING.md | 14 +- ci/buildkite/clickhouse-tests.sh | 2 +- ci/buildkite/download-fixtures.sh | 14 +- tensorzero-core/tests/e2e/clickhouse.rs | 10 +- .../tests/e2e/docker-compose.clickhouse.yml | 4 +- tensorzero-core/tests/e2e/docker-compose.yml | 2 + ui/fixtures/README.md | 23 +- ui/fixtures/check-fixtures.sh | 30 +-- ui/fixtures/docker-compose.e2e.ci.yml | 1 + ui/fixtures/docker-compose.e2e.yml | 1 + ui/fixtures/docker-compose.unit.yml | 1 + ui/fixtures/docker-compose.yml | 1 + ui/fixtures/download-large-fixtures-http.py | 182 ++++++++++++++++ ui/fixtures/download-large-fixtures.py | 91 ++------ ui/fixtures/download-small-fixtures-http.py | 172 +++++++++++++++ ui/fixtures/download-small-fixtures.py | 200 ++++++++++++------ ui/fixtures/load_fixtures.sh | 62 +++--- ui/fixtures/load_fixtures_postgres.sh | 12 +- .../regenerate-model-inference-cache.sh | 4 +- ui/fixtures/upload-large-fixtures.sh | 2 +- ui/fixtures/upload-small-fixtures.sh | 6 +- 27 files changed, 660 insertions(+), 226 deletions(-) create mode 100644 ui/fixtures/download-large-fixtures-http.py create mode 100644 ui/fixtures/download-small-fixtures-http.py diff --git a/.dockerignore b/.dockerignore index a9f5d41766c..b0208c00575 100644 --- a/.dockerignore +++ b/.dockerignore @@ -9,8 +9,8 @@ target-rust-analyzer **/.env* **/.vscode **/.claude -ui/fixtures/*.jsonl -ui/fixtures/s3-fixtures +ui/fixtures/small-fixtures +ui/fixtures/large-fixtures /recipes /examples gcp_jwt_key.json diff --git a/.github/workflows/general.yml b/.github/workflows/general.yml index 3ceacb028f5..5c64d6ba3ac 100644 --- a/.github/workflows/general.yml +++ b/.github/workflows/general.yml @@ -770,7 +770,7 @@ jobs: uses: namespacelabs/nscloud-cache-action@446d8f390563cd54ca27e8de5bdb816f63c0b706 with: path: | - ./ui/fixtures/s3-fixtures/ + ./ui/fixtures/large-fixtures/ - name: Install uv uses: astral-sh/setup-uv@ed21f2f24f8dd64503750218de024bcf64c7250a @@ -1280,6 +1280,8 @@ jobs: S3_ACCESS_KEY_ID: ${{ secrets.AWS_ACCESS_KEY_ID }} S3_SECRET_ACCESS_KEY: ${{ secrets.AWS_SECRET_ACCESS_KEY }} GCP_JWT_KEY: ${{ secrets.GCP_JWT_KEY }} + R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }} + R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }} check-production-deployment-docker-compose: needs: [build-gateway-container] diff --git a/.github/workflows/slash-command-regen-fixtures.yml b/.github/workflows/slash-command-regen-fixtures.yml index bb0c8e212a0..c9745cb6c64 100644 --- a/.github/workflows/slash-command-regen-fixtures.yml +++ b/.github/workflows/slash-command-regen-fixtures.yml @@ -85,13 +85,14 @@ jobs: # Upload with timestamped name to R2 AWS_ACCESS_KEY_ID=$R2_ACCESS_KEY_ID AWS_SECRET_ACCESS_KEY=$R2_SECRET_ACCESS_KEY \ aws s3 --endpoint-url https://19918a216783f0ac9e052233569aef60.r2.cloudflarestorage.com/ \ - cp model_inference_cache_e2e.jsonl "s3://tensorzero-fixtures/${REMOTE_FILE}" + cp small-fixtures/model_inference_cache_e2e.jsonl "s3://tensorzero-fixtures/${REMOTE_FILE}" - # Update download-small-fixtures.py to point to new remote filename + # Update download scripts to point to new remote filename sed -i "s|\"model_inference_cache_e2e[^\"]*\.jsonl\": \"model_inference_cache_e2e\.jsonl\"|\"${REMOTE_FILE}\": \"model_inference_cache_e2e.jsonl\"|" download-small-fixtures.py + sed -i "s|\"model_inference_cache_e2e[^\"]*\.jsonl\": \"model_inference_cache_e2e\.jsonl\"|\"${REMOTE_FILE}\": \"model_inference_cache_e2e.jsonl\"|" download-small-fixtures-http.py echo "Uploaded to R2 as ${REMOTE_FILE}" - echo "Updated download-small-fixtures.py to reference ${REMOTE_FILE}" + echo "Updated download scripts to reference ${REMOTE_FILE}" - name: Push to PR run: | diff --git a/.github/workflows/ui-tests-e2e-model-inference-cache.yml b/.github/workflows/ui-tests-e2e-model-inference-cache.yml index e7459d46d2f..361ca5cfbf0 100644 --- a/.github/workflows/ui-tests-e2e-model-inference-cache.yml +++ b/.github/workflows/ui-tests-e2e-model-inference-cache.yml @@ -8,9 +8,9 @@ on: S3_SECRET_ACCESS_KEY: required: true R2_ACCESS_KEY_ID: - required: false + required: true R2_SECRET_ACCESS_KEY: - required: false + required: true OPENAI_API_KEY: required: false FIREWORKS_ACCOUNT_ID: @@ -148,6 +148,8 @@ jobs: if: inputs.regen_cache == true env: TENSORZERO_CI: 1 + R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }} + R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }} run: | echo "FIREWORKS_ACCOUNT_ID=${{ secrets.FIREWORKS_ACCOUNT_ID }}" >> fixtures/.env-gateway echo "FIREWORKS_API_KEY=${{ secrets.FIREWORKS_API_KEY }}" >> fixtures/.env-gateway @@ -157,7 +159,7 @@ jobs: echo "S3_SECRET_ACCESS_KEY=${{ secrets.S3_SECRET_ACCESS_KEY }}" >> fixtures/.env-gateway echo "TOGETHER_API_KEY=${{ secrets.TOGETHER_API_KEY }}" >> fixtures/.env-gateway # Save the old cache file for comparison - cp ./fixtures/model_inference_cache_e2e.jsonl ./fixtures/model_inference_cache_e2e.jsonl.old + cp ./fixtures/small-fixtures/model_inference_cache_e2e.jsonl ./fixtures/small-fixtures/model_inference_cache_e2e.jsonl.old # Run up to 3 times if it fails. We need to retry *outside* of playwright, so that we wipe the # model inference cache each time. Retrying an individual playwright test without wiping the database # will pollute the cache with duplicate entries. @@ -172,8 +174,8 @@ jobs: fi done # Check that the old and new cache files have the same number of rows - OLD_ROWS=$(wc -l < ./fixtures/model_inference_cache_e2e.jsonl.old) - NEW_ROWS=$(wc -l < ./fixtures/model_inference_cache_e2e.jsonl) + OLD_ROWS=$(wc -l < ./fixtures/small-fixtures/model_inference_cache_e2e.jsonl.old) + NEW_ROWS=$(wc -l < ./fixtures/small-fixtures/model_inference_cache_e2e.jsonl) echo "Old cache file has $OLD_ROWS rows, new cache file has $NEW_ROWS rows" if [ "$OLD_ROWS" != "$NEW_ROWS" ]; then echo "ERROR: Row count mismatch between old and new cache files" @@ -198,7 +200,7 @@ jobs: with: name: model-inference-cache path: | - ui/fixtures/model_inference_cache_e2e.jsonl + ui/fixtures/small-fixtures/model_inference_cache_e2e.jsonl - name: Generate API key if: matrix.auth_enabled == true diff --git a/.github/workflows/ui-tests-e2e.yml b/.github/workflows/ui-tests-e2e.yml index 2e59a63d2c8..410c2335575 100644 --- a/.github/workflows/ui-tests-e2e.yml +++ b/.github/workflows/ui-tests-e2e.yml @@ -22,6 +22,10 @@ on: required: true GCP_JWT_KEY: required: true + R2_ACCESS_KEY_ID: + required: true + R2_SECRET_ACCESS_KEY: + required: true inputs: is_merge_group: required: true @@ -61,6 +65,9 @@ jobs: - name: Start docker containers without external network access working-directory: ui + env: + R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }} + R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }} run: | # Environment variables shared by the gateway and ui containers echo "TENSORZERO_GATEWAY_URL=http://gateway:3000" >> fixtures/.env @@ -143,6 +150,11 @@ jobs: - name: Start Docker containers and apply fixtures working-directory: ui + env: + R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }} + R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }} + # Fall back to HTTP download for fork PRs (secrets unavailable) + TENSORZERO_DOWNLOAD_FIXTURES_WITHOUT_CREDENTIALS: ${{ github.event.pull_request.head.repo.full_name != github.repository && '1' || '' }} run: | # We set all of the environment variables for both the gateway and ui containers here # The 'ui-tests-e2e' job tests that the UI container starts without some of these variables set, @@ -272,6 +284,9 @@ jobs: echo "S3_SECRET_ACCESS_KEY=${{ secrets.S3_SECRET_ACCESS_KEY }}" >> fixtures/.env - name: Start docker containers + env: + R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }} + R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }} run: | export TENSORZERO_SKIP_LARGE_FIXTURES=1 # Explicitly start services without mock-provider-api (TENSORZERO_INTERNAL_MOCK_PROVIDER_API is not set) @@ -375,6 +390,9 @@ jobs: echo "TOGETHER_API_KEY=not_used" >> fixtures/.env-gateway - name: Start docker containers + env: + R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }} + R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }} run: | export TENSORZERO_SKIP_LARGE_FIXTURES=1 docker compose -f ui/fixtures/docker-compose.e2e.yml -f ui/fixtures/docker-compose.ui.yml up -d @@ -413,3 +431,5 @@ jobs: secrets: S3_ACCESS_KEY_ID: ${{ secrets.AWS_ACCESS_KEY_ID }} S3_SECRET_ACCESS_KEY: ${{ secrets.AWS_SECRET_ACCESS_KEY }} + R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }} + R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }} diff --git a/.gitignore b/.gitignore index 36e4a9706e9..b8969ed81f4 100644 --- a/.gitignore +++ b/.gitignore @@ -47,7 +47,8 @@ node_modules/ !./claude/commands/ !./claude/commands/** -ui/fixtures/s3-fixtures +ui/fixtures/large-fixtures +ui/fixtures/small-fixtures ci/provider-proxy-cache # autogenerated bindings in the wrong / default place diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index e0ca76aa106..93b51e2d1f4 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -197,8 +197,18 @@ To set it up, follow these steps from the repository's root directory: 2. Build the internal N-API client for TensorZero using `pnpm -r build`. If you have changed your Rust code, you may also have to run `pnpm build-bindings`. 3. Create a `ui/fixtures/.env` following the `ui/fixtures/.env.example`. 4. Create a `ui/.env` following the `ui/.env.example`, or set the environment variables from that file in your shell before running the dev script. -5. Launch the dependencies: `docker compose -f ui/fixtures/docker-compose.yml up --build --force-recreate`. - You can omit these last 2 flags to skip the build step, but they ensure you're using the latest gateway. +5. Launch the dependencies: + + ```bash + # For local development without R2 credentials (downloads via public HTTP): + TENSORZERO_DOWNLOAD_FIXTURES_WITHOUT_CREDENTIALS=1 docker compose -f ui/fixtures/docker-compose.yml up --build --force-recreate + + # With R2 credentials (faster S3 sync, used in CI): + docker compose -f ui/fixtures/docker-compose.yml up --build --force-recreate + ``` + + You can omit the `--build --force-recreate` flags to skip the build step, but they ensure you're using the latest gateway. + 6. Launch the development server: `pnpm ui:dev` Separately, you can run headless tests with `pnpm ui:test` and Playwright tests with `pnpm ui:test:e2e` (the latter will require a `pnpm exec playwright install`). diff --git a/ci/buildkite/clickhouse-tests.sh b/ci/buildkite/clickhouse-tests.sh index 582c7db8129..4cc0038ad73 100644 --- a/ci/buildkite/clickhouse-tests.sh +++ b/ci/buildkite/clickhouse-tests.sh @@ -54,7 +54,7 @@ docker compose -f tensorzero-core/tests/e2e/docker-compose.clickhouse.yml pull echo "For test purposes: let's make sure the fixtures dir is populated" pwd ls ui/fixtures -ls ui/fixtures/s3-fixtures +ls ui/fixtures/large-fixtures # ------------------------------------------------------------------------------ # Run ClickHouse tests container via Docker Compose and capture exit code diff --git a/ci/buildkite/download-fixtures.sh b/ci/buildkite/download-fixtures.sh index 70a6e3f5104..ea0c462dcfb 100644 --- a/ci/buildkite/download-fixtures.sh +++ b/ci/buildkite/download-fixtures.sh @@ -1,6 +1,18 @@ #!/bin/bash set -euxo pipefail +export R2_ACCESS_KEY_ID=$(buildkite-agent secret get R2_ACCESS_KEY_ID) +if [ -z "$R2_ACCESS_KEY_ID" ]; then + echo "Error: R2_ACCESS_KEY_ID is not set" + exit 1 +fi + +export R2_SECRET_ACCESS_KEY=$(buildkite-agent secret get R2_SECRET_ACCESS_KEY) +if [ -z "$R2_SECRET_ACCESS_KEY" ]; then + echo "Error: R2_SECRET_ACCESS_KEY is not set" + exit 1 +fi + # Install `uv` curl -LsSf https://astral.sh/uv/0.6.17/install.sh | sh source $HOME/.local/bin/env @@ -8,4 +20,4 @@ source $HOME/.local/bin/env uv run ./ui/fixtures/download-large-fixtures.py uv run ./ui/fixtures/download-small-fixtures.py # Zip the fixtures -tar -czvf fixtures.tar.gz ./ui/fixtures/s3-fixtures ./ui/fixtures/*.jsonl +tar -czvf fixtures.tar.gz ./ui/fixtures/large-fixtures ./ui/fixtures/small-fixtures/*.jsonl diff --git a/tensorzero-core/tests/e2e/clickhouse.rs b/tensorzero-core/tests/e2e/clickhouse.rs index 7a414dcbedd..29def1b1716 100644 --- a/tensorzero-core/tests/e2e/clickhouse.rs +++ b/tensorzero-core/tests/e2e/clickhouse.rs @@ -154,9 +154,9 @@ async fn count_table_rows(clickhouse: &ClickHouseConnectionInfo, table: &str) -> async fn insert_large_fixtures(clickhouse: &ClickHouseConnectionInfo) { // Insert data so that we test the migration re-creates the tables properly. - let s3_fixtures_path = std::env::var("TENSORZERO_S3_FIXTURES_PATH") - .unwrap_or_else(|_| format!("{MANIFEST_PATH}/../ui/fixtures/s3-fixtures")); - let s3_fixtures_path = &s3_fixtures_path; + let large_fixtures_path = std::env::var("TENSORZERO_LARGE_FIXTURES_PATH") + .unwrap_or_else(|_| format!("{MANIFEST_PATH}/../ui/fixtures/large-fixtures")); + let large_fixtures_path = &large_fixtures_path; let database_url = clickhouse.database_url(); let database = clickhouse.database(); @@ -233,7 +233,7 @@ async fn insert_large_fixtures(clickhouse: &ClickHouseConnectionInfo) { "run", "--add-host=host.docker.internal:host-gateway", "-v", - &format!("{s3_fixtures_path}:/s3-fixtures"), + &format!("{large_fixtures_path}:/large-fixtures"), "clickhouse:25.4", "clickhouse-client", "--host", @@ -247,7 +247,7 @@ async fn insert_large_fixtures(clickhouse: &ClickHouseConnectionInfo) { "--query", &format!( r" - INSERT INTO {table} FROM INFILE '/s3-fixtures/{file}' FORMAT Parquet + INSERT INTO {table} FROM INFILE '/large-fixtures/{file}' FORMAT Parquet " ), ]); diff --git a/tensorzero-core/tests/e2e/docker-compose.clickhouse.yml b/tensorzero-core/tests/e2e/docker-compose.clickhouse.yml index f853a3c8a25..f91d6a67afd 100644 --- a/tensorzero-core/tests/e2e/docker-compose.clickhouse.yml +++ b/tensorzero-core/tests/e2e/docker-compose.clickhouse.yml @@ -11,7 +11,7 @@ services: - ./clickhouse-configs/${TENSORZERO_CLICKHOUSE_VERSION:-lts}/users.xml:/etc/clickhouse-server/users.d/users.xml - clickhouse-data:/var/lib/clickhouse # For S3 fixtures using file - - ../../../ui/fixtures/s3-fixtures:/var/lib/clickhouse/user_files + - ../../../ui/fixtures/large-fixtures:/var/lib/clickhouse/user_files ports: - "8123:8123" # HTTP port - "9000:9000" # Native port @@ -98,7 +98,7 @@ services: TENSORZERO_CLICKHOUSE_URL: http://chuser:chpassword@clickhouse:8123/tensorzero_e2e_tests BUILDKITE_COMMIT: ${BUILDKITE_COMMIT:-} TENSORZERO_GATEWAY_URL: http://gateway:3000 - TENSORZERO_S3_FIXTURES_PATH: /app/ui/fixtures/s3-fixtures + TENSORZERO_LARGE_FIXTURES_PATH: /app/ui/fixtures/large-fixtures TENSORZERO_CI: 1 volumes: - ./config/:/app/tensorzero-core/tests/e2e/config/:ro diff --git a/tensorzero-core/tests/e2e/docker-compose.yml b/tensorzero-core/tests/e2e/docker-compose.yml index 68bd239bacc..637f307ce4c 100644 --- a/tensorzero-core/tests/e2e/docker-compose.yml +++ b/tensorzero-core/tests/e2e/docker-compose.yml @@ -30,6 +30,7 @@ services: - CLICKHOUSE_HOST=clickhouse - CLICKHOUSE_PASSWORD=chpassword - CLICKHOUSE_USER=chuser + - TENSORZERO_DOWNLOAD_FIXTURES_WITHOUT_CREDENTIALS - TENSORZERO_SKIP_LARGE_FIXTURES - R2_ACCESS_KEY_ID - R2_SECRET_ACCESS_KEY @@ -63,6 +64,7 @@ services: - ~/.aws:/root/.aws environment: - TENSORZERO_POSTGRES_URL=postgres://postgres:postgres@postgres:5432/${TENSORZERO_E2E_TESTS_DATABASE:-tensorzero-e2e-tests} + - TENSORZERO_DOWNLOAD_FIXTURES_WITHOUT_CREDENTIALS - R2_ACCESS_KEY_ID - R2_SECRET_ACCESS_KEY depends_on: diff --git a/ui/fixtures/README.md b/ui/fixtures/README.md index 82bd039727b..f65dda1ef36 100644 --- a/ui/fixtures/README.md +++ b/ui/fixtures/README.md @@ -4,20 +4,35 @@ Most of our fixtures are stored in this directory, with the exception of some la ## Pulling fixtures from R2 -Our large fixtures are stored in Cloudflare R2: +Our fixtures are stored in Cloudflare R2. There are two ways to download them: + +### Local development (no credentials needed) + +Set `TENSORZERO_DOWNLOAD_FIXTURES_WITHOUT_CREDENTIALS=1` when running docker compose: + +```bash +TENSORZERO_DOWNLOAD_FIXTURES_WITHOUT_CREDENTIALS=1 docker compose -f docker-compose.yml up +``` + +This downloads fixtures via public HTTP URLs. You can also run the HTTP scripts directly: + +- **Parquet files** (large tables): `uv run ./download-large-fixtures-http.py` +- **JSONL files** (small tables): `uv run ./download-small-fixtures-http.py` + +### CI / With R2 credentials (faster) + +Set `R2_ACCESS_KEY_ID` and `R2_SECRET_ACCESS_KEY` environment variables. This uses `aws s3 sync` for faster, more reliable downloads: - **Parquet files** (large tables): `uv run ./download-large-fixtures.py` - **JSONL files** (small tables): `uv run ./download-small-fixtures.py` -Downloads use public URLs by default. In CI, `R2_ACCESS_KEY_ID` and `R2_SECRET_ACCESS_KEY` env vars enable authenticated access. - ## Writing new fixtures Large fixtures should _not_ be committed to the repository. Instead: **For parquet files:** -1. Add the new fixtures to `./s3-fixtures` +1. Add the new fixtures to `./large-fixtures` 2. Run `./upload-large-fixtures.sh` 3. List the new fixture files in `download-large-fixtures.py` diff --git a/ui/fixtures/check-fixtures.sh b/ui/fixtures/check-fixtures.sh index 9a68c56676a..fc0444aff07 100755 --- a/ui/fixtures/check-fixtures.sh +++ b/ui/fixtures/check-fixtures.sh @@ -22,20 +22,20 @@ echo "===============================================" # Define tables and their corresponding files declare -A all_tables -all_tables["JsonInference"]="json_inference_examples.jsonl ./s3-fixtures/large_json_inference_v2.parquet" -all_tables["BooleanMetricFeedback"]="boolean_metric_feedback_examples.jsonl ./s3-fixtures/large_chat_boolean_feedback.parquet ./s3-fixtures/large_json_boolean_feedback.parquet" -all_tables["BooleanMetricFeedbackByTargetId"]="boolean_metric_feedback_examples.jsonl ./s3-fixtures/large_chat_boolean_feedback.parquet ./s3-fixtures/large_json_boolean_feedback.parquet" -all_tables["FloatMetricFeedback"]="float_metric_feedback_examples.jsonl jaro_winkler_similarity_feedback.jsonl ./s3-fixtures/large_chat_float_feedback.parquet ./s3-fixtures/large_json_float_feedback.parquet" -all_tables["FloatMetricFeedbackByTargetId"]="float_metric_feedback_examples.jsonl jaro_winkler_similarity_feedback.jsonl ./s3-fixtures/large_chat_float_feedback.parquet ./s3-fixtures/large_json_float_feedback.parquet" -all_tables["CommentFeedback"]="comment_feedback_examples.jsonl ./s3-fixtures/large_chat_comment_feedback.parquet ./s3-fixtures/large_json_comment_feedback.parquet" -all_tables["DemonstrationFeedback"]="demonstration_feedback_examples.jsonl ./s3-fixtures/large_chat_demonstration_feedback.parquet ./s3-fixtures/large_json_demonstration_feedback.parquet" -all_tables["ChatInference"]="chat_inference_examples.jsonl ./s3-fixtures/large_chat_inference_v2.parquet" -all_tables["ModelInference"]="model_inference_examples.jsonl ./s3-fixtures/large_chat_model_inference_v2.parquet ./s3-fixtures/large_json_model_inference_v2.parquet" -all_tables["ChatInferenceDatapoint FINAL"]="chat_inference_datapoint_examples.jsonl" -all_tables["JsonInferenceDatapoint FINAL"]="json_inference_datapoint_examples.jsonl" -all_tables["ModelInferenceCache"]="model_inference_cache_e2e.jsonl" -all_tables["DynamicEvaluationRun"]="dynamic_evaluation_run_examples.jsonl" -all_tables["DynamicEvaluationRunEpisode"]="dynamic_evaluation_run_episode_examples.jsonl" +all_tables["JsonInference"]="./small-fixtures/json_inference_examples.jsonl ./large-fixtures/large_json_inference_v2.parquet" +all_tables["BooleanMetricFeedback"]="./small-fixtures/boolean_metric_feedback_examples.jsonl ./large-fixtures/large_chat_boolean_feedback.parquet ./large-fixtures/large_json_boolean_feedback.parquet" +all_tables["BooleanMetricFeedbackByTargetId"]="./small-fixtures/boolean_metric_feedback_examples.jsonl ./large-fixtures/large_chat_boolean_feedback.parquet ./large-fixtures/large_json_boolean_feedback.parquet" +all_tables["FloatMetricFeedback"]="./small-fixtures/float_metric_feedback_examples.jsonl ./small-fixtures/jaro_winkler_similarity_feedback.jsonl ./large-fixtures/large_chat_float_feedback.parquet ./large-fixtures/large_json_float_feedback.parquet" +all_tables["FloatMetricFeedbackByTargetId"]="./small-fixtures/float_metric_feedback_examples.jsonl ./small-fixtures/jaro_winkler_similarity_feedback.jsonl ./large-fixtures/large_chat_float_feedback.parquet ./large-fixtures/large_json_float_feedback.parquet" +all_tables["CommentFeedback"]="./small-fixtures/comment_feedback_examples.jsonl ./large-fixtures/large_chat_comment_feedback.parquet ./large-fixtures/large_json_comment_feedback.parquet" +all_tables["DemonstrationFeedback"]="./small-fixtures/demonstration_feedback_examples.jsonl ./large-fixtures/large_chat_demonstration_feedback.parquet ./large-fixtures/large_json_demonstration_feedback.parquet" +all_tables["ChatInference"]="./small-fixtures/chat_inference_examples.jsonl ./large-fixtures/large_chat_inference_v2.parquet" +all_tables["ModelInference"]="./small-fixtures/model_inference_examples.jsonl ./large-fixtures/large_chat_model_inference_v2.parquet ./large-fixtures/large_json_model_inference_v2.parquet" +all_tables["ChatInferenceDatapoint FINAL"]="./small-fixtures/chat_inference_datapoint_examples.jsonl" +all_tables["JsonInferenceDatapoint FINAL"]="./small-fixtures/json_inference_datapoint_examples.jsonl" +all_tables["ModelInferenceCache"]="./small-fixtures/model_inference_cache_e2e.jsonl" +all_tables["DynamicEvaluationRun"]="./small-fixtures/dynamic_evaluation_run_examples.jsonl" +all_tables["DynamicEvaluationRunEpisode"]="./small-fixtures/dynamic_evaluation_run_episode_examples.jsonl" # Track if there's any mismatch mismatch=0 @@ -75,7 +75,7 @@ for table in "${!all_tables[@]}"; do echo " - $file: $file_count rows" fi else - echo " - WARNING: File $file not found in current directory or s3-fixtures/" + echo " - WARNING: File $file not found" mismatch=1 fi total_file_count=$(($total_file_count + $file_count)) diff --git a/ui/fixtures/docker-compose.e2e.ci.yml b/ui/fixtures/docker-compose.e2e.ci.yml index 4d3edfb8393..c855a7f9347 100644 --- a/ui/fixtures/docker-compose.e2e.ci.yml +++ b/ui/fixtures/docker-compose.e2e.ci.yml @@ -82,6 +82,7 @@ services: - CLICKHOUSE_HOST=clickhouse - CLICKHOUSE_PASSWORD=chpassword - CLICKHOUSE_USER=chuser + - TENSORZERO_DOWNLOAD_FIXTURES_WITHOUT_CREDENTIALS - TENSORZERO_SKIP_LARGE_FIXTURES=1 - R2_ACCESS_KEY_ID - R2_SECRET_ACCESS_KEY diff --git a/ui/fixtures/docker-compose.e2e.yml b/ui/fixtures/docker-compose.e2e.yml index 9c129e532ec..ba117dcef7b 100644 --- a/ui/fixtures/docker-compose.e2e.yml +++ b/ui/fixtures/docker-compose.e2e.yml @@ -83,6 +83,7 @@ services: - CLICKHOUSE_HOST=clickhouse - CLICKHOUSE_PASSWORD=chpassword - CLICKHOUSE_USER=chuser + - TENSORZERO_DOWNLOAD_FIXTURES_WITHOUT_CREDENTIALS - TENSORZERO_SKIP_LARGE_FIXTURES - R2_ACCESS_KEY_ID - R2_SECRET_ACCESS_KEY diff --git a/ui/fixtures/docker-compose.unit.yml b/ui/fixtures/docker-compose.unit.yml index ec4b053008a..0268f5609b0 100644 --- a/ui/fixtures/docker-compose.unit.yml +++ b/ui/fixtures/docker-compose.unit.yml @@ -81,6 +81,7 @@ services: - CLICKHOUSE_HOST=clickhouse - CLICKHOUSE_PASSWORD=chpassword - CLICKHOUSE_USER=chuser + - TENSORZERO_DOWNLOAD_FIXTURES_WITHOUT_CREDENTIALS - R2_ACCESS_KEY_ID - R2_SECRET_ACCESS_KEY depends_on: diff --git a/ui/fixtures/docker-compose.yml b/ui/fixtures/docker-compose.yml index a29323fee05..aa1d5280388 100644 --- a/ui/fixtures/docker-compose.yml +++ b/ui/fixtures/docker-compose.yml @@ -71,6 +71,7 @@ services: - CLICKHOUSE_HOST=clickhouse - CLICKHOUSE_PASSWORD=chpassword - CLICKHOUSE_USER=chuser + - TENSORZERO_DOWNLOAD_FIXTURES_WITHOUT_CREDENTIALS - R2_ACCESS_KEY_ID - R2_SECRET_ACCESS_KEY depends_on: diff --git a/ui/fixtures/download-large-fixtures-http.py b/ui/fixtures/download-large-fixtures-http.py new file mode 100644 index 00000000000..21f6a8e68de --- /dev/null +++ b/ui/fixtures/download-large-fixtures-http.py @@ -0,0 +1,182 @@ +# /// script +# dependencies = [ +# "requests", +# "parquet-tools", +# ] +# /// +# HTTP-only version for local development without R2 credentials. +# For CI with R2 credentials, use download-large-fixtures.py instead. + +import concurrent.futures +import hashlib +import os +import subprocess +import time +from pathlib import Path + +import requests + +# cd to directory of this file +os.chdir(os.path.dirname(os.path.abspath(__file__))) + + +# ============================================================================= +# Constants +# ============================================================================= + +PART_SIZE = 8388608 +FIXTURES = [ + "large_chat_inference_v2.parquet", + "large_chat_model_inference_v2.parquet", + "large_json_inference_v2.parquet", + "large_json_model_inference_v2.parquet", + "large_chat_boolean_feedback.parquet", + "large_chat_float_feedback.parquet", + "large_chat_comment_feedback.parquet", + "large_chat_demonstration_feedback.parquet", + "large_json_boolean_feedback.parquet", + "large_json_float_feedback.parquet", + "large_json_comment_feedback.parquet", + "large_json_demonstration_feedback.parquet", +] +R2_PUBLIC_BUCKET_URL = "https://pub-147e9850a60643208c411e70b636e956.r2.dev" +LARGE_FIXTURES_DIR = Path("./large-fixtures") + + +# ============================================================================= +# Utilities +# ============================================================================= + + +def calculate_etag(file_path): + """Calculate S3/R2 style ETag for a file.""" + file_size = os.path.getsize(file_path) + num_parts = (file_size + PART_SIZE - 1) // PART_SIZE + + if num_parts == 1: + # Single part upload - just MD5 of the file + with open(file_path, "rb") as f: + return hashlib.md5(f.read()).hexdigest() + else: + # Multipart upload - concatenate MD5s of each part + md5s = [] + with open(file_path, "rb") as f: + while True: + chunk = f.read(PART_SIZE) + if not chunk: + break + md5s.append(hashlib.md5(chunk).digest()) + + # Calculate MD5 of concatenated part MD5s + combined_md5 = hashlib.md5(b"".join(md5s)).hexdigest() + return f"{combined_md5}-{num_parts}" + + +def get_remote_etag(filename, retries: int = 3): + """Get ETag from R2 bucket via public URL with retry logic.""" + last_error = None + for attempt in range(retries): + try: + response = requests.head(f"{R2_PUBLIC_BUCKET_URL}/{filename}", timeout=30) + response.raise_for_status() + etag = response.headers.get("ETag") + if not etag: + raise Exception(f"Missing ETag header for {filename}") + return etag.strip('"') + except Exception as exc: + last_error = exc + if attempt < retries - 1: + sleep_time = 3**attempt + print( + f"Error fetching ETag for {filename} (attempt {attempt + 1} of {retries}): {exc}", + flush=True, + ) + print(f"Retrying in {sleep_time} seconds...", flush=True) + time.sleep(sleep_time) + + raise Exception(f"Failed to fetch ETag for {filename} after {retries} attempts") from last_error + + +# ============================================================================= +# HTTP download +# ============================================================================= + + +def download_file_http(filename, remote_etag): + """Download a single file from R2 via public HTTP URL.""" + RETRIES = 3 + for i in range(RETRIES): + try: + url = f"{R2_PUBLIC_BUCKET_URL}/{filename}" + response = requests.get(url, stream=True, timeout=300) + response.raise_for_status() + + local_file = LARGE_FIXTURES_DIR / filename + + with open(local_file, "wb") as f: + for chunk in response.iter_content(chunk_size=8192): + f.write(chunk) + + local_etag = calculate_etag(local_file) + if local_etag != remote_etag: + raise Exception(f"ETag mismatch after downloading: {local_etag} != {remote_etag}") + return + except Exception as e: + print( + f"Error downloading `{filename}` (attempt {i + 1} of {RETRIES}): {e}", + flush=True, + ) + time.sleep(1) + raise Exception(f"Failed to download `{filename}` after {RETRIES} attempts") + + +def download_fixtures_http(): + """Download all fixtures via public HTTP.""" + + def process_fixture(fixture): + local_file = LARGE_FIXTURES_DIR / fixture + remote_etag = get_remote_etag(fixture) + + if not local_file.exists(): + print(f"Downloading {fixture} (file doesn't exist locally)", flush=True) + download_file_http(fixture, remote_etag) + return + + local_etag = calculate_etag(local_file) + + if local_etag != remote_etag: + print(f"Downloading {fixture} (ETag mismatch)", flush=True) + print(f"Local ETag: {local_etag}", flush=True) + print(f"Remote ETag: {remote_etag}", flush=True) + download_file_http(fixture, remote_etag) + else: + print(f"Skipping {fixture} (up to date)", flush=True) + + # Use ThreadPoolExecutor to download files in parallel + with concurrent.futures.ThreadPoolExecutor() as executor: + # Loop over the results to propagate exceptions + for result in executor.map(process_fixture, FIXTURES): + assert result is None + + +# ============================================================================= +# Main +# ============================================================================= + + +def main(): + LARGE_FIXTURES_DIR.mkdir(exist_ok=True) + print("Downloading fixtures via public HTTP...", flush=True) + download_fixtures_http() + + for fixture in FIXTURES: + print(f"Fixture {fixture}:", flush=True) + subprocess.run( + ["parquet-tools", "inspect", LARGE_FIXTURES_DIR / fixture], + check=True, + stderr=subprocess.STDOUT, + ) + + +if __name__ == "__main__": + main() diff --git a/ui/fixtures/download-large-fixtures.py b/ui/fixtures/download-large-fixtures.py index f62edd25d44..05c42b65980 100644 --- a/ui/fixtures/download-large-fixtures.py +++ b/ui/fixtures/download-large-fixtures.py @@ -4,8 +4,8 @@ # "parquet-tools", # ] # /// +# For local development without R2 credentials, use download-large-fixtures-http.py instead. -import concurrent.futures import hashlib import os import subprocess @@ -39,7 +39,7 @@ ] R2_PUBLIC_BUCKET_URL = "https://pub-147e9850a60643208c411e70b636e956.r2.dev" R2_S3_ENDPOINT_URL = "https://19918a216783f0ac9e052233569aef60.r2.cloudflarestorage.com/" -S3_FIXTURES_DIR = Path("./s3-fixtures") +LARGE_FIXTURES_DIR = Path("./large-fixtures") # ============================================================================= @@ -100,7 +100,7 @@ def verify_etags(): """Verify ETags of all downloaded fixtures match remote.""" mismatches = [] for fixture in FIXTURES: - local_file = S3_FIXTURES_DIR / fixture + local_file = LARGE_FIXTURES_DIR / fixture if not local_file.exists(): raise Exception(f"Fixture {fixture} not found after sync") @@ -137,7 +137,7 @@ def sync_fixtures_from_r2(retries: int = 3) -> None: "180", "sync", "s3://tensorzero-fixtures/", - str(S3_FIXTURES_DIR), + str(LARGE_FIXTURES_DIR), # Only download the files in `FIXTURES` "--exclude", "*", @@ -145,7 +145,7 @@ def sync_fixtures_from_r2(retries: int = 3) -> None: ] env = { - **os.environ, # Preserve PATH and other environment variables + "PATH": os.environ.get("PATH", ""), "AWS_ACCESS_KEY_ID": os.environ["R2_ACCESS_KEY_ID"], "AWS_SECRET_ACCESS_KEY": os.environ["R2_SECRET_ACCESS_KEY"], "AWS_MAX_ATTEMPTS": "15", @@ -183,90 +183,27 @@ def sync_fixtures_from_r2(retries: int = 3) -> None: raise Exception(f"Fixture sync failed after {retries} attempts") from last_error -# ============================================================================= -# Public HTTP fallback (used locally without R2 credentials) -# ============================================================================= - - -def download_file_http(filename, remote_etag): - """Download a single file from R2 via public HTTP URL.""" - RETRIES = 3 - for i in range(RETRIES): - try: - url = f"{R2_PUBLIC_BUCKET_URL}/{filename}" - response = requests.get(url, stream=True, timeout=300) - response.raise_for_status() - - local_file = S3_FIXTURES_DIR / filename - - with open(local_file, "wb") as f: - for chunk in response.iter_content(chunk_size=8192): - f.write(chunk) - - local_etag = calculate_etag(local_file) - if local_etag != remote_etag: - raise Exception(f"ETag mismatch after downloading: {local_etag} != {remote_etag}") - return - except Exception as e: - print( - f"Error downloading `{filename}` (attempt {i + 1} of {RETRIES}): {e}", - flush=True, - ) - time.sleep(1) - raise Exception(f"Failed to download `{filename}` after {RETRIES} attempts") - - -def download_fixtures_http(): - """Download all fixtures via public HTTP (fallback when no R2 credentials).""" - - def process_fixture(fixture): - local_file = S3_FIXTURES_DIR / fixture - remote_etag = get_remote_etag(fixture) - - if not local_file.exists(): - print(f"Downloading {fixture} (file doesn't exist locally)", flush=True) - download_file_http(fixture, remote_etag) - return - - local_etag = calculate_etag(local_file) - - if local_etag != remote_etag: - print(f"Downloading {fixture} (ETag mismatch)", flush=True) - print(f"Local ETag: {local_etag}", flush=True) - print(f"Remote ETag: {remote_etag}", flush=True) - download_file_http(fixture, remote_etag) - else: - print(f"Skipping {fixture} (up to date)", flush=True) - - # Use ThreadPoolExecutor to download files in parallel - with concurrent.futures.ThreadPoolExecutor() as executor: - # Loop over the results to propagate exceptions - for result in executor.map(process_fixture, FIXTURES): - assert result is None - - # ============================================================================= # Main # ============================================================================= def main(): - S3_FIXTURES_DIR.mkdir(exist_ok=True) + LARGE_FIXTURES_DIR.mkdir(exist_ok=True) - if os.environ.get("R2_ACCESS_KEY_ID") and os.environ.get("R2_SECRET_ACCESS_KEY"): - print("R2 credentials found, downloading fixtures using `aws s3 sync`", flush=True) - sync_fixtures_from_r2() - else: - print( - "WARNING: `R2_ACCESS_KEY_ID` or `R2_SECRET_ACCESS_KEY` not set. Falling back to slow public HTTP downloads.", - flush=True, + if not os.environ.get("R2_ACCESS_KEY_ID") or not os.environ.get("R2_SECRET_ACCESS_KEY"): + raise Exception( + "R2_ACCESS_KEY_ID and R2_SECRET_ACCESS_KEY must be set. " + "For local development without R2 credentials, use download-large-fixtures-http.py instead." ) - download_fixtures_http() + + print("R2 credentials found, downloading fixtures using `aws s3 sync`", flush=True) + sync_fixtures_from_r2() for fixture in FIXTURES: print(f"Fixture {fixture}:", flush=True) subprocess.run( - ["parquet-tools", "inspect", S3_FIXTURES_DIR / fixture], + ["parquet-tools", "inspect", LARGE_FIXTURES_DIR / fixture], check=True, stderr=subprocess.STDOUT, ) diff --git a/ui/fixtures/download-small-fixtures-http.py b/ui/fixtures/download-small-fixtures-http.py new file mode 100644 index 00000000000..5b6c2f79e94 --- /dev/null +++ b/ui/fixtures/download-small-fixtures-http.py @@ -0,0 +1,172 @@ +# /// script +# dependencies = [ +# "requests", +# ] +# /// +# HTTP-only version for local development without R2 credentials. +# For CI with R2 credentials, use download-small-fixtures.py instead. + +import concurrent.futures +import hashlib +import os +import time +from pathlib import Path + +import requests + +# cd to directory of this file +os.chdir(os.path.dirname(os.path.abspath(__file__))) + + +# ============================================================================= +# Constants +# ============================================================================= + +PART_SIZE = 8388608 + +# Map of remote filename (in R2) -> local filename +# When a file needs updating, add version suffix to remote name (e.g., "file_v2.jsonl") +# and keep the local name unchanged (e.g., "file.jsonl") +FIXTURES = { + "model_inference_examples.jsonl": "model_inference_examples.jsonl", + "chat_inference_examples_20260123.jsonl": "chat_inference_examples.jsonl", + "json_inference_examples.jsonl": "json_inference_examples.jsonl", + "boolean_metric_feedback_examples.jsonl": "boolean_metric_feedback_examples.jsonl", + "float_metric_feedback_examples.jsonl": "float_metric_feedback_examples.jsonl", + "demonstration_feedback_examples.jsonl": "demonstration_feedback_examples.jsonl", + "model_inference_cache_e2e_20260122_183412.jsonl": "model_inference_cache_e2e.jsonl", + "json_inference_datapoint_examples.jsonl": "json_inference_datapoint_examples.jsonl", + "chat_inference_datapoint_examples.jsonl": "chat_inference_datapoint_examples.jsonl", + "dynamic_evaluation_run_episode_examples.jsonl": "dynamic_evaluation_run_episode_examples.jsonl", + "jaro_winkler_similarity_feedback.jsonl": "jaro_winkler_similarity_feedback.jsonl", + "comment_feedback_examples.jsonl": "comment_feedback_examples.jsonl", + "dynamic_evaluation_run_examples.jsonl": "dynamic_evaluation_run_examples.jsonl", +} + +R2_PUBLIC_BUCKET_URL = "https://pub-147e9850a60643208c411e70b636e956.r2.dev" +FIXTURES_DIR = Path("./small-fixtures") + + +# ============================================================================= +# Utilities +# ============================================================================= + + +def calculate_etag(file_path): + """Calculate S3/R2 style ETag for a file.""" + file_size = os.path.getsize(file_path) + num_parts = (file_size + PART_SIZE - 1) // PART_SIZE + + if num_parts == 1: + # Single part upload - just MD5 of the file + with open(file_path, "rb") as f: + return hashlib.md5(f.read()).hexdigest() + else: + # Multipart upload - concatenate MD5s of each part + md5s = [] + with open(file_path, "rb") as f: + while True: + chunk = f.read(PART_SIZE) + if not chunk: + break + md5s.append(hashlib.md5(chunk).digest()) + + # Calculate MD5 of concatenated part MD5s + combined_md5 = hashlib.md5(b"".join(md5s)).hexdigest() + return f"{combined_md5}-{num_parts}" + + +def get_remote_etag(remote_filename, retries=3): + """Get ETag from R2 bucket via public URL.""" + for i in range(retries): + try: + response = requests.head(f"{R2_PUBLIC_BUCKET_URL}/{remote_filename}", timeout=30) + response.raise_for_status() + return response.headers.get("ETag", "").strip('"') + except Exception as e: + if i < retries - 1: + print( + f"Error getting ETag for `{remote_filename}` (attempt {i + 1} of {retries}): {e}", + flush=True, + ) + time.sleep(1) + else: + raise + + +# ============================================================================= +# HTTP download +# ============================================================================= + + +def download_file_http(remote_filename, local_filename, remote_etag): + """Download a single file from R2 via public HTTP URL.""" + RETRIES = 3 + for i in range(RETRIES): + try: + url = f"{R2_PUBLIC_BUCKET_URL}/{remote_filename}" + response = requests.get(url, stream=True, timeout=300) + response.raise_for_status() + + local_file = FIXTURES_DIR / local_filename + + with open(local_file, "wb") as f: + for chunk in response.iter_content(chunk_size=8192): + f.write(chunk) + + local_etag = calculate_etag(local_file) + if local_etag != remote_etag: + raise Exception(f"ETag mismatch after downloading: {local_etag} != {remote_etag}") + return + except Exception as e: + print( + f"Error downloading `{remote_filename}` (attempt {i + 1} of {RETRIES}): {e}", + flush=True, + ) + time.sleep(1) + raise Exception(f"Failed to download `{remote_filename}` after {RETRIES} attempts") + + +def download_fixtures_http(): + """Download all fixtures via public HTTP.""" + + def process_fixture(item): + remote_filename, local_filename = item + local_file = FIXTURES_DIR / local_filename + remote_etag = get_remote_etag(remote_filename) + + if not local_file.exists(): + print(f"Downloading {remote_filename} (file doesn't exist locally)", flush=True) + download_file_http(remote_filename, local_filename, remote_etag) + return + + local_etag = calculate_etag(local_file) + + if local_etag != remote_etag: + print(f"Downloading {remote_filename} (ETag mismatch)", flush=True) + print(f"Local ETag: {local_etag}", flush=True) + print(f"Remote ETag: {remote_etag}", flush=True) + download_file_http(remote_filename, local_filename, remote_etag) + else: + print(f"Skipping {remote_filename} (up to date)", flush=True) + + # Use ThreadPoolExecutor to download files in parallel + with concurrent.futures.ThreadPoolExecutor() as executor: + # Loop over the results to propagate exceptions + for result in executor.map(process_fixture, FIXTURES.items()): + assert result is None + + +# ============================================================================= +# Main +# ============================================================================= + + +def main(): + FIXTURES_DIR.mkdir(exist_ok=True) + print("Downloading fixtures via public HTTP...", flush=True) + download_fixtures_http() + + +if __name__ == "__main__": + main() diff --git a/ui/fixtures/download-small-fixtures.py b/ui/fixtures/download-small-fixtures.py index a6a42e87608..4c5bc95805b 100644 --- a/ui/fixtures/download-small-fixtures.py +++ b/ui/fixtures/download-small-fixtures.py @@ -3,8 +3,8 @@ # "requests", # ] # /// +# For local development without R2 credentials, use download-small-fixtures-http.py instead. -import concurrent.futures import hashlib import os import subprocess @@ -16,7 +16,11 @@ # cd to directory of this file os.chdir(os.path.dirname(os.path.abspath(__file__))) + +# ============================================================================= # Constants +# ============================================================================= + PART_SIZE = 8388608 # Map of remote filename (in R2) -> local filename @@ -38,9 +42,14 @@ "dynamic_evaluation_run_examples.jsonl": "dynamic_evaluation_run_examples.jsonl", } -R2_BUCKET = "https://pub-147e9850a60643208c411e70b636e956.r2.dev" -# Download directly to the fixtures directory (not a subdirectory) -FIXTURES_DIR = Path(".") +R2_PUBLIC_BUCKET_URL = "https://pub-147e9850a60643208c411e70b636e956.r2.dev" +R2_S3_ENDPOINT_URL = "https://19918a216783f0ac9e052233569aef60.r2.cloudflarestorage.com/" +FIXTURES_DIR = Path("./small-fixtures") + + +# ============================================================================= +# Shared utilities +# ============================================================================= def calculate_etag(file_path): @@ -67,85 +76,138 @@ def calculate_etag(file_path): return f"{combined_md5}-{num_parts}" -def get_remote_etag(remote_filename): - """Get ETag from R2 bucket.""" - response = requests.head(f"{R2_BUCKET}/{remote_filename}") - return response.headers.get("ETag", "").strip('"') - - -def download_file(remote_filename, local_filename, remote_etag): - """Download file from R2 bucket.""" - RETRIES = 3 - for i in range(RETRIES): +def get_remote_etag(remote_filename, retries=3): + """Get ETag from R2 bucket via public URL.""" + for i in range(retries): try: - url = f"{R2_BUCKET}/{remote_filename}" - response = requests.get(url, stream=True) + response = requests.head(f"{R2_PUBLIC_BUCKET_URL}/{remote_filename}", timeout=30) response.raise_for_status() + return response.headers.get("ETag", "").strip('"') + except Exception as e: + if i < retries - 1: + print( + f"Error getting ETag for `{remote_filename}` (attempt {i + 1} of {retries}): {e}", + flush=True, + ) + time.sleep(1) + else: + raise + + +def verify_etags(): + """Verify ETags of all downloaded fixtures match remote.""" + mismatches = [] + for remote_filename in FIXTURES.keys(): + local_file = FIXTURES_DIR / remote_filename + if not local_file.exists(): + raise Exception(f"Fixture {remote_filename} not found after sync") - local_file = FIXTURES_DIR / local_filename - - with open(local_file, "wb") as f: - for chunk in response.iter_content(chunk_size=8192): - f.write(chunk) + local_etag = calculate_etag(local_file) + remote_etag = get_remote_etag(remote_filename) - local_etag = calculate_etag(local_file) - if local_etag != remote_etag: - raise Exception(f"ETag mismatch after downloading: {local_etag} != {remote_etag}") - return - except Exception as e: + if local_etag != remote_etag: + mismatches.append(f"{remote_filename}: local={local_etag}, remote={remote_etag}") + else: + print(f"ETag OK: {remote_filename}", flush=True) + + if mismatches: + raise Exception("ETag mismatches after sync:\n" + "\n".join(mismatches)) + + +def rename_fixtures(): + """Rename downloaded fixtures from remote names to local names.""" + for remote_filename, local_filename in FIXTURES.items(): + if remote_filename != local_filename: + src = FIXTURES_DIR / remote_filename + dst = FIXTURES_DIR / local_filename + if src.exists(): + src.rename(dst) + print(f"Renamed {remote_filename} -> {local_filename}", flush=True) + + +# ============================================================================= +# Authenticated path: S3 sync (used in CI with R2 credentials) +# ============================================================================= + + +def sync_fixtures_from_r2(retries: int = 3) -> None: + """Sync fixtures from R2 using aws s3 sync with retry logic.""" + cmd = [ + "aws", + "s3", + "--region", + "auto", + "--endpoint-url", + R2_S3_ENDPOINT_URL, + "--no-progress", + "--cli-connect-timeout", + "30", + "--cli-read-timeout", + "180", + "sync", + "s3://tensorzero-fixtures/", + str(FIXTURES_DIR), + # Only download the files in `FIXTURES` + "--exclude", + "*", + *[arg for f in FIXTURES.keys() for arg in ("--include", f)], + ] + + env = { + "PATH": os.environ.get("PATH", ""), + "AWS_ACCESS_KEY_ID": os.environ["R2_ACCESS_KEY_ID"], + "AWS_SECRET_ACCESS_KEY": os.environ["R2_SECRET_ACCESS_KEY"], + "AWS_MAX_ATTEMPTS": "15", + "AWS_RETRY_MODE": "adaptive", + } + + for attempt in range(retries): + print(f"Running aws s3 sync (attempt {attempt + 1} of {retries})...", flush=True) + result = subprocess.run(cmd, env=env) + + if result.returncode == 0: + print("Sync completed successfully. Verifying ETags...", flush=True) + try: + verify_etags() + rename_fixtures() + return + except Exception as e: + print( + f"Verification failed (attempt {attempt + 1} of {retries}): {e}", + flush=True, + ) + if attempt >= retries - 1: + raise + else: print( - f"Error downloading `{remote_filename}` (attempt {i + 1} of {RETRIES}): {e}", + f"aws s3 sync failed with exit code {result.returncode} (attempt {attempt + 1} of {retries})", flush=True, ) - time.sleep(1) - raise Exception(f"Failed to download `{remote_filename}` after {RETRIES} attempts") + if attempt < retries - 1: + sleep_time = 3**attempt # Exponential backoff: 1, 3, 9 seconds + print(f"Retrying in {sleep_time} seconds...", flush=True) + time.sleep(sleep_time) -def main(): - if os.environ.get("R2_ACCESS_KEY_ID") is not None and os.environ.get("R2_SECRET_ACCESS_KEY"): - print("R2_ACCESS_KEY_ID set, downloading fixtures using 'aws s3 cp'", flush=True) - for remote_filename, local_filename in FIXTURES.items(): - local_file = FIXTURES_DIR / local_filename - subprocess.check_call( - f"aws s3 --region auto --endpoint-url https://19918a216783f0ac9e052233569aef60.r2.cloudflarestorage.com/ " - f"cp s3://tensorzero-fixtures/{remote_filename} {local_file}", - env={ - "AWS_ACCESS_KEY_ID": os.environ["R2_ACCESS_KEY_ID"], - "AWS_SECRET_ACCESS_KEY": os.environ["R2_SECRET_ACCESS_KEY"], - "PATH": os.environ.get("PATH", ""), - }, - shell=True, - ) - return + raise Exception(f"aws s3 sync failed after {retries} attempts") - def process_fixture(item): - remote_filename, local_filename = item - local_file = FIXTURES_DIR / local_filename - remote_etag = get_remote_etag(remote_filename) - if not local_file.exists(): - print( - f"Downloading {remote_filename} (file doesn't exist locally)", - flush=True, - ) - download_file(remote_filename, local_filename, remote_etag) - return +# ============================================================================= +# Main +# ============================================================================= - local_etag = calculate_etag(local_file) - if local_etag != remote_etag: - print(f"Downloading {remote_filename} (ETag mismatch)", flush=True) - print(f"Local ETag: {local_etag}", flush=True) - print(f"Remote ETag: {remote_etag}", flush=True) - download_file(remote_filename, local_filename, remote_etag) - else: - print(f"Skipping {remote_filename} (up to date)", flush=True) +def main(): + FIXTURES_DIR.mkdir(exist_ok=True) + + if not os.environ.get("R2_ACCESS_KEY_ID") or not os.environ.get("R2_SECRET_ACCESS_KEY"): + raise Exception( + "R2_ACCESS_KEY_ID and R2_SECRET_ACCESS_KEY must be set. " + "For local development without R2 credentials, use download-small-fixtures-http.py instead." + ) - # Use ThreadPoolExecutor to download files in parallel - with concurrent.futures.ThreadPoolExecutor() as executor: - # Loop over the results to propagate exceptions - for result in executor.map(process_fixture, FIXTURES.items()): - assert result is None + print("R2 credentials found, downloading fixtures using `aws s3 sync`", flush=True) + sync_fixtures_from_r2() if __name__ == "__main__": diff --git a/ui/fixtures/load_fixtures.sh b/ui/fixtures/load_fixtures.sh index 889a8c4300d..b8d4a0b51af 100755 --- a/ui/fixtures/load_fixtures.sh +++ b/ui/fixtures/load_fixtures.sh @@ -38,21 +38,25 @@ clickhouse-client --host $CLICKHOUSE_HOST_VAR --user $CLICKHOUSE_USER_VAR --pass clickhouse-client --host $CLICKHOUSE_HOST_VAR --user $CLICKHOUSE_USER_VAR --password $CLICKHOUSE_PASSWORD_VAR $CLICKHOUSE_SECURE_FLAG --database "$DATABASE_NAME" --query "TRUNCATE TABLE DeploymentID" # Download JSONL fixtures from R2 -uv run ./download-small-fixtures.py +if [ "${TENSORZERO_DOWNLOAD_FIXTURES_WITHOUT_CREDENTIALS:-}" = "1" ]; then + uv run ./download-small-fixtures-http.py +else + uv run ./download-small-fixtures.py +fi -clickhouse-client --host $CLICKHOUSE_HOST_VAR --user $CLICKHOUSE_USER_VAR --password $CLICKHOUSE_PASSWORD_VAR $CLICKHOUSE_SECURE_FLAG --database "$DATABASE_NAME" --query "INSERT INTO JsonInference FORMAT JSONEachRow" < json_inference_examples.jsonl -clickhouse-client --host $CLICKHOUSE_HOST_VAR --user $CLICKHOUSE_USER_VAR --password $CLICKHOUSE_PASSWORD_VAR $CLICKHOUSE_SECURE_FLAG --database "$DATABASE_NAME" --query "INSERT INTO ChatInference FORMAT JSONEachRow" < chat_inference_examples.jsonl -clickhouse-client --host $CLICKHOUSE_HOST_VAR --user $CLICKHOUSE_USER_VAR --password $CLICKHOUSE_PASSWORD_VAR $CLICKHOUSE_SECURE_FLAG --database "$DATABASE_NAME" --query "INSERT INTO BooleanMetricFeedback FORMAT JSONEachRow" < boolean_metric_feedback_examples.jsonl -clickhouse-client --host $CLICKHOUSE_HOST_VAR --user $CLICKHOUSE_USER_VAR --password $CLICKHOUSE_PASSWORD_VAR $CLICKHOUSE_SECURE_FLAG --database "$DATABASE_NAME" --query "INSERT INTO FloatMetricFeedback FORMAT JSONEachRow" < float_metric_feedback_examples.jsonl -clickhouse-client --host $CLICKHOUSE_HOST_VAR --user $CLICKHOUSE_USER_VAR --password $CLICKHOUSE_PASSWORD_VAR $CLICKHOUSE_SECURE_FLAG --database "$DATABASE_NAME" --query "INSERT INTO FloatMetricFeedback FORMAT JSONEachRow" < jaro_winkler_similarity_feedback.jsonl -clickhouse-client --host $CLICKHOUSE_HOST_VAR --user $CLICKHOUSE_USER_VAR --password $CLICKHOUSE_PASSWORD_VAR $CLICKHOUSE_SECURE_FLAG --database "$DATABASE_NAME" --query "INSERT INTO CommentFeedback FORMAT JSONEachRow" < comment_feedback_examples.jsonl -clickhouse-client --host $CLICKHOUSE_HOST_VAR --user $CLICKHOUSE_USER_VAR --password $CLICKHOUSE_PASSWORD_VAR $CLICKHOUSE_SECURE_FLAG --database "$DATABASE_NAME" --query "INSERT INTO DemonstrationFeedback FORMAT JSONEachRow" < demonstration_feedback_examples.jsonl -clickhouse-client --host $CLICKHOUSE_HOST_VAR --user $CLICKHOUSE_USER_VAR --password $CLICKHOUSE_PASSWORD_VAR $CLICKHOUSE_SECURE_FLAG --database "$DATABASE_NAME" --query "INSERT INTO ModelInference FORMAT JSONEachRow" < model_inference_examples.jsonl -clickhouse-client --host $CLICKHOUSE_HOST_VAR --user $CLICKHOUSE_USER_VAR --password $CLICKHOUSE_PASSWORD_VAR $CLICKHOUSE_SECURE_FLAG --database "$DATABASE_NAME" --query "INSERT INTO ChatInferenceDatapoint FORMAT JSONEachRow" < chat_inference_datapoint_examples.jsonl -clickhouse-client --host $CLICKHOUSE_HOST_VAR --user $CLICKHOUSE_USER_VAR --password $CLICKHOUSE_PASSWORD_VAR $CLICKHOUSE_SECURE_FLAG --database "$DATABASE_NAME" --query "INSERT INTO JsonInferenceDatapoint FORMAT JSONEachRow" < json_inference_datapoint_examples.jsonl -clickhouse-client --host $CLICKHOUSE_HOST_VAR --user $CLICKHOUSE_USER_VAR --password $CLICKHOUSE_PASSWORD_VAR $CLICKHOUSE_SECURE_FLAG --database "$DATABASE_NAME" --query "INSERT INTO DynamicEvaluationRun FORMAT JSONEachRow" < dynamic_evaluation_run_examples.jsonl -clickhouse-client --host $CLICKHOUSE_HOST_VAR --user $CLICKHOUSE_USER_VAR --password $CLICKHOUSE_PASSWORD_VAR $CLICKHOUSE_SECURE_FLAG --database "$DATABASE_NAME" --query "INSERT INTO DynamicEvaluationRunEpisode FORMAT JSONEachRow" < dynamic_evaluation_run_episode_examples.jsonl -clickhouse-client --host $CLICKHOUSE_HOST_VAR --user $CLICKHOUSE_USER_VAR --password $CLICKHOUSE_PASSWORD_VAR $CLICKHOUSE_SECURE_FLAG --database "$DATABASE_NAME" --query "INSERT INTO ModelInferenceCache FORMAT JSONEachRow" < model_inference_cache_e2e.jsonl +clickhouse-client --host $CLICKHOUSE_HOST_VAR --user $CLICKHOUSE_USER_VAR --password $CLICKHOUSE_PASSWORD_VAR $CLICKHOUSE_SECURE_FLAG --database "$DATABASE_NAME" --query "INSERT INTO JsonInference FORMAT JSONEachRow" < ./small-fixtures/json_inference_examples.jsonl +clickhouse-client --host $CLICKHOUSE_HOST_VAR --user $CLICKHOUSE_USER_VAR --password $CLICKHOUSE_PASSWORD_VAR $CLICKHOUSE_SECURE_FLAG --database "$DATABASE_NAME" --query "INSERT INTO ChatInference FORMAT JSONEachRow" < ./small-fixtures/chat_inference_examples.jsonl +clickhouse-client --host $CLICKHOUSE_HOST_VAR --user $CLICKHOUSE_USER_VAR --password $CLICKHOUSE_PASSWORD_VAR $CLICKHOUSE_SECURE_FLAG --database "$DATABASE_NAME" --query "INSERT INTO BooleanMetricFeedback FORMAT JSONEachRow" < ./small-fixtures/boolean_metric_feedback_examples.jsonl +clickhouse-client --host $CLICKHOUSE_HOST_VAR --user $CLICKHOUSE_USER_VAR --password $CLICKHOUSE_PASSWORD_VAR $CLICKHOUSE_SECURE_FLAG --database "$DATABASE_NAME" --query "INSERT INTO FloatMetricFeedback FORMAT JSONEachRow" < ./small-fixtures/float_metric_feedback_examples.jsonl +clickhouse-client --host $CLICKHOUSE_HOST_VAR --user $CLICKHOUSE_USER_VAR --password $CLICKHOUSE_PASSWORD_VAR $CLICKHOUSE_SECURE_FLAG --database "$DATABASE_NAME" --query "INSERT INTO FloatMetricFeedback FORMAT JSONEachRow" < ./small-fixtures/jaro_winkler_similarity_feedback.jsonl +clickhouse-client --host $CLICKHOUSE_HOST_VAR --user $CLICKHOUSE_USER_VAR --password $CLICKHOUSE_PASSWORD_VAR $CLICKHOUSE_SECURE_FLAG --database "$DATABASE_NAME" --query "INSERT INTO CommentFeedback FORMAT JSONEachRow" < ./small-fixtures/comment_feedback_examples.jsonl +clickhouse-client --host $CLICKHOUSE_HOST_VAR --user $CLICKHOUSE_USER_VAR --password $CLICKHOUSE_PASSWORD_VAR $CLICKHOUSE_SECURE_FLAG --database "$DATABASE_NAME" --query "INSERT INTO DemonstrationFeedback FORMAT JSONEachRow" < ./small-fixtures/demonstration_feedback_examples.jsonl +clickhouse-client --host $CLICKHOUSE_HOST_VAR --user $CLICKHOUSE_USER_VAR --password $CLICKHOUSE_PASSWORD_VAR $CLICKHOUSE_SECURE_FLAG --database "$DATABASE_NAME" --query "INSERT INTO ModelInference FORMAT JSONEachRow" < ./small-fixtures/model_inference_examples.jsonl +clickhouse-client --host $CLICKHOUSE_HOST_VAR --user $CLICKHOUSE_USER_VAR --password $CLICKHOUSE_PASSWORD_VAR $CLICKHOUSE_SECURE_FLAG --database "$DATABASE_NAME" --query "INSERT INTO ChatInferenceDatapoint FORMAT JSONEachRow" < ./small-fixtures/chat_inference_datapoint_examples.jsonl +clickhouse-client --host $CLICKHOUSE_HOST_VAR --user $CLICKHOUSE_USER_VAR --password $CLICKHOUSE_PASSWORD_VAR $CLICKHOUSE_SECURE_FLAG --database "$DATABASE_NAME" --query "INSERT INTO JsonInferenceDatapoint FORMAT JSONEachRow" < ./small-fixtures/json_inference_datapoint_examples.jsonl +clickhouse-client --host $CLICKHOUSE_HOST_VAR --user $CLICKHOUSE_USER_VAR --password $CLICKHOUSE_PASSWORD_VAR $CLICKHOUSE_SECURE_FLAG --database "$DATABASE_NAME" --query "INSERT INTO DynamicEvaluationRun FORMAT JSONEachRow" < ./small-fixtures/dynamic_evaluation_run_examples.jsonl +clickhouse-client --host $CLICKHOUSE_HOST_VAR --user $CLICKHOUSE_USER_VAR --password $CLICKHOUSE_PASSWORD_VAR $CLICKHOUSE_SECURE_FLAG --database "$DATABASE_NAME" --query "INSERT INTO DynamicEvaluationRunEpisode FORMAT JSONEachRow" < ./small-fixtures/dynamic_evaluation_run_episode_examples.jsonl +clickhouse-client --host $CLICKHOUSE_HOST_VAR --user $CLICKHOUSE_USER_VAR --password $CLICKHOUSE_PASSWORD_VAR $CLICKHOUSE_SECURE_FLAG --database "$DATABASE_NAME" --query "INSERT INTO ModelInferenceCache FORMAT JSONEachRow" < ./small-fixtures/model_inference_cache_e2e.jsonl clickhouse-client --host $CLICKHOUSE_HOST_VAR --user $CLICKHOUSE_USER_VAR --password $CLICKHOUSE_PASSWORD_VAR $CLICKHOUSE_SECURE_FLAG --database "$DATABASE_NAME" --query "INSERT INTO DeploymentID VALUES ('fixture', 0, 0, 4294967295)" # If TENSORZERO_SKIP_LARGE_FIXTURES equals 1, exit @@ -63,19 +67,23 @@ if [ "${TENSORZERO_SKIP_LARGE_FIXTURES:-}" = "1" ]; then fi uv run python --version -uv run ./download-large-fixtures.py -clickhouse-client --host $CLICKHOUSE_HOST_VAR --user $CLICKHOUSE_USER_VAR --password $CLICKHOUSE_PASSWORD_VAR $CLICKHOUSE_SECURE_FLAG --database "$DATABASE_NAME" --query "INSERT INTO ChatInference FROM INFILE './s3-fixtures/large_chat_inference_v2.parquet' SETTINGS input_format_parquet_use_native_reader_v3=0 FORMAT Parquet" -clickhouse-client --host $CLICKHOUSE_HOST_VAR --user $CLICKHOUSE_USER_VAR --password $CLICKHOUSE_PASSWORD_VAR $CLICKHOUSE_SECURE_FLAG --database "$DATABASE_NAME" --query "INSERT INTO JsonInference FROM INFILE './s3-fixtures/large_json_inference_v2.parquet' SETTINGS input_format_parquet_use_native_reader_v3=0 FORMAT Parquet" -clickhouse-client --host $CLICKHOUSE_HOST_VAR --user $CLICKHOUSE_USER_VAR --password $CLICKHOUSE_PASSWORD_VAR $CLICKHOUSE_SECURE_FLAG --database "$DATABASE_NAME" --query "INSERT INTO ModelInference FROM INFILE './s3-fixtures/large_chat_model_inference_v2.parquet' SETTINGS input_format_parquet_use_native_reader_v3=0 FORMAT Parquet" -clickhouse-client --host $CLICKHOUSE_HOST_VAR --user $CLICKHOUSE_USER_VAR --password $CLICKHOUSE_PASSWORD_VAR $CLICKHOUSE_SECURE_FLAG --database "$DATABASE_NAME" --query "INSERT INTO ModelInference FROM INFILE './s3-fixtures/large_json_model_inference_v2.parquet' SETTINGS input_format_parquet_use_native_reader_v3=0 FORMAT Parquet" -clickhouse-client --host $CLICKHOUSE_HOST_VAR --user $CLICKHOUSE_USER_VAR --password $CLICKHOUSE_PASSWORD_VAR $CLICKHOUSE_SECURE_FLAG --database "$DATABASE_NAME" --query "INSERT INTO BooleanMetricFeedback FROM INFILE './s3-fixtures/large_chat_boolean_feedback.parquet' SETTINGS input_format_parquet_use_native_reader_v3=0 FORMAT Parquet" -clickhouse-client --host $CLICKHOUSE_HOST_VAR --user $CLICKHOUSE_USER_VAR --password $CLICKHOUSE_PASSWORD_VAR $CLICKHOUSE_SECURE_FLAG --database "$DATABASE_NAME" --query "INSERT INTO BooleanMetricFeedback FROM INFILE './s3-fixtures/large_json_boolean_feedback.parquet' SETTINGS input_format_parquet_use_native_reader_v3=0 FORMAT Parquet" -clickhouse-client --host $CLICKHOUSE_HOST_VAR --user $CLICKHOUSE_USER_VAR --password $CLICKHOUSE_PASSWORD_VAR $CLICKHOUSE_SECURE_FLAG --database "$DATABASE_NAME" --query "INSERT INTO FloatMetricFeedback FROM INFILE './s3-fixtures/large_chat_float_feedback.parquet' SETTINGS input_format_parquet_use_native_reader_v3=0 FORMAT Parquet" -clickhouse-client --host $CLICKHOUSE_HOST_VAR --user $CLICKHOUSE_USER_VAR --password $CLICKHOUSE_PASSWORD_VAR $CLICKHOUSE_SECURE_FLAG --database "$DATABASE_NAME" --query "INSERT INTO FloatMetricFeedback FROM INFILE './s3-fixtures/large_json_float_feedback.parquet' SETTINGS input_format_parquet_use_native_reader_v3=0 FORMAT Parquet" -clickhouse-client --host $CLICKHOUSE_HOST_VAR --user $CLICKHOUSE_USER_VAR --password $CLICKHOUSE_PASSWORD_VAR $CLICKHOUSE_SECURE_FLAG --database "$DATABASE_NAME" --query "INSERT INTO CommentFeedback FROM INFILE './s3-fixtures/large_chat_comment_feedback.parquet' SETTINGS input_format_parquet_use_native_reader_v3=0 FORMAT Parquet" -clickhouse-client --host $CLICKHOUSE_HOST_VAR --user $CLICKHOUSE_USER_VAR --password $CLICKHOUSE_PASSWORD_VAR $CLICKHOUSE_SECURE_FLAG --database "$DATABASE_NAME" --query "INSERT INTO CommentFeedback FROM INFILE './s3-fixtures/large_json_comment_feedback.parquet' SETTINGS input_format_parquet_use_native_reader_v3=0 FORMAT Parquet" -clickhouse-client --host $CLICKHOUSE_HOST_VAR --user $CLICKHOUSE_USER_VAR --password $CLICKHOUSE_PASSWORD_VAR $CLICKHOUSE_SECURE_FLAG --database "$DATABASE_NAME" --query "INSERT INTO DemonstrationFeedback FROM INFILE './s3-fixtures/large_chat_demonstration_feedback.parquet' SETTINGS input_format_parquet_use_native_reader_v3=0 FORMAT Parquet" -clickhouse-client --host $CLICKHOUSE_HOST_VAR --user $CLICKHOUSE_USER_VAR --password $CLICKHOUSE_PASSWORD_VAR $CLICKHOUSE_SECURE_FLAG --database "$DATABASE_NAME" --query "INSERT INTO DemonstrationFeedback FROM INFILE './s3-fixtures/large_json_demonstration_feedback.parquet' SETTINGS input_format_parquet_use_native_reader_v3=0 FORMAT Parquet" +if [ "${TENSORZERO_DOWNLOAD_FIXTURES_WITHOUT_CREDENTIALS:-}" = "1" ]; then + uv run ./download-large-fixtures-http.py +else + uv run ./download-large-fixtures.py +fi +clickhouse-client --host $CLICKHOUSE_HOST_VAR --user $CLICKHOUSE_USER_VAR --password $CLICKHOUSE_PASSWORD_VAR $CLICKHOUSE_SECURE_FLAG --database "$DATABASE_NAME" --query "INSERT INTO ChatInference FROM INFILE './large-fixtures/large_chat_inference_v2.parquet' SETTINGS input_format_parquet_use_native_reader_v3=0 FORMAT Parquet" +clickhouse-client --host $CLICKHOUSE_HOST_VAR --user $CLICKHOUSE_USER_VAR --password $CLICKHOUSE_PASSWORD_VAR $CLICKHOUSE_SECURE_FLAG --database "$DATABASE_NAME" --query "INSERT INTO JsonInference FROM INFILE './large-fixtures/large_json_inference_v2.parquet' SETTINGS input_format_parquet_use_native_reader_v3=0 FORMAT Parquet" +clickhouse-client --host $CLICKHOUSE_HOST_VAR --user $CLICKHOUSE_USER_VAR --password $CLICKHOUSE_PASSWORD_VAR $CLICKHOUSE_SECURE_FLAG --database "$DATABASE_NAME" --query "INSERT INTO ModelInference FROM INFILE './large-fixtures/large_chat_model_inference_v2.parquet' SETTINGS input_format_parquet_use_native_reader_v3=0 FORMAT Parquet" +clickhouse-client --host $CLICKHOUSE_HOST_VAR --user $CLICKHOUSE_USER_VAR --password $CLICKHOUSE_PASSWORD_VAR $CLICKHOUSE_SECURE_FLAG --database "$DATABASE_NAME" --query "INSERT INTO ModelInference FROM INFILE './large-fixtures/large_json_model_inference_v2.parquet' SETTINGS input_format_parquet_use_native_reader_v3=0 FORMAT Parquet" +clickhouse-client --host $CLICKHOUSE_HOST_VAR --user $CLICKHOUSE_USER_VAR --password $CLICKHOUSE_PASSWORD_VAR $CLICKHOUSE_SECURE_FLAG --database "$DATABASE_NAME" --query "INSERT INTO BooleanMetricFeedback FROM INFILE './large-fixtures/large_chat_boolean_feedback.parquet' SETTINGS input_format_parquet_use_native_reader_v3=0 FORMAT Parquet" +clickhouse-client --host $CLICKHOUSE_HOST_VAR --user $CLICKHOUSE_USER_VAR --password $CLICKHOUSE_PASSWORD_VAR $CLICKHOUSE_SECURE_FLAG --database "$DATABASE_NAME" --query "INSERT INTO BooleanMetricFeedback FROM INFILE './large-fixtures/large_json_boolean_feedback.parquet' SETTINGS input_format_parquet_use_native_reader_v3=0 FORMAT Parquet" +clickhouse-client --host $CLICKHOUSE_HOST_VAR --user $CLICKHOUSE_USER_VAR --password $CLICKHOUSE_PASSWORD_VAR $CLICKHOUSE_SECURE_FLAG --database "$DATABASE_NAME" --query "INSERT INTO FloatMetricFeedback FROM INFILE './large-fixtures/large_chat_float_feedback.parquet' SETTINGS input_format_parquet_use_native_reader_v3=0 FORMAT Parquet" +clickhouse-client --host $CLICKHOUSE_HOST_VAR --user $CLICKHOUSE_USER_VAR --password $CLICKHOUSE_PASSWORD_VAR $CLICKHOUSE_SECURE_FLAG --database "$DATABASE_NAME" --query "INSERT INTO FloatMetricFeedback FROM INFILE './large-fixtures/large_json_float_feedback.parquet' SETTINGS input_format_parquet_use_native_reader_v3=0 FORMAT Parquet" +clickhouse-client --host $CLICKHOUSE_HOST_VAR --user $CLICKHOUSE_USER_VAR --password $CLICKHOUSE_PASSWORD_VAR $CLICKHOUSE_SECURE_FLAG --database "$DATABASE_NAME" --query "INSERT INTO CommentFeedback FROM INFILE './large-fixtures/large_chat_comment_feedback.parquet' SETTINGS input_format_parquet_use_native_reader_v3=0 FORMAT Parquet" +clickhouse-client --host $CLICKHOUSE_HOST_VAR --user $CLICKHOUSE_USER_VAR --password $CLICKHOUSE_PASSWORD_VAR $CLICKHOUSE_SECURE_FLAG --database "$DATABASE_NAME" --query "INSERT INTO CommentFeedback FROM INFILE './large-fixtures/large_json_comment_feedback.parquet' SETTINGS input_format_parquet_use_native_reader_v3=0 FORMAT Parquet" +clickhouse-client --host $CLICKHOUSE_HOST_VAR --user $CLICKHOUSE_USER_VAR --password $CLICKHOUSE_PASSWORD_VAR $CLICKHOUSE_SECURE_FLAG --database "$DATABASE_NAME" --query "INSERT INTO DemonstrationFeedback FROM INFILE './large-fixtures/large_chat_demonstration_feedback.parquet' SETTINGS input_format_parquet_use_native_reader_v3=0 FORMAT Parquet" +clickhouse-client --host $CLICKHOUSE_HOST_VAR --user $CLICKHOUSE_USER_VAR --password $CLICKHOUSE_PASSWORD_VAR $CLICKHOUSE_SECURE_FLAG --database "$DATABASE_NAME" --query "INSERT INTO DemonstrationFeedback FROM INFILE './large-fixtures/large_json_demonstration_feedback.parquet' SETTINGS input_format_parquet_use_native_reader_v3=0 FORMAT Parquet" # Give ClickHouse some time to make the writes visible sleep 2 diff --git a/ui/fixtures/load_fixtures_postgres.sh b/ui/fixtures/load_fixtures_postgres.sh index f42e7a2cfd6..e3c1c09e008 100755 --- a/ui/fixtures/load_fixtures_postgres.sh +++ b/ui/fixtures/load_fixtures_postgres.sh @@ -61,15 +61,19 @@ EOF fi # Download JSONL fixtures from R2 if not present -if [ ! -f "chat_inference_examples.jsonl" ] || [ ! -f "json_inference_examples.jsonl" ]; then +if [ ! -f "small-fixtures/chat_inference_examples.jsonl" ] || [ ! -f "small-fixtures/json_inference_examples.jsonl" ]; then echo "Downloading small fixtures..." - uv run ./download-small-fixtures.py + if [ "${TENSORZERO_DOWNLOAD_FIXTURES_WITHOUT_CREDENTIALS:-}" = "1" ]; then + uv run ./download-small-fixtures-http.py + else + uv run ./download-small-fixtures.py + fi fi # Chat Inferences # Note: input, output, tool_params, inference_params are JSONB in our schema # created_at is derived from the UUIDv7 id using tensorzero.uuid_v7_to_timestamp() -load_jsonl "chat_inference_examples.jsonl" "tensorzero.chat_inferences" " +load_jsonl "small-fixtures/chat_inference_examples.jsonl" "tensorzero.chat_inferences" " INSERT INTO tensorzero.chat_inferences ( id, function_name, variant_name, episode_id, input, output, tool_params, inference_params, @@ -103,7 +107,7 @@ ON CONFLICT (id, created_at) DO NOTHING; # JSON Inferences # Note: input, output, output_schema, inference_params, auxiliary_content are JSONB in our schema # created_at is derived from the UUIDv7 id using tensorzero.uuid_v7_to_timestamp() -load_jsonl "json_inference_examples.jsonl" "tensorzero.json_inferences" " +load_jsonl "small-fixtures/json_inference_examples.jsonl" "tensorzero.json_inferences" " INSERT INTO tensorzero.json_inferences ( id, function_name, variant_name, episode_id, input, output, output_schema, inference_params, diff --git a/ui/fixtures/regenerate-model-inference-cache.sh b/ui/fixtures/regenerate-model-inference-cache.sh index 09c877cbeda..5f4dd2378f4 100755 --- a/ui/fixtures/regenerate-model-inference-cache.sh +++ b/ui/fixtures/regenerate-model-inference-cache.sh @@ -18,5 +18,5 @@ export TENSORZERO_PLAYWRIGHT_NO_WEBSERVER=1 export TENSORZERO_GATEWAY_URL=http://localhost:3000 pnpm test-e2e -j 1 --grep-invert "@credentials" --max-failures 1 # Remove the existing file if it exists (it may have restricted permissions from R2 download) -rm -f ./fixtures/model_inference_cache_e2e.jsonl -docker run --add-host=host.docker.internal:host-gateway clickhouse/clickhouse-server clickhouse-client --host host.docker.internal --user chuser --password chpassword --database tensorzero_ui_fixtures 'SELECT * FROM ModelInferenceCache ORDER BY long_cache_key ASC FORMAT JSONEachRow' > ./fixtures/model_inference_cache_e2e.jsonl +rm -f ./fixtures/small-fixtures/model_inference_cache_e2e.jsonl +docker run --add-host=host.docker.internal:host-gateway clickhouse/clickhouse-server clickhouse-client --host host.docker.internal --user chuser --password chpassword --database tensorzero_ui_fixtures 'SELECT * FROM ModelInferenceCache ORDER BY long_cache_key ASC FORMAT JSONEachRow' > ./fixtures/small-fixtures/model_inference_cache_e2e.jsonl diff --git a/ui/fixtures/upload-large-fixtures.sh b/ui/fixtures/upload-large-fixtures.sh index 604e723ab5b..0455ca8a34a 100755 --- a/ui/fixtures/upload-large-fixtures.sh +++ b/ui/fixtures/upload-large-fixtures.sh @@ -1,6 +1,6 @@ #!/bin/bash set -euxo pipefail -cd $(dirname $0)/s3-fixtures +cd $(dirname $0)/large-fixtures aws s3 --endpoint-url https://19918a216783f0ac9e052233569aef60.r2.cloudflarestorage.com/ sync . s3://tensorzero-fixtures --checksum-algorithm CRC32 diff --git a/ui/fixtures/upload-small-fixtures.sh b/ui/fixtures/upload-small-fixtures.sh index d2878543f1d..34712f6c527 100755 --- a/ui/fixtures/upload-small-fixtures.sh +++ b/ui/fixtures/upload-small-fixtures.sh @@ -24,12 +24,12 @@ JSONL_FILES=( # Upload each file for file in "${JSONL_FILES[@]}"; do - if [ -f "$file" ]; then + if [ -f "small-fixtures/$file" ]; then echo "Uploading $file..." aws s3 --endpoint-url https://19918a216783f0ac9e052233569aef60.r2.cloudflarestorage.com/ \ - cp "$file" "s3://tensorzero-fixtures/${file}" + cp "small-fixtures/$file" "s3://tensorzero-fixtures/${file}" else - echo "Warning: $file not found, skipping" + echo "Warning: small-fixtures/$file not found, skipping" fi done From 406549a18aa0a35d2faf6c481a247a30d3e9c12d Mon Sep 17 00:00:00 2001 From: Gabriel Bianconi <1275491+GabrielBianconi@users.noreply.github.com> Date: Tue, 27 Jan 2026 18:56:12 -0500 Subject: [PATCH 44/46] Fix OpenAI provider tool serialization bug (#5924) --- .../src/providers/openai/responses.rs | 44 ++++++++++++++----- 1 file changed, 32 insertions(+), 12 deletions(-) diff --git a/tensorzero-core/src/providers/openai/responses.rs b/tensorzero-core/src/providers/openai/responses.rs index 8772f65fdbb..67c3d06d103 100644 --- a/tensorzero-core/src/providers/openai/responses.rs +++ b/tensorzero-core/src/providers/openai/responses.rs @@ -295,12 +295,12 @@ impl<'a> OpenAITool<'a> { pub fn into_openai_responses_tool(self) -> OpenAIResponsesTool<'a> { match self { OpenAITool::Function { function, strict } => { - OpenAIResponsesTool::Function(OpenAIResponsesFunctionTool { - name: function.name, - description: function.description, - parameters: function.parameters, + OpenAIResponsesTool::Function(OpenAIResponsesFunctionTool::new( + function.name, + function.description, + function.parameters, strict, - }) + )) } OpenAITool::Custom { custom: custom_tool, @@ -310,21 +310,39 @@ impl<'a> OpenAITool<'a> { } #[derive(Debug, Serialize)] -#[serde(tag = "type", rename_all = "snake_case")] +#[serde(untagged)] pub enum OpenAIResponsesTool<'a> { Function(OpenAIResponsesFunctionTool<'a>), - BuiltIn(&'a Value), Custom(OpenAIResponsesCustomTool<'a>), + Provider(&'a Value), // example: {"type": "web_search"} } #[derive(Serialize, Debug)] pub struct OpenAIResponsesFunctionTool<'a> { + r#type: &'static str, // must be "function" name: &'a str, description: Option<&'a str>, parameters: &'a Value, strict: bool, } +impl<'a> OpenAIResponsesFunctionTool<'a> { + pub fn new( + name: &'a str, + description: Option<&'a str>, + parameters: &'a Value, + strict: bool, + ) -> Self { + Self { + r#type: "function", + name, + description, + parameters, + strict, + } + } +} + /// Custom tool format for the Responses API (flattened structure) #[derive(Clone, Debug, Serialize)] #[serde(tag = "type", rename_all = "snake_case")] @@ -338,6 +356,7 @@ pub enum OpenAIResponsesCustomToolFormat { #[derive(Serialize, Debug)] pub struct OpenAIResponsesCustomTool<'a> { + r#type: &'static str, // must be "custom" name: &'a str, #[serde(skip_serializing_if = "Option::is_none")] description: Option<&'a str>, @@ -348,6 +367,7 @@ pub struct OpenAIResponsesCustomTool<'a> { impl<'a> From<&'a OpenAICustomTool> for OpenAIResponsesCustomTool<'a> { fn from(tool: &'a OpenAICustomTool) -> Self { OpenAIResponsesCustomTool { + r#type: "custom", name: &tool.name, description: tool.description.as_deref(), format: tool.format.as_ref().map(|f| match f { @@ -414,7 +434,7 @@ impl<'a> OpenAIResponsesRequest<'a> { openai_model: &'a str, request: &'a ModelInferenceRequest<'_>, include_encrypted_reasoning: bool, - built_in_tools: &'a [Value], + provider_tools: &'a [Value], tensorzero_model_name: &str, tensorzero_model_provider_name: &str, ) -> Result, Error> { @@ -435,15 +455,15 @@ impl<'a> OpenAIResponsesRequest<'a> { .collect() }) .unwrap_or_default(); - // If we have built_in_tools we should extend the list with them - tools.extend(built_in_tools.iter().map(OpenAIResponsesTool::BuiltIn)); + // If we have `provider_tools` we should extend the list with them + tools.extend(provider_tools.iter().map(OpenAIResponsesTool::Provider)); if let Some(tc) = request.tool_config.as_ref() { let provider_tools = tc.get_scoped_provider_tools(tensorzero_model_name, tensorzero_model_provider_name); tools.extend( provider_tools .iter() - .map(|t| OpenAIResponsesTool::BuiltIn(&t.tool)), + .map(|t| OpenAIResponsesTool::Provider(&t.tool)), ); } @@ -453,7 +473,7 @@ impl<'a> OpenAIResponsesRequest<'a> { .as_ref() .and_then(|tc| tc.allowed_tools.as_dynamic_allowed_tools()); - // For now, we don't allow selecting any built-in tools + // For now, we don't allow selecting any provider tools let tool_choice = request.tool_config.as_ref().map(|tool_config| { // If we have allowed_tools, create an AllowedTools variant if let Some(allowed_tool_names) = &allowed_tools_list { From 9e8be1094800be32f036375c9d244c45efb53128 Mon Sep 17 00:00:00 2001 From: Aaron Hill Date: Tue, 27 Jan 2026 20:12:52 -0500 Subject: [PATCH 45/46] Give evaluations tests more time to run (#5922) These consistently take 4 or 5 retries to pass in ci, always due to timeouts --- .config/nextest.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.config/nextest.toml b/.config/nextest.toml index 5bde8e94a34..f45a8cfc511 100644 --- a/.config/nextest.toml +++ b/.config/nextest.toml @@ -13,7 +13,7 @@ default-filter = "not (test(no_aws_credentials) | test(export_bindings_) | test( [[profile.default.overrides]] filter = 'test(evaluations)' -slow-timeout = { period = "20s", terminate-after = 3 } +slow-timeout = { period = "20s", terminate-after = 5 } # Profiles config # We use these profiles to define our major test groups. From c4ca3d611a53496f2fd834d5e29e96b626c95bfe Mon Sep 17 00:00:00 2001 From: Gabriel Bianconi <1275491+GabrielBianconi@users.noreply.github.com> Date: Tue, 27 Jan 2026 20:41:38 -0500 Subject: [PATCH 46/46] Improve reasoning for streaming OpenAI Responses API (#5861) * Streaming Responses API include_encrypted_reasoning * Allow cursoragent in CLA workflow Co-authored-by: gabriel * Improve reasoning for streaming OpenAI Responses API * Fix * Fix * Fix * Fix * Fix --------- Co-authored-by: Cursor Agent --- tensorzero-core/src/endpoints/inference.rs | 16 +- .../src/inference/types/streams.rs | 305 +++++++++++--- tensorzero-core/src/providers/openai/mod.rs | 3 +- .../src/providers/openai/responses.rs | 392 +++++++++++++++++- tensorzero-core/src/variant/mixture_of_n.rs | 2 +- tensorzero-core/tests/e2e/providers/common.rs | 24 +- .../tests/e2e/providers/openai/mod.rs | 19 +- .../tests/e2e/providers/reasoning.rs | 44 +- 8 files changed, 692 insertions(+), 113 deletions(-) diff --git a/tensorzero-core/src/endpoints/inference.rs b/tensorzero-core/src/endpoints/inference.rs index 60edb5954f2..20ec4b4f981 100644 --- a/tensorzero-core/src/endpoints/inference.rs +++ b/tensorzero-core/src/endpoints/inference.rs @@ -1272,7 +1272,7 @@ fn should_stream_chunk_in_create_stream( // We already handled `include_original_response` above raw_chunk: _, // We never actually stream this field, so we don't need it - thought: _, + thought_chunks: _, // We don't care about streaming the following fields in isolation provider_latency: _, } = c; @@ -2059,7 +2059,7 @@ mod tests { use crate::inference::types::{ ApiType, Base64File, ChatInferenceResultChunk, ContentBlockChunk, ContentBlockOutput, File, InputMessageContent, JsonInferenceResultChunk, Latency, ModelInferenceResponseWithMetadata, - ObjectStoragePointer, RequestMessagesOrBatch, Role, Text, TextChunk, UrlFile, + ObjectStoragePointer, RequestMessagesOrBatch, Role, Text, TextChunk, ThoughtChunk, UrlFile, storage::{StorageKind, StoragePath}, usage::RawUsageEntry, }; @@ -2132,7 +2132,15 @@ mod tests { // Test case 2: Valid JSON ProviderInferenceResponseChunk let chunk = InferenceResultChunk::Json(JsonInferenceResultChunk { raw: Some("Test content".to_string()), - thought: Some("Thought 1".to_string()), + thought_chunks: vec![ThoughtChunk { + id: "0".to_string(), + text: Some("Thought 1".to_string()), + signature: None, + summary_id: None, + summary_text: None, + provider_type: None, + extra_data: None, + }], usage: None, raw_usage: None, raw_response: None, @@ -2749,7 +2757,7 @@ mod tests { let chunk = InferenceResultChunk::Json(JsonInferenceResultChunk { raw: Some(r#"{"key": "value"}"#.to_string()), - thought: None, + thought_chunks: vec![], usage: Some(Usage { input_tokens: Some(30), output_tokens: Some(20), diff --git a/tensorzero-core/src/inference/types/streams.rs b/tensorzero-core/src/inference/types/streams.rs index 097724681db..4ee9604cd30 100644 --- a/tensorzero-core/src/inference/types/streams.rs +++ b/tensorzero-core/src/inference/types/streams.rs @@ -87,55 +87,47 @@ pub struct UnknownChunk { pub provider_name: Option, } -#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)] +#[derive(Clone, Debug, Default, PartialEq)] pub struct ChatInferenceResultChunk { pub content: Vec, - #[serde(skip_serializing_if = "Option::is_none")] pub usage: Option, - #[serde(skip_serializing_if = "Option::is_none")] pub raw_usage: Option>, /// Raw responses from previous model inferences (e.g., best-of-n candidates). /// Used for artificial chunks emitted at stream start. - #[serde(skip_serializing_if = "Option::is_none")] pub raw_response: Option>, /// Time elapsed between making the request to the model provider and receiving this chunk. /// Important: this is NOT latency from the start of the TensorZero request. /// None for artificial chunks created by TensorZero (e.g., usage/finish_reason chunks). - #[serde(skip_serializing_if = "Option::is_none")] pub provider_latency: Option, /// Raw response string for the current chunk from the model provider. /// Empty for fake streams (non-streaming converted to streaming). - #[serde(skip_serializing_if = "String::is_empty")] pub raw_chunk: String, pub finish_reason: Option, } -#[derive(Clone, Debug, Default, PartialEq, Serialize)] +#[derive(Clone, Debug, Default, PartialEq)] pub struct JsonInferenceResultChunk { pub raw: Option, - pub thought: Option, - #[serde(skip_serializing_if = "Option::is_none")] + /// Thought content from reasoning models (e.g., Claude's extended thinking). + /// Populated during streaming from `ContentBlockChunk::Thought` variants. + /// Used internally to collect the final `Thought` block; not streamed directly to clients. + pub thought_chunks: Vec, pub usage: Option, - #[serde(skip_serializing_if = "Option::is_none")] pub raw_usage: Option>, /// Raw responses from previous model inferences (e.g., best-of-n candidates). /// Used for artificial chunks emitted at stream start. - #[serde(skip_serializing_if = "Option::is_none")] pub raw_response: Option>, /// Time elapsed between making the request to the model provider and receiving this chunk. /// Important: this is NOT latency from the start of the TensorZero request. /// None for artificial chunks created by TensorZero (e.g., usage/finish_reason chunks). - #[serde(skip_serializing_if = "Option::is_none")] pub provider_latency: Option, /// Raw response string for the current chunk from the model provider. /// Empty for fake streams (non-streaming converted to streaming). - #[serde(skip_serializing_if = "String::is_empty")] pub raw_chunk: String, pub finish_reason: Option, } -#[derive(Clone, Debug, PartialEq, Serialize)] -#[serde(tag = "type")] +#[derive(Clone, Debug, PartialEq)] pub enum InferenceResultChunk { Chat(ChatInferenceResultChunk), Json(JsonInferenceResultChunk), @@ -227,21 +219,15 @@ impl From for ChatInferenceResultChunk { impl From for JsonInferenceResultChunk { fn from(chunk: ProviderInferenceResponseChunk) -> Self { let mut raw = None; - let mut thought = None; - for content in &chunk.content { + let mut thought_chunks = Vec::new(); + for content in chunk.content { match content { ContentBlockChunk::ToolCall(tool_call) => { raw = Some(tool_call.raw_arguments.to_owned()); } ContentBlockChunk::Text(text_chunk) => raw = Some(text_chunk.text.to_owned()), ContentBlockChunk::Thought(thought_chunk) => { - // Take text if present, otherwise fall back to summary_text - // (OpenRouter's reasoning.summary type uses summary_text instead of text) - if thought_chunk.text.is_some() { - thought.clone_from(&thought_chunk.text); - } else { - thought.clone_from(&thought_chunk.summary_text); - } + thought_chunks.push(thought_chunk); } ContentBlockChunk::Unknown(_) => { // Unknown chunks are ignored for JSON functions @@ -251,7 +237,7 @@ impl From for JsonInferenceResultChunk { } Self { raw, - thought, + thought_chunks, usage: chunk.usage, raw_usage: chunk.raw_usage, raw_response: None, // Only populated via artificial chunks from TensorZero @@ -594,27 +580,124 @@ pub async fn collect_chunks(args: CollectChunksArgs) -> Result { - existing_thought - .text - .get_or_insert_default() - .push_str(&thought); + for thought_chunk in chunk.thought_chunks { + let ThoughtChunk { + id: _, + text, + signature, + summary_id, + summary_text, + provider_type, + extra_data, + } = thought_chunk; + + // Handle text if present + if let Some(text) = text { + match blocks.get_mut(&(ContentBlockOutputType::Thought, String::new())) { + Some(ContentBlockOutput::Thought(existing_thought)) => { + existing_thought + .text + .get_or_insert_default() + .push_str(&text); + } + _ => { + blocks.insert( + (ContentBlockOutputType::Thought, String::new()), + ContentBlockOutput::Thought(Thought { + text: Some(text), + summary: None, + signature: None, + provider_type: provider_type.clone(), + extra_data: extra_data.clone(), + }), + ); + } } - // If there is no thought block, create one - _ => { - blocks.insert( - (ContentBlockOutputType::Thought, String::new()), - ContentBlockOutput::Thought(Thought { - text: Some(thought), - summary: None, - signature: None, - provider_type: None, - extra_data: None, - }), - ); + } + + // Handle signature if present + if let Some(signature) = signature { + match blocks.get_mut(&(ContentBlockOutputType::Thought, String::new())) { + Some(ContentBlockOutput::Thought(existing_thought)) => { + match &mut existing_thought.signature { + Some(existing) => existing.push_str(&signature), + None => { + existing_thought.signature = Some(signature); + } + } + } + _ => { + blocks.insert( + (ContentBlockOutputType::Thought, String::new()), + ContentBlockOutput::Thought(Thought { + text: None, + summary: None, + signature: Some(signature), + provider_type: provider_type.clone(), + extra_data: extra_data.clone(), + }), + ); + } + } + } + + // Handle summary if present (summary_id and summary_text both required) + if summary_id.is_some() && summary_text.is_none() { + tracing::error!( + "Summary id is present but summary text is missing for JSON thought" + ); + } + if summary_id.is_none() && summary_text.is_some() { + tracing::error!( + "Summary text is present but summary id is missing for JSON thought" + ); + } + if let (Some(summary_id), Some(summary_text)) = (summary_id, summary_text) { + // For JSON mode, we use a single thought block with id "" + // so we track summary ids in thought_summaries with key "" + let summary_set = thought_summaries.entry(String::new()).or_default(); + let (index, _) = summary_set.insert_full(summary_id); + + match blocks.get_mut(&(ContentBlockOutputType::Thought, String::new())) { + Some(ContentBlockOutput::Thought(existing_thought)) => { + match &mut existing_thought.summary { + Some(existing) => { + if index >= existing.len() { + existing.resize( + index + 1, + ThoughtSummaryBlock::SummaryText { + text: String::new(), + }, + ); + } + match &mut existing[index] { + ThoughtSummaryBlock::SummaryText { text } => { + text.push_str(&summary_text); + } + } + } + None => { + existing_thought.summary = + Some(vec![ThoughtSummaryBlock::SummaryText { + text: summary_text, + }]); + } + } + } + _ => { + blocks.insert( + (ContentBlockOutputType::Thought, String::new()), + ContentBlockOutput::Thought(Thought { + text: None, + summary: Some(vec![ThoughtSummaryBlock::SummaryText { + text: summary_text, + }]), + signature: None, + provider_type, + extra_data, + }), + ); + } } } } @@ -1051,7 +1134,15 @@ mod tests { let chunks = vec![ InferenceResultChunk::Json(JsonInferenceResultChunk { raw: Some("{\"name\":".to_string()), - thought: Some("Thought 1".to_string()), + thought_chunks: vec![ThoughtChunk { + id: "0".to_string(), + text: Some("Thought 1".to_string()), + signature: None, + summary_id: None, + summary_text: None, + provider_type: None, + extra_data: None, + }], usage: Some(usage1), raw_usage: None, raw_response: None, @@ -1061,7 +1152,15 @@ mod tests { }), InferenceResultChunk::Json(JsonInferenceResultChunk { raw: Some("\"John\",\"age\":30}".to_string()), - thought: Some("Thought 2".to_string()), + thought_chunks: vec![ThoughtChunk { + id: "0".to_string(), + text: Some("Thought 2".to_string()), + signature: None, + summary_id: None, + summary_text: None, + provider_type: None, + extra_data: None, + }], usage: Some(usage2), raw_usage: None, raw_response: None, @@ -1140,7 +1239,15 @@ mod tests { let chunks = vec![ InferenceResultChunk::Json(JsonInferenceResultChunk { raw: Some("{\"name\":".to_string()), - thought: Some("Thought 1".to_string()), + thought_chunks: vec![ThoughtChunk { + id: "0".to_string(), + text: Some("Thought 1".to_string()), + signature: None, + summary_id: None, + summary_text: None, + provider_type: None, + extra_data: None, + }], usage: Some(model_inference_usage), raw_usage: None, raw_response: None, @@ -1150,7 +1257,7 @@ mod tests { }), InferenceResultChunk::Json(JsonInferenceResultChunk { raw: Some("\"John\"}".to_string()), - thought: None, + thought_chunks: vec![], usage: None, raw_usage: None, raw_response: None, @@ -1225,7 +1332,7 @@ mod tests { let chunks = vec![ InferenceResultChunk::Json(JsonInferenceResultChunk { raw: Some("{\"name\":\"John\",".to_string()), - thought: None, + thought_chunks: vec![], usage: Some(model_inference_usage), raw_usage: None, raw_response: None, @@ -1235,7 +1342,15 @@ mod tests { }), InferenceResultChunk::Json(JsonInferenceResultChunk { raw: Some(String::new()), - thought: Some("Thought 2".to_string()), + thought_chunks: vec![ThoughtChunk { + id: "0".to_string(), + text: Some("Thought 2".to_string()), + signature: None, + summary_id: None, + summary_text: None, + provider_type: None, + extra_data: None, + }], usage: None, raw_usage: None, raw_response: None, @@ -1245,7 +1360,7 @@ mod tests { }), InferenceResultChunk::Json(JsonInferenceResultChunk { raw: Some("\"age\":30}".to_string()), - thought: None, + thought_chunks: vec![], usage: None, raw_usage: None, raw_response: None, @@ -1357,7 +1472,15 @@ mod tests { let chunks = vec![ InferenceResultChunk::Json(JsonInferenceResultChunk { raw: Some("{\"name\":".to_string()), - thought: Some("Thought 1".to_string()), + thought_chunks: vec![ThoughtChunk { + id: "0".to_string(), + text: Some("Thought 1".to_string()), + signature: None, + summary_id: None, + summary_text: None, + provider_type: None, + extra_data: None, + }], usage: Some(usage1), raw_usage: None, raw_response: None, @@ -1367,7 +1490,15 @@ mod tests { }), InferenceResultChunk::Json(JsonInferenceResultChunk { raw: Some("\"John\",\"age\":30}".to_string()), - thought: Some("Thought 2".to_string()), + thought_chunks: vec![ThoughtChunk { + id: "0".to_string(), + text: Some("Thought 2".to_string()), + signature: None, + summary_id: None, + summary_text: None, + provider_type: None, + extra_data: None, + }], usage: Some(usage2), raw_usage: None, raw_response: None, @@ -1475,7 +1606,15 @@ mod tests { let chunks = vec![ InferenceResultChunk::Json(JsonInferenceResultChunk { raw: Some("{\"name\":".to_string()), - thought: Some("Thought 1".to_string()), + thought_chunks: vec![ThoughtChunk { + id: "0".to_string(), + text: Some("Thought 1".to_string()), + signature: None, + summary_id: None, + summary_text: None, + provider_type: None, + extra_data: None, + }], usage: Some(usage1), raw_usage: None, raw_response: None, @@ -1485,7 +1624,15 @@ mod tests { }), InferenceResultChunk::Json(JsonInferenceResultChunk { raw: Some("\"John\",\"age\":30}".to_string()), - thought: Some("Thought 2".to_string()), + thought_chunks: vec![ThoughtChunk { + id: "0".to_string(), + text: Some("Thought 2".to_string()), + signature: None, + summary_id: None, + summary_text: None, + provider_type: None, + extra_data: None, + }], usage: Some(usage2), raw_usage: None, raw_response: None, @@ -2366,7 +2513,7 @@ mod tests { let result = JsonInferenceResultChunk::from(tool_chunk); assert_eq!(result.raw, Some("{\"key\": \"value\"}".to_string())); - assert_eq!(result.thought, None); + assert!(result.thought_chunks.is_empty()); assert_eq!(result.raw_response, None); assert_eq!(result.raw_chunk, "raw response"); assert_eq!(result.provider_latency, Some(Duration::from_secs(1))); @@ -2393,7 +2540,7 @@ mod tests { let result = JsonInferenceResultChunk::from(text_chunk); assert_eq!(result.raw, Some("some text".to_string())); - assert_eq!(result.thought, None); + assert!(result.thought_chunks.is_empty()); // Test case for Thought content let thought_chunk = ProviderInferenceResponseChunk { @@ -2415,7 +2562,18 @@ mod tests { let result = JsonInferenceResultChunk::from(thought_chunk); assert_eq!(result.raw, None); - assert_eq!(result.thought, Some("thinking...".to_string())); + assert_eq!( + result.thought_chunks, + vec![ThoughtChunk { + id: "123".to_string(), + text: Some("thinking...".to_string()), + summary_id: None, + summary_text: None, + signature: None, + provider_type: None, + extra_data: None, + }] + ); assert_eq!(result.finish_reason, None); // Test case for multiple content blocks - should use last raw content let mixed_chunk = ProviderInferenceResponseChunk { @@ -2448,7 +2606,18 @@ mod tests { let result = JsonInferenceResultChunk::from(mixed_chunk); assert_eq!(result.raw, Some("final content".to_string())); - assert_eq!(result.thought, Some("final thought".to_string())); + assert_eq!( + result.thought_chunks, + vec![ThoughtChunk { + id: "789".to_string(), + text: Some("final thought".to_string()), + summary_id: None, + summary_text: None, + signature: None, + provider_type: None, + extra_data: None, + }] + ); // Test case for empty content let empty_chunk = ProviderInferenceResponseChunk { @@ -2462,7 +2631,7 @@ mod tests { let result = JsonInferenceResultChunk::from(empty_chunk); assert_eq!(result.raw, None); - assert_eq!(result.thought, None); + assert!(result.thought_chunks.is_empty()); assert_eq!(result.finish_reason, None); // Test case: ThoughtChunk with only summary_text (no text) @@ -2486,9 +2655,17 @@ mod tests { let result = JsonInferenceResultChunk::from(summary_chunk); assert_eq!( - result.thought, - Some("This is a summary".to_string()), - "summary_text should be extracted when text is None" + result.thought_chunks, + vec![ThoughtChunk { + id: "0".to_string(), + text: None, + summary_id: Some("0".to_string()), + summary_text: Some("This is a summary".to_string()), + signature: None, + provider_type: Some("openrouter".to_string()), + extra_data: None, + }], + "ThoughtChunk should be preserved when text is None" ); } } diff --git a/tensorzero-core/src/providers/openai/mod.rs b/tensorzero-core/src/providers/openai/mod.rs index c0d76e59388..2b852f02fc2 100644 --- a/tensorzero-core/src/providers/openai/mod.rs +++ b/tensorzero-core/src/providers/openai/mod.rs @@ -551,12 +551,11 @@ impl InferenceProvider for OpenAIProvider { let request_url = get_responses_url(self.api_base.as_ref().unwrap_or(&OPENAI_DEFAULT_BASE_URL))?; - // TODO - support encrypted reasoning in streaming let request_body = serde_json::to_value( OpenAIResponsesRequest::new( &self.model_name, request, - false, + self.include_encrypted_reasoning, &self.provider_tools, model_name, provider_name, diff --git a/tensorzero-core/src/providers/openai/responses.rs b/tensorzero-core/src/providers/openai/responses.rs index 67c3d06d103..8778ccaa3cd 100644 --- a/tensorzero-core/src/providers/openai/responses.rs +++ b/tensorzero-core/src/providers/openai/responses.rs @@ -1397,6 +1397,20 @@ pub(super) fn openai_responses_to_tensorzero_chunk( } else if item_type == "message" { // Skip message items - these are handled by other events return Ok(None); + } else if item_type == "reasoning" { + // Emit reasoning items as Unknown - final signature comes from response.completed + return Ok(Some(ProviderInferenceResponseChunk::new( + vec![ContentBlockChunk::Unknown(UnknownChunk { + id: output_index.to_string(), + data: item, + model_name: Some(model_name.to_string()), + provider_name: Some(provider_name.to_string()), + })], + None, + raw_message, + message_latency, + None, + ))); } else { // Unknown item type - return as unknown chunk return Ok(Some(ProviderInferenceResponseChunk::new( @@ -1417,7 +1431,7 @@ pub(super) fn openai_responses_to_tensorzero_chunk( Ok(None) } - // Completed event - extract usage and finish reason + // Completed event - extract usage, finish reason, and final reasoning signature OpenAIResponsesStreamEvent::ResponseCompleted { response } => { // Filter out null values to avoid creating RawUsageEntry with null data let usage_value = response.get("usage").filter(|v| !v.is_null()).cloned(); @@ -1453,9 +1467,35 @@ pub(super) fn openai_responses_to_tensorzero_chunk( Some(FinishReason::Stop) }; + // Extract final signature from reasoning items in output array + // This is the authoritative signature for multi-turn conversations + let mut content_blocks = vec![]; + if let Some(output) = response.get("output").and_then(|o| o.as_array()) { + for (index, item) in output.iter().enumerate() { + if let Some(item_type) = item.get("type").and_then(|t| t.as_str()) + && item_type == "reasoning" + { + let signature = item + .get("encrypted_content") + .and_then(|v| v.as_str()) + .map(|s| s.to_string()); + + content_blocks.push(ContentBlockChunk::Thought(ThoughtChunk { + id: index.to_string(), + text: None, + signature, + summary_id: None, + summary_text: None, + provider_type: Some(PROVIDER_TYPE.to_string()), + extra_data: None, + })); + } + } + } + Ok(Some(match raw_usage { Some(entries) => ProviderInferenceResponseChunk::new_with_raw_usage( - vec![], + content_blocks, usage, raw_message, message_latency, @@ -1463,7 +1503,7 @@ pub(super) fn openai_responses_to_tensorzero_chunk( Some(entries), ), None => ProviderInferenceResponseChunk::new( - vec![], + content_blocks, usage, raw_message, message_latency, @@ -1557,10 +1597,31 @@ pub(super) fn openai_responses_to_tensorzero_chunk( })) } + // Output item done - emit reasoning as Unknown, skip others + // Final signature comes from response.completed + OpenAIResponsesStreamEvent::ResponseOutputItemDone { item, output_index } => { + if let Some(item_type) = item.get("type").and_then(|t| t.as_str()) + && item_type == "reasoning" + { + return Ok(Some(ProviderInferenceResponseChunk::new( + vec![ContentBlockChunk::Unknown(UnknownChunk { + id: output_index.to_string(), + data: item, + model_name: Some(model_name.to_string()), + provider_name: Some(provider_name.to_string()), + })], + None, + raw_message, + message_latency, + None, + ))); + } + Ok(None) + } + // Lifecycle and other events we don't need to process OpenAIResponsesStreamEvent::ResponseCreated { .. } | OpenAIResponsesStreamEvent::ResponseInProgress { .. } - | OpenAIResponsesStreamEvent::ResponseOutputItemDone { .. } | OpenAIResponsesStreamEvent::ResponseContentPartAdded { .. } | OpenAIResponsesStreamEvent::ResponseContentPartDone { .. } | OpenAIResponsesStreamEvent::ResponseOutputTextDone { .. } @@ -1952,6 +2013,194 @@ mod tests { assert_eq!(tool_name, Some("get_weather".to_string())); } + #[test] + fn test_output_item_added_reasoning_emits_unknown() { + // Reasoning items in .added event are emitted as Unknown + // Final signature comes from response.completed + let item_json = serde_json::json!({ + "type": "reasoning", + "id": "rs_abc123", + "encrypted_content": "encrypted_thought_signature_data", + "summary": [] + }); + + let event = OpenAIResponsesStreamEvent::ResponseOutputItemAdded { + item: item_json.clone(), + output_index: 0, + }; + + let mut tool_id = None; + let mut tool_name = None; + + let result = openai_responses_to_tensorzero_chunk( + "raw_json".to_string(), + event, + Duration::from_millis(100), + &mut tool_id, + &mut tool_name, + false, + "test_model", + "test_provider", + "", + Uuid::now_v7(), + PROVIDER_TYPE, + ) + .unwrap() + .unwrap(); + + // Should emit Unknown chunk + assert_eq!( + result.content.len(), + 1, + "Expected exactly one content block" + ); + match &result.content[0] { + ContentBlockChunk::Unknown(unknown) => { + assert_eq!(unknown.id, "0", "Expected output_index as id"); + assert_eq!(unknown.data, item_json, "Expected item data"); + } + _ => panic!("Expected Unknown chunk"), + } + } + + #[test] + fn test_output_item_done_reasoning_emits_unknown() { + // Test that ResponseOutputItemDone with reasoning type emits Unknown + // Final signature comes from response.completed + let item_json = serde_json::json!({ + "type": "reasoning", + "id": "rs_done_123", + "encrypted_content": "encrypted_signature_from_done_event", + "summary": [{"type": "summary_text", "text": "Some reasoning summary"}] + }); + + let event = OpenAIResponsesStreamEvent::ResponseOutputItemDone { + item: item_json.clone(), + output_index: 1, + }; + + let mut tool_id = None; + let mut tool_name = None; + + let result = openai_responses_to_tensorzero_chunk( + "raw_json".to_string(), + event, + Duration::from_millis(100), + &mut tool_id, + &mut tool_name, + false, + "test_model", + "test_provider", + "", + Uuid::now_v7(), + PROVIDER_TYPE, + ) + .unwrap() + .unwrap(); + + // Should emit Unknown chunk + assert_eq!( + result.content.len(), + 1, + "Expected exactly one content block" + ); + match &result.content[0] { + ContentBlockChunk::Unknown(unknown) => { + assert_eq!(unknown.id, "1", "Expected output_index as id"); + assert_eq!(unknown.data, item_json, "Expected item data"); + } + _ => panic!("Expected Unknown chunk, got {:?}", result.content[0]), + } + } + + #[test] + fn test_output_item_done_reasoning_without_encrypted_content_emits_unknown() { + // Test that ResponseOutputItemDone with reasoning but no encrypted_content emits Unknown + let item_json = serde_json::json!({ + "type": "reasoning", + "id": "rs_done_456", + "summary": [{"type": "summary_text", "text": "Summary only"}] + }); + + let event = OpenAIResponsesStreamEvent::ResponseOutputItemDone { + item: item_json.clone(), + output_index: 0, + }; + + let mut tool_id = None; + let mut tool_name = None; + + let result = openai_responses_to_tensorzero_chunk( + "raw_json".to_string(), + event, + Duration::from_millis(100), + &mut tool_id, + &mut tool_name, + false, + "test_model", + "test_provider", + "", + Uuid::now_v7(), + PROVIDER_TYPE, + ) + .unwrap() + .unwrap(); + + // Should emit Unknown chunk even without encrypted_content + assert_eq!( + result.content.len(), + 1, + "Expected exactly one content block" + ); + match &result.content[0] { + ContentBlockChunk::Unknown(unknown) => { + assert_eq!(unknown.id, "0", "Expected output_index as id"); + assert_eq!(unknown.data, item_json, "Expected item data"); + } + _ => panic!("Expected Unknown chunk"), + } + } + + #[test] + fn test_output_item_done_non_reasoning_returns_none() { + // Test that ResponseOutputItemDone with non-reasoning type returns None + let item_json = serde_json::json!({ + "type": "function_call", + "id": "fc_done_789", + "name": "some_function", + "arguments": "{}" + }); + + let event = OpenAIResponsesStreamEvent::ResponseOutputItemDone { + item: item_json, + output_index: 0, + }; + + let mut tool_id = None; + let mut tool_name = None; + + let result = openai_responses_to_tensorzero_chunk( + "raw_json".to_string(), + event, + Duration::from_millis(100), + &mut tool_id, + &mut tool_name, + false, + "test_model", + "test_provider", + "", + Uuid::now_v7(), + PROVIDER_TYPE, + ) + .unwrap(); + + // Should return None for non-reasoning items + assert!( + result.is_none(), + "Expected None for non-reasoning item.done event" + ); + } + #[test] fn test_function_call_done_conversion() { let event = OpenAIResponsesStreamEvent::ResponseFunctionCallArgumentsDone { @@ -2269,6 +2518,141 @@ mod tests { assert_eq!(result.finish_reason, Some(FinishReason::Stop)); } + #[test] + fn test_response_completed_with_reasoning_signature() { + // Test that ResponseCompleted extracts final signature from reasoning items + let usage_json = serde_json::json!({ + "input_tokens": 100, + "output_tokens": 200 + }); + let response_json = serde_json::json!({ + "id": "resp_123", + "status": "completed", + "usage": usage_json, + "output": [ + { + "type": "reasoning", + "id": "rs_final", + "encrypted_content": "final_authoritative_signature", + "summary": [{"type": "summary_text", "text": "Reasoning summary"}] + }, + { + "type": "message", + "id": "msg_123", + "content": [{"type": "output_text", "text": "Hello"}] + } + ] + }); + + let event = OpenAIResponsesStreamEvent::ResponseCompleted { + response: response_json, + }; + + let mut tool_id = None; + let mut tool_name = None; + + let model_inference_id = Uuid::now_v7(); + let result = openai_responses_to_tensorzero_chunk( + "raw_json".to_string(), + event, + Duration::from_millis(100), + &mut tool_id, + &mut tool_name, + false, + "test_model", + "test_provider", + "", + model_inference_id, + PROVIDER_TYPE, + ) + .unwrap() + .unwrap(); + + // Should have one ThoughtChunk with the final signature + assert_eq!( + result.content.len(), + 1, + "Expected one content block for reasoning item" + ); + match &result.content[0] { + ContentBlockChunk::Thought(thought_chunk) => { + assert_eq!( + thought_chunk.id, "0", + "Expected index 0 for first output item" + ); + assert_eq!( + thought_chunk.signature, + Some("final_authoritative_signature".to_string()), + "Expected final signature from response.completed" + ); + assert_eq!(thought_chunk.text, None); + assert_eq!(thought_chunk.summary_text, None); + } + _ => panic!("Expected Thought chunk, got {:?}", result.content[0]), + } + + // Should also have usage and finish_reason + assert_eq!( + result.usage, + Some(Usage { + input_tokens: Some(100), + output_tokens: Some(200), + }) + ); + assert_eq!(result.finish_reason, Some(FinishReason::Stop)); + } + + #[test] + fn test_response_completed_reasoning_without_encrypted_content() { + // Test that ResponseCompleted emits ThoughtChunk even without encrypted_content + let response_json = serde_json::json!({ + "id": "resp_123", + "status": "completed", + "output": [ + { + "type": "reasoning", + "id": "rs_no_sig", + "summary": [{"type": "summary_text", "text": "Summary only"}] + } + ] + }); + + let event = OpenAIResponsesStreamEvent::ResponseCompleted { + response: response_json, + }; + + let mut tool_id = None; + let mut tool_name = None; + + let result = openai_responses_to_tensorzero_chunk( + "raw_json".to_string(), + event, + Duration::from_millis(100), + &mut tool_id, + &mut tool_name, + false, + "test_model", + "test_provider", + "", + Uuid::now_v7(), + PROVIDER_TYPE, + ) + .unwrap() + .unwrap(); + + // Should have ThoughtChunk with None signature + assert_eq!(result.content.len(), 1); + match &result.content[0] { + ContentBlockChunk::Thought(thought_chunk) => { + assert_eq!( + thought_chunk.signature, None, + "Expected None signature when encrypted_content is missing" + ); + } + _ => panic!("Expected Thought chunk"), + } + } + #[test] fn test_response_incomplete_conversion() { let usage_json = serde_json::json!({ diff --git a/tensorzero-core/src/variant/mixture_of_n.rs b/tensorzero-core/src/variant/mixture_of_n.rs index 2a2808a6952..bd3f4e00397 100644 --- a/tensorzero-core/src/variant/mixture_of_n.rs +++ b/tensorzero-core/src/variant/mixture_of_n.rs @@ -452,7 +452,7 @@ fn make_stream_from_non_stream( } InferenceResult::Json(json) => Ok(InferenceResultChunk::Json(JsonInferenceResultChunk { raw: json.output.raw, - thought: None, + thought_chunks: Vec::new(), usage, raw_usage: raw_usage_entries, raw_response: None, // Not used for fused stream chunks diff --git a/tensorzero-core/tests/e2e/providers/common.rs b/tensorzero-core/tests/e2e/providers/common.rs index 591e6a113f5..06c98ac6393 100644 --- a/tensorzero-core/tests/e2e/providers/common.rs +++ b/tensorzero-core/tests/e2e/providers/common.rs @@ -166,10 +166,10 @@ macro_rules! generate_provider_tests { use $crate::providers::common::test_extra_body_with_provider; use $crate::providers::common::test_inference_extra_body_with_provider; use $crate::providers::common::test_assistant_prefill_inference_request_with_provider; - use $crate::providers::reasoning::test_reasoning_inference_request_simple_with_provider; - use $crate::providers::reasoning::test_streaming_reasoning_inference_request_simple_with_provider; - use $crate::providers::reasoning::test_reasoning_inference_request_with_provider_json_mode; - use $crate::providers::reasoning::test_streaming_reasoning_inference_request_with_provider_json_mode; + use $crate::providers::reasoning::test_reasoning_inference_request_simple_nonstreaming_with_provider; + use $crate::providers::reasoning::test_reasoning_inference_request_simple_streaming_with_provider; + use $crate::providers::reasoning::test_reasoning_inference_request_json_mode_nonstreaming_with_provider; + use $crate::providers::reasoning::test_reasoning_inference_request_json_mode_streaming_with_provider; use $crate::providers::common::test_short_inference_request_with_provider; use $crate::providers::common::test_multi_turn_parallel_tool_use_inference_request_with_provider; use $crate::providers::common::test_multi_turn_parallel_tool_use_streaming_inference_request_with_provider; @@ -223,18 +223,18 @@ macro_rules! generate_provider_tests { } #[tokio::test] - async fn test_reasoning_inference_request_simple() { + async fn test_reasoning_inference_request_simple_nonstreaming() { let providers = $func().await.reasoning_inference; for provider in providers { - test_reasoning_inference_request_simple_with_provider(provider).await; + test_reasoning_inference_request_simple_nonstreaming_with_provider(provider).await; } } #[tokio::test] - async fn test_streaming_reasoning_inference_request_simple() { + async fn test_reasoning_inference_request_simple_streaming() { let providers = $func().await.reasoning_inference; for provider in providers { - test_streaming_reasoning_inference_request_simple_with_provider(provider).await; + test_reasoning_inference_request_simple_streaming_with_provider(provider).await; } } @@ -491,18 +491,18 @@ macro_rules! generate_provider_tests { } #[tokio::test] - async fn test_reasoning_inference_request_json_mode() { + async fn test_reasoning_inference_request_json_mode_nonstreaming() { let providers = $func().await.reasoning_inference; for provider in providers { - test_reasoning_inference_request_with_provider_json_mode(provider).await; + test_reasoning_inference_request_json_mode_nonstreaming_with_provider(provider).await; } } #[tokio::test] - async fn test_streaming_reasoning_inference_request_json_mode() { + async fn test_reasoning_inference_request_json_mode_streaming() { let providers = $func().await.reasoning_inference; for provider in providers { - test_streaming_reasoning_inference_request_with_provider_json_mode(provider).await; + test_reasoning_inference_request_json_mode_streaming_with_provider(provider).await; } } diff --git a/tensorzero-core/tests/e2e/providers/openai/mod.rs b/tensorzero-core/tests/e2e/providers/openai/mod.rs index 8fb346ab304..28107f864c3 100644 --- a/tensorzero-core/tests/e2e/providers/openai/mod.rs +++ b/tensorzero-core/tests/e2e/providers/openai/mod.rs @@ -258,14 +258,13 @@ async fn get_providers() -> E2ETestProviders { use_modal_headers: false, }]; - // TODO (#5717): we don't parse streaming thought signatures correctly for OpenAI Responses API - // let reasoning_providers = vec![E2ETestProvider { - // supports_batch_inference: false, - // variant_name: "openai-gpt-5-mini".to_string(), - // model_name: "gpt-5-mini-responses".into(), - // model_provider_name: "openai".into(), - // credentials: HashMap::new(), - // }]; + let reasoning_providers = vec![E2ETestProvider { + supports_batch_inference: false, + variant_name: "openai-gpt-5-mini".to_string(), + model_name: "gpt-5-mini-responses".into(), + model_provider_name: "openai".into(), + credentials: HashMap::new(), + }]; let reasoning_usage_providers = vec![ E2ETestProvider { @@ -288,9 +287,7 @@ async fn get_providers() -> E2ETestProviders { simple_inference: standard_providers.clone(), extra_body_inference: extra_body_providers, bad_auth_extra_headers, - // TODO (#5717): we don't parse streaming thought signatures correctly for OpenAI Responses API - // reasoning_inference: reasoning_providers, - reasoning_inference: vec![], + reasoning_inference: reasoning_providers, reasoning_usage_inference: reasoning_usage_providers, cache_input_tokens_inference: standard_providers.clone(), embeddings: embedding_providers, diff --git a/tensorzero-core/tests/e2e/providers/reasoning.rs b/tensorzero-core/tests/e2e/providers/reasoning.rs index c80a446b588..25983ba1d2f 100644 --- a/tensorzero-core/tests/e2e/providers/reasoning.rs +++ b/tensorzero-core/tests/e2e/providers/reasoning.rs @@ -19,7 +19,9 @@ use tensorzero_core::inference::types::extra_headers::UnfilteredInferenceExtraHe use tensorzero_core::inference::types::{StoredContentBlock, StoredRequestMessage, Text}; use uuid::Uuid; -pub async fn test_reasoning_inference_request_simple_with_provider(provider: E2ETestProvider) { +pub async fn test_reasoning_inference_request_simple_nonstreaming_with_provider( + provider: E2ETestProvider, +) { let episode_id = Uuid::now_v7(); let extra_headers = if provider.is_modal_provider() { get_modal_extra_headers() @@ -85,7 +87,9 @@ pub async fn test_reasoning_inference_request_simple_with_provider(provider: E2E "thought" => { found_thought = true; } - _ => panic!("Unexpected content block type: {block_type}"), + _ => { + // Skip unknown content block types (e.g., raw reasoning data from OpenAI Responses API) + } } } @@ -156,7 +160,9 @@ pub async fn test_reasoning_inference_request_simple_with_provider(provider: E2E "thought" => { found_thought = true; } - _ => panic!("Unexpected content block type: {block_type}"), + _ => { + // Skip unknown content block types (e.g., raw reasoning data from OpenAI Responses API) + } } } @@ -270,7 +276,7 @@ pub async fn test_reasoning_inference_request_simple_with_provider(provider: E2E assert_eq!(id, inference_id); } -pub async fn test_streaming_reasoning_inference_request_simple_with_provider( +pub async fn test_reasoning_inference_request_simple_streaming_with_provider( provider: E2ETestProvider, ) { use reqwest_eventsource::{Event, RequestBuilderExt}; @@ -453,7 +459,9 @@ pub async fn test_streaming_reasoning_inference_request_simple_with_provider( .push_str(thought_text); } } - _ => panic!("Unexpected content block type: {block_type}"), + _ => { + // Skip unknown content block types (e.g., raw reasoning data from OpenAI Responses API) + } } } @@ -581,7 +589,9 @@ pub async fn test_streaming_reasoning_inference_request_simple_with_provider( assert_eq!(id, inference_id); } -pub async fn test_reasoning_inference_request_with_provider_json_mode(provider: E2ETestProvider) { +pub async fn test_reasoning_inference_request_json_mode_nonstreaming_with_provider( + provider: E2ETestProvider, +) { // Direct Anthropic uses output_format for json_mode=strict // AWS Bedrock and GCP Vertex Anthropic use json_mode=off (prompt-based JSON) to avoid prefill conflicts @@ -801,7 +811,7 @@ pub async fn test_reasoning_inference_request_with_provider_json_mode(provider: ); } -pub async fn test_streaming_reasoning_inference_request_with_provider_json_mode( +pub async fn test_reasoning_inference_request_json_mode_streaming_with_provider( provider: E2ETestProvider, ) { // OpenAI O1 doesn't support streaming responses @@ -1079,12 +1089,16 @@ pub async fn test_streaming_reasoning_inference_request_with_provider_json_mode( StoredContentBlock::Thought(thought) => thought, _ => panic!("Expected a thought block"), }; - assert!( - thought - .text - .as_ref() - .unwrap() - .to_lowercase() - .contains("tokyo") - ); + // If text is present, check it contains "tokyo"; otherwise ensure signature exists + if let Some(text) = &thought.text { + assert!( + text.to_lowercase().contains("tokyo"), + "Expected thought text to contain 'tokyo', got: {text}" + ); + } else { + assert!( + thought.signature.is_some(), + "Expected either thought text or signature to be present" + ); + } }