From b80bfe96501dbc3a074993faa5d12581505b20cb Mon Sep 17 00:00:00 2001 From: Shuyang Li Date: Tue, 16 Dec 2025 17:23:34 -0500 Subject: [PATCH] Migrate countWorkflowEvaluationProjects --- gateway/src/routes/internal.rs | 10 +++- ...tWorkflowEvaluationProjectCountResponse.ts | 6 ++ .../tensorzero-node/lib/bindings/index.ts | 1 + .../clickhouse/workflow_evaluation_queries.rs | 59 ++++++++++++++++++- .../src/db/workflow_evaluation_queries.rs | 3 + .../internal/get_project_count.rs | 51 ++++++++++++++++ .../internal/get_projects.rs | 24 ++++++-- .../workflow_evaluations/internal/mod.rs | 3 + .../workflow_evaluations/internal/types.rs | 24 ++++---- .../e2e/db/workflow_evaluation_queries.rs | 16 +++++ .../internal/workflow_evaluations.rs | 24 +++++++- ui/app/routes/index.tsx | 7 +-- ui/app/routes/workflow_evaluations/route.tsx | 9 +-- .../clickhouse/workflow_evaluations.server.ts | 15 ----- .../clickhouse/workflow_evaluations.test.ts | 8 --- ui/app/utils/tensorzero/tensorzero.ts | 20 +++++++ 16 files changed, 222 insertions(+), 58 deletions(-) create mode 100644 internal/tensorzero-node/lib/bindings/GetWorkflowEvaluationProjectCountResponse.ts create mode 100644 tensorzero-core/src/endpoints/workflow_evaluations/internal/get_project_count.rs diff --git a/gateway/src/routes/internal.rs b/gateway/src/routes/internal.rs index 63802fddbd..03657eb3a0 100644 --- a/gateway/src/routes/internal.rs +++ b/gateway/src/routes/internal.rs @@ -95,6 +95,10 @@ pub fn build_internal_non_otel_enabled_routes() -> Router { "/internal/evaluations/run-stats", get(endpoints::internal::evaluations::get_evaluation_run_stats_handler), ) + .route( + "/internal/evaluations/datapoint-count", + get(endpoints::internal::evaluations::count_datapoints_handler), + ) .route( "/internal/evaluations/runs", get(endpoints::internal::evaluations::list_evaluation_runs_handler), @@ -105,8 +109,10 @@ pub fn build_internal_non_otel_enabled_routes() -> Router { get(endpoints::workflow_evaluations::internal::get_workflow_evaluation_projects_handler), ) .route( - "/internal/evaluations/datapoint-count", - get(endpoints::internal::evaluations::count_datapoints_handler), + "/internal/workflow-evaluations/projects/count", + get( + endpoints::workflow_evaluations::internal::get_workflow_evaluation_project_count_handler, + ), ) .route( "/internal/models/usage", diff --git a/internal/tensorzero-node/lib/bindings/GetWorkflowEvaluationProjectCountResponse.ts b/internal/tensorzero-node/lib/bindings/GetWorkflowEvaluationProjectCountResponse.ts new file mode 100644 index 0000000000..ac4a1cf257 --- /dev/null +++ b/internal/tensorzero-node/lib/bindings/GetWorkflowEvaluationProjectCountResponse.ts @@ -0,0 +1,6 @@ +// This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually. + +/** + * Response containing the count of workflow evaluation projects. + */ +export type GetWorkflowEvaluationProjectCountResponse = { count: number }; diff --git a/internal/tensorzero-node/lib/bindings/index.ts b/internal/tensorzero-node/lib/bindings/index.ts index 649cd47832..36f143371b 100644 --- a/internal/tensorzero-node/lib/bindings/index.ts +++ b/internal/tensorzero-node/lib/bindings/index.ts @@ -123,6 +123,7 @@ export * from "./GetInferencesResponse"; export * from "./GetModelInferencesResponse"; export * from "./GetModelLatencyResponse"; export * from "./GetModelUsageResponse"; +export * from "./GetWorkflowEvaluationProjectCountResponse"; export * from "./GetWorkflowEvaluationProjectsResponse"; export * from "./GoogleAIStudioGeminiProvider"; export * from "./GroqProvider"; diff --git a/tensorzero-core/src/db/clickhouse/workflow_evaluation_queries.rs b/tensorzero-core/src/db/clickhouse/workflow_evaluation_queries.rs index a08dcfeb1d..06267db148 100644 --- a/tensorzero-core/src/db/clickhouse/workflow_evaluation_queries.rs +++ b/tensorzero-core/src/db/clickhouse/workflow_evaluation_queries.rs @@ -5,10 +5,10 @@ use std::collections::HashMap; use async_trait::async_trait; use super::ClickHouseConnectionInfo; -use super::select_queries::parse_json_rows; +use super::select_queries::{parse_count, parse_json_rows}; use crate::db::workflow_evaluation_queries::WorkflowEvaluationProjectRow; use crate::db::workflow_evaluation_queries::WorkflowEvaluationQueries; -use crate::error::Error; +use crate::error::{Error, ErrorDetails}; #[async_trait] impl WorkflowEvaluationQueries for ClickHouseConnectionInfo { @@ -41,6 +41,26 @@ impl WorkflowEvaluationQueries for ClickHouseConnectionInfo { parse_json_rows(response.response.as_str()) } + + async fn count_workflow_evaluation_projects(&self) -> Result { + let query = r" + SELECT + toUInt32(countDistinct(project_name)) as count + FROM DynamicEvaluationRunByProjectName + WHERE project_name IS NOT NULL + FORMAT JSONEachRow + " + .to_string(); + + let response = self.run_query_synchronous_no_params(query).await?; + let count = parse_count(response.response.as_str())?; + + u32::try_from(count).map_err(|error| { + Error::new(ErrorDetails::ClickHouseDeserialization { + message: format!("Failed to convert workflow evaluation project count: {error}"), + }) + }) + } } #[cfg(test)] @@ -160,4 +180,39 @@ mod tests { assert_eq!(result[1].name, "project2"); assert_eq!(result[2].name, "project3"); } + + #[tokio::test] + async fn test_count_workflow_evaluation_projects() { + let mut mock_clickhouse_client = MockClickHouseClient::new(); + + mock_clickhouse_client + .expect_run_query_synchronous() + .withf(|query, params| { + assert_query_contains( + query, + " + SELECT toUInt32(countDistinct(project_name)) as count + FROM DynamicEvaluationRunByProjectName + WHERE project_name IS NOT NULL + FORMAT JSONEachRow", + ); + assert!(params.is_empty()); + true + }) + .returning(|_, _| { + Ok(ClickHouseResponse { + response: r#"{"count":2}"#.to_string(), + metadata: ClickHouseResponseMetadata { + read_rows: 1, + written_rows: 0, + }, + }) + }); + + let conn = ClickHouseConnectionInfo::new_mock(Arc::new(mock_clickhouse_client)); + + let count = conn.count_workflow_evaluation_projects().await.unwrap(); + + assert_eq!(count, 2); + } } diff --git a/tensorzero-core/src/db/workflow_evaluation_queries.rs b/tensorzero-core/src/db/workflow_evaluation_queries.rs index 6e352d26d1..c14e3ab3d6 100644 --- a/tensorzero-core/src/db/workflow_evaluation_queries.rs +++ b/tensorzero-core/src/db/workflow_evaluation_queries.rs @@ -27,4 +27,7 @@ pub trait WorkflowEvaluationQueries { limit: u32, offset: u32, ) -> Result, Error>; + + /// Counts workflow evaluation projects. + async fn count_workflow_evaluation_projects(&self) -> Result; } diff --git a/tensorzero-core/src/endpoints/workflow_evaluations/internal/get_project_count.rs b/tensorzero-core/src/endpoints/workflow_evaluations/internal/get_project_count.rs new file mode 100644 index 0000000000..2fa62a00d5 --- /dev/null +++ b/tensorzero-core/src/endpoints/workflow_evaluations/internal/get_project_count.rs @@ -0,0 +1,51 @@ +//! Handler for getting the workflow evaluation project count. + +use axum::Json; +use axum::extract::State; +use tracing::instrument; + +use super::types::GetWorkflowEvaluationProjectCountResponse; +use crate::db::workflow_evaluation_queries::WorkflowEvaluationQueries; +use crate::error::Error; +use crate::utils::gateway::{AppState, AppStateData}; + +/// Handler for `GET /internal/workflow-evaluations/projects/count` +#[axum::debug_handler(state = AppStateData)] +#[instrument(name = "workflow_evaluations.count_projects", skip_all)] +pub async fn get_workflow_evaluation_project_count_handler( + State(app_state): AppState, +) -> Result, Error> { + let response = + get_workflow_evaluation_project_count(&app_state.clickhouse_connection_info).await?; + + Ok(Json(response)) +} + +/// Core business logic for retrieving the workflow evaluation project count. +pub async fn get_workflow_evaluation_project_count( + clickhouse: &impl WorkflowEvaluationQueries, +) -> Result { + let count = clickhouse.count_workflow_evaluation_projects().await?; + Ok(GetWorkflowEvaluationProjectCountResponse { count }) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::db::workflow_evaluation_queries::MockWorkflowEvaluationQueries; + + #[tokio::test] + async fn test_get_workflow_evaluation_project_count() { + let mut mock_clickhouse = MockWorkflowEvaluationQueries::new(); + mock_clickhouse + .expect_count_workflow_evaluation_projects() + .times(1) + .returning(|| Box::pin(async move { Ok(7) })); + + let response = get_workflow_evaluation_project_count(&mock_clickhouse) + .await + .unwrap(); + + assert_eq!(response.count, 7); + } +} diff --git a/tensorzero-core/src/endpoints/workflow_evaluations/internal/get_projects.rs b/tensorzero-core/src/endpoints/workflow_evaluations/internal/get_projects.rs index bf3936d5bc..28a573c564 100644 --- a/tensorzero-core/src/endpoints/workflow_evaluations/internal/get_projects.rs +++ b/tensorzero-core/src/endpoints/workflow_evaluations/internal/get_projects.rs @@ -2,16 +2,24 @@ use axum::Json; use axum::extract::{Query, State}; +use serde::Deserialize; use tracing::instrument; -use super::types::{ - GetWorkflowEvaluationProjectsParams, GetWorkflowEvaluationProjectsResponse, - WorkflowEvaluationProject, -}; +use super::types::{GetWorkflowEvaluationProjectsResponse, WorkflowEvaluationProject}; use crate::db::workflow_evaluation_queries::WorkflowEvaluationQueries; use crate::error::Error; use crate::utils::gateway::{AppState, AppStateData}; +/// Query parameters for getting workflow evaluation projects. +#[derive(Debug, Deserialize)] +pub struct GetWorkflowEvaluationProjectsParams { + pub limit: Option, + pub offset: Option, +} + +const DEFAULT_GET_WORKFLOW_EVALUATION_PROJECTS_LIMIT: u32 = 100; +const DEFAULT_GET_WORKFLOW_EVALUATION_PROJECTS_OFFSET: u32 = 0; + /// Handler for `GET /internal/workflow-evaluations/projects` /// /// Returns a paginated list of workflow evaluation projects. @@ -23,8 +31,12 @@ pub async fn get_workflow_evaluation_projects_handler( ) -> Result, Error> { let response = get_workflow_evaluation_projects( &app_state.clickhouse_connection_info, - params.limit, - params.offset, + params + .limit + .unwrap_or(DEFAULT_GET_WORKFLOW_EVALUATION_PROJECTS_LIMIT), + params + .offset + .unwrap_or(DEFAULT_GET_WORKFLOW_EVALUATION_PROJECTS_OFFSET), ) .await?; diff --git a/tensorzero-core/src/endpoints/workflow_evaluations/internal/mod.rs b/tensorzero-core/src/endpoints/workflow_evaluations/internal/mod.rs index fda9050fef..9eec7625cf 100644 --- a/tensorzero-core/src/endpoints/workflow_evaluations/internal/mod.rs +++ b/tensorzero-core/src/endpoints/workflow_evaluations/internal/mod.rs @@ -2,8 +2,11 @@ //! //! These endpoints support the UI for viewing and managing workflow evaluation runs and results. +mod get_project_count; mod get_projects; mod types; +pub use get_project_count::get_workflow_evaluation_project_count; +pub use get_project_count::get_workflow_evaluation_project_count_handler; pub use get_projects::get_workflow_evaluation_projects_handler; pub use types::*; diff --git a/tensorzero-core/src/endpoints/workflow_evaluations/internal/types.rs b/tensorzero-core/src/endpoints/workflow_evaluations/internal/types.rs index 60b13c54fe..4e96efe730 100644 --- a/tensorzero-core/src/endpoints/workflow_evaluations/internal/types.rs +++ b/tensorzero-core/src/endpoints/workflow_evaluations/internal/types.rs @@ -7,19 +7,6 @@ use serde::{Deserialize, Serialize}; // Get Workflow Evaluation Projects // ============================================================================= -/// Query parameters for getting workflow evaluation projects. -#[derive(Debug, Deserialize)] -pub struct GetWorkflowEvaluationProjectsParams { - #[serde(default = "default_limit")] - pub limit: u32, - #[serde(default)] - pub offset: u32, -} - -fn default_limit() -> u32 { - 100 -} - /// Response containing a list of workflow evaluation projects. #[derive(Debug, Serialize, Deserialize, ts_rs::TS)] #[ts(export)] @@ -35,3 +22,14 @@ pub struct WorkflowEvaluationProject { pub count: u32, pub last_updated: DateTime, } + +// ============================================================================= +// Get Workflow Evaluation Project Count +// ============================================================================= + +/// Response containing the count of workflow evaluation projects. +#[derive(Debug, Serialize, Deserialize, ts_rs::TS)] +#[ts(export)] +pub struct GetWorkflowEvaluationProjectCountResponse { + pub count: u32, +} diff --git a/tensorzero-core/tests/e2e/db/workflow_evaluation_queries.rs b/tensorzero-core/tests/e2e/db/workflow_evaluation_queries.rs index 7b93abe0f1..a012f0efb4 100644 --- a/tensorzero-core/tests/e2e/db/workflow_evaluation_queries.rs +++ b/tensorzero-core/tests/e2e/db/workflow_evaluation_queries.rs @@ -38,3 +38,19 @@ async fn test_list_workflow_evaluation_projects_with_fixture_data() { "Expected at least one of the fixture projects to be present. Found: {project_names:?}", ); } + +/// Ensures workflow evaluation project counts are returned from ClickHouse. +#[tokio::test] +async fn test_count_workflow_evaluation_projects_with_fixture_data() { + let clickhouse = get_clickhouse().await; + + let result = clickhouse + .count_workflow_evaluation_projects() + .await + .unwrap(); + + assert!( + result >= 2, + "Expected at least 2 workflow evaluation projects from fixtures, got {result}", + ); +} diff --git a/tensorzero-core/tests/e2e/endpoints/internal/workflow_evaluations.rs b/tensorzero-core/tests/e2e/endpoints/internal/workflow_evaluations.rs index b9ea00a0ba..081dd55179 100644 --- a/tensorzero-core/tests/e2e/endpoints/internal/workflow_evaluations.rs +++ b/tensorzero-core/tests/e2e/endpoints/internal/workflow_evaluations.rs @@ -1,7 +1,9 @@ //! E2E tests for the workflow evaluation endpoints. use reqwest::Client; -use tensorzero_core::endpoints::workflow_evaluations::internal::GetWorkflowEvaluationProjectsResponse; +use tensorzero_core::endpoints::workflow_evaluations::internal::{ + GetWorkflowEvaluationProjectCountResponse, GetWorkflowEvaluationProjectsResponse, +}; use crate::common::get_gateway_endpoint; @@ -58,3 +60,23 @@ async fn test_get_workflow_evaluation_projects_with_pagination() { response.projects.len() ); } + +#[tokio::test(flavor = "multi_thread")] +async fn test_get_workflow_evaluation_project_count_endpoint() { + let http_client = Client::new(); + let url = get_gateway_endpoint("/internal/workflow-evaluations/projects/count"); + + let resp = http_client.get(url).send().await.unwrap(); + assert!( + resp.status().is_success(), + "get_workflow_evaluation_project_count request failed: status={:?}", + resp.status() + ); + + let response: GetWorkflowEvaluationProjectCountResponse = resp.json().await.unwrap(); + + assert!( + response.count > 0, + "Expected workflow evaluation project count to be greater than 0" + ); +} diff --git a/ui/app/routes/index.tsx b/ui/app/routes/index.tsx index 2bfd32145d..f4f4314ca7 100644 --- a/ui/app/routes/index.tsx +++ b/ui/app/routes/index.tsx @@ -22,10 +22,7 @@ import { import { countInferencesByFunction } from "~/utils/clickhouse/inference.server"; import { getConfig, getAllFunctionConfigs } from "~/utils/config/index.server"; import type { Route } from "./+types/index"; -import { - countWorkflowEvaluationProjects, - countWorkflowEvaluationRuns, -} from "~/utils/clickhouse/workflow_evaluations.server"; +import { countWorkflowEvaluationRuns } from "~/utils/clickhouse/workflow_evaluations.server"; import { getTensorZeroClient } from "~/utils/tensorzero.server"; export const handle: RouteHandle = { @@ -111,7 +108,7 @@ export async function loader() { const numEvaluationRunsPromise = httpClient.countEvaluationRuns(); const numWorkflowEvaluationRunsPromise = countWorkflowEvaluationRuns(); const numWorkflowEvaluationRunProjectsPromise = - countWorkflowEvaluationProjects(); + httpClient.countWorkflowEvaluationProjects(); const configPromise = getConfig(); const functionConfigsPromise = getAllFunctionConfigs(); const numModelsUsedPromise = httpClient diff --git a/ui/app/routes/workflow_evaluations/route.tsx b/ui/app/routes/workflow_evaluations/route.tsx index 60a3d824a4..225e031336 100644 --- a/ui/app/routes/workflow_evaluations/route.tsx +++ b/ui/app/routes/workflow_evaluations/route.tsx @@ -10,7 +10,6 @@ import { import { getWorkflowEvaluationRuns, countWorkflowEvaluationRuns, - countWorkflowEvaluationProjects, } from "~/utils/clickhouse/workflow_evaluations.server"; import WorkflowEvaluationRunsTable from "./WorkflowEvaluationRunsTable"; import WorkflowEvaluationProjectsTable from "./WorkflowEvaluationProjectsTable"; @@ -24,6 +23,7 @@ export async function loader({ request }: Route.LoaderArgs) { const runLimit = parseInt(searchParams.get("runLimit") || "15"); const projectOffset = parseInt(searchParams.get("projectOffset") || "0"); const projectLimit = parseInt(searchParams.get("projectLimit") || "15"); + const tensorZeroClient = getTensorZeroClient(); const [ workflowEvaluationRuns, count, @@ -32,11 +32,8 @@ export async function loader({ request }: Route.LoaderArgs) { ] = await Promise.all([ getWorkflowEvaluationRuns(runLimit, runOffset), countWorkflowEvaluationRuns(), - getTensorZeroClient().getWorkflowEvaluationProjects( - projectLimit, - projectOffset, - ), - countWorkflowEvaluationProjects(), + tensorZeroClient.getWorkflowEvaluationProjects(projectLimit, projectOffset), + tensorZeroClient.countWorkflowEvaluationProjects(), ]); const workflowEvaluationProjects = workflowEvaluationProjectsResponse.projects; diff --git a/ui/app/utils/clickhouse/workflow_evaluations.server.ts b/ui/app/utils/clickhouse/workflow_evaluations.server.ts index 5ca5b27182..d448021531 100644 --- a/ui/app/utils/clickhouse/workflow_evaluations.server.ts +++ b/ui/app/utils/clickhouse/workflow_evaluations.server.ts @@ -315,21 +315,6 @@ export async function countWorkflowEvaluationRunEpisodes( return rows[0].count; } -export async function countWorkflowEvaluationProjects(): Promise { - const query = ` - SELECT toUInt32(countDistinct(project_name)) AS count - FROM DynamicEvaluationRunByProjectName - WHERE project_name IS NOT NULL -`; - const result = await getClickhouseClient().query({ - query, - format: "JSONEachRow", - }); - const rows = await result.json<{ count: number }[]>(); - const parsedRows = rows.map((row) => CountSchema.parse(row)); - return parsedRows[0].count; -} - export async function searchWorkflowEvaluationRuns( limit: number, offset: number, diff --git a/ui/app/utils/clickhouse/workflow_evaluations.test.ts b/ui/app/utils/clickhouse/workflow_evaluations.test.ts index abccd853c9..59017b8ad4 100644 --- a/ui/app/utils/clickhouse/workflow_evaluations.test.ts +++ b/ui/app/utils/clickhouse/workflow_evaluations.test.ts @@ -5,7 +5,6 @@ import { describe, test, expect } from "vitest"; import { - countWorkflowEvaluationProjects, countWorkflowEvaluationRunEpisodes, countWorkflowEvaluationRuns as countWorkflowEvaluationRuns, getWorkflowEvaluationRunEpisodesByTaskName, @@ -396,13 +395,6 @@ describe("countWorkflowEvaluationRunEpisodes", () => { }); }); -describe("countWorkflowEvaluationProjects", () => { - test("should return correct number of projects", async () => { - const count = await countWorkflowEvaluationProjects(); - expect(count).toBe(2); - }); -}); - describe("searchWorkflowEvaluationRuns", () => { test("should return correct runs by display name with project name set", async () => { const runs = await searchWorkflowEvaluationRuns( diff --git a/ui/app/utils/tensorzero/tensorzero.ts b/ui/app/utils/tensorzero/tensorzero.ts index e1b862e087..2fc916cdc3 100644 --- a/ui/app/utils/tensorzero/tensorzero.ts +++ b/ui/app/utils/tensorzero/tensorzero.ts @@ -26,6 +26,7 @@ import type { DeleteDatapointsResponse, GetModelLatencyResponse, GetModelUsageResponse, + GetWorkflowEvaluationProjectCountResponse, GetWorkflowEvaluationProjectsResponse, InferenceWithFeedbackStatsResponse, GetDatapointsRequest, @@ -946,6 +947,25 @@ export class TensorZeroClient { return (await response.json()) as GetWorkflowEvaluationProjectsResponse; } + /** + * Counts workflow evaluation projects. + * @returns A promise that resolves with the workflow evaluation project count + * @throws Error if the request fails + */ + async countWorkflowEvaluationProjects(): Promise { + const response = await this.fetch( + "/internal/workflow-evaluations/projects/count", + { method: "GET" }, + ); + if (!response.ok) { + const message = await this.getErrorText(response); + this.handleHttpError({ message, response }); + } + const body = + (await response.json()) as GetWorkflowEvaluationProjectCountResponse; + return body.count; + } + /** * Lists inference metadata with optional cursor-based pagination and filtering. * @param params - Optional pagination and filter parameters