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/.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 2de3c9cc02e..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. @@ -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)' @@ -155,9 +162,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/.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/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/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" 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 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 c8ef0e0c554..5c64d6ba3ac 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 @@ -646,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 @@ -706,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 @@ -765,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 @@ -776,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 @@ -797,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 @@ -918,8 +923,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, build-fixtures-container] if: ${{ github.event_name == 'merge_group' || (github.event.pull_request.head.repo.full_name == github.repository) }} steps: @@ -950,13 +959,34 @@ 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 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 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: | 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 - 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)' @@ -965,6 +995,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: @@ -1045,6 +1192,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 @@ -1085,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 @@ -1092,6 +1253,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: @@ -1116,17 +1280,29 @@ 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-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: | @@ -1187,6 +1363,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: @@ -1242,17 +1419,16 @@ jobs: check-all-general-jobs-passed: permissions: {} if: always() && github.repository == 'tensorzero/tensorzero' - needs: - [ + 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, 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, @@ -1268,6 +1444,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/.github/workflows/merge-queue.yml b/.github/workflows/merge-queue.yml index b55141eaa91..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 @@ -89,6 +127,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. @@ -148,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 @@ -181,6 +235,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. @@ -205,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: @@ -290,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() @@ -301,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 @@ -387,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() @@ -426,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/.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/.github/workflows/slash-command-regen-fixtures.yml b/.github/workflows/slash-command-regen-fixtures.yml index 6162f84c52c..c9745cb6c64 100644 --- a/.github/workflows/slash-command-regen-fixtures.yml +++ b/.github/workflows/slash-command-regen-fixtures.yml @@ -43,11 +43,20 @@ jobs: fi sleep $((10 * attempt)) done - shell: bash - 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 - name: Regenerate fixtures id: e2e_tests run: | @@ -76,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 500c859c74c..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: @@ -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 @@ -77,13 +77,22 @@ jobs: fi sleep $((10 * attempt)) done - shell: bash - 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 - name: Download container images uses: actions/download-artifact@018cc2cf5baa6db3ef3c5f8a56943fffe632ef53 @@ -139,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 @@ -148,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. @@ -163,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" @@ -189,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 @@ -256,6 +267,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-e2e.yml b/.github/workflows/ui-tests-e2e.yml index 57e7deee10e..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 @@ -105,13 +112,22 @@ jobs: fi sleep $((10 * attempt)) done - shell: bash - 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 - name: Download container images uses: actions/download-artifact@018cc2cf5baa6db3ef3c5f8a56943fffe632ef53 @@ -134,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, @@ -220,13 +241,22 @@ jobs: fi sleep $((10 * attempt)) done - shell: bash - 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 - name: Download container images uses: actions/download-artifact@018cc2cf5baa6db3ef3c5f8a56943fffe632ef53 @@ -254,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) @@ -308,13 +341,22 @@ jobs: fi sleep $((10 * attempt)) done - shell: bash - 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 - name: Download container images uses: actions/download-artifact@018cc2cf5baa6db3ef3c5f8a56943fffe632ef53 @@ -348,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 @@ -386,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/.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/.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/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/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/Cargo.lock b/Cargo.lock index c92b91d91a9..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" @@ -315,7 +334,7 @@ checksum = "c08606f8c3cbf4ce6ec8e28fb0014a2c086708fe954eaa885384a6165172e7e8" [[package]] name = "autopilot-client" -version = "2026.1.5" +version = "2026.1.6" dependencies = [ "chrono", "durable-tools-spawn", @@ -340,7 +359,7 @@ dependencies = [ [[package]] name = "autopilot-tools" -version = "2026.1.5" +version = "2026.1.6" dependencies = [ "anyhow", "async-trait", @@ -364,7 +383,7 @@ dependencies = [ [[package]] name = "autopilot-worker" -version = "2026.1.5" +version = "2026.1.6" dependencies = [ "anyhow", "async-trait", @@ -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", ] @@ -1357,9 +1371,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", ] @@ -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", @@ -1518,14 +1534,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" @@ -1556,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" @@ -1598,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" @@ -1632,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]] @@ -1761,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" @@ -1784,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" @@ -1796,6 +1870,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" @@ -1826,7 +1915,7 @@ dependencies = [ [[package]] name = "durable-tools" -version = "2026.1.5" +version = "2026.1.6" dependencies = [ "anyhow", "async-trait", @@ -1856,7 +1945,7 @@ dependencies = [ [[package]] name = "durable-tools-spawn" -version = "2026.1.5" +version = "2026.1.6" dependencies = [ "async-trait", "durable", @@ -1939,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" @@ -1972,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]] @@ -1988,7 +2087,7 @@ dependencies = [ [[package]] name = "evaluations" -version = "2026.1.5" +version = "2026.1.6" dependencies = [ "anyhow", "async-trait", @@ -2091,7 +2190,7 @@ dependencies = [ [[package]] name = "feedback-load-test" -version = "2026.1.5" +version = "2026.1.6" dependencies = [ "anyhow", "async-trait", @@ -2316,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" @@ -2337,9 +2427,10 @@ dependencies = [ [[package]] name = "gateway" -version = "2026.1.5" +version = "2026.1.6" dependencies = [ "async-stream", + "autopilot-tools", "autopilot-worker", "axum", "brotli", @@ -2536,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]] @@ -2605,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", ] @@ -2644,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" @@ -3066,7 +3152,7 @@ checksum = "9375e112e4b463ec1b1c6c011953545c65a30164fbab5b581df32b3abf0dcb88" dependencies = [ "console", "portable-atomic", - "unicode-width 0.2.2", + "unicode-width", "unit-prefix", "web-time 1.1.0", ] @@ -3103,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" @@ -3145,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" @@ -3256,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" @@ -3301,9 +3402,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", @@ -3358,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" @@ -3370,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" @@ -3396,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]] @@ -3565,7 +3681,7 @@ dependencies = [ [[package]] name = "minijinja-utils" -version = "2026.1.5" +version = "2026.1.6" dependencies = [ "minijinja", "serde", @@ -3588,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" @@ -3607,6 +3711,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a69bcab0ad47271a0234d9422b131806bf3968021e5dc9328caf2d4cd58557fc" dependencies = [ "libc", + "log", "wasi", "windows-sys 0.61.2", ] @@ -3710,15 +3815,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", ] @@ -3730,12 +3837,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", @@ -3744,24 +3851,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", ] @@ -3773,10 +3878,10 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "650eef8c711430f1a879fdd01d4745a7deea475becfb90269c06775983bbf086" [[package]] -name = "no-std-compat" -version = "0.4.1" +name = "nohash-hasher" +version = "0.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b93853da6d84c2e3c7d730d6473e8817692dd89be387eb01b94d7f108ecb5b8c" +checksum = "2bf50223579dc7cdcfb3bfcacf7069ff68243f8c363f62ffa99cf000a6b9c451" [[package]] name = "nom" @@ -3815,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]] @@ -3931,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" @@ -4095,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]] @@ -4308,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", ] @@ -4376,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", @@ -4511,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", @@ -4617,7 +4731,7 @@ dependencies = [ "once_cell", "socket2", "tracing", - "windows-sys 0.60.2", + "windows-sys 0.52.0", ] [[package]] @@ -4730,28 +4844,70 @@ dependencies = [ [[package]] name = "ratatui" -version = "0.27.0" +version = "0.30.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d16546c5b5962abf8ce6e2881e722b4e0ae3b6f1a08a26ae3573c55853ca68d3" +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 = "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]] name = "rate-limit-load-test" -version = "2026.1.5" +version = "2026.1.6" dependencies = [ "anyhow", "async-trait", @@ -5149,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", @@ -5247,7 +5405,7 @@ dependencies = [ "errno", "libc", "linux-raw-sys", - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -5306,7 +5464,7 @@ dependencies = [ "security-framework", "security-framework-sys", "webpki-root-certs", - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -5648,7 +5806,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b75a19a7a740b25bc7944bdee6172368f988763b744e3d4dfe753f6b4ece40cc" dependencies = [ "libc", - "mio 0.8.11", + "mio", "signal-hook", ] @@ -5838,7 +5996,7 @@ dependencies = [ "cfg-if", "dotenvy", "either", - "heck 0.5.0", + "heck", "hex", "proc-macro2", "quote", @@ -5962,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" @@ -6007,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]] @@ -6044,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", @@ -6121,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]] @@ -6177,7 +6305,7 @@ dependencies = [ "getrandom 0.3.4", "once_cell", "rustix", - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -6206,7 +6334,7 @@ dependencies = [ [[package]] name = "tensorzero-auth" -version = "2026.1.5" +version = "2026.1.6" dependencies = [ "axum", "chrono", @@ -6226,7 +6354,7 @@ dependencies = [ [[package]] name = "tensorzero-core" -version = "2026.1.5" +version = "2026.1.6" dependencies = [ "arc-swap", "async-stream", @@ -6248,7 +6376,7 @@ dependencies = [ "chrono", "clap", "clarabel", - "dashmap 6.1.0", + "dashmap", "derive_builder 0.20.2", "durable", "durable-tools", @@ -6306,7 +6434,7 @@ dependencies = [ "serde_path_to_error", "sha2", "sqlx", - "strum 0.27.2", + "strum", "tempfile", "tensorzero", "tensorzero-auth", @@ -6347,7 +6475,7 @@ dependencies = [ [[package]] name = "tensorzero-node" -version = "2026.1.5" +version = "2026.1.6" dependencies = [ "napi", "napi-build", @@ -6360,7 +6488,7 @@ dependencies = [ [[package]] name = "tensorzero-optimizers" -version = "2026.1.5" +version = "2026.1.6" dependencies = [ "axum", "base64 0.22.1", @@ -6394,7 +6522,7 @@ dependencies = [ [[package]] name = "tensorzero-python" -version = "2026.1.5" +version = "2026.1.6" dependencies = [ "evaluations", "futures", @@ -6412,7 +6540,7 @@ dependencies = [ [[package]] name = "tensorzero-types" -version = "2026.1.5" +version = "2026.1.6" dependencies = [ "aws-smithy-types", "infer", @@ -6432,7 +6560,7 @@ dependencies = [ [[package]] name = "tensorzero-types-providers" -version = "2026.1.5" +version = "2026.1.6" dependencies = [ "serde", "serde_json", @@ -6440,7 +6568,7 @@ dependencies = [ [[package]] name = "tensorzero-unsafe-helpers" -version = "2026.1.5" +version = "2026.1.6" [[package]] name = "termcolor" @@ -6457,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" @@ -6528,7 +6666,9 @@ checksum = "f9e442fc33d7fdb45aa9bfeb312c095964abdf596f7567261062b2a7107aaabd" dependencies = [ "deranged", "itoa", + "libc", "num-conv", + "num_threads", "powerfmt", "serde_core", "time-core", @@ -6584,7 +6724,7 @@ checksum = "72a2903cd7736441aac9df9d7688bd0ce48edccaadf181c3b90be801e81d3d86" dependencies = [ "bytes", "libc", - "mio 1.1.1", + "mio", "parking_lot", "pin-project-lite", "signal-hook-registry", @@ -6990,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]] @@ -7057,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" @@ -7196,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" @@ -7403,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]] @@ -7491,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" @@ -7551,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" @@ -7605,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" @@ -7629,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" @@ -7653,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" @@ -7689,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" @@ -7713,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" @@ -7737,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" @@ -7761,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 7e3930c6052..da92195eb44 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" @@ -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/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/ci/buildkite/test-clickhouse-cloud.sh b/ci/buildkite/test-clickhouse-cloud.sh index e564cdab544..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 @@ -95,7 +101,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 @@ -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/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..." 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/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/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/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/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/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) 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" }, ] diff --git a/clients/rust/src/lib.rs b/clients/rust/src/lib.rs index e4af20551ac..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 // ================================================================ @@ -598,12 +616,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) @@ -1098,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/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. 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/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/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..e03e18f20d3 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")?; @@ -372,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"); @@ -385,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, )); @@ -555,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/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/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..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 @@ -40,13 +41,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 +87,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 +112,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(()) /// } /// @@ -126,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?; @@ -139,7 +162,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 +195,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/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 05c27a6ea46..a1eec9f2733 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; @@ -15,11 +16,13 @@ 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; mod write_config; +pub use auto_reject_tool_call::AutoRejectToolCallTool; pub use create_datapoints::{CreateDatapointsTool, CreateDatapointsToolParams}; pub use create_datapoints_from_inferences::{ CreateDatapointsFromInferencesTool, CreateDatapointsFromInferencesToolParams, @@ -39,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/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..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; @@ -56,7 +57,7 @@ mock! { &self, session_id: Uuid, params: durable_tools::ListEventsParams, - ) -> Result; + ) -> Result; async fn list_autopilot_sessions( &self, @@ -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/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..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; @@ -437,7 +438,7 @@ mod tests { &self, session_id: Uuid, params: durable_tools::ListEventsParams, - ) -> Result; + ) -> Result; async fn list_autopilot_sessions( &self, @@ -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/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/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..a12d2a03473 100644 --- a/internal/durable-tools/src/tensorzero_client/client_ext.rs +++ b/internal/durable-tools/src/tensorzero_client/client_ext.rs @@ -6,13 +6,15 @@ use async_trait::async_trait; use autopilot_client::AutopilotError; +use autopilot_client::GatewayListEventsResponse; use tensorzero::{ Client, ClientExt, ClientInferenceParams, ClientMode, CreateDatapointRequest, 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; @@ -30,9 +32,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 +152,7 @@ impl TensorZeroClient for Client { &self, session_id: Uuid, params: ListEventsParams, - ) -> Result { + ) -> Result { match self.mode() { ClientMode::HTTPGateway(http) => { let mut url = http @@ -367,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 146015a3858..d7f27bdf29e 100644 --- a/internal/durable-tools/src/tensorzero_client/embedded.rs +++ b/internal/durable-tools/src/tensorzero_client/embedded.rs @@ -4,13 +4,14 @@ //! 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, 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; @@ -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 @@ -277,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 6e6969dec34..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; @@ -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( @@ -191,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-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/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/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/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/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..c297584fc64 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"; @@ -203,11 +208,11 @@ export * from "./LatestFeedbackIdByMetricResponse"; export * from "./LaunchOptimizationParams"; export * from "./LaunchOptimizationWorkflowParams"; export * from "./ListDatapointsRequest"; +export * from "./ListDatasetsRequest"; 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 +320,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/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/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/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/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/inference_queries.rs b/tensorzero-core/src/db/clickhouse/inference_queries.rs index 9191807444a..86270069531 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,193 @@ 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)) + } + + 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. @@ -1683,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 6423e48d5c7..d6b8c82b56e 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,46 @@ 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 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] @@ -126,6 +169,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/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/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/db/inferences.rs b/tensorzero-core/src/db/inferences.rs index 981eddc9294..b1e80e073af 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,52 @@ 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>; + + /// 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/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/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/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/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) -} 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, 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/feedback/human_feedback.rs b/tensorzero-core/src/endpoints/feedback/human_feedback.rs index f3faef525d8..172a99efaa5 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, InferenceQueries}; use crate::{ config::MetricConfigLevel, db::clickhouse::{ClickHouseConnectionInfo, TableName}, @@ -39,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, @@ -100,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, - 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 = '{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/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)) } } diff --git a/tensorzero-core/src/endpoints/inference.rs b/tensorzero-core/src/endpoints/inference.rs index 60380ed6aa5..20ec4b4f981 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; @@ -1264,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; @@ -2051,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, }; @@ -2124,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, @@ -2741,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/endpoints/internal/autopilot.rs b/tensorzero-core/src/endpoints/internal/autopilot.rs index f11ae34a9d6..016e2abe879 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) => { @@ -280,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)] @@ -291,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 { @@ -306,6 +316,7 @@ mod tests { postgres_connection_info, ValkeyConnectionInfo::Disabled, TaskTracker::new(), + CancellationToken::new(), ) .unwrap() } 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/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/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 // ============================================================================ diff --git a/tensorzero-core/src/error/mod.rs b/tensorzero-core/src/error/mod.rs index 10178ceb17f..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 @@ -493,6 +519,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 +548,6 @@ pub enum ErrorDetails { message: String, }, PostgresQuery { - function_name: Option, message: String, }, PostgresResult { @@ -716,6 +745,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, @@ -793,7 +823,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, @@ -872,6 +902,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, @@ -880,11 +911,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, @@ -1474,6 +1505,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 +1556,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}") } @@ -1680,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 @@ -1708,7 +1728,6 @@ impl From for Error { fn from(err: sqlx::Error) -> Self { Self::new(ErrorDetails::PostgresQuery { message: err.to_string(), - function_name: None, }) } } @@ -1757,6 +1776,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 +1803,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/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/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 8772f65fdbb..8778ccaa3cd 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 { @@ -1377,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( @@ -1397,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(); @@ -1433,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, @@ -1443,7 +1503,7 @@ pub(super) fn openai_responses_to_tensorzero_chunk( Some(entries), ), None => ProviderInferenceResponseChunk::new( - vec![], + content_blocks, usage, raw_message, message_latency, @@ -1537,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 { .. } @@ -1932,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 { @@ -2249,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/utils/gateway.rs b/tensorzero-core/src/utils/gateway.rs index dfd16d5c32c..aa648a557f9 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; @@ -59,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: (), } @@ -68,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 @@ -155,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. @@ -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 } @@ -257,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: (), } @@ -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 { @@ -329,9 +341,9 @@ impl GatewayHandle { autopilot_client, deployment_id, rate_limiting_manager, + shutdown_token: cancel_token, _private: (), }, - cancel_token, drop_wrapper, _private: (), }) @@ -349,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()), @@ -367,6 +380,7 @@ impl AppStateData { autopilot_client: None, deployment_id: None, rate_limiting_manager, + shutdown_token, _private: (), }) } @@ -464,6 +478,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 +514,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) } @@ -529,6 +550,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) => { @@ -550,10 +572,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") { @@ -589,40 +612,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) } } @@ -677,6 +714,7 @@ pub async fn start_openai_compatible_gateway( clickhouse_url, postgres_url, valkey_url, + HashSet::new(), // available_tools )) .await?; @@ -923,6 +961,7 @@ mod tests { postgres: PostgresConfig { enabled: Some(false), connection_pool_size: 20, + inference_retention_days: None, }, ..Default::default() })); @@ -941,6 +980,7 @@ mod tests { postgres: PostgresConfig { enabled: Some(false), connection_pool_size: 20, + inference_retention_days: None, }, ..Default::default() })); @@ -964,6 +1004,7 @@ mod tests { postgres: PostgresConfig { enabled: None, connection_pool_size: 20, + inference_retention_days: None, }, ..Default::default() })); @@ -982,6 +1023,7 @@ mod tests { postgres: PostgresConfig { enabled: Some(true), connection_pool_size: 20, + inference_retention_days: None, }, ..Default::default() })); @@ -1000,6 +1042,7 @@ mod tests { postgres: PostgresConfig { enabled: Some(true), connection_pool_size: 20, + inference_retention_days: None, }, ..Default::default() })); @@ -1016,6 +1059,7 @@ mod tests { postgres: PostgresConfig { enabled: Some(false), connection_pool_size: 20, + inference_retention_days: None, }, rate_limiting: Default::default(), ..Default::default() @@ -1031,6 +1075,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/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/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) 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/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/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) 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/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" + ); +} 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.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.live.yml b/tensorzero-core/tests/e2e/docker-compose.live.yml index 1dd060184f9..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,41 @@ 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"] + interval: 5s + timeout: 1s + 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.marker" ] + test: ["CMD", "test", "-f", "/load_complete_postgres.marker"] interval: 5s timeout: 1s retries: 72 # Retry for up to 6 minutes @@ -210,6 +256,8 @@ services: depends_on: fixtures: condition: service_completed_successfully + fixtures-postgres: + condition: service_completed_successfully clickhouse: condition: service_healthy postgres: 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 98a56420c02..637f307ce4c 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: ../../.. @@ -29,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 @@ -38,9 +40,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 @@ -57,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: @@ -64,9 +72,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 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 tensorzero ValkeyConnectionInfo::Disabled, TensorzeroHttpClient::new_testing().unwrap(), None, + HashSet::new(), // available_tools ) .await .unwrap(); 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/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/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) { 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 1dd20ad12af..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, @@ -2002,9 +1999,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/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" + ); + } } 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" }} 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" + ); +} 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}" diff --git a/tensorzero-core/tests/e2e/tests.rs b/tensorzero-core/tests/e2e/tests.rs index 0e751929fdb..1d80223077d 100644 --- a/tensorzero-core/tests/e2e/tests.rs +++ b/tensorzero-core/tests/e2e/tests.rs @@ -10,6 +10,7 @@ mod best_of_n; mod built_in; mod cache; mod clickhouse; +#[macro_use] mod common; mod config; mod datasets; 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 => { 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.** 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 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..0e1f13505d0 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) { @@ -165,8 +165,9 @@ function summarizeEvent(event: Event): EventSummary { } return {}; case "error": - // TODO: handle errors - return {}; + return { + description: payload.message, + }; case "unknown": return {}; default: @@ -174,7 +175,7 @@ function summarizeEvent(event: Event): EventSummary { } } -function renderEventTitle(event: Event) { +function renderEventTitle(event: GatewayEvent) { const { payload } = event; switch (payload.type) { @@ -321,7 +322,7 @@ function EventItem({ event, isPending = false, }: { - event: Event; + event: GatewayEvent; isPending?: boolean; }) { const summary = summarizeEvent(event); @@ -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", )} 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/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 f8a244883a7..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/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.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/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 ; } 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/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) => } + + ); } diff --git a/ui/app/routes/index.tsx b/ui/app/routes/index.tsx index f48b8a4a2e8..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 diff --git a/ui/app/tailwind.css b/ui/app/tailwind.css index e6c8beec50c..1ec5e1803f7 100644 --- a/ui/app/tailwind.css +++ b/ui/app/tailwind.css @@ -160,6 +160,8 @@ --color-card-highlight: hsl(var(--card-highlight)); --color-card-highlight-icon: hsl(var(--card-highlight-icon)); + --color-card-highlight-icon-bg: hsl(var(--card-highlight-icon-bg)); + --color-card-highlight-border: hsl(var(--card-highlight-border)); --animate-accordion-down: accordion-down 0.2s ease-out; --animate-accordion-up: accordion-up 0.2s ease-out; @@ -341,6 +343,8 @@ --menu-highlight-badge-foreground: var(--orange-800); --card-highlight: var(--orange-50); --card-highlight-icon: var(--orange-600); + --card-highlight-icon-bg: var(--orange-100); + --card-highlight-border: var(--orange-100); } .dark { @@ -390,6 +394,8 @@ --menu-highlight-badge-foreground: var(--orange-200); --card-highlight: var(--orange-900); --card-highlight-icon: var(--orange-200); + --card-highlight-icon-bg: var(--orange-800); + --card-highlight-border: var(--orange-800); } } 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; } } 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/error-boundaries.spec.ts b/ui/e2e_tests/error-boundaries.spec.ts index 170029a3023..9be3450adc9 100644 --- a/ui/e2e_tests/error-boundaries.spec.ts +++ b/ui/e2e_tests/error-boundaries.spec.ts @@ -12,7 +12,7 @@ test.describe("Error Boundaries", () => { ).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(); 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 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/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 20f8112526d..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 @@ -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,8 +37,14 @@ "large_json_comment_feedback.parquet", "large_json_demonstration_feedback.parquet", ] -R2_BUCKET = "https://pub-147e9850a60643208c411e70b636e956.r2.dev" -S3_FIXTURES_DIR = Path("./s3-fixtures") +R2_PUBLIC_BUCKET_URL = "https://pub-147e9850a60643208c411e70b636e956.r2.dev" +R2_S3_ENDPOINT_URL = "https://19918a216783f0ac9e052233569aef60.r2.cloudflarestorage.com/" +LARGE_FIXTURES_DIR = Path("./large-fixtures") + + +# ============================================================================= +# Shared utilities +# ============================================================================= def calculate_etag(file_path): @@ -61,85 +71,139 @@ 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 download_file(filename, remote_etag): - """Download file from R2 bucket.""" - RETRIES = 3 - for i in range(RETRIES): +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: - url = f"{R2_BUCKET}/{filename}" - response = requests.get(url, stream=True) + 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 = LARGE_FIXTURES_DIR / fixture + if not local_file.exists(): + raise Exception(f"Fixture {fixture} not found after sync") - 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) + remote_etag = get_remote_etag(fixture) - 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"{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(LARGE_FIXTURES_DIR), + # Only download the files in `FIXTURES` + "--exclude", + "*", + *[arg for f in FIXTURES 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", + } + + 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) + + 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"Error downloading `{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 `{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(): - # 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 + raise Exception(f"Fixture sync failed after {retries} attempts") from last_error - 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(fixture, remote_etag) - return +# ============================================================================= +# Main +# ============================================================================= - 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(fixture, remote_etag) - else: - print(f"Skipping {fixture} (up to date)", flush=True) +def main(): + LARGE_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-large-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): - assert result is None + 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 811a3738015..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 @@ -24,7 +28,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", @@ -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 25d19dab31c..b8d4a0b51af 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,23 +35,28 @@ 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 +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 @@ -80,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 0087c186852..e3c1c09e008 100755 --- a/ui/fixtures/load_fixtures_postgres.sh +++ b/ui/fixtures/load_fixtures_postgres.sh @@ -1,6 +1,146 @@ #!/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 "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, + 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" < ./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/reset-dev-postgres.sh b/ui/fixtures/reset-dev-postgres.sh new file mode 100755 index 00000000000..7c5977ddb6c --- /dev/null +++ b/ui/fixtures/reset-dev-postgres.sh @@ -0,0 +1,37 @@ +#!/bin/bash +set -e + +# Development database reset script +# Drops and recreates the dev database, then runs migrations + +DB_NAME="${POSTGRES_DEV_DB:-postgres_migration_dev}" +DB_HOST="${POSTGRES_HOST:-localhost}" +DB_PORT="${POSTGRES_PORT:-5432}" +DB_USER="${POSTGRES_USER:-postgres}" +DB_PASSWORD="${POSTGRES_PASSWORD:-postgres}" + +echo "==> 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-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 b1c93112aae..34712f6c527 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" @@ -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 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",