Port countInferencesForFunction and countInferencesForVariant (tensorzero#5069)

* Port countInferencesForFunction and countInferencesForVariant

* Make frontend call the count inferences API

Author: shuyangli

gateway/src/routes/internal.rs

@@ -21,6 +21,10 @@ pub fn build_internal_non_otel_enabled_routes() -> Router<AppStateData> {
             "/internal/functions/{function_name}/variant_sampling_probabilities",
             get(endpoints::variant_probabilities::get_variant_sampling_probabilities_by_function_handler),
         )
+        .route(
+            "/internal/functions/{function_name}/inference-stats",
+            get(endpoints::internal::inference_stats::get_inference_stats_handler),
+        )
         .route(
             "/internal/ui-config",
             get(endpoints::ui::get_config::ui_config_handler),

internal/tensorzero-node/lib/bindings/InferenceStatsResponse.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.
+
+/**
+ * Response containing inference statistics
+ */
+export type InferenceStatsResponse = {
+  /**
+   * The count of inferences for the function (and optionally variant)
+   */
+  inference_count: bigint;
+};

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

@@ -148,6 +148,7 @@ export * from "./InferenceOutputSource";
 export * from "./InferenceParams";
 export * from "./InferenceResponse";
 export * from "./InferenceResponseToolCall";
+export * from "./InferenceStatsResponse";
 export * from "./Input";
 export * from "./InputMessage";
 export * from "./InputMessageContent";

tensorzero-core/src/db/clickhouse/inference_stats.rs

@@ -0,0 +1,140 @@
+//! ClickHouse queries for inference statistics.
+
+use std::collections::HashMap;
+
+use super::ClickHouseConnectionInfo;
+use super::select_queries::parse_count;
+use crate::error::Error;
+use crate::function::FunctionConfigType;
+
+/// Parameters for counting inferences for a function.
+#[derive(Debug)]
+pub struct CountInferencesParams<'a> {
+    pub function_name: &'a str,
+    pub function_type: FunctionConfigType,
+    pub variant_name: Option<&'a str>,
+}
+
+impl ClickHouseConnectionInfo {
+    /// Counts the number of inferences for a function, optionally filtered by variant.
+    pub async fn count_inferences_for_function(
+        &self,
+        params: CountInferencesParams<'_>,
+    ) -> Result<u64, Error> {
+        let (query, params) = build_count_inferences_query(&params);
+        let response = self.run_query_synchronous(query, &params).await?;
+        parse_count(&response.response)
+    }
+}
+
+/// Builds the SQL query for counting inferences.
+fn build_count_inferences_query<'a>(
+    params: &'a CountInferencesParams<'a>,
+) -> (String, HashMap<&'a str, &'a str>) {
+    let mut query_params = HashMap::new();
+    query_params.insert("function_name", params.function_name);
+
+    let table_name = match params.function_type {
+        FunctionConfigType::Chat => "ChatInference",
+        FunctionConfigType::Json => "JsonInference",
+    };
+
+    let query = match params.variant_name {
+        Some(variant_name) => {
+            query_params.insert("variant_name", variant_name);
+            format!(
+                "SELECT COUNT() AS count
+         FROM {table_name}
+         WHERE function_name = {{function_name:String}}
+           AND variant_name = {{variant_name:String}}
+         FORMAT JSONEachRow"
+            )
+        }
+        None => {
+            format!(
+                "SELECT COUNT() AS count
+         FROM {table_name}
+         WHERE function_name = {{function_name:String}}
+         FORMAT JSONEachRow"
+            )
+        }
+    };
+
+    (query, query_params)
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::db::clickhouse::query_builder::test_util::{
+        assert_query_contains, assert_query_does_not_contain,
+    };
+
+    #[test]
+    fn test_build_count_inferences_query_chat_no_variant() {
+        let params = CountInferencesParams {
+            function_name: "write_haiku",
+            function_type: FunctionConfigType::Chat,
+            variant_name: None,
+        };
+        let (query, query_params) = build_count_inferences_query(&params);
+        assert_query_contains(&query, "FROM ChatInference");
+        assert_query_contains(&query, "function_name = {function_name:String}");
+        assert_query_does_not_contain(&query, "variant_name");
+        assert_eq!(query_params.len(), 1);
+        assert_eq!(query_params.get("function_name"), Some(&"write_haiku"));
+    }
+
+    #[test]
+    fn test_build_count_inferences_query_json_no_variant() {
+        let params = CountInferencesParams {
+            function_name: "extract_entities",
+            function_type: FunctionConfigType::Json,
+            variant_name: None,
+        };
+        let (query, query_params) = build_count_inferences_query(&params);
+        assert_query_contains(&query, "FROM JsonInference");
+        assert_query_contains(&query, "function_name = {function_name:String}");
+        assert_query_does_not_contain(&query, "variant_name");
+        assert_eq!(query_params.len(), 1);
+        assert_eq!(query_params.get("function_name"), Some(&"extract_entities"));
+    }
+
+    #[test]
+    fn test_build_count_inferences_query_chat_with_variant() {
+        let params = CountInferencesParams {
+            function_name: "write_haiku",
+            function_type: FunctionConfigType::Chat,
+            variant_name: Some("initial_prompt_gpt4o_mini"),
+        };
+        let (query, query_params) = build_count_inferences_query(&params);
+        assert_query_contains(&query, "FROM ChatInference");
+        assert_query_contains(&query, "function_name = {function_name:String}");
+        assert_query_contains(&query, "variant_name = {variant_name:String}");
+        assert_eq!(query_params.len(), 2);
+        assert_eq!(query_params.get("function_name"), Some(&"write_haiku"));
+        assert_eq!(
+            query_params.get("variant_name"),
+            Some(&"initial_prompt_gpt4o_mini")
+        );
+    }
+
+    #[test]
+    fn test_build_count_inferences_query_json_with_variant() {
+        let params = CountInferencesParams {
+            function_name: "extract_entities",
+            function_type: FunctionConfigType::Json,
+            variant_name: Some("gpt4o_initial_prompt"),
+        };
+        let (query, query_params) = build_count_inferences_query(&params);
+        assert_query_contains(&query, "FROM JsonInference");
+        assert_query_contains(&query, "function_name = {function_name:String}");
+        assert_query_contains(&query, "variant_name = {variant_name:String}");
+        assert_eq!(query_params.len(), 2);
+        assert_eq!(query_params.get("function_name"), Some(&"extract_entities"));
+        assert_eq!(
+            query_params.get("variant_name"),
+            Some(&"gpt4o_initial_prompt")
+        );
+    }
+}

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

@@ -30,6 +30,7 @@ pub mod clickhouse_client; // Public because tests will use clickhouse_client::F
 pub mod dataset_queries;
 pub mod feedback;
 pub mod inference_queries;
+pub mod inference_stats;
 pub mod migration_manager;
 pub mod query_builder;
 mod select_queries;

tensorzero-core/src/endpoints/internal/inference_stats.rs

@@ -0,0 +1,146 @@
+//! Inference statistics endpoint for getting inference counts.
+
+use axum::extract::{Path, Query, State};
+use axum::{Json, debug_handler};
+use serde::{Deserialize, Serialize};
+use tracing::instrument;
+
+use crate::db::clickhouse::inference_stats::CountInferencesParams;
+use crate::error::{Error, ErrorDetails};
+use crate::utils::gateway::{AppState, AppStateData};
+
+/// Query parameters for the inference stats endpoint
+#[derive(Debug, Deserialize)]
+pub struct InferenceStatsQueryParams {
+    /// Optional variant name to filter by
+    pub variant_name: Option<String>,
+}
+
+/// Response containing inference statistics
+#[derive(Debug, Serialize, Deserialize, ts_rs::TS)]
+#[ts(export)]
+pub struct InferenceStatsResponse {
+    /// The count of inferences for the function (and optionally variant)
+    pub inference_count: u64,
+}
+
+/// HTTP handler for the inference stats endpoint
+#[debug_handler(state = AppStateData)]
+#[instrument(
+    name = "get_inference_stats_handler",
+    skip_all,
+    fields(
+        function_name = %function_name,
+    )
+)]
+pub async fn get_inference_stats_handler(
+    State(app_state): AppState,
+    Path(function_name): Path<String>,
+    Query(params): Query<InferenceStatsQueryParams>,
+) -> Result<Json<InferenceStatsResponse>, Error> {
+    Ok(Json(
+        get_inference_stats(app_state, &function_name, params).await?,
+    ))
+}
+
+/// Core business logic for getting inference statistics
+async fn get_inference_stats(
+    AppStateData {
+        config,
+        clickhouse_connection_info,
+        ..
+    }: AppStateData,
+    function_name: &str,
+    params: InferenceStatsQueryParams,
+) -> Result<InferenceStatsResponse, Error> {
+    // Get the function config to determine the function type
+    let function = config.get_function(function_name)?;
+
+    // If variant_name is provided, validate that it exists
+    if let Some(ref variant_name) = params.variant_name
+        && !function.variants().contains_key(variant_name)
+    {
+        return Err(ErrorDetails::UnknownVariant {
+            name: variant_name.clone(),
+        }
+        .into());
+    }
+
+    let count_params = CountInferencesParams {
+        function_name,
+        function_type: function.config_type(),
+        variant_name: params.variant_name.as_deref(),
+    };
+
+    let inference_count = clickhouse_connection_info
+        .count_inferences_for_function(count_params)
+        .await?;
+
+    Ok(InferenceStatsResponse { inference_count })
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::config::{Config, ConfigFileGlob};
+    use crate::testing::get_unit_test_gateway_handle;
+    use std::io::Write;
+    use std::sync::Arc;
+    use tempfile::NamedTempFile;
+
+    #[tokio::test]
+    async fn test_get_inference_stats_function_not_found() {
+        let config = Arc::new(Config::default());
+        let gateway_handle = get_unit_test_gateway_handle(config);
+
+        let params = InferenceStatsQueryParams { variant_name: None };
+
+        let result = get_inference_stats(
+            gateway_handle.app_state.clone(),
+            "nonexistent_function",
+            params,
+        )
+        .await;
+
+        assert!(result.is_err());
+        let err = result.unwrap_err();
+        assert!(err.to_string().contains("nonexistent_function"));
+    }
+
+    #[tokio::test]
+    async fn test_get_inference_stats_variant_not_found() {
+        // Create a config with a function but without the requested variant
+        let config_str = r#"
+            [functions.test_function]
+            type = "chat"
+
+            [functions.test_function.variants.variant_a]
+            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 gateway_handle = get_unit_test_gateway_handle(Arc::new(config));
+
+        let params = InferenceStatsQueryParams {
+            variant_name: Some("nonexistent_variant".to_string()),
+        };
+
+        let result =
+            get_inference_stats(gateway_handle.app_state.clone(), "test_function", params).await;
+
+        assert!(result.is_err());
+        let err = result.unwrap_err();
+        assert!(err.to_string().contains("nonexistent_variant"));
+    }
+}

tensorzero-core/src/endpoints/internal/mod.rs

@@ -0,0 +1 @@
+pub mod inference_stats;

tensorzero-core/src/endpoints/mod.rs

@@ -8,6 +8,7 @@ pub mod embeddings;
 pub mod fallback;
 pub mod feedback;
 pub mod inference;
+pub mod internal;
 pub mod object_storage;
 pub mod openai_compatible;
 pub mod shared_types;