Merge branch 'main' into feat/storybook-ci-tests

Author: rrrodzilla

.github/workflows/caching-hack.yml

@@ -54,8 +54,9 @@ jobs:
           shared-key: "build-gateway-cache"
       - name: Install uv
         if: runner.os == 'windows'
-        run: |
-          powershell -ExecutionPolicy ByPass -c "irm https://astral.sh/uv/install.ps1 | iex"
+        uses: astral-sh/setup-uv@ed21f2f24f8dd64503750218de024bcf64c7250a
+        with:
+          version: "0.6.17"
       - name: Install Python 3.10
         if: runner.os == 'windows'
         run: |

.github/workflows/general.yml

@@ -281,7 +281,9 @@ jobs:
         shell: bash
 
       - name: Install uv
-        run: curl -LsSf --retry 2 --retry-delay 10 --retry-max-time 60 https://astral.sh/uv/0.6.17/install.sh | sh
+        uses: astral-sh/setup-uv@ed21f2f24f8dd64503750218de024bcf64c7250a
+        with:
+          version: "0.6.17"
 
       - name: Generate Python schemas
         run: pnpm generate-python-schemas
@@ -317,7 +319,9 @@ jobs:
           cache-provider: "buildjet"
           save-if: ${{ github.event_name == 'merge_group' }}
       - name: Install uv
-        run: powershell -ExecutionPolicy ByPass -c "irm https://astral.sh/uv/install.ps1 | iex"
+        uses: astral-sh/setup-uv@ed21f2f24f8dd64503750218de024bcf64c7250a
+        with:
+          version: "0.6.17"
       - name: Install Python 3.10
         run: uv python install 3.10
       - name: Set PYO3_PYTHON
@@ -384,7 +388,9 @@ jobs:
         run: rustc --version
 
       - name: Install uv
-        run: curl -LsSf --retry 2 --retry-delay 10 --retry-max-time 60 https://astral.sh/uv/0.6.17/install.sh | sh
+        uses: astral-sh/setup-uv@ed21f2f24f8dd64503750218de024bcf64c7250a
+        with:
+          version: "0.6.17"
 
       - name: uv-lock
         run: |
@@ -628,7 +634,9 @@ jobs:
         shell: bash
 
       - name: Install uv
-        run: curl -LsSf --retry 2 --retry-delay 10 --retry-max-time 60 https://astral.sh/uv/0.6.17/install.sh | sh
+        uses: astral-sh/setup-uv@ed21f2f24f8dd64503750218de024bcf64c7250a
+        with:
+          version: "0.6.17"
 
       - name: Install JS dependencies
         run: pnpm install --frozen-lockfile
@@ -780,7 +788,9 @@ jobs:
             ./ui/fixtures/s3-fixtures/
 
       - name: Install uv
-        run: curl -LsSf --retry 2 --retry-delay 10 --retry-max-time 60 https://astral.sh/uv/0.6.17/install.sh | sh
+        uses: astral-sh/setup-uv@ed21f2f24f8dd64503750218de024bcf64c7250a
+        with:
+          version: "0.6.17"
 
       - name: Download container images
         run: |
@@ -969,7 +979,11 @@ jobs:
           tool: cargo-nextest
 
       - name: Install uv
-        run: curl -LsSf --retry 2 --retry-delay 10 --retry-max-time 60 https://astral.sh/uv/0.6.17/install.sh | sh
+        uses: astral-sh/setup-uv@ed21f2f24f8dd64503750218de024bcf64c7250a
+        with:
+          version: "0.6.17"
+          # Without this, the post-install hook tries to prune and save the cache, but ./ci/free-disk-space.sh already nukes /opt/hostedtoolcache/*, so it fails.
+          enable-cache: "false"
 
       - name: Download ClickHouse fixtures
         run: uv run ./ui/fixtures/download-fixtures.py

.github/workflows/merge-queue.yml

@@ -221,7 +221,9 @@ jobs:
         run: rustc --version
 
       - name: Install uv
-        run: curl -LsSf --retry 2 --retry-delay 10 --retry-max-time 60 https://astral.sh/uv/0.6.17/install.sh | sh
+        uses: astral-sh/setup-uv@ed21f2f24f8dd64503750218de024bcf64c7250a
+        with:
+          version: "0.6.17"
 
       - name: Install pnpm
         run: |

.github/workflows/python-client-build.yml

@@ -28,7 +28,9 @@ jobs:
           python-version: "3.10"
 
       - name: Install uv
-        run: curl -LsSf --retry 2 --retry-delay 10 --retry-max-time 60 https://astral.sh/uv/0.6.17/install.sh | sh
+        uses: astral-sh/setup-uv@ed21f2f24f8dd64503750218de024bcf64c7250a
+        with:
+          version: "0.6.17"
 
       - name: Install Rust toolchain
         run: |
@@ -96,7 +98,9 @@ jobs:
         with:
           python-version: "3.10"
       - name: Install uv
-        run: curl -LsSf --retry 2 --retry-delay 10 --retry-max-time 60 https://astral.sh/uv/0.6.17/install.sh | sh
+        uses: astral-sh/setup-uv@ed21f2f24f8dd64503750218de024bcf64c7250a
+        with:
+          version: "0.6.17"
       - name: Build wheels
         uses: PyO3/maturin-action@86b9d133d34bc1b40018696f782949dac11bd380
         with:
@@ -145,8 +149,9 @@ jobs:
           python-version: "3.10"
           architecture: ${{ matrix.platform.target }}
       - name: Install uv
-        run: curl -LsSf --retry 2 --retry-delay 10 --retry-max-time 60 https://astral.sh/uv/0.6.17/install.sh | sh
-        shell: bash
+        uses: astral-sh/setup-uv@ed21f2f24f8dd64503750218de024bcf64c7250a
+        with:
+          version: "0.6.17"
       - name: Build wheels
         uses: PyO3/maturin-action@86b9d133d34bc1b40018696f782949dac11bd380
         with:
@@ -199,7 +204,9 @@ jobs:
           python-version: "3.10"
 
       - name: Install uv
-        run: curl -LsSf --retry 2 --retry-delay 10 --retry-max-time 60 https://astral.sh/uv/0.6.17/install.sh | sh
+        uses: astral-sh/setup-uv@ed21f2f24f8dd64503750218de024bcf64c7250a
+        with:
+          version: "0.6.17"
 
       - name: Build wheels
         uses: PyO3/maturin-action@86b9d133d34bc1b40018696f782949dac11bd380
@@ -235,7 +242,9 @@ jobs:
           name: wheels-sdist
           path: ./clients/python/dist
       - name: Install uv
-        run: curl -LsSf --retry 2 --retry-delay 10 --retry-max-time 60 https://astral.sh/uv/0.6.17/install.sh | sh
+        uses: astral-sh/setup-uv@ed21f2f24f8dd64503750218de024bcf64c7250a
+        with:
+          version: "0.6.17"
       - name: Check local sdist
         run: |
           ./ci/check-local-wheel.sh ./clients/python/dist/tensorzero-*.tar.gz

.github/workflows/tensorzero-ci-bot-failure-diagnosis.yml

@@ -40,7 +40,9 @@ jobs:
         uses: pnpm/action-setup@v4
 
       - name: Install uv
-        uses: astral-sh/setup-uv@v7
+        uses: astral-sh/setup-uv@ed21f2f24f8dd64503750218de024bcf64c7250a
+        with:
+          version: "0.6.17"
 
       - name: Install mini-swe-agent
         run: 'uv tool install mini-swe-agent --from "git+https://github.com/virajmehta/mini-swe-agent.git@main"'

.github/workflows/upgrade-tensorzero-deps.yml

@@ -28,7 +28,9 @@ jobs:
           token: ${{ secrets.GITHUB_TOKEN }}
 
       - name: Install uv
-        run: curl -LsSf --retry 2 --retry-delay 10 --retry-max-time 60 https://astral.sh/uv/0.6.17/install.sh | sh
+        uses: astral-sh/setup-uv@ed21f2f24f8dd64503750218de024bcf64c7250a
+        with:
+          version: "0.6.17"
 
       - name: Install jq
         run: sudo apt-get update && sudo apt-get install -y jq

clients/python/tensorzero/generated_types.py

@@ -113,10 +113,10 @@ class DeleteDatapointsResponse:
 @dataclass(kw_only=True)
 class DemonstrationFeedbackFilter:
     """
-    Filter by whether an inference has demonstration feedback.
+    Filter by whether an inference has a demonstration.
     """
 
-    has_demonstration_feedback: bool
+    has_demonstration: bool
 
 
 Detail = Literal["low", "high", "auto"]
@@ -470,7 +470,7 @@ class InferenceFilterBooleanMetric(BooleanMetricFilter):
 @dataclass(kw_only=True)
 class InferenceFilterDemonstrationFeedback(DemonstrationFeedbackFilter):
     """
-    Filter by whether an inference has demonstration feedback.
+    Filter by whether an inference has a demonstration.
     """
 
     type: Literal["demonstration_feedback"] = "demonstration_feedback"

clients/rust/src/lib.rs

@@ -45,8 +45,8 @@ pub use tensorzero_core::db::clickhouse::query_builder::{
     TimeFilter,
 };
 pub use tensorzero_core::db::datasets::{
-    CountDatapointsForDatasetFunctionParams, DatasetQueries, DatasetQueryParams,
-    GetDatapointParams, GetDatapointsParams, GetDatasetMetadataParams,
+    DatasetQueries, DatasetQueryParams, GetDatapointParams, GetDatapointsParams,
+    GetDatasetMetadataParams,
 };
 pub use tensorzero_core::db::inferences::{InferenceOutputSource, ListInferencesParams};
 pub use tensorzero_core::db::stored_datapoint::{

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}/metrics",
+            get(endpoints::functions::internal::get_function_metrics_handler),
+        )
         .route(
             "/internal/functions/{function_name}/inference-stats",
             get(endpoints::internal::inference_stats::get_inference_stats_handler),
@@ -41,6 +45,14 @@ pub fn build_internal_non_otel_enabled_routes() -> Router<AppStateData> {
             "/internal/ui-config",
             get(endpoints::ui::get_config::ui_config_handler),
         )
+        .route(
+            "/internal/episodes",
+            get(endpoints::episodes::internal::list_episodes_handler),
+        )
+        .route(
+            "/internal/episodes/bounds",
+            get(endpoints::episodes::internal::query_episode_table_bounds_handler),
+        )
         .route(
             "/internal/datasets/{dataset_name}/datapoints",
             post(endpoints::datasets::insert_from_existing_datapoint_handler),
@@ -53,6 +65,10 @@ pub fn build_internal_non_otel_enabled_routes() -> Router<AppStateData> {
             "/internal/datasets/{dataset_name}/datapoints/{datapoint_id}",
             put(endpoints::datasets::update_datapoint_handler),
         )
+        .route(
+            "/internal/datasets/{dataset_name}/datapoints/count",
+            get(endpoints::datasets::internal::get_datapoint_count_handler),
+        )
         .route(
             "/internal/object_storage",
             get(endpoints::object_storage::get_object_handler),