@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 39910f2b25
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 4, 05:00 UTC

[Rust Frontend] Move code from `vllm-frontend-rs` (#43283)

Signed-off-by: Bugen Zhao <i@bugenzhao.com> Signed-off-by: Nick Hill <nickhill123@gmail.com> Signed-off-by: Eric Curtin <eric.curtin@docker.com> Signed-off-by: Dev-X25874 <283057883+Dev-X25874@users.noreply.github.com> Signed-off-by: Will.hou <1205157517@qq.com> Signed-off-by: Will.hou <willamhou@ceresman.com> Co-authored-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: Eric Curtin <eric.curtin@docker.com> Co-authored-by: Dev-X25874 <283057883+Dev-X25874@users.noreply.github.com> Co-authored-by: Will.hou <1205157517@qq.com> Co-authored-by: Will.hou <willamhou@ceresman.com> Please see https://github.com/Inferact/vllm-frontend-rs for full original commit history.

Bugen Zhao committed May 22, 2026 at 08:21 UTC 39910f2b25aacc09f5e7f166cdf0030b19f8b9e8
284 files changed +72569 -90
.buildkite/image_build/image_build.sh
-7
@@ -223,13 +223,6 @@ echo "CACHE_FROM_MAIN: ${CACHE_FROM_MAIN}"
223
224 check_and_skip_if_image_exists
225
226 -# The rust frontend lives in a git submodule under rust/. Buildkite's default
227 -# checkout does not recurse submodules, and the Dockerfile only sees what's in
228 -# the build context, so initialize the submodule here before invoking bake.
229 -echo "--- :git: Initializing git submodules"
230 -git submodule sync --recursive
231 -git submodule update --init --recursive
232 -
226 echo "--- :docker: Setting up Docker buildx bake"
227 echo "Target: ${TARGET}"
228 echo "vLLM bake file: ${VLLM_BAKE_FILE_PATH}"
.buildkite/image_build/image_build_cpu.sh
-6
@@ -21,12 +21,6 @@ else
21 exit 0
22 fi
23
24 -# The rust frontend lives in a git submodule under rust/. Buildkite's default
25 -# checkout does not recurse submodules, and the Dockerfile only sees what's in
26 -# the build context, so initialize the submodule here before building.
27 -git submodule sync --recursive
28 -git submodule update --init --recursive
29 -
24 # build
25 docker build --file docker/Dockerfile.cpu \
26 --build-arg max_jobs=16 \
.buildkite/image_build/image_build_cpu_arm64.sh
-6
@@ -21,12 +21,6 @@ else
21 exit 0
22 fi
23
24 -# The rust frontend lives in a git submodule under rust/. Buildkite's default
25 -# checkout does not recurse submodules, and the Dockerfile only sees what's in
26 -# the build context, so initialize the submodule here before building.
27 -git submodule sync --recursive
28 -git submodule update --init --recursive
29 -
24 # build
25 docker build --file docker/Dockerfile.cpu \
26 --build-arg max_jobs=16 \
.buildkite/scripts/run-rust-frontend-cargo-ci.sh new
+156
@@ -0,0 +1,156 @@
1 +#!/usr/bin/env bash
2 +set -euo pipefail
3 +
4 +MODE="${1:-}"
5 +
6 +if [[ "$MODE" != "style-clippy" && "$MODE" != "test" ]]; then
7 + echo "Usage: $0 {style-clippy|test}" >&2
8 + exit 2
9 +fi
10 +
11 +ROOT_DIR="$(git rev-parse --show-toplevel)"
12 +cd "$ROOT_DIR"
13 +
14 +export CARGO_TERM_COLOR="${CARGO_TERM_COLOR:-always}"
15 +export CARGO_HOME="${CARGO_HOME:-$HOME/.cargo}"
16 +export RUSTUP_HOME="${RUSTUP_HOME:-$HOME/.rustup}"
17 +export PATH="$CARGO_HOME/bin:$PATH"
18 +
19 +log_section() {
20 + echo "--- $*"
21 +}
22 +
23 +install_protoc() {
24 + if command -v protoc >/dev/null 2>&1; then
25 + return
26 + fi
27 +
28 + local version="${PROTOC_VERSION:-31.1}"
29 + local arch
30 + case "$(uname -m)" in
31 + x86_64)
32 + arch="x86_64"
33 + ;;
34 + aarch64|arm64)
35 + arch="aarch_64"
36 + ;;
37 + *)
38 + echo "Unsupported protoc architecture: $(uname -m)" >&2
39 + return 1
40 + ;;
41 + esac
42 +
43 + local url="https://github.com/protocolbuffers/protobuf/releases/download/v${version}/protoc-${version}-linux-${arch}.zip"
44 + local tmp_dir
45 + tmp_dir="$(mktemp -d)"
46 +
47 + log_section "Installing protoc ${version}"
48 + curl -L --proto '=https' --tlsv1.2 -sSf "$url" -o "$tmp_dir/protoc.zip"
49 + mkdir -p "$CARGO_HOME/bin"
50 + unzip -q "$tmp_dir/protoc.zip" bin/protoc 'include/*' -d "$CARGO_HOME"
51 + chmod +x "$CARGO_HOME/bin/protoc"
52 + rm -rf "$tmp_dir"
53 +}
54 +
55 +rust_toolchain() {
56 + awk -F '"' '/channel[[:space:]]*=/ { print $2; exit }' rust-toolchain.toml
57 +}
58 +
59 +install_rust_toolchain() {
60 + log_section "Installing Rust toolchain"
61 + if ! command -v rustup >/dev/null 2>&1; then
62 + curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs \
63 + | sh -s -- -y --profile minimal --default-toolchain none
64 + fi
65 +
66 + local toolchain
67 + toolchain="$(rust_toolchain)"
68 + rustup toolchain install "$toolchain" --profile minimal --component rustfmt,clippy
69 + rustup component add --toolchain "$toolchain" rustfmt clippy
70 +}
71 +
72 +install_cargo_binstall() {
73 + if command -v cargo-binstall >/dev/null 2>&1; then
74 + return
75 + fi
76 +
77 + log_section "Installing cargo-binstall"
78 + curl -L --proto '=https' --tlsv1.2 -sSf \
79 + https://raw.githubusercontent.com/cargo-bins/cargo-binstall/main/install-from-binstall-release.sh \
80 + | bash
81 +}
82 +
83 +install_cargo_sort() {
84 + if command -v cargo-sort >/dev/null 2>&1; then
85 + return
86 + fi
87 +
88 + log_section "Installing cargo-sort"
89 + install_cargo_binstall
90 + cargo binstall --no-confirm cargo-sort
91 +}
92 +
93 +install_cargo_nextest() {
94 + if command -v cargo-nextest >/dev/null 2>&1; then
95 + return
96 + fi
97 +
98 + log_section "Installing cargo-nextest"
99 + install_cargo_binstall
100 + cargo binstall --no-confirm --secure cargo-nextest
101 +}
102 +
103 +install_uv() {
104 + if command -v uv >/dev/null 2>&1; then
105 + return
106 + fi
107 +
108 + log_section "Installing uv"
109 + curl -LsSf --proto '=https' --tlsv1.2 https://astral.sh/uv/install.sh \
110 + | env UV_INSTALL_DIR="$CARGO_HOME/bin" sh
111 +}
112 +
113 +run_style_clippy() {
114 + install_cargo_sort
115 +
116 + log_section "Checking Rust formatting"
117 + cargo fmt --manifest-path rust/Cargo.toml --all -- --check
118 +
119 + log_section "Checking Cargo.toml ordering"
120 + cargo sort --workspace --check rust
121 +
122 + log_section "Running clippy"
123 + cargo clippy \
124 + --manifest-path rust/Cargo.toml \
125 + --workspace \
126 + --all-targets \
127 + --all-features \
128 + --locked \
129 + -- \
130 + -D warnings
131 +}
132 +
133 +run_tests() {
134 + install_uv
135 + install_cargo_nextest
136 +
137 + log_section "Running cargo nextest"
138 + cargo nextest run \
139 + --manifest-path rust/Cargo.toml \
140 + --workspace \
141 + --all-features \
142 + --locked \
143 + --no-fail-fast
144 +}
145 +
146 +install_protoc
147 +install_rust_toolchain
148 +
149 +case "$MODE" in
150 + style-clippy)
151 + run_style_clippy
152 + ;;
153 + test)
154 + run_tests
155 + ;;
156 +esac
.buildkite/test_areas/rust_frontend.yaml
+1 -1
@@ -1,4 +1,4 @@
1 -group: Rust Frontend
1 +group: Rust Frontend E2E
2 depends_on:
3 - image-build
4 steps:
.buildkite/test_areas/rust_frontend_cargo.yaml new
+30
@@ -0,0 +1,30 @@
1 +group: Rust Frontend Cargo
2 +depends_on: []
3 +steps:
4 +- label: Rust Frontend Cargo Style + Clippy
5 + key: rust-frontend-cargo-style-clippy
6 + depends_on: []
7 + timeout_in_minutes: 30
8 + device: cpu-medium
9 + no_plugin: true
10 + source_file_dependencies:
11 + - rust/
12 + - rust-toolchain.toml
13 + - .buildkite/test_areas/rust_frontend_cargo.yaml
14 + - .buildkite/scripts/run-rust-frontend-cargo-ci.sh
15 + commands:
16 + - .buildkite/scripts/run-rust-frontend-cargo-ci.sh style-clippy
17 +
18 +- label: Rust Frontend Cargo Tests
19 + key: rust-frontend-cargo-tests
20 + depends_on: []
21 + timeout_in_minutes: 30
22 + device: cpu-medium
23 + no_plugin: true
24 + source_file_dependencies:
25 + - rust/
26 + - rust-toolchain.toml
27 + - .buildkite/test_areas/rust_frontend_cargo.yaml
28 + - .buildkite/scripts/run-rust-frontend-cargo-ci.sh
29 + commands:
30 + - .buildkite/scripts/run-rust-frontend-cargo-ci.sh test
.gitmodules deleted
-3
@@ -1,3 +0,0 @@
1 -[submodule "rust"]
2 - path = rust
3 - url = https://github.com/Inferact/vllm-frontend-rs.git
.pre-commit-config.yaml
+26
@@ -256,6 +256,32 @@ repos:
256 entry: python tools/pre_commit/check_boolean_context_manager.py
257 language: python
258 types: [python]
259 + # Rust hooks. These shell out to `cargo`; tools/pre_commit/rust-check.sh
260 + # skips with a warning when cargo is not installed.
261 + - id: rust-cargo-autoinherit
262 + name: Rust - Normalize Cargo manifests with autoinherit
263 + entry: tools/pre_commit/rust-check.sh autoinherit --prefer-simple-dotted
264 + language: script
265 + pass_filenames: false
266 + require_serial: true
267 + stages: [pre-commit] # Only run locally as Buildkite will cover this
268 + files: ^rust/(Cargo\.toml|src/.*/Cargo\.toml)$
269 + - id: rust-cargo-sort
270 + name: Rust - Sort Cargo manifest sections
271 + entry: tools/pre_commit/rust-check.sh sort --workspace
272 + language: script
273 + pass_filenames: false
274 + require_serial: true
275 + stages: [pre-commit] # Only run locally as Buildkite will cover this
276 + files: ^rust/(Cargo\.toml|src/.*/Cargo\.toml)$
277 + - id: rust-cargo-fmt
278 + name: Rust - Format code
279 + entry: tools/pre_commit/rust-check.sh fmt
280 + language: script
281 + pass_filenames: false
282 + require_serial: true
283 + stages: [pre-commit] # Only run locally as Buildkite will cover this
284 + files: ^rust/.*(\.rs|Cargo\.toml|rustfmt\.toml)$
285 # Keep `suggestion` last
286 - id: suggestion
287 name: Suggestion
build_rust.sh
+3 -2
@@ -9,10 +9,10 @@ set -euo pipefail
9
10 REPO_ROOT="$(cd "$(dirname "$0")" && pwd)"
11 RUST_DIR="$REPO_ROOT/rust"
12 -TARGET_PATH="$REPO_ROOT/vllm/vllm-rs"
12 +TARGET_PATH="${VLLM_RS_TARGET_PATH:-$REPO_ROOT/vllm/vllm-rs}"
13
14 # Read the required toolchain from rust-toolchain.toml.
15 -TOOLCHAIN=$(grep '^channel' "$RUST_DIR/rust-toolchain.toml" | sed 's/.*= *"\(.*\)"/\1/')
15 +TOOLCHAIN=$(grep '^channel' "$REPO_ROOT/rust-toolchain.toml" | sed 's/.*= *"\(.*\)"/\1/')
16
17 # Ensure rustup and the required toolchain are available.
18 if ! command -v rustup &>/dev/null; then
@@ -39,5 +39,6 @@ cargo +"$TOOLCHAIN" build "${PROFILE_ARGS[@]}" \
39 --bin vllm-rs \
40 --features native-tls-vendored
41
42 +mkdir -p "$(dirname "$TARGET_PATH")"
43 cp "$RUST_DIR/target/$PROFILE_DIR/vllm-rs" "$TARGET_PATH"
44 echo "Installed vllm-rs to $TARGET_PATH"
docker/Dockerfile
+4 -12
@@ -256,7 +256,7 @@ RUN if [ "${BUILD_OS}" = "manylinux" ]; then \
256 COPY tools/install_protoc.sh /tmp/install_protoc.sh
257 RUN /tmp/install_protoc.sh && rm /tmp/install_protoc.sh
258
259 -# Install rustup; the toolchain itself is pinned by rust/rust-toolchain.toml.
259 +# Install rustup; the toolchain itself is pinned by rust-toolchain.toml.
260 RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | \
261 sh -s -- -y --profile minimal --default-toolchain none
262 ENV PATH="/root/.cargo/bin:${PATH}"
@@ -265,14 +265,8 @@ WORKDIR /workspace
265
266 # Copy only the rust workspace — the binary is the sole artifact we need.
267 COPY rust rust
268 -
269 -# Fail loudly if the rust submodule was not initialized on the host before
270 -# `docker build`. Without this check, cargo would emit a confusing error.
271 -RUN if [ ! -f rust/Cargo.toml ]; then \
272 - echo "ERROR: rust/ submodule is not initialized."; \
273 - echo "Run 'git submodule update --init --recursive' on the host before building."; \
274 - exit 1; \
275 - fi
268 +COPY rust-toolchain.toml rust-toolchain.toml
269 +COPY build_rust.sh build_rust.sh
270
271 # Cap cargo parallelism to avoid exhausting the CI host's open-file limit
272 # (rustc spawns enough concurrent processes to hit RLIMIT_NOFILE otherwise).
@@ -284,9 +278,7 @@ ENV CARGO_BUILD_JOBS=4
278 RUN --mount=type=cache,target=/root/.cargo/registry \
279 --mount=type=cache,target=/root/.cargo/git \
280 --mount=type=cache,target=/workspace/rust/target \
287 - cd rust \
288 - && cargo build --release --bin vllm-rs --features native-tls-vendored \
289 - && cp target/release/vllm-rs /workspace/vllm-rs
281 + VLLM_RS_TARGET_PATH=/workspace/vllm-rs bash build_rust.sh
282 #################### RUST BUILD IMAGE ####################
283
284 #################### CSRC BUILD IMAGE ####################
docker/Dockerfile.cpu
+4 -12
@@ -96,7 +96,7 @@ RUN apt-get update -y \
96 COPY tools/install_protoc.sh /tmp/install_protoc.sh
97 RUN /tmp/install_protoc.sh && rm /tmp/install_protoc.sh
98
99 -# Install rustup; the toolchain itself is pinned by rust/rust-toolchain.toml.
99 +# Install rustup; the toolchain itself is pinned by rust-toolchain.toml.
100 RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | \
101 sh -s -- -y --profile minimal --default-toolchain none
102 ENV PATH="/root/.cargo/bin:${PATH}"
@@ -105,14 +105,8 @@ WORKDIR /workspace
105
106 # Copy only the rust workspace — the binary is the sole artifact we need.
107 COPY rust rust
108 -
109 -# Fail loudly if the rust submodule was not initialized on the host before
110 -# `docker build`. Without this check, cargo would emit a confusing error.
111 -RUN if [ ! -f rust/Cargo.toml ]; then \
112 - echo "ERROR: rust/ submodule is not initialized."; \
113 - echo "Run 'git submodule update --init --recursive' on the host before building."; \
114 - exit 1; \
115 - fi
108 +COPY rust-toolchain.toml rust-toolchain.toml
109 +COPY build_rust.sh build_rust.sh
110
111 # Cap cargo parallelism to avoid exhausting the CI host's open-file limit
112 # (rustc spawns enough concurrent processes to hit RLIMIT_NOFILE otherwise).
@@ -124,9 +118,7 @@ ENV CARGO_BUILD_JOBS=4
118 RUN --mount=type=cache,target=/root/.cargo/registry \
119 --mount=type=cache,target=/root/.cargo/git \
120 --mount=type=cache,target=/workspace/rust/target \
127 - cd rust \
128 - && cargo build --release --bin vllm-rs --features native-tls-vendored \
129 - && cp target/release/vllm-rs /workspace/vllm-rs
121 + VLLM_RS_TARGET_PATH=/workspace/vllm-rs bash build_rust.sh
122
123 ######################### BUILD IMAGE #########################
124 FROM base AS vllm-build
docker/Dockerfile.nightly_torch
+4 -12
@@ -108,7 +108,7 @@ RUN apt-get update -y \
108 COPY tools/install_protoc.sh /tmp/install_protoc.sh
109 RUN /tmp/install_protoc.sh && rm /tmp/install_protoc.sh
110
111 -# Install rustup; the toolchain itself is pinned by rust/rust-toolchain.toml.
111 +# Install rustup; the toolchain itself is pinned by rust-toolchain.toml.
112 RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | \
113 sh -s -- -y --profile minimal --default-toolchain none
114 ENV PATH="/root/.cargo/bin:${PATH}"
@@ -116,14 +116,8 @@ ENV PATH="/root/.cargo/bin:${PATH}"
116 WORKDIR /workspace
117
118 COPY rust rust
119 -
120 -# Fail loudly if the rust submodule was not initialized on the host before
121 -# `docker build`.
122 -RUN if [ ! -f rust/Cargo.toml ]; then \
123 - echo "ERROR: rust/ submodule is not initialized."; \
124 - echo "Run 'git submodule update --init --recursive' on the host before building."; \
125 - exit 1; \
126 - fi
119 +COPY rust-toolchain.toml rust-toolchain.toml
120 +COPY build_rust.sh build_rust.sh
121
122 # Cap cargo parallelism to avoid exhausting the CI host's open-file limit
123 # (rustc spawns enough concurrent processes to hit RLIMIT_NOFILE otherwise).
@@ -132,9 +126,7 @@ ENV CARGO_BUILD_JOBS=4
126 RUN --mount=type=cache,target=/root/.cargo/registry \
127 --mount=type=cache,target=/root/.cargo/git \
128 --mount=type=cache,target=/workspace/rust/target \
135 - cd rust \
136 - && cargo build --release --bin vllm-rs --features native-tls-vendored \
137 - && cp target/release/vllm-rs /workspace/vllm-rs
129 + VLLM_RS_TARGET_PATH=/workspace/vllm-rs bash build_rust.sh
130 #################### RUST BUILD IMAGE ####################
131
132 #################### WHEEL BUILD IMAGE ####################
docker/Dockerfile.rocm
+4 -15
@@ -106,7 +106,6 @@ ONBUILD RUN git clone ${VLLM_REPO} \
106 && cd vllm \
107 && git fetch -v --prune -- origin ${VLLM_BRANCH} \
108 && git checkout FETCH_HEAD \
109 - && git submodule update --init --recursive \
109 && if [ ${VLLM_REPO} != "https://github.com/vllm-project/vllm.git" ] ; then \
110 git remote add upstream "https://github.com/vllm-project/vllm.git" \
111 && git fetch upstream ; fi
@@ -120,16 +119,7 @@ FROM fetch_vllm_${REMOTE_VLLM} AS fetch_vllm
119 FROM fetch_vllm AS rust-build
120 ARG COMMON_WORKDIR
121
123 -# Fail loudly if the rust submodule was not initialized on the host before
124 -# `docker build`. The rust frontend source is brought in via the fetch_vllm
125 -# stage, so an uninitialized submodule would otherwise produce a confusing
126 -# cargo failure.
127 -RUN if [ ! -f ${COMMON_WORKDIR}/vllm/rust/Cargo.toml ]; then \
128 - echo "ERROR: rust/ submodule is not initialized."; \
129 - echo "Run 'git submodule update --init --recursive' on the host before building."; \
130 - exit 1; \
131 - fi
132 -
122 +# protoc is used by tonic-build/prost-build.
123 RUN apt-get update -q -y && apt-get install -q -y --no-install-recommends \
124 ca-certificates curl unzip \
125 && rm -rf /var/lib/apt/lists/*
@@ -137,7 +127,7 @@ RUN apt-get update -q -y && apt-get install -q -y --no-install-recommends \
127 COPY tools/install_protoc.sh /tmp/install_protoc.sh
128 RUN /tmp/install_protoc.sh && rm /tmp/install_protoc.sh
129
140 -# Install rustup; the toolchain itself is pinned by rust/rust-toolchain.toml.
130 +# Install rustup; the toolchain itself is pinned by rust-toolchain.toml.
131 RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | \
132 sh -s -- -y --profile minimal --default-toolchain none
133 ENV PATH="/root/.cargo/bin:${PATH}"
@@ -150,9 +140,8 @@ ENV CARGO_BUILD_JOBS=4
140 # so it persists into the image layer for later COPY --from=rust-build.
141 RUN --mount=type=cache,target=/root/.cargo/registry \
142 --mount=type=cache,target=/root/.cargo/git \
153 - cd ${COMMON_WORKDIR}/vllm/rust \
154 - && cargo build --release --bin vllm-rs --features native-tls-vendored \
155 - && cp target/release/vllm-rs /tmp/vllm-rs
143 + cd ${COMMON_WORKDIR}/vllm \
144 + && VLLM_RS_TARGET_PATH=/tmp/vllm-rs bash build_rust.sh
145
146 # -----------------------
147 # vLLM build stages
docker/Dockerfile.xpu
+4 -12
@@ -12,7 +12,7 @@ RUN apt-get update -y \
12 COPY tools/install_protoc.sh /tmp/install_protoc.sh
13 RUN /tmp/install_protoc.sh && rm /tmp/install_protoc.sh
14
15 -# Install rustup; the toolchain itself is pinned by rust/rust-toolchain.toml.
15 +# Install rustup; the toolchain itself is pinned by rust-toolchain.toml.
16 RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | \
17 sh -s -- -y --profile minimal --default-toolchain none
18 ENV PATH="/root/.cargo/bin:${PATH}"
@@ -20,14 +20,8 @@ ENV PATH="/root/.cargo/bin:${PATH}"
20 WORKDIR /workspace
21
22 COPY rust rust
23 -
24 -# Fail loudly if the rust submodule was not initialized on the host before
25 -# `docker build`.
26 -RUN if [ ! -f rust/Cargo.toml ]; then \
27 - echo "ERROR: rust/ submodule is not initialized."; \
28 - echo "Run 'git submodule update --init --recursive' on the host before building."; \
29 - exit 1; \
30 - fi
23 +COPY rust-toolchain.toml rust-toolchain.toml
24 +COPY build_rust.sh build_rust.sh
25
26 # Cap cargo parallelism to avoid exhausting the CI host's open-file limit
27 # (rustc spawns enough concurrent processes to hit RLIMIT_NOFILE otherwise).
@@ -36,9 +30,7 @@ ENV CARGO_BUILD_JOBS=4
30 RUN --mount=type=cache,target=/root/.cargo/registry \
31 --mount=type=cache,target=/root/.cargo/git \
32 --mount=type=cache,target=/workspace/rust/target \
39 - cd rust \
40 - && cargo build --release --bin vllm-rs --features native-tls-vendored \
41 - && cp target/release/vllm-rs /workspace/vllm-rs
33 + VLLM_RS_TARGET_PATH=/workspace/vllm-rs bash build_rust.sh
34
35 FROM intel/deep-learning-essentials:2025.3.2-0-devel-ubuntu24.04 AS vllm-base
36
pyproject.toml
+4 -1
@@ -127,7 +127,10 @@ extend-exclude = ["tests/models/fixtures/*", "tests/prompts/*", "tests/tokenizer
127 "vllm/third_party/*", "vllm/entrypoints/serve/instrumentator/static/*",
128 "tests/entrypoints/speech_to_text/transcription/test_transcription_validation.py",
129 "docs/governance/process.md", "docs/assets/contributing/vllm_bench_serve_timeline.html",
130 - "tests/v1/engine/test_fast_incdec_prefix_err.py", ".git/*", "csrc/cpu/sgl-kernels/*"]
130 + "tests/v1/engine/test_fast_incdec_prefix_err.py", ".git/*", "csrc/cpu/sgl-kernels/*",
131 + "rust/src/chat/src/renderer/deepseek_v32/fixtures/*",
132 + "rust/src/tool-parser/src/gemma4.rs", "rust/src/text/src/output/decoded.rs",
133 + "rust/src/tokenizer/src/incremental.rs", "rust/src/reasoning-parser/src/tests.rs"]
134 ignore-hidden = false
135
136 [tool.typos.default]
requirements/xpu.txt
+1
@@ -5,6 +5,7 @@ ray>=2.9
5 cmake>=3.26.1
6 packaging>=24.2
7 setuptools-scm>=8
8 +setuptools-rust>=1.9.0
9 setuptools>=77.0.3,<81.0.0
10 setuptools-rust>=1.9.0
11 wheel
rust deleted
-1
@@ -1 +0,0 @@
1 -Subproject commit ad6771ac093420632253ca43bfc81007beb2d6ce
rust/.gitattributes new
+2
@@ -0,0 +1,2 @@
1 +src/chat/src/renderer/deepseek_v32/fixtures/** linguist-generated=true
2 +src/chat/tests/templates/** linguist-vendored=true
rust/.gitignore new
+3
@@ -0,0 +1,3 @@
1 +/target
2 +AGENTS.override.md
3 +.vscode
rust/AGENTS.md new
+37
@@ -0,0 +1,37 @@
1 +# Alternative Frontend to vLLM Engine in Rust
2 +
3 +This project aims to implement an alternative frontend to the vLLM Engine in Rust, providing a more efficient and robust interface for interacting with the engine. Currently it's still in the very early stage and is actively evolving.
4 +
5 +## Coding Styles
6 +
7 +- Always use workspace dependencies for Cargo crates.
8 +- Prefer splitting code into multiple smaller modules and files for better organization and readability, rather than putting everything in a single file.
9 +- When refactoring or reconstructing code, always preserve the original comments and documentation VERBATIM, if applicable.
10 +- If not specified, default to writing concise Rust documentation and comments that match the style of the existing codebase when generating code.
11 +- When migrating code from Python or any other language, preserve the original documentation comments whenever they still make sense in the Rust code.
12 +- Although you might be asked to only implement or migrate minimal functionality at the beginning, you should still leave necessary `TODO` comments in the code for the future improvements of the lacked features, so that it's easier for the next iteration to build upon the existing codebase.
13 +- When writing parsers with `winnow`:
14 + - Prefer a declarative parser shape over imperative step-by-step parsing, as long as it's more readable and maintainable.
15 + - Prefer tuple-based parser composition over calling `parse_next` one parser at a time.
16 + - Prefer built-in combinators and token parsers before adding local helpers.
17 + - Add short documentation comments like `Parse a ..` to all local parser/combinator functions.
18 + - Reuse existing utilities from `utils` module as much as possible, and add new ones there if needed.
19 +- Rust error handling:
20 + - Never call `to_string()` directly on an error value.
21 + - Use `ToReportString` or `AsReport` by `thiserror-ext` instead.
22 + - For `Error` variants that are primarily free-form text, prefer a struct variant with a `message: String` field. `thiserror_ext::Macro` will auto-derive `foo!(...)` and `bail_foo!(...)` helper macros from that shape.
23 + - Use `foo!(...)` when you need to construct an error value inside an expression, such as `Err(foo!(...))`, `.ok_or_else(|| foo!(...))`, or `Err::<(), _>(foo!(...))?`.
24 + - Use `bail_foo!(...)` only in statement positions where you want to exit the current `Result`-returning function immediately. Prefer it over `return Err(foo!(...))` in those cases.
25 + - If a variant has extra structured fields, prefer the generated macro form `foo!(field = value, "message")` rather than manually writing `Error::Foo { ... }`.
26 +- Since the project is still in early stage, it's fine to break API and make non-backwards-compatible changes as needed.
27 +- Currently the project is only targeting Unix-like platforms, so it's fine to use Unix-specific APIs without extra compatibility layers like `cfg(unix)`
28 +
29 +## Testing
30 +
31 +- Prefer snapshot testing with the `expect-test` crate over writing multiple `assert_eq!` statements on individual fields. Use `expect_test::expect![[...]].assert_debug_eq(...)` to snapshot the `Debug` output of the entire struct.
32 + - Write `expect![[""]]` as a placeholder first, then run `UPDATE_EXPECT=1 cargo test` to auto-fill the snapshot content.
33 + - For values containing non-deterministic data (e.g., UUIDs), set them to a fixed value like `"<placeholder>"` before snapshotting.
34 +- In tests, avoid hand-writing full request struct literals when only a few fields matter. Prefer test fixtures such as `for_test()` with struct update syntax, so newly added fields do not force mechanical edits across many tests.
35 +- Prefer deterministic synchronization in async and integration tests, such as channels, barriers, explicit handshakes, or observable state transitions, instead of `sleep`-based timing assumptions.
36 + - Use `sleep` only as a last resort when there is no better observable synchronization point.
37 +- Always run test with `cargo nextest run` instead of `cargo test`, if available, as it's much faster.
rust/CLAUDE.md new
+4
@@ -0,0 +1,4 @@
1 +# CLAUDE.md
2 +
3 +First, check @AGENTS.override.md if exists.
4 +Then, follow instructions in @AGENTS.md.
rust/Cargo.lock new
+6611
@@ -0,0 +1,6611 @@
1 +# This file is automatically @generated by Cargo.
2 +# It is not intended for manual editing.
3 +version = 4
4 +
5 +[[package]]
6 +name = "adler2"
7 +version = "2.0.1"
8 +source = "registry+https://github.com/rust-lang/crates.io-index"
9 +checksum = "320119579fcad9c21884f5c4861d16174d0e06250625266f50fe6898340abefa"
10 +
11 +[[package]]
12 +name = "ahash"
13 +version = "0.8.12"
14 +source = "registry+https://github.com/rust-lang/crates.io-index"
15 +checksum = "5a15f179cd60c4584b8a8c596927aadc462e27f2ca70c04e0071964a73ba7a75"
16 +dependencies = [
17 + "cfg-if",
18 + "getrandom 0.3.4",
19 + "once_cell",
20 + "serde",
21 + "version_check",
22 + "zerocopy",
23 +]
24 +
25 +[[package]]
26 +name = "aho-corasick"
27 +version = "1.1.4"
28 +source = "registry+https://github.com/rust-lang/crates.io-index"
29 +checksum = "ddd31a130427c27518df266943a5308ed92d4b226cc639f5a8f1002816174301"
30 +dependencies = [
31 + "memchr",
32 +]
33 +
34 +[[package]]
35 +name = "aligned"
36 +version = "0.4.3"
37 +source = "registry+https://github.com/rust-lang/crates.io-index"
38 +checksum = "ee4508988c62edf04abd8d92897fca0c2995d907ce1dfeaf369dac3716a40685"
39 +dependencies = [
40 + "as-slice",
41 +]
42 +
43 +[[package]]
44 +name = "aligned-vec"
45 +version = "0.6.4"
46 +source = "registry+https://github.com/rust-lang/crates.io-index"
47 +checksum = "dc890384c8602f339876ded803c97ad529f3842aba97f6392b3dba0dd171769b"
48 +dependencies = [
49 + "equator",
50 +]
51 +
52 +[[package]]
53 +name = "android_system_properties"
54 +version = "0.1.5"
55 +source = "registry+https://github.com/rust-lang/crates.io-index"
56 +checksum = "819e7219dbd41043ac279b19830f2efc897156490d7fd6ea916720117ee66311"
57 +dependencies = [
58 + "libc",
59 +]
60 +
61 +[[package]]
62 +name = "anes"
63 +version = "0.1.6"
64 +source = "registry+https://github.com/rust-lang/crates.io-index"
65 +checksum = "4b46cbb362ab8752921c97e041f5e366ee6297bd428a31275b9fcf1e380f7299"
66 +
67 +[[package]]
68 +name = "anstream"
69 +version = "0.6.21"
70 +source = "registry+https://github.com/rust-lang/crates.io-index"
71 +checksum = "43d5b281e737544384e969a5ccad3f1cdd24b48086a0fc1b2a5262a26b8f4f4a"
72 +dependencies = [
73 + "anstyle",
74 + "anstyle-parse 0.2.7",
75 + "anstyle-query",
76 + "anstyle-wincon",
77 + "colorchoice",
78 + "is_terminal_polyfill",
79 + "utf8parse",
80 +]
81 +
82 +[[package]]
83 +name = "anstream"
84 +version = "1.0.0"
85 +source = "registry+https://github.com/rust-lang/crates.io-index"
86 +checksum = "824a212faf96e9acacdbd09febd34438f8f711fb84e09a8916013cd7815ca28d"
87 +dependencies = [
88 + "anstyle",
89 + "anstyle-parse 1.0.0",
90 + "anstyle-query",
91 + "anstyle-wincon",
92 + "colorchoice",
93 + "is_terminal_polyfill",
94 + "utf8parse",
95 +]
96 +
97 +[[package]]
98 +name = "anstyle"
99 +version = "1.0.13"
100 +source = "registry+https://github.com/rust-lang/crates.io-index"
101 +checksum = "5192cca8006f1fd4f7237516f40fa183bb07f8fbdfedaa0036de5ea9b0b45e78"
102 +
103 +[[package]]
104 +name = "anstyle-parse"
105 +version = "0.2.7"
106 +source = "registry+https://github.com/rust-lang/crates.io-index"
107 +checksum = "4e7644824f0aa2c7b9384579234ef10eb7efb6a0deb83f9630a49594dd9c15c2"
108 +dependencies = [
109 + "utf8parse",
110 +]
111 +
112 +[[package]]
113 +name = "anstyle-parse"
114 +version = "1.0.0"
115 +source = "registry+https://github.com/rust-lang/crates.io-index"
116 +checksum = "52ce7f38b242319f7cabaa6813055467063ecdc9d355bbb4ce0c68908cd8130e"
117 +dependencies = [
118 + "utf8parse",
119 +]
120 +
121 +[[package]]
122 +name = "anstyle-query"
123 +version = "1.1.5"
124 +source = "registry+https://github.com/rust-lang/crates.io-index"
125 +checksum = "40c48f72fd53cd289104fc64099abca73db4166ad86ea0b4341abe65af83dadc"
126 +dependencies = [
127 + "windows-sys 0.61.2",
128 +]
129 +
130 +[[package]]
131 +name = "anstyle-wincon"
132 +version = "3.0.11"
133 +source = "registry+https://github.com/rust-lang/crates.io-index"
134 +checksum = "291e6a250ff86cd4a820112fb8898808a366d8f9f58ce16d1f538353ad55747d"
135 +dependencies = [
136 + "anstyle",
137 + "once_cell_polyfill",
138 + "windows-sys 0.61.2",
139 +]
140 +
141 +[[package]]
142 +name = "anyhow"
143 +version = "1.0.102"
144 +source = "registry+https://github.com/rust-lang/crates.io-index"
145 +checksum = "7f202df86484c868dbad7eaa557ef785d5c66295e41b460ef922eca0723b842c"
146 +
147 +[[package]]
148 +name = "arbitrary"
149 +version = "1.4.2"
150 +source = "registry+https://github.com/rust-lang/crates.io-index"
151 +checksum = "c3d036a3c4ab069c7b410a2ce876bd74808d2d0888a82667669f8e783a898bf1"
152 +
153 +[[package]]
154 +name = "arc-swap"
155 +version = "1.9.0"
156 +source = "registry+https://github.com/rust-lang/crates.io-index"
157 +checksum = "a07d1f37ff60921c83bdfc7407723bdefe89b44b98a9b772f225c8f9d67141a6"
158 +dependencies = [
159 + "rustversion",
160 +]
161 +
162 +[[package]]
163 +name = "arg_enum_proc_macro"
164 +version = "0.3.4"
165 +source = "registry+https://github.com/rust-lang/crates.io-index"
166 +checksum = "0ae92a5119aa49cdbcf6b9f893fe4e1d98b04ccbf82ee0584ad948a44a734dea"
167 +dependencies = [
168 + "proc-macro2",
169 + "quote",
170 + "syn 2.0.117",
171 +]
172 +
173 +[[package]]
174 +name = "arrayref"
175 +version = "0.3.9"
176 +source = "registry+https://github.com/rust-lang/crates.io-index"
177 +checksum = "76a2e8124351fda1ef8aaaa3bbd7ebbcb486bbcd4225aca0aa0d84bb2db8fecb"
178 +
179 +[[package]]
180 +name = "arrayvec"
181 +version = "0.7.6"
182 +source = "registry+https://github.com/rust-lang/crates.io-index"
183 +checksum = "7c02d123df017efcdfbd739ef81735b36c5ba83ec3c59c80a9d7ecc718f92e50"
184 +
185 +[[package]]
186 +name = "as-slice"
187 +version = "0.2.1"
188 +source = "registry+https://github.com/rust-lang/crates.io-index"
189 +checksum = "516b6b4f0e40d50dcda9365d53964ec74560ad4284da2e7fc97122cd83174516"
190 +dependencies = [
191 + "stable_deref_trait",
192 +]
193 +
194 +[[package]]
195 +name = "async-io"
196 +version = "2.6.0"
197 +source = "registry+https://github.com/rust-lang/crates.io-index"
198 +checksum = "456b8a8feb6f42d237746d4b3e9a178494627745c3c56c6ea55d92ba50d026fc"
199 +dependencies = [
200 + "autocfg",
201 + "cfg-if",
202 + "concurrent-queue",
203 + "futures-io",
204 + "futures-lite",
205 + "parking",
206 + "polling",
207 + "rustix",
208 + "slab",
209 + "windows-sys 0.61.2",
210 +]
211 +
212 +[[package]]
213 +name = "async-openai"
214 +version = "0.33.1"
215 +source = "registry+https://github.com/rust-lang/crates.io-index"
216 +checksum = "cc48c3deb4ad9a2ee8c8e364c79eb0f74e69e17ed7e883d55988b90ea44fe986"
217 +dependencies = [
218 + "async-openai-macros",
219 + "backoff",
220 + "base64 0.22.1",
221 + "bytes",
222 + "derive_builder",
223 + "eventsource-stream",
224 + "futures",
225 + "getrandom 0.3.4",
226 + "hex",
227 + "hmac",
228 + "rand 0.9.2",
229 + "reqwest",
230 + "reqwest-eventsource",
231 + "secrecy",
232 + "serde",
233 + "serde_json",
234 + "serde_urlencoded",
235 + "sha2",
236 + "thiserror 2.0.18",
237 + "tokio",
238 + "tokio-stream",
239 + "tokio-tungstenite",
240 + "tokio-util",
241 + "tracing",
242 + "url",
243 +]
244 +
245 +[[package]]
246 +name = "async-openai-macros"
247 +version = "0.1.1"
248 +source = "registry+https://github.com/rust-lang/crates.io-index"
249 +checksum = "81872a8e595e8ceceab71c6ba1f9078e313b452a1e31934e6763ef5d308705e4"
250 +dependencies = [
251 + "proc-macro2",
252 + "quote",
253 + "syn 2.0.117",
254 +]
255 +
256 +[[package]]
257 +name = "async-trait"
258 +version = "0.1.89"
259 +source = "registry+https://github.com/rust-lang/crates.io-index"
260 +checksum = "9035ad2d096bed7955a320ee7e2230574d28fd3c3a0f186cbea1ff3c7eed5dbb"
261 +dependencies = [
262 + "proc-macro2",
263 + "quote",
264 + "syn 2.0.117",
265 +]
266 +
267 +[[package]]
268 +name = "asynchronous-codec"
269 +version = "0.7.0"
270 +source = "registry+https://github.com/rust-lang/crates.io-index"
271 +checksum = "a860072022177f903e59730004fb5dc13db9275b79bb2aef7ba8ce831956c233"
272 +dependencies = [
273 + "bytes",
274 + "futures-sink",
275 + "futures-util",
276 + "memchr",
277 + "pin-project-lite",
278 +]
279 +
280 +[[package]]
281 +name = "asynk-strim"
282 +version = "0.1.5"
283 +source = "registry+https://github.com/rust-lang/crates.io-index"
284 +checksum = "52697735bdaac441a29391a9e97102c74c6ef0f9b60a40cf109b1b404e29d2f6"
285 +dependencies = [
286 + "futures-core",
287 + "pin-project-lite",
288 +]
289 +
290 +[[package]]
291 +name = "asynk-strim-attr"
292 +version = "0.1.0"
293 +source = "registry+https://github.com/rust-lang/crates.io-index"
294 +checksum = "b6ccb67be092524ce594e599332719f1cd6d64dcaed8d46f1e8726d466c10bcb"
295 +dependencies = [
296 + "asynk-strim",
297 + "asynk-strim-attr-macro",
298 + "futures-core",
299 +]
300 +
301 +[[package]]
302 +name = "asynk-strim-attr-macro"
303 +version = "0.1.0"
304 +source = "registry+https://github.com/rust-lang/crates.io-index"
305 +checksum = "34e40a2181bb16fb68e25c49c8b3e25bbb9a808bf8f9f83bc596ac4ad70c86a1"
306 +dependencies = [
307 + "proc-macro-crate",
308 + "proc-macro2",
309 + "quote",
310 + "syn 2.0.117",
311 +]
312 +
313 +[[package]]
314 +name = "atomic-waker"
315 +version = "1.1.2"
316 +source = "registry+https://github.com/rust-lang/crates.io-index"
317 +checksum = "1505bd5d3d116872e7271a6d4e16d81d0c8570876c8de68093a09ac269d8aac0"
318 +
319 +[[package]]
320 +name = "autocfg"
321 +version = "1.5.0"
322 +source = "registry+https://github.com/rust-lang/crates.io-index"
323 +checksum = "c08606f8c3cbf4ce6ec8e28fb0014a2c086708fe954eaa885384a6165172e7e8"
324 +
325 +[[package]]
326 +name = "av-scenechange"
327 +version = "0.14.1"
328 +source = "registry+https://github.com/rust-lang/crates.io-index"
329 +checksum = "0f321d77c20e19b92c39e7471cf986812cbb46659d2af674adc4331ef3f18394"
330 +dependencies = [
331 + "aligned",
332 + "anyhow",
333 + "arg_enum_proc_macro",
334 + "arrayvec",
335 + "log",
336 + "num-rational",
337 + "num-traits",
338 + "pastey",
339 + "rayon",
340 + "thiserror 2.0.18",
341 + "v_frame",
342 + "y4m",
343 +]
344 +
345 +[[package]]
346 +name = "av1-grain"
347 +version = "0.2.5"
348 +source = "registry+https://github.com/rust-lang/crates.io-index"
349 +checksum = "8cfddb07216410377231960af4fcab838eaa12e013417781b78bd95ee22077f8"
350 +dependencies = [
351 + "anyhow",
352 + "arrayvec",
353 + "log",
354 + "nom 8.0.0",
355 + "num-rational",
356 + "v_frame",
357 +]
358 +
359 +[[package]]
360 +name = "avif-serialize"
361 +version = "0.8.8"
362 +source = "registry+https://github.com/rust-lang/crates.io-index"
363 +checksum = "375082f007bd67184fb9c0374614b29f9aaa604ec301635f72338bb65386a53d"
364 +dependencies = [
365 + "arrayvec",
366 +]
367 +
368 +[[package]]
369 +name = "axum"
370 +version = "0.8.8"
371 +source = "registry+https://github.com/rust-lang/crates.io-index"
372 +checksum = "8b52af3cb4058c895d37317bb27508dccc8e5f2d39454016b297bf4a400597b8"
373 +dependencies = [
374 + "axum-core",
375 + "bytes",
376 + "form_urlencoded",
377 + "futures-util",
378 + "http",
379 + "http-body",
380 + "http-body-util",
381 + "hyper",
382 + "hyper-util",
383 + "itoa",
384 + "matchit",
385 + "memchr",
386 + "mime",
387 + "percent-encoding",
388 + "pin-project-lite",
389 + "serde_core",
390 + "serde_json",
391 + "serde_path_to_error",
392 + "serde_urlencoded",
393 + "sync_wrapper",
394 + "tokio",
395 + "tower",
396 + "tower-layer",
397 + "tower-service",
398 + "tracing",
399 +]
400 +
401 +[[package]]
402 +name = "axum-core"
403 +version = "0.5.6"
404 +source = "registry+https://github.com/rust-lang/crates.io-index"
405 +checksum = "08c78f31d7b1291f7ee735c1c6780ccde7785daae9a9206026862dab7d8792d1"
406 +dependencies = [
407 + "bytes",
408 + "futures-core",
409 + "http",
410 + "http-body",
411 + "http-body-util",
412 + "mime",
413 + "pin-project-lite",
414 + "sync_wrapper",
415 + "tower-layer",
416 + "tower-service",
417 + "tracing",
418 +]
419 +
420 +[[package]]
421 +name = "backoff"
422 +version = "0.4.0"
423 +source = "registry+https://github.com/rust-lang/crates.io-index"
424 +checksum = "b62ddb9cb1ec0a098ad4bbf9344d0713fa193ae1a80af55febcff2627b6a00c1"
425 +dependencies = [
426 + "futures-core",
427 + "getrandom 0.2.17",
428 + "instant",
429 + "pin-project-lite",
430 + "rand 0.8.5",
431 + "tokio",
432 +]
433 +
434 +[[package]]
435 +name = "base64"
436 +version = "0.13.1"
437 +source = "registry+https://github.com/rust-lang/crates.io-index"
438 +checksum = "9e1b586273c5702936fe7b7d6896644d8be71e6314cfe09d3167c95f712589e8"
439 +
440 +[[package]]
441 +name = "base64"
442 +version = "0.22.1"
443 +source = "registry+https://github.com/rust-lang/crates.io-index"
444 +checksum = "72b3254f16251a8381aa12e40e3c4d2f0199f8c6508fbecb9d91f575e0fbb8c6"
445 +
446 +[[package]]
447 +name = "base64ct"
448 +version = "1.8.3"
449 +source = "registry+https://github.com/rust-lang/crates.io-index"
450 +checksum = "2af50177e190e07a26ab74f8b1efbfe2ef87da2116221318cb1c2e82baf7de06"
451 +
452 +[[package]]
453 +name = "bit-set"
454 +version = "0.5.3"
455 +source = "registry+https://github.com/rust-lang/crates.io-index"
456 +checksum = "0700ddab506f33b20a03b13996eccd309a48e5ff77d0d95926aa0210fb4e95f1"
457 +dependencies = [
458 + "bit-vec 0.6.3",
459 +]
460 +
461 +[[package]]
462 +name = "bit-set"
463 +version = "0.8.0"
464 +source = "registry+https://github.com/rust-lang/crates.io-index"
465 +checksum = "08807e080ed7f9d5433fa9b275196cfc35414f66a0c79d864dc51a0d825231a3"
466 +dependencies = [
467 + "bit-vec 0.8.0",
468 +]
469 +
470 +[[package]]
471 +name = "bit-vec"
472 +version = "0.6.3"
473 +source = "registry+https://github.com/rust-lang/crates.io-index"
474 +checksum = "349f9b6a179ed607305526ca489b34ad0a41aed5f7980fa90eb03160b69598fb"
475 +
476 +[[package]]
477 +name = "bit-vec"
478 +version = "0.8.0"
479 +source = "registry+https://github.com/rust-lang/crates.io-index"
480 +checksum = "5e764a1d40d510daf35e07be9eb06e75770908c27d411ee6c92109c9840eaaf7"
481 +
482 +[[package]]
483 +name = "bit_field"
484 +version = "0.10.3"
485 +source = "registry+https://github.com/rust-lang/crates.io-index"
486 +checksum = "1e4b40c7323adcfc0a41c4b88143ed58346ff65a288fc144329c5c45e05d70c6"
487 +
488 +[[package]]
489 +name = "bitflags"
490 +version = "2.11.0"
491 +source = "registry+https://github.com/rust-lang/crates.io-index"
492 +checksum = "843867be96c8daad0d758b57df9392b6d8d271134fce549de6ce169ff98a92af"
493 +
494 +[[package]]
495 +name = "bitstream-io"
496 +version = "4.10.0"
497 +source = "registry+https://github.com/rust-lang/crates.io-index"
498 +checksum = "7eff00be299a18769011411c9def0d827e8f2d7bf0c3dbf53633147a8867fd1f"
499 +dependencies = [
500 + "no_std_io2",
501 +]
502 +
503 +[[package]]
504 +name = "blake3"
505 +version = "1.8.5"
506 +source = "registry+https://github.com/rust-lang/crates.io-index"
507 +checksum = "0aa83c34e62843d924f905e0f5c866eb1dd6545fc4d719e803d9ba6030371fce"
508 +dependencies = [
509 + "arrayref",
510 + "arrayvec",
511 + "cc",
512 + "cfg-if",
513 + "constant_time_eq",
514 + "cpufeatures 0.3.0",
515 +]
516 +
517 +[[package]]
518 +name = "block-buffer"
519 +version = "0.10.4"
520 +source = "registry+https://github.com/rust-lang/crates.io-index"
521 +checksum = "3078c7629b62d3f0439517fa394996acacc5cbc91c5a20d8c658e77abd503a71"
522 +dependencies = [
523 + "generic-array",
524 +]
525 +
526 +[[package]]
527 +name = "bstr"
528 +version = "1.12.1"
529 +source = "registry+https://github.com/rust-lang/crates.io-index"
530 +checksum = "63044e1ae8e69f3b5a92c736ca6269b8d12fa7efe39bf34ddb06d102cf0e2cab"
531 +dependencies = [
532 + "memchr",
533 + "regex-automata",
534 + "serde",
535 +]
536 +
537 +[[package]]
538 +name = "built"
539 +version = "0.8.0"
540 +source = "registry+https://github.com/rust-lang/crates.io-index"
541 +checksum = "f4ad8f11f288f48ca24471bbd51ac257aaeaaa07adae295591266b792902ae64"
542 +
543 +[[package]]
544 +name = "bumpalo"
545 +version = "3.20.2"
546 +source = "registry+https://github.com/rust-lang/crates.io-index"
547 +checksum = "5d20789868f4b01b2f2caec9f5c4e0213b41e3e5702a50157d699ae31ced2fcb"
548 +
549 +[[package]]
550 +name = "bytemuck"
551 +version = "1.25.0"
552 +source = "registry+https://github.com/rust-lang/crates.io-index"
553 +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec"
554 +dependencies = [
555 + "bytemuck_derive",
556 +]
557 +
558 +[[package]]
559 +name = "bytemuck_derive"
560 +version = "1.10.2"
561 +source = "registry+https://github.com/rust-lang/crates.io-index"
562 +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff"
563 +dependencies = [
564 + "proc-macro2",
565 + "quote",
566 + "syn 2.0.117",
567 +]
568 +
569 +[[package]]
570 +name = "byteorder"
571 +version = "1.5.0"
572 +source = "registry+https://github.com/rust-lang/crates.io-index"
573 +checksum = "1fd0f2584146f6f2ef48085050886acf353beff7305ebd1ae69500e27c67f64b"
574 +
575 +[[package]]
576 +name = "byteorder-lite"
577 +version = "0.1.0"
578 +source = "registry+https://github.com/rust-lang/crates.io-index"
579 +checksum = "8f1fe948ff07f4bd06c30984e69f5b4899c516a3ef74f34df92a2df2ab535495"
580 +
581 +[[package]]
582 +name = "bytes"
583 +version = "1.11.1"
584 +source = "registry+https://github.com/rust-lang/crates.io-index"
585 +checksum = "1e748733b7cbc798e1434b6ac524f0c1ff2ab456fe201501e6497c8417a4fc33"
586 +dependencies = [
587 + "serde",
588 +]
589 +
590 +[[package]]
591 +name = "cast"
592 +version = "0.3.0"
593 +source = "registry+https://github.com/rust-lang/crates.io-index"
594 +checksum = "37b2a672a2cb129a2e41c10b1224bb368f9f37a2b16b612598138befd7b37eb5"
595 +
596 +[[package]]
597 +name = "castaway"
598 +version = "0.2.4"
599 +source = "registry+https://github.com/rust-lang/crates.io-index"
600 +checksum = "dec551ab6e7578819132c713a93c022a05d60159dc86e7a7050223577484c55a"
601 +dependencies = [
602 + "rustversion",
603 +]
604 +
605 +[[package]]
606 +name = "cc"
607 +version = "1.2.56"
608 +source = "registry+https://github.com/rust-lang/crates.io-index"
609 +checksum = "aebf35691d1bfb0ac386a69bac2fde4dd276fb618cf8bf4f5318fe285e821bb2"
610 +dependencies = [
611 + "find-msvc-tools",
612 + "jobserver",
613 + "libc",
614 + "shlex",
615 +]
616 +
617 +[[package]]
618 +name = "cfg-if"
619 +version = "1.0.4"
620 +source = "registry+https://github.com/rust-lang/crates.io-index"
621 +checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801"
622 +
623 +[[package]]
624 +name = "cfg_aliases"
625 +version = "0.2.1"
626 +source = "registry+https://github.com/rust-lang/crates.io-index"
627 +checksum = "613afe47fcd5fac7ccf1db93babcb082c5994d996f20b8b159f2ad1658eb5724"
628 +
629 +[[package]]
630 +name = "chrono"
631 +version = "0.4.44"
632 +source = "registry+https://github.com/rust-lang/crates.io-index"
633 +checksum = "c673075a2e0e5f4a1dde27ce9dee1ea4558c7ffe648f576438a20ca1d2acc4b0"
634 +dependencies = [
635 + "iana-time-zone",
636 + "js-sys",
637 + "num-traits",
638 + "serde",
639 + "wasm-bindgen",
640 + "windows-link",
641 +]
642 +
643 +[[package]]
644 +name = "ciborium"
645 +version = "0.2.2"
646 +source = "registry+https://github.com/rust-lang/crates.io-index"
647 +checksum = "42e69ffd6f0917f5c029256a24d0161db17cea3997d185db0d35926308770f0e"
648 +dependencies = [
649 + "ciborium-io",
650 + "ciborium-ll",
651 + "serde",
652 +]
653 +
654 +[[package]]
655 +name = "ciborium-io"
656 +version = "0.2.2"
657 +source = "registry+https://github.com/rust-lang/crates.io-index"
658 +checksum = "05afea1e0a06c9be33d539b876f1ce3692f4afea2cb41f740e7743225ed1c757"
659 +
660 +[[package]]
661 +name = "ciborium-ll"
662 +version = "0.2.2"
663 +source = "registry+https://github.com/rust-lang/crates.io-index"
664 +checksum = "57663b653d948a338bfb3eeba9bb2fd5fcfaecb9e199e87e1eda4d9e8b240fd9"
665 +dependencies = [
666 + "ciborium-io",
667 + "half",
668 +]
669 +
670 +[[package]]
671 +name = "clap"
672 +version = "4.5.60"
673 +source = "registry+https://github.com/rust-lang/crates.io-index"
674 +checksum = "2797f34da339ce31042b27d23607e051786132987f595b02ba4f6a6dffb7030a"
675 +dependencies = [
676 + "clap_builder",
677 + "clap_derive",
678 +]
679 +
680 +[[package]]
681 +name = "clap_builder"
682 +version = "4.5.60"
683 +source = "registry+https://github.com/rust-lang/crates.io-index"
684 +checksum = "24a241312cea5059b13574bb9b3861cabf758b879c15190b37b6d6fd63ab6876"
685 +dependencies = [
686 + "anstream 0.6.21",
687 + "anstyle",
688 + "clap_lex",
689 + "strsim",
690 +]
691 +
692 +[[package]]
693 +name = "clap_derive"
694 +version = "4.5.55"
695 +source = "registry+https://github.com/rust-lang/crates.io-index"
696 +checksum = "a92793da1a46a5f2a02a6f4c46c6496b28c43638adea8306fcb0caa1634f24e5"
697 +dependencies = [
698 + "heck",
699 + "proc-macro2",
700 + "quote",
701 + "syn 2.0.117",
702 +]
703 +
704 +[[package]]
705 +name = "clap_lex"
706 +version = "1.0.0"
707 +source = "registry+https://github.com/rust-lang/crates.io-index"
708 +checksum = "3a822ea5bc7590f9d40f1ba12c0dc3c2760f3482c6984db1573ad11031420831"
709 +
710 +[[package]]
711 +name = "color_quant"
712 +version = "1.1.0"
713 +source = "registry+https://github.com/rust-lang/crates.io-index"
714 +checksum = "3d7b894f5411737b7867f4827955924d7c254fc9f4d91a6aad6b097804b1018b"
715 +
716 +[[package]]
717 +name = "colorchoice"
718 +version = "1.0.4"
719 +source = "registry+https://github.com/rust-lang/crates.io-index"
720 +checksum = "b05b61dc5112cbb17e4b6cd61790d9845d13888356391624cbe7e41efeac1e75"
721 +
722 +[[package]]
723 +name = "compact_str"
724 +version = "0.9.0"
725 +source = "registry+https://github.com/rust-lang/crates.io-index"
726 +checksum = "3fdb1325a1cece981e8a296ab8f0f9b63ae357bd0784a9faaf548cc7b480707a"
727 +dependencies = [
728 + "castaway",
729 + "cfg-if",
730 + "itoa",
731 + "rustversion",
732 + "ryu",
733 + "serde",
734 + "static_assertions",
735 +]
736 +
737 +[[package]]
738 +name = "concurrent-queue"
739 +version = "2.5.0"
740 +source = "registry+https://github.com/rust-lang/crates.io-index"
741 +checksum = "4ca0197aee26d1ae37445ee532fefce43251d24cc7c166799f4d46817f1d3973"
742 +dependencies = [
743 + "crossbeam-utils",
744 +]
745 +
746 +[[package]]
747 +name = "console"
748 +version = "0.15.11"
749 +source = "registry+https://github.com/rust-lang/crates.io-index"
750 +checksum = "054ccb5b10f9f2cbf51eb355ca1d05c2d279ce1804688d0db74b4733a5aeafd8"
751 +dependencies = [
752 + "encode_unicode",
753 + "libc",
754 + "once_cell",
755 + "unicode-width",
756 + "windows-sys 0.59.0",
757 +]
758 +
759 +[[package]]
760 +name = "console"
761 +version = "0.16.2"
762 +source = "registry+https://github.com/rust-lang/crates.io-index"
763 +checksum = "03e45a4a8926227e4197636ba97a9fc9b00477e9f4bd711395687c5f0734bec4"
764 +dependencies = [
765 + "encode_unicode",
766 + "libc",
767 + "once_cell",
768 + "unicode-width",
769 + "windows-sys 0.61.2",
770 +]
771 +
772 +[[package]]
773 +name = "constant_time_eq"
774 +version = "0.4.2"
775 +source = "registry+https://github.com/rust-lang/crates.io-index"
776 +checksum = "3d52eff69cd5e647efe296129160853a42795992097e8af39800e1060caeea9b"
777 +
778 +[[package]]
779 +name = "cookie"
780 +version = "0.18.1"
781 +source = "registry+https://github.com/rust-lang/crates.io-index"
782 +checksum = "4ddef33a339a91ea89fb53151bd0a4689cfce27055c291dfa69945475d22c747"
783 +dependencies = [
784 + "percent-encoding",
785 + "time",
786 + "version_check",
787 +]
788 +
789 +[[package]]
790 +name = "cookie_store"
791 +version = "0.22.1"
792 +source = "registry+https://github.com/rust-lang/crates.io-index"
793 +checksum = "15b2c103cf610ec6cae3da84a766285b42fd16aad564758459e6ecf128c75206"
794 +dependencies = [
795 + "cookie",
796 + "document-features",
797 + "idna",
798 + "indexmap 2.13.0",
799 + "log",
800 + "serde",
801 + "serde_derive",
802 + "serde_json",
803 + "time",
804 + "url",
805 +]
806 +
807 +[[package]]
808 +name = "core-foundation"
809 +version = "0.9.4"
810 +source = "registry+https://github.com/rust-lang/crates.io-index"
811 +checksum = "91e195e091a93c46f7102ec7818a2aa394e1e1771c3ab4825963fa03e45afb8f"
812 +dependencies = [
813 + "core-foundation-sys",
814 + "libc",
815 +]
816 +
817 +[[package]]
818 +name = "core-foundation"
819 +version = "0.10.1"
820 +source = "registry+https://github.com/rust-lang/crates.io-index"
821 +checksum = "b2a6cd9ae233e7f62ba4e9353e81a88df7fc8a5987b8d445b4d90c879bd156f6"
822 +dependencies = [
823 + "core-foundation-sys",
824 + "libc",
825 +]
826 +
827 +[[package]]
828 +name = "core-foundation-sys"
829 +version = "0.8.7"
830 +source = "registry+https://github.com/rust-lang/crates.io-index"
831 +checksum = "773648b94d0e5d620f64f280777445740e61fe701025087ec8b57f45c791888b"
832 +
833 +[[package]]
834 +name = "cpufeatures"
835 +version = "0.2.17"
836 +source = "registry+https://github.com/rust-lang/crates.io-index"
837 +checksum = "59ed5838eebb26a2bb2e58f6d5b5316989ae9d08bab10e0e6d103e656d1b0280"
838 +dependencies = [
839 + "libc",
840 +]
841 +
842 +[[package]]
843 +name = "cpufeatures"
844 +version = "0.3.0"
845 +source = "registry+https://github.com/rust-lang/crates.io-index"
846 +checksum = "8b2a41393f66f16b0823bb79094d54ac5fbd34ab292ddafb9a0456ac9f87d201"
847 +dependencies = [
848 + "libc",
849 +]
850 +
851 +[[package]]
852 +name = "crc32fast"
853 +version = "1.5.0"
854 +source = "registry+https://github.com/rust-lang/crates.io-index"
855 +checksum = "9481c1c90cbf2ac953f07c8d4a58aa3945c425b7185c9154d67a65e4230da511"
856 +dependencies = [
857 + "cfg-if",
858 +]
859 +
860 +[[package]]
861 +name = "criterion"
862 +version = "0.5.1"
863 +source = "registry+https://github.com/rust-lang/crates.io-index"
864 +checksum = "f2b12d017a929603d80db1831cd3a24082f8137ce19c69e6447f54f5fc8d692f"
865 +dependencies = [
866 + "anes",
867 + "cast",
868 + "ciborium",
869 + "clap",
870 + "criterion-plot",
871 + "is-terminal",
872 + "itertools 0.10.5",
873 + "num-traits",
874 + "once_cell",
875 + "oorandom",
876 + "plotters",
877 + "rayon",
878 + "regex",
879 + "serde",
880 + "serde_derive",
881 + "serde_json",
882 + "tinytemplate",
883 + "walkdir",
884 +]
885 +
886 +[[package]]
887 +name = "criterion-plot"
888 +version = "0.5.0"
889 +source = "registry+https://github.com/rust-lang/crates.io-index"
890 +checksum = "6b50826342786a51a89e2da3a28f1c32b06e387201bc2d19791f622c673706b1"
891 +dependencies = [
892 + "cast",
893 + "itertools 0.10.5",
894 +]
895 +
896 +[[package]]
897 +name = "crossbeam-deque"
898 +version = "0.8.6"
899 +source = "registry+https://github.com/rust-lang/crates.io-index"
900 +checksum = "9dd111b7b7f7d55b72c0a6ae361660ee5853c9af73f70c3c2ef6858b950e2e51"
901 +dependencies = [
902 + "crossbeam-epoch",
903 + "crossbeam-utils",
904 +]
905 +
906 +[[package]]
907 +name = "crossbeam-epoch"
908 +version = "0.9.18"
909 +source = "registry+https://github.com/rust-lang/crates.io-index"
910 +checksum = "5b82ac4a3c2ca9c3460964f020e1402edd5753411d7737aa39c3714ad1b5420e"
911 +dependencies = [
912 + "crossbeam-utils",
913 +]
914 +
915 +[[package]]
916 +name = "crossbeam-queue"
917 +version = "0.3.12"
918 +source = "registry+https://github.com/rust-lang/crates.io-index"
919 +checksum = "0f58bbc28f91df819d0aa2a2c00cd19754769c2fad90579b3592b1c9ba7a3115"
920 +dependencies = [
921 + "crossbeam-utils",
922 +]
923 +
924 +[[package]]
925 +name = "crossbeam-utils"
926 +version = "0.8.21"
927 +source = "registry+https://github.com/rust-lang/crates.io-index"
928 +checksum = "d0a5c400df2834b80a4c3327b3aad3a4c4cd4de0629063962b03235697506a28"
929 +
930 +[[package]]
931 +name = "crunchy"
932 +version = "0.2.4"
933 +source = "registry+https://github.com/rust-lang/crates.io-index"
934 +checksum = "460fbee9c2c2f33933d720630a6a0bac33ba7053db5344fac858d4b8952d77d5"
935 +
936 +[[package]]
937 +name = "crypto-common"
938 +version = "0.1.7"
939 +source = "registry+https://github.com/rust-lang/crates.io-index"
940 +checksum = "78c8292055d1c1df0cce5d180393dc8cce0abec0a7102adb6c7b1eef6016d60a"
941 +dependencies = [
942 + "generic-array",
943 + "typenum",
944 +]
945 +
946 +[[package]]
947 +name = "daachorse"
948 +version = "1.0.0"
949 +source = "registry+https://github.com/rust-lang/crates.io-index"
950 +checksum = "63b7ef7a4be509357f4804d0a22e830daddb48f19fd604e4ad32ddce04a94c36"
951 +
952 +[[package]]
953 +name = "darling"
954 +version = "0.20.11"
955 +source = "registry+https://github.com/rust-lang/crates.io-index"
956 +checksum = "fc7f46116c46ff9ab3eb1597a45688b6715c6e628b5c133e288e709a29bcb4ee"
957 +dependencies = [
958 + "darling_core 0.20.11",
959 + "darling_macro 0.20.11",
960 +]
961 +
962 +[[package]]
963 +name = "darling"
964 +version = "0.23.0"
965 +source = "registry+https://github.com/rust-lang/crates.io-index"
966 +checksum = "25ae13da2f202d56bd7f91c25fba009e7717a1e4a1cc98a76d844b65ae912e9d"
967 +dependencies = [
968 + "darling_core 0.23.0",
969 + "darling_macro 0.23.0",
970 +]
971 +
972 +[[package]]
973 +name = "darling_core"
974 +version = "0.20.11"
975 +source = "registry+https://github.com/rust-lang/crates.io-index"
976 +checksum = "0d00b9596d185e565c2207a0b01f8bd1a135483d02d9b7b0a54b11da8d53412e"
977 +dependencies = [
978 + "fnv",
979 + "ident_case",
980 + "proc-macro2",
981 + "quote",
982 + "strsim",
983 + "syn 2.0.117",
984 +]
985 +
986 +[[package]]
987 +name = "darling_core"
988 +version = "0.23.0"
989 +source = "registry+https://github.com/rust-lang/crates.io-index"
990 +checksum = "9865a50f7c335f53564bb694ef660825eb8610e0a53d3e11bf1b0d3df31e03b0"
991 +dependencies = [
992 + "ident_case",
993 + "proc-macro2",
994 + "quote",
995 + "strsim",
996 + "syn 2.0.117",
997 +]
998 +
999 +[[package]]
1000 +name = "darling_macro"
1001 +version = "0.20.11"
1002 +source = "registry+https://github.com/rust-lang/crates.io-index"
1003 +checksum = "fc34b93ccb385b40dc71c6fceac4b2ad23662c7eeb248cf10d529b7e055b6ead"
1004 +dependencies = [
1005 + "darling_core 0.20.11",
1006 + "quote",
1007 + "syn 2.0.117",
1008 +]
1009 +
1010 +[[package]]
1011 +name = "darling_macro"
1012 +version = "0.23.0"
1013 +source = "registry+https://github.com/rust-lang/crates.io-index"
1014 +checksum = "ac3984ec7bd6cfa798e62b4a642426a5be0e68f9401cfc2a01e3fa9ea2fcdb8d"
1015 +dependencies = [
1016 + "darling_core 0.23.0",
1017 + "quote",
1018 + "syn 2.0.117",
1019 +]
1020 +
1021 +[[package]]
1022 +name = "dary_heap"
1023 +version = "0.3.8"
1024 +source = "registry+https://github.com/rust-lang/crates.io-index"
1025 +checksum = "06d2e3287df1c007e74221c49ca10a95d557349e54b3a75dc2fb14712c751f04"
1026 +dependencies = [
1027 + "serde",
1028 +]
1029 +
1030 +[[package]]
1031 +name = "der"
1032 +version = "0.8.0"
1033 +source = "registry+https://github.com/rust-lang/crates.io-index"
1034 +checksum = "71fd89660b2dc699704064e59e9dba0147b903e85319429e131620d022be411b"
1035 +dependencies = [
1036 + "pem-rfc7468",
1037 + "zeroize",
1038 +]
1039 +
1040 +[[package]]
1041 +name = "deranged"
1042 +version = "0.5.8"
1043 +source = "registry+https://github.com/rust-lang/crates.io-index"
1044 +checksum = "7cd812cc2bc1d69d4764bd80df88b4317eaef9e773c75226407d9bc0876b211c"
1045 +dependencies = [
1046 + "powerfmt",
1047 + "serde_core",
1048 +]
1049 +
1050 +[[package]]
1051 +name = "derive_builder"
1052 +version = "0.20.2"
1053 +source = "registry+https://github.com/rust-lang/crates.io-index"
1054 +checksum = "507dfb09ea8b7fa618fcf76e953f4f5e192547945816d5358edffe39f6f94947"
1055 +dependencies = [
1056 + "derive_builder_macro",
1057 +]
1058 +
1059 +[[package]]
1060 +name = "derive_builder_core"
1061 +version = "0.20.2"
1062 +source = "registry+https://github.com/rust-lang/crates.io-index"
1063 +checksum = "2d5bcf7b024d6835cfb3d473887cd966994907effbe9227e8c8219824d06c4e8"
1064 +dependencies = [
1065 + "darling 0.20.11",
1066 + "proc-macro2",
1067 + "quote",
1068 + "syn 2.0.117",
1069 +]
1070 +
1071 +[[package]]
1072 +name = "derive_builder_macro"
1073 +version = "0.20.2"
1074 +source = "registry+https://github.com/rust-lang/crates.io-index"
1075 +checksum = "ab63b0e2bf4d5928aff72e83a7dace85d7bba5fe12dcc3c5a572d78caffd3f3c"
1076 +dependencies = [
1077 + "derive_builder_core",
1078 + "syn 2.0.117",
1079 +]
1080 +
1081 +[[package]]
1082 +name = "derive_more"
1083 +version = "1.0.0"
1084 +source = "registry+https://github.com/rust-lang/crates.io-index"
1085 +checksum = "4a9b99b9cbbe49445b21764dc0625032a89b145a2642e67603e1c936f5458d05"
1086 +dependencies = [
1087 + "derive_more-impl",
1088 +]
1089 +
1090 +[[package]]
1091 +name = "derive_more-impl"
1092 +version = "1.0.0"
1093 +source = "registry+https://github.com/rust-lang/crates.io-index"
1094 +checksum = "cb7330aeadfbe296029522e6c40f315320aba36fc43a5b3632f3795348f3bd22"
1095 +dependencies = [
1096 + "proc-macro2",
1097 + "quote",
1098 + "syn 2.0.117",
1099 + "unicode-xid",
1100 +]
1101 +
1102 +[[package]]
1103 +name = "digest"
1104 +version = "0.10.7"
1105 +source = "registry+https://github.com/rust-lang/crates.io-index"
1106 +checksum = "9ed9a281f7bc9b7576e61468ba615a66a5c8cfdff42420a70aa82701a3b1e292"
1107 +dependencies = [
1108 + "block-buffer",
1109 + "crypto-common",
1110 + "subtle",
1111 +]
1112 +
1113 +[[package]]
1114 +name = "dirs"
1115 +version = "6.0.0"
1116 +source = "registry+https://github.com/rust-lang/crates.io-index"
1117 +checksum = "c3e8aa94d75141228480295a7d0e7feb620b1a5ad9f12bc40be62411e38cce4e"
1118 +dependencies = [
1119 + "dirs-sys",
1120 +]
1121 +
1122 +[[package]]
1123 +name = "dirs-sys"
1124 +version = "0.5.0"
1125 +source = "registry+https://github.com/rust-lang/crates.io-index"
1126 +checksum = "e01a3366d27ee9890022452ee61b2b63a67e6f13f58900b651ff5665f0bb1fab"
1127 +dependencies = [
1128 + "libc",
1129 + "option-ext",
1130 + "redox_users",
1131 + "windows-sys 0.61.2",
1132 +]
1133 +
1134 +[[package]]
1135 +name = "displaydoc"
1136 +version = "0.2.5"
1137 +source = "registry+https://github.com/rust-lang/crates.io-index"
1138 +checksum = "97369cbbc041bc366949bc74d34658d6cda5621039731c6310521892a3a20ae0"
1139 +dependencies = [
1140 + "proc-macro2",
1141 + "quote",
1142 + "syn 2.0.117",
1143 +]
1144 +
1145 +[[package]]
1146 +name = "dissimilar"
1147 +version = "1.0.11"
1148 +source = "registry+https://github.com/rust-lang/crates.io-index"
1149 +checksum = "aeda16ab4059c5fd2a83f2b9c9e9c981327b18aa8e3b313f7e6563799d4f093e"
1150 +
1151 +[[package]]
1152 +name = "document-features"
1153 +version = "0.2.12"
1154 +source = "registry+https://github.com/rust-lang/crates.io-index"
1155 +checksum = "d4b8a88685455ed29a21542a33abd9cb6510b6b129abadabdcef0f4c55bc8f61"
1156 +dependencies = [
1157 + "litrs",
1158 +]
1159 +
1160 +[[package]]
1161 +name = "dtoa"
1162 +version = "1.0.11"
1163 +source = "registry+https://github.com/rust-lang/crates.io-index"
1164 +checksum = "4c3cf4824e2d5f025c7b531afcb2325364084a16806f6d47fbc1f5fbd9960590"
1165 +
1166 +[[package]]
1167 +name = "dyn-clone"
1168 +version = "1.0.20"
1169 +source = "registry+https://github.com/rust-lang/crates.io-index"
1170 +checksum = "d0881ea181b1df73ff77ffaaf9c7544ecc11e82fba9b5f27b262a3c73a332555"
1171 +
1172 +[[package]]
1173 +name = "easy-ext"
1174 +version = "1.0.3"
1175 +source = "registry+https://github.com/rust-lang/crates.io-index"
1176 +checksum = "8072bec12b909b65aec01fa6518f387cfbf3427d4475409ad622898cd347522c"
1177 +
1178 +[[package]]
1179 +name = "educe"
1180 +version = "0.6.0"
1181 +source = "registry+https://github.com/rust-lang/crates.io-index"
1182 +checksum = "1d7bc049e1bd8cdeb31b68bbd586a9464ecf9f3944af3958a7a9d0f8b9799417"
1183 +dependencies = [
1184 + "enum-ordinalize",
1185 + "proc-macro2",
1186 + "quote",
1187 + "syn 2.0.117",
1188 +]
1189 +
1190 +[[package]]
1191 +name = "either"
1192 +version = "1.15.0"
1193 +source = "registry+https://github.com/rust-lang/crates.io-index"
1194 +checksum = "48c757948c5ede0e46177b7add2e67155f70e33c07fea8284df6576da70b3719"
1195 +
1196 +[[package]]
1197 +name = "encode_unicode"
1198 +version = "1.0.0"
1199 +source = "registry+https://github.com/rust-lang/crates.io-index"
1200 +checksum = "34aa73646ffb006b8f5147f3dc182bd4bcb190227ce861fc4a4844bf8e3cb2c0"
1201 +
1202 +[[package]]
1203 +name = "encoding_rs"
1204 +version = "0.8.35"
1205 +source = "registry+https://github.com/rust-lang/crates.io-index"
1206 +checksum = "75030f3c4f45dafd7586dd6780965a8c7e8e285a5ecb86713e63a79c5b2766f3"
1207 +dependencies = [
1208 + "cfg-if",
1209 +]
1210 +
1211 +[[package]]
1212 +name = "enum-as-inner"
1213 +version = "0.7.0"
1214 +source = "registry+https://github.com/rust-lang/crates.io-index"
1215 +checksum = "0359ee92f81184d7985519e474bda2a5738476334edd3746c9b1265c067afe70"
1216 +dependencies = [
1217 + "heck",
1218 + "proc-macro2",
1219 + "quote",
1220 + "syn 2.0.117",
1221 +]
1222 +
1223 +[[package]]
1224 +name = "enum-ordinalize"
1225 +version = "4.3.2"
1226 +source = "registry+https://github.com/rust-lang/crates.io-index"
1227 +checksum = "4a1091a7bb1f8f2c4b28f1fe2cef4980ca2d410a3d727d67ecc3178c9b0800f0"
1228 +dependencies = [
1229 + "enum-ordinalize-derive",
1230 +]
1231 +
1232 +[[package]]
1233 +name = "enum-ordinalize-derive"
1234 +version = "4.3.2"
1235 +source = "registry+https://github.com/rust-lang/crates.io-index"
1236 +checksum = "8ca9601fb2d62598ee17836250842873a413586e5d7ed88b356e38ddbb0ec631"
1237 +dependencies = [
1238 + "proc-macro2",
1239 + "quote",
1240 + "syn 2.0.117",
1241 +]
1242 +
1243 +[[package]]
1244 +name = "env_filter"
1245 +version = "1.0.1"
1246 +source = "registry+https://github.com/rust-lang/crates.io-index"
1247 +checksum = "32e90c2accc4b07a8456ea0debdc2e7587bdd890680d71173a15d4ae604f6eef"
1248 +dependencies = [
1249 + "log",
1250 + "regex",
1251 +]
1252 +
1253 +[[package]]
1254 +name = "env_logger"
1255 +version = "0.11.10"
1256 +source = "registry+https://github.com/rust-lang/crates.io-index"
1257 +checksum = "0621c04f2196ac3f488dd583365b9c09be011a4ab8b9f37248ffcc8f6198b56a"
1258 +dependencies = [
1259 + "anstream 1.0.0",
1260 + "anstyle",
1261 + "env_filter",
1262 + "jiff",
1263 + "log",
1264 +]
1265 +
1266 +[[package]]
1267 +name = "equator"
1268 +version = "0.4.2"
1269 +source = "registry+https://github.com/rust-lang/crates.io-index"
1270 +checksum = "4711b213838dfee0117e3be6ac926007d7f433d7bbe33595975d4190cb07e6fc"
1271 +dependencies = [
1272 + "equator-macro",
1273 +]
1274 +
1275 +[[package]]
1276 +name = "equator-macro"
1277 +version = "0.4.2"
1278 +source = "registry+https://github.com/rust-lang/crates.io-index"
1279 +checksum = "44f23cf4b44bfce11a86ace86f8a73ffdec849c9fd00a386a53d278bd9e81fb3"
1280 +dependencies = [
1281 + "proc-macro2",
1282 + "quote",
1283 + "syn 2.0.117",
1284 +]
1285 +
1286 +[[package]]
1287 +name = "equivalent"
1288 +version = "1.0.2"
1289 +source = "registry+https://github.com/rust-lang/crates.io-index"
1290 +checksum = "877a4ace8713b0bcf2a4e7eec82529c029f1d0619886d18145fea96c3ffe5c0f"
1291 +
1292 +[[package]]
1293 +name = "errno"
1294 +version = "0.3.14"
1295 +source = "registry+https://github.com/rust-lang/crates.io-index"
1296 +checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb"
1297 +dependencies = [
1298 + "libc",
1299 + "windows-sys 0.61.2",
1300 +]
1301 +
1302 +[[package]]
1303 +name = "esaxx-rs"
1304 +version = "0.1.10"
1305 +source = "registry+https://github.com/rust-lang/crates.io-index"
1306 +checksum = "d817e038c30374a4bcb22f94d0a8a0e216958d4c3dcde369b1439fec4bdda6e6"
1307 +dependencies = [
1308 + "cc",
1309 +]
1310 +
1311 +[[package]]
1312 +name = "eventsource-stream"
1313 +version = "0.2.3"
1314 +source = "registry+https://github.com/rust-lang/crates.io-index"
1315 +checksum = "74fef4569247a5f429d9156b9d0a2599914385dd189c539334c625d8099d90ab"
1316 +dependencies = [
1317 + "futures-core",
1318 + "nom 7.1.3",
1319 + "pin-project-lite",
1320 +]
1321 +
1322 +[[package]]
1323 +name = "expect-test"
1324 +version = "1.5.1"
1325 +source = "registry+https://github.com/rust-lang/crates.io-index"
1326 +checksum = "63af43ff4431e848fb47472a920f14fa71c24de13255a5692e93d4e90302acb0"
1327 +dependencies = [
1328 + "dissimilar",
1329 + "once_cell",
1330 +]
1331 +
1332 +[[package]]
1333 +name = "exr"
1334 +version = "1.74.0"
1335 +source = "registry+https://github.com/rust-lang/crates.io-index"
1336 +checksum = "4300e043a56aa2cb633c01af81ca8f699a321879a7854d3896a0ba89056363be"
1337 +dependencies = [
1338 + "bit_field",
1339 + "half",
1340 + "lebe",
1341 + "miniz_oxide",
1342 + "rayon-core",
1343 + "smallvec",
1344 + "zune-inflate",
1345 +]
1346 +
1347 +[[package]]
1348 +name = "fancy-regex"
1349 +version = "0.13.0"
1350 +source = "registry+https://github.com/rust-lang/crates.io-index"
1351 +checksum = "531e46835a22af56d1e3b66f04844bed63158bc094a628bec1d321d9b4c44bf2"
1352 +dependencies = [
1353 + "bit-set 0.5.3",
1354 + "regex-automata",
1355 + "regex-syntax",
1356 +]
1357 +
1358 +[[package]]
1359 +name = "fancy-regex"
1360 +version = "0.17.0"
1361 +source = "registry+https://github.com/rust-lang/crates.io-index"
1362 +checksum = "72cf461f865c862bb7dc573f643dd6a2b6842f7c30b07882b56bd148cc2761b8"
1363 +dependencies = [
1364 + "bit-set 0.8.0",
1365 + "regex-automata",
1366 + "regex-syntax",
1367 +]
1368 +
1369 +[[package]]
1370 +name = "fast_image_resize"
1371 +version = "6.0.0"
1372 +source = "registry+https://github.com/rust-lang/crates.io-index"
1373 +checksum = "12dd43e5011e8d8411a3215a0d57a2ec5c68282fb90eb5d7221fab0113442174"
1374 +dependencies = [
1375 + "bytemuck",
1376 + "cfg-if",
1377 + "document-features",
1378 + "image",
1379 + "num-traits",
1380 + "thiserror 2.0.18",
1381 +]
1382 +
1383 +[[package]]
1384 +name = "fastokens"
1385 +version = "0.2.0"
1386 +source = "registry+https://github.com/rust-lang/crates.io-index"
1387 +checksum = "796a262ed47d1458a4b40d0ed831c927e6f54d5b9c1de2683bb4ac9b04f4c7cc"
1388 +dependencies = [
1389 + "daachorse",
1390 + "fancy-regex 0.17.0",
1391 + "hf-hub 0.4.3",
1392 + "icu_normalizer",
1393 + "memchr",
1394 + "pcre2",
1395 + "rayon",
1396 + "serde",
1397 + "serde_json",
1398 + "strum",
1399 + "thiserror 2.0.18",
1400 +]
1401 +
1402 +[[package]]
1403 +name = "fastrand"
1404 +version = "2.3.0"
1405 +source = "registry+https://github.com/rust-lang/crates.io-index"
1406 +checksum = "37909eebbb50d72f9059c3b6d82c0463f2ff062c9e95845c43a6c9c0355411be"
1407 +
1408 +[[package]]
1409 +name = "fax"
1410 +version = "0.2.6"
1411 +source = "registry+https://github.com/rust-lang/crates.io-index"
1412 +checksum = "f05de7d48f37cd6730705cbca900770cab77a89f413d23e100ad7fad7795a0ab"
1413 +dependencies = [
1414 + "fax_derive",
1415 +]
1416 +
1417 +[[package]]
1418 +name = "fax_derive"
1419 +version = "0.2.0"
1420 +source = "registry+https://github.com/rust-lang/crates.io-index"
1421 +checksum = "a0aca10fb742cb43f9e7bb8467c91aa9bcb8e3ffbc6a6f7389bb93ffc920577d"
1422 +dependencies = [
1423 + "proc-macro2",
1424 + "quote",
1425 + "syn 2.0.117",
1426 +]
1427 +
1428 +[[package]]
1429 +name = "fdeflate"
1430 +version = "0.3.7"
1431 +source = "registry+https://github.com/rust-lang/crates.io-index"
1432 +checksum = "1e6853b52649d4ac5c0bd02320cddc5ba956bdb407c4b75a2c6b75bf51500f8c"
1433 +dependencies = [
1434 + "simd-adler32",
1435 +]
1436 +
1437 +[[package]]
1438 +name = "find-msvc-tools"
1439 +version = "0.1.9"
1440 +source = "registry+https://github.com/rust-lang/crates.io-index"
1441 +checksum = "5baebc0774151f905a1a2cc41989300b1e6fbb29aff0ceffa1064fdd3088d582"
1442 +
1443 +[[package]]
1444 +name = "fixedbitset"
1445 +version = "0.5.7"
1446 +source = "registry+https://github.com/rust-lang/crates.io-index"
1447 +checksum = "1d674e81391d1e1ab681a28d99df07927c6d4aa5b027d7da16ba32d1d21ecd99"
1448 +
1449 +[[package]]
1450 +name = "flate2"
1451 +version = "1.1.9"
1452 +source = "registry+https://github.com/rust-lang/crates.io-index"
1453 +checksum = "843fba2746e448b37e26a819579957415c8cef339bf08564fe8b7ddbd959573c"
1454 +dependencies = [
1455 + "crc32fast",
1456 + "miniz_oxide",
1457 +]
1458 +
1459 +[[package]]
1460 +name = "fnv"
1461 +version = "1.0.7"
1462 +source = "registry+https://github.com/rust-lang/crates.io-index"
1463 +checksum = "3f9eec918d3f24069decb9af1554cad7c880e2da24a9afd88aca000531ab82c1"
1464 +
1465 +[[package]]
1466 +name = "foldhash"
1467 +version = "0.1.5"
1468 +source = "registry+https://github.com/rust-lang/crates.io-index"
1469 +checksum = "d9c4f5dac5e15c24eb999c26181a6ca40b39fe946cbe4c263c7209467bc83af2"
1470 +
1471 +[[package]]
1472 +name = "foreign-types"
1473 +version = "0.3.2"
1474 +source = "registry+https://github.com/rust-lang/crates.io-index"
1475 +checksum = "f6f339eb8adc052cd2ca78910fda869aefa38d22d5cb648e6485e4d3fc06f3b1"
1476 +dependencies = [
1477 + "foreign-types-shared",
1478 +]
1479 +
1480 +[[package]]
1481 +name = "foreign-types-shared"
1482 +version = "0.1.1"
1483 +source = "registry+https://github.com/rust-lang/crates.io-index"
1484 +checksum = "00b0228411908ca8685dba7fc2cdd70ec9990a6e753e89b6ac91a84c40fbaf4b"
1485 +
1486 +[[package]]
1487 +name = "form_urlencoded"
1488 +version = "1.2.2"
1489 +source = "registry+https://github.com/rust-lang/crates.io-index"
1490 +checksum = "cb4cb245038516f5f85277875cdaa4f7d2c9a0fa0468de06ed190163b1581fcf"
1491 +dependencies = [
1492 + "percent-encoding",
1493 +]
1494 +
1495 +[[package]]
1496 +name = "fslock"
1497 +version = "0.2.1"
1498 +source = "registry+https://github.com/rust-lang/crates.io-index"
1499 +checksum = "04412b8935272e3a9bae6f48c7bfff74c2911f60525404edfdd28e49884c3bfb"
1500 +dependencies = [
1501 + "libc",
1502 + "winapi",
1503 +]
1504 +
1505 +[[package]]
1506 +name = "futures"
1507 +version = "0.3.32"
1508 +source = "registry+https://github.com/rust-lang/crates.io-index"
1509 +checksum = "8b147ee9d1f6d097cef9ce628cd2ee62288d963e16fb287bd9286455b241382d"
1510 +dependencies = [
1511 + "futures-channel",
1512 + "futures-core",
1513 + "futures-executor",
1514 + "futures-io",
1515 + "futures-sink",
1516 + "futures-task",
1517 + "futures-util",
1518 +]
1519 +
1520 +[[package]]
1521 +name = "futures-channel"
1522 +version = "0.3.32"
1523 +source = "registry+https://github.com/rust-lang/crates.io-index"
1524 +checksum = "07bbe89c50d7a535e539b8c17bc0b49bdb77747034daa8087407d655f3f7cc1d"
1525 +dependencies = [
1526 + "futures-core",
1527 + "futures-sink",
1528 +]
1529 +
1530 +[[package]]
1531 +name = "futures-core"
1532 +version = "0.3.32"
1533 +source = "registry+https://github.com/rust-lang/crates.io-index"
1534 +checksum = "7e3450815272ef58cec6d564423f6e755e25379b217b0bc688e295ba24df6b1d"
1535 +
1536 +[[package]]
1537 +name = "futures-executor"
1538 +version = "0.3.32"
1539 +source = "registry+https://github.com/rust-lang/crates.io-index"
1540 +checksum = "baf29c38818342a3b26b5b923639e7b1f4a61fc5e76102d4b1981c6dc7a7579d"
1541 +dependencies = [
1542 + "futures-core",
1543 + "futures-task",
1544 + "futures-util",
1545 +]
1546 +
1547 +[[package]]
1548 +name = "futures-io"
1549 +version = "0.3.32"
1550 +source = "registry+https://github.com/rust-lang/crates.io-index"
1551 +checksum = "cecba35d7ad927e23624b22ad55235f2239cfa44fd10428eecbeba6d6a717718"
1552 +
1553 +[[package]]
1554 +name = "futures-lite"
1555 +version = "2.6.1"
1556 +source = "registry+https://github.com/rust-lang/crates.io-index"
1557 +checksum = "f78e10609fe0e0b3f4157ffab1876319b5b0db102a2c60dc4626306dc46b44ad"
1558 +dependencies = [
1559 + "futures-core",
1560 + "pin-project-lite",
1561 +]
1562 +
1563 +[[package]]
1564 +name = "futures-macro"
1565 +version = "0.3.32"
1566 +source = "registry+https://github.com/rust-lang/crates.io-index"
1567 +checksum = "e835b70203e41293343137df5c0664546da5745f82ec9b84d40be8336958447b"
1568 +dependencies = [
1569 + "proc-macro2",
1570 + "quote",
1571 + "syn 2.0.117",
1572 +]
1573 +
1574 +[[package]]
1575 +name = "futures-sink"
1576 +version = "0.3.32"
1577 +source = "registry+https://github.com/rust-lang/crates.io-index"
1578 +checksum = "c39754e157331b013978ec91992bde1ac089843443c49cbc7f46150b0fad0893"
1579 +
1580 +[[package]]
1581 +name = "futures-task"
1582 +version = "0.3.32"
1583 +source = "registry+https://github.com/rust-lang/crates.io-index"
1584 +checksum = "037711b3d59c33004d3856fbdc83b99d4ff37a24768fa1be9ce3538a1cde4393"
1585 +
1586 +[[package]]
1587 +name = "futures-timer"
1588 +version = "3.0.3"
1589 +source = "registry+https://github.com/rust-lang/crates.io-index"
1590 +checksum = "f288b0a4f20f9a56b5d1da57e2227c661b7b16168e2f72365f57b63326e29b24"
1591 +
1592 +[[package]]
1593 +name = "futures-util"
1594 +version = "0.3.32"
1595 +source = "registry+https://github.com/rust-lang/crates.io-index"
1596 +checksum = "389ca41296e6190b48053de0321d02a77f32f8a5d2461dd38762c0593805c6d6"
1597 +dependencies = [
1598 + "futures-channel",
1599 + "futures-core",
1600 + "futures-io",
1601 + "futures-macro",
1602 + "futures-sink",
1603 + "futures-task",
1604 + "memchr",
1605 + "pin-project-lite",
1606 + "slab",
1607 +]
1608 +
1609 +[[package]]
1610 +name = "generic-array"
1611 +version = "0.14.7"
1612 +source = "registry+https://github.com/rust-lang/crates.io-index"
1613 +checksum = "85649ca51fd72272d7821adaf274ad91c288277713d9c18820d8499a7ff69e9a"
1614 +dependencies = [
1615 + "typenum",
1616 + "version_check",
1617 +]
1618 +
1619 +[[package]]
1620 +name = "getopts"
1621 +version = "0.2.24"
1622 +source = "registry+https://github.com/rust-lang/crates.io-index"
1623 +checksum = "cfe4fbac503b8d1f88e6676011885f34b7174f46e59956bba534ba83abded4df"
1624 +dependencies = [
1625 + "unicode-width",
1626 +]
1627 +
1628 +[[package]]
1629 +name = "getrandom"
1630 +version = "0.2.17"
1631 +source = "registry+https://github.com/rust-lang/crates.io-index"
1632 +checksum = "ff2abc00be7fca6ebc474524697ae276ad847ad0a6b3faa4bcb027e9a4614ad0"
1633 +dependencies = [
1634 + "cfg-if",
1635 + "js-sys",
1636 + "libc",
1637 + "wasi",
1638 + "wasm-bindgen",
1639 +]
1640 +
1641 +[[package]]
1642 +name = "getrandom"
1643 +version = "0.3.4"
1644 +source = "registry+https://github.com/rust-lang/crates.io-index"
1645 +checksum = "899def5c37c4fd7b2664648c28120ecec138e4d395b459e5ca34f9cce2dd77fd"
1646 +dependencies = [
1647 + "cfg-if",
1648 + "js-sys",
1649 + "libc",
1650 + "r-efi 5.3.0",
1651 + "wasip2",
1652 + "wasm-bindgen",
1653 +]
1654 +
1655 +[[package]]
1656 +name = "getrandom"
1657 +version = "0.4.2"
1658 +source = "registry+https://github.com/rust-lang/crates.io-index"
1659 +checksum = "0de51e6874e94e7bf76d726fc5d13ba782deca734ff60d5bb2fb2607c7406555"
1660 +dependencies = [
1661 + "cfg-if",
1662 + "libc",
1663 + "r-efi 6.0.0",
1664 + "wasip2",
1665 + "wasip3",
1666 +]
1667 +
1668 +[[package]]
1669 +name = "gif"
1670 +version = "0.14.2"
1671 +source = "registry+https://github.com/rust-lang/crates.io-index"
1672 +checksum = "ee8cfcc411d9adbbaba82fb72661cc1bcca13e8bba98b364e62b2dba8f960159"
1673 +dependencies = [
1674 + "color_quant",
1675 + "weezl",
1676 +]
1677 +
1678 +[[package]]
1679 +name = "h2"
1680 +version = "0.4.13"
1681 +source = "registry+https://github.com/rust-lang/crates.io-index"
1682 +checksum = "2f44da3a8150a6703ed5d34e164b875fd14c2cdab9af1252a9a1020bde2bdc54"
1683 +dependencies = [
1684 + "atomic-waker",
1685 + "bytes",
1686 + "fnv",
1687 + "futures-core",
1688 + "futures-sink",
1689 + "http",
1690 + "indexmap 2.13.0",
1691 + "slab",
1692 + "tokio",
1693 + "tokio-util",
1694 + "tracing",
1695 +]
1696 +
1697 +[[package]]
1698 +name = "half"
1699 +version = "2.7.1"
1700 +source = "registry+https://github.com/rust-lang/crates.io-index"
1701 +checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b"
1702 +dependencies = [
1703 + "bytemuck",
1704 + "cfg-if",
1705 + "crunchy",
1706 + "zerocopy",
1707 +]
1708 +
1709 +[[package]]
1710 +name = "hashbrown"
1711 +version = "0.12.3"
1712 +source = "registry+https://github.com/rust-lang/crates.io-index"
1713 +checksum = "8a9ee70c43aaf417c914396645a0fa852624801b24ebb7ae78fe8272889ac888"
1714 +
1715 +[[package]]
1716 +name = "hashbrown"
1717 +version = "0.14.5"
1718 +source = "registry+https://github.com/rust-lang/crates.io-index"
1719 +checksum = "e5274423e17b7c9fc20b6e7e208532f9b19825d82dfd615708b70edd83df41f1"
1720 +dependencies = [
1721 + "ahash",
1722 +]
1723 +
1724 +[[package]]
1725 +name = "hashbrown"
1726 +version = "0.15.5"
1727 +source = "registry+https://github.com/rust-lang/crates.io-index"
1728 +checksum = "9229cfe53dfd69f0609a49f65461bd93001ea1ef889cd5529dd176593f5338a1"
1729 +dependencies = [
1730 + "foldhash",
1731 +]
1732 +
1733 +[[package]]
1734 +name = "hashbrown"
1735 +version = "0.16.1"
1736 +source = "registry+https://github.com/rust-lang/crates.io-index"
1737 +checksum = "841d1cc9bed7f9236f321df977030373f4a4163ae1a7dbfe1a51a2c1a51d9100"
1738 +
1739 +[[package]]
1740 +name = "heck"
1741 +version = "0.5.0"
1742 +source = "registry+https://github.com/rust-lang/crates.io-index"
1743 +checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea"
1744 +
1745 +[[package]]
1746 +name = "hermit-abi"
1747 +version = "0.5.2"
1748 +source = "registry+https://github.com/rust-lang/crates.io-index"
1749 +checksum = "fc0fef456e4baa96da950455cd02c081ca953b141298e41db3fc7e36b1da849c"
1750 +
1751 +[[package]]
1752 +name = "hex"
1753 +version = "0.4.3"
1754 +source = "registry+https://github.com/rust-lang/crates.io-index"
1755 +checksum = "7f24254aa9a54b5c858eaee2f5bccdb46aaf0e486a595ed5fd8f86ba55232a70"
1756 +
1757 +[[package]]
1758 +name = "hf-hub"
1759 +version = "0.4.3"
1760 +source = "registry+https://github.com/rust-lang/crates.io-index"
1761 +checksum = "629d8f3bbeda9d148036d6b0de0a3ab947abd08ce90626327fc3547a49d59d97"
1762 +dependencies = [
1763 + "dirs",
1764 + "http",
1765 + "indicatif 0.17.11",
1766 + "libc",
1767 + "log",
1768 + "rand 0.9.2",
1769 + "reqwest",
1770 + "serde",
1771 + "serde_json",
1772 + "thiserror 2.0.18",
1773 + "ureq 2.12.1",
1774 + "windows-sys 0.60.2",
1775 +]
1776 +
1777 +[[package]]
1778 +name = "hf-hub"
1779 +version = "0.5.0"
1780 +source = "registry+https://github.com/rust-lang/crates.io-index"
1781 +checksum = "aef3982638978efa195ff11b305f51f1f22f4f0a6cabee7af79b383ebee6a213"
1782 +dependencies = [
1783 + "dirs",
1784 + "futures",
1785 + "http",
1786 + "indicatif 0.18.4",
1787 + "libc",
1788 + "log",
1789 + "native-tls",
1790 + "num_cpus",
1791 + "rand 0.9.2",
1792 + "reqwest",
1793 + "serde",
1794 + "serde_json",
1795 + "thiserror 2.0.18",
1796 + "tokio",
1797 + "ureq 3.3.0",
1798 + "windows-sys 0.61.2",
1799 +]
1800 +
1801 +[[package]]
1802 +name = "hmac"
1803 +version = "0.12.1"
1804 +source = "registry+https://github.com/rust-lang/crates.io-index"
1805 +checksum = "6c49c37c09c17a53d937dfbb742eb3a961d65a994e6bcdcf37e7399d0cc8ab5e"
1806 +dependencies = [
1807 + "digest",
1808 +]
1809 +
1810 +[[package]]
1811 +name = "hound"
1812 +version = "3.5.1"
1813 +source = "registry+https://github.com/rust-lang/crates.io-index"
1814 +checksum = "62adaabb884c94955b19907d60019f4e145d091c75345379e70d1ee696f7854f"
1815 +
1816 +[[package]]
1817 +name = "http"
1818 +version = "1.4.0"
1819 +source = "registry+https://github.com/rust-lang/crates.io-index"
1820 +checksum = "e3ba2a386d7f85a81f119ad7498ebe444d2e22c2af0b86b069416ace48b3311a"
1821 +dependencies = [
1822 + "bytes",
1823 + "itoa",
1824 +]
1825 +
1826 +[[package]]
1827 +name = "http-body"
1828 +version = "1.0.1"
1829 +source = "registry+https://github.com/rust-lang/crates.io-index"
1830 +checksum = "1efedce1fb8e6913f23e0c92de8e62cd5b772a67e7b3946df930a62566c93184"
1831 +dependencies = [
1832 + "bytes",
1833 + "http",
1834 +]
1835 +
1836 +[[package]]
1837 +name = "http-body-util"
1838 +version = "0.1.3"
1839 +source = "registry+https://github.com/rust-lang/crates.io-index"
1840 +checksum = "b021d93e26becf5dc7e1b75b1bed1fd93124b374ceb73f43d4d4eafec896a64a"
1841 +dependencies = [
1842 + "bytes",
1843 + "futures-core",
1844 + "http",
1845 + "http-body",
1846 + "pin-project-lite",
1847 +]
1848 +
1849 +[[package]]
1850 +name = "httparse"
1851 +version = "1.10.1"
1852 +source = "registry+https://github.com/rust-lang/crates.io-index"
1853 +checksum = "6dbf3de79e51f3d586ab4cb9d5c3e2c14aa28ed23d180cf89b4df0454a69cc87"
1854 +
1855 +[[package]]
1856 +name = "httpdate"
1857 +version = "1.0.3"
1858 +source = "registry+https://github.com/rust-lang/crates.io-index"
1859 +checksum = "df3b46402a9d5adb4c86a0cf463f42e19994e3ee891101b1841f30a545cb49a9"
1860 +
1861 +[[package]]
1862 +name = "hyper"
1863 +version = "1.8.1"
1864 +source = "registry+https://github.com/rust-lang/crates.io-index"
1865 +checksum = "2ab2d4f250c3d7b1c9fcdff1cece94ea4e2dfbec68614f7b87cb205f24ca9d11"
1866 +dependencies = [
1867 + "atomic-waker",
1868 + "bytes",
1869 + "futures-channel",
1870 + "futures-core",
1871 + "h2",
1872 + "http",
1873 + "http-body",
1874 + "httparse",
1875 + "httpdate",
1876 + "itoa",
1877 + "pin-project-lite",
1878 + "pin-utils",
1879 + "smallvec",
1880 + "tokio",
1881 + "want",
1882 +]
1883 +
1884 +[[package]]
1885 +name = "hyper-rustls"
1886 +version = "0.27.7"
1887 +source = "registry+https://github.com/rust-lang/crates.io-index"
1888 +checksum = "e3c93eb611681b207e1fe55d5a71ecf91572ec8a6705cdb6857f7d8d5242cf58"
1889 +dependencies = [
1890 + "http",
1891 + "hyper",
1892 + "hyper-util",
1893 + "rustls",
1894 + "rustls-native-certs",
1895 + "rustls-pki-types",
1896 + "tokio",
1897 + "tokio-rustls",
1898 + "tower-service",
1899 + "webpki-roots 1.0.6",
1900 +]
1901 +
1902 +[[package]]
1903 +name = "hyper-timeout"
1904 +version = "0.5.2"
1905 +source = "registry+https://github.com/rust-lang/crates.io-index"
1906 +checksum = "2b90d566bffbce6a75bd8b09a05aa8c2cb1fabb6cb348f8840c9e4c90a0d83b0"
1907 +dependencies = [
1908 + "hyper",
1909 + "hyper-util",
1910 + "pin-project-lite",
1911 + "tokio",
1912 + "tower-service",
1913 +]
1914 +
1915 +[[package]]
1916 +name = "hyper-tls"
1917 +version = "0.6.0"
1918 +source = "registry+https://github.com/rust-lang/crates.io-index"
1919 +checksum = "70206fc6890eaca9fde8a0bf71caa2ddfc9fe045ac9e5c70df101a7dbde866e0"
1920 +dependencies = [
1921 + "bytes",
1922 + "http-body-util",
1923 + "hyper",
1924 + "hyper-util",
1925 + "native-tls",
1926 + "tokio",
1927 + "tokio-native-tls",
1928 + "tower-service",
1929 +]
1930 +
1931 +[[package]]
1932 +name = "hyper-util"
1933 +version = "0.1.20"
1934 +source = "registry+https://github.com/rust-lang/crates.io-index"
1935 +checksum = "96547c2556ec9d12fb1578c4eaf448b04993e7fb79cbaad930a656880a6bdfa0"
1936 +dependencies = [
1937 + "base64 0.22.1",
1938 + "bytes",
1939 + "futures-channel",
1940 + "futures-util",
1941 + "http",
1942 + "http-body",
1943 + "hyper",
1944 + "ipnet",
1945 + "libc",
1946 + "percent-encoding",
1947 + "pin-project-lite",
1948 + "socket2",
1949 + "system-configuration",
1950 + "tokio",
1951 + "tower-service",
1952 + "tracing",
1953 + "windows-registry",
1954 +]
1955 +
1956 +[[package]]
1957 +name = "iana-time-zone"
1958 +version = "0.1.65"
1959 +source = "registry+https://github.com/rust-lang/crates.io-index"
1960 +checksum = "e31bc9ad994ba00e440a8aa5c9ef0ec67d5cb5e5cb0cc7f8b744a35b389cc470"
1961 +dependencies = [
1962 + "android_system_properties",
1963 + "core-foundation-sys",
1964 + "iana-time-zone-haiku",
1965 + "js-sys",
1966 + "log",
1967 + "wasm-bindgen",
1968 + "windows-core",
1969 +]
1970 +
1971 +[[package]]
1972 +name = "iana-time-zone-haiku"
1973 +version = "0.1.2"
1974 +source = "registry+https://github.com/rust-lang/crates.io-index"
1975 +checksum = "f31827a206f56af32e590ba56d5d2d085f558508192593743f16b2306495269f"
1976 +dependencies = [
1977 + "cc",
1978 +]
1979 +
1980 +[[package]]
1981 +name = "icu_collections"
1982 +version = "2.1.1"
1983 +source = "registry+https://github.com/rust-lang/crates.io-index"
1984 +checksum = "4c6b649701667bbe825c3b7e6388cb521c23d88644678e83c0c4d0a621a34b43"
1985 +dependencies = [
1986 + "displaydoc",
1987 + "potential_utf",
1988 + "yoke",
1989 + "zerofrom",
1990 + "zerovec",
1991 +]
1992 +
1993 +[[package]]
1994 +name = "icu_locale_core"
1995 +version = "2.1.1"
1996 +source = "registry+https://github.com/rust-lang/crates.io-index"
1997 +checksum = "edba7861004dd3714265b4db54a3c390e880ab658fec5f7db895fae2046b5bb6"
1998 +dependencies = [
1999 + "displaydoc",
2000 + "litemap",
2001 + "tinystr",
2002 + "writeable",
2003 + "zerovec",
2004 +]
2005 +
2006 +[[package]]
2007 +name = "icu_normalizer"
2008 +version = "2.1.1"
2009 +source = "registry+https://github.com/rust-lang/crates.io-index"
2010 +checksum = "5f6c8828b67bf8908d82127b2054ea1b4427ff0230ee9141c54251934ab1b599"
2011 +dependencies = [
2012 + "icu_collections",
2013 + "icu_normalizer_data",
2014 + "icu_properties",
2015 + "icu_provider",
2016 + "smallvec",
2017 + "utf16_iter",
2018 + "utf8_iter",
2019 + "write16",
2020 + "zerovec",
2021 +]
2022 +
2023 +[[package]]
2024 +name = "icu_normalizer_data"
2025 +version = "2.1.1"
2026 +source = "registry+https://github.com/rust-lang/crates.io-index"
2027 +checksum = "7aedcccd01fc5fe81e6b489c15b247b8b0690feb23304303a9e560f37efc560a"
2028 +
2029 +[[package]]
2030 +name = "icu_properties"
2031 +version = "2.1.2"
2032 +source = "registry+https://github.com/rust-lang/crates.io-index"
2033 +checksum = "020bfc02fe870ec3a66d93e677ccca0562506e5872c650f893269e08615d74ec"
2034 +dependencies = [
2035 + "icu_collections",
2036 + "icu_locale_core",
2037 + "icu_properties_data",
2038 + "icu_provider",
2039 + "zerotrie",
2040 + "zerovec",
2041 +]
2042 +
2043 +[[package]]
2044 +name = "icu_properties_data"
2045 +version = "2.1.2"
2046 +source = "registry+https://github.com/rust-lang/crates.io-index"
2047 +checksum = "616c294cf8d725c6afcd8f55abc17c56464ef6211f9ed59cccffe534129c77af"
2048 +
2049 +[[package]]
2050 +name = "icu_provider"
2051 +version = "2.1.1"
2052 +source = "registry+https://github.com/rust-lang/crates.io-index"
2053 +checksum = "85962cf0ce02e1e0a629cc34e7ca3e373ce20dda4c4d7294bbd0bf1fdb59e614"
2054 +dependencies = [
2055 + "displaydoc",
2056 + "icu_locale_core",
2057 + "writeable",
2058 + "yoke",
2059 + "zerofrom",
2060 + "zerotrie",
2061 + "zerovec",
2062 +]
2063 +
2064 +[[package]]
2065 +name = "id-arena"
2066 +version = "2.3.0"
2067 +source = "registry+https://github.com/rust-lang/crates.io-index"
2068 +checksum = "3d3067d79b975e8844ca9eb072e16b31c3c1c36928edf9c6789548c524d0d954"
2069 +
2070 +[[package]]
2071 +name = "ident_case"
2072 +version = "1.0.1"
2073 +source = "registry+https://github.com/rust-lang/crates.io-index"
2074 +checksum = "b9e0384b61958566e926dc50660321d12159025e767c18e043daf26b70104c39"
2075 +
2076 +[[package]]
2077 +name = "idna"
2078 +version = "1.1.0"
2079 +source = "registry+https://github.com/rust-lang/crates.io-index"
2080 +checksum = "3b0875f23caa03898994f6ddc501886a45c7d3d62d04d2d90788d47be1b1e4de"
2081 +dependencies = [
2082 + "idna_adapter",
2083 + "smallvec",
2084 + "utf8_iter",
2085 +]
2086 +
2087 +[[package]]
2088 +name = "idna_adapter"
2089 +version = "1.2.1"
2090 +source = "registry+https://github.com/rust-lang/crates.io-index"
2091 +checksum = "3acae9609540aa318d1bc588455225fb2085b9ed0c4f6bd0d9d5bcd86f1a0344"
2092 +dependencies = [
2093 + "icu_normalizer",
2094 + "icu_properties",
2095 +]
2096 +
2097 +[[package]]
2098 +name = "image"
2099 +version = "0.25.10"
2100 +source = "registry+https://github.com/rust-lang/crates.io-index"
2101 +checksum = "85ab80394333c02fe689eaf900ab500fbd0c2213da414687ebf995a65d5a6104"
2102 +dependencies = [
2103 + "bytemuck",
2104 + "byteorder-lite",
2105 + "color_quant",
2106 + "exr",
2107 + "gif",
2108 + "image-webp",
2109 + "moxcms",
2110 + "num-traits",
2111 + "png",
2112 + "qoi",
2113 + "ravif",
2114 + "rayon",
2115 + "rgb",
2116 + "tiff",
2117 + "zune-core",
2118 + "zune-jpeg",
2119 +]
2120 +
2121 +[[package]]
2122 +name = "image-webp"
2123 +version = "0.2.4"
2124 +source = "registry+https://github.com/rust-lang/crates.io-index"
2125 +checksum = "525e9ff3e1a4be2fbea1fdf0e98686a6d98b4d8f937e1bf7402245af1909e8c3"
2126 +dependencies = [
2127 + "byteorder-lite",
2128 + "quick-error",
2129 +]
2130 +
2131 +[[package]]
2132 +name = "imgref"
2133 +version = "1.12.0"
2134 +source = "registry+https://github.com/rust-lang/crates.io-index"
2135 +checksum = "e7c5cedc30da3a610cac6b4ba17597bdf7152cf974e8aab3afb3d54455e371c8"
2136 +
2137 +[[package]]
2138 +name = "indexmap"
2139 +version = "1.9.3"
2140 +source = "registry+https://github.com/rust-lang/crates.io-index"
2141 +checksum = "bd070e393353796e801d209ad339e89596eb4c8d430d18ede6a1cced8fafbd99"
2142 +dependencies = [
2143 + "autocfg",
2144 + "hashbrown 0.12.3",
2145 + "serde",
2146 +]
2147 +
2148 +[[package]]
2149 +name = "indexmap"
2150 +version = "2.13.0"
2151 +source = "registry+https://github.com/rust-lang/crates.io-index"
2152 +checksum = "7714e70437a7dc3ac8eb7e6f8df75fd8eb422675fc7678aff7364301092b1017"
2153 +dependencies = [
2154 + "equivalent",
2155 + "hashbrown 0.16.1",
2156 + "serde",
2157 + "serde_core",
2158 +]
2159 +
2160 +[[package]]
2161 +name = "indicatif"
2162 +version = "0.17.11"
2163 +source = "registry+https://github.com/rust-lang/crates.io-index"
2164 +checksum = "183b3088984b400f4cfac3620d5e076c84da5364016b4f49473de574b2586235"
2165 +dependencies = [
2166 + "console 0.15.11",
2167 + "number_prefix",
2168 + "portable-atomic",
2169 + "unicode-width",
2170 + "web-time",
2171 +]
2172 +
2173 +[[package]]
2174 +name = "indicatif"
2175 +version = "0.18.4"
2176 +source = "registry+https://github.com/rust-lang/crates.io-index"
2177 +checksum = "25470f23803092da7d239834776d653104d551bc4d7eacaf31e6837854b8e9eb"
2178 +dependencies = [
2179 + "console 0.16.2",
2180 + "portable-atomic",
2181 + "unicode-width",
2182 + "unit-prefix",
2183 + "web-time",
2184 +]
2185 +
2186 +[[package]]
2187 +name = "instant"
2188 +version = "0.1.13"
2189 +source = "registry+https://github.com/rust-lang/crates.io-index"
2190 +checksum = "e0242819d153cba4b4b05a5a8f2a7e9bbf97b6055b2a002b395c96b5ff3c0222"
2191 +dependencies = [
2192 + "cfg-if",
2193 +]
2194 +
2195 +[[package]]
2196 +name = "interpolate_name"
2197 +version = "0.2.4"
2198 +source = "registry+https://github.com/rust-lang/crates.io-index"
2199 +checksum = "c34819042dc3d3971c46c2190835914dfbe0c3c13f61449b2997f4e9722dfa60"
2200 +dependencies = [
2201 + "proc-macro2",
2202 + "quote",
2203 + "syn 2.0.117",
2204 +]
2205 +
2206 +[[package]]
2207 +name = "ipnet"
2208 +version = "2.12.0"
2209 +source = "registry+https://github.com/rust-lang/crates.io-index"
2210 +checksum = "d98f6fed1fde3f8c21bc40a1abb88dd75e67924f9cffc3ef95607bad8017f8e2"
2211 +
2212 +[[package]]
2213 +name = "iri-string"
2214 +version = "0.7.10"
2215 +source = "registry+https://github.com/rust-lang/crates.io-index"
2216 +checksum = "c91338f0783edbd6195decb37bae672fd3b165faffb89bf7b9e6942f8b1a731a"
2217 +dependencies = [
2218 + "memchr",
2219 + "serde",
2220 +]
2221 +
2222 +[[package]]
2223 +name = "is-macro"
2224 +version = "0.3.7"
2225 +source = "registry+https://github.com/rust-lang/crates.io-index"
2226 +checksum = "1d57a3e447e24c22647738e4607f1df1e0ec6f72e16182c4cd199f647cdfb0e4"
2227 +dependencies = [
2228 + "heck",
2229 + "proc-macro2",
2230 + "quote",
2231 + "syn 2.0.117",
2232 +]
2233 +
2234 +[[package]]
2235 +name = "is-terminal"
2236 +version = "0.4.17"
2237 +source = "registry+https://github.com/rust-lang/crates.io-index"
2238 +checksum = "3640c1c38b8e4e43584d8df18be5fc6b0aa314ce6ebf51b53313d4306cca8e46"
2239 +dependencies = [
2240 + "hermit-abi",
2241 + "libc",
2242 + "windows-sys 0.61.2",
2243 +]
2244 +
2245 +[[package]]
2246 +name = "is_terminal_polyfill"
2247 +version = "1.70.2"
2248 +source = "registry+https://github.com/rust-lang/crates.io-index"
2249 +checksum = "a6cb138bb79a146c1bd460005623e142ef0181e3d0219cb493e02f7d08a35695"
2250 +
2251 +[[package]]
2252 +name = "itertools"
2253 +version = "0.10.5"
2254 +source = "registry+https://github.com/rust-lang/crates.io-index"
2255 +checksum = "b0fd2260e829bddf4cb6ea802289de2f86d6a7a690192fbe91b3f46e0f2c8473"
2256 +dependencies = [
2257 + "either",
2258 +]
2259 +
2260 +[[package]]
2261 +name = "itertools"
2262 +version = "0.11.0"
2263 +source = "registry+https://github.com/rust-lang/crates.io-index"
2264 +checksum = "b1c173a5686ce8bfa551b3563d0c2170bf24ca44da99c7ca4bfdab5418c3fe57"
2265 +dependencies = [
2266 + "either",
2267 +]
2268 +
2269 +[[package]]
2270 +name = "itertools"
2271 +version = "0.14.0"
2272 +source = "registry+https://github.com/rust-lang/crates.io-index"
2273 +checksum = "2b192c782037fadd9cfa75548310488aabdbf3d2da73885b31bd0abd03351285"
2274 +dependencies = [
2275 + "either",
2276 +]
2277 +
2278 +[[package]]
2279 +name = "itoa"
2280 +version = "1.0.17"
2281 +source = "registry+https://github.com/rust-lang/crates.io-index"
2282 +checksum = "92ecc6618181def0457392ccd0ee51198e065e016d1d527a7ac1b6dc7c1f09d2"
2283 +
2284 +[[package]]
2285 +name = "jiff"
2286 +version = "0.2.23"
2287 +source = "registry+https://github.com/rust-lang/crates.io-index"
2288 +checksum = "1a3546dc96b6d42c5f24902af9e2538e82e39ad350b0c766eb3fbf2d8f3d8359"
2289 +dependencies = [
2290 + "jiff-static",
2291 + "log",
2292 + "portable-atomic",
2293 + "portable-atomic-util",
2294 + "serde_core",
2295 +]
2296 +
2297 +[[package]]
2298 +name = "jiff-static"
2299 +version = "0.2.23"
2300 +source = "registry+https://github.com/rust-lang/crates.io-index"
2301 +checksum = "2a8c8b344124222efd714b73bb41f8b5120b27a7cc1c75593a6ff768d9d05aa4"
2302 +dependencies = [
2303 + "proc-macro2",
2304 + "quote",
2305 + "syn 2.0.117",
2306 +]
2307 +
2308 +[[package]]
2309 +name = "jobserver"
2310 +version = "0.1.34"
2311 +source = "registry+https://github.com/rust-lang/crates.io-index"
2312 +checksum = "9afb3de4395d6b3e67a780b6de64b51c978ecf11cb9a462c66be7d4ca9039d33"
2313 +dependencies = [
2314 + "getrandom 0.3.4",
2315 + "libc",
2316 +]
2317 +
2318 +[[package]]
2319 +name = "js-sys"
2320 +version = "0.3.91"
2321 +source = "registry+https://github.com/rust-lang/crates.io-index"
2322 +checksum = "b49715b7073f385ba4bc528e5747d02e66cb39c6146efb66b781f131f0fb399c"
2323 +dependencies = [
2324 + "once_cell",
2325 + "wasm-bindgen",
2326 +]
2327 +
2328 +[[package]]
2329 +name = "lalrpop-util"
2330 +version = "0.20.2"
2331 +source = "registry+https://github.com/rust-lang/crates.io-index"
2332 +checksum = "507460a910eb7b32ee961886ff48539633b788a36b65692b95f225b844c82553"
2333 +
2334 +[[package]]
2335 +name = "lazy_static"
2336 +version = "1.5.0"
2337 +source = "registry+https://github.com/rust-lang/crates.io-index"
2338 +checksum = "bbd2bcb4c963f2ddae06a2efc7e9f3591312473c50c6685e1f298068316e66fe"
2339 +
2340 +[[package]]
2341 +name = "leb128fmt"
2342 +version = "0.1.0"
2343 +source = "registry+https://github.com/rust-lang/crates.io-index"
2344 +checksum = "09edd9e8b54e49e587e4f6295a7d29c3ea94d469cb40ab8ca70b288248a81db2"
2345 +
2346 +[[package]]
2347 +name = "lebe"
2348 +version = "0.5.3"
2349 +source = "registry+https://github.com/rust-lang/crates.io-index"
2350 +checksum = "7a79a3332a6609480d7d0c9eab957bca6b455b91bb84e66d19f5ff66294b85b8"
2351 +
2352 +[[package]]
2353 +name = "libc"
2354 +version = "0.2.183"
2355 +source = "registry+https://github.com/rust-lang/crates.io-index"
2356 +checksum = "b5b646652bf6661599e1da8901b3b9522896f01e736bad5f723fe7a3a27f899d"
2357 +
2358 +[[package]]
2359 +name = "libfuzzer-sys"
2360 +version = "0.4.12"
2361 +source = "registry+https://github.com/rust-lang/crates.io-index"
2362 +checksum = "f12a681b7dd8ce12bff52488013ba614b869148d54dd79836ab85aafdd53f08d"
2363 +dependencies = [
2364 + "arbitrary",
2365 + "cc",
2366 +]
2367 +
2368 +[[package]]
2369 +name = "libm"
2370 +version = "0.2.16"
2371 +source = "registry+https://github.com/rust-lang/crates.io-index"
2372 +checksum = "b6d2cec3eae94f9f509c767b45932f1ada8350c4bdb85af2fcab4a3c14807981"
2373 +
2374 +[[package]]
2375 +name = "libredox"
2376 +version = "0.1.14"
2377 +source = "registry+https://github.com/rust-lang/crates.io-index"
2378 +checksum = "1744e39d1d6a9948f4f388969627434e31128196de472883b39f148769bfe30a"
2379 +dependencies = [
2380 + "libc",
2381 +]
2382 +
2383 +[[package]]
2384 +name = "linux-raw-sys"
2385 +version = "0.12.1"
2386 +source = "registry+https://github.com/rust-lang/crates.io-index"
2387 +checksum = "32a66949e030da00e8c7d4434b251670a91556f4144941d37452769c25d58a53"
2388 +
2389 +[[package]]
2390 +name = "litemap"
2391 +version = "0.8.1"
2392 +source = "registry+https://github.com/rust-lang/crates.io-index"
2393 +checksum = "6373607a59f0be73a39b6fe456b8192fcc3585f602af20751600e974dd455e77"
2394 +
2395 +[[package]]
2396 +name = "litrs"
2397 +version = "1.0.0"
2398 +source = "registry+https://github.com/rust-lang/crates.io-index"
2399 +checksum = "11d3d7f243d5c5a8b9bb5d6dd2b1602c0cb0b9db1621bafc7ed66e35ff9fe092"
2400 +
2401 +[[package]]
2402 +name = "llm-multimodal"
2403 +version = "1.5.0"
2404 +source = "git+https://github.com/vllm-project/llm-multimodal?rev=5b558989844d1c7af3e43d0f604069ffd9c06320#5b558989844d1c7af3e43d0f604069ffd9c06320"
2405 +dependencies = [
2406 + "base64 0.22.1",
2407 + "blake3",
2408 + "bytes",
2409 + "fast_image_resize",
2410 + "image",
2411 + "ndarray 0.17.2",
2412 + "once_cell",
2413 + "reqwest",
2414 + "serde",
2415 + "serde_json",
2416 + "thiserror 2.0.18",
2417 + "tokio",
2418 + "url",
2419 +]
2420 +
2421 +[[package]]
2422 +name = "lock_api"
2423 +version = "0.4.14"
2424 +source = "registry+https://github.com/rust-lang/crates.io-index"
2425 +checksum = "224399e74b87b5f3557511d98dff8b14089b3dadafcab6bb93eab67d3aace965"
2426 +dependencies = [
2427 + "scopeguard",
2428 +]
2429 +
2430 +[[package]]
2431 +name = "log"
2432 +version = "0.4.29"
2433 +source = "registry+https://github.com/rust-lang/crates.io-index"
2434 +checksum = "5e5032e24019045c762d3c0f28f5b6b8bbf38563a65908389bf7978758920897"
2435 +
2436 +[[package]]
2437 +name = "loop9"
2438 +version = "0.1.5"
2439 +source = "registry+https://github.com/rust-lang/crates.io-index"
2440 +checksum = "0fae87c125b03c1d2c0150c90365d7d6bcc53fb73a9acaef207d2d065860f062"
2441 +dependencies = [
2442 + "imgref",
2443 +]
2444 +
2445 +[[package]]
2446 +name = "lru-slab"
2447 +version = "0.1.2"
2448 +source = "registry+https://github.com/rust-lang/crates.io-index"
2449 +checksum = "112b39cec0b298b6c1999fee3e31427f74f676e4cb9879ed1a121b43661a4154"
2450 +
2451 +[[package]]
2452 +name = "macro_rules_attribute"
2453 +version = "0.2.2"
2454 +source = "registry+https://github.com/rust-lang/crates.io-index"
2455 +checksum = "65049d7923698040cd0b1ddcced9b0eb14dd22c5f86ae59c3740eab64a676520"
2456 +dependencies = [
2457 + "macro_rules_attribute-proc_macro",
2458 + "paste",
2459 +]
2460 +
2461 +[[package]]
2462 +name = "macro_rules_attribute-proc_macro"
2463 +version = "0.2.2"
2464 +source = "registry+https://github.com/rust-lang/crates.io-index"
2465 +checksum = "670fdfda89751bc4a84ac13eaa63e205cf0fd22b4c9a5fbfa085b63c1f1d3a30"
2466 +
2467 +[[package]]
2468 +name = "malachite"
2469 +version = "0.4.22"
2470 +source = "registry+https://github.com/rust-lang/crates.io-index"
2471 +checksum = "2fbdf9cb251732db30a7200ebb6ae5d22fe8e11397364416617d2c2cf0c51cb5"
2472 +dependencies = [
2473 + "malachite-base",
2474 + "malachite-nz",
2475 + "malachite-q",
2476 +]
2477 +
2478 +[[package]]
2479 +name = "malachite-base"
2480 +version = "0.4.22"
2481 +source = "registry+https://github.com/rust-lang/crates.io-index"
2482 +checksum = "5ea0ed76adf7defc1a92240b5c36d5368cfe9251640dcce5bd2d0b7c1fd87aeb"
2483 +dependencies = [
2484 + "hashbrown 0.14.5",
2485 + "itertools 0.11.0",
2486 + "libm",
2487 + "ryu",
2488 +]
2489 +
2490 +[[package]]
2491 +name = "malachite-bigint"
2492 +version = "0.2.3"
2493 +source = "registry+https://github.com/rust-lang/crates.io-index"
2494 +checksum = "d149aaa2965d70381709d9df4c7ee1fc0de1c614a4efc2ee356f5e43d68749f8"
2495 +dependencies = [
2496 + "derive_more",
2497 + "malachite",
2498 + "num-integer",
2499 + "num-traits",
2500 + "paste",
2501 +]
2502 +
2503 +[[package]]
2504 +name = "malachite-nz"
2505 +version = "0.4.22"
2506 +source = "registry+https://github.com/rust-lang/crates.io-index"
2507 +checksum = "34a79feebb2bc9aa7762047c8e5495269a367da6b5a90a99882a0aeeac1841f7"
2508 +dependencies = [
2509 + "itertools 0.11.0",
2510 + "libm",
2511 + "malachite-base",
2512 +]
2513 +
2514 +[[package]]
2515 +name = "malachite-q"
2516 +version = "0.4.22"
2517 +source = "registry+https://github.com/rust-lang/crates.io-index"
2518 +checksum = "50f235d5747b1256b47620f5640c2a17a88c7569eebdf27cd9cb130e1a619191"
2519 +dependencies = [
2520 + "itertools 0.11.0",
2521 + "malachite-base",
2522 + "malachite-nz",
2523 +]
2524 +
2525 +[[package]]
2526 +name = "matchers"
2527 +version = "0.2.0"
2528 +source = "registry+https://github.com/rust-lang/crates.io-index"
2529 +checksum = "d1525a2a28c7f4fa0fc98bb91ae755d1e2d1505079e05539e35bc876b5d65ae9"
2530 +dependencies = [
2531 + "regex-automata",
2532 +]
2533 +
2534 +[[package]]
2535 +name = "matchit"
2536 +version = "0.8.4"
2537 +source = "registry+https://github.com/rust-lang/crates.io-index"
2538 +checksum = "47e1ffaa40ddd1f3ed91f717a33c8c0ee23fff369e3aa8772b9605cc1d22f4c3"
2539 +
2540 +[[package]]
2541 +name = "matrixmultiply"
2542 +version = "0.3.10"
2543 +source = "registry+https://github.com/rust-lang/crates.io-index"
2544 +checksum = "a06de3016e9fae57a36fd14dba131fccf49f74b40b7fbdb472f96e361ec71a08"
2545 +dependencies = [
2546 + "autocfg",
2547 + "rawpointer",
2548 +]
2549 +
2550 +[[package]]
2551 +name = "maybe-rayon"
2552 +version = "0.1.1"
2553 +source = "registry+https://github.com/rust-lang/crates.io-index"
2554 +checksum = "8ea1f30cedd69f0a2954655f7188c6a834246d2bcf1e315e2ac40c4b24dc9519"
2555 +dependencies = [
2556 + "cfg-if",
2557 + "rayon",
2558 +]
2559 +
2560 +[[package]]
2561 +name = "memchr"
2562 +version = "2.8.0"
2563 +source = "registry+https://github.com/rust-lang/crates.io-index"
2564 +checksum = "f8ca58f447f06ed17d5fc4043ce1b10dd205e060fb3ce5b979b8ed8e59ff3f79"
2565 +
2566 +[[package]]
2567 +name = "memo-map"
2568 +version = "0.3.3"
2569 +source = "registry+https://github.com/rust-lang/crates.io-index"
2570 +checksum = "38d1115007560874e373613744c6fba374c17688327a71c1476d1a5954cc857b"
2571 +
2572 +[[package]]
2573 +name = "mime"
2574 +version = "0.3.17"
2575 +source = "registry+https://github.com/rust-lang/crates.io-index"
2576 +checksum = "6877bb514081ee2a7ff5ef9de3281f14a4dd4bceac4c09388074a6b5df8a139a"
2577 +
2578 +[[package]]
2579 +name = "mime_guess"
2580 +version = "2.0.5"
2581 +source = "registry+https://github.com/rust-lang/crates.io-index"
2582 +checksum = "f7c44f8e672c00fe5308fa235f821cb4198414e1c77935c1ab6948d3fd78550e"
2583 +dependencies = [
2584 + "mime",
2585 + "unicase",
2586 +]
2587 +
2588 +[[package]]
2589 +name = "minijinja"
2590 +version = "2.18.0"
2591 +source = "registry+https://github.com/rust-lang/crates.io-index"
2592 +checksum = "328251e58ad8e415be6198888fc207502727dc77945806421ab34f35bf012e7d"
2593 +dependencies = [
2594 + "memo-map",
2595 + "serde",
2596 + "serde_json",
2597 +]
2598 +
2599 +[[package]]
2600 +name = "minijinja-contrib"
2601 +version = "2.18.0"
2602 +source = "registry+https://github.com/rust-lang/crates.io-index"
2603 +checksum = "8c6302e47d2b51f9fc978268ff7f5a014de5caa2ad48440309fd10ee711480d7"
2604 +dependencies = [
2605 + "minijinja",
2606 + "serde",
2607 +]
2608 +
2609 +[[package]]
2610 +name = "minimal-lexical"
2611 +version = "0.2.1"
2612 +source = "registry+https://github.com/rust-lang/crates.io-index"
2613 +checksum = "68354c5c6bd36d73ff3feceb05efa59b6acb7626617f4962be322a825e61f79a"
2614 +
2615 +[[package]]
2616 +name = "miniz_oxide"
2617 +version = "0.8.9"
2618 +source = "registry+https://github.com/rust-lang/crates.io-index"
2619 +checksum = "1fa76a2c86f704bdb222d66965fb3d63269ce38518b83cb0575fca855ebb6316"
2620 +dependencies = [
2621 + "adler2",
2622 + "simd-adler32",
2623 +]
2624 +
2625 +[[package]]
2626 +name = "mio"
2627 +version = "1.1.1"
2628 +source = "registry+https://github.com/rust-lang/crates.io-index"
2629 +checksum = "a69bcab0ad47271a0234d9422b131806bf3968021e5dc9328caf2d4cd58557fc"
2630 +dependencies = [
2631 + "libc",
2632 + "wasi",
2633 + "windows-sys 0.61.2",
2634 +]
2635 +
2636 +[[package]]
2637 +name = "monostate"
2638 +version = "0.1.18"
2639 +source = "registry+https://github.com/rust-lang/crates.io-index"
2640 +checksum = "3341a273f6c9d5bef1908f17b7267bbab0e95c9bf69a0d4dcf8e9e1b2c76ef67"
2641 +dependencies = [
2642 + "monostate-impl",
2643 + "serde",
2644 + "serde_core",
2645 +]
2646 +
2647 +[[package]]
2648 +name = "monostate-impl"
2649 +version = "0.1.18"
2650 +source = "registry+https://github.com/rust-lang/crates.io-index"
2651 +checksum = "e4db6d5580af57bf992f59068d4ea26fd518574ff48d7639b255a36f9de6e7e9"
2652 +dependencies = [
2653 + "proc-macro2",
2654 + "quote",
2655 + "syn 2.0.117",
2656 +]
2657 +
2658 +[[package]]
2659 +name = "moxcms"
2660 +version = "0.8.1"
2661 +source = "registry+https://github.com/rust-lang/crates.io-index"
2662 +checksum = "bb85c154ba489f01b25c0d36ae69a87e4a1c73a72631fc6c0eb6dde34a73e44b"
2663 +dependencies = [
2664 + "num-traits",
2665 + "pxfm",
2666 +]
2667 +
2668 +[[package]]
2669 +name = "multimap"
2670 +version = "0.10.1"
2671 +source = "registry+https://github.com/rust-lang/crates.io-index"
2672 +checksum = "1d87ecb2933e8aeadb3e3a02b828fed80a7528047e68b4f424523a0981a3a084"
2673 +
2674 +[[package]]
2675 +name = "native-tls"
2676 +version = "0.2.18"
2677 +source = "registry+https://github.com/rust-lang/crates.io-index"
2678 +checksum = "465500e14ea162429d264d44189adc38b199b62b1c21eea9f69e4b73cb03bbf2"
2679 +dependencies = [
2680 + "libc",
2681 + "log",
2682 + "openssl",
2683 + "openssl-probe",
2684 + "openssl-sys",
2685 + "schannel",
2686 + "security-framework",
2687 + "security-framework-sys",
2688 + "tempfile",
2689 +]
2690 +
2691 +[[package]]
2692 +name = "ndarray"
2693 +version = "0.16.1"
2694 +source = "registry+https://github.com/rust-lang/crates.io-index"
2695 +checksum = "882ed72dce9365842bf196bdeedf5055305f11fc8c03dee7bb0194a6cad34841"
2696 +dependencies = [
2697 + "matrixmultiply",
2698 + "num-complex",
2699 + "num-integer",
2700 + "num-traits",
2701 + "portable-atomic",
2702 + "portable-atomic-util",
2703 + "rawpointer",
2704 +]
2705 +
2706 +[[package]]
2707 +name = "ndarray"
2708 +version = "0.17.2"
2709 +source = "registry+https://github.com/rust-lang/crates.io-index"
2710 +checksum = "520080814a7a6b4a6e9070823bb24b4531daac8c4627e08ba5de8c5ef2f2752d"
2711 +dependencies = [
2712 + "matrixmultiply",
2713 + "num-complex",
2714 + "num-integer",
2715 + "num-traits",
2716 + "portable-atomic",
2717 + "portable-atomic-util",
2718 + "rawpointer",
2719 +]
2720 +
2721 +[[package]]
2722 +name = "new_debug_unreachable"
2723 +version = "1.0.6"
2724 +source = "registry+https://github.com/rust-lang/crates.io-index"
2725 +checksum = "650eef8c711430f1a879fdd01d4745a7deea475becfb90269c06775983bbf086"
2726 +
2727 +[[package]]
2728 +name = "no_std_io2"
2729 +version = "0.9.3"
2730 +source = "registry+https://github.com/rust-lang/crates.io-index"
2731 +checksum = "b51ed7824b6e07d354605f4abb3d9d300350701299da96642ee084f5ce631550"
2732 +dependencies = [
2733 + "memchr",
2734 +]
2735 +
2736 +[[package]]
2737 +name = "nom"
2738 +version = "7.1.3"
2739 +source = "registry+https://github.com/rust-lang/crates.io-index"
2740 +checksum = "d273983c5a657a70a3e8f2a01329822f3b8c8172b73826411a55751e404a0a4a"
2741 +dependencies = [
2742 + "memchr",
2743 + "minimal-lexical",
2744 +]
2745 +
2746 +[[package]]
2747 +name = "nom"
2748 +version = "8.0.0"
2749 +source = "registry+https://github.com/rust-lang/crates.io-index"
2750 +checksum = "df9761775871bdef83bee530e60050f7e54b1105350d6884eb0fb4f46c2f9405"
2751 +dependencies = [
2752 + "memchr",
2753 +]
2754 +
2755 +[[package]]
2756 +name = "noop_proc_macro"
2757 +version = "0.3.0"
2758 +source = "registry+https://github.com/rust-lang/crates.io-index"
2759 +checksum = "0676bb32a98c1a483ce53e500a81ad9c3d5b3f7c920c28c24e9cb0980d0b5bc8"
2760 +
2761 +[[package]]
2762 +name = "nu-ansi-term"
2763 +version = "0.50.3"
2764 +source = "registry+https://github.com/rust-lang/crates.io-index"
2765 +checksum = "7957b9740744892f114936ab4a57b3f487491bbeafaf8083688b16841a4240e5"
2766 +dependencies = [
2767 + "windows-sys 0.61.2",
2768 +]
2769 +
2770 +[[package]]
2771 +name = "num-bigint"
2772 +version = "0.4.6"
2773 +source = "registry+https://github.com/rust-lang/crates.io-index"
2774 +checksum = "a5e44f723f1133c9deac646763579fdb3ac745e418f2a7af9cd0c431da1f20b9"
2775 +dependencies = [
2776 + "num-integer",
2777 + "num-traits",
2778 +]
2779 +
2780 +[[package]]
2781 +name = "num-complex"
2782 +version = "0.4.6"
2783 +source = "registry+https://github.com/rust-lang/crates.io-index"
2784 +checksum = "73f88a1307638156682bada9d7604135552957b7818057dcef22705b4d509495"
2785 +dependencies = [
2786 + "num-traits",
2787 +]
2788 +
2789 +[[package]]
2790 +name = "num-conv"
2791 +version = "0.2.0"
2792 +source = "registry+https://github.com/rust-lang/crates.io-index"
2793 +checksum = "cf97ec579c3c42f953ef76dbf8d55ac91fb219dde70e49aa4a6b7d74e9919050"
2794 +
2795 +[[package]]
2796 +name = "num-derive"
2797 +version = "0.4.2"
2798 +source = "registry+https://github.com/rust-lang/crates.io-index"
2799 +checksum = "ed3955f1a9c7c0c15e092f9c887db08b1fc683305fdf6eb6684f22555355e202"
2800 +dependencies = [
2801 + "proc-macro2",
2802 + "quote",
2803 + "syn 2.0.117",
2804 +]
2805 +
2806 +[[package]]
2807 +name = "num-integer"
2808 +version = "0.1.46"
2809 +source = "registry+https://github.com/rust-lang/crates.io-index"
2810 +checksum = "7969661fd2958a5cb096e56c8e1ad0444ac2bbcd0061bd28660485a44879858f"
2811 +dependencies = [
2812 + "num-traits",
2813 +]
2814 +
2815 +[[package]]
2816 +name = "num-rational"
2817 +version = "0.4.2"
2818 +source = "registry+https://github.com/rust-lang/crates.io-index"
2819 +checksum = "f83d14da390562dca69fc84082e73e548e1ad308d24accdedd2720017cb37824"
2820 +dependencies = [
2821 + "num-bigint",
2822 + "num-integer",
2823 + "num-traits",
2824 +]
2825 +
2826 +[[package]]
2827 +name = "num-traits"
2828 +version = "0.2.19"
2829 +source = "registry+https://github.com/rust-lang/crates.io-index"
2830 +checksum = "071dfc062690e90b734c0b2273ce72ad0ffa95f0c74596bc250dcfd960262841"
2831 +dependencies = [
2832 + "autocfg",
2833 +]
2834 +
2835 +[[package]]
2836 +name = "num_cpus"
2837 +version = "1.17.0"
2838 +source = "registry+https://github.com/rust-lang/crates.io-index"
2839 +checksum = "91df4bbde75afed763b708b7eee1e8e7651e02d97f6d5dd763e89367e957b23b"
2840 +dependencies = [
2841 + "hermit-abi",
2842 + "libc",
2843 +]
2844 +
2845 +[[package]]
2846 +name = "num_threads"
2847 +version = "0.1.7"
2848 +source = "registry+https://github.com/rust-lang/crates.io-index"
2849 +checksum = "5c7398b9c8b70908f6371f47ed36737907c87c52af34c268fed0bf0ceb92ead9"
2850 +dependencies = [
2851 + "libc",
2852 +]
2853 +
2854 +[[package]]
2855 +name = "number_prefix"
2856 +version = "0.4.0"
2857 +source = "registry+https://github.com/rust-lang/crates.io-index"
2858 +checksum = "830b246a0e5f20af87141b25c173cd1b609bd7779a4617d6ec582abaf90870f3"
2859 +
2860 +[[package]]
2861 +name = "once_cell"
2862 +version = "1.21.3"
2863 +source = "registry+https://github.com/rust-lang/crates.io-index"
2864 +checksum = "42f5e15c9953c5e4ccceeb2e7382a716482c34515315f7b03532b8b4e8393d2d"
2865 +
2866 +[[package]]
2867 +name = "once_cell_polyfill"
2868 +version = "1.70.2"
2869 +source = "registry+https://github.com/rust-lang/crates.io-index"
2870 +checksum = "384b8ab6d37215f3c5301a95a4accb5d64aa607f1fcb26a11b5303878451b4fe"
2871 +
2872 +[[package]]
2873 +name = "onig"
2874 +version = "6.5.1"
2875 +source = "registry+https://github.com/rust-lang/crates.io-index"
2876 +checksum = "336b9c63443aceef14bea841b899035ae3abe89b7c486aaf4c5bd8aafedac3f0"
2877 +dependencies = [
2878 + "bitflags",
2879 + "libc",
2880 + "once_cell",
2881 + "onig_sys",
2882 +]
2883 +
2884 +[[package]]
2885 +name = "onig_sys"
2886 +version = "69.9.1"
2887 +source = "registry+https://github.com/rust-lang/crates.io-index"
2888 +checksum = "c7f86c6eef3d6df15f23bcfb6af487cbd2fed4e5581d58d5bf1f5f8b7f6727dc"
2889 +dependencies = [
2890 + "cc",
2891 + "pkg-config",
2892 +]
2893 +
2894 +[[package]]
2895 +name = "oorandom"
2896 +version = "11.1.5"
2897 +source = "registry+https://github.com/rust-lang/crates.io-index"
2898 +checksum = "d6790f58c7ff633d8771f42965289203411a5e5c68388703c06e14f24770b41e"
2899 +
2900 +[[package]]
2901 +name = "openai-harmony"
2902 +version = "0.0.8"
2903 +source = "registry+https://github.com/rust-lang/crates.io-index"
2904 +checksum = "e77e82af451fc95deeb728a40b84db8ee82d341e136c268de415123a560b9b72"
2905 +dependencies = [
2906 + "anyhow",
2907 + "base64 0.22.1",
2908 + "bstr",
2909 + "clap",
2910 + "fancy-regex 0.13.0",
2911 + "futures",
2912 + "image",
2913 + "regex",
2914 + "reqwest",
2915 + "rustc-hash 1.1.0",
2916 + "serde",
2917 + "serde_json",
2918 + "serde_with",
2919 + "sha1",
2920 + "sha2",
2921 + "thiserror 2.0.18",
2922 +]
2923 +
2924 +[[package]]
2925 +name = "openai-protocol"
2926 +version = "1.6.0"
2927 +source = "registry+https://github.com/rust-lang/crates.io-index"
2928 +checksum = "7b8d41ed865b7a26b6b2d2a519b774460e5ddc50eeeb4ac2a2409c8817ec2de9"
2929 +dependencies = [
2930 + "bitflags",
2931 + "chrono",
2932 + "rand 0.9.2",
2933 + "schemars 0.8.22",
2934 + "serde",
2935 + "serde_json",
2936 + "serde_with",
2937 + "tokio",
2938 + "tracing",
2939 + "url",
2940 + "validator",
2941 +]
2942 +
2943 +[[package]]
2944 +name = "openssl"
2945 +version = "0.10.76"
2946 +source = "registry+https://github.com/rust-lang/crates.io-index"
2947 +checksum = "951c002c75e16ea2c65b8c7e4d3d51d5530d8dfa7d060b4776828c88cfb18ecf"
2948 +dependencies = [
2949 + "bitflags",
2950 + "cfg-if",
2951 + "foreign-types",
2952 + "libc",
2953 + "once_cell",
2954 + "openssl-macros",
2955 + "openssl-sys",
2956 +]
2957 +
2958 +[[package]]
2959 +name = "openssl-macros"
2960 +version = "0.1.1"
2961 +source = "registry+https://github.com/rust-lang/crates.io-index"
2962 +checksum = "a948666b637a0f465e8564c73e89d4dde00d72d4d473cc972f390fc3dcee7d9c"
2963 +dependencies = [
2964 + "proc-macro2",
2965 + "quote",
2966 + "syn 2.0.117",
2967 +]
2968 +
2969 +[[package]]
2970 +name = "openssl-probe"
2971 +version = "0.2.1"
2972 +source = "registry+https://github.com/rust-lang/crates.io-index"
2973 +checksum = "7c87def4c32ab89d880effc9e097653c8da5d6ef28e6b539d313baaacfbafcbe"
2974 +
2975 +[[package]]
2976 +name = "openssl-src"
2977 +version = "300.5.5+3.5.5"
2978 +source = "registry+https://github.com/rust-lang/crates.io-index"
2979 +checksum = "3f1787d533e03597a7934fd0a765f0d28e94ecc5fb7789f8053b1e699a56f709"
2980 +dependencies = [
2981 + "cc",
2982 +]
2983 +
2984 +[[package]]
2985 +name = "openssl-sys"
2986 +version = "0.9.112"
2987 +source = "registry+https://github.com/rust-lang/crates.io-index"
2988 +checksum = "57d55af3b3e226502be1526dfdba67ab0e9c96fc293004e79576b2b9edb0dbdb"
2989 +dependencies = [
2990 + "cc",
2991 + "libc",
2992 + "openssl-src",
2993 + "pkg-config",
2994 + "vcpkg",
2995 +]
2996 +
2997 +[[package]]
2998 +name = "option-ext"
2999 +version = "0.2.0"
3000 +source = "registry+https://github.com/rust-lang/crates.io-index"
3001 +checksum = "04744f49eae99ab78e0d5c0b603ab218f515ea8cfe5a456d7629ad883a3b6e7d"
3002 +
3003 +[[package]]
3004 +name = "parking"
3005 +version = "2.2.1"
3006 +source = "registry+https://github.com/rust-lang/crates.io-index"
3007 +checksum = "f38d5652c16fde515bb1ecef450ab0f6a219d619a7274976324d5e377f7dceba"
3008 +
3009 +[[package]]
3010 +name = "parking_lot"
3011 +version = "0.12.5"
3012 +source = "registry+https://github.com/rust-lang/crates.io-index"
3013 +checksum = "93857453250e3077bd71ff98b6a65ea6621a19bb0f559a85248955ac12c45a1a"
3014 +dependencies = [
3015 + "lock_api",
3016 + "parking_lot_core",
3017 +]
3018 +
3019 +[[package]]
3020 +name = "parking_lot_core"
3021 +version = "0.9.12"
3022 +source = "registry+https://github.com/rust-lang/crates.io-index"
3023 +checksum = "2621685985a2ebf1c516881c026032ac7deafcda1a2c9b7850dc81e3dfcb64c1"
3024 +dependencies = [
3025 + "cfg-if",
3026 + "libc",
3027 + "redox_syscall",
3028 + "smallvec",
3029 + "windows-link",
3030 +]
3031 +
3032 +[[package]]
3033 +name = "paste"
3034 +version = "1.0.15"
3035 +source = "registry+https://github.com/rust-lang/crates.io-index"
3036 +checksum = "57c0d7b74b563b49d38dae00a0c37d4d6de9b432382b2892f0574ddcae73fd0a"
3037 +
3038 +[[package]]
3039 +name = "pastey"
3040 +version = "0.1.1"
3041 +source = "registry+https://github.com/rust-lang/crates.io-index"
3042 +checksum = "35fb2e5f958ec131621fdd531e9fc186ed768cbe395337403ae56c17a74c68ec"
3043 +
3044 +[[package]]
3045 +name = "pcre2"
3046 +version = "0.2.11"
3047 +source = "registry+https://github.com/rust-lang/crates.io-index"
3048 +checksum = "9e970b0fcce0c7ee6ef662744ff711f21ccd6f11b7cf03cd187a80e89797fc67"
3049 +dependencies = [
3050 + "libc",
3051 + "log",
3052 + "pcre2-sys",
3053 +]
3054 +
3055 +[[package]]
3056 +name = "pcre2-sys"
3057 +version = "0.2.10"
3058 +source = "registry+https://github.com/rust-lang/crates.io-index"
3059 +checksum = "18b9073c1a2549bd409bf4a32c94d903bb1a09bf845bc306ae148897fa0760a4"
3060 +dependencies = [
3061 + "cc",
3062 + "libc",
3063 + "pkg-config",
3064 +]
3065 +
3066 +[[package]]
3067 +name = "pem-rfc7468"
3068 +version = "1.0.0"
3069 +source = "registry+https://github.com/rust-lang/crates.io-index"
3070 +checksum = "a6305423e0e7738146434843d1694d621cce767262b2a86910beab705e4493d9"
3071 +dependencies = [
3072 + "base64ct",
3073 +]
3074 +
3075 +[[package]]
3076 +name = "percent-encoding"
3077 +version = "2.3.2"
3078 +source = "registry+https://github.com/rust-lang/crates.io-index"
3079 +checksum = "9b4f627cb1b25917193a259e49bdad08f671f8d9708acfd5fe0a8c1455d87220"
3080 +
3081 +[[package]]
3082 +name = "petgraph"
3083 +version = "0.8.3"
3084 +source = "registry+https://github.com/rust-lang/crates.io-index"
3085 +checksum = "8701b58ea97060d5e5b155d383a69952a60943f0e6dfe30b04c287beb0b27455"
3086 +dependencies = [
3087 + "fixedbitset",
3088 + "hashbrown 0.15.5",
3089 + "indexmap 2.13.0",
3090 +]
3091 +
3092 +[[package]]
3093 +name = "phf"
3094 +version = "0.11.3"
3095 +source = "registry+https://github.com/rust-lang/crates.io-index"
3096 +checksum = "1fd6780a80ae0c52cc120a26a1a42c1ae51b247a253e4e06113d23d2c2edd078"
3097 +dependencies = [
3098 + "phf_shared",
3099 +]
3100 +
3101 +[[package]]
3102 +name = "phf_codegen"
3103 +version = "0.11.3"
3104 +source = "registry+https://github.com/rust-lang/crates.io-index"
3105 +checksum = "aef8048c789fa5e851558d709946d6d79a8ff88c0440c587967f8e94bfb1216a"
3106 +dependencies = [
3107 + "phf_generator",
3108 + "phf_shared",
3109 +]
3110 +
3111 +[[package]]
3112 +name = "phf_generator"
3113 +version = "0.11.3"
3114 +source = "registry+https://github.com/rust-lang/crates.io-index"
3115 +checksum = "3c80231409c20246a13fddb31776fb942c38553c51e871f8cbd687a4cfb5843d"
3116 +dependencies = [
3117 + "phf_shared",
3118 + "rand 0.8.5",
3119 +]
3120 +
3121 +[[package]]
3122 +name = "phf_shared"
3123 +version = "0.11.3"
3124 +source = "registry+https://github.com/rust-lang/crates.io-index"
3125 +checksum = "67eabc2ef2a60eb7faa00097bd1ffdb5bd28e62bf39990626a582201b7a754e5"
3126 +dependencies = [
3127 + "siphasher",
3128 +]
3129 +
3130 +[[package]]
3131 +name = "pin-project"
3132 +version = "1.1.11"
3133 +source = "registry+https://github.com/rust-lang/crates.io-index"
3134 +checksum = "f1749c7ed4bcaf4c3d0a3efc28538844fb29bcdd7d2b67b2be7e20ba861ff517"
3135 +dependencies = [
3136 + "pin-project-internal",
3137 +]
3138 +
3139 +[[package]]
3140 +name = "pin-project-internal"
3141 +version = "1.1.11"
3142 +source = "registry+https://github.com/rust-lang/crates.io-index"
3143 +checksum = "d9b20ed30f105399776b9c883e68e536ef602a16ae6f596d2c473591d6ad64c6"
3144 +dependencies = [
3145 + "proc-macro2",
3146 + "quote",
3147 + "syn 2.0.117",
3148 +]
3149 +
3150 +[[package]]
3151 +name = "pin-project-lite"
3152 +version = "0.2.17"
3153 +source = "registry+https://github.com/rust-lang/crates.io-index"
3154 +checksum = "a89322df9ebe1c1578d689c92318e070967d1042b512afbe49518723f4e6d5cd"
3155 +
3156 +[[package]]
3157 +name = "pin-utils"
3158 +version = "0.1.0"
3159 +source = "registry+https://github.com/rust-lang/crates.io-index"
3160 +checksum = "8b870d8c151b6f2fb93e84a13146138f05d02ed11c7e7c54f8826aaaf7c9f184"
3161 +
3162 +[[package]]
3163 +name = "pkg-config"
3164 +version = "0.3.32"
3165 +source = "registry+https://github.com/rust-lang/crates.io-index"
3166 +checksum = "7edddbd0b52d732b21ad9a5fab5c704c14cd949e5e9a1ec5929a24fded1b904c"
3167 +
3168 +[[package]]
3169 +name = "plotters"
3170 +version = "0.3.7"
3171 +source = "registry+https://github.com/rust-lang/crates.io-index"
3172 +checksum = "5aeb6f403d7a4911efb1e33402027fc44f29b5bf6def3effcc22d7bb75f2b747"
3173 +dependencies = [
3174 + "num-traits",
3175 + "plotters-backend",
3176 + "plotters-svg",
3177 + "wasm-bindgen",
3178 + "web-sys",
3179 +]
3180 +
3181 +[[package]]
3182 +name = "plotters-backend"
3183 +version = "0.3.7"
3184 +source = "registry+https://github.com/rust-lang/crates.io-index"
3185 +checksum = "df42e13c12958a16b3f7f4386b9ab1f3e7933914ecea48da7139435263a4172a"
3186 +
3187 +[[package]]
3188 +name = "plotters-svg"
3189 +version = "0.3.7"
3190 +source = "registry+https://github.com/rust-lang/crates.io-index"
3191 +checksum = "51bae2ac328883f7acdfea3d66a7c35751187f870bc81f94563733a154d7a670"
3192 +dependencies = [
3193 + "plotters-backend",
3194 +]
3195 +
3196 +[[package]]
3197 +name = "png"
3198 +version = "0.18.1"
3199 +source = "registry+https://github.com/rust-lang/crates.io-index"
3200 +checksum = "60769b8b31b2a9f263dae2776c37b1b28ae246943cf719eb6946a1db05128a61"
3201 +dependencies = [
3202 + "bitflags",
3203 + "crc32fast",
3204 + "fdeflate",
3205 + "flate2",
3206 + "miniz_oxide",
3207 +]
3208 +
3209 +[[package]]
3210 +name = "polling"
3211 +version = "3.11.0"
3212 +source = "registry+https://github.com/rust-lang/crates.io-index"
3213 +checksum = "5d0e4f59085d47d8241c88ead0f274e8a0cb551f3625263c05eb8dd897c34218"
3214 +dependencies = [
3215 + "cfg-if",
3216 + "concurrent-queue",
3217 + "hermit-abi",
3218 + "pin-project-lite",
3219 + "rustix",
3220 + "windows-sys 0.61.2",
3221 +]
3222 +
3223 +[[package]]
3224 +name = "portable-atomic"
3225 +version = "1.13.1"
3226 +source = "registry+https://github.com/rust-lang/crates.io-index"
3227 +checksum = "c33a9471896f1c69cecef8d20cbe2f7accd12527ce60845ff44c153bb2a21b49"
3228 +
3229 +[[package]]
3230 +name = "portable-atomic-util"
3231 +version = "0.2.6"
3232 +source = "registry+https://github.com/rust-lang/crates.io-index"
3233 +checksum = "091397be61a01d4be58e7841595bd4bfedb15f1cd54977d79b8271e94ed799a3"
3234 +dependencies = [
3235 + "portable-atomic",
3236 +]
3237 +
3238 +[[package]]
3239 +name = "potential_utf"
3240 +version = "0.1.4"
3241 +source = "registry+https://github.com/rust-lang/crates.io-index"
3242 +checksum = "b73949432f5e2a09657003c25bca5e19a0e9c84f8058ca374f49e0ebe605af77"
3243 +dependencies = [
3244 + "zerovec",
3245 +]
3246 +
3247 +[[package]]
3248 +name = "powerfmt"
3249 +version = "0.2.0"
3250 +source = "registry+https://github.com/rust-lang/crates.io-index"
3251 +checksum = "439ee305def115ba05938db6eb1644ff94165c5ab5e9420d1c1bcedbba909391"
3252 +
3253 +[[package]]
3254 +name = "ppv-lite86"
3255 +version = "0.2.21"
3256 +source = "registry+https://github.com/rust-lang/crates.io-index"
3257 +checksum = "85eae3c4ed2f50dcfe72643da4befc30deadb458a9b590d720cde2f2b1e97da9"
3258 +dependencies = [
3259 + "zerocopy",
3260 +]
3261 +
3262 +[[package]]
3263 +name = "prettyplease"
3264 +version = "0.2.37"
3265 +source = "registry+https://github.com/rust-lang/crates.io-index"
3266 +checksum = "479ca8adacdd7ce8f1fb39ce9ecccbfe93a3f1344b3d0d97f20bc0196208f62b"
3267 +dependencies = [
3268 + "proc-macro2",
3269 + "syn 2.0.117",
3270 +]
3271 +
3272 +[[package]]
3273 +name = "primal-check"
3274 +version = "0.3.4"
3275 +source = "registry+https://github.com/rust-lang/crates.io-index"
3276 +checksum = "dc0d895b311e3af9902528fbb8f928688abbd95872819320517cc24ca6b2bd08"
3277 +dependencies = [
3278 + "num-integer",
3279 +]
3280 +
3281 +[[package]]
3282 +name = "proc-macro-crate"
3283 +version = "3.5.0"
3284 +source = "registry+https://github.com/rust-lang/crates.io-index"
3285 +checksum = "e67ba7e9b2b56446f1d419b1d807906278ffa1a658a8a5d8a39dcb1f5a78614f"
3286 +dependencies = [
3287 + "toml_edit",
3288 +]
3289 +
3290 +[[package]]
3291 +name = "proc-macro-error-attr2"
3292 +version = "2.0.0"
3293 +source = "registry+https://github.com/rust-lang/crates.io-index"
3294 +checksum = "96de42df36bb9bba5542fe9f1a054b8cc87e172759a1868aa05c1f3acc89dfc5"
3295 +dependencies = [
3296 + "proc-macro2",
3297 + "quote",
3298 +]
3299 +
3300 +[[package]]
3301 +name = "proc-macro-error2"
3302 +version = "2.0.1"
3303 +source = "registry+https://github.com/rust-lang/crates.io-index"
3304 +checksum = "11ec05c52be0a07b08061f7dd003e7d7092e0472bc731b4af7bb1ef876109802"
3305 +dependencies = [
3306 + "proc-macro-error-attr2",
3307 + "proc-macro2",
3308 + "quote",
3309 + "syn 2.0.117",
3310 +]
3311 +
3312 +[[package]]
3313 +name = "proc-macro2"
3314 +version = "1.0.106"
3315 +source = "registry+https://github.com/rust-lang/crates.io-index"
3316 +checksum = "8fd00f0bb2e90d81d1044c2b32617f68fcb9fa3bb7640c23e9c748e53fb30934"
3317 +dependencies = [
3318 + "unicode-ident",
3319 +]
3320 +
3321 +[[package]]
3322 +name = "profiling"
3323 +version = "1.0.17"
3324 +source = "registry+https://github.com/rust-lang/crates.io-index"
3325 +checksum = "3eb8486b569e12e2c32ad3e204dbaba5e4b5b216e9367044f25f1dba42341773"
3326 +dependencies = [
3327 + "profiling-procmacros",
3328 +]
3329 +
3330 +[[package]]
3331 +name = "profiling-procmacros"
3332 +version = "1.0.17"
3333 +source = "registry+https://github.com/rust-lang/crates.io-index"
3334 +checksum = "52717f9a02b6965224f95ca2a81e2e0c5c43baacd28ca057577988930b6c3d5b"
3335 +dependencies = [
3336 + "quote",
3337 + "syn 2.0.117",
3338 +]
3339 +
3340 +[[package]]
3341 +name = "prometheus-client"
3342 +version = "0.24.0"
3343 +source = "registry+https://github.com/rust-lang/crates.io-index"
3344 +checksum = "e4500adecd7af8e0e9f4dbce15cfee07ce913fbf6ad605cc468b83f2d531ee94"
3345 +dependencies = [
3346 + "dtoa",
3347 + "itoa",
3348 + "parking_lot",
3349 + "prometheus-client-derive-encode",
3350 +]
3351 +
3352 +[[package]]
3353 +name = "prometheus-client-derive-encode"
3354 +version = "0.5.0"
3355 +source = "registry+https://github.com/rust-lang/crates.io-index"
3356 +checksum = "9adf1691c04c0a5ff46ff8f262b58beb07b0dbb61f96f9f54f6cbd82106ed87f"
3357 +dependencies = [
3358 + "proc-macro2",
3359 + "quote",
3360 + "syn 2.0.117",
3361 +]
3362 +
3363 +[[package]]
3364 +name = "prost"
3365 +version = "0.14.3"
3366 +source = "registry+https://github.com/rust-lang/crates.io-index"
3367 +checksum = "d2ea70524a2f82d518bce41317d0fae74151505651af45faf1ffbd6fd33f0568"
3368 +dependencies = [
3369 + "bytes",
3370 + "prost-derive",
3371 +]
3372 +
3373 +[[package]]
3374 +name = "prost-build"
3375 +version = "0.14.3"
3376 +source = "registry+https://github.com/rust-lang/crates.io-index"
3377 +checksum = "343d3bd7056eda839b03204e68deff7d1b13aba7af2b2fd16890697274262ee7"
3378 +dependencies = [
3379 + "heck",
3380 + "itertools 0.14.0",
3381 + "log",
3382 + "multimap",
3383 + "petgraph",
3384 + "prettyplease",
3385 + "prost",
3386 + "prost-types",
3387 + "pulldown-cmark",
3388 + "pulldown-cmark-to-cmark",
3389 + "regex",
3390 + "syn 2.0.117",
3391 + "tempfile",
3392 +]
3393 +
3394 +[[package]]
3395 +name = "prost-derive"
3396 +version = "0.14.3"
3397 +source = "registry+https://github.com/rust-lang/crates.io-index"
3398 +checksum = "27c6023962132f4b30eb4c172c91ce92d933da334c59c23cddee82358ddafb0b"
3399 +dependencies = [
3400 + "anyhow",
3401 + "itertools 0.14.0",
3402 + "proc-macro2",
3403 + "quote",
3404 + "syn 2.0.117",
3405 +]
3406 +
3407 +[[package]]
3408 +name = "prost-types"
3409 +version = "0.14.3"
3410 +source = "registry+https://github.com/rust-lang/crates.io-index"
3411 +checksum = "8991c4cbdb8bc5b11f0b074ffe286c30e523de90fee5ba8132f1399f23cb3dd7"
3412 +dependencies = [
3413 + "prost",
3414 +]
3415 +
3416 +[[package]]
3417 +name = "pulldown-cmark"
3418 +version = "0.13.3"
3419 +source = "registry+https://github.com/rust-lang/crates.io-index"
3420 +checksum = "7c3a14896dfa883796f1cb410461aef38810ea05f2b2c33c5aded3649095fdad"
3421 +dependencies = [
3422 + "bitflags",
3423 + "memchr",
3424 + "unicase",
3425 +]
3426 +
3427 +[[package]]
3428 +name = "pulldown-cmark-to-cmark"
3429 +version = "22.0.0"
3430 +source = "registry+https://github.com/rust-lang/crates.io-index"
3431 +checksum = "50793def1b900256624a709439404384204a5dc3a6ec580281bfaac35e882e90"
3432 +dependencies = [
3433 + "pulldown-cmark",
3434 +]
3435 +
3436 +[[package]]
3437 +name = "pxfm"
3438 +version = "0.1.29"
3439 +source = "registry+https://github.com/rust-lang/crates.io-index"
3440 +checksum = "e0c5ccf5294c6ccd63a74f1565028353830a9c2f5eb0c682c355c471726a6e3f"
3441 +
3442 +[[package]]
3443 +name = "qoi"
3444 +version = "0.4.1"
3445 +source = "registry+https://github.com/rust-lang/crates.io-index"
3446 +checksum = "7f6d64c71eb498fe9eae14ce4ec935c555749aef511cca85b5568910d6e48001"
3447 +dependencies = [
3448 + "bytemuck",
3449 +]
3450 +
3451 +[[package]]
3452 +name = "quick-error"
3453 +version = "2.0.1"
3454 +source = "registry+https://github.com/rust-lang/crates.io-index"
3455 +checksum = "a993555f31e5a609f617c12db6250dedcac1b0a85076912c436e6fc9b2c8e6a3"
3456 +
3457 +[[package]]
3458 +name = "quinn"
3459 +version = "0.11.9"
3460 +source = "registry+https://github.com/rust-lang/crates.io-index"
3461 +checksum = "b9e20a958963c291dc322d98411f541009df2ced7b5a4f2bd52337638cfccf20"
3462 +dependencies = [
3463 + "bytes",
3464 + "cfg_aliases",
3465 + "pin-project-lite",
3466 + "quinn-proto",
3467 + "quinn-udp",
3468 + "rustc-hash 2.1.1",
3469 + "rustls",
3470 + "socket2",
3471 + "thiserror 2.0.18",
3472 + "tokio",
3473 + "tracing",
3474 + "web-time",
3475 +]
3476 +
3477 +[[package]]
3478 +name = "quinn-proto"
3479 +version = "0.11.14"
3480 +source = "registry+https://github.com/rust-lang/crates.io-index"
3481 +checksum = "434b42fec591c96ef50e21e886936e66d3cc3f737104fdb9b737c40ffb94c098"
3482 +dependencies = [
3483 + "bytes",
3484 + "getrandom 0.3.4",
3485 + "lru-slab",
3486 + "rand 0.9.2",
3487 + "ring",
3488 + "rustc-hash 2.1.1",
3489 + "rustls",
3490 + "rustls-pki-types",
3491 + "slab",
3492 + "thiserror 2.0.18",
3493 + "tinyvec",
3494 + "tracing",
3495 + "web-time",
3496 +]
3497 +
3498 +[[package]]
3499 +name = "quinn-udp"
3500 +version = "0.5.14"
3501 +source = "registry+https://github.com/rust-lang/crates.io-index"
3502 +checksum = "addec6a0dcad8a8d96a771f815f0eaf55f9d1805756410b39f5fa81332574cbd"
3503 +dependencies = [
3504 + "cfg_aliases",
3505 + "libc",
3506 + "once_cell",
3507 + "socket2",
3508 + "tracing",
3509 + "windows-sys 0.60.2",
3510 +]
3511 +
3512 +[[package]]
3513 +name = "quote"
3514 +version = "1.0.45"
3515 +source = "registry+https://github.com/rust-lang/crates.io-index"
3516 +checksum = "41f2619966050689382d2b44f664f4bc593e129785a36d6ee376ddf37259b924"
3517 +dependencies = [
3518 + "proc-macro2",
3519 +]
3520 +
3521 +[[package]]
3522 +name = "r-efi"
3523 +version = "5.3.0"
3524 +source = "registry+https://github.com/rust-lang/crates.io-index"
3525 +checksum = "69cdb34c158ceb288df11e18b4bd39de994f6657d83847bdffdbd7f346754b0f"
3526 +
3527 +[[package]]
3528 +name = "r-efi"
3529 +version = "6.0.0"
3530 +source = "registry+https://github.com/rust-lang/crates.io-index"
3531 +checksum = "f8dcc9c7d52a811697d2151c701e0d08956f92b0e24136cf4cf27b57a6a0d9bf"
3532 +
3533 +[[package]]
3534 +name = "rand"
3535 +version = "0.8.5"
3536 +source = "registry+https://github.com/rust-lang/crates.io-index"
3537 +checksum = "34af8d1a0e25924bc5b7c43c079c942339d8f0a8b57c39049bef581b46327404"
3538 +dependencies = [
3539 + "libc",
3540 + "rand_chacha 0.3.1",
3541 + "rand_core 0.6.4",
3542 +]
3543 +
3544 +[[package]]
3545 +name = "rand"
3546 +version = "0.9.2"
3547 +source = "registry+https://github.com/rust-lang/crates.io-index"
3548 +checksum = "6db2770f06117d490610c7488547d543617b21bfa07796d7a12f6f1bd53850d1"
3549 +dependencies = [
3550 + "rand_chacha 0.9.0",
3551 + "rand_core 0.9.5",
3552 +]
3553 +
3554 +[[package]]
3555 +name = "rand_chacha"
3556 +version = "0.3.1"
3557 +source = "registry+https://github.com/rust-lang/crates.io-index"
3558 +checksum = "e6c10a63a0fa32252be49d21e7709d4d4baf8d231c2dbce1eaa8141b9b127d88"
3559 +dependencies = [
3560 + "ppv-lite86",
3561 + "rand_core 0.6.4",
3562 +]
3563 +
3564 +[[package]]
3565 +name = "rand_chacha"
3566 +version = "0.9.0"
3567 +source = "registry+https://github.com/rust-lang/crates.io-index"
3568 +checksum = "d3022b5f1df60f26e1ffddd6c66e8aa15de382ae63b3a0c1bfc0e4d3e3f325cb"
3569 +dependencies = [
3570 + "ppv-lite86",
3571 + "rand_core 0.9.5",
3572 +]
3573 +
3574 +[[package]]
3575 +name = "rand_core"
3576 +version = "0.6.4"
3577 +source = "registry+https://github.com/rust-lang/crates.io-index"
3578 +checksum = "ec0be4795e2f6a28069bec0b5ff3e2ac9bafc99e6a9a7dc3547996c5c816922c"
3579 +dependencies = [
3580 + "getrandom 0.2.17",
3581 +]
3582 +
3583 +[[package]]
3584 +name = "rand_core"
3585 +version = "0.9.5"
3586 +source = "registry+https://github.com/rust-lang/crates.io-index"
3587 +checksum = "76afc826de14238e6e8c374ddcc1fa19e374fd8dd986b0d2af0d02377261d83c"
3588 +dependencies = [
3589 + "getrandom 0.3.4",
3590 +]
3591 +
3592 +[[package]]
3593 +name = "rav1e"
3594 +version = "0.8.1"
3595 +source = "registry+https://github.com/rust-lang/crates.io-index"
3596 +checksum = "43b6dd56e85d9483277cde964fd1bdb0428de4fec5ebba7540995639a21cb32b"
3597 +dependencies = [
3598 + "aligned-vec",
3599 + "arbitrary",
3600 + "arg_enum_proc_macro",
3601 + "arrayvec",
3602 + "av-scenechange",
3603 + "av1-grain",
3604 + "bitstream-io",
3605 + "built",
3606 + "cfg-if",
3607 + "interpolate_name",
3608 + "itertools 0.14.0",
3609 + "libc",
3610 + "libfuzzer-sys",
3611 + "log",
3612 + "maybe-rayon",
3613 + "new_debug_unreachable",
3614 + "noop_proc_macro",
3615 + "num-derive",
3616 + "num-traits",
3617 + "paste",
3618 + "profiling",
3619 + "rand 0.9.2",
3620 + "rand_chacha 0.9.0",
3621 + "simd_helpers",
3622 + "thiserror 2.0.18",
3623 + "v_frame",
3624 + "wasm-bindgen",
3625 +]
3626 +
3627 +[[package]]
3628 +name = "ravif"
3629 +version = "0.13.0"
3630 +source = "registry+https://github.com/rust-lang/crates.io-index"
3631 +checksum = "e52310197d971b0f5be7fe6b57530dcd27beb35c1b013f29d66c1ad73fbbcc45"
3632 +dependencies = [
3633 + "avif-serialize",
3634 + "imgref",
3635 + "loop9",
3636 + "quick-error",
3637 + "rav1e",
3638 + "rayon",
3639 + "rgb",
3640 +]
3641 +
3642 +[[package]]
3643 +name = "rawpointer"
3644 +version = "0.2.1"
3645 +source = "registry+https://github.com/rust-lang/crates.io-index"
3646 +checksum = "60a357793950651c4ed0f3f52338f53b2f809f32d83a07f72909fa13e4c6c1e3"
3647 +
3648 +[[package]]
3649 +name = "rayon"
3650 +version = "1.11.0"
3651 +source = "registry+https://github.com/rust-lang/crates.io-index"
3652 +checksum = "368f01d005bf8fd9b1206fb6fa653e6c4a81ceb1466406b81792d87c5677a58f"
3653 +dependencies = [
3654 + "either",
3655 + "rayon-core",
3656 +]
3657 +
3658 +[[package]]
3659 +name = "rayon-cond"
3660 +version = "0.4.0"
3661 +source = "registry+https://github.com/rust-lang/crates.io-index"
3662 +checksum = "2964d0cf57a3e7a06e8183d14a8b527195c706b7983549cd5462d5aa3747438f"
3663 +dependencies = [
3664 + "either",
3665 + "itertools 0.14.0",
3666 + "rayon",
3667 +]
3668 +
3669 +[[package]]
3670 +name = "rayon-core"
3671 +version = "1.13.0"
3672 +source = "registry+https://github.com/rust-lang/crates.io-index"
3673 +checksum = "22e18b0f0062d30d4230b2e85ff77fdfe4326feb054b9783a3460d8435c8ab91"
3674 +dependencies = [
3675 + "crossbeam-deque",
3676 + "crossbeam-utils",
3677 +]
3678 +
3679 +[[package]]
3680 +name = "realfft"
3681 +version = "3.5.0"
3682 +source = "registry+https://github.com/rust-lang/crates.io-index"
3683 +checksum = "f821338fddb99d089116342c46e9f1fbf3828dba077674613e734e01d6ea8677"
3684 +dependencies = [
3685 + "rustfft",
3686 +]
3687 +
3688 +[[package]]
3689 +name = "redox_syscall"
3690 +version = "0.5.18"
3691 +source = "registry+https://github.com/rust-lang/crates.io-index"
3692 +checksum = "ed2bf2547551a7053d6fdfafda3f938979645c44812fbfcda098faae3f1a362d"
3693 +dependencies = [
3694 + "bitflags",
3695 +]
3696 +
3697 +[[package]]
3698 +name = "redox_users"
3699 +version = "0.5.2"
3700 +source = "registry+https://github.com/rust-lang/crates.io-index"
3701 +checksum = "a4e608c6638b9c18977b00b475ac1f28d14e84b27d8d42f70e0bf1e3dec127ac"
3702 +dependencies = [
3703 + "getrandom 0.2.17",
3704 + "libredox",
3705 + "thiserror 2.0.18",
3706 +]
3707 +
3708 +[[package]]
3709 +name = "ref-cast"
3710 +version = "1.0.25"
3711 +source = "registry+https://github.com/rust-lang/crates.io-index"
3712 +checksum = "f354300ae66f76f1c85c5f84693f0ce81d747e2c3f21a45fef496d89c960bf7d"
3713 +dependencies = [
3714 + "ref-cast-impl",
3715 +]
3716 +
3717 +[[package]]
3718 +name = "ref-cast-impl"
3719 +version = "1.0.25"
3720 +source = "registry+https://github.com/rust-lang/crates.io-index"
3721 +checksum = "b7186006dcb21920990093f30e3dea63b7d6e977bf1256be20c3563a5db070da"
3722 +dependencies = [
3723 + "proc-macro2",
3724 + "quote",
3725 + "syn 2.0.117",
3726 +]
3727 +
3728 +[[package]]
3729 +name = "regex"
3730 +version = "1.12.3"
3731 +source = "registry+https://github.com/rust-lang/crates.io-index"
3732 +checksum = "e10754a14b9137dd7b1e3e5b0493cc9171fdd105e0ab477f51b72e7f3ac0e276"
3733 +dependencies = [
3734 + "aho-corasick",
3735 + "memchr",
3736 + "regex-automata",
3737 + "regex-syntax",
3738 +]
3739 +
3740 +[[package]]
3741 +name = "regex-automata"
3742 +version = "0.4.14"
3743 +source = "registry+https://github.com/rust-lang/crates.io-index"
3744 +checksum = "6e1dd4122fc1595e8162618945476892eefca7b88c52820e74af6262213cae8f"
3745 +dependencies = [
3746 + "aho-corasick",
3747 + "memchr",
3748 + "regex-syntax",
3749 +]
3750 +
3751 +[[package]]
3752 +name = "regex-syntax"
3753 +version = "0.8.10"
3754 +source = "registry+https://github.com/rust-lang/crates.io-index"
3755 +checksum = "dc897dd8d9e8bd1ed8cdad82b5966c3e0ecae09fb1907d58efaa013543185d0a"
3756 +
3757 +[[package]]
3758 +name = "reqwest"
3759 +version = "0.12.28"
3760 +source = "registry+https://github.com/rust-lang/crates.io-index"
3761 +checksum = "eddd3ca559203180a307f12d114c268abf583f59b03cb906fd0b3ff8646c1147"
3762 +dependencies = [
3763 + "base64 0.22.1",
3764 + "bytes",
3765 + "encoding_rs",
3766 + "futures-channel",
3767 + "futures-core",
3768 + "futures-util",
3769 + "h2",
3770 + "http",
3771 + "http-body",
3772 + "http-body-util",
3773 + "hyper",
3774 + "hyper-rustls",
3775 + "hyper-tls",
3776 + "hyper-util",
3777 + "js-sys",
3778 + "log",
3779 + "mime",
3780 + "mime_guess",
3781 + "native-tls",
3782 + "percent-encoding",
3783 + "pin-project-lite",
3784 + "quinn",
3785 + "rustls",
3786 + "rustls-native-certs",
3787 + "rustls-pki-types",
3788 + "serde",
3789 + "serde_json",
3790 + "serde_urlencoded",
3791 + "sync_wrapper",
3792 + "tokio",
3793 + "tokio-native-tls",
3794 + "tokio-rustls",
3795 + "tokio-util",
3796 + "tower",
3797 + "tower-http",
3798 + "tower-service",
3799 + "url",
3800 + "wasm-bindgen",
3801 + "wasm-bindgen-futures",
3802 + "wasm-streams",
3803 + "web-sys",
3804 + "webpki-roots 1.0.6",
3805 +]
3806 +
3807 +[[package]]
3808 +name = "reqwest-eventsource"
3809 +version = "0.6.0"
3810 +source = "registry+https://github.com/rust-lang/crates.io-index"
3811 +checksum = "632c55746dbb44275691640e7b40c907c16a2dc1a5842aa98aaec90da6ec6bde"
3812 +dependencies = [
3813 + "eventsource-stream",
3814 + "futures-core",
3815 + "futures-timer",
3816 + "mime",
3817 + "nom 7.1.3",
3818 + "pin-project-lite",
3819 + "reqwest",
3820 + "thiserror 1.0.69",
3821 +]
3822 +
3823 +[[package]]
3824 +name = "rgb"
3825 +version = "0.8.53"
3826 +source = "registry+https://github.com/rust-lang/crates.io-index"
3827 +checksum = "47b34b781b31e5d73e9fbc8689c70551fd1ade9a19e3e28cfec8580a79290cc4"
3828 +
3829 +[[package]]
3830 +name = "ring"
3831 +version = "0.17.14"
3832 +source = "registry+https://github.com/rust-lang/crates.io-index"
3833 +checksum = "a4689e6c2294d81e88dc6261c768b63bc4fcdb852be6d1352498b114f61383b7"
3834 +dependencies = [
3835 + "cc",
3836 + "cfg-if",
3837 + "getrandom 0.2.17",
3838 + "libc",
3839 + "untrusted",
3840 + "windows-sys 0.52.0",
3841 +]
3842 +
3843 +[[package]]
3844 +name = "riptoken"
3845 +version = "0.3.0"
3846 +source = "registry+https://github.com/rust-lang/crates.io-index"
3847 +checksum = "196b37c4dd48f99b51e8aeaa3d0c343df62c7bb5b66cd6735aa072524fbe8665"
3848 +dependencies = [
3849 + "fancy-regex 0.17.0",
3850 + "rayon",
3851 + "regex",
3852 + "regex-automata",
3853 + "rustc-hash 2.1.1",
3854 +]
3855 +
3856 +[[package]]
3857 +name = "rmp"
3858 +version = "0.8.15"
3859 +source = "registry+https://github.com/rust-lang/crates.io-index"
3860 +checksum = "4ba8be72d372b2c9b35542551678538b562e7cf86c3315773cae48dfbfe7790c"
3861 +dependencies = [
3862 + "num-traits",
3863 +]
3864 +
3865 +[[package]]
3866 +name = "rmp-serde"
3867 +version = "1.3.1"
3868 +source = "registry+https://github.com/rust-lang/crates.io-index"
3869 +checksum = "72f81bee8c8ef9b577d1681a70ebbc962c232461e397b22c208c43c04b67a155"
3870 +dependencies = [
3871 + "rmp",
3872 + "serde",
3873 +]
3874 +
3875 +[[package]]
3876 +name = "rmpv"
3877 +version = "1.3.1"
3878 +source = "registry+https://github.com/rust-lang/crates.io-index"
3879 +checksum = "7a4e1d4b9b938a26d2996af33229f0ca0956c652c1375067f0b45291c1df8417"
3880 +dependencies = [
3881 + "rmp",
3882 + "serde",
3883 + "serde_bytes",
3884 +]
3885 +
3886 +[[package]]
3887 +name = "rubato"
3888 +version = "0.16.2"
3889 +source = "registry+https://github.com/rust-lang/crates.io-index"
3890 +checksum = "5258099699851cfd0082aeb645feb9c084d9a5e1f1b8d5372086b989fc5e56a1"
3891 +dependencies = [
3892 + "num-complex",
3893 + "num-integer",
3894 + "num-traits",
3895 + "realfft",
3896 +]
3897 +
3898 +[[package]]
3899 +name = "rustc-hash"
3900 +version = "1.1.0"
3901 +source = "registry+https://github.com/rust-lang/crates.io-index"
3902 +checksum = "08d43f7aa6b08d49f382cde6a7982047c3426db949b1424bc4b7ec9ae12c6ce2"
3903 +
3904 +[[package]]
3905 +name = "rustc-hash"
3906 +version = "2.1.1"
3907 +source = "registry+https://github.com/rust-lang/crates.io-index"
3908 +checksum = "357703d41365b4b27c590e3ed91eabb1b663f07c4c084095e60cbed4362dff0d"
3909 +
3910 +[[package]]
3911 +name = "rustfft"
3912 +version = "6.4.1"
3913 +source = "registry+https://github.com/rust-lang/crates.io-index"
3914 +checksum = "21db5f9893e91f41798c88680037dba611ca6674703c1a18601b01a72c8adb89"
3915 +dependencies = [
3916 + "num-complex",
3917 + "num-integer",
3918 + "num-traits",
3919 + "primal-check",
3920 + "strength_reduce",
3921 + "transpose",
3922 +]
3923 +
3924 +[[package]]
3925 +name = "rustix"
3926 +version = "1.1.4"
3927 +source = "registry+https://github.com/rust-lang/crates.io-index"
3928 +checksum = "b6fe4565b9518b83ef4f91bb47ce29620ca828bd32cb7e408f0062e9930ba190"
3929 +dependencies = [
3930 + "bitflags",
3931 + "errno",
3932 + "libc",
3933 + "linux-raw-sys",
3934 + "windows-sys 0.61.2",
3935 +]
3936 +
3937 +[[package]]
3938 +name = "rustls"
3939 +version = "0.23.37"
3940 +source = "registry+https://github.com/rust-lang/crates.io-index"
3941 +checksum = "758025cb5fccfd3bc2fd74708fd4682be41d99e5dff73c377c0646c6012c73a4"
3942 +dependencies = [
3943 + "log",
3944 + "once_cell",
3945 + "ring",
3946 + "rustls-pki-types",
3947 + "rustls-webpki",
3948 + "subtle",
3949 + "zeroize",
3950 +]
3951 +
3952 +[[package]]
3953 +name = "rustls-native-certs"
3954 +version = "0.8.3"
3955 +source = "registry+https://github.com/rust-lang/crates.io-index"
3956 +checksum = "612460d5f7bea540c490b2b6395d8e34a953e52b491accd6c86c8164c5932a63"
3957 +dependencies = [
3958 + "openssl-probe",
3959 + "rustls-pki-types",
3960 + "schannel",
3961 + "security-framework",
3962 +]
3963 +
3964 +[[package]]
3965 +name = "rustls-pki-types"
3966 +version = "1.14.0"
3967 +source = "registry+https://github.com/rust-lang/crates.io-index"
3968 +checksum = "be040f8b0a225e40375822a563fa9524378b9d63112f53e19ffff34df5d33fdd"
3969 +dependencies = [
3970 + "web-time",
3971 + "zeroize",
3972 +]
3973 +
3974 +[[package]]
3975 +name = "rustls-webpki"
3976 +version = "0.103.9"
3977 +source = "registry+https://github.com/rust-lang/crates.io-index"
3978 +checksum = "d7df23109aa6c1567d1c575b9952556388da57401e4ace1d15f79eedad0d8f53"
3979 +dependencies = [
3980 + "ring",
3981 + "rustls-pki-types",
3982 + "untrusted",
3983 +]
3984 +
3985 +[[package]]
3986 +name = "rustpython-ast"
3987 +version = "0.4.0"
3988 +source = "registry+https://github.com/rust-lang/crates.io-index"
3989 +checksum = "4cdaf8ee5c1473b993b398c174641d3aa9da847af36e8d5eb8291930b72f31a5"
3990 +dependencies = [
3991 + "is-macro",
3992 + "malachite-bigint",
3993 + "rustpython-parser-core",
3994 + "static_assertions",
3995 +]
3996 +
3997 +[[package]]
3998 +name = "rustpython-parser"
3999 +version = "0.4.0"
4000 +source = "registry+https://github.com/rust-lang/crates.io-index"
4001 +checksum = "868f724daac0caf9bd36d38caf45819905193a901e8f1c983345a68e18fb2abb"
4002 +dependencies = [
4003 + "anyhow",
4004 + "is-macro",
4005 + "itertools 0.11.0",
4006 + "lalrpop-util",
4007 + "log",
4008 + "malachite-bigint",
4009 + "num-traits",
4010 + "phf",
4011 + "phf_codegen",
4012 + "rustc-hash 1.1.0",
4013 + "rustpython-ast",
4014 + "rustpython-parser-core",
4015 + "tiny-keccak",
4016 + "unic-emoji-char",
4017 + "unic-ucd-ident",
4018 + "unicode_names2",
4019 +]
4020 +
4021 +[[package]]
4022 +name = "rustpython-parser-core"
4023 +version = "0.4.0"
4024 +source = "registry+https://github.com/rust-lang/crates.io-index"
4025 +checksum = "b4b6c12fa273825edc7bccd9a734f0ad5ba4b8a2f4da5ff7efe946f066d0f4ad"
4026 +dependencies = [
4027 + "is-macro",
4028 + "memchr",
4029 + "rustpython-parser-vendored",
4030 +]
4031 +
4032 +[[package]]
4033 +name = "rustpython-parser-vendored"
4034 +version = "0.4.0"
4035 +source = "registry+https://github.com/rust-lang/crates.io-index"
4036 +checksum = "04fcea49a4630a3a5d940f4d514dc4f575ed63c14c3e3ed07146634aed7f67a6"
4037 +dependencies = [
4038 + "memchr",
4039 + "once_cell",
4040 +]
4041 +
4042 +[[package]]
4043 +name = "rustversion"
4044 +version = "1.0.22"
4045 +source = "registry+https://github.com/rust-lang/crates.io-index"
4046 +checksum = "b39cdef0fa800fc44525c84ccb54a029961a8215f9619753635a9c0d2538d46d"
4047 +
4048 +[[package]]
4049 +name = "ryu"
4050 +version = "1.0.23"
4051 +source = "registry+https://github.com/rust-lang/crates.io-index"
4052 +checksum = "9774ba4a74de5f7b1c1451ed6cd5285a32eddb5cccb8cc655a4e50009e06477f"
4053 +
4054 +[[package]]
4055 +name = "saa"
4056 +version = "5.5.0"
4057 +source = "registry+https://github.com/rust-lang/crates.io-index"
4058 +checksum = "16c7f49c9d5caa3bf4b3106900484b447b9253fe99670ceb81cb6cb5027855e1"
4059 +
4060 +[[package]]
4061 +name = "same-file"
4062 +version = "1.0.6"
4063 +source = "registry+https://github.com/rust-lang/crates.io-index"
4064 +checksum = "93fc1dc3aaa9bfed95e02e6eadabb4baf7e3078b0bd1b4d7b6b0b68378900502"
4065 +dependencies = [
4066 + "winapi-util",
4067 +]
4068 +
4069 +[[package]]
4070 +name = "scc"
4071 +version = "2.4.0"
4072 +source = "registry+https://github.com/rust-lang/crates.io-index"
4073 +checksum = "46e6f046b7fef48e2660c57ed794263155d713de679057f2d0c169bfc6e756cc"
4074 +dependencies = [
4075 + "sdd 3.0.10",
4076 +]
4077 +
4078 +[[package]]
4079 +name = "scc"
4080 +version = "3.6.9"
4081 +source = "registry+https://github.com/rust-lang/crates.io-index"
4082 +checksum = "45bb5ce9efd4a6e7b0f86c2697fe4c1d78d1f4e6d988c54b752d577cafe22fe8"
4083 +dependencies = [
4084 + "saa",
4085 + "sdd 4.7.3",
4086 +]
4087 +
4088 +[[package]]
4089 +name = "schannel"
4090 +version = "0.1.29"
4091 +source = "registry+https://github.com/rust-lang/crates.io-index"
4092 +checksum = "91c1b7e4904c873ef0710c1f407dde2e6287de2bebc1bbbf7d430bb7cbffd939"
4093 +dependencies = [
4094 + "windows-sys 0.61.2",
4095 +]
4096 +
4097 +[[package]]
4098 +name = "schemars"
4099 +version = "0.8.22"
4100 +source = "registry+https://github.com/rust-lang/crates.io-index"
4101 +checksum = "3fbf2ae1b8bc8e02df939598064d22402220cd5bbcca1c76f7d6a310974d5615"
4102 +dependencies = [
4103 + "dyn-clone",
4104 + "schemars_derive",
4105 + "serde",
4106 + "serde_json",
4107 +]
4108 +
4109 +[[package]]
4110 +name = "schemars"
4111 +version = "0.9.0"
4112 +source = "registry+https://github.com/rust-lang/crates.io-index"
4113 +checksum = "4cd191f9397d57d581cddd31014772520aa448f65ef991055d7f61582c65165f"
4114 +dependencies = [
4115 + "dyn-clone",
4116 + "ref-cast",
4117 + "serde",
4118 + "serde_json",
4119 +]
4120 +
4121 +[[package]]
4122 +name = "schemars"
4123 +version = "1.2.1"
4124 +source = "registry+https://github.com/rust-lang/crates.io-index"
4125 +checksum = "a2b42f36aa1cd011945615b92222f6bf73c599a102a300334cd7f8dbeec726cc"
4126 +dependencies = [
4127 + "dyn-clone",
4128 + "ref-cast",
4129 + "serde",
4130 + "serde_json",
4131 +]
4132 +
4133 +[[package]]
4134 +name = "schemars_derive"
4135 +version = "0.8.22"
4136 +source = "registry+https://github.com/rust-lang/crates.io-index"
4137 +checksum = "32e265784ad618884abaea0600a9adf15393368d840e0222d101a072f3f7534d"
4138 +dependencies = [
4139 + "proc-macro2",
4140 + "quote",
4141 + "serde_derive_internals",
4142 + "syn 2.0.117",
4143 +]
4144 +
4145 +[[package]]
4146 +name = "scopeguard"
4147 +version = "1.2.0"
4148 +source = "registry+https://github.com/rust-lang/crates.io-index"
4149 +checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49"
4150 +
4151 +[[package]]
4152 +name = "sdd"
4153 +version = "3.0.10"
4154 +source = "registry+https://github.com/rust-lang/crates.io-index"
4155 +checksum = "490dcfcbfef26be6800d11870ff2df8774fa6e86d047e3e8c8a76b25655e41ca"
4156 +
4157 +[[package]]
4158 +name = "sdd"
4159 +version = "4.7.3"
4160 +source = "registry+https://github.com/rust-lang/crates.io-index"
4161 +checksum = "b21a75f5913ab130e4b369fb8693be25f29b983e2ecad4279df9bfa5dd8aaf3e"
4162 +
4163 +[[package]]
4164 +name = "secrecy"
4165 +version = "0.10.3"
4166 +source = "registry+https://github.com/rust-lang/crates.io-index"
4167 +checksum = "e891af845473308773346dc847b2c23ee78fe442e0472ac50e22a18a93d3ae5a"
4168 +dependencies = [
4169 + "serde",
4170 + "zeroize",
4171 +]
4172 +
4173 +[[package]]
4174 +name = "security-framework"
4175 +version = "3.7.0"
4176 +source = "registry+https://github.com/rust-lang/crates.io-index"
4177 +checksum = "b7f4bc775c73d9a02cde8bf7b2ec4c9d12743edf609006c7facc23998404cd1d"
4178 +dependencies = [
4179 + "bitflags",
4180 + "core-foundation 0.10.1",
4181 + "core-foundation-sys",
4182 + "libc",
4183 + "security-framework-sys",
4184 +]
4185 +
4186 +[[package]]
4187 +name = "security-framework-sys"
4188 +version = "2.17.0"
4189 +source = "registry+https://github.com/rust-lang/crates.io-index"
4190 +checksum = "6ce2691df843ecc5d231c0b14ece2acc3efb62c0a398c7e1d875f3983ce020e3"
4191 +dependencies = [
4192 + "core-foundation-sys",
4193 + "libc",
4194 +]
4195 +
4196 +[[package]]
4197 +name = "semver"
4198 +version = "1.0.27"
4199 +source = "registry+https://github.com/rust-lang/crates.io-index"
4200 +checksum = "d767eb0aabc880b29956c35734170f26ed551a859dbd361d140cdbeca61ab1e2"
4201 +
4202 +[[package]]
4203 +name = "serde"
4204 +version = "1.0.228"
4205 +source = "registry+https://github.com/rust-lang/crates.io-index"
4206 +checksum = "9a8e94ea7f378bd32cbbd37198a4a91436180c5bb472411e48b5ec2e2124ae9e"
4207 +dependencies = [
4208 + "serde_core",
4209 + "serde_derive",
4210 +]
4211 +
4212 +[[package]]
4213 +name = "serde-json-fmt"
4214 +version = "0.1.0"
4215 +source = "registry+https://github.com/rust-lang/crates.io-index"
4216 +checksum = "a4a33b7a5f52a26d520099339add40c48baf2e5ada194c8cc1b18cafa2b5e419"
4217 +dependencies = [
4218 + "serde",
4219 + "serde_json",
4220 + "smartstring",
4221 +]
4222 +
4223 +[[package]]
4224 +name = "serde_bytes"
4225 +version = "0.11.19"
4226 +source = "registry+https://github.com/rust-lang/crates.io-index"
4227 +checksum = "a5d440709e79d88e51ac01c4b72fc6cb7314017bb7da9eeff678aa94c10e3ea8"
4228 +dependencies = [
4229 + "serde",
4230 + "serde_core",
4231 +]
4232 +
4233 +[[package]]
4234 +name = "serde_core"
4235 +version = "1.0.228"
4236 +source = "registry+https://github.com/rust-lang/crates.io-index"
4237 +checksum = "41d385c7d4ca58e59fc732af25c3983b67ac852c1a25000afe1175de458b67ad"
4238 +dependencies = [
4239 + "serde_derive",
4240 +]
4241 +
4242 +[[package]]
4243 +name = "serde_default"
4244 +version = "0.2.0"
4245 +source = "registry+https://github.com/rust-lang/crates.io-index"
4246 +checksum = "486b028b311aaaea83e0ba65a3e6e3cbef381e74e9d0bd6263faefd1fb503c1d"
4247 +dependencies = [
4248 + "darling 0.20.11",
4249 + "proc-macro2",
4250 + "quote",
4251 + "syn 2.0.117",
4252 +]
4253 +
4254 +[[package]]
4255 +name = "serde_derive"
4256 +version = "1.0.228"
4257 +source = "registry+https://github.com/rust-lang/crates.io-index"
4258 +checksum = "d540f220d3187173da220f885ab66608367b6574e925011a9353e4badda91d79"
4259 +dependencies = [
4260 + "proc-macro2",
4261 + "quote",
4262 + "syn 2.0.117",
4263 +]
4264 +
4265 +[[package]]
4266 +name = "serde_derive_internals"
4267 +version = "0.29.1"
4268 +source = "registry+https://github.com/rust-lang/crates.io-index"
4269 +checksum = "18d26a20a969b9e3fdf2fc2d9f21eda6c40e2de84c9408bb5d3b05d499aae711"
4270 +dependencies = [
4271 + "proc-macro2",
4272 + "quote",
4273 + "syn 2.0.117",
4274 +]
4275 +
4276 +[[package]]
4277 +name = "serde_json"
4278 +version = "1.0.149"
4279 +source = "registry+https://github.com/rust-lang/crates.io-index"
4280 +checksum = "83fc039473c5595ace860d8c4fafa220ff474b3fc6bfdb4293327f1a37e94d86"
4281 +dependencies = [
4282 + "indexmap 2.13.0",
4283 + "itoa",
4284 + "memchr",
4285 + "serde",
4286 + "serde_core",
4287 + "zmij",
4288 +]
4289 +
4290 +[[package]]
4291 +name = "serde_path_to_error"
4292 +version = "0.1.20"
4293 +source = "registry+https://github.com/rust-lang/crates.io-index"
4294 +checksum = "10a9ff822e371bb5403e391ecd83e182e0e77ba7f6fe0160b795797109d1b457"
4295 +dependencies = [
4296 + "itoa",
4297 + "serde",
4298 + "serde_core",
4299 +]
4300 +
4301 +[[package]]
4302 +name = "serde_repr"
4303 +version = "0.1.20"
4304 +source = "registry+https://github.com/rust-lang/crates.io-index"
4305 +checksum = "175ee3e80ae9982737ca543e96133087cbd9a485eecc3bc4de9c1a37b47ea59c"
4306 +dependencies = [
4307 + "proc-macro2",
4308 + "quote",
4309 + "syn 2.0.117",
4310 +]
4311 +
4312 +[[package]]
4313 +name = "serde_tuple"
4314 +version = "1.1.3"
4315 +source = "registry+https://github.com/rust-lang/crates.io-index"
4316 +checksum = "6af196b9c06f0aa5555ab980c01a2527b0f67517da8d68b1731b9d4764846a6f"
4317 +dependencies = [
4318 + "serde",
4319 + "serde_tuple_macros",
4320 +]
4321 +
4322 +[[package]]
4323 +name = "serde_tuple_macros"
4324 +version = "1.1.3"
4325 +source = "registry+https://github.com/rust-lang/crates.io-index"
4326 +checksum = "ec3a1e7d2eadec84deabd46ae061bf480a91a6bce74d25dad375bd656f2e19d8"
4327 +dependencies = [
4328 + "proc-macro2",
4329 + "quote",
4330 + "syn 2.0.117",
4331 +]
4332 +
4333 +[[package]]
4334 +name = "serde_urlencoded"
4335 +version = "0.7.1"
4336 +source = "registry+https://github.com/rust-lang/crates.io-index"
4337 +checksum = "d3491c14715ca2294c4d6a88f15e84739788c1d030eed8c110436aafdaa2f3fd"
4338 +dependencies = [
4339 + "form_urlencoded",
4340 + "itoa",
4341 + "ryu",
4342 + "serde",
4343 +]
4344 +
4345 +[[package]]
4346 +name = "serde_with"
4347 +version = "3.18.0"
4348 +source = "registry+https://github.com/rust-lang/crates.io-index"
4349 +checksum = "dd5414fad8e6907dbdd5bc441a50ae8d6e26151a03b1de04d89a5576de61d01f"
4350 +dependencies = [
4351 + "base64 0.22.1",
4352 + "chrono",
4353 + "hex",
4354 + "indexmap 1.9.3",
4355 + "indexmap 2.13.0",
4356 + "schemars 0.9.0",
4357 + "schemars 1.2.1",
4358 + "serde_core",
4359 + "serde_json",
4360 + "serde_with_macros",
4361 + "time",
4362 +]
4363 +
4364 +[[package]]
4365 +name = "serde_with_macros"
4366 +version = "3.18.0"
4367 +source = "registry+https://github.com/rust-lang/crates.io-index"
4368 +checksum = "d3db8978e608f1fe7357e211969fd9abdcae80bac1ba7a3369bb7eb6b404eb65"
4369 +dependencies = [
4370 + "darling 0.23.0",
4371 + "proc-macro2",
4372 + "quote",
4373 + "syn 2.0.117",
4374 +]
4375 +
4376 +[[package]]
4377 +name = "serial_test"
4378 +version = "3.4.0"
4379 +source = "registry+https://github.com/rust-lang/crates.io-index"
4380 +checksum = "911bd979bf1070a3f3aa7b691a3b3e9968f339ceeec89e08c280a8a22207a32f"
4381 +dependencies = [
4382 + "fslock",
4383 + "futures-executor",
4384 + "futures-util",
4385 + "log",
4386 + "once_cell",
4387 + "parking_lot",
4388 + "scc 2.4.0",
4389 + "serial_test_derive",
4390 +]
4391 +
4392 +[[package]]
4393 +name = "serial_test_derive"
4394 +version = "3.4.0"
4395 +source = "registry+https://github.com/rust-lang/crates.io-index"
4396 +checksum = "0a7d91949b85b0d2fb687445e448b40d322b6b3e4af6b44a29b21d9a5f33e6d9"
4397 +dependencies = [
4398 + "proc-macro2",
4399 + "quote",
4400 + "syn 2.0.117",
4401 +]
4402 +
4403 +[[package]]
4404 +name = "sha1"
4405 +version = "0.10.6"
4406 +source = "registry+https://github.com/rust-lang/crates.io-index"
4407 +checksum = "e3bf829a2d51ab4a5ddf1352d8470c140cadc8301b2ae1789db023f01cedd6ba"
4408 +dependencies = [
4409 + "cfg-if",
4410 + "cpufeatures 0.2.17",
4411 + "digest",
4412 +]
4413 +
4414 +[[package]]
4415 +name = "sha2"
4416 +version = "0.10.9"
4417 +source = "registry+https://github.com/rust-lang/crates.io-index"
4418 +checksum = "a7507d819769d01a365ab707794a4084392c824f54a7a6a7862f8c3d0892b283"
4419 +dependencies = [
4420 + "cfg-if",
4421 + "cpufeatures 0.2.17",
4422 + "digest",
4423 +]
4424 +
4425 +[[package]]
4426 +name = "sharded-slab"
4427 +version = "0.1.7"
4428 +source = "registry+https://github.com/rust-lang/crates.io-index"
4429 +checksum = "f40ca3c46823713e0d4209592e8d6e826aa57e928f09752619fc696c499637f6"
4430 +dependencies = [
4431 + "lazy_static",
4432 +]
4433 +
4434 +[[package]]
4435 +name = "shlex"
4436 +version = "1.3.0"
4437 +source = "registry+https://github.com/rust-lang/crates.io-index"
4438 +checksum = "0fda2ff0d084019ba4d7c6f371c95d8fd75ce3524c3cb8fb653a3023f6323e64"
4439 +
4440 +[[package]]
4441 +name = "signal-hook-registry"
4442 +version = "1.4.8"
4443 +source = "registry+https://github.com/rust-lang/crates.io-index"
4444 +checksum = "c4db69cba1110affc0e9f7bcd48bbf87b3f4fc7c61fc9155afd4c469eb3d6c1b"
4445 +dependencies = [
4446 + "errno",
4447 + "libc",
4448 +]
4449 +
4450 +[[package]]
4451 +name = "simd-adler32"
4452 +version = "0.3.8"
4453 +source = "registry+https://github.com/rust-lang/crates.io-index"
4454 +checksum = "e320a6c5ad31d271ad523dcf3ad13e2767ad8b1cb8f047f75a8aeaf8da139da2"
4455 +
4456 +[[package]]
4457 +name = "simd_helpers"
4458 +version = "0.1.0"
4459 +source = "registry+https://github.com/rust-lang/crates.io-index"
4460 +checksum = "95890f873bec569a0362c235787f3aca6e1e887302ba4840839bcc6459c42da6"
4461 +dependencies = [
4462 + "quote",
4463 +]
4464 +
4465 +[[package]]
4466 +name = "siphasher"
4467 +version = "1.0.2"
4468 +source = "registry+https://github.com/rust-lang/crates.io-index"
4469 +checksum = "b2aa850e253778c88a04c3d7323b043aeda9d3e30d5971937c1855769763678e"
4470 +
4471 +[[package]]
4472 +name = "slab"
4473 +version = "0.4.12"
4474 +source = "registry+https://github.com/rust-lang/crates.io-index"
4475 +checksum = "0c790de23124f9ab44544d7ac05d60440adc586479ce501c1d6d7da3cd8c9cf5"
4476 +
4477 +[[package]]
4478 +name = "smallvec"
4479 +version = "1.15.1"
4480 +source = "registry+https://github.com/rust-lang/crates.io-index"
4481 +checksum = "67b1b7a3b5fe4f1376887184045fcf45c69e92af734b7aaddc05fb777b6fbd03"
4482 +
4483 +[[package]]
4484 +name = "smartstring"
4485 +version = "1.0.1"
4486 +source = "registry+https://github.com/rust-lang/crates.io-index"
4487 +checksum = "3fb72c633efbaa2dd666986505016c32c3044395ceaf881518399d2f4127ee29"
4488 +dependencies = [
4489 + "autocfg",
4490 + "static_assertions",
4491 + "version_check",
4492 +]
4493 +
4494 +[[package]]
4495 +name = "socket2"
4496 +version = "0.6.3"
4497 +source = "registry+https://github.com/rust-lang/crates.io-index"
4498 +checksum = "3a766e1110788c36f4fa1c2b71b387a7815aa65f88ce0229841826633d93723e"
4499 +dependencies = [
4500 + "libc",
4501 + "windows-sys 0.61.2",
4502 +]
4503 +
4504 +[[package]]
4505 +name = "socks"
4506 +version = "0.3.4"
4507 +source = "registry+https://github.com/rust-lang/crates.io-index"
4508 +checksum = "f0c3dbbd9ae980613c6dd8e28a9407b50509d3803b57624d5dfe8315218cd58b"
4509 +dependencies = [
4510 + "byteorder",
4511 + "libc",
4512 + "winapi",
4513 +]
4514 +
4515 +[[package]]
4516 +name = "spm_precompiled"
4517 +version = "0.1.4"
4518 +source = "registry+https://github.com/rust-lang/crates.io-index"
4519 +checksum = "5851699c4033c63636f7ea4cf7b7c1f1bf06d0cc03cfb42e711de5a5c46cf326"
4520 +dependencies = [
4521 + "base64 0.13.1",
4522 + "nom 7.1.3",
4523 + "serde",
4524 + "unicode-segmentation",
4525 +]
4526 +
4527 +[[package]]
4528 +name = "stable_deref_trait"
4529 +version = "1.2.1"
4530 +source = "registry+https://github.com/rust-lang/crates.io-index"
4531 +checksum = "6ce2be8dc25455e1f91df71bfa12ad37d7af1092ae736f3a6cd0e37bc7810596"
4532 +
4533 +[[package]]
4534 +name = "static_assertions"
4535 +version = "1.1.0"
4536 +source = "registry+https://github.com/rust-lang/crates.io-index"
4537 +checksum = "a2eb9349b6444b326872e140eb1cf5e7c522154d69e7a0ffb0fb81c06b37543f"
4538 +
4539 +[[package]]
4540 +name = "strength_reduce"
4541 +version = "0.2.4"
4542 +source = "registry+https://github.com/rust-lang/crates.io-index"
4543 +checksum = "fe895eb47f22e2ddd4dabc02bce419d2e643c8e3b585c78158b349195bc24d82"
4544 +
4545 +[[package]]
4546 +name = "strsim"
4547 +version = "0.11.1"
4548 +source = "registry+https://github.com/rust-lang/crates.io-index"
4549 +checksum = "7da8b5736845d9f2fcb837ea5d9e2628564b3b043a70948a3f0b778838c5fb4f"
4550 +
4551 +[[package]]
4552 +name = "strum"
4553 +version = "0.27.2"
4554 +source = "registry+https://github.com/rust-lang/crates.io-index"
4555 +checksum = "af23d6f6c1a224baef9d3f61e287d2761385a5b88fdab4eb4c6f11aeb54c4bcf"
4556 +dependencies = [
4557 + "strum_macros",
4558 +]
4559 +
4560 +[[package]]
4561 +name = "strum_macros"
4562 +version = "0.27.2"
4563 +source = "registry+https://github.com/rust-lang/crates.io-index"
4564 +checksum = "7695ce3845ea4b33927c055a39dc438a45b059f7c1b3d91d38d10355fb8cbca7"
4565 +dependencies = [
4566 + "heck",
4567 + "proc-macro2",
4568 + "quote",
4569 + "syn 2.0.117",
4570 +]
4571 +
4572 +[[package]]
4573 +name = "subenum"
4574 +version = "1.1.3"
4575 +source = "registry+https://github.com/rust-lang/crates.io-index"
4576 +checksum = "ec3d08fe7078c57309d5c3d938e50eba95ba1d33b9c3a101a8465fc6861a5416"
4577 +dependencies = [
4578 + "heck",
4579 + "proc-macro2",
4580 + "quote",
4581 + "syn 2.0.117",
4582 +]
4583 +
4584 +[[package]]
4585 +name = "subtle"
4586 +version = "2.6.1"
4587 +source = "registry+https://github.com/rust-lang/crates.io-index"
4588 +checksum = "13c2bddecc57b384dee18652358fb23172facb8a2c51ccc10d74c157bdea3292"
4589 +
4590 +[[package]]
4591 +name = "syn"
4592 +version = "1.0.109"
4593 +source = "registry+https://github.com/rust-lang/crates.io-index"
4594 +checksum = "72b64191b275b66ffe2469e8af2c1cfe3bafa67b529ead792a6d0160888b4237"
4595 +dependencies = [
4596 + "proc-macro2",
4597 + "quote",
4598 + "unicode-ident",
4599 +]
4600 +
4601 +[[package]]
4602 +name = "syn"
4603 +version = "2.0.117"
4604 +source = "registry+https://github.com/rust-lang/crates.io-index"
4605 +checksum = "e665b8803e7b1d2a727f4023456bbbbe74da67099c585258af0ad9c5013b9b99"
4606 +dependencies = [
4607 + "proc-macro2",
4608 + "quote",
4609 + "unicode-ident",
4610 +]
4611 +
4612 +[[package]]
4613 +name = "sync_wrapper"
4614 +version = "1.0.2"
4615 +source = "registry+https://github.com/rust-lang/crates.io-index"
4616 +checksum = "0bf256ce5efdfa370213c1dabab5935a12e49f2c58d15e9eac2870d3b4f27263"
4617 +dependencies = [
4618 + "futures-core",
4619 +]
4620 +
4621 +[[package]]
4622 +name = "synstructure"
4623 +version = "0.13.2"
4624 +source = "registry+https://github.com/rust-lang/crates.io-index"
4625 +checksum = "728a70f3dbaf5bab7f0c4b1ac8d7ae5ea60a4b5549c8a5914361c99147a709d2"
4626 +dependencies = [
4627 + "proc-macro2",
4628 + "quote",
4629 + "syn 2.0.117",
4630 +]
4631 +
4632 +[[package]]
4633 +name = "system-configuration"
4634 +version = "0.7.0"
4635 +source = "registry+https://github.com/rust-lang/crates.io-index"
4636 +checksum = "a13f3d0daba03132c0aa9767f98351b3488edc2c100cda2d2ec2b04f3d8d3c8b"
4637 +dependencies = [
4638 + "bitflags",
4639 + "core-foundation 0.9.4",
4640 + "system-configuration-sys",
4641 +]
4642 +
4643 +[[package]]
4644 +name = "system-configuration-sys"
4645 +version = "0.6.0"
4646 +source = "registry+https://github.com/rust-lang/crates.io-index"
4647 +checksum = "8e1d1b10ced5ca923a1fcb8d03e96b8d3268065d724548c0211415ff6ac6bac4"
4648 +dependencies = [
4649 + "core-foundation-sys",
4650 + "libc",
4651 +]
4652 +
4653 +[[package]]
4654 +name = "task-local"
4655 +version = "0.1.1"
4656 +source = "registry+https://github.com/rust-lang/crates.io-index"
4657 +checksum = "2972044a9e5e448a506a7ff6f0d03b566d8ef4cd6918a58fc59835a0f8666626"
4658 +dependencies = [
4659 + "pin-project-lite",
4660 +]
4661 +
4662 +[[package]]
4663 +name = "tekken-rs"
4664 +version = "0.1.1"
4665 +source = "registry+https://github.com/rust-lang/crates.io-index"
4666 +checksum = "49623843103837f53f7ebe8cfafc19ccff28ff0e15e7c4b9f6ad21e36fbfde3a"
4667 +dependencies = [
4668 + "anyhow",
4669 + "base64 0.22.1",
4670 + "env_logger",
4671 + "hound",
4672 + "log",
4673 + "ndarray 0.16.1",
4674 + "regex",
4675 + "rubato",
4676 + "rustc-hash 1.1.0",
4677 + "rustfft",
4678 + "serde",
4679 + "serde_json",
4680 + "thiserror 2.0.18",
4681 + "tiktoken-rs 0.7.0",
4682 +]
4683 +
4684 +[[package]]
4685 +name = "tempfile"
4686 +version = "3.27.0"
4687 +source = "registry+https://github.com/rust-lang/crates.io-index"
4688 +checksum = "32497e9a4c7b38532efcdebeef879707aa9f794296a4f0244f6f69e9bc8574bd"
4689 +dependencies = [
4690 + "fastrand",
4691 + "getrandom 0.4.2",
4692 + "once_cell",
4693 + "rustix",
4694 + "windows-sys 0.61.2",
4695 +]
4696 +
4697 +[[package]]
4698 +name = "thiserror"
4699 +version = "1.0.69"
4700 +source = "registry+https://github.com/rust-lang/crates.io-index"
4701 +checksum = "b6aaf5339b578ea85b50e080feb250a3e8ae8cfcdff9a461c9ec2904bc923f52"
4702 +dependencies = [
4703 + "thiserror-impl 1.0.69",
4704 +]
4705 +
4706 +[[package]]
4707 +name = "thiserror"
4708 +version = "2.0.18"
4709 +source = "registry+https://github.com/rust-lang/crates.io-index"
4710 +checksum = "4288b5bcbc7920c07a1149a35cf9590a2aa808e0bc1eafaade0b80947865fbc4"
4711 +dependencies = [
4712 + "thiserror-impl 2.0.18",
4713 +]
4714 +
4715 +[[package]]
4716 +name = "thiserror-ext"
4717 +version = "0.3.0"
4718 +source = "registry+https://github.com/rust-lang/crates.io-index"
4719 +checksum = "5fb7e61141f4141832ca9aad63c3c90023843f944a1975460abdacc64d03f534"
4720 +dependencies = [
4721 + "thiserror 2.0.18",
4722 + "thiserror-ext-derive",
4723 +]
4724 +
4725 +[[package]]
4726 +name = "thiserror-ext-derive"
4727 +version = "0.3.0"
4728 +source = "registry+https://github.com/rust-lang/crates.io-index"
4729 +checksum = "2b5042dd3b562d1d57711be902006a0003fa2781b81d5b2bec07416be31586ff"
4730 +dependencies = [
4731 + "either",
4732 + "proc-macro2",
4733 + "quote",
4734 + "syn 2.0.117",
4735 +]
4736 +
4737 +[[package]]
4738 +name = "thiserror-impl"
4739 +version = "1.0.69"
4740 +source = "registry+https://github.com/rust-lang/crates.io-index"
4741 +checksum = "4fee6c4efc90059e10f81e6d42c60a18f76588c3d74cb83a0b242a2b6c7504c1"
4742 +dependencies = [
4743 + "proc-macro2",
4744 + "quote",
4745 + "syn 2.0.117",
4746 +]
4747 +
4748 +[[package]]
4749 +name = "thiserror-impl"
4750 +version = "2.0.18"
4751 +source = "registry+https://github.com/rust-lang/crates.io-index"
4752 +checksum = "ebc4ee7f67670e9b64d05fa4253e753e016c6c95ff35b89b7941d6b856dec1d5"
4753 +dependencies = [
4754 + "proc-macro2",
4755 + "quote",
4756 + "syn 2.0.117",
4757 +]
4758 +
4759 +[[package]]
4760 +name = "thread_local"
4761 +version = "1.1.9"
4762 +source = "registry+https://github.com/rust-lang/crates.io-index"
4763 +checksum = "f60246a4944f24f6e018aa17cdeffb7818b76356965d03b07d6a9886e8962185"
4764 +dependencies = [
4765 + "cfg-if",
4766 +]
4767 +
4768 +[[package]]
4769 +name = "tiff"
4770 +version = "0.11.3"
4771 +source = "registry+https://github.com/rust-lang/crates.io-index"
4772 +checksum = "b63feaf3343d35b6ca4d50483f94843803b0f51634937cc2ec519fc32232bc52"
4773 +dependencies = [
4774 + "fax",
4775 + "flate2",
4776 + "half",
4777 + "quick-error",
4778 + "weezl",
4779 + "zune-jpeg",
4780 +]
4781 +
4782 +[[package]]
4783 +name = "tiktoken-rs"
4784 +version = "0.7.0"
4785 +source = "registry+https://github.com/rust-lang/crates.io-index"
4786 +checksum = "25563eeba904d770acf527e8b370fe9a5547bacd20ff84a0b6c3bc41288e5625"
4787 +dependencies = [
4788 + "anyhow",
4789 + "base64 0.22.1",
4790 + "bstr",
4791 + "fancy-regex 0.13.0",
4792 + "lazy_static",
4793 + "regex",
4794 + "rustc-hash 1.1.0",
4795 +]
4796 +
4797 +[[package]]
4798 +name = "tiktoken-rs"
4799 +version = "0.9.1"
4800 +source = "registry+https://github.com/rust-lang/crates.io-index"
4801 +checksum = "3a19830747d9034cd9da43a60eaa8e552dfda7712424aebf187b7a60126bae0d"
4802 +dependencies = [
4803 + "anyhow",
4804 + "base64 0.22.1",
4805 + "bstr",
4806 + "fancy-regex 0.13.0",
4807 + "lazy_static",
4808 + "regex",
4809 + "rustc-hash 1.1.0",
4810 +]
4811 +
4812 +[[package]]
4813 +name = "time"
4814 +version = "0.3.47"
4815 +source = "registry+https://github.com/rust-lang/crates.io-index"
4816 +checksum = "743bd48c283afc0388f9b8827b976905fb217ad9e647fae3a379a9283c4def2c"
4817 +dependencies = [
4818 + "deranged",
4819 + "itoa",
4820 + "libc",
4821 + "num-conv",
4822 + "num_threads",
4823 + "powerfmt",
4824 + "serde_core",
4825 + "time-core",
4826 + "time-macros",
4827 +]
4828 +
4829 +[[package]]
4830 +name = "time-core"
4831 +version = "0.1.8"
4832 +source = "registry+https://github.com/rust-lang/crates.io-index"
4833 +checksum = "7694e1cfe791f8d31026952abf09c69ca6f6fa4e1a1229e18988f06a04a12dca"
4834 +
4835 +[[package]]
4836 +name = "time-macros"
4837 +version = "0.2.27"
4838 +source = "registry+https://github.com/rust-lang/crates.io-index"
4839 +checksum = "2e70e4c5a0e0a8a4823ad65dfe1a6930e4f4d756dcd9dd7939022b5e8c501215"
4840 +dependencies = [
4841 + "num-conv",
4842 + "time-core",
4843 +]
4844 +
4845 +[[package]]
4846 +name = "tiny-keccak"
4847 +version = "2.0.2"
4848 +source = "registry+https://github.com/rust-lang/crates.io-index"
4849 +checksum = "2c9d3793400a45f954c52e73d068316d76b6f4e36977e3fcebb13a2721e80237"
4850 +dependencies = [
4851 + "crunchy",
4852 +]
4853 +
4854 +[[package]]
4855 +name = "tinystr"
4856 +version = "0.8.2"
4857 +source = "registry+https://github.com/rust-lang/crates.io-index"
4858 +checksum = "42d3e9c45c09de15d06dd8acf5f4e0e399e85927b7f00711024eb7ae10fa4869"
4859 +dependencies = [
4860 + "displaydoc",
4861 + "zerovec",
4862 +]
4863 +
4864 +[[package]]
4865 +name = "tinytemplate"
4866 +version = "1.2.1"
4867 +source = "registry+https://github.com/rust-lang/crates.io-index"
4868 +checksum = "be4d6b5f19ff7664e8c98d03e2139cb510db9b0a60b55f8e8709b689d939b6bc"
4869 +dependencies = [
4870 + "serde",
4871 + "serde_json",
4872 +]
4873 +
4874 +[[package]]
4875 +name = "tinyvec"
4876 +version = "1.11.0"
4877 +source = "registry+https://github.com/rust-lang/crates.io-index"
4878 +checksum = "3e61e67053d25a4e82c844e8424039d9745781b3fc4f32b8d55ed50f5f667ef3"
4879 +dependencies = [
4880 + "tinyvec_macros",
4881 +]
4882 +
4883 +[[package]]
4884 +name = "tinyvec_macros"
4885 +version = "0.1.1"
4886 +source = "registry+https://github.com/rust-lang/crates.io-index"
4887 +checksum = "1f3ccbac311fea05f86f61904b462b55fb3df8837a366dfc601a0161d0532f20"
4888 +
4889 +[[package]]
4890 +name = "tokenizers"
4891 +version = "0.22.2"
4892 +source = "registry+https://github.com/rust-lang/crates.io-index"
4893 +checksum = "b238e22d44a15349529690fb07bd645cf58149a1b1e44d6cb5bd1641ff1a6223"
4894 +dependencies = [
4895 + "ahash",
4896 + "aho-corasick",
4897 + "compact_str",
4898 + "dary_heap",
4899 + "derive_builder",
4900 + "esaxx-rs",
4901 + "getrandom 0.3.4",
4902 + "indicatif 0.18.4",
4903 + "itertools 0.14.0",
4904 + "log",
4905 + "macro_rules_attribute",
4906 + "monostate",
4907 + "onig",
4908 + "paste",
4909 + "rand 0.9.2",
4910 + "rayon",
4911 + "rayon-cond",
4912 + "regex",
4913 + "regex-syntax",
4914 + "serde",
4915 + "serde_json",
4916 + "spm_precompiled",
4917 + "thiserror 2.0.18",
4918 + "unicode-normalization-alignments",
4919 + "unicode-segmentation",
4920 + "unicode_categories",
4921 +]
4922 +
4923 +[[package]]
4924 +name = "tokio"
4925 +version = "1.50.0"
4926 +source = "registry+https://github.com/rust-lang/crates.io-index"
4927 +checksum = "27ad5e34374e03cfffefc301becb44e9dc3c17584f414349ebe29ed26661822d"
4928 +dependencies = [
4929 + "bytes",
4930 + "libc",
4931 + "mio",
4932 + "parking_lot",
4933 + "pin-project-lite",
4934 + "signal-hook-registry",
4935 + "socket2",
4936 + "tokio-macros",
4937 + "windows-sys 0.61.2",
4938 +]
4939 +
4940 +[[package]]
4941 +name = "tokio-macros"
4942 +version = "2.6.1"
4943 +source = "registry+https://github.com/rust-lang/crates.io-index"
4944 +checksum = "5c55a2eff8b69ce66c84f85e1da1c233edc36ceb85a2058d11b0d6a3c7e7569c"
4945 +dependencies = [
4946 + "proc-macro2",
4947 + "quote",
4948 + "syn 2.0.117",
4949 +]
4950 +
4951 +[[package]]
4952 +name = "tokio-native-tls"
4953 +version = "0.3.1"
4954 +source = "registry+https://github.com/rust-lang/crates.io-index"
4955 +checksum = "bbae76ab933c85776efabc971569dd6119c580d8f5d448769dec1764bf796ef2"
4956 +dependencies = [
4957 + "native-tls",
4958 + "tokio",
4959 +]
4960 +
4961 +[[package]]
4962 +name = "tokio-rustls"
4963 +version = "0.26.4"
4964 +source = "registry+https://github.com/rust-lang/crates.io-index"
4965 +checksum = "1729aa945f29d91ba541258c8df89027d5792d85a8841fb65e8bf0f4ede4ef61"
4966 +dependencies = [
4967 + "rustls",
4968 + "tokio",
4969 +]
4970 +
4971 +[[package]]
4972 +name = "tokio-stream"
4973 +version = "0.1.18"
4974 +source = "registry+https://github.com/rust-lang/crates.io-index"
4975 +checksum = "32da49809aab5c3bc678af03902d4ccddea2a87d028d86392a4b1560c6906c70"
4976 +dependencies = [
4977 + "futures-core",
4978 + "pin-project-lite",
4979 + "tokio",
4980 +]
4981 +
4982 +[[package]]
4983 +name = "tokio-tungstenite"
4984 +version = "0.28.0"
4985 +source = "registry+https://github.com/rust-lang/crates.io-index"
4986 +checksum = "d25a406cddcc431a75d3d9afc6a7c0f7428d4891dd973e4d54c56b46127bf857"
4987 +dependencies = [
4988 + "futures-util",
4989 + "log",
4990 + "tokio",
4991 + "tungstenite",
4992 +]
4993 +
4994 +[[package]]
4995 +name = "tokio-util"
4996 +version = "0.7.18"
4997 +source = "registry+https://github.com/rust-lang/crates.io-index"
4998 +checksum = "9ae9cec805b01e8fc3fd2fe289f89149a9b66dd16786abd8b19cfa7b48cb0098"
4999 +dependencies = [

This file is too large to show in full.

rust/Cargo.toml new
+129
@@ -0,0 +1,129 @@
1 +[workspace]
2 +members = [
3 + "src/chat",
4 + "src/cmd",
5 + "src/engine-core-client",
6 + "src/llm",
7 + "src/managed-engine",
8 + "src/metrics",
9 + "src/reasoning-parser",
10 + "src/server",
11 + "src/text",
12 + "src/tokenizer",
13 + "src/tool-parser",
14 +]
15 +resolver = "3"
16 +
17 +[workspace.package]
18 +version = "0.1.0"
19 +edition = "2024"
20 +license = "Apache-2.0"
21 +
22 +[workspace.dependencies]
23 +anyhow = "1.0.100"
24 +arc-swap = "1.9.0"
25 +async-openai = "0.33.1"
26 +async-trait = "0.1.89"
27 +asynk-strim-attr = "0.1.0"
28 +axum = "0.8.8"
29 +base64 = "0.22.1"
30 +bytemuck = { version = "1.25.0", features = ["extern_crate_alloc"] }
31 +byteorder = "1.5.0"
32 +bytes = "1.11.1"
33 +clap = { version = "4.5.38", features = ["derive", "env"] }
34 +criterion = "0.5.1"
35 +easy-ext = "1.0.3"
36 +educe = "0.6.0"
37 +enum-as-inner = "0.7.0"
38 +expect-test = "1.5.1"
39 +fastokens = "0.2.0"
40 +futures = "0.3.31"
41 +half = { version = "2.7.1", features = ["bytemuck"] }
42 +hex = "0.4.3"
43 +hf-hub = { version = "0.5.0", features = ["tokio"] }
44 +http-body = "1.0.1"
45 +itertools = "0.14.0"
46 +libc = "0.2.177"
47 +llm-multimodal = { git = "https://github.com/vllm-project/llm-multimodal", rev = "5b558989844d1c7af3e43d0f604069ffd9c06320" }
48 +minijinja = { version = "2.0", features = ["unstable_machinery", "json", "builtins", "loader", "loop_controls"] }
49 +minijinja-contrib = { version = "2.0", features = ["pycompat"] }
50 +native-tls-vendored = { package = "native-tls", version = "0.2.18", features = ["vendored"] }
51 +ndarray = { version = "0.16.1", features = ["serde"] }
52 +openai-harmony = "0.0.8"
53 +openai-protocol = "1.6.0"
54 +parking_lot = "0.12.5"
55 +prometheus-client = "0.24.0"
56 +prometheus-client-derive-encode = "0.5.0"
57 +prost = "0.14.3"
58 +prost-types = "0.14.3"
59 +reasoning-parser = "1.2.2"
60 +reqwest = { version = "0.12.8", default-features = false, features = ["rustls-tls"] }
61 +riptoken = { version = "0.3.0", default-features = false }
62 +rmp-serde = "1.3.1"
63 +rmpv = { version = "1.3.1", features = ["with-serde"] }
64 +rustc-hash = "1.1.0"
65 +serde = { version = "1.0.228", features = ["derive"] }
66 +serde-json-fmt = "0.1.0"
67 +serde_default = "0.2.0"
68 +serde_json = "1.0.145"
69 +serde_repr = "0.1.20"
70 +serde_tuple = "1.1.3"
71 +serde_with = "3.18.0"
72 +serial_test = "3.2.0"
73 +socket2 = "0.6.3"
74 +subenum = "1.1.3"
75 +task-local = "0.1.1"
76 +tekken = { package = "tekken-rs", version = "0.1.1", default-features = false }
77 +tempfile = "3.23.0"
78 +thiserror = "2.0.16"
79 +thiserror-ext = "0.3.0"
80 +tiktoken-rs = "0.9.1"
81 +time = { version = "0.3.47", features = ["formatting", "local-offset", "macros"] }
82 +tokenizers = "0.22.0"
83 +tokio = { version = "1.47.1", features = [
84 + "macros",
85 + "net",
86 + "rt-multi-thread",
87 + "sync",
88 + "time",
89 +] }
90 +tokio-stream = "0.1"
91 +tokio-util = { version = "0.7.18", features = ["rt"] }
92 +tonic = "0.14.5"
93 +tonic-build = "0.14.5"
94 +tonic-prost = "0.14.5"
95 +tonic-prost-build = "0.14.5"
96 +tool-parser = "1.2.0"
97 +tower = { version = "0.5.3", features = ["util"] }
98 +tower-http = { version = "0.6.8", features = ["trace"] }
99 +tracing = { version = "0.1.44", features = ["release_max_level_debug"] }
100 +tracing-futures = { version = "0.2.5", features = ["futures-03"] }
101 +tracing-subscriber = { version = "0.3.20", features = ["env-filter", "fmt"] }
102 +trait-set = "0.3.0"
103 +uuid = { version = "1.22.0", features = ["v4"] }
104 +validator = { version = "0.20.0", features = ["derive"] }
105 +vllm-chat = { path = "src/chat" }
106 +vllm-engine-core-client = { path = "src/engine-core-client" }
107 +vllm-llm = { path = "src/llm" }
108 +vllm-managed-engine = { path = "src/managed-engine" }
109 +vllm-metrics = { path = "src/metrics" }
110 +vllm-reasoning-parser = { path = "src/reasoning-parser" }
111 +vllm-server = { path = "src/server" }
112 +vllm-text = { path = "src/text" }
113 +vllm-tokenizer = { path = "src/tokenizer" }
114 +vllm-tool-parser = { path = "src/tool-parser" }
115 +winnow = "1.0.2"
116 +zeromq = { version = "0.6.0", default-features = false, features = [
117 + "tokio-runtime",
118 + "all-transport",
119 +] }
120 +
121 +[workspace.lints.clippy]
122 +too_many_arguments = "allow"
123 +
124 +[profile.dev]
125 +panic = "abort"
126 +
127 +[profile.release]
128 +lto = "thin"
129 +panic = "abort"
rust/README.md new
+89
@@ -0,0 +1,89 @@
1 +# vllm-frontend-rs
2 +
3 +This is a Rust drop-in alternative frontend for vLLM. The current goal is to rebuild the northbound serving layer in Rust while still talking to the core Python vLLM engine process(es) via ZMQ over the existing engine boundary.
4 +
5 +It should still be considered experimental, and is not feature-complete. We are working to add more functionality from the python front-end.
6 +
7 +See <https://github.com/Inferact/vllm-frontend-rs> for the original commit history before it was moved into the main vllm repo.
8 +
9 +## Architecture
10 +
11 +The component is organized as a Cargo workspace with several crates, layered bottom-up:
12 +
13 +```text
14 +┌─────────────────────────────────┐
15 +│ vllm-cmd / vllm-rs │ CLI entrypoint:
16 +│ │ Python vLLM frontend subprocess
17 +│ │ Rust managed-engine serve mode
18 +├─────────────────────────────────┤
19 +│ vllm-server │ OpenAI-compatible HTTP API (axum)
20 +├─────────────────────────────────┤
21 +│ vllm-chat │ Chat completions: template rendering,
22 +│ │ structured assistant events,
23 +│ │ reasoning & tool parsing
24 +├─────────────────────────────────┤
25 +│ vllm-text │ Tokenizer & incremental detokenizer
26 +├─────────────────────────────────┤
27 +│ vllm-llm │ Thin token-in/token-out facade over
28 +│ │ the engine client
29 +├─────────────────────────────────┤
30 +│ vllm-engine-core-client │ ZMQ transport + MessagePack protocol
31 +│ │ for the headless vLLM engine
32 +└─────────────────────────────────┘
33 +```
34 +
35 +`vllm-rs` integrates into Python `vllm` as a Rust frontend subprocess.
36 +Python owns process startup and launches the Rust API server as a Python-supervised worker, while
37 +passing the inherited listening socket and transport addresses into `vllm-rs`.
38 +
39 +For example:
40 +
41 +```bash
42 +VLLM_USE_RUST_FRONTEND=1 vllm serve Qwen/Qwen3-0.6B
43 +```
44 +
45 +### External Engine
46 +
47 +`vllm-rs serve` can be run standalone with `--data-parallel-size-local 0` when the Python engines
48 +are started elsewhere and this node should run only the Rust frontend. The frontend still uses
49 +the global `--data-parallel-size` to determine how many engines it expects to join the shared handshake.
50 +
51 +```bash
52 +vllm serve Qwen/Qwen3-0.6B \
53 + --headless \
54 + --data-parallel-address 127.0.0.1 \
55 + --data-parallel-rpc-port 62100 \
56 + --data-parallel-size 1 \
57 + --data-parallel-size-local 1
58 +```
59 +
60 +Then start the Rust frontend-only server:
61 +
62 +```bash
63 +vllm-rs serve Qwen/Qwen3-0.6B \
64 + --data-parallel-address 127.0.0.1 \
65 + --data-parallel-rpc-port 62100 \
66 + --data-parallel-size 1 \
67 + --data-parallel-size-local 0
68 +```
69 +
70 +To build the `vllm-rs` in isolation:
71 +
72 +```bash
73 +# from the local checkout
74 +cargo install --path src/cmd --bin vllm-rs
75 +```
76 +
77 +### Example Request
78 +
79 +After either startup path, you can use any OpenAI-compatible client:
80 +
81 +```bash
82 +curl http://127.0.0.1:8000/v1/chat/completions \
83 + -H "Content-Type: application/json" \
84 + -d '{
85 + "model": "Qwen/Qwen3-0.6B",
86 + "messages": [{"role": "user", "content": "What is the capital of France?"}],
87 + "stream": true
88 + }'
89 +```
rust/proto/vllm_grpc.proto new
+196
@@ -0,0 +1,196 @@
1 +syntax = "proto3";
2 +package vllm;
3 +
4 +import "google/protobuf/struct.proto";
5 +
6 +
7 +service Generate {
8 + // Generates text given a prompt
9 + rpc Generate (GenerateRequest) returns (GenerateResponse) {}
10 + // Generates text given a prompt, streaming the outputs
11 + rpc GenerateStream (GenerateRequest) returns (stream GenerateResponse) {}
12 +}
13 +
14 +// ======================================================================================
15 +// Generate Request
16 +// ======================================================================================
17 +
18 +message GenerateRequest {
19 + string request_id = 1;
20 + string model = 2;
21 +
22 + oneof prompt {
23 + string text = 3;
24 + TokenIds token_ids = 4;
25 + }
26 +
27 + // Temperature, defaults to model-specific default or 0
28 + optional float temperature = 5;
29 + // Parameters controlling random sampling, not applicable if temperature == 0
30 + RandomSampling sampling = 6;
31 + // Parameters for conditionally penalizing/boosting
32 + // candidate tokens during decoding
33 + DecodingParameters decoding = 7;
34 + // Parameters controlling when generation should stop
35 + StoppingCriteria stopping = 8;
36 + // Flags to control what is returned in the response
37 + ResponseOptions response = 9;
38 + // Parameters controlling KV cache/distribution
39 + KVCacheParameters kv = 10;
40 +
41 + // Truncate prompt tokens; default (0) means no truncation
42 + uint32 truncate_prompt_tokens = 11;
43 +
44 + int32 priority = 12;
45 +}
46 +
47 +message RandomSampling {
48 + uint32 num_sequences = 1; // "n", default (0) means 1
49 + uint32 top_k = 2; // 0 means default
50 + float top_p = 3; // 0 means default
51 + float min_p = 4; // 0 means default
52 + optional int64 seed = 5;
53 +}
54 +
55 +message DecodingParameters {
56 + // Penalties
57 + float presence_penalty = 1; // Default (0.0) means no penalty
58 + float frequency_penalty = 2; // Default (0.0) means no penalty
59 + float repetition_penalty = 3; // Default (0.0) means no penalty
60 + map<uint32, float> logit_bias = 4;
61 + repeated uint32 allowed_token_ids = 5;
62 +
63 + message StringChoices {
64 + repeated string choices = 1;
65 + }
66 +
67 + // Control structured outputs
68 + oneof structured_output {
69 + string json = 6;
70 + string regex = 7;
71 + StringChoices choice = 8;
72 + string grammar = 9;
73 + bool json_object = 10;
74 + string structural_tag = 11;
75 + }
76 +}
77 +
78 +message StoppingCriteria {
79 + // Default (0) is currently 20
80 + uint32 max_new_tokens = 1;
81 + // Default (0) means no minimum
82 + uint32 min_new_tokens = 2;
83 +
84 + repeated uint32 stop_token_ids = 3;
85 + repeated string stop_strings = 4;
86 + bool include_stop_strings = 5;
87 +
88 + bool ignore_eos = 6;
89 +}
90 +
91 +message ResponseOptions {
92 + // Prompt options
93 + bool prompt_token_ids = 1;
94 + bool prompt_logprobs = 2;
95 + optional CandidateTokens prompt_candidates = 3;
96 +
97 + // Output options; output_text defaults to true
98 + optional bool output_text = 4;
99 + bool output_token_ids = 5;
100 + bool output_logprobs = 6;
101 + optional CandidateTokens output_candidates = 7;
102 +}
103 +
104 +message KVCacheParameters {
105 + bool bypass_prefix_cache = 1;
106 + string cache_salt = 2;
107 +
108 + // KV Connector transfer parameters
109 + google.protobuf.Struct kv_transfer_params = 3;
110 +}
111 +
112 +// Controls which extra candidate tokens at each position should be returned
113 +message CandidateTokens {
114 + oneof select {
115 + uint32 top_n = 1;
116 + TokenIds token_ids = 2;
117 + bool all = 3;
118 + }
119 +}
120 +
121 +// ======================================================================================
122 +// Generate Response
123 +// ======================================================================================
124 +
125 +message GenerateResponse {
126 + // Only present in first response
127 + optional PromptInfo prompt_info = 1;
128 + SequenceOutput outputs = 2;
129 +}
130 +
131 +message SequenceOutput {
132 + // Index of output sequence for num_sequences > 1.
133 + uint32 index = 1;
134 +
135 + string text = 2;
136 + uint32 num_tokens = 3; // Number of tokens in this chunk
137 + repeated uint32 token_ids = 4; // If requested
138 + repeated float logprobs = 5; // If requested
139 + repeated uint32 ranks = 6; // If logprobs were requested
140 + repeated CandidateTokenInfo candidate_tokens = 7; // If requested
141 +
142 + // Only present in final output for this sequence
143 + optional FinishInfo finish_info = 8;
144 +}
145 +
146 +// Prompt info, returned in the first response
147 +message PromptInfo {
148 + uint32 num_prompt_tokens = 1;
149 + repeated uint32 token_ids = 2; // If requested
150 + repeated float logprobs = 3; // If requested
151 + repeated uint32 ranks = 4; // If logprobs were requested
152 + repeated CandidateTokenInfo candidate_tokens = 5;
153 +}
154 +
155 +// Finish info, returned in the final response
156 +message FinishInfo {
157 + uint32 num_output_tokens = 1;
158 +
159 + enum FinishReason {
160 + NOT_FINISHED = 0; // Possibly more tokens to be streamed
161 + LENGTH = 1; // Finished due to length constraint
162 + STOP = 2; // Stop string/token or EOS encountered
163 + ABORTED = 3; // Request aborted/cancelled
164 + }
165 +
166 + FinishReason finish_reason = 2;
167 + // One of these will be set when finish_reason == STOP
168 + oneof stop_reason {
169 + uint32 stop_token_id = 3;
170 + uint32 eos_token_id = 4;
171 + string stop_string = 5;
172 + }
173 +
174 + google.protobuf.Struct kv_transfer_params = 6;
175 + //uint64 seed = 7;
176 +}
177 +
178 +// Info for candidate tokens other than the input/sampled
179 +// token at a given position
180 +message CandidateTokenInfo {
181 + message TokenInfo {
182 + uint32 id = 1;
183 + float logprob = 2;
184 + uint32 rank = 3;
185 + // string text = 4;
186 + // bytes token_bytes = 5;
187 + }
188 + // Candidate token infos at this position
189 + repeated TokenInfo tokens = 1;
190 +}
191 +
192 +// Token ids used for prompt
193 +message TokenIds {
194 + repeated uint32 ids = 1;
195 +}
196 +
rust/rustfmt.toml new
+3
@@ -0,0 +1,3 @@
1 +style_edition = "2024"
2 +chain_width = 80
3 +use_field_init_shorthand = true
rust/rustfmt.unstable.toml new
+20
@@ -0,0 +1,20 @@
1 +# Optional local formatting profile. CI and pre-commit use rustfmt.toml.
2 +# Apply manually with:
3 +# cargo +nightly fmt -- --config-path rustfmt.unstable.toml
4 +
5 +style_edition = "2024"
6 +chain_width = 80
7 +comment_width = 100
8 +use_field_init_shorthand = true
9 +
10 +# Unstable features go here.
11 +unstable_features = true
12 +
13 +format_code_in_doc_comments = true
14 +format_macro_matchers = true
15 +normalize_comments = true
16 +normalize_doc_attributes = true
17 +imports_granularity = "Module"
18 +group_imports = "StdExternalCrate"
19 +reorder_impl_items = true
20 +wrap_comments = true
rust/src/chat/Cargo.toml new
+52
@@ -0,0 +1,52 @@
1 +[package]
2 +name = "vllm-chat"
3 +version.workspace = true
4 +edition.workspace = true
5 +license.workspace = true
6 +
7 +[dependencies]
8 +anyhow.workspace = true
9 +asynk-strim-attr.workspace = true
10 +easy-ext.workspace = true
11 +futures.workspace = true
12 +half.workspace = true
13 +itertools.workspace = true
14 +llm-multimodal.workspace = true
15 +minijinja.workspace = true
16 +minijinja-contrib.workspace = true
17 +openai-harmony.workspace = true
18 +reqwest.workspace = true
19 +serde.workspace = true
20 +serde-json-fmt.workspace = true
21 +serde_json.workspace = true
22 +serde_with.workspace = true
23 +subenum.workspace = true
24 +thiserror.workspace = true
25 +thiserror-ext.workspace = true
26 +tokio.workspace = true
27 +tracing.workspace = true
28 +trait-set.workspace = true
29 +uuid.workspace = true
30 +vllm-engine-core-client.workspace = true
31 +vllm-llm.workspace = true
32 +vllm-reasoning-parser.workspace = true
33 +vllm-text.workspace = true
34 +vllm-tokenizer.workspace = true
35 +vllm-tool-parser.workspace = true
36 +
37 +[dev-dependencies]
38 +anyhow.workspace = true
39 +bytes.workspace = true
40 +clap.workspace = true
41 +expect-test.workspace = true
42 +rmp-serde.workspace = true
43 +serial_test = { workspace = true, features = ["file_locks"] }
44 +tempfile.workspace = true
45 +tokio.workspace = true
46 +tracing-subscriber.workspace = true
47 +uuid.workspace = true
48 +vllm-engine-core-client = { workspace = true, features = ["test-util"] }
49 +zeromq.workspace = true
50 +
51 +[lints]
52 +workspace = true
rust/src/chat/examples/README.md new
+39
@@ -0,0 +1,39 @@
1 +# Chat Smoke Test
2 +
3 +Start a fresh headless `vllm` engine:
4 +
5 +```bash
6 +source ../vllm/.venv/bin/activate
7 +HF_HUB_OFFLINE=1 \
8 +VLLM_LOGGING_LEVEL=DEBUG \
9 +VLLM_CPU_KVCACHE_SPACE=2 \
10 +VLLM_HOST_IP=127.0.0.1 \
11 +VLLM_LOOPBACK_IP=127.0.0.1 \
12 +python3 -m vllm.entrypoints.cli.main serve Qwen/Qwen3-0.6B \
13 + --headless \
14 + --data-parallel-address 127.0.0.1 \
15 + --data-parallel-rpc-port 62100 \
16 + --data-parallel-size-local 1 \
17 + --max-model-len 512 \
18 + --dtype float16
19 +```
20 +
21 +Run the Rust chat smoke test through the `vllm-chat` interface:
22 +
23 +```bash
24 +cargo run -p vllm-chat --example external_engine_chat_qwen -- \
25 + --handshake-address tcp://127.0.0.1:62100 \
26 + --host 127.0.0.1 \
27 + --prompt 'What is the capital of France? Answer with one word.'
28 +```
29 +
30 +The example now defaults to `Qwen/Qwen3-0.6B`. The current `vllm-chat`
31 +request model stays text-first and supports either plain string content or
32 +OpenAI-style text blocks, while the output side now emits structured assistant
33 +events and automatically separates reasoning blocks for supported models. Tool
34 +use and multimodal inputs are still out of scope. It uses the Rust
35 +`tokenizers` library for the tokenizer itself, plus standard Hugging Face
36 +config files to load the chat template and EOS metadata.
37 +
38 +IMPORTANT: Restart `vllm` each time you run the smoke test. The current headless
39 +engine cannot safely handle frontend reconnects after the client shuts down.
rust/src/chat/examples/external_engine_chat_qwen.rs new
+178
@@ -0,0 +1,178 @@
1 +use std::time::Duration;
2 +
3 +use anyhow::{Context, Result, bail};
4 +use clap::Parser;
5 +use futures::StreamExt as _;
6 +use tracing_subscriber::EnvFilter;
7 +use vllm_chat::{
8 + AssistantBlockKind, AssistantMessageExt as _, ChatEvent, ChatLlm, ChatMessage, ChatRequest,
9 + ChatRole, SamplingParams, load_model_backends,
10 +};
11 +use vllm_engine_core_client::{EngineCoreClient, EngineCoreClientConfig, TransportMode};
12 +use vllm_llm::Llm;
13 +use vllm_text::TextLlm;
14 +
15 +#[derive(Debug, Parser)]
16 +#[command(about = "Smoke-test the Rust chat facade against an external Qwen vLLM engine.")]
17 +struct Args {
18 + #[arg(long)]
19 + handshake_address: String,
20 + #[arg(long, default_value_t = 1)]
21 + engine_count: usize,
22 + #[arg(long, default_value = "Qwen/Qwen3-0.6B")]
23 + model: String,
24 + #[arg(long, default_value = "127.0.0.1")]
25 + host: String,
26 + #[arg(long, default_value_t = 30)]
27 + ready_timeout_secs: u64,
28 + #[arg(long)]
29 + prompt: String,
30 +}
31 +
32 +const CLIENT_INDEX: u32 = 0;
33 +const OUTPUT_TIMEOUT_SECS: u64 = 120;
34 +
35 +fn unique_request_id() -> String {
36 + format!("rust-chat-smoke-{}", uuid::Uuid::new_v4())
37 +}
38 +
39 +fn init_tracing() {
40 + let filter = EnvFilter::try_from_default_env()
41 + .unwrap_or_else(|_| EnvFilter::new("vllm_engine_core_client=debug"));
42 + let _ = tracing_subscriber::fmt().with_env_filter(filter).try_init();
43 +}
44 +
45 +#[tokio::main(flavor = "multi_thread")]
46 +async fn main() -> Result<()> {
47 + init_tracing();
48 + let args = Args::parse();
49 + let loaded = load_model_backends(&args.model, Default::default())
50 + .await
51 + .with_context(|| format!("failed to load backends for {}", args.model))?;
52 + let text_backend = loaded.text_backend;
53 + let chat_backend = loaded.chat_backend;
54 +
55 + let ready_timeout = Duration::from_secs(args.ready_timeout_secs);
56 + let output_timeout = Duration::from_secs(OUTPUT_TIMEOUT_SECS);
57 + let request_id = unique_request_id();
58 + let client = EngineCoreClient::connect(EngineCoreClientConfig {
59 + transport_mode: TransportMode::HandshakeOwner {
60 + handshake_address: args.handshake_address.clone(),
61 + advertised_host: args.host.clone(),
62 + engine_count: args.engine_count,
63 + ready_timeout,
64 + local_input_address: None,
65 + local_output_address: None,
66 + },
67 + coordinator_mode: None,
68 + model_name: args.model.clone(),
69 + client_index: CLIENT_INDEX,
70 + })
71 + .await
72 + .context("failed to connect to external vLLM engine")?;
73 +
74 + println!("model={}", args.model);
75 + println!("tokenizer_source=tokenizers + hf-hub");
76 + println!("chat_template_source=tokenizer_config.json or adjacent chat template file");
77 + println!("handshake_address={}", args.handshake_address);
78 + println!("engine_count={}", args.engine_count);
79 + println!("input_address={}", client.input_address());
80 + println!("output_address={}", client.output_address());
81 + println!("engine_identities={:x?}", client.engine_identities());
82 +
83 + let llm = Llm::new(client);
84 + let chat = ChatLlm::new(TextLlm::new(llm, text_backend), chat_backend);
85 +
86 + let request = ChatRequest {
87 + messages: vec![ChatMessage::text(ChatRole::User, args.prompt.clone())],
88 + sampling_params: SamplingParams {
89 + temperature: Some(0.0),
90 + ..Default::default()
91 + },
92 + request_id: request_id.clone(),
93 + ..ChatRequest::for_test()
94 + };
95 +
96 + println!("request_id={request_id}");
97 + println!("prompt={}", args.prompt);
98 +
99 + let mut stream = chat.chat(request).await.context("failed to submit chat request")?;
100 + let output = tokio::time::timeout(output_timeout, async {
101 + let mut final_reasoning = String::new();
102 + let mut final_text = String::new();
103 + let mut final_output_token_count = 0usize;
104 + let mut finish_reason = None;
105 + let mut saw_start = false;
106 + let mut saw_stream_output = false;
107 +
108 + while let Some(event) = stream.next().await.transpose()? {
109 + match event {
110 + ChatEvent::Start { .. } => {
111 + saw_start = true;
112 + }
113 + ChatEvent::BlockStart { kind, .. } => {
114 + if saw_stream_output {
115 + println!();
116 + }
117 + match kind {
118 + AssistantBlockKind::Reasoning => print!("[reasoning] "),
119 + AssistantBlockKind::Text => print!("[answer] "),
120 + AssistantBlockKind::ToolCall => {}
121 + }
122 + saw_stream_output = true;
123 + }
124 + ChatEvent::ToolCallStart { name, .. } => {
125 + if saw_stream_output {
126 + println!();
127 + }
128 + print!("[tool:{name}] ");
129 + saw_stream_output = true;
130 + }
131 + ChatEvent::LogprobsDelta { .. } => {}
132 + ChatEvent::Done {
133 + message,
134 + output_token_count,
135 + finish_reason: reason,
136 + ..
137 + } => {
138 + final_reasoning = message.reasoning().unwrap_or_default();
139 + final_text = message.text();
140 + final_output_token_count = output_token_count;
141 + finish_reason = Some(reason);
142 + break;
143 + }
144 + ChatEvent::BlockDelta { kind, delta, .. } => match kind {
145 + AssistantBlockKind::Reasoning | AssistantBlockKind::Text => {
146 + print!("{delta}");
147 + }
148 + AssistantBlockKind::ToolCall => {}
149 + },
150 + ChatEvent::ToolCallArgumentsDelta { delta, .. } => print!("{delta}"),
151 + ChatEvent::BlockEnd { .. } | ChatEvent::ToolCallEnd { .. } => {}
152 + }
153 + }
154 +
155 + println!();
156 +
157 + if !saw_start {
158 + bail!("chat stream ended without a start event");
159 + }
160 + Ok::<_, anyhow::Error>((
161 + final_reasoning,
162 + final_text,
163 + final_output_token_count,
164 + finish_reason,
165 + ))
166 + })
167 + .await
168 + .context("timed out waiting for chat output")??;
169 +
170 + chat.shutdown().await.context("failed to shut down chat client")?;
171 +
172 + println!("final_reasoning={:?}", output.0);
173 + println!("final_text={:?}", output.1);
174 + println!("final_output_token_count={:?}", output.2);
175 + println!("finish_reason={:?}", output.3);
176 +
177 + Ok(())
178 +}
rust/src/chat/src/backend/hf.rs new
+308
@@ -0,0 +1,308 @@
1 +use std::sync::Arc;
2 +
3 +use tracing::info;
4 +use vllm_text::backend::hf::{HfTextBackend, ResolvedModelFiles, load_model_config};
5 +use vllm_text::tokenizer::DynTokenizer;
6 +use vllm_text::{DynTextBackend, TextBackend as _};
7 +
8 +use crate::backend::{
9 + ChatBackend, DynChatBackend, LoadModelBackendsOptions, LoadedModelBackends,
10 + NewChatOutputProcessorOptions,
11 +};
12 +use crate::error::Result;
13 +use crate::multimodal::MultimodalModelInfo;
14 +use crate::output::{
15 + DefaultChatOutputProcessor, HarmonyChatOutputProcessor, validate_harmony_parser_overrides,
16 +};
17 +use crate::renderer::hf::{HfChatRenderer, MultimodalRenderInfo};
18 +use crate::renderer::{DeepSeekV4ChatRenderer, DeepSeekV32ChatRenderer, DynChatRenderer};
19 +use crate::request::ChatRequest;
20 +use crate::{DynChatOutputProcessor, RendererSelection};
21 +
22 +/// [`ChatBackend`] implementation built on Hugging Face model files.
23 +pub struct HfChatBackend {
24 + model_id: String,
25 + model_type: String,
26 + tokenizer: DynTokenizer,
27 + chat_renderer: DynChatRenderer,
28 + multimodal_model_info: Option<MultimodalModelInfo>,
29 +}
30 +
31 +impl HfChatBackend {
32 + /// Load the chat backend from resolved Hugging Face model files.
33 + pub fn from_resolved_model_files(
34 + files: ResolvedModelFiles,
35 + model_id: String,
36 + options: LoadModelBackendsOptions,
37 + tokenizer: DynTokenizer,
38 + ) -> Result<Self> {
39 + let model_config = load_model_config(files.config_path.as_deref())?;
40 + let model_type = model_config.model_type().unwrap_or_default();
41 + let multimodal_model_info = MultimodalModelInfo::from_paths(
42 + model_id.clone(),
43 + (!model_type.is_empty()).then_some(model_type.to_string()),
44 + files.config_path.as_deref(),
45 + files.preprocessor_config_path.as_deref(),
46 + tokenizer.clone(),
47 + )?;
48 + let multimodal_render_info = resolve_multimodal_render_info(multimodal_model_info.as_ref());
49 +
50 + let renderer = options.renderer.resolve(model_type);
51 + let chat_renderer: DynChatRenderer = match renderer {
52 + RendererSelection::Auto => unreachable!("renderer auto should be resolved above"),
53 + RendererSelection::Hf => Arc::new(HfChatRenderer::load(
54 + &files,
55 + options,
56 + multimodal_render_info,
57 + )?),
58 + RendererSelection::DeepSeekV32 => Arc::new(DeepSeekV32ChatRenderer::new()),
59 + RendererSelection::DeepSeekV4 => Arc::new(DeepSeekV4ChatRenderer::new()),
60 + };
61 +
62 + info!(
63 + model_id,
64 + model_type,
65 + %renderer,
66 + "loaded chat backend with Hugging Face model files"
67 + );
68 +
69 + Ok(Self {
70 + model_id,
71 + model_type: model_type.to_string(),
72 + tokenizer,
73 + chat_renderer,
74 + multimodal_model_info,
75 + })
76 + }
77 +}
78 +
79 +impl ChatBackend for HfChatBackend {
80 + fn chat_renderer(&self) -> DynChatRenderer {
81 + self.chat_renderer.clone()
82 + }
83 +
84 + fn multimodal_model_info(&self) -> Option<&MultimodalModelInfo> {
85 + self.multimodal_model_info.as_ref()
86 + }
87 +
88 + fn new_chat_output_processor(
89 + &self,
90 + request: &mut ChatRequest,
91 + options: NewChatOutputProcessorOptions<'_>,
92 + ) -> Result<DynChatOutputProcessor> {
93 + if self.model_type == "gpt_oss" {
94 + validate_harmony_parser_overrides(options.tool_call_parser, options.reasoning_parser)?;
95 + return Ok(Box::new(HarmonyChatOutputProcessor::new(request)?));
96 + }
97 +
98 + Ok(Box::new(DefaultChatOutputProcessor::new(
99 + request,
100 + &self.model_id,
101 + self.tokenizer.clone(),
102 + options.tool_call_parser,
103 + options.reasoning_parser,
104 + )?))
105 + }
106 +}
107 +
108 +/// Load the Hugging Face text and chat backends for the given model id.
109 +pub(super) async fn load_model_backends(
110 + model_id: &str,
111 + options: LoadModelBackendsOptions,
112 +) -> Result<LoadedModelBackends> {
113 + let files = ResolvedModelFiles::new(model_id).await?;
114 + let text_backend =
115 + HfTextBackend::from_resolved_model_files(files.clone(), model_id.to_string())?;
116 + let tokenizer = text_backend.tokenizer();
117 + let text_backend: DynTextBackend = Arc::new(text_backend);
118 +
119 + let chat_backend: DynChatBackend = Arc::new(HfChatBackend::from_resolved_model_files(
120 + files,
121 + model_id.to_string(),
122 + options,
123 + tokenizer,
124 + )?);
125 +
126 + Ok(LoadedModelBackends {
127 + text_backend,
128 + chat_backend,
129 + })
130 +}
131 +
132 +fn resolve_multimodal_render_info(
133 + info: Option<&MultimodalModelInfo>,
134 +) -> Option<MultimodalRenderInfo> {
135 + info.map(|info| MultimodalRenderInfo {
136 + placeholder_token: info.placeholder_token().to_string(),
137 + })
138 +}
139 +
140 +#[cfg(test)]
141 +mod tests {
142 + use std::collections::HashMap;
143 + use std::path::PathBuf;
144 + use std::sync::Arc;
145 +
146 + use tempfile::tempdir;
147 + use vllm_text::backend::hf::TokenizerSource;
148 + use vllm_text::tokenizer::{DynTokenizer, Tokenizer};
149 +
150 + use super::HfChatBackend;
151 + use crate::RendererSelection;
152 + use crate::backend::{ChatBackend, LoadModelBackendsOptions};
153 + use crate::request::{ChatContent, ChatMessage, ChatRequest};
154 +
155 + fn request_with_user_text(text: &str) -> ChatRequest {
156 + ChatRequest {
157 + request_id: "renderer-selection-test".to_string(),
158 + messages: vec![ChatMessage::User {
159 + content: ChatContent::Text(text.to_string()),
160 + }],
161 + ..ChatRequest::for_test()
162 + }
163 + }
164 +
165 + fn write_json(path: &std::path::Path, content: &str) {
166 + std::fs::write(path, content).unwrap();
167 + }
168 +
169 + fn resolved_files(
170 + config_json: &str,
171 + tokenizer_config_json: &str,
172 + ) -> vllm_text::backend::hf::ResolvedModelFiles {
173 + let dir = tempdir().unwrap();
174 + let root = dir.keep();
175 + let config_path = root.join("config.json");
176 + let tokenizer_config_path = root.join("tokenizer_config.json");
177 + write_json(&config_path, config_json);
178 + write_json(&tokenizer_config_path, tokenizer_config_json);
179 +
180 + vllm_text::backend::hf::ResolvedModelFiles {
181 + tokenizer: TokenizerSource::HuggingFace(PathBuf::from("/tmp/unused-tokenizer.json")),
182 + tokenizer_config_path: Some(tokenizer_config_path),
183 + generation_config_path: None,
184 + preprocessor_config_path: None,
185 + chat_template_path: None,
186 + config_path: Some(config_path),
187 + }
188 + }
189 +
190 + struct TestTokenizer;
191 +
192 + impl Tokenizer for TestTokenizer {
193 + fn encode(
194 + &self,
195 + _text: &str,
196 + _add_special_tokens: bool,
197 + ) -> vllm_text::tokenizer::Result<Vec<u32>> {
198 + Ok(Vec::new())
199 + }
200 +
201 + fn decode(
202 + &self,
203 + _token_ids: &[u32],
204 + _skip_special_tokens: bool,
205 + ) -> vllm_text::tokenizer::Result<String> {
206 + Ok(String::new())
207 + }
208 +
209 + fn token_to_id(&self, _token: &str) -> Option<u32> {
210 + None
211 + }
212 + }
213 +
214 + fn test_tokenizer() -> DynTokenizer {
215 + Arc::new(TestTokenizer)
216 + }
217 +
218 + fn render_prompt(
219 + renderer: RendererSelection,
220 + config_json: &str,
221 + tokenizer_config_json: &str,
222 + ) -> String {
223 + let backend = HfChatBackend::from_resolved_model_files(
224 + resolved_files(config_json, tokenizer_config_json),
225 + "test-model".to_string(),
226 + LoadModelBackendsOptions {
227 + renderer,
228 + chat_template_content_format: Default::default(),
229 + chat_template: None,
230 + default_chat_template_kwargs: HashMap::new(),
231 + },
232 + test_tokenizer(),
233 + )
234 + .unwrap();
235 +
236 + backend
237 + .chat_renderer()
238 + .render(&request_with_user_text("hello"))
239 + .unwrap()
240 + .prompt
241 + .into_text()
242 + .expect("renderer should return text prompt")
243 + }
244 +
245 + #[test]
246 + fn auto_uses_deepseek_renderer_for_deepseek_v32_model_type() {
247 + let prompt = render_prompt(
248 + RendererSelection::Auto,
249 + r#"{"model_type":"deepseek_v32"}"#,
250 + r#"{}"#,
251 + );
252 +
253 + assert_eq!(
254 + prompt,
255 + "<｜begin▁of▁sentence｜><｜User｜>hello<｜Assistant｜></think>"
256 + );
257 + }
258 +
259 + #[test]
260 + fn auto_uses_hf_renderer_for_other_model_types() {
261 + let prompt = render_prompt(
262 + RendererSelection::Auto,
263 + r#"{"model_type":"qwen2"}"#,
264 + r#"{"chat_template":"{{ messages[0].content }}"}"#,
265 + );
266 +
267 + assert_eq!(prompt, "hello");
268 + }
269 +
270 + #[test]
271 + fn explicit_deepseek_renderer_overrides_generic_model_type() {
272 + let prompt = render_prompt(
273 + RendererSelection::DeepSeekV32,
274 + r#"{"model_type":"qwen2"}"#,
275 + r#"{"chat_template":"{{ messages[0].content }}"}"#,
276 + );
277 +
278 + assert_eq!(
279 + prompt,
280 + "<｜begin▁of▁sentence｜><｜User｜>hello<｜Assistant｜></think>"
281 + );
282 + }
283 +
284 + #[test]
285 + fn explicit_hf_renderer_overrides_deepseek_v32_model_type() {
286 + let prompt = render_prompt(
287 + RendererSelection::Hf,
288 + r#"{"model_type":"deepseek_v32"}"#,
289 + r#"{"chat_template":"{{ messages[0].content }}"}"#,
290 + );
291 +
292 + assert_eq!(prompt, "hello");
293 + }
294 +
295 + #[test]
296 + fn auto_uses_nested_text_config_model_type() {
297 + let prompt = render_prompt(
298 + RendererSelection::Auto,
299 + r#"{"text_config":{"model_type":"deepseek_v32","num_attention_heads":32}}"#,
300 + r#"{}"#,
301 + );
302 +
303 + assert_eq!(
304 + prompt,
305 + "<｜begin▁of▁sentence｜><｜User｜>hello<｜Assistant｜></think>"
306 + );
307 + }
308 +}
rust/src/chat/src/backend/mod.rs new
+86
@@ -0,0 +1,86 @@
1 +use std::collections::HashMap;
2 +use std::sync::Arc;
3 +
4 +use serde_json::Value;
5 +use vllm_text::{DynTextBackend, TextBackend};
6 +
7 +use crate::error::Result;
8 +use crate::multimodal::MultimodalModelInfo;
9 +use crate::output::DynChatOutputProcessor;
10 +use crate::renderer::DynChatRenderer;
11 +use crate::request::ChatRequest;
12 +use crate::{ChatTemplateContentFormatOption, ParserSelection, RendererSelection};
13 +
14 +pub mod hf;
15 +
16 +/// Options for creating a new chat output processor.
17 +pub struct NewChatOutputProcessorOptions<'a> {
18 + pub tool_call_parser: &'a ParserSelection,
19 + pub reasoning_parser: &'a ParserSelection,
20 +}
21 +
22 +/// Minimal prompt-processing backend needed by `vllm-chat`.
23 +pub trait ChatBackend: Send + Sync {
24 + /// Return the renderer used for chat-prompt construction.
25 + fn chat_renderer(&self) -> DynChatRenderer;
26 +
27 + /// Return model files/config needed for request-scoped multimodal
28 + /// preprocessing, if supported.
29 + fn multimodal_model_info(&self) -> Option<&MultimodalModelInfo> {
30 + None
31 + }
32 +
33 + /// Create a request-scoped output processor after request-level adjustments
34 + /// are applied.
35 + fn new_chat_output_processor(
36 + &self,
37 + request: &mut ChatRequest,
38 + options: NewChatOutputProcessorOptions<'_>,
39 + ) -> Result<DynChatOutputProcessor>;
40 +}
41 +
42 +/// Shared trait-object form of [`ChatBackend`].
43 +pub type DynChatBackend = Arc<dyn ChatBackend>;
44 +
45 +/// Convenience trait for backends that can serve both raw text generation and
46 +/// chat templating.
47 +///
48 +/// This is mainly useful in tests and small examples, where one mock/backend
49 +/// often implements both sides and callers want `ChatLlm` to wire the shared
50 +/// object into `TextLlm` automatically.
51 +pub trait ChatTextBackend: ChatBackend + TextBackend {}
52 +
53 +impl<T> ChatTextBackend for T where T: ChatBackend + TextBackend + ?Sized {}
54 +
55 +/// Shared trait-object form of [`ChatTextBackend`].
56 +pub type DynChatTextBackend = Arc<dyn ChatTextBackend>;
57 +
58 +/// Frontend-side chat backend loading options.
59 +#[derive(Debug, Clone, PartialEq, Eq, Default)]
60 +pub struct LoadModelBackendsOptions {
61 + /// Which chat renderer implementation to use.
62 + pub renderer: RendererSelection,
63 + /// How to serialize `message.content` when rendering the chat template.
64 + pub chat_template_content_format: ChatTemplateContentFormatOption,
65 + /// Optional server-default chat template override, provided either as an
66 + /// inline template or as a path to a template file.
67 + pub chat_template: Option<String>,
68 + /// Optional server-default keyword arguments merged into every
69 + /// chat-template render before request-level `chat_template_kwargs`.
70 + pub default_chat_template_kwargs: HashMap<String, Value>,
71 +}
72 +
73 +/// Shared backends loaded from a model id.
74 +pub struct LoadedModelBackends {
75 + pub text_backend: DynTextBackend,
76 + pub chat_backend: DynChatBackend,
77 +}
78 +
79 +/// Load text and chat backends for the given model id.
80 +pub async fn load_model_backends(
81 + model_id: &str,
82 + options: LoadModelBackendsOptions,
83 +) -> Result<LoadedModelBackends> {
84 + // Currently, we only have HuggingFace backends.
85 + hf::load_model_backends(model_id, options).await
86 +}
rust/src/chat/src/error.rs new
+80
@@ -0,0 +1,80 @@
1 +use thiserror::Error;
2 +use thiserror_ext::Macro;
3 +
4 +type BoxedError = Box<dyn std::error::Error + Send + Sync>;
5 +
6 +#[derive(Debug, Error, Macro)]
7 +#[thiserror_ext(macro(path = "crate::error"))]
8 +pub enum Error {
9 + #[error("chat request must contain at least one message")]
10 + EmptyMessages,
11 + #[error("cannot continue the final message when the last message is not from the assistant")]
12 + ContinueFinalAssistantWithoutFinalAssistant,
13 + #[error("chat template is required but none was configured")]
14 + MissingChatTemplate,
15 + #[error("chat template error: {0}")]
16 + ChatTemplate(String),
17 + #[error("multimodal input is not supported by this chat renderer")]
18 + UnsupportedMultimodalRenderer,
19 + #[error("unsupported multimodal content: {0}")]
20 + UnsupportedMultimodalContent(&'static str),
21 + #[error("multimodal preprocessing error: {0}")]
22 + Multimodal(#[message] String),
23 + #[error("{kind} parsing is not available for model `{model_id}`")]
24 + ParserUnavailableForModel {
25 + kind: &'static str,
26 + model_id: String,
27 + },
28 + #[error("{kind} parsing is disabled by frontend configuration")]
29 + ParserDisabled { kind: &'static str },
30 + #[error(
31 + "{kind} parser `{name}` is not registered{}",
32 + available_parser_hint(.available_names)
33 + )]
34 + ParserUnavailableByName {
35 + kind: &'static str,
36 + name: String,
37 + available_names: Vec<String>,
38 + },
39 + #[error("failed to initialize {kind} parser `{name}`")]
40 + ParserInitialization {
41 + kind: &'static str,
42 + name: String,
43 + #[source]
44 + error: BoxedError,
45 + },
46 + #[error(
47 + "gpt_oss uses native Harmony output parsing; generic {kind} parser override `{selection}` is not supported"
48 + )]
49 + HarmonyParserOverrideUnsupported {
50 + kind: &'static str,
51 + selection: String,
52 + },
53 + #[error("harmony output parsing failed")]
54 + HarmonyOutputParsing {
55 + #[source]
56 + error: BoxedError,
57 + },
58 + #[error(
59 + "this model's maximum context length is {max_model_len} tokens, \
60 + but the prompt contains {prompt_len} input tokens"
61 + )]
62 + PromptTooLong { max_model_len: u32, prompt_len: u32 },
63 + #[error("chat request stream `{request_id}` closed before terminal output")]
64 + StreamClosedBeforeTerminalOutput { request_id: String },
65 + #[error("tool call stream state is inconsistent: {message}")]
66 + ToolCallStreamInvariant { message: String },
67 + #[error(transparent)]
68 + Text(#[from] vllm_text::Error),
69 +}
70 +
71 +pub type Result<T> = std::result::Result<T, Error>;
72 +
73 +/// Format the available-parser suffix used in user-facing error messages.
74 +fn available_parser_hint(available_names: &[String]) -> String {
75 + if available_names.is_empty() {
76 + String::new()
77 + } else {
78 + format!(" (choose from: {})", available_names.join(", "))
79 + }
80 +}
rust/src/chat/src/event.rs new
+183
@@ -0,0 +1,183 @@
1 +use std::ops::Deref;
2 +use std::sync::Arc;
3 +
4 +use serde::{Deserialize, Serialize};
5 +use vllm_text::{DecodedLogprobs, DecodedPromptLogprobs};
6 +
7 +use crate::FinishReason;
8 +
9 +/// One finalized assistant tool call.
10 +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
11 +pub struct AssistantToolCall {
12 + pub id: String,
13 + pub name: String,
14 + pub arguments: String,
15 +}
16 +
17 +/// Semantic kind of one assistant output block.
18 +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
19 +pub enum AssistantBlockKind {
20 + /// Visible final-answer text.
21 + Text,
22 + /// Extracted reasoning content.
23 + Reasoning,
24 + /// One finalized tool call.
25 + ToolCall,
26 +}
27 +
28 +/// One structured assistant output block.
29 +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
30 +pub enum AssistantContentBlock {
31 + /// Visible final-answer text.
32 + Text { text: String },
33 + /// Extracted reasoning content.
34 + Reasoning { text: String },
35 + /// One finalized tool call.
36 + ToolCall(AssistantToolCall),
37 +}
38 +
39 +impl AssistantContentBlock {
40 + /// Return the semantic kind of this block.
41 + pub fn kind(&self) -> AssistantBlockKind {
42 + match self {
43 + Self::Text { .. } => AssistantBlockKind::Text,
44 + Self::Reasoning { .. } => AssistantBlockKind::Reasoning,
45 + Self::ToolCall(..) => AssistantBlockKind::ToolCall,
46 + }
47 + }
48 +
49 + /// Return this block as one finalized tool call, if applicable.
50 + pub fn as_tool_call(&self) -> Option<&AssistantToolCall> {
51 + match self {
52 + Self::ToolCall(call) => Some(call),
53 + _ => None,
54 + }
55 + }
56 +}
57 +
58 +#[easy_ext::ext(AssistantMessageExt)]
59 +impl [AssistantContentBlock] {
60 + /// Concatenate all visible final-answer text blocks.
61 + pub fn text(&self) -> String {
62 + self.iter()
63 + .filter_map(|block| match block {
64 + AssistantContentBlock::Text { text } => Some(text.as_str()),
65 + _ => None,
66 + })
67 + .collect()
68 + }
69 +
70 + /// Concatenate all extracted reasoning blocks, if any.
71 + pub fn reasoning(&self) -> Option<String> {
72 + Some(
73 + self.iter()
74 + .filter_map(|block| match block {
75 + AssistantContentBlock::Reasoning { text } => Some(text.as_str()),
76 + _ => None,
77 + })
78 + .collect(),
79 + )
80 + .filter(|s: &String| !s.is_empty())
81 + }
82 +
83 + /// Return whether this assistant message contains any non-empty reasoning
84 + /// text blocks.
85 + pub fn has_reasoning(&self) -> bool {
86 + self.iter().any(|block| match block {
87 + AssistantContentBlock::Reasoning { text } => !text.is_empty(),
88 + _ => false,
89 + })
90 + }
91 +
92 + /// Return finalized assistant tool calls in encounter order.
93 + pub fn tool_calls(&self) -> impl Iterator<Item = &AssistantToolCall> {
94 + self.iter().filter_map(AssistantContentBlock::as_tool_call)
95 + }
96 +
97 + /// Return whether this assistant message contains any tool-call blocks.
98 + pub fn has_tool_calls(&self) -> bool {
99 + self.iter().any(|block| matches!(block, AssistantContentBlock::ToolCall(_)))
100 + }
101 +}
102 +
103 +/// Final structured assistant message assembled from the event stream.
104 +#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
105 +pub struct AssistantMessage {
106 + pub content: Vec<AssistantContentBlock>,
107 +}
108 +
109 +impl Deref for AssistantMessage {
110 + type Target = [AssistantContentBlock];
111 +
112 + fn deref(&self) -> &Self::Target {
113 + &self.content
114 + }
115 +}
116 +
117 +impl AssistantMessage {
118 + /// Push one new block to the end of the message content.
119 + pub(crate) fn push_block(&mut self, block: AssistantContentBlock) {
120 + self.content.push(block);
121 + }
122 +}
123 +
124 +/// Streamed chat event emitted by [`crate::ChatEventStream`].
125 +#[derive(Debug, Clone, PartialEq)]
126 +pub enum ChatEvent {
127 + /// The request was accepted, streaming has started, and prompt metadata is
128 + /// ready.
129 + Start {
130 + /// The actual prompt token IDs for this request.
131 + prompt_token_ids: Arc<[u32]>,
132 + /// Once-only prompt logprobs metadata, when requested.
133 + prompt_logprobs: Option<DecodedPromptLogprobs>,
134 + },
135 + /// A new assistant output block has started.
136 + BlockStart {
137 + index: usize,
138 + kind: AssistantBlockKind,
139 + },
140 + /// A newly observed delta for one open assistant output block.
141 + BlockDelta {
142 + index: usize,
143 + kind: AssistantBlockKind,
144 + delta: String,
145 + },
146 + /// Per-decoded-update sample metadata: logprobs and/or output token IDs.
147 + LogprobsDelta {
148 + logprobs: Option<DecodedLogprobs>,
149 + token_ids: Vec<u32>,
150 + },
151 + /// One assistant output block has ended.
152 + BlockEnd {
153 + index: usize,
154 + block: AssistantContentBlock,
155 + },
156 + /// One tool call has started.
157 + ToolCallStart {
158 + index: usize,
159 + id: String,
160 + name: String,
161 + },
162 + /// One incremental tool-call arguments delta for the currently open tool
163 + /// call.
164 + ToolCallArgumentsDelta { index: usize, delta: String },
165 + /// One tool call has ended.
166 + ToolCallEnd {
167 + index: usize,
168 + call: AssistantToolCall,
169 + },
170 + /// Terminal event carrying the final assembled assistant message and finish
171 + /// metadata.
172 + Done {
173 + message: AssistantMessage,
174 + /// Number of prompt tokens actually sent to the engine after chat
175 + /// template rendering and tokenization.
176 + prompt_token_count: usize,
177 + /// Number of output tokens generated.
178 + output_token_count: usize,
179 + finish_reason: FinishReason,
180 + /// Connector-specific KV transfer parameters for disaggregated serving.
181 + kv_transfer_params: Option<serde_json::Value>,
182 + },
183 +}
rust/src/chat/src/lib.rs new
+249
@@ -0,0 +1,249 @@
1 +//! Minimal chat facade above [`vllm_text`].
2 +//!
3 +//! This crate keeps the northbound boundary intentionally small:
4 +//! `messages -> rendered prompt -> tokenized prompt -> engine request ->
5 +//! streamed structured assistant events`. The request side remains text-first,
6 +//! while the response side can emit structured reasoning and final-answer
7 +//! blocks. It is closer to vLLM's internal chat-rendering flow than to a full
8 +//! OpenAI-compatible surface.
9 +
10 +pub use backend::hf::HfChatBackend;
11 +pub use backend::{
12 + ChatBackend, ChatTextBackend, DynChatBackend, DynChatTextBackend, LoadModelBackendsOptions,
13 + LoadedModelBackends, NewChatOutputProcessorOptions, load_model_backends,
14 +};
15 +pub use error::{Error, Result};
16 +pub use event::{
17 + AssistantBlockKind, AssistantContentBlock, AssistantMessage, AssistantMessageExt,
18 + AssistantToolCall, ChatEvent,
19 +};
20 +use futures::{StreamExt, TryStreamExt as _};
21 +pub use output::{
22 + ChatOutputProcessor, DefaultChatOutputProcessor, DynChatOutputProcessor,
23 + HarmonyChatOutputProcessor,
24 +};
25 +pub use parser::ParserSelection;
26 +pub use parser::reasoning::{
27 + ReasoningDelta, ReasoningError, ReasoningParser, ReasoningParserFactory,
28 +};
29 +pub use parser::tool::{ToolParser, ToolParserError, ToolParserFactory};
30 +pub use renderer::hf::ChatTemplateContentFormatOption;
31 +pub use renderer::{
32 + ChatRenderer, DeepSeekV4ChatRenderer, DeepSeekV32ChatRenderer, DynChatRenderer, RenderedPrompt,
33 + RendererSelection,
34 +};
35 +pub use request::{
36 + ChatContent, ChatContentPart, ChatMessage, ChatOptions, ChatRequest, ChatRole, ChatTool,
37 + ChatToolChoice, GenerationPromptMode, ReasoningEffort, SamplingParams,
38 +};
39 +pub use stream::{ChatEventStream, ChatEventStreamTrait, CollectedAssistantMessage};
40 +pub use vllm_llm::FinishReason;
41 +
42 +mod backend;
43 +mod error;
44 +mod event;
45 +pub mod multimodal;
46 +mod output;
47 +mod parser;
48 +mod renderer;
49 +mod request;
50 +mod stream;
51 +
52 +use vllm_engine_core_client::EngineCoreClient;
53 +use vllm_engine_core_client::protocol::ModelDtype;
54 +use vllm_llm::Llm;
55 +use vllm_text::{TextLlm, TextRequest};
56 +
57 +/// Validate explicit parser override names without starting request processing.
58 +pub fn validate_parser_overrides(
59 + tool_call_parser: &ParserSelection,
60 + reasoning_parser: &ParserSelection,
61 +) -> Result<()> {
62 + let tool_parser_factory = ToolParserFactory::global();
63 + if let ParserSelection::Explicit(name) = tool_call_parser
64 + && !tool_parser_factory.contains(name)
65 + {
66 + return Err(Error::ParserUnavailableByName {
67 + kind: "tool",
68 + name: name.clone(),
69 + available_names: tool_parser_factory.list(),
70 + });
71 + }
72 +
73 + let reasoning_parser_factory = ReasoningParserFactory::global();
74 + if let ParserSelection::Explicit(name) = reasoning_parser
75 + && !reasoning_parser_factory.contains(name)
76 + {
77 + return Err(Error::ParserUnavailableByName {
78 + kind: "reasoning",
79 + name: name.clone(),
80 + available_names: reasoning_parser_factory.list(),
81 + });
82 + }
83 +
84 + Ok(())
85 +}
86 +
87 +/// Structured chat facade above [`TextLlm`].
88 +///
89 +/// This layer stays above raw text semantics: it takes care of chat-template
90 +/// rendering, exposes structured assistant events, and adds chat-specific
91 +/// request semantics such as tool calls.
92 +pub struct ChatLlm {
93 + text: TextLlm,
94 + backend: DynChatBackend,
95 + /// Effective model dtype reported by the engine.
96 + model_dtype: Option<ModelDtype>,
97 + /// Tool-call parser selection.
98 + tool_call_parser: ParserSelection,
99 + /// Reasoning parser selection.
100 + reasoning_parser: ParserSelection,
101 +}
102 +
103 +impl ChatLlm {
104 + /// Create a new chat facade from a text-generation facade plus a chat
105 + /// backend.
106 + pub fn new(text: TextLlm, backend: DynChatBackend) -> Self {
107 + let model_dtype = text.engine_core_client().model_dtype();
108 +
109 + Self {
110 + text,
111 + backend,
112 + model_dtype,
113 + tool_call_parser: ParserSelection::Auto,
114 + reasoning_parser: ParserSelection::Auto,
115 + }
116 + }
117 +
118 + /// Convenience constructor for one shared backend object that implements
119 + /// both text and chat responsibilities.
120 + pub fn from_shared_backend(llm: Llm, backend: DynChatTextBackend) -> Self {
121 + let text = TextLlm::new(llm, backend.clone());
122 + Self::new(text, backend)
123 + }
124 +
125 + /// Set tool-call parser selection.
126 + pub fn with_tool_call_parser(mut self, selection: ParserSelection) -> Self {
127 + self.tool_call_parser = selection;
128 + self
129 + }
130 +
131 + /// Set reasoning parser selection.
132 + pub fn with_reasoning_parser(mut self, selection: ParserSelection) -> Self {
133 + self.reasoning_parser = selection;
134 + self
135 + }
136 +
137 + /// Override the effective model dtype used for multimodal tensor encoding.
138 + pub fn with_model_dtype(mut self, model_dtype: Option<ModelDtype>) -> Self {
139 + self.model_dtype = model_dtype;
140 + self
141 + }
142 +
143 + /// Expose the underlying text facade for raw text-generation routes such as
144 + /// `/v1/completions`.
145 + pub fn text(&self) -> &TextLlm {
146 + &self.text
147 + }
148 +
149 + /// Return the model ID reported by the underlying text backend.
150 + pub fn model_id(&self) -> &str {
151 + self.text.model_id()
152 + }
153 +
154 + /// Expose the underlying engine-core client for low-level utility/admin
155 + /// calls.
156 + pub fn engine_core_client(&self) -> &EngineCoreClient {
157 + self.text.engine_core_client()
158 + }
159 +
160 + /// Render, tokenize, and submit one chat request.
161 + pub async fn chat(&self, mut request: ChatRequest) -> Result<ChatEventStream> {
162 + request.validate()?;
163 +
164 + let output_processor = self.backend.new_chat_output_processor(
165 + &mut request,
166 + NewChatOutputProcessorOptions {
167 + tool_call_parser: &self.tool_call_parser,
168 + reasoning_parser: &self.reasoning_parser,
169 + },
170 + )?;
171 + let rendered = self.backend.chat_renderer().render(&request)?;
172 +
173 + let (prompt, mm_features) = multimodal::finalize_rendered_prompt(
174 + &request,
175 + rendered,
176 + self.backend.multimodal_model_info(),
177 + self.model_dtype,
178 + )
179 + .await?;
180 +
181 + let text_request = TextRequest {
182 + request_id: request.request_id.clone(),
183 + prompt,
184 + mm_features,
185 + sampling_params: request.sampling_params,
186 + decode_options: request.decode_options,
187 + intermediate: request.intermediate,
188 + priority: request.priority,
189 + cache_salt: request.cache_salt,
190 + add_special_tokens: request.add_special_tokens,
191 + data_parallel_rank: request.data_parallel_rank,
192 + };
193 + let decoded_stream = self.text.generate(text_request).await?.map_err(Error::from).boxed();
194 +
195 + let structured_stream = output_processor.process(decoded_stream)?;
196 +
197 + Ok(ChatEventStream::new(request.request_id, structured_stream))
198 + }
199 +
200 + /// Shut down the underlying LLM client and its background tasks.
201 + pub async fn shutdown(self) -> Result<()> {
202 + self.text.shutdown().await?;
203 + Ok(())
204 + }
205 +}
206 +
207 +#[cfg(test)]
208 +mod tests {
209 + use thiserror_ext::AsReport;
210 +
211 + use super::{ParserSelection, validate_parser_overrides};
212 + use crate::parser::reasoning::names;
213 +
214 + #[test]
215 + fn validate_parser_overrides_accepts_registered_names() {
216 + validate_parser_overrides(
217 + &ParserSelection::Explicit("llama3_json".to_string()),
218 + &ParserSelection::Explicit(names::QWEN3.to_string()),
219 + )
220 + .unwrap();
221 + }
222 +
223 + #[test]
224 + fn validate_parser_overrides_accepts_auto_and_none() {
225 + validate_parser_overrides(&ParserSelection::Auto, &ParserSelection::None).unwrap();
226 + }
227 +
228 + #[test]
229 + fn validate_parser_overrides_rejects_unknown_tool_parser() {
230 + let error = validate_parser_overrides(
231 + &ParserSelection::Explicit("definitely_missing_tool_parser".to_string()),
232 + &ParserSelection::Auto,
233 + )
234 + .unwrap_err();
235 +
236 + expect_test::expect!["tool parser `definitely_missing_tool_parser` is not registered (choose from: deepseek_v3, deepseek_v31, deepseek_v32, deepseek_v4, gemma4, glm45, glm47, hermes, kimi_k2, llama3_json, llama4_json, minimax_m2, mistral, qwen3_coder, qwen3_xml)"].assert_eq(&error.to_report_string());
237 + }
238 +
239 + #[test]
240 + fn validate_parser_overrides_rejects_unknown_reasoning_parser() {
241 + let error = validate_parser_overrides(
242 + &ParserSelection::Auto,
243 + &ParserSelection::Explicit("definitely_missing_reasoning_parser".to_string()),
244 + )
245 + .unwrap_err();
246 +
247 + expect_test::expect!["reasoning parser `definitely_missing_reasoning_parser` is not registered (choose from: cohere_cmd, deepseek_r1, deepseek_v3, deepseek_v4, gemma4, glm45, kimi, kimi_k2, minimax_m2, nemotron_v3, qwen3, step3)"].assert_eq(&error.to_report_string());
248 + }
249 +}
rust/src/chat/src/multimodal.rs new
+775
@@ -0,0 +1,775 @@
1 +//! Chat-layer multimodal image preparation.
2 +//!
3 +//! This module owns the narrow image-only multimodal path for chat requests:
4 +//! it extracts image parts from structured chat messages, fetches and
5 +//! preprocesses them through `llm-multimodal`, expands rendered prompt
6 +//! placeholders after tokenization, and builds the engine-facing
7 +//! `MmFeatures` payload.
8 +//!
9 +//! Raw media stays above `vllm-text`; this module lowers it into token IDs and
10 +//! opaque tensor payloads before the request is handed to text generation.
11 +
12 +use std::collections::{HashMap, HashSet};
13 +use std::fs;
14 +use std::path::Path;
15 +use std::sync::{Arc, LazyLock, Once};
16 +
17 +use itertools::izip;
18 +use llm_multimodal::{
19 + AsyncMultiModalTracker, FieldLayout, ImagePreProcessor, ImageProcessorRegistry, MediaConnector,
20 + MediaConnectorConfig, MediaContentPart, Modality, ModelMetadata, ModelProcessorSpec,
21 + ModelRegistry, PreProcessorConfig, PreprocessedImages, PromptReplacement, TokenResolver,
22 + TrackedMedia,
23 +};
24 +use tracing::warn;
25 +use vllm_engine_core_client::protocol::ModelDtype;
26 +use vllm_engine_core_client::protocol::multimodal::{
27 + MmBatchedField, MmFeatureSpec, MmFeatures, MmField, MmFieldElem, MmFlatField, MmKwargsItem,
28 + MmSharedField, MmSlice, PlaceholderRange, SliceSpec,
29 +};
30 +use vllm_engine_core_client::protocol::tensor::WireTensor;
31 +use vllm_text::Prompt;
32 +use vllm_text::tokenizer::{DynTokenizer, Tokenizer};
33 +
34 +use crate::error::{Error, Result, bail_multimodal, multimodal};
35 +use crate::renderer::RenderedPrompt;
36 +use crate::request::{ChatContent, ChatContentPart, ChatMessage, ChatRequest};
37 +
38 +mod tensor;
39 +
40 +/// Resolved multimodal support for one loaded model.
41 +#[derive(Clone)]
42 +pub struct MultimodalModelInfo {
43 + context: MultimodalModelContext,
44 + spec: ResolvedMultimodalSpec,
45 + image_processor: ResolvedImageProcessor,
46 + media_connector: Arc<MediaConnector>,
47 +}
48 +
49 +/// Model metadata and tokenizer access shared by all multimodal specs.
50 +#[derive(Clone)]
51 +struct MultimodalModelContext {
52 + model_id: String,
53 + model_type: Option<String>,
54 + config: serde_json::Value,
55 + tokenizer: TokenizerResolver,
56 +}
57 +
58 +impl MultimodalModelContext {
59 + fn metadata(&self) -> ModelMetadata<'_> {
60 + ModelMetadata {
61 + model_id: &self.model_id,
62 + tokenizer: &self.tokenizer,
63 + config: &self.config,
64 + }
65 + }
66 +
67 + fn tokenizer(&self) -> &dyn Tokenizer {
68 + self.tokenizer.0.as_ref()
69 + }
70 +
71 + /// Resolve a static model processor spec for one loaded model.
72 + fn resolve_model_spec(&self) -> Option<&'static dyn ModelProcessorSpec> {
73 + static REGISTRY: LazyLock<ModelRegistry> = LazyLock::new(ModelRegistry::new);
74 + REGISTRY.lookup(&self.metadata())
75 + }
76 +
77 + /// Resolve a static image preprocessor for one loaded model.
78 + fn resolve_image_processor(&self) -> Option<&'static dyn ImagePreProcessor> {
79 + static REGISTRY: LazyLock<ImageProcessorRegistry> =
80 + LazyLock::new(ImageProcessorRegistry::with_defaults);
81 + REGISTRY.find(&self.model_id, self.model_type.as_deref())
82 + }
83 +}
84 +
85 +/// Static model-specific prompt and tensor-layout behavior.
86 +#[derive(Clone)]
87 +struct ResolvedMultimodalSpec {
88 + raw: &'static dyn ModelProcessorSpec,
89 + placeholder_token: String,
90 + placeholder_marker_token_id: u32,
91 + placeholder_embed_token_id: u32,
92 + field_layouts: HashMap<String, FieldLayout>,
93 + keep_on_cpu_keys: HashSet<String>,
94 +}
95 +
96 +impl ResolvedMultimodalSpec {
97 + fn new(raw: &'static dyn ModelProcessorSpec, context: &MultimodalModelContext) -> Result<Self> {
98 + let metadata = context.metadata();
99 + let placeholder_token =
100 + raw.placeholder_token(&metadata).map_err(|error| multimodal!("{error}"))?;
101 + // This is the rendered prompt marker, so resolve it from the token
102 + // string itself. Do not use `ModelProcessorSpec::placeholder_token_id()`:
103 + // for some specs that ID is the replacement vision/patch token,
104 + // not necessarily the token ID of `placeholder_token`.
105 + let placeholder_marker_token_id =
106 + context.tokenizer().token_to_id(&placeholder_token).ok_or_else(|| {
107 + multimodal!(
108 + "placeholder token `{placeholder_token}` is not in the tokenizer vocabulary"
109 + )
110 + })?;
111 + let placeholder_embed_token_id =
112 + raw.placeholder_token_id(&metadata).map_err(|error| multimodal!("{error}"))? as u32;
113 +
114 + Ok(Self {
115 + raw,
116 + placeholder_token,
117 + placeholder_marker_token_id,
118 + placeholder_embed_token_id,
119 + field_layouts: raw.field_layouts(),
120 + keep_on_cpu_keys: raw.keep_on_cpu_keys().into_iter().collect(),
121 + })
122 + }
123 +
124 + fn prompt_replacements(
125 + &self,
126 + context: &MultimodalModelContext,
127 + preprocessed: &PreprocessedImages,
128 + ) -> Result<Vec<PromptReplacement>> {
129 + self.raw
130 + .prompt_replacements(&context.metadata(), preprocessed)
131 + .map_err(|error| multimodal!("{error}"))
132 + }
133 +}
134 +
135 +/// Static image preprocessor plus its loaded config.
136 +#[derive(Clone)]
137 +struct ResolvedImageProcessor {
138 + raw: &'static dyn ImagePreProcessor,
139 + config: PreProcessorConfig,
140 +}
141 +
142 +/// Request-scoped fetched media, kept together with tracker UUID metadata.
143 +struct FetchedImageMedia {
144 + frames: Vec<Arc<llm_multimodal::ImageFrame>>,
145 + uuids: Vec<Option<String>>,
146 +}
147 +
148 +impl MultimodalModelInfo {
149 + /// Load and resolve multimodal support from model files.
150 + ///
151 + /// Returns `Ok(Some(_))` only when both the model spec and image processor
152 + /// are registered. File read/parse failures are real errors; unsupported
153 + /// model families are logged and returned as `Ok(None)`.
154 + pub fn from_paths(
155 + model_id: String,
156 + model_type: Option<String>,
157 + config_path: Option<&Path>,
158 + preprocessor_config_path: Option<&Path>,
159 + tokenizer: DynTokenizer,
160 + ) -> Result<Option<Self>> {
161 + let config = match config_path {
162 + Some(path) => {
163 + let text = fs::read_to_string(path)
164 + .map_err(|error| multimodal!("failed to read config.json: {error}"))?;
165 + serde_json::from_str(&text)
166 + .map_err(|error| multimodal!("failed to parse config.json: {error}"))?
167 + }
168 + None => serde_json::Value::Object(Default::default()),
169 + };
170 + let preprocessor_config = match preprocessor_config_path {
171 + Some(path) => {
172 + let text = fs::read_to_string(path).map_err(|error| {
173 + multimodal!("failed to read preprocessor_config.json: {error}")
174 + })?;
175 + PreProcessorConfig::from_json(&text).map_err(|error| {
176 + multimodal!("failed to parse preprocessor_config.json: {error}")
177 + })?
178 + }
179 + None => PreProcessorConfig::default(),
180 + };
181 +
182 + let context = MultimodalModelContext {
183 + model_id,
184 + model_type,
185 + config,
186 + tokenizer: TokenizerResolver(tokenizer),
187 + };
188 +
189 + let Some(spec) = context.resolve_model_spec() else {
190 + warn!(
191 + model_id = context.model_id,
192 + model_type = context.model_type,
193 + "multimodal model spec is not registered; disabling multimodal support for this model"
194 + );
195 + return Ok(None);
196 + };
197 + let spec = ResolvedMultimodalSpec::new(spec, &context)?;
198 +
199 + let Some(image_processor) = context.resolve_image_processor() else {
200 + warn!(
201 + model_id = context.model_id,
202 + model_type = context.model_type,
203 + "image processor is not registered; disabling multimodal support for this model"
204 + );
205 + return Ok(None);
206 + };
207 +
208 + let media_connector = Arc::new(
209 + MediaConnector::new(reqwest::Client::new(), MediaConnectorConfig::default())
210 + .map_err(|error| multimodal!("{error}"))?,
211 + );
212 +
213 + Ok(Some(Self {
214 + context,
215 + spec,
216 + image_processor: ResolvedImageProcessor {
217 + raw: image_processor,
218 + config: preprocessor_config,
219 + },
220 + media_connector,
221 + }))
222 + }
223 +
224 + /// Return the template-visible placeholder token for this model.
225 + ///
226 + /// The HF renderer uses this token while flattening image content in string
227 + /// content format.
228 + pub(crate) fn placeholder_token(&self) -> &str {
229 + &self.spec.placeholder_token
230 + }
231 +}
232 +
233 +/// Finalize a rendered chat prompt into text-generation input.
234 +///
235 +/// Text-only requests pass through unchanged as `Prompt::Text`. Multimodal
236 +/// requests are tokenized in chat, their image placeholders are expanded, and
237 +/// preprocessed image features are attached for engine-core transport.
238 +pub(crate) async fn finalize_rendered_prompt(
239 + request: &ChatRequest,
240 + rendered: RenderedPrompt,
241 + info: Option<&MultimodalModelInfo>,
242 + model_dtype: Option<ModelDtype>,
243 +) -> Result<(Prompt, Option<MmFeatures>)> {
244 + if !request.has_multimodal() {
245 + return Ok((rendered.prompt, None));
246 + }
247 + let info = info.ok_or(Error::UnsupportedMultimodalRenderer)?;
248 + let Prompt::Text(prompt) = rendered.prompt else {
249 + bail_multimodal!("multimodal chat renderer must return a text prompt before expansion");
250 + };
251 + let media_parts = extract_media_parts(request)?;
252 + let model_dtype = model_dtype.unwrap_or_else(|| {
253 + static WARN_ONCE: Once = Once::new();
254 + WARN_ONCE.call_once(|| {
255 + warn!(
256 + "engine handshake did not report model dtype; \
257 + falling back to float32 for multimodal tensor encoding"
258 + );
259 + });
260 + ModelDtype::Float32
261 + });
262 +
263 + let mut prompt_token_ids = info
264 + .context
265 + .tokenizer()
266 + .encode(&prompt, request.add_special_tokens)
267 + .map_err(|error| multimodal!("{error}"))?;
268 + let prepared = info.prepare_multimodal(media_parts, &mut prompt_token_ids, model_dtype).await?;
269 +
270 + Ok((Prompt::TokenIds(prompt_token_ids), Some(prepared)))
271 +}
272 +
273 +/// Extract image media parts from chat messages in message/content order.
274 +///
275 +/// Assistant history is skipped because generated assistant blocks are already
276 +/// represented as text for prompt rendering in this crate.
277 +fn extract_media_parts(request: &ChatRequest) -> Result<Vec<MediaContentPart>> {
278 + let mut all_parts = Vec::new();
279 + for message in &request.messages {
280 + let content = match message {
281 + ChatMessage::System { content }
282 + | ChatMessage::Developer { content, .. }
283 + | ChatMessage::User { content }
284 + | ChatMessage::ToolResponse { content, .. } => content,
285 + ChatMessage::Assistant { .. } => continue,
286 + };
287 + let ChatContent::Parts(parts) = content else {
288 + continue;
289 + };
290 + for part in parts {
291 + match part {
292 + ChatContentPart::Text { .. } => {}
293 + ChatContentPart::ImageUrl {
294 + image_url,
295 + detail,
296 + uuid,
297 + } => all_parts.push(MediaContentPart::ImageUrl {
298 + url: image_url.clone(),
299 + detail: *detail,
300 + uuid: uuid.clone(),
301 + }),
302 + }
303 + }
304 + }
305 + Ok(all_parts)
306 +}
307 +
308 +impl MultimodalModelInfo {
309 + /// Run media fetch, image preprocessing, prompt expansion, and feature
310 + /// build.
311 + ///
312 + /// `prompt_token_ids` is mutated in place because placeholder expansion
313 + /// changes both the final prompt and the offsets recorded in
314 + /// `PlaceholderRange`.
315 + async fn prepare_multimodal(
316 + &self,
317 + media_parts: Vec<MediaContentPart>,
318 + prompt_token_ids: &mut Vec<u32>,
319 + model_dtype: ModelDtype,
320 + ) -> Result<MmFeatures> {
321 + if media_parts.is_empty() {
322 + return Ok(Vec::new());
323 + }
324 + let media_parts_len = media_parts.len();
325 +
326 + let fetched = self.fetch_images(media_parts).await?;
327 + let preprocessed = self.preprocess_images(&fetched.frames).await?;
328 + let replacements = self.spec.prompt_replacements(&self.context, &preprocessed)?;
329 + let ranges = self.expand_prompt_tokens(prompt_token_ids, replacements)?;
330 +
331 + let features = self.build_features(preprocessed, fetched, ranges, model_dtype)?;
332 + if features.len() != media_parts_len {
333 + bail_multimodal!(
334 + "number of built multimodal features {} does not match number of media parts {}",
335 + features.len(),
336 + media_parts_len
337 + );
338 + }
339 + Ok(features)
340 + }
341 +
342 + /// Fetch all image parts and preserve their request-order UUID metadata.
343 + async fn fetch_images(&self, media_parts: Vec<MediaContentPart>) -> Result<FetchedImageMedia> {
344 + let mut tracker = AsyncMultiModalTracker::new(Arc::clone(&self.media_connector));
345 + for part in media_parts {
346 + tracker.push_part(part).map_err(|error| multimodal!("{error}"))?;
347 + }
348 +
349 + let tracker_output = tracker.finalize().await.map_err(|error| multimodal!("{error}"))?;
350 + let images = tracker_output.data.get(&Modality::Image).cloned().unwrap_or_default();
351 + let uuids = tracker_output.uuids.get(&Modality::Image).cloned().unwrap_or_default();
352 +
353 + let frames = images
354 + .into_iter()
355 + .map(|media| match media {
356 + TrackedMedia::Image(frame) => Ok(frame),
357 + _ => Err(Error::UnsupportedMultimodalContent("non-image")),
358 + })
359 + .collect::<Result<Vec<_>>>()?;
360 +
361 + Ok(FetchedImageMedia { frames, uuids })
362 + }
363 +
364 + /// Preprocess fetched image frames with the model's resolved image
365 + /// processor.
366 + ///
367 + /// The processor work is CPU-heavy relative to request wiring, so it runs
368 + /// in a blocking task and returns owned tensors ready for wire
369 + /// conversion.
370 + async fn preprocess_images(
371 + &self,
372 + image_frames: &[Arc<llm_multimodal::ImageFrame>],
373 + ) -> Result<PreprocessedImages> {
374 + let config = self.image_processor.config.clone();
375 + let processor = self.image_processor.raw;
376 + let images = image_frames.iter().map(|frame| frame.data().clone()).collect::<Vec<_>>();
377 +
378 + tokio::task::spawn_blocking(move || {
379 + processor.preprocess(&images, &config).map_err(|error| multimodal!("{error}"))
380 + })
381 + .await
382 + .map_err(|error| multimodal!("image preprocessing task failed: {error}"))?
383 + }
384 +
385 + /// Replace rendered placeholder markers with model-specific replacement
386 + /// tokens.
387 + ///
388 + /// Replacements are consumed in order, matching the original media-part
389 + /// order. The returned ranges point into the already-expanded prompt.
390 + fn expand_prompt_tokens(
391 + &self,
392 + prompt_token_ids: &mut Vec<u32>,
393 + replacements: Vec<PromptReplacement>,
394 + ) -> Result<Vec<PlaceholderRange>> {
395 + let mut cursor = 0;
396 + let mut ranges = Vec::with_capacity(replacements.len());
397 + for replacement in replacements {
398 + if replacement.modality != Modality::Image {
399 + bail_multimodal!(
400 + "unsupported prompt replacement modality `{}`",
401 + replacement.modality
402 + );
403 + }
404 + let offset = find_next_token(
405 + prompt_token_ids,
406 + self.spec.placeholder_marker_token_id,
407 + cursor,
408 + )
409 + .ok_or_else(|| {
410 + multimodal!(
411 + "placeholder token `{}` was not found in tokenized prompt",
412 + self.spec.placeholder_token
413 + )
414 + })?;
415 +
416 + if replacement.tokens.is_empty() {
417 + bail_multimodal!(
418 + "placeholder token `{}` expanded to no tokens",
419 + self.spec.placeholder_token
420 + );
421 + }
422 + let replacement_len = replacement.tokens.len();
423 + let replacement_tokens =
424 + replacement.tokens.iter().map(|&token| token as u32).collect::<Vec<_>>();
425 + let is_embed = {
426 + let mask = replacement_tokens
427 + .iter()
428 + .map(|&token| token == self.spec.placeholder_embed_token_id)
429 + .collect::<Vec<_>>();
430 + WireTensor::from_bool(vec![replacement_len], mask).map_err(Error::Multimodal)?
431 + };
432 +
433 + prompt_token_ids.splice(offset..offset + 1, replacement_tokens);
434 + ranges.push(PlaceholderRange {
435 + offset,
436 + length: replacement_len,
437 + is_embed: Some(is_embed),
438 + });
439 + cursor = offset + replacement_len;
440 + }
441 + Ok(ranges)
442 + }
443 +
444 + /// Convert preprocessed image tensors into engine-core multimodal features.
445 + ///
446 + /// One `MmFeatureSpec` is produced per image. Tensor fields are
447 + /// sliced according to the model spec's field layout declarations.
448 + fn build_features(
449 + &self,
450 + preprocessed: PreprocessedImages,
451 + images: FetchedImageMedia,
452 + ranges: Vec<PlaceholderRange>,
453 + model_dtype: ModelDtype,
454 + ) -> Result<MmFeatures> {
455 + let len = images.frames.len();
456 + let tensors = tensor::collect_tensors(preprocessed, model_dtype)?;
457 +
458 + let mut features = Vec::with_capacity(images.frames.len());
459 + for (index, (frame, uuid, range)) in izip!(images.frames, images.uuids, ranges).enumerate()
460 + {
461 + let mut data = MmKwargsItem::new();
462 + for (key, tensor) in &tensors {
463 + let keep_on_cpu = self.spec.keep_on_cpu_keys.contains(key);
464 + let (value, field) = match self.spec.field_layouts.get(key) {
465 + Some(FieldLayout::Batched) => (
466 + tensor.batched_value_at(index)?,
467 + MmField::Batched(MmBatchedField { keep_on_cpu }),
468 + ),
469 + Some(FieldLayout::Flat { sizes_key }) => {
470 + let sizes = tensors.get(sizes_key).ok_or_else(|| {
471 + multimodal!("flat tensor sizes key `{sizes_key}` is missing")
472 + })?;
473 + let (start, end) = tensor::flat_range_for_index(sizes, sizes_key, index)?;
474 + (
475 + tensor.flat_value_range(start, end)?,
476 + MmField::Flat(MmFlatField {
477 + slices: vec![MmSlice::Slice(SliceSpec {
478 + start: Some(0),
479 + stop: Some((end - start) as isize),
480 + step: None,
481 + })],
482 + dim: 0,
483 + keep_on_cpu,
484 + }),
485 + )
486 + }
487 + None => (
488 + tensor.clone(),
489 + MmField::Shared(MmSharedField {
490 + batch_size: len,
491 + keep_on_cpu,
492 + }),
493 + ),
494 + };
495 +
496 + data.insert(
497 + key.clone(),
498 + MmFieldElem {
499 + data: Some(value.try_into()?),
500 + field,
501 + },
502 + );
503 + }
504 +
505 + let hash = frame.hash.clone();
506 + features.push(MmFeatureSpec {
507 + data: Some(data),
508 + modality: "image".to_string(),
509 + identifier: uuid.unwrap_or_else(|| hash.clone()),
510 + mm_position: range,
511 + mm_hash: Some(hash),
512 + });
513 + }
514 +
515 + Ok(features)
516 + }
517 +}
518 +
519 +/// Find `needle` in `haystack`, starting at `start`.
520 +///
521 +/// This is intentionally order-preserving rather than a global replace: each
522 +/// image consumes the next placeholder occurrence.
523 +fn find_next_token(haystack: &[u32], needle: u32, start: usize) -> Option<usize> {
524 + haystack
525 + .get(start..)?
526 + .iter()
527 + .position(|token| *token == needle)
528 + .map(|offset| start + offset)
529 +}
530 +
531 +/// Adapter from the frontend tokenizer trait to `llm-multimodal`.
532 +#[derive(Clone)]
533 +struct TokenizerResolver(DynTokenizer);
534 +
535 +impl TokenResolver for TokenizerResolver {
536 + fn token_to_id(&self, token: &str) -> Option<u32> {
537 + self.0.token_to_id(token)
538 + }
539 +
540 + fn id_to_token(&self, id: u32) -> Option<String> {
541 + self.0.id_to_token(id)
542 + }
543 +}
544 +
545 +#[cfg(test)]
546 +mod tests {
547 + use std::sync::Arc;
548 +
549 + use llm_multimodal::TokenId;
550 + use vllm_engine_core_client::protocol::tensor::WireArrayData;
551 + use vllm_text::tokenizer::{IncrementalDecoder, Tokenizer, TokenizerError};
552 +
553 + use super::*;
554 +
555 + const LLAMA4_IMAGE_START_ID: u32 = 200088;
556 + const LLAMA4_IMAGE_END_ID: u32 = 200089;
557 + const LLAMA4_IMAGE_ID: u32 = 200090;
558 + const LLAMA4_PATCH_ID: u32 = 200092;
559 + const LLAMA4_TILE_X_SEPARATOR_ID: u32 = 200093;
560 + const LLAMA4_TILE_Y_SEPARATOR_ID: u32 = 200094;
561 +
562 + struct TestTokenizer;
563 +
564 + impl Tokenizer for TestTokenizer {
565 + fn encode(
566 + &self,
567 + text: &str,
568 + _add_special_tokens: bool,
569 + ) -> std::result::Result<Vec<u32>, TokenizerError> {
570 + Ok(match text {
571 + "<|image|>" => vec![LLAMA4_IMAGE_ID],
572 + text => text.bytes().map(u32::from).collect(),
573 + })
574 + }
575 +
576 + fn decode(
577 + &self,
578 + _token_ids: &[u32],
579 + _skip_special_tokens: bool,
580 + ) -> std::result::Result<String, TokenizerError> {
581 + Ok(String::new())
582 + }
583 +
584 + fn token_to_id(&self, token: &str) -> Option<u32> {
585 + match token {
586 + "<|image_start|>" => Some(LLAMA4_IMAGE_START_ID),
587 + "<|image_end|>" => Some(LLAMA4_IMAGE_END_ID),
588 + "<|image|>" => Some(LLAMA4_IMAGE_ID),
589 + "<|patch|>" => Some(LLAMA4_PATCH_ID),
590 + "<|tile_x_separator|>" => Some(LLAMA4_TILE_X_SEPARATOR_ID),
591 + "<|tile_y_separator|>" => Some(LLAMA4_TILE_Y_SEPARATOR_ID),
592 + _ => None,
593 + }
594 + }
595 +
596 + fn id_to_token(&self, id: u32) -> Option<String> {
597 + match id {
598 + LLAMA4_IMAGE_START_ID => Some("<|image_start|>".to_string()),
599 + LLAMA4_IMAGE_END_ID => Some("<|image_end|>".to_string()),
600 + LLAMA4_IMAGE_ID => Some("<|image|>".to_string()),
601 + LLAMA4_PATCH_ID => Some("<|patch|>".to_string()),
602 + LLAMA4_TILE_X_SEPARATOR_ID => Some("<|tile_x_separator|>".to_string()),
603 + LLAMA4_TILE_Y_SEPARATOR_ID => Some("<|tile_y_separator|>".to_string()),
604 + _ => None,
605 + }
606 + }
607 +
608 + fn create_decode_stream(
609 + &self,
610 + _prompt_token_ids: &[u32],
611 + _skip_special_tokens: bool,
612 + _min_bytes_to_buffer: usize,
613 + ) -> Box<dyn IncrementalDecoder + '_> {
614 + unreachable!("not used")
615 + }
616 + }
617 +
618 + fn test_info(model_type: &str, config: serde_json::Value) -> MultimodalModelInfo {
619 + let context = MultimodalModelContext {
620 + model_id: format!("{model_type}-test"),
621 + model_type: Some(model_type.to_string()),
622 + config,
623 + tokenizer: TokenizerResolver(Arc::new(TestTokenizer)),
624 + };
625 + let spec = context
626 + .resolve_model_spec()
627 + .unwrap_or_else(|| panic!("{model_type} spec should match"));
628 + let spec = ResolvedMultimodalSpec::new(spec, &context).unwrap();
629 + let raw_image_processor = context
630 + .resolve_image_processor()
631 + .unwrap_or_else(|| panic!("{model_type} image processor should match"));
632 + let media_connector = Arc::new(
633 + MediaConnector::new(reqwest::Client::new(), MediaConnectorConfig::default()).unwrap(),
634 + );
635 +
636 + MultimodalModelInfo {
637 + context,
638 + spec,
639 + image_processor: ResolvedImageProcessor {
640 + raw: raw_image_processor,
641 + config: PreProcessorConfig::default(),
642 + },
643 + media_connector,
644 + }
645 + }
646 +
647 + fn llama4_info() -> MultimodalModelInfo {
648 + let config = serde_json::json!({
649 + "model_type": "llama4",
650 + "image_token_index": LLAMA4_PATCH_ID,
651 + "vision_config": {"image_size": 336, "patch_size": 14}
652 + });
653 + test_info("llama4", config)
654 + }
655 +
656 + fn llama4_single_tile_replacement() -> PromptReplacement {
657 + PromptReplacement::sequence(
658 + Modality::Image,
659 + "<|image|>",
660 + vec![
661 + LLAMA4_IMAGE_START_ID as TokenId,
662 + LLAMA4_IMAGE_ID as TokenId,
663 + LLAMA4_PATCH_ID as TokenId,
664 + LLAMA4_PATCH_ID as TokenId,
665 + LLAMA4_IMAGE_END_ID as TokenId,
666 + ],
667 + )
668 + }
669 +
670 + fn llama4_multi_tile_replacement() -> PromptReplacement {
671 + PromptReplacement::sequence(
672 + Modality::Image,
673 + "<|image|>",
674 + vec![
675 + LLAMA4_IMAGE_START_ID as TokenId,
676 + LLAMA4_PATCH_ID as TokenId,
677 + LLAMA4_TILE_X_SEPARATOR_ID as TokenId,
678 + LLAMA4_PATCH_ID as TokenId,
679 + LLAMA4_TILE_Y_SEPARATOR_ID as TokenId,
680 + LLAMA4_IMAGE_ID as TokenId,
681 + LLAMA4_PATCH_ID as TokenId,
682 + LLAMA4_IMAGE_END_ID as TokenId,
683 + ],
684 + )
685 + }
686 +
687 + fn assert_bool_mask(range: &PlaceholderRange, expected: &[bool]) {
688 + let tensor = range.is_embed.as_ref().expect("is_embed mask");
689 + assert_eq!(tensor.dtype, "bool");
690 + assert_eq!(tensor.shape, vec![expected.len()]);
691 + assert_eq!(
692 + tensor.data,
693 + WireArrayData::RawView(expected.iter().map(|value| u8::from(*value)).collect())
694 + );
695 + }
696 +
697 + #[test]
698 + fn expand_prompt_tokens_marks_only_llama4_patch_tokens_as_embed() {
699 + let info = llama4_info();
700 + let mut prompt_token_ids = vec![1, LLAMA4_IMAGE_ID, 2];
701 + let replacements = vec![llama4_multi_tile_replacement()];
702 +
703 + let ranges = info.expand_prompt_tokens(&mut prompt_token_ids, replacements).unwrap();
704 +
705 + assert_eq!(
706 + prompt_token_ids,
707 + vec![
708 + 1,
709 + LLAMA4_IMAGE_START_ID,
710 + LLAMA4_PATCH_ID,
711 + LLAMA4_TILE_X_SEPARATOR_ID,
712 + LLAMA4_PATCH_ID,
713 + LLAMA4_TILE_Y_SEPARATOR_ID,
714 + LLAMA4_IMAGE_ID,
715 + LLAMA4_PATCH_ID,
716 + LLAMA4_IMAGE_END_ID,
717 + 2,
718 + ]
719 + );
720 + assert_eq!(ranges[0].offset, 1);
721 + assert_eq!(ranges[0].length, 8);
722 + assert_bool_mask(
723 + &ranges[0],
724 + &[false, true, false, true, false, false, true, false],
725 + );
726 + }
727 +
728 + #[test]
729 + fn expand_prompt_tokens_errors_when_placeholder_missing() {
730 + let info = llama4_info();
731 + let mut prompt_token_ids = vec![1, 2, 3];
732 + let replacements = vec![llama4_single_tile_replacement()];
733 +
734 + let error = info.expand_prompt_tokens(&mut prompt_token_ids, replacements).unwrap_err();
735 +
736 + assert!(matches!(error, Error::Multimodal(message) if message.contains("not found")));
737 + }
738 +
739 + #[test]
740 + fn expand_prompt_tokens_skips_llama4_image_marker_inside_replacement() {
741 + let info = llama4_info();
742 + let mut prompt_token_ids = vec![1, LLAMA4_IMAGE_ID, 2, LLAMA4_IMAGE_ID, 3];
743 + let replacements = vec![
744 + llama4_single_tile_replacement(),
745 + llama4_single_tile_replacement(),
746 + ];
747 +
748 + let ranges = info.expand_prompt_tokens(&mut prompt_token_ids, replacements).unwrap();
749 +
750 + assert_eq!(
751 + prompt_token_ids,
752 + vec![
753 + 1,
754 + LLAMA4_IMAGE_START_ID,
755 + LLAMA4_IMAGE_ID,
756 + LLAMA4_PATCH_ID,
757 + LLAMA4_PATCH_ID,
758 + LLAMA4_IMAGE_END_ID,
759 + 2,
760 + LLAMA4_IMAGE_START_ID,
761 + LLAMA4_IMAGE_ID,
762 + LLAMA4_PATCH_ID,
763 + LLAMA4_PATCH_ID,
764 + LLAMA4_IMAGE_END_ID,
765 + 3,
766 + ]
767 + );
768 + assert_eq!(ranges[0].offset, 1);
769 + assert_eq!(ranges[0].length, 5);
770 + assert_bool_mask(&ranges[0], &[false, false, true, true, false]);
771 + assert_eq!(ranges[1].offset, 7);
772 + assert_eq!(ranges[1].length, 5);
773 + assert_bool_mask(&ranges[1], &[false, false, true, true, false]);
774 + }
775 +}
rust/src/chat/src/multimodal/tensor.rs new
+342
@@ -0,0 +1,342 @@
1 +use std::collections::HashMap;
2 +
3 +use half::{bf16, f16};
4 +use llm_multimodal::{ModelSpecificValue, PreprocessedImages};
5 +use vllm_engine_core_client::protocol::ModelDtype;
6 +use vllm_engine_core_client::protocol::multimodal::MmKwargValue as ProtocolKwargValue;
7 +use vllm_engine_core_client::protocol::tensor::{ShapeExt as _, WireTensor};
8 +
9 +use crate::error::{Error, Result, bail_multimodal, multimodal};
10 +
11 +/// Representation for multimodal kwarg values for transformation.
12 +#[derive(Debug, Clone)]
13 +pub(super) enum KwargValue {
14 + /// Float tensor with row-major flat data and shape.
15 + F32Tensor { data: Vec<f32>, shape: Vec<usize> },
16 + /// Float16 tensor with row-major flat data and shape.
17 + F16Tensor { data: Vec<f16>, shape: Vec<usize> },
18 + /// BFloat16 tensor with row-major flat data and shape.
19 + Bf16Tensor { data: Vec<bf16>, shape: Vec<usize> },
20 + /// Signed integer tensor with row-major flat data and shape.
21 + I64Tensor { data: Vec<i64>, shape: Vec<usize> },
22 + /// Unsigned integer tensor with row-major flat data and shape.
23 + U32Tensor { data: Vec<u32>, shape: Vec<usize> },
24 + /// Non-tensor kwarg value that is shared or copied as-is.
25 + Passthrough(ProtocolKwargValue),
26 +}
27 +
28 +/// Collect `pixel_values` and model-specific outputs into one tensor map.
29 +pub(super) fn collect_tensors(
30 + preprocessed: PreprocessedImages,
31 + float_dtype: ModelDtype,
32 +) -> Result<HashMap<String, KwargValue>> {
33 + let PreprocessedImages {
34 + pixel_values,
35 + model_specific,
36 + ..
37 + } = preprocessed;
38 +
39 + let pixel_values = {
40 + let shape = pixel_values.shape().to_vec();
41 + let data = pixel_values.into_iter().collect();
42 + KwargValue::from_f32_tensor(data, shape, float_dtype)?
43 + };
44 +
45 + let mut tensors = HashMap::new();
46 + tensors.insert("pixel_values".to_string(), pixel_values);
47 + for (key, value) in model_specific {
48 + tensors.insert(key, KwargValue::from_model_specific(value, float_dtype)?);
49 + }
50 + Ok(tensors)
51 +}
52 +
53 +impl KwargValue {
54 + fn from_model_specific(value: ModelSpecificValue, float_dtype: ModelDtype) -> Result<Self> {
55 + use ProtocolKwargValue::*;
56 +
57 + Ok(match value {
58 + ModelSpecificValue::Tensor { data, shape } => {
59 + Self::from_f32_tensor(data, shape, float_dtype)?
60 + }
61 + ModelSpecificValue::IntTensor { data, shape } => Self::I64Tensor { data, shape },
62 + ModelSpecificValue::UintTensor { data, shape } => Self::U32Tensor { data, shape },
63 + ModelSpecificValue::Int(value) => Self::Passthrough(Int(value)),
64 + ModelSpecificValue::Float(value) => Self::Passthrough(Float(value)),
65 + ModelSpecificValue::IntVec(values) => {
66 + Self::Passthrough(List(values.into_iter().map(Int).collect()))
67 + }
68 + ModelSpecificValue::UintVec(values) => Self::Passthrough(List(
69 + values.into_iter().map(|value| Int(value as i64)).collect(),
70 + )),
71 + ModelSpecificValue::FloatVec(values) => Self::Passthrough(List(
72 + values.into_iter().map(|value| Float(value as f64)).collect(),
73 + )),
74 + ModelSpecificValue::TupleVec(values) => Self::Passthrough(List(
75 + values
76 + .into_iter()
77 + .map(|(height, width)| List(vec![Int(height as i64), Int(width as i64)]))
78 + .collect(),
79 + )),
80 + ModelSpecificValue::Bool(value) => Self::Passthrough(Int(i64::from(value))),
81 + })
82 + }
83 +
84 + /// Convert a float tensor to the target float dtype if needed, keeping the
85 + /// same shape.
86 + fn from_f32_tensor(data: Vec<f32>, shape: Vec<usize>, float_dtype: ModelDtype) -> Result<Self> {
87 + match float_dtype {
88 + ModelDtype::Float16 => Ok(Self::F16Tensor {
89 + data: data.into_iter().map(f16::from_f32).collect(),
90 + shape,
91 + }),
92 + ModelDtype::BFloat16 => Ok(Self::Bf16Tensor {
93 + data: data.into_iter().map(bf16::from_f32).collect(),
94 + shape,
95 + }),
96 + ModelDtype::Float32 => Ok(Self::F32Tensor { data, shape }),
97 + }
98 + }
99 +}
100 +
101 +impl TryFrom<KwargValue> for ProtocolKwargValue {
102 + type Error = Error;
103 +
104 + fn try_from(value: KwargValue) -> Result<Self> {
105 + match value {
106 + KwargValue::F32Tensor { data, shape } => Ok(Self::Tensor(
107 + WireTensor::from_f32(shape, data).map_err(Error::Multimodal)?,
108 + )),
109 + KwargValue::F16Tensor { data, shape } => Ok(Self::Tensor(
110 + WireTensor::from_f16(shape, data).map_err(Error::Multimodal)?,
111 + )),
112 + KwargValue::Bf16Tensor { data, shape } => Ok(Self::Tensor(
113 + WireTensor::from_bf16(shape, data).map_err(Error::Multimodal)?,
114 + )),
115 + KwargValue::I64Tensor { data, shape } => Ok(Self::Tensor(
116 + WireTensor::from_i64(shape, data).map_err(Error::Multimodal)?,
117 + )),
118 + KwargValue::U32Tensor { data, shape } => Ok(Self::Tensor(
119 + WireTensor::from_u32(shape, data).map_err(Error::Multimodal)?,
120 + )),
121 + KwargValue::Passthrough(value) => Ok(value),
122 + }
123 + }
124 +}
125 +
126 +impl KwargValue {
127 + /// Extract one image from a batched tensor field.
128 + ///
129 + /// Batched fields use their first axis as image index and drop that axis in
130 + /// the per-feature value, matching vLLM's batched-field semantics.
131 + pub(super) fn batched_value_at(&self, index: usize) -> Result<Self> {
132 + match self {
133 + Self::F32Tensor { data, shape } => {
134 + let (shape, data) = slice_first_axis_range(shape, data, index, index + 1, true)?;
135 + Ok(Self::F32Tensor { data, shape })
136 + }
137 + Self::F16Tensor { data, shape } => {
138 + let (shape, data) = slice_first_axis_range(shape, data, index, index + 1, true)?;
139 + Ok(Self::F16Tensor { data, shape })
140 + }
141 + Self::Bf16Tensor { data, shape } => {
142 + let (shape, data) = slice_first_axis_range(shape, data, index, index + 1, true)?;
143 + Ok(Self::Bf16Tensor { data, shape })
144 + }
145 + Self::I64Tensor { data, shape } => {
146 + let (shape, data) = slice_first_axis_range(shape, data, index, index + 1, true)?;
147 + Ok(Self::I64Tensor { data, shape })
148 + }
149 + Self::U32Tensor { data, shape } => {
150 + let (shape, data) = slice_first_axis_range(shape, data, index, index + 1, true)?;
151 + Ok(Self::U32Tensor { data, shape })
152 + }
153 + Self::Passthrough(value) => Ok(Self::Passthrough(value.clone())),
154 + }
155 + }
156 +
157 + /// Extract one image's variable-length range from a flat tensor field.
158 + ///
159 + /// Flat fields keep the first axis as the sliced length for this image.
160 + pub(super) fn flat_value_range(&self, start: usize, end: usize) -> Result<Self> {
161 + match self {
162 + Self::F32Tensor { data, shape } => {
163 + let (shape, data) = slice_first_axis_range(shape, data, start, end, false)?;
164 + Ok(Self::F32Tensor { data, shape })
165 + }
166 + Self::F16Tensor { data, shape } => {
167 + let (shape, data) = slice_first_axis_range(shape, data, start, end, false)?;
168 + Ok(Self::F16Tensor { data, shape })
169 + }
170 + Self::Bf16Tensor { data, shape } => {
171 + let (shape, data) = slice_first_axis_range(shape, data, start, end, false)?;
172 + Ok(Self::Bf16Tensor { data, shape })
173 + }
174 + Self::I64Tensor { data, shape } => {
175 + let (shape, data) = slice_first_axis_range(shape, data, start, end, false)?;
176 + Ok(Self::I64Tensor { data, shape })
177 + }
178 + Self::U32Tensor { data, shape } => {
179 + let (shape, data) = slice_first_axis_range(shape, data, start, end, false)?;
180 + Ok(Self::U32Tensor { data, shape })
181 + }
182 + Self::Passthrough(value) => Ok(Self::Passthrough(value.clone())),
183 + }
184 + }
185 +}
186 +
187 +/// Compute the first-axis range for one image in a flat tensor.
188 +///
189 +/// `sizes_key` names a companion tensor whose entries are cumulative slice
190 +/// sizes per image.
191 +pub(super) fn flat_range_for_index(
192 + sizes: &KwargValue,
193 + sizes_key: &str,
194 + index: usize,
195 +) -> Result<(usize, usize)> {
196 + let sizes = tensor_as_usize_vec(sizes)?;
197 + let size = *sizes.get(index).ok_or_else(|| {
198 + multimodal!("flat tensor sizes key `{sizes_key}` has no entry for image {index}")
199 + })?;
200 + let start = sizes[..index].iter().sum::<usize>();
201 + Ok((start, start + size))
202 +}
203 +
204 +/// Read a tensor value as per-image sizes for flat slicing.
205 +fn tensor_as_usize_vec(tensor: &KwargValue) -> Result<Vec<usize>> {
206 + match tensor {
207 + KwargValue::I64Tensor { data, .. } => data
208 + .iter()
209 + .map(|value| {
210 + usize::try_from(*value)
211 + .map_err(|_| multimodal!("negative flat tensor size `{value}`"))
212 + })
213 + .collect(),
214 + KwargValue::U32Tensor { data, .. } => {
215 + Ok(data.iter().map(|value| *value as usize).collect())
216 + }
217 + _ => Err(multimodal!("flat tensor sizes must be int64 or uint32")),
218 + }
219 +}
220 +
221 +/// Slice a flat row-major tensor along its first axis.
222 +fn slice_first_axis_range<T: Clone>(
223 + shape: &[usize],
224 + data: &[T],
225 + start: usize,
226 + end: usize,
227 + drop_axis: bool,
228 +) -> Result<(Vec<usize>, Vec<T>)> {
229 + let first_dim = *shape.first().ok_or_else(|| multimodal!("tensor has no first dimension"))?;
230 + if start > end || end > first_dim {
231 + bail_multimodal!("invalid tensor slice {start}..{end} for first dimension {first_dim}");
232 + }
233 + let expected_len = shape
234 + .checked_numel()
235 + .ok_or_else(|| multimodal!("tensor shape {shape:?} has too many elements"))?;
236 + if expected_len != data.len() {
237 + bail_multimodal!(
238 + "tensor shape {shape:?} expects {expected_len} elements, got {}",
239 + data.len()
240 + );
241 + }
242 + let stride = shape[1..].iter().product::<usize>();
243 + let data_start = start * stride;
244 + let data_end = end * stride;
245 + let out_shape = if drop_axis {
246 + shape[1..].to_vec()
247 + } else {
248 + let mut shape = shape.to_vec();
249 + shape[0] = end - start;
250 + shape
251 + };
252 + Ok((out_shape, data[data_start..data_end].to_vec()))
253 +}
254 +
255 +#[cfg(test)]
256 +mod tests {
257 + use super::*;
258 +
259 + #[test]
260 + fn batched_value_at_drops_first_axis() {
261 + let value = KwargValue::F32Tensor {
262 + data: vec![1.0, 2.0, 3.0, 4.0],
263 + shape: vec![2, 2],
264 + };
265 +
266 + let value = value.batched_value_at(1).unwrap();
267 +
268 + assert!(matches!(
269 + value,
270 + KwargValue::F32Tensor { data, shape }
271 + if shape == vec![2] && data == vec![3.0, 4.0]
272 + ));
273 + }
274 +
275 + #[test]
276 + fn flat_value_range_keeps_first_axis() {
277 + let value = KwargValue::U32Tensor {
278 + data: (0..10).collect(),
279 + shape: vec![5, 2],
280 + };
281 +
282 + let value = value.flat_value_range(1, 3).unwrap();
283 +
284 + assert!(matches!(
285 + value,
286 + KwargValue::U32Tensor { data, shape }
287 + if shape == vec![2, 2] && data == vec![2, 3, 4, 5]
288 + ));
289 + }
290 +
291 + #[test]
292 + fn flat_range_for_index_uses_size_tensor() {
293 + let sizes = KwargValue::I64Tensor {
294 + data: vec![2, 3, 4],
295 + shape: vec![3],
296 + };
297 +
298 + assert_eq!(
299 + flat_range_for_index(&sizes, "image_grid_thw", 1).unwrap(),
300 + (2, 5)
301 + );
302 + }
303 +
304 + #[test]
305 + fn slice_first_axis_range_errors_on_shape_data_mismatch() {
306 + let error = slice_first_axis_range(&[2, 2], &[1.0_f32, 2.0, 3.0], 0, 1, true).unwrap_err();
307 +
308 + assert!(
309 + matches!(error, Error::Multimodal(message) if message.contains("expects 4 elements"))
310 + );
311 + }
312 +
313 + #[test]
314 + fn bfloat16_tensor_wire_uses_bfloat16_dtype() {
315 + let value =
316 + KwargValue::from_f32_tensor(vec![1.0, -1.0], vec![2], ModelDtype::BFloat16).unwrap();
317 +
318 + let ProtocolKwargValue::Tensor(tensor) = ProtocolKwargValue::try_from(value).unwrap()
319 + else {
320 + panic!("expected tensor");
321 + };
322 +
323 + assert_eq!(tensor.dtype, "bfloat16");
324 + assert_eq!(tensor.shape, vec![2]);
325 + assert_eq!(tensor.data.into_raw_view().unwrap().len(), 4);
326 + }
327 +
328 + #[test]
329 + fn float16_tensor_wire_uses_float16_dtype() {
330 + let value =
331 + KwargValue::from_f32_tensor(vec![1.0, -1.0], vec![2], ModelDtype::Float16).unwrap();
332 +
333 + let ProtocolKwargValue::Tensor(tensor) = ProtocolKwargValue::try_from(value).unwrap()
334 + else {
335 + panic!("expected tensor");
336 + };
337 +
338 + assert_eq!(tensor.dtype, "float16");
339 + assert_eq!(tensor.shape, vec![2]);
340 + assert_eq!(tensor.data.into_raw_view().unwrap().len(), 4);
341 + }
342 +}
rust/src/chat/src/output/default/mod.rs new
+166
@@ -0,0 +1,166 @@
1 +//! Default output processing pipeline.
2 +
3 +mod reasoning;
4 +mod tool;
5 +
6 +use std::sync::Once;
7 +
8 +use futures::{Stream, StreamExt as _};
9 +use tracing::info;
10 +use trait_set::trait_set;
11 +use vllm_text::tokenizer::DynTokenizer;
12 +
13 +use self::reasoning::reasoning_event_stream;
14 +use self::tool::tool_event_stream;
15 +use super::structured::structured_chat_event_stream;
16 +use crate::error::Result;
17 +use crate::output::{
18 + AssistantEvent, ChatOutputProcessor, ContentEvent, DynChatEventStream,
19 + DynDecodedTextEventStream,
20 +};
21 +use crate::parser::ParserSelection;
22 +use crate::parser::reasoning::{ReasoningParser, ReasoningParserFactory};
23 +use crate::parser::tool::{ToolParser, ToolParserFactory};
24 +use crate::request::{ChatRequest, ChatToolChoice};
25 +use crate::{Error, Result as ChatResult};
26 +
27 +trait_set! {
28 + trait ContentEventStream = Stream<Item = Result<ContentEvent>> + Send + 'static;
29 +}
30 +
31 +/// Default request-scoped output processor used by Hugging Face style chat
32 +/// backends.
33 +///
34 +/// This implementation assumes the backend already emitted decoded text deltas,
35 +/// then optionally layers reasoning parsing and tool-call parsing before
36 +/// assembling final structured chat events.
37 +pub struct DefaultChatOutputProcessor {
38 + reasoning_parser: Option<Box<dyn ReasoningParser>>,
39 + tool_parser: Option<Box<dyn ToolParser>>,
40 +}
41 +
42 +impl DefaultChatOutputProcessor {
43 + /// Build the default output processor and apply any parser-specific request
44 + /// adjustments.
45 + ///
46 + /// Parser resolution happens here so that request validation, prompt
47 + /// rendering, and streaming all observe the same parser-adjusted
48 + /// request state.
49 + pub fn new(
50 + request: &mut ChatRequest,
51 + model_id: &str,
52 + tokenizer: DynTokenizer,
53 + tool_call_parser: &ParserSelection,
54 + reasoning_parser: &ParserSelection,
55 + ) -> ChatResult<Self> {
56 + let tool_parsing_enabled =
57 + matches!(request.tool_choice, ChatToolChoice::Auto) && !request.tools.is_empty();
58 + let tool_parser = if tool_parsing_enabled {
59 + Some(Self::resolve_tool_parser(
60 + request,
61 + model_id,
62 + tool_call_parser,
63 + )?)
64 + } else {
65 + None
66 + };
67 + let reasoning_parser = Self::resolve_optional_reasoning_parser(
68 + request,
69 + model_id,
70 + tokenizer,
71 + reasoning_parser,
72 + )?;
73 +
74 + Ok(Self {
75 + reasoning_parser,
76 + tool_parser,
77 + })
78 + }
79 +
80 + /// Build the plain-text-only default output processor.
81 + ///
82 + /// This keeps the default structured chat-event assembly but disables both
83 + /// reasoning parsing and tool-call parsing completely, so that all
84 + /// content is treated as opaque text.
85 + pub fn plain_text_only() -> Self {
86 + Self {
87 + reasoning_parser: None,
88 + tool_parser: None,
89 + }
90 + }
91 +
92 + fn resolve_tool_parser(
93 + request: &mut ChatRequest,
94 + model_id: &str,
95 + selection: &ParserSelection,
96 + ) -> ChatResult<Box<dyn ToolParser>> {
97 + let factory = ToolParserFactory::global();
98 + let parser_name = match selection {
99 + ParserSelection::Auto => factory.resolve_name_for_model(model_id).ok_or_else(|| {
100 + Error::ParserUnavailableForModel {
101 + kind: "tool",
102 + model_id: model_id.to_string(),
103 + }
104 + })?,
105 + ParserSelection::None => return Err(Error::ParserDisabled { kind: "tool" }),
106 + ParserSelection::Explicit(name) => name.as_str(),
107 + };
108 +
109 + let parser = factory.create(parser_name, &request.tools)?;
110 +
111 + if parser.preserve_special_tokens() {
112 + request.decode_options.skip_special_tokens = false;
113 + }
114 +
115 + TOOL_PARSER_LOG_ONCE.call_once(|| info!(parser_name, "using tool parser"));
116 + Ok(parser)
117 + }
118 +
119 + fn resolve_optional_reasoning_parser(
120 + request: &mut ChatRequest,
121 + model_id: &str,
122 + tokenizer: DynTokenizer,
123 + selection: &ParserSelection,
124 + ) -> ChatResult<Option<Box<dyn ReasoningParser>>> {
125 + let factory = ReasoningParserFactory::global();
126 + let parser_name = match selection {
127 + ParserSelection::Auto => factory.resolve_name_for_model(model_id),
128 + ParserSelection::None => None,
129 + ParserSelection::Explicit(name) => Some(name.as_str()),
130 + };
131 +
132 + let Some(parser_name) = parser_name else {
133 + REASONING_PARSER_LOG_ONCE.call_once(|| info!("reasoning parsing disabled"));
134 + return Ok(None);
135 + };
136 +
137 + let parser = factory.create(parser_name, tokenizer)?;
138 +
139 + if parser.preserve_special_tokens() {
140 + request.decode_options.skip_special_tokens = false;
141 + }
142 +
143 + REASONING_PARSER_LOG_ONCE.call_once(|| info!(parser_name, "using reasoning parser"));
144 + Ok(Some(parser))
145 + }
146 +}
147 +
148 +static TOOL_PARSER_LOG_ONCE: Once = Once::new();
149 +static REASONING_PARSER_LOG_ONCE: Once = Once::new();
150 +
151 +impl ChatOutputProcessor for DefaultChatOutputProcessor {
152 + /// Transforms a raw generate-output token stream into structured chat
153 + /// events through three sequential stages once text decoding has
154 + /// already happened:
155 + ///
156 + /// 1. [`reasoning_event_stream`] — reasoning/content separation
157 + /// 2. [`tool_event_stream`] — tool-call parsing
158 + /// 3. [`structured_chat_event_stream`] — final block assembly
159 + fn process(self: Box<Self>, decoded: DynDecodedTextEventStream) -> Result<DynChatEventStream> {
160 + let reasoning = reasoning_event_stream(decoded, self.reasoning_parser);
161 + let tool = tool_event_stream(reasoning, self.tool_parser);
162 + let structured = structured_chat_event_stream(tool);
163 +
164 + Ok(structured.boxed())
165 + }
166 +}
rust/src/chat/src/output/default/reasoning.rs new
+504
@@ -0,0 +1,504 @@
1 +//! Adapts decoded text updates into reasoning-aware assistant deltas.
2 +//!
3 +//! This stage sits between low-level token decoding and final block assembly.
4 +//! It is the only place in the new pipeline that understands reasoning
5 +//! separation: `decoded.rs` still only produces plain text deltas, while later
6 +//! stages consume the semantic `Text` / `Reasoning` split emitted here.
7 +
8 +use asynk_strim_attr::{TryYielder, try_stream};
9 +use futures::{StreamExt as _, pin_mut};
10 +use thiserror_ext::AsReport;
11 +use tracing::warn;
12 +use vllm_text::output::DecodedTextEvent;
13 +
14 +use super::ContentEvent;
15 +use crate::Result;
16 +use crate::error::Error;
17 +use crate::event::AssistantBlockKind;
18 +use crate::output::DecodedTextEventStream;
19 +use crate::parser::reasoning::{ReasoningDelta, ReasoningParser};
20 +
21 +/// Per-stream reasoning parsing state.
22 +struct ReasoningState {
23 + /// Reasoning parser for the current model family.
24 + parser: Box<dyn ReasoningParser>,
25 + /// Whether reasoning parsing has already failed for this stream.
26 + parser_failed: bool,
27 +}
28 +
29 +impl ReasoningState {
30 + /// Create one fresh reasoning-adaptation state for a new streamed response.
31 + fn new(parser: Box<dyn ReasoningParser>) -> Self {
32 + Self {
33 + parser,
34 + parser_failed: false,
35 + }
36 + }
37 +
38 + /// Convert one decoded text delta into zero or more semantic assistant
39 + /// deltas.
40 + fn process_delta(&mut self, delta: String) -> Vec<ContentEvent> {
41 + // If the parser has already failed, skip parsing and return plain text deltas.
42 + if self.parser_failed {
43 + return vec![ContentEvent::TextDelta {
44 + kind: AssistantBlockKind::Text,
45 + delta,
46 + }];
47 + }
48 +
49 + let mut events = Vec::new();
50 +
51 + match self.parser.push(&delta) {
52 + Ok(result) => {
53 + push_reasoning_delta(&mut events, result);
54 + }
55 + Err(error) => {
56 + if !self.parser_failed {
57 + warn!(
58 + error = %error.as_report(),
59 + "reasoning parser failed; falling back to plain text deltas"
60 + );
61 + self.parser_failed = true;
62 + }
63 + push_text_delta(&mut events, AssistantBlockKind::Text, delta);
64 + }
65 + }
66 +
67 + events
68 + }
69 +
70 + /// Initialize parser state once prompt token IDs are available.
71 + fn initialize(&mut self, prompt_token_ids: &[u32]) {
72 + if self.parser_failed {
73 + return;
74 + }
75 +
76 + match self.parser.initialize(prompt_token_ids) {
77 + Ok(()) => {}
78 + Err(error) => {
79 + warn!(
80 + error = %error.as_report(),
81 + "failed to initialize reasoning parser; falling back to plain text deltas"
82 + );
83 + self.parser_failed = true;
84 + }
85 + }
86 + }
87 +
88 + /// Flush any parser-held partial delimiter state at end of stream.
89 + fn finish(&mut self) -> Vec<ContentEvent> {
90 + if self.parser_failed {
91 + return Vec::new();
92 + }
93 +
94 + match self.parser.finish() {
95 + Ok(result) => {
96 + let mut events = Vec::new();
97 + push_reasoning_delta(&mut events, result);
98 + events
99 + }
100 + Err(error) => {
101 + warn!(error = %error.as_report(), "failed to flush reasoning parser state");
102 + Vec::new()
103 + }
104 + }
105 + }
106 +}
107 +
108 +/// Push one semantic text delta if it is non-empty.
109 +fn push_text_delta(events: &mut Vec<ContentEvent>, kind: AssistantBlockKind, delta: String) {
110 + if delta.is_empty() {
111 + return;
112 + }
113 + events.push(ContentEvent::TextDelta { kind, delta });
114 +}
115 +
116 +/// Convert one parsed reasoning delta into zero or more content events.
117 +fn push_reasoning_delta(events: &mut Vec<ContentEvent>, delta: ReasoningDelta) {
118 + if let Some(reasoning) = delta.reasoning {
119 + push_text_delta(events, AssistantBlockKind::Reasoning, reasoning);
120 + }
121 + if let Some(content) = delta.content {
122 + push_text_delta(events, AssistantBlockKind::Text, content);
123 + }
124 +}
125 +
126 +/// Wrap one decoded-text stream into the internal reasoning event stream.
127 +#[try_stream]
128 +pub(crate) async fn reasoning_event_stream(
129 + decoded_stream: impl DecodedTextEventStream,
130 + reasoning_parser: Option<Box<dyn ReasoningParser>>,
131 + mut y: TryYielder<ContentEvent, Error>,
132 +) -> Result<()> {
133 + pin_mut!(decoded_stream);
134 +
135 + // Without a parser, pass through as plain text deltas.
136 + let Some(reasoning_parser) = reasoning_parser else {
137 + while let Some(event) = decoded_stream.next().await.transpose()? {
138 + for next in ContentEvent::from_decoded_plain_text(event) {
139 + y.yield_ok(next).await;
140 + }
141 + }
142 + return Ok(());
143 + };
144 +
145 + let mut state = ReasoningState::new(reasoning_parser);
146 +
147 + while let Some(event) = decoded_stream.next().await.transpose()? {
148 + match event {
149 + DecodedTextEvent::Start {
150 + prompt_token_ids,
151 + prompt_logprobs,
152 + } => {
153 + state.initialize(&prompt_token_ids);
154 + y.yield_ok(ContentEvent::Start {
155 + prompt_token_ids,
156 + prompt_logprobs,
157 + })
158 + .await;
159 + }
160 + DecodedTextEvent::TextDelta {
161 + delta,
162 + token_ids,
163 + logprobs,
164 + finished,
165 + } => {
166 + for next in state.process_delta(delta) {
167 + y.yield_ok(next).await;
168 + }
169 + if logprobs.is_some() || !token_ids.is_empty() {
170 + y.yield_ok(ContentEvent::LogprobsDelta {
171 + logprobs,
172 + token_ids,
173 + })
174 + .await;
175 + }
176 + if let Some(finished) = finished {
177 + for next in state.finish() {
178 + y.yield_ok(next).await;
179 + }
180 + y.yield_ok(ContentEvent::Done {
181 + prompt_token_count: finished.prompt_token_count,
182 + output_token_count: finished.output_token_count,
183 + finish_reason: finished.finish_reason,
184 + kv_transfer_params: finished.kv_transfer_params,
185 + })
186 + .await;
187 + }
188 + }
189 + }
190 + }
191 + Ok(())
192 +}
193 +
194 +#[cfg(test)]
195 +mod tests {
196 +
197 + use std::sync::Arc;
198 +
199 + use futures::{StreamExt as _, stream};
200 + use vllm_llm::FinishReason;
201 + use vllm_text::output::{
202 + DecodedLogprobs, DecodedPositionLogprobs, DecodedTextEvent, DecodedTokenLogprob,
203 + };
204 + use vllm_tokenizer::{DynTokenizer, Tokenizer};
205 +
206 + use super::super::ContentEvent;
207 + use super::reasoning_event_stream;
208 + use crate::event::AssistantBlockKind;
209 + use crate::parser::reasoning::{
210 + ReasoningDelta, ReasoningError, ReasoningParser, ReasoningParserFactory, names,
211 + };
212 +
213 + struct FakeTokenizer;
214 +
215 + impl Tokenizer for FakeTokenizer {
216 + fn encode(
217 + &self,
218 + text: &str,
219 + _add_special_tokens: bool,
220 + ) -> vllm_tokenizer::Result<Vec<u32>> {
221 + Ok(text.chars().map(u32::from).collect())
222 + }
223 +
224 + fn decode(
225 + &self,
226 + token_ids: &[u32],
227 + _skip_special_tokens: bool,
228 + ) -> vllm_tokenizer::Result<String> {
229 + Ok(token_ids
230 + .iter()
231 + .map(|token_id| char::from_u32(*token_id).unwrap_or('\u{FFFD}'))
232 + .collect())
233 + }
234 +
235 + fn token_to_id(&self, token: &str) -> Option<u32> {
236 + match token {
237 + "<think>" => Some(1),
238 + "</think>" => Some(2),
239 + _ => None,
240 + }
241 + }
242 + }
243 +
244 + struct FailingReasoningParser {
245 + fail_next: bool,
246 + }
247 +
248 + impl ReasoningParser for FailingReasoningParser {
249 + fn create(_tokenizer: DynTokenizer) -> Result<Box<dyn ReasoningParser>, ReasoningError>
250 + where
251 + Self: Sized + 'static,
252 + {
253 + Ok(Box::new(Self { fail_next: true }))
254 + }
255 +
256 + fn push(&mut self, _text: &str) -> Result<ReasoningDelta, ReasoningError> {
257 + if self.fail_next {
258 + self.fail_next = false;
259 + return Err(ReasoningError::MissingToken {
260 + token: "<think>".to_string(),
261 + });
262 + }
263 + Ok(ReasoningDelta::default())
264 + }
265 + }
266 +
267 + fn test_reasoning_parser(factory: &mut ReasoningParserFactory) -> Box<dyn ReasoningParser> {
268 + factory.register_parser::<FailingReasoningParser>("failing");
269 +
270 + factory.create("failing", Arc::new(FakeTokenizer)).unwrap()
271 + }
272 +
273 + #[tokio::test]
274 + async fn reasoning_parser_failure_falls_back_to_plain_text() {
275 + let mut factory = ReasoningParserFactory::new();
276 + let events = stream::iter(vec![
277 + Ok(DecodedTextEvent::Start {
278 + prompt_token_ids: vec![1, 2, 3].into(),
279 + prompt_logprobs: None,
280 + }),
281 + Ok(DecodedTextEvent::TextDelta {
282 + delta: "abc".to_string(),
283 + token_ids: vec![],
284 + logprobs: None,
285 + finished: None,
286 + }),
287 + Ok(DecodedTextEvent::TextDelta {
288 + delta: "def".to_string(),
289 + token_ids: vec![],
290 + logprobs: None,
291 + finished: Some(vllm_text::Finished {
292 + prompt_token_count: 3,
293 + output_token_count: 0,
294 + finish_reason: FinishReason::stop_eos(),
295 + kv_transfer_params: None,
296 + }),
297 + }),
298 + ]);
299 +
300 + let collected = reasoning_event_stream(events, Some(test_reasoning_parser(&mut factory)))
301 + .collect::<Vec<_>>()
302 + .await;
303 +
304 + let events = collected
305 + .into_iter()
306 + .collect::<crate::Result<Vec<_>>>()
307 + .expect("reasoning stream should not fail");
308 +
309 + assert_eq!(
310 + events,
311 + vec![
312 + ContentEvent::Start {
313 + prompt_token_ids: vec![1, 2, 3].into(),
314 + prompt_logprobs: None,
315 + },
316 + ContentEvent::TextDelta {
317 + kind: AssistantBlockKind::Text,
318 + delta: "abc".to_string(),
319 + },
320 + ContentEvent::TextDelta {
321 + kind: AssistantBlockKind::Text,
322 + delta: "def".to_string(),
323 + },
324 + ContentEvent::Done {
325 + prompt_token_count: 3,
326 + output_token_count: 0,
327 + finish_reason: FinishReason::stop_eos(),
328 + kv_transfer_params: None,
329 + },
330 + ]
331 + );
332 + }
333 +
334 + #[tokio::test]
335 + async fn reasoning_stream_preserves_logprobs_delta() {
336 + let events = stream::iter(vec![
337 + Ok(DecodedTextEvent::Start {
338 + prompt_token_ids: vec![1].into(),
339 + prompt_logprobs: None,
340 + }),
341 + Ok(DecodedTextEvent::TextDelta {
342 + delta: "abc".to_string(),
343 + token_ids: vec![],
344 + logprobs: Some(DecodedLogprobs {
345 + positions: vec![DecodedPositionLogprobs {
346 + entries: vec![DecodedTokenLogprob {
347 + token_id: 0,
348 + token: "a".to_string(),
349 + logprob: -0.1,
350 + rank: 1,
351 + }],
352 + }],
353 + }),
354 + finished: None,
355 + }),
356 + ]);
357 +
358 + let collected = reasoning_event_stream(events, None)
359 + .collect::<Vec<_>>()
360 + .await
361 + .into_iter()
362 + .collect::<crate::Result<Vec<_>>>()
363 + .unwrap();
364 +
365 + assert_eq!(
366 + collected,
367 + vec![
368 + ContentEvent::Start {
369 + prompt_token_ids: vec![1].into(),
370 + prompt_logprobs: None,
371 + },
372 + ContentEvent::TextDelta {
373 + kind: AssistantBlockKind::Text,
374 + delta: "abc".to_string(),
375 + },
376 + ContentEvent::LogprobsDelta {
377 + logprobs: Some(DecodedLogprobs {
378 + positions: vec![DecodedPositionLogprobs {
379 + entries: vec![DecodedTokenLogprob {
380 + token_id: 0,
381 + token: "a".to_string(),
382 + logprob: -0.1,
383 + rank: 1,
384 + }],
385 + }],
386 + }),
387 + token_ids: vec![],
388 + },
389 + ]
390 + );
391 + }
392 +
393 + #[tokio::test]
394 + async fn qwen3_parser_uses_prompt_end_marker_to_switch_to_content() {
395 + let tokenizer = Arc::new(FakeTokenizer);
396 + let events = stream::iter(vec![
397 + Ok(DecodedTextEvent::Start {
398 + prompt_token_ids: vec![2].into(),
399 + prompt_logprobs: None,
400 + }),
401 + Ok(DecodedTextEvent::TextDelta {
402 + delta: "thought ".to_string(),
403 + token_ids: vec![],
404 + logprobs: None,
405 + finished: None,
406 + }),
407 + Ok(DecodedTextEvent::TextDelta {
408 + delta: "done</think>OK".to_string(),
409 + token_ids: vec![],
410 + logprobs: None,
411 + finished: None,
412 + }),
413 + ]);
414 +
415 + let factory = ReasoningParserFactory::new();
416 + let collected = reasoning_event_stream(
417 + events,
418 + Some(factory.create(names::QWEN3, tokenizer).unwrap()),
419 + )
420 + .collect::<Vec<_>>()
421 + .await;
422 +
423 + let events = collected
424 + .into_iter()
425 + .collect::<crate::Result<Vec<_>>>()
426 + .expect("reasoning stream should not fail");
427 +
428 + assert_eq!(
429 + events,
430 + vec![
431 + ContentEvent::Start {
432 + prompt_token_ids: vec![2].into(),
433 + prompt_logprobs: None,
434 + },
435 + ContentEvent::TextDelta {
436 + kind: AssistantBlockKind::Text,
437 + delta: "thought ".to_string(),
438 + },
439 + ContentEvent::TextDelta {
440 + kind: AssistantBlockKind::Text,
441 + delta: "done</think>OK".to_string(),
442 + },
443 + ]
444 + );
445 + }
446 +
447 + #[tokio::test]
448 + async fn qwen3_parser_tolerates_prompt_prefill_reasoning() {
449 + let tokenizer = Arc::new(FakeTokenizer);
450 + let events = stream::iter(vec![
451 + Ok(DecodedTextEvent::Start {
452 + prompt_token_ids: vec![1].into(),
453 + prompt_logprobs: None,
454 + }),
455 + Ok(DecodedTextEvent::TextDelta {
456 + delta: "thought ".to_string(),
457 + token_ids: vec![],
458 + logprobs: None,
459 + finished: None,
460 + }),
461 + Ok(DecodedTextEvent::TextDelta {
462 + delta: "done</think>OK".to_string(),
463 + token_ids: vec![],
464 + logprobs: None,
465 + finished: None,
466 + }),
467 + ]);
468 +
469 + let factory = ReasoningParserFactory::new();
470 + let collected = reasoning_event_stream(
471 + events,
472 + Some(factory.create(names::QWEN3, tokenizer).unwrap()),
473 + )
474 + .collect::<Vec<_>>()
475 + .await;
476 +
477 + let events = collected
478 + .into_iter()
479 + .collect::<crate::Result<Vec<_>>>()
480 + .expect("reasoning stream should not fail");
481 +
482 + assert_eq!(
483 + events,
484 + vec![
485 + ContentEvent::Start {
486 + prompt_token_ids: vec![1].into(),
487 + prompt_logprobs: None,
488 + },
489 + ContentEvent::TextDelta {
490 + kind: AssistantBlockKind::Reasoning,
491 + delta: "thought ".to_string(),
492 + },
493 + ContentEvent::TextDelta {
494 + kind: AssistantBlockKind::Reasoning,
495 + delta: "done".to_string(),
496 + },
497 + ContentEvent::TextDelta {
498 + kind: AssistantBlockKind::Text,
499 + delta: "OK".to_string(),
500 + },
501 + ]
502 + );
503 + }
504 +}
rust/src/chat/src/output/default/tool.rs new
+625
@@ -0,0 +1,625 @@
1 +//! Adapts plain assistant text deltas into tool-call-aware assistant updates.
2 +//!
3 +//! This stage runs after reasoning separation and before final block assembly.
4 +//! It only inspects normal assistant text, leaves reasoning deltas untouched,
5 +//! and translates incremental tool parsing output into internal tool-call
6 +//! events while preserving plain-text fallback behavior.
7 +
8 +use asynk_strim_attr::{TryYielder, try_stream};
9 +use futures::{StreamExt as _, pin_mut};
10 +use thiserror_ext::AsReport;
11 +use tracing::warn;
12 +
13 +use super::{AssistantEvent, ContentEvent, ContentEventStream};
14 +use crate::Result;
15 +use crate::error::Error;
16 +use crate::event::AssistantBlockKind;
17 +use crate::output::generate_tool_call_id;
18 +use crate::parser::tool::{ToolCallDelta, ToolParseResult, ToolParser};
19 +
20 +/// Per-stream tool parsing state.
21 +struct ToolState {
22 + /// Parser for the current model family.
23 + parser: Box<dyn ToolParser>,
24 + /// Whether tool parsing has already failed for this stream.
25 + parser_failed: bool,
26 + /// The parser-local index of the currently open tool call, if any.
27 + // NOTE: We only allow single open tool call at a time right now, since that's what all
28 + // supported parsers currently emit. Change this to a `BTreeMap` if we need to support multiple
29 + // interleaved calls in the future.
30 + open_call_index: Option<usize>,
31 +}
32 +
33 +impl ToolState {
34 + /// Create one fresh tool-parsing state for a new streamed response.
35 + fn new(parser: Box<dyn ToolParser>) -> Self {
36 + Self {
37 + parser,
38 + parser_failed: false,
39 + open_call_index: None,
40 + }
41 + }
42 +
43 + /// Convert one semantic assistant text delta into zero or more tool-aware
44 + /// internal events.
45 + fn process_text_delta(
46 + &mut self,
47 + kind: AssistantBlockKind,
48 + delta: String,
49 + ) -> Result<Vec<AssistantEvent>> {
50 + let mut events = Vec::new();
51 +
52 + // Only normal assistant text is eligible for tool parsing. Reasoning
53 + // blocks and plain-text fallback should pass through unchanged.
54 + if kind != AssistantBlockKind::Text || self.parser_failed {
55 + self.open_call_index = None;
56 + events.push(AssistantEvent::TextDelta { kind, delta });
57 + return Ok(events);
58 + }
59 +
60 + let parse_result = self.parser.push(&delta);
61 +
62 + match parse_result {
63 + Ok(result) => self.process_parse_result(kind, result, &mut events)?,
64 + Err(error) => {
65 + if !self.parser_failed {
66 + warn!(
67 + error = %error.as_report(),
68 + "tool parser failed; falling back to plain text deltas"
69 + );
70 + self.parser_failed = true;
71 + }
72 + self.open_call_index = None;
73 + events.push(AssistantEvent::TextDelta { kind, delta });
74 + }
75 + }
76 +
77 + Ok(events)
78 + }
79 +
80 + /// Apply one parsed tool result to the current stream state.
81 + fn process_parse_result(
82 + &mut self,
83 + kind: AssistantBlockKind,
84 + result: ToolParseResult,
85 + events: &mut Vec<AssistantEvent>,
86 + ) -> Result<()> {
87 + // When we are not currently streaming a tool call, preserve plain
88 + // text first and then surface any new tool call items.
89 + if self.open_call_index.is_none() {
90 + push_text_delta(events, kind, result.normal_text);
91 + self.process_tool_items(result.calls, events)?;
92 + } else {
93 + // Once a tool call is open, prioritize tool deltas first. If the
94 + // parser emits normal text again, close the tool call and resume
95 + // plain text output.
96 + self.process_tool_items(result.calls, events)?;
97 + if !result.normal_text.is_empty() {
98 + self.open_call_index = None;
99 + push_text_delta(events, kind, result.normal_text);
100 + }
101 + }
102 + Ok(())
103 + }
104 +
105 + /// Apply one batch of parsed tool-call deltas emitted by the parser.
106 + fn process_tool_items(
107 + &mut self,
108 + items: Vec<ToolCallDelta>,
109 + events: &mut Vec<AssistantEvent>,
110 + ) -> Result<()> {
111 + for item in items {
112 + if let Some(name) = item.name {
113 + let is_new_tool = match self.open_call_index {
114 + Some(open_call_index) => open_call_index != item.tool_index,
115 + None => true,
116 + };
117 + if is_new_tool {
118 + let id = generate_tool_call_id();
119 + self.open_call_index = Some(item.tool_index);
120 + events.push(AssistantEvent::ToolCallStart { id, name });
121 + }
122 + }
123 +
124 + if item.arguments.is_empty() {
125 + // No arguments delta to apply.
126 + continue;
127 + }
128 + let Some(open_call_index) = self.open_call_index else {
129 + return Err(Error::ToolCallStreamInvariant {
130 + message: format!(
131 + "received arguments for tool index {} before any tool-call start",
132 + item.tool_index
133 + ),
134 + });
135 + };
136 + if open_call_index != item.tool_index {
137 + return Err(Error::ToolCallStreamInvariant {
138 + message: format!(
139 + "received arguments for tool index {} while tool index {} is open",
140 + item.tool_index, open_call_index
141 + ),
142 + });
143 + }
144 +
145 + events.push(AssistantEvent::ToolCallArgumentsDelta {
146 + delta: item.arguments,
147 + });
148 + }
149 + Ok(())
150 + }
151 +
152 + /// Flush parser state at end-of-stream and close any remaining open calls.
153 + fn finish(&mut self) -> Result<Vec<AssistantEvent>> {
154 + let mut events = Vec::new();
155 +
156 + if self.parser_failed {
157 + return Ok(events);
158 + }
159 +
160 + match self.parser.finish() {
161 + Ok(result) => {
162 + self.process_parse_result(AssistantBlockKind::Text, result, &mut events)?
163 + }
164 + Err(error) => {
165 + warn!(
166 + error = %error.as_report(),
167 + "tool parser finish failed; closing open tool calls with buffered state"
168 + );
169 + self.parser_failed = true;
170 + }
171 + }
172 +
173 + Ok(events)
174 + }
175 +}
176 +
177 +/// Push one plain-text delta if it is non-empty.
178 +fn push_text_delta(events: &mut Vec<AssistantEvent>, kind: AssistantBlockKind, delta: String) {
179 + if delta.is_empty() {
180 + return;
181 + }
182 + events.push(AssistantEvent::TextDelta { kind, delta });
183 +}
184 +
185 +/// Wrap one semantic assistant stream into the internal tool-aware assistant
186 +/// stream.
187 +#[try_stream]
188 +pub(crate) async fn tool_event_stream(
189 + stream: impl ContentEventStream,
190 + parser: Option<Box<dyn ToolParser>>,
191 + mut y: TryYielder<AssistantEvent, Error>,
192 +) -> Result<()> {
193 + // Without a parser, pass through the input stream unchanged.
194 + let Some(parser) = parser else {
195 + pin_mut!(stream);
196 + while let Some(event) = stream.next().await.transpose()? {
197 + y.yield_ok(event.into()).await;
198 + }
199 + return Ok(());
200 + };
201 +
202 + pin_mut!(stream);
203 + let mut state = ToolState::new(parser);
204 +
205 + while let Some(event) = stream.next().await.transpose()? {
206 + match event {
207 + ContentEvent::Start {
208 + prompt_token_ids,
209 + prompt_logprobs,
210 + } => {
211 + y.yield_ok(AssistantEvent::Start {
212 + prompt_token_ids,
213 + prompt_logprobs,
214 + })
215 + .await;
216 + }
217 + ContentEvent::TextDelta { kind, delta } => {
218 + for next in state.process_text_delta(kind, delta)? {
219 + y.yield_ok(next).await;
220 + }
221 + }
222 + ContentEvent::LogprobsDelta {
223 + logprobs,
224 + token_ids,
225 + } => {
226 + y.yield_ok(AssistantEvent::LogprobsDelta {
227 + logprobs,
228 + token_ids,
229 + })
230 + .await;
231 + }
232 + ContentEvent::Done {
233 + prompt_token_count,
234 + output_token_count,
235 + finish_reason,
236 + kv_transfer_params,
237 + } => {
238 + for next in state.finish()? {
239 + y.yield_ok(next).await;
240 + }
241 +
242 + y.yield_ok(AssistantEvent::Done {
243 + prompt_token_count,
244 + output_token_count,
245 + finish_reason,
246 + kv_transfer_params,
247 + })
248 + .await;
249 + }
250 + }
251 + }
252 + Ok(())
253 +}
254 +
255 +#[cfg(test)]
256 +mod tests {
257 +
258 + use futures::{StreamExt as _, stream};
259 + use vllm_llm::FinishReason;
260 + use vllm_text::{DecodedLogprobs, DecodedPositionLogprobs, DecodedTokenLogprob};
261 + use vllm_tool_parser::Result;
262 +
263 + use super::super::{AssistantEvent, ContentEvent};
264 + use super::tool_event_stream;
265 + use crate::error::Error;
266 + use crate::event::{AssistantBlockKind, AssistantMessageExt as _};
267 + use crate::output::structured::structured_chat_event_stream;
268 + use crate::parser::tool::{ToolParseResult, ToolParser, ToolParserError};
269 + use crate::request::ChatTool;
270 + use crate::stream::ChatEventStream;
271 +
272 + struct FailingParser {
273 + fail_next: bool,
274 + }
275 +
276 + struct ScriptedParser {
277 + push_results: Vec<ToolParseResult>,
278 + finish_result: ToolParseResult,
279 + }
280 +
281 + impl ToolParser for FailingParser {
282 + fn create(_tools: &[ChatTool]) -> vllm_tool_parser::Result<Box<dyn ToolParser>>
283 + where
284 + Self: Sized + 'static,
285 + {
286 + Ok(Box::new(Self { fail_next: false }))
287 + }
288 +
289 + fn push(&mut self, _chunk: &str) -> Result<ToolParseResult> {
290 + if self.fail_next {
291 + self.fail_next = false;
292 + return Err(ToolParserError::ParsingFailed {
293 + message: "boom".to_string(),
294 + });
295 + }
296 +
297 + Ok(ToolParseResult::default())
298 + }
299 + }
300 +
301 + impl ToolParser for ScriptedParser {
302 + fn create(_tools: &[ChatTool]) -> vllm_tool_parser::Result<Box<dyn ToolParser>>
303 + where
304 + Self: Sized + 'static,
305 + {
306 + Ok(Box::new(Self {
307 + push_results: Vec::new(),
308 + finish_result: ToolParseResult::default(),
309 + }))
310 + }
311 +
312 + fn push(&mut self, _chunk: &str) -> Result<ToolParseResult> {
313 + Ok(self.push_results.pop().unwrap_or_default())
314 + }
315 +
316 + fn finish(&mut self) -> Result<ToolParseResult> {
317 + Ok(std::mem::take(&mut self.finish_result))
318 + }
319 + }
320 +
321 + #[tokio::test]
322 + async fn tool_parser_failure_falls_back_to_plain_text() {
323 + let events = stream::iter(vec![
324 + Ok(ContentEvent::Start {
325 + prompt_token_ids: vec![1, 2, 3].into(),
326 + prompt_logprobs: None,
327 + }),
328 + Ok(ContentEvent::TextDelta {
329 + kind: AssistantBlockKind::Text,
330 + delta: "abc".to_string(),
331 + }),
332 + Ok(ContentEvent::TextDelta {
333 + kind: AssistantBlockKind::Text,
334 + delta: "def".to_string(),
335 + }),
336 + Ok(ContentEvent::Done {
337 + prompt_token_count: 3,
338 + output_token_count: 0,
339 + finish_reason: FinishReason::stop_eos(),
340 + kv_transfer_params: None,
341 + }),
342 + ]);
343 +
344 + let collected =
345 + tool_event_stream(events, Some(Box::new(FailingParser { fail_next: true })))
346 + .collect::<Vec<_>>()
347 + .await;
348 +
349 + let events = collected
350 + .into_iter()
351 + .collect::<crate::Result<Vec<_>>>()
352 + .expect("tool stream should not fail");
353 +
354 + assert_eq!(
355 + events,
356 + vec![
357 + AssistantEvent::Start {
358 + prompt_token_ids: vec![1, 2, 3].into(),
359 + prompt_logprobs: None,
360 + },
361 + AssistantEvent::TextDelta {
362 + kind: AssistantBlockKind::Text,
363 + delta: "abc".to_string(),
364 + },
365 + AssistantEvent::TextDelta {
366 + kind: AssistantBlockKind::Text,
367 + delta: "def".to_string(),
368 + },
369 + AssistantEvent::Done {
370 + prompt_token_count: 3,
371 + output_token_count: 0,
372 + finish_reason: FinishReason::stop_eos(),
373 + kv_transfer_params: None,
374 + },
375 + ]
376 + );
377 +
378 + let message = ChatEventStream::new(
379 + "req_fallback".to_string(),
380 + Box::pin(structured_chat_event_stream(stream::iter(
381 + events.into_iter().map(Ok),
382 + ))),
383 + )
384 + .collect_message()
385 + .await
386 + .expect("collect_message should succeed");
387 + assert_eq!(message.message.text(), "abcdef");
388 + assert!(message.message.tool_calls().next().is_none());
389 + }
390 +
391 + #[tokio::test]
392 + async fn tool_stream_preserves_logprobs_delta() {
393 + let events = stream::iter(vec![
394 + Ok(ContentEvent::Start {
395 + prompt_token_ids: vec![1].into(),
396 + prompt_logprobs: None,
397 + }),
398 + Ok(ContentEvent::LogprobsDelta {
399 + logprobs: Some(DecodedLogprobs {
400 + positions: vec![DecodedPositionLogprobs {
401 + entries: vec![DecodedTokenLogprob {
402 + token_id: 0,
403 + token: "a".to_string(),
404 + logprob: -0.2,
405 + rank: 1,
406 + }],
407 + }],
408 + }),
409 + token_ids: vec![],
410 + }),
411 + Ok(ContentEvent::Done {
412 + prompt_token_count: 1,
413 + output_token_count: 0,
414 + finish_reason: FinishReason::stop_eos(),
415 + kv_transfer_params: None,
416 + }),
417 + ]);
418 + let events = tool_event_stream(events, Some(Box::new(FailingParser { fail_next: false })))
419 + .collect::<Vec<_>>()
420 + .await
421 + .into_iter()
422 + .collect::<crate::Result<Vec<_>>>()
423 + .unwrap();
424 +
425 + assert_eq!(
426 + events,
427 + vec![
428 + AssistantEvent::Start {
429 + prompt_token_ids: vec![1].into(),
430 + prompt_logprobs: None,
431 + },
432 + AssistantEvent::LogprobsDelta {
433 + logprobs: Some(DecodedLogprobs {
434 + positions: vec![DecodedPositionLogprobs {
435 + entries: vec![DecodedTokenLogprob {
436 + token_id: 0,
437 + token: "a".to_string(),
438 + logprob: -0.2,
439 + rank: 1,
440 + }],
441 + }],
442 + }),
443 + token_ids: vec![],
444 + },
445 + AssistantEvent::Done {
446 + prompt_token_count: 1,
447 + output_token_count: 0,
448 + finish_reason: FinishReason::stop_eos(),
449 + kv_transfer_params: None,
450 + },
451 + ]
452 + );
453 + }
454 +
455 + #[tokio::test]
456 + async fn tool_stream_rejects_interleaved_tool_indices() {
457 + let events = stream::iter(vec![
458 + Ok(ContentEvent::TextDelta {
459 + kind: AssistantBlockKind::Text,
460 + delta: "ignored".to_string(),
461 + }),
462 + Ok(ContentEvent::Done {
463 + prompt_token_count: 1,
464 + output_token_count: 1,
465 + finish_reason: FinishReason::stop_eos(),
466 + kv_transfer_params: None,
467 + }),
468 + ]);
469 +
470 + let parser = ScriptedParser {
471 + push_results: vec![ToolParseResult {
472 + normal_text: String::new(),
473 + calls: vec![
474 + crate::parser::tool::ToolCallDelta {
475 + tool_index: 0,
476 + name: Some("first".to_string()),
477 + arguments: String::new(),
478 + },
479 + crate::parser::tool::ToolCallDelta {
480 + tool_index: 1,
481 + name: None,
482 + arguments: "{}".to_string(),
483 + },
484 + ],
485 + }],
486 + finish_result: ToolParseResult::default(),
487 + };
488 +
489 + let err = tool_event_stream(events, Some(Box::new(parser)))
490 + .collect::<Vec<_>>()
491 + .await
492 + .into_iter()
493 + .find_map(|result| result.err())
494 + .expect("expected invariant error");
495 +
496 + assert!(matches!(err, Error::ToolCallStreamInvariant { .. }));
497 + }
498 +
499 + #[tokio::test]
500 + async fn tool_stream_resets_open_tool_when_normal_text_interrupts_it() {
501 + let events = stream::iter(vec![
502 + Ok(ContentEvent::TextDelta {
503 + kind: AssistantBlockKind::Text,
504 + delta: "start".to_string(),
505 + }),
506 + Ok(ContentEvent::TextDelta {
507 + kind: AssistantBlockKind::Text,
508 + delta: "text".to_string(),
509 + }),
510 + Ok(ContentEvent::TextDelta {
511 + kind: AssistantBlockKind::Text,
512 + delta: "args".to_string(),
513 + }),
514 + ]);
515 +
516 + let parser = ScriptedParser {
517 + push_results: vec![
518 + ToolParseResult {
519 + normal_text: String::new(),
520 + calls: vec![crate::parser::tool::ToolCallDelta {
521 + tool_index: 0,
522 + name: None,
523 + arguments: "}".to_string(),
524 + }],
525 + },
526 + ToolParseResult {
527 + normal_text: "plain text".to_string(),
528 + calls: Vec::new(),
529 + },
530 + ToolParseResult {
531 + normal_text: String::new(),
532 + calls: vec![crate::parser::tool::ToolCallDelta {
533 + tool_index: 0,
534 + name: Some("first".to_string()),
535 + arguments: "{".to_string(),
536 + }],
537 + },
538 + ],
539 + finish_result: ToolParseResult::default(),
540 + };
541 +
542 + let err = tool_event_stream(events, Some(Box::new(parser)))
543 + .collect::<Vec<_>>()
544 + .await
545 + .into_iter()
546 + .find_map(|result| result.err())
547 + .expect("expected invariant error");
548 +
549 + assert!(matches!(
550 + err,
551 + Error::ToolCallStreamInvariant { message }
552 + if message == "received arguments for tool index 0 before any tool-call start"
553 + ));
554 + }
555 +
556 + #[tokio::test]
557 + async fn tool_stream_emits_start_and_args_for_terminal_text() {
558 + let events = stream::iter(vec![
559 + Ok(ContentEvent::Start {
560 + prompt_token_ids: vec![1].into(),
561 + prompt_logprobs: None,
562 + }),
563 + Ok(ContentEvent::TextDelta {
564 + kind: AssistantBlockKind::Text,
565 + delta: "ignored".to_string(),
566 + }),
567 + Ok(ContentEvent::Done {
568 + prompt_token_count: 1,
569 + output_token_count: 1,
570 + finish_reason: FinishReason::stop_eos(),
571 + kv_transfer_params: None,
572 + }),
573 + ]);
574 +
575 + let parser = ScriptedParser {
576 + push_results: vec![ToolParseResult {
577 + normal_text: String::new(),
578 + calls: vec![
579 + crate::parser::tool::ToolCallDelta {
580 + tool_index: 0,
581 + name: Some("first".to_string()),
582 + arguments: r#"{"a":1}"#.to_string(),
583 + },
584 + crate::parser::tool::ToolCallDelta {
585 + tool_index: 1,
586 + name: Some("second".to_string()),
587 + arguments: r#"{"b":2}"#.to_string(),
588 + },
589 + ],
590 + }],
591 + finish_result: ToolParseResult::default(),
592 + };
593 +
594 + let events = tool_event_stream(events, Some(Box::new(parser)))
595 + .collect::<Vec<_>>()
596 + .await
597 + .into_iter()
598 + .collect::<crate::Result<Vec<_>>>()
599 + .unwrap();
600 +
601 + assert!(matches!(events[1], AssistantEvent::ToolCallStart { .. }));
602 + assert!(matches!(
603 + events[2],
604 + AssistantEvent::ToolCallArgumentsDelta { .. }
605 + ));
606 + assert!(matches!(events[3], AssistantEvent::ToolCallStart { .. }));
607 + assert!(matches!(
608 + events[4],
609 + AssistantEvent::ToolCallArgumentsDelta { .. }
610 + ));
611 + let collected = ChatEventStream::new(
612 + "req_final_only".to_string(),
613 + Box::pin(structured_chat_event_stream(stream::iter(
614 + events.into_iter().map(Ok),
615 + ))),
616 + )
617 + .collect_message()
618 + .await
619 + .unwrap();
620 + let tool_calls = collected.message.tool_calls().collect::<Vec<_>>();
621 + assert_eq!(tool_calls.len(), 2);
622 + assert_eq!(tool_calls[0].name, "first");
623 + assert_eq!(tool_calls[1].name, "second");
624 + }
625 +}
rust/src/chat/src/output/harmony/mod.rs new
+430
@@ -0,0 +1,430 @@
1 +//! Native Harmony output processing for `gpt_oss`.
2 +//!
3 +//! Unlike the default text-first pipeline, this processor consumes
4 +//! `DecodedTextEvent` token IDs directly and lets the official `openai-harmony`
5 +//! parser recover the structured assistant message shape at token granularity.
6 +
7 +use std::sync::LazyLock;
8 +
9 +use anyhow::Context;
10 +use asynk_strim_attr::{TryYielder, try_stream};
11 +use futures::StreamExt as _;
12 +use openai_harmony::chat::{Content as HarmonyContent, Message as HarmonyMessage, Role};
13 +use openai_harmony::{
14 + HarmonyEncoding, HarmonyEncodingName, StreamableParser, load_harmony_encoding,
15 +};
16 +use thiserror_ext::AsReport;
17 +use vllm_text::output::DecodedTextEvent;
18 +
19 +use crate::Result as ChatResult;
20 +use crate::error::{Error, Result};
21 +use crate::event::AssistantBlockKind;
22 +use crate::output::{
23 + AssistantEvent, ChatOutputProcessor, DynChatEventStream, DynDecodedTextEventStream,
24 + generate_tool_call_id,
25 +};
26 +use crate::parser::ParserSelection;
27 +use crate::request::ChatRequest;
28 +
29 +/// Request-scoped Harmony output processor used for `model_type == "gpt_oss"`.
30 +///
31 +/// This processor keeps the existing northbound `ChatEvent` shape, but swaps
32 +/// the parsed-assistant backend from generic text/reasoning/tool parsers to the
33 +/// official Harmony token parser.
34 +#[derive(Debug)]
35 +pub struct HarmonyChatOutputProcessor {
36 + encoding: &'static HarmonyEncoding,
37 + tool_calls_enabled: bool,
38 +}
39 +
40 +#[derive(Debug, Clone, PartialEq, Eq, Hash)]
41 +struct HarmonyGroupKey {
42 + serial: usize,
43 + channel: Option<String>,
44 + recipient: Option<String>,
45 +}
46 +
47 +#[derive(Debug)]
48 +struct HarmonyGroup {
49 + key: HarmonyGroupKey,
50 + text: String,
51 +}
52 +
53 +#[derive(Debug)]
54 +struct OpenHarmonyToolCall {
55 + recipient: String,
56 +}
57 +
58 +struct HarmonyState {
59 + /// Incremental Harmony parser over assistant token IDs.
60 + parser: StreamableParser,
61 + /// Whether tool-call content should surface as structured tool events.
62 + tool_calls_enabled: bool,
63 + /// Count of completed visible assistant messages for newline insertion.
64 + completed_visible_messages: usize,
65 + /// Count of completed reasoning messages for newline insertion.
66 + completed_reasoning_messages: usize,
67 + /// The current visible text/reasoning group, if any.
68 + current_text_group: Option<HarmonyGroupKey>,
69 + /// The currently open Harmony tool recipient, if any.
70 + open_tool_call: Option<OpenHarmonyToolCall>,
71 +}
72 +
73 +impl HarmonyChatOutputProcessor {
74 + /// Build one request-scoped Harmony processor after backend policy checks.
75 + pub fn new(request: &ChatRequest) -> ChatResult<Self> {
76 + Ok(Self {
77 + encoding: harmony_encoding()?,
78 + tool_calls_enabled: request.tool_parsing_enabled(),
79 + })
80 + }
81 +}
82 +
83 +/// Validate that the generic parser selections are compatible with native
84 +/// Harmony output parsing.
85 +///
86 +/// `gpt_oss` uses a model-specific token-level parser, so any generic
87 +/// reasoning/tool parser override is rejected instead of being silently
88 +/// ignored.
89 +pub(crate) fn validate_harmony_parser_overrides(
90 + tool_call_parser: &ParserSelection,
91 + reasoning_parser: &ParserSelection,
92 +) -> ChatResult<()> {
93 + validate_harmony_override("tool", tool_call_parser)?;
94 + validate_harmony_override("reasoning", reasoning_parser)?;
95 + Ok(())
96 +}
97 +
98 +fn validate_harmony_override(kind: &'static str, selection: &ParserSelection) -> ChatResult<()> {
99 + if matches!(selection, ParserSelection::Auto) {
100 + return Ok(());
101 + }
102 +
103 + Err(Error::HarmonyParserOverrideUnsupported {
104 + kind,
105 + selection: selection.to_string(),
106 + })
107 +}
108 +
109 +impl ChatOutputProcessor for HarmonyChatOutputProcessor {
110 + fn process(self: Box<Self>, decoded: DynDecodedTextEventStream) -> Result<DynChatEventStream> {
111 + let assistant =
112 + harmony_assistant_event_stream(decoded, self.encoding, self.tool_calls_enabled);
113 + Ok(crate::output::structured::structured_chat_event_stream(assistant).boxed())
114 + }
115 +}
116 +
117 +impl HarmonyState {
118 + /// Create one fresh Harmony streaming state for a new assistant response.
119 + fn new(encoding: HarmonyEncoding, tool_calls_enabled: bool) -> Result<Self> {
120 + Ok(Self {
121 + parser: StreamableParser::new(encoding, Some(Role::Assistant))
122 + .map_err(harmony_output_parsing_error)?,
123 + tool_calls_enabled,
124 + completed_visible_messages: 0,
125 + completed_reasoning_messages: 0,
126 + current_text_group: None,
127 + open_tool_call: None,
128 + })
129 + }
130 +
131 + fn process_token_ids(&mut self, token_ids: &[u32]) -> Result<Vec<AssistantEvent>> {
132 + let mut events = Vec::new();
133 + let mut pending_group: Option<HarmonyGroup> = None;
134 +
135 + for &token_id in token_ids {
136 + let completed_before = self.parser.messages().len();
137 + self.parser.process(token_id).map_err(harmony_output_parsing_error)?;
138 + let completed_after = self.parser.messages().len();
139 +
140 + if let Some(delta) = self
141 + .parser
142 + .last_content_delta()
143 + .map_err(harmony_output_parsing_error)?
144 + .filter(|delta| !delta.is_empty())
145 + {
146 + let key = HarmonyGroupKey {
147 + serial: completed_after,
148 + channel: self.parser.current_channel(),
149 + recipient: self.parser.current_recipient(),
150 + };
151 +
152 + match pending_group.as_mut() {
153 + Some(group) if group.key == key => group.text.push_str(&delta),
154 + _ => {
155 + if let Some(group) = pending_group.take() {
156 + self.emit_group(group, &mut events);
157 + }
158 + pending_group = Some(HarmonyGroup { key, text: delta });
159 + }
160 + }
161 + }
162 +
163 + if completed_after > completed_before {
164 + if let Some(group) = pending_group.take() {
165 + self.emit_group(group, &mut events);
166 + }
167 +
168 + for serial in completed_before..completed_after {
169 + let key = {
170 + let message = &self.parser.messages()[serial];
171 + HarmonyGroupKey {
172 + serial,
173 + channel: message.channel.clone(),
174 + recipient: message.recipient.clone(),
175 + }
176 + };
177 + self.handle_completed_message(key);
178 + }
179 + }
180 + }
181 +
182 + if let Some(group) = pending_group {
183 + self.emit_group(group, &mut events);
184 + }
185 +
186 + Ok(events)
187 + }
188 +
189 + /// Flush Harmony parser state at EOS and emit any newly finalized assistant
190 + /// events.
191 + fn process_eos(&mut self) -> Result<Vec<AssistantEvent>> {
192 + let completed_before = self.parser.messages().len();
193 + let pending_key = HarmonyGroupKey {
194 + serial: completed_before,
195 + channel: self.parser.current_channel(),
196 + recipient: self.parser.current_recipient(),
197 + };
198 + let pending_content =
199 + self.parser.current_content().map_err(harmony_output_parsing_error)?;
200 +
201 + self.parser.process_eos().map_err(harmony_output_parsing_error)?;
202 +
203 + let completed_after = self.parser.messages().len();
204 + let mut events = Vec::new();
205 +
206 + if completed_after == completed_before {
207 + return Ok(events);
208 + }
209 +
210 + let final_message = &self.parser.messages()[completed_before];
211 + let final_text = harmony_message_text(final_message);
212 + let tail = final_text.strip_prefix(&pending_content).unwrap_or(final_text).to_string();
213 + if !tail.is_empty() {
214 + self.emit_group(
215 + HarmonyGroup {
216 + key: pending_key,
217 + text: tail,
218 + },
219 + &mut events,
220 + );
221 + }
222 +
223 + for serial in completed_before..completed_after {
224 + let key = {
225 + let message = &self.parser.messages()[serial];
226 + HarmonyGroupKey {
227 + serial,
228 + channel: message.channel.clone(),
229 + recipient: message.recipient.clone(),
230 + }
231 + };
232 + self.handle_completed_message(key);
233 + }
234 +
235 + Ok(events)
236 + }
237 +
238 + /// Flush one coalesced Harmony content group into internal assistant
239 + /// events.
240 + fn emit_group(&mut self, group: HarmonyGroup, events: &mut Vec<AssistantEvent>) {
241 + let channel = group.key.channel.as_deref();
242 + let recipient = group.key.recipient.as_deref();
243 +
244 + if let Some(kind) = text_block_kind(channel, recipient) {
245 + self.open_tool_call = None;
246 +
247 + if self.current_text_group.as_ref() != Some(&group.key) {
248 + let needs_newline = match kind {
249 + AssistantBlockKind::Text => self.completed_visible_messages > 0,
250 + AssistantBlockKind::Reasoning => self.completed_reasoning_messages > 0,
251 + AssistantBlockKind::ToolCall => false,
252 + };
253 +
254 + if needs_newline {
255 + events.push(AssistantEvent::TextDelta {
256 + kind,
257 + delta: "\n".to_string(),
258 + });
259 + }
260 +
261 + self.current_text_group = Some(group.key.clone());
262 + }
263 +
264 + events.push(AssistantEvent::TextDelta {
265 + kind,
266 + delta: group.text,
267 + });
268 + return;
269 + }
270 +
271 + self.current_text_group = None;
272 +
273 + let Some(tool_name) = tool_name(channel, recipient) else {
274 + return;
275 + };
276 + if !self.tool_calls_enabled {
277 + return;
278 + }
279 +
280 + let recipient = recipient.expect("tool groups always have recipient").to_string();
281 + let opens_same_call = match self.open_tool_call.as_ref() {
282 + Some(open_call) => open_call.recipient == recipient,
283 + None => false,
284 + };
285 + if !opens_same_call {
286 + let id = generate_tool_call_id();
287 + self.open_tool_call = Some(OpenHarmonyToolCall { recipient });
288 + events.push(AssistantEvent::ToolCallStart {
289 + id,
290 + name: tool_name.to_string(),
291 + });
292 + }
293 +
294 + if !group.text.is_empty() {
295 + events.push(AssistantEvent::ToolCallArgumentsDelta { delta: group.text });
296 + }
297 + }
298 +
299 + /// Update newline and open-tool state after one Harmony message completes.
300 + fn handle_completed_message(&mut self, key: HarmonyGroupKey) {
301 + if self.current_text_group.as_ref() == Some(&key) {
302 + self.current_text_group = None;
303 + }
304 +
305 + let channel = key.channel.as_deref();
306 + let recipient = key.recipient.as_deref();
307 + let kind = text_block_kind(channel, recipient);
308 +
309 + if kind == Some(AssistantBlockKind::Text) {
310 + self.completed_visible_messages += 1;
311 + } else if kind == Some(AssistantBlockKind::Reasoning) {
312 + self.completed_reasoning_messages += 1;
313 + } else if tool_name(channel, recipient).is_some() {
314 + self.open_tool_call = None;
315 + }
316 + }
317 +}
318 +
319 +/// Convert decoded token updates into internal assistant events with Harmony
320 +/// parsing.
321 +#[try_stream]
322 +async fn harmony_assistant_event_stream(
323 + decoded: DynDecodedTextEventStream,
324 + encoding: &'static HarmonyEncoding,
325 + tool_calls_enabled: bool,
326 + mut y: TryYielder<AssistantEvent, Error>,
327 +) -> Result<()> {
328 + let mut state = HarmonyState::new(encoding.clone(), tool_calls_enabled)?;
329 + futures::pin_mut!(decoded);
330 +
331 + while let Some(event) = decoded.next().await.transpose()? {
332 + match event {
333 + DecodedTextEvent::Start {
334 + prompt_token_ids,
335 + prompt_logprobs,
336 + } => {
337 + y.yield_ok(AssistantEvent::Start {
338 + prompt_token_ids,
339 + prompt_logprobs,
340 + })
341 + .await;
342 + }
343 + DecodedTextEvent::TextDelta {
344 + delta: _, // harmony takes raw token IDs as input, so we ignore text deltas here
345 + token_ids,
346 + logprobs,
347 + finished,
348 + } => {
349 + for event in state.process_token_ids(&token_ids)? {
350 + y.yield_ok(event).await;
351 + }
352 +
353 + if finished.is_some() {
354 + for event in state.process_eos()? {
355 + y.yield_ok(event).await;
356 + }
357 + }
358 +
359 + if logprobs.is_some() || !token_ids.is_empty() {
360 + y.yield_ok(AssistantEvent::LogprobsDelta {
361 + logprobs,
362 + token_ids,
363 + })
364 + .await;
365 + }
366 +
367 + if let Some(finished) = finished {
368 + y.yield_ok(AssistantEvent::Done {
369 + prompt_token_count: finished.prompt_token_count,
370 + output_token_count: finished.output_token_count,
371 + finish_reason: finished.finish_reason,
372 + kv_transfer_params: finished.kv_transfer_params,
373 + })
374 + .await;
375 + }
376 + }
377 + }
378 + }
379 + Ok(())
380 +}
381 +
382 +/// Lazily load the shared GPT-OSS Harmony encoding once per process.
383 +fn harmony_encoding() -> Result<&'static HarmonyEncoding> {
384 + static ENCODING: LazyLock<anyhow::Result<HarmonyEncoding>> = LazyLock::new(|| {
385 + load_harmony_encoding(HarmonyEncodingName::HarmonyGptOss)
386 + .context("failed to load harmony encoding for gpt-oss")
387 + });
388 +
389 + ENCODING.as_ref().map_err(|error| Error::HarmonyOutputParsing {
390 + error: error.to_report_string().into(),
391 + })
392 +}
393 +
394 +fn harmony_output_parsing_error(
395 + error: impl Into<Box<dyn std::error::Error + Send + Sync>>,
396 +) -> Error {
397 + Error::HarmonyOutputParsing {
398 + error: error.into(),
399 + }
400 +}
401 +
402 +/// Return the decoded text payload from one parsed Harmony message.
403 +fn harmony_message_text(message: &HarmonyMessage) -> &str {
404 + let [HarmonyContent::Text(text)] = message.content.as_slice() else {
405 + unreachable!("Harmony parser emits one text content block per parsed message")
406 + };
407 + &text.text
408 +}
409 +
410 +/// Map one Harmony `(channel, recipient)` pair to a visible assistant block
411 +/// kind.
412 +fn text_block_kind(channel: Option<&str>, recipient: Option<&str>) -> Option<AssistantBlockKind> {
413 + match (channel, recipient) {
414 + (Some("final"), _) => Some(AssistantBlockKind::Text),
415 + (Some("analysis"), None) => Some(AssistantBlockKind::Reasoning),
416 + (Some("commentary"), None) => Some(AssistantBlockKind::Text),
417 + _ => None,
418 + }
419 +}
420 +
421 +/// Extract the tool name from a Harmony tool-recipient field, if present.
422 +fn tool_name<'a>(channel: Option<&str>, recipient: Option<&'a str>) -> Option<&'a str> {
423 + match (channel, recipient) {
424 + (Some("commentary" | "analysis"), Some(recipient)) => recipient.strip_prefix("functions."),
425 + _ => None,
426 + }
427 +}
428 +
429 +#[cfg(test)]
430 +mod tests;
rust/src/chat/src/output/harmony/tests.rs new
+351
@@ -0,0 +1,351 @@
1 +//! Harmony output tests share the upstream `openai-harmony` tiktoken cache.
2 +//!
3 +//! Use a file lock for tests that load the encoding so `cargo nextest` cannot
4 +//! start multiple processes that concurrently populate the same cache file.
5 +
6 +use std::sync::Arc;
7 +
8 +use futures::executor::block_on;
9 +use futures::{TryStreamExt as _, stream};
10 +use openai_harmony::chat::{Message, Role};
11 +use serial_test::file_serial;
12 +use vllm_text::output::{DecodedLogprobs, DecodedPositionLogprobs, DecodedTextEvent, Finished};
13 +
14 +use super::*;
15 +use crate::output::ChatOutputProcessor;
16 +use crate::request::{ChatRequest, ChatTool, ChatToolChoice};
17 +use crate::{AssistantMessageExt, ChatEvent, FinishReason};
18 +
19 +fn assistant_prefix() -> Vec<u32> {
20 + harmony_encoding()
21 + .unwrap()
22 + .render_conversation_for_completion(std::iter::empty::<&Message>(), Role::Assistant, None)
23 + .unwrap()
24 +}
25 +
26 +fn completion_tokens(messages: &[Message]) -> Vec<u32> {
27 + let encoding = harmony_encoding().unwrap();
28 + let prefix = assistant_prefix();
29 + let rendered = encoding.render_conversation(messages.iter(), None).unwrap();
30 + assert!(rendered.starts_with(&prefix));
31 + rendered[prefix.len()..].to_vec()
32 +}
33 +
34 +fn text_message(channel: &str, text: &str) -> Message {
35 + Message::from_role_and_content(Role::Assistant, text).with_channel(channel)
36 +}
37 +
38 +fn tool_message(name: &str, arguments: &str, channel: &str) -> Message {
39 + Message::from_role_and_content(Role::Assistant, arguments)
40 + .with_channel(channel)
41 + .with_recipient(format!("functions.{name}"))
42 + .with_content_type("json")
43 +}
44 +
45 +fn decoded_start() -> DecodedTextEvent {
46 + DecodedTextEvent::Start {
47 + prompt_token_ids: Arc::<[u32]>::from([]),
48 + prompt_logprobs: None,
49 + }
50 +}
51 +
52 +fn finished() -> Finished {
53 + Finished {
54 + prompt_token_count: 0,
55 + output_token_count: 0,
56 + finish_reason: FinishReason::stop_eos(),
57 + kv_transfer_params: None,
58 + }
59 +}
60 +
61 +async fn collect_events(
62 + processor: HarmonyChatOutputProcessor,
63 + events: Vec<DecodedTextEvent>,
64 +) -> Vec<ChatEvent> {
65 + Box::new(processor)
66 + .process(Box::pin(stream::iter(events.into_iter().map(Ok))))
67 + .unwrap()
68 + .try_collect()
69 + .await
70 + .unwrap()
71 +}
72 +
73 +fn request_with_tools() -> ChatRequest {
74 + ChatRequest {
75 + tool_choice: ChatToolChoice::Auto,
76 + tools: vec![ChatTool {
77 + name: "get_weather".to_string(),
78 + description: Some("Get weather".to_string()),
79 + parameters: serde_json::json!({
80 + "type": "object",
81 + "properties": {"city": {"type": "string"}},
82 + "required": ["city"]
83 + }),
84 + strict: None,
85 + }],
86 + ..ChatRequest::for_test()
87 + }
88 +}
89 +
90 +#[test]
91 +#[file_serial(harmony_tiktoken_cache)]
92 +fn interrupted_final_message_is_preserved() {
93 + let tokens = completion_tokens(&[text_message("final", "hello")]);
94 + let events = block_on(collect_events(
95 + HarmonyChatOutputProcessor::new(&ChatRequest::for_test()).unwrap(),
96 + vec![
97 + decoded_start(),
98 + DecodedTextEvent::TextDelta {
99 + delta: String::new(),
100 + token_ids: tokens[..tokens.len() - 1].to_vec(),
101 + logprobs: None,
102 + finished: Some(finished()),
103 + },
104 + ],
105 + ));
106 +
107 + assert_eq!(
108 + events.last(),
109 + Some(&ChatEvent::Done {
110 + message: crate::AssistantMessage {
111 + content: vec![crate::AssistantContentBlock::Text {
112 + text: "hello".to_string(),
113 + }],
114 + },
115 + prompt_token_count: 0,
116 + output_token_count: 0,
117 + finish_reason: FinishReason::stop_eos(),
118 + kv_transfer_params: None,
119 + })
120 + );
121 +}
122 +
123 +#[test]
124 +#[file_serial(harmony_tiktoken_cache)]
125 +fn eos_flush_preserves_trailing_replacement_text() {
126 + let mut tokens = completion_tokens(&[text_message("final", "Hi")]);
127 + tokens.pop();
128 + tokens.push(u32::MAX);
129 +
130 + let events = block_on(collect_events(
131 + HarmonyChatOutputProcessor::new(&ChatRequest::for_test()).unwrap(),
132 + vec![
133 + decoded_start(),
134 + DecodedTextEvent::TextDelta {
135 + delta: String::new(),
136 + token_ids: tokens,
137 + logprobs: None,
138 + finished: Some(finished()),
139 + },
140 + ],
141 + ));
142 +
143 + let ChatEvent::Done { message, .. } = events.last().unwrap() else {
144 + panic!("expected done");
145 + };
146 + assert_eq!(message.text(), format!("Hi{}", char::REPLACEMENT_CHARACTER));
147 +}
148 +
149 +#[test]
150 +#[file_serial(harmony_tiktoken_cache)]
151 +fn interrupted_analysis_message_is_preserved() {
152 + let tokens = completion_tokens(&[text_message("analysis", "think")]);
153 + let events = block_on(collect_events(
154 + HarmonyChatOutputProcessor::new(&ChatRequest::for_test()).unwrap(),
155 + vec![
156 + decoded_start(),
157 + DecodedTextEvent::TextDelta {
158 + delta: String::new(),
159 + token_ids: tokens[..tokens.len() - 1].to_vec(),
160 + logprobs: None,
161 + finished: Some(finished()),
162 + },
163 + ],
164 + ));
165 +
166 + assert_eq!(
167 + events.last(),
168 + Some(&ChatEvent::Done {
169 + message: crate::AssistantMessage {
170 + content: vec![crate::AssistantContentBlock::Reasoning {
171 + text: "think".to_string(),
172 + }],
173 + },
174 + prompt_token_count: 0,
175 + output_token_count: 0,
176 + finish_reason: FinishReason::stop_eos(),
177 + kv_transfer_params: None,
178 + })
179 + );
180 +}
181 +
182 +#[test]
183 +#[file_serial(harmony_tiktoken_cache)]
184 +fn commentary_preamble_is_visible_but_commentary_tool_payload_is_not() {
185 + let tokens = completion_tokens(&[
186 + text_message("commentary", "Let me check."),
187 + tool_message("get_weather", r#"{"city":"Paris"}"#, "commentary"),
188 + ]);
189 + let events = block_on(collect_events(
190 + HarmonyChatOutputProcessor::new(&request_with_tools()).unwrap(),
191 + vec![
192 + decoded_start(),
193 + DecodedTextEvent::TextDelta {
194 + delta: String::new(),
195 + token_ids: tokens,
196 + logprobs: None,
197 + finished: Some(finished()),
198 + },
199 + ],
200 + ));
201 +
202 + let done = events.last().unwrap();
203 + let ChatEvent::Done { message, .. } = done else {
204 + panic!("expected done");
205 + };
206 + assert_eq!(message.text(), "Let me check.");
207 + assert_eq!(message.tool_calls().count(), 1);
208 +}
209 +
210 +#[test]
211 +#[file_serial(harmony_tiktoken_cache)]
212 +fn multiple_messages_get_newline_separators() {
213 + let tokens = completion_tokens(&[
214 + text_message("analysis", "first think"),
215 + text_message("analysis", "second think"),
216 + text_message("final", "first answer"),
217 + text_message("final", "second answer"),
218 + ]);
219 + let events = block_on(collect_events(
220 + HarmonyChatOutputProcessor::new(&ChatRequest::for_test()).unwrap(),
221 + vec![
222 + decoded_start(),
223 + DecodedTextEvent::TextDelta {
224 + delta: String::new(),
225 + token_ids: tokens,
226 + logprobs: None,
227 + finished: Some(finished()),
228 + },
229 + ],
230 + ));
231 +
232 + let ChatEvent::Done { message, .. } = events.last().unwrap() else {
233 + panic!("expected done");
234 + };
235 + assert_eq!(
236 + message.reasoning().as_deref(),
237 + Some("first think\nsecond think")
238 + );
239 + assert_eq!(message.text(), "first answer\nsecond answer");
240 +}
241 +
242 +#[test]
243 +#[file_serial(harmony_tiktoken_cache)]
244 +fn tool_calls_stream_arguments_and_finish_with_local_id_shape() {
245 + let tokens = completion_tokens(&[tool_message(
246 + "get_weather",
247 + r#"{"city":"Paris"}"#,
248 + "commentary",
249 + )]);
250 + let midpoint = tokens.len() / 2;
251 + let events = block_on(collect_events(
252 + HarmonyChatOutputProcessor::new(&request_with_tools()).unwrap(),
253 + vec![
254 + decoded_start(),
255 + DecodedTextEvent::TextDelta {
256 + delta: String::new(),
257 + token_ids: tokens[..midpoint].to_vec(),
258 + logprobs: None,
259 + finished: None,
260 + },
261 + DecodedTextEvent::TextDelta {
262 + delta: String::new(),
263 + token_ids: tokens[midpoint..].to_vec(),
264 + logprobs: None,
265 + finished: Some(finished()),
266 + },
267 + ],
268 + ));
269 +
270 + let mut saw_start = None;
271 + let mut saw_args = String::new();
272 + let mut saw_end = None;
273 + for event in &events {
274 + match event {
275 + ChatEvent::ToolCallStart { id, name, .. } => {
276 + assert!(id.starts_with("call_"));
277 + assert_eq!(name, "get_weather");
278 + saw_start = Some(id.clone());
279 + }
280 + ChatEvent::ToolCallArgumentsDelta { delta, .. } => saw_args.push_str(delta),
281 + ChatEvent::ToolCallEnd { call, .. } => {
282 + saw_end = Some(call.clone());
283 + }
284 + _ => {}
285 + }
286 + }
287 +
288 + let start_id = saw_start.expect("tool start");
289 + assert_eq!(saw_args, r#"{"city":"Paris"}"#);
290 + let end = saw_end.expect("tool end");
291 + assert_eq!(end.id, start_id);
292 + assert_eq!(end.arguments, r#"{"city":"Paris"}"#);
293 +}
294 +
295 +#[test]
296 +#[file_serial(harmony_tiktoken_cache)]
297 +fn semantic_events_precede_same_update_logprobs() {
298 + let tokens = completion_tokens(&[text_message("final", "hello")]);
299 + let events = block_on(collect_events(
300 + HarmonyChatOutputProcessor::new(&ChatRequest::for_test()).unwrap(),
301 + vec![
302 + decoded_start(),
303 + DecodedTextEvent::TextDelta {
304 + delta: String::new(),
305 + token_ids: tokens,
306 + logprobs: Some(DecodedLogprobs {
307 + positions: vec![DecodedPositionLogprobs { entries: vec![] }],
308 + }),
309 + finished: Some(finished()),
310 + },
311 + ],
312 + ));
313 +
314 + let block_delta_index = events
315 + .iter()
316 + .position(|event| matches!(event, ChatEvent::BlockDelta { .. }))
317 + .unwrap();
318 + let logprobs_index = events
319 + .iter()
320 + .position(|event| matches!(event, ChatEvent::LogprobsDelta { .. }))
321 + .unwrap();
322 + assert!(block_delta_index < logprobs_index);
323 +}
324 +
325 +#[test]
326 +fn rejects_generic_parser_overrides() {
327 + let reasoning_error =
328 + validate_harmony_parser_overrides(&ParserSelection::Auto, &ParserSelection::None)
329 + .unwrap_err();
330 + assert_eq!(
331 + reasoning_error.to_string(),
332 + "gpt_oss uses native Harmony output parsing; generic reasoning parser override `none` is not supported"
333 + );
334 +
335 + let tool_error = validate_harmony_parser_overrides(
336 + &ParserSelection::Explicit("json".to_string()),
337 + &ParserSelection::Auto,
338 + )
339 + .unwrap_err();
340 + assert_eq!(
341 + tool_error.to_string(),
342 + "gpt_oss uses native Harmony output parsing; generic tool parser override `json` is not supported"
343 + );
344 +}
345 +
346 +#[test]
347 +#[file_serial(harmony_tiktoken_cache)]
348 +fn allows_auto_auto_only() {
349 + validate_harmony_parser_overrides(&ParserSelection::Auto, &ParserSelection::Auto).unwrap();
350 + let _ = HarmonyChatOutputProcessor::new(&ChatRequest::for_test()).unwrap();
351 +}
rust/src/chat/src/output/mod.rs new
+135
@@ -0,0 +1,135 @@
1 +use std::pin::Pin;
2 +use std::sync::Arc;
3 +
4 +use futures::Stream;
5 +use subenum::subenum;
6 +use trait_set::trait_set;
7 +use uuid::Uuid;
8 +use vllm_text::output::{DecodedLogprobs, DecodedPromptLogprobs, DecodedTextEvent};
9 +
10 +use crate::FinishReason;
11 +use crate::error::Result;
12 +use crate::event::{AssistantBlockKind, ChatEvent};
13 +
14 +mod default;
15 +mod harmony;
16 +mod structured;
17 +
18 +pub use default::DefaultChatOutputProcessor;
19 +pub use harmony::HarmonyChatOutputProcessor;
20 +pub(crate) use harmony::validate_harmony_parser_overrides;
21 +
22 +/// Internal assistant event before final assembly.
23 +///
24 +/// - [`ContentEvent`]: subenum after reasoning parsing, carries only text content.
25 +/// - [`AssistantEvent`]: full event after tool parsing, adds tool-call variants.
26 +#[subenum(ContentEvent)]
27 +#[derive(Debug, Clone, PartialEq)]
28 +pub(crate) enum AssistantEvent {
29 + #[subenum(ContentEvent)]
30 + Start {
31 + prompt_token_ids: Arc<[u32]>,
32 + prompt_logprobs: Option<DecodedPromptLogprobs>,
33 + },
34 + #[subenum(ContentEvent)]
35 + TextDelta {
36 + kind: AssistantBlockKind,
37 + delta: String,
38 + },
39 + /// Per-decoded-update sample metadata: logprobs and/or output token IDs.
40 + #[subenum(ContentEvent)]
41 + LogprobsDelta {
42 + logprobs: Option<DecodedLogprobs>,
43 + token_ids: Vec<u32>,
44 + },
45 + /// The start of a new tool call, with its declared name and generated ID.
46 + ToolCallStart { id: String, name: String },
47 + /// A delta for the arguments of the currently open tool call. Must follow a
48 + /// `ToolCallStart`.
49 + ToolCallArgumentsDelta { delta: String },
50 + #[subenum(ContentEvent)]
51 + Done {
52 + prompt_token_count: usize,
53 + output_token_count: usize,
54 + finish_reason: FinishReason,
55 + /// Connector-specific KV transfer parameters for disaggregated serving.
56 + kv_transfer_params: Option<serde_json::Value>,
57 + },
58 +}
59 +
60 +impl ContentEvent {
61 + /// Convert a [`DecodedTextEvent`] into one or more [`ContentEvent`] values
62 + /// by treating all text as plain (non-reasoning) content.
63 + fn from_decoded_plain_text(event: DecodedTextEvent) -> Vec<Self> {
64 + match event {
65 + DecodedTextEvent::Start {
66 + prompt_token_ids,
67 + prompt_logprobs,
68 + } => vec![Self::Start {
69 + prompt_token_ids,
70 + prompt_logprobs,
71 + }],
72 + DecodedTextEvent::TextDelta {
73 + delta,
74 + token_ids,
75 + logprobs,
76 + finished,
77 + } => {
78 + let mut events = Vec::new();
79 + if !delta.is_empty() {
80 + events.push(Self::TextDelta {
81 + kind: AssistantBlockKind::Text,
82 + delta,
83 + });
84 + }
85 + if logprobs.is_some() || !token_ids.is_empty() {
86 + events.push(Self::LogprobsDelta {
87 + logprobs,
88 + token_ids,
89 + });
90 + }
91 + if let Some(finished) = finished {
92 + events.push(Self::Done {
93 + prompt_token_count: finished.prompt_token_count,
94 + output_token_count: finished.output_token_count,
95 + finish_reason: finished.finish_reason,
96 + kv_transfer_params: finished.kv_transfer_params,
97 + });
98 + }
99 + events
100 + }
101 + }
102 + }
103 +}
104 +
105 +/// Boxed stream of decoded text events coming from [`vllm_text`].
106 +pub type DynDecodedTextEventStream = Pin<Box<dyn Stream<Item = Result<DecodedTextEvent>> + Send>>;
107 +/// Boxed stream of structured chat events exposed by [`crate::ChatLlm`].
108 +pub type DynChatEventStream = Pin<Box<dyn Stream<Item = Result<ChatEvent>> + Send>>;
109 +
110 +/// Request-scoped output processor from decoded text events into structured
111 +/// chat events.
112 +pub trait ChatOutputProcessor: Send {
113 + /// Consume decoded text stream and return the structured chat-event stream.
114 + fn process(self: Box<Self>, decoded: DynDecodedTextEventStream) -> Result<DynChatEventStream>;
115 +}
116 +
117 +/// Trait-object form of [`ChatOutputProcessor`].
118 +pub type DynChatOutputProcessor = Box<dyn ChatOutputProcessor>;
119 +
120 +trait_set! {
121 + /// Boxed-stream constraint for decoded text updates.
122 + pub(crate) trait DecodedTextEventStream = Stream<Item = Result<DecodedTextEvent>> + Send + 'static;
123 + /// Boxed-stream constraint for internal assistant events.
124 + pub(crate) trait AssistantEventStream = Stream<Item = Result<AssistantEvent>> + Send + 'static;
125 + /// Boxed-stream constraint for public chat events.
126 + pub(crate) trait ChatEventStream = Stream<Item = Result<ChatEvent>> + Send + 'static;
127 +}
128 +
129 +/// Generate the northbound tool-call ID using the OpenAI-style `call_<id>`
130 +/// format.
131 +// TODO: support other ID scheme like Kimi-K2's
132 +// `functions.{name}:{global_index}`.
133 +pub(crate) fn generate_tool_call_id() -> String {
134 + format!("call_{}", &Uuid::new_v4().simple().to_string()[..24])
135 +}
rust/src/chat/src/output/structured.rs new
+508
@@ -0,0 +1,508 @@
1 +//! Adapts parsed assistant updates into structured chat events.
2 +//!
3 +//! This module remains the final assembly stage in `vllm-chat`. Token-to-text
4 +//! decoding still lives in `decoded.rs`, while reasoning separation and tool
5 +//! parsing are handled earlier by their own adapters. This stage consumes those
6 +//! parsed deltas and assembles higher-level assistant content blocks.
7 +
8 +use asynk_strim_attr::{TryYielder, try_stream};
9 +use futures::{StreamExt as _, pin_mut};
10 +use vllm_text::DecodedLogprobs;
11 +
12 +use super::{AssistantEvent, AssistantEventStream};
13 +use crate::error::Error;
14 +use crate::event::{
15 + AssistantBlockKind, AssistantContentBlock, AssistantMessage, AssistantToolCall, ChatEvent,
16 +};
17 +use crate::{FinishReason, Result};
18 +
19 +/// One currently open assistant text-like block being assembled from streamed
20 +/// deltas.
21 +struct OpenTextBlock {
22 + /// Stable position of this block in the final assistant message.
23 + index: usize,
24 + /// Semantic kind of the block being assembled.
25 + kind: AssistantBlockKind,
26 + /// Accumulated text payload for the block.
27 + text: String,
28 +}
29 +
30 +/// One currently open assistant tool call being assembled from streamed deltas.
31 +struct OpenToolCall {
32 + /// Stable ordinal of this tool call in the assistant tool-call list.
33 + index: usize,
34 + /// Stable tool-call ID exposed northbound.
35 + id: String,
36 + /// Function name.
37 + name: String,
38 + /// Incremental JSON arguments accumulated so far.
39 + arguments: String,
40 +}
41 +
42 +/// Per-stream block assembly state.
43 +///
44 +/// The adapter maintains at most one open text block and one open tool call,
45 +/// and appends deltas to them until the semantic kind changes or the stream
46 +/// terminates.
47 +struct StructuredEventState {
48 + /// Final assistant message assembled so far.
49 + message: AssistantMessage,
50 + /// Currently open text or reasoning block, if any.
51 + open_text_block: Option<OpenTextBlock>,
52 + /// Currently open tool call, if any.
53 + open_tool_call: Option<OpenToolCall>,
54 + /// Next OpenAI-compatible tool-call ordinal.
55 + next_tool_call_index: usize,
56 +}
57 +
58 +impl StructuredEventState {
59 + /// Create one fresh assembly state for a new streamed response.
60 + fn new() -> Self {
61 + Self {
62 + message: AssistantMessage::default(),
63 + open_text_block: None,
64 + open_tool_call: None,
65 + next_tool_call_index: 0,
66 + }
67 + }
68 +
69 + /// Convert one parsed text delta into zero or more structured chat events.
70 + fn process_text_delta(
71 + &mut self,
72 + kind: AssistantBlockKind,
73 + delta: String,
74 + ) -> Result<Vec<ChatEvent>> {
75 + let mut events = Vec::new();
76 + self.close_open_tool_call(&mut events);
77 + self.push_text_delta(kind, delta, &mut events);
78 + Ok(events)
79 + }
80 +
81 + /// Forward per-update sample metadata without attaching it to text blocks.
82 + fn process_logprobs_delta(
83 + &mut self,
84 + logprobs: Option<DecodedLogprobs>,
85 + token_ids: Vec<u32>,
86 + ) -> Result<Vec<ChatEvent>> {
87 + Ok(vec![ChatEvent::LogprobsDelta {
88 + logprobs,
89 + token_ids,
90 + }])
91 + }
92 +
93 + /// Start one new tool call, closing any incompatible open block first.
94 + fn start_tool_call(&mut self, id: String, name: String) -> Result<Vec<ChatEvent>> {
95 + let mut events = Vec::new();
96 + self.close_open_text_block(&mut events);
97 + self.close_open_tool_call(&mut events);
98 +
99 + let index = self.next_tool_call_index;
100 + self.next_tool_call_index += 1;
101 + self.open_tool_call = Some(OpenToolCall {
102 + index,
103 + id: id.clone(),
104 + name: name.clone(),
105 + arguments: String::new(),
106 + });
107 + events.push(ChatEvent::ToolCallStart { index, id, name });
108 + Ok(events)
109 + }
110 +
111 + /// Append one incremental tool-call arguments delta.
112 + fn push_tool_call_arguments(&mut self, delta: String) -> Result<Vec<ChatEvent>> {
113 + let mut events = Vec::new();
114 + let Some(open_tool_call) = self.open_tool_call.as_mut() else {
115 + return Err(Error::ToolCallStreamInvariant {
116 + message: "received tool-call arguments delta without an open tool call".to_string(),
117 + });
118 + };
119 + open_tool_call.arguments.push_str(&delta);
120 + events.push(ChatEvent::ToolCallArgumentsDelta {
121 + index: open_tool_call.index,
122 + delta,
123 + });
124 + Ok(events)
125 + }
126 +
127 + /// Close any open block and emit the terminal `Done` event.
128 + fn finish(
129 + &mut self,
130 + prompt_token_count: usize,
131 + output_token_count: usize,
132 + finish_reason: FinishReason,
133 + kv_transfer_params: Option<serde_json::Value>,
134 + ) -> Result<Vec<ChatEvent>> {
135 + let mut events = Vec::new();
136 + self.close_open_text_block(&mut events);
137 + self.close_open_tool_call(&mut events);
138 + events.push(ChatEvent::Done {
139 + message: self.message.clone(),
140 + prompt_token_count,
141 + output_token_count,
142 + finish_reason,
143 + kv_transfer_params,
144 + });
145 + Ok(events)
146 + }
147 +
148 + /// Append one semantic text delta to the current block, or open a new block
149 + /// when the semantic kind changes.
150 + fn push_text_delta(
151 + &mut self,
152 + kind: AssistantBlockKind,
153 + delta: String,
154 + events: &mut Vec<ChatEvent>,
155 + ) {
156 + if delta.is_empty() {
157 + return;
158 + }
159 +
160 + match self.open_text_block.as_mut() {
161 + // If there's a currently open block of the same kind, append to it.
162 + Some(open_block) if open_block.kind == kind => {
163 + open_block.text.push_str(&delta);
164 + events.push(ChatEvent::BlockDelta {
165 + index: open_block.index,
166 + kind,
167 + delta,
168 + });
169 + }
170 + // Otherwise, close the currently open block (if any) and start a
171 + // new one.
172 + _ => {
173 + self.close_open_text_block(events);
174 + let index = self.message.content.len();
175 + self.open_text_block = Some(OpenTextBlock {
176 + index,
177 + kind,
178 + text: delta.clone(),
179 + });
180 + events.push(ChatEvent::BlockStart { index, kind });
181 + events.push(ChatEvent::BlockDelta { index, kind, delta });
182 + }
183 + }
184 + }
185 +
186 + /// Finalize the currently open text block, if present.
187 + fn close_open_text_block(&mut self, events: &mut Vec<ChatEvent>) {
188 + let Some(open_block) = self.open_text_block.take() else {
189 + return;
190 + };
191 +
192 + let block = match open_block.kind {
193 + AssistantBlockKind::Text => AssistantContentBlock::Text {
194 + text: open_block.text,
195 + },
196 + AssistantBlockKind::Reasoning => AssistantContentBlock::Reasoning {
197 + text: open_block.text,
198 + },
199 + AssistantBlockKind::ToolCall => {
200 + unreachable!("tool calls must not be assembled as text blocks")
201 + }
202 + };
203 + self.message.push_block(block.clone());
204 + events.push(ChatEvent::BlockEnd {
205 + index: open_block.index,
206 + block,
207 + });
208 + }
209 +
210 + /// Finalize the currently open tool call, if present.
211 + fn close_open_tool_call(&mut self, events: &mut Vec<ChatEvent>) {
212 + let Some(open_tool_call) = self.open_tool_call.take() else {
213 + return;
214 + };
215 +
216 + let call = AssistantToolCall {
217 + id: open_tool_call.id,
218 + name: open_tool_call.name,
219 + arguments: open_tool_call.arguments,
220 + };
221 + self.message.push_block(AssistantContentBlock::ToolCall(call.clone()));
222 + events.push(ChatEvent::ToolCallEnd {
223 + index: open_tool_call.index,
224 + call,
225 + });
226 + }
227 +}
228 +
229 +/// Wrap one parsed assistant stream into the public structured chat event
230 +/// stream.
231 +#[try_stream]
232 +pub(crate) async fn structured_chat_event_stream(
233 + stream: impl AssistantEventStream,
234 + mut y: TryYielder<ChatEvent, Error>,
235 +) -> Result<()> {
236 + pin_mut!(stream);
237 +
238 + let mut state = StructuredEventState::new();
239 +
240 + while let Some(event) = stream.next().await.transpose()? {
241 + match event {
242 + AssistantEvent::Start {
243 + prompt_token_ids,
244 + prompt_logprobs,
245 + } => {
246 + y.yield_ok(ChatEvent::Start {
247 + prompt_token_ids,
248 + prompt_logprobs,
249 + })
250 + .await;
251 + }
252 + AssistantEvent::TextDelta { kind, delta } => {
253 + for next in state.process_text_delta(kind, delta)? {
254 + y.yield_ok(next).await;
255 + }
256 + }
257 + AssistantEvent::LogprobsDelta {
258 + logprobs,
259 + token_ids,
260 + } => {
261 + for next in state.process_logprobs_delta(logprobs, token_ids)? {
262 + y.yield_ok(next).await;
263 + }
264 + }
265 + AssistantEvent::ToolCallStart { id, name } => {
266 + for next in state.start_tool_call(id, name)? {
267 + y.yield_ok(next).await;
268 + }
269 + }
270 + AssistantEvent::ToolCallArgumentsDelta { delta } => {
271 + for next in state.push_tool_call_arguments(delta)? {
272 + y.yield_ok(next).await;
273 + }
274 + }
275 + AssistantEvent::Done {
276 + prompt_token_count,
277 + output_token_count,
278 + finish_reason,
279 + kv_transfer_params,
280 + } => {
281 + for next in state.finish(
282 + prompt_token_count,
283 + output_token_count,
284 + finish_reason,
285 + kv_transfer_params,
286 + )? {
287 + y.yield_ok(next).await;
288 + }
289 + }
290 + }
291 + }
292 + Ok(())
293 +}
294 +
295 +#[cfg(test)]
296 +mod tests {
297 + use futures::{StreamExt as _, stream};
298 +
299 + use super::structured_chat_event_stream;
300 + use crate::FinishReason;
301 + use crate::error::Error;
302 + use crate::event::{AssistantBlockKind, AssistantMessageExt as _, ChatEvent};
303 + use crate::output::AssistantEvent;
304 +
305 + #[tokio::test]
306 + async fn structured_stream_closes_tool_call_on_done() {
307 + let events = stream::iter(vec![
308 + Ok(AssistantEvent::ToolCallStart {
309 + id: "call_1".to_string(),
310 + name: "get_weather".to_string(),
311 + }),
312 + Ok(AssistantEvent::ToolCallArgumentsDelta {
313 + delta: r#"{"city":"Paris"}"#.to_string(),
314 + }),
315 + Ok(AssistantEvent::Done {
316 + prompt_token_count: 1,
317 + output_token_count: 1,
318 + finish_reason: FinishReason::stop_eos(),
319 + kv_transfer_params: None,
320 + }),
321 + ]);
322 +
323 + let events = structured_chat_event_stream(events)
324 + .collect::<Vec<_>>()
325 + .await
326 + .into_iter()
327 + .collect::<crate::Result<Vec<_>>>()
328 + .unwrap();
329 +
330 + assert!(matches!(events[0], ChatEvent::ToolCallStart { .. }));
331 + assert!(matches!(
332 + events[1],
333 + ChatEvent::ToolCallArgumentsDelta { .. }
334 + ));
335 + let ChatEvent::ToolCallEnd { call, .. } = &events[2] else {
336 + panic!("expected tool call end");
337 + };
338 + assert_eq!(call.name, "get_weather");
339 + assert_eq!(call.arguments, r#"{"city":"Paris"}"#);
340 + let ChatEvent::Done { message, .. } = &events[3] else {
341 + panic!("expected done");
342 + };
343 + let tool_calls = message.tool_calls().collect::<Vec<_>>();
344 + assert_eq!(tool_calls.len(), 1);
345 + assert_eq!(tool_calls[0].id, "call_1");
346 + assert_eq!(tool_calls[0].arguments, r#"{"city":"Paris"}"#);
347 + }
348 +
349 + #[tokio::test]
350 + async fn structured_stream_closes_previous_tool_call_on_next_start() {
351 + let events = stream::iter(vec![
352 + Ok(AssistantEvent::ToolCallStart {
353 + id: "call_1".to_string(),
354 + name: "first".to_string(),
355 + }),
356 + Ok(AssistantEvent::ToolCallArgumentsDelta {
357 + delta: r#"{"a":1}"#.to_string(),
358 + }),
359 + Ok(AssistantEvent::ToolCallStart {
360 + id: "call_2".to_string(),
361 + name: "second".to_string(),
362 + }),
363 + Ok(AssistantEvent::ToolCallArgumentsDelta {
364 + delta: r#"{"b":2}"#.to_string(),
365 + }),
366 + Ok(AssistantEvent::Done {
367 + prompt_token_count: 1,
368 + output_token_count: 1,
369 + finish_reason: FinishReason::stop_eos(),
370 + kv_transfer_params: None,
371 + }),
372 + ]);
373 +
374 + let events = structured_chat_event_stream(events)
375 + .collect::<Vec<_>>()
376 + .await
377 + .into_iter()
378 + .collect::<crate::Result<Vec<_>>>()
379 + .unwrap();
380 +
381 + assert!(matches!(events[0], ChatEvent::ToolCallStart { .. }));
382 + assert!(matches!(
383 + events[1],
384 + ChatEvent::ToolCallArgumentsDelta { .. }
385 + ));
386 + let ChatEvent::ToolCallEnd { call, .. } = &events[2] else {
387 + panic!("expected first tool call end");
388 + };
389 + assert_eq!(call.name, "first");
390 + assert!(matches!(events[3], ChatEvent::ToolCallStart { .. }));
391 + let ChatEvent::Done { message, .. } = &events[6] else {
392 + panic!("expected done");
393 + };
394 + let tool_calls = message.tool_calls().collect::<Vec<_>>();
395 + assert_eq!(tool_calls.len(), 2);
396 + assert_eq!(tool_calls[0].name, "first");
397 + assert_eq!(tool_calls[1].name, "second");
398 + }
399 +
400 + #[tokio::test]
401 + async fn structured_stream_numbers_tool_calls_independent_of_text_blocks() {
402 + let events = stream::iter(vec![
403 + Ok(AssistantEvent::TextDelta {
404 + kind: AssistantBlockKind::Text,
405 + delta: "before".to_string(),
406 + }),
407 + Ok(AssistantEvent::ToolCallStart {
408 + id: "call_1".to_string(),
409 + name: "get_weather".to_string(),
410 + }),
411 + Ok(AssistantEvent::ToolCallArgumentsDelta {
412 + delta: r#"{"city":"Paris"}"#.to_string(),
413 + }),
414 + Ok(AssistantEvent::Done {
415 + prompt_token_count: 1,
416 + output_token_count: 1,
417 + finish_reason: FinishReason::stop_eos(),
418 + kv_transfer_params: None,
419 + }),
420 + ]);
421 +
422 + let events = structured_chat_event_stream(events)
423 + .collect::<Vec<_>>()
424 + .await
425 + .into_iter()
426 + .collect::<crate::Result<Vec<_>>>()
427 + .unwrap();
428 +
429 + assert!(matches!(
430 + events[0],
431 + ChatEvent::BlockStart {
432 + index: 0,
433 + kind: AssistantBlockKind::Text,
434 + }
435 + ));
436 + assert!(matches!(events[2], ChatEvent::BlockEnd { index: 0, .. }));
437 + assert!(matches!(
438 + events[3],
439 + ChatEvent::ToolCallStart { index: 0, .. }
440 + ));
441 + assert!(matches!(
442 + events[4],
443 + ChatEvent::ToolCallArgumentsDelta { index: 0, .. }
444 + ));
445 + assert!(matches!(events[5], ChatEvent::ToolCallEnd { index: 0, .. }));
446 + }
447 +
448 + #[tokio::test]
449 + async fn structured_stream_closes_tool_call_before_text() {
450 + let events = stream::iter(vec![
451 + Ok(AssistantEvent::ToolCallStart {
452 + id: "call_1".to_string(),
453 + name: "get_weather".to_string(),
454 + }),
455 + Ok(AssistantEvent::ToolCallArgumentsDelta {
456 + delta: r#"{"city":"Paris"}"#.to_string(),
457 + }),
458 + Ok(AssistantEvent::TextDelta {
459 + kind: AssistantBlockKind::Text,
460 + delta: "done".to_string(),
461 + }),
462 + Ok(AssistantEvent::Done {
463 + prompt_token_count: 1,
464 + output_token_count: 1,
465 + finish_reason: FinishReason::stop_eos(),
466 + kv_transfer_params: None,
467 + }),
468 + ]);
469 +
470 + let events = structured_chat_event_stream(events)
471 + .collect::<Vec<_>>()
472 + .await
473 + .into_iter()
474 + .collect::<crate::Result<Vec<_>>>()
475 + .unwrap();
476 +
477 + assert!(matches!(events[2], ChatEvent::ToolCallEnd { .. }));
478 + assert!(matches!(
479 + events[3],
480 + ChatEvent::BlockStart {
481 + kind: AssistantBlockKind::Text,
482 + ..
483 + }
484 + ));
485 + let ChatEvent::Done { message, .. } = &events[6] else {
486 + panic!("expected done");
487 + };
488 + assert_eq!(message.text(), "done");
489 + assert_eq!(message.tool_calls().count(), 1);
490 + }
491 +
492 + #[tokio::test]
493 + async fn structured_stream_rejects_arguments_without_open_tool_call() {
494 + let events = stream::iter(vec![Ok(AssistantEvent::ToolCallArgumentsDelta {
495 + delta: "{}".to_string(),
496 + })]);
497 +
498 + let err = structured_chat_event_stream(events)
499 + .collect::<Vec<_>>()
500 + .await
501 + .into_iter()
502 + .next()
503 + .expect("expected one event")
504 + .expect_err("expected invariant error");
505 +
506 + assert!(matches!(err, Error::ToolCallStreamInvariant { .. }));
507 + }
508 +}
rust/src/chat/src/parser/mod.rs new
+107
@@ -0,0 +1,107 @@
1 +pub mod reasoning;
2 +pub mod tool;
3 +
4 +use std::collections::HashMap;
5 +use std::convert::Infallible;
6 +use std::fmt;
7 +use std::str::FromStr;
8 +
9 +use serde_with::DeserializeFromStr;
10 +
11 +/// Specify which reasoning or tool-call parser implementation to use.
12 +#[derive(Debug, Clone, PartialEq, Eq, Default, DeserializeFromStr)]
13 +pub enum ParserSelection {
14 + /// Use model-based auto-detection.
15 + #[default]
16 + Auto,
17 + /// Disable the parser entirely.
18 + None,
19 + /// Force one specific parser implementation by name.
20 + Explicit(String),
21 +}
22 +
23 +impl ParserSelection {
24 + pub const AUTO_LITERAL: &str = "auto";
25 + pub const NONE_LITERAL: &str = "none";
26 +}
27 +
28 +impl FromStr for ParserSelection {
29 + type Err = Infallible;
30 +
31 + fn from_str(value: &str) -> Result<Self, Self::Err> {
32 + Ok(if value.eq_ignore_ascii_case(Self::AUTO_LITERAL) {
33 + Self::Auto
34 + } else if value.eq_ignore_ascii_case(Self::NONE_LITERAL) {
35 + Self::None
36 + } else {
37 + Self::Explicit(value.to_owned())
38 + })
39 + }
40 +}
41 +
42 +impl fmt::Display for ParserSelection {
43 + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
44 + match self {
45 + Self::Auto => f.write_str(Self::AUTO_LITERAL),
46 + Self::None => f.write_str(Self::NONE_LITERAL),
47 + Self::Explicit(name) => f.write_str(name),
48 + }
49 + }
50 +}
51 +
52 +/// Registry and model matcher for reasoning and tool parsers.
53 +#[derive(Clone)]
54 +pub struct ParserFactory<C> {
55 + creators: HashMap<String, C>,
56 + patterns: Vec<(String, String)>,
57 +}
58 +
59 +impl<C> Default for ParserFactory<C> {
60 + fn default() -> Self {
61 + Self {
62 + creators: HashMap::new(),
63 + patterns: Vec::new(),
64 + }
65 + }
66 +}
67 +
68 +impl<C> ParserFactory<C> {
69 + /// Register a creator for a parser by an exact name.
70 + pub fn register_creator(&mut self, name: &str, creator: C) -> &mut Self {
71 + self.creators.insert(name.to_string(), creator);
72 + self
73 + }
74 +
75 + /// Add a case-insensitive substring match from model ID to parser name.
76 + pub fn register_pattern(&mut self, pattern: &str, parser_name: &str) -> &mut Self {
77 + self.patterns.push((pattern.to_lowercase(), parser_name.to_string()));
78 + self
79 + }
80 +
81 + /// Return the first registered parser name matching the given model ID.
82 + pub fn resolve_name_for_model(&self, model_id: &str) -> Option<&str> {
83 + let model_lower = model_id.to_lowercase();
84 + self.patterns
85 + .iter()
86 + .find(|(pattern, _)| model_lower.contains(pattern))
87 + .map(|(_, parser_name)| parser_name.as_str())
88 + }
89 +
90 + /// Return true if the exact parser name is registered.
91 + pub fn contains(&self, name: &str) -> bool {
92 + self.creators.contains_key(name)
93 + }
94 +
95 + /// Return all registered parser names sorted for stable display.
96 + pub fn list(&self) -> Vec<String> {
97 + let mut names: Vec<_> = self.creators.keys().cloned().collect();
98 + names.sort_unstable();
99 + names
100 + }
101 +
102 + /// Get the constructor for a parser by its exact registered name, or return
103 + /// None if not found.
104 + pub fn creator(&self, name: &str) -> Option<&C> {
105 + self.creators.get(name)
106 + }
107 +}
rust/src/chat/src/parser/reasoning/mod.rs new
+120
@@ -0,0 +1,120 @@
1 +//! Reasoning parser registration and selection boundary for `vllm-chat`.
2 +
3 +use std::sync::LazyLock;
4 +
5 +pub use vllm_reasoning_parser::{
6 + CohereCmdReasoningParser, DeepSeekR1ReasoningParser, DeepSeekV3ReasoningParser,
7 + DeepSeekV4ReasoningParser, Gemma4ReasoningParser, Glm45ReasoningParser, KimiK2ReasoningParser,
8 + KimiReasoningParser, MiniMaxM2ReasoningParser, NemotronV3ReasoningParser, Qwen3ReasoningParser,
9 + ReasoningDelta, ReasoningError, ReasoningParser, Step3ReasoningParser,
10 +};
11 +use vllm_tokenizer::DynTokenizer;
12 +
13 +use crate::parser::ParserFactory;
14 +
15 +/// Canonical public names for registered reasoning parsers.
16 +pub mod names {
17 + pub const COHERE_CMD: &str = "cohere_cmd";
18 + pub const DEEPSEEK_R1: &str = "deepseek_r1";
19 + pub const DEEPSEEK_V3: &str = "deepseek_v3";
20 + pub const DEEPSEEK_V4: &str = "deepseek_v4";
21 + pub const GEMMA4: &str = "gemma4";
22 + pub const GLM45: &str = "glm45";
23 + pub const KIMI: &str = "kimi";
24 + pub const KIMI_K2: &str = "kimi_k2";
25 + pub const MINIMAX_M2: &str = "minimax_m2";
26 + pub const NEMOTRON_V3: &str = "nemotron_v3";
27 + pub const QWEN3: &str = "qwen3";
28 + pub const STEP3: &str = "step3";
29 +}
30 +
31 +/// Constructor signature for one registered reasoning parser implementation.
32 +type ReasoningParserCreator =
33 + fn(DynTokenizer) -> vllm_reasoning_parser::Result<Box<dyn ReasoningParser>>;
34 +
35 +/// Registry and model matcher for reasoning parsers.
36 +pub type ReasoningParserFactory = ParserFactory<ReasoningParserCreator>;
37 +
38 +impl ReasoningParserFactory {
39 + /// Get the global reasoning parser factory with built-in registrations and
40 + /// model mappings.
41 + pub fn global() -> &'static Self {
42 + static INSTANCE: LazyLock<ReasoningParserFactory> =
43 + LazyLock::new(ReasoningParserFactory::new);
44 + &INSTANCE
45 + }
46 +
47 + /// Create the default registry with built-in parser names and model
48 + /// mappings.
49 + pub fn new() -> Self {
50 + let mut factory = Self::default();
51 +
52 + factory
53 + .register_parser::<CohereCmdReasoningParser>(names::COHERE_CMD)
54 + .register_parser::<DeepSeekR1ReasoningParser>(names::DEEPSEEK_R1)
55 + .register_parser::<DeepSeekV3ReasoningParser>(names::DEEPSEEK_V3)
56 + .register_parser::<DeepSeekV4ReasoningParser>(names::DEEPSEEK_V4)
57 + .register_parser::<Gemma4ReasoningParser>(names::GEMMA4)
58 + .register_parser::<Glm45ReasoningParser>(names::GLM45)
59 + .register_parser::<KimiReasoningParser>(names::KIMI)
60 + .register_parser::<KimiK2ReasoningParser>(names::KIMI_K2)
61 + .register_parser::<MiniMaxM2ReasoningParser>(names::MINIMAX_M2)
62 + .register_parser::<NemotronV3ReasoningParser>(names::NEMOTRON_V3)
63 + .register_parser::<Qwen3ReasoningParser>(names::QWEN3)
64 + .register_parser::<Step3ReasoningParser>(names::STEP3);
65 +
66 + factory
67 + .register_pattern("deepseek-r1", names::DEEPSEEK_R1)
68 + .register_pattern("deepseek-v4", names::DEEPSEEK_V4)
69 + .register_pattern("deepseek_v4", names::DEEPSEEK_V4)
70 + .register_pattern("deepseek-v3", names::DEEPSEEK_V3)
71 + .register_pattern("gemma-4", names::GEMMA4)
72 + .register_pattern("gemma4", names::GEMMA4)
73 + .register_pattern("qwen", names::QWEN3)
74 + .register_pattern("glm-5", names::GLM45)
75 + .register_pattern("glm-4.7", names::GLM45)
76 + .register_pattern("glm-4.6", names::GLM45)
77 + .register_pattern("glm-4.5", names::GLM45)
78 + .register_pattern("kimi-k2", names::KIMI_K2)
79 + .register_pattern("kimi", names::KIMI)
80 + .register_pattern("step3", names::STEP3)
81 + .register_pattern("minimax", names::MINIMAX_M2)
82 + .register_pattern("mm-m2", names::MINIMAX_M2)
83 + .register_pattern("cohere", names::COHERE_CMD)
84 + .register_pattern("command", names::COHERE_CMD)
85 + .register_pattern("nano", names::NEMOTRON_V3)
86 + .register_pattern("nemotron", names::NEMOTRON_V3);
87 +
88 + factory
89 + }
90 +
91 + /// Register one parser type that exposes a static `create()` constructor.
92 + pub fn register_parser<T>(&mut self, name: &str) -> &mut Self
93 + where
94 + T: ReasoningParser + 'static,
95 + {
96 + self.register_creator(name, T::create)
97 + }
98 +
99 + /// Construct a parser from an exact name.
100 + pub fn create(
101 + &self,
102 + name: &str,
103 + tokenizer: DynTokenizer,
104 + ) -> crate::Result<Box<dyn ReasoningParser>> {
105 + let creator = self.creator(name).ok_or_else(|| crate::Error::ParserUnavailableByName {
106 + kind: "reasoning",
107 + name: name.to_string(),
108 + available_names: self.list(),
109 + })?;
110 +
111 + creator(tokenizer).map_err(|error| crate::Error::ParserInitialization {
112 + kind: "reasoning",
113 + name: name.to_string(),
114 + error: error.into(),
115 + })
116 + }
117 +}
118 +
119 +#[cfg(test)]
120 +mod tests;
rust/src/chat/src/parser/reasoning/tests.rs new
+61
@@ -0,0 +1,61 @@
1 +use std::sync::Arc;
2 +
3 +use vllm_tokenizer::Tokenizer;
4 +
5 +use super::{ReasoningParserFactory, names};
6 +
7 +struct FakeTokenizer;
8 +
9 +impl Tokenizer for FakeTokenizer {
10 + fn encode(&self, text: &str, _add_special_tokens: bool) -> vllm_tokenizer::Result<Vec<u32>> {
11 + Ok(text.chars().map(u32::from).collect())
12 + }
13 +
14 + fn decode(
15 + &self,
16 + token_ids: &[u32],
17 + _skip_special_tokens: bool,
18 + ) -> vllm_tokenizer::Result<String> {
19 + Ok(token_ids
20 + .iter()
21 + .map(|token_id| char::from_u32(*token_id).unwrap_or('\u{FFFD}'))
22 + .collect())
23 + }
24 +
25 + fn token_to_id(&self, _token: &str) -> Option<u32> {
26 + None
27 + }
28 +}
29 +
30 +#[test]
31 +fn factory_contains_and_lists_registered_parsers() {
32 + let factory = ReasoningParserFactory::new();
33 + assert!(factory.contains(names::QWEN3));
34 + assert!(factory.contains(names::DEEPSEEK_V4));
35 + assert!(factory.list().contains(&names::QWEN3.to_string()));
36 + assert!(factory.list().contains(&names::DEEPSEEK_V4.to_string()));
37 +}
38 +
39 +#[test]
40 +fn factory_resolves_deepseek_v4_to_qwen3_alias() {
41 + let factory = ReasoningParserFactory::new();
42 + assert_eq!(
43 + factory.resolve_name_for_model("deepseek-ai/DeepSeek-V4"),
44 + Some(names::DEEPSEEK_V4)
45 + );
46 + assert_eq!(
47 + factory.resolve_name_for_model("deepseek_v4"),
48 + Some(names::DEEPSEEK_V4)
49 + );
50 +}
51 +
52 +#[test]
53 +fn factory_rejects_unknown_parser_names() {
54 + let tokenizer = Arc::new(FakeTokenizer);
55 + let factory = ReasoningParserFactory::new();
56 + let error = match factory.create("missing", tokenizer) {
57 + Ok(_) => panic!("expected parser lookup to fail"),
58 + Err(error) => error,
59 + };
60 + assert!(error.to_string().contains("choose from"));
61 +}
rust/src/chat/src/parser/tool/mod.rs new
+140
@@ -0,0 +1,140 @@
1 +//! Tool parser registration and selection boundary for `vllm-chat`.
2 +
3 +use std::sync::LazyLock;
4 +
5 +pub use vllm_tool_parser::{
6 + DeepSeekV3ToolParser, DeepSeekV4ToolParser, DeepSeekV31ToolParser, DeepSeekV32ToolParser,
7 + Gemma4ToolParser, Glm45MoeToolParser, Glm47MoeToolParser, HermesToolParser, KimiK2ToolParser,
8 + Llama3JsonToolParser, MinimaxM2ToolParser, MistralToolParser, Qwen3CoderToolParser,
9 + Qwen3XmlToolParser, ToolCallDelta, ToolParseResult, ToolParser, ToolParserError,
10 +};
11 +
12 +use crate::parser::ParserFactory;
13 +use crate::request::ChatTool;
14 +
15 +/// Canonical public names for registered tool parsers.
16 +pub mod names {
17 + pub const DEEPSEEK_V3: &str = "deepseek_v3";
18 + pub const DEEPSEEK_V31: &str = "deepseek_v31";
19 + pub const DEEPSEEK_V32: &str = "deepseek_v32";
20 + pub const DEEPSEEK_V4: &str = "deepseek_v4";
21 + pub const GLM45: &str = "glm45";
22 + pub const GLM47: &str = "glm47";
23 + pub const GEMMA4: &str = "gemma4";
24 + pub const HERMES: &str = "hermes";
25 + pub const KIMI_K2: &str = "kimi_k2";
26 + pub const LLAMA3_JSON: &str = "llama3_json";
27 + pub const LLAMA4_JSON: &str = "llama4_json";
28 + pub const MINIMAX_M2: &str = "minimax_m2";
29 + pub const MISTRAL: &str = "mistral";
30 + pub const QWEN3_CODER: &str = "qwen3_coder";
31 + pub const QWEN3_XML: &str = "qwen3_xml";
32 +}
33 +
34 +/// Constructor signature for one registered tool parser implementation.
35 +type ToolParserCreator = fn(&[ChatTool]) -> vllm_tool_parser::Result<Box<dyn ToolParser>>;
36 +
37 +/// Registry and model matcher for tool parsers.
38 +pub type ToolParserFactory = ParserFactory<ToolParserCreator>;
39 +
40 +impl ToolParserFactory {
41 + /// Get the global tool parser factory with built-in registrations and model
42 + /// mappings.
43 + pub fn global() -> &'static Self {
44 + static INSTANCE: LazyLock<ToolParserFactory> = LazyLock::new(ToolParserFactory::new);
45 + &INSTANCE
46 + }
47 +
48 + /// Create the default registry with built-in parser names and model
49 + /// mappings.
50 + pub fn new() -> Self {
51 + let mut factory = Self::default();
52 +
53 + factory
54 + .register_parser::<DeepSeekV3ToolParser>(names::DEEPSEEK_V3)
55 + .register_parser::<DeepSeekV31ToolParser>(names::DEEPSEEK_V31)
56 + .register_parser::<DeepSeekV32ToolParser>(names::DEEPSEEK_V32)
57 + .register_parser::<DeepSeekV4ToolParser>(names::DEEPSEEK_V4)
58 + .register_parser::<Glm45MoeToolParser>(names::GLM45)
59 + .register_parser::<Glm47MoeToolParser>(names::GLM47)
60 + .register_parser::<Gemma4ToolParser>(names::GEMMA4)
61 + .register_parser::<HermesToolParser>(names::HERMES)
62 + .register_parser::<KimiK2ToolParser>(names::KIMI_K2)
63 + .register_parser::<Llama3JsonToolParser>(names::LLAMA3_JSON)
64 + .register_parser::<Llama3JsonToolParser>(names::LLAMA4_JSON)
65 + .register_parser::<MinimaxM2ToolParser>(names::MINIMAX_M2)
66 + .register_parser::<MistralToolParser>(names::MISTRAL)
67 + .register_parser::<Qwen3XmlToolParser>(names::QWEN3_XML)
68 + .register_parser::<Qwen3CoderToolParser>(names::QWEN3_CODER);
69 +
70 + factory
71 + .register_pattern("mistral-", names::MISTRAL)
72 + .register_pattern("mixtral-", names::MISTRAL)
73 + .register_pattern("qwen3-coder", names::QWEN3_CODER)
74 + .register_pattern("qwen2.5-coder", names::QWEN3_CODER)
75 + .register_pattern("qwen3.5", names::QWEN3_CODER)
76 + .register_pattern("qwen", names::QWEN3_XML)
77 + .register_pattern("hermes", names::HERMES)
78 + .register_pattern("llama-4", names::LLAMA4_JSON)
79 + .register_pattern("llama-3.2", names::LLAMA3_JSON)
80 + .register_pattern("llama-3.1", names::LLAMA3_JSON)
81 + .register_pattern("deepseek-r1", names::DEEPSEEK_V3)
82 + .register_pattern("deepseek-v4", names::DEEPSEEK_V4)
83 + .register_pattern("deepseek_v4", names::DEEPSEEK_V4)
84 + .register_pattern("deepseek-v3.2", names::DEEPSEEK_V32)
85 + .register_pattern("deepseek-v3.1", names::DEEPSEEK_V31)
86 + .register_pattern("deepseek-v3", names::DEEPSEEK_V3)
87 + .register_pattern("glm-5", names::GLM47)
88 + .register_pattern("glm-4.7", names::GLM47)
89 + .register_pattern("glm-4.6", names::GLM45)
90 + .register_pattern("glm-4.5", names::GLM45)
91 + .register_pattern("gemma4", names::GEMMA4)
92 + .register_pattern("gemma-4", names::GEMMA4)
93 + .register_pattern("kimi-k2", names::KIMI_K2)
94 + .register_pattern("minimax", names::MINIMAX_M2)
95 + .register_pattern("mm-m2", names::MINIMAX_M2);
96 +
97 + factory
98 + }
99 +
100 + /// Register one parser type that exposes a static `create()` constructor.
101 + pub fn register_parser<T>(&mut self, name: &str) -> &mut Self
102 + where
103 + T: ToolParser + 'static,
104 + {
105 + self.register_creator(name, T::create)
106 + }
107 +
108 + /// Construct a parser from an exact name.
109 + pub fn create(&self, name: &str, tools: &[ChatTool]) -> crate::Result<Box<dyn ToolParser>> {
110 + let creator = self.creator(name).ok_or_else(|| crate::Error::ParserUnavailableByName {
111 + kind: "tool",
112 + name: name.to_string(),
113 + available_names: self.list(),
114 + })?;
115 +
116 + creator(tools).map_err(|error| crate::Error::ParserInitialization {
117 + kind: "tool",
118 + name: name.to_string(),
119 + error: error.into(),
120 + })
121 + }
122 +
123 + /// Resolve a parser from model ID and then construct it.
124 + pub fn create_for_model(
125 + &self,
126 + model_id: &str,
127 + tools: &[ChatTool],
128 + ) -> crate::Result<Box<dyn ToolParser>> {
129 + let name = self.resolve_name_for_model(model_id).ok_or_else(|| {
130 + crate::Error::ParserUnavailableForModel {
131 + kind: "tool",
132 + model_id: model_id.to_string(),
133 + }
134 + })?;
135 + self.create(name, tools)
136 + }
137 +}
138 +
139 +#[cfg(test)]
140 +mod tests;
rust/src/chat/src/parser/tool/tests.rs new
+152
@@ -0,0 +1,152 @@
1 +use vllm_tool_parser::Result;
2 +
3 +use super::{ToolParseResult, ToolParser, ToolParserFactory, names};
4 +use crate::Error;
5 +use crate::request::ChatTool;
6 +
7 +struct FakeToolParser;
8 +
9 +impl ToolParser for FakeToolParser {
10 + fn create(_tools: &[ChatTool]) -> Result<Box<dyn ToolParser>>
11 + where
12 + Self: Sized + 'static,
13 + {
14 + Ok(Box::new(Self))
15 + }
16 +
17 + fn preserve_special_tokens(&self) -> bool {
18 + true
19 + }
20 +
21 + fn push(&mut self, _chunk: &str) -> Result<ToolParseResult> {
22 + Ok(ToolParseResult::default())
23 + }
24 +}
25 +
26 +#[test]
27 +fn default_factory_starts_empty() {
28 + let factory = ToolParserFactory::default();
29 + assert!(factory.list().is_empty());
30 +}
31 +
32 +#[test]
33 +fn factory_contains_and_creates_registered_parsers() {
34 + let mut factory = ToolParserFactory::default();
35 + factory.register_parser::<FakeToolParser>("fake");
36 +
37 + assert!(factory.contains("fake"));
38 + assert!(factory.list().contains(&"fake".to_string()));
39 + factory.create("fake", &[]).unwrap();
40 +}
41 +
42 +#[test]
43 +fn factory_rejects_unknown_parser_names() {
44 + let factory = ToolParserFactory::default();
45 + let error = match factory.create("missing", &[]) {
46 + Ok(_) => panic!("expected parser lookup to fail"),
47 + Err(error) => error,
48 + };
49 + assert!(matches!(error, Error::ParserUnavailableByName { .. }));
50 +}
51 +
52 +#[test]
53 +fn factory_rejects_unknown_models() {
54 + let factory = ToolParserFactory::default();
55 + let error = match factory.create_for_model("definitely-unknown-model", &[]) {
56 + Ok(_) => panic!("expected model lookup to fail"),
57 + Err(error) => error,
58 + };
59 + assert!(matches!(error, Error::ParserUnavailableForModel { .. }));
60 +}
61 +
62 +#[test]
63 +fn factory_creates_registered_parser_for_model() {
64 + let mut factory = ToolParserFactory::default();
65 + factory
66 + .register_parser::<FakeToolParser>("fake")
67 + .register_pattern("fake-model", "fake");
68 +
69 + factory.create_for_model("my-fake-model-v1", &[]).unwrap();
70 +}
71 +
72 +#[test]
73 +fn factory_new_resolves_default_patterns() {
74 + let factory = ToolParserFactory::new();
75 +
76 + assert_eq!(
77 + factory.resolve_name_for_model("Qwen/Qwen3.5-0.8B"),
78 + Some(names::QWEN3_CODER)
79 + );
80 + assert_eq!(
81 + factory.resolve_name_for_model("Qwen/Qwen3-0.6B"),
82 + Some(names::QWEN3_XML)
83 + );
84 + assert_eq!(
85 + factory.resolve_name_for_model("Qwen/Qwen3-Coder-30B"),
86 + Some(names::QWEN3_CODER)
87 + );
88 + assert_eq!(
89 + factory.resolve_name_for_model("meta-llama-4-maverick"),
90 + Some(names::LLAMA4_JSON)
91 + );
92 + assert_eq!(
93 + factory.resolve_name_for_model("meta-llama-3.2-3b-instruct"),
94 + Some(names::LLAMA3_JSON)
95 + );
96 + assert_eq!(
97 + factory.resolve_name_for_model("meta-llama/Llama-3.1-8B-Instruct"),
98 + Some(names::LLAMA3_JSON)
99 + );
100 + assert_eq!(
101 + factory.resolve_name_for_model("deepseek-ai/DeepSeek-V4"),
102 + Some(names::DEEPSEEK_V4)
103 + );
104 + assert_eq!(
105 + factory.resolve_name_for_model("deepseek-ai/DeepSeek-V3.2-Exp"),
106 + Some(names::DEEPSEEK_V32)
107 + );
108 + assert_eq!(
109 + factory.resolve_name_for_model("deepseek-ai/DeepSeek-V4-Chat"),
110 + Some(names::DEEPSEEK_V4)
111 + );
112 + assert_eq!(
113 + factory.resolve_name_for_model("deepseek_v4"),
114 + Some(names::DEEPSEEK_V4)
115 + );
116 + assert_eq!(
117 + factory.resolve_name_for_model("deepseek-ai/DeepSeek-R1-0528"),
118 + Some(names::DEEPSEEK_V3)
119 + );
120 + assert_eq!(
121 + factory.resolve_name_for_model("deepseek-ai/DeepSeek-V3.1"),
122 + Some(names::DEEPSEEK_V31)
123 + );
124 + assert_eq!(
125 + factory.resolve_name_for_model("zai-org/GLM-5-32B-Chat"),
126 + Some(names::GLM47)
127 + );
128 + assert_eq!(
129 + factory.resolve_name_for_model("zai-org/GLM-5.1-32B-Instruct"),
130 + Some(names::GLM47)
131 + );
132 + assert_eq!(
133 + factory.resolve_name_for_model("glm-4.7"),
134 + Some(names::GLM47)
135 + );
136 + assert_eq!(
137 + factory.resolve_name_for_model("google/gemma-4-27b-it"),
138 + Some(names::GEMMA4)
139 + );
140 + assert_eq!(
141 + factory.resolve_name_for_model("NousResearch/Hermes-3-Llama-3.1-8B"),
142 + Some(names::HERMES)
143 + );
144 + assert_eq!(
145 + factory.resolve_name_for_model("MiniMax/MiniMax-M2-01"),
146 + Some(names::MINIMAX_M2)
147 + );
148 + assert_eq!(
149 + factory.resolve_name_for_model("org/mm-m2-base"),
150 + Some(names::MINIMAX_M2)
151 + );
152 +}
rust/src/chat/src/renderer/deepseek_v32/encoding.rs new
+555
@@ -0,0 +1,555 @@
1 +//! DeepSeek V3.2 prompt renderer.
2 +
3 +use std::collections::{HashMap, HashSet};
4 +use std::fmt::Write as _;
5 +
6 +use serde::Serialize;
7 +use serde_json::Value;
8 +use serde_json_fmt::JsonFormat;
9 +
10 +use crate::error::{Error, Result};
11 +use crate::request::{ChatContent, ChatMessage, ChatRequest, ChatRole, ChatTool};
12 +use crate::{AssistantContentBlock, AssistantMessageExt, AssistantToolCall};
13 +
14 +const BOS_TOKEN: &str = "<｜begin▁of▁sentence｜>";
15 +const EOS_TOKEN: &str = "<｜end▁of▁sentence｜>";
16 +const THINKING_START_TOKEN: &str = "<think>";
17 +const THINKING_END_TOKEN: &str = "</think>";
18 +const DSML_TOKEN: &str = "｜DSML｜";
19 +
20 +/// DeepSeek uses `"chat"` vs `"thinking"` mode names. Keep the split explicit
21 +/// here so the render branches stay easy to read.
22 +#[derive(Debug, Clone, Copy, PartialEq, Eq)]
23 +enum ThinkingMode {
24 + Chat,
25 + Thinking,
26 +}
27 +
28 +/// Tool schema shape rendered inside the `<functions>` block.
29 +#[serde_with::skip_serializing_none]
30 +#[derive(Debug, Serialize)]
31 +struct RenderedToolSchema<'a> {
32 + name: &'a str,
33 + description: Option<&'a str>,
34 + parameters: &'a Value,
35 + strict: Option<bool>,
36 +}
37 +
38 +/// Render one chat request into the final prompt string.
39 +pub(super) fn render_request(request: &ChatRequest) -> Result<String> {
40 + let thinking_mode = match request.enable_thinking()?.unwrap_or(false) {
41 + true => ThinkingMode::Thinking,
42 + false => ThinkingMode::Chat,
43 + };
44 + let drop_thinking = matches!(
45 + request.messages.last().map(ChatMessage::role),
46 + Some(ChatRole::User | ChatRole::Developer)
47 + );
48 + let render_offset = isize::from(request.tool_parsing_enabled());
49 + let last_user_render_index =
50 + find_last_user_render_index(request.messages.as_slice(), render_offset);
51 + let last_user_actual_index = find_last_user_actual_index(request.messages.as_slice());
52 + let mut prompt = String::from(BOS_TOKEN);
53 +
54 + if request.tool_parsing_enabled() {
55 + render_system_message(&mut prompt, None, &request.tools)?;
56 + }
57 +
58 + for (message_index, message) in request.messages.iter().enumerate() {
59 + render_message(
60 + &mut prompt,
61 + request.messages.as_slice(),
62 + message_index,
63 + message,
64 + render_offset,
65 + last_user_render_index,
66 + last_user_actual_index,
67 + thinking_mode,
68 + drop_thinking,
69 + )?;
70 + }
71 +
72 + Ok(prompt)
73 +}
74 +
75 +/// Find the last user-like turn in render order.
76 +///
77 +/// `render_offset` is `1` when a synthetic tool-only system turn is rendered
78 +/// before the real request messages, and `0` otherwise.
79 +fn find_last_user_render_index(messages: &[ChatMessage], render_offset: isize) -> isize {
80 + messages
81 + .iter()
82 + .rposition(|message| matches!(message.role(), ChatRole::User | ChatRole::Developer))
83 + .map(|index| index as isize + render_offset)
84 + .unwrap_or(-1)
85 +}
86 +
87 +/// Render one real request message, using `render_offset` to account for any
88 +/// synthetic tool-only system turn that was already emitted before the loop.
89 +fn render_message(
90 + out: &mut String,
91 + messages: &[ChatMessage],
92 + message_index: usize,
93 + message: &ChatMessage,
94 + render_offset: isize,
95 + last_user_render_index: isize,
96 + last_user_actual_index: usize,
97 + thinking_mode: ThinkingMode,
98 + drop_thinking: bool,
99 +) -> Result<()> {
100 + let render_index = message_index as isize + render_offset;
101 + let opens_thinking = render_index == last_user_render_index;
102 + let after_last_user_turn = render_index > last_user_render_index;
103 + let after_or_at_last_user_turn = render_index >= last_user_render_index;
104 +
105 + match message {
106 + ChatMessage::System { content } => render_system_message(out, Some(content), &[]),
107 + ChatMessage::Developer { content, tools } => render_developer_message(
108 + out,
109 + content,
110 + tools.as_deref().unwrap_or(&[]),
111 + thinking_mode == ThinkingMode::Thinking && opens_thinking,
112 + ),
113 + ChatMessage::User { content } => render_user_message(
114 + out,
115 + content,
116 + thinking_mode == ThinkingMode::Thinking && opens_thinking,
117 + ),
118 + ChatMessage::Assistant { content } => render_assistant_message(
119 + out,
120 + thinking_mode == ThinkingMode::Thinking && after_last_user_turn,
121 + content,
122 + should_keep_assistant_reasoning(
123 + message_index,
124 + last_user_actual_index,
125 + thinking_mode,
126 + drop_thinking,
127 + ),
128 + // TODO: Respect `continue_final_message` and map it to DeepSeek's
129 + // prefix-style final-assistant continuation behavior.
130 + false,
131 + ),
132 + ChatMessage::ToolResponse { content, .. } => render_tool_message(
133 + out,
134 + messages,
135 + message_index,
136 + thinking_mode == ThinkingMode::Thinking && after_or_at_last_user_turn,
137 + content,
138 + ),
139 + }
140 +}
141 +
142 +/// Historical assistant reasoning is dropped in thinking mode when the final
143 +/// request turn is a new user-like message.
144 +fn should_keep_assistant_reasoning(
145 + actual_index: usize,
146 + last_user_actual_index: usize,
147 + thinking_mode: ThinkingMode,
148 + drop_thinking: bool,
149 +) -> bool {
150 + !(thinking_mode == ThinkingMode::Thinking
151 + && drop_thinking
152 + && actual_index < last_user_actual_index)
153 +}
154 +
155 +/// Return the last user/developer turn in the real request message list.
156 +fn find_last_user_actual_index(messages: &[ChatMessage]) -> usize {
157 + messages
158 + .iter()
159 + .rposition(|message| matches!(message.role(), ChatRole::User | ChatRole::Developer))
160 + .unwrap_or(usize::MAX)
161 +}
162 +
163 +/// Render a system turn, optionally followed by the tool preamble.
164 +fn render_system_message(
165 + out: &mut String,
166 + content: Option<&ChatContent>,
167 + tools: &[ChatTool],
168 +) -> Result<()> {
169 + if let Some(content) = content {
170 + write_chat_content(out, content)?;
171 + }
172 + if !tools.is_empty() {
173 + out.push_str("\n\n");
174 + render_tools(out, tools)?;
175 + }
176 + Ok(())
177 +}
178 +
179 +/// Developer messages are wrapped into the same user-like turn shape as real
180 +/// user messages, but can also carry message-local tools.
181 +fn render_developer_message(
182 + out: &mut String,
183 + content: &ChatContent,
184 + tools: &[ChatTool],
185 + opens_thinking: bool,
186 +) -> Result<()> {
187 + if content.is_empty() {
188 + return Err(Error::ChatTemplate(
189 + "invalid DeepSeek V3.2 developer message: empty content".to_string(),
190 + ));
191 + }
192 +
193 + out.push_str("<｜User｜>");
194 + if !tools.is_empty() {
195 + out.push_str("\n\n");
196 + render_tools(out, tools)?;
197 + }
198 + out.push_str("\n\n# The user's message is: ");
199 + write_chat_content(out, content)?;
200 + write_user_like_suffix(out, opens_thinking);
201 + Ok(())
202 +}
203 +
204 +/// Plain user turns share the same wrapper shape as developer turns without the
205 +/// developer-specific preamble.
206 +fn render_user_message(
207 + out: &mut String,
208 + content: &ChatContent,
209 + opens_thinking: bool,
210 +) -> Result<()> {
211 + out.push_str("<｜User｜>");
212 + write_chat_content(out, content)?;
213 + write_user_like_suffix(out, opens_thinking);
214 + Ok(())
215 +}
216 +
217 +/// Shared trailing wrapper used by both real user turns and native developer
218 +/// turns after their content has already been written.
219 +// TODO: respect `add_generation_prompt` option
220 +fn write_user_like_suffix(out: &mut String, opens_thinking: bool) {
221 + out.push_str("<｜Assistant｜>");
222 + if opens_thinking {
223 + out.push_str(THINKING_START_TOKEN);
224 + } else {
225 + out.push_str(THINKING_END_TOKEN);
226 + }
227 +}
228 +
229 +/// Render one tool result turn and decide whether it opens or closes the shared
230 +/// `<function_results>` block for the preceding assistant tool-call message.
231 +fn render_tool_message(
232 + out: &mut String,
233 + messages: &[ChatMessage],
234 + message_index: usize,
235 + resumes_thinking: bool,
236 + _content: &ChatContent,
237 +) -> Result<()> {
238 + let (block_start, block_end) = tool_response_block_bounds(messages, message_index);
239 + if message_index != block_start {
240 + return Ok(());
241 + }
242 +
243 + let Some(prev_assistant_idx) = previous_assistant_actual_index(messages, block_start) else {
244 + return Err(Error::ChatTemplate(
245 + "invalid DeepSeek V3.2 tool message: missing previous assistant message".to_string(),
246 + ));
247 + };
248 +
249 + let ChatMessage::Assistant {
250 + content: assistant_content,
251 + } = &messages[prev_assistant_idx]
252 + else {
253 + return Err(Error::ChatTemplate(
254 + "invalid DeepSeek V3.2 tool message: previous non-tool message is not assistant"
255 + .to_string(),
256 + ));
257 + };
258 +
259 + let assistant_tool_calls = assistant_content.tool_calls().collect::<Vec<_>>();
260 + if assistant_tool_calls.is_empty() {
261 + return Err(Error::ChatTemplate(
262 + "invalid DeepSeek V3.2 tool message: previous assistant message has no tool calls"
263 + .to_string(),
264 + ));
265 + }
266 +
267 + let mut expected_tool_call_ids = HashSet::with_capacity(assistant_tool_calls.len());
268 + for tool_call in &assistant_tool_calls {
269 + if !expected_tool_call_ids.insert(tool_call.id.as_str()) {
270 + return Err(Error::ChatTemplate(
271 + "invalid DeepSeek V3.2 assistant tool calls: duplicate tool_call_id".to_string(),
272 + ));
273 + }
274 + }
275 +
276 + let mut tool_results_by_id = HashMap::with_capacity(assistant_tool_calls.len());
277 + for message in &messages[block_start..block_end] {
278 + let ChatMessage::ToolResponse {
279 + content,
280 + tool_call_id,
281 + } = message
282 + else {
283 + unreachable!("tool response block should only contain tool messages");
284 + };
285 +
286 + if !expected_tool_call_ids.contains(tool_call_id.as_str()) {
287 + return Err(Error::ChatTemplate(format!(
288 + "invalid DeepSeek V3.2 tool message: unknown tool_call_id `{tool_call_id}`"
289 + )));
290 + }
291 +
292 + if tool_results_by_id.insert(tool_call_id.as_str(), content).is_some() {
293 + return Err(Error::ChatTemplate(format!(
294 + "invalid DeepSeek V3.2 tool message: duplicate tool_call_id `{tool_call_id}`"
295 + )));
296 + }
297 + }
298 +
299 + if tool_results_by_id.len() != assistant_tool_calls.len() {
300 + return Err(Error::ChatTemplate(
301 + "invalid DeepSeek V3.2 tool messages: missing tool result for assistant tool call"
302 + .to_string(),
303 + ));
304 + }
305 +
306 + out.push_str("\n\n<function_results>");
307 + for tool_call in assistant_tool_calls {
308 + let content = tool_results_by_id
309 + .get(tool_call.id.as_str())
310 + .expect("validated tool_call_id set should be complete");
311 + out.push_str("\n<result>");
312 + write_chat_content(out, content)?;
313 + out.push_str("</result>");
314 + }
315 +
316 + out.push_str("\n</function_results>");
317 + out.push_str("\n\n");
318 + if resumes_thinking {
319 + out.push_str(THINKING_START_TOKEN);
320 + } else {
321 + out.push_str(THINKING_END_TOKEN);
322 + }
323 +
324 + Ok(())
325 +}
326 +
327 +/// Return the contiguous tool-response block containing `actual_index`.
328 +fn tool_response_block_bounds(messages: &[ChatMessage], actual_index: usize) -> (usize, usize) {
329 + let mut block_start = actual_index;
330 + while block_start > 0 && matches!(messages[block_start - 1], ChatMessage::ToolResponse { .. }) {
331 + block_start -= 1;
332 + }
333 +
334 + let mut block_end = actual_index + 1;
335 + while block_end < messages.len()
336 + && matches!(messages[block_end], ChatMessage::ToolResponse { .. })
337 + {
338 + block_end += 1;
339 + }
340 +
341 + (block_start, block_end)
342 +}
343 +
344 +/// Return the most recent assistant turn before `actual_index`.
345 +fn previous_assistant_actual_index(messages: &[ChatMessage], actual_index: usize) -> Option<usize> {
346 + messages[..actual_index]
347 + .iter()
348 + .rposition(|message| matches!(message, ChatMessage::Assistant { .. }))
349 +}
350 +
351 +/// Render one assistant turn, including optional reasoning, DSML tool calls,
352 +/// and the trailing EOS marker.
353 +fn render_assistant_message(
354 + out: &mut String,
355 + after_last_user_turn: bool,
356 + content: &[AssistantContentBlock],
357 + keep_reasoning: bool,
358 + prefix: bool,
359 +) -> Result<()> {
360 + let has_reasoning = keep_reasoning && content.has_reasoning();
361 + let has_tool_calls = content.has_tool_calls();
362 +
363 + if !has_tool_calls && prefix {
364 + write_assistant_text(out, content);
365 + return Ok(());
366 + }
367 +
368 + if after_last_user_turn {
369 + if !has_reasoning && !has_tool_calls {
370 + return Err(Error::ChatTemplate(
371 + "invalid DeepSeek V3.2 assistant message after last user message: expected reasoning or tool calls"
372 + .to_string(),
373 + ));
374 + }
375 +
376 + if has_reasoning {
377 + write_assistant_reasoning(out, content);
378 + }
379 + out.push_str(THINKING_END_TOKEN);
380 + }
381 +
382 + write_assistant_text(out, content);
383 +
384 + if has_tool_calls {
385 + out.push_str("\n\n<｜DSML｜function_calls>\n");
386 + for (index, tool_call) in content.tool_calls().enumerate() {
387 + if index > 0 {
388 + out.push('\n');
389 + }
390 + render_tool_call(out, tool_call)?;
391 + }
392 + out.push_str("\n</｜DSML｜function_calls>");
393 + }
394 +
395 + out.push_str(EOS_TOKEN);
396 + Ok(())
397 +}
398 +
399 +/// Render one assistant tool call in DSML XML-like format.
400 +fn render_tool_call(out: &mut String, tool_call: &AssistantToolCall) -> Result<()> {
401 + writeln!(out, "<{DSML_TOKEN}invoke name=\"{}\">", tool_call.name)
402 + .expect("writing to String cannot fail");
403 + encode_arguments_to_dsml(out, tool_call)?;
404 + write!(out, "\n</{DSML_TOKEN}invoke>").expect("writing to String cannot fail");
405 + Ok(())
406 +}
407 +
408 +/// Convert one assistant tool-call arguments object into DSML parameter form.
409 +///
410 +/// String values are emitted raw with `string="true"`, while all other JSON
411 +/// values are rendered with JSON syntax and `string="false"`.
412 +fn encode_arguments_to_dsml(out: &mut String, tool_call: &AssistantToolCall) -> Result<()> {
413 + let arguments: Value = serde_json::from_str(&tool_call.arguments).map_err(|error| {
414 + Error::ChatTemplate(format!(
415 + "assistant tool call has invalid JSON arguments for DeepSeek V3.2: {error}"
416 + ))
417 + })?;
418 + let Some(arguments) = arguments.as_object() else {
419 + return Err(Error::ChatTemplate(
420 + "assistant tool call arguments for DeepSeek V3.2 must be a JSON object".to_string(),
421 + ));
422 + };
423 +
424 + let mut wrote_parameter = false;
425 + for (key, value) in arguments {
426 + if wrote_parameter {
427 + out.push('\n');
428 + }
429 +
430 + let is_string = matches!(value, Value::String(_));
431 + write!(
432 + out,
433 + "<{DSML_TOKEN}parameter name=\"{key}\" string=\"{}\">",
434 + if is_string { "true" } else { "false" }
435 + )
436 + .expect("writing to String cannot fail");
437 +
438 + match value {
439 + Value::String(value) => out.push_str(value),
440 + value => out.push_str(&json_dumps(value)?),
441 + }
442 +
443 + write!(out, "</{DSML_TOKEN}parameter>").expect("writing to String cannot fail");
444 + wrote_parameter = true;
445 + }
446 +
447 + Ok(())
448 +}
449 +
450 +/// Render the full tool preamble shown to the model.
451 +fn render_tools(out: &mut String, tools: &[ChatTool]) -> Result<()> {
452 + out.push_str(
453 + r#"## Tools
454 +
455 +You have access to a set of tools you can use to answer the user's question.
456 +You can invoke functions by writing a "<｜DSML｜function_calls>" block like the following as part of your reply to the user:
457 +<｜DSML｜function_calls>
458 +<｜DSML｜invoke name="$FUNCTION_NAME">
459 +<｜DSML｜parameter name="$PARAMETER_NAME" string="true|false">$PARAMETER_VALUE</｜DSML｜parameter>
460 +...
461 +</｜DSML｜invoke>
462 +<｜DSML｜invoke name="$FUNCTION_NAME2">
463 +...
464 +</｜DSML｜invoke>
465 +</｜DSML｜function_calls>
466 +
467 +String and scalar parameters should be specified as is without any escaping or quotes, while lists and objects should use JSON format. The "string" attribute should be set to "true" for string type parameters and "false" for other types (numbers, booleans, arrays, objects).
468 +
469 +If the thinking_mode is enabled, then after function results you should strongly consider outputting a thinking block. Here is an example:
470 +
471 +<｜DSML｜function_calls>
472 +...
473 +</｜DSML｜function_calls>
474 +
475 +<function_results>
476 +...
477 +</function_results>
478 +
479 +<think>...thinking about results</think>
480 +
481 +Here are the functions available in JSONSchema format:
482 +<functions>
483 +"#,
484 + );
485 +
486 + for (index, tool) in tools.iter().enumerate() {
487 + if index > 0 {
488 + out.push('\n');
489 + }
490 + render_tool_schema(out, tool)?;
491 + }
492 +
493 + out.push_str("\n</functions>\n");
494 + Ok(())
495 +}
496 +
497 +/// Serialize one typed tool schema into the JSON shape embedded inside
498 +/// `<functions>`.
499 +fn render_tool_schema(out: &mut String, tool: &ChatTool) -> Result<()> {
500 + out.push_str(&json_dumps(&RenderedToolSchema {
501 + name: &tool.name,
502 + description: tool.description.as_deref(),
503 + parameters: &tool.parameters,
504 + strict: tool.strict,
505 + })?);
506 + Ok(())
507 +}
508 +
509 +/// Write chat content directly into the destination buffer without flattening
510 +/// it into an intermediate `String`.
511 +fn write_chat_content(out: &mut String, content: &ChatContent) -> Result<()> {
512 + match content {
513 + ChatContent::Text(text) => out.push_str(text),
514 + ChatContent::Parts(parts) => {
515 + for part in parts {
516 + out.push_str(part.as_text()?);
517 + }
518 + }
519 + }
520 + Ok(())
521 +}
522 +
523 +/// Write all reasoning blocks in encounter order.
524 +fn write_assistant_reasoning(out: &mut String, content: &[AssistantContentBlock]) {
525 + for block in content {
526 + if let AssistantContentBlock::Reasoning { text } = block {
527 + out.push_str(text);
528 + }
529 + }
530 +}
531 +
532 +/// Write all visible assistant text blocks in encounter order.
533 +fn write_assistant_text(out: &mut String, content: &[AssistantContentBlock]) {
534 + for block in content {
535 + if let AssistantContentBlock::Text { text } = block {
536 + out.push_str(text);
537 + }
538 + }
539 +}
540 +
541 +/// Compact JSON serialization used by this renderer for exact prompt text.
542 +fn json_dumps<T: Serialize>(value: &T) -> Result<String> {
543 + JsonFormat::new()
544 + .comma(", ")
545 + .expect("literal comma separator is valid JSON")
546 + .colon(": ")
547 + .expect("literal colon separator is valid JSON")
548 + .ascii(false)
549 + .format_to_string(value)
550 + .map_err(|error| {
551 + Error::ChatTemplate(format!(
552 + "failed to serialize DeepSeek V3.2 JSON payload: {error}"
553 + ))
554 + })
555 +}
rust/src/chat/src/renderer/deepseek_v32/fixtures/test_input.json new
+149
@@ -0,0 +1,149 @@
1 +{
2 + "tools": [
3 + {
4 + "type": "function",
5 + "function": {
6 + "name": "get_datetime",
7 + "description": "Get the current date and time",
8 + "parameters": {
9 + "type": "object",
10 + "properties": {
11 + "timezone": {
12 + "type": "string",
13 + "description": "The timezone, e.g. Asia/Shanghai, UTC"
14 + }
15 + },
16 + "required": ["timezone"]
17 + }
18 + }
19 + },
20 + {
21 + "type": "function",
22 + "function": {
23 + "name": "get_weather",
24 + "description": "Get the weather for a specific date and location",
25 + "parameters": {
26 + "type": "object",
27 + "properties": {
28 + "location": {
29 + "type": "string",
30 + "description": "The city name, e.g. Beijing, Hangzhou"
31 + },
32 + "date": {
33 + "type": "string",
34 + "description": "The date in YYYY-MM-DD format"
35 + }
36 + },
37 + "required": ["location", "date"]
38 + }
39 + }
40 + }
41 + ],
42 + "messages": [
43 + {
44 + "role": "system",
45 + "content": "You are a helpful Assistant."
46 + },
47 + {
48 + "role": "user",
49 + "content": "明天杭州和北京的天气怎么样？"
50 + },
51 + {
52 + "role": "assistant",
53 + "reasoning_content": "用户询问明天的天气，我需要先获取当前日期来计算明天的日期📅",
54 + "tool_calls": [
55 + {
56 + "id": "call_xK9mN3pL2qR8vT5wY6hZ1aB4",
57 + "type": "function",
58 + "function": {
59 + "arguments": "{\"timezone\": \"Asia/Shanghai\"}",
60 + "name": "get_datetime"
61 + }
62 + }
63 + ]
64 + },
65 + {
66 + "tool_call_id": "call_xK9mN3pL2qR8vT5wY6hZ1aB4",
67 + "role": "tool",
68 + "content": "{\"current_date\": \"2024-01-15\", \"current_time\": \"14:30:00\", \"timezone\": \"Asia/Shanghai\"}"
69 + },
70 + {
71 + "role": "assistant",
72 + "reasoning_content": "现在知道今天是2024-01-15，明天就是2024-01-16。接下来查询杭州和北京明天的天气🌤️",
73 + "tool_calls": [
74 + {
75 + "id": "call_bN7kR9mX3pQ2wL5vY8jZ4cD6",
76 + "type": "function",
77 + "function": {
78 + "arguments": "{\"location\": \"Hangzhou\", \"date\": \"2024-01-16\"}",
79 + "name": "get_weather"
80 + }
81 + },
82 + {
83 + "id": "call_dP9mL7kX5rT4yN3wZ2hV8eF1",
84 + "type": "function",
85 + "function": {
86 + "arguments": "{\"location\": \"Beijing\", \"date\": \"2024-01-16\"}",
87 + "name": "get_weather"
88 + }
89 + }
90 + ]
91 + },
92 + {
93 + "tool_call_id": "call_bN7kR9mX3pQ2wL5vY8jZ4cD6",
94 + "role": "tool",
95 + "content": "{\"location\": \"Hangzhou\", \"date\": \"2024-01-16\", \"temperature_high\": \"12\", \"temperature_low\": \"5\", \"weather\": \"多云\", \"humidity\": \"65%\"}"
96 + },
97 + {
98 + "tool_call_id": "call_dP9mL7kX5rT4yN3wZ2hV8eF1",
99 + "role": "tool",
100 + "content": "{\"location\": \"Beijing\", \"date\": \"2024-01-16\", \"temperature_high\": \"-2\", \"temperature_low\": \"-8\", \"weather\": \"晴\", \"humidity\": \"30%\"}"
101 + },
102 + {
103 + "role": "assistant",
104 + "reasoning_content": "已获取两个城市明天的天气信息，现在整理给用户✨",
105 + "content": "根据查询结果，明天（2024年1月16日）的天气情况如下：\n\n**杭州**：\n- 天气：多云\n- 最高温度：12°C\n- 最低温度：5°C\n- 湿度：65%\n\n**北京**：\n- 天气：晴\n- 最高温度：-2°C\n- 最低温度：-8°C\n- 湿度：30%\n\n杭州明天会比较温暖但有些多云，而北京会很冷但是晴天。建议在北京的朋友要注意保暖！"
106 + },
107 + {
108 + "role": "user",
109 + "content": "谢谢！那后天呢？"
110 + },
111 + {
112 + "role": "assistant",
113 + "reasoning_content": "用户现在问后天的天气，后天是2024-01-17，我可以直接查询（因为已知今天日期）🗓️",
114 + "tool_calls": [
115 + {
116 + "id": "call_fR3nK8mV7pL4xT2yW9jB5gH3",
117 + "type": "function",
118 + "function": {
119 + "arguments": "{\"location\": \"Hangzhou\", \"date\": \"2024-01-17\"}",
120 + "name": "get_weather"
121 + }
122 + },
123 + {
124 + "id": "call_hT5pN2kY9rV6zL3wX1mD7jK8",
125 + "type": "function",
126 + "function": {
127 + "arguments": "{\"location\": \"Beijing\", \"date\": \"2024-01-17\"}",
128 + "name": "get_weather"
129 + }
130 + }
131 + ]
132 + },
133 + {
134 + "tool_call_id": "call_fR3nK8mV7pL4xT2yW9jB5gH3",
135 + "role": "tool",
136 + "content": "{\"location\": \"Hangzhou\", \"date\": \"2024-01-17\", \"temperature_high\": \"15\", \"temperature_low\": \"8\", \"weather\": \"小雨\", \"humidity\": \"80%\"}"
137 + },
138 + {
139 + "tool_call_id": "call_hT5pN2kY9rV6zL3wX1mD7jK8",
140 + "role": "tool",
141 + "content": "{\"location\": \"Beijing\", \"date\": \"2024-01-17\", \"temperature_high\": \"0\", \"temperature_low\": \"-6\", \"weather\": \"多云\", \"humidity\": \"45%\"}"
142 + },
143 + {
144 + "role": "assistant",
145 + "reasoning_content": "获取到后天的天气数据，整理回复给用户📝",
146 + "content": "后天（2024年1月17日）的天气情况：\n\n**杭州**：\n- 天气：小雨\n- 最高温度：15°C\n- 最低温度：8°C\n- 湿度：80%\n\n**北京**：\n- 天气：多云\n- 最高温度：0°C\n- 最低温度：-6°C\n- 湿度：45%\n\n杭州后天会有小雨，温度略有回升，记得带伞。北京会稍微暖和一点，但依然很冷，请继续做好保暖措施。"
147 + }
148 + ]
149 +}
rust/src/chat/src/renderer/deepseek_v32/fixtures/test_input_search_w_date.json new
+732
@@ -0,0 +1,732 @@
1 +{
2 + "messages": [
3 + {
4 + "role": "developer",
5 + "content": "帮我调研一下，目前有哪些针对search agent的benchmark？详细介绍各自的特点、使用场景、例题。\n\n<system-reminder>\n## Today’s Date\n2025-11-27, Thursday.\n</system-reminder>",
6 + "tools": [
7 + {
8 + "type": "function",
9 + "function": {
10 + "name": "search",
11 + "description": "Searches for information related to query and displays topn results.",
12 + "parameters": {
13 + "type": "object",
14 + "properties": {
15 + "query": {
16 + "type": "string",
17 + "description": "The search query string"
18 + },
19 + "topn": {
20 + "type": "integer",
21 + "description": "Number of top results to display",
22 + "default": 10
23 + },
24 + "source": {
25 + "type": "string",
26 + "description": "Source to search within",
27 + "enum": [
28 + "web",
29 + "news"
30 + ],
31 + "default": "web"
32 + }
33 + },
34 + "required": [
35 + "query"
36 + ],
37 + "additionalProperties": false,
38 + "$schema": "http://json-schema.org/draft-07/schema#"
39 + }
40 + }
41 + },
42 + {
43 + "type": "function",
44 + "function": {
45 + "name": "open",
46 + "description": "Opens the link id from the page indicated by cursor starting at line number loc, showing num_lines lines. Valid link ids are displayed with the formatting: 【{id}†.*】. If cursor is not provided, the most recent page is implied. If id is a string, it is treated as a fully qualified URL associated with source. If loc is not provided, the viewport will be positioned at the beginning of the document or centered on the most relevant passage, if available. Use this function without id to scroll to a new location of an opened page.",
47 + "parameters": {
48 + "type": "object",
49 + "properties": {
50 + "id": {
51 + "description": "Link ID (number) or fully qualified URL (string)",
52 + "anyOf": [
53 + {
54 + "type": "integer"
55 + },
56 + {
57 + "type": "string"
58 + }
59 + ],
60 + "default": -1
61 + },
62 + "cursor": {
63 + "type": "integer",
64 + "description": "Page cursor indicator",
65 + "default": -1
66 + },
67 + "loc": {
68 + "type": "integer",
69 + "description": "Starting line number",
70 + "default": -1
71 + },
72 + "num_lines": {
73 + "type": "integer",
74 + "description": "Number of lines to show",
75 + "default": -1
76 + },
77 + "view_source": {
78 + "type": "boolean",
79 + "description": "Whether to view source",
80 + "default": false
81 + },
82 + "source": {
83 + "type": "string",
84 + "description": "Source associated with the URL",
85 + "enum": [
86 + "web",
87 + "news"
88 + ],
89 + "default": "web"
90 + }
91 + },
92 + "additionalProperties": false,
93 + "$schema": "http://json-schema.org/draft-07/schema#"
94 + }
95 + }
96 + },
97 + {
98 + "type": "function",
99 + "function": {
100 + "name": "find",
101 + "description": "Finds exact matches of pattern in the current page, or the page given by cursor.",
102 + "parameters": {
103 + "type": "object",
104 + "properties": {
105 + "pattern": {
106 + "type": "string",
107 + "description": "The exact pattern to search for"
108 + },
109 + "cursor": {
110 + "type": "integer",
111 + "description": "Page cursor to search in",
112 + "default": -1
113 + }
114 + },
115 + "required": [
116 + "pattern"
117 + ],
118 + "additionalProperties": false,
119 + "$schema": "http://json-schema.org/draft-07/schema#"
120 + }
121 + }
122 + }
123 + ]
124 + },
125 + {
126 + "role": "assistant",
127 + "content": "",
128 + "reasoning_content": "用户想调研目前针对搜索智能体（search agent）的基准测试（benchmark），要求详细介绍各自的特点、使用场景和例题。我需要提供一份全面的概述，涵盖现有的基准测试，例如WebGPT、WebShop、ALFWorld、HotpotQA、FEVER、TriviaQA、Natural Questions、SearchQA、MS MARCO、BEIR、LoTTE、MIRACL、Zero-shot IR基准等。同时，需要区分这些基准测试是针对检索任务、问答任务，还是具体针对搜索智能体的评估。此外，还应考虑那些专门为评估搜索智能体（即能够与环境交互、执行搜索、浏览网页、使用工具的智能体）而设计的基准测试，例如WebGPT、WebShop、ALFWorld、Mind2Web、WebArena、VisualWebArena、SearchBench、WebVLN、WebNav等。我需要详细介绍每个基准测试的特点、使用场景，并举例说明。为了确保信息的全面性和时效性，我将同时进行两项搜索：一项是针对“search agent benchmark”的通用搜索，以获取广泛的基准测试信息；另一项是针对“搜索智能体 基准测试”的中文搜索，以获取更符合用户语言习惯和可能更具体的中文资料。",
129 + "tool_calls": [
130 + {
131 + "id": "fixture-tool-call-1-0",
132 + "type": "function",
133 + "function": {
134 + "name": "search",
135 + "arguments": "{\"query\": \"search agent benchmark 2024\", \"topn\": 10, \"source\": \"web\"}"
136 + }
137 + },
138 + {
139 + "id": "fixture-tool-call-1-1",
140 + "type": "function",
141 + "function": {
142 + "name": "search",
143 + "arguments": "{\"query\": \"搜索智能体 基准测试\", \"topn\": 10, \"source\": \"web\"}"
144 + }
145 + }
146 + ]
147 + },
148 + {
149 + "role": "tool",
150 + "content": "[0] search agent benchmark 2024 (Search_Results/search agent benchmark 2024)\n**viewing lines [0 - 39] of 39**\n\nL0: \nL1: URL: Search_Results/search agent benchmark 2024\nL2: # Search Results\nL3: \nL4: * 【0†HSCodeComp: A Realistic and Expert-level Benchmark for ...; publish_date:\nL5: none†arxiv.org】 Oct 22, 2025 — To fill this gap, we introduce HSCodeComp, the \nL6: first realistic, expert-level e-commerce benchmark designed to evaluate deep \nL7: search agents in ...\nL8: * 【1†open-compass/GTA - A Benchmark for General Tool Agents; publish_date: \nL9: none†github.com】 GTA is a benchmark to evaluate the tool-use capability of LLM-\nL10: based agents in real-world scenarios. It features three main aspects.\nL11: * 【2†Benchmarking real-time trust scoring across five AI Agent ...; \nL12: publish_date: none†cleanlab.ai】 Aug 20, 2025 — This article evaluates 5 AI Agent\nL13: architectures over the BOLAA (ICLR 2024) benchmark, and assesses the effects of\nL14: adding automated trust ...\nL15: * 【3†10 AI agent benchmarks; publish_date: none†www.evidentlyai.com】 Jul 11, \nL16: 2025 — We put together 10 AI agent benchmarks designed to assess how well \nL17: different LLMs perform as agents in real-world scenarios, ...\nL18: * 【4†A state-of-the-art search API purpose-built for agents; publish_date: \nL19: none†parallel.ai】 Jul 31, 2025 — To evaluate real-world performance of the \nL20: Parallel Search MCP Server, we created the WISER-Search benchmark which blends \nL21: WISER-Fresh (queries ...\nL22: * 【5†AI Agent Benchmarks are Broken; publish_date: none†medium.com】 We break \nL23: down the failure modes in current AI agent benchmarks and introduce a checklist \nL24: that minimizes the gamability of AI agent benchmarks.\nL25: * 【6†Benchmarks and Tree Search for Multimodal LLM Web Agents; publish_date: \nL26: none†dpfried.github.io】 2024, When is Tree Search Useful? ○ Dealing with \nL27: destructive actions. ○ Some things on the web are very difficult to undo, e.g., \nL28: ordering an item. 56.\nL29: * 【7†-Bench: Benchmarking AI agents for the real-world; publish_date: \nL30: none†sierra.ai】 Jun 20, 2024 — τ-bench measures an agent's ability to interact \nL31: with (simulated) human users and programmatic APIs while following domain-\nL32: specific policies in a consistent ...\nL33: * 【8†Browser Use = state of the art Web Agent; publish_date: none†browser-\nL34: use.com】 Dec 15, 2024 — Browser Use has achieved state-of-the-art performance on\nL35: the WebVoyager benchmark, with an impressive 89.1% success rate across 586 \nL36: diverse web tasks.\nL37: * 【9†FutureSearch Benchmarks; publish_date: none†evals.futuresearch.ai】 Find \nL38: the original source of a given claim. Example: From , more than 8 out of 1000 \nL39: users clicked on a phishing link monthly in 2024, up 190% vs 2023.",
151 + "tool_call_id": "fixture-tool-call-1-0"
152 + },
153 + {
154 + "role": "tool",
155 + "content": "[1] 搜索智能体 基准测试 (Search_Results/搜索智能体 基准测试)\n**viewing lines [0 - 33] of 33**\n\nL0: \nL1: URL: Search_Results/搜索智能体 基准测试\nL2: # Search Results\nL3: \nL4: * 【0†WideSearch：揭示AI 智能体缺失的「广度」能力; publish_date: none†zhuanlan.zhihu.com】 Aug \nL5: 16, 2025 — 为系统评估智能体在该任务上的能力，论文构建了第一个专门的基准测试 WideSearch ，包含200 个源于真实世界、横跨18 \nL6: 个领域的高质量任务。 通过对超过10 个 ...\nL7: * 【1†GAIA: 一个严苛的智能体基准- HuggingFace; publish_date: none†www.cnblogs.com】 Jul 9,\nL8: 2024 — 我们使用一个用库构建的代码智能体 在GAIA 基准上进行测试，这可以说是最困难、最全面的智能体基准测试……最终我们取得了第一名的成绩！ \nL9: GAIA: 一个严苛的 ...\nL10: * 【2†AI搜索智能体遭遇新挑战：滑铁卢大学团队提出更公平透明的 ...; publish_date: none†www.techwalker.com】 \nL11: Aug 14, 2025 — \nL12: 目前评测AI搜索智能体主要依靠BrowseComp这样的基准测试，它就像一场实时的开卷考试，让AI在真实的网络环境中搜索信息来回答复杂问题。听起来很合理 ...\nL13: * 【3†Agentic AI基础设施实践经验系列（六）：Agent质量评估 - AWS; publish_date: \nL14: none†aws.amazon.com】 Sep 19, 2025 — TAU-bench \nL15: 是一个评估AI智能体在真实世界环境中可靠性的基准测试。它评估智能体是否能够在动态的多轮对话中与用户进行交互，理解需求并完成任务。T-bench ...\nL16: * 【4†DeepAgent：能自己找工具的通用推理智能体 - 高瓴人工智能学院; publish_date: none†ai.ruc.edu.cn】 \nL17: Nov 6, 2025 — 在八大基准测试中，DeepAgent在绝大多数任务上全面领先所有基线模型。 \nL18: 开放环境优势:在最具挑战的“开放工具检索”场景下(如ToolBench)，其成功率达到64%，远 ...\nL19: * 【5†BrowseComp：为浏览智能体设计的简单而具挑战性的基准测试; publish_date: none†blog.csdn.net】 Sep \nL20: 22, 2025 — 该基准测试由OpenAI团队开发，旨在推动更可信赖和可靠的AI代理研究。 核心特点. \nL21: 挑战性问题设计：BrowseComp的问题设计遵循严格的难度标准：. 人类创建者确保 ...\nL22: * 【6†什么是GAIA？; publish_date: none†huggingface.co】 什么是GAIA？ GAIA \nL23: 是一个用于评估AI助手在需要核心能力组合的真实世界任务上的表现的基准，这些核心能力包括推理、多模态理解、网页浏览和熟练的工具使用。\nL24: * 【7†SuperCLUE-Agent: Agent智能体中文原生任务能力测评基准; publish_date: \nL25: none†www.cluebenchmarks.com】 AI \nL26: Agent是一个能够自主理解、规划决策和执行复杂任务的智能体，有代表性的Agent比如AutoGPT、AgentGPT等。 \nL27: 现有关于Agent能力的测评，主要是在英文场景或任务的测评。目前还 ...\nL28: * 【8†Gaia2 与ARE：赋能社区的智能体评测; publish_date: none†hub.baai.ac.cn】 Sep 24, 2025 — \nL29: 快速评估任意智能体：基于真实或模拟数据，测试不同规则、工具、内容和验证方式下的表现 · 测试智能体的 工具调用与编排能力：可结合本地应用或MCP 工具 · 生成\nL30: ...\nL31: * 【9†许多AI 智能体评测基准并不可靠; publish_date: none†my.oschina.net】 Aug 13, 2025 — \nL32: 当前AI智能体基准测试存在严重缺陷，亟需更严谨的评估框架。作者提出AI智能体基准测试核查清单（ABC），包含43个项目，旨在识别基准测试潜在陷阱，构建可靠评估 \nL33: ...",
156 + "tool_call_id": "fixture-tool-call-1-1"
157 + },
158 + {
159 + "role": "assistant",
160 + "content": "",
161 + "reasoning_content": "第一个搜索结果（针对“search agent benchmark 2024”）显示了一些相关条目，例如“HSCodeComp: A Realistic and Expert-level Benchmark for ...”、“open-compass/GTA - A Benchmark for General Tool Agents”、“10 AI agent benchmarks”、“WISER-Search benchmark”、“Benchmarks and Tree Search for Multimodal LLM Web Agents”、“τ-bench”、“WebVoyager benchmark”、“FutureSearch Benchmarks”。这些可能都与搜索智能体相关。第二个搜索结果（针对“搜索智能体 基准测试”）显示了一些中文资料，例如“WideSearch”、“BrowseComp”、“GAIA”、“SuperCLUE-Agent”等。这些都可能与搜索智能体相关。为了全面了解这些基准测试，我们需要同时打开并分析其中几个最有潜力的结果。具体来说，我们将打开第一个搜索结果中的“HSCodeComp”（索引0）和“open-compass/GTA”（索引1），以及第二个搜索结果中的“WideSearch”（索引0）和“BrowseComp”（索引5）。",
162 + "tool_calls": [
163 + {
164 + "id": "fixture-tool-call-4-0",
165 + "type": "function",
166 + "function": {
167 + "name": "open",
168 + "arguments": "{\"cursor\": 0, \"id\": 0}"
169 + }
170 + },
171 + {
172 + "id": "fixture-tool-call-4-1",
173 + "type": "function",
174 + "function": {
175 + "name": "open",
176 + "arguments": "{\"cursor\": 0, \"id\": 1}"
177 + }
178 + },
179 + {
180 + "id": "fixture-tool-call-4-2",
181 + "type": "function",
182 + "function": {
183 + "name": "open",
184 + "arguments": "{\"cursor\": 1, \"id\": 0}"
185 + }
186 + },
187 + {
188 + "id": "fixture-tool-call-4-3",
189 + "type": "function",
190 + "function": {
191 + "name": "open",
192 + "arguments": "{\"cursor\": 1, \"id\": 5}"
193 + }
194 + }
195 + ]
196 + },
197 + {
198 + "role": "tool",
199 + "content": "[2] (https://ar5iv.org/html/2510.19631v1)\n**viewing lines [0 - 119] of 983**\n\nL0: \nL1: URL: https://ar5iv.org/html/2510.19631v1\nL2: # HSCodeComp: A Realistic and Expert-level Benchmark for Deep Search Agents in \nL3: Hierarchical Rule Application\nL4: \nL5: Yiqian Yang† Tian Lan† Qianghuai Jia∗ Li Zhu Hui Jiang Hang Zhu Longyue Wang \nL6: Weihua Luo Kaifu Zhang\nL7: \nL8: Alibaba International Digital Commerce∗* Corresponding Author: Qianghuai Jia \nL9: (qianghuai.jqh@alibaba-inc.com)\nL10: †\\dagger Equal Contribution: Yiqian Yang\nL11: \nL12: Tian Lan\nL13: \nL14: ###### Abstract\nL15: \nL16: Abstract\nL17: \nL18: Effective deep search agents must not only access open-domain and domain-\nL19: specific knowledge but also apply complex rules—such as legal clauses, medical \nL20: manuals and tariff rules. These rules often feature vague boundaries and \nL21: implicit logic relationships, making precise application challenging for agents.\nL22: However, this critical capability is largely overlooked by current agent \nL23: benchmarks. To fill this gap, we introduce HSCodeComp, the first realistic, \nL24: expert-level e-commerce benchmark designed to evaluate deep search agents in \nL25: hierarchical rule application. In this task, the deep reasoning process of \nL26: agents is guided by these rules to predict 10-digit Harmonized System Code \nL27: (HSCode) of products with noisy but realistic descriptions. These codes, \nL28: established by the World Customs Organization, are vital for global supply chain\nL29: efficiency. Built from real-world data collected from large-scale e-commerce \nL30: platforms, our proposed HSCodeComp comprises 632 product entries spanning \nL31: diverse product categories, with these HSCodes annotated by several human \nL32: experts. Extensive experimental results on several state-of-the-art LLMs, open-\nL33: source, and closed-source agents reveal a huge performance gap: best agent \nL34: achieves only 46.8% 10-digit accuracy, far below human experts at 95.0%. \nL35: Besides, detailed analysis demonstrates the challenges of hierarchical rule \nL36: application, and test-time scaling fails to improve performance further.\nL37: \nL38: ## 1 Introduction\nL39: \nL40: Deep search agents have demonstrated significant value in solving complex real-\nL41: world problems, where robust external knowledge utilization constitutes a \nL42: critical capability [Wu et al., 2025, Tao et al., 2025, Li et al., 2025b]. To \nL43: evaluate this capability, numerous established benchmarks are proposed to assess\nL44: agents in utilizing open-domain data (e.g., GAIA [Mialon et al., 2023b] and \nL45: BrowseComp [Wei et al., 2025]) and domain-specific data (e.g., WebMall [Peeters \nL46: et al., 2025a], FinSearchComp [Hu et al., 2025a] and MedBrowseComp [Yu et al., \nL47: 2025b]).\nL48: \nL49: Beyond open-domain and domain-specific data, agents also need to effectively \nL50: apply rules that encode human expert knowledge, particularly in scenarios like \nL51: law, medical and e-commerce [Li et al., 2025a, Chen et al., 2025b, Yao et al., \nL52: 2022, Chollet et al., 2025]. For instance, legal case adjudication require \nL53: interpreting abstract legal provisions, and accurate e-commerce product \nL54: classification in depends on tariff rules [Grainger, 2024]. Previous works have \nL55: defined rule application as using specific logical rules with supporting facts \nL56: to derive conclusions [Wang et al., 2024, Servantez et al., 2024]. In contrast, \nL57: we define it as a core capability for deep search agents, where human-written \nL58: rules are systematically applied to guide complex reasoning and decision-making \nL59: [Sadowski and Chudziak, 2025]. Building on this observation, we categorize \nL60: knowledge data for deep search agents into three levels (Figure 1, left), with \nL61: increasing knowledge complexity: (1) Level 1: Open-domain Data - Tests \nL62: understanding and deep reasoning abilities of agents on long-form web content. \nL63: Established benchmarks include GAIA [Mialon et al., 2023b] and BrowseComp [Wei \nL64: et al., 2025]; (2) Level 2: Structured Data - Assesses agents to precisely \nL65: utilize structured data such as databases and knowledge graphs, as seen in \nL66: domain-specific benchmarks like WebMall [Peeters et al., 2025a], MedBrowseComp \nL67: [Chen et al., 2025b] and FinSearchComp [Hu et al., 2025a]; (3) Level 3: Rule \nL68: Data - Evaluates agents to apply complex and abstract rules [Chollet et al., \nL69: 2025]. This level presents two key challenges: (a) making accurate decisions \nL70: when rules contain vague natural language descriptions [Sadowski and Chudziak, \nL71: 2025]; and (b) reasoning about logical dependencies among rules, such as \nL72: exception clauses and cross-category relationships [Guha et al., 2023]. Despite \nL73: the importance of rule application in real-world scenarios, current agent \nL74: benchmarks largely overlook its evaluation.\nL75: \nL76: To fill this gap, we introduce HSCodeComp (short for the Harmonized System Code \nL77: (HSCode) Competition), the first realistic, expert-level e-commerce benchmark \nL78: designed to evaluate agents in predicting complete 10-digit Harmonized System \nL79: Code (HSCode) of the product, using hierarchical rules (e.g., eWTP tariff \nL80: rules111https://www.ewtp.com/web/smart/hscode). HSCodes organize products \nL81: through a hierarchical structure spanning over 5,000 distinct codes across \nL82: multiple classification levels, representing the global standard for classifying\nL83: traded international goods, established by the World Customs Organization and \nL84: implemented across more than 200 countries for customs clearance and tariff \nL85: determination [Grainger, 2024, Nath et al., 2025]. Built from the data of the \nL86: large-scale e-commerce platforms, our proposed HSCodeComp comprises 632 \nL87: carefully curated product entries, encompassing 27 unique HS chapters and 32 \nL88: distinct first-level categories. These HSCodes have been rigorously annotated by\nL89: multiple e-commerce domain experts, ensuring that HSCodeComp is expert-level. \nL90: Accurately predicting the exact 10-digit HSCode presents significant challenges:\nL91: agents must perform multi-hop hierarchical reasoning with complex tariff rules \nL92: while processing noisy but realistic product descriptions that often contain \nL93: abbreviations, language variations, or incomplete information.\nL94: \nL95: Extensive experiments on the state-of-the-art baselines, including 14 advanced \nL96: foundation models, 6 advanced open-source agent systems and 3 closed-source \nL97: agent systems, demonstrate that HSCode prediction task remains a substantial \nL98: challenge for current AI approaches. As shown in the Figure 1 (right), even the \nL99: best-performing system (SmolAgent [Roucher et al., 2025] with GPT-5) achieves \nL100: only 46.8% accuracy, substantially below the 95.0% accuracy attained by human \nL101: experts. Further detailed analysis reveals that existing agent systems lack \nL102: critical capabilities required for this complex hierarchical rule applications. \nL103: Notably, test-time scaling approach—which has proven effective in other \nL104: reasoning tasks [Guo et al., 2025, Liu et al., 2025]—fail to improve performance\nL105: on HSCodeComp. These observations demonstrate the challenging nature of our \nL106: proposed HSCodeComp, highlighting the need for more effective designs of agent \nL107: systems. To facilitate future research, we will publicly release codes and the \nL108: benchmark dataset of HSCodeComp.\nL109: \nL110: ## 2 Related Works\nL111: \nL112: ### 2.1 Previous Works in HSCode Prediction\nL113: \nL114: Previous works treat HSCode prediction as the e-commerce text classification \nL115: task [Grainger, 2024], using pre-trained BERT models [Liao et al., 2024, Shubham\nL116: et al., 2022] or Large Language Models (LLMs) prompting [Hussain and Ahmed, \nL117: 2023]. However, these approaches fail to leverage domain-specific knowledge, \nL118: especially the rules written by human experts [Hussain and Ahmed, 2023, Judy, \nL119: 2024]. Besides, existing HSCode benchmarks face two critical limitations [Judy, ",
200 + "tool_call_id": "fixture-tool-call-4-0"
201 + },
202 + {
203 + "role": "tool",
204 + "content": "Error fetching URL `https://github.com/open-compass/GTA`",
205 + "tool_call_id": "fixture-tool-call-4-1"
206 + },
207 + {
208 + "role": "tool",
209 + "content": "Error fetching URL `https://zhuanlan.zhihu.com/p/1939316761775301093`",
210 + "tool_call_id": "fixture-tool-call-4-2"
211 + },
212 + {
213 + "role": "tool",
214 + "content": "[3] (https://blog.csdn.net/qq_41472205/article/details/151974603)\n**viewing lines [0 - 167] of 187**\n\nL0: \nL1: URL: https://blog.csdn.net/qq_41472205/article/details/151974603\nL2: # BrowseComp：为浏览智能体设计的简单而具挑战性的基准测试\nL3: \nL4: BrowseComp：AI浏览能力评估基准\nL5: \nL6: 最新推荐文章于 2025-11-12 13:40:20 发布\nL7: \nL8: 原创 于 2025-09-22 22:33:04 发布 · 1.3k 阅读\nL9: \nL10: · 9\nL11: · 25 · \nL12: CC 4.0 BY-SA版权\nL13: \nL14: 版权声明：本文为博主原创文章，遵循 CC 4.0 BY-SA 版权协议，转载请附上原文出处链接和本声明。\nL15: \nL16: ## BrowseComp：为浏览智能体设计的简单而具挑战性的基准测试\nL17: \nL18: 在人工智能从基础聊天机器人向推理器和智能体发展的进程中，具备浏览互联网能力的人工智能模型正变得越来越重要。今天，我们将介绍一个名为BrowseComp的创新基准\nL19: 测试，它专门设计用于评估AI代理在复杂网络浏览任务中的能力。\nL20: \nL21: ### 什么是BrowseComp？\nL22: \nL23: BrowseComp（全称Browsing Competition）是一个包含1,266个挑战性问题的基准测试集，专门用于衡量AI代理在互联网上持续导航、寻找难\nL24: 以找到的纠缠信息的能力。该基准测试由OpenAI团队开发，旨在推动更可信赖和可靠的AI代理研究。\nL25: \nL26: #### 核心特点\nL27: \nL28: 挑战性问题设计：BrowseComp的问题设计遵循严格的难度标准：\nL29: \nL30: - 人类创建者确保问题在10分钟内无法被人解决\nL31: - 现有模型（包括带浏览功能的ChatGPT和早期版本的OpenAI Deep Research）无法解决\nL32: - 通过5次简单Google搜索无法在结果首页找到答案\nL33: \nL34: 简单易验证：尽管问题极具挑战性，但答案形式简单——都是短字符串，便于自动验证模型输出的正确性。\nL35: \nL36: ### 为什么需要BrowseComp？\nL37: \nL38: #### 现有基准的局限性\nL39: \nL40: 传统的信息检索基准（如TriviaQA、HotpotQA等）主要关注易于查找的信息，随着语言模型的进步，这些基准已经趋于饱和。而BrowseComp专注于那些需\nL41: 要浏览大量网站才能解决的\"硬核\"问题。\nL42: \nL43: #### 模拟真实挑战\nL44: \nL45: BrowseComp问题通常采用\"逆向设计\"方法：创建者从一个已知事实出发，构建一个搜索空间巨大但验证简单的问题。例如：\nL46: \nL47: “找出2018-2023年间在EMNLP会议上发表、第一作者本科毕业于达特茅斯学院、第四作者本科毕业于宾夕法尼亚大学的科学论文标题”\nL48: \nL49: 这类问题验证简单，但解决起来需要检查数千篇论文并调查每位作者的背景。\nL50: \nL51: ### 数据集特点\nL52: \nL53: #### 主题多样性\nL54: \nL55: BrowseComp涵盖了广泛的主题领域（如图2所示），包括历史、科学、文化等。创建者被鼓励基于个人兴趣设计问题，这有助于提高数据质量和参与度。\nL56: \nL57: #### 质量保证\nL58: \nL59: 为确保答案的唯一性，创建者需要：\nL60: \nL61: - 对问题内容有足够了解，确信没有其他有效答案\nL62: - 如果不确定，则添加更多约束条件\nL63: - 接受其他创建者的验证反馈\nL64: \nL65: ### 人类表现基准\nL66: \nL67: 为了衡量BrowseComp的难度，研究人员让人类创建者尝试解决问题（不能解答自己创建的问题）。结果显示：\nL68: \nL69: - **70.8%**的问题在2小时搜索后人类选择放弃\nL70: - **29.2%**的问题被成功解决\nL71: - 在解决的问题中，**86.4%**的人类答案与参考答案一致\nL72: \nL73: 这表明BrowseComp确实极具挑战性，即使是熟悉数据集的人类专家也难以在有限时间内解决大部分问题。\nL74: \nL75: ### AI模型表现评估\nL76: \nL77: #### 各模型对比\nL78: \nL79: 研究人员评估了多种模型在BrowseComp上的表现：\nL80: \nL81: 模型 | 准确率(%) | 校准误差(%) \nL82: ---|---|---\nL83: GPT-4o | 0.6 | 69 \nL84: GPT-4o（带浏览） | 1.9 | 82 \nL85: GPT-4.5 | 0.9 | 68 \nL86: OpenAI o1 | 9.9 | 65 \nL87: Deep Research | 51.5 | 91 \nL88: \nL89: #### 关键发现\nL90: \nL91: - 基础模型表现不佳：GPT-4o和GPT-4.5准确率接近零，凸显了基准的难度\nL92: - 浏览功能带来有限提升：启用浏览功能的GPT-4o准确率略有提高，但仍很低\nL93: - 推理能力的重要性：OpenAI o1虽然没有浏览能力，但凭借更强的推理能力获得较高准确率\nL94: - 专业模型的优势：专门为持久网络浏览训练的Deep Research模型解决了约一半的问题\nL95: \nL96: #### 计算资源与性能关系\nL97: \nL98: 研究表明，BrowseComp性能随测试时计算资源的增加而平滑提升（如图1所示）。这与智能体模型的特性一致——更多计算资源允许模型浏览更多网站，从而提高找到正确\nL99: 答案的机会。\nL100: \nL101: ### 进阶策略分析\nL102: \nL103: #### 聚合策略的效果\nL104: \nL105: 通过让模型多次尝试同一问题并采用投票策略，可以显著提升性能：\nL106: \nL107: - 多数投票：选择样本中最常见的答案\nL108: - 加权投票：根据模型置信度加权投票\nL109: - 最佳选择：选择置信度最高的答案\nL110: \nL111: 这些方法将Deep Research的性能提升了15-25%，表明模型通常能够识别自己的正确答案。\nL112: \nL113: #### 任务难度分布\nL114: \nL115: 分析显示，BrowseComp中的任务难度分布广泛：\nL116: \nL117: - 16%的任务被Deep Research完美解决（100%通过率）\nL118: - 14%的任务完全失败（0%通过率）\nL119: - 其余任务处于中间难度水平\nL120: \nL121: ### BrowseComp的意义与局限性\nL122: \nL123: #### 作为评估工具的价值\nL124: \nL125: BrowseComp可被视为浏览智能体的\"编程竞赛\"——虽然不全面，但对核心浏览能力提供了有用的衡量。正如在编程竞赛中表现出色的模型很可能在其他编码任务中表现良\nL126: 好，能够解决BrowseComp的模型在定位难以查找信息方面应该具有很强的能力。\nL127: \nL128: #### 已知局限性\nL129: \nL130: - 不涉及真实用户查询分布：如生成长回答或解决查询歧义等挑战\nL131: - 答案唯一性假设：虽然尽力确保，但不能完全保证没有其他有效答案\nL132: - 单模态限制：当前仅关注文本信息检索，未来可扩展至图像、视频等模态\nL133: \nL134: ### 结论与展望\nL135: \nL136: BrowseComp为评估AI代理的网络浏览能力提供了一个具有挑战性且易于使用的基准。它强调的持久性、创造性和战略性浏览能力，正是未来AI智能体在实际应用中所需\nL137: 的核心技能。\nL138: \nL139: 随着更多研究团队在BrowseComp上测试他们的模型，我们期待看到浏览智能体技术的快速进步，最终使AI能够更可靠地帮助我们获取那些深藏在互联网角落的宝贵信息。\nL140: \nL141: 访问地址：BrowseComp已在GitHub上开源，地址为 https://github.com/openai/simple-evals\nL142: \nL143: 添加红包\nL144: \nL145: 请填写红包祝福语或标题\nL146: \nL147: 个\nL148: \nL149: 红包个数最小为10个\nL150: \nL151: 元\nL152: \nL153: 红包金额最低5元\nL154: \nL155: 当前余额3.43元 前往充值 >\nL156: \nL157: 需支付：10.00元\nL158: \nL159: 成就一亿技术人!\nL160: \nL161: hope_wisdom\nL162: 发出的红包 \nL163: 实付元\nL164: \nL165: 点击重新获取\nL166: \nL167: 扫码支付",
215 + "tool_call_id": "fixture-tool-call-4-3"
216 + },
217 + {
218 + "role": "assistant",
219 + "content": "",
220 + "reasoning_content": "打开的HSCodeComp页面详细介绍了这个基准测试，它是一个用于评估深度搜索智能体在分层规则应用方面能力的基准测试，特点包括：基于真实世界电商数据、任务复杂、需要多跳层次推理、使用嘈杂的产品描述，并且标注由专家完成。使用场景是评估智能体在规则应用方面的能力，例如关税代码分类。例题是预测产品的10位HS编码。打开的BrowseComp页面详细介绍了这个基准测试，它是一个用于评估浏览智能体（即能够在互联网上导航、搜索信息的智能体）的基准测试。其特点包括：包含1266个具有挑战性的问题，需要持续导航和搜索难以查找的信息；答案形式简单，便于自动验证；问题设计确保人类在10分钟内无法解决，且现有模型也无法解决；涵盖了广泛的主题领域。使用场景是评估智能体在复杂网络浏览任务中的能力。例题是：“找出2018-2023年间在EMNLP会议上发表、第一作者本科毕业于达特茅斯学院、第四作者本科毕业于宾夕法尼亚大学的科学论文标题”。另外两个打开的页面（open-compass/GTA和WideSearch）出现了错误，可能链接已失效或需要其他访问方式。为了进一步获取其他相关基准测试的信息，我们需要同时进行多项搜索。考虑到用户可能对搜索智能体基准测试有广泛需求，我们将同时搜索“WebVoyager benchmark”、“WebArena benchmark”、“GAIA benchmark”、“ToolBench benchmark”、“HotpotQA benchmark”、“FEVER benchmark”、“TriviaQA benchmark”、“Natural Questions benchmark”、“MS MARCO benchmark”、“BEIR benchmark”、“MIRACL benchmark”、“Zero-shot IR benchmark”、“WebGPT benchmark”、“WebShop benchmark”、“ALFWorld benchmark”、“Mind2Web benchmark”、“VisualWebArena benchmark”、“SearchBench benchmark”、“WebVLN benchmark”、“WebNav benchmark”。这些搜索将帮助我们全面了解现有基准测试的概况。",
221 + "tool_calls": [
222 + {
223 + "id": "fixture-tool-call-9-0",
224 + "type": "function",
225 + "function": {
226 + "name": "search",
227 + "arguments": "{\"query\": \"WebVoyager benchmark\", \"topn\": 10, \"source\": \"web\"}"
228 + }
229 + },
230 + {
231 + "id": "fixture-tool-call-9-1",
232 + "type": "function",
233 + "function": {
234 + "name": "search",
235 + "arguments": "{\"query\": \"WebArena benchmark\", \"topn\": 10, \"source\": \"web\"}"
236 + }
237 + },
238 + {
239 + "id": "fixture-tool-call-9-2",
240 + "type": "function",
241 + "function": {
242 + "name": "search",
243 + "arguments": "{\"query\": \"GAIA benchmark\", \"topn\": 10, \"source\": \"web\"}"
244 + }
245 + },
246 + {
247 + "id": "fixture-tool-call-9-3",
248 + "type": "function",
249 + "function": {
250 + "name": "search",
251 + "arguments": "{\"query\": \"ToolBench benchmark\", \"topn\": 10, \"source\": \"web\"}"
252 + }
253 + },
254 + {
255 + "id": "fixture-tool-call-9-4",
256 + "type": "function",
257 + "function": {
258 + "name": "search",
259 + "arguments": "{\"query\": \"HotpotQA benchmark\", \"topn\": 10, \"source\": \"web\"}"
260 + }
261 + },
262 + {
263 + "id": "fixture-tool-call-9-5",
264 + "type": "function",
265 + "function": {
266 + "name": "search",
267 + "arguments": "{\"query\": \"FEVER benchmark\", \"topn\": 10, \"source\": \"web\"}"
268 + }
269 + },
270 + {
271 + "id": "fixture-tool-call-9-6",
272 + "type": "function",
273 + "function": {
274 + "name": "search",
275 + "arguments": "{\"query\": \"TriviaQA benchmark\", \"topn\": 10, \"source\": \"web\"}"
276 + }
277 + },
278 + {
279 + "id": "fixture-tool-call-9-7",
280 + "type": "function",
281 + "function": {
282 + "name": "search",
283 + "arguments": "{\"query\": \"Natural Questions benchmark\", \"topn\": 10, \"source\": \"web\"}"
284 + }
285 + },
286 + {
287 + "id": "fixture-tool-call-9-8",
288 + "type": "function",
289 + "function": {
290 + "name": "search",
291 + "arguments": "{\"query\": \"MS MARCO benchmark\", \"topn\": 10, \"source\": \"web\"}"
292 + }
293 + },
294 + {
295 + "id": "fixture-tool-call-9-9",
296 + "type": "function",
297 + "function": {
298 + "name": "search",
299 + "arguments": "{\"query\": \"BEIR benchmark\", \"topn\": 10, \"source\": \"web\"}"
300 + }
301 + },
302 + {
303 + "id": "fixture-tool-call-9-10",
304 + "type": "function",
305 + "function": {
306 + "name": "search",
307 + "arguments": "{\"query\": \"MIRACL benchmark\", \"topn\": 10, \"source\": \"web\"}"
308 + }
309 + },
310 + {
311 + "id": "fixture-tool-call-9-11",
312 + "type": "function",
313 + "function": {
314 + "name": "search",
315 + "arguments": "{\"query\": \"Zero-shot IR benchmark\", \"topn\": 10, \"source\": \"web\"}"
316 + }
317 + },
318 + {
319 + "id": "fixture-tool-call-9-12",
320 + "type": "function",
321 + "function": {
322 + "name": "search",
323 + "arguments": "{\"query\": \"WebGPT benchmark\", \"topn\": 10, \"source\": \"web\"}"
324 + }
325 + },
326 + {
327 + "id": "fixture-tool-call-9-13",
328 + "type": "function",
329 + "function": {
330 + "name": "search",
331 + "arguments": "{\"query\": \"WebShop benchmark\", \"topn\": 10, \"source\": \"web\"}"
332 + }
333 + },
334 + {
335 + "id": "fixture-tool-call-9-14",
336 + "type": "function",
337 + "function": {
338 + "name": "search",
339 + "arguments": "{\"query\": \"ALFWorld benchmark\", \"topn\": 10, \"source\": \"web\"}"
340 + }
341 + },
342 + {
343 + "id": "fixture-tool-call-9-15",
344 + "type": "function",
345 + "function": {
346 + "name": "search",
347 + "arguments": "{\"query\": \"Mind2Web benchmark\", \"topn\": 10, \"source\": \"web\"}"
348 + }
349 + },
350 + {
351 + "id": "fixture-tool-call-9-16",
352 + "type": "function",
353 + "function": {
354 + "name": "search",
355 + "arguments": "{\"query\": \"VisualWebArena benchmark\", \"topn\": 10, \"source\": \"web\"}"
356 + }
357 + },
358 + {
359 + "id": "fixture-tool-call-9-17",
360 + "type": "function",
361 + "function": {
362 + "name": "search",
363 + "arguments": "{\"query\": \"SearchBench benchmark\", \"topn\": 10, \"source\": \"web\"}"
364 + }
365 + },
366 + {
367 + "id": "fixture-tool-call-9-18",
368 + "type": "function",
369 + "function": {
370 + "name": "search",
371 + "arguments": "{\"query\": \"WebVLN benchmark\", \"topn\": 10, \"source\": \"web\"}"
372 + }
373 + },
374 + {
375 + "id": "fixture-tool-call-9-19",
376 + "type": "function",
377 + "function": {
378 + "name": "search",
379 + "arguments": "{\"query\": \"WebNav benchmark\", \"topn\": 10, \"source\": \"web\"}"
380 + }
381 + }
382 + ]
383 + },
384 + {
385 + "role": "tool",
386 + "content": "[4] WebVoyager benchmark (Search_Results/WebVoyager benchmark)\n**viewing lines [0 - 40] of 40**\n\nL0: \nL1: URL: Search_Results/WebVoyager benchmark\nL2: # Search Results\nL3: \nL4: * 【0†WebVoyager Benchmark; publish_date: none†www.trykura.com】 Performance \nL5: measured across various websites in WebVoyager shows Kura consistently \nL6: outperforming other agents, with particularly strong results on e-commerce ...\nL7: * 【1†WebVoyager: Building an End-to-End Web Agent with ...; publish_date: \nL8: none†github.com】 WebVoyager is an innovative Large Multimodal Model (LMM) \nL9: powered web agent that can complete user instructions end-to-end by interacting \nL10: with real-world ...\nL11: * 【2†AI Browser Agent Leaderboard | Steel.dev; publish_date: \nL12: none†leaderboard.steel.dev】 See how various AI browser agents stack up based on \nL13: their accuracy in completing web-based tasks on the WebVoyager benchmark.\nL14: * 【3†[2401.13919] WebVoyager: Building an End-to-End Web ...; publish_date: \nL15: none†arxiv.org】 by H He · 2024 · Cited by 282 — We show that WebVoyager achieves\nL16: a 59.1% task success rate on our benchmark, significantly surpassing the \nL17: performance of both GPT-4 (All ...\nL18: * 【4†Our Agent-E SOTA Results on the WebVoyager Benchmark; publish_date: \nL19: none†www.emergence.ai】 Jul 11, 2024 — WebVoyager is a benchmark that tests an \nL20: agent's capabilities for navigation on dynamic live websites. It is more \nL21: representative than WebArena [4] ...\nL22: * 【5†Browser Use = state of the art Web Agent; publish_date: none†browser-\nL23: use.com】 Dec 15, 2024 — Browser Use has achieved state-of-the-art performance on\nL24: the WebVoyager benchmark, with an impressive 89.1% success rate across 586 \nL25: diverse web tasks.\nL26: * 【6†Magnitude achieves SOTA 94% on WebVoyager benchmark; publish_date: \nL27: none†github.com】 Magnitude achieves state-of-the-art performance with 93.9% \nL28: success rate on WebVoyager, beating all other browser agents.\nL29: * 【7†WebVoyager: Autonomous Web Agent Benchmark; publish_date: \nL30: none†www.emergentmind.com】 3 days ago — WebVoyager Benchmark is a comprehensive \nL31: evaluation suite for autonomous web agents, featuring 643 tasks across 15 \nL32: popular websites.\nL33: * 【8†WebVoyager Benchmark Results; publish_date: none†www.browserable.ai】 \nL34: Browserable has achieved 90.4% on the WebVoyager benchmark. This is best-in-\nL35: class performance across all web agents. This was done across 567 web tasks \nL36: which ...\nL37: * 【9†89% achieved on WebVoyager using Anchor + Browser Use; publish_date: \nL38: none†www.reddit.com】 Thanks to the amazing work from the browser-use open-source\nL39: community and the built-in support from Anchor Browser, we've hit an 89% score \nL40: on WebVoyager.",
387 + "tool_call_id": "fixture-tool-call-9-0"
388 + },
389 + {
390 + "role": "tool",
391 + "content": "[5] WebArena benchmark (Search_Results/WebArena benchmark)\n**viewing lines [0 - 42] of 42**\n\nL0: \nL1: URL: Search_Results/WebArena benchmark\nL2: # Search Results\nL3: \nL4: * 【0†WebArena: A Realistic Web Environment for Building ...; publish_date: \nL5: none†webarena.dev】 Our benchmark is implemented in our fully interactable \nL6: highly-realistic WebArena environment. It features diverse tasks human may \nL7: encounter in their daily ...\nL8: * 【1†[2307.13854] WebArena: A Realistic Web Environment for ...; publish_date:\nL9: none†arxiv.org】 by S Zhou · 2023 · Cited by 637 — Building upon our \nL10: environment, we release a set of benchmark tasks focusing on evaluating the \nL11: functional correctness of task completions.\nL12: * 【2†WebArena: A Realistic Web Environment for Building ...; publish_date: \nL13: none†www.cmu.edu】 WebArena introduces a benchmark on interpreting high-level \nL14: realistic natural language command to concrete web-based interactions. We \nL15: provide annotated programs ...\nL16: * 【3†GitHub - web-arena-x/webarena: Code repo for ...; publish_date: \nL17: none†github.com】 [12/20/2024] Check out our new benchmark on even more \nL18: consequential tasks, including terminal use and coding, TheAgentCompany. \nL19: [12/21/2023] We release the ...\nL20: * 【4†WebArena Benchmark and the State of Agentic AI; publish_date: \nL21: none†medium.com】 In short, WebArena established a new standard for realism and \nL22: complexity in web agent evaluation, forcing AI agents to operate in dynamic, \nL23: high- ...\nL24: * 【5†WebArena: A Realistic Web Environment for Building ...; publish_date: \nL25: none†huggingface.co】 Jul 25, 2023 — WebArena, a realistic and reproducible \nL26: environment, evaluates the performance of autonomous agents performing complex \nL27: tasks on websites using ...\nL28: * 【6†WebArena Benchmark: Evaluating Web Agents; publish_date: \nL29: none†www.emergentmind.com】 Jun 30, 2025 — WebArena Benchmark is a self-contained\nL30: suite that evaluates autonomous agents on realistic, multi-step web tasks using\nL31: natural language ...\nL32: * 【7†VisualWebArena is a benchmark for multimodal agents.; publish_date: \nL33: none†github.com】 VisualWebArena is a realistic and diverse benchmark for \nL34: evaluating multimodal autonomous language agents. It comprises of a set of \nL35: diverse and complex web-based ...\nL36: * 【8†WebDev Arena Leaderboard - LMArena; publish_date: none†web.lmarena.ai】 \nL37: WebDev Arena is a real-time AI coding competition where models go head-to-head \nL38: in web development challenges, developed by LMArena.\nL39: * 【9†WebArena: A Realistic Web Environment for Building ...; publish_date: \nL40: none†arxiv.org】 Apr 16, 2024 — We use this benchmark to evaluate several agents \nL41: that can follow NL command and perform web-based tasks (§4). These agents are \nL42: implemented in a ...",
392 + "tool_call_id": "fixture-tool-call-9-1"
393 + },
394 + {
395 + "role": "tool",
396 + "content": "[6] GAIA benchmark (Search_Results/GAIA benchmark)\n**viewing lines [0 - 41] of 41**\n\nL0: \nL1: URL: Search_Results/GAIA benchmark\nL2: # Search Results\nL3: \nL4: * 【0†GAIA Leaderboard - a Hugging Face Space by ...; publish_date: \nL5: none†huggingface.co】 GAIA is a benchmark which aims at evaluating next-\nL6: generation LLMs (LLMs with augmented capabilities due to added tooling, \nL7: efficient prompting, access to search ...\nL8: * 【1†[2311.12983] GAIA: a benchmark for General AI Assistants; publish_date: \nL9: none†arxiv.org】 by G Mialon · 2023 · Cited by 367 — GAIA proposes real-world \nL10: questions that require a set of fundamental abilities such as reasoning, multi-\nL11: modality handling, web browsing, and generally tool-use ...\nL12: * 【2†GAIA benchmark; publish_date: none†huggingface.co】 This is the \nL13: organisation page for all things related to GAIA, a benchmark for General AI \nL14: Assistants. You can find all the information and links on the GAIA ...\nL15: * 【3†GAIA: A Benchmark for General AI Assistants; publish_date: \nL16: none†ukgovernmentbeis.github.io】 This is an Inspect AI implementation of the \nL17: GAIA (General AI Assistants) benchmark, consisting of 450 questions testing tool\nL18: use on realistic assistant tasks.\nL19: * 【4†GAIA: a benchmark for general AI assistants | Research; publish_date: \nL20: none†ai.meta.com】 May 6, 2024 — GAIA proposes real-world questions that require \nL21: a set of fundamental abilities such as reasoning, multi-modality handling, web \nL22: browsing, and generally tool-use ...\nL23: * 【5†HAL: GAIA Leaderboard; publish_date: none†hal.cs.princeton.edu】 GAIA is a\nL24: benchmark for General AI Assistants that requires a set of fundamental \nL25: abilities such as reasoning, multi-modality handling, web browsing, and tool- \nL26: ...\nL27: * 【6†GAIA: The LLM Agent Benchmark Everyone's Talking About; publish_date: \nL28: none†towardsdatascience.com】 May 29, 2025 — GAIA stands for General AI \nL29: Assistants benchmark [1]. This benchmark was introduced to specifically evaluate\nL30: LLM agents on their ability to act as general- ...\nL31: * 【7†GAIA: a benchmark for General AI Assistants; publish_date: \nL32: none†openreview.net】 by G Mialon · Cited by 367 — GAIA proposes real-world \nL33: questions that require a set of fundamental abilities such as reasoning, multi-\nL34: modality handling, web browsing, and generally tool-use ...\nL35: * 【8†Rethinking AI Evaluation: Introducing the GAIA Benchmark; publish_date: \nL36: none†medium.com】 The authors introduce GAIA, a benchmark designed to assess the \nL37: robustness of AI systems across a variety of practical tasks.\nL38: * 【9†H2O.ai Tops the General AI Assistant (GAIA) Test; publish_date: \nL39: none†h2o.ai】 Mar 17, 2025 — Our h2oGPTe Agent has once again claimed the #1 spot\nL40: on the prestigious GAIA (General AI Assistants) benchmark with an impressive \nL41: 75% accuracy rate.",
397 + "tool_call_id": "fixture-tool-call-9-2"
398 + },
399 + {
400 + "role": "tool",
401 + "content": "[7] ToolBench benchmark (Search_Results/ToolBench benchmark)\n**viewing lines [0 - 40] of 40**\n\nL0: \nL1: URL: Search_Results/ToolBench benchmark\nL2: # Search Results\nL3: \nL4: * 【0†ToolBench, an evaluation suite for LLM tool manipulation ...; \nL5: publish_date: none†github.com】 The ToolBench is a benchmark consisting of \nL6: diverse software tools for real-world tasks. We also provide easy-to-use \nL7: infrastructure in this repository.\nL8: * 【1†OpenBMB/ToolBench; publish_date: none†github.com】 [2023/7/27] New version\nL9: ToolBench is released. ✨Here is an overview of the dataset construction, \nL10: training, and evaluation. ✨✨Features:.\nL11: * 【2†Towards Stable Large-Scale Benchmarking on Tool ...; publish_date: \nL12: none†arxiv.org】 by Z Guo · 2024 · Cited by 100 — We introduce StableToolBench, a\nL13: benchmark evolving from ToolBench, proposing a virtual API server and stable \nL14: evaluation system.\nL15: * 【3†StableToolBench - Zhicheng Guo; publish_date: none†zhichengg.github.io】 \nL16: We introduce StableToolBench, a benchmark evolving from ToolBench, proposing a \nL17: virtual API server and stable evaluation system.\nL18: * 【4†ToolBench | EvalScope - Read the Docs; publish_date: \nL19: none†evalscope.readthedocs.io】 We evaluate the effectiveness of the ToolBench \nL20: benchmark: ToolBench (Qin et al., 2023b). The task involves integrating API \nL21: calls to complete tasks.\nL22: * 【5†Towards Stable Large-Scale Benchmarking on Tool ...; publish_date: \nL23: none†aclanthology.org】 by Z Guo · 2024 · Cited by 100 — We introduce \nL24: StableToolBench, a benchmark evolving from ToolBench, proposing a virtual API \nL25: server and stable evaluation system.\nL26: * 【6†ML-Tool-Bench: Tool-Augmented Planning for ML Tasks; publish_date: \nL27: none†openreview.net】 Sep 18, 2025 — In this work, we introduce a comprehensive \nL28: benchmark for evaluating tool-augmented ML agents using a curated set of 61 \nL29: specialized tools and 15 ...\nL30: * 【7†-Bench: Benchmarking AI agents for the real-world; publish_date: \nL31: none†sierra.ai】 Jun 20, 2024 — τ-bench measures an agent's ability to interact \nL32: with (simulated) human users and programmatic APIs while following domain-\nL33: specific policies in a consistent ...\nL34: * 【8†ToolEval Leaderboard; publish_date: none†openbmb.github.io】 ToolEval is \nL35: an automatic evaluator build for tool learning which incorporates two evaluation\nL36: metrics, Pass Rate and Win Rate(Preference).\nL37: * 【9†What is the best benchmark dataset for multi-step tool-use?; \nL38: publish_date: none†www.reddit.com】 I'm a newbie trying to evaluate the \nL39: performance of different prompts strategies for multi-step tool-using, wondering\nL40: what is the recommended benchmark dataset ...",
402 + "tool_call_id": "fixture-tool-call-9-3"
403 + },
404 + {
405 + "role": "tool",
406 + "content": "[8] HotpotQA benchmark (Search_Results/HotpotQA benchmark)\n**viewing lines [0 - 39] of 39**\n\nL0: \nL1: URL: Search_Results/HotpotQA benchmark\nL2: # Search Results\nL3: \nL4: * 【0†HotpotQA Homepage; publish_date: none†hotpotqa.github.io】 HotpotQA is a \nL5: question answering dataset featuring natural, multi-hop questions, with strong \nL6: supervision for supporting facts to enable more explainable ...See more\nL7: * 【1†HotpotQA: A Dataset for Diverse, Explainable Multi-hop ...; publish_date:\nL8: none†arxiv.org】 by Z Yang · 2018 · Cited by 3834 — HotpotQA is a dataset with \nL9: 113k Wikipedia-based question-answer pairs requiring multi-document reasoning, \nL10: diverse questions, sentence-level ...\nL11: * 【2†hotpotqa/hotpot_qa · Datasets at Hugging Face; publish_date: \nL12: none†huggingface.co】 HotpotQA is a new dataset with 113k Wikipedia-based \nL13: question-answer pairs with four key features: (1) the questions require finding \nL14: and reasoning over multiple ...See more\nL15: * 【3†Why You Should Stop Using HotpotQA for AI Agents ...; publish_date: \nL16: none†qipeng.me】 Jul 1, 2025 — HotpotQA pioneered a class of AI tasks that \nL17: requires the AI system to autonomously perform multiple steps of reasoning in an\nL18: open-domain setting.See more\nL19: * 【4†hotpotqa/hotpot; publish_date: none†github.com】 A dataset for diverse, \nL20: explainable multi-hop question answering. This repository contains the baseline \nL21: model code, as well as the entire pipeline of running ...See more\nL22: * 【5†HotpotQA: Multi-Hop QA Benchmark; publish_date: \nL23: none†www.emergentmind.com】 Sep 10, 2025 — HotpotQA is a large-scale multi-hop \nL24: question answering benchmark featuring 112,779 Wikipedia-based Q&A pairs with \nL25: detailed, sentence-level ...See more\nL26: * 【6†HotpotQA Dataset | Papers With Code; publish_date: \nL27: none†paperswithcode.com】 HotpotQA is a question answering dataset collected on \nL28: the English Wikipedia, containing about 113K crowd-sourced questions.See more\nL29: * 【7†HotpotQA: A Dataset for Diverse, Explainable Multi-hop ...; publish_date:\nL30: none†aclanthology.org】 by Z Yang · 2018 · Cited by 3834 — HotpotQA is a dataset\nL31: with 113k Wikipedia-based question-answer pairs requiring multi-document \nL32: reasoning, diverse questions, sentence-level facts, and factoid ...\nL33: * 【8†Benchmark BM25S: HotpotQA; publish_date: none†www.kaggle.com】 Explore and\nL34: run machine learning code with Kaggle Notebooks | Using data from No attached \nL35: data sources.\nL36: * 【9†mteb/hotpotqa · Datasets at Hugging Face; publish_date: \nL37: none†huggingface.co】 HotpotQA is a question answering dataset featuring natural,\nL38: multi-hop questions, with strong supervision for supporting facts to enable \nL39: more explainable ...See more",
407 + "tool_call_id": "fixture-tool-call-9-4"
408 + },
409 + {
410 + "role": "tool",
411 + "content": "[9] FEVER benchmark (Search_Results/FEVER benchmark)\n**viewing lines [0 - 40] of 40**\n\nL0: \nL1: URL: Search_Results/FEVER benchmark\nL2: # Search Results\nL3: \nL4: * 【0†Fever.ai; publish_date: none†fever.ai】 We are pleased to announce that \nL5: FEVER9 will be co-located with EACL 2026. In this year's workshop, we will \nL6: introduce a new shared task focused on automated fact ...\nL7: * 【1†a Large-scale Dataset for Fact Extraction and VERification; publish_date:\nL8: none†aclanthology.org】 by J Thorne · 2018 · Cited by 2315 — In this paper we \nL9: introduce a new publicly available dataset for verification against textual \nL10: sources, FEVER: Fact Extraction.\nL11: * 【2†awslabs/fever: FEVER (Fact Extraction and VERification) ...; \nL12: publish_date: none†github.com】 In this paper we introduce a new publicly \nL13: available dataset for verification against textual sources, FEVER: Fact \nL14: Extraction and VERification.\nL15: * 【3†FEVER: Fact Extraction and VERification; publish_date: \nL16: none†www.amazon.science】 The best accuracy we achieve on labeling a claim \nL17: accompanied by the correct evidence is 31.87%, while if we ignore the evidence \nL18: we achieve 50.91%. Thus we ...\nL19: * 【4†FEVER Dataset; publish_date: none†fever.ai】 FEVER (Fact Extraction and \nL20: VERification) consists of 185,445 claims generated by altering sentences \nL21: extracted from Wikipedia and subsequently verified ...\nL22: * 【5†mteb/fever · Datasets at Hugging Face; publish_date: none†huggingface.co】\nL23: FEVER. An MTEB dataset. Massive Text Embedding Benchmark. FEVER (Fact \nL24: Extraction and VERification) consists of 185,445 claims generated by altering \nL25: sentences ...\nL26: * 【6†FEVEROUS: Fact Extraction and VERification Over ...; publish_date: \nL27: none†datasets-benchmarks-proceedings.neurips.cc】 by R Aly · Cited by 359 — In \nL28: this paper we introduce a novel dataset and benchmark, Fact Extraction and \nL29: VERification Over. Unstructured and Structured information (FEVEROUS), which ...\nL30: * 【7†a large-scale dataset for Fact Extraction and VERification; publish_date:\nL31: none†arxiv.org】 by J Thorne · 2018 · Cited by 2315 — In this paper we introduce\nL32: a new publicly available dataset for verification against textual sources, \nL33: FEVER: Fact Extraction and VERification.\nL34: * 【8†FEVER Resources; publish_date: none†fever.ai】 2018 FEVER: a large-scale \nL35: dataset for Fact Extraction and VERification .bib James Thorne, Andreas Vlachos,\nL36: Christos Christodoulopoulos, Arpit Mittal\nL37: * 【9†a Large-scale Dataset for Fact Extraction and VERification; publish_date:\nL38: none†www.semanticscholar.org】 This paper introduces a new publicly available \nL39: dataset for verification against textual sources, FEVER, which consists of \nL40: 185,445 claims generated by ...",
412 + "tool_call_id": "fixture-tool-call-9-5"
413 + },
414 + {
415 + "role": "tool",
416 + "content": "[10] TriviaQA benchmark (Search_Results/TriviaQA benchmark)\n**viewing lines [0 - 35] of 35**\n\nL0: \nL1: URL: Search_Results/TriviaQA benchmark\nL2: # Search Results\nL3: \nL4: * 【0†TriviaQA; publish_date: none†nlp.cs.washington.edu】 TriviaQA is a reading\nL5: comprehension dataset containing over 650K question-answer-evidence triples. \nL6: TriviaQA includes 95K question-answer pairs authored ...\nL7: * 【1†TriviaQA: A Large Scale Distantly Supervised Challenge ...; publish_date:\nL8: none†aclanthology.org】 by M Joshi · 2017 · Cited by 3451 — We present TriviaQA,\nL9: a challenging reading comprehension dataset containing over 650K question-\nL10: answer-evidence triples. TriviaQA includes 95K question ...\nL11: * 【2†mandarjoshi/trivia_qa · Datasets at Hugging Face; publish_date: \nL12: none†huggingface.co】 TriviaqQA is a reading comprehension dataset containing \nL13: over 650K question-answer-evidence triples. TriviaqQA includes 95K question-\nL14: answer pairs authored by ...\nL15: * 【3†[1705.03551] TriviaQA: A Large Scale Distantly Supervised ...; \nL16: publish_date: none†arxiv.org】 by M Joshi · 2017 · Cited by 3451 — We present \nL17: TriviaQA, a challenging reading comprehension dataset containing over 650K \nL18: question-answer-evidence triples.\nL19: * 【4†TriviaQA; publish_date: none†epoch.ai】 An open-domain question answering \nL20: benchmark with challenging trivia questions paired with evidence documents.\nL21: * 【5†TriviaQA Leaderboard; publish_date: none†llm-stats.com】 What is the \nL22: TriviaQA benchmark? A large-scale reading comprehension dataset containing over \nL23: 650K question-answer-evidence triples. TriviaQA includes 95K ...\nL24: * 【6†Code for the TriviaQA reading comprehension dataset; publish_date: \nL25: none†github.com】 A large scale distantly supervised challenge dataset for \nL26: reading comprehension. In Association for Computational Linguistics (ACL) 2017, \nL27: Vancouver, Canada.\nL28: * 【7†TriviaQA - Model Benchmarks - The Regularizer; publish_date: \nL29: none†www.theregularizer.com】 May 4, 2025 — Compare the performance of different \nL30: AI models across standardized benchmarks. Higher scores generally indicate \nL31: better performance, but context ...\nL32: * 【8†TriviaQA: A Large Scale Distantly Supervised Challenge ...; publish_date:\nL33: none†www.cs.utexas.edu】 by M Joshi · Cited by 3445 — We present TriviaQA, a \nL34: challenging reading comprehension dataset contain- ing over 650K question-\nL35: answer-evidence triples. TriviaQA includes 95K question-.",
417 + "tool_call_id": "fixture-tool-call-9-6"
418 + },
419 + {
420 + "role": "tool",
421 + "content": "[11] Natural Questions benchmark (Search_Results/Natural Questions benchmark)\n**viewing lines [0 - 39] of 39**\n\nL0: \nL1: URL: Search_Results/Natural Questions benchmark\nL2: # Search Results\nL3: \nL4: * 【0†Natural Questions: a Benchmark for Question Answering ...; publish_date: \nL5: none†research.google】 by T Kwiatkowski · Cited by 4339 — We present the Natural \nL6: Questions corpus, a question answering dataset. Questions consist of real \nL7: anonymized, aggregated queries issued to the Google search ...\nL8: * 【1†Natural Questions: A Benchmark for Question Answering ...; publish_date: \nL9: none†aclanthology.org】 by T Kwiatkowski · Cited by 4308 — Abstract. We present \nL10: the Natural Questions corpus, a question answering data set. Questions consist \nL11: of real anonymized, aggregated queries issued.\nL12: * 【2†Google's Natural Questions; publish_date: none†ai.google.com】 Natural \nL13: Questions. A Benchmark for Question Answering Research. View examples · Download\nL14: dataset. Open Domain Question Answering. A core goal in artificial ...\nL15: * 【3†google-research-datasets/natural-questions; publish_date: \nL16: none†github.com】 Natural Questions (NQ) contains real user questions issued to \nL17: Google search, and answers found from Wikipedia by annotators. NQ is designed \nL18: for the training and ...\nL19: * 【4†Natural Questions: A Benchmark for Question Answering ...; publish_date: \nL20: none†direct.mit.edu】 Aug 1, 2019 — We present the Natural Questions corpus, a \nL21: question answering data set. Questions consist of real anonymized, aggregated \nL22: queries issued to the Google search ...\nL23: * 【5†ir_datasets : Natural Questions; publish_date: none†ir-datasets.com】 \nL24: Google Natural Questions is a Q&A dataset containing long, short, and Yes/No \nL25: answers from Wikipedia. ir_datasets frames this around an ad-hoc ranking setting\nL26: ...\nL27: * 【6†sentence-transformers/natural-questions · Datasets at ...; publish_date: \nL28: none†huggingface.co】 This dataset is a collection of question-answer pairs from \nL29: the Natural Questions dataset. See Natural Questions for additional information.\nL30: * 【7†Google's Natural Questions; publish_date: none†ai.google.com】 Natural \nL31: Questions contains 307K training examples, 8K examples for development, and a \nL32: further 8K examples for testing. In the paper, we demonstrate a human ...\nL33: * 【8†A Benchmark for Question Answering Research; publish_date: \nL34: none†www.researchgate.net】 Jul 27, 2025 — We present the Natural Questions \nL35: corpus, a question answering data set. Questions consist of real anonymized, \nL36: aggregated queries issued to the Google search ...\nL37: * 【9†natural-questions; publish_date: none†docs.unity.rc.umass.edu】 Sep 4, \nL38: 2025 — “Natural questions: a benchmark for question answering research.” \nL39: Transactions of the Association for Computational Linguistics 7 (2019): ...",
422 + "tool_call_id": "fixture-tool-call-9-7"
423 + },
424 + {
425 + "role": "tool",
426 + "content": "[12] MS MARCO benchmark (Search_Results/MS MARCO benchmark)\n**viewing lines [0 - 41] of 41**\n\nL0: \nL1: URL: Search_Results/MS MARCO benchmark\nL2: # Search Results\nL3: \nL4: * 【0†MS MARCO - Microsoft Open Source; publish_date: none†microsoft.github.io】\nL5: The MS MARCO datasets are intended for non-commercial research purposes only to\nL6: promote advancement in the field of artificial intelligence and related areas, \nL7: ...\nL8: * 【1†microsoft/ms_marco · Datasets at Hugging Face; publish_date: \nL9: none†huggingface.co】 Starting with a paper released at NIPS 2016, MS MARCO is a \nL10: collection of datasets focused on deep learning in search. The first dataset was\nL11: a question ...\nL12: * 【2†Benchmarking Ranking Models in the Large-Data Regime; publish_date: \nL13: none†arxiv.org】 by N Craswell · 2021 · Cited by 89 — This paper uses the MS \nL14: MARCO and TREC Deep Learning Track as our case study, comparing it to the case \nL15: of TREC ad hoc ranking in the 1990s.\nL16: * 【3†Benchmarking Ranking Models in the Large-Data Regime; publish_date: \nL17: none†www.microsoft.com】 This paper uses the MS MARCO and TREC Deep Learning \nL18: Track as our case study, comparing it to the case of TREC ad hoc ranking in the \nL19: 1990s. We show how the ...\nL20: * 【4†Datasets for Document and Passage Ranking Leadboards; publish_date: \nL21: none†microsoft.github.io】 The MS MARCO document and passage ranking leaderboards\nL22: complements the TREC Deep Learning Track by providing on-going evaluation of \nL23: submissions using pre- ...\nL24: * 【5†MS MARCO: Benchmarking Ranking Models in the Large- ...; publish_date: \nL25: none†dl.acm.org】 Jul 11, 2021 — This paper uses the MS MARCO and TREC Deep \nL26: Learning Track as our case study, comparing it to the case of TREC ad hoc \nL27: ranking in the 1990s.\nL28: * 【6†ir_datasets : MSMARCO (passage); publish_date: none†ir-datasets.com】 A \nL29: passage ranking benchmark with a collection of 8.8 million passages and question\nL30: queries. Most relevance judgments are shallow.\nL31: * 【7†MS MARCO; publish_date: none†sbert.net】 MS MARCO Passage Ranking is a \nL32: large dataset to train models for information retrieval. It consists of about \nL33: 500k real search queries from Bing search engine ...\nL34: * 【8†MS MARCO: A Human Generated MAchine Reading ...; publish_date: \nL35: none†arxiv.org】 by P Bajaj · 2016 · Cited by 1151 — We introduce a large scale \nL36: MAchine Reading COmprehension dataset, which we name MS MARCO. The dataset \nL37: comprises of 1,010,916 anonymized ...\nL38: * 【9†MS MARCO Web Search: A Large-scale Information-rich ...; publish_date: \nL39: none†www.microsoft.com】 May 13, 2024 — MS MARCO Web Search offers a retrieval \nL40: benchmark with three web retrieval challenge tasks that demands innovations in \nL41: both machine learning and ...",
427 + "tool_call_id": "fixture-tool-call-9-8"
428 + },
429 + {
430 + "role": "tool",
431 + "content": "[13] BEIR benchmark (Search_Results/BEIR benchmark)\n**viewing lines [0 - 37] of 37**\n\nL0: \nL1: URL: Search_Results/BEIR benchmark\nL2: # Search Results\nL3: \nL4: * 【0†详细介绍文本检索基准BEIR: A Heterogeneous Benchmark ...; publish_date: \nL5: none†blog.csdn.net】 2023年1月1日 — \nL6: BEIR旨在为所有不同的检索任务提供一个一站式的零样本评估基准。为了构建一个全面的评估基准，选择方法对于收集具有理想属性的任务和数据集至关重要。对于 ...\nL7: * 【1†beir-cellar/beir; publish_date: none†github.com】 BEIR is a heterogeneous \nL8: benchmark containing diverse IR tasks. It also provides a common and easy \nL9: framework for evaluation of your NLP-based retrieval models ...\nL10: * 【2†BEIR: A Heterogenous Benchmark for Zero-shot Evaluation ...; \nL11: publish_date: none†arxiv.org】 作者：N Thakur · 2021 · 被引用次数：1480 — We introduce \nL12: Benchmarking-IR (BEIR), a robust and heterogeneous evaluation benchmark for \nL13: information retrieval.\nL14: * 【3†BeIR; publish_date: none†huggingface.co】 BEIR (Benchmarking IR) consists \nL15: of a homogenous benchmark for diverse sentence or passage level IR tasks. It \nL16: provides a common and easy framework for the cross ...\nL17: * 【4†论文分享：BEIR A Heterogeneous Benchmark for Zero-shot ...; publish_date: \nL18: none†zhuanlan.zhihu.com】 2022年10月3日 — 分享论文，夹带个人理解的分享，建议结合原论文看。 1 研究背景. \nL19: 本论文主要关注的领域是query-document检索（下文简称qd检索），即根据query去文档库里 ...\nL20: * 【5†Benchmarking IR Information Retrieval (BEIR); publish_date: \nL21: none†zilliz.com】 BEIR is a benchmark designed for evaluating the versatility and\nL22: robustness of information retrieval models. It features 18 diverse datasets \nL23: from domains like ...\nL24: * 【6†BEIR (Benchmarking IR) - OpenDataLab; publish_date: none†opendatalab.com】\nL25: 简介-Introduction. BEIR（Benchmarking \nL26: IR）是包含不同信息检索（IR）任务的异构基准。通过BEIR，可以系统地研究多种神经检索方法的零样本泛化能力。\nL27: * 【7†What is the BEIR benchmark and how is it used?; publish_date: \nL28: none†milvus.io】 The BEIR (Benchmarking Information Retrieval) benchmark is a \nL29: standardized framework designed to evaluate the effectiveness of search and \nL30: retrieval algorithms.\nL31: * 【8†BEIR Benchmark数据集卡片; publish_date: none†www.atyun.com】 BEIR \nL32: Benchmark数据集卡片. 数据集简介. BEIR是一个异构评测基准，由18个多样化的数据集构建而成，代表了9个信息检索任务：. 事实查证： FEVER ，\nL33: Climate-FEVER ， SciFact ...\nL34: * 【9†Evaluating search relevance part 1 - The BEIR benchmark; publish_date: \nL35: none†www.elastic.co】 2024年7月16日 — Learn to evaluate your search system in the \nL36: context of better understanding the BEIR benchmark, with tips & techniques to \nL37: improve your ...",
432 + "tool_call_id": "fixture-tool-call-9-9"
433 + },
434 + {
435 + "role": "tool",
436 + "content": "[14] MIRACL benchmark (Search_Results/MIRACL benchmark)\n**viewing lines [0 - 41] of 41**\n\nL0: \nL1: URL: Search_Results/MIRACL benchmark\nL2: # Search Results\nL3: \nL4: * 【0†MIRACL | Multilingual Information Retrieval Across a ...; publish_date: \nL5: none†project-miracl.github.io】 MIRACL (Multilingual Information Retrieval Across\nL6: a Continuum of Languages) is an WSDM 2023 Cup challenge that focuses on search \nL7: across 18 different ...\nL8: * 【1†project-miracl/miracl: A large-scale multilingual dataset for ...; \nL9: publish_date: none†github.com】 A large-scale multilingual dataset for \nL10: Information Retrieval. Thorough human-annotations across 18 diverse languages.\nL11: * 【2†A Large, multilingual, visual document retrieval benchmark; publish_date:\nL12: none†arxiv.org】 by R Osmulski · 2025 · Cited by 2 — MIRACL-VISION is a \nL13: challenging, representative, multilingual evaluation benchmark for visual \nL14: retrieval pipelines and will help the community build robust ...\nL15: * 【3†miracl/miracl · Datasets at Hugging Face; publish_date: \nL16: none†huggingface.co】 MIRACL (Multilingual Information Retrieval Across a \nL17: Continuum of Languages) is a multilingual retrieval dataset that focuses on \nL18: search across 18 different ...\nL19: * 【4†MIRACL: A Multilingual Retrieval Dataset Covering 18 ...; publish_date: \nL20: none†direct.mit.edu】 by X Zhang · 2023 · Cited by 131 — MIRACL is a multilingual\nL21: dataset for ad hoc retrieval across 18 languages that collectively encompass \nL22: over three billion native speakers around the world.\nL23: * 【5†(PDF) MIRACL-VISION: A Large, multilingual, visual ...; publish_date: \nL24: none†www.researchgate.net】 May 23, 2025 — MIRACL-VISION covers 18 languages, and\nL25: is an extension of the MIRACL dataset, a popular benchmark to evaluate text-\nL26: based multilingual retrieval ...\nL27: * 【6†A Large, multilingual, visual document retrieval benchmark; publish_date:\nL28: none†arxiv.org】 by R Osmulski · 2025 · Cited by 2 — MIRACL-VISION is a \nL29: challenging, representative, multilingual evaluation benchmark for visual \nL30: retrieval pipelines and will help the community ...\nL31: * 【7†ir_datasets : MIRACL; publish_date: none†ir-datasets.com】 \nL32: \"miracl/ar/test-a\". The held-out test set (version a) for Arabic. \nL33: queriesdocsCitationMetadata. 936 queries. Language: ar. Query type: \nL34: GenericQuery: (namedtuple).\nL35: * 【8†Evaluate on MIRACL — BGE documentation; publish_date: none†bge-model.com】\nL36: MIRACL (Multilingual Information Retrieval Across a Continuum of Languages) is \nL37: an WSDM 2023 Cup challenge that focuses on search across 18 different languages.\nL38: * 【9†MIRACL - Alpha's Tech Garden; publish_date: \nL39: none†techgarden.alphasmanifesto.com】 MIRACL (Multilingual Information Retrieval \nL40: Across a Continuum of Languages) is a multilingual dataset we have built for the\nL41: WSDM 2023 Cup ...",
437 + "tool_call_id": "fixture-tool-call-9-10"
438 + },
439 + {
440 + "role": "tool",
441 + "content": "[15] Zero-shot IR benchmark (Search_Results/Zero-shot IR benchmark)\n**viewing lines [0 - 40] of 40**\n\nL0: \nL1: URL: Search_Results/Zero-shot IR benchmark\nL2: # Search Results\nL3: \nL4: * 【0†BEIR: A Heterogenous Benchmark for Zero-shot Evaluation ...; \nL5: publish_date: none†arxiv.org】 by N Thakur · 2021 · Cited by 1480 — We introduce \nL6: Benchmarking-IR (BEIR), a robust and heterogeneous evaluation benchmark for \nL7: information retrieval.See more\nL8: * 【1†beir-cellar/beir; publish_date: none†github.com】 BEIR: A Heterogenous \nL9: Benchmark for Zero-shot Evaluation of Information Retrieval Models (NeurIPS \nL10: 2021, Datasets and Benchmarks Track); Resources for Brewing ...See more\nL11: * 【2†Benchmarking IR Information Retrieval (BEIR); publish_date: \nL12: none†zilliz.com】 BEIR is a tool to evaluate how well Information Retrieval \nL13: systems perform across many tasks and types of information, and is a standard \nL14: benchmark.\nL15: * 【3†BEIR: A Heterogeneous Benchmark for Zero-shot ...; publish_date: \nL16: none†datasets-benchmarks-proceedings.neurips.cc】 by N Thakur · Cited by 1480 — \nL17: BEIR is a robust, heterogeneous benchmark for information retrieval, using 18 \nL18: datasets and 9 tasks to evaluate model generalization.\nL19: * 【4†BEIR; publish_date: none†eval.ai】 BEIR is a heterogeneous zero-shot \nL20: retrieval benchmark containing 18 datasets from diverse text retrieval tasks and\nL21: domains.See more\nL22: * 【5†[2409.15763] IRSC: A Zero-shot Evaluation Benchmark for ...; \nL23: publish_date: none†arxiv.org】 by H Lin · 2024 · Cited by 2 — This paper \nL24: introduces the IRSC benchmark for evaluating the performance of embedding models\nL25: in multilingual RAG tasks.See more\nL26: * 【6†FactIR: A Real-World Zero-shot Open-Domain Retrieval ...; publish_date: \nL27: none†dl.acm.org】 May 23, 2025 — In this paper, we present a real-world retrieval\nL28: benchmark FactIR, derived from Factiverse production logs, enhanced with human \nL29: annotations. We ...See more\nL30: * 【7†UniIR: Training and Benchmarking Universal Multimodal ...; publish_date: \nL31: none†tiger-ai-lab.github.io】 At test time, we evaluated the zero-shot \nL32: performance of all fine-tuned models, as well as SoTA pre-trained retrievers on \nL33: the three held-out datasets. UniIR ...See more\nL34: * 【8†Zero-Shot BEIR Tasks; publish_date: none†www.emergentmind.com】 Aug 26, \nL35: 2025 — Zero-Shot BEIR Tasks are evaluation methodologies that assess IR models' \nL36: ability to generalize to unseen query domains without task-specific ...See more\nL37: * 【9†BEIR-PL: Zero Shot Information Retrieval Benchmark for ...; publish_date:\nL38: none†aclanthology.org】 by K Wojtasik · 2024 · Cited by 12 — BEIR-PL is a new \nL39: benchmark with 13 datasets for Polish Information Retrieval, created to advance \nL40: research in this area.",
442 + "tool_call_id": "fixture-tool-call-9-11"
443 + },
444 + {
445 + "role": "tool",
446 + "content": "[16] WebGPT benchmark (Search_Results/WebGPT benchmark)\n**viewing lines [0 - 38] of 38**\n\nL0: \nL1: URL: Search_Results/WebGPT benchmark\nL2: # Search Results\nL3: \nL4: * 【0†WebGPT: Improving the factual accuracy of language ...; publish_date: \nL5: none†openai.com】 Dec 16, 2021 — Our models outperform GPT‑3 on TruthfulQA and \nL6: exhibit more favourable scaling properties. However, our models lag behind human\nL7: performance, ...\nL8: * 【1†A Simple Yet Challenging Benchmark for Browsing Agents; publish_date: \nL9: none†arxiv.org】 by J Wei · 2025 · Cited by 124 — Abstract. We present \nL10: BrowseComp, a simple yet challenging benchmark for measuring the ability for \nL11: agents to browse the web.\nL12: * 【2†openai/webgpt_comparisons · Datasets at Hugging Face; publish_date: \nL13: none†huggingface.co】 This is the dataset of all comparisons that were marked as \nL14: suitable for reward modeling by the end of the WebGPT project. There are 19,578 \nL15: comparisons in total.\nL16: * 【3†Evaluation & Limitations of WebGPT, WebVoyager & Agent-E; publish_date: \nL17: none†deepsense.ai】 Oct 14, 2024 — WebArena benchmark features 812 tasks \nL18: evaluated using metrics such as Exact Match, Must Include, and Fuzzy Match, \nL19: focusing on outcomes rather ...\nL20: * 【4†OpenAI Announces Question-Answering AI WebGPT; publish_date: \nL21: none†www.infoq.com】 Jan 25, 2022 — On the TriviaQA benchmark, WebGPT \nL22: outperformed GPT-3, producing answers that were true 75% of the time, and \"both \nL23: true and informative\" 54% of ...\nL24: * 【5†WebGPT: Improving the factual accuracy of language models ...; \nL25: publish_date: none†kargarisaac.medium.com】 The top-performing model generated \nL26: answers that were preferred over 56% of the time compared to answers produced by\nL27: human demonstrators, with ...\nL28: * 【6†Browser-assisted question-answering with human feedback; publish_date: \nL29: none†www.alphaxiv.org】 WebGPT represents a significant advancement in long-form \nL30: question answering by combining the language generation capabilities of GPT-3 \nL31: with real-time web ...\nL32: * 【7†Benchmarking Open-Source Large Language Models, GPT-4 ...; publish_date: \nL33: none†ai.nejm.org】 by S Wu · 2024 · Cited by 69 — We show that the current widely\nL34: used open-source LLMs have poor zero-shot reasoning ability in nephrology \nL35: compared with GPT-4 and Claude 2.\nL36: * 【8†0hq/WebGPT: Run GPT model on ...; publish_date: none†github.com】 WebGPT \nL37: is a vanilla JS and HTML implementation of a transformer model, intended as a \nL38: proof-of-concept as well as educational resource.",
447 + "tool_call_id": "fixture-tool-call-9-12"
448 + },
449 + {
450 + "role": "tool",
451 + "content": "[17] WebShop benchmark (Search_Results/WebShop benchmark)\n**viewing lines [0 - 41] of 41**\n\nL0: \nL1: URL: Search_Results/WebShop benchmark\nL2: # Search Results\nL3: \nL4: * 【0†WebShop: Towards Scalable Real-World Web Interaction ...; publish_date: \nL5: none†arxiv.org】 by S Yao · 2022 · Cited by 710 — To bridge this gap, we develop \nL6: WebShop -- a simulated e-commerce website environment with 1.18 million real-\nL7: world products and 12,087 crowd- ...\nL8: * 【1†WebShop; publish_date: none†webshop-pnlp.github.io】 To bridge this gap, \nL9: we develop WebShop – a simulated e-commerce website environment with 1.18 \nL10: million real-world products and 12,087 crowd-sourced text ...\nL11: * 【2†princeton-nlp/WebShop; publish_date: none†github.com】 WebShop is a \nL12: simulated e-commerce website environment with 1.18 million real-world products \nL13: and 12,087 crowd-sourced text instructions. In this environment, an ...\nL14: * 【3†WebShop: Towards Scalable Real-World Web Interaction ...; publish_date: \nL15: none†papers.nips.cc】 by S Yao · 2022 · Cited by 710 — We collect over 1,600 \nL16: human trajectories to first validate the benchmark, then train and evaluate a \nL17: diverse range of agents using reinforcement learning, ...\nL18: * 【4†WebShop: Towards Scalable Real-World Web Interaction ...; publish_date: \nL19: none†proceedings.neurips.cc】 by S Yao · 2022 · Cited by 709 — We have developed \nL20: WebShop, a new web-based benchmark for sequential decision making and language \nL21: grounding, modeled on interaction with an e-commerce website.\nL22: * 【5†Webshop & Benchmark Analysis | Documentation Infinity; publish_date: \nL23: none†docs.fact-finder.com】 Aug 15, 2025 — Evaluation of your shop based on \nL24: different categories in comparison, to your competitors/industry. Recommended \nL25: when doing a shop relaunch.\nL26: * 【6†A Multi-Shop Benchmark for Evaluating Web Agents; publish_date: \nL27: none†arxiv.org】 by R Peeters · 2025 · Cited by 2 — Compared to existing \nL28: e-commerce benchmarks, such as WebShop or ShoppingBench, WebMall introduces \nL29: comparison-shopping tasks across multiple shops ...\nL30: * 【7†WebShop: towards scalable real-world web interaction with ...; \nL31: publish_date: none†dl.acm.org】 by S Yao · 2022 · Cited by 710 — To bridge this \nL32: gap, we develop WebShop - a simulated e-commerce website environment with 1.18 \nL33: million real-world products and 12, 087 crowd- ...\nL34: * 【8†[PDF] WebShop: Towards Scalable Real-World Web ...; publish_date: \nL35: none†www.semanticscholar.org】 It is shown that agents trained on WebShop exhibit\nL36: non-trivial sim-to-real transfer when evaluated on amazon.com and ebay.com, \nL37: indicating the potential ...\nL38: * 【9†X-WebAgentBench: A Multilingual Interactive Web ...; publish_date: \nL39: none†aclanthology.org】 by P Wang · 2025 · Cited by 3 — (2023) based on the \nL40: English WebShop benchmark (Yao et al., 2022), while the multilingual task scores\nL41: are ob- tained through evaluation on our own benchmark.",
452 + "tool_call_id": "fixture-tool-call-9-13"
453 + },
454 + {
455 + "role": "tool",
456 + "content": "[18] ALFWorld benchmark (Search_Results/ALFWorld benchmark)\n**viewing lines [0 - 31] of 31**\n\nL0: \nL1: URL: Search_Results/ALFWorld benchmark\nL2: # Search Results\nL3: \nL4: * 【0†ALFWorld; publish_date: none†alfworld.github.io】 ALFWorld contains \nL5: interactive TextWorld environments (Côté et. al) that parallel embodied worlds \nL6: in the ALFRED dataset (Shridhar et. al).\nL7: * 【1†ALFWorld: Aligning Text and Embodied Environments for ...; publish_date: \nL8: none†arxiv.org】 by M Shridhar · 2020 · Cited by 674 — ALFWorld enables the \nL9: creation of a new BUTLER agent whose abstract knowledge, learned in TextWorld, \nL10: corresponds directly to concrete, visually grounded actions.\nL11: * 【2†ALFWorld: Aligning Text and Embodied Environments ...; publish_date: \nL12: none†github.com】 ALFWorld contains interactive TextWorld environments (Côté et. \nL13: al) that parallel embodied worlds in the ALFRED dataset (Shridhar et. al).\nL14: * 【3†alfworld - benchmark's activity; publish_date: none†huggingface.co】 MM-\nL15: IQ: Benchmarking Human-Like Abstraction and Reasoning in Multimodal Models Paper\nL16: • 2502.00698 • Published Feb 1 • 24\nL17: * 【4†Tackling AlfWorld with Action Attention and Common ...; publish_date: \nL18: none†neurips.cc】 On the Alfworld benchmark for indoor instruction following, we \nL19: achieve a significantly higher success rate (50% over the baseline) with our \nL20: novel object ...\nL21: * 【5†ALFWORLD: ALIGNING TEXT AND EMBODIED ...; publish_date: \nL22: none†openreview.net】 by M Shridhar · Cited by 674 — The ALFRED dataset (Shridhar\nL23: et al., 2020), set in the THOR simulator (Kolve et al., 2017), is a benchmark \nL24: for learning to com- plete embodied household tasks ...\nL25: * 【6†AlfWorld; publish_date: none†primo.ai】 Mar 23, 2024 — A simulator that \nL26: enables agents to learn abstract, text based policies in TextWorld (Côté et al.,\nL27: 2018) and then execute goals from the ALFRED benchmark.\nL28: * 【7†AlfWorld performance across 134 tasks showing cumulative...; \nL29: publish_date: none†www.researchgate.net】 In the AlfWorld benchmark, we defined \nL30: hallucination as the occurrence of two or more consecutive identical actions in \nL31: which the environment responded with ...",
457 + "tool_call_id": "fixture-tool-call-9-14"
458 + },
459 + {
460 + "role": "tool",
461 + "content": "[19] Mind2Web benchmark (Search_Results/Mind2Web benchmark)\n**viewing lines [0 - 40] of 40**\n\nL0: \nL1: URL: Search_Results/Mind2Web benchmark\nL2: # Search Results\nL3: \nL4: * 【0†Mind2Web: Towards a Generalist Agent for the Web; publish_date: none†osu-\nL5: nlp-group.github.io】 Mind2Web is a dataset for developing and evaluating \nL6: generalist agents for the web that can follow language instructions to complete \nL7: complex tasks on any ...\nL8: * 【1†Online-Mind2Web Leaderboard; publish_date: none†huggingface.co】 Online-\nL9: Mind2Web is a benchmark designed to evaluate the real-world performance of web \nL10: agents on live websites, featuring 300 tasks across 136 popular sites ...\nL11: * 【2†Mind2Web: Towards a Generalist Agent for the Web; publish_date: \nL12: none†github.com】 Mind2Web is the first dataset for developing and evaluating \nL13: generalist agents for the web that can follow language instructions to complete \nL14: complex tasks on any ...\nL15: * 【3†HAL: Online Mind2Web Leaderboard; publish_date: \nL16: none†hal.cs.princeton.edu】 Online Mind2Web leaderboard for evaluating AI agents'\nL17: ability to complete tasks on real, changing webpages.\nL18: * 【4†[2506.21506] Mind2Web 2: Evaluating Agentic Search with ...; \nL19: publish_date: none†arxiv.org】 by B Gou · 2025 · Cited by 11 — In this paper, we \nL20: introduce Mind2Web 2, a benchmark of 130 realistic, high-quality, and long-\nL21: horizon tasks that require real-time web browsing and extensive ...\nL22: * 【5†Mind2Web 2: Evaluating Agentic Search with Agent-as-a-Judge; \nL23: publish_date: none†osu-nlp-group.github.io】 We introduce Mind2Web 2, a benchmark\nL24: of 130 realistic, high-quality, long-horizon tasks that require real-time web \nL25: browsing and extensive information ...\nL26: * 【6†Mind2Web: The Benchmark for AI Agent Evaluation and ...; publish_date: \nL27: none†www.enhans.ai】 Sep 26, 2025 — Mind2Web is a globally recognized web-based \nL28: AI Agent evaluation benchmark introduced by the NLP group at Ohio State \nL29: University at NeurIPS 2023.\nL30: * 【7†Evaluating AI Web Agents: Insights from the WebCanvas ...; publish_date: \nL31: none†medium.com】 Thanks to the comprehensive WebCanvas Benchmark, which \nL32: incorporates a robust Mind2Web-Live data set of 542 live web tasks and 2,439 ...\nL33: * 【8†Mind2Web: Towards a Generalist Agent for the Web; publish_date: \nL34: none†proceedings.neurips.cc】 by X Deng · 2023 · Cited by 760 — We introduce \nL35: Mind2Web, the first dataset for developing and evaluating generalist agents for \nL36: the web that can follow language instructions to complete complex ...\nL37: * 【9†Mind2Web: Towards a Generalist Agent for the Web; publish_date: \nL38: none†openreview.net】 by X Deng · Cited by 760 — We introduce Mind2Web, the first\nL39: dataset for developing and evaluating generalist agents for the web that can \nL40: follow language instructions to complete complex ...",
462 + "tool_call_id": "fixture-tool-call-9-15"
463 + },
464 + {
465 + "role": "tool",
466 + "content": "[20] VisualWebArena benchmark (Search_Results/VisualWebArena benchmark)\n**viewing lines [0 - 38] of 38**\n\nL0: \nL1: URL: Search_Results/VisualWebArena benchmark\nL2: # Search Results\nL3: \nL4: * 【0†VisualWebArena is a benchmark for multimodal agents.; publish_date: \nL5: none†github.com】 VisualWebArena is a realistic and diverse benchmark for \nL6: evaluating multimodal autonomous language agents.\nL7: * 【1†Evaluating Multimodal Agents on Realistic Visual Web Tasks; publish_date:\nL8: none†arxiv.org】 by JY Koh · 2024 · Cited by 363 — To bridge this gap, we \nL9: introduce VisualWebArena, a benchmark designed to assess the performance of \nL10: multimodal web agents on realistic \\textit{ ...\nL11: * 【2†Evaluating Multimodal Agents on Realistic Visual Web Tasks; publish_date:\nL12: none†jykoh.com】 To bridge this gap, we introduce VisualWebArena, a benchmark \nL13: designed to assess the performance of multimodal web agents on realistic \nL14: visually grounded tasks.\nL15: * 【3†VisualWebArena: Evaluating Multimodal Agents on ...; publish_date: \nL16: none†arxiv.org】 VisualWebArena is a research benchmark to measure and evaluate \nL17: the progress of multimodal agents. It is primarily meant to act as a self-\nL18: contained sandbox ...\nL19: * 【4†Evaluating Multimodal Agents on Realistic Visual Web Tasks; publish_date:\nL20: none†aclanthology.org】 by JY Koh · 2024 · Cited by 363 — To bridge this gap, we\nL21: introduce VisualWebArena, a benchmark designed to assess the performance of \nL22: multimodal web agents on *realistic visually grounded tasks*.\nL23: * 【5†Evaluating Multimodal Agents on Realistic Visual Web Tasks; publish_date:\nL24: none†www.semanticscholar.org】 VisualWebArena: Evaluating Multimodal Agents on \nL25: Realistic Visual Web Tasks ... MMInA, a multihop and multimodal benchmark to \nL26: evaluate the embodied agents ...\nL27: * 【6†CMU Researchers Introduce VisualWebArena: An AI ...; publish_date: \nL28: none†www.marktechpost.com】 Feb 9, 2024 — VisualWebArena, a benchmark designed \nL29: and developed to evaluate the performance of multimodal web agents on realistic \nL30: and visually stimulating challenges.\nL31: * 【7†Evaluating Multimodal Agents on Realistic Visual Web Tasks; publish_date:\nL32: none†www.themoonlight.io】 The paper \"VisualWebArena: Evaluating Multimodal \nL33: Agents on Realistic Visually Grounded Web Tasks\" introduces a new benchmark, \nL34: **VisualWebArena**, ...\nL35: * 【8†WebArena: A Realistic Web Environment for Building ...; publish_date: \nL36: none†webarena.dev】 Our benchmark is implemented in our fully interactable \nL37: highly-realistic WebArena environment. It features diverse tasks human may \nL38: encounter in their daily ...",
467 + "tool_call_id": "fixture-tool-call-9-16"
468 + },
469 + {
470 + "role": "tool",
471 + "content": "[21] SearchBench benchmark (Search_Results/SearchBench benchmark)\n**viewing lines [0 - 40] of 40**\n\nL0: \nL1: URL: Search_Results/SearchBench benchmark\nL2: # Search Results\nL3: \nL4: * 【0†Talc-AI/search-bench; publish_date: none†github.com】 A practical \nL5: benchmark that focuses on every day helpfulness of LLM products, not just the \nL6: underlying models. Searchbench is a benchmark that addresses these ...\nL7: * 【1†Evaluating LLMs' Ability to Reason About Search Problems; publish_date: \nL8: none†arxiv.org】 These capabilities are essential for robust reasoning, making \nL9: SearchBench a valuable benchmark for evaluating LLMs' reasoning capabilities as \nL10: they continue to ...\nL11: * 【2†NasimBrz/SearchBench · Datasets at Hugging Face; publish_date: \nL12: none†huggingface.co】 Dataset Summary. SearchBench is a benchmark designed to \nL13: evaluate Language Models' (LLMs) ability to solve state-based problems that \nL14: require combinatorial search ...\nL15: * 【3†Evaluating LLMs' Ability to Reason About Search Problems; publish_date: \nL16: none†openreview.net】 2025年10月22日 — To further investigate this, we introduce a \nL17: new benchmark, SearchBench, which contains 11 unique search problems inspired by\nL18: intuitive puzzles.\nL19: * 【4†Navigating the Labyrinth: Evaluating and Enhancing LLMs' ...; \nL20: publish_date: none†hub.baai.ac.cn】 2024年6月17日 — \nL21: 论文提出了一个新的基准测试SearchBench，包含11种独特的搜索问题类型，并自动化生成任意数量的实例和分析解决方案的可行性、正确性和最优性。论文使用A* \nL22: ...\nL23: * 【5†Towards Unified Text-based Person Retrieval: A Large- ...; publish_date: \nL24: none†blog.csdn.net】 2023年10月17日 — ... Search \nL25: Benchmark（面向统一的基于文本的人物检索：一个大规模的多属性和语言搜索基准）; 研究背景. 相关工作; BENCHMARK. 论文方法分析. 网络框架;\nL26: 1、APTM ...\nL27: * 【6†Desearch-ai/ai-search-benchmark; publish_date: none†github.com】 The \nL28: SearchBench repository addresses common issues with traditional benchmarks by \nL29: focusing on practical, everyday use cases rather than theoretical limits. It ...\nL30: * 【7†o1 results for 3 benchmarks: PlanBench, SearchBench, ...; publish_date: \nL31: none†www.reddit.com】 o1 results for 3 benchmarks: PlanBench, SearchBench, and \nL32: Summary of a Haystack. AI. PlanBench: Paper \"LLMs Still Can't Plan; Can LRMs? A \nL33: ...\nL34: * 【8†Evaluating LLMs' Ability to Reason About Search Problems; publish_date: \nL35: none†ui.adsabs.harvard.edu】 To further investigate this, we introduce a new \nL36: benchmark, SearchBench, which contains 11 unique search problems inspired by \nL37: intuitive puzzles. Each SearchBench ...\nL38: * 【9†Introducing SearchBench; publish_date: none†www.tag1consulting.com】 \nL39: Toward this goal, over the weekend I launched a new project called SearchBench, \nL40: a Drupal module for benchmarking Drupal's search performance. As the module ...",
472 + "tool_call_id": "fixture-tool-call-9-17"
473 + },
474 + {
475 + "role": "tool",
476 + "content": "[22] WebVLN benchmark (Search_Results/WebVLN benchmark)\n**viewing lines [0 - 42] of 42**\n\nL0: \nL1: URL: Search_Results/WebVLN benchmark\nL2: # Search Results\nL3: \nL4: * 【0†WebVLN: Vision-and-Language Navigation on Websites; publish_date: \nL5: none†ojs.aaai.org】 by Q Chen · 2024 · Cited by 35 — the WebVLN-v1 dataset, where\nL6: the performance is far from saturation, highlighting the utility of our \nL7: WebVLN-v1 as a benchmark to assess progress in this field.\nL8: * 【1†[2312.15820] WebVLN: Vision-and-Language Navigation on Websites; \nL9: publish_date: none†ar5iv.labs.arxiv.org】 Experimental results show that WebVLN-\nL10: Net outperforms current VLN and web-related navigation methods. ... Code is \nL11: available at: https://github.com/WebVLN/WebVLN.\nL12: * 【2†WebVLN: Vision-and-Language Navigation on Websites; publish_date: \nL13: none†github.com】 Experimental results show that WebVLN-Net outperforms current \nL14: VLN and web-related navigation methods. We believe that the introduction of the \nL15: new WebVLN task ...\nL16: * 【3†Vision-and-Language Navigation in the Real-World; publish_date: \nL17: none†digital.library.adelaide.edu.au】 By leveraging our proposed WebVLN-v1 \nL18: dataset, experimental results showcase the superior performance of WebVLN-Net \nL19: compared to existing VLN and web-related ...\nL20: * 【4†WebVLN: Vision-and-Language Navigation on Websites; publish_date: \nL21: none†www.researchgate.net】 Experimental results show that WebVLN-Net outperforms\nL22: current VLN and web-related navigation methods. We believe that the \nL23: introduction of the newWebVLN task and ...\nL24: * 【5†[PDF] WebVLN: Vision-and-Language Navigation on Websites; publish_date: \nL25: none†www.semanticscholar.org】 A new task named Vision-and-Language Navigation on\nL26: Websites (WebVLN), where question-based instructions are used to train an \nL27: agent, emulating how users ...\nL28: * 【6†WebVLN: Vision-and-Language Navigation on Websites; publish_date: \nL29: none†arxiv.org】 by Q Chen · 2023 · Cited by 35 — Experimental results show that \nL30: WebVLN-Net outperforms current VLN and web-related navigation methods. We \nL31: believe that the introduction of the ...\nL32: * 【7†Human-Aware Vision-and-Language Navigation; publish_date: \nL33: none†proceedings.neurips.cc】 by H Li · 2024 · Cited by 19 — Vision-and-Language \nL34: Navigation (VLN) [2, 7, 9, 40] has emerged as a key benchmark for evaluating. \nL35: Sim2Real transfer [23], showing impressive performance in ...\nL36: * 【8†LiveBench; publish_date: none†livebench.ai】 Introducing LiveBench: a \nL37: benchmark for LLMs designed with test set contamination and objective evaluation\nL38: in mind.\nL39: * 【9†MG-VLN: Benchmarking Multi-Goal and Long-Horizon ...; publish_date: \nL40: none†ieeexplore.ieee.org】 by J Zhang · 2024 — This task aims to provide a \nL41: simulation benchmark to guide the design of lifelong and long-horizon navigation\nL42: robots.",
477 + "tool_call_id": "fixture-tool-call-9-18"
478 + },
479 + {
480 + "role": "tool",
481 + "content": "[23] WebNav benchmark (Search_Results/WebNav benchmark)\n**viewing lines [0 - 36] of 36**\n\nL0: \nL1: URL: Search_Results/WebNav benchmark\nL2: # Search Results\nL3: \nL4: * 【0†WebNav: A New Large-Scale Task for Natural Language ...; publish_date: \nL5: none†github.com】 WebNav is a benchmark task for evaluating an agent with \nL6: abilities to understand natural language and plan on partially observed \nL7: environments.\nL8: * 【1†[1602.02261] End-to-End Goal-Driven Web Navigation; publish_date: \nL9: none†arxiv.org】 by R Nogueira · 2016 · Cited by 39 — We propose a goal-driven \nL10: web navigation as a benchmark task for evaluating an agent with abilities to \nL11: understand natural language and plan on partially ...\nL12: * 【2†nyu-dl/dl4ir-webnav; publish_date: none†github.com】 WebNav is a benchmark\nL13: task for evaluating an agent with abilities to understand natural language and \nL14: plan on partially observed environments.\nL15: * 【3†WebNav: A New Large-Scale Task for Natural Language ...; publish_date: \nL16: none†www.researchgate.net】 We propose a goal-driven web navigation as a \nL17: benchmark task for evaluating an agent with abilities to understand natural \nL18: language and plan on partially ...\nL19: * 【4†WebNav: An Intelligent Agent for Voice-Controlled Web ...; publish_date: \nL20: none†arxiv.org】 In this paper, we introduce WebNav, an innovative multi-modal \nL21: agent designed to automate web tasks based on natural language.\nL22: * 【5†WebCanvas: Benchmarking Web Agents in Online ...; publish_date: \nL23: none†openreview.net】 by Y Pan · Cited by 78 — TL;DR: We introduce WebCanvas, an \nL24: online evaluation framework for web agents designed to address the dynamic \nL25: nature of web interactions.\nL26: * 【6†WebNav: An Intelligent Agent for Voice-Controlled Web ...; publish_date: \nL27: none†huggingface.co】 Mar 17, 2025 — Preliminary evaluations show that WebNav \nL28: outperforms traditional screen readers in response time and task completion \nL29: accuracy for the visually ...\nL30: * 【7†Web Bench - A new way to compare AI Browser Agents; publish_date: \nL31: none†www.skyvern.com】 May 29, 2025 — TL;DR: Web Bench is a new dataset to \nL32: evaluate web browsing agents that consists of 5,750 tasks on 452 different \nL33: websites, with 2,454 tasks ...\nL34: * 【8†Web Framework Benchmarks; publish_date: none†www.techempower.com】 In the \nL35: following tests, we have measured the performance of several web application \nL36: platforms, full-stack frameworks, and micro-frameworks.",
482 + "tool_call_id": "fixture-tool-call-9-19"
483 + },
484 + {
485 + "role": "assistant",
486 + "content": "",
487 + "reasoning_content": "我们收到了大量关于各种基准测试的搜索结果。现在，我们需要整理一份针对搜索智能体的基准测试列表。这些基准测试可以分为几类：网络浏览智能体基准测试（例如BrowseComp、WebVoyager、WebArena、VisualWebArena、Mind2Web、WebVLN、WebNav）、工具使用智能体基准测试（例如ToolBench、GTA）、问答/检索基准测试（例如HotpotQA、FEVER、TriviaQA、Natural Questions、MS MARCO、BEIR、MIRACL），以及特定领域的搜索基准测试（例如HSCodeComp、WebShop、ALFWorld）。此外，还有一些通用智能体基准测试，例如GAIA、SuperCLUE-Agent等。用户要求“详细介绍各自的特点、使用场景、例题”。因此，我需要为每个基准测试提供详细的描述，包括其特点、使用场景，并举例说明。为了全面获取这些信息，我将同时打开以下基准测试的相关页面：WebVoyager（索引3）、WebArena（索引1）、GAIA（索引1）、ToolBench（索引0）、HotpotQA（索引0）、FEVER（索引1）、TriviaQA（索引0）、Natural Questions（索引0）、MS MARCO（索引0）、BEIR（索引1）、MIRACL（索引0）、WebShop（索引0）、ALFWorld（索引0）、Mind2Web（索引0）、VisualWebArena（索引0）、SearchBench（索引1）、WebVLN（索引1）和WebNav（索引0）。",
488 + "tool_calls": [
489 + {
490 + "id": "fixture-tool-call-30-0",
491 + "type": "function",
492 + "function": {
493 + "name": "open",
494 + "arguments": "{\"cursor\": 4, \"id\": 3}"
495 + }
496 + },
497 + {
498 + "id": "fixture-tool-call-30-1",
499 + "type": "function",
500 + "function": {
501 + "name": "open",
502 + "arguments": "{\"cursor\": 5, \"id\": 1}"
503 + }
504 + },
505 + {
506 + "id": "fixture-tool-call-30-2",
507 + "type": "function",
508 + "function": {
509 + "name": "open",
510 + "arguments": "{\"cursor\": 6, \"id\": 1}"
511 + }
512 + },
513 + {
514 + "id": "fixture-tool-call-30-3",
515 + "type": "function",
516 + "function": {
517 + "name": "open",
518 + "arguments": "{\"cursor\": 7, \"id\": 0}"
519 + }
520 + },
521 + {
522 + "id": "fixture-tool-call-30-4",
523 + "type": "function",
524 + "function": {
525 + "name": "open",
526 + "arguments": "{\"cursor\": 8, \"id\": 0}"
527 + }
528 + },
529 + {
530 + "id": "fixture-tool-call-30-5",
531 + "type": "function",
532 + "function": {
533 + "name": "open",
534 + "arguments": "{\"cursor\": 9, \"id\": 1}"
535 + }
536 + },
537 + {
538 + "id": "fixture-tool-call-30-6",
539 + "type": "function",
540 + "function": {
541 + "name": "open",
542 + "arguments": "{\"cursor\": 10, \"id\": 0}"
543 + }
544 + },
545 + {
546 + "id": "fixture-tool-call-30-7",
547 + "type": "function",
548 + "function": {
549 + "name": "open",
550 + "arguments": "{\"cursor\": 11, \"id\": 0}"
551 + }
552 + },
553 + {
554 + "id": "fixture-tool-call-30-8",
555 + "type": "function",
556 + "function": {
557 + "name": "open",
558 + "arguments": "{\"cursor\": 12, \"id\": 0}"
559 + }
560 + },
561 + {
562 + "id": "fixture-tool-call-30-9",
563 + "type": "function",
564 + "function": {
565 + "name": "open",
566 + "arguments": "{\"cursor\": 13, \"id\": 1}"
567 + }
568 + },
569 + {
570 + "id": "fixture-tool-call-30-10",
571 + "type": "function",
572 + "function": {
573 + "name": "open",
574 + "arguments": "{\"cursor\": 14, \"id\": 0}"
575 + }
576 + },
577 + {
578 + "id": "fixture-tool-call-30-11",
579 + "type": "function",
580 + "function": {
581 + "name": "open",
582 + "arguments": "{\"cursor\": 17, \"id\": 0}"
583 + }
584 + },
585 + {
586 + "id": "fixture-tool-call-30-12",
587 + "type": "function",
588 + "function": {
589 + "name": "open",
590 + "arguments": "{\"cursor\": 18, \"id\": 0}"
591 + }
592 + },
593 + {
594 + "id": "fixture-tool-call-30-13",
595 + "type": "function",
596 + "function": {
597 + "name": "open",
598 + "arguments": "{\"cursor\": 19, \"id\": 0}"
599 + }
600 + },
601 + {
602 + "id": "fixture-tool-call-30-14",
603 + "type": "function",
604 + "function": {
605 + "name": "open",
606 + "arguments": "{\"cursor\": 20, \"id\": 0}"
607 + }
608 + },
609 + {
610 + "id": "fixture-tool-call-30-15",
611 + "type": "function",
612 + "function": {
613 + "name": "open",
614 + "arguments": "{\"cursor\": 21, \"id\": 1}"
615 + }
616 + },
617 + {
618 + "id": "fixture-tool-call-30-16",
619 + "type": "function",
620 + "function": {
621 + "name": "open",
622 + "arguments": "{\"cursor\": 22, \"id\": 1}"
623 + }
624 + },
625 + {
626 + "id": "fixture-tool-call-30-17",
627 + "type": "function",
628 + "function": {
629 + "name": "open",
630 + "arguments": "{\"cursor\": 23, \"id\": 0}"
631 + }
632 + }
633 + ]
634 + },
635 + {
636 + "role": "tool",
637 + "content": "[24] (https://ar5iv.org/abs/2401.13919)\n**viewing lines [0 - 129] of 825**\n\nL0: \nL1: URL: https://ar5iv.org/abs/2401.13919\nL2: # WebVoyager : Building an End-to-End Web Agent with\nL3: Large Multimodal Models\nL4: \nL5: Hongliang He1,3 , Wenlin Yao2, Kaixin Ma2, Wenhao Yu2, Yong Dai2,\nL6: \nL7: Hongming Zhang2,\nL8: Zhenzhong Lan3,\nL9: Dong Yu2\nL10: 1Zhejiang University,\nL11: 2Tencent AI Lab,\nL12: 3Westlake University\nL13: hehongliang@westlake.edu.cn, wenlinyao@global.tencent.com\nL14: Work done during the internship at Tencent AI Lab.\nL15: \nL16: ###### Abstract\nL17: \nL18: The advancement of large language models (LLMs) leads to a new era marked by the\nL19: development of autonomous applications in the real world, which drives \nL20: innovation in the creation of advanced web-based agents. Existing web agents \nL21: typically only handle one input modality and are evaluated only in simplified \nL22: web simulators or static web snapshots, greatly limiting their applicability in \nL23: real-world scenarios. To bridge this gap, we introduce WebVoyager, an innovative\nL24: Large Multimodal Model (LMM) powered web agent that can complete user \nL25: instructions end-to-end by interacting with real-world websites. Moreover, we \nL26: propose a new evaluation protocol for web agents to address the challenges of \nL27: automatic evaluation of open-ended web agent tasks, leveraging the robust \nL28: multimodal comprehension capabilities of GPT-4V. We create a new benchmark by \nL29: gathering real-world tasks from 15 widely used websites to evaluate our agents. \nL30: We show that WebVoyager achieves a 55.7% task success rate, significantly \nL31: surpassing the performance of both GPT-4 (All Tools) and the WebVoyager (text-\nL32: only) setups, underscoring the exceptional capability of WebVoyager in practical\nL33: applications. We found that our proposed automatic evaluation achieves 85.3% \nL34: agreement with human judgment, paving the way for further development of web \nL35: agents in a real-world setting.111Our code and data will be released at \nL36: https://github.com/MinorJerry/WebVoyager\nL37: \nL38: ## 1 Introduction\nL39: \nL40: The recent advancement of large language models (LLMs), such as ChatGPT and \nL41: GPT-4 (OpenAI, 2023), have sparked significant interest in developing LLM-based \nL42: autonomous agents (AutoGPT, 2022) for complex task execution (Qin et al., 2023; \nL43: Schick et al., 2023). Recent studies have explored the construction of text-\nL44: based web browsing environments and how to instruct large language model agents \nL45: to perform web navigation (Nakano et al., 2021; Gur et al., 2023; Zhou et al., \nL46: 2023; Lu et al., 2023). The primary challenge in these works lies in managing \nL47: complex and verbose HTML texts, and solutions include simplifying and \nL48: structuring HTML (Nakano et al., 2021; Zhou et al., 2023; Gur et al., 2023; Deng\nL49: et al., 2023).\nL50: \nL51: However, existing approaches overlook a critical functionality of browsing: \nL52: rendering HTML into visual webpages. Particularly, vision capability is crucial \nL53: for utilizing tools like web browsers, as rendered web pages are inherently \nL54: designed with user experience (UX), emphasizing intuitive information and \nL55: structured presentation. This design principle of rendering makes visual \nL56: analysis more effective than mere HTML representation. At present, large \nL57: multimodal models (LMMs), particularly GPT-4V(ision) (OpenAI, 2023) and Gemini \nL58: (Team et al., 2023), demonstrate a remarkable ability to integrate intricate \nL59: visual cues with textual information. Existing studies such as Pix2Struct (Lee \nL60: et al., 2023) and WebArena (Zhou et al., 2023), have initiated explorations into\nL61: using screenshots as inputs for decision-making in web navigation, yet these \nL62: are preliminary and do not represent a deep exploration. Therefore, building \nL63: multimodal web agents to leverage the environment rendered by browsers through \nL64: screenshots, thus mimicking human web browsing behavior, is now a viable \nL65: approach to enhance web navigation efficiency.\nL66: \nL67: We introduce WebVoyager, a multimodal web agent designed to handle web tasks \nL68: online in an end-to-end manner, which denotes managing the process from start to\nL69: finish autonomously without intermediate human intervention. We construct an \nL70: online environment using Selenium for WebVoyager, feeding it with screenshots \nL71: and textual content in interactive web elements. Inspired by Set-of-Mark \nL72: Prompting (Yang et al., 2023a), we mark interactive web elements on screenshots \nL73: (see Figure 2) to facilitate decision-making for WebVoyager. As a pioneer in \nL74: combining vision and text information during web navigation, we advocate that \nL75: autonomous end-to-end task completion, multimodal capabilities and online \nL76: navigation constitute the essential trajectory toward the genuine intelligence \nL77: of web agents.\nL78: \nL79: Another challenge arises when it comes to evaluating an end-to-end web agent \nL80: with online navigation. Existing benchmarks, such as Mind2Web (Deng et al., \nL81: 2023), primarily focus on stepwise and offline evaluation, where agents follow \nL82: predefined “golden” trajectory for action selection. This approach, however, may\nL83: not fully account for the variety of viable strategies to accomplish a task, as\nL84: it only reflects one possible plan. This limitation could lead to a biased \nL85: evaluation and difficulties in fairly comparing different methods. To more \nL86: accurately gauge the capabilities of web agents in end-to-end task completion, \nL87: we save screenshots throughout the online navigation process, and then use \nL88: GPT-4V to evaluate these trajectories and the final results automatically. Human\nL89: evaluations are also conducted to verify the results and confirm the \nL90: reliability of GPT-4V as the evaluator.\nL91: \nL92: We conduct evaluations on a collected dataset, which is semi-automatically \nL93: generated using a self-instruct (Wang et al., 2022) method, comprising 300 web \nL94: tasks from 15 commonly accessed websites. Additionally, we extract 90 web-\nL95: related tasks of level 1 and level 2 from the GAIA (Mialon et al., 2023) to \nL96: enrich our evaluation. We compare our WebVoyager with 1) GPT-4 (All \nL97: Tools)222GPT-4 (All Tools) is an integrated tool-based agent released by OpenAI \nL98: in Oct. 2023. See https://chat.openai.com/, and 2) WebVoyager in a text-only \nL99: setting, employing the accessibility tree proposed in WebArena (Zhou et al., \nL100: 2023) to describe web pages. The results show that WebVoyager achieves a Task \nL101: Success Rate of 55.7%, significantly outperforming GPT-4 (All Tools) with a rate\nL102: of 32.7% and the text-only setting with a rate of 39.0%, demonstrating the \nL103: effectiveness of our method. Furthermore, we report the consistency between \nL104: human-human and human-GPT4V to ensure credibility. Our main contributions are as\nL105: follows:\nL106: \nL107: - •\nL108: We employ a multimodal web agent that integrates textual and visual information \nL109: to address web tasks end-to-end and introduce a generalist planning approach for\nL110: navigation.\nL111: \nL112: - •\nL113: We build an online web browsing environment, offering a variety of tasks \nL114: centered on widely used websites and introducing a method for expanding these \nL115: tasks.\nL116: \nL117: - •\nL118: We conduct manual evaluations of navigation trajectories and propose an \nL119: automated evaluation protocol using GPT-4V. We present a comprehensive analysis \nL120: of the evaluation results and show that GPT-4V can serve as a reliable evaluator\nL121: for online agents.\nL122: \nL123: ## 2 Related Work\nL124: \nL125: ### 2.1 Web Agents\nL126: \nL127: Autonomous web navigation (Shi et al., 2017; Yang et al., 2023b) requires an \nL128: agent to follow instructions, construct plans, comprehend complex web \nL129: structures, and decompose tasks into step-by-step decisions (Weng, 2023). To ",
638 + "tool_call_id": "fixture-tool-call-30-0"
639 + },
640 + {
641 + "role": "tool",
642 + "content": "[25] (https://ar5iv.org/abs/2307.13854)\n**viewing lines [0 - 124] of 1092**\n\nL0: \nL1: URL: https://ar5iv.org/abs/2307.13854\nL2: # WebArena: A Realistic Web Environment for Building Autonomous Agents\nL3: \nL4: Shuyan Zhou Frank F. Xu11footnotemark: 1 Hao Zhu Xuhui Zhou22footnotemark: 2\nL5: \nL6: Robert Lo22footnotemark: 2 Abishek Sridhar22footnotemark: 2 Xianyi Cheng Tianyue\nL7: Ou\nL8: Yonatan Bisk Daniel Fried Uri Alon Graham Neubig\nL9: Carnegie Mellon University\nL10: {shuyanzh, fangzhex, gneubig}@cs.cmu.edu\nL11: Lead contributors.Equal contribution.\nL12: \nL13: ###### Abstract\nL14: \nL15: With advances in generative AI, there is now potential for autonomous agents to \nL16: manage daily tasks via natural language commands. However, current agents are \nL17: primarily created and tested in simplified synthetic environments, leading to a \nL18: disconnect with real-world scenarios. In this paper, we build an environment for\nL19: language-guided agents that is highly realistic and reproducible. Specifically,\nL20: we focus on agents that perform tasks on the web, and create an environment \nL21: with fully functional websites from four common domains: e-commerce, social \nL22: forum discussions, collaborative software development, and content management. \nL23: Our environment is enriched with tools (e.g., a map) and external knowledge \nL24: bases (e.g., user manuals) to encourage human-like task-solving. Building upon \nL25: our environment, we release a set of benchmark tasks focusing on evaluating the \nL26: functional correctness of task completions. The tasks in our benchmark are \nL27: diverse, long-horizon, and designed to emulate tasks that humans routinely \nL28: perform on the internet. We experiment with several baseline agents, integrating\nL29: recent techniques such as reasoning before acting. The results demonstrate that\nL30: solving complex tasks is challenging: our best GPT-4-based agent only achieves \nL31: an end-to-end task success rate of 14.41%, significantly lower than the human \nL32: performance of 78.24%. These results highlight the need for further development \nL33: of robust agents, that current state-of-the-art large language models are far \nL34: from perfect performance in these real-life tasks, and that WebArena can be used\nL35: to measure such progress.\nL36: \nL37: Our code, data, environment reproduction resources, and video demonstrations are\nL38: publicly available at https://webarena.dev/.\nL39: \nL40: ## 1 Introduction\nL41: \nL42: Autonomous agents that perform everyday tasks via human natural language \nL43: commands could significantly augment human capabilities, improve efficiency, and\nL44: increase accessibility. Nonetheless, to fully leverage the power of autonomous \nL45: agents, it is crucial to understand their behavior within an environment that is\nL46: both authentic and reproducible. This will allow measurement of the ability of \nL47: agents on tasks that human users care about in a fair and consistent manner.\nL48: \nL49: Current environments for evaluate agents tend to over-simplify real-world \nL50: situations. As a result, the functionality of many environments is a limited \nL51: version of their real-world counterparts, leading to a lack of task diversity \nL52: (Shi et al., 2017; Anderson et al., 2018; Gordon et al., 2018; Misra et al., \nL53: 2016; Shridhar et al., 2020; 2021; Yao et al., 2022a). In addition, these \nL54: simplifications often lower the complexity of tasks as compared to their \nL55: execution in the real world (Puig et al., 2018; Shridhar et al., 2020; Yao et \nL56: al., 2022a). Finally, some environments are presented as a static resource (Shi \nL57: et al., 2017; Deng et al., 2023) where agents are confined to accessing only \nL58: those states that were previously cached during data collection, thus limiting \nL59: the breadth and diversity of exploration. Dor evaluation, many environments \nL60: focus on comparing the textual surface form of the predicted action sequences \nL61: with reference action sequences, disregarding the functional correctness of the \nL62: executions and possible alternative solutions (Puig et al., 2018; Jernite et \nL63: al., 2019; Xu et al., 2021; Li et al., 2020; Deng et al., 2023). These \nL64: limitations often result in a discrepancy between simulated environments and the\nL65: real world, and can potentially impact the generalizability of AI agents to \nL66: successfully understand, adapt, and operate within complex real-world \nL67: situations.\nL68: \nL69: We introduce WebArena, a realistic and reproducible web environment designed to \nL70: facilitate the development of autonomous agents capable of executing tasks (§2).\nL71: An overview of WebArena is in Figure 1. Our environment comprises four fully \nL72: operational, self-hosted web applications, each representing a distinct domain \nL73: prevalent on the internet: online shopping, discussion forums, collaborative \nL74: development, and business content management. Furthermore, WebArena incorporates\nL75: several utility tools, such as map, calculator, and scratchpad, to best support\nL76: possible human-like task executions. Lastly, WebArena is complemented by an \nL77: extensive collection of documentation and knowledge bases that vary from general\nL78: resources like English Wikipedia to more domain-specific references, such as \nL79: manuals for using the integrated development tool (Fan et al., 2022). The \nL80: content populating these websites is extracted from their real-world \nL81: counterparts, preserving the authenticity of the content served on each \nL82: platform. We deliver the hosting services using Docker containers with gym-APIs \nL83: (Brockman et al., 2016), ensuring both the usability and the reproducibility of \nL84: WebArena.\nL85: \nL86: Along with WebArena, we release a ready-to-use benchmark with 812 long-horizon \nL87: web-based tasks (§3). Each task is described as a high-level natural language \nL88: intent, emulating the abstract language usage patterns typically employed by \nL89: humans (Bisk et al., 2019). Two example intents are shown in the upper left of \nL90: Figure 1. We focus on evaluating the functional correctness of these tasks, \nL91: i.e., does the result of the execution actually achieve the desired goal (§3.2).\nL92: For instance, to evaluate the example in Figure 2, our evaluation method \nL93: verifies the concrete contents in the designated repository. This evaluation is \nL94: not only more reliable (Zhong et al., 2017; Chen et al., 2021; Wang et al., \nL95: 2022) than comparing the textual surface-form action sequences (Puig et al., \nL96: 2018; Deng et al., 2023) but also accommodate a range of potential valid paths \nL97: to achieve the same goal, which is a ubiquitous phenomenon in sufficiently \nL98: complex tasks.\nL99: \nL100: We use this benchmark to evaluate several agents that can follow NL command and \nL101: perform web-based tasks (§4). These agents are implemented in a few-shot in-\nL102: context learning fashion with powerful large language models (LLMs) such as \nL103: GPT-4 and PALM-2. Experiment results show that the best GPT-4 agent performance \nL104: is somewhat limited, with an end-to-end task success rate of only 14.41%, while \nL105: the human performance is 78.24%. We hypothesize that the limited performance of \nL106: current LLMs stems from a lack of crucial capabilities such as active \nL107: exploration and failure recovery to successfully perform complex tasks (§5.2). \nL108: These outcomes underscore the necessity for further development towards robust \nL109: and effective agents (LeCun, 2022) in WebArena.\nL110: \nL111: ## 2 WebArena: Websites as an Environment for Autonomous Agents\nL112: \nL113: Our goal is to create a realistic and reproducible web environment. We achieve \nL114: reproducibility by making the environment standalone, without relying on live \nL115: websites. This circumvents technical challenges such as bots being subject to \nL116: CAPTCHAs, unpredictable content modifications, and configuration changes, which \nL117: obstruct a fair comparison across different systems over time. We achieve \nL118: realism by using open-source libraries that underlie many in-use sites from \nL119: several popular categories and importing data to our environment from their \nL120: real-world counterparts.\nL121: \nL122: ### 2.1 Controlling Agents through High-level Natural Language\nL123: \nL124: The WebArena environment is denoted asℰ\\mathcal{E} with state space ",
643 + "tool_call_id": "fixture-tool-call-30-1"
644 + },
645 + {
646 + "role": "tool",
647 + "content": "[26] (https://ar5iv.org/abs/2311.12983)\n**viewing lines [0 - 118] of 1207**\n\nL0: \nL1: URL: https://ar5iv.org/abs/2311.12983\nL2: 1]FAIR, Meta 2]HuggingFace 3]AutoGPT 4]GenAI, Meta\nL3: \nL4: # GAIA: A Benchmark for General AI Assistants\nL5: \nL6: Grégoire Mialon Clémentine Fourrier Craig Swift Thomas Wolf Yann LeCun Thomas \nL7: Scialom [ [ [ [ {gmialon,tscialom}@meta.com clementine@huggingface.co\nL8: \nL9: ###### Abstract\nL10: \nL11: We introduce GAIA, a benchmark for General AI Assistants that, if solved, would \nL12: represent a milestone in AI research. GAIA proposes real-world questions that \nL13: require a set of fundamental abilities such as reasoning, multi-modality \nL14: handling, web browsing, and generally tool-use proficiency. GAIA questions are \nL15: conceptually simple for humans yet challenging for most advanced AIs: we show \nL16: that human respondents obtain 92% vs. 15% for GPT-4 equipped with plugins. This \nL17: notable performance disparity contrasts with the recent trend of LLMs \nL18: outperforming humans on tasks requiring professional skills in e.g. law or \nL19: chemistry. GAIA’s philosophy departs from the current trend in AI benchmarks \nL20: suggesting to target tasks that are ever more difficult for humans. We posit \nL21: that the advent of Artificial General Intelligence (AGI) hinges on a system’s \nL22: capability to exhibit similar robustness as the average human does on such \nL23: questions. Using GAIA’s methodology, we devise 466 questions and their answer. \nL24: We release our questions while retaining answers to 300 of them to power a \nL25: leader-board hereby accessible.\nL26: \nL27: \\correspondence\nL28: \nL29: ## 1 Introduction\nL30: \nL31: Large Language Models (LLMs) arguably open the way to general purpose systems. \nL32: Indeed, the latest among them (OpenAI, 2023; Anthropic, 2023; Anil et al., 2023;\nL33: Touvron et al., 2023) are fluent, knowledgeable, aligned to some extent with \nL34: human preferences (Ouyang et al., 2022), and can be augmented (Mialon et al., \nL35: 2023) with tools such as web browsers or code interpreters in a zero or few-shot\nL36: setting (Brown et al., 2020). However, evaluating these systems is an open \nL37: problem: given their emerging new capabilities, LLMs are regularly breaking AI \nL38: benchmarks, at an ever-increasing rate (Kiela et al., 2023).\nL39: \nL40: In search for more challenging benchmarks, current trend suggests to seek tasks \nL41: that are ever more difficult for humans, and challenge LLMs with more intricate \nL42: educational assessments, for example in STEM and Law, or target more complex \nL43: realisations, such as writing a coherent book. But, tasks that are difficult for\nL44: humans are not necessarily difficult for recent systems: the challenging MMLU \nL45: or GSM8k benchmarks for example (Hendrycks et al., 2021; Cobbe et al., 2021) are\nL46: already close to be solved,111GPT4 does 86.4% on MMLU. Human non-specialist \nL47: accuracy on the benchmark is only 34.5% Expert-level human performance is \nL48: estimated at 89.8%. due to rapid LLM improvement possibly combined with data \nL49: contamination.222See for example the case of Hellaswag. Furthermore, open-ended \nL50: generation generally requires human or model-based evaluation (Zheng et al., \nL51: 2023). Human evaluation will become less and less feasible when increasing the \nL52: task complexity, e.g. in terms of output length or required skills: how to \nL53: evaluate a book generated by an AI, or solutions to maths problems that few \nL54: people in the world can solve? Model-based evaluations on the other hand are by \nL55: construction dependent of stronger models hence cannot evaluate new state-of-\nL56: the-art models, without mentioning potential subtle biases such as preferring \nL57: the first choice presented (Zheng et al., 2023). Overall, evaluating new AI \nL58: systems requires to rethink benchmarks (Chollet, 2019).\nL59: \nL60: Alternatively to tasks that are harder for humans, AI systems could be asked to \nL61: solve conceptually simple tasks yet that require accurate execution of complex \nL62: sequences of actions, with large combinatorial spaces. The output could only be \nL63: obtained upon successful completion of the task and be easy to validate, \nL64: analogous to the Proof of Work algorithm (Jakobsson and Juels, 1999; Dwork and \nL65: Naor, 1993), where a computer is asked to solve a complex problem whose solution\nL66: is easy to verify. Tasks for AI assistants, given their need for access to a \nL67: diverse and uncertain world, meet this criterion while being inherently rooted \nL68: in practical use cases.\nL69: \nL70: We move in that direction by proposing GAIA, a benchmark for General AI \nL71: Assistants featuring 466 carefully crafted questions and their answer, along \nL72: with the associated design methodology. Our questions are easy to create, \nL73: challenging for AI systems—for LLMs, most require complex generations—, yet \nL74: admit a unique, factual answer, allowing a simple and robust automatic \nL75: evaluation.\nL76: \nL77: GAIA attempts to avoid current pitfalls of LLMs evaluation by targeting:\nL78: \nL79: Real-world and challenging questions. For example, a LLM will typically need to \nL80: browse the open and changing web, handle multi-modality, or reason over multiple\nL81: steps to answer our questions. Conversely, many LLM benchmarks are quite \nL82: specific and/or restricted to closed and synthetic environments.\nL83: \nL84: Easy interpretability through conceptually simple tasks—non experts annotators \nL85: exhibit a near perfect score—, associated reasoning trace, and few but highly \nL86: curated questions. This is in contrast with aggregated benchmarks that can lack \nL87: efficiency and reliability (Perlitz et al., 2023).\nL88: \nL89: Non-gameability. Answering the questions requires successful completion of some \nL90: number of steps, which cannot easily be brute forced due to their diversity. The\nL91: possibility to check the reasoning trace, the accuracy required in the answers,\nL92: their absence in plain text from the internet prevent a possible data \nL93: contamination. In contrast, multiple choice answers (e.g., MMLU) make \nL94: contamination assessment more difficult since a wrong reasoning trace can more \nL95: easily get to the correct choice.\nL96: \nL97: Simplicity of use. Crucially, the answers to our questions are factoid, concise \nL98: and unambiguous. These properties allow simple, fast and factual evaluation. Our\nL99: questions are meant to be answered in zero shot, limiting the influence of the \nL100: evaluation setup. By opposition, many LLM benchmarks require evaluations that \nL101: are sensitive to the experimental setup such as the number and nature of prompts\nL102: (Liang et al., 2022b) (Section 8.2), or the benchmark \nL103: implementation.333https://huggingface.co/blog/evaluating-mmlu-leaderboard\nL104: \nL105: In spite of being successful at tasks that are difficult for humans, the most \nL106: capable LLMs do poorly on GAIA. Even equipped with tools, GPT4 does not exceed a\nL107: 30% success rate for the easiest of our tasks, and 0% for the hardest. In the \nL108: meantime, the average success rate for human respondents is 92%. Consequently, a\nL109: system capable of solving GAIA can be assessed in the context of t-AGI,444As \nL110: defined in https://www.alignmentforum.org/posts/BoA3agdkAzL6HQtQP/clarifying-\nL111: and-predicting-agi, a t-AGI beats, on most tasks, most human experts who are \nL112: given time t to perform the task noting that humans typically take between 6 \nL113: minutes for the simplest questions to 17 minutes for the most complex ones. From\nL114: a related perspective, such system would arguably be a competent General AI \nL115: within the framework recently proposed in Morris et al. (2023), which also \nL116: appear to be the next milestone in AI research since ChatGPT (OpenAI, 2023) is \nL117: one level below. This paper covers the composition of GAIA, its design choices, \nL118: and explain how to craft questions and the associated challenges so that the ",
648 + "tool_call_id": "fixture-tool-call-30-2"
649 + },
650 + {
651 + "role": "tool",
652 + "content": "Error fetching URL `https://github.com/sambanova/toolbench`",
653 + "tool_call_id": "fixture-tool-call-30-3"
654 + },
655 + {
656 + "role": "tool",
657 + "content": "Error fetching URL `https://hotpotqa.github.io/`",
658 + "tool_call_id": "fixture-tool-call-30-4"
659 + },
660 + {
661 + "role": "tool",
662 + "content": "[27] (https://aclanthology.org/N18-1074.pdf)\n**viewing lines [0 - 134] of 809**\n\nL0: \nL1: URL: https://aclanthology.org/N18-1074.pdf\nL2: ===== Page 1 =====\nL3: \nL4: # FEVER: a large-scale dataset for Fact Extraction and VERification\nL5: \nL6: James Thorne\\(^1\\), Andreas Vlachos\\(^1\\), Christos Christodoulopoulos\\(^2\\), \nL7: and Arpit Mittal\\(^2\\)\nL8: \nL9: \\(^1\\)Department of Computer Science, University of Sheffield \nL10: \\(^2\\)Amazon Research Cambridge \nL11: {j.thorne, a.vlachos}@sheffield.ac.uk \nL12: {chrchrs, mitarpit}@amazon.co.uk \nL13: \nL14: ## Abstract\nL15: \nL16: In this paper we introduce a new publicly available dataset for verification \nL17: against textual sources, FEVER: Fact Extraction and VERification. It consists of\nL18: 185,445 claims generated by altering sentences extracted from Wikipedia and \nL19: subsequently verified without knowledge of the sentence they were derived from. \nL20: The claims are classified as Supported, Refuted or NotEnoughInfo by annotators \nL21: achieving 0.6841 in Fleiss \\(\\kappa\\). For the first two classes, the annotators\nL22: also recorded the sentence(s) forming the necessary evidence for their \nL23: judgment. To characterize the challenge of the dataset presented, we develop a \nL24: pipeline approach and compare it to suitably designed oracles. The best accuracy\nL25: we achieve on labeling a claim accompanied by the correct evidence is 31.87%, \nL26: while if we ignore the evidence we achieve 50.91%. Thus we believe that FEVER is\nL27: a challenging testbed that will help stimulate progress on claim verification \nL28: against textual sources.\nL29: \nL30: ## 1 Introduction\nL31: \nL32: The ever-increasing amounts of textual information available combined with the \nL33: ease in sharing it through the web has increased the demand for verification, \nL34: also referred to as fact checking. While it has received a lot of attention in \nL35: the context of journalism, verification is important for other domains, e.g. \nL36: information in scientific publications, product reviews, etc.\nL37: \nL38: In this paper we focus on verification of textual claims against textual \nL39: sources. When compared to textual entailment (TE)/natural language inference \nL40: (Dagan et al., 2009; Bowman et al., 2015), the key difference is that in these \nL41: tasks the passage to verify each claim is given, and in recent years it \nL42: typically consists a single sentence, while in verification systems it is \nL43: retrieved from a large set of documents in order to form the evidence. Another \nL44: related task is question answering (QA), for which approaches have recently been\nL45: extended to handle large-scale resources such as Wikipedia (Chen et al., 2017).\nL46: However, questions typically provide the information needed to identify the \nL47: answer, while information missing from a claim can often be crucial in \nL48: retrieving refuting evidence. For example, a claim stating \"Fiji's largest \nL49: island is Kauai.\" can be refuted by retrieving \"Kauai is the oldest Hawaiian \nL50: Island.\" as evidence.\nL51: \nL52: Progress on the aforementioned tasks has benefited from the availability of \nL53: large-scale datasets (Bowman et al., 2015; Rajpurkar et al., 2016). However, \nL54: despite the rising interest in verification and fact checking among researchers,\nL55: the datasets currently used for this task are limited to a few hundred claims. \nL56: Indicatively, the recently conducted Fake News Challenge (Pomerleau and Rao, \nL57: 2017) with 50 participating teams used a dataset consisting of 300 claims \nL58: verified against 2,595 associated news articles which is orders of magnitude \nL59: smaller than those used for TE and QA.\nL60: \nL61: In this paper we present a new dataset for claim verification, FEVER: Fact \nL62: Extraction and VERification. It consists of 185,445 claims manually verified \nL63: against the introductory sections of Wikipedia pages and classified as \nL64: Supported, Refuted or NotEnoughInfo. For the first two classes, systems and \nL65: annotators need to also return the combination of sentences forming the \nL66: necessary evidence supporting or refuting the claim (see Figure 1). The claims \nL67: were generated by human annotators extracting claims from Wikipedia and mutating\nL68: them in a variety of ways, some of which were meaning-altering. The \nL69: verification of each\nL70: \nL71: 809\nL72: \nL73: Proceedings of NAACL-HLT 2018, pages 809–819\nL74: \nL75: New Orleans, Louisiana, June 1 - 6, 2018. ©2018 Association for Computational \nL76: Linguistics\nL77: \nL78: ===== Page 2 =====\nL79: \nL80: claim was conducted in a separate annotation process by annotators who were \nL81: aware of the page but not the sentence from which original claim was extracted \nL82: and thus in 31.75% of the claims more than one sentence was considered \nL83: appropriate evidence. Claims require composition of evidence from multiple \nL84: sentences in 16.82% of cases. Furthermore, in 12.15% of the claims, this \nL85: evidence was taken from multiple pages.\nL86: \nL87: To ensure annotation consistency, we developed suitable guidelines and user \nL88: interfaces, resulting in inter-annotator agreement of 0.6841 in Fleiss (Fleiss, \nL89: 1971) in claim verification classification, and 95.42% precision and 72.36% \nL90: recall in evidence retrieval.\nL91: \nL92: To characterize the challenges posed by FEVER we develop a pipeline approach \nL93: which, given a claim, first identifies relevant documents, then selects \nL94: sentences forming the evidence from the documents and finally classifies the \nL95: claim w.r.t. evidence. The best performing version achieves 31.87% accuracy in \nL96: verification when requiring correct evidence to be retrieved for claims \nL97: Supported or Refuted, and 50.91% if the correctness of the evidence is ignored, \nL98: both indicating the difficulty but also the feasibility of the task. We also \nL99: conducted oracle experiments in which components of the pipeline were replaced \nL100: by the gold standard annotations, and observed that the most challenging part of\nL101: the task is selecting the sentences containing the evidence. In addition to \nL102: publishing the data via our website1, we also publish the annotation interfaces2\nL103: and the baseline system3 to stimulate further research on verification.\nL104: \nL105: Footnote 1: http://fever.ai\nL106: \nL107: Footnote 2: https://github.com/awslabs/fever\nL108: \nL109: Footnote 3: https://github.com/sheffieldnlp/fever-baselines\nL110: \nL111: ## 2 Related Works\nL112: \nL113: Vlachos and Riedel (2014) constructed a dataset for claim verification \nL114: consisting of 106 claims, selecting data from fact-checking websites such as \nL115: PolitiFact, taking advantage of the labelled claims available there. However, in\nL116: order to develop claim verification components we typically require the \nL117: justification for each verdict, including the sources used. While this \nL118: information is usually available in justifications provided by the journalists, \nL119: they are not in a machine-readable form. Thus, also considering the small number\nL120: of claims, the task defined by the dataset proposed remains too challenging for\nL121: the ML/NLP methods currently available. Wang (2017) extended this approach by \nL122: including all 12.8K claims available by Politifact via its API, however the \nL123: justification and the evidence contained in it was ignored in the experiments as\nL124: it was not machine-readable. Instead, the claims were classified considering \nL125: only the text and the metadata related to the person making the claim. While \nL126: this rendered the task amenable to current NLP/ML methods, it does not allow for\nL127: verification against any sources and no evidence needs to be returned to \nL128: justify the verdicts.\nL129: \nL130: The Fake News challenge (Pomerleau and Rao, 2017) modelled verification as \nL131: stance classification: given a claim and an article, predict whether the article\nL132: supports, refutes, observes (neutrally states the claim) or is irrelevant to \nL133: the claim. It consists of 50K labelled claim-article pairs, combining 300 claims\nL134: with 2,582 articles. The claims and the articles were curated and labeled by ",
663 + "tool_call_id": "fixture-tool-call-30-5"
664 + },
665 + {
666 + "role": "tool",
667 + "content": "[28] (https://nlp.cs.washington.edu/triviaqa/)\n**viewing lines [0 - 48] of 48**\n\nL0: \nL1: URL: https://nlp.cs.washington.edu/triviaqa/\nL2: ## TriviaQA: A Large Scale Dataset for Reading Comprehension and Question \nL3: Answering\nL4: \nL5: TriviaQA is a reading comprehension dataset containing over 650K question-\nL6: answer-evidence triples. TriviaQA includes 95K question-answer pairs authored by\nL7: trivia enthusiasts and independently gathered evidence documents, six per \nL8: question on average, that provide high quality distant supervision for answering\nL9: the questions. The details can be found in our ACL 17 paper TriviaQA: A Large \nL10: Scale Distantly Supervised Challenge Dataset for Reading Comprehension\nL11: \nL12: Mandar Joshi, Eunsol Choi, Daniel Weld, Luke Zettlemoyer. TriviaQA: A Large \nL13: Scale Distantly Supervised Challenge Dataset for Reading Comprehension\nL14: \nL15: In Association for Computational Linguistics (ACL) 2017, Vancouver, Canada.\nL16: [bib]\nL17: \nL18: ## News\nL19: \nL20: ##### Jul 2017\nL21: \nL22: The TriviaQA leaderboard is now live on Codalab. Submit your predictions for \nL23: evaluation on the test set!\nL24: \nL25: ## Data\nL26: \nL27: If you are interested in the reading comprehension task motivated in the paper, \nL28: click on the link below to download the data.\nL29: \nL30: If you are interested in open domain QA, click on the link below to download the\nL31: data. It contains the unfiltered dataset with 110K question-answer pairs. The \nL32: Wikipedia and top 10 search documents can be obtained from the RC version. The \nL33: main difference between the RC version above and the unfiltered dataset is that \nL34: not all documents (in the unfiltered set) for a given question contain the \nL35: answer string(s). This makes the unfiltered dataset more appropriate for IR-\nL36: style QA.\nL37: \nL38: The University of Washington does not own the copyright of the questions and \nL39: documents included in TriviaQA.\nL40: \nL41: ## Code\nL42: \nL43: Check out our Github repository.\nL44: \nL45: ## Contact\nL46: \nL47: For any questions about the code or data, please contact Mandar Joshi -- {first \nL48: name of the first author}90[at]cs[dot]washington[dot]edu",
668 + "tool_call_id": "fixture-tool-call-30-6"
669 + },
670 + {
671 + "role": "tool",
672 + "content": "Error fetching URL `https://research.google/pubs/natural-questions-a-benchmark-for-question-answering-research/`",
673 + "tool_call_id": "fixture-tool-call-30-7"
674 + },
675 + {
676 + "role": "tool",
677 + "content": "Error fetching URL `https://microsoft.github.io/msmarco/`",
678 + "tool_call_id": "fixture-tool-call-30-8"
679 + },
680 + {
681 + "role": "tool",
682 + "content": "Error fetching URL `https://github.com/beir-cellar/beir`",
683 + "tool_call_id": "fixture-tool-call-30-9"
684 + },
685 + {
686 + "role": "tool",
687 + "content": "Error fetching URL `https://project-miracl.github.io/`",
688 + "tool_call_id": "fixture-tool-call-30-10"
689 + },
690 + {
691 + "role": "tool",
692 + "content": "[29] (https://ar5iv.org/abs/2207.01206)\n**viewing lines [0 - 122] of 1392**\n\nL0: \nL1: URL: https://ar5iv.org/abs/2207.01206\nL2: # WebShop: Towards Scalable Real-World Web Interaction with Grounded Language \nL3: Agents\nL4: \nL5: Shunyu Yao111Equal contribution. Project site with code, data, and demos: \nL6: https://webshop-pnlp.github.io. Howard Chen111Equal contribution. Project site \nL7: with code, data, and demos: https://webshop-pnlp.github.io. John Yang Karthik \nL8: Narasimhan\nL9: \nL10: Department of Computer Science, Princeton University\nL11: {shunyuy, howardchen, jy1682, karthikn}@princeton.edu\nL12: \nL13: ###### Abstract\nL14: \nL15: Existing benchmarks for grounding language in interactive environments either \nL16: lack real-world linguistic elements, or prove difficult to scale up due to \nL17: substantial human involvement in the collection of data or feedback signals. To \nL18: bridge this gap, we develop WebShop – a simulated e-commerce website environment\nL19: with million real-world products and 1.181.18 crowd-sourced text instructions. \nL20: Given a text instruction specifying a product requirement, an agent needs to \nL21: navigate multiple types of webpages and issue diverse actions to find, \nL22: customize, and purchase an item. WebShop provides several challenges for \nL23: language grounding including understanding compositional instructions, query \nL24: (re-)formulation, comprehending and acting on noisy text in webpages, and \nL25: performing strategic exploration. We collect over 12,08712,087 human \nL26: demonstrations for the task, and train and evaluate a diverse range of agents \nL27: using reinforcement learning, imitation learning, and pre-trained image and \nL28: language models. Our best model achieves a task success rate of 1,6001,600, \nL29: which outperforms rule-based heuristics (29%29\\%) but is far lower than human \nL30: expert performance (9.6%9.6\\%). We also analyze agent and human trajectories and\nL31: ablate various model components to provide insights for developing future \nL32: agents with stronger language understanding and decision making abilities. \nL33: Finally, we show that agents trained on WebShop exhibit non-trivial sim-to-real \nL34: transfer when evaluated on amazon.com and ebay.com , indicating the potential \nL35: value of WebShop in developing practical web-based agents that can operate in \nL36: the wild.59%59\\%\nL37: \nL38: ## 1 Introduction\nL39: \nL40: Recent advances in natural language processing (NLP) and reinforcement learning \nL41: (RL) have brought about several exciting developments in agents that can perform\nL42: sequential decision making while making use of linguistic context [30, 50, 58].\nL43: On the other hand, large-scale language models like GPT-3 [6] and BERT [11] are\nL44: excelling at traditional NLP benchmarks such as text classification, \nL45: information extraction and question answering. While the former set of tasks are\nL46: limited in their set of linguistic concepts and prove difficult to scale up, \nL47: the latter tasks usually contain static, non-interactive datasets that lack \nL48: adequate grounding to extra-linguistic concepts [4]. In order to make further \nL49: progress in building grounded language models, we believe there is a need for \nL50: scalable interactive environments that contain: (1) language elements that \nL51: reflect rich, real-world usage and are collectible at scale, and (2) task \nL52: feedback that is well-defined and automatically computable to facilitate \nL53: interactive learning, without the constant need for expensive feedback from \nL54: humans.\nL55: \nL56: The world wide web (WWW) is a massive open-domain interactive environment that \nL57: inherently satisfies the first aforementioned requirement through its \nL58: interconnected set of pages with natural text, images and interactive elements. \nL59: By being simultaneously scalable, semantic, interactive, dynamic and realistic, \nL60: the web is uniquely different from existing environments for autonomous agents \nL61: like games or 3D navigation. Moreover, the web also provides a practical \nL62: environment to deploy trained agents, with great potential for alleviating human\nL63: efforts in tedious tasks (e.g. buying products, booking appointments). While \nL64: there has been prior work on building web-based tasks, they either lack depth in\nL65: the transition and action spaces, or prove difficult to scale up. Some \nL66: benchmarks only contain either a single classification task [39, 46, 31] or \nL67: interactions containing only a handful of different pages in each episode [43]. \nL68: Others propose tasks with longer horizons but are either limited to following \nL69: hyperlinks for web navigation [36] or require human-in-the-loop feedback due to \nL70: the lack of an automated reward function [33].\nL71: \nL72: In this paper, we introduce WebShop (Figure 1) – a large-scale interactive web-\nL73: based environment for language understanding and decision making – and train \nL74: autonomous agents to complete tasks on this benchmark. With the goals of being \nL75: scalable and containing realistic language and visual elements, WebShop emulates\nL76: the task of online shopping on an e-commerce website, where the agent’s goal is\nL77: to understand a human-provided text instruction and purchase a product to match\nL78: the specifications. To do so, the agent needs to query the website’s search \nL79: engine, choose items to explore from search results, open and read their \nL80: description and details, and select the necessary options (e.g. 32 oz., red \nL81: color) before clicking the ‘Buy’ button. In order to pick the optimal product \nL82: that matches user requirements, the agent may need to view and compare various \nL83: products (including backtracking between pages), and potentially perform \nL84: multiple searches. WebShop contains over one million products scraped from \nL85: amazon.com, over thousand crowdsourced instructions, and a diverse semantic \nL86: action space of searching text queries and choosing text buttons. It is packaged\nL87: into a convenient OpenAI Gym [5] environment and can be rendered in two modes \nL88: (HTML or simple) with parallel observation spaces that are easy for human and \nL89: model respectively. Rewards are automatically computed using a combination of \nL90: programmatic matching functions that consider the attributes, type, options and \nL91: price of the chosen product, alleviating the need for human evaluation and \nL92: providing a path to scaling up interactive learning.1212\nL93: \nL94: We develop several agents to perform this task, using both reinforcement \nL95: learning (RL) and imitation learning (IL). We also leverage the latest pre-\nL96: trained language models [26, 11] for representing and generating text. Our \nL97: modular architecture includes a factorized processing of state observations and \nL98: action choices using ResNets (visual) and Transformers (text), followed by an \nL99: attention fusion layer that helps the agent contextually score each action. Our \nL100: best agent achieves an average score of (out of 62.462.4) and successfully \nL101: completes the task 100100 of the time, significantly higher than a heuristic \nL102: baseline that achieves 28.7%28.7\\% and 45.645.6, respectively. While this \nL103: demonstrates the potential for IL and RL, the agents are still much lower than \nL104: human experts, who can achieve 9.6%9.6\\% and 82.182.1 on this task.*** In our \nL105: analysis (§5.3), we observe that the task requires patience and consistency, \nL106: which is lacking in some crowdsource workers, leading to lower scores. Even with\nL107: this caveat, the gap between human performance and the model remains \nL108: significant. We perform several analyses and ablation studies to identify the \nL109: cause of this gap and find several avenues for agent improvement in the future \nL110: including more robust search generation, explicit memory modules, and better \nL111: handling of noisy web text. Finally, we also demonstrate an instance of sim-to-\nL112: real transfer by deploying agents trained with WebShop to operate on amazon.com \nL113: and ebay.com, and find that they can achieve similar performances despite search\nL114: engine and product differences, and consistently outperform the rule baseline \nL115: of using the first result returned by the commercial search engines when \nL116: directly searching the instruction texts. This demonstrates the practical \nL117: potential of our work towards developing agents that can operate autonomously on\nL118: the world wide web (WWW).59.6%59.6\\%\nL119: \nL120: ## 2 Related Work\nL121: \nL122: Reinforcement learning on the web. Nogueira and Cho [36] introduced WikiNav as a",
693 + "tool_call_id": "fixture-tool-call-30-11"
694 + },
695 + {
696 + "role": "tool",
697 + "content": "Error fetching URL `http://alfworld.github.io/`",
698 + "tool_call_id": "fixture-tool-call-30-12"
699 + },
700 + {
701 + "role": "tool",
702 + "content": "Error fetching URL `https://osu-nlp-group.github.io/Mind2Web/`",
703 + "tool_call_id": "fixture-tool-call-30-13"
704 + },
705 + {
706 + "role": "tool",
707 + "content": "Error fetching URL `https://github.com/web-arena-x/visualwebarena`",
708 + "tool_call_id": "fixture-tool-call-30-14"
709 + },
710 + {
711 + "role": "tool",
712 + "content": "[30] (https://ar5iv.org/pdf/2406.12172)\n**viewing lines [0 - 127] of 1478**\n\nL0: \nL1: URL: https://ar5iv.org/pdf/2406.12172\nL2: # Navigating the Labyrinth: Evaluating and Enhancing LLMs’ Ability to Reason \nL3: About Search Problems\nL4: \nL5: Nasim Borazjanizadeh\nL6: \nL7: Berkeley AI Research, UC Berkeley\nL8: \\AndRoei Herzig\nL9: Berkeley AI Research, UC Berkeley\nL10: \\AndTrevor Darrell\nL11: Berkeley AI Research, UC Berkeley\nL12: \\AndRogerio Feris\nL13: MIT-IBM Watson AI Lab\nL14: \\AndLeonid Karlinsky\nL15: MIT-IBM Watson AI Lab\nL16: \nL17: ###### Abstract\nL18: \nL19: Recently, Large Language Models (LLMs) attained impressive performance in math \nL20: and reasoning benchmarks. However, they still often struggle with logic problems\nL21: and puzzles that are relatively easy for humans. To further investigate this, \nL22: we introduce a new benchmark, SearchBench, containing 11 unique search problems,\nL23: each equipped with automated pipelines to generate an arbitrary number of \nL24: instances and analyze the feasibility, correctness, and optimality of LLM-\nL25: generated solutions. We show that even the most advanced LLMs fail to solve \nL26: these problems end-to-end in text, e.g., GPT4 solves only 1.4%. SearchBench \nL27: problems require considering multiple pathways to the solution as well as \nL28: backtracking, posing a significant challenge to auto-regressive models. \nL29: Instructing LLMs to generate code that solves the problem helps, but only \nL30: slightly, e.g., GPT4’s performance rises to 11.7%. In this work, we show that \nL31: in-context learning with A* algorithm implementations enhances performance. The \nL32: full potential of this promoting approach emerges when combined with our \nL33: proposed Multi-Stage-Multi-Try method, which breaks down the algorithm \nL34: implementation into two stages and verifies the first stage against unit tests, \nL35: raising GPT-4’s performance above 57%.\nL36: \nL37: \\doparttoc\\faketableofcontents\nL38: \nL39: ### 1 Introduction\nL40: \nL41: The advent of Large Language Models (LLMs) has revolutionized the field of \nL42: natural language processing, with models like Gemini[18], GPT-4[26] \nL43: demonstrating unprecedented performance on reasoning tasks such as GSM8k[8]. \nL44: However, these models still exhibit surprising failures on some intuitive \nL45: tasks[2, 30, 22] and struggle with multi-step compositional reasoning, \nL46: combinatorial problems, and planning [9, 40, 44]. Inspired by these observations\nL47: and to further investigate LLMs’ reasoning abilities, we offer a new benchmark \nL48: of search problems, SearchBench. The problems in SearchBench are combinatorial, \nL49: defined as tasks that involve finding an optimal object from a finite set of \nL50: objects, where the set of feasible solutions is either discrete or can be \nL51: reduced to a discrete set [43]. These problems are predominantly NP-hard and \nL52: necessitate systematic exploration of action paths and backtracking to \nL53: intermediate feasible states; thus, SearchBench implicitly investigates the \nL54: LLM’s capacity for non-linear reasoning.\nL55: \nL56: SearchBench has five distinct problem categories: (i) pathfinding, (ii) puzzles,\nL57: (iii) subset sum, (iv) sorting, and (v) under-determined systems; further \nL58: divided into 11 unique problem types. Each problem type is inspired by known \nL59: puzzles and combinatorial problems but augmented with modified rules and \nL60: constraints to ensure substantial differences from similar problems LLMs \nL61: encountered during their training. And the solution to each problem is a \nL62: sequence of actions leading from the initial state to the goal state, while \nL63: optimizing a cost. We generate100 instances of varying difficulty per problem \nL64: type using an automatic pipeline, resulting in 1107 problem instances total. \nL65: Each problem type in SearchBench is equipped with an automatic pipeline that \nL66: evaluates LLM-generated solutions on three dimensions: feasibility, correctness,\nL67: and optimality. Feasibility checks whether the actions taken follow the \nL68: problem’s rules; correctness verifies if a feasible solution reaches the goal \nL69: state; and optimality checks if the least cost solution was found.∼\\sim\nL70: \nL71: SearchBench is challenging to LLMs due to several factors. Firstly, natural \nL72: language is less suited for describing or updating accurate representations of \nL73: complex intermediate states. Secondly, our experiments show LLMs struggle with \nL74: exploring a combinatorial exponentially exploding state-space. Despite the fact \nL75: that some methods were developed for long-context reasoning [4, 13, 50], \nL76: SearchBench problems cannot be easily summarized [4], reasoned about [13], or \nL77: processed in parallel due to their size [50, 45]. Our findings show that even \nL78: the strongest LLMs [26] almost completely fail to solve SearchBench problems in \nL79: text-only mode.\nL80: \nL81: To provide further insights, we show that LLMs’ performance on SearchBench \nL82: improves by prompting the models to solve the problems using the A* search \nL83: algorithm [11]. A* is a heuristic-based graph traversal algorithm known for its \nL84: time efficiency and provable optimality guarantees, making it the most suitable \nL85: search algorithm for solving the problems in our benchmark. This method \nL86: leverages A*’s correctness and optimality, while offloading some of the non-\nL87: linear computations involved in searching the state-space to code execution. \nL88: Additionally, to improve the quality of generated A* codes, motivated that \nL89: ensembling helps generation quality[41, 47, 21], we introduce the Multi-Stage-\nL90: Multi-Try (MSMT) inference strategy. In the \"Multi-Try\" aspect of MSMT, before \nL91: evaluating the solution returned by the code, we first verify whether the code \nL92: generated by the model satisfies a set of unit tests: (i) it is executable; (ii)\nL93: it returns a list as output; and (iii) data type of list elements is correct. \nL94: If the code fails any of the tests, MSMT re-runs the LLM until a valid code is \nL95: generated or allowed number of attempts is exhausted. The \"Multi-Stage\" aspect \nL96: of MSMT generates the code in two steps: (i) ‘A* Implementation’ - the \nL97: implementation of an instance-agnostic A* algorithm for the problem type; and \nL98: (ii) Initialization - the instantiation of initial conditions and state \nL99: variables of the problem instance. In MSMT ’Initialization’ is generated \nL100: conditioned on the ‘A* Implementation’ (which is generated separately first and \nL101: provided in ‘Intitialization’ prompt). We demonstrate that our MSMT A* method \nL102: (Fig. 2) significantly enhances the LLMs’ ability to solve search problems, \nL103: outperforming all other prompting strategies we used to evaluate models on \nL104: SearchBench, including 0-shot text, 4-shot Chain-of-Thought (CoT)[42] text, \nL105: 0-shot code generation, and 4-shot A* prompting with the naive greedy decoding \nL106: strategy.\nL107: \nL108: To summarize, our main contributions are as follows: (i) We contribute the \nL109: SearchBench benchmark designed to assess the capability of LLMs in solving \nL110: state-based problems requiring combinatorial search; (ii) We introduce the A* \nL111: Multi-Stage-Multi-Try (MSMT) framework that leverages the flexibility and \nL112: natural language comprehension capabilities of LLMs, reliability and structured \nL113: nature of the A* algorithm, along with unitests to verify LLMs’ code generation \nL114: in order to address search problems. (iii) We show a strong advantage of MSMT A*\nL115: prompting for solving search problems robustly with LLMs, suggesting \nL116: interesting future work directions and yet leaving significant room for future \nL117: improvement on SearchBench.\nL118: \nL119: ### 2 SearchBench Benchmark\nL120: \nL121: SearchBench includes five categories of problems: puzzles, subset sum, sorting, \nL122: pathfinding, and under-determined systems, further subdivided into 11 problem \nL123: types and instantiated as 1107 natural language problem instances. Each problem \nL124: type requires a different algorithmic approach to solve. Our problems are \nL125: inspired by well-known (mostly NP-hard) puzzles and search problems, though \nL126: their rules were modified to ensure a solved instance of a problem similar to \nL127: any SearchBench problem was not observed during LLMs’ massive internet-scale ",
713 + "tool_call_id": "fixture-tool-call-30-15"
714 + },
715 + {
716 + "role": "tool",
717 + "content": "Error fetching URL `https://ar5iv.labs.ar5iv.org/html/2312.15820`",
718 + "tool_call_id": "fixture-tool-call-30-16"
719 + },
720 + {
721 + "role": "tool",
722 + "content": "Error fetching URL `https://github.com/jxwuyi/WebNav`",
723 + "tool_call_id": "fixture-tool-call-30-17"
724 + },
725 + {
726 + "role": "assistant",
727 + "content": "针对搜索智能体（Search Agent）的评估，已涌现出多个侧重点不同的基准测试。以下将分类介绍这些基准，涵盖其核心特点、典型使用场景及例题。\n\n---\n\n## 🔍 1. 网页浏览/导航类基准\n这类基准要求智能体在真实或模拟的网页环境中执行多步操作，以完成复杂任务。\n\n| 基准 | 核心特点 | 使用场景 | 例题 |\n|------|----------|----------|------|\n| **BrowseComp**（OpenAI）【3†L16-L35】 | - 1,266 个挑战性问题，需持续导航、查找难以直接搜索的信息<br>- 问题设计保证人类在 10 分钟内无法解决，且现有模型（包括带浏览的 ChatGPT）也难以完成<br>- 答案为短字符串，便于自动验证 | 评估智能体在复杂网络浏览任务中的持久性、创造性搜索能力 | “找出 2018‑2023 年间在 EMNLP 会议上发表、第一作者本科毕业于达特茅斯学院、第四作者本科毕业于宾夕法尼亚大学的科学论文标题”【3†L45-L49】 |\n| **WebVoyager**【24†L18-L33】 | - 多模态（视觉+文本）端到端网页智能体基准<br>- 包含 300 个真实任务，覆盖 15 个常用网站（如 Amazon、Wikipedia）<br>- 使用 GPT‑4V 进行自动评估，与人工评估一致性达 85.3% | 评估智能体在真实网站上的端到端任务完成能力（多模态输入） | “在 Amazon 上找到‘无线蓝牙耳机’并按价格从低到高排序，然后选择最便宜的一款加入购物车” |\n| **WebArena**【25†L15-L35】 | - 完全可交互的、高度真实的网页环境（4 个领域：电商、论坛、协作开发、内容管理）<br>- 812 项长视野任务，评估功能正确性而非表面动作匹配<br>- 最佳 GPT‑4 智能体成功率仅 14.41%，远低于人类的 78.24% | 评估自主智能体在真实网页环境中的复杂任务执行能力 | “在电商网站中购买一双‘黑色运动鞋’，价格低于 $80，并将收货地址修改为纽约”【25†L86-L92】 |\n| **Mind2Web**【19†L4-L14】 | - 首个面向通用网页智能体的数据集，任务来自真实网站<br>- 支持在线评估（Online‑Mind2Web），包含 300 个实时任务、136 个流行网站<br>- 强调多步指令跟随与复杂交互 | 评估智能体在任意网站上的通用任务完成能力 | “在 GitHub 上创建一个名为 ‘test‑repo’ 的公开仓库，并添加一个 README.md 文件” |\n| **VisualWebArena**【20†L4-L10】 | - 多模态网页智能体基准，任务依赖视觉信息（截图）<br>- 包含丰富的视觉接地任务，如识别图片、按钮位置等<br>- 评估智能体在真实视觉网页环境中的表现 | 评估多模态智能体在视觉丰富的网页任务中的能力 | “在纽约时报网站找到今天头条新闻的配图，并描述图中主要物体” |\n| **WebVLN**【22†L4-L11】 | - 网站上的视觉‑语言导航任务，智能体需根据问题指令在网站中导航<br>- 数据集 WebVLN‑v1 包含多种网站导航场景<br>- 评估导航成功率与任务完成度 | 评估智能体在网站中基于视觉与语言指令的导航能力 | “找到‘联系我们’页面，并提取其中的邮箱地址” |\n| **WebNav**【23†L4-L11】 | - 目标驱动的网页导航基准，环境为部分观测<br>- 智能体需理解自然语言指令并在网页中规划动作序列<br>- 早期经典任务，常用于测试基础导航能力 | 评估智能体在部分观测网页环境中的规划与指令跟随能力 | “导航到 Wikipedia 上关于‘人工智能’的页面，并摘录第一段内容” |\n\n---\n\n## 🛠️ 2. 工具使用/通用助手类基准\n这类基准测试智能体调用外部工具（如搜索、计算、API）来解决实际问题的能力。\n\n| 基准 | 核心特点 | 使用场景 | 例题 |\n|------|----------|----------|------|\n| **GAIA**【26†L10-L25】 | - 466 个真实世界问题，需要推理、多模态处理、网页浏览、工具使用等综合能力<br>- 对人类简单（92% 准确率），但对最强 GPT‑4 仅 15%<br>- 答案唯一、事实性，支持自动评估 | 评估通用 AI 助手在需要多步工具使用的真实任务中的表现 | “2022 年法国总统选举的获胜者出生在哪个城市？”【26†L10-L25】 |\n| **ToolBench**（OpenBMB）【7†L4-L10】 | - 包含多样软件工具的真实任务基准<br>- 提供易用的基础设施（虚拟 API 服务器等）<br>- 评估 LLM 调用工具完成复杂任务的能力 | 评估智能体在工具调用与组合上的熟练度 | “使用计算器 API 计算 (3.14 * 7.2)^2，然后使用搜索 API 查找‘圆周率’的最新精确值” |\n| **HSCodeComp**【2†L16-L35】 | - 首个针对深度搜索智能体在分层规则应用上的专家级电商基准<br>- 任务：根据嘈杂的商品描述预测 10 位 HS 编码（关税规则）<br>- 632 个商品条目，标注由多位电商专家完成，最佳智能体准确率仅 46.8%（人类 95%） | 评估智能体在复杂规则（如关税分类）下的深层推理与搜索能力 | “给定商品描述‘男士纯棉针织T恤，领口为V领’，预测其 10 位 HS 编码”【2†L24-L32】 |\n\n---\n\n## ❓ 3. 问答/检索类基准\n这类基准主要评估智能体在开放域或特定领域的信息检索与问答能力。\n\n| 基准 | 核心特点 | 使用场景 | 例题 |\n|------|----------|----------|------|\n| **HotpotQA**【8†L4-L10】 | - 113k 个基于 Wikipedia 的多跳问答对<br>- 需要跨多个文档推理，并提供句子级支持事实<br>- 涵盖多样的问题类型（比较、列举、因果等） | 评估模型在需要多跳推理的开放域问答中的能力 | “《了不起的盖茨比》的作者还写过哪些小说？” |\n| **FEVER**【27†L16-L28】 | - 185,445 个基于 Wikipedia 的声明，需分类为 Supported/Refuted/NotEnoughInfo<br>- 要求提供证据句子（可多句、多页面）<br>- 挑战性高（最佳系统仅 31.87% 准确率） | 评估系统在事实核查与证据检索上的能力 | “斐济最大的岛屿是考艾岛。”【27†L48-L50】 |\n| **TriviaQA**【28†L5-L10】 | - 超过 650k 个问答‑证据三元组，问题由琐事爱好者编写<br>- 每个问题平均提供 6 篇证据文档，适合远程监督阅读<br>- 包含阅读理解和开放域 QA 两种任务设置 | 评估模型在开放域琐事问答中的检索与阅读理解能力 | “哪位演员在《星球大战：原力觉醒》中扮演凯洛·伦？” |\n| **Natural Questions**【11†L4-L8】 | - 来自 Google 搜索的真实用户问题，答案来自 Wikipedia<br>- 包含长答案（段落）和短答案（实体/日期等）<br>- 训练集 307k，开发/测试集各 8k | 评估开放域问答系统对真实用户查询的响应能力 | “谁写了《傲慢与偏见》？”（答案：“简·奥斯汀”） |\n| **MS MARCO**【12†L4-L10】 | - 基于 Bing 真实搜索查询的检索/问答基准<br>- 包含段落排序、文档排序、问答等任务<br>- 数据规模大（约 500k 查询，880 万段落） | 训练与评估检索模型在真实搜索场景下的表现 | “如何冲泡咖啡？”（系统需检索相关段落回答） |\n| **BEIR**【13†L4-L10】 | - 异构零样本检索基准，包含 18 个数据集、9 种任务（事实核查、QA、对话检索等）<br>- 评估模型在未见任务上的泛化能力<br>- 已成为检索模型的标准评估套件 | 评估检索模型在零样本设置下的跨任务泛化能力 | 给定查询“全球变暖的主要原因”，在文档集合中检索相关文档 |\n| **MIRACL**【14†L4-L10】 | - 多语言检索基准，覆盖 18 种语言（包括中文、阿拉伯语等）<br>- 专注于 ad‑hoc 检索，所有查询与文档均经过人工标注<br>- 提供多语言检索挑战（WSDM 2023 Cup） | 评估多语言检索模型在不同语言下的检索效果 | 中文查询：“全球变暖的原因”，检索相关中文文档 |\n\n---\n\n## 🧩 4. 组合搜索/推理类基准\n这类基准专门测试智能体在组合搜索问题上的逻辑推理与规划能力。\n\n| 基准 | 核心特点 | 使用场景 | 例题 |\n|------|----------|----------|------|\n| **SearchBench**【30†L22-L36】 | - 11 种独特的搜索问题类型（路径查找、谜题、子集和、排序、欠定系统等）<br>- 自动生成任意数量实例，并评估解决方案的可行性、正确性、最优性<br>- GPT‑4 在纯文本模式下仅解决 1.4%，通过 A* 代码生成提升至 11.7% | 评估 LLM 在组合搜索问题上的推理与规划能力 | “在一个 5x5 迷宫中，从起点 (0,0) 到终点 (4,4) 的最短路径是什么？（障碍物位置已知）”【30†L56-L64】 |\n\n---\n\n## 🛒 5. 特定领域/仿真环境基准\n这类基准针对特定领域（如电商、家庭环境）设计，测试智能体在仿真环境中的交互能力。\n\n| 基准 | 核心特点 | 使用场景 | 例题 |\n|------|----------|----------|------|\n| **WebShop**【29†L15-L30】 | - 模拟电商网站环境，包含 118 万真实商品、12,087 条众包指令<br>- 智能体需根据文本指令搜索、筛选、购买商品<br>- 奖励自动计算，最佳模型成功率 29%（人类 59%） | 评估语言智能体在电商任务中的导航、搜索与决策能力 | “找到一款价格低于 $50 的无线蓝牙耳机，评分至少 4 星，并加入购物车”【29†L20-L25】 |\n| **ALFWorld**【18†L4-L10】 | - 将文本环境（TextWorld）与具身环境（ALFRED）对齐的基准<br>- 智能体先在文本中学习抽象策略，然后在视觉环境中执行<br>- 包含 134 个家庭任务（如“在厨房找到杯子并拿到客厅”） | 评估智能体在文本与具身环境之间的迁移学习与规划能力 | “在厨房找到‘杯子’，把它拿到客厅放在桌子上” |\n\n---\n\n## 📊 总结\n以上基准覆盖了搜索智能体的不同能力维度：从**网页浏览、工具使用、问答检索**，到**组合搜索推理**以及**特定领域仿真**。选择合适的基准取决于您要评估的具体能力（如多模态交互、工具调用、多跳推理等）。这些基准大多提供公开的数据集与评估代码，可直接用于实验。\n\n> 注：本文基于 2025 年 11 月前的公开资料整理，部分基准可能已有更新版本，建议访问对应官网获取最新信息。",
728 + "reasoning_content": "在同时打开多个基准测试相关页面后，我们获得了大量信息。WebVoyager 是一个多模态网页智能体基准测试，专注于在真实网站上进行端到端任务执行，其特点包括多模态输入（屏幕截图和文本）、在15个常用网站上的300个任务，以及使用GPT-4V进行自动评估。使用场景是评估网页智能体在真实世界中的导航和任务完成能力。例题可能是在亚马逊上找到特定产品并将其加入购物车。\n\nWebArena 是一个真实且可复现的网页环境，用于评估自主智能体。其特点包括四个功能完整的网站（电子商务、论坛、协作开发、内容管理）、812项长视野任务，以及专注于功能正确性的评估。使用场景是评估智能体在复杂网页任务中的表现。例题可能是在电子商务网站上购买符合特定标准的商品。\n\nGAIA 是一个针对通用AI助手的基准测试，要求具备推理、多模态处理、网页浏览和工具使用等能力。其特点包括466个问题，对人类来说概念简单但对AI具有挑战性，答案简洁明确，支持自动评估。使用场景是评估AI助手在真实世界问题中的表现。例题可能是“找出2022年法国总统选举中获胜者的出生城市”。\n\nToolBench 是一个用于评估LLM工具操作能力的基准测试。其特点包括多样化的软件工具、真实世界任务，以及易于使用的基础设施。使用场景是评估智能体使用工具完成任务的能力。例题可能是使用计算器API计算表达式，或使用搜索API查找信息。\n\nHotpotQA 是一个多跳问答基准测试，其特点包括基于Wikipedia的113k个问题-答案对，需要多文档推理，并提供句子级别的支持事实。使用场景是评估模型在回答需要多跳推理的问题时的能力。例题可能是“《了不起的盖茨比》的作者还写过哪些小说？”。\n\nFEVER 是一个事实提取和验证基准测试，其特点包括185,445个基于Wikipedia的声明，需要分类为支持、反驳或信息不足，并需要提供证据句子。使用场景是评估系统验证声称事实的能力。例题可能是“斐济最大的岛屿是考艾岛”，系统需要检索证据来支持或反驳这一说法。\n\nTriviaQA 是一个阅读理解基准测试，其特点包括超过65万个问题-答案-证据三元组，问题由琐事爱好者编写，并附带证据文档。使用场景是评估模型在开放域问答中的能力。例题可能是“哪位演员在《星球大战：原力觉醒》中扮演凯洛·伦？”。\n\nNatural Questions 是一个问答基准测试，其特点包括来自谷歌搜索的真实问题，答案来自Wikipedia，并包含长答案和短答案。使用场景是评估开放域问答系统。例题可能是“谁写了《傲慢与偏见》？”，答案可能是“简·奥斯汀”。\n\nMS MARCO 是一个用于深度学习搜索的基准测试，其特点包括来自Bing的真实搜索查询、人工生成的答案，以及用于段落排序和文档排序的任务。使用场景是训练和评估检索模型。例题可能是“如何冲泡咖啡？”，系统需要检索相关段落来回答。\n\nBEIR 是一个用于零样本评估的异构信息检索基准测试，其特点包括18个数据集、9项任务，涵盖事实核查、问答、对话检索等。使用场景是评估检索模型在零样本设置下的泛化能力。例题可能是给定一个查询，在文档集合中检索相关文档。\n\nMIRACL 是一个多语言信息检索基准测试，其特点涵盖18种语言，专注于ad-hoc检索，并包含人工标注的相关性判断。使用场景是评估多语言检索模型。例题可能是用中文查询“全球变暖的原因”，并检索相关文档。\n\nWebShop 是一个模拟电子商务网站环境，用于评估语言智能体。其特点包括118万个真实产品、12,087条众包文本指令，以及需要导航、搜索和购买的任务。使用场景是评估智能体在电子商务任务中的表现。例题可能是“找到一款价格低于50美元的无线蓝牙耳机”。\n\nALFWorld 是一个将文本和具身环境对齐的基准测试，用于在文本世界中训练智能体，然后将其转移到具身环境中。其特点包括基于文本的环境，与ALFRED数据集并行，并专注于指令跟随。使用场景是评估智能体在家庭任务中的规划能力。例题可能是“在厨房里找到一个杯子，把它拿到客厅”。\n\nMind2Web 是一个用于评估网页通用智能体的基准测试，其特点包括来自真实网站的任务、需要多步交互，并专注于语言指令跟随。使用场景是评估智能体在复杂网页任务中的表现。例题可能是“在GitHub上创建一个新的仓库，命名为‘test-repo’”。\n\nVisualWebArena 是一个用于多模态智能体的基准测试，其特点包括真实视觉网页任务、需要理解视觉和文本元素，并评估功能正确性。使用场景是评估多模态智能体在视觉接地网页任务中的表现。例题可能是“在纽约时报网站上找到今天头条新闻的图片”。\n\nSearchBench 是一个用于评估LLM解决搜索问题能力的基准测试，其特点包括11种搜索问题类型、自动生成实例，并评估可行性、正确性和最优性。使用场景是评估LLM在组合搜索问题中的推理能力。例题可能是“在一个迷宫中找到从起点到终点的最短路径”。\n\nWebVLN 是一个网站上的视觉语言导航基准测试，其特点包括基于问题的指令、需要导航网站，并评估导航成功率。使用场景是评估智能体在网站中基于视觉和语言指令进行导航的能力。例题可能是“找到‘联系我们’页面并提取电子邮件地址”。\n\nWebNav 是一个用于目标驱动网页导航的基准测试，其特点包括部分观察环境、需要规划和理解自然语言指令。使用场景是评估智能体在网页导航中的能力。例题可能是“导航到维基百科上关于‘人工智能’的页面”。\n\n现在，我需要整理一份针对搜索智能体的基准测试列表。考虑到用户可能更关注中文资料，我会优先考虑那些有中文描述的基准测试，例如BrowseComp、HSCodeComp、GAIA、WebVoyager、WebArena、Mind2Web、WebShop、ALFWorld、HotpotQA、FEVER、TriviaQA、Natural Questions、MS MARCO、BEIR、MIRACL、SearchBench、WebVLN、WebNav。我会为每个基准测试提供特点、使用场景和例题。同时，我会引用相关来源，例如HSCodeComp论文、BrowseComp博客文章、WebVoyager论文、WebArena论文、GAIA论文、WebShop论文、SearchBench论文等。我的回答将用中文呈现，并确保内容全面、结构清晰。现在开始撰写回答。",
729 + "tool_calls": []
730 + }
731 + ]
732 +}
rust/src/chat/src/renderer/deepseek_v32/fixtures/test_input_search_wo_date.json new
+533
@@ -0,0 +1,533 @@
1 +{
2 + "messages": [
3 + {
4 + "role": "developer",
5 + "content": "In a year before 2023, between December 1 and 15 (inclusive), a press release reported that as of the 5th of that month, a country had received a total of 16,201,670 vaccine doses, with 7,583,134 administered. That same year, a company's annual report revealed that its nearly 2.2 million members and supporters contributed over $49 million, with 12% of total revenue coming from contributed services. The following year, the annual report stated that the organization had played a role in rescinding a previous administration’s plan to open most of the 22 million acres of the National Petroleum Reserve in Alaska to oil and gas drilling and that contributed services had decreased by 2% to 10%. Who was the emeritus board member that year? Name only, no titles.",
6 + "tools": [
7 + {
8 + "type": "function",
9 + "function": {
10 + "name": "search",
11 + "description": "Searches for information related to query and displays topn results.",
12 + "parameters": {
13 + "type": "object",
14 + "properties": {
15 + "query": {
16 + "type": "string",
17 + "description": "The search query string"
18 + },
19 + "topn": {
20 + "type": "integer",
21 + "description": "Number of top results to display",
22 + "default": 10
23 + },
24 + "source": {
25 + "type": "string",
26 + "description": "Source to search within",
27 + "enum": [
28 + "web",
29 + "news"
30 + ],
31 + "default": "web"
32 + }
33 + },
34 + "required": [
35 + "query"
36 + ],
37 + "additionalProperties": false,
38 + "$schema": "http://json-schema.org/draft-07/schema#"
39 + }
40 + }
41 + },
42 + {
43 + "type": "function",
44 + "function": {
45 + "name": "open",
46 + "description": "Opens the link id from the page indicated by cursor starting at line number loc, showing num_lines lines. Valid link ids are displayed with the formatting: 【{id}†.*】. If cursor is not provided, the most recent page is implied. If id is a string, it is treated as a fully qualified URL associated with source. If loc is not provided, the viewport will be positioned at the beginning of the document or centered on the most relevant passage, if available. Use this function without id to scroll to a new location of an opened page.",
47 + "parameters": {
48 + "type": "object",
49 + "properties": {
50 + "id": {
51 + "description": "Link ID (number) or fully qualified URL (string)",
52 + "anyOf": [
53 + {
54 + "type": "integer"
55 + },
56 + {
57 + "type": "string"
58 + }
59 + ],
60 + "default": -1
61 + },
62 + "cursor": {
63 + "type": "integer",
64 + "description": "Page cursor indicator",
65 + "default": -1
66 + },
67 + "loc": {
68 + "type": "integer",
69 + "description": "Starting line number",
70 + "default": -1
71 + },
72 + "num_lines": {
73 + "type": "integer",
74 + "description": "Number of lines to show",
75 + "default": -1
76 + },
77 + "view_source": {
78 + "type": "boolean",
79 + "description": "Whether to view source",
80 + "default": false
81 + },
82 + "source": {
83 + "type": "string",
84 + "description": "Source associated with the URL",
85 + "enum": [
86 + "web",
87 + "news"
88 + ],
89 + "default": "web"
90 + }
91 + },
92 + "additionalProperties": false,
93 + "$schema": "http://json-schema.org/draft-07/schema#"
94 + }
95 + }
96 + },
97 + {
98 + "type": "function",
99 + "function": {
100 + "name": "find",
101 + "description": "Finds exact matches of pattern in the current page, or the page given by cursor.",
102 + "parameters": {
103 + "type": "object",
104 + "properties": {
105 + "pattern": {
106 + "type": "string",
107 + "description": "The exact pattern to search for"
108 + },
109 + "cursor": {
110 + "type": "integer",
111 + "description": "Page cursor to search in",
112 + "default": -1
113 + }
114 + },
115 + "required": [
116 + "pattern"
117 + ],
118 + "additionalProperties": false,
119 + "$schema": "http://json-schema.org/draft-07/schema#"
120 + }
121 + }
122 + },
123 + {
124 + "type": "function",
125 + "function": {
126 + "name": "python",
127 + "description": "Use this tool to execute Python code. The code will not be shown to the user. This tool should be used for internal reasoning, but not for code that is intended to be visible to the user (e.g. when creating plots, tables, or files). When you send a message containing Python code to python, it will be executed in a stateful Jupyter notebook environment. python will respond with the output of the execution or time out after 120.0 seconds.",
128 + "parameters": {
129 + "type": "object",
130 + "properties": {
131 + "code": {
132 + "type": "string",
133 + "description": "The python code to execute in the jupyter environment."
134 + }
135 + },
136 + "required": [
137 + "code"
138 + ],
139 + "additionalProperties": false,
140 + "$schema": "http://json-schema.org/draft-07/schema#"
141 + }
142 + }
143 + }
144 + ]
145 + },
146 + {
147 + "role": "assistant",
148 + "content": "",
149 + "reasoning_content": "The user asks: \"In a year before 2023, between December 1 and 15 (inclusive), a press release reported that as of the 5th of that month, a country had received a total of 16,201,670 vaccine doses, with 7,583,134 administered. That same year, a company's annual report revealed that its nearly 2.2 million members and supporters contributed over $49 million, with 12% of total revenue coming from contributed services. The following year, the annual report stated that the organization had played a role in rescinding a previous administration’s plan to open most of the 22 million acres of the National Petroleum Reserve in Alaska to oil and gas drilling and that contributed services had decreased by 2% to 10%. Who was the emeritus board member that year? Name only, no titles.\"\n\nWe need to parse the puzzle.\n\nWe have clues about a press release about vaccine doses for a country, dated between Dec 1-15 inclusive, reporting as of 5th of that month total doses received 16,201,670 and administered 7,583,134. That is likely a COVID-19 vaccine update for some country. Which country? Might be a smaller country. Could be something like \"Uganda\", \"Zimbabwe\", \"Kenya\", \"Philippines\", etc. Let's search.\n\nSearch for \"16,201,670 vaccine doses\" \"7,583,134 administered\".",
150 + "tool_calls": [
151 + {
152 + "id": "fixture-tool-call-1-0",
153 + "type": "function",
154 + "function": {
155 + "name": "search",
156 + "arguments": "{\"query\": \"\\\"16,201,670\\\" vaccine doses\", \"topn\": 10, \"source\": \"web\"}"
157 + }
158 + }
159 + ]
160 + },
161 + {
162 + "role": "tool",
163 + "content": "[0] \"16,201,670\" vaccine doses (Search_Results/\"16,201,670\" vaccine doses)\n**viewing lines [0 - 38] of 38**\n\nL0: \nL1: URL: Search_Results/\"16,201,670\" vaccine doses\nL2: # Search Results\nL3: \nL4: * 【0†Kenya's Economy is Showing Resilience as Output Rises ...; publish_date: \nL5: none†www.worldbank.org】 Dec 14, 2021 — As of December 5, 2021, Kenya had \nL6: received a total of 16,201,670 vaccines, with 7,583,134 administered. While \nL7: vaccine acceptance is ...\nL8: * 【1†MINISTRY OF HEALTH; publish_date: none†covidhub.mediacouncil.or.ke】 Dec \nL9: 1, 2021 — Total Doses Received 16,201,670. Table 10 gives the total vaccines \nL10: received since the start of Covid -19 vaccination exercise in the country.\nL11: * 【2†Output Result Page; publish_date: none†open.unicef.org】 ... 16,201,670 \nL12: doses of multiple vaccines nationwide and full vaccination of 15.5 per cent with\nL13: two doses of COVID-19 vaccine as of 31 December 2021.\nL14: * 【3†rebased GDP; publish_date: none†documents1.worldbank.org】 Dec 7, 2021 — \nL15: As of December 5, 2021,. Kenya had received a total of 16,201,670 vaccines, \nL16: with. 7,583,134 administered. Vaccine acceptance is reportedly high.\nL17: * 【4†Integrated Annual Report; publish_date: none†www.co-opbank.co.ke】 May 27,\nL18: 2022 — ... doses of Covid-19 vaccines and administered close to 17 million ... \nL19: 16,201,670, huku 7,583,134 zikiwa tayari zimedungwa watu. Bado kuna ...\nL20: * 【5†World Bank lifts Kenya's growth prospect to 5% in 2021; publish_date: \nL21: none†www.africa-press.net】 ... 16,201,670 vaccines, with 7,583,134 administered.\nL22: While vaccine acceptance is reportedly high there is still a long way to go \nL23: towards the government's ...\nL24: * 【6†2020 Annual Report United States; publish_date: none†www.hi-us.org】 \nL25: including medicine, vaccines, Covid-19 prevention kits, and kits of ... \nL26: 16,201,670. 9,630,584. Prepaid expenses. 56,707. 22,601. Total current assets.\nL27: * 【7†HOSPITAL CENTRAL DE LA DEFENSA; publish_date: \nL28: none†contrataciondelestado.es】 16201670 X. URPA-URPA. 0. HCD-000393/00. KONTRON.\nL29: 7000 ... LASER SHOT LBP-1120. CNEUR-C. NEUROFISIOLOGÍA. 0 ... VICTOREEN CT HEAD\nL30: DOSE PHANTON. 43650-014.\nL31: * 【8†body metal concentrations: Topics by ...; publish_date: \nL32: none†www.science.gov】 PMID:16201670. Design Characteristics and Tobacco Metal \nL33: Concentrations in Filtered Cigars. PubMed. Caruso, Rosalie V; O'Connor, Richard \nL34: J; Travers, Mark J ...\nL35: * 【9†Public Offices, Personal Demands: Capability in ...; publish_date: \nL36: none†dokumen.pub】 ... (16201670).9 The Institutiones Politicae consists of two \nL37: books. Book one deals with the origin, nature and goal of the commonwealth on \nL38: the one hand, and the ...",
164 + "tool_call_id": "fixture-tool-call-1-0"
165 + },
166 + {
167 + "role": "assistant",
168 + "content": "",
169 + "reasoning_content": "Result shows Kenya. So the press release is about Kenya's vaccine doses as of Dec 5, 2021. The press release likely by World Bank or something. The search result shows \"World Bank lifts Kenya's growth prospect...\" maybe a press release dated Dec 14, 2021. But the clue says \"between December 1 and 15 (inclusive), a press release reported that as of the 5th of that month, a country had received a total of 16,201,670 vaccine doses, with 7,583,134 administered.\" That matches Kenya on Dec 5, 2021. So the year is 2021.\n\nNow the next clue: \"That same year, a company's annual report revealed that its nearly 2.2 million members and supporters contributed over $49 million, with 12% of total revenue coming from contributed services.\"\n\nThis sounds like an environmental or conservation organization. \"Nearly 2.2 million members and supporters contributed over $49 million.\" That might be something like \"The Nature Conservancy\", \"World Wildlife Fund\", \"Sierra Club\", \"Environmental Defense Fund\", \"Natural Resources Defense Council\", etc. Let's search.\n\nSearch for \"nearly 2.2 million members and supporters contributed over $49 million\" \"12% of total revenue coming from contributed services\".",
170 + "tool_calls": [
171 + {
172 + "id": "fixture-tool-call-3-0",
173 + "type": "function",
174 + "function": {
175 + "name": "search",
176 + "arguments": "{\"query\": \"\\\"nearly 2.2 million members\\\" \\\"contributed over $49 million\\\"\", \"topn\": 10, \"source\": \"web\"}"
177 + }
178 + }
179 + ]
180 + },
181 + {
182 + "role": "tool",
183 + "content": "[1] \"nearly 2.2 million members\" \"contributed over $49 million\" (Search_Results/\"nearly 2.2 million members\" \"contributed over $49 million\")\n**viewing lines [0 - 41] of 41**\n\nL0: \nL1: URL: Search_Results/\"nearly 2.2 million members\" \"contributed over $49 million\"\nL2: # Search Results\nL3: \nL4: * 【0†The Global Good Fund, Microsoft and Goodwill join forces ...; \nL5: publish_date: none†globalgoodfund.org】 Sep 28, 2022 — ... more than 2.2 million \nL6: people. They additionally contribute close to $49 billion in local, state and \nL7: federal tax revenues. When COVID hit ...\nL8: * 【1†Almost 22 billion American tax dollars spent to wipe out a ...; \nL9: publish_date: none†www.facebook.com】 US military funding for Israel's war crimes\nL10: in Lebanon and Gaza has now cost US taxpayers over $22 billion. When millions \nL11: struggle to afford the ...\nL12: * 【2†Corporate America has largely abandoned its post-January ...; \nL13: publish_date: none†www.citizensforethics.org】 Jul 29, 2025 — Since the January 6\nL14: insurrection, over 2,000 corporate and industry group PACs have given over $174\nL15: million to members of the Sedition ...\nL16: * 【3†Audit shows millions in questionable taxpayer spending at ...; \nL17: publish_date: none†www.aol.com】 18 hours ago — ... nearly doubled from 1.3 \nL18: million to about 2.2 million. That is more than one in four Washington state \nL19: residents receiving Medicaid, and the ...\nL20: * 【4†Incarceration and Poverty in the United States - AAF; publish_date: \nL21: none†www.americanactionforum.org】 Jun 30, 2020 — The United States currently \nL22: incarcerates 2.2 million people, nearly half of whom are non-violent drug \nL23: offenders, accused people held pre ...\nL24: * 【5†How the NRA Spends Revenue (2023); publish_date: none†paddockpost.com】 \nL25: Apr 25, 2025 — In 2023, the NRA reported total revenue of $176 million (compared\nL26: to $211 million in 2022, $227 million in 2021, $282 million in 2020, $292 ...\nL27: * 【6†NCA Newsroom; publish_date: none†www.ncausa.org】 ... 2.2 million U.S. \nL28: jobs—operating in every U.S. state and territory and contributing nearly $350 \nL29: billion to the U.S. economy every year. For media inquiries ...\nL30: * 【7†The Early Childhood Educator Workforce; publish_date: \nL31: none†cscce.berkeley.edu】 About 2.2 million adults are paid to care for and \nL32: educate more than 9.7 million children between birth and age five in center-\nL33: based and home-based settings.\nL34: * 【8†Bringing Out the Worst in Congress: CRA By the Numbers ...; publish_date:\nL35: none†progressivereform.org】 Jul 14, 2025 — In fact, the top ten recipients (out\nL36: of 25) account for almost $7 million, showing how these contributions are \nL37: targeted and strategic. Real ...\nL38: * 【9†Club 100: Over 100 Americans gave more than $100000 to ...; publish_date:\nL39: none†publicintegrity.org】 Jan 7, 2022 — Through Internet investments during the\nL40: 1990s, he attracted 2.2 million online accounts valued at $174 million. Schwab \nL41: has been a vocal ...",
184 + "tool_call_id": "fixture-tool-call-3-0"
185 + },
186 + {
187 + "role": "assistant",
188 + "content": "",
189 + "reasoning_content": "Not that. Try search for \"2.2 million members and supporters contributed over $49 million\".",
190 + "tool_calls": [
191 + {
192 + "id": "fixture-tool-call-5-0",
193 + "type": "function",
194 + "function": {
195 + "name": "search",
196 + "arguments": "{\"query\": \"\\\"2.2 million members\\\" \\\"contributed over $49\\\"\", \"topn\": 10, \"source\": \"web\"}"
197 + }
198 + }
199 + ]
200 + },
201 + {
202 + "role": "tool",
203 + "content": "[2] \"2.2 million members\" \"contributed over $49\" (Search_Results/\"2.2 million members\" \"contributed over $49\")\n**viewing lines [0 - 2] of 2**\n\nL0: \nL1: URL: Search_Results/\"2.2 million members\" \"contributed over $49\"\nL2: # Search Results",
204 + "tool_call_id": "fixture-tool-call-5-0"
205 + },
206 + {
207 + "role": "assistant",
208 + "content": "",
209 + "reasoning_content": "Maybe it's an annual report of a non-profit. Search for \"annual report 2.2 million members contributed $49 million\".",
210 + "tool_calls": [
211 + {
212 + "id": "fixture-tool-call-7-0",
213 + "type": "function",
214 + "function": {
215 + "name": "search",
216 + "arguments": "{\"query\": \"\\\"2.2 million members\\\" \\\"annual report\\\" $49 million\", \"topn\": 10, \"source\": \"web\"}"
217 + }
218 + }
219 + ]
220 + },
221 + {
222 + "role": "tool",
223 + "content": "[3] \"2.2 million members\" \"annual report\" $49 million (Search_Results/\"2.2 million members\" \"annual report\" $49 million)\n**viewing lines [0 - 36] of 36**\n\nL0: \nL1: URL: Search_Results/\"2.2 million members\" \"annual report\" $49 million\nL2: # Search Results\nL3: \nL4: * 【0†20-F; publish_date: none†www.sec.gov】 ANNUAL REPORT PURSUANT TO SECTION \nL5: ... Our membership grew from 2.1 million members as at December 31, 2023 to 2.2 \nL6: million members as at December 31, 2024.\nL7: * 【1†Oportun Reports Fourth Quarter and Full Year 2023 ...; publish_date: \nL8: none†investor.oportun.com】 Mar 12, 2024 — Oportun (Nasdaq: OPRT) is a mission-\nL9: driven fintech that puts its 2.2 million members' financial goals within reach. \nL10: ... annual report on ...\nL11: * 【2†2 0 21; publish_date: none†www.annualreports.com】 ANNUAL REPORT. 2. 0. \nL12: 21. 2. 0. 21. Page 2. 2. DEFENDERS OF WILDLIFE. 2. 0. 21. 2. 0. 21 ... In 2021, \nL13: Defenders of Wildlife's nearly 2.2 million members and.\nL14: * 【3†Annual report and accounts 2020; publish_date: none†www.3i.com】 \nL15: Disclaimer. The Annual report and accounts have been prepared solely to provide \nL16: information to shareholders. ... 2.2 million members. In December 2019, we sold \nL17: ...\nL18: * 【4†united states securities and exchange commission; publish_date: \nL19: none†s21.q4cdn.com】 Dec 14, 2018 — ☒ ANNUAL REPORT ... approximately 5.7 million\nL20: Members, including approximately 2.2 million Members from the acquisition of \nL21: Best Doctors.\nL22: * 【5†securities and exchange commission; publish_date: none†www.sec.gov】 ... \nL23: Annual Report on Form 10-K filed with the SEC on March 10, 2020. The ... 2.2 \nL24: million members at June 30, 2020. The following table highlights the ...\nL25: * 【6†2022-23 Annual Comprehensive Financial Report - CalPERS; publish_date: \nL26: none†www.calpers.ca.gov】 Nov 17, 2023 — ... annual report is consistent with \nL27: these accounting principles ... 2.2 million members and health benefit services \nL28: to over 1.5 million ...\nL29: * 【7†Medical Centers Report; publish_date: \nL30: none†regents.universityofcalifornia.edu】 Oct 8, 2019 — which will allow its 2.2 \nL31: million members to receive advanced and ... annual report of the University of \nL32: California. For the Medical ...\nL33: * 【8†Working TogeTher to make HealtH care better; publish_date: \nL34: none†www.unitedhealthgroup.com】 Feb 8, 2012 — ANNUAL REPORT PURSUANT TO SECTION \nL35: 13 OR 15(d) OF THE ... approximately 2.2 million members enrolled in its \nL36: Medicare. Advantage ...",
224 + "tool_call_id": "fixture-tool-call-7-0"
225 + },
226 + {
227 + "role": "assistant",
228 + "content": "",
229 + "reasoning_content": "Result 2 shows \"Defenders of Wildlife\". That might be the organization. Let's open.",
230 + "tool_calls": [
231 + {
232 + "id": "fixture-tool-call-9-0",
233 + "type": "function",
234 + "function": {
235 + "name": "open",
236 + "arguments": "{\"cursor\": 3, \"id\": 2}"
237 + }
238 + }
239 + ]
240 + },
241 + {
242 + "role": "tool",
243 + "content": "[4] www.annualreports.com (https://www.annualreports.com/HostedData/AnnualReportArchive/d/defenders-wildlife_2021.pdf)\n**viewing lines [0 - 53] of 261**\n\nL0: \nL1: URL: https://www.annualreports.com/HostedData/AnnualReportArchive/d/defenders-\nL2: wildlife_2021.pdf\nL3: ===== Page 1 ===== 2021 ANNUAL REPORT ===== Page 2 ===== # DEFENDERS OF WILDLIFE\nL4: made important progress for imperiled species and vital landscapes across the \nL5: United States in 2021. \\--- **LAWYERS** immediately **challenged** the premature\nL6: and reckless decision to strip **gray wolves** of federal **Endangered Species \nL7: Act (ESA)** protections. For many decades, Defenders has led the effort to \nL8: protect and restore the gray wolf, and we will continue to fight the \nL9: unscientific and hostile anti-wolf policies that impede conservation progress \nL10: and will carry on our unrelenting battle to restore federal protections for this\nL11: iconic keystone species. \\--- **LOBBYISTS** worked around the clock to keep \nL12: wildlife and climate priorities in the **Infrastructure Investment and Jobs \nL13: Act**. We also continue fighting to keep important wildlife and habitat funding \nL14: in relevant **appropriations bills**. \\--- 2 DEFENDERS OF WILDLIFE ===== Page 3 \nL15: ===== POLICY EXPERTS pushed forward on the urgent need for a National \nL16: Biodiversity Strategy (NBS), an all-of-government approach to address the \nL17: unprecedented loss of wildlife and habitat we are experiencing. We have coupled \nL18: this with our new campaign to expand the National Wildlife Refuge System to \nL19: preserve our nation’s only lands set aside for wildlife. By defending, funding \nL20: and expanding our national wildlife refuges, we will directly address \nL21: biodiversity loss and climate change while promoting increased equitable access \nL22: to nature. FIELD TEAMS were on the ground helping to recover imperiled species. \nL23: From panthers and sea turtles in Florida to wolves, bison and black-footed \nL24: ferrets in Montana, Defenders’ conservation experts were in the field saving \nL25: wildlife all over the country. CONSERVATION INNOVATION EXPERTS provided \nL26: comprehensive analyses to guide policy and inform conservation strategies to \nL27: reach the goal of protecting 30% of our terrestrial and marine systems by 2030 \nL28: (“30x30”). Defenders’ Center for Conservation Innovation (CCI) produced a report\nL29: which details actions we need to take to achieve 30x30 while protecting \nL30: biodiversity and addressing the climate crisis. DEFENDERS.ORG ===== Page 4 =====\nL31: WE HAVE ACCOMPLISHED MUCH THIS YEAR WORKING WITH AN ADMINISTRATION THAT VALUES \nL32: SCIENCE AND CONSERVATION. Our many successes include the return of protections \nL33: to the Tongass National Forest in Alaska and the suspension of oil and gas \nL34: leasing permits for the Arctic National Wildlife Refuge. Defenders also played a\nL35: lead role in successfully defending the Migratory Bird Protection Act from \nL36: attack and securing critical protections for migratory birds. Throughout 2021 \nL37: Defenders made important progress for the wildlife and wild places we all love, \nL38: yet our nation’s wildlife still face unprecedented challenges. More and more \nL39: species, including manatees, Mexican gray wolves, polar bears and North Atlantic\nL40: right whales, face the very real threat of extinction because of climate \nL41: change, habitat loss, pollution and inadequate protections. In our work we \nL42: continue to demonstrate success by following the science, building partnerships,\nL43: leading with innovative solutions and focusing on collaborative and practical \nL44: approaches in defense of nature. Together we can make a lasting difference for \nL45: wildlife and habitats. Your support is essential to our work and we couldn’t do \nL46: it without you. Thank you for being an important part of our efforts to protect \nL47: imperiled wildlife across our nation. Jamie Rappaport Clark President and CEO \nL48: John Dayton Board Chair OUR NEW TELEVISION SERIES, WILDLIFE NATION WITH JEFF \nL49: CORWIN, premiered in October and airs every weekend on ABC stations nationwide. \nL50: This groundbreaking series takes Jeff across the country to meet Defenders’ \nL51: experts along with our science, tribal and community partners to educate and \nL52: inspire viewers to join with Defenders to protect wildlife. While highlighting \nL53: some of our important conservation work and raising awareness of important ",
244 + "tool_call_id": "fixture-tool-call-9-0"
245 + },
246 + {
247 + "role": "assistant",
248 + "content": "",
249 + "reasoning_content": "Search within PDF for \"2.2 million members\".",
250 + "tool_calls": [
251 + {
252 + "id": "fixture-tool-call-11-0",
253 + "type": "function",
254 + "function": {
255 + "name": "find",
256 + "arguments": "{\"cursor\": 4, \"pattern\": \"2.2\"}"
257 + }
258 + }
259 + ]
260 + },
261 + {
262 + "role": "tool",
263 + "content": "[5] Find results for text: `2.2` in `www.annualreports.com` (https://www.annualreports.com/HostedData/AnnualReportArchive/d/defenders-wildlife_2021.pdf/find?pattern=2.2)\n**viewing lines [0 - 4] of 4**\n\nL0: # 【0†match at L187】\nL1: 2021, Defenders of Wildlife’s nearly 2.2 million members and supporters \nL2: contributed more than $49 million for wildlife and wildlife habitat. This figure\nL3: includes all those who generously donated their time and expertise. The \nL4: steadfast support of our donors allows Defenders to sustain our program and ",
264 + "tool_call_id": "fixture-tool-call-11-0"
265 + },
266 + {
267 + "role": "assistant",
268 + "content": "",
269 + "reasoning_content": "Open around L180.",
270 + "tool_calls": [
271 + {
272 + "id": "fixture-tool-call-13-0",
273 + "type": "function",
274 + "function": {
275 + "name": "open",
276 + "arguments": "{\"cursor\": 4, \"loc\": 170}"
277 + }
278 + }
279 + ]
280 + },
281 + {
282 + "role": "tool",
283 + "content": "[6] www.annualreports.com (https://www.annualreports.com/HostedData/AnnualReportArchive/d/defenders-wildlife_2021.pdf)\n**viewing lines [170 - 215] of 261**\n\nL170: into the wild. **10 DEFENDERS OF WILDLIFE** ===== Page 11 ===== AFTER ANOTHER \nL171: SUCCESSFUL DEFENDERS LAWSUIT ON BEHALF OF THE CRITICALLY ENDANGERED RED WOLF, \nL172: FWS reversed its decision to limit the recovery area and committed to a robust \nL173: release strategy. After years of inaction, FWS released eight wolves to the wild\nL174: in North Carolina and plan to release nine more wolves in the spring of 2022. \nL175: This is an incredible win for this critically endangered species whose \nL176: population has dwindled down to single digits in the wild because of \nL177: mismanagement, vehicle strikes and poaching. DEFENDERS CONTINUED TO LEAD EFFORTS\nL178: TO PROTECT THE FLORIDA MANATEE, a beloved species that suffered the deadliest \nL179: year on record in 2021, tragically surpassing 1,000 deaths because of water \nL180: pollution and lack of warm water habitat. Defenders led advocacy and education \nL181: aimed at restoring the natural flow of the dammed Ocklawaha River, which would \nL182: provide critical warm-water habitat that manatees need to survive. Defenders’ \nL183: legal team continued to fight for manatees in the courts, holding government \nL184: agencies accountable for protecting critical habitat and addressing the \nL185: devastating water pollution that is killing the seagrass and causing manatees to\nL186: starve. DAVID TES | SAM FRENZY DRAW DEFENDERS.ORG 11 ===== Page 12 ===== In \nL187: 2021, Defenders of Wildlife’s nearly 2.2 million members and supporters \nL188: contributed more than $49 million for wildlife and wildlife habitat. This figure\nL189: includes all those who generously donated their time and expertise. The \nL190: steadfast support of our donors allows Defenders to sustain our program and \nL191: public education efforts in the field, the courts and on Capitol Hill. 2021 \nL192: SOURCES OF FUNDS Grants and contributions $29,057 Bequests, trusts and split \nL193: interests $7,692 Income from investments, annuity reserve funds and trusts \nL194: $3,354 Royalties and other income $3,576 Contributed services $6,140 Total \nL195: Revenue $49,819 USES OF FUNDS Biodiversity conservation $22,420 Constituency \nL196: mobilization $16,324 Fundraising $1,211 Management and general $5,865 Total \nL197: Expenses $45,820 Change in net assets $3,999 Net assets, start of year $41,145 \nL198: Net Assets, End of the Year $45,144 Dollars are in thousands. 12 DEFENDERS OF \nL199: WILDLIFE Grants and contributions 58% Income from investments 7% Requests, \nL200: trusts and split interests 15% Royalties and other income 7% Contributed \nL201: services 12% Program and support services 97% Fundraising 3% ===== Page 13 =====\nL202: DEFENDERS.ORG 13 ===== Page 14 ===== John Dayton, Esq.* Chair Texas Sheryl \nL203: Sachs* Vice Chair Connecticut Whitfield Bailey* Secretary Tennessee Mark Caylor*\nL204: Treasurer District of Columbia Sylvia Baca* California Donald Barry, Esq. \nL205: Oregon Dinah Bear, Esq. Arizona Kim O’Keefe Beck, Ph.D California Gabriela \nL206: Chavarria, Ph.D Washington Jamie Rappaport Clark* President & CEO Virginia \nL207: Caroline D. Gabel Maryland Keith M. Harper, Esq. Maryland John Donovan, Esq. \nL208: Montana and California James M. Hecker, Esq. Virginia Judith Ponsikoff, Ph.D* \nL209: California Terry L. Root, Ph.D Florida BOARD EMERITUS Terry C. Pelster, Esq. New\nL210: York *Executive Committee member* ST A F F & A D V I S O R S SCIENCE ADVISORS \nL211: Gabriela Chavarria, Ph.D Burke Museum, University of Washington Leah Gerber, \nL212: Ph.D Arizona State University Kathayoon Khalil, Ph.D Oregon Zoo Rurik List, Ph.D\nL213: Universidad Autónoma Metropolitana-Lerma Barry Noon, Ph.D Colorado State \nL214: University Terry L. Root, Ph.D Stanford University Mark Schwartz, Ph.D \nL215: University of California-Davis David S. Wilcox, Ph.D Princeton University Rae ",
284 + "tool_call_id": "fixture-tool-call-13-0"
285 + },
286 + {
287 + "role": "assistant",
288 + "content": "",
289 + "reasoning_content": "Scrolling further for \"Contributed services 12%\". Already seen. So the annual report is Defenders of Wildlife 2021 annual report.\n\nNext clue: \"The following year, the annual report stated that the organization had played a role in rescinding a previous administration’s plan to open most of the 22 million acres of the National Petroleum Reserve in Alaska to oil and gas drilling and that contributed services had decreased by 2% to 10%.\"\n\nThus, the following year (2022) annual report of Defenders of Wildlife includes mention of rescinding plan to open most of 22 million acres of NPR-A to oil and gas drilling. Also contributed services decreased to 10% (from 12% in 2021). So we need to find the 2022 annual report for Defenders of Wildlife.\n\nSearch for \"Defenders of Wildlife 2022 annual report National Petroleum Reserve Alaska\".",
290 + "tool_calls": [
291 + {
292 + "id": "fixture-tool-call-15-0",
293 + "type": "function",
294 + "function": {
295 + "name": "search",
296 + "arguments": "{\"query\": \"Defenders of Wildlife 2022 annual report National Petroleum Reserve Alaska\", \"topn\": 10, \"source\": \"web\"}"
297 + }
298 + }
299 + ]
300 + },
301 + {
302 + "role": "tool",
303 + "content": "[7] Defenders of Wildlife 2022 annual report National Petroleum Reserve Alaska (Search_Results/Defenders of Wildlife 2022 annual report National Petroleum Reserve Alaska)\n**viewing lines [0 - 37] of 37**\n\nL0: \nL1: URL: Search_Results/Defenders of Wildlife 2022 annual report National Petroleum \nL2: Reserve Alaska\nL3: # Search Results\nL4: \nL5: * 【0†CELEBRATING YEARS; publish_date: none†www.annualreports.com】 With less \nL6: than 340 right whales left in the wild, Defenders is fighting tirelessly to end \nL7: deadly lobster gear entanglements and vessel strikes that are driving ...\nL8: * 【1†Financials; publish_date: none†defenders.org】 We invite you to explore \nL9: the reports below to learn more about our activities and accomplishments, and \nL10: how we put your money to work for wildlife.\nL11: * 【2†Alaska Program Looks Back on 2022; publish_date: none†defenders.org】 Feb \nL12: 9, 2023 — Thanks to a lawsuit joined by Defenders, seven million acres were \nL13: returned to protection within the National Petroleum Reserve-Alaska (NPR-A), ...\nL14: * 【3†Defenders-of-Wildlife-2022-Financial-Statement. ...; publish_date: \nL15: none†defenders.org】 We have audited the accompanying consolidated financial \nL16: statements of Defenders of Wildlife and Affiliated Defenders of Wildlife Action \nL17: Fund (collectively, ...\nL18: * 【4†2022 Annual Report; publish_date: none†alaskaconservation.org】 Jun 13, \nL19: 2023 — In 2022, we focused on three landscapes: the Arctic. National Wildlife \nL20: Refuge, Bristol Bay, and the Tongass National Forest. In March 2022,.\nL21: * 【5†Assessment of ecological and cultural values within the ...; \nL22: publish_date: none†www.blm.gov】 This document was written to provide technical \nL23: information regarding the ecological importance of the National Petroleum \nL24: Reserve – Alaska (NPR-A). Several ...\nL25: * 【6†Accomplishments Report; publish_date: none†defenders.org】 National \nL26: Petroleum Reserve-Alaska. Identified for exceptional wildlife and cultural \nL27: values, including critical habitat for polar bears and other species ...\nL28: * 【7†2022 annual report; publish_date: none†dory-\nL29: plantain-s2zc.squarespace.com】 These projects are made possible through \nL30: collaborations with Defenders of. Wildlife, Cook Inletkeeper, Trustees for \nL31: Alaska,. Environmental Investigation Agency, ...\nL32: * 【8†23IMPACT REPORT; publish_date: none†www.annualreports.com】 Defenders of \nL33: Wildlife made incredible progress protecting wildlife and wild places in 2023, \nL34: helping shape a brighter future for imperiled species and vital ...\nL35: * 【9†From Leasing to Land Protections; publish_date: none†defenders-cci.org】 \nL36: Abstract. When you hear the name, National Petroleum Reserve - Alaska (NPR-A, or\nL37: Reserve), you might think of a barren land filled with oil wells and ...",
304 + "tool_call_id": "fixture-tool-call-15-0"
305 + },
306 + {
307 + "role": "assistant",
308 + "content": "",
309 + "reasoning_content": "Open result 0.",
310 + "tool_calls": [
311 + {
312 + "id": "fixture-tool-call-17-0",
313 + "type": "function",
314 + "function": {
315 + "name": "open",
316 + "arguments": "{\"cursor\": 7, \"id\": 0}"
317 + }
318 + }
319 + ]
320 + },
321 + {
322 + "role": "tool",
323 + "content": "[8] www.annualreports.com (https://www.annualreports.com/HostedData/AnnualReportArchive/d/defenders-wildlife_2022.pdf)\n**viewing lines [0 - 53] of 289**\n\nL0: \nL1: URL: https://www.annualreports.com/HostedData/AnnualReportArchive/d/defenders-\nL2: wildlife_2022.pdf\nL3: ===== Page 1 ===== 2022 Impact Report C E L E B R A T I N G Y E A R S ===== Page\nL4: 2 ===== 2022 Defenders of Wildlife made important progress for imperiled \nL5: species and vital landscapes across the United States in 2022. GRAY WOLF | JIM \nL6: GUMMERAL MAY STOCK PRIOR Lawyers successfully challenged the previous \nL7: administration’s decision to delist the gray wolf and restored critical federal \nL8: protections under the Endangered Species Act. This latest triumph in court is \nL9: part of our ongoing battle to protect and restore gray wolves throughout their \nL10: historical range and shield them from persecution by extremist legislators in \nL11: Idaho, Montana and Wyoming. TWO MORE FATALIZED GRAY SWALLETS TO SEA TO SHARE \nL12: ALLIANCE Lobbyists worked around the clock to expand funding for wildlife \nL13: conservation in the FY2022 federal spending bill, which included $31 million (a \nL14: 44% increase) for the Bureau of Land Management’s Threatened and Endangered \nL15: Species Program, $2.5 million (an 81% increase) for the U.S. Department of \nL16: Agriculture Wildlife Services’ Nonlethal Initiative to prevent human-wildlife \nL17: conflicts and $21 million (a 320% increase) for North Atlantic right whale \nL18: conservation. 2 DEFENDERS OF WILDLIFE ===== Page 3 ===== **Policy Experts** \nL19: played a crucial role in securing international trade protections for 100 \nL20: species of sharks and rays, all 158 species of glass frogs and 73 species of \nL21: reptiles, including 21 species of desert horned lizards, at the Convention on \nL22: International Trade in Endangered Species (CITES) in Panama. \\--- **Field \nL23: Teams** worked tirelessly to protect and restore imperiled species across the \nL24: country. From Florida manatees and red wolves in the Southeast to belugas and \nL25: grizzly bears in Alaska, Defenders’ conservation experts were on the ground \nL26: saving species that need our help to survive and thrive. \\--- **Conservation \nL27: Innovation Experts** published more than 10 peer-reviewed studies on topics that\nL28: include the Cook Inlet beluga whale, golden-cheeked warbler, global parrot \nL29: biodiversity, the Endangered Species Act, the effects of mountaintop removal \nL30: mining on endangered species, the ecological importance of panthers and the \nL31: implementation of “30x30” – the globally recognized goal to which President \nL32: Biden committed the U.S. to conserve 30% of our imperiled lands and waters by \nL33: 2030. \\--- **DEFENDERS.ORG** ===== Page 4 ===== THANK YOU Defenders celebrated \nL34: our 75th anniversary in 2022—an exciting milestone that we attribute to our \nL35: unwavering dedication to our wildlife conservation mission. From helping to pass\nL36: the Endangered Species Act and other bedrock conservation laws to leading the \nL37: advocacy on the reintroduction of wolves to Yellowstone National Park and other \nL38: parts of the West, Defenders is proud to inspire, lead and ultimately celebrate \nL39: some of the most consequential conservation wins for wildlife and wild places. \nL40: It is wonderfully fitting that we also mark 2022 as a year when we successfully \nL41: defended, yet again, the original focal species and symbol of our organization, \nL42: the gray wolf. While the future for gray wolves and many other species remains \nL43: uncertain, we will never stop fighting to protect the imperiled wildlife that \nL44: need us and to restore vibrant and balanced ecosystems. Our mission is more \nL45: important than ever due to the escalating biodiversity and climate crises. \nL46: Defenders is determined to build on our accomplishments and seize new \nL47: opportunities to address these challenges through our multifaceted approach to \nL48: conservation. We focus on developing innovative and science-based solutions, \nL49: coordinating on-the-ground species and habitat restoration, promoting human-\nL50: wildlife coexistence, protecting and strengthening our bedrock environmental \nL51: laws, investing in impactful collaborations and partnerships, building a more \nL52: diverse and powerful conservation movement and challenging attacks on wildlife \nL53: through advocacy, policy and litigation. Defenders’ continued success is made ",
324 + "tool_call_id": "fixture-tool-call-17-0"
325 + },
326 + {
327 + "role": "assistant",
328 + "content": "",
329 + "reasoning_content": "Search within PDF for \"National Petroleum Reserve\".",
330 + "tool_calls": [
331 + {
332 + "id": "fixture-tool-call-19-0",
333 + "type": "function",
334 + "function": {
335 + "name": "find",
336 + "arguments": "{\"cursor\": 8, \"pattern\": \"Petroleum\"}"
337 + }
338 + }
339 + ]
340 + },
341 + {
342 + "role": "tool",
343 + "content": "[9] Find results for text: `petroleum` in `www.annualreports.com` (https://www.annualreports.com/HostedData/AnnualReportArchive/d/defenders-wildlife_2022.pdf/find?pattern=petroleum)\n**viewing lines [0 - 4] of 4**\n\nL0: # 【0†match at L201】\nL1: the 22 million acres of the National Petroleum Reserve in Alaska to oil and gas \nL2: drilling and removed protections for important wildlife habitat. Protections \nL3: have now been restored for nearly half of the reserve’s pristine lands, which \nL4: are vital habitat for shorebirds, denning polar bears and tens of thousands of ",
344 + "tool_call_id": "fixture-tool-call-19-0"
345 + },
346 + {
347 + "role": "assistant",
348 + "content": "",
349 + "reasoning_content": "Open around L190.",
350 + "tool_calls": [
351 + {
352 + "id": "fixture-tool-call-21-0",
353 + "type": "function",
354 + "function": {
355 + "name": "open",
356 + "arguments": "{\"cursor\": 8, \"loc\": 180}"
357 + }
358 + }
359 + ]
360 + },
361 + {
362 + "role": "tool",
363 + "content": "[10] www.annualreports.com (https://www.annualreports.com/HostedData/AnnualReportArchive/d/defenders-wildlife_2022.pdf)\n**viewing lines [180 - 227] of 289**\n\nL180: the sixth successful transfer of bison to the Assiniboine and Sioux Tribes of \nL181: Fort Peck since 2019. \\--- **SWIFT FIX KITS | © ROSIMA PAELARINTSKIMMA MADDIAL \nL182: 200 AND CONSERVATION BIOLOGY INSTITUTE** \\--- **Celebrating our third year** of \nL183: a collaborative program with the Aaniih and Nakoda Tribes and others to restore \nL184: swift foxes to the Fort Belknap Indian Reservation in Montana, Defenders helped \nL185: with the release of 28 more swift foxes. With over 100 foxes reintroduced \nL186: through this program, monitoring efforts show that they are reproducing in the \nL187: wild—a critical measure of success for a self-sustaining population. \\--- \nL188: **Defenders continued to lead the way** for conserving and recovering the \nL189: endangered black-footed ferret, supporting the black-footed ferret survey for \nL190: the Fort Belknap Indian community. Thirty-six ferrets were vaccinated against \nL191: sylvatic plague and two dozen kits were released in the wild. \\--- **10 \nL192: DEFENDERS OF WILDLIFE** ===== Page 11 ===== Defenders helped to bring hope for \nL193: recovery for the endangered military macaw, adding 11 fledglings to a growing \nL194: wild population in Puerta Vallarta, Mexico, that is under pressure from habitat \nL195: loss and poachers for the illegal pet trade. Accord- ing to our recent report, \nL196: the 2008 parrot trade ban that Defenders fought to achieve is working. \nL197: Preventing more than 30,000 parrots from being illegally trapped each year, the \nL198: trade ban has resulted in a 47% decrease in the illegal trade of parrots and an \nL199: 88% decrease in U.S. seizures of Mexican parrots. As a result of a Defenders \nL200: lawsuit, BLM rescinded the previous administration’s plan that opened most of \nL201: the 22 million acres of the National Petroleum Reserve in Alaska to oil and gas \nL202: drilling and removed protections for important wildlife habitat. Protections \nL203: have now been restored for nearly half of the reserve’s pristine lands, which \nL204: are vital habitat for shorebirds, denning polar bears and tens of thousands of \nL205: caribou in the Teshekpuk caribou herd. Our new storymap illustrates why the \nL206: reserve is so important to wildlife. Defenders marked the 20th anniversary of \nL207: our Sea Otter Awareness Week. In celebration of this milestone, we sponsored \nL208: state legislation, signed by Gov. Gavin Newsom, that formally recognizes the \nL209: event’s anniversary and acknowledges the vital role that sea otters play in \nL210: California’s nearshore ecosystems. This annual celebration provides a unique \nL211: opportunity to raise awareness of the important role sea otters play in \nL212: maintaining ecosystem health, our need to coexist with sea otters and the \nL213: efforts of Defenders and others in advancing sea otter conservation. \nL214: DEFENDERS.ORG ===== Page 12 ===== FINANCIAL REPORT In 2022, Defenders of \nL215: Wildlife’s nearly 2.2 million members and supporters contributed more than $41 \nL216: million for wildlife and wildlife habitat. This figure includes all those who \nL217: generously donated their time and expertise. The steadfast support of our donors\nL218: allows Defenders to sustain our program and public education efforts in the \nL219: field, the courts and on Capitol Hill. 2022 SOURCES OF FUNDS Grants and \nL220: contributions $28,946 Bequests, trusts and split interests $13,928 Income from \nL221: investments, annuity reserve funds and trusts (-$7,495) Royalties and other \nL222: income $744 Contributed services $5,047 Total Revenue $41,170 USES OF FUNDS \nL223: Biodiversity conservation $18,970 Constituency mobilization $16,528 Fundraising \nL224: $1,029 Management and general $6,354 Total Expenses $42,881 Change in net assets\nL225: (-$1,711) Net assets, start of year $45,144 Net Assets, End of the Year $43,433\nL226: Dollars are in thousands. 12 DEFENDERS OF WILDLIFE Grants and contributions 59%\nL227: Contributed services 10% Royalties and other income 2% Bequests, trusts and ",
364 + "tool_call_id": "fixture-tool-call-21-0"
365 + },
366 + {
367 + "role": "assistant",
368 + "content": "",
369 + "reasoning_content": "Scrolling a bit more for \"Contributed services 10%\". Already seen. So indeed the 2022 annual report shows contributed services 10%, down from 12% in 2021.\n\nThus the organization is Defenders of Wildlife.\n\nThe question: \"Who was the emeritus board member that year? Name only, no titles.\"\n\nWe need to identify the emeritus board member listed in the 2022 annual report. From the 2021 annual report, we saw \"BOARD EMERITUS Terry C. Pelster, Esq.\" However, that is for 2021. The question asks \"Who was the emeritus board member that year?\" That year refers to the following year (2022). So we need to check the 2022 annual report for the emeritus board member.\n\nOpen the 2022 PDF near the board list. Search for \"Emeritus\".",
370 + "tool_calls": [
371 + {
372 + "id": "fixture-tool-call-23-0",
373 + "type": "function",
374 + "function": {
375 + "name": "find",
376 + "arguments": "{\"cursor\": 8, \"pattern\": \"Emeritus\"}"
377 + }
378 + }
379 + ]
380 + },
381 + {
382 + "role": "tool",
383 + "content": "[11] Find results for text: `emeritus` in `www.annualreports.com` (https://www.annualreports.com/HostedData/AnnualReportArchive/d/defenders-wildlife_2022.pdf/find?pattern=emeritus)\n**viewing lines [0 - 10] of 10**\n\nL0: # 【0†match at L237】\nL1: Sisk, Ph.D British Columbia, Canada BOARD EMERITUS Terry C. Pelster, Esq. New \nL2: York *Executive Committee member* STAFF & ADVISORS SCIENCE ADVISORS Gabriela \nL3: Chavarria, Ph.D Burke Museum, University of Washington Leah Gerber, Ph.D Arizona\nL4: State University Kathayoon Khalil, Ph.D Oregon Zoo Rurik List, Ph.D Universidad\nL5: \nL6: # 【1†match at L243】\nL7: Davis Thomas D. Sisk, Ph.D Emeritus Professor Northern Arizona University David \nL8: S. Wilcox, Ph.D Princeton University Rae Wynn-Grant, Ph.D National Geographic \nL9: Society SENIOR STAFF Jamie Rappaport Clark President & CEO Thu Pham Chief of \nL10: Staff James Stofan Senior Vice President, Operations McCrystle Adams Vice ",
384 + "tool_call_id": "fixture-tool-call-23-0"
385 + },
386 + {
387 + "role": "assistant",
388 + "content": "",
389 + "reasoning_content": "Open around L230.",
390 + "tool_calls": [
391 + {
392 + "id": "fixture-tool-call-25-0",
393 + "type": "function",
394 + "function": {
395 + "name": "open",
396 + "arguments": "{\"cursor\": 8, \"loc\": 220}"
397 + }
398 + }
399 + ]
400 + },
401 + {
402 + "role": "tool",
403 + "content": "[12] www.annualreports.com (https://www.annualreports.com/HostedData/AnnualReportArchive/d/defenders-wildlife_2022.pdf)\n**viewing lines [220 - 267] of 289**\n\nL220: contributions $28,946 Bequests, trusts and split interests $13,928 Income from \nL221: investments, annuity reserve funds and trusts (-$7,495) Royalties and other \nL222: income $744 Contributed services $5,047 Total Revenue $41,170 USES OF FUNDS \nL223: Biodiversity conservation $18,970 Constituency mobilization $16,528 Fundraising \nL224: $1,029 Management and general $6,354 Total Expenses $42,881 Change in net assets\nL225: (-$1,711) Net assets, start of year $45,144 Net Assets, End of the Year $43,433\nL226: Dollars are in thousands. 12 DEFENDERS OF WILDLIFE Grants and contributions 59%\nL227: Contributed services 10% Royalties and other income 2% Bequests, trusts and \nL228: split interests 29% TOTAL REVENUE Income from investments -15% Program and \nL229: support services 98% USE OF FUNDS Fundraising 2% ===== Page 13 ===== \nL230: DEFENDERS.ORG 13 ===== Page 14 ===== BOARD OF DIRECTORS John Dayton, Esq.* Chair\nL231: Texas Sheryl Sachs* Vice Chair Connecticut Dinah Bear, Esq.* Secretary Arizona \nL232: Mark Caylor* Treasurer District of Columbia Susan Alexander Connecticut Sylvia \nL233: Baca* California Whitefield Bailey Tennessee Donald Barry, Esq. Oregon Kim \nL234: O’Keefe Beck, Ph.D California Gabriela Chavarria, Ph.D Washington Jamie \nL235: Rappaport Clark* President & CEO Virginia John Donovan, Esq.* Montana and \nL236: California Caroline D. Gabel Maryland Terry L. Root, Ph.D New Mexico Thomas D. \nL237: Sisk, Ph.D British Columbia, Canada BOARD EMERITUS Terry C. Pelster, Esq. New \nL238: York *Executive Committee member* STAFF & ADVISORS SCIENCE ADVISORS Gabriela \nL239: Chavarria, Ph.D Burke Museum, University of Washington Leah Gerber, Ph.D Arizona\nL240: State University Kathayoon Khalil, Ph.D Oregon Zoo Rurik List, Ph.D Universidad\nL241: Autónoma Metropolitana-Lerma Barry Noon, Ph.D Colorado State University Terry \nL242: L. Root, Ph.D Stanford University Mark Schwartz, Ph.D University of California-\nL243: Davis Thomas D. Sisk, Ph.D Emeritus Professor Northern Arizona University David \nL244: S. Wilcox, Ph.D Princeton University Rae Wynn-Grant, Ph.D National Geographic \nL245: Society SENIOR STAFF Jamie Rappaport Clark President & CEO Thu Pham Chief of \nL246: Staff James Stofan Senior Vice President, Operations McCrystle Adams Vice \nL247: President, Conservation Law Matt Braughler Vice President, Philanthropy Rachel \nL248: Brittin Vice President, Communications Shawn Cantrell Vice President, Field \nL249: Conservation Programs Robert L. Dewey Vice President, Government Relations & \nL250: External Affairs Monica Goldberg Vice President, Landscape Conservation Sharon \nL251: Morris-Smith Vice President of Human Resources Elizabeth Nichols Vice President,\nL252: Finance/CFO Anthony Pegues Vice President, Diversity, Equity, Inclusion, \nL253: Justice Lindsay Rosa, Ph.D Vice President, Conservation Research and Innovation \nL254: Brandace Swimley Vice President, Integrated Marketing 14 DEFENDERS OF WILDLIFE \nL255: ===== Page 15 ===== There are many ways to support Defenders of Wildlife and \nL256: protect our imperiled wildlife. To learn more or donate, contact us at \nL257: 1-800-385-9712, visit Defenders.org/donate or choose one of the following \nL258: options. ➤ Transfer appreciated assets directly to Defenders and maximize \nL259: potential tax benefits. Defenders.org/ donate-stocks-help-save-wildlife ➤ Become\nL260: a Leadership Council member with a single donation of $1,000 or more and enjoy \nL261: more personal involvement with Defenders. Defenders.org/leadershipcouncil ➤ If \nL262: you are 70.5 or older, roll over a portion of your IRA to benefit wildlife. \nL263: Remember, these gifts must come directly from your IRA account. \nL264: Defenders.org/planned-giving/gifts-your-retirement-plan ➤ Get started on a \nL265: contribution from your donor-advised fund Defenders.org/donor-advised-funds ➤ \nL266: Check with your personnel office or that of your spouse/partner about your \nL267: company’s Corporate Matching Gifts program or visit us at ",
404 + "tool_call_id": "fixture-tool-call-25-0"
405 + },
406 + {
407 + "role": "assistant",
408 + "content": "",
409 + "reasoning_content": "Thus the emeritus board member is Terry C. Pelster. However, the question might ask for \"Who was the emeritus board member that year?\" Possibly referring to the year of the annual report that mentioned the rescinded plan and contributed services decrease. That is 2022. So the answer is Terry C. Pelster.\n\nBut let's verify that the 2022 annual report indeed includes the mention of contributed services decreasing to 10%. We saw that. So answer is Terry C. Pelster.\n\nBut maybe the question is tricky: \"Who was the emeritus board member that year?\" Could be the emeritus board member listed in the 2021 annual report? However, the question says \"That same year, a company's annual report revealed that its nearly 2.2 million members and supporters contributed over $49 million, with 12% of total revenue coming from contributed services. The following year, the annual report stated that the organization had played a role in rescinding a previous administration’s plan to open most of the 22 million acres of the National Petroleum Reserve in Alaska to oil and gas drilling and that contributed services had decreased by 2% to 10%.\"\n\nThus the \"following year\" annual report is 2022. The question likely asks: \"Who was the emeritus board member that year?\" Meaning the emeritus board member listed in that annual report (2022). So answer: Terry C. Pelster.\n\nBut let's confirm that the 2021 annual report also shows emeritus board member Terry C. Pelster. Indeed we saw that. But perhaps the 2022 report might have a different emeritus board member. We saw it's the same. So answer is Terry C. Pelster.\n\nHowever, the question might be about a different organization. Let's verify the details.\n\nFirst clue: Press release about Kenya vaccine doses on Dec 5, 2021. That is likely a World Bank press release. Search for \"World Bank Kenya vaccine doses December 2021 press release\".",
410 + "tool_calls": [
411 + {
412 + "id": "fixture-tool-call-27-0",
413 + "type": "function",
414 + "function": {
415 + "name": "search",
416 + "arguments": "{\"query\": \"As of December 5, 2021, Kenya had received a total of 16,201,670 vaccines\", \"topn\": 10, \"source\": \"web\"}"
417 + }
418 + }
419 + ]
420 + },
421 + {
422 + "role": "tool",
423 + "content": "[13] As of December 5, 2021, Kenya had received a total of 16,201,670 vaccines (Search_Results/As of December 5, 2021, Kenya had received a total of 16,201,670 vaccines)\n**viewing lines [0 - 41] of 41**\n\nL0: \nL1: URL: Search_Results/As of December 5, 2021, Kenya had received a total of \nL2: 16,201,670 vaccines\nL3: # Search Results\nL4: \nL5: * 【0†Kenya's Economy is Showing Resilience as Output Rises ...; publish_date: \nL6: none†www.worldbank.org】 Dec 14, 2021 — As of December 5, 2021, Kenya had \nL7: received a total of 16,201,670 vaccines, with 7,583,134 administered. While \nL8: vaccine acceptance is ...\nL9: * 【1†Unmet need for COVID-19 vaccination coverage in Kenya - PMC; \nL10: publish_date: none†pmc.ncbi.nlm.nih.gov】 by SK Muchiri · 2022 · Cited by 42 — As\nL11: of December 2021, six counties had a vaccination coverage of less than 5%. \nL12: These counties include Garissa, Mandera, Marsabit, Tana River, Turkana, and ...\nL13: * 【2†MINISTRY OF HEALTH; publish_date: none†covidhub.mediacouncil.or.ke】 Dec \nL14: 1, 2021 — • Total Covid-19 Vaccines Received to date- 16,201,670 ... Table 10: \nL15: Vaccine Logistics Received in the Country as at 5th, December 2021.\nL16: * 【3†COVID-19 vaccination refusal trends in Kenya over 2021 - PMC; \nL17: publish_date: none†pmc.ncbi.nlm.nih.gov】 by RT Rego · 2023 · Cited by 21 — We \nL18: assessed vaccine refusal over time in Kenya, and characterized factors \nL19: associated with changes in vaccination refusal.\nL20: * 【4†Ciheb-Kenya on the Front Lines of the COVID-19 Vaccine ...; publish_date:\nL21: none†www.ciheb.org】 Of the 721,509 persons vaccinated, 409,711 are aged 58 \nL22: years and above, 143,050 are healthcare workers, 109,593 are teachers, and \nL23: 59,155 are security officers.\nL24: * 【5†proportion of fully vaccinated adults 23.3pc.; publish_date: \nL25: none†www.facebook.com】 COVID-19: KENYA administers 150632 new vaccinations; \nL26: total jabs given now 13782963; proportion of fully vaccinated adults 23.3pc.\nL27: * 【6†National COVID-19 Vaccines Deployment and ...; publish_date: \nL28: none†www.parliament.go.ke】 Kenya has a well-established immunization program and\nL29: a robust disease surveillance and response unit (DSRU) which serves as an early\nL30: warning system, to.\nL31: * 【7†Kenya: Vaccine Requirements Violate Rights; publish_date: \nL32: none†www.hrw.org】 Dec 13, 2021 — Kenya, with an estimated adult population of \nL33: 27.2 million and a total population of 55 million, has received approximately 23\nL34: million doses as ...\nL35: * 【8†Government gives 20 million a month to get vaccinated; publish_date: \nL36: none†www.bbc.com】 Nov 22, 2021 — Although less than 10% of Kenyans are \nL37: vaccinated, the government wants to avoid a surge over Christmas.\nL38: * 【9†ministry of health kenya covid-19 immunization status ...; publish_date: \nL39: none†www.health.go.ke】 May 15, 2022 — By the end of 2021 Kenya had administered \nL40: 10.1 million doses of 5 types of Covid-19 vaccines. During this period 4.2 \nL41: million people were ...",
424 + "tool_call_id": "fixture-tool-call-27-0"
425 + },
426 + {
427 + "role": "assistant",
428 + "content": "",
429 + "reasoning_content": "Open result 0.",
430 + "tool_calls": [
431 + {
432 + "id": "fixture-tool-call-29-0",
433 + "type": "function",
434 + "function": {
435 + "name": "open",
436 + "arguments": "{\"cursor\": 13, \"id\": 0}"
437 + }
438 + }
439 + ]
440 + },
441 + {
442 + "role": "tool",
443 + "content": "[14] www.worldbank.org (https://www.worldbank.org/en/news/press-release/2021/12/14/kenya-s-economy-is-showing-resilience-as-output-rises-above-pre-pandemic-levels-driven-by-a-rebound-in-the-services-sect)\n**viewing lines [0 - 52] of 61**\n\nL0: \nL1: URL: https://www.worldbank.org/en/news/press-release/2021/12/14/kenya-s-economy-\nL2: is-showing-resilience-as-output-rises-above-pre-pandemic-levels-driven-by-a-\nL3: rebound-in-the-services-sect\nL4: NAIROBI, December 14, 2021 — Kenya’s economy has demonstrated resilience to the \nL5: COVID-19 shock, with output in the first half of the year rising above pre-\nL6: pandemic levels. In 2021 as a whole, gross domestic product (GDP) is expected to\nL7: grow by 5%, one of the faster recoveries among Sub-Saharan African countries. \nL8: Overall economic performance is expected to be robust at 4.9% per year in \nL9: 2022-23, similar to the pre-pandemic pace (5% average annual growth from 2010 to\nL10: 2019). According to the 24th edition of the Kenya Economic Update, “From \nL11: Recovery to Better Jobs,” growth has been supported by rebounds in industry and,\nL12: especially, services. Agricultural output, however, fell by 0.5% year on year \nL13: in the first half of 2021 following a particularly strong performance in 2020, \nL14: partly due to below-average rains. Demand-side recovery has been supported by a \nL15: revival in private consumption, against a backdrop of improving employment \nL16: conditions and household incomes. “Kenya’s economy has shown considerable \nL17: resilience to the enormous shock of the pandemic, and this year is expected to \nL18: post one of the stronger growth rebounds in the region thanks to diversified \nL19: sources of growth and sound economic policies and management,” said Keith \nL20: Hansen, World Bank Country Director for Kenya. “However, poverty has increased, \nL21: and the buffers and coping mechanisms of households, firms, and the public \nL22: finances have been depleted.” Economic activity in Kenya has continued to adapt \nL23: to the pandemic and associated restrictions. A mix of containment measures, such\nL24: as a nightly curfew, were in effect through most of 2021, while more \nL25: economically disruptive measures such as lockdowns and travel restrictions were \nL26: phased, limiting the impact on economic activities. The vaccine rollout, which \nL27: had a slow start due to supply constraints, has picked up as new shipments of \nL28: vaccines have arrived, particularly since September. This has supported economic\nL29: recovery and growth through the third quarter of 2021. As of December 5, 2021, \nL30: Kenya had received a total of 16,201,670 vaccines, with 7,583,134 administered. \nL31: While vaccine acceptance is reportedly high there is still a long way to go \nL32: towards the government’s target of fully inoculating the adult population of \nL33: about 30 million by the end of 2022.As of December 6, 2021, about 10% of adults \nL34: (2.9 million people) had been fully vaccinated while another 16% (4.9 million \nL35: people) had received their first dose. The report also projects robust growth in\nL36: the medium-term. This outlook takes into account that some sub-sectors, such as\nL37: education, have bounced back strongly, but others such as international \nL38: tourism, have only partially and face a much more protracted recovery. The \nL39: anticipated further recovery of hotels and restaurants, trade, transport, and \nL40: other services, depends on substantial vaccination progress to help prevent new \nL41: waves of infections and associated containment measures. “Recent economic \nL42: performance has been strong and the outlook is positive, but in Kenya, as \nL43: everywhere, the future course of the pandemic remains a major source of \nL44: uncertainty, said Alex Sienaert, Senior Economist for Kenya. “To support further\nL45: recovery, reduce debt distress risks, and rebuild space for social and \nL46: development spending, it is critical that Kenya continues to implement medium-\nL47: term fiscal consolidation plans.” In addition to pandemic-related risks to the \nL48: outlook, a second key domestic risk factor stems from the drought conditions \nL49: which are affecting parts of the country and already causing severe hardship. \nL50: Should the drought intensify or spread, this would weigh on the near-term \nL51: economic outlook. Weaker global growth, higher-than-anticipated energy prices, \nL52: and tighter external financing conditions are the primary external risks. As ",
444 + "tool_call_id": "fixture-tool-call-29-0"
445 + },
446 + {
447 + "role": "assistant",
448 + "content": "",
449 + "reasoning_content": "Thus first clue confirmed.\n\nSecond clue: Defenders of Wildlife 2021 annual report. Confirmed.\n\nThird clue: Defenders of Wildlife 2022 annual report. Confirmed.\n\nThus the answer is Terry C. Pelster.\n\nBut let's verify that the 2022 annual report indeed shows \"Contributed services 10%\". We saw that. Also confirm that 2021 report shows 12%. We saw that.\n\nThus the question: \"Who was the emeritus board member that year?\" Refers to 2022. Answer: Terry C. Pelster.\n\nHowever, the question might be asking for the emeritus board member in the 2021 report? Let's read the question carefully:\n\n\"In a year before 2023, between December 1 and 15 (inclusive), a press release reported that as of the 5th of that month, a country had received a total of 16,201,670 vaccine doses, with 7,583,134 administered. That same year, a company's annual report revealed that its nearly 2.2 million members and supporters contributed over $49 million, with 12% of total revenue coming from contributed services. The following year, the annual report stated that the organization had played a role in rescinding a previous administration’s plan to open most of the 22 million acres of the National Petroleum Reserve in Alaska to oil and gas drilling and that contributed services had decreased by 2% to 10%. Who was the emeritus board member that year? Name only, no titles.\"\n\nThe phrase \"that year\" could refer to \"the following year\" i.e., the year of the annual report that mentioned the rescinded plan. Or could refer to \"that same year\" (the year of the press release and the first annual report). However, the phrase \"that year\" appears at the end: \"Who was the emeritus board member that year?\" It likely refers to the year of the annual report that is being discussed in the preceding sentence. The preceding sentence talks about the following year's annual report. So indeed \"that year\" is the following year (2022).\n\nThus answer: Terry C. Pelster.\n\nBut to be thorough, we should check the 2022 annual report for any other emeritus board members. The search only showed \"BOARD EMERITUS Terry C. Pelster, Esq.\" Maybe there are multiple emeritus board members. But likely only one.\n\nCheck the 2021 annual report for emeritus board member. It also shows Terry C. Pelster. So maybe the question is asking for the emeritus board member in the 2021 report. But why would they ask that? The question is about connecting the clues to identify the organization and then find the emeritus board member. The clues point to Defenders of Wildlife. The question likely asks for the emeritus board member as of the 2022 report.\n\nNevertheless, we should confirm that the 2022 report is indeed the one that mentions the rescinded plan. We saw that. Also confirm that contributed services decreased from 12% to 10%. We saw that.\n\nThus answer is Terry C. Pelster.\n\nBut let's search for \"Defenders of Wildlife emeritus board member 2022 Terry Pelster\".",
450 + "tool_calls": [
451 + {
452 + "id": "fixture-tool-call-31-0",
453 + "type": "function",
454 + "function": {
455 + "name": "search",
456 + "arguments": "{\"query\": \"Terry C. Pelster Defenders of Wildlife emeritus board member\", \"topn\": 10, \"source\": \"web\"}"
457 + }
458 + }
459 + ]
460 + },
461 + {
462 + "role": "tool",
463 + "content": "[15] Terry C. Pelster Defenders of Wildlife emeritus board member (Search_Results/Terry C. Pelster Defenders of Wildlife emeritus board member)\n**viewing lines [0 - 31] of 31**\n\nL0: \nL1: URL: Search_Results/Terry C. Pelster Defenders of Wildlife emeritus board member\nL2: # Search Results\nL3: \nL4: * 【0†2024 impact report; publish_date: none†defenders.org】 Terry Root, Ph.D. \nL5: New Mexico. Thomas D. Sisk, Ph.D. British Columbia, Canada. BOARD EMERITUS. \nL6: Terry C. Pelster, Esq. New York. SCIENCE ADVISORS. Leah Gerber, Ph ...\nL7: * 【1†Defenders of Wildlife; publish_date: none†www.wrongkindofgreen.org】 \nL8: Directors ; Victor M. Sher*, Chair – California ; Terry C. Pelster*, Vice Chair \nL9: – New York ; Richard Kopcho*, Treasurer – California ; Adelaide P. Gomer*, \nL10: Secretary ...\nL11: * 【2†Exhibit 12; publish_date: none†www.uschamber.com】 DECLARATION OF TERRY \nL12: PELSTER. I, Terry C. Pelster, declare as follows: 1. I am a current member of \nL13: Defenders of Wildlife (“Defenders”) and have been a member.\nL14: * 【3†2020 ANNUAL REPORT; publish_date: none†defenders.org】 Terry L. Root. \nL15: Florida. BOARD. EMERITUS. Terry C. Pelster. New York. Alan Steinberg. Florida. \nL16: SCIENCE. ADVISORS. Gabriela Chavarria,. Ph.D. Denver Museum of.\nL17: * 【4†2019 Annual Report; publish_date: none†defenders.org】 BOARD EMERITUS. \nL18: Terry C. Pelster. New York. Alan Steinberg. Florida. SCIENCE ADVISORS. Leah \nL19: Gerber, Ph.D. Arizona State University. Kathayoon Khalil, Ph.D.\nL20: * 【5†annual report; publish_date: none†defenders.org】 * Executive Committee \nL21: member. BoARD EMERITuS. Terry C. Pelster. New York. Alan Steinberg. Florida. \nL22: PRESIDENT. EMERITuS. Rodger Schlickeisen. Virginia. SCIENCE ...\nL23: * 【6†BEING THERE FOR BELUGAS; publish_date: none†wendeenicole.com】 Mar 2, 2023\nL24: — Defenders of Wildlife is a national, nonprofit member- ship ... Terry C. \nL25: Pelster, Esq. New York. *Executive Committee member outfront.\nL26: * 【7†Annual Reports; publish_date: none†www.annualreports.com】 WAYS TO GIVE \nL27: There are many ways to support Defenders of Wildlife and protect our imperiled \nL28: wildlife. ... Terry C. Pelster. New York. Alan Steinberg. Florida.\nL29: * 【8†2014 Annual Report; publish_date: none†defenders.org】 BOARD OF DIRECTORS \nL30: defenders of wildlife is a national, nonprofit ... BOARD EMERITUS. Terry C. \nL31: Pelster. New York. Alan Steinberg. Florida. Page 3. PO.",
464 + "tool_call_id": "fixture-tool-call-31-0"
465 + },
466 + {
467 + "role": "assistant",
468 + "content": "",
469 + "reasoning_content": "Thus Terry C. Pelster is indeed the emeritus board member.\n\nNow the question might want just \"Terry C. Pelster\". Or maybe \"Terry Pelster\". Or \"Terry C. Pelster, Esq.\" but instructions say \"Name only, no titles.\" So likely \"Terry C. Pelster\". However, maybe they want \"Terry Pelster\". Better to include full name as listed: \"Terry C. Pelster\".\n\nLet's read the question again: \"Who was the emeritus board member that year? Name only, no titles.\" So we should answer just the name, no \"Esq.\" or \"Board Emeritus\". So \"Terry C. Pelster\". Or maybe \"Terry Pelster\". But the document shows \"Terry C. Pelster, Esq.\" So the name is \"Terry C. Pelster\". \n\nBetter to verify that the 2022 annual report indeed lists \"BOARD EMERITUS Terry C. Pelster, Esq.\" We saw that line.\n\nThus answer: Terry C. Pelster.\n\nBut let's double-check that the question might refer to the 2021 annual report. The phrase \"that year\" could be ambiguous. Let's parse the timeline:\n\n- Year 1: 2021. Press release in Dec 2021 about Kenya vaccine doses. Same year (2021) annual report of Defenders of Wildlife shows 12% contributed services.\n\n- Year 2: 2022. Annual report mentions rescinded plan and contributed services decreased to 10%.\n\nThe question says \"The following year, the annual report stated that the organization had played a role in rescinding a previous administration’s plan ... and that contributed services had decreased by 2% to 10%.\" This is year 2.\n\nThen asks \"Who was the emeritus board member that year?\" That likely refers to year 2.\n\nThus answer: Terry C. Pelster.\n\nHowever, perhaps the emeritus board member changed between 2021 and 2022. We need to check 2022 annual report for any other emeritus board member. We saw only Terry C. Pelster. Also check 2021 annual report: same. So likely unchanged.\n\nThus answer is Terry C. Pelster.\n\nNevertheless, we should confirm that the 2022 annual report is indeed the one that includes the phrase \"BLM rescinded the previous administration’s plan that opened most of the 22 million acres of the National Petroleum Reserve in Alaska to oil and gas drilling\". We saw that in line 200-202.\n\nThus answer is Terry C. Pelster.\n\nNow we need to provide the answer in the final message. Only the name.\n\nThus final answer: Terry C. Pelster.\n\nBut let's be thorough and search for \"Defenders of Wildlife 2022 annual report contributed services 10%\".",
470 + "tool_calls": [
471 + {
472 + "id": "fixture-tool-call-33-0",
473 + "type": "function",
474 + "function": {
475 + "name": "find",
476 + "arguments": "{\"cursor\": 8, \"pattern\": \"Contributed services\"}"
477 + }
478 + }
479 + ]
480 + },
481 + {
482 + "role": "tool",
483 + "content": "[16] Find results for text: `contributed services` in `www.annualreports.com` (https://www.annualreports.com/HostedData/AnnualReportArchive/d/defenders-wildlife_2022.pdf/find?pattern=contributed services)\n**viewing lines [0 - 10] of 10**\n\nL0: # 【0†match at L222】\nL1: income $744 Contributed services $5,047 Total Revenue $41,170 USES OF FUNDS \nL2: Biodiversity conservation $18,970 Constituency mobilization $16,528 Fundraising \nL3: $1,029 Management and general $6,354 Total Expenses $42,881 Change in net assets\nL4: (-$1,711) Net assets, start of year $45,144 Net Assets, End of the Year $43,433\nL5: \nL6: # 【1†match at L227】\nL7: Contributed services 10% Royalties and other income 2% Bequests, trusts and \nL8: split interests 29% TOTAL REVENUE Income from investments -15% Program and \nL9: support services 98% USE OF FUNDS Fundraising 2% ===== Page 13 ===== \nL10: DEFENDERS.ORG 13 ===== Page 14 ===== BOARD OF DIRECTORS John Dayton, Esq.* Chair",
484 + "tool_call_id": "fixture-tool-call-33-0"
485 + },
486 + {
487 + "role": "assistant",
488 + "content": "",
489 + "reasoning_content": "Thus confirmed.\n\nNow we should verify that the 2021 annual report shows contributed services 12%.",
490 + "tool_calls": [
491 + {
492 + "id": "fixture-tool-call-35-0",
493 + "type": "function",
494 + "function": {
495 + "name": "find",
496 + "arguments": "{\"cursor\": 4, \"pattern\": \"Contributed services\"}"
497 + }
498 + }
499 + ]
500 + },
501 + {
502 + "role": "tool",
503 + "content": "[17] Find results for text: `contributed services` in `www.annualreports.com` (https://www.annualreports.com/HostedData/AnnualReportArchive/d/defenders-wildlife_2021.pdf/find?pattern=contributed services)\n**viewing lines [0 - 4] of 4**\n\nL0: # 【0†match at L194】\nL1: $3,354 Royalties and other income $3,576 Contributed services $6,140 Total \nL2: Revenue $49,819 USES OF FUNDS Biodiversity conservation $22,420 Constituency \nL3: mobilization $16,324 Fundraising $1,211 Management and general $5,865 Total \nL4: Expenses $45,820 Change in net assets $3,999 Net assets, start of year $41,145 ",
504 + "tool_call_id": "fixture-tool-call-35-0"
505 + },
506 + {
507 + "role": "assistant",
508 + "content": "",
509 + "reasoning_content": "Open around L190-200 to see percentage.",
510 + "tool_calls": [
511 + {
512 + "id": "fixture-tool-call-37-0",
513 + "type": "function",
514 + "function": {
515 + "name": "open",
516 + "arguments": "{\"cursor\": 4, \"loc\": 185}"
517 + }
518 + }
519 + ]
520 + },
521 + {
522 + "role": "tool",
523 + "content": "[18] www.annualreports.com (https://www.annualreports.com/HostedData/AnnualReportArchive/d/defenders-wildlife_2021.pdf)\n**viewing lines [185 - 232] of 261**\n\nL185: devastating water pollution that is killing the seagrass and causing manatees to\nL186: starve. DAVID TES | SAM FRENZY DRAW DEFENDERS.ORG 11 ===== Page 12 ===== In \nL187: 2021, Defenders of Wildlife’s nearly 2.2 million members and supporters \nL188: contributed more than $49 million for wildlife and wildlife habitat. This figure\nL189: includes all those who generously donated their time and expertise. The \nL190: steadfast support of our donors allows Defenders to sustain our program and \nL191: public education efforts in the field, the courts and on Capitol Hill. 2021 \nL192: SOURCES OF FUNDS Grants and contributions $29,057 Bequests, trusts and split \nL193: interests $7,692 Income from investments, annuity reserve funds and trusts \nL194: $3,354 Royalties and other income $3,576 Contributed services $6,140 Total \nL195: Revenue $49,819 USES OF FUNDS Biodiversity conservation $22,420 Constituency \nL196: mobilization $16,324 Fundraising $1,211 Management and general $5,865 Total \nL197: Expenses $45,820 Change in net assets $3,999 Net assets, start of year $41,145 \nL198: Net Assets, End of the Year $45,144 Dollars are in thousands. 12 DEFENDERS OF \nL199: WILDLIFE Grants and contributions 58% Income from investments 7% Requests, \nL200: trusts and split interests 15% Royalties and other income 7% Contributed \nL201: services 12% Program and support services 97% Fundraising 3% ===== Page 13 =====\nL202: DEFENDERS.ORG 13 ===== Page 14 ===== John Dayton, Esq.* Chair Texas Sheryl \nL203: Sachs* Vice Chair Connecticut Whitfield Bailey* Secretary Tennessee Mark Caylor*\nL204: Treasurer District of Columbia Sylvia Baca* California Donald Barry, Esq. \nL205: Oregon Dinah Bear, Esq. Arizona Kim O’Keefe Beck, Ph.D California Gabriela \nL206: Chavarria, Ph.D Washington Jamie Rappaport Clark* President & CEO Virginia \nL207: Caroline D. Gabel Maryland Keith M. Harper, Esq. Maryland John Donovan, Esq. \nL208: Montana and California James M. Hecker, Esq. Virginia Judith Ponsikoff, Ph.D* \nL209: California Terry L. Root, Ph.D Florida BOARD EMERITUS Terry C. Pelster, Esq. New\nL210: York *Executive Committee member* ST A F F & A D V I S O R S SCIENCE ADVISORS \nL211: Gabriela Chavarria, Ph.D Burke Museum, University of Washington Leah Gerber, \nL212: Ph.D Arizona State University Kathayoon Khalil, Ph.D Oregon Zoo Rurik List, Ph.D\nL213: Universidad Autónoma Metropolitana-Lerma Barry Noon, Ph.D Colorado State \nL214: University Terry L. Root, Ph.D Stanford University Mark Schwartz, Ph.D \nL215: University of California-Davis David S. Wilcox, Ph.D Princeton University Rae \nL216: Wynn-Grant, Ph.D National Geographic Society SENIOR STAFF Jamie Rappaport Clark \nL217: President & CEO Thu Pham Chief of Staff Renee Stone Senior Vice President, \nL218: Conservation Programs & General Counsel James Stefan Senior Vice President, \nL219: Operations Matt Braughler Vice President, Philanthropy Rachel Brittin Vice \nL220: President, Communications Shawn Cantrell Vice President, Field Conservation \nL221: Programs Robert L. Dewey Vice President, Government Relations & External Affairs\nL222: Monica Goldberg Vice President, Landscape Conservation Sharon Morris-Smith Vice\nL223: President of Human Resources Elizabeth Nichols Vice President, Finance/CFO \nL224: Anthony Pegues Vice President, Diversity, Equity, Inclusion, Justice Lindsay \nL225: Rosa, Ph.D Director, Center for Conservation Innovation Michael P. Senatore Vice\nL226: President, Conservation Law Brandace Swimley Vice President, Integrated \nL227: Marketing 14 DEFENDERS OF WILDLIFE ===== Page 15 ===== THERE ARE MANY WAYS TO \nL228: SUPPORT DEFENDERS OF WILDLIFE and protect our imperiled wildlife. To learn more \nL229: or donate, contact us at 1-800-385-9712, visit defenders.org/donate or choose \nL230: one of the following options. ➤ Transfer appreciated assets directly to \nL231: Defenders and maximize potential tax benefits. Defenders.org/donate-stocks-help-\nL232: save-wildlife ➤ Become a Leadership Council member with a single donation of ",
524 + "tool_call_id": "fixture-tool-call-37-0"
525 + },
526 + {
527 + "role": "assistant",
528 + "content": "Terry C. Pelster",
529 + "reasoning_content": "Thus everything matches.\n\nThus answer: Terry C. Pelster.\n\nNow we need to output only the name.\n\nThus final answer: Terry C. Pelster.\n\nBut perhaps the question expects \"Terry Pelster\". However, the document includes middle initial \"C.\" It might be better to include the full name as listed: \"Terry C. Pelster\". The instruction \"Name only, no titles\" suggests we should not include \"Esq.\" or \"Board Emeritus\". But including middle initial is likely okay.\n\nTo be safe, we could output \"Terry C. Pelster\". Let's see similar puzzles: often they expect the full name as appears in the document.\n\nThus answer: Terry C. Pelster.\n\nNow we will provide final answer.",
530 + "tool_calls": []
531 + }
532 + ]
533 +}
rust/src/chat/src/renderer/deepseek_v32/fixtures/test_output_search_w_date.txt new
+2455
@@ -0,0 +1,2455 @@
1 +<｜begin▁of▁sentence｜><｜User｜>
2 +
3 +## Tools
4 +
5 +You have access to a set of tools you can use to answer the user's question.
6 +You can invoke functions by writing a "<｜DSML｜function_calls>" block like the following as part of your reply to the user:
7 +<｜DSML｜function_calls>
8 +<｜DSML｜invoke name="$FUNCTION_NAME">
9 +<｜DSML｜parameter name="$PARAMETER_NAME" string="true|false">$PARAMETER_VALUE</｜DSML｜parameter>
10 +...
11 +</｜DSML｜invoke>
12 +<｜DSML｜invoke name="$FUNCTION_NAME2">
13 +...
14 +</｜DSML｜invoke>
15 +</｜DSML｜function_calls>
16 +
17 +String and scalar parameters should be specified as is without any escaping or quotes, while lists and objects should use JSON format. The "string" attribute should be set to "true" for string type parameters and "false" for other types (numbers, booleans, arrays, objects).
18 +
19 +If the thinking_mode is enabled, then after function results you should strongly consider outputting a thinking block. Here is an example:
20 +
21 +<｜DSML｜function_calls>
22 +...
23 +</｜DSML｜function_calls>
24 +
25 +<function_results>
26 +...
27 +</function_results>
28 +
29 +<think>...thinking about results</think>
30 +
31 +Here are the functions available in JSONSchema format:
32 +<functions>
33 +{"name": "search", "description": "Searches for information related to query and displays topn results.", "parameters": {"type": "object", "properties": {"query": {"type": "string", "description": "The search query string"}, "topn": {"type": "integer", "description": "Number of top results to display", "default": 10}, "source": {"type": "string", "description": "Source to search within", "enum": ["web", "news"], "default": "web"}}, "required": ["query"], "additionalProperties": false, "$schema": "http://json-schema.org/draft-07/schema#"}}
34 +{"name": "open", "description": "Opens the link id from the page indicated by cursor starting at line number loc, showing num_lines lines. Valid link ids are displayed with the formatting: 【{id}†.*】. If cursor is not provided, the most recent page is implied. If id is a string, it is treated as a fully qualified URL associated with source. If loc is not provided, the viewport will be positioned at the beginning of the document or centered on the most relevant passage, if available. Use this function without id to scroll to a new location of an opened page.", "parameters": {"type": "object", "properties": {"id": {"description": "Link ID (number) or fully qualified URL (string)", "anyOf": [{"type": "integer"}, {"type": "string"}], "default": -1}, "cursor": {"type": "integer", "description": "Page cursor indicator", "default": -1}, "loc": {"type": "integer", "description": "Starting line number", "default": -1}, "num_lines": {"type": "integer", "description": "Number of lines to show", "default": -1}, "view_source": {"type": "boolean", "description": "Whether to view source", "default": false}, "source": {"type": "string", "description": "Source associated with the URL", "enum": ["web", "news"], "default": "web"}}, "additionalProperties": false, "$schema": "http://json-schema.org/draft-07/schema#"}}
35 +{"name": "find", "description": "Finds exact matches of pattern in the current page, or the page given by cursor.", "parameters": {"type": "object", "properties": {"pattern": {"type": "string", "description": "The exact pattern to search for"}, "cursor": {"type": "integer", "description": "Page cursor to search in", "default": -1}}, "required": ["pattern"], "additionalProperties": false, "$schema": "http://json-schema.org/draft-07/schema#"}}
36 +</functions>
37 +
38 +
39 +# The user's message is: 帮我调研一下，目前有哪些针对search agent的benchmark？详细介绍各自的特点、使用场景、例题。
40 +
41 +<system-reminder>
42 +## Today’s Date
43 +2025-11-27, Thursday.
44 +</system-reminder><｜Assistant｜><think>用户想调研目前针对搜索智能体（search agent）的基准测试（benchmark），要求详细介绍各自的特点、使用场景和例题。我需要提供一份全面的概述，涵盖现有的基准测试，例如WebGPT、WebShop、ALFWorld、HotpotQA、FEVER、TriviaQA、Natural Questions、SearchQA、MS MARCO、BEIR、LoTTE、MIRACL、Zero-shot IR基准等。同时，需要区分这些基准测试是针对检索任务、问答任务，还是具体针对搜索智能体的评估。此外，还应考虑那些专门为评估搜索智能体（即能够与环境交互、执行搜索、浏览网页、使用工具的智能体）而设计的基准测试，例如WebGPT、WebShop、ALFWorld、Mind2Web、WebArena、VisualWebArena、SearchBench、WebVLN、WebNav等。我需要详细介绍每个基准测试的特点、使用场景，并举例说明。为了确保信息的全面性和时效性，我将同时进行两项搜索：一项是针对“search agent benchmark”的通用搜索，以获取广泛的基准测试信息；另一项是针对“搜索智能体 基准测试”的中文搜索，以获取更符合用户语言习惯和可能更具体的中文资料。</think>
45 +
46 +<｜DSML｜function_calls>
47 +<｜DSML｜invoke name="search">
48 +<｜DSML｜parameter name="query" string="true">search agent benchmark 2024</｜DSML｜parameter>
49 +<｜DSML｜parameter name="topn" string="false">10</｜DSML｜parameter>
50 +<｜DSML｜parameter name="source" string="true">web</｜DSML｜parameter>
51 +</｜DSML｜invoke>
52 +<｜DSML｜invoke name="search">
53 +<｜DSML｜parameter name="query" string="true">搜索智能体 基准测试</｜DSML｜parameter>
54 +<｜DSML｜parameter name="topn" string="false">10</｜DSML｜parameter>
55 +<｜DSML｜parameter name="source" string="true">web</｜DSML｜parameter>
56 +</｜DSML｜invoke>
57 +</｜DSML｜function_calls><｜end▁of▁sentence｜>
58 +
59 +<function_results>
60 +<result>[0] search agent benchmark 2024 (Search_Results/search agent benchmark 2024)
61 +**viewing lines [0 - 39] of 39**
62 +
63 +L0:
64 +L1: URL: Search_Results/search agent benchmark 2024
65 +L2: # Search Results
66 +L3:
67 +L4: * 【0†HSCodeComp: A Realistic and Expert-level Benchmark for ...; publish_date:
68 +L5: none†arxiv.org】 Oct 22, 2025 — To fill this gap, we introduce HSCodeComp, the
69 +L6: first realistic, expert-level e-commerce benchmark designed to evaluate deep
70 +L7: search agents in ...
71 +L8: * 【1†open-compass/GTA - A Benchmark for General Tool Agents; publish_date:
72 +L9: none†github.com】 GTA is a benchmark to evaluate the tool-use capability of LLM-
73 +L10: based agents in real-world scenarios. It features three main aspects.
74 +L11: * 【2†Benchmarking real-time trust scoring across five AI Agent ...;
75 +L12: publish_date: none†cleanlab.ai】 Aug 20, 2025 — This article evaluates 5 AI Agent
76 +L13: architectures over the BOLAA (ICLR 2024) benchmark, and assesses the effects of
77 +L14: adding automated trust ...
78 +L15: * 【3†10 AI agent benchmarks; publish_date: none†www.evidentlyai.com】 Jul 11,
79 +L16: 2025 — We put together 10 AI agent benchmarks designed to assess how well
80 +L17: different LLMs perform as agents in real-world scenarios, ...
81 +L18: * 【4†A state-of-the-art search API purpose-built for agents; publish_date:
82 +L19: none†parallel.ai】 Jul 31, 2025 — To evaluate real-world performance of the
83 +L20: Parallel Search MCP Server, we created the WISER-Search benchmark which blends
84 +L21: WISER-Fresh (queries ...
85 +L22: * 【5†AI Agent Benchmarks are Broken; publish_date: none†medium.com】 We break
86 +L23: down the failure modes in current AI agent benchmarks and introduce a checklist
87 +L24: that minimizes the gamability of AI agent benchmarks.
88 +L25: * 【6†Benchmarks and Tree Search for Multimodal LLM Web Agents; publish_date:
89 +L26: none†dpfried.github.io】 2024, When is Tree Search Useful? ○ Dealing with
90 +L27: destructive actions. ○ Some things on the web are very difficult to undo, e.g.,
91 +L28: ordering an item. 56.
92 +L29: * 【7†-Bench: Benchmarking AI agents for the real-world; publish_date:
93 +L30: none†sierra.ai】 Jun 20, 2024 — τ-bench measures an agent's ability to interact
94 +L31: with (simulated) human users and programmatic APIs while following domain-
95 +L32: specific policies in a consistent ...
96 +L33: * 【8†Browser Use = state of the art Web Agent; publish_date: none†browser-
97 +L34: use.com】 Dec 15, 2024 — Browser Use has achieved state-of-the-art performance on
98 +L35: the WebVoyager benchmark, with an impressive 89.1% success rate across 586
99 +L36: diverse web tasks.
100 +L37: * 【9†FutureSearch Benchmarks; publish_date: none†evals.futuresearch.ai】 Find
101 +L38: the original source of a given claim. Example: From , more than 8 out of 1000
102 +L39: users clicked on a phishing link monthly in 2024, up 190% vs 2023.</result>
103 +<result>[1] 搜索智能体 基准测试 (Search_Results/搜索智能体 基准测试)
104 +**viewing lines [0 - 33] of 33**
105 +
106 +L0:
107 +L1: URL: Search_Results/搜索智能体 基准测试
108 +L2: # Search Results
109 +L3:
110 +L4: * 【0†WideSearch：揭示AI 智能体缺失的「广度」能力; publish_date: none†zhuanlan.zhihu.com】 Aug
111 +L5: 16, 2025 — 为系统评估智能体在该任务上的能力，论文构建了第一个专门的基准测试 WideSearch ，包含200 个源于真实世界、横跨18
112 +L6: 个领域的高质量任务。 通过对超过10 个 ...
113 +L7: * 【1†GAIA: 一个严苛的智能体基准- HuggingFace; publish_date: none†www.cnblogs.com】 Jul 9,
114 +L8: 2024 — 我们使用一个用库构建的代码智能体 在GAIA 基准上进行测试，这可以说是最困难、最全面的智能体基准测试……最终我们取得了第一名的成绩！
115 +L9: GAIA: 一个严苛的 ...
116 +L10: * 【2†AI搜索智能体遭遇新挑战：滑铁卢大学团队提出更公平透明的 ...; publish_date: none†www.techwalker.com】
117 +L11: Aug 14, 2025 —
118 +L12: 目前评测AI搜索智能体主要依靠BrowseComp这样的基准测试，它就像一场实时的开卷考试，让AI在真实的网络环境中搜索信息来回答复杂问题。听起来很合理 ...
119 +L13: * 【3†Agentic AI基础设施实践经验系列（六）：Agent质量评估 - AWS; publish_date:
120 +L14: none†aws.amazon.com】 Sep 19, 2025 — TAU-bench
121 +L15: 是一个评估AI智能体在真实世界环境中可靠性的基准测试。它评估智能体是否能够在动态的多轮对话中与用户进行交互，理解需求并完成任务。T-bench ...
122 +L16: * 【4†DeepAgent：能自己找工具的通用推理智能体 - 高瓴人工智能学院; publish_date: none†ai.ruc.edu.cn】
123 +L17: Nov 6, 2025 — 在八大基准测试中，DeepAgent在绝大多数任务上全面领先所有基线模型。
124 +L18: 开放环境优势:在最具挑战的“开放工具检索”场景下(如ToolBench)，其成功率达到64%，远 ...
125 +L19: * 【5†BrowseComp：为浏览智能体设计的简单而具挑战性的基准测试; publish_date: none†blog.csdn.net】 Sep
126 +L20: 22, 2025 — 该基准测试由OpenAI团队开发，旨在推动更可信赖和可靠的AI代理研究。 核心特点.
127 +L21: 挑战性问题设计：BrowseComp的问题设计遵循严格的难度标准：. 人类创建者确保 ...
128 +L22: * 【6†什么是GAIA？; publish_date: none†huggingface.co】 什么是GAIA？ GAIA
129 +L23: 是一个用于评估AI助手在需要核心能力组合的真实世界任务上的表现的基准，这些核心能力包括推理、多模态理解、网页浏览和熟练的工具使用。
130 +L24: * 【7†SuperCLUE-Agent: Agent智能体中文原生任务能力测评基准; publish_date:
131 +L25: none†www.cluebenchmarks.com】 AI
132 +L26: Agent是一个能够自主理解、规划决策和执行复杂任务的智能体，有代表性的Agent比如AutoGPT、AgentGPT等。
133 +L27: 现有关于Agent能力的测评，主要是在英文场景或任务的测评。目前还 ...
134 +L28: * 【8†Gaia2 与ARE：赋能社区的智能体评测; publish_date: none†hub.baai.ac.cn】 Sep 24, 2025 —
135 +L29: 快速评估任意智能体：基于真实或模拟数据，测试不同规则、工具、内容和验证方式下的表现 · 测试智能体的 工具调用与编排能力：可结合本地应用或MCP 工具 · 生成
136 +L30: ...
137 +L31: * 【9†许多AI 智能体评测基准并不可靠; publish_date: none†my.oschina.net】 Aug 13, 2025 —
138 +L32: 当前AI智能体基准测试存在严重缺陷，亟需更严谨的评估框架。作者提出AI智能体基准测试核查清单（ABC），包含43个项目，旨在识别基准测试潜在陷阱，构建可靠评估
139 +L33: ...</result>
140 +</function_results>
141 +
142 +<think>第一个搜索结果（针对“search agent benchmark 2024”）显示了一些相关条目，例如“HSCodeComp: A Realistic and Expert-level Benchmark for ...”、“open-compass/GTA - A Benchmark for General Tool Agents”、“10 AI agent benchmarks”、“WISER-Search benchmark”、“Benchmarks and Tree Search for Multimodal LLM Web Agents”、“τ-bench”、“WebVoyager benchmark”、“FutureSearch Benchmarks”。这些可能都与搜索智能体相关。第二个搜索结果（针对“搜索智能体 基准测试”）显示了一些中文资料，例如“WideSearch”、“BrowseComp”、“GAIA”、“SuperCLUE-Agent”等。这些都可能与搜索智能体相关。为了全面了解这些基准测试，我们需要同时打开并分析其中几个最有潜力的结果。具体来说，我们将打开第一个搜索结果中的“HSCodeComp”（索引0）和“open-compass/GTA”（索引1），以及第二个搜索结果中的“WideSearch”（索引0）和“BrowseComp”（索引5）。</think>
143 +
144 +<｜DSML｜function_calls>
145 +<｜DSML｜invoke name="open">
146 +<｜DSML｜parameter name="cursor" string="false">0</｜DSML｜parameter>
147 +<｜DSML｜parameter name="id" string="false">0</｜DSML｜parameter>
148 +</｜DSML｜invoke>
149 +<｜DSML｜invoke name="open">
150 +<｜DSML｜parameter name="cursor" string="false">0</｜DSML｜parameter>
151 +<｜DSML｜parameter name="id" string="false">1</｜DSML｜parameter>
152 +</｜DSML｜invoke>
153 +<｜DSML｜invoke name="open">
154 +<｜DSML｜parameter name="cursor" string="false">1</｜DSML｜parameter>
155 +<｜DSML｜parameter name="id" string="false">0</｜DSML｜parameter>
156 +</｜DSML｜invoke>
157 +<｜DSML｜invoke name="open">
158 +<｜DSML｜parameter name="cursor" string="false">1</｜DSML｜parameter>
159 +<｜DSML｜parameter name="id" string="false">5</｜DSML｜parameter>
160 +</｜DSML｜invoke>
161 +</｜DSML｜function_calls><｜end▁of▁sentence｜>
162 +
163 +<function_results>
164 +<result>[2] (https://ar5iv.org/html/2510.19631v1)
165 +**viewing lines [0 - 119] of 983**
166 +
167 +L0:
168 +L1: URL: https://ar5iv.org/html/2510.19631v1
169 +L2: # HSCodeComp: A Realistic and Expert-level Benchmark for Deep Search Agents in
170 +L3: Hierarchical Rule Application
171 +L4:
172 +L5: Yiqian Yang† Tian Lan† Qianghuai Jia∗ Li Zhu Hui Jiang Hang Zhu Longyue Wang
173 +L6: Weihua Luo Kaifu Zhang
174 +L7:
175 +L8: Alibaba International Digital Commerce∗* Corresponding Author: Qianghuai Jia
176 +L9: (qianghuai.jqh@alibaba-inc.com)
177 +L10: †\dagger Equal Contribution: Yiqian Yang
178 +L11:
179 +L12: Tian Lan
180 +L13:
181 +L14: ###### Abstract
182 +L15:
183 +L16: Abstract
184 +L17:
185 +L18: Effective deep search agents must not only access open-domain and domain-
186 +L19: specific knowledge but also apply complex rules—such as legal clauses, medical
187 +L20: manuals and tariff rules. These rules often feature vague boundaries and
188 +L21: implicit logic relationships, making precise application challenging for agents.
189 +L22: However, this critical capability is largely overlooked by current agent
190 +L23: benchmarks. To fill this gap, we introduce HSCodeComp, the first realistic,
191 +L24: expert-level e-commerce benchmark designed to evaluate deep search agents in
192 +L25: hierarchical rule application. In this task, the deep reasoning process of
193 +L26: agents is guided by these rules to predict 10-digit Harmonized System Code
194 +L27: (HSCode) of products with noisy but realistic descriptions. These codes,
195 +L28: established by the World Customs Organization, are vital for global supply chain
196 +L29: efficiency. Built from real-world data collected from large-scale e-commerce
197 +L30: platforms, our proposed HSCodeComp comprises 632 product entries spanning
198 +L31: diverse product categories, with these HSCodes annotated by several human
199 +L32: experts. Extensive experimental results on several state-of-the-art LLMs, open-
200 +L33: source, and closed-source agents reveal a huge performance gap: best agent
201 +L34: achieves only 46.8% 10-digit accuracy, far below human experts at 95.0%.
202 +L35: Besides, detailed analysis demonstrates the challenges of hierarchical rule
203 +L36: application, and test-time scaling fails to improve performance further.
204 +L37:
205 +L38: ## 1 Introduction
206 +L39:
207 +L40: Deep search agents have demonstrated significant value in solving complex real-
208 +L41: world problems, where robust external knowledge utilization constitutes a
209 +L42: critical capability [Wu et al., 2025, Tao et al., 2025, Li et al., 2025b]. To
210 +L43: evaluate this capability, numerous established benchmarks are proposed to assess
211 +L44: agents in utilizing open-domain data (e.g., GAIA [Mialon et al., 2023b] and
212 +L45: BrowseComp [Wei et al., 2025]) and domain-specific data (e.g., WebMall [Peeters
213 +L46: et al., 2025a], FinSearchComp [Hu et al., 2025a] and MedBrowseComp [Yu et al.,
214 +L47: 2025b]).
215 +L48:
216 +L49: Beyond open-domain and domain-specific data, agents also need to effectively
217 +L50: apply rules that encode human expert knowledge, particularly in scenarios like
218 +L51: law, medical and e-commerce [Li et al., 2025a, Chen et al., 2025b, Yao et al.,
219 +L52: 2022, Chollet et al., 2025]. For instance, legal case adjudication require
220 +L53: interpreting abstract legal provisions, and accurate e-commerce product
221 +L54: classification in depends on tariff rules [Grainger, 2024]. Previous works have
222 +L55: defined rule application as using specific logical rules with supporting facts
223 +L56: to derive conclusions [Wang et al., 2024, Servantez et al., 2024]. In contrast,
224 +L57: we define it as a core capability for deep search agents, where human-written
225 +L58: rules are systematically applied to guide complex reasoning and decision-making
226 +L59: [Sadowski and Chudziak, 2025]. Building on this observation, we categorize
227 +L60: knowledge data for deep search agents into three levels (Figure 1, left), with
228 +L61: increasing knowledge complexity: (1) Level 1: Open-domain Data - Tests
229 +L62: understanding and deep reasoning abilities of agents on long-form web content.
230 +L63: Established benchmarks include GAIA [Mialon et al., 2023b] and BrowseComp [Wei
231 +L64: et al., 2025]; (2) Level 2: Structured Data - Assesses agents to precisely
232 +L65: utilize structured data such as databases and knowledge graphs, as seen in
233 +L66: domain-specific benchmarks like WebMall [Peeters et al., 2025a], MedBrowseComp
234 +L67: [Chen et al., 2025b] and FinSearchComp [Hu et al., 2025a]; (3) Level 3: Rule
235 +L68: Data - Evaluates agents to apply complex and abstract rules [Chollet et al.,
236 +L69: 2025]. This level presents two key challenges: (a) making accurate decisions
237 +L70: when rules contain vague natural language descriptions [Sadowski and Chudziak,
238 +L71: 2025]; and (b) reasoning about logical dependencies among rules, such as
239 +L72: exception clauses and cross-category relationships [Guha et al., 2023]. Despite
240 +L73: the importance of rule application in real-world scenarios, current agent
241 +L74: benchmarks largely overlook its evaluation.
242 +L75:
243 +L76: To fill this gap, we introduce HSCodeComp (short for the Harmonized System Code
244 +L77: (HSCode) Competition), the first realistic, expert-level e-commerce benchmark
245 +L78: designed to evaluate agents in predicting complete 10-digit Harmonized System
246 +L79: Code (HSCode) of the product, using hierarchical rules (e.g., eWTP tariff
247 +L80: rules111https://www.ewtp.com/web/smart/hscode). HSCodes organize products
248 +L81: through a hierarchical structure spanning over 5,000 distinct codes across
249 +L82: multiple classification levels, representing the global standard for classifying
250 +L83: traded international goods, established by the World Customs Organization and
251 +L84: implemented across more than 200 countries for customs clearance and tariff
252 +L85: determination [Grainger, 2024, Nath et al., 2025]. Built from the data of the
253 +L86: large-scale e-commerce platforms, our proposed HSCodeComp comprises 632
254 +L87: carefully curated product entries, encompassing 27 unique HS chapters and 32
255 +L88: distinct first-level categories. These HSCodes have been rigorously annotated by
256 +L89: multiple e-commerce domain experts, ensuring that HSCodeComp is expert-level.
257 +L90: Accurately predicting the exact 10-digit HSCode presents significant challenges:
258 +L91: agents must perform multi-hop hierarchical reasoning with complex tariff rules
259 +L92: while processing noisy but realistic product descriptions that often contain
260 +L93: abbreviations, language variations, or incomplete information.
261 +L94:
262 +L95: Extensive experiments on the state-of-the-art baselines, including 14 advanced
263 +L96: foundation models, 6 advanced open-source agent systems and 3 closed-source
264 +L97: agent systems, demonstrate that HSCode prediction task remains a substantial
265 +L98: challenge for current AI approaches. As shown in the Figure 1 (right), even the
266 +L99: best-performing system (SmolAgent [Roucher et al., 2025] with GPT-5) achieves
267 +L100: only 46.8% accuracy, substantially below the 95.0% accuracy attained by human
268 +L101: experts. Further detailed analysis reveals that existing agent systems lack
269 +L102: critical capabilities required for this complex hierarchical rule applications.
270 +L103: Notably, test-time scaling approach—which has proven effective in other
271 +L104: reasoning tasks [Guo et al., 2025, Liu et al., 2025]—fail to improve performance
272 +L105: on HSCodeComp. These observations demonstrate the challenging nature of our
273 +L106: proposed HSCodeComp, highlighting the need for more effective designs of agent
274 +L107: systems. To facilitate future research, we will publicly release codes and the
275 +L108: benchmark dataset of HSCodeComp.
276 +L109:
277 +L110: ## 2 Related Works
278 +L111:
279 +L112: ### 2.1 Previous Works in HSCode Prediction
280 +L113:
281 +L114: Previous works treat HSCode prediction as the e-commerce text classification
282 +L115: task [Grainger, 2024], using pre-trained BERT models [Liao et al., 2024, Shubham
283 +L116: et al., 2022] or Large Language Models (LLMs) prompting [Hussain and Ahmed,
284 +L117: 2023]. However, these approaches fail to leverage domain-specific knowledge,
285 +L118: especially the rules written by human experts [Hussain and Ahmed, 2023, Judy,
286 +L119: 2024]. Besides, existing HSCode benchmarks face two critical limitations [Judy, </result>
287 +<result>Error fetching URL `https://github.com/open-compass/GTA`</result>
288 +<result>Error fetching URL `https://zhuanlan.zhihu.com/p/1939316761775301093`</result>
289 +<result>[3] (https://blog.csdn.net/qq_41472205/article/details/151974603)
290 +**viewing lines [0 - 167] of 187**
291 +
292 +L0:
293 +L1: URL: https://blog.csdn.net/qq_41472205/article/details/151974603
294 +L2: # BrowseComp：为浏览智能体设计的简单而具挑战性的基准测试
295 +L3:
296 +L4: BrowseComp：AI浏览能力评估基准
297 +L5:
298 +L6: 最新推荐文章于 2025-11-12 13:40:20 发布
299 +L7:
300 +L8: 原创 于 2025-09-22 22:33:04 发布 · 1.3k 阅读
301 +L9:
302 +L10: · 9
303 +L11: · 25 ·
304 +L12: CC 4.0 BY-SA版权
305 +L13:
306 +L14: 版权声明：本文为博主原创文章，遵循 CC 4.0 BY-SA 版权协议，转载请附上原文出处链接和本声明。
307 +L15:
308 +L16: ## BrowseComp：为浏览智能体设计的简单而具挑战性的基准测试
309 +L17:
310 +L18: 在人工智能从基础聊天机器人向推理器和智能体发展的进程中，具备浏览互联网能力的人工智能模型正变得越来越重要。今天，我们将介绍一个名为BrowseComp的创新基准
311 +L19: 测试，它专门设计用于评估AI代理在复杂网络浏览任务中的能力。
312 +L20:
313 +L21: ### 什么是BrowseComp？
314 +L22:
315 +L23: BrowseComp（全称Browsing Competition）是一个包含1,266个挑战性问题的基准测试集，专门用于衡量AI代理在互联网上持续导航、寻找难
316 +L24: 以找到的纠缠信息的能力。该基准测试由OpenAI团队开发，旨在推动更可信赖和可靠的AI代理研究。
317 +L25:
318 +L26: #### 核心特点
319 +L27:
320 +L28: 挑战性问题设计：BrowseComp的问题设计遵循严格的难度标准：
321 +L29:
322 +L30: - 人类创建者确保问题在10分钟内无法被人解决
323 +L31: - 现有模型（包括带浏览功能的ChatGPT和早期版本的OpenAI Deep Research）无法解决
324 +L32: - 通过5次简单Google搜索无法在结果首页找到答案
325 +L33:
326 +L34: 简单易验证：尽管问题极具挑战性，但答案形式简单——都是短字符串，便于自动验证模型输出的正确性。
327 +L35:
328 +L36: ### 为什么需要BrowseComp？
329 +L37:
330 +L38: #### 现有基准的局限性
331 +L39:
332 +L40: 传统的信息检索基准（如TriviaQA、HotpotQA等）主要关注易于查找的信息，随着语言模型的进步，这些基准已经趋于饱和。而BrowseComp专注于那些需
333 +L41: 要浏览大量网站才能解决的"硬核"问题。
334 +L42:
335 +L43: #### 模拟真实挑战
336 +L44:
337 +L45: BrowseComp问题通常采用"逆向设计"方法：创建者从一个已知事实出发，构建一个搜索空间巨大但验证简单的问题。例如：
338 +L46:
339 +L47: “找出2018-2023年间在EMNLP会议上发表、第一作者本科毕业于达特茅斯学院、第四作者本科毕业于宾夕法尼亚大学的科学论文标题”
340 +L48:
341 +L49: 这类问题验证简单，但解决起来需要检查数千篇论文并调查每位作者的背景。
342 +L50:
343 +L51: ### 数据集特点
344 +L52:
345 +L53: #### 主题多样性
346 +L54:
347 +L55: BrowseComp涵盖了广泛的主题领域（如图2所示），包括历史、科学、文化等。创建者被鼓励基于个人兴趣设计问题，这有助于提高数据质量和参与度。
348 +L56:
349 +L57: #### 质量保证
350 +L58:
351 +L59: 为确保答案的唯一性，创建者需要：
352 +L60:
353 +L61: - 对问题内容有足够了解，确信没有其他有效答案
354 +L62: - 如果不确定，则添加更多约束条件
355 +L63: - 接受其他创建者的验证反馈
356 +L64:
357 +L65: ### 人类表现基准
358 +L66:
359 +L67: 为了衡量BrowseComp的难度，研究人员让人类创建者尝试解决问题（不能解答自己创建的问题）。结果显示：
360 +L68:
361 +L69: - **70.8%**的问题在2小时搜索后人类选择放弃
362 +L70: - **29.2%**的问题被成功解决
363 +L71: - 在解决的问题中，**86.4%**的人类答案与参考答案一致
364 +L72:
365 +L73: 这表明BrowseComp确实极具挑战性，即使是熟悉数据集的人类专家也难以在有限时间内解决大部分问题。
366 +L74:
367 +L75: ### AI模型表现评估
368 +L76:
369 +L77: #### 各模型对比
370 +L78:
371 +L79: 研究人员评估了多种模型在BrowseComp上的表现：
372 +L80:
373 +L81: 模型 | 准确率(%) | 校准误差(%)
374 +L82: ---|---|---
375 +L83: GPT-4o | 0.6 | 69
376 +L84: GPT-4o（带浏览） | 1.9 | 82
377 +L85: GPT-4.5 | 0.9 | 68
378 +L86: OpenAI o1 | 9.9 | 65
379 +L87: Deep Research | 51.5 | 91
380 +L88:
381 +L89: #### 关键发现
382 +L90:
383 +L91: - 基础模型表现不佳：GPT-4o和GPT-4.5准确率接近零，凸显了基准的难度
384 +L92: - 浏览功能带来有限提升：启用浏览功能的GPT-4o准确率略有提高，但仍很低
385 +L93: - 推理能力的重要性：OpenAI o1虽然没有浏览能力，但凭借更强的推理能力获得较高准确率
386 +L94: - 专业模型的优势：专门为持久网络浏览训练的Deep Research模型解决了约一半的问题
387 +L95:
388 +L96: #### 计算资源与性能关系
389 +L97:
390 +L98: 研究表明，BrowseComp性能随测试时计算资源的增加而平滑提升（如图1所示）。这与智能体模型的特性一致——更多计算资源允许模型浏览更多网站，从而提高找到正确
391 +L99: 答案的机会。
392 +L100:
393 +L101: ### 进阶策略分析
394 +L102:
395 +L103: #### 聚合策略的效果
396 +L104:
397 +L105: 通过让模型多次尝试同一问题并采用投票策略，可以显著提升性能：
398 +L106:
399 +L107: - 多数投票：选择样本中最常见的答案
400 +L108: - 加权投票：根据模型置信度加权投票
401 +L109: - 最佳选择：选择置信度最高的答案
402 +L110:
403 +L111: 这些方法将Deep Research的性能提升了15-25%，表明模型通常能够识别自己的正确答案。
404 +L112:
405 +L113: #### 任务难度分布
406 +L114:
407 +L115: 分析显示，BrowseComp中的任务难度分布广泛：
408 +L116:
409 +L117: - 16%的任务被Deep Research完美解决（100%通过率）
410 +L118: - 14%的任务完全失败（0%通过率）
411 +L119: - 其余任务处于中间难度水平
412 +L120:
413 +L121: ### BrowseComp的意义与局限性
414 +L122:
415 +L123: #### 作为评估工具的价值
416 +L124:
417 +L125: BrowseComp可被视为浏览智能体的"编程竞赛"——虽然不全面，但对核心浏览能力提供了有用的衡量。正如在编程竞赛中表现出色的模型很可能在其他编码任务中表现良
418 +L126: 好，能够解决BrowseComp的模型在定位难以查找信息方面应该具有很强的能力。
419 +L127:
420 +L128: #### 已知局限性
421 +L129:
422 +L130: - 不涉及真实用户查询分布：如生成长回答或解决查询歧义等挑战
423 +L131: - 答案唯一性假设：虽然尽力确保，但不能完全保证没有其他有效答案
424 +L132: - 单模态限制：当前仅关注文本信息检索，未来可扩展至图像、视频等模态
425 +L133:
426 +L134: ### 结论与展望
427 +L135:
428 +L136: BrowseComp为评估AI代理的网络浏览能力提供了一个具有挑战性且易于使用的基准。它强调的持久性、创造性和战略性浏览能力，正是未来AI智能体在实际应用中所需
429 +L137: 的核心技能。
430 +L138:
431 +L139: 随着更多研究团队在BrowseComp上测试他们的模型，我们期待看到浏览智能体技术的快速进步，最终使AI能够更可靠地帮助我们获取那些深藏在互联网角落的宝贵信息。
432 +L140:
433 +L141: 访问地址：BrowseComp已在GitHub上开源，地址为 https://github.com/openai/simple-evals
434 +L142:
435 +L143: 添加红包
436 +L144:
437 +L145: 请填写红包祝福语或标题
438 +L146:
439 +L147: 个
440 +L148:
441 +L149: 红包个数最小为10个
442 +L150:
443 +L151: 元
444 +L152:
445 +L153: 红包金额最低5元
446 +L154:
447 +L155: 当前余额3.43元 前往充值 >
448 +L156:
449 +L157: 需支付：10.00元
450 +L158:
451 +L159: 成就一亿技术人!
452 +L160:
453 +L161: hope_wisdom
454 +L162: 发出的红包
455 +L163: 实付元
456 +L164:
457 +L165: 点击重新获取
458 +L166:
459 +L167: 扫码支付</result>
460 +</function_results>
461 +
462 +<think>打开的HSCodeComp页面详细介绍了这个基准测试，它是一个用于评估深度搜索智能体在分层规则应用方面能力的基准测试，特点包括：基于真实世界电商数据、任务复杂、需要多跳层次推理、使用嘈杂的产品描述，并且标注由专家完成。使用场景是评估智能体在规则应用方面的能力，例如关税代码分类。例题是预测产品的10位HS编码。打开的BrowseComp页面详细介绍了这个基准测试，它是一个用于评估浏览智能体（即能够在互联网上导航、搜索信息的智能体）的基准测试。其特点包括：包含1266个具有挑战性的问题，需要持续导航和搜索难以查找的信息；答案形式简单，便于自动验证；问题设计确保人类在10分钟内无法解决，且现有模型也无法解决；涵盖了广泛的主题领域。使用场景是评估智能体在复杂网络浏览任务中的能力。例题是：“找出2018-2023年间在EMNLP会议上发表、第一作者本科毕业于达特茅斯学院、第四作者本科毕业于宾夕法尼亚大学的科学论文标题”。另外两个打开的页面（open-compass/GTA和WideSearch）出现了错误，可能链接已失效或需要其他访问方式。为了进一步获取其他相关基准测试的信息，我们需要同时进行多项搜索。考虑到用户可能对搜索智能体基准测试有广泛需求，我们将同时搜索“WebVoyager benchmark”、“WebArena benchmark”、“GAIA benchmark”、“ToolBench benchmark”、“HotpotQA benchmark”、“FEVER benchmark”、“TriviaQA benchmark”、“Natural Questions benchmark”、“MS MARCO benchmark”、“BEIR benchmark”、“MIRACL benchmark”、“Zero-shot IR benchmark”、“WebGPT benchmark”、“WebShop benchmark”、“ALFWorld benchmark”、“Mind2Web benchmark”、“VisualWebArena benchmark”、“SearchBench benchmark”、“WebVLN benchmark”、“WebNav benchmark”。这些搜索将帮助我们全面了解现有基准测试的概况。</think>
463 +
464 +<｜DSML｜function_calls>
465 +<｜DSML｜invoke name="search">
466 +<｜DSML｜parameter name="query" string="true">WebVoyager benchmark</｜DSML｜parameter>
467 +<｜DSML｜parameter name="topn" string="false">10</｜DSML｜parameter>
468 +<｜DSML｜parameter name="source" string="true">web</｜DSML｜parameter>
469 +</｜DSML｜invoke>
470 +<｜DSML｜invoke name="search">
471 +<｜DSML｜parameter name="query" string="true">WebArena benchmark</｜DSML｜parameter>
472 +<｜DSML｜parameter name="topn" string="false">10</｜DSML｜parameter>
473 +<｜DSML｜parameter name="source" string="true">web</｜DSML｜parameter>
474 +</｜DSML｜invoke>
475 +<｜DSML｜invoke name="search">
476 +<｜DSML｜parameter name="query" string="true">GAIA benchmark</｜DSML｜parameter>
477 +<｜DSML｜parameter name="topn" string="false">10</｜DSML｜parameter>
478 +<｜DSML｜parameter name="source" string="true">web</｜DSML｜parameter>
479 +</｜DSML｜invoke>
480 +<｜DSML｜invoke name="search">
481 +<｜DSML｜parameter name="query" string="true">ToolBench benchmark</｜DSML｜parameter>
482 +<｜DSML｜parameter name="topn" string="false">10</｜DSML｜parameter>
483 +<｜DSML｜parameter name="source" string="true">web</｜DSML｜parameter>
484 +</｜DSML｜invoke>
485 +<｜DSML｜invoke name="search">
486 +<｜DSML｜parameter name="query" string="true">HotpotQA benchmark</｜DSML｜parameter>
487 +<｜DSML｜parameter name="topn" string="false">10</｜DSML｜parameter>
488 +<｜DSML｜parameter name="source" string="true">web</｜DSML｜parameter>
489 +</｜DSML｜invoke>
490 +<｜DSML｜invoke name="search">
491 +<｜DSML｜parameter name="query" string="true">FEVER benchmark</｜DSML｜parameter>
492 +<｜DSML｜parameter name="topn" string="false">10</｜DSML｜parameter>
493 +<｜DSML｜parameter name="source" string="true">web</｜DSML｜parameter>
494 +</｜DSML｜invoke>
495 +<｜DSML｜invoke name="search">
496 +<｜DSML｜parameter name="query" string="true">TriviaQA benchmark</｜DSML｜parameter>
497 +<｜DSML｜parameter name="topn" string="false">10</｜DSML｜parameter>
498 +<｜DSML｜parameter name="source" string="true">web</｜DSML｜parameter>
499 +</｜DSML｜invoke>
500 +<｜DSML｜invoke name="search">
501 +<｜DSML｜parameter name="query" string="true">Natural Questions benchmark</｜DSML｜parameter>
502 +<｜DSML｜parameter name="topn" string="false">10</｜DSML｜parameter>
503 +<｜DSML｜parameter name="source" string="true">web</｜DSML｜parameter>
504 +</｜DSML｜invoke>
505 +<｜DSML｜invoke name="search">
506 +<｜DSML｜parameter name="query" string="true">MS MARCO benchmark</｜DSML｜parameter>
507 +<｜DSML｜parameter name="topn" string="false">10</｜DSML｜parameter>
508 +<｜DSML｜parameter name="source" string="true">web</｜DSML｜parameter>
509 +</｜DSML｜invoke>
510 +<｜DSML｜invoke name="search">
511 +<｜DSML｜parameter name="query" string="true">BEIR benchmark</｜DSML｜parameter>
512 +<｜DSML｜parameter name="topn" string="false">10</｜DSML｜parameter>
513 +<｜DSML｜parameter name="source" string="true">web</｜DSML｜parameter>
514 +</｜DSML｜invoke>
515 +<｜DSML｜invoke name="search">
516 +<｜DSML｜parameter name="query" string="true">MIRACL benchmark</｜DSML｜parameter>
517 +<｜DSML｜parameter name="topn" string="false">10</｜DSML｜parameter>
518 +<｜DSML｜parameter name="source" string="true">web</｜DSML｜parameter>
519 +</｜DSML｜invoke>
520 +<｜DSML｜invoke name="search">
521 +<｜DSML｜parameter name="query" string="true">Zero-shot IR benchmark</｜DSML｜parameter>
522 +<｜DSML｜parameter name="topn" string="false">10</｜DSML｜parameter>
523 +<｜DSML｜parameter name="source" string="true">web</｜DSML｜parameter>
524 +</｜DSML｜invoke>
525 +<｜DSML｜invoke name="search">
526 +<｜DSML｜parameter name="query" string="true">WebGPT benchmark</｜DSML｜parameter>
527 +<｜DSML｜parameter name="topn" string="false">10</｜DSML｜parameter>
528 +<｜DSML｜parameter name="source" string="true">web</｜DSML｜parameter>
529 +</｜DSML｜invoke>
530 +<｜DSML｜invoke name="search">
531 +<｜DSML｜parameter name="query" string="true">WebShop benchmark</｜DSML｜parameter>
532 +<｜DSML｜parameter name="topn" string="false">10</｜DSML｜parameter>
533 +<｜DSML｜parameter name="source" string="true">web</｜DSML｜parameter>
534 +</｜DSML｜invoke>
535 +<｜DSML｜invoke name="search">
536 +<｜DSML｜parameter name="query" string="true">ALFWorld benchmark</｜DSML｜parameter>
537 +<｜DSML｜parameter name="topn" string="false">10</｜DSML｜parameter>
538 +<｜DSML｜parameter name="source" string="true">web</｜DSML｜parameter>
539 +</｜DSML｜invoke>
540 +<｜DSML｜invoke name="search">
541 +<｜DSML｜parameter name="query" string="true">Mind2Web benchmark</｜DSML｜parameter>
542 +<｜DSML｜parameter name="topn" string="false">10</｜DSML｜parameter>
543 +<｜DSML｜parameter name="source" string="true">web</｜DSML｜parameter>
544 +</｜DSML｜invoke>
545 +<｜DSML｜invoke name="search">
546 +<｜DSML｜parameter name="query" string="true">VisualWebArena benchmark</｜DSML｜parameter>
547 +<｜DSML｜parameter name="topn" string="false">10</｜DSML｜parameter>
548 +<｜DSML｜parameter name="source" string="true">web</｜DSML｜parameter>
549 +</｜DSML｜invoke>
550 +<｜DSML｜invoke name="search">
551 +<｜DSML｜parameter name="query" string="true">SearchBench benchmark</｜DSML｜parameter>
552 +<｜DSML｜parameter name="topn" string="false">10</｜DSML｜parameter>
553 +<｜DSML｜parameter name="source" string="true">web</｜DSML｜parameter>
554 +</｜DSML｜invoke>
555 +<｜DSML｜invoke name="search">
556 +<｜DSML｜parameter name="query" string="true">WebVLN benchmark</｜DSML｜parameter>
557 +<｜DSML｜parameter name="topn" string="false">10</｜DSML｜parameter>
558 +<｜DSML｜parameter name="source" string="true">web</｜DSML｜parameter>
559 +</｜DSML｜invoke>
560 +<｜DSML｜invoke name="search">
561 +<｜DSML｜parameter name="query" string="true">WebNav benchmark</｜DSML｜parameter>
562 +<｜DSML｜parameter name="topn" string="false">10</｜DSML｜parameter>
563 +<｜DSML｜parameter name="source" string="true">web</｜DSML｜parameter>
564 +</｜DSML｜invoke>
565 +</｜DSML｜function_calls><｜end▁of▁sentence｜>
566 +
567 +<function_results>
568 +<result>[4] WebVoyager benchmark (Search_Results/WebVoyager benchmark)
569 +**viewing lines [0 - 40] of 40**
570 +
571 +L0:
572 +L1: URL: Search_Results/WebVoyager benchmark
573 +L2: # Search Results
574 +L3:
575 +L4: * 【0†WebVoyager Benchmark; publish_date: none†www.trykura.com】 Performance
576 +L5: measured across various websites in WebVoyager shows Kura consistently
577 +L6: outperforming other agents, with particularly strong results on e-commerce ...
578 +L7: * 【1†WebVoyager: Building an End-to-End Web Agent with ...; publish_date:
579 +L8: none†github.com】 WebVoyager is an innovative Large Multimodal Model (LMM)
580 +L9: powered web agent that can complete user instructions end-to-end by interacting
581 +L10: with real-world ...
582 +L11: * 【2†AI Browser Agent Leaderboard | Steel.dev; publish_date:
583 +L12: none†leaderboard.steel.dev】 See how various AI browser agents stack up based on
584 +L13: their accuracy in completing web-based tasks on the WebVoyager benchmark.
585 +L14: * 【3†[2401.13919] WebVoyager: Building an End-to-End Web ...; publish_date:
586 +L15: none†arxiv.org】 by H He · 2024 · Cited by 282 — We show that WebVoyager achieves
587 +L16: a 59.1% task success rate on our benchmark, significantly surpassing the
588 +L17: performance of both GPT-4 (All ...
589 +L18: * 【4†Our Agent-E SOTA Results on the WebVoyager Benchmark; publish_date:
590 +L19: none†www.emergence.ai】 Jul 11, 2024 — WebVoyager is a benchmark that tests an
591 +L20: agent's capabilities for navigation on dynamic live websites. It is more
592 +L21: representative than WebArena [4] ...
593 +L22: * 【5†Browser Use = state of the art Web Agent; publish_date: none†browser-
594 +L23: use.com】 Dec 15, 2024 — Browser Use has achieved state-of-the-art performance on
595 +L24: the WebVoyager benchmark, with an impressive 89.1% success rate across 586
596 +L25: diverse web tasks.
597 +L26: * 【6†Magnitude achieves SOTA 94% on WebVoyager benchmark; publish_date:
598 +L27: none†github.com】 Magnitude achieves state-of-the-art performance with 93.9%
599 +L28: success rate on WebVoyager, beating all other browser agents.
600 +L29: * 【7†WebVoyager: Autonomous Web Agent Benchmark; publish_date:
601 +L30: none†www.emergentmind.com】 3 days ago — WebVoyager Benchmark is a comprehensive
602 +L31: evaluation suite for autonomous web agents, featuring 643 tasks across 15
603 +L32: popular websites.
604 +L33: * 【8†WebVoyager Benchmark Results; publish_date: none†www.browserable.ai】
605 +L34: Browserable has achieved 90.4% on the WebVoyager benchmark. This is best-in-
606 +L35: class performance across all web agents. This was done across 567 web tasks
607 +L36: which ...
608 +L37: * 【9†89% achieved on WebVoyager using Anchor + Browser Use; publish_date:
609 +L38: none†www.reddit.com】 Thanks to the amazing work from the browser-use open-source
610 +L39: community and the built-in support from Anchor Browser, we've hit an 89% score
611 +L40: on WebVoyager.</result>
612 +<result>[5] WebArena benchmark (Search_Results/WebArena benchmark)
613 +**viewing lines [0 - 42] of 42**
614 +
615 +L0:
616 +L1: URL: Search_Results/WebArena benchmark
617 +L2: # Search Results
618 +L3:
619 +L4: * 【0†WebArena: A Realistic Web Environment for Building ...; publish_date:
620 +L5: none†webarena.dev】 Our benchmark is implemented in our fully interactable
621 +L6: highly-realistic WebArena environment. It features diverse tasks human may
622 +L7: encounter in their daily ...
623 +L8: * 【1†[2307.13854] WebArena: A Realistic Web Environment for ...; publish_date:
624 +L9: none†arxiv.org】 by S Zhou · 2023 · Cited by 637 — Building upon our
625 +L10: environment, we release a set of benchmark tasks focusing on evaluating the
626 +L11: functional correctness of task completions.
627 +L12: * 【2†WebArena: A Realistic Web Environment for Building ...; publish_date:
628 +L13: none†www.cmu.edu】 WebArena introduces a benchmark on interpreting high-level
629 +L14: realistic natural language command to concrete web-based interactions. We
630 +L15: provide annotated programs ...
631 +L16: * 【3†GitHub - web-arena-x/webarena: Code repo for ...; publish_date:
632 +L17: none†github.com】 [12/20/2024] Check out our new benchmark on even more
633 +L18: consequential tasks, including terminal use and coding, TheAgentCompany.
634 +L19: [12/21/2023] We release the ...
635 +L20: * 【4†WebArena Benchmark and the State of Agentic AI; publish_date:
636 +L21: none†medium.com】 In short, WebArena established a new standard for realism and
637 +L22: complexity in web agent evaluation, forcing AI agents to operate in dynamic,
638 +L23: high- ...
639 +L24: * 【5†WebArena: A Realistic Web Environment for Building ...; publish_date:
640 +L25: none†huggingface.co】 Jul 25, 2023 — WebArena, a realistic and reproducible
641 +L26: environment, evaluates the performance of autonomous agents performing complex
642 +L27: tasks on websites using ...
643 +L28: * 【6†WebArena Benchmark: Evaluating Web Agents; publish_date:
644 +L29: none†www.emergentmind.com】 Jun 30, 2025 — WebArena Benchmark is a self-contained
645 +L30: suite that evaluates autonomous agents on realistic, multi-step web tasks using
646 +L31: natural language ...
647 +L32: * 【7†VisualWebArena is a benchmark for multimodal agents.; publish_date:
648 +L33: none†github.com】 VisualWebArena is a realistic and diverse benchmark for
649 +L34: evaluating multimodal autonomous language agents. It comprises of a set of
650 +L35: diverse and complex web-based ...
651 +L36: * 【8†WebDev Arena Leaderboard - LMArena; publish_date: none†web.lmarena.ai】
652 +L37: WebDev Arena is a real-time AI coding competition where models go head-to-head
653 +L38: in web development challenges, developed by LMArena.
654 +L39: * 【9†WebArena: A Realistic Web Environment for Building ...; publish_date:
655 +L40: none†arxiv.org】 Apr 16, 2024 — We use this benchmark to evaluate several agents
656 +L41: that can follow NL command and perform web-based tasks (§4). These agents are
657 +L42: implemented in a ...</result>
658 +<result>[6] GAIA benchmark (Search_Results/GAIA benchmark)
659 +**viewing lines [0 - 41] of 41**
660 +
661 +L0:
662 +L1: URL: Search_Results/GAIA benchmark
663 +L2: # Search Results
664 +L3:
665 +L4: * 【0†GAIA Leaderboard - a Hugging Face Space by ...; publish_date:
666 +L5: none†huggingface.co】 GAIA is a benchmark which aims at evaluating next-
667 +L6: generation LLMs (LLMs with augmented capabilities due to added tooling,
668 +L7: efficient prompting, access to search ...
669 +L8: * 【1†[2311.12983] GAIA: a benchmark for General AI Assistants; publish_date:
670 +L9: none†arxiv.org】 by G Mialon · 2023 · Cited by 367 — GAIA proposes real-world
671 +L10: questions that require a set of fundamental abilities such as reasoning, multi-
672 +L11: modality handling, web browsing, and generally tool-use ...
673 +L12: * 【2†GAIA benchmark; publish_date: none†huggingface.co】 This is the
674 +L13: organisation page for all things related to GAIA, a benchmark for General AI
675 +L14: Assistants. You can find all the information and links on the GAIA ...
676 +L15: * 【3†GAIA: A Benchmark for General AI Assistants; publish_date:
677 +L16: none†ukgovernmentbeis.github.io】 This is an Inspect AI implementation of the
678 +L17: GAIA (General AI Assistants) benchmark, consisting of 450 questions testing tool
679 +L18: use on realistic assistant tasks.
680 +L19: * 【4†GAIA: a benchmark for general AI assistants | Research; publish_date:
681 +L20: none†ai.meta.com】 May 6, 2024 — GAIA proposes real-world questions that require
682 +L21: a set of fundamental abilities such as reasoning, multi-modality handling, web
683 +L22: browsing, and generally tool-use ...
684 +L23: * 【5†HAL: GAIA Leaderboard; publish_date: none†hal.cs.princeton.edu】 GAIA is a
685 +L24: benchmark for General AI Assistants that requires a set of fundamental
686 +L25: abilities such as reasoning, multi-modality handling, web browsing, and tool-
687 +L26: ...
688 +L27: * 【6†GAIA: The LLM Agent Benchmark Everyone's Talking About; publish_date:
689 +L28: none†towardsdatascience.com】 May 29, 2025 — GAIA stands for General AI
690 +L29: Assistants benchmark [1]. This benchmark was introduced to specifically evaluate
691 +L30: LLM agents on their ability to act as general- ...
692 +L31: * 【7†GAIA: a benchmark for General AI Assistants; publish_date:
693 +L32: none†openreview.net】 by G Mialon · Cited by 367 — GAIA proposes real-world
694 +L33: questions that require a set of fundamental abilities such as reasoning, multi-
695 +L34: modality handling, web browsing, and generally tool-use ...
696 +L35: * 【8†Rethinking AI Evaluation: Introducing the GAIA Benchmark; publish_date:
697 +L36: none†medium.com】 The authors introduce GAIA, a benchmark designed to assess the
698 +L37: robustness of AI systems across a variety of practical tasks.
699 +L38: * 【9†H2O.ai Tops the General AI Assistant (GAIA) Test; publish_date:
700 +L39: none†h2o.ai】 Mar 17, 2025 — Our h2oGPTe Agent has once again claimed the #1 spot
701 +L40: on the prestigious GAIA (General AI Assistants) benchmark with an impressive
702 +L41: 75% accuracy rate.</result>
703 +<result>[7] ToolBench benchmark (Search_Results/ToolBench benchmark)
704 +**viewing lines [0 - 40] of 40**
705 +
706 +L0:
707 +L1: URL: Search_Results/ToolBench benchmark
708 +L2: # Search Results
709 +L3:
710 +L4: * 【0†ToolBench, an evaluation suite for LLM tool manipulation ...;
711 +L5: publish_date: none†github.com】 The ToolBench is a benchmark consisting of
712 +L6: diverse software tools for real-world tasks. We also provide easy-to-use
713 +L7: infrastructure in this repository.
714 +L8: * 【1†OpenBMB/ToolBench; publish_date: none†github.com】 [2023/7/27] New version
715 +L9: ToolBench is released. ✨Here is an overview of the dataset construction,
716 +L10: training, and evaluation. ✨✨Features:.
717 +L11: * 【2†Towards Stable Large-Scale Benchmarking on Tool ...; publish_date:
718 +L12: none†arxiv.org】 by Z Guo · 2024 · Cited by 100 — We introduce StableToolBench, a
719 +L13: benchmark evolving from ToolBench, proposing a virtual API server and stable
720 +L14: evaluation system.
721 +L15: * 【3†StableToolBench - Zhicheng Guo; publish_date: none†zhichengg.github.io】
722 +L16: We introduce StableToolBench, a benchmark evolving from ToolBench, proposing a
723 +L17: virtual API server and stable evaluation system.
724 +L18: * 【4†ToolBench | EvalScope - Read the Docs; publish_date:
725 +L19: none†evalscope.readthedocs.io】 We evaluate the effectiveness of the ToolBench
726 +L20: benchmark: ToolBench (Qin et al., 2023b). The task involves integrating API
727 +L21: calls to complete tasks.
728 +L22: * 【5†Towards Stable Large-Scale Benchmarking on Tool ...; publish_date:
729 +L23: none†aclanthology.org】 by Z Guo · 2024 · Cited by 100 — We introduce
730 +L24: StableToolBench, a benchmark evolving from ToolBench, proposing a virtual API
731 +L25: server and stable evaluation system.
732 +L26: * 【6†ML-Tool-Bench: Tool-Augmented Planning for ML Tasks; publish_date:
733 +L27: none†openreview.net】 Sep 18, 2025 — In this work, we introduce a comprehensive
734 +L28: benchmark for evaluating tool-augmented ML agents using a curated set of 61
735 +L29: specialized tools and 15 ...
736 +L30: * 【7†-Bench: Benchmarking AI agents for the real-world; publish_date:
737 +L31: none†sierra.ai】 Jun 20, 2024 — τ-bench measures an agent's ability to interact
738 +L32: with (simulated) human users and programmatic APIs while following domain-
739 +L33: specific policies in a consistent ...
740 +L34: * 【8†ToolEval Leaderboard; publish_date: none†openbmb.github.io】 ToolEval is
741 +L35: an automatic evaluator build for tool learning which incorporates two evaluation
742 +L36: metrics, Pass Rate and Win Rate(Preference).
743 +L37: * 【9†What is the best benchmark dataset for multi-step tool-use?;
744 +L38: publish_date: none†www.reddit.com】 I'm a newbie trying to evaluate the
745 +L39: performance of different prompts strategies for multi-step tool-using, wondering
746 +L40: what is the recommended benchmark dataset ...</result>
747 +<result>[8] HotpotQA benchmark (Search_Results/HotpotQA benchmark)
748 +**viewing lines [0 - 39] of 39**
749 +
750 +L0:
751 +L1: URL: Search_Results/HotpotQA benchmark
752 +L2: # Search Results
753 +L3:
754 +L4: * 【0†HotpotQA Homepage; publish_date: none†hotpotqa.github.io】 HotpotQA is a
755 +L5: question answering dataset featuring natural, multi-hop questions, with strong
756 +L6: supervision for supporting facts to enable more explainable ...See more
757 +L7: * 【1†HotpotQA: A Dataset for Diverse, Explainable Multi-hop ...; publish_date:
758 +L8: none†arxiv.org】 by Z Yang · 2018 · Cited by 3834 — HotpotQA is a dataset with
759 +L9: 113k Wikipedia-based question-answer pairs requiring multi-document reasoning,
760 +L10: diverse questions, sentence-level ...
761 +L11: * 【2†hotpotqa/hotpot_qa · Datasets at Hugging Face; publish_date:
762 +L12: none†huggingface.co】 HotpotQA is a new dataset with 113k Wikipedia-based
763 +L13: question-answer pairs with four key features: (1) the questions require finding
764 +L14: and reasoning over multiple ...See more
765 +L15: * 【3†Why You Should Stop Using HotpotQA for AI Agents ...; publish_date:
766 +L16: none†qipeng.me】 Jul 1, 2025 — HotpotQA pioneered a class of AI tasks that
767 +L17: requires the AI system to autonomously perform multiple steps of reasoning in an
768 +L18: open-domain setting.See more
769 +L19: * 【4†hotpotqa/hotpot; publish_date: none†github.com】 A dataset for diverse,
770 +L20: explainable multi-hop question answering. This repository contains the baseline
771 +L21: model code, as well as the entire pipeline of running ...See more
772 +L22: * 【5†HotpotQA: Multi-Hop QA Benchmark; publish_date:
773 +L23: none†www.emergentmind.com】 Sep 10, 2025 — HotpotQA is a large-scale multi-hop
774 +L24: question answering benchmark featuring 112,779 Wikipedia-based Q&A pairs with
775 +L25: detailed, sentence-level ...See more
776 +L26: * 【6†HotpotQA Dataset | Papers With Code; publish_date:
777 +L27: none†paperswithcode.com】 HotpotQA is a question answering dataset collected on
778 +L28: the English Wikipedia, containing about 113K crowd-sourced questions.See more
779 +L29: * 【7†HotpotQA: A Dataset for Diverse, Explainable Multi-hop ...; publish_date:
780 +L30: none†aclanthology.org】 by Z Yang · 2018 · Cited by 3834 — HotpotQA is a dataset
781 +L31: with 113k Wikipedia-based question-answer pairs requiring multi-document
782 +L32: reasoning, diverse questions, sentence-level facts, and factoid ...
783 +L33: * 【8†Benchmark BM25S: HotpotQA; publish_date: none†www.kaggle.com】 Explore and
784 +L34: run machine learning code with Kaggle Notebooks | Using data from No attached
785 +L35: data sources.
786 +L36: * 【9†mteb/hotpotqa · Datasets at Hugging Face; publish_date:
787 +L37: none†huggingface.co】 HotpotQA is a question answering dataset featuring natural,
788 +L38: multi-hop questions, with strong supervision for supporting facts to enable
789 +L39: more explainable ...See more</result>
790 +<result>[9] FEVER benchmark (Search_Results/FEVER benchmark)
791 +**viewing lines [0 - 40] of 40**
792 +
793 +L0:
794 +L1: URL: Search_Results/FEVER benchmark
795 +L2: # Search Results
796 +L3:
797 +L4: * 【0†Fever.ai; publish_date: none†fever.ai】 We are pleased to announce that
798 +L5: FEVER9 will be co-located with EACL 2026. In this year's workshop, we will
799 +L6: introduce a new shared task focused on automated fact ...
800 +L7: * 【1†a Large-scale Dataset for Fact Extraction and VERification; publish_date:
801 +L8: none†aclanthology.org】 by J Thorne · 2018 · Cited by 2315 — In this paper we
802 +L9: introduce a new publicly available dataset for verification against textual
803 +L10: sources, FEVER: Fact Extraction.
804 +L11: * 【2†awslabs/fever: FEVER (Fact Extraction and VERification) ...;
805 +L12: publish_date: none†github.com】 In this paper we introduce a new publicly
806 +L13: available dataset for verification against textual sources, FEVER: Fact
807 +L14: Extraction and VERification.
808 +L15: * 【3†FEVER: Fact Extraction and VERification; publish_date:
809 +L16: none†www.amazon.science】 The best accuracy we achieve on labeling a claim
810 +L17: accompanied by the correct evidence is 31.87%, while if we ignore the evidence
811 +L18: we achieve 50.91%. Thus we ...
812 +L19: * 【4†FEVER Dataset; publish_date: none†fever.ai】 FEVER (Fact Extraction and
813 +L20: VERification) consists of 185,445 claims generated by altering sentences
814 +L21: extracted from Wikipedia and subsequently verified ...
815 +L22: * 【5†mteb/fever · Datasets at Hugging Face; publish_date: none†huggingface.co】
816 +L23: FEVER. An MTEB dataset. Massive Text Embedding Benchmark. FEVER (Fact
817 +L24: Extraction and VERification) consists of 185,445 claims generated by altering
818 +L25: sentences ...
819 +L26: * 【6†FEVEROUS: Fact Extraction and VERification Over ...; publish_date:
820 +L27: none†datasets-benchmarks-proceedings.neurips.cc】 by R Aly · Cited by 359 — In
821 +L28: this paper we introduce a novel dataset and benchmark, Fact Extraction and
822 +L29: VERification Over. Unstructured and Structured information (FEVEROUS), which ...
823 +L30: * 【7†a large-scale dataset for Fact Extraction and VERification; publish_date:
824 +L31: none†arxiv.org】 by J Thorne · 2018 · Cited by 2315 — In this paper we introduce
825 +L32: a new publicly available dataset for verification against textual sources,
826 +L33: FEVER: Fact Extraction and VERification.
827 +L34: * 【8†FEVER Resources; publish_date: none†fever.ai】 2018 FEVER: a large-scale
828 +L35: dataset for Fact Extraction and VERification .bib James Thorne, Andreas Vlachos,
829 +L36: Christos Christodoulopoulos, Arpit Mittal
830 +L37: * 【9†a Large-scale Dataset for Fact Extraction and VERification; publish_date:
831 +L38: none†www.semanticscholar.org】 This paper introduces a new publicly available
832 +L39: dataset for verification against textual sources, FEVER, which consists of
833 +L40: 185,445 claims generated by ...</result>
834 +<result>[10] TriviaQA benchmark (Search_Results/TriviaQA benchmark)
835 +**viewing lines [0 - 35] of 35**
836 +
837 +L0:
838 +L1: URL: Search_Results/TriviaQA benchmark
839 +L2: # Search Results
840 +L3:
841 +L4: * 【0†TriviaQA; publish_date: none†nlp.cs.washington.edu】 TriviaQA is a reading
842 +L5: comprehension dataset containing over 650K question-answer-evidence triples.
843 +L6: TriviaQA includes 95K question-answer pairs authored ...
844 +L7: * 【1†TriviaQA: A Large Scale Distantly Supervised Challenge ...; publish_date:
845 +L8: none†aclanthology.org】 by M Joshi · 2017 · Cited by 3451 — We present TriviaQA,
846 +L9: a challenging reading comprehension dataset containing over 650K question-
847 +L10: answer-evidence triples. TriviaQA includes 95K question ...
848 +L11: * 【2†mandarjoshi/trivia_qa · Datasets at Hugging Face; publish_date:
849 +L12: none†huggingface.co】 TriviaqQA is a reading comprehension dataset containing
850 +L13: over 650K question-answer-evidence triples. TriviaqQA includes 95K question-
851 +L14: answer pairs authored by ...
852 +L15: * 【3†[1705.03551] TriviaQA: A Large Scale Distantly Supervised ...;
853 +L16: publish_date: none†arxiv.org】 by M Joshi · 2017 · Cited by 3451 — We present
854 +L17: TriviaQA, a challenging reading comprehension dataset containing over 650K
855 +L18: question-answer-evidence triples.
856 +L19: * 【4†TriviaQA; publish_date: none†epoch.ai】 An open-domain question answering
857 +L20: benchmark with challenging trivia questions paired with evidence documents.
858 +L21: * 【5†TriviaQA Leaderboard; publish_date: none†llm-stats.com】 What is the
859 +L22: TriviaQA benchmark? A large-scale reading comprehension dataset containing over
860 +L23: 650K question-answer-evidence triples. TriviaQA includes 95K ...
861 +L24: * 【6†Code for the TriviaQA reading comprehension dataset; publish_date:
862 +L25: none†github.com】 A large scale distantly supervised challenge dataset for
863 +L26: reading comprehension. In Association for Computational Linguistics (ACL) 2017,
864 +L27: Vancouver, Canada.
865 +L28: * 【7†TriviaQA - Model Benchmarks - The Regularizer; publish_date:
866 +L29: none†www.theregularizer.com】 May 4, 2025 — Compare the performance of different
867 +L30: AI models across standardized benchmarks. Higher scores generally indicate
868 +L31: better performance, but context ...
869 +L32: * 【8†TriviaQA: A Large Scale Distantly Supervised Challenge ...; publish_date:
870 +L33: none†www.cs.utexas.edu】 by M Joshi · Cited by 3445 — We present TriviaQA, a
871 +L34: challenging reading comprehension dataset contain- ing over 650K question-
872 +L35: answer-evidence triples. TriviaQA includes 95K question-.</result>
873 +<result>[11] Natural Questions benchmark (Search_Results/Natural Questions benchmark)
874 +**viewing lines [0 - 39] of 39**
875 +
876 +L0:
877 +L1: URL: Search_Results/Natural Questions benchmark
878 +L2: # Search Results
879 +L3:
880 +L4: * 【0†Natural Questions: a Benchmark for Question Answering ...; publish_date:
881 +L5: none†research.google】 by T Kwiatkowski · Cited by 4339 — We present the Natural
882 +L6: Questions corpus, a question answering dataset. Questions consist of real
883 +L7: anonymized, aggregated queries issued to the Google search ...
884 +L8: * 【1†Natural Questions: A Benchmark for Question Answering ...; publish_date:
885 +L9: none†aclanthology.org】 by T Kwiatkowski · Cited by 4308 — Abstract. We present
886 +L10: the Natural Questions corpus, a question answering data set. Questions consist
887 +L11: of real anonymized, aggregated queries issued.
888 +L12: * 【2†Google's Natural Questions; publish_date: none†ai.google.com】 Natural
889 +L13: Questions. A Benchmark for Question Answering Research. View examples · Download
890 +L14: dataset. Open Domain Question Answering. A core goal in artificial ...
891 +L15: * 【3†google-research-datasets/natural-questions; publish_date:
892 +L16: none†github.com】 Natural Questions (NQ) contains real user questions issued to
893 +L17: Google search, and answers found from Wikipedia by annotators. NQ is designed
894 +L18: for the training and ...
895 +L19: * 【4†Natural Questions: A Benchmark for Question Answering ...; publish_date:
896 +L20: none†direct.mit.edu】 Aug 1, 2019 — We present the Natural Questions corpus, a
897 +L21: question answering data set. Questions consist of real anonymized, aggregated
898 +L22: queries issued to the Google search ...
899 +L23: * 【5†ir_datasets : Natural Questions; publish_date: none†ir-datasets.com】
900 +L24: Google Natural Questions is a Q&A dataset containing long, short, and Yes/No
901 +L25: answers from Wikipedia. ir_datasets frames this around an ad-hoc ranking setting
902 +L26: ...
903 +L27: * 【6†sentence-transformers/natural-questions · Datasets at ...; publish_date:
904 +L28: none†huggingface.co】 This dataset is a collection of question-answer pairs from
905 +L29: the Natural Questions dataset. See Natural Questions for additional information.
906 +L30: * 【7†Google's Natural Questions; publish_date: none†ai.google.com】 Natural
907 +L31: Questions contains 307K training examples, 8K examples for development, and a
908 +L32: further 8K examples for testing. In the paper, we demonstrate a human ...
909 +L33: * 【8†A Benchmark for Question Answering Research; publish_date:
910 +L34: none†www.researchgate.net】 Jul 27, 2025 — We present the Natural Questions
911 +L35: corpus, a question answering data set. Questions consist of real anonymized,
912 +L36: aggregated queries issued to the Google search ...
913 +L37: * 【9†natural-questions; publish_date: none†docs.unity.rc.umass.edu】 Sep 4,
914 +L38: 2025 — “Natural questions: a benchmark for question answering research.”
915 +L39: Transactions of the Association for Computational Linguistics 7 (2019): ...</result>
916 +<result>[12] MS MARCO benchmark (Search_Results/MS MARCO benchmark)
917 +**viewing lines [0 - 41] of 41**
918 +
919 +L0:
920 +L1: URL: Search_Results/MS MARCO benchmark
921 +L2: # Search Results
922 +L3:
923 +L4: * 【0†MS MARCO - Microsoft Open Source; publish_date: none†microsoft.github.io】
924 +L5: The MS MARCO datasets are intended for non-commercial research purposes only to
925 +L6: promote advancement in the field of artificial intelligence and related areas,
926 +L7: ...
927 +L8: * 【1†microsoft/ms_marco · Datasets at Hugging Face; publish_date:
928 +L9: none†huggingface.co】 Starting with a paper released at NIPS 2016, MS MARCO is a
929 +L10: collection of datasets focused on deep learning in search. The first dataset was
930 +L11: a question ...
931 +L12: * 【2†Benchmarking Ranking Models in the Large-Data Regime; publish_date:
932 +L13: none†arxiv.org】 by N Craswell · 2021 · Cited by 89 — This paper uses the MS
933 +L14: MARCO and TREC Deep Learning Track as our case study, comparing it to the case
934 +L15: of TREC ad hoc ranking in the 1990s.
935 +L16: * 【3†Benchmarking Ranking Models in the Large-Data Regime; publish_date:
936 +L17: none†www.microsoft.com】 This paper uses the MS MARCO and TREC Deep Learning
937 +L18: Track as our case study, comparing it to the case of TREC ad hoc ranking in the
938 +L19: 1990s. We show how the ...
939 +L20: * 【4†Datasets for Document and Passage Ranking Leadboards; publish_date:
940 +L21: none†microsoft.github.io】 The MS MARCO document and passage ranking leaderboards
941 +L22: complements the TREC Deep Learning Track by providing on-going evaluation of
942 +L23: submissions using pre- ...
943 +L24: * 【5†MS MARCO: Benchmarking Ranking Models in the Large- ...; publish_date:
944 +L25: none†dl.acm.org】 Jul 11, 2021 — This paper uses the MS MARCO and TREC Deep
945 +L26: Learning Track as our case study, comparing it to the case of TREC ad hoc
946 +L27: ranking in the 1990s.
947 +L28: * 【6†ir_datasets : MSMARCO (passage); publish_date: none†ir-datasets.com】 A
948 +L29: passage ranking benchmark with a collection of 8.8 million passages and question
949 +L30: queries. Most relevance judgments are shallow.
950 +L31: * 【7†MS MARCO; publish_date: none†sbert.net】 MS MARCO Passage Ranking is a
951 +L32: large dataset to train models for information retrieval. It consists of about
952 +L33: 500k real search queries from Bing search engine ...
953 +L34: * 【8†MS MARCO: A Human Generated MAchine Reading ...; publish_date:
954 +L35: none†arxiv.org】 by P Bajaj · 2016 · Cited by 1151 — We introduce a large scale
955 +L36: MAchine Reading COmprehension dataset, which we name MS MARCO. The dataset
956 +L37: comprises of 1,010,916 anonymized ...
957 +L38: * 【9†MS MARCO Web Search: A Large-scale Information-rich ...; publish_date:
958 +L39: none†www.microsoft.com】 May 13, 2024 — MS MARCO Web Search offers a retrieval
959 +L40: benchmark with three web retrieval challenge tasks that demands innovations in
960 +L41: both machine learning and ...</result>
961 +<result>[13] BEIR benchmark (Search_Results/BEIR benchmark)
962 +**viewing lines [0 - 37] of 37**
963 +
964 +L0:
965 +L1: URL: Search_Results/BEIR benchmark
966 +L2: # Search Results
967 +L3:
968 +L4: * 【0†详细介绍文本检索基准BEIR: A Heterogeneous Benchmark ...; publish_date:
969 +L5: none†blog.csdn.net】 2023年1月1日 —
970 +L6: BEIR旨在为所有不同的检索任务提供一个一站式的零样本评估基准。为了构建一个全面的评估基准，选择方法对于收集具有理想属性的任务和数据集至关重要。对于 ...
971 +L7: * 【1†beir-cellar/beir; publish_date: none†github.com】 BEIR is a heterogeneous
972 +L8: benchmark containing diverse IR tasks. It also provides a common and easy
973 +L9: framework for evaluation of your NLP-based retrieval models ...
974 +L10: * 【2†BEIR: A Heterogenous Benchmark for Zero-shot Evaluation ...;
975 +L11: publish_date: none†arxiv.org】 作者：N Thakur · 2021 · 被引用次数：1480 — We introduce
976 +L12: Benchmarking-IR (BEIR), a robust and heterogeneous evaluation benchmark for
977 +L13: information retrieval.
978 +L14: * 【3†BeIR; publish_date: none†huggingface.co】 BEIR (Benchmarking IR) consists
979 +L15: of a homogenous benchmark for diverse sentence or passage level IR tasks. It
980 +L16: provides a common and easy framework for the cross ...
981 +L17: * 【4†论文分享：BEIR A Heterogeneous Benchmark for Zero-shot ...; publish_date:
982 +L18: none†zhuanlan.zhihu.com】 2022年10月3日 — 分享论文，夹带个人理解的分享，建议结合原论文看。 1 研究背景.
983 +L19: 本论文主要关注的领域是query-document检索（下文简称qd检索），即根据query去文档库里 ...
984 +L20: * 【5†Benchmarking IR Information Retrieval (BEIR); publish_date:
985 +L21: none†zilliz.com】 BEIR is a benchmark designed for evaluating the versatility and
986 +L22: robustness of information retrieval models. It features 18 diverse datasets
987 +L23: from domains like ...
988 +L24: * 【6†BEIR (Benchmarking IR) - OpenDataLab; publish_date: none†opendatalab.com】
989 +L25: 简介-Introduction. BEIR（Benchmarking
990 +L26: IR）是包含不同信息检索（IR）任务的异构基准。通过BEIR，可以系统地研究多种神经检索方法的零样本泛化能力。
991 +L27: * 【7†What is the BEIR benchmark and how is it used?; publish_date:
992 +L28: none†milvus.io】 The BEIR (Benchmarking Information Retrieval) benchmark is a
993 +L29: standardized framework designed to evaluate the effectiveness of search and
994 +L30: retrieval algorithms.
995 +L31: * 【8†BEIR Benchmark数据集卡片; publish_date: none†www.atyun.com】 BEIR
996 +L32: Benchmark数据集卡片. 数据集简介. BEIR是一个异构评测基准，由18个多样化的数据集构建而成，代表了9个信息检索任务：. 事实查证： FEVER ，
997 +L33: Climate-FEVER ， SciFact ...
998 +L34: * 【9†Evaluating search relevance part 1 - The BEIR benchmark; publish_date:
999 +L35: none†www.elastic.co】 2024年7月16日 — Learn to evaluate your search system in the
1000 +L36: context of better understanding the BEIR benchmark, with tips & techniques to
1001 +L37: improve your ...</result>
1002 +<result>[14] MIRACL benchmark (Search_Results/MIRACL benchmark)
1003 +**viewing lines [0 - 41] of 41**
1004 +
1005 +L0:
1006 +L1: URL: Search_Results/MIRACL benchmark
1007 +L2: # Search Results
1008 +L3:
1009 +L4: * 【0†MIRACL | Multilingual Information Retrieval Across a ...; publish_date:
1010 +L5: none†project-miracl.github.io】 MIRACL (Multilingual Information Retrieval Across
1011 +L6: a Continuum of Languages) is an WSDM 2023 Cup challenge that focuses on search
1012 +L7: across 18 different ...
1013 +L8: * 【1†project-miracl/miracl: A large-scale multilingual dataset for ...;
1014 +L9: publish_date: none†github.com】 A large-scale multilingual dataset for
1015 +L10: Information Retrieval. Thorough human-annotations across 18 diverse languages.
1016 +L11: * 【2†A Large, multilingual, visual document retrieval benchmark; publish_date:
1017 +L12: none†arxiv.org】 by R Osmulski · 2025 · Cited by 2 — MIRACL-VISION is a
1018 +L13: challenging, representative, multilingual evaluation benchmark for visual
1019 +L14: retrieval pipelines and will help the community build robust ...
1020 +L15: * 【3†miracl/miracl · Datasets at Hugging Face; publish_date:
1021 +L16: none†huggingface.co】 MIRACL (Multilingual Information Retrieval Across a
1022 +L17: Continuum of Languages) is a multilingual retrieval dataset that focuses on
1023 +L18: search across 18 different ...
1024 +L19: * 【4†MIRACL: A Multilingual Retrieval Dataset Covering 18 ...; publish_date:
1025 +L20: none†direct.mit.edu】 by X Zhang · 2023 · Cited by 131 — MIRACL is a multilingual
1026 +L21: dataset for ad hoc retrieval across 18 languages that collectively encompass
1027 +L22: over three billion native speakers around the world.
1028 +L23: * 【5†(PDF) MIRACL-VISION: A Large, multilingual, visual ...; publish_date:
1029 +L24: none†www.researchgate.net】 May 23, 2025 — MIRACL-VISION covers 18 languages, and
1030 +L25: is an extension of the MIRACL dataset, a popular benchmark to evaluate text-
1031 +L26: based multilingual retrieval ...
1032 +L27: * 【6†A Large, multilingual, visual document retrieval benchmark; publish_date:
1033 +L28: none†arxiv.org】 by R Osmulski · 2025 · Cited by 2 — MIRACL-VISION is a
1034 +L29: challenging, representative, multilingual evaluation benchmark for visual
1035 +L30: retrieval pipelines and will help the community ...
1036 +L31: * 【7†ir_datasets : MIRACL; publish_date: none†ir-datasets.com】
1037 +L32: "miracl/ar/test-a". The held-out test set (version a) for Arabic.
1038 +L33: queriesdocsCitationMetadata. 936 queries. Language: ar. Query type:
1039 +L34: GenericQuery: (namedtuple).
1040 +L35: * 【8†Evaluate on MIRACL — BGE documentation; publish_date: none†bge-model.com】
1041 +L36: MIRACL (Multilingual Information Retrieval Across a Continuum of Languages) is
1042 +L37: an WSDM 2023 Cup challenge that focuses on search across 18 different languages.
1043 +L38: * 【9†MIRACL - Alpha's Tech Garden; publish_date:
1044 +L39: none†techgarden.alphasmanifesto.com】 MIRACL (Multilingual Information Retrieval
1045 +L40: Across a Continuum of Languages) is a multilingual dataset we have built for the
1046 +L41: WSDM 2023 Cup ...</result>
1047 +<result>[15] Zero-shot IR benchmark (Search_Results/Zero-shot IR benchmark)
1048 +**viewing lines [0 - 40] of 40**
1049 +
1050 +L0:
1051 +L1: URL: Search_Results/Zero-shot IR benchmark
1052 +L2: # Search Results
1053 +L3:
1054 +L4: * 【0†BEIR: A Heterogenous Benchmark for Zero-shot Evaluation ...;
1055 +L5: publish_date: none†arxiv.org】 by N Thakur · 2021 · Cited by 1480 — We introduce
1056 +L6: Benchmarking-IR (BEIR), a robust and heterogeneous evaluation benchmark for
1057 +L7: information retrieval.See more
1058 +L8: * 【1†beir-cellar/beir; publish_date: none†github.com】 BEIR: A Heterogenous
1059 +L9: Benchmark for Zero-shot Evaluation of Information Retrieval Models (NeurIPS
1060 +L10: 2021, Datasets and Benchmarks Track); Resources for Brewing ...See more
1061 +L11: * 【2†Benchmarking IR Information Retrieval (BEIR); publish_date:
1062 +L12: none†zilliz.com】 BEIR is a tool to evaluate how well Information Retrieval
1063 +L13: systems perform across many tasks and types of information, and is a standard
1064 +L14: benchmark.
1065 +L15: * 【3†BEIR: A Heterogeneous Benchmark for Zero-shot ...; publish_date:
1066 +L16: none†datasets-benchmarks-proceedings.neurips.cc】 by N Thakur · Cited by 1480 —
1067 +L17: BEIR is a robust, heterogeneous benchmark for information retrieval, using 18
1068 +L18: datasets and 9 tasks to evaluate model generalization.
1069 +L19: * 【4†BEIR; publish_date: none†eval.ai】 BEIR is a heterogeneous zero-shot
1070 +L20: retrieval benchmark containing 18 datasets from diverse text retrieval tasks and
1071 +L21: domains.See more
1072 +L22: * 【5†[2409.15763] IRSC: A Zero-shot Evaluation Benchmark for ...;
1073 +L23: publish_date: none†arxiv.org】 by H Lin · 2024 · Cited by 2 — This paper
1074 +L24: introduces the IRSC benchmark for evaluating the performance of embedding models
1075 +L25: in multilingual RAG tasks.See more
1076 +L26: * 【6†FactIR: A Real-World Zero-shot Open-Domain Retrieval ...; publish_date:
1077 +L27: none†dl.acm.org】 May 23, 2025 — In this paper, we present a real-world retrieval
1078 +L28: benchmark FactIR, derived from Factiverse production logs, enhanced with human
1079 +L29: annotations. We ...See more
1080 +L30: * 【7†UniIR: Training and Benchmarking Universal Multimodal ...; publish_date:
1081 +L31: none†tiger-ai-lab.github.io】 At test time, we evaluated the zero-shot
1082 +L32: performance of all fine-tuned models, as well as SoTA pre-trained retrievers on
1083 +L33: the three held-out datasets. UniIR ...See more
1084 +L34: * 【8†Zero-Shot BEIR Tasks; publish_date: none†www.emergentmind.com】 Aug 26,
1085 +L35: 2025 — Zero-Shot BEIR Tasks are evaluation methodologies that assess IR models'
1086 +L36: ability to generalize to unseen query domains without task-specific ...See more
1087 +L37: * 【9†BEIR-PL: Zero Shot Information Retrieval Benchmark for ...; publish_date:
1088 +L38: none†aclanthology.org】 by K Wojtasik · 2024 · Cited by 12 — BEIR-PL is a new
1089 +L39: benchmark with 13 datasets for Polish Information Retrieval, created to advance
1090 +L40: research in this area.</result>
1091 +<result>[16] WebGPT benchmark (Search_Results/WebGPT benchmark)
1092 +**viewing lines [0 - 38] of 38**
1093 +
1094 +L0:
1095 +L1: URL: Search_Results/WebGPT benchmark
1096 +L2: # Search Results
1097 +L3:
1098 +L4: * 【0†WebGPT: Improving the factual accuracy of language ...; publish_date:
1099 +L5: none†openai.com】 Dec 16, 2021 — Our models outperform GPT‑3 on TruthfulQA and
1100 +L6: exhibit more favourable scaling properties. However, our models lag behind human
1101 +L7: performance, ...
1102 +L8: * 【1†A Simple Yet Challenging Benchmark for Browsing Agents; publish_date:
1103 +L9: none†arxiv.org】 by J Wei · 2025 · Cited by 124 — Abstract. We present
1104 +L10: BrowseComp, a simple yet challenging benchmark for measuring the ability for
1105 +L11: agents to browse the web.
1106 +L12: * 【2†openai/webgpt_comparisons · Datasets at Hugging Face; publish_date:
1107 +L13: none†huggingface.co】 This is the dataset of all comparisons that were marked as
1108 +L14: suitable for reward modeling by the end of the WebGPT project. There are 19,578
1109 +L15: comparisons in total.
1110 +L16: * 【3†Evaluation & Limitations of WebGPT, WebVoyager & Agent-E; publish_date:
1111 +L17: none†deepsense.ai】 Oct 14, 2024 — WebArena benchmark features 812 tasks
1112 +L18: evaluated using metrics such as Exact Match, Must Include, and Fuzzy Match,
1113 +L19: focusing on outcomes rather ...
1114 +L20: * 【4†OpenAI Announces Question-Answering AI WebGPT; publish_date:
1115 +L21: none†www.infoq.com】 Jan 25, 2022 — On the TriviaQA benchmark, WebGPT
1116 +L22: outperformed GPT-3, producing answers that were true 75% of the time, and "both
1117 +L23: true and informative" 54% of ...
1118 +L24: * 【5†WebGPT: Improving the factual accuracy of language models ...;
1119 +L25: publish_date: none†kargarisaac.medium.com】 The top-performing model generated
1120 +L26: answers that were preferred over 56% of the time compared to answers produced by
1121 +L27: human demonstrators, with ...
1122 +L28: * 【6†Browser-assisted question-answering with human feedback; publish_date:
1123 +L29: none†www.alphaxiv.org】 WebGPT represents a significant advancement in long-form
1124 +L30: question answering by combining the language generation capabilities of GPT-3
1125 +L31: with real-time web ...
1126 +L32: * 【7†Benchmarking Open-Source Large Language Models, GPT-4 ...; publish_date:
1127 +L33: none†ai.nejm.org】 by S Wu · 2024 · Cited by 69 — We show that the current widely
1128 +L34: used open-source LLMs have poor zero-shot reasoning ability in nephrology
1129 +L35: compared with GPT-4 and Claude 2.
1130 +L36: * 【8†0hq/WebGPT: Run GPT model on ...; publish_date: none†github.com】 WebGPT
1131 +L37: is a vanilla JS and HTML implementation of a transformer model, intended as a
1132 +L38: proof-of-concept as well as educational resource.</result>
1133 +<result>[17] WebShop benchmark (Search_Results/WebShop benchmark)
1134 +**viewing lines [0 - 41] of 41**
1135 +
1136 +L0:
1137 +L1: URL: Search_Results/WebShop benchmark
1138 +L2: # Search Results
1139 +L3:
1140 +L4: * 【0†WebShop: Towards Scalable Real-World Web Interaction ...; publish_date:
1141 +L5: none†arxiv.org】 by S Yao · 2022 · Cited by 710 — To bridge this gap, we develop
1142 +L6: WebShop -- a simulated e-commerce website environment with 1.18 million real-
1143 +L7: world products and 12,087 crowd- ...
1144 +L8: * 【1†WebShop; publish_date: none†webshop-pnlp.github.io】 To bridge this gap,
1145 +L9: we develop WebShop – a simulated e-commerce website environment with 1.18
1146 +L10: million real-world products and 12,087 crowd-sourced text ...
1147 +L11: * 【2†princeton-nlp/WebShop; publish_date: none†github.com】 WebShop is a
1148 +L12: simulated e-commerce website environment with 1.18 million real-world products
1149 +L13: and 12,087 crowd-sourced text instructions. In this environment, an ...
1150 +L14: * 【3†WebShop: Towards Scalable Real-World Web Interaction ...; publish_date:
1151 +L15: none†papers.nips.cc】 by S Yao · 2022 · Cited by 710 — We collect over 1,600
1152 +L16: human trajectories to first validate the benchmark, then train and evaluate a
1153 +L17: diverse range of agents using reinforcement learning, ...
1154 +L18: * 【4†WebShop: Towards Scalable Real-World Web Interaction ...; publish_date:
1155 +L19: none†proceedings.neurips.cc】 by S Yao · 2022 · Cited by 709 — We have developed
1156 +L20: WebShop, a new web-based benchmark for sequential decision making and language
1157 +L21: grounding, modeled on interaction with an e-commerce website.
1158 +L22: * 【5†Webshop & Benchmark Analysis | Documentation Infinity; publish_date:
1159 +L23: none†docs.fact-finder.com】 Aug 15, 2025 — Evaluation of your shop based on
1160 +L24: different categories in comparison, to your competitors/industry. Recommended
1161 +L25: when doing a shop relaunch.
1162 +L26: * 【6†A Multi-Shop Benchmark for Evaluating Web Agents; publish_date:
1163 +L27: none†arxiv.org】 by R Peeters · 2025 · Cited by 2 — Compared to existing
1164 +L28: e-commerce benchmarks, such as WebShop or ShoppingBench, WebMall introduces
1165 +L29: comparison-shopping tasks across multiple shops ...
1166 +L30: * 【7†WebShop: towards scalable real-world web interaction with ...;
1167 +L31: publish_date: none†dl.acm.org】 by S Yao · 2022 · Cited by 710 — To bridge this
1168 +L32: gap, we develop WebShop - a simulated e-commerce website environment with 1.18
1169 +L33: million real-world products and 12, 087 crowd- ...
1170 +L34: * 【8†[PDF] WebShop: Towards Scalable Real-World Web ...; publish_date:
1171 +L35: none†www.semanticscholar.org】 It is shown that agents trained on WebShop exhibit
1172 +L36: non-trivial sim-to-real transfer when evaluated on amazon.com and ebay.com,
1173 +L37: indicating the potential ...
1174 +L38: * 【9†X-WebAgentBench: A Multilingual Interactive Web ...; publish_date:
1175 +L39: none†aclanthology.org】 by P Wang · 2025 · Cited by 3 — (2023) based on the
1176 +L40: English WebShop benchmark (Yao et al., 2022), while the multilingual task scores
1177 +L41: are ob- tained through evaluation on our own benchmark.</result>
1178 +<result>[18] ALFWorld benchmark (Search_Results/ALFWorld benchmark)
1179 +**viewing lines [0 - 31] of 31**
1180 +
1181 +L0:
1182 +L1: URL: Search_Results/ALFWorld benchmark
1183 +L2: # Search Results
1184 +L3:
1185 +L4: * 【0†ALFWorld; publish_date: none†alfworld.github.io】 ALFWorld contains
1186 +L5: interactive TextWorld environments (Côté et. al) that parallel embodied worlds
1187 +L6: in the ALFRED dataset (Shridhar et. al).
1188 +L7: * 【1†ALFWorld: Aligning Text and Embodied Environments for ...; publish_date:
1189 +L8: none†arxiv.org】 by M Shridhar · 2020 · Cited by 674 — ALFWorld enables the
1190 +L9: creation of a new BUTLER agent whose abstract knowledge, learned in TextWorld,
1191 +L10: corresponds directly to concrete, visually grounded actions.
1192 +L11: * 【2†ALFWorld: Aligning Text and Embodied Environments ...; publish_date:
1193 +L12: none†github.com】 ALFWorld contains interactive TextWorld environments (Côté et.
1194 +L13: al) that parallel embodied worlds in the ALFRED dataset (Shridhar et. al).
1195 +L14: * 【3†alfworld - benchmark's activity; publish_date: none†huggingface.co】 MM-
1196 +L15: IQ: Benchmarking Human-Like Abstraction and Reasoning in Multimodal Models Paper
1197 +L16: • 2502.00698 • Published Feb 1 • 24
1198 +L17: * 【4†Tackling AlfWorld with Action Attention and Common ...; publish_date:
1199 +L18: none†neurips.cc】 On the Alfworld benchmark for indoor instruction following, we
1200 +L19: achieve a significantly higher success rate (50% over the baseline) with our
1201 +L20: novel object ...
1202 +L21: * 【5†ALFWORLD: ALIGNING TEXT AND EMBODIED ...; publish_date:
1203 +L22: none†openreview.net】 by M Shridhar · Cited by 674 — The ALFRED dataset (Shridhar
1204 +L23: et al., 2020), set in the THOR simulator (Kolve et al., 2017), is a benchmark
1205 +L24: for learning to com- plete embodied household tasks ...
1206 +L25: * 【6†AlfWorld; publish_date: none†primo.ai】 Mar 23, 2024 — A simulator that
1207 +L26: enables agents to learn abstract, text based policies in TextWorld (Côté et al.,
1208 +L27: 2018) and then execute goals from the ALFRED benchmark.
1209 +L28: * 【7†AlfWorld performance across 134 tasks showing cumulative...;
1210 +L29: publish_date: none†www.researchgate.net】 In the AlfWorld benchmark, we defined
1211 +L30: hallucination as the occurrence of two or more consecutive identical actions in
1212 +L31: which the environment responded with ...</result>
1213 +<result>[19] Mind2Web benchmark (Search_Results/Mind2Web benchmark)
1214 +**viewing lines [0 - 40] of 40**
1215 +
1216 +L0:
1217 +L1: URL: Search_Results/Mind2Web benchmark
1218 +L2: # Search Results
1219 +L3:
1220 +L4: * 【0†Mind2Web: Towards a Generalist Agent for the Web; publish_date: none†osu-
1221 +L5: nlp-group.github.io】 Mind2Web is a dataset for developing and evaluating
1222 +L6: generalist agents for the web that can follow language instructions to complete
1223 +L7: complex tasks on any ...
1224 +L8: * 【1†Online-Mind2Web Leaderboard; publish_date: none†huggingface.co】 Online-
1225 +L9: Mind2Web is a benchmark designed to evaluate the real-world performance of web
1226 +L10: agents on live websites, featuring 300 tasks across 136 popular sites ...
1227 +L11: * 【2†Mind2Web: Towards a Generalist Agent for the Web; publish_date:
1228 +L12: none†github.com】 Mind2Web is the first dataset for developing and evaluating
1229 +L13: generalist agents for the web that can follow language instructions to complete
1230 +L14: complex tasks on any ...
1231 +L15: * 【3†HAL: Online Mind2Web Leaderboard; publish_date:
1232 +L16: none†hal.cs.princeton.edu】 Online Mind2Web leaderboard for evaluating AI agents'
1233 +L17: ability to complete tasks on real, changing webpages.
1234 +L18: * 【4†[2506.21506] Mind2Web 2: Evaluating Agentic Search with ...;
1235 +L19: publish_date: none†arxiv.org】 by B Gou · 2025 · Cited by 11 — In this paper, we
1236 +L20: introduce Mind2Web 2, a benchmark of 130 realistic, high-quality, and long-
1237 +L21: horizon tasks that require real-time web browsing and extensive ...
1238 +L22: * 【5†Mind2Web 2: Evaluating Agentic Search with Agent-as-a-Judge;
1239 +L23: publish_date: none†osu-nlp-group.github.io】 We introduce Mind2Web 2, a benchmark
1240 +L24: of 130 realistic, high-quality, long-horizon tasks that require real-time web
1241 +L25: browsing and extensive information ...
1242 +L26: * 【6†Mind2Web: The Benchmark for AI Agent Evaluation and ...; publish_date:
1243 +L27: none†www.enhans.ai】 Sep 26, 2025 — Mind2Web is a globally recognized web-based
1244 +L28: AI Agent evaluation benchmark introduced by the NLP group at Ohio State
1245 +L29: University at NeurIPS 2023.
1246 +L30: * 【7†Evaluating AI Web Agents: Insights from the WebCanvas ...; publish_date:
1247 +L31: none†medium.com】 Thanks to the comprehensive WebCanvas Benchmark, which
1248 +L32: incorporates a robust Mind2Web-Live data set of 542 live web tasks and 2,439 ...
1249 +L33: * 【8†Mind2Web: Towards a Generalist Agent for the Web; publish_date:
1250 +L34: none†proceedings.neurips.cc】 by X Deng · 2023 · Cited by 760 — We introduce
1251 +L35: Mind2Web, the first dataset for developing and evaluating generalist agents for
1252 +L36: the web that can follow language instructions to complete complex ...
1253 +L37: * 【9†Mind2Web: Towards a Generalist Agent for the Web; publish_date:
1254 +L38: none†openreview.net】 by X Deng · Cited by 760 — We introduce Mind2Web, the first
1255 +L39: dataset for developing and evaluating generalist agents for the web that can
1256 +L40: follow language instructions to complete complex ...</result>
1257 +<result>[20] VisualWebArena benchmark (Search_Results/VisualWebArena benchmark)
1258 +**viewing lines [0 - 38] of 38**
1259 +
1260 +L0:
1261 +L1: URL: Search_Results/VisualWebArena benchmark
1262 +L2: # Search Results
1263 +L3:
1264 +L4: * 【0†VisualWebArena is a benchmark for multimodal agents.; publish_date:
1265 +L5: none†github.com】 VisualWebArena is a realistic and diverse benchmark for
1266 +L6: evaluating multimodal autonomous language agents.
1267 +L7: * 【1†Evaluating Multimodal Agents on Realistic Visual Web Tasks; publish_date:
1268 +L8: none†arxiv.org】 by JY Koh · 2024 · Cited by 363 — To bridge this gap, we
1269 +L9: introduce VisualWebArena, a benchmark designed to assess the performance of
1270 +L10: multimodal web agents on realistic \textit{ ...
1271 +L11: * 【2†Evaluating Multimodal Agents on Realistic Visual Web Tasks; publish_date:
1272 +L12: none†jykoh.com】 To bridge this gap, we introduce VisualWebArena, a benchmark
1273 +L13: designed to assess the performance of multimodal web agents on realistic
1274 +L14: visually grounded tasks.
1275 +L15: * 【3†VisualWebArena: Evaluating Multimodal Agents on ...; publish_date:
1276 +L16: none†arxiv.org】 VisualWebArena is a research benchmark to measure and evaluate
1277 +L17: the progress of multimodal agents. It is primarily meant to act as a self-
1278 +L18: contained sandbox ...
1279 +L19: * 【4†Evaluating Multimodal Agents on Realistic Visual Web Tasks; publish_date:
1280 +L20: none†aclanthology.org】 by JY Koh · 2024 · Cited by 363 — To bridge this gap, we
1281 +L21: introduce VisualWebArena, a benchmark designed to assess the performance of
1282 +L22: multimodal web agents on *realistic visually grounded tasks*.
1283 +L23: * 【5†Evaluating Multimodal Agents on Realistic Visual Web Tasks; publish_date:
1284 +L24: none†www.semanticscholar.org】 VisualWebArena: Evaluating Multimodal Agents on
1285 +L25: Realistic Visual Web Tasks ... MMInA, a multihop and multimodal benchmark to
1286 +L26: evaluate the embodied agents ...
1287 +L27: * 【6†CMU Researchers Introduce VisualWebArena: An AI ...; publish_date:
1288 +L28: none†www.marktechpost.com】 Feb 9, 2024 — VisualWebArena, a benchmark designed
1289 +L29: and developed to evaluate the performance of multimodal web agents on realistic
1290 +L30: and visually stimulating challenges.
1291 +L31: * 【7†Evaluating Multimodal Agents on Realistic Visual Web Tasks; publish_date:
1292 +L32: none†www.themoonlight.io】 The paper "VisualWebArena: Evaluating Multimodal
1293 +L33: Agents on Realistic Visually Grounded Web Tasks" introduces a new benchmark,
1294 +L34: **VisualWebArena**, ...
1295 +L35: * 【8†WebArena: A Realistic Web Environment for Building ...; publish_date:
1296 +L36: none†webarena.dev】 Our benchmark is implemented in our fully interactable
1297 +L37: highly-realistic WebArena environment. It features diverse tasks human may
1298 +L38: encounter in their daily ...</result>
1299 +<result>[21] SearchBench benchmark (Search_Results/SearchBench benchmark)
1300 +**viewing lines [0 - 40] of 40**
1301 +
1302 +L0:
1303 +L1: URL: Search_Results/SearchBench benchmark
1304 +L2: # Search Results
1305 +L3:
1306 +L4: * 【0†Talc-AI/search-bench; publish_date: none†github.com】 A practical
1307 +L5: benchmark that focuses on every day helpfulness of LLM products, not just the
1308 +L6: underlying models. Searchbench is a benchmark that addresses these ...
1309 +L7: * 【1†Evaluating LLMs' Ability to Reason About Search Problems; publish_date:
1310 +L8: none†arxiv.org】 These capabilities are essential for robust reasoning, making
1311 +L9: SearchBench a valuable benchmark for evaluating LLMs' reasoning capabilities as
1312 +L10: they continue to ...
1313 +L11: * 【2†NasimBrz/SearchBench · Datasets at Hugging Face; publish_date:
1314 +L12: none†huggingface.co】 Dataset Summary. SearchBench is a benchmark designed to
1315 +L13: evaluate Language Models' (LLMs) ability to solve state-based problems that
1316 +L14: require combinatorial search ...
1317 +L15: * 【3†Evaluating LLMs' Ability to Reason About Search Problems; publish_date:
1318 +L16: none†openreview.net】 2025年10月22日 — To further investigate this, we introduce a
1319 +L17: new benchmark, SearchBench, which contains 11 unique search problems inspired by
1320 +L18: intuitive puzzles.
1321 +L19: * 【4†Navigating the Labyrinth: Evaluating and Enhancing LLMs' ...;
1322 +L20: publish_date: none†hub.baai.ac.cn】 2024年6月17日 —
1323 +L21: 论文提出了一个新的基准测试SearchBench，包含11种独特的搜索问题类型，并自动化生成任意数量的实例和分析解决方案的可行性、正确性和最优性。论文使用A*
1324 +L22: ...
1325 +L23: * 【5†Towards Unified Text-based Person Retrieval: A Large- ...; publish_date:
1326 +L24: none†blog.csdn.net】 2023年10月17日 — ... Search
1327 +L25: Benchmark（面向统一的基于文本的人物检索：一个大规模的多属性和语言搜索基准）; 研究背景. 相关工作; BENCHMARK. 论文方法分析. 网络框架;
1328 +L26: 1、APTM ...
1329 +L27: * 【6†Desearch-ai/ai-search-benchmark; publish_date: none†github.com】 The
1330 +L28: SearchBench repository addresses common issues with traditional benchmarks by
1331 +L29: focusing on practical, everyday use cases rather than theoretical limits. It ...
1332 +L30: * 【7†o1 results for 3 benchmarks: PlanBench, SearchBench, ...; publish_date:
1333 +L31: none†www.reddit.com】 o1 results for 3 benchmarks: PlanBench, SearchBench, and
1334 +L32: Summary of a Haystack. AI. PlanBench: Paper "LLMs Still Can't Plan; Can LRMs? A
1335 +L33: ...
1336 +L34: * 【8†Evaluating LLMs' Ability to Reason About Search Problems; publish_date:
1337 +L35: none†ui.adsabs.harvard.edu】 To further investigate this, we introduce a new
1338 +L36: benchmark, SearchBench, which contains 11 unique search problems inspired by
1339 +L37: intuitive puzzles. Each SearchBench ...
1340 +L38: * 【9†Introducing SearchBench; publish_date: none†www.tag1consulting.com】
1341 +L39: Toward this goal, over the weekend I launched a new project called SearchBench,
1342 +L40: a Drupal module for benchmarking Drupal's search performance. As the module ...</result>
1343 +<result>[22] WebVLN benchmark (Search_Results/WebVLN benchmark)
1344 +**viewing lines [0 - 42] of 42**
1345 +
1346 +L0:
1347 +L1: URL: Search_Results/WebVLN benchmark
1348 +L2: # Search Results
1349 +L3:
1350 +L4: * 【0†WebVLN: Vision-and-Language Navigation on Websites; publish_date:
1351 +L5: none†ojs.aaai.org】 by Q Chen · 2024 · Cited by 35 — the WebVLN-v1 dataset, where
1352 +L6: the performance is far from saturation, highlighting the utility of our
1353 +L7: WebVLN-v1 as a benchmark to assess progress in this field.
1354 +L8: * 【1†[2312.15820] WebVLN: Vision-and-Language Navigation on Websites;
1355 +L9: publish_date: none†ar5iv.labs.arxiv.org】 Experimental results show that WebVLN-
1356 +L10: Net outperforms current VLN and web-related navigation methods. ... Code is
1357 +L11: available at: https://github.com/WebVLN/WebVLN.
1358 +L12: * 【2†WebVLN: Vision-and-Language Navigation on Websites; publish_date:
1359 +L13: none†github.com】 Experimental results show that WebVLN-Net outperforms current
1360 +L14: VLN and web-related navigation methods. We believe that the introduction of the
1361 +L15: new WebVLN task ...
1362 +L16: * 【3†Vision-and-Language Navigation in the Real-World; publish_date:
1363 +L17: none†digital.library.adelaide.edu.au】 By leveraging our proposed WebVLN-v1
1364 +L18: dataset, experimental results showcase the superior performance of WebVLN-Net
1365 +L19: compared to existing VLN and web-related ...
1366 +L20: * 【4†WebVLN: Vision-and-Language Navigation on Websites; publish_date:
1367 +L21: none†www.researchgate.net】 Experimental results show that WebVLN-Net outperforms
1368 +L22: current VLN and web-related navigation methods. We believe that the
1369 +L23: introduction of the newWebVLN task and ...
1370 +L24: * 【5†[PDF] WebVLN: Vision-and-Language Navigation on Websites; publish_date:
1371 +L25: none†www.semanticscholar.org】 A new task named Vision-and-Language Navigation on
1372 +L26: Websites (WebVLN), where question-based instructions are used to train an
1373 +L27: agent, emulating how users ...
1374 +L28: * 【6†WebVLN: Vision-and-Language Navigation on Websites; publish_date:
1375 +L29: none†arxiv.org】 by Q Chen · 2023 · Cited by 35 — Experimental results show that
1376 +L30: WebVLN-Net outperforms current VLN and web-related navigation methods. We
1377 +L31: believe that the introduction of the ...
1378 +L32: * 【7†Human-Aware Vision-and-Language Navigation; publish_date:
1379 +L33: none†proceedings.neurips.cc】 by H Li · 2024 · Cited by 19 — Vision-and-Language
1380 +L34: Navigation (VLN) [2, 7, 9, 40] has emerged as a key benchmark for evaluating.
1381 +L35: Sim2Real transfer [23], showing impressive performance in ...
1382 +L36: * 【8†LiveBench; publish_date: none†livebench.ai】 Introducing LiveBench: a
1383 +L37: benchmark for LLMs designed with test set contamination and objective evaluation
1384 +L38: in mind.
1385 +L39: * 【9†MG-VLN: Benchmarking Multi-Goal and Long-Horizon ...; publish_date:
1386 +L40: none†ieeexplore.ieee.org】 by J Zhang · 2024 — This task aims to provide a
1387 +L41: simulation benchmark to guide the design of lifelong and long-horizon navigation
1388 +L42: robots.</result>
1389 +<result>[23] WebNav benchmark (Search_Results/WebNav benchmark)
1390 +**viewing lines [0 - 36] of 36**
1391 +
1392 +L0:
1393 +L1: URL: Search_Results/WebNav benchmark
1394 +L2: # Search Results
1395 +L3:
1396 +L4: * 【0†WebNav: A New Large-Scale Task for Natural Language ...; publish_date:
1397 +L5: none†github.com】 WebNav is a benchmark task for evaluating an agent with
1398 +L6: abilities to understand natural language and plan on partially observed
1399 +L7: environments.
1400 +L8: * 【1†[1602.02261] End-to-End Goal-Driven Web Navigation; publish_date:
1401 +L9: none†arxiv.org】 by R Nogueira · 2016 · Cited by 39 — We propose a goal-driven
1402 +L10: web navigation as a benchmark task for evaluating an agent with abilities to
1403 +L11: understand natural language and plan on partially ...
1404 +L12: * 【2†nyu-dl/dl4ir-webnav; publish_date: none†github.com】 WebNav is a benchmark
1405 +L13: task for evaluating an agent with abilities to understand natural language and
1406 +L14: plan on partially observed environments.
1407 +L15: * 【3†WebNav: A New Large-Scale Task for Natural Language ...; publish_date:
1408 +L16: none†www.researchgate.net】 We propose a goal-driven web navigation as a
1409 +L17: benchmark task for evaluating an agent with abilities to understand natural
1410 +L18: language and plan on partially ...
1411 +L19: * 【4†WebNav: An Intelligent Agent for Voice-Controlled Web ...; publish_date:
1412 +L20: none†arxiv.org】 In this paper, we introduce WebNav, an innovative multi-modal
1413 +L21: agent designed to automate web tasks based on natural language.
1414 +L22: * 【5†WebCanvas: Benchmarking Web Agents in Online ...; publish_date:
1415 +L23: none†openreview.net】 by Y Pan · Cited by 78 — TL;DR: We introduce WebCanvas, an
1416 +L24: online evaluation framework for web agents designed to address the dynamic
1417 +L25: nature of web interactions.
1418 +L26: * 【6†WebNav: An Intelligent Agent for Voice-Controlled Web ...; publish_date:
1419 +L27: none†huggingface.co】 Mar 17, 2025 — Preliminary evaluations show that WebNav
1420 +L28: outperforms traditional screen readers in response time and task completion
1421 +L29: accuracy for the visually ...
1422 +L30: * 【7†Web Bench - A new way to compare AI Browser Agents; publish_date:
1423 +L31: none†www.skyvern.com】 May 29, 2025 — TL;DR: Web Bench is a new dataset to
1424 +L32: evaluate web browsing agents that consists of 5,750 tasks on 452 different
1425 +L33: websites, with 2,454 tasks ...
1426 +L34: * 【8†Web Framework Benchmarks; publish_date: none†www.techempower.com】 In the
1427 +L35: following tests, we have measured the performance of several web application
1428 +L36: platforms, full-stack frameworks, and micro-frameworks.</result>
1429 +</function_results>
1430 +
1431 +<think>我们收到了大量关于各种基准测试的搜索结果。现在，我们需要整理一份针对搜索智能体的基准测试列表。这些基准测试可以分为几类：网络浏览智能体基准测试（例如BrowseComp、WebVoyager、WebArena、VisualWebArena、Mind2Web、WebVLN、WebNav）、工具使用智能体基准测试（例如ToolBench、GTA）、问答/检索基准测试（例如HotpotQA、FEVER、TriviaQA、Natural Questions、MS MARCO、BEIR、MIRACL），以及特定领域的搜索基准测试（例如HSCodeComp、WebShop、ALFWorld）。此外，还有一些通用智能体基准测试，例如GAIA、SuperCLUE-Agent等。用户要求“详细介绍各自的特点、使用场景、例题”。因此，我需要为每个基准测试提供详细的描述，包括其特点、使用场景，并举例说明。为了全面获取这些信息，我将同时打开以下基准测试的相关页面：WebVoyager（索引3）、WebArena（索引1）、GAIA（索引1）、ToolBench（索引0）、HotpotQA（索引0）、FEVER（索引1）、TriviaQA（索引0）、Natural Questions（索引0）、MS MARCO（索引0）、BEIR（索引1）、MIRACL（索引0）、WebShop（索引0）、ALFWorld（索引0）、Mind2Web（索引0）、VisualWebArena（索引0）、SearchBench（索引1）、WebVLN（索引1）和WebNav（索引0）。</think>
1432 +
1433 +<｜DSML｜function_calls>
1434 +<｜DSML｜invoke name="open">
1435 +<｜DSML｜parameter name="cursor" string="false">4</｜DSML｜parameter>
1436 +<｜DSML｜parameter name="id" string="false">3</｜DSML｜parameter>
1437 +</｜DSML｜invoke>
1438 +<｜DSML｜invoke name="open">
1439 +<｜DSML｜parameter name="cursor" string="false">5</｜DSML｜parameter>
1440 +<｜DSML｜parameter name="id" string="false">1</｜DSML｜parameter>
1441 +</｜DSML｜invoke>
1442 +<｜DSML｜invoke name="open">
1443 +<｜DSML｜parameter name="cursor" string="false">6</｜DSML｜parameter>
1444 +<｜DSML｜parameter name="id" string="false">1</｜DSML｜parameter>
1445 +</｜DSML｜invoke>
1446 +<｜DSML｜invoke name="open">
1447 +<｜DSML｜parameter name="cursor" string="false">7</｜DSML｜parameter>
1448 +<｜DSML｜parameter name="id" string="false">0</｜DSML｜parameter>
1449 +</｜DSML｜invoke>
1450 +<｜DSML｜invoke name="open">
1451 +<｜DSML｜parameter name="cursor" string="false">8</｜DSML｜parameter>
1452 +<｜DSML｜parameter name="id" string="false">0</｜DSML｜parameter>
1453 +</｜DSML｜invoke>
1454 +<｜DSML｜invoke name="open">
1455 +<｜DSML｜parameter name="cursor" string="false">9</｜DSML｜parameter>
1456 +<｜DSML｜parameter name="id" string="false">1</｜DSML｜parameter>
1457 +</｜DSML｜invoke>
1458 +<｜DSML｜invoke name="open">
1459 +<｜DSML｜parameter name="cursor" string="false">10</｜DSML｜parameter>
1460 +<｜DSML｜parameter name="id" string="false">0</｜DSML｜parameter>
1461 +</｜DSML｜invoke>
1462 +<｜DSML｜invoke name="open">
1463 +<｜DSML｜parameter name="cursor" string="false">11</｜DSML｜parameter>
1464 +<｜DSML｜parameter name="id" string="false">0</｜DSML｜parameter>
1465 +</｜DSML｜invoke>
1466 +<｜DSML｜invoke name="open">
1467 +<｜DSML｜parameter name="cursor" string="false">12</｜DSML｜parameter>
1468 +<｜DSML｜parameter name="id" string="false">0</｜DSML｜parameter>
1469 +</｜DSML｜invoke>
1470 +<｜DSML｜invoke name="open">
1471 +<｜DSML｜parameter name="cursor" string="false">13</｜DSML｜parameter>
1472 +<｜DSML｜parameter name="id" string="false">1</｜DSML｜parameter>
1473 +</｜DSML｜invoke>
1474 +<｜DSML｜invoke name="open">
1475 +<｜DSML｜parameter name="cursor" string="false">14</｜DSML｜parameter>
1476 +<｜DSML｜parameter name="id" string="false">0</｜DSML｜parameter>
1477 +</｜DSML｜invoke>
1478 +<｜DSML｜invoke name="open">
1479 +<｜DSML｜parameter name="cursor" string="false">17</｜DSML｜parameter>
1480 +<｜DSML｜parameter name="id" string="false">0</｜DSML｜parameter>
1481 +</｜DSML｜invoke>
1482 +<｜DSML｜invoke name="open">
1483 +<｜DSML｜parameter name="cursor" string="false">18</｜DSML｜parameter>
1484 +<｜DSML｜parameter name="id" string="false">0</｜DSML｜parameter>
1485 +</｜DSML｜invoke>
1486 +<｜DSML｜invoke name="open">
1487 +<｜DSML｜parameter name="cursor" string="false">19</｜DSML｜parameter>
1488 +<｜DSML｜parameter name="id" string="false">0</｜DSML｜parameter>
1489 +</｜DSML｜invoke>
1490 +<｜DSML｜invoke name="open">
1491 +<｜DSML｜parameter name="cursor" string="false">20</｜DSML｜parameter>
1492 +<｜DSML｜parameter name="id" string="false">0</｜DSML｜parameter>
1493 +</｜DSML｜invoke>
1494 +<｜DSML｜invoke name="open">
1495 +<｜DSML｜parameter name="cursor" string="false">21</｜DSML｜parameter>
1496 +<｜DSML｜parameter name="id" string="false">1</｜DSML｜parameter>
1497 +</｜DSML｜invoke>
1498 +<｜DSML｜invoke name="open">
1499 +<｜DSML｜parameter name="cursor" string="false">22</｜DSML｜parameter>
1500 +<｜DSML｜parameter name="id" string="false">1</｜DSML｜parameter>
1501 +</｜DSML｜invoke>
1502 +<｜DSML｜invoke name="open">
1503 +<｜DSML｜parameter name="cursor" string="false">23</｜DSML｜parameter>
1504 +<｜DSML｜parameter name="id" string="false">0</｜DSML｜parameter>
1505 +</｜DSML｜invoke>
1506 +</｜DSML｜function_calls><｜end▁of▁sentence｜>
1507 +
1508 +<function_results>
1509 +<result>[24] (https://ar5iv.org/abs/2401.13919)
1510 +**viewing lines [0 - 129] of 825**
1511 +
1512 +L0:
1513 +L1: URL: https://ar5iv.org/abs/2401.13919
1514 +L2: # WebVoyager : Building an End-to-End Web Agent with
1515 +L3: Large Multimodal Models
1516 +L4:
1517 +L5: Hongliang He1,3 , Wenlin Yao2, Kaixin Ma2, Wenhao Yu2, Yong Dai2,
1518 +L6:
1519 +L7: Hongming Zhang2,
1520 +L8: Zhenzhong Lan3,
1521 +L9: Dong Yu2
1522 +L10: 1Zhejiang University,
1523 +L11: 2Tencent AI Lab,
1524 +L12: 3Westlake University
1525 +L13: hehongliang@westlake.edu.cn, wenlinyao@global.tencent.com
1526 +L14: Work done during the internship at Tencent AI Lab.
1527 +L15:
1528 +L16: ###### Abstract
1529 +L17:
1530 +L18: The advancement of large language models (LLMs) leads to a new era marked by the
1531 +L19: development of autonomous applications in the real world, which drives
1532 +L20: innovation in the creation of advanced web-based agents. Existing web agents
1533 +L21: typically only handle one input modality and are evaluated only in simplified
1534 +L22: web simulators or static web snapshots, greatly limiting their applicability in
1535 +L23: real-world scenarios. To bridge this gap, we introduce WebVoyager, an innovative
1536 +L24: Large Multimodal Model (LMM) powered web agent that can complete user
1537 +L25: instructions end-to-end by interacting with real-world websites. Moreover, we
1538 +L26: propose a new evaluation protocol for web agents to address the challenges of
1539 +L27: automatic evaluation of open-ended web agent tasks, leveraging the robust
1540 +L28: multimodal comprehension capabilities of GPT-4V. We create a new benchmark by
1541 +L29: gathering real-world tasks from 15 widely used websites to evaluate our agents.
1542 +L30: We show that WebVoyager achieves a 55.7% task success rate, significantly
1543 +L31: surpassing the performance of both GPT-4 (All Tools) and the WebVoyager (text-
1544 +L32: only) setups, underscoring the exceptional capability of WebVoyager in practical
1545 +L33: applications. We found that our proposed automatic evaluation achieves 85.3%
1546 +L34: agreement with human judgment, paving the way for further development of web
1547 +L35: agents in a real-world setting.111Our code and data will be released at
1548 +L36: https://github.com/MinorJerry/WebVoyager
1549 +L37:
1550 +L38: ## 1 Introduction
1551 +L39:
1552 +L40: The recent advancement of large language models (LLMs), such as ChatGPT and
1553 +L41: GPT-4 (OpenAI, 2023), have sparked significant interest in developing LLM-based
1554 +L42: autonomous agents (AutoGPT, 2022) for complex task execution (Qin et al., 2023;
1555 +L43: Schick et al., 2023). Recent studies have explored the construction of text-
1556 +L44: based web browsing environments and how to instruct large language model agents
1557 +L45: to perform web navigation (Nakano et al., 2021; Gur et al., 2023; Zhou et al.,
1558 +L46: 2023; Lu et al., 2023). The primary challenge in these works lies in managing
1559 +L47: complex and verbose HTML texts, and solutions include simplifying and
1560 +L48: structuring HTML (Nakano et al., 2021; Zhou et al., 2023; Gur et al., 2023; Deng
1561 +L49: et al., 2023).
1562 +L50:
1563 +L51: However, existing approaches overlook a critical functionality of browsing:
1564 +L52: rendering HTML into visual webpages. Particularly, vision capability is crucial
1565 +L53: for utilizing tools like web browsers, as rendered web pages are inherently
1566 +L54: designed with user experience (UX), emphasizing intuitive information and
1567 +L55: structured presentation. This design principle of rendering makes visual
1568 +L56: analysis more effective than mere HTML representation. At present, large
1569 +L57: multimodal models (LMMs), particularly GPT-4V(ision) (OpenAI, 2023) and Gemini
1570 +L58: (Team et al., 2023), demonstrate a remarkable ability to integrate intricate
1571 +L59: visual cues with textual information. Existing studies such as Pix2Struct (Lee
1572 +L60: et al., 2023) and WebArena (Zhou et al., 2023), have initiated explorations into
1573 +L61: using screenshots as inputs for decision-making in web navigation, yet these
1574 +L62: are preliminary and do not represent a deep exploration. Therefore, building
1575 +L63: multimodal web agents to leverage the environment rendered by browsers through
1576 +L64: screenshots, thus mimicking human web browsing behavior, is now a viable
1577 +L65: approach to enhance web navigation efficiency.
1578 +L66:
1579 +L67: We introduce WebVoyager, a multimodal web agent designed to handle web tasks
1580 +L68: online in an end-to-end manner, which denotes managing the process from start to
1581 +L69: finish autonomously without intermediate human intervention. We construct an
1582 +L70: online environment using Selenium for WebVoyager, feeding it with screenshots
1583 +L71: and textual content in interactive web elements. Inspired by Set-of-Mark
1584 +L72: Prompting (Yang et al., 2023a), we mark interactive web elements on screenshots
1585 +L73: (see Figure 2) to facilitate decision-making for WebVoyager. As a pioneer in
1586 +L74: combining vision and text information during web navigation, we advocate that
1587 +L75: autonomous end-to-end task completion, multimodal capabilities and online
1588 +L76: navigation constitute the essential trajectory toward the genuine intelligence
1589 +L77: of web agents.
1590 +L78:
1591 +L79: Another challenge arises when it comes to evaluating an end-to-end web agent
1592 +L80: with online navigation. Existing benchmarks, such as Mind2Web (Deng et al.,
1593 +L81: 2023), primarily focus on stepwise and offline evaluation, where agents follow
1594 +L82: predefined “golden” trajectory for action selection. This approach, however, may
1595 +L83: not fully account for the variety of viable strategies to accomplish a task, as
1596 +L84: it only reflects one possible plan. This limitation could lead to a biased
1597 +L85: evaluation and difficulties in fairly comparing different methods. To more
1598 +L86: accurately gauge the capabilities of web agents in end-to-end task completion,
1599 +L87: we save screenshots throughout the online navigation process, and then use
1600 +L88: GPT-4V to evaluate these trajectories and the final results automatically. Human
1601 +L89: evaluations are also conducted to verify the results and confirm the
1602 +L90: reliability of GPT-4V as the evaluator.
1603 +L91:
1604 +L92: We conduct evaluations on a collected dataset, which is semi-automatically
1605 +L93: generated using a self-instruct (Wang et al., 2022) method, comprising 300 web
1606 +L94: tasks from 15 commonly accessed websites. Additionally, we extract 90 web-
1607 +L95: related tasks of level 1 and level 2 from the GAIA (Mialon et al., 2023) to
1608 +L96: enrich our evaluation. We compare our WebVoyager with 1) GPT-4 (All
1609 +L97: Tools)222GPT-4 (All Tools) is an integrated tool-based agent released by OpenAI
1610 +L98: in Oct. 2023. See https://chat.openai.com/, and 2) WebVoyager in a text-only
1611 +L99: setting, employing the accessibility tree proposed in WebArena (Zhou et al.,
1612 +L100: 2023) to describe web pages. The results show that WebVoyager achieves a Task
1613 +L101: Success Rate of 55.7%, significantly outperforming GPT-4 (All Tools) with a rate
1614 +L102: of 32.7% and the text-only setting with a rate of 39.0%, demonstrating the
1615 +L103: effectiveness of our method. Furthermore, we report the consistency between
1616 +L104: human-human and human-GPT4V to ensure credibility. Our main contributions are as
1617 +L105: follows:
1618 +L106:
1619 +L107: - •
1620 +L108: We employ a multimodal web agent that integrates textual and visual information
1621 +L109: to address web tasks end-to-end and introduce a generalist planning approach for
1622 +L110: navigation.
1623 +L111:
1624 +L112: - •
1625 +L113: We build an online web browsing environment, offering a variety of tasks
1626 +L114: centered on widely used websites and introducing a method for expanding these
1627 +L115: tasks.
1628 +L116:
1629 +L117: - •
1630 +L118: We conduct manual evaluations of navigation trajectories and propose an
1631 +L119: automated evaluation protocol using GPT-4V. We present a comprehensive analysis
1632 +L120: of the evaluation results and show that GPT-4V can serve as a reliable evaluator
1633 +L121: for online agents.
1634 +L122:
1635 +L123: ## 2 Related Work
1636 +L124:
1637 +L125: ### 2.1 Web Agents
1638 +L126:
1639 +L127: Autonomous web navigation (Shi et al., 2017; Yang et al., 2023b) requires an
1640 +L128: agent to follow instructions, construct plans, comprehend complex web
1641 +L129: structures, and decompose tasks into step-by-step decisions (Weng, 2023). To </result>
1642 +<result>[25] (https://ar5iv.org/abs/2307.13854)
1643 +**viewing lines [0 - 124] of 1092**
1644 +
1645 +L0:
1646 +L1: URL: https://ar5iv.org/abs/2307.13854
1647 +L2: # WebArena: A Realistic Web Environment for Building Autonomous Agents
1648 +L3:
1649 +L4: Shuyan Zhou Frank F. Xu11footnotemark: 1 Hao Zhu Xuhui Zhou22footnotemark: 2
1650 +L5:
1651 +L6: Robert Lo22footnotemark: 2 Abishek Sridhar22footnotemark: 2 Xianyi Cheng Tianyue
1652 +L7: Ou
1653 +L8: Yonatan Bisk Daniel Fried Uri Alon Graham Neubig
1654 +L9: Carnegie Mellon University
1655 +L10: {shuyanzh, fangzhex, gneubig}@cs.cmu.edu
1656 +L11: Lead contributors.Equal contribution.
1657 +L12:
1658 +L13: ###### Abstract
1659 +L14:
1660 +L15: With advances in generative AI, there is now potential for autonomous agents to
1661 +L16: manage daily tasks via natural language commands. However, current agents are
1662 +L17: primarily created and tested in simplified synthetic environments, leading to a
1663 +L18: disconnect with real-world scenarios. In this paper, we build an environment for
1664 +L19: language-guided agents that is highly realistic and reproducible. Specifically,
1665 +L20: we focus on agents that perform tasks on the web, and create an environment
1666 +L21: with fully functional websites from four common domains: e-commerce, social
1667 +L22: forum discussions, collaborative software development, and content management.
1668 +L23: Our environment is enriched with tools (e.g., a map) and external knowledge
1669 +L24: bases (e.g., user manuals) to encourage human-like task-solving. Building upon
1670 +L25: our environment, we release a set of benchmark tasks focusing on evaluating the
1671 +L26: functional correctness of task completions. The tasks in our benchmark are
1672 +L27: diverse, long-horizon, and designed to emulate tasks that humans routinely
1673 +L28: perform on the internet. We experiment with several baseline agents, integrating
1674 +L29: recent techniques such as reasoning before acting. The results demonstrate that
1675 +L30: solving complex tasks is challenging: our best GPT-4-based agent only achieves
1676 +L31: an end-to-end task success rate of 14.41%, significantly lower than the human
1677 +L32: performance of 78.24%. These results highlight the need for further development
1678 +L33: of robust agents, that current state-of-the-art large language models are far
1679 +L34: from perfect performance in these real-life tasks, and that WebArena can be used
1680 +L35: to measure such progress.
1681 +L36:
1682 +L37: Our code, data, environment reproduction resources, and video demonstrations are
1683 +L38: publicly available at https://webarena.dev/.
1684 +L39:
1685 +L40: ## 1 Introduction
1686 +L41:
1687 +L42: Autonomous agents that perform everyday tasks via human natural language
1688 +L43: commands could significantly augment human capabilities, improve efficiency, and
1689 +L44: increase accessibility. Nonetheless, to fully leverage the power of autonomous
1690 +L45: agents, it is crucial to understand their behavior within an environment that is
1691 +L46: both authentic and reproducible. This will allow measurement of the ability of
1692 +L47: agents on tasks that human users care about in a fair and consistent manner.
1693 +L48:
1694 +L49: Current environments for evaluate agents tend to over-simplify real-world
1695 +L50: situations. As a result, the functionality of many environments is a limited
1696 +L51: version of their real-world counterparts, leading to a lack of task diversity
1697 +L52: (Shi et al., 2017; Anderson et al., 2018; Gordon et al., 2018; Misra et al.,
1698 +L53: 2016; Shridhar et al., 2020; 2021; Yao et al., 2022a). In addition, these
1699 +L54: simplifications often lower the complexity of tasks as compared to their
1700 +L55: execution in the real world (Puig et al., 2018; Shridhar et al., 2020; Yao et
1701 +L56: al., 2022a). Finally, some environments are presented as a static resource (Shi
1702 +L57: et al., 2017; Deng et al., 2023) where agents are confined to accessing only
1703 +L58: those states that were previously cached during data collection, thus limiting
1704 +L59: the breadth and diversity of exploration. Dor evaluation, many environments
1705 +L60: focus on comparing the textual surface form of the predicted action sequences
1706 +L61: with reference action sequences, disregarding the functional correctness of the
1707 +L62: executions and possible alternative solutions (Puig et al., 2018; Jernite et
1708 +L63: al., 2019; Xu et al., 2021; Li et al., 2020; Deng et al., 2023). These
1709 +L64: limitations often result in a discrepancy between simulated environments and the
1710 +L65: real world, and can potentially impact the generalizability of AI agents to
1711 +L66: successfully understand, adapt, and operate within complex real-world
1712 +L67: situations.
1713 +L68:
1714 +L69: We introduce WebArena, a realistic and reproducible web environment designed to
1715 +L70: facilitate the development of autonomous agents capable of executing tasks (§2).
1716 +L71: An overview of WebArena is in Figure 1. Our environment comprises four fully
1717 +L72: operational, self-hosted web applications, each representing a distinct domain
1718 +L73: prevalent on the internet: online shopping, discussion forums, collaborative
1719 +L74: development, and business content management. Furthermore, WebArena incorporates
1720 +L75: several utility tools, such as map, calculator, and scratchpad, to best support
1721 +L76: possible human-like task executions. Lastly, WebArena is complemented by an
1722 +L77: extensive collection of documentation and knowledge bases that vary from general
1723 +L78: resources like English Wikipedia to more domain-specific references, such as
1724 +L79: manuals for using the integrated development tool (Fan et al., 2022). The
1725 +L80: content populating these websites is extracted from their real-world
1726 +L81: counterparts, preserving the authenticity of the content served on each
1727 +L82: platform. We deliver the hosting services using Docker containers with gym-APIs
1728 +L83: (Brockman et al., 2016), ensuring both the usability and the reproducibility of
1729 +L84: WebArena.
1730 +L85:
1731 +L86: Along with WebArena, we release a ready-to-use benchmark with 812 long-horizon
1732 +L87: web-based tasks (§3). Each task is described as a high-level natural language
1733 +L88: intent, emulating the abstract language usage patterns typically employed by
1734 +L89: humans (Bisk et al., 2019). Two example intents are shown in the upper left of
1735 +L90: Figure 1. We focus on evaluating the functional correctness of these tasks,
1736 +L91: i.e., does the result of the execution actually achieve the desired goal (§3.2).
1737 +L92: For instance, to evaluate the example in Figure 2, our evaluation method
1738 +L93: verifies the concrete contents in the designated repository. This evaluation is
1739 +L94: not only more reliable (Zhong et al., 2017; Chen et al., 2021; Wang et al.,
1740 +L95: 2022) than comparing the textual surface-form action sequences (Puig et al.,
1741 +L96: 2018; Deng et al., 2023) but also accommodate a range of potential valid paths
1742 +L97: to achieve the same goal, which is a ubiquitous phenomenon in sufficiently
1743 +L98: complex tasks.
1744 +L99:
1745 +L100: We use this benchmark to evaluate several agents that can follow NL command and
1746 +L101: perform web-based tasks (§4). These agents are implemented in a few-shot in-
1747 +L102: context learning fashion with powerful large language models (LLMs) such as
1748 +L103: GPT-4 and PALM-2. Experiment results show that the best GPT-4 agent performance
1749 +L104: is somewhat limited, with an end-to-end task success rate of only 14.41%, while
1750 +L105: the human performance is 78.24%. We hypothesize that the limited performance of
1751 +L106: current LLMs stems from a lack of crucial capabilities such as active
1752 +L107: exploration and failure recovery to successfully perform complex tasks (§5.2).
1753 +L108: These outcomes underscore the necessity for further development towards robust
1754 +L109: and effective agents (LeCun, 2022) in WebArena.
1755 +L110:
1756 +L111: ## 2 WebArena: Websites as an Environment for Autonomous Agents
1757 +L112:
1758 +L113: Our goal is to create a realistic and reproducible web environment. We achieve
1759 +L114: reproducibility by making the environment standalone, without relying on live
1760 +L115: websites. This circumvents technical challenges such as bots being subject to
1761 +L116: CAPTCHAs, unpredictable content modifications, and configuration changes, which
1762 +L117: obstruct a fair comparison across different systems over time. We achieve
1763 +L118: realism by using open-source libraries that underlie many in-use sites from
1764 +L119: several popular categories and importing data to our environment from their
1765 +L120: real-world counterparts.
1766 +L121:
1767 +L122: ### 2.1 Controlling Agents through High-level Natural Language
1768 +L123:
1769 +L124: The WebArena environment is denoted asℰ\mathcal{E} with state space </result>
1770 +<result>[26] (https://ar5iv.org/abs/2311.12983)
1771 +**viewing lines [0 - 118] of 1207**
1772 +
1773 +L0:
1774 +L1: URL: https://ar5iv.org/abs/2311.12983
1775 +L2: 1]FAIR, Meta 2]HuggingFace 3]AutoGPT 4]GenAI, Meta
1776 +L3:
1777 +L4: # GAIA: A Benchmark for General AI Assistants
1778 +L5:
1779 +L6: Grégoire Mialon Clémentine Fourrier Craig Swift Thomas Wolf Yann LeCun Thomas
1780 +L7: Scialom [ [ [ [ {gmialon,tscialom}@meta.com clementine@huggingface.co
1781 +L8:
1782 +L9: ###### Abstract
1783 +L10:
1784 +L11: We introduce GAIA, a benchmark for General AI Assistants that, if solved, would
1785 +L12: represent a milestone in AI research. GAIA proposes real-world questions that
1786 +L13: require a set of fundamental abilities such as reasoning, multi-modality
1787 +L14: handling, web browsing, and generally tool-use proficiency. GAIA questions are
1788 +L15: conceptually simple for humans yet challenging for most advanced AIs: we show
1789 +L16: that human respondents obtain 92% vs. 15% for GPT-4 equipped with plugins. This
1790 +L17: notable performance disparity contrasts with the recent trend of LLMs
1791 +L18: outperforming humans on tasks requiring professional skills in e.g. law or
1792 +L19: chemistry. GAIA’s philosophy departs from the current trend in AI benchmarks
1793 +L20: suggesting to target tasks that are ever more difficult for humans. We posit
1794 +L21: that the advent of Artificial General Intelligence (AGI) hinges on a system’s
1795 +L22: capability to exhibit similar robustness as the average human does on such
1796 +L23: questions. Using GAIA’s methodology, we devise 466 questions and their answer.
1797 +L24: We release our questions while retaining answers to 300 of them to power a
1798 +L25: leader-board hereby accessible.
1799 +L26:
1800 +L27: \correspondence
1801 +L28:
1802 +L29: ## 1 Introduction
1803 +L30:
1804 +L31: Large Language Models (LLMs) arguably open the way to general purpose systems.
1805 +L32: Indeed, the latest among them (OpenAI, 2023; Anthropic, 2023; Anil et al., 2023;
1806 +L33: Touvron et al., 2023) are fluent, knowledgeable, aligned to some extent with
1807 +L34: human preferences (Ouyang et al., 2022), and can be augmented (Mialon et al.,
1808 +L35: 2023) with tools such as web browsers or code interpreters in a zero or few-shot
1809 +L36: setting (Brown et al., 2020). However, evaluating these systems is an open
1810 +L37: problem: given their emerging new capabilities, LLMs are regularly breaking AI
1811 +L38: benchmarks, at an ever-increasing rate (Kiela et al., 2023).
1812 +L39:
1813 +L40: In search for more challenging benchmarks, current trend suggests to seek tasks
1814 +L41: that are ever more difficult for humans, and challenge LLMs with more intricate
1815 +L42: educational assessments, for example in STEM and Law, or target more complex
1816 +L43: realisations, such as writing a coherent book. But, tasks that are difficult for
1817 +L44: humans are not necessarily difficult for recent systems: the challenging MMLU
1818 +L45: or GSM8k benchmarks for example (Hendrycks et al., 2021; Cobbe et al., 2021) are
1819 +L46: already close to be solved,111GPT4 does 86.4% on MMLU. Human non-specialist
1820 +L47: accuracy on the benchmark is only 34.5% Expert-level human performance is
1821 +L48: estimated at 89.8%. due to rapid LLM improvement possibly combined with data
1822 +L49: contamination.222See for example the case of Hellaswag. Furthermore, open-ended
1823 +L50: generation generally requires human or model-based evaluation (Zheng et al.,
1824 +L51: 2023). Human evaluation will become less and less feasible when increasing the
1825 +L52: task complexity, e.g. in terms of output length or required skills: how to
1826 +L53: evaluate a book generated by an AI, or solutions to maths problems that few
1827 +L54: people in the world can solve? Model-based evaluations on the other hand are by
1828 +L55: construction dependent of stronger models hence cannot evaluate new state-of-
1829 +L56: the-art models, without mentioning potential subtle biases such as preferring
1830 +L57: the first choice presented (Zheng et al., 2023). Overall, evaluating new AI
1831 +L58: systems requires to rethink benchmarks (Chollet, 2019).
1832 +L59:
1833 +L60: Alternatively to tasks that are harder for humans, AI systems could be asked to
1834 +L61: solve conceptually simple tasks yet that require accurate execution of complex
1835 +L62: sequences of actions, with large combinatorial spaces. The output could only be
1836 +L63: obtained upon successful completion of the task and be easy to validate,
1837 +L64: analogous to the Proof of Work algorithm (Jakobsson and Juels, 1999; Dwork and
1838 +L65: Naor, 1993), where a computer is asked to solve a complex problem whose solution
1839 +L66: is easy to verify. Tasks for AI assistants, given their need for access to a
1840 +L67: diverse and uncertain world, meet this criterion while being inherently rooted
1841 +L68: in practical use cases.
1842 +L69:
1843 +L70: We move in that direction by proposing GAIA, a benchmark for General AI
1844 +L71: Assistants featuring 466 carefully crafted questions and their answer, along
1845 +L72: with the associated design methodology. Our questions are easy to create,
1846 +L73: challenging for AI systems—for LLMs, most require complex generations—, yet
1847 +L74: admit a unique, factual answer, allowing a simple and robust automatic
1848 +L75: evaluation.
1849 +L76:
1850 +L77: GAIA attempts to avoid current pitfalls of LLMs evaluation by targeting:
1851 +L78:
1852 +L79: Real-world and challenging questions. For example, a LLM will typically need to
1853 +L80: browse the open and changing web, handle multi-modality, or reason over multiple
1854 +L81: steps to answer our questions. Conversely, many LLM benchmarks are quite
1855 +L82: specific and/or restricted to closed and synthetic environments.
1856 +L83:
1857 +L84: Easy interpretability through conceptually simple tasks—non experts annotators
1858 +L85: exhibit a near perfect score—, associated reasoning trace, and few but highly
1859 +L86: curated questions. This is in contrast with aggregated benchmarks that can lack
1860 +L87: efficiency and reliability (Perlitz et al., 2023).
1861 +L88:
1862 +L89: Non-gameability. Answering the questions requires successful completion of some
1863 +L90: number of steps, which cannot easily be brute forced due to their diversity. The
1864 +L91: possibility to check the reasoning trace, the accuracy required in the answers,
1865 +L92: their absence in plain text from the internet prevent a possible data
1866 +L93: contamination. In contrast, multiple choice answers (e.g., MMLU) make
1867 +L94: contamination assessment more difficult since a wrong reasoning trace can more
1868 +L95: easily get to the correct choice.
1869 +L96:
1870 +L97: Simplicity of use. Crucially, the answers to our questions are factoid, concise
1871 +L98: and unambiguous. These properties allow simple, fast and factual evaluation. Our
1872 +L99: questions are meant to be answered in zero shot, limiting the influence of the
1873 +L100: evaluation setup. By opposition, many LLM benchmarks require evaluations that
1874 +L101: are sensitive to the experimental setup such as the number and nature of prompts
1875 +L102: (Liang et al., 2022b) (Section 8.2), or the benchmark
1876 +L103: implementation.333https://huggingface.co/blog/evaluating-mmlu-leaderboard
1877 +L104:
1878 +L105: In spite of being successful at tasks that are difficult for humans, the most
1879 +L106: capable LLMs do poorly on GAIA. Even equipped with tools, GPT4 does not exceed a
1880 +L107: 30% success rate for the easiest of our tasks, and 0% for the hardest. In the
1881 +L108: meantime, the average success rate for human respondents is 92%. Consequently, a
1882 +L109: system capable of solving GAIA can be assessed in the context of t-AGI,444As
1883 +L110: defined in https://www.alignmentforum.org/posts/BoA3agdkAzL6HQtQP/clarifying-
1884 +L111: and-predicting-agi, a t-AGI beats, on most tasks, most human experts who are
1885 +L112: given time t to perform the task noting that humans typically take between 6
1886 +L113: minutes for the simplest questions to 17 minutes for the most complex ones. From
1887 +L114: a related perspective, such system would arguably be a competent General AI
1888 +L115: within the framework recently proposed in Morris et al. (2023), which also
1889 +L116: appear to be the next milestone in AI research since ChatGPT (OpenAI, 2023) is
1890 +L117: one level below. This paper covers the composition of GAIA, its design choices,
1891 +L118: and explain how to craft questions and the associated challenges so that the </result>
1892 +<result>Error fetching URL `https://github.com/sambanova/toolbench`</result>
1893 +<result>Error fetching URL `https://hotpotqa.github.io/`</result>
1894 +<result>[27] (https://aclanthology.org/N18-1074.pdf)
1895 +**viewing lines [0 - 134] of 809**
1896 +
1897 +L0:
1898 +L1: URL: https://aclanthology.org/N18-1074.pdf
1899 +L2: ===== Page 1 =====
1900 +L3:
1901 +L4: # FEVER: a large-scale dataset for Fact Extraction and VERification
1902 +L5:
1903 +L6: James Thorne\(^1\), Andreas Vlachos\(^1\), Christos Christodoulopoulos\(^2\),
1904 +L7: and Arpit Mittal\(^2\)
1905 +L8:
1906 +L9: \(^1\)Department of Computer Science, University of Sheffield
1907 +L10: \(^2\)Amazon Research Cambridge
1908 +L11: {j.thorne, a.vlachos}@sheffield.ac.uk
1909 +L12: {chrchrs, mitarpit}@amazon.co.uk
1910 +L13:
1911 +L14: ## Abstract
1912 +L15:
1913 +L16: In this paper we introduce a new publicly available dataset for verification
1914 +L17: against textual sources, FEVER: Fact Extraction and VERification. It consists of
1915 +L18: 185,445 claims generated by altering sentences extracted from Wikipedia and
1916 +L19: subsequently verified without knowledge of the sentence they were derived from.
1917 +L20: The claims are classified as Supported, Refuted or NotEnoughInfo by annotators
1918 +L21: achieving 0.6841 in Fleiss \(\kappa\). For the first two classes, the annotators
1919 +L22: also recorded the sentence(s) forming the necessary evidence for their
1920 +L23: judgment. To characterize the challenge of the dataset presented, we develop a
1921 +L24: pipeline approach and compare it to suitably designed oracles. The best accuracy
1922 +L25: we achieve on labeling a claim accompanied by the correct evidence is 31.87%,
1923 +L26: while if we ignore the evidence we achieve 50.91%. Thus we believe that FEVER is
1924 +L27: a challenging testbed that will help stimulate progress on claim verification
1925 +L28: against textual sources.
1926 +L29:
1927 +L30: ## 1 Introduction
1928 +L31:
1929 +L32: The ever-increasing amounts of textual information available combined with the
1930 +L33: ease in sharing it through the web has increased the demand for verification,
1931 +L34: also referred to as fact checking. While it has received a lot of attention in
1932 +L35: the context of journalism, verification is important for other domains, e.g.
1933 +L36: information in scientific publications, product reviews, etc.
1934 +L37:
1935 +L38: In this paper we focus on verification of textual claims against textual
1936 +L39: sources. When compared to textual entailment (TE)/natural language inference
1937 +L40: (Dagan et al., 2009; Bowman et al., 2015), the key difference is that in these
1938 +L41: tasks the passage to verify each claim is given, and in recent years it
1939 +L42: typically consists a single sentence, while in verification systems it is
1940 +L43: retrieved from a large set of documents in order to form the evidence. Another
1941 +L44: related task is question answering (QA), for which approaches have recently been
1942 +L45: extended to handle large-scale resources such as Wikipedia (Chen et al., 2017).
1943 +L46: However, questions typically provide the information needed to identify the
1944 +L47: answer, while information missing from a claim can often be crucial in
1945 +L48: retrieving refuting evidence. For example, a claim stating "Fiji's largest
1946 +L49: island is Kauai." can be refuted by retrieving "Kauai is the oldest Hawaiian
1947 +L50: Island." as evidence.
1948 +L51:
1949 +L52: Progress on the aforementioned tasks has benefited from the availability of
1950 +L53: large-scale datasets (Bowman et al., 2015; Rajpurkar et al., 2016). However,
1951 +L54: despite the rising interest in verification and fact checking among researchers,
1952 +L55: the datasets currently used for this task are limited to a few hundred claims.
1953 +L56: Indicatively, the recently conducted Fake News Challenge (Pomerleau and Rao,
1954 +L57: 2017) with 50 participating teams used a dataset consisting of 300 claims
1955 +L58: verified against 2,595 associated news articles which is orders of magnitude
1956 +L59: smaller than those used for TE and QA.
1957 +L60:
1958 +L61: In this paper we present a new dataset for claim verification, FEVER: Fact
1959 +L62: Extraction and VERification. It consists of 185,445 claims manually verified
1960 +L63: against the introductory sections of Wikipedia pages and classified as
1961 +L64: Supported, Refuted or NotEnoughInfo. For the first two classes, systems and
1962 +L65: annotators need to also return the combination of sentences forming the
1963 +L66: necessary evidence supporting or refuting the claim (see Figure 1). The claims
1964 +L67: were generated by human annotators extracting claims from Wikipedia and mutating
1965 +L68: them in a variety of ways, some of which were meaning-altering. The
1966 +L69: verification of each
1967 +L70:
1968 +L71: 809
1969 +L72:
1970 +L73: Proceedings of NAACL-HLT 2018, pages 809–819
1971 +L74:
1972 +L75: New Orleans, Louisiana, June 1 - 6, 2018. ©2018 Association for Computational
1973 +L76: Linguistics
1974 +L77:
1975 +L78: ===== Page 2 =====
1976 +L79:
1977 +L80: claim was conducted in a separate annotation process by annotators who were
1978 +L81: aware of the page but not the sentence from which original claim was extracted
1979 +L82: and thus in 31.75% of the claims more than one sentence was considered
1980 +L83: appropriate evidence. Claims require composition of evidence from multiple
1981 +L84: sentences in 16.82% of cases. Furthermore, in 12.15% of the claims, this
1982 +L85: evidence was taken from multiple pages.
1983 +L86:
1984 +L87: To ensure annotation consistency, we developed suitable guidelines and user
1985 +L88: interfaces, resulting in inter-annotator agreement of 0.6841 in Fleiss (Fleiss,
1986 +L89: 1971) in claim verification classification, and 95.42% precision and 72.36%
1987 +L90: recall in evidence retrieval.
1988 +L91:
1989 +L92: To characterize the challenges posed by FEVER we develop a pipeline approach
1990 +L93: which, given a claim, first identifies relevant documents, then selects
1991 +L94: sentences forming the evidence from the documents and finally classifies the
1992 +L95: claim w.r.t. evidence. The best performing version achieves 31.87% accuracy in
1993 +L96: verification when requiring correct evidence to be retrieved for claims
1994 +L97: Supported or Refuted, and 50.91% if the correctness of the evidence is ignored,
1995 +L98: both indicating the difficulty but also the feasibility of the task. We also
1996 +L99: conducted oracle experiments in which components of the pipeline were replaced
1997 +L100: by the gold standard annotations, and observed that the most challenging part of
1998 +L101: the task is selecting the sentences containing the evidence. In addition to
1999 +L102: publishing the data via our website1, we also publish the annotation interfaces2
2000 +L103: and the baseline system3 to stimulate further research on verification.
2001 +L104:
2002 +L105: Footnote 1: http://fever.ai
2003 +L106:
2004 +L107: Footnote 2: https://github.com/awslabs/fever
2005 +L108:
2006 +L109: Footnote 3: https://github.com/sheffieldnlp/fever-baselines
2007 +L110:
2008 +L111: ## 2 Related Works
2009 +L112:
2010 +L113: Vlachos and Riedel (2014) constructed a dataset for claim verification
2011 +L114: consisting of 106 claims, selecting data from fact-checking websites such as
2012 +L115: PolitiFact, taking advantage of the labelled claims available there. However, in
2013 +L116: order to develop claim verification components we typically require the
2014 +L117: justification for each verdict, including the sources used. While this
2015 +L118: information is usually available in justifications provided by the journalists,
2016 +L119: they are not in a machine-readable form. Thus, also considering the small number
2017 +L120: of claims, the task defined by the dataset proposed remains too challenging for
2018 +L121: the ML/NLP methods currently available. Wang (2017) extended this approach by
2019 +L122: including all 12.8K claims available by Politifact via its API, however the
2020 +L123: justification and the evidence contained in it was ignored in the experiments as
2021 +L124: it was not machine-readable. Instead, the claims were classified considering
2022 +L125: only the text and the metadata related to the person making the claim. While
2023 +L126: this rendered the task amenable to current NLP/ML methods, it does not allow for
2024 +L127: verification against any sources and no evidence needs to be returned to
2025 +L128: justify the verdicts.
2026 +L129:
2027 +L130: The Fake News challenge (Pomerleau and Rao, 2017) modelled verification as
2028 +L131: stance classification: given a claim and an article, predict whether the article
2029 +L132: supports, refutes, observes (neutrally states the claim) or is irrelevant to
2030 +L133: the claim. It consists of 50K labelled claim-article pairs, combining 300 claims
2031 +L134: with 2,582 articles. The claims and the articles were curated and labeled by </result>
2032 +<result>[28] (https://nlp.cs.washington.edu/triviaqa/)
2033 +**viewing lines [0 - 48] of 48**
2034 +
2035 +L0:
2036 +L1: URL: https://nlp.cs.washington.edu/triviaqa/
2037 +L2: ## TriviaQA: A Large Scale Dataset for Reading Comprehension and Question
2038 +L3: Answering
2039 +L4:
2040 +L5: TriviaQA is a reading comprehension dataset containing over 650K question-
2041 +L6: answer-evidence triples. TriviaQA includes 95K question-answer pairs authored by
2042 +L7: trivia enthusiasts and independently gathered evidence documents, six per
2043 +L8: question on average, that provide high quality distant supervision for answering
2044 +L9: the questions. The details can be found in our ACL 17 paper TriviaQA: A Large
2045 +L10: Scale Distantly Supervised Challenge Dataset for Reading Comprehension
2046 +L11:
2047 +L12: Mandar Joshi, Eunsol Choi, Daniel Weld, Luke Zettlemoyer. TriviaQA: A Large
2048 +L13: Scale Distantly Supervised Challenge Dataset for Reading Comprehension
2049 +L14:
2050 +L15: In Association for Computational Linguistics (ACL) 2017, Vancouver, Canada.
2051 +L16: [bib]
2052 +L17:
2053 +L18: ## News
2054 +L19:
2055 +L20: ##### Jul 2017
2056 +L21:
2057 +L22: The TriviaQA leaderboard is now live on Codalab. Submit your predictions for
2058 +L23: evaluation on the test set!
2059 +L24:
2060 +L25: ## Data
2061 +L26:
2062 +L27: If you are interested in the reading comprehension task motivated in the paper,
2063 +L28: click on the link below to download the data.
2064 +L29:
2065 +L30: If you are interested in open domain QA, click on the link below to download the
2066 +L31: data. It contains the unfiltered dataset with 110K question-answer pairs. The
2067 +L32: Wikipedia and top 10 search documents can be obtained from the RC version. The
2068 +L33: main difference between the RC version above and the unfiltered dataset is that
2069 +L34: not all documents (in the unfiltered set) for a given question contain the
2070 +L35: answer string(s). This makes the unfiltered dataset more appropriate for IR-
2071 +L36: style QA.
2072 +L37:
2073 +L38: The University of Washington does not own the copyright of the questions and
2074 +L39: documents included in TriviaQA.
2075 +L40:
2076 +L41: ## Code
2077 +L42:
2078 +L43: Check out our Github repository.
2079 +L44:
2080 +L45: ## Contact
2081 +L46:
2082 +L47: For any questions about the code or data, please contact Mandar Joshi -- {first
2083 +L48: name of the first author}90[at]cs[dot]washington[dot]edu</result>
2084 +<result>Error fetching URL `https://research.google/pubs/natural-questions-a-benchmark-for-question-answering-research/`</result>
2085 +<result>Error fetching URL `https://microsoft.github.io/msmarco/`</result>
2086 +<result>Error fetching URL `https://github.com/beir-cellar/beir`</result>
2087 +<result>Error fetching URL `https://project-miracl.github.io/`</result>
2088 +<result>[29] (https://ar5iv.org/abs/2207.01206)
2089 +**viewing lines [0 - 122] of 1392**
2090 +
2091 +L0:
2092 +L1: URL: https://ar5iv.org/abs/2207.01206
2093 +L2: # WebShop: Towards Scalable Real-World Web Interaction with Grounded Language
2094 +L3: Agents
2095 +L4:
2096 +L5: Shunyu Yao111Equal contribution. Project site with code, data, and demos:
2097 +L6: https://webshop-pnlp.github.io. Howard Chen111Equal contribution. Project site
2098 +L7: with code, data, and demos: https://webshop-pnlp.github.io. John Yang Karthik
2099 +L8: Narasimhan
2100 +L9:
2101 +L10: Department of Computer Science, Princeton University
2102 +L11: {shunyuy, howardchen, jy1682, karthikn}@princeton.edu
2103 +L12:
2104 +L13: ###### Abstract
2105 +L14:
2106 +L15: Existing benchmarks for grounding language in interactive environments either
2107 +L16: lack real-world linguistic elements, or prove difficult to scale up due to
2108 +L17: substantial human involvement in the collection of data or feedback signals. To
2109 +L18: bridge this gap, we develop WebShop – a simulated e-commerce website environment
2110 +L19: with million real-world products and 1.181.18 crowd-sourced text instructions.
2111 +L20: Given a text instruction specifying a product requirement, an agent needs to
2112 +L21: navigate multiple types of webpages and issue diverse actions to find,
2113 +L22: customize, and purchase an item. WebShop provides several challenges for
2114 +L23: language grounding including understanding compositional instructions, query
2115 +L24: (re-)formulation, comprehending and acting on noisy text in webpages, and
2116 +L25: performing strategic exploration. We collect over 12,08712,087 human
2117 +L26: demonstrations for the task, and train and evaluate a diverse range of agents
2118 +L27: using reinforcement learning, imitation learning, and pre-trained image and
2119 +L28: language models. Our best model achieves a task success rate of 1,6001,600,
2120 +L29: which outperforms rule-based heuristics (29%29\%) but is far lower than human
2121 +L30: expert performance (9.6%9.6\%). We also analyze agent and human trajectories and
2122 +L31: ablate various model components to provide insights for developing future
2123 +L32: agents with stronger language understanding and decision making abilities.
2124 +L33: Finally, we show that agents trained on WebShop exhibit non-trivial sim-to-real
2125 +L34: transfer when evaluated on amazon.com and ebay.com , indicating the potential
2126 +L35: value of WebShop in developing practical web-based agents that can operate in
2127 +L36: the wild.59%59\%
2128 +L37:
2129 +L38: ## 1 Introduction
2130 +L39:
2131 +L40: Recent advances in natural language processing (NLP) and reinforcement learning
2132 +L41: (RL) have brought about several exciting developments in agents that can perform
2133 +L42: sequential decision making while making use of linguistic context [30, 50, 58].
2134 +L43: On the other hand, large-scale language models like GPT-3 [6] and BERT [11] are
2135 +L44: excelling at traditional NLP benchmarks such as text classification,
2136 +L45: information extraction and question answering. While the former set of tasks are
2137 +L46: limited in their set of linguistic concepts and prove difficult to scale up,
2138 +L47: the latter tasks usually contain static, non-interactive datasets that lack
2139 +L48: adequate grounding to extra-linguistic concepts [4]. In order to make further
2140 +L49: progress in building grounded language models, we believe there is a need for
2141 +L50: scalable interactive environments that contain: (1) language elements that
2142 +L51: reflect rich, real-world usage and are collectible at scale, and (2) task
2143 +L52: feedback that is well-defined and automatically computable to facilitate
2144 +L53: interactive learning, without the constant need for expensive feedback from
2145 +L54: humans.
2146 +L55:
2147 +L56: The world wide web (WWW) is a massive open-domain interactive environment that
2148 +L57: inherently satisfies the first aforementioned requirement through its
2149 +L58: interconnected set of pages with natural text, images and interactive elements.
2150 +L59: By being simultaneously scalable, semantic, interactive, dynamic and realistic,
2151 +L60: the web is uniquely different from existing environments for autonomous agents
2152 +L61: like games or 3D navigation. Moreover, the web also provides a practical
2153 +L62: environment to deploy trained agents, with great potential for alleviating human
2154 +L63: efforts in tedious tasks (e.g. buying products, booking appointments). While
2155 +L64: there has been prior work on building web-based tasks, they either lack depth in
2156 +L65: the transition and action spaces, or prove difficult to scale up. Some
2157 +L66: benchmarks only contain either a single classification task [39, 46, 31] or
2158 +L67: interactions containing only a handful of different pages in each episode [43].
2159 +L68: Others propose tasks with longer horizons but are either limited to following
2160 +L69: hyperlinks for web navigation [36] or require human-in-the-loop feedback due to
2161 +L70: the lack of an automated reward function [33].
2162 +L71:
2163 +L72: In this paper, we introduce WebShop (Figure 1) – a large-scale interactive web-
2164 +L73: based environment for language understanding and decision making – and train
2165 +L74: autonomous agents to complete tasks on this benchmark. With the goals of being
2166 +L75: scalable and containing realistic language and visual elements, WebShop emulates
2167 +L76: the task of online shopping on an e-commerce website, where the agent’s goal is
2168 +L77: to understand a human-provided text instruction and purchase a product to match
2169 +L78: the specifications. To do so, the agent needs to query the website’s search
2170 +L79: engine, choose items to explore from search results, open and read their
2171 +L80: description and details, and select the necessary options (e.g. 32 oz., red
2172 +L81: color) before clicking the ‘Buy’ button. In order to pick the optimal product
2173 +L82: that matches user requirements, the agent may need to view and compare various
2174 +L83: products (including backtracking between pages), and potentially perform
2175 +L84: multiple searches. WebShop contains over one million products scraped from
2176 +L85: amazon.com, over thousand crowdsourced instructions, and a diverse semantic
2177 +L86: action space of searching text queries and choosing text buttons. It is packaged
2178 +L87: into a convenient OpenAI Gym [5] environment and can be rendered in two modes
2179 +L88: (HTML or simple) with parallel observation spaces that are easy for human and
2180 +L89: model respectively. Rewards are automatically computed using a combination of
2181 +L90: programmatic matching functions that consider the attributes, type, options and
2182 +L91: price of the chosen product, alleviating the need for human evaluation and
2183 +L92: providing a path to scaling up interactive learning.1212
2184 +L93:
2185 +L94: We develop several agents to perform this task, using both reinforcement
2186 +L95: learning (RL) and imitation learning (IL). We also leverage the latest pre-
2187 +L96: trained language models [26, 11] for representing and generating text. Our
2188 +L97: modular architecture includes a factorized processing of state observations and
2189 +L98: action choices using ResNets (visual) and Transformers (text), followed by an
2190 +L99: attention fusion layer that helps the agent contextually score each action. Our
2191 +L100: best agent achieves an average score of (out of 62.462.4) and successfully
2192 +L101: completes the task 100100 of the time, significantly higher than a heuristic
2193 +L102: baseline that achieves 28.7%28.7\% and 45.645.6, respectively. While this
2194 +L103: demonstrates the potential for IL and RL, the agents are still much lower than
2195 +L104: human experts, who can achieve 9.6%9.6\% and 82.182.1 on this task.*** In our
2196 +L105: analysis (§5.3), we observe that the task requires patience and consistency,
2197 +L106: which is lacking in some crowdsource workers, leading to lower scores. Even with
2198 +L107: this caveat, the gap between human performance and the model remains
2199 +L108: significant. We perform several analyses and ablation studies to identify the
2200 +L109: cause of this gap and find several avenues for agent improvement in the future
2201 +L110: including more robust search generation, explicit memory modules, and better
2202 +L111: handling of noisy web text. Finally, we also demonstrate an instance of sim-to-
2203 +L112: real transfer by deploying agents trained with WebShop to operate on amazon.com
2204 +L113: and ebay.com, and find that they can achieve similar performances despite search
2205 +L114: engine and product differences, and consistently outperform the rule baseline
2206 +L115: of using the first result returned by the commercial search engines when
2207 +L116: directly searching the instruction texts. This demonstrates the practical
2208 +L117: potential of our work towards developing agents that can operate autonomously on
2209 +L118: the world wide web (WWW).59.6%59.6\%
2210 +L119:
2211 +L120: ## 2 Related Work
2212 +L121:
2213 +L122: Reinforcement learning on the web. Nogueira and Cho [36] introduced WikiNav as a</result>
2214 +<result>Error fetching URL `http://alfworld.github.io/`</result>
2215 +<result>Error fetching URL `https://osu-nlp-group.github.io/Mind2Web/`</result>
2216 +<result>Error fetching URL `https://github.com/web-arena-x/visualwebarena`</result>
2217 +<result>[30] (https://ar5iv.org/pdf/2406.12172)
2218 +**viewing lines [0 - 127] of 1478**
2219 +
2220 +L0:
2221 +L1: URL: https://ar5iv.org/pdf/2406.12172
2222 +L2: # Navigating the Labyrinth: Evaluating and Enhancing LLMs’ Ability to Reason
2223 +L3: About Search Problems
2224 +L4:
2225 +L5: Nasim Borazjanizadeh
2226 +L6:
2227 +L7: Berkeley AI Research, UC Berkeley
2228 +L8: \AndRoei Herzig
2229 +L9: Berkeley AI Research, UC Berkeley
2230 +L10: \AndTrevor Darrell
2231 +L11: Berkeley AI Research, UC Berkeley
2232 +L12: \AndRogerio Feris
2233 +L13: MIT-IBM Watson AI Lab
2234 +L14: \AndLeonid Karlinsky
2235 +L15: MIT-IBM Watson AI Lab
2236 +L16:
2237 +L17: ###### Abstract
2238 +L18:
2239 +L19: Recently, Large Language Models (LLMs) attained impressive performance in math
2240 +L20: and reasoning benchmarks. However, they still often struggle with logic problems
2241 +L21: and puzzles that are relatively easy for humans. To further investigate this,
2242 +L22: we introduce a new benchmark, SearchBench, containing 11 unique search problems,
2243 +L23: each equipped with automated pipelines to generate an arbitrary number of
2244 +L24: instances and analyze the feasibility, correctness, and optimality of LLM-
2245 +L25: generated solutions. We show that even the most advanced LLMs fail to solve
2246 +L26: these problems end-to-end in text, e.g., GPT4 solves only 1.4%. SearchBench
2247 +L27: problems require considering multiple pathways to the solution as well as
2248 +L28: backtracking, posing a significant challenge to auto-regressive models.
2249 +L29: Instructing LLMs to generate code that solves the problem helps, but only
2250 +L30: slightly, e.g., GPT4’s performance rises to 11.7%. In this work, we show that
2251 +L31: in-context learning with A* algorithm implementations enhances performance. The
2252 +L32: full potential of this promoting approach emerges when combined with our
2253 +L33: proposed Multi-Stage-Multi-Try method, which breaks down the algorithm
2254 +L34: implementation into two stages and verifies the first stage against unit tests,
2255 +L35: raising GPT-4’s performance above 57%.
2256 +L36:
2257 +L37: \doparttoc\faketableofcontents
2258 +L38:
2259 +L39: ### 1 Introduction
2260 +L40:
2261 +L41: The advent of Large Language Models (LLMs) has revolutionized the field of
2262 +L42: natural language processing, with models like Gemini[18], GPT-4[26]
2263 +L43: demonstrating unprecedented performance on reasoning tasks such as GSM8k[8].
2264 +L44: However, these models still exhibit surprising failures on some intuitive
2265 +L45: tasks[2, 30, 22] and struggle with multi-step compositional reasoning,
2266 +L46: combinatorial problems, and planning [9, 40, 44]. Inspired by these observations
2267 +L47: and to further investigate LLMs’ reasoning abilities, we offer a new benchmark
2268 +L48: of search problems, SearchBench. The problems in SearchBench are combinatorial,
2269 +L49: defined as tasks that involve finding an optimal object from a finite set of
2270 +L50: objects, where the set of feasible solutions is either discrete or can be
2271 +L51: reduced to a discrete set [43]. These problems are predominantly NP-hard and
2272 +L52: necessitate systematic exploration of action paths and backtracking to
2273 +L53: intermediate feasible states; thus, SearchBench implicitly investigates the
2274 +L54: LLM’s capacity for non-linear reasoning.
2275 +L55:
2276 +L56: SearchBench has five distinct problem categories: (i) pathfinding, (ii) puzzles,
2277 +L57: (iii) subset sum, (iv) sorting, and (v) under-determined systems; further
2278 +L58: divided into 11 unique problem types. Each problem type is inspired by known
2279 +L59: puzzles and combinatorial problems but augmented with modified rules and
2280 +L60: constraints to ensure substantial differences from similar problems LLMs
2281 +L61: encountered during their training. And the solution to each problem is a
2282 +L62: sequence of actions leading from the initial state to the goal state, while
2283 +L63: optimizing a cost. We generate100 instances of varying difficulty per problem
2284 +L64: type using an automatic pipeline, resulting in 1107 problem instances total.
2285 +L65: Each problem type in SearchBench is equipped with an automatic pipeline that
2286 +L66: evaluates LLM-generated solutions on three dimensions: feasibility, correctness,
2287 +L67: and optimality. Feasibility checks whether the actions taken follow the
2288 +L68: problem’s rules; correctness verifies if a feasible solution reaches the goal
2289 +L69: state; and optimality checks if the least cost solution was found.∼\sim
2290 +L70:
2291 +L71: SearchBench is challenging to LLMs due to several factors. Firstly, natural
2292 +L72: language is less suited for describing or updating accurate representations of
2293 +L73: complex intermediate states. Secondly, our experiments show LLMs struggle with
2294 +L74: exploring a combinatorial exponentially exploding state-space. Despite the fact
2295 +L75: that some methods were developed for long-context reasoning [4, 13, 50],
2296 +L76: SearchBench problems cannot be easily summarized [4], reasoned about [13], or
2297 +L77: processed in parallel due to their size [50, 45]. Our findings show that even
2298 +L78: the strongest LLMs [26] almost completely fail to solve SearchBench problems in
2299 +L79: text-only mode.
2300 +L80:
2301 +L81: To provide further insights, we show that LLMs’ performance on SearchBench
2302 +L82: improves by prompting the models to solve the problems using the A* search
2303 +L83: algorithm [11]. A* is a heuristic-based graph traversal algorithm known for its
2304 +L84: time efficiency and provable optimality guarantees, making it the most suitable
2305 +L85: search algorithm for solving the problems in our benchmark. This method
2306 +L86: leverages A*’s correctness and optimality, while offloading some of the non-
2307 +L87: linear computations involved in searching the state-space to code execution.
2308 +L88: Additionally, to improve the quality of generated A* codes, motivated that
2309 +L89: ensembling helps generation quality[41, 47, 21], we introduce the Multi-Stage-
2310 +L90: Multi-Try (MSMT) inference strategy. In the "Multi-Try" aspect of MSMT, before
2311 +L91: evaluating the solution returned by the code, we first verify whether the code
2312 +L92: generated by the model satisfies a set of unit tests: (i) it is executable; (ii)
2313 +L93: it returns a list as output; and (iii) data type of list elements is correct.
2314 +L94: If the code fails any of the tests, MSMT re-runs the LLM until a valid code is
2315 +L95: generated or allowed number of attempts is exhausted. The "Multi-Stage" aspect
2316 +L96: of MSMT generates the code in two steps: (i) ‘A* Implementation’ - the
2317 +L97: implementation of an instance-agnostic A* algorithm for the problem type; and
2318 +L98: (ii) Initialization - the instantiation of initial conditions and state
2319 +L99: variables of the problem instance. In MSMT ’Initialization’ is generated
2320 +L100: conditioned on the ‘A* Implementation’ (which is generated separately first and
2321 +L101: provided in ‘Intitialization’ prompt). We demonstrate that our MSMT A* method
2322 +L102: (Fig. 2) significantly enhances the LLMs’ ability to solve search problems,
2323 +L103: outperforming all other prompting strategies we used to evaluate models on
2324 +L104: SearchBench, including 0-shot text, 4-shot Chain-of-Thought (CoT)[42] text,
2325 +L105: 0-shot code generation, and 4-shot A* prompting with the naive greedy decoding
2326 +L106: strategy.
2327 +L107:
2328 +L108: To summarize, our main contributions are as follows: (i) We contribute the
2329 +L109: SearchBench benchmark designed to assess the capability of LLMs in solving
2330 +L110: state-based problems requiring combinatorial search; (ii) We introduce the A*
2331 +L111: Multi-Stage-Multi-Try (MSMT) framework that leverages the flexibility and
2332 +L112: natural language comprehension capabilities of LLMs, reliability and structured
2333 +L113: nature of the A* algorithm, along with unitests to verify LLMs’ code generation
2334 +L114: in order to address search problems. (iii) We show a strong advantage of MSMT A*
2335 +L115: prompting for solving search problems robustly with LLMs, suggesting
2336 +L116: interesting future work directions and yet leaving significant room for future
2337 +L117: improvement on SearchBench.
2338 +L118:
2339 +L119: ### 2 SearchBench Benchmark
2340 +L120:
2341 +L121: SearchBench includes five categories of problems: puzzles, subset sum, sorting,
2342 +L122: pathfinding, and under-determined systems, further subdivided into 11 problem
2343 +L123: types and instantiated as 1107 natural language problem instances. Each problem
2344 +L124: type requires a different algorithmic approach to solve. Our problems are
2345 +L125: inspired by well-known (mostly NP-hard) puzzles and search problems, though
2346 +L126: their rules were modified to ensure a solved instance of a problem similar to
2347 +L127: any SearchBench problem was not observed during LLMs’ massive internet-scale </result>
2348 +<result>Error fetching URL `https://ar5iv.labs.ar5iv.org/html/2312.15820`</result>
2349 +<result>Error fetching URL `https://github.com/jxwuyi/WebNav`</result>
2350 +</function_results>
2351 +
2352 +<think>在同时打开多个基准测试相关页面后，我们获得了大量信息。WebVoyager 是一个多模态网页智能体基准测试，专注于在真实网站上进行端到端任务执行，其特点包括多模态输入（屏幕截图和文本）、在15个常用网站上的300个任务，以及使用GPT-4V进行自动评估。使用场景是评估网页智能体在真实世界中的导航和任务完成能力。例题可能是在亚马逊上找到特定产品并将其加入购物车。
2353 +
2354 +WebArena 是一个真实且可复现的网页环境，用于评估自主智能体。其特点包括四个功能完整的网站（电子商务、论坛、协作开发、内容管理）、812项长视野任务，以及专注于功能正确性的评估。使用场景是评估智能体在复杂网页任务中的表现。例题可能是在电子商务网站上购买符合特定标准的商品。
2355 +
2356 +GAIA 是一个针对通用AI助手的基准测试，要求具备推理、多模态处理、网页浏览和工具使用等能力。其特点包括466个问题，对人类来说概念简单但对AI具有挑战性，答案简洁明确，支持自动评估。使用场景是评估AI助手在真实世界问题中的表现。例题可能是“找出2022年法国总统选举中获胜者的出生城市”。
2357 +
2358 +ToolBench 是一个用于评估LLM工具操作能力的基准测试。其特点包括多样化的软件工具、真实世界任务，以及易于使用的基础设施。使用场景是评估智能体使用工具完成任务的能力。例题可能是使用计算器API计算表达式，或使用搜索API查找信息。
2359 +
2360 +HotpotQA 是一个多跳问答基准测试，其特点包括基于Wikipedia的113k个问题-答案对，需要多文档推理，并提供句子级别的支持事实。使用场景是评估模型在回答需要多跳推理的问题时的能力。例题可能是“《了不起的盖茨比》的作者还写过哪些小说？”。
2361 +
2362 +FEVER 是一个事实提取和验证基准测试，其特点包括185,445个基于Wikipedia的声明，需要分类为支持、反驳或信息不足，并需要提供证据句子。使用场景是评估系统验证声称事实的能力。例题可能是“斐济最大的岛屿是考艾岛”，系统需要检索证据来支持或反驳这一说法。
2363 +
2364 +TriviaQA 是一个阅读理解基准测试，其特点包括超过65万个问题-答案-证据三元组，问题由琐事爱好者编写，并附带证据文档。使用场景是评估模型在开放域问答中的能力。例题可能是“哪位演员在《星球大战：原力觉醒》中扮演凯洛·伦？”。
2365 +
2366 +Natural Questions 是一个问答基准测试，其特点包括来自谷歌搜索的真实问题，答案来自Wikipedia，并包含长答案和短答案。使用场景是评估开放域问答系统。例题可能是“谁写了《傲慢与偏见》？”，答案可能是“简·奥斯汀”。
2367 +
2368 +MS MARCO 是一个用于深度学习搜索的基准测试，其特点包括来自Bing的真实搜索查询、人工生成的答案，以及用于段落排序和文档排序的任务。使用场景是训练和评估检索模型。例题可能是“如何冲泡咖啡？”，系统需要检索相关段落来回答。
2369 +
2370 +BEIR 是一个用于零样本评估的异构信息检索基准测试，其特点包括18个数据集、9项任务，涵盖事实核查、问答、对话检索等。使用场景是评估检索模型在零样本设置下的泛化能力。例题可能是给定一个查询，在文档集合中检索相关文档。
2371 +
2372 +MIRACL 是一个多语言信息检索基准测试，其特点涵盖18种语言，专注于ad-hoc检索，并包含人工标注的相关性判断。使用场景是评估多语言检索模型。例题可能是用中文查询“全球变暖的原因”，并检索相关文档。
2373 +
2374 +WebShop 是一个模拟电子商务网站环境，用于评估语言智能体。其特点包括118万个真实产品、12,087条众包文本指令，以及需要导航、搜索和购买的任务。使用场景是评估智能体在电子商务任务中的表现。例题可能是“找到一款价格低于50美元的无线蓝牙耳机”。
2375 +
2376 +ALFWorld 是一个将文本和具身环境对齐的基准测试，用于在文本世界中训练智能体，然后将其转移到具身环境中。其特点包括基于文本的环境，与ALFRED数据集并行，并专注于指令跟随。使用场景是评估智能体在家庭任务中的规划能力。例题可能是“在厨房里找到一个杯子，把它拿到客厅”。
2377 +
2378 +Mind2Web 是一个用于评估网页通用智能体的基准测试，其特点包括来自真实网站的任务、需要多步交互，并专注于语言指令跟随。使用场景是评估智能体在复杂网页任务中的表现。例题可能是“在GitHub上创建一个新的仓库，命名为‘test-repo’”。
2379 +
2380 +VisualWebArena 是一个用于多模态智能体的基准测试，其特点包括真实视觉网页任务、需要理解视觉和文本元素，并评估功能正确性。使用场景是评估多模态智能体在视觉接地网页任务中的表现。例题可能是“在纽约时报网站上找到今天头条新闻的图片”。
2381 +
2382 +SearchBench 是一个用于评估LLM解决搜索问题能力的基准测试，其特点包括11种搜索问题类型、自动生成实例，并评估可行性、正确性和最优性。使用场景是评估LLM在组合搜索问题中的推理能力。例题可能是“在一个迷宫中找到从起点到终点的最短路径”。
2383 +
2384 +WebVLN 是一个网站上的视觉语言导航基准测试，其特点包括基于问题的指令、需要导航网站，并评估导航成功率。使用场景是评估智能体在网站中基于视觉和语言指令进行导航的能力。例题可能是“找到‘联系我们’页面并提取电子邮件地址”。
2385 +
2386 +WebNav 是一个用于目标驱动网页导航的基准测试，其特点包括部分观察环境、需要规划和理解自然语言指令。使用场景是评估智能体在网页导航中的能力。例题可能是“导航到维基百科上关于‘人工智能’的页面”。
2387 +
2388 +现在，我需要整理一份针对搜索智能体的基准测试列表。考虑到用户可能更关注中文资料，我会优先考虑那些有中文描述的基准测试，例如BrowseComp、HSCodeComp、GAIA、WebVoyager、WebArena、Mind2Web、WebShop、ALFWorld、HotpotQA、FEVER、TriviaQA、Natural Questions、MS MARCO、BEIR、MIRACL、SearchBench、WebVLN、WebNav。我会为每个基准测试提供特点、使用场景和例题。同时，我会引用相关来源，例如HSCodeComp论文、BrowseComp博客文章、WebVoyager论文、WebArena论文、GAIA论文、WebShop论文、SearchBench论文等。我的回答将用中文呈现，并确保内容全面、结构清晰。现在开始撰写回答。</think>针对搜索智能体（Search Agent）的评估，已涌现出多个侧重点不同的基准测试。以下将分类介绍这些基准，涵盖其核心特点、典型使用场景及例题。
2389 +
2390 +---
2391 +
2392 +## 🔍 1. 网页浏览/导航类基准
2393 +这类基准要求智能体在真实或模拟的网页环境中执行多步操作，以完成复杂任务。
2394 +
2395 +| 基准 | 核心特点 | 使用场景 | 例题 |
2396 +|------|----------|----------|------|
2397 +| **BrowseComp**（OpenAI）【3†L16-L35】 | - 1,266 个挑战性问题，需持续导航、查找难以直接搜索的信息<br>- 问题设计保证人类在 10 分钟内无法解决，且现有模型（包括带浏览的 ChatGPT）也难以完成<br>- 答案为短字符串，便于自动验证 | 评估智能体在复杂网络浏览任务中的持久性、创造性搜索能力 | “找出 2018‑2023 年间在 EMNLP 会议上发表、第一作者本科毕业于达特茅斯学院、第四作者本科毕业于宾夕法尼亚大学的科学论文标题”【3†L45-L49】 |
2398 +| **WebVoyager**【24†L18-L33】 | - 多模态（视觉+文本）端到端网页智能体基准<br>- 包含 300 个真实任务，覆盖 15 个常用网站（如 Amazon、Wikipedia）<br>- 使用 GPT‑4V 进行自动评估，与人工评估一致性达 85.3% | 评估智能体在真实网站上的端到端任务完成能力（多模态输入） | “在 Amazon 上找到‘无线蓝牙耳机’并按价格从低到高排序，然后选择最便宜的一款加入购物车” |
2399 +| **WebArena**【25†L15-L35】 | - 完全可交互的、高度真实的网页环境（4 个领域：电商、论坛、协作开发、内容管理）<br>- 812 项长视野任务，评估功能正确性而非表面动作匹配<br>- 最佳 GPT‑4 智能体成功率仅 14.41%，远低于人类的 78.24% | 评估自主智能体在真实网页环境中的复杂任务执行能力 | “在电商网站中购买一双‘黑色运动鞋’，价格低于 $80，并将收货地址修改为纽约”【25†L86-L92】 |
2400 +| **Mind2Web**【19†L4-L14】 | - 首个面向通用网页智能体的数据集，任务来自真实网站<br>- 支持在线评估（Online‑Mind2Web），包含 300 个实时任务、136 个流行网站<br>- 强调多步指令跟随与复杂交互 | 评估智能体在任意网站上的通用任务完成能力 | “在 GitHub 上创建一个名为 ‘test‑repo’ 的公开仓库，并添加一个 README.md 文件” |
2401 +| **VisualWebArena**【20†L4-L10】 | - 多模态网页智能体基准，任务依赖视觉信息（截图）<br>- 包含丰富的视觉接地任务，如识别图片、按钮位置等<br>- 评估智能体在真实视觉网页环境中的表现 | 评估多模态智能体在视觉丰富的网页任务中的能力 | “在纽约时报网站找到今天头条新闻的配图，并描述图中主要物体” |
2402 +| **WebVLN**【22†L4-L11】 | - 网站上的视觉‑语言导航任务，智能体需根据问题指令在网站中导航<br>- 数据集 WebVLN‑v1 包含多种网站导航场景<br>- 评估导航成功率与任务完成度 | 评估智能体在网站中基于视觉与语言指令的导航能力 | “找到‘联系我们’页面，并提取其中的邮箱地址” |
2403 +| **WebNav**【23†L4-L11】 | - 目标驱动的网页导航基准，环境为部分观测<br>- 智能体需理解自然语言指令并在网页中规划动作序列<br>- 早期经典任务，常用于测试基础导航能力 | 评估智能体在部分观测网页环境中的规划与指令跟随能力 | “导航到 Wikipedia 上关于‘人工智能’的页面，并摘录第一段内容” |
2404 +
2405 +---
2406 +
2407 +## 🛠️ 2. 工具使用/通用助手类基准
2408 +这类基准测试智能体调用外部工具（如搜索、计算、API）来解决实际问题的能力。
2409 +
2410 +| 基准 | 核心特点 | 使用场景 | 例题 |
2411 +|------|----------|----------|------|
2412 +| **GAIA**【26†L10-L25】 | - 466 个真实世界问题，需要推理、多模态处理、网页浏览、工具使用等综合能力<br>- 对人类简单（92% 准确率），但对最强 GPT‑4 仅 15%<br>- 答案唯一、事实性，支持自动评估 | 评估通用 AI 助手在需要多步工具使用的真实任务中的表现 | “2022 年法国总统选举的获胜者出生在哪个城市？”【26†L10-L25】 |
2413 +| **ToolBench**（OpenBMB）【7†L4-L10】 | - 包含多样软件工具的真实任务基准<br>- 提供易用的基础设施（虚拟 API 服务器等）<br>- 评估 LLM 调用工具完成复杂任务的能力 | 评估智能体在工具调用与组合上的熟练度 | “使用计算器 API 计算 (3.14 * 7.2)^2，然后使用搜索 API 查找‘圆周率’的最新精确值” |
2414 +| **HSCodeComp**【2†L16-L35】 | - 首个针对深度搜索智能体在分层规则应用上的专家级电商基准<br>- 任务：根据嘈杂的商品描述预测 10 位 HS 编码（关税规则）<br>- 632 个商品条目，标注由多位电商专家完成，最佳智能体准确率仅 46.8%（人类 95%） | 评估智能体在复杂规则（如关税分类）下的深层推理与搜索能力 | “给定商品描述‘男士纯棉针织T恤，领口为V领’，预测其 10 位 HS 编码”【2†L24-L32】 |
2415 +
2416 +---
2417 +
2418 +## ❓ 3. 问答/检索类基准
2419 +这类基准主要评估智能体在开放域或特定领域的信息检索与问答能力。
2420 +
2421 +| 基准 | 核心特点 | 使用场景 | 例题 |
2422 +|------|----------|----------|------|
2423 +| **HotpotQA**【8†L4-L10】 | - 113k 个基于 Wikipedia 的多跳问答对<br>- 需要跨多个文档推理，并提供句子级支持事实<br>- 涵盖多样的问题类型（比较、列举、因果等） | 评估模型在需要多跳推理的开放域问答中的能力 | “《了不起的盖茨比》的作者还写过哪些小说？” |
2424 +| **FEVER**【27†L16-L28】 | - 185,445 个基于 Wikipedia 的声明，需分类为 Supported/Refuted/NotEnoughInfo<br>- 要求提供证据句子（可多句、多页面）<br>- 挑战性高（最佳系统仅 31.87% 准确率） | 评估系统在事实核查与证据检索上的能力 | “斐济最大的岛屿是考艾岛。”【27†L48-L50】 |
2425 +| **TriviaQA**【28†L5-L10】 | - 超过 650k 个问答‑证据三元组，问题由琐事爱好者编写<br>- 每个问题平均提供 6 篇证据文档，适合远程监督阅读<br>- 包含阅读理解和开放域 QA 两种任务设置 | 评估模型在开放域琐事问答中的检索与阅读理解能力 | “哪位演员在《星球大战：原力觉醒》中扮演凯洛·伦？” |
2426 +| **Natural Questions**【11†L4-L8】 | - 来自 Google 搜索的真实用户问题，答案来自 Wikipedia<br>- 包含长答案（段落）和短答案（实体/日期等）<br>- 训练集 307k，开发/测试集各 8k | 评估开放域问答系统对真实用户查询的响应能力 | “谁写了《傲慢与偏见》？”（答案：“简·奥斯汀”） |
2427 +| **MS MARCO**【12†L4-L10】 | - 基于 Bing 真实搜索查询的检索/问答基准<br>- 包含段落排序、文档排序、问答等任务<br>- 数据规模大（约 500k 查询，880 万段落） | 训练与评估检索模型在真实搜索场景下的表现 | “如何冲泡咖啡？”（系统需检索相关段落回答） |
2428 +| **BEIR**【13†L4-L10】 | - 异构零样本检索基准，包含 18 个数据集、9 种任务（事实核查、QA、对话检索等）<br>- 评估模型在未见任务上的泛化能力<br>- 已成为检索模型的标准评估套件 | 评估检索模型在零样本设置下的跨任务泛化能力 | 给定查询“全球变暖的主要原因”，在文档集合中检索相关文档 |
2429 +| **MIRACL**【14†L4-L10】 | - 多语言检索基准，覆盖 18 种语言（包括中文、阿拉伯语等）<br>- 专注于 ad‑hoc 检索，所有查询与文档均经过人工标注<br>- 提供多语言检索挑战（WSDM 2023 Cup） | 评估多语言检索模型在不同语言下的检索效果 | 中文查询：“全球变暖的原因”，检索相关中文文档 |
2430 +
2431 +---
2432 +
2433 +## 🧩 4. 组合搜索/推理类基准
2434 +这类基准专门测试智能体在组合搜索问题上的逻辑推理与规划能力。
2435 +
2436 +| 基准 | 核心特点 | 使用场景 | 例题 |
2437 +|------|----------|----------|------|
2438 +| **SearchBench**【30†L22-L36】 | - 11 种独特的搜索问题类型（路径查找、谜题、子集和、排序、欠定系统等）<br>- 自动生成任意数量实例，并评估解决方案的可行性、正确性、最优性<br>- GPT‑4 在纯文本模式下仅解决 1.4%，通过 A* 代码生成提升至 11.7% | 评估 LLM 在组合搜索问题上的推理与规划能力 | “在一个 5x5 迷宫中，从起点 (0,0) 到终点 (4,4) 的最短路径是什么？（障碍物位置已知）”【30†L56-L64】 |
2439 +
2440 +---
2441 +
2442 +## 🛒 5. 特定领域/仿真环境基准
2443 +这类基准针对特定领域（如电商、家庭环境）设计，测试智能体在仿真环境中的交互能力。
2444 +
2445 +| 基准 | 核心特点 | 使用场景 | 例题 |
2446 +|------|----------|----------|------|
2447 +| **WebShop**【29†L15-L30】 | - 模拟电商网站环境，包含 118 万真实商品、12,087 条众包指令<br>- 智能体需根据文本指令搜索、筛选、购买商品<br>- 奖励自动计算，最佳模型成功率 29%（人类 59%） | 评估语言智能体在电商任务中的导航、搜索与决策能力 | “找到一款价格低于 $50 的无线蓝牙耳机，评分至少 4 星，并加入购物车”【29†L20-L25】 |
2448 +| **ALFWorld**【18†L4-L10】 | - 将文本环境（TextWorld）与具身环境（ALFRED）对齐的基准<br>- 智能体先在文本中学习抽象策略，然后在视觉环境中执行<br>- 包含 134 个家庭任务（如“在厨房找到杯子并拿到客厅”） | 评估智能体在文本与具身环境之间的迁移学习与规划能力 | “在厨房找到‘杯子’，把它拿到客厅放在桌子上” |
2449 +
2450 +---
2451 +
2452 +## 📊 总结
2453 +以上基准覆盖了搜索智能体的不同能力维度：从**网页浏览、工具使用、问答检索**，到**组合搜索推理**以及**特定领域仿真**。选择合适的基准取决于您要评估的具体能力（如多模态交互、工具调用、多跳推理等）。这些基准大多提供公开的数据集与评估代码，可直接用于实验。
2454 +
2455 +> 注：本文基于 2025 年 11 月前的公开资料整理，部分基准可能已有更新版本，建议访问对应官网获取最新信息。<｜end▁of▁sentence｜>
\ No newline at end of file
rust/src/chat/src/renderer/deepseek_v32/fixtures/test_output_search_wo_date.txt new
+1069
@@ -0,0 +1,1069 @@
1 +<｜begin▁of▁sentence｜><｜User｜>
2 +
3 +## Tools
4 +
5 +You have access to a set of tools you can use to answer the user's question.
6 +You can invoke functions by writing a "<｜DSML｜function_calls>" block like the following as part of your reply to the user:
7 +<｜DSML｜function_calls>
8 +<｜DSML｜invoke name="$FUNCTION_NAME">
9 +<｜DSML｜parameter name="$PARAMETER_NAME" string="true|false">$PARAMETER_VALUE</｜DSML｜parameter>
10 +...
11 +</｜DSML｜invoke>
12 +<｜DSML｜invoke name="$FUNCTION_NAME2">
13 +...
14 +</｜DSML｜invoke>
15 +</｜DSML｜function_calls>
16 +
17 +String and scalar parameters should be specified as is without any escaping or quotes, while lists and objects should use JSON format. The "string" attribute should be set to "true" for string type parameters and "false" for other types (numbers, booleans, arrays, objects).
18 +
19 +If the thinking_mode is enabled, then after function results you should strongly consider outputting a thinking block. Here is an example:
20 +
21 +<｜DSML｜function_calls>
22 +...
23 +</｜DSML｜function_calls>
24 +
25 +<function_results>
26 +...
27 +</function_results>
28 +
29 +<think>...thinking about results</think>
30 +
31 +Here are the functions available in JSONSchema format:
32 +<functions>
33 +{"name": "search", "description": "Searches for information related to query and displays topn results.", "parameters": {"type": "object", "properties": {"query": {"type": "string", "description": "The search query string"}, "topn": {"type": "integer", "description": "Number of top results to display", "default": 10}, "source": {"type": "string", "description": "Source to search within", "enum": ["web", "news"], "default": "web"}}, "required": ["query"], "additionalProperties": false, "$schema": "http://json-schema.org/draft-07/schema#"}}
34 +{"name": "open", "description": "Opens the link id from the page indicated by cursor starting at line number loc, showing num_lines lines. Valid link ids are displayed with the formatting: 【{id}†.*】. If cursor is not provided, the most recent page is implied. If id is a string, it is treated as a fully qualified URL associated with source. If loc is not provided, the viewport will be positioned at the beginning of the document or centered on the most relevant passage, if available. Use this function without id to scroll to a new location of an opened page.", "parameters": {"type": "object", "properties": {"id": {"description": "Link ID (number) or fully qualified URL (string)", "anyOf": [{"type": "integer"}, {"type": "string"}], "default": -1}, "cursor": {"type": "integer", "description": "Page cursor indicator", "default": -1}, "loc": {"type": "integer", "description": "Starting line number", "default": -1}, "num_lines": {"type": "integer", "description": "Number of lines to show", "default": -1}, "view_source": {"type": "boolean", "description": "Whether to view source", "default": false}, "source": {"type": "string", "description": "Source associated with the URL", "enum": ["web", "news"], "default": "web"}}, "additionalProperties": false, "$schema": "http://json-schema.org/draft-07/schema#"}}
35 +{"name": "find", "description": "Finds exact matches of pattern in the current page, or the page given by cursor.", "parameters": {"type": "object", "properties": {"pattern": {"type": "string", "description": "The exact pattern to search for"}, "cursor": {"type": "integer", "description": "Page cursor to search in", "default": -1}}, "required": ["pattern"], "additionalProperties": false, "$schema": "http://json-schema.org/draft-07/schema#"}}
36 +{"name": "python", "description": "Use this tool to execute Python code. The code will not be shown to the user. This tool should be used for internal reasoning, but not for code that is intended to be visible to the user (e.g. when creating plots, tables, or files). When you send a message containing Python code to python, it will be executed in a stateful Jupyter notebook environment. python will respond with the output of the execution or time out after 120.0 seconds.", "parameters": {"type": "object", "properties": {"code": {"type": "string", "description": "The python code to execute in the jupyter environment."}}, "required": ["code"], "additionalProperties": false, "$schema": "http://json-schema.org/draft-07/schema#"}}
37 +</functions>
38 +
39 +
40 +# The user's message is: In a year before 2023, between December 1 and 15 (inclusive), a press release reported that as of the 5th of that month, a country had received a total of 16,201,670 vaccine doses, with 7,583,134 administered. That same year, a company's annual report revealed that its nearly 2.2 million members and supporters contributed over $49 million, with 12% of total revenue coming from contributed services. The following year, the annual report stated that the organization had played a role in rescinding a previous administration’s plan to open most of the 22 million acres of the National Petroleum Reserve in Alaska to oil and gas drilling and that contributed services had decreased by 2% to 10%. Who was the emeritus board member that year? Name only, no titles.<｜Assistant｜><think>The user asks: "In a year before 2023, between December 1 and 15 (inclusive), a press release reported that as of the 5th of that month, a country had received a total of 16,201,670 vaccine doses, with 7,583,134 administered. That same year, a company's annual report revealed that its nearly 2.2 million members and supporters contributed over $49 million, with 12% of total revenue coming from contributed services. The following year, the annual report stated that the organization had played a role in rescinding a previous administration’s plan to open most of the 22 million acres of the National Petroleum Reserve in Alaska to oil and gas drilling and that contributed services had decreased by 2% to 10%. Who was the emeritus board member that year? Name only, no titles."
41 +
42 +We need to parse the puzzle.
43 +
44 +We have clues about a press release about vaccine doses for a country, dated between Dec 1-15 inclusive, reporting as of 5th of that month total doses received 16,201,670 and administered 7,583,134. That is likely a COVID-19 vaccine update for some country. Which country? Might be a smaller country. Could be something like "Uganda", "Zimbabwe", "Kenya", "Philippines", etc. Let's search.
45 +
46 +Search for "16,201,670 vaccine doses" "7,583,134 administered".</think>
47 +
48 +<｜DSML｜function_calls>
49 +<｜DSML｜invoke name="search">
50 +<｜DSML｜parameter name="query" string="true">"16,201,670" vaccine doses</｜DSML｜parameter>
51 +<｜DSML｜parameter name="topn" string="false">10</｜DSML｜parameter>
52 +<｜DSML｜parameter name="source" string="true">web</｜DSML｜parameter>
53 +</｜DSML｜invoke>
54 +</｜DSML｜function_calls><｜end▁of▁sentence｜>
55 +
56 +<function_results>
57 +<result>[0] "16,201,670" vaccine doses (Search_Results/"16,201,670" vaccine doses)
58 +**viewing lines [0 - 38] of 38**
59 +
60 +L0:
61 +L1: URL: Search_Results/"16,201,670" vaccine doses
62 +L2: # Search Results
63 +L3:
64 +L4: * 【0†Kenya's Economy is Showing Resilience as Output Rises ...; publish_date:
65 +L5: none†www.worldbank.org】 Dec 14, 2021 — As of December 5, 2021, Kenya had
66 +L6: received a total of 16,201,670 vaccines, with 7,583,134 administered. While
67 +L7: vaccine acceptance is ...
68 +L8: * 【1†MINISTRY OF HEALTH; publish_date: none†covidhub.mediacouncil.or.ke】 Dec
69 +L9: 1, 2021 — Total Doses Received 16,201,670. Table 10 gives the total vaccines
70 +L10: received since the start of Covid -19 vaccination exercise in the country.
71 +L11: * 【2†Output Result Page; publish_date: none†open.unicef.org】 ... 16,201,670
72 +L12: doses of multiple vaccines nationwide and full vaccination of 15.5 per cent with
73 +L13: two doses of COVID-19 vaccine as of 31 December 2021.
74 +L14: * 【3†rebased GDP; publish_date: none†documents1.worldbank.org】 Dec 7, 2021 —
75 +L15: As of December 5, 2021,. Kenya had received a total of 16,201,670 vaccines,
76 +L16: with. 7,583,134 administered. Vaccine acceptance is reportedly high.
77 +L17: * 【4†Integrated Annual Report; publish_date: none†www.co-opbank.co.ke】 May 27,
78 +L18: 2022 — ... doses of Covid-19 vaccines and administered close to 17 million ...
79 +L19: 16,201,670, huku 7,583,134 zikiwa tayari zimedungwa watu. Bado kuna ...
80 +L20: * 【5†World Bank lifts Kenya's growth prospect to 5% in 2021; publish_date:
81 +L21: none†www.africa-press.net】 ... 16,201,670 vaccines, with 7,583,134 administered.
82 +L22: While vaccine acceptance is reportedly high there is still a long way to go
83 +L23: towards the government's ...
84 +L24: * 【6†2020 Annual Report United States; publish_date: none†www.hi-us.org】
85 +L25: including medicine, vaccines, Covid-19 prevention kits, and kits of ...
86 +L26: 16,201,670. 9,630,584. Prepaid expenses. 56,707. 22,601. Total current assets.
87 +L27: * 【7†HOSPITAL CENTRAL DE LA DEFENSA; publish_date:
88 +L28: none†contrataciondelestado.es】 16201670 X. URPA-URPA. 0. HCD-000393/00. KONTRON.
89 +L29: 7000 ... LASER SHOT LBP-1120. CNEUR-C. NEUROFISIOLOGÍA. 0 ... VICTOREEN CT HEAD
90 +L30: DOSE PHANTON. 43650-014.
91 +L31: * 【8†body metal concentrations: Topics by ...; publish_date:
92 +L32: none†www.science.gov】 PMID:16201670. Design Characteristics and Tobacco Metal
93 +L33: Concentrations in Filtered Cigars. PubMed. Caruso, Rosalie V; O'Connor, Richard
94 +L34: J; Travers, Mark J ...
95 +L35: * 【9†Public Offices, Personal Demands: Capability in ...; publish_date:
96 +L36: none†dokumen.pub】 ... (16201670).9 The Institutiones Politicae consists of two
97 +L37: books. Book one deals with the origin, nature and goal of the commonwealth on
98 +L38: the one hand, and the ...</result>
99 +</function_results>
100 +
101 +<think>Result shows Kenya. So the press release is about Kenya's vaccine doses as of Dec 5, 2021. The press release likely by World Bank or something. The search result shows "World Bank lifts Kenya's growth prospect..." maybe a press release dated Dec 14, 2021. But the clue says "between December 1 and 15 (inclusive), a press release reported that as of the 5th of that month, a country had received a total of 16,201,670 vaccine doses, with 7,583,134 administered." That matches Kenya on Dec 5, 2021. So the year is 2021.
102 +
103 +Now the next clue: "That same year, a company's annual report revealed that its nearly 2.2 million members and supporters contributed over $49 million, with 12% of total revenue coming from contributed services."
104 +
105 +This sounds like an environmental or conservation organization. "Nearly 2.2 million members and supporters contributed over $49 million." That might be something like "The Nature Conservancy", "World Wildlife Fund", "Sierra Club", "Environmental Defense Fund", "Natural Resources Defense Council", etc. Let's search.
106 +
107 +Search for "nearly 2.2 million members and supporters contributed over $49 million" "12% of total revenue coming from contributed services".</think>
108 +
109 +<｜DSML｜function_calls>
110 +<｜DSML｜invoke name="search">
111 +<｜DSML｜parameter name="query" string="true">"nearly 2.2 million members" "contributed over $49 million"</｜DSML｜parameter>
112 +<｜DSML｜parameter name="topn" string="false">10</｜DSML｜parameter>
113 +<｜DSML｜parameter name="source" string="true">web</｜DSML｜parameter>
114 +</｜DSML｜invoke>
115 +</｜DSML｜function_calls><｜end▁of▁sentence｜>
116 +
117 +<function_results>
118 +<result>[1] "nearly 2.2 million members" "contributed over $49 million" (Search_Results/"nearly 2.2 million members" "contributed over $49 million")
119 +**viewing lines [0 - 41] of 41**
120 +
121 +L0:
122 +L1: URL: Search_Results/"nearly 2.2 million members" "contributed over $49 million"
123 +L2: # Search Results
124 +L3:
125 +L4: * 【0†The Global Good Fund, Microsoft and Goodwill join forces ...;
126 +L5: publish_date: none†globalgoodfund.org】 Sep 28, 2022 — ... more than 2.2 million
127 +L6: people. They additionally contribute close to $49 billion in local, state and
128 +L7: federal tax revenues. When COVID hit ...
129 +L8: * 【1†Almost 22 billion American tax dollars spent to wipe out a ...;
130 +L9: publish_date: none†www.facebook.com】 US military funding for Israel's war crimes
131 +L10: in Lebanon and Gaza has now cost US taxpayers over $22 billion. When millions
132 +L11: struggle to afford the ...
133 +L12: * 【2†Corporate America has largely abandoned its post-January ...;
134 +L13: publish_date: none†www.citizensforethics.org】 Jul 29, 2025 — Since the January 6
135 +L14: insurrection, over 2,000 corporate and industry group PACs have given over $174
136 +L15: million to members of the Sedition ...
137 +L16: * 【3†Audit shows millions in questionable taxpayer spending at ...;
138 +L17: publish_date: none†www.aol.com】 18 hours ago — ... nearly doubled from 1.3
139 +L18: million to about 2.2 million. That is more than one in four Washington state
140 +L19: residents receiving Medicaid, and the ...
141 +L20: * 【4†Incarceration and Poverty in the United States - AAF; publish_date:
142 +L21: none†www.americanactionforum.org】 Jun 30, 2020 — The United States currently
143 +L22: incarcerates 2.2 million people, nearly half of whom are non-violent drug
144 +L23: offenders, accused people held pre ...
145 +L24: * 【5†How the NRA Spends Revenue (2023); publish_date: none†paddockpost.com】
146 +L25: Apr 25, 2025 — In 2023, the NRA reported total revenue of $176 million (compared
147 +L26: to $211 million in 2022, $227 million in 2021, $282 million in 2020, $292 ...
148 +L27: * 【6†NCA Newsroom; publish_date: none†www.ncausa.org】 ... 2.2 million U.S.
149 +L28: jobs—operating in every U.S. state and territory and contributing nearly $350
150 +L29: billion to the U.S. economy every year. For media inquiries ...
151 +L30: * 【7†The Early Childhood Educator Workforce; publish_date:
152 +L31: none†cscce.berkeley.edu】 About 2.2 million adults are paid to care for and
153 +L32: educate more than 9.7 million children between birth and age five in center-
154 +L33: based and home-based settings.
155 +L34: * 【8†Bringing Out the Worst in Congress: CRA By the Numbers ...; publish_date:
156 +L35: none†progressivereform.org】 Jul 14, 2025 — In fact, the top ten recipients (out
157 +L36: of 25) account for almost $7 million, showing how these contributions are
158 +L37: targeted and strategic. Real ...
159 +L38: * 【9†Club 100: Over 100 Americans gave more than $100000 to ...; publish_date:
160 +L39: none†publicintegrity.org】 Jan 7, 2022 — Through Internet investments during the
161 +L40: 1990s, he attracted 2.2 million online accounts valued at $174 million. Schwab
162 +L41: has been a vocal ...</result>
163 +</function_results>
164 +
165 +<think>Not that. Try search for "2.2 million members and supporters contributed over $49 million".</think>
166 +
167 +<｜DSML｜function_calls>
168 +<｜DSML｜invoke name="search">
169 +<｜DSML｜parameter name="query" string="true">"2.2 million members" "contributed over $49"</｜DSML｜parameter>
170 +<｜DSML｜parameter name="topn" string="false">10</｜DSML｜parameter>
171 +<｜DSML｜parameter name="source" string="true">web</｜DSML｜parameter>
172 +</｜DSML｜invoke>
173 +</｜DSML｜function_calls><｜end▁of▁sentence｜>
174 +
175 +<function_results>
176 +<result>[2] "2.2 million members" "contributed over $49" (Search_Results/"2.2 million members" "contributed over $49")
177 +**viewing lines [0 - 2] of 2**
178 +
179 +L0:
180 +L1: URL: Search_Results/"2.2 million members" "contributed over $49"
181 +L2: # Search Results</result>
182 +</function_results>
183 +
184 +<think>Maybe it's an annual report of a non-profit. Search for "annual report 2.2 million members contributed $49 million".</think>
185 +
186 +<｜DSML｜function_calls>
187 +<｜DSML｜invoke name="search">
188 +<｜DSML｜parameter name="query" string="true">"2.2 million members" "annual report" $49 million</｜DSML｜parameter>
189 +<｜DSML｜parameter name="topn" string="false">10</｜DSML｜parameter>
190 +<｜DSML｜parameter name="source" string="true">web</｜DSML｜parameter>
191 +</｜DSML｜invoke>
192 +</｜DSML｜function_calls><｜end▁of▁sentence｜>
193 +
194 +<function_results>
195 +<result>[3] "2.2 million members" "annual report" $49 million (Search_Results/"2.2 million members" "annual report" $49 million)
196 +**viewing lines [0 - 36] of 36**
197 +
198 +L0:
199 +L1: URL: Search_Results/"2.2 million members" "annual report" $49 million
200 +L2: # Search Results
201 +L3:
202 +L4: * 【0†20-F; publish_date: none†www.sec.gov】 ANNUAL REPORT PURSUANT TO SECTION
203 +L5: ... Our membership grew from 2.1 million members as at December 31, 2023 to 2.2
204 +L6: million members as at December 31, 2024.
205 +L7: * 【1†Oportun Reports Fourth Quarter and Full Year 2023 ...; publish_date:
206 +L8: none†investor.oportun.com】 Mar 12, 2024 — Oportun (Nasdaq: OPRT) is a mission-
207 +L9: driven fintech that puts its 2.2 million members' financial goals within reach.
208 +L10: ... annual report on ...
209 +L11: * 【2†2 0 21; publish_date: none†www.annualreports.com】 ANNUAL REPORT. 2. 0.
210 +L12: 21. 2. 0. 21. Page 2. 2. DEFENDERS OF WILDLIFE. 2. 0. 21. 2. 0. 21 ... In 2021,
211 +L13: Defenders of Wildlife's nearly 2.2 million members and.
212 +L14: * 【3†Annual report and accounts 2020; publish_date: none†www.3i.com】
213 +L15: Disclaimer. The Annual report and accounts have been prepared solely to provide
214 +L16: information to shareholders. ... 2.2 million members. In December 2019, we sold
215 +L17: ...
216 +L18: * 【4†united states securities and exchange commission; publish_date:
217 +L19: none†s21.q4cdn.com】 Dec 14, 2018 — ☒ ANNUAL REPORT ... approximately 5.7 million
218 +L20: Members, including approximately 2.2 million Members from the acquisition of
219 +L21: Best Doctors.
220 +L22: * 【5†securities and exchange commission; publish_date: none†www.sec.gov】 ...
221 +L23: Annual Report on Form 10-K filed with the SEC on March 10, 2020. The ... 2.2
222 +L24: million members at June 30, 2020. The following table highlights the ...
223 +L25: * 【6†2022-23 Annual Comprehensive Financial Report - CalPERS; publish_date:
224 +L26: none†www.calpers.ca.gov】 Nov 17, 2023 — ... annual report is consistent with
225 +L27: these accounting principles ... 2.2 million members and health benefit services
226 +L28: to over 1.5 million ...
227 +L29: * 【7†Medical Centers Report; publish_date:
228 +L30: none†regents.universityofcalifornia.edu】 Oct 8, 2019 — which will allow its 2.2
229 +L31: million members to receive advanced and ... annual report of the University of
230 +L32: California. For the Medical ...
231 +L33: * 【8†Working TogeTher to make HealtH care better; publish_date:
232 +L34: none†www.unitedhealthgroup.com】 Feb 8, 2012 — ANNUAL REPORT PURSUANT TO SECTION
233 +L35: 13 OR 15(d) OF THE ... approximately 2.2 million members enrolled in its
234 +L36: Medicare. Advantage ...</result>
235 +</function_results>
236 +
237 +<think>Result 2 shows "Defenders of Wildlife". That might be the organization. Let's open.</think>
238 +
239 +<｜DSML｜function_calls>
240 +<｜DSML｜invoke name="open">
241 +<｜DSML｜parameter name="cursor" string="false">3</｜DSML｜parameter>
242 +<｜DSML｜parameter name="id" string="false">2</｜DSML｜parameter>
243 +</｜DSML｜invoke>
244 +</｜DSML｜function_calls><｜end▁of▁sentence｜>
245 +
246 +<function_results>
247 +<result>[4] www.annualreports.com (https://www.annualreports.com/HostedData/AnnualReportArchive/d/defenders-wildlife_2021.pdf)
248 +**viewing lines [0 - 53] of 261**
249 +
250 +L0:
251 +L1: URL: https://www.annualreports.com/HostedData/AnnualReportArchive/d/defenders-
252 +L2: wildlife_2021.pdf
253 +L3: ===== Page 1 ===== 2021 ANNUAL REPORT ===== Page 2 ===== # DEFENDERS OF WILDLIFE
254 +L4: made important progress for imperiled species and vital landscapes across the
255 +L5: United States in 2021. \--- **LAWYERS** immediately **challenged** the premature
256 +L6: and reckless decision to strip **gray wolves** of federal **Endangered Species
257 +L7: Act (ESA)** protections. For many decades, Defenders has led the effort to
258 +L8: protect and restore the gray wolf, and we will continue to fight the
259 +L9: unscientific and hostile anti-wolf policies that impede conservation progress
260 +L10: and will carry on our unrelenting battle to restore federal protections for this
261 +L11: iconic keystone species. \--- **LOBBYISTS** worked around the clock to keep
262 +L12: wildlife and climate priorities in the **Infrastructure Investment and Jobs
263 +L13: Act**. We also continue fighting to keep important wildlife and habitat funding
264 +L14: in relevant **appropriations bills**. \--- 2 DEFENDERS OF WILDLIFE ===== Page 3
265 +L15: ===== POLICY EXPERTS pushed forward on the urgent need for a National
266 +L16: Biodiversity Strategy (NBS), an all-of-government approach to address the
267 +L17: unprecedented loss of wildlife and habitat we are experiencing. We have coupled
268 +L18: this with our new campaign to expand the National Wildlife Refuge System to
269 +L19: preserve our nation’s only lands set aside for wildlife. By defending, funding
270 +L20: and expanding our national wildlife refuges, we will directly address
271 +L21: biodiversity loss and climate change while promoting increased equitable access
272 +L22: to nature. FIELD TEAMS were on the ground helping to recover imperiled species.
273 +L23: From panthers and sea turtles in Florida to wolves, bison and black-footed
274 +L24: ferrets in Montana, Defenders’ conservation experts were in the field saving
275 +L25: wildlife all over the country. CONSERVATION INNOVATION EXPERTS provided
276 +L26: comprehensive analyses to guide policy and inform conservation strategies to
277 +L27: reach the goal of protecting 30% of our terrestrial and marine systems by 2030
278 +L28: (“30x30”). Defenders’ Center for Conservation Innovation (CCI) produced a report
279 +L29: which details actions we need to take to achieve 30x30 while protecting
280 +L30: biodiversity and addressing the climate crisis. DEFENDERS.ORG ===== Page 4 =====
281 +L31: WE HAVE ACCOMPLISHED MUCH THIS YEAR WORKING WITH AN ADMINISTRATION THAT VALUES
282 +L32: SCIENCE AND CONSERVATION. Our many successes include the return of protections
283 +L33: to the Tongass National Forest in Alaska and the suspension of oil and gas
284 +L34: leasing permits for the Arctic National Wildlife Refuge. Defenders also played a
285 +L35: lead role in successfully defending the Migratory Bird Protection Act from
286 +L36: attack and securing critical protections for migratory birds. Throughout 2021
287 +L37: Defenders made important progress for the wildlife and wild places we all love,
288 +L38: yet our nation’s wildlife still face unprecedented challenges. More and more
289 +L39: species, including manatees, Mexican gray wolves, polar bears and North Atlantic
290 +L40: right whales, face the very real threat of extinction because of climate
291 +L41: change, habitat loss, pollution and inadequate protections. In our work we
292 +L42: continue to demonstrate success by following the science, building partnerships,
293 +L43: leading with innovative solutions and focusing on collaborative and practical
294 +L44: approaches in defense of nature. Together we can make a lasting difference for
295 +L45: wildlife and habitats. Your support is essential to our work and we couldn’t do
296 +L46: it without you. Thank you for being an important part of our efforts to protect
297 +L47: imperiled wildlife across our nation. Jamie Rappaport Clark President and CEO
298 +L48: John Dayton Board Chair OUR NEW TELEVISION SERIES, WILDLIFE NATION WITH JEFF
299 +L49: CORWIN, premiered in October and airs every weekend on ABC stations nationwide.
300 +L50: This groundbreaking series takes Jeff across the country to meet Defenders’
301 +L51: experts along with our science, tribal and community partners to educate and
302 +L52: inspire viewers to join with Defenders to protect wildlife. While highlighting
303 +L53: some of our important conservation work and raising awareness of important </result>
304 +</function_results>
305 +
306 +<think>Search within PDF for "2.2 million members".</think>
307 +
308 +<｜DSML｜function_calls>
309 +<｜DSML｜invoke name="find">
310 +<｜DSML｜parameter name="cursor" string="false">4</｜DSML｜parameter>
311 +<｜DSML｜parameter name="pattern" string="true">2.2</｜DSML｜parameter>
312 +</｜DSML｜invoke>
313 +</｜DSML｜function_calls><｜end▁of▁sentence｜>
314 +
315 +<function_results>
316 +<result>[5] Find results for text: `2.2` in `www.annualreports.com` (https://www.annualreports.com/HostedData/AnnualReportArchive/d/defenders-wildlife_2021.pdf/find?pattern=2.2)
317 +**viewing lines [0 - 4] of 4**
318 +
319 +L0: # 【0†match at L187】
320 +L1: 2021, Defenders of Wildlife’s nearly 2.2 million members and supporters
321 +L2: contributed more than $49 million for wildlife and wildlife habitat. This figure
322 +L3: includes all those who generously donated their time and expertise. The
323 +L4: steadfast support of our donors allows Defenders to sustain our program and </result>
324 +</function_results>
325 +
326 +<think>Open around L180.</think>
327 +
328 +<｜DSML｜function_calls>
329 +<｜DSML｜invoke name="open">
330 +<｜DSML｜parameter name="cursor" string="false">4</｜DSML｜parameter>
331 +<｜DSML｜parameter name="loc" string="false">170</｜DSML｜parameter>
332 +</｜DSML｜invoke>
333 +</｜DSML｜function_calls><｜end▁of▁sentence｜>
334 +
335 +<function_results>
336 +<result>[6] www.annualreports.com (https://www.annualreports.com/HostedData/AnnualReportArchive/d/defenders-wildlife_2021.pdf)
337 +**viewing lines [170 - 215] of 261**
338 +
339 +L170: into the wild. **10 DEFENDERS OF WILDLIFE** ===== Page 11 ===== AFTER ANOTHER
340 +L171: SUCCESSFUL DEFENDERS LAWSUIT ON BEHALF OF THE CRITICALLY ENDANGERED RED WOLF,
341 +L172: FWS reversed its decision to limit the recovery area and committed to a robust
342 +L173: release strategy. After years of inaction, FWS released eight wolves to the wild
343 +L174: in North Carolina and plan to release nine more wolves in the spring of 2022.
344 +L175: This is an incredible win for this critically endangered species whose
345 +L176: population has dwindled down to single digits in the wild because of
346 +L177: mismanagement, vehicle strikes and poaching. DEFENDERS CONTINUED TO LEAD EFFORTS
347 +L178: TO PROTECT THE FLORIDA MANATEE, a beloved species that suffered the deadliest
348 +L179: year on record in 2021, tragically surpassing 1,000 deaths because of water
349 +L180: pollution and lack of warm water habitat. Defenders led advocacy and education
350 +L181: aimed at restoring the natural flow of the dammed Ocklawaha River, which would
351 +L182: provide critical warm-water habitat that manatees need to survive. Defenders’
352 +L183: legal team continued to fight for manatees in the courts, holding government
353 +L184: agencies accountable for protecting critical habitat and addressing the
354 +L185: devastating water pollution that is killing the seagrass and causing manatees to
355 +L186: starve. DAVID TES | SAM FRENZY DRAW DEFENDERS.ORG 11 ===== Page 12 ===== In
356 +L187: 2021, Defenders of Wildlife’s nearly 2.2 million members and supporters
357 +L188: contributed more than $49 million for wildlife and wildlife habitat. This figure
358 +L189: includes all those who generously donated their time and expertise. The
359 +L190: steadfast support of our donors allows Defenders to sustain our program and
360 +L191: public education efforts in the field, the courts and on Capitol Hill. 2021
361 +L192: SOURCES OF FUNDS Grants and contributions $29,057 Bequests, trusts and split
362 +L193: interests $7,692 Income from investments, annuity reserve funds and trusts
363 +L194: $3,354 Royalties and other income $3,576 Contributed services $6,140 Total
364 +L195: Revenue $49,819 USES OF FUNDS Biodiversity conservation $22,420 Constituency
365 +L196: mobilization $16,324 Fundraising $1,211 Management and general $5,865 Total
366 +L197: Expenses $45,820 Change in net assets $3,999 Net assets, start of year $41,145
367 +L198: Net Assets, End of the Year $45,144 Dollars are in thousands. 12 DEFENDERS OF
368 +L199: WILDLIFE Grants and contributions 58% Income from investments 7% Requests,
369 +L200: trusts and split interests 15% Royalties and other income 7% Contributed
370 +L201: services 12% Program and support services 97% Fundraising 3% ===== Page 13 =====
371 +L202: DEFENDERS.ORG 13 ===== Page 14 ===== John Dayton, Esq.* Chair Texas Sheryl
372 +L203: Sachs* Vice Chair Connecticut Whitfield Bailey* Secretary Tennessee Mark Caylor*
373 +L204: Treasurer District of Columbia Sylvia Baca* California Donald Barry, Esq.
374 +L205: Oregon Dinah Bear, Esq. Arizona Kim O’Keefe Beck, Ph.D California Gabriela
375 +L206: Chavarria, Ph.D Washington Jamie Rappaport Clark* President & CEO Virginia
376 +L207: Caroline D. Gabel Maryland Keith M. Harper, Esq. Maryland John Donovan, Esq.
377 +L208: Montana and California James M. Hecker, Esq. Virginia Judith Ponsikoff, Ph.D*
378 +L209: California Terry L. Root, Ph.D Florida BOARD EMERITUS Terry C. Pelster, Esq. New
379 +L210: York *Executive Committee member* ST A F F & A D V I S O R S SCIENCE ADVISORS
380 +L211: Gabriela Chavarria, Ph.D Burke Museum, University of Washington Leah Gerber,
381 +L212: Ph.D Arizona State University Kathayoon Khalil, Ph.D Oregon Zoo Rurik List, Ph.D
382 +L213: Universidad Autónoma Metropolitana-Lerma Barry Noon, Ph.D Colorado State
383 +L214: University Terry L. Root, Ph.D Stanford University Mark Schwartz, Ph.D
384 +L215: University of California-Davis David S. Wilcox, Ph.D Princeton University Rae </result>
385 +</function_results>
386 +
387 +<think>Scrolling further for "Contributed services 12%". Already seen. So the annual report is Defenders of Wildlife 2021 annual report.
388 +
389 +Next clue: "The following year, the annual report stated that the organization had played a role in rescinding a previous administration’s plan to open most of the 22 million acres of the National Petroleum Reserve in Alaska to oil and gas drilling and that contributed services had decreased by 2% to 10%."
390 +
391 +Thus, the following year (2022) annual report of Defenders of Wildlife includes mention of rescinding plan to open most of 22 million acres of NPR-A to oil and gas drilling. Also contributed services decreased to 10% (from 12% in 2021). So we need to find the 2022 annual report for Defenders of Wildlife.
392 +
393 +Search for "Defenders of Wildlife 2022 annual report National Petroleum Reserve Alaska".</think>
394 +
395 +<｜DSML｜function_calls>
396 +<｜DSML｜invoke name="search">
397 +<｜DSML｜parameter name="query" string="true">Defenders of Wildlife 2022 annual report National Petroleum Reserve Alaska</｜DSML｜parameter>
398 +<｜DSML｜parameter name="topn" string="false">10</｜DSML｜parameter>
399 +<｜DSML｜parameter name="source" string="true">web</｜DSML｜parameter>
400 +</｜DSML｜invoke>
401 +</｜DSML｜function_calls><｜end▁of▁sentence｜>
402 +
403 +<function_results>
404 +<result>[7] Defenders of Wildlife 2022 annual report National Petroleum Reserve Alaska (Search_Results/Defenders of Wildlife 2022 annual report National Petroleum Reserve Alaska)
405 +**viewing lines [0 - 37] of 37**
406 +
407 +L0:
408 +L1: URL: Search_Results/Defenders of Wildlife 2022 annual report National Petroleum
409 +L2: Reserve Alaska
410 +L3: # Search Results
411 +L4:
412 +L5: * 【0†CELEBRATING YEARS; publish_date: none†www.annualreports.com】 With less
413 +L6: than 340 right whales left in the wild, Defenders is fighting tirelessly to end
414 +L7: deadly lobster gear entanglements and vessel strikes that are driving ...
415 +L8: * 【1†Financials; publish_date: none†defenders.org】 We invite you to explore
416 +L9: the reports below to learn more about our activities and accomplishments, and
417 +L10: how we put your money to work for wildlife.
418 +L11: * 【2†Alaska Program Looks Back on 2022; publish_date: none†defenders.org】 Feb
419 +L12: 9, 2023 — Thanks to a lawsuit joined by Defenders, seven million acres were
420 +L13: returned to protection within the National Petroleum Reserve-Alaska (NPR-A), ...
421 +L14: * 【3†Defenders-of-Wildlife-2022-Financial-Statement. ...; publish_date:
422 +L15: none†defenders.org】 We have audited the accompanying consolidated financial
423 +L16: statements of Defenders of Wildlife and Affiliated Defenders of Wildlife Action
424 +L17: Fund (collectively, ...
425 +L18: * 【4†2022 Annual Report; publish_date: none†alaskaconservation.org】 Jun 13,
426 +L19: 2023 — In 2022, we focused on three landscapes: the Arctic. National Wildlife
427 +L20: Refuge, Bristol Bay, and the Tongass National Forest. In March 2022,.
428 +L21: * 【5†Assessment of ecological and cultural values within the ...;
429 +L22: publish_date: none†www.blm.gov】 This document was written to provide technical
430 +L23: information regarding the ecological importance of the National Petroleum
431 +L24: Reserve – Alaska (NPR-A). Several ...
432 +L25: * 【6†Accomplishments Report; publish_date: none†defenders.org】 National
433 +L26: Petroleum Reserve-Alaska. Identified for exceptional wildlife and cultural
434 +L27: values, including critical habitat for polar bears and other species ...
435 +L28: * 【7†2022 annual report; publish_date: none†dory-
436 +L29: plantain-s2zc.squarespace.com】 These projects are made possible through
437 +L30: collaborations with Defenders of. Wildlife, Cook Inletkeeper, Trustees for
438 +L31: Alaska,. Environmental Investigation Agency, ...
439 +L32: * 【8†23IMPACT REPORT; publish_date: none†www.annualreports.com】 Defenders of
440 +L33: Wildlife made incredible progress protecting wildlife and wild places in 2023,
441 +L34: helping shape a brighter future for imperiled species and vital ...
442 +L35: * 【9†From Leasing to Land Protections; publish_date: none†defenders-cci.org】
443 +L36: Abstract. When you hear the name, National Petroleum Reserve - Alaska (NPR-A, or
444 +L37: Reserve), you might think of a barren land filled with oil wells and ...</result>
445 +</function_results>
446 +
447 +<think>Open result 0.</think>
448 +
449 +<｜DSML｜function_calls>
450 +<｜DSML｜invoke name="open">
451 +<｜DSML｜parameter name="cursor" string="false">7</｜DSML｜parameter>
452 +<｜DSML｜parameter name="id" string="false">0</｜DSML｜parameter>
453 +</｜DSML｜invoke>
454 +</｜DSML｜function_calls><｜end▁of▁sentence｜>
455 +
456 +<function_results>
457 +<result>[8] www.annualreports.com (https://www.annualreports.com/HostedData/AnnualReportArchive/d/defenders-wildlife_2022.pdf)
458 +**viewing lines [0 - 53] of 289**
459 +
460 +L0:
461 +L1: URL: https://www.annualreports.com/HostedData/AnnualReportArchive/d/defenders-
462 +L2: wildlife_2022.pdf
463 +L3: ===== Page 1 ===== 2022 Impact Report C E L E B R A T I N G Y E A R S ===== Page
464 +L4: 2 ===== 2022 Defenders of Wildlife made important progress for imperiled
465 +L5: species and vital landscapes across the United States in 2022. GRAY WOLF | JIM
466 +L6: GUMMERAL MAY STOCK PRIOR Lawyers successfully challenged the previous
467 +L7: administration’s decision to delist the gray wolf and restored critical federal
468 +L8: protections under the Endangered Species Act. This latest triumph in court is
469 +L9: part of our ongoing battle to protect and restore gray wolves throughout their
470 +L10: historical range and shield them from persecution by extremist legislators in
471 +L11: Idaho, Montana and Wyoming. TWO MORE FATALIZED GRAY SWALLETS TO SEA TO SHARE
472 +L12: ALLIANCE Lobbyists worked around the clock to expand funding for wildlife
473 +L13: conservation in the FY2022 federal spending bill, which included $31 million (a
474 +L14: 44% increase) for the Bureau of Land Management’s Threatened and Endangered
475 +L15: Species Program, $2.5 million (an 81% increase) for the U.S. Department of
476 +L16: Agriculture Wildlife Services’ Nonlethal Initiative to prevent human-wildlife
477 +L17: conflicts and $21 million (a 320% increase) for North Atlantic right whale
478 +L18: conservation. 2 DEFENDERS OF WILDLIFE ===== Page 3 ===== **Policy Experts**
479 +L19: played a crucial role in securing international trade protections for 100
480 +L20: species of sharks and rays, all 158 species of glass frogs and 73 species of
481 +L21: reptiles, including 21 species of desert horned lizards, at the Convention on
482 +L22: International Trade in Endangered Species (CITES) in Panama. \--- **Field
483 +L23: Teams** worked tirelessly to protect and restore imperiled species across the
484 +L24: country. From Florida manatees and red wolves in the Southeast to belugas and
485 +L25: grizzly bears in Alaska, Defenders’ conservation experts were on the ground
486 +L26: saving species that need our help to survive and thrive. \--- **Conservation
487 +L27: Innovation Experts** published more than 10 peer-reviewed studies on topics that
488 +L28: include the Cook Inlet beluga whale, golden-cheeked warbler, global parrot
489 +L29: biodiversity, the Endangered Species Act, the effects of mountaintop removal
490 +L30: mining on endangered species, the ecological importance of panthers and the
491 +L31: implementation of “30x30” – the globally recognized goal to which President
492 +L32: Biden committed the U.S. to conserve 30% of our imperiled lands and waters by
493 +L33: 2030. \--- **DEFENDERS.ORG** ===== Page 4 ===== THANK YOU Defenders celebrated
494 +L34: our 75th anniversary in 2022—an exciting milestone that we attribute to our
495 +L35: unwavering dedication to our wildlife conservation mission. From helping to pass
496 +L36: the Endangered Species Act and other bedrock conservation laws to leading the
497 +L37: advocacy on the reintroduction of wolves to Yellowstone National Park and other
498 +L38: parts of the West, Defenders is proud to inspire, lead and ultimately celebrate
499 +L39: some of the most consequential conservation wins for wildlife and wild places.
500 +L40: It is wonderfully fitting that we also mark 2022 as a year when we successfully
501 +L41: defended, yet again, the original focal species and symbol of our organization,
502 +L42: the gray wolf. While the future for gray wolves and many other species remains
503 +L43: uncertain, we will never stop fighting to protect the imperiled wildlife that
504 +L44: need us and to restore vibrant and balanced ecosystems. Our mission is more
505 +L45: important than ever due to the escalating biodiversity and climate crises.
506 +L46: Defenders is determined to build on our accomplishments and seize new
507 +L47: opportunities to address these challenges through our multifaceted approach to
508 +L48: conservation. We focus on developing innovative and science-based solutions,
509 +L49: coordinating on-the-ground species and habitat restoration, promoting human-
510 +L50: wildlife coexistence, protecting and strengthening our bedrock environmental
511 +L51: laws, investing in impactful collaborations and partnerships, building a more
512 +L52: diverse and powerful conservation movement and challenging attacks on wildlife
513 +L53: through advocacy, policy and litigation. Defenders’ continued success is made </result>
514 +</function_results>
515 +
516 +<think>Search within PDF for "National Petroleum Reserve".</think>
517 +
518 +<｜DSML｜function_calls>
519 +<｜DSML｜invoke name="find">
520 +<｜DSML｜parameter name="cursor" string="false">8</｜DSML｜parameter>
521 +<｜DSML｜parameter name="pattern" string="true">Petroleum</｜DSML｜parameter>
522 +</｜DSML｜invoke>
523 +</｜DSML｜function_calls><｜end▁of▁sentence｜>
524 +
525 +<function_results>
526 +<result>[9] Find results for text: `petroleum` in `www.annualreports.com` (https://www.annualreports.com/HostedData/AnnualReportArchive/d/defenders-wildlife_2022.pdf/find?pattern=petroleum)
527 +**viewing lines [0 - 4] of 4**
528 +
529 +L0: # 【0†match at L201】
530 +L1: the 22 million acres of the National Petroleum Reserve in Alaska to oil and gas
531 +L2: drilling and removed protections for important wildlife habitat. Protections
532 +L3: have now been restored for nearly half of the reserve’s pristine lands, which
533 +L4: are vital habitat for shorebirds, denning polar bears and tens of thousands of </result>
534 +</function_results>
535 +
536 +<think>Open around L190.</think>
537 +
538 +<｜DSML｜function_calls>
539 +<｜DSML｜invoke name="open">
540 +<｜DSML｜parameter name="cursor" string="false">8</｜DSML｜parameter>
541 +<｜DSML｜parameter name="loc" string="false">180</｜DSML｜parameter>
542 +</｜DSML｜invoke>
543 +</｜DSML｜function_calls><｜end▁of▁sentence｜>
544 +
545 +<function_results>
546 +<result>[10] www.annualreports.com (https://www.annualreports.com/HostedData/AnnualReportArchive/d/defenders-wildlife_2022.pdf)
547 +**viewing lines [180 - 227] of 289**
548 +
549 +L180: the sixth successful transfer of bison to the Assiniboine and Sioux Tribes of
550 +L181: Fort Peck since 2019. \--- **SWIFT FIX KITS | © ROSIMA PAELARINTSKIMMA MADDIAL
551 +L182: 200 AND CONSERVATION BIOLOGY INSTITUTE** \--- **Celebrating our third year** of
552 +L183: a collaborative program with the Aaniih and Nakoda Tribes and others to restore
553 +L184: swift foxes to the Fort Belknap Indian Reservation in Montana, Defenders helped
554 +L185: with the release of 28 more swift foxes. With over 100 foxes reintroduced
555 +L186: through this program, monitoring efforts show that they are reproducing in the
556 +L187: wild—a critical measure of success for a self-sustaining population. \---
557 +L188: **Defenders continued to lead the way** for conserving and recovering the
558 +L189: endangered black-footed ferret, supporting the black-footed ferret survey for
559 +L190: the Fort Belknap Indian community. Thirty-six ferrets were vaccinated against
560 +L191: sylvatic plague and two dozen kits were released in the wild. \--- **10
561 +L192: DEFENDERS OF WILDLIFE** ===== Page 11 ===== Defenders helped to bring hope for
562 +L193: recovery for the endangered military macaw, adding 11 fledglings to a growing
563 +L194: wild population in Puerta Vallarta, Mexico, that is under pressure from habitat
564 +L195: loss and poachers for the illegal pet trade. Accord- ing to our recent report,
565 +L196: the 2008 parrot trade ban that Defenders fought to achieve is working.
566 +L197: Preventing more than 30,000 parrots from being illegally trapped each year, the
567 +L198: trade ban has resulted in a 47% decrease in the illegal trade of parrots and an
568 +L199: 88% decrease in U.S. seizures of Mexican parrots. As a result of a Defenders
569 +L200: lawsuit, BLM rescinded the previous administration’s plan that opened most of
570 +L201: the 22 million acres of the National Petroleum Reserve in Alaska to oil and gas
571 +L202: drilling and removed protections for important wildlife habitat. Protections
572 +L203: have now been restored for nearly half of the reserve’s pristine lands, which
573 +L204: are vital habitat for shorebirds, denning polar bears and tens of thousands of
574 +L205: caribou in the Teshekpuk caribou herd. Our new storymap illustrates why the
575 +L206: reserve is so important to wildlife. Defenders marked the 20th anniversary of
576 +L207: our Sea Otter Awareness Week. In celebration of this milestone, we sponsored
577 +L208: state legislation, signed by Gov. Gavin Newsom, that formally recognizes the
578 +L209: event’s anniversary and acknowledges the vital role that sea otters play in
579 +L210: California’s nearshore ecosystems. This annual celebration provides a unique
580 +L211: opportunity to raise awareness of the important role sea otters play in
581 +L212: maintaining ecosystem health, our need to coexist with sea otters and the
582 +L213: efforts of Defenders and others in advancing sea otter conservation.
583 +L214: DEFENDERS.ORG ===== Page 12 ===== FINANCIAL REPORT In 2022, Defenders of
584 +L215: Wildlife’s nearly 2.2 million members and supporters contributed more than $41
585 +L216: million for wildlife and wildlife habitat. This figure includes all those who
586 +L217: generously donated their time and expertise. The steadfast support of our donors
587 +L218: allows Defenders to sustain our program and public education efforts in the
588 +L219: field, the courts and on Capitol Hill. 2022 SOURCES OF FUNDS Grants and
589 +L220: contributions $28,946 Bequests, trusts and split interests $13,928 Income from
590 +L221: investments, annuity reserve funds and trusts (-$7,495) Royalties and other
591 +L222: income $744 Contributed services $5,047 Total Revenue $41,170 USES OF FUNDS
592 +L223: Biodiversity conservation $18,970 Constituency mobilization $16,528 Fundraising
593 +L224: $1,029 Management and general $6,354 Total Expenses $42,881 Change in net assets
594 +L225: (-$1,711) Net assets, start of year $45,144 Net Assets, End of the Year $43,433
595 +L226: Dollars are in thousands. 12 DEFENDERS OF WILDLIFE Grants and contributions 59%
596 +L227: Contributed services 10% Royalties and other income 2% Bequests, trusts and </result>
597 +</function_results>
598 +
599 +<think>Scrolling a bit more for "Contributed services 10%". Already seen. So indeed the 2022 annual report shows contributed services 10%, down from 12% in 2021.
600 +
601 +Thus the organization is Defenders of Wildlife.
602 +
603 +The question: "Who was the emeritus board member that year? Name only, no titles."
604 +
605 +We need to identify the emeritus board member listed in the 2022 annual report. From the 2021 annual report, we saw "BOARD EMERITUS Terry C. Pelster, Esq." However, that is for 2021. The question asks "Who was the emeritus board member that year?" That year refers to the following year (2022). So we need to check the 2022 annual report for the emeritus board member.
606 +
607 +Open the 2022 PDF near the board list. Search for "Emeritus".</think>
608 +
609 +<｜DSML｜function_calls>
610 +<｜DSML｜invoke name="find">
611 +<｜DSML｜parameter name="cursor" string="false">8</｜DSML｜parameter>
612 +<｜DSML｜parameter name="pattern" string="true">Emeritus</｜DSML｜parameter>
613 +</｜DSML｜invoke>
614 +</｜DSML｜function_calls><｜end▁of▁sentence｜>
615 +
616 +<function_results>
617 +<result>[11] Find results for text: `emeritus` in `www.annualreports.com` (https://www.annualreports.com/HostedData/AnnualReportArchive/d/defenders-wildlife_2022.pdf/find?pattern=emeritus)
618 +**viewing lines [0 - 10] of 10**
619 +
620 +L0: # 【0†match at L237】
621 +L1: Sisk, Ph.D British Columbia, Canada BOARD EMERITUS Terry C. Pelster, Esq. New
622 +L2: York *Executive Committee member* STAFF & ADVISORS SCIENCE ADVISORS Gabriela
623 +L3: Chavarria, Ph.D Burke Museum, University of Washington Leah Gerber, Ph.D Arizona
624 +L4: State University Kathayoon Khalil, Ph.D Oregon Zoo Rurik List, Ph.D Universidad
625 +L5:
626 +L6: # 【1†match at L243】
627 +L7: Davis Thomas D. Sisk, Ph.D Emeritus Professor Northern Arizona University David
628 +L8: S. Wilcox, Ph.D Princeton University Rae Wynn-Grant, Ph.D National Geographic
629 +L9: Society SENIOR STAFF Jamie Rappaport Clark President & CEO Thu Pham Chief of
630 +L10: Staff James Stofan Senior Vice President, Operations McCrystle Adams Vice </result>
631 +</function_results>
632 +
633 +<think>Open around L230.</think>
634 +
635 +<｜DSML｜function_calls>
636 +<｜DSML｜invoke name="open">
637 +<｜DSML｜parameter name="cursor" string="false">8</｜DSML｜parameter>
638 +<｜DSML｜parameter name="loc" string="false">220</｜DSML｜parameter>
639 +</｜DSML｜invoke>
640 +</｜DSML｜function_calls><｜end▁of▁sentence｜>
641 +
642 +<function_results>
643 +<result>[12] www.annualreports.com (https://www.annualreports.com/HostedData/AnnualReportArchive/d/defenders-wildlife_2022.pdf)
644 +**viewing lines [220 - 267] of 289**
645 +
646 +L220: contributions $28,946 Bequests, trusts and split interests $13,928 Income from
647 +L221: investments, annuity reserve funds and trusts (-$7,495) Royalties and other
648 +L222: income $744 Contributed services $5,047 Total Revenue $41,170 USES OF FUNDS
649 +L223: Biodiversity conservation $18,970 Constituency mobilization $16,528 Fundraising
650 +L224: $1,029 Management and general $6,354 Total Expenses $42,881 Change in net assets
651 +L225: (-$1,711) Net assets, start of year $45,144 Net Assets, End of the Year $43,433
652 +L226: Dollars are in thousands. 12 DEFENDERS OF WILDLIFE Grants and contributions 59%
653 +L227: Contributed services 10% Royalties and other income 2% Bequests, trusts and
654 +L228: split interests 29% TOTAL REVENUE Income from investments -15% Program and
655 +L229: support services 98% USE OF FUNDS Fundraising 2% ===== Page 13 =====
656 +L230: DEFENDERS.ORG 13 ===== Page 14 ===== BOARD OF DIRECTORS John Dayton, Esq.* Chair
657 +L231: Texas Sheryl Sachs* Vice Chair Connecticut Dinah Bear, Esq.* Secretary Arizona
658 +L232: Mark Caylor* Treasurer District of Columbia Susan Alexander Connecticut Sylvia
659 +L233: Baca* California Whitefield Bailey Tennessee Donald Barry, Esq. Oregon Kim
660 +L234: O’Keefe Beck, Ph.D California Gabriela Chavarria, Ph.D Washington Jamie
661 +L235: Rappaport Clark* President & CEO Virginia John Donovan, Esq.* Montana and
662 +L236: California Caroline D. Gabel Maryland Terry L. Root, Ph.D New Mexico Thomas D.
663 +L237: Sisk, Ph.D British Columbia, Canada BOARD EMERITUS Terry C. Pelster, Esq. New
664 +L238: York *Executive Committee member* STAFF & ADVISORS SCIENCE ADVISORS Gabriela
665 +L239: Chavarria, Ph.D Burke Museum, University of Washington Leah Gerber, Ph.D Arizona
666 +L240: State University Kathayoon Khalil, Ph.D Oregon Zoo Rurik List, Ph.D Universidad
667 +L241: Autónoma Metropolitana-Lerma Barry Noon, Ph.D Colorado State University Terry
668 +L242: L. Root, Ph.D Stanford University Mark Schwartz, Ph.D University of California-
669 +L243: Davis Thomas D. Sisk, Ph.D Emeritus Professor Northern Arizona University David
670 +L244: S. Wilcox, Ph.D Princeton University Rae Wynn-Grant, Ph.D National Geographic
671 +L245: Society SENIOR STAFF Jamie Rappaport Clark President & CEO Thu Pham Chief of
672 +L246: Staff James Stofan Senior Vice President, Operations McCrystle Adams Vice
673 +L247: President, Conservation Law Matt Braughler Vice President, Philanthropy Rachel
674 +L248: Brittin Vice President, Communications Shawn Cantrell Vice President, Field
675 +L249: Conservation Programs Robert L. Dewey Vice President, Government Relations &
676 +L250: External Affairs Monica Goldberg Vice President, Landscape Conservation Sharon
677 +L251: Morris-Smith Vice President of Human Resources Elizabeth Nichols Vice President,
678 +L252: Finance/CFO Anthony Pegues Vice President, Diversity, Equity, Inclusion,
679 +L253: Justice Lindsay Rosa, Ph.D Vice President, Conservation Research and Innovation
680 +L254: Brandace Swimley Vice President, Integrated Marketing 14 DEFENDERS OF WILDLIFE
681 +L255: ===== Page 15 ===== There are many ways to support Defenders of Wildlife and
682 +L256: protect our imperiled wildlife. To learn more or donate, contact us at
683 +L257: 1-800-385-9712, visit Defenders.org/donate or choose one of the following
684 +L258: options. ➤ Transfer appreciated assets directly to Defenders and maximize
685 +L259: potential tax benefits. Defenders.org/ donate-stocks-help-save-wildlife ➤ Become
686 +L260: a Leadership Council member with a single donation of $1,000 or more and enjoy
687 +L261: more personal involvement with Defenders. Defenders.org/leadershipcouncil ➤ If
688 +L262: you are 70.5 or older, roll over a portion of your IRA to benefit wildlife.
689 +L263: Remember, these gifts must come directly from your IRA account.
690 +L264: Defenders.org/planned-giving/gifts-your-retirement-plan ➤ Get started on a
691 +L265: contribution from your donor-advised fund Defenders.org/donor-advised-funds ➤
692 +L266: Check with your personnel office or that of your spouse/partner about your
693 +L267: company’s Corporate Matching Gifts program or visit us at </result>
694 +</function_results>
695 +
696 +<think>Thus the emeritus board member is Terry C. Pelster. However, the question might ask for "Who was the emeritus board member that year?" Possibly referring to the year of the annual report that mentioned the rescinded plan and contributed services decrease. That is 2022. So the answer is Terry C. Pelster.
697 +
698 +But let's verify that the 2022 annual report indeed includes the mention of contributed services decreasing to 10%. We saw that. So answer is Terry C. Pelster.
699 +
700 +But maybe the question is tricky: "Who was the emeritus board member that year?" Could be the emeritus board member listed in the 2021 annual report? However, the question says "That same year, a company's annual report revealed that its nearly 2.2 million members and supporters contributed over $49 million, with 12% of total revenue coming from contributed services. The following year, the annual report stated that the organization had played a role in rescinding a previous administration’s plan to open most of the 22 million acres of the National Petroleum Reserve in Alaska to oil and gas drilling and that contributed services had decreased by 2% to 10%."
701 +
702 +Thus the "following year" annual report is 2022. The question likely asks: "Who was the emeritus board member that year?" Meaning the emeritus board member listed in that annual report (2022). So answer: Terry C. Pelster.
703 +
704 +But let's confirm that the 2021 annual report also shows emeritus board member Terry C. Pelster. Indeed we saw that. But perhaps the 2022 report might have a different emeritus board member. We saw it's the same. So answer is Terry C. Pelster.
705 +
706 +However, the question might be about a different organization. Let's verify the details.
707 +
708 +First clue: Press release about Kenya vaccine doses on Dec 5, 2021. That is likely a World Bank press release. Search for "World Bank Kenya vaccine doses December 2021 press release".</think>
709 +
710 +<｜DSML｜function_calls>
711 +<｜DSML｜invoke name="search">
712 +<｜DSML｜parameter name="query" string="true">As of December 5, 2021, Kenya had received a total of 16,201,670 vaccines</｜DSML｜parameter>
713 +<｜DSML｜parameter name="topn" string="false">10</｜DSML｜parameter>
714 +<｜DSML｜parameter name="source" string="true">web</｜DSML｜parameter>
715 +</｜DSML｜invoke>
716 +</｜DSML｜function_calls><｜end▁of▁sentence｜>
717 +
718 +<function_results>
719 +<result>[13] As of December 5, 2021, Kenya had received a total of 16,201,670 vaccines (Search_Results/As of December 5, 2021, Kenya had received a total of 16,201,670 vaccines)
720 +**viewing lines [0 - 41] of 41**
721 +
722 +L0:
723 +L1: URL: Search_Results/As of December 5, 2021, Kenya had received a total of
724 +L2: 16,201,670 vaccines
725 +L3: # Search Results
726 +L4:
727 +L5: * 【0†Kenya's Economy is Showing Resilience as Output Rises ...; publish_date:
728 +L6: none†www.worldbank.org】 Dec 14, 2021 — As of December 5, 2021, Kenya had
729 +L7: received a total of 16,201,670 vaccines, with 7,583,134 administered. While
730 +L8: vaccine acceptance is ...
731 +L9: * 【1†Unmet need for COVID-19 vaccination coverage in Kenya - PMC;
732 +L10: publish_date: none†pmc.ncbi.nlm.nih.gov】 by SK Muchiri · 2022 · Cited by 42 — As
733 +L11: of December 2021, six counties had a vaccination coverage of less than 5%.
734 +L12: These counties include Garissa, Mandera, Marsabit, Tana River, Turkana, and ...
735 +L13: * 【2†MINISTRY OF HEALTH; publish_date: none†covidhub.mediacouncil.or.ke】 Dec
736 +L14: 1, 2021 — • Total Covid-19 Vaccines Received to date- 16,201,670 ... Table 10:
737 +L15: Vaccine Logistics Received in the Country as at 5th, December 2021.
738 +L16: * 【3†COVID-19 vaccination refusal trends in Kenya over 2021 - PMC;
739 +L17: publish_date: none†pmc.ncbi.nlm.nih.gov】 by RT Rego · 2023 · Cited by 21 — We
740 +L18: assessed vaccine refusal over time in Kenya, and characterized factors
741 +L19: associated with changes in vaccination refusal.
742 +L20: * 【4†Ciheb-Kenya on the Front Lines of the COVID-19 Vaccine ...; publish_date:
743 +L21: none†www.ciheb.org】 Of the 721,509 persons vaccinated, 409,711 are aged 58
744 +L22: years and above, 143,050 are healthcare workers, 109,593 are teachers, and
745 +L23: 59,155 are security officers.
746 +L24: * 【5†proportion of fully vaccinated adults 23.3pc.; publish_date:
747 +L25: none†www.facebook.com】 COVID-19: KENYA administers 150632 new vaccinations;
748 +L26: total jabs given now 13782963; proportion of fully vaccinated adults 23.3pc.
749 +L27: * 【6†National COVID-19 Vaccines Deployment and ...; publish_date:
750 +L28: none†www.parliament.go.ke】 Kenya has a well-established immunization program and
751 +L29: a robust disease surveillance and response unit (DSRU) which serves as an early
752 +L30: warning system, to.
753 +L31: * 【7†Kenya: Vaccine Requirements Violate Rights; publish_date:
754 +L32: none†www.hrw.org】 Dec 13, 2021 — Kenya, with an estimated adult population of
755 +L33: 27.2 million and a total population of 55 million, has received approximately 23
756 +L34: million doses as ...
757 +L35: * 【8†Government gives 20 million a month to get vaccinated; publish_date:
758 +L36: none†www.bbc.com】 Nov 22, 2021 — Although less than 10% of Kenyans are
759 +L37: vaccinated, the government wants to avoid a surge over Christmas.
760 +L38: * 【9†ministry of health kenya covid-19 immunization status ...; publish_date:
761 +L39: none†www.health.go.ke】 May 15, 2022 — By the end of 2021 Kenya had administered
762 +L40: 10.1 million doses of 5 types of Covid-19 vaccines. During this period 4.2
763 +L41: million people were ...</result>
764 +</function_results>
765 +
766 +<think>Open result 0.</think>
767 +
768 +<｜DSML｜function_calls>
769 +<｜DSML｜invoke name="open">
770 +<｜DSML｜parameter name="cursor" string="false">13</｜DSML｜parameter>
771 +<｜DSML｜parameter name="id" string="false">0</｜DSML｜parameter>
772 +</｜DSML｜invoke>
773 +</｜DSML｜function_calls><｜end▁of▁sentence｜>
774 +
775 +<function_results>
776 +<result>[14] www.worldbank.org (https://www.worldbank.org/en/news/press-release/2021/12/14/kenya-s-economy-is-showing-resilience-as-output-rises-above-pre-pandemic-levels-driven-by-a-rebound-in-the-services-sect)
777 +**viewing lines [0 - 52] of 61**
778 +
779 +L0:
780 +L1: URL: https://www.worldbank.org/en/news/press-release/2021/12/14/kenya-s-economy-
781 +L2: is-showing-resilience-as-output-rises-above-pre-pandemic-levels-driven-by-a-
782 +L3: rebound-in-the-services-sect
783 +L4: NAIROBI, December 14, 2021 — Kenya’s economy has demonstrated resilience to the
784 +L5: COVID-19 shock, with output in the first half of the year rising above pre-
785 +L6: pandemic levels. In 2021 as a whole, gross domestic product (GDP) is expected to
786 +L7: grow by 5%, one of the faster recoveries among Sub-Saharan African countries.
787 +L8: Overall economic performance is expected to be robust at 4.9% per year in
788 +L9: 2022-23, similar to the pre-pandemic pace (5% average annual growth from 2010 to
789 +L10: 2019). According to the 24th edition of the Kenya Economic Update, “From
790 +L11: Recovery to Better Jobs,” growth has been supported by rebounds in industry and,
791 +L12: especially, services. Agricultural output, however, fell by 0.5% year on year
792 +L13: in the first half of 2021 following a particularly strong performance in 2020,
793 +L14: partly due to below-average rains. Demand-side recovery has been supported by a
794 +L15: revival in private consumption, against a backdrop of improving employment
795 +L16: conditions and household incomes. “Kenya’s economy has shown considerable
796 +L17: resilience to the enormous shock of the pandemic, and this year is expected to
797 +L18: post one of the stronger growth rebounds in the region thanks to diversified
798 +L19: sources of growth and sound economic policies and management,” said Keith
799 +L20: Hansen, World Bank Country Director for Kenya. “However, poverty has increased,
800 +L21: and the buffers and coping mechanisms of households, firms, and the public
801 +L22: finances have been depleted.” Economic activity in Kenya has continued to adapt
802 +L23: to the pandemic and associated restrictions. A mix of containment measures, such
803 +L24: as a nightly curfew, were in effect through most of 2021, while more
804 +L25: economically disruptive measures such as lockdowns and travel restrictions were
805 +L26: phased, limiting the impact on economic activities. The vaccine rollout, which
806 +L27: had a slow start due to supply constraints, has picked up as new shipments of
807 +L28: vaccines have arrived, particularly since September. This has supported economic
808 +L29: recovery and growth through the third quarter of 2021. As of December 5, 2021,
809 +L30: Kenya had received a total of 16,201,670 vaccines, with 7,583,134 administered.
810 +L31: While vaccine acceptance is reportedly high there is still a long way to go
811 +L32: towards the government’s target of fully inoculating the adult population of
812 +L33: about 30 million by the end of 2022.As of December 6, 2021, about 10% of adults
813 +L34: (2.9 million people) had been fully vaccinated while another 16% (4.9 million
814 +L35: people) had received their first dose. The report also projects robust growth in
815 +L36: the medium-term. This outlook takes into account that some sub-sectors, such as
816 +L37: education, have bounced back strongly, but others such as international
817 +L38: tourism, have only partially and face a much more protracted recovery. The
818 +L39: anticipated further recovery of hotels and restaurants, trade, transport, and
819 +L40: other services, depends on substantial vaccination progress to help prevent new
820 +L41: waves of infections and associated containment measures. “Recent economic
821 +L42: performance has been strong and the outlook is positive, but in Kenya, as
822 +L43: everywhere, the future course of the pandemic remains a major source of
823 +L44: uncertainty, said Alex Sienaert, Senior Economist for Kenya. “To support further
824 +L45: recovery, reduce debt distress risks, and rebuild space for social and
825 +L46: development spending, it is critical that Kenya continues to implement medium-
826 +L47: term fiscal consolidation plans.” In addition to pandemic-related risks to the
827 +L48: outlook, a second key domestic risk factor stems from the drought conditions
828 +L49: which are affecting parts of the country and already causing severe hardship.
829 +L50: Should the drought intensify or spread, this would weigh on the near-term
830 +L51: economic outlook. Weaker global growth, higher-than-anticipated energy prices,
831 +L52: and tighter external financing conditions are the primary external risks. As </result>
832 +</function_results>
833 +
834 +<think>Thus first clue confirmed.
835 +
836 +Second clue: Defenders of Wildlife 2021 annual report. Confirmed.
837 +
838 +Third clue: Defenders of Wildlife 2022 annual report. Confirmed.
839 +
840 +Thus the answer is Terry C. Pelster.
841 +
842 +But let's verify that the 2022 annual report indeed shows "Contributed services 10%". We saw that. Also confirm that 2021 report shows 12%. We saw that.
843 +
844 +Thus the question: "Who was the emeritus board member that year?" Refers to 2022. Answer: Terry C. Pelster.
845 +
846 +However, the question might be asking for the emeritus board member in the 2021 report? Let's read the question carefully:
847 +
848 +"In a year before 2023, between December 1 and 15 (inclusive), a press release reported that as of the 5th of that month, a country had received a total of 16,201,670 vaccine doses, with 7,583,134 administered. That same year, a company's annual report revealed that its nearly 2.2 million members and supporters contributed over $49 million, with 12% of total revenue coming from contributed services. The following year, the annual report stated that the organization had played a role in rescinding a previous administration’s plan to open most of the 22 million acres of the National Petroleum Reserve in Alaska to oil and gas drilling and that contributed services had decreased by 2% to 10%. Who was the emeritus board member that year? Name only, no titles."
849 +
850 +The phrase "that year" could refer to "the following year" i.e., the year of the annual report that mentioned the rescinded plan. Or could refer to "that same year" (the year of the press release and the first annual report). However, the phrase "that year" appears at the end: "Who was the emeritus board member that year?" It likely refers to the year of the annual report that is being discussed in the preceding sentence. The preceding sentence talks about the following year's annual report. So indeed "that year" is the following year (2022).
851 +
852 +Thus answer: Terry C. Pelster.
853 +
854 +But to be thorough, we should check the 2022 annual report for any other emeritus board members. The search only showed "BOARD EMERITUS Terry C. Pelster, Esq." Maybe there are multiple emeritus board members. But likely only one.
855 +
856 +Check the 2021 annual report for emeritus board member. It also shows Terry C. Pelster. So maybe the question is asking for the emeritus board member in the 2021 report. But why would they ask that? The question is about connecting the clues to identify the organization and then find the emeritus board member. The clues point to Defenders of Wildlife. The question likely asks for the emeritus board member as of the 2022 report.
857 +
858 +Nevertheless, we should confirm that the 2022 report is indeed the one that mentions the rescinded plan. We saw that. Also confirm that contributed services decreased from 12% to 10%. We saw that.
859 +
860 +Thus answer is Terry C. Pelster.
861 +
862 +But let's search for "Defenders of Wildlife emeritus board member 2022 Terry Pelster".</think>
863 +
864 +<｜DSML｜function_calls>
865 +<｜DSML｜invoke name="search">
866 +<｜DSML｜parameter name="query" string="true">Terry C. Pelster Defenders of Wildlife emeritus board member</｜DSML｜parameter>
867 +<｜DSML｜parameter name="topn" string="false">10</｜DSML｜parameter>
868 +<｜DSML｜parameter name="source" string="true">web</｜DSML｜parameter>
869 +</｜DSML｜invoke>
870 +</｜DSML｜function_calls><｜end▁of▁sentence｜>
871 +
872 +<function_results>
873 +<result>[15] Terry C. Pelster Defenders of Wildlife emeritus board member (Search_Results/Terry C. Pelster Defenders of Wildlife emeritus board member)
874 +**viewing lines [0 - 31] of 31**
875 +
876 +L0:
877 +L1: URL: Search_Results/Terry C. Pelster Defenders of Wildlife emeritus board member
878 +L2: # Search Results
879 +L3:
880 +L4: * 【0†2024 impact report; publish_date: none†defenders.org】 Terry Root, Ph.D.
881 +L5: New Mexico. Thomas D. Sisk, Ph.D. British Columbia, Canada. BOARD EMERITUS.
882 +L6: Terry C. Pelster, Esq. New York. SCIENCE ADVISORS. Leah Gerber, Ph ...
883 +L7: * 【1†Defenders of Wildlife; publish_date: none†www.wrongkindofgreen.org】
884 +L8: Directors ; Victor M. Sher*, Chair – California ; Terry C. Pelster*, Vice Chair
885 +L9: – New York ; Richard Kopcho*, Treasurer – California ; Adelaide P. Gomer*,
886 +L10: Secretary ...
887 +L11: * 【2†Exhibit 12; publish_date: none†www.uschamber.com】 DECLARATION OF TERRY
888 +L12: PELSTER. I, Terry C. Pelster, declare as follows: 1. I am a current member of
889 +L13: Defenders of Wildlife (“Defenders”) and have been a member.
890 +L14: * 【3†2020 ANNUAL REPORT; publish_date: none†defenders.org】 Terry L. Root.
891 +L15: Florida. BOARD. EMERITUS. Terry C. Pelster. New York. Alan Steinberg. Florida.
892 +L16: SCIENCE. ADVISORS. Gabriela Chavarria,. Ph.D. Denver Museum of.
893 +L17: * 【4†2019 Annual Report; publish_date: none†defenders.org】 BOARD EMERITUS.
894 +L18: Terry C. Pelster. New York. Alan Steinberg. Florida. SCIENCE ADVISORS. Leah
895 +L19: Gerber, Ph.D. Arizona State University. Kathayoon Khalil, Ph.D.
896 +L20: * 【5†annual report; publish_date: none†defenders.org】 * Executive Committee
897 +L21: member. BoARD EMERITuS. Terry C. Pelster. New York. Alan Steinberg. Florida.
898 +L22: PRESIDENT. EMERITuS. Rodger Schlickeisen. Virginia. SCIENCE ...
899 +L23: * 【6†BEING THERE FOR BELUGAS; publish_date: none†wendeenicole.com】 Mar 2, 2023
900 +L24: — Defenders of Wildlife is a national, nonprofit member- ship ... Terry C.
901 +L25: Pelster, Esq. New York. *Executive Committee member outfront.
902 +L26: * 【7†Annual Reports; publish_date: none†www.annualreports.com】 WAYS TO GIVE
903 +L27: There are many ways to support Defenders of Wildlife and protect our imperiled
904 +L28: wildlife. ... Terry C. Pelster. New York. Alan Steinberg. Florida.
905 +L29: * 【8†2014 Annual Report; publish_date: none†defenders.org】 BOARD OF DIRECTORS
906 +L30: defenders of wildlife is a national, nonprofit ... BOARD EMERITUS. Terry C.
907 +L31: Pelster. New York. Alan Steinberg. Florida. Page 3. PO.</result>
908 +</function_results>
909 +
910 +<think>Thus Terry C. Pelster is indeed the emeritus board member.
911 +
912 +Now the question might want just "Terry C. Pelster". Or maybe "Terry Pelster". Or "Terry C. Pelster, Esq." but instructions say "Name only, no titles." So likely "Terry C. Pelster". However, maybe they want "Terry Pelster". Better to include full name as listed: "Terry C. Pelster".
913 +
914 +Let's read the question again: "Who was the emeritus board member that year? Name only, no titles." So we should answer just the name, no "Esq." or "Board Emeritus". So "Terry C. Pelster". Or maybe "Terry Pelster". But the document shows "Terry C. Pelster, Esq." So the name is "Terry C. Pelster".
915 +
916 +Better to verify that the 2022 annual report indeed lists "BOARD EMERITUS Terry C. Pelster, Esq." We saw that line.
917 +
918 +Thus answer: Terry C. Pelster.
919 +
920 +But let's double-check that the question might refer to the 2021 annual report. The phrase "that year" could be ambiguous. Let's parse the timeline:
921 +
922 +- Year 1: 2021. Press release in Dec 2021 about Kenya vaccine doses. Same year (2021) annual report of Defenders of Wildlife shows 12% contributed services.
923 +
924 +- Year 2: 2022. Annual report mentions rescinded plan and contributed services decreased to 10%.
925 +
926 +The question says "The following year, the annual report stated that the organization had played a role in rescinding a previous administration’s plan ... and that contributed services had decreased by 2% to 10%." This is year 2.
927 +
928 +Then asks "Who was the emeritus board member that year?" That likely refers to year 2.
929 +
930 +Thus answer: Terry C. Pelster.
931 +
932 +However, perhaps the emeritus board member changed between 2021 and 2022. We need to check 2022 annual report for any other emeritus board member. We saw only Terry C. Pelster. Also check 2021 annual report: same. So likely unchanged.
933 +
934 +Thus answer is Terry C. Pelster.
935 +
936 +Nevertheless, we should confirm that the 2022 annual report is indeed the one that includes the phrase "BLM rescinded the previous administration’s plan that opened most of the 22 million acres of the National Petroleum Reserve in Alaska to oil and gas drilling". We saw that in line 200-202.
937 +
938 +Thus answer is Terry C. Pelster.
939 +
940 +Now we need to provide the answer in the final message. Only the name.
941 +
942 +Thus final answer: Terry C. Pelster.
943 +
944 +But let's be thorough and search for "Defenders of Wildlife 2022 annual report contributed services 10%".</think>
945 +
946 +<｜DSML｜function_calls>
947 +<｜DSML｜invoke name="find">
948 +<｜DSML｜parameter name="cursor" string="false">8</｜DSML｜parameter>
949 +<｜DSML｜parameter name="pattern" string="true">Contributed services</｜DSML｜parameter>
950 +</｜DSML｜invoke>
951 +</｜DSML｜function_calls><｜end▁of▁sentence｜>
952 +
953 +<function_results>
954 +<result>[16] Find results for text: `contributed services` in `www.annualreports.com` (https://www.annualreports.com/HostedData/AnnualReportArchive/d/defenders-wildlife_2022.pdf/find?pattern=contributed services)
955 +**viewing lines [0 - 10] of 10**
956 +
957 +L0: # 【0†match at L222】
958 +L1: income $744 Contributed services $5,047 Total Revenue $41,170 USES OF FUNDS
959 +L2: Biodiversity conservation $18,970 Constituency mobilization $16,528 Fundraising
960 +L3: $1,029 Management and general $6,354 Total Expenses $42,881 Change in net assets
961 +L4: (-$1,711) Net assets, start of year $45,144 Net Assets, End of the Year $43,433
962 +L5:
963 +L6: # 【1†match at L227】
964 +L7: Contributed services 10% Royalties and other income 2% Bequests, trusts and
965 +L8: split interests 29% TOTAL REVENUE Income from investments -15% Program and
966 +L9: support services 98% USE OF FUNDS Fundraising 2% ===== Page 13 =====
967 +L10: DEFENDERS.ORG 13 ===== Page 14 ===== BOARD OF DIRECTORS John Dayton, Esq.* Chair</result>
968 +</function_results>
969 +
970 +<think>Thus confirmed.
971 +
972 +Now we should verify that the 2021 annual report shows contributed services 12%.</think>
973 +
974 +<｜DSML｜function_calls>
975 +<｜DSML｜invoke name="find">
976 +<｜DSML｜parameter name="cursor" string="false">4</｜DSML｜parameter>
977 +<｜DSML｜parameter name="pattern" string="true">Contributed services</｜DSML｜parameter>
978 +</｜DSML｜invoke>
979 +</｜DSML｜function_calls><｜end▁of▁sentence｜>
980 +
981 +<function_results>
982 +<result>[17] Find results for text: `contributed services` in `www.annualreports.com` (https://www.annualreports.com/HostedData/AnnualReportArchive/d/defenders-wildlife_2021.pdf/find?pattern=contributed services)
983 +**viewing lines [0 - 4] of 4**
984 +
985 +L0: # 【0†match at L194】
986 +L1: $3,354 Royalties and other income $3,576 Contributed services $6,140 Total
987 +L2: Revenue $49,819 USES OF FUNDS Biodiversity conservation $22,420 Constituency
988 +L3: mobilization $16,324 Fundraising $1,211 Management and general $5,865 Total
989 +L4: Expenses $45,820 Change in net assets $3,999 Net assets, start of year $41,145 </result>
990 +</function_results>
991 +
992 +<think>Open around L190-200 to see percentage.</think>
993 +
994 +<｜DSML｜function_calls>
995 +<｜DSML｜invoke name="open">
996 +<｜DSML｜parameter name="cursor" string="false">4</｜DSML｜parameter>
997 +<｜DSML｜parameter name="loc" string="false">185</｜DSML｜parameter>
998 +</｜DSML｜invoke>
999 +</｜DSML｜function_calls><｜end▁of▁sentence｜>
1000 +
1001 +<function_results>
1002 +<result>[18] www.annualreports.com (https://www.annualreports.com/HostedData/AnnualReportArchive/d/defenders-wildlife_2021.pdf)
1003 +**viewing lines [185 - 232] of 261**
1004 +
1005 +L185: devastating water pollution that is killing the seagrass and causing manatees to
1006 +L186: starve. DAVID TES | SAM FRENZY DRAW DEFENDERS.ORG 11 ===== Page 12 ===== In
1007 +L187: 2021, Defenders of Wildlife’s nearly 2.2 million members and supporters
1008 +L188: contributed more than $49 million for wildlife and wildlife habitat. This figure
1009 +L189: includes all those who generously donated their time and expertise. The
1010 +L190: steadfast support of our donors allows Defenders to sustain our program and
1011 +L191: public education efforts in the field, the courts and on Capitol Hill. 2021
1012 +L192: SOURCES OF FUNDS Grants and contributions $29,057 Bequests, trusts and split
1013 +L193: interests $7,692 Income from investments, annuity reserve funds and trusts
1014 +L194: $3,354 Royalties and other income $3,576 Contributed services $6,140 Total
1015 +L195: Revenue $49,819 USES OF FUNDS Biodiversity conservation $22,420 Constituency
1016 +L196: mobilization $16,324 Fundraising $1,211 Management and general $5,865 Total
1017 +L197: Expenses $45,820 Change in net assets $3,999 Net assets, start of year $41,145
1018 +L198: Net Assets, End of the Year $45,144 Dollars are in thousands. 12 DEFENDERS OF
1019 +L199: WILDLIFE Grants and contributions 58% Income from investments 7% Requests,
1020 +L200: trusts and split interests 15% Royalties and other income 7% Contributed
1021 +L201: services 12% Program and support services 97% Fundraising 3% ===== Page 13 =====
1022 +L202: DEFENDERS.ORG 13 ===== Page 14 ===== John Dayton, Esq.* Chair Texas Sheryl
1023 +L203: Sachs* Vice Chair Connecticut Whitfield Bailey* Secretary Tennessee Mark Caylor*
1024 +L204: Treasurer District of Columbia Sylvia Baca* California Donald Barry, Esq.
1025 +L205: Oregon Dinah Bear, Esq. Arizona Kim O’Keefe Beck, Ph.D California Gabriela
1026 +L206: Chavarria, Ph.D Washington Jamie Rappaport Clark* President & CEO Virginia
1027 +L207: Caroline D. Gabel Maryland Keith M. Harper, Esq. Maryland John Donovan, Esq.
1028 +L208: Montana and California James M. Hecker, Esq. Virginia Judith Ponsikoff, Ph.D*
1029 +L209: California Terry L. Root, Ph.D Florida BOARD EMERITUS Terry C. Pelster, Esq. New
1030 +L210: York *Executive Committee member* ST A F F & A D V I S O R S SCIENCE ADVISORS
1031 +L211: Gabriela Chavarria, Ph.D Burke Museum, University of Washington Leah Gerber,
1032 +L212: Ph.D Arizona State University Kathayoon Khalil, Ph.D Oregon Zoo Rurik List, Ph.D
1033 +L213: Universidad Autónoma Metropolitana-Lerma Barry Noon, Ph.D Colorado State
1034 +L214: University Terry L. Root, Ph.D Stanford University Mark Schwartz, Ph.D
1035 +L215: University of California-Davis David S. Wilcox, Ph.D Princeton University Rae
1036 +L216: Wynn-Grant, Ph.D National Geographic Society SENIOR STAFF Jamie Rappaport Clark
1037 +L217: President & CEO Thu Pham Chief of Staff Renee Stone Senior Vice President,
1038 +L218: Conservation Programs & General Counsel James Stefan Senior Vice President,
1039 +L219: Operations Matt Braughler Vice President, Philanthropy Rachel Brittin Vice
1040 +L220: President, Communications Shawn Cantrell Vice President, Field Conservation
1041 +L221: Programs Robert L. Dewey Vice President, Government Relations & External Affairs
1042 +L222: Monica Goldberg Vice President, Landscape Conservation Sharon Morris-Smith Vice
1043 +L223: President of Human Resources Elizabeth Nichols Vice President, Finance/CFO
1044 +L224: Anthony Pegues Vice President, Diversity, Equity, Inclusion, Justice Lindsay
1045 +L225: Rosa, Ph.D Director, Center for Conservation Innovation Michael P. Senatore Vice
1046 +L226: President, Conservation Law Brandace Swimley Vice President, Integrated
1047 +L227: Marketing 14 DEFENDERS OF WILDLIFE ===== Page 15 ===== THERE ARE MANY WAYS TO
1048 +L228: SUPPORT DEFENDERS OF WILDLIFE and protect our imperiled wildlife. To learn more
1049 +L229: or donate, contact us at 1-800-385-9712, visit defenders.org/donate or choose
1050 +L230: one of the following options. ➤ Transfer appreciated assets directly to
1051 +L231: Defenders and maximize potential tax benefits. Defenders.org/donate-stocks-help-
1052 +L232: save-wildlife ➤ Become a Leadership Council member with a single donation of </result>
1053 +</function_results>
1054 +
1055 +<think>Thus everything matches.
1056 +
1057 +Thus answer: Terry C. Pelster.
1058 +
1059 +Now we need to output only the name.
1060 +
1061 +Thus final answer: Terry C. Pelster.
1062 +
1063 +But perhaps the question expects "Terry Pelster". However, the document includes middle initial "C." It might be better to include the full name as listed: "Terry C. Pelster". The instruction "Name only, no titles" suggests we should not include "Esq." or "Board Emeritus". But including middle initial is likely okay.
1064 +
1065 +To be safe, we could output "Terry C. Pelster". Let's see similar puzzles: often they expect the full name as appears in the document.
1066 +
1067 +Thus answer: Terry C. Pelster.
1068 +
1069 +Now we will provide final answer.</think>Terry C. Pelster<｜end▁of▁sentence｜>
\ No newline at end of file
rust/src/chat/src/renderer/deepseek_v32/fixtures/test_output_vllm_parity.txt new
+112
@@ -0,0 +1,112 @@
1 +<｜begin▁of▁sentence｜>
2 +
3 +## Tools
4 +
5 +You have access to a set of tools you can use to answer the user's question.
6 +You can invoke functions by writing a "<｜DSML｜function_calls>" block like the following as part of your reply to the user:
7 +<｜DSML｜function_calls>
8 +<｜DSML｜invoke name="$FUNCTION_NAME">
9 +<｜DSML｜parameter name="$PARAMETER_NAME" string="true|false">$PARAMETER_VALUE</｜DSML｜parameter>
10 +...
11 +</｜DSML｜invoke>
12 +<｜DSML｜invoke name="$FUNCTION_NAME2">
13 +...
14 +</｜DSML｜invoke>
15 +</｜DSML｜function_calls>
16 +
17 +String and scalar parameters should be specified as is without any escaping or quotes, while lists and objects should use JSON format. The "string" attribute should be set to "true" for string type parameters and "false" for other types (numbers, booleans, arrays, objects).
18 +
19 +If the thinking_mode is enabled, then after function results you should strongly consider outputting a thinking block. Here is an example:
20 +
21 +<｜DSML｜function_calls>
22 +...
23 +</｜DSML｜function_calls>
24 +
25 +<function_results>
26 +...
27 +</function_results>
28 +
29 +<think>...thinking about results</think>
30 +
31 +Here are the functions available in JSONSchema format:
32 +<functions>
33 +{"name": "get_datetime", "description": "Get the current date and time", "parameters": {"type": "object", "properties": {"timezone": {"type": "string", "description": "The timezone, e.g. Asia/Shanghai, UTC"}}, "required": ["timezone"]}}
34 +{"name": "get_weather", "description": "Get the weather for a specific date and location", "parameters": {"type": "object", "properties": {"location": {"type": "string", "description": "The city name, e.g. Beijing, Hangzhou"}, "date": {"type": "string", "description": "The date in YYYY-MM-DD format"}}, "required": ["location", "date"]}}
35 +</functions>
36 +You are a helpful Assistant.<｜User｜>明天杭州和北京的天气怎么样？<｜Assistant｜></think>
37 +
38 +<｜DSML｜function_calls>
39 +<｜DSML｜invoke name="get_datetime">
40 +<｜DSML｜parameter name="timezone" string="true">Asia/Shanghai</｜DSML｜parameter>
41 +</｜DSML｜invoke>
42 +</｜DSML｜function_calls><｜end▁of▁sentence｜>
43 +
44 +<function_results>
45 +<result>{"current_date": "2024-01-15", "current_time": "14:30:00", "timezone": "Asia/Shanghai"}</result>
46 +</function_results>
47 +
48 +</think>
49 +
50 +<｜DSML｜function_calls>
51 +<｜DSML｜invoke name="get_weather">
52 +<｜DSML｜parameter name="location" string="true">Hangzhou</｜DSML｜parameter>
53 +<｜DSML｜parameter name="date" string="true">2024-01-16</｜DSML｜parameter>
54 +</｜DSML｜invoke>
55 +<｜DSML｜invoke name="get_weather">
56 +<｜DSML｜parameter name="location" string="true">Beijing</｜DSML｜parameter>
57 +<｜DSML｜parameter name="date" string="true">2024-01-16</｜DSML｜parameter>
58 +</｜DSML｜invoke>
59 +</｜DSML｜function_calls><｜end▁of▁sentence｜>
60 +
61 +<function_results>
62 +<result>{"location": "Hangzhou", "date": "2024-01-16", "temperature_high": "12", "temperature_low": "5", "weather": "多云", "humidity": "65%"}</result>
63 +<result>{"location": "Beijing", "date": "2024-01-16", "temperature_high": "-2", "temperature_low": "-8", "weather": "晴", "humidity": "30%"}</result>
64 +</function_results>
65 +
66 +</think>根据查询结果，明天（2024年1月16日）的天气情况如下：
67 +
68 +**杭州**：
69 +- 天气：多云
70 +- 最高温度：12°C
71 +- 最低温度：5°C
72 +- 湿度：65%
73 +
74 +**北京**：
75 +- 天气：晴
76 +- 最高温度：-2°C
77 +- 最低温度：-8°C
78 +- 湿度：30%
79 +
80 +杭州明天会比较温暖但有些多云，而北京会很冷但是晴天。建议在北京的朋友要注意保暖！<｜end▁of▁sentence｜><｜User｜>谢谢！那后天呢？<｜Assistant｜><think>用户现在问后天的天气，后天是2024-01-17，我可以直接查询（因为已知今天日期）🗓️</think>
81 +
82 +<｜DSML｜function_calls>
83 +<｜DSML｜invoke name="get_weather">
84 +<｜DSML｜parameter name="location" string="true">Hangzhou</｜DSML｜parameter>
85 +<｜DSML｜parameter name="date" string="true">2024-01-17</｜DSML｜parameter>
86 +</｜DSML｜invoke>
87 +<｜DSML｜invoke name="get_weather">
88 +<｜DSML｜parameter name="location" string="true">Beijing</｜DSML｜parameter>
89 +<｜DSML｜parameter name="date" string="true">2024-01-17</｜DSML｜parameter>
90 +</｜DSML｜invoke>
91 +</｜DSML｜function_calls><｜end▁of▁sentence｜>
92 +
93 +<function_results>
94 +<result>{"location": "Hangzhou", "date": "2024-01-17", "temperature_high": "15", "temperature_low": "8", "weather": "小雨", "humidity": "80%"}</result>
95 +<result>{"location": "Beijing", "date": "2024-01-17", "temperature_high": "0", "temperature_low": "-6", "weather": "多云", "humidity": "45%"}</result>
96 +</function_results>
97 +
98 +<think>获取到后天的天气数据，整理回复给用户📝</think>后天（2024年1月17日）的天气情况：
99 +
100 +**杭州**：
101 +- 天气：小雨
102 +- 最高温度：15°C
103 +- 最低温度：8°C
104 +- 湿度：80%
105 +
106 +**北京**：
107 +- 天气：多云
108 +- 最高温度：0°C
109 +- 最低温度：-6°C
110 +- 湿度：45%
111 +
112 +杭州后天会有小雨，温度略有回升，记得带伞。北京会稍微暖和一点，但依然很冷，请继续做好保暖措施。<｜end▁of▁sentence｜>
\ No newline at end of file
rust/src/chat/src/renderer/deepseek_v32/mod.rs new
+31
@@ -0,0 +1,31 @@
1 +mod encoding;
2 +
3 +use vllm_text::Prompt;
4 +
5 +use super::{ChatRenderer, RenderedPrompt};
6 +use crate::Result;
7 +use crate::request::ChatRequest;
8 +
9 +/// Dedicated DeepSeek V3.2 renderer.
10 +#[derive(Debug, Clone, Copy, Default)]
11 +pub struct DeepSeekV32ChatRenderer;
12 +
13 +impl DeepSeekV32ChatRenderer {
14 + /// Create the dedicated DeepSeek V3.2 renderer.
15 + pub fn new() -> Self {
16 + Self
17 + }
18 +}
19 +
20 +impl ChatRenderer for DeepSeekV32ChatRenderer {
21 + fn render(&self, request: &ChatRequest) -> Result<RenderedPrompt> {
22 + request.validate()?;
23 +
24 + Ok(RenderedPrompt {
25 + prompt: Prompt::Text(encoding::render_request(request)?),
26 + })
27 + }
28 +}
29 +
30 +#[cfg(test)]
31 +mod tests;
rust/src/chat/src/renderer/deepseek_v32/tests.rs new
+422
@@ -0,0 +1,422 @@
1 +use std::fs;
2 +use std::path::PathBuf;
3 +
4 +use expect_test::{ExpectFile, expect, expect_file};
5 +use serde::Deserialize;
6 +use serde_json::{Value, json};
7 +use thiserror_ext::AsReport;
8 +
9 +use super::DeepSeekV32ChatRenderer;
10 +use crate::error::Error;
11 +use crate::event::{AssistantContentBlock, AssistantToolCall};
12 +use crate::request::{
13 + ChatContentPart, ChatMessage, ChatRequest, ChatTool, ChatToolChoice, GenerationPromptMode,
14 +};
15 +use crate::{ChatRenderer, ChatRole};
16 +
17 +#[derive(Debug, Deserialize)]
18 +struct FixtureRequest {
19 + #[serde(default)]
20 + tools: Vec<FixtureTool>,
21 + messages: Vec<FixtureMessage>,
22 +}
23 +
24 +#[derive(Debug, Deserialize)]
25 +struct FixtureTool {
26 + function: FixtureToolFunction,
27 +}
28 +
29 +#[derive(Debug, Deserialize)]
30 +struct FixtureToolFunction {
31 + name: String,
32 + description: Option<String>,
33 + parameters: Value,
34 + #[serde(default)]
35 + strict: Option<bool>,
36 +}
37 +
38 +#[derive(Debug, Deserialize)]
39 +#[serde(tag = "role", rename_all = "snake_case")]
40 +enum FixtureMessage {
41 + System {
42 + content: String,
43 + },
44 + Developer {
45 + content: String,
46 + #[serde(default)]
47 + tools: Vec<FixtureTool>,
48 + },
49 + User {
50 + content: String,
51 + },
52 + Assistant {
53 + #[serde(default)]
54 + content: String,
55 + #[serde(default)]
56 + reasoning_content: String,
57 + #[serde(default)]
58 + tool_calls: Vec<FixtureToolCall>,
59 + },
60 + Tool {
61 + content: String,
62 + #[serde(default)]
63 + tool_call_id: Option<String>,
64 + },
65 +}
66 +
67 +#[derive(Debug, Deserialize)]
68 +struct FixtureToolCall {
69 + #[serde(default)]
70 + id: Option<String>,
71 + function: FixtureToolCallFunction,
72 +}
73 +
74 +#[derive(Debug, Deserialize)]
75 +struct FixtureToolCallFunction {
76 + name: String,
77 + arguments: String,
78 +}
79 +
80 +fn render_request(request: &ChatRequest) -> String {
81 + DeepSeekV32ChatRenderer::new()
82 + .render(request)
83 + .unwrap()
84 + .prompt
85 + .into_text()
86 + .expect("deepseek renderer should return text prompt")
87 +}
88 +
89 +fn render_result(request: &ChatRequest) -> Result<String, Error> {
90 + DeepSeekV32ChatRenderer::new().render(request).map(|rendered| {
91 + rendered
92 + .prompt
93 + .into_text()
94 + .expect("deepseek renderer should return text prompt")
95 + })
96 +}
97 +
98 +fn thinking_request(messages: Vec<ChatMessage>) -> ChatRequest {
99 + let mut request = ChatRequest {
100 + request_id: "deepseek-v32-small-test".to_string(),
101 + messages,
102 + ..ChatRequest::for_test()
103 + };
104 + if matches!(
105 + request.messages.last().map(ChatMessage::role),
106 + Some(ChatRole::Assistant)
107 + ) {
108 + request.chat_options.generation_prompt_mode = GenerationPromptMode::NoGenerationPrompt;
109 + }
110 + request
111 + .chat_options
112 + .template_kwargs
113 + .insert("thinking".to_string(), Value::Bool(true));
114 + request
115 +}
116 +
117 +fn fixture_request(input_name: &str) -> ChatRequest {
118 + let fixture = fs::read_to_string(fixture_path(input_name)).unwrap();
119 + let fixture: FixtureRequest = serde_json::from_str(&fixture).unwrap();
120 + let mut request = ChatRequest {
121 + request_id: "deepseek-v32-fixture".to_string(),
122 + messages: fixture
123 + .messages
124 + .into_iter()
125 + .enumerate()
126 + .map(|(index, message)| match message {
127 + FixtureMessage::System { content } => ChatMessage::system(content),
128 + FixtureMessage::Developer { content, tools } => ChatMessage::developer(
129 + content,
130 + (!tools.is_empty()).then(|| to_chat_tools(&tools)),
131 + ),
132 + FixtureMessage::User { content } => ChatMessage::user(content),
133 + FixtureMessage::Assistant {
134 + content,
135 + reasoning_content,
136 + tool_calls,
137 + } => {
138 + let mut blocks = Vec::new();
139 + if !reasoning_content.is_empty() {
140 + blocks.push(AssistantContentBlock::Reasoning {
141 + text: reasoning_content,
142 + });
143 + }
144 + if !content.is_empty() {
145 + blocks.push(AssistantContentBlock::Text { text: content });
146 + }
147 + blocks.extend(tool_calls.into_iter().enumerate().map(
148 + |(tool_index, tool_call)| {
149 + AssistantContentBlock::ToolCall(AssistantToolCall {
150 + id: tool_call.id.unwrap_or_else(|| {
151 + format!("fixture-tool-call-{index}-{tool_index}")
152 + }),
153 + name: tool_call.function.name,
154 + arguments: tool_call.function.arguments,
155 + })
156 + },
157 + ));
158 + ChatMessage::assistant_blocks(blocks)
159 + }
160 + FixtureMessage::Tool {
161 + content,
162 + tool_call_id,
163 + } => ChatMessage::tool_response(
164 + content,
165 + tool_call_id.unwrap_or_else(|| format!("fixture-tool-response-{index}")),
166 + ),
167 + })
168 + .collect(),
169 + tools: to_chat_tools(&fixture.tools),
170 + tool_choice: if fixture.tools.is_empty() {
171 + ChatToolChoice::None
172 + } else {
173 + ChatToolChoice::Auto
174 + },
175 + ..ChatRequest::for_test()
176 + };
177 + if matches!(
178 + request.messages.last().map(ChatMessage::role),
179 + Some(ChatRole::Assistant)
180 + ) {
181 + request.chat_options.generation_prompt_mode = GenerationPromptMode::NoGenerationPrompt;
182 + }
183 + request
184 + .chat_options
185 + .template_kwargs
186 + .insert("thinking".to_string(), Value::Bool(true));
187 + request
188 +}
189 +
190 +fn to_chat_tools(tools: &[FixtureTool]) -> Vec<ChatTool> {
191 + tools
192 + .iter()
193 + .map(|tool| ChatTool {
194 + name: tool.function.name.clone(),
195 + description: tool.function.description.clone(),
196 + parameters: tool.function.parameters.clone(),
197 + strict: tool.function.strict,
198 + })
199 + .collect()
200 +}
201 +
202 +fn fixture_path(name: &str) -> PathBuf {
203 + PathBuf::from(env!("CARGO_MANIFEST_DIR"))
204 + .join("src/renderer/deepseek_v32")
205 + .join("fixtures")
206 + .join(name)
207 +}
208 +
209 +fn assert_fixture(input_name: &str, expected: ExpectFile) {
210 + let request = fixture_request(input_name);
211 + let rendered = render_request(&request);
212 + expected.assert_eq(&rendered);
213 +}
214 +
215 +#[test]
216 +fn renders_vllm_parity_prompt_for_request_level_tools_fixture() {
217 + assert_fixture(
218 + "test_input.json",
219 + expect_file!["fixtures/test_output_vllm_parity.txt"],
220 + );
221 +}
222 +
223 +#[test]
224 +fn renders_official_search_fixture_without_date() {
225 + assert_fixture(
226 + "test_input_search_wo_date.json",
227 + expect_file!["fixtures/test_output_search_wo_date.txt"],
228 + );
229 +}
230 +
231 +#[test]
232 +fn renders_official_search_fixture_with_date() {
233 + assert_fixture(
234 + "test_input_search_w_date.json",
235 + expect_file!["fixtures/test_output_search_w_date.txt"],
236 + );
237 +}
238 +
239 +#[test]
240 +fn request_level_tools_are_lowered_as_synthetic_leading_system_message() {
241 + let mut request = ChatRequest {
242 + request_id: "deepseek-v32-tools".to_string(),
243 + messages: vec![
244 + ChatMessage::system("System prompt."),
245 + ChatMessage::text(ChatRole::User, "Hello"),
246 + ],
247 + tools: vec![ChatTool {
248 + name: "lookup".to_string(),
249 + description: Some("Look things up".to_string()),
250 + parameters: json!({
251 + "type": "object",
252 + "properties": {
253 + "query": {
254 + "type": "string"
255 + }
256 + },
257 + "required": ["query"]
258 + }),
259 + strict: None,
260 + }],
261 + tool_choice: ChatToolChoice::Auto,
262 + ..ChatRequest::for_test()
263 + };
264 + request
265 + .chat_options
266 + .template_kwargs
267 + .insert("thinking".to_string(), Value::Bool(true));
268 +
269 + let rendered = render_request(&request);
270 +
271 + assert!(rendered.starts_with("<｜begin▁of▁sentence｜>\n\n## Tools\n"));
272 + assert!(rendered.contains("</functions>\nSystem prompt."));
273 + assert!(rendered.ends_with("<｜User｜>Hello<｜Assistant｜><think>"));
274 +}
275 +
276 +#[test]
277 +fn developer_turn_is_treated_as_last_user_like_turn() {
278 + let request = thinking_request(vec![ChatMessage::developer("Follow policy.", None)]);
279 +
280 + let rendered = render_request(&request);
281 +
282 + assert!(rendered.contains("# The user's message is: Follow policy."));
283 + assert!(rendered.ends_with("<｜Assistant｜><think>"));
284 +}
285 +
286 +#[test]
287 +fn historical_assistant_reasoning_is_dropped_before_final_user_turn() {
288 + let request = thinking_request(vec![
289 + ChatMessage::assistant_blocks(vec![
290 + AssistantContentBlock::Reasoning {
291 + text: "internal reasoning".to_string(),
292 + },
293 + AssistantContentBlock::Text {
294 + text: "Visible answer.".to_string(),
295 + },
296 + ]),
297 + ChatMessage::user("What about the next one?"),
298 + ]);
299 +
300 + let rendered = render_request(&request);
301 +
302 + assert!(!rendered.contains("internal reasoning"));
303 + assert!(rendered.contains("Visible answer.<｜end▁of▁sentence｜>"));
304 + assert!(rendered.ends_with("<｜User｜>What about the next one?<｜Assistant｜><think>"));
305 +}
306 +
307 +#[test]
308 +fn historical_assistant_reasoning_is_dropped_before_final_developer_turn() {
309 + let request = thinking_request(vec![
310 + ChatMessage::assistant_blocks(vec![
311 + AssistantContentBlock::Reasoning {
312 + text: "internal reasoning".to_string(),
313 + },
314 + AssistantContentBlock::Text {
315 + text: "Visible answer.".to_string(),
316 + },
317 + ]),
318 + ChatMessage::developer("Follow the rubric.", None),
319 + ]);
320 +
321 + let rendered = render_request(&request);
322 +
323 + assert!(!rendered.contains("internal reasoning"));
324 + assert!(rendered.contains("Visible answer.<｜end▁of▁sentence｜>"));
325 + assert!(rendered.ends_with(
326 + "<｜User｜>\n\n# The user's message is: Follow the rubric.<｜Assistant｜><think>"
327 + ));
328 +}
329 +
330 +#[test]
331 +fn tool_results_after_last_user_resume_thinking() {
332 + let request = thinking_request(vec![
333 + ChatMessage::user("Check the weather."),
334 + ChatMessage::assistant_blocks(vec![AssistantContentBlock::ToolCall(AssistantToolCall {
335 + id: "call-weather".to_string(),
336 + name: "weather".to_string(),
337 + arguments: "{\"city\":\"Hangzhou\"}".to_string(),
338 + })]),
339 + ChatMessage::tool_response("{\"ok\":true}", "call-weather"),
340 + ]);
341 +
342 + let rendered = render_request(&request);
343 +
344 + assert!(rendered.contains(
345 + "<｜User｜>Check the weather.<｜Assistant｜><think></think>\n\n<｜DSML｜function_calls>"
346 + ));
347 + assert!(rendered.ends_with("</function_results>\n\n<think>"));
348 +}
349 +
350 +#[test]
351 +fn tool_results_follow_assistant_tool_call_id_order() {
352 + let request = thinking_request(vec![
353 + ChatMessage::user("Check two cities."),
354 + ChatMessage::assistant_blocks(vec![
355 + AssistantContentBlock::ToolCall(AssistantToolCall {
356 + id: "call-hangzhou".to_string(),
357 + name: "weather".to_string(),
358 + arguments: "{\"city\":\"Hangzhou\"}".to_string(),
359 + }),
360 + AssistantContentBlock::ToolCall(AssistantToolCall {
361 + id: "call-beijing".to_string(),
362 + name: "weather".to_string(),
363 + arguments: "{\"city\":\"Beijing\"}".to_string(),
364 + }),
365 + ]),
366 + ChatMessage::tool_response("{\"city\":\"Beijing\"}", "call-beijing"),
367 + ChatMessage::tool_response("{\"city\":\"Hangzhou\"}", "call-hangzhou"),
368 + ]);
369 +
370 + let rendered = render_request(&request);
371 +
372 + assert!(rendered.contains(
373 + "<function_results>\n<result>{\"city\":\"Hangzhou\"}</result>\n<result>{\"city\":\"Beijing\"}</result>\n</function_results>"
374 + ));
375 +}
376 +
377 +#[test]
378 +fn tool_results_require_matching_tool_call_ids() {
379 + let request = thinking_request(vec![
380 + ChatMessage::user("Check the weather."),
381 + ChatMessage::assistant_blocks(vec![AssistantContentBlock::ToolCall(AssistantToolCall {
382 + id: "call-weather".to_string(),
383 + name: "weather".to_string(),
384 + arguments: "{\"city\":\"Hangzhou\"}".to_string(),
385 + })]),
386 + ChatMessage::tool_response("{\"ok\":true}", "call-unknown"),
387 + ]);
388 +
389 + let error = render_result(&request).unwrap_err();
390 +
391 + expect!["chat template error: invalid DeepSeek V3.2 tool message: unknown tool_call_id `call-unknown`"]
392 + .assert_eq(&error.to_report_string());
393 +}
394 +
395 +#[test]
396 +fn assistant_after_last_user_requires_reasoning_or_tool_calls() {
397 + let request = thinking_request(vec![
398 + ChatMessage::user("Hello"),
399 + ChatMessage::assistant_text("Hi there."),
400 + ]);
401 +
402 + let error = render_result(&request).unwrap_err();
403 +
404 + expect!["chat template error: invalid DeepSeek V3.2 assistant message after last user message: expected reasoning or tool calls"]
405 + .assert_eq(&error.to_report_string());
406 +}
407 +#[test]
408 +fn render_rejects_multimodal_input() {
409 + let request = ChatRequest {
410 + messages: vec![ChatMessage::user(vec![ChatContentPart::image_url(
411 + "data:image/png;base64,test",
412 + )])],
413 + ..ChatRequest::for_test()
414 + };
415 +
416 + let error = DeepSeekV32ChatRenderer::new().render(&request).unwrap_err();
417 +
418 + assert!(matches!(
419 + error,
420 + Error::UnsupportedMultimodalContent("image_url")
421 + ));
422 +}
rust/src/chat/src/renderer/deepseek_v4/encoding.rs new
+558
@@ -0,0 +1,558 @@
1 +//! DeepSeek V4 prompt renderer.
2 +//!
3 +//! Original Python implementation:
4 +//! <https://github.com/vllm-project/vllm/blob/main/vllm/tokenizers/deepseek_v4_encoding.py>
5 +
6 +use std::collections::HashMap;
7 +use std::fmt::Write as _;
8 +
9 +use serde::Serialize;
10 +use serde_json::Value;
11 +use serde_json_fmt::JsonFormat;
12 +
13 +use crate::error::{Error, Result};
14 +use crate::request::{ChatContent, ChatMessage, ChatRequest, ChatTool, ReasoningEffort};
15 +use crate::{AssistantContentBlock, AssistantMessageExt, AssistantToolCall};
16 +
17 +const BOS_TOKEN: &str = "<｜begin▁of▁sentence｜>";
18 +const EOS_TOKEN: &str = "<｜end▁of▁sentence｜>";
19 +const THINKING_START_TOKEN: &str = "<think>";
20 +const THINKING_END_TOKEN: &str = "</think>";
21 +const DSML_TOKEN: &str = "｜DSML｜";
22 +const USER_SP_TOKEN: &str = "<｜User｜>";
23 +const ASSISTANT_SP_TOKEN: &str = "<｜Assistant｜>";
24 +const REASONING_EFFORT_MAX: &str = concat!(
25 + "Reasoning Effort: Absolute maximum with no shortcuts permitted.\n",
26 + "You MUST be very thorough in your thinking and comprehensively decompose the problem to resolve the root cause, rigorously stress-testing your logic against all potential paths, edge cases, and adversarial scenarios.\n",
27 + "Explicitly write out your entire deliberation process, documenting every intermediate step, considered alternative, and rejected hypothesis to ensure absolutely no assumption is left unchecked.\n\n",
28 +);
29 +
30 +#[derive(Debug, Clone, Copy, PartialEq, Eq)]
31 +enum ThinkingMode {
32 + Chat,
33 + Thinking,
34 +}
35 +
36 +#[serde_with::skip_serializing_none]
37 +#[derive(Debug, Serialize)]
38 +struct RenderedToolSchema<'a> {
39 + name: &'a str,
40 + description: Option<&'a str>,
41 + parameters: &'a Value,
42 + strict: Option<bool>,
43 +}
44 +
45 +/// Render one chat request into the final prompt string.
46 +pub(super) fn render_request(request: &ChatRequest) -> Result<String> {
47 + let (thinking_mode, max_reasoning_effort) = resolve_thinking_options(request)?;
48 + let request_tools = request_tools(request);
49 + let synthetic_tool_system = needs_synthetic_tool_system(request, request_tools);
50 + let drop_thinking = request.parse_template_bool("drop_thinking")?.unwrap_or(true)
51 + && !rendered_tools_present(request, request_tools);
52 + let last_user_render_index =
53 + find_last_user_render_index(request.messages.as_slice(), synthetic_tool_system);
54 + let mut out = String::from(BOS_TOKEN);
55 + if thinking_mode == ThinkingMode::Thinking && max_reasoning_effort {
56 + out.push_str(REASONING_EFFORT_MAX);
57 + }
58 +
59 + let mut request_tools_attached = false;
60 + let mut render_index = 0isize;
61 + if synthetic_tool_system {
62 + render_system_message(&mut out, None, request_tools)?;
63 + request_tools_attached = true;
64 + render_index += 1;
65 + }
66 +
67 + for (message_index, message) in request.messages.iter().enumerate() {
68 + if is_following_tool_response(request.messages.as_slice(), message_index) {
69 + continue;
70 + }
71 +
72 + let current_render_index = render_index;
73 + render_index += 1;
74 +
75 + match message {
76 + ChatMessage::System { content } => {
77 + let tools = if !request_tools_attached {
78 + request_tools_attached = true;
79 + request_tools
80 + } else {
81 + &[]
82 + };
83 + render_system_message(&mut out, Some(content), tools)?;
84 + }
85 + ChatMessage::Developer { content, tools } => {
86 + render_developer_message(&mut out, content, tools.as_deref().unwrap_or(&[]))?;
87 + }
88 + ChatMessage::User { content } => render_user_message(&mut out, content)?,
89 + ChatMessage::Assistant { content } => {
90 + // Mirror Python: thinking block (reasoning + </think>) is
91 + // emitted whenever thinking is active and reasoning isn't
92 + // dropped - i.e. drop_thinking is off OR this turn lies
93 + // strictly after the last user turn.
94 + let emit_thinking_block = thinking_mode == ThinkingMode::Thinking
95 + && (!drop_thinking || current_render_index > last_user_render_index);
96 + let append_eos = !(message_index + 1 == request.messages.len()
97 + && request.chat_options.continue_final_message());
98 + render_assistant_message(&mut out, emit_thinking_block, append_eos, content)?;
99 + }
100 + ChatMessage::ToolResponse { .. } => {
101 + render_tool_response_block(&mut out, request.messages.as_slice(), message_index)?;
102 + }
103 + }
104 +
105 + if is_user_like_entry(message)
106 + && next_rendered_entry_is_assistant_or_end(request.messages.as_slice(), message_index)
107 + {
108 + write_assistant_transition(
109 + &mut out,
110 + thinking_mode,
111 + drop_thinking,
112 + current_render_index >= last_user_render_index,
113 + );
114 + }
115 + }
116 +
117 + Ok(out)
118 +}
119 +
120 +/// Resolve DeepSeek V4's thinking controls. Unlike the Python tokenizer
121 +/// wrapper, the Rust renderer only consumes the typed top-level
122 +/// `reasoning_effort`; the generic template-kwargs map is left for HF
123 +/// templates.
124 +fn resolve_thinking_options(request: &ChatRequest) -> Result<(ThinkingMode, bool)> {
125 + let mut thinking_mode = match request.enable_thinking()?.unwrap_or(false) {
126 + true => ThinkingMode::Thinking,
127 + false => ThinkingMode::Chat,
128 + };
129 + let mut max_reasoning_effort = false;
130 +
131 + match request.chat_options.reasoning_effort {
132 + Some(ReasoningEffort::None) => thinking_mode = ThinkingMode::Chat,
133 + Some(ReasoningEffort::Max | ReasoningEffort::XHigh) => max_reasoning_effort = true,
134 + Some(_) | None => {}
135 + }
136 +
137 + Ok((thinking_mode, max_reasoning_effort))
138 +}
139 +
140 +/// Return request-level tools only when native tool parsing is enabled.
141 +fn request_tools(request: &ChatRequest) -> &[ChatTool] {
142 + if request.tool_parsing_enabled() {
143 + request.tools.as_slice()
144 + } else {
145 + &[]
146 + }
147 +}
148 +
149 +/// Return whether request tools need a synthetic leading system entry.
150 +fn needs_synthetic_tool_system(request: &ChatRequest, request_tools: &[ChatTool]) -> bool {
151 + !request_tools.is_empty()
152 + && !request
153 + .messages
154 + .iter()
155 + .any(|message| matches!(message, ChatMessage::System { .. }))
156 +}
157 +
158 +/// Return whether any rendered message carries tool schemas.
159 +fn rendered_tools_present(request: &ChatRequest, request_tools: &[ChatTool]) -> bool {
160 + !request_tools.is_empty()
161 + || request.messages.iter().any(|message| {
162 + matches!(
163 + message,
164 + ChatMessage::Developer {
165 + tools: Some(tools),
166 + ..
167 + } if !tools.is_empty()
168 + )
169 + })
170 +}
171 +
172 +/// Find the last user-like turn after inline tool-response merging.
173 +fn find_last_user_render_index(messages: &[ChatMessage], synthetic_tool_system: bool) -> isize {
174 + let mut render_index = isize::from(synthetic_tool_system);
175 + let mut last_user_index = -1;
176 +
177 + for (message_index, message) in messages.iter().enumerate() {
178 + if is_following_tool_response(messages, message_index) {
179 + continue;
180 + }
181 +
182 + if is_user_like_entry(message) {
183 + last_user_index = render_index;
184 + }
185 + render_index += 1;
186 + }
187 +
188 + last_user_index
189 +}
190 +
191 +/// Return whether this tool message is already covered by a previous tool run.
192 +fn is_following_tool_response(messages: &[ChatMessage], message_index: usize) -> bool {
193 + matches!(messages[message_index], ChatMessage::ToolResponse { .. })
194 + && message_index > 0
195 + && matches!(
196 + messages[message_index - 1],
197 + ChatMessage::ToolResponse { .. }
198 + )
199 +}
200 +
201 +/// Return whether one rendered entry should be treated as user-like.
202 +fn is_user_like_entry(message: &ChatMessage) -> bool {
203 + matches!(
204 + message,
205 + ChatMessage::Developer { .. } | ChatMessage::User { .. } | ChatMessage::ToolResponse { .. }
206 + )
207 +}
208 +
209 +/// Return whether the next rendered entry is assistant, or there is no next
210 +/// entry.
211 +fn next_rendered_entry_is_assistant_or_end(messages: &[ChatMessage], message_index: usize) -> bool {
212 + let mut next_index = message_index + 1;
213 + if matches!(messages[message_index], ChatMessage::ToolResponse { .. }) {
214 + while next_index < messages.len()
215 + && matches!(messages[next_index], ChatMessage::ToolResponse { .. })
216 + {
217 + next_index += 1;
218 + }
219 + }
220 +
221 + messages
222 + .get(next_index)
223 + .map(|message| matches!(message, ChatMessage::Assistant { .. }))
224 + .unwrap_or(true)
225 +}
226 +
227 +/// Render the tool preamble shown to the model, V4 flavor.
228 +fn render_tools(out: &mut String, tools: &[ChatTool]) -> Result<()> {
229 + out.push_str(
230 + r#"## Tools
231 +
232 +You have access to a set of tools to help answer the user's question. You can invoke tools by writing a "<｜DSML｜tool_calls>" block like the following:
233 +
234 +<｜DSML｜tool_calls>
235 +<｜DSML｜invoke name="$TOOL_NAME">
236 +<｜DSML｜parameter name="$PARAMETER_NAME" string="true|false">$PARAMETER_VALUE</｜DSML｜parameter>
237 +...
238 +</｜DSML｜invoke>
239 +<｜DSML｜invoke name="$TOOL_NAME2">
240 +...
241 +</｜DSML｜invoke>
242 +</｜DSML｜tool_calls>
243 +
244 +String parameters should be specified as is and set `string="true"`. For all other types (numbers, booleans, arrays, objects), pass the value in JSON format and set `string="false"`.
245 +
246 +If thinking_mode is enabled (triggered by <think>), you MUST output your complete reasoning inside <think>...</think> BEFORE any tool calls or final response.
247 +
248 +Otherwise, output directly after </think> with tool calls or final response.
249 +
250 +### Available Tool Schemas
251 +
252 +"#,
253 + );
254 +
255 + for (index, tool) in tools.iter().enumerate() {
256 + if index > 0 {
257 + out.push('\n');
258 + }
259 + render_tool_schema(out, tool)?;
260 + }
261 +
262 + out.push_str(
263 + "\n\nYou MUST strictly follow the above defined tool name and parameter schemas to invoke tool calls.\n",
264 + );
265 + Ok(())
266 +}
267 +
268 +/// Serialize one typed tool schema into the JSON shape embedded in the prompt.
269 +fn render_tool_schema(out: &mut String, tool: &ChatTool) -> Result<()> {
270 + out.push_str(&json_dumps(&RenderedToolSchema {
271 + name: &tool.name,
272 + description: tool.description.as_deref(),
273 + parameters: &tool.parameters,
274 + strict: tool.strict,
275 + })?);
276 + Ok(())
277 +}
278 +
279 +/// Render a system turn, optionally followed by the V4 tool preamble.
280 +fn render_system_message(
281 + out: &mut String,
282 + content: Option<&ChatContent>,
283 + tools: &[ChatTool],
284 +) -> Result<()> {
285 + if let Some(content) = content {
286 + write_chat_content(out, content)?;
287 + }
288 + if !tools.is_empty() {
289 + out.push_str("\n\n");
290 + render_tools(out, tools)?;
291 + }
292 + Ok(())
293 +}
294 +
295 +/// Developer messages are rendered as user-like turns with optional tools.
296 +fn render_developer_message(
297 + out: &mut String,
298 + content: &ChatContent,
299 + tools: &[ChatTool],
300 +) -> Result<()> {
301 + if content.is_empty() {
302 + return Err(Error::ChatTemplate(
303 + "invalid DeepSeek V4 developer message: empty content".to_string(),
304 + ));
305 + }
306 +
307 + out.push_str(USER_SP_TOKEN);
308 + write_chat_content(out, content)?;
309 + if !tools.is_empty() {
310 + out.push_str("\n\n");
311 + render_tools(out, tools)?;
312 + }
313 + Ok(())
314 +}
315 +
316 +/// Render one plain user turn.
317 +fn render_user_message(out: &mut String, content: &ChatContent) -> Result<()> {
318 + out.push_str(USER_SP_TOKEN);
319 + write_chat_content(out, content)?;
320 + Ok(())
321 +}
322 +
323 +/// Render a contiguous tool-response run as one synthetic user turn.
324 +fn render_tool_response_block(
325 + out: &mut String,
326 + messages: &[ChatMessage],
327 + message_index: usize,
328 +) -> Result<()> {
329 + let (block_start, block_end) = tool_response_block_bounds(messages, message_index);
330 + let sorted_indices = sorted_tool_response_indices(messages, block_start, block_end);
331 +
332 + out.push_str(USER_SP_TOKEN);
333 + for (offset, message_index) in sorted_indices.iter().enumerate() {
334 + if offset > 0 {
335 + out.push_str("\n\n");
336 + }
337 + let ChatMessage::ToolResponse { content, .. } = &messages[*message_index] else {
338 + unreachable!("tool response block should only contain tool messages");
339 + };
340 + write_tool_result(out, content)?;
341 + }
342 +
343 + Ok(())
344 +}
345 +
346 +/// Return the contiguous tool-response block containing `actual_index`.
347 +fn tool_response_block_bounds(messages: &[ChatMessage], actual_index: usize) -> (usize, usize) {
348 + let mut block_start = actual_index;
349 + while block_start > 0 && matches!(messages[block_start - 1], ChatMessage::ToolResponse { .. }) {
350 + block_start -= 1;
351 + }
352 +
353 + let mut block_end = actual_index + 1;
354 + while block_end < messages.len()
355 + && matches!(messages[block_end], ChatMessage::ToolResponse { .. })
356 + {
357 + block_end += 1;
358 + }
359 +
360 + (block_start, block_end)
361 +}
362 +
363 +fn sorted_tool_response_indices(
364 + messages: &[ChatMessage],
365 + block_start: usize,
366 + block_end: usize,
367 +) -> Vec<usize> {
368 + let Some(tool_call_order) = last_tool_call_order_before(messages, block_start) else {
369 + return (block_start..block_end).collect();
370 + };
371 +
372 + let mut indices = (block_start..block_end).collect::<Vec<_>>();
373 + indices.sort_by_key(|index| {
374 + let ChatMessage::ToolResponse { tool_call_id, .. } = &messages[*index] else {
375 + unreachable!("tool response block should only contain tool messages");
376 + };
377 + tool_call_order.get(tool_call_id.as_str()).copied().unwrap_or(0)
378 + });
379 + indices
380 +}
381 +
382 +fn last_tool_call_order_before(
383 + messages: &[ChatMessage],
384 + message_index: usize,
385 +) -> Option<HashMap<&str, usize>> {
386 + let mut tool_call_order = None;
387 + for message in &messages[..message_index] {
388 + if let ChatMessage::Assistant { content } = message {
389 + let order = content
390 + .tool_calls()
391 + .enumerate()
392 + .map(|(index, tool_call)| (tool_call.id.as_str(), index))
393 + .collect::<HashMap<_, _>>();
394 + if !order.is_empty() {
395 + tool_call_order = Some(order);
396 + }
397 + }
398 + }
399 + tool_call_order
400 +}
401 +
402 +/// Render one tool response payload inside a V4 `<tool_result>` block.
403 +fn write_tool_result(out: &mut String, content: &ChatContent) -> Result<()> {
404 + out.push_str("<tool_result>");
405 + write_chat_content(out, content)?;
406 + out.push_str("</tool_result>");
407 + Ok(())
408 +}
409 +
410 +/// Append the assistant transition token after a user-like turn.
411 +fn write_assistant_transition(
412 + out: &mut String,
413 + thinking_mode: ThinkingMode,
414 + drop_thinking: bool,
415 + opens_thinking: bool,
416 +) {
417 + out.push_str(ASSISTANT_SP_TOKEN);
418 + if thinking_mode == ThinkingMode::Thinking && (!drop_thinking || opens_thinking) {
419 + out.push_str(THINKING_START_TOKEN);
420 + } else {
421 + out.push_str(THINKING_END_TOKEN);
422 + }
423 +}
424 +
425 +/// Render one assistant turn, including optional reasoning, DSML tool calls,
426 +/// and the trailing EOS marker.
427 +fn render_assistant_message(
428 + out: &mut String,
429 + emit_thinking_block: bool,
430 + append_eos: bool,
431 + content: &[AssistantContentBlock],
432 +) -> Result<()> {
433 + let has_tool_calls = content.has_tool_calls();
434 +
435 + if emit_thinking_block {
436 + if content.has_reasoning() {
437 + write_assistant_reasoning(out, content);
438 + }
439 + out.push_str(THINKING_END_TOKEN);
440 + }
441 +
442 + write_assistant_text(out, content);
443 +
444 + if has_tool_calls {
445 + out.push_str("\n\n<｜DSML｜tool_calls>\n");
446 + for (index, tool_call) in content.tool_calls().enumerate() {
447 + if index > 0 {
448 + out.push('\n');
449 + }
450 + render_tool_call(out, tool_call)?;
451 + }
452 + out.push_str("\n</｜DSML｜tool_calls>");
453 + }
454 +
455 + if append_eos {
456 + out.push_str(EOS_TOKEN);
457 + }
458 + Ok(())
459 +}
460 +
461 +/// Render one assistant tool call in DSML XML-like format.
462 +fn render_tool_call(out: &mut String, tool_call: &AssistantToolCall) -> Result<()> {
463 + writeln!(out, "<{DSML_TOKEN}invoke name=\"{}\">", tool_call.name)
464 + .expect("writing to String cannot fail");
465 + encode_arguments_to_dsml(out, tool_call)?;
466 + write!(out, "\n</{DSML_TOKEN}invoke>").expect("writing to String cannot fail");
467 + Ok(())
468 +}
469 +
470 +/// Convert one assistant tool-call arguments object into DSML parameter form.
471 +///
472 +/// String values are emitted raw with `string="true"`, while all other JSON
473 +/// values are rendered with JSON syntax and `string="false"`.
474 +fn encode_arguments_to_dsml(out: &mut String, tool_call: &AssistantToolCall) -> Result<()> {
475 + let arguments: Value = serde_json::from_str(&tool_call.arguments).map_err(|error| {
476 + Error::ChatTemplate(format!(
477 + "assistant tool call has invalid JSON arguments for DeepSeek V4: {error}"
478 + ))
479 + })?;
480 + let Some(arguments) = arguments.as_object() else {
481 + return Err(Error::ChatTemplate(
482 + "assistant tool call arguments for DeepSeek V4 must be a JSON object".to_string(),
483 + ));
484 + };
485 +
486 + let mut wrote_parameter = false;
487 + for (key, value) in arguments {
488 + if wrote_parameter {
489 + out.push('\n');
490 + }
491 +
492 + let is_string = matches!(value, Value::String(_));
493 + write!(
494 + out,
495 + "<{DSML_TOKEN}parameter name=\"{key}\" string=\"{}\">",
496 + if is_string { "true" } else { "false" }
497 + )
498 + .expect("writing to String cannot fail");
499 +
500 + match value {
501 + Value::String(value) => out.push_str(value),
502 + value => out.push_str(&json_dumps(value)?),
503 + }
504 +
505 + write!(out, "</{DSML_TOKEN}parameter>").expect("writing to String cannot fail");
506 + wrote_parameter = true;
507 + }
508 +
509 + Ok(())
510 +}
511 +
512 +/// Write chat content directly into the destination buffer without flattening
513 +/// it into an intermediate `String`.
514 +fn write_chat_content(out: &mut String, content: &ChatContent) -> Result<()> {
515 + match content {
516 + ChatContent::Text(text) => out.push_str(text),
517 + ChatContent::Parts(parts) => {
518 + for part in parts {
519 + out.push_str(part.as_text()?);
520 + }
521 + }
522 + }
523 + Ok(())
524 +}
525 +
526 +/// Write all reasoning blocks in encounter order.
527 +fn write_assistant_reasoning(out: &mut String, content: &[AssistantContentBlock]) {
528 + for block in content {
529 + if let AssistantContentBlock::Reasoning { text } = block {
530 + out.push_str(text);
531 + }
532 + }
533 +}
534 +
535 +/// Write all visible assistant text blocks in encounter order.
536 +fn write_assistant_text(out: &mut String, content: &[AssistantContentBlock]) {
537 + for block in content {
538 + if let AssistantContentBlock::Text { text } = block {
539 + out.push_str(text);
540 + }
541 + }
542 +}
543 +
544 +/// Compact JSON serialization used by this renderer for exact prompt text.
545 +fn json_dumps<T: Serialize>(value: &T) -> Result<String> {
546 + JsonFormat::new()
547 + .comma(", ")
548 + .expect("literal comma separator is valid JSON")
549 + .colon(": ")
550 + .expect("literal colon separator is valid JSON")
551 + .ascii(false)
552 + .format_to_string(value)
553 + .map_err(|error| {
554 + Error::ChatTemplate(format!(
555 + "failed to serialize DeepSeek V4 JSON payload: {error}"
556 + ))
557 + })
558 +}
rust/src/chat/src/renderer/deepseek_v4/fixtures/test_input_1.json new
+81
@@ -0,0 +1,81 @@
1 +{
2 + "tools": [
3 + {
4 + "type": "function",
5 + "function": {
6 + "name": "get_weather",
7 + "description": "Get the weather for a specific location",
8 + "parameters": {
9 + "type": "object",
10 + "properties": {
11 + "location": {
12 + "type": "string",
13 + "description": "The city name"
14 + },
15 + "unit": {
16 + "type": "string",
17 + "enum": ["celsius", "fahrenheit"],
18 + "description": "Temperature unit"
19 + }
20 + },
21 + "required": ["location"]
22 + }
23 + }
24 + },
25 + {
26 + "type": "function",
27 + "function": {
28 + "name": "search",
29 + "description": "Search the web for information",
30 + "parameters": {
31 + "type": "object",
32 + "properties": {
33 + "query": {
34 + "type": "string",
35 + "description": "Search query"
36 + },
37 + "num_results": {
38 + "type": "integer",
39 + "description": "Number of results to return"
40 + }
41 + },
42 + "required": ["query"]
43 + }
44 + }
45 + }
46 + ],
47 + "messages": [
48 + {
49 + "role": "system",
50 + "content": "You are a helpful assistant."
51 + },
52 + {
53 + "role": "user",
54 + "content": "What's the weather in Beijing?"
55 + },
56 + {
57 + "role": "assistant",
58 + "reasoning_content": "The user wants to know the weather in Beijing. I should use the get_weather tool.",
59 + "tool_calls": [
60 + {
61 + "id": "call_001",
62 + "type": "function",
63 + "function": {
64 + "name": "get_weather",
65 + "arguments": "{\"location\": \"Beijing\", \"unit\": \"celsius\"}"
66 + }
67 + }
68 + ]
69 + },
70 + {
71 + "role": "tool",
72 + "tool_call_id": "call_001",
73 + "content": "{\"temperature\": 22, \"condition\": \"sunny\", \"humidity\": 45}"
74 + },
75 + {
76 + "role": "assistant",
77 + "reasoning_content": "Got the weather data. Let me format a nice response.",
78 + "content": "The weather in Beijing is currently sunny with a temperature of 22°C and 45% humidity."
79 + }
80 + ]
81 +}
rust/src/chat/src/renderer/deepseek_v4/fixtures/test_input_2.json new
+24
@@ -0,0 +1,24 @@
1 +[
2 + {
3 + "role": "system",
4 + "content": "You are a helpful assistant."
5 + },
6 + {
7 + "role": "user",
8 + "content": "Hello"
9 + },
10 + {
11 + "role": "assistant",
12 + "reasoning_content": "The user said hello, I should greet back.",
13 + "content": "Hi there! How can I help you?"
14 + },
15 + {
16 + "role": "user",
17 + "content": "What is the capital of France?"
18 + },
19 + {
20 + "role": "assistant",
21 + "reasoning_content": "The user asks about the capital of France. It is Paris.",
22 + "content": "The capital of France is Paris."
23 + }
24 +]
\ No newline at end of file
rust/src/chat/src/renderer/deepseek_v4/fixtures/test_output_1.txt new
+36
@@ -0,0 +1,36 @@
1 +<｜begin▁of▁sentence｜>You are a helpful assistant.
2 +
3 +## Tools
4 +
5 +You have access to a set of tools to help answer the user's question. You can invoke tools by writing a "<｜DSML｜tool_calls>" block like the following:
6 +
7 +<｜DSML｜tool_calls>
8 +<｜DSML｜invoke name="$TOOL_NAME">
9 +<｜DSML｜parameter name="$PARAMETER_NAME" string="true|false">$PARAMETER_VALUE</｜DSML｜parameter>
10 +...
11 +</｜DSML｜invoke>
12 +<｜DSML｜invoke name="$TOOL_NAME2">
13 +...
14 +</｜DSML｜invoke>
15 +</｜DSML｜tool_calls>
16 +
17 +String parameters should be specified as is and set `string="true"`. For all other types (numbers, booleans, arrays, objects), pass the value in JSON format and set `string="false"`.
18 +
19 +If thinking_mode is enabled (triggered by <think>), you MUST output your complete reasoning inside <think>...</think> BEFORE any tool calls or final response.
20 +
21 +Otherwise, output directly after </think> with tool calls or final response.
22 +
23 +### Available Tool Schemas
24 +
25 +{"name": "get_weather", "description": "Get the weather for a specific location", "parameters": {"type": "object", "properties": {"location": {"type": "string", "description": "The city name"}, "unit": {"type": "string", "enum": ["celsius", "fahrenheit"], "description": "Temperature unit"}}, "required": ["location"]}}
26 +{"name": "search", "description": "Search the web for information", "parameters": {"type": "object", "properties": {"query": {"type": "string", "description": "Search query"}, "num_results": {"type": "integer", "description": "Number of results to return"}}, "required": ["query"]}}
27 +
28 +You MUST strictly follow the above defined tool name and parameter schemas to invoke tool calls.
29 +<｜User｜>What's the weather in Beijing?<｜Assistant｜><think>The user wants to know the weather in Beijing. I should use the get_weather tool.</think>
30 +
31 +<｜DSML｜tool_calls>
32 +<｜DSML｜invoke name="get_weather">
33 +<｜DSML｜parameter name="location" string="true">Beijing</｜DSML｜parameter>
34 +<｜DSML｜parameter name="unit" string="true">celsius</｜DSML｜parameter>
35 +</｜DSML｜invoke>
36 +</｜DSML｜tool_calls><｜end▁of▁sentence｜><｜User｜><tool_result>{"temperature": 22, "condition": "sunny", "humidity": 45}</tool_result><｜Assistant｜><think>Got the weather data. Let me format a nice response.</think>The weather in Beijing is currently sunny with a temperature of 22°C and 45% humidity.<｜end▁of▁sentence｜>
\ No newline at end of file
rust/src/chat/src/renderer/deepseek_v4/fixtures/test_output_2.txt new
+1
@@ -0,0 +1 @@
1 +<｜begin▁of▁sentence｜>You are a helpful assistant.<｜User｜>Hello<｜Assistant｜></think>Hi there! How can I help you?<｜end▁of▁sentence｜><｜User｜>What is the capital of France?<｜Assistant｜><think>The user asks about the capital of France. It is Paris.</think>The capital of France is Paris.<｜end▁of▁sentence｜>
\ No newline at end of file
rust/src/chat/src/renderer/deepseek_v4/mod.rs new
+30
@@ -0,0 +1,30 @@
1 +mod encoding;
2 +
3 +use vllm_text::Prompt;
4 +
5 +use super::{ChatRenderer, RenderedPrompt};
6 +use crate::Result;
7 +use crate::request::ChatRequest;
8 +
9 +/// Dedicated DeepSeek V4 renderer.
10 +#[derive(Debug, Clone, Copy, Default)]
11 +pub struct DeepSeekV4ChatRenderer;
12 +
13 +impl DeepSeekV4ChatRenderer {
14 + pub fn new() -> Self {
15 + Self
16 + }
17 +}
18 +
19 +impl ChatRenderer for DeepSeekV4ChatRenderer {
20 + fn render(&self, request: &ChatRequest) -> Result<RenderedPrompt> {
21 + request.validate()?;
22 +
23 + Ok(RenderedPrompt {
24 + prompt: Prompt::Text(encoding::render_request(request)?),
25 + })
26 + }
27 +}
28 +
29 +#[cfg(test)]
30 +mod tests;
rust/src/chat/src/renderer/deepseek_v4/tests.rs new
+369
@@ -0,0 +1,369 @@
1 +use std::fs;
2 +use std::path::PathBuf;
3 +
4 +use expect_test::{ExpectFile, expect, expect_file};
5 +use serde::Deserialize;
6 +use serde_json::Value;
7 +
8 +use super::DeepSeekV4ChatRenderer;
9 +use crate::event::{AssistantContentBlock, AssistantToolCall};
10 +use crate::request::{
11 + ChatMessage, ChatRequest, ChatTool, ChatToolChoice, GenerationPromptMode, ReasoningEffort,
12 +};
13 +use crate::{ChatRenderer, ChatRole};
14 +
15 +#[derive(Debug, Deserialize)]
16 +#[serde(untagged)]
17 +enum FixtureFile {
18 + WithTools(FixtureRequest),
19 + MessagesOnly(Vec<FixtureMessage>),
20 +}
21 +
22 +#[derive(Debug, Deserialize)]
23 +struct FixtureRequest {
24 + #[serde(default)]
25 + tools: Vec<FixtureTool>,
26 + messages: Vec<FixtureMessage>,
27 +}
28 +
29 +impl FixtureFile {
30 + fn into_parts(self) -> (Vec<FixtureTool>, Vec<FixtureMessage>) {
31 + match self {
32 + Self::WithTools(req) => (req.tools, req.messages),
33 + Self::MessagesOnly(messages) => (Vec::new(), messages),
34 + }
35 + }
36 +}
37 +
38 +#[derive(Debug, Deserialize)]
39 +struct FixtureTool {
40 + function: FixtureToolFunction,
41 +}
42 +
43 +#[derive(Debug, Deserialize)]
44 +struct FixtureToolFunction {
45 + name: String,
46 + description: Option<String>,
47 + parameters: Value,
48 + #[serde(default)]
49 + strict: Option<bool>,
50 +}
51 +
52 +#[derive(Debug, Deserialize)]
53 +#[serde(tag = "role", rename_all = "snake_case")]
54 +enum FixtureMessage {
55 + System {
56 + content: String,
57 + },
58 + Developer {
59 + content: String,
60 + #[serde(default)]
61 + tools: Vec<FixtureTool>,
62 + },
63 + User {
64 + content: String,
65 + },
66 + Assistant {
67 + #[serde(default)]
68 + content: String,
69 + #[serde(default)]
70 + reasoning_content: String,
71 + #[serde(default)]
72 + tool_calls: Vec<FixtureToolCall>,
73 + },
74 + Tool {
75 + content: String,
76 + #[serde(default)]
77 + tool_call_id: Option<String>,
78 + },
79 +}
80 +
81 +#[derive(Debug, Deserialize)]
82 +struct FixtureToolCall {
83 + #[serde(default)]
84 + id: Option<String>,
85 + function: FixtureToolCallFunction,
86 +}
87 +
88 +#[derive(Debug, Deserialize)]
89 +struct FixtureToolCallFunction {
90 + name: String,
91 + arguments: String,
92 +}
93 +
94 +fn render_request(request: &ChatRequest) -> String {
95 + DeepSeekV4ChatRenderer::new()
96 + .render(request)
97 + .unwrap()
98 + .prompt
99 + .into_text()
100 + .expect("deepseek v4 renderer should return text prompt")
101 +}
102 +
103 +fn fixture_request(input_name: &str) -> ChatRequest {
104 + let fixture = fs::read_to_string(fixture_path(input_name)).unwrap();
105 + let fixture: FixtureFile = serde_json::from_str(&fixture).unwrap();
106 + let (fixture_tools, fixture_messages) = fixture.into_parts();
107 + let mut request = ChatRequest {
108 + request_id: "deepseek-v4-fixture".to_string(),
109 + messages: fixture_messages
110 + .into_iter()
111 + .enumerate()
112 + .map(|(index, message)| match message {
113 + FixtureMessage::System { content } => ChatMessage::system(content),
114 + FixtureMessage::Developer { content, tools } => ChatMessage::developer(
115 + content,
116 + (!tools.is_empty()).then(|| to_chat_tools(&tools)),
117 + ),
118 + FixtureMessage::User { content } => ChatMessage::user(content),
119 + FixtureMessage::Assistant {
120 + content,
121 + reasoning_content,
122 + tool_calls,
123 + } => {
124 + let mut blocks = Vec::new();
125 + if !reasoning_content.is_empty() {
126 + blocks.push(AssistantContentBlock::Reasoning {
127 + text: reasoning_content,
128 + });
129 + }
130 + if !content.is_empty() {
131 + blocks.push(AssistantContentBlock::Text { text: content });
132 + }
133 + blocks.extend(tool_calls.into_iter().enumerate().map(
134 + |(tool_index, tool_call)| {
135 + AssistantContentBlock::ToolCall(AssistantToolCall {
136 + id: tool_call.id.unwrap_or_else(|| {
137 + format!("fixture-tool-call-{index}-{tool_index}")
138 + }),
139 + name: tool_call.function.name,
140 + arguments: tool_call.function.arguments,
141 + })
142 + },
143 + ));
144 + ChatMessage::assistant_blocks(blocks)
145 + }
146 + FixtureMessage::Tool {
147 + content,
148 + tool_call_id,
149 + } => ChatMessage::tool_response(
150 + content,
151 + tool_call_id.unwrap_or_else(|| format!("fixture-tool-response-{index}")),
152 + ),
153 + })
154 + .collect(),
155 + tools: to_chat_tools(&fixture_tools),
156 + tool_choice: if fixture_tools.is_empty() {
157 + ChatToolChoice::None
158 + } else {
159 + ChatToolChoice::Auto
160 + },
161 + ..ChatRequest::for_test()
162 + };
163 + if matches!(
164 + request.messages.last().map(ChatMessage::role),
165 + Some(ChatRole::Assistant)
166 + ) {
167 + request.chat_options.generation_prompt_mode = GenerationPromptMode::NoGenerationPrompt;
168 + }
169 + request
170 + .chat_options
171 + .template_kwargs
172 + .insert("thinking".to_string(), Value::Bool(true));
173 + request
174 +}
175 +
176 +fn to_chat_tools(tools: &[FixtureTool]) -> Vec<ChatTool> {
177 + tools
178 + .iter()
179 + .map(|tool| ChatTool {
180 + name: tool.function.name.clone(),
181 + description: tool.function.description.clone(),
182 + parameters: tool.function.parameters.clone(),
183 + strict: tool.function.strict,
184 + })
185 + .collect()
186 +}
187 +
188 +fn fixture_path(name: &str) -> PathBuf {
189 + PathBuf::from(env!("CARGO_MANIFEST_DIR"))
190 + .join("src/renderer/deepseek_v4")
191 + .join("fixtures")
192 + .join(name)
193 +}
194 +
195 +fn assert_fixture(input_name: &str, expected: ExpectFile) {
196 + let request = fixture_request(input_name);
197 + let rendered = render_request(&request);
198 + expected.assert_eq(&rendered);
199 +}
200 +
201 +#[test]
202 +fn renders_v4_fixture_1_tool_call_round_trip() {
203 + assert_fixture(
204 + "test_input_1.json",
205 + expect_file!["fixtures/test_output_1.txt"],
206 + );
207 +}
208 +
209 +#[test]
210 +fn renders_v4_fixture_2_multi_turn_drop_thinking() {
211 + assert_fixture(
212 + "test_input_2.json",
213 + expect_file!["fixtures/test_output_2.txt"],
214 + );
215 +}
216 +
217 +#[test]
218 +fn reasoning_effort_max_adds_prefix_when_thinking_is_enabled() {
219 + let mut request = ChatRequest {
220 + messages: vec![ChatMessage::user("solve it")],
221 + ..ChatRequest::for_test()
222 + };
223 + request
224 + .chat_options
225 + .template_kwargs
226 + .insert("thinking".to_string(), Value::Bool(true));
227 + request.chat_options.reasoning_effort = Some(ReasoningEffort::Max);
228 +
229 + let rendered = render_request(&request);
230 +
231 + expect![[r#"
232 + <｜begin▁of▁sentence｜>Reasoning Effort: Absolute maximum with no shortcuts permitted.
233 + You MUST be very thorough in your thinking and comprehensively decompose the problem to resolve the root cause, rigorously stress-testing your logic against all potential paths, edge cases, and adversarial scenarios.
234 + Explicitly write out your entire deliberation process, documenting every intermediate step, considered alternative, and rejected hypothesis to ensure absolutely no assumption is left unchecked.
235 +
236 + <｜User｜>solve it<｜Assistant｜><think>"#]]
237 + .assert_eq(&rendered);
238 +}
239 +
240 +#[test]
241 +fn reasoning_effort_none_disables_thinking() {
242 + let mut request = ChatRequest {
243 + messages: vec![ChatMessage::user("answer directly")],
244 + ..ChatRequest::for_test()
245 + };
246 + request
247 + .chat_options
248 + .template_kwargs
249 + .insert("thinking".to_string(), Value::Bool(true));
250 + request.chat_options.reasoning_effort = Some(ReasoningEffort::None);
251 +
252 + let rendered = render_request(&request);
253 +
254 + expect!["<｜begin▁of▁sentence｜><｜User｜>answer directly<｜Assistant｜></think>"]
255 + .assert_eq(&rendered);
256 +}
257 +
258 +#[test]
259 +fn reasoning_effort_template_kwarg_is_ignored() {
260 + let mut request = ChatRequest {
261 + messages: vec![ChatMessage::user("solve it")],
262 + ..ChatRequest::for_test()
263 + };
264 + request
265 + .chat_options
266 + .template_kwargs
267 + .insert("thinking".to_string(), Value::Bool(true));
268 + request.chat_options.template_kwargs.insert(
269 + "reasoning_effort".to_string(),
270 + Value::String("max".to_string()),
271 + );
272 +
273 + let rendered = render_request(&request);
274 +
275 + expect!["<｜begin▁of▁sentence｜><｜User｜>solve it<｜Assistant｜><think>"].assert_eq(&rendered);
276 +}
277 +
278 +#[test]
279 +fn tool_results_are_sorted_by_previous_assistant_tool_call_order() {
280 + let request = ChatRequest {
281 + messages: vec![
282 + ChatMessage::assistant_blocks(vec![
283 + AssistantContentBlock::ToolCall(AssistantToolCall {
284 + id: "second".to_string(),
285 + name: "second_tool".to_string(),
286 + arguments: "{}".to_string(),
287 + }),
288 + AssistantContentBlock::ToolCall(AssistantToolCall {
289 + id: "first".to_string(),
290 + name: "first_tool".to_string(),
291 + arguments: "{}".to_string(),
292 + }),
293 + ]),
294 + ChatMessage::tool_response("first result", "first"),
295 + ChatMessage::tool_response("second result", "second"),
296 + ],
297 + ..ChatRequest::for_test()
298 + };
299 +
300 + let rendered = render_request(&request);
301 +
302 + expect![[r#"
303 + <｜begin▁of▁sentence｜>
304 +
305 + <｜DSML｜tool_calls>
306 + <｜DSML｜invoke name="second_tool">
307 +
308 + </｜DSML｜invoke>
309 + <｜DSML｜invoke name="first_tool">
310 +
311 + </｜DSML｜invoke>
312 + </｜DSML｜tool_calls><｜end▁of▁sentence｜><｜User｜><tool_result>second result</tool_result>
313 +
314 + <tool_result>first result</tool_result><｜Assistant｜></think>"#]]
315 + .assert_eq(&rendered);
316 +}
317 +
318 +#[test]
319 +fn drop_thinking_false_keeps_prior_assistant_reasoning() {
320 + let mut request = ChatRequest {
321 + messages: vec![
322 + ChatMessage::assistant_blocks(vec![
323 + AssistantContentBlock::Reasoning {
324 + text: "old reasoning".to_string(),
325 + },
326 + AssistantContentBlock::Text {
327 + text: "old answer".to_string(),
328 + },
329 + ]),
330 + ChatMessage::user("next"),
331 + ],
332 + ..ChatRequest::for_test()
333 + };
334 + request
335 + .chat_options
336 + .template_kwargs
337 + .insert("thinking".to_string(), Value::Bool(true));
338 + request
339 + .chat_options
340 + .template_kwargs
341 + .insert("drop_thinking".to_string(), Value::Bool(false));
342 +
343 + let rendered = render_request(&request);
344 +
345 + expect!(
346 + "<｜begin▁of▁sentence｜>old reasoning</think>old answer<｜end▁of▁sentence｜><｜User｜>next<｜Assistant｜><think>"
347 + )
348 + .assert_eq(&rendered);
349 +}
350 +
351 +#[test]
352 +fn continue_final_assistant_omits_final_eos() {
353 + let request = ChatRequest {
354 + messages: vec![
355 + ChatMessage::user("write"),
356 + ChatMessage::assistant_text("partial answer"),
357 + ],
358 + chat_options: crate::request::ChatOptions {
359 + generation_prompt_mode: GenerationPromptMode::ContinueFinalAssistant,
360 + ..Default::default()
361 + },
362 + ..ChatRequest::for_test()
363 + };
364 +
365 + let rendered = render_request(&request);
366 +
367 + expect!["<｜begin▁of▁sentence｜><｜User｜>write<｜Assistant｜></think>partial answer"]
368 + .assert_eq(&rendered);
369 +}
rust/src/chat/src/renderer/hf/error.rs new
+15
@@ -0,0 +1,15 @@
1 +use thiserror::Error as ThisError;
2 +
3 +#[derive(Debug, ThisError)]
4 +pub(crate) enum TemplateError {
5 + #[error("failed to render jinja template")]
6 + Jinja(#[from] minijinja::Error),
7 + #[error("failed to read chat template file")]
8 + ReadTemplateFile(#[source] std::io::Error),
9 + #[error("chat template looks like a file path but does not exist")]
10 + MissingTemplatePath,
11 + #[error("failed to parse chat_template.json")]
12 + ParseTemplateJson(#[source] serde_json::Error),
13 + #[error("chat_template.json does not contain a valid template")]
14 + InvalidTemplateJson,
15 +}
rust/src/chat/src/renderer/hf/format.rs new
+400
@@ -0,0 +1,400 @@
1 +use std::collections::{HashSet, VecDeque};
2 +use std::fmt;
3 +use std::str::FromStr;
4 +
5 +use minijinja::machinery::ast::{Expr, ForLoop, Set, Stmt};
6 +use minijinja::machinery::{WhitespaceConfig, parse};
7 +use minijinja::syntax::SyntaxConfig;
8 +use serde_with::DeserializeFromStr;
9 +
10 +/// Chat template content format.
11 +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
12 +pub enum ChatTemplateContentFormat {
13 + /// Content is a simple string.
14 + #[default]
15 + String,
16 + /// Content is a list of structured parts (OpenAI format).
17 + OpenAi,
18 +}
19 +
20 +/// Configurable chat-template content format selection.
21 +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, DeserializeFromStr)]
22 +pub enum ChatTemplateContentFormatOption {
23 + /// Detect the format from the template source.
24 + #[default]
25 + Auto,
26 + /// Always flatten content into plain strings before rendering.
27 + String,
28 + /// Always pass content through in OpenAI-compatible structured form.
29 + OpenAi,
30 +}
31 +
32 +impl ChatTemplateContentFormatOption {
33 + pub const AUTO_LITERAL: &str = "auto";
34 + pub const OPENAI_LITERAL: &str = "openai";
35 + pub const STRING_LITERAL: &str = "string";
36 +}
37 +
38 +impl FromStr for ChatTemplateContentFormatOption {
39 + type Err = String;
40 +
41 + fn from_str(value: &str) -> Result<Self, Self::Err> {
42 + if value.eq_ignore_ascii_case(Self::AUTO_LITERAL) {
43 + Ok(Self::Auto)
44 + } else if value.eq_ignore_ascii_case(Self::STRING_LITERAL) {
45 + Ok(Self::String)
46 + } else if value.eq_ignore_ascii_case(Self::OPENAI_LITERAL) {
47 + Ok(Self::OpenAi)
48 + } else {
49 + Err(format!(
50 + "invalid content format `{value}`; expected one of: {}, {}, {}",
51 + Self::AUTO_LITERAL,
52 + Self::STRING_LITERAL,
53 + Self::OPENAI_LITERAL
54 + ))
55 + }
56 + }
57 +}
58 +
59 +impl fmt::Display for ChatTemplateContentFormatOption {
60 + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
61 + match self {
62 + Self::Auto => f.write_str(Self::AUTO_LITERAL),
63 + Self::String => f.write_str(Self::STRING_LITERAL),
64 + Self::OpenAi => f.write_str(Self::OPENAI_LITERAL),
65 + }
66 + }
67 +}
68 +
69 +fn is_var_access(expr: &Expr, varname: &str) -> bool {
70 + matches!(expr, Expr::Var(v) if v.id == varname)
71 +}
72 +
73 +fn is_const_str(expr: &Expr, value: &str) -> bool {
74 + matches!(expr, Expr::Const(c) if c.value.as_str() == Some(value))
75 +}
76 +
77 +fn is_attr_access(expr: &Expr, varname: &str, key: &str) -> bool {
78 + match expr {
79 + Expr::GetItem(g) => is_var_access(&g.expr, varname) && is_const_str(&g.subscript_expr, key),
80 + Expr::GetAttr(g) => is_var_access(&g.expr, varname) && g.name == key,
81 + _ => false,
82 + }
83 +}
84 +
85 +fn is_var_or_elems_access(expr: &Expr, varname: &str, key: Option<&str>) -> bool {
86 + match expr {
87 + Expr::Filter(f) => {
88 + f.expr.as_ref().is_some_and(|inner| is_var_or_elems_access(inner, varname, key))
89 + }
90 + Expr::Test(t) => is_var_or_elems_access(&t.expr, varname, key),
91 + Expr::Slice(s) => is_var_or_elems_access(&s.expr, varname, key),
92 + _ => key.map_or_else(
93 + || is_var_access(expr, varname),
94 + |key| is_attr_access(expr, varname, key),
95 + ),
96 + }
97 +}
98 +
99 +fn visit_stmt<'a>(
100 + stmt: &'a Stmt<'a>,
101 + assignments: &mut Vec<&'a Set<'a>>,
102 + loops: &mut Vec<&'a ForLoop<'a>>,
103 +) {
104 + match stmt {
105 + Stmt::Template(t) => {
106 + for child in &t.children {
107 + visit_stmt(child, assignments, loops);
108 + }
109 + }
110 + Stmt::ForLoop(fl) => {
111 + loops.push(fl);
112 + for child in &fl.body {
113 + visit_stmt(child, assignments, loops);
114 + }
115 + for child in &fl.else_body {
116 + visit_stmt(child, assignments, loops);
117 + }
118 + }
119 + Stmt::IfCond(ic) => {
120 + for child in &ic.true_body {
121 + visit_stmt(child, assignments, loops);
122 + }
123 + for child in &ic.false_body {
124 + visit_stmt(child, assignments, loops);
125 + }
126 + }
127 + Stmt::WithBlock(wb) => {
128 + for child in &wb.body {
129 + visit_stmt(child, assignments, loops);
130 + }
131 + }
132 + Stmt::Set(set_stmt) => assignments.push(set_stmt),
133 + Stmt::SetBlock(sb) => {
134 + for child in &sb.body {
135 + visit_stmt(child, assignments, loops);
136 + }
137 + }
138 + Stmt::AutoEscape(ae) => {
139 + for child in &ae.body {
140 + visit_stmt(child, assignments, loops);
141 + }
142 + }
143 + Stmt::FilterBlock(fb) => {
144 + for child in &fb.body {
145 + visit_stmt(child, assignments, loops);
146 + }
147 + }
148 + Stmt::Block(b) => {
149 + for child in &b.body {
150 + visit_stmt(child, assignments, loops);
151 + }
152 + }
153 + Stmt::Macro(m) => {
154 + for child in &m.body {
155 + visit_stmt(child, assignments, loops);
156 + }
157 + }
158 + Stmt::CallBlock(cb) => {
159 + for child in &cb.macro_decl.body {
160 + visit_stmt(child, assignments, loops);
161 + }
162 + }
163 + _ => {}
164 + }
165 +}
166 +
167 +fn collect_assignments_and_loops<'a>(
168 + root: &'a Stmt<'a>,
169 +) -> (Vec<&'a Set<'a>>, Vec<&'a ForLoop<'a>>) {
170 + let mut assignments = Vec::new();
171 + let mut loops = Vec::new();
172 + visit_stmt(root, &mut assignments, &mut loops);
173 + (assignments, loops)
174 +}
175 +
176 +fn iter_nodes_assign_var_or_elems(root: &Stmt<'_>, varname: &str) -> Vec<String> {
177 + let (assignments, _) = collect_assignments_and_loops(root);
178 +
179 + let mut discovered = vec![varname.to_string()];
180 + let mut seen = HashSet::from([varname.to_string()]);
181 + let mut related = VecDeque::from([varname.to_string()]);
182 +
183 + while let Some(related_varname) = related.pop_front() {
184 + for assign in &assignments {
185 + let Expr::Var(lhs) = &assign.target else {
186 + continue;
187 + };
188 +
189 + if is_var_or_elems_access(&assign.expr, &related_varname, None) {
190 + let lhs_name = lhs.id.to_string();
191 + if seen.insert(lhs_name.clone()) {
192 + discovered.push(lhs_name.clone());
193 + if lhs_name != related_varname {
194 + related.push_back(lhs_name);
195 + }
196 + }
197 + }
198 + }
199 + }
200 +
201 + discovered
202 +}
203 +
204 +fn iter_nodes_assign_messages_item(root: &Stmt<'_>) -> Vec<String> {
205 + let message_varnames = iter_nodes_assign_var_or_elems(root, "messages");
206 + let (_, loops) = collect_assignments_and_loops(root);
207 +
208 + let mut discovered = Vec::new();
209 + let mut seen = HashSet::new();
210 +
211 + for loop_ast in loops {
212 + let Expr::Var(target) = &loop_ast.target else {
213 + continue;
214 + };
215 +
216 + if message_varnames
217 + .iter()
218 + .any(|varname| is_var_or_elems_access(&loop_ast.iter, varname, None))
219 + {
220 + let target_name = target.id.to_string();
221 + if seen.insert(target_name.clone()) {
222 + discovered.push(target_name);
223 + }
224 + }
225 + }
226 +
227 + discovered
228 +}
229 +
230 +fn has_content_item_loop(root: &Stmt<'_>) -> bool {
231 + let message_varnames = iter_nodes_assign_messages_item(root);
232 + let (_, loops) = collect_assignments_and_loops(root);
233 +
234 + loops.into_iter().any(|loop_ast| {
235 + matches!(loop_ast.target, Expr::Var(_))
236 + && message_varnames
237 + .iter()
238 + .any(|varname| is_var_or_elems_access(&loop_ast.iter, varname, Some("content")))
239 + })
240 +}
241 +
242 +/// Detect the content format expected by a Jinja2 chat template based on AST
243 +/// analysis.
244 +pub fn detect_chat_template_content_format(template: &str) -> ChatTemplateContentFormat {
245 + let ast = match parse(
246 + template,
247 + "template",
248 + SyntaxConfig {},
249 + WhitespaceConfig::default(),
250 + ) {
251 + Ok(ast) => ast,
252 + Err(_) => return ChatTemplateContentFormat::String,
253 + };
254 +
255 + if has_content_item_loop(&ast) {
256 + ChatTemplateContentFormat::OpenAi
257 + } else {
258 + ChatTemplateContentFormat::String
259 + }
260 +}
261 +
262 +#[cfg(test)]
263 +mod tests {
264 + use std::fs;
265 + use std::path::{Path, PathBuf};
266 +
267 + use expect_test::expect;
268 +
269 + use super::{ChatTemplateContentFormat, detect_chat_template_content_format};
270 +
271 + fn detect(template: &str) -> ChatTemplateContentFormat {
272 + detect_chat_template_content_format(template)
273 + }
274 +
275 + fn vllm_examples_dir() -> PathBuf {
276 + Path::new(env!("CARGO_MANIFEST_DIR"))
277 + .join("tests/templates/vllm_examples")
278 + .canonicalize()
279 + .expect("vLLM example template directory should exist locally")
280 + }
281 +
282 + fn read_vllm_example(relative_path: &str) -> String {
283 + fs::read_to_string(vllm_examples_dir().join(relative_path))
284 + .unwrap_or_else(|_| panic!("failed to read vLLM example template: {relative_path}"))
285 + }
286 +
287 + fn iter_vllm_example_template_paths() -> impl Iterator<Item = PathBuf> {
288 + let mut paths = fs::read_dir(vllm_examples_dir())
289 + .expect("failed to read vLLM example template directory")
290 + .map(|entry| entry.expect("failed to read vLLM example template dir entry").path())
291 + .filter(|path| path.extension().is_some_and(|ext| ext == "jinja"))
292 + .collect::<Vec<_>>();
293 + paths.sort();
294 + paths.into_iter()
295 + }
296 +
297 + #[test]
298 + fn detects_string_template_without_content_loop() {
299 + assert_eq!(
300 + detect("{% for message in messages %}{{ message.content }}{% endfor %}"),
301 + ChatTemplateContentFormat::String
302 + );
303 + }
304 +
305 + #[test]
306 + fn detects_openai_template_with_direct_content_loop() {
307 + assert_eq!(
308 + detect(
309 + "{% for message in messages %}{% for content in message['content'] %}{{ content }}{% endfor %}{% endfor %}"
310 + ),
311 + ChatTemplateContentFormat::OpenAi
312 + );
313 + }
314 +
315 + #[test]
316 + fn detects_openai_template_with_messages_alias() {
317 + assert_eq!(
318 + detect(
319 + "{% set msgs = messages %}{% for message in msgs %}{% for content in message.content %}{{ content }}{% endfor %}{% endfor %}"
320 + ),
321 + ChatTemplateContentFormat::OpenAi
322 + );
323 + }
324 +
325 + #[test]
326 + fn does_not_detect_content_alias_loop_as_openai() {
327 + assert_eq!(
328 + detect(
329 + "{% for message in messages %}{% set parts = message.content %}{% for item in parts %}{{ item }}{% endfor %}{% endfor %}"
330 + ),
331 + ChatTemplateContentFormat::String
332 + );
333 + }
334 +
335 + #[test]
336 + fn does_not_treat_length_or_index_access_as_openai() {
337 + assert_eq!(
338 + detect("{% for message in messages %}{{ message.content|length }}{% endfor %}"),
339 + ChatTemplateContentFormat::String
340 + );
341 + assert_eq!(
342 + detect("{% for message in messages %}{{ message.content[0] }}{% endfor %}"),
343 + ChatTemplateContentFormat::String
344 + );
345 + }
346 +
347 + #[test]
348 + fn matches_vllm_example_template_formats() {
349 + let snapshot = iter_vllm_example_template_paths()
350 + .map(|path| {
351 + let file_name = path
352 + .file_name()
353 + .and_then(|name| name.to_str())
354 + .expect("template file name should be valid UTF-8");
355 + let template = read_vllm_example(file_name);
356 + let format = detect(&template);
357 + format!("{file_name:50} => {format:?}")
358 + })
359 + .collect::<Vec<_>>()
360 + .join("\n");
361 +
362 + expect![[r#"
363 + template_alpaca.jinja => String
364 + template_baichuan.jinja => String
365 + template_chatglm.jinja => String
366 + template_chatglm2.jinja => String
367 + template_chatml.jinja => String
368 + template_falcon.jinja => String
369 + template_falcon_180b.jinja => String
370 + template_inkbot.jinja => String
371 + template_teleflm.jinja => String
372 + tool_chat_template_deepseekr1.jinja => String
373 + tool_chat_template_deepseekv3.jinja => String
374 + tool_chat_template_deepseekv31.jinja => String
375 + tool_chat_template_functiongemma.jinja => String
376 + tool_chat_template_gemma3_pythonic.jinja => OpenAi
377 + tool_chat_template_gemma4.jinja => OpenAi
378 + tool_chat_template_glm4.jinja => String
379 + tool_chat_template_granite.jinja => String
380 + tool_chat_template_granite_20b_fc.jinja => String
381 + tool_chat_template_hermes.jinja => String
382 + tool_chat_template_hunyuan_a13b.jinja => String
383 + tool_chat_template_internlm2_tool.jinja => String
384 + tool_chat_template_llama3.1_json.jinja => OpenAi
385 + tool_chat_template_llama3.2_json.jinja => OpenAi
386 + tool_chat_template_llama3.2_pythonic.jinja => String
387 + tool_chat_template_llama4_json.jinja => OpenAi
388 + tool_chat_template_llama4_pythonic.jinja => OpenAi
389 + tool_chat_template_minimax_m1.jinja => OpenAi
390 + tool_chat_template_mistral.jinja => String
391 + tool_chat_template_mistral3.jinja => OpenAi
392 + tool_chat_template_mistral_parallel.jinja => String
393 + tool_chat_template_phi4_mini.jinja => String
394 + tool_chat_template_qwen3coder.jinja => String
395 + tool_chat_template_toolace.jinja => String
396 + tool_chat_template_xlam_llama.jinja => String
397 + tool_chat_template_xlam_qwen.jinja => String"#]]
398 + .assert_eq(&snapshot);
399 + }
400 +}
rust/src/chat/src/renderer/hf/mod.rs new
+970
@@ -0,0 +1,970 @@
1 +use std::collections::HashMap;
2 +
3 +use serde::Serialize;
4 +use serde_json::Value;
5 +use thiserror_ext::AsReport as _;
6 +use tracing::{info, trace, warn};
7 +use vllm_text::Prompt;
8 +use vllm_text::backend::hf::{
9 + HfSpecialTokens, HfTokenizerConfig, ResolvedModelFiles, load_tokenizer_config,
10 +};
11 +
12 +use self::format::{
13 + ChatTemplateContentFormat, ChatTemplateContentFormatOption as ContentFormatOption,
14 +};
15 +use self::template::{CompiledChatTemplate, TemplateContext};
16 +use super::{ChatRenderer, RenderedPrompt};
17 +use crate::error::Result;
18 +use crate::request::{ChatContent, ChatContentPart, ChatMessage, ChatRequest};
19 +use crate::{
20 + AssistantContentBlock, AssistantMessageExt, ChatTool, Error, LoadModelBackendsOptions,
21 +};
22 +
23 +mod error;
24 +mod format;
25 +mod template;
26 +mod tojson;
27 +
28 +pub use template::{load_chat_template, resolve_chat_template};
29 +
30 +pub use self::format::ChatTemplateContentFormatOption;
31 +
32 +#[derive(Debug, Clone)]
33 +pub struct MultimodalRenderInfo {
34 + pub placeholder_token: String,
35 +}
36 +
37 +/// Hugging Face chat-template renderer backed by the local Jinja chat-template
38 +/// state.
39 +pub struct HfChatRenderer {
40 + default_template: Option<CompiledChatTemplate>,
41 + default_template_kwargs: HashMap<String, Value>,
42 + content_format: ContentFormatOption,
43 + special_tokens: Option<HfSpecialTokens>,
44 + multimodal: Option<MultimodalRenderInfo>,
45 +}
46 +
47 +impl HfChatRenderer {
48 + /// Create a renderer from the given template string.
49 + pub fn new(
50 + template: Option<String>,
51 + default_template_kwargs: HashMap<String, Value>,
52 + content_format: ContentFormatOption,
53 + ) -> Result<Self> {
54 + Ok(Self {
55 + default_template: template
56 + .map(|template| {
57 + CompiledChatTemplate::new(template, content_format)
58 + .map_err(|error| Error::ChatTemplate(error.to_report_string()))
59 + })
60 + .transpose()?,
61 + default_template_kwargs,
62 + content_format,
63 + special_tokens: None,
64 + multimodal: None,
65 + })
66 + }
67 +
68 + pub fn with_special_tokens(mut self, special_tokens: Option<HfSpecialTokens>) -> Self {
69 + self.special_tokens = special_tokens;
70 + self
71 + }
72 +
73 + pub fn with_multimodal(mut self, multimodal: Option<MultimodalRenderInfo>) -> Self {
74 + self.multimodal = multimodal;
75 + self
76 + }
77 +
78 + /// Create a renderer from the given model files and loading options.
79 + pub fn load(
80 + files: &ResolvedModelFiles,
81 + options: LoadModelBackendsOptions,
82 + multimodal: Option<MultimodalRenderInfo>,
83 + ) -> Result<Self> {
84 + let HfTokenizerConfig {
85 + special_tokens,
86 + chat_template,
87 + ..
88 + } = load_tokenizer_config(files.tokenizer_config_path.as_deref())?;
89 + let mut template = chat_template;
90 + let special_tokens = (!special_tokens.is_empty()).then_some(special_tokens);
91 +
92 + if let Some(configured_template) = options.chat_template.as_deref() {
93 + template = Some(
94 + resolve_chat_template(configured_template)
95 + .map_err(|error| Error::ChatTemplate(error.to_report_string()))?,
96 + );
97 + info!("using configured chat template override");
98 + } else if let Some(chat_template_path) = files.chat_template_path.as_deref() {
99 + // If independent chat template file(s) exist and contain non-empty content,
100 + // they take priority over template entries in the tokenizer config
101 + let file_template = load_chat_template(chat_template_path)
102 + .map_err(|error| Error::ChatTemplate(error.to_report_string()))?;
103 +
104 + if file_template.as_ref().is_some_and(|t| !t.trim().is_empty()) {
105 + info!(
106 + path = %chat_template_path.display(),
107 + "loaded dedicated chat template file, overriding tokenizer_config chat_template"
108 + );
109 + template = file_template;
110 + } else {
111 + warn!(
112 + path = %chat_template_path.display(),
113 + "ignoring empty dedicated chat template file and falling back to tokenizer_config chat_template"
114 + );
115 + }
116 + }
117 +
118 + Ok(Self::new(
119 + template,
120 + options.default_chat_template_kwargs,
121 + options.chat_template_content_format,
122 + )?
123 + .with_special_tokens(special_tokens)
124 + .with_multimodal(multimodal))
125 + }
126 +
127 + /// Apply the chat template to one chat request, rendering the prompt string
128 + /// to be tokenized and submitted to the model.
129 + ///
130 + /// If the request carries a per-request `chat_template` override, a
131 + /// temporary template is compiled from that string and used instead of
132 + /// the model's default.
133 + fn apply_chat_template(&self, request: &ChatRequest) -> Result<RenderedPrompt> {
134 + let override_template = request
135 + .chat_options
136 + .chat_template
137 + .as_ref()
138 + .map(|template| {
139 + CompiledChatTemplate::new(template.clone(), self.content_format)
140 + .map_err(|error| Error::ChatTemplate(error.to_report_string()))
141 + })
142 + .transpose()?;
143 + let template = override_template
144 + .as_ref()
145 + .or(self.default_template.as_ref())
146 + .ok_or(Error::MissingChatTemplate)?;
147 +
148 + self.apply_chat_template_inner(template, request)
149 + }
150 +
151 + fn apply_chat_template_inner(
152 + &self,
153 + effective_template: &CompiledChatTemplate,
154 + request: &ChatRequest,
155 + ) -> Result<RenderedPrompt> {
156 + let messages = to_template_messages(
157 + &request.messages,
158 + effective_template.content_format(),
159 + self.multimodal.as_ref(),
160 + )?;
161 + let tools = request.tool_parsing_enabled().then(|| to_template_tools(&request.tools));
162 + trace!(
163 + message_count = messages.len(),
164 + content_format = ?effective_template.content_format(),
165 + ?messages,
166 + ?tools,
167 + "applying chat template"
168 + );
169 +
170 + let mut merged_template_kwargs = self.default_template_kwargs.clone();
171 + merged_template_kwargs.extend(request.chat_options.template_kwargs.clone());
172 + let prompt = effective_template
173 + .apply(TemplateContext {
174 + messages: &messages,
175 + add_generation_prompt: request.chat_options.add_generation_prompt(),
176 + continue_final_message: request.chat_options.continue_final_message(),
177 + tools: tools.as_deref(),
178 + documents: request.documents.as_deref(),
179 + template_kwargs: Some(&merged_template_kwargs),
180 + special_tokens: self.special_tokens.as_ref(),
181 + reasoning_effort: request.chat_options.reasoning_effort,
182 + })
183 + .map_err(|error| Error::ChatTemplate(error.to_report_string()))?;
184 +
185 + trace!(
186 + prompt_len = prompt.len(),
187 + prompt, "rendered chat template prompt"
188 + );
189 +
190 + Ok(RenderedPrompt {
191 + prompt: Prompt::Text(prompt),
192 + })
193 + }
194 +}
195 +
196 +impl ChatRenderer for HfChatRenderer {
197 + fn render(&self, request: &ChatRequest) -> Result<RenderedPrompt> {
198 + self.apply_chat_template(request)
199 + }
200 +}
201 +
202 +/// Chat message in the JSON shape expected by Jinja chat templates.
203 +// TODO: borrow more fields directly from the original `ChatMessage`.
204 +#[serde_with::skip_serializing_none]
205 +#[derive(Debug, Serialize)]
206 +struct TemplateMessage {
207 + role: &'static str,
208 + content: TemplateContent,
209 + // Developer-role messages may provide message-local tools in the same shape
210 + // as top-level request tools.
211 + tools: Option<Vec<TemplateTool>>,
212 + // Reasoning-capable HF templates are inconsistent on the exact field name,
213 + // so expose both variants for compatibility.
214 + reasoning: Option<String>,
215 + reasoning_content: Option<String>,
216 + // Function-call-capable templates commonly expect assistant tool calls
217 + // under this OpenAI-compatible field name.
218 + tool_calls: Option<Vec<TemplateToolCall>>,
219 + // Tool-role messages refer back to the assistant call they are answering.
220 + tool_call_id: Option<String>,
221 +}
222 +
223 +/// Chat content in the two shapes HF templates commonly expect.
224 +#[derive(Debug, Serialize)]
225 +#[serde(untagged)]
226 +enum TemplateContent {
227 + String(String),
228 + OpenAi(Vec<TemplateContentPart>),
229 +}
230 +
231 +#[derive(Debug, Serialize)]
232 +#[serde(tag = "type", rename_all = "snake_case")]
233 +enum TemplateContentPart {
234 + Text { text: String },
235 + Image,
236 +}
237 +
238 +#[derive(Debug, Serialize)]
239 +struct TemplateToolCall {
240 + id: String,
241 + r#type: &'static str, // always "function"
242 + function: TemplateToolFunction,
243 +}
244 +
245 +#[derive(Debug, Serialize)]
246 +struct TemplateToolFunction {
247 + name: String,
248 + arguments: Value,
249 +}
250 +
251 +#[derive(Debug, Serialize)]
252 +pub(super) struct TemplateTool {
253 + #[serde(rename = "type")]
254 + tool_type: &'static str,
255 + function: TemplateToolDefinition,
256 +}
257 +
258 +#[derive(Debug, Serialize)]
259 +struct TemplateToolDefinition {
260 + name: String,
261 + description: Option<String>,
262 + parameters: Value,
263 + strict: Option<bool>,
264 +}
265 +
266 +/// Convert chat messages into the JSON shape expected by Jinja chat templates.
267 +fn to_template_messages(
268 + messages: &[ChatMessage],
269 + content_format: ChatTemplateContentFormat,
270 + multimodal: Option<&MultimodalRenderInfo>,
271 +) -> Result<Vec<TemplateMessage>> {
272 + messages
273 + .iter()
274 + .map(|message| to_template_message(message, content_format, multimodal))
275 + .collect()
276 +}
277 +
278 +fn to_template_message(
279 + message: &ChatMessage,
280 + content_format: ChatTemplateContentFormat,
281 + multimodal: Option<&MultimodalRenderInfo>,
282 +) -> Result<TemplateMessage> {
283 + Ok(match message {
284 + ChatMessage::System { content } => TemplateMessage {
285 + role: "system",
286 + content: to_template_content(content, content_format, multimodal)?,
287 + tools: None,
288 + reasoning: None,
289 + reasoning_content: None,
290 + tool_calls: None,
291 + tool_call_id: None,
292 + },
293 + ChatMessage::Developer { content, tools } => TemplateMessage {
294 + role: "developer",
295 + content: to_template_content(content, content_format, multimodal)?,
296 + tools: tools.as_deref().map(to_template_tools),
297 + reasoning: None,
298 + reasoning_content: None,
299 + tool_calls: None,
300 + tool_call_id: None,
301 + },
302 + ChatMessage::User { content } => TemplateMessage {
303 + role: "user",
304 + content: to_template_content(content, content_format, multimodal)?,
305 + tools: None,
306 + reasoning: None,
307 + reasoning_content: None,
308 + tool_calls: None,
309 + tool_call_id: None,
310 + },
311 + ChatMessage::Assistant { content } => {
312 + let text = content.text();
313 + let reasoning = content.reasoning();
314 + let tool_calls = to_template_tool_calls(content)?;
315 + let content =
316 + to_template_content(&ChatContent::Text(text), content_format, multimodal)?;
317 + TemplateMessage {
318 + role: "assistant",
319 + content,
320 + tools: None,
321 + reasoning: reasoning.clone(),
322 + reasoning_content: reasoning,
323 + tool_calls,
324 + tool_call_id: None,
325 + }
326 + }
327 + ChatMessage::ToolResponse {
328 + content,
329 + tool_call_id,
330 + } => TemplateMessage {
331 + role: "tool",
332 + content: to_template_content(content, content_format, multimodal)?,
333 + tools: None,
334 + reasoning: None,
335 + reasoning_content: None,
336 + tool_calls: None,
337 + tool_call_id: Some(tool_call_id.clone()),
338 + },
339 + })
340 +}
341 +
342 +fn to_template_tool_calls(
343 + content: &[AssistantContentBlock],
344 +) -> Result<Option<Vec<TemplateToolCall>>> {
345 + let mut tool_calls = Vec::new();
346 +
347 + for tool_call in content.tool_calls() {
348 + let arguments = serde_json::from_str::<Value>(&tool_call.arguments).map_err(|error| {
349 + Error::ChatTemplate(format!(
350 + "assistant tool call `{}` has invalid JSON arguments: {}",
351 + tool_call.id,
352 + error.as_report()
353 + ))
354 + })?;
355 +
356 + tool_calls.push(TemplateToolCall {
357 + id: tool_call.id.clone(),
358 + r#type: "function",
359 + function: TemplateToolFunction {
360 + name: tool_call.name.clone(),
361 + arguments,
362 + },
363 + });
364 + }
365 +
366 + Ok((!tool_calls.is_empty()).then_some(tool_calls))
367 +}
368 +
369 +fn to_template_content(
370 + content: &ChatContent,
371 + content_format: ChatTemplateContentFormat,
372 + multimodal: Option<&MultimodalRenderInfo>,
373 +) -> Result<TemplateContent> {
374 + Ok(match content_format {
375 + ChatTemplateContentFormat::String => {
376 + TemplateContent::String(to_template_string_content(content, multimodal)?)
377 + }
378 + ChatTemplateContentFormat::OpenAi => {
379 + TemplateContent::OpenAi(to_template_openai_content(content, multimodal)?)
380 + }
381 + })
382 +}
383 +
384 +fn to_template_openai_content(
385 + content: &ChatContent,
386 + multimodal: Option<&MultimodalRenderInfo>,
387 +) -> Result<Vec<TemplateContentPart>> {
388 + match content {
389 + ChatContent::Text(text) => Ok(vec![TemplateContentPart::Text { text: text.clone() }]),
390 + ChatContent::Parts(parts) => parts
391 + .iter()
392 + .map(|part| match part {
393 + ChatContentPart::Text { text } => {
394 + Ok(TemplateContentPart::Text { text: text.clone() })
395 + }
396 + // All multimodal contents are normalized to `{ "type": <modality> }`.
397 + ChatContentPart::ImageUrl { .. } => {
398 + multimodal.ok_or(Error::UnsupportedMultimodalContent("image_url"))?;
399 + Ok(TemplateContentPart::Image)
400 + }
401 + })
402 + .collect(),
403 + }
404 +}
405 +
406 +fn to_template_string_content(
407 + content: &ChatContent,
408 + multimodal: Option<&MultimodalRenderInfo>,
409 +) -> Result<String> {
410 + match content {
411 + ChatContent::Text(text) => Ok(text.clone()),
412 + ChatContent::Parts(parts) => {
413 + let mut out = String::new();
414 + for part in parts {
415 + match part {
416 + ChatContentPart::Text { text } => out.push_str(text),
417 + ChatContentPart::ImageUrl { .. } => {
418 + let multimodal =
419 + multimodal.ok_or(Error::UnsupportedMultimodalContent("image_url"))?;
420 + out.push_str(&multimodal.placeholder_token);
421 + }
422 + }
423 + }
424 + Ok(out)
425 + }
426 + }
427 +}
428 +
429 +fn to_template_tools(tools: &[ChatTool]) -> Vec<TemplateTool> {
430 + tools
431 + .iter()
432 + .map(|tool| TemplateTool {
433 + tool_type: "function",
434 + function: TemplateToolDefinition {
435 + name: tool.name.clone(),
436 + description: tool.description.clone(),
437 + parameters: tool.parameters.clone(),
438 + strict: tool.strict,
439 + },
440 + })
441 + .collect()
442 +}
443 +
444 +#[cfg(test)]
445 +mod tests {
446 + use std::collections::HashMap;
447 +
448 + use expect_test::expect;
449 + use serde_json::Value;
450 + use vllm_text::Prompt;
451 + use vllm_text::backend::hf::{HfSpecialTokens, NamedSpecialToken};
452 +
453 + use super::{ChatTemplateContentFormatOption, HfChatRenderer, MultimodalRenderInfo};
454 + use crate::request::{
455 + ChatContentPart, ChatMessage, ChatRequest, ChatRole, ChatTool, ChatToolChoice,
456 + GenerationPromptMode, ReasoningEffort,
457 + };
458 + use crate::{AssistantContentBlock, ChatRenderer, Error, Result};
459 +
460 + const QWEN3_0_6B_TEMPLATE: &str = include_str!("../../../tests/templates/qwen3.jinja");
461 + const QWEN3_5_0_8B_TEMPLATE: &str = include_str!("../../../tests/templates/qwen35.jinja");
462 +
463 + fn sample_request(messages: Vec<ChatMessage>) -> ChatRequest {
464 + ChatRequest {
465 + messages,
466 + request_id: "render-test".to_string(),
467 + ..ChatRequest::for_test()
468 + }
469 + }
470 +
471 + fn render(template: Option<&str>, request: &ChatRequest) -> Result<String> {
472 + HfChatRenderer::new(
473 + template.map(str::to_owned),
474 + HashMap::new(),
475 + ChatTemplateContentFormatOption::Auto,
476 + )?
477 + .render(request)?
478 + .prompt
479 + .into_text()
480 + .map_err(|_| unreachable!("HF renderer should return text prompt"))
481 + }
482 +
483 + fn render_mm(
484 + template: &str,
485 + request: &ChatRequest,
486 + content_format: ChatTemplateContentFormatOption,
487 + ) -> Result<crate::RenderedPrompt> {
488 + HfChatRenderer::new(Some(template.to_string()), HashMap::new(), content_format)?
489 + .with_multimodal(Some(MultimodalRenderInfo {
490 + placeholder_token: "<image>".to_string(),
491 + }))
492 + .render(request)
493 + }
494 +
495 + fn image_request() -> ChatRequest {
496 + sample_request(vec![ChatMessage::user(vec![
497 + ChatContentPart::text("a"),
498 + ChatContentPart::image_url("data:image/png;base64,test"),
499 + ChatContentPart::text("b"),
500 + ])])
501 + }
502 +
503 + #[test]
504 + fn string_content_format_replaces_image_with_placeholder_text() {
505 + let rendered = render_mm(
506 + "{{ messages[0].content }}",
507 + &image_request(),
508 + ChatTemplateContentFormatOption::String,
509 + )
510 + .unwrap();
511 +
512 + assert_eq!(rendered.prompt, Prompt::Text("a<image>b".to_string()));
513 + }
514 +
515 + #[test]
516 + fn openai_content_format_normalizes_image_url_for_template() {
517 + let rendered = render_mm(
518 + "{% for item in messages[0].content %}{% if item.type == 'image' %}<|image_pad|>{% else %}{{ item.text }}{% endif %}{% endfor %}",
519 + &image_request(),
520 + ChatTemplateContentFormatOption::OpenAi,
521 + )
522 + .unwrap();
523 +
524 + assert_eq!(rendered.prompt, Prompt::Text("a<|image_pad|>b".to_string()));
525 + }
526 +
527 + #[test]
528 + fn chat_template_supports_pycompat_templates() {
529 + let request = sample_request(vec![ChatMessage::text(ChatRole::User, "<think>hello")]);
530 +
531 + let rendered = render(
532 + Some(
533 + "{% for message in messages %}{% if message.content.startswith('<think>') %}think{% else %}plain{% endif %}{% endfor %}",
534 + ),
535 + &request,
536 + )
537 + .unwrap();
538 +
539 + assert_eq!(rendered, "think");
540 + }
541 +
542 + #[test]
543 + fn chat_template_passes_continue_final_message_to_template() {
544 + let mut request = sample_request(vec![ChatMessage::text(
545 + ChatRole::Assistant,
546 + "The capital of",
547 + )]);
548 +
549 + assert_eq!(
550 + render(
551 + Some("{% if continue_final_message %}continue{% else %}new{% endif %}"),
552 + &request,
553 + )
554 + .unwrap(),
555 + "new"
556 + );
557 +
558 + request.chat_options.generation_prompt_mode = GenerationPromptMode::ContinueFinalAssistant;
559 +
560 + assert_eq!(
561 + render(
562 + Some("{% if continue_final_message %}continue{% else %}new{% endif %}"),
563 + &request,
564 + )
565 + .unwrap(),
566 + "continue"
567 + );
568 + }
569 +
570 + #[test]
571 + fn chat_template_flattens_text_parts_for_string_templates() {
572 + let request = sample_request(vec![ChatMessage::user(vec![
573 + ChatContentPart::text("hello"),
574 + ChatContentPart::text(" world"),
575 + ])]);
576 +
577 + let rendered = render(Some("{{ messages[0].content }}"), &request).unwrap();
578 +
579 + assert_eq!(rendered, "hello world");
580 + }
581 +
582 + #[test]
583 + fn chat_template_exposes_developer_tools() {
584 + let request = sample_request(vec![ChatMessage::developer(
585 + "policy",
586 + Some(vec![ChatTool {
587 + name: "get_weather".to_string(),
588 + description: Some("Get weather".to_string()),
589 + parameters: serde_json::json!({
590 + "type": "object",
591 + "properties": {"city": {"type": "string"}},
592 + "required": ["city"],
593 + }),
594 + strict: Some(true),
595 + }]),
596 + )]);
597 +
598 + let rendered = render(
599 + Some("{{ messages[0].role }}|{{ messages[0].content }}|{{ messages[0].tools[0].function.name }}|{{ messages[0].tools[0].function.parameters.required[0] }}"),
600 + &request,
601 + )
602 + .unwrap();
603 +
604 + assert_eq!(rendered, "developer|policy|get_weather|city");
605 + }
606 +
607 + #[test]
608 + fn chat_template_keeps_string_text_for_openai_detected_templates() {
609 + let request = sample_request(vec![ChatMessage::text(ChatRole::User, "hello")]);
610 +
611 + let rendered = render(
612 + Some(
613 + "{%- for message in messages %}{%- if message.content is string %}{%- set content = message.content %}{{ content }}{%- endif %}{%- endfor %}",
614 + ),
615 + &request,
616 + )
617 + .unwrap();
618 +
619 + assert_eq!(rendered, "hello");
620 + }
621 +
622 + #[test]
623 + fn chat_template_emits_openai_text_blocks_for_structured_templates() {
624 + let request = sample_request(vec![ChatMessage::user(vec![
625 + ChatContentPart::text("hello"),
626 + ChatContentPart::text("world"),
627 + ])]);
628 +
629 + let rendered = render(
630 + Some(
631 + "{%- for message in messages %}{%- for item in message.content %}{{ item.text }}|{%- endfor %}{%- endfor %}",
632 + ),
633 + &request,
634 + )
635 + .unwrap();
636 +
637 + assert_eq!(rendered, "hello|world|");
638 + }
639 +
640 + #[test]
641 + fn chat_template_per_request_override() {
642 + let mut request = sample_request(vec![ChatMessage::text(ChatRole::User, "hello")]);
643 +
644 + // Default template renders one way.
645 + let default_rendered = render(Some("{{ messages[0].content }}"), &request).unwrap();
646 + assert_eq!(default_rendered, "hello");
647 +
648 + // Per-request override replaces the default template entirely.
649 + request.chat_options.chat_template = Some("override:{{ messages[0].content }}".to_string());
650 + let overridden = render(Some("{{ messages[0].content }}"), &request).unwrap();
651 + assert_eq!(overridden, "override:hello");
652 + }
653 +
654 + #[test]
655 + fn chat_template_per_request_override_without_default_template() {
656 + let mut request = sample_request(vec![ChatMessage::text(ChatRole::User, "hello")]);
657 + request.chat_options.chat_template = Some("override:{{ messages[0].content }}".to_string());
658 +
659 + let rendered = render(None, &request).unwrap();
660 +
661 + assert_eq!(rendered, "override:hello");
662 + }
663 +
664 + #[test]
665 + fn chat_template_requires_a_template() {
666 + let request = sample_request(vec![ChatMessage::text(ChatRole::User, "hello")]);
667 + let error = render(None, &request).unwrap_err();
668 +
669 + assert!(matches!(error, Error::MissingChatTemplate));
670 + }
671 +
672 + #[test]
673 + fn chat_template_injects_special_tokens_into_context() {
674 + let request = sample_request(vec![ChatMessage::text(ChatRole::User, "hello")]);
675 + let special_tokens = HfSpecialTokens {
676 + bos_token: Some(NamedSpecialToken::Text("<bos>".to_string())),
677 + ..Default::default()
678 + };
679 +
680 + let rendered = HfChatRenderer::new(
681 + Some("{{ bos_token }}|{{ bos_token is defined }}".to_string()),
682 + HashMap::new(),
683 + ChatTemplateContentFormatOption::Auto,
684 + )
685 + .unwrap()
686 + .with_special_tokens(Some(special_tokens))
687 + .apply_chat_template(&request)
688 + .unwrap();
689 +
690 + assert_eq!(rendered.prompt, Prompt::Text("<bos>|true".to_string()));
691 + }
692 +
693 + #[test]
694 + fn chat_template_exposes_assistant_reasoning_separately() {
695 + let request = sample_request(vec![ChatMessage::assistant_blocks(vec![
696 + AssistantContentBlock::Reasoning {
697 + text: "inner".to_string(),
698 + },
699 + AssistantContentBlock::Text {
700 + text: "outer".to_string(),
701 + },
702 + ])]);
703 +
704 + let rendered = render(
705 + Some("{{ messages[0].reasoning_content }}|{{ messages[0].content }}"),
706 + &request,
707 + )
708 + .unwrap();
709 +
710 + assert_eq!(rendered, "inner|outer");
711 + }
712 +
713 + #[test]
714 + fn chat_template_forces_string_content_format_when_configured() {
715 + let request = sample_request(vec![ChatMessage::user(vec![
716 + ChatContentPart::text("hello"),
717 + ChatContentPart::text(" world"),
718 + ])]);
719 +
720 + let rendered = HfChatRenderer::new(
721 + Some(
722 + "{%- if messages[0].content is string -%}{{ messages[0].content }}{%- else -%}{%- for item in messages[0].content %}{{ item.text }}|{%- endfor -%}{%- endif -%}".to_string(),
723 + ),
724 + HashMap::new(),
725 + ChatTemplateContentFormatOption::String,
726 + )
727 + .unwrap()
728 + .render(&request)
729 + .unwrap()
730 + .prompt;
731 +
732 + assert_eq!(rendered, Prompt::Text("hello world".to_string()));
733 + }
734 +
735 + #[test]
736 + fn chat_template_forces_openai_content_format_when_configured() {
737 + let request = sample_request(vec![ChatMessage::user(vec![
738 + ChatContentPart::text("hello"),
739 + ChatContentPart::text(" world"),
740 + ])]);
741 +
742 + let rendered = HfChatRenderer::new(
743 + Some("{{ messages[0].content[0].text }}{{ messages[0].content[1].text }}".to_string()),
744 + HashMap::new(),
745 + ChatTemplateContentFormatOption::OpenAi,
746 + )
747 + .unwrap()
748 + .render(&request)
749 + .unwrap()
750 + .prompt;
751 +
752 + assert_eq!(rendered, Prompt::Text("hello world".to_string()));
753 + }
754 +
755 + #[test]
756 + fn chat_template_merges_default_template_kwargs_before_request_kwargs() {
757 + let mut request = sample_request(vec![ChatMessage::text(ChatRole::User, "hello")]);
758 + request
759 + .chat_options
760 + .template_kwargs
761 + .insert("enable_thinking".to_string(), Value::Bool(true));
762 +
763 + let renderer = HfChatRenderer::new(
764 + Some("{{ enable_thinking }}|{{ default_only }}".to_string()),
765 + HashMap::from([
766 + ("enable_thinking".to_string(), Value::Bool(false)),
767 + ("default_only".to_string(), Value::String("x".to_string())),
768 + ]),
769 + ChatTemplateContentFormatOption::Auto,
770 + )
771 + .unwrap();
772 +
773 + let rendered = renderer.render(&request).unwrap().prompt;
774 +
775 + assert_eq!(rendered, Prompt::Text("true|x".to_string()));
776 + }
777 +
778 + #[test]
779 + fn chat_template_reasoning_effort_overrides_template_kwargs() {
780 + let mut request = sample_request(vec![ChatMessage::text(ChatRole::User, "hello")]);
781 + request.chat_options.reasoning_effort = Some(ReasoningEffort::Max);
782 + request.chat_options.template_kwargs.insert(
783 + "reasoning_effort".to_string(),
784 + Value::String("low".to_string()),
785 + );
786 +
787 + let renderer = HfChatRenderer::new(
788 + Some("{{ reasoning_effort }}".to_string()),
789 + HashMap::from([(
790 + "reasoning_effort".to_string(),
791 + Value::String("medium".to_string()),
792 + )]),
793 + ChatTemplateContentFormatOption::Auto,
794 + )
795 + .unwrap();
796 +
797 + let rendered = renderer.render(&request).unwrap().prompt;
798 +
799 + assert_eq!(rendered, Prompt::Text("max".to_string()));
800 + }
801 +
802 + #[test]
803 + fn qwen3_template_omits_reasoning_for_historical_assistant_messages() {
804 + let request = sample_request(vec![
805 + ChatMessage::text(
806 + ChatRole::User,
807 + "Hi. Tell me about the capital of France in short",
808 + ),
809 + ChatMessage::assistant_blocks(vec![
810 + AssistantContentBlock::Reasoning {
811 + text: "\nOkay, the user is asking... I think that's all.\n".to_string(),
812 + },
813 + AssistantContentBlock::Text {
814 + text: "Paris is the capital of France.".to_string(),
815 + },
816 + ]),
817 + ChatMessage::text(ChatRole::User, "Tell me about Paris more."),
818 + ]);
819 +
820 + let rendered = render(Some(QWEN3_0_6B_TEMPLATE), &request).unwrap();
821 +
822 + expect![[r#"
823 + <|im_start|>user
824 + Hi. Tell me about the capital of France in short<|im_end|>
825 + <|im_start|>assistant
826 + Paris is the capital of France.<|im_end|>
827 + <|im_start|>user
828 + Tell me about Paris more.<|im_end|>
829 + <|im_start|>assistant
830 + "#]]
831 + .assert_eq(&rendered);
832 + }
833 +
834 + #[test]
835 + fn qwen3_template_keeps_reasoning_after_the_last_user_query() {
836 + let mut request = sample_request(vec![
837 + ChatMessage::text(ChatRole::User, "What is 1 + 1?"),
838 + ChatMessage::assistant_blocks(vec![
839 + AssistantContentBlock::Reasoning {
840 + text: "need simple arithmetic".to_string(),
841 + },
842 + AssistantContentBlock::Text {
843 + text: "2".to_string(),
844 + },
845 + ]),
846 + ]);
847 + request.chat_options.generation_prompt_mode = GenerationPromptMode::NoGenerationPrompt;
848 +
849 + let rendered = render(Some(QWEN3_0_6B_TEMPLATE), &request).unwrap();
850 +
851 + expect![[r#"
852 + <|im_start|>user
853 + What is 1 + 1?<|im_end|>
854 + <|im_start|>assistant
855 + <think>
856 + need simple arithmetic
857 + </think>
858 +
859 + 2<|im_end|>
860 + "#]]
861 + .assert_eq(&rendered);
862 + }
863 +
864 + #[test]
865 + fn chat_template_exposes_tools_to_templates_when_auto_enabled() {
866 + let mut request = sample_request(vec![ChatMessage::text(ChatRole::User, "hello")]);
867 + request.tools = vec![ChatTool {
868 + name: "get_weather".to_string(),
869 + description: Some("Get weather".to_string()),
870 + parameters: serde_json::json!({
871 + "type": "object",
872 + "properties": {"city": {"type": "string"}},
873 + "required": ["city"],
874 + }),
875 + strict: None,
876 + }];
877 + request.tool_choice = ChatToolChoice::Auto;
878 +
879 + let rendered = render(
880 + Some("{{ tools[0].function.name }}|{{ tools[0].function.parameters.required[0] }}"),
881 + &request,
882 + )
883 + .unwrap();
884 +
885 + assert_eq!(rendered, "get_weather|city");
886 + }
887 +
888 + #[test]
889 + fn chat_template_exposes_assistant_tool_calls_and_tool_messages() {
890 + let request = sample_request(vec![
891 + ChatMessage::assistant_blocks(vec![AssistantContentBlock::ToolCall(
892 + crate::AssistantToolCall {
893 + id: "call_1".to_string(),
894 + name: "get_weather".to_string(),
895 + arguments: r#"{"city":"Paris"}"#.to_string(),
896 + },
897 + )]),
898 + ChatMessage::tool_response("Sunny", "call_1"),
899 + ]);
900 +
901 + let rendered = render(
902 + Some(
903 + "{{ messages[0].tool_calls[0].function.name }}|{{ messages[0].tool_calls[0].function.arguments.city }}|{{ messages[1].tool_call_id }}|{{ messages[1].content }}",
904 + ),
905 + &request,
906 + )
907 + .unwrap();
908 +
909 + assert_eq!(rendered, "get_weather|Paris|call_1|Sunny");
910 + }
911 +
912 + #[test]
913 + fn qwen35_template_renders_prefilled_reasoning_start_when_thinking_enabled() {
914 + let mut request = sample_request(vec![ChatMessage::text(ChatRole::User, "hello")]);
915 + request
916 + .chat_options
917 + .template_kwargs
918 + .insert("enable_thinking".to_string(), Value::Bool(true));
919 +
920 + let rendered = render(Some(QWEN3_5_0_8B_TEMPLATE), &request).unwrap();
921 +
922 + expect![[r#"
923 + <|im_start|>user
924 + hello<|im_end|>
925 + <|im_start|>assistant
926 + <think>
927 + "#]]
928 + .assert_eq(&rendered);
929 + }
930 +
931 + #[test]
932 + fn qwen35_template_renders_closed_empty_reasoning_span_when_thinking_disabled() {
933 + let mut request = sample_request(vec![ChatMessage::text(ChatRole::User, "hello")]);
934 + request
935 + .chat_options
936 + .template_kwargs
937 + .insert("enable_thinking".to_string(), Value::Bool(false));
938 +
939 + let rendered = render(Some(QWEN3_5_0_8B_TEMPLATE), &request).unwrap();
940 +
941 + expect![[r#"
942 + <|im_start|>user
943 + hello<|im_end|>
944 + <|im_start|>assistant
945 + <think>
946 +
947 + </think>
948 +
949 + "#]]
950 + .assert_eq(&rendered);
951 + }
952 +
953 + #[test]
954 + fn qwen35_template_omits_assistant_reasoning_prefill_without_generation_prompt() {
955 + let mut request = sample_request(vec![ChatMessage::text(ChatRole::User, "hello")]);
956 + request.chat_options.generation_prompt_mode = GenerationPromptMode::NoGenerationPrompt;
957 + request
958 + .chat_options
959 + .template_kwargs
960 + .insert("enable_thinking".to_string(), Value::Bool(true));
961 +
962 + let rendered = render(Some(QWEN3_5_0_8B_TEMPLATE), &request).unwrap();
963 +
964 + expect![[r#"
965 + <|im_start|>user
966 + hello<|im_end|>
967 + "#]]
968 + .assert_eq(&rendered);
969 + }
970 +}
rust/src/chat/src/renderer/hf/template.rs new
+316
@@ -0,0 +1,316 @@
1 +//! Chat template support for tokenizers using Jinja2 templates.
2 +//!
3 +//! This module is inlined from SMG's tokenizer crate with local adaptations:
4 +//! - thinking-related detection/state is removed
5 +//! - special tokens are wired to `vllm_text::backends::hf::HfSpecialTokens`
6 +
7 +use std::collections::HashMap;
8 +use std::fs;
9 +use std::path::Path;
10 +
11 +use minijinja::Environment;
12 +use serde::{Deserialize, Serialize};
13 +use serde_json::{self};
14 +use vllm_text::backend::hf::HfSpecialTokens;
15 +
16 +use super::error::TemplateError;
17 +use super::format::{
18 + ChatTemplateContentFormat, ChatTemplateContentFormatOption, detect_chat_template_content_format,
19 +};
20 +use super::tojson::hf_tojson_filter;
21 +use crate::renderer::hf::{TemplateMessage, TemplateTool};
22 +use crate::request::ReasoningEffort;
23 +
24 +type Result<T> = std::result::Result<T, TemplateError>;
25 +
26 +/// Build a pre-configured environment with the given template string.
27 +fn build_environment(template: String) -> Result<Environment<'static>> {
28 + let mut env = Environment::new();
29 +
30 + env.set_trim_blocks(true);
31 + env.set_lstrip_blocks(true);
32 +
33 + env.add_template_owned("chat".to_owned(), template)?;
34 +
35 + env.set_unknown_method_callback(minijinja_contrib::pycompat::unknown_method_callback);
36 + env.add_filter("tojson", hf_tojson_filter);
37 +
38 + Ok(env)
39 +}
40 +
41 +#[serde_with::skip_serializing_none]
42 +#[derive(Default, Serialize)]
43 +pub(super) struct TemplateContext<'a> {
44 + pub(super) messages: &'a [TemplateMessage],
45 + pub(super) add_generation_prompt: bool,
46 + pub(super) continue_final_message: bool,
47 + pub(super) tools: Option<&'a [TemplateTool]>,
48 + pub(super) documents: Option<&'a [serde_json::Value]>,
49 + #[serde(flatten)]
50 + pub(super) special_tokens: Option<&'a HfSpecialTokens>,
51 + #[serde(flatten)]
52 + pub(super) template_kwargs: Option<&'a HashMap<String, serde_json::Value>>,
53 + // By putting top-level `reasoning_effort` after `template_kwargs`, this overrides any
54 + // `reasoning_effort` value that might be present there.
55 + pub(super) reasoning_effort: Option<ReasoningEffort>,
56 +}
57 +
58 +/// Load chat template from a file (`.jinja` or `.json` containing Jinja).
59 +pub fn load_chat_template(template_path: &Path) -> Result<Option<String>> {
60 + let content = fs::read_to_string(template_path).map_err(TemplateError::ReadTemplateFile)?;
61 +
62 + if template_path.extension().is_some_and(|ext| ext == "json") {
63 + #[derive(Deserialize)]
64 + #[serde(untagged)]
65 + enum ChatTemplateFile {
66 + String(String),
67 + Object { chat_template: String },
68 + }
69 +
70 + let json_value =
71 + serde_json::from_str(&content).map_err(TemplateError::ParseTemplateJson)?;
72 + let json_template =
73 + serde_json::from_value(json_value).map_err(|_| TemplateError::InvalidTemplateJson)?;
74 +
75 + return Ok(Some(match json_template {
76 + ChatTemplateFile::String(template) => template,
77 + ChatTemplateFile::Object { chat_template } => chat_template,
78 + }));
79 + }
80 +
81 + let template = content.trim().replace("\\n", "\n");
82 + Ok(Some(template))
83 +}
84 +
85 +/// Resolve a configured chat template value into a template string.
86 +pub fn resolve_chat_template(chat_template: &str) -> Result<String> {
87 + let path = Path::new(chat_template);
88 + if path.exists() {
89 + return load_chat_template(path).map(|template| template.unwrap_or_default());
90 + }
91 +
92 + const JINJA_CHARS: [char; 3] = ['{', '}', '\n'];
93 + if chat_template.chars().any(|c| JINJA_CHARS.contains(&c)) {
94 + return Ok(chat_template.to_string());
95 + }
96 +
97 + Err(TemplateError::MissingTemplatePath)
98 +}
99 +
100 +/// One compiled chat template with its Jinja environment and detected content
101 +/// format.
102 +pub(super) struct CompiledChatTemplate {
103 + /// Cached, fully-configured environment for one compiled template.
104 + env: Environment<'static>,
105 + content_format: ChatTemplateContentFormat,
106 +}
107 +
108 +impl CompiledChatTemplate {
109 + /// Compile the given chat template string into a [`CompiledChatTemplate`].
110 + pub fn new(template: String, content_format: ChatTemplateContentFormatOption) -> Result<Self> {
111 + let content_format = match content_format {
112 + ChatTemplateContentFormatOption::Auto => detect_chat_template_content_format(&template),
113 + ChatTemplateContentFormatOption::String => ChatTemplateContentFormat::String,
114 + ChatTemplateContentFormatOption::OpenAi => ChatTemplateContentFormat::OpenAi,
115 + };
116 + let env = build_environment(template)?;
117 + Ok(Self {
118 + env,
119 + content_format,
120 + })
121 + }
122 +
123 + /// Apply the compiled template to the given context and return the rendered
124 + /// prompt.
125 + pub fn apply(&self, ctx: TemplateContext<'_>) -> Result<String> {
126 + let tmpl = self.env.get_template("chat")?;
127 + tmpl.render(ctx).map_err(TemplateError::from)
128 + }
129 +
130 + pub fn content_format(&self) -> ChatTemplateContentFormat {
131 + self.content_format
132 + }
133 +}
134 +
135 +#[cfg(test)]
136 +mod tests {
137 + use std::fs;
138 +
139 + use tempfile::TempDir;
140 + use vllm_text::backend::hf::{HfSpecialTokens, NamedSpecialToken};
141 +
142 + use super::*;
143 +
144 + #[test]
145 + fn test_chat_template_state_valid_template() {
146 + let template = CompiledChatTemplate::new(
147 + "{{ messages }}".to_string(),
148 + ChatTemplateContentFormatOption::Auto,
149 + )
150 + .unwrap();
151 + assert_eq!(template.content_format(), ChatTemplateContentFormat::String);
152 + let result = template.apply(TemplateContext::default()).unwrap();
153 + assert_eq!(result, "[]");
154 + }
155 +
156 + #[test]
157 + fn test_chat_template_state_invalid_template() {
158 + let result = CompiledChatTemplate::new(
159 + "{% invalid".to_string(),
160 + ChatTemplateContentFormatOption::Auto,
161 + );
162 + assert!(result.is_err());
163 + let err = result.err().unwrap().to_string();
164 + assert!(
165 + err.contains("failed to render jinja template"),
166 + "Error should explain parse failure, got: {err}"
167 + );
168 + }
169 +
170 + #[test]
171 + fn test_special_tokens_injected_into_context() {
172 + let template = "{{ bos_token }}hello{{ eos_token }}";
173 + let template =
174 + CompiledChatTemplate::new(template.to_string(), ChatTemplateContentFormatOption::Auto)
175 + .unwrap();
176 +
177 + let special_tokens = HfSpecialTokens {
178 + bos_token: Some(NamedSpecialToken::Text("<s>".to_string())),
179 + eos_token: Some(NamedSpecialToken::Text("</s>".to_string())),
180 + ..Default::default()
181 + };
182 +
183 + let result = template
184 + .apply(TemplateContext {
185 + special_tokens: Some(&special_tokens),
186 + ..Default::default()
187 + })
188 + .unwrap();
189 +
190 + assert_eq!(result, "<s>hello</s>");
191 + }
192 +
193 + #[test]
194 + fn test_special_tokens_undefined_when_not_provided() {
195 + let template = "{% if bos_token is defined %}{{ bos_token }}{% endif %}hello";
196 + let template =
197 + CompiledChatTemplate::new(template.to_string(), ChatTemplateContentFormatOption::Auto)
198 + .unwrap();
199 +
200 + let result = template.apply(TemplateContext::default()).unwrap();
201 + assert_eq!(result, "hello");
202 + }
203 +
204 + #[test]
205 + fn test_special_tokens_partial() {
206 + let template =
207 + "{{ bos_token }}hello{% if eos_token is defined %}{{ eos_token }}{% endif %}";
208 + let template =
209 + CompiledChatTemplate::new(template.to_string(), ChatTemplateContentFormatOption::Auto)
210 + .unwrap();
211 +
212 + let special_tokens = HfSpecialTokens {
213 + bos_token: Some(NamedSpecialToken::Text("<s>".to_string())),
214 + eos_token: None,
215 + ..Default::default()
216 + };
217 +
218 + let result = template
219 + .apply(TemplateContext {
220 + special_tokens: Some(&special_tokens),
221 + ..Default::default()
222 + })
223 + .unwrap();
224 +
225 + assert_eq!(result, "<s>hello");
226 + }
227 +
228 + #[test]
229 + fn test_tojson_filter_supports_indent_and_sort_keys() {
230 + let template = CompiledChatTemplate::new(
231 + "{{ payload | tojson(indent=2, sort_keys=true) }}".to_string(),
232 + ChatTemplateContentFormatOption::Auto,
233 + )
234 + .unwrap();
235 + let mut kwargs = HashMap::new();
236 + kwargs.insert("payload".to_string(), serde_json::json!({"b": 1, "a": 2}));
237 +
238 + let result = template
239 + .apply(TemplateContext {
240 + template_kwargs: Some(&kwargs),
241 + ..Default::default()
242 + })
243 + .unwrap();
244 +
245 + assert_eq!(result, "{\n \"a\": 2,\n \"b\": 1\n}");
246 + }
247 +
248 + #[test]
249 + fn test_load_chat_template_from_file_jinja() {
250 + let dir = TempDir::new().unwrap();
251 + let path = dir.path().join("chat_template.jinja");
252 + fs::write(&path, "{{ messages }}").unwrap();
253 +
254 + let template = load_chat_template(&path).unwrap();
255 +
256 + assert_eq!(template.as_deref(), Some("{{ messages }}"));
257 + }
258 +
259 + #[test]
260 + fn test_resolve_chat_template_from_inline_literal() {
261 + let template = resolve_chat_template("{{ messages }}").unwrap();
262 +
263 + assert_eq!(template, "{{ messages }}");
264 + }
265 +
266 + #[test]
267 + fn test_resolve_chat_template_from_existing_file() {
268 + let dir = TempDir::new().unwrap();
269 + let path = dir.path().join("chat_template.jinja");
270 + fs::write(&path, "{{ messages }}").unwrap();
271 +
272 + let template = resolve_chat_template(path.to_str().unwrap()).unwrap();
273 +
274 + assert_eq!(template, "{{ messages }}");
275 + }
276 +
277 + #[test]
278 + fn test_resolve_chat_template_rejects_missing_path_like_value() {
279 + let error = resolve_chat_template("missing_template.jinja").unwrap_err();
280 +
281 + assert!(matches!(error, TemplateError::MissingTemplatePath));
282 + }
283 +
284 + #[test]
285 + fn test_chat_template_state_respects_explicit_content_format_override() {
286 + let template = CompiledChatTemplate::new(
287 + "{% for item in messages[0].content %}{{ item.text }}{% endfor %}".to_string(),
288 + ChatTemplateContentFormatOption::String,
289 + )
290 + .unwrap();
291 +
292 + assert_eq!(template.content_format(), ChatTemplateContentFormat::String);
293 + }
294 +
295 + #[test]
296 + fn test_load_chat_template_from_file_json_string() {
297 + let dir = TempDir::new().unwrap();
298 + let path = dir.path().join("chat_template.json");
299 + fs::write(&path, "\"{{ messages }}\"").unwrap();
300 +
301 + let template = load_chat_template(&path).unwrap();
302 +
303 + assert_eq!(template.as_deref(), Some("{{ messages }}"));
304 + }
305 +
306 + #[test]
307 + fn test_load_chat_template_from_file_json_object() {
308 + let dir = TempDir::new().unwrap();
309 + let path = dir.path().join("chat_template.json");
310 + fs::write(&path, r#"{"chat_template":"{{ messages }}"}"#).unwrap();
311 +
312 + let template = load_chat_template(&path).unwrap();
313 +
314 + assert_eq!(template.as_deref(), Some("{{ messages }}"));
315 + }
316 +}
rust/src/chat/src/renderer/hf/tojson.rs new
+277
@@ -0,0 +1,277 @@
1 +use minijinja::value::{Kwargs, ViaDeserialize};
2 +use minijinja::{Error as MinijinjaError, ErrorKind, Value};
3 +use serde::Deserialize;
4 +use serde_json::{self, Value as JsonValue};
5 +use serde_json_fmt::{JsonFormat, JsonSyntaxError};
6 +use thiserror_ext::AsReport;
7 +
8 +/// Hugging Face-compatible `tojson` filter for chat templates.
9 +///
10 +/// We cannot use MiniJinja's built-in filter directly because HF relies on
11 +/// Python `json.dumps` semantics:
12 +/// - no HTML escaping
13 +/// - extra kwargs such as `ensure_ascii`, `separators`, and `sort_keys`
14 +/// - Python-style `indent` handling
15 +pub(super) fn hf_tojson_filter(
16 + value: Value,
17 + kwargs: Kwargs,
18 +) -> std::result::Result<Value, MinijinjaError> {
19 + let ensure_ascii = kwargs.get::<Option<bool>>("ensure_ascii")?.unwrap_or(false);
20 + let indent = parse_indent(
21 + kwargs.get::<Option<ViaDeserialize<IndentArg>>>("indent")?.map(|value| value.0),
22 + );
23 + let separators = parse_separators(
24 + kwargs
25 + .get::<Option<ViaDeserialize<SeparatorsArg>>>("separators")?
26 + .map(|value| value.0),
27 + indent.is_some(),
28 + );
29 + let sort_keys = kwargs.get::<Option<bool>>("sort_keys")?.unwrap_or(false);
30 +
31 + kwargs.assert_all_used()?;
32 +
33 + let json_value: serde_json::Value = serde_json::to_value(&value).map_err(|e| {
34 + MinijinjaError::new(
35 + ErrorKind::InvalidOperation,
36 + format!("Failed to convert to JSON value: {e}"),
37 + )
38 + })?;
39 +
40 + let json_str = {
41 + let value_to_serialize = if sort_keys {
42 + &sort_json_keys(&json_value)
43 + } else {
44 + &json_value
45 + };
46 +
47 + build_json_format(indent, separators.0, separators.1, ensure_ascii)?
48 + .format_to_string(value_to_serialize)
49 + .map_err(|e| {
50 + MinijinjaError::new(
51 + ErrorKind::InvalidOperation,
52 + format!("Failed to serialize JSON: {}", e.as_report()),
53 + )
54 + })?
55 + };
56 +
57 + Ok(Value::from_safe_string(json_str))
58 +}
59 +
60 +#[derive(Deserialize)]
61 +#[serde(untagged)]
62 +enum IndentArg {
63 + // Python `json.dumps` accepts bool, int, and string indentation styles.
64 + Bool(bool),
65 + Integer(i64),
66 + String(String),
67 +}
68 +
69 +fn parse_indent(value: Option<IndentArg>) -> Option<String> {
70 + match value? {
71 + IndentArg::Bool(indent) => Some(if indent {
72 + " ".to_owned()
73 + } else {
74 + String::new()
75 + }),
76 + IndentArg::Integer(indent) => Some(if indent > 0 {
77 + " ".repeat(indent as usize)
78 + } else {
79 + String::new()
80 + }),
81 + IndentArg::String(indent) => Some(indent),
82 + }
83 +}
84 +
85 +#[derive(Deserialize)]
86 +struct SeparatorsArg((String, String));
87 +
88 +fn parse_separators(value: Option<SeparatorsArg>, pretty: bool) -> (String, String) {
89 + let Some(SeparatorsArg((item_separator, key_separator))) = value else {
90 + let default_item_separator = if pretty { "," } else { ", " };
91 + let default_key_separator = ": ";
92 +
93 + return (
94 + default_item_separator.to_owned(),
95 + default_key_separator.to_owned(),
96 + );
97 + };
98 +
99 + (item_separator, key_separator)
100 +}
101 +
102 +fn build_json_format(
103 + indent: Option<String>,
104 + item_separator: String,
105 + key_separator: String,
106 + ensure_ascii: bool,
107 +) -> std::result::Result<JsonFormat, MinijinjaError> {
108 + JsonFormat::new()
109 + .indent(indent)
110 + .map_err(map_json_syntax_error("indent"))?
111 + .comma(item_separator)
112 + .map_err(map_json_syntax_error("separators (item)"))?
113 + .colon(key_separator)
114 + .map_err(map_json_syntax_error("separators (key)"))
115 + .map(|format| format.ascii(ensure_ascii))
116 +}
117 +
118 +fn map_json_syntax_error(
119 + field: &'static str,
120 +) -> impl FnOnce(JsonSyntaxError) -> MinijinjaError + Copy {
121 + move |error| {
122 + MinijinjaError::new(
123 + ErrorKind::InvalidOperation,
124 + format!("invalid {field} value for tojson: {error}"),
125 + )
126 + }
127 +}
128 +
129 +/// Recursively sort all object keys in a JSON value.
130 +fn sort_json_keys(value: &JsonValue) -> JsonValue {
131 + match value {
132 + JsonValue::Object(map) => {
133 + let mut sorted: serde_json::Map<String, JsonValue> = serde_json::Map::new();
134 + let mut keys: Vec<_> = map.keys().collect();
135 + keys.sort();
136 + for key in keys {
137 + sorted.insert(key.clone(), sort_json_keys(&map[key]));
138 + }
139 + JsonValue::Object(sorted)
140 + }
141 + JsonValue::Array(arr) => JsonValue::Array(arr.iter().map(sort_json_keys).collect()),
142 + _ => value.clone(),
143 + }
144 +}
145 +
146 +#[cfg(test)]
147 +mod tests {
148 + use expect_test::expect;
149 + use minijinja::Environment;
150 + use serde_json::json;
151 + use thiserror_ext::AsReport;
152 +
153 + use super::hf_tojson_filter;
154 +
155 + fn render(template: &str, payload: serde_json::Value) -> String {
156 + let mut env = Environment::new();
157 + env.add_filter("tojson", hf_tojson_filter);
158 + env.render_str(template, json!({ "payload": payload })).unwrap()
159 + }
160 +
161 + fn render_error(template: &str, payload: serde_json::Value) -> minijinja::Error {
162 + let mut env = Environment::new();
163 + env.add_filter("tojson", hf_tojson_filter);
164 + env.render_str(template, json!({ "payload": payload })).unwrap_err()
165 + }
166 +
167 + #[test]
168 + fn tojson_does_not_html_escape_like_minijinja_builtin() {
169 + let rendered = render("{{ payload|tojson }}", json!("<tag>&'"));
170 + assert_eq!(rendered, "\"<tag>&'\"");
171 + }
172 +
173 + #[test]
174 + fn tojson_supports_sort_keys_recursively() {
175 + let rendered = render(
176 + "{{ payload|tojson(sort_keys=true) }}",
177 + json!({
178 + "z": {"b": 1, "a": 2},
179 + "a": 0
180 + }),
181 + );
182 +
183 + assert_eq!(rendered, "{\"a\": 0, \"z\": {\"a\": 2, \"b\": 1}}");
184 + }
185 +
186 + #[test]
187 + fn tojson_supports_indent() {
188 + let rendered = render("{{ payload|tojson(indent=2) }}", json!([1, 2]));
189 +
190 + assert_eq!(rendered, "[\n 1,\n 2\n]");
191 + }
192 +
193 + #[test]
194 + fn tojson_supports_ensure_ascii_false() {
195 + let rendered = render("{{ payload|tojson(ensure_ascii=false) }}", json!("中文"));
196 + assert_eq!(rendered, "\"中文\"");
197 + }
198 +
199 + #[test]
200 + fn tojson_supports_ensure_ascii_true() {
201 + let rendered = render("{{ payload|tojson(ensure_ascii=true) }}", json!("中文"));
202 + assert_eq!(rendered, "\"\\u4e2d\\u6587\"");
203 + }
204 +
205 + #[test]
206 + fn tojson_supports_separators() {
207 + let rendered = render(
208 + "{{ payload|tojson(separators=[',', ':']) }}",
209 + json!({
210 + "x": [1, 2]
211 + }),
212 + );
213 +
214 + assert_eq!(rendered, "{\"x\":[1,2]}");
215 + }
216 +
217 + #[test]
218 + fn tojson_supports_negative_indent_as_newline_only() {
219 + let rendered = render("{{ payload|tojson(indent=-1) }}", json!([1, 2]));
220 + assert_eq!(rendered, "[\n1,\n2\n]");
221 + }
222 +
223 + #[test]
224 + fn tojson_supports_string_indent() {
225 + let rendered = render("{{ payload|tojson(indent=' ') }}", json!([1, 2]));
226 + assert_eq!(rendered, "[\n 1,\n 2\n]");
227 + }
228 +
229 + #[test]
230 + fn tojson_supports_boolean_indent() {
231 + let rendered_true = render("{{ payload|tojson(indent=true) }}", json!([1, 2]));
232 + assert_eq!(rendered_true, "[\n 1,\n 2\n]");
233 +
234 + let rendered_false = render("{{ payload|tojson(indent=false) }}", json!([1, 2]));
235 + assert_eq!(rendered_false, "[\n1,\n2\n]");
236 + }
237 +
238 + #[test]
239 + fn tojson_combines_indent_sort_keys_separators_and_ensure_ascii() {
240 + let rendered = render(
241 + "{{ payload|tojson(ensure_ascii=true, sort_keys=true, separators=[',', ':'], indent=' ') }}",
242 + json!({
243 + "b": "<中>",
244 + "a": [1, 2]
245 + }),
246 + );
247 +
248 + assert_eq!(
249 + rendered,
250 + "{\n \"a\":[\n 1,\n 2\n ],\n \"b\":\"<\\u4e2d>\"\n}"
251 + );
252 + }
253 +
254 + #[test]
255 + fn tojson_rejects_invalid_indent() {
256 + let error = render_error("{{ payload|tojson(indent='-->') }}", json!({"a": 1}));
257 + expect!["invalid operation: invalid indent value for tojson: string contains unexpected character '-' (in <string>:1)"]
258 + .assert_eq(&error.to_report_string());
259 + }
260 +
261 + #[test]
262 + fn tojson_rejects_invalid_separator_shape() {
263 + let error = render_error("{{ payload|tojson(separators=':,') }}", json!({"a": 1}));
264 + expect!["cannot deserialize: invalid type: string \":,\", expected a tuple of size 2 (in <string>:1)"]
265 + .assert_eq(&error.to_report_string());
266 + }
267 +
268 + #[test]
269 + fn tojson_rejects_invalid_key_separator() {
270 + let error = render_error(
271 + "{{ payload|tojson(separators=[',', '=>']) }}",
272 + json!({"a": 1}),
273 + );
274 + expect!["invalid operation: invalid separators (key) value for tojson: string contains unexpected character '=' (in <string>:1)"]
275 + .assert_eq(&error.to_report_string());
276 + }
277 +}
rust/src/chat/src/renderer/mod.rs new
+31
@@ -0,0 +1,31 @@
1 +use std::sync::Arc;
2 +
3 +use vllm_text::Prompt;
4 +
5 +use crate::error::Result;
6 +use crate::request::ChatRequest;
7 +
8 +pub mod deepseek_v32;
9 +pub mod deepseek_v4;
10 +pub mod hf;
11 +mod selection;
12 +
13 +pub use deepseek_v4::DeepSeekV4ChatRenderer;
14 +pub use deepseek_v32::DeepSeekV32ChatRenderer;
15 +pub use selection::RendererSelection;
16 +
17 +/// Rendered chat prompt submitted to the text backend.
18 +#[derive(Debug, Clone, PartialEq, Eq)]
19 +pub struct RenderedPrompt {
20 + pub prompt: Prompt,
21 +}
22 +
23 +/// Minimal chat-prompt renderer used by `vllm-chat`.
24 +pub trait ChatRenderer: Send + Sync {
25 + /// Render one chat request into the text prompt submitted to the text
26 + /// backend.
27 + fn render(&self, request: &ChatRequest) -> Result<RenderedPrompt>;
28 +}
29 +
30 +/// Shared trait-object form of [`ChatRenderer`].
31 +pub type DynChatRenderer = Arc<dyn ChatRenderer>;
rust/src/chat/src/renderer/selection.rs new
+109
@@ -0,0 +1,109 @@
1 +use std::fmt;
2 +use std::str::FromStr;
3 +
4 +use serde_with::DeserializeFromStr;
5 +
6 +/// Specify which chat renderer implementation to use.
7 +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, DeserializeFromStr)]
8 +pub enum RendererSelection {
9 + /// Use model-based auto-detection.
10 + #[default]
11 + Auto,
12 + /// Force the generic Hugging Face chat-template renderer.
13 + Hf,
14 + /// Force the DeepSeek V3.2 renderer.
15 + DeepSeekV32,
16 + /// Force the DeepSeek V4 renderer.
17 + DeepSeekV4,
18 +}
19 +
20 +impl RendererSelection {
21 + pub const AUTO_LITERAL: &str = "auto";
22 + pub const DEEPSEEK_V32_LITERAL: &str = "deepseek_v32";
23 + pub const DEEPSEEK_V4_LITERAL: &str = "deepseek_v4";
24 + pub const HF_LITERAL: &str = "hf";
25 +
26 + /// Resolve the renderer selection using the given model type string, if
27 + /// it's `Auto`.
28 + pub fn resolve(self, model_type: &str) -> Self {
29 + match self {
30 + Self::Auto => match model_type {
31 + Self::DEEPSEEK_V32_LITERAL => Self::DeepSeekV32,
32 + Self::DEEPSEEK_V4_LITERAL => Self::DeepSeekV4,
33 + _ => Self::Hf,
34 + },
35 + selection => selection,
36 + }
37 + }
38 +}
39 +
40 +impl FromStr for RendererSelection {
41 + type Err = String;
42 +
43 + fn from_str(value: &str) -> Result<Self, Self::Err> {
44 + if value.eq_ignore_ascii_case(Self::AUTO_LITERAL) {
45 + Ok(Self::Auto)
46 + } else if value.eq_ignore_ascii_case(Self::HF_LITERAL) {
47 + Ok(Self::Hf)
48 + } else if value.eq_ignore_ascii_case(Self::DEEPSEEK_V32_LITERAL) {
49 + Ok(Self::DeepSeekV32)
50 + } else if value.eq_ignore_ascii_case(Self::DEEPSEEK_V4_LITERAL) {
51 + Ok(Self::DeepSeekV4)
52 + } else {
53 + Err(format!(
54 + "unknown renderer `{value}` (expected one of: auto, hf, deepseek_v32, deepseek_v4)"
55 + ))
56 + }
57 + }
58 +}
59 +
60 +impl fmt::Display for RendererSelection {
61 + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
62 + match self {
63 + Self::Auto => f.write_str(Self::AUTO_LITERAL),
64 + Self::Hf => f.write_str(Self::HF_LITERAL),
65 + Self::DeepSeekV32 => f.write_str(Self::DEEPSEEK_V32_LITERAL),
66 + Self::DeepSeekV4 => f.write_str(Self::DEEPSEEK_V4_LITERAL),
67 + }
68 + }
69 +}
70 +
71 +#[cfg(test)]
72 +mod tests {
73 + use super::RendererSelection;
74 +
75 + #[test]
76 + fn renderer_selection_parses_known_values() {
77 + assert_eq!(
78 + "auto".parse::<RendererSelection>().unwrap(),
79 + RendererSelection::Auto
80 + );
81 + assert_eq!(
82 + "hf".parse::<RendererSelection>().unwrap(),
83 + RendererSelection::Hf
84 + );
85 + assert_eq!(
86 + "deepseek_v32".parse::<RendererSelection>().unwrap(),
87 + RendererSelection::DeepSeekV32
88 + );
89 + assert_eq!(
90 + "deepseek_v4".parse::<RendererSelection>().unwrap(),
91 + RendererSelection::DeepSeekV4
92 + );
93 + }
94 +
95 + #[test]
96 + fn renderer_selection_display_round_trips() {
97 + for selection in [
98 + RendererSelection::Auto,
99 + RendererSelection::Hf,
100 + RendererSelection::DeepSeekV32,
101 + RendererSelection::DeepSeekV4,
102 + ] {
103 + assert_eq!(
104 + selection.to_string().parse::<RendererSelection>().unwrap(),
105 + selection
106 + );
107 + }
108 + }
109 +}
rust/src/chat/src/request.rs new
+662
@@ -0,0 +1,662 @@
1 +use std::collections::HashMap;
2 +
3 +use llm_multimodal::ImageDetail;
4 +use serde::{Deserialize, Serialize};
5 +use serde_json::Value;
6 +pub use vllm_text::SamplingParams;
7 +use vllm_text::TextDecodeOptions;
8 +pub use vllm_tool_parser::Tool as ChatTool;
9 +
10 +use crate::AssistantMessageExt;
11 +use crate::error::{Error, Result};
12 +use crate::event::{AssistantContentBlock, AssistantMessage};
13 +
14 +/// Role label for one text-only chat message.
15 +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
16 +#[serde(rename_all = "snake_case")]
17 +pub enum ChatRole {
18 + System,
19 + Developer,
20 + User,
21 + Assistant,
22 + ToolResponse,
23 +}
24 +
25 +/// One text-only chat content part in OpenAI-style block format.
26 +#[serde_with::skip_serializing_none]
27 +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
28 +#[serde(tag = "type", rename_all = "snake_case")]
29 +pub enum ChatContentPart {
30 + /// One plain-text content block.
31 + Text { text: String },
32 + /// One image URL/data URL content block.
33 + ImageUrl {
34 + image_url: String,
35 + detail: Option<ImageDetail>,
36 + uuid: Option<String>,
37 + },
38 + // ImageData...
39 + // ImageEmbeds...
40 +}
41 +
42 +impl ChatContentPart {
43 + /// Construct one text content part with plain string content.
44 + pub fn text(text: impl Into<String>) -> Self {
45 + Self::Text { text: text.into() }
46 + }
47 +
48 + /// Construct one image URL content part with the given URL string.
49 + pub fn image_url(image_url: impl Into<String>) -> Self {
50 + Self::ImageUrl {
51 + image_url: image_url.into(),
52 + detail: None,
53 + uuid: None,
54 + }
55 + }
56 +
57 + /// Return the text content of this part when it's a text block, or an
58 + /// "unsupported multimodal content" error otherwise.
59 + pub(crate) fn as_text(&self) -> Result<&str> {
60 + match self {
61 + Self::Text { text } => Ok(text),
62 + Self::ImageUrl { .. } => Err(Error::UnsupportedMultimodalContent("image_url")),
63 + }
64 + }
65 +
66 + /// Return whether this part is a text block with empty content.
67 + pub(crate) fn is_empty_text(&self) -> bool {
68 + matches!(self, Self::Text { text } if text.is_empty())
69 + }
70 +
71 + /// Return whether this part contains any multimodal content.
72 + pub(crate) fn is_multimodal(&self) -> bool {
73 + match self {
74 + Self::Text { .. } => false,
75 + Self::ImageUrl { .. } => true,
76 + }
77 + }
78 +}
79 +
80 +/// Text-only chat content.
81 +///
82 +/// This supports either a simple string or an OpenAI-style list of text blocks.
83 +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
84 +#[serde(untagged)]
85 +pub enum ChatContent {
86 + /// Simple text content.
87 + Text(String),
88 + /// OpenAI-style blocks.
89 + Parts(Vec<ChatContentPart>),
90 +}
91 +
92 +impl ChatContent {
93 + /// Flatten the text content into one plain string without adding
94 + /// separators.
95 + // TODO: this method will be truly fallible once we add non-text content parts.
96 + pub fn try_flatten_to_text(&self) -> Result<String> {
97 + Ok(match self {
98 + Self::Text(text) => text.clone(),
99 + Self::Parts(parts) => {
100 + parts.iter().map(ChatContentPart::as_text).collect::<Result<Vec<_>>>()?.concat()
101 + }
102 + })
103 + }
104 +
105 + /// Return whether there's no text content or only empty text blocks.
106 + pub fn is_empty(&self) -> bool {
107 + match self {
108 + Self::Text(text) => text.is_empty(),
109 + Self::Parts(parts) => parts.iter().all(ChatContentPart::is_empty_text),
110 + }
111 + }
112 +
113 + /// Return whether this content contains any multimodal parts.
114 + pub fn has_multimodal(&self) -> bool {
115 + match self {
116 + Self::Text(_) => false,
117 + Self::Parts(parts) => parts.iter().any(ChatContentPart::is_multimodal),
118 + }
119 + }
120 +}
121 +
122 +impl From<String> for ChatContent {
123 + fn from(value: String) -> Self {
124 + Self::Text(value)
125 + }
126 +}
127 +
128 +impl From<&str> for ChatContent {
129 + fn from(value: &str) -> Self {
130 + Self::Text(value.to_string())
131 + }
132 +}
133 +
134 +impl From<Vec<ChatContentPart>> for ChatContent {
135 + fn from(value: Vec<ChatContentPart>) -> Self {
136 + Self::Parts(value)
137 + }
138 +}
139 +
140 +/// One chat message.
141 +///
142 +/// Original Python API reference:
143 +/// <https://github.com/vllm-project/vllm/blob/bc2c0c86efb28e77677a3cfb8687e976914a313a/vllm/entrypoints/chat_utils.py#L309-L333>
144 +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
145 +#[serde(tag = "role", rename_all = "snake_case")]
146 +pub enum ChatMessage {
147 + /// System message content.
148 + System { content: ChatContent },
149 + /// Developer message content plus optional message-local tools.
150 + Developer {
151 + content: ChatContent,
152 + tools: Option<Vec<ChatTool>>,
153 + },
154 + /// User message content.
155 + User { content: ChatContent },
156 + /// Assistant history content assembled from structured assistant blocks.
157 + Assistant { content: Vec<AssistantContentBlock> },
158 + /// Tool response content associated with one prior assistant tool call.
159 + ToolResponse {
160 + content: ChatContent,
161 + tool_call_id: String,
162 + },
163 +}
164 +
165 +impl ChatMessage {
166 + /// Construct one chat message with plain string content.
167 + pub fn text(role: ChatRole, text: impl Into<String>) -> Self {
168 + let content: String = text.into();
169 +
170 + match role {
171 + ChatRole::System => Self::system(content),
172 + ChatRole::Developer => Self::developer(content, None),
173 + ChatRole::User => Self::user(content),
174 + ChatRole::Assistant => Self::assistant_text(content),
175 + ChatRole::ToolResponse => {
176 + panic!(
177 + "tool response messages require a tool_call_id; \
178 + use ChatMessage::tool_response() instead"
179 + )
180 + }
181 + }
182 + }
183 +
184 + /// Construct one chat message with system role.
185 + pub fn system(content: impl Into<ChatContent>) -> Self {
186 + Self::System {
187 + content: content.into(),
188 + }
189 + }
190 +
191 + /// Construct one chat message with developer role.
192 + pub fn developer(content: impl Into<ChatContent>, tools: Option<Vec<ChatTool>>) -> Self {
193 + Self::Developer {
194 + content: content.into(),
195 + tools,
196 + }
197 + }
198 +
199 + /// Construct one chat message with user role.
200 + pub fn user(content: impl Into<ChatContent>) -> Self {
201 + Self::User {
202 + content: content.into(),
203 + }
204 + }
205 +
206 + /// Construct one chat message with assistant role and plain string content.
207 + pub fn assistant_text(text: impl Into<String>) -> Self {
208 + Self::Assistant {
209 + content: vec![AssistantContentBlock::Text { text: text.into() }],
210 + }
211 + }
212 +
213 + /// Construct one chat message with assistant role and structured content
214 + /// blocks.
215 + pub fn assistant_blocks(content: Vec<AssistantContentBlock>) -> Self {
216 + Self::Assistant { content }
217 + }
218 +
219 + /// Construct one tool-role message.
220 + pub fn tool_response(content: impl Into<ChatContent>, tool_call_id: impl Into<String>) -> Self {
221 + Self::ToolResponse {
222 + content: content.into(),
223 + tool_call_id: tool_call_id.into(),
224 + }
225 + }
226 +
227 + /// Return the chat role of this message.
228 + pub fn role(&self) -> ChatRole {
229 + match self {
230 + Self::System { .. } => ChatRole::System,
231 + Self::Developer { .. } => ChatRole::Developer,
232 + Self::User { .. } => ChatRole::User,
233 + Self::Assistant { .. } => ChatRole::Assistant,
234 + Self::ToolResponse { .. } => ChatRole::ToolResponse,
235 + }
236 + }
237 +
238 + /// Concatenate the visible text carried by this message.
239 + pub fn text_content(&self) -> Result<String> {
240 + match self {
241 + Self::System { content }
242 + | Self::Developer { content, .. }
243 + | Self::User { content }
244 + | Self::ToolResponse { content, .. } => content.try_flatten_to_text(),
245 + Self::Assistant { content } => Ok(content.text()),
246 + }
247 + }
248 +
249 + /// Concatenate assistant reasoning text when present.
250 + pub fn reasoning_content(&self) -> Option<String> {
251 + match self {
252 + Self::Assistant { content } => content.reasoning(),
253 + Self::System { .. }
254 + | Self::Developer { .. }
255 + | Self::User { .. }
256 + | Self::ToolResponse { .. } => None,
257 + }
258 + }
259 +
260 + /// Return whether this message contains any multimodal content.
261 + pub fn has_multimodal(&self) -> bool {
262 + match self {
263 + Self::System { content }
264 + | Self::Developer { content, .. }
265 + | Self::User { content }
266 + | Self::ToolResponse { content, .. } => content.has_multimodal(),
267 + Self::Assistant { .. } => false,
268 + }
269 + }
270 +}
271 +
272 +impl From<AssistantMessage> for ChatMessage {
273 + fn from(value: AssistantMessage) -> Self {
274 + Self::Assistant {
275 + content: value.content,
276 + }
277 + }
278 +}
279 +
280 +/// Controls how prompt rendering should end after the existing chat history.
281 +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
282 +#[serde(rename_all = "snake_case")]
283 +pub enum GenerationPromptMode {
284 + /// Append a generation prompt for a new assistant turn.
285 + ///
286 + /// Equivalent to `add_generation_prompt = true` and `continue_final_message
287 + /// = false`.
288 + #[default]
289 + StartNewAssistant,
290 + /// Leave the final assistant message open so generation continues it.
291 + ///
292 + /// Equivalent to `add_generation_prompt = false` and
293 + /// `continue_final_message = true`.
294 + ContinueFinalAssistant,
295 + /// Render the existing chat history without adding any trailing generation
296 + /// prompt.
297 + ///
298 + /// Equivalent to `add_generation_prompt = false` and
299 + /// `continue_final_message = false`.
300 + NoGenerationPrompt,
301 +}
302 +
303 +/// Effort level for reasoning models.
304 +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
305 +#[serde(rename_all = "lowercase")]
306 +pub enum ReasoningEffort {
307 + None,
308 + Minimal,
309 + Low,
310 + Medium,
311 + High,
312 + XHigh,
313 + Max,
314 +}
315 +
316 +impl ReasoningEffort {
317 + pub fn as_str(self) -> &'static str {
318 + match self {
319 + Self::None => "none",
320 + Self::Minimal => "minimal",
321 + Self::Low => "low",
322 + Self::Medium => "medium",
323 + Self::High => "high",
324 + Self::XHigh => "xhigh",
325 + Self::Max => "max",
326 + }
327 + }
328 +}
329 +
330 +/// Chat-template-related request options.
331 +///
332 +/// These are the small subset of chat controls that currently affect prompt
333 +/// rendering in `vllm-chat`.
334 +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
335 +pub struct ChatOptions {
336 + /// Controls whether rendering starts a new assistant turn, continues the
337 + /// final assistant message, or emits no trailing generation prompt at
338 + /// all.
339 + pub generation_prompt_mode: GenerationPromptMode,
340 +
341 + /// Per-request Jinja chat template override. When set, this template is
342 + /// used instead of the model's default chat template.
343 + pub chat_template: Option<String>,
344 +
345 + /// Effort level exposed to chat templates for reasoning models.
346 + pub reasoning_effort: Option<ReasoningEffort>,
347 +
348 + /// Additional keyword arguments exposed to the chat template.
349 + pub template_kwargs: HashMap<String, Value>,
350 +}
351 +
352 +impl Default for ChatOptions {
353 + fn default() -> Self {
354 + Self {
355 + generation_prompt_mode: GenerationPromptMode::StartNewAssistant,
356 + chat_template: None,
357 + reasoning_effort: None,
358 + template_kwargs: HashMap::new(),
359 + }
360 + }
361 +}
362 +
363 +impl ChatOptions {
364 + /// Whether to add a generation prompt for a new assistant turn after the
365 + /// existing chat history.
366 + pub fn add_generation_prompt(&self) -> bool {
367 + matches!(
368 + self.generation_prompt_mode,
369 + GenerationPromptMode::StartNewAssistant
370 + )
371 + }
372 +
373 + /// Whether to leave the final assistant message open so generation
374 + /// continues it.
375 + pub fn continue_final_message(&self) -> bool {
376 + matches!(
377 + self.generation_prompt_mode,
378 + GenerationPromptMode::ContinueFinalAssistant
379 + )
380 + }
381 +}
382 +
383 +/// Tool-choice semantics supported by `vllm-chat`.
384 +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
385 +#[serde(rename_all = "snake_case")]
386 +pub enum ChatToolChoice {
387 + Auto,
388 + #[default]
389 + None,
390 +}
391 +
392 +/// One chat request ready to be rendered into a prompt and lowered into a
393 +/// generate request.
394 +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
395 +pub struct ChatRequest {
396 + /// Stable caller-supplied request ID.
397 + pub request_id: String,
398 + /// Ordered chat history to render.
399 + pub messages: Vec<ChatMessage>,
400 + /// User-facing sampling parameters accepted by `vllm-chat`.
401 + pub sampling_params: SamplingParams,
402 + /// Chat-specific rendering options.
403 + pub chat_options: ChatOptions,
404 + /// Function tools made available to the model for this request.
405 + pub tools: Vec<ChatTool>,
406 + /// Tool-choice behavior for this request.
407 + pub tool_choice: ChatToolChoice,
408 + /// Text decode options for incremental detokenization.
409 + pub decode_options: TextDecodeOptions,
410 + /// Whether to emit intermediate northbound content deltas before the
411 + /// terminal result.
412 + ///
413 + /// If `false`, callers only observe the terminal accumulated assistant
414 + /// output. If `true`, callers may receive zero or more incremental
415 + /// content events before the final terminal one.
416 + pub intermediate: bool,
417 + /// Request scheduling priority (lower means earlier handling; default 0).
418 + pub priority: i32,
419 + /// Documents for RAG (retrieval-augmented generation), passed to the chat
420 + /// template.
421 + pub documents: Option<Vec<Value>>,
422 + /// Salt for prefix cache isolation in multi-user environments.
423 + pub cache_salt: Option<String>,
424 + /// Whether to add special tokens (e.g. BOS) during prompt tokenization.
425 + pub add_special_tokens: bool,
426 + /// Override data parallel rank.
427 + #[serde(default)]
428 + pub data_parallel_rank: Option<u32>,
429 +}
430 +
431 +impl ChatRequest {
432 + /// Return one minimal valid request fixture for tests.
433 + pub fn for_test() -> Self {
434 + Self {
435 + request_id: "test-request".to_string(),
436 + messages: vec![ChatMessage::text(ChatRole::User, "test")],
437 + sampling_params: SamplingParams::default(),
438 + chat_options: ChatOptions::default(),
439 + tools: Vec::new(),
440 + tool_choice: ChatToolChoice::None,
441 + decode_options: TextDecodeOptions::default(),
442 + intermediate: true,
443 + priority: 0,
444 + documents: None,
445 + cache_salt: None,
446 + add_special_tokens: false,
447 + data_parallel_rank: None,
448 + }
449 + }
450 +
451 + /// Validate basic request invariants before rendering.
452 + pub fn validate(&self) -> Result<()> {
453 + if self.messages.is_empty() {
454 + return Err(Error::EmptyMessages);
455 + }
456 + match (
457 + self.chat_options.generation_prompt_mode,
458 + self.messages.last().map(ChatMessage::role),
459 + ) {
460 + (GenerationPromptMode::ContinueFinalAssistant, Some(ChatRole::Assistant)) => {}
461 + (GenerationPromptMode::ContinueFinalAssistant, _) => {
462 + return Err(Error::ContinueFinalAssistantWithoutFinalAssistant);
463 + }
464 + (GenerationPromptMode::NoGenerationPrompt, _)
465 + | (GenerationPromptMode::StartNewAssistant, _) => {}
466 + }
467 + Ok(())
468 + }
469 +
470 + /// Return true if this request contains any multimodal content in its
471 + /// messages.
472 + pub fn has_multimodal(&self) -> bool {
473 + self.messages.iter().any(ChatMessage::has_multimodal)
474 + }
475 +
476 + /// Return true if this request should enable tool parsing based on the tool
477 + /// choice and tool list.
478 + pub(crate) fn tool_parsing_enabled(&self) -> bool {
479 + matches!(self.tool_choice, ChatToolChoice::Auto) && !self.tools.is_empty()
480 + }
481 +
482 + /// Return the request-level thinking toggle when explicitly requested.
483 + ///
484 + /// We currently accept the two request kwargs `thinking` and
485 + /// `enable_thinking`. Both must be booleans when present. If both are
486 + /// present, they must have the same value. If neither key is provided,
487 + /// return `None`.
488 + pub(crate) fn enable_thinking(&self) -> Result<Option<bool>> {
489 + let thinking = self.parse_template_bool("thinking")?;
490 + let enable_thinking = self.parse_template_bool("enable_thinking")?;
491 +
492 + match (thinking, enable_thinking) {
493 + (None, None) => Ok(None),
494 + (Some(thinking), Some(enable_thinking)) if thinking != enable_thinking => {
495 + Err(Error::ChatTemplate(
496 + "template kwargs `thinking` and `enable_thinking` must match when both are set"
497 + .to_string(),
498 + ))
499 + }
500 + (Some(thinking), _) => Ok(Some(thinking)),
501 + (None, Some(enable_thinking)) => Ok(Some(enable_thinking)),
502 + }
503 + }
504 +
505 + pub(crate) fn parse_template_bool(&self, key: &str) -> Result<Option<bool>> {
506 + match self.chat_options.template_kwargs.get(key) {
507 + None => Ok(None),
508 + Some(Value::Bool(value)) => Ok(Some(*value)),
509 + Some(other) => Err(Error::ChatTemplate(format!(
510 + "template kwarg `{key}` must be a boolean, got {other}"
511 + ))),
512 + }
513 + }
514 +}
515 +
516 +impl ChatRole {
517 + /// Return the chat-template role string used by the current text-only chat
518 + /// backend.
519 + pub fn as_str(&self) -> &'static str {
520 + match self {
521 + Self::System => "system",
522 + Self::Developer => "developer",
523 + Self::User => "user",
524 + Self::Assistant => "assistant",
525 + Self::ToolResponse => "tool_response",
526 + }
527 + }
528 +}
529 +
530 +#[cfg(test)]
531 +mod tests {
532 + use serde_json::{json, to_value};
533 +
534 + use super::{ChatContent, ChatContentPart, ChatMessage, ChatRequest, ChatRole, ChatTool};
535 + use crate::Error;
536 + use crate::event::AssistantContentBlock;
537 +
538 + #[test]
539 + fn chat_content_deserializes_from_raw_string() {
540 + let content: ChatContent = serde_json::from_value(json!("hello")).unwrap();
541 + assert_eq!(content, ChatContent::Text("hello".to_string()));
542 + }
543 +
544 + #[test]
545 + fn chat_content_deserializes_from_openai_text_blocks() {
546 + let content: ChatContent =
547 + serde_json::from_value(json!([{ "type": "text", "text": "hello" }])).unwrap();
548 + assert_eq!(
549 + content,
550 + ChatContent::Parts(vec![ChatContentPart::text("hello")])
551 + );
552 + }
553 +
554 + #[test]
555 + fn chat_content_from_string_like_values_builds_text() {
556 + assert_eq!(
557 + ChatContent::from("hello"),
558 + ChatContent::Text("hello".to_string())
559 + );
560 + assert_eq!(
561 + ChatContent::from("hello".to_string()),
562 + ChatContent::Text("hello".to_string())
563 + );
564 + }
565 +
566 + #[test]
567 + fn chat_content_try_flattens_text_parts_without_separators() {
568 + let content = ChatContent::Parts(vec![
569 + ChatContentPart::text("hello"),
570 + ChatContentPart::text(" world"),
571 + ]);
572 + assert_eq!(content.try_flatten_to_text().unwrap(), "hello world");
573 + }
574 +
575 + #[test]
576 + fn assistant_message_collects_visible_and_reasoning_text() {
577 + let message = ChatMessage::assistant_blocks(vec![
578 + AssistantContentBlock::Reasoning {
579 + text: "inner".to_string(),
580 + },
581 + AssistantContentBlock::Text {
582 + text: "outer".to_string(),
583 + },
584 + ]);
585 +
586 + assert_eq!(message.role(), ChatRole::Assistant);
587 + assert_eq!(message.text_content().unwrap(), "outer");
588 + assert_eq!(message.reasoning_content().as_deref(), Some("inner"));
589 + }
590 +
591 + #[test]
592 + fn developer_message_round_trips_through_serde() {
593 + let message = ChatMessage::developer(
594 + "hello",
595 + Some(vec![ChatTool {
596 + name: "get_weather".to_string(),
597 + description: Some("Get weather".to_string()),
598 + parameters: json!({
599 + "type": "object",
600 + "properties": {"city": {"type": "string"}},
601 + }),
602 + strict: Some(true),
603 + }]),
604 + );
605 +
606 + let value = to_value(&message).unwrap();
607 + let decoded: ChatMessage = serde_json::from_value(value).unwrap();
608 + assert_eq!(decoded, message);
609 + }
610 +
611 + #[test]
612 + fn enable_thinking_is_none_when_no_kwargs_are_present() {
613 + let request = ChatRequest::for_test();
614 + assert_eq!(request.enable_thinking().unwrap(), None);
615 + }
616 +
617 + #[test]
618 + fn enable_thinking_accepts_matching_duplicate_kwargs() {
619 + let mut request = ChatRequest::for_test();
620 + request.chat_options.template_kwargs.insert("thinking".to_string(), json!(true));
621 + request
622 + .chat_options
623 + .template_kwargs
624 + .insert("enable_thinking".to_string(), json!(true));
625 +
626 + assert_eq!(request.enable_thinking().unwrap(), Some(true));
627 + }
628 +
629 + #[test]
630 + fn enable_thinking_rejects_non_boolean_kwargs() {
631 + let mut request = ChatRequest::for_test();
632 + request
633 + .chat_options
634 + .template_kwargs
635 + .insert("thinking".to_string(), json!("yes"));
636 +
637 + assert!(matches!(
638 + request.enable_thinking(),
639 + Err(Error::ChatTemplate(message))
640 + if message.contains("`thinking` must be a boolean")
641 + ));
642 + }
643 +
644 + #[test]
645 + fn enable_thinking_rejects_conflicting_duplicate_kwargs() {
646 + let mut request = ChatRequest::for_test();
647 + request
648 + .chat_options
649 + .template_kwargs
650 + .insert("thinking".to_string(), json!(false));
651 + request
652 + .chat_options
653 + .template_kwargs
654 + .insert("enable_thinking".to_string(), json!(true));
655 +
656 + assert!(matches!(
657 + request.enable_thinking(),
658 + Err(Error::ChatTemplate(message))
659 + if message.contains("`thinking` and `enable_thinking` must match")
660 + ));
661 + }
662 +}
rust/src/chat/src/stream.rs new
+237
@@ -0,0 +1,237 @@
1 +use std::pin::Pin;
2 +use std::sync::Arc;
3 +use std::task::{Context, Poll};
4 +
5 +use futures::Stream;
6 +use trait_set::trait_set;
7 +use vllm_text::{DecodedLogprobs, DecodedPositionLogprobs, DecodedPromptLogprobs};
8 +
9 +use crate::FinishReason;
10 +use crate::error::{Error, Result};
11 +use crate::event::{AssistantContentBlock, AssistantMessage, ChatEvent};
12 +
13 +/// Final structured assistant message plus terminal stream metadata.
14 +#[derive(Debug, Clone, PartialEq)]
15 +pub struct CollectedAssistantMessage {
16 + pub message: AssistantMessage,
17 + pub prompt_token_count: usize,
18 + pub prompt_token_ids: Arc<[u32]>,
19 + pub prompt_logprobs: Option<DecodedPromptLogprobs>,
20 + pub logprobs: Option<DecodedLogprobs>,
21 + pub token_ids: Vec<u32>,
22 + pub output_token_count: usize,
23 + pub finish_reason: FinishReason,
24 + /// Connector-specific KV transfer parameters for disaggregated serving.
25 + pub kv_transfer_params: Option<serde_json::Value>,
26 +}
27 +
28 +/// Per-request stream of chat events.
29 +pub struct ChatEventStream {
30 + request_id: String,
31 + inner: Pin<Box<dyn Stream<Item = Result<ChatEvent>> + Send>>,
32 +}
33 +
34 +impl ChatEventStream {
35 + pub(crate) fn new(request_id: String, inner: impl crate::output::ChatEventStream) -> Self {
36 + Self {
37 + request_id,
38 + inner: Box::pin(inner),
39 + }
40 + }
41 +
42 + /// Return the request ID associated with this stream.
43 + pub fn request_id(&self) -> &str {
44 + &self.request_id
45 + }
46 +
47 + /// Collect the stream to completion and return the final assembled
48 + /// assistant message.
49 + pub async fn collect_message(mut self) -> Result<CollectedAssistantMessage> {
50 + use futures::StreamExt as _;
51 +
52 + let mut message = AssistantMessage::default();
53 + let mut prompt_logprobs = None;
54 + let mut prompt_token_ids: Arc<[u32]> = Arc::from([]);
55 + let mut logprob_positions: Vec<DecodedPositionLogprobs> = Vec::new();
56 + let mut token_ids: Vec<u32> = Vec::new();
57 + while let Some(event) = self.next().await.transpose()? {
58 + match event {
59 + ChatEvent::Start {
60 + prompt_logprobs: start_prompt_logprobs,
61 + prompt_token_ids: start_prompt_token_ids,
62 + } => {
63 + prompt_logprobs = start_prompt_logprobs;
64 + prompt_token_ids = start_prompt_token_ids;
65 + }
66 + ChatEvent::BlockEnd { block, .. } => message.push_block(block),
67 + ChatEvent::LogprobsDelta {
68 + logprobs,
69 + token_ids: delta_ids,
70 + } => {
71 + if let Some(logprobs) = logprobs {
72 + logprob_positions.extend(logprobs.positions);
73 + }
74 + token_ids.extend(delta_ids);
75 + }
76 + ChatEvent::Done {
77 + message: done,
78 + prompt_token_count,
79 + output_token_count,
80 + finish_reason,
81 + kv_transfer_params,
82 + } => {
83 + return Ok(CollectedAssistantMessage {
84 + message: done,
85 + prompt_token_count,
86 + prompt_token_ids,
87 + prompt_logprobs,
88 + logprobs: (!logprob_positions.is_empty()).then_some(DecodedLogprobs {
89 + positions: logprob_positions,
90 + }),
91 + token_ids,
92 + output_token_count,
93 + finish_reason,
94 + kv_transfer_params,
95 + });
96 + }
97 + ChatEvent::ToolCallEnd { call, .. } => {
98 + message.push_block(AssistantContentBlock::ToolCall(call));
99 + }
100 + ChatEvent::BlockStart { .. }
101 + | ChatEvent::BlockDelta { .. }
102 + | ChatEvent::ToolCallStart { .. }
103 + | ChatEvent::ToolCallArgumentsDelta { .. } => {}
104 + }
105 + }
106 +
107 + // Note: this is actually unreachable, as the underlying stream always emit an
108 + // error on unexpected close.
109 + Err(Error::StreamClosedBeforeTerminalOutput {
110 + request_id: self.request_id,
111 + })
112 + }
113 +}
114 +
115 +impl Stream for ChatEventStream {
116 + type Item = Result<ChatEvent>;
117 +
118 + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
119 + Pin::new(&mut self.inner).poll_next(cx)
120 + }
121 +}
122 +
123 +trait_set! {
124 + pub trait ChatEventStreamTrait = Stream<Item = Result<ChatEvent>> + Send + 'static;
125 +}
126 +
127 +#[cfg(test)]
128 +mod tests {
129 +
130 + use futures::stream;
131 + use vllm_llm::FinishReason;
132 + use vllm_text::{
133 + DecodedLogprobs, DecodedPositionLogprobs, DecodedPromptLogprobs, DecodedTokenLogprob,
134 + };
135 +
136 + use super::{ChatEventStream, CollectedAssistantMessage};
137 + use crate::error::Error;
138 + use crate::event::ChatEvent;
139 +
140 + #[tokio::test]
141 + async fn collect_message_requires_terminal_done_event() {
142 + let stream = ChatEventStream::new(
143 + "chat-missing-done".to_string(),
144 + stream::iter([Ok(ChatEvent::Start {
145 + prompt_token_ids: vec![].into(),
146 + prompt_logprobs: None,
147 + })]),
148 + );
149 +
150 + let error = stream.collect_message().await.expect_err("missing done");
151 + assert!(matches!(
152 + error,
153 + Error::StreamClosedBeforeTerminalOutput { request_id }
154 + if request_id == "chat-missing-done"
155 + ));
156 + }
157 +
158 + #[tokio::test]
159 + async fn collect_message_retains_prompt_and_sample_logprobs() {
160 + let stream = ChatEventStream::new(
161 + "chat-logprobs".to_string(),
162 + stream::iter(vec![
163 + Ok(ChatEvent::Start {
164 + prompt_token_ids: vec![10, 11].into(),
165 + prompt_logprobs: Some(DecodedPromptLogprobs {
166 + first_token_id: 0,
167 + first_token: "o".to_string(),
168 + scored_positions: vec![DecodedPositionLogprobs {
169 + entries: vec![DecodedTokenLogprob {
170 + token_id: 0,
171 + token: "p".to_string(),
172 + logprob: -0.1,
173 + rank: 1,
174 + }],
175 + }],
176 + }),
177 + }),
178 + Ok(ChatEvent::LogprobsDelta {
179 + logprobs: Some(DecodedLogprobs {
180 + positions: vec![DecodedPositionLogprobs {
181 + entries: vec![DecodedTokenLogprob {
182 + token_id: 0,
183 + token: "a".to_string(),
184 + logprob: -0.2,
185 + rank: 1,
186 + }],
187 + }],
188 + }),
189 + token_ids: vec![],
190 + }),
191 + Ok(ChatEvent::Done {
192 + message: Default::default(),
193 + prompt_token_count: 2,
194 + output_token_count: 1,
195 + finish_reason: FinishReason::stop_eos(),
196 + kv_transfer_params: None,
197 + }),
198 + ]),
199 + );
200 +
201 + let collected = stream.collect_message().await.unwrap();
202 + assert_eq!(
203 + collected,
204 + CollectedAssistantMessage {
205 + message: Default::default(),
206 + prompt_token_count: 2,
207 + prompt_token_ids: vec![10, 11].into(),
208 + prompt_logprobs: Some(DecodedPromptLogprobs {
209 + first_token_id: 0,
210 + first_token: "o".to_string(),
211 + scored_positions: vec![DecodedPositionLogprobs {
212 + entries: vec![DecodedTokenLogprob {
213 + token_id: 0,
214 + token: "p".to_string(),
215 + logprob: -0.1,
216 + rank: 1,
217 + }],
218 + }],
219 + }),
220 + logprobs: Some(DecodedLogprobs {
221 + positions: vec![DecodedPositionLogprobs {
222 + entries: vec![DecodedTokenLogprob {
223 + token_id: 0,
224 + token: "a".to_string(),
225 + logprob: -0.2,
226 + rank: 1,
227 + }],
228 + }],
229 + }),
230 + token_ids: vec![],
231 + output_token_count: 1,
232 + finish_reason: FinishReason::stop_eos(),
233 + kv_transfer_params: None,
234 + }
235 + );
236 + }
237 +}
rust/src/chat/tests/chat.rs new
+1671
@@ -0,0 +1,1671 @@
1 +use std::collections::BTreeSet;
2 +use std::fmt;
3 +use std::sync::Arc;
4 +use std::time::Duration;
5 +
6 +use futures::StreamExt as _;
7 +use tokio::time::timeout;
8 +use vllm_chat::{
9 + AssistantBlockKind, AssistantContentBlock, AssistantMessageExt as _, ChatBackend, ChatEvent,
10 + ChatLlm, ChatMessage, ChatRenderer, ChatRequest, ChatRole, ChatTextBackend, ChatTool,
11 + ChatToolChoice, DefaultChatOutputProcessor, DynChatOutputProcessor, DynChatRenderer,
12 + FinishReason, GenerationPromptMode, NewChatOutputProcessorOptions, ParserSelection,
13 + RenderedPrompt, SamplingParams,
14 +};
15 +use vllm_engine_core_client::protocol::logprobs::{
16 + Logprobs, MaybeWireLogprobs, PositionLogprobs, TokenLogprob,
17 +};
18 +use vllm_engine_core_client::protocol::{
19 + EngineCoreFinishReason, EngineCoreOutput, EngineCoreOutputs, EngineCoreRequest, StopReason,
20 +};
21 +use vllm_engine_core_client::test_utils::{IpcNamespace, spawn_mock_engine_task};
22 +use vllm_engine_core_client::{EngineCoreClient, EngineCoreClientConfig};
23 +use vllm_llm::Llm;
24 +use vllm_text::tokenizer::{DynTokenizer, Tokenizer};
25 +use vllm_text::{
26 + DecodedLogprobs, DecodedPositionLogprobs, DecodedPromptLogprobs, DecodedTokenLogprob, Prompt,
27 + TextBackend,
28 +};
29 +use zeromq::prelude::{SocketRecv, SocketSend};
30 +use zeromq::{DealerSocket, PushSocket, ZmqMessage};
31 +
32 +const SPECIAL_STOP_TOKEN_ID: u32 = 256;
33 +
34 +fn request_output(
35 + request_id: &str,
36 + new_token_ids: Vec<u32>,
37 + finish_reason: Option<EngineCoreFinishReason>,
38 + stop_reason: Option<StopReason>,
39 +) -> EngineCoreOutput {
40 + EngineCoreOutput {
41 + request_id: request_id.to_string(),
42 + new_token_ids,
43 + new_logprobs: None,
44 + new_prompt_logprobs_tensors: None,
45 + pooling_output: None,
46 + finish_reason,
47 + stop_reason,
48 + events: None,
49 + kv_transfer_params: None,
50 + trace_headers: None,
51 + prefill_stats: None,
52 + routed_experts: None,
53 + num_nans_in_logits: 0,
54 + }
55 +}
56 +
57 +fn request_output_with_logprobs(
58 + request_id: &str,
59 + new_token_ids: Vec<u32>,
60 + finish_reason: Option<EngineCoreFinishReason>,
61 + stop_reason: Option<StopReason>,
62 + new_logprobs: Option<Logprobs>,
63 + new_prompt_logprobs_tensors: Option<Logprobs>,
64 +) -> EngineCoreOutput {
65 + EngineCoreOutput {
66 + request_id: request_id.to_string(),
67 + new_token_ids,
68 + new_logprobs: new_logprobs.map(MaybeWireLogprobs::Direct),
69 + new_prompt_logprobs_tensors: new_prompt_logprobs_tensors.map(MaybeWireLogprobs::Direct),
70 + pooling_output: None,
71 + finish_reason,
72 + stop_reason,
73 + events: None,
74 + kv_transfer_params: None,
75 + trace_headers: None,
76 + prefill_stats: None,
77 + routed_experts: None,
78 + num_nans_in_logits: 0,
79 + }
80 +}
81 +
82 +fn sample_logprobs_for_token(token_id: u32, alternate_token_id: u32) -> Logprobs {
83 + Logprobs {
84 + positions: vec![PositionLogprobs {
85 + entries: vec![
86 + TokenLogprob {
87 + token_id,
88 + logprob: -0.1,
89 + rank: 1,
90 + },
91 + TokenLogprob {
92 + token_id: alternate_token_id,
93 + logprob: -0.2,
94 + rank: 2,
95 + },
96 + ],
97 + }],
98 + }
99 +}
100 +
101 +fn prompt_logprobs_for_hi() -> Logprobs {
102 + Logprobs {
103 + positions: vec![PositionLogprobs {
104 + entries: vec![
105 + TokenLogprob {
106 + token_id: b'i' as u32,
107 + logprob: -0.3,
108 + rank: 1,
109 + },
110 + TokenLogprob {
111 + token_id: b'!' as u32,
112 + logprob: -0.4,
113 + rank: 2,
114 + },
115 + ],
116 + }],
117 + }
118 +}
119 +
120 +fn bytes_to_token_ids(bytes: &[u8]) -> Vec<u32> {
121 + bytes.iter().map(|byte| u32::from(*byte)).collect()
122 +}
123 +
124 +fn bytes_with_special_stop_token(bytes: &[u8]) -> Vec<u32> {
125 + let mut token_ids = bytes_to_token_ids(bytes);
126 + token_ids.push(SPECIAL_STOP_TOKEN_ID);
127 + token_ids
128 +}
129 +
130 +async fn send_outputs(push: &mut PushSocket, outputs: EngineCoreOutputs) {
131 + push.send(ZmqMessage::from(rmp_serde::to_vec_named(&outputs).unwrap()))
132 + .await
133 + .unwrap();
134 +}
135 +
136 +async fn recv_engine_message(dealer: &mut DealerSocket) -> Vec<bytes::Bytes> {
137 + dealer.recv().await.unwrap().into_vec()
138 +}
139 +
140 +async fn connect_chat_llm_with_ipc(
141 + config: EngineCoreClientConfig,
142 + ipc: &IpcNamespace,
143 + backend: Arc<dyn ChatTextBackend>,
144 +) -> ChatLlm {
145 + let client = EngineCoreClient::connect(config.with_local_input_output_addresses(
146 + Some(ipc.input_endpoint()),
147 + Some(ipc.output_endpoint()),
148 + ))
149 + .await
150 + .unwrap();
151 + ChatLlm::from_shared_backend(
152 + Llm::new(client).with_request_id_randomization(false),
153 + backend,
154 + )
155 +}
156 +
157 +#[derive(Clone)]
158 +struct FakeChatBackend {
159 + has_template: bool,
160 + model_id: String,
161 +}
162 +
163 +#[derive(Debug)]
164 +struct FakeChatTokenizer;
165 +
166 +impl Tokenizer for FakeChatTokenizer {
167 + fn encode(&self, text: &str, _add_special_tokens: bool) -> vllm_tokenizer::Result<Vec<u32>> {
168 + Ok(text.bytes().map(u32::from).collect())
169 + }
170 +
171 + fn decode(
172 + &self,
173 + token_ids: &[u32],
174 + skip_special_tokens: bool,
175 + ) -> vllm_tokenizer::Result<String> {
176 + let bytes = token_ids
177 + .iter()
178 + .filter_map(|id| {
179 + if skip_special_tokens && *id == SPECIAL_STOP_TOKEN_ID {
180 + None
181 + } else {
182 + Some(*id as u8)
183 + }
184 + })
185 + .collect::<Vec<_>>();
186 + Ok(String::from_utf8_lossy(&bytes).into_owned())
187 + }
188 +
189 + fn token_to_id(&self, token: &str) -> Option<u32> {
190 + match token {
191 + "<think>" => Some(0xF001),
192 + "</think>" => Some(0xF002),
193 + "<|START_THINKING|>" => Some(0xF003),
194 + "<|END_THINKING|>" => Some(0xF004),
195 + "◁think▷" => Some(0xF005),
196 + "◁/think▷" => Some(0xF006),
197 + _ => None,
198 + }
199 + }
200 +}
201 +
202 +impl fmt::Debug for FakeChatBackend {
203 + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
204 + f.debug_struct("FakeChatBackend").finish_non_exhaustive()
205 + }
206 +}
207 +
208 +impl FakeChatBackend {
209 + fn new() -> Self {
210 + Self {
211 + has_template: true,
212 + model_id: "test-model".to_string(),
213 + }
214 + }
215 +
216 + fn without_template() -> Self {
217 + Self {
218 + has_template: false,
219 + model_id: "test-model".to_string(),
220 + }
221 + }
222 +
223 + fn with_model_id(model_id: impl Into<String>) -> Self {
224 + Self {
225 + has_template: true,
226 + model_id: model_id.into(),
227 + }
228 + }
229 +}
230 +
231 +impl TextBackend for FakeChatBackend {
232 + fn tokenizer(&self) -> DynTokenizer {
233 + Arc::new(FakeChatTokenizer)
234 + }
235 +
236 + fn model_id(&self) -> &str {
237 + &self.model_id
238 + }
239 +}
240 +
241 +impl ChatBackend for FakeChatBackend {
242 + fn chat_renderer(&self) -> DynChatRenderer {
243 + Arc::new(self.clone())
244 + }
245 +
246 + fn new_chat_output_processor(
247 + &self,
248 + request: &mut ChatRequest,
249 + options: NewChatOutputProcessorOptions<'_>,
250 + ) -> vllm_chat::Result<DynChatOutputProcessor> {
251 + Ok(Box::new(DefaultChatOutputProcessor::new(
252 + request,
253 + &self.model_id,
254 + self.tokenizer(),
255 + options.tool_call_parser,
256 + options.reasoning_parser,
257 + )?))
258 + }
259 +}
260 +
261 +impl ChatRenderer for FakeChatBackend {
262 + fn render(&self, request: &ChatRequest) -> vllm_chat::Result<RenderedPrompt> {
263 + if !self.has_template {
264 + return Err(vllm_chat::Error::MissingChatTemplate);
265 + }
266 +
267 + let mut prompt = String::new();
268 + for message in &request.messages {
269 + prompt.push_str(message.role().as_str());
270 + prompt.push_str(": ");
271 + prompt.push_str(&message.text_content()?);
272 + prompt.push('\n');
273 + }
274 + if request.chat_options.add_generation_prompt() {
275 + prompt.push_str("assistant:");
276 + }
277 +
278 + Ok(RenderedPrompt {
279 + prompt: Prompt::Text(prompt),
280 + })
281 + }
282 +}
283 +
284 +#[derive(Clone, Debug)]
285 +struct FailingDecodeBackend {
286 + inner: FakeChatBackend,
287 +}
288 +
289 +#[derive(Debug)]
290 +struct FailingDecodeTokenizer;
291 +
292 +impl Tokenizer for FailingDecodeTokenizer {
293 + fn encode(&self, text: &str, add_special_tokens: bool) -> vllm_tokenizer::Result<Vec<u32>> {
294 + FakeChatTokenizer.encode(text, add_special_tokens)
295 + }
296 +
297 + fn decode(
298 + &self,
299 + token_ids: &[u32],
300 + skip_special_tokens: bool,
301 + ) -> vllm_tokenizer::Result<String> {
302 + if token_ids.contains(&(b'i' as u32)) {
303 + return Err(vllm_tokenizer::TokenizerError("decode failed".to_string()));
304 + }
305 + FakeChatTokenizer.decode(token_ids, skip_special_tokens)
306 + }
307 +
308 + fn token_to_id(&self, token: &str) -> Option<u32> {
309 + FakeChatTokenizer.token_to_id(token)
310 + }
311 +}
312 +
313 +impl TextBackend for FailingDecodeBackend {
314 + fn tokenizer(&self) -> DynTokenizer {
315 + Arc::new(FailingDecodeTokenizer)
316 + }
317 +
318 + fn model_id(&self) -> &str {
319 + self.inner.model_id()
320 + }
321 +}
322 +
323 +impl ChatBackend for FailingDecodeBackend {
324 + fn chat_renderer(&self) -> DynChatRenderer {
325 + Arc::new(self.clone())
326 + }
327 +
328 + fn new_chat_output_processor(
329 + &self,
330 + _request: &mut ChatRequest,
331 + _options: NewChatOutputProcessorOptions<'_>,
332 + ) -> vllm_chat::Result<DynChatOutputProcessor> {
333 + Ok(Box::new(DefaultChatOutputProcessor::plain_text_only()))
334 + }
335 +}
336 +
337 +impl ChatRenderer for FailingDecodeBackend {
338 + fn render(&self, request: &ChatRequest) -> vllm_chat::Result<RenderedPrompt> {
339 + self.inner.render(request)
340 + }
341 +}
342 +
343 +/// Skip `LogprobsDelta` events that carry only token_ids (no logprobs),
344 +/// returning the next semantically interesting event.
345 +async fn next_semantic<S>(stream: &mut S) -> Option<Result<ChatEvent, vllm_chat::Error>>
346 +where
347 + S: futures::Stream<Item = Result<ChatEvent, vllm_chat::Error>> + Unpin,
348 +{
349 + loop {
350 + match stream.next().await {
351 + Some(Ok(ChatEvent::LogprobsDelta { logprobs: None, .. })) => continue,
352 + other => return other,
353 + }
354 + }
355 +}
356 +
357 +fn sample_request(request_id: &str) -> ChatRequest {
358 + ChatRequest {
359 + messages: vec![
360 + ChatMessage::text(ChatRole::System, "You are terse."),
361 + ChatMessage::text(ChatRole::User, "Say hi"),
362 + ],
363 + sampling_params: SamplingParams {
364 + max_tokens: Some(8),
365 + ..Default::default()
366 + },
367 + request_id: request_id.to_string(),
368 + ..ChatRequest::for_test()
369 + }
370 +}
371 +
372 +fn sample_tool_request(request_id: &str) -> ChatRequest {
373 + let mut request = sample_request(request_id);
374 + request.tools = vec![ChatTool {
375 + name: "get_weather".to_string(),
376 + description: Some("Get weather".to_string()),
377 + parameters: serde_json::json!({
378 + "type": "object",
379 + "properties": {"city": {"type": "string"}},
380 + "required": ["city"],
381 + }),
382 + strict: None,
383 + }];
384 + request.tool_choice = ChatToolChoice::Auto;
385 + request
386 +}
387 +
388 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
389 +async fn chat_streams_text_events() {
390 + let ipc = IpcNamespace::new().unwrap();
391 + let handshake_address = ipc.handshake_endpoint();
392 + let engine_id = b"engine-chat".to_vec();
393 +
394 + let (shutdown_tx, engine_task) = spawn_mock_engine_task(
395 + handshake_address.clone(),
396 + engine_id.clone(),
397 + |dealer, push| {
398 + Box::pin(async move {
399 + let add = recv_engine_message(dealer).await;
400 + assert_eq!(add[0].as_ref(), &[0x00]);
401 + let request: EngineCoreRequest = rmp_serde::from_slice(&add[1]).unwrap();
402 + assert_eq!(request.request_id, "chat-1");
403 + // more fields here in the future
404 + assert_eq!(
405 + String::from_utf8(
406 + request
407 + .prompt_token_ids
408 + .clone()
409 + .unwrap()
410 + .into_iter()
411 + .map(|id| id as u8)
412 + .collect()
413 + )
414 + .unwrap(),
415 + "system: You are terse.\nuser: Say hi\nassistant:"
416 + );
417 + send_outputs(
418 + push,
419 + EngineCoreOutputs {
420 + outputs: vec![
421 + request_output("chat-1", vec![b'H' as u32], None, None),
422 + request_output(
423 + "chat-1",
424 + vec![b'i' as u32, b'!' as u32],
425 + Some(EngineCoreFinishReason::Stop),
426 + Some(StopReason::TokenId(b'!' as u32)),
427 + ),
428 + ],
429 + finished_requests: Some(BTreeSet::from(["chat-1".to_string()])),
430 + ..Default::default()
431 + },
432 + )
433 + .await;
434 + })
435 + },
436 + );
437 +
438 + let backend: Arc<dyn ChatTextBackend> = Arc::new(FakeChatBackend::new());
439 + let chat = connect_chat_llm_with_ipc(
440 + EngineCoreClientConfig::new_single(handshake_address).with_model_name("test-model"),
441 + &ipc,
442 + backend,
443 + )
444 + .await;
445 +
446 + let mut stream = chat.chat(sample_request("chat-1")).await.unwrap();
447 +
448 + match next_semantic(&mut stream).await.unwrap().unwrap() {
449 + ChatEvent::Start {
450 + prompt_token_ids,
451 + prompt_logprobs: None,
452 + } => {
453 + assert_eq!(
454 + prompt_token_ids.len(),
455 + "system: You are terse.\nuser: Say hi\nassistant:".len()
456 + );
457 + assert!(!prompt_token_ids.is_empty());
458 + }
459 + other => panic!("expected Start, got {other:?}"),
460 + }
461 + assert_eq!(
462 + next_semantic(&mut stream).await.unwrap().unwrap(),
463 + ChatEvent::BlockStart {
464 + index: 0,
465 + kind: AssistantBlockKind::Text,
466 + }
467 + );
468 + assert_eq!(
469 + next_semantic(&mut stream).await.unwrap().unwrap(),
470 + ChatEvent::BlockDelta {
471 + index: 0,
472 + kind: AssistantBlockKind::Text,
473 + delta: "H".to_string(),
474 + }
475 + );
476 + assert_eq!(
477 + next_semantic(&mut stream).await.unwrap().unwrap(),
478 + ChatEvent::BlockDelta {
479 + index: 0,
480 + kind: AssistantBlockKind::Text,
481 + delta: "i".to_string(),
482 + }
483 + );
484 + assert_eq!(
485 + next_semantic(&mut stream).await.unwrap().unwrap(),
486 + ChatEvent::BlockEnd {
487 + index: 0,
488 + block: AssistantContentBlock::Text {
489 + text: "Hi".to_string(),
490 + },
491 + }
492 + );
493 +
494 + match next_semantic(&mut stream).await {
495 + Some(Ok(ChatEvent::Done {
496 + message,
497 + output_token_count,
498 + finish_reason,
499 + ..
500 + })) => {
501 + assert_eq!(message.text(), "Hi");
502 + assert_eq!(output_token_count, 3);
503 + assert_eq!(
504 + finish_reason,
505 + FinishReason::Stop(Some(StopReason::TokenId(b'!' as u32)))
506 + );
507 + }
508 + other => panic!("unexpected final event: {other:?}"),
509 + }
510 + assert!(next_semantic(&mut stream).await.is_none());
511 +
512 + let _ = shutdown_tx.send(());
513 + engine_task.await.unwrap();
514 + chat.shutdown().await.unwrap();
515 +}
516 +
517 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
518 +async fn chat_stream_waits_for_complete_utf8_before_emitting() {
519 + let ipc = IpcNamespace::new().unwrap();
520 + let handshake_address = ipc.handshake_endpoint();
521 + let engine_id = b"engine-chat-utf8".to_vec();
522 +
523 + let (shutdown_tx, engine_task) = spawn_mock_engine_task(
524 + handshake_address.clone(),
525 + engine_id.clone(),
526 + |dealer, push| {
527 + Box::pin(async move {
528 + let _ = recv_engine_message(dealer).await;
529 + send_outputs(
530 + push,
531 + EngineCoreOutputs {
532 + outputs: vec![
533 + request_output("chat-utf8", bytes_to_token_ids(&[0xe4]), None, None),
534 + request_output(
535 + "chat-utf8",
536 + bytes_to_token_ids(&[0xbd, 0xa0, b'!']),
537 + Some(EngineCoreFinishReason::Stop),
538 + Some(StopReason::TokenId(b'!' as u32)),
539 + ),
540 + ],
541 + finished_requests: Some(BTreeSet::from(["chat-utf8".to_string()])),
542 + ..Default::default()
543 + },
544 + )
545 + .await;
546 + })
547 + },
548 + );
549 +
550 + let backend: Arc<dyn ChatTextBackend> = Arc::new(FakeChatBackend::new());
551 + let chat = connect_chat_llm_with_ipc(
552 + EngineCoreClientConfig::new_single(handshake_address),
553 + &ipc,
554 + backend,
555 + )
556 + .await;
557 +
558 + let mut stream = chat.chat(sample_request("chat-utf8")).await.unwrap();
559 +
560 + assert!(matches!(
561 + next_semantic(&mut stream).await,
562 + Some(Ok(ChatEvent::Start {
563 + prompt_logprobs: None,
564 + ..
565 + }))
566 + ));
567 + assert_eq!(
568 + next_semantic(&mut stream).await.unwrap().unwrap(),
569 + ChatEvent::BlockStart {
570 + index: 0,
571 + kind: AssistantBlockKind::Text,
572 + }
573 + );
574 + assert_eq!(
575 + next_semantic(&mut stream).await.unwrap().unwrap(),
576 + ChatEvent::BlockDelta {
577 + index: 0,
578 + kind: AssistantBlockKind::Text,
579 + delta: "你".to_string(),
580 + }
581 + );
582 + assert_eq!(
583 + next_semantic(&mut stream).await.unwrap().unwrap(),
584 + ChatEvent::BlockEnd {
585 + index: 0,
586 + block: AssistantContentBlock::Text {
587 + text: "你".to_string(),
588 + },
589 + }
590 + );
591 +
592 + match next_semantic(&mut stream).await {
593 + Some(Ok(ChatEvent::Done {
594 + message,
595 + output_token_count,
596 + ..
597 + })) => {
598 + assert_eq!(message.text(), "你");
599 + assert_eq!(output_token_count, 4);
600 + }
601 + other => panic!("unexpected final event: {other:?}"),
602 + }
603 +
604 + let _ = shutdown_tx.send(());
605 + engine_task.await.unwrap();
606 + chat.shutdown().await.unwrap();
607 +}
608 +
609 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
610 +async fn chat_stream_flushes_held_text_on_finish() {
611 + let ipc = IpcNamespace::new().unwrap();
612 + let handshake_address = ipc.handshake_endpoint();
613 + let engine_id = b"engine-chat-final-flush".to_vec();
614 +
615 + let (shutdown_tx, engine_task) = spawn_mock_engine_task(
616 + handshake_address.clone(),
617 + engine_id.clone(),
618 + |dealer, push| {
619 + Box::pin(async move {
620 + let _ = recv_engine_message(dealer).await;
621 + send_outputs(
622 + push,
623 + EngineCoreOutputs {
624 + outputs: vec![request_output(
625 + "chat-final-flush",
626 + bytes_to_token_ids(b"ok st"),
627 + Some(EngineCoreFinishReason::Length),
628 + None,
629 + )],
630 + finished_requests: Some(BTreeSet::from(["chat-final-flush".to_string()])),
631 + ..Default::default()
632 + },
633 + )
634 + .await;
635 + })
636 + },
637 + );
638 +
639 + let backend: Arc<dyn ChatTextBackend> = Arc::new(FakeChatBackend::new());
640 + let chat = connect_chat_llm_with_ipc(
641 + EngineCoreClientConfig::new_single(handshake_address),
642 + &ipc,
643 + backend,
644 + )
645 + .await;
646 +
647 + let mut stream = chat.chat(sample_request("chat-final-flush")).await.unwrap();
648 +
649 + assert!(matches!(
650 + next_semantic(&mut stream).await,
651 + Some(Ok(ChatEvent::Start {
652 + prompt_logprobs: None,
653 + ..
654 + }))
655 + ));
656 + assert_eq!(
657 + next_semantic(&mut stream).await.unwrap().unwrap(),
658 + ChatEvent::BlockStart {
659 + index: 0,
660 + kind: AssistantBlockKind::Text,
661 + }
662 + );
663 + assert_eq!(
664 + next_semantic(&mut stream).await.unwrap().unwrap(),
665 + ChatEvent::BlockDelta {
666 + index: 0,
667 + kind: AssistantBlockKind::Text,
668 + delta: "ok st".to_string(),
669 + }
670 + );
671 + assert_eq!(
672 + next_semantic(&mut stream).await.unwrap().unwrap(),
673 + ChatEvent::BlockEnd {
674 + index: 0,
675 + block: AssistantContentBlock::Text {
676 + text: "ok st".to_string(),
677 + },
678 + }
679 + );
680 +
681 + match next_semantic(&mut stream).await {
682 + Some(Ok(ChatEvent::Done {
683 + message,
684 + output_token_count,
685 + finish_reason,
686 + ..
687 + })) => {
688 + assert_eq!(message.text(), "ok st");
689 + assert_eq!(output_token_count, 5);
690 + assert_eq!(finish_reason, FinishReason::Length);
691 + }
692 + other => panic!("unexpected final event: {other:?}"),
693 + }
694 +
695 + let _ = shutdown_tx.send(());
696 + engine_task.await.unwrap();
697 + chat.shutdown().await.unwrap();
698 +}
699 +
700 +#[test]
701 +fn chat_request_rejects_conflicting_generation_modes() {
702 + let mut request = sample_request("chat-2");
703 + request.chat_options.generation_prompt_mode = GenerationPromptMode::ContinueFinalAssistant;
704 + let error = request.validate().unwrap_err();
705 +
706 + assert!(matches!(
707 + error,
708 + vllm_chat::Error::ContinueFinalAssistantWithoutFinalAssistant
709 + ));
710 +}
711 +
712 +#[test]
713 +fn chat_request_accepts_continue_final_assistant_mode_with_final_assistant() {
714 + let mut request = sample_request("chat-2b");
715 + request.messages = vec![ChatMessage::assistant_text("hello")];
716 + request.chat_options.generation_prompt_mode = GenerationPromptMode::ContinueFinalAssistant;
717 +
718 + request.validate().unwrap();
719 +}
720 +
721 +#[test]
722 +fn backend_requires_a_template() {
723 + let request = sample_request("chat-3");
724 + let backend = FakeChatBackend::without_template();
725 + let error = backend.chat_renderer().render(&request).unwrap_err();
726 + assert!(matches!(error, vllm_chat::Error::MissingChatTemplate));
727 +}
728 +
729 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
730 +async fn chat_stream_reports_decode_failure_as_error_event() {
731 + let ipc = IpcNamespace::new().unwrap();
732 + let handshake_address = ipc.handshake_endpoint();
733 + let engine_id = b"engine-chat-decode-fail".to_vec();
734 +
735 + let (shutdown_tx, engine_task) = spawn_mock_engine_task(
736 + handshake_address.clone(),
737 + engine_id.clone(),
738 + |dealer, push| {
739 + Box::pin(async move {
740 + let _ = recv_engine_message(dealer).await;
741 + send_outputs(
742 + push,
743 + EngineCoreOutputs {
744 + outputs: vec![request_output("chat-4", vec![b'i' as u32], None, None)],
745 + ..Default::default()
746 + },
747 + )
748 + .await;
749 + })
750 + },
751 + );
752 +
753 + let backend: Arc<dyn ChatTextBackend> = Arc::new(FailingDecodeBackend {
754 + inner: FakeChatBackend::new(),
755 + });
756 + let chat = connect_chat_llm_with_ipc(
757 + EngineCoreClientConfig::new_single(handshake_address),
758 + &ipc,
759 + backend,
760 + )
761 + .await;
762 +
763 + let mut stream = chat.chat(sample_request("chat-4")).await.unwrap();
764 + assert_eq!(stream.request_id(), "chat-4");
765 + assert!(matches!(
766 + next_semantic(&mut stream).await,
767 + Some(Ok(ChatEvent::Start {
768 + prompt_logprobs: None,
769 + ..
770 + }))
771 + ));
772 +
773 + match timeout(Duration::from_secs(2), stream.next()).await.unwrap() {
774 + Some(Err(vllm_chat::Error::Text(vllm_text::Error::Tokenizer(message)))) => {
775 + assert_eq!(message, "decode failed");
776 + }
777 + other => panic!("unexpected event after close: {other:?}"),
778 + }
779 +
780 + let _ = shutdown_tx.send(());
781 + engine_task.await.unwrap();
782 + chat.shutdown().await.unwrap();
783 +}
784 +
785 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
786 +async fn chat_stream_preserves_terminal_stop_token_when_requested() {
787 + let ipc = IpcNamespace::new().unwrap();
788 + let handshake_address = ipc.handshake_endpoint();
789 + let engine_id = b"engine-chat-include-stop".to_vec();
790 +
791 + let (shutdown_tx, engine_task) = spawn_mock_engine_task(
792 + handshake_address.clone(),
793 + engine_id.clone(),
794 + |dealer, push| {
795 + Box::pin(async move {
796 + let _ = recv_engine_message(dealer).await;
797 + send_outputs(
798 + push,
799 + EngineCoreOutputs {
800 + outputs: vec![request_output(
801 + "chat-include-stop",
802 + vec![b'H' as u32, b'i' as u32, b'!' as u32],
803 + Some(EngineCoreFinishReason::Stop),
804 + Some(StopReason::TokenId(b'!' as u32)),
805 + )],
806 + finished_requests: Some(BTreeSet::from(["chat-include-stop".to_string()])),
807 + ..Default::default()
808 + },
809 + )
810 + .await;
811 + })
812 + },
813 + );
814 +
815 + let backend: Arc<dyn ChatTextBackend> = Arc::new(FakeChatBackend::new());
816 + let chat = connect_chat_llm_with_ipc(
817 + EngineCoreClientConfig::new_single(handshake_address),
818 + &ipc,
819 + backend,
820 + )
821 + .await;
822 +
823 + let mut request = sample_request("chat-include-stop");
824 + request.decode_options.include_stop_str_in_output = true;
825 + let mut stream = chat.chat(request).await.unwrap();
826 +
827 + assert!(matches!(
828 + next_semantic(&mut stream).await,
829 + Some(Ok(ChatEvent::Start {
830 + prompt_logprobs: None,
831 + ..
832 + }))
833 + ));
834 + assert_eq!(
835 + next_semantic(&mut stream).await.unwrap().unwrap(),
836 + ChatEvent::BlockStart {
837 + index: 0,
838 + kind: AssistantBlockKind::Text,
839 + }
840 + );
841 + assert_eq!(
842 + next_semantic(&mut stream).await.unwrap().unwrap(),
843 + ChatEvent::BlockDelta {
844 + index: 0,
845 + kind: AssistantBlockKind::Text,
846 + delta: "Hi!".to_string(),
847 + }
848 + );
849 + assert_eq!(
850 + next_semantic(&mut stream).await.unwrap().unwrap(),
851 + ChatEvent::BlockEnd {
852 + index: 0,
853 + block: AssistantContentBlock::Text {
854 + text: "Hi!".to_string(),
855 + },
856 + }
857 + );
858 +
859 + match next_semantic(&mut stream).await {
860 + Some(Ok(ChatEvent::Done {
861 + message,
862 + output_token_count,
863 + ..
864 + })) => {
865 + assert_eq!(message.text(), "Hi!");
866 + assert_eq!(output_token_count, 3);
867 + }
868 + other => panic!("unexpected final event: {other:?}"),
869 + }
870 +
871 + let _ = shutdown_tx.send(());
872 + engine_task.await.unwrap();
873 + chat.shutdown().await.unwrap();
874 +}
875 +
876 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
877 +async fn chat_stream_separates_reasoning_blocks_automatically() {
878 + let ipc = IpcNamespace::new().unwrap();
879 + let handshake_address = ipc.handshake_endpoint();
880 + let engine_id = b"engine-chat-reasoning".to_vec();
881 +
882 + let (shutdown_tx, engine_task) = spawn_mock_engine_task(
883 + handshake_address.clone(),
884 + engine_id.clone(),
885 + |dealer, push| {
886 + Box::pin(async move {
887 + let _ = recv_engine_message(dealer).await;
888 + send_outputs(
889 + push,
890 + EngineCoreOutputs {
891 + outputs: vec![
892 + request_output(
893 + "chat-reasoning",
894 + bytes_to_token_ids(b"<think>"),
895 + None,
896 + None,
897 + ),
898 + request_output(
899 + "chat-reasoning",
900 + bytes_to_token_ids(b"reason "),
901 + None,
902 + None,
903 + ),
904 + request_output(
905 + "chat-reasoning",
906 + bytes_to_token_ids(b"more</think>"),
907 + None,
908 + None,
909 + ),
910 + request_output(
911 + "chat-reasoning",
912 + bytes_to_token_ids(b"answer"),
913 + Some(EngineCoreFinishReason::Length),
914 + None,
915 + ),
916 + ],
917 + finished_requests: Some(BTreeSet::from(["chat-reasoning".to_string()])),
918 + ..Default::default()
919 + },
920 + )
921 + .await;
922 + })
923 + },
924 + );
925 +
926 + let backend: Arc<dyn ChatTextBackend> =
927 + Arc::new(FakeChatBackend::with_model_id("Qwen/Qwen3-0.6B"));
928 + let chat = connect_chat_llm_with_ipc(
929 + EngineCoreClientConfig::new_single(handshake_address),
930 + &ipc,
931 + backend,
932 + )
933 + .await;
934 +
935 + let mut stream = chat.chat(sample_request("chat-reasoning")).await.unwrap();
936 +
937 + assert!(matches!(
938 + next_semantic(&mut stream).await,
939 + Some(Ok(ChatEvent::Start {
940 + prompt_logprobs: None,
941 + ..
942 + }))
943 + ));
944 + assert_eq!(
945 + next_semantic(&mut stream).await.unwrap().unwrap(),
946 + ChatEvent::BlockStart {
947 + index: 0,
948 + kind: AssistantBlockKind::Reasoning,
949 + }
950 + );
951 + assert_eq!(
952 + next_semantic(&mut stream).await.unwrap().unwrap(),
953 + ChatEvent::BlockDelta {
954 + index: 0,
955 + kind: AssistantBlockKind::Reasoning,
956 + delta: "reason ".to_string(),
957 + }
958 + );
959 + assert_eq!(
960 + next_semantic(&mut stream).await.unwrap().unwrap(),
961 + ChatEvent::BlockDelta {
962 + index: 0,
963 + kind: AssistantBlockKind::Reasoning,
964 + delta: "more".to_string(),
965 + }
966 + );
967 + assert_eq!(
968 + next_semantic(&mut stream).await.unwrap().unwrap(),
969 + ChatEvent::BlockEnd {
970 + index: 0,
971 + block: AssistantContentBlock::Reasoning {
972 + text: "reason more".to_string(),
973 + },
974 + }
975 + );
976 + assert_eq!(
977 + next_semantic(&mut stream).await.unwrap().unwrap(),
978 + ChatEvent::BlockStart {
979 + index: 1,
980 + kind: AssistantBlockKind::Text,
981 + }
982 + );
983 + assert_eq!(
984 + next_semantic(&mut stream).await.unwrap().unwrap(),
985 + ChatEvent::BlockDelta {
986 + index: 1,
987 + kind: AssistantBlockKind::Text,
988 + delta: "answer".to_string(),
989 + }
990 + );
991 + assert_eq!(
992 + next_semantic(&mut stream).await.unwrap().unwrap(),
993 + ChatEvent::BlockEnd {
994 + index: 1,
995 + block: AssistantContentBlock::Text {
996 + text: "answer".to_string(),
997 + },
998 + }
999 + );
1000 +
1001 + match next_semantic(&mut stream).await {
1002 + Some(Ok(ChatEvent::Done {
1003 + message,
1004 + finish_reason,
1005 + ..
1006 + })) => {
1007 + assert_eq!(message.reasoning().unwrap(), "reason more");
1008 + assert_eq!(message.text(), "answer");
1009 + assert_eq!(finish_reason, FinishReason::Length);
1010 + }
1011 + other => panic!("unexpected final event: {other:?}"),
1012 + }
1013 +
1014 + let _ = shutdown_tx.send(());
1015 + engine_task.await.unwrap();
1016 + chat.shutdown().await.unwrap();
1017 +}
1018 +
1019 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1020 +async fn chat_collectors_return_structured_message_and_visible_text() {
1021 + let ipc = IpcNamespace::new().unwrap();
1022 + let handshake_address = ipc.handshake_endpoint();
1023 + let engine_id = b"engine-chat-collect".to_vec();
1024 +
1025 + let (shutdown_tx, engine_task) = spawn_mock_engine_task(
1026 + handshake_address.clone(),
1027 + engine_id.clone(),
1028 + |dealer, push| {
1029 + Box::pin(async move {
1030 + let _ = recv_engine_message(dealer).await;
1031 + send_outputs(
1032 + push,
1033 + EngineCoreOutputs {
1034 + outputs: vec![request_output(
1035 + "chat-collect",
1036 + bytes_to_token_ids(b"<think>inner</think>outer"),
1037 + Some(EngineCoreFinishReason::Length),
1038 + None,
1039 + )],
1040 + finished_requests: Some(BTreeSet::from(["chat-collect".to_string()])),
1041 + ..Default::default()
1042 + },
1043 + )
1044 + .await;
1045 + })
1046 + },
1047 + );
1048 +
1049 + let backend: Arc<dyn ChatTextBackend> =
1050 + Arc::new(FakeChatBackend::with_model_id("Qwen/Qwen3-0.6B"));
1051 + let chat = connect_chat_llm_with_ipc(
1052 + EngineCoreClientConfig::new_single(handshake_address.clone()),
1053 + &ipc,
1054 + backend.clone(),
1055 + )
1056 + .await;
1057 +
1058 + let message = chat
1059 + .chat(sample_request("chat-collect"))
1060 + .await
1061 + .unwrap()
1062 + .collect_message()
1063 + .await
1064 + .unwrap();
1065 + assert_eq!(message.message.reasoning().unwrap(), "inner");
1066 + assert_eq!(message.message.text(), "outer");
1067 + assert_eq!(message.finish_reason, FinishReason::Length);
1068 + assert_eq!(
1069 + message.prompt_token_count,
1070 + "system: You are terse.\nuser: Say hi\nassistant:".len()
1071 + );
1072 + assert_eq!(
1073 + message.output_token_count,
1074 + "<think>inner</think>outer".len()
1075 + );
1076 +
1077 + let _ = shutdown_tx.send(());
1078 + engine_task.await.unwrap();
1079 + chat.shutdown().await.unwrap();
1080 +}
1081 +
1082 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1083 +async fn chat_explicitly_disables_reasoning_parser() {
1084 + let ipc = IpcNamespace::new().unwrap();
1085 + let handshake_address = ipc.handshake_endpoint();
1086 + let engine_id = b"engine-chat-reasoning-disabled".to_vec();
1087 +
1088 + let (shutdown_tx, engine_task) = spawn_mock_engine_task(
1089 + handshake_address.clone(),
1090 + engine_id.clone(),
1091 + |dealer, push| {
1092 + Box::pin(async move {
1093 + let _ = recv_engine_message(dealer).await;
1094 + send_outputs(
1095 + push,
1096 + EngineCoreOutputs {
1097 + outputs: vec![
1098 + request_output(
1099 + "chat-reasoning-disabled",
1100 + bytes_to_token_ids(b"<think>"),
1101 + None,
1102 + None,
1103 + ),
1104 + request_output(
1105 + "chat-reasoning-disabled",
1106 + bytes_to_token_ids(b"reason "),
1107 + None,
1108 + None,
1109 + ),
1110 + request_output(
1111 + "chat-reasoning-disabled",
1112 + bytes_to_token_ids(b"more</think>"),
1113 + None,
1114 + None,
1115 + ),
1116 + request_output(
1117 + "chat-reasoning-disabled",
1118 + bytes_to_token_ids(b"answer"),
1119 + Some(EngineCoreFinishReason::Length),
1120 + None,
1121 + ),
1122 + ],
1123 + finished_requests: Some(BTreeSet::from([
1124 + "chat-reasoning-disabled".to_string()
1125 + ])),
1126 + ..Default::default()
1127 + },
1128 + )
1129 + .await;
1130 + })
1131 + },
1132 + );
1133 +
1134 + let backend: Arc<dyn ChatTextBackend> =
1135 + Arc::new(FakeChatBackend::with_model_id("Qwen/Qwen3-0.6B"));
1136 + let chat = connect_chat_llm_with_ipc(
1137 + EngineCoreClientConfig::new_single(handshake_address),
1138 + &ipc,
1139 + backend,
1140 + )
1141 + .await
1142 + .with_reasoning_parser(ParserSelection::None);
1143 +
1144 + let message = chat
1145 + .chat(sample_request("chat-reasoning-disabled"))
1146 + .await
1147 + .unwrap()
1148 + .collect_message()
1149 + .await
1150 + .unwrap();
1151 + assert_eq!(message.message.reasoning(), None);
1152 + assert_eq!(message.message.text(), "<think>reason more</think>answer");
1153 + assert_eq!(message.finish_reason, FinishReason::Length);
1154 +
1155 + let _ = shutdown_tx.send(());
1156 + engine_task.await.unwrap();
1157 + chat.shutdown().await.unwrap();
1158 +}
1159 +
1160 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1161 +async fn chat_stream_parses_tool_calls_automatically() {
1162 + let ipc = IpcNamespace::new().unwrap();
1163 + let handshake_address = ipc.handshake_endpoint();
1164 + let engine_id = b"engine-chat-tool".to_vec();
1165 +
1166 + let (shutdown_tx, engine_task) = spawn_mock_engine_task(
1167 + handshake_address.clone(),
1168 + engine_id.clone(),
1169 + |dealer, push| {
1170 + Box::pin(async move {
1171 + let _ = recv_engine_message(dealer).await;
1172 + send_outputs(
1173 + push,
1174 + EngineCoreOutputs {
1175 + outputs: vec![
1176 + request_output(
1177 + "chat-tool",
1178 + bytes_to_token_ids(b"<think>Need tool.</think>"),
1179 + None,
1180 + None,
1181 + ),
1182 + request_output(
1183 + "chat-tool",
1184 + bytes_to_token_ids(b"<tool_call>\n{\"name\":\"get_weather\", "),
1185 + None,
1186 + None,
1187 + ),
1188 + request_output(
1189 + "chat-tool",
1190 + bytes_to_token_ids(
1191 + b"\"arguments\":{\"city\":\"Paris\"}}\n</tool_call>",
1192 + ),
1193 + Some(EngineCoreFinishReason::Stop),
1194 + Some(StopReason::TokenId(SPECIAL_STOP_TOKEN_ID)),
1195 + ),
1196 + ],
1197 + finished_requests: Some(BTreeSet::from(["chat-tool".to_string()])),
1198 + ..Default::default()
1199 + },
1200 + )
1201 + .await;
1202 + })
1203 + },
1204 + );
1205 +
1206 + let backend: Arc<dyn ChatTextBackend> =
1207 + Arc::new(FakeChatBackend::with_model_id("Qwen/Qwen3-0.6B"));
1208 + let chat = connect_chat_llm_with_ipc(
1209 + EngineCoreClientConfig::new_single(handshake_address),
1210 + &ipc,
1211 + backend,
1212 + )
1213 + .await;
1214 + let mut stream = chat.chat(sample_tool_request("chat-tool")).await.unwrap();
1215 +
1216 + let mut saw_tool_start = false;
1217 + let mut saw_tool_args = false;
1218 + let mut saw_tool_end = false;
1219 +
1220 + while let Some(event) = stream.next().await {
1221 + match event.unwrap() {
1222 + ChatEvent::Start { .. } => {}
1223 + ChatEvent::LogprobsDelta { .. } => {}
1224 + ChatEvent::ToolCallStart { name, .. } => {
1225 + saw_tool_start = true;
1226 + assert_eq!(name, "get_weather");
1227 + }
1228 + ChatEvent::ToolCallArgumentsDelta { delta, .. } => {
1229 + saw_tool_args = true;
1230 + assert!(delta.contains("Paris"), "{delta}");
1231 + }
1232 + ChatEvent::ToolCallEnd { call, .. } => {
1233 + saw_tool_end = true;
1234 + assert_eq!(call.name, "get_weather");
1235 + assert_eq!(call.arguments, r#"{"city":"Paris"}"#);
1236 + }
1237 + ChatEvent::Done {
1238 + message,
1239 + finish_reason,
1240 + ..
1241 + } => {
1242 + assert_eq!(
1243 + finish_reason,
1244 + FinishReason::Stop(Some(StopReason::TokenId(SPECIAL_STOP_TOKEN_ID)))
1245 + );
1246 + assert_eq!(message.text(), "");
1247 + let tool_calls = message.tool_calls().collect::<Vec<_>>();
1248 + assert_eq!(tool_calls.len(), 1);
1249 + assert_eq!(tool_calls[0].name, "get_weather");
1250 + assert_eq!(tool_calls[0].arguments, r#"{"city":"Paris"}"#);
1251 + break;
1252 + }
1253 + ChatEvent::BlockStart { .. }
1254 + | ChatEvent::BlockDelta { .. }
1255 + | ChatEvent::BlockEnd { .. } => {}
1256 + }
1257 + }
1258 +
1259 + assert!(saw_tool_start);
1260 + assert!(saw_tool_args);
1261 + assert!(saw_tool_end);
1262 +
1263 + let _ = shutdown_tx.send(());
1264 + engine_task.await.unwrap();
1265 + chat.shutdown().await.unwrap();
1266 +}
1267 +
1268 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1269 +async fn chat_collect_message_preserves_tool_call_arguments_in_final_only_mode() {
1270 + let ipc = IpcNamespace::new().unwrap();
1271 + let handshake_address = ipc.handshake_endpoint();
1272 + let engine_id = b"engine-chat-final-only-tool".to_vec();
1273 +
1274 + let (shutdown_tx, engine_task) = spawn_mock_engine_task(
1275 + handshake_address.clone(),
1276 + engine_id.clone(),
1277 + |dealer, push| {
1278 + Box::pin(async move {
1279 + let _ = recv_engine_message(dealer).await;
1280 + send_outputs(
1281 + push,
1282 + EngineCoreOutputs {
1283 + outputs: vec![
1284 + request_output(
1285 + "chat-final-only-tool",
1286 + bytes_to_token_ids(b"<think>Need tool.</think>"),
1287 + None,
1288 + None,
1289 + ),
1290 + request_output(
1291 + "chat-final-only-tool",
1292 + bytes_to_token_ids(b"<tool_call>\n{\"name\":\"get_weather\", "),
1293 + None,
1294 + None,
1295 + ),
1296 + request_output(
1297 + "chat-final-only-tool",
1298 + bytes_with_special_stop_token(
1299 + b"\"arguments\":{\"city\":\"Paris\"}}\n</tool_call>",
1300 + ),
1301 + Some(EngineCoreFinishReason::Stop),
1302 + Some(StopReason::TokenId(SPECIAL_STOP_TOKEN_ID)),
1303 + ),
1304 + ],
1305 + finished_requests: Some(BTreeSet::from([
1306 + "chat-final-only-tool".to_string()
1307 + ])),
1308 + ..Default::default()
1309 + },
1310 + )
1311 + .await;
1312 + })
1313 + },
1314 + );
1315 +
1316 + let backend: Arc<dyn ChatTextBackend> =
1317 + Arc::new(FakeChatBackend::with_model_id("Qwen/Qwen3-0.6B"));
1318 + let chat = connect_chat_llm_with_ipc(
1319 + EngineCoreClientConfig::new_single(handshake_address),
1320 + &ipc,
1321 + backend,
1322 + )
1323 + .await;
1324 + let mut request = sample_tool_request("chat-final-only-tool");
1325 + request.intermediate = false;
1326 +
1327 + let message = chat.chat(request).await.unwrap().collect_message().await.unwrap();
1328 +
1329 + assert_eq!(
1330 + message.finish_reason,
1331 + FinishReason::Stop(Some(StopReason::TokenId(SPECIAL_STOP_TOKEN_ID)))
1332 + );
1333 + assert_eq!(message.message.tool_calls().count(), 1);
1334 + assert_eq!(
1335 + message.message.tool_calls().next().unwrap().arguments,
1336 + r#"{"city":"Paris"}"#
1337 + );
1338 +
1339 + let _ = shutdown_tx.send(());
1340 + engine_task.await.unwrap();
1341 + chat.shutdown().await.unwrap();
1342 +}
1343 +
1344 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1345 +async fn chat_stream_and_collect_preserve_prompt_and_sample_logprobs() {
1346 + let ipc = IpcNamespace::new().unwrap();
1347 + let handshake_address = ipc.handshake_endpoint();
1348 + let engine_id = b"engine-chat-logprobs".to_vec();
1349 +
1350 + let (shutdown_tx, engine_task) = spawn_mock_engine_task(
1351 + handshake_address.clone(),
1352 + engine_id.clone(),
1353 + |dealer, push| {
1354 + Box::pin(async move {
1355 + for _ in 0..2 {
1356 + let add = recv_engine_message(dealer).await;
1357 + let request: EngineCoreRequest = rmp_serde::from_slice(&add[1]).unwrap();
1358 + send_outputs(
1359 + push,
1360 + EngineCoreOutputs {
1361 + outputs: vec![
1362 + request_output_with_logprobs(
1363 + &request.request_id,
1364 + vec![b'H' as u32],
1365 + None,
1366 + None,
1367 + Some(sample_logprobs_for_token(b'H' as u32, b'h' as u32)),
1368 + Some(prompt_logprobs_for_hi()),
1369 + ),
1370 + request_output_with_logprobs(
1371 + &request.request_id,
1372 + vec![b'i' as u32],
1373 + Some(EngineCoreFinishReason::Length),
1374 + None,
1375 + Some(sample_logprobs_for_token(b'i' as u32, b'I' as u32)),
1376 + None,
1377 + ),
1378 + ],
1379 + finished_requests: Some(BTreeSet::from([request.request_id])),
1380 + ..Default::default()
1381 + },
1382 + )
1383 + .await;
1384 + }
1385 + })
1386 + },
1387 + );
1388 +
1389 + let backend: Arc<dyn ChatTextBackend> = Arc::new(FakeChatBackend::new());
1390 + let chat = connect_chat_llm_with_ipc(
1391 + EngineCoreClientConfig::new_single(handshake_address.clone()),
1392 + &ipc,
1393 + backend,
1394 + )
1395 + .await;
1396 +
1397 + let mut request = sample_request("chat-logprobs");
1398 + request.sampling_params.logprobs = Some(1);
1399 + request.sampling_params.prompt_logprobs = Some(1);
1400 +
1401 + let mut stream = chat.chat(request.clone()).await.unwrap();
1402 + match next_semantic(&mut stream).await.unwrap().unwrap() {
1403 + ChatEvent::Start {
1404 + prompt_token_ids,
1405 + prompt_logprobs,
1406 + } => {
1407 + assert_eq!(
1408 + prompt_token_ids.len(),
1409 + "system: You are terse.\nuser: Say hi\nassistant:".len()
1410 + );
1411 + assert!(!prompt_token_ids.is_empty());
1412 + assert_eq!(
1413 + prompt_logprobs,
1414 + Some(DecodedPromptLogprobs {
1415 + first_token_id: b's' as u32,
1416 + first_token: "s".to_string(),
1417 + scored_positions: vec![DecodedPositionLogprobs {
1418 + entries: vec![
1419 + DecodedTokenLogprob {
1420 + token_id: b'i' as u32,
1421 + token: "i".to_string(),
1422 + logprob: -0.3,
1423 + rank: 1,
1424 + },
1425 + DecodedTokenLogprob {
1426 + token_id: b'!' as u32,
1427 + token: "!".to_string(),
1428 + logprob: -0.4,
1429 + rank: 1,
1430 + },
1431 + ],
1432 + }],
1433 + })
1434 + );
1435 + }
1436 + other => panic!("expected Start, got {other:?}"),
1437 + }
1438 + assert_eq!(
1439 + next_semantic(&mut stream).await.unwrap().unwrap(),
1440 + ChatEvent::BlockStart {
1441 + index: 0,
1442 + kind: AssistantBlockKind::Text,
1443 + }
1444 + );
1445 + assert_eq!(
1446 + next_semantic(&mut stream).await.unwrap().unwrap(),
1447 + ChatEvent::BlockDelta {
1448 + index: 0,
1449 + kind: AssistantBlockKind::Text,
1450 + delta: "H".to_string(),
1451 + }
1452 + );
1453 + assert_eq!(
1454 + next_semantic(&mut stream).await.unwrap().unwrap(),
1455 + ChatEvent::LogprobsDelta {
1456 + logprobs: Some(DecodedLogprobs {
1457 + positions: vec![DecodedPositionLogprobs {
1458 + entries: vec![
1459 + DecodedTokenLogprob {
1460 + token_id: b'H' as u32,
1461 + token: "H".to_string(),
1462 + logprob: -0.1,
1463 + rank: 1,
1464 + },
1465 + DecodedTokenLogprob {
1466 + token_id: b'h' as u32,
1467 + token: "h".to_string(),
1468 + logprob: -0.2,
1469 + rank: 1,
1470 + },
1471 + ],
1472 + }],
1473 + }),
1474 + token_ids: vec![b'H' as u32],
1475 + }
1476 + );
1477 + while !matches!(
1478 + next_semantic(&mut stream).await,
1479 + Some(Ok(ChatEvent::Done { .. }))
1480 + ) {}
1481 +
1482 + request.request_id = "chat-logprobs-collect".to_string();
1483 + let collected = chat.chat(request).await.unwrap().collect_message().await.unwrap();
1484 + assert_eq!(collected.message.text(), "Hi");
1485 + assert_eq!(
1486 + collected.prompt_logprobs,
1487 + Some(DecodedPromptLogprobs {
1488 + first_token_id: b's' as u32,
1489 + first_token: "s".to_string(),
1490 + scored_positions: vec![DecodedPositionLogprobs {
1491 + entries: vec![
1492 + DecodedTokenLogprob {
1493 + token_id: b'i' as u32,
1494 + token: "i".to_string(),
1495 + logprob: -0.3,
1496 + rank: 1,
1497 + },
1498 + DecodedTokenLogprob {
1499 + token_id: b'!' as u32,
1500 + token: "!".to_string(),
1501 + logprob: -0.4,
1502 + rank: 1,
1503 + },
1504 + ],
1505 + }],
1506 + })
1507 + );
1508 + assert_eq!(
1509 + collected.logprobs,
1510 + Some(DecodedLogprobs {
1511 + positions: vec![
1512 + DecodedPositionLogprobs {
1513 + entries: vec![
1514 + DecodedTokenLogprob {
1515 + token_id: b'H' as u32,
1516 + token: "H".to_string(),
1517 + logprob: -0.1,
1518 + rank: 1,
1519 + },
1520 + DecodedTokenLogprob {
1521 + token_id: b'h' as u32,
1522 + token: "h".to_string(),
1523 + logprob: -0.2,
1524 + rank: 1,
1525 + },
1526 + ],
1527 + },
1528 + DecodedPositionLogprobs {
1529 + entries: vec![
1530 + DecodedTokenLogprob {
1531 + token_id: b'i' as u32,
1532 + token: "i".to_string(),
1533 + logprob: -0.1,
1534 + rank: 1,
1535 + },
1536 + DecodedTokenLogprob {
1537 + token_id: b'I' as u32,
1538 + token: "I".to_string(),
1539 + logprob: -0.2,
1540 + rank: 1,
1541 + },
1542 + ],
1543 + },
1544 + ],
1545 + })
1546 + );
1547 +
1548 + let _ = shutdown_tx.send(());
1549 + engine_task.await.unwrap();
1550 + chat.shutdown().await.unwrap();
1551 +}
1552 +
1553 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1554 +async fn chat_rejects_unknown_tool_parser_before_engine_request() {
1555 + let ipc = IpcNamespace::new().unwrap();
1556 + let handshake_address = ipc.handshake_endpoint();
1557 + let engine_id = b"engine-chat-tool-no-model".to_vec();
1558 + let (shutdown_tx, engine_task) =
1559 + spawn_mock_engine_task(handshake_address.clone(), engine_id, |dealer, _| {
1560 + Box::pin(async move {
1561 + assert!(
1562 + timeout(Duration::from_millis(100), recv_engine_message(dealer)).await.is_err(),
1563 + "chat request should fail before any engine request is sent"
1564 + );
1565 + })
1566 + });
1567 +
1568 + let backend: Arc<dyn ChatTextBackend> = Arc::new(FakeChatBackend::new());
1569 + let chat = connect_chat_llm_with_ipc(
1570 + EngineCoreClientConfig::new_single(handshake_address),
1571 + &ipc,
1572 + backend,
1573 + )
1574 + .await
1575 + .with_tool_call_parser(ParserSelection::Explicit(
1576 + "definitely_missing_tool_parser".into(),
1577 + ));
1578 + let error = match chat.chat(sample_tool_request("chat-tool-no-model")).await {
1579 + Ok(_) => panic!("unknown explicit tool parser should fail"),
1580 + Err(error) => error,
1581 + };
1582 +
1583 + assert!(matches!(
1584 + error,
1585 + vllm_chat::Error::ParserUnavailableByName { name, .. }
1586 + if name == "definitely_missing_tool_parser"
1587 + ));
1588 +
1589 + let _ = shutdown_tx.send(());
1590 + engine_task.await.unwrap();
1591 + chat.shutdown().await.unwrap();
1592 +}
1593 +
1594 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1595 +async fn chat_rejects_unknown_reasoning_parser_before_engine_request() {
1596 + let ipc = IpcNamespace::new().unwrap();
1597 + let handshake_address = ipc.handshake_endpoint();
1598 + let engine_id = b"engine-chat-reasoning-no-model".to_vec();
1599 + let (shutdown_tx, engine_task) =
1600 + spawn_mock_engine_task(handshake_address.clone(), engine_id, |dealer, _| {
1601 + Box::pin(async move {
1602 + assert!(
1603 + timeout(Duration::from_millis(100), recv_engine_message(dealer)).await.is_err(),
1604 + "chat request should fail before any engine request is sent"
1605 + );
1606 + })
1607 + });
1608 +
1609 + let backend: Arc<dyn ChatTextBackend> = Arc::new(FakeChatBackend::new());
1610 + let chat = connect_chat_llm_with_ipc(
1611 + EngineCoreClientConfig::new_single(handshake_address),
1612 + &ipc,
1613 + backend,
1614 + )
1615 + .await
1616 + .with_reasoning_parser(ParserSelection::Explicit(
1617 + "definitely_missing_reasoning_parser".into(),
1618 + ));
1619 + let error = match chat.chat(sample_request("chat-reasoning-no-model")).await {
1620 + Ok(_) => panic!("unknown explicit reasoning parser should fail"),
1621 + Err(error) => error,
1622 + };
1623 +
1624 + assert!(matches!(
1625 + error,
1626 + vllm_chat::Error::ParserUnavailableByName { name, .. }
1627 + if name == "definitely_missing_reasoning_parser"
1628 + ));
1629 +
1630 + let _ = shutdown_tx.send(());
1631 + engine_task.await.unwrap();
1632 + chat.shutdown().await.unwrap();
1633 +}
1634 +
1635 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1636 +async fn chat_rejects_tool_requests_when_tool_parser_is_disabled() {
1637 + let ipc = IpcNamespace::new().unwrap();
1638 + let handshake_address = ipc.handshake_endpoint();
1639 + let engine_id = b"engine-chat-tool-parser-disabled".to_vec();
1640 + let (shutdown_tx, engine_task) =
1641 + spawn_mock_engine_task(handshake_address.clone(), engine_id, |dealer, _| {
1642 + Box::pin(async move {
1643 + assert!(
1644 + timeout(Duration::from_millis(100), recv_engine_message(dealer)).await.is_err(),
1645 + "chat request should fail before any engine request is sent"
1646 + );
1647 + })
1648 + });
1649 +
1650 + let backend: Arc<dyn ChatTextBackend> = Arc::new(FakeChatBackend::new());
1651 + let chat = connect_chat_llm_with_ipc(
1652 + EngineCoreClientConfig::new_single(handshake_address),
1653 + &ipc,
1654 + backend,
1655 + )
1656 + .await
1657 + .with_tool_call_parser(ParserSelection::None);
1658 + let error = match chat.chat(sample_tool_request("chat-tool-parser-disabled")).await {
1659 + Ok(_) => panic!("tool requests should fail when tool parsing is disabled"),
1660 + Err(error) => error,
1661 + };
1662 +
1663 + assert!(matches!(
1664 + error,
1665 + vllm_chat::Error::ParserDisabled { kind: "tool" }
1666 + ));
1667 +
1668 + let _ = shutdown_tx.send(());
1669 + engine_task.await.unwrap();
1670 + chat.shutdown().await.unwrap();
1671 +}
rust/src/chat/tests/templates/qwen3.jinja new
+89
@@ -0,0 +1,89 @@
1 +{%- if tools %}
2 + {{- '<|im_start|>system\n' }}
3 + {%- if messages[0].role == 'system' %}
4 + {{- messages[0].content + '\n\n' }}
5 + {%- endif %}
6 + {{- "# Tools\n\nYou may call one or more functions to assist with the user query.\n\nYou are provided with function signatures within <tools></tools> XML tags:\n<tools>" }}
7 + {%- for tool in tools %}
8 + {{- "\n" }}
9 + {{- tool | tojson }}
10 + {%- endfor %}
11 + {{- "\n</tools>\n\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\n<tool_call>\n{\"name\": <function-name>, \"arguments\": <args-json-object>}\n</tool_call><|im_end|>\n" }}
12 +{%- else %}
13 + {%- if messages[0].role == 'system' %}
14 + {{- '<|im_start|>system\n' + messages[0].content + '<|im_end|>\n' }}
15 + {%- endif %}
16 +{%- endif %}
17 +{%- set ns = namespace(multi_step_tool=true, last_query_index=messages|length - 1) %}
18 +{%- for message in messages[::-1] %}
19 + {%- set index = (messages|length - 1) - loop.index0 %}
20 + {%- if ns.multi_step_tool and message.role == "user" and message.content is string and not(message.content.startswith('<tool_response>') and message.content.endswith('</tool_response>')) %}
21 + {%- set ns.multi_step_tool = false %}
22 + {%- set ns.last_query_index = index %}
23 + {%- endif %}
24 +{%- endfor %}
25 +{%- for message in messages %}
26 + {%- if message.content is string %}
27 + {%- set content = message.content %}
28 + {%- else %}
29 + {%- set content = '' %}
30 + {%- endif %}
31 + {%- if (message.role == "user") or (message.role == "system" and not loop.first) %}
32 + {{- '<|im_start|>' + message.role + '\n' + content + '<|im_end|>' + '\n' }}
33 + {%- elif message.role == "assistant" %}
34 + {%- set reasoning_content = '' %}
35 + {%- if message.reasoning_content is string %}
36 + {%- set reasoning_content = message.reasoning_content %}
37 + {%- else %}
38 + {%- if '</think>' in content %}
39 + {%- set reasoning_content = content.split('</think>')[0].rstrip('\n').split('<think>')[-1].lstrip('\n') %}
40 + {%- set content = content.split('</think>')[-1].lstrip('\n') %}
41 + {%- endif %}
42 + {%- endif %}
43 + {%- if loop.index0 > ns.last_query_index %}
44 + {%- if loop.last or (not loop.last and reasoning_content) %}
45 + {{- '<|im_start|>' + message.role + '\n<think>\n' + reasoning_content.strip('\n') + '\n</think>\n\n' + content.lstrip('\n') }}
46 + {%- else %}
47 + {{- '<|im_start|>' + message.role + '\n' + content }}
48 + {%- endif %}
49 + {%- else %}
50 + {{- '<|im_start|>' + message.role + '\n' + content }}
51 + {%- endif %}
52 + {%- if message.tool_calls %}
53 + {%- for tool_call in message.tool_calls %}
54 + {%- if (loop.first and content) or (not loop.first) %}
55 + {{- '\n' }}
56 + {%- endif %}
57 + {%- if tool_call.function %}
58 + {%- set tool_call = tool_call.function %}
59 + {%- endif %}
60 + {{- '<tool_call>\n{"name": "' }}
61 + {{- tool_call.name }}
62 + {{- '", "arguments": ' }}
63 + {%- if tool_call.arguments is string %}
64 + {{- tool_call.arguments }}
65 + {%- else %}
66 + {{- tool_call.arguments | tojson }}
67 + {%- endif %}
68 + {{- '}\n</tool_call>' }}
69 + {%- endfor %}
70 + {%- endif %}
71 + {{- '<|im_end|>\n' }}
72 + {%- elif message.role == "tool" %}
73 + {%- if loop.first or (messages[loop.index0 - 1].role != "tool") %}
74 + {{- '<|im_start|>user' }}
75 + {%- endif %}
76 + {{- '\n<tool_response>\n' }}
77 + {{- content }}
78 + {{- '\n</tool_response>' }}
79 + {%- if loop.last or (messages[loop.index0 + 1].role != "tool") %}
80 + {{- '<|im_end|>\n' }}
81 + {%- endif %}
82 + {%- endif %}
83 +{%- endfor %}
84 +{%- if add_generation_prompt %}
85 + {{- '<|im_start|>assistant\n' }}
86 + {%- if enable_thinking is defined and enable_thinking is false %}
87 + {{- '<think>\n\n</think>\n\n' }}
88 + {%- endif %}
89 +{%- endif %}
rust/src/chat/tests/templates/qwen35.jinja new
+154
@@ -0,0 +1,154 @@
1 +{%- set image_count = namespace(value=0) %}
2 +{%- set video_count = namespace(value=0) %}
3 +{%- macro render_content(content, do_vision_count, is_system_content=false) %}
4 + {%- if content is string %}
5 + {{- content }}
6 + {%- elif content is iterable and content is not mapping %}
7 + {%- for item in content %}
8 + {%- if 'image' in item or 'image_url' in item or item.type == 'image' %}
9 + {%- if is_system_content %}
10 + {{- raise_exception('System message cannot contain images.') }}
11 + {%- endif %}
12 + {%- if do_vision_count %}
13 + {%- set image_count.value = image_count.value + 1 %}
14 + {%- endif %}
15 + {%- if add_vision_id %}
16 + {{- 'Picture ' ~ image_count.value ~ ': ' }}
17 + {%- endif %}
18 + {{- '<|vision_start|><|image_pad|><|vision_end|>' }}
19 + {%- elif 'video' in item or item.type == 'video' %}
20 + {%- if is_system_content %}
21 + {{- raise_exception('System message cannot contain videos.') }}
22 + {%- endif %}
23 + {%- if do_vision_count %}
24 + {%- set video_count.value = video_count.value + 1 %}
25 + {%- endif %}
26 + {%- if add_vision_id %}
27 + {{- 'Video ' ~ video_count.value ~ ': ' }}
28 + {%- endif %}
29 + {{- '<|vision_start|><|video_pad|><|vision_end|>' }}
30 + {%- elif 'text' in item %}
31 + {{- item.text }}
32 + {%- else %}
33 + {{- raise_exception('Unexpected item type in content.') }}
34 + {%- endif %}
35 + {%- endfor %}
36 + {%- elif content is none or content is undefined %}
37 + {{- '' }}
38 + {%- else %}
39 + {{- raise_exception('Unexpected content type.') }}
40 + {%- endif %}
41 +{%- endmacro %}
42 +{%- if not messages %}
43 + {{- raise_exception('No messages provided.') }}
44 +{%- endif %}
45 +{%- if tools and tools is iterable and tools is not mapping %}
46 + {{- '<|im_start|>system\n' }}
47 + {{- "# Tools\n\nYou have access to the following functions:\n\n<tools>" }}
48 + {%- for tool in tools %}
49 + {{- "\n" }}
50 + {{- tool | tojson }}
51 + {%- endfor %}
52 + {{- "\n</tools>" }}
53 + {{- '\n\nIf you choose to call a function ONLY reply in the following format with NO suffix:\n\n<tool_call>\n<function=example_function_name>\n<parameter=example_parameter_1>\nvalue_1\n</parameter>\n<parameter=example_parameter_2>\nThis is the value for the second parameter\nthat can span\nmultiple lines\n</parameter>\n</function>\n</tool_call>\n\n<IMPORTANT>\nReminder:\n- Function calls MUST follow the specified format: an inner <function=...></function> block must be nested within <tool_call></tool_call> XML tags\n- Required parameters MUST be specified\n- You may provide optional reasoning for your function call in natural language BEFORE the function call, but NOT after\n- If there is no function call available, answer the question like normal with your current knowledge and do not tell the user about function calls\n</IMPORTANT>' }}
54 + {%- if messages[0].role == 'system' %}
55 + {%- set content = render_content(messages[0].content, false, true)|trim %}
56 + {%- if content %}
57 + {{- '\n\n' + content }}
58 + {%- endif %}
59 + {%- endif %}
60 + {{- '<|im_end|>\n' }}
61 +{%- else %}
62 + {%- if messages[0].role == 'system' %}
63 + {%- set content = render_content(messages[0].content, false, true)|trim %}
64 + {{- '<|im_start|>system\n' + content + '<|im_end|>\n' }}
65 + {%- endif %}
66 +{%- endif %}
67 +{%- set ns = namespace(multi_step_tool=true, last_query_index=messages|length - 1) %}
68 +{%- for message in messages[::-1] %}
69 + {%- set index = (messages|length - 1) - loop.index0 %}
70 + {%- if ns.multi_step_tool and message.role == "user" %}
71 + {%- set content = render_content(message.content, false)|trim %}
72 + {%- if not(content.startswith('<tool_response>') and content.endswith('</tool_response>')) %}
73 + {%- set ns.multi_step_tool = false %}
74 + {%- set ns.last_query_index = index %}
75 + {%- endif %}
76 + {%- endif %}
77 +{%- endfor %}
78 +{%- if ns.multi_step_tool %}
79 + {{- raise_exception('No user query found in messages.') }}
80 +{%- endif %}
81 +{%- for message in messages %}
82 + {%- set content = render_content(message.content, true)|trim %}
83 + {%- if message.role == "system" %}
84 + {%- if not loop.first %}
85 + {{- raise_exception('System message must be at the beginning.') }}
86 + {%- endif %}
87 + {%- elif message.role == "user" %}
88 + {{- '<|im_start|>' + message.role + '\n' + content + '<|im_end|>' + '\n' }}
89 + {%- elif message.role == "assistant" %}
90 + {%- set reasoning_content = '' %}
91 + {%- if message.reasoning_content is string %}
92 + {%- set reasoning_content = message.reasoning_content %}
93 + {%- else %}
94 + {%- if '</think>' in content %}
95 + {%- set reasoning_content = content.split('</think>')[0].rstrip('\n').split('<think>')[-1].lstrip('\n') %}
96 + {%- set content = content.split('</think>')[-1].lstrip('\n') %}
97 + {%- endif %}
98 + {%- endif %}
99 + {%- set reasoning_content = reasoning_content|trim %}
100 + {%- if loop.index0 > ns.last_query_index %}
101 + {{- '<|im_start|>' + message.role + '\n<think>\n' + reasoning_content + '\n</think>\n\n' + content }}
102 + {%- else %}
103 + {{- '<|im_start|>' + message.role + '\n' + content }}
104 + {%- endif %}
105 + {%- if message.tool_calls and message.tool_calls is iterable and message.tool_calls is not mapping %}
106 + {%- for tool_call in message.tool_calls %}
107 + {%- if tool_call.function is defined %}
108 + {%- set tool_call = tool_call.function %}
109 + {%- endif %}
110 + {%- if loop.first %}
111 + {%- if content|trim %}
112 + {{- '\n\n<tool_call>\n<function=' + tool_call.name + '>\n' }}
113 + {%- else %}
114 + {{- '<tool_call>\n<function=' + tool_call.name + '>\n' }}
115 + {%- endif %}
116 + {%- else %}
117 + {{- '\n<tool_call>\n<function=' + tool_call.name + '>\n' }}
118 + {%- endif %}
119 + {%- if tool_call.arguments is defined %}
120 + {%- for args_name, args_value in tool_call.arguments|items %}
121 + {{- '<parameter=' + args_name + '>\n' }}
122 + {%- set args_value = args_value | tojson | safe if args_value is mapping or (args_value is sequence and args_value is not string) else args_value | string %}
123 + {{- args_value }}
124 + {{- '\n</parameter>\n' }}
125 + {%- endfor %}
126 + {%- endif %}
127 + {{- '</function>\n</tool_call>' }}
128 + {%- endfor %}
129 + {%- endif %}
130 + {{- '<|im_end|>\n' }}
131 + {%- elif message.role == "tool" %}
132 + {%- if loop.previtem and loop.previtem.role != "tool" %}
133 + {{- '<|im_start|>user' }}
134 + {%- endif %}
135 + {{- '\n<tool_response>\n' }}
136 + {{- content }}
137 + {{- '\n</tool_response>' }}
138 + {%- if not loop.last and loop.nextitem.role != "tool" %}
139 + {{- '<|im_end|>\n' }}
140 + {%- elif loop.last %}
141 + {{- '<|im_end|>\n' }}
142 + {%- endif %}
143 + {%- else %}
144 + {{- raise_exception('Unexpected message role.') }}
145 + {%- endif %}
146 +{%- endfor %}
147 +{%- if add_generation_prompt %}
148 + {{- '<|im_start|>assistant\n' }}
149 + {%- if enable_thinking is defined and enable_thinking is true %}
150 + {{- '<think>\n' }}
151 + {%- else %}
152 + {{- '<think>\n\n</think>\n\n' }}
153 + {%- endif %}
154 +{%- endif %}
rust/src/chat/tests/templates/vllm_examples/README.md new
+6
@@ -0,0 +1,6 @@
1 +# vLLM Example Templates
2 +
3 +These fixtures are copied from `vllm/examples/`.
4 +
5 +They are currently used by `src/chat/src/renderers/hf/format.rs` tests to keep
6 +our chat-template content format detection aligned with Python vLLM behavior.
rust/src/chat/tests/templates/vllm_examples/template_alpaca.jinja new
+29
@@ -0,0 +1,29 @@
1 +{{ (messages|selectattr('role', 'equalto', 'system')|list|last).content|trim if (messages|selectattr('role', 'equalto', 'system')|list) else '' }}
2 +
3 +{% for message in messages %}
4 +{% if message['role'] == 'user' %}
5 +### Instruction:
6 +{{ message['content']|trim -}}
7 +{% if not loop.last %}
8 +
9 +
10 +{% endif %}
11 +{% elif message['role'] == 'assistant' %}
12 +### Response:
13 +{{ message['content']|trim -}}
14 +{% if not loop.last %}
15 +
16 +
17 +{% endif %}
18 +{% elif message['role'] == 'user_context' %}
19 +### Input:
20 +{{ message['content']|trim -}}
21 +{% if not loop.last %}
22 +
23 +
24 +{% endif %}
25 +{% endif %}
26 +{% endfor %}
27 +{% if add_generation_prompt and messages[-1]['role'] != 'assistant' %}
28 +### Response:
29 +{% endif %}
\ No newline at end of file
rust/src/chat/tests/templates/vllm_examples/template_baichuan.jinja new
+13
@@ -0,0 +1,13 @@
1 +{{ (messages|selectattr('role', 'equalto', 'system')|list|last).content|trim if (messages|selectattr('role', 'equalto', 'system')|list) else '' }}
2 +
3 +{%- for message in messages -%}
4 + {%- if message['role'] == 'user' -%}
5 + {{- '<reserved_106>' + message['content'] -}}
6 + {%- elif message['role'] == 'assistant' -%}
7 + {{- '<reserved_107>' + message['content'] -}}
8 + {%- endif -%}
9 +{%- endfor -%}
10 +
11 +{%- if add_generation_prompt and messages[-1]['role'] != 'assistant' -%}
12 + {{- '<reserved_107>' -}}
13 +{% endif %}
\ No newline at end of file
rust/src/chat/tests/templates/vllm_examples/template_chatglm.jinja new
+18
@@ -0,0 +1,18 @@
1 +{%- set counter = namespace(index=0) -%}
2 +{%- for message in messages -%}
3 + {%- if message['role'] == 'user' -%}
4 + {{- '[Round ' + counter.index|string + ']\n问：' + message['content'] -}}
5 + {%- set counter.index = counter.index + 1 -%}
6 + {%- endif -%}
7 + {%- if message['role'] == 'assistant' -%}
8 + {{- '\n答：' + message['content'] -}}
9 + {%- if (loop.last and add_generation_prompt) or not loop.last -%}
10 + {{- '\n' -}}
11 + {%- endif -%}
12 + {%- endif -%}
13 +{%- endfor -%}
14 +
15 +
16 +{%- if add_generation_prompt and messages[-1]['role'] != 'assistant' -%}
17 + {{- '\n答：' -}}
18 +{%- endif -%}
\ No newline at end of file
rust/src/chat/tests/templates/vllm_examples/template_chatglm2.jinja new
+18
@@ -0,0 +1,18 @@
1 +{%- set counter = namespace(index=1) -%}
2 +{%- for message in messages -%}
3 + {%- if message['role'] == 'user' -%}
4 + {{- '[Round ' + counter.index|string + ']\n\n问：' + message['content'] -}}
5 + {%- set counter.index = counter.index + 1 -%}
6 + {%- endif -%}
7 + {%- if message['role'] == 'assistant' -%}
8 + {{- '\n\n答：' + message['content'] -}}
9 + {%- if (loop.last and add_generation_prompt) or not loop.last -%}
10 + {{- '\n\n' -}}
11 + {%- endif -%}
12 + {%- endif -%}
13 +{%- endfor -%}
14 +
15 +
16 +{%- if add_generation_prompt and messages[-1]['role'] != 'assistant' -%}
17 + {{- '\n\n答：' -}}
18 +{%- endif -%}
\ No newline at end of file
rust/src/chat/tests/templates/vllm_examples/template_chatml.jinja new
+2
@@ -0,0 +1,2 @@
1 +{% for message in messages %}{{'<|im_start|>' + message['role'] + '\n' + message['content']}}{% if (loop.last and add_generation_prompt) or not loop.last %}{{ '<|im_end|>' + '\n'}}{% endif %}{% endfor %}
2 +{% if add_generation_prompt and messages[-1]['role'] != 'assistant' %}{{ '<|im_start|>assistant\n' }}{% endif %}
\ No newline at end of file
rust/src/chat/tests/templates/vllm_examples/template_falcon.jinja new
+15
@@ -0,0 +1,15 @@
1 +{%- for message in messages -%}
2 + {%- if message['role'] == 'user' -%}
3 + {{- 'User: ' + message['content'] -}}
4 + {%- elif message['role'] == 'assistant' -%}
5 + {{- 'Assistant: ' + message['content'] -}}
6 + {%- endif -%}
7 + {%- if (loop.last and add_generation_prompt) or not loop.last -%}
8 + {{- '\n' -}}
9 + {%- endif -%}
10 +{%- endfor -%}
11 +
12 +
13 +{%- if add_generation_prompt and messages[-1]['role'] != 'assistant' -%}
14 + {{- 'Assistant:' -}}
15 +{% endif %}
\ No newline at end of file
rust/src/chat/tests/templates/vllm_examples/template_falcon_180b.jinja new
+17
@@ -0,0 +1,17 @@
1 +{%- for message in messages -%}
2 + {%- if message['role'] == 'system' -%}
3 + {{- 'System: ' + message['content'] -}}
4 + {%- elif message['role'] == 'user' -%}
5 + {{- 'User: ' + message['content'] -}}
6 + {%- elif message['role'] == 'assistant' -%}
7 + {{- 'Falcon: ' + message['content'] -}}
8 + {%- endif -%}
9 + {%- if (loop.last and add_generation_prompt) or not loop.last -%}
10 + {{- '\n' -}}
11 + {%- endif -%}
12 +{%- endfor -%}
13 +
14 +
15 +{%- if add_generation_prompt and messages[-1]['role'] != 'assistant' -%}
16 + {{- 'Falcon:' -}}
17 +{% endif %}
\ No newline at end of file
rust/src/chat/tests/templates/vllm_examples/template_inkbot.jinja new
+30
@@ -0,0 +1,30 @@
1 +<#meta#>
2 +- Date: {{ (messages|selectattr('role', 'equalto', 'meta-current_date')|list|last).content|trim if (messages|selectattr('role', 'equalto', 'meta-current_date')|list) else '' }}
3 +- Task: {{ (messages|selectattr('role', 'equalto', 'meta-task_name')|list|last).content|trim if (messages|selectattr('role', 'equalto', 'meta-task_name')|list) else '' }}
4 +<#system#>
5 +{{ (messages|selectattr('role', 'equalto', 'system')|list|last).content|trim if (messages|selectattr('role', 'equalto', 'system')|list) else '' }}
6 +<#chat#>
7 +{% for message in messages %}
8 +{% if message['role'] == 'user' %}
9 +<#user#>
10 +{{ message['content']|trim -}}
11 +{% if not loop.last %}
12 +
13 +{% endif %}
14 +{% elif message['role'] == 'assistant' %}
15 +<#bot#>
16 +{{ message['content']|trim -}}
17 +{% if not loop.last %}
18 +
19 +{% endif %}
20 +{% elif message['role'] == 'user_context' %}
21 +<#user_context#>
22 +{{ message['content']|trim -}}
23 +{% if not loop.last %}
24 +
25 +{% endif %}
26 +{% endif %}
27 +{% endfor %}
28 +{% if add_generation_prompt and messages[-1]['role'] != 'assistant' %}
29 +<#bot#>
30 +{% endif %}
\ No newline at end of file
rust/src/chat/tests/templates/vllm_examples/template_teleflm.jinja new
+12
@@ -0,0 +1,12 @@
1 +{%- for message in messages %}
2 + {%- if message['role'] == 'user' %}
3 + {{- '<_user>' + message['content']|trim }}
4 + {%- elif message['role'] == 'system' %}
5 + {{- '<_system>' + message['content']|trim }}
6 + {%- elif message['role'] == 'assistant' %}
7 + {{- '<_bot>' + message['content'] }}
8 + {%- endif %}
9 +{%- endfor %}
10 +{%- if add_generation_prompt %}
11 + {{- '<_bot>' }}
12 +{%- endif %}
rust/src/chat/tests/templates/vllm_examples/tool_chat_template_deepseekr1.jinja new
+92
@@ -0,0 +1,92 @@
1 +{% if not add_generation_prompt is defined %}
2 + {% set add_generation_prompt = false %}
3 +{% endif %}
4 +{% set ns = namespace(is_first=false, is_tool=false, is_output_first=true, system_prompt='', is_first_sp=true, is_last_user=false) %}
5 +{%- for message in messages %}
6 + {%- if message['role'] == 'system' %}
7 + {%- if ns.is_first_sp %}
8 + {% set ns.system_prompt = ns.system_prompt + message['content'] %}
9 + {% set ns.is_first_sp = false %}
10 + {%- else %}
11 + {% set ns.system_prompt = ns.system_prompt + '\n\n' + message['content'] %}
12 + {%- endif %}
13 + {%- endif %}
14 +{%- endfor -%}
15 +
16 +{#- Adapted from https://github.com/sgl-project/sglang/blob/main/examples/chat_template/tool_chat_template_deepseekr1.jinja #}
17 +{% if tools is defined and tools is not none %}
18 + {% set tool_ns = namespace(text='You are a helpful assistant with tool calling capabilities. '
19 + 'When a tool call is needed, you MUST use the following format to issue the call:\n'
20 + '<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>function<｜tool▁sep｜>FUNCTION_NAME\n'
21 + '```json\n{"param1": "value1", "param2": "value2"}\n```<｜tool▁call▁end｜><｜tool▁calls▁end｜>\n\n'
22 + 'Make sure the JSON is valid.'
23 + '## Tools\n\n### Function\n\nYou have the following functions available:\n\n') %}
24 + {% for tool in tools %}
25 + {% set tool_ns.text = tool_ns.text + '\n```json\n' + (tool | tojson) + '\n```\n' %}
26 + {% endfor %}
27 + {% set ns.system_prompt = ns.system_prompt + '\n\n' + tool_ns.text %}
28 +{% endif %}
29 +
30 +{{- bos_token }}
31 +{{- ns.system_prompt }}
32 +{%- for message in messages %}
33 + {% set content = message['content'] %}
34 + {%- if message['role'] == 'user' %}
35 + {%- set ns.is_tool = false -%}
36 + {%- set ns.is_first = false -%}
37 + {%- set ns.is_last_user = true -%}
38 + {{'<｜User｜>' + content + '<｜Assistant｜>'}}
39 + {%- endif %}
40 + {%- if message['role'] == 'assistant' %}
41 + {% if '</think>' in content %}
42 + {% set content = content.split('</think>')[-1] %}
43 + {% endif %}
44 + {% endif %}
45 + {%- if message['role'] == 'assistant' and message['tool_calls'] is defined and message['tool_calls'] is not none %}
46 + {%- set ns.is_last_user = false -%}
47 + {%- if ns.is_tool %}
48 + {{- '<｜tool▁outputs▁end｜>'}}
49 + {%- endif %}
50 + {%- set ns.is_first = false %}
51 + {%- set ns.is_tool = false -%}
52 + {%- set ns.is_output_first = true %}
53 + {%- for tool in message['tool_calls'] %}
54 + {%- if not ns.is_first %}
55 + {%- if content is none %}
56 + {{- '<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>' + tool['type'] + '<｜tool▁sep｜>' + tool['function']['name'] + '\n' + '```json' + '\n' + tool['function']['arguments']|tojson + '\n' + '```' + '<｜tool▁call▁end｜>'}}
57 + {%- else %}
58 + {{- content + '<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>' + tool['type'] + '<｜tool▁sep｜>' + tool['function']['name'] + '\n' + '```json' + '\n' + tool['function']['arguments']|tojson + '\n' + '```' + '<｜tool▁call▁end｜>'}}
59 + {%- endif %}
60 + {%- set ns.is_first = true -%}
61 + {%- else %}
62 + {{- '\n' + '<｜tool▁call▁begin｜>' + tool['type'] + '<｜tool▁sep｜>' + tool['function']['name'] + '\n' + '```json' + '\n' + tool['function']['arguments']|tojson + '\n' + '```' + '<｜tool▁call▁end｜>'}}
63 + {%- endif %}
64 + {%- endfor %}
65 + {{- '<｜tool▁calls▁end｜><｜end▁of▁sentence｜>'}}
66 + {%- endif %}
67 + {%- if message['role'] == 'assistant' and (message['tool_calls'] is not defined or message['tool_calls'] is none)%}
68 + {%- set ns.is_last_user = false -%}
69 + {%- if ns.is_tool %}
70 + {{- '<｜tool▁outputs▁end｜>' + content + '<｜end▁of▁sentence｜>'}}
71 + {%- set ns.is_tool = false -%}
72 + {%- else %}
73 + {{- content + '<｜end▁of▁sentence｜>'}}
74 + {%- endif %}
75 + {%- endif %}
76 + {%- if message['role'] == 'tool' %}
77 + {%- set ns.is_last_user = false -%}
78 + {%- set ns.is_tool = true -%}
79 + {%- if ns.is_output_first %}
80 + {{- '<｜tool▁outputs▁begin｜><｜tool▁output▁begin｜>' + content + '<｜tool▁output▁end｜>'}}
81 + {%- set ns.is_output_first = false %}
82 + {%- else %}
83 + {{- '\n<｜tool▁output▁begin｜>' + content + '<｜tool▁output▁end｜>'}}
84 + {%- endif %}
85 + {%- endif %}
86 +{%- endfor -%}
87 +{% if ns.is_tool %}
88 + {{- '<｜tool▁outputs▁end｜>'}}
89 +{%- endif %}
90 +{% if add_generation_prompt and not ns.is_last_user and not ns.is_tool %}
91 + {{- '<｜Assistant｜>'}}
92 +{%- endif %}
\ No newline at end of file
rust/src/chat/tests/templates/vllm_examples/tool_chat_template_deepseekv3.jinja new
+96
@@ -0,0 +1,96 @@
1 +{% if not add_generation_prompt is defined %}
2 + {% set add_generation_prompt = false %}
3 +{% endif %}
4 +
5 +{% set ns = namespace(is_first=false, is_tool=false, is_output_first=true, system_prompt='', is_first_sp=true, is_last_user=false) %}
6 +
7 +{%- for message in messages %}
8 + {%- if message['role'] == 'system' %}
9 + {%- if ns.is_first_sp %}
10 + {% set ns.system_prompt = ns.system_prompt + message['content'] %}
11 + {% set ns.is_first_sp = false %}
12 + {%- else %}
13 + {% set ns.system_prompt = ns.system_prompt + '\n\n' + message['content'] %}
14 + {%- endif %}
15 + {%- endif %}
16 +{%- endfor %}
17 +
18 +{{ bos_token }}
19 +{{ ns.system_prompt }}
20 +{%- if tools %}
21 + {{"\n\n# Tools\n\nYou may call one or more functions to assist with the user query." }}
22 + {%- for tool in tools %}
23 + {{- "\n" }}
24 + {{- tool | tojson }}
25 + {%- endfor %}
26 + {{"\n</tools>\n\n"}}
27 +
28 + {{"For function call returns, you should first print <｜tool▁calls▁begin｜>"}}
29 +
30 + {{"For each function call, you should return object like:\n" }}
31 + {{"<｜tool▁call▁begin｜>function<｜tool▁sep｜><function_name>\n```json\n<function_arguments_in_json_format>\n```<｜tool▁call▁end｜>"}}
32 +
33 + {{"At the end of function call returns, you should print <｜tool▁calls▁end｜><｜end▁of▁sentence｜>"}}
34 +{%- endif %}
35 +
36 +{%- for message in messages %}
37 + {%- if message['role'] == 'user' %}
38 + {%- set ns.is_tool = false -%}
39 + {%- set ns.is_first = false -%}
40 + {%- set ns.is_last_user = true -%}
41 + {{'<｜User｜>' + message['content'] + '<｜Assistant｜>'}}
42 + {%- endif %}
43 +
44 + {%- if message['role'] == 'assistant' and message['tool_calls'] is defined and message['tool_calls'] is not none %}
45 + {%- set ns.is_last_user = false -%}
46 + {%- if ns.is_tool %}
47 + {{'<｜tool▁outputs▁end｜>'}}
48 + {%- endif %}
49 + {%- set ns.is_first = false %}
50 + {%- set ns.is_tool = false -%}
51 + {%- set ns.is_output_first = true %}
52 +
53 + {%- for tool in message['tool_calls'] %}
54 + {%- if not ns.is_first %}
55 + {%- if message['content'] is none %}
56 + {{'<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>' + tool['type'] + '<｜tool▁sep｜>' + tool['function']['name'] + '\n' + '```json' + '\n' + tool['function']['arguments']|tojson + '\n' + '```' + '<｜tool▁call▁end｜>'}}
57 + {%- else %}
58 + {{message['content'] + '<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>' + tool['type'] + '<｜tool▁sep｜>' + tool['function']['name'] + '\n' + '```json' + '\n' + tool['function']['arguments']|tojson + '\n' + '```' + '<｜tool▁call▁end｜>'}}
59 + {%- endif %}
60 + {%- set ns.is_first = true -%}
61 + {%- else %}
62 + {{'\n' + '<｜tool▁call▁begin｜>' + tool['type'] + '<｜tool▁sep｜>' + tool['function']['name'] + '\n' + '```json' + '\n' + tool['function']['arguments']|tojson + '\n' + '```' + '<｜tool▁call▁end｜>'}}
63 + {%- endif %}
64 + {%- endfor %}
65 + {{'<｜tool▁calls▁end｜><｜end▁of▁sentence｜>'}}
66 + {%- endif %}
67 + {%- if message['role'] == 'assistant' and (message['tool_calls'] is not defined or message['tool_calls'] is none)%}
68 + {%- set ns.is_last_user = false -%}
69 + {%- if ns.is_tool %}
70 + {{'<｜tool▁outputs▁end｜>' + message['content'] + '<｜end▁of▁sentence｜>'}}
71 + {%- set ns.is_tool = false -%}
72 + {%- else %}
73 + {% set content = message['content'] %}
74 + {{content + '<｜end▁of▁sentence｜>'}}
75 + {%- endif %}
76 + {%- endif %}
77 +
78 + {%- if message['role'] == 'tool' %}
79 + {%- set ns.is_last_user = false -%}
80 + {%- set ns.is_tool = true -%}
81 + {%- if ns.is_output_first %}
82 + {{'<｜tool▁outputs▁begin｜><｜tool▁output▁begin｜>' + message['content'] + '<｜tool▁output▁end｜>'}}
83 + {%- set ns.is_output_first = false %}
84 + {%- else %}
85 + {{'\n<｜tool▁output▁begin｜>' + message['content'] + '<｜tool▁output▁end｜>'}}
86 + {%- endif %}
87 + {%- endif %}
88 +{%- endfor -%}
89 +
90 +{% if ns.is_tool %}
91 + {{'<｜tool▁outputs▁end｜>'}}
92 +{% endif %}
93 +
94 +{% if add_generation_prompt and not ns.is_last_user and not ns.is_tool %}
95 + {{'<｜Assistant｜>'}}
96 +{% endif %}
rust/src/chat/tests/templates/vllm_examples/tool_chat_template_deepseekv31.jinja new
+91
@@ -0,0 +1,91 @@
1 +{% if not add_generation_prompt is defined %}
2 + {% set add_generation_prompt = false %}
3 +{% endif %}
4 +{% if not thinking is defined %}
5 + {% set thinking = false %}
6 +{% endif %}
7 +{% set ns = namespace(is_first=false, is_tool=false, system_prompt='', is_first_sp=true, is_last_user=false) %}
8 +{%- for message in messages %}
9 + {%- if message['role'] == 'system' %}
10 + {%- if ns.is_first_sp %}
11 + {% set ns.system_prompt = ns.system_prompt + message['content'] %}
12 + {% set ns.is_first_sp = false %}
13 + {%- else %}
14 + {% set ns.system_prompt = ns.system_prompt + '\n\n' + message['content'] %}
15 + {%- endif %}
16 + {%- endif %}
17 +{%- endfor %}
18 +
19 +{% if tools is defined and tools is not none %}
20 + {% set tool_ns = namespace(text='## Tools\nYou have access to the following tools:\n') %}
21 + {% for tool in tools %}
22 + {% set tool_ns.text = tool_ns.text + '\n### ' + tool.function.name + '\nDescription: ' + tool.function.description + '\n\nParameters: ' + (tool.function.parameters | tojson) + '\n' %}
23 + {% endfor %}
24 + {% set tool_ns.text = tool_ns.text + "\nIMPORTANT: ALWAYS adhere to this exact format for tool use:\n<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>tool_call_name<｜tool▁sep｜>tool_call_arguments<｜tool▁call▁end｜>{{additional_tool_calls}}<｜tool▁calls▁end｜>\n\nWhere:\n\n- `tool_call_name` must be an exact match to one of the available tools\n- `tool_call_arguments` must be valid JSON that strictly follows the tool's Parameters Schema\n- For multiple tool calls, chain them directly without separators or spaces\n" %}
25 + {% set ns.system_prompt = ns.system_prompt + '\n\n' + tool_ns.text %}
26 +{% endif %}
27 +
28 +{{ bos_token }}{{ ns.system_prompt }}
29 +{%- for message in messages %}
30 + {%- if message['role'] == 'user' %}
31 + {%- set ns.is_tool = false -%}
32 + {%- set ns.is_first = false -%}
33 + {%- set ns.is_last_user = true -%}
34 + {{'<｜User｜>' + message['content']}}
35 + {%- endif %}
36 + {%- if message['role'] == 'assistant' and message['tool_calls'] is defined and message['tool_calls'] is not none %}
37 + {%- if ns.is_last_user %}
38 + {{'<｜Assistant｜></think>'}}
39 + {%- endif %}
40 + {%- set ns.is_last_user = false -%}
41 + {%- set ns.is_first = false %}
42 + {%- set ns.is_tool = false -%}
43 + {%- for tool in message['tool_calls'] %}
44 + {%- if not ns.is_first %}
45 + {%- if message['content'] is none %}
46 + {{'<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>'+ tool['function']['name'] + '<｜tool▁sep｜>' + tool['function']['arguments']|tojson + '<｜tool▁call▁end｜>'}}
47 + {%- else %}
48 + {{message['content'] + '<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>' + tool['function']['name'] + '<｜tool▁sep｜>' + tool['function']['arguments']|tojson + '<｜tool▁call▁end｜>'}}
49 + {%- endif %}
50 + {%- set ns.is_first = true -%}
51 + {%- else %}
52 + {{'<｜tool▁call▁begin｜>'+ tool['function']['name'] + '<｜tool▁sep｜>' + tool['function']['arguments']|tojson + '<｜tool▁call▁end｜>'}}
53 + {%- endif %}
54 + {%- endfor %}
55 + {{'<｜tool▁calls▁end｜><｜end▁of▁sentence｜>'}}
56 + {%- endif %}
57 + {%- if message['role'] == 'assistant' and (message['tool_calls'] is not defined or message['tool_calls'] is none) %}
58 + {%- if ns.is_last_user %}
59 + {{'<｜Assistant｜>'}}
60 + {%- if message['prefix'] is defined and message['prefix'] and thinking %}
61 + {{'<think>'}}
62 + {%- else %}
63 + {{'</think>'}}
64 + {%- endif %}
65 + {%- endif %}
66 + {%- set ns.is_last_user = false -%}
67 + {%- if ns.is_tool %}
68 + {{message['content'] + '<｜end▁of▁sentence｜>'}}
69 + {%- set ns.is_tool = false -%}
70 + {%- else %}
71 + {%- set content = message['content'] -%}
72 + {%- if '</think>' in content %}
73 + {%- set content = content.split('</think>', 1)[1] -%}
74 + {%- endif %}
75 + {{content + '<｜end▁of▁sentence｜>'}}
76 + {%- endif %}
77 + {%- endif %}
78 + {%- if message['role'] == 'tool' %}
79 + {%- set ns.is_last_user = false -%}
80 + {%- set ns.is_tool = true -%}
81 + {{'<｜tool▁output▁begin｜>' + message['content'] + '<｜tool▁output▁end｜>'}}
82 + {%- endif %}
83 +{%- endfor -%}
84 +{%- if add_generation_prompt and ns.is_last_user and not ns.is_tool %}
85 + {{'<｜Assistant｜>'}}
86 + {%- if not thinking %}
87 + {{'</think>'}}
88 + {%- else %}
89 + {{'<think>'}}
90 + {%- endif %}
91 +{% endif %}
rust/src/chat/tests/templates/vllm_examples/tool_chat_template_functiongemma.jinja new
+54
@@ -0,0 +1,54 @@
1 +{%- set ns = namespace(developer_content='', has_tools=false) -%}
2 +
3 +{%- if tools is defined and tools | length > 0 -%}
4 + {%- set ns.has_tools = true -%}
5 +{%- endif -%}
6 +
7 +{%- for message in messages -%}
8 + {%- if message.role == 'developer' or message.role == 'system' -%}
9 +<start_of_turn>user
10 +{{ message.content }}
11 +{%- if ns.has_tools %}
12 +
13 +Available functions:
14 +{%- for tool in tools %}
15 +{%- if tool.type == 'function' %}
16 +
17 +Function: {{ tool.function.name }}
18 +Description: {{ tool.function.description | default('No description provided') }}
19 +Parameters: {{ tool.function.parameters | tojson }}
20 +{%- endif %}
21 +{%- endfor %}
22 +{%- endif %}
23 +<end_of_turn>
24 + {%- elif message.role == 'user' -%}
25 +<start_of_turn>user
26 +{{ message.content }}<end_of_turn>
27 + {%- elif message.role == 'assistant' -%}
28 + {%- if message.tool_calls is defined and message.tool_calls | length > 0 -%}
29 +<start_of_turn>model
30 +{%- for tool_call in message.tool_calls %}
31 +<start_function_call>call:{{ tool_call.function.name }}{
32 +{%- set args = tool_call.function.arguments -%}
33 +{%- if args is string -%}
34 +{%- set args = args | fromjson -%}
35 +{%- endif -%}
36 +{%- for key, value in args.items() -%}
37 +{{ key }}:<escape>{{ value }}<escape>{% if not loop.last %},{% endif %}
38 +{%- endfor -%}
39 +}<end_function_call>
40 +{%- endfor %}
41 +<end_of_turn>
42 + {%- else -%}
43 +<start_of_turn>model
44 +{{ message.content }}<end_of_turn>
45 + {%- endif -%}
46 + {%- elif message.role == 'tool' -%}
47 +<start_of_turn>user
48 +Function result for {{ message.name | default('function') }}: {{ message.content }}<end_of_turn>
49 + {%- endif -%}
50 +{%- endfor -%}
51 +
52 +{%- if add_generation_prompt -%}
53 +<start_of_turn>model
54 +{%- endif -%}
rust/src/chat/tests/templates/vllm_examples/tool_chat_template_gemma3_pythonic.jinja new
+123
@@ -0,0 +1,123 @@
1 +{#- Begin-of-sequence token to start the model prompt -#}
2 +{{ bos_token }}
3 +{#- Extracts the system message. Gemma does not support system messages so it will be prepended to first user message. -#}
4 +{%- if messages[0]['role'] == 'system' -%}
5 + {%- if messages[0]['content'] is string -%}
6 + {%- set first_user_prefix = messages[0]['content'] + '\n\n' -%}
7 + {%- else -%}
8 + {%- set first_user_prefix = messages[0]['content'][0]['text'] + '\n\n' -%}
9 + {%- endif -%}
10 + {%- set loop_messages = messages[1:] -%}
11 +{%- else -%}
12 + {%- set first_user_prefix = "" -%}
13 + {%- set loop_messages = messages -%}
14 +{%- endif -%}
15 +{#- Set tools to none if not defined for this ChatCompletion request (helps avoid errors later) -#}
16 +{%- if not tools is defined %}
17 + {%- set tools = none %}
18 +{%- endif %}
19 +{#- Validate alternating user/assistant messages (excluding 'tool' messages and ones with tool_calls) -#}
20 +{%- for message in loop_messages | rejectattr("role", "equalto", "tool") | selectattr("tool_calls", "undefined") -%}
21 + {%- if (message['role'] == 'user') != (loop.index0 % 2 == 0) %}
22 + {{ raise_exception("Conversation roles must alternate user/assistant/user/assistant/...") }}
23 + {%- endif -%}
24 +{%- endfor -%}
25 +
26 +{#- Main loop over all messages in the conversation history -#}
27 +{%- for message in loop_messages -%}
28 + {#- Normalize roles for model prompt formatting -#}
29 + {%- if (message['role'] == 'assistant') -%}
30 + {%- set role = "model" -%}
31 + {%- elif (message['role'] == 'tool') -%}
32 + {%- set role = "user" -%}
33 + {%- else -%}
34 + {%- set role = message['role'] -%}
35 + {%- endif -%}
36 + {#- Mark the start of a message block with the appropriate role -#}
37 + {{ '<start_of_turn>' + role + '\n' -}}
38 +
39 + {#- Insert system message content (if present) at the beginning of the first message. -#}
40 + {%- if loop.first -%}
41 + {{ first_user_prefix }}
42 + {#- Append system message with tool information if using tools in message request. -#}
43 + {%- if tools is not none -%}
44 + {{- "Tools (functions) are available. If you decide to invoke one or more of the tools, you must respond with a python list of the function calls.\n" -}}
45 + {{- "Example Format: [func_name1(params_name1=params_value1, params_name2=params_value2...), func_name2(params)] \n" -}}
46 + {{- "Do not use variables. DO NOT USE MARKDOWN SYNTAX. You SHOULD NOT include any other text in the response if you call a function. If none of the functions can be used, point it out. If you lack the parameters required by the function, also point it out.\n" -}}
47 + {{- "Here is a list of functions in JSON format that you can invoke.\n" -}}
48 + {{- tools | tojson(indent=4) -}}
49 + {{- "\n\n" -}}
50 + {%- endif -%}
51 + {%- endif -%}
52 +
53 + {#- Format model tool calls (turns where model indicates they want to call a tool) -#}
54 + {%- if 'tool_calls' in message -%}
55 + {#- Opening bracket for tool call list. -#}
56 + {{- '[' -}}
57 + {#- For each tool call -#}
58 + {%- for tool_call in message.tool_calls -%}
59 + {#- Get tool call function. -#}
60 + {%- if tool_call.function is defined -%}
61 + {%- set tool_call = tool_call.function -%}
62 + {%- endif -%}
63 + {#- Function name & opening parenthesis. -#}
64 + {{- tool_call.name + '(' -}}
65 +
66 + {#-- Handle arguments as list (positional) or dict (named) --#}
67 + {#-- Named arguments (dict) --#}
68 + {%- if tool_call.arguments is iterable and tool_call.arguments is mapping -%}
69 + {%- set first = true -%}
70 + {%- for key, val in tool_call.arguments.items() -%}
71 + {%- if not first %}, {% endif -%}
72 + {{ key }}={{ val | tojson }}
73 + {%- set first = false -%}
74 + {%- endfor -%}
75 + {#-- Positional arguments (list) --#}
76 + {%- elif tool_call.arguments is iterable -%}
77 + {{- tool_call.arguments | map('tojson') | join(', ') -}}
78 + {#-- Fallback: single positional value --#}
79 + {%- else -%}
80 + {{- tool_call.arguments | tojson -}}
81 + {#-- Closing parenthesis. --#}
82 + {%- endif -%}
83 + {{- ')' -}}
84 + {#-- If more than one tool call, place comma and move to formatting next tool call --#}
85 + {%- if not loop.last -%}, {% endif -%}
86 + {%- endfor -%}
87 + {#- Closing bracket for tool call list. -#}
88 + {{- ']' -}}
89 + {%- endif -%}
90 +
91 + {#- Tool response start tag (for messages from a tool) -#}
92 + {%- if (message['role'] == 'tool') -%}
93 + {{ '<tool_response>\n' -}}
94 + {%- endif -%}
95 +
96 + {#- Render the message content: handle plain string or multimodal content like image/text -#}
97 + {%- if message['content'] is string -%}
98 + {{ message['content'] | trim }}
99 + {%- elif message['content'] is iterable -%}
100 + {%- for item in message['content'] -%}
101 + {%- if item['type'] == 'image' -%}
102 + {{ '<start_of_image>' }}
103 + {%- elif item['type'] == 'text' -%}
104 + {{ item['text'] | trim }}
105 + {%- endif -%}
106 + {%- endfor -%}
107 + {%- else -%}
108 + {{ raise_exception("Invalid content type") }}
109 + {%- endif -%}
110 +
111 + {#- Tool response end tag -#}
112 + {%- if (message['role'] == 'tool') -%}
113 + {{ '</tool_response>' -}}
114 + {%- endif -%}
115 +
116 + {#- Mark end of a single turn -#}
117 + {{ '<end_of_turn>\n' }}
118 +{%- endfor -%}
119 +
120 +{#- If generation is to be triggered, add model prompt prefix -#}
121 +{%- if add_generation_prompt -%}
122 + {{'<start_of_turn>model\n'}}
123 +{%- endif -%}
\ No newline at end of file
rust/src/chat/tests/templates/vllm_examples/tool_chat_template_gemma4.jinja new
+331
@@ -0,0 +1,331 @@
1 +{%- macro format_parameters(properties, required) -%}
2 + {%- set standard_keys = ['description', 'type', 'properties', 'required', 'nullable'] -%}
3 + {%- set ns = namespace(found_first=false) -%}
4 + {%- for key, value in properties | dictsort -%}
5 + {%- set add_comma = false -%}
6 + {%- if key not in standard_keys -%}
7 + {%- if ns.found_first %},{% endif -%}
8 + {%- set ns.found_first = true -%}
9 + {{ key }}:{
10 + {%- if value['description'] -%}
11 + description:<|"|>{{ value['description'] }}<|"|>
12 + {%- set add_comma = true -%}
13 + {%- endif -%}
14 + {%- if value['nullable'] %}
15 + {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}
16 + nullable:true
17 + {%- endif -%}
18 + {%- if value['type'] | upper == 'STRING' -%}
19 + {%- if value['enum'] -%}
20 + {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}
21 + enum:{{ format_argument(value['enum']) }}
22 + {%- endif -%}
23 + {%- elif value['type'] | upper == 'OBJECT' -%}
24 + ,properties:{
25 + {%- if value['properties'] is defined and value['properties'] is mapping -%}
26 + {{- format_parameters(value['properties'], value['required'] | default([])) -}}
27 + {%- elif value is mapping -%}
28 + {{- format_parameters(value, value['required'] | default([])) -}}
29 + {%- endif -%}
30 + }
31 + {%- if value['required'] -%}
32 + ,required:[
33 + {%- for item in value['required'] | default([]) -%}
34 + <|"|>{{- item -}}<|"|>
35 + {%- if not loop.last %},{% endif -%}
36 + {%- endfor -%}
37 + ]
38 + {%- endif -%}
39 + {%- elif value['type'] | upper == 'ARRAY' -%}
40 + {%- if value['items'] is mapping and value['items'] -%}
41 + ,items:{
42 + {%- set ns_items = namespace(found_first=false) -%}
43 + {%- for item_key, item_value in value['items'] | dictsort -%}
44 + {%- if item_value is not none -%}
45 + {%- if ns_items.found_first %},{% endif -%}
46 + {%- set ns_items.found_first = true -%}
47 + {%- if item_key == 'properties' -%}
48 + properties:{
49 + {%- if item_value is mapping -%}
50 + {{- format_parameters(item_value, value['items']['required'] | default([])) -}}
51 + {%- endif -%}
52 + }
53 + {%- elif item_key == 'required' -%}
54 + required:[
55 + {%- for req_item in item_value -%}
56 + <|"|>{{- req_item -}}<|"|>
57 + {%- if not loop.last %},{% endif -%}
58 + {%- endfor -%}
59 + ]
60 + {%- elif item_key == 'type' -%}
61 + {%- if item_value is string -%}
62 + type:{{ format_argument(item_value | upper) }}
63 + {%- else -%}
64 + type:{{ format_argument(item_value | map('upper') | list) }}
65 + {%- endif -%}
66 + {%- else -%}
67 + {{ item_key }}:{{ format_argument(item_value) }}
68 + {%- endif -%}
69 + {%- endif -%}
70 + {%- endfor -%}
71 + }
72 + {%- endif -%}
73 + {%- endif -%}
74 + {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}
75 + type:<|"|>{{ value['type'] | upper }}<|"|>}
76 + {%- endif -%}
77 + {%- endfor -%}
78 +{%- endmacro -%}
79 +{%- macro format_function_declaration(tool_data) -%}
80 + declaration:{{- tool_data['function']['name'] -}}{description:<|"|>{{- tool_data['function']['description'] -}}<|"|>
81 + {%- set params = tool_data['function']['parameters'] -%}
82 + {%- if params -%}
83 + ,parameters:{
84 + {%- if params['properties'] -%}
85 + properties:{ {{- format_parameters(params['properties'], params['required']) -}} },
86 + {%- endif -%}
87 + {%- if params['required'] -%}
88 + required:[
89 + {%- for item in params['required'] -%}
90 + <|"|>{{- item -}}<|"|>
91 + {{- ',' if not loop.last -}}
92 + {%- endfor -%}
93 + ],
94 + {%- endif -%}
95 + {%- if params['type'] -%}
96 + type:<|"|>{{- params['type'] | upper -}}<|"|>}
97 + {%- endif -%}
98 + {%- endif -%}
99 + {%- if 'response' in tool_data['function'] -%}
100 + {%- set response_declaration = tool_data['function']['response'] -%}
101 + ,response:{
102 + {%- if response_declaration['description'] -%}
103 + description:<|"|>{{- response_declaration['description'] -}}<|"|>,
104 + {%- endif -%}
105 + {%- if response_declaration['type'] | upper == 'OBJECT' -%}
106 + type:<|"|>{{- response_declaration['type'] | upper -}}<|"|>}
107 + {%- endif -%}
108 + {%- endif -%}
109 + }
110 +{%- endmacro -%}
111 +{%- macro format_argument(argument, escape_keys=True) -%}
112 + {%- if argument is string -%}
113 + {{- '<|"|>' + argument + '<|"|>' -}}
114 + {%- elif argument is boolean -%}
115 + {{- 'true' if argument else 'false' -}}
116 + {%- elif argument is mapping -%}
117 + {{- '{' -}}
118 + {%- set ns = namespace(found_first=false) -%}
119 + {%- for key, value in argument | dictsort -%}
120 + {%- if ns.found_first %},{% endif -%}
121 + {%- set ns.found_first = true -%}
122 + {%- if escape_keys -%}
123 + {{- '<|"|>' + key + '<|"|>' -}}
124 + {%- else -%}
125 + {{- key -}}
126 + {%- endif -%}
127 + :{{- format_argument(value, escape_keys=escape_keys) -}}
128 + {%- endfor -%}
129 + {{- '}' -}}
130 + {%- elif argument is sequence -%}
131 + {{- '[' -}}
132 + {%- for item in argument -%}
133 + {{- format_argument(item, escape_keys=escape_keys) -}}
134 + {%- if not loop.last %},{% endif -%}
135 + {%- endfor -%}
136 + {{- ']' -}}
137 + {%- else -%}
138 + {{- argument -}}
139 + {%- endif -%}
140 +{%- endmacro -%}
141 +{%- macro strip_thinking(text) -%}
142 + {%- set ns = namespace(result='') -%}
143 + {%- for part in text.split('<channel|>') -%}
144 + {%- if '<|channel>' in part -%}
145 + {%- set ns.result = ns.result + part.split('<|channel>')[0] -%}
146 + {%- else -%}
147 + {%- set ns.result = ns.result + part -%}
148 + {%- endif -%}
149 + {%- endfor -%}
150 + {{- ns.result | trim -}}
151 +{%- endmacro -%}
152 +
153 +{%- macro format_tool_response_block(tool_name, response) -%}
154 + {{- '<|tool_response>' -}}
155 + {%- if response is mapping -%}
156 + {{- 'response:' + tool_name + '{' -}}
157 + {%- for key, value in response | dictsort -%}
158 + {{- key -}}:{{- format_argument(value, escape_keys=False) -}}
159 + {%- if not loop.last %},{% endif -%}
160 + {%- endfor -%}
161 + {{- '}' -}}
162 + {%- else -%}
163 + {{- 'response:' + tool_name + '{value:' + format_argument(response, escape_keys=False) + '}' -}}
164 + {%- endif -%}
165 + {{- '<tool_response|>' -}}
166 +{%- endmacro -%}
167 +
168 +{%- set ns = namespace(prev_message_type=None) -%}
169 +{%- set loop_messages = messages -%}
170 +{{ bos_token }}
171 +{%- if (enable_thinking is defined and enable_thinking) or tools or messages[0]['role'] in ['system', 'developer'] -%}
172 + {{- '<|turn>system\n' -}}
173 +
174 + {%- if enable_thinking is defined and enable_thinking -%}
175 + {{- '<|think|>' -}}
176 + {%- set ns.prev_message_type = 'think' -%}
177 + {%- endif -%}
178 +
179 + {%- if messages[0]['role'] in ['system', 'developer'] -%}
180 + {{- messages[0]['content'] | trim -}}
181 + {%- set loop_messages = messages[1:] -%}
182 + {%- endif -%}
183 +
184 + {%- if tools -%}
185 + {%- for tool in tools %}
186 + {{- '<|tool>' -}}
187 + {{- format_function_declaration(tool) | trim -}}
188 + {{- '<tool|>' -}}
189 + {%- endfor %}
190 + {%- set ns.prev_message_type = 'tool' -%}
191 + {%- endif -%}
192 +
193 + {{- '<turn|>\n' -}}
194 +{%- endif %}
195 +
196 +{%- set ns_turn = namespace(last_user_idx=-1) -%}
197 +{%- for i in range(loop_messages | length) -%}
198 + {%- if loop_messages[i]['role'] == 'user' -%}
199 + {%- set ns_turn.last_user_idx = i -%}
200 + {%- endif -%}
201 +{%- endfor -%}
202 +
203 +{%- for message in loop_messages -%}
204 + {%- if message['role'] != 'tool' -%}
205 + {%- set ns.prev_message_type = None -%}
206 + {%- set role = 'model' if message['role'] == 'assistant' else message['role'] -%}
207 + {#- OpenAI may emit multiple assistant messages in one tool loop (user → asst → tool → asst → tool).
208 + Only the first of those should open <|turn>model; later ones continue the same model turn. -#}
209 + {%- set prev_nt = namespace(role=None, found=false) -%}
210 + {%- if loop.index0 > 0 -%}
211 + {%- for j in range(loop.index0 - 1, -1, -1) -%}
212 + {%- if not prev_nt.found -%}
213 + {%- if loop_messages[j]['role'] != 'tool' -%}
214 + {%- set prev_nt.role = loop_messages[j]['role'] -%}
215 + {%- set prev_nt.found = true -%}
216 + {%- endif -%}
217 + {%- endif -%}
218 + {%- endfor -%}
219 + {%- endif -%}
220 + {%- set continue_same_model_turn = (role == 'model' and prev_nt.role == 'assistant') -%}
221 + {%- if not continue_same_model_turn -%}
222 + {{- '<|turn>' + role + '\n' }}
223 + {%- endif -%}
224 +
225 + {%- if message.get('reasoning') and loop.index0 > ns_turn.last_user_idx and message.get('tool_calls') -%}
226 + {{- '<|channel>thought\n' + message['reasoning'] + '\n<channel|>'}}
227 + {%- endif -%}
228 +
229 + {%- if message['tool_calls'] -%}
230 + {%- for tool_call in message['tool_calls'] -%}
231 + {%- set function = tool_call['function'] -%}
232 + {{- '<|tool_call>call:' + function['name'] + '{' -}}
233 + {%- if function['arguments'] is mapping -%}
234 + {%- set ns_args = namespace(found_first=false) -%}
235 + {%- for key, value in function['arguments'] | dictsort -%}
236 + {%- if ns_args.found_first %},{% endif -%}
237 + {%- set ns_args.found_first = true -%}
238 + {{- key -}}:{{- format_argument(value, escape_keys=False) -}}
239 + {%- endfor -%}
240 + {%- elif function['arguments'] is string -%}
241 + {{- function['arguments'] -}}
242 + {%- endif -%}
243 + {{- '}<tool_call|>' -}}
244 + {%- endfor -%}
245 + {%- set ns.prev_message_type = 'tool_call' -%}
246 + {%- endif -%}
247 +
248 + {%- set ns_tr_out = namespace(flag=false) -%}
249 + {%- if message.get('tool_responses') -%}
250 + {#- Legacy: tool_responses embedded on the assistant message -#}
251 + {%- for tool_response in message['tool_responses'] -%}
252 + {{- format_tool_response_block(tool_response['name'] | default('unknown'), tool_response['response']) -}}
253 + {%- set ns_tr_out.flag = true -%}
254 + {%- set ns.prev_message_type = 'tool_response' -%}
255 + {%- endfor -%}
256 + {%- elif message.get('tool_calls') -%}
257 + {#- OpenAI Chat Completions: consecutive following messages with role "tool" (no break/continue; range scan) -#}
258 + {%- set ns_tool_scan = namespace(stopped=false) -%}
259 + {%- for k in range(loop.index0 + 1, loop_messages | length) -%}
260 + {%- if ns_tool_scan.stopped -%}
261 + {%- elif loop_messages[k]['role'] != 'tool' -%}
262 + {%- set ns_tool_scan.stopped = true -%}
263 + {%- else -%}
264 + {%- set follow = loop_messages[k] -%}
265 + {%- set ns_tname = namespace(name=follow.get('name') | default('unknown')) -%}
266 + {%- for tc in message['tool_calls'] -%}
267 + {%- if tc.get('id') == follow.get('tool_call_id') -%}
268 + {%- set ns_tname.name = tc['function']['name'] -%}
269 + {%- endif -%}
270 + {%- endfor -%}
271 + {%- set tool_body = follow.get('content') -%}
272 + {%- if tool_body is string -%}
273 + {{- format_tool_response_block(ns_tname.name, tool_body) -}}
274 + {%- elif tool_body is sequence and tool_body is not string -%}
275 + {%- set ns_txt = namespace(s='') -%}
276 + {%- for part in tool_body -%}
277 + {%- if part.get('type') == 'text' -%}
278 + {%- set ns_txt.s = ns_txt.s + (part.get('text') | default('')) -%}
279 + {%- endif -%}
280 + {%- endfor -%}
281 + {{- format_tool_response_block(ns_tname.name, ns_txt.s) -}}
282 + {%- else -%}
283 + {{- format_tool_response_block(ns_tname.name, tool_body) -}}
284 + {%- endif -%}
285 + {%- set ns_tr_out.flag = true -%}
286 + {%- set ns.prev_message_type = 'tool_response' -%}
287 + {%- endif -%}
288 + {%- endfor -%}
289 + {%- endif -%}
290 +
291 + {%- if message['content'] is string -%}
292 + {%- if role == 'model' -%}
293 + {{- strip_thinking(message['content']) -}}
294 + {%- else -%}
295 + {{- message['content'] | trim -}}
296 + {%- endif -%}
297 + {%- elif message['content'] is sequence -%}
298 + {%- for item in message['content'] -%}
299 + {%- if item['type'] == 'text' -%}
300 + {%- if role == 'model' -%}
301 + {{- strip_thinking(item['text']) -}}
302 + {%- else -%}
303 + {{- item['text'] | trim -}}
304 + {%- endif -%}
305 + {%- elif item['type'] == 'image' -%}
306 + {{- '\n\n<|image|>\n\n' -}}
307 + {%- set ns.prev_message_type = 'image' -%}
308 + {%- elif item['type'] == 'audio' -%}
309 + {{- '<|audio|>' -}}
310 + {%- set ns.prev_message_type = 'audio' -%}
311 + {%- elif item['type'] == 'video' -%}
312 + {{- '\n\n<|video|>\n\n' -}}
313 + {%- set ns.prev_message_type = 'video' -%}
314 + {%- endif -%}
315 + {%- endfor -%}
316 + {%- endif -%}
317 +
318 + {%- if not (ns_tr_out.flag and not message.get('content')) -%}
319 + {{- '<turn|>\n' -}}
320 + {%- endif -%}
321 + {%- endif -%}
322 +{%- endfor -%}
323 +
324 +{%- if add_generation_prompt -%}
325 + {%- if ns.prev_message_type != 'tool_response' -%}
326 + {{- '<|turn>model\n' -}}
327 + {%- endif -%}
328 + {%- if not enable_thinking | default(false) -%}
329 + {{- '<|channel>thought\n<channel|>' -}}
330 + {%- endif -%}
331 +{%- endif -%}
rust/src/chat/tests/templates/vllm_examples/tool_chat_template_glm4.jinja new
+54
@@ -0,0 +1,54 @@
1 +{%- set counter = namespace(index=0) -%}
2 +{%- if not tools is defined %}
3 + {%- set tools = none %}
4 +{%- endif %}
5 +
6 +{%- if messages and messages[0]['role'] == 'system' %}
7 + {%- set system_message = messages[0]['content']|trim %}
8 + {%- set messages = messages[1:] %}
9 +{%- else %}
10 + {%- set system_message = "You are a helpful assistant." %}
11 +{%- endif %}
12 +
13 +{%- if tools is not none %}
14 + {%- set tool_instruction %}
15 +You have access to the following tools. When you need to call a tool, you MUST use the following format:
16 +
17 +<tool_call>function_name
18 +<arg_key>parameter_name</arg_key>
19 +<arg_value>parameter_value</arg_value>
20 +</tool_call>
21 +
22 +Important rules:
23 +- Always wrap tool calls with <tool_call>...</tool_call> tags
24 +- Put the function name on the first line after <tool_call>
25 +- Use <arg_key> and <arg_value> tags for each parameter
26 +- If a parameter value is a string, keep it as-is. If it's a number or boolean, convert it appropriately
27 +- You can make multiple tool calls if needed
28 +- If no tool is suitable, respond with regular text
29 +
30 +Available tools:
31 +{% endset %}
32 + {{- tool_instruction + "\n\n" }}
33 + {%- for t in tools %}
34 + {{- t | tojson(indent=4) }}
35 + {{- "\n\n" }}
36 + {%- endfor %}
37 +{%- endif %}
38 +
39 +{%- for message in messages -%}
40 + {%- if message['role'] == 'user' -%}
41 + {{- '[Round ' + counter.index|string + ']\n问：' + message['content'] -}}
42 + {%- set counter.index = counter.index + 1 -%}
43 + {%- endif -%}
44 + {%- if message['role'] == 'assistant' -%}
45 + {{- '\n答：' + message['content'] -}}
46 + {%- if (loop.last and add_generation_prompt) or not loop.last -%}
47 + {{- '\n' -}}
48 + {%- endif -%}
49 + {%- endif -%}
50 +{%- endfor -%}
51 +
52 +{%- if add_generation_prompt and messages[-1]['role'] != 'assistant' -%}
53 + {{- '\n答：' -}}
54 +{%- endif -%}
rust/src/chat/tests/templates/vllm_examples/tool_chat_template_granite.jinja new
+36
@@ -0,0 +1,36 @@
1 +{%- if tools %}
2 + {{- '<|start_of_role|>available_tools<|end_of_role|>
3 +' }}
4 + {%- for tool in tools %}
5 + {{- tool | tojson(indent=4) }}
6 + {%- if not loop.last %}
7 + {{- '
8 +
9 +' }}
10 + {%- endif %}
11 + {%- endfor %}
12 + {{- '<|end_of_text|>
13 +' }}
14 +{%- endif %}
15 +
16 +{%- for message in messages %}
17 + {%- if message['role'] == 'system' %}
18 + {{- '<|start_of_role|>system<|end_of_role|>' + message['content'] + '<|end_of_text|>
19 +' }}
20 + {%- elif message['role'] == 'user' %}
21 + {{- '<|start_of_role|>user<|end_of_role|>' + message['content'] + '<|end_of_text|>
22 +' }}
23 + {%- elif message['role'] == 'assistant_tool_call' or (message['role'] == 'assistant' and message.tool_calls is defined) %}
24 + {{- '<|start_of_role|>assistant<|end_of_role|><|tool_call|>' + message.tool_calls|map(attribute='function')|list|tojson(indent=4) + '<|end_of_text|>
25 +' }}
26 + {%- elif message['role'] == 'assistant' %}
27 + {{- '<|start_of_role|>assistant<|end_of_role|>' + message['content'] + '<|end_of_text|>
28 +' }}
29 + {%- elif message['role'] == 'tool_response' or message['role'] == 'tool' %}
30 + {{- '<|start_of_role|>tool_response<|end_of_role|>' + message['content'] + '<|end_of_text|>
31 +' }}
32 + {%- endif %}
33 + {%- if loop.last and add_generation_prompt %}
34 + {{- '<|start_of_role|>assistant<|end_of_role|>' }}
35 + {%- endif %}
36 +{%- endfor %}
rust/src/chat/tests/templates/vllm_examples/tool_chat_template_granite_20b_fc.jinja new
+130
@@ -0,0 +1,130 @@
1 +{%- macro json_to_python_type(json_spec) %}
2 + {%- set basic_type_map = {
3 + "string": "str",
4 + "number": "float",
5 + "integer": "int",
6 + "boolean": "bool"
7 +} %}
8 +
9 + {%- if basic_type_map[json_spec.type] is defined %}
10 + {{- basic_type_map[json_spec.type] }}
11 + {%- elif json_spec.type == "array" %}
12 + {{- "list[" + json_to_python_type(json_spec|items) + "]" }}
13 + {%- elif json_spec.type == "object" %}
14 + {%- if json_spec.additionalProperties is defined %}
15 + {{- "dict[str, " + json_to_python_type(json_spec.additionalProperties) + ']' }}
16 + {%- else %}
17 + {{- "dict" }}
18 + {%- endif %}
19 + {%- elif json_spec.type is iterable %}
20 + {{- "Union[" }}
21 + {%- for t in json_spec.type %}
22 + {{- json_to_python_type({"type": t}) }}
23 + {%- if not loop.last %}
24 + {{- "," }}
25 + {%- endif %}
26 + {%- endfor %}
27 + {{- "]" }}
28 + {%- else %}
29 + {{- "Any" }}
30 + {%- endif %}
31 +{%- endmacro %}
32 +
33 +{%- if not full_function_description is defined %}
34 + {%- set full_function_description = false %}
35 +{%- endif %}
36 +
37 +{%- macro full_description(tool) %}
38 + {{- tool.name + '(' }}
39 + {%- if tool.parameters is defined %}
40 + {%- for param_name, param_fields in tool.parameters.properties|items %}
41 + {{- param_name + ": " + json_to_python_type(param_fields) }}
42 + {%- if not loop.last %}
43 + {{- ", " }}
44 + {%- endif %}
45 + {%- endfor %}
46 + {%- endif %}
47 + {{- ")" }}
48 + {%- if tool.return is defined %}
49 + {{- " -> " + json_to_python_type(tool.return) }}
50 + {%- endif %}
51 + {{- " - " + tool.description + "\n\n" }}
52 + {%- if tool.parameters is defined %}
53 + {%- for param_name, param_fields in tool.parameters.properties|items %}
54 + {%- if loop.first %}
55 + {{- " Args:\n" }}
56 + {%- endif %}
57 + {{- " " + param_name + "(" + json_to_python_type(param_fields) + "): " + param_fields.description|trim }}
58 + {%- endfor %}
59 + {%- endif %}
60 + {%- if tool.return is defined and tool.return.description is defined %}
61 + {{- "\n Returns:\n " + tool.return.description }}
62 + {%- endif %}
63 + {{- '"' }}
64 +{%- endmacro %}
65 +
66 +{%- macro simple_description(tool) %}
67 + {{- tool.description }}
68 +{%- endmacro %}
69 +
70 +{%- macro function_description(tool) %}
71 + {%- if full_function_description %}
72 + {{- full_description(tool) }}
73 + {%- else %}
74 + {{- simple_description(tool) }}
75 + {%- endif %}
76 +{%- endmacro %}
77 +
78 +{%- if messages[0]["role"] == "system" %}
79 + {%- set sys_prompt = messages[0]["content"] %}
80 + {%- set loop_messages = messages[1:] %}
81 +{%- else %}
82 + {%- set loop_messages = messages %}
83 + {% set sys_prompt = 'You are a helpful assistant with access to the following function calls. Your task is to understand the given conversation with function calls and responses and generate natural language response as the ASSISTANT to continue the conversation. You may use the following function calls to understand how to respond to the user query.' %}
84 +{%- endif %}
85 +
86 +{{ 'SYSTEM: ' + sys_prompt }}
87 +{% if tools is iterable and tools | length > 0 %}
88 +<|function_call_library|>
89 + {%- for tool in tools %}
90 + {%- if tool.function is defined %}
91 + {%- set tool = tool.function %}
92 + {%- endif %}
93 + {{- '{"name": "' + tool.name + '", ' }}
94 + {{- '"description": "' + function_description(tool) }}
95 + {{- ', "parameters": ' }}
96 + {%- if not tool.parameters is defined or tool.parameters.properties | length == 0 %}
97 + {{- "{}" }}
98 + {%- else %}
99 + {{- tool.parameters|tojson }}
100 + {%- endif %}
101 + {{- "}" }}
102 + {%- if not loop.last %}
103 + {{- "\n" }}
104 + {%- endif %}
105 + {%- endfor %}
106 +If none of the functions are relevant or the given question lacks the parameters required by the function, please output \"<function_call> {\"name\": \"no_function\", \"arguments\": {}}\".
107 +{%- endif %}
108 +
109 +
110 +
111 +{% for message in messages %}
112 + {% if message['role'] == 'user' %}
113 + {{- '\nUSER: ' + message['content'] }}
114 + {% elif message['role'] == 'assistant' and message.tool_calls is defined %}
115 + {{- '\nASSISTANT:' }}
116 + {% for tc in message.tool_calls %}
117 + {{- '<function_call> ' + {'name': tc.function.name, 'arguments': tc.function.arguments}|tojson }}
118 + {% endfor %}
119 + {{- '<|endoftext|>' }}
120 + {% elif message['role'] == 'assistant' %}
121 + {{- '\nASSISTANT: ' + message['content'] + ' <|endoftext|>' }}
122 + {% elif message['role'] == 'tool' %}
123 + {{- '<function_response> ' + message['content'] }}
124 + {%- else %}
125 + {{- raise_exception("Unexpected combination of role and message content") }}
126 + {% endif %}
127 + {% if loop.last and add_generation_prompt %}
128 + {{- '\nASSISTANT: ' }}
129 + {% endif %}
130 +{% endfor %}
rust/src/chat/tests/templates/vllm_examples/tool_chat_template_hermes.jinja new
+130
@@ -0,0 +1,130 @@
1 +{%- macro json_to_python_type(json_spec) %}
2 + {%- set basic_type_map = {
3 + "string": "str",
4 + "number": "float",
5 + "integer": "int",
6 + "boolean": "bool"
7 +} %}
8 +
9 + {%- if basic_type_map[json_spec.type] is defined %}
10 + {{- basic_type_map[json_spec.type] }}
11 + {%- elif json_spec.type == "array" %}
12 + {{- "list[" + json_to_python_type(json_spec|items) + "]" }}
13 + {%- elif json_spec.type == "object" %}
14 + {%- if json_spec.additionalProperties is defined %}
15 + {{- "dict[str, " + json_to_python_type(json_spec.additionalProperties) + ']' }}
16 + {%- else %}
17 + {{- "dict" }}
18 + {%- endif %}
19 + {%- elif json_spec.type is iterable %}
20 + {{- "Union[" }}
21 + {%- for t in json_spec.type %}
22 + {{- json_to_python_type({"type": t}) }}
23 + {%- if not loop.last %}
24 + {{- "," }}
25 + {%- endif %}
26 + {%- endfor %}
27 + {{- "]" }}
28 + {%- else %}
29 + {{- "Any" }}
30 + {%- endif %}
31 +{%- endmacro %}
32 +
33 +
34 +{{- bos_token }}
35 +{{- "<|im_start|>system\nYou are a function calling AI model. You are provided with function signatures within <tools></tools> XML tags. You may call one or more functions to assist with the user query. Don't make assumptions about what values to plug into functions. Here are the available tools: <tools> " }}
36 +{%- if tools is iterable and tools | length > 0 %}
37 + {%- for tool in tools %}
38 + {%- if tool.function is defined %}
39 + {%- set tool = tool.function %}
40 + {%- endif %}
41 + {{- '{"type": "function", "function": ' }}
42 + {{- '{"name": "' + tool.name + '", ' }}
43 + {{- '"description": "' + tool.name + '(' }}
44 + {%- for param_name, param_fields in tool.parameters.properties|items %}
45 + {{- param_name + ": " + json_to_python_type(param_fields) }}
46 + {%- if not loop.last %}
47 + {{- ", " }}
48 + {%- endif %}
49 + {%- endfor %}
50 + {{- ")" }}
51 + {%- if tool.return is defined %}
52 + {{- " -> " + json_to_python_type(tool.return) }}
53 + {%- endif %}
54 + {{- " - " + tool.description + "\n\n" }}
55 + {%- for param_name, param_fields in tool.parameters.properties|items %}
56 + {%- if loop.first %}
57 + {{- " Args:\n" }}
58 + {%- endif %}
59 + {{- " " + param_name + "(" + json_to_python_type(param_fields) + "): " + param_fields.description|trim }}
60 + {%- endfor %}
61 + {%- if tool.return is defined and tool.return.description is defined %}
62 + {{- "\n Returns:\n " + tool.return.description }}
63 + {%- endif %}
64 + {{- '"' }}
65 + {{- ', "parameters": ' }}
66 + {%- if tool.parameters.properties | length == 0 %}
67 + {{- "{}" }}
68 + {%- else %}
69 + {{- tool.parameters|tojson }}
70 + {%- endif %}
71 + {{- "}" }}
72 + {%- if not loop.last %}
73 + {{- "\n" }}
74 + {%- endif %}
75 + {%- endfor %}
76 +{%- endif %}
77 +{{- " </tools>" }}
78 +{{- 'Use the following pydantic model json schema for each tool call you will make: {"properties": {"name": {"title": "Name", "type": "string"}, "arguments": {"title": "Arguments", "type": "object"}}, "required": ["name", "arguments"], "title": "FunctionCall", "type": "object"}}
79 +' }}
80 +{{- "For each function call return a json object with function name and arguments within <tool_call></tool_call> XML tags as follows:
81 +" }}
82 +{{- "<tool_call>
83 +" }}
84 +{{- '{"name": <function-name>, "arguments": <args-dict>}
85 +' }}
86 +{{- '</tool_call><|im_end|>' }}
87 +{%- for message in messages %}
88 + {%- if message.role == "user" or message.role == "system" or (message.role == "assistant" and message.tool_calls is not defined) %}
89 + {{- '<|im_start|>' + message.role + '\n' + message.content + '<|im_end|>' + '\n' }}
90 + {%- elif message.role == "assistant" and message.tool_calls is defined %}
91 + {{- '<|im_start|>' + message.role }}
92 + {%- for tool_call in message.tool_calls %}
93 + {{- '\n<tool_call>\n' }}
94 + {%- if tool_call.function is defined %}
95 + {%- set tool_call = tool_call.function %}
96 + {%- endif %}
97 + {{- '{' }}
98 + {{- '"name": "' }}
99 + {{- tool_call.name }}
100 + {{- '"' }}
101 + {%- if tool_call.arguments is defined %}
102 + {{- ', ' }}
103 + {{- '"arguments": ' }}
104 + {{- tool_call.arguments|tojson }}
105 + {%- endif %}
106 + {{- '}' }}
107 + {{- '\n</tool_call>' }}
108 + {%- endfor %}
109 + {{- '<|im_end|>\n' }}
110 + {%- elif message.role == "tool" %}
111 + {%- if loop.previtem and loop.previtem.role != "tool" %}
112 + {{- '<|im_start|>tool\n' }}
113 + {%- endif %}
114 + {{- '<tool_response>\n' }}
115 + {{- message.content }}
116 + {%- if not loop.last %}
117 + {{- '\n</tool_response>\n' }}
118 + {%- else %}
119 + {{- '\n</tool_response>' }}
120 + {%- endif %}
121 + {%- if not loop.last and loop.nextitem.role != "tool" %}
122 + {{- '<|im_end|>' }}
123 + {%- elif loop.last %}
124 + {{- '<|im_end|>' }}
125 + {%- endif %}
126 + {%- endif %}
127 +{%- endfor %}
128 +{%- if add_generation_prompt %}
129 + {{- '<|im_start|>assistant\n' }}
130 +{%- endif %}
rust/src/chat/tests/templates/vllm_examples/tool_chat_template_hunyuan_a13b.jinja new
+113
@@ -0,0 +1,113 @@
1 +{% set loop_messages = messages %}
2 +{% if tools %}
3 + {% set weekday_map = {'Monday': '星期一', 'Tuesday': '星期二', 'Wednesday': '星期三', 'Thursday': '星期四', 'Friday': '星期五', 'Saturday': '星期六', 'Sunday': '星期日'} %}
4 + {% set weekday_cn = weekday_map[strftime_now('%A')] %}
5 + {% set datetime_str = strftime_now('%Y-%m-%d %H:%M:%S') %}
6 + {% set datetime_str = datetime_str + ' ' + weekday_cn %}
7 + {% for message in loop_messages %}
8 + {% if 'content' in message %}
9 + {% set content = message['content'] %}
10 + {% else %}
11 + {% set content = '' %}
12 + {% endif %}
13 + {% if loop.index0 == 0 %}
14 + {% set content_tmp = '你是一位函数组合专家。你会得到一个问题和一组可能的函数。根据问题，你需要进行一个或多个函数/工具调用以实现目的。
15 +如果没有一个函数可以使用，请直接使用自然语言回复用户，以助手：开头。
16 +如果给定的问题缺少函数所需的参数，请使用自然语言进行提问，向用户询问必要信息，以助手：开头。
17 +如果调用结果已经足够回答用户问题，请对历史结果进行总结，使用自然语言回复用户，以助手：开头。
18 +你应该只在工具调用部分返回函数调用。如果你决定调用任何函数，你必须将其格式化为<tool_calls>[{"name": "func_name1", "arguments": {"argument1": "value1", "argument2": "value2"}},...]</tool_calls>。你不应该在回复中包含任何其他文本。以下是你可以调用的函数列表，格式为JSON。
19 +' %}
20 + {% set content_tmp = content_tmp + '
21 +' + tools | tojson + '
22 +' %}
23 + {% if message['role'] == 'system' %}
24 + {% set content_tmp = content_tmp + '
25 +额外要求：
26 +' + content + '
27 +
28 +如果你决定返回函数调用，请将其格式化为<tool_calls>[{"name": "func_name1", "arguments": {"argument1": "value1", "argument2": "value2"}},...]</tool_calls>，不得包含其他文本。如果额外要求里有格式要求，请忽略，以此处为准。
29 +否则，请参考开头说的三种情况，以助手：开头进行回复。
30 +
31 +如果额外要求里有时间信息，就以额外要求里的时间为准，否则，参考当前时间：' + datetime_str %}
32 + {% set content = '<|startoftext|>' + content_tmp + '<|extra_4|>' %}
33 + {% elif message['role'] == 'user' %}
34 + {% set content_tmp = content_tmp + '
35 +如果你决定返回函数调用，请将其格式化为<tool_calls>[{"name": "func_name1", "arguments": {"argument1": "value1", "argument2": "value2"}},...]</tool_calls>，不得包含其他文本。
36 +否则，请参考开头说的三种情况，以助手：开头进行回复。
37 +
38 +当前时间：' + datetime_str %}
39 + {% set content_tmp = '<|startoftext|>' + content_tmp + '<|extra_4|>'%}
40 + {% set content = content_tmp + '用户：' + content + '<|extra_0|>' %}
41 + {% endif %}
42 + {% else %}
43 + {% if message['role'] == 'user' %}
44 + {% set content = '用户：' + content + '<|extra_0|>' %}
45 + {% elif message['role'] == 'assistant' %}
46 + {% if 'tool_calls' in message %}
47 + {% set tool_calls = message['tool_calls'] %}
48 + {% set ns = namespace(tool_calls="[") %}
49 + {% for tool_call in tool_calls %}
50 + {% set function = tool_call['function'] %}
51 + {% set name = function['name'] %}
52 + {% set ns.tool_calls = ns.tool_calls + '{"name": "' + name + '", '%}
53 + {% set arguments = function['arguments'] %}
54 + {% if arguments is not string %}
55 + {% set arguments = arguments | tojson %}
56 + {% endif %}
57 + {% set ns.tool_calls = ns.tool_calls + '"arguments": ' + arguments + '}' %}
58 + {% if not loop.last %}
59 + {% set ns.tool_calls = ns.tool_calls + ', '%}
60 + {% endif %}
61 + {% endfor %}
62 + {% set ns.tool_calls = ns.tool_calls + ']' %}
63 + {% set content = content + '<tool_calls>' + ns.tool_calls + '</tool_calls>' %}
64 + {% else %}
65 + {% set content = '助手：' + content %}
66 + {% endif %}
67 + {% set content = content + '<|eos|>' %}
68 + {% elif message['role'] == 'tool' %}
69 + {% if content is not string %}
70 + {set content = content | tojson }
71 + {% endif %}
72 + {% set content = '<tool_response>' + content + '</tool_response>' %}
73 + {% set content = content + '<|extra_0|>' %}
74 + {% endif %}
75 + {% endif %}
76 + {{- content -}}
77 + {% endfor %}
78 +{% else %}
79 + {% set context = {'has_head': true} %}
80 + {% for message in loop_messages %}
81 + {% if 'content' in message %}
82 + {% set content = message['content'] %}
83 + {% else %}
84 + {% set content = '' %}
85 + {% endif %}
86 + {% if loop.index0 == 0 %}
87 + {% if content == '' %}
88 + {% set _ = context.update({'has_head': false}) %}
89 + {% elif message['role'] == 'system' %}
90 + {% set content = '<|startoftext|>' + content + '<|extra_4|>' %}
91 + {% endif %}
92 + {% endif %}
93 + {% if message['role'] == 'user' %}
94 + {% if loop.index0 == 1 and not context.has_head %}
95 + {% set content = '<|startoftext|>' + content %}
96 + {% endif %}
97 + {% if loop.index0 == 1 and context.has_head %}
98 + {% set content = content + '<|extra_0|>' %}
99 + {% else %}
100 + {% set content = '<|startoftext|>' + content + '<|extra_0|>' %}
101 + {% endif %}
102 + {% elif message['role'] == 'assistant' %}
103 + {% set content = content + '<|eos|>' %}
104 + {% elif message['role'] == 'tool' %}
105 + {% set content = content + '<|extra_0|>' %}
106 + {% endif %}
107 + {{- content -}}
108 + {% endfor %}
109 +{% endif %}
110 +{%- if enable_thinking is defined and enable_thinking is false %}
111 + {{- '<think>\n\n</think>\n' }}
112 +{%- endif %}
113 +
rust/src/chat/tests/templates/vllm_examples/tool_chat_template_internlm2_tool.jinja new
+60
@@ -0,0 +1,60 @@
1 +{%- if messages[0]["role"] == "system" %}
2 + {%- set system_message = messages[0]["content"] %}
3 + {%- set loop_messages = messages[1:] %}
4 +{%- else %}
5 + {%- set loop_messages = messages %}
6 +{%- endif %}
7 +
8 +{%- if not tools is defined %}
9 + {%- set tools = none %}
10 +{%- endif %}
11 +
12 +{{- bos_token }}
13 +{%- if system_message is defined %}
14 +{{- "<|im_start|>system\n" + system_message + "<|im_end|>\n" }}
15 +{%- endif %}
16 +
17 +{%- if tools is not none %}
18 + {{- "<|im_start|>system name=<|plugin|>\n[" }}
19 + {%- for tool in tools %}
20 + {{- tool.function|tojson }}
21 + {%- if not loop.last %}
22 + {{- ", " }}
23 + {%- else %}
24 + {{- "]" }}
25 + {%- endif %}
26 + {%- endfor %}
27 + {{- "<|im_end|>\n" }}
28 +{%- endif %}
29 +
30 +{%- for message in loop_messages %}
31 + {%- if message["role"] == "user" %}
32 + {{- "<|im_start|>user\n" + message["content"] + "<|im_end|>\n"}}
33 + {%- elif message.tool_calls is defined and message.tool_calls is not none %}
34 + {%- set content = message["content"] if message["content"] else "" %}
35 + {{- "<|im_start|>assistant\n" + content }}
36 + {%- for tool_call in message.tool_calls %}
37 + {%- set function=tool_call.function %}
38 + {{- "<|action_start|><|plugin|>\n" }}
39 + {{- '{"name": "' + function.name + '", '}}
40 + {{- '"arguments": ' + function.arguments|tojson + '}' }}
41 + {{- "<|action_end|>" }}
42 + {%- endfor %}
43 + {{- "<|im_end|>\n" }}
44 + {%- elif message["role"] == "assistant" %}
45 + {{- "<|im_start|>assistant\n" + message["content"] + "<|im_end|>\n"}}
46 + {%- elif message["role"] == "tool_results" or message["role"] == "tool" or message["role"] == "function" %}
47 + {%- if message.content is defined and message.content.content is defined %}
48 + {%- set content = message.content.content %}
49 + {%- else %}
50 + {%- set content = message.content %}
51 + {%- endif %}
52 + {{- "<|im_start|>environment name=<|plugin|>\n" + content|string + "<|im_end|>\n" }}
53 + {%- else %}
54 + {{- raise_exception("Only user and assistant and tool_results and tool and function roles are supported, with the exception of an initial optional system message!") }}
55 + {%- endif %}
56 +{%- endfor %}
57 +
58 +{%- if add_generation_prompt %}
59 +{{- '<|im_start|>assistant\n' }}
60 +{%- endif %}
\ No newline at end of file
rust/src/chat/tests/templates/vllm_examples/tool_chat_template_llama3.1_json.jinja new
+120
@@ -0,0 +1,120 @@
1 +{{- bos_token }}
2 +{%- if custom_tools is defined %}
3 + {%- set tools = custom_tools %}
4 +{%- endif %}
5 +{%- if not tools_in_user_message is defined %}
6 + {#- Llama 3.1 doesn't pass all tests if the tools are in the system prompt #}
7 + {%- set tools_in_user_message = true %}
8 +{%- endif %}
9 +{%- if not date_string is defined %}
10 + {%- if strftime_now is defined %}
11 + {%- set date_string = strftime_now("%d %b %Y") %}
12 + {%- else %}
13 + {%- set date_string = "26 Jul 2024" %}
14 + {%- endif %}
15 +{%- endif %}
16 +{%- if not tools is defined %}
17 + {%- set tools = none %}
18 +{%- endif %}
19 +
20 +{#- This block extracts the system message, so we can slot it into the right place. #}
21 +{%- if messages[0]['role'] == 'system' %}
22 + {%- if messages[0]['content'] is string %}
23 + {%- set system_message = messages[0]['content']|trim %}
24 + {%- else %}
25 + {%- set system_message = messages[0]['content'][0]['text']|trim %}
26 + {%- endif %}
27 + {%- set messages = messages[1:] %}
28 +{%- else %}
29 + {%- if tools is not none %}
30 + {%- set system_message = "You are a helpful assistant with tool calling capabilities. Only reply with a tool call if the function exists in the library provided by the user. If it doesn't exist, just reply directly in natural language. When you receive a tool call response, use the output to format an answer to the original user question." %}
31 + {%- else %}
32 + {%- set system_message = "" %}
33 + {%- endif %}
34 +{%- endif %}
35 +
36 +{#- System message #}
37 +{{- "<|start_header_id|>system<|end_header_id|>\n\n" }}
38 +{%- if tools is not none %}
39 + {{- "Environment: ipython\n" }}
40 +{%- endif %}
41 +{{- "Cutting Knowledge Date: December 2023\n" }}
42 +{{- "Today Date: " + date_string + "\n\n" }}
43 +{%- if tools is not none and not tools_in_user_message %}
44 + {{- "You have access to the following functions. To call a function, please respond with JSON for a function call. " }}
45 + {{- 'Respond in the format {"name": function name, "parameters": dictionary of argument name and its value}. ' }}
46 + {{- "Do not use variables.\n\n" }}
47 + {%- for t in tools %}
48 + {{- t | tojson(indent=4) }}
49 + {{- "\n\n" }}
50 + {%- endfor %}
51 +{%- endif %}
52 +{{- system_message }}
53 +{{- "<|eot_id|>" }}
54 +
55 +{#- Custom tools are passed in a user message with some extra guidance #}
56 +{%- if tools_in_user_message and not tools is none %}
57 + {#- Extract the first user message so we can plug it in here #}
58 + {%- if messages | length != 0 %}
59 + {%- if messages[0]['content'] is string %}
60 + {%- set first_user_message = messages[0]['content']|trim %}
61 + {%- else %}
62 + {%- set first_user_message = messages[0]['content'] | selectattr('type', 'equalto', 'text') | map(attribute='text') | map('trim') | join('\n') %}
63 + {%- endif %}
64 + {%- set messages = messages[1:] %}
65 + {%- else %}
66 + {{- raise_exception("Cannot put tools in the first user message when there's no first user message!") }}
67 + {%- endif %}
68 + {{- '<|start_header_id|>user<|end_header_id|>\n\n' -}}
69 + {{- "Given the following functions, please respond with a JSON for a function call " }}
70 + {{- "with its proper arguments that best answers the given prompt.\n\n" }}
71 + {{- 'Respond in the format {"name": function name, "parameters": dictionary of argument name and its value}. ' }}
72 + {{- "Do not use variables.\n\n" }}
73 + {%- for t in tools %}
74 + {{- t | tojson(indent=4) }}
75 + {{- "\n\n" }}
76 + {%- endfor %}
77 + {{- first_user_message + "<|eot_id|>"}}
78 +{%- endif %}
79 +
80 +{%- for message in messages %}
81 + {%- if not (message.role == 'ipython' or message.role == 'tool' or 'tool_calls' in message) %}
82 + {{- '<|start_header_id|>' + message['role'] + '<|end_header_id|>\n\n' }}
83 + {%- if message['content'] is string %}
84 + {{- message['content'] | trim}}
85 + {%- else %}
86 + {%- for content in message['content'] %}
87 + {%- if content['type'] == 'text' %}
88 + {{- content['text'] | trim }}
89 + {%- endif %}
90 + {%- endfor %}
91 + {%- endif %}
92 + {{- '<|eot_id|>' }}
93 + {%- elif 'tool_calls' in message %}
94 + {%- if not message.tool_calls|length == 1 %}
95 + {{- raise_exception("This model only supports single tool-calls at once!") }}
96 + {%- endif %}
97 + {%- set tool_call = message.tool_calls[0].function %}
98 + {{- '<|start_header_id|>assistant<|end_header_id|>\n\n' -}}
99 + {{- '{"name": "' + tool_call.name + '", ' }}
100 + {{- '"parameters": ' }}
101 + {{- tool_call.arguments | tojson }}
102 + {{- "}" }}
103 + {{- "<|eot_id|>" }}
104 + {%- elif message.role == "tool" or message.role == "ipython" %}
105 + {{- "<|start_header_id|>ipython<|end_header_id|>\n\n" }}
106 + {%- if message.content is string %}
107 + {{- { "output": message.content } | tojson }}
108 + {%- else %}
109 + {%- for content in message['content'] %}
110 + {%- if content['type'] == 'text' %}
111 + {{- { "output": content['text'] } | tojson }}
112 + {%- endif %}
113 + {%- endfor %}
114 + {%- endif %}
115 + {{- "<|eot_id|>" }}
116 + {%- endif %}
117 +{%- endfor %}
118 +{%- if add_generation_prompt %}
119 + {{- '<|start_header_id|>assistant<|end_header_id|>\n\n' }}
120 +{%- endif %}
rust/src/chat/tests/templates/vllm_examples/tool_chat_template_llama3.2_json.jinja new
+133
@@ -0,0 +1,133 @@
1 +{{- bos_token }}
2 +{%- if custom_tools is defined %}
3 + {%- set tools = custom_tools %}
4 +{%- endif %}
5 +{%- if not tools_in_user_message is defined %}
6 + {%- set tools_in_user_message = false %}
7 +{%- endif %}
8 +{%- if not date_string is defined %}
9 + {%- if strftime_now is defined %}
10 + {%- set date_string = strftime_now("%d %b %Y") %}
11 + {%- else %}
12 + {%- set date_string = "26 Jul 2024" %}
13 + {%- endif %}
14 +{%- endif %}
15 +{%- if not tools is defined %}
16 + {%- set tools = none %}
17 +{%- endif %}
18 +
19 +{#- Find out if there are any images #}
20 +{% set image_ns = namespace(has_images=false) %}
21 +{%- for message in messages %}
22 + {%- for content in message['content'] %}
23 + {%- if content['type'] == 'image' %}
24 + {%- set image_ns.has_images = true %}
25 + {%- endif %}
26 + {%- endfor %}
27 +{%- endfor %}
28 +
29 +{#- This block extracts the system message, so we can slot it into the right place. #}
30 +{%- if messages[0]['role'] == 'system' %}
31 + {%- if messages[0]['content'] is string %}
32 + {%- set system_message = messages[0]['content']|trim %}
33 + {%- else %}
34 + {%- set system_message = messages[0]['content'][0]['text']|trim %}
35 + {%- endif %}
36 + {%- set messages = messages[1:] %}
37 +{%- else %}
38 + {%- if tools is not none %}
39 + {%- set system_message = "You are a helpful assistant with tool calling capabilities. Only reply with a tool call if the function exists in the library provided by the user. If it doesn't exist, just reply directly in natural language. When you receive a tool call response, use the output to format an answer to the original user question." %}
40 + {%- else %}
41 + {%- set system_message = "" %}
42 + {%- endif %}
43 +{%- endif %}
44 +
45 +{#- System message if there are no images, if the user supplied one, or if tools are used (default tool system message) #}
46 +{%- if system_message or not image_ns.has_images %}
47 + {{- "<|start_header_id|>system<|end_header_id|>\n\n" }}
48 + {%- if tools is not none %}
49 + {{- "Environment: ipython\n" }}
50 + {%- endif %}
51 + {{- "Cutting Knowledge Date: December 2023\n" }}
52 + {{- "Today Date: " + date_string + "\n\n" }}
53 + {%- if tools is not none and not tools_in_user_message %}
54 + {{- "You have access to the following functions. To call a function, please respond with JSON for a function call. " }}
55 + {{- 'Respond in the format {"name": function name, "parameters": dictionary of argument name and its value}. ' }}
56 + {{- "Do not use variables.\n\n" }}
57 + {%- for t in tools %}
58 + {{- t | tojson(indent=4) }}
59 + {{- "\n\n" }}
60 + {%- endfor %}
61 + {%- endif %}
62 + {{- system_message }}
63 + {{- "<|eot_id|>" }}
64 +{%- endif %}
65 +
66 +{#- Custom tools are passed in a user message with some extra guidance #}
67 +{%- if tools_in_user_message and not tools is none %}
68 + {#- Extract the first user message so we can plug it in here #}
69 + {%- if messages | length != 0 %}
70 + {%- if messages[0]['content'] is string %}
71 + {%- set first_user_message = messages[0]['content']|trim %}
72 + {%- else %}
73 + {%- set first_user_message = messages[0]['content'] | selectattr('type', 'equalto', 'text') | map(attribute='text') | map('trim') | join('\n') %}
74 + {%- endif %}
75 + {%- set messages = messages[1:] %}
76 + {%- else %}
77 + {{- raise_exception("Cannot put tools in the first user message when there's no first user message!") }}
78 + {%- endif %}
79 + {{- '<|start_header_id|>user<|end_header_id|>\n\n' -}}
80 + {{- "Given the following functions, please respond with a JSON for a function call " }}
81 + {{- "with its proper arguments that best answers the given prompt.\n\n" }}
82 + {{- 'Respond in the format {"name": function name, "parameters": dictionary of argument name and its value}. ' }}
83 + {{- "Do not use variables.\n\n" }}
84 + {%- for t in tools %}
85 + {{- t | tojson(indent=4) }}
86 + {{- "\n\n" }}
87 + {%- endfor %}
88 + {{- first_user_message + "<|eot_id|>"}}
89 +{%- endif %}
90 +
91 +{%- for message in messages %}
92 + {%- if not (message.role == 'ipython' or message.role == 'tool' or 'tool_calls' in message) %}
93 + {{- '<|start_header_id|>' + message['role'] + '<|end_header_id|>\n\n' }}
94 + {%- if message['content'] is string %}
95 + {{- message['content'] | trim}}
96 + {%- else %}
97 + {%- for content in message['content'] %}
98 + {%- if content['type'] == 'image' %}
99 + {{- '<|image|>' }}
100 + {%- elif content['type'] == 'text' %}
101 + {{- content['text'] | trim }}
102 + {%- endif %}
103 + {%- endfor %}
104 + {%- endif %}
105 + {{- '<|eot_id|>' }}
106 + {%- elif 'tool_calls' in message %}
107 + {%- if not message.tool_calls|length == 1 %}
108 + {{- raise_exception("This model only supports single tool-calls at once!") }}
109 + {%- endif %}
110 + {%- set tool_call = message.tool_calls[0].function %}
111 + {{- '<|start_header_id|>assistant<|end_header_id|>\n\n' -}}
112 + {{- '{"name": "' + tool_call.name + '", ' }}
113 + {{- '"parameters": ' }}
114 + {{- tool_call.arguments | tojson }}
115 + {{- "}" }}
116 + {{- "<|eot_id|>" }}
117 + {%- elif message.role == "tool" or message.role == "ipython" %}
118 + {{- "<|start_header_id|>ipython<|end_header_id|>\n\n" }}
119 + {%- if message.content is string %}
120 + {{- { "output": message.content } | tojson }}
121 + {%- else %}
122 + {%- for content in message['content'] %}
123 + {%- if content['type'] == 'text' %}
124 + {{- { "output": content['text'] } | tojson }}
125 + {%- endif %}
126 + {%- endfor %}
127 + {%- endif %}
128 + {{- "<|eot_id|>" }}
129 + {%- endif %}
130 +{%- endfor %}
131 +{%- if add_generation_prompt %}
132 + {{- '<|start_header_id|>assistant<|end_header_id|>\n\n' }}
133 +{%- endif %}
rust/src/chat/tests/templates/vllm_examples/tool_chat_template_llama3.2_pythonic.jinja new
+98
@@ -0,0 +1,98 @@
1 +{{- bos_token }}
2 +{%- if custom_tools is defined %}
3 + {%- set tools = custom_tools %}
4 +{%- endif %}
5 +{%- if not tools_in_user_message is defined %}
6 + {%- set tools_in_user_message = false %}
7 +{%- endif %}
8 +{%- if not date_string is defined %}
9 + {%- if strftime_now is defined %}
10 + {%- set date_string = strftime_now("%d %b %Y") %}
11 + {%- else %}
12 + {%- set date_string = "26 Jul 2024" %}
13 + {%- endif %}
14 +{%- endif %}
15 +{%- if not tools is defined %}
16 + {%- set tools = none %}
17 +{%- endif %}
18 +
19 +{#- This block extracts the system message, so we can slot it into the right place. #}
20 +{%- if messages[0]['role'] == 'system' %}
21 + {%- set system_message = messages[0]['content']|trim %}
22 + {%- set messages = messages[1:] %}
23 +{%- else %}
24 + {%- set system_message = "You are a helpful assistant with tool calling capabilities. Only reply with a tool call if the function exists in the library provided by the user. If it doesn't exist, just reply directly in natural language. When you receive a tool call response, use the output to format an answer to the original user question." %}
25 +{%- endif %}
26 +
27 +{#- System message #}
28 +{{- "<|start_header_id|>system<|end_header_id|>\n\n" }}
29 +{%- if tools is not none %}
30 + {{- "Environment: ipython\n" }}
31 +{%- endif %}
32 +{{- "Cutting Knowledge Date: December 2023\n" }}
33 +{{- "Today Date: " + date_string + "\n\n" }}
34 +{%- if tools is not none and not tools_in_user_message %}
35 + {{- "You have access to the following functions. To call functions, please respond with a python list of the calls. " }}
36 + {{- 'Respond in the format [func_name1(params_name1=params_value1, params_name2=params_value2...), func_name2(params)] ' }}
37 + {{- "Do not use variables.\n\n" }}
38 + {%- for t in tools %}
39 + {{- t | tojson(indent=4) }}
40 + {{- "\n\n" }}
41 + {%- endfor %}
42 +{%- endif %}
43 +{{- system_message }}
44 +{{- "<|eot_id|>" }}
45 +
46 +{#- Custom tools are passed in a user message with some extra guidance #}
47 +{%- if tools_in_user_message and not tools is none %}
48 + {#- Extract the first user message so we can plug it in here #}
49 + {%- if messages | length != 0 %}
50 + {%- set first_user_message = messages[0]['content']|trim %}
51 + {%- set messages = messages[1:] %}
52 + {%- else %}
53 + {{- raise_exception("Cannot put tools in the first user message when there's no first user message!") }}
54 + {%- endif %}
55 + {{- '<|start_header_id|>user<|end_header_id|>\n\n' -}}
56 + {{- "Given the following functions, please respond with a python list for function calls " }}
57 + {{- "with their proper arguments to best answer the given prompt.\n\n" }}
58 + {{- 'Respond in the format [func_name1(params_name1=params_value1, params_name2=params_value2...), func_name2(params)] ' }}
59 + {{- "Do not use variables.\n\n" }}
60 + {%- for t in tools %}
61 + {{- t | tojson(indent=4) }}
62 + {{- "\n\n" }}
63 + {%- endfor %}
64 + {{- first_user_message + "<|eot_id|>"}}
65 +{%- endif %}
66 +
67 +{%- for message in messages %}
68 + {%- if not (message.role == 'ipython' or message.role == 'tool' or 'tool_calls' in message) %}
69 + {{- '<|start_header_id|>' + message['role'] + '<|end_header_id|>\n\n'+ message['content'] | trim + '<|eot_id|>' }}
70 + {%- elif 'tool_calls' in message %}
71 + {{- '<|start_header_id|>assistant<|end_header_id|>\n\n[' -}}
72 + {%- for tool_call in message.tool_calls %}
73 + {%- if tool_call.function is defined %}
74 + {%- set tool_call = tool_call.function %}
75 + {%- endif %}
76 + {{- tool_call.name + '(' -}}
77 + {%- for param in tool_call.arguments %}
78 + {{- param + '=' -}}
79 + {{- "%s" | format(tool_call.arguments[param]) -}}
80 + {% if not loop.last %}, {% endif %}
81 + {%- endfor %}
82 + {{- ')' -}}
83 + {% if not loop.last %}, {% endif %}
84 + {%- endfor %}
85 + {{- ']<|eot_id|>' -}}
86 + {%- elif message.role == "tool" or message.role == "ipython" %}
87 + {{- "<|start_header_id|>ipython<|end_header_id|>\n\n" }}
88 + {%- if message.content is mapping %}
89 + {{- message.content | tojson }}
90 + {%- else %}
91 + {{- { "output": message.content } | tojson }}
92 + {%- endif %}
93 + {{- "<|eot_id|>" }}
94 + {%- endif %}
95 +{%- endfor %}
96 +{%- if add_generation_prompt %}
97 + {{- '<|start_header_id|>assistant<|end_header_id|>\n\n' }}
98 +{%- endif %}
rust/src/chat/tests/templates/vllm_examples/tool_chat_template_llama4_json.jinja new
+116
@@ -0,0 +1,116 @@
1 +{%- macro is_array_of_type_objects(var) -%}
2 + {%- if var is iterable and var is not string -%}
3 + {%- set valid = true -%}
4 + {%- for item in var -%}
5 + {%- if 'type' not in item -%}
6 + {%- set valid = false -%}
7 + {%- break -%}
8 + {%- endif -%}
9 + {%- endfor -%}
10 + {{ valid }}
11 + {%- else -%}
12 + {{ false }}
13 + {%- endif -%}
14 +{%- endmacro %}
15 +
16 +{%- macro render_message(message) %}
17 + {%- if message['content'] is string %}
18 + {{- message['content']|trim }}
19 + {%- elif is_array_of_type_objects(data) == 'True' %}
20 + {%- for content in message['content'] %}
21 + {%- if content['type'] == 'image' %}
22 + {{- '<|image|>' }}
23 + {%- elif content['type'] == 'text' %}
24 + {{- content['text']|trim }}
25 + {%- endif %}
26 + {%- endfor %}
27 + {%- else %}
28 + {{- message['content']|tojson }}
29 + {%- endif %}
30 +{%- endmacro %}
31 +
32 +{{- bos_token }}
33 +{%- if custom_tools is defined %}
34 + {%- set tools = custom_tools %}
35 +{%- endif %}
36 +{%- if not tools_in_user_message is defined %}
37 + {%- set tools_in_user_message = true %}
38 +{%- endif %}
39 +{%- if not tools is defined %}
40 + {%- set tools = none %}
41 +{%- endif %}
42 +
43 +{#- This block extracts the system message, so we can slot it into the right place. #}
44 +{%- if messages[0]['role'] == 'system' %}
45 + {%- set system_message = messages[0] %}
46 + {%- set messages = messages[1:] %}
47 +{%- else %}
48 + {%- set system_message = ({ "content": "You are a helpful assistant with tool calling "
49 + "capabilities. Only reply with a tool call if the function exists in the "
50 + "library provided by the user. If it doesn't exist, just reply directly in "
51 + "natural language. When you receive a tool call response, use the output to "
52 + "format an answer to the original user question."}) %}
53 +{%- endif %}
54 +
55 +{%- set tool_lib_preamble = 'Tools: You have access to the following tools. You might need to use one '
56 + 'or more function/tool calls to fulfill the task. \n'
57 + 'If none are needed, then proceed to the response.\n\n'
58 + 'Tool Call Syntax: You can call tools using the following syntax:\n'
59 + '{"name": function name, "parameters": dictionary of argument name and its value}.\n'
60 + 'Separate multiple function calls by "; ". Do not use variables.\n'
61 + 'Do not include anything else when calling the tools with the syntax above.\n\n'
62 + 'Here is a list of functions in JSON format that you can invoke.\n' %}
63 +
64 +{{- "<|header_start|>system<|header_end|>\n\n" }}
65 +{%- if tools is not none and not tools_in_user_message %}
66 + {{- tool_lib_preamble }}
67 + {%- for t in tools %}
68 + {{- t | tojson(indent=4) }}
69 + {{- "\n\n" }}
70 + {%- endfor %}
71 +{%- endif %}
72 +{{- render_message(system_message) }}
73 +{{ "<|eot|>\n" }}
74 +
75 +{#- Custom tools are passed in a user message with some extra guidance #}
76 +{%- if tools_in_user_message and not tools is none %}
77 + {#- Extract the first user message so we can plug it in here #}
78 + {%- if messages | length != 0 %}
79 + {%- set first_user_message = messages[0] %}
80 + {%- set messages = messages[1:] %}
81 + {%- else %}
82 + {{- raise_exception("Cannot put tools in the first user message when there's no first user message!") }}
83 + {%- endif %}
84 + {{- '<|header_start|>user<|header_end|>\n\n' }}
85 + {{- tool_lib_preamble }}
86 + {%- for t in tools %}
87 + {{- t | tojson(indent=4) }}
88 + {{- "\n\n" }}
89 + {%- endfor %}
90 + {{- render_message(first_user_message) + "\n<|eot|>"}}
91 +{%- endif %}
92 +
93 +{%- for message in messages %}
94 + {%- if not (message.role == 'ipython' or message.role == 'tool' or 'tool_calls' in message) %}
95 + {{- '<|header_start|>' + message['role'] + '<|header_end|>\n\n' }}
96 + {{- render_message(message) }}
97 + {{- "\n<|eot|>" }}
98 + {%- elif 'tool_calls' in message and message.tool_calls|length > 0 %}
99 + {{- '\n<|header_start|>assistant<|header_end|>\n\n' -}}
100 + {{- render_message(message) }}
101 + {%- for tool_call in message.tool_calls %}
102 + {{- '{"name": "' + tool_call.function.name + '", ' }}
103 + {{- '"parameters": ' }}
104 + {{- tool_call.function.arguments | tojson }}
105 + {{- "}" }}
106 + {%- endfor %}
107 + {{- "\n<|eot|>" }}
108 + {%- elif message.role == "tool" or message.role == "ipython" %}
109 + {{- "\n<|header_start|>ipython<|header_end|>\n\n" }}
110 + {{- render_message(message) }}
111 + {{- "\n<|eom|>" }}
112 + {%- endif %}
113 +{%- endfor %}
114 +{%- if add_generation_prompt %}
115 + {{- '\n<|header_start|>assistant<|header_end|>\n\n' }}
116 +{%- endif %}
rust/src/chat/tests/templates/vllm_examples/tool_chat_template_llama4_pythonic.jinja new
+111
@@ -0,0 +1,111 @@
1 +{{- bos_token }}
2 +{%- if custom_tools is defined and custom_tools%}
3 + {%- set tools = custom_tools %}
4 +{%- endif %}
5 +{%- if tools is defined and tools %}
6 + {%- set tool_definition = tool_definition ~ (tools | tojson(indent=4)) %}
7 +{%- else %}
8 + {%- set tools = none %}
9 +{%- endif %}
10 +
11 +
12 +{#- This block extracts the system message, so we can slot it into the right place. #}
13 +{%- if messages[0]['role'] == 'system' %}
14 + {%- set user_provided_system_message = true %}
15 + {%- if messages[0]['content'] is string %}
16 + {%- set system_message = messages[0]['content']|trim %}
17 + {%- else %}
18 + {%- set system_message = messages[0]['content'][0]['text']|trim %}
19 + {%- endif %}
20 + {%- set messages = messages[1:] %}
21 +{%- else %}
22 + {%- if tools is not none %}
23 + {#- Since not system_message was provided by user, if tool is provided, system_message is now default tool system message #}
24 + {#- This system message is from llama website:https://www.llama.com/docs/model-cards-and-prompt-formats/llama4/ #}
25 + {%- set system_message = "You are a helpful assistant and an expert in function composition. You can answer general questions using your internal knowledge OR invoke functions when necessary. Follow these strict guidelines:\n\n1. FUNCTION CALLS:\n- ONLY use functions that are EXPLICITLY listed in the function list below\n- If NO functions are listed (empty function list []), respond ONLY with internal knowledge or \"I don't have access to [Unavailable service] information\"\n- If a function is not in the list, respond ONLY with internal knowledge or \"I don't have access to [Unavailable service] information\"\n- If ALL required parameters are present AND the query EXACTLY matches a listed function's purpose: output ONLY the function call(s)\n- Use exact format: [func_name1(param1=value1, param2=value2), func_name2(...)]\nExamples:\nCORRECT: [get_weather(location=\"Vancouver\"), calculate_route(start=\"Boston\", end=\"New York\")] <- Only if get_weather and calculate_route are in function list\nINCORRECT: get_weather(location=\"New York\")\nINCORRECT: Let me check the weather: [get_weather(location=\"New York\")]\nINCORRECT: [get_events(location=\"Singapore\")] <- If function not in list\n\n2. RESPONSE RULES:\n- For pure function requests matching a listed function: ONLY output the function call(s)\n- For knowledge questions: ONLY output text\n- For missing parameters: ONLY request the specific missing parameters\n- For unavailable services (not in function list): output ONLY with internal knowledge or \"I don't have access to [Unavailable service] information\". Do NOT execute a function call.\n- If the query asks for information beyond what a listed function provides: output ONLY with internal knowledge about your limitations\n- NEVER combine text and function calls in the same response\n- NEVER suggest alternative functions when the requested service is unavailable\n- NEVER create or invent new functions not listed below\n\n3. STRICT BOUNDARIES:\n- ONLY use functions from the list below - no exceptions\n- NEVER use a function as an alternative to unavailable information\n- NEVER call functions not present in the function list\n- NEVER add explanatory text to function calls\n- NEVER respond with empty brackets\n- Use proper Python/JSON syntax for function calls\n- Check the function list carefully before responding\n\n4. TOOL RESPONSE HANDLING:\n- When receiving tool responses: provide concise, natural language responses\n- Don't repeat tool response verbatim\n- Don't add supplementary information\n\nHere is a list of functions in JSON format that you can invoke:\n" %}
26 + {%- else %}
27 + {%- set system_message = "" %}
28 + {%- endif %}
29 +{%- endif %}
30 +{#- Now writing the system message: use the user provided system message if user_provided_system_message, else default tool system message if tools presented #}
31 +{%- if system_message %}
32 + {#- always use user provided system message to override default tool system message #}
33 + {{- "<|header_start|>system<|header_end|>\n\n" }}
34 + {{- system_message }}
35 + {%- if user_provided_system_message and tools %}
36 + {{- "\nHere is a list of functions in JSON format that you can invoke. Use exact format: [func_name1(param1=value1, param2=value2), func_name2(...)]\n" }}
37 + {{- tool_definition -}}
38 + {%- elif tool_definition %}
39 + {{- tool_definition -}}
40 + {%- endif %}
41 + {{- "<|eot|>" }}
42 +{%- endif %}
43 +
44 +{#- Now deal with all other messages #}
45 +{%- for message in messages %}
46 + {#- Base case: messages that are not from tool role and has empty tool_call list #}
47 + {%- if not (message.role == 'ipython' or message.role == 'tool' or ('tool_calls' in message and message.tool_calls|length != 0 )) %}
48 + {{- '<|header_start|>' + message['role'] + '<|header_end|>\n\n' }}
49 + {%- if message['content'] is string %}
50 + {{- message['content'] }}
51 + {%- else %}
52 + {%- for content in message['content'] %}
53 + {%- if content['type'] == 'image' %}
54 + {{- '<|image|>' }}
55 + {%- elif content['type'] == 'text' %}
56 + {{- content['text'] | trim }}
57 + {%- endif %}
58 + {%- endfor %}
59 + {%- endif %}
60 + {{- "<|eot|>" }}
61 + {#- Tool case: messages has non-empty tool_call list, must from assistant #}
62 + {%- elif 'tool_calls' in message %}
63 + {#- assume tool_calls are always coming from assistant #}
64 + {%- if message.role == 'assistant' %}
65 + {{- '<|header_start|>assistant<|header_end|>\n\n' -}}
66 + {%- if message['content'] is string %}
67 + {{- message['content'] }}
68 + {%- else %}
69 + {%- for content in message['content'] %}
70 + {%- if content['type'] == 'image' %}
71 + {{- '<|image|>' }}
72 + {%- elif content['type'] == 'text' %}
73 + {{- content['text'] }}
74 + {%- endif %}
75 + {%- endfor %}
76 + {%- endif %}
77 + {{- "[" }}
78 + {%- for tool_call in message.tool_calls %}
79 + {%- if tool_call.function is defined %}
80 + {%- set tool_call = tool_call.function %}
81 + {%- endif %}
82 + {{- tool_call.name + '(' -}}
83 + {%- for param in tool_call.arguments %}
84 + {{- param + '="' -}}
85 + {{- "%s" | format(tool_call.arguments[param]) -}}
86 + {{- '"' -}}
87 + {% if not loop.last %}, {% endif %}
88 + {%- endfor %}
89 + {{- ')' -}}
90 + {% if not loop.last %}, {% endif %}
91 + {%- endfor %}
92 + {{- "]<|eot|>" }}
93 +{%- endif %}
94 +{#- Tool_response case: messages are from tool_response #}
95 + {%- elif message.role == "tool" or message.role == "ipython" %}
96 + {{- "<|header_start|>ipython<|header_end|>\n\n" }}
97 + {%- if message.content is string %}
98 + {{- message.content | tojson }}
99 + {%- else %}
100 + {%- for content in message['content'] %}
101 + {%- if content['type'] == 'text' %}
102 + {{- content['text'] | tojson }}
103 + {%- endif %}
104 + {%- endfor %}
105 + {%- endif %}
106 + {{- "<|eot|>" }}
107 + {%- endif %}
108 +{%- endfor %}
109 +{%- if add_generation_prompt %}
110 + {{- '<|header_start|>assistant<|header_end|>\n\n' }}
111 +{%- endif %}
rust/src/chat/tests/templates/vllm_examples/tool_chat_template_minimax_m1.jinja new
+91
@@ -0,0 +1,91 @@
1 +{{ '<begin_of_document>' -}}
2 +{%- if custom_tools is defined %}
3 + {%- set tools = custom_tools %}
4 +{%- endif %}
5 +{%- if not tools is defined %}
6 + {%- set tools = none %}
7 +{%- endif %}
8 +
9 +{#- Extract system message #}
10 +{% set ns = namespace(system_prompt='') -%}
11 +{%- if messages[0]['role'] == 'system' %}
12 + {%- if messages[0]['content'] is string %}
13 + {%- set ns.system_prompt = messages[0]['content']|trim %}
14 + {%- else %}
15 + {%- set ns.system_prompt = messages[0]['content'][0]['text']|trim %}
16 + {%- endif %}
17 + {%- set messages = messages[1:] %}
18 +{%- else %}
19 + {%- if tools is not none %}
20 + {%- set ns.system_prompt = "You are a helpful assistant created by Minimax based on MiniMax-M1 model." %}
21 + {%- else %}
22 + {%- set ns.system_prompt = "You are a helpful assistant created by Minimax based on MiniMax-M1 model." %}
23 + {%- endif %}
24 +{%- endif %}
25 +
26 +{#- System message #}
27 +{%- if ns.system_prompt != '' %}
28 +{{ '<beginning_of_sentence>system ai_setting=assistant\n' + ns.system_prompt + '<end_of_sentence>\n' -}}
29 +{%- endif %}
30 +
31 +{#- Tools configuration #}
32 +{%- if tools is not none %}
33 +{{ '<beginning_of_sentence>system tool_setting=tools\nYou are provided with these tools:\n<tools>\n' -}}
34 +{%- for tool in tools %}
35 +{{ tool | tojson ~ '\n' -}}
36 +{%- endfor %}
37 +{{ '</tools>\n\nIf you need to call tools, please respond with <tool_calls></tool_calls> XML tags, and provide tool-name and json-object of arguments, following the format below:\n<tool_calls>\n{"name": <tool-name>, "arguments": <args-json-object>}\n...\n</tool_calls><end_of_sentence>\n' -}}
38 +{%- endif %}
39 +
40 +{#- Process messages #}
41 +{%- for message in messages %}
42 + {%- if not (message.role == 'ipython' or message.role == 'tool' or 'tool_calls' in message) %}
43 + {%- if message['role'] == 'user' %}
44 +{{ '<beginning_of_sentence>user name=user\n' -}}
45 +{%- if message['content'] is string %}
46 +{{ message['content']|trim -}}
47 +{%- else %}
48 +{%- for content in message['content'] %}
49 +{%- if content['type'] == 'text' %}
50 +{{ content['text']|trim -}}
51 +{%- endif %}
52 +{%- endfor %}
53 +{%- endif %}
54 +{{ '<end_of_sentence>\n' -}}
55 + {%- elif message['role'] == 'assistant' %}
56 +{{ '<beginning_of_sentence>ai name=assistant\n' -}}
57 +{%- if message['content'] is string %}
58 +{{ message['content']|trim -}}
59 +{%- else %}
60 +{%- for content in message['content'] | selectattr('type', 'equalto', 'text') %}
61 +{{ content['text']|trim -}}
62 +{%- endfor %}
63 +{%- endif %}
64 +{{ '<end_of_sentence>\n' -}}
65 + {%- endif %}
66 + {%- elif 'tool_calls' in message %}
67 +{{ '<beginning_of_sentence>ai name=assistant\n<tool_calls>\n' -}}
68 +{%- for tool_call in message.tool_calls %}
69 +{{ '{"name": "' + tool_call.function.name + '", "arguments": ' + tool_call.function.arguments | tojson + '}\n' -}}
70 +{%- endfor %}
71 +{{ '</tool_calls><end_of_sentence>\n' -}}
72 + {%- elif message.role == "tool" or message.role == "ipython" %}
73 +{{ '<beginning_of_sentence>tool name=tools\n' -}}
74 +{%- if message.content is string %}
75 +{{ 'tool result: ' + message.content + '\n\n' -}}
76 +{%- else %}
77 +{%- for content in message['content'] %}
78 +{%- if content['type'] == 'text' %}
79 +{{ 'tool result: ' + content['text'] + '\n\n' -}}
80 +{%- elif content.get('name') %}
81 +{{ 'tool name: ' + content['name'] + '\ntool result: ' + content['text'] + '\n\n' -}}
82 +{%- endif %}
83 +{%- endfor %}
84 +{%- endif %}
85 +{{ '<end_of_sentence>\n' -}}
86 + {%- endif %}
87 +{%- endfor %}
88 +
89 +{%- if add_generation_prompt %}
90 +{{ '<beginning_of_sentence>ai name=assistant\n' -}}
91 +{%- endif %}
\ No newline at end of file
rust/src/chat/tests/templates/vllm_examples/tool_chat_template_mistral.jinja new
+86
@@ -0,0 +1,86 @@
1 +{%- if messages[0]["role"] == "system" %}
2 + {%- set system_message = messages[0]["content"] %}
3 + {%- set loop_messages = messages[1:] %}
4 +{%- else %}
5 + {%- set loop_messages = messages %}
6 +{%- endif %}
7 +{%- if not tools is defined %}
8 + {%- set tools = none %}
9 +{%- endif %}
10 +{%- set user_messages = loop_messages | selectattr("role", "equalto", "user") | list %}
11 +
12 +{%- for message in loop_messages | rejectattr("role", "equalto", "tool") | rejectattr("role", "equalto", "tool_results") | selectattr("tool_calls", "undefined") %}
13 + {%- if (message["role"] == "user") != (loop.index0 % 2 == 0) %}
14 + {{- raise_exception("After the optional system message, conversation roles must alternate user/assistant/user/assistant/...") }}
15 + {%- endif %}
16 +{%- endfor %}
17 +
18 +{{- bos_token }}
19 +{%- for message in loop_messages %}
20 + {%- if message["role"] == "user" %}
21 + {%- if tools is not none and (message == user_messages[-1]) %}
22 + {{- "[AVAILABLE_TOOLS] [" }}
23 + {%- for tool in tools %}
24 + {%- set tool = tool.function %}
25 + {{- '{"type": "function", "function": {' }}
26 + {%- for key, val in tool.items() if key != "return" %}
27 + {%- if val is string %}
28 + {{- '"' + key + '": "' + val + '"' }}
29 + {%- else %}
30 + {{- '"' + key + '": ' + val|tojson }}
31 + {%- endif %}
32 + {%- if not loop.last %}
33 + {{- ", " }}
34 + {%- endif %}
35 + {%- endfor %}
36 + {{- "}}" }}
37 + {%- if not loop.last %}
38 + {{- ", " }}
39 + {%- else %}
40 + {{- "]" }}
41 + {%- endif %}
42 + {%- endfor %}
43 + {{- "[/AVAILABLE_TOOLS]" }}
44 + {%- endif %}
45 + {%- if loop.last and system_message is defined %}
46 + {{- "[INST] " + system_message + "\n\n" + message["content"] + "[/INST]" }}
47 + {%- else %}
48 + {{- "[INST] " + message["content"] + "[/INST]" }}
49 + {%- endif %}
50 + {%- elif message["role"] == "tool_calls" or message.tool_calls is defined %}
51 + {%- if message.tool_calls is defined %}
52 + {%- set tool_calls = message.tool_calls %}
53 + {%- else %}
54 + {%- set tool_calls = message.content %}
55 + {%- endif %}
56 + {{- "[TOOL_CALLS] [" }}
57 + {%- for tool_call in tool_calls %}
58 + {%- set out = tool_call.function|tojson %}
59 + {{- out[:-1] }}
60 + {%- if not tool_call.id is defined or tool_call.id|length < 9 %}
61 + {{- raise_exception("Tool call IDs should be alphanumeric strings with length >= 9! (1)" + tool_call.id) }}
62 + {%- endif %}
63 + {{- ', "id": "' + tool_call.id[-9:] + '"}' }}
64 + {%- if not loop.last %}
65 + {{- ", " }}
66 + {%- else %}
67 + {{- "]" + eos_token }}
68 + {%- endif %}
69 + {%- endfor %}
70 + {%- elif message["role"] == "assistant" %}
71 + {{- " " + message["content"] + eos_token }}
72 + {%- elif message["role"] == "tool_results" or message["role"] == "tool" %}
73 + {%- if message.content is defined and message.content.content is defined %}
74 + {%- set content = message.content.content %}
75 + {%- else %}
76 + {%- set content = message.content %}
77 + {%- endif %}
78 + {{- '[TOOL_RESULTS] {"content": ' + content|string + ", " }}
79 + {%- if not message.tool_call_id is defined or message.tool_call_id|length < 9 %}
80 + {{- raise_exception("Tool call IDs should be alphanumeric strings with length >= 9! (2)" + message.tool_call_id) }}
81 + {%- endif %}
82 + {{- '"call_id": "' + message.tool_call_id[-9:] + '"}[/TOOL_RESULTS]' }}
83 + {%- else %}
84 + {{- raise_exception("Only user and assistant roles are supported, with the exception of an initial optional system message!") }}
85 + {%- endif %}
86 +{%- endfor %}
rust/src/chat/tests/templates/vllm_examples/tool_chat_template_mistral3.jinja new
+126
@@ -0,0 +1,126 @@
1 +{%- set today = strftime_now("%Y-%m-%d") %}
2 +{%- set default_system_message = "You are Mistral Small 3, a Large Language Model (LLM) created by Mistral AI, a French startup headquartered in Paris.\nYour knowledge base was last updated on 2023-10-01. The current date is " + today + ".\n\nWhen you're not sure about some information, you say that you don't have the information and don't make up anything.\nIf the user's question is not clear, ambiguous, or does not provide enough context for you to accurately answer the question, you do not try to answer it right away and you rather ask the user to clarify their request (e.g. \"What are some good restaurants around me?\" => \"Where are you?\" or \"When is the next flight to Tokyo\" => \"Where do you travel from?\")" %}
3 +
4 +{{- bos_token }}
5 +
6 +{%- if messages[0]['role'] == 'system' %}
7 + {%- if messages[0]['content'] is string %}
8 + {%- set system_message = messages[0]['content'] %}
9 + {%- set loop_messages = messages[1:] %}
10 + {%- else %}
11 + {%- set system_message = messages[0]['content'][0]['text'] %}
12 + {%- set loop_messages = messages[1:] %}
13 + {%- endif %}
14 +{%- else %}
15 + {%- set system_message = default_system_message %}
16 + {%- set loop_messages = messages %}
17 +{%- endif %}
18 +{%- if not tools is defined %}
19 + {%- set tools = none %}
20 +{%- elif tools is not none %}
21 + {%- set parallel_tool_prompt = "You are a helpful assistant that can call tools. If you call one or more tools, format them in a single JSON array or objects, where each object is a tool call, not as separate objects outside of an array or multiple arrays. Use the format [{\"name\": tool call name, \"arguments\": tool call arguments}, additional tool calls] if you call more than one tool. If you call tools, do not attempt to interpret them or otherwise provide a response until you receive a tool call result that you can interpret for the user." %}
22 + {%- if system_message is defined %}
23 + {%- set system_message = parallel_tool_prompt + "\n\n" + system_message %}
24 + {%- else %}
25 + {%- set system_message = parallel_tool_prompt %}
26 + {%- endif %}
27 +{%- endif %}
28 +{{- '[SYSTEM_PROMPT]' + system_message + '[/SYSTEM_PROMPT]' }}
29 +
30 +{%- set user_messages = loop_messages | selectattr("role", "equalto", "user") | list %}
31 +
32 +{%- set filtered_messages = [] %}
33 +{%- for message in loop_messages %}
34 + {%- if message["role"] not in ["tool", "tool_results"] and not message.get("tool_calls") %}
35 + {%- set filtered_messages = filtered_messages + [message] %}
36 + {%- endif %}
37 +{%- endfor %}
38 +
39 +{%- for message in filtered_messages %}
40 + {%- if (message["role"] == "user") != (loop.index0 % 2 == 0) %}
41 + {{- raise_exception("After the optional system message, conversation roles must alternate user/assistant/user/assistant/...") }}
42 + {%- endif %}
43 +{%- endfor %}
44 +
45 +{%- for message in loop_messages %}
46 + {%- if message["role"] == "user" %}
47 + {%- if tools is not none and (message == user_messages[-1]) %}
48 + {{- "[AVAILABLE_TOOLS] [" }}
49 + {%- for tool in tools %}
50 + {%- set tool = tool.function %}
51 + {{- '{"type": "function", "function": {' }}
52 + {%- for key, val in tool.items() if key != "return" %}
53 + {%- if val is string %}
54 + {{- '"' + key + '": "' + val + '"' }}
55 + {%- else %}
56 + {{- '"' + key + '": ' + val|tojson }}
57 + {%- endif %}
58 + {%- if not loop.last %}
59 + {{- ", " }}
60 + {%- endif %}
61 + {%- endfor %}
62 + {{- "}}" }}
63 + {%- if not loop.last %}
64 + {{- ", " }}
65 + {%- else %}
66 + {{- "]" }}
67 + {%- endif %}
68 + {%- endfor %}
69 + {{- "[/AVAILABLE_TOOLS]" }}
70 + {%- endif %}
71 + {%- if message['content'] is string %}
72 + {{- '[INST]' + message['content'] + '[/INST]' }}
73 + {%- else %}
74 + {{- '[INST]' }}
75 + {%- for block in message['content'] %}
76 + {%- if block['type'] == 'text' %}
77 + {{- block['text'] }}
78 + {%- elif block['type'] == 'image' or block['type'] == 'image_url' %}
79 + {{- '[IMG]' }}
80 + {%- else %}
81 + {{- raise_exception('Only text and image blocks are supported in message content!') }}
82 + {%- endif %}
83 + {%- endfor %}
84 + {{- '[/INST]' }}
85 + {%- endif %}
86 + {%- elif message["role"] == "tool_calls" or message.tool_calls is defined %}
87 + {%- if message.tool_calls is defined %}
88 + {%- set tool_calls = message.tool_calls %}
89 + {%- else %}
90 + {%- set tool_calls = message.content %}
91 + {%- endif %}
92 + {{- "[TOOL_CALLS] [" }}
93 + {%- for tool_call in tool_calls %}
94 + {%- set out = tool_call.function|tojson %}
95 + {{- out[:-1] }}
96 + {%- if not tool_call.id is defined or tool_call.id|length < 9 %}
97 + {{- raise_exception("Tool call IDs should be alphanumeric strings with length >= 9! (1)" + tool_call.id) }}
98 + {%- endif %}
99 + {{- ', "id": "' + tool_call.id[-9:] + '"}' }}
100 + {%- if not loop.last %}
101 + {{- ", " }}
102 + {%- else %}
103 + {{- "]" + eos_token }}
104 + {%- endif %}
105 + {%- endfor %}
106 + {%- elif message['role'] == 'assistant' %}
107 + {%- if message['content'] is string %}
108 + {{- message['content'] + eos_token }}
109 + {%- else %}
110 + {{- message['content'][0]['text'] + eos_token }}
111 + {%- endif %}
112 + {%- elif message["role"] == "tool_results" or message["role"] == "tool" %}
113 + {%- if message.content is defined and message.content.content is defined %}
114 + {%- set content = message.content.content %}
115 + {%- else %}
116 + {%- set content = message.content %}
117 + {%- endif %}
118 + {{- '[TOOL_RESULTS] {"content": ' + content|string + ", " }}
119 + {%- if not message.tool_call_id is defined or message.tool_call_id|length < 9 %}
120 + {{- raise_exception("Tool call IDs should be alphanumeric strings with length >= 9! (2)" + message.tool_call_id) }}
121 + {%- endif %}
122 + {{- '"call_id": "' + message.tool_call_id[-9:] + '"}[/TOOL_RESULTS]' }}
123 + {%- else %}
124 + {{- raise_exception("Only user and assistant roles are supported, with the exception of an initial optional system message!") }}
125 + {%- endif %}
126 +{%- endfor %}
rust/src/chat/tests/templates/vllm_examples/tool_chat_template_mistral_parallel.jinja new
+93
@@ -0,0 +1,93 @@
1 +{%- if messages[0]["role"] == "system" %}
2 + {%- set system_message = messages[0]["content"] %}
3 + {%- set loop_messages = messages[1:] %}
4 +{%- else %}
5 + {%- set loop_messages = messages %}
6 +{%- endif %}
7 +{%- if not tools is defined %}
8 + {%- set tools = none %}
9 +{%- elif tools is not none %}
10 + {%- set parallel_tool_prompt = "You are a helpful assistant that can call tools. If you call one or more tools, format them in a single JSON array or objects, where each object is a tool call, not as separate objects outside of an array or multiple arrays. Use the format [{\"name\": tool call name, \"arguments\": tool call arguments}, additional tool calls] if you call more than one tool. If you call tools, do not attempt to interpret them or otherwise provide a response until you receive a tool call result that you can interpret for the user." %}
11 + {%- if system_message is defined %}
12 + {%- set system_message = parallel_tool_prompt + "\n\n" + system_message %}
13 + {%- else %}
14 + {%- set system_message = parallel_tool_prompt %}
15 + {%- endif %}
16 +{%- endif %}
17 +{%- set user_messages = loop_messages | selectattr("role", "equalto", "user") | list %}
18 +
19 +{%- for message in loop_messages | rejectattr("role", "equalto", "tool") | rejectattr("role", "equalto", "tool_results") | selectattr("tool_calls", "undefined") %}
20 + {%- if (message["role"] == "user") != (loop.index0 % 2 == 0) %}
21 + {{- raise_exception("After the optional system message, conversation roles must alternate user/assistant/user/assistant/...") }}
22 + {%- endif %}
23 +{%- endfor %}
24 +
25 +{{- bos_token }}
26 +{%- for message in loop_messages %}
27 + {%- if message["role"] == "user" %}
28 + {%- if tools is not none and (message == user_messages[-1]) %}
29 + {{- "[AVAILABLE_TOOLS] [" }}
30 + {%- for tool in tools %}
31 + {%- set tool = tool.function %}
32 + {{- '{"type": "function", "function": {' }}
33 + {%- for key, val in tool.items() if key != "return" %}
34 + {%- if val is string %}
35 + {{- '"' + key + '": "' + val + '"' }}
36 + {%- else %}
37 + {{- '"' + key + '": ' + val|tojson }}
38 + {%- endif %}
39 + {%- if not loop.last %}
40 + {{- ", " }}
41 + {%- endif %}
42 + {%- endfor %}
43 + {{- "}}" }}
44 + {%- if not loop.last %}
45 + {{- ", " }}
46 + {%- else %}
47 + {{- "]" }}
48 + {%- endif %}
49 + {%- endfor %}
50 + {{- "[/AVAILABLE_TOOLS]" }}
51 + {%- endif %}
52 + {%- if loop.last and system_message is defined %}
53 + {{- "[INST] " + system_message + "\n\n" + message["content"] + "[/INST]" }}
54 + {%- else %}
55 + {{- "[INST] " + message["content"] + "[/INST]" }}
56 + {%- endif %}
57 + {%- elif message["role"] == "tool_calls" or message.tool_calls is defined %}
58 + {%- if message.tool_calls is defined %}
59 + {%- set tool_calls = message.tool_calls %}
60 + {%- else %}
61 + {%- set tool_calls = message.content %}
62 + {%- endif %}
63 + {{- "[TOOL_CALLS] [" }}
64 + {%- for tool_call in tool_calls %}
65 + {%- set out = tool_call.function|tojson %}
66 + {{- out[:-1] }}
67 + {%- if not tool_call.id is defined or tool_call.id|length < 9 %}
68 + {{- raise_exception("Tool call IDs should be alphanumeric strings with length >= 9! (1)" + tool_call.id) }}
69 + {%- endif %}
70 + {{- ', "id": "' + tool_call.id[-9:] + '"}' }}
71 + {%- if not loop.last %}
72 + {{- ", " }}
73 + {%- else %}
74 + {{- "]" + eos_token }}
75 + {%- endif %}
76 + {%- endfor %}
77 + {%- elif message["role"] == "assistant" %}
78 + {{- " " + message["content"] + eos_token }}
79 + {%- elif message["role"] == "tool_results" or message["role"] == "tool" %}
80 + {%- if message.content is defined and message.content.content is defined %}
81 + {%- set content = message.content.content %}
82 + {%- else %}
83 + {%- set content = message.content %}
84 + {%- endif %}
85 + {{- '[TOOL_RESULTS] {"content": ' + content|string + ", " }}
86 + {%- if not message.tool_call_id is defined or message.tool_call_id|length < 9 %}
87 + {{- raise_exception("Tool call IDs should be alphanumeric strings with length >= 9! (2)" + message.tool_call_id) }}
88 + {%- endif %}
89 + {{- '"call_id": "' + message.tool_call_id[-9:] + '"}[/TOOL_RESULTS]' }}
90 + {%- else %}
91 + {{- raise_exception("Only user and assistant roles are supported, with the exception of an initial optional system message!") }}
92 + {%- endif %}
93 +{%- endfor %}
rust/src/chat/tests/templates/vllm_examples/tool_chat_template_phi4_mini.jinja new
+62
@@ -0,0 +1,62 @@
1 +{%- if messages and messages[0]['role'] == 'system' %}
2 + {%- set system_message = messages[0]['content']|trim %}
3 + {%- set messages = messages[1:] %}
4 +{%- else %}
5 + {%- set system_message = "You are a helpful assistant." %}
6 +{%- endif %}
7 +
8 +{%- if messages %}
9 +<|system|>
10 +{{ system_message }}
11 +{%- if tools %}
12 +In addition to plain text responses, you can choose to call one or more of the provided functions.
13 +
14 +Use the following rule to decide when to call a function:
15 + * if the response can be generated from your internal knowledge (e.g., as in the case of queries like "What is the capital of Poland?"), do so
16 + * if you need external information that can be obtained by calling one or more of the provided functions, generate a function calls
17 +
18 +If you decide to call functions:
19 + * prefix function calls with functools marker (no closing marker required)
20 + * all function calls should be generated in a single JSON list formatted as functools[{"name": [function name], "arguments": [function arguments as JSON]}, ...]
21 + * follow the provided JSON schema. Do not hallucinate arguments or values. Do to blindly copy values from the provided samples
22 + * respect the argument type formatting. E.g., if the type is number and format is float, write value 7 as 7.0
23 + * make sure you pick the right functions that match the user intent
24 +
25 +
26 + {%- for t in tools %}
27 + {{- t | tojson(indent=4) }}
28 + {{- "\n\n" }}
29 + {%- endfor %}
30 +{%- endif %}<|end|>
31 +
32 + {%- for message in messages %}
33 + {%- if message.role != "system" %}
34 +<|{{ message.role }}|>
35 + {%- if message.content and message.role == "tools" %}
36 +{"result": {{ message.content }}}
37 + {%- elif message.content %}
38 +{{ message.content }}
39 + {%- elif message.tool_calls %}
40 + {%- for call in message.tool_calls %}
41 +{"name": "{{ call.function.name }}", "arguments": {{ call.function.arguments }}}
42 + {%- if not loop.last %},{% endif %}
43 + {%- endfor %}
44 + {%- endif %}<|end|>
45 + {%- endif %}
46 + {%- endfor %}<|assistant|>
47 +
48 +{%- else %}
49 + {%- if system_message %}
50 +<|system|>
51 +
52 +{{ system_message }}<|end|>
53 + {%- endif %}
54 + {%- if prompt %}
55 +<|user|>
56 +
57 +{{ prompt }}<|end|>
58 + {%- endif %}<|assistant|>
59 +
60 +{%- endif %}
61 +{{ response }}
62 +{%- if response %}<|user|>{% endif %}
\ No newline at end of file
rust/src/chat/tests/templates/vllm_examples/tool_chat_template_qwen3coder.jinja new
+117
@@ -0,0 +1,117 @@
1 +{% macro render_extra_keys(json_dict, handled_keys) %}
2 + {%- if json_dict is mapping %}
3 + {%- for json_key in json_dict if json_key not in handled_keys %}
4 + {%- if json_dict[json_key] is mapping or (json_dict[json_key] is sequence and json_dict[json_key] is not string) %}
5 + {{- '\n<' ~ json_key ~ '>' ~ (json_dict[json_key] | tojson | safe) ~ '</' ~ json_key ~ '>' }}
6 + {%- else %}
7 + {{-'\n<' ~ json_key ~ '>' ~ (json_dict[json_key] | string) ~ '</' ~ json_key ~ '>' }}
8 + {%- endif %}
9 + {%- endfor %}
10 + {%- endif %}
11 +{% endmacro %}
12 +
13 +{%- if messages[0]["role"] == "system" %}
14 + {%- set system_message = messages[0]["content"] %}
15 + {%- set loop_messages = messages[1:] %}
16 +{%- else %}
17 + {%- set loop_messages = messages %}
18 +{%- endif %}
19 +
20 +{%- if not tools is defined %}
21 + {%- set tools = [] %}
22 +{%- endif %}
23 +
24 +{%- if system_message is defined %}
25 + {{- "<|im_start|>system\n" + system_message }}
26 +{%- else %}
27 + {%- if tools is iterable and tools | length > 0 %}
28 + {{- "<|im_start|>system\nYou are Qwen, a helpful AI assistant that can interact with a computer to solve tasks." }}
29 + {%- endif %}
30 +{%- endif %}
31 +{%- if tools is iterable and tools | length > 0 %}
32 + {{- "\n\n# Tools\n\nYou have access to the following functions:\n\n" }}
33 + {{- "<tools>" }}
34 + {%- for tool in tools %}
35 + {%- if tool.function is defined %}
36 + {%- set tool = tool.function %}
37 + {%- endif %}
38 + {{- "\n<function>\n<name>" ~ tool.name ~ "</name>" }}
39 + {%- if tool.description is defined %}
40 + {{- '\n<description>' ~ (tool.description | trim) ~ '</description>' }}
41 + {%- endif %}
42 + {{- '\n<parameters>' }}
43 + {%- if tool.parameters is defined and tool.parameters is mapping and tool.parameters.properties is defined and tool.parameters.properties is mapping %}
44 + {%- for param_name, param_fields in tool.parameters.properties|items %}
45 + {{- '\n<parameter>' }}
46 + {{- '\n<name>' ~ param_name ~ '</name>' }}
47 + {%- if param_fields.type is defined %}
48 + {{- '\n<type>' ~ (param_fields.type | string) ~ '</type>' }}
49 + {%- endif %}
50 + {%- if param_fields.description is defined %}
51 + {{- '\n<description>' ~ (param_fields.description | trim) ~ '</description>' }}
52 + {%- endif %}
53 + {%- set handled_keys = ['name', 'type', 'description'] %}
54 + {{- render_extra_keys(param_fields, handled_keys) }}
55 + {{- '\n</parameter>' }}
56 + {%- endfor %}
57 + {%- endif %}
58 + {% set handled_keys = ['type', 'properties'] %}
59 + {{- render_extra_keys(tool.parameters, handled_keys) }}
60 + {{- '\n</parameters>' }}
61 + {%- set handled_keys = ['type', 'name', 'description', 'parameters'] %}
62 + {{- render_extra_keys(tool, handled_keys) }}
63 + {{- '\n</function>' }}
64 + {%- endfor %}
65 + {{- "\n</tools>" }}
66 + {{- '\n\nIf you choose to call a function ONLY reply in the following format with NO suffix:\n\n<tool_call>\n<function=example_function_name>\n<parameter=example_parameter_1>\nvalue_1\n</parameter>\n<parameter=example_parameter_2>\nThis is the value for the second parameter\nthat can span\nmultiple lines\n</parameter>\n</function>\n</tool_call>\n\n<IMPORTANT>\nReminder:\n- Function calls MUST follow the specified format: an inner <function=...></function> block must be nested within <tool_call></tool_call> XML tags\n- Required parameters MUST be specified\n- You may provide optional reasoning for your function call in natural language BEFORE the function call, but NOT after\n- If there is no function call available, answer the question like normal with your current knowledge and do not tell the user about function calls\n</IMPORTANT>' }}
67 +{%- endif %}
68 +{%- if system_message is defined %}
69 + {{- '<|im_end|>\n' }}
70 +{%- else %}
71 + {%- if tools is iterable and tools | length > 0 %}
72 + {{- '<|im_end|>\n' }}
73 + {%- endif %}
74 +{%- endif %}
75 +{%- for message in loop_messages %}
76 + {%- if message.role == "assistant" and message.tool_calls is defined and message.tool_calls is iterable and message.tool_calls | length > 0 %}
77 + {{- '<|im_start|>' + message.role }}
78 + {%- if message.content is defined and message.content is string and message.content | trim | length > 0 %}
79 + {{- '\n' + message.content | trim + '\n' }}
80 + {%- endif %}
81 + {%- for tool_call in message.tool_calls %}
82 + {%- if tool_call.function is defined %}
83 + {%- set tool_call = tool_call.function %}
84 + {%- endif %}
85 + {{- '\n<tool_call>\n<function=' + tool_call.name + '>\n' }}
86 + {%- if tool_call.arguments is defined %}
87 + {%- for args_name, args_value in tool_call.arguments|items %}
88 + {{- '<parameter=' + args_name + '>\n' }}
89 + {%- set args_value = args_value | tojson | safe if args_value is mapping or (args_value is sequence and args_value is not string) else args_value | string %}
90 + {{- args_value }}
91 + {{- '\n</parameter>\n' }}
92 + {%- endfor %}
93 + {%- endif %}
94 + {{- '</function>\n</tool_call>' }}
95 + {%- endfor %}
96 + {{- '<|im_end|>\n' }}
97 + {%- elif message.role == "user" or message.role == "system" or message.role == "assistant" %}
98 + {{- '<|im_start|>' + message.role + '\n' + message.content + '<|im_end|>' + '\n' }}
99 + {%- elif message.role == "tool" %}
100 + {%- if loop.previtem and loop.previtem.role != "tool" %}
101 + {{- '<|im_start|>user\n' }}
102 + {%- endif %}
103 + {{- '<tool_response>\n' }}
104 + {{- message.content }}
105 + {{- '\n</tool_response>\n' }}
106 + {%- if not loop.last and loop.nextitem.role != "tool" %}
107 + {{- '<|im_end|>\n' }}
108 + {%- elif loop.last %}
109 + {{- '<|im_end|>\n' }}
110 + {%- endif %}
111 + {%- else %}
112 + {{- '<|im_start|>' + message.role + '\n' + message.content + '<|im_end|>\n' }}
113 + {%- endif %}
114 +{%- endfor %}
115 +{%- if add_generation_prompt %}
116 + {{- '<|im_start|>assistant\n' }}
117 +{%- endif %}
rust/src/chat/tests/templates/vllm_examples/tool_chat_template_toolace.jinja new
+65
@@ -0,0 +1,65 @@
1 +{{- bos_token }}
2 +
3 +{%- if custom_tools is defined %}
4 + {%- set tools = custom_tools %}
5 +{%- endif %}
6 +{%- if not tools is defined %}
7 + {%- set tools = none %}
8 +{%- endif %}
9 +
10 +{#- This block extracts the system message, so we can slot it into the right place. #}
11 +{%- if messages[0]['role'] == 'system' %}
12 + {%- set system_message = messages[0]['content']|trim %}
13 + {%- set messages = messages[1:] %}
14 +{%- else %}
15 + {%- set system_message = "You are a helpful assistant with tool calling capabilities. Only reply with a tool call if the function exists in the library provided by the user. If it doesn't exist, just reply directly in natural language." %}
16 +{%- endif %}
17 +
18 +{{- "<|start_header_id|>system<|end_header_id|>\n\n" }}
19 +{%- if tools is not none and not tools_in_user_message %}
20 + {{- "You are an expert in composing functions. You are given a question and a set of possible functions. Based on the question, you will need to make one or more function/tool calls to achieve the purpose.\n" }}
21 + {{- "If none of the function can be used, point it out. If the given question lacks the parameters required by the function, also point it out.\n" }}
22 + {{- "You should only return the function call in tools call sections.\n\n" }}
23 + {{- "If you decide to invoke any of the function(s), you MUST put it in the format of [func_name1(params_name1=params_value1, params_name2=params_value2...), func_name2(params)]\n" }}
24 + {{- "You SHOULD NOT include any other text in the response.\n" }}
25 + {{- "Here is a list of functions in JSON format that you can invoke.\n" }}
26 + {%- for t in tools %}
27 + {{- t | tojson(indent=4) }}
28 + {{- "\n\n" }}
29 + {%- endfor %}
30 + {{- "\n" }}
31 +{%- endif %}
32 +{{- system_message }}
33 +{{- "<|eot_id|>" }}
34 +
35 +{%- for message in messages %}
36 + {%- if not (message.role == 'ipython' or message.role == 'tool' or 'tool_calls' in message) %}
37 + {{- '<|start_header_id|>' + message['role'] + '<|end_header_id|>\n\n'+ message['content'] | trim + '<|eot_id|>' }}
38 + {%- elif 'tool_calls' in message %}
39 + {{- '<|start_header_id|>assistant<|end_header_id|>\n\n[' -}}
40 + {%- for tool_call in message.tool_calls %}
41 + {%- if tool_call.function is defined %}
42 + {%- set tool_call = tool_call.function %}
43 + {%- endif %}
44 + {{- tool_call.name + '(' -}}
45 + {%- for param in tool_call.arguments %}
46 + {{- param + '=' -}}
47 + {{- "%s" | format(tool_call.arguments[param]) -}}
48 + {% if not loop.last %}, {% endif %}
49 + {%- endfor %}
50 + {{- ')' -}}
51 + {% if not loop.last %}, {% endif %}
52 + {%- endfor %}
53 + {{- ']<|eot_id|>' -}}
54 + {%- elif message.role == "tool" or message.role == "ipython" %}
55 + {{- "<|start_header_id|>ipython<|end_header_id|>\n\n" }}
56 + {%- if message.content is mapping %}
57 + {{- message.content | tojson }}
58 + {%- else %}
59 + {{- { "output": message.content } | tojson }}
60 + {%- endif %}
61 + {{- "<|eot_id|>" }}
62 + {%- endif %}
63 +{%- endfor %}
64 +
65 +{{- '<|start_header_id|>assistant<|end_header_id|>\n\n' }}
rust/src/chat/tests/templates/vllm_examples/tool_chat_template_xlam_llama.jinja new
+77
@@ -0,0 +1,77 @@
1 +{{- bos_token }}
2 +{%- if custom_tools is defined %}
3 + {%- set tools = custom_tools %}
4 +{%- endif %}
5 +{%- if not tools_in_user_message is defined %}
6 + {%- set tools_in_user_message = true %}
7 +{%- endif %}
8 +{%- if not tools is defined %}
9 + {%- set tools = none %}
10 +{%- endif %}
11 +
12 +{#- Extract system message #}
13 +{{- "<|start_header_id|>system<|end_header_id|>\n\n" }}
14 +{%- if messages[0]['role'] == 'system' %}
15 + {%- set system_message = messages[0]['content'] | trim %}
16 + {%- set messages = messages[1:] %}
17 + {{- system_message + "\n" }}
18 +{%- else %}
19 + {%- set system_message = "You are a helpful assistant. You are developed by Salesforce xLAM team." %}
20 + {% set format_instruction %}You have access to a set of tools. When using tools, make calls in a single JSON array:
21 +
22 +[{"name": "tool_call_name", "arguments": {"arg1": "value1", "arg2": "value2"}}, ... (additional parallel tool calls as needed)]
23 +
24 +If no tool is suitable, state that explicitly. If the user's input lacks required parameters, ask for clarification. Do not interpret or respond until tool results are returned. Once they are available, process them or make additional calls if needed. For tasks that don't require tools, such as casual conversation or general advice, respond directly in plain text. The available tools are:{% endset %}
25 + {{- system_message + "\n" }}
26 + {%- if tools is not none %}
27 + {{- format_instruction + "\n\n" }}
28 + {%- endif %}
29 +{%- endif %}
30 +
31 +
32 +{%- if tools is not none %}
33 + {%- for t in tools %}
34 + {{- t | tojson(indent=4) }}
35 + {{- "\n\n" }}
36 + {%- endfor %}
37 +{%- endif %}
38 +{{- "<|eot_id|>" }}
39 +
40 +{%- for message in messages %}
41 + {%- if not (message.role == 'ipython' or message.role == 'tool' or 'tool_calls' in message) %}
42 + {{- '<|start_header_id|>' + message['role'] + '<|end_header_id|>\n\n'+ message['content'] | trim + '<|eot_id|>' }}
43 + {%- elif 'tool_calls' in message %}
44 + {{- '<|start_header_id|>assistant<|end_header_id|>\n\n' -}}
45 + {%- if message['tool_calls'] %}
46 + {{- "[" }}
47 + {%- for tool_call_function in message.tool_calls %}
48 + {%- set tool_call = tool_call_function.function %}
49 + {{- '{"name": "' + tool_call.name + '", ' }}
50 + {{- '"arguments": ' }}
51 + {{- tool_call.arguments | tojson }}
52 + {{- "}" }}
53 + {%- if not loop.last %}
54 + {{- ", " }}
55 + {%- endif %}
56 + {%- endfor %}
57 + {{- "]" }}
58 + {{- "<|eot_id|>" }}
59 + {%- elif message['content'] %}
60 + {{- message['content'] | trim + '<|eot_id|>' }}
61 + {%- else %}
62 + {{- "[]\n" + '<|eot_id|>' }}
63 + {%- endif %}
64 + {%- elif message.role == "tool" or message.role == "ipython" %}
65 + {{- "<|start_header_id|>" + "ipython" + "<|end_header_id|>\n\n" }}
66 + {%- set content = message["content"] %}
67 + {%- if content is mapping or (content is iterable and content is not string) %}
68 + {{- content | tojson }}
69 + {%- else %}
70 + {{- content }}
71 + {%- endif %}
72 + {{- "<|eot_id|>" }}
73 + {%- endif %}
74 +{%- endfor %}
75 +{%- if add_generation_prompt %}
76 + {{- '<|start_header_id|>assistant<|end_header_id|>\n\n' }}
77 +{%- endif %}
\ No newline at end of file
rust/src/chat/tests/templates/vllm_examples/tool_chat_template_xlam_qwen.jinja new
+66
@@ -0,0 +1,66 @@
1 +{# System message #}
2 +{{- "<|im_start|>system\n" }}
3 +{%- if messages[0]['role'] == 'system' %}
4 + {%- set system_message = messages[0]['content'] | trim %}
5 + {%- set messages = messages[1:] %}
6 + {{- system_message + "\n" }}
7 +{%- else %}
8 + {%- set system_message = "You are a helpful assistant. You are developed by Salesforce xLAM team." %}
9 + {% set format_instruction %}You have access to a set of tools. When using tools, make calls in a single JSON array:
10 +
11 +[{"name": "tool_call_name", "arguments": {"arg1": "value1", "arg2": "value2"}}, ... (additional parallel tool calls as needed)]
12 +
13 +If no tool is suitable, state that explicitly. If the user's input lacks required parameters, ask for clarification. Do not interpret or respond until tool results are returned. Once they are available, process them or make additional calls if needed. For tasks that don't require tools, such as casual conversation or general advice, respond directly in plain text. The available tools are:{% endset %}
14 + {{- system_message + "\n" }}
15 + {%- if tools is not none %}
16 + {{- format_instruction + "\n\n" }}
17 + {%- endif %}
18 +{%- endif %}
19 +
20 +{%- if tools is not none %}
21 + {%- for func in tools %}
22 + {{- func | tojson(indent=4) }}
23 + {{- "\n\n" }}
24 + {%- endfor %}
25 +{%- endif %}
26 +{{- "<|im_end|>\n" }}
27 +{%- for message in messages %}
28 + {%- if message['role'] == 'tool' %}
29 + {{- "<|im_start|>tool\n" }}
30 + {%- if message.content is defined and message.content.content is defined %}
31 + {%- set content = message.content.content %}
32 + {%- else %}
33 + {%- set content = message.content %}
34 + {%- endif %}
35 + {%- if content is mapping or content is iterable and content is not string %}
36 + {{- content | tojson }}
37 + {%- else %}
38 + {{- content }}
39 + {%- endif %}
40 + {{- "<|im_end|>\n" }}
41 + {%- elif 'tool_calls' in message %}
42 + {{- "<|im_start|>assistant\n" }}
43 + {%- if message['tool_calls'] %}
44 + {{- "[" }}
45 + {%- for tool_call in message.tool_calls %}
46 + {%- set out = tool_call.function | tojson %}
47 + {{- out }}
48 + {%- if not loop.last %}
49 + {{- ", " }}
50 + {%- endif %}
51 + {%- endfor %}
52 + {{- "]"}}
53 + {%- elif message['content'] %}
54 + {{- message['content'] | trim }}
55 + {%- else %}
56 + {{- "[]\n" }}
57 + {%- endif %}
58 + {{- "<|im_end|>\n" }}
59 + {%- else %}
60 + {{- "<|im_start|>" + message['role'] + "\n" + message['content'] | trim + "<|im_end|>\n" }}
61 + {%- endif %}
62 +{%- endfor %}
63 +
64 +{%- if add_generation_prompt %}
65 + {{- "<|im_start|>assistant\n" }}
66 +{%- endif %}
rust/src/cmd/Cargo.toml new
+39
@@ -0,0 +1,39 @@
1 +[package]
2 +name = "vllm-cmd"
3 +version.workspace = true
4 +edition.workspace = true
5 +license.workspace = true
6 +
7 +[[bin]]
8 +name = "vllm-rs"
9 +path = "src/main.rs"
10 +
11 +[features]
12 +default = []
13 +native-tls-vendored = ["dep:native-tls-vendored"]
14 +
15 +[dependencies]
16 +anyhow.workspace = true
17 +clap.workspace = true
18 +educe.workspace = true
19 +itertools.workspace = true
20 +native-tls-vendored = { workspace = true, optional = true }
21 +serde.workspace = true
22 +serde_json.workspace = true
23 +serde_with.workspace = true
24 +thiserror-ext.workspace = true
25 +time.workspace = true
26 +tokio = { workspace = true, features = ["signal"] }
27 +tokio-util.workspace = true
28 +tracing.workspace = true
29 +tracing-subscriber.workspace = true
30 +uuid.workspace = true
31 +vllm-engine-core-client.workspace = true
32 +vllm-managed-engine.workspace = true
33 +vllm-server.workspace = true
34 +
35 +[dev-dependencies]
36 +expect-test.workspace = true
37 +
38 +[lints]
39 +workspace = true
rust/src/cmd/examples/README.md new
+44
@@ -0,0 +1,44 @@
1 +# `vllm-rs` CLI Quick Start
2 +
3 +Start Qwen3 with one managed `vllm-rs serve` command from the repo root:
4 +
5 +```bash
6 +HF_HUB_OFFLINE=1 \
7 +VLLM_CPU_KVCACHE_SPACE=2 \
8 +VLLM_HOST_IP=127.0.0.1 \
9 +VLLM_LOOPBACK_IP=127.0.0.1 \
10 +cargo run --bin vllm-rs -- serve \
11 + Qwen/Qwen3-0.6B \
12 + --python ../vllm/.venv/bin/python \
13 + --max-model-len 512 \
14 + -- \
15 + --dtype float16
16 +```
17 +
18 +This launches:
19 +
20 +- a managed headless Python `vllm` engine
21 +- the Rust OpenAI-compatible frontend on `127.0.0.1:8000`
22 +
23 +All Python engine arguments must be placed after `--`. Arguments before `--` are parsed by the Rust
24 +frontend itself.
25 +
26 +You can then send OpenAI-style requests to the Rust frontend:
27 +
28 +```bash
29 +curl http://127.0.0.1:8000/v1/chat/completions \
30 + -H "Content-Type: application/json" \
31 + -d '{
32 + "model": "Qwen/Qwen3-0.6B",
33 + "messages": [{"role": "user", "content": "What is the capital of France?"}],
34 + "stream": true
35 + }'
36 +```
37 +
38 +If you already started headless `vllm` yourself, use `frontend` instead:
39 +
40 +```bash
41 +cargo run --bin vllm-rs -- frontend \
42 + --handshake-address tcp://127.0.0.1:62100 \
43 + Qwen/Qwen3-0.6B
44 +```
rust/src/cmd/src/cli.rs new
+434
@@ -0,0 +1,434 @@
1 +//! CLI argument definitions for the `vllm-rs` binary.
2 +//!
3 +//! Python vLLM references:
4 +//! - Engine args: <https://github.com/vllm-project/vllm/blob/bc2c0c86efb28e77677a3cfb8687e976914a313a/vllm/engine/arg_utils.py#L657-L1311>
5 +//! - Environment variables: <https://github.com/vllm-project/vllm/blob/bc2c0c86efb28e77677a3cfb8687e976914a313a/vllm/envs.py#L472>
6 +
7 +mod unsupported;
8 +
9 +use std::collections::HashMap;
10 +use std::ffi::{OsStr, OsString};
11 +use std::path::PathBuf;
12 +use std::time::Duration;
13 +
14 +use clap::{Args, Parser, Subcommand};
15 +use educe::Educe;
16 +use serde::Deserialize;
17 +use serde::de::DeserializeOwned;
18 +use serde_json::Value;
19 +use thiserror_ext::AsReport as _;
20 +use uuid::Uuid;
21 +use vllm_engine_core_client::TransportMode;
22 +use vllm_managed_engine::ManagedEngineConfig;
23 +use vllm_managed_engine::cli::{ManagedEngineArgs, repartition_managed_engine_args};
24 +use vllm_server::{
25 + ChatTemplateContentFormatOption, Config, CoordinatorMode, HttpListenerMode, ParserSelection,
26 + RendererSelection,
27 +};
28 +
29 +use crate::cli::unsupported::UnsupportedArgs;
30 +
31 +/// Top-level parser for the `vllm-rs` binary.
32 +#[derive(Debug, Parser)]
33 +#[command(
34 + name = "vllm-rs",
35 + about = "Rust frontend and managed-engine CLI for vLLM."
36 +)]
37 +pub struct Cli {
38 + #[command(subcommand)]
39 + pub command: Command,
40 +}
41 +
42 +impl Cli {
43 + pub fn parse() -> Self {
44 + Self::try_parse_from(std::env::args_os()).unwrap_or_else(|error| error.exit())
45 + }
46 +
47 + pub fn try_parse_from<I, T>(itr: I) -> Result<Self, clap::Error>
48 + where
49 + I: IntoIterator<Item = T>,
50 + T: Into<OsString>,
51 + {
52 + let args: Vec<OsString> = itr.into_iter().map(Into::into).collect();
53 + let repartitioned_args = repartition_managed_engine_args::<Self>(&args, Some("serve"))?;
54 + <Self as Parser>::try_parse_from(&repartitioned_args).inspect(|cli| {
55 + if let Command::Serve(serve) = &cli.command
56 + && serve.debug_cli
57 + {
58 + println!(
59 + "Original CLI args: {}\n",
60 + args.join(OsStr::new(" ")).display()
61 + );
62 + println!(
63 + "Repartitioned CLI args: {}\n",
64 + repartitioned_args.join(OsStr::new(" ")).display()
65 + );
66 + println!(
67 + "Passthrough Python args: {}",
68 + serve.managed_engine.python_args.join(" ")
69 + );
70 + std::process::exit(0);
71 + }
72 + })
73 + }
74 +}
75 +
76 +/// Supported top-level CLI commands.
77 +#[derive(Debug, Subcommand, PartialEq, Eq)]
78 +pub enum Command {
79 + /// Run the Rust OpenAI frontend as a Python-supervised worker.
80 + Frontend(FrontendArgs),
81 + /// Launch a managed Python headless engine, then run the Rust OpenAI
82 + /// frontend.
83 + Serve(ServeArgs),
84 +}
85 +
86 +/// Runtime arguments shared by the external-engine and managed-engine paths.
87 +#[derive(Educe, Clone, Args, PartialEq, Eq, Deserialize)]
88 +#[educe(Debug)]
89 +pub struct SharedRuntimeArgs {
90 + #[serde(rename = "model_tag")]
91 + /// Model identifier or local model directory used for backend loading and
92 + /// public model ID.
93 + pub model: String,
94 +
95 + /// Maximum time to wait for the expected engines to register on the
96 + /// frontend transport.
97 + #[arg(
98 + long = "engine-ready-timeout-secs",
99 + env = "VLLM_ENGINE_READY_TIMEOUT_S",
100 + default_value_t = default_engine_ready_timeout_secs()
101 + )]
102 + #[serde(default = "default_engine_ready_timeout_secs")]
103 + pub engine_ready_timeout_secs: u64,
104 +
105 + /// Select the tool call parser depending on the model that you're using.
106 + /// Use `auto` to infer from the model or `none` to disable parsing.
107 + #[arg(long, default_value_t)]
108 + #[serde(default)]
109 + pub tool_call_parser: ParserSelection,
110 + /// Select the reasoning parser depending on the model that you're using.
111 + /// Use `auto` to infer from the model or `none` to disable parsing.
112 + #[arg(long, default_value_t)]
113 + #[serde(default)]
114 + pub reasoning_parser: ParserSelection,
115 + /// Select the chat renderer implementation.
116 + #[arg(long = "tokenizer-mode", default_value_t)]
117 + #[serde(default, rename = "tokenizer_mode")]
118 + pub renderer: RendererSelection,
119 + /// Override the maximum model context length. When set, the frontend uses
120 + /// this value instead of the model's `max_position_embeddings` from
121 + /// `config.json`.
122 + #[arg(long)]
123 + pub max_model_len: Option<u32>,
124 + /// TCP port for the gRPC Generate service. When not set, no gRPC server is
125 + /// started.
126 + #[arg(long)]
127 + #[serde(default)]
128 + pub grpc_port: Option<u16>,
129 + /// Maximum time to wait for active requests to drain during shutdown.
130 + #[arg(long, default_value_t = 0)]
131 + #[serde(default)]
132 + pub shutdown_timeout: u64,
133 +
134 + /// The file path to the chat template, or the template in single-line form
135 + /// for the specified model.
136 + #[arg(long)]
137 + #[serde(default)]
138 + pub chat_template: Option<String>,
139 +
140 + /// Default keyword arguments to pass to the chat template renderer.
141 + ///
142 + /// These will be merged with request-level chat_template_kwargs, with
143 + /// request values taking precedence. Useful for setting default
144 + /// behavior for reasoning models.
145 + ///
146 + /// Example: `{"enable_thinking": false}` to disable thinking mode by
147 + /// default for Qwen3/DeepSeek models.
148 + #[arg(long, value_parser = parse_json::<HashMap<String, Value>>, value_name = "JSON")]
149 + #[serde(default)]
150 + pub default_chat_template_kwargs: Option<HashMap<String, Value>>,
151 +
152 + /// The format to render message content within a chat template.
153 + ///
154 + /// * "auto" detects the format from the template
155 + /// * "string" renders content as a string. Example: `"Hello World"`
156 + /// * "openai" renders content as a list of dictionaries, similar to OpenAI schema. Example:
157 + /// `[{"type": "text", "text": "Hello world!"}]`
158 + #[arg(long, default_value_t)]
159 + #[serde(default)]
160 + pub chat_template_content_format: ChatTemplateContentFormatOption,
161 +
162 + /// Log a summary line for each completed request, including prompt/output
163 + /// token counts and finish reason.
164 + #[arg(long)]
165 + #[serde(default)]
166 + pub enable_log_requests: bool,
167 +
168 + /// Disable periodic logging of engine statistics (throughput, queue depth,
169 + /// cache usage).
170 + #[arg(long)]
171 + #[serde(default)]
172 + pub disable_log_stats: bool,
173 +
174 + /// The model name(s) used in the API. If multiple names are provided, the
175 + /// server will respond to any of the provided names. The model name in the
176 + /// model field of a response will be the first name in this list. If not
177 + /// specified, the model name will be the same as the `--model` argument.
178 + /// Noted that this name(s) will also be used in `model_name` tag
179 + /// content of prometheus metrics, if multiple names provided, metrics
180 + /// tag will take the first one.
181 + #[arg(long, num_args = 0..)]
182 + #[serde(default)]
183 + pub served_model_name: Vec<String>,
184 +
185 + /// Unsupported Python vLLM frontend arguments recognized but not yet
186 + /// implemented in Rust.
187 + #[educe(Debug(ignore))]
188 + #[command(flatten)]
189 + #[serde(default, flatten)]
190 + pub unsupported: UnsupportedArgs,
191 +}
192 +
193 +impl SharedRuntimeArgs {
194 + /// Maximum time to wait for the expected engines to register on the
195 + /// frontend transport.
196 + pub fn ready_timeout(&self) -> Duration {
197 + Duration::from_secs(self.engine_ready_timeout_secs)
198 + }
199 +
200 + /// Maximum time to wait for active requests to drain during shutdown.
201 + pub fn shutdown_timeout(&self) -> Duration {
202 + Duration::from_secs(self.shutdown_timeout)
203 + }
204 +
205 + /// Build the OpenAI-server config for the Python-bootstrap worker contract.
206 + ///
207 + /// The resulting config binds the Python-supplied transport addresses and
208 + /// inherits an already open HTTP listener from the supervisor process.
209 + fn into_bootstrapped_config(
210 + self,
211 + listen_fd: i32,
212 + input_address: String,
213 + output_address: String,
214 + coordinator_address: Option<String>,
215 + engine_count: usize,
216 + ) -> Config {
217 + let ready_timeout = self.ready_timeout();
218 + let shutdown_timeout = self.shutdown_timeout();
219 +
220 + Config {
221 + transport_mode: TransportMode::Bootstrapped {
222 + input_address,
223 + output_address,
224 + engine_count,
225 + ready_timeout,
226 + },
227 + coordinator_mode: match coordinator_address {
228 + Some(address) => CoordinatorMode::External { address },
229 + None => CoordinatorMode::None,
230 + },
231 + model: self.model,
232 + served_model_name: self.served_model_name,
233 + listener_mode: HttpListenerMode::InheritedFd { fd: listen_fd },
234 + tool_call_parser: self.tool_call_parser,
235 + reasoning_parser: self.reasoning_parser,
236 + renderer: self.renderer,
237 + chat_template: self.chat_template,
238 + default_chat_template_kwargs: self.default_chat_template_kwargs,
239 + chat_template_content_format: self.chat_template_content_format,
240 + enable_log_requests: self.enable_log_requests,
241 + disable_log_stats: self.disable_log_stats,
242 + grpc_port: self.grpc_port,
243 + shutdown_timeout,
244 + }
245 + }
246 +
247 + /// Build the OpenAI-server config for the managed `serve` path that still
248 + /// owns the startup handshake and binds its own HTTP listener.
249 + fn into_managed_config(
250 + self,
251 + listener_mode: HttpListenerMode,
252 + handshake_address: String,
253 + advertised_host: String,
254 + engine_count: usize,
255 + local_input_address: Option<String>,
256 + local_output_address: Option<String>,
257 + ) -> Config {
258 + let ready_timeout = self.ready_timeout();
259 + let shutdown_timeout = self.shutdown_timeout();
260 +
261 + Config {
262 + transport_mode: TransportMode::HandshakeOwner {
263 + handshake_address,
264 + advertised_host,
265 + engine_count,
266 + ready_timeout,
267 + local_input_address,
268 + local_output_address,
269 + },
270 + coordinator_mode: CoordinatorMode::MaybeInProc,
271 + model: self.model,
272 + served_model_name: self.served_model_name,
273 + listener_mode,
274 + tool_call_parser: self.tool_call_parser,
275 + reasoning_parser: self.reasoning_parser,
276 + renderer: self.renderer,
277 + chat_template: self.chat_template,
278 + default_chat_template_kwargs: self.default_chat_template_kwargs,
279 + chat_template_content_format: self.chat_template_content_format,
280 + enable_log_requests: self.enable_log_requests,
281 + disable_log_stats: self.disable_log_stats,
282 + grpc_port: self.grpc_port,
283 + shutdown_timeout,
284 + }
285 + }
286 +}
287 +
288 +fn default_engine_ready_timeout_secs() -> u64 {
289 + 600
290 +}
291 +
292 +fn parse_json<T: DeserializeOwned>(value: &str) -> Result<T, String> {
293 + serde_json::from_str(value).map_err(|e| format!("invalid JSON object: {}", e.as_report()))
294 +}
295 +
296 +fn parse_runtime_args_json(value: &str) -> Result<SharedRuntimeArgs, String> {
297 + let args: SharedRuntimeArgs = serde_json::from_str(value)
298 + .map_err(|e| format!("invalid JSON arguments: {}", e.as_report()))?;
299 + args.unsupported.check()?;
300 + Ok(args)
301 +}
302 +
303 +/// Arguments for running the Rust frontend as a Python-bootstrapped worker.
304 +#[derive(Educe, Clone, Args, PartialEq, Eq)]
305 +#[educe(Debug)]
306 +pub struct FrontendArgs {
307 + /// Inherited listening socket file descriptor passed by the Python
308 + /// supervisor.
309 + #[arg(long)]
310 + pub listen_fd: i32,
311 + /// Frontend input ROUTER socket address that the Python engines will
312 + /// connect to.
313 + #[arg(long)]
314 + pub input_address: String,
315 + /// Frontend output PULL socket address that the Python engines will push
316 + /// responses to.
317 + #[arg(long)]
318 + pub output_address: String,
319 + /// Optional Python-owned frontend-side DP coordinator socket address for
320 + /// external coordinator mode in the bootstrapped frontend path, i.e.,
321 + /// `stats_update_address`.
322 + #[arg(long)]
323 + pub coordinator_address: Option<String>,
324 + /// Total number of data-parallel engines expected for this frontend.
325 + #[arg(long, default_value_t = 1)]
326 + pub engine_count: usize,
327 +
328 + /// Shared frontend arguments as one JSON object.
329 + #[arg(long = "args-json", value_parser = parse_runtime_args_json, value_name = "JSON")]
330 + pub runtime: SharedRuntimeArgs,
331 +}
332 +
333 +impl FrontendArgs {
334 + /// Convert the CLI arguments into the OpenAI server's runtime config.
335 + pub fn into_config(self) -> Config {
336 + self.runtime.into_bootstrapped_config(
337 + self.listen_fd,
338 + self.input_address,
339 + self.output_address,
340 + self.coordinator_address,
341 + self.engine_count,
342 + )
343 + }
344 +}
345 +
346 +/// Arguments for the managed-engine mode that spawns Python on behalf of the
347 +/// user.
348 +#[derive(Educe, Clone, Args, PartialEq, Eq)]
349 +#[educe(Debug)]
350 +#[command(override_usage = "vllm-rs serve <MODEL> [OPTIONS] [-- <PYTHON_ARGS>...]")]
351 +pub struct ServeArgs {
352 + /// Only launch the managed Python headless engine and do not start the Rust
353 + /// frontend.
354 + #[arg(long)]
355 + pub headless: bool,
356 + /// HTTP bind host for the OpenAI-compatible server.
357 + #[arg(long, default_value = "127.0.0.1")]
358 + pub host: String,
359 + /// HTTP bind port for the OpenAI-compatible server.
360 + #[arg(long, default_value_t = 8000)]
361 + pub port: u16,
362 + /// Unix domain socket path. If set, host and port arguments are ignored.
363 + #[arg(long)]
364 + pub uds: Option<String>,
365 +
366 + /// Flag to print debug information about CLI argument parsing and exit.
367 + #[educe(Debug(ignore))]
368 + #[arg(long, hide = true, env = "VLLM_RS_DEBUG_CLI")]
369 + pub debug_cli: bool,
370 +
371 + /// Shared frontend arguments.
372 + #[command(flatten)]
373 + pub runtime: SharedRuntimeArgs,
374 +
375 + /// Managed Python headless-engine arguments.
376 + #[command(flatten)]
377 + pub managed_engine: ManagedEngineArgs,
378 +}
379 +
380 +impl ServeArgs {
381 + /// Build the OpenAI-server runtime config used after the managed Python
382 + /// engine starts.
383 + pub fn to_frontend_config(&self, handshake_address: String) -> Config {
384 + // Prefer IPC sockets for local engine input/output.
385 + let (local_input_address, local_output_address) =
386 + self.managed_engine.frontend_local_only().then(frontend_ipc_addresses).unzip();
387 + let listener_mode = match &self.uds {
388 + Some(path) => HttpListenerMode::BindUnix { path: path.clone() },
389 + None => HttpListenerMode::BindTcp {
390 + host: self.host.clone(),
391 + port: self.port,
392 + },
393 + };
394 +
395 + self.runtime.clone().into_managed_config(
396 + listener_mode,
397 + handshake_address,
398 + self.managed_engine.handshake_host.clone(),
399 + self.managed_engine.data_parallel_size,
400 + local_input_address,
401 + local_output_address,
402 + )
403 + }
404 +
405 + /// Build the managed Python-engine spawn configuration with the given
406 + /// handshake port.
407 + pub fn to_managed_engine_config(&self, handshake_port: u16) -> ManagedEngineConfig {
408 + self.managed_engine.clone().into_config(
409 + self.runtime.model.clone(),
410 + self.runtime.max_model_len,
411 + handshake_port,
412 + )
413 + }
414 +}
415 +
416 +/// Allocate fresh IPC endpoints for one managed frontend instance.
417 +fn frontend_ipc_addresses() -> (String, String) {
418 + let preferred_base_path = std::env::var_os("VLLM_RPC_BASE_PATH")
419 + .map(PathBuf::from)
420 + .unwrap_or_else(std::env::temp_dir);
421 + let input_name = format!("vllm-rs-i-{}", Uuid::new_v4().simple());
422 + let output_name = format!("vllm-rs-o-{}", Uuid::new_v4().simple());
423 +
424 + let input = preferred_base_path.join(input_name);
425 + let output = preferred_base_path.join(output_name);
426 +
427 + (
428 + format!("ipc://{}", input.to_string_lossy()),
429 + format!("ipc://{}", output.to_string_lossy()),
430 + )
431 +}
432 +
433 +#[cfg(test)]
434 +mod tests;
rust/src/cmd/src/cli/tests.rs new
+905
@@ -0,0 +1,905 @@
1 +use expect_test::expect;
2 +use vllm_engine_core_client::TransportMode;
3 +use vllm_server::{Config, HttpListenerMode, ParserSelection, RendererSelection};
4 +
5 +use super::{Cli, Command};
6 +
7 +#[test]
8 +fn serve_args_forward_python_flags_with_separator() {
9 + let cli = Cli::try_parse_from([
10 + "vllm-rs",
11 + "serve",
12 + "Qwen/Qwen3-0.6B",
13 + "--python",
14 + "../vllm/.venv/bin/python",
15 + "--max-model-len",
16 + "512",
17 + "--",
18 + "--dtype",
19 + "float16",
20 + ])
21 + .unwrap();
22 +
23 + expect![[r#"
24 + Cli {
25 + command: Serve(
26 + ServeArgs {
27 + headless: false,
28 + host: "127.0.0.1",
29 + port: 8000,
30 + uds: None,
31 + runtime: SharedRuntimeArgs {
32 + model: "Qwen/Qwen3-0.6B",
33 + engine_ready_timeout_secs: 600,
34 + tool_call_parser: Auto,
35 + reasoning_parser: Auto,
36 + renderer: Auto,
37 + max_model_len: Some(
38 + 512,
39 + ),
40 + grpc_port: None,
41 + shutdown_timeout: 0,
42 + chat_template: None,
43 + default_chat_template_kwargs: None,
44 + chat_template_content_format: Auto,
45 + enable_log_requests: false,
46 + disable_log_stats: false,
47 + served_model_name: [],
48 + },
49 + managed_engine: ManagedEngineArgs {
50 + python: "../vllm/.venv/bin/python",
51 + handshake_host: "127.0.0.1",
52 + handshake_port: None,
53 + data_parallel_size: 1,
54 + data_parallel_size_local: None,
55 + python_args: [
56 + "--dtype",
57 + "float16",
58 + ],
59 + },
60 + },
61 + ),
62 + }
63 + "#]]
64 + .assert_debug_eq(&cli);
65 +}
66 +
67 +#[test]
68 +fn serve_args_auto_forward_python_flags_without_separator() {
69 + let cli = Cli::try_parse_from([
70 + "vllm-rs",
71 + "serve",
72 + "Qwen/Qwen3-0.6B",
73 + "--python",
74 + "python3",
75 + "--quantization",
76 + "awq",
77 + ])
78 + .unwrap();
79 +
80 + let Command::Serve(args) = cli.command else {
81 + panic!("expected serve args");
82 + };
83 + assert_eq!(
84 + args.managed_engine.python_args,
85 + vec!["--quantization", "awq"]
86 + );
87 +}
88 +
89 +#[test]
90 +fn serve_args_auto_forward_python_multi_char_alias_without_separator() {
91 + let cli = Cli::try_parse_from(["vllm-rs", "serve", "Qwen/Qwen3-0.6B", "-tp", "2"]).unwrap();
92 +
93 + let Command::Serve(args) = cli.command else {
94 + panic!("expected serve args");
95 + };
96 + assert_eq!(
97 + args.managed_engine.python_args,
98 + vec!["--tensor-parallel-size", "2"]
99 + );
100 +}
101 +
102 +#[test]
103 +fn serve_args_accept_explicit_deepseek_v32_renderer() {
104 + let cli = Cli::try_parse_from([
105 + "vllm-rs",
106 + "serve",
107 + "Qwen/Qwen3-0.6B",
108 + "--tokenizer-mode",
109 + "deepseek_v32",
110 + ])
111 + .unwrap();
112 +
113 + let Command::Serve(args) = cli.command else {
114 + panic!("expected serve args");
115 + };
116 + assert_eq!(args.runtime.renderer, RendererSelection::DeepSeekV32);
117 +}
118 +
119 +#[test]
120 +fn serve_args_reject_unknown_renderer_value() {
121 + let error = Cli::try_parse_from([
122 + "vllm-rs",
123 + "serve",
124 + "Qwen/Qwen3-0.6B",
125 + "--tokenizer-mode",
126 + "definitely_missing",
127 + ])
128 + .unwrap_err();
129 +
130 + expect![[r#"
131 + error: invalid value 'definitely_missing' for '--tokenizer-mode <RENDERER>': unknown renderer `definitely_missing` (expected one of: auto, hf, deepseek_v32, deepseek_v4)
132 +
133 + For more information, try '--help'.
134 + "#]]
135 + .assert_eq(&error.to_string());
136 +}
137 +
138 +#[test]
139 +fn serve_args_reject_unsupported_flag_arg() {
140 + let error = Cli::try_parse_from(["vllm-rs", "serve", "Qwen/Qwen3-0.6B", "--allow-credentials"])
141 + .unwrap_err();
142 +
143 + expect![[r#"
144 + error: invalid value 'true' for '--allow-credentials [<ALLOW_CREDENTIALS>]': argument is not implemented in Rust frontend yet
145 +
146 + Remove this unsupported argument to continue.
147 +
148 + Alternatively, if you intend to pass it only to the Python engine, put it after `--` (e.g., `-- <arg>`).
149 + This may lead to unexpected behavior as the Rust frontend will completely ignore that argument.
150 +
151 + For more information, try '--help'.
152 + "#]]
153 + .assert_eq(&error.to_string());
154 +}
155 +
156 +#[test]
157 +fn serve_args_reject_unsupported_no_flag_alias() {
158 + let error = Cli::try_parse_from([
159 + "vllm-rs",
160 + "serve",
161 + "Qwen/Qwen3-0.6B",
162 + "--no-enable-log-deltas",
163 + ])
164 + .unwrap_err();
165 +
166 + expect![[r#"
167 + error: invalid value 'true' for '--enable-log-deltas [<ENABLE_LOG_DELTAS>]': argument is not implemented in Rust frontend yet
168 +
169 + Remove this unsupported argument to continue.
170 +
171 + Alternatively, if you intend to pass it only to the Python engine, put it after `--` (e.g., `-- <arg>`).
172 + This may lead to unexpected behavior as the Rust frontend will completely ignore that argument.
173 +
174 + For more information, try '--help'.
175 + "#]]
176 + .assert_eq(&error.to_string());
177 +}
178 +
179 +#[test]
180 +fn frontend_args_accept_json() {
181 + let cli = Cli::try_parse_from([
182 + "vllm-rs",
183 + "frontend",
184 + "--listen-fd",
185 + "3",
186 + "--input-address",
187 + "ipc:///tmp/input.sock",
188 + "--output-address",
189 + "ipc:///tmp/output.sock",
190 + "--coordinator-address",
191 + "tcp://127.0.0.1:7000",
192 + "--args-json",
193 + r#"{"model_tag":"Qwen/Qwen3-0.6B","engine_count":2}"#,
194 + ])
195 + .unwrap();
196 +
197 + expect![[r#"
198 + Cli {
199 + command: Frontend(
200 + FrontendArgs {
201 + listen_fd: 3,
202 + input_address: "ipc:///tmp/input.sock",
203 + output_address: "ipc:///tmp/output.sock",
204 + coordinator_address: Some(
205 + "tcp://127.0.0.1:7000",
206 + ),
207 + engine_count: 1,
208 + runtime: SharedRuntimeArgs {
209 + model: "Qwen/Qwen3-0.6B",
210 + engine_ready_timeout_secs: 600,
211 + tool_call_parser: Auto,
212 + reasoning_parser: Auto,
213 + renderer: Auto,
214 + max_model_len: None,
215 + grpc_port: None,
216 + shutdown_timeout: 0,
217 + chat_template: None,
218 + default_chat_template_kwargs: None,
219 + chat_template_content_format: Auto,
220 + enable_log_requests: false,
221 + disable_log_stats: false,
222 + served_model_name: [],
223 + },
224 + },
225 + ),
226 + }
227 + "#]]
228 + .assert_debug_eq(&cli);
229 +}
230 +
231 +#[test]
232 +fn frontend_args_json_applies_defaults() {
233 + let cli = Cli::try_parse_from([
234 + "vllm-rs",
235 + "frontend",
236 + "--listen-fd",
237 + "3",
238 + "--input-address",
239 + "ipc:///tmp/input.sock",
240 + "--output-address",
241 + "ipc:///tmp/output.sock",
242 + "--args-json",
243 + r#"{"model_tag":"Qwen/Qwen3-0.6B"}"#,
244 + ])
245 + .unwrap();
246 +
247 + let Command::Frontend(args) = cli.command else {
248 + panic!("expected frontend args");
249 + };
250 + assert_eq!(args.runtime.model, "Qwen/Qwen3-0.6B");
251 + assert_eq!(args.runtime.engine_ready_timeout_secs, 600);
252 + assert_eq!(args.runtime.tool_call_parser, ParserSelection::Auto);
253 + assert_eq!(args.runtime.reasoning_parser, ParserSelection::Auto);
254 + assert_eq!(args.runtime.renderer, RendererSelection::Auto);
255 + assert_eq!(args.runtime.max_model_len, None);
256 + assert_eq!(args.runtime.shutdown_timeout, 0);
257 +}
258 +
259 +#[test]
260 +fn frontend_args_json_accepts_supported_non_default_fields() {
261 + let cli = Cli::try_parse_from([
262 + "vllm-rs",
263 + "frontend",
264 + "--listen-fd",
265 + "3",
266 + "--input-address",
267 + "ipc:///tmp/input.sock",
268 + "--output-address",
269 + "ipc:///tmp/output.sock",
270 + "--args-json",
271 + r#"{"model_tag":"Qwen/Qwen3-0.6B","engine_ready_timeout_secs":42,"tool_call_parser":"hermes","reasoning_parser":"qwen3_thinking","tokenizer_mode":"deepseek_v32","max_model_len":8192,"shutdown_timeout":3}"#,
272 + ])
273 + .unwrap();
274 +
275 + let Command::Frontend(args) = cli.command else {
276 + panic!("expected frontend args");
277 + };
278 + assert_eq!(args.runtime.engine_ready_timeout_secs, 42);
279 + assert_eq!(
280 + args.runtime.tool_call_parser,
281 + ParserSelection::Explicit("hermes".to_string())
282 + );
283 + assert_eq!(
284 + args.runtime.reasoning_parser,
285 + ParserSelection::Explicit("qwen3_thinking".to_string())
286 + );
287 + assert_eq!(args.runtime.renderer, RendererSelection::DeepSeekV32);
288 + assert_eq!(args.runtime.max_model_len, Some(8192));
289 + assert_eq!(args.runtime.shutdown_timeout, 3);
290 +}
291 +
292 +#[test]
293 +fn serve_args_accept_none_reasoning_parser() {
294 + let cli = Cli::try_parse_from([
295 + "vllm-rs",
296 + "serve",
297 + "Qwen/Qwen3-0.6B",
298 + "--reasoning-parser",
299 + "none",
300 + ])
301 + .unwrap();
302 +
303 + let Command::Serve(args) = cli.command else {
304 + panic!("expected serve args");
305 + };
306 + assert_eq!(args.runtime.reasoning_parser, ParserSelection::None);
307 + assert_eq!(args.runtime.tool_call_parser, ParserSelection::Auto);
308 +}
309 +
310 +#[test]
311 +fn frontend_args_json_ignores_unknown_fields() {
312 + let cli = Cli::try_parse_from([
313 + "vllm-rs",
314 + "frontend",
315 + "--listen-fd",
316 + "3",
317 + "--input-address",
318 + "ipc:///tmp/input.sock",
319 + "--output-address",
320 + "ipc:///tmp/output.sock",
321 + "--args-json",
322 + r#"{"model_tag":"Qwen/Qwen3-0.6B","uds":"/tmp/vllm.sock","nested_unknown":{"x":1}}"#,
323 + ])
324 + .unwrap();
325 +
326 + let Command::Frontend(args) = cli.command else {
327 + panic!("expected frontend args");
328 + };
329 + assert_eq!(args.runtime.model, "Qwen/Qwen3-0.6B");
330 +}
331 +
332 +#[test]
333 +fn frontend_args_json_accepts_noop_fields() {
334 + let cli = Cli::try_parse_from([
335 + "vllm-rs",
336 + "frontend",
337 + "--listen-fd",
338 + "3",
339 + "--input-address",
340 + "ipc:///tmp/input.sock",
341 + "--output-address",
342 + "ipc:///tmp/output.sock",
343 + "--args-json",
344 + r#"{"model_tag":"Qwen/Qwen3-0.6B","api_server_count":2}"#,
345 + ])
346 + .unwrap();
347 +
348 + let Command::Frontend(args) = cli.command else {
349 + panic!("expected frontend args");
350 + };
351 + assert_eq!(args.runtime.model, "Qwen/Qwen3-0.6B");
352 +}
353 +
354 +#[test]
355 +fn frontend_args_json_rejects_unsupported_fields() {
356 + let error = Cli::try_parse_from([
357 + "vllm-rs",
358 + "frontend",
359 + "--listen-fd",
360 + "3",
361 + "--input-address",
362 + "ipc:///tmp/input.sock",
363 + "--output-address",
364 + "ipc:///tmp/output.sock",
365 + "--args-json",
366 + r#"{"model_tag":"Qwen/Qwen3-0.6B","allow_credentials":true}"#,
367 + ])
368 + .unwrap_err();
369 +
370 + expect![[r#"
371 + error: invalid value '{"model_tag":"Qwen/Qwen3-0.6B","allow_credentials":true}' for '--args-json <JSON>':
372 + The following arguments are not implemented in Rust frontend yet:
373 + - allow_credentials
374 +
375 + Remove these arguments to continue.
376 +
377 + For more information, try '--help'.
378 + "#]].assert_eq(&error.to_string());
379 +}
380 +
381 +#[test]
382 +fn frontend_args_json_aggregates_multiple_unsupported_fields() {
383 + let error = Cli::try_parse_from([
384 + "vllm-rs",
385 + "frontend",
386 + "--listen-fd",
387 + "3",
388 + "--input-address",
389 + "ipc:///tmp/input.sock",
390 + "--output-address",
391 + "ipc:///tmp/output.sock",
392 + "--args-json",
393 + r#"{"model_tag":"Qwen/Qwen3-0.6B","allow_credentials":true,"api_key":"secret"}"#,
394 + ])
395 + .unwrap_err();
396 +
397 + expect![[r#"
398 + error: invalid value '{"model_tag":"Qwen/Qwen3-0.6B","allow_credentials":true,"api_key":"secret"}' for '--args-json <JSON>':
399 + The following arguments are not implemented in Rust frontend yet:
400 + - allow_credentials
401 + - api_key
402 +
403 + Remove these arguments to continue.
404 +
405 + For more information, try '--help'.
406 + "#]].assert_eq(&error.to_string());
407 +}
408 +
409 +#[test]
410 +fn frontend_args_json_rejects_malformed_json() {
411 + let error = Cli::try_parse_from([
412 + "vllm-rs",
413 + "frontend",
414 + "--listen-fd",
415 + "3",
416 + "--input-address",
417 + "ipc:///tmp/input.sock",
418 + "--output-address",
419 + "ipc:///tmp/output.sock",
420 + "--args-json",
421 + r#"{"model_tag":"Qwen/Qwen3-0.6B""#,
422 + ])
423 + .unwrap_err();
424 +
425 + expect![[r#"
426 + error: invalid value '{"model_tag":"Qwen/Qwen3-0.6B"' for '--args-json <JSON>': invalid JSON arguments: EOF while parsing an object at line 1 column 30
427 +
428 + For more information, try '--help'.
429 + "#]].assert_eq(&error.to_string());
430 +}
431 +
432 +#[test]
433 +fn serve_args_reject_flags_before_model() {
434 + let error = Cli::try_parse_from(["vllm-rs", "serve", "--python", "python3", "Qwen/Qwen3-0.6B"])
435 + .unwrap_err();
436 +
437 + expect![[r#"
438 + error: the model must appear immediately after the command
439 +
440 + Usage: vllm-rs serve <MODEL> [OPTIONS] [-- <PYTHON_ARGS>...]
441 +
442 + For more information, try '--help'.
443 + "#]]
444 + .assert_eq(&error.to_string());
445 +}
446 +
447 +#[test]
448 +fn serve_args_accept_headless_mode() {
449 + let cli = Cli::try_parse_from(["vllm-rs", "serve", "Qwen/Qwen3-0.6B", "--headless"]).unwrap();
450 +
451 + let Command::Serve(args) = cli.command else {
452 + panic!("expected serve args");
453 + };
454 + assert!(args.headless);
455 +}
456 +
457 +#[test]
458 +fn serve_args_keep_python_passthrough_flags_after_separator() {
459 + let cli = Cli::try_parse_from([
460 + "vllm-rs",
461 + "serve",
462 + "Qwen/Qwen3-0.6B",
463 + "--python",
464 + "python3",
465 + "--",
466 + "--tensor-parallel-size",
467 + "2",
468 + "--dtype",
469 + "float16",
470 + ])
471 + .unwrap();
472 +
473 + let Command::Serve(args) = cli.command else {
474 + panic!("expected serve args");
475 + };
476 + assert_eq!(
477 + args.managed_engine.python_args,
478 + vec!["--tensor-parallel-size", "2", "--dtype", "float16"]
479 + );
480 +}
481 +
482 +#[test]
483 +fn serve_args_keep_python_multi_char_alias_after_separator() {
484 + let cli = Cli::try_parse_from([
485 + "vllm-rs",
486 + "serve",
487 + "Qwen/Qwen3-0.6B",
488 + "--python",
489 + "python3",
490 + "--",
491 + "-tp",
492 + "2",
493 + "--dtype",
494 + "float16",
495 + ])
496 + .unwrap();
497 +
498 + let Command::Serve(args) = cli.command else {
499 + panic!("expected serve args");
500 + };
501 + assert_eq!(
502 + args.managed_engine.python_args,
503 + vec!["-tp", "2", "--dtype", "float16"]
504 + );
505 +}
506 +
507 +#[test]
508 +fn serve_args_keep_frontend_arg_after_separator() {
509 + let cli = Cli::try_parse_from([
510 + "vllm-rs",
511 + "serve",
512 + "Qwen/Qwen3-0.6B",
513 + "--",
514 + "--uds",
515 + "/tmp/vllm.sock",
516 + ])
517 + .unwrap();
518 +
519 + let Command::Serve(args) = cli.command else {
520 + panic!("expected serve args");
521 + };
522 + assert_eq!(
523 + args.managed_engine.python_args,
524 + vec!["--uds", "/tmp/vllm.sock"]
525 + );
526 +}
527 +
528 +#[test]
529 +fn serve_args_keep_python_multi_char_engine_aliases_after_separator() {
530 + let cli = Cli::try_parse_from([
531 + "vllm-rs",
532 + "serve",
533 + "Qwen/Qwen3-0.6B",
534 + "--",
535 + "-dpr",
536 + "1",
537 + "-dpl",
538 + "2",
539 + ])
540 + .unwrap();
541 +
542 + let Command::Serve(args) = cli.command else {
543 + panic!("expected serve args");
544 + };
545 + assert_eq!(
546 + args.managed_engine.python_args,
547 + vec!["-dpr", "1", "-dpl", "2"]
548 + );
549 +}
550 +
551 +#[test]
552 +fn serve_args_auto_forward_unknown_flags_without_separator() {
553 + let cli = Cli::try_parse_from(["vllm-rs", "serve", "Qwen/Qwen3-0.6B", "--foo", "bar"]).unwrap();
554 +
555 + let Command::Serve(args) = cli.command else {
556 + panic!("expected serve args");
557 + };
558 + assert_eq!(args.managed_engine.python_args, vec!["--foo", "bar"]);
559 +}
560 +
561 +#[test]
562 +fn serve_args_auto_forward_negative_value_without_separator() {
563 + let cli = Cli::try_parse_from([
564 + "vllm-rs",
565 + "serve",
566 + "Qwen/Qwen3-0.6B",
567 + "--num-gpu-blocks-override",
568 + "-1",
569 + ])
570 + .unwrap();
571 +
572 + let Command::Serve(args) = cli.command else {
573 + panic!("expected serve args");
574 + };
575 + assert_eq!(
576 + args.managed_engine.python_args,
577 + vec!["--num-gpu-blocks-override", "-1"]
578 + );
579 +}
580 +
581 +#[test]
582 +fn serve_args_accept_handshake_aliases() {
583 + let cli = Cli::try_parse_from([
584 + "vllm-rs",
585 + "serve",
586 + "Qwen/Qwen3-0.6B",
587 + "--python",
588 + "python3",
589 + "--handshake-host",
590 + "10.99.48.128",
591 + "--handshake-port",
592 + "13345",
593 + "--data-parallel-size",
594 + "4",
595 + ])
596 + .unwrap();
597 +
598 + expect![[r#"
599 + Cli {
600 + command: Serve(
601 + ServeArgs {
602 + headless: false,
603 + host: "127.0.0.1",
604 + port: 8000,
605 + uds: None,
606 + runtime: SharedRuntimeArgs {
607 + model: "Qwen/Qwen3-0.6B",
608 + engine_ready_timeout_secs: 600,
609 + tool_call_parser: Auto,
610 + reasoning_parser: Auto,
611 + renderer: Auto,
612 + max_model_len: None,
613 + grpc_port: None,
614 + shutdown_timeout: 0,
615 + chat_template: None,
616 + default_chat_template_kwargs: None,
617 + chat_template_content_format: Auto,
618 + enable_log_requests: false,
619 + disable_log_stats: false,
620 + served_model_name: [],
621 + },
622 + managed_engine: ManagedEngineArgs {
623 + python: "python3",
624 + handshake_host: "10.99.48.128",
625 + handshake_port: Some(
626 + 13345,
627 + ),
628 + data_parallel_size: 4,
629 + data_parallel_size_local: None,
630 + python_args: [],
631 + },
632 + },
633 + ),
634 + }
635 + "#]]
636 + .assert_debug_eq(&cli);
637 +}
638 +
639 +#[test]
640 +fn serve_args_accept_data_parallel_primary_flags() {
641 + let cli = Cli::try_parse_from([
642 + "vllm-rs",
643 + "serve",
644 + "Qwen/Qwen3-0.6B",
645 + "--data-parallel-address",
646 + "10.99.48.128",
647 + "--data-parallel-rpc-port",
648 + "13345",
649 + "--data-parallel-size",
650 + "4",
651 + ])
652 + .unwrap();
653 +
654 + let Command::Serve(args) = cli.command else {
655 + panic!("expected serve args");
656 + };
657 + assert!(!args.headless);
658 + assert_eq!(args.managed_engine.handshake_host, "10.99.48.128");
659 + assert_eq!(args.managed_engine.handshake_port, Some(13345));
660 + assert_eq!(args.managed_engine.data_parallel_size, 4);
661 +}
662 +
663 +#[test]
664 +fn serve_frontend_config_uses_dp_address_as_advertised_host() {
665 + let cli = Cli::try_parse_from([
666 + "vllm-rs",
667 + "serve",
668 + "Qwen/Qwen3-0.6B",
669 + "--handshake-host",
670 + "10.99.48.128",
671 + "--data-parallel-size",
672 + "4",
673 + ])
674 + .unwrap();
675 +
676 + let Command::Serve(args) = cli.command else {
677 + panic!("expected serve args");
678 + };
679 + let config = args.to_frontend_config("tcp://10.99.48.128:29550".to_string());
680 +
681 + let TransportMode::HandshakeOwner {
682 + handshake_address,
683 + advertised_host,
684 + engine_count,
685 + ready_timeout,
686 + local_input_address,
687 + local_output_address,
688 + } = &config.transport_mode
689 + else {
690 + panic!("expected handshake-owned transport");
691 + };
692 +
693 + assert_eq!(handshake_address, "tcp://10.99.48.128:29550");
694 + assert_eq!(advertised_host, "10.99.48.128");
695 + assert_eq!(*engine_count, 4);
696 + assert!(
697 + local_input_address
698 + .as_deref()
699 + .is_some_and(|address| address.starts_with("ipc://"))
700 + );
701 + assert!(
702 + local_output_address
703 + .as_deref()
704 + .is_some_and(|address| address.starts_with("ipc://"))
705 + );
706 + assert_ne!(local_input_address, local_output_address);
707 +
708 + expect![[r#"
709 + Config {
710 + transport_mode: HandshakeOwner {
711 + handshake_address: "tcp://10.99.48.128:29550",
712 + advertised_host: "10.99.48.128",
713 + engine_count: 4,
714 + ready_timeout: 600s,
715 + local_input_address: Some(
716 + "<ipc input>",
717 + ),
718 + local_output_address: Some(
719 + "<ipc output>",
720 + ),
721 + },
722 + coordinator_mode: MaybeInProc,
723 + model: "Qwen/Qwen3-0.6B",
724 + served_model_name: [],
725 + listener_mode: BindTcp {
726 + host: "127.0.0.1",
727 + port: 8000,
728 + },
729 + tool_call_parser: Auto,
730 + reasoning_parser: Auto,
731 + renderer: Auto,
732 + chat_template: None,
733 + default_chat_template_kwargs: None,
734 + chat_template_content_format: Auto,
735 + enable_log_requests: false,
736 + disable_log_stats: false,
737 + grpc_port: None,
738 + shutdown_timeout: 0ns,
739 + }
740 + "#]]
741 + .assert_debug_eq(&Config {
742 + transport_mode: TransportMode::HandshakeOwner {
743 + handshake_address: handshake_address.clone(),
744 + advertised_host: advertised_host.clone(),
745 + engine_count: *engine_count,
746 + ready_timeout: *ready_timeout,
747 + local_input_address: Some("<ipc input>".to_string()),
748 + local_output_address: Some("<ipc output>".to_string()),
749 + },
750 + ..config.clone()
751 + });
752 +}
753 +
754 +#[test]
755 +fn serve_frontend_config_keeps_tcp_transport_for_non_local_only_topology() {
756 + let cli = Cli::try_parse_from([
757 + "vllm-rs",
758 + "serve",
759 + "Qwen/Qwen3-0.6B",
760 + "--data-parallel-address",
761 + "10.99.48.128",
762 + "--data-parallel-size",
763 + "4",
764 + "--data-parallel-size-local",
765 + "2",
766 + ])
767 + .unwrap();
768 +
769 + let Command::Serve(args) = cli.command else {
770 + panic!("expected serve args");
771 + };
772 + let config = args.to_frontend_config("tcp://10.99.48.128:29550".to_string());
773 +
774 + expect![[r#"
775 + Config {
776 + transport_mode: HandshakeOwner {
777 + handshake_address: "tcp://10.99.48.128:29550",
778 + advertised_host: "10.99.48.128",
779 + engine_count: 4,
780 + ready_timeout: 600s,
781 + local_input_address: None,
782 + local_output_address: None,
783 + },
784 + coordinator_mode: MaybeInProc,
785 + model: "Qwen/Qwen3-0.6B",
786 + served_model_name: [],
787 + listener_mode: BindTcp {
788 + host: "127.0.0.1",
789 + port: 8000,
790 + },
791 + tool_call_parser: Auto,
792 + reasoning_parser: Auto,
793 + renderer: Auto,
794 + chat_template: None,
795 + default_chat_template_kwargs: None,
796 + chat_template_content_format: Auto,
797 + enable_log_requests: false,
798 + disable_log_stats: false,
799 + grpc_port: None,
800 + shutdown_timeout: 0ns,
801 + }
802 + "#]]
803 + .assert_debug_eq(&config);
804 +}
805 +
806 +#[test]
807 +fn frontend_args_reject_legacy_handshake_flags() {
808 + let error = Cli::try_parse_from([
809 + "vllm-rs",
810 + "frontend",
811 + "--listen-fd",
812 + "3",
813 + "--input-address",
814 + "ipc:///tmp/input.sock",
815 + "--output-address",
816 + "ipc:///tmp/output.sock",
817 + "--args-json",
818 + r#"{"model_tag":"Qwen/Qwen3-0.6B"}"#,
819 + "--handshake-address",
820 + "tcp://127.0.0.1:62100",
821 + ])
822 + .unwrap_err();
823 +
824 + assert!(error.to_string().contains("--handshake-address"));
825 +}
826 +
827 +#[test]
828 +fn frontend_config_uses_external_coordinator_when_coordinator_address_is_present() {
829 + let cli = Cli::try_parse_from([
830 + "vllm-rs",
831 + "frontend",
832 + "--listen-fd",
833 + "3",
834 + "--input-address",
835 + "ipc:///tmp/input.sock",
836 + "--output-address",
837 + "ipc:///tmp/output.sock",
838 + "--coordinator-address",
839 + "tcp://127.0.0.1:7000",
840 + "--engine-count",
841 + "2",
842 + "--args-json",
843 + r#"{"model_tag":"Qwen/Qwen3-0.6B"}"#,
844 + ])
845 + .unwrap();
846 +
847 + let Command::Frontend(args) = cli.command else {
848 + panic!("expected frontend args");
849 + };
850 + let config = args.into_config();
851 +
852 + expect![[r#"
853 + Config {
854 + transport_mode: Bootstrapped {
855 + input_address: "ipc:///tmp/input.sock",
856 + output_address: "ipc:///tmp/output.sock",
857 + engine_count: 2,
858 + ready_timeout: 600s,
859 + },
860 + coordinator_mode: External {
861 + address: "tcp://127.0.0.1:7000",
862 + },
863 + model: "Qwen/Qwen3-0.6B",
864 + served_model_name: [],
865 + listener_mode: InheritedFd {
866 + fd: 3,
867 + },
868 + tool_call_parser: Auto,
869 + reasoning_parser: Auto,
870 + renderer: Auto,
871 + chat_template: None,
872 + default_chat_template_kwargs: None,
873 + chat_template_content_format: Auto,
874 + enable_log_requests: false,
875 + disable_log_stats: false,
876 + grpc_port: None,
877 + shutdown_timeout: 0ns,
878 + }
879 + "#]]
880 + .assert_debug_eq(&config);
881 +}
882 +
883 +#[test]
884 +fn serve_frontend_config_uses_unix_listener_when_uds_is_present() {
885 + let cli = Cli::try_parse_from([
886 + "vllm-rs",
887 + "serve",
888 + "Qwen/Qwen3-0.6B",
889 + "--uds",
890 + "/tmp/vllm.sock",
891 + ])
892 + .unwrap();
893 +
894 + let Command::Serve(args) = cli.command else {
895 + panic!("expected serve args");
896 + };
897 + let config = args.to_frontend_config("tcp://127.0.0.1:29550".to_string());
898 +
899 + assert_eq!(
900 + config.listener_mode,
901 + HttpListenerMode::BindUnix {
902 + path: "/tmp/vllm.sock".to_string(),
903 + }
904 + );
905 +}
rust/src/cmd/src/cli/unsupported.rs new
+660
@@ -0,0 +1,660 @@
1 +#![allow(clippy::doc_lazy_continuation)]
2 +
3 +use std::fmt::Display;
4 +use std::str::FromStr;
5 +
6 +use clap::Args;
7 +use clap::builder::{TypedValueParser, ValueParserFactory};
8 +use itertools::Itertools;
9 +use serde::{Deserialize, Deserializer, Serialize};
10 +
11 +/// Marker type for frontend-owned `serve` arguments that `vllm-rs` recognizes
12 +/// but does not support yet.
13 +///
14 +/// When passed as JSON args, it can be deserialized from any value, and
15 +/// serializes back to the original value.
16 +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
17 +#[serde(transparent)]
18 +pub struct Unsupported(pub serde_json::Value);
19 +
20 +impl FromStr for Unsupported {
21 + type Err = String;
22 +
23 + fn from_str(_s: &str) -> Result<Self, Self::Err> {
24 + Err("argument is not implemented in Rust frontend yet
25 +
26 +Remove this unsupported argument to continue.
27 +
28 +Alternatively, if you intend to pass it only to the Python engine, put it after `--` (e.g., `-- <arg>`).
29 +This may lead to unexpected behavior as the Rust frontend will completely ignore that argument."
30 + .to_string())
31 + }
32 +}
33 +
34 +/// Marker type for no-op arguments that are accepted by the Rust frontend but
35 +/// have no effect.
36 +///
37 +/// When passed as JSON args, it can be deserialized from any value, but always
38 +/// serializes back to `null`.
39 +#[derive(Clone, Debug, PartialEq, Eq, Serialize)]
40 +pub struct Noop;
41 +
42 +impl<'de> Deserialize<'de> for Noop {
43 + fn deserialize<D>(_deserializer: D) -> Result<Self, D::Error>
44 + where
45 + D: Deserializer<'de>,
46 + {
47 + Ok(Noop)
48 + }
49 +}
50 +
51 +impl ValueParserFactory for Noop {
52 + type Parser = NoopValueParser;
53 +
54 + fn value_parser() -> Self::Parser {
55 + NoopValueParser
56 + }
57 +}
58 +
59 +#[derive(Copy, Clone, Debug)]
60 +pub struct NoopValueParser;
61 +
62 +#[track_caller]
63 +fn noop_warn(arg: impl Display) {
64 + tracing::warn!("argument '{arg}' currently has no effect in Rust frontend, ignoring");
65 +}
66 +
67 +impl TypedValueParser for NoopValueParser {
68 + type Value = Noop;
69 +
70 + fn parse_ref(
71 + &self,
72 + _cmd: &clap::Command,
73 + arg: Option<&clap::Arg>,
74 + _value: &std::ffi::OsStr,
75 + ) -> Result<Self::Value, clap::Error> {
76 + if let Some(arg) = arg {
77 + noop_warn(arg);
78 + }
79 + Ok(Noop)
80 + }
81 +}
82 +
83 +/// Frontend-owned Python `serve` arguments that `vllm-rs` recognizes but does
84 +/// not support yet.
85 +#[serde_with::skip_serializing_none]
86 +#[derive(Debug, Clone, PartialEq, Eq, Default, Args, Serialize, Deserialize)]
87 +#[command(next_help_heading = "Options not implemented in Rust frontend yet")]
88 +pub struct UnsupportedArgs {
89 + #[command(flatten)]
90 + #[serde(default, flatten)]
91 + top_level: TopLevelUnsupportedArgs,
92 + #[command(flatten)]
93 + #[serde(default, flatten)]
94 + engine: EngineUnsupportedArgs,
95 + #[command(flatten)]
96 + #[serde(default, flatten)]
97 + server: ServerUnsupportedArgs,
98 +}
99 +
100 +impl UnsupportedArgs {
101 + /// Check whether any unsupported arguments are set, and if so, return an
102 + /// error listing them. Also warn about any no-op arguments that are set
103 + /// but will be ignored.
104 + pub(crate) fn check(&self) -> Result<(), String> {
105 + let value = serde_json::to_value(self).unwrap();
106 + let map = value.as_object().unwrap();
107 + let mut unsupported = Vec::new();
108 +
109 + for (key, value) in map {
110 + if value.is_null() {
111 + noop_warn(key);
112 + } else {
113 + unsupported.push(key.as_str());
114 + }
115 + }
116 +
117 + if !unsupported.is_empty() {
118 + unsupported.sort_unstable();
119 + let bullets = unsupported.into_iter().map(|key| format!("- {key}")).join("\n");
120 + return Err(format!(
121 + "
122 +The following arguments are not implemented in Rust frontend yet:
123 +{bullets}
124 +
125 +Remove these arguments to continue."
126 + ));
127 + }
128 +
129 + Ok(())
130 + }
131 +}
132 +
133 +/// Frontend-owned Python `vllm serve` top-level arguments that `vllm-rs`
134 +/// recognizes but does not support yet.
135 +///
136 +/// Source of truth in Python vLLM:
137 +/// - `vllm.entrypoints.openai.cli_args.make_arg_parser(...)`
138 +/// - `vllm.entrypoints.cli.serve.ServeSubcommand.subparser_init(...)`
139 +///
140 +/// These are not part of `EngineArgs`, `AsyncEngineArgs`, `BaseFrontendArgs`,
141 +/// or `FrontendArgs`. They live on the `serve` command itself and control
142 +/// managed-engine / multi-process orchestration rather than the shared frontend
143 +/// runtime config.
144 +#[serde_with::skip_serializing_none]
145 +#[derive(Debug, Clone, PartialEq, Eq, Default, Args, Serialize, Deserialize)]
146 +pub struct TopLevelUnsupportedArgs {
147 + /// How many API server processes to run. Defaults to data_parallel_size if
148 + /// not specified.
149 + #[arg(long, hide = true)]
150 + pub api_server_count: Option<Noop>,
151 +
152 + /// Read CLI options from a config file. Must be a YAML with the following
153 + /// options: https://docs.vllm.ai/en/latest/configuration/serve_args.html
154 + #[arg(long)]
155 + pub config: Option<Unsupported>,
156 +
157 + /// Launch a gRPC server instead of the HTTP OpenAI-compatible server.
158 + /// Requires: pip install vllm[grpc].
159 + #[arg(long, default_missing_value = "true", num_args = 0..=1)]
160 + pub grpc: Option<Unsupported>,
161 +}
162 +
163 +/// Frontend-owned Python engine arguments that `vllm-rs` recognizes but does
164 +/// not support yet.
165 +///
166 +/// Source of truth in Python vLLM:
167 +/// - `vllm.engine.arg_utils.EngineArgs.add_cli_args(...)`
168 +/// - `vllm.engine.arg_utils.AsyncEngineArgs.add_cli_args(...)`
169 +///
170 +/// These arguments are declared through the Python engine-args surface, but
171 +/// they are still frontend-owned: the API server / AsyncLLM layer reads them
172 +/// for tokenizer setup, request validation, routing, logging, and other
173 +/// frontend behavior, so Rust must recognize them rather than treating them as
174 +/// pure engine passthrough.
175 +#[serde_with::skip_serializing_none]
176 +#[derive(Debug, Clone, PartialEq, Eq, Default, Args, Serialize, Deserialize)]
177 +pub struct EngineUnsupportedArgs {
178 + /// Name or path of the Hugging Face tokenizer to use. If unspecified, model
179 + /// name or path will be used.
180 + #[arg(long)]
181 + pub tokenizer: Option<Unsupported>,
182 +
183 + /// Name or path of the Hugging Face config to use. If unspecified, model
184 + /// name or path will be used.
185 + #[arg(long)]
186 + pub hf_config_path: Option<Unsupported>,
187 +
188 + /// Allowing API requests to read local images or videos from directories
189 + /// specified by the server file system. This is a security risk. Should
190 + /// only be enabled in trusted environments.
191 + #[arg(long)]
192 + pub allowed_local_media_path: Option<Unsupported>,
193 +
194 + /// If set, only media URLs that belong to this domain can be used for
195 + /// multi-modal inputs.
196 + #[arg(long)]
197 + pub allowed_media_domains: Option<Unsupported>,
198 +
199 + /// The specific revision to use for the tokenizer on the Hugging Face Hub.
200 + /// It can be a branch name, a tag name, or a commit id. If unspecified,
201 + /// will use the default version.
202 + #[arg(long)]
203 + pub tokenizer_revision: Option<Unsupported>,
204 +
205 + /// Maximum number of log probabilities to return when `logprobs` is
206 + /// specified in `SamplingParams`. The default value comes the default for
207 + /// the OpenAI Chat Completions API. -1 means no cap, i.e. all
208 + /// (output_length * vocab_size) logprobs are allowed to be returned and
209 + /// it may cause OOM.
210 + #[arg(long)]
211 + pub max_logprobs: Option<Unsupported>,
212 +
213 + /// Skip initialization of tokenizer and detokenizer. Expects valid
214 + /// `prompt_token_ids` and `None` for prompt from the input. The generated
215 + /// output will contain token ids.
216 + #[arg(
217 + long,
218 + visible_alias = "no-skip-tokenizer-init",
219 + default_missing_value = "true",
220 + num_args = 0..=1
221 + )]
222 + pub skip_tokenizer_init: Option<Unsupported>,
223 +
224 + /// If `True`, enables passing text embeddings as inputs via the
225 + /// `prompt_embeds` key.
226 + ///
227 + /// WARNING: The vLLM engine may crash if incorrect shape of embeddings is
228 + /// passed. Only enable this flag for trusted users!
229 + #[arg(
230 + long,
231 + visible_alias = "no-enable-prompt-embeds",
232 + default_missing_value = "true",
233 + num_args = 0..=1
234 + )]
235 + pub enable_prompt_embeds: Option<Unsupported>,
236 +
237 + /// The token to use as HTTP bearer authorization for remote files. If
238 + /// `True`, will use the token generated when running `hf auth login`
239 + /// (stored in `~/.cache/huggingface/token`).
240 + #[arg(long, default_missing_value = "true", num_args = 0..=1)]
241 + pub hf_token: Option<Unsupported>,
242 +
243 + /// If a dictionary, contains arguments to be forwarded to the Hugging Face
244 + /// config. If a callable, it is called to update the HuggingFace config.
245 + #[arg(long)]
246 + pub hf_overrides: Option<Unsupported>,
247 +
248 + /// The folder path to the generation config. Defaults to `"auto"`, the
249 + /// generation config will be loaded from model path. If set to `"vllm"`, no
250 + /// generation config is loaded, vLLM defaults will be used. If set to a
251 + /// folder path, the generation config will be loaded from the specified
252 + /// folder path. If `max_new_tokens` is specified in generation config,
253 + /// then it sets a server-wide limit on the number of output tokens for
254 + /// all requests.
255 + #[arg(long)]
256 + pub generation_config: Option<Unsupported>,
257 +
258 + /// IOProcessor plugin name to load at model startup
259 + #[arg(long)]
260 + pub io_processor_plugin: Option<Unsupported>,
261 +
262 + /// Path to a dynamically reasoning parser plugin that can be dynamically
263 + /// loaded and registered.
264 + #[arg(long)]
265 + pub reasoning_parser_plugin: Option<Unsupported>,
266 +
267 + /// Rank of the data parallel group.
268 + #[arg(long, env = "VLLM_DP_RANK")]
269 + pub data_parallel_rank: Option<Unsupported>,
270 +
271 + /// Whether to use "hybrid" DP LB mode. Applies only to online serving
272 + /// and when data_parallel_size > 0. Enables running an AsyncLLM
273 + /// and API server on a "per-node" basis where vLLM load balances
274 + /// between local data parallel ranks, but an external LB balances
275 + /// between vLLM nodes/replicas. Set explicitly in conjunction with
276 + /// --data-parallel-start-rank.
277 + #[arg(
278 + long,
279 + visible_alias = "no-data-parallel-hybrid-lb",
280 + default_missing_value = "true",
281 + num_args = 0..=1
282 + )]
283 + pub data_parallel_hybrid_lb: Option<Unsupported>,
284 +
285 + /// Whether to use "external" DP LB mode. Applies only to online serving
286 + /// and when data_parallel_size > 0. This is useful for a "one-pod-per-rank"
287 + /// wide-EP setup in Kubernetes. Set implicitly when --data-parallel-rank
288 + /// is provided explicitly to vllm serve.
289 + #[arg(
290 + long,
291 + visible_alias = "no-data-parallel-external-lb",
292 + default_missing_value = "true",
293 + num_args = 0..=1
294 + )]
295 + pub data_parallel_external_lb: Option<Unsupported>,
296 +
297 + /// This feature is work in progress and no prefill optimization takes place
298 + /// with this flag enabled currently.
299 + #[arg(
300 + long,
301 + visible_alias = "no-kv-sharing-fast-prefill",
302 + default_missing_value = "true",
303 + num_args = 0..=1
304 + )]
305 + pub kv_sharing_fast_prefill: Option<Unsupported>,
306 +
307 + /// The maximum number of input items and options allowed per
308 + /// prompt for each modality.
309 + #[arg(long)]
310 + pub limit_mm_per_prompt: Option<Unsupported>,
311 +
312 + /// Additional args passed to process media inputs, keyed by modalities.
313 + #[arg(long)]
314 + pub media_io_kwargs: Option<Unsupported>,
315 +
316 + /// Arguments to be forwarded to the model's processor for multi-modal data,
317 + /// e.g., image processor.
318 + #[arg(long)]
319 + pub mm_processor_kwargs: Option<Unsupported>,
320 +
321 + /// The size (in GiB) of the multi-modal processor cache.
322 + #[arg(long)]
323 + pub mm_processor_cache_gb: Option<Unsupported>,
324 +
325 + /// Type of cache to use for the multi-modal preprocessor/mapper.
326 + #[arg(long)]
327 + pub mm_processor_cache_type: Option<Unsupported>,
328 +
329 + /// If True, enable handling of LoRA adapters.
330 + #[arg(
331 + long,
332 + visible_alias = "no-enable-lora",
333 + default_missing_value = "true",
334 + num_args = 0..=1
335 + )]
336 + pub enable_lora: Option<Unsupported>,
337 +
338 + /// Dictionary mapping specific modalities to LoRA model paths.
339 + #[arg(long)]
340 + pub default_mm_loras: Option<Unsupported>,
341 +
342 + /// Target URL to which OpenTelemetry traces will be sent.
343 + #[arg(long)]
344 + pub otlp_traces_endpoint: Option<Unsupported>,
345 +
346 + /// It makes sense to set this only if `--otlp-traces-endpoint` is set.
347 + #[arg(long)]
348 + pub collect_detailed_traces: Option<Unsupported>,
349 +
350 + /// The interval (or buffer size) for streaming in terms of token length.
351 + #[arg(long)]
352 + pub stream_interval: Option<Unsupported>,
353 +
354 + /// Structured outputs configuration.
355 + #[arg(long)]
356 + pub structured_outputs_config: Option<Noop>,
357 +
358 + /// Log aggregate rather than per-engine statistics when using data
359 + /// parallelism.
360 + #[arg(long, default_missing_value = "true", num_args = 0..=1)]
361 + pub aggregate_engine_logging: Option<Unsupported>,
362 +}
363 +
364 +/// Frontend-owned Python OpenAI server arguments that `vllm-rs` recognizes but
365 +/// does not support yet.
366 +///
367 +/// Source of truth in Python vLLM:
368 +/// - `vllm.entrypoints.openai.cli_args.BaseFrontendArgs`
369 +/// - `vllm.entrypoints.openai.cli_args.FrontendArgs`
370 +///
371 +/// These are not engine args. They belong to the Python OpenAI-compatible
372 +/// frontend / API-server layer itself, for example chat-template configuration,
373 +/// tool/frontend behavior, TLS / CORS / HTTP server settings, and other
374 +/// northbound server knobs.
375 +#[serde_with::skip_serializing_none]
376 +#[derive(Debug, Clone, PartialEq, Eq, Default, Args, Serialize, Deserialize)]
377 +pub struct ServerUnsupportedArgs {
378 + /// LoRA modules configurations in either 'name=path' format or JSON format
379 + /// or JSON list format. Example (old format): `'name=path'` Example (new
380 + /// format): `{"name": "name", "path": "lora_path",
381 + /// "base_model_name": "id"}`
382 + #[arg(long)]
383 + pub lora_modules: Option<Unsupported>,
384 +
385 + /// Whether to trust the chat template provided in the request. If False,
386 + /// the server will always use the chat template specified by
387 + /// `--chat-template` or the ones from tokenizer.
388 + #[arg(
389 + long,
390 + visible_alias = "no-trust-request-chat-template",
391 + default_missing_value = "true",
392 + num_args = 0..=1
393 + )]
394 + pub trust_request_chat_template: Option<Unsupported>,
395 +
396 + /// The role name to return if `request.add_generation_prompt=true`.
397 + #[arg(long)]
398 + pub response_role: Option<Unsupported>,
399 +
400 + /// When `--max-logprobs` is specified, represents single tokens as
401 + /// strings of the form 'token_id:{token_id}' so that tokens that are not
402 + /// JSON-encodable can be identified.
403 + #[arg(
404 + long,
405 + visible_alias = "no-return-tokens-as-token-ids",
406 + default_missing_value = "true",
407 + num_args = 0..=1
408 + )]
409 + pub return_tokens_as_token_ids: Option<Unsupported>,
410 +
411 + /// Enable auto tool choice for supported models. Use `--tool-call-parser`
412 + /// to specify which parser to use.
413 + #[arg(
414 + long,
415 + visible_alias = "no-enable-auto-tool-choice",
416 + default_missing_value = "true",
417 + num_args = 0..=1,
418 + hide = true
419 + )]
420 + pub enable_auto_tool_choice: Option<Noop>,
421 +
422 + /// If specified, exclude tool definitions in prompts when
423 + /// tool_choice='none'.
424 + #[arg(
425 + long,
426 + visible_alias = "no-exclude-tools-when-tool-choice-none",
427 + default_missing_value = "true",
428 + num_args = 0..=1
429 + )]
430 + pub exclude_tools_when_tool_choice_none: Option<Unsupported>,
431 +
432 + /// Special the tool parser plugin write to parse the model-generated tool
433 + /// into OpenAI API format, the name register in this plugin can be used in
434 + /// `--tool-call-parser`.
435 + #[arg(long)]
436 + pub tool_parser_plugin: Option<Unsupported>,
437 +
438 + /// Comma-separated list of host:port pairs (IPv4, IPv6, or hostname).
439 + /// Examples: 127.0.0.1:8000, [::1]:8000, localhost:1234. Or `demo` for
440 + /// built-in demo tools (browser and Python code interpreter). WARNING:
441 + /// The `demo` Python tool executes model-generated code in Docker without
442 + /// network isolation by default. See the security guide for more
443 + /// information.
444 + #[arg(long)]
445 + pub tool_server: Option<Unsupported>,
446 +
447 + /// Path to logging config JSON file for both vllm and uvicorn
448 + #[arg(long, /* env = "VLLM_LOGGING_CONFIG_PATH" */)]
449 + pub log_config_file: Option<Unsupported>,
450 +
451 + /// Max number of prompt characters or prompt ID numbers being printed in
452 + /// log. The default of None means unlimited.
453 + #[arg(long)]
454 + pub max_log_len: Option<Unsupported>,
455 +
456 + /// If set to True, enable prompt_tokens_details in usage.
457 + #[arg(
458 + long,
459 + visible_alias = "no-enable-prompt-tokens-details",
460 + default_missing_value = "true",
461 + num_args = 0..=1
462 + )]
463 + pub enable_prompt_tokens_details: Option<Unsupported>,
464 +
465 + /// If set to True, enable tracking server_load_metrics in the app state.
466 + #[arg(
467 + long,
468 + visible_alias = "no-enable-server-load-tracking",
469 + default_missing_value = "true",
470 + num_args = 0..=1
471 + )]
472 + pub enable_server_load_tracking: Option<Noop>,
473 +
474 + /// If set to True, including usage on every request.
475 + #[arg(
476 + long,
477 + visible_alias = "no-enable-force-include-usage",
478 + default_missing_value = "true",
479 + num_args = 0..=1
480 + )]
481 + pub enable_force_include_usage: Option<Unsupported>,
482 +
483 + /// Enable the `/tokenizer_info` endpoint. May expose chat
484 + /// templates and other tokenizer configuration.
485 + #[arg(
486 + long,
487 + visible_alias = "no-enable-tokenizer-info-endpoint",
488 + default_missing_value = "true",
489 + num_args = 0..=1
490 + )]
491 + pub enable_tokenizer_info_endpoint: Option<Unsupported>,
492 +
493 + /// If set to True, log model outputs (generations).
494 + /// Requires `--enable-log-requests`. As with `--enable-log-requests`,
495 + /// information is only logged at INFO level at maximum.
496 + #[arg(
497 + long,
498 + visible_alias = "no-enable-log-outputs",
499 + default_missing_value = "true",
500 + num_args = 0..=1
501 + )]
502 + pub enable_log_outputs: Option<Unsupported>,
503 +
504 + /// If set to False, output deltas will not be logged. Relevant only if
505 + /// --enable-log-outputs is set.
506 + #[arg(
507 + long,
508 + visible_alias = "no-enable-log-deltas",
509 + default_missing_value = "true",
510 + num_args = 0..=1
511 + )]
512 + pub enable_log_deltas: Option<Unsupported>,
513 +
514 + /// If set to True, log the stack trace of error responses
515 + #[arg(
516 + long,
517 + // env = "VLLM_SERVER_DEV_MODE",
518 + visible_alias = "no-log-error-stack",
519 + default_missing_value = "true",
520 + num_args = 0..=1
521 + )]
522 + pub log_error_stack: Option<Unsupported>,
523 +
524 + /// If set to True, only enable the Tokens In<>Out endpoint.
525 + /// This is intended for use in a Disaggregated Everything setup.
526 + #[arg(
527 + long,
528 + visible_alias = "no-tokens-only",
529 + default_missing_value = "true",
530 + num_args = 0..=1
531 + )]
532 + pub tokens_only: Option<Unsupported>,
533 +
534 + /// Log level for uvicorn.
535 + #[arg(long)]
536 + pub uvicorn_log_level: Option<Unsupported>,
537 +
538 + /// Disable uvicorn access log.
539 + #[arg(
540 + long,
541 + visible_alias = "no-disable-uvicorn-access-log",
542 + default_missing_value = "true",
543 + num_args = 0..=1
544 + )]
545 + pub disable_uvicorn_access_log: Option<Noop>,
546 +
547 + /// Comma-separated list of endpoint paths to exclude from uvicorn access
548 + /// logs. This is useful to reduce log noise from high-frequency endpoints
549 + /// like health checks. Example: "/health,/metrics,/ping".
550 + /// When set, access logs for requests to these paths will be suppressed
551 + /// while keeping logs for other endpoints.
552 + #[arg(long)]
553 + pub disable_access_log_for_endpoints: Option<Noop>,
554 +
555 + /// Allow credentials.
556 + #[arg(
557 + long,
558 + visible_alias = "no-allow-credentials",
559 + default_missing_value = "true",
560 + num_args = 0..=1
561 + )]
562 + pub allow_credentials: Option<Unsupported>,
563 +
564 + /// Allowed origins.
565 + #[arg(long)]
566 + pub allowed_origins: Option<Unsupported>,
567 +
568 + /// Allowed methods.
569 + #[arg(long)]
570 + pub allowed_methods: Option<Unsupported>,
571 +
572 + /// Allowed headers.
573 + #[arg(long)]
574 + pub allowed_headers: Option<Unsupported>,
575 +
576 + /// If provided, the server will require one of these keys to be presented
577 + /// in the header.
578 + #[arg(long)]
579 + pub api_key: Option<Unsupported>,
580 +
581 + /// The file path to the SSL key file.
582 + #[arg(long)]
583 + pub ssl_keyfile: Option<Unsupported>,
584 +
585 + /// The file path to the SSL cert file.
586 + #[arg(long)]
587 + pub ssl_certfile: Option<Unsupported>,
588 +
589 + /// The CA certificates file.
590 + #[arg(long)]
591 + pub ssl_ca_certs: Option<Unsupported>,
592 +
593 + /// Refresh SSL Context when SSL certificate files change
594 + #[arg(
595 + long,
596 + visible_alias = "no-enable-ssl-refresh",
597 + default_missing_value = "true",
598 + num_args = 0..=1
599 + )]
600 + pub enable_ssl_refresh: Option<Unsupported>,
601 +
602 + /// Whether client certificate is required (see stdlib ssl module's).
603 + #[arg(long)]
604 + pub ssl_cert_reqs: Option<Unsupported>,
605 +
606 + /// SSL cipher suites for HTTPS (TLS 1.2 and below only).
607 + /// Example: 'ECDHE-RSA-AES256-GCM-SHA384:ECDHE-RSA-CHACHA20-POLY1305'
608 + #[arg(long)]
609 + pub ssl_ciphers: Option<Unsupported>,
610 +
611 + /// FastAPI root_path when app is behind a path based routing proxy.
612 + #[arg(long)]
613 + pub root_path: Option<Unsupported>,
614 +
615 + /// Additional ASGI middleware to apply to the app. We accept multiple
616 + /// --middleware arguments. The value should be an import path. If a
617 + /// function is provided, vLLM will add it to the server using
618 + /// `@app.middleware('http')`. If a class is provided, vLLM will
619 + /// add it to the server using `app.add_middleware()`.
620 + #[arg(long)]
621 + pub middleware: Option<Unsupported>,
622 +
623 + /// If specified, API server will add X-Request-Id header to responses.
624 + #[arg(
625 + long,
626 + visible_alias = "no-enable-request-id-headers",
627 + default_missing_value = "true",
628 + num_args = 0..=1
629 + )]
630 + pub enable_request_id_headers: Option<Unsupported>,
631 +
632 + /// Disable FastAPI's OpenAPI schema, Swagger UI, and ReDoc endpoint.
633 + #[arg(
634 + long,
635 + visible_alias = "no-disable-fastapi-docs",
636 + default_missing_value = "true",
637 + num_args = 0..=1
638 + )]
639 + pub disable_fastapi_docs: Option<Unsupported>,
640 +
641 + /// Maximum size (bytes) of an incomplete HTTP event (header or body) for
642 + /// h11 parser. Helps mitigate header abuse. Default: 4194304 (4 MB).
643 + #[arg(long)]
644 + pub h11_max_incomplete_event_size: Option<Unsupported>,
645 +
646 + /// Maximum number of HTTP headers allowed in a request for h11 parser.
647 + /// Helps mitigate header abuse. Default: 256.
648 + #[arg(long)]
649 + pub h11_max_header_count: Option<Unsupported>,
650 +
651 + /// Enable offline FastAPI documentation for air-gapped environments.
652 + /// Uses vendored static assets bundled with vLLM.
653 + #[arg(
654 + long,
655 + visible_alias = "no-enable-offline-docs",
656 + default_missing_value = "true",
657 + num_args = 0..=1
658 + )]
659 + pub enable_offline_docs: Option<Unsupported>,
660 +}
rust/src/cmd/src/logging.rs new
+328
@@ -0,0 +1,328 @@
1 +use std::{env, fmt, process};
2 +
3 +use time::UtcOffset;
4 +use time::macros::format_description;
5 +use tracing::level_filters::LevelFilter;
6 +use tracing::{Event, Level, Subscriber};
7 +use tracing_subscriber::Layer as _;
8 +use tracing_subscriber::filter::Targets;
9 +use tracing_subscriber::fmt::format::{FormatEvent, FormatFields, Writer};
10 +use tracing_subscriber::fmt::time::FormatTime;
11 +use tracing_subscriber::fmt::{FmtContext, FormattedFields};
12 +use tracing_subscriber::layer::SubscriberExt as _;
13 +use tracing_subscriber::registry::LookupSpan;
14 +use tracing_subscriber::util::SubscriberInitExt as _;
15 +
16 +const CYAN: &str = "\x1b[0;36m";
17 +const GREY: &str = "\x1b[90m";
18 +const GREEN: &str = "\x1b[32m";
19 +const YELLOW: &str = "\x1b[33m";
20 +const RED: &str = "\x1b[31m";
21 +const WHITE: &str = "\x1b[37m";
22 +const RESET: &str = "\x1b[0m";
23 +const VLLM_TIME_FORMAT: &[time::format_description::FormatItem<'static>] =
24 + format_description!("[month]-[day] [hour]:[minute]:[second]");
25 +
26 +const PROCESS_LABEL: &str = "RustFrontend";
27 +
28 +/// Install the process-wide vLLM-style tracing subscriber for the CLI binary.
29 +pub(crate) fn init_tracing() {
30 + let filter = build_targets_filter(
31 + env::var("VLLM_LOGGING_LEVEL").ok().as_deref(),
32 + env::var("RUST_LOG").ok().as_deref(),
33 + );
34 + let formatter = VllmEventFormatter::new();
35 +
36 + let _ = tracing_subscriber::registry()
37 + .with(tracing_subscriber::fmt::layer().event_format(formatter).with_filter(filter))
38 + .try_init();
39 +}
40 +
41 +/// Build the CLI log filter by merging the vLLM-style default level with
42 +/// Rust-style target overrides.
43 +///
44 +/// Precedence:
45 +/// - Start from `VLLM_LOGGING_LEVEL` as the default level for all targets.
46 +/// - If `RUST_LOG` contains a global default level such as `warn`, it overrides
47 +/// `VLLM_LOGGING_LEVEL`.
48 +/// - Any explicit target directives in `RUST_LOG`, such as `hyper=info`, override whichever default
49 +/// level is active for those targets only.
50 +fn build_targets_filter(vllm_logging_level: Option<&str>, rust_log: Option<&str>) -> Targets {
51 + let mut filter =
52 + Targets::new().with_default(map_python_log_level(vllm_logging_level.unwrap_or("INFO")));
53 +
54 + if let Some(rust_log) = rust_log
55 + && !rust_log.is_empty()
56 + {
57 + let rust_log_targets: Targets = rust_log.parse().expect("failed to parse `RUST_LOG`");
58 + if let Some(default_level) = rust_log_targets.default_level() {
59 + filter = filter.with_default(default_level);
60 + }
61 + filter = filter.with_targets(rust_log_targets);
62 + }
63 +
64 + filter
65 +}
66 +
67 +#[derive(Debug, Clone, Copy)]
68 +struct VllmLocalTimer {
69 + local_offset: UtcOffset,
70 +}
71 +
72 +impl Default for VllmLocalTimer {
73 + fn default() -> Self {
74 + let local_offset = UtcOffset::current_local_offset().unwrap_or(UtcOffset::UTC);
75 + Self { local_offset }
76 + }
77 +}
78 +
79 +impl FormatTime for VllmLocalTimer {
80 + fn format_time(&self, w: &mut Writer<'_>) -> fmt::Result {
81 + let now = time::OffsetDateTime::now_utc().to_offset(self.local_offset);
82 + let formatted = now.format(VLLM_TIME_FORMAT).map_err(|_| fmt::Error)?;
83 + w.write_str(&formatted)
84 + }
85 +}
86 +
87 +#[derive(Debug, Clone)]
88 +struct VllmEventFormatter {
89 + prefix: String,
90 + timer: VllmLocalTimer,
91 +}
92 +
93 +impl VllmEventFormatter {
94 + fn new() -> Self {
95 + Self {
96 + prefix: format!("({} pid={})", PROCESS_LABEL, process::id()),
97 + timer: VllmLocalTimer::default(),
98 + }
99 + }
100 +
101 + fn write_process_prefix(&self, writer: &mut Writer<'_>, ansi: bool) -> fmt::Result {
102 + write_colored(writer, ansi, Some(CYAN), &self.prefix)?;
103 + writer.write_char(' ')
104 + }
105 +
106 + fn write_level(&self, writer: &mut Writer<'_>, level: &Level, ansi: bool) -> fmt::Result {
107 + let (text, color) = match *level {
108 + Level::TRACE => ("TRACE", WHITE),
109 + Level::DEBUG => ("DEBUG", WHITE),
110 + Level::INFO => ("INFO", GREEN),
111 + Level::WARN => ("WARNING", YELLOW),
112 + Level::ERROR => ("ERROR", RED),
113 + };
114 + write_colored(writer, ansi, Some(color), text)
115 + }
116 +
117 + fn write_timestamp(&self, writer: &mut Writer<'_>, ansi: bool) -> fmt::Result {
118 + if ansi {
119 + writer.write_str(GREY)?;
120 + }
121 + if self.timer.format_time(writer).is_err() {
122 + writer.write_str("<unknown time>")?;
123 + }
124 + if ansi {
125 + writer.write_str(RESET)?;
126 + }
127 + Ok(())
128 + }
129 +
130 + fn write_location(
131 + &self,
132 + writer: &mut Writer<'_>,
133 + file: Option<&str>,
134 + line: Option<u32>,
135 + full_path: bool,
136 + ansi: bool,
137 + ) -> fmt::Result {
138 + let Some(file) = file else {
139 + return Ok(());
140 + };
141 + let file = if full_path {
142 + file
143 + } else {
144 + shorten_file_path(file)
145 + };
146 + if ansi {
147 + writer.write_str(GREY)?;
148 + }
149 + match line {
150 + Some(line) => write!(writer, "[{file}:{line}]")?,
151 + None => write!(writer, "[{file}]")?,
152 + }
153 + if ansi {
154 + writer.write_str(RESET)?;
155 + }
156 + Ok(())
157 + }
158 +
159 + fn write_scope<S, N>(&self, ctx: &FmtContext<'_, S, N>, writer: &mut Writer<'_>) -> fmt::Result
160 + where
161 + S: Subscriber + for<'lookup> LookupSpan<'lookup>,
162 + N: for<'writer> FormatFields<'writer> + 'static,
163 + {
164 + let Some(scope) = ctx.event_scope() else {
165 + return Ok(());
166 + };
167 +
168 + let mut seen = false;
169 + for span in scope.from_root() {
170 + if seen {
171 + writer.write_str(":")?;
172 + }
173 + seen = true;
174 + writer.write_str(span.metadata().name())?;
175 +
176 + let ext = span.extensions();
177 + if let Some(fields) = ext.get::<FormattedFields<N>>()
178 + && !fields.is_empty()
179 + {
180 + write!(writer, "{{{fields}}}")?;
181 + }
182 + }
183 +
184 + if seen {
185 + writer.write_str(": ")?;
186 + }
187 +
188 + Ok(())
189 + }
190 +}
191 +
192 +impl<S, N> FormatEvent<S, N> for VllmEventFormatter
193 +where
194 + S: Subscriber + for<'lookup> LookupSpan<'lookup>,
195 + N: for<'writer> FormatFields<'writer> + 'static,
196 +{
197 + fn format_event(
198 + &self,
199 + ctx: &FmtContext<'_, S, N>,
200 + mut writer: Writer<'_>,
201 + event: &Event<'_>,
202 + ) -> fmt::Result {
203 + let meta = event.metadata();
204 + let ansi = writer.has_ansi_escapes();
205 +
206 + self.write_process_prefix(&mut writer, ansi)?;
207 + self.write_level(&mut writer, meta.level(), ansi)?;
208 + writer.write_char(' ')?;
209 + self.write_timestamp(&mut writer, ansi)?;
210 + writer.write_char(' ')?;
211 + // Use the full file path only when DEBUG (or more verbose) is enabled anywhere,
212 + // independent of the level of this particular event. Filenames alone are often
213 + // ambiguous, but full paths are too noisy for normal INFO-level operation.
214 + let full_path = LevelFilter::current() >= LevelFilter::DEBUG;
215 + self.write_location(&mut writer, meta.file(), meta.line(), full_path, ansi)?;
216 + writer.write_char(' ')?;
217 + self.write_scope(ctx, &mut writer)?;
218 + ctx.format_fields(writer.by_ref(), event)?;
219 + writer.write_char('\n')
220 + }
221 +}
222 +
223 +/// Shorten a source file path for log output while preserving enough context
224 +/// for common Rust entrypoint and module filenames.
225 +///
226 +/// - For `mod.rs`, keep the parent directory as `parent/mod.rs`.
227 +/// - For `src/lib.rs` and `src/main.rs`, keep one additional component as `crate/src/lib.rs` or
228 +/// `crate/src/main.rs` when available.
229 +/// - Other files are displayed as just the basename.
230 +fn shorten_file_path(file: &str) -> &str {
231 + let mut parts = file.rsplit('/');
232 + let name = parts.next().unwrap_or(file);
233 + let parent = parts.next();
234 + let grandparent = parts.next();
235 +
236 + let Some(parent) = parent else {
237 + return file;
238 + };
239 +
240 + if name == "mod.rs" {
241 + return &file[file.len() - parent.len() - 1 - name.len()..];
242 + }
243 +
244 + if !matches!(name, "lib.rs" | "main.rs") || parent != "src" {
245 + return name;
246 + }
247 + let Some(grandparent) = grandparent else {
248 + return file;
249 + };
250 +
251 + &file[file.len() - grandparent.len() - 1 - parent.len() - 1 - name.len()..]
252 +}
253 +
254 +fn write_colored(
255 + writer: &mut Writer<'_>,
256 + ansi: bool,
257 + color: Option<&str>,
258 + text: &str,
259 +) -> fmt::Result {
260 + if ansi {
261 + if let Some(color) = color {
262 + writer.write_str(color)?;
263 + }
264 + writer.write_str(text)?;
265 + if color.is_some() {
266 + writer.write_str(RESET)?;
267 + }
268 + return Ok(());
269 + }
270 +
271 + writer.write_str(text)
272 +}
273 +
274 +/// Map a Python logging level name to the corresponding Rust tracing level.
275 +fn map_python_log_level(level: &str) -> LevelFilter {
276 + match level.to_ascii_uppercase().as_str() {
277 + "CRITICAL" | "FATAL" => LevelFilter::ERROR,
278 + "ERROR" => LevelFilter::ERROR,
279 + "WARNING" | "WARN" => LevelFilter::WARN,
280 + "INFO" => LevelFilter::INFO,
281 + "DEBUG" => LevelFilter::DEBUG,
282 + "NOTSET" => LevelFilter::TRACE,
283 + _ => LevelFilter::INFO,
284 + }
285 +}
286 +
287 +#[cfg(test)]
288 +mod tests {
289 + use super::*;
290 +
291 + #[test]
292 + fn rust_log_target_overrides_are_merged_with_vllm_default_level() {
293 + let filter = build_targets_filter(Some("DEBUG"), Some("hyper=warn,tower=error"));
294 +
295 + assert_eq!(filter.to_string(), "tower=error,hyper=warn,debug");
296 + }
297 +
298 + #[test]
299 + fn rust_log_default_level_overrides_vllm_default_level() {
300 + let filter = build_targets_filter(Some("DEBUG"), Some("warn,hyper=info"));
301 +
302 + assert_eq!(filter.to_string(), "hyper=info,warn");
303 + }
304 +
305 + #[test]
306 + fn invalid_vllm_level_falls_back_to_info() {
307 + let filter = build_targets_filter(Some("bogus"), None);
308 +
309 + assert_eq!(filter.to_string(), "info");
310 + }
311 +
312 + #[test]
313 + fn location_path_uses_filename_for_non_ambiguous_files() {
314 + assert_eq!(shorten_file_path("src/cmd/src/logging.rs"), "logging.rs");
315 + assert_eq!(shorten_file_path("src/chat/lib.rs"), "lib.rs");
316 + assert_eq!(shorten_file_path("src/chat/main.rs"), "main.rs");
317 + assert_eq!(shorten_file_path("src/chat/src/xmod.rs"), "xmod.rs");
318 + }
319 +
320 + #[test]
321 + fn location_path_keeps_more_context_for_common_entrypoint_filenames() {
322 + assert_eq!(shorten_file_path("src/lib.rs"), "src/lib.rs");
323 + assert_eq!(shorten_file_path("src/chat/src/lib.rs"), "chat/src/lib.rs");
324 + assert_eq!(shorten_file_path("src/cmd/src/main.rs"), "cmd/src/main.rs");
325 + assert_eq!(shorten_file_path("mod.rs"), "mod.rs");
326 + assert_eq!(shorten_file_path("src/chat/src/tool/mod.rs"), "tool/mod.rs");
327 + }
328 +}
rust/src/cmd/src/main.rs new
+189
@@ -0,0 +1,189 @@
1 +mod cli;
2 +mod logging;
3 +
4 +use std::env;
5 +use std::process::ExitStatus;
6 +
7 +use anyhow::{Context, Result, anyhow, bail};
8 +use tokio_util::sync::CancellationToken;
9 +use tracing::{info, warn};
10 +use vllm_managed_engine::ManagedEngineHandle;
11 +
12 +use crate::cli::{Cli, Command};
13 +
14 +const TOKIO_WORKER_THREADS_ENV: &str = "TOKIO_WORKER_THREADS";
15 +const DEFAULT_MAX_TOKIO_WORKER_THREADS: usize = 32;
16 +
17 +/// Cap the default number of Tokio worker threads if the user did not
18 +/// explicitly set `TOKIO_WORKER_THREADS` to avoid spawning too many threads on
19 +/// machines with a large number of CPUs, which may lead to excessive context
20 +/// switching and degraded performance.
21 +fn tokio_worker_threads() -> Option<usize> {
22 + if env::var_os(TOKIO_WORKER_THREADS_ENV).is_some() {
23 + return None;
24 + }
25 +
26 + std::thread::available_parallelism()
27 + .map(|parallelism| {
28 + let available = parallelism.get();
29 + let worker_threads = available.min(DEFAULT_MAX_TOKIO_WORKER_THREADS);
30 + if worker_threads < available {
31 + info!(
32 + available_parallelism = available,
33 + capped_worker_threads = worker_threads,
34 + "capping tokio worker threads, set {TOKIO_WORKER_THREADS_ENV} to override"
35 + );
36 + }
37 + worker_threads
38 + })
39 + .ok()
40 +}
41 +
42 +/// Reason that caused a managed `serve` session to stop.
43 +#[derive(Debug)]
44 +enum ShutdownReason {
45 + Signal,
46 + Server(anyhow::Error),
47 + EngineExited(ExitStatus),
48 +}
49 +
50 +/// Cancellation token tripped by Ctrl-C or SIGTERM.
51 +fn shutdown_signal() -> CancellationToken {
52 + let token = CancellationToken::new();
53 + let shutdown = token.clone();
54 +
55 + tokio::spawn(async move {
56 + let ctrl_c = async {
57 + tokio::signal::ctrl_c().await.expect("failed to install Ctrl-C signal handler");
58 + };
59 +
60 + let sigterm = async {
61 + tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate())
62 + .expect("failed to install SIGTERM signal handler")
63 + .recv()
64 + .await;
65 + };
66 +
67 + tokio::select! {
68 + _ = ctrl_c => info!("received shutdown signal (Ctrl-C), shutting down..."),
69 + _ = sigterm => info!("received shutdown signal (SIGTERM), shutting down..."),
70 + }
71 +
72 + shutdown.cancel();
73 + });
74 +
75 + token
76 +}
77 +
78 +fn main() -> Result<()> {
79 + logging::init_tracing();
80 + let cli = Cli::parse();
81 +
82 + let mut runtime = tokio::runtime::Builder::new_multi_thread();
83 + runtime.enable_all();
84 + if let Some(worker_threads) = tokio_worker_threads() {
85 + runtime.worker_threads(worker_threads);
86 + }
87 +
88 + runtime
89 + .build()
90 + .context("failed to build Tokio runtime")?
91 + .block_on(async_main(cli))
92 +}
93 +
94 +async fn async_main(cli: Cli) -> Result<()> {
95 + match cli.command {
96 + Command::Frontend(args) => vllm_server::serve(args.into_config(), shutdown_signal()).await,
97 + Command::Serve(args) => {
98 + let handshake_port = args.managed_engine.resolve_handshake_port()?;
99 +
100 + if args.managed_engine.data_parallel_size_local == Some(0) {
101 + if args.headless {
102 + bail!("cannot combine `--headless` with `--data-parallel-size-local 0`");
103 + }
104 +
105 + let handshake_address = args.managed_engine.handshake_address(handshake_port);
106 + info!(
107 + %handshake_address,
108 + engine_count = args.managed_engine.data_parallel_size,
109 + "running Rust frontend without a managed local Python engine"
110 + );
111 + let config = args.to_frontend_config(handshake_address);
112 + return vllm_server::serve(config, shutdown_signal()).await;
113 + }
114 +
115 + let shutdown_timeout = args.runtime.shutdown_timeout();
116 + let engine_config = args.to_managed_engine_config(handshake_port);
117 + let handshake_address = engine_config.handshake_address();
118 +
119 + let engine = ManagedEngineHandle::spawn(engine_config)
120 + .await
121 + .context("failed to start managed Python headless engine")?;
122 +
123 + let shutdown = shutdown_signal();
124 +
125 + let mut serve_task = if args.headless {
126 + info!("running managed Python headless engine without Rust frontend");
127 + let shutdown = shutdown.clone();
128 + tokio::spawn(async move {
129 + shutdown.cancelled().await;
130 + Ok(())
131 + })
132 + } else {
133 + let config = args.to_frontend_config(handshake_address);
134 + let shutdown = shutdown.clone();
135 + tokio::spawn(async move {
136 + let result = vllm_server::serve(config, shutdown).await;
137 + if result.is_ok() {
138 + info!("OpenAI server shut down gracefully");
139 + }
140 + result
141 + })
142 + };
143 +
144 + let shutdown_reason = tokio::select! {
145 + biased;
146 +
147 + // Received shutdown signal via Ctrl-C or SIGTERM.
148 + _ = shutdown.cancelled() => ShutdownReason::Signal,
149 +
150 + // Engine process exited unexpectedly.
151 + status = engine.wait_for_exit() => {
152 + warn!(%status, "managed Python headless engine exited, shutting down...");
153 + ShutdownReason::EngineExited(status)
154 + }
155 +
156 + // Serve task exited unexpectedly.
157 + serve_result = &mut serve_task => {
158 + let serve_result = serve_result.context("serve task join failed")?;
159 + match serve_result {
160 + Ok(()) => ShutdownReason::Server(anyhow!("OpenAI server shut down unexpectedly without error")),
161 + Err(error) => ShutdownReason::Server(error),
162 + }
163 + }
164 + };
165 + // Regardless of the shutdown reason, broadcast shutdown signal here to ensure
166 + // that all serving tasks are notified.
167 + shutdown.cancel();
168 +
169 + // Shutdown begins. Terminate the managed engine first.
170 + engine.shutdown(shutdown_timeout).await?;
171 + info!("managed engine shut down gracefully");
172 + // Wait for the API server to shut down gracefully by draining in-flight
173 + // requests.
174 + if !matches!(shutdown_reason, ShutdownReason::Server(_)) {
175 + serve_task.await.context("serve task join failed")??;
176 + }
177 +
178 + match shutdown_reason {
179 + ShutdownReason::Signal => Ok(()),
180 + ShutdownReason::Server(error) => {
181 + Err(error.context("OpenAI server shut down unexpectedly"))
182 + }
183 + ShutdownReason::EngineExited(status) => Err(anyhow!(
184 + "managed Python headless engine exited unexpectedly with status {status}"
185 + )),
186 + }
187 + }
188 + }
189 +}
rust/src/engine-core-client/Cargo.toml new
+49
@@ -0,0 +1,49 @@
1 +[package]
2 +name = "vllm-engine-core-client"
3 +version.workspace = true
4 +edition.workspace = true
5 +license.workspace = true
6 +
7 +[features]
8 +test-util = ["dep:tempfile"]
9 +
10 +[dependencies]
11 +arc-swap.workspace = true
12 +bytemuck.workspace = true
13 +byteorder.workspace = true
14 +bytes.workspace = true
15 +easy-ext.workspace = true
16 +enum-as-inner.workspace = true
17 +futures.workspace = true
18 +half.workspace = true
19 +hex.workspace = true
20 +itertools.workspace = true
21 +parking_lot.workspace = true
22 +rmp-serde.workspace = true
23 +rmpv.workspace = true
24 +serde.workspace = true
25 +serde_default.workspace = true
26 +serde_json.workspace = true
27 +serde_repr.workspace = true
28 +serde_tuple.workspace = true
29 +serde_with.workspace = true
30 +task-local.workspace = true
31 +tempfile = { workspace = true, optional = true }
32 +thiserror.workspace = true
33 +thiserror-ext.workspace = true
34 +tokio.workspace = true
35 +tokio-util.workspace = true
36 +tracing.workspace = true
37 +vllm-metrics.workspace = true
38 +zeromq.workspace = true
39 +
40 +[dev-dependencies]
41 +anyhow.workspace = true
42 +clap.workspace = true
43 +expect-test.workspace = true
44 +hex.workspace = true
45 +tempfile.workspace = true
46 +tracing-subscriber.workspace = true
47 +
48 +[lints]
49 +workspace = true
rust/src/engine-core-client/examples/README.md new
+53
@@ -0,0 +1,53 @@
1 +# Engine-Core Smoke Tests
2 +
3 +Start headless `vllm`:
4 +
5 +```bash
6 +source ../vllm/.venv/bin/activate
7 +HF_HUB_OFFLINE=1 \
8 +VLLM_LOGGING_LEVEL=DEBUG \
9 +VLLM_CPU_KVCACHE_SPACE=2 \
10 +VLLM_HOST_IP=127.0.0.1 \
11 +VLLM_LOOPBACK_IP=127.0.0.1 \
12 +python3 -m vllm.entrypoints.cli.main serve Qwen/Qwen3-0.6B \
13 + --headless \
14 + --enable-sleep-mode \
15 + --data-parallel-address 127.0.0.1 \
16 + --data-parallel-rpc-port 62100 \
17 + --data-parallel-size-local 1 \
18 + --max-model-len 512 \
19 + --dtype float16
20 +```
21 +
22 +Run the Rust smoke test through the `vllm-engine-core-client` utility interface:
23 +
24 +```bash
25 +cargo run -p vllm-engine-core-client --example external_engine_utility_call -- \
26 + --handshake-address tcp://127.0.0.1:62100 \
27 + --host 127.0.0.1
28 +```
29 +
30 +If your current engine setup does not support sleep mode, skip the `sleep` / `wake_up` part of the
31 +smoke:
32 +
33 +```bash
34 +cargo run -p vllm-engine-core-client --example external_engine_utility_call -- \
35 + --handshake-address tcp://127.0.0.1:62100 \
36 + --host 127.0.0.1 \
37 + --skip-sleep-wake
38 +```
39 +
40 +Run the Rust smoke test for sample logprobs decoding through the raw engine-core request path:
41 +
42 +```bash
43 +cargo run -p vllm-engine-core-client --example external_engine_logprobs -- \
44 + --handshake-address tcp://127.0.0.1:62100 \
45 + --host 127.0.0.1
46 +```
47 +
48 +This smoke requests a small generated-token `logprobs` payload plus prompt logprobs over a much
49 +longer prompt, so it exercises both the inline and aux-frame decode paths against a real engine.
50 +The Rust client decodes those payloads into semantic per-position records rather than exposing the
51 +raw ndarray/tensor wire shape.
52 +
53 +IMPORTANT: You must restart `vllm` each time you run the smoke test, as the vLLM engine cannot manage frontend closures and subsequent reconnects. In other words, do not reuse existing `vllm` instances, if any.
rust/src/engine-core-client/examples/external_engine_logprobs.rs new
+190
@@ -0,0 +1,190 @@
1 +use std::time::Duration;
2 +
3 +use anyhow::{Context, Result, bail};
4 +use clap::Parser;
5 +use futures::StreamExt as _;
6 +use tokio::time::timeout;
7 +use tracing_subscriber::EnvFilter;
8 +use vllm_engine_core_client::protocol::{
9 + EngineCoreFinishReason, EngineCoreRequest, EngineCoreSamplingParams,
10 +};
11 +use vllm_engine_core_client::{
12 + EngineCoreClient, EngineCoreClientConfig, EngineCoreStreamOutput, TransportMode,
13 +};
14 +
15 +const BASE_PROMPT_TOKEN_IDS: &[u32] = &[20841, 448, 6896, 25, 23811];
16 +
17 +#[derive(Debug, Parser)]
18 +#[command(about = "Smoke-test engine-core sample logprobs against an external vLLM engine.")]
19 +struct Args {
20 + #[arg(long)]
21 + handshake_address: String,
22 + #[arg(long, default_value_t = 1)]
23 + engine_count: usize,
24 + #[arg(long, default_value = "Qwen/Qwen3-0.6B")]
25 + model: String,
26 + #[arg(long, default_value = "127.0.0.1")]
27 + host: String,
28 + #[arg(long, default_value_t = 0)]
29 + client_index: u32,
30 + #[arg(long, default_value_t = 30)]
31 + ready_timeout_secs: u64,
32 + #[arg(long, default_value_t = 120)]
33 + output_timeout_secs: u64,
34 + #[arg(long, default_value_t = 1)]
35 + max_tokens: u32,
36 + #[arg(long, default_value_t = 2)]
37 + logprobs: i32,
38 + #[arg(long, default_value_t = 1)]
39 + prompt_logprobs: i32,
40 + #[arg(long, default_value_t = 96)]
41 + prompt_repeats: usize,
42 +}
43 +
44 +fn init_tracing() {
45 + let filter = EnvFilter::try_from_default_env()
46 + .unwrap_or_else(|_| EnvFilter::new("vllm_engine_core_client=debug"));
47 + let _ = tracing_subscriber::fmt().with_env_filter(filter).try_init();
48 +}
49 +
50 +fn unique_request_id() -> String {
51 + let nanos = std::time::SystemTime::now()
52 + .duration_since(std::time::UNIX_EPOCH)
53 + .expect("system clock should be after unix epoch")
54 + .as_nanos();
55 + format!("rust-engine-core-logprobs-{nanos}")
56 +}
57 +
58 +fn build_prompt_token_ids(prompt_repeats: usize) -> Vec<u32> {
59 + let repeats = prompt_repeats.max(1);
60 + BASE_PROMPT_TOKEN_IDS.repeat(repeats)
61 +}
62 +
63 +fn build_request(
64 + request_id: String,
65 + prompt_token_ids: Vec<u32>,
66 + max_tokens: u32,
67 + logprobs: i32,
68 + prompt_logprobs: i32,
69 + client_index: u32,
70 +) -> EngineCoreRequest {
71 + EngineCoreRequest {
72 + request_id,
73 + prompt_token_ids: Some(prompt_token_ids),
74 + sampling_params: Some(EngineCoreSamplingParams {
75 + max_tokens,
76 + logprobs: Some(logprobs),
77 + prompt_logprobs: Some(prompt_logprobs),
78 + ..EngineCoreSamplingParams::for_test()
79 + }),
80 + arrival_time: 0.0,
81 + client_index,
82 + ..EngineCoreRequest::default()
83 + }
84 +}
85 +
86 +async fn wait_for_final_output(
87 + mut stream: vllm_engine_core_client::EngineCoreOutputStream,
88 +) -> Result<EngineCoreStreamOutput> {
89 + while let Some(output) = stream.next().await {
90 + let output = output.context("failed to receive engine-core output")?;
91 + if output.finished() {
92 + return Ok(output);
93 + }
94 + }
95 + bail!("request stream ended without a final output")
96 +}
97 +
98 +#[tokio::main(flavor = "multi_thread")]
99 +async fn main() -> Result<()> {
100 + init_tracing();
101 + let args = Args::parse();
102 + let ready_timeout = Duration::from_secs(args.ready_timeout_secs);
103 + let output_timeout = Duration::from_secs(args.output_timeout_secs);
104 + let request_id = unique_request_id();
105 + let prompt_token_ids = build_prompt_token_ids(args.prompt_repeats);
106 + let client = EngineCoreClient::connect(EngineCoreClientConfig {
107 + transport_mode: TransportMode::HandshakeOwner {
108 + handshake_address: args.handshake_address.clone(),
109 + advertised_host: args.host.clone(),
110 + engine_count: args.engine_count,
111 + ready_timeout,
112 + local_input_address: None,
113 + local_output_address: None,
114 + },
115 + coordinator_mode: None,
116 + model_name: args.model.clone(),
117 + client_index: args.client_index,
118 + })
119 + .await
120 + .context("failed to connect to external vLLM engine")?;
121 +
122 + println!("model={}", args.model);
123 + println!("handshake_address={}", args.handshake_address);
124 + println!("engine_count={}", args.engine_count);
125 + println!("input_address={}", client.input_address());
126 + println!("output_address={}", client.output_address());
127 + println!("engine_identities={:x?}", client.engine_identities());
128 +
129 + let request = build_request(
130 + request_id.clone(),
131 + prompt_token_ids.clone(),
132 + args.max_tokens,
133 + args.logprobs,
134 + args.prompt_logprobs,
135 + args.client_index,
136 + );
137 + println!("request_id={request_id}");
138 + println!("prompt_len={}", prompt_token_ids.len());
139 + println!("base_prompt_len={}", BASE_PROMPT_TOKEN_IDS.len());
140 + println!("prompt_repeats={}", args.prompt_repeats);
141 + println!("requested_logprobs={}", args.logprobs);
142 + println!("requested_prompt_logprobs={}", args.prompt_logprobs);
143 +
144 + let stream = client.call(request).await.context("failed to submit engine-core request")?;
145 + let output = timeout(output_timeout, wait_for_final_output(stream))
146 + .await
147 + .context("timed out waiting for final output")??;
148 +
149 + let finish_reason = output.finish_reason;
150 + let token_ids = output.new_token_ids.clone();
151 + let logprobs = output
152 + .new_logprobs
153 + .as_ref()
154 + .and_then(|value| value.as_direct())
155 + .context("engine output did not include decoded sample logprobs")?;
156 + let prompt_logprobs = output
157 + .new_prompt_logprobs_tensors
158 + .as_ref()
159 + .and_then(|value| value.as_direct())
160 + .context("engine output did not include decoded prompt logprobs")?;
161 +
162 + println!("token_ids={token_ids:?}");
163 + println!("finish_reason={finish_reason:?}");
164 + println!("new_logprobs={logprobs:#?}");
165 + println!("new_prompt_logprobs_tensors={prompt_logprobs:#?}");
166 +
167 + client.shutdown().await.context("failed to shut down engine-core client")?;
168 +
169 + if finish_reason != Some(EngineCoreFinishReason::Length) {
170 + bail!("unexpected finish_reason: expected Length, got {finish_reason:?}");
171 + }
172 + if token_ids.is_empty() {
173 + bail!("engine returned no generated token ids");
174 + }
175 + if logprobs.is_empty() {
176 + bail!("decoded logprobs payload is unexpectedly empty");
177 + }
178 + if prompt_logprobs.is_empty() {
179 + bail!("decoded prompt logprobs payload is unexpectedly empty");
180 + }
181 + if prompt_logprobs.len() + 1 < prompt_token_ids.len() {
182 + bail!(
183 + "prompt logprobs rows look too short: prompt_len={}, rows={}",
184 + prompt_token_ids.len(),
185 + prompt_logprobs.len()
186 + );
187 + }
188 +
189 + Ok(())
190 +}
rust/src/engine-core-client/examples/external_engine_utility_call.rs new
+143
@@ -0,0 +1,143 @@
1 +use std::time::Duration;
2 +
3 +use anyhow::{Context, Result, bail};
4 +use clap::Parser;
5 +use tracing_subscriber::EnvFilter;
6 +use vllm_engine_core_client::{EngineCoreClient, EngineCoreClientConfig, TransportMode};
7 +
8 +#[derive(Debug, Parser)]
9 +#[command(about = "Smoke-test EngineCoreClient utility calls against an external vLLM engine.")]
10 +struct Args {
11 + #[arg(long)]
12 + handshake_address: String,
13 + #[arg(long, default_value_t = 1)]
14 + engine_count: usize,
15 + #[arg(long, default_value = "Qwen/Qwen3-0.6B")]
16 + model: String,
17 + #[arg(long, default_value = "127.0.0.1")]
18 + host: String,
19 + #[arg(long, default_value_t = 0)]
20 + client_index: u32,
21 + #[arg(long, default_value_t = 30)]
22 + ready_timeout_secs: u64,
23 + #[arg(
24 + long,
25 + default_value_t = false,
26 + help = "Expected initial result of is_sleeping() before running the smoke steps."
27 + )]
28 + expected_is_sleeping: bool,
29 + #[arg(long, default_value_t = false)]
30 + reset_running_requests: bool,
31 + #[arg(long, default_value_t = false)]
32 + reset_external: bool,
33 + #[arg(long, default_value_t = 1)]
34 + sleep_level: u32,
35 + #[arg(long, default_value = "abort")]
36 + sleep_mode: String,
37 + #[arg(
38 + long,
39 + default_value_t = false,
40 + help = "Skip sleep/wake_up calls when the engine was not started with sleep-mode support."
41 + )]
42 + skip_sleep_wake: bool,
43 +}
44 +
45 +fn init_tracing() {
46 + let filter = EnvFilter::try_from_default_env()
47 + .unwrap_or_else(|_| EnvFilter::new("vllm_engine_core_client=debug"));
48 + let _ = tracing_subscriber::fmt().with_env_filter(filter).try_init();
49 +}
50 +
51 +#[tokio::main(flavor = "multi_thread")]
52 +async fn main() -> Result<()> {
53 + init_tracing();
54 + let args = Args::parse();
55 + let client = EngineCoreClient::connect(EngineCoreClientConfig {
56 + transport_mode: TransportMode::HandshakeOwner {
57 + handshake_address: args.handshake_address.clone(),
58 + advertised_host: args.host.clone(),
59 + engine_count: args.engine_count,
60 + ready_timeout: Duration::from_secs(args.ready_timeout_secs),
61 + local_input_address: None,
62 + local_output_address: None,
63 + },
64 + coordinator_mode: None,
65 + model_name: args.model.clone(),
66 + client_index: args.client_index,
67 + })
68 + .await
69 + .context("failed to connect to external vLLM engine")?;
70 +
71 + println!("model={}", args.model);
72 + println!("handshake_address={}", args.handshake_address);
73 + println!("engine_count={}", args.engine_count);
74 + println!("input_address={}", client.input_address());
75 + println!("output_address={}", client.output_address());
76 + println!("engine_identities={:x?}", client.engine_identities());
77 +
78 + let initial_is_sleeping =
79 + client.is_sleeping().await.context("failed to call is_sleeping utility")?;
80 +
81 + println!("is_sleeping(initial)={initial_is_sleeping}");
82 +
83 + if initial_is_sleeping != args.expected_is_sleeping {
84 + bail!(
85 + "unexpected initial is_sleeping state: expected {}, got {}",
86 + args.expected_is_sleeping,
87 + initial_is_sleeping
88 + );
89 + }
90 +
91 + let reset_prefix_cache = client
92 + .reset_prefix_cache(args.reset_running_requests, args.reset_external)
93 + .await
94 + .context("failed to call reset_prefix_cache utility")?;
95 + println!("reset_prefix_cache={reset_prefix_cache}");
96 +
97 + client.reset_mm_cache().await.context("failed to call reset_mm_cache utility")?;
98 + println!("reset_mm_cache=ok");
99 +
100 + client
101 + .reset_encoder_cache()
102 + .await
103 + .context("failed to call reset_encoder_cache utility")?;
104 + println!("reset_encoder_cache=ok");
105 +
106 + if args.skip_sleep_wake {
107 + println!("sleep_wake=skipped");
108 + } else {
109 + client.sleep(args.sleep_level, &args.sleep_mode).await.with_context(|| {
110 + format!(
111 + "failed to call sleep utility with level={} mode={}",
112 + args.sleep_level, args.sleep_mode
113 + )
114 + })?;
115 + println!(
116 + "sleep=ok level={} mode={}",
117 + args.sleep_level, args.sleep_mode
118 + );
119 +
120 + let sleeping_after_sleep =
121 + client.is_sleeping().await.context("failed to call is_sleeping after sleep")?;
122 + println!("is_sleeping(after_sleep)={sleeping_after_sleep}");
123 +
124 + if !sleeping_after_sleep {
125 + bail!("engine should report sleeping=true after sleep()");
126 + }
127 +
128 + client.wake_up(None).await.context("failed to call wake_up utility")?;
129 + println!("wake_up=ok");
130 +
131 + let sleeping_after_wake =
132 + client.is_sleeping().await.context("failed to call is_sleeping after wake_up")?;
133 + println!("is_sleeping(after_wake)={sleeping_after_wake}");
134 +
135 + if sleeping_after_wake {
136 + bail!("engine should report sleeping=false after wake_up()");
137 + }
138 + }
139 +
140 + client.shutdown().await.context("failed to shut down engine-core client")?;
141 +
142 + Ok(())
143 +}
rust/src/engine-core-client/src/client.rs new
+673
@@ -0,0 +1,673 @@
1 +use std::sync::Arc;
2 +use std::time::Duration;
3 +
4 +use futures::future::{join_all, try_join_all};
5 +use tokio::sync::mpsc;
6 +use tokio_util::task::AbortOnDropHandle;
7 +use tracing::{debug, info, trace};
8 +
9 +use crate::client::imp::{ClientInner, run_abort_loop, run_output_dispatcher_loop};
10 +use crate::coordinator::CoordinatorHandle;
11 +use crate::error::{Error, Result};
12 +use crate::protocol::handshake::EngineCoreReadyResponse;
13 +use crate::protocol::{
14 + EngineCoreRequest, EngineCoreRequestType, EngineCoreUtilityRequest, ModelDtype,
15 +};
16 +use crate::transport::{self, ConnectedEngine};
17 +
18 +pub(crate) mod imp;
19 +mod state;
20 +mod stream;
21 +
22 +pub use stream::{EngineCoreOutputStream, EngineCoreStreamOutput};
23 +
24 +/// How the frontend acquires its request/response transport with Python
25 +/// `EngineCoreProc`s.
26 +#[derive(Debug, Clone, PartialEq, Eq)]
27 +pub enum TransportMode {
28 + /// The Rust process owns the startup handshake and allocates or binds the
29 + /// frontend transport addresses itself before replying to engine
30 + /// `HELLO` messages.
31 + HandshakeOwner {
32 + /// Shared handshake endpoint that engines dial during startup.
33 + handshake_address: String,
34 + /// Host/IP that engines should use to connect back to the frontend
35 + /// transport sockets.
36 + advertised_host: String,
37 + /// Total number of engines expected to join this transport.
38 + engine_count: usize,
39 + /// Maximum time to wait for each startup phase to complete.
40 + ready_timeout: Duration,
41 + /// Optional explicit bind address for the input ROUTER socket.
42 + local_input_address: Option<String>,
43 + /// Optional explicit bind address for the output PULL socket.
44 + local_output_address: Option<String>,
45 + },
46 +
47 + /// The Python supervisor has already chosen the frontend transport
48 + /// addresses, and the Rust process only needs to bind them and wait for
49 + /// engine registration frames.
50 + Bootstrapped {
51 + /// Input ROUTER socket address that engines will connect to for
52 + /// requests.
53 + input_address: String,
54 + /// Output PULL socket address that engines will connect to for
55 + /// responses.
56 + output_address: String,
57 + /// Total number of engines expected to register on this transport.
58 + engine_count: usize,
59 + /// Maximum time to wait for all expected engines to register.
60 + ready_timeout: Duration,
61 + },
62 +}
63 +
64 +/// Which coordinator implementation should be active when one is present for a
65 +/// frontend client.
66 +#[derive(Debug, Clone, PartialEq, Eq)]
67 +pub enum CoordinatorMode {
68 + /// Run the Rust in-process coordinator for managed `serve` deployments.
69 + InProc,
70 + /// Connect to an external coordinator owned by another process.
71 + External { address: String },
72 +}
73 +
74 +/// Configuration for connecting a Rust frontend client to an already running
75 +/// Python `EngineCoreProc`.
76 +#[derive(Debug, Clone, PartialEq, Eq)]
77 +pub struct EngineCoreClientConfig {
78 + /// Frontend-to-engine transport setup.
79 + pub transport_mode: TransportMode,
80 + /// Frontend-side coordinator behavior, or `None` when requests should flow
81 + /// directly to engines without any coordinator involvement.
82 + pub coordinator_mode: Option<CoordinatorMode>,
83 + /// Model name used for frontend-side metrics labels.
84 + pub model_name: String,
85 + /// Frontend client index stamped onto every request.
86 + pub client_index: u32,
87 +}
88 +
89 +impl EngineCoreClientConfig {
90 + /// Create a new client config with the given handshake address, expecting a
91 + /// single engine, and default values for all other fields.
92 + pub fn new_single(handshake_address: impl Into<String>) -> Self {
93 + Self {
94 + transport_mode: TransportMode::HandshakeOwner {
95 + handshake_address: handshake_address.into(),
96 + advertised_host: "127.0.0.1".to_string(),
97 + engine_count: 1,
98 + ready_timeout: Duration::from_secs(30),
99 + local_input_address: None,
100 + local_output_address: None,
101 + },
102 + coordinator_mode: None,
103 + model_name: String::new(),
104 + client_index: 0,
105 + }
106 + }
107 +
108 + /// Set the model name used by frontend-side metrics and diagnostics.
109 + pub fn with_model_name(mut self, model_name: impl Into<String>) -> Self {
110 + self.model_name = model_name.into();
111 + self
112 + }
113 +
114 + /// Override the client index stamped onto every outgoing request.
115 + pub fn with_client_index(mut self, client_index: u32) -> Self {
116 + self.client_index = client_index;
117 + self
118 + }
119 +
120 + /// Override the optional coordinator mode for this client config.
121 + pub fn with_coordinator_mode(mut self, coordinator_mode: Option<CoordinatorMode>) -> Self {
122 + self.coordinator_mode = coordinator_mode;
123 + self
124 + }
125 +
126 + /// Override the locally bound input/output addresses for handshake-owned
127 + /// transport mode.
128 + ///
129 + /// This is primarily used by tests that want deterministic IPC endpoints
130 + /// while still exercising the handshake-owned startup path.
131 + pub fn with_local_input_output_addresses(
132 + mut self,
133 + local_input_address: Option<String>,
134 + local_output_address: Option<String>,
135 + ) -> Self {
136 + let TransportMode::HandshakeOwner {
137 + local_input_address: current_input,
138 + local_output_address: current_output,
139 + ..
140 + } = &mut self.transport_mode
141 + else {
142 + panic!("local input/output overrides are only valid in handshake-owned mode");
143 + };
144 + *current_input = local_input_address;
145 + *current_output = local_output_address;
146 + self
147 + }
148 +}
149 +
150 +/// The reason a request stream is being aborted when its output stream is
151 +/// dropped.
152 +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
153 +pub enum AbortCause {
154 + /// The consumer dropped the stream before the request reached a terminal
155 + /// engine output.
156 + #[default]
157 + DroppedStream,
158 + /// The frontend matched a stop string locally and intentionally stopped
159 + /// consuming the stream.
160 + StopStringMatched,
161 +}
162 +
163 +task_local::task_local! {
164 + static ABORT_CAUSE: AbortCause;
165 +}
166 +
167 +impl AbortCause {
168 + /// Return the abort cause currently associated with this task, or
169 + /// [`AbortCause::DroppedStream`] by default.
170 + pub fn current() -> Self {
171 + ABORT_CAUSE.try_get().unwrap_or_default()
172 + }
173 +
174 + /// Drop one value while marking the drop as happening for this abort cause.
175 + pub fn drop_as<T>(self, value: T) {
176 + ABORT_CAUSE.sync_scope(self, move || drop(value));
177 + }
178 +}
179 +
180 +/// Internal auto-abort work item sent from stream `Drop` handlers to the abort
181 +/// worker.
182 +#[derive(Debug, Clone)]
183 +pub(crate) struct AbortRequest {
184 + request_id: String,
185 + cause: AbortCause,
186 +}
187 +
188 +/// Default ZMQ-based implementation that talks directly to a Python
189 +/// `EngineCoreProc`.
190 +pub struct EngineCoreClient {
191 + config: EngineCoreClientConfig,
192 + input_address: String,
193 + output_address: String,
194 + engines: Vec<ConnectedEngine>,
195 + inner: Arc<ClientInner>,
196 + coordinator: Option<CoordinatorHandle>,
197 + abort_tx: mpsc::UnboundedSender<AbortRequest>,
198 +
199 + // Background tasks
200 + output_task: AbortOnDropHandle<()>,
201 + dispatcher_task: AbortOnDropHandle<()>,
202 + abort_task: AbortOnDropHandle<()>,
203 + coordinator_output_task: Option<AbortOnDropHandle<()>>,
204 + coordinator_task: Option<AbortOnDropHandle<()>>,
205 +}
206 +
207 +impl EngineCoreClient {
208 + /// Connect to Python `EngineCoreProc`s using the configured
209 + /// transport/coordinator modes.
210 + ///
211 + /// In handshake-owned mode this method drives the full engine startup
212 + /// handshake. In bootstrapped mode it binds the provided frontend
213 + /// sockets and waits for the expected engine registration frames.
214 + pub async fn connect(config: EngineCoreClientConfig) -> Result<Self> {
215 + let connected = match &config.transport_mode {
216 + TransportMode::HandshakeOwner {
217 + handshake_address,
218 + advertised_host,
219 + engine_count,
220 + ready_timeout,
221 + local_input_address,
222 + local_output_address,
223 + } => {
224 + let enable_inproc_coordinator = match config.coordinator_mode {
225 + None => false,
226 + Some(CoordinatorMode::InProc) => true,
227 + Some(CoordinatorMode::External { .. }) => {
228 + return Err(Error::UnsupportedExternalCoordinator);
229 + }
230 + };
231 +
232 + transport::connect_handshake(
233 + handshake_address,
234 + *engine_count,
235 + advertised_host,
236 + local_input_address.as_deref(),
237 + local_output_address.as_deref(),
238 + enable_inproc_coordinator,
239 + *ready_timeout,
240 + )
241 + .await?
242 + }
243 +
244 + TransportMode::Bootstrapped {
245 + input_address,
246 + output_address,
247 + engine_count,
248 + ready_timeout,
249 + } => {
250 + if let Some(CoordinatorMode::InProc) = config.coordinator_mode {
251 + panic!("cannot use in-process coordinator with bootstrapped transport mode")
252 + }
253 +
254 + transport::connect_bootstrapped(
255 + input_address,
256 + output_address,
257 + *engine_count,
258 + *ready_timeout,
259 + )
260 + .await?
261 + }
262 + };
263 +
264 + Self::from_connected(config, connected).await
265 + }
266 +
267 + /// Create a new client instance from the connected transport state after
268 + /// the startup handshake completes.
269 + async fn from_connected(
270 + config: EngineCoreClientConfig,
271 + connected: transport::ConnectedTransport,
272 + ) -> Result<Self> {
273 + let (output_tx, output_rx) = mpsc::channel(64);
274 + let (abort_tx, abort_rx) = mpsc::unbounded_channel();
275 + let engines = connected.engines;
276 + let inner = Arc::new(ClientInner::new(
277 + connected.input_send,
278 + config.model_name.clone(),
279 + &engines,
280 + ));
281 + let output_task = AbortOnDropHandle::new(tokio::spawn(transport::run_output_loop(
282 + connected.output_socket,
283 + output_tx,
284 + )));
285 + let dispatcher_task = AbortOnDropHandle::new(tokio::spawn(run_output_dispatcher_loop(
286 + inner.clone(),
287 + output_rx,
288 + )));
289 + let abort_task =
290 + AbortOnDropHandle::new(tokio::spawn(run_abort_loop(inner.clone(), abort_rx)));
291 +
292 + // If any engine reported a dp_stats_address in its ready response, use it
293 + // as the external coordinator address.
294 + let dp_stats_address: Option<String> = engines
295 + .iter()
296 + .filter_map(|e| e.ready_response.as_ref())
297 + .find_map(|r| r.dp_stats_address.clone());
298 +
299 + let (coordinator, coordinator_output_task, coordinator_task) =
300 + if let Some(coordinator_transport) = connected.coordinator {
301 + let (handle, runner) =
302 + CoordinatorHandle::new_inproc(coordinator_transport.input_socket);
303 + let (coordinator_output_tx, coordinator_output_rx) = mpsc::channel(64);
304 + let coordinator_output_task =
305 + AbortOnDropHandle::new(tokio::spawn(transport::run_output_loop(
306 + coordinator_transport.output_socket,
307 + coordinator_output_tx,
308 + )));
309 + let coordinator_task = AbortOnDropHandle::new(tokio::spawn(
310 + runner.run(coordinator_output_rx, inner.clone()),
311 + ));
312 + (
313 + Some(handle),
314 + Some(coordinator_output_task),
315 + Some(coordinator_task),
316 + )
317 + } else if let Some(address) =
318 + dp_stats_address.as_deref().or(match config.coordinator_mode.as_ref() {
319 + Some(CoordinatorMode::External { address }) => Some(address.as_str()),
320 + _ => None,
321 + })
322 + {
323 + let (handle, service) = CoordinatorHandle::connect_external(address).await?;
324 + let coordinator_task =
325 + AbortOnDropHandle::new(tokio::spawn(service.run(inner.clone())));
326 + (Some(handle), None, Some(coordinator_task))
327 + } else {
328 + (None, None, None)
329 + };
330 +
331 + Ok(Self {
332 + config,
333 + input_address: connected.input_address,
334 + output_address: connected.output_address,
335 + engines,
336 + inner,
337 + coordinator,
338 + abort_tx,
339 + output_task,
340 + dispatcher_task,
341 + abort_task,
342 + coordinator_output_task,
343 + coordinator_task,
344 + })
345 + }
346 +
347 + /// Return the address of the input socket that the client uses to send
348 + /// requests to the engine.
349 + pub fn input_address(&self) -> &str {
350 + &self.input_address
351 + }
352 +
353 + /// Return the address of the output socket that the client listens on for
354 + /// engine responses.
355 + pub fn output_address(&self) -> &str {
356 + &self.output_address
357 + }
358 +
359 + /// Return the number of engines connected to this client.
360 + pub fn engine_count(&self) -> usize {
361 + self.engines.len()
362 + }
363 +
364 + /// Return the engine identities of all engines connected to this client.
365 + pub fn engine_identities(&self) -> Vec<&[u8]> {
366 + self.engines.iter().map(|engine| &*engine.engine_id).collect()
367 + }
368 +
369 + /// Return the ready responses received from all engines on the input
370 + /// socket.
371 + pub fn ready_responses(&self) -> Vec<&EngineCoreReadyResponse> {
372 + self.engines
373 + .iter()
374 + .filter_map(|engine| engine.ready_response.as_ref())
375 + .collect()
376 + }
377 +
378 + /// Return the engine-reported effective model dtype, when available.
379 + pub fn model_dtype(&self) -> Option<ModelDtype> {
380 + self.engines
381 + .iter()
382 + .filter_map(|engine| engine.ready_response.as_ref())
383 + .find_map(|response| response.dtype)
384 + }
385 +
386 + /// Return the total number of GPU blocks summed across all connected
387 + /// engines.
388 + pub fn total_num_gpu_blocks(&self) -> u64 {
389 + self.engines
390 + .iter()
391 + .filter_map(|engine| engine.ready_response.as_ref())
392 + .map(|r| r.num_gpu_blocks)
393 + .sum()
394 + }
395 +
396 + /// Return the minimum engine-reported `max_model_len` across all engines.
397 + ///
398 + /// This is the auto-fitted value after KV cache profiling and may differ
399 + /// from the originally configured value.
400 + pub fn max_model_len(&self) -> Option<u32> {
401 + self.engines
402 + .iter()
403 + .filter_map(|e| e.ready_response.as_ref())
404 + .map(|r| r.max_model_len as u32)
405 + .min()
406 + }
407 +
408 + /// Get the model name associated with this client used for metrics
409 + /// labeling.
410 + pub fn model_name(&self) -> &str {
411 + self.inner.model_name()
412 + }
413 +
414 + /// Return whether the client still considers the engine healthy.
415 + pub fn is_healthy(&self) -> bool {
416 + self.inner.is_healthy()
417 + }
418 +
419 + /// Return the first persistent health error observed by the client, if any.
420 + pub fn health_error(&self) -> Option<Arc<Error>> {
421 + self.inner.health_error()
422 + }
423 +}
424 +
425 +// Client API implementation.
426 +impl EngineCoreClient {
427 + /// Add a new request to the engine and return a per-request raw output
428 + /// stream.
429 + pub async fn call(&self, mut req: EngineCoreRequest) -> Result<EngineCoreOutputStream> {
430 + req.client_index = self.config.client_index;
431 + req.validate()?;
432 + trace!(
433 + request_id = %req.request_id,
434 + client_index = req.client_index,
435 + current_wave = req.current_wave,
436 + request = ?req,
437 + "sending add request"
438 + );
439 +
440 + let request_id = req.request_id.clone();
441 + let data_parallel_rank = req.data_parallel_rank;
442 + let (engine_id, rx) =
443 + self.inner.register_request(request_id.clone(), data_parallel_rank)?;
444 +
445 + let result: Result<()> = async {
446 + if let Some(coordinator) = self.coordinator.as_ref() {
447 + let snapshot = coordinator.snapshot();
448 + req.current_wave = snapshot.current_wave;
449 + if !snapshot.engines_running {
450 + coordinator.notify_first_request(engine_id.clone())?;
451 + }
452 + }
453 +
454 + debug!(
455 + request_id = req.request_id,
456 + ?engine_id,
457 + "registered request to engine"
458 + );
459 +
460 + self.inner.send_to_engine(&engine_id, EngineCoreRequestType::Add, &req).await?;
461 + Ok(())
462 + }
463 + .await;
464 +
465 + // Failed to send the request to the engine, roll back the registration.
466 + if let Err(error) = result {
467 + self.inner.rollback_request(&request_id);
468 + return Err(error);
469 + }
470 +
471 + Ok(EngineCoreOutputStream::new(
472 + request_id,
473 + self.abort_tx.clone(),
474 + rx,
475 + ))
476 + }
477 +
478 + /// Abort currently in-flight requests by request ID.
479 + pub async fn abort(&self, ids: &[String]) -> Result<()> {
480 + let abortable = self.inner.abortable_request_ids(ids)?;
481 +
482 + trace!(request_ids = ?ids, abortable_request_ids = ?abortable, "sending abort request ids");
483 +
484 + if abortable.is_empty() {
485 + return Ok(());
486 + }
487 +
488 + for (engine_id, request_ids) in abortable {
489 + self.inner.do_abort_requests(&engine_id, &request_ids).await?;
490 + }
491 + Ok(())
492 + }
493 +
494 + /// Call a typed utility method on all connected engines, returning one
495 + /// decoded result per connected engine if all calls succeed or an error
496 + /// if any call fails.
497 + ///
498 + /// Callers should pass utility arguments using Rust tuple semantics so the
499 + /// encoded payload matches Python's `(client_index, call_id,
500 + /// method_name, args)` contract: `()`, `(arg,)`, `(arg1, arg2)`, etc.
501 + pub async fn call_utility<T, A>(&self, method: &str, args: A) -> Result<Vec<T>>
502 + where
503 + T: serde::de::DeserializeOwned,
504 + A: serde::Serialize + std::fmt::Debug,
505 + {
506 + trace!(
507 + method,
508 + client_index = self.config.client_index,
509 + engine_count = self.engines.len(),
510 + "sending utility request"
511 + );
512 +
513 + // Phase 1: allocate one call id per engine and build the per-engine
514 + // request payloads up-front. Any failure here (registry closed, encode
515 + // error) must roll back the call ids already allocated so they do not
516 + // leak in the utility registry until shutdown.
517 + let mut pending_calls = Vec::with_capacity(self.engines.len());
518 + let mut prepared_sends = Vec::with_capacity(self.engines.len());
519 + for engine in &self.engines {
520 + let (call_id, rx) = match self.inner.allocate_and_register_utility_call() {
521 + Ok(pair) => pair,
522 + Err(err) => {
523 + self.inner.unregister_utility_calls(pending_calls.iter().map(|(id, _)| *id));
524 + return Err(err);
525 + }
526 + };
527 + let request = match EngineCoreUtilityRequest::new(
528 + self.config.client_index,
529 + call_id,
530 + method,
531 + &args,
532 + ) {
533 + Ok(request) => request,
534 + Err(err) => {
535 + self.inner.unregister_utility_calls(
536 + pending_calls.iter().map(|(id, _)| *id).chain(std::iter::once(call_id)),
537 + );
538 + return Err(err);
539 + }
540 + };
541 + pending_calls.push((call_id, rx));
542 + prepared_sends.push((&engine.engine_id, request));
543 + }
544 +
545 + // Phase 2: dispatch every utility request concurrently. `try_join_all`
546 + // fails fast on the first transport error and drops the remaining send
547 + // futures; any engines that already received the request will reply,
548 + // but those replies are simply dropped because we roll back the call
549 + // ids below.
550 + let send_futures = prepared_sends.iter().map(|(engine_id, request)| {
551 + self.inner.send_to_engine(engine_id, EngineCoreRequestType::Utility, request)
552 + });
553 + if let Err(err) = try_join_all(send_futures).await {
554 + self.inner.unregister_utility_calls(pending_calls.iter().map(|(id, _)| *id));
555 + return Err(err);
556 + }
557 +
558 + // Phase 3: wait for all engines to respond and preserve the per-engine
559 + // result list.
560 + let futures = pending_calls.into_iter().map(|(call_id, rx)| async move {
561 + rx.await
562 + .map_err(|_| Error::UtilityCallClosed {
563 + method: method.to_string(),
564 + call_id,
565 + })??
566 + .into_typed_result(method)
567 + });
568 + try_join_all(futures).await
569 + }
570 +
571 + /// Execute `collective_rpc` on all engines and flatten all engine results
572 + /// into one list.
573 + pub async fn collective_rpc<A, K>(
574 + &self,
575 + method: &str,
576 + timeout: Option<f64>,
577 + args: A,
578 + kwargs: K,
579 + ) -> Result<Vec<rmpv::Value>>
580 + where
581 + A: serde::Serialize + std::fmt::Debug,
582 + K: serde::Serialize + std::fmt::Debug,
583 + {
584 + let results = self
585 + .call_utility::<rmpv::Value, _>("collective_rpc", (method, timeout, args, kwargs))
586 + .await?;
587 +
588 + Ok(results
589 + .into_iter()
590 + .flat_map(|result| match result {
591 + // Each engine's `collective_rpc` result is itself the worker-level result list.
592 + rmpv::Value::Array(results) => results,
593 + other => vec![other],
594 + })
595 + .collect())
596 + }
597 +
598 + /// Return whether the engine is currently sleeping at any level.
599 + pub async fn is_sleeping(&self) -> Result<bool> {
600 + // TODO: we only return the result of the first engine here.
601 + Ok(self.call_utility("is_sleeping", ()).await?[0])
602 + }
603 +
604 + /// Reset the multi-modal cache.
605 + pub async fn reset_mm_cache(&self) -> Result<()> {
606 + self.call_utility::<(), _>("reset_mm_cache", ()).await?;
607 + Ok(())
608 + }
609 +
610 + /// Reset the encoder cache.
611 + pub async fn reset_encoder_cache(&self) -> Result<()> {
612 + self.call_utility::<(), _>("reset_encoder_cache", ()).await?;
613 + Ok(())
614 + }
615 +
616 + /// Reset the prefix cache and optionally the external connector cache.
617 + pub async fn reset_prefix_cache(
618 + &self,
619 + reset_running_requests: bool,
620 + reset_connector: bool,
621 + ) -> Result<bool> {
622 + // TODO: we only return the result of the first engine here.
623 + Ok(self
624 + .call_utility(
625 + "reset_prefix_cache",
626 + (reset_running_requests, reset_connector),
627 + )
628 + .await?[0])
629 + }
630 +
631 + /// Put the engine to sleep.
632 + pub async fn sleep(&self, level: u32, mode: &str) -> Result<()> {
633 + self.call_utility::<(), _>("sleep", (level, mode)).await?;
634 + Ok(())
635 + }
636 +
637 + /// Wake the engine from sleep, optionally limiting the wake-up to specific
638 + /// tags.
639 + pub async fn wake_up(&self, tags: Option<Vec<String>>) -> Result<()> {
640 + self.call_utility::<(), _>("wake_up", (tags,)).await?;
641 + Ok(())
642 + }
643 +
644 + /// Shut down local client tasks and close transport state.
645 + pub async fn shutdown(self) -> Result<()> {
646 + let Self {
647 + inner,
648 + abort_tx,
649 + output_task,
650 + dispatcher_task,
651 + abort_task,
652 + coordinator_output_task,
653 + coordinator_task,
654 + ..
655 + } = self;
656 +
657 + info!("shutting down engine-core client");
658 + inner.shutdown();
659 + drop(abort_tx);
660 +
661 + // Abort all client tasks first, then await them.
662 + // Note the aborting orders here.
663 + let mut tasks = vec![abort_task, dispatcher_task, output_task];
664 + tasks.extend(coordinator_task);
665 + tasks.extend(coordinator_output_task);
666 +
667 + tasks.iter().for_each(|t| t.abort());
668 + join_all(tasks).await;
669 +
670 + info!("engine-core client shut down");
671 + Ok(())
672 + }
673 +}
rust/src/engine-core-client/src/client/imp.rs new
+410
@@ -0,0 +1,410 @@
1 +use std::collections::BTreeMap;
2 +use std::slice;
3 +use std::sync::Arc;
4 +
5 +use arc_swap::ArcSwapOption;
6 +use parking_lot::Mutex;
7 +use thiserror_ext::AsReport as _;
8 +use tokio::sync::mpsc;
9 +use tracing::{debug, info, trace, warn};
10 +use vllm_metrics::METRICS;
11 +use zeromq::RouterSendHalf;
12 +
13 +use crate::client::state::{OutputReceiver, RequestRegistry, UtilityReceiver, UtilityRegistry};
14 +use crate::client::stream::EngineCoreStreamOutput;
15 +use crate::client::{AbortCause, AbortRequest};
16 +use crate::error::{client_closed, dispatcher_closed, unexpected_dispatcher_output};
17 +use crate::metrics::record_scheduler_stats;
18 +use crate::protocol::stats::SchedulerStats;
19 +use crate::protocol::{
20 + ClassifiedEngineCoreOutputs, EngineCoreOutput, EngineCoreOutputs, EngineCoreRequestType,
21 + UtilityOutput, encode_msgpack,
22 +};
23 +use crate::transport::{ConnectedEngine, EngineId};
24 +use crate::{Error, Result, transport};
25 +
26 +pub(crate) struct ClientInner {
27 + input_send: RouterSendHalf,
28 + model_name: String,
29 + request_reg: Mutex<RequestRegistry>,
30 + utility_reg: Mutex<UtilityRegistry>,
31 + health_error: ArcSwapOption<Error>,
32 +}
33 +
34 +impl ClientInner {
35 + /// Create a new instance with the given input send half after the startup
36 + /// handshake completes.
37 + pub fn new(
38 + input_send: RouterSendHalf,
39 + model_name: String,
40 + engines: &[ConnectedEngine],
41 + ) -> Self {
42 + Self {
43 + input_send,
44 + model_name,
45 + request_reg: Mutex::new(RequestRegistry::new(engines)),
46 + utility_reg: Mutex::new(UtilityRegistry::default()),
47 + health_error: ArcSwapOption::empty(),
48 + }
49 + }
50 +
51 + /// Get the model name associated with this client used for metrics
52 + /// labeling.
53 + pub fn model_name(&self) -> &str {
54 + &self.model_name
55 + }
56 +
57 + /// Register a newly added request. Return the selected engine id and the
58 + /// per-request output channel bound to its `request_id`.
59 + ///
60 + /// When `data_parallel_rank` is provided, the request is routed to that
61 + /// specific engine rank, bypassing load balancing.
62 + pub fn register_request(
63 + &self,
64 + request_id: String,
65 + data_parallel_rank: Option<u32>,
66 + ) -> Result<(EngineId, OutputReceiver)> {
67 + let mut registry = self.request_reg.lock();
68 + if registry.is_closed() {
69 + return Err(self.closed_error());
70 + }
71 + registry.register(request_id, data_parallel_rank)
72 + }
73 +
74 + /// Allocate the next utility `call_id` and register its waiting receiver.
75 + pub fn allocate_and_register_utility_call(&self) -> Result<(i64, UtilityReceiver)> {
76 + let mut registry = self.utility_reg.lock();
77 + if registry.is_closed() {
78 + return Err(self.closed_error());
79 + }
80 + Ok(registry.allocate_and_register())
81 + }
82 +
83 + /// Undo a batch of utility call allocations when the fan-out send fails
84 + /// partway through. Silently ignores unknown call ids so callers can pass
85 + /// the full set without first filtering successful sends.
86 + pub fn unregister_utility_calls(&self, call_ids: impl IntoIterator<Item = i64>) {
87 + self.utility_reg.lock().unregister_many(call_ids);
88 + }
89 +
90 + /// Undo a request registration when `add_request()` fails.
91 + pub fn rollback_request(&self, request_id: &str) {
92 + let _ = self.request_reg.lock().remove(request_id);
93 + }
94 +
95 + /// Filter the given request IDs to the subset that are still tracked as
96 + /// active and can be aborted, grouped by the engine that originally
97 + /// accepted them.
98 + pub fn abortable_request_ids(
99 + &self,
100 + request_ids: &[String],
101 + ) -> Result<BTreeMap<EngineId, Vec<String>>> {
102 + let registry = self.request_reg.lock();
103 + if registry.is_closed() {
104 + return Err(self.closed_error());
105 + }
106 + Ok(registry.abortable_request_ids(request_ids))
107 + }
108 +
109 + /// Obtain the stream sender for one output. If it indicates the request is
110 + /// finished, it will be removed from the registry.
111 + pub fn take_sender_for_output(
112 + &self,
113 + output: &EngineCoreOutput,
114 + ) -> Option<mpsc::UnboundedSender<Result<EngineCoreStreamOutput>>> {
115 + self.request_reg.lock().sender_for_output(output)
116 + }
117 +
118 + /// Remove a batch of requests that have finished or aborted, returning
119 + /// their stream senders.
120 + pub fn finish_requests<'a>(
121 + &self,
122 + request_ids: impl IntoIterator<Item = &'a String>,
123 + ) -> Vec<mpsc::UnboundedSender<Result<EngineCoreStreamOutput>>> {
124 + self.request_reg.lock().finish_many(request_ids)
125 + }
126 +
127 + /// Apply one scheduler stats update for the given engine to the local
128 + /// routing state. Returns `false` if the engine is unknown to the
129 + /// client.
130 + pub fn apply_scheduler_stats(&self, engine_index: u32, stats: &SchedulerStats) -> bool {
131 + self.request_reg.lock().apply_scheduler_stats(engine_index, stats)
132 + }
133 +
134 + /// Close all active request streams and utility calls with the first
135 + /// persistent health error.
136 + pub fn close_registries(&self, error: Arc<Error>) {
137 + let persistent_error = self.record_health_error(error);
138 + let request_senders = self.request_reg.lock().close();
139 + let utility_senders = self.utility_reg.lock().close();
140 +
141 + // Notify all ongoing requests that the client is closed.
142 + for sender in request_senders {
143 + let _ = sender.send(Err(Error::Shared(persistent_error.clone())));
144 + }
145 + for sender in utility_senders {
146 + let _ = sender.send(Err(Error::Shared(persistent_error.clone())));
147 + }
148 + }
149 +
150 + /// Return the first persistent health error observed by the client, if any.
151 + pub fn health_error(&self) -> Option<Arc<Error>> {
152 + self.health_error.load_full()
153 + }
154 +
155 + /// Return whether the client still considers the engine healthy.
156 + pub fn is_healthy(&self) -> bool {
157 + self.health_error.load().is_none()
158 + }
159 +
160 + /// Resolve one utility output to the waiting caller. Returns `true` if a
161 + /// waiting caller existed.
162 + pub fn resolve_utility_output(&self, output: UtilityOutput) -> bool {
163 + match self.utility_reg.lock().resolve(&output.call_id) {
164 + Some(sender) => {
165 + sender.send(Ok(output)).unwrap_or_default();
166 + true
167 + }
168 + None => false,
169 + }
170 + }
171 +
172 + /// Send the given message to the engine. The request should be first
173 + /// registered via `register_request()` to ensure the request stream is
174 + /// tracked.
175 + pub async fn send_to_engine<T>(
176 + &self,
177 + engine_id: &EngineId,
178 + request_type: EngineCoreRequestType,
179 + payload: &T,
180 + ) -> Result<()>
181 + where
182 + T: serde::Serialize + std::fmt::Debug,
183 + {
184 + // TODO: for `EngineCoreRequest`, split outbound tensor raw views into aux
185 + // frames instead of always producing a single msgpack frame.
186 + let payload = encode_msgpack(payload)?;
187 + let mut input_send = self.input_send.clone();
188 + transport::send_message(&mut input_send, engine_id, request_type.to_frame(), payload)
189 + .await?;
190 + Ok(())
191 + }
192 +
193 + /// Handle an abort request by sending the abort message to the engine.
194 + pub async fn do_abort_requests(
195 + &self,
196 + engine_id: &EngineId,
197 + request_ids: &[String],
198 + ) -> Result<()> {
199 + self.send_to_engine(engine_id, EngineCoreRequestType::Abort, &request_ids).await
200 + }
201 +
202 + /// Shut down by closing all active request streams and utility calls with a
203 + /// sticky client closed error.
204 + pub fn shutdown(&self) {
205 + self.close_registries(Arc::new(client_closed!("engine-core client shut down")));
206 + }
207 +
208 + /// Remove the request from the active registry for auto-abort and return
209 + /// the engine that the request was originally routed to, if it is still
210 + /// active.
211 + pub fn take_auto_abort_target(&self, request_id: &str) -> Option<EngineId> {
212 + let mut registry = self.request_reg.lock();
213 + let (_, engine_id) = registry.remove(request_id)?;
214 + if registry.is_closed() {
215 + return None;
216 + }
217 + Some(engine_id)
218 + }
219 +
220 + /// Publish the first persistent health error and return the sticky error
221 + /// recorded for this client. Later failures do not overwrite the first
222 + /// one so `/health` and post-close callers observe a stable cause.
223 + fn record_health_error(&self, error: Arc<Error>) -> Arc<Error> {
224 + if let Some(existing) = self.health_error.load_full() {
225 + return existing;
226 + }
227 + self.health_error
228 + .rcu(|current| current.clone().unwrap_or_else(|| error.clone()));
229 + self.health_error
230 + .load_full()
231 + .expect("health error must be recorded before registries close")
232 + }
233 +
234 + /// Assert there is a recorded health error and return a `Shared` variant
235 + /// wrapping it for error returns when the client is already closed.
236 + fn closed_error(&self) -> Error {
237 + Error::Shared(self.health_error.load_full().expect(
238 + "closed registry must have a recorded health error before rejecting new operations",
239 + ))
240 + }
241 +}
242 +
243 +/// Background loop that listens for request IDs to abort and sends abort
244 +/// messages to the engine. This is used to implement the auto-abort behavior
245 +/// when a request stream is dropped without being properly terminated.
246 +pub(crate) async fn run_abort_loop(
247 + inner: Arc<ClientInner>,
248 + mut abort_rx: mpsc::UnboundedReceiver<AbortRequest>,
249 +) {
250 + // TODO: receive and abort requests in batch
251 + while let Some(AbortRequest { request_id, cause }) = abort_rx.recv().await {
252 + let Some(engine_id) = inner.take_auto_abort_target(&request_id) else {
253 + debug!(request_id, "skip auto-abort for inactive request");
254 + continue;
255 + };
256 +
257 + match cause {
258 + AbortCause::DroppedStream => {
259 + info!(request_id, "auto-aborting request due to dropped stream")
260 + }
261 + AbortCause::StopStringMatched => {
262 + debug!(
263 + request_id,
264 + "auto-aborting request due to stop string matched"
265 + )
266 + }
267 + }
268 +
269 + if let Err(error) = inner.do_abort_requests(&engine_id, slice::from_ref(&request_id)).await
270 + {
271 + warn!(
272 + request_id,
273 + ?engine_id,
274 + error = %error.as_report(),
275 + "failed to auto-abort dropped request stream"
276 + );
277 + }
278 + }
279 +}
280 +
281 +/// Background loop that listens for engine-core outputs and dispatches them to
282 +/// the corresponding request streams based on their `request_id`.
283 +pub(crate) async fn run_output_dispatcher_loop(
284 + inner: Arc<ClientInner>,
285 + mut output_rx: mpsc::Receiver<Result<EngineCoreOutputs>>,
286 +) {
287 + let result: Result<()> = async {
288 + loop {
289 + let outputs = match output_rx.recv().await {
290 + Some(outputs) => outputs,
291 + None => Err(dispatcher_closed!(
292 + "engine-core output dispatcher channel closed"
293 + )),
294 + }?;
295 +
296 + match outputs.classify() {
297 + ClassifiedEngineCoreOutputs::RequestBatch(batch) => {
298 + for output in batch.outputs {
299 + let request_id = output.request_id.clone();
300 + let Some(sender) = inner.take_sender_for_output(&output) else {
301 + debug!(request_id, "dropping output for inactive request");
302 + continue;
303 + };
304 +
305 + let wrapped_output = EngineCoreStreamOutput {
306 + engine_index: batch.engine_index,
307 + timestamp: batch.timestamp,
308 + output,
309 + };
310 + if sender.send(Ok(wrapped_output)).is_err() {
311 + debug!(request_id, "request output stream receiver dropped");
312 + }
313 + }
314 +
315 + // The sender for normally-finished requests should have already been removed
316 + // from the registry when their final output was dispatched
317 + // above. This serves as a safety net to capture any
318 + // requests marked as finished by the engine.
319 + if let Some(finished_requests) = batch.finished_requests.as_ref() {
320 + for request_id in finished_requests {
321 + trace!(request_id, "request completed via finished_requests");
322 + }
323 + drop(inner.finish_requests(finished_requests));
324 + }
325 +
326 + if let Some(scheduler_stats) = batch.scheduler_stats.as_ref() {
327 + if !inner.apply_scheduler_stats(batch.engine_index, scheduler_stats) {
328 + debug!(
329 + engine_index = batch.engine_index,
330 + "dropping scheduler stats for unknown engine"
331 + );
332 + }
333 + record_scheduler_stats(
334 + &METRICS.scheduler,
335 + inner.model_name(),
336 + batch.engine_index,
337 + scheduler_stats,
338 + );
339 + }
340 + }
341 + ClassifiedEngineCoreOutputs::Utility(utility) => {
342 + let call_id = utility.output.call_id;
343 + if inner.resolve_utility_output(utility.output) {
344 + trace!(
345 + call_id,
346 + engine_index = utility.engine_index,
347 + "resolved utility output"
348 + );
349 + } else {
350 + warn!(
351 + call_id,
352 + engine_index = utility.engine_index,
353 + "dropping output for unexpected utility call"
354 + );
355 + }
356 + }
357 + other @ (ClassifiedEngineCoreOutputs::DpControl { .. }
358 + | ClassifiedEngineCoreOutputs::Other(_)) => {
359 + Err::<(), _>(unexpected_dispatcher_output!(
360 + "received unexpected output on main dispatcher path: {other:?}"
361 + ))?;
362 + }
363 + }
364 + }
365 + }
366 + .await;
367 + let Err(error) = result else { return };
368 +
369 + warn!(error = %error.as_report(), "output dispatcher exiting with error");
370 + inner.close_registries(Arc::new(error));
371 +}
372 +
373 +#[cfg(test)]
374 +mod tests {
375 + use zeromq::{RouterSocket, Socket};
376 +
377 + use super::*;
378 +
379 + async fn test_inner() -> ClientInner {
380 + let mut socket = RouterSocket::new();
381 + socket.bind("tcp://127.0.0.1:0").await.unwrap();
382 + let (send, _) = socket.split();
383 + ClientInner::new(
384 + send,
385 + "test-model".to_string(),
386 + &[ConnectedEngine {
387 + engine_id: EngineId::from(b"engine-0"),
388 + ready_response: None,
389 + }],
390 + )
391 + }
392 +
393 + #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
394 + async fn close_registries_records_first_health_error_only() {
395 + let inner = test_inner().await;
396 +
397 + inner.close_registries(Arc::new(Error::EngineCoreDead));
398 + assert!(!inner.is_healthy());
399 + assert!(matches!(
400 + inner.health_error().as_deref(),
401 + Some(Error::EngineCoreDead)
402 + ));
403 +
404 + inner.close_registries(Arc::new(client_closed!("shutdown")));
405 + assert!(matches!(
406 + inner.health_error().as_deref(),
407 + Some(Error::EngineCoreDead)
408 + ));
409 + }
410 +}
rust/src/engine-core-client/src/client/state.rs new
+670
@@ -0,0 +1,670 @@
1 +use std::collections::BTreeMap;
2 +use std::sync::atomic::{AtomicI64, Ordering};
3 +
4 +use tokio::sync::{mpsc, oneshot};
5 +use tracing::trace;
6 +
7 +use crate::EngineId;
8 +use crate::client::stream::EngineCoreStreamOutput;
9 +use crate::error::{Error, Result};
10 +use crate::protocol::stats::SchedulerStats;
11 +use crate::protocol::{EngineCoreOutput, UtilityOutput};
12 +use crate::transport::ConnectedEngine;
13 +
14 +pub type OutputSender = mpsc::UnboundedSender<Result<EngineCoreStreamOutput>>;
15 +pub type OutputReceiver = mpsc::UnboundedReceiver<Result<EngineCoreStreamOutput>>;
16 +pub type UtilitySender = oneshot::Sender<Result<UtilityOutput>>;
17 +pub type UtilityReceiver = oneshot::Receiver<Result<UtilityOutput>>;
18 +
19 +#[derive(Debug)]
20 +struct TrackedRequest {
21 + sender: OutputSender,
22 + engine_id: EngineId,
23 +}
24 +
25 +/// The latest real scheduler-side load snapshot observed from one engine.
26 +///
27 +/// These counters come from `scheduler_stats` on the normal engine output path
28 +/// and are the preferred routing signal once available.
29 +#[derive(Debug, Clone, Copy, PartialEq, Eq)]
30 +struct EngineLoadSnapshot {
31 + /// Requests still counted on the scheduler's waiting side.
32 + waiting: usize,
33 + /// Requests currently counted on the scheduler's running side.
34 + running: usize,
35 +}
36 +
37 +#[derive(Debug, Default)]
38 +struct EngineRoutingState {
39 + /// Requests admitted by this frontend that have not finished yet.
40 + ///
41 + /// This is used both as the bootstrap fallback before real scheduler stats
42 + /// exist and as a lower bound afterwards so asynchronous scheduler
43 + /// snapshots cannot erase frontend admission history.
44 + inflight: usize,
45 + /// The latest real scheduler snapshot received from this engine, if any.
46 + last_scheduler_stats: Option<EngineLoadSnapshot>,
47 +}
48 +
49 +impl EngineRoutingState {
50 + /// Compute the routing score used to pick the least-loaded engine.
51 + ///
52 + /// Scheduler stats can raise the load estimate above the frontend-local
53 + /// view, but they should not lower it below requests this frontend has
54 + /// already admitted. Waiting requests still get the same extra penalty
55 + /// as the original `waiting * 4 + running` score.
56 + fn routing_score(&self) -> usize {
57 + const WAITING_WEIGHT: usize = 4;
58 +
59 + let Some(stats) = self.last_scheduler_stats else {
60 + return self.inflight;
61 + };
62 +
63 + let scheduler_total = stats.running + stats.waiting;
64 + self.inflight.max(scheduler_total) + stats.waiting * (WAITING_WEIGHT - 1)
65 + }
66 +
67 + /// Replace the local routing view with a fresh real scheduler snapshot.
68 + fn apply_scheduler_counts(&mut self, next: EngineLoadSnapshot) {
69 + self.last_scheduler_stats = Some(next);
70 + }
71 +}
72 +
73 +/// Internal registry for tracking active requests and their output stream
74 +/// senders.
75 +///
76 +/// This is used to route incoming outputs to the correct request stream, and to
77 +/// ensure proper cleanup of senders when requests finish or the client shuts
78 +/// down.
79 +#[derive(Debug)]
80 +pub struct RequestRegistry {
81 + closed: bool,
82 + requests: BTreeMap<String, TrackedRequest>,
83 + routing_per_engine: BTreeMap<EngineId, EngineRoutingState>,
84 +}
85 +
86 +impl RequestRegistry {
87 + pub fn new(engines: &[ConnectedEngine]) -> Self {
88 + Self {
89 + closed: false,
90 + requests: BTreeMap::default(),
91 + routing_per_engine: engines
92 + .iter()
93 + .map(|engine| (engine.engine_id.clone(), EngineRoutingState::default()))
94 + .collect(),
95 + }
96 + }
97 +
98 + /// Register a newly added request. Create the per-request output channel
99 + /// bound to its `request_id` and return the selected engine id.
100 + ///
101 + /// When `data_parallel_rank` is provided, the request is routed directly to
102 + /// the engine at that rank index, bypassing load balancing. Otherwise
103 + /// the engine with the fewest in-flight requests is chosen.
104 + pub fn register(
105 + &mut self,
106 + request_id: String,
107 + data_parallel_rank: Option<u32>,
108 + ) -> Result<(EngineId, OutputReceiver)> {
109 + if self.requests.contains_key(&request_id) {
110 + return Err(Error::DuplicateRequestId { request_id });
111 + }
112 +
113 + let engine_id = self.choose_engine_for_request(data_parallel_rank)?;
114 + let (tx, rx) = mpsc::unbounded_channel();
115 + self.requests.insert(
116 + request_id,
117 + TrackedRequest {
118 + sender: tx,
119 + engine_id: engine_id.clone(),
120 + },
121 + );
122 +
123 + let state = self
124 + .routing_per_engine
125 + .get_mut(&engine_id)
126 + .expect("request registry must track all known engines");
127 + state.inflight += 1;
128 +
129 + Ok((engine_id, rx))
130 + }
131 +
132 + fn choose_engine_for_request(&mut self, data_parallel_rank: Option<u32>) -> Result<EngineId> {
133 + if let Some(rank) = data_parallel_rank {
134 + // Route to the engine at the specified rank index.
135 + let engine_id = EngineId::from_engine_index(rank);
136 + return self
137 + .routing_per_engine
138 + .contains_key(&engine_id)
139 + .then_some(engine_id)
140 + .ok_or_else(|| Error::InvalidDataParallelRank {
141 + rank,
142 + num_engines: self.routing_per_engine.len() as u32,
143 + });
144 + }
145 +
146 + Ok(self
147 + .routing_per_engine
148 + .iter()
149 + .min_by_key(|(_, state)| state.routing_score())
150 + .map(|(engine_id, _)| engine_id.clone())
151 + .expect("request registry must contain at least one engine"))
152 + }
153 +
154 + /// Filter the given request IDs to the subset that are still tracked as
155 + /// active and can be aborted, grouped by engine.
156 + pub fn abortable_request_ids(&self, request_ids: &[String]) -> BTreeMap<EngineId, Vec<String>> {
157 + let mut by_engine = BTreeMap::new();
158 + for request_id in request_ids {
159 + let Some(tracked) = self.requests.get(request_id.as_str()) else {
160 + continue;
161 + };
162 + by_engine
163 + .entry(tracked.engine_id.clone())
164 + .or_insert_with(Vec::new)
165 + .push(request_id.clone());
166 + }
167 + by_engine
168 + }
169 +
170 + /// Obtain the stream sender for one output. If it indicates the request is
171 + /// finished, it will be removed from the registry.
172 + pub fn sender_for_output(&mut self, output: &EngineCoreOutput) -> Option<OutputSender> {
173 + if output.finished() {
174 + self.remove(output.request_id.as_str()).map(|tracked| tracked.0)
175 + } else {
176 + self.requests
177 + .get(output.request_id.as_str())
178 + .map(|tracked| tracked.sender.clone())
179 + }
180 + }
181 +
182 + /// Remove a batch of requests that have finished or aborted, returning
183 + /// their stream senders.
184 + pub fn finish_many<'a>(
185 + &mut self,
186 + request_ids: impl IntoIterator<Item = &'a String>,
187 + ) -> Vec<OutputSender> {
188 + request_ids
189 + .into_iter()
190 + .filter_map(|request_id| self.remove(request_id.as_str()).map(|tracked| tracked.0))
191 + .collect()
192 + }
193 +
194 + /// Apply one scheduler stats update for the given engine to the local
195 + /// routing state. Returns `false` if the engine is unknown to the
196 + /// client.
197 + pub fn apply_scheduler_stats(&mut self, engine_index: u32, stats: &SchedulerStats) -> bool {
198 + self.apply_scheduler_counts(
199 + engine_index,
200 + EngineLoadSnapshot {
201 + waiting: stats.num_waiting_reqs as usize,
202 + running: stats.num_running_reqs as usize,
203 + },
204 + )
205 + }
206 +
207 + /// Mark the registry as closed, detach and return all tracked senders.
208 + pub fn close(&mut self) -> Vec<OutputSender> {
209 + if self.closed {
210 + return Vec::new();
211 + }
212 +
213 + self.closed = true;
214 + std::mem::take(&mut self.requests)
215 + .into_values()
216 + .map(|tracked| tracked.sender)
217 + .collect()
218 + }
219 +
220 + /// Remove one request from the local registry. Returns the tracked entry if
221 + /// it exists.
222 + #[must_use]
223 + pub fn remove(&mut self, request_id: &str) -> Option<(OutputSender, EngineId)> {
224 + let tracked = self.requests.remove(request_id)?;
225 + self.routing_per_engine
226 + .get_mut(&tracked.engine_id)
227 + .expect("request registry must track all known engines")
228 + .inflight -= 1;
229 + Some((tracked.sender, tracked.engine_id))
230 + }
231 +
232 + fn apply_scheduler_counts(&mut self, engine_index: u32, next: EngineLoadSnapshot) -> bool {
233 + let engine_id = EngineId::from_engine_index(engine_index);
234 + let Some(state) = self.routing_per_engine.get_mut(&engine_id) else {
235 + return false;
236 + };
237 +
238 + let previous = state.last_scheduler_stats;
239 + if previous != Some(next) {
240 + trace!(
241 + ?engine_id,
242 + previous_waiting = previous.map(|stats| stats.waiting),
243 + previous_running = previous.map(|stats| stats.running),
244 + waiting = next.waiting,
245 + running = next.running,
246 + "updated scheduler routing counts",
247 + );
248 + }
249 +
250 + state.apply_scheduler_counts(next);
251 + true
252 + }
253 +
254 + #[cfg(test)]
255 + pub fn contains(&self, request_id: &str) -> bool {
256 + self.requests.contains_key(request_id)
257 + }
258 +
259 + pub fn is_closed(&self) -> bool {
260 + self.closed
261 + }
262 +}
263 +
264 +/// Internal registry for tracking active utility calls and their waiting
265 +/// receivers.
266 +#[derive(Debug)]
267 +pub struct UtilityRegistry {
268 + closed: bool,
269 + next_call_id: AtomicI64,
270 + utility_calls: BTreeMap<i64, UtilitySender>,
271 +}
272 +
273 +impl Default for UtilityRegistry {
274 + fn default() -> Self {
275 + Self {
276 + closed: false,
277 + next_call_id: AtomicI64::new(1),
278 + utility_calls: BTreeMap::default(),
279 + }
280 + }
281 +}
282 +
283 +impl UtilityRegistry {
284 + /// Allocate the next utility `call_id` and register a newly added utility
285 + /// call.
286 + pub fn allocate_and_register(&mut self) -> (i64, UtilityReceiver) {
287 + let call_id = self.next_call_id.fetch_add(1, Ordering::Relaxed);
288 + let (tx, rx) = oneshot::channel();
289 + self.utility_calls.insert(call_id, tx);
290 + (call_id, rx)
291 + }
292 +
293 + /// Resolve a utility output to its waiting receiver.
294 + pub fn resolve(&mut self, call_id: &i64) -> Option<UtilitySender> {
295 + self.utility_calls.remove(call_id)
296 + }
297 +
298 + /// Drop a batch of registered utility calls without delivering a result.
299 + /// Used to roll back allocations when the dispatch fan-out fails before
300 + /// every engine could accept the request.
301 + pub fn unregister_many(&mut self, call_ids: impl IntoIterator<Item = i64>) {
302 + for call_id in call_ids {
303 + self.utility_calls.remove(&call_id);
304 + }
305 + }
306 +
307 + /// Mark the registry as closed, detach and return all tracked senders.
308 + pub fn close(&mut self) -> Vec<UtilitySender> {
309 + if self.closed {
310 + return Vec::new();
311 + }
312 +
313 + self.closed = true;
314 + std::mem::take(&mut self.utility_calls).into_values().collect()
315 + }
316 +
317 + #[cfg(test)]
318 + pub fn contains(&self, call_id: i64) -> bool {
319 + self.utility_calls.contains_key(&call_id)
320 + }
321 +
322 + pub fn is_closed(&self) -> bool {
323 + self.closed
324 + }
325 +}
326 +
327 +#[cfg(test)]
328 +mod tests {
329 + use super::{EngineRoutingState, RequestRegistry, UtilityRegistry};
330 + use crate::EngineId;
331 + use crate::client::state::EngineLoadSnapshot;
332 + use crate::protocol::{EngineCoreFinishReason, EngineCoreOutput};
333 + use crate::transport::ConnectedEngine;
334 +
335 + #[test]
336 + fn registry_rejects_duplicate_request_ids() {
337 + let mut registry = RequestRegistry::new(&[ConnectedEngine {
338 + engine_id: EngineId::from(b"engine-0"),
339 + ready_response: None,
340 + }]);
341 + registry.register("req-1".to_string(), None).unwrap();
342 + let error = registry.register("req-1".to_string(), None).unwrap_err();
343 + assert!(matches!(
344 + error,
345 + crate::error::Error::DuplicateRequestId { request_id } if request_id == "req-1"
346 + ));
347 + }
348 +
349 + #[test]
350 + fn registry_removes_finished_request_on_output() {
351 + let mut registry = RequestRegistry::new(&[ConnectedEngine {
352 + engine_id: EngineId::from(b"engine-0"),
353 + ready_response: None,
354 + }]);
355 + registry.register("req-1".to_string(), None).unwrap();
356 +
357 + let sender = registry.sender_for_output(&EngineCoreOutput {
358 + request_id: "req-1".to_string(),
359 + finish_reason: Some(EngineCoreFinishReason::Length),
360 + ..Default::default()
361 + });
362 +
363 + assert!(sender.is_some());
364 + assert!(!registry.contains("req-1"));
365 + }
366 +
367 + #[test]
368 + fn registry_closes_all_requests_on_failure() {
369 + let mut registry = RequestRegistry::new(&[ConnectedEngine {
370 + engine_id: EngineId::from(b"engine-0"),
371 + ready_response: None,
372 + }]);
373 + registry.register("req-1".to_string(), None).unwrap();
374 + registry.register("req-2".to_string(), None).unwrap();
375 +
376 + let senders = registry.close();
377 +
378 + assert_eq!(senders.len(), 2);
379 + assert!(registry.is_closed());
380 + }
381 +
382 + #[test]
383 + fn registry_tracks_engine_id_per_request() {
384 + let engine_0 = EngineId::from_engine_index(0);
385 + let engine_1 = EngineId::from_engine_index(1);
386 + let mut registry = RequestRegistry::new(&[
387 + ConnectedEngine {
388 + engine_id: engine_0.clone(),
389 + ready_response: None,
390 + },
391 + ConnectedEngine {
392 + engine_id: engine_1.clone(),
393 + ready_response: None,
394 + },
395 + ]);
396 + let (chosen_0, _) = registry.register("req-1".to_string(), None).unwrap();
397 + let (chosen_1, _) = registry.register("req-2".to_string(), None).unwrap();
398 + let (chosen_0_again, _) = registry.register("req-3".to_string(), None).unwrap();
399 +
400 + assert_eq!(chosen_0, engine_0);
401 + assert_eq!(chosen_1, engine_1);
402 + assert_eq!(chosen_0_again, engine_0);
403 +
404 + let grouped = registry.abortable_request_ids(&[
405 + "req-1".to_string(),
406 + "req-2".to_string(),
407 + "req-3".to_string(),
408 + ]);
409 + assert_eq!(
410 + grouped.get(&engine_0).unwrap(),
411 + &vec!["req-1".to_string(), "req-3".to_string()]
412 + );
413 + assert_eq!(grouped.get(&engine_1).unwrap(), &vec!["req-2".to_string()]);
414 + }
415 +
416 + #[test]
417 + fn registry_uses_inflight_as_waiting_fallback_before_stats_arrive() {
418 + let engine_0 = EngineId::from_engine_index(0);
419 + let engine_1 = EngineId::from_engine_index(1);
420 + let mut registry = RequestRegistry::new(&[
421 + ConnectedEngine {
422 + engine_id: engine_0.clone(),
423 + ready_response: None,
424 + },
425 + ConnectedEngine {
426 + engine_id: engine_1.clone(),
427 + ready_response: None,
428 + },
429 + ]);
430 +
431 + let (chosen_0, _) = registry.register("req-1".to_string(), None).unwrap();
432 + let (chosen_1, _) = registry.register("req-2".to_string(), None).unwrap();
433 + let (chosen_0_again, _) = registry.register("req-3".to_string(), None).unwrap();
434 +
435 + assert_eq!(chosen_0, engine_0);
436 + assert_eq!(chosen_1, engine_1);
437 + assert_eq!(chosen_0_again, engine_0);
438 + }
439 +
440 + #[test]
441 + fn routing_score_uses_inflight_before_stats_arrive() {
442 + let state = EngineRoutingState {
443 + inflight: 3,
444 + last_scheduler_stats: None,
445 + };
446 +
447 + assert_eq!(state.routing_score(), 3);
448 + }
449 +
450 + #[test]
451 + fn routing_score_uses_inflight_as_scheduler_stats_lower_bound() {
452 + let state = EngineRoutingState {
453 + inflight: 7,
454 + last_scheduler_stats: Some(EngineLoadSnapshot {
455 + waiting: 0,
456 + running: 2,
457 + }),
458 + };
459 +
460 + assert_eq!(state.routing_score(), 7);
461 + }
462 +
463 + #[test]
464 + fn routing_score_keeps_extra_waiting_penalty() {
465 + let state = EngineRoutingState {
466 + inflight: 1,
467 + last_scheduler_stats: Some(EngineLoadSnapshot {
468 + waiting: 3,
469 + running: 2,
470 + }),
471 + };
472 +
473 + assert_eq!(state.routing_score(), 14);
474 + }
475 +
476 + #[test]
477 + fn registry_prefers_real_scheduler_stats_over_inflight() {
478 + let engine_0 = EngineId::from_engine_index(0);
479 + let engine_1 = EngineId::from_engine_index(1);
480 + let mut registry = RequestRegistry::new(&[
481 + ConnectedEngine {
482 + engine_id: engine_0.clone(),
483 + ready_response: None,
484 + },
485 + ConnectedEngine {
486 + engine_id: engine_1.clone(),
487 + ready_response: None,
488 + },
489 + ]);
490 +
491 + assert!(registry.apply_scheduler_counts(
492 + 0,
493 + EngineLoadSnapshot {
494 + waiting: 3,
495 + running: 2
496 + }
497 + ));
498 + assert!(registry.apply_scheduler_counts(
499 + 1,
500 + EngineLoadSnapshot {
501 + waiting: 0,
502 + running: 1
503 + }
504 + ));
505 +
506 + let (chosen, _) = registry.register("req-stats".to_string(), None).unwrap();
507 + assert_eq!(chosen, engine_1);
508 + }
509 +
510 + #[test]
511 + fn register_with_data_parallel_rank_routes_to_specified_engine() {
512 + let engine_0 = EngineId::from_engine_index(0);
513 + let engine_1 = EngineId::from_engine_index(1);
514 + let engine_2 = EngineId::from_engine_index(2);
515 + let mut registry = RequestRegistry::new(&[
516 + ConnectedEngine {
517 + engine_id: engine_0.clone(),
518 + ready_response: None,
519 + },
520 + ConnectedEngine {
521 + engine_id: engine_1.clone(),
522 + ready_response: None,
523 + },
524 + ConnectedEngine {
525 + engine_id: engine_2.clone(),
526 + ready_response: None,
527 + },
528 + ]);
529 +
530 + // Explicitly target rank 2 (third engine).
531 + let (chosen, _) = registry.register("req-1".to_string(), Some(2)).unwrap();
532 + assert_eq!(chosen, engine_2);
533 +
534 + // Explicitly target rank 0 (first engine).
535 + let (chosen, _) = registry.register("req-2".to_string(), Some(0)).unwrap();
536 + assert_eq!(chosen, engine_0);
537 +
538 + // Explicitly target rank 1.
539 + let (chosen, _) = registry.register("req-3".to_string(), Some(1)).unwrap();
540 + assert_eq!(chosen, engine_1);
541 + }
542 +
543 + #[test]
544 + fn register_with_data_parallel_rank_bypasses_load_balancing() {
545 + let engine_0 = EngineId::from_engine_index(0);
546 + let engine_1 = EngineId::from_engine_index(1);
547 + let mut registry = RequestRegistry::new(&[
548 + ConnectedEngine {
549 + engine_id: engine_0.clone(),
550 + ready_response: None,
551 + },
552 + ConnectedEngine {
553 + engine_id: engine_1.clone(),
554 + ready_response: None,
555 + },
556 + ]);
557 +
558 + // Load-balance: first two go to engine_0 and engine_1.
559 + registry.register("req-lb-0".to_string(), None).unwrap();
560 +
561 + // Now engine_0 has 1 in-flight. Without dp_rank, next would go to engine_1.
562 + // But with dp_rank=0, it should still go to engine_0.
563 + let (chosen, _) = registry.register("req-dp".to_string(), Some(0)).unwrap();
564 + assert_eq!(chosen, engine_0);
565 + }
566 +
567 + #[test]
568 + fn register_with_out_of_range_rank_returns_error() {
569 + let mut registry = RequestRegistry::new(&[
570 + ConnectedEngine {
571 + engine_id: EngineId::from_engine_index(0),
572 + ready_response: None,
573 + },
574 + ConnectedEngine {
575 + engine_id: EngineId::from_engine_index(1),
576 + ready_response: None,
577 + },
578 + ]);
579 +
580 + let error = registry.register("req-1".to_string(), Some(2)).unwrap_err();
581 + assert!(matches!(
582 + error,
583 + crate::error::Error::InvalidDataParallelRank {
584 + rank: 2,
585 + num_engines: 2,
586 + }
587 + ));
588 + }
589 +
590 + #[test]
591 + fn register_with_rank_on_single_engine_only_accepts_zero() {
592 + let engine_0 = EngineId::from_engine_index(0);
593 + let mut registry = RequestRegistry::new(&[ConnectedEngine {
594 + engine_id: engine_0.clone(),
595 + ready_response: None,
596 + }]);
597 +
598 + let (chosen, _) = registry.register("req-ok".to_string(), Some(0)).unwrap();
599 + assert_eq!(chosen, engine_0);
600 +
601 + let error = registry.register("req-bad".to_string(), Some(1)).unwrap_err();
602 + assert!(matches!(
603 + error,
604 + crate::error::Error::InvalidDataParallelRank {
605 + rank: 1,
606 + num_engines: 1,
607 + }
608 + ));
609 + }
610 +
611 + #[test]
612 + fn utility_registry_tracks_and_removes_call_ids() {
613 + let mut registry = UtilityRegistry::default();
614 + let (call_id_1, _) = registry.allocate_and_register();
615 + let (call_id_2, _) = registry.allocate_and_register();
616 +
617 + assert_eq!(call_id_1, 1);
618 + assert_eq!(call_id_2, 2);
619 + assert!(registry.contains(1));
620 + assert!(registry.contains(2));
621 + assert!(registry.resolve(&1).is_some());
622 + assert!(!registry.contains(1));
623 + assert!(registry.contains(2));
624 + }
625 +
626 + #[test]
627 + fn utility_registry_closes_all_waiters_on_failure() {
628 + let mut registry = UtilityRegistry::default();
629 + registry.allocate_and_register();
630 + registry.allocate_and_register();
631 +
632 + let senders = registry.close();
633 +
634 + assert_eq!(senders.len(), 2);
635 + assert!(!registry.contains(1));
636 + assert!(!registry.contains(2));
637 + assert!(registry.is_closed());
638 + }
639 +
640 + #[test]
641 + fn utility_registry_unregister_many_drops_pending_calls() {
642 + use tokio::sync::oneshot::error::TryRecvError;
643 +
644 + let mut registry = UtilityRegistry::default();
645 + let (call_id_1, mut rx_1) = registry.allocate_and_register();
646 + let (call_id_2, mut rx_2) = registry.allocate_and_register();
647 + let (call_id_3, _rx_3) = registry.allocate_and_register();
648 +
649 + // Drop two of the three allocated calls; the third stays pending.
650 + registry.unregister_many([call_id_1, call_id_2]);
651 +
652 + assert!(!registry.contains(call_id_1));
653 + assert!(!registry.contains(call_id_2));
654 + assert!(registry.contains(call_id_3));
655 + // The receivers must observe the sender being dropped (channel closed).
656 + assert!(matches!(rx_1.try_recv(), Err(TryRecvError::Closed)));
657 + assert!(matches!(rx_2.try_recv(), Err(TryRecvError::Closed)));
658 + }
659 +
660 + #[test]
661 + fn utility_registry_unregister_many_ignores_unknown_call_ids() {
662 + let mut registry = UtilityRegistry::default();
663 + let (call_id, _rx) = registry.allocate_and_register();
664 +
665 + // Unknown call ids are silently ignored — caller doesn't care which were live.
666 + registry.unregister_many([call_id, 42, 9999]);
667 +
668 + assert!(!registry.contains(call_id));
669 + }
670 +}
rust/src/engine-core-client/src/client/stream.rs new
+153
@@ -0,0 +1,153 @@
1 +use std::ops::Deref;
2 +use std::pin::Pin;
3 +use std::task::{Context, Poll};
4 +
5 +use futures::Stream;
6 +use futures::stream::FusedStream;
7 +use thiserror_ext::AsReport as _;
8 +use tokio::sync::mpsc;
9 +use tracing::{debug, error, warn};
10 +
11 +use crate::client::AbortRequest;
12 +use crate::client::state::OutputReceiver;
13 +use crate::protocol::{EngineCoreFinishReason, EngineCoreOutput};
14 +use crate::{AbortCause, Error, Result};
15 +
16 +#[derive(Debug, Clone, Copy, PartialEq, Eq)]
17 +enum State {
18 + Running,
19 + Finished,
20 + ClosedWithError,
21 + UnexpectedClose,
22 +}
23 +
24 +/// One request-scoped engine-core output plus the enclosing batch metadata.
25 +#[derive(Debug, Clone, PartialEq)]
26 +pub struct EngineCoreStreamOutput {
27 + pub engine_index: u32,
28 + pub timestamp: f64,
29 + pub output: EngineCoreOutput,
30 +}
31 +
32 +impl Deref for EngineCoreStreamOutput {
33 + type Target = EngineCoreOutput;
34 +
35 + fn deref(&self) -> &Self::Target {
36 + &self.output
37 + }
38 +}
39 +
40 +/// Stream of raw engine-core outputs for one request.
41 +///
42 +/// The stream yields only [`EngineCoreStreamOutput`] values whose embedded
43 +/// output `request_id` matches the originating `add_request()` call. Normal
44 +/// request completion is expected to include a final output object whose
45 +/// `finish_reason` is non-`None`.
46 +pub struct EngineCoreOutputStream {
47 + request_id: String,
48 + abort_tx: mpsc::UnboundedSender<AbortRequest>,
49 + state: State,
50 + rx: OutputReceiver,
51 +}
52 +
53 +impl EngineCoreOutputStream {
54 + pub(crate) fn new(
55 + request_id: String,
56 + abort_tx: mpsc::UnboundedSender<AbortRequest>,
57 + rx: OutputReceiver,
58 + ) -> Self {
59 + Self {
60 + request_id,
61 + abort_tx,
62 + state: State::Running,
63 + rx,
64 + }
65 + }
66 +
67 + /// Return the engine-core `request_id` bound to this stream.
68 + pub fn request_id(&self) -> &str {
69 + &self.request_id
70 + }
71 +}
72 +
73 +impl Stream for EngineCoreOutputStream {
74 + type Item = Result<EngineCoreStreamOutput>;
75 +
76 + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
77 + if self.is_terminated() {
78 + return Poll::Ready(None);
79 + }
80 +
81 + match Pin::new(&mut self.rx).poll_recv(cx) {
82 + Poll::Pending => Poll::Pending,
83 + Poll::Ready(Some(item)) => {
84 + match &item {
85 + Ok(output) => {
86 + // If the output indicates the request is finished, mark the stream as
87 + // terminated with cleanly-finished state and expect no more outputs to
88 + // come.
89 + if output.finished() {
90 + if output.finish_reason == Some(EngineCoreFinishReason::Error) {
91 + error!(
92 + self.request_id,
93 + "request failed with an internal error during generation"
94 + );
95 + }
96 + debug!(self.request_id, "request completed via final output");
97 + self.state = State::Finished;
98 + }
99 + }
100 + Err(error) => {
101 + // If we get an error from the output stream, mark the stream as terminated
102 + // with an error.
103 + warn!(self.request_id, error = %error.as_report(), "request encountered an error");
104 + self.state = State::ClosedWithError;
105 + }
106 + }
107 + Poll::Ready(Some(item))
108 + }
109 + Poll::Ready(None) => {
110 + // If we get a `None` without seeing a finished output, this is an unexpected
111 + // close from the engine side. Mark the stream as terminated
112 + // with an unexpected close state and send an error down the
113 + // stream to notify the caller.
114 + warn!(self.request_id, "request stream closed unexpectedly");
115 + self.state = State::UnexpectedClose;
116 +
117 + Poll::Ready(Some(Err(Error::RequestStreamClosed {
118 + request_id: self.request_id.clone(),
119 + })))
120 + }
121 + }
122 + }
123 +}
124 +
125 +impl FusedStream for EngineCoreOutputStream {
126 + fn is_terminated(&self) -> bool {
127 + !matches!(self.state, State::Running)
128 + }
129 +}
130 +
131 +impl Drop for EngineCoreOutputStream {
132 + fn drop(&mut self) {
133 + if self.is_terminated() {
134 + // If it's terminated, it means that the request either finished cleanly, or
135 + // encountered an error or unexpected close from the engine. In any
136 + // case, the request stream is already considered inactive and
137 + // there's no need to abort it on the engine side.
138 + return;
139 + }
140 +
141 + let abort_req = AbortRequest {
142 + request_id: self.request_id.clone(),
143 + cause: AbortCause::current(),
144 + };
145 +
146 + if self.abort_tx.send(abort_req).is_err() {
147 + warn!(
148 + request_id = self.request_id,
149 + "auto-abort worker already shut down; skip auto-abort"
150 + );
151 + }
152 + }
153 +}
rust/src/engine-core-client/src/coordinator/bootstrap.rs new
+91
@@ -0,0 +1,91 @@
1 +use std::time::Duration;
2 +
3 +use bytes::Bytes;
4 +use zeromq::prelude::{Socket, SocketRecv, SocketSend};
5 +use zeromq::{PullSocket, XPubSocket, ZmqMessage};
6 +
7 +use crate::error::{Error, Result, bail_unexpected_handshake_message};
8 +
9 +/// Engine-facing sockets owned by the in-process coordinator.
10 +pub(crate) struct CoordinatorBootstrap {
11 + pub input_address: String,
12 + pub output_address: String,
13 + pub input_socket: XPubSocket,
14 + pub output_socket: PullSocket,
15 +}
16 +
17 +impl CoordinatorBootstrap {
18 + /// Bind the engine-facing coordinator sockets on the given host.
19 + pub(crate) async fn bind(local_host: &str) -> Result<Self> {
20 + let mut input_socket = XPubSocket::new();
21 + let input_address = input_socket.bind(&format!("tcp://{local_host}:0")).await?.to_string();
22 +
23 + let mut output_socket = PullSocket::new();
24 + let output_address =
25 + output_socket.bind(&format!("tcp://{local_host}:0")).await?.to_string();
26 +
27 + Ok(Self {
28 + input_address,
29 + output_address,
30 + input_socket,
31 + output_socket,
32 + })
33 + }
34 +
35 + /// Complete the engine-facing startup gate before engines are allowed to
36 + /// send handshake READY.
37 + pub(crate) async fn wait_for_startup_gate(
38 + &mut self,
39 + engine_count: usize,
40 + ready_timeout: Duration,
41 + ) -> Result<()> {
42 + wait_for_engine_subscriptions(&mut self.input_socket, engine_count, ready_timeout).await?;
43 + send_ready_to_engines(&mut self.input_socket).await?;
44 + Ok(())
45 + }
46 +}
47 +
48 +/// Wait until all engines subscribe to the coordinator broadcast socket.
49 +async fn wait_for_engine_subscriptions(
50 + input_socket: &mut XPubSocket,
51 + engine_count: usize,
52 + ready_timeout: Duration,
53 +) -> Result<()> {
54 + let mut received = 0;
55 + while received < engine_count {
56 + let message =
57 + tokio::time::timeout(ready_timeout, input_socket.recv()).await.map_err(|_| {
58 + Error::HandshakeTimeout {
59 + stage: "coordinator engine subscriptions",
60 + timeout: ready_timeout,
61 + }
62 + })??;
63 + if message.len() != 1 {
64 + bail_unexpected_handshake_message!(
65 + "expected 1 frame for coordinator subscription, got {}",
66 + message.len()
67 + );
68 + }
69 +
70 + let frame = message
71 + .into_vec()
72 + .into_iter()
73 + .next()
74 + .expect("single-frame coordinator subscription message");
75 + if frame.as_ref() != [0x01] {
76 + bail_unexpected_handshake_message!(
77 + "expected coordinator subscription frame [0x01], got {:?}",
78 + frame.as_ref()
79 + );
80 + }
81 + received += 1;
82 + }
83 +
84 + Ok(())
85 +}
86 +
87 +/// Send the coordinator READY marker to all subscribed engines.
88 +async fn send_ready_to_engines(input_socket: &mut XPubSocket) -> Result<()> {
89 + input_socket.send(ZmqMessage::from(Bytes::from_static(b"READY"))).await?;
90 + Ok(())
91 +}
rust/src/engine-core-client/src/coordinator/external.rs new
+167
@@ -0,0 +1,167 @@
1 +use std::sync::Arc;
2 +
3 +use serde_tuple::{Deserialize_tuple, Serialize_tuple};
4 +use thiserror_ext::AsReport;
5 +use tokio::sync::mpsc;
6 +use tracing::{debug, warn};
7 +use zeromq::prelude::{SocketRecv, SocketSend};
8 +use zeromq::{XSubSocket, ZmqMessage};
9 +
10 +use crate::client::imp::ClientInner;
11 +use crate::coordinator::handle::{CoordinatorCommand, CoordinatorState};
12 +use crate::error::{Error, Result, bail_unexpected_coordinator_output};
13 +use crate::protocol::{OpaqueValue, decode_msgpack, encode_msgpack};
14 +
15 +/// Frontend-to-coordinator wakeup message sent when the first request arrives
16 +/// while all engines are paused.
17 +///
18 +/// This matches the frontend-side msgpack tuple sent by Python
19 +/// `DPAsyncMPClient._ensure_stats_update_task` to the coordinator front socket.
20 +///
21 +/// Original Python definition:
22 +/// <https://github.com/vllm-project/vllm/blob/694449050f8dac3d9853e97e518b4a43ec52106a/vllm/v1/engine/core_client.py#L1230-L1236>
23 +#[derive(Debug, Clone, PartialEq, Eq, Serialize_tuple, Deserialize_tuple)]
24 +struct CoordinatorWakeupMessage {
25 + /// Engine index that already has the triggering request and should be
26 + /// excluded from the coordinator's `START_DP_WAVE` rebroadcast.
27 + exclude_engine_index: u32,
28 + /// DP wave number observed by the frontend when the request was admitted.
29 + wave: u32,
30 +}
31 +
32 +/// Coordinator-to-frontend state publish received on the front-side coordinator
33 +/// socket.
34 +///
35 +/// This matches the msgpack tuple periodically published by Python
36 +/// `DPCoordinatorProc.run_coordinator` to all connected frontends.
37 +///
38 +/// Original Python definitions:
39 +/// <https://github.com/vllm-project/vllm/blob/694449050f8dac3d9853e97e518b4a43ec52106a/vllm/v1/engine/coordinator.py#L282-L283>
40 +/// <https://github.com/vllm-project/vllm/blob/694449050f8dac3d9853e97e518b4a43ec52106a/vllm/v1/engine/coordinator.py#L445-L447>
41 +#[derive(Debug, Clone, PartialEq, Deserialize_tuple)]
42 +struct CoordinatorStateUpdate {
43 + /// Global per-engine request counts published by the coordinator.
44 + ///
45 + /// The Rust bootstrapped external-coordinator path preserves this field for
46 + /// wire compatibility but intentionally ignores it for routing decisions.
47 + counts: OpaqueValue,
48 + /// Current global DP wave number stamped onto newly admitted requests.
49 + wave: u32,
50 + /// Whether engines are currently running (`true`) or paused (`false`).
51 + engines_running: bool,
52 +}
53 +
54 +/// Background half of an external Python-owned coordinator connection.
55 +///
56 +/// This owns the command receiver and one frontend-facing XSUB socket. It
57 +/// mirrors the subset of Python's coordinator protocol needed by the Rust
58 +/// bootstrapped frontend: receive `(counts, wave, running)` publishes, ignore
59 +/// `counts`, and send `(exclude_engine_index, wave)` wakeup messages when the
60 +/// first request arrives while engines are paused.
61 +pub(crate) struct ExternalCoordinatorService {
62 + state: Arc<CoordinatorState>,
63 + command_rx: mpsc::UnboundedReceiver<CoordinatorCommand>,
64 + socket: XSubSocket,
65 +}
66 +
67 +impl ExternalCoordinatorService {
68 + pub(super) fn new(
69 + state: Arc<CoordinatorState>,
70 + command_rx: mpsc::UnboundedReceiver<CoordinatorCommand>,
71 + socket: XSubSocket,
72 + ) -> Self {
73 + Self {
74 + state,
75 + command_rx,
76 + socket,
77 + }
78 + }
79 +
80 + /// Apply one frontend-originated command to the external coordinator state
81 + /// machine.
82 + async fn handle_command(&mut self, command: CoordinatorCommand) -> Result<()> {
83 + match command {
84 + CoordinatorCommand::FirstRequest {
85 + target_engine_id,
86 + wave,
87 + } => {
88 + let target_engine_index = target_engine_id.engine_index().ok_or_else(|| {
89 + Error::UnsupportedCoordinatorEngineId {
90 + engine_id: target_engine_id.to_vec(),
91 + }
92 + })?;
93 + debug!(
94 + wave,
95 + exclude_engine_index = target_engine_index,
96 + "notifying external coordinator about first request while engines were paused"
97 + );
98 + let payload = encode_msgpack(&CoordinatorWakeupMessage {
99 + exclude_engine_index: target_engine_index,
100 + wave,
101 + })?;
102 + self.socket.send(ZmqMessage::from(payload)).await?;
103 + }
104 + }
105 + Ok(())
106 + }
107 +
108 + /// Apply one publish received from the xsub socket containing a coordinator
109 + /// state update.
110 + async fn handle_publish(&mut self, message: ZmqMessage) -> Result<()> {
111 + let frames = message.into_vec();
112 + if frames.len() != 1 {
113 + bail_unexpected_coordinator_output!(
114 + "received malformed external coordinator publish with {} frame(s)",
115 + frames.len()
116 + );
117 + }
118 +
119 + let update: CoordinatorStateUpdate = decode_msgpack(&frames[0])?;
120 +
121 + let mut state = self.state.lock();
122 + let previous_wave = state.current_wave;
123 + let previous_engines_running = state.engines_running;
124 + state.current_wave = update.wave;
125 + state.engines_running = update.engines_running;
126 + debug!(
127 + previous_wave,
128 + wave = update.wave,
129 + previous_engines_running,
130 + engines_running = update.engines_running,
131 + "applied external coordinator state update"
132 + );
133 + Ok(())
134 + }
135 +
136 + /// Drive the coordinator event loop until either side of the control plane
137 + /// is closed or a fatal error is observed.
138 + pub(crate) async fn run(mut self, inner: Arc<ClientInner>) {
139 + let result: Result<()> = async {
140 + loop {
141 + tokio::select! {
142 + // Received frontend-originated command from the handle.
143 + command = self.command_rx.recv() => {
144 + let Some(command) = command else {
145 + warn!("external coordinator command channel closed, shutting down service");
146 + return Ok(());
147 + };
148 + self.handle_command(command).await?;
149 + }
150 + // Received publish from the external coordinator socket.
151 + publish = self.socket.recv() => {
152 + let publish = publish.map_err(Error::from)?;
153 + self.handle_publish(publish).await?;
154 + }
155 + }
156 + }
157 + }
158 + .await;
159 + let Err(error) = result else { return };
160 +
161 + warn!(
162 + error = %error.as_report(),
163 + "external coordinator service exiting with error"
164 + );
165 + inner.close_registries(Arc::new(error));
166 + }
167 +}
rust/src/engine-core-client/src/coordinator/handle.rs new
+123
@@ -0,0 +1,123 @@
1 +use std::sync::Arc;
2 +
3 +use parking_lot::Mutex;
4 +use tokio::sync::mpsc;
5 +use zeromq::prelude::Socket;
6 +use zeromq::{XPubSocket, XSubSocket};
7 +
8 +use crate::coordinator::external::ExternalCoordinatorService;
9 +use crate::coordinator::inproc::InProcCoordinatorRunner;
10 +use crate::error::{Error, Result, bail_control_closed};
11 +use crate::transport::EngineId;
12 +
13 +/// Snapshot to the coordinator state for request routing and stamping.
14 +#[derive(Debug, Clone, Copy, PartialEq, Eq)]
15 +pub(crate) struct CoordinatorStateSnapshot {
16 + /// The current DP wave, which will be stamped on outgoing requests.
17 + pub current_wave: u32,
18 + /// Whether the engines are currently running or paused, which determines if
19 + /// the frontend must trigger a new wave on the next request.
20 + pub engines_running: bool,
21 +}
22 +
23 +/// Shared in-process coordinator state.
24 +pub(crate) type CoordinatorState = Mutex<CoordinatorStateSnapshot>;
25 +
26 +/// Commands sent from the frontend request path into the background runner.
27 +#[derive(Debug)]
28 +pub(crate) enum CoordinatorCommand {
29 + /// The first request arrived while all engines were paused.
30 + ///
31 + /// The coordinator should broadcast `START_DP_WAVE` with the current wave
32 + /// and the target engine index as the excluded engine.
33 + FirstRequest {
34 + target_engine_id: EngineId,
35 + wave: u32,
36 + },
37 +}
38 +
39 +/// Frontend-facing coordinator handle used by `EngineCoreClient::call()`.
40 +///
41 +/// This side stays intentionally small: it can read the latest wave snapshot
42 +/// and enqueue a `FirstRequest` transition when the request path observes the
43 +/// system in the paused state.
44 +#[derive(Clone)]
45 +pub(crate) struct CoordinatorHandle {
46 + state: Arc<CoordinatorState>,
47 + command_tx: mpsc::UnboundedSender<CoordinatorCommand>,
48 +}
49 +
50 +impl CoordinatorHandle {
51 + fn new_parts() -> (
52 + Self,
53 + Arc<CoordinatorState>,
54 + mpsc::UnboundedReceiver<CoordinatorCommand>,
55 + ) {
56 + let state = Arc::new(Mutex::new(CoordinatorStateSnapshot {
57 + current_wave: 0,
58 + engines_running: false,
59 + }));
60 + let (command_tx, command_rx) = mpsc::unbounded_channel();
61 + (
62 + Self {
63 + state: state.clone(),
64 + command_tx,
65 + },
66 + state,
67 + command_rx,
68 + )
69 + }
70 +
71 + /// Build the paired frontend handle and background runner around one
72 + /// engine-facing coordinator broadcast socket.
73 + pub(crate) fn new_inproc(coordinator_input: XPubSocket) -> (Self, InProcCoordinatorRunner) {
74 + let (handle, state, command_rx) = Self::new_parts();
75 + (
76 + handle,
77 + InProcCoordinatorRunner::new(state, command_rx, coordinator_input),
78 + )
79 + }
80 +
81 + /// Build the paired frontend handle and background service around an
82 + /// external Python-owned frontend-side coordinator socket.
83 + pub(crate) async fn connect_external(
84 + coordinator_address: &str,
85 + ) -> Result<(Self, ExternalCoordinatorService)> {
86 + let (handle, state, command_rx) = Self::new_parts();
87 + let mut socket = XSubSocket::new();
88 + socket.connect(coordinator_address).await?;
89 + socket.subscribe("").await?;
90 + Ok((
91 + handle,
92 + ExternalCoordinatorService::new(state, command_rx, socket),
93 + ))
94 + }
95 +
96 + /// Snapshot the coordinator state for request routing and stamping.
97 + pub(crate) fn snapshot(&self) -> CoordinatorStateSnapshot {
98 + *self.state.lock()
99 + }
100 +
101 + /// Notify the runner that a new request arrived while engines were paused.
102 + ///
103 + /// The handle flips `engines_running` optimistically so concurrent request
104 + /// submissions coalesce behind one `START_DP_WAVE` broadcast instead of all
105 + /// trying to trigger the wave independently.
106 + pub(crate) fn notify_first_request(&self, target_engine_id: EngineId) -> Result<()> {
107 + let mut state = self.state.lock();
108 + if state.engines_running {
109 + return Ok(());
110 + }
111 +
112 + let command = CoordinatorCommand::FirstRequest {
113 + target_engine_id,
114 + wave: state.current_wave,
115 + };
116 + if self.command_tx.send(command).is_err() {
117 + bail_control_closed!("in-process coordinator command channel already shut down");
118 + }
119 +
120 + state.engines_running = true;
121 + Ok(())
122 + }
123 +}
rust/src/engine-core-client/src/coordinator/inproc.rs new
+204
@@ -0,0 +1,204 @@
1 +use std::sync::Arc;
2 +
3 +use serde_tuple::{Deserialize_tuple, Serialize_tuple};
4 +use thiserror_ext::AsReport;
5 +use tokio::sync::mpsc;
6 +use tracing::{debug, warn};
7 +use zeromq::prelude::SocketSend;
8 +use zeromq::{XPubSocket, ZmqMessage};
9 +
10 +use crate::client::imp::ClientInner;
11 +use crate::coordinator::handle::{CoordinatorCommand, CoordinatorState};
12 +use crate::error::{Error, Result, bail_unexpected_coordinator_output};
13 +use crate::protocol::{
14 + ClassifiedEngineCoreOutputs, DpControlMessage, EngineCoreOutputs, EngineCoreRequestType,
15 + encode_msgpack,
16 +};
17 +
18 +/// Coordinator-to-engine `START_DP_WAVE` control payload encoded on the
19 +/// engine-facing coordinator socket.
20 +///
21 +/// This matches the msgpack tuple broadcast by Python
22 +/// `DPCoordinatorProc._send_start_wave`.
23 +///
24 +/// Original Python definition:
25 +/// <https://github.com/vllm-project/vllm/blob/694449050f8dac3d9853e97e518b4a43ec52106a/vllm/v1/engine/coordinator.py#L453-L459>
26 +#[derive(Debug, Clone, PartialEq, Eq, Serialize_tuple, Deserialize_tuple)]
27 +struct StartDpWaveMessage {
28 + /// DP wave number that all engines should start processing.
29 + wave: u32,
30 + /// Engine index that already received the triggering request and should not
31 + /// receive an extra wakeup notification.
32 + exclude_engine_index: u32,
33 +}
34 +
35 +/// Background half of the in-process coordinator.
36 +///
37 +/// This owns the command receiver and the engine-facing coordinator input
38 +/// socket. It is the single place where wave transitions are serialized and
39 +/// where `START_DP_WAVE` broadcasts are emitted.
40 +pub(crate) struct InProcCoordinatorRunner {
41 + state: Arc<CoordinatorState>,
42 + command_rx: mpsc::UnboundedReceiver<CoordinatorCommand>,
43 + coordinator_input: XPubSocket,
44 +}
45 +
46 +impl InProcCoordinatorRunner {
47 + pub(super) fn new(
48 + state: Arc<CoordinatorState>,
49 + command_rx: mpsc::UnboundedReceiver<CoordinatorCommand>,
50 + coordinator_input: XPubSocket,
51 + ) -> Self {
52 + Self {
53 + state,
54 + command_rx,
55 + coordinator_input,
56 + }
57 + }
58 +
59 + /// Broadcast Python-compatible `START_DP_WAVE` to all connected engines.
60 + async fn broadcast_start_wave(&mut self, wave: u32, exclude_engine_index: u32) -> Result<()> {
61 + let payload = encode_msgpack(&StartDpWaveMessage {
62 + wave,
63 + exclude_engine_index,
64 + })?;
65 + self.coordinator_input
66 + .send(
67 + ZmqMessage::try_from(vec![
68 + EngineCoreRequestType::StartDpWave.to_frame(),
69 + payload.into(),
70 + ])
71 + .expect("coordinator START_DP_WAVE message must contain two frames"),
72 + )
73 + .await?;
74 + Ok(())
75 + }
76 +
77 + /// Apply one frontend-originated command to the coordinator state machine.
78 + async fn handle_command(&mut self, command: CoordinatorCommand) -> Result<()> {
79 + match command {
80 + CoordinatorCommand::FirstRequest {
81 + target_engine_id,
82 + wave,
83 + } => {
84 + let target_engine_index = target_engine_id.engine_index().ok_or_else(|| {
85 + Error::UnsupportedCoordinatorEngineId {
86 + engine_id: target_engine_id.to_vec(),
87 + }
88 + })?;
89 + self.state.lock().current_wave = wave;
90 + debug!(
91 + wave,
92 + exclude_engine_index = target_engine_index,
93 + "starting DP wave after first request while engines were paused"
94 + );
95 + self.broadcast_start_wave(wave, target_engine_index).await?;
96 + }
97 + }
98 + Ok(())
99 + }
100 +
101 + /// Apply one engine-originated control output to the coordinator state
102 + /// machine.
103 + async fn handle_outputs(&mut self, outputs: EngineCoreOutputs) -> Result<()> {
104 + match outputs.classify() {
105 + ClassifiedEngineCoreOutputs::RequestBatch(batch)
106 + if batch.outputs.is_empty() && batch.finished_requests.is_none() =>
107 + {
108 + // Stats-only output for coordinator.
109 + // Ignore since the Rust coordinator doesn't track stats for
110 + // routing decisions.
111 + }
112 + ClassifiedEngineCoreOutputs::DpControl {
113 + engine_index,
114 + control,
115 + ..
116 + } => match control {
117 + // The engines signals they completed the current wave and are now paused.
118 + // Advance the current wave and mark the state as paused.
119 + DpControlMessage::WaveComplete(wave) => {
120 + let mut state = self.state.lock();
121 + if wave >= state.current_wave {
122 + let next_wave = wave + 1;
123 + debug!(
124 + wave,
125 + next_wave,
126 + "DP wave finished; pausing engines and advancing coordinator state"
127 + );
128 + state.current_wave = wave + 1;
129 + state.engines_running = false;
130 + }
131 + }
132 + // An engine requests to start the wave.
133 + // Rebroadcast the wave to all engines except for the originated one.
134 + DpControlMessage::StartWave(wave) => {
135 + let should_broadcast = {
136 + let mut state = self.state.lock();
137 + if wave > state.current_wave
138 + || (wave == state.current_wave && !state.engines_running)
139 + {
140 + state.current_wave = wave;
141 + state.engines_running = true;
142 + true
143 + } else {
144 + false
145 + }
146 + };
147 + if should_broadcast {
148 + debug!(
149 + wave,
150 + exclude_engine_index = engine_index,
151 + "starting DP wave after stale-wave notification from engine"
152 + );
153 + self.broadcast_start_wave(wave, engine_index).await?;
154 + }
155 + }
156 + },
157 + other => {
158 + bail_unexpected_coordinator_output!(
159 + "received non-control output on coordinator path: {other:?}"
160 + );
161 + }
162 + }
163 + Ok(())
164 + }
165 +
166 + /// Drive the coordinator event loop until either side of the control plane
167 + /// is closed or a fatal error is observed.
168 + ///
169 + /// Any fatal error closes the main client registries so request streams and
170 + /// future calls observe a stable shutdown cause.
171 + pub(crate) async fn run(
172 + mut self,
173 + mut output_rx: mpsc::Receiver<Result<EngineCoreOutputs>>,
174 + inner: Arc<ClientInner>,
175 + ) {
176 + let result: Result<()> = async {
177 + loop {
178 + tokio::select! {
179 + // Received frontend-originated command from the handle.
180 + command = self.command_rx.recv() => {
181 + let Some(command) = command else {
182 + warn!("coordinator command channel closed, shutting down coordinator runner");
183 + return Ok(());
184 + };
185 + self.handle_command(command).await?;
186 + }
187 + // Received engine-originated control output from the coordinator socket.
188 + outputs = output_rx.recv() => {
189 + let Some(outputs) = outputs else {
190 + warn!("coordinator output channel closed, shutting down coordinator runner");
191 + return Ok(());
192 + };
193 + self.handle_outputs(outputs?).await?;
194 + }
195 + }
196 + }
197 + }
198 + .await;
199 + let Err(error) = result else { return };
200 +
201 + warn!(error = %error.as_report(), "coordinator runner exiting with error");
202 + inner.close_registries(Arc::new(error));
203 + }
204 +}
rust/src/engine-core-client/src/coordinator/mod.rs new
+7
@@ -0,0 +1,7 @@
1 +mod bootstrap;
2 +mod external;
3 +mod handle;
4 +mod inproc;
5 +
6 +pub(crate) use bootstrap::CoordinatorBootstrap;
7 +pub(crate) use handle::CoordinatorHandle;
rust/src/engine-core-client/src/error.rs new
+88
@@ -0,0 +1,88 @@
1 +use std::sync::Arc;
2 +use std::time::Duration;
3 +
4 +use thiserror::Error;
5 +use thiserror_ext::Macro;
6 +
7 +pub type Result<T> = std::result::Result<T, Error>;
8 +
9 +/// Public error type for the Rust engine-core client.
10 +#[derive(Debug, Error, Macro)]
11 +pub enum Error {
12 + #[error("messagepack encode failed for {target_type}: {message}")]
13 + Encode {
14 + target_type: &'static str,
15 + message: String,
16 + },
17 + #[error("messagepack decode failed for {target_type}: {message}")]
18 + Decode {
19 + target_type: &'static str,
20 + message: String,
21 + },
22 + #[error("messagepack value decode failed")]
23 + ValueDecode(#[from] rmpv::decode::Error),
24 + #[error("messagepack ext value decode failed: {message}")]
25 + ExtValueDecode { message: String },
26 + #[error("io error")]
27 + Io(#[from] std::io::Error),
28 + #[error("transport error")]
29 + Transport(#[from] zeromq::ZmqError),
30 + #[error("engine core reported fatal failure")]
31 + EngineCoreDead,
32 + #[error("startup handshake timed out while waiting for {stage} after {timeout:?}")]
33 + HandshakeTimeout {
34 + stage: &'static str,
35 + timeout: Duration,
36 + },
37 + #[error("engine input registration timed out after {timeout:?}")]
38 + InputRegistrationTimeout { timeout: Duration },
39 + #[error("unexpected engine id in startup handshake: expected {expected:?}, got {actual:?}")]
40 + UnexpectedHandshakeIdentity { expected: Vec<u8>, actual: Vec<u8> },
41 + #[error("unexpected startup handshake message: {message}")]
42 + UnexpectedHandshakeMessage { message: String },
43 + #[error("unexpected non-control output on coordinator path: {message}")]
44 + UnexpectedCoordinatorOutput { message: String },
45 + #[error("unexpected output on main dispatcher path: {message}")]
46 + UnexpectedDispatcherOutput { message: String },
47 + #[error("coordinator requires a Python-compatible two-byte engine id, got {engine_id:?}")]
48 + UnsupportedCoordinatorEngineId { engine_id: Vec<u8> },
49 + #[error("unsupported auxiliary frame(s): expected 1 frame, got {frame_count}")]
50 + UnsupportedAuxFrames { frame_count: usize },
51 + #[error("external coordinator mode is not implemented yet")]
52 + UnsupportedExternalCoordinator,
53 + #[error("unsupported field `{field}` in {context}")]
54 + UnsupportedField {
55 + context: &'static str,
56 + field: &'static str,
57 + },
58 + #[error("engine control channel closed unexpectedly: {message}")]
59 + ControlClosed { message: String },
60 + #[error("request `{request_id}` is already in flight")]
61 + DuplicateRequestId { request_id: String },
62 + #[error("data parallel rank {rank} is out of range for {num_engines} engine(s)")]
63 + InvalidDataParallelRank { rank: u32, num_engines: u32 },
64 + #[error("engine-core output dispatcher closed: {message}")]
65 + DispatcherClosed { message: String },
66 + #[error("engine-core client is closed: {message}")]
67 + ClientClosed { message: String },
68 + #[error("request output stream for `{request_id}` closed unexpectedly")]
69 + RequestStreamClosed { request_id: String },
70 + #[error("utility call `{method}` failed (call_id={call_id}): {message}")]
71 + UtilityCallFailed {
72 + method: String,
73 + call_id: i64,
74 + message: String,
75 + },
76 + #[error("utility call `{method}` returned an invalid result (call_id={call_id}): {message}")]
77 + UtilityResultDecode {
78 + method: String,
79 + call_id: i64,
80 + message: String,
81 + },
82 + #[error("utility call `{method}` closed unexpectedly (call_id={call_id})")]
83 + UtilityCallClosed { method: String, call_id: i64 },
84 +
85 + /// A special variant to allow cloning the same error.
86 + #[error(transparent)]
87 + Shared(Arc<Self>),
88 +}
rust/src/engine-core-client/src/lib.rs new
+18
@@ -0,0 +1,18 @@
1 +mod client;
2 +mod coordinator;
3 +mod error;
4 +mod metrics;
5 +pub mod protocol;
6 +#[cfg(any(test, feature = "test-util"))]
7 +pub mod test_utils;
8 +mod transport;
9 +
10 +pub use client::{
11 + AbortCause, CoordinatorMode, EngineCoreClient, EngineCoreClientConfig, EngineCoreOutputStream,
12 + EngineCoreStreamOutput, TransportMode,
13 +};
14 +pub use error::{Error, Result};
15 +pub use transport::{ENGINE_CORE_DEAD_SENTINEL, EngineId};
16 +
17 +#[cfg(test)]
18 +mod tests;
rust/src/engine-core-client/src/metrics.rs new
+131
@@ -0,0 +1,131 @@
1 +use vllm_metrics::{EngineLabels, EnginePositionLabels, SchedulerMetrics, WaitingReasonLabels};
2 +
3 +use crate::protocol::stats::SchedulerStats;
4 +
5 +const WAITING_REASON_CAPACITY: &str = "capacity";
6 +const WAITING_REASON_DEFERRED: &str = "deferred";
7 +
8 +/// Record the scheduler-stats-backed metrics for one engine at one point in
9 +/// time.
10 +pub(crate) fn record_scheduler_stats(
11 + metrics: &SchedulerMetrics,
12 + model_name: impl Into<String>,
13 + engine: u32,
14 + stats: &SchedulerStats,
15 +) {
16 + let model_name = model_name.into();
17 + let labels = EngineLabels {
18 + model_name: model_name.clone(),
19 + engine,
20 + };
21 +
22 + // Scheduler state gauges.
23 + metrics.scheduler_running.get_or_create(&labels).set(stats.num_running_reqs);
24 + metrics
25 + .scheduler_waiting
26 + .get_or_create(&labels)
27 + .set(stats.num_waiting_reqs + stats.num_skipped_waiting_reqs);
28 + metrics
29 + .scheduler_waiting_by_reason
30 + .get_or_create(&WaitingReasonLabels {
31 + model_name: model_name.clone(),
32 + engine,
33 + reason: WAITING_REASON_CAPACITY,
34 + })
35 + .set(stats.num_waiting_reqs);
36 + metrics
37 + .scheduler_waiting_by_reason
38 + .get_or_create(&WaitingReasonLabels {
39 + model_name: model_name.clone(),
40 + engine,
41 + reason: WAITING_REASON_DEFERRED,
42 + })
43 + .set(stats.num_skipped_waiting_reqs);
44 + metrics.kv_cache_usage.get_or_create(&labels).set(stats.kv_cache_usage);
45 +
46 + // Prefix-cache counters, including the connector-backed external cache path.
47 + metrics
48 + .prefix_cache_queries
49 + .get_or_create(&labels)
50 + .inc_by(stats.prefix_cache_stats.base.queries);
51 + metrics
52 + .prefix_cache_hits
53 + .get_or_create(&labels)
54 + .inc_by(stats.prefix_cache_stats.base.hits);
55 +
56 + if let Some(connector_prefix_cache_stats) = &stats.connector_prefix_cache_stats {
57 + metrics
58 + .external_prefix_cache_queries
59 + .get_or_create(&labels)
60 + .inc_by(connector_prefix_cache_stats.base.queries);
61 + metrics
62 + .external_prefix_cache_hits
63 + .get_or_create(&labels)
64 + .inc_by(connector_prefix_cache_stats.base.hits);
65 + }
66 +
67 + // Speculative decoding counters.
68 + if let Some(spec_decoding_stats) = &stats.spec_decoding_stats {
69 + metrics
70 + .spec_decode_num_drafts
71 + .get_or_create(&labels)
72 + .inc_by(spec_decoding_stats.num_drafts);
73 + metrics
74 + .spec_decode_num_draft_tokens
75 + .get_or_create(&labels)
76 + .inc_by(spec_decoding_stats.num_draft_tokens);
77 + metrics
78 + .spec_decode_num_accepted_tokens
79 + .get_or_create(&labels)
80 + .inc_by(spec_decoding_stats.num_accepted_tokens);
81 +
82 + for (position, accepted_tokens) in
83 + spec_decoding_stats.num_accepted_tokens_per_pos.iter().copied().enumerate()
84 + {
85 + metrics
86 + .spec_decode_num_accepted_tokens_per_pos
87 + .get_or_create(&EnginePositionLabels {
88 + model_name: model_name.clone(),
89 + engine,
90 + position: position as u32,
91 + })
92 + .inc_by(accepted_tokens);
93 + }
94 + }
95 +
96 + // Per-engine performance / MFU counters.
97 + if let Some(perf_stats) = &stats.perf_stats
98 + && (perf_stats.num_flops_per_gpu != 0
99 + || perf_stats.num_read_bytes_per_gpu != 0
100 + || perf_stats.num_write_bytes_per_gpu != 0)
101 + {
102 + metrics
103 + .estimated_flops_per_gpu
104 + .get_or_create(&labels)
105 + .inc_by(perf_stats.num_flops_per_gpu);
106 + metrics
107 + .estimated_read_bytes_per_gpu
108 + .get_or_create(&labels)
109 + .inc_by(perf_stats.num_read_bytes_per_gpu);
110 + metrics
111 + .estimated_write_bytes_per_gpu
112 + .get_or_create(&labels)
113 + .inc_by(perf_stats.num_write_bytes_per_gpu);
114 + }
115 +
116 + // Sampled KV-cache residency histograms.
117 + if !stats.kv_cache_eviction_events.is_empty() {
118 + let kv_block_lifetime_seconds = metrics.kv_block_lifetime_seconds.get_or_create(&labels);
119 + let kv_block_idle_before_evict_seconds =
120 + metrics.kv_block_idle_before_evict_seconds.get_or_create(&labels);
121 + let kv_block_reuse_gap_seconds = metrics.kv_block_reuse_gap_seconds.get_or_create(&labels);
122 +
123 + for event in &stats.kv_cache_eviction_events {
124 + kv_block_lifetime_seconds.observe(event.lifetime_seconds);
125 + kv_block_idle_before_evict_seconds.observe(event.idle_seconds);
126 + for reuse_gap_seconds in &event.reuse_gaps_seconds {
127 + kv_block_reuse_gap_seconds.observe(*reuse_gap_seconds);
128 + }
129 + }
130 + }
131 +}
rust/src/engine-core-client/src/protocol/classified_outputs.rs new
+251
@@ -0,0 +1,251 @@
1 +use std::collections::BTreeSet;
2 +
3 +use enum_as_inner::EnumAsInner;
4 +
5 +use super::{EngineCoreOutput, EngineCoreOutputs, UtilityOutput};
6 +use crate::protocol::stats::SchedulerStats;
7 +
8 +/// Data-parallel control notifications multiplexed through `EngineCoreOutputs`.
9 +#[derive(Debug, Clone, PartialEq, Eq)]
10 +pub enum DpControlMessage {
11 + WaveComplete(u32),
12 + StartWave(u32),
13 +}
14 +
15 +#[derive(Debug, Clone, PartialEq)]
16 +pub struct RequestBatchOutputs {
17 + pub engine_index: u32,
18 + pub outputs: Vec<EngineCoreOutput>,
19 + pub scheduler_stats: Option<Box<SchedulerStats>>,
20 + pub timestamp: f64,
21 + pub finished_requests: Option<BTreeSet<String>>,
22 +}
23 +
24 +#[derive(Debug, Clone, PartialEq)]
25 +pub struct UtilityCallOutput {
26 + pub engine_index: u32,
27 + pub timestamp: f64,
28 + pub output: UtilityOutput,
29 +}
30 +
31 +/// Semantic classification of a raw `EngineCoreOutputs` message.
32 +///
33 +/// Python currently uses one product-shaped wire struct for several distinct
34 +/// output families. This enum exposes those families more explicitly without
35 +/// changing the wire format.
36 +#[derive(Debug, Clone, PartialEq, EnumAsInner)]
37 +pub enum ClassifiedEngineCoreOutputs {
38 + RequestBatch(RequestBatchOutputs),
39 + Utility(UtilityCallOutput),
40 + DpControl {
41 + engine_index: u32,
42 + timestamp: f64,
43 + control: DpControlMessage,
44 + },
45 + /// Fallback for wire-shape combinations that do not map cleanly onto the
46 + /// current semantic families.
47 + Other(EngineCoreOutputs),
48 +}
49 +
50 +impl EngineCoreOutputs {
51 + /// Classify the raw wire message into a more semantic Rust enum.
52 + pub fn classify(self) -> ClassifiedEngineCoreOutputs {
53 + let has_request_payload = !self.outputs.is_empty()
54 + || self.scheduler_stats.is_some()
55 + || self.finished_requests.is_some();
56 +
57 + match (
58 + has_request_payload,
59 + &self.utility_output,
60 + &self.wave_complete,
61 + &self.start_wave,
62 + ) {
63 + (true, None, None, None) => {
64 + ClassifiedEngineCoreOutputs::RequestBatch(RequestBatchOutputs {
65 + engine_index: self.engine_index,
66 + outputs: self.outputs,
67 + scheduler_stats: self.scheduler_stats,
68 + timestamp: self.timestamp,
69 + finished_requests: self.finished_requests,
70 + })
71 + }
72 + (false, Some(_), None, None) => {
73 + ClassifiedEngineCoreOutputs::Utility(UtilityCallOutput {
74 + engine_index: self.engine_index,
75 + timestamp: self.timestamp,
76 + output: self.utility_output.unwrap(),
77 + })
78 + }
79 + (false, None, Some(_), None) => ClassifiedEngineCoreOutputs::DpControl {
80 + engine_index: self.engine_index,
81 + timestamp: self.timestamp,
82 + control: DpControlMessage::WaveComplete(self.wave_complete.unwrap()),
83 + },
84 + (false, None, None, Some(_)) => ClassifiedEngineCoreOutputs::DpControl {
85 + engine_index: self.engine_index,
86 + timestamp: self.timestamp,
87 + control: DpControlMessage::StartWave(self.start_wave.unwrap()),
88 + },
89 + _ => ClassifiedEngineCoreOutputs::Other(self),
90 + }
91 + }
92 +}
93 +
94 +#[cfg(test)]
95 +mod tests {
96 + use std::collections::BTreeSet;
97 +
98 + use super::*;
99 + use crate::protocol::EngineCoreOutput;
100 +
101 + #[test]
102 + fn engine_core_outputs_classify_request_batch() {
103 + let outputs = EngineCoreOutputs {
104 + outputs: vec![EngineCoreOutput {
105 + request_id: "req-1".to_string(),
106 + new_token_ids: vec![7],
107 + ..Default::default()
108 + }],
109 + finished_requests: Some(BTreeSet::from(["req-1".to_string()])),
110 + ..Default::default()
111 + };
112 +
113 + expect_test::expect![[r#"
114 + RequestBatch(
115 + RequestBatchOutputs {
116 + engine_index: 0,
117 + outputs: [
118 + EngineCoreOutput {
119 + request_id: "req-1",
120 + new_token_ids: [
121 + 7,
122 + ],
123 + new_logprobs: None,
124 + new_prompt_logprobs_tensors: None,
125 + pooling_output: None,
126 + finish_reason: None,
127 + stop_reason: None,
128 + events: None,
129 + kv_transfer_params: None,
130 + trace_headers: None,
131 + prefill_stats: None,
132 + routed_experts: None,
133 + num_nans_in_logits: 0,
134 + },
135 + ],
136 + scheduler_stats: None,
137 + timestamp: 0.0,
138 + finished_requests: Some(
139 + {
140 + "req-1",
141 + },
142 + ),
143 + },
144 + )
145 + "#]]
146 + .assert_debug_eq(&outputs.classify());
147 + }
148 +
149 + #[test]
150 + fn engine_core_outputs_classify_utility() {
151 + let outputs = EngineCoreOutputs {
152 + utility_output: Some(UtilityOutput {
153 + call_id: 42,
154 + failure_message: None,
155 + result: None,
156 + }),
157 + ..Default::default()
158 + };
159 +
160 + expect_test::expect![[r#"
161 + Utility(
162 + UtilityCallOutput {
163 + engine_index: 0,
164 + timestamp: 0.0,
165 + output: UtilityOutput {
166 + call_id: 42,
167 + failure_message: None,
168 + result: None,
169 + },
170 + },
171 + )
172 + "#]]
173 + .assert_debug_eq(&outputs.classify());
174 + }
175 +
176 + #[test]
177 + fn engine_core_outputs_classify_control() {
178 + let outputs = EngineCoreOutputs {
179 + start_wave: Some(3),
180 + ..Default::default()
181 + };
182 +
183 + expect_test::expect![[r#"
184 + DpControl {
185 + engine_index: 0,
186 + timestamp: 0.0,
187 + control: StartWave(
188 + 3,
189 + ),
190 + }
191 + "#]]
192 + .assert_debug_eq(&outputs.classify());
193 + }
194 +
195 + #[test]
196 + fn engine_core_outputs_classify_mixed_shape_as_raw() {
197 + let outputs = EngineCoreOutputs {
198 + outputs: vec![EngineCoreOutput {
199 + request_id: "req-1".to_string(),
200 + new_token_ids: vec![7],
201 + ..Default::default()
202 + }],
203 + utility_output: Some(UtilityOutput {
204 + call_id: 1,
205 + failure_message: None,
206 + result: None,
207 + }),
208 + ..Default::default()
209 + };
210 +
211 + expect_test::expect![[r#"
212 + Other(
213 + EngineCoreOutputs {
214 + engine_index: 0,
215 + outputs: [
216 + EngineCoreOutput {
217 + request_id: "req-1",
218 + new_token_ids: [
219 + 7,
220 + ],
221 + new_logprobs: None,
222 + new_prompt_logprobs_tensors: None,
223 + pooling_output: None,
224 + finish_reason: None,
225 + stop_reason: None,
226 + events: None,
227 + kv_transfer_params: None,
228 + trace_headers: None,
229 + prefill_stats: None,
230 + routed_experts: None,
231 + num_nans_in_logits: 0,
232 + },
233 + ],
234 + scheduler_stats: None,
235 + timestamp: 0.0,
236 + utility_output: Some(
237 + UtilityOutput {
238 + call_id: 1,
239 + failure_message: None,
240 + result: None,
241 + },
242 + ),
243 + finished_requests: None,
244 + wave_complete: None,
245 + start_wave: None,
246 + },
247 + )
248 + "#]]
249 + .assert_debug_eq(&outputs.classify());
250 + }
251 +}
rust/src/engine-core-client/src/protocol/dtype.rs new
+40
@@ -0,0 +1,40 @@
1 +use serde::{Deserialize, Serialize};
2 +
3 +/// Effective model dtype reported by the engine after config resolution.
4 +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
5 +pub enum ModelDtype {
6 + #[serde(rename = "float16")]
7 + Float16,
8 + #[serde(rename = "bfloat16")]
9 + BFloat16,
10 + #[serde(rename = "float32")]
11 + Float32,
12 +}
13 +
14 +impl ModelDtype {
15 + pub fn as_str(self) -> &'static str {
16 + match self {
17 + Self::Float16 => "float16",
18 + Self::BFloat16 => "bfloat16",
19 + Self::Float32 => "float32",
20 + }
21 + }
22 +}
23 +
24 +#[cfg(test)]
25 +mod tests {
26 + use super::ModelDtype;
27 +
28 + #[test]
29 + fn serde_uses_protocol_dtype_strings() {
30 + assert_eq!(
31 + serde_json::to_value(ModelDtype::Float16).unwrap(),
32 + serde_json::json!("float16")
33 + );
34 + assert_eq!(
35 + serde_json::from_value::<ModelDtype>(serde_json::json!("bfloat16")).unwrap(),
36 + ModelDtype::BFloat16
37 + );
38 + assert_eq!(ModelDtype::Float32.as_str(), "float32");
39 + }
40 +}
rust/src/engine-core-client/src/protocol/handshake.rs new
+90
@@ -0,0 +1,90 @@
1 +use std::collections::BTreeMap;
2 +
3 +use serde::{Deserialize, Serialize};
4 +
5 +use crate::protocol::{ModelDtype, OpaqueValue};
6 +
7 +/// Decoded engine startup-handshake payload sent on the handshake socket.
8 +///
9 +/// Original Python payload construction:
10 +/// <https://github.com/vllm-project/vllm/blob/c8d98f81f6/vllm/v1/engine/core.py#L1000-L1035>
11 +#[derive(Debug, Clone, Default, Serialize, Deserialize)]
12 +pub struct ReadyMessage {
13 + #[serde(default)]
14 + pub status: Option<String>,
15 + #[serde(default)]
16 + pub local: Option<bool>,
17 + #[serde(default)]
18 + pub headless: Option<bool>,
19 + #[serde(default)]
20 + pub parallel_config_hash: Option<String>,
21 +}
22 +
23 +/// Post-initialization configuration sent from each engine on the input socket
24 +/// registration message, after the handshake completes.
25 +///
26 +/// Contains values that may differ from the original config (e.g.
27 +/// `max_model_len` after KV cache auto-fitting, `num_gpu_blocks` after
28 +/// profiling).
29 +///
30 +/// Original Python definition:
31 +/// <https://github.com/vllm-project/vllm/blob/c8d98f81f6/vllm/v1/engine/__init__.py#L67-L77>
32 +#[derive(Debug, Clone, Serialize, Deserialize)]
33 +pub struct EngineCoreReadyResponse {
34 + /// Engine-reported maximum model context length (auto-fitted after
35 + /// KV cache profiling and may differ from the original config value).
36 + pub max_model_len: u64,
37 + /// Number of GPU blocks available for KV cache on this engine.
38 + pub num_gpu_blocks: u64,
39 + /// DP coordinator stats publish address, if applicable.
40 + pub dp_stats_address: Option<String>,
41 + /// Effective model dtype after Python vLLM resolves `--dtype`.
42 + // TODO: This is currently not wired up on the engine side. After it's added, remove `Option`
43 + // and `serde(default)`.
44 + #[serde(default)]
45 + pub dtype: Option<ModelDtype>,
46 +}
47 +
48 +/// Frontend-owned ZMQ addresses that are sent to the engine during startup
49 +/// handshake initialization.
50 +///
51 +/// Original Python definition (`EngineZmqAddresses`):
52 +/// <https://github.com/vllm-project/vllm/blob/f22d6e026798a74e6542a52ef776c054f2de572a/vllm/v1/engine/utils.py#L53-L67>
53 +#[derive(Debug, Clone, Serialize, Deserialize)]
54 +pub struct HandshakeAddresses {
55 + pub inputs: Vec<String>,
56 + pub outputs: Vec<String>,
57 + pub coordinator_input: Option<String>,
58 + pub coordinator_output: Option<String>,
59 + pub frontend_stats_publish_address: Option<String>,
60 +}
61 +
62 +/// Startup handshake payload sent from the frontend to initialize an engine
63 +/// after receiving `HELLO`.
64 +///
65 +/// Original Python definition (`EngineHandshakeMetadata`):
66 +/// <https://github.com/vllm-project/vllm/blob/f22d6e026798a74e6542a52ef776c054f2de572a/vllm/v1/engine/utils.py#L69-L77>
67 +#[derive(Debug, Clone, Serialize, Deserialize)]
68 +pub struct HandshakeInitMessage {
69 + pub addresses: HandshakeAddresses,
70 + pub parallel_config: BTreeMap<String, OpaqueValue>,
71 +}
72 +
73 +#[cfg(test)]
74 +mod tests {
75 + use super::EngineCoreReadyResponse;
76 + use crate::protocol::ModelDtype;
77 +
78 + #[test]
79 + fn ready_response_accepts_effective_dtype() {
80 + let response: EngineCoreReadyResponse = serde_json::from_value(serde_json::json!({
81 + "max_model_len": 4096,
82 + "num_gpu_blocks": 2,
83 + "dp_stats_address": null,
84 + "dtype": "bfloat16"
85 + }))
86 + .unwrap();
87 +
88 + assert_eq!(response.dtype, Some(ModelDtype::BFloat16));
89 + }
90 +}
rust/src/engine-core-client/src/protocol/logprobs.rs new
+330
@@ -0,0 +1,330 @@
1 +mod array;
2 +#[cfg(test)]
3 +mod tests;
4 +mod wire;
5 +
6 +use std::ops::{Deref, DerefMut};
7 +
8 +use enum_as_inner::EnumAsInner;
9 +use serde::{Deserialize, Deserializer, Serialize};
10 +
11 +use self::wire::*;
12 +use super::{EngineCoreOutput, EngineCoreOutputs, decode_msgpack};
13 +use crate::error::{Error, Result, bail_ext_value_decode, ext_value_decode};
14 +use crate::protocol::tensor::{WireArrayData, WireNdArray};
15 +
16 +/// One token candidate and its logprob metadata for a single sequence position.
17 +///
18 +/// The first entry in a [`PositionLogprobs`] is always the sampled/selected
19 +/// token for that position. Any remaining entries follow the engine's returned
20 +/// top-k candidate order.
21 +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
22 +pub struct TokenLogprob {
23 + pub token_id: u32,
24 + pub logprob: f32,
25 + /// The sampled/selected token uses its actual vocab rank. Remaining entries
26 + /// use 1-based top-k ranks matching the engine's returned candidate
27 + /// order.
28 + pub rank: u32,
29 +}
30 +
31 +/// Logprob payload for one sequence position.
32 +///
33 +/// This is the semantic Rust representation used by the public client API after
34 +/// the lower-level ndarray/tensor wire payload has been decoded.
35 +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
36 +pub struct PositionLogprobs {
37 + pub entries: Vec<TokenLogprob>,
38 +}
39 +
40 +impl PositionLogprobs {
41 + /// Convert one decoded logprobs row into this per-position form by grouping
42 + /// each token/logprob pair together with the sampled/selected token's
43 + /// actual vocab rank.
44 + fn from_decoded_row(token_ids: &[u32], logprobs: &[f32], sampled_rank: u32) -> Result<Self> {
45 + if token_ids.len() != logprobs.len() {
46 + bail_ext_value_decode!(
47 + "logprobs row length mismatch: token_ids={}, logprobs={}",
48 + token_ids.len(),
49 + logprobs.len()
50 + );
51 + }
52 + if sampled_rank == 0 {
53 + bail_ext_value_decode!("token_ranks must be >= 1 for decoded engine-core logprobs");
54 + }
55 +
56 + let mut entries = Vec::with_capacity(token_ids.len());
57 + for (index, (&token_id, &logprob)) in token_ids.iter().zip(logprobs.iter()).enumerate() {
58 + let rank = if index == 0 {
59 + sampled_rank
60 + } else {
61 + index as u32
62 + };
63 + entries.push(TokenLogprob {
64 + token_id,
65 + logprob,
66 + rank,
67 + });
68 + }
69 + Ok(Self { entries })
70 + }
71 +}
72 +
73 +/// Decoded per-request logprobs payload for one engine-core output.
74 +///
75 +/// Unlike the Python wire payload, this public Rust type is already fully
76 +/// semantic: one [`PositionLogprobs`] per scored position, each containing the
77 +/// sampled/selected token plus any returned top-k alternatives for that same
78 +/// position.
79 +///
80 +/// The Python engine still sends logprobs as ndarray/tensor-shaped wire tuples.
81 +/// Rust resolves that lower-level representation during decode and exposes only
82 +/// this per-position form to callers.
83 +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
84 +pub struct Logprobs {
85 + /// One decoded logprobs record per scored position in this engine-core
86 + /// output.
87 + pub positions: Vec<PositionLogprobs>,
88 +}
89 +
90 +impl Logprobs {
91 + /// Returns the number of scored positions in this payload.
92 + pub fn len(&self) -> usize {
93 + self.positions.len()
94 + }
95 +
96 + /// Returns whether the payload contains no scored positions.
97 + pub fn is_empty(&self) -> bool {
98 + self.positions.is_empty()
99 + }
100 +}
101 +
102 +/// Output field wrapper that is initially deserialized from the Python wire
103 +/// shape, then resolved into [`Logprobs`] before the decoded message is
104 +/// returned to callers.
105 +#[derive(Clone, PartialEq, Debug, EnumAsInner)]
106 +pub enum MaybeWireLogprobs {
107 + /// The logprobs are still in the wire format and need to be resolved by
108 + /// looking up aux frames and decoding raw views. Should only be used
109 + /// internally during deserialization.
110 + Wire(Box<WireLogprobs>),
111 + /// The actual decoded logprobs value,
112 + Direct(Logprobs),
113 +}
114 +
115 +impl Deref for MaybeWireLogprobs {
116 + type Target = Logprobs;
117 +
118 + fn deref(&self) -> &Self::Target {
119 + match self {
120 + Self::Wire(_) => panic!("Logprobs is still in wire format"),
121 + Self::Direct(value) => value,
122 + }
123 + }
124 +}
125 +
126 +impl DerefMut for MaybeWireLogprobs {
127 + fn deref_mut(&mut self) -> &mut Self::Target {
128 + match self {
129 + Self::Wire(_) => panic!("Logprobs is still in wire format"),
130 + Self::Direct(value) => value,
131 + }
132 + }
133 +}
134 +
135 +impl<'de> Deserialize<'de> for MaybeWireLogprobs {
136 + fn deserialize<D>(deserializer: D) -> std::result::Result<Self, D::Error>
137 + where
138 + D: Deserializer<'de>,
139 + {
140 + // When deserializing, it's always in the wire form.
141 + WireLogprobs::deserialize(deserializer).map(|v| Self::Wire(Box::new(v)))
142 + }
143 +}
144 +
145 +impl Serialize for MaybeWireLogprobs {
146 + fn serialize<S>(&self, serializer: S) -> std::result::Result<S::Ok, S::Error>
147 + where
148 + S: serde::Serializer,
149 + {
150 + // For testing purposes only. We don't actually serialize it into aux frames.
151 + match self {
152 + Self::Wire(value) => value.serialize(serializer),
153 + Self::Direct(value) => WireLogprobs::from_direct(value)
154 + .map_err(serde::ser::Error::custom)?
155 + .serialize(serializer),
156 + }
157 + }
158 +}
159 +
160 +impl MaybeWireLogprobs {
161 + /// Resolve the wire representation into decoded logprobs by looking up aux
162 + /// frames and decoding raw views as needed.
163 + fn resolve<Frame>(self, frames: &[Frame], field_prefix: &str) -> Result<Self>
164 + where
165 + Frame: AsRef<[u8]>,
166 + {
167 + match self {
168 + Self::Direct(value) => Ok(Self::Direct(value)),
169 + Self::Wire(value) => value.resolve(frames, field_prefix).map(Self::Direct),
170 + }
171 + }
172 +}
173 +
174 +impl EngineCoreOutputs {
175 + /// Resolve all wire-format fields in-place by looking up aux frames and
176 + /// decoding raw-view payloads as needed.
177 + fn resolve_in_place<Frame>(&mut self, frames: &[Frame]) -> Result<()>
178 + where
179 + Frame: AsRef<[u8]>,
180 + {
181 + for output in &mut self.outputs {
182 + output.resolve_in_place(frames)?;
183 + }
184 + Ok(())
185 + }
186 +}
187 +
188 +impl EngineCoreOutput {
189 + /// Resolve all wire-format fields in-place by looking up aux frames and
190 + /// decoding raw-view payloads as needed.
191 + fn resolve_in_place<Frame>(&mut self, frames: &[Frame]) -> Result<()>
192 + where
193 + Frame: AsRef<[u8]>,
194 + {
195 + self.new_logprobs = (self.new_logprobs.take())
196 + .map(|value| value.resolve(frames, "new_logprobs"))
197 + .transpose()?;
198 + self.new_prompt_logprobs_tensors = (self.new_prompt_logprobs_tensors.take())
199 + .map(|value| value.resolve(frames, "new_prompt_logprobs_tensors"))
200 + .transpose()?;
201 + Ok(())
202 + }
203 +}
204 +
205 +impl WireLogprobs {
206 + /// Convert semantic per-position logprobs into the Python wire tuple shape.
207 + ///
208 + /// This exists mainly so Rust-side tests can inject semantic logprobs into
209 + /// mocked engine-core outputs without manually building ndarray
210 + /// raw-view tuples.
211 + fn from_direct(value: &Logprobs) -> std::result::Result<Self, String> {
212 + let rows = value.positions.len();
213 + let cols = value.positions.first().map(|position| position.entries.len()).unwrap_or(0);
214 +
215 + let mut token_ids = Vec::with_capacity(rows.saturating_mul(cols).saturating_mul(8));
216 + let mut logprobs = Vec::with_capacity(rows.saturating_mul(cols).saturating_mul(4));
217 + let mut token_ranks = Vec::with_capacity(rows.saturating_mul(8));
218 +
219 + for (row_index, position) in value.positions.iter().enumerate() {
220 + if position.entries.len() != cols {
221 + return Err(format!(
222 + "logprobs row {row_index} length mismatch: expected {cols}, got {}",
223 + position.entries.len()
224 + ));
225 + }
226 + let Some((sampled, _)) = position.entries.split_first() else {
227 + return Err(format!("logprobs row {row_index} is empty"));
228 + };
229 +
230 + token_ranks.extend_from_slice(&(sampled.rank as i64).to_le_bytes());
231 + for entry in &position.entries {
232 + token_ids.extend_from_slice(&(entry.token_id as i64).to_le_bytes());
233 + logprobs.extend_from_slice(&entry.logprob.to_le_bytes());
234 + }
235 + }
236 +
237 + Ok(Self {
238 + logprob_token_ids: WireNdArray {
239 + dtype: "<i8".to_string(),
240 + shape: vec![rows, cols],
241 + data: WireArrayData::RawView(token_ids),
242 + },
243 + logprobs: WireNdArray {
244 + dtype: "<f4".to_string(),
245 + shape: vec![rows, cols],
246 + data: WireArrayData::RawView(logprobs),
247 + },
248 + token_ranks: WireNdArray {
249 + dtype: "<i8".to_string(),
250 + shape: vec![rows],
251 + data: WireArrayData::RawView(token_ranks),
252 + },
253 + cu_num_generated_tokens: None,
254 + })
255 + }
256 +
257 + /// Resolve the wire-format logprobs into semantic [`Logprobs`] records by
258 + /// looking up aux frames, decoding raw views, and grouping each row
259 + /// into one [`PositionLogprobs`].
260 + fn resolve<Frame>(self, frames: &[Frame], field_prefix: &str) -> Result<Logprobs>
261 + where
262 + Frame: AsRef<[u8]>,
263 + {
264 + if let Some(indices) = self.cu_num_generated_tokens {
265 + bail_ext_value_decode!(
266 + "{field_prefix}.cu_num_generated_tokens: \
267 + expected None for per-request engine-core logprobs payload, got {indices:?}"
268 + );
269 + }
270 +
271 + let token_ids = array::decode_array2_u32(
272 + self.logprob_token_ids,
273 + &format!("{field_prefix}.logprob_token_ids"),
274 + frames,
275 + )?;
276 + let logprobs =
277 + array::decode_array2_f32(self.logprobs, &format!("{field_prefix}.logprobs"), frames)?;
278 + let token_ranks = array::decode_array1_u32(
279 + self.token_ranks,
280 + &format!("{field_prefix}.token_ranks"),
281 + frames,
282 + )?;
283 +
284 + if token_ids.rows != logprobs.rows || token_ids.cols != logprobs.cols {
285 + bail_ext_value_decode!(
286 + "{field_prefix}: row shape mismatch between token ids ({}, {}) and logprobs ({}, {})",
287 + token_ids.rows,
288 + token_ids.cols,
289 + logprobs.rows,
290 + logprobs.cols
291 + );
292 + }
293 + if token_ids.rows != token_ranks.len() {
294 + bail_ext_value_decode!(
295 + "{field_prefix}: token_ranks length {} does not match row count {}",
296 + token_ranks.len(),
297 + token_ids.rows
298 + );
299 + }
300 +
301 + let mut positions = Vec::with_capacity(token_ids.rows);
302 + for ((token_ids_row, logprobs_row), sampled_rank) in token_ids
303 + .data
304 + .chunks(token_ids.cols)
305 + .zip(logprobs.data.chunks(logprobs.cols))
306 + .zip(token_ranks)
307 + {
308 + positions.push(PositionLogprobs::from_decoded_row(
309 + token_ids_row,
310 + logprobs_row,
311 + sampled_rank,
312 + )?);
313 + }
314 +
315 + Ok(Logprobs { positions })
316 + }
317 +}
318 +
319 +/// Decode one ordinary or multipart engine-core output message into the strong
320 +/// typed public protocol shape.
321 +pub fn decode_engine_core_outputs<Frame>(frames: &[Frame]) -> Result<EngineCoreOutputs>
322 +where
323 + Frame: AsRef<[u8]>,
324 +{
325 + let first_frame = frames.first().ok_or_else(|| ext_value_decode!("missing output frame"))?;
326 +
327 + let mut outputs: EngineCoreOutputs = decode_msgpack(first_frame.as_ref())?;
328 + outputs.resolve_in_place(frames)?;
329 + Ok(outputs)
330 +}
rust/src/engine-core-client/src/protocol/logprobs/array.rs new
+306
@@ -0,0 +1,306 @@
1 +use std::io::Cursor;
2 +
3 +use byteorder::{BigEndian, LittleEndian, NativeEndian, ReadBytesExt};
4 +use itertools::Itertools as _;
5 +
6 +use crate::error::{Error, Result, ext_value_decode};
7 +use crate::protocol::tensor::{ShapeExt as _, WireArrayData, WireNdArray};
8 +
9 +#[derive(Debug, Clone, Copy, PartialEq, Eq)]
10 +pub(super) enum ScalarType {
11 + I32,
12 + I64,
13 + F32,
14 +}
15 +
16 +#[derive(Debug, Clone, Copy, PartialEq, Eq)]
17 +pub(super) enum Endianness {
18 + Little,
19 + Big,
20 + Native,
21 +}
22 +
23 +#[derive(Debug, Clone, PartialEq)]
24 +pub(super) struct DecodedArray2<T> {
25 + pub rows: usize,
26 + pub cols: usize,
27 + pub data: Vec<T>,
28 +}
29 +
30 +pub(super) fn decode_array2_u32<Frame>(
31 + value: WireNdArray,
32 + field: &str,
33 + frames: &[Frame],
34 +) -> Result<DecodedArray2<u32>>
35 +where
36 + Frame: AsRef<[u8]>,
37 +{
38 + let (shape, bytes, scalar, endianness) =
39 + decode_array_metadata(value, field, frames, &[ScalarType::I32, ScalarType::I64])?;
40 + if shape.len() != 2 {
41 + return Err(decode_error(
42 + field,
43 + &format!("expected rank-2 array, got rank {}", shape.len()),
44 + ));
45 + }
46 +
47 + let data = match scalar {
48 + ScalarType::I32 => decode_i32_vec(&bytes, endianness, field)?
49 + .into_iter()
50 + .map(|value| convert_to_u32(value, field))
51 + .try_collect()?,
52 + ScalarType::I64 => decode_i64_vec(&bytes, endianness, field)?
53 + .into_iter()
54 + .map(|value| convert_to_u32(value, field))
55 + .try_collect()?,
56 + ScalarType::F32 => unreachable!("scalar validation should reject f32"),
57 + };
58 + Ok(DecodedArray2 {
59 + rows: shape[0],
60 + cols: shape[1],
61 + data,
62 + })
63 +}
64 +
65 +pub(super) fn decode_array1_u32<Frame>(
66 + value: WireNdArray,
67 + field: &str,
68 + frames: &[Frame],
69 +) -> Result<Vec<u32>>
70 +where
71 + Frame: AsRef<[u8]>,
72 +{
73 + let (shape, bytes, scalar, endianness) =
74 + decode_array_metadata(value, field, frames, &[ScalarType::I32, ScalarType::I64])?;
75 + if shape.len() != 1 {
76 + return Err(decode_error(
77 + field,
78 + &format!("expected rank-1 array, got rank {}", shape.len()),
79 + ));
80 + }
81 +
82 + let data = match scalar {
83 + ScalarType::I32 => decode_i32_vec(&bytes, endianness, field)?
84 + .into_iter()
85 + .map(|value| convert_to_u32(value, field))
86 + .try_collect()?,
87 + ScalarType::I64 => decode_i64_vec(&bytes, endianness, field)?
88 + .into_iter()
89 + .map(|value| convert_to_u32(value, field))
90 + .try_collect()?,
91 + ScalarType::F32 => unreachable!("scalar validation should reject f32"),
92 + };
93 + Ok(data)
94 +}
95 +
96 +pub(super) fn decode_array2_f32<Frame>(
97 + value: WireNdArray,
98 + field: &str,
99 + frames: &[Frame],
100 +) -> Result<DecodedArray2<f32>>
101 +where
102 + Frame: AsRef<[u8]>,
103 +{
104 + let (shape, bytes, _, endianness) =
105 + decode_array_metadata(value, field, frames, &[ScalarType::F32])?;
106 + if shape.len() != 2 {
107 + return Err(decode_error(
108 + field,
109 + &format!("expected rank-2 array, got rank {}", shape.len()),
110 + ));
111 + }
112 +
113 + let data = decode_f32_vec(&bytes, endianness, field)?;
114 + Ok(DecodedArray2 {
115 + rows: shape[0],
116 + cols: shape[1],
117 + data,
118 + })
119 +}
120 +
121 +pub(super) fn decode_array_metadata<Frame>(
122 + value: WireNdArray,
123 + field: &str,
124 + frames: &[Frame],
125 + expected_scalars: &[ScalarType],
126 +) -> Result<(Vec<usize>, Vec<u8>, ScalarType, Endianness)>
127 +where
128 + Frame: AsRef<[u8]>,
129 +{
130 + let WireNdArray { dtype, shape, data } = value;
131 + let (scalar, endianness) = parse_dtype(&dtype, field)?;
132 + if !expected_scalars.contains(&scalar) {
133 + return Err(decode_error(
134 + field,
135 + &format!("expected dtype in {:?}, got {}", expected_scalars, dtype),
136 + ));
137 + }
138 +
139 + let bytes = resolve_array_bytes(data, field, frames)?;
140 + validate_byte_length(shape.as_slice(), bytes.len(), field, scalar)?;
141 + Ok((shape, bytes, scalar, endianness))
142 +}
143 +
144 +pub(super) fn parse_dtype(dtype: &str, field: &str) -> Result<(ScalarType, Endianness)> {
145 + let (endianness, body) = match dtype.as_bytes().first().copied() {
146 + Some(b'<') => (Endianness::Little, &dtype[1..]),
147 + Some(b'>') => (Endianness::Big, &dtype[1..]),
148 + Some(b'=') => (Endianness::Native, &dtype[1..]),
149 + Some(b'|') => (Endianness::Native, &dtype[1..]),
150 + _ => (Endianness::Native, dtype),
151 + };
152 +
153 + let scalar = match body {
154 + "i4" | "int32" => ScalarType::I32,
155 + "i8" | "int64" => ScalarType::I64,
156 + "f4" | "float32" => ScalarType::F32,
157 + _ => {
158 + return Err(decode_error(
159 + field,
160 + &format!("unsupported dtype string {dtype:?}"),
161 + ));
162 + }
163 + };
164 + Ok((scalar, endianness))
165 +}
166 +
167 +pub(super) fn resolve_array_bytes<Frame>(
168 + value: WireArrayData,
169 + field: &str,
170 + frames: &[Frame],
171 +) -> Result<Vec<u8>>
172 +where
173 + Frame: AsRef<[u8]>,
174 +{
175 + match value {
176 + WireArrayData::RawView(bytes) => Ok(bytes),
177 + WireArrayData::AuxIndex(index) => {
178 + let frame = frames.get(index).ok_or_else(|| {
179 + decode_error(
180 + field,
181 + &format!(
182 + "aux frame index {index} out of range for {} frames",
183 + frames.len()
184 + ),
185 + )
186 + })?;
187 + Ok(frame.as_ref().to_vec())
188 + }
189 + }
190 +}
191 +
192 +pub(super) fn validate_byte_length(
193 + shape: &[usize],
194 + byte_len: usize,
195 + field: &str,
196 + scalar: ScalarType,
197 +) -> Result<()> {
198 + let element_count = shape
199 + .checked_numel()
200 + .ok_or_else(|| decode_error(field, "shape element count overflowed usize"))?;
201 + let element_size = match scalar {
202 + ScalarType::I32 | ScalarType::F32 => 4,
203 + ScalarType::I64 => 8,
204 + };
205 + let expected = element_count
206 + .checked_mul(element_size)
207 + .ok_or_else(|| decode_error(field, "byte length overflowed usize"))?;
208 + if expected != byte_len {
209 + return Err(decode_error(
210 + field,
211 + &format!("byte length mismatch: expected {expected}, got {byte_len}"),
212 + ));
213 + }
214 + Ok(())
215 +}
216 +
217 +pub(super) fn decode_i32_vec(
218 + bytes: &[u8],
219 + endianness: Endianness,
220 + field: &str,
221 +) -> Result<Vec<i32>> {
222 + if !bytes.len().is_multiple_of(4) {
223 + return Err(decode_error(
224 + field,
225 + &format!("byte length {} is not divisible by 4", bytes.len()),
226 + ));
227 + }
228 + let mut cursor = Cursor::new(bytes);
229 + let mut values = Vec::with_capacity(bytes.len() / 4);
230 + while (cursor.position() as usize) < bytes.len() {
231 + let value = match endianness {
232 + Endianness::Little => cursor.read_i32::<LittleEndian>(),
233 + Endianness::Big => cursor.read_i32::<BigEndian>(),
234 + Endianness::Native => cursor.read_i32::<NativeEndian>(),
235 + }
236 + .map_err(|error| decode_error(field, &format!("failed to read i32 payload: {error}")))?;
237 + values.push(value);
238 + }
239 + Ok(values)
240 +}
241 +
242 +pub(super) fn decode_f32_vec(
243 + bytes: &[u8],
244 + endianness: Endianness,
245 + field: &str,
246 +) -> Result<Vec<f32>> {
247 + if !bytes.len().is_multiple_of(4) {
248 + return Err(decode_error(
249 + field,
250 + &format!("byte length {} is not divisible by 4", bytes.len()),
251 + ));
252 + }
253 + let mut cursor = Cursor::new(bytes);
254 + let mut values = Vec::with_capacity(bytes.len() / 4);
255 + while (cursor.position() as usize) < bytes.len() {
256 + let value = match endianness {
257 + Endianness::Little => cursor.read_f32::<LittleEndian>(),
258 + Endianness::Big => cursor.read_f32::<BigEndian>(),
259 + Endianness::Native => cursor.read_f32::<NativeEndian>(),
260 + }
261 + .map_err(|error| decode_error(field, &format!("failed to read f32 payload: {error}")))?;
262 + values.push(value);
263 + }
264 + Ok(values)
265 +}
266 +
267 +pub(super) fn decode_i64_vec(
268 + bytes: &[u8],
269 + endianness: Endianness,
270 + field: &str,
271 +) -> Result<Vec<i64>> {
272 + if !bytes.len().is_multiple_of(8) {
273 + return Err(decode_error(
274 + field,
275 + &format!("byte length {} is not divisible by 8", bytes.len()),
276 + ));
277 + }
278 + let mut cursor = Cursor::new(bytes);
279 + let mut values = Vec::with_capacity(bytes.len() / 8);
280 + while (cursor.position() as usize) < bytes.len() {
281 + let value = match endianness {
282 + Endianness::Little => cursor.read_i64::<LittleEndian>(),
283 + Endianness::Big => cursor.read_i64::<BigEndian>(),
284 + Endianness::Native => cursor.read_i64::<NativeEndian>(),
285 + }
286 + .map_err(|error| decode_error(field, &format!("failed to read i64 payload: {error}")))?;
287 + values.push(value);
288 + }
289 + Ok(values)
290 +}
291 +
292 +fn convert_to_u32<I>(value: I, field: &str) -> Result<u32>
293 +where
294 + I: TryInto<u32> + std::fmt::Display + Copy,
295 +{
296 + value.try_into().map_err(|_| {
297 + decode_error(
298 + field,
299 + &format!("expected non-negative token id/rank that fits in u32, got {value}"),
300 + )
301 + })
302 +}
303 +
304 +pub(super) fn decode_error(field: &str, reason: &str) -> Error {
305 + ext_value_decode!("{field}: {reason}")
306 +}
rust/src/engine-core-client/src/protocol/logprobs/tests.rs new
+302
@@ -0,0 +1,302 @@
1 +use std::collections::BTreeSet;
2 +
3 +use bytes::Bytes;
4 +use rmpv::Value;
5 +
6 +use super::{Logprobs, PositionLogprobs, TokenLogprob, decode_engine_core_outputs};
7 +use crate::protocol::EngineCoreFinishReason;
8 +
9 +fn encode_value(value: &Value) -> Vec<u8> {
10 + let mut out = Vec::new();
11 + rmpv::encode::write_value(&mut out, value).unwrap();
12 + out
13 +}
14 +
15 +fn output_wire_with_custom_fields(
16 + new_logprobs: Option<Value>,
17 + prompt_logprobs: Option<Value>,
18 +) -> Value {
19 + Value::Array(vec![
20 + Value::from(0),
21 + Value::Array(vec![Value::Array(vec![
22 + Value::from("req-1"),
23 + Value::Array(vec![Value::from(7), Value::from(8)]),
24 + new_logprobs.unwrap_or(Value::Nil),
25 + prompt_logprobs.unwrap_or(Value::Nil),
26 + Value::Nil,
27 + Value::from(EngineCoreFinishReason::Length as u8),
28 + ])]),
29 + Value::Nil,
30 + Value::from(0.0),
31 + Value::Nil,
32 + Value::Array(vec![Value::from("req-1")]),
33 + ])
34 +}
35 +
36 +fn ndarray_value(dtype: &str, shape: &[usize], data: Value) -> Value {
37 + Value::Array(vec![
38 + Value::from(dtype),
39 + Value::Array(shape.iter().copied().map(Value::from).collect()),
40 + data,
41 + ])
42 +}
43 +
44 +fn inline_logprobs_value() -> Value {
45 + let ids = Value::Ext(
46 + 3,
47 + vec![
48 + 1, 0, 0, 0, 0, 0, 0, 0, 2, 0, 0, 0, 0, 0, 0, 0, 3, 0, 0, 0, 0, 0, 0, 0, 4, 0, 0, 0, 0,
49 + 0, 0, 0, 5, 0, 0, 0, 0, 0, 0, 0, 6, 0, 0, 0, 0, 0, 0, 0,
50 + ],
51 + );
52 + let probs = Value::Ext(
53 + 3,
54 + vec![
55 + 0, 0, 128, 63, 0, 0, 0, 64, 0, 0, 64, 64, 0, 0, 128, 64, 0, 0, 160, 64, 0, 0, 192, 64,
56 + ],
57 + );
58 + let ranks = Value::Ext(3, vec![1, 0, 0, 0, 0, 0, 0, 0, 2, 0, 0, 0, 0, 0, 0, 0]);
59 + Value::Array(vec![
60 + ndarray_value("<i8", &[2, 3], ids),
61 + ndarray_value("<f4", &[2, 3], probs),
62 + ndarray_value("<i8", &[2], ranks),
63 + Value::Nil,
64 + ])
65 +}
66 +
67 +fn inline_prompt_logprobs_value() -> Value {
68 + let ids = Value::Ext(
69 + 3,
70 + vec![
71 + 10, 0, 0, 0, 0, 0, 0, 0, 11, 0, 0, 0, 0, 0, 0, 0, 12, 0, 0, 0, 0, 0, 0, 0, 13, 0, 0, 0,
72 + 0, 0, 0, 0, 14, 0, 0, 0, 0, 0, 0, 0, 15, 0, 0, 0, 0, 0, 0, 0,
73 + ],
74 + );
75 + let probs = Value::Ext(
76 + 3,
77 + vec![
78 + 0, 0, 32, 65, 0, 0, 48, 65, 0, 0, 64, 65, 0, 0, 80, 65, 0, 0, 96, 65, 0, 0, 112, 65,
79 + ],
80 + );
81 + let ranks = Value::Ext(3, vec![3, 0, 0, 0, 0, 0, 0, 0, 4, 0, 0, 0, 0, 0, 0, 0]);
82 + Value::Array(vec![
83 + ndarray_value("int64", &[2, 3], ids),
84 + ndarray_value("float32", &[2, 3], probs),
85 + ndarray_value("int64", &[2], ranks),
86 + Value::Nil,
87 + ])
88 +}
89 +
90 +fn expected_sample_logprobs() -> Logprobs {
91 + Logprobs {
92 + positions: vec![
93 + PositionLogprobs {
94 + entries: vec![
95 + TokenLogprob {
96 + token_id: 1,
97 + logprob: 1.0,
98 + rank: 1,
99 + },
100 + TokenLogprob {
101 + token_id: 2,
102 + logprob: 2.0,
103 + rank: 1,
104 + },
105 + TokenLogprob {
106 + token_id: 3,
107 + logprob: 3.0,
108 + rank: 2,
109 + },
110 + ],
111 + },
112 + PositionLogprobs {
113 + entries: vec![
114 + TokenLogprob {
115 + token_id: 4,
116 + logprob: 4.0,
117 + rank: 2,
118 + },
119 + TokenLogprob {
120 + token_id: 5,
121 + logprob: 5.0,
122 + rank: 1,
123 + },
124 + TokenLogprob {
125 + token_id: 6,
126 + logprob: 6.0,
127 + rank: 2,
128 + },
129 + ],
130 + },
131 + ],
132 + }
133 +}
134 +
135 +fn expected_prompt_logprobs() -> Logprobs {
136 + Logprobs {
137 + positions: vec![
138 + PositionLogprobs {
139 + entries: vec![
140 + TokenLogprob {
141 + token_id: 10,
142 + logprob: 10.0,
143 + rank: 3,
144 + },
145 + TokenLogprob {
146 + token_id: 11,
147 + logprob: 11.0,
148 + rank: 1,
149 + },
150 + TokenLogprob {
151 + token_id: 12,
152 + logprob: 12.0,
153 + rank: 2,
154 + },
155 + ],
156 + },
157 + PositionLogprobs {
158 + entries: vec![
159 + TokenLogprob {
160 + token_id: 13,
161 + logprob: 13.0,
162 + rank: 4,
163 + },
164 + TokenLogprob {
165 + token_id: 14,
166 + logprob: 14.0,
167 + rank: 1,
168 + },
169 + TokenLogprob {
170 + token_id: 15,
171 + logprob: 15.0,
172 + rank: 2,
173 + },
174 + ],
175 + },
176 + ],
177 + }
178 +}
179 +
180 +#[test]
181 +fn decodes_inline_new_logprobs() {
182 + let frames = vec![Bytes::from(encode_value(&output_wire_with_custom_fields(
183 + Some(inline_logprobs_value()),
184 + None,
185 + )))];
186 + let decoded = decode_engine_core_outputs(&frames).unwrap();
187 +
188 + let logprobs = decoded.outputs[0].new_logprobs.clone().unwrap().into_direct().unwrap();
189 + assert_eq!(logprobs, expected_sample_logprobs());
190 + assert_eq!(
191 + decoded.finished_requests,
192 + Some(BTreeSet::from(["req-1".to_string()]))
193 + );
194 +}
195 +
196 +#[test]
197 +fn decodes_multipart_new_logprobs() {
198 + let frames = vec![
199 + Bytes::from(encode_value(&output_wire_with_custom_fields(
200 + Some(Value::Array(vec![
201 + ndarray_value("<i8", &[2, 3], Value::from(1)),
202 + ndarray_value("<f4", &[2, 3], Value::from(2)),
203 + ndarray_value("<i8", &[2], Value::from(3)),
204 + Value::Nil,
205 + ])),
206 + None,
207 + ))),
208 + Bytes::from_static(&[
209 + 1, 0, 0, 0, 0, 0, 0, 0, 2, 0, 0, 0, 0, 0, 0, 0, 3, 0, 0, 0, 0, 0, 0, 0, 4, 0, 0, 0, 0,
210 + 0, 0, 0, 5, 0, 0, 0, 0, 0, 0, 0, 6, 0, 0, 0, 0, 0, 0, 0,
211 + ]),
212 + Bytes::from_static(&[
213 + 0, 0, 128, 63, 0, 0, 0, 64, 0, 0, 64, 64, 0, 0, 128, 64, 0, 0, 160, 64, 0, 0, 192, 64,
214 + ]),
215 + Bytes::from_static(&[1, 0, 0, 0, 0, 0, 0, 0, 2, 0, 0, 0, 0, 0, 0, 0]),
216 + ];
217 + let decoded = decode_engine_core_outputs(&frames).unwrap();
218 +
219 + let logprobs = decoded.outputs[0].new_logprobs.clone().unwrap().into_direct().unwrap();
220 + assert_eq!(logprobs, expected_sample_logprobs());
221 +}
222 +
223 +#[test]
224 +fn decodes_inline_prompt_logprobs() {
225 + let frames = vec![Bytes::from(encode_value(&output_wire_with_custom_fields(
226 + None,
227 + Some(inline_prompt_logprobs_value()),
228 + )))];
229 + let decoded = decode_engine_core_outputs(&frames).unwrap();
230 +
231 + let logprobs = decoded.outputs[0]
232 + .new_prompt_logprobs_tensors
233 + .clone()
234 + .unwrap()
235 + .into_direct()
236 + .unwrap();
237 + assert_eq!(logprobs, expected_prompt_logprobs());
238 +}
239 +
240 +#[test]
241 +fn decodes_big_endian_payloads() {
242 + let frames = vec![Bytes::from(encode_value(&output_wire_with_custom_fields(
243 + Some(Value::Array(vec![
244 + ndarray_value(">i4", &[1, 2], Value::Ext(3, vec![0, 0, 0, 1, 0, 0, 0, 2])),
245 + ndarray_value(
246 + ">f4",
247 + &[1, 2],
248 + Value::Ext(3, vec![63, 128, 0, 0, 64, 0, 0, 0]),
249 + ),
250 + ndarray_value(">i4", &[1], Value::Ext(3, vec![0, 0, 0, 3])),
251 + Value::Nil,
252 + ])),
253 + None,
254 + )))];
255 + let decoded = decode_engine_core_outputs(&frames).unwrap();
256 + let logprobs = decoded.outputs[0].new_logprobs.clone().unwrap().into_direct().unwrap();
257 + assert_eq!(
258 + logprobs,
259 + Logprobs {
260 + positions: vec![PositionLogprobs {
261 + entries: vec![
262 + TokenLogprob {
263 + token_id: 1,
264 + logprob: 1.0,
265 + rank: 3,
266 + },
267 + TokenLogprob {
268 + token_id: 2,
269 + logprob: 2.0,
270 + rank: 1,
271 + },
272 + ],
273 + }],
274 + }
275 + );
276 +}
277 +
278 +#[test]
279 +fn rejects_non_none_cu_num_generated_tokens() {
280 + let frames = vec![Bytes::from(encode_value(&output_wire_with_custom_fields(
281 + Some(Value::Array(vec![
282 + ndarray_value("<i8", &[1, 1], Value::Ext(3, vec![1, 0, 0, 0, 0, 0, 0, 0])),
283 + ndarray_value("<f4", &[1, 1], Value::Ext(3, vec![0, 0, 128, 63])),
284 + ndarray_value("<i8", &[1], Value::Ext(3, vec![1, 0, 0, 0, 0, 0, 0, 0])),
285 + Value::Array(vec![Value::from(0usize), Value::from(1usize)]),
286 + ])),
287 + None,
288 + )))];
289 +
290 + let error = decode_engine_core_outputs(&frames).unwrap_err();
291 + let crate::error::Error::ExtValueDecode { message } = &error else {
292 + panic!("expected ValueDecodeExt");
293 + };
294 + assert_eq!(
295 + message,
296 + "new_logprobs.cu_num_generated_tokens: expected None for per-request engine-core logprobs payload, got [0, 1]"
297 + );
298 + assert_eq!(
299 + error.to_string(),
300 + "messagepack ext value decode failed: new_logprobs.cu_num_generated_tokens: expected None for per-request engine-core logprobs payload, got [0, 1]"
301 + );
302 +}
rust/src/engine-core-client/src/protocol/logprobs/wire.rs new
+31
@@ -0,0 +1,31 @@
1 +use serde_tuple::{Deserialize_tuple, Serialize_tuple};
2 +
3 +use crate::protocol::tensor::WireNdArray;
4 +
5 +/// Python wire representation of `LogprobsLists` / `LogprobsTensors` before
6 +/// aux-frame references and raw-view payloads are resolved.
7 +///
8 +/// This mirrors the tuple shape emitted by Python engine-core so serde can
9 +/// first deserialize the raw wire payload before the Rust client converts it
10 +/// into semantic per-position logprobs records.
11 +///
12 +/// Original Python definition:
13 +/// <https://github.com/vllm-project/vllm/blob/f22d6e026798a74e6542a52ef776c054f2de572a/vllm/v1/outputs.py#L23-L56>
14 +#[derive(Debug, Clone, PartialEq, Serialize_tuple, Deserialize_tuple)]
15 +pub struct WireLogprobs {
16 + /// Wire array with shape `[num_positions, max_num_logprobs + 1]`.
17 + pub logprob_token_ids: WireNdArray,
18 + /// Wire array with shape `[num_positions, max_num_logprobs + 1]`.
19 + pub logprobs: WireNdArray,
20 + /// Wire array with shape `[num_positions]`.
21 + ///
22 + /// Python uses the field name `sampled_token_ranks` for sample logprobs and
23 + /// `selected_token_ranks` for prompt logprobs. Rust keeps one neutral field
24 + /// because both payloads share the same wire representation.
25 + pub token_ranks: WireNdArray,
26 + /// Preserved only for wire compatibility with batch-level Python tensors.
27 + /// Scheduler-sliced per-request outputs should emit `None` here, and
28 + /// the semantic Rust decoder rejects any other value.
29 + #[serde(default)]
30 + pub cu_num_generated_tokens: Option<Vec<usize>>,
31 +}
rust/src/engine-core-client/src/protocol/mod.rs new
+756
@@ -0,0 +1,756 @@
1 +use std::any::type_name;
2 +use std::collections::{BTreeMap, BTreeSet, HashMap};
3 +use std::io::Cursor;
4 +
5 +use bytes::Bytes;
6 +use rmpv::Value;
7 +use serde::{Deserialize, Serialize};
8 +use serde_default::DefaultFromSerde;
9 +use serde_repr::{Deserialize_repr, Serialize_repr};
10 +use serde_tuple::{Deserialize_tuple, Serialize_tuple};
11 +use thiserror_ext::AsReport;
12 +
13 +use crate::error::{Error, Result};
14 +use crate::protocol::logprobs::MaybeWireLogprobs;
15 +use crate::protocol::multimodal::MmFeatures;
16 +use crate::protocol::stats::{PrefillStats, SchedulerStats};
17 +
18 +// TODO: This module currently mixes reusable frontend-facing semantic types
19 +// (for example `FinishReason`, `StopReason`, `RequestOutputKind`, and future
20 +// cleaned-up frontend sampling types) with engine-core-specific wire DTOs and
21 +// handshake/control messages. While the Rust frontend is still evolving
22 +// quickly, keep them co-located here for iteration speed. Once the higher-level
23 +// API boundary stabilizes, move the truly reusable semantic types into a
24 +// lower-level common crate and keep the engine transport/wire messages here.
25 +
26 +/// Dynamic msgpack value used for schema positions that are preserved but not
27 +/// yet strongly typed in the early-stage Rust client.
28 +pub type OpaqueValue = Value;
29 +
30 +fn default_opaque_value_nil() -> OpaqueValue {
31 + Value::Nil
32 +}
33 +
34 +fn is_false(v: &bool) -> bool {
35 + !v
36 +}
37 +
38 +mod classified_outputs;
39 +pub mod dtype;
40 +pub mod handshake;
41 +pub mod logprobs;
42 +pub mod multimodal;
43 +pub mod stats;
44 +pub mod tensor;
45 +pub use classified_outputs::{
46 + ClassifiedEngineCoreOutputs, DpControlMessage, RequestBatchOutputs, UtilityCallOutput,
47 +};
48 +pub use dtype::ModelDtype;
49 +pub use logprobs::decode_engine_core_outputs;
50 +
51 +/// Request types are encoded as single-byte protocol constants so they can be
52 +/// sent over the ZMQ socket without an extra encoding step.
53 +///
54 +/// Original Python definition:
55 +/// <https://github.com/vllm-project/vllm/blob/f22d6e026798a74e6542a52ef776c054f2de572a/vllm/v1/engine/__init__.py#L217-L228>
56 +#[derive(Debug, Clone, Copy, PartialEq, Eq)]
57 +#[repr(u8)]
58 +pub enum EngineCoreRequestType {
59 + Add = 0,
60 + Abort = 1,
61 + StartDpWave = 2,
62 + Utility = 3,
63 +}
64 +
65 +impl EngineCoreRequestType {
66 + pub fn to_frame(self) -> Bytes {
67 + Bytes::from_static(match self {
68 + Self::Add => b"\x00",
69 + Self::Abort => b"\x01",
70 + Self::StartDpWave => b"\x02",
71 + Self::Utility => b"\x03",
72 + })
73 + }
74 +}
75 +
76 +/// Reason a request finished: stop, length, abort, error, or repetition.
77 +///
78 +/// This mirrors the Python enum and uses integer encoding for compact wire
79 +/// representation.
80 +///
81 +/// Original Python definition:
82 +/// <https://github.com/vllm-project/vllm/blob/f22d6e026798a74e6542a52ef776c054f2de572a/vllm/v1/engine/__init__.py#L41-L63>
83 +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize_repr, Deserialize_repr)]
84 +#[repr(u8)]
85 +pub enum EngineCoreFinishReason {
86 + /// A stop string was emitted.
87 + Stop = 0,
88 + /// `max_tokens` or `max_model_len` was reached.
89 + Length = 1,
90 + /// The request was aborted by the client.
91 + Abort = 2,
92 + /// A retryable request-level internal error occurred.
93 + Error = 3,
94 + /// A repetitive token pattern was detected.
95 + Repetition = 4,
96 +}
97 +
98 +/// Event types emitted by engine-core for one request.
99 +///
100 +/// Original Python definition:
101 +/// <https://github.com/vllm-project/vllm/blob/f22d6e026798a74e6542a52ef776c054f2de572a/vllm/v1/engine/__init__.py#L113-L118>
102 +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize_repr, Deserialize_repr)]
103 +#[repr(u8)]
104 +pub enum EngineCoreEventType {
105 + Queued = 1,
106 + Scheduled = 2,
107 + Preempted = 3,
108 +}
109 +
110 +/// A timestamped engine-core event associated with one request.
111 +///
112 +/// Original Python definition:
113 +/// <https://github.com/vllm-project/vllm/blob/f22d6e026798a74e6542a52ef776c054f2de572a/vllm/v1/engine/__init__.py#L121-L130>
114 +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
115 +pub struct EngineCoreEvent {
116 + pub r#type: EngineCoreEventType,
117 + pub timestamp: f64,
118 +}
119 +
120 +/// Controls how intermediate outputs are returned to the frontend.
121 +///
122 +/// `Cumulative = 0` is intentionally not supported in Rust frontend.
123 +///
124 +/// Original Python definition:
125 +/// <https://github.com/vllm-project/vllm/blob/f22d6e026798a74e6542a52ef776c054f2de572a/vllm/sampling_params.py#L146-L152>
126 +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize_repr, Deserialize_repr)]
127 +#[repr(u8)]
128 +pub enum RequestOutputKind {
129 + /// Return only token deltas in each update.
130 + #[default]
131 + Delta = 1,
132 + /// Suppress intermediate updates and return only the final output.
133 + FinalOnly = 2,
134 +}
135 +
136 +/// The stop reason associated with a finished output.
137 +///
138 +/// Python models this as the union-typed `stop_reason: int | str | None`
139 +/// field on `EngineCoreOutput`; the Rust client narrows it into a tagged enum.
140 +///
141 +/// Original Python field:
142 +/// <https://github.com/vllm-project/vllm/blob/f22d6e026798a74e6542a52ef776c054f2de572a/vllm/v1/engine/__init__.py#L155>
143 +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
144 +#[serde(untagged)]
145 +pub enum StopReason {
146 + TokenId(u32),
147 + Text(String),
148 +}
149 +
150 +/// Parameters for configuring structured outputs (guided decoding).
151 +///
152 +/// Exactly one constraint field (`json`, `regex`, `choice`, `grammar`,
153 +/// `json_object`, or `structural_tag`) should be set. The engine-core
154 +/// backend selects the appropriate grammar compiler based on which field
155 +/// is present.
156 +///
157 +/// Original Python definition:
158 +/// <https://github.com/vllm-project/vllm/blob/f22d6e026798a74e6542a52ef776c054f2de572a/vllm/sampling_params.py#L36-L107>
159 +#[serde_with::skip_serializing_none]
160 +#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
161 +#[serde(default)]
162 +pub struct StructuredOutputsParams {
163 + /// JSON schema (as a dict/object or JSON string) constraining the output.
164 + pub json: Option<serde_json::Value>,
165 + /// Regular expression the output must match.
166 + pub regex: Option<String>,
167 + /// List of allowed output strings (the model must produce one of these).
168 + pub choice: Option<Vec<String>>,
169 + /// Context-free grammar (in EBNF-like notation) the output must conform to.
170 + pub grammar: Option<String>,
171 + /// When `true`, output must be valid JSON (free-form, no schema).
172 + pub json_object: Option<bool>,
173 + /// Disable any additional whitespace in guided JSON output.
174 + #[serde(skip_serializing_if = "crate::protocol::is_false")]
175 + pub disable_any_whitespace: bool,
176 + /// Disable `additionalProperties` in JSON schema output.
177 + #[serde(skip_serializing_if = "crate::protocol::is_false")]
178 + pub disable_additional_properties: bool,
179 + /// Custom whitespace pattern for guided JSON output.
180 + pub whitespace_pattern: Option<String>,
181 + /// Structural tag configuration (JSON-encoded string).
182 + pub structural_tag: Option<String>,
183 +}
184 +
185 +/// Engine-core-facing sampling parameters for text generation.
186 +///
187 +/// This is the normalized southbound subset used by the Rust frontend when it
188 +/// talks to Python engine-core over the wire. User-facing request semantics
189 +/// such as `stop` strings, `n`, `ignore_eos`, and output aggregation mode are
190 +/// intentionally handled by higher layers before values reach this DTO.
191 +///
192 +/// Original Python definition:
193 +/// <https://github.com/vllm-project/vllm/blob/f22d6e026798a74e6542a52ef776c054f2de572a/vllm/sampling_params.py#L155-L291>
194 +#[serde_with::skip_serializing_none]
195 +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
196 +pub struct EngineCoreSamplingParams {
197 + /// Controls randomness. Lower values are more deterministic; zero means
198 + /// greedy sampling.
199 + pub temperature: f32,
200 + /// Cumulative probability threshold for nucleus sampling.
201 + pub top_p: f32,
202 + /// Maximum number of top tokens to consider. `0` means all tokens.
203 + pub top_k: u32,
204 + /// Random seed used by the sampler when present.
205 + pub seed: Option<i64>,
206 + /// Maximum number of tokens to generate per output sequence.
207 + pub max_tokens: u32,
208 + /// Minimum number of tokens to generate before EOS or stop-token handling.
209 + pub min_tokens: u32,
210 + /// Number of log probabilities to return per generated token.
211 + ///
212 + /// `None` disables sample logprobs. `-1` requests the full vocabulary.
213 + pub logprobs: Option<i32>,
214 + /// Number of log probabilities to return per prompt token.
215 + ///
216 + /// `None` disables prompt logprobs. `-1` requests the full vocabulary.
217 + pub prompt_logprobs: Option<i32>,
218 + /// Minimum probability threshold for token sampling.
219 + pub min_p: f32,
220 + /// Frequency penalty applied by the sampler.
221 + pub frequency_penalty: f32,
222 + /// Presence penalty applied by the sampler.
223 + pub presence_penalty: f32,
224 + /// Repetition penalty applied by the sampler.
225 + pub repetition_penalty: f32,
226 + /// Token IDs that stop generation.
227 + pub stop_token_ids: Vec<u32>,
228 + /// Primary EOS token ID used by engine-core's dedicated EOS stop path.
229 + ///
230 + /// This mirrors Python's internal `_eos_token_id` field and is derived by
231 + /// the frontend from tokenizer/model metadata rather than supplied directly
232 + /// by end users.
233 + #[serde(rename = "_eos_token_id")]
234 + pub eos_token_id: Option<u32>,
235 + /// Complete stop-token set used by engine-core for `min_tokens` masking.
236 + ///
237 + /// This mirrors Python's internal `_all_stop_token_ids` field and should
238 + /// contain explicit `stop_token_ids` plus any frontend-derived EOS token
239 + /// IDs.
240 + #[serde(rename = "_all_stop_token_ids")]
241 + pub all_stop_token_ids: BTreeSet<u32>,
242 + /// Logit biases to apply during sampling.
243 + /// Keys are token IDs
244 + #[serde(default)]
245 + pub logit_bias: Option<HashMap<u32, f32>>,
246 + /// Restrict output to these token IDs only.
247 + #[serde(default)]
248 + pub allowed_token_ids: Option<Vec<u32>>,
249 + /// Tokenized bad words to avoid during generation.
250 + #[serde(default, rename = "_bad_words_token_ids")]
251 + pub bad_words_token_ids: Option<Vec<Vec<u32>>>,
252 + /// Parameters for configuring structured outputs (guided decoding).
253 + #[serde(default)]
254 + pub structured_outputs: Option<StructuredOutputsParams>,
255 + /// Specific token IDs for which log probabilities should be returned at
256 + /// each position.
257 + ///
258 + /// When set, the engine returns logprobs for exactly these tokens in
259 + /// addition to the sampled/scored token. Mutually exclusive with the
260 + /// `logprobs` count field in practice.
261 + #[serde(default)]
262 + pub logprob_token_ids: Option<Vec<u32>>,
263 + /// If `Some(true)`, the request will not attempt to read from the prefix
264 + /// cache; newly computed blocks may still populate the cache. `None`
265 + /// defers to engine-core defaults.
266 + #[serde(default)]
267 + pub skip_reading_prefix_cache: Option<bool>,
268 + /// Additional request parameters for custom extensions (from `vllm_xargs`).
269 + #[serde(default)]
270 + pub extra_args: Option<HashMap<String, serde_json::Value>>,
271 +}
272 +
273 +impl EngineCoreSamplingParams {
274 + /// Constructs a default sampling params for testing purposes only.
275 + pub fn for_test() -> Self {
276 + Self {
277 + temperature: 1.0,
278 + top_p: 1.0,
279 + top_k: 0,
280 + seed: None,
281 + max_tokens: 65536,
282 + min_tokens: 0,
283 + logprobs: None,
284 + prompt_logprobs: None,
285 + min_p: 0.0,
286 + frequency_penalty: 0.0,
287 + presence_penalty: 0.0,
288 + repetition_penalty: 1.0,
289 + stop_token_ids: Vec::new(),
290 + eos_token_id: None,
291 + all_stop_token_ids: BTreeSet::new(),
292 + logit_bias: None,
293 + allowed_token_ids: None,
294 + bad_words_token_ids: None,
295 + structured_outputs: None,
296 + logprob_token_ids: None,
297 + skip_reading_prefix_cache: None,
298 + extra_args: None,
299 + }
300 + }
301 +}
302 +
303 +/// Engine-core add-request payload sent from frontend to engine.
304 +///
305 +/// Original Python definition:
306 +/// <https://github.com/vllm-project/vllm/blob/3f5bd482f5c1a5dbdffbbf68d624e20bb7032013/vllm/v1/engine/__init__.py#L80-L129>
307 +#[derive(Debug, Clone, PartialEq, Serialize_tuple, Deserialize_tuple, DefaultFromSerde)]
308 +pub struct EngineCoreRequest {
309 + pub request_id: String,
310 + pub prompt_token_ids: Option<Vec<u32>>,
311 + /// Multimodal features attached to the request.
312 + pub mm_features: Option<MmFeatures>,
313 + pub sampling_params: Option<EngineCoreSamplingParams>,
314 + /// Pooling parameters are preserved in the schema but not yet strongly
315 + /// typed.
316 + pub pooling_params: Option<OpaqueValue>,
317 + pub arrival_time: f64,
318 + #[serde(default)]
319 + pub lora_request: Option<OpaqueValue>,
320 + #[serde(default)]
321 + pub cache_salt: Option<String>,
322 + #[serde(default)]
323 + pub data_parallel_rank: Option<u32>,
324 + /// Unsupported in the first-stage Rust client because Python uses a custom
325 + /// tensor/aux-frame encoding path for this field.
326 + #[serde(default)]
327 + pub prompt_embeds: Option<OpaqueValue>,
328 + /// Per-position mask for mixed-mode inputs (e.g. chat completion with
329 + /// `prompt_embeds` content parts). `Some(true)` means real token id;
330 + /// `Some(false)` means the position uses a pre-computed entry from
331 + /// `prompt_embeds`. `None` for pure-tokens and pure-embeds requests.
332 + #[serde(default)]
333 + pub prompt_is_token_ids: Option<Vec<bool>>,
334 + /// Index of the client, used to ensure outputs are sent back to the same
335 + /// client when scaling out the frontend.
336 + #[serde(default)]
337 + pub client_index: u32,
338 + /// In DP mode, indicates which wave this request is expected to belong to.
339 + #[serde(default)]
340 + pub current_wave: u32,
341 + #[serde(default)]
342 + pub priority: i32,
343 + #[serde(default)]
344 + pub trace_headers: Option<BTreeMap<String, String>>,
345 + #[serde(default)]
346 + pub resumable: bool,
347 + /// Original user-provided request ID, used for output reporting and aborts.
348 + #[serde(default)]
349 + pub external_req_id: Option<String>,
350 + #[serde(default)]
351 + pub reasoning_ended: Option<bool>,
352 + /// Opaque reasoning-parser kwargs forwarded from the frontend to the
353 + /// structured-output backend.
354 + #[serde(default)]
355 + pub reasoning_parser_kwargs: Option<OpaqueValue>,
356 + /// If `true`, the request should be added to the scheduler's waiting queue
357 + /// and immediately aborted, so connector-side cleanup runs via the
358 + /// standard `request_finished` hook.
359 + #[serde(default)]
360 + pub abort_immediately: bool,
361 +}
362 +
363 +impl EngineCoreRequest {
364 + /// Validate fields intentionally not supported in the first-stage client.
365 + pub fn validate(&self) -> Result<()> {
366 + if self.prompt_embeds.is_some() {
367 + return Err(Error::UnsupportedField {
368 + context: "EngineCoreRequest",
369 + field: "prompt_embeds",
370 + });
371 + }
372 + Ok(())
373 + }
374 +}
375 +
376 +/// Engine-core utility call payload sent from frontend to engine.
377 +///
378 +/// Original Python payload shape:
379 +/// `(client_index, call_id, method_name, args)`
380 +#[derive(Debug, Clone, PartialEq, Serialize_tuple)]
381 +pub struct EngineCoreUtilityRequest {
382 + pub client_index: u32,
383 + pub call_id: i64,
384 + pub method_name: String,
385 + pub args: OpaqueValue,
386 +}
387 +
388 +impl EngineCoreUtilityRequest {
389 + /// Create a new utility request with the given strongly typed arguments,
390 + /// encoding them into the expected msgpack value format.
391 + pub fn new<T>(
392 + client_index: u32,
393 + call_id: i64,
394 + method_name: impl Into<String>,
395 + args: T,
396 + ) -> Result<Self>
397 + where
398 + T: Serialize + std::fmt::Debug,
399 + {
400 + let args = rmpv::ext::to_value(&args).map_err(|error| Error::Encode {
401 + target_type: type_name::<T>(),
402 + message: format!(
403 + "failed to encode utility args `{args:?}`: {}",
404 + error.to_report_string()
405 + ),
406 + })?;
407 + let args = match args {
408 + Value::Nil => Value::Array(Vec::new()),
409 + other => other,
410 + };
411 +
412 + Ok(Self {
413 + client_index,
414 + call_id,
415 + method_name: method_name.into(),
416 + args,
417 + })
418 + }
419 +}
420 +
421 +/// Engine-core output for a single request.
422 +///
423 +/// Original Python definition:
424 +/// <https://github.com/vllm-project/vllm/blob/d3af8c18317c0dc008d42e4367fbb9045cfb7bf6/vllm/v1/engine/__init__.py#L154-L184>
425 +#[derive(Debug, Clone, PartialEq, Serialize_tuple, Deserialize_tuple, DefaultFromSerde)]
426 +pub struct EngineCoreOutput {
427 + pub request_id: String,
428 + pub new_token_ids: Vec<u32>,
429 + /// Decoded sample logprobs for the newly generated positions in this
430 + /// output.
431 + #[serde(default)]
432 + pub new_logprobs: Option<MaybeWireLogprobs>,
433 + /// Decoded prompt logprobs for the scored prompt positions emitted in this
434 + /// output.
435 + #[serde(default)]
436 + pub new_prompt_logprobs_tensors: Option<MaybeWireLogprobs>,
437 + #[serde(default)]
438 + pub pooling_output: Option<OpaqueValue>,
439 + #[serde(default)]
440 + pub finish_reason: Option<EngineCoreFinishReason>,
441 + #[serde(default)]
442 + pub stop_reason: Option<StopReason>,
443 + #[serde(default)]
444 + pub events: Option<Vec<EngineCoreEvent>>,
445 + #[serde(default)]
446 + pub kv_transfer_params: Option<serde_json::Value>,
447 + #[serde(default)]
448 + pub trace_headers: Option<OpaqueValue>,
449 + /// Breakdown of the scheduled prefill computation, set on the first output
450 + /// of a newly scheduled prefill and elided for subsequent decode outputs.
451 + #[serde(default)]
452 + pub prefill_stats: Option<PrefillStats>,
453 + #[serde(default)]
454 + pub routed_experts: Option<OpaqueValue>,
455 + /// Number of NaNs seen in logits. Values above zero indicate corruption.
456 + #[serde(default)]
457 + pub num_nans_in_logits: u32,
458 +}
459 +
460 +impl EngineCoreOutput {
461 + /// Returns whether this output is terminal for the request.
462 + pub fn finished(&self) -> bool {
463 + self.finish_reason.is_some()
464 + }
465 +}
466 +
467 +/// Result of a utility call.
468 +///
469 +/// Original Python definition:
470 +/// <https://github.com/vllm-project/vllm/blob/f22d6e026798a74e6542a52ef776c054f2de572a/vllm/v1/engine/__init__.py#L174-L183>
471 +#[derive(Debug, Clone, PartialEq, Serialize_tuple, Deserialize_tuple, DefaultFromSerde)]
472 +pub struct UtilityOutput {
473 + pub call_id: i64,
474 + /// Non-`None` implies the call failed and `result` should be ignored.
475 + #[serde(default)]
476 + pub failure_message: Option<String>,
477 + #[serde(default)]
478 + pub result: Option<UtilityResultEnvelope>,
479 +}
480 +
481 +/// Python `UtilityResult` wrapper carried inside `UtilityOutput.result`.
482 +///
483 +/// Upstream reference:
484 +/// <https://github.com/vllm-project/vllm/blob/bc2c0c86efb28e77677a3cfb8687e976914a313a/vllm/v1/serial_utils.py#L178-L185>
485 +#[derive(Debug, Clone, PartialEq, Serialize_tuple, Deserialize_tuple)]
486 +pub struct UtilityResultEnvelope {
487 + /// Recursive type information encoded on Python side, serving as the hint
488 + /// for deserialization. We don't care it here as in Rust frontend all
489 + /// utility calls are strongly-typed.
490 + #[serde(default)]
491 + type_info: Option<OpaqueValue>,
492 + /// The actual utility result.
493 + #[serde(default = "default_opaque_value_nil")]
494 + result: OpaqueValue,
495 +}
496 +
497 +impl UtilityResultEnvelope {
498 + /// Create a utility result envelope without type information.
499 + pub fn without_type_info(result: OpaqueValue) -> Self {
500 + Self {
501 + type_info: None,
502 + result,
503 + }
504 + }
505 +}
506 +
507 +impl UtilityOutput {
508 + /// Decode the typed result of a utility call.
509 + pub fn into_typed_result<T>(self, method: &str) -> Result<T>
510 + where
511 + T: serde::de::DeserializeOwned,
512 + {
513 + if let Some(message) = self.failure_message {
514 + return Err(Error::UtilityCallFailed {
515 + method: method.to_string(),
516 + call_id: self.call_id,
517 + message,
518 + });
519 + }
520 +
521 + let result = self.result.map(|e| e.result).unwrap_or(Value::Nil);
522 +
523 + rmpv::ext::from_value(result).map_err(|error| Error::UtilityResultDecode {
524 + method: method.to_string(),
525 + call_id: self.call_id,
526 + message: error.to_report_string(),
527 + })
528 + }
529 +}
530 +
531 +/// Batch of engine-core outputs returned to a frontend client.
532 +///
533 +/// Original Python definition:
534 +/// <https://github.com/vllm-project/vllm/blob/f22d6e026798a74e6542a52ef776c054f2de572a/vllm/v1/engine/__init__.py#L186-L214>
535 +#[derive(Debug, Clone, PartialEq, Serialize_tuple, Deserialize_tuple, DefaultFromSerde)]
536 +pub struct EngineCoreOutputs {
537 + #[serde(default)]
538 + pub engine_index: u32,
539 + /// Outputs grouped for this client in the current engine tick.
540 + #[serde(default)]
541 + pub outputs: Vec<EngineCoreOutput>,
542 + #[serde(default)]
543 + pub scheduler_stats: Option<Box<SchedulerStats>>,
544 + #[serde(default)]
545 + pub timestamp: f64,
546 + #[serde(default)]
547 + pub utility_output: Option<UtilityOutput>,
548 + #[serde(default)]
549 + pub finished_requests: Option<BTreeSet<String>>,
550 + /// In DP mode, signals that the current wave finished and engines are
551 + /// paused.
552 + #[serde(default)]
553 + pub wave_complete: Option<u32>,
554 + /// In DP mode, signals that a request arrived for an old wave and the next
555 + /// wave needs to start in other engines.
556 + #[serde(default)]
557 + pub start_wave: Option<u32>,
558 +}
559 +
560 +/// Encode a Rust value into msgpack using the protocol crate's serde model.
561 +pub fn encode_msgpack<T>(value: &T) -> Result<Vec<u8>>
562 +where
563 + T: Serialize + std::fmt::Debug,
564 +{
565 + rmp_serde::to_vec_named(value).map_err(|error| Error::Encode {
566 + target_type: type_name::<T>(),
567 + message: format!(
568 + "failed to encode value `{:?}`: {}",
569 + value,
570 + error.to_report_string()
571 + ),
572 + })
573 +}
574 +
575 +/// Decode a msgpack payload into a strongly typed protocol value, with enhanced
576 +/// error reporting.
577 +pub fn decode_msgpack<T>(bytes: &[u8]) -> Result<T>
578 +where
579 + T: for<'de> Deserialize<'de>,
580 +{
581 + fn decode_value_preview(bytes: &[u8]) -> String {
582 + match decode_value(bytes) {
583 + Ok(value) => format!("{value}"),
584 + Err(error) => format!("<value decode failed: {error}>"),
585 + }
586 + }
587 +
588 + rmp_serde::from_slice(bytes).map_err(|error| Error::Decode {
589 + target_type: type_name::<T>(),
590 + message: format!("{error}; value fallback: {}", decode_value_preview(bytes)),
591 + })
592 +}
593 +
594 +pub fn decode_value(bytes: &[u8]) -> Result<Value> {
595 + Ok(rmpv::decode::read_value(&mut Cursor::new(bytes))?)
596 +}
597 +
598 +#[cfg(test)]
599 +mod tests {
600 + use std::collections::BTreeSet;
601 +
602 + use super::*;
603 +
604 + fn utility_result_value<T>(value: T) -> UtilityResultEnvelope
605 + where
606 + T: Serialize,
607 + {
608 + UtilityResultEnvelope::without_type_info(rmpv::ext::to_value(value).unwrap())
609 + }
610 +
611 + #[test]
612 + fn engine_core_request_serializes_as_full_array() {
613 + let request = EngineCoreRequest {
614 + request_id: "req-1".to_string(),
615 + prompt_token_ids: Some(vec![1, 2, 3]),
616 + sampling_params: Some(EngineCoreSamplingParams {
617 + max_tokens: 8,
618 + ..EngineCoreSamplingParams::for_test()
619 + }),
620 + arrival_time: 1234.5,
621 + client_index: 7,
622 + ..EngineCoreRequest::default()
623 + };
624 +
625 + let encoded = encode_msgpack(&request).unwrap();
626 + let value = decode_value(&encoded).unwrap();
627 + let array = match value {
628 + Value::Array(array) => array,
629 + other => panic!("expected array, got {other:?}"),
630 + };
631 +
632 + assert_eq!(array.len(), 20);
633 + assert_eq!(array[0], Value::from("req-1"));
634 + assert_eq!(array[2], Value::Nil);
635 + assert_eq!(array[4], Value::Nil);
636 + assert_eq!(array[10], Value::Nil);
637 + assert_eq!(array[11], Value::from(7));
638 + }
639 +
640 + #[test]
641 + fn engine_core_outputs_roundtrip_finished_fields() {
642 + let outputs = EngineCoreOutputs {
643 + outputs: vec![EngineCoreOutput {
644 + request_id: "req-1".to_string(),
645 + new_token_ids: vec![42],
646 + new_logprobs: None,
647 + new_prompt_logprobs_tensors: None,
648 + pooling_output: None,
649 + finish_reason: Some(EngineCoreFinishReason::Length),
650 + stop_reason: Some(StopReason::Text("stop".to_string())),
651 + events: None,
652 + kv_transfer_params: None,
653 + trace_headers: None,
654 + prefill_stats: None,
655 + routed_experts: None,
656 + num_nans_in_logits: 0,
657 + }],
658 + finished_requests: Some(BTreeSet::from(["req-1".to_string()])),
659 + ..Default::default()
660 + };
661 +
662 + let encoded = encode_msgpack(&outputs).unwrap();
663 + let decoded: EngineCoreOutputs = decode_msgpack(&encoded).unwrap();
664 +
665 + assert_eq!(decoded.outputs.len(), 1);
666 + assert_eq!(
667 + decoded.outputs[0].finish_reason,
668 + Some(EngineCoreFinishReason::Length)
669 + );
670 + assert_eq!(
671 + decoded.finished_requests,
672 + Some(BTreeSet::from(["req-1".to_string()]))
673 + );
674 + }
675 +
676 + #[test]
677 + fn utility_request_serializes_as_tuple_payload() {
678 + let request = EngineCoreUtilityRequest::new(7, 42, "is_sleeping", ()).unwrap();
679 +
680 + let encoded = encode_msgpack(&request).unwrap();
681 + let value = decode_value(&encoded).unwrap();
682 + let array = match value {
683 + Value::Array(array) => array,
684 + other => panic!("expected utility request array, got {other:?}"),
685 + };
686 +
687 + assert_eq!(array.len(), 4);
688 + assert_eq!(array[0], Value::from(7));
689 + assert_eq!(array[1], Value::from(42));
690 + assert_eq!(array[2], Value::from("is_sleeping"));
691 + assert_eq!(array[3], Value::Array(Vec::new()));
692 + }
693 +
694 + #[test]
695 + fn utility_output_decodes_typed_result() {
696 + let output = UtilityOutput {
697 + call_id: 9,
698 + failure_message: None,
699 + result: Some(utility_result_value(true)),
700 + };
701 +
702 + assert!(output.into_typed_result::<bool>("is_sleeping").unwrap());
703 + }
704 +
705 + #[test]
706 + fn utility_output_reports_failure_message() {
707 + let error = UtilityOutput {
708 + call_id: 9,
709 + failure_message: Some("boom".to_string()),
710 + result: None,
711 + }
712 + .into_typed_result::<bool>("is_sleeping")
713 + .unwrap_err();
714 +
715 + assert!(matches!(
716 + error,
717 + Error::UtilityCallFailed {
718 + method,
719 + call_id,
720 + message
721 + } if method == "is_sleeping" && call_id == 9 && message == "boom"
722 + ));
723 + }
724 +
725 + #[test]
726 + fn utility_output_decodes_missing_result_as_unit() {
727 + UtilityOutput {
728 + call_id: 3,
729 + failure_message: None,
730 + result: None,
731 + }
732 + .into_typed_result::<()>("reset_mm_cache")
733 + .unwrap();
734 + }
735 +
736 + #[test]
737 + fn utility_output_decodes_nil_result_as_unit() {
738 + UtilityOutput {
739 + call_id: 4,
740 + failure_message: None,
741 + result: Some(UtilityResultEnvelope::without_type_info(Value::Nil)),
742 + }
743 + .into_typed_result::<()>("sleep")
744 + .unwrap();
745 + }
746 +
747 + #[test]
748 + fn decode_msgpack_includes_type_name_and_value_fallback() {
749 + let error = decode_msgpack::<u64>(
750 + &rmp_serde::to_vec_named(&BTreeMap::from([("status", "READY")])).unwrap(),
751 + )
752 + .unwrap_err();
753 +
754 + expect_test::expect![[r#"messagepack decode failed for u64: wrong msgpack marker FixMap(1); value fallback: {"status": "READY"}"#]].assert_eq(&error.to_report_string());
755 + }
756 +}
rust/src/engine-core-client/src/protocol/multimodal.rs new
+282
@@ -0,0 +1,282 @@
1 +use std::collections::BTreeMap;
2 +
3 +use serde::{Deserialize, Serialize};
4 +use serde_tuple::{Deserialize_tuple, Serialize_tuple};
5 +
6 +use super::tensor::WireTensor;
7 +
8 +/// Multimodal feature payload accepted from higher-level frontend code.
9 +///
10 +/// Original Python definition:
11 +/// <https://github.com/vllm-project/vllm/blob/5a0a8fc1ea7542394ff315138bd5677b7b53bca1/vllm/v1/engine/__init__.py#L88>
12 +pub type MmFeatures = Vec<MmFeatureSpec>;
13 +
14 +/// Represents a single multimodal input with its processed data and metadata.
15 +///
16 +/// Used to track multimodal data through processing and caching. A request
17 +/// containing multiple multimodal items will have one `MmFeatureSpec`
18 +/// per item.
19 +///
20 +/// Original Python definition:
21 +/// <https://github.com/vllm-project/vllm/blob/5a0a8fc1ea7542394ff315138bd5677b7b53bca1/vllm/multimodal/inputs.py#L301-L332>
22 +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
23 +pub struct MmFeatureSpec {
24 + /// Represents multimodal data for this feature.
25 + ///
26 + /// Can be `None` if the item is cached, to skip IPC between API server
27 + /// and engine core processes.
28 + pub data: Option<MmKwargsItem>,
29 +
30 + /// The input modality, e.g., `"image"`, `"audio"`, `"video"`.
31 + pub modality: String,
32 +
33 + /// The hash for caching encoder outputs (with LoRA prefix if applicable).
34 + pub identifier: String,
35 +
36 + /// The location of the `modality` tokens corresponding to this item
37 + /// in the prompt, e.g., `PlaceholderRange(offset=2, length=336)`.
38 + pub mm_position: PlaceholderRange,
39 +
40 + /// The hash for caching processor outputs (without LoRA prefix).
41 + #[serde(default)]
42 + pub mm_hash: Option<String>,
43 +}
44 +
45 +/// Placeholder location information for multi-modal data.
46 +///
47 +/// Original Python definition:
48 +/// <https://github.com/vllm-project/vllm/blob/5a0a8fc1ea7542394ff315138bd5677b7b53bca1/vllm/multimodal/inputs.py#L118-L145>
49 +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
50 +pub struct PlaceholderRange {
51 + /// The start index of the placeholder in the prompt.
52 + pub offset: usize,
53 +
54 + /// The length of the placeholder.
55 + pub length: usize,
56 +
57 + /// A boolean mask of shape `(length,)` indicating which positions
58 + /// between `offset` and `offset + length` to assign embeddings to.
59 + /// `None` means all positions.
60 + #[serde(default)]
61 + pub is_embed: Option<WireTensor>,
62 +}
63 +
64 +/// A dictionary of processed keyword arguments to pass to the model,
65 +/// corresponding to a single item in `MultiModalDataItems`.
66 +///
67 +/// Original Python definition:
68 +/// <https://github.com/vllm-project/vllm/blob/5a0a8fc1ea7542394ff315138bd5677b7b53bca1/vllm/multimodal/inputs.py#L854-L871>
69 +pub type MmKwargsItem = BTreeMap<String, MmFieldElem>;
70 +
71 +/// Represents a processed keyword argument to pass to a model for a
72 +/// `MmKwargsItem`.
73 +///
74 +/// Original Python definition:
75 +/// <https://github.com/vllm-project/vllm/blob/5a0a8fc1ea7542394ff315138bd5677b7b53bca1/vllm/multimodal/inputs.py#L348-L369>
76 +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
77 +pub struct MmFieldElem {
78 + /// The processed value of this field in `MmKwargsItem`, i.e. the
79 + /// keyword argument value to be passed to the model.
80 + ///
81 + /// It may be set to `None` if it is determined that the item is cached
82 + /// in `EngineCore`.
83 + pub data: Option<MmKwargValue>,
84 +
85 + /// Defines how to combine this field's processed values with others in
86 + /// order to batch multi-modal items together for model inference.
87 + pub field: MmField,
88 +}
89 +
90 +/// Processed multimodal keyword argument value.
91 +///
92 +/// Original Python definition (`NestedTensors`) and wire encoding:
93 +/// <https://github.com/vllm-project/vllm/blob/5a0a8fc1ea7542394ff315138bd5677b7b53bca1/vllm/multimodal/inputs.py#L218-L226>
94 +/// <https://github.com/vllm-project/vllm/blob/5a0a8fc1ea7542394ff315138bd5677b7b53bca1/vllm/v1/serial_utils.py#L292-L299>
95 +/// <https://github.com/vllm-project/vllm/blob/5a0a8fc1ea7542394ff315138bd5677b7b53bca1/vllm/v1/serial_utils.py#L456-L465>
96 +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
97 +#[serde(untagged)]
98 +pub enum MmKwargValue {
99 + Tensor(WireTensor),
100 + Int(i64),
101 + Float(f64),
102 + List(Vec<MmKwargValue>),
103 +}
104 +
105 +/// Defines how to interpret tensor data belonging to a keyword argument for
106 +/// `MultiModalKwargsItems`, and vice versa.
107 +///
108 +/// Original Python definitions and wire encoding:
109 +/// <https://github.com/vllm-project/vllm/blob/5a0a8fc1ea7542394ff315138bd5677b7b53bca1/vllm/multimodal/inputs.py#L385-L630>
110 +/// <https://github.com/vllm-project/vllm/blob/5a0a8fc1ea7542394ff315138bd5677b7b53bca1/vllm/v1/serial_utils.py#L301-L310>
111 +/// <https://github.com/vllm-project/vllm/blob/5a0a8fc1ea7542394ff315138bd5677b7b53bca1/vllm/v1/serial_utils.py#L440-L454>
112 +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
113 +#[serde(try_from = "MmFieldWire", into = "MmFieldWire")]
114 +pub enum MmField {
115 + Batched(MmBatchedField),
116 + Flat(MmFlatField),
117 + Shared(MmSharedField),
118 +}
119 +
120 +/// Info: `MultiModalFieldConfig.batched`.
121 +///
122 +/// Original Python definition:
123 +/// <https://github.com/vllm-project/vllm/blob/5a0a8fc1ea7542394ff315138bd5677b7b53bca1/vllm/multimodal/inputs.py#L385-L502>
124 +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
125 +#[serde(deny_unknown_fields)]
126 +pub struct MmBatchedField {
127 + /// If `True`, then this field is excluded from being moved to the
128 + /// accelerator when multimodal items are grouped and batched.
129 + pub keep_on_cpu: bool,
130 +}
131 +
132 +/// Info: `MultiModalFieldConfig.flat` and
133 +/// `MultiModalFieldConfig.flat_from_sizes`.
134 +///
135 +/// Original Python definition:
136 +/// <https://github.com/vllm-project/vllm/blob/5a0a8fc1ea7542394ff315138bd5677b7b53bca1/vllm/multimodal/inputs.py#L385-L397>
137 +/// <https://github.com/vllm-project/vllm/blob/5a0a8fc1ea7542394ff315138bd5677b7b53bca1/vllm/multimodal/inputs.py#L505-L603>
138 +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
139 +#[serde(deny_unknown_fields)]
140 +pub struct MmFlatField {
141 + /// For each multi-modal item, a slice (`dim=0`) or a tuple of slices
142 + /// (`dim>0`) that is used to extract the data corresponding to it.
143 + pub slices: Vec<MmSlice>,
144 +
145 + /// The dimension to extract data, default to 0.
146 + pub dim: i32,
147 +
148 + /// If `True`, then this field is excluded from being moved to the
149 + /// accelerator when multimodal items are grouped and batched.
150 + pub keep_on_cpu: bool,
151 +}
152 +
153 +/// Info: `MultiModalFieldConfig.shared`.
154 +///
155 +/// Original Python definition:
156 +/// <https://github.com/vllm-project/vllm/blob/5a0a8fc1ea7542394ff315138bd5677b7b53bca1/vllm/multimodal/inputs.py#L385-L397>
157 +/// <https://github.com/vllm-project/vllm/blob/5a0a8fc1ea7542394ff315138bd5677b7b53bca1/vllm/multimodal/inputs.py#L606-L630>
158 +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
159 +#[serde(deny_unknown_fields)]
160 +pub struct MmSharedField {
161 + pub batch_size: usize,
162 +
163 + /// If `True`, then this field is excluded from being moved to the
164 + /// accelerator when multimodal items are grouped and batched.
165 + pub keep_on_cpu: bool,
166 +}
167 +
168 +/// Python slice encoded as `(start, stop, step)`.
169 +///
170 +/// Original Python wire encoding:
171 +/// <https://github.com/vllm-project/vllm/blob/5a0a8fc1ea7542394ff315138bd5677b7b53bca1/vllm/v1/serial_utils.py#L199-L204>
172 +#[derive(Debug, Clone, PartialEq, Eq, Serialize_tuple, Deserialize_tuple)]
173 +pub struct SliceSpec {
174 + pub start: Option<isize>,
175 + pub stop: Option<isize>,
176 + pub step: Option<isize>,
177 +}
178 +
179 +/// A single slice or a tuple of slices used by `MmFlatField`.
180 +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
181 +#[serde(untagged)]
182 +pub enum MmSlice {
183 + Slice(SliceSpec),
184 + Slices(Vec<SliceSpec>),
185 +}
186 +
187 +#[derive(Debug, Clone, PartialEq, Eq, Serialize_tuple, Deserialize_tuple)]
188 +struct MmFieldWire {
189 + name: String,
190 + inner: MmFieldWireInner,
191 +}
192 +
193 +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
194 +#[serde(untagged)]
195 +enum MmFieldWireInner {
196 + Batched(MmBatchedField),
197 + Flat(MmFlatField),
198 + Shared(MmSharedField),
199 +}
200 +
201 +impl TryFrom<MmFieldWire> for MmField {
202 + type Error = String;
203 +
204 + fn try_from(value: MmFieldWire) -> Result<Self, Self::Error> {
205 + match (value.name.as_str(), value.inner) {
206 + ("batched", MmFieldWireInner::Batched(kwargs)) => Ok(Self::Batched(kwargs)),
207 + ("flat", MmFieldWireInner::Flat(kwargs)) => Ok(Self::Flat(kwargs)),
208 + ("shared", MmFieldWireInner::Shared(kwargs)) => Ok(Self::Shared(kwargs)),
209 + (name, _) => Err(format!(
210 + "mismatched or unknown multimodal field factory {name:?}"
211 + )),
212 + }
213 + }
214 +}
215 +
216 +impl From<MmField> for MmFieldWire {
217 + fn from(value: MmField) -> Self {
218 + match value {
219 + MmField::Batched(kwargs) => Self {
220 + name: "batched".to_string(),
221 + inner: MmFieldWireInner::Batched(kwargs),
222 + },
223 + MmField::Flat(kwargs) => Self {
224 + name: "flat".to_string(),
225 + inner: MmFieldWireInner::Flat(kwargs),
226 + },
227 + MmField::Shared(kwargs) => Self {
228 + name: "shared".to_string(),
229 + inner: MmFieldWireInner::Shared(kwargs),
230 + },
231 + }
232 + }
233 +}
234 +
235 +#[cfg(test)]
236 +mod tests {
237 + use std::io::Cursor;
238 +
239 + use rmpv::Value;
240 +
241 + use super::*;
242 +
243 + fn encode_value<T: Serialize>(value: &T) -> Value {
244 + let bytes = rmp_serde::to_vec_named(value).expect("encode value");
245 + rmpv::decode::read_value(&mut Cursor::new(bytes)).expect("decode value")
246 + }
247 +
248 + #[test]
249 + fn multimodal_field_serializes_to_python_factory_tuple() {
250 + let field = MmField::Flat(MmFlatField {
251 + slices: vec![MmSlice::Slice(SliceSpec {
252 + start: Some(0),
253 + stop: Some(1200),
254 + step: None,
255 + })],
256 + dim: 0,
257 + keep_on_cpu: false,
258 + });
259 +
260 + let value = encode_value(&field);
261 + let Value::Array(items) = value else {
262 + panic!("field should encode as tuple array");
263 + };
264 + assert_eq!(items.len(), 2);
265 + assert_eq!(items[0].as_str(), Some("flat"));
266 +
267 + let Value::Map(kwargs) = &items[1] else {
268 + panic!("field kwargs should encode as map");
269 + };
270 + assert!(kwargs.iter().any(|(key, _)| key.as_str() == Some("slices")));
271 + assert!(kwargs.iter().any(|(key, _)| key.as_str() == Some("dim")));
272 + assert!(kwargs.iter().any(|(key, _)| key.as_str() == Some("keep_on_cpu")));
273 + }
274 +
275 + #[test]
276 + fn multimodal_field_round_trips_python_factory_tuple() {
277 + let field = MmField::Batched(MmBatchedField { keep_on_cpu: true });
278 + let encoded = rmp_serde::to_vec_named(&field).expect("encode field");
279 + let decoded: MmField = rmp_serde::from_slice(&encoded).expect("decode field");
280 + assert_eq!(decoded, field);
281 + }
282 +}
rust/src/engine-core-client/src/protocol/stats.rs new
+192
@@ -0,0 +1,192 @@
1 +use std::collections::BTreeMap;
2 +
3 +use serde::{Deserialize, Serialize};
4 +
5 +use crate::protocol::OpaqueValue;
6 +
7 +/// Stores cache hit statistics.
8 +///
9 +/// Original Python definition:
10 +/// <https://github.com/vllm-project/vllm/blob/bc2c0c86efb28e77677a3cfb8687e976914a313a/vllm/v1/metrics/stats.py#L18-L32>
11 +#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
12 +pub struct BaseCacheStats {
13 + /// Whether the cache was reset.
14 + pub reset: bool,
15 + /// The number of requests in this update.
16 + pub requests: u64,
17 + /// The number of queries in these requests.
18 + pub queries: u64,
19 + /// The number of hits in these requests.
20 + pub hits: u64,
21 +}
22 +
23 +/// Stores prefix cache hit statistics.
24 +/// - `reset`: Whether `reset_prefix_cache` was invoked.
25 +/// - `queries`: Refers to the number of tokens that were queried.
26 +///
27 +/// Original Python definition:
28 +/// <https://github.com/vllm-project/vllm/blob/bc2c0c86efb28e77677a3cfb8687e976914a313a/vllm/v1/metrics/stats.py#L114-L143>
29 +#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
30 +pub struct PrefixCacheStats {
31 + /// Embedded base cache counters and reset flag.
32 + #[serde(flatten)]
33 + pub base: BaseCacheStats,
34 + /// The number of previously preempted requests in this update.
35 + pub preempted_requests: u64,
36 + /// The `queries` number for preempted requests.
37 + pub preempted_queries: u64,
38 + /// The `hits` number for preempted requests.
39 + pub preempted_hits: u64,
40 +}
41 +
42 +/// Single KV cache block eviction sample.
43 +///
44 +/// Original Python definition:
45 +/// <https://github.com/vllm-project/vllm/blob/bc2c0c86efb28e77677a3cfb8687e976914a313a/vllm/v1/metrics/stats.py#L161-L167>
46 +#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
47 +pub struct KvCacheEvictionEvent {
48 + /// Lifetime from allocation to eviction.
49 + pub lifetime_seconds: f64,
50 + /// Idle time observed before eviction.
51 + pub idle_seconds: f64,
52 + /// Time gaps between consecutive accesses before eviction.
53 + pub reuse_gaps_seconds: Vec<f64>,
54 +}
55 +
56 +/// Per-step iteration decoding stats from scheduler.
57 +///
58 +/// Each scheduler step, statistics on spec decoding performance are aggregated
59 +/// across requests by the scheduler and returned to the frontend in
60 +/// `EngineCoreOutputs -> SchedulerStats`.
61 +///
62 +/// Original Python definition:
63 +/// <https://github.com/vllm-project/vllm/blob/bc2c0c86efb28e77677a3cfb8687e976914a313a/vllm/v1/spec_decode/metrics.py#L16-L44>
64 +#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
65 +pub struct SpecDecodingStats {
66 + /// Configured speculative token count for this scheduler.
67 + pub num_spec_tokens: u64,
68 + /// Number of drafted speculative decoding attempts.
69 + pub num_drafts: u64,
70 + /// Number of drafted tokens.
71 + pub num_draft_tokens: u64,
72 + /// Number of accepted drafted tokens.
73 + pub num_accepted_tokens: u64,
74 + /// Accepted drafted tokens counted by draft position.
75 + pub num_accepted_tokens_per_pos: Vec<u64>,
76 +}
77 +
78 +/// Breakdown of a scheduled prefill computation.
79 +///
80 +/// Python models this as a plain `@dataclass`, so it is serialized by msgspec
81 +/// as a map (named fields) rather than in the array-like form used by
82 +/// `EngineCoreOutput` itself.
83 +///
84 +/// Original Python definition:
85 +/// <https://github.com/vllm-project/vllm/blob/d3af8c18317c0dc008d42e4367fbb9045cfb7bf6/vllm/v1/metrics/stats.py#L242-L273>
86 +#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
87 +pub struct PrefillStats {
88 + /// Total number of tokens to be prefilled.
89 + #[serde(default)]
90 + pub num_prompt_tokens: u32,
91 + /// Tokens to be prefilled locally (actual compute work).
92 + #[serde(default)]
93 + pub num_computed_tokens: u32,
94 + /// Tokens to be prefilled without actual compute work.
95 + #[serde(default)]
96 + pub num_cached_tokens: u32,
97 + /// Tokens to be prefilled from local prefix cache.
98 + #[serde(default)]
99 + pub num_local_cached_tokens: u32,
100 + /// Tokens to be prefilled from external KV transfer.
101 + #[serde(default)]
102 + pub num_external_cached_tokens: u32,
103 +}
104 +
105 +/// Stats for debugging the metrics calculation.
106 +///
107 +/// Original Python definition:
108 +/// <https://github.com/vllm-project/vllm/blob/bc2c0c86efb28e77677a3cfb8687e976914a313a/vllm/v1/metrics/perf.py#L46-L55>
109 +#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
110 +pub struct DebugPerfStats {
111 + /// Time spent calculating these stats.
112 + pub calc_duration: f64,
113 + /// Number of prefill requests included in the sampled batch.
114 + pub num_prefill_requests: u64,
115 + /// Number of decode requests included in the sampled batch.
116 + pub num_decode_requests: u64,
117 + /// Optional execution-context breakdown used for debugging.
118 + pub context_breakdown: Option<BTreeMap<String, u64>>,
119 + /// Optional per-component FLOPs breakdown.
120 + pub num_flops_per_gpu_breakdown: Option<BTreeMap<String, u64>>,
121 + /// Optional per-component memory-read breakdown.
122 + pub num_read_bytes_per_gpu_breakdown: Option<BTreeMap<String, u64>>,
123 + /// Optional per-component memory-write breakdown.
124 + pub num_write_bytes_per_gpu_breakdown: Option<BTreeMap<String, u64>>,
125 +}
126 +
127 +/// Original Python definition:
128 +/// <https://github.com/vllm-project/vllm/blob/bc2c0c86efb28e77677a3cfb8687e976914a313a/vllm/v1/metrics/perf.py#L58-L63>
129 +#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
130 +pub struct PerfStats {
131 + /// Estimated floating point operations per GPU.
132 + pub num_flops_per_gpu: u64,
133 + /// Estimated bytes read from memory per GPU.
134 + pub num_read_bytes_per_gpu: u64,
135 + /// Estimated bytes written to memory per GPU.
136 + pub num_write_bytes_per_gpu: u64,
137 + /// Optional debug-only perf derivation details.
138 + pub debug_stats: Option<DebugPerfStats>,
139 +}
140 +
141 +/// Original Python definition:
142 +/// <https://github.com/vllm-project/vllm/blob/bc2c0c86efb28e77677a3cfb8687e976914a313a/vllm/compilation/cuda_graph.py#L28-L33>
143 +#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
144 +pub struct CudagraphStat {
145 + /// Number of real tokens in the captured batch before padding.
146 + pub num_unpadded_tokens: u64,
147 + /// Number of padded tokens in the captured batch.
148 + pub num_padded_tokens: u64,
149 + /// Number of padding positions added for capture/runtime shape alignment.
150 + pub num_paddings: u64,
151 + /// Runtime mode string associated with this CUDA graph sample.
152 + pub runtime_mode: String,
153 +}
154 +
155 +/// Stats associated with the scheduler.
156 +///
157 +/// Original Python definition:
158 +/// <https://github.com/vllm-project/vllm/blob/906a8c15d0beded2cca8b357877266c631340f74/vllm/v1/metrics/stats.py#L170-L196>
159 +#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
160 +pub struct SchedulerStats {
161 + /// Number of requests in model execution batches.
162 + pub num_running_reqs: u64,
163 + /// Length of the "waiting" request queue.
164 + pub num_waiting_reqs: u64,
165 + /// Length of the "skipped waiting" queue.
166 + #[serde(default)]
167 + pub num_skipped_waiting_reqs: u64,
168 + /// Internal DP load-balancing step counter.
169 + pub step_counter: u64,
170 + /// Internal DP load-balancing wave number.
171 + pub current_wave: u64,
172 + /// KV-cache usage. `1.0` means 100% usage.
173 + pub kv_cache_usage: f64,
174 + /// Local prefix cache statistics.
175 + pub prefix_cache_stats: PrefixCacheStats,
176 + /// External connector prefix cache statistics, when configured.
177 + pub connector_prefix_cache_stats: Option<PrefixCacheStats>,
178 + /// Sampled KV cache eviction events for residency metrics.
179 + pub kv_cache_eviction_events: Vec<KvCacheEvictionEvent>,
180 + /// Speculative decoding scheduler stats, when enabled.
181 + pub spec_decoding_stats: Option<SpecDecodingStats>,
182 + /// Connector-specific KV transfer stats, kept opaque for now.
183 + pub kv_connector_stats: Option<BTreeMap<String, OpaqueValue>>,
184 + /// Waiting request counts per LoRA adapter.
185 + pub waiting_lora_adapters: BTreeMap<String, u64>,
186 + /// Running request counts per LoRA adapter.
187 + pub running_lora_adapters: BTreeMap<String, u64>,
188 + /// CUDA graph runtime stats when graph metrics are enabled.
189 + pub cudagraph_stats: Option<CudagraphStat>,
190 + /// Estimated MFU/performance stats, when enabled.
191 + pub perf_stats: Option<PerfStats>,
192 +}
rust/src/engine-core-client/src/protocol/tensor.rs new
+255
@@ -0,0 +1,255 @@
1 +use bytemuck::allocation::pod_collect_to_vec;
2 +use enum_as_inner::EnumAsInner;
3 +use half::{bf16, f16};
4 +use rmpv::Value;
5 +use serde::{Deserialize, Deserializer, Serialize, Serializer};
6 +use serde_tuple::{Deserialize_tuple, Serialize_tuple};
7 +
8 +/// Tensors and ndarrays are encoded with this extension type in Python.
9 +///
10 +/// Original Python definition:
11 +/// <https://github.com/vllm-project/vllm/blob/5a0a8fc1ea7542394ff315138bd5677b7b53bca1/vllm/v1/serial_utils.py#L41-L43>
12 +const CUSTOM_TYPE_RAW_VIEW: i8 = 3;
13 +
14 +#[easy_ext::ext(ShapeExt)]
15 +impl [usize] {
16 + /// Returned the total number of elements implied by this shape, or `None`
17 + /// if the product of the dimensions overflows `usize`.
18 + pub fn checked_numel(&self) -> Option<usize> {
19 + self.iter().try_fold(1usize, |acc, dim| acc.checked_mul(*dim))
20 + }
21 +}
22 +
23 +/// Python ndarray/tensor wire tuple encoded as `(dtype, shape, data)`.
24 +///
25 +/// This matches the custom msgpack representation built by Python
26 +/// `serial_utils.encode_ndarray` / `encode_tensor`.
27 +///
28 +/// Original Python wire encoders:
29 +/// <https://github.com/vllm-project/vllm/blob/5a0a8fc1ea7542394ff315138bd5677b7b53bca1/vllm/v1/serial_utils.py#L237-L273>
30 +/// <https://github.com/vllm-project/vllm/blob/5a0a8fc1ea7542394ff315138bd5677b7b53bca1/vllm/v1/serial_utils.py#L389-L425>
31 +#[derive(Debug, Clone, PartialEq, Serialize_tuple, Deserialize_tuple)]
32 +pub struct WireNdArray {
33 + pub dtype: String,
34 + pub shape: Vec<usize>,
35 + pub data: WireArrayData,
36 +}
37 +
38 +impl WireNdArray {
39 + /// Build a float32 tensor/ndarray backed by native-endian raw-view bytes.
40 + pub fn from_f32(shape: Vec<usize>, data: Vec<f32>) -> Result<Self, String> {
41 + validate_element_count(&shape, data.len())?;
42 + Ok(Self {
43 + dtype: "float32".to_string(),
44 + shape,
45 + data: WireArrayData::RawView(pod_collect_to_vec::<f32, u8>(&data)),
46 + })
47 + }
48 +
49 + /// Build a float16 tensor/ndarray backed by native-endian raw-view bytes.
50 + pub fn from_f16(shape: Vec<usize>, data: Vec<f16>) -> Result<Self, String> {
51 + validate_element_count(&shape, data.len())?;
52 + Ok(Self {
53 + dtype: "float16".to_string(),
54 + shape,
55 + data: WireArrayData::RawView(pod_collect_to_vec::<f16, u8>(&data)),
56 + })
57 + }
58 +
59 + /// Build a bfloat16 tensor/ndarray backed by native-endian raw-view bytes.
60 + pub fn from_bf16(shape: Vec<usize>, data: Vec<bf16>) -> Result<Self, String> {
61 + validate_element_count(&shape, data.len())?;
62 + Ok(Self {
63 + dtype: "bfloat16".to_string(),
64 + shape,
65 + data: WireArrayData::RawView(pod_collect_to_vec::<bf16, u8>(&data)),
66 + })
67 + }
68 +
69 + /// Build an int64 tensor/ndarray backed by native-endian raw-view bytes.
70 + pub fn from_i64(shape: Vec<usize>, data: Vec<i64>) -> Result<Self, String> {
71 + validate_element_count(&shape, data.len())?;
72 + Ok(Self {
73 + dtype: "int64".to_string(),
74 + shape,
75 + data: WireArrayData::RawView(pod_collect_to_vec::<i64, u8>(&data)),
76 + })
77 + }
78 +
79 + /// Build a uint32 tensor/ndarray backed by native-endian raw-view bytes.
80 + pub fn from_u32(shape: Vec<usize>, data: Vec<u32>) -> Result<Self, String> {
81 + validate_element_count(&shape, data.len())?;
82 + Ok(Self {
83 + dtype: "uint32".to_string(),
84 + shape,
85 + data: WireArrayData::RawView(pod_collect_to_vec::<u32, u8>(&data)),
86 + })
87 + }
88 +
89 + /// Build a bool tensor/ndarray backed by raw-view bytes.
90 + ///
91 + /// This matches `torch.bool` storage: one byte per element, not a packed
92 + /// bitmap. Values are canonicalized as `false -> 0` and `true -> 1`.
93 + pub fn from_bool(shape: Vec<usize>, data: Vec<bool>) -> Result<Self, String> {
94 + validate_element_count(&shape, data.len())?;
95 + Ok(Self {
96 + dtype: "bool".to_string(),
97 + shape,
98 + data: WireArrayData::RawView(data.iter().map(|value| u8::from(*value)).collect()),
99 + })
100 + }
101 +
102 + /// Build a tensor/ndarray from already-encoded raw-view bytes.
103 + ///
104 + /// Use this as an escape hatch when the caller already owns bytes that
105 + /// match the requested `dtype` and `shape`.
106 + pub fn from_raw(dtype: impl Into<String>, shape: Vec<usize>, data: Vec<u8>) -> Self {
107 + Self {
108 + dtype: dtype.into(),
109 + shape,
110 + data: WireArrayData::RawView(data),
111 + }
112 + }
113 +}
114 +
115 +/// Validate that the number of elements implied by the shape matches the length
116 +/// of the data.
117 +fn validate_element_count(shape: &[usize], len: usize) -> Result<(), String> {
118 + let expected = shape
119 + .checked_numel()
120 + .ok_or_else(|| format!("tensor shape product overflows usize: {shape:?}"))?;
121 + if expected == len {
122 + Ok(())
123 + } else {
124 + Err(format!(
125 + "tensor data length {len} does not match shape {shape:?} product {expected}"
126 + ))
127 + }
128 +}
129 +
130 +/// Python tensor wire tuple encoded as `(dtype, shape, data)`.
131 +///
132 +/// This is the same wire shape as [`WireNdArray`]; multimodal request payloads
133 +/// use it for `torch.Tensor` values.
134 +pub type WireTensor = WireNdArray;
135 +
136 +/// Python array/tensor payload reference inside [`WireNdArray`].
137 +///
138 +/// The data can be either an inline msgpack raw-view extension or an index into
139 +/// the multipart aux-frame list carried alongside the primary msgpack frame.
140 +///
141 +/// Original Python wire encoders:
142 +/// <https://github.com/vllm-project/vllm/blob/5a0a8fc1ea7542394ff315138bd5677b7b53bca1/vllm/v1/serial_utils.py#L237-L273>
143 +/// <https://github.com/vllm-project/vllm/blob/5a0a8fc1ea7542394ff315138bd5677b7b53bca1/vllm/v1/serial_utils.py#L389-L425>
144 +#[derive(Debug, Clone, PartialEq, EnumAsInner)]
145 +pub enum WireArrayData {
146 + /// The index of the aux frame where the raw bytes of this array/tensor are
147 + /// stored.
148 + AuxIndex(usize),
149 + /// The raw bytes of this array/tensor.
150 + RawView(Vec<u8>),
151 +}
152 +
153 +impl<'de> Deserialize<'de> for WireArrayData {
154 + fn deserialize<D>(deserializer: D) -> std::result::Result<Self, D::Error>
155 + where
156 + D: Deserializer<'de>,
157 + {
158 + let value = Value::deserialize(deserializer)?;
159 + match value {
160 + Value::Ext(tag, bytes) if tag == CUSTOM_TYPE_RAW_VIEW => Ok(Self::RawView(bytes)),
161 + Value::Ext(tag, _) => Err(serde::de::Error::custom(format!(
162 + "unsupported extension type code {tag}"
163 + ))),
164 + Value::Integer(index) => {
165 + index.as_u64().map(|index| Self::AuxIndex(index as usize)).ok_or_else(|| {
166 + serde::de::Error::custom("aux frame index must be a non-negative integer")
167 + })
168 + }
169 + other => Err(serde::de::Error::custom(format!(
170 + "expected raw-view ext or aux frame index, got {other:?}"
171 + ))),
172 + }
173 + }
174 +}
175 +
176 +impl Serialize for WireArrayData {
177 + fn serialize<S>(&self, serializer: S) -> std::result::Result<S::Ok, S::Error>
178 + where
179 + S: Serializer,
180 + {
181 + // TODO: outbound request serialization currently only supports inline
182 + // raw-view bytes. Emitting aux frames needs transport-level plumbing;
183 + // serializing `AuxIndex` here only preserves an already-built reference.
184 + match self {
185 + Self::AuxIndex(index) => serializer.serialize_u64(*index as u64),
186 + Self::RawView(bytes) => {
187 + Value::Ext(CUSTOM_TYPE_RAW_VIEW, bytes.clone()).serialize(serializer)
188 + }
189 + }
190 + }
191 +}
192 +
193 +#[cfg(test)]
194 +mod tests {
195 + use super::*;
196 +
197 + #[test]
198 + fn constructors_build_raw_view_tensors() {
199 + let f32_tensor = WireNdArray::from_f32(vec![2], vec![1.0, 2.5]).unwrap();
200 + assert_eq!(f32_tensor.dtype, "float32");
201 + assert_eq!(f32_tensor.shape, vec![2]);
202 + assert_eq!(
203 + f32_tensor.data.into_raw_view().expect("raw view"),
204 + [1.0_f32, 2.5].into_iter().flat_map(f32::to_ne_bytes).collect::<Vec<_>>()
205 + );
206 +
207 + let f16_tensor =
208 + WireNdArray::from_f16(vec![2], vec![f16::from_f32(1.0), f16::from_f32(2.5)]).unwrap();
209 + assert_eq!(f16_tensor.dtype, "float16");
210 + assert_eq!(f16_tensor.shape, vec![2]);
211 + assert_eq!(f16_tensor.data.into_raw_view().expect("raw view").len(), 4);
212 +
213 + let bf16_tensor =
214 + WireNdArray::from_bf16(vec![2], vec![bf16::from_f32(1.0), bf16::from_f32(2.5)])
215 + .unwrap();
216 + assert_eq!(bf16_tensor.dtype, "bfloat16");
217 + assert_eq!(bf16_tensor.shape, vec![2]);
218 + assert_eq!(bf16_tensor.data.into_raw_view().expect("raw view").len(), 4);
219 +
220 + let i64_tensor = WireNdArray::from_i64(vec![1], vec![-7]).unwrap();
221 + assert_eq!(i64_tensor.dtype, "int64");
222 + assert_eq!(
223 + i64_tensor.data.into_raw_view().expect("raw view"),
224 + (-7_i64).to_ne_bytes()
225 + );
226 +
227 + let u32_tensor = WireNdArray::from_u32(vec![1], vec![42]).unwrap();
228 + assert_eq!(u32_tensor.dtype, "uint32");
229 + assert_eq!(
230 + u32_tensor.data.into_raw_view().expect("raw view"),
231 + 42_u32.to_ne_bytes()
232 + );
233 +
234 + let bool_tensor = WireNdArray::from_bool(vec![2], vec![false, true]).unwrap();
235 + assert_eq!(bool_tensor.dtype, "bool");
236 + assert_eq!(
237 + bool_tensor.data.into_raw_view().expect("raw view"),
238 + vec![0, 1]
239 + );
240 +
241 + let raw_tensor = WireNdArray::from_raw("custom", vec![3], vec![1, 2, 3]);
242 + assert_eq!(raw_tensor.dtype, "custom");
243 + assert_eq!(raw_tensor.shape, vec![3]);
244 + assert_eq!(
245 + raw_tensor.data.into_raw_view().expect("raw view"),
246 + vec![1, 2, 3]
247 + );
248 + }
249 +
250 + #[test]
251 + fn constructors_validate_shape_product() {
252 + let err = WireNdArray::from_f32(vec![2, 2], vec![1.0, 2.0]).unwrap_err();
253 + assert!(err.contains("does not match shape"));
254 + }
255 +}
rust/src/engine-core-client/src/test_utils.rs new
+295
@@ -0,0 +1,295 @@
1 +use std::future::Future;
2 +use std::path::Path;
3 +use std::pin::Pin;
4 +use std::time::Duration;
5 +
6 +use tempfile::TempDir;
7 +use tokio::sync::oneshot;
8 +use zeromq::prelude::{Socket, SocketRecv, SocketSend};
9 +use zeromq::util::PeerIdentity;
10 +use zeromq::{DealerSocket, PushSocket, SocketOptions, SubSocket, ZmqMessage};
11 +
12 +use crate::EngineId;
13 +use crate::protocol::ModelDtype;
14 +use crate::protocol::handshake::{EngineCoreReadyResponse, HandshakeInitMessage, ReadyMessage};
15 +
16 +/// Per-test IPC endpoint namespace backed by a unique temporary directory.
17 +///
18 +/// Using one directory per test avoids endpoint collisions without requiring
19 +/// ad-hoc unique-name generation at each call site.
20 +#[derive(Debug)]
21 +pub struct IpcNamespace {
22 + dir: TempDir,
23 +}
24 +
25 +impl IpcNamespace {
26 + /// Create a fresh namespace for one test case.
27 + pub fn new() -> std::io::Result<Self> {
28 + Ok(Self {
29 + dir: TempDir::new()?,
30 + })
31 + }
32 +
33 + /// Build one `ipc://...` endpoint under this namespace.
34 + pub fn endpoint(&self, name: impl AsRef<Path>) -> String {
35 + let path = self.dir.path().join(name);
36 + format!("ipc://{}", path.to_string_lossy())
37 + }
38 +
39 + /// Endpoint used for the initial READY/HELLO handshake.
40 + pub fn handshake_endpoint(&self) -> String {
41 + self.endpoint("handshake.sock")
42 + }
43 +
44 + /// Endpoint used for engine-core request traffic.
45 + pub fn input_endpoint(&self) -> String {
46 + self.endpoint("input.sock")
47 + }
48 +
49 + /// Endpoint used for engine-core output traffic.
50 + pub fn output_endpoint(&self) -> String {
51 + self.endpoint("output.sock")
52 + }
53 +}
54 +
55 +/// Construct a standard local READY message used by mock engines in tests.
56 +fn ready_message(status: &str) -> ReadyMessage {
57 + ReadyMessage {
58 + status: Some(status.to_string()),
59 + local: Some(true),
60 + headless: Some(true),
61 + parallel_config_hash: None,
62 + }
63 +}
64 +
65 +/// Construct a default ready response payload for mock engine input
66 +/// registration.
67 +fn ready_response_payload() -> Vec<u8> {
68 + rmp_serde::to_vec_named(&EngineCoreReadyResponse {
69 + max_model_len: 4096,
70 + num_gpu_blocks: 0,
71 + dp_stats_address: None,
72 + dtype: Some(ModelDtype::Float32),
73 + })
74 + .expect("encode ready response payload")
75 +}
76 +
77 +/// Coordinator-side sockets connected by one mock engine when coordinator mode
78 +/// is enabled.
79 +pub struct MockCoordinatorConnections {
80 + /// Subscription socket that receives coordinator broadcasts such as
81 + /// `START_DP_WAVE`.
82 + pub input_sub: SubSocket,
83 + /// Push socket used to send coordinator-only `EngineCoreOutputs` back to
84 + /// the frontend.
85 + pub output_push: PushSocket,
86 +}
87 +
88 +/// Fully connected mock engine transport state used by tests.
89 +pub struct MockEngineConnections {
90 + /// Decoded INIT message sent by the frontend during handshake.
91 + pub init: HandshakeInitMessage,
92 + /// Socket used to receive frontend requests.
93 + pub dealer: DealerSocket,
94 + /// Socket used to publish normal request outputs back to the frontend.
95 + pub push: PushSocket,
96 + /// Optional coordinator sockets when the client enabled the in-process
97 + /// coordinator.
98 + pub coordinator: Option<MockCoordinatorConnections>,
99 +}
100 +
101 +/// Complete the engine-core handshake and connect mock input/output sockets
102 +/// plus optional coordinator sockets.
103 +pub async fn setup_mock_engine_connections(
104 + engine_handshake: String,
105 + engine_id: impl Into<EngineId>,
106 +) -> MockEngineConnections {
107 + // Wait for the client to bind the handshake socket before connecting.
108 + // A fixed sleep is racy under CI load; instead poll for the socket file.
109 + let socket_path = engine_handshake
110 + .strip_prefix("ipc://")
111 + .expect("handshake address must be ipc://");
112 + for _ in 0..100 {
113 + if Path::new(socket_path).exists() {
114 + break;
115 + }
116 + tokio::time::sleep(Duration::from_millis(20)).await;
117 + }
118 +
119 + let peer_identity = PeerIdentity::try_from(engine_id.into()).expect("peer id");
120 +
121 + let mut options = SocketOptions::default();
122 + options.peer_identity(peer_identity.clone());
123 + let mut handshake = DealerSocket::with_options(options);
124 + handshake
125 + .connect(&engine_handshake)
126 + .await
127 + .expect("connect mock engine handshake socket");
128 + handshake
129 + .send(ZmqMessage::from(
130 + rmp_serde::to_vec_named(&ready_message("HELLO")).expect("encode HELLO ready message"),
131 + ))
132 + .await
133 + .expect("send HELLO ready message");
134 +
135 + let init_frames = handshake.recv().await.expect("receive handshake init message").into_vec();
136 + assert_eq!(init_frames.len(), 1);
137 + let init: HandshakeInitMessage =
138 + rmp_serde::from_slice(init_frames[0].as_ref()).expect("decode handshake init message");
139 +
140 + let mut input_options = SocketOptions::default();
141 + input_options.peer_identity(peer_identity);
142 + let mut dealer = DealerSocket::with_options(input_options);
143 + dealer
144 + .connect(&init.addresses.inputs[0])
145 + .await
146 + .expect("connect mock engine input socket");
147 + dealer
148 + .send(ZmqMessage::from(ready_response_payload()))
149 + .await
150 + .expect("send mock engine input ready frame");
151 +
152 + let mut push = PushSocket::new();
153 + push.connect(&init.addresses.outputs[0])
154 + .await
155 + .expect("connect mock engine output socket");
156 +
157 + let coordinator = match (
158 + init.addresses.coordinator_input.as_deref(),
159 + init.addresses.coordinator_output.as_deref(),
160 + ) {
161 + (Some(coordinator_input), Some(coordinator_output)) => {
162 + let mut input_sub = SubSocket::new();
163 + input_sub
164 + .connect(coordinator_input)
165 + .await
166 + .expect("connect mock engine coordinator input socket");
167 + input_sub
168 + .subscribe("")
169 + .await
170 + .expect("subscribe mock engine coordinator input socket");
171 +
172 + let mut output_push = PushSocket::new();
173 + output_push
174 + .connect(coordinator_output)
175 + .await
176 + .expect("connect mock engine coordinator output socket");
177 +
178 + let ready =
179 + input_sub.recv().await.expect("receive coordinator READY marker").into_vec();
180 + assert_eq!(ready.len(), 1);
181 + assert_eq!(ready[0].as_ref(), b"READY");
182 +
183 + Some(MockCoordinatorConnections {
184 + input_sub,
185 + output_push,
186 + })
187 + }
188 + (None, None) => None,
189 + _ => panic!("coordinator handshake addresses must be both present or both absent"),
190 + };
191 +
192 + handshake
193 + .send(ZmqMessage::from(
194 + rmp_serde::to_vec_named(&ready_message("READY")).expect("encode READY ready message"),
195 + ))
196 + .await
197 + .expect("send READY ready message");
198 +
199 + MockEngineConnections {
200 + init,
201 + dealer,
202 + push,
203 + coordinator,
204 + }
205 +}
206 +
207 +/// Connect one mock engine directly to already-bootstrapped frontend
208 +/// input/output sockets.
209 +pub async fn setup_bootstrapped_mock_engine(
210 + input_address: String,
211 + output_address: String,
212 + engine_id: impl Into<EngineId>,
213 +) -> (DealerSocket, PushSocket) {
214 + for endpoint in [&input_address, &output_address] {
215 + if let Some(socket_path) = endpoint.strip_prefix("ipc://") {
216 + for _ in 0..100 {
217 + if Path::new(socket_path).exists() {
218 + break;
219 + }
220 + tokio::time::sleep(Duration::from_millis(20)).await;
221 + }
222 + }
223 + }
224 +
225 + let peer_identity = PeerIdentity::try_from(engine_id.into()).expect("peer id");
226 + let mut input_options = SocketOptions::default();
227 + input_options.peer_identity(peer_identity);
228 + let mut dealer = DealerSocket::with_options(input_options);
229 + dealer.connect(&input_address).await.expect("connect mock engine input socket");
230 + dealer
231 + .send(ZmqMessage::from(ready_response_payload()))
232 + .await
233 + .expect("send mock engine input ready frame");
234 +
235 + let mut push = PushSocket::new();
236 + push.connect(&output_address).await.expect("connect mock engine output socket");
237 +
238 + (dealer, push)
239 +}
240 +
241 +/// Complete the engine-core handshake and connect mock input/output sockets.
242 +///
243 +/// This returns the decoded handshake init message plus the `DealerSocket` used
244 +/// to receive client requests and the `PushSocket` used to send engine outputs
245 +/// back to the client.
246 +pub async fn setup_mock_engine_with_init(
247 + engine_handshake: String,
248 + engine_id: impl Into<EngineId>,
249 +) -> (HandshakeInitMessage, DealerSocket, PushSocket) {
250 + let MockEngineConnections {
251 + init, dealer, push, ..
252 + } = setup_mock_engine_connections(engine_handshake, engine_id).await;
253 + (init, dealer, push)
254 +}
255 +
256 +/// Complete the engine-core handshake and connect mock input/output sockets.
257 +///
258 +/// This returns the `DealerSocket` used to receive client requests and the
259 +/// `PushSocket` used to send engine outputs back to the client.
260 +pub async fn setup_mock_engine(
261 + engine_handshake: String,
262 + engine_id: impl Into<EngineId>,
263 +) -> (DealerSocket, PushSocket) {
264 + let (_, dealer, push) = setup_mock_engine_with_init(engine_handshake, engine_id).await;
265 + (dealer, push)
266 +}
267 +
268 +/// Spawn a mock engine task and keep its sockets alive until the returned
269 +/// shutdown sender is triggered by the test.
270 +///
271 +/// The script borrows the connected sockets mutably while it runs. After the
272 +/// script completes, this helper keeps the sockets alive until the test
273 +/// explicitly signals shutdown.
274 +pub fn spawn_mock_engine_task<F>(
275 + engine_handshake: String,
276 + engine_id: impl Into<EngineId>,
277 + run: F,
278 +) -> (oneshot::Sender<()>, tokio::task::JoinHandle<()>)
279 +where
280 + F: for<'a> FnOnce(
281 + &'a mut DealerSocket,
282 + &'a mut PushSocket,
283 + ) -> Pin<Box<dyn Future<Output = ()> + Send + 'a>>
284 + + Send
285 + + 'static,
286 +{
287 + let (shutdown_tx, shutdown_rx) = oneshot::channel();
288 + let engine_id = engine_id.into();
289 + let engine_task = tokio::spawn(async move {
290 + let (mut dealer, mut push) = setup_mock_engine(engine_handshake, engine_id).await;
291 + run(&mut dealer, &mut push).await;
292 + let _ = shutdown_rx.await;
293 + });
294 + (shutdown_tx, engine_task)
295 +}
rust/src/engine-core-client/src/tests/client.rs new
+2559
@@ -0,0 +1,2559 @@
1 +use std::collections::{BTreeMap, BTreeSet};
2 +use std::convert::TryFrom;
3 +use std::io::Cursor;
4 +use std::path::PathBuf;
5 +use std::process::Command;
6 +use std::sync::Once;
7 +use std::time::Duration;
8 +
9 +use futures::StreamExt;
10 +use rmpv::Value;
11 +use thiserror_ext::AsReport as _;
12 +use tokio::sync::{mpsc, oneshot};
13 +use tokio::time::timeout;
14 +use tracing_subscriber::EnvFilter;
15 +use zeromq::prelude::{Socket, SocketRecv, SocketSend};
16 +use zeromq::util::PeerIdentity;
17 +use zeromq::{DealerSocket, PushSocket, SocketOptions, SubSocket, XPubSocket, ZmqMessage};
18 +
19 +use crate::protocol::handshake::{HandshakeInitMessage, ReadyMessage};
20 +use crate::protocol::logprobs::MaybeWireLogprobs;
21 +use crate::protocol::multimodal::{
22 + MmFeatureSpec, MmField, MmFieldElem, MmFlatField, MmKwargValue, MmSlice, PlaceholderRange,
23 + SliceSpec,
24 +};
25 +use crate::protocol::stats::SchedulerStats;
26 +use crate::protocol::tensor::WireTensor;
27 +use crate::protocol::{
28 + EngineCoreFinishReason, EngineCoreOutput, EngineCoreOutputs, EngineCoreRequest,
29 + EngineCoreRequestType, EngineCoreSamplingParams, UtilityOutput, UtilityResultEnvelope,
30 + decode_engine_core_outputs,
31 +};
32 +use crate::test_utils::{
33 + IpcNamespace, setup_bootstrapped_mock_engine, setup_mock_engine_connections,
34 + setup_mock_engine_with_init, spawn_mock_engine_task,
35 +};
36 +use crate::{
37 + CoordinatorMode, ENGINE_CORE_DEAD_SENTINEL, EngineCoreClient, EngineCoreClientConfig, EngineId,
38 + Error, TransportMode,
39 +};
40 +
41 +static TRACING: Once = Once::new();
42 +
43 +fn expect_sample_logprobs(actual: &MaybeWireLogprobs) {
44 + expect_test::expect![[r#"
45 + Logprobs {
46 + positions: [
47 + PositionLogprobs {
48 + entries: [
49 + TokenLogprob {
50 + token_id: 1,
51 + logprob: 1.0,
52 + rank: 1,
53 + },
54 + TokenLogprob {
55 + token_id: 2,
56 + logprob: 2.0,
57 + rank: 1,
58 + },
59 + TokenLogprob {
60 + token_id: 3,
61 + logprob: 3.0,
62 + rank: 2,
63 + },
64 + ],
65 + },
66 + PositionLogprobs {
67 + entries: [
68 + TokenLogprob {
69 + token_id: 4,
70 + logprob: 4.0,
71 + rank: 2,
72 + },
73 + TokenLogprob {
74 + token_id: 5,
75 + logprob: 5.0,
76 + rank: 1,
77 + },
78 + TokenLogprob {
79 + token_id: 6,
80 + logprob: 6.0,
81 + rank: 2,
82 + },
83 + ],
84 + },
85 + ],
86 + }
87 + "#]]
88 + .assert_debug_eq(actual.as_direct().expect("logprobs resolved"));
89 +}
90 +
91 +fn expect_prompt_logprobs(actual: &MaybeWireLogprobs) {
92 + expect_test::expect![[r#"
93 + Logprobs {
94 + positions: [
95 + PositionLogprobs {
96 + entries: [
97 + TokenLogprob {
98 + token_id: 10,
99 + logprob: 10.0,
100 + rank: 3,
101 + },
102 + TokenLogprob {
103 + token_id: 11,
104 + logprob: 11.0,
105 + rank: 1,
106 + },
107 + TokenLogprob {
108 + token_id: 12,
109 + logprob: 12.0,
110 + rank: 2,
111 + },
112 + ],
113 + },
114 + PositionLogprobs {
115 + entries: [
116 + TokenLogprob {
117 + token_id: 13,
118 + logprob: 13.0,
119 + rank: 4,
120 + },
121 + TokenLogprob {
122 + token_id: 14,
123 + logprob: 14.0,
124 + rank: 1,
125 + },
126 + TokenLogprob {
127 + token_id: 15,
128 + logprob: 15.0,
129 + rank: 2,
130 + },
131 + ],
132 + },
133 + ],
134 + }
135 + "#]]
136 + .assert_debug_eq(actual.as_direct().expect("prompt logprobs resolved"));
137 +}
138 +
139 +fn sample_request() -> EngineCoreRequest {
140 + sample_request_with_id("req-1")
141 +}
142 +
143 +fn sample_request_with_id(request_id: &str) -> EngineCoreRequest {
144 + EngineCoreRequest {
145 + request_id: request_id.to_string(),
146 + prompt_token_ids: Some(vec![11, 22]),
147 + sampling_params: Some(EngineCoreSamplingParams {
148 + temperature: 0.8,
149 + top_p: 0.9,
150 + top_k: 8,
151 + max_tokens: 32,
152 + min_tokens: 1,
153 + stop_token_ids: vec![151643],
154 + eos_token_id: Some(151645),
155 + all_stop_token_ids: BTreeSet::from([151643, 151645]),
156 + ..EngineCoreSamplingParams::for_test()
157 + }),
158 + arrival_time: 42.5,
159 + ..EngineCoreRequest::default()
160 + }
161 +}
162 +
163 +fn sample_multimodal_request() -> EngineCoreRequest {
164 + EngineCoreRequest {
165 + request_id: "req-mm".to_string(),
166 + prompt_token_ids: Some(vec![101, 102, 103, 104]),
167 + mm_features: Some(vec![MmFeatureSpec {
168 + data: Some(BTreeMap::from([(
169 + "pixel_values".to_string(),
170 + MmFieldElem {
171 + data: Some(MmKwargValue::Tensor(
172 + WireTensor::from_f32(vec![2, 2], vec![1.0, 2.0, 3.5, 4.25])
173 + .expect("valid tensor shape"),
174 + )),
175 + field: MmField::Flat(MmFlatField {
176 + slices: vec![MmSlice::Slice(SliceSpec {
177 + start: Some(0),
178 + stop: Some(2),
179 + step: None,
180 + })],
181 + dim: 0,
182 + keep_on_cpu: false,
183 + }),
184 + },
185 + )])),
186 + modality: "image".to_string(),
187 + identifier: "mm-cache-key".to_string(),
188 + mm_position: PlaceholderRange {
189 + offset: 1,
190 + length: 2,
191 + is_embed: None,
192 + },
193 + mm_hash: Some("processor-hash".to_string()),
194 + }]),
195 + sampling_params: None,
196 + pooling_params: None,
197 + arrival_time: 43.5,
198 + ..EngineCoreRequest::default()
199 + }
200 +}
201 +
202 +fn ready_message(status: &str) -> ReadyMessage {
203 + ReadyMessage {
204 + status: Some(status.to_string()),
205 + local: Some(true),
206 + headless: Some(true),
207 + parallel_config_hash: None,
208 + }
209 +}
210 +
211 +fn request_output(
212 + request_id: &str,
213 + new_token_ids: Vec<u32>,
214 + finish_reason: Option<EngineCoreFinishReason>,
215 +) -> EngineCoreOutput {
216 + EngineCoreOutput {
217 + request_id: request_id.to_string(),
218 + new_token_ids,
219 + new_logprobs: None,
220 + new_prompt_logprobs_tensors: None,
221 + pooling_output: None,
222 + finish_reason,
223 + stop_reason: None,
224 + events: None,
225 + kv_transfer_params: None,
226 + trace_headers: None,
227 + prefill_stats: None,
228 + routed_experts: None,
229 + num_nans_in_logits: 0,
230 + }
231 +}
232 +
233 +async fn send_outputs(push: &mut PushSocket, outputs: EngineCoreOutputs) {
234 + push.send(ZmqMessage::from(rmp_serde::to_vec_named(&outputs).unwrap()))
235 + .await
236 + .unwrap();
237 +}
238 +
239 +async fn send_output_frames(push: &mut PushSocket, frames: Vec<bytes::Bytes>) {
240 + push.send(ZmqMessage::try_from(frames).unwrap()).await.unwrap();
241 +}
242 +
243 +async fn recv_engine_message(dealer: &mut DealerSocket) -> Vec<bytes::Bytes> {
244 + dealer.recv().await.unwrap().into_vec()
245 +}
246 +
247 +async fn recv_start_dp_wave(sub: &mut SubSocket) -> (u32, u32) {
248 + let frames = sub.recv().await.unwrap().into_vec();
249 + assert_eq!(frames.len(), 2);
250 + assert_eq!(
251 + frames[0].as_ref(),
252 + EngineCoreRequestType::StartDpWave.to_frame().as_ref()
253 + );
254 + rmp_serde::from_slice(&frames[1]).expect("decode START_DP_WAVE payload")
255 +}
256 +
257 +async fn connect_client_with_ipc(
258 + config: EngineCoreClientConfig,
259 + ipc: &IpcNamespace,
260 +) -> EngineCoreClient {
261 + EngineCoreClient::connect(
262 + config.with_local_input_output_addresses(
263 + Some(ipc.input_endpoint()),
264 + Some(ipc.output_endpoint()),
265 + ),
266 + )
267 + .await
268 + .unwrap()
269 +}
270 +
271 +fn handshake_test_config(
272 + handshake_address: String,
273 + engine_count: usize,
274 + model_name: &str,
275 + ready_timeout: Duration,
276 + client_index: u32,
277 + coordinator_mode: Option<CoordinatorMode>,
278 +) -> EngineCoreClientConfig {
279 + EngineCoreClientConfig {
280 + transport_mode: TransportMode::HandshakeOwner {
281 + handshake_address,
282 + advertised_host: "127.0.0.1".to_string(),
283 + engine_count,
284 + ready_timeout,
285 + local_input_address: None,
286 + local_output_address: None,
287 + },
288 + coordinator_mode,
289 + model_name: model_name.to_string(),
290 + client_index,
291 + }
292 +}
293 +
294 +fn bootstrapped_test_config(
295 + input_address: String,
296 + output_address: String,
297 + engine_count: usize,
298 + ready_timeout: Duration,
299 + client_index: u32,
300 + coordinator_mode: Option<CoordinatorMode>,
301 +) -> EngineCoreClientConfig {
302 + EngineCoreClientConfig {
303 + transport_mode: TransportMode::Bootstrapped {
304 + input_address,
305 + output_address,
306 + engine_count,
307 + ready_timeout,
308 + },
309 + coordinator_mode,
310 + model_name: "test-model".to_string(),
311 + client_index,
312 + }
313 +}
314 +
315 +async fn recv_xpub_message(xpub: &mut XPubSocket) -> Vec<bytes::Bytes> {
316 + xpub.recv().await.unwrap().into_vec()
317 +}
318 +
319 +async fn recv_xpub_subscription(xpub: &mut XPubSocket) {
320 + let frames = recv_xpub_message(xpub).await;
321 + assert_eq!(frames.len(), 1);
322 + assert_eq!(frames[0].as_ref(), b"\x01");
323 +}
324 +
325 +async fn recv_external_coordinator_wakeup(xpub: &mut XPubSocket) -> (u32, u32) {
326 + let frames = recv_xpub_message(xpub).await;
327 + assert_eq!(frames.len(), 1);
328 + rmp_serde::from_slice(&frames[0]).expect("decode external coordinator wakeup")
329 +}
330 +
331 +async fn send_external_coordinator_publish<T: serde::Serialize>(
332 + xpub: &mut XPubSocket,
333 + payload: &T,
334 +) {
335 + xpub.send(ZmqMessage::from(rmp_serde::to_vec_named(payload).unwrap()))
336 + .await
337 + .unwrap();
338 +}
339 +
340 +fn spawn_mock_engine_task_with_init<F>(
341 + engine_handshake: String,
342 + engine_id: impl Into<EngineId>,
343 + run: F,
344 +) -> (
345 + oneshot::Receiver<HandshakeInitMessage>,
346 + oneshot::Sender<()>,
347 + tokio::task::JoinHandle<()>,
348 +)
349 +where
350 + F: for<'a> FnOnce(
351 + &'a mut DealerSocket,
352 + &'a mut PushSocket,
353 + )
354 + -> std::pin::Pin<Box<dyn std::future::Future<Output = ()> + Send + 'a>>
355 + + Send
356 + + 'static,
357 +{
358 + let (shutdown_tx, shutdown_rx) = oneshot::channel();
359 + let (init_tx, init_rx) = oneshot::channel();
360 + let engine_id = engine_id.into();
361 + let engine_task = tokio::spawn(async move {
362 + let (init, mut dealer, mut push) =
363 + setup_mock_engine_with_init(engine_handshake, engine_id).await;
364 + let _ = init_tx.send(init);
365 + run(&mut dealer, &mut push).await;
366 + let _ = shutdown_rx.await;
367 + });
368 + (init_rx, shutdown_tx, engine_task)
369 +}
370 +
371 +fn init_tracing() {
372 + TRACING.call_once(|| {
373 + let filter = EnvFilter::try_from_default_env()
374 + .unwrap_or_else(|_| EnvFilter::new("vllm_engine_core_client=debug"));
375 + let _ = tracing_subscriber::fmt().with_test_writer().with_env_filter(filter).try_init();
376 + });
377 +}
378 +
379 +fn is_dispatcher_closed(error: &Error) -> bool {
380 + match error {
381 + Error::DispatcherClosed { .. } => true,
382 + Error::Shared(error) => is_dispatcher_closed(error),
383 + _ => false,
384 + }
385 +}
386 +
387 +fn is_engine_core_dead(error: &Error) -> bool {
388 + match error {
389 + Error::EngineCoreDead => true,
390 + Error::Shared(error) => is_engine_core_dead(error),
391 + _ => false,
392 + }
393 +}
394 +
395 +fn is_decode_error(error: &Error) -> bool {
396 + match error {
397 + Error::Decode { .. } | Error::ExtValueDecode { .. } => true,
398 + Error::Shared(error) => is_decode_error(error),
399 + _ => false,
400 + }
401 +}
402 +
403 +fn is_unexpected_dispatcher_output(error: &Error) -> bool {
404 + match error {
405 + Error::UnexpectedDispatcherOutput { .. } => true,
406 + Error::Shared(error) => is_unexpected_dispatcher_output(error),
407 + _ => false,
408 + }
409 +}
410 +
411 +fn decode_value(bytes: &[u8]) -> Value {
412 + rmpv::decode::read_value(&mut Cursor::new(bytes)).unwrap()
413 +}
414 +
415 +fn encode_value(value: &Value) -> Vec<u8> {
416 + let mut out = Vec::new();
417 + rmpv::encode::write_value(&mut out, value).unwrap();
418 + out
419 +}
420 +
421 +fn ndarray_value(dtype: &str, shape: &[usize], data: Value) -> Value {
422 + Value::Array(vec![
423 + Value::from(dtype),
424 + Value::Array(shape.iter().copied().map(Value::from).collect()),
425 + data,
426 + ])
427 +}
428 +
429 +fn multipart_logprob_output_frames(request_id: &str) -> Vec<bytes::Bytes> {
430 + let main = Value::Array(vec![
431 + Value::from(0),
432 + Value::Array(vec![Value::Array(vec![
433 + Value::from(request_id),
434 + Value::Array(vec![Value::from(7), Value::from(8)]),
435 + Value::Array(vec![
436 + ndarray_value("<i4", &[2, 3], Value::from(1)),
437 + ndarray_value("<f4", &[2, 3], Value::from(2)),
438 + ndarray_value("<i4", &[2], Value::from(3)),
439 + Value::Nil,
440 + ]),
441 + Value::Nil,
442 + Value::Nil,
443 + Value::from(EngineCoreFinishReason::Length as u8),
444 + ])]),
445 + Value::Nil,
446 + Value::from(0.0),
447 + Value::Nil,
448 + Value::Array(vec![Value::from(request_id)]),
449 + ]);
450 +
451 + vec![
452 + bytes::Bytes::from(encode_value(&main)),
453 + bytes::Bytes::from_static(&[
454 + 1, 0, 0, 0, 2, 0, 0, 0, 3, 0, 0, 0, 4, 0, 0, 0, 5, 0, 0, 0, 6, 0, 0, 0,
455 + ]),
456 + bytes::Bytes::from_static(&[
457 + 0, 0, 128, 63, 0, 0, 0, 64, 0, 0, 64, 64, 0, 0, 128, 64, 0, 0, 160, 64, 0, 0, 192, 64,
458 + ]),
459 + bytes::Bytes::from_static(&[1, 0, 0, 0, 2, 0, 0, 0]),
460 + ]
461 +}
462 +
463 +fn utility_result_value<T>(value: T) -> UtilityResultEnvelope
464 +where
465 + T: serde::Serialize,
466 +{
467 + UtilityResultEnvelope::without_type_info(rmpv::ext::to_value(value).unwrap())
468 +}
469 +
470 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
471 +async fn coordinator_handshake_includes_engine_control_addresses() {
472 + init_tracing();
473 + let ipc = IpcNamespace::new().unwrap();
474 + let handshake_address = ipc.handshake_endpoint();
475 + let engine_id = [0x00, 0x00];
476 +
477 + let (init_tx, init_rx) = oneshot::channel();
478 + let (shutdown_tx, shutdown_rx) = oneshot::channel();
479 + let engine_task = tokio::spawn(async move {
480 + let connections = setup_mock_engine_connections(handshake_address, &engine_id).await;
481 + let _ = init_tx.send(connections.init.clone());
482 + let _ = shutdown_rx.await;
483 + });
484 +
485 + let client = connect_client_with_ipc(
486 + handshake_test_config(
487 + ipc.handshake_endpoint(),
488 + 1,
489 + "test-model",
490 + Duration::from_secs(2),
491 + 0,
492 + Some(CoordinatorMode::InProc),
493 + ),
494 + &ipc,
495 + )
496 + .await;
497 +
498 + let init = init_rx.await.unwrap();
499 + assert!(init.addresses.coordinator_input.is_some());
500 + assert!(init.addresses.coordinator_output.is_some());
501 + assert!(init.addresses.frontend_stats_publish_address.is_none());
502 +
503 + let _ = shutdown_tx.send(());
504 + engine_task.await.unwrap();
505 + client.shutdown().await.unwrap();
506 +}
507 +
508 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
509 +async fn coordinator_wave_control_tracks_pause_running_and_rebroadcasts() {
510 + init_tracing();
511 + let ipc = IpcNamespace::new().unwrap();
512 + let handshake_address = ipc.handshake_endpoint();
513 +
514 + let (shutdown0_tx, shutdown0_rx) = oneshot::channel();
515 + let engine0_task = tokio::spawn({
516 + let handshake_address = handshake_address.clone();
517 + async move {
518 + let mut engine = setup_mock_engine_connections(handshake_address, &[0x00, 0x00]).await;
519 + let mut coordinator =
520 + engine.coordinator.take().expect("coordinator sockets should be present");
521 +
522 + let (wave, exclude_engine) = recv_start_dp_wave(&mut coordinator.input_sub).await;
523 + assert_eq!((wave, exclude_engine), (0, 0));
524 +
525 + let add = recv_engine_message(&mut engine.dealer).await;
526 + assert_eq!(add[0].as_ref(), &[0x00]);
527 + let request: EngineCoreRequest = rmp_serde::from_slice(&add[1]).unwrap();
528 + assert_eq!(request.request_id, "req-1");
529 + assert_eq!(request.current_wave, 0);
530 +
531 + assert!(
532 + timeout(
533 + Duration::from_millis(200),
534 + recv_start_dp_wave(&mut coordinator.input_sub)
535 + )
536 + .await
537 + .is_err()
538 + );
539 +
540 + send_outputs(
541 + &mut engine.push,
542 + EngineCoreOutputs {
543 + engine_index: 0,
544 + outputs: vec![request_output(
545 + "req-1",
546 + vec![],
547 + Some(EngineCoreFinishReason::Length),
548 + )],
549 + finished_requests: Some(BTreeSet::from(["req-1".to_string()])),
550 + ..Default::default()
551 + },
552 + )
553 + .await;
554 +
555 + send_outputs(
556 + &mut coordinator.output_push,
557 + EngineCoreOutputs {
558 + engine_index: 0,
559 + wave_complete: Some(0),
560 + ..Default::default()
561 + },
562 + )
563 + .await;
564 +
565 + let (wave, exclude_engine) = recv_start_dp_wave(&mut coordinator.input_sub).await;
566 + assert_eq!((wave, exclude_engine), (1, 0));
567 +
568 + let add = recv_engine_message(&mut engine.dealer).await;
569 + assert_eq!(add[0].as_ref(), &[0x00]);
570 + let request: EngineCoreRequest = rmp_serde::from_slice(&add[1]).unwrap();
571 + assert_eq!(request.request_id, "req-3");
572 + assert_eq!(request.current_wave, 1);
573 +
574 + send_outputs(
575 + &mut engine.push,
576 + EngineCoreOutputs {
577 + engine_index: 0,
578 + outputs: vec![request_output(
579 + "req-3",
580 + vec![],
581 + Some(EngineCoreFinishReason::Length),
582 + )],
583 + finished_requests: Some(BTreeSet::from(["req-3".to_string()])),
584 + ..Default::default()
585 + },
586 + )
587 + .await;
588 +
589 + let _ = shutdown0_rx.await;
590 + }
591 + });
592 +
593 + let (shutdown1_tx, shutdown1_rx) = oneshot::channel();
594 + let engine1_task = tokio::spawn({
595 + let handshake_address = handshake_address.clone();
596 + async move {
597 + let mut engine = setup_mock_engine_connections(handshake_address, &[0x01, 0x00]).await;
598 + let mut coordinator =
599 + engine.coordinator.take().expect("coordinator sockets should be present");
600 +
601 + let (wave, exclude_engine) = recv_start_dp_wave(&mut coordinator.input_sub).await;
602 + assert_eq!((wave, exclude_engine), (0, 0));
603 +
604 + let add = recv_engine_message(&mut engine.dealer).await;
605 + assert_eq!(add[0].as_ref(), &[0x00]);
606 + let request: EngineCoreRequest = rmp_serde::from_slice(&add[1]).unwrap();
607 + assert_eq!(request.request_id, "req-2");
608 + assert_eq!(request.current_wave, 0);
609 +
610 + assert!(
611 + timeout(
612 + Duration::from_millis(200),
613 + recv_start_dp_wave(&mut coordinator.input_sub)
614 + )
615 + .await
616 + .is_err()
617 + );
618 +
619 + send_outputs(
620 + &mut engine.push,
621 + EngineCoreOutputs {
622 + engine_index: 1,
623 + outputs: vec![request_output(
624 + "req-2",
625 + vec![],
626 + Some(EngineCoreFinishReason::Length),
627 + )],
628 + finished_requests: Some(BTreeSet::from(["req-2".to_string()])),
629 + ..Default::default()
630 + },
631 + )
632 + .await;
633 +
634 + let (wave, exclude_engine) = recv_start_dp_wave(&mut coordinator.input_sub).await;
635 + assert_eq!((wave, exclude_engine), (1, 0));
636 +
637 + assert!(
638 + timeout(
639 + Duration::from_millis(200),
640 + recv_engine_message(&mut engine.dealer)
641 + )
642 + .await
643 + .is_err()
644 + );
645 +
646 + let _ = shutdown1_rx.await;
647 + }
648 + });
649 +
650 + let client = connect_client_with_ipc(
651 + handshake_test_config(
652 + handshake_address,
653 + 2,
654 + "test-model",
655 + Duration::from_secs(2),
656 + 0,
657 + Some(CoordinatorMode::InProc),
658 + ),
659 + &ipc,
660 + )
661 + .await;
662 +
663 + let mut stream_1 = client.call(sample_request_with_id("req-1")).await.unwrap();
664 + let mut stream_2 = client.call(sample_request_with_id("req-2")).await.unwrap();
665 +
666 + let final_1 = timeout(Duration::from_secs(1), stream_1.next())
667 + .await
668 + .unwrap()
669 + .unwrap()
670 + .unwrap();
671 + assert_eq!(final_1.request_id, "req-1");
672 + assert_eq!(final_1.finish_reason, Some(EngineCoreFinishReason::Length));
673 + assert!(timeout(Duration::from_secs(1), stream_1.next()).await.unwrap().is_none());
674 +
675 + let final_2 = timeout(Duration::from_secs(1), stream_2.next())
676 + .await
677 + .unwrap()
678 + .unwrap()
679 + .unwrap();
680 + assert_eq!(final_2.request_id, "req-2");
681 + assert_eq!(final_2.finish_reason, Some(EngineCoreFinishReason::Length));
682 + assert!(timeout(Duration::from_secs(1), stream_2.next()).await.unwrap().is_none());
683 +
684 + tokio::time::sleep(Duration::from_millis(100)).await;
685 +
686 + let mut stream_3 = client.call(sample_request_with_id("req-3")).await.unwrap();
687 + let final_3 = timeout(Duration::from_secs(1), stream_3.next())
688 + .await
689 + .unwrap()
690 + .unwrap()
691 + .unwrap();
692 + assert_eq!(final_3.request_id, "req-3");
693 + assert_eq!(final_3.finish_reason, Some(EngineCoreFinishReason::Length));
694 + assert!(timeout(Duration::from_secs(1), stream_3.next()).await.unwrap().is_none());
695 +
696 + tokio::time::sleep(Duration::from_millis(100)).await;
697 +
698 + let _ = shutdown0_tx.send(());
699 + let _ = shutdown1_tx.send(());
700 + engine0_task.await.unwrap();
701 + engine1_task.await.unwrap();
702 + client.shutdown().await.unwrap();
703 +}
704 +
705 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
706 +async fn coordinator_rebroadcasts_engine_start_wave_control() {
707 + init_tracing();
708 + let ipc = IpcNamespace::new().unwrap();
709 + let handshake_address = ipc.handshake_endpoint();
710 +
711 + let (shutdown0_tx, shutdown0_rx) = oneshot::channel();
712 + let engine0_task = tokio::spawn({
713 + let handshake_address = handshake_address.clone();
714 + async move {
715 + let mut engine = setup_mock_engine_connections(handshake_address, &[0x00, 0x00]).await;
716 + let mut coordinator =
717 + engine.coordinator.take().expect("coordinator sockets should be present");
718 +
719 + let (wave, exclude_engine) = recv_start_dp_wave(&mut coordinator.input_sub).await;
720 + assert_eq!((wave, exclude_engine), (4, 1));
721 +
722 + let _ = shutdown0_rx.await;
723 + }
724 + });
725 +
726 + let (shutdown1_tx, shutdown1_rx) = oneshot::channel();
727 + let engine1_task = tokio::spawn({
728 + let handshake_address = handshake_address.clone();
729 + async move {
730 + let mut engine = setup_mock_engine_connections(handshake_address, &[0x01, 0x00]).await;
731 + let mut coordinator =
732 + engine.coordinator.take().expect("coordinator sockets should be present");
733 +
734 + send_outputs(
735 + &mut coordinator.output_push,
736 + EngineCoreOutputs {
737 + engine_index: 1,
738 + start_wave: Some(4),
739 + ..Default::default()
740 + },
741 + )
742 + .await;
743 +
744 + let (wave, exclude_engine) = recv_start_dp_wave(&mut coordinator.input_sub).await;
745 + assert_eq!((wave, exclude_engine), (4, 1));
746 +
747 + let _ = shutdown1_rx.await;
748 + }
749 + });
750 +
751 + let client = connect_client_with_ipc(
752 + handshake_test_config(
753 + handshake_address,
754 + 2,
755 + "test-model",
756 + Duration::from_secs(2),
757 + 0,
758 + Some(CoordinatorMode::InProc),
759 + ),
760 + &ipc,
761 + )
762 + .await;
763 +
764 + tokio::time::sleep(Duration::from_millis(200)).await;
765 +
766 + let _ = shutdown0_tx.send(());
767 + let _ = shutdown1_tx.send(());
768 + engine0_task.await.unwrap();
769 + engine1_task.await.unwrap();
770 + client.shutdown().await.unwrap();
771 +}
772 +
773 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
774 +async fn coordinator_accepts_stats_only_outputs() {
775 + init_tracing();
776 + let ipc = IpcNamespace::new().unwrap();
777 + let handshake_address = ipc.handshake_endpoint();
778 +
779 + let (shutdown_tx, shutdown_rx) = oneshot::channel();
780 + let engine_task = tokio::spawn(async move {
781 + let mut engine = setup_mock_engine_connections(handshake_address, &[0x00, 0x00]).await;
782 + let mut coordinator =
783 + engine.coordinator.take().expect("coordinator sockets should be present");
784 +
785 + let (wave, exclude_engine) = recv_start_dp_wave(&mut coordinator.input_sub).await;
786 + assert_eq!((wave, exclude_engine), (0, 0));
787 +
788 + send_outputs(
789 + &mut coordinator.output_push,
790 + EngineCoreOutputs {
791 + engine_index: 0,
792 + scheduler_stats: Some(Box::new(SchedulerStats {
793 + num_running_reqs: 1,
794 + current_wave: 0,
795 + ..Default::default()
796 + })),
797 + ..Default::default()
798 + },
799 + )
800 + .await;
801 +
802 + let add = recv_engine_message(&mut engine.dealer).await;
803 + assert_eq!(add[0].as_ref(), &[0x00]);
804 + let request: EngineCoreRequest = rmp_serde::from_slice(&add[1]).unwrap();
805 + assert_eq!(request.request_id, "req-stats");
806 +
807 + send_outputs(
808 + &mut engine.push,
809 + EngineCoreOutputs {
810 + engine_index: 0,
811 + outputs: vec![request_output(
812 + "req-stats",
813 + vec![],
814 + Some(EngineCoreFinishReason::Length),
815 + )],
816 + finished_requests: Some(BTreeSet::from(["req-stats".to_string()])),
817 + ..Default::default()
818 + },
819 + )
820 + .await;
821 +
822 + let _ = shutdown_rx.await;
823 + });
824 +
825 + let client = connect_client_with_ipc(
826 + handshake_test_config(
827 + ipc.handshake_endpoint(),
828 + 1,
829 + "test-model",
830 + Duration::from_secs(2),
831 + 0,
832 + Some(CoordinatorMode::InProc),
833 + ),
834 + &ipc,
835 + )
836 + .await;
837 +
838 + let mut stream = client.call(sample_request_with_id("req-stats")).await.unwrap();
839 + let final_output =
840 + timeout(Duration::from_secs(1), stream.next()).await.unwrap().unwrap().unwrap();
841 + assert_eq!(final_output.request_id, "req-stats");
842 + assert_eq!(
843 + final_output.finish_reason,
844 + Some(EngineCoreFinishReason::Length)
845 + );
846 + assert!(client.is_healthy());
847 +
848 + let _ = shutdown_tx.send(());
849 + engine_task.await.unwrap();
850 + client.shutdown().await.unwrap();
851 +}
852 +
853 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
854 +async fn client_fail_closes_when_main_output_path_receives_dp_control() {
855 + init_tracing();
856 + let ipc = IpcNamespace::new().unwrap();
857 + let handshake_address = ipc.handshake_endpoint();
858 + let engine_id = b"engine-0".to_vec();
859 +
860 + let (shutdown_tx, engine_task) = spawn_mock_engine_task(
861 + handshake_address.clone(),
862 + engine_id.clone(),
863 + |dealer, push| {
864 + Box::pin(async move {
865 + let add_1 = recv_engine_message(dealer).await;
866 + assert_eq!(add_1[0].as_ref(), &[0x00]);
867 + let request_1: EngineCoreRequest = rmp_serde::from_slice(&add_1[1]).unwrap();
868 + assert_eq!(request_1.client_index, 7);
869 + assert_eq!(request_1.request_id, "req-1");
870 +
871 + let add_2 = recv_engine_message(dealer).await;
872 + assert_eq!(add_2[0].as_ref(), &[0x00]);
873 + let request_2: EngineCoreRequest = rmp_serde::from_slice(&add_2[1]).unwrap();
874 + assert_eq!(request_2.client_index, 7);
875 + assert_eq!(request_2.request_id, "req-2");
876 +
877 + send_outputs(
878 + push,
879 + EngineCoreOutputs {
880 + utility_output: Some(UtilityOutput {
881 + call_id: 1,
882 + failure_message: None,
883 + result: None,
884 + }),
885 + ..Default::default()
886 + },
887 + )
888 + .await;
889 + send_outputs(
890 + push,
891 + EngineCoreOutputs {
892 + start_wave: Some(3),
893 + ..Default::default()
894 + },
895 + )
896 + .await;
897 + send_outputs(
898 + push,
899 + EngineCoreOutputs {
900 + outputs: vec![request_output("req-1", vec![999], None)],
901 + ..Default::default()
902 + },
903 + )
904 + .await;
905 +
906 + tokio::time::sleep(Duration::from_millis(50)).await;
907 + })
908 + },
909 + );
910 +
911 + let client = connect_client_with_ipc(
912 + handshake_test_config(
913 + handshake_address,
914 + 1,
915 + "test-model",
916 + Duration::from_secs(2),
917 + 7,
918 + None,
919 + ),
920 + &ipc,
921 + )
922 + .await;
923 + assert_eq!(client.engine_identities()[0], b"engine-0");
924 + assert!(client.ready_responses()[0].max_model_len > 0);
925 +
926 + let mut stream_1 = client.call(sample_request_with_id("req-1")).await.unwrap();
927 + let mut stream_2 = client.call(sample_request_with_id("req-2")).await.unwrap();
928 +
929 + let error_2 = timeout(Duration::from_secs(1), stream_2.next())
930 + .await
931 + .unwrap()
932 + .unwrap()
933 + .unwrap_err();
934 + assert!(is_unexpected_dispatcher_output(&error_2));
935 +
936 + let error_1 = timeout(Duration::from_secs(1), stream_1.next())
937 + .await
938 + .unwrap()
939 + .unwrap()
940 + .unwrap_err();
941 + assert!(is_unexpected_dispatcher_output(&error_1));
942 +
943 + assert!(matches!(
944 + client.health_error().as_deref(),
945 + Some(error) if is_unexpected_dispatcher_output(error)
946 + ));
947 +
948 + let _ = shutdown_tx.send(());
949 + engine_task.await.unwrap();
950 + client.shutdown().await.unwrap();
951 +}
952 +
953 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
954 +async fn client_fail_closes_when_main_output_path_receives_mixed_shape_output() {
955 + init_tracing();
956 + let ipc = IpcNamespace::new().unwrap();
957 + let handshake_address = ipc.handshake_endpoint();
958 + let engine_id = b"engine-0".to_vec();
959 +
960 + let (shutdown_tx, engine_task) = spawn_mock_engine_task(
961 + handshake_address.clone(),
962 + engine_id.clone(),
963 + |dealer, push| {
964 + Box::pin(async move {
965 + let add_1 = recv_engine_message(dealer).await;
966 + assert_eq!(add_1[0].as_ref(), &[0x00]);
967 + let request_1: EngineCoreRequest = rmp_serde::from_slice(&add_1[1]).unwrap();
968 + assert_eq!(request_1.client_index, 7);
969 + assert_eq!(request_1.request_id, "req-1");
970 +
971 + let add_2 = recv_engine_message(dealer).await;
972 + assert_eq!(add_2[0].as_ref(), &[0x00]);
973 + let request_2: EngineCoreRequest = rmp_serde::from_slice(&add_2[1]).unwrap();
974 + assert_eq!(request_2.client_index, 7);
975 + assert_eq!(request_2.request_id, "req-2");
976 +
977 + send_outputs(
978 + push,
979 + EngineCoreOutputs {
980 + utility_output: Some(UtilityOutput {
981 + call_id: 1,
982 + failure_message: None,
983 + result: None,
984 + }),
985 + outputs: vec![request_output("req-1", vec![999], None)],
986 + ..Default::default()
987 + },
988 + )
989 + .await;
990 +
991 + tokio::time::sleep(Duration::from_millis(50)).await;
992 + })
993 + },
994 + );
995 +
996 + let client = connect_client_with_ipc(
997 + handshake_test_config(
998 + handshake_address,
999 + 1,
1000 + "test-model",
1001 + Duration::from_secs(2),
1002 + 7,
1003 + None,
1004 + ),
1005 + &ipc,
1006 + )
1007 + .await;
1008 + assert_eq!(client.engine_identities()[0], b"engine-0");
1009 + assert!(client.ready_responses()[0].max_model_len > 0);
1010 +
1011 + let mut stream_1 = client.call(sample_request_with_id("req-1")).await.unwrap();
1012 + let mut stream_2 = client.call(sample_request_with_id("req-2")).await.unwrap();
1013 +
1014 + let error_2 = timeout(Duration::from_secs(1), stream_2.next())
1015 + .await
1016 + .unwrap()
1017 + .unwrap()
1018 + .unwrap_err();
1019 + assert!(is_unexpected_dispatcher_output(&error_2));
1020 +
1021 + let error_1 = timeout(Duration::from_secs(1), stream_1.next())
1022 + .await
1023 + .unwrap()
1024 + .unwrap()
1025 + .unwrap_err();
1026 + assert!(is_unexpected_dispatcher_output(&error_1));
1027 +
1028 + assert!(matches!(
1029 + client.health_error().as_deref(),
1030 + Some(error) if is_unexpected_dispatcher_output(error)
1031 + ));
1032 +
1033 + let _ = shutdown_tx.send(());
1034 + engine_task.await.unwrap();
1035 + client.shutdown().await.unwrap();
1036 +}
1037 +
1038 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1039 +async fn duplicate_request_ids_are_rejected_without_sending_a_second_add() {
1040 + init_tracing();
1041 + let ipc = IpcNamespace::new().unwrap();
1042 + let handshake_address = ipc.handshake_endpoint();
1043 + let engine_id = b"engine-dup".to_vec();
1044 +
1045 + let (shutdown_tx, engine_task) = spawn_mock_engine_task(
1046 + handshake_address.clone(),
1047 + engine_id.clone(),
1048 + |dealer, push| {
1049 + Box::pin(async move {
1050 + let add_1 = recv_engine_message(dealer).await;
1051 + assert_eq!(add_1[0].as_ref(), &[0x00]);
1052 + let request_1: EngineCoreRequest = rmp_serde::from_slice(&add_1[1]).unwrap();
1053 + assert_eq!(request_1.request_id, "req-1");
1054 +
1055 + assert!(timeout(Duration::from_millis(200), dealer.recv()).await.is_err());
1056 +
1057 + send_outputs(
1058 + push,
1059 + EngineCoreOutputs {
1060 + outputs: vec![request_output(
1061 + "req-1",
1062 + vec![],
1063 + Some(EngineCoreFinishReason::Length),
1064 + )],
1065 + finished_requests: Some(BTreeSet::from(["req-1".to_string()])),
1066 + ..Default::default()
1067 + },
1068 + )
1069 + .await;
1070 + })
1071 + },
1072 + );
1073 +
1074 + let client = connect_client_with_ipc(
1075 + handshake_test_config(
1076 + handshake_address,
1077 + 1,
1078 + "test-model",
1079 + Duration::from_secs(2),
1080 + 0,
1081 + None,
1082 + ),
1083 + &ipc,
1084 + )
1085 + .await;
1086 +
1087 + let mut stream = client.call(sample_request()).await.unwrap();
1088 + let error = match client.call(sample_request()).await {
1089 + Ok(_) => panic!("expected duplicate request error"),
1090 + Err(error) => error,
1091 + };
1092 + assert!(matches!(
1093 + error,
1094 + Error::DuplicateRequestId { request_id } if request_id == "req-1"
1095 + ));
1096 +
1097 + let final_output =
1098 + timeout(Duration::from_secs(1), stream.next()).await.unwrap().unwrap().unwrap();
1099 + assert_eq!(
1100 + final_output.finish_reason,
1101 + Some(EngineCoreFinishReason::Length)
1102 + );
1103 + assert!(timeout(Duration::from_secs(1), stream.next()).await.unwrap().is_none());
1104 + let _ = shutdown_tx.send(());
1105 + engine_task.await.unwrap();
1106 + client.shutdown().await.unwrap();
1107 +}
1108 +
1109 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1110 +async fn finished_requests_without_final_output_is_treated_as_unexpected_close() {
1111 + init_tracing();
1112 + let ipc = IpcNamespace::new().unwrap();
1113 + let handshake_address = ipc.handshake_endpoint();
1114 + let engine_id = b"engine-finished-only".to_vec();
1115 +
1116 + let (shutdown_tx, engine_task) = spawn_mock_engine_task(
1117 + handshake_address.clone(),
1118 + engine_id.clone(),
1119 + |dealer, push| {
1120 + Box::pin(async move {
1121 + let add = recv_engine_message(dealer).await;
1122 + assert_eq!(add[0].as_ref(), &[0x00]);
1123 +
1124 + send_outputs(
1125 + push,
1126 + EngineCoreOutputs {
1127 + finished_requests: Some(BTreeSet::from(["req-1".to_string()])),
1128 + ..Default::default()
1129 + },
1130 + )
1131 + .await;
1132 +
1133 + assert!(timeout(Duration::from_millis(200), dealer.recv()).await.is_err());
1134 + let _ = push;
1135 + })
1136 + },
1137 + );
1138 +
1139 + let client = connect_client_with_ipc(
1140 + handshake_test_config(
1141 + handshake_address,
1142 + 1,
1143 + "test-model",
1144 + Duration::from_secs(2),
1145 + 0,
1146 + None,
1147 + ),
1148 + &ipc,
1149 + )
1150 + .await;
1151 +
1152 + let mut stream = client.call(sample_request()).await.unwrap();
1153 + let error = timeout(Duration::from_secs(1), stream.next())
1154 + .await
1155 + .unwrap()
1156 + .unwrap()
1157 + .unwrap_err();
1158 + assert!(matches!(
1159 + error,
1160 + Error::RequestStreamClosed { request_id } if request_id == "req-1"
1161 + ));
1162 + assert!(timeout(Duration::from_secs(1), stream.next()).await.unwrap().is_none());
1163 +
1164 + let _ = shutdown_tx.send(());
1165 + engine_task.await.unwrap();
1166 + client.shutdown().await.unwrap();
1167 +}
1168 +
1169 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1170 +async fn dropping_a_live_stream_triggers_abort() {
1171 + init_tracing();
1172 + let ipc = IpcNamespace::new().unwrap();
1173 + let handshake_address = ipc.handshake_endpoint();
1174 + let engine_id = b"engine-drop".to_vec();
1175 +
1176 + let (shutdown_tx, engine_task) = spawn_mock_engine_task(
1177 + handshake_address.clone(),
1178 + engine_id.clone(),
1179 + |dealer, push| {
1180 + Box::pin(async move {
1181 + let add = recv_engine_message(dealer).await;
1182 + assert_eq!(add[0].as_ref(), &[0x00]);
1183 + send_outputs(
1184 + push,
1185 + EngineCoreOutputs {
1186 + outputs: vec![request_output("req-1", vec![99], None)],
1187 + ..Default::default()
1188 + },
1189 + )
1190 + .await;
1191 +
1192 + let abort =
1193 + timeout(Duration::from_secs(1), recv_engine_message(dealer)).await.unwrap();
1194 + assert_eq!(abort[0].as_ref(), &[0x01]);
1195 + let aborted_ids: Vec<String> = rmp_serde::from_slice(&abort[1]).unwrap();
1196 + assert_eq!(aborted_ids, vec!["req-1".to_string()]);
1197 + })
1198 + },
1199 + );
1200 +
1201 + let client = connect_client_with_ipc(
1202 + handshake_test_config(
1203 + handshake_address,
1204 + 1,
1205 + "test-model",
1206 + Duration::from_secs(2),
1207 + 0,
1208 + None,
1209 + ),
1210 + &ipc,
1211 + )
1212 + .await;
1213 +
1214 + let mut stream = client.call(sample_request()).await.unwrap();
1215 + let first = timeout(Duration::from_secs(1), stream.next()).await.unwrap().unwrap().unwrap();
1216 + assert_eq!(first.new_token_ids, vec![99]);
1217 + drop(stream);
1218 +
1219 + let _ = shutdown_tx.send(());
1220 + engine_task.await.unwrap();
1221 + client.shutdown().await.unwrap();
1222 +}
1223 +
1224 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1225 +async fn dispatcher_failure_propagates_to_streams_and_future_calls() {
1226 + init_tracing();
1227 + let ipc = IpcNamespace::new().unwrap();
1228 + let handshake_address = ipc.handshake_endpoint();
1229 + let engine_id = b"engine-fail".to_vec();
1230 +
1231 + let (shutdown_tx, engine_task) = spawn_mock_engine_task(
1232 + handshake_address.clone(),
1233 + engine_id.clone(),
1234 + |dealer, push| {
1235 + Box::pin(async move {
1236 + let _ = recv_engine_message(dealer).await;
1237 + let _ = recv_engine_message(dealer).await;
1238 +
1239 + push.send(ZmqMessage::from(vec![0xc1])).await.unwrap();
1240 + })
1241 + },
1242 + );
1243 +
1244 + let client = connect_client_with_ipc(
1245 + handshake_test_config(
1246 + handshake_address,
1247 + 1,
1248 + "test-model",
1249 + Duration::from_secs(2),
1250 + 0,
1251 + None,
1252 + ),
1253 + &ipc,
1254 + )
1255 + .await;
1256 +
1257 + let mut stream_1 = client.call(sample_request_with_id("req-1")).await.unwrap();
1258 + let mut stream_2 = client.call(sample_request_with_id("req-2")).await.unwrap();
1259 +
1260 + let error_1 = timeout(Duration::from_secs(1), stream_1.next())
1261 + .await
1262 + .unwrap()
1263 + .unwrap()
1264 + .unwrap_err();
1265 + let error_2 = timeout(Duration::from_secs(1), stream_2.next())
1266 + .await
1267 + .unwrap()
1268 + .unwrap()
1269 + .unwrap_err();
1270 + assert!(is_decode_error(&error_1));
1271 + assert!(is_decode_error(&error_2));
1272 + assert!(is_decode_error(
1273 + client.health_error().as_deref().expect("health error recorded")
1274 + ));
1275 +
1276 + let abort_error = client.abort(&["req-1".to_string()]).await.unwrap_err();
1277 + assert!(is_decode_error(&abort_error));
1278 +
1279 + let add_error = match client.call(sample_request_with_id("req-3")).await {
1280 + Ok(_) => panic!("expected dispatcher closed error"),
1281 + Err(error) => error,
1282 + };
1283 + assert!(is_decode_error(&add_error));
1284 +
1285 + let _ = shutdown_tx.send(());
1286 + engine_task.await.unwrap();
1287 + client.shutdown().await.unwrap();
1288 +}
1289 +
1290 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1291 +async fn is_sleeping_wrapper_sends_typed_request_and_returns_typed_response() {
1292 + init_tracing();
1293 + let ipc = IpcNamespace::new().unwrap();
1294 + let handshake_address = ipc.handshake_endpoint();
1295 + let engine_id = b"engine-utility-success".to_vec();
1296 +
1297 + let (shutdown_tx, engine_task) = spawn_mock_engine_task(
1298 + handshake_address.clone(),
1299 + engine_id.clone(),
1300 + |dealer, push| {
1301 + Box::pin(async move {
1302 + let utility = recv_engine_message(dealer).await;
1303 + assert_eq!(utility[0].as_ref(), &[0x03]);
1304 +
1305 + let payload = decode_value(&utility[1]);
1306 + let array = match payload {
1307 + Value::Array(array) => array,
1308 + other => panic!("expected utility payload array, got {other:?}"),
1309 + };
1310 + assert_eq!(array.len(), 4);
1311 + assert_eq!(array[0], Value::from(5));
1312 + let call_id = array[1].as_i64().expect("call_id");
1313 + assert_eq!(array[2], Value::from("is_sleeping"));
1314 + assert_eq!(array[3], Value::Array(Vec::new()));
1315 +
1316 + send_outputs(
1317 + push,
1318 + EngineCoreOutputs {
1319 + utility_output: Some(UtilityOutput {
1320 + call_id,
1321 + failure_message: None,
1322 + result: Some(utility_result_value(true)),
1323 + }),
1324 + ..Default::default()
1325 + },
1326 + )
1327 + .await;
1328 + })
1329 + },
1330 + );
1331 +
1332 + let client = connect_client_with_ipc(
1333 + handshake_test_config(
1334 + handshake_address,
1335 + 1,
1336 + "test-model",
1337 + Duration::from_secs(2),
1338 + 5,
1339 + None,
1340 + ),
1341 + &ipc,
1342 + )
1343 + .await;
1344 +
1345 + let result = client.is_sleeping().await.unwrap();
1346 + assert!(result);
1347 +
1348 + let _ = shutdown_tx.send(());
1349 + engine_task.await.unwrap();
1350 + client.shutdown().await.unwrap();
1351 +}
1352 +
1353 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1354 +async fn call_utility_failure_message_surfaces_as_error() {
1355 + init_tracing();
1356 + let ipc = IpcNamespace::new().unwrap();
1357 + let handshake_address = ipc.handshake_endpoint();
1358 + let engine_id = b"engine-utility-fail".to_vec();
1359 +
1360 + let (shutdown_tx, engine_task) = spawn_mock_engine_task(
1361 + handshake_address.clone(),
1362 + engine_id.clone(),
1363 + |dealer, push| {
1364 + Box::pin(async move {
1365 + let utility = recv_engine_message(dealer).await;
1366 + assert_eq!(utility[0].as_ref(), &[0x03]);
1367 + let payload = decode_value(&utility[1]);
1368 + let call_id =
1369 + payload.as_array().and_then(|array| array[1].as_i64()).expect("call_id");
1370 +
1371 + send_outputs(
1372 + push,
1373 + EngineCoreOutputs {
1374 + utility_output: Some(UtilityOutput {
1375 + call_id,
1376 + failure_message: Some("boom".to_string()),
1377 + result: None,
1378 + }),
1379 + ..Default::default()
1380 + },
1381 + )
1382 + .await;
1383 + })
1384 + },
1385 + );
1386 +
1387 + let client = connect_client_with_ipc(
1388 + handshake_test_config(
1389 + handshake_address,
1390 + 1,
1391 + "test-model",
1392 + Duration::from_secs(2),
1393 + 0,
1394 + None,
1395 + ),
1396 + &ipc,
1397 + )
1398 + .await;
1399 +
1400 + let error = client.call_utility::<bool, _>("is_sleeping", ()).await.unwrap_err();
1401 + assert!(matches!(
1402 + error,
1403 + Error::UtilityCallFailed {
1404 + method,
1405 + message,
1406 + ..
1407 + } if method == "is_sleeping" && message == "boom"
1408 + ));
1409 +
1410 + let _ = shutdown_tx.send(());
1411 + engine_task.await.unwrap();
1412 + client.shutdown().await.unwrap();
1413 +}
1414 +
1415 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1416 +async fn dispatcher_failure_propagates_to_waiting_utility_calls() {
1417 + init_tracing();
1418 + let ipc = IpcNamespace::new().unwrap();
1419 + let handshake_address = ipc.handshake_endpoint();
1420 + let engine_id = b"engine-utility-dispatcher-fail".to_vec();
1421 +
1422 + let (shutdown_tx, engine_task) = spawn_mock_engine_task(
1423 + handshake_address.clone(),
1424 + engine_id.clone(),
1425 + |dealer, push| {
1426 + Box::pin(async move {
1427 + let utility = recv_engine_message(dealer).await;
1428 + assert_eq!(utility[0].as_ref(), &[0x03]);
1429 +
1430 + push.send(ZmqMessage::from(vec![0xc1])).await.unwrap();
1431 + })
1432 + },
1433 + );
1434 +
1435 + let client = connect_client_with_ipc(
1436 + handshake_test_config(
1437 + handshake_address,
1438 + 1,
1439 + "test-model",
1440 + Duration::from_secs(2),
1441 + 0,
1442 + None,
1443 + ),
1444 + &ipc,
1445 + )
1446 + .await;
1447 +
1448 + let error = client.call_utility::<bool, _>("is_sleeping", ()).await.unwrap_err();
1449 + assert!(is_decode_error(&error));
1450 + assert!(is_decode_error(
1451 + client.health_error().as_deref().expect("health error recorded")
1452 + ));
1453 +
1454 + let _ = shutdown_tx.send(());
1455 + engine_task.await.unwrap();
1456 + client.shutdown().await.unwrap();
1457 +}
1458 +
1459 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1460 +async fn connect_times_out_without_ready_message() {
1461 + init_tracing();
1462 + let ipc = IpcNamespace::new().unwrap();
1463 + let handshake_address = ipc.handshake_endpoint();
1464 + let engine_handshake = handshake_address.clone();
1465 + let engine_task = tokio::spawn(async move {
1466 + tokio::time::sleep(Duration::from_millis(50)).await;
1467 +
1468 + let mut options = SocketOptions::default();
1469 + options.peer_identity(PeerIdentity::try_from(b"engine-timeout".to_vec()).unwrap());
1470 + let mut handshake = DealerSocket::with_options(options);
1471 + handshake.connect(&engine_handshake).await.unwrap();
1472 + handshake
1473 + .send(ZmqMessage::from(
1474 + rmp_serde::to_vec_named(&ready_message("HELLO")).unwrap(),
1475 + ))
1476 + .await
1477 + .unwrap();
1478 +
1479 + let _ = handshake.recv().await.unwrap();
1480 + });
1481 +
1482 + let result = EngineCoreClient::connect(
1483 + handshake_test_config(
1484 + handshake_address,
1485 + 1,
1486 + "test-model",
1487 + Duration::from_millis(100),
1488 + 0,
1489 + None,
1490 + )
1491 + .with_local_input_output_addresses(Some(ipc.input_endpoint()), Some(ipc.output_endpoint())),
1492 + )
1493 + .await;
1494 +
1495 + let error = match result {
1496 + Ok(_) => panic!("expected ready timeout"),
1497 + Err(error) => error,
1498 + };
1499 +
1500 + let message = error.to_report_string();
1501 + assert!(message.contains("timed out"));
1502 + assert!(message.contains("READY"));
1503 + engine_task.await.unwrap();
1504 +}
1505 +
1506 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1507 +async fn engine_core_dead_sentinel_marks_client_unhealthy_and_sticks() {
1508 + init_tracing();
1509 + let ipc = IpcNamespace::new().unwrap();
1510 + let handshake_address = ipc.handshake_endpoint();
1511 + let engine_id = b"engine-dead".to_vec();
1512 +
1513 + let (shutdown_tx, engine_task) = spawn_mock_engine_task(
1514 + handshake_address.clone(),
1515 + engine_id.clone(),
1516 + |_dealer, push| {
1517 + Box::pin(async move {
1518 + push.send(ZmqMessage::from(ENGINE_CORE_DEAD_SENTINEL.to_vec())).await.unwrap();
1519 + })
1520 + },
1521 + );
1522 +
1523 + let client = connect_client_with_ipc(
1524 + handshake_test_config(
1525 + handshake_address,
1526 + 1,
1527 + "test-model",
1528 + Duration::from_secs(2),
1529 + 0,
1530 + None,
1531 + ),
1532 + &ipc,
1533 + )
1534 + .await;
1535 +
1536 + timeout(Duration::from_secs(2), async {
1537 + while client.is_healthy() {
1538 + tokio::task::yield_now().await;
1539 + }
1540 + })
1541 + .await
1542 + .expect("wait for unhealthy client");
1543 +
1544 + assert!(!client.is_healthy());
1545 + assert!(matches!(
1546 + client.health_error().as_deref(),
1547 + Some(Error::EngineCoreDead)
1548 + ));
1549 +
1550 + let error = client.call_utility::<bool, _>("is_sleeping", ()).await.unwrap_err();
1551 + assert!(
1552 + is_dispatcher_closed(&error) || is_engine_core_dead(&error),
1553 + "unexpected error: {error:?}"
1554 + );
1555 + assert!(matches!(
1556 + client.health_error().as_deref(),
1557 + Some(Error::EngineCoreDead)
1558 + ));
1559 +
1560 + let _ = shutdown_tx.send(());
1561 + engine_task.await.unwrap();
1562 + client.shutdown().await.unwrap();
1563 +}
1564 +
1565 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1566 +async fn output_loop_failure_marks_client_unhealthy_and_records_first_error() {
1567 + init_tracing();
1568 + let ipc = IpcNamespace::new().unwrap();
1569 + let handshake_address = ipc.handshake_endpoint();
1570 + let engine_id = b"engine-output-failure".to_vec();
1571 +
1572 + let (shutdown_tx, engine_task) = spawn_mock_engine_task(
1573 + handshake_address.clone(),
1574 + engine_id.clone(),
1575 + |_dealer, push| {
1576 + Box::pin(async move {
1577 + send_output_frames(
1578 + push,
1579 + vec![
1580 + bytes::Bytes::from_static(b"frame-1"),
1581 + bytes::Bytes::from_static(b"frame-2"),
1582 + ],
1583 + )
1584 + .await;
1585 + })
1586 + },
1587 + );
1588 +
1589 + let client = connect_client_with_ipc(
1590 + handshake_test_config(
1591 + handshake_address,
1592 + 1,
1593 + "test-model",
1594 + Duration::from_secs(2),
1595 + 0,
1596 + None,
1597 + ),
1598 + &ipc,
1599 + )
1600 + .await;
1601 +
1602 + timeout(Duration::from_secs(2), async {
1603 + while client.is_healthy() {
1604 + let _ = client.call_utility::<bool, _>("is_sleeping", ()).await;
1605 + tokio::task::yield_now().await;
1606 + }
1607 + })
1608 + .await
1609 + .expect("wait for unhealthy client");
1610 +
1611 + assert!(!client.is_healthy());
1612 + assert!(is_decode_error(
1613 + client.health_error().as_deref().expect("health error recorded")
1614 + ));
1615 +
1616 + let _ = shutdown_tx.send(());
1617 + engine_task.await.unwrap();
1618 + client.shutdown().await.unwrap();
1619 +}
1620 +
1621 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1622 +async fn client_decodes_multipart_logprob_outputs() {
1623 + init_tracing();
1624 + let ipc = IpcNamespace::new().unwrap();
1625 + let handshake_address = ipc.handshake_endpoint();
1626 + let engine_id = b"engine-multipart-logprobs".to_vec();
1627 +
1628 + let (shutdown_tx, engine_task) = spawn_mock_engine_task(
1629 + handshake_address.clone(),
1630 + engine_id.clone(),
1631 + |dealer, push| {
1632 + Box::pin(async move {
1633 + let add = recv_engine_message(dealer).await;
1634 + assert_eq!(add[0].as_ref(), &[0x00]);
1635 + let request: EngineCoreRequest = rmp_serde::from_slice(&add[1]).unwrap();
1636 + assert_eq!(request.request_id, "req-1");
1637 +
1638 + send_output_frames(push, multipart_logprob_output_frames("req-1")).await;
1639 + })
1640 + },
1641 + );
1642 +
1643 + let client = connect_client_with_ipc(
1644 + handshake_test_config(
1645 + handshake_address,
1646 + 1,
1647 + "test-model",
1648 + Duration::from_secs(2),
1649 + 0,
1650 + None,
1651 + ),
1652 + &ipc,
1653 + )
1654 + .await;
1655 +
1656 + let stream = client.call(sample_request()).await.unwrap();
1657 + let outputs = stream.collect::<Vec<_>>().await;
1658 + assert_eq!(outputs.len(), 1);
1659 +
1660 + let output = outputs.into_iter().next().unwrap().unwrap();
1661 + assert_eq!(output.output.new_token_ids, vec![7, 8]);
1662 + assert_eq!(
1663 + output.output.finish_reason,
1664 + Some(EngineCoreFinishReason::Length)
1665 + );
1666 + expect_sample_logprobs(output.output.new_logprobs.as_ref().expect("logprobs decoded"));
1667 +
1668 + let _ = shutdown_tx.send(());
1669 + engine_task.await.unwrap();
1670 + client.shutdown().await.unwrap();
1671 +}
1672 +
1673 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1674 +async fn multi_engine_client_shares_transport_and_routes_by_inflight_count() {
1675 + init_tracing();
1676 + let ipc = IpcNamespace::new().unwrap();
1677 + let handshake_address = ipc.handshake_endpoint();
1678 + let (engine_0_seen_tx, mut engine_0_seen_rx) = mpsc::unbounded_channel();
1679 + let (engine_1_seen_tx, engine_1_seen_rx) = oneshot::channel();
1680 + let (finish_req_1_tx, finish_req_1_rx) = oneshot::channel();
1681 + let (finish_req_2_tx, finish_req_2_rx) = oneshot::channel();
1682 + let (finish_req_3_tx, finish_req_3_rx) = oneshot::channel();
1683 +
1684 + let (init_rx_0, shutdown_tx_0, engine_task_0) = spawn_mock_engine_task_with_init(
1685 + handshake_address.clone(),
1686 + b"engine-0".to_vec(),
1687 + |dealer, push| {
1688 + Box::pin(async move {
1689 + let add_1 = recv_engine_message(dealer).await;
1690 + assert_eq!(add_1[0].as_ref(), &[0x00]);
1691 + let request_1: EngineCoreRequest = rmp_serde::from_slice(&add_1[1]).unwrap();
1692 + assert_eq!(request_1.request_id, "req-1");
1693 + engine_0_seen_tx.send(request_1.request_id.clone()).unwrap();
1694 + finish_req_1_rx.await.unwrap();
1695 + send_outputs(
1696 + push,
1697 + EngineCoreOutputs {
1698 + engine_index: 0,
1699 + outputs: vec![request_output(
1700 + &request_1.request_id,
1701 + vec![10],
1702 + Some(EngineCoreFinishReason::Length),
1703 + )],
1704 + finished_requests: Some(BTreeSet::from([request_1.request_id.clone()])),
1705 + ..Default::default()
1706 + },
1707 + )
1708 + .await;
1709 +
1710 + let add_3 = recv_engine_message(dealer).await;
1711 + assert_eq!(add_3[0].as_ref(), &[0x00]);
1712 + let request_3: EngineCoreRequest = rmp_serde::from_slice(&add_3[1]).unwrap();
1713 + assert_eq!(request_3.request_id, "req-3");
1714 + engine_0_seen_tx.send(request_3.request_id.clone()).unwrap();
1715 + finish_req_3_rx.await.unwrap();
1716 + send_outputs(
1717 + push,
1718 + EngineCoreOutputs {
1719 + engine_index: 0,
1720 + outputs: vec![request_output(
1721 + &request_3.request_id,
1722 + vec![30],
1723 + Some(EngineCoreFinishReason::Length),
1724 + )],
1725 + finished_requests: Some(BTreeSet::from([request_3.request_id.clone()])),
1726 + ..Default::default()
1727 + },
1728 + )
1729 + .await;
1730 + })
1731 + },
1732 + );
1733 + let (init_rx_1, shutdown_tx_1, engine_task_1) = spawn_mock_engine_task_with_init(
1734 + handshake_address.clone(),
1735 + b"engine-1".to_vec(),
1736 + |dealer, push| {
1737 + Box::pin(async move {
1738 + let add_2 = recv_engine_message(dealer).await;
1739 + assert_eq!(add_2[0].as_ref(), &[0x00]);
1740 + let request_2: EngineCoreRequest = rmp_serde::from_slice(&add_2[1]).unwrap();
1741 + assert_eq!(request_2.request_id, "req-2");
1742 + let _ = engine_1_seen_tx.send(request_2.request_id.clone());
1743 + finish_req_2_rx.await.unwrap();
1744 + send_outputs(
1745 + push,
1746 + EngineCoreOutputs {
1747 + engine_index: 1,
1748 + outputs: vec![request_output(
1749 + &request_2.request_id,
1750 + vec![20],
1751 + Some(EngineCoreFinishReason::Length),
1752 + )],
1753 + finished_requests: Some(BTreeSet::from([request_2.request_id.clone()])),
1754 + ..Default::default()
1755 + },
1756 + )
1757 + .await;
1758 + })
1759 + },
1760 + );
1761 +
1762 + let client = connect_client_with_ipc(
1763 + handshake_test_config(
1764 + handshake_address.clone(),
1765 + 2,
1766 + "test-model",
1767 + Duration::from_secs(2),
1768 + 0,
1769 + None,
1770 + ),
1771 + &ipc,
1772 + )
1773 + .await;
1774 +
1775 + let init_0 = timeout(Duration::from_secs(1), init_rx_0).await.unwrap().unwrap();
1776 + let init_1 = timeout(Duration::from_secs(1), init_rx_1).await.unwrap().unwrap();
1777 + assert_eq!(init_0.addresses.inputs, vec![ipc.input_endpoint()]);
1778 + assert_eq!(init_1.addresses.inputs, vec![ipc.input_endpoint()]);
1779 + assert_eq!(init_0.addresses.outputs, vec![ipc.output_endpoint()]);
1780 + assert_eq!(init_1.addresses.outputs, vec![ipc.output_endpoint()]);
1781 +
1782 + assert_eq!(client.input_address(), ipc.input_endpoint());
1783 + assert_eq!(client.output_address(), ipc.output_endpoint());
1784 + assert_eq!(client.engine_count(), 2);
1785 + assert_eq!(
1786 + client.engine_identities(),
1787 + vec![b"engine-0".as_slice(), b"engine-1".as_slice()]
1788 + );
1789 + assert_eq!(client.ready_responses().len(), 2);
1790 + assert_eq!(client.engine_identities()[0], b"engine-0");
1791 +
1792 + let mut stream_1 = client.call(sample_request_with_id("req-1")).await.unwrap();
1793 + let mut stream_2 = client.call(sample_request_with_id("req-2")).await.unwrap();
1794 + assert_eq!(
1795 + timeout(Duration::from_secs(1), engine_0_seen_rx.recv()).await.unwrap().unwrap(),
1796 + "req-1"
1797 + );
1798 + assert_eq!(
1799 + timeout(Duration::from_secs(1), engine_1_seen_rx).await.unwrap().unwrap(),
1800 + "req-2"
1801 + );
1802 +
1803 + let _ = finish_req_1_tx.send(());
1804 + let final_1 = timeout(Duration::from_secs(1), stream_1.next())
1805 + .await
1806 + .unwrap()
1807 + .unwrap()
1808 + .unwrap();
1809 + assert_eq!(final_1.engine_index, 0);
1810 + assert_eq!(final_1.new_token_ids, vec![10]);
1811 + assert_eq!(final_1.finish_reason, Some(EngineCoreFinishReason::Length));
1812 +
1813 + let mut stream_3 = client.call(sample_request_with_id("req-3")).await.unwrap();
1814 + assert_eq!(
1815 + timeout(Duration::from_secs(1), engine_0_seen_rx.recv()).await.unwrap().unwrap(),
1816 + "req-3"
1817 + );
1818 +
1819 + let _ = finish_req_3_tx.send(());
1820 + let final_3 = timeout(Duration::from_secs(1), stream_3.next())
1821 + .await
1822 + .unwrap()
1823 + .unwrap()
1824 + .unwrap();
1825 + assert_eq!(final_3.engine_index, 0);
1826 + assert_eq!(final_3.new_token_ids, vec![30]);
1827 + assert_eq!(final_3.finish_reason, Some(EngineCoreFinishReason::Length));
1828 +
1829 + let _ = finish_req_2_tx.send(());
1830 + let final_2 = timeout(Duration::from_secs(1), stream_2.next())
1831 + .await
1832 + .unwrap()
1833 + .unwrap()
1834 + .unwrap();
1835 + assert_eq!(final_2.engine_index, 1);
1836 + assert_eq!(final_2.new_token_ids, vec![20]);
1837 + assert_eq!(final_2.finish_reason, Some(EngineCoreFinishReason::Length));
1838 +
1839 + assert!(timeout(Duration::from_secs(1), stream_1.next()).await.unwrap().is_none());
1840 + assert!(timeout(Duration::from_secs(1), stream_2.next()).await.unwrap().is_none());
1841 + assert!(timeout(Duration::from_secs(1), stream_3.next()).await.unwrap().is_none());
1842 +
1843 + let _ = shutdown_tx_0.send(());
1844 + let _ = shutdown_tx_1.send(());
1845 + engine_task_0.await.unwrap();
1846 + engine_task_1.await.unwrap();
1847 + client.shutdown().await.unwrap();
1848 +}
1849 +
1850 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1851 +async fn multi_engine_abort_is_grouped_and_utility_fans_out_to_all_engines() {
1852 + init_tracing();
1853 + let ipc = IpcNamespace::new().unwrap();
1854 + let handshake_address = ipc.handshake_endpoint();
1855 +
1856 + let (shutdown_tx_0, engine_task_0) = spawn_mock_engine_task(
1857 + handshake_address.clone(),
1858 + b"engine-0".to_vec(),
1859 + |dealer, push| {
1860 + Box::pin(async move {
1861 + let utility = recv_engine_message(dealer).await;
1862 + assert_eq!(utility[0].as_ref(), &[0x03]);
1863 + let payload = decode_value(&utility[1]);
1864 + let array = match payload {
1865 + Value::Array(array) => array,
1866 + other => panic!("expected utility payload array, got {other:?}"),
1867 + };
1868 + let call_id = array[1].as_i64().expect("call_id");
1869 + assert_eq!(array[2], Value::from("is_sleeping"));
1870 + send_outputs(
1871 + push,
1872 + EngineCoreOutputs {
1873 + utility_output: Some(UtilityOutput {
1874 + call_id,
1875 + failure_message: None,
1876 + result: Some(utility_result_value(true)),
1877 + }),
1878 + ..Default::default()
1879 + },
1880 + )
1881 + .await;
1882 +
1883 + let add_1 = recv_engine_message(dealer).await;
1884 + assert_eq!(add_1[0].as_ref(), &[0x00]);
1885 + let request_1: EngineCoreRequest = rmp_serde::from_slice(&add_1[1]).unwrap();
1886 + assert_eq!(request_1.request_id, "req-1");
1887 +
1888 + let abort = recv_engine_message(dealer).await;
1889 + assert_eq!(abort[0].as_ref(), &[0x01]);
1890 + let aborted_ids: Vec<String> = rmp_serde::from_slice(&abort[1]).unwrap();
1891 + assert_eq!(aborted_ids, vec!["req-1".to_string()]);
1892 + send_outputs(
1893 + push,
1894 + EngineCoreOutputs {
1895 + engine_index: 0,
1896 + outputs: vec![request_output(
1897 + "req-1",
1898 + vec![],
1899 + Some(EngineCoreFinishReason::Abort),
1900 + )],
1901 + finished_requests: Some(BTreeSet::from(["req-1".to_string()])),
1902 + ..Default::default()
1903 + },
1904 + )
1905 + .await;
1906 + })
1907 + },
1908 + );
1909 + tokio::time::sleep(Duration::from_millis(50)).await;
1910 + let (shutdown_tx_1, engine_task_1) = spawn_mock_engine_task(
1911 + handshake_address.clone(),
1912 + b"engine-1".to_vec(),
1913 + |dealer, push| {
1914 + Box::pin(async move {
1915 + let utility = recv_engine_message(dealer).await;
1916 + assert_eq!(utility[0].as_ref(), &[0x03]);
1917 + let payload = decode_value(&utility[1]);
1918 + let array = match payload {
1919 + Value::Array(array) => array,
1920 + other => panic!("expected utility payload array, got {other:?}"),
1921 + };
1922 + let call_id = array[1].as_i64().expect("call_id");
1923 + assert_eq!(array[2], Value::from("is_sleeping"));
1924 + send_outputs(
1925 + push,
1926 + EngineCoreOutputs {
1927 + utility_output: Some(UtilityOutput {
1928 + call_id,
1929 + failure_message: None,
1930 + result: Some(utility_result_value(true)),
1931 + }),
1932 + ..Default::default()
1933 + },
1934 + )
1935 + .await;
1936 +
1937 + let add_2 = recv_engine_message(dealer).await;
1938 + assert_eq!(add_2[0].as_ref(), &[0x00]);
1939 + let request_2: EngineCoreRequest = rmp_serde::from_slice(&add_2[1]).unwrap();
1940 + assert_eq!(request_2.request_id, "req-2");
1941 +
1942 + let abort = recv_engine_message(dealer).await;
1943 + assert_eq!(abort[0].as_ref(), &[0x01]);
1944 + let aborted_ids: Vec<String> = rmp_serde::from_slice(&abort[1]).unwrap();
1945 + assert_eq!(aborted_ids, vec!["req-2".to_string()]);
1946 + send_outputs(
1947 + push,
1948 + EngineCoreOutputs {
1949 + engine_index: 1,
1950 + outputs: vec![request_output(
1951 + "req-2",
1952 + vec![],
1953 + Some(EngineCoreFinishReason::Abort),
1954 + )],
1955 + finished_requests: Some(BTreeSet::from(["req-2".to_string()])),
1956 + ..Default::default()
1957 + },
1958 + )
1959 + .await;
1960 + })
1961 + },
1962 + );
1963 +
1964 + let client = connect_client_with_ipc(
1965 + handshake_test_config(
1966 + handshake_address,
1967 + 2,
1968 + "test-model",
1969 + Duration::from_secs(2),
1970 + 5,
1971 + None,
1972 + ),
1973 + &ipc,
1974 + )
1975 + .await;
1976 +
1977 + assert!(client.is_sleeping().await.unwrap());
1978 +
1979 + let mut stream_1 = client.call(sample_request_with_id("req-1")).await.unwrap();
1980 + let mut stream_2 = client.call(sample_request_with_id("req-2")).await.unwrap();
1981 +
1982 + client
1983 + .abort(&[
1984 + "req-2".to_string(),
1985 + "req-1".to_string(),
1986 + "unknown".to_string(),
1987 + ])
1988 + .await
1989 + .unwrap();
1990 +
1991 + let final_1 = timeout(Duration::from_secs(1), stream_1.next())
1992 + .await
1993 + .unwrap()
1994 + .unwrap()
1995 + .unwrap();
1996 + assert_eq!(final_1.engine_index, 0);
1997 + assert_eq!(final_1.finish_reason, Some(EngineCoreFinishReason::Abort));
1998 +
1999 + let final_2 = timeout(Duration::from_secs(1), stream_2.next())
2000 + .await
2001 + .unwrap()
2002 + .unwrap()
2003 + .unwrap();
2004 + assert_eq!(final_2.engine_index, 1);
2005 + assert_eq!(final_2.finish_reason, Some(EngineCoreFinishReason::Abort));
2006 +
2007 + let _ = shutdown_tx_0.send(());
2008 + let _ = shutdown_tx_1.send(());
2009 + engine_task_0.await.unwrap();
2010 + engine_task_1.await.unwrap();
2011 + client.shutdown().await.unwrap();
2012 +}
2013 +
2014 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
2015 +async fn collective_rpc_flattens_results_from_all_engines() {
2016 + init_tracing();
2017 + let ipc = IpcNamespace::new().unwrap();
2018 + let handshake_address = ipc.handshake_endpoint();
2019 +
2020 + let (shutdown_tx_0, engine_task_0) = spawn_mock_engine_task(
2021 + handshake_address.clone(),
2022 + b"engine-0".to_vec(),
2023 + |dealer, push| {
2024 + Box::pin(async move {
2025 + let utility = recv_engine_message(dealer).await;
2026 + assert_eq!(utility[0].as_ref(), &[0x03]);
2027 + let payload = decode_value(&utility[1]);
2028 + let array = match payload {
2029 + Value::Array(array) => array,
2030 + other => panic!("expected utility payload array, got {other:?}"),
2031 + };
2032 + let call_id = array[1].as_i64().expect("call_id");
2033 + assert_eq!(array[2], Value::from("collective_rpc"));
2034 +
2035 + send_outputs(
2036 + push,
2037 + EngineCoreOutputs {
2038 + utility_output: Some(UtilityOutput {
2039 + call_id,
2040 + failure_message: None,
2041 + result: Some(utility_result_value(vec!["engine-0-worker"])),
2042 + }),
2043 + ..Default::default()
2044 + },
2045 + )
2046 + .await;
2047 + })
2048 + },
2049 + );
2050 + tokio::time::sleep(Duration::from_millis(50)).await;
2051 + let (shutdown_tx_1, engine_task_1) = spawn_mock_engine_task(
2052 + handshake_address.clone(),
2053 + b"engine-1".to_vec(),
2054 + |dealer, push| {
2055 + Box::pin(async move {
2056 + let utility = recv_engine_message(dealer).await;
2057 + assert_eq!(utility[0].as_ref(), &[0x03]);
2058 + let payload = decode_value(&utility[1]);
2059 + let array = match payload {
2060 + Value::Array(array) => array,
2061 + other => panic!("expected utility payload array, got {other:?}"),
2062 + };
2063 + let call_id = array[1].as_i64().expect("call_id");
2064 + assert_eq!(array[2], Value::from("collective_rpc"));
2065 +
2066 + send_outputs(
2067 + push,
2068 + EngineCoreOutputs {
2069 + utility_output: Some(UtilityOutput {
2070 + call_id,
2071 + failure_message: None,
2072 + result: Some(utility_result_value(vec!["engine-1-worker"])),
2073 + }),
2074 + ..Default::default()
2075 + },
2076 + )
2077 + .await;
2078 + })
2079 + },
2080 + );
2081 +
2082 + let client = connect_client_with_ipc(
2083 + handshake_test_config(
2084 + handshake_address,
2085 + 2,
2086 + "test-model",
2087 + Duration::from_secs(2),
2088 + 5,
2089 + None,
2090 + ),
2091 + &ipc,
2092 + )
2093 + .await;
2094 +
2095 + let results = client
2096 + .collective_rpc(
2097 + "get_model_name",
2098 + Option::<f64>::None,
2099 + Vec::<String>::new(),
2100 + BTreeMap::<String, String>::new(),
2101 + )
2102 + .await
2103 + .unwrap();
2104 + assert_eq!(
2105 + results,
2106 + vec![
2107 + Value::from("engine-0-worker"),
2108 + Value::from("engine-1-worker")
2109 + ]
2110 + );
2111 +
2112 + let _ = shutdown_tx_0.send(());
2113 + let _ = shutdown_tx_1.send(());
2114 + engine_task_0.await.unwrap();
2115 + engine_task_1.await.unwrap();
2116 + client.shutdown().await.unwrap();
2117 +}
2118 +
2119 +#[test]
2120 +fn python_msgpack_fixtures_match_rust_encoding() {
2121 + init_tracing();
2122 + let script = PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("src/tests/python_compat.py");
2123 + let output = Command::new(&script)
2124 + .output()
2125 + .unwrap_or_else(|error| panic!("failed to execute {:?}: {error}", script));
2126 + assert!(
2127 + output.status.success(),
2128 + "python fixture script failed: status={:?}\nstdout:\n{}\nstderr:\n{}",
2129 + output.status.code(),
2130 + String::from_utf8_lossy(&output.stdout),
2131 + String::from_utf8_lossy(&output.stderr),
2132 + );
2133 +
2134 + let stdout = String::from_utf8(output.stdout).unwrap();
2135 + let mut lines = stdout.lines();
2136 + let request_hex = lines.next().expect("missing request fixture line");
2137 + let multimodal_request_hex = lines.next().expect("missing multimodal request fixture line");
2138 + let outputs_hex = lines.next().expect("missing outputs fixture line");
2139 + let inline_logprobs_frames = lines.next().expect("missing inline logprobs fixture line");
2140 + let multipart_logprobs_frames = lines.next().expect("missing multipart logprobs fixture line");
2141 + let inline_prompt_frames = lines.next().expect("missing inline prompt logprobs fixture line");
2142 + let multipart_prompt_frames =
2143 + lines.next().expect("missing multipart prompt logprobs fixture line");
2144 +
2145 + let request_bytes = hex::decode(request_hex).unwrap();
2146 + let multimodal_request_bytes = hex::decode(multimodal_request_hex).unwrap();
2147 + let outputs_bytes = hex::decode(outputs_hex).unwrap();
2148 +
2149 + let decoded_request: EngineCoreRequest = rmp_serde::from_slice(&request_bytes).unwrap();
2150 + let expected_request = sample_request();
2151 + assert_eq!(decoded_request, expected_request);
2152 +
2153 + let decoded_multimodal_request: EngineCoreRequest =
2154 + rmp_serde::from_slice(&multimodal_request_bytes).unwrap();
2155 + assert_eq!(decoded_multimodal_request, sample_multimodal_request());
2156 +
2157 + // The decode assertion above proves Python wire -> Rust struct. Also compare
2158 + // Rust struct -> wire for the multimodal subtree, which is the frontend's
2159 + // production direction when sending requests to Python EngineCore.
2160 + let expected_multimodal_request = sample_multimodal_request();
2161 + let decode_value = |bytes: &[u8]| {
2162 + rmpv::decode::read_value(&mut Cursor::new(bytes)).expect("decode msgpack value")
2163 + };
2164 + let extract_mm_features = |value: Value| match value {
2165 + Value::Array(items) => items.get(2).cloned().expect("request mm_features slot"),
2166 + other => panic!("request should encode as tuple array, got {other:?}"),
2167 + };
2168 + let python_mm_features = extract_mm_features(decode_value(&multimodal_request_bytes));
2169 + let rust_mm_features =
2170 + decode_value(&rmp_serde::to_vec_named(&expected_multimodal_request.mm_features).unwrap());
2171 + assert_eq!(python_mm_features, rust_mm_features);
2172 +
2173 + let decoded_outputs: EngineCoreOutputs = rmp_serde::from_slice(&outputs_bytes).unwrap();
2174 + expect_test::expect![[r#"
2175 + EngineCoreOutputs {
2176 + engine_index: 0,
2177 + outputs: [
2178 + EngineCoreOutput {
2179 + request_id: "req-1",
2180 + new_token_ids: [
2181 + 7,
2182 + 8,
2183 + ],
2184 + new_logprobs: None,
2185 + new_prompt_logprobs_tensors: None,
2186 + pooling_output: None,
2187 + finish_reason: Some(
2188 + Length,
2189 + ),
2190 + stop_reason: None,
2191 + events: None,
2192 + kv_transfer_params: None,
2193 + trace_headers: None,
2194 + prefill_stats: None,
2195 + routed_experts: None,
2196 + num_nans_in_logits: 0,
2197 + },
2198 + ],
2199 + scheduler_stats: None,
2200 + timestamp: 0.0,
2201 + utility_output: None,
2202 + finished_requests: Some(
2203 + {
2204 + "req-1",
2205 + },
2206 + ),
2207 + wave_complete: None,
2208 + start_wave: None,
2209 + }
2210 + "#]]
2211 + .assert_debug_eq(&decoded_outputs);
2212 +
2213 + let decode_frames = |line: &str| {
2214 + line.split_whitespace()
2215 + .map(|frame| bytes::Bytes::from(hex::decode(frame).unwrap()))
2216 + .collect::<Vec<_>>()
2217 + };
2218 +
2219 + let inline_logprobs =
2220 + decode_engine_core_outputs(&decode_frames(inline_logprobs_frames)).unwrap();
2221 + expect_sample_logprobs(
2222 + inline_logprobs.outputs[0]
2223 + .new_logprobs
2224 + .as_ref()
2225 + .expect("inline logprobs decoded"),
2226 + );
2227 +
2228 + let multipart_logprobs =
2229 + decode_engine_core_outputs(&decode_frames(multipart_logprobs_frames)).unwrap();
2230 + expect_sample_logprobs(
2231 + multipart_logprobs.outputs[0]
2232 + .new_logprobs
2233 + .as_ref()
2234 + .expect("multipart logprobs decoded"),
2235 + );
2236 +
2237 + let inline_prompt = decode_engine_core_outputs(&decode_frames(inline_prompt_frames)).unwrap();
2238 + expect_prompt_logprobs(
2239 + inline_prompt.outputs[0]
2240 + .new_prompt_logprobs_tensors
2241 + .as_ref()
2242 + .expect("inline prompt logprobs decoded"),
2243 + );
2244 +
2245 + let multipart_prompt =
2246 + decode_engine_core_outputs(&decode_frames(multipart_prompt_frames)).unwrap();
2247 + expect_prompt_logprobs(
2248 + multipart_prompt.outputs[0]
2249 + .new_prompt_logprobs_tensors
2250 + .as_ref()
2251 + .expect("multipart prompt logprobs decoded"),
2252 + );
2253 +}
2254 +
2255 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
2256 +async fn bootstrapped_connects_after_single_engine_registration() {
2257 + init_tracing();
2258 + let ipc = IpcNamespace::new().unwrap();
2259 + let input_address = ipc.input_endpoint();
2260 + let output_address = ipc.output_endpoint();
2261 +
2262 + let client_task = tokio::spawn({
2263 + let input_address = input_address.clone();
2264 + let output_address = output_address.clone();
2265 + async move {
2266 + EngineCoreClient::connect(bootstrapped_test_config(
2267 + input_address,
2268 + output_address,
2269 + 1,
2270 + Duration::from_secs(2),
2271 + 0,
2272 + None,
2273 + ))
2274 + .await
2275 + .unwrap()
2276 + }
2277 + });
2278 +
2279 + let (_dealer, _push) =
2280 + setup_bootstrapped_mock_engine(input_address, output_address, &[0x00, 0x00]).await;
2281 + let client = client_task.await.unwrap();
2282 +
2283 + assert_eq!(client.engine_count(), 1);
2284 + let engine_ids =
2285 + client.engine_identities().into_iter().map(|id| id.to_vec()).collect::<Vec<_>>();
2286 + assert_eq!(engine_ids, vec![vec![0x00, 0x00]]);
2287 +
2288 + client.shutdown().await.unwrap();
2289 +}
2290 +
2291 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
2292 +async fn bootstrapped_connects_with_contiguous_engine_ids() {
2293 + init_tracing();
2294 + let ipc = IpcNamespace::new().unwrap();
2295 + let input_address = ipc.input_endpoint();
2296 + let output_address = ipc.output_endpoint();
2297 +
2298 + let client_task = tokio::spawn({
2299 + let input_address = input_address.clone();
2300 + let output_address = output_address.clone();
2301 + async move {
2302 + EngineCoreClient::connect(bootstrapped_test_config(
2303 + input_address,
2304 + output_address,
2305 + 2,
2306 + Duration::from_secs(2),
2307 + 0,
2308 + None,
2309 + ))
2310 + .await
2311 + .unwrap()
2312 + }
2313 + });
2314 +
2315 + let (_dealer0, _push0) = setup_bootstrapped_mock_engine(
2316 + input_address.clone(),
2317 + output_address.clone(),
2318 + &[0x00, 0x00],
2319 + )
2320 + .await;
2321 + let (_dealer1, _push1) =
2322 + setup_bootstrapped_mock_engine(input_address, output_address, &[0x01, 0x00]).await;
2323 + let client = client_task.await.unwrap();
2324 +
2325 + assert_eq!(client.engine_count(), 2);
2326 + let engine_ids =
2327 + client.engine_identities().into_iter().map(|id| id.to_vec()).collect::<Vec<_>>();
2328 + assert_eq!(engine_ids, vec![vec![0x00, 0x00], vec![0x01, 0x00]]);
2329 +
2330 + client.shutdown().await.unwrap();
2331 +}
2332 +
2333 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
2334 +async fn bootstrapped_connect_times_out_without_registration() {
2335 + init_tracing();
2336 + let ipc = IpcNamespace::new().unwrap();
2337 + let result = EngineCoreClient::connect(bootstrapped_test_config(
2338 + ipc.input_endpoint(),
2339 + ipc.output_endpoint(),
2340 + 1,
2341 + Duration::from_millis(100),
2342 + 0,
2343 + None,
2344 + ))
2345 + .await;
2346 +
2347 + let error = match result {
2348 + Ok(_) => panic!("bootstrapped connect should time out"),
2349 + Err(error) => error,
2350 + };
2351 + assert!(matches!(error, Error::InputRegistrationTimeout { .. }));
2352 +}
2353 +
2354 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
2355 +async fn bootstrapped_external_coordinator_connects_and_subscribes() {
2356 + init_tracing();
2357 + let ipc = IpcNamespace::new().unwrap();
2358 + let input_address = ipc.input_endpoint();
2359 + let output_address = ipc.output_endpoint();
2360 + let coordinator_address = ipc.endpoint("stats.sock");
2361 +
2362 + let mut stats_socket = XPubSocket::new();
2363 + stats_socket.bind(&coordinator_address).await.unwrap();
2364 +
2365 + let client_task = tokio::spawn({
2366 + let input_address = input_address.clone();
2367 + let output_address = output_address.clone();
2368 + let coordinator_address = coordinator_address.clone();
2369 + async move {
2370 + EngineCoreClient::connect(bootstrapped_test_config(
2371 + input_address,
2372 + output_address,
2373 + 1,
2374 + Duration::from_secs(2),
2375 + 0,
2376 + Some(CoordinatorMode::External {
2377 + address: coordinator_address,
2378 + }),
2379 + ))
2380 + .await
2381 + .unwrap()
2382 + }
2383 + });
2384 +
2385 + let (_dealer, _push) =
2386 + setup_bootstrapped_mock_engine(input_address, output_address, &[0x00, 0x00]).await;
2387 + let client = client_task.await.unwrap();
2388 +
2389 + timeout(
2390 + Duration::from_secs(1),
2391 + recv_xpub_subscription(&mut stats_socket),
2392 + )
2393 + .await
2394 + .unwrap();
2395 +
2396 + client.shutdown().await.unwrap();
2397 +}
2398 +
2399 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
2400 +async fn bootstrapped_external_coordinator_updates_wave_ignores_counts_and_sends_one_wakeup() {
2401 + init_tracing();
2402 + let ipc = IpcNamespace::new().unwrap();
2403 + let input_address = ipc.input_endpoint();
2404 + let output_address = ipc.output_endpoint();
2405 + let coordinator_address = ipc.endpoint("stats.sock");
2406 +
2407 + let mut stats_socket = XPubSocket::new();
2408 + stats_socket.bind(&coordinator_address).await.unwrap();
2409 +
2410 + let client_task = tokio::spawn({
2411 + let input_address = input_address.clone();
2412 + let output_address = output_address.clone();
2413 + let coordinator_address = coordinator_address.clone();
2414 + async move {
2415 + EngineCoreClient::connect(bootstrapped_test_config(
2416 + input_address,
2417 + output_address,
2418 + 1,
2419 + Duration::from_secs(2),
2420 + 0,
2421 + Some(CoordinatorMode::External {
2422 + address: coordinator_address,
2423 + }),
2424 + ))
2425 + .await
2426 + .unwrap()
2427 + }
2428 + });
2429 +
2430 + let (mut dealer, mut push) =
2431 + setup_bootstrapped_mock_engine(input_address, output_address, &[0x00, 0x00]).await;
2432 + let client = client_task.await.unwrap();
2433 + recv_xpub_subscription(&mut stats_socket).await;
2434 +
2435 + send_external_coordinator_publish(&mut stats_socket, &(vec![(11_u32, 3_u32)], 7_u32, false))
2436 + .await;
2437 + tokio::time::sleep(Duration::from_millis(50)).await;
2438 +
2439 + let mut stream = client.call(sample_request()).await.unwrap();
2440 +
2441 + let wakeup = timeout(
2442 + Duration::from_secs(1),
2443 + recv_external_coordinator_wakeup(&mut stats_socket),
2444 + )
2445 + .await
2446 + .unwrap();
2447 + assert_eq!(wakeup, (0, 7));
2448 +
2449 + assert!(
2450 + timeout(
2451 + Duration::from_millis(200),
2452 + recv_external_coordinator_wakeup(&mut stats_socket)
2453 + )
2454 + .await
2455 + .is_err()
2456 + );
2457 +
2458 + let add = recv_engine_message(&mut dealer).await;
2459 + assert_eq!(add[0].as_ref(), &[0x00]);
2460 + let request: EngineCoreRequest = rmp_serde::from_slice(&add[1]).unwrap();
2461 + assert_eq!(request.request_id, "req-1");
2462 + assert_eq!(request.current_wave, 7);
2463 + assert!(client.is_healthy());
2464 +
2465 + send_outputs(
2466 + &mut push,
2467 + EngineCoreOutputs {
2468 + engine_index: 0,
2469 + outputs: vec![request_output(
2470 + "req-1",
2471 + vec![],
2472 + Some(EngineCoreFinishReason::Length),
2473 + )],
2474 + finished_requests: Some(BTreeSet::from(["req-1".to_string()])),
2475 + ..Default::default()
2476 + },
2477 + )
2478 + .await;
2479 +
2480 + let final_output = timeout(Duration::from_secs(1), stream.next()).await.unwrap();
2481 + assert!(final_output.is_some());
2482 +
2483 + client.shutdown().await.unwrap();
2484 +}
2485 +
2486 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
2487 +async fn bootstrapped_external_coordinator_running_state_suppresses_wakeup() {
2488 + init_tracing();
2489 + let ipc = IpcNamespace::new().unwrap();
2490 + let input_address = ipc.input_endpoint();
2491 + let output_address = ipc.output_endpoint();
2492 + let coordinator_address = ipc.endpoint("stats.sock");
2493 +
2494 + let mut stats_socket = XPubSocket::new();
2495 + stats_socket.bind(&coordinator_address).await.unwrap();
2496 +
2497 + let client_task = tokio::spawn({
2498 + let input_address = input_address.clone();
2499 + let output_address = output_address.clone();
2500 + let coordinator_address = coordinator_address.clone();
2501 + async move {
2502 + EngineCoreClient::connect(bootstrapped_test_config(
2503 + input_address,
2504 + output_address,
2505 + 1,
2506 + Duration::from_secs(2),
2507 + 0,
2508 + Some(CoordinatorMode::External {
2509 + address: coordinator_address,
2510 + }),
2511 + ))
2512 + .await
2513 + .unwrap()
2514 + }
2515 + });
2516 +
2517 + let (mut dealer, mut push) =
2518 + setup_bootstrapped_mock_engine(input_address, output_address, &[0x00, 0x00]).await;
2519 + let client = client_task.await.unwrap();
2520 + recv_xpub_subscription(&mut stats_socket).await;
2521 +
2522 + send_external_coordinator_publish(&mut stats_socket, &(Value::Nil, 5_u32, true)).await;
2523 + tokio::time::sleep(Duration::from_millis(50)).await;
2524 +
2525 + let mut stream = client.call(sample_request()).await.unwrap();
2526 +
2527 + assert!(
2528 + timeout(
2529 + Duration::from_millis(200),
2530 + recv_external_coordinator_wakeup(&mut stats_socket)
2531 + )
2532 + .await
2533 + .is_err()
2534 + );
2535 +
2536 + let add = recv_engine_message(&mut dealer).await;
2537 + let request: EngineCoreRequest = rmp_serde::from_slice(&add[1]).unwrap();
2538 + assert_eq!(request.current_wave, 5);
2539 +
2540 + send_outputs(
2541 + &mut push,
2542 + EngineCoreOutputs {
2543 + engine_index: 0,
2544 + outputs: vec![request_output(
2545 + "req-1",
2546 + vec![],
2547 + Some(EngineCoreFinishReason::Length),
2548 + )],
2549 + finished_requests: Some(BTreeSet::from(["req-1".to_string()])),
2550 + ..Default::default()
2551 + },
2552 + )
2553 + .await;
2554 +
2555 + let final_output = timeout(Duration::from_secs(1), stream.next()).await.unwrap();
2556 + assert!(final_output.is_some());
2557 +
2558 + client.shutdown().await.unwrap();
2559 +}
rust/src/engine-core-client/src/tests/mod.rs new
+1
@@ -0,0 +1 @@
1 +mod client;
rust/src/engine-core-client/src/tests/python_compat.py new
+356
@@ -0,0 +1,356 @@
1 +#!/usr/bin/env -S uv run
2 +# SPDX-License-Identifier: Apache-2.0
3 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
4 +# /// script
5 +# requires-python = ">=3.11"
6 +# dependencies = [
7 +# "msgspec>=0.19,<1",
8 +# "msgpack>=1,<2",
9 +# "numpy>=2,<3",
10 +# ]
11 +# ///
12 +
13 +from enum import Enum, IntEnum
14 +
15 +import msgpack
16 +import msgspec
17 +import numpy as np
18 +
19 +
20 +class RequestOutputKind(Enum):
21 + DELTA = 1
22 + FINAL_ONLY = 2
23 +
24 +
25 +class FinishReason(IntEnum):
26 + STOP = 0
27 + LENGTH = 1
28 + ABORT = 2
29 + ERROR = 3
30 + REPETITION = 4
31 +
32 +
33 +class EngineCoreSamplingParams(msgspec.Struct, dict=True):
34 + temperature: float = 1.0
35 + top_p: float = 1.0
36 + top_k: int = 0
37 + seed: int | None = None
38 + max_tokens: int = 65536
39 + min_tokens: int = 0
40 + min_p: float = 0.0
41 + frequency_penalty: float = 0.0
42 + presence_penalty: float = 0.0
43 + repetition_penalty: float = 1.0
44 + stop_token_ids: list[int] = []
45 + _eos_token_id: int | None = None
46 + _all_stop_token_ids: set[int] = set()
47 + output_kind: RequestOutputKind = RequestOutputKind.DELTA
48 +
49 +
50 +class EngineCoreRequest(
51 + msgspec.Struct,
52 + array_like=True,
53 + omit_defaults=True,
54 +):
55 + request_id: str
56 + prompt_token_ids: list[int] | None
57 + mm_features: object | None
58 + sampling_params: EngineCoreSamplingParams | None
59 + pooling_params: object | None
60 + arrival_time: float
61 + lora_request: object | None = None
62 + cache_salt: str | None = None
63 + data_parallel_rank: int | None = None
64 + prompt_embeds: object | None = None
65 + prompt_is_token_ids: list[bool] | None = None
66 + client_index: int = 0
67 + current_wave: int = 0
68 + priority: int = 0
69 + trace_headers: dict[str, str] | None = None
70 + resumable: bool = False
71 + external_req_id: str | None = None
72 + reasoning_ended: bool | None = None
73 + reasoning_parser_kwargs: dict[str, object] | None = None
74 + abort_immediately: bool = False
75 +
76 +
77 +class EngineCoreOutput(
78 + msgspec.Struct,
79 + array_like=True,
80 + omit_defaults=True,
81 +):
82 + request_id: str
83 + new_token_ids: list[int]
84 + new_logprobs: object | None = None
85 + new_prompt_logprobs_tensors: object | None = None
86 + pooling_output: object | None = None
87 + finish_reason: FinishReason | None = None
88 + stop_reason: int | str | None = None
89 + events: object | None = None
90 + kv_transfer_params: object | None = None
91 + trace_headers: object | None = None
92 + prefill_stats: object | None = None
93 + routed_experts: object | None = None
94 + num_nans_in_logits: int = 0
95 +
96 +
97 +class EngineCoreOutputs(
98 + msgspec.Struct,
99 + array_like=True,
100 + omit_defaults=True,
101 +):
102 + engine_index: int = 0
103 + outputs: list[EngineCoreOutput] = []
104 + scheduler_stats: object | None = None
105 + timestamp: float = 0.0
106 + utility_output: object | None = None
107 + finished_requests: set[str] | None = None
108 + wave_complete: int | None = None
109 + start_wave: int | None = None
110 +
111 +
112 +request = EngineCoreRequest(
113 + request_id="req-1",
114 + prompt_token_ids=[11, 22],
115 + mm_features=None,
116 + sampling_params=EngineCoreSamplingParams(
117 + temperature=0.8,
118 + top_p=0.9,
119 + top_k=8,
120 + seed=None,
121 + max_tokens=32,
122 + min_tokens=1,
123 + min_p=0.0,
124 + frequency_penalty=0.0,
125 + presence_penalty=0.0,
126 + repetition_penalty=1.0,
127 + stop_token_ids=[151643],
128 + _eos_token_id=151645,
129 + _all_stop_token_ids={151643, 151645},
130 + output_kind=RequestOutputKind.FINAL_ONLY,
131 + ),
132 + pooling_params=None,
133 + arrival_time=42.5,
134 + client_index=0,
135 +)
136 +
137 +multimodal_tensor = np.array([[1.0, 2.0], [3.5, 4.25]], dtype=np.float32)
138 +multimodal_features = [
139 + {
140 + "data": {
141 + "pixel_values": {
142 + "data": [
143 + "float32",
144 + [2, 2],
145 + msgpack.ExtType(3, multimodal_tensor.tobytes()),
146 + ],
147 + "field": [
148 + "flat",
149 + {
150 + "slices": [[0, 2, None]],
151 + "dim": 0,
152 + "keep_on_cpu": False,
153 + },
154 + ],
155 + }
156 + },
157 + "modality": "image",
158 + "identifier": "mm-cache-key",
159 + "mm_position": {
160 + "offset": 1,
161 + "length": 2,
162 + "is_embed": None,
163 + },
164 + "mm_hash": "processor-hash",
165 + }
166 +]
167 +multimodal_request_wire = [
168 + "req-mm",
169 + [101, 102, 103, 104],
170 + multimodal_features,
171 + None,
172 + None,
173 + 43.5,
174 +]
175 +
176 +outputs = EngineCoreOutputs(
177 + outputs=[
178 + EngineCoreOutput(
179 + request_id="req-1",
180 + new_token_ids=[7, 8],
181 + finish_reason=FinishReason.LENGTH,
182 + )
183 + ],
184 + finished_requests={"req-1"},
185 +)
186 +
187 +
188 +def encode_ndarray(
189 + array: np.ndarray,
190 + buffers: list[bytes],
191 + *,
192 + size_threshold: int = 256,
193 +):
194 + arr_data = array.data if array.flags.c_contiguous else array.tobytes()
195 + if not array.shape or array.nbytes < size_threshold:
196 + data = msgpack.ExtType(3, bytes(arr_data))
197 + else:
198 + data = len(buffers)
199 + buffers.append(bytes(arr_data))
200 + return [array.dtype.str, list(array.shape), data]
201 +
202 +
203 +def encode_tensor_like(
204 + dtype: str,
205 + shape: list[int],
206 + payload: bytes,
207 + buffers: list[bytes],
208 + *,
209 + size_threshold: int = 256,
210 +):
211 + if len(payload) < size_threshold:
212 + data = msgpack.ExtType(3, payload)
213 + else:
214 + data = len(buffers)
215 + buffers.append(payload)
216 + return [dtype, shape, data]
217 +
218 +
219 +def encode_output_frames(obj, *, size_threshold: int = 256) -> list[bytes]:
220 + buffers = [b""]
221 +
222 + def transform(value):
223 + if isinstance(value, np.ndarray):
224 + return encode_ndarray(value, buffers, size_threshold=size_threshold)
225 + if (
226 + isinstance(value, tuple)
227 + and len(value) == 3
228 + and value[0] in ("int32", "int64", "float32")
229 + ):
230 + dtype, shape, payload = value
231 + return encode_tensor_like(
232 + dtype,
233 + shape,
234 + payload,
235 + buffers,
236 + size_threshold=size_threshold,
237 + )
238 + if type(value) is list:
239 + return [transform(v) for v in value]
240 + if type(value) is tuple:
241 + return [transform(v) for v in value]
242 + if type(value) is dict:
243 + return {k: transform(v) for k, v in value.items()}
244 + return value
245 +
246 + buffers[0] = msgpack.packb(transform(obj), use_bin_type=True)
247 + return buffers
248 +
249 +
250 +def engine_output_wire(
251 + request_id: str,
252 + *,
253 + new_logprobs=None,
254 + new_prompt_logprobs_tensors=None,
255 +):
256 + return [
257 + request_id,
258 + [7, 8],
259 + new_logprobs,
260 + new_prompt_logprobs_tensors,
261 + None,
262 + int(FinishReason.LENGTH),
263 + ]
264 +
265 +
266 +def engine_outputs_wire(output):
267 + return [0, [output], None, 0.0, None, ["req-1"]]
268 +
269 +
270 +inline_logprobs = engine_outputs_wire(
271 + engine_output_wire(
272 + "req-1",
273 + new_logprobs=(
274 + np.array([[1, 2, 3], [4, 5, 6]], dtype=np.int64),
275 + np.array([[1, 2, 3], [4, 5, 6]], dtype=np.float32),
276 + np.array([1, 2], dtype=np.int64),
277 + None,
278 + ),
279 + )
280 +)
281 +
282 +multipart_logprobs = engine_outputs_wire(
283 + engine_output_wire(
284 + "req-1",
285 + new_logprobs=(
286 + np.array([[1, 2, 3], [4, 5, 6]], dtype=np.int64),
287 + np.array([[1, 2, 3], [4, 5, 6]], dtype=np.float32),
288 + np.array([1, 2], dtype=np.int64),
289 + None,
290 + ),
291 + )
292 +)
293 +
294 +inline_prompt_logprobs = engine_outputs_wire(
295 + engine_output_wire(
296 + "req-1",
297 + new_prompt_logprobs_tensors=(
298 + (
299 + "int64",
300 + [2, 3],
301 + np.array([[10, 11, 12], [13, 14, 15]], dtype=np.int64).tobytes(),
302 + ),
303 + (
304 + "float32",
305 + [2, 3],
306 + np.array(
307 + [[10, 11, 12], [13, 14, 15]],
308 + dtype=np.float32,
309 + ).tobytes(),
310 + ),
311 + ("int64", [2], np.array([3, 4], dtype=np.int64).tobytes()),
312 + None,
313 + ),
314 + )
315 +)
316 +
317 +multipart_prompt_logprobs = engine_outputs_wire(
318 + engine_output_wire(
319 + "req-1",
320 + new_prompt_logprobs_tensors=(
321 + (
322 + "int64",
323 + [2, 3],
324 + np.array([[10, 11, 12], [13, 14, 15]], dtype=np.int64).tobytes(),
325 + ),
326 + (
327 + "float32",
328 + [2, 3],
329 + np.array(
330 + [[10, 11, 12], [13, 14, 15]],
331 + dtype=np.float32,
332 + ).tobytes(),
333 + ),
334 + ("int64", [2], np.array([3, 4], dtype=np.int64).tobytes()),
335 + None,
336 + ),
337 + )
338 +)
339 +
340 +print(msgspec.msgpack.encode(request).hex())
341 +print(msgpack.packb(multimodal_request_wire, use_bin_type=True).hex())
342 +print(msgspec.msgpack.encode(outputs).hex())
343 +print(" ".join(frame.hex() for frame in encode_output_frames(inline_logprobs)))
344 +print(
345 + " ".join(
346 + frame.hex()
347 + for frame in encode_output_frames(multipart_logprobs, size_threshold=1)
348 + )
349 +)
350 +print(" ".join(frame.hex() for frame in encode_output_frames(inline_prompt_logprobs)))
351 +print(
352 + " ".join(
353 + frame.hex()
354 + for frame in encode_output_frames(multipart_prompt_logprobs, size_threshold=1)
355 + )
356 +)
rust/src/engine-core-client/src/transport.rs new
+606
@@ -0,0 +1,606 @@
1 +use std::collections::{BTreeMap, BTreeSet};
2 +use std::fmt::Debug;
3 +use std::ops::Deref;
4 +use std::time::Duration;
5 +
6 +use bytes::Bytes;
7 +use enum_as_inner::EnumAsInner;
8 +use thiserror_ext::AsReport;
9 +use tokio::sync::mpsc;
10 +use tokio::time::timeout;
11 +use tracing::{debug, error, info, trace, warn};
12 +use zeromq::prelude::{Socket, SocketRecv, SocketSend};
13 +use zeromq::util::PeerIdentity;
14 +use zeromq::{PullSocket, RouterSendHalf, RouterSocket, ZmqError, ZmqMessage};
15 +
16 +use crate::coordinator::CoordinatorBootstrap;
17 +use crate::error::{Error, Result, bail_unexpected_handshake_message};
18 +use crate::protocol::handshake::{
19 + EngineCoreReadyResponse, HandshakeAddresses, HandshakeInitMessage, ReadyMessage,
20 +};
21 +use crate::protocol::{
22 + EngineCoreOutputs, decode_engine_core_outputs, decode_msgpack, encode_msgpack,
23 +};
24 +
25 +/// Dedicated single-frame sentinel emitted by Python `EngineCoreProc` when the
26 +/// engine dies.
27 +pub const ENGINE_CORE_DEAD_SENTINEL: &[u8] = b"ENGINE_CORE_DEAD";
28 +
29 +/// Opaque routing identity of one engine on the frontend transport.
30 +#[derive(Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
31 +pub struct EngineId(Bytes);
32 +
33 +impl Debug for EngineId {
34 + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
35 + // Display the engine id as a hex string for easier debugging.
36 + write!(f, "EngineId({})", hex::encode(&self.0))
37 + }
38 +}
39 +
40 +impl EngineId {
41 + /// Convert the engine id into a ZMQ frame for sending.
42 + pub fn to_frame(&self) -> Bytes {
43 + self.0.clone()
44 + }
45 +
46 + /// Convert the engine id into a ZMQ frame for sending.
47 + pub fn into_frame(self) -> Bytes {
48 + self.0
49 + }
50 +
51 + /// Parse the Python-compatible engine index encoded in the routing
52 + /// identity.
53 + ///
54 + /// Python `EngineCoreProc` currently uses a two-byte little-endian engine
55 + /// index as its ROUTER/DEALER identity. Coordinator control messages
56 + /// such as `START_DP_WAVE(exclude_engine_index)` need that engine-side
57 + /// index rather than any frontend-local ordering.
58 + pub fn engine_index(&self) -> Option<u32> {
59 + if self.len() != 2 {
60 + return None;
61 + }
62 + Some(u16::from_le_bytes([self[0], self[1]]) as u32)
63 + }
64 +
65 + /// Construct an engine id from the Python-compatible engine index encoding
66 + /// (two-byte little-endian).
67 + pub fn from_engine_index(value: u32) -> Self {
68 + Self(Bytes::copy_from_slice(&(value as u16).to_le_bytes()))
69 + }
70 +}
71 +
72 +impl Deref for EngineId {
73 + type Target = [u8];
74 +
75 + fn deref(&self) -> &Self::Target {
76 + self.0.as_ref()
77 + }
78 +}
79 +
80 +impl From<Vec<u8>> for EngineId {
81 + fn from(value: Vec<u8>) -> Self {
82 + Self(Bytes::from(value))
83 + }
84 +}
85 +
86 +impl<const N: usize> From<&[u8; N]> for EngineId {
87 + fn from(value: &[u8; N]) -> Self {
88 + Self(Bytes::copy_from_slice(value))
89 + }
90 +}
91 +
92 +impl TryFrom<EngineId> for PeerIdentity {
93 + type Error = ZmqError;
94 +
95 + fn try_from(value: EngineId) -> std::result::Result<Self, Self::Error> {
96 + PeerIdentity::try_from(value.into_frame())
97 + }
98 +}
99 +
100 +/// Per-engine handshake result collected while bootstrapping one shared
101 +/// transport.
102 +#[derive(Clone, Debug)]
103 +pub struct ConnectedEngine {
104 + /// The identity of the connected engine.
105 + pub engine_id: EngineId,
106 + /// Post-initialization configuration received from the engine on the input
107 + /// socket registration message. `None` until the registration is received.
108 + pub ready_response: Option<EngineCoreReadyResponse>,
109 +}
110 +
111 +/// Represents the connected shared transport plus all registered engines after
112 +/// a successful multi-engine startup handshake.
113 +pub struct ConnectedTransport {
114 + /// The local address of the shared input socket that all engines connect to
115 + /// for receiving requests.
116 + pub input_address: String,
117 + /// The local address of the shared output socket that all engines connect
118 + /// to for sending responses.
119 + pub output_address: String,
120 + /// All engines connected through the startup handshake.
121 + pub engines: Vec<ConnectedEngine>,
122 + /// Optional engine-facing coordinator transport used for in-process wave
123 + /// coordination.
124 + pub coordinator: Option<CoordinatorBootstrap>,
125 +
126 + /// The sending half of the shared input socket.
127 + pub input_send: RouterSendHalf,
128 + /// The shared output socket for receiving responses from all engines.
129 + pub output_socket: PullSocket,
130 +}
131 +
132 +#[derive(Clone, Debug, EnumAsInner)]
133 +enum EngineStartupState {
134 + HelloReceived,
135 + ReadyReceived,
136 +}
137 +
138 +/// Connect to one or more engines through the startup handshake protocol,
139 +/// returning the shared data-plane transport plus the registered engines.
140 +pub async fn connect_handshake(
141 + handshake_address: &str,
142 + engine_count: usize,
143 + local_host: &str,
144 + local_input_address: Option<&str>,
145 + local_output_address: Option<&str>,
146 + enable_inproc_coordinator: bool,
147 + ready_timeout: Duration,
148 +) -> Result<ConnectedTransport> {
149 + if engine_count == 0 {
150 + bail_unexpected_handshake_message!("expected engine_count >= 1");
151 + }
152 +
153 + info!(
154 + engine_count,
155 + handshake_address, "waiting for engines to connect"
156 + );
157 +
158 + // 1. Bind shared local input/output sockets first so every engine receives the same data-plane
159 + // addresses during handshake.
160 + debug!(
161 + local_host,
162 + ?ready_timeout,
163 + engine_count,
164 + "binding shared transport sockets"
165 + );
166 + let (input_address, mut input_socket, output_address, output_socket) =
167 + bind_local_sockets(local_host, local_input_address, local_output_address).await?;
168 + info!(%input_address, %output_address, "bound local transport sockets");
169 +
170 + let mut coordinator = if enable_inproc_coordinator {
171 + Some(CoordinatorBootstrap::bind(local_host).await?)
172 + } else {
173 + None
174 + };
175 +
176 + // 2. Bind the shared handshake socket once. All engines connect to this socket with their own
177 + // identities, and startup order does not matter.
178 + let mut handshake_socket = RouterSocket::new();
179 + handshake_socket.bind(handshake_address).await?;
180 +
181 + let mut engines = BTreeMap::new();
182 +
183 + // 3. Receive HELLO from every engine and send a matching INIT. When coordinator mode is
184 + // enabled, the engines will not emit READY until the coordinator barrier below completes.
185 + while engines.len() < engine_count {
186 + debug!(
187 + handshake_address,
188 + connected = engines.len(),
189 + waiting_for = engine_count,
190 + "waiting for engine HELLO"
191 + );
192 + let message = timeout(ready_timeout, handshake_socket.recv()).await.map_err(|_| {
193 + Error::HandshakeTimeout {
194 + stage: "HELLO",
195 + timeout: ready_timeout,
196 + }
197 + })??;
198 + let (engine_id, handshake_message) = decode_handshake_message(message, None)?;
199 + match handshake_message.status.as_deref() {
200 + Some("HELLO") => {
201 + if engines.contains_key(&engine_id) {
202 + bail_unexpected_handshake_message!(
203 + "duplicate engine id {engine_id:?} observed during startup handshake"
204 + );
205 + }
206 + debug!(handshake_address, ?engine_id, "received HELLO from engine");
207 +
208 + send_init_message(
209 + &mut handshake_socket,
210 + &engine_id,
211 + &input_address,
212 + &output_address,
213 + coordinator.as_ref(),
214 + )
215 + .await?;
216 + debug!(handshake_address, ?engine_id, "sent INIT to engine");
217 +
218 + engines.insert(engine_id.clone(), EngineStartupState::HelloReceived);
219 + }
220 + Some("READY") => {
221 + if coordinator.is_some() {
222 + bail_unexpected_handshake_message!(
223 + "received READY for engine id {engine_id:?} before coordinator startup gate completed"
224 + );
225 + }
226 + let state = match engines.get_mut(&engine_id) {
227 + Some(state) if !state.is_ready_received() => state,
228 + _ => {
229 + bail_unexpected_handshake_message!(
230 + "received READY for unexpected or duplicate engine id {engine_id:?}"
231 + );
232 + }
233 + };
234 + debug!(
235 + handshake_address,
236 + ?engine_id,
237 + ?handshake_message,
238 + "received overlapping READY from engine during HELLO phase"
239 + );
240 + *state = EngineStartupState::ReadyReceived;
241 + }
242 + other => {
243 + bail_unexpected_handshake_message!("unexpected handshake status {other:?}");
244 + }
245 + }
246 + }
247 +
248 + // 4. Optional coordinator startup gate. Without coordinator there is nothing to do.
249 + if let Some(coordinator) = coordinator.as_mut() {
250 + coordinator.wait_for_startup_gate(engine_count, ready_timeout).await?;
251 + }
252 +
253 + // 5. After the optional gate has opened, every engine may now send READY.
254 + while engines.values().any(|state| !state.is_ready_received()) {
255 + debug!(
256 + handshake_address,
257 + connected = engines.len(),
258 + ready = engines.values().filter(|state| state.is_ready_received()).count(),
259 + waiting_for = engine_count,
260 + "waiting for engine READY"
261 + );
262 + let message = timeout(ready_timeout, handshake_socket.recv()).await.map_err(|_| {
263 + Error::HandshakeTimeout {
264 + stage: "READY",
265 + timeout: ready_timeout,
266 + }
267 + })??;
268 + let (engine_id, handshake_message) = decode_handshake_message(message, None)?;
269 + match handshake_message.status.as_deref() {
270 + Some("READY") => {
271 + let state = match engines.get_mut(&engine_id) {
272 + Some(state) if !state.is_ready_received() => state,
273 + _ => {
274 + bail_unexpected_handshake_message!(
275 + "received READY for unexpected or duplicate engine id {engine_id:?}"
276 + );
277 + }
278 + };
279 + debug!(
280 + handshake_address,
281 + ?engine_id,
282 + ?handshake_message,
283 + "received READY from engine"
284 + );
285 + *state = EngineStartupState::ReadyReceived;
286 + }
287 + Some("HELLO") => {
288 + bail_unexpected_handshake_message!(
289 + "received duplicate HELLO for engine id {engine_id:?} after INIT phase completed"
290 + );
291 + }
292 + other => {
293 + bail_unexpected_handshake_message!("unexpected handshake status {other:?}");
294 + }
295 + }
296 + }
297 +
298 + // 4. Wait for every engine to connect to the shared input socket and register itself. The
299 + // `ready_response` is a placeholder; it is populated for each engine by
300 + // `wait_for_input_registrations` below.
301 + let mut engines: Vec<_> = engines
302 + .into_keys()
303 + .map(|engine_id| ConnectedEngine {
304 + engine_id,
305 + ready_response: None,
306 + })
307 + .collect();
308 +
309 + wait_for_input_registrations(&mut input_socket, &mut engines, ready_timeout).await?;
310 + debug!(
311 + engine_count = engines.len(),
312 + "all engines registered on shared input socket"
313 + );
314 +
315 + info!(engine_count = engines.len(), "engines connected");
316 +
317 + let (input_send, _) = input_socket.split();
318 +
319 + Ok(ConnectedTransport {
320 + input_address,
321 + output_address,
322 + input_send,
323 + output_socket,
324 + engines,
325 + coordinator,
326 + })
327 +}
328 +
329 +/// Bind to Python-supplied frontend transport addresses and wait for
330 +/// already-initialized engines to register themselves on the input socket.
331 +///
332 +/// This path mirrors Python's externally managed `AsyncMPClient` bootstrap
333 +/// model: the addresses are already fixed by the supervisor, and engine
334 +/// identities are synthesized from contiguous rank order instead of being
335 +/// discovered through a Rust-owned handshake.
336 +pub async fn connect_bootstrapped(
337 + input_address: &str,
338 + output_address: &str,
339 + engine_count: usize,
340 + ready_timeout: Duration,
341 +) -> Result<ConnectedTransport> {
342 + if engine_count == 0 {
343 + bail_unexpected_handshake_message!("expected engine_count >= 1");
344 + }
345 +
346 + let mut input_socket = RouterSocket::new();
347 + let input_address = input_socket.bind(input_address).await?.to_string();
348 +
349 + let mut output_socket = PullSocket::new();
350 + let output_address = output_socket.bind(output_address).await?.to_string();
351 +
352 + // TODO: follow start rank
353 + let mut engines = (0..engine_count)
354 + .map(|index| ConnectedEngine {
355 + engine_id: EngineId::from((index as u16).to_le_bytes().to_vec()),
356 + ready_response: None,
357 + })
358 + .collect::<Vec<_>>();
359 +
360 + wait_for_input_registrations(&mut input_socket, &mut engines, ready_timeout).await?;
361 + info!(
362 + engine_count = engines.len(),
363 + "bootstrapped engines connected"
364 + );
365 +
366 + let (input_send, _) = input_socket.split();
367 +
368 + Ok(ConnectedTransport {
369 + input_address,
370 + output_address,
371 + engines,
372 + coordinator: None,
373 + input_send,
374 + output_socket,
375 + })
376 +}
377 +
378 +/// Bind new input and output sockets.
379 +async fn bind_local_sockets(
380 + local_host: &str,
381 + local_input_address: Option<&str>,
382 + local_output_address: Option<&str>,
383 +) -> Result<(String, RouterSocket, String, PullSocket)> {
384 + let mut input_socket = RouterSocket::new();
385 + let input_bind_address = local_input_address
386 + .map(str::to_owned)
387 + .unwrap_or_else(|| format!("tcp://{local_host}:0"));
388 + let input_address = input_socket.bind(&input_bind_address).await?.to_string();
389 +
390 + let mut output_socket = PullSocket::new();
391 + let output_bind_address = local_output_address
392 + .map(str::to_owned)
393 + .unwrap_or_else(|| format!("tcp://{local_host}:0"));
394 + let output_address = output_socket.bind(&output_bind_address).await?.to_string();
395 +
396 + Ok((input_address, input_socket, output_address, output_socket))
397 +}
398 +
399 +/// Decode a handshake message and validate its structure and identity.
400 +fn decode_handshake_message(
401 + message: ZmqMessage,
402 + expected_id: Option<&EngineId>,
403 +) -> Result<(EngineId, ReadyMessage)> {
404 + if message.len() != 2 {
405 + bail_unexpected_handshake_message!("expected 2 frames, got {}", message.len());
406 + }
407 +
408 + let frames = message.into_vec();
409 + let actual_id = EngineId(frames[0].clone());
410 + if let Some(expected_id) = expected_id
411 + && actual_id != *expected_id
412 + {
413 + return Err(Error::UnexpectedHandshakeIdentity {
414 + expected: expected_id.to_vec(),
415 + actual: actual_id.to_vec(),
416 + });
417 + }
418 +
419 + let handshake_message: ReadyMessage = decode_msgpack(&frames[1])?;
420 + Ok((actual_id, handshake_message))
421 +}
422 +
423 +/// Send an INIT message to the engine with the local socket addresses for the
424 +/// engine to connect to, using the handshake socket.
425 +async fn send_init_message(
426 + handshake_socket: &mut RouterSocket,
427 + engine_id: &EngineId,
428 + input_address: &str,
429 + output_address: &str,
430 + coordinator: Option<&CoordinatorBootstrap>,
431 +) -> Result<()> {
432 + let init_message = HandshakeInitMessage {
433 + addresses: HandshakeAddresses {
434 + inputs: vec![input_address.to_string()],
435 + outputs: vec![output_address.to_string()],
436 + coordinator_input: coordinator.map(|c| c.input_address.clone()),
437 + coordinator_output: coordinator.map(|c| c.output_address.clone()),
438 + frontend_stats_publish_address: None,
439 + },
440 + parallel_config: Default::default(),
441 + };
442 + let payload = encode_msgpack(&init_message)?;
443 + let message = ZmqMessage::try_from(vec![engine_id.to_frame(), Bytes::from(payload)])
444 + .expect("handshake router messages must contain identity and payload");
445 + handshake_socket.send(message).await?;
446 + Ok(())
447 +}
448 +
449 +/// Receive the input registration message from each engine and validate its
450 +/// identity.
451 +///
452 +/// Each registration contains 2 frames: `[identity, ready-payload]`.
453 +///
454 +/// Since vLLM commit `c8d98f81f676552c263f35bbde55e6edbe81b4e8` ("[Core]
455 +/// Simplify API server handshake"), the payload is a msgpack-encoded
456 +/// [`EngineCoreReadyResponse`] carrying post-initialization values such as
457 +/// `max_model_len`.
458 +///
459 +/// Older engines sent an empty second frame here just to establish the
460 +/// ROUTER/DEALER backchannel, with no structured payload on the input socket.
461 +/// We continue to tolerate that legacy shape so the frontend can still connect
462 +/// to slightly older local engine checkouts.
463 +async fn wait_for_input_registrations(
464 + input_socket: &mut RouterSocket,
465 + engines: &mut [ConnectedEngine],
466 + ready_timeout: Duration,
467 +) -> Result<()> {
468 + let mut pending = engines.iter().map(|e| e.engine_id.clone()).collect::<BTreeSet<_>>();
469 +
470 + while !pending.is_empty() {
471 + let registration = timeout(ready_timeout, input_socket.recv()).await.map_err(|_| {
472 + Error::InputRegistrationTimeout {
473 + timeout: ready_timeout,
474 + }
475 + })??;
476 +
477 + if registration.len() != 2 {
478 + bail_unexpected_handshake_message!(
479 + "expected 2 frames for engine input registration, got {}",
480 + registration.len()
481 + );
482 + }
483 +
484 + let frames = registration.into_vec();
485 + let actual_id = EngineId(frames[0].clone());
486 + if !pending.remove(&actual_id) {
487 + bail_unexpected_handshake_message!(
488 + "received input registration for unexpected engine id {actual_id:?}"
489 + );
490 + }
491 +
492 + let ready_response = if frames[1].is_empty() {
493 + debug!(
494 + ?actual_id,
495 + "received legacy empty input registration from engine"
496 + );
497 + None
498 + } else {
499 + let ready_response: EngineCoreReadyResponse = decode_msgpack(&frames[1])?;
500 + debug!(
501 + ?actual_id,
502 + ?ready_response,
503 + "received input registration from engine"
504 + );
505 + Some(ready_response)
506 + };
507 +
508 + // Store the ready response in the corresponding engine entry.
509 + if let Some(engine) = engines.iter_mut().find(|e| e.engine_id == actual_id) {
510 + engine.ready_response = ready_response;
511 + }
512 + }
513 +
514 + Ok(())
515 +}
516 +
517 +/// Send an encoded message to the engine through the input socket.
518 +pub async fn send_message(
519 + input_send: &mut RouterSendHalf,
520 + engine_id: &EngineId,
521 + request_type: Bytes,
522 + payload: Vec<u8>,
523 +) -> Result<()> {
524 + let message = ZmqMessage::try_from(vec![
525 + engine_id.to_frame(),
526 + request_type,
527 + Bytes::from(payload),
528 + ])
529 + .expect("router messages must contain identity and payload");
530 +
531 + trace!(
532 + ?engine_id,
533 + frame_count = message.len(),
534 + "sending ZMQ message"
535 + );
536 + input_send.send(message).await?;
537 + Ok(())
538 +}
539 +
540 +/// Run the output loop to receive messages from the engine and send them to the
541 +/// provided channel.
542 +pub async fn run_output_loop(
543 + mut output_socket: PullSocket,
544 + tx: mpsc::Sender<Result<EngineCoreOutputs>>,
545 +) {
546 + loop {
547 + let message = match output_socket.recv().await {
548 + Ok(message) => message,
549 + Err(error) => {
550 + // If we fail to receive a message from the engine, it's likely that the engine
551 + // has crashed or become unreachable, so we should notify the
552 + // client and shut down the output loop.
553 + error!(error = %error.as_report(), "failed to receive output message");
554 + let _ = tx.send(Err(Error::Transport(error))).await;
555 + return;
556 + }
557 + };
558 +
559 + let frame_count = message.len();
560 + trace!(frame_count, "received output message");
561 + let frames = message.into_vec();
562 + let frame = frames.first().expect("output message must have at least one frame");
563 + let frame_len = frame.len();
564 + if frame.as_ref() == ENGINE_CORE_DEAD_SENTINEL {
565 + warn!("received ENGINE_CORE_DEAD sentinel from engine");
566 + let _ = tx.send(Err(Error::EngineCoreDead)).await;
567 + return;
568 + }
569 + let decoded = match decode_engine_core_outputs(&frames) {
570 + Ok(decoded) => {
571 + trace!(frame_len, outputs = ?decoded, "decoded output message");
572 + Ok(decoded)
573 + }
574 + Err(error) => {
575 + // If we fail to decode the message from the engine, notify the client but keep
576 + // the output loop running to continue processing future
577 + // messages from the engine.
578 + warn!(frame_len, error = %error.as_report(), "failed to decode output message");
579 + Err(error)
580 + }
581 + };
582 +
583 + if tx.send(decoded).await.is_err() {
584 + // If we fail to send the decoded message to the client, it's likely that the
585 + // client has shut down, so we should shut down the output loop as
586 + // well.
587 + warn!("output loop rx dropped, shutting down output loop");
588 + return;
589 + }
590 + }
591 +}
592 +
593 +#[cfg(test)]
594 +mod tests {
595 + use super::bind_local_sockets;
596 +
597 + #[tokio::test]
598 + async fn bind_local_sockets_resolves_zero_port_bindings() {
599 + let (input_address, _input_socket, output_address, _output_socket) =
600 + bind_local_sockets("127.0.0.1", None, None).await.expect("bind local sockets");
601 +
602 + assert!(input_address.starts_with("tcp://127.0.0.1:"));
603 + assert!(output_address.starts_with("tcp://127.0.0.1:"));
604 + assert_ne!(input_address, output_address);
605 + }
606 +}
rust/src/llm/Cargo.toml new
+38
@@ -0,0 +1,38 @@
1 +[package]
2 +name = "vllm-llm"
3 +version.workspace = true
4 +edition.workspace = true
5 +license.workspace = true
6 +
7 +[features]
8 +test-util = []
9 +
10 +[dependencies]
11 +easy-ext.workspace = true
12 +enum-as-inner.workspace = true
13 +futures.workspace = true
14 +serde.workspace = true
15 +serde_json.workspace = true
16 +thiserror.workspace = true
17 +tokio.workspace = true
18 +tokio-util.workspace = true
19 +tracing.workspace = true
20 +uuid.workspace = true
21 +vllm-engine-core-client.workspace = true
22 +vllm-metrics.workspace = true
23 +
24 +[dev-dependencies]
25 +anyhow.workspace = true
26 +bytes.workspace = true
27 +clap.workspace = true
28 +expect-test.workspace = true
29 +rmp-serde.workspace = true
30 +tokio.workspace = true
31 +tracing-subscriber.workspace = true
32 +uuid.workspace = true
33 +vllm-engine-core-client = { workspace = true, features = ["test-util"] }
34 +vllm-metrics.workspace = true
35 +zeromq.workspace = true
36 +
37 +[lints]
38 +workspace = true
rust/src/llm/examples/README.md new
+29
@@ -0,0 +1,29 @@
1 +# LLM Smoke Test
2 +
3 +Start headless `vllm`:
4 +
5 +```bash
6 +source ../vllm/.venv/bin/activate
7 +HF_HUB_OFFLINE=1 \
8 +VLLM_LOGGING_LEVEL=DEBUG \
9 +VLLM_CPU_KVCACHE_SPACE=2 \
10 +VLLM_HOST_IP=127.0.0.1 \
11 +VLLM_LOOPBACK_IP=127.0.0.1 \
12 +python3 -m vllm.entrypoints.cli.main serve Qwen/Qwen3-0.6B \
13 + --headless \
14 + --data-parallel-address 127.0.0.1 \
15 + --data-parallel-rpc-port 62100 \
16 + --data-parallel-size-local 1 \
17 + --max-model-len 512 \
18 + --dtype float16
19 +```
20 +
21 +Run the Rust smoke test through the `vllm-llm` generate interface:
22 +
23 +```bash
24 +cargo run -p vllm-llm --example external_engine_smoke -- \
25 + --handshake-address tcp://127.0.0.1:62100 \
26 + --host 127.0.0.1
27 +```
28 +
29 +IMPORTANT: You must restart `vllm` each time you run the smoke test, as the vLLM engine cannot manage frontend closures and subsequent reconnects. In other words, do not reuse existing `vllm` instances, if any.
rust/src/llm/examples/external_engine_smoke.rs new
+144
@@ -0,0 +1,144 @@
1 +use std::time::Duration;
2 +
3 +use anyhow::{Context, Result, bail};
4 +use clap::Parser;
5 +use futures::StreamExt as _;
6 +use tokio::time::timeout;
7 +use tracing_subscriber::EnvFilter;
8 +use vllm_engine_core_client::protocol::EngineCoreSamplingParams;
9 +use vllm_engine_core_client::{EngineCoreClient, EngineCoreClientConfig, TransportMode};
10 +use vllm_llm::{FinishReason, GenerateOutputStream, GenerateRequest, Llm};
11 +
12 +const PROMPT_TOKEN_IDS: &[u32] = &[20841, 448, 6896, 25, 23811];
13 +
14 +#[derive(Debug, Parser)]
15 +#[command(about = "Smoke-test the Rust LLM facade against an external vLLM engine.")]
16 +struct Args {
17 + #[arg(long)]
18 + handshake_address: String,
19 + #[arg(long, default_value_t = 1)]
20 + engine_count: usize,
21 + #[arg(long, default_value = "Qwen/Qwen3-0.6B")]
22 + model: String,
23 + #[arg(long, default_value = "127.0.0.1")]
24 + host: String,
25 + #[arg(long, default_value_t = 0)]
26 + client_index: u32,
27 + #[arg(long, default_value_t = 30)]
28 + ready_timeout_secs: u64,
29 + #[arg(long, default_value_t = 120)]
30 + output_timeout_secs: u64,
31 + #[arg(long, default_value_t = 5)]
32 + max_tokens: u32,
33 +}
34 +
35 +fn unique_request_id() -> String {
36 + format!("rust-llm-smoke-{}", uuid::Uuid::new_v4())
37 +}
38 +
39 +fn init_tracing() {
40 + let filter = EnvFilter::try_from_default_env()
41 + .unwrap_or_else(|_| EnvFilter::new("vllm_engine_core_client=debug"));
42 + let _ = tracing_subscriber::fmt().with_env_filter(filter).try_init();
43 +}
44 +
45 +fn build_request(request_id: String, max_tokens: u32) -> GenerateRequest {
46 + GenerateRequest {
47 + request_id,
48 + prompt_token_ids: PROMPT_TOKEN_IDS.to_vec(),
49 + sampling_params: EngineCoreSamplingParams {
50 + max_tokens,
51 + ..EngineCoreSamplingParams::for_test()
52 + },
53 + mm_features: None,
54 + arrival_time: None,
55 + cache_salt: None,
56 + trace_headers: None,
57 + priority: 0,
58 + data_parallel_rank: None,
59 + reasoning_ended: None,
60 + lora_request: None,
61 + }
62 +}
63 +
64 +#[derive(Debug)]
65 +struct CompletedRequest {
66 + token_ids: Vec<u32>,
67 + finish_reason: FinishReason,
68 +}
69 +
70 +async fn wait_for_request_completion(mut stream: GenerateOutputStream) -> Result<CompletedRequest> {
71 + let output = match stream.next().await {
72 + Some(output) => output.context("failed to receive request output")?,
73 + None => bail!("request stream ended without a final output"),
74 + };
75 +
76 + let none = stream.next().await;
77 + assert!(
78 + none.is_none(),
79 + "expected final-only stream to end after the final output"
80 + );
81 +
82 + let finish_reason = output.finish_reason.expect("final-only output must have a finish reason");
83 + let token_ids = output.token_ids;
84 +
85 + Ok(CompletedRequest {
86 + token_ids,
87 + finish_reason,
88 + })
89 +}
90 +
91 +async fn wait_for_timeout(
92 + stream: GenerateOutputStream,
93 + output_timeout: Duration,
94 +) -> Result<CompletedRequest> {
95 + timeout(output_timeout, wait_for_request_completion(stream))
96 + .await
97 + .context("timed out waiting for request output")?
98 +}
99 +
100 +#[tokio::main(flavor = "multi_thread")]
101 +async fn main() -> Result<()> {
102 + init_tracing();
103 + let args = Args::parse();
104 + let ready_timeout = Duration::from_secs(args.ready_timeout_secs);
105 + let output_timeout = Duration::from_secs(args.output_timeout_secs);
106 + let request_id = unique_request_id();
107 + let client = EngineCoreClient::connect(EngineCoreClientConfig {
108 + transport_mode: TransportMode::HandshakeOwner {
109 + handshake_address: args.handshake_address.clone(),
110 + advertised_host: args.host.clone(),
111 + engine_count: args.engine_count,
112 + ready_timeout,
113 + local_input_address: None,
114 + local_output_address: None,
115 + },
116 + coordinator_mode: None,
117 + model_name: args.model.clone(),
118 + client_index: args.client_index,
119 + })
120 + .await
121 + .context("failed to connect to external vLLM engine")?;
122 +
123 + println!("model={}", args.model);
124 + println!("handshake_address={}", args.handshake_address);
125 + println!("engine_count={}", args.engine_count);
126 + println!("input_address={}", client.input_address());
127 + println!("output_address={}", client.output_address());
128 + println!("engine_identities={:x?}", client.engine_identities());
129 +
130 + let llm = Llm::new(client);
131 + let request = build_request(request_id.clone(), args.max_tokens);
132 + println!("request_id={request_id}");
133 + println!("prompt_token_ids={PROMPT_TOKEN_IDS:?}");
134 +
135 + let stream = llm.generate(request).await.context("failed to submit generate request")?;
136 + let output = wait_for_timeout(stream, output_timeout).await?;
137 +
138 + llm.shutdown().await.context("failed to shut down llm client")?;
139 +
140 + println!("token_ids={:?}", output.token_ids);
141 + println!("finish_reason={:?}", output.finish_reason);
142 +
143 + Ok(())
144 +}
rust/src/llm/src/error.rs new
+12
@@ -0,0 +1,12 @@
1 +use thiserror::Error;
2 +
3 +pub type Result<T> = std::result::Result<T, Error>;
4 +
5 +/// Public error type for the Rust `llm` facade.
6 +#[derive(Debug, Error)]
7 +pub enum Error {
8 + #[error("generate request `{request_id}` has an empty prompt_token_ids")]
9 + EmptyPromptTokenIds { request_id: String },
10 + #[error("engine-core error")]
11 + EngineCoreClient(#[from] vllm_engine_core_client::Error),
12 +}
rust/src/llm/src/lib.rs new
+100
@@ -0,0 +1,100 @@
1 +use tracing::Span;
2 +use vllm_engine_core_client::EngineCoreClient;
3 +
4 +mod error;
5 +mod log_stats;
6 +mod output;
7 +mod request;
8 +mod request_metrics;
9 +
10 +pub use error::{Error, Result};
11 +pub use output::{
12 + CollectedGenerateOutput, FinishReason, GenerateOutput, GenerateOutputStream,
13 + GenerateOutputStreamExt, GeneratePromptInfo,
14 +};
15 +pub use request::GenerateRequest;
16 +pub use vllm_engine_core_client::protocol::logprobs::{Logprobs, PositionLogprobs, TokenLogprob};
17 +
18 +use crate::log_stats::StatsLogger;
19 +use crate::request_metrics::RequestMetricsTracker;
20 +
21 +/// Thin generate-only facade over [`EngineCoreClient`].
22 +///
23 +/// This mirrors the narrow public shape of Python `AsyncLLM.generate()` and
24 +/// `abort()`, but keeps the boundary close to raw engine-core requests and
25 +/// outputs.
26 +pub struct Llm {
27 + client: EngineCoreClient,
28 + randomize_request_id: bool,
29 + stats_logger: Option<StatsLogger>,
30 +}
31 +
32 +impl Llm {
33 + /// Create a new minimal LLM facade from an already connected engine-core
34 + /// client.
35 + pub fn new(client: EngineCoreClient) -> Self {
36 + Self {
37 + client,
38 + randomize_request_id: true,
39 + stats_logger: None,
40 + }
41 + }
42 +
43 + /// Enable or disable periodic stats logging.
44 + pub fn with_log_stats(mut self, enabled: bool) -> Self {
45 + if enabled {
46 + let stats_logger = StatsLogger::start(
47 + self.client.model_name().to_string(),
48 + self.client.engine_count(),
49 + );
50 + self.stats_logger = Some(stats_logger);
51 + } else {
52 + self.stats_logger = None;
53 + }
54 + self
55 + }
56 +
57 + /// Control whether external request ids are randomized before reaching
58 + /// engine-core.
59 + pub fn with_request_id_randomization(mut self, enabled: bool) -> Self {
60 + self.randomize_request_id = enabled;
61 + self
62 + }
63 +
64 + /// Expose the underlying engine-core client for low-level utility/admin
65 + /// calls.
66 + pub fn engine_core_client(&self) -> &EngineCoreClient {
67 + &self.client
68 + }
69 +
70 + /// Submit one tokenized generate request and return a per-request output
71 + /// stream.
72 + pub async fn generate(&self, req: GenerateRequest) -> Result<GenerateOutputStream> {
73 + let prepared = req.prepare(self.randomize_request_id)?;
74 + let prompt_token_ids = prepared.prompt_token_ids().into();
75 +
76 + // Record internal engine-core request ID in the current tracing span.
77 + Span::current().record("engine_request_id", &prepared.engine_request.request_id);
78 +
79 + let request_metrics = RequestMetricsTracker::new(
80 + self.client.model_name().to_string(),
81 + prepared.engine_request.arrival_time,
82 + prepared.prompt_token_ids().len() as u32,
83 + (prepared.engine_request.sampling_params.as_ref()).map(|p| p.max_tokens),
84 + 1,
85 + );
86 + let stream = self.client.call(prepared.engine_request).await?;
87 +
88 + Ok(GenerateOutputStream::new(
89 + prompt_token_ids,
90 + stream,
91 + request_metrics,
92 + ))
93 + }
94 +
95 + /// Shut down the underlying engine-core client and its background tasks.
96 + pub async fn shutdown(self) -> Result<()> {
97 + self.client.shutdown().await?;
98 + Ok(())
99 + }
100 +}
rust/src/llm/src/log_stats.rs new
+199
@@ -0,0 +1,199 @@
1 +use std::fmt::Write;
2 +use std::time::{Duration, Instant};
3 +
4 +use tokio_util::task::AbortOnDropHandle;
5 +use tracing::{debug, info};
6 +use vllm_metrics::{
7 + EngineLabels, F64Gauge, METRICS, PromptTokenSourceLabels, U64Counter, U64Gauge,
8 +};
9 +
10 +const LOG_STATS_INTERVAL: Duration = Duration::from_secs(10);
11 +
12 +/// Cached, cloned metric handles for one engine. Each clone shares the same
13 +/// underlying `Arc<Atomic*>` as the prometheus `Family` entry, so reads go
14 +/// straight to the atomic with no lock.
15 +struct EngineMetrics {
16 + // Counters for throughput deltas.
17 + prompt_tokens_computed: U64Counter,
18 + generation_tokens: U64Counter,
19 + prefix_cache_queries: U64Counter,
20 + prefix_cache_hits: U64Counter,
21 +
22 + // Gauges for instantaneous scheduler state.
23 + scheduler_running: U64Gauge,
24 + scheduler_waiting: U64Gauge,
25 + kv_cache_usage: F64Gauge,
26 +}
27 +
28 +/// Accumulated snapshot values from the last logging interval, used to compute
29 +/// deltas.
30 +struct CounterSnapshot {
31 + prompt_tokens: u64,
32 + generation_tokens: u64,
33 + prefix_cache_queries: u64,
34 + prefix_cache_hits: u64,
35 +}
36 +
37 +/// Periodic stats logger that mirrors Python vLLM's `LoggingStatLogger`.
38 +///
39 +/// Spawns a background task that logs throughput and scheduler state at a fixed
40 +/// interval. When idle (both current and previous throughputs are zero), logs
41 +/// at DEBUG level. When load drops to zero, emits one final INFO-level line
42 +/// before going quiet.
43 +pub(crate) struct StatsLogger {
44 + _task: AbortOnDropHandle<()>,
45 +}
46 +
47 +impl StatsLogger {
48 + /// Start the background stats logging task.
49 + pub(crate) fn start(model_name: String, engine_count: usize) -> Self {
50 + let task = AbortOnDropHandle::new(tokio::spawn(async move {
51 + run_stats_logger(model_name, engine_count).await;
52 + }));
53 + Self { _task: task }
54 + }
55 +}
56 +
57 +/// Resolve and clone all metric handles once so the hot path is lock-free.
58 +fn resolve_engine_metrics(model_name: &str, engine_count: usize) -> Vec<EngineMetrics> {
59 + let m = &METRICS;
60 + (0..engine_count as u32)
61 + .map(|engine| {
62 + let el = EngineLabels {
63 + model_name: model_name.to_string(),
64 + engine,
65 + };
66 + let pt = PromptTokenSourceLabels {
67 + model_name: model_name.to_string(),
68 + engine,
69 + source: "local_compute",
70 + };
71 + EngineMetrics {
72 + // Use "local_compute" source for prompt throughput (excludes
73 + // cached/transferred tokens), matching Python's
74 + // `iteration_stats.prompt_token_stats.computed`.
75 + prompt_tokens_computed: m.request.prompt_tokens_by_source.get_or_create_owned(&pt),
76 + generation_tokens: m.request.generation_tokens.get_or_create_owned(&el),
77 + prefix_cache_queries: m.scheduler.prefix_cache_queries.get_or_create_owned(&el),
78 + prefix_cache_hits: m.scheduler.prefix_cache_hits.get_or_create_owned(&el),
79 + scheduler_running: m.scheduler.scheduler_running.get_or_create_owned(&el),
80 + scheduler_waiting: m.scheduler.scheduler_waiting.get_or_create_owned(&el),
81 + kv_cache_usage: m.scheduler.kv_cache_usage.get_or_create_owned(&el),
82 + }
83 + })
84 + .collect()
85 +}
86 +
87 +async fn run_stats_logger(model_name: String, engine_count: usize) {
88 + let engines = resolve_engine_metrics(&model_name, engine_count);
89 +
90 + let mut interval = tokio::time::interval(LOG_STATS_INTERVAL);
91 + interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
92 + // The first tick fires immediately; skip it so the first log is after one full
93 + // interval.
94 + interval.tick().await;
95 +
96 + let mut prev = read_counters(&engines);
97 + let mut last_log_time = Instant::now();
98 + let mut last_prompt_throughput: f64 = 0.0;
99 + let mut last_generation_throughput: f64 = 0.0;
100 +
101 + let mut msg = String::new();
102 + loop {
103 + interval.tick().await;
104 +
105 + let now = Instant::now();
106 + let elapsed = now.duration_since(last_log_time).as_secs_f64();
107 + if elapsed <= 0.0 {
108 + continue;
109 + }
110 +
111 + let curr = read_counters(&engines);
112 +
113 + let prompt_throughput =
114 + curr.prompt_tokens.wrapping_sub(prev.prompt_tokens) as f64 / elapsed;
115 + let generation_throughput =
116 + curr.generation_tokens.wrapping_sub(prev.generation_tokens) as f64 / elapsed;
117 +
118 + // Idle = both current and previous throughputs are zero.
119 + let is_idle = prompt_throughput == 0.0
120 + && generation_throughput == 0.0
121 + && last_prompt_throughput == 0.0
122 + && last_generation_throughput == 0.0;
123 +
124 + // Read scheduler gauges (aggregate across engines).
125 + let (num_running, num_waiting, kv_cache_usage) = read_scheduler_gauges(&engines);
126 +
127 + // Compute prefix cache hit rate over this interval.
128 + let delta_queries = curr.prefix_cache_queries.wrapping_sub(prev.prefix_cache_queries);
129 + let prefix_cache_hit_rate = if delta_queries > 0 {
130 + let delta_hits = curr.prefix_cache_hits.wrapping_sub(prev.prefix_cache_hits);
131 + delta_hits as f64 / delta_queries as f64 * 100.0
132 + } else {
133 + 0.0
134 + };
135 +
136 + // Build the log line.
137 + msg.clear();
138 + write!(
139 + msg,
140 + "Avg prompt tput: {prompt_throughput:.1} toks/s, \
141 + Avg generation tput: {generation_throughput:.1} toks/s, \
142 + Reqs Running: {num_running}, \
143 + Waiting: {num_waiting}, \
144 + GPU KV cache used: {:.1}%, \
145 + Prefix cache hit rate: {prefix_cache_hit_rate:.1}%",
146 + kv_cache_usage * 100.0,
147 + )
148 + .unwrap();
149 +
150 + if is_idle {
151 + debug!("{msg}");
152 + } else {
153 + info!("{msg}");
154 + }
155 +
156 + last_prompt_throughput = prompt_throughput;
157 + last_generation_throughput = generation_throughput;
158 + last_log_time = now;
159 + prev = curr;
160 + }
161 +}
162 +
163 +/// Read the current cumulative counter values for throughput computation.
164 +fn read_counters(engines: &[EngineMetrics]) -> CounterSnapshot {
165 + let mut snap = CounterSnapshot {
166 + prompt_tokens: 0,
167 + generation_tokens: 0,
168 + prefix_cache_queries: 0,
169 + prefix_cache_hits: 0,
170 + };
171 + for e in engines {
172 + snap.prompt_tokens += e.prompt_tokens_computed.get();
173 + snap.generation_tokens += e.generation_tokens.get();
174 + snap.prefix_cache_queries += e.prefix_cache_queries.get();
175 + snap.prefix_cache_hits += e.prefix_cache_hits.get();
176 + }
177 + snap
178 +}
179 +
180 +/// Read the current scheduler gauge values, aggregated across engines.
181 +fn read_scheduler_gauges(engines: &[EngineMetrics]) -> (u64, u64, f64) {
182 + let mut num_running = 0u64;
183 + let mut num_waiting = 0u64;
184 + let mut kv_cache_usage_sum = 0.0f64;
185 +
186 + for e in engines {
187 + num_running += e.scheduler_running.get();
188 + num_waiting += e.scheduler_waiting.get();
189 + kv_cache_usage_sum += e.kv_cache_usage.get();
190 + }
191 +
192 + let kv_cache_usage = if !engines.is_empty() {
193 + kv_cache_usage_sum / engines.len() as f64
194 + } else {
195 + 0.0
196 + };
197 +
198 + (num_running, num_waiting, kv_cache_usage)
199 +}
rust/src/llm/src/output.rs new
+346
@@ -0,0 +1,346 @@
1 +use std::future::Future;
2 +use std::pin::Pin;
3 +use std::sync::Arc;
4 +use std::task::{Context, Poll, ready};
5 +
6 +use enum_as_inner::EnumAsInner;
7 +use futures::stream::FusedStream;
8 +use futures::{Stream, StreamExt as _, pin_mut};
9 +use serde::{Deserialize, Serialize};
10 +use vllm_engine_core_client::protocol::logprobs::Logprobs;
11 +use vllm_engine_core_client::protocol::{EngineCoreFinishReason, StopReason};
12 +use vllm_engine_core_client::{AbortCause, EngineCoreOutputStream};
13 +
14 +use crate::error::Result;
15 +use crate::request_metrics::{RequestMetricsTracker, current_unix_timestamp_secs};
16 +
17 +/// Final raw token output plus terminal stream metadata.
18 +#[derive(Debug, Clone, PartialEq)]
19 +pub struct CollectedGenerateOutput {
20 + pub request_id: String,
21 + pub prompt_token_ids: Vec<u32>,
22 + pub prompt_logprobs: Option<Logprobs>,
23 + pub token_ids: Vec<u32>,
24 + pub logprobs: Option<Logprobs>,
25 + pub finish_reason: FinishReason,
26 + /// Connector-specific KV transfer parameters for disaggregated serving.
27 + pub kv_transfer_params: Option<serde_json::Value>,
28 +}
29 +
30 +/// Prompt-scoped metadata emitted only once on the first [`GenerateOutput`] for
31 +/// one request.
32 +#[derive(Debug, Clone, PartialEq)]
33 +pub struct GeneratePromptInfo {
34 + /// Original prompt token IDs for this request.
35 + pub prompt_token_ids: Arc<[u32]>,
36 + /// Prompt logprobs returned by engine-core for scored prompt positions,
37 + /// when requested.
38 + pub prompt_logprobs: Option<Logprobs>,
39 +}
40 +
41 +/// The reason a request finished.
42 +///
43 +/// This is a higher-level abstraction over engine-core's finish and stop
44 +/// reasons.
45 +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, EnumAsInner)]
46 +pub enum FinishReason {
47 + /// Generation stopped for a stop string, stop token, or EOS.
48 + ///
49 + /// The inner stop reason is present for explicit stop strings or stop
50 + /// tokens, and absent for EOS-driven stops.
51 + Stop(Option<StopReason>),
52 + /// `max_tokens` or `max_model_len` was reached.
53 + Length,
54 + /// The request was aborted by the client.
55 + Abort,
56 + /// A retryable request-level internal error occurred.
57 + Error,
58 + /// A repetitive token pattern was detected.
59 + Repetition,
60 +}
61 +
62 +impl FinishReason {
63 + /// Construct a stop finish reason caused by EOS rather than an explicit
64 + /// stop string/token.
65 + pub fn stop_eos() -> Self {
66 + Self::Stop(None)
67 + }
68 +
69 + /// Returns a human-readable string for this finish reason, used for metrics
70 + /// and reporting.
71 + pub fn as_str(&self) -> &'static str {
72 + match self {
73 + Self::Stop(_) => "stop",
74 + Self::Length => "length",
75 + Self::Abort => "abort",
76 + Self::Error => "error",
77 + Self::Repetition => "repetition",
78 + }
79 + }
80 +
81 + /// If this is a stop finish reason, returns the inner stop reason if it
82 + /// exists.
83 + pub fn as_stop_reason(&self) -> Option<&StopReason> {
84 + match self {
85 + Self::Stop(stop_reason) => stop_reason.as_ref(),
86 + _ => None,
87 + }
88 + }
89 +
90 + /// If this is a stop finish reason, returns the inner stop reason if it
91 + /// exists.
92 + pub fn into_stop_reason(self) -> Option<StopReason> {
93 + match self {
94 + Self::Stop(stop_reason) => stop_reason,
95 + _ => None,
96 + }
97 + }
98 +}
99 +
100 +fn finish_reason_from_engine(
101 + finish_reason: Option<EngineCoreFinishReason>,
102 + stop_reason: Option<StopReason>,
103 +) -> Option<FinishReason> {
104 + finish_reason.map(|reason| match reason {
105 + EngineCoreFinishReason::Stop => FinishReason::Stop(stop_reason),
106 + EngineCoreFinishReason::Length => FinishReason::Length,
107 + EngineCoreFinishReason::Abort => FinishReason::Abort,
108 + EngineCoreFinishReason::Error => FinishReason::Error,
109 + EngineCoreFinishReason::Repetition => FinishReason::Repetition,
110 + })
111 +}
112 +
113 +/// Token and logprob output item returned by [`GenerateOutputStream`].
114 +///
115 +/// Original Python output reference:
116 +/// <https://github.com/vllm-project/vllm/blob/bc2c0c86efb28e77677a3cfb8687e976914a313a/vllm/outputs.py#L85-L143>
117 +#[derive(Debug, Clone, PartialEq)]
118 +pub struct GenerateOutput {
119 + /// Unique ID of the request that produced this output.
120 + pub request_id: String,
121 + /// One-time prompt metadata emitted only on the first output for this
122 + /// request.
123 + pub prompt_info: Option<GeneratePromptInfo>,
124 + /// Newly produced token IDs for this step.
125 + pub token_ids: Vec<u32>,
126 + /// Sample logprobs for the generated positions in this step.
127 + pub logprobs: Option<Logprobs>,
128 + /// Terminal finish reason, when this is the final output for the request.
129 + pub finish_reason: Option<FinishReason>,
130 + /// Connector-specific KV transfer parameters for disaggregated serving.
131 + pub kv_transfer_params: Option<serde_json::Value>,
132 +}
133 +
134 +impl GenerateOutput {
135 + /// Returns the prompt token IDs when this output carries
136 + /// [`GeneratePromptInfo`].
137 + ///
138 + /// Only the first output for a request can return `Some`; all later outputs
139 + /// return `None`.
140 + pub fn prompt_token_ids(&self) -> Option<&Arc<[u32]>> {
141 + self.prompt_info.as_ref().map(|info| &info.prompt_token_ids)
142 + }
143 +
144 + /// Returns the prompt logprobs when this output carries
145 + /// [`GeneratePromptInfo`].
146 + ///
147 + /// Only the first output for a request can return `Some`; all later outputs
148 + /// return `None`.
149 + pub fn prompt_logprobs(&self) -> Option<&Logprobs> {
150 + self.prompt_info.as_ref().and_then(|info| info.prompt_logprobs.as_ref())
151 + }
152 +
153 + /// Returns whether this output is terminal for the request.
154 + pub fn finished(&self) -> bool {
155 + self.finish_reason.is_some()
156 + }
157 +}
158 +
159 +#[cfg(any(test, feature = "test-util"))]
160 +impl GenerateOutput {
161 + /// Build a [`GenerateOutput`] for tests.
162 + pub fn for_test(
163 + prompt_token_ids: Option<Arc<[u32]>>,
164 + token_ids: Vec<u32>,
165 + finish_reason: Option<FinishReason>,
166 + ) -> Self {
167 + Self {
168 + request_id: String::new(),
169 + prompt_info: prompt_token_ids.map(|ids| GeneratePromptInfo {
170 + prompt_token_ids: ids,
171 + prompt_logprobs: None,
172 + }),
173 + token_ids,
174 + logprobs: None,
175 + finish_reason,
176 + kv_transfer_params: None,
177 + }
178 + }
179 +}
180 +
181 +/// Stream of per-request generate outputs for one request.
182 +///
183 +/// - A normal termination of the stream represents a clean completion of the request.
184 +/// - For errors, unexpected closes, or explicit aborts, the stream terminates with an error.
185 +pub struct GenerateOutputStream {
186 + pending_prompt_info: Option<GeneratePromptInfo>,
187 + raw_stream: EngineCoreOutputStream,
188 + request_metrics: RequestMetricsTracker,
189 +}
190 +
191 +impl GenerateOutputStream {
192 + /// Create a new generate output stream by adapting one raw engine-core
193 + /// output stream.
194 + pub(crate) fn new(
195 + prompt_token_ids: Arc<[u32]>,
196 + raw_stream: EngineCoreOutputStream,
197 + request_metrics: RequestMetricsTracker,
198 + ) -> Self {
199 + Self {
200 + pending_prompt_info: Some(GeneratePromptInfo {
201 + prompt_token_ids,
202 + prompt_logprobs: None,
203 + }),
204 + raw_stream,
205 + request_metrics,
206 + }
207 + }
208 +
209 + /// Return the internal engine request ID bound to this stream.
210 + pub fn request_id(&self) -> &str {
211 + self.raw_stream.request_id()
212 + }
213 +}
214 +
215 +impl Stream for GenerateOutputStream {
216 + type Item = Result<GenerateOutput>;
217 +
218 + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
219 + let raw = match ready!(Pin::new(&mut self.raw_stream).poll_next(cx)) {
220 + Some(Ok(raw)) => raw,
221 + Some(Err(error)) => return Poll::Ready(Some(Err(error.into()))),
222 + None => return Poll::Ready(None),
223 + };
224 +
225 + let received_at = current_unix_timestamp_secs();
226 + self.request_metrics.observe_output(
227 + raw.engine_index,
228 + raw.timestamp,
229 + received_at,
230 + &raw.output,
231 + );
232 +
233 + let raw = raw.output;
234 +
235 + // Populate the one-time prompt info on the first output.
236 + if let Some(info) = &mut self.pending_prompt_info
237 + && info.prompt_logprobs.is_none()
238 + {
239 + info.prompt_logprobs =
240 + raw.new_prompt_logprobs_tensors.map(|value| value.into_direct().unwrap());
241 + }
242 +
243 + let logprobs = raw.new_logprobs.map(|value| value.into_direct().unwrap());
244 +
245 + let finish_reason = finish_reason_from_engine(raw.finish_reason, raw.stop_reason);
246 + if let Some(finish_reason) = finish_reason.as_ref() {
247 + self.request_metrics.record_finished(received_at, finish_reason.clone());
248 + }
249 +
250 + let output = GenerateOutput {
251 + request_id: raw.request_id,
252 + prompt_info: self.pending_prompt_info.take(),
253 + token_ids: raw.new_token_ids,
254 + logprobs,
255 + finish_reason,
256 + kv_transfer_params: raw.kv_transfer_params,
257 + };
258 +
259 + Poll::Ready(Some(Ok(output)))
260 + }
261 +}
262 +
263 +impl FusedStream for GenerateOutputStream {
264 + fn is_terminated(&self) -> bool {
265 + self.raw_stream.is_terminated()
266 + }
267 +}
268 +
269 +impl Drop for GenerateOutputStream {
270 + fn drop(&mut self) {
271 + if self.raw_stream.is_terminated() {
272 + // Already terminated cleanly, no need to record abort metrics.
273 + return;
274 + }
275 +
276 + // If the user or the upper layer drops a live generate stream,
277 + // `EngineCoreOutputStream::Drop` will trigger an engine-side abort. Record the
278 + // matching terminal request metrics here so frontend-driven aborts are still
279 + // visible as `finished_reason=...` instead of disappearing from observability
280 + // entirely.
281 + let finish_reason = match AbortCause::current() {
282 + AbortCause::DroppedStream => FinishReason::Abort,
283 + AbortCause::StopStringMatched => FinishReason::Stop(None),
284 + };
285 +
286 + self.request_metrics
287 + .record_finished(current_unix_timestamp_secs(), finish_reason);
288 + }
289 +}
290 +
291 +#[allow(clippy::manual_async_fn, reason = "specify `Send` bound")]
292 +#[easy_ext::ext(GenerateOutputStreamExt)]
293 +impl<T: Stream<Item = Result<GenerateOutput>> + Send> T {
294 + /// Collect the raw generate stream to completion and return the final token
295 + /// output.
296 + pub fn collect_output(self) -> impl Future<Output = Result<CollectedGenerateOutput>> + Send {
297 + async move {
298 + let stream = self;
299 + pin_mut!(stream);
300 + let mut prompt_token_ids = None;
301 + let mut prompt_logprobs = None;
302 + let mut collected: Option<CollectedGenerateOutput> = None;
303 +
304 + while let Some(output) = stream.next().await.transpose()? {
305 + if let Some(info) = output.prompt_info {
306 + if prompt_token_ids.is_none() {
307 + prompt_token_ids = Some(info.prompt_token_ids.to_vec());
308 + }
309 + if prompt_logprobs.is_none() {
310 + prompt_logprobs = info.prompt_logprobs;
311 + }
312 + }
313 +
314 + if let Some(existing) = collected.as_mut() {
315 + existing.token_ids.extend(output.token_ids);
316 + if let Some(step_logprobs) = output.logprobs {
317 + if let Some(collected_logprobs) = existing.logprobs.as_mut() {
318 + collected_logprobs.positions.extend(step_logprobs.positions);
319 + } else {
320 + existing.logprobs = Some(step_logprobs);
321 + }
322 + }
323 + } else {
324 + collected = Some(CollectedGenerateOutput {
325 + request_id: output.request_id,
326 + prompt_token_ids: prompt_token_ids.take().unwrap_or_default(),
327 + prompt_logprobs: prompt_logprobs.take(),
328 + token_ids: output.token_ids,
329 + logprobs: output.logprobs,
330 + finish_reason: FinishReason::Error,
331 + kv_transfer_params: None,
332 + });
333 + }
334 +
335 + if let Some(finish_reason) = output.finish_reason {
336 + let mut collected = collected.expect("terminal output must exist");
337 + collected.finish_reason = finish_reason;
338 + collected.kv_transfer_params = output.kv_transfer_params;
339 + return Ok(collected);
340 + }
341 + }
342 +
343 + unreachable!("generate stream should yield an error instead of closing early")
344 + }
345 + }
346 +}
rust/src/llm/src/request.rs new
+201
@@ -0,0 +1,201 @@
1 +use std::collections::BTreeMap;
2 +use std::time::{SystemTime, UNIX_EPOCH};
3 +
4 +use uuid::Uuid;
5 +use vllm_engine_core_client::protocol::multimodal::MmFeatures;
6 +use vllm_engine_core_client::protocol::{EngineCoreRequest, EngineCoreSamplingParams, OpaqueValue};
7 +
8 +use crate::error::{Error, Result};
9 +
10 +/// Tokenized decoder-only generate request accepted by [`crate::Llm`].
11 +///
12 +/// This is the first-stage Rust subset of the inputs that eventually flow into
13 +/// Python `AsyncLLM.generate()`. The boundary is intentionally above
14 +/// [`EngineCoreRequest`], but below higher-level text and multimodal
15 +/// preprocessing.
16 +///
17 +/// Original Python API reference:
18 +/// <https://github.com/vllm-project/vllm/blob/bc2c0c86efb28e77677a3cfb8687e976914a313a/vllm/engine/protocol.py#L67-L84>
19 +#[derive(Debug, Clone, PartialEq)]
20 +pub struct GenerateRequest {
21 + /// Unique ID of the request.
22 + pub request_id: String,
23 + /// Token IDs of the prompt.
24 + pub prompt_token_ids: Vec<u32>,
25 + /// Sampling parameters forwarded to engine-core.
26 + pub sampling_params: EngineCoreSamplingParams,
27 + /// Optional multimodal features already prepared by `vllm-chat`.
28 + pub mm_features: Option<MmFeatures>,
29 +
30 + // Fields below are currently likely unused by callers.
31 + pub arrival_time: Option<f64>,
32 + pub cache_salt: Option<String>,
33 + pub trace_headers: Option<BTreeMap<String, String>>,
34 + pub priority: i32,
35 + pub data_parallel_rank: Option<u32>,
36 + pub reasoning_ended: Option<bool>,
37 + pub lora_request: Option<OpaqueValue>,
38 +}
39 +
40 +#[derive(Debug)]
41 +pub(crate) struct PreparedGenerateRequest {
42 + pub engine_request: EngineCoreRequest,
43 +}
44 +
45 +impl GenerateRequest {
46 + /// Validate and lower this request into the raw engine-core request format.
47 + pub(crate) fn prepare(self, randomize_request_id: bool) -> Result<PreparedGenerateRequest> {
48 + if self.prompt_token_ids.is_empty() {
49 + return Err(Error::EmptyPromptTokenIds {
50 + request_id: self.request_id,
51 + });
52 + }
53 + let GenerateRequest {
54 + request_id,
55 + prompt_token_ids,
56 + sampling_params,
57 + mm_features,
58 + arrival_time,
59 + cache_salt,
60 + trace_headers,
61 + priority,
62 + data_parallel_rank,
63 + reasoning_ended,
64 + lora_request,
65 + } = self;
66 +
67 + let external_request_id = request_id;
68 + let engine_request_id = if randomize_request_id {
69 + let random_suffix = Uuid::new_v4().simple().to_string();
70 + format!("{external_request_id}-{}", &random_suffix[..8])
71 + } else {
72 + external_request_id.clone()
73 + };
74 +
75 + Ok(PreparedGenerateRequest {
76 + engine_request: EngineCoreRequest {
77 + request_id: engine_request_id,
78 + prompt_token_ids: Some(prompt_token_ids),
79 + mm_features,
80 + sampling_params: Some(sampling_params),
81 + pooling_params: None,
82 + arrival_time: arrival_time.unwrap_or_else(current_unix_timestamp_secs),
83 + lora_request,
84 + cache_salt,
85 + data_parallel_rank,
86 + prompt_embeds: None,
87 + prompt_is_token_ids: None,
88 + client_index: 0,
89 + current_wave: 0,
90 + priority,
91 + trace_headers,
92 + resumable: false,
93 + external_req_id: Some(external_request_id),
94 + reasoning_ended,
95 + reasoning_parser_kwargs: None,
96 + abort_immediately: false,
97 + },
98 + })
99 + }
100 +}
101 +
102 +impl PreparedGenerateRequest {
103 + /// Return the original prompt token IDs copied into the raw engine request.
104 + pub fn prompt_token_ids(&self) -> &[u32] {
105 + self.engine_request
106 + .prompt_token_ids
107 + .as_ref()
108 + .expect("prepared request must have prompt token ids")
109 + }
110 +}
111 +
112 +fn current_unix_timestamp_secs() -> f64 {
113 + SystemTime::now()
114 + .duration_since(UNIX_EPOCH)
115 + .expect("system clock is before unix epoch")
116 + .as_secs_f64()
117 +}
118 +
119 +#[cfg(test)]
120 +mod tests {
121 + use std::collections::BTreeMap;
122 +
123 + use vllm_engine_core_client::protocol::EngineCoreSamplingParams;
124 +
125 + use super::GenerateRequest;
126 + use crate::error::Error;
127 +
128 + fn sample_request() -> GenerateRequest {
129 + GenerateRequest {
130 + request_id: "req-1".to_string(),
131 + prompt_token_ids: vec![11, 22, 33],
132 + sampling_params: EngineCoreSamplingParams::for_test(),
133 + mm_features: None,
134 + arrival_time: Some(42.5),
135 + cache_salt: Some("salt".to_string()),
136 + trace_headers: Some(BTreeMap::from([(
137 + "x-trace-id".to_string(),
138 + "abc".to_string(),
139 + )])),
140 + priority: 3,
141 + data_parallel_rank: Some(2),
142 + reasoning_ended: Some(true),
143 + lora_request: None,
144 + }
145 + }
146 +
147 + #[test]
148 + fn prepare_builds_engine_core_request() {
149 + let prepared = sample_request().prepare(true).unwrap();
150 +
151 + assert_eq!(prepared.prompt_token_ids(), &[11, 22, 33]);
152 +
153 + let request = prepared.engine_request;
154 + assert_eq!(request.external_req_id.as_deref(), Some("req-1"));
155 + assert!(request.request_id.starts_with("req-1-"));
156 + assert_ne!(request.request_id, "req-1");
157 + assert_eq!(request.prompt_token_ids.as_deref(), Some(&[11, 22, 33][..]));
158 + assert_eq!(request.arrival_time, 42.5);
159 + assert_eq!(request.cache_salt.as_deref(), Some("salt"));
160 + assert_eq!(request.data_parallel_rank, Some(2));
161 + assert_eq!(
162 + request.trace_headers,
163 + Some(BTreeMap::from([(
164 + "x-trace-id".to_string(),
165 + "abc".to_string(),
166 + )]))
167 + );
168 + assert_eq!(request.reasoning_ended, Some(true));
169 + }
170 +
171 + #[test]
172 + fn prepare_rejects_empty_prompt_tokens() {
173 + let mut request = sample_request();
174 + request.prompt_token_ids.clear();
175 +
176 + let error = request.prepare(true).unwrap_err();
177 + assert!(matches!(
178 + error,
179 + Error::EmptyPromptTokenIds { request_id } if request_id == "req-1"
180 + ));
181 + }
182 +
183 + #[test]
184 + fn prepare_can_preserve_external_request_id() {
185 + let prepared = sample_request().prepare(false).unwrap();
186 +
187 + let request = prepared.engine_request;
188 + assert_eq!(request.external_req_id.as_deref(), Some("req-1"));
189 + assert_eq!(request.request_id, "req-1");
190 + }
191 +
192 + #[test]
193 + fn prepare_forwards_multimodal_features() {
194 + let mut request = sample_request();
195 + request.mm_features = Some(Vec::new());
196 +
197 + let prepared = request.prepare(false).unwrap();
198 +
199 + assert_eq!(prepared.engine_request.mm_features, Some(Vec::new()));
200 + }
201 +}
rust/src/llm/src/request_metrics.rs new
+391
@@ -0,0 +1,391 @@
1 +use std::time::{SystemTime, UNIX_EPOCH};
2 +
3 +use vllm_engine_core_client::protocol::stats::PrefillStats;
4 +use vllm_engine_core_client::protocol::{EngineCoreEvent, EngineCoreEventType, EngineCoreOutput};
5 +use vllm_metrics::{
6 + EngineLabels, FinishedReasonLabels, METRICS, PromptTokenSourceLabels, RequestMetrics,
7 +};
8 +
9 +use crate::FinishReason;
10 +
11 +fn metrics() -> &'static RequestMetrics {
12 + &METRICS.request
13 +}
14 +
15 +const PROMPT_TOKEN_SOURCE_LOCAL_COMPUTE: &str = "local_compute";
16 +const PROMPT_TOKEN_SOURCE_LOCAL_CACHE_HIT: &str = "local_cache_hit";
17 +const PROMPT_TOKEN_SOURCE_EXTERNAL_KV_TRANSFER: &str = "external_kv_transfer";
18 +
19 +/// Request-scoped metrics state tracked across streamed engine-core updates.
20 +///
21 +/// This is the Rust-side counterpart of the Python frontend's request-lifecycle
22 +/// bookkeeping, centered on `RequestStateStats` and the per-output/per-finished
23 +/// update flow.
24 +///
25 +/// Original Python definitions:
26 +/// <https://github.com/vllm-project/vllm/blob/bc2c0c86efb28e77677a3cfb8687e976914a313a/vllm/v1/metrics/stats.py#L200-L237>
27 +///
28 +/// Original Python update flow:
29 +/// <https://github.com/vllm-project/vllm/blob/bc2c0c86efb28e77677a3cfb8687e976914a313a/vllm/v1/engine/output_processor.py#L600-L677>
30 +#[derive(Debug, Clone)]
31 +pub(crate) struct RequestMetricsTracker {
32 + model_name: String,
33 + arrival_time: f64,
34 + prompt_len: u32,
35 + max_tokens_param: Option<u32>,
36 + n_param: u32,
37 + is_prefilling: bool,
38 + queued_ts: f64,
39 + scheduled_ts: f64,
40 + first_token_ts: f64,
41 + last_token_ts: f64,
42 + first_token_latency: f64,
43 + num_generation_tokens: u32,
44 + latest_num_cached_tokens: u32,
45 + last_seen_engine_index: u32,
46 +}
47 +
48 +impl RequestMetricsTracker {
49 + /// Create the per-request tracker from the normalized `llm`-layer request
50 + /// context.
51 + pub(crate) fn new(
52 + model_name: String,
53 + arrival_time: f64,
54 + prompt_len: u32,
55 + max_tokens_param: Option<u32>,
56 + n_param: u32,
57 + ) -> Self {
58 + Self {
59 + model_name,
60 + arrival_time,
61 + prompt_len,
62 + max_tokens_param,
63 + n_param,
64 + is_prefilling: true,
65 + queued_ts: 0.0,
66 + scheduled_ts: 0.0,
67 + first_token_ts: 0.0,
68 + last_token_ts: 0.0,
69 + first_token_latency: 0.0,
70 + num_generation_tokens: 0,
71 + latest_num_cached_tokens: 0,
72 + last_seen_engine_index: 0,
73 + }
74 + }
75 +
76 + /// Update request-lifecycle state from one engine-core output item.
77 + ///
78 + /// Original Python stats logic:
79 + /// <https://github.com/vllm-project/vllm/blob/bc2c0c86efb28e77677a3cfb8687e976914a313a/vllm/v1/metrics/stats.py#L331-L384>
80 + pub(crate) fn observe_output(
81 + &mut self,
82 + engine_index: u32,
83 + batch_timestamp: f64,
84 + received_at: f64,
85 + output: &EngineCoreOutput,
86 + ) {
87 + self.last_seen_engine_index = engine_index;
88 + if let Some(prefill_stats) = &output.prefill_stats {
89 + self.latest_num_cached_tokens = prefill_stats.num_cached_tokens;
90 + }
91 + self.num_generation_tokens += output.new_token_ids.len() as u32;
92 + metrics()
93 + .generation_tokens
94 + .get_or_create(&engine_labels(&self.model_name, engine_index))
95 + .inc_by(output.new_token_ids.len() as u64);
96 +
97 + if let Some(events) = &output.events {
98 + self.observe_events(engine_index, events);
99 + }
100 +
101 + if self.is_prefilling {
102 + if let Some(prefill_stats) = &output.prefill_stats {
103 + record_prompt_tokens(&self.model_name, engine_index, prefill_stats);
104 + }
105 + self.first_token_latency = received_at - self.arrival_time;
106 + observe_time_to_first_token_seconds(
107 + &self.model_name,
108 + engine_index,
109 + self.first_token_latency,
110 + );
111 + self.first_token_ts = batch_timestamp;
112 + self.is_prefilling = false;
113 + } else if self.last_token_ts > 0.0 {
114 + observe_inter_token_latency_seconds(
115 + &self.model_name,
116 + engine_index,
117 + batch_timestamp - self.last_token_ts,
118 + );
119 + }
120 +
121 + self.last_token_ts = batch_timestamp;
122 + }
123 +
124 + /// Emit the terminal request metrics once a finished output has been
125 + /// observed.
126 + ///
127 + /// Original Python finished-request stats:
128 + /// <https://github.com/vllm-project/vllm/blob/bc2c0c86efb28e77677a3cfb8687e976914a313a/vllm/v1/metrics/stats.py#L222-L237>
129 + pub(crate) fn record_finished(&self, received_at: f64, finish_reason: FinishReason) {
130 + let labels = engine_labels(&self.model_name, self.last_seen_engine_index);
131 + let prefill_kv_computed_tokens =
132 + self.prompt_len.saturating_sub(self.latest_num_cached_tokens);
133 + let e2e_latency_seconds = received_at - self.arrival_time;
134 + let queue_time_seconds = diff_or_zero(self.scheduled_ts, self.queued_ts);
135 + let prefill_time_seconds = diff_or_zero(self.first_token_ts, self.scheduled_ts);
136 + let decode_time_seconds = diff_or_zero(self.last_token_ts, self.first_token_ts);
137 + let inference_time_seconds = diff_or_zero(self.last_token_ts, self.scheduled_ts);
138 + let time_per_output_token_seconds = if self.num_generation_tokens > 1 {
139 + diff_or_zero(self.last_token_ts, self.first_token_ts)
140 + / (self.num_generation_tokens - 1) as f64
141 + } else {
142 + 0.0
143 + };
144 +
145 + record_request_success(&self.model_name, self.last_seen_engine_index, finish_reason);
146 + metrics()
147 + .request_prompt_tokens
148 + .get_or_create(&labels)
149 + .observe(self.prompt_len as f64);
150 + metrics()
151 + .request_generation_tokens
152 + .get_or_create(&labels)
153 + .observe(self.num_generation_tokens as f64);
154 + metrics()
155 + .request_max_num_generation_tokens
156 + .get_or_create(&labels)
157 + .observe(self.num_generation_tokens as f64);
158 + if let Some(max_tokens_param) = self.max_tokens_param {
159 + metrics()
160 + .request_params_max_tokens
161 + .get_or_create(&labels)
162 + .observe(max_tokens_param as f64);
163 + }
164 + metrics().request_params_n.get_or_create(&labels).observe(self.n_param as f64);
165 + metrics()
166 + .request_prefill_kv_computed_tokens
167 + .get_or_create(&labels)
168 + .observe(prefill_kv_computed_tokens as f64);
169 + metrics()
170 + .e2e_request_latency_seconds
171 + .get_or_create(&labels)
172 + .observe(e2e_latency_seconds);
173 + metrics()
174 + .request_queue_time_seconds
175 + .get_or_create(&labels)
176 + .observe(queue_time_seconds);
177 + metrics()
178 + .request_prefill_time_seconds
179 + .get_or_create(&labels)
180 + .observe(prefill_time_seconds);
181 + metrics()
182 + .request_decode_time_seconds
183 + .get_or_create(&labels)
184 + .observe(decode_time_seconds);
185 + metrics()
186 + .request_inference_time_seconds
187 + .get_or_create(&labels)
188 + .observe(inference_time_seconds);
189 + metrics()
190 + .request_time_per_output_token_seconds
191 + .get_or_create(&labels)
192 + .observe(time_per_output_token_seconds);
193 + }
194 +
195 + fn observe_events(&mut self, engine_index: u32, events: &[EngineCoreEvent]) {
196 + for event in events {
197 + match event.r#type {
198 + EngineCoreEventType::Queued => {
199 + self.queued_ts = event.timestamp;
200 + }
201 + EngineCoreEventType::Scheduled => {
202 + if self.scheduled_ts == 0.0 {
203 + self.scheduled_ts = event.timestamp;
204 + }
205 + }
206 + EngineCoreEventType::Preempted => {
207 + metrics()
208 + .num_preemptions
209 + .get_or_create(&engine_labels(&self.model_name, engine_index))
210 + .inc();
211 + }
212 + }
213 + }
214 + }
215 +}
216 +
217 +fn engine_labels(model_name: &str, engine: u32) -> EngineLabels {
218 + EngineLabels {
219 + model_name: model_name.to_string(),
220 + engine,
221 + }
222 +}
223 +
224 +fn observe_time_to_first_token_seconds(model_name: &str, engine: u32, seconds: f64) {
225 + metrics()
226 + .time_to_first_token_seconds
227 + .get_or_create(&engine_labels(model_name, engine))
228 + .observe(seconds);
229 +}
230 +
231 +fn observe_inter_token_latency_seconds(model_name: &str, engine: u32, seconds: f64) {
232 + metrics()
233 + .inter_token_latency_seconds
234 + .get_or_create(&engine_labels(model_name, engine))
235 + .observe(seconds);
236 +}
237 +
238 +fn record_request_success(model_name: &str, engine: u32, finish_reason: FinishReason) {
239 + metrics()
240 + .request_success
241 + .get_or_create(&FinishedReasonLabels {
242 + model_name: model_name.to_string(),
243 + engine,
244 + finished_reason: finish_reason.as_str(),
245 + })
246 + .inc();
247 +}
248 +
249 +fn prompt_token_source_labels(
250 + model_name: &str,
251 + engine: u32,
252 + source: &'static str,
253 +) -> PromptTokenSourceLabels {
254 + PromptTokenSourceLabels {
255 + model_name: model_name.to_string(),
256 + engine,
257 + source,
258 + }
259 +}
260 +
261 +fn record_prompt_tokens(model_name: &str, engine: u32, prefill_stats: &PrefillStats) {
262 + let computed = prefill_stats.num_computed_tokens as u64;
263 + let local_cache_hit = prefill_stats.num_local_cached_tokens as u64;
264 + let external_kv_transfer = prefill_stats.num_external_cached_tokens as u64;
265 +
266 + metrics()
267 + .prompt_tokens
268 + .get_or_create(&engine_labels(model_name, engine))
269 + .inc_by(prefill_stats.num_prompt_tokens as u64);
270 + metrics()
271 + .prompt_tokens_by_source
272 + .get_or_create(&prompt_token_source_labels(
273 + model_name,
274 + engine,
275 + PROMPT_TOKEN_SOURCE_LOCAL_COMPUTE,
276 + ))
277 + .inc_by(computed);
278 + metrics()
279 + .prompt_tokens_by_source
280 + .get_or_create(&prompt_token_source_labels(
281 + model_name,
282 + engine,
283 + PROMPT_TOKEN_SOURCE_LOCAL_CACHE_HIT,
284 + ))
285 + .inc_by(local_cache_hit);
286 + metrics()
287 + .prompt_tokens_by_source
288 + .get_or_create(&prompt_token_source_labels(
289 + model_name,
290 + engine,
291 + PROMPT_TOKEN_SOURCE_EXTERNAL_KV_TRANSFER,
292 + ))
293 + .inc_by(external_kv_transfer);
294 + metrics()
295 + .prompt_tokens_cached
296 + .get_or_create(&engine_labels(model_name, engine))
297 + .inc_by(prefill_stats.num_cached_tokens as u64);
298 +}
299 +
300 +fn diff_or_zero(end: f64, start: f64) -> f64 {
301 + if end > 0.0 && start > 0.0 && end >= start {
302 + end - start
303 + } else {
304 + 0.0
305 + }
306 +}
307 +
308 +/// Return the current wall-clock time in seconds since the Unix epoch.
309 +///
310 +/// This is used for frontend-side latency measurements such as TTFT and E2E,
311 +/// matching the Python frontend's use of wall-clock request arrival/iteration
312 +/// timestamps rather than engine-core's monotonic scheduler timestamps.
313 +///
314 +/// Original Python request timestamp source:
315 +/// <https://github.com/vllm-project/vllm/blob/bc2c0c86efb28e77677a3cfb8687e976914a313a/vllm/v1/metrics/stats.py#L206-L216>
316 +pub(crate) fn current_unix_timestamp_secs() -> f64 {
317 + SystemTime::now()
318 + .duration_since(UNIX_EPOCH)
319 + .expect("system clock is before unix epoch")
320 + .as_secs_f64()
321 +}
322 +
323 +#[cfg(test)]
324 +mod tests {
325 + use vllm_engine_core_client::protocol::stats::PrefillStats;
326 + use vllm_engine_core_client::protocol::{EngineCoreEvent, EngineCoreEventType};
327 +
328 + use super::{RequestMetricsTracker, diff_or_zero};
329 +
330 + #[test]
331 + fn tracker_updates_timing_state_across_prefill_decode_and_finish() {
332 + let mut tracker = RequestMetricsTracker::new("model".to_string(), 100.0, 64, Some(128), 1);
333 +
334 + tracker.observe_output(
335 + 2,
336 + 10.0,
337 + 100.2,
338 + &vllm_engine_core_client::protocol::EngineCoreOutput {
339 + request_id: "req-1".to_string(),
340 + new_token_ids: vec![1],
341 + finish_reason: None,
342 + events: Some(vec![
343 + EngineCoreEvent {
344 + r#type: EngineCoreEventType::Queued,
345 + timestamp: 8.0,
346 + },
347 + EngineCoreEvent {
348 + r#type: EngineCoreEventType::Scheduled,
349 + timestamp: 9.0,
350 + },
351 + ]),
352 + prefill_stats: Some(PrefillStats {
353 + num_prompt_tokens: 64,
354 + num_computed_tokens: 60,
355 + num_cached_tokens: 4,
356 + num_local_cached_tokens: 4,
357 + num_external_cached_tokens: 0,
358 + }),
359 + ..Default::default()
360 + },
361 + );
362 + tracker.observe_output(
363 + 2,
364 + 11.5,
365 + 100.4,
366 + &vllm_engine_core_client::protocol::EngineCoreOutput {
367 + request_id: "req-1".to_string(),
368 + new_token_ids: vec![2, 3],
369 + finish_reason: None,
370 + events: Some(vec![EngineCoreEvent {
371 + r#type: EngineCoreEventType::Preempted,
372 + timestamp: 10.5,
373 + }]),
374 + ..Default::default()
375 + },
376 + );
377 +
378 + assert!(!tracker.is_prefilling);
379 + assert_eq!(tracker.last_seen_engine_index, 2);
380 + assert_eq!(tracker.num_generation_tokens, 3);
381 + assert_eq!(tracker.queued_ts, 8.0);
382 + assert_eq!(tracker.scheduled_ts, 9.0);
383 + assert_eq!(tracker.first_token_ts, 10.0);
384 + assert_eq!(tracker.last_token_ts, 11.5);
385 + assert!((tracker.first_token_latency - 0.2).abs() < 1e-9);
386 + assert_eq!(
387 + diff_or_zero(tracker.last_token_ts, tracker.first_token_ts),
388 + 1.5
389 + );
390 + }
391 +}
rust/src/llm/tests/generate.rs new
+749
@@ -0,0 +1,749 @@
1 +use std::collections::BTreeSet;
2 +use std::sync::Once;
3 +use std::time::Duration;
4 +
5 +use futures::StreamExt as _;
6 +use tokio::time::timeout;
7 +use tracing_subscriber::EnvFilter;
8 +use uuid::Uuid;
9 +use vllm_engine_core_client::protocol::logprobs::{
10 + Logprobs, MaybeWireLogprobs, PositionLogprobs, TokenLogprob,
11 +};
12 +use vllm_engine_core_client::protocol::stats::PrefillStats;
13 +use vllm_engine_core_client::protocol::{
14 + EngineCoreEvent, EngineCoreEventType, EngineCoreFinishReason, EngineCoreOutput,
15 + EngineCoreOutputs, EngineCoreRequest, EngineCoreSamplingParams,
16 +};
17 +use vllm_engine_core_client::test_utils::{IpcNamespace, spawn_mock_engine_task};
18 +use vllm_engine_core_client::{EngineCoreClient, EngineCoreClientConfig};
19 +use vllm_llm::{
20 + Error, FinishReason, GenerateOutputStreamExt as _, GeneratePromptInfo, GenerateRequest, Llm,
21 +};
22 +use vllm_metrics::METRICS;
23 +use zeromq::prelude::{SocketRecv, SocketSend};
24 +use zeromq::{DealerSocket, PushSocket, ZmqMessage};
25 +
26 +static TRACING: Once = Once::new();
27 +
28 +fn request_output(
29 + request_id: &str,
30 + new_token_ids: Vec<u32>,
31 + finish_reason: Option<EngineCoreFinishReason>,
32 +) -> EngineCoreOutput {
33 + request_output_with_events(request_id, new_token_ids, finish_reason, None)
34 +}
35 +
36 +fn request_output_with_events(
37 + request_id: &str,
38 + new_token_ids: Vec<u32>,
39 + finish_reason: Option<EngineCoreFinishReason>,
40 + events: Option<Vec<EngineCoreEvent>>,
41 +) -> EngineCoreOutput {
42 + EngineCoreOutput {
43 + request_id: request_id.to_string(),
44 + new_token_ids,
45 + new_logprobs: None,
46 + new_prompt_logprobs_tensors: None,
47 + pooling_output: None,
48 + finish_reason,
49 + stop_reason: None,
50 + events,
51 + kv_transfer_params: None,
52 + trace_headers: None,
53 + prefill_stats: None,
54 + routed_experts: None,
55 + num_nans_in_logits: 0,
56 + }
57 +}
58 +
59 +fn request_output_with_logprobs(
60 + request_id: &str,
61 + new_token_ids: Vec<u32>,
62 + finish_reason: Option<EngineCoreFinishReason>,
63 + new_logprobs: Option<Logprobs>,
64 + prompt_logprobs: Option<Logprobs>,
65 +) -> EngineCoreOutput {
66 + EngineCoreOutput {
67 + request_id: request_id.to_string(),
68 + new_token_ids,
69 + new_logprobs: new_logprobs.map(MaybeWireLogprobs::Direct),
70 + new_prompt_logprobs_tensors: prompt_logprobs.map(MaybeWireLogprobs::Direct),
71 + pooling_output: None,
72 + finish_reason,
73 + stop_reason: None,
74 + events: None,
75 + kv_transfer_params: None,
76 + trace_headers: None,
77 + prefill_stats: None,
78 + routed_experts: None,
79 + num_nans_in_logits: 0,
80 + }
81 +}
82 +
83 +fn request_output_with_logprobs_and_kv(
84 + request_id: &str,
85 + new_token_ids: Vec<u32>,
86 + finish_reason: Option<EngineCoreFinishReason>,
87 + new_logprobs: Option<Logprobs>,
88 + prompt_logprobs: Option<Logprobs>,
89 + kv_transfer_params: Option<serde_json::Value>,
90 +) -> EngineCoreOutput {
91 + EngineCoreOutput {
92 + request_id: request_id.to_string(),
93 + new_token_ids,
94 + new_logprobs: new_logprobs.map(MaybeWireLogprobs::Direct),
95 + new_prompt_logprobs_tensors: prompt_logprobs.map(MaybeWireLogprobs::Direct),
96 + pooling_output: None,
97 + finish_reason,
98 + stop_reason: None,
99 + events: None,
100 + kv_transfer_params,
101 + trace_headers: None,
102 + prefill_stats: None,
103 + routed_experts: None,
104 + num_nans_in_logits: 0,
105 + }
106 +}
107 +
108 +fn logprobs_for_position(
109 + sampled_token_id: u32,
110 + sampled_logprob: f32,
111 + sampled_rank: u32,
112 + top_token_id: u32,
113 + top_logprob: f32,
114 +) -> Logprobs {
115 + Logprobs {
116 + positions: vec![PositionLogprobs {
117 + entries: vec![
118 + TokenLogprob {
119 + token_id: sampled_token_id,
120 + logprob: sampled_logprob,
121 + rank: sampled_rank,
122 + },
123 + TokenLogprob {
124 + token_id: top_token_id,
125 + logprob: top_logprob,
126 + rank: 1,
127 + },
128 + ],
129 + }],
130 + }
131 +}
132 +
133 +fn prompt_logprobs() -> Logprobs {
134 + Logprobs {
135 + positions: vec![
136 + PositionLogprobs {
137 + entries: vec![
138 + TokenLogprob {
139 + token_id: 11,
140 + logprob: -0.1,
141 + rank: 2,
142 + },
143 + TokenLogprob {
144 + token_id: 7,
145 + logprob: -0.05,
146 + rank: 1,
147 + },
148 + ],
149 + },
150 + PositionLogprobs {
151 + entries: vec![
152 + TokenLogprob {
153 + token_id: 22,
154 + logprob: -0.2,
155 + rank: 3,
156 + },
157 + TokenLogprob {
158 + token_id: 8,
159 + logprob: -0.1,
160 + rank: 1,
161 + },
162 + ],
163 + },
164 + ],
165 + }
166 +}
167 +
168 +fn sample_generate_request(request_id: &str, max_tokens: u32) -> GenerateRequest {
169 + GenerateRequest {
170 + request_id: request_id.to_string(),
171 + prompt_token_ids: vec![11, 22],
172 + sampling_params: EngineCoreSamplingParams {
173 + max_tokens,
174 + ..EngineCoreSamplingParams::for_test()
175 + },
176 + mm_features: None,
177 + arrival_time: Some(42.5),
178 + cache_salt: None,
179 + trace_headers: None,
180 + priority: 0,
181 + data_parallel_rank: None,
182 + reasoning_ended: None,
183 + lora_request: None,
184 + }
185 +}
186 +
187 +async fn send_outputs(push: &mut PushSocket, outputs: EngineCoreOutputs) {
188 + push.send(ZmqMessage::from(rmp_serde::to_vec_named(&outputs).unwrap()))
189 + .await
190 + .unwrap();
191 +}
192 +
193 +async fn recv_engine_message(dealer: &mut DealerSocket) -> Vec<bytes::Bytes> {
194 + dealer.recv().await.unwrap().into_vec()
195 +}
196 +
197 +async fn connect_async_llm_with_ipc(
198 + handshake_address: String,
199 + client_index: u32,
200 + model_name: &str,
201 + ipc: &IpcNamespace,
202 +) -> Llm {
203 + let client = EngineCoreClient::connect(
204 + EngineCoreClientConfig::new_single(handshake_address)
205 + .with_model_name(model_name)
206 + .with_client_index(client_index)
207 + .with_local_input_output_addresses(
208 + Some(ipc.input_endpoint()),
209 + Some(ipc.output_endpoint()),
210 + ),
211 + )
212 + .await
213 + .unwrap();
214 + Llm::new(client)
215 +}
216 +
217 +fn request_metrics_model_name(prefix: &str) -> String {
218 + format!("{prefix}-{}", Uuid::new_v4().simple())
219 +}
220 +
221 +fn init_tracing() {
222 + TRACING.call_once(|| {
223 + let filter = EnvFilter::try_from_default_env()
224 + .unwrap_or_else(|_| EnvFilter::new("vllm_engine_core_client=debug"));
225 + let _ = tracing_subscriber::fmt().with_test_writer().with_env_filter(filter).try_init();
226 + });
227 +}
228 +
229 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
230 +async fn generate_streams_outputs() {
231 + init_tracing();
232 + let ipc = IpcNamespace::new().unwrap();
233 + let handshake_address = ipc.handshake_endpoint();
234 + let engine_id = b"engine-delta".to_vec();
235 +
236 + let (shutdown_tx, engine_task) = spawn_mock_engine_task(
237 + handshake_address.clone(),
238 + engine_id.clone(),
239 + |dealer, push| {
240 + Box::pin(async move {
241 + let add = recv_engine_message(dealer).await;
242 + assert_eq!(add[0].as_ref(), &[0x00]);
243 + let request: EngineCoreRequest = rmp_serde::from_slice(&add[1]).unwrap();
244 + assert_eq!(request.external_req_id.as_deref(), Some("req-delta"));
245 + assert!(request.request_id.starts_with("req-delta-"));
246 + assert_ne!(request.request_id, "req-delta");
247 + assert_eq!(request.client_index, 7);
248 + assert_eq!(request.prompt_token_ids, Some(vec![11, 22]));
249 +
250 + send_outputs(
251 + push,
252 + EngineCoreOutputs {
253 + outputs: vec![
254 + request_output_with_logprobs(
255 + &request.request_id,
256 + vec![1, 2],
257 + None,
258 + Some(logprobs_for_position(1, -0.3, 4, 9, -0.1)),
259 + Some(prompt_logprobs()),
260 + ),
261 + request_output_with_logprobs(
262 + &request.request_id,
263 + vec![3],
264 + Some(EngineCoreFinishReason::Length),
265 + Some(logprobs_for_position(3, -0.4, 5, 10, -0.2)),
266 + None,
267 + ),
268 + ],
269 + finished_requests: Some(BTreeSet::from([request.request_id.clone()])),
270 + ..Default::default()
271 + },
272 + )
273 + .await;
274 + })
275 + },
276 + );
277 +
278 + let llm = connect_async_llm_with_ipc(handshake_address, 7, "test-model", &ipc).await;
279 + let mut stream = llm.generate(sample_generate_request("req-delta", 3)).await.unwrap();
280 + let internal_id = stream.request_id().to_string();
281 +
282 + let first = stream.next().await.unwrap().unwrap();
283 + assert_eq!(first.request_id, internal_id);
284 + assert_eq!(
285 + first.prompt_info,
286 + Some(GeneratePromptInfo {
287 + prompt_token_ids: vec![11, 22].into(),
288 + prompt_logprobs: Some(prompt_logprobs()),
289 + })
290 + );
291 + assert_eq!(first.token_ids, vec![1, 2]);
292 + assert_eq!(
293 + first.logprobs,
294 + Some(logprobs_for_position(1, -0.3, 4, 9, -0.1))
295 + );
296 + assert_eq!(first.finish_reason, None);
297 +
298 + let second = stream.next().await.unwrap().unwrap();
299 + assert_eq!(second.prompt_info, None);
300 + assert_eq!(second.token_ids, vec![3]);
301 + assert_eq!(
302 + second.logprobs,
303 + Some(logprobs_for_position(3, -0.4, 5, 10, -0.2))
304 + );
305 + assert_eq!(second.finish_reason, Some(FinishReason::Length));
306 + assert!(stream.next().await.is_none());
307 +
308 + let _ = shutdown_tx.send(());
309 + engine_task.await.unwrap();
310 + llm.shutdown().await.unwrap();
311 +}
312 +
313 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
314 +async fn collect_output_aggregates_raw_tokens_logprobs_and_terminal_metadata() {
315 + init_tracing();
316 + let ipc = IpcNamespace::new().unwrap();
317 + let handshake_address = ipc.handshake_endpoint();
318 + let engine_id = b"engine-collect-output".to_vec();
319 +
320 + let (shutdown_tx, engine_task) = spawn_mock_engine_task(
321 + handshake_address.clone(),
322 + engine_id.clone(),
323 + |dealer, push| {
324 + Box::pin(async move {
325 + let add = recv_engine_message(dealer).await;
326 + let request: EngineCoreRequest = rmp_serde::from_slice(&add[1]).unwrap();
327 + assert_eq!(request.external_req_id.as_deref(), Some("req-collect"));
328 + assert!(request.request_id.starts_with("req-collect-"));
329 +
330 + send_outputs(
331 + push,
332 + EngineCoreOutputs {
333 + engine_index: 0,
334 + outputs: vec![
335 + request_output_with_logprobs(
336 + &request.request_id,
337 + vec![33],
338 + None,
339 + Some(logprobs_for_position(33, -0.1, 1, 99, -0.2)),
340 + Some(prompt_logprobs()),
341 + ),
342 + request_output_with_logprobs_and_kv(
343 + &request.request_id,
344 + vec![44],
345 + Some(EngineCoreFinishReason::Stop),
346 + Some(logprobs_for_position(44, -0.3, 1, 88, -0.4)),
347 + None,
348 + Some(serde_json::json!({"connector": "x"})),
349 + ),
350 + ],
351 + scheduler_stats: None,
352 + timestamp: 0.0,
353 + utility_output: None,
354 + finished_requests: None,
355 + wave_complete: None,
356 + start_wave: None,
357 + },
358 + )
359 + .await;
360 + })
361 + },
362 + );
363 +
364 + let llm = connect_async_llm_with_ipc(handshake_address, 7, "test-model", &ipc).await;
365 + let stream = llm.generate(sample_generate_request("req-collect", 4)).await.unwrap();
366 + let internal_id = stream.request_id().to_string();
367 + let collected = stream.collect_output().await.unwrap();
368 +
369 + let _ = shutdown_tx.send(());
370 + engine_task.await.unwrap();
371 +
372 + assert_eq!(collected.request_id, internal_id);
373 + assert_eq!(collected.prompt_token_ids, vec![11, 22]);
374 + assert_eq!(collected.token_ids, vec![33, 44]);
375 + assert_eq!(collected.finish_reason, FinishReason::stop_eos());
376 + assert_eq!(collected.prompt_logprobs, Some(prompt_logprobs()));
377 + assert_eq!(
378 + collected.logprobs.as_ref().map(|lp| lp.positions.len()),
379 + Some(2)
380 + );
381 + assert_eq!(
382 + collected.kv_transfer_params,
383 + Some(serde_json::json!({"connector": "x"}))
384 + );
385 +}
386 +
387 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
388 +async fn generate_propagates_unexpected_close_errors() {
389 + let ipc = IpcNamespace::new().unwrap();
390 + let handshake_address = ipc.handshake_endpoint();
391 + let engine_id = b"engine-close".to_vec();
392 +
393 + let (shutdown_tx, engine_task) = spawn_mock_engine_task(
394 + handshake_address.clone(),
395 + engine_id.clone(),
396 + |dealer, push| {
397 + Box::pin(async move {
398 + let add = recv_engine_message(dealer).await;
399 + assert_eq!(add[0].as_ref(), &[0x00]);
400 + let request: EngineCoreRequest = rmp_serde::from_slice(&add[1]).unwrap();
401 +
402 + send_outputs(
403 + push,
404 + EngineCoreOutputs {
405 + finished_requests: Some(BTreeSet::from([request.request_id])),
406 + ..Default::default()
407 + },
408 + )
409 + .await;
410 + })
411 + },
412 + );
413 +
414 + let llm = connect_async_llm_with_ipc(handshake_address, 0, "test-model", &ipc).await;
415 + let mut stream = llm.generate(sample_generate_request("req-close", 1)).await.unwrap();
416 + let internal_id = stream.request_id().to_string();
417 +
418 + let error = stream.next().await.unwrap().unwrap_err();
419 + assert!(matches!(
420 + error,
421 + Error::EngineCoreClient(vllm_engine_core_client::Error::RequestStreamClosed {
422 + request_id
423 + }) if request_id == internal_id
424 + ));
425 + assert!(stream.next().await.is_none());
426 +
427 + let _ = shutdown_tx.send(());
428 + engine_task.await.unwrap();
429 + llm.shutdown().await.unwrap();
430 +}
431 +
432 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
433 +async fn dropping_a_live_generate_stream_triggers_abort() {
434 + let ipc = IpcNamespace::new().unwrap();
435 + let handshake_address = ipc.handshake_endpoint();
436 + let engine_id = b"engine-drop".to_vec();
437 +
438 + let (shutdown_tx, engine_task) = spawn_mock_engine_task(
439 + handshake_address.clone(),
440 + engine_id.clone(),
441 + |dealer, push| {
442 + Box::pin(async move {
443 + let add = recv_engine_message(dealer).await;
444 + assert_eq!(add[0].as_ref(), &[0x00]);
445 + let request: EngineCoreRequest = rmp_serde::from_slice(&add[1]).unwrap();
446 + assert_eq!(request.external_req_id.as_deref(), Some("req-drop"));
447 + assert!(request.request_id.starts_with("req-drop-"));
448 +
449 + send_outputs(
450 + push,
451 + EngineCoreOutputs {
452 + outputs: vec![request_output(&request.request_id, vec![99], None)],
453 + ..Default::default()
454 + },
455 + )
456 + .await;
457 +
458 + let abort =
459 + timeout(Duration::from_secs(1), recv_engine_message(dealer)).await.unwrap();
460 + assert_eq!(abort[0].as_ref(), &[0x01]);
461 + let aborted_ids: Vec<String> = rmp_serde::from_slice(&abort[1]).unwrap();
462 + assert_eq!(aborted_ids, vec![request.request_id]);
463 + })
464 + },
465 + );
466 +
467 + let llm = connect_async_llm_with_ipc(handshake_address, 0, "test-model", &ipc).await;
468 + let mut stream = llm.generate(sample_generate_request("req-drop", 4)).await.unwrap();
469 +
470 + let output = stream.next().await.unwrap().unwrap();
471 + assert_eq!(output.token_ids, vec![99]);
472 + drop(stream);
473 +
474 + let _ = shutdown_tx.send(());
475 + engine_task.await.unwrap();
476 + llm.shutdown().await.unwrap();
477 +}
478 +
479 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
480 +async fn duplicate_external_request_ids_are_randomized_before_reaching_engine_core_client() {
481 + let ipc = IpcNamespace::new().unwrap();
482 + let handshake_address = ipc.handshake_endpoint();
483 + let engine_id = b"engine-dup".to_vec();
484 +
485 + let (shutdown_tx, engine_task) = spawn_mock_engine_task(
486 + handshake_address.clone(),
487 + engine_id.clone(),
488 + |dealer, push| {
489 + Box::pin(async move {
490 + let add_1 = recv_engine_message(dealer).await;
491 + assert_eq!(add_1[0].as_ref(), &[0x00]);
492 + let request_1: EngineCoreRequest = rmp_serde::from_slice(&add_1[1]).unwrap();
493 + assert_eq!(request_1.external_req_id.as_deref(), Some("req-dup"));
494 + assert!(request_1.request_id.starts_with("req-dup-"));
495 +
496 + let add_2 = recv_engine_message(dealer).await;
497 + assert_eq!(add_2[0].as_ref(), &[0x00]);
498 + let request_2: EngineCoreRequest = rmp_serde::from_slice(&add_2[1]).unwrap();
499 + assert_eq!(request_2.external_req_id.as_deref(), Some("req-dup"));
500 + assert!(request_2.request_id.starts_with("req-dup-"));
501 + assert_ne!(request_1.request_id, request_2.request_id);
502 +
503 + send_outputs(
504 + push,
505 + EngineCoreOutputs {
506 + outputs: vec![request_output(
507 + &request_1.request_id,
508 + vec![],
509 + Some(EngineCoreFinishReason::Length),
510 + )],
511 + finished_requests: Some(BTreeSet::from([request_1.request_id.clone()])),
512 + ..Default::default()
513 + },
514 + )
515 + .await;
516 +
517 + send_outputs(
518 + push,
519 + EngineCoreOutputs {
520 + outputs: vec![request_output(
521 + &request_2.request_id,
522 + vec![],
523 + Some(EngineCoreFinishReason::Length),
524 + )],
525 + finished_requests: Some(BTreeSet::from([request_2.request_id])),
526 + ..Default::default()
527 + },
528 + )
529 + .await;
530 + })
531 + },
532 + );
533 +
534 + let llm = connect_async_llm_with_ipc(handshake_address, 0, "test-model", &ipc).await;
535 + let stream_1 = llm.generate(sample_generate_request("req-dup", 1)).await.unwrap();
536 + let stream_2 = llm.generate(sample_generate_request("req-dup", 1)).await.unwrap();
537 + let internal_id_1 = stream_1.request_id().to_string();
538 + let internal_id_2 = stream_2.request_id().to_string();
539 + let collected_1 = stream_1.collect_output().await.unwrap();
540 + let collected_2 = stream_2.collect_output().await.unwrap();
541 + assert_eq!(collected_1.request_id, internal_id_1);
542 + assert_eq!(collected_2.request_id, internal_id_2);
543 + let _ = shutdown_tx.send(());
544 + engine_task.await.unwrap();
545 + llm.shutdown().await.unwrap();
546 +}
547 +
548 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
549 +async fn generate_records_request_metrics_in_prometheus_output() {
550 + let ipc = IpcNamespace::new().unwrap();
551 + let handshake_address = ipc.handshake_endpoint();
552 + let engine_id = b"engine-metrics".to_vec();
553 + let model_name = request_metrics_model_name("metrics-model");
554 +
555 + let (shutdown_tx, engine_task) = spawn_mock_engine_task(
556 + handshake_address.clone(),
557 + engine_id.clone(),
558 + |dealer, push| {
559 + Box::pin(async move {
560 + let add = recv_engine_message(dealer).await;
561 + assert_eq!(add[0].as_ref(), &[0x00]);
562 + let request: EngineCoreRequest = rmp_serde::from_slice(&add[1]).unwrap();
563 +
564 + send_outputs(
565 + push,
566 + EngineCoreOutputs {
567 + engine_index: 4,
568 + timestamp: 10.0,
569 + outputs: vec![EngineCoreOutput {
570 + prefill_stats: Some(PrefillStats {
571 + num_prompt_tokens: 2,
572 + num_computed_tokens: 2,
573 + ..Default::default()
574 + }),
575 + ..request_output_with_events(
576 + &request.request_id,
577 + vec![1],
578 + None,
579 + Some(vec![
580 + EngineCoreEvent {
581 + r#type: EngineCoreEventType::Queued,
582 + timestamp: 8.0,
583 + },
584 + EngineCoreEvent {
585 + r#type: EngineCoreEventType::Scheduled,
586 + timestamp: 9.0,
587 + },
588 + ]),
589 + )
590 + }],
591 + ..Default::default()
592 + },
593 + )
594 + .await;
595 +
596 + send_outputs(
597 + push,
598 + EngineCoreOutputs {
599 + engine_index: 4,
600 + timestamp: 11.5,
601 + outputs: vec![request_output_with_events(
602 + &request.request_id,
603 + vec![2, 3],
604 + Some(EngineCoreFinishReason::Length),
605 + Some(vec![EngineCoreEvent {
606 + r#type: EngineCoreEventType::Preempted,
607 + timestamp: 10.5,
608 + }]),
609 + )],
610 + finished_requests: Some(BTreeSet::from([request.request_id])),
611 + ..Default::default()
612 + },
613 + )
614 + .await;
615 + })
616 + },
617 + );
618 +
619 + let llm = connect_async_llm_with_ipc(handshake_address, 0, &model_name, &ipc).await;
620 + let mut request = sample_generate_request("req-metrics", 8);
621 + request.arrival_time = None;
622 + let mut stream = llm.generate(request).await.unwrap();
623 +
624 + assert_eq!(stream.next().await.unwrap().unwrap().token_ids, vec![1]);
625 + let final_output = stream.next().await.unwrap().unwrap();
626 + assert_eq!(final_output.token_ids, vec![2, 3]);
627 + assert_eq!(final_output.finish_reason, Some(FinishReason::Length));
628 + assert!(stream.next().await.is_none());
629 +
630 + let rendered = METRICS.render().unwrap();
631 + assert!(rendered.contains(&format!(
632 + "vllm:request_success_total{{model_name=\"{model_name}\",engine=\"4\",finished_reason=\"length\"}} 1"
633 + )));
634 + assert!(rendered.contains(&format!(
635 + "vllm:prompt_tokens_total{{model_name=\"{model_name}\",engine=\"4\"}} 2"
636 + )));
637 + assert!(rendered.contains(&format!(
638 + "vllm:prompt_tokens_by_source_total{{model_name=\"{model_name}\",engine=\"4\",source=\"local_compute\"}} 2"
639 + )));
640 + assert!(rendered.contains(&format!(
641 + "vllm:prompt_tokens_by_source_total{{model_name=\"{model_name}\",engine=\"4\",source=\"local_cache_hit\"}} 0"
642 + )));
643 + assert!(rendered.contains(&format!(
644 + "vllm:prompt_tokens_by_source_total{{model_name=\"{model_name}\",engine=\"4\",source=\"external_kv_transfer\"}} 0"
645 + )));
646 + assert!(rendered.contains(&format!(
647 + "vllm:prompt_tokens_cached_total{{model_name=\"{model_name}\",engine=\"4\"}} 0"
648 + )));
649 + assert!(rendered.contains(&format!(
650 + "vllm:generation_tokens_total{{model_name=\"{model_name}\",engine=\"4\"}} 3"
651 + )));
652 + assert!(rendered.contains(&format!(
653 + "vllm:num_preemptions_total{{model_name=\"{model_name}\",engine=\"4\"}} 1"
654 + )));
655 + assert!(rendered.contains(&format!(
656 + "vllm:time_to_first_token_seconds_count{{model_name=\"{model_name}\",engine=\"4\"}} 1"
657 + )));
658 + assert!(rendered.contains(&format!(
659 + "vllm:inter_token_latency_seconds_count{{model_name=\"{model_name}\",engine=\"4\"}} 1"
660 + )));
661 + assert!(rendered.contains(&format!(
662 + "vllm:e2e_request_latency_seconds_count{{model_name=\"{model_name}\",engine=\"4\"}} 1"
663 + )));
664 + assert!(rendered.contains(&format!(
665 + "vllm:request_prompt_tokens_count{{model_name=\"{model_name}\",engine=\"4\"}} 1"
666 + )));
667 + assert!(rendered.contains(&format!(
668 + "vllm:request_generation_tokens_count{{model_name=\"{model_name}\",engine=\"4\"}} 1"
669 + )));
670 + assert!(rendered.contains(&format!(
671 + "vllm:request_prefill_kv_computed_tokens_count{{model_name=\"{model_name}\",engine=\"4\"}} 1"
672 + )));
673 +
674 + let _ = shutdown_tx.send(());
675 + engine_task.await.unwrap();
676 + llm.shutdown().await.unwrap();
677 +}
678 +
679 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
680 +async fn dropping_stream_records_abort_terminal_request_metrics() {
681 + let ipc = IpcNamespace::new().unwrap();
682 + let handshake_address = ipc.handshake_endpoint();
683 + let engine_id = b"engine-metrics-drop".to_vec();
684 + let model_name = request_metrics_model_name("metrics-drop-model");
685 +
686 + let (shutdown_tx, engine_task) = spawn_mock_engine_task(
687 + handshake_address.clone(),
688 + engine_id.clone(),
689 + |dealer, push| {
690 + Box::pin(async move {
691 + let add = recv_engine_message(dealer).await;
692 + assert_eq!(add[0].as_ref(), &[0x00]);
693 + let request: EngineCoreRequest = rmp_serde::from_slice(&add[1]).unwrap();
694 + assert_eq!(request.external_req_id.as_deref(), Some("req-metrics-drop"));
695 + assert!(request.request_id.starts_with("req-metrics-drop-"));
696 +
697 + send_outputs(
698 + push,
699 + EngineCoreOutputs {
700 + engine_index: 5,
701 + timestamp: 10.0,
702 + outputs: vec![request_output_with_events(
703 + &request.request_id,
704 + vec![99],
705 + None,
706 + Some(vec![
707 + EngineCoreEvent {
708 + r#type: EngineCoreEventType::Queued,
709 + timestamp: 8.0,
710 + },
711 + EngineCoreEvent {
712 + r#type: EngineCoreEventType::Scheduled,
713 + timestamp: 9.0,
714 + },
715 + ]),
716 + )],
717 + ..Default::default()
718 + },
719 + )
720 + .await;
721 +
722 + let abort =
723 + timeout(Duration::from_secs(1), recv_engine_message(dealer)).await.unwrap();
724 + assert_eq!(abort[0].as_ref(), &[0x01]);
725 + let aborted_ids: Vec<String> = rmp_serde::from_slice(&abort[1]).unwrap();
726 + assert_eq!(aborted_ids, vec![request.request_id]);
727 + })
728 + },
729 + );
730 +
731 + let llm = connect_async_llm_with_ipc(handshake_address, 0, &model_name, &ipc).await;
732 + let mut request = sample_generate_request("req-metrics-drop", 8);
733 + request.arrival_time = None;
734 + let mut stream = llm.generate(request).await.unwrap();
735 + assert_eq!(stream.next().await.unwrap().unwrap().token_ids, vec![99]);
736 + drop(stream);
737 +
738 + let _ = shutdown_tx.send(());
739 + engine_task.await.unwrap();
740 + let rendered = METRICS.render().unwrap();
741 + assert!(rendered.contains(&format!(
742 + "vllm:request_success_total{{model_name=\"{model_name}\",engine=\"5\",finished_reason=\"abort\"}} 1"
743 + )));
744 + assert!(rendered.contains(&format!(
745 + "vllm:e2e_request_latency_seconds_count{{model_name=\"{model_name}\",engine=\"5\"}} 1"
746 + )));
747 +
748 + llm.shutdown().await.unwrap();
749 +}
rust/src/managed-engine/Cargo.toml new
+18
@@ -0,0 +1,18 @@
1 +[package]
2 +name = "vllm-managed-engine"
3 +version.workspace = true
4 +edition.workspace = true
5 +license.workspace = true
6 +
7 +[dependencies]
8 +anyhow.workspace = true
9 +clap.workspace = true
10 +libc.workspace = true
11 +tokio = { workspace = true, features = ["process"] }
12 +tracing.workspace = true
13 +
14 +[dev-dependencies]
15 +expect-test.workspace = true
16 +
17 +[lints]
18 +workspace = true
rust/src/managed-engine/src/cli.rs new
+349
@@ -0,0 +1,349 @@
1 +use std::collections::HashSet;
2 +use std::ffi::OsString;
3 +
4 +use clap::error::ErrorKind;
5 +use clap::{Args, CommandFactory};
6 +
7 +use crate::{ManagedEngineConfig, allocate_handshake_port};
8 +
9 +/// Managed Python headless-engine CLI arguments.
10 +#[derive(Debug, Clone, Args, PartialEq, Eq)]
11 +pub struct ManagedEngineArgs {
12 + /// Python executable used to launch the managed headless vLLM engine.
13 + #[arg(long, env = "VLLM_RS_PYTHON", default_value = "python3")]
14 + pub python: String,
15 + /// Host/IP used both for the managed-engine handshake endpoint and the
16 + /// frontend-advertised input/output ZMQ socket addresses.
17 + #[arg(
18 + long = "data-parallel-address",
19 + visible_alias = "handshake-host",
20 + default_value = "127.0.0.1"
21 + )]
22 + pub handshake_host: String,
23 + /// Optional TCP port for the managed-engine handshake / data-parallel RPC
24 + /// endpoint.
25 + ///
26 + /// When omitted, the CLI allocates an ephemeral port automatically.
27 + #[arg(
28 + long = "data-parallel-rpc-port",
29 + visible_alias = "handshake-port",
30 + value_parser = clap::value_parser!(u16).range(1..)
31 + )]
32 + pub handshake_port: Option<u16>,
33 + /// Number of data parallel replicas across the whole deployment.
34 + #[arg(long, default_value_t = 1)]
35 + pub data_parallel_size: usize,
36 + /// Number of data parallel replicas to run on this node.
37 + #[arg(long)]
38 + pub data_parallel_size_local: Option<usize>,
39 +
40 + /// Additional arguments forwarded to `python -m vllm.entrypoints.cli.main
41 + /// serve ...`.
42 + ///
43 + /// Arguments after an explicit `--` are forwarded verbatim. Before `--`,
44 + /// `vllm-rs serve` automatically keeps recognized frontend options on
45 + /// the Rust side and forwards everything else to Python.
46 + #[arg(
47 + last = true,
48 + allow_hyphen_values = true,
49 + help_heading = "Passthrough arguments"
50 + )]
51 + pub python_args: Vec<String>,
52 +}
53 +
54 +impl ManagedEngineArgs {
55 + /// Build the handshake address shared by the Rust frontend and managed
56 + /// Python engine.
57 + pub fn handshake_address(&self, handshake_port: u16) -> String {
58 + format!("tcp://{}:{}", self.handshake_host, handshake_port)
59 + }
60 +
61 + /// Resolve the handshake port, either from the CLI argument (if specified)
62 + /// or by allocating a fresh port.
63 + pub fn resolve_handshake_port(&self) -> anyhow::Result<u16> {
64 + self.handshake_port
65 + .map(Ok)
66 + .unwrap_or_else(|| allocate_handshake_port(&self.handshake_host))
67 + }
68 +
69 + /// Build the managed Python-engine spawn configuration.
70 + pub fn into_config(
71 + self,
72 + model: String,
73 + max_model_len: Option<u32>,
74 + handshake_port: u16,
75 + ) -> ManagedEngineConfig {
76 + let mut python_args = self.python_args;
77 + // Manually forward some args to the Python engine.
78 + if let Some(max_model_len) = max_model_len {
79 + python_args.push("--max-model-len".to_string());
80 + python_args.push(max_model_len.to_string());
81 + }
82 + if let Some(data_parallel_size_local) = self.data_parallel_size_local {
83 + python_args.push("--data-parallel-size-local".to_string());
84 + python_args.push(data_parallel_size_local.to_string());
85 + }
86 +
87 + ManagedEngineConfig {
88 + python: self.python,
89 + model,
90 + handshake_host: self.handshake_host,
91 + handshake_port,
92 + data_parallel_size: self.data_parallel_size,
93 + python_args,
94 + }
95 + }
96 +
97 + /// Return the number of engines that the Rust frontend should expect to
98 + /// coordinate with.
99 + fn local_engine_count(&self) -> usize {
100 + self.data_parallel_size_local.unwrap_or(self.data_parallel_size)
101 + }
102 +
103 + /// Return whether the managed Rust frontend only needs to communicate with
104 + /// colocated engines.
105 + pub fn frontend_local_only(&self) -> bool {
106 + self.data_parallel_size_local != Some(0)
107 + && self.local_engine_count() == self.data_parallel_size
108 + }
109 +}
110 +
111 +/// Python `argparse` accepts these multi-character single-dash aliases, but
112 +/// `clap` cannot model them directly.
113 +const PYTHON_MULTI_CHAR_ALIASES: &[(&str, &str)] = &[
114 + ("-asc", "--api-server-count"),
115 + ("-pp", "--pipeline-parallel-size"),
116 + ("-tp", "--tensor-parallel-size"),
117 + ("-dcp", "--decode-context-parallel-size"),
118 + ("-pcp", "--prefill-context-parallel-size"),
119 + ("-dp", "--data-parallel-size"),
120 + ("-dpn", "--data-parallel-rank"),
121 + ("-dpr", "--data-parallel-start-rank"),
122 + ("-dpl", "--data-parallel-size-local"),
123 + ("-dpa", "--data-parallel-address"),
124 + ("-dpp", "--data-parallel-rpc-port"),
125 + ("-dpb", "--data-parallel-backend"),
126 + ("-dph", "--data-parallel-hybrid-lb"),
127 + ("-dpe", "--data-parallel-external-lb"),
128 + ("-ep", "--enable-expert-parallel"),
129 + ("-cc", "--compilation-config"),
130 + ("-ac", "--attention-config"),
131 +];
132 +
133 +/// Repartition managed-engine argv so Rust-owned flags stay before `--`, while
134 +/// everything else is forwarded to Python.
135 +pub fn repartition_managed_engine_args<C>(
136 + args: &[OsString],
137 + subcommand: Option<&str>,
138 +) -> Result<Vec<OsString>, clap::Error>
139 +where
140 + C: CommandFactory,
141 +{
142 + let command = C::command();
143 + let (prefix, real_args, command) = match subcommand {
144 + Some(subcommand) => {
145 + if !matches_subcommand(args, subcommand) {
146 + return Ok(args.to_vec());
147 + };
148 +
149 + let subcommand = command
150 + .find_subcommand(subcommand)
151 + .expect("managed-engine subcommand should exist");
152 +
153 + (args[..2].to_vec(), &args[2..], subcommand)
154 + }
155 + None => {
156 + let Some(program) = args.first() else {
157 + return Ok(args.to_vec());
158 + };
159 +
160 + (vec![program.clone()], &args[1..], &command)
161 + }
162 + };
163 +
164 + let mut repartitioned = prefix;
165 + repartitioned.extend(repartition_real_managed_engine_args(real_args, command)?);
166 + Ok(repartitioned)
167 +}
168 +
169 +fn repartition_real_managed_engine_args(
170 + args: &[OsString],
171 + command: &clap::Command,
172 +) -> Result<Vec<OsString>, clap::Error> {
173 + let Some(model) = args.first() else {
174 + return Ok(args.to_vec());
175 + };
176 +
177 + let model = model.to_string_lossy();
178 + if is_help_flag(&model) {
179 + return Ok(args.to_vec());
180 + }
181 + if model == "--" || is_option_like(&model) {
182 + return Err(build_missing_model_error(command));
183 + }
184 +
185 + let (long_flags, short_flags) = collect_option_names(command);
186 + let (front_args, explicit_passthrough, had_separator) = split_managed_engine_args(&args[1..]);
187 + let normalized_front_args = normalize_python_arg_aliases(front_args);
188 +
189 + let mut frontend_chunks = Vec::new();
190 + let mut python_chunks = Vec::new();
191 + let mut current_chunk = Vec::new();
192 +
193 + for arg in normalized_front_args {
194 + let text = arg.to_string_lossy();
195 + if is_option_like(&text) && !current_chunk.is_empty() {
196 + push_chunk(
197 + &mut frontend_chunks,
198 + &mut python_chunks,
199 + std::mem::take(&mut current_chunk),
200 + &long_flags,
201 + &short_flags,
202 + );
203 + }
204 + current_chunk.push(arg);
205 + }
206 + if !current_chunk.is_empty() {
207 + push_chunk(
208 + &mut frontend_chunks,
209 + &mut python_chunks,
210 + current_chunk,
211 + &long_flags,
212 + &short_flags,
213 + );
214 + }
215 +
216 + let mut repartitioned = vec![args[0].clone()];
217 + repartitioned.extend(frontend_chunks);
218 + if had_separator || !python_chunks.is_empty() || !explicit_passthrough.is_empty() {
219 + repartitioned.push("--".into());
220 + repartitioned.extend(python_chunks);
221 + repartitioned.extend(explicit_passthrough.iter().cloned());
222 + }
223 +
224 + Ok(repartitioned)
225 +}
226 +
227 +fn matches_subcommand(args: &[OsString], subcommand: &str) -> bool {
228 + args.get(1)
229 + .and_then(|arg| arg.to_str())
230 + .is_some_and(|candidate| candidate == subcommand)
231 +}
232 +
233 +fn split_managed_engine_args(args: &[OsString]) -> (&[OsString], &[OsString], bool) {
234 + if let Some(index) = args.iter().position(|arg| arg == "--") {
235 + (&args[..index], &args[index + 1..], true)
236 + } else {
237 + (args, &[], false)
238 + }
239 +}
240 +
241 +fn normalize_python_arg_aliases(args: &[OsString]) -> Vec<OsString> {
242 + args.iter()
243 + .map(|arg| {
244 + let text = arg.to_string_lossy();
245 + normalize_python_multi_char_alias(&text)
246 + .map(Into::into)
247 + .unwrap_or_else(|| arg.clone())
248 + })
249 + .collect()
250 +}
251 +
252 +fn normalize_python_multi_char_alias(arg: &str) -> Option<String> {
253 + find_python_multi_char_alias(arg).map(|canonical| match arg.split_once('=') {
254 + Some((_, value)) => format!("{canonical}={value}"),
255 + None => canonical.to_string(),
256 + })
257 +}
258 +
259 +fn find_python_multi_char_alias(arg: &str) -> Option<&'static str> {
260 + PYTHON_MULTI_CHAR_ALIASES.iter().find_map(|&(alias, canonical)| {
261 + (arg == alias || arg.starts_with(&format!("{alias}="))).then_some(canonical)
262 + })
263 +}
264 +
265 +fn push_chunk(
266 + frontend_chunks: &mut Vec<OsString>,
267 + python_chunks: &mut Vec<OsString>,
268 + chunk: Vec<OsString>,
269 + long_flags: &HashSet<String>,
270 + short_flags: &HashSet<char>,
271 +) {
272 + if chunk_head_is_frontend_owned(&chunk, long_flags, short_flags) {
273 + frontend_chunks.extend(chunk);
274 + } else {
275 + python_chunks.extend(chunk);
276 + }
277 +}
278 +
279 +fn chunk_head_is_frontend_owned(
280 + chunk: &[OsString],
281 + long_flags: &HashSet<String>,
282 + short_flags: &HashSet<char>,
283 +) -> bool {
284 + let Some(head) = chunk.first() else {
285 + return false;
286 + };
287 + let head = head.to_string_lossy();
288 +
289 + if let Some(rest) = head.strip_prefix("--") {
290 + let name = rest.split_once('=').map_or(rest, |(name, _)| name);
291 + return long_flags.contains(name);
292 + }
293 +
294 + let Some(rest) = head.strip_prefix('-') else {
295 + return false;
296 + };
297 + let Some(short) = rest.chars().next() else {
298 + return false;
299 + };
300 + short_flags.contains(&short)
301 +}
302 +
303 +fn collect_option_names(command: &clap::Command) -> (HashSet<String>, HashSet<char>) {
304 + let mut long_flags = HashSet::new();
305 + let mut short_flags = HashSet::new();
306 + for arg in command.get_arguments() {
307 + if let Some(names) = arg.get_long_and_visible_aliases() {
308 + long_flags.extend(names.into_iter().map(str::to_owned));
309 + }
310 + if let Some(short) = arg.get_short() {
311 + short_flags.insert(short);
312 + }
313 + if let Some(short_aliases) = arg.get_visible_short_aliases() {
314 + short_flags.extend(short_aliases);
315 + }
316 + }
317 +
318 + long_flags.insert("help".to_string());
319 + short_flags.insert('h');
320 +
321 + (long_flags, short_flags)
322 +}
323 +
324 +fn is_option_like(arg: &str) -> bool {
325 + if arg == "--" {
326 + return false;
327 + }
328 +
329 + if let Some(rest) = arg.strip_prefix("--") {
330 + return rest.chars().next().is_some_and(char::is_alphabetic);
331 + }
332 +
333 + if let Some(rest) = arg.strip_prefix('-') {
334 + return rest.chars().next().is_some_and(char::is_alphabetic);
335 + }
336 +
337 + false
338 +}
339 +
340 +fn is_help_flag(arg: &str) -> bool {
341 + arg == "-h" || arg == "--help"
342 +}
343 +
344 +fn build_missing_model_error(command: &clap::Command) -> clap::Error {
345 + command.clone().error(
346 + ErrorKind::MissingRequiredArgument,
347 + "the model must appear immediately after the command",
348 + )
349 +}
rust/src/managed-engine/src/lib.rs new
+4
@@ -0,0 +1,4 @@
1 +pub mod cli;
2 +mod process;
3 +
4 +pub use process::{ManagedEngineConfig, ManagedEngineHandle, allocate_handshake_port};
rust/src/managed-engine/src/process.rs new
+263
@@ -0,0 +1,263 @@
1 +use std::io;
2 +use std::net::TcpListener;
3 +use std::process::{Command as StdCommand, ExitStatus, Stdio};
4 +use std::sync::Arc;
5 +use std::sync::atomic::{AtomicBool, Ordering};
6 +use std::time::Duration;
7 +
8 +use anyhow::{Context, Result};
9 +use tokio::process::{Child, Command};
10 +use tokio::sync::Mutex;
11 +use tokio::time::interval;
12 +use tracing::info;
13 +
14 +const CHILD_POLL_INTERVAL: Duration = Duration::from_millis(200);
15 +const MIN_SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(5);
16 +
17 +/// Allocate one ephemeral TCP port for the managed headless-engine handshake on
18 +/// the given host.
19 +pub fn allocate_handshake_port(host: &str) -> Result<u16> {
20 + let listener = TcpListener::bind((host, 0)).context("failed to allocate handshake port")?;
21 + let port = listener
22 + .local_addr()
23 + .context("failed to inspect allocated handshake listener address")?
24 + .port();
25 + Ok(port)
26 +}
27 +
28 +/// Spawn configuration for one managed headless Python vLLM engine.
29 +#[derive(Debug, Clone, PartialEq, Eq)]
30 +pub struct ManagedEngineConfig {
31 + /// Python executable used to launch `vllm.entrypoints.cli.main`.
32 + pub python: String,
33 + /// Model identifier passed to `vllm ... serve <model>`.
34 + pub model: String,
35 + /// Host portion of the headless-engine handshake endpoint.
36 + pub handshake_host: String,
37 + /// Port portion of the headless-engine handshake endpoint.
38 + pub handshake_port: u16,
39 + /// Number of data parallel replicas across the whole deployment.
40 + ///
41 + /// The per-node replica count is forwarded separately in `python_args` as
42 + /// `--data-parallel-size-local`.
43 + pub data_parallel_size: usize,
44 + /// Extra CLI arguments forwarded verbatim to Python vLLM.
45 + pub python_args: Vec<String>,
46 +}
47 +
48 +impl ManagedEngineConfig {
49 + /// Render the handshake address that the Rust frontend should dial.
50 + pub fn handshake_address(&self) -> String {
51 + format!("tcp://{}:{}", self.handshake_host, self.handshake_port)
52 + }
53 +
54 + /// Build the concrete Python command line for the managed headless engine.
55 + pub fn to_command(&self) -> StdCommand {
56 + let mut command = StdCommand::new(&self.python);
57 + command
58 + .arg("-m")
59 + .arg("vllm.entrypoints.cli.main")
60 + .arg("serve")
61 + .arg(&self.model)
62 + .arg("--headless")
63 + .arg("--data-parallel-address")
64 + .arg(&self.handshake_host)
65 + .arg("--data-parallel-rpc-port")
66 + .arg(self.handshake_port.to_string())
67 + .arg("--data-parallel-size")
68 + .arg(self.data_parallel_size.to_string())
69 + .args(&self.python_args);
70 + command
71 + }
72 +}
73 +
74 +/// RAII-style handle for one managed Python headless engine subprocess.
75 +#[derive(Clone)]
76 +pub struct ManagedEngineHandle {
77 + child: Arc<Mutex<Child>>,
78 + shutdown_started: Arc<AtomicBool>,
79 +}
80 +
81 +impl ManagedEngineHandle {
82 + /// Spawn one managed Python headless engine and return a handle for
83 + /// monitoring it.
84 + pub async fn spawn(config: ManagedEngineConfig) -> Result<Self> {
85 + let command = config.to_command();
86 + info!(
87 + handshake_address = %config.handshake_address(),
88 + ?command,
89 + "starting managed Python headless engine"
90 + );
91 +
92 + let mut command = Command::from(command);
93 + command.stdin(Stdio::null()).stdout(Stdio::inherit()).stderr(Stdio::inherit());
94 +
95 + process_group::configure(&mut command);
96 +
97 + let child = command.spawn().context("failed to spawn managed engine")?;
98 +
99 + Ok(Self {
100 + child: Arc::new(Mutex::new(child)),
101 + shutdown_started: Arc::new(AtomicBool::new(false)),
102 + })
103 + }
104 +
105 + /// Poll whether the managed engine has exited yet.
106 + pub async fn try_wait(&self) -> Option<ExitStatus> {
107 + let mut child = self.child.lock().await;
108 + child.try_wait().expect("failed to poll the status of managed engine")
109 + }
110 +
111 + /// Wait until the managed engine exits.
112 + pub async fn wait_for_exit(&self) -> ExitStatus {
113 + let mut interval = interval(CHILD_POLL_INTERVAL);
114 + loop {
115 + interval.tick().await;
116 + if let Some(status) = self.try_wait().await {
117 + return status;
118 + }
119 + }
120 + }
121 +
122 + /// Terminate the managed engine process group and wait for it to stop.
123 + pub async fn shutdown(&self, timeout: Duration) -> Result<()> {
124 + if self.shutdown_started.swap(true, Ordering::SeqCst) {
125 + return Ok(());
126 + }
127 +
128 + let Some(pid) = self.child.lock().await.id() else {
129 + return Ok(());
130 + };
131 +
132 + // Enforce a minimum shutdown timeout to give the engine process enough time to
133 + // clean up.
134 + let shutdown_timeout = std::cmp::max(timeout, MIN_SHUTDOWN_TIMEOUT);
135 +
136 + // First, try to gracefully terminate.
137 + info!(
138 + pid,
139 + ?shutdown_timeout,
140 + "shutting down managed engine with SIGTERM"
141 + );
142 + process_group::terminate(pid)?;
143 +
144 + // Wait for the process to exit on its own.
145 + if tokio::time::timeout(shutdown_timeout, self.wait_for_exit()).await.is_ok() {
146 + return Ok(());
147 + }
148 +
149 + // If it doesn't exit within the timeout, force kill it.
150 + info!(
151 + pid,
152 + "managed engine did not exit within timeout, sending SIGKILL"
153 + );
154 + process_group::kill(pid)?;
155 +
156 + let _ = self.wait_for_exit().await;
157 + Ok(())
158 + }
159 +}
160 +
161 +/// Process group helper functions for managing the Python subprocess and its
162 +/// potential children in a platform-aware way.
163 +mod process_group {
164 + use super::*;
165 +
166 + /// Place the Python child into its own process group so `serve` can tear
167 + /// down the whole subtree rather than just the immediate shell process.
168 + pub fn configure(command: &mut Command) {
169 + unsafe {
170 + command.pre_exec(|| {
171 + if libc::setpgid(0, 0) != 0 {
172 + return Err(io::Error::last_os_error());
173 + }
174 + Ok(())
175 + });
176 + }
177 + }
178 +
179 + /// Send SIGTERM to the managed Python process group.
180 + pub fn terminate(pid: u32) -> Result<()> {
181 + signal(pid, libc::SIGTERM)
182 + }
183 +
184 + /// Send SIGKILL to the managed Python process group.
185 + pub fn kill(pid: u32) -> Result<()> {
186 + signal(pid, libc::SIGKILL)
187 + }
188 +
189 + /// Deliver one signal to the managed Python process group.
190 + fn signal(pid: u32, signal: i32) -> Result<()> {
191 + let rc = unsafe { libc::kill(-(pid as i32), signal) };
192 + if rc == 0 {
193 + return Ok(());
194 + }
195 +
196 + let error = io::Error::last_os_error();
197 + if matches!(error.raw_os_error(), Some(code) if code == libc::ESRCH) {
198 + return Ok(());
199 + }
200 + Err(error).context("failed to signal managed engine process group")
201 + }
202 +}
203 +
204 +#[cfg(test)]
205 +mod tests {
206 + use expect_test::expect;
207 +
208 + use super::{ManagedEngineConfig, allocate_handshake_port};
209 +
210 + #[test]
211 + fn command_snapshot() {
212 + let config = ManagedEngineConfig {
213 + python: "python3".to_string(),
214 + model: "Qwen/Qwen3-0.6B".to_string(),
215 + handshake_host: "127.0.0.1".to_string(),
216 + handshake_port: 62100,
217 + data_parallel_size: 4,
218 + python_args: vec![
219 + "--data-parallel-size-local".to_string(),
220 + "2".to_string(),
221 + "--data-parallel-start-rank".to_string(),
222 + "2".to_string(),
223 + "--dtype".to_string(),
224 + "float16".to_string(),
225 + "--max-model-len".to_string(),
226 + "512".to_string(),
227 + ],
228 + };
229 + let command = config.to_command();
230 + let args = command.get_args().collect::<Vec<_>>();
231 +
232 + expect![[r#"
233 + [
234 + "-m",
235 + "vllm.entrypoints.cli.main",
236 + "serve",
237 + "Qwen/Qwen3-0.6B",
238 + "--headless",
239 + "--data-parallel-address",
240 + "127.0.0.1",
241 + "--data-parallel-rpc-port",
242 + "62100",
243 + "--data-parallel-size",
244 + "4",
245 + "--data-parallel-size-local",
246 + "2",
247 + "--data-parallel-start-rank",
248 + "2",
249 + "--dtype",
250 + "float16",
251 + "--max-model-len",
252 + "512",
253 + ]
254 + "#]]
255 + .assert_debug_eq(&args);
256 + }
257 +
258 + #[test]
259 + fn allocate_handshake_port_returns_non_zero_port() {
260 + let port = allocate_handshake_port("127.0.0.1").unwrap();
261 + assert_ne!(port, 0);
262 + }
263 +}
rust/src/metrics/Cargo.toml new
+11
@@ -0,0 +1,11 @@
1 +[package]
2 +name = "vllm-metrics"
3 +version.workspace = true
4 +edition.workspace = true
5 +license.workspace = true
6 +
7 +[dependencies]
8 +prometheus-client.workspace = true
9 +
10 +[lints]
11 +workspace = true
rust/src/metrics/src/api_server.rs new
+78
@@ -0,0 +1,78 @@
1 +use prometheus_client::encoding::EncodeLabelSet;
2 +use prometheus_client::metrics::family::Family;
3 +use prometheus_client::metrics::histogram::Histogram;
4 +use prometheus_client::registry::Registry;
5 +
6 +use crate::U64Counter;
7 +
8 +const HTTP_REQUEST_DURATION_BUCKETS: [f64; 3] = [0.1, 0.5, 1.0];
9 +const HTTP_REQUEST_DURATION_HIGHR_BUCKETS: [f64; 21] = [
10 + 0.01, 0.025, 0.05, 0.075, 0.1, 0.25, 0.5, 0.75, 1.0, 1.5, 2.0, 2.5, 3.0, 3.5, 4.0, 4.5, 5.0,
11 + 7.5, 10.0, 30.0, 60.0,
12 +];
13 +
14 +fn http_request_duration_histogram() -> Histogram {
15 + Histogram::new(HTTP_REQUEST_DURATION_BUCKETS.iter().copied())
16 +}
17 +
18 +fn http_request_duration_highr_histogram() -> Histogram {
19 + Histogram::new(HTTP_REQUEST_DURATION_HIGHR_BUCKETS.iter().copied())
20 +}
21 +
22 +#[derive(Clone, Debug, Hash, PartialEq, Eq, EncodeLabelSet)]
23 +pub struct HttpRequestLabels {
24 + pub method: String,
25 + pub status: &'static str,
26 + pub handler: String,
27 +}
28 +
29 +#[derive(Clone, Debug, Hash, PartialEq, Eq, EncodeLabelSet)]
30 +pub struct HttpHandlerLabels {
31 + pub method: String,
32 + pub handler: String,
33 +}
34 +
35 +pub(crate) type HttpRequestCounterFamily = Family<HttpRequestLabels, U64Counter>;
36 +pub(crate) type HttpHandlerHistogramFamily =
37 + Family<HttpHandlerLabels, Histogram, fn() -> Histogram>;
38 +
39 +/// API-server Prometheus families exported from the HTTP middleware layer.
40 +pub struct ApiServerMetrics {
41 + pub http_requests: HttpRequestCounterFamily,
42 + pub http_request_duration_seconds: HttpHandlerHistogramFamily,
43 + pub http_request_duration_highr_seconds: Histogram,
44 +}
45 +
46 +impl ApiServerMetrics {
47 + /// Register the API-server metric families into the shared registry.
48 + pub(crate) fn register(registry: &mut Registry) -> Self {
49 + let http_requests = HttpRequestCounterFamily::default();
50 + registry.register(
51 + "http_requests",
52 + "Total number of HTTP requests by method, status, and handler.",
53 + http_requests.clone(),
54 + );
55 +
56 + let http_request_duration_seconds = HttpHandlerHistogramFamily::new_with_constructor(
57 + http_request_duration_histogram as fn() -> Histogram,
58 + );
59 + registry.register(
60 + "http_request_duration_seconds",
61 + "Duration of HTTP requests in seconds grouped by method and handler.",
62 + http_request_duration_seconds.clone(),
63 + );
64 +
65 + let http_request_duration_highr_seconds = http_request_duration_highr_histogram();
66 + registry.register(
67 + "http_request_duration_highr_seconds",
68 + "High-resolution duration of HTTP requests in seconds.",
69 + http_request_duration_highr_seconds.clone(),
70 + );
71 +
72 + Self {
73 + http_requests,
74 + http_request_duration_seconds,
75 + http_request_duration_highr_seconds,
76 + }
77 + }
78 +}
rust/src/metrics/src/lib.rs new
+76
@@ -0,0 +1,76 @@
1 +use std::fmt;
2 +use std::sync::LazyLock;
3 +use std::sync::atomic::AtomicU64;
4 +
5 +use prometheus_client::encoding::text::encode;
6 +use prometheus_client::metrics::counter::Counter;
7 +use prometheus_client::metrics::family::Family;
8 +use prometheus_client::metrics::gauge::Gauge;
9 +use prometheus_client::metrics::histogram::Histogram;
10 +use prometheus_client::registry::Registry;
11 +
12 +mod api_server;
13 +mod request;
14 +mod scheduler;
15 +
16 +pub use api_server::*;
17 +pub use request::*;
18 +pub use scheduler::*;
19 +
20 +// Note: `prometheus-client` appends the `_total` suffix automatically when
21 +// encoding counters, so all counter family registration names in this crate
22 +// must use the base metric name without a trailing `_total`.
23 +pub type U64Counter = Counter<u64, AtomicU64>;
24 +pub type U64Gauge = Gauge<u64, AtomicU64>;
25 +pub type F64Gauge = Gauge<f64, AtomicU64>;
26 +pub(crate) type HistogramFamily = Family<EngineLabels, Histogram, fn() -> Histogram>;
27 +
28 +/// Shared Prometheus registry for frontend metrics.
29 +///
30 +/// Original Python definition:
31 +/// <https://github.com/vllm-project/vllm/blob/bc2c0c86efb28e77677a3cfb8687e976914a313a/vllm/v1/metrics/loggers.py#L389-L1004>
32 +pub struct Metrics {
33 + registry: Registry,
34 + pub scheduler: SchedulerMetrics,
35 + pub request: RequestMetrics,
36 + pub api_server: ApiServerMetrics,
37 +}
38 +
39 +impl Metrics {
40 + /// Construct a new metrics registry.
41 + pub fn new() -> Self {
42 + let mut registry = Registry::default();
43 + let scheduler = SchedulerMetrics::register(&mut registry);
44 + let request = RequestMetrics::register(&mut registry);
45 + let api_server = ApiServerMetrics::register(&mut registry);
46 +
47 + Self {
48 + registry,
49 + scheduler,
50 + request,
51 + api_server,
52 + }
53 + }
54 +
55 + /// Render the current metrics registry into Prometheus/OpenMetrics text
56 + /// format.
57 + pub fn render(&self) -> Result<String, fmt::Error> {
58 + let mut output = String::new();
59 + encode(&mut output, &self.registry)?;
60 + Ok(output)
61 + }
62 +
63 + /// Return the registry owned by this metrics object.
64 + pub fn registry(&self) -> &Registry {
65 + &self.registry
66 + }
67 +}
68 +
69 +impl Default for Metrics {
70 + fn default() -> Self {
71 + Self::new()
72 + }
73 +}
74 +
75 +/// Process-global metrics registry shared by the frontend crates.
76 +pub static METRICS: LazyLock<Metrics> = LazyLock::new(Metrics::new);
rust/src/metrics/src/request.rs new
+298
@@ -0,0 +1,298 @@
1 +use prometheus_client::encoding::EncodeLabelSet;
2 +use prometheus_client::metrics::family::Family;
3 +use prometheus_client::metrics::histogram::Histogram;
4 +use prometheus_client::registry::Registry;
5 +
6 +use crate::{EngineLabels, HistogramFamily, U64Counter};
7 +
8 +const TTFT_BUCKETS: [f64; 22] = [
9 + 0.001, 0.005, 0.01, 0.02, 0.04, 0.06, 0.08, 0.1, 0.25, 0.5, 0.75, 1.0, 2.5, 5.0, 7.5, 10.0,
10 + 20.0, 40.0, 80.0, 160.0, 640.0, 2560.0,
11 +];
12 +const ITL_BUCKETS: [f64; 19] = [
13 + 0.01, 0.025, 0.05, 0.075, 0.1, 0.15, 0.2, 0.3, 0.4, 0.5, 0.75, 1.0, 2.5, 5.0, 7.5, 10.0, 20.0,
14 + 40.0, 80.0,
15 +];
16 +const REQUEST_LATENCY_BUCKETS: [f64; 21] = [
17 + 0.3, 0.5, 0.8, 1.0, 1.5, 2.0, 2.5, 5.0, 10.0, 15.0, 20.0, 30.0, 40.0, 50.0, 60.0, 120.0, 240.0,
18 + 480.0, 960.0, 1920.0, 7680.0,
19 +];
20 +const REQUEST_PARAMS_N_BUCKETS: [f64; 5] = [1.0, 2.0, 5.0, 10.0, 20.0];
21 +
22 +fn build_1_2_5_buckets(max_value: u32) -> Vec<f64> {
23 + let mut buckets = Vec::new();
24 + let mut exponent = 0;
25 + loop {
26 + for mantissa in [1_u32, 2, 5] {
27 + let value = mantissa * 10_u32.pow(exponent);
28 + if value <= max_value {
29 + buckets.push(value as f64);
30 + } else {
31 + if buckets.last().copied() != Some(max_value as f64) {
32 + buckets.push(max_value as f64);
33 + }
34 + return buckets;
35 + }
36 + }
37 + exponent += 1;
38 + }
39 +}
40 +
41 +fn time_to_first_token_histogram() -> Histogram {
42 + Histogram::new(TTFT_BUCKETS.iter().copied())
43 +}
44 +
45 +fn inter_token_latency_histogram() -> Histogram {
46 + Histogram::new(ITL_BUCKETS.iter().copied())
47 +}
48 +
49 +fn request_time_per_output_token_histogram() -> Histogram {
50 + Histogram::new(ITL_BUCKETS.iter().copied())
51 +}
52 +
53 +fn request_latency_histogram() -> Histogram {
54 + Histogram::new(REQUEST_LATENCY_BUCKETS.iter().copied())
55 +}
56 +
57 +fn request_token_count_histogram() -> Histogram {
58 + // TODO: determine max value based on `max_model_len`.
59 + Histogram::new(build_1_2_5_buckets(131_072))
60 +}
61 +
62 +fn request_params_n_histogram() -> Histogram {
63 + Histogram::new(REQUEST_PARAMS_N_BUCKETS.iter().copied())
64 +}
65 +
66 +#[derive(Clone, Debug, Hash, PartialEq, Eq, EncodeLabelSet)]
67 +pub struct FinishedReasonLabels {
68 + pub model_name: String,
69 + pub engine: u32,
70 + pub finished_reason: &'static str,
71 +}
72 +
73 +#[derive(Clone, Debug, Hash, PartialEq, Eq, EncodeLabelSet)]
74 +pub struct PromptTokenSourceLabels {
75 + pub model_name: String,
76 + pub engine: u32,
77 + pub source: &'static str,
78 +}
79 +
80 +pub(crate) type FinishedReasonCounterFamily = Family<FinishedReasonLabels, U64Counter>;
81 +pub(crate) type PromptTokenSourceCounterFamily = Family<PromptTokenSourceLabels, U64Counter>;
82 +
83 +/// Request-lifecycle Prometheus families exported from the `llm` layer.
84 +pub struct RequestMetrics {
85 + // Request-derived counters.
86 + pub num_preemptions: Family<EngineLabels, U64Counter>,
87 + pub prompt_tokens: Family<EngineLabels, U64Counter>,
88 + pub prompt_tokens_by_source: PromptTokenSourceCounterFamily,
89 + pub prompt_tokens_cached: Family<EngineLabels, U64Counter>,
90 + pub generation_tokens: Family<EngineLabels, U64Counter>,
91 +
92 + // We intentionally don't support iteration-level histograms for now, since it seems to make
93 + // more sense if the engine maintains these metrics and frontend simply forwards.
94 + //
95 + // pub iteration_tokens_total: HistogramFamily,
96 +
97 + // Request lifecycle counters and histograms.
98 + pub request_success: FinishedReasonCounterFamily,
99 + pub request_prompt_tokens: HistogramFamily,
100 + pub request_generation_tokens: HistogramFamily,
101 + pub request_max_num_generation_tokens: HistogramFamily,
102 + pub request_params_max_tokens: HistogramFamily,
103 + pub request_params_n: HistogramFamily,
104 + pub request_prefill_kv_computed_tokens: HistogramFamily,
105 + pub time_to_first_token_seconds: HistogramFamily,
106 + pub inter_token_latency_seconds: HistogramFamily,
107 + pub e2e_request_latency_seconds: HistogramFamily,
108 + pub request_queue_time_seconds: HistogramFamily,
109 + pub request_prefill_time_seconds: HistogramFamily,
110 + pub request_decode_time_seconds: HistogramFamily,
111 + pub request_inference_time_seconds: HistogramFamily,
112 + pub request_time_per_output_token_seconds: HistogramFamily,
113 +}
114 +
115 +impl RequestMetrics {
116 + /// Register the request-oriented metric families into the shared registry.
117 + pub(crate) fn register(registry: &mut Registry) -> Self {
118 + // Request-derived counters.
119 + let num_preemptions = Family::default();
120 + registry.register(
121 + "vllm:num_preemptions",
122 + "Cumulative number of preemption events.",
123 + num_preemptions.clone(),
124 + );
125 +
126 + let prompt_tokens = Family::default();
127 + registry.register(
128 + "vllm:prompt_tokens",
129 + "Number of prefill tokens processed.",
130 + prompt_tokens.clone(),
131 + );
132 +
133 + let prompt_tokens_by_source = Family::default();
134 + registry.register(
135 + "vllm:prompt_tokens_by_source",
136 + "Number of prompt tokens by source.",
137 + prompt_tokens_by_source.clone(),
138 + );
139 +
140 + let prompt_tokens_cached = Family::default();
141 + registry.register(
142 + "vllm:prompt_tokens_cached",
143 + "Number of prompt tokens with prefix cache hits.",
144 + prompt_tokens_cached.clone(),
145 + );
146 +
147 + let generation_tokens = Family::default();
148 + registry.register(
149 + "vllm:generation_tokens",
150 + "Number of generation tokens processed.",
151 + generation_tokens.clone(),
152 + );
153 +
154 + // Request lifecycle counters and histograms.
155 + let request_success = Family::default();
156 + registry.register(
157 + "vllm:request_success",
158 + "Count of successfully processed requests.",
159 + request_success.clone(),
160 + );
161 +
162 + let request_prompt_tokens =
163 + Family::new_with_constructor(request_token_count_histogram as fn() -> Histogram);
164 + registry.register(
165 + "vllm:request_prompt_tokens",
166 + "Number of prefill tokens processed.",
167 + request_prompt_tokens.clone(),
168 + );
169 +
170 + let request_generation_tokens =
171 + Family::new_with_constructor(request_token_count_histogram as fn() -> Histogram);
172 + registry.register(
173 + "vllm:request_generation_tokens",
174 + "Number of generation tokens processed.",
175 + request_generation_tokens.clone(),
176 + );
177 +
178 + let request_max_num_generation_tokens =
179 + Family::new_with_constructor(request_token_count_histogram as fn() -> Histogram);
180 + registry.register(
181 + "vllm:request_max_num_generation_tokens",
182 + "Histogram of maximum number of requested generation tokens.",
183 + request_max_num_generation_tokens.clone(),
184 + );
185 +
186 + let request_params_max_tokens =
187 + Family::new_with_constructor(request_token_count_histogram as fn() -> Histogram);
188 + registry.register(
189 + "vllm:request_params_max_tokens",
190 + "Histogram of the max_tokens request parameter.",
191 + request_params_max_tokens.clone(),
192 + );
193 +
194 + let request_params_n =
195 + Family::new_with_constructor(request_params_n_histogram as fn() -> Histogram);
196 + registry.register(
197 + "vllm:request_params_n",
198 + "Histogram of the n request parameter.",
199 + request_params_n.clone(),
200 + );
201 +
202 + let request_prefill_kv_computed_tokens =
203 + Family::new_with_constructor(request_token_count_histogram as fn() -> Histogram);
204 + registry.register(
205 + "vllm:request_prefill_kv_computed_tokens",
206 + "Histogram of new KV tokens computed during prefill (excluding cached tokens).",
207 + request_prefill_kv_computed_tokens.clone(),
208 + );
209 +
210 + let time_to_first_token_seconds =
211 + Family::new_with_constructor(time_to_first_token_histogram as fn() -> Histogram);
212 + registry.register(
213 + "vllm:time_to_first_token_seconds",
214 + "Histogram of time to first token in seconds.",
215 + time_to_first_token_seconds.clone(),
216 + );
217 +
218 + let inter_token_latency_seconds =
219 + Family::new_with_constructor(inter_token_latency_histogram as fn() -> Histogram);
220 + registry.register(
221 + "vllm:inter_token_latency_seconds",
222 + "Histogram of inter-token latency in seconds.",
223 + inter_token_latency_seconds.clone(),
224 + );
225 +
226 + let e2e_request_latency_seconds =
227 + Family::new_with_constructor(request_latency_histogram as fn() -> Histogram);
228 + registry.register(
229 + "vllm:e2e_request_latency_seconds",
230 + "Histogram of e2e request latency in seconds.",
231 + e2e_request_latency_seconds.clone(),
232 + );
233 +
234 + let request_queue_time_seconds =
235 + Family::new_with_constructor(request_latency_histogram as fn() -> Histogram);
236 + registry.register(
237 + "vllm:request_queue_time_seconds",
238 + "Histogram of time spent in WAITING phase for request.",
239 + request_queue_time_seconds.clone(),
240 + );
241 +
242 + let request_prefill_time_seconds =
243 + Family::new_with_constructor(request_latency_histogram as fn() -> Histogram);
244 + registry.register(
245 + "vllm:request_prefill_time_seconds",
246 + "Histogram of time spent in PREFILL phase for request.",
247 + request_prefill_time_seconds.clone(),
248 + );
249 +
250 + let request_decode_time_seconds =
251 + Family::new_with_constructor(request_latency_histogram as fn() -> Histogram);
252 + registry.register(
253 + "vllm:request_decode_time_seconds",
254 + "Histogram of time spent in DECODE phase for request.",
255 + request_decode_time_seconds.clone(),
256 + );
257 +
258 + let request_inference_time_seconds =
259 + Family::new_with_constructor(request_latency_histogram as fn() -> Histogram);
260 + registry.register(
261 + "vllm:request_inference_time_seconds",
262 + "Histogram of time spent in RUNNING phase for request.",
263 + request_inference_time_seconds.clone(),
264 + );
265 +
266 + let request_time_per_output_token_seconds = Family::new_with_constructor(
267 + request_time_per_output_token_histogram as fn() -> Histogram,
268 + );
269 + registry.register(
270 + "vllm:request_time_per_output_token_seconds",
271 + "Histogram of time_per_output_token_seconds per request.",
272 + request_time_per_output_token_seconds.clone(),
273 + );
274 +
275 + Self {
276 + num_preemptions,
277 + prompt_tokens,
278 + prompt_tokens_by_source,
279 + prompt_tokens_cached,
280 + generation_tokens,
281 + request_success,
282 + request_prompt_tokens,
283 + request_generation_tokens,
284 + request_max_num_generation_tokens,
285 + request_params_max_tokens,
286 + request_params_n,
287 + request_prefill_kv_computed_tokens,
288 + time_to_first_token_seconds,
289 + inter_token_latency_seconds,
290 + e2e_request_latency_seconds,
291 + request_queue_time_seconds,
292 + request_prefill_time_seconds,
293 + request_decode_time_seconds,
294 + request_inference_time_seconds,
295 + request_time_per_output_token_seconds,
296 + }
297 + }
298 +}
rust/src/metrics/src/scheduler.rs new
+272
@@ -0,0 +1,272 @@
1 +use prometheus_client::encoding::EncodeLabelSet;
2 +use prometheus_client::metrics::family::Family;
3 +use prometheus_client::metrics::histogram::Histogram;
4 +use prometheus_client::registry::Registry;
5 +
6 +use crate::{F64Gauge, HistogramFamily, U64Counter, U64Gauge};
7 +
8 +const KV_CACHE_RESIDENCY_BUCKETS: [f64; 21] = [
9 + 0.001, 0.002, 0.005, 0.01, 0.02, 0.05, 0.1, 0.2, 0.5, 1.0, 2.0, 5.0, 10.0, 20.0, 30.0, 60.0,
10 + 120.0, 300.0, 600.0, 1200.0, 1800.0,
11 +];
12 +
13 +fn kv_block_lifetime_histogram() -> Histogram {
14 + Histogram::new(KV_CACHE_RESIDENCY_BUCKETS.iter().copied())
15 +}
16 +
17 +fn kv_block_idle_before_evict_histogram() -> Histogram {
18 + Histogram::new(KV_CACHE_RESIDENCY_BUCKETS.iter().copied())
19 +}
20 +
21 +fn kv_block_reuse_gap_histogram() -> Histogram {
22 + Histogram::new(KV_CACHE_RESIDENCY_BUCKETS.iter().copied())
23 +}
24 +
25 +#[derive(Clone, Debug, Hash, PartialEq, Eq, EncodeLabelSet)]
26 +pub struct EngineLabels {
27 + pub model_name: String,
28 + pub engine: u32,
29 +}
30 +
31 +#[derive(Clone, Debug, Hash, PartialEq, Eq, EncodeLabelSet)]
32 +pub struct EnginePositionLabels {
33 + pub model_name: String,
34 + pub engine: u32,
35 + pub position: u32,
36 +}
37 +
38 +#[derive(Clone, Debug, Hash, PartialEq, Eq, EncodeLabelSet)]
39 +pub struct WaitingReasonLabels {
40 + pub model_name: String,
41 + pub engine: u32,
42 + pub reason: &'static str,
43 +}
44 +
45 +/// Scheduler/batch-scoped Prometheus families exported from `SchedulerStats`.
46 +pub struct SchedulerMetrics {
47 + // Scheduler state gauges.
48 + pub scheduler_running: Family<EngineLabels, U64Gauge>,
49 + pub scheduler_waiting: Family<EngineLabels, U64Gauge>,
50 + pub scheduler_waiting_by_reason: Family<WaitingReasonLabels, U64Gauge>,
51 + pub kv_cache_usage: Family<EngineLabels, F64Gauge>,
52 +
53 + // Prefix-cache counters, including the connector-backed external cache path.
54 + pub prefix_cache_queries: Family<EngineLabels, U64Counter>,
55 + pub prefix_cache_hits: Family<EngineLabels, U64Counter>,
56 + pub external_prefix_cache_queries: Family<EngineLabels, U64Counter>,
57 + pub external_prefix_cache_hits: Family<EngineLabels, U64Counter>,
58 +
59 + // Speculative decoding counters.
60 + pub spec_decode_num_drafts: Family<EngineLabels, U64Counter>,
61 + pub spec_decode_num_draft_tokens: Family<EngineLabels, U64Counter>,
62 + pub spec_decode_num_accepted_tokens: Family<EngineLabels, U64Counter>,
63 + pub spec_decode_num_accepted_tokens_per_pos: Family<EnginePositionLabels, U64Counter>,
64 +
65 + // Per-engine performance / MFU counters.
66 + pub estimated_flops_per_gpu: Family<EngineLabels, U64Counter>,
67 + pub estimated_read_bytes_per_gpu: Family<EngineLabels, U64Counter>,
68 + pub estimated_write_bytes_per_gpu: Family<EngineLabels, U64Counter>,
69 +
70 + // Sampled KV-cache residency histograms.
71 + pub kv_block_lifetime_seconds: HistogramFamily,
72 + pub kv_block_idle_before_evict_seconds: HistogramFamily,
73 + pub kv_block_reuse_gap_seconds: HistogramFamily,
74 +}
75 +
76 +impl SchedulerMetrics {
77 + /// Register the scheduler-oriented metric families into the shared
78 + /// registry.
79 + pub(crate) fn register(registry: &mut Registry) -> Self {
80 + // Scheduler state gauges.
81 + let scheduler_running = Family::default();
82 + registry.register(
83 + "vllm:num_requests_running",
84 + "Number of requests in model execution batches",
85 + scheduler_running.clone(),
86 + );
87 +
88 + let scheduler_waiting = Family::default();
89 + registry.register(
90 + "vllm:num_requests_waiting",
91 + "Number of requests waiting to be processed",
92 + scheduler_waiting.clone(),
93 + );
94 +
95 + let scheduler_waiting_by_reason = Family::default();
96 + registry.register(
97 + "vllm:num_requests_waiting_by_reason",
98 + "Number of waiting requests by reason. \
99 + Reason labels: 'capacity' = waiting for scheduling capacity; \
100 + 'deferred' = deferred by transient constraints (LoRA budget, KV transfer, \
101 + blocked status). Sum of all reasons equals vllm:num_requests_waiting.",
102 + scheduler_waiting_by_reason.clone(),
103 + );
104 +
105 + let kv_cache_usage = Family::default();
106 + registry.register(
107 + "vllm:kv_cache_usage_perc",
108 + "KV-cache usage. 1 means 100 percent usage",
109 + kv_cache_usage.clone(),
110 + );
111 +
112 + // Prefix-cache counters, including the connector-backed external cache path.
113 + let prefix_cache_queries = Family::default();
114 + registry.register(
115 + "vllm:prefix_cache_queries",
116 + "Prefix cache queries, in terms of number of queried tokens",
117 + prefix_cache_queries.clone(),
118 + );
119 +
120 + let prefix_cache_hits = Family::default();
121 + registry.register(
122 + "vllm:prefix_cache_hits",
123 + "Prefix cache hits, in terms of number of cached tokens.",
124 + prefix_cache_hits.clone(),
125 + );
126 +
127 + let external_prefix_cache_queries = Family::default();
128 + registry.register(
129 + "vllm:external_prefix_cache_queries",
130 + "External prefix cache queries from KV connector cross-instance cache sharing, in terms of number of queried tokens.",
131 + external_prefix_cache_queries.clone(),
132 + );
133 +
134 + let external_prefix_cache_hits = Family::default();
135 + registry.register(
136 + "vllm:external_prefix_cache_hits",
137 + "External prefix cache hits from KV connector cross-instance cache sharing, in terms of number of cached tokens.",
138 + external_prefix_cache_hits.clone(),
139 + );
140 +
141 + // Speculative decoding counters.
142 + let spec_decode_num_drafts = Family::default();
143 + registry.register(
144 + "vllm:spec_decode_num_drafts",
145 + "Number of spec decoding drafts.",
146 + spec_decode_num_drafts.clone(),
147 + );
148 +
149 + let spec_decode_num_draft_tokens = Family::default();
150 + registry.register(
151 + "vllm:spec_decode_num_draft_tokens",
152 + "Number of draft tokens.",
153 + spec_decode_num_draft_tokens.clone(),
154 + );
155 +
156 + let spec_decode_num_accepted_tokens = Family::default();
157 + registry.register(
158 + "vllm:spec_decode_num_accepted_tokens",
159 + "Number of accepted tokens.",
160 + spec_decode_num_accepted_tokens.clone(),
161 + );
162 +
163 + let spec_decode_num_accepted_tokens_per_pos = Family::default();
164 + registry.register(
165 + "vllm:spec_decode_num_accepted_tokens_per_pos",
166 + "Accepted tokens per draft position.",
167 + spec_decode_num_accepted_tokens_per_pos.clone(),
168 + );
169 +
170 + // Per-engine performance / MFU counters.
171 + let estimated_flops_per_gpu = Family::default();
172 + registry.register(
173 + "vllm:estimated_flops_per_gpu",
174 + "Estimated number of floating point operations per GPU (for Model Flops Utilization calculations).",
175 + estimated_flops_per_gpu.clone(),
176 + );
177 +
178 + let estimated_read_bytes_per_gpu = Family::default();
179 + registry.register(
180 + "vllm:estimated_read_bytes_per_gpu",
181 + "Estimated number of bytes read from memory per GPU (for Model Flops Utilization calculations).",
182 + estimated_read_bytes_per_gpu.clone(),
183 + );
184 +
185 + let estimated_write_bytes_per_gpu = Family::default();
186 + registry.register(
187 + "vllm:estimated_write_bytes_per_gpu",
188 + "Estimated number of bytes written to memory per GPU (for Model Flops Utilization calculations).",
189 + estimated_write_bytes_per_gpu.clone(),
190 + );
191 +
192 + // Sampled KV-cache residency histograms.
193 + let kv_block_lifetime_seconds =
194 + Family::new_with_constructor(kv_block_lifetime_histogram as fn() -> Histogram);
195 + registry.register(
196 + "vllm:kv_block_lifetime_seconds",
197 + "Histogram of KV cache block lifetime from allocation to eviction. Sampled metrics (controlled by --kv-cache-metrics-sample).",
198 + kv_block_lifetime_seconds.clone(),
199 + );
200 +
201 + let kv_block_idle_before_evict_seconds =
202 + Family::new_with_constructor(kv_block_idle_before_evict_histogram as fn() -> Histogram);
203 + registry.register(
204 + "vllm:kv_block_idle_before_evict_seconds",
205 + "Histogram of idle time before KV cache block eviction. Sampled metrics (controlled by --kv-cache-metrics-sample).",
206 + kv_block_idle_before_evict_seconds.clone(),
207 + );
208 +
209 + let kv_block_reuse_gap_seconds =
210 + Family::new_with_constructor(kv_block_reuse_gap_histogram as fn() -> Histogram);
211 + registry.register(
212 + "vllm:kv_block_reuse_gap_seconds",
213 + "Histogram of time gaps between consecutive KV cache block accesses. Only the most recent accesses are recorded (ring buffer). Sampled metrics (controlled by --kv-cache-metrics-sample).",
214 + kv_block_reuse_gap_seconds.clone(),
215 + );
216 +
217 + Self {
218 + scheduler_running,
219 + scheduler_waiting,
220 + scheduler_waiting_by_reason,
221 + kv_cache_usage,
222 + prefix_cache_queries,
223 + prefix_cache_hits,
224 + external_prefix_cache_queries,
225 + external_prefix_cache_hits,
226 + spec_decode_num_drafts,
227 + spec_decode_num_draft_tokens,
228 + spec_decode_num_accepted_tokens,
229 + spec_decode_num_accepted_tokens_per_pos,
230 + estimated_flops_per_gpu,
231 + estimated_read_bytes_per_gpu,
232 + estimated_write_bytes_per_gpu,
233 + kv_block_lifetime_seconds,
234 + kv_block_idle_before_evict_seconds,
235 + kv_block_reuse_gap_seconds,
236 + }
237 + }
238 +}
239 +
240 +#[cfg(test)]
241 +mod tests {
242 + use crate::{EngineLabels, Metrics};
243 +
244 + #[test]
245 + fn perf_counters_render_with_a_single_total_suffix() {
246 + let metrics = Metrics::new();
247 + let labels = EngineLabels {
248 + model_name: "model".to_string(),
249 + engine: 0,
250 + };
251 +
252 + metrics.scheduler.estimated_flops_per_gpu.get_or_create(&labels).inc();
253 + metrics.scheduler.estimated_read_bytes_per_gpu.get_or_create(&labels).inc();
254 + metrics.scheduler.estimated_write_bytes_per_gpu.get_or_create(&labels).inc();
255 +
256 + let rendered = metrics.render().unwrap();
257 + assert!(
258 + rendered.contains(
259 + "vllm:estimated_flops_per_gpu_total{model_name=\"model\",engine=\"0\"} 1"
260 + )
261 + );
262 + assert!(rendered.contains(
263 + "vllm:estimated_read_bytes_per_gpu_total{model_name=\"model\",engine=\"0\"} 1"
264 + ));
265 + assert!(rendered.contains(
266 + "vllm:estimated_write_bytes_per_gpu_total{model_name=\"model\",engine=\"0\"} 1"
267 + ));
268 + assert!(!rendered.contains("vllm:estimated_flops_per_gpu_total_total"));
269 + assert!(!rendered.contains("vllm:estimated_read_bytes_per_gpu_total_total"));
270 + assert!(!rendered.contains("vllm:estimated_write_bytes_per_gpu_total_total"));
271 + }
272 +}
rust/src/reasoning-parser/Cargo.toml new
+12
@@ -0,0 +1,12 @@
1 +[package]
2 +name = "vllm-reasoning-parser"
3 +version.workspace = true
4 +edition.workspace = true
5 +license.workspace = true
6 +
7 +[dependencies]
8 +thiserror.workspace = true
9 +vllm-tokenizer.workspace = true
10 +
11 +[lints]
12 +workspace = true
rust/src/reasoning-parser/src/cohere_cmd.rs new
+45
@@ -0,0 +1,45 @@
1 +use vllm_tokenizer::DynTokenizer;
2 +
3 +use super::{DelimitedReasoningParser, ReasoningDelta, ReasoningParser, Result};
4 +
5 +/// Reasoning parser for Cohere Command models that use explicit START/END tags.
6 +pub struct CohereCmdReasoningParser {
7 + inner: DelimitedReasoningParser,
8 +}
9 +
10 +impl CohereCmdReasoningParser {
11 + /// Create a Cohere Command parser backed by the shared delimited state
12 + /// machine.
13 + pub fn new(tokenizer: DynTokenizer) -> Result<Self> {
14 + Ok(Self {
15 + inner: DelimitedReasoningParser::new(
16 + tokenizer,
17 + "<|START_THINKING|>",
18 + "<|END_THINKING|>",
19 + false,
20 + )?,
21 + })
22 + }
23 +}
24 +
25 +impl ReasoningParser for CohereCmdReasoningParser {
26 + fn create(tokenizer: DynTokenizer) -> Result<Box<dyn ReasoningParser>>
27 + where
28 + Self: Sized + 'static,
29 + {
30 + Ok(Box::new(Self::new(tokenizer)?))
31 + }
32 +
33 + fn initialize(&mut self, prompt_token_ids: &[u32]) -> Result<()> {
34 + self.inner.initialize(prompt_token_ids);
35 + Ok(())
36 + }
37 +
38 + fn push(&mut self, delta: &str) -> Result<ReasoningDelta> {
39 + Ok(self.inner.push(delta))
40 + }
41 +
42 + fn finish(&mut self) -> Result<ReasoningDelta> {
43 + Ok(self.inner.finish())
44 + }
45 +}
rust/src/reasoning-parser/src/deepseek_r1.rs new
+44
@@ -0,0 +1,44 @@
1 +use vllm_tokenizer::DynTokenizer;
2 +
3 +use super::{DelimitedReasoningParser, ReasoningDelta, ReasoningParser, Result};
4 +
5 +/// Reasoning parser for DeepSeek R1 style outputs.
6 +///
7 +/// DeepSeek R1 may begin generating directly inside a reasoning span and only
8 +/// emit the closing `</think>` delimiter, so the no-boundary fallback defaults
9 +/// to `in_reasoning = true`.
10 +pub struct DeepSeekR1ReasoningParser {
11 + inner: DelimitedReasoningParser,
12 +}
13 +
14 +impl DeepSeekR1ReasoningParser {
15 + /// Create a DeepSeek R1 parser backed by the shared delimited state
16 + /// machine.
17 + pub fn new(tokenizer: DynTokenizer) -> Result<Self> {
18 + Ok(Self {
19 + inner: DelimitedReasoningParser::new(tokenizer, "<think>", "</think>", true)?,
20 + })
21 + }
22 +}
23 +
24 +impl ReasoningParser for DeepSeekR1ReasoningParser {
25 + fn create(tokenizer: DynTokenizer) -> Result<Box<dyn ReasoningParser>>
26 + where
27 + Self: Sized + 'static,
28 + {
29 + Ok(Box::new(Self::new(tokenizer)?))
30 + }
31 +
32 + fn initialize(&mut self, prompt_token_ids: &[u32]) -> Result<()> {
33 + self.inner.initialize(prompt_token_ids);
34 + Ok(())
35 + }
36 +
37 + fn push(&mut self, delta: &str) -> Result<ReasoningDelta> {
38 + Ok(self.inner.push(delta))
39 + }
40 +
41 + fn finish(&mut self) -> Result<ReasoningDelta> {
42 + Ok(self.inner.finish())
43 + }
44 +}
rust/src/reasoning-parser/src/delimited.rs new
+161
@@ -0,0 +1,161 @@
1 +use vllm_tokenizer::{DynTokenizer, Tokenizer};
2 +
3 +use super::{ReasoningDelta, ReasoningError, Result};
4 +
5 +/// Shared incremental state machine for tag-delimited reasoning protocols.
6 +///
7 +/// This helper is intentionally not a public parser type. Model-family parser
8 +/// wrappers own one `DelimitedReasoningParser` internally and expose the
9 +/// request-facing [`super::ReasoningParser`] trait.
10 +///
11 +/// The shared state machine stays generic by deriving its initial
12 +/// `current_in_reasoning` state from the prompt token boundary instead of
13 +/// hardcoding model-family conventions. That means families with the same
14 +/// delimiters can often reuse this implementation even if their chat templates
15 +/// prefill different prompts.
16 +pub(crate) struct DelimitedReasoningParser {
17 + tokenizer: DynTokenizer,
18 + current_in_reasoning: bool,
19 + buffer: String,
20 + start_token: String,
21 + end_token: String,
22 + start_token_id: u32,
23 + end_token_id: u32,
24 + default_in_reasoning: bool,
25 +}
26 +
27 +impl DelimitedReasoningParser {
28 + /// Create one delimited parser state machine.
29 + ///
30 + /// `default_in_reasoning` is only used when prompt initialization sees no
31 + /// reasoning boundary token at all. If the prompt contains either the
32 + /// start or end delimiter, that prompt boundary always wins.
33 + pub(crate) fn new(
34 + tokenizer: DynTokenizer,
35 + start_token: &'static str,
36 + end_token: &'static str,
37 + default_in_reasoning: bool,
38 + ) -> Result<Self> {
39 + let start_token_id =
40 + tokenizer.token_to_id(start_token).ok_or_else(|| ReasoningError::MissingToken {
41 + token: start_token.to_string(),
42 + })?;
43 + let end_token_id =
44 + tokenizer.token_to_id(end_token).ok_or_else(|| ReasoningError::MissingToken {
45 + token: end_token.to_string(),
46 + })?;
47 +
48 + Ok(Self {
49 + tokenizer,
50 + current_in_reasoning: default_in_reasoning,
51 + buffer: String::new(),
52 + start_token: start_token.to_string(),
53 + end_token: end_token.to_string(),
54 + start_token_id,
55 + end_token_id,
56 + default_in_reasoning,
57 + })
58 + }
59 +
60 + /// Initialize the starting state from prompt token IDs.
61 + pub(crate) fn initialize(&mut self, prompt_token_ids: &[u32]) {
62 + self.current_in_reasoning = last_reasoning_boundary(
63 + prompt_token_ids,
64 + self.start_token_id,
65 + self.end_token_id,
66 + self.tokenizer.as_ref(),
67 + )
68 + .unwrap_or(self.default_in_reasoning);
69 + }
70 +
71 + /// Parse one decoded text delta and return its reasoning/content split.
72 + pub(crate) fn push(&mut self, delta: &str) -> ReasoningDelta {
73 + self.buffer.push_str(delta);
74 +
75 + let partial_suffix_len = self.partial_suffix_len(&self.buffer);
76 + let stable_len = self.buffer.len() - partial_suffix_len;
77 + let pending_suffix = self.buffer.split_off(stable_len);
78 + let stable_text = std::mem::replace(&mut self.buffer, pending_suffix);
79 +
80 + self.parse_stable_text(&stable_text)
81 + }
82 +
83 + /// Flush any buffered partial delimiter suffix at end of stream.
84 + pub(crate) fn finish(&mut self) -> ReasoningDelta {
85 + let stable_text = std::mem::take(&mut self.buffer);
86 + self.parse_stable_text(&stable_text)
87 + }
88 +
89 + /// Parse text that is known not to end with a partial delimiter suffix.
90 + fn parse_stable_text(&mut self, mut stable: &str) -> ReasoningDelta {
91 + let mut delta = ReasoningDelta::default();
92 +
93 + while !stable.is_empty() {
94 + if self.current_in_reasoning {
95 + if let Some(end_idx) = stable.find(&self.end_token) {
96 + delta.push_reasoning(&stable[..end_idx]);
97 + stable = &stable[end_idx + self.end_token.len()..];
98 + self.current_in_reasoning = false;
99 + } else {
100 + delta.push_reasoning(stable);
101 + break;
102 + }
103 + } else if let Some(start_idx) = stable.find(&self.start_token) {
104 + delta.push_content(&stable[..start_idx]);
105 + stable = &stable[start_idx + self.start_token.len()..];
106 + self.current_in_reasoning = true;
107 + } else {
108 + delta.push_content(stable);
109 + break;
110 + }
111 + }
112 +
113 + delta
114 + }
115 +
116 + /// Return the longest trailing suffix that could still complete a
117 + /// delimiter.
118 + fn partial_suffix_len(&self, text: &str) -> usize {
119 + let mut best = 0;
120 + for idx in text.char_indices().map(|(idx, _)| idx).skip(1) {
121 + let suffix = &text[idx..];
122 + if self.start_token.starts_with(suffix) && self.start_token != suffix {
123 + best = best.max(text.len() - idx);
124 + }
125 + if self.end_token.starts_with(suffix) && self.end_token != suffix {
126 + best = best.max(text.len() - idx);
127 + }
128 + }
129 +
130 + if self.start_token.starts_with(text) && self.start_token != text {
131 + best = best.max(text.len());
132 + }
133 + if self.end_token.starts_with(text) && self.end_token != text {
134 + best = best.max(text.len());
135 + }
136 +
137 + best
138 + }
139 +}
140 +
141 +/// Determine the reasoning state implied by the last prompt boundary, if any.
142 +fn last_reasoning_boundary(
143 + prompt_token_ids: &[u32],
144 + start_token_id: u32,
145 + end_token_id: u32,
146 + tokenizer: &dyn Tokenizer,
147 +) -> Option<bool> {
148 + for token_id in prompt_token_ids.iter().rev() {
149 + if *token_id == start_token_id {
150 + return Some(true);
151 + }
152 + if *token_id == end_token_id {
153 + return Some(false);
154 + }
155 + if tokenizer.is_special_id(*token_id) {
156 + return None;
157 + }
158 + }
159 +
160 + None
161 +}
rust/src/reasoning-parser/src/gemma4.rs new
+273
@@ -0,0 +1,273 @@
1 +use vllm_tokenizer::DynTokenizer;
2 +
3 +use super::{DelimitedReasoningParser, ReasoningDelta, ReasoningParser, Result};
4 +
5 +const THOUGHT_PREFIX: &str = "thought\n";
6 +
7 +/// Reasoning parser for Google Gemma4 thinking models.
8 +///
9 +/// Gemma4 emits reasoning inside `<|channel> ... <channel|>` spans and adds a
10 +/// structural `thought\n` label at the beginning of the reasoning channel.
11 +/// This parser keeps the delimiter handling in the shared delimited parser and
12 +/// only layers on Gemma4-specific request adjustment plus prefix stripping.
13 +///
14 +/// Original Python implementation:
15 +/// <https://github.com/vllm-project/vllm/blob/18b1c77211d8f6fe800bcfb89524d2b598708032/vllm/reasoning/gemma4_reasoning_parser.py#L23>
16 +pub struct Gemma4ReasoningParser {
17 + inner: DelimitedReasoningParser,
18 + reasoning_text: String,
19 + prefix_stripped: bool,
20 +}
21 +
22 +impl Gemma4ReasoningParser {
23 + /// Create a Gemma4 parser.
24 + pub fn new(tokenizer: DynTokenizer) -> Result<Self> {
25 + Ok(Self {
26 + inner: DelimitedReasoningParser::new(tokenizer, "<|channel>", "<channel|>", false)?,
27 + reasoning_text: String::new(),
28 + prefix_stripped: false,
29 + })
30 + }
31 +
32 + /// Apply Gemma4's `thought\n` stripping rule to one reasoning delta.
33 + ///
34 + /// Early reasoning text is buffered until we can decide whether it begins
35 + /// with the structural channel label.
36 + fn strip_thought_prefix(&mut self, reasoning: &str) -> Option<String> {
37 + if self.prefix_stripped {
38 + return Some(reasoning.to_string());
39 + }
40 +
41 + self.reasoning_text.push_str(reasoning);
42 +
43 + if self.reasoning_text.starts_with(THOUGHT_PREFIX) {
44 + let prefix_len = THOUGHT_PREFIX.len();
45 + let previous_len = self.reasoning_text.len() - reasoning.len();
46 + if previous_len >= prefix_len {
47 + self.reasoning_text.clear();
48 + self.prefix_stripped = true;
49 + return Some(reasoning.to_string());
50 + }
51 +
52 + let prefix_chars_in_delta = prefix_len - previous_len;
53 + let stripped = &reasoning[prefix_chars_in_delta.min(reasoning.len())..];
54 + if stripped.is_empty() {
55 + if self.reasoning_text.len() >= prefix_len {
56 + self.reasoning_text.clear();
57 + self.prefix_stripped = true;
58 + }
59 + return None;
60 + }
61 +
62 + self.reasoning_text.clear();
63 + self.prefix_stripped = true;
64 + return Some(stripped.to_string());
65 + }
66 +
67 + if THOUGHT_PREFIX.starts_with(&self.reasoning_text) {
68 + return None;
69 + }
70 +
71 + self.prefix_stripped = true;
72 + Some(std::mem::take(&mut self.reasoning_text))
73 + }
74 +
75 + /// Apply Gemma4-specific reasoning post-processing to one parsed delta.
76 + fn post_process(&mut self, mut result: ReasoningDelta) -> ReasoningDelta {
77 + if let Some(reasoning) = result.reasoning.take() {
78 + result.reasoning =
79 + self.strip_thought_prefix(&reasoning).filter(|text| !text.is_empty());
80 + }
81 + result
82 + }
83 +}
84 +
85 +impl ReasoningParser for Gemma4ReasoningParser {
86 + fn create(tokenizer: DynTokenizer) -> Result<Box<dyn ReasoningParser>>
87 + where
88 + Self: Sized + 'static,
89 + {
90 + Ok(Box::new(Self::new(tokenizer)?))
91 + }
92 +
93 + fn preserve_special_tokens(&self) -> bool {
94 + true
95 + }
96 +
97 + fn initialize(&mut self, prompt_token_ids: &[u32]) -> Result<()> {
98 + self.inner.initialize(prompt_token_ids);
99 + self.reasoning_text.clear();
100 + self.prefix_stripped = false;
101 + Ok(())
102 + }
103 +
104 + fn push(&mut self, delta: &str) -> Result<ReasoningDelta> {
105 + let result = self.inner.push(delta);
106 + Ok(self.post_process(result))
107 + }
108 +
109 + fn finish(&mut self) -> Result<ReasoningDelta> {
110 + let result = self.inner.finish();
111 + Ok(self.post_process(result))
112 + }
113 +}
114 +
115 +#[cfg(test)]
116 +mod tests {
117 + use std::sync::Arc;
118 +
119 + use vllm_tokenizer::Tokenizer;
120 +
121 + use super::Gemma4ReasoningParser;
122 + use crate::ReasoningParser;
123 +
124 + struct FakeTokenizer;
125 +
126 + impl Tokenizer for FakeTokenizer {
127 + fn encode(
128 + &self,
129 + text: &str,
130 + _add_special_tokens: bool,
131 + ) -> vllm_tokenizer::Result<Vec<u32>> {
132 + Ok(text.chars().map(u32::from).collect())
133 + }
134 +
135 + fn decode(
136 + &self,
137 + token_ids: &[u32],
138 + _skip_special_tokens: bool,
139 + ) -> vllm_tokenizer::Result<String> {
140 + Ok(token_ids
141 + .iter()
142 + .map(|token_id| char::from_u32(*token_id).unwrap_or('\u{FFFD}'))
143 + .collect())
144 + }
145 +
146 + fn token_to_id(&self, token: &str) -> Option<u32> {
147 + match token {
148 + "<|channel>" => Some(1000),
149 + "<channel|>" => Some(1001),
150 + _ => None,
151 + }
152 + }
153 + }
154 +
155 + fn run_streaming(output: &[&str]) -> (Option<String>, Option<String>) {
156 + let tokenizer = Arc::new(FakeTokenizer);
157 + let mut parser = Gemma4ReasoningParser::new(tokenizer).unwrap();
158 + let mut reasoning = String::new();
159 + let mut content = String::new();
160 +
161 + for delta in output {
162 + let result = parser.push(delta).unwrap();
163 + if let Some(next) = result.reasoning {
164 + reasoning.push_str(&next);
165 + }
166 + if let Some(next) = result.content {
167 + content.push_str(&next);
168 + }
169 + }
170 +
171 + let final_delta = parser.finish().unwrap();
172 + if let Some(next) = final_delta.reasoning {
173 + reasoning.push_str(&next);
174 + }
175 + if let Some(next) = final_delta.content {
176 + content.push_str(&next);
177 + }
178 +
179 + (
180 + (!reasoning.is_empty()).then_some(reasoning),
181 + (!content.is_empty()).then_some(content),
182 + )
183 + }
184 +
185 + #[test]
186 + fn gemma4_reasoning_streaming_handles_channel_delimited_outputs() {
187 + let cases = [
188 + (
189 + "no_reasoning",
190 + vec!["This is content"],
191 + None,
192 + Some("This is content"),
193 + ),
194 + (
195 + "reasoning_and_content",
196 + vec!["<|channel>This is a reasoning section<channel|>This is the rest"],
197 + Some("This is a reasoning section"),
198 + Some("This is the rest"),
199 + ),
200 + (
201 + "complete_reasoning",
202 + vec!["<|channel>This is a reasoning section<channel|>"],
203 + Some("This is a reasoning section"),
204 + None,
205 + ),
206 + (
207 + "multiple_lines",
208 + vec!["<|channel>This\nThat<channel|>This is the rest\nThat"],
209 + Some("This\nThat"),
210 + Some("This is the rest\nThat"),
211 + ),
212 + (
213 + "no_end",
214 + vec!["<|channel>This is a reasoning section"],
215 + Some("This is a reasoning section"),
216 + None,
217 + ),
218 + ("empty", vec![""], None, None),
219 + (
220 + "newline_around_reasoning",
221 + vec!["Before\n<|channel>This is a reasoning section<channel|>\nThis is the rest"],
222 + Some("This is a reasoning section"),
223 + Some("Before\n\nThis is the rest"),
224 + ),
225 + (
226 + "thought_prefix",
227 + vec!["<|channel>thought\nActual reasoning here<channel|>Final answer"],
228 + Some("Actual reasoning here"),
229 + Some("Final answer"),
230 + ),
231 + (
232 + "thought_prefix_only",
233 + vec!["<|channel>thought\n<channel|>"],
234 + None,
235 + None,
236 + ),
237 + (
238 + "thought_prefix_multiline",
239 + vec!["<|channel>thought\nLine1\nLine2<channel|>Answer"],
240 + Some("Line1\nLine2"),
241 + Some("Answer"),
242 + ),
243 + (
244 + "thought_prefix_diverge",
245 + vec!["<|channel>thousand reasons<channel|>Done"],
246 + Some("thousand reasons"),
247 + Some("Done"),
248 + ),
249 + ];
250 +
251 + for (name, output, expected_reasoning, expected_content) in cases {
252 + let (reasoning, content) = run_streaming(&output);
253 + assert_eq!(reasoning.as_deref(), expected_reasoning, "{name}");
254 + assert_eq!(content.as_deref(), expected_content, "{name}");
255 + }
256 + }
257 +
258 + #[test]
259 + fn gemma4_strips_thought_prefix_even_when_split_across_deltas() {
260 + let (reasoning, content) =
261 + run_streaming(&["<|channel>thou", "ght", "\nabc", "<channel|>done"]);
262 + assert_eq!(reasoning.as_deref(), Some("abc"));
263 + assert_eq!(content.as_deref(), Some("done"));
264 + }
265 +
266 + #[test]
267 + fn gemma4_preserves_special_tokens() {
268 + let tokenizer = Arc::new(FakeTokenizer);
269 + let parser = Gemma4ReasoningParser::new(tokenizer).unwrap();
270 +
271 + assert!(parser.preserve_special_tokens());
272 + }
273 +}
rust/src/reasoning-parser/src/kimi.rs new
+39
@@ -0,0 +1,39 @@
1 +use vllm_tokenizer::DynTokenizer;
2 +
3 +use super::{DelimitedReasoningParser, ReasoningDelta, ReasoningParser, Result};
4 +
5 +/// Reasoning parser for legacy Kimi models that use Unicode thinking tags.
6 +pub struct KimiReasoningParser {
7 + inner: DelimitedReasoningParser,
8 +}
9 +
10 +impl KimiReasoningParser {
11 + /// Create a Kimi parser backed by the shared delimited state machine.
12 + pub fn new(tokenizer: DynTokenizer) -> Result<Self> {
13 + Ok(Self {
14 + inner: DelimitedReasoningParser::new(tokenizer, "◁think▷", "◁/think▷", false)?,
15 + })
16 + }
17 +}
18 +
19 +impl ReasoningParser for KimiReasoningParser {
20 + fn create(tokenizer: DynTokenizer) -> Result<Box<dyn ReasoningParser>>
21 + where
22 + Self: Sized + 'static,
23 + {
24 + Ok(Box::new(Self::new(tokenizer)?))
25 + }
26 +
27 + fn initialize(&mut self, prompt_token_ids: &[u32]) -> Result<()> {
28 + self.inner.initialize(prompt_token_ids);
29 + Ok(())
30 + }
31 +
32 + fn push(&mut self, delta: &str) -> Result<ReasoningDelta> {
33 + Ok(self.inner.push(delta))
34 + }
35 +
36 + fn finish(&mut self) -> Result<ReasoningDelta> {
37 + Ok(self.inner.finish())
38 + }
39 +}
rust/src/reasoning-parser/src/lib.rs new
+129
@@ -0,0 +1,129 @@
1 +//! Streaming reasoning parsers for chat completions.
2 +//!
3 +//! The key design choice here is that parser initialization prefers the
4 +//! *actual rendered prompt state* over model-family conventions. When a stream
5 +//! starts, each parser receives the prompt token IDs and inspects the last
6 +//! reasoning boundary that is already present in the prompt. In practice this
7 +//! is a more faithful signal than hardcoding assumptions such as "this model
8 +//! always starts in reasoning" or "this model always emits `<think>` itself".
9 +//!
10 +//! That prompt-first initialization lets multiple model families share the
11 +//! same incremental parser implementation even when older Python parsers split
12 +//! them apart. If two families use the same textual delimiters and differ
13 +//! mostly in how their chat templates prefill `<think>` / `</think>`, they can
14 +//! usually reuse one parser here because the prompt token IDs already tell us
15 +//! which state the stream is entering with.
16 +
17 +mod cohere_cmd;
18 +mod deepseek_r1;
19 +mod delimited;
20 +mod gemma4;
21 +mod kimi;
22 +mod qwen3;
23 +
24 +use thiserror::Error;
25 +use vllm_tokenizer::DynTokenizer;
26 +
27 +pub use self::cohere_cmd::CohereCmdReasoningParser;
28 +pub use self::deepseek_r1::DeepSeekR1ReasoningParser;
29 +pub(crate) use self::delimited::DelimitedReasoningParser;
30 +pub use self::gemma4::Gemma4ReasoningParser;
31 +pub use self::kimi::KimiReasoningParser;
32 +pub use self::qwen3::Qwen3ReasoningParser;
33 +
34 +/// DeepSeek V3 currently shares the standard `<think>...</think>` parser.
35 +pub type DeepSeekV3ReasoningParser = Qwen3ReasoningParser;
36 +/// DeepSeek V4 currently shares the standard `<think>...</think>` parser.
37 +pub type DeepSeekV4ReasoningParser = Qwen3ReasoningParser;
38 +/// GLM45 currently shares the standard `<think>...</think>` parser.
39 +pub type Glm45ReasoningParser = Qwen3ReasoningParser;
40 +/// Kimi K2 currently shares the standard `<think>...</think>` parser.
41 +// TODO: kimi k2 may implicitly end reasoning by starting a tool call section
42 +// using <|tool_calls_section_begin|>, we should support that.
43 +pub type KimiK2ReasoningParser = Qwen3ReasoningParser;
44 +/// MiniMax M2 currently shares the standard `<think>...</think>` parser.
45 +pub type MiniMaxM2ReasoningParser = Qwen3ReasoningParser;
46 +/// Nemotron V3 currently shares the standard `<think>...</think>` parser.
47 +pub type NemotronV3ReasoningParser = Qwen3ReasoningParser;
48 +/// Step3 currently shares the standard `<think>...</think>` parser.
49 +pub type Step3ReasoningParser = Qwen3ReasoningParser;
50 +
51 +/// Result alias for reasoning parser operations.
52 +pub type Result<T> = std::result::Result<T, ReasoningError>;
53 +
54 +/// One parsed streaming delta split into reasoning and visible content.
55 +#[derive(Debug, Default, Clone, PartialEq, Eq)]
56 +pub struct ReasoningDelta {
57 + pub reasoning: Option<String>,
58 + pub content: Option<String>,
59 +}
60 +
61 +impl ReasoningDelta {
62 + /// Return true when this delta carries neither reasoning nor content text.
63 + pub fn is_empty(&self) -> bool {
64 + self.reasoning.is_none() && self.content.is_none()
65 + }
66 +
67 + /// Append text to the reasoning portion, creating it on first use.
68 + pub(crate) fn push_reasoning(&mut self, text: &str) {
69 + if text.is_empty() {
70 + return;
71 + }
72 + match &mut self.reasoning {
73 + Some(existing) => existing.push_str(text),
74 + None => self.reasoning = Some(text.to_string()),
75 + }
76 + }
77 +
78 + /// Append text to the visible content portion, creating it on first use.
79 + pub(crate) fn push_content(&mut self, text: &str) {
80 + if text.is_empty() {
81 + return;
82 + }
83 + match &mut self.content {
84 + Some(existing) => existing.push_str(text),
85 + None => self.content = Some(text.to_string()),
86 + }
87 + }
88 +}
89 +
90 +/// Incremental parser that splits decoded text deltas into reasoning and
91 +/// content.
92 +pub trait ReasoningParser: Send {
93 + /// Construct a boxed parser instance for one request stream.
94 + fn create(tokenizer: DynTokenizer) -> Result<Box<dyn ReasoningParser>>
95 + where
96 + Self: Sized + 'static;
97 +
98 + /// Initialize parser state from prompt token IDs before output deltas
99 + /// arrive.
100 + fn initialize(&mut self, _prompt_token_ids: &[u32]) -> Result<()> {
101 + Ok(())
102 + }
103 +
104 + /// Return whether decoded output must preserve tokenizer special tokens.
105 + ///
106 + /// Some model families emit reasoning sentinels as special tokens. Those
107 + /// parsers need `skip_special_tokens = false` while parsing is enabled.
108 + fn preserve_special_tokens(&self) -> bool {
109 + false
110 + }
111 +
112 + /// Feed one decoded text delta into the parser.
113 + fn push(&mut self, delta: &str) -> Result<ReasoningDelta>;
114 +
115 + /// Flush any buffered partial delimiter state at end of stream.
116 + fn finish(&mut self) -> Result<ReasoningDelta> {
117 + Ok(ReasoningDelta::default())
118 + }
119 +}
120 +
121 +/// Errors produced while creating or running reasoning parsers.
122 +#[derive(Debug, Error)]
123 +pub enum ReasoningError {
124 + #[error("tokenizer is missing reasoning delimiter token `{token}`")]
125 + MissingToken { token: String },
126 +}
127 +
128 +#[cfg(test)]
129 +mod tests;
rust/src/reasoning-parser/src/qwen3.rs new
+43
@@ -0,0 +1,43 @@
1 +use vllm_tokenizer::DynTokenizer;
2 +
3 +use super::{DelimitedReasoningParser, ReasoningDelta, ReasoningParser, Result};
4 +
5 +/// Reasoning parser for the Qwen3/Qwen3.5 family.
6 +///
7 +/// This parser uses standard `<think>...</think>` delimiters and defaults to
8 +/// waiting for an explicit start token when prompt initialization finds no
9 +/// reasoning boundary.
10 +pub struct Qwen3ReasoningParser {
11 + inner: DelimitedReasoningParser,
12 +}
13 +
14 +impl Qwen3ReasoningParser {
15 + /// Create a Qwen3 parser backed by the shared delimited state machine.
16 + pub fn new(tokenizer: DynTokenizer) -> Result<Self> {
17 + Ok(Self {
18 + inner: DelimitedReasoningParser::new(tokenizer, "<think>", "</think>", false)?,
19 + })
20 + }
21 +}
22 +
23 +impl ReasoningParser for Qwen3ReasoningParser {
24 + fn create(tokenizer: DynTokenizer) -> Result<Box<dyn ReasoningParser>>
25 + where
26 + Self: Sized + 'static,
27 + {
28 + Ok(Box::new(Self::new(tokenizer)?))
29 + }
30 +
31 + fn initialize(&mut self, prompt_token_ids: &[u32]) -> Result<()> {
32 + self.inner.initialize(prompt_token_ids);
33 + Ok(())
34 + }
35 +
36 + fn push(&mut self, delta: &str) -> Result<ReasoningDelta> {
37 + Ok(self.inner.push(delta))
38 + }
39 +
40 + fn finish(&mut self) -> Result<ReasoningDelta> {
41 + Ok(self.inner.finish())
42 + }
43 +}
rust/src/reasoning-parser/src/tests.rs new
+161
@@ -0,0 +1,161 @@
1 +use std::sync::Arc;
2 +
3 +use vllm_tokenizer::Tokenizer;
4 +
5 +use super::{
6 + DeepSeekR1ReasoningParser, DelimitedReasoningParser, Qwen3ReasoningParser, ReasoningParser,
7 +};
8 +
9 +struct FakeTokenizer;
10 +
11 +impl Tokenizer for FakeTokenizer {
12 + fn encode(&self, text: &str, _add_special_tokens: bool) -> vllm_tokenizer::Result<Vec<u32>> {
13 + Ok(text.chars().map(u32::from).collect())
14 + }
15 +
16 + fn decode(
17 + &self,
18 + token_ids: &[u32],
19 + _skip_special_tokens: bool,
20 + ) -> vllm_tokenizer::Result<String> {
21 + Ok(token_ids
22 + .iter()
23 + .map(|token_id| char::from_u32(*token_id).unwrap_or('\u{FFFD}'))
24 + .collect())
25 + }
26 +
27 + fn token_to_id(&self, token: &str) -> Option<u32> {
28 + match token {
29 + "<think>" => Some(1),
30 + "</think>" => Some(2),
31 + "<|START_THINKING|>" => Some(3),
32 + "<|END_THINKING|>" => Some(4),
33 + "◁think▷" => Some(5),
34 + "◁/think▷" => Some(6),
35 + _ => None,
36 + }
37 + }
38 +
39 + fn is_special_id(&self, token_id: u32) -> bool {
40 + token_id == 7
41 + }
42 +}
43 +
44 +#[test]
45 +fn delimited_content_only_stream() {
46 + let tokenizer = Arc::new(FakeTokenizer);
47 + let mut parser =
48 + DelimitedReasoningParser::new(tokenizer, "<think>", "</think>", false).unwrap();
49 +
50 + assert_eq!(
51 + parser.push("plain content").content.as_deref(),
52 + Some("plain content")
53 + );
54 +}
55 +
56 +#[test]
57 +fn delimited_single_chunk_with_reasoning_and_content() {
58 + let tokenizer = Arc::new(FakeTokenizer);
59 + let mut parser =
60 + DelimitedReasoningParser::new(tokenizer, "<think>", "</think>", false).unwrap();
61 +
62 + let delta = parser.push("<think>reason</think>answer");
63 + assert_eq!(delta.reasoning.as_deref(), Some("reason"));
64 + assert_eq!(delta.content.as_deref(), Some("answer"));
65 +}
66 +
67 +#[test]
68 +fn delimited_partial_tokens_across_chunks() {
69 + let tokenizer = Arc::new(FakeTokenizer);
70 + let mut parser =
71 + DelimitedReasoningParser::new(tokenizer, "<think>", "</think>", false).unwrap();
72 +
73 + assert!(parser.push("<thi").is_empty());
74 + let delta = parser.push("nk>reason</think>answer");
75 + assert_eq!(delta.reasoning.as_deref(), Some("reason"));
76 + assert_eq!(delta.content.as_deref(), Some("answer"));
77 +}
78 +
79 +#[test]
80 +fn delimited_finish_flushes_buffer() {
81 + let tokenizer = Arc::new(FakeTokenizer);
82 + let mut parser =
83 + DelimitedReasoningParser::new(tokenizer, "<think>", "</think>", false).unwrap();
84 + parser.initialize(&[1]);
85 +
86 + let delta = parser.push("unfinished</thi");
87 + assert_eq!(delta.reasoning.as_deref(), Some("unfinished"));
88 + let final_delta = parser.finish();
89 + assert_eq!(final_delta.reasoning.as_deref(), Some("</thi"));
90 +}
91 +
92 +#[test]
93 +fn qwen3_without_prompt_markers_expects_start_token() {
94 + let tokenizer = Arc::new(FakeTokenizer);
95 + let mut parser = Qwen3ReasoningParser::new(tokenizer).unwrap();
96 +
97 + let delta = parser.push("reason</think>answer").unwrap();
98 + assert_eq!(delta.reasoning, None);
99 + assert_eq!(delta.content.as_deref(), Some("reason</think>answer"));
100 +}
101 +
102 +#[test]
103 +fn qwen3_prompt_end_marker_starts_in_content() {
104 + let tokenizer = Arc::new(FakeTokenizer);
105 + let mut parser = Qwen3ReasoningParser::new(tokenizer).unwrap();
106 + parser.initialize(&[2]).unwrap();
107 +
108 + let delta = parser.push("answer").unwrap();
109 + assert_eq!(delta.reasoning, None);
110 + assert_eq!(delta.content.as_deref(), Some("answer"));
111 +}
112 +
113 +#[test]
114 +fn qwen3_tolerates_old_and_new_formats() {
115 + let tokenizer = Arc::new(FakeTokenizer);
116 +
117 + let mut old_parser = Qwen3ReasoningParser::new(tokenizer.clone()).unwrap();
118 + let old = old_parser.push("<think>reason</think>answer").unwrap();
119 + assert_eq!(old.reasoning.as_deref(), Some("reason"));
120 + assert_eq!(old.content.as_deref(), Some("answer"));
121 +
122 + let mut new_parser = Qwen3ReasoningParser::new(tokenizer).unwrap();
123 + new_parser.initialize(&[1]).unwrap();
124 + let new = new_parser.push("reason</think>answer").unwrap();
125 + assert_eq!(new.reasoning.as_deref(), Some("reason"));
126 + assert_eq!(new.content.as_deref(), Some("answer"));
127 +}
128 +
129 +#[test]
130 +fn qwen3_stops_scanning_at_last_special_token() {
131 + let tokenizer = Arc::new(FakeTokenizer);
132 + let mut parser = Qwen3ReasoningParser::new(tokenizer).unwrap();
133 +
134 + parser.initialize(&[1, 7]).unwrap();
135 +
136 + let delta = parser.push("answer").unwrap();
137 + assert_eq!(delta.reasoning, None);
138 + assert_eq!(delta.content.as_deref(), Some("answer"));
139 +}
140 +
141 +#[test]
142 +fn deepseek_r1_defaults_to_reasoning_without_prompt_boundary() {
143 + let tokenizer = Arc::new(FakeTokenizer);
144 + let mut parser = DeepSeekR1ReasoningParser::new(tokenizer).unwrap();
145 +
146 + let delta = parser.push("reason</think>answer").unwrap();
147 + assert_eq!(delta.reasoning.as_deref(), Some("reason"));
148 + assert_eq!(delta.content.as_deref(), Some("answer"));
149 +}
150 +
151 +#[test]
152 +fn deepseek_r1_stops_scanning_at_last_special_token() {
153 + let tokenizer = Arc::new(FakeTokenizer);
154 + let mut parser = DeepSeekR1ReasoningParser::new(tokenizer).unwrap();
155 +
156 + parser.initialize(&[2, 7]).unwrap();
157 +
158 + let delta = parser.push("reason</think>answer").unwrap();
159 + assert_eq!(delta.reasoning.as_deref(), Some("reason"));
160 + assert_eq!(delta.content.as_deref(), Some("answer"));
161 +}
rust/src/server/Cargo.toml new
+57
@@ -0,0 +1,57 @@
1 +[package]
2 +name = "vllm-server"
3 +version.workspace = true
4 +edition.workspace = true
5 +license.workspace = true
6 +
7 +[dependencies]
8 +anyhow.workspace = true
9 +asynk-strim-attr.workspace = true
10 +axum.workspace = true
11 +futures.workspace = true
12 +http-body.workspace = true
13 +itertools.workspace = true
14 +libc.workspace = true
15 +llm-multimodal.workspace = true
16 +prost.workspace = true
17 +prost-types.workspace = true
18 +rmpv.workspace = true
19 +serde.workspace = true
20 +serde_json.workspace = true
21 +serde_with.workspace = true
22 +socket2.workspace = true
23 +thiserror-ext.workspace = true
24 +tokio.workspace = true
25 +tokio-stream.workspace = true
26 +tokio-util.workspace = true
27 +tonic.workspace = true
28 +tonic-prost.workspace = true
29 +tower-http.workspace = true
30 +tracing.workspace = true
31 +tracing-futures.workspace = true
32 +tracing-subscriber.workspace = true
33 +uuid.workspace = true
34 +validator.workspace = true
35 +vllm-chat.workspace = true
36 +vllm-engine-core-client.workspace = true
37 +vllm-llm.workspace = true
38 +vllm-metrics.workspace = true
39 +vllm-text.workspace = true
40 +
41 +[build-dependencies]
42 +tonic-prost-build.workspace = true
43 +
44 +[dev-dependencies]
45 +anyhow.workspace = true
46 +async-openai = { workspace = true, features = ["full"] }
47 +bytes.workspace = true
48 +clap.workspace = true
49 +expect-test.workspace = true
50 +rmp-serde.workspace = true
51 +serial_test.workspace = true
52 +tower.workspace = true
53 +vllm-engine-core-client = { workspace = true, features = ["test-util"] }
54 +zeromq.workspace = true
55 +
56 +[lints]
57 +workspace = true
rust/src/server/build.rs new
+12
@@ -0,0 +1,12 @@
1 +fn main() -> Result<(), Box<dyn std::error::Error>> {
2 + let manifest_dir = std::env::var("CARGO_MANIFEST_DIR").unwrap();
3 + let proto_dir = format!("{manifest_dir}/../../proto");
4 +
5 + tonic_prost_build::configure()
6 + .build_server(true)
7 + .build_client(true)
8 + .protoc_arg("--experimental_allow_proto3_optional") // be compatible with old compilers
9 + .compile_protos(&[format!("{proto_dir}/vllm_grpc.proto")], &[proto_dir])?;
10 +
11 + Ok(())
12 +}
rust/src/server/examples/README.md new
+39
@@ -0,0 +1,39 @@
1 +# Server Smoke Test
2 +
3 +Start a fresh headless `vllm` engine:
4 +
5 +```bash
6 +source ../vllm/.venv/bin/activate
7 +HF_HUB_OFFLINE=1 \
8 +VLLM_LOGGING_LEVEL=DEBUG \
9 +VLLM_CPU_KVCACHE_SPACE=2 \
10 +VLLM_HOST_IP=127.0.0.1 \
11 +VLLM_LOOPBACK_IP=127.0.0.1 \
12 +python3 -m vllm.entrypoints.cli.main serve Qwen/Qwen3-0.6B \
13 + --headless \
14 + --data-parallel-address 127.0.0.1 \
15 + --data-parallel-rpc-port 62100 \
16 + --data-parallel-size-local 1 \
17 + --max-model-len 512 \
18 + --dtype float16
19 +```
20 +
21 +Run the Rust server smoke test:
22 +
23 +```bash
24 +cargo run -p vllm-server --example external_engine_openai_qwen -- \
25 + --handshake-address tcp://127.0.0.1:62100
26 +```
27 +
28 +The example starts the Rust OpenAI-compatible server on an ephemeral local port,
29 +connects to it via the `async-openai` Rust client, lists models, and then checks
30 +that a streamed chat completion yields the assistant role chunk, final-answer
31 +content chunks, and a terminal finish chunk. This example intentionally uses
32 +`async-openai`'s standard typed `create_stream` API instead of BYOT, so it does
33 +not inspect the nonstandard `reasoning_content` field even though the Rust
34 +server may emit it for reasoning-capable models such as Qwen3. For reasoning
35 +behavior itself, use the `vllm-chat` smoke test or the `vllm-server`
36 +route tests.
37 +
38 +IMPORTANT: Restart `vllm` each time you run the smoke test. The current headless
39 +engine cannot safely handle frontend reconnects after the client shuts down.
rust/src/server/examples/external_engine_openai_qwen.rs new
+215
@@ -0,0 +1,215 @@
1 +use std::time::Duration;
2 +
3 +use anyhow::{Context, Result, bail};
4 +use async_openai::Client;
5 +use async_openai::config::OpenAIConfig;
6 +use async_openai::types::chat::{
7 + ChatCompletionRequestUserMessageArgs, CreateChatCompletionRequest,
8 + CreateChatCompletionRequestArgs,
9 +};
10 +use async_openai::types::models::ListModelResponse;
11 +use clap::Parser;
12 +use futures::StreamExt as _;
13 +use tokio_util::sync::CancellationToken;
14 +use tracing_subscriber::EnvFilter;
15 +use vllm_engine_core_client::TransportMode;
16 +use vllm_server::{
17 + ChatTemplateContentFormatOption, Config, CoordinatorMode, HttpListenerMode, ParserSelection,
18 + RendererSelection, serve,
19 +};
20 +
21 +#[derive(Debug, Parser)]
22 +#[command(
23 + about = "Smoke-test the Rust OpenAI server with async-openai against an external Qwen vLLM engine."
24 +)]
25 +struct Args {
26 + #[arg(long)]
27 + handshake_address: String,
28 + #[arg(long, default_value_t = 1)]
29 + engine_count: usize,
30 + #[arg(long, default_value = "Qwen/Qwen3-0.6B")]
31 + model: String,
32 + #[arg(long, default_value = "127.0.0.1")]
33 + host: String,
34 + #[arg(long, default_value_t = 30)]
35 + ready_timeout_secs: u64,
36 + #[arg(
37 + long,
38 + default_value = "What is the capital of France? Answer with one word."
39 + )]
40 + prompt: String,
41 +}
42 +
43 +#[tokio::main(flavor = "multi_thread")]
44 +async fn main() -> Result<()> {
45 + init_tracing();
46 + let args = Args::parse();
47 + let port = unique_local_port()?;
48 + let config = Config {
49 + transport_mode: TransportMode::HandshakeOwner {
50 + handshake_address: args.handshake_address,
51 + advertised_host: args.host,
52 + engine_count: args.engine_count,
53 + ready_timeout: Duration::from_secs(args.ready_timeout_secs),
54 + local_input_address: None,
55 + local_output_address: None,
56 + },
57 + coordinator_mode: CoordinatorMode::MaybeInProc,
58 + model: args.model,
59 + served_model_name: vec![],
60 + listener_mode: HttpListenerMode::BindTcp {
61 + host: "127.0.0.1".to_string(),
62 + port,
63 + },
64 + tool_call_parser: ParserSelection::Auto,
65 + reasoning_parser: ParserSelection::Auto,
66 + renderer: RendererSelection::Auto,
67 + chat_template: None,
68 + default_chat_template_kwargs: None,
69 + chat_template_content_format: ChatTemplateContentFormatOption::Auto,
70 + enable_log_requests: false,
71 + disable_log_stats: false,
72 + grpc_port: None,
73 + shutdown_timeout: Duration::ZERO,
74 + };
75 +
76 + let bind_address = format!("127.0.0.1:{port}");
77 + let shutdown = CancellationToken::new();
78 + let server_config = config.clone();
79 + let server_shutdown = shutdown.clone();
80 + let server_task = tokio::spawn(async move { serve(server_config, server_shutdown).await });
81 +
82 + let client = Client::with_config(
83 + OpenAIConfig::new()
84 + .with_api_key("unused")
85 + .with_api_base(format!("http://{bind_address}/v1")),
86 + );
87 +
88 + print_models(&client).await?;
89 + let final_text = stream_completion(&client, &config.model, &args.prompt).await?;
90 +
91 + println!();
92 + println!("final_text={final_text:?}");
93 +
94 + shutdown_server(server_task, shutdown).await
95 +}
96 +
97 +fn init_tracing() {
98 + let filter = EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("info"));
99 + let _ = tracing_subscriber::fmt().with_env_filter(filter).try_init();
100 +}
101 +
102 +fn unique_local_port() -> Result<u16> {
103 + let listener = std::net::TcpListener::bind("127.0.0.1:0")
104 + .context("failed to allocate local smoke-test port")?;
105 + let port = listener.local_addr().context("failed to read local smoke-test port")?.port();
106 + drop(listener);
107 + Ok(port)
108 +}
109 +
110 +async fn print_models(client: &Client<OpenAIConfig>) -> Result<()> {
111 + let models = wait_for_models(client).await?;
112 + let model_ids = models.data.into_iter().map(|model| model.id).collect::<Vec<_>>();
113 + println!("models={model_ids:?}");
114 + Ok(())
115 +}
116 +
117 +async fn wait_for_models(client: &Client<OpenAIConfig>) -> Result<ListModelResponse> {
118 + let mut last_error = None;
119 + for _ in 0..240 {
120 + match client.models().list().await {
121 + Ok(models) => return Ok(models),
122 + Err(error) => {
123 + last_error = Some(error);
124 + tokio::time::sleep(Duration::from_millis(500)).await;
125 + }
126 + }
127 + }
128 +
129 + match last_error {
130 + Some(error) => Err(error).context("OpenAI server did not become ready in time"),
131 + None => bail!("OpenAI server readiness loop finished without a result"),
132 + }
133 +}
134 +
135 +async fn stream_completion(
136 + client: &Client<OpenAIConfig>,
137 + model: &str,
138 + prompt: &str,
139 +) -> Result<String> {
140 + // Keep this smoke test on async-openai's standard `create_stream` path so it
141 + // exercises the ordinary typed chat-completions client without BYOT
142 + // request/response types.
143 + //
144 + // The current async-openai chat-completions stream delta type does not expose
145 + // our OpenAI-compatible `reasoning_content` extension field, so this
146 + // example only validates the assistant role chunk, visible `content`
147 + // deltas, and terminal finish chunk. Reasoning coverage lives in our own
148 + // route tests and in the `vllm-chat` smoke example.
149 + let request: CreateChatCompletionRequest = CreateChatCompletionRequestArgs::default()
150 + .model(model)
151 + .stream(true)
152 + .temperature(0.0)
153 + .max_completion_tokens(128u32)
154 + .messages([ChatCompletionRequestUserMessageArgs::default()
155 + .content(prompt)
156 + .build()
157 + .context("failed to build user chat message")?
158 + .into()])
159 + .build()
160 + .context("failed to build chat completion request")?;
161 +
162 + let mut stream = client
163 + .chat()
164 + .create_stream(request)
165 + .await
166 + .context("failed to create streaming chat completion")?;
167 +
168 + let mut final_text = String::new();
169 + let mut saw_role = false;
170 + let mut saw_finish_reason = false;
171 + let mut saw_text = false;
172 +
173 + while let Some(chunk) = stream.next().await {
174 + let chunk = chunk.context("streaming chat completion failed")?;
175 + for choice in chunk.choices {
176 + if choice.delta.role.is_some() {
177 + saw_role = true;
178 + }
179 + if let Some(delta) = choice.delta.content {
180 + if !saw_text {
181 + print!("[answer] ");
182 + }
183 + print!("{delta}");
184 + final_text.push_str(&delta);
185 + saw_text = true;
186 + }
187 + if choice.finish_reason.is_some() {
188 + saw_finish_reason = true;
189 + }
190 + }
191 + }
192 +
193 + if !saw_role {
194 + bail!("stream ended without an assistant role chunk");
195 + }
196 + if !saw_finish_reason {
197 + bail!("stream ended without a terminal finish reason");
198 + }
199 + if final_text.is_empty() {
200 + bail!("stream ended without any content deltas");
201 + }
202 +
203 + Ok(final_text)
204 +}
205 +
206 +async fn shutdown_server(
207 + server_task: tokio::task::JoinHandle<anyhow::Result<()>>,
208 + shutdown: CancellationToken,
209 +) -> Result<()> {
210 + shutdown.cancel();
211 + server_task
212 + .await
213 + .context("server task join failed")?
214 + .context("server task failed")
215 +}
rust/src/server/src/config.rs new
+110
@@ -0,0 +1,110 @@
1 +use std::collections::HashMap;
2 +use std::time::Duration;
3 +
4 +use anyhow::Result;
5 +use serde_json::Value;
6 +use vllm_chat::{ChatTemplateContentFormatOption, ParserSelection, RendererSelection};
7 +use vllm_engine_core_client::{CoordinatorMode as EngineCoreCoordinatorMode, TransportMode};
8 +
9 +/// How the HTTP server obtains its listening socket.
10 +#[derive(Debug, Clone, PartialEq, Eq)]
11 +pub enum HttpListenerMode {
12 + /// Bind a fresh TCP listener on the given host/port.
13 + BindTcp { host: String, port: u16 },
14 + /// Bind a fresh Unix domain listener on the given filesystem path.
15 + BindUnix { path: String },
16 + /// Adopt an already-open listening socket inherited from a supervisor
17 + /// process.
18 + InheritedFd { fd: i32 },
19 +}
20 +
21 +/// Which coordinator implementation should be active when one is present for a
22 +/// frontend client.
23 +#[derive(Debug, Clone, PartialEq, Eq)]
24 +pub enum CoordinatorMode {
25 + /// Do not run a coordinator at all.
26 + None,
27 + /// Run the Rust in-process coordinator for managed `serve` deployments, if
28 + /// there are multiple engines and the model is MoE.
29 + MaybeInProc,
30 + /// Connect to an external coordinator owned by another process.
31 + External { address: String },
32 +}
33 +
34 +/// Normalized runtime configuration for the minimal OpenAI-compatible server.
35 +#[derive(Debug, Clone, PartialEq, Eq)]
36 +pub struct Config {
37 + /// Frontend-to-engine transport setup.
38 + pub transport_mode: TransportMode,
39 + /// Requested frontend-side coordinator behavior.
40 + pub coordinator_mode: CoordinatorMode,
41 + /// Backend model identifier used for engine-core loading.
42 + pub model: String,
43 + /// Model name(s) exposed to clients via the OpenAI API. When non-empty,
44 + /// the first entry is used as the primary ID in responses and all entries
45 + /// are accepted in requests. When empty, falls back to `model`.
46 + pub served_model_name: Vec<String>,
47 + /// HTTP listener setup.
48 + pub listener_mode: HttpListenerMode,
49 + /// Tool-call parser selection.
50 + pub tool_call_parser: ParserSelection,
51 + /// Reasoning parser selection.
52 + pub reasoning_parser: ParserSelection,
53 + /// Chat renderer selection.
54 + pub renderer: RendererSelection,
55 + /// Server-default chat template override, as a file path or inline
56 + /// template.
57 + pub chat_template: Option<String>,
58 + /// Server-default keyword arguments merged into every chat-template render.
59 + pub default_chat_template_kwargs: Option<HashMap<String, Value>>,
60 + /// How to serialize `message.content` for chat-template rendering.
61 + pub chat_template_content_format: ChatTemplateContentFormatOption,
62 + /// Log a summary line for each completed request.
63 + pub enable_log_requests: bool,
64 + /// When `true`, suppress periodic stats logging (throughput, queue depth,
65 + /// cache usage).
66 + pub disable_log_stats: bool,
67 + /// TCP port for the gRPC Generate service. When `None`, no gRPC server is
68 + /// started.
69 + pub grpc_port: Option<u16>,
70 + /// Maximum time to wait for active HTTP/gRPC requests to drain on shutdown.
71 + pub shutdown_timeout: Duration,
72 +}
73 +
74 +impl Config {
75 + /// Validate frontend configuration that can be checked before engine
76 + /// startup.
77 + pub fn validate(&self) -> Result<()> {
78 + vllm_chat::validate_parser_overrides(&self.tool_call_parser, &self.reasoning_parser)?;
79 +
80 + Ok(())
81 + }
82 +
83 + /// Return the number of engines implied by the configured transport mode.
84 + pub fn engine_count(&self) -> usize {
85 + match &self.transport_mode {
86 + TransportMode::HandshakeOwner { engine_count, .. }
87 + | TransportMode::Bootstrapped { engine_count, .. } => *engine_count,
88 + }
89 + }
90 +
91 + /// Resolve the effective coordinator mode.
92 + pub fn effective_coordinator_mode(
93 + &self,
94 + model_is_moe: bool,
95 + ) -> Option<EngineCoreCoordinatorMode> {
96 + match &self.coordinator_mode {
97 + CoordinatorMode::None => None,
98 + CoordinatorMode::MaybeInProc => {
99 + if model_is_moe && self.engine_count() > 1 {
100 + Some(EngineCoreCoordinatorMode::InProc)
101 + } else {
102 + None
103 + }
104 + }
105 + CoordinatorMode::External { address } => Some(EngineCoreCoordinatorMode::External {
106 + address: address.clone(),
107 + }),
108 + }
109 + }
110 +}
rust/src/server/src/error.rs new
+74
@@ -0,0 +1,74 @@
1 +use axum::Json;
2 +use axum::http::StatusCode;
3 +use axum::response::{IntoResponse, Response};
4 +use thiserror_ext::{Construct, Macro};
5 +
6 +use crate::routes::openai::utils::types::{ErrorDetail, ErrorResponse};
7 +
8 +/// Small OpenAI-style error family used by the minimal HTTP layer.
9 +#[derive(Debug, Construct, Macro)]
10 +pub enum ApiError {
11 + /// The request is syntactically valid OpenAI JSON but asks for unsupported
12 + /// behavior.
13 + InvalidRequest {
14 + message: String,
15 + param: Option<&'static str>,
16 + },
17 + /// The requested model name does not match the single configured model.
18 + ModelNotFound { model: String },
19 + /// The request body could not be parsed as valid JSON.
20 + JsonParseError { message: String },
21 + /// An unexpected internal failure happened before streaming started.
22 + ServerError { message: String },
23 +}
24 +
25 +impl ApiError {
26 + /// Return the HTTP status code associated with this API error.
27 + pub fn status_code(&self) -> StatusCode {
28 + match self {
29 + Self::InvalidRequest { .. } => StatusCode::BAD_REQUEST,
30 + Self::ModelNotFound { .. } => StatusCode::NOT_FOUND,
31 + Self::ServerError { .. } => StatusCode::INTERNAL_SERVER_ERROR,
32 + Self::JsonParseError { .. } => StatusCode::BAD_REQUEST,
33 + }
34 + }
35 +
36 + /// Convert this error into the standard OpenAI-compatible JSON error
37 + /// payload.
38 + pub fn to_error_response(&self) -> ErrorResponse {
39 + let error = match self {
40 + Self::InvalidRequest { message, param } => ErrorDetail {
41 + message: message.clone(),
42 + error_type: "invalid_request_error".to_string(),
43 + param: param.map(|p| p.to_string()),
44 + code: Some("invalid_request_error".to_string()),
45 + },
46 + Self::ModelNotFound { model } => ErrorDetail {
47 + message: format!("The model `{model}` does not exist."),
48 + error_type: "invalid_request_error".to_string(),
49 + param: Some("model".to_string()),
50 + code: Some("model_not_found".to_string()),
51 + },
52 + Self::ServerError { message } => ErrorDetail {
53 + message: message.clone(),
54 + error_type: "server_error".to_string(),
55 + param: None,
56 + code: Some("server_error".to_string()),
57 + },
58 + Self::JsonParseError { message } => ErrorDetail {
59 + message: message.clone(),
60 + error_type: "invalid_request_error".to_string(),
61 + param: None,
62 + code: Some("json_parse_error".to_string()),
63 + },
64 + };
65 +
66 + ErrorResponse { error }
67 + }
68 +}
69 +
70 +impl IntoResponse for ApiError {
71 + fn into_response(self) -> Response {
72 + (self.status_code(), Json(self.to_error_response())).into_response()
73 + }
74 +}
rust/src/server/src/grpc/convert.rs new
+679
@@ -0,0 +1,679 @@
1 +//! Conversion between gRPC protobuf types and internal `vllm-text`
2 +//! request/response types.
3 +
4 +use tonic::Status;
5 +use uuid::Uuid;
6 +use vllm_engine_core_client::protocol::{StopReason, StructuredOutputsParams};
7 +use vllm_text::{
8 + DecodedLogprobs, DecodedPromptLogprobs, FinishReason, Finished, Prompt, SamplingParams,
9 + TextDecodeOptions, TextRequest,
10 +};
11 +
12 +use super::pb;
13 +
14 +// ========================================================================================
15 +// Request conversion
16 +// ========================================================================================
17 +
18 +/// Convert a gRPC `GenerateRequest` into the internal `TextRequest`.
19 +///
20 +/// If `req.model` is non-empty, it must match one of `served_model_names`;
21 +/// otherwise the request is rejected with `NotFound`. An empty string is
22 +/// treated as "unset" (proto3 default) and accepted.
23 +pub fn to_text_request(
24 + req: pb::GenerateRequest,
25 + stream: bool,
26 + served_model_names: &[String],
27 +) -> Result<TextRequest, Status> {
28 + if !req.model.is_empty() && !served_model_names.iter().any(|n| n == &req.model) {
29 + return Err(Status::not_found(format!(
30 + "model `{}` not found",
31 + req.model
32 + )));
33 + }
34 +
35 + if req.truncate_prompt_tokens != 0 {
36 + return Err(Status::invalid_argument(
37 + "truncate_prompt_tokens is not supported",
38 + ));
39 + }
40 +
41 + let prompt = match req.prompt {
42 + Some(pb::generate_request::Prompt::Text(text)) => Prompt::Text(text),
43 + Some(pb::generate_request::Prompt::TokenIds(ids)) => Prompt::TokenIds(ids.ids),
44 + None => return Err(Status::invalid_argument("prompt is required")),
45 + };
46 +
47 + let request_id = if req.request_id.is_empty() {
48 + Uuid::new_v4().to_string()
49 + } else {
50 + req.request_id
51 + };
52 +
53 + let sampling = req.sampling.as_ref();
54 + let decoding = req.decoding.as_ref();
55 + let stopping = req.stopping.as_ref();
56 + let response = req.response.as_ref();
57 + let kv = req.kv.as_ref();
58 +
59 + let mut sampling_params =
60 + build_sampling_params(req.temperature, sampling, decoding, stopping, response)?;
61 +
62 + // Thread KVCacheParameters → SamplingParams fields.
63 + if let Some(kv) = kv {
64 + // Thread kv_transfer_params through vllm_xargs, matching the HTTP route
65 + // convention.
66 + if let Some(kv_struct) = kv.kv_transfer_params.as_ref() {
67 + let kv_json = proto_struct_to_json(kv_struct);
68 + let map = sampling_params.vllm_xargs.get_or_insert_with(Default::default);
69 + map.insert("kv_transfer_params".to_string(), kv_json);
70 + }
71 + if kv.bypass_prefix_cache {
72 + sampling_params.skip_reading_prefix_cache = Some(true);
73 + }
74 + }
75 +
76 + let decode_options = TextDecodeOptions {
77 + skip_special_tokens: true,
78 + include_stop_str_in_output: stopping.is_some_and(|s| s.include_stop_strings),
79 + stop_strings: stopping.map(|s| &s.stop_strings).filter(|ss| !ss.is_empty()).cloned(),
80 + min_tokens: stopping.map_or(0, |s| s.min_new_tokens),
81 + };
82 +
83 + Ok(TextRequest {
84 + request_id,
85 + prompt,
86 + mm_features: None,
87 + sampling_params,
88 + decode_options,
89 + intermediate: stream,
90 + priority: req.priority,
91 + cache_salt: kv.map(|k| &k.cache_salt).filter(|s| !s.is_empty()).cloned(),
92 + add_special_tokens: true,
93 + data_parallel_rank: None,
94 + })
95 +}
96 +
97 +fn build_sampling_params(
98 + temperature: Option<f32>,
99 + sampling: Option<&pb::RandomSampling>,
100 + decoding: Option<&pb::DecodingParameters>,
101 + stopping: Option<&pb::StoppingCriteria>,
102 + response: Option<&pb::ResponseOptions>,
103 +) -> Result<SamplingParams, Status> {
104 + // Temperature is a top-level GenerateRequest field. Default to greedy (0.0) for
105 + // the gRPC API when the caller does not specify a value. This differs from
106 + // the HTTP/OpenAI API (which defaults to 1.0) and matches the convention of
107 + // programmatic generation APIs.
108 + let temperature = temperature.or(Some(0.0));
109 + let mut params = SamplingParams {
110 + temperature,
111 + ..SamplingParams::default()
112 + };
113 +
114 + // RandomSampling: for every remaining sampling field the protobuf default (`0`)
115 + // is treated as "unset" and leaves the resolved value to the lowering
116 + // stage, which falls back to the model-provided default or a
117 + // neutral/disabled value otherwise.
118 + if let Some(s) = sampling {
119 + // num_sequences (n > 1) is not supported yet by the TextLlm layer; the response
120 + // path also hardcodes SequenceOutput.index = 0, so accepting >1 would silently
121 + // truncate output cardinality. Reject explicitly.
122 + if s.num_sequences > 1 {
123 + return Err(Status::invalid_argument(
124 + "num_sequences > 1 is not supported",
125 + ));
126 + }
127 + if s.top_k != 0 {
128 + params.top_k = Some(s.top_k);
129 + }
130 + if s.top_p != 0.0 {
131 + params.top_p = Some(s.top_p);
132 + }
133 + if s.min_p != 0.0 {
134 + params.min_p = Some(s.min_p);
135 + }
136 + params.seed = s.seed;
137 + }
138 +
139 + // DecodingParameters
140 + if let Some(d) = decoding {
141 + if d.presence_penalty != 0.0 {
142 + params.presence_penalty = Some(d.presence_penalty);
143 + }
144 + if d.frequency_penalty != 0.0 {
145 + params.frequency_penalty = Some(d.frequency_penalty);
146 + }
147 + if d.repetition_penalty != 0.0 {
148 + params.repetition_penalty = Some(d.repetition_penalty);
149 + }
150 + if !d.logit_bias.is_empty() {
151 + params.logit_bias = Some(d.logit_bias.clone());
152 + }
153 + if !d.allowed_token_ids.is_empty() {
154 + params.allowed_token_ids = Some(d.allowed_token_ids.clone());
155 + }
156 + params.structured_outputs = convert_structured_output(d)?;
157 + }
158 +
159 + // StoppingCriteria
160 + if let Some(s) = stopping {
161 + if s.max_new_tokens != 0 {
162 + params.max_tokens = Some(s.max_new_tokens);
163 + }
164 + if s.min_new_tokens != 0 {
165 + params.min_tokens = Some(s.min_new_tokens);
166 + }
167 + if !s.stop_token_ids.is_empty() {
168 + params.stop_token_ids = Some(s.stop_token_ids.clone());
169 + }
170 + params.ignore_eos = s.ignore_eos;
171 + }
172 +
173 + // ResponseOptions → logprobs
174 + if let Some(r) = response {
175 + if r.output_logprobs {
176 + let (count, token_ids) = candidate_logprob_spec(r.output_candidates.as_ref());
177 + params.logprobs = Some(count);
178 + params.logprob_token_ids = token_ids;
179 + }
180 + if r.prompt_logprobs {
181 + // The engine-core protocol has only one shared `logprob_token_ids` field
182 + // for output and prompt logprobs, so a per-token-id selector for prompt
183 + // candidates can't be honored independently. Reject it instead of silently
184 + // dropping the list.
185 + if matches!(
186 + r.prompt_candidates.as_ref().and_then(|c| c.select.as_ref()),
187 + Some(pb::candidate_tokens::Select::TokenIds(_))
188 + ) {
189 + return Err(Status::invalid_argument(
190 + "prompt_candidates token_ids selector is not supported",
191 + ));
192 + }
193 + let (count, _) = candidate_logprob_spec(r.prompt_candidates.as_ref());
194 + params.prompt_logprobs = Some(count);
195 + }
196 + }
197 +
198 + Ok(params)
199 +}
200 +
201 +/// Map the proto `CandidateTokens` selector to a `(logprobs_count,
202 +/// logprob_token_ids)` pair.
203 +///
204 +/// - `top_n(k)` → `(k, None)` — return top-k candidates by probability
205 +/// - `all` → `(-1, None)` — return the full vocabulary
206 +/// - `token_ids(n)` → `(1, Some(vec of n token ids))` — return logprobs for specific tokens (the
207 +/// count `n` is stored in the proto as the number of token IDs that follow, but the actual IDs
208 +/// are carried via `logprob_token_ids` on `SamplingParams`)
209 +/// - absent → `(1, None)` — just the sampled/scored token
210 +fn candidate_logprob_spec(candidates: Option<&pb::CandidateTokens>) -> (i32, Option<Vec<u32>>) {
211 + match candidates.and_then(|c| c.select.as_ref()) {
212 + Some(pb::candidate_tokens::Select::TopN(n)) => (*n as i32, None),
213 + Some(pb::candidate_tokens::Select::All(true)) => (-1, None),
214 + Some(pb::candidate_tokens::Select::TokenIds(ids)) => (1, Some(ids.ids.clone())),
215 + _ => (1, None),
216 + }
217 +}
218 +
219 +fn convert_structured_output(
220 + d: &pb::DecodingParameters,
221 +) -> Result<Option<StructuredOutputsParams>, Status> {
222 + let so = match d.structured_output.as_ref() {
223 + None => return Ok(None),
224 + Some(so) => so,
225 + };
226 + use pb::decoding_parameters::StructuredOutput;
227 + let params = match so {
228 + StructuredOutput::Json(schema) => {
229 + let json: serde_json::Value = serde_json::from_str(schema)
230 + .map_err(|e| Status::invalid_argument(format!("invalid json schema: {e}")))?;
231 + StructuredOutputsParams {
232 + json: Some(json),
233 + ..Default::default()
234 + }
235 + }
236 + StructuredOutput::Regex(regex) => StructuredOutputsParams {
237 + regex: Some(regex.clone()),
238 + ..Default::default()
239 + },
240 + StructuredOutput::Choice(choices) => StructuredOutputsParams {
241 + choice: Some(choices.choices.clone()),
242 + ..Default::default()
243 + },
244 + StructuredOutput::Grammar(grammar) => StructuredOutputsParams {
245 + grammar: Some(grammar.clone()),
246 + ..Default::default()
247 + },
248 + StructuredOutput::JsonObject(true) => StructuredOutputsParams {
249 + json_object: Some(true),
250 + ..Default::default()
251 + },
252 + StructuredOutput::JsonObject(false) => return Ok(None),
253 + StructuredOutput::StructuralTag(tag) => StructuredOutputsParams {
254 + structural_tag: Some(tag.clone()),
255 + ..Default::default()
256 + },
257 + };
258 + Ok(Some(params))
259 +}
260 +
261 +// ========================================================================================
262 +// Response conversion
263 +// ========================================================================================
264 +
265 +/// Convert a `DecodedTextEvent::Start` into the prompt info portion of a gRPC
266 +/// response.
267 +pub fn to_prompt_info(
268 + prompt_token_ids: &[u32],
269 + prompt_logprobs: Option<&DecodedPromptLogprobs>,
270 + opts: &ResponseOpts,
271 +) -> pb::PromptInfo {
272 + let token_ids = if opts.prompt_token_ids {
273 + prompt_token_ids.to_vec()
274 + } else {
275 + vec![]
276 + };
277 +
278 + let (logprobs, ranks, candidate_tokens) = match prompt_logprobs {
279 + Some(plp) if opts.prompt_logprobs => prompt_logprobs_to_proto(plp),
280 + _ => (vec![], vec![], vec![]),
281 + };
282 +
283 + pb::PromptInfo {
284 + num_prompt_tokens: prompt_token_ids.len() as u32,
285 + token_ids,
286 + logprobs,
287 + ranks,
288 + candidate_tokens,
289 + }
290 +}
291 +
292 +/// Convert a `DecodedTextEvent::TextDelta` into a gRPC `SequenceOutput`.
293 +pub fn to_sequence_output(
294 + delta: &str,
295 + token_ids: &[u32],
296 + logprobs: Option<&DecodedLogprobs>,
297 + finished: Option<&Finished>,
298 + opts: &ResponseOpts,
299 +) -> pb::SequenceOutput {
300 + let (lp_values, rank_values, candidates) = match logprobs {
301 + Some(lp) if opts.output_logprobs => output_logprobs_to_proto(lp),
302 + _ => (vec![], vec![], vec![]),
303 + };
304 +
305 + pb::SequenceOutput {
306 + index: 0, // TODO: multi-sequence (n > 1) not supported
307 + text: if opts.output_text {
308 + delta.to_string()
309 + } else {
310 + String::new()
311 + },
312 + num_tokens: token_ids.len() as u32,
313 + token_ids: if opts.output_token_ids {
314 + token_ids.to_vec()
315 + } else {
316 + vec![]
317 + },
318 + logprobs: lp_values,
319 + ranks: rank_values,
320 + candidate_tokens: candidates,
321 + finish_info: finished.map(|f| to_finish_info(f, token_ids)),
322 + }
323 +}
324 +
325 +fn to_finish_info(finished: &Finished, token_ids: &[u32]) -> pb::FinishInfo {
326 + use pb::finish_info::FinishReason as PbFinishReason;
327 +
328 + let (finish_reason, stop_reason) = match &finished.finish_reason {
329 + FinishReason::Stop(reason) => {
330 + let sr = match reason {
331 + Some(StopReason::TokenId(id)) => {
332 + Some(pb::finish_info::StopReason::StopTokenId(*id))
333 + }
334 + Some(StopReason::Text(s)) => {
335 + Some(pb::finish_info::StopReason::StopString(s.clone()))
336 + }
337 + // EOS-driven stop: engine-core matched the primary EOS token id but did not
338 + // echo it back as a `stop_reason`. The matched token is, by construction, the
339 + // last token of the terminal output batch (see vllm's `check_stop` in
340 + // vllm/v1/core/sched/utils.py), so we recover it from there.
341 + None => token_ids.last().copied().map(pb::finish_info::StopReason::EosTokenId),
342 + };
343 + (PbFinishReason::Stop as i32, sr)
344 + }
345 + FinishReason::Length => (PbFinishReason::Length as i32, None),
346 + FinishReason::Abort | FinishReason::Error | FinishReason::Repetition => {
347 + (PbFinishReason::Aborted as i32, None)
348 + }
349 + };
350 +
351 + pb::FinishInfo {
352 + num_output_tokens: finished.output_token_count as u32,
353 + finish_reason,
354 + stop_reason,
355 + kv_transfer_params: finished.kv_transfer_params.as_ref().and_then(json_to_proto_struct),
356 + }
357 +}
358 +
359 +// ========================================================================================
360 +// Logprobs helpers
361 +// ========================================================================================
362 +
363 +/// Convert output logprobs to the flat proto representation.
364 +///
365 +/// Returns (logprob_values, ranks, candidate_tokens) — all parallel arrays
366 +/// indexed by position.
367 +fn output_logprobs_to_proto(
368 + lp: &DecodedLogprobs,
369 +) -> (Vec<f32>, Vec<u32>, Vec<pb::CandidateTokenInfo>) {
370 + positions_to_proto(&lp.positions)
371 +}
372 +
373 +/// Convert prompt logprobs to the flat proto representation.
374 +fn prompt_logprobs_to_proto(
375 + plp: &DecodedPromptLogprobs,
376 +) -> (Vec<f32>, Vec<u32>, Vec<pb::CandidateTokenInfo>) {
377 + // The proto PromptInfo has flat parallel arrays covering all prompt positions.
378 + // DecodedPromptLogprobs has first_token separately + scored_positions for the
379 + // rest. The first prompt position has no scores, so we emit zeros for it.
380 + let (mut logprobs, mut ranks, mut candidates) = positions_to_proto(&plp.scored_positions);
381 + logprobs.insert(0, 0.0);
382 + ranks.insert(0, 0);
383 + candidates.insert(0, pb::CandidateTokenInfo { tokens: vec![] });
384 + (logprobs, ranks, candidates)
385 +}
386 +
387 +/// Shared helper: convert a slice of decoded position logprobs to flat proto
388 +/// arrays.
389 +fn positions_to_proto(
390 + positions: &[vllm_text::DecodedPositionLogprobs],
391 +) -> (Vec<f32>, Vec<u32>, Vec<pb::CandidateTokenInfo>) {
392 + let mut logprobs = Vec::with_capacity(positions.len());
393 + let mut ranks = Vec::with_capacity(positions.len());
394 + let mut candidates = Vec::with_capacity(positions.len());
395 +
396 + for pos in positions {
397 + // First entry is the sampled/scored token.
398 + if let Some(first) = pos.entries.first() {
399 + logprobs.push(first.logprob);
400 + ranks.push(first.rank);
401 + }
402 +
403 + // Extra candidates beyond the first.
404 + let entries = pos.entries.iter().skip(1);
405 + candidates.push(pb::CandidateTokenInfo {
406 + tokens: entries
407 + .map(|e| pb::candidate_token_info::TokenInfo {
408 + id: e.token_id,
409 + logprob: e.logprob,
410 + rank: e.rank,
411 + })
412 + .collect(),
413 + });
414 + }
415 +
416 + (logprobs, ranks, candidates)
417 +}
418 +
419 +// ========================================================================================
420 +// KV transfer params conversion (serde_json::Value ↔ prost_types::Struct)
421 +// ========================================================================================
422 +
423 +fn proto_struct_to_json(s: &prost_types::Struct) -> serde_json::Value {
424 + serde_json::Value::Object(
425 + s.fields.iter().map(|(k, v)| (k.clone(), proto_value_to_json(v))).collect(),
426 + )
427 +}
428 +
429 +fn proto_value_to_json(v: &prost_types::Value) -> serde_json::Value {
430 + use prost_types::value::Kind;
431 + match v.kind.as_ref() {
432 + None | Some(Kind::NullValue(_)) => serde_json::Value::Null,
433 + Some(Kind::BoolValue(b)) => serde_json::Value::Bool(*b),
434 + Some(Kind::NumberValue(n)) => serde_json::json!(*n),
435 + Some(Kind::StringValue(s)) => serde_json::Value::String(s.clone()),
436 + Some(Kind::ListValue(list)) => {
437 + serde_json::Value::Array(list.values.iter().map(proto_value_to_json).collect())
438 + }
439 + Some(Kind::StructValue(s)) => proto_struct_to_json(s),
440 + }
441 +}
442 +
443 +fn json_to_proto_struct(value: &serde_json::Value) -> Option<prost_types::Struct> {
444 + match value {
445 + serde_json::Value::Object(map) => Some(prost_types::Struct {
446 + fields: map.iter().map(|(k, v)| (k.clone(), json_to_proto_value(v))).collect(),
447 + }),
448 + _ => None,
449 + }
450 +}
451 +
452 +fn json_to_proto_value(v: &serde_json::Value) -> prost_types::Value {
453 + use prost_types::value::Kind;
454 + let kind = match v {
455 + serde_json::Value::Null => Kind::NullValue(0),
456 + serde_json::Value::Bool(b) => Kind::BoolValue(*b),
457 + serde_json::Value::Number(n) => Kind::NumberValue(n.as_f64().unwrap_or(0.0)),
458 + serde_json::Value::String(s) => Kind::StringValue(s.clone()),
459 + serde_json::Value::Array(arr) => Kind::ListValue(prost_types::ListValue {
460 + values: arr.iter().map(json_to_proto_value).collect(),
461 + }),
462 + serde_json::Value::Object(map) => Kind::StructValue(prost_types::Struct {
463 + fields: map.iter().map(|(k, v)| (k.clone(), json_to_proto_value(v))).collect(),
464 + }),
465 + };
466 + prost_types::Value { kind: Some(kind) }
467 +}
468 +
469 +// ========================================================================================
470 +// Options extracted from the request for response building
471 +// ========================================================================================
472 +
473 +/// Response-shaping options extracted from the proto `ResponseOptions`.
474 +#[derive(Default)]
475 +pub struct ResponseOpts {
476 + pub prompt_token_ids: bool,
477 + pub prompt_logprobs: bool,
478 + pub output_text: bool,
479 + pub output_token_ids: bool,
480 + pub output_logprobs: bool,
481 +}
482 +
483 +impl ResponseOpts {
484 + pub fn from_proto(r: Option<&pb::ResponseOptions>) -> Self {
485 + match r {
486 + Some(r) => Self {
487 + prompt_token_ids: r.prompt_token_ids,
488 + prompt_logprobs: r.prompt_logprobs,
489 + output_text: r.output_text.unwrap_or(true),
490 + output_token_ids: r.output_token_ids,
491 + output_logprobs: r.output_logprobs,
492 + },
493 + None => Self {
494 + output_text: true,
495 + ..Default::default()
496 + },
497 + }
498 + }
499 +}
500 +
501 +#[cfg(test)]
502 +mod tests {
503 + use vllm_engine_core_client::protocol::StopReason;
504 + use vllm_text::{FinishReason, Finished, Prompt};
505 +
506 + use super::pb::finish_info::{FinishReason as PbFinishReason, StopReason as PbStopReason};
507 + use super::{ResponseOpts, pb, to_finish_info, to_sequence_output, to_text_request};
508 +
509 + fn base_request() -> pb::GenerateRequest {
510 + pb::GenerateRequest {
511 + request_id: "req".to_string(),
512 + model: "test-model".to_string(),
513 + prompt: Some(pb::generate_request::Prompt::Text("hi".to_string())),
514 + ..Default::default()
515 + }
516 + }
517 +
518 + #[test]
519 + fn temperature_propagates_from_top_level_request_field() {
520 + let req = pb::GenerateRequest {
521 + temperature: Some(0.7),
522 + ..base_request()
523 + };
524 + let text = to_text_request(req, false, &["test-model".to_string()]).expect("convert ok");
525 + assert_eq!(text.sampling_params.temperature, Some(0.7));
526 + }
527 +
528 + #[test]
529 + fn unset_temperature_defaults_to_greedy() {
530 + let text = to_text_request(base_request(), false, &["test-model".to_string()])
531 + .expect("convert ok");
532 + // The gRPC API defaults to greedy (0.0) when temperature is not specified.
533 + assert_eq!(text.sampling_params.temperature, Some(0.0));
534 + }
535 +
536 + #[test]
537 + fn absent_seed_is_none() {
538 + let req = pb::GenerateRequest {
539 + sampling: Some(pb::RandomSampling {
540 + seed: None,
541 + ..Default::default()
542 + }),
543 + ..base_request()
544 + };
545 + let text = to_text_request(req, false, &["test-model".to_string()]).expect("convert ok");
546 + assert_eq!(text.sampling_params.seed, None);
547 + }
548 +
549 + #[test]
550 + fn zero_seed_is_valid() {
551 + let req = pb::GenerateRequest {
552 + sampling: Some(pb::RandomSampling {
553 + seed: Some(0),
554 + ..Default::default()
555 + }),
556 + ..base_request()
557 + };
558 + let text = to_text_request(req, false, &["test-model".to_string()]).expect("convert ok");
559 + assert_eq!(text.sampling_params.seed, Some(0));
560 + }
561 +
562 + #[test]
563 + fn bypass_prefix_cache_maps_to_skip_reading_prefix_cache() {
564 + let req = pb::GenerateRequest {
565 + kv: Some(pb::KvCacheParameters {
566 + bypass_prefix_cache: true,
567 + ..Default::default()
568 + }),
569 + ..base_request()
570 + };
571 + let text = to_text_request(req, false, &["test-model".to_string()]).expect("convert ok");
572 + assert_eq!(text.sampling_params.skip_reading_prefix_cache, Some(true));
573 + }
574 +
575 + #[test]
576 + fn bypass_prefix_cache_false_leaves_field_unset() {
577 + let req = pb::GenerateRequest {
578 + kv: Some(pb::KvCacheParameters {
579 + bypass_prefix_cache: false,
580 + ..Default::default()
581 + }),
582 + ..base_request()
583 + };
584 + let text = to_text_request(req, false, &["test-model".to_string()]).expect("convert ok");
585 + assert_eq!(text.sampling_params.skip_reading_prefix_cache, None);
586 + // Prompt conversion still succeeds and reaches the expected variant.
587 + assert!(matches!(text.prompt, Prompt::Text(s) if s == "hi"));
588 + }
589 +
590 + fn finished(reason: FinishReason) -> Finished {
591 + Finished {
592 + prompt_token_count: 0,
593 + output_token_count: 0,
594 + finish_reason: reason,
595 + kv_transfer_params: None,
596 + }
597 + }
598 +
599 + #[test]
600 + fn eos_stop_reports_last_output_token_as_eos_id() {
601 + let fin = finished(FinishReason::Stop(None));
602 + let token_ids = [1_u32, 2, 3, 151643];
603 +
604 + let info = to_finish_info(&fin, &token_ids);
605 +
606 + assert_eq!(info.finish_reason, PbFinishReason::Stop as i32);
607 + assert_eq!(info.stop_reason, Some(PbStopReason::EosTokenId(151643)));
608 + }
609 +
610 + #[test]
611 + fn eos_stop_with_empty_token_ids_leaves_stop_reason_unset() {
612 + let fin = finished(FinishReason::Stop(None));
613 +
614 + let info = to_finish_info(&fin, &[]);
615 +
616 + assert_eq!(info.finish_reason, PbFinishReason::Stop as i32);
617 + assert_eq!(info.stop_reason, None);
618 + }
619 +
620 + #[test]
621 + fn explicit_stop_token_id_is_preserved() {
622 + let fin = finished(FinishReason::Stop(Some(StopReason::TokenId(42))));
623 + // Terminal token list should be ignored when an explicit stop reason is
624 + // present.
625 + let info = to_finish_info(&fin, &[7, 42]);
626 +
627 + assert_eq!(info.finish_reason, PbFinishReason::Stop as i32);
628 + assert_eq!(info.stop_reason, Some(PbStopReason::StopTokenId(42)));
629 + }
630 +
631 + #[test]
632 + fn explicit_stop_string_is_preserved() {
633 + let fin = finished(FinishReason::Stop(Some(StopReason::Text("</stop>".into()))));
634 +
635 + let info = to_finish_info(&fin, &[1, 2, 3]);
636 +
637 + assert_eq!(info.finish_reason, PbFinishReason::Stop as i32);
638 + assert_eq!(
639 + info.stop_reason,
640 + Some(PbStopReason::StopString("</stop>".into()))
641 + );
642 + }
643 +
644 + #[test]
645 + fn length_finish_has_no_stop_reason() {
646 + let fin = finished(FinishReason::Length);
647 +
648 + let info = to_finish_info(&fin, &[1, 2, 3]);
649 +
650 + assert_eq!(info.finish_reason, PbFinishReason::Length as i32);
651 + assert_eq!(info.stop_reason, None);
652 + }
653 +
654 + #[test]
655 + fn abort_finish_is_mapped_to_aborted() {
656 + let fin = finished(FinishReason::Abort);
657 +
658 + let info = to_finish_info(&fin, &[]);
659 +
660 + assert_eq!(info.finish_reason, PbFinishReason::Aborted as i32);
661 + assert_eq!(info.stop_reason, None);
662 + }
663 +
664 + #[test]
665 + fn to_sequence_output_threads_token_ids_into_eos_id() {
666 + let fin = finished(FinishReason::Stop(None));
667 + let opts = ResponseOpts {
668 + output_text: true,
669 + output_token_ids: true,
670 + ..Default::default()
671 + };
672 +
673 + let out = to_sequence_output("hello", &[10, 20, 30], None, Some(&fin), &opts);
674 +
675 + let finish = out.finish_info.expect("finish_info should be present");
676 + assert_eq!(finish.finish_reason, PbFinishReason::Stop as i32);
677 + assert_eq!(finish.stop_reason, Some(PbStopReason::EosTokenId(30)));
678 + }
679 +}
rust/src/server/src/grpc/mod.rs new
+158
@@ -0,0 +1,158 @@
1 +//! gRPC Generate service backed by the shared [`vllm_text::TextLlm`] facade.
2 +
3 +mod convert;
4 +
5 +use std::pin::Pin;
6 +use std::sync::Arc;
7 +
8 +use futures::{Stream, StreamExt as _};
9 +use thiserror_ext::AsReport as _;
10 +use tokio::sync::mpsc;
11 +use tokio_stream::wrappers::ReceiverStream;
12 +use tonic::{Request, Response, Status};
13 +use tracing::info;
14 +use vllm_text::{DecodedTextEvent, TextOutputStreamExt as _};
15 +
16 +use self::convert::ResponseOpts;
17 +use crate::state::AppState;
18 +
19 +/// Generated protobuf/gRPC types for the `vllm` package.
20 +pub mod pb {
21 + tonic::include_proto!("vllm");
22 +}
23 +
24 +pub use pb::generate_server::GenerateServer;
25 +
26 +#[cfg(test)]
27 +mod tests;
28 +
29 +/// gRPC Generate service implementation backed by the shared application state.
30 +pub struct GenerateServiceImpl {
31 + state: Arc<AppState>,
32 +}
33 +
34 +impl GenerateServiceImpl {
35 + pub fn new(state: Arc<AppState>) -> Self {
36 + Self { state }
37 + }
38 +}
39 +
40 +#[tonic::async_trait]
41 +impl pb::generate_server::Generate for GenerateServiceImpl {
42 + type GenerateStreamStream =
43 + Pin<Box<dyn Stream<Item = Result<pb::GenerateResponse, Status>> + Send>>;
44 +
45 + /// Unary generate: collect all output and return a single response.
46 + async fn generate(
47 + &self,
48 + request: Request<pb::GenerateRequest>,
49 + ) -> Result<Response<pb::GenerateResponse>, Status> {
50 + let proto_req = request.into_inner();
51 + let response_opts = ResponseOpts::from_proto(proto_req.response.as_ref());
52 + let text_request =
53 + convert::to_text_request(proto_req, false, self.state.served_model_names())?;
54 +
55 + let request_id = text_request.request_id.clone();
56 + info!(%request_id, "grpc generate (unary)");
57 +
58 + let stream = self.state.chat.text().generate(text_request).await;
59 + let stream = stream.map_err(|e| Status::internal(e.to_report_string()))?;
60 +
61 + let collected = stream
62 + .collect_output()
63 + .await
64 + .map_err(|e| Status::internal(e.to_report_string()))?;
65 +
66 + // Build the single aggregated response.
67 + let prompt_info = convert::to_prompt_info(
68 + &collected.prompt_token_ids,
69 + collected.prompt_logprobs.as_ref(),
70 + &response_opts,
71 + );
72 +
73 + let finish_info = vllm_text::Finished {
74 + prompt_token_count: collected.prompt_token_ids.len(),
75 + output_token_count: collected.token_ids.len(),
76 + finish_reason: collected.finish_reason,
77 + kv_transfer_params: collected.kv_transfer_params,
78 + };
79 +
80 + let outputs = convert::to_sequence_output(
81 + &collected.text,
82 + &collected.token_ids,
83 + collected.logprobs.as_ref(),
84 + Some(&finish_info),
85 + &response_opts,
86 + );
87 +
88 + Ok(Response::new(pb::GenerateResponse {
89 + prompt_info: Some(prompt_info),
90 + outputs: Some(outputs),
91 + }))
92 + }
93 +
94 + /// Streaming generate: yield incremental responses as tokens are produced.
95 + async fn generate_stream(
96 + &self,
97 + request: Request<pb::GenerateRequest>,
98 + ) -> Result<Response<Self::GenerateStreamStream>, Status> {
99 + let proto_req = request.into_inner();
100 + let response_opts = ResponseOpts::from_proto(proto_req.response.as_ref());
101 + let text_request =
102 + convert::to_text_request(proto_req, true, self.state.served_model_names())?;
103 +
104 + let request_id = text_request.request_id.clone();
105 + info!(%request_id, "grpc generate (stream)");
106 +
107 + let stream = self.state.chat.text().generate(text_request).await;
108 + let stream = stream.map_err(|e| Status::internal(e.to_report_string()))?;
109 +
110 + let (tx, rx) = mpsc::channel(32);
111 +
112 + tokio::spawn(async move {
113 + futures::pin_mut!(stream);
114 + while let Some(event) = stream.next().await {
115 + let response = match event {
116 + Err(e) => Err(Status::internal(e.to_report_string())),
117 + Ok(DecodedTextEvent::Start {
118 + prompt_token_ids,
119 + prompt_logprobs,
120 + }) => {
121 + let prompt_info = convert::to_prompt_info(
122 + &prompt_token_ids,
123 + prompt_logprobs.as_ref(),
124 + &response_opts,
125 + );
126 + Ok(pb::GenerateResponse {
127 + prompt_info: Some(prompt_info),
128 + outputs: None,
129 + })
130 + }
131 + Ok(DecodedTextEvent::TextDelta {
132 + delta,
133 + token_ids,
134 + logprobs,
135 + finished,
136 + }) => Ok(pb::GenerateResponse {
137 + prompt_info: None,
138 + outputs: Some(convert::to_sequence_output(
139 + &delta,
140 + &token_ids,
141 + logprobs.as_ref(),
142 + finished.as_ref(),
143 + &response_opts,
144 + )),
145 + }),
146 + };
147 +
148 + if tx.send(response).await.is_err() {
149 + // Client disconnected.
150 + break;
151 + }
152 + }
153 + });
154 +
155 + let response_stream = ReceiverStream::new(rx);
156 + Ok(Response::new(Box::pin(response_stream)))
157 + }
158 +}
rust/src/server/src/grpc/tests.rs new
+662
@@ -0,0 +1,662 @@
1 +use std::future::Future;
2 +use std::pin::Pin;
3 +use std::sync::Arc;
4 +use std::task::{Context, Poll};
5 +
6 +use futures::StreamExt as _;
7 +use serial_test::serial;
8 +use tonic::transport::Server as TonicServer;
9 +use vllm_chat::{
10 + ChatBackend, ChatLlm, ChatRenderer, ChatRequest, ChatTextBackend, DefaultChatOutputProcessor,
11 + DynChatOutputProcessor, DynChatRenderer, NewChatOutputProcessorOptions, RenderedPrompt,
12 +};
13 +use vllm_engine_core_client::protocol::{
14 + EngineCoreFinishReason, EngineCoreOutput, EngineCoreOutputs, EngineCoreRequest,
15 +};
16 +use vllm_engine_core_client::test_utils::{IpcNamespace, spawn_mock_engine_task};
17 +use vllm_engine_core_client::{EngineCoreClient, EngineCoreClientConfig, EngineId};
18 +use vllm_llm::Llm;
19 +use vllm_text::tokenizer::{DynTokenizer, Tokenizer};
20 +use vllm_text::{Prompt, TextBackend};
21 +use zeromq::prelude::{SocketRecv, SocketSend};
22 +use zeromq::{DealerSocket, PushSocket, ZmqMessage};
23 +
24 +use super::pb::generate_client::GenerateClient;
25 +use super::{GenerateServer, GenerateServiceImpl, pb};
26 +use crate::state::AppState;
27 +
28 +// ========================================================================================
29 +// Helpers (mirrors the patterns in routes/tests.rs)
30 +// ========================================================================================
31 +
32 +type TestFuture<'a> = Pin<Box<dyn Future<Output = ()> + Send + 'a>>;
33 +
34 +fn boxed_test_future<'a>(future: impl Future<Output = ()> + Send + 'a) -> TestFuture<'a> {
35 + Box::pin(future)
36 +}
37 +
38 +struct MockEngineTask {
39 + shutdown_tx: Option<tokio::sync::oneshot::Sender<()>>,
40 + join_handle: Option<tokio::task::JoinHandle<()>>,
41 +}
42 +
43 +impl MockEngineTask {
44 + fn new(
45 + (shutdown_tx, join_handle): (
46 + tokio::sync::oneshot::Sender<()>,
47 + tokio::task::JoinHandle<()>,
48 + ),
49 + ) -> Self {
50 + Self {
51 + shutdown_tx: Some(shutdown_tx),
52 + join_handle: Some(join_handle),
53 + }
54 + }
55 +}
56 +
57 +impl Future for MockEngineTask {
58 + type Output = Result<(), tokio::task::JoinError>;
59 +
60 + fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
61 + if let Some(shutdown_tx) = self.shutdown_tx.take() {
62 + let _ = shutdown_tx.send(());
63 + }
64 + match self.join_handle.as_mut() {
65 + Some(join_handle) => Pin::new(join_handle).poll(cx),
66 + None => Poll::Ready(Ok(())),
67 + }
68 + }
69 +}
70 +
71 +impl Drop for MockEngineTask {
72 + fn drop(&mut self) {
73 + if let Some(join_handle) = &self.join_handle {
74 + join_handle.abort();
75 + }
76 + }
77 +}
78 +
79 +fn request_output(
80 + request_id: &str,
81 + new_token_ids: Vec<u32>,
82 + finish_reason: Option<EngineCoreFinishReason>,
83 +) -> EngineCoreOutput {
84 + EngineCoreOutput {
85 + request_id: request_id.to_string(),
86 + new_token_ids,
87 + new_logprobs: None,
88 + new_prompt_logprobs_tensors: None,
89 + pooling_output: None,
90 + finish_reason,
91 + stop_reason: None,
92 + events: None,
93 + kv_transfer_params: None,
94 + trace_headers: None,
95 + prefill_stats: None,
96 + routed_experts: None,
97 + num_nans_in_logits: 0,
98 + }
99 +}
100 +
101 +fn engine_outputs_for_request(
102 + request_id: &str,
103 + output_specs: Vec<(Vec<u32>, Option<EngineCoreFinishReason>)>,
104 +) -> EngineCoreOutputs {
105 + EngineCoreOutputs {
106 + engine_index: 0,
107 + outputs: output_specs
108 + .into_iter()
109 + .map(|(token_ids, finish_reason)| request_output(request_id, token_ids, finish_reason))
110 + .collect(),
111 + scheduler_stats: None,
112 + timestamp: 0.0,
113 + utility_output: None,
114 + finished_requests: None,
115 + wave_complete: None,
116 + start_wave: None,
117 + }
118 +}
119 +
120 +fn default_stream_output_specs() -> Vec<(Vec<u32>, Option<EngineCoreFinishReason>)> {
121 + vec![
122 + (vec![b'h' as u32], None),
123 + (vec![b'i' as u32], None),
124 + (vec![b'!' as u32], Some(EngineCoreFinishReason::Stop)),
125 + ]
126 +}
127 +
128 +async fn send_outputs(push: &mut PushSocket, outputs: EngineCoreOutputs) {
129 + push.send(ZmqMessage::from(
130 + rmp_serde::to_vec_named(&outputs).expect("encode outputs"),
131 + ))
132 + .await
133 + .expect("send outputs");
134 +}
135 +
136 +async fn recv_engine_message(dealer: &mut DealerSocket) -> Vec<bytes::Bytes> {
137 + dealer.recv().await.expect("recv engine message").into_vec()
138 +}
139 +
140 +fn test_llm(client: EngineCoreClient) -> Llm {
141 + Llm::new(client).with_request_id_randomization(false)
142 +}
143 +
144 +#[derive(Clone, Debug)]
145 +struct FakeTextBackend;
146 +
147 +#[derive(Debug)]
148 +struct FakeTokenizer;
149 +
150 +impl Tokenizer for FakeTokenizer {
151 + fn encode(
152 + &self,
153 + text: &str,
154 + _add_special_tokens: bool,
155 + ) -> vllm_text::tokenizer::Result<Vec<u32>> {
156 + Ok(text.bytes().map(u32::from).collect())
157 + }
158 +
159 + fn decode(
160 + &self,
161 + token_ids: &[u32],
162 + _skip_special_tokens: bool,
163 + ) -> vllm_text::tokenizer::Result<String> {
164 + Ok(
165 + String::from_utf8_lossy(&token_ids.iter().map(|id| *id as u8).collect::<Vec<_>>())
166 + .into_owned(),
167 + )
168 + }
169 +
170 + fn token_to_id(&self, token: &str) -> Option<u32> {
171 + token.bytes().next().map(u32::from)
172 + }
173 +}
174 +
175 +impl TextBackend for FakeTextBackend {
176 + fn tokenizer(&self) -> DynTokenizer {
177 + Arc::new(FakeTokenizer)
178 + }
179 +
180 + fn model_id(&self) -> &str {
181 + "test-model"
182 + }
183 +}
184 +
185 +impl ChatBackend for FakeTextBackend {
186 + fn chat_renderer(&self) -> DynChatRenderer {
187 + Arc::new(self.clone())
188 + }
189 +
190 + fn new_chat_output_processor(
191 + &self,
192 + request: &mut ChatRequest,
193 + options: NewChatOutputProcessorOptions<'_>,
194 + ) -> vllm_chat::Result<DynChatOutputProcessor> {
195 + Ok(Box::new(DefaultChatOutputProcessor::new(
196 + request,
197 + self.model_id(),
198 + self.tokenizer(),
199 + options.tool_call_parser,
200 + options.reasoning_parser,
201 + )?))
202 + }
203 +}
204 +
205 +impl ChatRenderer for FakeTextBackend {
206 + fn render(&self, _request: &ChatRequest) -> vllm_chat::Result<RenderedPrompt> {
207 + Ok(RenderedPrompt {
208 + prompt: Prompt::Text(String::new()),
209 + })
210 + }
211 +}
212 +
213 +/// Spin up a gRPC server backed by a mock engine that serves a single request
214 +/// with the given output specs. Returns the client, the gRPC server task, and
215 +/// the mock engine task.
216 +async fn grpc_test_server(
217 + engine_id: impl Into<EngineId>,
218 + output_specs: Vec<(Vec<u32>, Option<EngineCoreFinishReason>)>,
219 +) -> (
220 + GenerateClient<tonic::transport::Channel>,
221 + tokio::task::JoinHandle<()>,
222 + MockEngineTask,
223 +) {
224 + let ipc = IpcNamespace::new().expect("create ipc namespace");
225 + let handshake_address = ipc.handshake_endpoint();
226 + let engine_id = engine_id.into();
227 +
228 + let engine_task = MockEngineTask::new(spawn_mock_engine_task(
229 + handshake_address.clone(),
230 + engine_id.clone(),
231 + move |dealer, push| {
232 + boxed_test_future(async move {
233 + let add = recv_engine_message(dealer).await;
234 + let request: EngineCoreRequest =
235 + rmp_serde::from_slice(&add[1]).expect("decode request");
236 + send_outputs(
237 + push,
238 + engine_outputs_for_request(&request.request_id, output_specs),
239 + )
240 + .await;
241 + })
242 + },
243 + ));
244 +
245 + let client = EngineCoreClient::connect(
246 + EngineCoreClientConfig::new_single(handshake_address)
247 + .with_model_name("test-model")
248 + .with_local_input_output_addresses(
249 + Some(ipc.input_endpoint()),
250 + Some(ipc.output_endpoint()),
251 + ),
252 + )
253 + .await
254 + .expect("connect client");
255 +
256 + let chat = ChatLlm::from_shared_backend(
257 + test_llm(client),
258 + Arc::new(FakeTextBackend) as Arc<dyn ChatTextBackend>,
259 + );
260 + let state = Arc::new(AppState::new(vec!["test-model".to_string()], chat));
261 + let svc = GenerateServer::new(GenerateServiceImpl::new(state));
262 +
263 + // Bind to an OS-assigned port.
264 + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.expect("bind grpc listener");
265 + let addr = listener.local_addr().expect("local addr");
266 +
267 + let server_task = tokio::spawn(async move {
268 + let incoming = tokio_stream::wrappers::TcpListenerStream::new(listener);
269 + TonicServer::builder()
270 + .add_service(svc)
271 + .serve_with_incoming(incoming)
272 + .await
273 + .expect("grpc server");
274 + });
275 +
276 + // Connect the client.
277 + let grpc_client = GenerateClient::connect(format!("http://{addr}"))
278 + .await
279 + .expect("connect grpc client");
280 +
281 + (grpc_client, server_task, engine_task)
282 +}
283 +
284 +// ========================================================================================
285 +// Tests
286 +// ========================================================================================
287 +
288 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
289 +#[serial]
290 +async fn unary_generate_returns_collected_text() {
291 + let (mut client, server_task, engine_task) =
292 + grpc_test_server(b"engine-grpc-unary", default_stream_output_specs()).await;
293 +
294 + let response = client
295 + .generate(pb::GenerateRequest {
296 + request_id: "test-unary-1".to_string(),
297 + model: "test-model".to_string(),
298 + prompt: Some(pb::generate_request::Prompt::Text("hello".to_string())),
299 + stopping: Some(pb::StoppingCriteria {
300 + max_new_tokens: 10,
301 + ..Default::default()
302 + }),
303 + response: Some(pb::ResponseOptions {
304 + output_text: Some(true),
305 + ..Default::default()
306 + }),
307 + ..Default::default()
308 + })
309 + .await
310 + .expect("unary generate")
311 + .into_inner();
312 +
313 + // Unary collects all tokens into one response.
314 + let outputs = response.outputs.expect("outputs present");
315 + assert_eq!(outputs.text, "hi");
316 +
317 + let finish = outputs.finish_info.expect("finish_info present");
318 + assert_eq!(
319 + finish.finish_reason,
320 + pb::finish_info::FinishReason::Stop as i32
321 + );
322 + assert_eq!(finish.num_output_tokens, 3);
323 +
324 + let prompt = response.prompt_info.expect("prompt_info present");
325 + assert_eq!(prompt.num_prompt_tokens, 5); // "hello" = 5 bytes
326 +
327 + engine_task.await.expect("mock engine task");
328 + server_task.abort();
329 +}
330 +
331 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
332 +#[serial]
333 +async fn unary_generate_with_token_ids_prompt() {
334 + let (mut client, server_task, engine_task) =
335 + grpc_test_server(b"engine-grpc-token-ids", default_stream_output_specs()).await;
336 +
337 + let response = client
338 + .generate(pb::GenerateRequest {
339 + request_id: "test-token-ids".to_string(),
340 + model: "test-model".to_string(),
341 + prompt: Some(pb::generate_request::Prompt::TokenIds(pb::TokenIds {
342 + ids: vec![1, 2, 3],
343 + })),
344 + stopping: Some(pb::StoppingCriteria {
345 + max_new_tokens: 10,
346 + ..Default::default()
347 + }),
348 + ..Default::default()
349 + })
350 + .await
351 + .expect("unary generate with token ids")
352 + .into_inner();
353 +
354 + let outputs = response.outputs.expect("outputs present");
355 + assert_eq!(outputs.text, "hi");
356 + assert_eq!(
357 + response.prompt_info.expect("prompt_info").num_prompt_tokens,
358 + 3
359 + );
360 +
361 + engine_task.await.expect("mock engine task");
362 + server_task.abort();
363 +}
364 +
365 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
366 +#[serial]
367 +async fn unary_generate_returns_token_ids_when_requested() {
368 + let (mut client, server_task, engine_task) =
369 + grpc_test_server(b"engine-grpc-tok-resp", default_stream_output_specs()).await;
370 +
371 + let response = client
372 + .generate(pb::GenerateRequest {
373 + request_id: "test-tok-resp".to_string(),
374 + model: "test-model".to_string(),
375 + prompt: Some(pb::generate_request::Prompt::Text("hi".to_string())),
376 + stopping: Some(pb::StoppingCriteria {
377 + max_new_tokens: 10,
378 + ..Default::default()
379 + }),
380 + response: Some(pb::ResponseOptions {
381 + output_text: Some(true),
382 + output_token_ids: true,
383 + prompt_token_ids: true,
384 + ..Default::default()
385 + }),
386 + ..Default::default()
387 + })
388 + .await
389 + .expect("unary generate")
390 + .into_inner();
391 +
392 + let outputs = response.outputs.expect("outputs present");
393 + assert_eq!(
394 + outputs.token_ids,
395 + vec![b'h' as u32, b'i' as u32, b'!' as u32]
396 + );
397 +
398 + let prompt = response.prompt_info.expect("prompt_info present");
399 + assert_eq!(prompt.token_ids, vec![b'h' as u32, b'i' as u32]);
400 +
401 + engine_task.await.expect("mock engine task");
402 + server_task.abort();
403 +}
404 +
405 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
406 +#[serial]
407 +async fn unary_generate_missing_prompt_returns_invalid_argument() {
408 + let (mut client, server_task, _engine_task) =
409 + grpc_test_server(b"engine-grpc-no-prompt", default_stream_output_specs()).await;
410 +
411 + let status = client
412 + .generate(pb::GenerateRequest {
413 + request_id: "test-no-prompt".to_string(),
414 + model: "test-model".to_string(),
415 + prompt: None,
416 + ..Default::default()
417 + })
418 + .await
419 + .expect_err("should fail without prompt");
420 +
421 + assert_eq!(status.code(), tonic::Code::InvalidArgument);
422 + assert!(status.message().contains("prompt"));
423 +
424 + server_task.abort();
425 +}
426 +
427 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
428 +#[serial]
429 +async fn streaming_generate_yields_incremental_responses() {
430 + let (mut client, server_task, engine_task) =
431 + grpc_test_server(b"engine-grpc-stream", default_stream_output_specs()).await;
432 +
433 + let stream = client
434 + .generate_stream(pb::GenerateRequest {
435 + request_id: "test-stream-1".to_string(),
436 + model: "test-model".to_string(),
437 + prompt: Some(pb::generate_request::Prompt::Text("hello".to_string())),
438 + stopping: Some(pb::StoppingCriteria {
439 + max_new_tokens: 10,
440 + ..Default::default()
441 + }),
442 + response: Some(pb::ResponseOptions {
443 + output_text: Some(true),
444 + ..Default::default()
445 + }),
446 + ..Default::default()
447 + })
448 + .await
449 + .expect("streaming generate")
450 + .into_inner();
451 +
452 + let responses: Vec<pb::GenerateResponse> =
453 + stream.map(|r| r.expect("stream item")).collect().await;
454 +
455 + // First response carries prompt info, subsequent ones carry output deltas.
456 + assert!(
457 + responses.len() >= 2,
458 + "expected at least 2 streamed responses, got {}",
459 + responses.len()
460 + );
461 +
462 + // First message should have prompt info.
463 + let first = &responses[0];
464 + let prompt_info = first.prompt_info.as_ref().expect("first response has prompt_info");
465 + assert_eq!(prompt_info.num_prompt_tokens, 5); // "hello"
466 +
467 + // Collect all text deltas.
468 + let full_text: String = responses
469 + .iter()
470 + .filter_map(|r| r.outputs.as_ref())
471 + .map(|o| o.text.as_str())
472 + .collect();
473 + assert_eq!(full_text, "hi");
474 +
475 + // Last output response should have finish info.
476 + let last_output = responses
477 + .iter()
478 + .rev()
479 + .find_map(|r| r.outputs.as_ref())
480 + .expect("at least one output");
481 + let finish = last_output.finish_info.as_ref().expect("finish_info on last output");
482 + assert_eq!(
483 + finish.finish_reason,
484 + pb::finish_info::FinishReason::Stop as i32
485 + );
486 +
487 + engine_task.await.expect("mock engine task");
488 + server_task.abort();
489 +}
490 +
491 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
492 +#[serial]
493 +async fn streaming_generate_missing_prompt_returns_invalid_argument() {
494 + let (mut client, server_task, _engine_task) = grpc_test_server(
495 + b"engine-grpc-stream-no-prompt",
496 + default_stream_output_specs(),
497 + )
498 + .await;
499 +
500 + let status = client
501 + .generate_stream(pb::GenerateRequest {
502 + request_id: "test-stream-no-prompt".to_string(),
503 + model: "test-model".to_string(),
504 + prompt: None,
505 + ..Default::default()
506 + })
507 + .await
508 + .expect_err("should fail without prompt");
509 +
510 + assert_eq!(status.code(), tonic::Code::InvalidArgument);
511 +
512 + server_task.abort();
513 +}
514 +
515 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
516 +#[serial]
517 +async fn unary_generate_with_sampling_params() {
518 + let (mut client, server_task, engine_task) =
519 + grpc_test_server(b"engine-grpc-sampling", default_stream_output_specs()).await;
520 +
521 + let response = client
522 + .generate(pb::GenerateRequest {
523 + request_id: "test-sampling".to_string(),
524 + model: "test-model".to_string(),
525 + prompt: Some(pb::generate_request::Prompt::Text("test".to_string())),
526 + temperature: Some(0.7),
527 + sampling: Some(pb::RandomSampling {
528 + top_k: 50,
529 + top_p: 0.9,
530 + seed: Some(42),
531 + ..Default::default()
532 + }),
533 + stopping: Some(pb::StoppingCriteria {
534 + max_new_tokens: 5,
535 + ..Default::default()
536 + }),
537 + ..Default::default()
538 + })
539 + .await
540 + .expect("generate with sampling params")
541 + .into_inner();
542 +
543 + // Verify the request was accepted and produced output.
544 + let outputs = response.outputs.expect("outputs present");
545 + assert_eq!(outputs.text, "hi");
546 +
547 + engine_task.await.expect("mock engine task");
548 + server_task.abort();
549 +}
550 +
551 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
552 +#[serial]
553 +async fn unary_generate_rejects_wrong_model() {
554 + let (mut client, server_task, _engine_task) =
555 + grpc_test_server(b"engine-grpc-wrong-model", default_stream_output_specs()).await;
556 +
557 + let status = client
558 + .generate(pb::GenerateRequest {
559 + request_id: "test-wrong-model".to_string(),
560 + model: "other-model".to_string(),
561 + prompt: Some(pb::generate_request::Prompt::Text("hi".to_string())),
562 + stopping: Some(pb::StoppingCriteria {
563 + max_new_tokens: 10,
564 + ..Default::default()
565 + }),
566 + ..Default::default()
567 + })
568 + .await
569 + .expect_err("should fail with wrong model");
570 +
571 + assert_eq!(status.code(), tonic::Code::NotFound);
572 + assert!(status.message().contains("other-model"));
573 +
574 + server_task.abort();
575 +}
576 +
577 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
578 +#[serial]
579 +async fn streaming_generate_rejects_wrong_model() {
580 + let (mut client, server_task, _engine_task) = grpc_test_server(
581 + b"engine-grpc-stream-wrong-model",
582 + default_stream_output_specs(),
583 + )
584 + .await;
585 +
586 + let status = client
587 + .generate_stream(pb::GenerateRequest {
588 + request_id: "test-stream-wrong-model".to_string(),
589 + model: "other-model".to_string(),
590 + prompt: Some(pb::generate_request::Prompt::Text("hi".to_string())),
591 + stopping: Some(pb::StoppingCriteria {
592 + max_new_tokens: 10,
593 + ..Default::default()
594 + }),
595 + ..Default::default()
596 + })
597 + .await
598 + .expect_err("should fail with wrong model");
599 +
600 + assert_eq!(status.code(), tonic::Code::NotFound);
601 + assert!(status.message().contains("other-model"));
602 +
603 + server_task.abort();
604 +}
605 +
606 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
607 +#[serial]
608 +async fn unary_generate_accepts_empty_model() {
609 + let (mut client, server_task, engine_task) =
610 + grpc_test_server(b"engine-grpc-empty-model", default_stream_output_specs()).await;
611 +
612 + // Empty `model` (proto3 default) is treated as "unset" and should be accepted.
613 + let response = client
614 + .generate(pb::GenerateRequest {
615 + request_id: "test-empty-model".to_string(),
616 + model: String::new(),
617 + prompt: Some(pb::generate_request::Prompt::Text("hi".to_string())),
618 + stopping: Some(pb::StoppingCriteria {
619 + max_new_tokens: 10,
620 + ..Default::default()
621 + }),
622 + ..Default::default()
623 + })
624 + .await
625 + .expect("unary generate with empty model")
626 + .into_inner();
627 +
628 + let outputs = response.outputs.expect("outputs present");
629 + assert_eq!(outputs.text, "hi");
630 +
631 + engine_task.await.expect("mock engine task");
632 + server_task.abort();
633 +}
634 +
635 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
636 +#[serial]
637 +async fn unary_generate_output_text_defaults_to_true() {
638 + let (mut client, server_task, engine_task) =
639 + grpc_test_server(b"engine-grpc-default-text", default_stream_output_specs()).await;
640 +
641 + // No response options at all — output_text should default to true.
642 + let response = client
643 + .generate(pb::GenerateRequest {
644 + request_id: "test-default-text".to_string(),
645 + model: "test-model".to_string(),
646 + prompt: Some(pb::generate_request::Prompt::Text("x".to_string())),
647 + stopping: Some(pb::StoppingCriteria {
648 + max_new_tokens: 10,
649 + ..Default::default()
650 + }),
651 + ..Default::default()
652 + })
653 + .await
654 + .expect("unary generate")
655 + .into_inner();
656 +
657 + let outputs = response.outputs.expect("outputs present");
658 + assert_eq!(outputs.text, "hi");
659 +
660 + engine_task.await.expect("mock engine task");
661 + server_task.abort();
662 +}
rust/src/server/src/lib.rs new
+237
@@ -0,0 +1,237 @@
1 +//! Minimal OpenAI-compatible HTTP server above [`vllm_chat`].
2 +
3 +mod config;
4 +mod error;
5 +mod grpc;
6 +mod listener;
7 +mod middleware;
8 +mod routes;
9 +mod state;
10 +mod utils;
11 +
12 +use std::sync::{Arc, OnceLock};
13 +
14 +use anyhow::{Context as _, Result};
15 +use axum::serve::ListenerExt as _;
16 +pub use config::{Config, CoordinatorMode, HttpListenerMode};
17 +use tokio::net::TcpListener;
18 +use tokio::time::{Instant, sleep_until};
19 +use tokio_stream::wrappers::TcpListenerStream;
20 +use tokio_util::either::Either;
21 +use tokio_util::sync::CancellationToken;
22 +use tonic::transport::Server as TonicServer;
23 +use tracing::{info, trace, warn};
24 +use vllm_chat::{ChatLlm, LoadModelBackendsOptions, load_model_backends};
25 +pub use vllm_chat::{ChatTemplateContentFormatOption, ParserSelection, RendererSelection};
26 +use vllm_engine_core_client::{EngineCoreClient, EngineCoreClientConfig};
27 +use vllm_llm::Llm;
28 +use vllm_text::TextLlm;
29 +
30 +use crate::listener::Listener;
31 +use crate::routes::build_router;
32 +use crate::state::AppState;
33 +
34 +/// Build the shared application state for one configured model and one engine
35 +/// client.
36 +async fn build_state(config: &Config) -> Result<Arc<AppState>> {
37 + // Load both backends from the same model metadata so they stay in sync.
38 + let loaded = load_model_backends(
39 + &config.model,
40 + LoadModelBackendsOptions {
41 + renderer: config.renderer,
42 + chat_template: config.chat_template.clone(),
43 + chat_template_content_format: config.chat_template_content_format,
44 + default_chat_template_kwargs: config
45 + .default_chat_template_kwargs
46 + .clone()
47 + .unwrap_or_default(),
48 + },
49 + )
50 + .await
51 + .context("failed to create chat/text backends")?;
52 + let text_backend = loaded.text_backend;
53 + let chat_backend = loaded.chat_backend;
54 +
55 + let coordinator_mode = config.effective_coordinator_mode(text_backend.is_moe());
56 + info!(
57 + engine_count = config.engine_count(),
58 + model_is_moe = text_backend.is_moe(),
59 + ?coordinator_mode,
60 + "resolved coordinator mode"
61 + );
62 +
63 + let client = EngineCoreClient::connect(EngineCoreClientConfig {
64 + transport_mode: config.transport_mode.clone(),
65 + coordinator_mode,
66 + model_name: config.model.clone(),
67 + client_index: 0,
68 + })
69 + .await
70 + .context("failed to connect to engine core")?;
71 +
72 + let llm = Llm::new(client).with_log_stats(!config.disable_log_stats);
73 + let text = TextLlm::new(llm, text_backend);
74 +
75 + let chat = ChatLlm::new(text, chat_backend)
76 + .with_tool_call_parser(config.tool_call_parser.clone())
77 + .with_reasoning_parser(config.reasoning_parser.clone());
78 +
79 + // If no served names are specified, fall back to the backend model path so
80 + // that the API always has at least one valid model ID.
81 + let served_model_names = if config.served_model_name.is_empty() {
82 + vec![config.model.clone()]
83 + } else {
84 + config.served_model_name.clone()
85 + };
86 +
87 + Ok(Arc::new(
88 + AppState::new(served_model_names, chat).with_log_requests(config.enable_log_requests),
89 + ))
90 +}
91 +
92 +/// Run the OpenAI-compatible HTTP server until the supplied shutdown token is
93 +/// cancelled.
94 +///
95 +/// The server owns one `vllm-chat` facade, which in turn owns the lower
96 +/// `vllm-text` and `vllm-llm` layers, and shuts them down before returning.
97 +pub async fn serve(config: Config, shutdown: CancellationToken) -> Result<()> {
98 + config.validate().context("invalid OpenAI frontend configuration")?;
99 +
100 + // Also check shutdown during the (potentially long) startup handshake.
101 + let state = tokio::select! {
102 + result = build_state(&config) => result?,
103 + _ = shutdown.cancelled() => return Ok(()),
104 + };
105 + let listener = Listener::bind(&config.listener_mode)
106 + .await
107 + .context("failed to bind listener for OpenAI server")?;
108 + let bind_address = listener.local_addr()?;
109 + let model = state.primary_model_name().to_owned();
110 + let app = build_router(state.clone());
111 +
112 + // Optionally bind the gRPC Generate server on a separate port. Bind
113 + // synchronously here so bind errors (port in use, permission denied, ...)
114 + // surface before we start serving, rather than being deferred until
115 + // shutdown. The gRPC listener follows the same host as the HTTP listener so
116 + // that enabling --grpc-port does not accidentally expose the service on all
117 + // interfaces when HTTP is intentionally local-only.
118 + let grpc_setup = if let Some(grpc_port) = config.grpc_port {
119 + let grpc_host = match &config.listener_mode {
120 + HttpListenerMode::BindTcp { host, .. } => host.as_str(),
121 + HttpListenerMode::BindUnix { .. } | HttpListenerMode::InheritedFd { .. } => "0.0.0.0",
122 + };
123 + let grpc_listener = TcpListener::bind((grpc_host, grpc_port))
124 + .await
125 + .with_context(|| format!("failed to bind gRPC listener on {grpc_host}:{grpc_port}"))?;
126 + let addr = grpc_listener.local_addr()?;
127 + let svc = grpc::GenerateServer::new(grpc::GenerateServiceImpl::new(state.clone()));
128 + info!(%addr, "starting gRPC server");
129 + Some((grpc_listener, svc))
130 + } else {
131 + None
132 + };
133 +
134 + info!(%bind_address, %model, "starting OpenAI server");
135 +
136 + // Set TCP_NODELAY on accepted connections to reduce latency.
137 + // By `tap_io` we will do this on every accepted connection.
138 + let listener = listener.tap_io(|io| {
139 + if let Either::Left(tcp_stream) = io
140 + && let Err(err) = tcp_stream.set_nodelay(true)
141 + {
142 + trace!(error = %err, "failed to enable TCP_NODELAY on accepted HTTP connection");
143 + }
144 + });
145 +
146 + // Run HTTP and gRPC concurrently under a child token of the caller's shutdown
147 + // token. Caller cancellation propagates into both protocols; if either
148 + // protocol exits first, we cancel this child token so its sibling also
149 + // begins a graceful drain.
150 + let server_shutdown = shutdown.child_token();
151 + let force_shutdown = CancellationToken::new();
152 + let shutdown_deadline = Arc::new(OnceLock::new());
153 +
154 + // Spawn a task to trigger `force_shutdown` after shutdown deadline elapses.
155 + tokio::spawn({
156 + let shutdown = server_shutdown.clone();
157 + let force_shutdown = force_shutdown.clone();
158 + let shutdown_deadline = shutdown_deadline.clone();
159 + let shutdown_timeout = config.shutdown_timeout;
160 +
161 + async move {
162 + shutdown.cancelled().await;
163 + let deadline = Instant::now() + shutdown_timeout;
164 + let _ = shutdown_deadline.set(deadline);
165 +
166 + if shutdown_timeout.is_zero() {
167 + force_shutdown.cancel();
168 + } else {
169 + sleep_until(deadline).await;
170 + force_shutdown.cancel();
171 + }
172 + }
173 + });
174 +
175 + let http_fut = {
176 + let shutdown = server_shutdown.child_token();
177 + let server_shutdown = server_shutdown.clone();
178 + let force_shutdown = force_shutdown.clone();
179 + async move {
180 + let server =
181 + axum::serve(listener, app).with_graceful_shutdown(shutdown.cancelled_owned());
182 +
183 + let result = tokio::select! {
184 + result = server => {
185 + result.context("HTTP server failed")
186 + }
187 + _ = force_shutdown.cancelled() => {
188 + warn!("HTTP graceful shutdown deadline elapsed; aborting server");
189 + Ok(())
190 + }
191 + };
192 +
193 + server_shutdown.cancel();
194 + result
195 + }
196 + };
197 +
198 + let grpc_fut = {
199 + let shutdown = server_shutdown.child_token();
200 + let server_shutdown = server_shutdown.clone();
201 + let force_shutdown = force_shutdown.clone();
202 + async move {
203 + let Some((grpc_listener, svc)) = grpc_setup else {
204 + // No gRPC configured: just wait for shutdown so we do not race the
205 + // join! by resolving early and tripping the cancellation token.
206 + shutdown.cancelled().await;
207 + return Ok(());
208 + };
209 + let server = TonicServer::builder().add_service(svc).serve_with_incoming_shutdown(
210 + TcpListenerStream::new(grpc_listener),
211 + shutdown.cancelled_owned(),
212 + );
213 +
214 + let result = tokio::select! {
215 + result = server => {
216 + result.context("gRPC server failed")
217 + }
218 + _ = force_shutdown.cancelled() => {
219 + warn!("gRPC graceful shutdown deadline elapsed; aborting server");
220 + Ok(())
221 + }
222 + };
223 +
224 + server_shutdown.cancel();
225 + result
226 + }
227 + };
228 +
229 + let (http_res, grpc_res) = tokio::join!(http_fut, grpc_fut);
230 + http_res.and(grpc_res)?;
231 +
232 + let shutdown_deadline = shutdown_deadline
233 + .get()
234 + .copied()
235 + .unwrap_or_else(|| Instant::now() + config.shutdown_timeout);
236 + state.shutdown(shutdown_deadline).await
237 +}
rust/src/server/src/listener.rs new
+135
@@ -0,0 +1,135 @@
1 +//! Unified HTTP listener wrapper for the Rust frontend.
2 +//!
3 +//! This module hides the difference between TCP and Unix-domain listeners so
4 +//! the rest of the server can bind or inherit one socket and pass it to
5 +//! `axum::serve(...)` through a single type.
6 +
7 +use std::io::Result;
8 +use std::net::TcpListener as StdTcpListener;
9 +use std::os::fd::{FromRawFd, IntoRawFd, OwnedFd};
10 +use std::os::unix::net::UnixListener as StdUnixListener;
11 +
12 +use socket2::Socket;
13 +use tokio::net::{TcpListener, TcpStream, UnixListener, UnixStream};
14 +use tokio_util::either::Either;
15 +
16 +use crate::HttpListenerMode;
17 +
18 +/// Runtime listener type used by the OpenAI-compatible HTTP server, which is
19 +/// either a TCP listener or a Unix-domain listener.
20 +#[derive(Debug)]
21 +pub enum Listener {
22 + Tcp(TcpListener),
23 + Unix(UnixListener),
24 +}
25 +
26 +impl Listener {
27 + /// Bind or adopt the listener described by the frontend configuration.
28 + ///
29 + /// For inherited sockets, the concrete listener kind is detected from the
30 + /// socket family of the supplied file descriptor.
31 + pub async fn bind(mode: &HttpListenerMode) -> Result<Self> {
32 + match mode {
33 + HttpListenerMode::BindTcp { host, port } => {
34 + Ok(Self::Tcp(TcpListener::bind((host.as_str(), *port)).await?))
35 + }
36 + HttpListenerMode::BindUnix { path } => Ok(Self::Unix(UnixListener::bind(path)?)),
37 + HttpListenerMode::InheritedFd { fd } => Self::from_inherited_fd(*fd),
38 + }
39 + }
40 +
41 + /// Return a log-friendly local address string for either TCP or Unix
42 + /// sockets.
43 + pub fn local_addr(&self) -> Result<String> {
44 + match self {
45 + Self::Tcp(listener) => Ok(listener.local_addr()?.to_string()),
46 + Self::Unix(listener) => Ok(match listener.local_addr()?.as_pathname() {
47 + Some(path) => format!("unix:{}", path.display()),
48 + None => "unix:<unnamed>".to_string(),
49 + }),
50 + }
51 + }
52 +
53 + fn from_inherited_fd(fd: i32) -> Result<Self> {
54 + // SAFETY: We trust the caller to only pass valid listener fds, and we only use
55 + // this fd once to create a single listener.
56 + let owned_fd = unsafe { OwnedFd::from_raw_fd(fd) };
57 + let socket = Socket::from(owned_fd);
58 +
59 + // The Python supervisor pre-binds the socket to reserve the endpoint early, but
60 + // Rust is responsible for transitioning inherited stream sockets into
61 + // the listening state before accepting connections.
62 + socket.listen(libc::SOMAXCONN)?;
63 + socket.set_nonblocking(true)?;
64 +
65 + if socket.local_addr()?.is_unix() {
66 + let std_listener = unsafe { StdUnixListener::from_raw_fd(socket.into_raw_fd()) };
67 + Ok(Self::Unix(UnixListener::from_std(std_listener)?))
68 + } else {
69 + let std_listener = unsafe { StdTcpListener::from_raw_fd(socket.into_raw_fd()) };
70 + Ok(Self::Tcp(TcpListener::from_std(std_listener)?))
71 + }
72 + }
73 +}
74 +
75 +/// Allow the unified listener to plug directly into `axum::serve(...)`.
76 +impl axum::serve::Listener for Listener {
77 + type Addr = Either<std::net::SocketAddr, tokio::net::unix::SocketAddr>;
78 + type Io = Either<TcpStream, UnixStream>;
79 +
80 + async fn accept(&mut self) -> (Self::Io, Self::Addr) {
81 + match self {
82 + Self::Tcp(listener) => {
83 + let (io, addr) = listener.accept().await;
84 + (Either::Left(io), Either::Left(addr))
85 + }
86 + Self::Unix(listener) => {
87 + let (io, addr) = listener.accept().await;
88 + (Either::Right(io), Either::Right(addr))
89 + }
90 + }
91 + }
92 +
93 + fn local_addr(&self) -> Result<Self::Addr> {
94 + match self {
95 + Self::Tcp(listener) => listener.local_addr().map(Either::Left),
96 + Self::Unix(listener) => listener.local_addr().map(Either::Right),
97 + }
98 + }
99 +}
100 +
101 +#[cfg(test)]
102 +mod tests {
103 + use std::net::{Ipv4Addr, SocketAddrV4};
104 + use std::os::fd::IntoRawFd;
105 +
106 + use socket2::{Domain, SockAddr, Socket, Type};
107 + use uuid::Uuid;
108 +
109 + use super::Listener;
110 + use crate::HttpListenerMode;
111 +
112 + #[tokio::test(flavor = "current_thread")]
113 + async fn inherited_fd_detects_tcp_listener_without_uds_hint() {
114 + let socket = Socket::new(Domain::IPV4, Type::STREAM, None).unwrap();
115 + socket.bind(&SockAddr::from(SocketAddrV4::new(Ipv4Addr::LOCALHOST, 0))).unwrap();
116 + let fd = socket.into_raw_fd();
117 +
118 + let listener = Listener::bind(&HttpListenerMode::InheritedFd { fd }).await.unwrap();
119 +
120 + assert!(matches!(listener, Listener::Tcp(_)));
121 + }
122 +
123 + #[tokio::test(flavor = "current_thread")]
124 + async fn inherited_fd_detects_unix_listener_from_fd() {
125 + let path = std::env::temp_dir().join(format!("vllm-rs-{}.sock", Uuid::new_v4()));
126 + let socket = Socket::new(Domain::UNIX, Type::STREAM, None).unwrap();
127 + socket.bind(&SockAddr::unix(&path).unwrap()).unwrap();
128 + let fd = socket.into_raw_fd();
129 +
130 + let listener = Listener::bind(&HttpListenerMode::InheritedFd { fd }).await.unwrap();
131 +
132 + assert!(matches!(listener, Listener::Unix(_)));
133 + let _ = std::fs::remove_file(path);
134 + }
135 +}
rust/src/server/src/middleware/load.rs new
+110
@@ -0,0 +1,110 @@
1 +use std::pin::Pin;
2 +use std::sync::{Arc, Weak};
3 +use std::task::{Context, Poll};
4 +
5 +use axum::body::{Body, Bytes, HttpBody};
6 +use axum::extract::{MatchedPath, Request, State};
7 +use axum::middleware::Next;
8 +use axum::response::Response;
9 +use http_body::{Frame, SizeHint};
10 +
11 +use crate::state::AppState;
12 +
13 +/// Endpoints that will be tracked for server load.
14 +///
15 +/// Derived from the Python frontend's actual `@load_aware_call` coverage. This
16 +/// includes alias paths that delegate into decorated handlers, such as
17 +/// `/v1/rerank` and `/v2/rerank`.
18 +const TRACKED_HANDLERS: &[&str] = &[
19 + "/v1/responses",
20 + "/v1/responses/{response_id}",
21 + "/v1/responses/{response_id}/cancel",
22 + "/v1/messages",
23 + "/v1/messages/count_tokens",
24 + "/v1/chat/completions",
25 + "/v1/completions",
26 + "/v1/audio/transcriptions",
27 + "/v1/audio/translations",
28 + "/v1/embeddings",
29 + "/pooling",
30 + "/classify",
31 + "/score",
32 + "/v1/score",
33 + "/rerank",
34 + "/v1/rerank",
35 + "/v2/rerank",
36 + "/inference/v1/generate",
37 +];
38 +
39 +/// Track frontend-local in-flight inference requests for the `/load` endpoint.
40 +pub async fn track_server_load(
41 + State(state): State<Arc<AppState>>,
42 + req: Request,
43 + next: Next,
44 +) -> Response {
45 + let handler = req
46 + .extensions()
47 + .get::<MatchedPath>()
48 + .map_or_else(|| "none", |path| path.as_str());
49 +
50 + if !TRACKED_HANDLERS.contains(&handler) {
51 + return next.run(req).await;
52 + }
53 +
54 + state.increment_server_load();
55 + let guard = ServerLoadGuard {
56 + state: Arc::downgrade(&state),
57 + };
58 + let response = next.run(req).await;
59 +
60 + let (parts, body) = response.into_parts();
61 + Response::from_parts(
62 + parts,
63 + Body::new(LoadTrackedBody {
64 + inner: body,
65 + _guard: guard,
66 + }),
67 + )
68 +}
69 +
70 +/// A guard that decrements the server load when dropped.
71 +struct ServerLoadGuard {
72 + state: Weak<AppState>,
73 +}
74 +
75 +impl Drop for ServerLoadGuard {
76 + fn drop(&mut self) {
77 + if let Some(state) = self.state.upgrade() {
78 + state.decrement_server_load();
79 + }
80 + }
81 +}
82 +
83 +/// A wrapper around response bodies that tracks server load by holding a
84 +/// `ServerLoadGuard`, which will decrement the load when the body is fully
85 +/// consumed and dropped.
86 +struct LoadTrackedBody {
87 + inner: Body,
88 + _guard: ServerLoadGuard,
89 +}
90 +
91 +// Simply delegate all `HttpBody` methods to the inner body.
92 +impl HttpBody for LoadTrackedBody {
93 + type Data = Bytes;
94 + type Error = axum::Error;
95 +
96 + fn poll_frame(
97 + mut self: Pin<&mut Self>,
98 + cx: &mut Context<'_>,
99 + ) -> Poll<Option<Result<Frame<Self::Data>, Self::Error>>> {
100 + Pin::new(&mut self.inner).poll_frame(cx)
101 + }
102 +
103 + fn is_end_stream(&self) -> bool {
104 + self.inner.is_end_stream()
105 + }
106 +
107 + fn size_hint(&self) -> SizeHint {
108 + self.inner.size_hint()
109 + }
110 +}
rust/src/server/src/middleware/metrics.rs new
+79
@@ -0,0 +1,79 @@
1 +use std::time::Instant;
2 +
3 +use axum::extract::{MatchedPath, Request};
4 +use axum::middleware::Next;
5 +use axum::response::Response;
6 +use vllm_metrics::{HttpHandlerLabels, HttpRequestLabels, METRICS};
7 +
8 +/// Endpoints that will be excluded from HTTP metrics tracking.
9 +///
10 +/// Original Python definition:
11 +/// <https://github.com/vllm-project/vllm/blob/bc2c0c86efb28e77677a3cfb8687e976914a313a/vllm/entrypoints/serve/instrumentator/metrics.py#L28-L38>
12 +const EXCLUDED_HANDLERS: &[&str] = &[
13 + "/metrics",
14 + "/health",
15 + "/load",
16 + "/ping",
17 + "/version",
18 + "/server_info",
19 + // Rust frontend extra:
20 + "/reset_prefix_cache",
21 + "/reset_mm_cache",
22 + "/reset_encoder_cache",
23 + "/collective_rpc",
24 + "/sleep",
25 + "/wake_up",
26 + "/is_sleeping",
27 +];
28 +
29 +/// Record API-server HTTP metrics with Python-compatible
30 +/// (`PrometheusFastApiInstrumentator` style) family names and labels.
31 +pub async fn track_http_metrics(req: Request, next: Next) -> Response {
32 + let method = req.method().as_str().to_string();
33 + let handler = req
34 + .extensions()
35 + .get::<MatchedPath>()
36 + .map_or_else(|| "none".to_string(), |path| path.as_str().to_string());
37 + let excluded = EXCLUDED_HANDLERS.contains(&handler.as_str());
38 + let started_at = Instant::now();
39 +
40 + let response = next.run(req).await;
41 +
42 + if excluded {
43 + return response;
44 + }
45 +
46 + let elapsed = started_at.elapsed().as_secs_f64();
47 + let status = status_group(response.status().as_u16());
48 +
49 + let metrics = &METRICS.api_server;
50 +
51 + metrics
52 + .http_requests
53 + .get_or_create(&HttpRequestLabels {
54 + method: method.clone(),
55 + status,
56 + handler: handler.clone(),
57 + })
58 + .inc();
59 +
60 + metrics
61 + .http_request_duration_seconds
62 + .get_or_create(&HttpHandlerLabels { method, handler })
63 + .observe(elapsed);
64 +
65 + metrics.http_request_duration_highr_seconds.observe(elapsed);
66 +
67 + response
68 +}
69 +
70 +fn status_group(status: u16) -> &'static str {
71 + match status / 100 {
72 + 1 => "1xx",
73 + 2 => "2xx",
74 + 3 => "3xx",
75 + 4 => "4xx",
76 + 5 => "5xx",
77 + _ => "unknown",
78 + }
79 +}
rust/src/server/src/middleware/mod.rs new
+5
@@ -0,0 +1,5 @@
1 +mod load;
2 +mod metrics;
3 +
4 +pub use load::track_server_load;
5 +pub use metrics::track_http_metrics;
rust/src/server/src/routes.rs new
+68
@@ -0,0 +1,68 @@
1 +mod cache;
2 +mod collective_rpc;
3 +mod health;
4 +mod inference;
5 +mod load;
6 +mod metrics;
7 +pub(crate) mod openai;
8 +mod sleep;
9 +
10 +use std::sync::Arc;
11 +
12 +use axum::Router;
13 +use axum::middleware::{from_fn, from_fn_with_state};
14 +use axum::routing::{get, post};
15 +use tower_http::trace::TraceLayer;
16 +
17 +use crate::middleware;
18 +use crate::state::AppState;
19 +
20 +fn server_dev_mode_enabled() -> bool {
21 + std::env::var("VLLM_SERVER_DEV_MODE")
22 + .ok()
23 + .and_then(|value| value.parse::<i64>().ok())
24 + .is_some_and(|value| value != 0)
25 +}
26 +
27 +/// Build the minimal OpenAI-compatible router for one configured model.
28 +pub fn build_router(state: Arc<AppState>) -> Router {
29 + build_router_with_dev_mode(state, server_dev_mode_enabled())
30 +}
31 +
32 +fn build_router_with_dev_mode(state: Arc<AppState>, dev_mode_enabled: bool) -> Router {
33 + let mut router = Router::new()
34 + // Health & monitoring
35 + .route("/health", get(health::health))
36 + .route("/metrics", get(metrics::scrape))
37 + .route("/load", get(load::load))
38 + // OpenAI-compatible endpoints
39 + .route("/v1/models", get(openai::list_models))
40 + .route("/v1/completions", post(openai::completions))
41 + .route("/v1/chat/completions", post(openai::chat_completions))
42 + // vLLM specific inference endpoints
43 + .route("/inference/v1/generate", post(inference::generate));
44 +
45 + if dev_mode_enabled {
46 + // Development-only
47 + router = router
48 + .route("/reset_prefix_cache", post(cache::reset_prefix_cache))
49 + .route("/reset_mm_cache", post(cache::reset_mm_cache))
50 + .route("/reset_encoder_cache", post(cache::reset_encoder_cache))
51 + .route("/collective_rpc", post(collective_rpc::collective_rpc))
52 + .route("/sleep", post(sleep::sleep))
53 + .route("/wake_up", post(sleep::wake_up))
54 + .route("/is_sleeping", get(sleep::is_sleeping))
55 + }
56 +
57 + router
58 + .with_state(state.clone())
59 + .layer(from_fn_with_state(state, middleware::track_server_load))
60 + .layer(from_fn(middleware::track_http_metrics))
61 + .layer(TraceLayer::new_for_http())
62 +}
63 +
64 +#[cfg(test)]
65 +mod tests;
66 +
67 +#[cfg(test)]
68 +mod http_client_tests;
rust/src/server/src/routes/cache.rs new
+56
@@ -0,0 +1,56 @@
1 +use std::sync::Arc;
2 +
3 +use axum::extract::{Query, State};
4 +use axum::http::StatusCode;
5 +use serde::Deserialize;
6 +
7 +use crate::error::ApiError;
8 +use crate::state::AppState;
9 +use crate::utils::utility_call_error;
10 +
11 +#[derive(Debug, Default, Deserialize)]
12 +pub(crate) struct ResetPrefixCacheParams {
13 + #[serde(default)]
14 + reset_running_requests: bool,
15 + #[serde(default)]
16 + reset_external: bool,
17 +}
18 +
19 +/// Reset the local prefix cache and optionally the connector-managed external
20 +/// cache.
21 +pub async fn reset_prefix_cache(
22 + State(state): State<Arc<AppState>>,
23 + Query(params): Query<ResetPrefixCacheParams>,
24 +) -> Result<StatusCode, ApiError> {
25 + state
26 + .engine_core_client()
27 + .reset_prefix_cache(params.reset_running_requests, params.reset_external)
28 + .await
29 + .map_err(|error| utility_call_error("reset_prefix_cache", error))?;
30 +
31 + Ok(StatusCode::OK)
32 +}
33 +
34 +/// Reset the multi-modal cache.
35 +pub async fn reset_mm_cache(State(state): State<Arc<AppState>>) -> Result<StatusCode, ApiError> {
36 + state
37 + .engine_core_client()
38 + .reset_mm_cache()
39 + .await
40 + .map_err(|error| utility_call_error("reset_mm_cache", error))?;
41 +
42 + Ok(StatusCode::OK)
43 +}
44 +
45 +/// Reset the encoder cache.
46 +pub async fn reset_encoder_cache(
47 + State(state): State<Arc<AppState>>,
48 +) -> Result<StatusCode, ApiError> {
49 + state
50 + .engine_core_client()
51 + .reset_encoder_cache()
52 + .await
53 + .map_err(|error| utility_call_error("reset_encoder_cache", error))?;
54 +
55 + Ok(StatusCode::OK)
56 +}
rust/src/server/src/routes/collective_rpc.rs new
+51
@@ -0,0 +1,51 @@
1 +use std::collections::BTreeMap;
2 +use std::sync::Arc;
3 +
4 +use axum::Json;
5 +use axum::extract::State;
6 +use axum::extract::rejection::JsonRejection;
7 +use rmpv::Value as MsgpackValue;
8 +use serde::{Deserialize, Serialize};
9 +use serde_json::Value as JsonValue;
10 +
11 +use crate::error::ApiError;
12 +use crate::state::AppState;
13 +use crate::utils::utility_call_error;
14 +
15 +#[derive(Debug, Deserialize)]
16 +pub(crate) struct CollectiveRpcRequest {
17 + method: Option<String>,
18 + #[serde(default)]
19 + timeout: Option<f64>,
20 + #[serde(default)]
21 + args: Vec<JsonValue>,
22 + #[serde(default)]
23 + kwargs: BTreeMap<String, JsonValue>,
24 +}
25 +
26 +#[derive(Debug, Serialize)]
27 +pub(crate) struct CollectiveRpcResponse {
28 + results: Vec<MsgpackValue>,
29 +}
30 +
31 +/// Execute a development-only collective RPC on the connected engine(s).
32 +pub async fn collective_rpc(
33 + State(state): State<Arc<AppState>>,
34 + body: Result<Json<CollectiveRpcRequest>, JsonRejection>,
35 +) -> Result<Json<CollectiveRpcResponse>, ApiError> {
36 + let Json(body) = body.map_err(|error| ApiError::json_parse_error(error.body_text()))?;
37 + let method = body.method.ok_or_else(|| {
38 + ApiError::invalid_request(
39 + "Missing 'method' in request body".to_string(),
40 + Some("method"),
41 + )
42 + })?;
43 +
44 + let results = state
45 + .engine_core_client()
46 + .collective_rpc(&method, body.timeout, body.args, body.kwargs)
47 + .await
48 + .map_err(|error| utility_call_error("collective_rpc", error))?;
49 +
50 + Ok(Json(CollectiveRpcResponse { results }))
51 +}
rust/src/server/src/routes/health.rs new
+14
@@ -0,0 +1,14 @@
1 +use std::sync::Arc;
2 +
3 +use axum::extract::State;
4 +use axum::http::StatusCode;
5 +
6 +use crate::state::AppState;
7 +
8 +pub async fn health(State(state): State<Arc<AppState>>) -> StatusCode {
9 + if state.chat.engine_core_client().is_healthy() {
10 + StatusCode::OK
11 + } else {
12 + StatusCode::SERVICE_UNAVAILABLE
13 + }
14 +}
rust/src/server/src/routes/http_client_tests.rs new
+392
@@ -0,0 +1,392 @@
1 +//! Integration tests that exercise the OpenAI-compatible HTTP API through a
2 +//! real TCP connection using the `async-openai` client library, backed by a
3 +//! mock engine.
4 +
5 +use std::future::Future;
6 +use std::pin::Pin;
7 +use std::sync::Arc;
8 +use std::task::{Context, Poll};
9 +
10 +use async_openai::Client;
11 +use async_openai::config::OpenAIConfig;
12 +use async_openai::types::chat::{
13 + ChatCompletionRequestUserMessageArgs, CreateChatCompletionRequestArgs,
14 +};
15 +use futures::StreamExt as _;
16 +use serial_test::serial;
17 +use vllm_chat::{
18 + ChatBackend, ChatLlm, ChatRenderer, ChatRequest, ChatTextBackend, DefaultChatOutputProcessor,
19 + DynChatOutputProcessor, DynChatRenderer, NewChatOutputProcessorOptions, RenderedPrompt,
20 +};
21 +use vllm_engine_core_client::protocol::{
22 + EngineCoreFinishReason, EngineCoreOutput, EngineCoreOutputs, EngineCoreRequest,
23 +};
24 +use vllm_engine_core_client::test_utils::{IpcNamespace, spawn_mock_engine_task};
25 +use vllm_engine_core_client::{EngineCoreClient, EngineCoreClientConfig, EngineId};
26 +use vllm_llm::Llm;
27 +use vllm_text::tokenizer::{DynTokenizer, Tokenizer};
28 +use vllm_text::{Prompt, TextBackend};
29 +use zeromq::prelude::{SocketRecv, SocketSend};
30 +use zeromq::{DealerSocket, PushSocket, ZmqMessage};
31 +
32 +use crate::routes::build_router;
33 +use crate::state::AppState;
34 +
35 +// ========================================================================================
36 +// Test infrastructure (mirrors routes/tests.rs helpers)
37 +// ========================================================================================
38 +
39 +type TestFuture<'a> = Pin<Box<dyn Future<Output = ()> + Send + 'a>>;
40 +
41 +fn boxed_test_future<'a>(future: impl Future<Output = ()> + Send + 'a) -> TestFuture<'a> {
42 + Box::pin(future)
43 +}
44 +
45 +struct MockEngineTask {
46 + shutdown_tx: Option<tokio::sync::oneshot::Sender<()>>,
47 + join_handle: Option<tokio::task::JoinHandle<()>>,
48 +}
49 +
50 +impl MockEngineTask {
51 + fn new(
52 + (shutdown_tx, join_handle): (
53 + tokio::sync::oneshot::Sender<()>,
54 + tokio::task::JoinHandle<()>,
55 + ),
56 + ) -> Self {
57 + Self {
58 + shutdown_tx: Some(shutdown_tx),
59 + join_handle: Some(join_handle),
60 + }
61 + }
62 +}
63 +
64 +impl Future for MockEngineTask {
65 + type Output = Result<(), tokio::task::JoinError>;
66 +
67 + fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
68 + if let Some(shutdown_tx) = self.shutdown_tx.take() {
69 + let _ = shutdown_tx.send(());
70 + }
71 + match self.join_handle.as_mut() {
72 + Some(join_handle) => Pin::new(join_handle).poll(cx),
73 + None => Poll::Ready(Ok(())),
74 + }
75 + }
76 +}
77 +
78 +impl Drop for MockEngineTask {
79 + fn drop(&mut self) {
80 + if let Some(join_handle) = &self.join_handle {
81 + join_handle.abort();
82 + }
83 + }
84 +}
85 +
86 +fn request_output(
87 + request_id: &str,
88 + new_token_ids: Vec<u32>,
89 + finish_reason: Option<EngineCoreFinishReason>,
90 +) -> EngineCoreOutput {
91 + EngineCoreOutput {
92 + request_id: request_id.to_string(),
93 + new_token_ids,
94 + new_logprobs: None,
95 + new_prompt_logprobs_tensors: None,
96 + pooling_output: None,
97 + finish_reason,
98 + stop_reason: None,
99 + events: None,
100 + kv_transfer_params: None,
101 + trace_headers: None,
102 + prefill_stats: None,
103 + routed_experts: None,
104 + num_nans_in_logits: 0,
105 + }
106 +}
107 +
108 +fn engine_outputs_for_request(
109 + request_id: &str,
110 + output_specs: Vec<(Vec<u32>, Option<EngineCoreFinishReason>)>,
111 +) -> EngineCoreOutputs {
112 + EngineCoreOutputs {
113 + engine_index: 0,
114 + outputs: output_specs
115 + .into_iter()
116 + .map(|(token_ids, finish_reason)| request_output(request_id, token_ids, finish_reason))
117 + .collect(),
118 + scheduler_stats: None,
119 + timestamp: 0.0,
120 + utility_output: None,
121 + finished_requests: None,
122 + wave_complete: None,
123 + start_wave: None,
124 + }
125 +}
126 +
127 +fn default_stream_output_specs() -> Vec<(Vec<u32>, Option<EngineCoreFinishReason>)> {
128 + vec![
129 + (vec![b'h' as u32], None),
130 + (vec![b'i' as u32], None),
131 + (vec![b'!' as u32], Some(EngineCoreFinishReason::Stop)),
132 + ]
133 +}
134 +
135 +async fn send_outputs(push: &mut PushSocket, outputs: EngineCoreOutputs) {
136 + push.send(ZmqMessage::from(
137 + rmp_serde::to_vec_named(&outputs).expect("encode outputs"),
138 + ))
139 + .await
140 + .expect("send outputs");
141 +}
142 +
143 +async fn recv_engine_message(dealer: &mut DealerSocket) -> Vec<bytes::Bytes> {
144 + dealer.recv().await.expect("recv engine message").into_vec()
145 +}
146 +
147 +fn test_llm(client: EngineCoreClient) -> Llm {
148 + Llm::new(client).with_request_id_randomization(false)
149 +}
150 +
151 +#[derive(Clone, Debug)]
152 +struct FakeChatBackend;
153 +
154 +#[derive(Debug)]
155 +struct FakeChatTokenizer;
156 +
157 +impl Tokenizer for FakeChatTokenizer {
158 + fn encode(
159 + &self,
160 + text: &str,
161 + _add_special_tokens: bool,
162 + ) -> vllm_text::tokenizer::Result<Vec<u32>> {
163 + Ok(text.bytes().map(u32::from).collect())
164 + }
165 +
166 + fn decode(
167 + &self,
168 + token_ids: &[u32],
169 + _skip_special_tokens: bool,
170 + ) -> vllm_text::tokenizer::Result<String> {
171 + Ok(
172 + String::from_utf8_lossy(&token_ids.iter().map(|id| *id as u8).collect::<Vec<_>>())
173 + .into_owned(),
174 + )
175 + }
176 +
177 + fn token_to_id(&self, token: &str) -> Option<u32> {
178 + token.bytes().next().map(u32::from)
179 + }
180 +}
181 +
182 +impl TextBackend for FakeChatBackend {
183 + fn tokenizer(&self) -> DynTokenizer {
184 + Arc::new(FakeChatTokenizer)
185 + }
186 +
187 + fn model_id(&self) -> &str {
188 + "test-model"
189 + }
190 +}
191 +
192 +impl ChatBackend for FakeChatBackend {
193 + fn chat_renderer(&self) -> DynChatRenderer {
194 + Arc::new(self.clone())
195 + }
196 +
197 + fn new_chat_output_processor(
198 + &self,
199 + request: &mut ChatRequest,
200 + options: NewChatOutputProcessorOptions<'_>,
201 + ) -> vllm_chat::Result<DynChatOutputProcessor> {
202 + Ok(Box::new(DefaultChatOutputProcessor::new(
203 + request,
204 + self.model_id(),
205 + self.tokenizer(),
206 + options.tool_call_parser,
207 + options.reasoning_parser,
208 + )?))
209 + }
210 +}
211 +
212 +impl ChatRenderer for FakeChatBackend {
213 + fn render(&self, request: &ChatRequest) -> vllm_chat::Result<RenderedPrompt> {
214 + let mut prompt = String::new();
215 + for message in &request.messages {
216 + prompt.push_str(message.role().as_str());
217 + prompt.push_str(": ");
218 + prompt.push_str(&message.text_content()?);
219 + prompt.push('\n');
220 + }
221 + if request.chat_options.add_generation_prompt() {
222 + prompt.push_str("assistant:");
223 + }
224 + Ok(RenderedPrompt {
225 + prompt: Prompt::Text(prompt),
226 + })
227 + }
228 +}
229 +
230 +/// Spin up an HTTP server on a random port backed by a mock engine.
231 +/// Returns the `async-openai` client, the HTTP server task, and the mock engine
232 +/// task.
233 +async fn http_test_server(
234 + engine_id: impl Into<EngineId>,
235 + output_specs: Vec<(Vec<u32>, Option<EngineCoreFinishReason>)>,
236 +) -> (
237 + Client<OpenAIConfig>,
238 + tokio::task::JoinHandle<()>,
239 + MockEngineTask,
240 +) {
241 + let ipc = IpcNamespace::new().expect("create ipc namespace");
242 + let handshake_address = ipc.handshake_endpoint();
243 + let engine_id = engine_id.into();
244 +
245 + let engine_task = MockEngineTask::new(spawn_mock_engine_task(
246 + handshake_address.clone(),
247 + engine_id.clone(),
248 + move |dealer, push| {
249 + boxed_test_future(async move {
250 + let add = recv_engine_message(dealer).await;
251 + let request: EngineCoreRequest =
252 + rmp_serde::from_slice(&add[1]).expect("decode request");
253 + send_outputs(
254 + push,
255 + engine_outputs_for_request(&request.request_id, output_specs),
256 + )
257 + .await;
258 + })
259 + },
260 + ));
261 +
262 + let client = EngineCoreClient::connect(
263 + EngineCoreClientConfig::new_single(handshake_address)
264 + .with_model_name("test-model")
265 + .with_local_input_output_addresses(
266 + Some(ipc.input_endpoint()),
267 + Some(ipc.output_endpoint()),
268 + ),
269 + )
270 + .await
271 + .expect("connect client");
272 +
273 + let chat = ChatLlm::from_shared_backend(
274 + test_llm(client),
275 + Arc::new(FakeChatBackend) as Arc<dyn ChatTextBackend>,
276 + );
277 + let state = Arc::new(AppState::new(vec!["test-model".to_string()], chat));
278 + let app = build_router(state);
279 +
280 + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.expect("bind http listener");
281 + let addr = listener.local_addr().expect("local addr");
282 +
283 + let server_task = tokio::spawn(async move {
284 + axum::serve(listener, app).await.expect("http server");
285 + });
286 +
287 + let openai_client = Client::with_config(
288 + OpenAIConfig::new()
289 + .with_api_key("unused")
290 + .with_api_base(format!("http://{addr}/v1")),
291 + );
292 +
293 + (openai_client, server_task, engine_task)
294 +}
295 +
296 +// ========================================================================================
297 +// Tests
298 +// ========================================================================================
299 +
300 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
301 +#[serial]
302 +async fn list_models_via_http_client() {
303 + let (client, server_task, _engine_task) =
304 + http_test_server(b"engine-http-models", default_stream_output_specs()).await;
305 +
306 + let models = client.models().list().await.expect("list models");
307 + let model_ids: Vec<&str> = models.data.iter().map(|m| m.id.as_str()).collect();
308 + assert_eq!(model_ids, vec!["test-model"]);
309 +
310 + server_task.abort();
311 +}
312 +
313 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
314 +#[serial]
315 +async fn non_streaming_chat_via_http_client() {
316 + let (client, server_task, engine_task) =
317 + http_test_server(b"engine-http-chat", default_stream_output_specs()).await;
318 +
319 + let request = CreateChatCompletionRequestArgs::default()
320 + .model("test-model")
321 + .stream(false)
322 + .max_completion_tokens(10u32)
323 + .messages([ChatCompletionRequestUserMessageArgs::default()
324 + .content("hello")
325 + .build()
326 + .expect("build user message")
327 + .into()])
328 + .build()
329 + .expect("build request");
330 +
331 + let response = client.chat().create(request).await.expect("chat completion");
332 +
333 + assert_eq!(response.model, "test-model");
334 + assert_eq!(response.choices.len(), 1);
335 + let choice = &response.choices[0];
336 + // The stop token `!` is suppressed from text.
337 + assert_eq!(choice.message.content.as_deref(), Some("hi"));
338 + assert_eq!(
339 + choice.finish_reason,
340 + Some(async_openai::types::chat::FinishReason::Stop)
341 + );
342 +
343 + engine_task.await.expect("mock engine task");
344 + server_task.abort();
345 +}
346 +
347 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
348 +#[serial]
349 +async fn streaming_chat_via_http_client() {
350 + let (client, server_task, engine_task) =
351 + http_test_server(b"engine-http-stream", default_stream_output_specs()).await;
352 +
353 + let request = CreateChatCompletionRequestArgs::default()
354 + .model("test-model")
355 + .stream(true)
356 + .max_completion_tokens(10u32)
357 + .messages([ChatCompletionRequestUserMessageArgs::default()
358 + .content("hello")
359 + .build()
360 + .expect("build user message")
361 + .into()])
362 + .build()
363 + .expect("build request");
364 +
365 + let mut stream = client.chat().create_stream(request).await.expect("streaming chat completion");
366 +
367 + let mut full_text = String::new();
368 + let mut saw_role = false;
369 + let mut saw_finish_reason = false;
370 +
371 + while let Some(chunk) = stream.next().await {
372 + let chunk = chunk.expect("stream chunk");
373 + for choice in &chunk.choices {
374 + if choice.delta.role.is_some() {
375 + saw_role = true;
376 + }
377 + if let Some(ref delta) = choice.delta.content {
378 + full_text.push_str(delta);
379 + }
380 + if choice.finish_reason.is_some() {
381 + saw_finish_reason = true;
382 + }
383 + }
384 + }
385 +
386 + assert!(saw_role, "expected an assistant role chunk");
387 + assert!(saw_finish_reason, "expected a terminal finish reason");
388 + assert_eq!(full_text, "hi");
389 +
390 + engine_task.await.expect("mock engine task");
391 + server_task.abort();
392 +}
rust/src/server/src/routes/inference/generate.rs new
+215
@@ -0,0 +1,215 @@
1 +mod convert;
2 +mod types;
3 +mod validate;
4 +
5 +use std::collections::HashMap;
6 +use std::sync::Arc;
7 +
8 +use axum::Json;
9 +use axum::extract::State;
10 +use axum::http::HeaderMap;
11 +use axum::response::{IntoResponse, Response};
12 +use thiserror_ext::AsReport as _;
13 +use tracing::info;
14 +use tracing_futures::Instrument as _;
15 +use vllm_engine_core_client::protocol::logprobs::{Logprobs, PositionLogprobs};
16 +use vllm_llm::{CollectedGenerateOutput, GenerateOutputStreamExt as _};
17 +
18 +use self::convert::prepare_generate_request;
19 +use self::types::{GenerateLogprob, GenerateRequest, GenerateResponse, GenerateResponseChoice};
20 +use crate::error::{ApiError, server_error};
21 +use crate::routes::openai::utils::logprobs::clamp_logprob;
22 +use crate::routes::openai::utils::types::{ChatLogProbs, ChatLogProbsContent, TopLogProb};
23 +use crate::routes::openai::utils::validated_json::ValidatedJson;
24 +use crate::state::AppState;
25 +use crate::utils::resolve_request_context;
26 +
27 +/// Validate one token-in/token-out request and proxy it into the shared
28 +/// `vllm-text` stack.
29 +pub async fn generate(
30 + State(state): State<Arc<AppState>>,
31 + headers: HeaderMap,
32 + ValidatedJson(body): ValidatedJson<GenerateRequest>,
33 +) -> Response {
34 + let request_context = resolve_request_context(&headers, body.request_id.as_deref());
35 + let prepared = match prepare_generate_request(body, state.served_model_names(), request_context)
36 + {
37 + Ok(prepared) => prepared,
38 + Err(error) => return error.into_response(),
39 + };
40 + let request_span = tracing::info_span!(
41 + "generate",
42 + request_id = %prepared.request_id,
43 + engine_request_id = tracing::field::Empty,
44 + );
45 +
46 + let log_request = state.enable_log_requests;
47 + let include_logprobs = prepared.include_logprobs;
48 + let include_prompt_logprobs = prepared.include_prompt_logprobs;
49 +
50 + let raw_stream = match state
51 + .chat
52 + .text()
53 + .generate_raw(prepared.text_request)
54 + .instrument(request_span.clone())
55 + .await
56 + {
57 + Ok(stream) => stream,
58 + Err(error) => {
59 + return server_error!(
60 + "failed to submit raw generate request: {}",
61 + error.to_report_string()
62 + )
63 + .into_response();
64 + }
65 + };
66 +
67 + let collected = match raw_stream.collect_output().instrument(request_span.clone()).await {
68 + Ok(collected) => collected,
69 + Err(error) => {
70 + return server_error!(
71 + "failed to collect raw generate response: {}",
72 + error.to_report_string()
73 + )
74 + .into_response();
75 + }
76 + };
77 +
78 + if log_request {
79 + info!(
80 + parent: &request_span,
81 + prompt_tokens = collected.prompt_token_ids.len(),
82 + output_tokens = collected.token_ids.len(),
83 + finish_reason = collected.finish_reason.as_str(),
84 + "generate finished"
85 + );
86 + }
87 +
88 + let response = match collect_generate(
89 + collected,
90 + prepared.request_id,
91 + include_logprobs,
92 + include_prompt_logprobs,
93 + ) {
94 + Ok(response) => response,
95 + Err(error) => return error.into_response(),
96 + };
97 +
98 + Json(response).into_response()
99 +}
100 +
101 +fn collect_generate(
102 + collected: CollectedGenerateOutput,
103 + request_id: String,
104 + include_logprobs: bool,
105 + include_prompt_logprobs: bool,
106 +) -> Result<GenerateResponse, ApiError> {
107 + let logprobs = if include_logprobs {
108 + let logprobs = collected.logprobs.as_ref().ok_or_else(|| {
109 + ApiError::server_error(
110 + "raw generate response requested logprobs but generation returned none".to_string(),
111 + )
112 + })?;
113 + Some(raw_logprobs_to_openai_chat(logprobs)?)
114 + } else {
115 + None
116 + };
117 + let prompt_logprobs = if include_prompt_logprobs {
118 + let prompt_logprobs = collected.prompt_logprobs.as_ref().ok_or_else(|| {
119 + ApiError::server_error(
120 + "raw generate response requested prompt_logprobs but generation returned none"
121 + .to_string(),
122 + )
123 + })?;
124 + Some(raw_prompt_logprobs_to_maps(prompt_logprobs))
125 + } else {
126 + None
127 + };
128 +
129 + Ok(GenerateResponse {
130 + request_id,
131 + choices: vec![GenerateResponseChoice {
132 + index: 0,
133 + logprobs,
134 + finish_reason: Some(collected.finish_reason.as_str().to_string()),
135 + token_ids: collected.token_ids,
136 + }],
137 + prompt_logprobs,
138 + kv_transfer_params: collected.kv_transfer_params,
139 + })
140 +}
141 +
142 +fn raw_logprobs_to_openai_chat(logprobs: &Logprobs) -> Result<ChatLogProbs, ApiError> {
143 + let content = logprobs
144 + .positions
145 + .iter()
146 + .map(position_to_chat_logprobs_content)
147 + .collect::<Result<Vec<_>, _>>()?;
148 +
149 + Ok(ChatLogProbs {
150 + content: Some(content),
151 + })
152 +}
153 +
154 +fn raw_prompt_logprobs_to_maps(
155 + prompt_logprobs: &Logprobs,
156 +) -> Vec<Option<HashMap<u32, GenerateLogprob>>> {
157 + std::iter::once(None)
158 + .chain(
159 + prompt_logprobs
160 + .positions
161 + .iter()
162 + .map(|position| Some(position_to_logprob_map(position))),
163 + )
164 + .collect()
165 +}
166 +
167 +fn position_to_chat_logprobs_content(
168 + position: &PositionLogprobs,
169 +) -> Result<ChatLogProbsContent, ApiError> {
170 + let chosen = position.entries.first().ok_or_else(|| {
171 + ApiError::server_error(
172 + "raw generate logprobs position unexpectedly had no token candidates".to_string(),
173 + )
174 + })?;
175 + let token = format_token_id(chosen.token_id);
176 +
177 + Ok(ChatLogProbsContent {
178 + token: token.clone(),
179 + logprob: clamp_logprob(chosen.logprob),
180 + bytes: Some(token.as_bytes().to_vec()),
181 + top_logprobs: position
182 + .entries
183 + .iter()
184 + .map(|entry| {
185 + let token = format_token_id(entry.token_id);
186 + TopLogProb {
187 + token: token.clone(),
188 + logprob: clamp_logprob(entry.logprob),
189 + bytes: Some(token.into_bytes()),
190 + }
191 + })
192 + .collect(),
193 + })
194 +}
195 +
196 +fn position_to_logprob_map(position: &PositionLogprobs) -> HashMap<u32, GenerateLogprob> {
197 + position
198 + .entries
199 + .iter()
200 + .map(|entry| {
201 + (
202 + entry.token_id,
203 + GenerateLogprob {
204 + logprob: clamp_logprob(entry.logprob),
205 + rank: Some(entry.rank),
206 + decoded_token: Some(format_token_id(entry.token_id)),
207 + },
208 + )
209 + })
210 + .collect()
211 +}
212 +
213 +fn format_token_id(token_id: u32) -> String {
214 + format!("token_id:{token_id}")
215 +}
rust/src/server/src/routes/inference/generate/convert.rs new
+112
@@ -0,0 +1,112 @@
1 +use vllm_text::{Prompt, TextDecodeOptions, TextRequest};
2 +
3 +use super::types::GenerateRequest;
4 +use super::validate;
5 +use crate::error::ApiError;
6 +use crate::utils::{ResolvedRequestContext, merge_kv_transfer_params};
7 +
8 +/// Lowered generate request plus the response request ID.
9 +#[derive(Debug, Clone, PartialEq)]
10 +pub struct PreparedRequest {
11 + pub request_id: String,
12 + pub text_request: TextRequest,
13 + pub include_logprobs: bool,
14 + pub include_prompt_logprobs: bool,
15 +}
16 +
17 +/// Validate and lower one raw generate request into the internal
18 +/// text-generation format.
19 +pub fn prepare_generate_request(
20 + request: GenerateRequest,
21 + served_model_names: &[String],
22 + ctx: ResolvedRequestContext,
23 +) -> Result<PreparedRequest, ApiError> {
24 + validate::validate_request_compat(&request, served_model_names)?;
25 +
26 + let include_logprobs = request.sampling_params.logprobs.is_some();
27 + let include_prompt_logprobs = request.sampling_params.prompt_logprobs.is_some();
28 + let mut sampling_params = request.sampling_params;
29 + sampling_params.vllm_xargs = merge_kv_transfer_params(
30 + sampling_params.vllm_xargs,
31 + request.kv_transfer_params.as_ref(),
32 + );
33 +
34 + let text_request = TextRequest {
35 + request_id: ctx.request_id.clone(),
36 + prompt: Prompt::TokenIds(request.token_ids),
37 + mm_features: None,
38 + sampling_params,
39 + decode_options: TextDecodeOptions::default(),
40 + intermediate: false,
41 + priority: request.priority,
42 + cache_salt: request.cache_salt,
43 + add_special_tokens: false,
44 + data_parallel_rank: ctx.data_parallel_rank,
45 + };
46 +
47 + Ok(PreparedRequest {
48 + request_id: ctx.request_id,
49 + text_request,
50 + include_logprobs,
51 + include_prompt_logprobs,
52 + })
53 +}
54 +
55 +#[cfg(test)]
56 +mod tests {
57 + use serde_json::json;
58 + use vllm_text::Prompt;
59 +
60 + use super::prepare_generate_request;
61 + use crate::routes::inference::generate::types::GenerateRequest;
62 + use crate::utils::ResolvedRequestContext;
63 +
64 + #[test]
65 + fn prepare_generate_request_maps_token_prompt_and_sampling_params() {
66 + let request: GenerateRequest = serde_json::from_value(json!({
67 + "model": "Qwen/Qwen1.5-0.5B-Chat",
68 + "token_ids": [11, 22, 33],
69 + "priority": -3,
70 + "cache_salt": "salt",
71 + "sampling_params": {
72 + "max_tokens": 7,
73 + "logprobs": 2,
74 + "prompt_logprobs": 1,
75 + "ignore_eos": true
76 + },
77 + "kv_transfer_params": {
78 + "connector": "x"
79 + }
80 + }))
81 + .expect("parse request");
82 +
83 + let prepared = prepare_generate_request(
84 + request,
85 + &["Qwen/Qwen1.5-0.5B-Chat".to_string()],
86 + ResolvedRequestContext::default(),
87 + )
88 + .expect("prepare");
89 +
90 + assert_eq!(
91 + prepared.text_request.prompt,
92 + Prompt::TokenIds(vec![11, 22, 33])
93 + );
94 + assert_eq!(prepared.text_request.sampling_params.max_tokens, Some(7));
95 + assert_eq!(prepared.text_request.sampling_params.logprobs, Some(2));
96 + assert_eq!(
97 + prepared.text_request.sampling_params.prompt_logprobs,
98 + Some(1)
99 + );
100 + assert!(prepared.text_request.sampling_params.ignore_eos);
101 + assert_eq!(prepared.text_request.priority, -3);
102 + assert_eq!(prepared.text_request.cache_salt.as_deref(), Some("salt"));
103 + assert_eq!(
104 + prepared
105 + .text_request
106 + .sampling_params
107 + .vllm_xargs
108 + .and_then(|mut xargs| xargs.remove("kv_transfer_params")),
109 + Some(json!({"connector": "x"}))
110 + );
111 + }
112 +}
rust/src/server/src/routes/inference/generate/types.rs new
+57
@@ -0,0 +1,57 @@
1 +use std::collections::HashMap;
2 +
3 +use serde::{Deserialize, Serialize};
4 +use serde_json::{Map, Value};
5 +use validator::Validate;
6 +use vllm_text::SamplingParams;
7 +
8 +use crate::routes::openai::utils::types::{ChatLogProbs, Normalizable};
9 +
10 +/// vLLM-compatible request type for the token-in/token-out generate API.
11 +#[serde_with::skip_serializing_none]
12 +#[derive(Debug, Clone, Deserialize, Serialize, Validate)]
13 +pub struct GenerateRequest {
14 + pub request_id: Option<String>,
15 + pub model: Option<String>,
16 + pub token_ids: Vec<u32>,
17 + pub sampling_params: SamplingParams,
18 + #[serde(default)]
19 + pub stream: bool,
20 + pub cache_salt: Option<String>,
21 + #[serde(default)]
22 + pub priority: i32,
23 + pub kv_transfer_params: Option<HashMap<String, Value>>,
24 + #[serde(flatten)]
25 + pub other: Map<String, Value>,
26 +}
27 +
28 +impl Normalizable for GenerateRequest {}
29 +
30 +/// Mirrors the Python vLLM `GenerateResponseChoice` class.
31 +#[serde_with::skip_serializing_none]
32 +#[derive(Debug, Clone, Serialize)]
33 +pub(super) struct GenerateResponseChoice {
34 + pub index: u32,
35 + pub logprobs: Option<ChatLogProbs>,
36 + pub finish_reason: Option<String>,
37 + pub token_ids: Vec<u32>,
38 +}
39 +
40 +/// Mirrors the Python vLLM `GenerateResponse` class.
41 +#[serde_with::skip_serializing_none]
42 +#[derive(Debug, Clone, Serialize)]
43 +pub(super) struct GenerateResponse {
44 + pub request_id: String,
45 + pub choices: Vec<GenerateResponseChoice>,
46 + pub prompt_logprobs: Option<Vec<Option<HashMap<u32, GenerateLogprob>>>>,
47 + pub kv_transfer_params: Option<Value>,
48 +}
49 +
50 +/// Mirrors the Python vLLM `Logprob` class used in prompt-logprobs payloads.
51 +#[serde_with::skip_serializing_none]
52 +#[derive(Debug, Clone, Serialize)]
53 +pub(super) struct GenerateLogprob {
54 + pub logprob: f32,
55 + pub rank: Option<u32>,
56 + pub decoded_token: Option<String>,
57 +}
rust/src/server/src/routes/inference/generate/validate.rs new
+84
@@ -0,0 +1,84 @@
1 +use super::types::GenerateRequest;
2 +use crate::error::{ApiError, bail_invalid_request};
3 +
4 +/// Enforce the minimal compatibility contract for the Rust token generate
5 +/// route.
6 +pub(super) fn validate_request_compat(
7 + request: &GenerateRequest,
8 + served_model_names: &[String],
9 +) -> Result<(), ApiError> {
10 + if let Some(model) = request.model.as_ref()
11 + && !served_model_names.iter().any(|n| n == model)
12 + {
13 + return Err(ApiError::model_not_found(model.clone()));
14 + }
15 +
16 + if request.stream {
17 + bail_invalid_request!(param = "stream", "stream=true is not supported.");
18 + }
19 +
20 + if request.token_ids.is_empty() {
21 + bail_invalid_request!(
22 + param = "token_ids",
23 + "token_ids must contain at least one token ID."
24 + );
25 + }
26 +
27 + if request.sampling_params.max_tokens == Some(0) {
28 + bail_invalid_request!(
29 + param = "sampling_params",
30 + "max_tokens must be greater than 0."
31 + );
32 + }
33 +
34 + if let Some(prompt_logprobs) = request.sampling_params.prompt_logprobs
35 + && prompt_logprobs < 0
36 + && prompt_logprobs != -1
37 + {
38 + bail_invalid_request!(
39 + param = "sampling_params",
40 + "`prompt_logprobs` must be a non-negative value or -1."
41 + );
42 + }
43 +
44 + Ok(())
45 +}
46 +
47 +#[cfg(test)]
48 +mod tests {
49 + use serde_json::json;
50 +
51 + use super::validate_request_compat;
52 + use crate::routes::inference::generate::types::GenerateRequest;
53 +
54 + fn base_request() -> GenerateRequest {
55 + serde_json::from_value(json!({
56 + "model": "Qwen/Qwen1.5-0.5B-Chat",
57 + "token_ids": [11, 22],
58 + "sampling_params": {}
59 + }))
60 + .expect("parse request")
61 + }
62 +
63 + fn served(names: &[&str]) -> Vec<String> {
64 + names.iter().map(|s| s.to_string()).collect()
65 + }
66 +
67 + #[test]
68 + fn validate_request_compat_rejects_streaming() {
69 + let request = GenerateRequest {
70 + stream: true,
71 + ..base_request()
72 + };
73 + assert!(validate_request_compat(&request, &served(&["Qwen/Qwen1.5-0.5B-Chat"])).is_err());
74 + }
75 +
76 + #[test]
77 + fn validate_request_compat_rejects_empty_token_ids() {
78 + let request = GenerateRequest {
79 + token_ids: Vec::new(),
80 + ..base_request()
81 + };
82 + assert!(validate_request_compat(&request, &served(&["Qwen/Qwen1.5-0.5B-Chat"])).is_err());
83 + }
84 +}
rust/src/server/src/routes/inference/mod.rs new
+3
@@ -0,0 +1,3 @@
1 +pub mod generate;
2 +
3 +pub use generate::generate;
rust/src/server/src/routes/load.rs new
+18
@@ -0,0 +1,18 @@
1 +use std::sync::Arc;
2 +
3 +use axum::Json;
4 +use axum::extract::State;
5 +use serde::Serialize;
6 +
7 +use crate::state::AppState;
8 +
9 +#[derive(Serialize)]
10 +pub(crate) struct ServerLoadResponse {
11 + server_load: u64,
12 +}
13 +
14 +pub async fn load(State(state): State<Arc<AppState>>) -> Json<ServerLoadResponse> {
15 + Json(ServerLoadResponse {
16 + server_load: state.server_load(),
17 + })
18 +}
rust/src/server/src/routes/metrics.rs new
+26
@@ -0,0 +1,26 @@
1 +use axum::http::header::CONTENT_TYPE;
2 +use axum::http::{HeaderValue, StatusCode};
3 +use axum::response::{IntoResponse, Response};
4 +use thiserror_ext::AsReport;
5 +use vllm_metrics::METRICS;
6 +
7 +const OPENMETRICS_CONTENT_TYPE: &str = "application/openmetrics-text; version=1.0.0; charset=utf-8";
8 +
9 +pub async fn scrape() -> Response {
10 + match METRICS.render() {
11 + Ok(body) => (
12 + [(
13 + CONTENT_TYPE,
14 + HeaderValue::from_static(OPENMETRICS_CONTENT_TYPE),
15 + )],
16 + body,
17 + )
18 + .into_response(),
19 +
20 + Err(error) => (
21 + StatusCode::INTERNAL_SERVER_ERROR,
22 + format!("failed to render metrics: {}", error.as_report()),
23 + )
24 + .into_response(),
25 + }
26 +}
rust/src/server/src/routes/openai/chat_completions.rs new
+1046
@@ -0,0 +1,1046 @@
1 +pub mod convert;
2 +mod types;
3 +mod validate;
4 +
5 +use std::convert::Infallible;
6 +use std::result::Result;
7 +use std::sync::Arc;
8 +
9 +use asynk_strim_attr::{TryYielder, try_stream};
10 +use axum::Json;
11 +use axum::extract::State;
12 +use axum::http::HeaderMap;
13 +use axum::response::sse::{Event, Sse};
14 +use axum::response::{IntoResponse, Response};
15 +use futures::{Stream, StreamExt as _, pin_mut};
16 +use serde_json::Value;
17 +use thiserror_ext::AsReport as _;
18 +use tracing::{debug, error, info, trace};
19 +use tracing_futures::Instrument as _;
20 +use vllm_chat::{
21 + AssistantBlockKind, AssistantMessageExt as _, ChatEvent, ChatEventStream, ChatEventStreamTrait,
22 + CollectedAssistantMessage, FinishReason,
23 +};
24 +use vllm_engine_core_client::protocol::StopReason;
25 +
26 +use crate::error::{ApiError, bail_server_error, server_error};
27 +use crate::routes::openai::chat_completions::convert::prepare_chat_request;
28 +use crate::routes::openai::chat_completions::types::{
29 + AssistantRole, ChatCompletionChoice, ChatCompletionMessage, ChatCompletionRequest,
30 + ChatCompletionResponse, ChatCompletionStreamChoice, ChatCompletionStreamResponse,
31 + ChatMessageDelta,
32 +};
33 +use crate::routes::openai::utils::logprobs::{
34 + decoded_logprobs_to_openai_chat, decoded_prompt_logprobs_to_maps,
35 +};
36 +use crate::routes::openai::utils::types::{
37 + ChatLogProbs, FunctionCallDelta, FunctionCallResponse, ToolCall, ToolCallDelta, Usage,
38 +};
39 +use crate::routes::openai::utils::validated_json::ValidatedJson;
40 +use crate::state::AppState;
41 +use crate::utils::{resolve_request_context, unix_timestamp};
42 +
43 +/// Validate one chat completion request and proxy it into the shared
44 +/// `vllm-chat` stack.
45 +pub async fn chat_completions(
46 + State(state): State<Arc<AppState>>,
47 + headers: HeaderMap,
48 + ValidatedJson(body): ValidatedJson<ChatCompletionRequest>,
49 +) -> Response {
50 + let stream = body.stream;
51 + let request_context = resolve_request_context(&headers, body.request_id.as_deref());
52 +
53 + let prepared = match prepare_chat_request(body, state.served_model_names(), request_context) {
54 + Ok(prepared) => prepared,
55 + Err(error) => return error.into_response(),
56 + };
57 + let request_span = tracing::info_span!(
58 + "chat_completions",
59 + request_id = %prepared.request_id,
60 + engine_request_id = tracing::field::Empty,
61 + );
62 +
63 + let created = unix_timestamp();
64 + let log_request = state.enable_log_requests;
65 +
66 + let chat_stream =
67 + match state.chat.chat(prepared.chat_request).instrument(request_span.clone()).await {
68 + Ok(stream) => stream,
69 + Err(error) => {
70 + return server_error!(
71 + "failed to submit chat request: {}",
72 + error.to_report_string()
73 + )
74 + .into_response();
75 + }
76 + };
77 +
78 + if stream {
79 + let chunk_stream = chat_completion_chunk_stream(
80 + chat_stream,
81 + prepared.request_id,
82 + prepared.response_model,
83 + created,
84 + log_request,
85 + prepared.include_usage,
86 + prepared.requested_logprobs,
87 + prepared.echo,
88 + prepared.return_token_ids,
89 + prepared.return_tokens_as_token_ids,
90 + );
91 + let sse_stream = chat_completion_sse_stream(chunk_stream).instrument(request_span);
92 +
93 + Sse::new(sse_stream).into_response()
94 + } else {
95 + let response = match collect_chat_completion(
96 + chat_stream,
97 + prepared.request_id,
98 + prepared.response_model,
99 + created,
100 + prepared.requested_logprobs,
101 + prepared.include_prompt_logprobs,
102 + prepared.echo,
103 + prepared.return_token_ids,
104 + prepared.return_tokens_as_token_ids,
105 + )
106 + .instrument(request_span.clone())
107 + .await
108 + {
109 + Ok(response) => response,
110 + Err(error) => return error.into_response(),
111 + };
112 +
113 + if log_request {
114 + let usage = response.usage.as_ref();
115 + info!(
116 + parent: &request_span,
117 + model = %response.model,
118 + prompt_tokens = usage.map_or(0, |u| u.prompt_tokens),
119 + output_tokens = usage.and_then(|u| u.completion_tokens).unwrap_or(0),
120 + finish_reason = response.choices.first().and_then(|c| c.finish_reason.as_deref()).unwrap_or("unknown"),
121 + "chat completion finished"
122 + );
123 + }
124 +
125 + Json(response).into_response()
126 + }
127 +}
128 +
129 +async fn collect_chat_completion(
130 + stream: ChatEventStream,
131 + request_id: String,
132 + response_model: String,
133 + created: u64,
134 + requested_logprobs: bool,
135 + include_prompt_logprobs: bool,
136 + echo: Option<String>,
137 + return_token_ids: bool,
138 + return_tokens_as_token_ids: bool,
139 +) -> Result<ChatCompletionResponse, ApiError> {
140 + let collected = stream.collect_message().await.map_err(|error| {
141 + server_error!(
142 + "failed to collect chat completion response: {}",
143 + error.to_report_string()
144 + )
145 + })?;
146 + let CollectedAssistantMessage {
147 + message,
148 + prompt_token_count,
149 + prompt_token_ids,
150 + prompt_logprobs,
151 + logprobs,
152 + token_ids,
153 + output_token_count,
154 + finish_reason,
155 + kv_transfer_params,
156 + } = collected;
157 + let stop_reason = finish_reason.as_stop_reason().map(stop_reason_to_json);
158 + let saw_tool_calls = message.tool_calls().next().is_some();
159 + let finish_reason = chat_finish_reason_to_openai(&finish_reason, saw_tool_calls)?.to_string();
160 + let tool_calls = message
161 + .tool_calls()
162 + .map(|call| ToolCall {
163 + id: call.id.clone(),
164 + tool_type: "function".to_string(),
165 + function: FunctionCallResponse {
166 + name: call.name.clone(),
167 + arguments: Some(call.arguments.clone()),
168 + },
169 + })
170 + .collect::<Vec<_>>();
171 + let logprobs = if requested_logprobs {
172 + Some(decoded_logprobs_to_openai_chat(
173 + logprobs.as_ref().ok_or_else(|| {
174 + server_error!("chat response requested logprobs but generation returned none")
175 + })?,
176 + return_tokens_as_token_ids,
177 + )?)
178 + } else {
179 + None
180 + };
181 + let prompt_logprobs = if include_prompt_logprobs {
182 + Some(decoded_prompt_logprobs_to_maps(
183 + prompt_logprobs.as_ref().ok_or_else(|| {
184 + server_error!(
185 + "chat response requested prompt_logprobs but generation returned none"
186 + )
187 + })?,
188 + return_tokens_as_token_ids,
189 + ))
190 + } else {
191 + None
192 + };
193 + let usage = Usage::from_counts(prompt_token_count as u32, output_token_count as u32);
194 +
195 + Ok(ChatCompletionResponse {
196 + id: request_id,
197 + object: "chat.completion".to_string(),
198 + created,
199 + model: response_model,
200 + choices: vec![ChatCompletionChoice {
201 + index: 0,
202 + message: ChatCompletionMessage {
203 + role: AssistantRole,
204 + content: match &echo {
205 + Some(prefix) => Some(format!("{prefix}{}", message.text())),
206 + None => Some(message.text()).filter(|t| !t.is_empty()),
207 + },
208 + tool_calls: Some(tool_calls).filter(|calls| !calls.is_empty()),
209 + reasoning: message.reasoning(),
210 + },
211 + logprobs,
212 + finish_reason: Some(finish_reason),
213 + stop_reason,
214 + token_ids: return_token_ids.then_some(token_ids),
215 + }],
216 + usage: Some(usage),
217 + system_fingerprint: None,
218 + prompt_logprobs,
219 + prompt_token_ids: return_token_ids.then(|| prompt_token_ids.to_vec()),
220 + kv_transfer_params,
221 + })
222 +}
223 +
224 +/// Convert one internal chat event stream into OpenAI chat-completion chunks.
225 +#[try_stream]
226 +async fn chat_completion_chunk_stream(
227 + mut stream: impl ChatEventStreamTrait + Unpin,
228 + request_id: String,
229 + response_model: String,
230 + created: u64,
231 + log_request: bool,
232 + include_usage: bool,
233 + requested_logprobs: bool,
234 + echo: Option<String>,
235 + return_token_ids: bool,
236 + return_tokens_as_token_ids: bool,
237 + mut y: TryYielder<ChatCompletionStreamResponse, ApiError>,
238 +) -> Result<(), ApiError> {
239 + let mut saw_tool_calls = false;
240 +
241 + // If the client requested logprobs or token_ids, we need to buffer chunks until
242 + // we receive the separate `LogprobsDelta` event, so that we can emit one
243 + // combined chunk with both the semantic delta and its per-update metadata.
244 + let mut pending_chunk =
245 + (requested_logprobs || return_token_ids).then(PendingChatChunk::default);
246 +
247 + while let Some(next) = stream.next().await {
248 + match next {
249 + Ok(ChatEvent::Start {
250 + prompt_token_ids, ..
251 + }) => {
252 + let mut chunk = start_chunk(&request_id, &response_model, created);
253 + if return_token_ids {
254 + chunk.prompt_token_ids = Some(prompt_token_ids.to_vec());
255 + }
256 + y.yield_ok(chunk).await;
257 + // When echo=true, emit the last assistant message content as a delta chunk.
258 + if let Some(echo_text) = &echo {
259 + y.yield_ok(block_delta_chunk(
260 + &request_id,
261 + &response_model,
262 + created,
263 + AssistantBlockKind::Text,
264 + echo_text.clone(),
265 + ))
266 + .await;
267 + }
268 + }
269 + Ok(ChatEvent::BlockDelta { kind, delta, .. }) => {
270 + if let Some(pending_chunk) = pending_chunk.as_mut() {
271 + pending_chunk.push_block_delta(kind, delta);
272 + } else {
273 + y.yield_ok(block_delta_chunk(
274 + &request_id,
275 + &response_model,
276 + created,
277 + kind,
278 + delta,
279 + ))
280 + .await;
281 + }
282 + }
283 + Ok(ChatEvent::LogprobsDelta {
284 + logprobs,
285 + token_ids,
286 + }) => {
287 + let openai_logprobs = logprobs
288 + .as_ref()
289 + .map(|lp| decoded_logprobs_to_openai_chat(lp, return_tokens_as_token_ids))
290 + .transpose()?;
291 + let openai_token_ids =
292 + return_token_ids.then_some(token_ids).filter(|t| !t.is_empty());
293 + if let Some(pending_chunk) = pending_chunk.as_mut() {
294 + pending_chunk.logprobs = openai_logprobs;
295 + pending_chunk.token_ids = openai_token_ids;
296 + if let Some(chunk) =
297 + pending_chunk.take_chunk(&request_id, &response_model, created)
298 + {
299 + y.yield_ok(chunk).await;
300 + }
301 + } else if let Some(logprobs) = openai_logprobs {
302 + y.yield_ok(logprobs_only_chunk(
303 + &request_id,
304 + &response_model,
305 + created,
306 + logprobs,
307 + ))
308 + .await;
309 + }
310 + }
311 + Ok(ChatEvent::BlockStart { kind, .. }) => {
312 + debug!(?kind, "starting new block");
313 + }
314 + Ok(ChatEvent::BlockEnd { .. }) => {
315 + debug!("ending current block");
316 + }
317 + Ok(ChatEvent::ToolCallStart { index, id, name }) => {
318 + let tool_index = index as u32;
319 + saw_tool_calls = true;
320 + debug!(
321 + tool_call_id = %id,
322 + tool_call_name = %name,
323 + "starting new tool call"
324 + );
325 + if let Some(pending_chunk) = pending_chunk.as_mut() {
326 + pending_chunk.push_tool_call_start(tool_index, id, name);
327 + } else {
328 + y.yield_ok(tool_call_start_chunk(
329 + &request_id,
330 + &response_model,
331 + created,
332 + tool_index,
333 + id,
334 + name,
335 + ))
336 + .await;
337 + }
338 + }
339 + Ok(ChatEvent::ToolCallArgumentsDelta { index, delta }) => {
340 + let tool_index = index as u32;
341 + if let Some(pending_chunk) = pending_chunk.as_mut() {
342 + pending_chunk.push_tool_call_arguments(tool_index, delta);
343 + } else {
344 + y.yield_ok(tool_call_arguments_chunk(
345 + &request_id,
346 + &response_model,
347 + created,
348 + tool_index,
349 + delta,
350 + ))
351 + .await;
352 + }
353 + }
354 + Ok(ChatEvent::ToolCallEnd { .. }) => {
355 + debug!("ending current tool call");
356 + }
357 + Ok(ChatEvent::Done {
358 + prompt_token_count,
359 + finish_reason,
360 + output_token_count,
361 + ..
362 + }) => {
363 + if log_request {
364 + info!(
365 + stream = true,
366 + model = %response_model,
367 + prompt_tokens = prompt_token_count,
368 + output_tokens = output_token_count,
369 + finish_reason = finish_reason.as_str(),
370 + "chat completion finished"
371 + );
372 + }
373 +
374 + if let Some(pending_chunk) = pending_chunk.as_mut()
375 + && let Some(chunk) =
376 + pending_chunk.take_chunk(&request_id, &response_model, created)
377 + {
378 + y.yield_ok(chunk).await;
379 + }
380 +
381 + match final_chunk(
382 + &request_id,
383 + &response_model,
384 + created,
385 + finish_reason,
386 + saw_tool_calls,
387 + ) {
388 + Ok(chunk) => y.yield_ok(chunk).await,
389 + Err(error) => {
390 + error!(
391 + error = %error.to_error_response().error.message,
392 + "invalid terminal finish reason"
393 + );
394 + return Err(error);
395 + }
396 + }
397 +
398 + if include_usage {
399 + y.yield_ok(usage_chunk(
400 + &request_id,
401 + &response_model,
402 + created,
403 + Usage::from_counts(prompt_token_count as u32, output_token_count as u32),
404 + ))
405 + .await;
406 + }
407 +
408 + return Ok(());
409 + }
410 + Err(error) => {
411 + error!(
412 + error = %error.as_report(),
413 + "chat stream failed"
414 + );
415 + bail_server_error!("{}", error.to_report_string());
416 + }
417 + }
418 + }
419 + Ok(())
420 +}
421 +
422 +fn usage_chunk(
423 + request_id: &str,
424 + response_model: &str,
425 + created: u64,
426 + usage: Usage,
427 +) -> ChatCompletionStreamResponse {
428 + let mut chunk = ChatCompletionStreamResponse::new(request_id, response_model, created);
429 + chunk.usage = Some(usage);
430 + chunk
431 +}
432 +
433 +/// One in-flight chat-completions SSE chunk being assembled at the route layer.
434 +///
435 +/// `vllm-chat` emits semantic chat events first and `LogprobsDelta` separately,
436 +/// because one decoded update may be rewritten into multiple chat events.
437 +/// The OpenAI chat API, though, wants one streamed chunk to optionally carry
438 +/// both the delta and its logprobs.
439 +///
440 +/// This small buffer accumulates the semantic delta first, then attaches the
441 +/// following `LogprobsDelta` and flushes one combined chunk. It relies on the
442 +/// current `vllm-chat` invariant that all semantic events from one decoded
443 +/// update are emitted before that update's `LogprobsDelta`.
444 +#[derive(Debug, Default)]
445 +struct PendingChatChunk {
446 + /// The currently buffered OpenAI delta payload assembled from one or more
447 + /// chat semantic events belonging to the same decoded update.
448 + delta: ChatMessageDelta,
449 + /// The token-aligned logprobs for that same decoded update.
450 + logprobs: Option<ChatLogProbs>,
451 + /// Per-update output token IDs for the same decoded update.
452 + token_ids: Option<Vec<u32>>,
453 +}
454 +
455 +impl PendingChatChunk {
456 + /// Append one assistant text/reasoning block delta to the buffered OpenAI
457 + /// delta payload.
458 + fn push_block_delta(&mut self, kind: AssistantBlockKind, delta: String) {
459 + match kind {
460 + AssistantBlockKind::Text => append_delta_text(&mut self.delta.content, delta),
461 + AssistantBlockKind::Reasoning => append_delta_text(&mut self.delta.reasoning, delta),
462 + AssistantBlockKind::ToolCall => {
463 + unreachable!("tool calls must flow through dedicated tool-call chunks")
464 + }
465 + }
466 + }
467 +
468 + /// Append the OpenAI tool-call-start representation to the buffered delta.
469 + fn push_tool_call_start(&mut self, index: u32, id: String, name: String) {
470 + self.delta.tool_calls.get_or_insert_with(Vec::new).push(ToolCallDelta {
471 + index,
472 + id: Some(id),
473 + tool_type: Some("function".to_string()),
474 + function: Some(FunctionCallDelta {
475 + name: Some(name),
476 + arguments: None,
477 + }),
478 + });
479 + }
480 +
481 + /// Append one incremental tool-call arguments update to the buffered delta.
482 + fn push_tool_call_arguments(&mut self, index: u32, delta: String) {
483 + self.delta.tool_calls.get_or_insert_with(Vec::new).push(ToolCallDelta {
484 + index,
485 + id: None,
486 + tool_type: None,
487 + function: Some(FunctionCallDelta {
488 + name: None,
489 + arguments: Some(delta),
490 + }),
491 + });
492 + }
493 +
494 + /// Finalize the currently buffered SSE chunk, if it contains either a
495 + /// semantic delta or a logprobs payload.
496 + ///
497 + /// This may produce:
498 + /// - a combined delta + logprobs chunk
499 + /// - a delta-only chunk
500 + /// - a logprobs-only chunk
501 + ///
502 + /// The logprobs-only case is intentional: token-level metadata in one
503 + /// decoded update is correlated with the same update boundary, not
504 + /// necessarily with a visible/chat-semantic delta.
505 + fn take_chunk(
506 + &mut self,
507 + request_id: &str,
508 + response_model: &str,
509 + created: u64,
510 + ) -> Option<ChatCompletionStreamResponse> {
511 + let has_delta = self.delta.content.is_some()
512 + || self.delta.reasoning.is_some()
513 + || self.delta.tool_calls.is_some();
514 + let logprobs = self.logprobs.take();
515 + let token_ids = self.token_ids.take();
516 + if !has_delta && logprobs.is_none() && token_ids.is_none() {
517 + return None;
518 + }
519 +
520 + let mut chunk = ChatCompletionStreamResponse::new(request_id, response_model, created);
521 + chunk.choices.push(ChatCompletionStreamChoice {
522 + delta: self.take_delta(),
523 + logprobs,
524 + token_ids,
525 + ..Default::default()
526 + });
527 + Some(chunk)
528 + }
529 +
530 + /// Take the currently buffered OpenAI delta payload and leave this pending
531 + /// chunk empty for the next decoded update.
532 + fn take_delta(&mut self) -> ChatMessageDelta {
533 + ChatMessageDelta {
534 + role: self.delta.role.take(),
535 + content: self.delta.content.take(),
536 + tool_calls: self.delta.tool_calls.take(),
537 + reasoning: self.delta.reasoning.take(),
538 + }
539 + }
540 +}
541 +
542 +/// Append one text fragment to an optional OpenAI delta string field.
543 +fn append_delta_text(slot: &mut Option<String>, delta: String) {
544 + match slot {
545 + Some(existing) => existing.push_str(&delta),
546 + None => *slot = Some(delta),
547 + }
548 +}
549 +
550 +/// Convert one chunk stream into OpenAI-style SSE events.
551 +///
552 +/// OpenAI-style streaming errors are encoded as ordinary `data: {"error": ...}`
553 +/// events followed by `data: [DONE]`, so the transport stream itself stays
554 +/// infallible even when generation fails after the HTTP response has started.
555 +#[try_stream]
556 +async fn chat_completion_sse_stream(
557 + stream: impl Stream<Item = Result<ChatCompletionStreamResponse, ApiError>>,
558 + mut y: TryYielder<Event, Infallible>,
559 +) -> Result<(), Infallible> {
560 + pin_mut!(stream);
561 +
562 + while let Some(next) = stream.next().await {
563 + match next {
564 + Ok(chunk) => y.yield_ok(to_sse_event(&chunk)).await,
565 + Err(error) => {
566 + y.yield_ok(to_error_sse_event(&error)).await;
567 + break;
568 + }
569 + }
570 + }
571 +
572 + y.yield_ok(done_sse_event()).await;
573 + Ok(())
574 +}
575 +
576 +/// Serialize one OpenAI chunk payload into one SSE `data:` event.
577 +fn to_sse_event(chunk: &ChatCompletionStreamResponse) -> Event {
578 + let payload =
579 + serde_json::to_string(chunk).expect("ChatCompletionStreamResponse must serialize to JSON");
580 + trace!(payload, "chat completion emitting chunk");
581 + Event::default().data(payload)
582 +}
583 +
584 +/// Serialize one OpenAI error payload into one SSE `data:` event.
585 +fn to_error_sse_event(error: &ApiError) -> Event {
586 + let payload = serde_json::to_string(&error.to_error_response())
587 + .expect("ErrorResponse must serialize to JSON");
588 + trace!(payload, "chat completion emitting error");
589 + Event::default().data(payload)
590 +}
591 +
592 +/// Build the terminal OpenAI SSE sentinel event.
593 +fn done_sse_event() -> Event {
594 + trace!("chat completion emitting done");
595 + Event::default().data("[DONE]")
596 +}
597 +
598 +/// Build the initial assistant-role SSE chunk required by the OpenAI streaming
599 +/// protocol.
600 +fn start_chunk(
601 + request_id: &str,
602 + response_model: &str,
603 + created: u64,
604 +) -> ChatCompletionStreamResponse {
605 + let mut chunk = ChatCompletionStreamResponse::new(request_id, response_model, created);
606 + chunk.choices.push(ChatCompletionStreamChoice {
607 + delta: ChatMessageDelta {
608 + role: Some(AssistantRole),
609 + ..Default::default()
610 + },
611 + ..Default::default()
612 + });
613 + chunk
614 +}
615 +
616 +/// Build one content-delta SSE chunk from one internal assistant block delta.
617 +fn block_delta_chunk(
618 + request_id: &str,
619 + response_model: &str,
620 + created: u64,
621 + kind: AssistantBlockKind,
622 + delta: String,
623 +) -> ChatCompletionStreamResponse {
624 + let delta = match kind {
625 + AssistantBlockKind::Text => ChatMessageDelta {
626 + content: Some(delta),
627 + ..Default::default()
628 + },
629 + AssistantBlockKind::Reasoning => ChatMessageDelta {
630 + reasoning: Some(delta),
631 + ..Default::default()
632 + },
633 + AssistantBlockKind::ToolCall => {
634 + unreachable!("tool calls must flow through dedicated tool-call chunks")
635 + }
636 + };
637 +
638 + let mut chunk = ChatCompletionStreamResponse::new(request_id, response_model, created);
639 + chunk.choices.push(ChatCompletionStreamChoice {
640 + delta,
641 + ..Default::default()
642 + });
643 + chunk
644 +}
645 +
646 +fn tool_call_start_chunk(
647 + request_id: &str,
648 + response_model: &str,
649 + created: u64,
650 + tool_index: u32,
651 + id: String,
652 + name: String,
653 +) -> ChatCompletionStreamResponse {
654 + let mut chunk = ChatCompletionStreamResponse::new(request_id, response_model, created);
655 + chunk.choices.push(ChatCompletionStreamChoice {
656 + delta: ChatMessageDelta {
657 + tool_calls: Some(vec![ToolCallDelta {
658 + index: tool_index,
659 + id: Some(id),
660 + tool_type: Some("function".to_string()),
661 + function: Some(FunctionCallDelta {
662 + name: Some(name),
663 + arguments: None,
664 + }),
665 + }]),
666 + ..Default::default()
667 + },
668 + ..Default::default()
669 + });
670 + chunk
671 +}
672 +
673 +fn tool_call_arguments_chunk(
674 + request_id: &str,
675 + response_model: &str,
676 + created: u64,
677 + tool_index: u32,
678 + delta: String,
679 +) -> ChatCompletionStreamResponse {
680 + let mut chunk = ChatCompletionStreamResponse::new(request_id, response_model, created);
681 + chunk.choices.push(ChatCompletionStreamChoice {
682 + delta: ChatMessageDelta {
683 + tool_calls: Some(vec![ToolCallDelta {
684 + index: tool_index,
685 + id: None,
686 + tool_type: None,
687 + function: Some(FunctionCallDelta {
688 + name: None,
689 + arguments: Some(delta),
690 + }),
691 + }]),
692 + ..Default::default()
693 + },
694 + ..Default::default()
695 + });
696 + chunk
697 +}
698 +
699 +fn logprobs_only_chunk(
700 + request_id: &str,
701 + response_model: &str,
702 + created: u64,
703 + logprobs: ChatLogProbs,
704 +) -> ChatCompletionStreamResponse {
705 + let mut chunk = ChatCompletionStreamResponse::new(request_id, response_model, created);
706 + chunk.choices.push(ChatCompletionStreamChoice {
707 + logprobs: Some(logprobs),
708 + ..Default::default()
709 + });
710 + chunk
711 +}
712 +
713 +/// Build the terminal SSE chunk carrying the OpenAI finish reason.
714 +fn final_chunk(
715 + request_id: &str,
716 + response_model: &str,
717 + created: u64,
718 + finish_reason: FinishReason,
719 + saw_tool_calls: bool,
720 +) -> Result<ChatCompletionStreamResponse, ApiError> {
721 + let stop_reason = finish_reason.as_stop_reason().map(stop_reason_to_json);
722 + let finish_reason = chat_finish_reason_to_openai(&finish_reason, saw_tool_calls)?;
723 +
724 + debug!(
725 + finish_reason = %finish_reason,
726 + stop_reason = ?stop_reason,
727 + "chat stream finished"
728 + );
729 +
730 + let mut chunk = ChatCompletionStreamResponse::new(request_id, response_model, created);
731 + chunk.choices.push(ChatCompletionStreamChoice {
732 + finish_reason: Some(finish_reason.to_string()),
733 + stop_reason,
734 + ..Default::default()
735 + });
736 + Ok(chunk)
737 +}
738 +
739 +fn chat_finish_reason_to_openai(
740 + finish_reason: &FinishReason,
741 + saw_tool_calls: bool,
742 +) -> Result<&'static str, ApiError> {
743 + match finish_reason {
744 + FinishReason::Stop(_) if saw_tool_calls => Ok("tool_calls"),
745 + FinishReason::Stop(_) => Ok("stop"),
746 + FinishReason::Length => Ok("length"),
747 + FinishReason::Abort => Ok("abort"),
748 + FinishReason::Repetition => Ok("stop"),
749 + FinishReason::Error => {
750 + bail_server_error!("Internal server error");
751 + }
752 + }
753 +}
754 +
755 +/// Convert one internal stop reason into the OpenAI-compatible `stop_reason`
756 +/// JSON shape.
757 +fn stop_reason_to_json(stop_reason: &StopReason) -> Value {
758 + serde_json::to_value(stop_reason).expect("StopReason must serialize to JSON")
759 +}
760 +
761 +#[cfg(test)]
762 +mod tests {
763 + use futures::{StreamExt as _, stream};
764 + use serde_json::json;
765 + use vllm_chat::{AssistantBlockKind, AssistantToolCall, ChatEvent, FinishReason};
766 + use vllm_engine_core_client::protocol::StopReason;
767 + use vllm_text::{DecodedLogprobs, DecodedPositionLogprobs, DecodedTokenLogprob};
768 +
769 + use super::{block_delta_chunk, chat_completion_chunk_stream, final_chunk};
770 +
771 + #[test]
772 + fn text_chunk_uses_content_only_delta() {
773 + let chunk = block_delta_chunk(
774 + "chatcmpl-1",
775 + "model",
776 + 1,
777 + AssistantBlockKind::Text,
778 + "hello".to_string(),
779 + );
780 + assert_eq!(chunk.choices[0].delta.role, None);
781 + assert_eq!(chunk.choices[0].delta.content.as_deref(), Some("hello"));
782 + assert_eq!(chunk.choices[0].delta.reasoning, None);
783 + }
784 +
785 + #[test]
786 + fn reasoning_chunk_uses_reasoning_only_delta() {
787 + let chunk = block_delta_chunk(
788 + "chatcmpl-1",
789 + "model",
790 + 1,
791 + AssistantBlockKind::Reasoning,
792 + "thinking".to_string(),
793 + );
794 + assert_eq!(chunk.choices[0].delta.role, None);
795 + assert_eq!(chunk.choices[0].delta.content, None);
796 + assert_eq!(
797 + chunk.choices[0].delta.reasoning.as_deref(),
798 + Some("thinking")
799 + );
800 + }
801 +
802 + #[test]
803 + fn final_chunk_maps_stop_finish_reason_and_stop_reason() {
804 + let chunk = final_chunk(
805 + "chatcmpl-1",
806 + "model",
807 + 1,
808 + FinishReason::Stop(Some(StopReason::Text("stop".to_string()))),
809 + false,
810 + )
811 + .expect("finish reason is valid");
812 +
813 + assert_eq!(chunk.choices[0].finish_reason.as_deref(), Some("stop"));
814 + assert_eq!(chunk.choices[0].stop_reason, Some(json!("stop")));
815 + }
816 +
817 + #[test]
818 + fn final_chunk_maps_length_finish_reason() {
819 + let chunk = final_chunk("chatcmpl-1", "model", 1, FinishReason::Length, false)
820 + .expect("finish reason is valid");
821 +
822 + assert_eq!(chunk.choices[0].finish_reason.as_deref(), Some("length"));
823 + assert_eq!(chunk.choices[0].stop_reason, None);
824 + }
825 +
826 + #[test]
827 + fn final_chunk_maps_abort_finish_reason() {
828 + let chunk = final_chunk("chatcmpl-1", "model", 1, FinishReason::Abort, false)
829 + .expect("abort is a valid finish reason");
830 +
831 + assert_eq!(chunk.choices[0].finish_reason.as_deref(), Some("abort"));
832 + assert_eq!(chunk.choices[0].stop_reason, None);
833 + }
834 +
835 + #[test]
836 + fn final_chunk_rejects_error_finish_reason() {
837 + assert!(final_chunk("chatcmpl-1", "model", 1, FinishReason::Error, false).is_err());
838 + }
839 +
840 + #[test]
841 + fn final_chunk_maps_stop_to_tool_calls_when_tool_calls_were_streamed() {
842 + let chunk = final_chunk("chatcmpl-1", "model", 1, FinishReason::stop_eos(), true)
843 + .expect("finish reason is valid");
844 +
845 + assert_eq!(
846 + chunk.choices[0].finish_reason.as_deref(),
847 + Some("tool_calls")
848 + );
849 + }
850 +
851 + #[tokio::test]
852 + async fn chunk_stream_coalesces_text_delta_with_logprobs() {
853 + let stream = stream::iter(vec![
854 + Ok(ChatEvent::Start {
855 + prompt_token_ids: vec![].into(),
856 + prompt_logprobs: None,
857 + }),
858 + Ok(ChatEvent::BlockStart {
859 + index: 0,
860 + kind: AssistantBlockKind::Text,
861 + }),
862 + Ok(ChatEvent::BlockDelta {
863 + index: 0,
864 + kind: AssistantBlockKind::Text,
865 + delta: "hi".to_string(),
866 + }),
867 + Ok(ChatEvent::LogprobsDelta {
868 + logprobs: Some(DecodedLogprobs {
869 + positions: vec![DecodedPositionLogprobs {
870 + entries: vec![DecodedTokenLogprob {
871 + token_id: 0,
872 + token: "hi".to_string(),
873 + logprob: -0.1,
874 + rank: 1,
875 + }],
876 + }],
877 + }),
878 + token_ids: vec![],
879 + }),
880 + Ok(ChatEvent::Done {
881 + message: Default::default(),
882 + prompt_token_count: 1,
883 + output_token_count: 1,
884 + finish_reason: FinishReason::stop_eos(),
885 + kv_transfer_params: None,
886 + }),
887 + ]);
888 +
889 + let chunks = chat_completion_chunk_stream(
890 + stream,
891 + "chatcmpl-1".to_string(),
892 + "model".to_string(),
893 + 1,
894 + false,
895 + false,
896 + true,
897 + None,
898 + false,
899 + false,
900 + )
901 + .collect::<Vec<_>>()
902 + .await
903 + .into_iter()
904 + .collect::<Result<Vec<_>, _>>()
905 + .expect("stream chunks");
906 +
907 + assert_eq!(chunks.len(), 3);
908 + assert_eq!(chunks[1].choices[0].delta.content.as_deref(), Some("hi"));
909 + let logprobs = chunks[1].choices[0].logprobs.as_ref().expect("logprobs");
910 + let content = logprobs.content.as_ref().expect("logprobs content");
911 + assert_eq!(content[0].token, "hi");
912 + }
913 +
914 + #[tokio::test]
915 + async fn chunk_stream_coalesces_reasoning_delta_with_logprobs() {
916 + let stream = stream::iter(vec![
917 + Ok(ChatEvent::Start {
918 + prompt_token_ids: vec![].into(),
919 + prompt_logprobs: None,
920 + }),
921 + Ok(ChatEvent::BlockStart {
922 + index: 0,
923 + kind: AssistantBlockKind::Reasoning,
924 + }),
925 + Ok(ChatEvent::BlockDelta {
926 + index: 0,
927 + kind: AssistantBlockKind::Reasoning,
928 + delta: "think".to_string(),
929 + }),
930 + Ok(ChatEvent::LogprobsDelta {
931 + logprobs: Some(DecodedLogprobs {
932 + positions: vec![DecodedPositionLogprobs {
933 + entries: vec![DecodedTokenLogprob {
934 + token_id: 0,
935 + token: "think".to_string(),
936 + logprob: -0.1,
937 + rank: 1,
938 + }],
939 + }],
940 + }),
941 + token_ids: vec![],
942 + }),
943 + Ok(ChatEvent::Done {
944 + message: Default::default(),
945 + prompt_token_count: 1,
946 + output_token_count: 1,
947 + finish_reason: FinishReason::stop_eos(),
948 + kv_transfer_params: None,
949 + }),
950 + ]);
951 +
952 + let chunks = chat_completion_chunk_stream(
953 + stream,
954 + "chatcmpl-1".to_string(),
955 + "model".to_string(),
956 + 1,
957 + false,
958 + false,
959 + true,
960 + None,
961 + false,
962 + false,
963 + )
964 + .collect::<Vec<_>>()
965 + .await
966 + .into_iter()
967 + .collect::<Result<Vec<_>, _>>()
968 + .expect("stream chunks");
969 +
970 + assert_eq!(chunks.len(), 3);
971 + assert_eq!(
972 + chunks[1].choices[0].delta.reasoning.as_deref(),
973 + Some("think")
974 + );
975 + assert!(chunks[1].choices[0].logprobs.is_some());
976 + }
977 +
978 + #[tokio::test]
979 + async fn chunk_stream_preserves_tool_call_index_and_omits_id_from_arguments_delta() {
980 + let stream = stream::iter(vec![
981 + Ok(ChatEvent::Start {
982 + prompt_token_ids: vec![].into(),
983 + prompt_logprobs: None,
984 + }),
985 + Ok(ChatEvent::ToolCallStart {
986 + index: 3,
987 + id: "call_1".to_string(),
988 + name: "get_weather".to_string(),
989 + }),
990 + Ok(ChatEvent::ToolCallArgumentsDelta {
991 + index: 3,
992 + delta: r#"{"city":"Paris"}"#.to_string(),
993 + }),
994 + Ok(ChatEvent::ToolCallEnd {
995 + index: 3,
996 + call: AssistantToolCall {
997 + id: "call_1".to_string(),
998 + name: "get_weather".to_string(),
999 + arguments: r#"{"city":"Paris"}"#.to_string(),
1000 + },
1001 + }),
1002 + Ok(ChatEvent::Done {
1003 + message: Default::default(),
1004 + prompt_token_count: 1,
1005 + output_token_count: 1,
1006 + finish_reason: FinishReason::stop_eos(),
1007 + kv_transfer_params: None,
1008 + }),
1009 + ]);
1010 +
1011 + let chunks = chat_completion_chunk_stream(
1012 + stream,
1013 + "chatcmpl-1".to_string(),
1014 + "model".to_string(),
1015 + 1,
1016 + false,
1017 + false,
1018 + false,
1019 + None,
1020 + false,
1021 + false,
1022 + )
1023 + .collect::<Vec<_>>()
1024 + .await
1025 + .into_iter()
1026 + .collect::<Result<Vec<_>, _>>()
1027 + .expect("stream chunks");
1028 +
1029 + assert_eq!(
1030 + chunks[1].choices[0].delta.tool_calls.as_ref().unwrap()[0].index,
1031 + 3
1032 + );
1033 + assert_eq!(
1034 + chunks[1].choices[0].delta.tool_calls.as_ref().unwrap()[0].id,
1035 + Some("call_1".to_string())
1036 + );
1037 + assert_eq!(
1038 + chunks[2].choices[0].delta.tool_calls.as_ref().unwrap()[0].index,
1039 + 3
1040 + );
1041 + assert_eq!(
1042 + chunks[2].choices[0].delta.tool_calls.as_ref().unwrap()[0].id,
1043 + None
1044 + );
1045 + }
1046 +}
rust/src/server/src/routes/openai/chat_completions/convert.rs new
+992
@@ -0,0 +1,992 @@
1 +use itertools::Itertools as _;
2 +use vllm_chat::{
3 + AssistantContentBlock, AssistantToolCall, ChatContent, ChatContentPart,
4 + ChatMessage as VllmChatMessage, ChatOptions, ChatRequest, ChatTool, ChatToolChoice,
5 + GenerationPromptMode, SamplingParams,
6 +};
7 +
8 +use super::types::ChatCompletionRequest;
9 +use super::validate;
10 +use crate::error::{ApiError, bail_invalid_request};
11 +use crate::routes::openai::utils::structured_outputs::convert_from_response_format;
12 +use crate::routes::openai::utils::types::{
13 + ChatMessage, ContentPart, MessageContent, Tool, ToolChoice, ToolChoiceValue,
14 +};
15 +use crate::utils::{ResolvedRequestContext, convert_logit_bias, merge_kv_transfer_params};
16 +
17 +/// Lowered chat request plus the public response metadata carried by every SSE
18 +/// chunk.
19 +#[derive(Debug, Clone, PartialEq)]
20 +pub struct PreparedRequest {
21 + /// Stable OpenAI-style request ID, reused as the external chat request ID.
22 + pub request_id: String,
23 + /// Public model ID echoed back to the client.
24 + pub response_model: String,
25 + /// Whether the caller asked for the final streamed usage chunk.
26 + pub include_usage: bool,
27 + /// Whether the caller requested output logprobs on chat choices.
28 + pub requested_logprobs: bool,
29 + /// Whether the caller requested top-level prompt logprobs.
30 + pub include_prompt_logprobs: bool,
31 + /// Lowered chat request for `vllm-chat`.
32 + pub chat_request: ChatRequest,
33 + /// Last assistant-role message content to echo back when `echo=true`.
34 + pub echo: Option<String>,
35 + /// Whether to include token IDs alongside generated text.
36 + pub return_token_ids: bool,
37 + /// Whether to format logprob tokens as `token_id:{id}`.
38 + pub return_tokens_as_token_ids: bool,
39 +}
40 +
41 +/// Validate and lower one OpenAI chat completion request into the internal chat
42 +/// format.
43 +///
44 +/// `served_model_names` must be non-empty; the first entry is used as the
45 +/// `model` field in responses.
46 +pub(crate) fn prepare_chat_request(
47 + request: ChatCompletionRequest,
48 + served_model_names: &[String],
49 + ctx: ResolvedRequestContext,
50 +) -> Result<PreparedRequest, ApiError> {
51 + validate::validate_request_compat(&request, served_model_names)?;
52 +
53 + let request_id = format!("chatcmpl-{}", ctx.request_id);
54 + let echo = request
55 + .echo
56 + .then(|| extract_last_assistant_content(&request.messages))
57 + .flatten();
58 + let messages: Vec<_> = request.messages.into_iter().map(convert_message).try_collect()?;
59 + let generation_prompt_mode = normalize_generation_prompt_mode(
60 + request.add_generation_prompt,
61 + request.continue_final_message,
62 + &messages,
63 + )?;
64 +
65 + let template_kwargs = request.chat_template_kwargs.unwrap_or_default();
66 +
67 + let include_usage = (request.stream_options.as_ref())
68 + .and_then(|options| options.include_usage)
69 + .unwrap_or(false);
70 + let requested_logprobs = request.logprobs;
71 +
72 + // Auto-enable prompt logprobs for non-streaming echo, matching Python vLLM's
73 + // behavior.
74 + let top_logprobs = request.top_logprobs.unwrap_or(0);
75 + let prompt_logprobs = request
76 + .prompt_logprobs
77 + .or((request.echo && !request.stream).then_some(top_logprobs));
78 + let include_prompt_logprobs = prompt_logprobs.is_some();
79 +
80 + let structured_outputs = convert_from_response_format(
81 + request.response_format.as_ref(),
82 + &request.structured_outputs,
83 + )?;
84 +
85 + let chat_request = ChatRequest {
86 + request_id: request_id.clone(),
87 + messages,
88 + sampling_params: SamplingParams {
89 + temperature: request.temperature,
90 + top_p: request.top_p,
91 + top_k: request.top_k,
92 + seed: request.seed,
93 + max_tokens: request.max_completion_tokens,
94 + min_tokens: request.min_tokens,
95 + logprobs: request.logprobs.then_some(top_logprobs),
96 + prompt_logprobs,
97 + min_p: request.min_p,
98 + frequency_penalty: request.frequency_penalty,
99 + presence_penalty: request.presence_penalty,
100 + repetition_penalty: request.repetition_penalty,
101 + stop_token_ids: request.stop_token_ids,
102 + ignore_eos: request.ignore_eos,
103 + logit_bias: convert_logit_bias(request.logit_bias)?,
104 + allowed_token_ids: request.allowed_token_ids,
105 + bad_words: request.bad_words,
106 + logprob_token_ids: None,
107 + structured_outputs,
108 + skip_reading_prefix_cache: None,
109 + vllm_xargs: merge_kv_transfer_params(
110 + request.vllm_xargs,
111 + request.kv_transfer_params.as_ref(),
112 + ),
113 + },
114 + chat_options: ChatOptions {
115 + generation_prompt_mode,
116 + chat_template: request.chat_template,
117 + reasoning_effort: request.reasoning_effort,
118 + template_kwargs,
119 + },
120 + tools: convert_tools(request.tools)?,
121 + tool_choice: convert_tool_choice(request.tool_choice.as_ref())?,
122 + decode_options: vllm_text::output::TextDecodeOptions {
123 + skip_special_tokens: request.skip_special_tokens,
124 + include_stop_str_in_output: request.include_stop_str_in_output,
125 + stop_strings: request.stop.map(|stop| stop.into_vec()),
126 + min_tokens: request.min_tokens.unwrap_or(0),
127 + },
128 + intermediate: request.stream,
129 + priority: request.priority.unwrap_or(0),
130 + documents: request.documents,
131 + cache_salt: request.cache_salt,
132 + add_special_tokens: request.add_special_tokens,
133 + data_parallel_rank: ctx.data_parallel_rank,
134 + };
135 +
136 + Ok(PreparedRequest {
137 + request_id,
138 + response_model: served_model_names.first().cloned().unwrap_or_default(),
139 + include_usage,
140 + requested_logprobs,
141 + include_prompt_logprobs,
142 + chat_request,
143 + echo,
144 + return_token_ids: request.return_token_ids.unwrap_or(false),
145 + return_tokens_as_token_ids: request.return_tokens_as_token_ids.unwrap_or(false),
146 + })
147 +}
148 +
149 +fn normalize_generation_prompt_mode(
150 + add_generation_prompt: Option<bool>,
151 + continue_final_message: bool,
152 + messages: &[VllmChatMessage],
153 +) -> Result<GenerationPromptMode, ApiError> {
154 + if add_generation_prompt == Some(true) && continue_final_message {
155 + bail_invalid_request!(
156 + "Cannot set both `continue_final_message` and `add_generation_prompt` to True."
157 + );
158 + }
159 +
160 + let last_role = messages.last().map(VllmChatMessage::role);
161 + match (add_generation_prompt, continue_final_message, last_role) {
162 + (Some(true), true, _) => unreachable!("rejected above"),
163 + (_, true, Some(vllm_chat::ChatRole::Assistant)) => {
164 + Ok(GenerationPromptMode::ContinueFinalAssistant)
165 + }
166 + (_, true, _) => {
167 + bail_invalid_request!(
168 + "Cannot set `continue_final_message` to True when the last message is not from the assistant."
169 + );
170 + }
171 + (Some(false), false, _) => Ok(GenerationPromptMode::NoGenerationPrompt),
172 + (None | Some(true), false, _) => Ok(GenerationPromptMode::StartNewAssistant),
173 + }
174 +}
175 +
176 +/// Extract the text content of the last message if it has the assistant role.
177 +fn extract_last_assistant_content(messages: &[ChatMessage]) -> Option<String> {
178 + let ChatMessage::Assistant { content, .. } = messages.last()? else {
179 + return None;
180 + };
181 + let text = match content.as_ref()? {
182 + MessageContent::Text(text) => text.clone(),
183 + MessageContent::Parts(parts) => parts
184 + .iter()
185 + .filter_map(|p| match p {
186 + ContentPart::Text { text } => Some(text.as_str()),
187 + _ => None,
188 + })
189 + .collect::<Vec<_>>()
190 + .join("\n"),
191 + };
192 + (!text.is_empty()).then_some(text)
193 +}
194 +
195 +/// Lower one OpenAI chat message into the `vllm-chat` message shape.
196 +fn convert_message(message: ChatMessage) -> Result<VllmChatMessage, ApiError> {
197 + match message {
198 + ChatMessage::System { content, .. } => {
199 + Ok(VllmChatMessage::system(convert_content(content)?))
200 + }
201 + ChatMessage::User { content, .. } => Ok(VllmChatMessage::user(convert_content(content)?)),
202 + ChatMessage::Assistant {
203 + content,
204 + tool_calls,
205 + reasoning,
206 + name: _,
207 + } => {
208 + let mut blocks = Vec::new();
209 + if let Some(reasoning) = reasoning
210 + && !reasoning.is_empty()
211 + {
212 + blocks.push(AssistantContentBlock::Reasoning { text: reasoning });
213 + }
214 + if let Some(content) = content {
215 + blocks.extend(convert_assistant_text_blocks(content)?);
216 + }
217 + if let Some(tool_calls) = tool_calls {
218 + blocks.extend(convert_assistant_tool_calls(tool_calls)?);
219 + }
220 + if blocks.is_empty() {
221 + bail_invalid_request!(
222 + "Assistant messages must contain text, reasoning content, or tool_calls."
223 + );
224 + }
225 +
226 + Ok(VllmChatMessage::assistant_blocks(blocks))
227 + }
228 + ChatMessage::Tool {
229 + content,
230 + tool_call_id,
231 + } => Ok(VllmChatMessage::tool_response(
232 + convert_content(content)?,
233 + tool_call_id,
234 + )),
235 + ChatMessage::Function { .. } => {
236 + bail_invalid_request!("Function messages are not supported.")
237 + }
238 + ChatMessage::Developer {
239 + content,
240 + tools,
241 + name: _,
242 + } => Ok(VllmChatMessage::developer(
243 + convert_content(content)?,
244 + convert_message_tools(tools)?,
245 + )),
246 + }
247 +}
248 +
249 +/// Convert the given OpenAI message content value into the internal format in
250 +/// `vllm-chat`.
251 +fn convert_content(content: MessageContent) -> Result<ChatContent, ApiError> {
252 + match content {
253 + MessageContent::Text(text) => Ok(ChatContent::Text(text)),
254 + MessageContent::Parts(parts) => parts
255 + .into_iter()
256 + .map(|part| match part {
257 + ContentPart::Text { text } => Ok(ChatContentPart::text(text)),
258 + ContentPart::ImageUrl { image_url, uuid } => Ok(ChatContentPart::ImageUrl {
259 + image_url: image_url.url,
260 + detail: image_url.detail,
261 + uuid,
262 + }),
263 + _ => bail_invalid_request!("Only text and image_url content parts are supported."),
264 + })
265 + .try_collect()
266 + .map(ChatContent::Parts),
267 + }
268 +}
269 +
270 +/// Convert the given OpenAI assistant message content into the internal format
271 +/// in `vllm-chat`.
272 +fn convert_assistant_text_blocks(
273 + content: MessageContent,
274 +) -> Result<Vec<AssistantContentBlock>, ApiError> {
275 + match content {
276 + MessageContent::Text(text) => Ok(vec![AssistantContentBlock::Text { text }]),
277 + MessageContent::Parts(parts) => parts
278 + .into_iter()
279 + .map(|part| match part {
280 + ContentPart::Text { text } => Ok(AssistantContentBlock::Text { text }),
281 + _ => bail_invalid_request!(
282 + "Only text content parts are supported for assistant messages."
283 + ),
284 + })
285 + .try_collect(),
286 + }
287 +}
288 +
289 +fn convert_assistant_tool_calls(
290 + tool_calls: Vec<crate::routes::openai::utils::types::ToolCall>,
291 +) -> Result<Vec<AssistantContentBlock>, ApiError> {
292 + tool_calls
293 + .into_iter()
294 + .map(|tool_call| {
295 + if tool_call.tool_type != "function" {
296 + bail_invalid_request!("Only function tool calls are supported.");
297 + }
298 +
299 + Ok(AssistantContentBlock::ToolCall(AssistantToolCall {
300 + id: tool_call.id,
301 + name: tool_call.function.name,
302 + arguments: tool_call.function.arguments.unwrap_or_else(|| "{}".to_string()),
303 + }))
304 + })
305 + .collect()
306 +}
307 +
308 +fn convert_tools(tools: Option<Vec<Tool>>) -> Result<Vec<ChatTool>, ApiError> {
309 + tools
310 + .unwrap_or_default()
311 + .into_iter()
312 + .map(|tool| {
313 + if tool.tool_type != "function" {
314 + bail_invalid_request!("Only function tools are supported.");
315 + }
316 + Ok(ChatTool {
317 + name: tool.function.name,
318 + description: tool.function.description,
319 + parameters: tool.function.parameters,
320 + strict: tool.function.strict,
321 + })
322 + })
323 + .collect()
324 +}
325 +
326 +fn convert_message_tools(tools: Option<Vec<Tool>>) -> Result<Option<Vec<ChatTool>>, ApiError> {
327 + let tools = convert_tools(tools)?;
328 + Ok((!tools.is_empty()).then_some(tools))
329 +}
330 +
331 +fn convert_tool_choice(tool_choice: Option<&ToolChoice>) -> Result<ChatToolChoice, ApiError> {
332 + match tool_choice {
333 + None | Some(ToolChoice::Value(ToolChoiceValue::Auto)) => Ok(ChatToolChoice::Auto),
334 + Some(ToolChoice::Value(ToolChoiceValue::None)) => Ok(ChatToolChoice::None),
335 + _ => bail_invalid_request!("tool_choice={:?} is not supported yet.", tool_choice),
336 + }
337 +}
338 +
339 +#[cfg(test)]
340 +mod tests {
341 + use std::collections::HashMap;
342 +
343 + use axum::http::HeaderMap;
344 + use expect_test::expect;
345 + use llm_multimodal::ImageDetail;
346 + use serde_json::json;
347 + use vllm_chat::{
348 + AssistantContentBlock, AssistantToolCall, ChatContentPart, ChatMessage as VllmChatMessage,
349 + ChatTool as VllmChatTool, ChatToolChoice, GenerationPromptMode,
350 + SamplingParams as VllmSamplingParams,
351 + };
352 + use vllm_text::output::TextDecodeOptions;
353 +
354 + use super::prepare_chat_request;
355 + use crate::routes::openai::chat_completions::types::{
356 + AssistantRole, ChatCompletionMessage, ChatCompletionRequest,
357 + };
358 + use crate::routes::openai::utils::types::{
359 + ChatMessage, ContentPart, Function, FunctionCallResponse, ImageUrl, MessageContent, Tool,
360 + ToolCall, ToolChoice, ToolChoiceValue, VideoUrl,
361 + };
362 + use crate::utils::{ResolvedRequestContext, resolve_request_context};
363 +
364 + fn request_context(headers: &HeaderMap, request_id: Option<&str>) -> ResolvedRequestContext {
365 + resolve_request_context(headers, request_id)
366 + }
367 +
368 + fn served(names: &[&str]) -> Vec<String> {
369 + names.iter().map(|s| s.to_string()).collect()
370 + }
371 +
372 + fn base_request() -> ChatCompletionRequest {
373 + ChatCompletionRequest {
374 + model: "Qwen/Qwen1.5-0.5B-Chat".to_string(),
375 + messages: vec![ChatMessage::User {
376 + content: MessageContent::Text("hello".to_string()),
377 + name: None,
378 + }],
379 + stream: true,
380 + ..Default::default()
381 + }
382 + }
383 +
384 + #[test]
385 + fn prepare_chat_request_maps_text_parts() {
386 + let mut request = base_request();
387 + request.messages = vec![ChatMessage::Assistant {
388 + content: Some(MessageContent::Parts(vec![ContentPart::Text {
389 + text: "hello".to_string(),
390 + }])),
391 + name: None,
392 + tool_calls: None,
393 + reasoning: None,
394 + }];
395 + request.add_generation_prompt = Some(false);
396 + request.continue_final_message = true;
397 + request.skip_special_tokens = false;
398 + request.chat_template_kwargs = Some(HashMap::from([("foo".to_string(), json!("bar"))]));
399 +
400 + let prepared = prepare_chat_request(
401 + request,
402 + &served(&["Qwen/Qwen1.5-0.5B-Chat"]),
403 + ResolvedRequestContext::default(),
404 + )
405 + .expect("request is valid");
406 +
407 + assert!(prepared.request_id.starts_with("chatcmpl-"));
408 + assert_eq!(
409 + prepared.chat_request.messages,
410 + vec![VllmChatMessage::assistant_text("hello")]
411 + );
412 + assert_eq!(
413 + prepared.chat_request.sampling_params,
414 + VllmSamplingParams::default()
415 + );
416 + assert_eq!(
417 + prepared.chat_request.chat_options.generation_prompt_mode,
418 + GenerationPromptMode::ContinueFinalAssistant
419 + );
420 + assert_eq!(
421 + prepared.chat_request.chat_options.template_kwargs,
422 + HashMap::from([("foo".to_string(), json!("bar"))])
423 + );
424 + assert_eq!(
425 + prepared.chat_request.decode_options,
426 + TextDecodeOptions {
427 + skip_special_tokens: false,
428 + include_stop_str_in_output: false,
429 + stop_strings: None,
430 + min_tokens: 0,
431 + }
432 + );
433 + assert!(prepared.chat_request.tools.is_empty());
434 + assert_eq!(prepared.chat_request.tool_choice, ChatToolChoice::Auto);
435 + }
436 +
437 + #[test]
438 + fn prepare_chat_request_keeps_optional_sampling_fields_unset() {
439 + let prepared = prepare_chat_request(
440 + base_request(),
441 + &served(&["Qwen/Qwen1.5-0.5B-Chat"]),
442 + ResolvedRequestContext::default(),
443 + )
444 + .expect("request is valid");
445 +
446 + assert!(prepared.request_id.starts_with("chatcmpl-"));
447 + assert_eq!(
448 + prepared.chat_request.messages,
449 + vec![VllmChatMessage::user("hello")]
450 + );
451 + assert_eq!(
452 + prepared.chat_request.sampling_params,
453 + VllmSamplingParams::default()
454 + );
455 + assert_eq!(
456 + prepared.chat_request.chat_options.generation_prompt_mode,
457 + GenerationPromptMode::StartNewAssistant
458 + );
459 + assert_eq!(
460 + prepared.chat_request.decode_options,
461 + TextDecodeOptions {
462 + skip_special_tokens: true,
463 + include_stop_str_in_output: false,
464 + stop_strings: None,
465 + min_tokens: 0,
466 + }
467 + );
468 + assert!(prepared.chat_request.tools.is_empty());
469 + assert_eq!(prepared.chat_request.tool_choice, ChatToolChoice::Auto);
470 + }
471 +
472 + #[test]
473 + fn prepare_chat_request_preserves_sampling_passthrough_fields() {
474 + let request = ChatCompletionRequest {
475 + seed: Some(42),
476 + min_p: Some(0.2),
477 + frequency_penalty: Some(0.3),
478 + presence_penalty: Some(0.4),
479 + repetition_penalty: Some(1.1),
480 + ..base_request()
481 + };
482 +
483 + let prepared = prepare_chat_request(
484 + request,
485 + &served(&["Qwen/Qwen1.5-0.5B-Chat"]),
486 + ResolvedRequestContext::default(),
487 + )
488 + .expect("request is valid");
489 + let expected = VllmSamplingParams {
490 + seed: Some(42),
491 + min_p: Some(0.2),
492 + frequency_penalty: Some(0.3),
493 + presence_penalty: Some(0.4),
494 + repetition_penalty: Some(1.1),
495 + ..VllmSamplingParams::default()
496 + };
497 + assert_eq!(prepared.chat_request.sampling_params, expected);
498 + }
499 +
500 + #[test]
501 + fn prepare_chat_request_accepts_developer_messages() {
502 + let request = ChatCompletionRequest {
503 + messages: vec![ChatMessage::Developer {
504 + content: MessageContent::Text("hello".to_string()),
505 + tools: Some(vec![Tool {
506 + tool_type: "function".to_string(),
507 + function: Function {
508 + name: "get_weather".to_string(),
509 + description: Some("Get weather".to_string()),
510 + parameters: json!({
511 + "type": "object",
512 + "properties": {"city": {"type": "string"}},
513 + }),
514 + strict: Some(true),
515 + },
516 + }]),
517 + name: None,
518 + }],
519 + ..base_request()
520 + };
521 +
522 + let prepared = prepare_chat_request(
523 + request,
524 + &served(&["Qwen/Qwen1.5-0.5B-Chat"]),
525 + ResolvedRequestContext::default(),
526 + )
527 + .expect("request is valid");
528 +
529 + assert_eq!(
530 + prepared.chat_request.messages,
531 + vec![VllmChatMessage::developer(
532 + "hello",
533 + Some(vec![VllmChatTool {
534 + name: "get_weather".to_string(),
535 + description: Some("Get weather".to_string()),
536 + parameters: json!({
537 + "type": "object",
538 + "properties": {"city": {"type": "string"}},
539 + }),
540 + strict: Some(true),
541 + }]),
542 + )]
543 + );
544 + }
545 +
546 + #[test]
547 + fn prepare_chat_request_maps_image_url_content_parts() {
548 + let request = ChatCompletionRequest {
549 + messages: vec![ChatMessage::User {
550 + content: MessageContent::Parts(vec![
551 + ContentPart::Text {
552 + text: "describe ".to_string(),
553 + },
554 + ContentPart::ImageUrl {
555 + image_url: ImageUrl {
556 + url: "https://example.com/image.png".to_string(),
557 + detail: Some(ImageDetail::Low),
558 + },
559 + uuid: Some("image-1".to_string()),
560 + },
561 + ContentPart::Text {
562 + text: " briefly".to_string(),
563 + },
564 + ]),
565 + name: None,
566 + }],
567 + ..base_request()
568 + };
569 +
570 + let prepared = prepare_chat_request(
571 + request,
572 + &served(&["Qwen/Qwen1.5-0.5B-Chat"]),
573 + ResolvedRequestContext::default(),
574 + )
575 + .expect("request is valid");
576 +
577 + assert_eq!(
578 + prepared.chat_request.messages,
579 + vec![VllmChatMessage::user(vec![
580 + ChatContentPart::text("describe "),
581 + ChatContentPart::ImageUrl {
582 + image_url: "https://example.com/image.png".to_string(),
583 + detail: Some(ImageDetail::Low),
584 + uuid: Some("image-1".to_string()),
585 + },
586 + ChatContentPart::text(" briefly"),
587 + ])]
588 + );
589 + }
590 +
591 + #[test]
592 + fn prepare_chat_request_maps_developer_image_url_content_parts() {
593 + let request = ChatCompletionRequest {
594 + messages: vec![ChatMessage::Developer {
595 + content: MessageContent::Parts(vec![ContentPart::ImageUrl {
596 + image_url: ImageUrl {
597 + url: "https://example.com/image.png".to_string(),
598 + detail: None,
599 + },
600 + uuid: None,
601 + }]),
602 + tools: None,
603 + name: None,
604 + }],
605 + ..base_request()
606 + };
607 +
608 + let prepared = prepare_chat_request(
609 + request,
610 + &served(&["Qwen/Qwen1.5-0.5B-Chat"]),
611 + ResolvedRequestContext::default(),
612 + )
613 + .expect("request is valid");
614 +
615 + assert_eq!(
616 + prepared.chat_request.messages,
617 + vec![VllmChatMessage::developer(
618 + vec![ChatContentPart::image_url("https://example.com/image.png")],
619 + None,
620 + )]
621 + );
622 + }
623 +
624 + #[test]
625 + fn prepare_chat_request_rejects_video_content_parts() {
626 + let request = ChatCompletionRequest {
627 + messages: vec![ChatMessage::User {
628 + content: MessageContent::Parts(vec![ContentPart::VideoUrl {
629 + video_url: VideoUrl {
630 + url: "https://example.com/video.mp4".to_string(),
631 + },
632 + }]),
633 + name: None,
634 + }],
635 + ..base_request()
636 + };
637 +
638 + let error = prepare_chat_request(
639 + request,
640 + &served(&["Qwen/Qwen1.5-0.5B-Chat"]),
641 + ResolvedRequestContext::default(),
642 + )
643 + .unwrap_err();
644 +
645 + expect!["Only text and image_url content parts are supported."]
646 + .assert_eq(&error.to_error_response().error.message);
647 + }
648 +
649 + #[test]
650 + fn prepare_chat_request_rejects_assistant_image_url_content_parts() {
651 + let request = ChatCompletionRequest {
652 + messages: vec![ChatMessage::Assistant {
653 + content: Some(MessageContent::Parts(vec![ContentPart::ImageUrl {
654 + image_url: ImageUrl {
655 + url: "https://example.com/image.png".to_string(),
656 + detail: None,
657 + },
658 + uuid: None,
659 + }])),
660 + name: None,
661 + tool_calls: None,
662 + reasoning: None,
663 + }],
664 + ..base_request()
665 + };
666 +
667 + let error = prepare_chat_request(
668 + request,
669 + &served(&["Qwen/Qwen1.5-0.5B-Chat"]),
670 + ResolvedRequestContext::default(),
671 + )
672 + .unwrap_err();
673 +
674 + expect!["Only text content parts are supported for assistant messages."]
675 + .assert_eq(&error.to_error_response().error.message);
676 + }
677 +
678 + #[test]
679 + fn prepare_chat_request_accepts_assistant_reasoning_history() {
680 + let message = ChatCompletionMessage {
681 + role: AssistantRole,
682 + content: Some("answer".to_string()),
683 + tool_calls: None,
684 + reasoning: Some("inner".to_string()),
685 + };
686 + let message_json = serde_json::to_value(message).expect("message serializes");
687 +
688 + let request = ChatCompletionRequest {
689 + messages: vec![
690 + serde_json::from_value(message_json).expect("response message is valid history"),
691 + ],
692 + add_generation_prompt: Some(false),
693 + ..base_request()
694 + };
695 +
696 + let prepared = prepare_chat_request(
697 + request,
698 + &served(&["Qwen/Qwen1.5-0.5B-Chat"]),
699 + ResolvedRequestContext::default(),
700 + )
701 + .expect("request is valid");
702 + assert_eq!(
703 + prepared.chat_request.messages,
704 + vec![VllmChatMessage::assistant_blocks(vec![
705 + AssistantContentBlock::Reasoning {
706 + text: "inner".to_string(),
707 + },
708 + AssistantContentBlock::Text {
709 + text: "answer".to_string(),
710 + },
711 + ])]
712 + );
713 + assert!(prepared.chat_request.tools.is_empty());
714 + assert_eq!(prepared.chat_request.tool_choice, ChatToolChoice::Auto);
715 + }
716 +
717 + #[test]
718 + fn prepare_chat_request_accepts_legacy_reasoning_content_alias() {
719 + let request = ChatCompletionRequest {
720 + messages: vec![
721 + serde_json::from_value(json!({
722 + "role": "assistant",
723 + "content": "answer",
724 + "reasoning_content": "inner",
725 + }))
726 + .expect("legacy reasoning_content alias is accepted"),
727 + ],
728 + add_generation_prompt: Some(false),
729 + ..base_request()
730 + };
731 +
732 + let prepared = prepare_chat_request(
733 + request,
734 + &served(&["Qwen/Qwen1.5-0.5B-Chat"]),
735 + ResolvedRequestContext::default(),
736 + )
737 + .expect("request is valid");
738 + assert_eq!(
739 + prepared.chat_request.messages,
740 + vec![VllmChatMessage::assistant_blocks(vec![
741 + AssistantContentBlock::Reasoning {
742 + text: "inner".to_string(),
743 + },
744 + AssistantContentBlock::Text {
745 + text: "answer".to_string(),
746 + },
747 + ])]
748 + );
749 + }
750 +
751 + #[test]
752 + fn prepare_chat_request_accepts_tools_and_tool_history() {
753 + let request = ChatCompletionRequest {
754 + messages: vec![
755 + ChatMessage::Assistant {
756 + content: None,
757 + name: None,
758 + tool_calls: Some(vec![ToolCall {
759 + id: "call_1".to_string(),
760 + tool_type: "function".to_string(),
761 + function: FunctionCallResponse {
762 + name: "get_weather".to_string(),
763 + arguments: Some(r#"{"city":"Paris"}"#.to_string()),
764 + },
765 + }]),
766 + reasoning: None,
767 + },
768 + ChatMessage::Tool {
769 + content: MessageContent::Text("Sunny".to_string()),
770 + tool_call_id: "call_1".to_string(),
771 + },
772 + ],
773 + tools: Some(vec![Tool {
774 + tool_type: "function".to_string(),
775 + function: Function {
776 + name: "get_weather".to_string(),
777 + description: Some("Get weather".to_string()),
778 + parameters: json!({
779 + "type": "object",
780 + "properties": {"city": {"type": "string"}},
781 + }),
782 + strict: None,
783 + },
784 + }]),
785 + tool_choice: Some(ToolChoice::Value(ToolChoiceValue::None)),
786 + ..base_request()
787 + };
788 +
789 + let prepared = prepare_chat_request(
790 + request,
791 + &served(&["Qwen/Qwen1.5-0.5B-Chat"]),
792 + ResolvedRequestContext::default(),
793 + )
794 + .expect("request is valid");
795 + assert_eq!(
796 + prepared.chat_request.messages,
797 + vec![
798 + VllmChatMessage::assistant_blocks(vec![AssistantContentBlock::ToolCall(
799 + AssistantToolCall {
800 + id: "call_1".to_string(),
801 + name: "get_weather".to_string(),
802 + arguments: r#"{"city":"Paris"}"#.to_string(),
803 + },
804 + )]),
805 + VllmChatMessage::tool_response("Sunny", "call_1"),
806 + ]
807 + );
808 + assert_eq!(
809 + prepared.chat_request.tools,
810 + vec![VllmChatTool {
811 + name: "get_weather".to_string(),
812 + description: Some("Get weather".to_string()),
813 + parameters: json!({
814 + "type": "object",
815 + "properties": {"city": {"type": "string"}},
816 + }),
817 + strict: None,
818 + }]
819 + );
820 + assert_eq!(prepared.chat_request.tool_choice, ChatToolChoice::None);
821 + }
822 +
823 + #[test]
824 + fn prepare_chat_request_lowers_logprobs_fields() {
825 + let request = ChatCompletionRequest {
826 + stream: false,
827 + logprobs: true,
828 + prompt_logprobs: Some(2),
829 + ..base_request()
830 + };
831 +
832 + let prepared = prepare_chat_request(
833 + request,
834 + &served(&["Qwen/Qwen1.5-0.5B-Chat"]),
835 + ResolvedRequestContext::default(),
836 + )
837 + .expect("request is valid");
838 +
839 + assert!(prepared.requested_logprobs);
840 + assert!(prepared.include_prompt_logprobs);
841 + assert_eq!(prepared.chat_request.sampling_params.logprobs, Some(0));
842 + assert_eq!(
843 + prepared.chat_request.sampling_params.prompt_logprobs,
844 + Some(2)
845 + );
846 + }
847 +
848 + #[test]
849 + fn prepare_chat_request_keeps_prompt_logprobs_independent_from_echo() {
850 + let request = ChatCompletionRequest {
851 + logprobs: true,
852 + top_logprobs: Some(3),
853 + echo: true,
854 + ..base_request()
855 + };
856 +
857 + let prepared = prepare_chat_request(
858 + request,
859 + &served(&["Qwen/Qwen1.5-0.5B-Chat"]),
860 + ResolvedRequestContext::default(),
861 + )
862 + .expect("request is valid");
863 +
864 + assert_eq!(prepared.chat_request.sampling_params.logprobs, Some(3));
865 + assert_eq!(prepared.chat_request.sampling_params.prompt_logprobs, None);
866 + assert!(!prepared.include_prompt_logprobs);
867 + }
868 +
869 + #[test]
870 + fn prepare_chat_request_threads_data_parallel_rank() {
871 + let mut headers = HeaderMap::new();
872 + headers.insert("X-data-parallel-rank", "7".parse().unwrap());
873 + let prepared = prepare_chat_request(
874 + base_request(),
875 + &served(&["Qwen/Qwen1.5-0.5B-Chat"]),
876 + request_context(&headers, None),
877 + )
878 + .expect("request is valid");
879 + assert_eq!(prepared.chat_request.data_parallel_rank, Some(7));
880 + }
881 +
882 + #[test]
883 + fn prepare_chat_request_leaves_data_parallel_rank_none_when_absent() {
884 + let prepared = prepare_chat_request(
885 + base_request(),
886 + &served(&["Qwen/Qwen1.5-0.5B-Chat"]),
887 + ResolvedRequestContext::default(),
888 + )
889 + .expect("request is valid");
890 + assert_eq!(prepared.chat_request.data_parallel_rank, None);
891 + }
892 +
893 + #[test]
894 + fn prepare_chat_request_maps_no_generation_prompt_mode() {
895 + let mut request = base_request();
896 + request.add_generation_prompt = Some(false);
897 +
898 + let prepared = prepare_chat_request(
899 + request,
900 + &served(&["Qwen/Qwen1.5-0.5B-Chat"]),
901 + ResolvedRequestContext::default(),
902 + )
903 + .expect("request is valid");
904 +
905 + assert_eq!(
906 + prepared.chat_request.chat_options.generation_prompt_mode,
907 + GenerationPromptMode::NoGenerationPrompt
908 + );
909 + }
910 +
911 + #[test]
912 + fn prepare_chat_request_rejects_conflicting_explicit_generation_prompt_flags() {
913 + let mut request = base_request();
914 + request.add_generation_prompt = Some(true);
915 + request.continue_final_message = true;
916 +
917 + let error = prepare_chat_request(
918 + request,
919 + &served(&["Qwen/Qwen1.5-0.5B-Chat"]),
920 + ResolvedRequestContext::default(),
921 + )
922 + .unwrap_err();
923 +
924 + expect!["Cannot set both `continue_final_message` and `add_generation_prompt` to True."]
925 + .assert_eq(&error.to_error_response().error.message);
926 + }
927 +
928 + #[test]
929 + fn prepare_chat_request_accepts_continue_final_message_with_implicit_add_generation_prompt() {
930 + let mut request = base_request();
931 + request.messages = vec![ChatMessage::Assistant {
932 + content: Some(MessageContent::Text("hello".to_string())),
933 + name: None,
934 + tool_calls: None,
935 + reasoning: None,
936 + }];
937 + request.continue_final_message = true;
938 +
939 + let prepared = prepare_chat_request(
940 + request,
941 + &served(&["Qwen/Qwen1.5-0.5B-Chat"]),
942 + ResolvedRequestContext::default(),
943 + )
944 + .expect("request is valid");
945 +
946 + assert_eq!(
947 + prepared.chat_request.chat_options.generation_prompt_mode,
948 + GenerationPromptMode::ContinueFinalAssistant
949 + );
950 + }
951 +
952 + #[test]
953 + fn prepare_chat_request_rejects_continue_final_message_without_final_assistant() {
954 + let mut request = base_request();
955 + request.continue_final_message = true;
956 +
957 + let error = prepare_chat_request(
958 + request,
959 + &served(&["Qwen/Qwen1.5-0.5B-Chat"]),
960 + ResolvedRequestContext::default(),
961 + )
962 + .unwrap_err();
963 +
964 + expect!["Cannot set `continue_final_message` to True when the last message is not from the assistant."]
965 + .assert_eq(&error.to_error_response().error.message);
966 + }
967 +
968 + #[test]
969 + fn prepare_chat_request_allows_new_assistant_mode_after_final_assistant() {
970 + let request = ChatCompletionRequest {
971 + messages: vec![ChatMessage::Assistant {
972 + content: Some(MessageContent::Text("hello".to_string())),
973 + name: None,
974 + tool_calls: None,
975 + reasoning: None,
976 + }],
977 + ..base_request()
978 + };
979 +
980 + let prepared = prepare_chat_request(
981 + request,
982 + &served(&["Qwen/Qwen1.5-0.5B-Chat"]),
983 + ResolvedRequestContext::default(),
984 + )
985 + .expect("request is valid");
986 +
987 + assert_eq!(
988 + prepared.chat_request.chat_options.generation_prompt_mode,
989 + GenerationPromptMode::StartNewAssistant
990 + );
991 + }
992 +}
rust/src/server/src/routes/openai/chat_completions/types.rs new
+600
@@ -0,0 +1,600 @@
1 +use std::collections::HashMap;
2 +use std::fmt;
3 +
4 +use serde::{Deserialize, Serialize};
5 +use serde_json::Value;
6 +use serde_with::SerializeDisplay;
7 +use validator::Validate;
8 +use vllm_chat::ReasoningEffort;
9 +
10 +use crate::routes::openai::utils::structured_outputs::ResponseFormat;
11 +use crate::routes::openai::utils::types::{
12 + ChatLogProbs, ChatMessage, MessageContent, Normalizable, StreamOptions, StringOrArray, Tool,
13 + ToolCall, ToolCallDelta, ToolChoice, ToolChoiceValue, ToolReference, UNKNOWN_MODEL_ID, Usage,
14 + default_true, validate_stop, validate_top_p_value,
15 +};
16 +
17 +/// vLLM-compatible request type for the Chat Completions API.
18 +///
19 +/// Mirrors the Python vLLM `ChatCompletionRequest` class. The local copy keeps
20 +/// the request type route-owned so we can add vLLM-only fields directly instead
21 +/// of layering wrapper deserializers on top.
22 +#[serde_with::skip_serializing_none]
23 +#[derive(Debug, Clone, Deserialize, Serialize, Validate)]
24 +#[validate(schema(function = "validate_chat_cross_parameters"))]
25 +pub struct ChatCompletionRequest {
26 + // -------- Standard OpenAI API Parameters --------
27 + /// A list of messages comprising the conversation so far
28 + #[validate(custom(function = "validate_messages"))]
29 + pub messages: Vec<ChatMessage>,
30 +
31 + /// ID of the model to use
32 + #[serde(default = "default_model")]
33 + pub model: String,
34 +
35 + /// Number between -2.0 and 2.0. Positive values penalize new tokens based
36 + /// on their existing frequency in the text so far
37 + #[validate(range(min = -2.0, max = 2.0))]
38 + pub frequency_penalty: Option<f32>,
39 +
40 + /// Modify the likelihood of specified tokens appearing in the completion
41 + pub logit_bias: Option<HashMap<String, f32>>,
42 +
43 + /// Whether to return log probabilities of the output tokens
44 + #[serde(default)]
45 + pub logprobs: bool,
46 +
47 + /// An integer specifying the number of most likely tokens to return
48 + /// -1 means return all
49 + #[validate(range(min = -1))]
50 + pub top_logprobs: Option<i32>,
51 +
52 + /// Deprecated: Replaced by max_completion_tokens
53 + #[deprecated(note = "Use max_completion_tokens instead")]
54 + #[validate(range(min = 1))]
55 + pub max_tokens: Option<u32>,
56 +
57 + /// An upper bound for the number of tokens that can be generated for a
58 + /// completion
59 + #[validate(range(min = 1))]
60 + pub max_completion_tokens: Option<u32>,
61 +
62 + /// How many chat completion choices to generate for each input message
63 + #[validate(range(min = 1, max = 10))]
64 + pub n: Option<u32>,
65 +
66 + /// Number between -2.0 and 2.0. Positive values penalize new tokens based
67 + /// on whether they appear in the text so far
68 + #[validate(range(min = -2.0, max = 2.0))]
69 + pub presence_penalty: Option<f32>,
70 +
71 + /// An object specifying the format that the model must output
72 + pub response_format: Option<ResponseFormat>,
73 +
74 + /// If specified, our system will make a best effort to sample
75 + /// deterministically
76 + pub seed: Option<i64>,
77 +
78 + /// Up to 4 sequences where the API will stop generating further tokens
79 + #[validate(custom(function = "validate_stop"))]
80 + pub stop: Option<StringOrArray>,
81 +
82 + /// If set, partial message deltas will be sent
83 + #[serde(default)]
84 + pub stream: bool,
85 +
86 + /// Options for streaming response
87 + pub stream_options: Option<StreamOptions>,
88 +
89 + /// What sampling temperature to use, between 0 and 2
90 + #[validate(range(min = 0.0, max = 2.0))]
91 + pub temperature: Option<f32>,
92 +
93 + /// An alternative to sampling with temperature
94 + #[validate(custom(function = "validate_top_p_value"))]
95 + pub top_p: Option<f32>,
96 +
97 + /// A list of tools the model may call
98 + pub tools: Option<Vec<Tool>>,
99 +
100 + /// Controls which (if any) tool is called by the model
101 + pub tool_choice: Option<ToolChoice>,
102 +
103 + /// Effort level for reasoning models (none, minimal, low, medium, high,
104 + /// xhigh, max)
105 + pub reasoning_effort: Option<ReasoningEffort>,
106 +
107 + /// Whether to enable parallel function calling during tool use
108 + pub parallel_tool_calls: Option<bool>,
109 +
110 + /// A unique identifier representing your end-user
111 + pub user: Option<String>,
112 +
113 + // -------- vLLM Sampling Parameters --------
114 + /// Use beam search instead of sampling
115 + #[serde(default)]
116 + pub use_beam_search: bool,
117 +
118 + /// Top-k sampling parameter
119 + pub top_k: Option<u32>,
120 +
121 + /// Min-p nucleus sampling parameter
122 + #[validate(range(min = 0.0, max = 1.0))]
123 + pub min_p: Option<f32>,
124 +
125 + /// Repetition penalty for reducing repetitive text
126 + #[validate(range(min = 0.0, max = 2.0))]
127 + pub repetition_penalty: Option<f32>,
128 +
129 + /// Length penalty for beam search
130 + pub length_penalty: Option<f32>,
131 +
132 + /// Specific token IDs to use as stop conditions
133 + pub stop_token_ids: Option<Vec<u32>>,
134 +
135 + /// Include stop string in output
136 + #[serde(default)]
137 + pub include_stop_str_in_output: bool,
138 +
139 + /// Ignore end-of-sequence tokens during generation
140 + #[serde(default)]
141 + pub ignore_eos: bool,
142 +
143 + /// Minimum number of tokens to generate
144 + #[validate(range(min = 1))]
145 + pub min_tokens: Option<u32>,
146 +
147 + /// Skip special tokens during detokenization
148 + #[serde(default = "default_true")]
149 + pub skip_special_tokens: bool,
150 +
151 + /// Add spaces between special tokens during detokenization
152 + #[serde(default = "default_true")]
153 + pub spaces_between_special_tokens: bool,
154 +
155 + /// Truncate prompt tokens to this length
156 + pub truncate_prompt_tokens: Option<i64>,
157 +
158 + /// Number of prompt logprobs to return
159 + pub prompt_logprobs: Option<i32>,
160 +
161 + /// Restrict output to these token IDs only
162 + pub allowed_token_ids: Option<Vec<u32>>,
163 +
164 + /// List of bad words to avoid during generation
165 + pub bad_words: Option<Vec<String>>,
166 +
167 + // -------- Extra vLLM Parameters --------
168 + /// Token budget for reasoning/thinking
169 + pub thinking_token_budget: Option<u32>,
170 +
171 + /// Whether to include reasoning content in the response
172 + #[serde(default = "default_true")]
173 + pub include_reasoning: bool,
174 +
175 + /// If true, the new message will be prepended with the last message if they
176 + /// belong to the same role.
177 + #[serde(default)]
178 + pub echo: bool,
179 +
180 + /// Whether to add the generation prompt to the chat template.
181 + ///
182 + /// When omitted, the request follows the API default behavior, which is
183 + /// equivalent to `true` unless `continue_final_message=true` selects
184 + /// final assistant continuation instead.
185 + pub add_generation_prompt: Option<bool>,
186 +
187 + /// Continue generating from final assistant message
188 + #[serde(default)]
189 + pub continue_final_message: bool,
190 +
191 + /// Whether to add special tokens (e.g. BOS) to the prompt
192 + #[serde(default)]
193 + pub add_special_tokens: bool,
194 +
195 + /// Documents for RAG (retrieval-augmented generation)
196 + pub documents: Option<Vec<Value>>,
197 +
198 + /// Jinja chat template override
199 + pub chat_template: Option<String>,
200 +
201 + /// Additional keyword args passed to the chat template renderer
202 + pub chat_template_kwargs: Option<HashMap<String, Value>>,
203 +
204 + /// Additional kwargs for media IO connectors, keyed by modality
205 + pub media_io_kwargs: Option<HashMap<String, Value>>,
206 +
207 + /// Additional kwargs for the HF processor
208 + pub mm_processor_kwargs: Option<HashMap<String, Value>>,
209 +
210 + /// Additional kwargs for structured outputs
211 + pub structured_outputs: Option<Value>,
212 +
213 + /// Request scheduling priority (lower means earlier; default 0)
214 + pub priority: Option<i32>,
215 +
216 + /// External request ID used for response correlation.
217 + pub request_id: Option<String>,
218 +
219 + /// Tokens represented as strings of the form 'token_id:{token_id}' in
220 + /// logprobs
221 + pub return_tokens_as_token_ids: Option<bool>,
222 +
223 + /// Include token IDs alongside generated text
224 + pub return_token_ids: Option<bool>,
225 +
226 + /// Salt for prefix cache isolation in multi-user environments
227 + pub cache_salt: Option<String>,
228 +
229 + /// KV transfer parameters for disaggregated serving
230 + pub kv_transfer_params: Option<HashMap<String, Value>>,
231 +
232 + /// Additional request parameters with string or numeric values for custom
233 + /// extensions
234 + pub vllm_xargs: Option<HashMap<String, Value>>,
235 +
236 + /// Parameters for detecting repetitive N-gram patterns in output tokens
237 + pub repetition_detection: Option<Value>,
238 +}
239 +
240 +impl Default for ChatCompletionRequest {
241 + #[expect(deprecated)]
242 + fn default() -> Self {
243 + Self {
244 + messages: Vec::new(),
245 + model: default_model(),
246 + frequency_penalty: None,
247 + logit_bias: None,
248 + logprobs: false,
249 + top_logprobs: None,
250 + max_tokens: None,
251 + max_completion_tokens: None,
252 + n: None,
253 + presence_penalty: None,
254 + response_format: None,
255 + seed: None,
256 + stop: None,
257 + stream: false,
258 + stream_options: None,
259 + temperature: None,
260 + top_p: None,
261 + tools: None,
262 + tool_choice: None,
263 + reasoning_effort: None,
264 + thinking_token_budget: None,
265 + include_reasoning: true,
266 + parallel_tool_calls: None,
267 + user: None,
268 + use_beam_search: false,
269 + top_k: None,
270 + min_p: None,
271 + repetition_penalty: None,
272 + length_penalty: None,
273 + stop_token_ids: None,
274 + include_stop_str_in_output: false,
275 + ignore_eos: false,
276 + min_tokens: None,
277 + skip_special_tokens: true,
278 + spaces_between_special_tokens: true,
279 + truncate_prompt_tokens: None,
280 + prompt_logprobs: None,
281 + allowed_token_ids: None,
282 + bad_words: None,
283 + echo: false,
284 + add_generation_prompt: None,
285 + continue_final_message: false,
286 + add_special_tokens: false,
287 + documents: None,
288 + chat_template: None,
289 + chat_template_kwargs: None,
290 + media_io_kwargs: None,
291 + mm_processor_kwargs: None,
292 + structured_outputs: None,
293 + priority: None,
294 + request_id: None,
295 + return_tokens_as_token_ids: None,
296 + return_token_ids: None,
297 + cache_salt: None,
298 + kv_transfer_params: None,
299 + vllm_xargs: None,
300 + repetition_detection: None,
301 + }
302 + }
303 +}
304 +
305 +impl Normalizable for ChatCompletionRequest {
306 + /// Normalize the request by applying migrations and defaults.
307 + fn normalize(&mut self) {
308 + // Migrate deprecated max_tokens → max_completion_tokens
309 + #[expect(deprecated)]
310 + if self.max_completion_tokens.is_none() && self.max_tokens.is_some() {
311 + self.max_completion_tokens = self.max_tokens;
312 + self.max_tokens = None;
313 + }
314 +
315 + // Apply tool_choice defaults
316 + // If tools is None, leave tool_choice as None (don't set it)
317 + if self.tool_choice.is_none()
318 + && let Some(tools) = &self.tools
319 + {
320 + let choice_value = if tools.is_empty() {
321 + ToolChoiceValue::None
322 + } else {
323 + ToolChoiceValue::Auto
324 + };
325 + self.tool_choice = Some(ToolChoice::Value(choice_value));
326 + }
327 + }
328 +}
329 +
330 +/// Mirrors the Python vLLM `ChatCompletionResponse` class.
331 +#[serde_with::skip_serializing_none]
332 +#[derive(Debug, Clone, Serialize)]
333 +pub(super) struct ChatCompletionResponse {
334 + pub id: String,
335 + pub object: String,
336 + pub created: u64,
337 + pub model: String,
338 + pub choices: Vec<ChatCompletionChoice>,
339 + pub usage: Option<Usage>,
340 + pub system_fingerprint: Option<String>,
341 + pub prompt_logprobs: Option<Vec<Option<HashMap<String, f32>>>>,
342 + pub prompt_token_ids: Option<Vec<u32>>,
343 + pub kv_transfer_params: Option<Value>,
344 +}
345 +
346 +/// Mirrors the Python vLLM `ChatCompletionResponseChoice` class.
347 +#[serde_with::skip_serializing_none]
348 +#[derive(Debug, Clone, Serialize)]
349 +pub(super) struct ChatCompletionChoice {
350 + pub index: u32,
351 + pub message: ChatCompletionMessage,
352 + pub logprobs: Option<ChatLogProbs>,
353 + pub finish_reason: Option<String>,
354 + pub stop_reason: Option<Value>,
355 + pub token_ids: Option<Vec<u32>>,
356 +}
357 +
358 +/// A literal type for the "assistant" role, since the API only allows that
359 +/// specific value in responses.
360 +#[derive(Debug, Clone, Copy, PartialEq, Eq, SerializeDisplay)]
361 +pub(super) struct AssistantRole;
362 +
363 +impl fmt::Display for AssistantRole {
364 + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
365 + f.write_str("assistant")
366 + }
367 +}
368 +
369 +/// Mirrors the Python vLLM response `ChatMessage` class.
370 +#[serde_with::skip_serializing_none]
371 +#[derive(Debug, Clone, Serialize)]
372 +pub(super) struct ChatCompletionMessage {
373 + pub role: AssistantRole,
374 + pub content: Option<String>,
375 + pub tool_calls: Option<Vec<ToolCall>>,
376 + pub reasoning: Option<String>,
377 +}
378 +
379 +/// Mirrors the Python vLLM `ChatCompletionStreamResponse` class.
380 +#[serde_with::skip_serializing_none]
381 +#[derive(Debug, Clone, Serialize)]
382 +pub(super) struct ChatCompletionStreamResponse {
383 + pub id: String,
384 + pub object: String,
385 + pub created: u64,
386 + pub model: String,
387 + pub choices: Vec<ChatCompletionStreamChoice>,
388 + pub usage: Option<Usage>,
389 + pub prompt_token_ids: Option<Vec<u32>>,
390 +}
391 +
392 +impl ChatCompletionStreamResponse {
393 + /// Create a stream response with the standard envelope fields pre-filled.
394 + pub fn new(id: &str, model: &str, created: u64) -> Self {
395 + Self {
396 + id: id.to_string(),
397 + object: "chat.completion.chunk".to_string(),
398 + created,
399 + model: model.to_string(),
400 + choices: Vec::new(),
401 + usage: None,
402 + prompt_token_ids: None,
403 + }
404 + }
405 +}
406 +
407 +/// Mirrors the Python vLLM `ChatCompletionResponseStreamChoice` class.
408 +#[serde_with::skip_serializing_none]
409 +#[derive(Debug, Clone, Default, Serialize)]
410 +pub(super) struct ChatCompletionStreamChoice {
411 + pub index: u32,
412 + pub delta: ChatMessageDelta,
413 + pub logprobs: Option<ChatLogProbs>,
414 + pub finish_reason: Option<String>,
415 + pub stop_reason: Option<Value>,
416 + pub token_ids: Option<Vec<u32>>,
417 +}
418 +
419 +/// Mirrors the Python vLLM `DeltaMessage` class.
420 +#[serde_with::skip_serializing_none]
421 +#[derive(Debug, Clone, Default, Serialize)]
422 +pub(super) struct ChatMessageDelta {
423 + pub role: Option<AssistantRole>,
424 + pub content: Option<String>,
425 + pub tool_calls: Option<Vec<ToolCallDelta>>,
426 + pub reasoning: Option<String>,
427 +}
428 +
429 +fn default_model() -> String {
430 + UNKNOWN_MODEL_ID.to_string()
431 +}
432 +
433 +/// Validates messages array is not empty and has valid content
434 +fn validate_messages(messages: &[ChatMessage]) -> Result<(), validator::ValidationError> {
435 + if messages.is_empty() {
436 + return Err(validator::ValidationError::new("messages cannot be empty"));
437 + }
438 +
439 + for msg in messages {
440 + if let ChatMessage::User { content, .. } = msg {
441 + match content {
442 + MessageContent::Text(text) if text.is_empty() => {
443 + return Err(validator::ValidationError::new(
444 + "message content cannot be empty",
445 + ));
446 + }
447 + MessageContent::Parts(parts) if parts.is_empty() => {
448 + return Err(validator::ValidationError::new(
449 + "message content parts cannot be empty",
450 + ));
451 + }
452 + _ => {}
453 + }
454 + }
455 + }
456 + Ok(())
457 +}
458 +
459 +/// Schema-level validation for cross-field dependencies
460 +fn validate_chat_cross_parameters(
461 + req: &ChatCompletionRequest,
462 +) -> Result<(), validator::ValidationError> {
463 + // 1. Validate logprobs dependency
464 + if req.top_logprobs.is_some() && !req.logprobs {
465 + let mut e = validator::ValidationError::new("top_logprobs_requires_logprobs");
466 + e.message = Some("top_logprobs is only allowed when logprobs is enabled".into());
467 + return Err(e);
468 + }
469 +
470 + // 2. Validate stream_options dependency
471 + if req.stream_options.is_some() && !req.stream {
472 + let mut e = validator::ValidationError::new("stream_options_requires_stream");
473 + e.message = Some("stream_options can only be used when stream is true".into());
474 + return Err(e);
475 + }
476 +
477 + // 3. Validate token limits - min <= max
478 + if let (Some(min_tokens), Some(max_completion_tokens)) =
479 + (req.min_tokens, req.max_completion_tokens)
480 + && min_tokens > max_completion_tokens
481 + {
482 + let mut e = validator::ValidationError::new("min_tokens_exceeds_max_completion_tokens");
483 + e.message = Some("min_tokens cannot be greater than max_completion_tokens".into());
484 + return Err(e);
485 + }
486 +
487 + #[expect(deprecated, reason = "Local type still mirrors legacy upstream field")]
488 + if let (Some(min_tokens), Some(max_tokens)) = (req.min_tokens, req.max_tokens)
489 + && min_tokens > max_tokens
490 + {
491 + let mut e = validator::ValidationError::new("min_tokens_exceeds_max_tokens");
492 + e.message = Some("min_tokens cannot be greater than max_tokens".into());
493 + return Err(e);
494 + }
495 +
496 + // 4. Validate response format JSON schema name
497 + if let Some(ResponseFormat::JsonSchema { json_schema }) = &req.response_format
498 + && json_schema.name.is_empty()
499 + {
500 + let mut e = validator::ValidationError::new("json_schema_name_empty");
501 + e.message = Some("JSON schema name cannot be empty".into());
502 + return Err(e);
503 + }
504 +
505 + // 5. Validate tool_choice requires tools (except for "none")
506 + if let Some(ref tool_choice) = req.tool_choice {
507 + let has_tools = req.tools.as_ref().is_some_and(|t| !t.is_empty());
508 +
509 + // Check if tool_choice is anything other than "none"
510 + let is_some_choice = !matches!(tool_choice, ToolChoice::Value(ToolChoiceValue::None));
511 +
512 + if is_some_choice && !has_tools {
513 + let mut e = validator::ValidationError::new("tool_choice_requires_tools");
514 + e.message = Some("Invalid value for 'tool_choice': 'tool_choice' is only allowed when 'tools' are specified.".into());
515 + return Err(e);
516 + }
517 +
518 + // Additional validation when tools are present
519 + if let Some(tools) = req.tools.as_ref().filter(|t| !t.is_empty()) {
520 + match tool_choice {
521 + ToolChoice::Function { function, .. } => {
522 + // Validate that the specified function name exists in tools
523 + let function_exists = tools.iter().any(|tool| {
524 + tool.tool_type == "function" && tool.function.name == function.name
525 + });
526 +
527 + if !function_exists {
528 + let mut e =
529 + validator::ValidationError::new("tool_choice_function_not_found");
530 + e.message = Some(
531 + format!(
532 + "Invalid value for 'tool_choice': function '{}' not found in 'tools'.",
533 + function.name
534 + )
535 + .into(),
536 + );
537 + return Err(e);
538 + }
539 + }
540 + ToolChoice::AllowedTools {
541 + mode,
542 + tools: allowed_tools,
543 + ..
544 + } => {
545 + // Validate mode is "auto" or "required"
546 + if mode != "auto" && mode != "required" {
547 + let mut e = validator::ValidationError::new("tool_choice_invalid_mode");
548 + e.message = Some(format!(
549 + "Invalid value for 'tool_choice.mode': must be 'auto' or 'required', got '{mode}'."
550 + ).into());
551 + return Err(e);
552 + }
553 +
554 + // Validate that all ToolReferences are Function type (Chat API only supports
555 + // function tools)
556 + for tool_ref in allowed_tools {
557 + match tool_ref {
558 + ToolReference::Function { name } => {
559 + // Validate that the function exists in tools array
560 + let tool_exists = tools.iter().any(|tool| {
561 + tool.tool_type == "function" && tool.function.name == *name
562 + });
563 +
564 + if !tool_exists {
565 + let mut e = validator::ValidationError::new(
566 + "tool_choice_tool_not_found",
567 + );
568 + e.message = Some(
569 + format!(
570 + "Invalid value for 'tool_choice.tools': tool '{name}' not found in 'tools'."
571 + )
572 + .into(),
573 + );
574 + return Err(e);
575 + }
576 + }
577 + _ => {
578 + // Chat Completion API only supports function tools in tool_choice
579 + let mut e = validator::ValidationError::new(
580 + "tool_choice_invalid_tool_type",
581 + );
582 + e.message = Some(
583 + format!(
584 + "Invalid value for 'tool_choice.tools': Chat Completion API only supports function tools, got '{}'.",
585 + tool_ref.identifier()
586 + )
587 + .into(),
588 + );
589 + return Err(e);
590 + }
591 + }
592 + }
593 + }
594 + ToolChoice::Value(_) => {}
595 + }
596 + }
597 + }
598 +
599 + Ok(())
600 +}
rust/src/server/src/routes/openai/chat_completions/validate.rs new
+410
@@ -0,0 +1,410 @@
1 +use super::types::ChatCompletionRequest;
2 +use crate::error::{ApiError, bail_invalid_request};
3 +use crate::routes::openai::utils::types::{ChatMessage, Tool, ToolChoice, ToolChoiceValue};
4 +
5 +/// Enforce the minimal compatibility contract for the Rust OpenAI server.
6 +pub(super) fn validate_request_compat(
7 + request: &ChatCompletionRequest,
8 + served_model_names: &[String],
9 +) -> Result<(), ApiError> {
10 + if !served_model_names.iter().any(|n| n == &request.model) {
11 + return Err(ApiError::model_not_found(request.model.clone()));
12 + }
13 +
14 + if request.stream_options.is_some() && !request.stream {
15 + bail_invalid_request!(
16 + param = "stream_options",
17 + "stream_options are only supported when stream=true."
18 + );
19 + }
20 +
21 + if request.n.unwrap_or(1) > 1 {
22 + bail_invalid_request!(param = "n", "Only n=1 is supported.");
23 + }
24 +
25 + if request.top_logprobs.is_some() && !request.logprobs {
26 + bail_invalid_request!(
27 + param = "top_logprobs",
28 + "top_logprobs can only be used when logprobs=true."
29 + );
30 + }
31 +
32 + if let Some(prompt_logprobs) = request.prompt_logprobs {
33 + if prompt_logprobs < 0 && prompt_logprobs != -1 {
34 + bail_invalid_request!(
35 + param = "prompt_logprobs",
36 + "prompt_logprobs must be a non-negative value or -1."
37 + );
38 + }
39 +
40 + if request.stream && (prompt_logprobs > 0 || prompt_logprobs == -1) {
41 + bail_invalid_request!(
42 + param = "prompt_logprobs",
43 + "prompt_logprobs are not available when stream=true."
44 + );
45 + }
46 + }
47 +
48 + if let Some(tools) = request.tools.as_ref() {
49 + validate_function_tools(tools, "tools")?;
50 + }
51 +
52 + for message in &request.messages {
53 + if let ChatMessage::Developer {
54 + tools: Some(tools), ..
55 + } = message
56 + {
57 + validate_function_tools(tools, "messages[].tools")?;
58 + }
59 + }
60 +
61 + if let Some(tool_choice) = &request.tool_choice {
62 + match tool_choice {
63 + ToolChoice::Value(ToolChoiceValue::Auto | ToolChoiceValue::None) => {}
64 + ToolChoice::Value(ToolChoiceValue::Required) => {
65 + bail_invalid_request!(
66 + param = "tool_choice",
67 + "tool_choice=required is not supported yet."
68 + );
69 + }
70 + ToolChoice::Function { .. } => {
71 + bail_invalid_request!(
72 + param = "tool_choice",
73 + "Named function tool_choice is not supported yet."
74 + );
75 + }
76 + ToolChoice::AllowedTools { .. } => {
77 + bail_invalid_request!(
78 + param = "tool_choice",
79 + "allowed_tools tool_choice is not supported yet."
80 + );
81 + }
82 + }
83 + }
84 +
85 + if request.use_beam_search {
86 + bail_invalid_request!(
87 + param = "use_beam_search",
88 + "use_beam_search is not supported."
89 + );
90 + }
91 +
92 + // ---- Reject parameters that are accepted for deserialization but not yet
93 + // implemented ----
94 +
95 + if request.parallel_tool_calls.is_some() {
96 + bail_invalid_request!(
97 + param = "parallel_tool_calls",
98 + "parallel_tool_calls is not supported."
99 + );
100 + }
101 +
102 + reject_non_default(
103 + request.length_penalty.as_ref(),
104 + "length_penalty",
105 + "length_penalty is not supported.",
106 + )?;
107 + if !request.spaces_between_special_tokens {
108 + bail_invalid_request!(
109 + param = "spaces_between_special_tokens",
110 + "spaces_between_special_tokens is not supported."
111 + );
112 + }
113 + reject_non_default(
114 + request.truncate_prompt_tokens.as_ref(),
115 + "truncate_prompt_tokens",
116 + "truncate_prompt_tokens is not supported.",
117 + )?;
118 + reject_non_default(
119 + request.thinking_token_budget.as_ref(),
120 + "thinking_token_budget",
121 + "thinking_token_budget is not supported.",
122 + )?;
123 + if !request.include_reasoning {
124 + bail_invalid_request!(
125 + param = "include_reasoning",
126 + "include_reasoning is not supported."
127 + );
128 + }
129 + reject_non_default(
130 + request.media_io_kwargs.as_ref(),
131 + "media_io_kwargs",
132 + "media_io_kwargs is not supported.",
133 + )?;
134 + reject_non_default(
135 + request.mm_processor_kwargs.as_ref(),
136 + "mm_processor_kwargs",
137 + "mm_processor_kwargs is not supported.",
138 + )?;
139 + reject_non_default(
140 + request.repetition_detection.as_ref(),
141 + "repetition_detection",
142 + "repetition_detection is not supported.",
143 + )?;
144 +
145 + if let Some(options) = &request.stream_options
146 + && options.continuous_usage_stats.is_some()
147 + {
148 + bail_invalid_request!(
149 + param = "stream_options",
150 + "continuous_usage_stats is not supported."
151 + );
152 + }
153 +
154 + Ok(())
155 +}
156 +
157 +/// Reject one option unless it is entirely absent.
158 +fn reject_non_default<T>(
159 + value: Option<&T>,
160 + param: &'static str,
161 + message: &str,
162 +) -> Result<(), ApiError> {
163 + if value.is_some() {
164 + bail_invalid_request!(param = param, "{}", message);
165 + }
166 + Ok(())
167 +}
168 +
169 +fn validate_function_tools(tools: &[Tool], param: &'static str) -> Result<(), ApiError> {
170 + for tool in tools {
171 + if tool.tool_type != "function" {
172 + bail_invalid_request!(param = param, "Only function tools are supported.");
173 + }
174 + }
175 + Ok(())
176 +}
177 +
178 +#[cfg(test)]
179 +mod tests {
180 + use std::collections::HashMap;
181 +
182 + use serde_json::json;
183 + use vllm_chat::ReasoningEffort;
184 +
185 + use super::validate_request_compat;
186 + use crate::routes::openai::chat_completions::types::ChatCompletionRequest;
187 + use crate::routes::openai::utils::structured_outputs::ResponseFormat;
188 + use crate::routes::openai::utils::types::{
189 + ChatMessage, Function, FunctionChoice, MessageContent, StringOrArray, Tool, ToolChoice,
190 + ToolChoiceValue, ToolReference,
191 + };
192 +
193 + fn served(names: &[&str]) -> Vec<String> {
194 + names.iter().map(|s| s.to_string()).collect()
195 + }
196 +
197 + fn base_request() -> ChatCompletionRequest {
198 + ChatCompletionRequest {
199 + model: "Qwen/Qwen1.5-0.5B-Chat".to_string(),
200 + messages: vec![ChatMessage::User {
201 + content: MessageContent::Text("hello".to_string()),
202 + name: None,
203 + }],
204 + stream: true,
205 + ..Default::default()
206 + }
207 + }
208 +
209 + #[test]
210 + fn validate_request_compat_accepts_stop() {
211 + let request = ChatCompletionRequest {
212 + stop: Some(StringOrArray::String("stop".to_string())),
213 + ..base_request()
214 + };
215 +
216 + validate_request_compat(&request, &served(&["Qwen/Qwen1.5-0.5B-Chat"]))
217 + .expect("stop strings should be accepted");
218 + }
219 +
220 + #[test]
221 + fn validate_request_compat_accepts_non_zero_penalties_and_function_tools() {
222 + let request = ChatCompletionRequest {
223 + frequency_penalty: Some(0.5),
224 + presence_penalty: Some(0.25),
225 + min_p: Some(0.2),
226 + repetition_penalty: Some(1.1),
227 + seed: Some(7),
228 + ..base_request()
229 + };
230 + validate_request_compat(&request, &served(&["Qwen/Qwen1.5-0.5B-Chat"]))
231 + .expect("sampling fields should be accepted");
232 +
233 + let request = ChatCompletionRequest {
234 + tools: Some(vec![Tool {
235 + tool_type: "function".to_string(),
236 + function: Function {
237 + name: "tool".to_string(),
238 + description: None,
239 + parameters: json!({}),
240 + strict: None,
241 + },
242 + }]),
243 + ..base_request()
244 + };
245 + validate_request_compat(&request, &served(&["Qwen/Qwen1.5-0.5B-Chat"]))
246 + .expect("function tools should be accepted");
247 +
248 + let request = ChatCompletionRequest {
249 + messages: vec![ChatMessage::Developer {
250 + content: MessageContent::Text("policy".to_string()),
251 + tools: Some(vec![Tool {
252 + tool_type: "function".to_string(),
253 + function: Function {
254 + name: "tool".to_string(),
255 + description: None,
256 + parameters: json!({}),
257 + strict: None,
258 + },
259 + }]),
260 + name: None,
261 + }],
262 + ..base_request()
263 + };
264 + validate_request_compat(&request, &served(&["Qwen/Qwen1.5-0.5B-Chat"]))
265 + .expect("developer function tools should be accepted");
266 + }
267 +
268 + #[test]
269 + fn validate_request_compat_rejects_non_function_developer_tools() {
270 + let request = ChatCompletionRequest {
271 + messages: vec![ChatMessage::Developer {
272 + content: MessageContent::Text("policy".to_string()),
273 + tools: Some(vec![Tool {
274 + tool_type: "mcp".to_string(),
275 + function: Function {
276 + name: "tool".to_string(),
277 + description: None,
278 + parameters: json!({}),
279 + strict: None,
280 + },
281 + }]),
282 + name: None,
283 + }],
284 + ..base_request()
285 + };
286 +
287 + assert!(validate_request_compat(&request, &served(&["Qwen/Qwen1.5-0.5B-Chat"])).is_err());
288 + }
289 +
290 + #[test]
291 + fn validate_request_compat_accepts_output_logprobs() {
292 + let request = ChatCompletionRequest {
293 + logprobs: true,
294 + ..base_request()
295 + };
296 + validate_request_compat(&request, &served(&["Qwen/Qwen1.5-0.5B-Chat"]))
297 + .expect("logprobs should be accepted");
298 + }
299 +
300 + #[test]
301 + fn validate_request_compat_accepts_reasoning_effort() {
302 + let request = ChatCompletionRequest {
303 + reasoning_effort: Some(ReasoningEffort::Max),
304 + chat_template_kwargs: Some(HashMap::from([(
305 + "reasoning_effort".to_string(),
306 + json!("low"),
307 + )])),
308 + ..base_request()
309 + };
310 +
311 + validate_request_compat(&request, &served(&["Qwen/Qwen1.5-0.5B-Chat"]))
312 + .expect("reasoning_effort should be accepted");
313 + }
314 +
315 + #[test]
316 + fn validate_request_compat_rejects_top_logprobs_without_logprobs() {
317 + let request = ChatCompletionRequest {
318 + top_logprobs: Some(0),
319 + ..base_request()
320 + };
321 + assert!(validate_request_compat(&request, &served(&["Qwen/Qwen1.5-0.5B-Chat"])).is_err());
322 + }
323 +
324 + #[test]
325 + fn validate_request_compat_rejects_streaming_prompt_logprobs_requests() {
326 + let request = ChatCompletionRequest {
327 + prompt_logprobs: Some(1),
328 + ..base_request()
329 + };
330 + assert!(validate_request_compat(&request, &served(&["Qwen/Qwen1.5-0.5B-Chat"])).is_err());
331 +
332 + let request = ChatCompletionRequest {
333 + prompt_logprobs: Some(-1),
334 + ..base_request()
335 + };
336 + assert!(validate_request_compat(&request, &served(&["Qwen/Qwen1.5-0.5B-Chat"])).is_err());
337 + }
338 +
339 + #[test]
340 + fn validate_request_compat_rejects_invalid_prompt_logprobs_value() {
341 + let request = ChatCompletionRequest {
342 + stream: false,
343 + prompt_logprobs: Some(-2),
344 + ..base_request()
345 + };
346 + assert!(validate_request_compat(&request, &served(&["Qwen/Qwen1.5-0.5B-Chat"])).is_err());
347 + }
348 +
349 + #[test]
350 + fn validate_request_compat_accepts_response_format() {
351 + let request = ChatCompletionRequest {
352 + response_format: Some(ResponseFormat::Text),
353 + ..base_request()
354 + };
355 + validate_request_compat(&request, &served(&["Qwen/Qwen1.5-0.5B-Chat"]))
356 + .expect("response_format=text should be accepted");
357 +
358 + let request = ChatCompletionRequest {
359 + response_format: Some(ResponseFormat::JsonObject),
360 + ..base_request()
361 + };
362 + validate_request_compat(&request, &served(&["Qwen/Qwen1.5-0.5B-Chat"]))
363 + .expect("response_format=json_object should be accepted");
364 + }
365 +
366 + #[test]
367 + fn validate_request_compat_accepts_noop_tool_choice_none() {
368 + let request = ChatCompletionRequest {
369 + tool_choice: Some(ToolChoice::Value(ToolChoiceValue::None)),
370 + ..base_request()
371 + };
372 +
373 + validate_request_compat(&request, &served(&["Qwen/Qwen1.5-0.5B-Chat"]))
374 + .expect("tool_choice=none is ok");
375 + }
376 +
377 + #[test]
378 + fn validate_request_compat_rejects_required_and_named_tool_choices() {
379 + let required = ChatCompletionRequest {
380 + tool_choice: Some(ToolChoice::Value(ToolChoiceValue::Required)),
381 + ..base_request()
382 + };
383 + assert!(validate_request_compat(&required, &served(&["Qwen/Qwen1.5-0.5B-Chat"])).is_err());
384 +
385 + let named = ChatCompletionRequest {
386 + tool_choice: Some(ToolChoice::Function {
387 + tool_type: "function".to_string(),
388 + function: FunctionChoice {
389 + name: "tool".to_string(),
390 + },
391 + }),
392 + ..base_request()
393 + };
394 + assert!(validate_request_compat(&named, &served(&["Qwen/Qwen1.5-0.5B-Chat"])).is_err());
395 +
396 + let allowed_tools = ChatCompletionRequest {
397 + tool_choice: Some(ToolChoice::AllowedTools {
398 + tool_type: "allowed_tools".to_string(),
399 + mode: "auto".to_string(),
400 + tools: vec![ToolReference::Function {
401 + name: "tool".to_string(),
402 + }],
403 + }),
404 + ..base_request()
405 + };
406 + assert!(
407 + validate_request_compat(&allowed_tools, &served(&["Qwen/Qwen1.5-0.5B-Chat"])).is_err()
408 + );
409 + }
410 +}
rust/src/server/src/routes/openai/completions.rs new
+571
@@ -0,0 +1,571 @@
1 +mod convert;
2 +mod types;
3 +mod validate;
4 +
5 +use std::convert::Infallible;
6 +use std::result::Result;
7 +use std::sync::Arc;
8 +
9 +use asynk_strim_attr::{TryYielder, try_stream};
10 +use axum::Json;
11 +use axum::extract::State;
12 +use axum::http::HeaderMap;
13 +use axum::response::sse::{Event, Sse};
14 +use axum::response::{IntoResponse, Response};
15 +use futures::{Stream, StreamExt as _, pin_mut};
16 +use thiserror_ext::AsReport as _;
17 +use tracing::{debug, error, info, trace};
18 +use tracing_futures::Instrument as _;
19 +use vllm_text::{DecodedTextEvent, FinishReason, TextOutputStream, TextOutputStreamExt as _};
20 +
21 +use super::utils::logprobs::{
22 + collected_logprobs_to_openai, decoded_logprobs_to_openai, decoded_prompt_logprobs_to_maps,
23 + text_len,
24 +};
25 +use super::utils::types::Usage;
26 +use crate::error::{ApiError, bail_server_error, server_error};
27 +use crate::routes::openai::completions::convert::prepare_completion_request;
28 +use crate::routes::openai::completions::types::{
29 + CompletionChoice, CompletionRequest, CompletionResponse, CompletionSseChunk,
30 + CompletionStreamChoice, CompletionStreamResponse,
31 +};
32 +use crate::routes::openai::utils::types::LogProbs;
33 +use crate::routes::openai::utils::validated_json::ValidatedJson;
34 +use crate::state::AppState;
35 +use crate::utils::{resolve_request_context, unix_timestamp};
36 +
37 +/// Validate one completions request and proxy it into the shared `vllm-text`
38 +/// stack.
39 +pub async fn completions(
40 + State(state): State<Arc<AppState>>,
41 + headers: HeaderMap,
42 + ValidatedJson(body): ValidatedJson<CompletionRequest>,
43 +) -> Response {
44 + let stream = body.stream;
45 + let logprobs = body.logprobs;
46 + let request_context = resolve_request_context(&headers, body.request_id.as_deref());
47 +
48 + let prepared =
49 + match prepare_completion_request(body, state.served_model_names(), request_context) {
50 + Ok(prepared) => prepared,
51 + Err(error) => return error.into_response(),
52 + };
53 + let request_span = tracing::info_span!(
54 + "completions",
55 + request_id = %prepared.request_id,
56 + engine_request_id = tracing::field::Empty,
57 + );
58 +
59 + let created = unix_timestamp();
60 + let include_prompt_logprobs = prepared.text_request.sampling_params.prompt_logprobs.is_some();
61 + let log_request = state.enable_log_requests;
62 +
63 + let text_stream = match state
64 + .chat
65 + .text()
66 + .generate(prepared.text_request)
67 + .instrument(request_span.clone())
68 + .await
69 + {
70 + Ok(stream) => stream,
71 + Err(error) => {
72 + return server_error!(
73 + "failed to submit completion request: {}",
74 + error.to_report_string()
75 + )
76 + .into_response();
77 + }
78 + };
79 +
80 + if stream {
81 + let chunk_stream = completion_chunk_stream(
82 + text_stream,
83 + prepared.request_id,
84 + prepared.response_model,
85 + created,
86 + log_request,
87 + prepared.include_usage,
88 + prepared.echo,
89 + logprobs,
90 + prepared.return_token_ids,
91 + prepared.return_tokens_as_token_ids,
92 + );
93 + let sse_stream = completion_sse_stream(chunk_stream).instrument(request_span);
94 +
95 + Sse::new(sse_stream).into_response()
96 + } else {
97 + let response = match collect_completion(
98 + text_stream,
99 + prepared.request_id,
100 + prepared.response_model,
101 + created,
102 + prepared.echo,
103 + logprobs,
104 + include_prompt_logprobs,
105 + prepared.return_token_ids,
106 + prepared.return_tokens_as_token_ids,
107 + )
108 + .instrument(request_span.clone())
109 + .await
110 + {
111 + Ok(response) => response,
112 + Err(error) => return error.into_response(),
113 + };
114 +
115 + if log_request {
116 + let usage = response.usage.as_ref();
117 + info!(
118 + parent: &request_span,
119 + model = %response.model,
120 + prompt_tokens = usage.map_or(0, |u| u.prompt_tokens),
121 + output_tokens = usage.and_then(|u| u.completion_tokens).unwrap_or(0),
122 + finish_reason = response.choices.first().and_then(|c| c.finish_reason.as_deref()).unwrap_or("unknown"),
123 + "completion finished"
124 + );
125 + }
126 +
127 + Json(response).into_response()
128 + }
129 +}
130 +
131 +async fn collect_completion(
132 + stream: impl TextOutputStream,
133 + request_id: String,
134 + response_model: String,
135 + created: u64,
136 + echo: Option<String>,
137 + requested_logprobs: Option<u32>,
138 + include_prompt_logprobs: bool,
139 + return_token_ids: bool,
140 + return_tokens_as_token_ids: bool,
141 +) -> Result<CompletionResponse, ApiError> {
142 + let collected = stream
143 + .collect_output()
144 + .await
145 + .map_err(|error| server_error!("completion stream failed: {}", error.to_report_string()))?;
146 + let finish_reason = collected.finish_reason.clone();
147 + let stop_reason = finish_reason
148 + .as_stop_reason()
149 + .map(|sr| serde_json::to_value(sr).expect("StopReason must serialize to JSON"));
150 +
151 + let prompt_char_count = echo.as_ref().map(|prompt| text_len(prompt)).unwrap_or_default();
152 + let prompt_logprobs = if include_prompt_logprobs {
153 + let prompt_logprobs = collected.prompt_logprobs.as_ref().ok_or_else(|| {
154 + server_error!(
155 + "completion response requested prompt_logprobs but generation returned none"
156 + )
157 + })?;
158 + Some(prompt_logprobs)
159 + } else {
160 + None
161 + };
162 + let logprobs = if requested_logprobs.is_some() {
163 + Some(collected_logprobs_to_openai(
164 + &collected,
165 + echo.is_some(),
166 + prompt_char_count,
167 + return_tokens_as_token_ids,
168 + )?)
169 + } else {
170 + None
171 + };
172 + let prompt_logprobs =
173 + prompt_logprobs.map(|lp| decoded_prompt_logprobs_to_maps(lp, return_tokens_as_token_ids));
174 + let text = match &echo {
175 + None => collected.text,
176 + Some(prompt) => format!("{prompt}{}", collected.text),
177 + };
178 +
179 + Ok(CompletionResponse {
180 + id: request_id,
181 + object: "text_completion".to_string(),
182 + created,
183 + model: response_model,
184 + choices: vec![CompletionChoice {
185 + index: 0,
186 + text,
187 + logprobs,
188 + finish_reason: Some(completion_finish_reason_to_openai(finish_reason)?.into()),
189 + stop_reason,
190 + prompt_logprobs,
191 + token_ids: return_token_ids.then(|| collected.token_ids.clone()),
192 + prompt_token_ids: return_token_ids.then(|| collected.prompt_token_ids.to_vec()),
193 + }],
194 + usage: Some(Usage::from_counts(
195 + collected.prompt_token_ids.len() as u32,
196 + collected.token_ids.len() as u32,
197 + )),
198 + system_fingerprint: None,
199 + kv_transfer_params: collected.kv_transfer_params,
200 + })
201 +}
202 +
203 +/// Convert one internal decoded-text stream into OpenAI completions chunks.
204 +#[try_stream]
205 +async fn completion_chunk_stream(
206 + stream: impl TextOutputStream,
207 + request_id: String,
208 + response_model: String,
209 + created: u64,
210 + log_request: bool,
211 + include_usage: bool,
212 + echo: Option<String>,
213 + requested_logprobs: Option<u32>,
214 + return_token_ids: bool,
215 + return_tokens_as_token_ids: bool,
216 + mut y: TryYielder<CompletionSseChunk, ApiError>,
217 +) -> Result<(), ApiError> {
218 + pin_mut!(stream);
219 + let mut visible_text_len = 0_u32;
220 + let mut first_chunk = true;
221 +
222 + while let Some(next) = stream.next().await {
223 + match next {
224 + Ok(DecodedTextEvent::Start {
225 + prompt_token_ids, ..
226 + }) => {
227 + debug!("completion stream started");
228 + if let Some(prompt) = echo.as_ref() {
229 + visible_text_len = text_len(prompt);
230 + let mut chunk =
231 + delta_chunk(&request_id, &response_model, created, prompt.clone(), None);
232 + if return_token_ids && first_chunk {
233 + if let Some(choice) = chunk.choices.first_mut() {
234 + choice.prompt_token_ids = Some(prompt_token_ids.to_vec());
235 + }
236 + first_chunk = false;
237 + }
238 + y.yield_ok(CompletionSseChunk::Chunk(chunk)).await;
239 + } else if return_token_ids {
240 + // Emit a chunk with prompt_token_ids in the first streaming response
241 + let mut chunk =
242 + delta_chunk(&request_id, &response_model, created, String::new(), None);
243 + if let Some(choice) = chunk.choices.first_mut() {
244 + choice.prompt_token_ids = Some(prompt_token_ids.to_vec());
245 + }
246 + first_chunk = false;
247 + y.yield_ok(CompletionSseChunk::Chunk(chunk)).await;
248 + }
249 + }
250 + Ok(DecodedTextEvent::TextDelta {
251 + delta,
252 + token_ids,
253 + logprobs,
254 + finished,
255 + }) => {
256 + let delta_text_len = text_len(&delta);
257 + let logprobs = if requested_logprobs.is_some() {
258 + let decoded_logprobs = logprobs.as_ref().ok_or_else(|| {
259 + server_error!(
260 + "completion stream requested logprobs but generation returned none"
261 + )
262 + })?;
263 + Some(decoded_logprobs_to_openai(
264 + decoded_logprobs,
265 + visible_text_len,
266 + return_tokens_as_token_ids,
267 + )?)
268 + } else {
269 + None
270 + };
271 + let mut chunk = delta_chunk(&request_id, &response_model, created, delta, logprobs);
272 + if return_token_ids && let Some(choice) = chunk.choices.first_mut() {
273 + choice.token_ids = Some(token_ids);
274 + }
275 + y.yield_ok(CompletionSseChunk::Chunk(chunk)).await;
276 + visible_text_len = visible_text_len.saturating_add(delta_text_len);
277 +
278 + if let Some(finished) = finished {
279 + if log_request {
280 + info!(
281 + stream = true,
282 + model = %response_model,
283 + prompt_tokens = finished.prompt_token_count,
284 + output_tokens = finished.output_token_count,
285 + finish_reason = finished.finish_reason.as_str(),
286 + "completion finished"
287 + );
288 + }
289 + y.yield_ok(CompletionSseChunk::Chunk(final_chunk(
290 + &request_id,
291 + &response_model,
292 + created,
293 + finished.finish_reason,
294 + )?))
295 + .await;
296 +
297 + if include_usage {
298 + y.yield_ok(CompletionSseChunk::Usage(usage_chunk(
299 + &request_id,
300 + &response_model,
301 + created,
302 + Usage::from_counts(
303 + finished.prompt_token_count as u32,
304 + finished.output_token_count as u32,
305 + ),
306 + )))
307 + .await;
308 + }
309 + }
310 + }
311 + Err(error) => {
312 + error!(
313 + error = %error.as_report(),
314 + "completion stream failed"
315 + );
316 + bail_server_error!("{}", error.to_report_string());
317 + }
318 + }
319 + }
320 + Ok(())
321 +}
322 +
323 +fn delta_chunk(
324 + request_id: &str,
325 + response_model: &str,
326 + created: u64,
327 + text: String,
328 + logprobs: Option<LogProbs>,
329 +) -> CompletionStreamResponse {
330 + let mut chunk = CompletionStreamResponse::new(request_id, response_model, created);
331 + chunk.choices.push(CompletionStreamChoice {
332 + text,
333 + logprobs,
334 + ..Default::default()
335 + });
336 + chunk
337 +}
338 +
339 +fn final_chunk(
340 + request_id: &str,
341 + response_model: &str,
342 + created: u64,
343 + finish_reason: FinishReason,
344 +) -> Result<CompletionStreamResponse, ApiError> {
345 + let finish_reason = completion_finish_reason_to_openai(finish_reason)?;
346 +
347 + let mut chunk = CompletionStreamResponse::new(request_id, response_model, created);
348 + chunk.choices.push(CompletionStreamChoice {
349 + finish_reason: Some(finish_reason.to_string()),
350 + ..Default::default()
351 + });
352 + Ok(chunk)
353 +}
354 +
355 +fn completion_finish_reason_to_openai(
356 + finish_reason: FinishReason,
357 +) -> Result<&'static str, ApiError> {
358 + match finish_reason {
359 + FinishReason::Stop(_) | FinishReason::Repetition => Ok("stop"),
360 + FinishReason::Length => Ok("length"),
361 + FinishReason::Abort => Ok("abort"),
362 + FinishReason::Error => {
363 + bail_server_error!("Internal server error");
364 + }
365 + }
366 +}
367 +
368 +fn usage_chunk(
369 + request_id: &str,
370 + response_model: &str,
371 + created: u64,
372 + usage: Usage,
373 +) -> CompletionStreamResponse {
374 + let mut chunk = CompletionStreamResponse::new(request_id, response_model, created);
375 + chunk.usage = Some(usage);
376 + chunk
377 +}
378 +
379 +/// Convert one chunk stream into OpenAI-style SSE events.
380 +///
381 +/// OpenAI-style streaming errors are encoded as ordinary `data: {"error": ...}`
382 +/// events followed by `data: [DONE]`, so the transport stream itself stays
383 +/// infallible even when generation fails after the HTTP response has started.
384 +#[try_stream]
385 +async fn completion_sse_stream(
386 + stream: impl Stream<Item = Result<CompletionSseChunk, ApiError>>,
387 + mut y: TryYielder<Event, Infallible>,
388 +) -> Result<(), Infallible> {
389 + pin_mut!(stream);
390 +
391 + while let Some(next) = stream.next().await {
392 + match next {
393 + Ok(chunk) => y.yield_ok(to_sse_event(&chunk)).await,
394 + Err(error) => {
395 + y.yield_ok(to_error_sse_event(&error)).await;
396 + break;
397 + }
398 + }
399 + }
400 +
401 + y.yield_ok(done_sse_event()).await;
402 + Ok(())
403 +}
404 +
405 +/// Serialize one OpenAI chunk payload into one SSE `data:` event.
406 +fn to_sse_event(chunk: &CompletionSseChunk) -> Event {
407 + let payload = serde_json::to_string(chunk).expect("completion chunk must serialize to JSON");
408 + trace!(payload, "completion emitting chunk");
409 + Event::default().data(payload)
410 +}
411 +
412 +/// Serialize one OpenAI error payload into one SSE `data:` event.
413 +fn to_error_sse_event(error: &ApiError) -> Event {
414 + let payload = serde_json::to_string(&error.to_error_response())
415 + .expect("ErrorResponse must serialize to JSON");
416 + trace!(payload, "completion emitting error");
417 + Event::default().data(payload)
418 +}
419 +
420 +/// Build the terminal OpenAI SSE sentinel event.
421 +fn done_sse_event() -> Event {
422 + trace!("completion emitting done");
423 + Event::default().data("[DONE]")
424 +}
425 +
426 +#[cfg(test)]
427 +mod tests {
428 + use futures::{StreamExt as _, stream};
429 + use itertools::Itertools as _;
430 + use vllm_text::{
431 + DecodedLogprobs, DecodedPositionLogprobs, DecodedTextEvent, DecodedTokenLogprob,
432 + FinishReason, Finished,
433 + };
434 +
435 + use super::{CompletionSseChunk, completion_chunk_stream, final_chunk};
436 +
437 + #[test]
438 + fn final_chunk_maps_stop_finish_reason() {
439 + let chunk = final_chunk("cmpl-1", "model", 1, FinishReason::stop_eos())
440 + .expect("finish reason valid");
441 + assert_eq!(chunk.choices[0].finish_reason.as_deref(), Some("stop"));
442 + assert_eq!(chunk.choices[0].text, "");
443 + }
444 +
445 + #[test]
446 + fn final_chunk_maps_length_finish_reason() {
447 + let chunk =
448 + final_chunk("cmpl-1", "model", 1, FinishReason::Length).expect("finish reason valid");
449 + assert_eq!(chunk.choices[0].finish_reason.as_deref(), Some("length"));
450 + }
451 +
452 + #[test]
453 + fn final_chunk_maps_abort_finish_reason() {
454 + let chunk =
455 + final_chunk("cmpl-1", "model", 1, FinishReason::Abort).expect("finish reason valid");
456 + assert_eq!(chunk.choices[0].finish_reason.as_deref(), Some("abort"));
457 + }
458 +
459 + #[test]
460 + fn final_chunk_rejects_error_finish_reason() {
461 + assert!(final_chunk("cmpl-1", "model", 1, FinishReason::Error).is_err());
462 + }
463 +
464 + #[tokio::test]
465 + async fn completion_chunk_stream_maps_streaming_logprobs() {
466 + let stream = stream::iter(vec![
467 + Ok(DecodedTextEvent::Start {
468 + prompt_token_ids: vec![1, 2, 3, 4, 5].into(),
469 + prompt_logprobs: None,
470 + }),
471 + Ok(DecodedTextEvent::TextDelta {
472 + delta: "h".to_string(),
473 + token_ids: vec![b'h' as u32],
474 + logprobs: Some(DecodedLogprobs {
475 + positions: vec![DecodedPositionLogprobs {
476 + entries: vec![
477 + DecodedTokenLogprob {
478 + token_id: 0,
479 + token: "h".to_string(),
480 + logprob: -0.1,
481 + rank: 1,
482 + },
483 + DecodedTokenLogprob {
484 + token_id: 0,
485 + token: "H".to_string(),
486 + logprob: -0.2,
487 + rank: 1,
488 + },
489 + ],
490 + }],
491 + }),
492 + finished: None,
493 + }),
494 + Ok(DecodedTextEvent::TextDelta {
495 + delta: String::new(),
496 + token_ids: vec![b'!' as u32],
497 + logprobs: Some(DecodedLogprobs {
498 + positions: vec![DecodedPositionLogprobs {
499 + entries: vec![
500 + DecodedTokenLogprob {
501 + token_id: 0,
502 + token: "!".to_string(),
503 + logprob: -0.3,
504 + rank: 1,
505 + },
506 + DecodedTokenLogprob {
507 + token_id: 0,
508 + token: "?".to_string(),
509 + logprob: -0.4,
510 + rank: 1,
511 + },
512 + ],
513 + }],
514 + }),
515 + finished: Some(Finished {
516 + prompt_token_count: 5,
517 + output_token_count: 2,
518 + finish_reason: FinishReason::stop_eos(),
519 + kv_transfer_params: None,
520 + }),
521 + }),
522 + ]);
523 +
524 + let chunks = completion_chunk_stream(
525 + stream,
526 + "cmpl-1".to_string(),
527 + "model".to_string(),
528 + 1,
529 + false,
530 + false,
531 + None,
532 + Some(1),
533 + false,
534 + false,
535 + )
536 + .collect::<Vec<_>>()
537 + .await;
538 +
539 + let chunks: Vec<_> = chunks.into_iter().try_collect().expect("stream should succeed");
540 +
541 + match &chunks[0] {
542 + CompletionSseChunk::Chunk(chunk) => {
543 + assert_eq!(chunk.choices[0].text, "h");
544 + assert_eq!(
545 + chunk.choices[0].logprobs.as_ref().expect("logprobs").tokens,
546 + vec!["h".to_string()]
547 + );
548 + assert_eq!(
549 + chunk.choices[0].logprobs.as_ref().expect("logprobs").text_offset,
550 + vec![0]
551 + );
552 + }
553 + CompletionSseChunk::Usage(_) => panic!("expected regular chunk"),
554 + }
555 +
556 + match &chunks[1] {
557 + CompletionSseChunk::Chunk(chunk) => {
558 + assert_eq!(chunk.choices[0].text, "");
559 + assert_eq!(
560 + chunk.choices[0].logprobs.as_ref().expect("logprobs").tokens,
561 + vec!["!".to_string()]
562 + );
563 + assert_eq!(
564 + chunk.choices[0].logprobs.as_ref().expect("logprobs").text_offset,
565 + vec![1]
566 + );
567 + }
568 + CompletionSseChunk::Usage(_) => panic!("expected regular chunk"),
569 + }
570 + }
571 +}
rust/src/server/src/routes/openai/completions/convert.rs new
+352
@@ -0,0 +1,352 @@
1 +use vllm_text::{SamplingParams, TextDecodeOptions, TextRequest};
2 +
3 +use super::types::CompletionRequest;
4 +use crate::error::ApiError;
5 +use crate::routes::openai::completions::validate;
6 +use crate::routes::openai::utils::structured_outputs::convert_from_response_format_value;
7 +use crate::utils::{ResolvedRequestContext, convert_logit_bias, merge_kv_transfer_params};
8 +
9 +/// Lowered completion request plus the public response metadata carried by
10 +/// every SSE chunk.
11 +#[derive(Debug, Clone, PartialEq)]
12 +pub struct PreparedRequest {
13 + /// Stable OpenAI-style request ID, reused as the external text request ID.
14 + pub request_id: String,
15 + /// Public model ID echoed back to the client.
16 + pub response_model: String,
17 + /// Whether the caller asked for the final streamed usage chunk.
18 + pub include_usage: bool,
19 + /// Lowered text request for the shared `vllm-text` facade.
20 + pub text_request: TextRequest,
21 + /// Original text prompt that should be echoed back northbound when
22 + /// `echo=true`.
23 + pub echo: Option<String>,
24 + /// Whether to include token IDs alongside generated text.
25 + pub return_token_ids: bool,
26 + /// Whether to format logprob tokens as `token_id:{id}`.
27 + pub return_tokens_as_token_ids: bool,
28 +}
29 +
30 +/// Validate and lower one OpenAI completions request into the internal
31 +/// text-generation format.
32 +///
33 +/// `served_model_names` must be non-empty; the first entry is used as the
34 +/// `model` field in responses.
35 +pub(crate) fn prepare_completion_request(
36 + request: CompletionRequest,
37 + served_model_names: &[String],
38 + ctx: ResolvedRequestContext,
39 +) -> Result<PreparedRequest, ApiError> {
40 + validate::validate_request_compat(&request, served_model_names)?;
41 +
42 + let request_id = format!("cmpl-{}", ctx.request_id);
43 +
44 + let logprobs = match request.logprobs {
45 + Some(logprobs) => Some(i32::try_from(logprobs).map_err(|_| {
46 + ApiError::invalid_request(
47 + "`logprobs` must fit within a signed 32-bit integer.".to_string(),
48 + Some("logprobs"),
49 + )
50 + })?),
51 + None => None,
52 + };
53 + let prompt_logprobs = request.prompt_logprobs.or(if request.echo && !request.stream {
54 + logprobs
55 + } else {
56 + None
57 + });
58 + let include_usage = (request.stream_options.as_ref())
59 + .and_then(|options| options.include_usage)
60 + .unwrap_or(false);
61 + let echo = request.echo.then(|| request.prompt.as_text().cloned()).flatten();
62 +
63 + let structured_outputs =
64 + convert_from_response_format_value(&request.response_format, &request.structured_outputs)?;
65 +
66 + let text_request = TextRequest {
67 + request_id: request_id.clone(),
68 + prompt: request.prompt,
69 + mm_features: None,
70 + sampling_params: SamplingParams {
71 + temperature: request.temperature,
72 + top_p: request.top_p,
73 + top_k: request.top_k,
74 + seed: request.seed,
75 + max_tokens: request.max_tokens,
76 + min_tokens: request.min_tokens,
77 + logprobs,
78 + prompt_logprobs,
79 + min_p: request.min_p,
80 + frequency_penalty: request.frequency_penalty,
81 + presence_penalty: request.presence_penalty,
82 + repetition_penalty: request.repetition_penalty,
83 + stop_token_ids: request.stop_token_ids,
84 + ignore_eos: request.ignore_eos,
85 + logit_bias: convert_logit_bias(request.logit_bias)?,
86 + allowed_token_ids: request.allowed_token_ids,
87 + bad_words: None,
88 + logprob_token_ids: None,
89 + structured_outputs,
90 + skip_reading_prefix_cache: None,
91 + vllm_xargs: merge_kv_transfer_params(
92 + request.vllm_xargs,
93 + request.kv_transfer_params.as_ref(),
94 + ),
95 + },
96 + decode_options: TextDecodeOptions {
97 + skip_special_tokens: request.skip_special_tokens,
98 + include_stop_str_in_output: request.include_stop_str_in_output,
99 + stop_strings: request.stop.map(|stop| stop.into_vec()),
100 + min_tokens: request.min_tokens.unwrap_or(0),
101 + },
102 + intermediate: request.stream,
103 + priority: request.priority.unwrap_or(0),
104 + cache_salt: request.cache_salt,
105 + add_special_tokens: request.add_special_tokens,
106 + data_parallel_rank: ctx.data_parallel_rank,
107 + };
108 +
109 + Ok(PreparedRequest {
110 + request_id,
111 + response_model: served_model_names.first().cloned().unwrap_or_default(),
112 + include_usage,
113 + text_request,
114 + echo,
115 + return_token_ids: request.return_token_ids.unwrap_or(false),
116 + return_tokens_as_token_ids: request.return_tokens_as_token_ids.unwrap_or(false),
117 + })
118 +}
119 +
120 +#[cfg(test)]
121 +mod tests {
122 + use axum::http::HeaderMap;
123 + use serde_json::json;
124 + use vllm_text::Prompt;
125 +
126 + use super::prepare_completion_request;
127 + use crate::routes::openai::completions::types::CompletionRequest;
128 + use crate::utils::{ResolvedRequestContext, resolve_request_context};
129 +
130 + fn request_context(headers: &HeaderMap, request_id: Option<&str>) -> ResolvedRequestContext {
131 + resolve_request_context(headers, request_id)
132 + }
133 +
134 + fn served(names: &[&str]) -> Vec<String> {
135 + names.iter().map(|s| s.to_string()).collect()
136 + }
137 +
138 + fn base_request_json() -> serde_json::Value {
139 + json!({
140 + "model": "Qwen/Qwen1.5-0.5B-Chat",
141 + "prompt": "hello",
142 + "stream": true
143 + })
144 + }
145 +
146 + #[test]
147 + fn completion_http_request_deserializes_text_prompt() {
148 + let request: CompletionRequest =
149 + serde_json::from_value(base_request_json()).expect("parse request");
150 +
151 + assert_eq!(request.prompt, Prompt::Text("hello".to_string()));
152 + assert_eq!(request.model, "Qwen/Qwen1.5-0.5B-Chat");
153 + }
154 +
155 + #[test]
156 + fn completion_http_request_deserializes_token_id_prompt() {
157 + let request: CompletionRequest = serde_json::from_value(json!({
158 + "model": "Qwen/Qwen1.5-0.5B-Chat",
159 + "prompt": [11, 22, 33],
160 + "stream": true,
161 + "ignore_eos": true,
162 + "max_tokens": 7
163 + }))
164 + .expect("parse request");
165 +
166 + assert_eq!(request.prompt, Prompt::TokenIds(vec![11, 22, 33]));
167 + assert_eq!(request.max_tokens, Some(7));
168 + assert!(request.ignore_eos);
169 + }
170 +
171 + #[test]
172 + fn prepare_completion_request_maps_sampling_fields() {
173 + let request: CompletionRequest = serde_json::from_value(json!({
174 + "model": "Qwen/Qwen1.5-0.5B-Chat",
175 + "prompt": [11, 22, 33],
176 + "stream": true,
177 + "stream_options": {"include_usage": true},
178 + "max_tokens": 7,
179 + "logprobs": 2,
180 + "top_p": 0.9,
181 + "top_k": 42,
182 + "min_p": 0.1,
183 + "frequency_penalty": 0.2,
184 + "presence_penalty": 0.3,
185 + "repetition_penalty": 1.1,
186 + "ignore_eos": true,
187 + "skip_special_tokens": false
188 + }))
189 + .expect("parse request");
190 +
191 + let prepared = prepare_completion_request(
192 + request,
193 + &served(&["Qwen/Qwen1.5-0.5B-Chat"]),
194 + ResolvedRequestContext::default(),
195 + )
196 + .expect("prepare");
197 +
198 + assert!(prepared.include_usage);
199 + assert_eq!(
200 + prepared.text_request.prompt,
201 + Prompt::TokenIds(vec![11, 22, 33])
202 + );
203 + assert_eq!(prepared.text_request.sampling_params.max_tokens, Some(7));
204 + assert_eq!(prepared.text_request.sampling_params.logprobs, Some(2));
205 + assert_eq!(prepared.text_request.sampling_params.top_p, Some(0.9));
206 + assert_eq!(prepared.text_request.sampling_params.top_k, Some(42));
207 + assert_eq!(prepared.text_request.sampling_params.min_p, Some(0.1));
208 + assert_eq!(
209 + prepared.text_request.sampling_params.frequency_penalty,
210 + Some(0.2)
211 + );
212 + assert_eq!(
213 + prepared.text_request.sampling_params.presence_penalty,
214 + Some(0.3)
215 + );
216 + assert_eq!(
217 + prepared.text_request.sampling_params.repetition_penalty,
218 + Some(1.1)
219 + );
220 + assert!(prepared.text_request.sampling_params.ignore_eos);
221 + assert!(!prepared.text_request.decode_options.skip_special_tokens);
222 + }
223 +
224 + #[test]
225 + fn prepare_completion_request_accepts_text_echo() {
226 + let request: CompletionRequest = serde_json::from_value(json!({
227 + "model": "Qwen/Qwen1.5-0.5B-Chat",
228 + "prompt": "hello",
229 + "stream": true,
230 + "echo": true,
231 + "max_tokens": 7
232 + }))
233 + .expect("parse request");
234 +
235 + let prepared = prepare_completion_request(
236 + request,
237 + &served(&["Qwen/Qwen1.5-0.5B-Chat"]),
238 + ResolvedRequestContext::default(),
239 + )
240 + .expect("prepare");
241 +
242 + assert_eq!(prepared.echo, Some("hello".to_string()));
243 + assert_eq!(prepared.text_request.sampling_params.max_tokens, Some(7));
244 + }
245 +
246 + #[test]
247 + fn prepare_completion_request_enables_prompt_logprobs_for_non_stream_echo() {
248 + let request: CompletionRequest = serde_json::from_value(json!({
249 + "model": "Qwen/Qwen1.5-0.5B-Chat",
250 + "prompt": "hello",
251 + "echo": true,
252 + "stream": false,
253 + "logprobs": 3
254 + }))
255 + .expect("parse request");
256 +
257 + let prepared = prepare_completion_request(
258 + request,
259 + &served(&["Qwen/Qwen1.5-0.5B-Chat"]),
260 + ResolvedRequestContext::default(),
261 + )
262 + .expect("prepare");
263 +
264 + assert_eq!(prepared.text_request.sampling_params.logprobs, Some(3));
265 + assert_eq!(
266 + prepared.text_request.sampling_params.prompt_logprobs,
267 + Some(3)
268 + );
269 + }
270 +
271 + #[test]
272 + fn prepare_completion_request_rejects_token_id_prompt_echo() {
273 + let request: CompletionRequest = serde_json::from_value(json!({
274 + "model": "Qwen/Qwen1.5-0.5B-Chat",
275 + "prompt": [11, 22, 33],
276 + "stream": true,
277 + "echo": true
278 + }))
279 + .expect("parse request");
280 +
281 + assert!(
282 + prepare_completion_request(
283 + request,
284 + &served(&["Qwen/Qwen1.5-0.5B-Chat"]),
285 + ResolvedRequestContext::default(),
286 + )
287 + .is_err()
288 + );
289 + }
290 +
291 + #[test]
292 + fn prepare_completion_request_accepts_logprobs_fields() {
293 + let request: CompletionRequest = serde_json::from_value(json!({
294 + "model": "Qwen/Qwen1.5-0.5B-Chat",
295 + "prompt": "hello",
296 + "stream": false,
297 + "logprobs": 1,
298 + "prompt_logprobs": 2
299 + }))
300 + .expect("parse request");
301 +
302 + let prepared = prepare_completion_request(
303 + request,
304 + &served(&["Qwen/Qwen1.5-0.5B-Chat"]),
305 + ResolvedRequestContext::default(),
306 + )
307 + .expect("prepare");
308 + assert_eq!(prepared.text_request.sampling_params.logprobs, Some(1));
309 + assert_eq!(
310 + prepared.text_request.sampling_params.prompt_logprobs,
311 + Some(2)
312 + );
313 + }
314 +
315 + #[test]
316 + fn prepare_completion_request_threads_data_parallel_rank() {
317 + let request: CompletionRequest = serde_json::from_value(json!({
318 + "model": "Qwen/Qwen1.5-0.5B-Chat",
319 + "prompt": "hello",
320 + "stream": false,
321 + }))
322 + .expect("parse request");
323 +
324 + let mut headers = HeaderMap::new();
325 + headers.insert("X-data-parallel-rank", "3".parse().unwrap());
326 + let prepared = prepare_completion_request(
327 + request,
328 + &served(&["Qwen/Qwen1.5-0.5B-Chat"]),
329 + request_context(&headers, None),
330 + )
331 + .expect("prepare");
332 + assert_eq!(prepared.text_request.data_parallel_rank, Some(3));
333 + }
334 +
335 + #[test]
336 + fn prepare_completion_request_leaves_data_parallel_rank_none_when_absent() {
337 + let request: CompletionRequest = serde_json::from_value(json!({
338 + "model": "Qwen/Qwen1.5-0.5B-Chat",
339 + "prompt": "hello",
340 + "stream": false,
341 + }))
342 + .expect("parse request");
343 +
344 + let prepared = prepare_completion_request(
345 + request,
346 + &served(&["Qwen/Qwen1.5-0.5B-Chat"]),
347 + ResolvedRequestContext::default(),
348 + )
349 + .expect("prepare");
350 + assert_eq!(prepared.text_request.data_parallel_rank, None);
351 + }
352 +}
rust/src/server/src/routes/openai/completions/types.rs new
+253
@@ -0,0 +1,253 @@
1 +use std::collections::HashMap;
2 +
3 +use serde::{Deserialize, Serialize};
4 +use serde_json::{Map, Value};
5 +use validator::Validate;
6 +use vllm_text::Prompt;
7 +
8 +use crate::routes::openai::utils::types::{
9 + LogProbs, Normalizable, StreamOptions, StringOrArray, Usage, default_true, validate_stop,
10 +};
11 +
12 +/// Serde default for `CompletionRequest::max_tokens`, matching the Python vLLM
13 +/// / OpenAI default.
14 +fn default_completion_max_tokens() -> Option<u32> {
15 + Some(16)
16 +}
17 +
18 +/// vLLM-compatible request type for the Completions API.
19 +///
20 +/// Mirrors the Python vLLM `CompletionRequest` class. The local copy keeps the
21 +/// request type route-owned so we can accept token-id prompts via
22 +/// [`vllm_text::Prompt`] and add vLLM-only fields directly instead of layering
23 +/// wrapper deserializers on top.
24 +#[serde_with::skip_serializing_none]
25 +#[derive(Debug, Clone, Deserialize, Serialize, Validate)]
26 +pub struct CompletionRequest {
27 + // -------- Standard OpenAI API Parameters --------
28 + /// ID of the model to use
29 + pub model: String,
30 +
31 + /// The prompt(s) to generate completions for.
32 + ///
33 + /// We use [`vllm_text::Prompt`] here to support token-id input.
34 + pub prompt: Prompt,
35 +
36 + /// Echo back the prompt in addition to the completion
37 + #[serde(default)]
38 + pub echo: bool,
39 +
40 + /// Number between -2.0 and 2.0. Positive values penalize new tokens based
41 + /// on their existing frequency in the text so far
42 + pub frequency_penalty: Option<f32>,
43 +
44 + /// Modify the likelihood of specified tokens appearing in the completion
45 + pub logit_bias: Option<HashMap<String, f32>>,
46 +
47 + /// Include the log probabilities on the logprobs most likely tokens
48 + pub logprobs: Option<u32>,
49 +
50 + /// The maximum number of tokens to generate (defaults to 16 when absent,
51 + /// matching the Python vLLM / OpenAI API convention)
52 + #[serde(default = "default_completion_max_tokens")]
53 + pub max_tokens: Option<u32>,
54 +
55 + /// How many completions to generate for each prompt
56 + pub n: Option<u32>,
57 +
58 + /// Number between -2.0 and 2.0. Positive values penalize new tokens based
59 + /// on whether they appear in the text so far
60 + pub presence_penalty: Option<f32>,
61 +
62 + /// If specified, our system will make a best effort to sample
63 + /// deterministically
64 + pub seed: Option<i64>,
65 +
66 + /// Up to 4 sequences where the API will stop generating further tokens
67 + #[validate(custom(function = "validate_stop"))]
68 + pub stop: Option<StringOrArray>,
69 +
70 + /// Whether to stream back partial progress
71 + #[serde(default)]
72 + pub stream: bool,
73 +
74 + /// The suffix that comes after a completion of inserted text
75 + pub suffix: Option<String>,
76 +
77 + /// What sampling temperature to use, between 0 and 2
78 + pub temperature: Option<f32>,
79 +
80 + /// An alternative to sampling with temperature (nucleus sampling)
81 + pub top_p: Option<f32>,
82 +
83 + /// A unique identifier representing your end-user
84 + pub user: Option<String>,
85 +
86 + // -------- vLLM Sampling Parameters --------
87 + /// Options for streaming response
88 + pub stream_options: Option<StreamOptions>,
89 +
90 + /// Use beam search instead of sampling
91 + #[serde(default)]
92 + pub use_beam_search: bool,
93 +
94 + /// Top-k sampling parameter
95 + pub top_k: Option<u32>,
96 +
97 + /// Min-p nucleus sampling parameter
98 + pub min_p: Option<f32>,
99 +
100 + /// Repetition penalty for reducing repetitive text
101 + pub repetition_penalty: Option<f32>,
102 +
103 + /// Length penalty for beam search
104 + pub length_penalty: Option<f32>,
105 +
106 + /// Specific token IDs to use as stop conditions
107 + pub stop_token_ids: Option<Vec<u32>>,
108 +
109 + /// Include stop string in output
110 + #[serde(default)]
111 + pub include_stop_str_in_output: bool,
112 +
113 + /// Ignore end-of-sequence tokens during generation
114 + #[serde(default)]
115 + pub ignore_eos: bool,
116 +
117 + /// Minimum number of tokens to generate
118 + pub min_tokens: Option<u32>,
119 +
120 + /// Skip special tokens during detokenization
121 + #[serde(default = "default_true")]
122 + pub skip_special_tokens: bool,
123 +
124 + /// Add spaces between special tokens during detokenization
125 + #[serde(default = "default_true")]
126 + pub spaces_between_special_tokens: bool,
127 +
128 + /// Truncate prompt tokens to this length
129 + pub truncate_prompt_tokens: Option<i64>,
130 +
131 + /// Restrict output to these token IDs only
132 + pub allowed_token_ids: Option<Vec<u32>>,
133 +
134 + /// Number of prompt logprobs to return
135 + pub prompt_logprobs: Option<i32>,
136 +
137 + // -------- Extra vLLM Parameters --------
138 + /// Whether to add special tokens (e.g. BOS) to the prompt
139 + #[serde(default = "default_true")]
140 + pub add_special_tokens: bool,
141 +
142 + /// Format specification for structured output (JSON mode, JSON schema,
143 + /// etc.)
144 + pub response_format: Option<Value>,
145 +
146 + /// Additional kwargs for structured outputs
147 + pub structured_outputs: Option<Value>,
148 +
149 + /// Request scheduling priority (lower means earlier; default 0)
150 + pub priority: Option<i32>,
151 +
152 + /// External request ID used for response correlation.
153 + pub request_id: Option<String>,
154 +
155 + /// Tokens represented as strings of the form 'token_id:{token_id}' in
156 + /// logprobs
157 + pub return_tokens_as_token_ids: Option<bool>,
158 +
159 + /// Include token IDs alongside generated text
160 + pub return_token_ids: Option<bool>,
161 +
162 + /// Salt for prefix cache isolation in multi-user environments
163 + pub cache_salt: Option<String>,
164 +
165 + /// KV transfer parameters for disaggregated serving
166 + pub kv_transfer_params: Option<HashMap<String, Value>>,
167 +
168 + /// Additional request parameters with string or numeric values for custom
169 + /// extensions
170 + pub vllm_xargs: Option<HashMap<String, Value>>,
171 +
172 + /// Additional fields
173 + #[serde(flatten)]
174 + pub other: Map<String, Value>,
175 +}
176 +
177 +impl Normalizable for CompletionRequest {}
178 +
179 +/// Mirrors the Python vLLM `CompletionResponse` class.
180 +#[serde_with::skip_serializing_none]
181 +#[derive(Debug, Clone, Serialize)]
182 +pub(super) struct CompletionResponse {
183 + pub id: String,
184 + pub object: String,
185 + pub created: u64,
186 + pub model: String,
187 + pub choices: Vec<CompletionChoice>,
188 + pub usage: Option<Usage>,
189 + pub system_fingerprint: Option<String>,
190 + pub kv_transfer_params: Option<Value>,
191 +}
192 +
193 +/// Mirrors the Python vLLM `CompletionResponseChoice` class.
194 +#[serde_with::skip_serializing_none]
195 +#[derive(Debug, Clone, Serialize)]
196 +pub(super) struct CompletionChoice {
197 + pub index: u32,
198 + pub text: String,
199 + pub logprobs: Option<LogProbs>,
200 + pub finish_reason: Option<String>,
201 + pub stop_reason: Option<Value>,
202 + pub prompt_logprobs: Option<Vec<Option<HashMap<String, f32>>>>,
203 + pub token_ids: Option<Vec<u32>>,
204 + pub prompt_token_ids: Option<Vec<u32>>,
205 +}
206 +
207 +/// Mirrors the Python vLLM `CompletionStreamResponse` class.
208 +#[serde_with::skip_serializing_none]
209 +#[derive(Debug, Clone, Serialize)]
210 +pub(super) struct CompletionStreamResponse {
211 + pub id: String,
212 + pub object: String,
213 + pub created: u64,
214 + pub model: String,
215 + pub choices: Vec<CompletionStreamChoice>,
216 + pub usage: Option<Usage>,
217 +}
218 +
219 +impl CompletionStreamResponse {
220 + /// Create a stream response with the standard envelope fields pre-filled.
221 + pub fn new(id: &str, model: &str, created: u64) -> Self {
222 + Self {
223 + id: id.to_string(),
224 + object: "text_completion".to_string(),
225 + created,
226 + model: model.to_string(),
227 + choices: Vec::new(),
228 + usage: None,
229 + }
230 + }
231 +}
232 +
233 +/// Mirrors the Python vLLM `CompletionResponseStreamChoice` class.
234 +#[serde_with::skip_serializing_none]
235 +#[derive(Debug, Clone, Default, Serialize)]
236 +pub(super) struct CompletionStreamChoice {
237 + pub index: u32,
238 + pub text: String,
239 + pub logprobs: Option<LogProbs>,
240 + pub finish_reason: Option<String>,
241 + pub stop_reason: Option<Value>,
242 + pub token_ids: Option<Vec<u32>>,
243 + pub prompt_token_ids: Option<Vec<u32>>,
244 +}
245 +
246 +#[derive(Debug, Clone, Serialize)]
247 +#[serde(untagged)]
248 +pub(super) enum CompletionSseChunk {
249 + /// Ordinary OpenAI completions delta/final chunk.
250 + Chunk(CompletionStreamResponse),
251 + /// Final usage chunk emitted before `[DONE]` when `include_usage=true`.
252 + Usage(CompletionStreamResponse),
253 +}
rust/src/server/src/routes/openai/completions/validate.rs new
+178
@@ -0,0 +1,178 @@
1 +use vllm_text::Prompt;
2 +
3 +use super::types::CompletionRequest;
4 +use crate::error::{ApiError, bail_invalid_request};
5 +
6 +/// Enforce the minimal compatibility contract for the Rust OpenAI server.
7 +pub(super) fn validate_request_compat(
8 + request: &CompletionRequest,
9 + served_model_names: &[String],
10 +) -> Result<(), ApiError> {
11 + // This path is intentionally scoped to the minimum surface needed by
12 + // `vllm-bench` random workload compatibility, so unsupported legacy
13 + // completions features fail early here.
14 + if !served_model_names.iter().any(|n| n == &request.model) {
15 + return Err(ApiError::model_not_found(request.model.clone()));
16 + }
17 +
18 + if request.stream_options.is_some() && !request.stream {
19 + bail_invalid_request!(
20 + param = "stream_options",
21 + "stream_options are only supported when stream=true."
22 + );
23 + }
24 +
25 + if request.n.unwrap_or(1) > 1 {
26 + bail_invalid_request!(param = "n", "Only n=1 is supported.");
27 + }
28 +
29 + if request.max_tokens == Some(0) {
30 + bail_invalid_request!(param = "max_tokens", "max_tokens must be greater than 0.");
31 + }
32 +
33 + if request.echo && matches!(request.prompt, Prompt::TokenIds(_)) {
34 + bail_invalid_request!(
35 + param = "echo",
36 + "echo is not supported with token-ID prompts."
37 + );
38 + }
39 +
40 + if request.suffix.is_some() {
41 + bail_invalid_request!(param = "suffix", "suffix is not supported.");
42 + }
43 +
44 + if let Some(logprobs) = request.logprobs
45 + && logprobs > i32::MAX as u32
46 + {
47 + bail_invalid_request!(
48 + param = "logprobs",
49 + "`logprobs` must fit within a signed 32-bit integer."
50 + );
51 + }
52 +
53 + if let Some(prompt_logprobs) = request.prompt_logprobs {
54 + if request.stream && (prompt_logprobs > 0 || prompt_logprobs == -1) {
55 + bail_invalid_request!(
56 + param = "prompt_logprobs",
57 + "`prompt_logprobs` are not available when `stream=true`."
58 + );
59 + }
60 +
61 + if prompt_logprobs < 0 && prompt_logprobs != -1 {
62 + bail_invalid_request!(
63 + param = "prompt_logprobs",
64 + "`prompt_logprobs` must be a non-negative value or -1."
65 + );
66 + }
67 + }
68 +
69 + if request.use_beam_search {
70 + bail_invalid_request!(
71 + param = "use_beam_search",
72 + "use_beam_search is not supported."
73 + );
74 + }
75 +
76 + // ---- Reject parameters that are accepted for deserialization but not yet
77 + // implemented ----
78 +
79 + if request.length_penalty.is_some() {
80 + bail_invalid_request!(param = "length_penalty", "length_penalty is not supported.");
81 + }
82 + if !request.spaces_between_special_tokens {
83 + bail_invalid_request!(
84 + param = "spaces_between_special_tokens",
85 + "spaces_between_special_tokens is not supported."
86 + );
87 + }
88 + if request.truncate_prompt_tokens.is_some() {
89 + bail_invalid_request!(
90 + param = "truncate_prompt_tokens",
91 + "truncate_prompt_tokens is not supported."
92 + );
93 + }
94 +
95 + if let Some(options) = &request.stream_options
96 + && options.continuous_usage_stats.is_some()
97 + {
98 + bail_invalid_request!(
99 + param = "stream_options",
100 + "continuous_usage_stats is not supported."
101 + );
102 + }
103 +
104 + Ok(())
105 +}
106 +
107 +#[cfg(test)]
108 +mod tests {
109 + use serde_json::json;
110 +
111 + use super::validate_request_compat;
112 + use crate::routes::openai::completions::types::CompletionRequest;
113 +
114 + fn base_request() -> CompletionRequest {
115 + serde_json::from_value(json!({
116 + "model": "Qwen/Qwen1.5-0.5B-Chat",
117 + "prompt": "hello",
118 + "stream": true,
119 + }))
120 + .expect("parse request")
121 + }
122 +
123 + fn served_names(names: &[&str]) -> Vec<String> {
124 + names.iter().map(|s| s.to_string()).collect()
125 + }
126 +
127 + #[test]
128 + fn validate_request_compat_accepts_logprobs() {
129 + let request = CompletionRequest {
130 + logprobs: Some(1),
131 + ..base_request()
132 + };
133 + assert!(
134 + validate_request_compat(&request, &served_names(&["Qwen/Qwen1.5-0.5B-Chat"])).is_ok()
135 + );
136 + }
137 +
138 + #[test]
139 + fn validate_request_compat_accepts_any_served_name() {
140 + let request = base_request();
141 + assert!(
142 + validate_request_compat(
143 + &request,
144 + &served_names(&["other-alias", "Qwen/Qwen1.5-0.5B-Chat"])
145 + )
146 + .is_ok()
147 + );
148 + }
149 +
150 + #[test]
151 + fn validate_request_compat_rejects_unknown_model() {
152 + let request = base_request();
153 + assert!(validate_request_compat(&request, &served_names(&["other-model"])).is_err());
154 + }
155 +
156 + #[test]
157 + fn validate_request_compat_rejects_streaming_prompt_logprobs() {
158 + let request = CompletionRequest {
159 + prompt_logprobs: Some(1),
160 + ..base_request()
161 + };
162 + assert!(
163 + validate_request_compat(&request, &served_names(&["Qwen/Qwen1.5-0.5B-Chat"])).is_err()
164 + );
165 + }
166 +
167 + #[test]
168 + fn validate_request_compat_accepts_non_stream_prompt_logprobs() {
169 + let request = CompletionRequest {
170 + stream: false,
171 + prompt_logprobs: Some(-1),
172 + ..base_request()
173 + };
174 + assert!(
175 + validate_request_compat(&request, &served_names(&["Qwen/Qwen1.5-0.5B-Chat"])).is_ok()
176 + );
177 + }
178 +}
rust/src/server/src/routes/openai/mod.rs new
+8
@@ -0,0 +1,8 @@
1 +pub mod chat_completions;
2 +mod completions;
3 +mod models;
4 +pub(crate) mod utils;
5 +
6 +pub use chat_completions::chat_completions;
7 +pub use completions::completions;
8 +pub use models::list_models;
rust/src/server/src/routes/openai/models.rs new
+24
@@ -0,0 +1,24 @@
1 +use std::sync::Arc;
2 +
3 +use axum::Json;
4 +use axum::extract::State;
5 +
6 +use crate::routes::openai::utils::types::{ListModelsResponse, ModelObject};
7 +use crate::state::AppState;
8 +
9 +/// Return all configured served model names in OpenAI `list models` format.
10 +pub async fn list_models(State(state): State<Arc<AppState>>) -> Json<ListModelsResponse> {
11 + Json(ListModelsResponse {
12 + object: "list".to_string(),
13 + data: state
14 + .served_model_names()
15 + .iter()
16 + .map(|name| ModelObject {
17 + id: name.clone(),
18 + object: "model".to_string(),
19 + created: 0,
20 + owned_by: "vllm-frontend-rs".to_string(),
21 + })
22 + .collect(),
23 + })
24 +}
rust/src/server/src/routes/openai/utils/logprobs.rs new
+256
@@ -0,0 +1,256 @@
1 +use std::collections::HashMap;
2 +
3 +use itertools::Itertools as _;
4 +use vllm_text::{
5 + CollectedTextOutput, DecodedLogprobs, DecodedPositionLogprobs, DecodedPromptLogprobs,
6 + DecodedTokenLogprob,
7 +};
8 +
9 +use super::types::{ChatLogProbs, ChatLogProbsContent, LogProbs, TopLogProb};
10 +use crate::error::{ApiError, server_error};
11 +
12 +/// Convert decoded token-position logprobs into the OpenAI completions
13 +/// `logprobs` shape.
14 +pub fn decoded_logprobs_to_openai(
15 + logprobs: &DecodedLogprobs,
16 + initial_text_offset: u32,
17 + return_tokens_as_token_ids: bool,
18 +) -> Result<LogProbs, ApiError> {
19 + let mut text_offset = Vec::with_capacity(logprobs.positions.len());
20 + let mut token_logprobs = Vec::with_capacity(logprobs.positions.len());
21 + let mut tokens = Vec::with_capacity(logprobs.positions.len());
22 + let mut top_logprobs = Vec::with_capacity(logprobs.positions.len());
23 + let mut current_offset = initial_text_offset;
24 +
25 + for position in &logprobs.positions {
26 + let chosen = position.entries.first().ok_or_else(|| {
27 + server_error!("decoded logprobs position unexpectedly had no token candidates")
28 + })?;
29 +
30 + let token_str = format_token(chosen, return_tokens_as_token_ids);
31 + text_offset.push(current_offset);
32 + token_logprobs.push(Some(clamp_logprob(chosen.logprob)));
33 + current_offset = current_offset.saturating_add(text_len(&token_str));
34 + tokens.push(token_str);
35 + top_logprobs.push(Some(position_top_logprobs_map(
36 + position,
37 + return_tokens_as_token_ids,
38 + )));
39 + }
40 +
41 + Ok(LogProbs {
42 + tokens,
43 + token_logprobs,
44 + top_logprobs,
45 + text_offset,
46 + })
47 +}
48 +
49 +/// Convert decoded prompt logprobs into the OpenAI completions `logprobs`
50 +/// shape.
51 +///
52 +/// The first prompt token is included with `None` logprob metadata, matching
53 +/// Python vLLM's echoed completions behavior.
54 +pub fn decoded_prompt_logprobs_to_openai(
55 + prompt_logprobs: &DecodedPromptLogprobs,
56 + initial_text_offset: u32,
57 + return_tokens_as_token_ids: bool,
58 +) -> Result<LogProbs, ApiError> {
59 + let mut text_offset = Vec::with_capacity(prompt_logprobs.scored_positions.len() + 1);
60 + let mut token_logprobs = Vec::with_capacity(prompt_logprobs.scored_positions.len() + 1);
61 + let mut tokens = Vec::with_capacity(prompt_logprobs.scored_positions.len() + 1);
62 + let mut top_logprobs = Vec::with_capacity(prompt_logprobs.scored_positions.len() + 1);
63 + let mut current_offset = initial_text_offset;
64 +
65 + let first_token_str = if return_tokens_as_token_ids {
66 + format!("token_id:{}", prompt_logprobs.first_token_id)
67 + } else {
68 + prompt_logprobs.first_token.clone()
69 + };
70 + text_offset.push(current_offset);
71 + token_logprobs.push(None);
72 + current_offset = current_offset.saturating_add(text_len(&first_token_str));
73 + tokens.push(first_token_str);
74 + top_logprobs.push(None);
75 +
76 + for position in &prompt_logprobs.scored_positions {
77 + let chosen = position.entries.first().ok_or_else(|| {
78 + server_error!("decoded prompt logprobs position unexpectedly had no token candidates")
79 + })?;
80 +
81 + let token_str = format_token(chosen, return_tokens_as_token_ids);
82 + text_offset.push(current_offset);
83 + token_logprobs.push(Some(clamp_logprob(chosen.logprob)));
84 + current_offset = current_offset.saturating_add(text_len(&token_str));
85 + tokens.push(token_str);
86 + top_logprobs.push(Some(position_top_logprobs_map(
87 + position,
88 + return_tokens_as_token_ids,
89 + )));
90 + }
91 +
92 + Ok(LogProbs {
93 + tokens,
94 + token_logprobs,
95 + top_logprobs,
96 + text_offset,
97 + })
98 +}
99 +
100 +/// Convert decoded prompt logprobs into the vLLM-style prompt-logprobs response
101 +/// shape.
102 +pub fn decoded_prompt_logprobs_to_maps(
103 + prompt_logprobs: &DecodedPromptLogprobs,
104 + return_tokens_as_token_ids: bool,
105 +) -> Vec<Option<HashMap<String, f32>>> {
106 + std::iter::once(None)
107 + .chain(prompt_logprobs.scored_positions.iter().map(|position| {
108 + Some(position_top_logprobs_map(
109 + position,
110 + return_tokens_as_token_ids,
111 + ))
112 + }))
113 + .collect()
114 +}
115 +
116 +/// Convert decoded token-position logprobs into the OpenAI chat `logprobs`
117 +/// shape.
118 +pub fn decoded_logprobs_to_openai_chat(
119 + logprobs: &DecodedLogprobs,
120 + return_tokens_as_token_ids: bool,
121 +) -> Result<ChatLogProbs, ApiError> {
122 + let content = logprobs
123 + .positions
124 + .iter()
125 + .map(|pos| position_to_chat_logprobs_content(pos, return_tokens_as_token_ids))
126 + .try_collect()?;
127 +
128 + Ok(ChatLogProbs {
129 + content: Some(content),
130 + })
131 +}
132 +
133 +/// Count visible text positions using OpenAI completions' character-offset
134 +/// convention.
135 +pub fn text_len(text: &str) -> u32 {
136 + u32::try_from(text.chars().count()).unwrap_or(u32::MAX)
137 +}
138 +
139 +/// Concatenate two OpenAI-style completion logprobs payloads in token order.
140 +pub fn append_openai_logprobs(mut prefix: LogProbs, suffix: LogProbs) -> LogProbs {
141 + prefix.tokens.extend(suffix.tokens);
142 + prefix.token_logprobs.extend(suffix.token_logprobs);
143 + prefix.top_logprobs.extend(suffix.top_logprobs);
144 + prefix.text_offset.extend(suffix.text_offset);
145 + prefix
146 +}
147 +
148 +/// Build the non-stream completions `logprobs` payload from collected text
149 +/// output.
150 +///
151 +/// When `echoed_prompt` is true, the returned payload matches Python vLLM's
152 +/// echoed completions behavior by concatenating prompt and completion logprobs
153 +/// into one OpenAI `LogProbs` object.
154 +pub fn collected_logprobs_to_openai(
155 + collected: &CollectedTextOutput,
156 + echoed_prompt: bool,
157 + initial_completion_offset: u32,
158 + return_tokens_as_token_ids: bool,
159 +) -> Result<LogProbs, ApiError> {
160 + if echoed_prompt {
161 + let prompt_logprobs = collected.prompt_logprobs.as_ref().ok_or_else(|| {
162 + server_error!(
163 + "echoed completion logprobs require prompt logprobs but generation returned none"
164 + )
165 + })?;
166 + let prompt_logprobs =
167 + decoded_prompt_logprobs_to_openai(prompt_logprobs, 0, return_tokens_as_token_ids)?;
168 + let completion_start = prompt_logprobs
169 + .text_offset
170 + .last()
171 + .zip(prompt_logprobs.tokens.last())
172 + .map(|(&offset, token)| offset.saturating_add(text_len(token)))
173 + .unwrap_or(0);
174 + return match collected.logprobs.as_ref() {
175 + Some(completion_logprobs) => Ok(append_openai_logprobs(
176 + prompt_logprobs,
177 + decoded_logprobs_to_openai(
178 + completion_logprobs,
179 + completion_start,
180 + return_tokens_as_token_ids,
181 + )?,
182 + )),
183 + None => Ok(prompt_logprobs),
184 + };
185 + }
186 +
187 + let completion_logprobs = collected.logprobs.as_ref().ok_or_else(|| {
188 + server_error!("completion response requested logprobs but generation returned none")
189 + })?;
190 + decoded_logprobs_to_openai(
191 + completion_logprobs,
192 + initial_completion_offset,
193 + return_tokens_as_token_ids,
194 + )
195 +}
196 +
197 +/// Format a token entry as either its decoded string or `token_id:{id}`.
198 +fn format_token(entry: &DecodedTokenLogprob, as_token_id: bool) -> String {
199 + if as_token_id {
200 + format!("token_id:{}", entry.token_id)
201 + } else {
202 + entry.token.clone()
203 + }
204 +}
205 +
206 +fn position_top_logprobs_map(
207 + position: &DecodedPositionLogprobs,
208 + return_tokens_as_token_ids: bool,
209 +) -> HashMap<String, f32> {
210 + position
211 + .entries
212 + .iter()
213 + .map(|entry| {
214 + (
215 + format_token(entry, return_tokens_as_token_ids),
216 + clamp_logprob(entry.logprob),
217 + )
218 + })
219 + .collect()
220 +}
221 +
222 +fn position_to_chat_logprobs_content(
223 + position: &DecodedPositionLogprobs,
224 + return_tokens_as_token_ids: bool,
225 +) -> Result<ChatLogProbsContent, ApiError> {
226 + let chosen = position.entries.first().ok_or_else(|| {
227 + server_error!("decoded chat logprobs position unexpectedly had no token candidates")
228 + })?;
229 +
230 + let token_str = format_token(chosen, return_tokens_as_token_ids);
231 + Ok(ChatLogProbsContent {
232 + token: token_str.clone(),
233 + logprob: clamp_logprob(chosen.logprob),
234 + bytes: Some(token_bytes(&token_str)),
235 + top_logprobs: position
236 + .entries
237 + .iter()
238 + .map(|entry| {
239 + let t = format_token(entry, return_tokens_as_token_ids);
240 + TopLogProb {
241 + logprob: clamp_logprob(entry.logprob),
242 + bytes: Some(token_bytes(&t)),
243 + token: t,
244 + }
245 + })
246 + .collect(),
247 + })
248 +}
249 +
250 +fn token_bytes(token: &str) -> Vec<u8> {
251 + token.as_bytes().to_vec()
252 +}
253 +
254 +pub fn clamp_logprob(logprob: f32) -> f32 {
255 + logprob.max(-9999.0)
256 +}
rust/src/server/src/routes/openai/utils/mod.rs new
+4
@@ -0,0 +1,4 @@
1 +pub mod logprobs;
2 +pub mod structured_outputs;
3 +pub mod types;
4 +pub mod validated_json;
rust/src/server/src/routes/openai/utils/structured_outputs.rs new
+132
@@ -0,0 +1,132 @@
1 +use serde::{Deserialize, Serialize};
2 +use serde_json::Value;
3 +use vllm_engine_core_client::protocol::StructuredOutputsParams;
4 +
5 +use crate::error::ApiError;
6 +
7 +/// JSON schema specification nested inside a `json_schema` response format.
8 +///
9 +/// Mirrors the Python vLLM `JsonSchemaResponseFormat` class.
10 +#[serde_with::skip_serializing_none]
11 +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
12 +pub struct JsonSchemaFormat {
13 + pub name: String,
14 + #[serde(default)]
15 + pub description: Option<String>,
16 + /// The actual JSON schema object.
17 + #[serde(alias = "json_schema")]
18 + pub schema: Value,
19 + #[serde(default)]
20 + pub strict: Option<bool>,
21 +}
22 +
23 +/// Supported `response_format` types for chat and completion requests.
24 +///
25 +/// This is our own definition (rather than the `openai-protocol` crate's) so
26 +/// that we can support the vLLM-specific `structural_tag` variant.
27 +///
28 +/// Original Python definitions:
29 +/// <https://github.com/vllm-project/vllm/blob/f22d6e026/vllm/entrypoints/openai/engine/protocol.py#L116-L157>
30 +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
31 +#[serde(tag = "type", rename_all = "snake_case")]
32 +pub enum ResponseFormat {
33 + Text,
34 + JsonObject,
35 + JsonSchema {
36 + json_schema: JsonSchemaFormat,
37 + },
38 + /// vLLM-specific structural tag format. The entire object (including the
39 + /// `type` field) is JSON-serialized and passed as
40 + /// `StructuredOutputsParams.structural_tag`.
41 + ///
42 + /// We capture the payload as a catch-all map so both the legacy
43 + /// (`structures`/`triggers`) and current (`format`) shapes are
44 + /// preserved without needing typed structs.
45 + StructuralTag {
46 + #[serde(flatten)]
47 + extra: serde_json::Map<String, Value>,
48 + },
49 +}
50 +
51 +/// Convert an explicit `structured_outputs` JSON blob into
52 +/// [`StructuredOutputsParams`].
53 +fn deserialize_structured_outputs(
54 + raw: &serde_json::Value,
55 +) -> Result<StructuredOutputsParams, ApiError> {
56 + serde_json::from_value(raw.clone()).map_err(|e| {
57 + ApiError::invalid_request(
58 + format!("invalid structured_outputs: {e}"),
59 + Some("structured_outputs"),
60 + )
61 + })
62 +}
63 +
64 +/// Convert a typed [`ResponseFormat`] and/or raw `structured_outputs` blob into
65 +/// engine-core [`StructuredOutputsParams`].
66 +///
67 +/// Mirrors the Python vLLM conversion in
68 +/// `ChatCompletionRequest.to_sampling_params()`: <https://github.com/vllm-project/vllm/blob/f22d6e026/vllm/entrypoints/openai/chat_completion/protocol.py#L457-L487>
69 +pub fn convert_from_response_format(
70 + response_format: Option<&ResponseFormat>,
71 + structured_outputs: &Option<serde_json::Value>,
72 +) -> Result<Option<StructuredOutputsParams>, ApiError> {
73 + if let Some(raw) = structured_outputs {
74 + return Ok(Some(deserialize_structured_outputs(raw)?));
75 + }
76 +
77 + let Some(fmt) = response_format else {
78 + return Ok(None);
79 + };
80 + match fmt {
81 + ResponseFormat::Text => Ok(None),
82 + ResponseFormat::JsonObject => Ok(Some(StructuredOutputsParams {
83 + json_object: Some(true),
84 + ..Default::default()
85 + })),
86 + ResponseFormat::JsonSchema { json_schema } => Ok(Some(StructuredOutputsParams {
87 + json: Some(json_schema.schema.clone()),
88 + ..Default::default()
89 + })),
90 + ResponseFormat::StructuralTag { .. } => {
91 + // The Python frontend dumps the entire response_format object (including the
92 + // `type` field) as a JSON string for the engine-core backend.
93 + let tag_json = serde_json::to_string(fmt).map_err(|e| {
94 + ApiError::invalid_request(
95 + format!("failed to serialize structural_tag: {e}"),
96 + Some("response_format"),
97 + )
98 + })?;
99 + Ok(Some(StructuredOutputsParams {
100 + structural_tag: Some(tag_json),
101 + ..Default::default()
102 + }))
103 + }
104 + }
105 +}
106 +
107 +/// Convert raw `response_format` and/or `structured_outputs` JSON blobs into
108 +/// engine-core [`StructuredOutputsParams`].
109 +///
110 +/// Used by the completions endpoint which keeps both fields as opaque
111 +/// `serde_json::Value`.
112 +pub fn convert_from_response_format_value(
113 + response_format: &Option<serde_json::Value>,
114 + structured_outputs: &Option<serde_json::Value>,
115 +) -> Result<Option<StructuredOutputsParams>, ApiError> {
116 + if let Some(raw) = structured_outputs {
117 + return Ok(Some(deserialize_structured_outputs(raw)?));
118 + }
119 +
120 + let Some(raw) = response_format else {
121 + return Ok(None);
122 + };
123 +
124 + // Deserialize into our typed enum and delegate.
125 + let fmt: ResponseFormat = serde_json::from_value(raw.clone()).map_err(|e| {
126 + ApiError::invalid_request(
127 + format!("invalid response_format: {e}"),
128 + Some("response_format"),
129 + )
130 + })?;
131 + convert_from_response_format(Some(&fmt), &None)
132 +}
rust/src/server/src/routes/openai/utils/types.rs new
+425
@@ -0,0 +1,425 @@
1 +use std::collections::HashMap;
2 +use std::slice;
3 +
4 +use llm_multimodal::ImageDetail;
5 +use serde::{Deserialize, Serialize};
6 +use serde_json::Value;
7 +
8 +// ============================================================================
9 +// Constants
10 +// ============================================================================
11 +
12 +/// Default model identifier used when no model is specified.
13 +pub const UNKNOWN_MODEL_ID: &str = "unknown";
14 +
15 +// ============================================================================
16 +// Default value helpers
17 +// ============================================================================
18 +
19 +/// Helper function for serde default value (returns true).
20 +pub fn default_true() -> bool {
21 + true
22 +}
23 +
24 +// ============================================================================
25 +// String/Array Utilities
26 +// ============================================================================
27 +
28 +/// A type that can be either a single string or an array of strings.
29 +#[derive(Debug, Clone, PartialEq, Deserialize, Serialize)]
30 +#[serde(untagged)]
31 +pub enum StringOrArray {
32 + String(String),
33 + Array(Vec<String>),
34 +}
35 +
36 +impl StringOrArray {
37 + pub fn as_slice(&self) -> &[String] {
38 + match self {
39 + StringOrArray::String(s) => slice::from_ref(s),
40 + StringOrArray::Array(arr) => arr,
41 + }
42 + }
43 +
44 + #[allow(unused)]
45 + pub fn into_vec(self) -> Vec<String> {
46 + match self {
47 + StringOrArray::String(s) => vec![s],
48 + StringOrArray::Array(arr) => arr,
49 + }
50 + }
51 +}
52 +
53 +/// Validates stop sequences (non-empty strings)
54 +pub fn validate_stop(stop: &StringOrArray) -> Result<(), validator::ValidationError> {
55 + if stop.as_slice().iter().any(|s| s.is_empty()) {
56 + return Err(validator::ValidationError::new(
57 + "stop strings cannot be empty",
58 + ));
59 + }
60 + Ok(())
61 +}
62 +
63 +// ============================================================================
64 +// Validation helpers
65 +// ============================================================================
66 +
67 +/// Validates top_p: 0.0 < top_p <= 1.0.
68 +pub fn validate_top_p_value(top_p: f32) -> Result<(), validator::ValidationError> {
69 + if !(top_p > 0.0 && top_p <= 1.0) {
70 + return Err(validator::ValidationError::new(
71 + "top_p must be in (0, 1] - greater than 0.0 and at most 1.0",
72 + ));
73 + }
74 + Ok(())
75 +}
76 +
77 +// ============================================================================
78 +// Content Parts (for multimodal messages)
79 +// ============================================================================
80 +
81 +#[derive(Debug, Clone, Deserialize, Serialize, PartialEq)]
82 +#[serde(tag = "type")]
83 +pub enum ContentPart {
84 + #[serde(rename = "text")]
85 + Text { text: String },
86 + #[serde(rename = "image_url")]
87 + ImageUrl {
88 + image_url: ImageUrl,
89 + #[serde(skip_serializing_if = "Option::is_none")]
90 + uuid: Option<String>,
91 + },
92 + #[serde(rename = "video_url")]
93 + VideoUrl { video_url: VideoUrl },
94 +}
95 +
96 +#[serde_with::skip_serializing_none]
97 +#[derive(Debug, Clone, Deserialize, Serialize, PartialEq)]
98 +pub struct ImageUrl {
99 + pub url: String,
100 + pub detail: Option<ImageDetail>,
101 +}
102 +
103 +#[derive(Debug, Clone, Deserialize, Serialize, PartialEq)]
104 +pub struct VideoUrl {
105 + pub url: String,
106 +}
107 +
108 +// ============================================================================
109 +// Streaming
110 +// ============================================================================
111 +
112 +/// Mirrors the Python vLLM `StreamOptions` class.
113 +#[derive(Debug, Clone, Deserialize, Serialize)]
114 +pub struct StreamOptions {
115 + pub include_usage: Option<bool>,
116 + pub continuous_usage_stats: Option<bool>,
117 +}
118 +
119 +// ============================================================================
120 +// Tools and Function Calling
121 +// ============================================================================
122 +
123 +#[derive(Debug, Clone, Deserialize, Serialize)]
124 +pub struct Tool {
125 + #[serde(rename = "type")]
126 + pub tool_type: String,
127 + pub function: Function,
128 +}
129 +
130 +#[serde_with::skip_serializing_none]
131 +#[derive(Debug, Clone, Deserialize, Serialize)]
132 +pub struct Function {
133 + pub name: String,
134 + pub description: Option<String>,
135 + pub parameters: Value,
136 + /// Whether to enable strict schema adherence (OpenAI structured outputs).
137 + pub strict: Option<bool>,
138 +}
139 +
140 +#[derive(Debug, Clone, Deserialize, Serialize)]
141 +pub struct ToolCall {
142 + pub id: String,
143 + #[serde(rename = "type")]
144 + pub tool_type: String,
145 + pub function: FunctionCallResponse,
146 +}
147 +
148 +#[derive(Debug, Clone, Deserialize, Serialize)]
149 +pub struct FunctionCallResponse {
150 + pub name: String,
151 + #[serde(default)]
152 + pub arguments: Option<String>,
153 +}
154 +
155 +#[serde_with::skip_serializing_none]
156 +#[derive(Debug, Clone, Deserialize, Serialize)]
157 +pub struct ToolCallDelta {
158 + pub index: u32,
159 + pub id: Option<String>,
160 + #[serde(rename = "type")]
161 + pub tool_type: Option<String>,
162 + pub function: Option<FunctionCallDelta>,
163 +}
164 +
165 +#[serde_with::skip_serializing_none]
166 +#[derive(Debug, Clone, Deserialize, Serialize)]
167 +pub struct FunctionCallDelta {
168 + pub name: Option<String>,
169 + pub arguments: Option<String>,
170 +}
171 +
172 +/// Tool choice value for simple string options.
173 +#[derive(Debug, Clone, Deserialize, Serialize)]
174 +#[serde(rename_all = "snake_case")]
175 +pub enum ToolChoiceValue {
176 + Auto,
177 + Required,
178 + None,
179 +}
180 +
181 +/// Tool choice for the Chat Completion API.
182 +#[derive(Debug, Clone, Deserialize, Serialize)]
183 +#[serde(untagged)]
184 +pub enum ToolChoice {
185 + Value(ToolChoiceValue),
186 + Function {
187 + #[serde(rename = "type")]
188 + tool_type: String,
189 + function: FunctionChoice,
190 + },
191 + AllowedTools {
192 + #[serde(rename = "type")]
193 + tool_type: String,
194 + mode: String,
195 + tools: Vec<ToolReference>,
196 + },
197 +}
198 +
199 +impl Default for ToolChoice {
200 + fn default() -> Self {
201 + Self::Value(ToolChoiceValue::Auto)
202 + }
203 +}
204 +
205 +/// Function choice specification for `ToolChoice::Function`.
206 +#[derive(Debug, Clone, Deserialize, Serialize)]
207 +pub struct FunctionChoice {
208 + pub name: String,
209 +}
210 +
211 +/// Tool reference for `ToolChoice::AllowedTools`.
212 +#[derive(Debug, Clone, Deserialize, Serialize)]
213 +#[serde(tag = "type")]
214 +#[serde(rename_all = "snake_case")]
215 +pub enum ToolReference {
216 + #[serde(rename = "function")]
217 + Function { name: String },
218 + #[serde(rename = "mcp")]
219 + Mcp {
220 + server_label: String,
221 + #[serde(skip_serializing_if = "Option::is_none")]
222 + name: Option<String>,
223 + },
224 + #[serde(rename = "file_search")]
225 + FileSearch,
226 + #[serde(rename = "web_search_preview")]
227 + WebSearchPreview,
228 + #[serde(rename = "computer_use_preview")]
229 + ComputerUsePreview,
230 + #[serde(rename = "code_interpreter")]
231 + CodeInterpreter,
232 + #[serde(rename = "image_generation")]
233 + ImageGeneration,
234 +}
235 +
236 +impl ToolReference {
237 + /// Get a unique identifier for this tool reference.
238 + pub fn identifier(&self) -> String {
239 + match self {
240 + ToolReference::Function { name } => format!("function:{name}"),
241 + ToolReference::Mcp {
242 + server_label,
243 + name: Some(n),
244 + } => format!("mcp:{server_label}:{n}"),
245 + ToolReference::Mcp {
246 + server_label,
247 + name: _,
248 + } => format!("mcp:{server_label}"),
249 + ToolReference::FileSearch => "file_search".to_string(),
250 + ToolReference::WebSearchPreview => "web_search_preview".to_string(),
251 + ToolReference::ComputerUsePreview => "computer_use_preview".to_string(),
252 + ToolReference::CodeInterpreter => "code_interpreter".to_string(),
253 + ToolReference::ImageGeneration => "image_generation".to_string(),
254 + }
255 + }
256 +}
257 +
258 +// ============================================================================
259 +// Chat Messages
260 +// ============================================================================
261 +
262 +#[serde_with::skip_serializing_none]
263 +#[derive(Debug, Clone, Deserialize, Serialize)]
264 +#[serde(tag = "role")]
265 +pub enum ChatMessage {
266 + #[serde(rename = "system")]
267 + System {
268 + content: MessageContent,
269 + name: Option<String>,
270 + },
271 + #[serde(rename = "user")]
272 + User {
273 + content: MessageContent,
274 + name: Option<String>,
275 + },
276 + #[serde(rename = "assistant")]
277 + Assistant {
278 + content: Option<MessageContent>,
279 + name: Option<String>,
280 + tool_calls: Option<Vec<ToolCall>>,
281 + /// Reasoning content for reasoning-capable models.
282 + #[serde(alias = "reasoning_content")]
283 + #[serde(alias = "thinking")]
284 + reasoning: Option<String>,
285 + },
286 + #[serde(rename = "tool")]
287 + Tool {
288 + content: MessageContent,
289 + tool_call_id: String,
290 + },
291 + #[serde(rename = "function")]
292 + Function { content: String, name: String },
293 + #[serde(rename = "developer")]
294 + Developer {
295 + content: MessageContent,
296 + tools: Option<Vec<Tool>>,
297 + name: Option<String>,
298 + },
299 +}
300 +
301 +#[derive(Debug, Clone, Deserialize, Serialize, PartialEq)]
302 +#[serde(untagged)]
303 +pub enum MessageContent {
304 + Text(String),
305 + Parts(Vec<ContentPart>),
306 +}
307 +
308 +// ============================================================================
309 +// Usage and Logging
310 +// ============================================================================
311 +
312 +/// Mirrors the Python vLLM `UsageInfo` class.
313 +#[serde_with::skip_serializing_none]
314 +#[derive(Debug, Clone, Serialize)]
315 +pub struct Usage {
316 + pub prompt_tokens: u32,
317 + pub total_tokens: u32,
318 + pub completion_tokens: Option<u32>,
319 + pub prompt_tokens_details: Option<PromptTokenUsageInfo>,
320 +}
321 +
322 +impl Usage {
323 + /// Create a Usage from prompt and completion token counts.
324 + pub fn from_counts(prompt_tokens: u32, completion_tokens: u32) -> Self {
325 + Self {
326 + prompt_tokens,
327 + total_tokens: prompt_tokens + completion_tokens,
328 + completion_tokens: Some(completion_tokens),
329 + prompt_tokens_details: None,
330 + }
331 + }
332 +}
333 +
334 +/// Mirrors the Python vLLM `PromptTokenUsageInfo` class.
335 +#[serde_with::skip_serializing_none]
336 +#[derive(Debug, Clone, Serialize)]
337 +pub struct PromptTokenUsageInfo {
338 + pub cached_tokens: Option<u32>,
339 +}
340 +
341 +/// OpenAI completions-style logprobs.
342 +#[derive(Debug, Clone, Deserialize, Serialize)]
343 +pub struct LogProbs {
344 + pub tokens: Vec<String>,
345 + pub token_logprobs: Vec<Option<f32>>,
346 + pub top_logprobs: Vec<Option<HashMap<String, f32>>>,
347 + pub text_offset: Vec<u32>,
348 +}
349 +
350 +/// Mirrors the Python vLLM `ChatCompletionLogProbs` class.
351 +#[serde_with::skip_serializing_none]
352 +#[derive(Debug, Clone, Serialize)]
353 +pub struct ChatLogProbs {
354 + pub content: Option<Vec<ChatLogProbsContent>>,
355 +}
356 +
357 +/// Mirrors the Python vLLM `ChatCompletionLogProbsContent` class.
358 +#[serde_with::skip_serializing_none]
359 +#[derive(Debug, Clone, Serialize)]
360 +pub struct ChatLogProbsContent {
361 + pub token: String,
362 + pub logprob: f32,
363 + pub bytes: Option<Vec<u8>>,
364 + pub top_logprobs: Vec<TopLogProb>,
365 +}
366 +
367 +/// Mirrors the Python vLLM `ChatCompletionLogProb` class.
368 +#[serde_with::skip_serializing_none]
369 +#[derive(Debug, Clone, Serialize)]
370 +pub struct TopLogProb {
371 + pub token: String,
372 + pub logprob: f32,
373 + pub bytes: Option<Vec<u8>>,
374 +}
375 +
376 +// ============================================================================
377 +// Error Types
378 +// ============================================================================
379 +
380 +#[derive(Debug, Clone, Deserialize, Serialize)]
381 +pub struct ErrorResponse {
382 + pub error: ErrorDetail,
383 +}
384 +
385 +#[serde_with::skip_serializing_none]
386 +#[derive(Debug, Clone, Deserialize, Serialize)]
387 +pub struct ErrorDetail {
388 + pub message: String,
389 + #[serde(rename = "type")]
390 + pub error_type: String,
391 + pub param: Option<String>,
392 + pub code: Option<String>,
393 +}
394 +
395 +// ============================================================================
396 +// Model types
397 +// ============================================================================
398 +
399 +/// A single model entry in the `/v1/models` response.
400 +#[derive(Debug, Clone, Serialize, Deserialize)]
401 +pub struct ModelObject {
402 + pub id: String,
403 + pub object: String,
404 + pub created: i64,
405 + pub owned_by: String,
406 +}
407 +
408 +/// Response body for `GET /v1/models`.
409 +#[derive(Debug, Clone, Serialize, Deserialize)]
410 +pub struct ListModelsResponse {
411 + pub object: String,
412 + pub data: Vec<ModelObject>,
413 +}
414 +
415 +// ============================================================================
416 +// Normalizable trait
417 +// ============================================================================
418 +
419 +/// Trait for request types that need post-deserialization normalization.
420 +pub trait Normalizable {
421 + /// Normalize the request by applying defaults and transformations.
422 + fn normalize(&mut self) {
423 + // Default: no-op
424 + }
425 +}
rust/src/server/src/routes/openai/utils/validated_json.rs new
+55
@@ -0,0 +1,55 @@
1 +//! Validated JSON extractor for automatic request validation.
2 +//! Variation of https://github.com/lightseekorg/smg/blob/main/crates/protocols/src/validated.rs
3 +
4 +use axum::Json;
5 +use axum::extract::rejection::JsonRejection;
6 +use axum::extract::{FromRequest, Request};
7 +use serde::de::DeserializeOwned;
8 +use validator::Validate;
9 +
10 +use super::types::Normalizable;
11 +use crate::error::{ApiError, invalid_request};
12 +
13 +/// A JSON extractor that automatically validates and normalizes the request
14 +/// body.
15 +///
16 +/// This extractor deserializes the request body and automatically calls
17 +/// `.validate()` on types that implement the `Validate` trait. If validation
18 +/// fails, it returns [`ApiError::InvalidRequest`] with details about the
19 +/// validation errors.
20 +pub struct ValidatedJson<T>(pub T);
21 +
22 +impl<S, T> FromRequest<S> for ValidatedJson<T>
23 +where
24 + T: DeserializeOwned + Validate + Normalizable + Send,
25 + S: Send + Sync,
26 +{
27 + type Rejection = ApiError;
28 +
29 + async fn from_request(req: Request, state: &S) -> Result<Self, Self::Rejection> {
30 + let Json(mut data) = Json::<T>::from_request(req, state)
31 + .await
32 + .map_err(|err: JsonRejection| ApiError::json_parse_error(err.body_text()))?;
33 +
34 + data.normalize();
35 +
36 + data.validate()
37 + .map_err(|validation_errors| invalid_request!("{}", validation_errors))?;
38 +
39 + Ok(ValidatedJson(data))
40 + }
41 +}
42 +
43 +impl<T> std::ops::Deref for ValidatedJson<T> {
44 + type Target = T;
45 +
46 + fn deref(&self) -> &Self::Target {
47 + &self.0
48 + }
49 +}
50 +
51 +impl<T> std::ops::DerefMut for ValidatedJson<T> {
52 + fn deref_mut(&mut self) -> &mut Self::Target {
53 + &mut self.0
54 + }
55 +}
rust/src/server/src/routes/sleep.rs new
+78
@@ -0,0 +1,78 @@
1 +use std::sync::Arc;
2 +
3 +use axum::Json;
4 +use axum::extract::{Query, State};
5 +use axum::http::StatusCode;
6 +use serde::{Deserialize, Serialize};
7 +
8 +use crate::error::ApiError;
9 +use crate::state::AppState;
10 +use crate::utils::utility_call_error;
11 +
12 +#[derive(Serialize)]
13 +pub(crate) struct IsSleepingResponse {
14 + is_sleeping: bool,
15 +}
16 +
17 +#[derive(Debug, Deserialize)]
18 +pub(crate) struct SleepParams {
19 + #[serde(default = "default_sleep_level")]
20 + level: u32,
21 + #[serde(default = "default_sleep_mode")]
22 + mode: String,
23 +}
24 +
25 +#[derive(Debug, Default, Deserialize)]
26 +pub(crate) struct WakeUpParams {
27 + #[serde(default)]
28 + tags: Option<Vec<String>>,
29 +}
30 +
31 +const fn default_sleep_level() -> u32 {
32 + 1
33 +}
34 +
35 +fn default_sleep_mode() -> String {
36 + "abort".to_string()
37 +}
38 +
39 +/// Put the engine to sleep.
40 +pub async fn sleep(
41 + State(state): State<Arc<AppState>>,
42 + Query(params): Query<SleepParams>,
43 +) -> Result<StatusCode, ApiError> {
44 + state
45 + .engine_core_client()
46 + .sleep(params.level, &params.mode)
47 + .await
48 + .map_err(|error| utility_call_error("sleep", error))?;
49 +
50 + Ok(StatusCode::OK)
51 +}
52 +
53 +/// Wake the engine from sleep mode.
54 +pub async fn wake_up(
55 + State(state): State<Arc<AppState>>,
56 + Query(params): Query<WakeUpParams>,
57 +) -> Result<StatusCode, ApiError> {
58 + state
59 + .engine_core_client()
60 + .wake_up(params.tags)
61 + .await
62 + .map_err(|error| utility_call_error("wake_up", error))?;
63 +
64 + Ok(StatusCode::OK)
65 +}
66 +
67 +/// Return whether the engine is currently sleeping at any level.
68 +pub async fn is_sleeping(
69 + State(state): State<Arc<AppState>>,
70 +) -> Result<Json<IsSleepingResponse>, ApiError> {
71 + let is_sleeping = state
72 + .engine_core_client()
73 + .is_sleeping()
74 + .await
75 + .map_err(|error| utility_call_error("is_sleeping", error))?;
76 +
77 + Ok(Json(IsSleepingResponse { is_sleeping }))
78 +}
rust/src/server/src/routes/tests.rs new
+3679
@@ -0,0 +1,3679 @@
1 +// Route tests should use `Service::call` rather than `ServiceExt::oneshot`.
2 +// `oneshot` consumes the router and can drop `AppState` before a streaming
3 +// response body is fully drained, which closes the mock engine connection too
4 +// early and causes flaky `closed unexpectedly` failures.
5 +
6 +use std::collections::BTreeSet;
7 +use std::future::Future;
8 +use std::pin::Pin;
9 +use std::sync::Arc;
10 +use std::task::{Context, Poll};
11 +use std::time::Duration;
12 +use std::{fmt, fs};
13 +
14 +use axum::body::{Body, to_bytes};
15 +use axum::http::{Request, StatusCode};
16 +use bytes::Bytes;
17 +use futures::StreamExt as _;
18 +use rmpv::Value;
19 +use serde_json::json;
20 +use serial_test::serial;
21 +use tower::{Service as _, ServiceExt as _};
22 +use vllm_chat::{
23 + ChatBackend, ChatContent, ChatContentPart, ChatEvent, ChatLlm, ChatMessage, ChatRenderer,
24 + ChatRequest, ChatRole, ChatTextBackend, DefaultChatOutputProcessor, DynChatOutputProcessor,
25 + DynChatRenderer, NewChatOutputProcessorOptions, SamplingParams,
26 +};
27 +use vllm_engine_core_client::protocol::logprobs::{
28 + Logprobs, MaybeWireLogprobs, PositionLogprobs, TokenLogprob,
29 +};
30 +use vllm_engine_core_client::protocol::{
31 + EngineCoreFinishReason, EngineCoreOutput, EngineCoreOutputs, EngineCoreRequest, StopReason,
32 + UtilityOutput, UtilityResultEnvelope, decode_value,
33 +};
34 +use vllm_engine_core_client::test_utils::{IpcNamespace, spawn_mock_engine_task};
35 +use vllm_engine_core_client::{
36 + ENGINE_CORE_DEAD_SENTINEL, EngineCoreClient, EngineCoreClientConfig, EngineId,
37 +};
38 +use vllm_llm::Llm;
39 +use vllm_metrics::METRICS;
40 +use vllm_text::tokenizer::{DynTokenizer, Tokenizer};
41 +use vllm_text::{Prompt, TextBackend};
42 +use zeromq::prelude::{SocketRecv, SocketSend};
43 +use zeromq::{DealerSocket, PushSocket, ZmqMessage};
44 +
45 +use super::{build_router, build_router_with_dev_mode};
46 +use crate::routes::openai::chat_completions::convert::prepare_chat_request;
47 +use crate::state::AppState;
48 +
49 +fn request_output(
50 + request_id: &str,
51 + new_token_ids: Vec<u32>,
52 + finish_reason: Option<EngineCoreFinishReason>,
53 +) -> EngineCoreOutput {
54 + request_output_with_stop_reason(request_id, new_token_ids, finish_reason, None)
55 +}
56 +
57 +fn request_output_with_stop_reason(
58 + request_id: &str,
59 + new_token_ids: Vec<u32>,
60 + finish_reason: Option<EngineCoreFinishReason>,
61 + stop_reason: Option<StopReason>,
62 +) -> EngineCoreOutput {
63 + EngineCoreOutput {
64 + request_id: request_id.to_string(),
65 + new_token_ids,
66 + new_logprobs: None,
67 + new_prompt_logprobs_tensors: None,
68 + pooling_output: None,
69 + finish_reason,
70 + stop_reason,
71 + events: None,
72 + kv_transfer_params: None,
73 + trace_headers: None,
74 + prefill_stats: None,
75 + routed_experts: None,
76 + num_nans_in_logits: 0,
77 + }
78 +}
79 +
80 +fn request_output_with_logprobs(
81 + request_id: &str,
82 + new_token_ids: Vec<u32>,
83 + finish_reason: Option<EngineCoreFinishReason>,
84 + stop_reason: Option<StopReason>,
85 + new_logprobs: Option<Logprobs>,
86 + new_prompt_logprobs_tensors: Option<Logprobs>,
87 +) -> EngineCoreOutput {
88 + EngineCoreOutput {
89 + request_id: request_id.to_string(),
90 + new_token_ids,
91 + new_logprobs: new_logprobs.map(MaybeWireLogprobs::Direct),
92 + new_prompt_logprobs_tensors: new_prompt_logprobs_tensors.map(MaybeWireLogprobs::Direct),
93 + pooling_output: None,
94 + finish_reason,
95 + stop_reason,
96 + events: None,
97 + kv_transfer_params: None,
98 + trace_headers: None,
99 + prefill_stats: None,
100 + routed_experts: None,
101 + num_nans_in_logits: 0,
102 + }
103 +}
104 +
105 +fn request_output_with_logprobs_and_kv(
106 + request_id: &str,
107 + new_token_ids: Vec<u32>,
108 + finish_reason: Option<EngineCoreFinishReason>,
109 + stop_reason: Option<StopReason>,
110 + new_logprobs: Option<Logprobs>,
111 + new_prompt_logprobs_tensors: Option<Logprobs>,
112 + kv_transfer_params: Option<serde_json::Value>,
113 +) -> EngineCoreOutput {
114 + EngineCoreOutput {
115 + request_id: request_id.to_string(),
116 + new_token_ids,
117 + new_logprobs: new_logprobs.map(MaybeWireLogprobs::Direct),
118 + new_prompt_logprobs_tensors: new_prompt_logprobs_tensors.map(MaybeWireLogprobs::Direct),
119 + pooling_output: None,
120 + finish_reason,
121 + stop_reason,
122 + events: None,
123 + kv_transfer_params,
124 + trace_headers: None,
125 + prefill_stats: None,
126 + routed_experts: None,
127 + num_nans_in_logits: 0,
128 + }
129 +}
130 +
131 +fn bytes_to_token_ids(bytes: &[u8]) -> Vec<u32> {
132 + bytes.iter().map(|byte| u32::from(*byte)).collect()
133 +}
134 +
135 +fn default_stream_output_specs() -> Vec<(Vec<u32>, Option<EngineCoreFinishReason>)> {
136 + vec![
137 + (vec![b'h' as u32], None),
138 + (vec![b'i' as u32], None),
139 + (vec![b'!' as u32], Some(EngineCoreFinishReason::Stop)),
140 + ]
141 +}
142 +
143 +fn sse_data_payloads(text: &str) -> Vec<&str> {
144 + text.lines().filter_map(|line| line.strip_prefix("data: ")).collect()
145 +}
146 +
147 +type TestFuture<'a> = Pin<Box<dyn Future<Output = ()> + Send + 'a>>;
148 +
149 +fn boxed_test_future<'a>(future: impl Future<Output = ()> + Send + 'a) -> TestFuture<'a> {
150 + Box::pin(future)
151 +}
152 +
153 +struct MockEngineTask {
154 + shutdown_tx: Option<tokio::sync::oneshot::Sender<()>>,
155 + join_handle: Option<tokio::task::JoinHandle<()>>,
156 +}
157 +
158 +impl MockEngineTask {
159 + fn new(
160 + (shutdown_tx, join_handle): (
161 + tokio::sync::oneshot::Sender<()>,
162 + tokio::task::JoinHandle<()>,
163 + ),
164 + ) -> Self {
165 + Self {
166 + shutdown_tx: Some(shutdown_tx),
167 + join_handle: Some(join_handle),
168 + }
169 + }
170 +
171 + async fn finish(self) {
172 + self.await.expect("mock engine task");
173 + }
174 +
175 + fn abort(&self) {
176 + if let Some(join_handle) = &self.join_handle {
177 + join_handle.abort();
178 + }
179 + }
180 +
181 + async fn abort_and_join(mut self) {
182 + if let Some(join_handle) = self.join_handle.take() {
183 + join_handle.abort();
184 + let _ = join_handle.await;
185 + }
186 + }
187 +}
188 +
189 +impl Future for MockEngineTask {
190 + type Output = Result<(), tokio::task::JoinError>;
191 +
192 + fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
193 + if let Some(shutdown_tx) = self.shutdown_tx.take() {
194 + let _ = shutdown_tx.send(());
195 + }
196 + match self.join_handle.as_mut() {
197 + Some(join_handle) => Pin::new(join_handle).poll(cx),
198 + None => Poll::Ready(Ok(())),
199 + }
200 + }
201 +}
202 +
203 +impl Drop for MockEngineTask {
204 + fn drop(&mut self) {
205 + if let Some(join_handle) = &self.join_handle {
206 + join_handle.abort();
207 + }
208 + }
209 +}
210 +
211 +fn engine_outputs_for_request(
212 + request_id: &str,
213 + output_specs: Vec<(Vec<u32>, Option<EngineCoreFinishReason>)>,
214 +) -> EngineCoreOutputs {
215 + EngineCoreOutputs {
216 + engine_index: 0,
217 + outputs: output_specs
218 + .into_iter()
219 + .map(|(token_ids, finish_reason)| request_output(request_id, token_ids, finish_reason))
220 + .collect(),
221 + scheduler_stats: None,
222 + timestamp: 0.0,
223 + utility_output: None,
224 + finished_requests: None,
225 + wave_complete: None,
226 + start_wave: None,
227 + }
228 +}
229 +
230 +fn test_llm(client: EngineCoreClient) -> Llm {
231 + Llm::new(client).with_request_id_randomization(false)
232 +}
233 +
234 +fn sample_logprobs_for_token(token_id: u32, alternate_token_id: u32) -> Logprobs {
235 + Logprobs {
236 + positions: vec![PositionLogprobs {
237 + entries: vec![
238 + TokenLogprob {
239 + token_id,
240 + logprob: -0.1,
241 + rank: 1,
242 + },
243 + TokenLogprob {
244 + token_id: alternate_token_id,
245 + logprob: -0.2,
246 + rank: 1,
247 + },
248 + ],
249 + }],
250 + }
251 +}
252 +
253 +fn sample_logprobs_for_tokens(token_ids: &[u32]) -> Logprobs {
254 + Logprobs {
255 + positions: token_ids
256 + .iter()
257 + .map(|&token_id| PositionLogprobs {
258 + entries: vec![
259 + TokenLogprob {
260 + token_id,
261 + logprob: -0.1,
262 + rank: 1,
263 + },
264 + TokenLogprob {
265 + token_id: token_id.saturating_add(1),
266 + logprob: -0.2,
267 + rank: 2,
268 + },
269 + ],
270 + })
271 + .collect(),
272 + }
273 +}
274 +
275 +fn prompt_logprobs_for_hello() -> Logprobs {
276 + Logprobs {
277 + positions: vec![
278 + PositionLogprobs {
279 + entries: vec![
280 + TokenLogprob {
281 + token_id: b'e' as u32,
282 + logprob: -0.3,
283 + rank: 1,
284 + },
285 + TokenLogprob {
286 + token_id: b'a' as u32,
287 + logprob: -0.5,
288 + rank: 1,
289 + },
290 + ],
291 + },
292 + PositionLogprobs {
293 + entries: vec![
294 + TokenLogprob {
295 + token_id: b'l' as u32,
296 + logprob: -0.4,
297 + rank: 1,
298 + },
299 + TokenLogprob {
300 + token_id: b'r' as u32,
301 + logprob: -0.6,
302 + rank: 1,
303 + },
304 + ],
305 + },
306 + PositionLogprobs {
307 + entries: vec![
308 + TokenLogprob {
309 + token_id: b'l' as u32,
310 + logprob: -0.45,
311 + rank: 1,
312 + },
313 + TokenLogprob {
314 + token_id: b'i' as u32,
315 + logprob: -0.65,
316 + rank: 1,
317 + },
318 + ],
319 + },
320 + PositionLogprobs {
321 + entries: vec![
322 + TokenLogprob {
323 + token_id: b'o' as u32,
324 + logprob: -0.5,
325 + rank: 1,
326 + },
327 + TokenLogprob {
328 + token_id: b'u' as u32,
329 + logprob: -0.7,
330 + rank: 1,
331 + },
332 + ],
333 + },
334 + ],
335 + }
336 +}
337 +
338 +fn prompt_logprobs_for_tokens(token_ids: &[u32]) -> Logprobs {
339 + Logprobs {
340 + positions: token_ids
341 + .iter()
342 + .skip(1)
343 + .map(|&token_id| PositionLogprobs {
344 + entries: vec![
345 + TokenLogprob {
346 + token_id,
347 + logprob: -0.3,
348 + rank: 1,
349 + },
350 + TokenLogprob {
351 + token_id: token_id.saturating_add(1),
352 + logprob: -0.5,
353 + rank: 2,
354 + },
355 + ],
356 + })
357 + .collect(),
358 + }
359 +}
360 +
361 +fn utility_result_value<T>(value: T) -> UtilityResultEnvelope
362 +where
363 + T: serde::Serialize,
364 +{
365 + UtilityResultEnvelope::without_type_info(rmpv::ext::to_value(value).expect("encode result"))
366 +}
367 +
368 +fn utility_none_result() -> UtilityResultEnvelope {
369 + UtilityResultEnvelope::without_type_info(Value::Nil)
370 +}
371 +
372 +fn utility_outputs(call_id: i64, result: UtilityResultEnvelope) -> EngineCoreOutputs {
373 + EngineCoreOutputs {
374 + utility_output: Some(UtilityOutput {
375 + call_id,
376 + failure_message: None,
377 + result: Some(result),
378 + }),
379 + ..Default::default()
380 + }
381 +}
382 +
383 +async fn send_outputs(push: &mut PushSocket, outputs: EngineCoreOutputs) {
384 + push.send(ZmqMessage::from(
385 + rmp_serde::to_vec_named(&outputs).expect("encode outputs"),
386 + ))
387 + .await
388 + .expect("send outputs");
389 +}
390 +
391 +async fn recv_engine_message(dealer: &mut DealerSocket) -> Vec<Bytes> {
392 + dealer.recv().await.expect("recv engine message").into_vec()
393 +}
394 +
395 +#[derive(Clone)]
396 +struct FakeChatBackend {
397 + model_id: String,
398 + multimodal_model_info: Option<vllm_chat::multimodal::MultimodalModelInfo>,
399 +}
400 +
401 +#[derive(Debug)]
402 +struct FakeChatTokenizer;
403 +
404 +impl Tokenizer for FakeChatTokenizer {
405 + fn encode(
406 + &self,
407 + text: &str,
408 + _add_special_tokens: bool,
409 + ) -> vllm_text::tokenizer::Result<Vec<u32>> {
410 + let mut token_ids = Vec::new();
411 + let mut rest = text;
412 + while !rest.is_empty() {
413 + if let Some(stripped) = rest.strip_prefix("<image>") {
414 + token_ids.push(999);
415 + rest = stripped;
416 + continue;
417 + }
418 +
419 + let ch = rest.chars().next().expect("rest is not empty");
420 + let mut buf = [0; 4];
421 + token_ids.extend(ch.encode_utf8(&mut buf).bytes().map(u32::from));
422 + rest = &rest[ch.len_utf8()..];
423 + }
424 + Ok(token_ids)
425 + }
426 +
427 + fn decode(
428 + &self,
429 + token_ids: &[u32],
430 + _skip_special_tokens: bool,
431 + ) -> vllm_text::tokenizer::Result<String> {
432 + Ok(
433 + String::from_utf8_lossy(&token_ids.iter().map(|id| *id as u8).collect::<Vec<_>>())
434 + .into_owned(),
435 + )
436 + }
437 +
438 + fn token_to_id(&self, token: &str) -> Option<u32> {
439 + match token {
440 + "<image>" => Some(999),
441 + "<|image_pad|>" => Some(151655),
442 + "<think>" => Some(0xF001),
443 + "</think>" => Some(0xF002),
444 + "<|START_THINKING|>" => Some(0xF003),
445 + "<|END_THINKING|>" => Some(0xF004),
446 + "◁think▷" => Some(0xF005),
447 + "◁/think▷" => Some(0xF006),
448 + _ => None,
449 + }
450 + }
451 +
452 + fn id_to_token(&self, id: u32) -> Option<String> {
453 + match id {
454 + 999 => Some("<image>".to_string()),
455 + 151655 => Some("<|image_pad|>".to_string()),
456 + 0xF001 => Some("<think>".to_string()),
457 + 0xF002 => Some("</think>".to_string()),
458 + 0xF003 => Some("<|START_THINKING|>".to_string()),
459 + 0xF004 => Some("<|END_THINKING|>".to_string()),
460 + 0xF005 => Some("◁think▷".to_string()),
461 + 0xF006 => Some("◁/think▷".to_string()),
462 + _ => None,
463 + }
464 + }
465 +}
466 +
467 +impl FakeChatBackend {
468 + fn new() -> Self {
469 + Self {
470 + model_id: "test-model".to_string(),
471 + multimodal_model_info: None,
472 + }
473 + }
474 +
475 + fn with_model_id(model_id: impl Into<String>) -> Self {
476 + Self {
477 + model_id: model_id.into(),
478 + multimodal_model_info: None,
479 + }
480 + }
481 +
482 + fn with_multimodal_model_info(
483 + multimodal_model_info: vllm_chat::multimodal::MultimodalModelInfo,
484 + ) -> Self {
485 + Self {
486 + model_id: "test-model".to_string(),
487 + multimodal_model_info: Some(multimodal_model_info),
488 + }
489 + }
490 +}
491 +
492 +impl fmt::Debug for FakeChatBackend {
493 + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
494 + f.debug_struct("FakeChatBackend")
495 + .field("model_id", &self.model_id)
496 + .finish_non_exhaustive()
497 + }
498 +}
499 +
500 +impl TextBackend for FakeChatBackend {
501 + fn tokenizer(&self) -> DynTokenizer {
502 + Arc::new(FakeChatTokenizer)
503 + }
504 +
505 + fn model_id(&self) -> &str {
506 + &self.model_id
507 + }
508 +}
509 +
510 +impl ChatBackend for FakeChatBackend {
511 + fn chat_renderer(&self) -> DynChatRenderer {
512 + Arc::new(self.clone())
513 + }
514 +
515 + fn multimodal_model_info(&self) -> Option<&vllm_chat::multimodal::MultimodalModelInfo> {
516 + self.multimodal_model_info.as_ref()
517 + }
518 +
519 + fn new_chat_output_processor(
520 + &self,
521 + request: &mut ChatRequest,
522 + options: NewChatOutputProcessorOptions<'_>,
523 + ) -> vllm_chat::Result<DynChatOutputProcessor> {
524 + Ok(Box::new(DefaultChatOutputProcessor::new(
525 + request,
526 + &self.model_id,
527 + self.tokenizer(),
528 + options.tool_call_parser,
529 + options.reasoning_parser,
530 + )?))
531 + }
532 +}
533 +
534 +impl ChatRenderer for FakeChatBackend {
535 + fn render(&self, request: &ChatRequest) -> vllm_chat::Result<vllm_chat::RenderedPrompt> {
536 + let mut prompt = String::new();
537 + for message in &request.messages {
538 + prompt.push_str(message.role().as_str());
539 + prompt.push_str(": ");
540 + prompt.push_str(&render_fake_message_content(message)?);
541 + prompt.push('\n');
542 + }
543 + if request.chat_options.add_generation_prompt() {
544 + prompt.push_str("assistant:");
545 + }
546 + Ok(vllm_chat::RenderedPrompt {
547 + prompt: Prompt::Text(prompt),
548 + })
549 + }
550 +}
551 +
552 +fn render_fake_message_content(message: &ChatMessage) -> vllm_chat::Result<String> {
553 + match message {
554 + ChatMessage::System { content }
555 + | ChatMessage::Developer { content, .. }
556 + | ChatMessage::User { content }
557 + | ChatMessage::ToolResponse { content, .. } => render_fake_content(content),
558 + ChatMessage::Assistant { .. } => message.text_content(),
559 + }
560 +}
561 +
562 +fn render_fake_content(content: &ChatContent) -> vllm_chat::Result<String> {
563 + Ok(match content {
564 + ChatContent::Text(text) => text.clone(),
565 + ChatContent::Parts(parts) => {
566 + let mut out = String::new();
567 + for part in parts {
568 + match part {
569 + ChatContentPart::Text { text } => out.push_str(text),
570 + ChatContentPart::ImageUrl { .. } => out.push_str("<image>"),
571 + }
572 + }
573 + out
574 + }
575 + })
576 +}
577 +
578 +fn qwen_multimodal_model_info() -> vllm_chat::multimodal::MultimodalModelInfo {
579 + let config_path = std::env::temp_dir().join(format!(
580 + "vllm-server-qwen-config-{}.json",
581 + uuid::Uuid::new_v4()
582 + ));
583 + fs::write(
584 + &config_path,
585 + r#"{"model_type":"qwen2_vl","vision_token_id":151655}"#,
586 + )
587 + .expect("write qwen test config");
588 + let info = vllm_chat::multimodal::MultimodalModelInfo::from_paths(
589 + "qwen2-vl-test".to_string(),
590 + Some("qwen2_vl".to_string()),
591 + Some(&config_path),
592 + None,
593 + Arc::new(FakeChatTokenizer),
594 + )
595 + .expect("load multimodal info")
596 + .expect("qwen multimodal info is registered");
597 + let _ = fs::remove_file(config_path);
598 + info
599 +}
600 +
601 +#[derive(Clone, Debug)]
602 +struct FailingDecodeChatBackend;
603 +
604 +#[derive(Debug)]
605 +struct FailingDecodeTokenizer;
606 +
607 +impl Tokenizer for FailingDecodeTokenizer {
608 + fn encode(
609 + &self,
610 + text: &str,
611 + add_special_tokens: bool,
612 + ) -> vllm_text::tokenizer::Result<Vec<u32>> {
613 + FakeChatTokenizer.encode(text, add_special_tokens)
614 + }
615 +
616 + fn decode(
617 + &self,
618 + token_ids: &[u32],
619 + skip_special_tokens: bool,
620 + ) -> vllm_text::tokenizer::Result<String> {
621 + if token_ids.contains(&(b'i' as u32)) {
622 + return Err(vllm_text::tokenizer::TokenizerError(
623 + "forced decode failure for streaming test".to_string(),
624 + ));
625 + }
626 +
627 + FakeChatTokenizer.decode(token_ids, skip_special_tokens)
628 + }
629 +
630 + fn token_to_id(&self, token: &str) -> Option<u32> {
631 + FakeChatTokenizer.token_to_id(token)
632 + }
633 +}
634 +
635 +impl TextBackend for FailingDecodeChatBackend {
636 + fn tokenizer(&self) -> DynTokenizer {
637 + Arc::new(FailingDecodeTokenizer)
638 + }
639 +
640 + fn model_id(&self) -> &str {
641 + "test-model"
642 + }
643 +}
644 +
645 +impl ChatBackend for FailingDecodeChatBackend {
646 + fn chat_renderer(&self) -> DynChatRenderer {
647 + Arc::new(self.clone())
648 + }
649 +
650 + fn new_chat_output_processor(
651 + &self,
652 + _request: &mut ChatRequest,
653 + _options: NewChatOutputProcessorOptions<'_>,
654 + ) -> vllm_chat::Result<DynChatOutputProcessor> {
655 + Ok(Box::new(DefaultChatOutputProcessor::plain_text_only()))
656 + }
657 +}
658 +
659 +impl ChatRenderer for FailingDecodeChatBackend {
660 + fn render(&self, request: &ChatRequest) -> vllm_chat::Result<vllm_chat::RenderedPrompt> {
661 + FakeChatBackend::new().render(request)
662 + }
663 +}
664 +
665 +async fn test_models_with_engine_outputs_and_backend_inner(
666 + engine_id: impl Into<EngineId>,
667 + output_specs: Vec<(Vec<u32>, Option<EngineCoreFinishReason>)>,
668 + expected_prompt_token_ids: Option<Vec<u32>>,
669 + backend: Arc<dyn ChatTextBackend>,
670 +) -> (ChatLlm, MockEngineTask) {
671 + let ipc = IpcNamespace::new().expect("create ipc namespace");
672 + let handshake_address = ipc.handshake_endpoint();
673 + let engine_id = engine_id.into();
674 +
675 + let engine_task = MockEngineTask::new(spawn_mock_engine_task(
676 + handshake_address.clone(),
677 + engine_id.clone(),
678 + move |dealer, push| {
679 + boxed_test_future(async move {
680 + let add = recv_engine_message(dealer).await;
681 + let request: EngineCoreRequest =
682 + rmp_serde::from_slice(&add[1]).expect("decode request");
683 + if let Some(expected_prompt_token_ids) = expected_prompt_token_ids {
684 + assert_eq!(
685 + request.prompt_token_ids.as_deref(),
686 + Some(expected_prompt_token_ids.as_slice())
687 + );
688 + }
689 + send_outputs(
690 + push,
691 + engine_outputs_for_request(&request.request_id, output_specs),
692 + )
693 + .await;
694 + })
695 + },
696 + ));
697 +
698 + let client = EngineCoreClient::connect(
699 + EngineCoreClientConfig::new_single(handshake_address)
700 + .with_model_name("test-model")
701 + .with_local_input_output_addresses(
702 + Some(ipc.input_endpoint()),
703 + Some(ipc.output_endpoint()),
704 + ),
705 + )
706 + .await
707 + .expect("connect client");
708 +
709 + (
710 + ChatLlm::from_shared_backend(test_llm(client), backend),
711 + engine_task,
712 + )
713 +}
714 +
715 +async fn test_models_with_engine_outputs_and_backend(
716 + engine_id: impl Into<EngineId>,
717 + output_specs: Vec<(Vec<u32>, Option<EngineCoreFinishReason>)>,
718 + backend: Arc<dyn ChatTextBackend>,
719 +) -> (ChatLlm, MockEngineTask) {
720 + test_models_with_engine_outputs_and_backend_inner(engine_id, output_specs, None, backend).await
721 +}
722 +
723 +async fn test_chat_with_engine_outputs(
724 + engine_id: impl Into<EngineId>,
725 + output_specs: Vec<(Vec<u32>, Option<EngineCoreFinishReason>)>,
726 +) -> (ChatLlm, MockEngineTask) {
727 + test_models_with_engine_outputs_and_backend(
728 + engine_id,
729 + output_specs,
730 + Arc::new(FakeChatBackend::new()),
731 + )
732 + .await
733 +}
734 +
735 +async fn test_app() -> axum::Router {
736 + let (chat, _engine_task) = test_models_with_engine_outputs_and_backend(
737 + b"engine-openai",
738 + default_stream_output_specs(),
739 + Arc::new(FakeChatBackend::new()),
740 + )
741 + .await;
742 + build_router(Arc::new(AppState::new(
743 + vec!["Qwen/Qwen1.5-0.5B-Chat".to_string()],
744 + chat,
745 + )))
746 +}
747 +
748 +async fn test_health_app_with_engine_script<F>(
749 + script: F,
750 +) -> (axum::Router, Arc<AppState>, MockEngineTask)
751 +where
752 + F: for<'a> FnOnce(&'a mut PushSocket) -> TestFuture<'a> + Send + 'static,
753 +{
754 + let ipc = IpcNamespace::new().expect("create ipc namespace");
755 + let handshake_address = ipc.handshake_endpoint();
756 + let engine_id = b"engine-openai-health".to_vec();
757 +
758 + let engine_task = MockEngineTask::new(spawn_mock_engine_task(
759 + handshake_address.clone(),
760 + engine_id.clone(),
761 + move |_dealer, push| script(push),
762 + ));
763 +
764 + let client = EngineCoreClient::connect(
765 + EngineCoreClientConfig::new_single(handshake_address)
766 + .with_model_name("test-model")
767 + .with_local_input_output_addresses(
768 + Some(ipc.input_endpoint()),
769 + Some(ipc.output_endpoint()),
770 + ),
771 + )
772 + .await
773 + .expect("connect client");
774 +
775 + let chat = ChatLlm::from_shared_backend(test_llm(client), Arc::new(FakeChatBackend::new()));
776 + let state = Arc::new(AppState::new(
777 + vec!["Qwen/Qwen1.5-0.5B-Chat".to_string()],
778 + chat,
779 + ));
780 + (build_router(state.clone()), state, engine_task)
781 +}
782 +
783 +async fn test_admin_app_with_engine_script<F>(script: F) -> (axum::Router, MockEngineTask)
784 +where
785 + F: for<'a> FnOnce(&'a mut DealerSocket, &'a mut PushSocket) -> TestFuture<'a> + Send + 'static,
786 +{
787 + let ipc = IpcNamespace::new().expect("create ipc namespace");
788 + let handshake_address = ipc.handshake_endpoint();
789 + let engine_id = b"engine-openai-admin".to_vec();
790 +
791 + let engine_task = MockEngineTask::new(spawn_mock_engine_task(
792 + handshake_address.clone(),
793 + engine_id.clone(),
794 + move |dealer, push| script(dealer, push),
795 + ));
796 +
797 + let client = EngineCoreClient::connect(
798 + EngineCoreClientConfig::new_single(handshake_address)
799 + .with_model_name("test-model")
800 + .with_local_input_output_addresses(
801 + Some(ipc.input_endpoint()),
802 + Some(ipc.output_endpoint()),
803 + ),
804 + )
805 + .await
806 + .expect("connect client");
807 +
808 + let chat = ChatLlm::from_shared_backend(test_llm(client), Arc::new(FakeChatBackend::new()));
809 + (
810 + build_router_with_dev_mode(
811 + Arc::new(AppState::new(
812 + vec!["Qwen/Qwen1.5-0.5B-Chat".to_string()],
813 + chat,
814 + )),
815 + true,
816 + ),
817 + engine_task,
818 + )
819 +}
820 +
821 +async fn test_app_with_engine_handle() -> (axum::Router, MockEngineTask) {
822 + test_app_with_stream_output_specs(default_stream_output_specs()).await
823 +}
824 +
825 +async fn test_app_with_stream_output_specs(
826 + output_specs: Vec<(Vec<u32>, Option<EngineCoreFinishReason>)>,
827 +) -> (axum::Router, MockEngineTask) {
828 + let (chat, engine_task) = test_models_with_engine_outputs_and_backend(
829 + b"engine-openai",
830 + output_specs,
831 + Arc::new(FakeChatBackend::new()),
832 + )
833 + .await;
834 + (
835 + build_router(Arc::new(AppState::new(
836 + vec!["Qwen/Qwen1.5-0.5B-Chat".to_string()],
837 + chat,
838 + ))),
839 + engine_task,
840 + )
841 +}
842 +
843 +async fn test_app_with_backend_and_stream_output_specs(
844 + backend: Arc<dyn ChatTextBackend>,
845 + output_specs: Vec<(Vec<u32>, Option<EngineCoreFinishReason>)>,
846 +) -> (axum::Router, MockEngineTask) {
847 + let (chat, engine_task) =
848 + test_models_with_engine_outputs_and_backend(b"engine-openai", output_specs, backend).await;
849 + (
850 + build_router(Arc::new(AppState::new(
851 + vec!["Qwen/Qwen1.5-0.5B-Chat".to_string()],
852 + chat,
853 + ))),
854 + engine_task,
855 + )
856 +}
857 +
858 +async fn test_app_with_backend_and_engine_request_check<F>(
859 + backend: Arc<dyn ChatTextBackend>,
860 + check_request: F,
861 +) -> (axum::Router, MockEngineTask)
862 +where
863 + F: FnOnce(&EngineCoreRequest) + Send + 'static,
864 +{
865 + let ipc = IpcNamespace::new().expect("create ipc namespace");
866 + let handshake_address = ipc.handshake_endpoint();
867 + let engine_id = b"engine-openai-check-request".to_vec();
868 +
869 + let engine_task = MockEngineTask::new(spawn_mock_engine_task(
870 + handshake_address.clone(),
871 + engine_id.clone(),
872 + move |dealer, push| {
873 + boxed_test_future(async move {
874 + let add = recv_engine_message(dealer).await;
875 + let request: EngineCoreRequest =
876 + rmp_serde::from_slice(&add[1]).expect("decode request");
877 + check_request(&request);
878 + send_outputs(
879 + push,
880 + engine_outputs_for_request(&request.request_id, default_stream_output_specs()),
881 + )
882 + .await;
883 + })
884 + },
885 + ));
886 +
887 + let client = EngineCoreClient::connect(
888 + EngineCoreClientConfig::new_single(handshake_address)
889 + .with_model_name("test-model")
890 + .with_local_input_output_addresses(
891 + Some(ipc.input_endpoint()),
892 + Some(ipc.output_endpoint()),
893 + ),
894 + )
895 + .await
896 + .expect("connect client");
897 +
898 + let chat = ChatLlm::from_shared_backend(test_llm(client), backend);
899 + (
900 + build_router(Arc::new(AppState::new(
901 + vec!["Qwen/Qwen1.5-0.5B-Chat".to_string()],
902 + chat,
903 + ))),
904 + engine_task,
905 + )
906 +}
907 +
908 +async fn test_chat_with_engine_handle() -> (ChatLlm, MockEngineTask) {
909 + test_chat_with_engine_outputs(b"engine-openai-chat", default_stream_output_specs()).await
910 +}
911 +
912 +async fn server_load(app: &axum::Router) -> u64 {
913 + let response = app
914 + .clone()
915 + .call(
916 + Request::builder()
917 + .method("GET")
918 + .uri("/load")
919 + .body(Body::empty())
920 + .expect("build request"),
921 + )
922 + .await
923 + .expect("call app");
924 +
925 + assert_eq!(response.status(), StatusCode::OK);
926 + let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");
927 + let value: serde_json::Value = serde_json::from_slice(&body).expect("json body");
928 + value["server_load"].as_u64().expect("server_load")
929 +}
930 +
931 +async fn health_status(app: &axum::Router) -> (StatusCode, Bytes) {
932 + let response = app
933 + .clone()
934 + .call(
935 + Request::builder()
936 + .method("GET")
937 + .uri("/health")
938 + .body(Body::empty())
939 + .expect("build request"),
940 + )
941 + .await
942 + .expect("call app");
943 +
944 + let status = response.status();
945 + let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");
946 + (status, body)
947 +}
948 +
949 +fn metric_value(rendered: &str, metric: &str, labels: Option<&str>) -> Option<f64> {
950 + rendered.lines().find_map(|line| {
951 + let rest = line.strip_prefix(metric)?;
952 +
953 + match labels {
954 + Some(labels) => {
955 + let (encoded_labels, value) = rest.split_once("} ")?;
956 + if !encoded_labels.starts_with('{') {
957 + return None;
958 + }
959 + let expected_parts = labels.split(',');
960 + if expected_parts.into_iter().all(|part| encoded_labels.contains(part)) {
961 + value.parse::<f64>().ok()
962 + } else {
963 + None
964 + }
965 + }
966 + None => rest.strip_prefix(' ').and_then(|value| value.parse::<f64>().ok()),
967 + }
968 + })
969 +}
970 +
971 +fn metric_delta(
972 + rendered_before: &str,
973 + rendered_after: &str,
974 + metric: &str,
975 + labels: Option<&str>,
976 +) -> f64 {
977 + metric_value(rendered_after, metric, labels).unwrap_or(0.0)
978 + - metric_value(rendered_before, metric, labels).unwrap_or(0.0)
979 +}
980 +
981 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
982 +#[serial]
983 +async fn list_models_returns_configured_model() {
984 + let mut app = test_app().await;
985 + let response = app
986 + .call(Request::builder().uri("/v1/models").body(Body::empty()).expect("build request"))
987 + .await
988 + .expect("call app");
989 +
990 + assert_eq!(response.status(), StatusCode::OK);
991 + let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");
992 + let json: serde_json::Value = serde_json::from_slice(&body).expect("decode json");
993 + assert_eq!(json["data"][0]["id"], "Qwen/Qwen1.5-0.5B-Chat");
994 +}
995 +
996 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
997 +#[serial]
998 +async fn http_metrics_record_list_models_requests() {
999 + let mut app = test_app().await;
1000 + let before = METRICS.render().unwrap();
1001 +
1002 + let response = app
1003 + .call(
1004 + Request::builder()
1005 + .method("GET")
1006 + .uri("/v1/models")
1007 + .body(Body::empty())
1008 + .expect("build request"),
1009 + )
1010 + .await
1011 + .expect("call app");
1012 +
1013 + assert_eq!(response.status(), StatusCode::OK);
1014 +
1015 + let after = METRICS.render().unwrap();
1016 + assert_eq!(
1017 + metric_delta(
1018 + &before,
1019 + &after,
1020 + "http_requests_total",
1021 + Some("method=\"GET\",status=\"2xx\",handler=\"/v1/models\""),
1022 + ),
1023 + 1.0
1024 + );
1025 + assert_eq!(
1026 + metric_delta(
1027 + &before,
1028 + &after,
1029 + "http_request_duration_seconds_count",
1030 + Some("method=\"GET\",handler=\"/v1/models\""),
1031 + ),
1032 + 1.0
1033 + );
1034 + assert_eq!(
1035 + metric_delta(
1036 + &before,
1037 + &after,
1038 + "http_request_duration_highr_seconds_count",
1039 + None,
1040 + ),
1041 + 1.0
1042 + );
1043 +}
1044 +
1045 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1046 +#[serial]
1047 +async fn wrong_model_returns_not_found() {
1048 + let mut app = test_app().await;
1049 + let response = app
1050 + .call(
1051 + Request::builder()
1052 + .method("POST")
1053 + .uri("/v1/chat/completions")
1054 + .header("content-type", "application/json")
1055 + .body(Body::from(
1056 + json!({
1057 + "model": "wrong-model",
1058 + "stream": true,
1059 + "messages": [{"role": "user", "content": "hello"}]
1060 + })
1061 + .to_string(),
1062 + ))
1063 + .expect("build request"),
1064 + )
1065 + .await
1066 + .expect("call app");
1067 +
1068 + assert_eq!(response.status(), StatusCode::NOT_FOUND);
1069 +}
1070 +
1071 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1072 +#[serial]
1073 +async fn invalid_request_returns_openai_error() {
1074 + let mut app = test_app().await;
1075 + let response = app
1076 + .call(
1077 + Request::builder()
1078 + .method("POST")
1079 + .uri("/v1/chat/completions")
1080 + .header("content-type", "application/json")
1081 + .body(Body::from(
1082 + json!({
1083 + "model": "Qwen/Qwen1.5-0.5B-Chat",
1084 + "stream": false,
1085 + "stream_options": {"include_usage": true},
1086 + "messages": [{"role": "user", "content": "hello"}]
1087 + })
1088 + .to_string(),
1089 + ))
1090 + .expect("build request"),
1091 + )
1092 + .await
1093 + .expect("call app");
1094 +
1095 + assert_eq!(response.status(), StatusCode::BAD_REQUEST);
1096 + let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");
1097 + let json: serde_json::Value = serde_json::from_slice(&body).expect("decode json");
1098 + assert_eq!(json["error"]["type"], "invalid_request_error");
1099 +}
1100 +
1101 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1102 +#[serial]
1103 +async fn non_stream_chat_returns_json_response() {
1104 + let (app, engine_task) = test_app_with_engine_handle().await;
1105 + let response = app
1106 + .clone()
1107 + .call(
1108 + Request::builder()
1109 + .method("POST")
1110 + .uri("/v1/chat/completions")
1111 + .header("content-type", "application/json")
1112 + .body(Body::from(
1113 + json!({
1114 + "model": "Qwen/Qwen1.5-0.5B-Chat",
1115 + "stream": false,
1116 + "messages": [{"role": "user", "content": "hello"}]
1117 + })
1118 + .to_string(),
1119 + ))
1120 + .expect("build request"),
1121 + )
1122 + .await
1123 + .expect("call app");
1124 +
1125 + assert_eq!(response.status(), StatusCode::OK);
1126 + assert!(
1127 + response
1128 + .headers()
1129 + .get("content-type")
1130 + .and_then(|value| value.to_str().ok())
1131 + .is_some_and(|value| value.starts_with("application/json"))
1132 + );
1133 +
1134 + let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");
1135 + engine_task.await.expect("mock engine task");
1136 + let json: serde_json::Value = serde_json::from_slice(&body).expect("decode json");
1137 +
1138 + assert_eq!(json["object"], "chat.completion");
1139 + assert_eq!(json["choices"][0]["message"]["role"], "assistant");
1140 + assert_eq!(json["choices"][0]["message"]["content"], "hi");
1141 + assert_eq!(json["choices"][0]["finish_reason"], "stop");
1142 + assert_eq!(json["usage"]["prompt_tokens"], 22);
1143 + assert_eq!(json["usage"]["completion_tokens"], 3);
1144 + assert_eq!(json["usage"]["total_tokens"], 25);
1145 +}
1146 +
1147 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1148 +#[serial]
1149 +async fn non_stream_chat_image_url_reaches_engine_mm_features() {
1150 + let (app, engine_task) = test_app_with_backend_and_engine_request_check(
1151 + Arc::new(FakeChatBackend::with_multimodal_model_info(
1152 + qwen_multimodal_model_info(),
1153 + )),
1154 + |request| {
1155 + let prompt_token_ids = request.prompt_token_ids.as_ref().expect("prompt token ids");
1156 + assert!(prompt_token_ids.contains(&151655));
1157 +
1158 + let features = request.mm_features.as_ref().expect("multimodal features");
1159 + assert_eq!(features.len(), 1);
1160 + assert_eq!(features[0].modality, "image");
1161 + assert_eq!(features[0].identifier, "image-1");
1162 + assert!(features[0].mm_position.length > 0);
1163 + assert!(features[0].mm_position.is_embed.is_some());
1164 +
1165 + let data = features[0].data.as_ref().expect("feature data");
1166 + assert!(data.contains_key("pixel_values"));
1167 + assert!(data.contains_key("image_grid_thw"));
1168 + },
1169 + )
1170 + .await;
1171 +
1172 + let response = app
1173 + .clone()
1174 + .call(
1175 + Request::builder()
1176 + .method("POST")
1177 + .uri("/v1/chat/completions")
1178 + .header("content-type", "application/json")
1179 + .body(Body::from(
1180 + json!({
1181 + "model": "Qwen/Qwen1.5-0.5B-Chat",
1182 + "stream": false,
1183 + "messages": [{
1184 + "role": "user",
1185 + "content": [
1186 + {"type": "text", "text": "describe "},
1187 + {
1188 + "type": "image_url",
1189 + "image_url": {
1190 + "url": "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNk+A8AAQUBAScY42YAAAAASUVORK5CYII="
1191 + },
1192 + "uuid": "image-1"
1193 + }
1194 + ]
1195 + }]
1196 + })
1197 + .to_string(),
1198 + ))
1199 + .expect("build request"),
1200 + )
1201 + .await
1202 + .expect("call app");
1203 +
1204 + assert_eq!(response.status(), StatusCode::OK);
1205 + let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");
1206 + engine_task.await.expect("mock engine task");
1207 + let json: serde_json::Value = serde_json::from_slice(&body).expect("decode json");
1208 +
1209 + assert_eq!(json["object"], "chat.completion");
1210 + assert_eq!(json["choices"][0]["message"]["content"], "hi");
1211 +}
1212 +
1213 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1214 +#[serial]
1215 +async fn non_stream_chat_includes_logprobs_and_prompt_logprobs() {
1216 + let ipc = IpcNamespace::new().expect("create ipc namespace");
1217 + let handshake_address = ipc.handshake_endpoint();
1218 + let engine_id = b"engine-openai-chat-logprobs".to_vec();
1219 +
1220 + let engine_task = MockEngineTask::new(spawn_mock_engine_task(
1221 + handshake_address.clone(),
1222 + engine_id.clone(),
1223 + |dealer, push| {
1224 + boxed_test_future(async move {
1225 + let add = recv_engine_message(dealer).await;
1226 + let request: EngineCoreRequest =
1227 + rmp_serde::from_slice(&add[1]).expect("decode request");
1228 + let prompt_token_ids = request.prompt_token_ids.clone().expect("prompt token ids");
1229 +
1230 + send_outputs(
1231 + push,
1232 + EngineCoreOutputs {
1233 + engine_index: 0,
1234 + outputs: vec![request_output_with_logprobs(
1235 + &request.request_id,
1236 + bytes_to_token_ids(b"hi"),
1237 + Some(EngineCoreFinishReason::Stop),
1238 + None,
1239 + Some(sample_logprobs_for_tokens(&bytes_to_token_ids(b"hi"))),
1240 + Some(prompt_logprobs_for_tokens(&prompt_token_ids)),
1241 + )],
1242 + scheduler_stats: None,
1243 + timestamp: 0.0,
1244 + utility_output: None,
1245 + finished_requests: None,
1246 + wave_complete: None,
1247 + start_wave: None,
1248 + },
1249 + )
1250 + .await;
1251 + })
1252 + },
1253 + ));
1254 +
1255 + let client = EngineCoreClient::connect(
1256 + EngineCoreClientConfig::new_single(handshake_address)
1257 + .with_model_name("test-model")
1258 + .with_local_input_output_addresses(
1259 + Some(ipc.input_endpoint()),
1260 + Some(ipc.output_endpoint()),
1261 + ),
1262 + )
1263 + .await
1264 + .expect("connect client");
1265 + let chat = ChatLlm::from_shared_backend(test_llm(client), Arc::new(FakeChatBackend::new()));
1266 + let mut app = build_router(Arc::new(AppState::new(
1267 + vec!["Qwen/Qwen1.5-0.5B-Chat".to_string()],
1268 + chat,
1269 + )));
1270 +
1271 + let response = app
1272 + .call(
1273 + Request::builder()
1274 + .method("POST")
1275 + .uri("/v1/chat/completions")
1276 + .header("content-type", "application/json")
1277 + .body(Body::from(
1278 + json!({
1279 + "model": "Qwen/Qwen1.5-0.5B-Chat",
1280 + "stream": false,
1281 + "logprobs": true,
1282 + "prompt_logprobs": 1,
1283 + "messages": [{"role": "user", "content": "hello"}]
1284 + })
1285 + .to_string(),
1286 + ))
1287 + .expect("build request"),
1288 + )
1289 + .await
1290 + .expect("call app");
1291 +
1292 + assert_eq!(response.status(), StatusCode::OK);
1293 + let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");
1294 + engine_task.await.expect("mock engine task");
1295 + let json: serde_json::Value = serde_json::from_slice(&body).expect("decode json");
1296 +
1297 + assert_eq!(
1298 + json["choices"][0]["logprobs"]["content"][0]["token"],
1299 + json!("h")
1300 + );
1301 + assert_eq!(
1302 + json["choices"][0]["logprobs"]["content"][1]["token"],
1303 + json!("i")
1304 + );
1305 + assert_eq!(json["prompt_logprobs"][0], serde_json::Value::Null);
1306 + assert!(json["prompt_logprobs"][1].is_object());
1307 +}
1308 +
1309 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1310 +#[serial]
1311 +async fn happy_path_returns_sse_stream() {
1312 + let (app, engine_task) = test_app_with_engine_handle().await;
1313 + let before = METRICS.render().unwrap();
1314 + let response = app
1315 + .clone()
1316 + .call(
1317 + Request::builder()
1318 + .method("POST")
1319 + .uri("/v1/chat/completions")
1320 + .header("content-type", "application/json")
1321 + .body(Body::from(
1322 + json!({
1323 + "model": "Qwen/Qwen1.5-0.5B-Chat",
1324 + "stream": true,
1325 + "messages": [{"role": "user", "content": "hello"}]
1326 + })
1327 + .to_string(),
1328 + ))
1329 + .expect("build request"),
1330 + )
1331 + .await
1332 + .expect("call app");
1333 +
1334 + assert_eq!(response.status(), StatusCode::OK);
1335 + assert_eq!(
1336 + response.headers().get("content-type").and_then(|value| value.to_str().ok()),
1337 + Some("text/event-stream")
1338 + );
1339 +
1340 + let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");
1341 + engine_task.await.expect("mock engine task");
1342 + let text = String::from_utf8(body.to_vec()).expect("utf8 body");
1343 + let after = METRICS.render().unwrap();
1344 +
1345 + assert!(text.contains("\"role\":\"assistant\""), "{text}");
1346 + assert!(text.starts_with("data: "), "{text}");
1347 + assert_eq!(
1348 + metric_delta(
1349 + &before,
1350 + &after,
1351 + "http_requests_total",
1352 + Some("method=\"POST\",status=\"2xx\",handler=\"/v1/chat/completions\""),
1353 + ),
1354 + 1.0
1355 + );
1356 + assert_eq!(
1357 + metric_delta(
1358 + &before,
1359 + &after,
1360 + "http_request_duration_seconds_count",
1361 + Some("method=\"POST\",handler=\"/v1/chat/completions\""),
1362 + ),
1363 + 1.0
1364 + );
1365 + assert_eq!(
1366 + metric_delta(
1367 + &before,
1368 + &after,
1369 + "http_request_duration_highr_seconds_count",
1370 + None,
1371 + ),
1372 + 1.0
1373 + );
1374 +}
1375 +
1376 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1377 +#[serial]
1378 +async fn http_metrics_exclude_metrics_route() {
1379 + let mut app = test_app().await;
1380 + let before = METRICS.render().unwrap();
1381 +
1382 + let response = app
1383 + .call(
1384 + Request::builder()
1385 + .method("GET")
1386 + .uri("/metrics")
1387 + .body(Body::empty())
1388 + .expect("build request"),
1389 + )
1390 + .await
1391 + .expect("call app");
1392 +
1393 + assert_eq!(response.status(), StatusCode::OK);
1394 +
1395 + let after = METRICS.render().unwrap();
1396 + assert_eq!(
1397 + metric_delta(
1398 + &before,
1399 + &after,
1400 + "http_request_duration_highr_seconds_count",
1401 + None,
1402 + ),
1403 + 0.0
1404 + );
1405 + assert_eq!(
1406 + metric_value(
1407 + &after,
1408 + "http_requests_total",
1409 + Some("method=\"GET\",status=\"2xx\",handler=\"/metrics\""),
1410 + ),
1411 + None
1412 + );
1413 + assert_eq!(
1414 + metric_value(
1415 + &after,
1416 + "http_request_duration_seconds_count",
1417 + Some("method=\"GET\",handler=\"/metrics\""),
1418 + ),
1419 + None
1420 + );
1421 +}
1422 +
1423 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1424 +#[serial]
1425 +async fn http_metrics_group_error_statuses() {
1426 + let mut app = test_app().await;
1427 + let before = METRICS.render().unwrap();
1428 +
1429 + let response = app
1430 + .call(
1431 + Request::builder()
1432 + .method("POST")
1433 + .uri("/v1/chat/completions")
1434 + .header("content-type", "application/json")
1435 + .body(Body::from(
1436 + json!({
1437 + "model": "Qwen/Qwen1.5-0.5B-Chat",
1438 + "stream": false,
1439 + "stream_options": {"include_usage": true},
1440 + "messages": [{"role": "user", "content": "hello"}]
1441 + })
1442 + .to_string(),
1443 + ))
1444 + .expect("build request"),
1445 + )
1446 + .await
1447 + .expect("call app");
1448 +
1449 + assert_eq!(response.status(), StatusCode::BAD_REQUEST);
1450 +
1451 + let after = METRICS.render().unwrap();
1452 + assert_eq!(
1453 + metric_delta(
1454 + &before,
1455 + &after,
1456 + "http_requests_total",
1457 + Some("method=\"POST\",status=\"4xx\",handler=\"/v1/chat/completions\""),
1458 + ),
1459 + 1.0
1460 + );
1461 +}
1462 +
1463 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1464 +#[serial]
1465 +async fn load_endpoint_tracks_chat_stream_lifecycle() {
1466 + let (app, engine_task) = test_app_with_engine_handle().await;
1467 +
1468 + assert_eq!(server_load(&app).await, 0);
1469 +
1470 + let response = app
1471 + .clone()
1472 + .call(
1473 + Request::builder()
1474 + .method("POST")
1475 + .uri("/v1/chat/completions")
1476 + .header("content-type", "application/json")
1477 + .body(Body::from(
1478 + json!({
1479 + "model": "Qwen/Qwen1.5-0.5B-Chat",
1480 + "stream": true,
1481 + "messages": [{"role": "user", "content": "hello"}]
1482 + })
1483 + .to_string(),
1484 + ))
1485 + .expect("build request"),
1486 + )
1487 + .await
1488 + .expect("call app");
1489 +
1490 + assert_eq!(response.status(), StatusCode::OK);
1491 + assert_eq!(server_load(&app).await, 1);
1492 +
1493 + let _body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");
1494 + engine_task.await.expect("mock engine task");
1495 +
1496 + assert_eq!(server_load(&app).await, 0);
1497 +}
1498 +
1499 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1500 +#[serial]
1501 +async fn health_endpoint_returns_ok_with_empty_body_when_client_is_healthy() {
1502 + let (app, _state, engine_task) =
1503 + test_health_app_with_engine_script(|_push| boxed_test_future(async move {})).await;
1504 +
1505 + let (status, body) = health_status(&app).await;
1506 + assert_eq!(status, StatusCode::OK);
1507 + assert!(body.is_empty(), "expected empty body, got {:?}", body);
1508 +
1509 + engine_task.await.expect("mock engine task");
1510 +}
1511 +
1512 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1513 +#[serial]
1514 +async fn health_endpoint_returns_503_after_engine_core_dead_sentinel() {
1515 + let (app, state, engine_task) = test_health_app_with_engine_script(|push| {
1516 + boxed_test_future(async move {
1517 + push.send(ZmqMessage::from(ENGINE_CORE_DEAD_SENTINEL.to_vec()))
1518 + .await
1519 + .expect("send sentinel");
1520 + })
1521 + })
1522 + .await;
1523 +
1524 + tokio::time::timeout(Duration::from_secs(2), async {
1525 + while state.chat.engine_core_client().is_healthy() {
1526 + tokio::task::yield_now().await;
1527 + }
1528 + })
1529 + .await
1530 + .expect("wait for unhealthy client");
1531 +
1532 + let (status, body) = health_status(&app).await;
1533 + assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
1534 + assert!(body.is_empty(), "expected empty body, got {:?}", body);
1535 + assert!(matches!(
1536 + state.chat.engine_core_client().health_error().as_deref(),
1537 + Some(vllm_engine_core_client::Error::EngineCoreDead)
1538 + ));
1539 +
1540 + engine_task.await.expect("mock engine task");
1541 +}
1542 +
1543 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1544 +#[serial]
1545 +async fn load_endpoint_resets_when_stream_response_is_dropped() {
1546 + let (app, engine_task) = test_app_with_engine_handle().await;
1547 +
1548 + let response = app
1549 + .clone()
1550 + .call(
1551 + Request::builder()
1552 + .method("POST")
1553 + .uri("/v1/completions")
1554 + .header("content-type", "application/json")
1555 + .body(Body::from(
1556 + json!({
1557 + "model": "Qwen/Qwen1.5-0.5B-Chat",
1558 + "prompt": "hello",
1559 + "stream": true
1560 + })
1561 + .to_string(),
1562 + ))
1563 + .expect("build request"),
1564 + )
1565 + .await
1566 + .expect("call app");
1567 +
1568 + assert_eq!(response.status(), StatusCode::OK);
1569 + assert_eq!(server_load(&app).await, 1);
1570 +
1571 + drop(response);
1572 + tokio::task::yield_now().await;
1573 + engine_task.await.expect("mock engine task");
1574 +
1575 + assert_eq!(server_load(&app).await, 0);
1576 +}
1577 +
1578 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1579 +#[serial]
1580 +async fn stream_error_is_returned_as_openai_error_sse() {
1581 + let (app, engine_task) = test_app_with_backend_and_stream_output_specs(
1582 + Arc::new(FailingDecodeChatBackend),
1583 + default_stream_output_specs(),
1584 + )
1585 + .await;
1586 + let response = app
1587 + .clone()
1588 + .call(
1589 + Request::builder()
1590 + .method("POST")
1591 + .uri("/v1/chat/completions")
1592 + .header("content-type", "application/json")
1593 + .body(Body::from(
1594 + json!({
1595 + "model": "Qwen/Qwen1.5-0.5B-Chat",
1596 + "stream": true,
1597 + "stream_options": {"include_usage": true},
1598 + "messages": [{"role": "user", "content": "hello"}]
1599 + })
1600 + .to_string(),
1601 + ))
1602 + .expect("build request"),
1603 + )
1604 + .await
1605 + .expect("call app");
1606 +
1607 + assert_eq!(response.status(), StatusCode::OK);
1608 +
1609 + let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");
1610 + engine_task.await.expect("mock engine task");
1611 + let text = String::from_utf8(body.to_vec()).expect("utf8 body");
1612 +
1613 + assert!(text.contains("\"role\":\"assistant\""), "{text}");
1614 + assert!(text.contains("\"type\":\"server_error\""), "{text}");
1615 + assert!(
1616 + text.contains("forced decode failure for streaming test"),
1617 + "{text}"
1618 + );
1619 + assert!(!text.contains("\"usage\":"), "{text}");
1620 + assert!(text.trim_end().ends_with("data: [DONE]"), "{text}");
1621 +}
1622 +
1623 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1624 +#[serial]
1625 +async fn invalid_terminal_finish_reason_is_returned_as_openai_error_sse() {
1626 + let (app, engine_task) =
1627 + test_app_with_stream_output_specs(vec![(vec![], Some(EngineCoreFinishReason::Error))])
1628 + .await;
1629 + let response = app
1630 + .clone()
1631 + .call(
1632 + Request::builder()
1633 + .method("POST")
1634 + .uri("/v1/chat/completions")
1635 + .header("content-type", "application/json")
1636 + .body(Body::from(
1637 + json!({
1638 + "model": "Qwen/Qwen1.5-0.5B-Chat",
1639 + "stream": true,
1640 + "stream_options": {"include_usage": true},
1641 + "messages": [{"role": "user", "content": "hello"}]
1642 + })
1643 + .to_string(),
1644 + ))
1645 + .expect("build request"),
1646 + )
1647 + .await
1648 + .expect("call app");
1649 +
1650 + assert_eq!(response.status(), StatusCode::OK);
1651 +
1652 + let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");
1653 + engine_task.await.expect("mock engine task");
1654 + let text = String::from_utf8(body.to_vec()).expect("utf8 body");
1655 +
1656 + assert!(text.contains("\"role\":\"assistant\""), "{text}");
1657 + assert!(text.contains("\"type\":\"server_error\""), "{text}");
1658 + assert!(text.contains("Internal server error"), "{text}");
1659 + assert!(!text.contains("\"usage\":"), "{text}");
1660 + assert!(text.trim_end().ends_with("data: [DONE]"), "{text}");
1661 +}
1662 +
1663 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1664 +#[serial]
1665 +async fn include_usage_adds_final_usage_chunk_before_done() {
1666 + let (app, engine_task) = test_app_with_stream_output_specs(default_stream_output_specs()).await;
1667 + let response = app
1668 + .clone()
1669 + .call(
1670 + Request::builder()
1671 + .method("POST")
1672 + .uri("/v1/chat/completions")
1673 + .header("content-type", "application/json")
1674 + .body(Body::from(
1675 + json!({
1676 + "model": "Qwen/Qwen1.5-0.5B-Chat",
1677 + "stream": true,
1678 + "stream_options": {"include_usage": true},
1679 + "messages": [{"role": "user", "content": "hello"}]
1680 + })
1681 + .to_string(),
1682 + ))
1683 + .expect("build request"),
1684 + )
1685 + .await
1686 + .expect("call app");
1687 +
1688 + assert_eq!(response.status(), StatusCode::OK);
1689 +
1690 + let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");
1691 + engine_task.await.expect("mock engine task");
1692 + let text = String::from_utf8(body.to_vec()).expect("utf8 body");
1693 +
1694 + let payloads = sse_data_payloads(&text);
1695 + let finish_index = payloads
1696 + .iter()
1697 + .position(|payload| payload.contains("\"finish_reason\":\"stop\""))
1698 + .expect("finish chunk");
1699 + let usage_index = payloads
1700 + .iter()
1701 + .position(|payload| payload.contains("\"usage\":"))
1702 + .expect("usage chunk");
1703 + let done_index =
1704 + payloads.iter().position(|payload| *payload == "[DONE]").expect("done sentinel");
1705 +
1706 + assert!(finish_index < usage_index, "{text}");
1707 + assert!(usage_index < done_index, "{text}");
1708 +
1709 + let usage_chunk: serde_json::Value =
1710 + serde_json::from_str(payloads[usage_index]).expect("usage chunk json");
1711 + assert_eq!(usage_chunk["choices"], json!([]));
1712 + assert_eq!(usage_chunk["usage"]["prompt_tokens"], 22);
1713 + assert_eq!(usage_chunk["usage"]["completion_tokens"], 3);
1714 + assert_eq!(usage_chunk["usage"]["total_tokens"], 25);
1715 +}
1716 +
1717 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1718 +#[serial]
1719 +async fn stream_without_include_usage_keeps_existing_shape() {
1720 + let (app, engine_task) = test_app_with_stream_output_specs(default_stream_output_specs()).await;
1721 + let response = app
1722 + .clone()
1723 + .call(
1724 + Request::builder()
1725 + .method("POST")
1726 + .uri("/v1/chat/completions")
1727 + .header("content-type", "application/json")
1728 + .body(Body::from(
1729 + json!({
1730 + "model": "Qwen/Qwen1.5-0.5B-Chat",
1731 + "stream": true,
1732 + "messages": [{"role": "user", "content": "hello"}]
1733 + })
1734 + .to_string(),
1735 + ))
1736 + .expect("build request"),
1737 + )
1738 + .await
1739 + .expect("call app");
1740 +
1741 + assert_eq!(response.status(), StatusCode::OK);
1742 +
1743 + let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");
1744 + engine_task.await.expect("mock engine task");
1745 + let text = String::from_utf8(body.to_vec()).expect("utf8 body");
1746 +
1747 + assert!(!text.contains("\"usage\":"), "{text}");
1748 + assert!(text.contains("\"finish_reason\":\"stop\""), "{text}");
1749 + assert!(text.trim_end().ends_with("data: [DONE]"), "{text}");
1750 +}
1751 +
1752 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1753 +#[serial]
1754 +async fn completions_invalid_request_returns_openai_error() {
1755 + let mut app = test_app().await;
1756 + let response = app
1757 + .call(
1758 + Request::builder()
1759 + .method("POST")
1760 + .uri("/v1/completions")
1761 + .header("content-type", "application/json")
1762 + .body(Body::from(
1763 + json!({
1764 + "model": "Qwen/Qwen1.5-0.5B-Chat",
1765 + "prompt": "hello",
1766 + "stream": false,
1767 + "stream_options": {"include_usage": true}
1768 + })
1769 + .to_string(),
1770 + ))
1771 + .expect("build request"),
1772 + )
1773 + .await
1774 + .expect("call app");
1775 +
1776 + assert_eq!(response.status(), StatusCode::BAD_REQUEST);
1777 + let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");
1778 + let json: serde_json::Value = serde_json::from_slice(&body).expect("decode json");
1779 + assert_eq!(json["error"]["type"], "invalid_request_error");
1780 +}
1781 +
1782 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1783 +#[serial]
1784 +async fn non_stream_completions_return_json_response() {
1785 + let (app, engine_task) = test_app_with_engine_handle().await;
1786 + let response = app
1787 + .clone()
1788 + .call(
1789 + Request::builder()
1790 + .method("POST")
1791 + .uri("/v1/completions")
1792 + .header("content-type", "application/json")
1793 + .body(Body::from(
1794 + json!({
1795 + "model": "Qwen/Qwen1.5-0.5B-Chat",
1796 + "prompt": "hello",
1797 + "stream": false
1798 + })
1799 + .to_string(),
1800 + ))
1801 + .expect("build request"),
1802 + )
1803 + .await
1804 + .expect("call app");
1805 +
1806 + assert_eq!(response.status(), StatusCode::OK);
1807 + assert!(
1808 + response
1809 + .headers()
1810 + .get("content-type")
1811 + .and_then(|value| value.to_str().ok())
1812 + .is_some_and(|value| value.starts_with("application/json"))
1813 + );
1814 +
1815 + let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");
1816 + engine_task.await.expect("mock engine task");
1817 + let json: serde_json::Value = serde_json::from_slice(&body).expect("decode json");
1818 +
1819 + assert_eq!(json["object"], "text_completion");
1820 + assert_eq!(json["choices"][0]["text"], "hi");
1821 + assert_eq!(json["choices"][0]["finish_reason"], "stop");
1822 + assert_eq!(json["usage"]["completion_tokens"], 3);
1823 +}
1824 +
1825 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1826 +#[serial]
1827 +async fn non_stream_completions_echo_prepends_prompt_text() {
1828 + let (app, engine_task) = test_app_with_engine_handle().await;
1829 + let response = app
1830 + .clone()
1831 + .call(
1832 + Request::builder()
1833 + .method("POST")
1834 + .uri("/v1/completions")
1835 + .header("content-type", "application/json")
1836 + .body(Body::from(
1837 + json!({
1838 + "model": "Qwen/Qwen1.5-0.5B-Chat",
1839 + "prompt": "hello",
1840 + "echo": true,
1841 + "stream": false
1842 + })
1843 + .to_string(),
1844 + ))
1845 + .expect("build request"),
1846 + )
1847 + .await
1848 + .expect("call app");
1849 +
1850 + assert_eq!(response.status(), StatusCode::OK);
1851 +
1852 + let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");
1853 + engine_task.await.expect("mock engine task");
1854 + let json: serde_json::Value = serde_json::from_slice(&body).expect("decode json");
1855 +
1856 + assert_eq!(json["choices"][0]["text"], "hellohi");
1857 + assert_eq!(json["usage"]["prompt_tokens"], 5);
1858 + assert_eq!(json["usage"]["completion_tokens"], 3);
1859 +}
1860 +
1861 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1862 +#[serial]
1863 +async fn non_stream_completions_include_logprobs() {
1864 + let ipc = IpcNamespace::new().expect("create ipc namespace");
1865 + let handshake_address = ipc.handshake_endpoint();
1866 + let engine_id = b"engine-openai-completion-logprobs".to_vec();
1867 +
1868 + let engine_task = MockEngineTask::new(spawn_mock_engine_task(
1869 + handshake_address.clone(),
1870 + engine_id.clone(),
1871 + |dealer, push| {
1872 + boxed_test_future(async move {
1873 + let add = recv_engine_message(dealer).await;
1874 + let request: EngineCoreRequest =
1875 + rmp_serde::from_slice(&add[1]).expect("decode request");
1876 + send_outputs(
1877 + push,
1878 + EngineCoreOutputs {
1879 + engine_index: 0,
1880 + outputs: vec![
1881 + request_output_with_logprobs(
1882 + &request.request_id,
1883 + vec![b'h' as u32],
1884 + None,
1885 + None,
1886 + Some(sample_logprobs_for_token(b'h' as u32, b'H' as u32)),
1887 + None,
1888 + ),
1889 + request_output_with_logprobs(
1890 + &request.request_id,
1891 + vec![b'i' as u32],
1892 + Some(EngineCoreFinishReason::Stop),
1893 + None,
1894 + Some(sample_logprobs_for_token(b'i' as u32, b'I' as u32)),
1895 + None,
1896 + ),
1897 + ],
1898 + scheduler_stats: None,
1899 + timestamp: 0.0,
1900 + utility_output: None,
1901 + finished_requests: Some(BTreeSet::from([request.request_id.clone()])),
1902 + wave_complete: None,
1903 + start_wave: None,
1904 + },
1905 + )
1906 + .await;
1907 + })
1908 + },
1909 + ));
1910 +
1911 + let client = EngineCoreClient::connect(
1912 + EngineCoreClientConfig::new_single(handshake_address)
1913 + .with_model_name("test-model")
1914 + .with_local_input_output_addresses(
1915 + Some(ipc.input_endpoint()),
1916 + Some(ipc.output_endpoint()),
1917 + ),
1918 + )
1919 + .await
1920 + .expect("connect client");
1921 + let chat = ChatLlm::from_shared_backend(test_llm(client), Arc::new(FakeChatBackend::new()));
1922 + let mut app = build_router(Arc::new(AppState::new(
1923 + vec!["Qwen/Qwen1.5-0.5B-Chat".to_string()],
1924 + chat,
1925 + )));
1926 +
1927 + let response = app
1928 + .call(
1929 + Request::builder()
1930 + .method("POST")
1931 + .uri("/v1/completions")
1932 + .header("content-type", "application/json")
1933 + .body(Body::from(
1934 + json!({
1935 + "model": "Qwen/Qwen1.5-0.5B-Chat",
1936 + "prompt": "hello",
1937 + "stream": false,
1938 + "logprobs": 1
1939 + })
1940 + .to_string(),
1941 + ))
1942 + .expect("build request"),
1943 + )
1944 + .await
1945 + .expect("call app");
1946 +
1947 + assert_eq!(response.status(), StatusCode::OK);
1948 + let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");
1949 + engine_task.await.expect("mock engine task");
1950 + let json: serde_json::Value = serde_json::from_slice(&body).expect("decode json");
1951 +
1952 + assert_eq!(json["choices"][0]["logprobs"]["tokens"], json!(["h", "i"]));
1953 + assert_eq!(
1954 + json["choices"][0]["logprobs"]["token_logprobs"],
1955 + json!([-0.1, -0.1])
1956 + );
1957 + assert_eq!(json["choices"][0]["logprobs"]["text_offset"], json!([0, 1]));
1958 + assert_eq!(
1959 + json["choices"][0]["logprobs"]["top_logprobs"][0],
1960 + json!({"h": -0.1, "H": -0.2})
1961 + );
1962 +}
1963 +
1964 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1965 +#[serial]
1966 +async fn non_stream_completions_include_prompt_logprobs() {
1967 + let ipc = IpcNamespace::new().expect("create ipc namespace");
1968 + let handshake_address = ipc.handshake_endpoint();
1969 + let engine_id = b"engine-openai-completion-prompt-logprobs".to_vec();
1970 +
1971 + let engine_task = MockEngineTask::new(spawn_mock_engine_task(
1972 + handshake_address.clone(),
1973 + engine_id.clone(),
1974 + |dealer, push| {
1975 + boxed_test_future(async move {
1976 + let add = recv_engine_message(dealer).await;
1977 + let request: EngineCoreRequest =
1978 + rmp_serde::from_slice(&add[1]).expect("decode request");
1979 + send_outputs(
1980 + push,
1981 + EngineCoreOutputs {
1982 + engine_index: 0,
1983 + outputs: vec![request_output_with_logprobs(
1984 + &request.request_id,
1985 + vec![b'h' as u32, b'i' as u32, b'!' as u32],
1986 + Some(EngineCoreFinishReason::Stop),
1987 + None,
1988 + Some(Logprobs {
1989 + positions: vec![
1990 + sample_logprobs_for_token(b'h' as u32, b'H' as u32).positions
1991 + [0]
1992 + .clone(),
1993 + sample_logprobs_for_token(b'i' as u32, b'I' as u32).positions
1994 + [0]
1995 + .clone(),
1996 + sample_logprobs_for_token(b'!' as u32, b'?' as u32).positions
1997 + [0]
1998 + .clone(),
1999 + ],
2000 + }),
2001 + Some(prompt_logprobs_for_hello()),
2002 + )],
2003 + scheduler_stats: None,
2004 + timestamp: 0.0,
2005 + utility_output: None,
2006 + finished_requests: None,
2007 + wave_complete: None,
2008 + start_wave: None,
2009 + },
2010 + )
2011 + .await;
2012 + })
2013 + },
2014 + ));
2015 +
2016 + let client = EngineCoreClient::connect(
2017 + EngineCoreClientConfig::new_single(handshake_address)
2018 + .with_model_name("test-model")
2019 + .with_local_input_output_addresses(
2020 + Some(ipc.input_endpoint()),
2021 + Some(ipc.output_endpoint()),
2022 + ),
2023 + )
2024 + .await
2025 + .expect("connect client");
2026 + let chat = ChatLlm::from_shared_backend(test_llm(client), Arc::new(FakeChatBackend::new()));
2027 + let mut app = build_router(Arc::new(AppState::new(
2028 + vec!["Qwen/Qwen1.5-0.5B-Chat".to_string()],
2029 + chat,
2030 + )));
2031 +
2032 + let response = app
2033 + .call(
2034 + Request::builder()
2035 + .method("POST")
2036 + .uri("/v1/completions")
2037 + .header("content-type", "application/json")
2038 + .body(Body::from(
2039 + json!({
2040 + "model": "Qwen/Qwen1.5-0.5B-Chat",
2041 + "prompt": "hello",
2042 + "stream": false,
2043 + "echo": true,
2044 + "logprobs": 1
2045 + })
2046 + .to_string(),
2047 + ))
2048 + .expect("build request"),
2049 + )
2050 + .await
2051 + .expect("call app");
2052 +
2053 + assert_eq!(response.status(), StatusCode::OK);
2054 + let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");
2055 + engine_task.await.expect("mock engine task");
2056 + let json: serde_json::Value = serde_json::from_slice(&body).expect("decode json");
2057 +
2058 + assert_eq!(json["choices"][0]["text"], "hellohi");
2059 + assert_eq!(
2060 + json["choices"][0]["logprobs"]["tokens"],
2061 + json!(["h", "e", "l", "l", "o", "h", "i", "!"])
2062 + );
2063 + assert_eq!(
2064 + json["choices"][0]["logprobs"]["text_offset"],
2065 + json!([0, 1, 2, 3, 4, 5, 6, 7])
2066 + );
2067 + assert_eq!(
2068 + json["choices"][0]["logprobs"]["token_logprobs"],
2069 + json!([null, -0.3, -0.4, -0.45, -0.5, -0.1, -0.1, -0.1])
2070 + );
2071 + assert_eq!(
2072 + json["choices"][0]["prompt_logprobs"][0],
2073 + serde_json::Value::Null
2074 + );
2075 + assert_eq!(
2076 + json["choices"][0]["prompt_logprobs"][1],
2077 + json!({"a": -0.5, "e": -0.3})
2078 + );
2079 +}
2080 +
2081 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
2082 +#[serial]
2083 +async fn non_stream_chat_completions_still_succeed() {
2084 + let ipc = IpcNamespace::new().expect("create ipc namespace");
2085 + let handshake_address = ipc.handshake_endpoint();
2086 + let engine_id = b"engine-openai-chat-non-stream".to_vec();
2087 +
2088 + let engine_task = MockEngineTask::new(spawn_mock_engine_task(
2089 + handshake_address.clone(),
2090 + engine_id.clone(),
2091 + |dealer, push| {
2092 + boxed_test_future(async move {
2093 + let add = recv_engine_message(dealer).await;
2094 + let request: EngineCoreRequest =
2095 + rmp_serde::from_slice(&add[1]).expect("decode request");
2096 + send_outputs(
2097 + push,
2098 + engine_outputs_for_request(&request.request_id, default_stream_output_specs()),
2099 + )
2100 + .await;
2101 + })
2102 + },
2103 + ));
2104 +
2105 + let client = EngineCoreClient::connect(
2106 + EngineCoreClientConfig::new_single(handshake_address)
2107 + .with_model_name("test-model")
2108 + .with_local_input_output_addresses(
2109 + Some(ipc.input_endpoint()),
2110 + Some(ipc.output_endpoint()),
2111 + ),
2112 + )
2113 + .await
2114 + .expect("connect client");
2115 + let chat = ChatLlm::from_shared_backend(test_llm(client), Arc::new(FakeChatBackend::new()));
2116 + let mut app = build_router(Arc::new(AppState::new(
2117 + vec!["Qwen/Qwen1.5-0.5B-Chat".to_string()],
2118 + chat,
2119 + )));
2120 +
2121 + let response = app
2122 + .call(
2123 + Request::builder()
2124 + .method("POST")
2125 + .uri("/v1/chat/completions")
2126 + .header("content-type", "application/json")
2127 + .body(Body::from(
2128 + json!({
2129 + "model": "Qwen/Qwen1.5-0.5B-Chat",
2130 + "stream": false,
2131 + "messages": [{"role": "user", "content": "hello"}]
2132 + })
2133 + .to_string(),
2134 + ))
2135 + .expect("build request"),
2136 + )
2137 + .await
2138 + .expect("call app");
2139 +
2140 + assert_eq!(response.status(), StatusCode::OK);
2141 + engine_task.await.expect("mock engine task");
2142 +}
2143 +
2144 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
2145 +#[serial]
2146 +async fn non_stream_completions_still_succeed() {
2147 + let ipc = IpcNamespace::new().expect("create ipc namespace");
2148 + let handshake_address = ipc.handshake_endpoint();
2149 + let engine_id = b"engine-openai-completion-non-stream".to_vec();
2150 +
2151 + let engine_task = MockEngineTask::new(spawn_mock_engine_task(
2152 + handshake_address.clone(),
2153 + engine_id.clone(),
2154 + |dealer, push| {
2155 + boxed_test_future(async move {
2156 + let add = recv_engine_message(dealer).await;
2157 + let request: EngineCoreRequest =
2158 + rmp_serde::from_slice(&add[1]).expect("decode request");
2159 + send_outputs(
2160 + push,
2161 + engine_outputs_for_request(&request.request_id, default_stream_output_specs()),
2162 + )
2163 + .await;
2164 + })
2165 + },
2166 + ));
2167 +
2168 + let client = EngineCoreClient::connect(
2169 + EngineCoreClientConfig::new_single(handshake_address)
2170 + .with_model_name("test-model")
2171 + .with_local_input_output_addresses(
2172 + Some(ipc.input_endpoint()),
2173 + Some(ipc.output_endpoint()),
2174 + ),
2175 + )
2176 + .await
2177 + .expect("connect client");
2178 + let chat = ChatLlm::from_shared_backend(test_llm(client), Arc::new(FakeChatBackend::new()));
2179 + let mut app = build_router(Arc::new(AppState::new(
2180 + vec!["Qwen/Qwen1.5-0.5B-Chat".to_string()],
2181 + chat,
2182 + )));
2183 +
2184 + let response = app
2185 + .call(
2186 + Request::builder()
2187 + .method("POST")
2188 + .uri("/v1/completions")
2189 + .header("content-type", "application/json")
2190 + .body(Body::from(
2191 + json!({
2192 + "model": "Qwen/Qwen1.5-0.5B-Chat",
2193 + "prompt": "hello",
2194 + "stream": false
2195 + })
2196 + .to_string(),
2197 + ))
2198 + .expect("build request"),
2199 + )
2200 + .await
2201 + .expect("call app");
2202 +
2203 + assert_eq!(response.status(), StatusCode::OK);
2204 + engine_task.await.expect("mock engine task");
2205 +}
2206 +
2207 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
2208 +#[serial]
2209 +async fn chat_completions_header_request_id_takes_precedence() {
2210 + let ipc = IpcNamespace::new().expect("create ipc namespace");
2211 + let handshake_address = ipc.handshake_endpoint();
2212 + let engine_id = b"engine-chat-request-id-precedence".to_vec();
2213 +
2214 + let engine_task = MockEngineTask::new(spawn_mock_engine_task(
2215 + handshake_address.clone(),
2216 + engine_id.clone(),
2217 + |dealer, push| {
2218 + boxed_test_future(async move {
2219 + let add = recv_engine_message(dealer).await;
2220 + let request: EngineCoreRequest =
2221 + rmp_serde::from_slice(&add[1]).expect("decode request");
2222 + assert_eq!(
2223 + request.external_req_id.as_deref(),
2224 + Some("chatcmpl-header-req")
2225 + );
2226 + assert!(request.request_id.starts_with("chatcmpl-header-req-"));
2227 + assert_ne!(request.request_id, "chatcmpl-header-req");
2228 +
2229 + send_outputs(
2230 + push,
2231 + engine_outputs_for_request(&request.request_id, default_stream_output_specs()),
2232 + )
2233 + .await;
2234 + })
2235 + },
2236 + ));
2237 +
2238 + let client = EngineCoreClient::connect(
2239 + EngineCoreClientConfig::new_single(handshake_address)
2240 + .with_model_name("test-model")
2241 + .with_local_input_output_addresses(
2242 + Some(ipc.input_endpoint()),
2243 + Some(ipc.output_endpoint()),
2244 + ),
2245 + )
2246 + .await
2247 + .expect("connect client");
2248 + let chat = ChatLlm::from_shared_backend(Llm::new(client), Arc::new(FakeChatBackend::new()));
2249 + let mut app = build_router(Arc::new(AppState::new(
2250 + vec!["Qwen/Qwen1.5-0.5B-Chat".to_string()],
2251 + chat,
2252 + )));
2253 +
2254 + let response = app
2255 + .call(
2256 + Request::builder()
2257 + .method("POST")
2258 + .uri("/v1/chat/completions")
2259 + .header("content-type", "application/json")
2260 + .header("X-Request-Id", "header-req")
2261 + .body(Body::from(
2262 + json!({
2263 + "request_id": "body-req",
2264 + "model": "Qwen/Qwen1.5-0.5B-Chat",
2265 + "messages": [{"role": "user", "content": "hello"}]
2266 + })
2267 + .to_string(),
2268 + ))
2269 + .expect("build request"),
2270 + )
2271 + .await
2272 + .expect("call app");
2273 +
2274 + assert_eq!(response.status(), StatusCode::OK);
2275 + let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");
2276 + engine_task.await.expect("mock engine task");
2277 + let json: serde_json::Value = serde_json::from_slice(&body).expect("decode json");
2278 +
2279 + assert_eq!(json["id"], "chatcmpl-header-req");
2280 +}
2281 +
2282 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
2283 +#[serial]
2284 +async fn non_stream_raw_generate_returns_token_output_envelope() {
2285 + let ipc = IpcNamespace::new().expect("create ipc namespace");
2286 + let handshake_address = ipc.handshake_endpoint();
2287 + let engine_id = b"engine-raw-generate-non-stream".to_vec();
2288 +
2289 + let engine_task = MockEngineTask::new(spawn_mock_engine_task(
2290 + handshake_address.clone(),
2291 + engine_id.clone(),
2292 + |dealer, push| {
2293 + boxed_test_future(async move {
2294 + let add = recv_engine_message(dealer).await;
2295 + let request: EngineCoreRequest =
2296 + rmp_serde::from_slice(&add[1]).expect("decode request");
2297 + assert_eq!(request.prompt_token_ids.as_deref(), Some(&[11, 22][..]));
2298 + assert_eq!(request.external_req_id.as_deref(), Some("raw-req"));
2299 + assert!(request.request_id.starts_with("raw-req-"));
2300 + assert_ne!(request.request_id, "raw-req");
2301 +
2302 + send_outputs(
2303 + push,
2304 + EngineCoreOutputs {
2305 + engine_index: 0,
2306 + outputs: vec![
2307 + request_output_with_logprobs(
2308 + &request.request_id,
2309 + vec![33],
2310 + None,
2311 + None,
2312 + Some(sample_logprobs_for_token(33, 34)),
2313 + Some(prompt_logprobs_for_tokens(&[11, 22])),
2314 + ),
2315 + request_output_with_logprobs_and_kv(
2316 + &request.request_id,
2317 + vec![44],
2318 + Some(EngineCoreFinishReason::Stop),
2319 + None,
2320 + Some(sample_logprobs_for_token(44, 45)),
2321 + None,
2322 + Some(json!({"connector": "x"})),
2323 + ),
2324 + ],
2325 + scheduler_stats: None,
2326 + timestamp: 0.0,
2327 + utility_output: None,
2328 + finished_requests: None,
2329 + wave_complete: None,
2330 + start_wave: None,
2331 + },
2332 + )
2333 + .await;
2334 + })
2335 + },
2336 + ));
2337 +
2338 + let client = EngineCoreClient::connect(
2339 + EngineCoreClientConfig::new_single(handshake_address)
2340 + .with_model_name("test-model")
2341 + .with_local_input_output_addresses(
2342 + Some(ipc.input_endpoint()),
2343 + Some(ipc.output_endpoint()),
2344 + ),
2345 + )
2346 + .await
2347 + .expect("connect client");
2348 + let chat = ChatLlm::from_shared_backend(Llm::new(client), Arc::new(FakeChatBackend::new()));
2349 + let mut app = build_router(Arc::new(AppState::new(
2350 + vec!["Qwen/Qwen1.5-0.5B-Chat".to_string()],
2351 + chat,
2352 + )));
2353 +
2354 + let response = app
2355 + .call(
2356 + Request::builder()
2357 + .method("POST")
2358 + .uri("/inference/v1/generate")
2359 + .header("content-type", "application/json")
2360 + .body(Body::from(
2361 + json!({
2362 + "request_id": "raw-req",
2363 + "model": "Qwen/Qwen1.5-0.5B-Chat",
2364 + "token_ids": [11, 22],
2365 + "stream": false,
2366 + "sampling_params": {
2367 + "max_tokens": 2,
2368 + "logprobs": 1,
2369 + "prompt_logprobs": 1
2370 + }
2371 + })
2372 + .to_string(),
2373 + ))
2374 + .expect("build request"),
2375 + )
2376 + .await
2377 + .expect("call app");
2378 +
2379 + assert_eq!(response.status(), StatusCode::OK);
2380 + let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");
2381 + engine_task.await.expect("mock engine task");
2382 + let json: serde_json::Value = serde_json::from_slice(&body).expect("decode json");
2383 +
2384 + assert_eq!(json["request_id"], "raw-req");
2385 + assert_eq!(json["choices"][0]["index"], 0);
2386 + assert_eq!(json["choices"][0]["token_ids"], json!([33, 44]));
2387 + assert_eq!(json["choices"][0]["finish_reason"], "stop");
2388 + assert_eq!(
2389 + json["choices"][0]["logprobs"]["content"][0]["token"],
2390 + "token_id:33"
2391 + );
2392 + assert_eq!(
2393 + json["choices"][0]["logprobs"]["content"][1]["top_logprobs"][0]["token"],
2394 + "token_id:44"
2395 + );
2396 + assert_eq!(json["prompt_logprobs"][0], serde_json::Value::Null);
2397 + assert_eq!(
2398 + json["prompt_logprobs"][1]["22"]["decoded_token"],
2399 + "token_id:22"
2400 + );
2401 + assert_eq!(json["kv_transfer_params"], json!({"connector": "x"}));
2402 +}
2403 +
2404 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
2405 +#[serial]
2406 +async fn raw_generate_rejects_streaming() {
2407 + let mut app = test_app().await;
2408 +
2409 + let response = app
2410 + .call(
2411 + Request::builder()
2412 + .method("POST")
2413 + .uri("/inference/v1/generate")
2414 + .header("content-type", "application/json")
2415 + .body(Body::from(
2416 + json!({
2417 + "model": "Qwen/Qwen1.5-0.5B-Chat",
2418 + "token_ids": [11, 22],
2419 + "stream": true,
2420 + "sampling_params": {}
2421 + })
2422 + .to_string(),
2423 + ))
2424 + .expect("build request"),
2425 + )
2426 + .await
2427 + .expect("call app");
2428 +
2429 + assert_eq!(response.status(), StatusCode::BAD_REQUEST);
2430 + let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");
2431 + let json: serde_json::Value = serde_json::from_slice(&body).expect("decode json");
2432 + assert_eq!(json["error"]["param"], "stream");
2433 +}
2434 +
2435 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
2436 +#[serial]
2437 +async fn raw_generate_rejects_empty_token_ids() {
2438 + let mut app = test_app().await;
2439 +
2440 + let response = app
2441 + .call(
2442 + Request::builder()
2443 + .method("POST")
2444 + .uri("/inference/v1/generate")
2445 + .header("content-type", "application/json")
2446 + .body(Body::from(
2447 + json!({
2448 + "model": "Qwen/Qwen1.5-0.5B-Chat",
2449 + "token_ids": [],
2450 + "sampling_params": {}
2451 + })
2452 + .to_string(),
2453 + ))
2454 + .expect("build request"),
2455 + )
2456 + .await
2457 + .expect("call app");
2458 +
2459 + assert_eq!(response.status(), StatusCode::BAD_REQUEST);
2460 + let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");
2461 + let json: serde_json::Value = serde_json::from_slice(&body).expect("decode json");
2462 + assert_eq!(json["error"]["param"], "token_ids");
2463 +}
2464 +
2465 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
2466 +#[serial]
2467 +async fn raw_generate_rejects_wrong_model() {
2468 + let mut app = test_app().await;
2469 +
2470 + let response = app
2471 + .call(
2472 + Request::builder()
2473 + .method("POST")
2474 + .uri("/inference/v1/generate")
2475 + .header("content-type", "application/json")
2476 + .body(Body::from(
2477 + json!({
2478 + "model": "wrong-model",
2479 + "token_ids": [11, 22],
2480 + "sampling_params": {}
2481 + })
2482 + .to_string(),
2483 + ))
2484 + .expect("build request"),
2485 + )
2486 + .await
2487 + .expect("call app");
2488 +
2489 + assert_eq!(response.status(), StatusCode::NOT_FOUND);
2490 +}
2491 +
2492 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
2493 +#[serial]
2494 +async fn completions_happy_path_returns_sse_stream() {
2495 + let (app, engine_task) = test_app_with_engine_handle().await;
2496 + let response = app
2497 + .clone()
2498 + .call(
2499 + Request::builder()
2500 + .method("POST")
2501 + .uri("/v1/completions")
2502 + .header("content-type", "application/json")
2503 + .body(Body::from(
2504 + json!({
2505 + "model": "Qwen/Qwen1.5-0.5B-Chat",
2506 + "prompt": "hello",
2507 + "stream": true,
2508 + "stream_options": {"include_usage": true}
2509 + })
2510 + .to_string(),
2511 + ))
2512 + .expect("build request"),
2513 + )
2514 + .await
2515 + .expect("call app");
2516 +
2517 + assert_eq!(response.status(), StatusCode::OK);
2518 + assert_eq!(
2519 + response.headers().get("content-type").and_then(|value| value.to_str().ok()),
2520 + Some("text/event-stream")
2521 + );
2522 +
2523 + let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");
2524 + engine_task.await.expect("mock engine task");
2525 + let text = String::from_utf8(body.to_vec()).expect("utf8 body");
2526 + let payloads = sse_data_payloads(&text);
2527 + let usage_index = payloads
2528 + .iter()
2529 + .position(|payload| payload.contains("\"usage\":"))
2530 + .expect("usage chunk");
2531 + let done_index =
2532 + payloads.iter().position(|payload| *payload == "[DONE]").expect("done sentinel");
2533 +
2534 + assert!(
2535 + payloads.iter().any(|payload| payload.contains("\"text\":\"h\"")),
2536 + "{text}"
2537 + );
2538 + assert!(
2539 + payloads.iter().any(|payload| payload.contains("\"finish_reason\":\"stop\"")),
2540 + "{text}"
2541 + );
2542 + assert!(usage_index < done_index, "{text}");
2543 +
2544 + let usage_chunk: serde_json::Value =
2545 + serde_json::from_str(payloads[usage_index]).expect("usage chunk json");
2546 + assert_eq!(usage_chunk["choices"], json!([]));
2547 + assert_eq!(usage_chunk["usage"]["completion_tokens"], 3);
2548 +}
2549 +
2550 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
2551 +#[serial]
2552 +async fn completions_echo_stream_emits_separate_prompt_chunk() {
2553 + let (app, engine_task) = test_app_with_engine_handle().await;
2554 + let response = app
2555 + .clone()
2556 + .call(
2557 + Request::builder()
2558 + .method("POST")
2559 + .uri("/v1/completions")
2560 + .header("content-type", "application/json")
2561 + .body(Body::from(
2562 + json!({
2563 + "model": "Qwen/Qwen1.5-0.5B-Chat",
2564 + "prompt": "hello",
2565 + "echo": true,
2566 + "stream": true,
2567 + "stream_options": {"include_usage": true}
2568 + })
2569 + .to_string(),
2570 + ))
2571 + .expect("build request"),
2572 + )
2573 + .await
2574 + .expect("call app");
2575 +
2576 + assert_eq!(response.status(), StatusCode::OK);
2577 +
2578 + let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");
2579 + engine_task.await.expect("mock engine task");
2580 + let text = String::from_utf8(body.to_vec()).expect("utf8 body");
2581 + let payloads = sse_data_payloads(&text);
2582 + let hello_index = payloads
2583 + .iter()
2584 + .position(|payload| payload.contains("\"text\":\"hello\""))
2585 + .expect("prompt echo chunk");
2586 + let h_index = payloads
2587 + .iter()
2588 + .position(|payload| payload.contains("\"text\":\"h\""))
2589 + .expect("first generation chunk");
2590 +
2591 + assert!(hello_index < h_index, "{text}");
2592 + assert!(
2593 + payloads.iter().any(|payload| payload.contains("\"text\":\"i\"")),
2594 + "{text}"
2595 + );
2596 +
2597 + let usage_chunk: serde_json::Value = serde_json::from_str(
2598 + payloads
2599 + .iter()
2600 + .find(|payload| payload.contains("\"usage\":"))
2601 + .expect("usage chunk"),
2602 + )
2603 + .expect("usage chunk json");
2604 + assert_eq!(usage_chunk["usage"]["prompt_tokens"], 5);
2605 + assert_eq!(usage_chunk["usage"]["completion_tokens"], 3);
2606 +}
2607 +
2608 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
2609 +#[serial]
2610 +async fn chat_harness_streams_text_events() {
2611 + let (chat, engine_task) = test_chat_with_engine_handle().await;
2612 + let mut stream = chat
2613 + .chat(ChatRequest {
2614 + messages: vec![ChatMessage::text(ChatRole::User, "hello")],
2615 + sampling_params: SamplingParams {
2616 + max_tokens: Some(8),
2617 + ..Default::default()
2618 + },
2619 + request_id: "chat-harness".to_string(),
2620 + ..ChatRequest::for_test()
2621 + })
2622 + .await
2623 + .expect("submit chat request");
2624 +
2625 + let mut saw_text = false;
2626 + let mut saw_done = false;
2627 + while let Some(event) = stream.next().await {
2628 + match event.expect("chat event") {
2629 + ChatEvent::BlockDelta { .. } => saw_text = true,
2630 + ChatEvent::Done { .. } => {
2631 + saw_done = true;
2632 + break;
2633 + }
2634 + ChatEvent::Start { .. }
2635 + | ChatEvent::LogprobsDelta { .. }
2636 + | ChatEvent::BlockStart { .. }
2637 + | ChatEvent::BlockEnd { .. }
2638 + | ChatEvent::ToolCallStart { .. }
2639 + | ChatEvent::ToolCallArgumentsDelta { .. }
2640 + | ChatEvent::ToolCallEnd { .. } => {}
2641 + }
2642 + }
2643 + engine_task.await.expect("mock engine task");
2644 +
2645 + assert!(saw_text);
2646 + assert!(saw_done);
2647 +}
2648 +
2649 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
2650 +#[serial]
2651 +async fn prepared_openai_request_streams_text_events() {
2652 + let (chat, engine_task) = test_chat_with_engine_handle().await;
2653 + let prepared = prepare_chat_request(
2654 + serde_json::from_value(json!({
2655 + "model": "Qwen/Qwen1.5-0.5B-Chat",
2656 + "stream": true,
2657 + "messages": [{"role": "user", "content": "hello"}]
2658 + }))
2659 + .expect("decode request"),
2660 + &["Qwen/Qwen1.5-0.5B-Chat".to_string()],
2661 + crate::utils::ResolvedRequestContext::default(),
2662 + )
2663 + .expect("prepare request");
2664 +
2665 + let mut stream = chat.chat(prepared.chat_request).await.expect("submit chat request");
2666 +
2667 + let mut saw_text = false;
2668 + let mut saw_done = false;
2669 + while let Some(event) = stream.next().await {
2670 + match event.expect("chat event") {
2671 + ChatEvent::BlockDelta { .. } => saw_text = true,
2672 + ChatEvent::Done { .. } => {
2673 + saw_done = true;
2674 + break;
2675 + }
2676 + ChatEvent::Start { .. }
2677 + | ChatEvent::LogprobsDelta { .. }
2678 + | ChatEvent::BlockStart { .. }
2679 + | ChatEvent::BlockEnd { .. }
2680 + | ChatEvent::ToolCallStart { .. }
2681 + | ChatEvent::ToolCallArgumentsDelta { .. }
2682 + | ChatEvent::ToolCallEnd { .. } => {}
2683 + }
2684 + }
2685 + engine_task.await.expect("mock engine task");
2686 +
2687 + assert!(saw_text);
2688 + assert!(saw_done);
2689 +}
2690 +
2691 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
2692 +#[serial]
2693 +async fn reasoning_blocks_are_mapped_to_reasoning_sse_chunks() {
2694 + let (app, engine_task) = test_app_with_backend_and_stream_output_specs(
2695 + Arc::new(FakeChatBackend::with_model_id("Qwen/Qwen3-0.6B")),
2696 + vec![
2697 + (bytes_to_token_ids(b"<think>"), None),
2698 + (bytes_to_token_ids(b"think "), None),
2699 + (bytes_to_token_ids(b"more</think>"), None),
2700 + (
2701 + bytes_to_token_ids(b"answer"),
2702 + Some(EngineCoreFinishReason::Length),
2703 + ),
2704 + ],
2705 + )
2706 + .await;
2707 +
2708 + let response = app
2709 + .clone()
2710 + .call(
2711 + Request::builder()
2712 + .method("POST")
2713 + .uri("/v1/chat/completions")
2714 + .header("content-type", "application/json")
2715 + .body(Body::from(
2716 + json!({
2717 + "model": "Qwen/Qwen1.5-0.5B-Chat",
2718 + "stream": true,
2719 + "messages": [{"role": "user", "content": "hello"}]
2720 + })
2721 + .to_string(),
2722 + ))
2723 + .expect("build request"),
2724 + )
2725 + .await
2726 + .expect("call app");
2727 +
2728 + assert_eq!(response.status(), StatusCode::OK);
2729 + let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");
2730 + engine_task.await.expect("mock engine task");
2731 + let text = String::from_utf8(body.to_vec()).expect("utf8 body");
2732 +
2733 + assert!(text.contains("\"reasoning\":\"think \""), "{text}");
2734 + assert!(text.contains("\"reasoning\":\"more\""), "{text}");
2735 + assert!(text.contains("\"content\":\"answer\""), "{text}");
2736 +}
2737 +
2738 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
2739 +#[serial]
2740 +async fn tool_calls_are_mapped_to_tool_call_sse_chunks() {
2741 + let (app, engine_task) = test_app_with_backend_and_stream_output_specs(
2742 + Arc::new(FakeChatBackend::with_model_id("Qwen/Qwen3-0.6B")),
2743 + vec![
2744 + (bytes_to_token_ids(b"<think>Need tool.</think>"), None),
2745 + (
2746 + bytes_to_token_ids(b"<tool_call>\n{\"name\":\"get_weather\", "),
2747 + None,
2748 + ),
2749 + (
2750 + bytes_to_token_ids(b"\"arguments\":{\"city\":\"Paris\"}}\n</tool_call>"),
2751 + Some(EngineCoreFinishReason::Stop),
2752 + ),
2753 + ],
2754 + )
2755 + .await;
2756 +
2757 + let response = app
2758 + .clone()
2759 + .call(
2760 + Request::builder()
2761 + .method("POST")
2762 + .uri("/v1/chat/completions")
2763 + .header("content-type", "application/json")
2764 + .body(Body::from(
2765 + json!({
2766 + "model": "Qwen/Qwen1.5-0.5B-Chat",
2767 + "stream": true,
2768 + "messages": [{"role": "user", "content": "hello"}],
2769 + "tools": [{
2770 + "type": "function",
2771 + "function": {
2772 + "name": "get_weather",
2773 + "description": "Get weather",
2774 + "parameters": {
2775 + "type": "object",
2776 + "properties": {"city": {"type": "string"}}
2777 + }
2778 + }
2779 + }]
2780 + })
2781 + .to_string(),
2782 + ))
2783 + .expect("build request"),
2784 + )
2785 + .await
2786 + .expect("call app");
2787 +
2788 + assert_eq!(response.status(), StatusCode::OK);
2789 + let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");
2790 + engine_task.await.expect("mock engine task");
2791 + let text = String::from_utf8(body.to_vec()).expect("utf8 body");
2792 +
2793 + assert!(text.contains("\"tool_calls\":"), "{text}");
2794 + assert!(text.contains("\"name\":\"get_weather\""), "{text}");
2795 + assert!(
2796 + text.contains("\"arguments\":\"{\\\"city\\\":\\\"Paris\\\"}\""),
2797 + "{text}"
2798 + );
2799 + assert!(text.contains("\"finish_reason\":\"tool_calls\""), "{text}");
2800 +}
2801 +
2802 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
2803 +#[serial]
2804 +async fn tool_call_sse_chunks_can_carry_logprobs() {
2805 + let ipc = IpcNamespace::new().expect("create ipc namespace");
2806 + let handshake_address = ipc.handshake_endpoint();
2807 + let engine_id = b"engine-openai-chat-tools-logprobs".to_vec();
2808 +
2809 + let engine_task = MockEngineTask::new(spawn_mock_engine_task(
2810 + handshake_address.clone(),
2811 + engine_id.clone(),
2812 + |dealer, push| {
2813 + Box::pin(async move {
2814 + let add = recv_engine_message(dealer).await;
2815 + let request: EngineCoreRequest =
2816 + rmp_serde::from_slice(&add[1]).expect("decode request");
2817 +
2818 + send_outputs(
2819 + push,
2820 + EngineCoreOutputs {
2821 + engine_index: 0,
2822 + outputs: vec![request_output_with_logprobs(
2823 + &request.request_id,
2824 + bytes_to_token_ids(b"<think>Need tool.</think>"),
2825 + None,
2826 + None,
2827 + Some(sample_logprobs_for_tokens(&bytes_to_token_ids(
2828 + b"<think>Need tool.</think>",
2829 + ))),
2830 + None,
2831 + )],
2832 + scheduler_stats: None,
2833 + timestamp: 0.0,
2834 + utility_output: None,
2835 + finished_requests: None,
2836 + wave_complete: None,
2837 + start_wave: None,
2838 + },
2839 + )
2840 + .await;
2841 + send_outputs(
2842 + push,
2843 + EngineCoreOutputs {
2844 + engine_index: 0,
2845 + outputs: vec![request_output_with_logprobs(
2846 + &request.request_id,
2847 + bytes_to_token_ids(b"<tool_call>\n{\"name\":\"get_weather\", "),
2848 + None,
2849 + None,
2850 + Some(sample_logprobs_for_tokens(&bytes_to_token_ids(
2851 + b"<tool_call>\n{\"name\":\"get_weather\", ",
2852 + ))),
2853 + None,
2854 + )],
2855 + scheduler_stats: None,
2856 + timestamp: 0.0,
2857 + utility_output: None,
2858 + finished_requests: None,
2859 + wave_complete: None,
2860 + start_wave: None,
2861 + },
2862 + )
2863 + .await;
2864 + send_outputs(
2865 + push,
2866 + EngineCoreOutputs {
2867 + engine_index: 0,
2868 + outputs: vec![request_output_with_logprobs(
2869 + &request.request_id,
2870 + bytes_to_token_ids(
2871 + b"\"arguments\":{\"city\":\"Paris\"}}\n</tool_call>",
2872 + ),
2873 + Some(EngineCoreFinishReason::Stop),
2874 + None,
2875 + Some(sample_logprobs_for_tokens(&bytes_to_token_ids(
2876 + b"\"arguments\":{\"city\":\"Paris\"}}\n</tool_call>",
2877 + ))),
2878 + None,
2879 + )],
2880 + scheduler_stats: None,
2881 + timestamp: 0.0,
2882 + utility_output: None,
2883 + finished_requests: Some(BTreeSet::from([request.request_id.clone()])),
2884 + wave_complete: None,
2885 + start_wave: None,
2886 + },
2887 + )
2888 + .await;
2889 + })
2890 + },
2891 + ));
2892 +
2893 + let client = EngineCoreClient::connect(
2894 + EngineCoreClientConfig::new_single(handshake_address)
2895 + .with_model_name("test-model")
2896 + .with_local_input_output_addresses(
2897 + Some(ipc.input_endpoint()),
2898 + Some(ipc.output_endpoint()),
2899 + ),
2900 + )
2901 + .await
2902 + .expect("connect client");
2903 + let chat = ChatLlm::from_shared_backend(
2904 + test_llm(client),
2905 + Arc::new(FakeChatBackend::with_model_id("Qwen/Qwen3-0.6B")),
2906 + );
2907 + let app = build_router(Arc::new(AppState::new(
2908 + vec!["Qwen/Qwen1.5-0.5B-Chat".to_string()],
2909 + chat,
2910 + )));
2911 +
2912 + let response = app
2913 + .clone()
2914 + .call(
2915 + Request::builder()
2916 + .method("POST")
2917 + .uri("/v1/chat/completions")
2918 + .header("content-type", "application/json")
2919 + .body(Body::from(
2920 + json!({
2921 + "model": "Qwen/Qwen1.5-0.5B-Chat",
2922 + "stream": true,
2923 + "logprobs": true,
2924 + "messages": [{"role": "user", "content": "hello"}],
2925 + "tools": [{
2926 + "type": "function",
2927 + "function": {
2928 + "name": "get_weather",
2929 + "description": "Get weather",
2930 + "parameters": {
2931 + "type": "object",
2932 + "properties": {"city": {"type": "string"}}
2933 + }
2934 + }
2935 + }]
2936 + })
2937 + .to_string(),
2938 + ))
2939 + .expect("build request"),
2940 + )
2941 + .await
2942 + .expect("call app");
2943 +
2944 + assert_eq!(response.status(), StatusCode::OK);
2945 + let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");
2946 + engine_task.finish().await;
2947 + let text = String::from_utf8(body.to_vec()).expect("utf8 body");
2948 +
2949 + assert!(text.contains("\"tool_calls\":"), "{text}");
2950 + assert!(text.contains("\"logprobs\":{\"content\":"), "{text}");
2951 +}
2952 +
2953 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
2954 +#[serial]
2955 +async fn streaming_chat_prompt_logprobs_are_rejected() {
2956 + let (app, engine_task) = test_app_with_engine_handle().await;
2957 + let response = app
2958 + .clone()
2959 + .call(
2960 + Request::builder()
2961 + .method("POST")
2962 + .uri("/v1/chat/completions")
2963 + .header("content-type", "application/json")
2964 + .body(Body::from(
2965 + json!({
2966 + "model": "Qwen/Qwen1.5-0.5B-Chat",
2967 + "stream": true,
2968 + "prompt_logprobs": 1,
2969 + "messages": [{"role": "user", "content": "hello"}]
2970 + })
2971 + .to_string(),
2972 + ))
2973 + .expect("build request"),
2974 + )
2975 + .await
2976 + .expect("call app");
2977 +
2978 + assert_eq!(response.status(), StatusCode::BAD_REQUEST);
2979 + engine_task.abort();
2980 +}
2981 +
2982 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
2983 +#[serial]
2984 +async fn reset_prefix_cache_route_sends_expected_utility_call() {
2985 + let (app, engine_task) = test_admin_app_with_engine_script(|dealer, push| {
2986 + boxed_test_future(async move {
2987 + let utility = recv_engine_message(dealer).await;
2988 + assert_eq!(utility[0].as_ref(), &[0x03]);
2989 +
2990 + let payload = decode_value(&utility[1]).expect("decode utility payload");
2991 + let array = payload.as_array().expect("utility payload array");
2992 + let call_id = array[1].as_i64().expect("call id");
2993 +
2994 + assert_eq!(array[2], Value::from("reset_prefix_cache"));
2995 + assert_eq!(
2996 + array[3],
2997 + Value::Array(vec![Value::from(true), Value::from(true)])
2998 + );
2999 +
3000 + send_outputs(push, utility_outputs(call_id, utility_result_value(true))).await;
3001 + })
3002 + })
3003 + .await;
3004 +
3005 + let response = app
3006 + .clone()
3007 + .call(
3008 + Request::builder()
3009 + .method("POST")
3010 + .uri("/reset_prefix_cache?reset_running_requests=true&reset_external=true")
3011 + .body(Body::empty())
3012 + .expect("build request"),
3013 + )
3014 + .await
3015 + .expect("call app");
3016 +
3017 + let status = response.status();
3018 + let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");
3019 + assert_eq!(status, StatusCode::OK, "{}", String::from_utf8_lossy(&body));
3020 + assert!(body.is_empty());
3021 + engine_task.await.expect("mock engine task");
3022 +}
3023 +
3024 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
3025 +#[serial]
3026 +async fn reset_mm_cache_route_sends_expected_utility_call() {
3027 + let (app, engine_task) = test_admin_app_with_engine_script(|dealer, push| {
3028 + boxed_test_future(async move {
3029 + let utility = recv_engine_message(dealer).await;
3030 + assert_eq!(utility[0].as_ref(), &[0x03]);
3031 +
3032 + let payload = decode_value(&utility[1]).expect("decode utility payload");
3033 + let array = payload.as_array().expect("utility payload array");
3034 + let call_id = array[1].as_i64().expect("call id");
3035 +
3036 + assert_eq!(array[2], Value::from("reset_mm_cache"));
3037 + assert_eq!(array[3], Value::Array(Vec::new()));
3038 +
3039 + send_outputs(push, utility_outputs(call_id, utility_none_result())).await;
3040 + })
3041 + })
3042 + .await;
3043 +
3044 + let response = app
3045 + .clone()
3046 + .call(
3047 + Request::builder()
3048 + .method("POST")
3049 + .uri("/reset_mm_cache")
3050 + .body(Body::empty())
3051 + .expect("build request"),
3052 + )
3053 + .await
3054 + .expect("call app");
3055 +
3056 + let status = response.status();
3057 + let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");
3058 + assert_eq!(status, StatusCode::OK, "{}", String::from_utf8_lossy(&body));
3059 + assert!(body.is_empty());
3060 + engine_task.await.expect("mock engine task");
3061 +}
3062 +
3063 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
3064 +#[serial]
3065 +async fn reset_encoder_cache_route_sends_expected_utility_call() {
3066 + let (app, engine_task) = test_admin_app_with_engine_script(|dealer, push| {
3067 + boxed_test_future(async move {
3068 + let utility = recv_engine_message(dealer).await;
3069 + assert_eq!(utility[0].as_ref(), &[0x03]);
3070 +
3071 + let payload = decode_value(&utility[1]).expect("decode utility payload");
3072 + let array = payload.as_array().expect("utility payload array");
3073 + let call_id = array[1].as_i64().expect("call id");
3074 +
3075 + assert_eq!(array[2], Value::from("reset_encoder_cache"));
3076 + assert_eq!(array[3], Value::Array(Vec::new()));
3077 +
3078 + send_outputs(push, utility_outputs(call_id, utility_none_result())).await;
3079 + })
3080 + })
3081 + .await;
3082 +
3083 + let response = app
3084 + .clone()
3085 + .call(
3086 + Request::builder()
3087 + .method("POST")
3088 + .uri("/reset_encoder_cache")
3089 + .body(Body::empty())
3090 + .expect("build request"),
3091 + )
3092 + .await
3093 + .expect("call app");
3094 +
3095 + let status = response.status();
3096 + let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");
3097 + assert_eq!(status, StatusCode::OK, "{}", String::from_utf8_lossy(&body));
3098 + assert!(body.is_empty());
3099 + engine_task.await.expect("mock engine task");
3100 +}
3101 +
3102 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
3103 +#[serial]
3104 +async fn collective_rpc_route_sends_expected_utility_call_and_returns_results() {
3105 + let (app, engine_task) = test_admin_app_with_engine_script(|dealer, push| {
3106 + boxed_test_future(async move {
3107 + let utility = recv_engine_message(dealer).await;
3108 + assert_eq!(utility[0].as_ref(), &[0x03]);
3109 +
3110 + let payload = decode_value(&utility[1]).expect("decode utility payload");
3111 + let array = payload.as_array().expect("utility payload array");
3112 + let call_id = array[1].as_i64().expect("call id");
3113 +
3114 + assert_eq!(array[2], Value::from("collective_rpc"));
3115 + assert_eq!(
3116 + array[3],
3117 + Value::Array(vec![
3118 + Value::from("echo_args_kwargs"),
3119 + Value::from(1.5_f64),
3120 + Value::Array(vec![Value::from("arg1"), Value::from("arg2")]),
3121 + Value::Map(vec![
3122 + (Value::from("key1"), Value::from("value1")),
3123 + (Value::from("key2"), Value::from("value2")),
3124 + ]),
3125 + ])
3126 + );
3127 +
3128 + send_outputs(
3129 + push,
3130 + utility_outputs(
3131 + call_id,
3132 + UtilityResultEnvelope::without_type_info(Value::Array(vec![Value::Map(vec![
3133 + (
3134 + Value::from("args"),
3135 + Value::Array(vec![Value::from("arg1"), Value::from("arg2")]),
3136 + ),
3137 + (
3138 + Value::from("kwargs"),
3139 + Value::Map(vec![
3140 + (Value::from("key1"), Value::from("value1")),
3141 + (Value::from("key2"), Value::from("value2")),
3142 + ]),
3143 + ),
3144 + (Value::from("total_items"), Value::from(4_u64)),
3145 + ])])),
3146 + ),
3147 + )
3148 + .await;
3149 + })
3150 + })
3151 + .await;
3152 +
3153 + let response = app
3154 + .clone()
3155 + .call(
3156 + Request::builder()
3157 + .method("POST")
3158 + .uri("/collective_rpc")
3159 + .header("content-type", "application/json")
3160 + .body(Body::from(
3161 + r#"{"method":"echo_args_kwargs","args":["arg1","arg2"],"kwargs":{"key1":"value1","key2":"value2"},"timeout":1.5}"#,
3162 + ))
3163 + .expect("build request"),
3164 + )
3165 + .await
3166 + .expect("call app");
3167 +
3168 + assert_eq!(response.status(), StatusCode::OK);
3169 + let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");
3170 + engine_task.await.expect("mock engine task");
3171 +
3172 + assert_eq!(
3173 + serde_json::from_slice::<serde_json::Value>(&body).expect("decode json"),
3174 + json!({
3175 + "results": [{
3176 + "args": ["arg1", "arg2"],
3177 + "kwargs": {
3178 + "key1": "value1",
3179 + "key2": "value2"
3180 + },
3181 + "total_items": 4
3182 + }]
3183 + })
3184 + );
3185 +}
3186 +
3187 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
3188 +#[serial]
3189 +async fn sleep_route_uses_python_compatible_default_query_values() {
3190 + let (app, engine_task) = test_admin_app_with_engine_script(|dealer, push| {
3191 + boxed_test_future(async move {
3192 + let utility = recv_engine_message(dealer).await;
3193 + assert_eq!(utility[0].as_ref(), &[0x03]);
3194 +
3195 + let payload = decode_value(&utility[1]).expect("decode utility payload");
3196 + let array = payload.as_array().expect("utility payload array");
3197 + let call_id = array[1].as_i64().expect("call id");
3198 +
3199 + assert_eq!(array[2], Value::from("sleep"));
3200 + assert_eq!(
3201 + array[3],
3202 + Value::Array(vec![Value::from(1_u64), Value::from("abort")])
3203 + );
3204 +
3205 + send_outputs(push, utility_outputs(call_id, utility_none_result())).await;
3206 + })
3207 + })
3208 + .await;
3209 +
3210 + let response = app
3211 + .clone()
3212 + .call(
3213 + Request::builder()
3214 + .method("POST")
3215 + .uri("/sleep")
3216 + .body(Body::empty())
3217 + .expect("build request"),
3218 + )
3219 + .await
3220 + .expect("call app");
3221 +
3222 + let status = response.status();
3223 + let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");
3224 + assert_eq!(status, StatusCode::OK, "{}", String::from_utf8_lossy(&body));
3225 + assert!(body.is_empty());
3226 + engine_task.await.expect("mock engine task");
3227 +}
3228 +
3229 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
3230 +#[serial]
3231 +async fn wake_up_route_without_tags_sends_none() {
3232 + let (app, engine_task) = test_admin_app_with_engine_script(|dealer, push| {
3233 + boxed_test_future(async move {
3234 + let utility = recv_engine_message(dealer).await;
3235 + assert_eq!(utility[0].as_ref(), &[0x03]);
3236 +
3237 + let payload = decode_value(&utility[1]).expect("decode utility payload");
3238 + let array = payload.as_array().expect("utility payload array");
3239 + let call_id = array[1].as_i64().expect("call id");
3240 +
3241 + assert_eq!(array[2], Value::from("wake_up"));
3242 + assert_eq!(array[3], Value::Array(vec![Value::Nil]));
3243 +
3244 + send_outputs(push, utility_outputs(call_id, utility_none_result())).await;
3245 + })
3246 + })
3247 + .await;
3248 +
3249 + let response = app
3250 + .clone()
3251 + .call(
3252 + Request::builder()
3253 + .method("POST")
3254 + .uri("/wake_up")
3255 + .body(Body::empty())
3256 + .expect("build request"),
3257 + )
3258 + .await
3259 + .expect("call app");
3260 +
3261 + let status = response.status();
3262 + let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");
3263 + assert_eq!(status, StatusCode::OK, "{}", String::from_utf8_lossy(&body));
3264 + assert!(body.is_empty());
3265 + engine_task.await.expect("mock engine task");
3266 +}
3267 +
3268 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
3269 +#[serial]
3270 +async fn is_sleeping_route_returns_json_payload() {
3271 + let (app, engine_task) = test_admin_app_with_engine_script(|dealer, push| {
3272 + boxed_test_future(async move {
3273 + let utility = recv_engine_message(dealer).await;
3274 + assert_eq!(utility[0].as_ref(), &[0x03]);
3275 +
3276 + let payload = decode_value(&utility[1]).expect("decode utility payload");
3277 + let array = payload.as_array().expect("utility payload array");
3278 + let call_id = array[1].as_i64().expect("call id");
3279 +
3280 + assert_eq!(array[2], Value::from("is_sleeping"));
3281 + assert_eq!(array[3], Value::Array(Vec::new()));
3282 +
3283 + send_outputs(push, utility_outputs(call_id, utility_result_value(true))).await;
3284 + })
3285 + })
3286 + .await;
3287 +
3288 + let response = app
3289 + .clone()
3290 + .call(
3291 + Request::builder()
3292 + .method("GET")
3293 + .uri("/is_sleeping")
3294 + .body(Body::empty())
3295 + .expect("build request"),
3296 + )
3297 + .await
3298 + .expect("call app");
3299 +
3300 + assert_eq!(response.status(), StatusCode::OK);
3301 + let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");
3302 + engine_task.await.expect("mock engine task");
3303 +
3304 + assert_eq!(
3305 + serde_json::from_slice::<serde_json::Value>(&body).expect("decode json"),
3306 + json!({ "is_sleeping": true })
3307 + );
3308 +}
3309 +
3310 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
3311 +#[serial]
3312 +async fn admin_routes_are_hidden_when_dev_mode_is_disabled() {
3313 + let (chat, engine_task) = test_chat_with_engine_handle().await;
3314 + let app = build_router_with_dev_mode(
3315 + Arc::new(AppState::new(
3316 + vec!["Qwen/Qwen1.5-0.5B-Chat".to_string()],
3317 + chat,
3318 + )),
3319 + false,
3320 + );
3321 +
3322 + for (method, uri) in [
3323 + ("GET", "/is_sleeping"),
3324 + ("POST", "/sleep"),
3325 + ("POST", "/wake_up"),
3326 + ("POST", "/collective_rpc"),
3327 + ("POST", "/reset_prefix_cache"),
3328 + ("POST", "/reset_mm_cache"),
3329 + ("POST", "/reset_encoder_cache"),
3330 + ] {
3331 + let response = app
3332 + .clone()
3333 + .call(
3334 + Request::builder()
3335 + .method(method)
3336 + .uri(uri)
3337 + .body(Body::empty())
3338 + .expect("build request"),
3339 + )
3340 + .await
3341 + .expect("call app");
3342 +
3343 + assert_eq!(response.status(), StatusCode::NOT_FOUND, "{method} {uri}");
3344 + }
3345 +
3346 + engine_task.abort_and_join().await;
3347 +}
3348 +
3349 +// ========================= Stop string tests =========================
3350 +
3351 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
3352 +#[serial]
3353 +async fn non_stream_completions_stop_string_excluded_from_output() {
3354 + // Engine generates "say world" but stop string "wor" truncates output to "say
3355 + // ".
3356 + let output_specs = vec![
3357 + (bytes_to_token_ids(b"say"), None),
3358 + (
3359 + bytes_to_token_ids(b" world"),
3360 + Some(EngineCoreFinishReason::Length),
3361 + ),
3362 + ];
3363 + let (app, engine_task) = test_app_with_stream_output_specs(output_specs).await;
3364 +
3365 + let response = app
3366 + .clone()
3367 + .oneshot(
3368 + Request::builder()
3369 + .method("POST")
3370 + .uri("/v1/completions")
3371 + .header("content-type", "application/json")
3372 + .body(Body::from(
3373 + json!({
3374 + "model": "Qwen/Qwen1.5-0.5B-Chat",
3375 + "prompt": "hello",
3376 + "stream": false,
3377 + "stop": ["wor"]
3378 + })
3379 + .to_string(),
3380 + ))
3381 + .expect("build request"),
3382 + )
3383 + .await
3384 + .expect("call app");
3385 +
3386 + assert_eq!(response.status(), StatusCode::OK);
3387 +
3388 + let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");
3389 + engine_task.await.expect("mock engine task");
3390 + let json: serde_json::Value = serde_json::from_slice(&body).expect("decode json");
3391 +
3392 + assert_eq!(json["choices"][0]["text"], "say ");
3393 + assert_eq!(json["choices"][0]["finish_reason"], "stop");
3394 + assert_eq!(json["choices"][0]["stop_reason"], "wor");
3395 +}
3396 +
3397 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
3398 +#[serial]
3399 +async fn non_stream_completions_stop_string_included_in_output() {
3400 + // Same tokens but include_stop_str_in_output=true includes the stop string in
3401 + // the output.
3402 + let output_specs = vec![
3403 + (bytes_to_token_ids(b"say"), None),
3404 + (
3405 + bytes_to_token_ids(b" world"),
3406 + Some(EngineCoreFinishReason::Length),
3407 + ),
3408 + ];
3409 + let (app, engine_task) = test_app_with_stream_output_specs(output_specs).await;
3410 +
3411 + let response = app
3412 + .clone()
3413 + .oneshot(
3414 + Request::builder()
3415 + .method("POST")
3416 + .uri("/v1/completions")
3417 + .header("content-type", "application/json")
3418 + .body(Body::from(
3419 + json!({
3420 + "model": "Qwen/Qwen1.5-0.5B-Chat",
3421 + "prompt": "hello",
3422 + "stream": false,
3423 + "stop": ["wor"],
3424 + "include_stop_str_in_output": true
3425 + })
3426 + .to_string(),
3427 + ))
3428 + .expect("build request"),
3429 + )
3430 + .await
3431 + .expect("call app");
3432 +
3433 + assert_eq!(response.status(), StatusCode::OK);
3434 +
3435 + let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");
3436 + engine_task.await.expect("mock engine task");
3437 + let json: serde_json::Value = serde_json::from_slice(&body).expect("decode json");
3438 +
3439 + assert_eq!(json["choices"][0]["text"], "say wor");
3440 + assert_eq!(json["choices"][0]["finish_reason"], "stop");
3441 + assert_eq!(json["choices"][0]["stop_reason"], "wor");
3442 +}
3443 +
3444 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
3445 +#[serial]
3446 +async fn stream_completions_stop_string_excluded_from_output() {
3447 + let output_specs = vec![
3448 + (bytes_to_token_ids(b"say"), None),
3449 + (
3450 + bytes_to_token_ids(b" world"),
3451 + Some(EngineCoreFinishReason::Length),
3452 + ),
3453 + ];
3454 + let (app, engine_task) = test_app_with_stream_output_specs(output_specs).await;
3455 +
3456 + let response = app
3457 + .clone()
3458 + .oneshot(
3459 + Request::builder()
3460 + .method("POST")
3461 + .uri("/v1/completions")
3462 + .header("content-type", "application/json")
3463 + .body(Body::from(
3464 + json!({
3465 + "model": "Qwen/Qwen1.5-0.5B-Chat",
3466 + "prompt": "hello",
3467 + "stream": true,
3468 + "stop": ["wor"]
3469 + })
3470 + .to_string(),
3471 + ))
3472 + .expect("build request"),
3473 + )
3474 + .await
3475 + .expect("call app");
3476 +
3477 + assert_eq!(response.status(), StatusCode::OK);
3478 +
3479 + let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");
3480 + engine_task.await.expect("mock engine task");
3481 + let text = String::from_utf8(body.to_vec()).expect("utf8 body");
3482 + let payloads = sse_data_payloads(&text);
3483 +
3484 + // Collect all text deltas from the SSE chunks.
3485 + let mut full_text = String::new();
3486 + for payload in &payloads {
3487 + if *payload == "[DONE]" {
3488 + continue;
3489 + }
3490 + let chunk: serde_json::Value = serde_json::from_str(payload).expect("json chunk");
3491 + if let Some(text) = chunk["choices"][0]["text"].as_str() {
3492 + full_text.push_str(text);
3493 + }
3494 + }
3495 +
3496 + // The concatenated text deltas should equal "say " (stop string excluded).
3497 + assert_eq!(full_text, "say ", "full streamed text: {text}");
3498 +
3499 + // The final chunk should have finish_reason "stop".
3500 + assert!(
3501 + payloads.iter().any(|p| p.contains("\"finish_reason\":\"stop\"")),
3502 + "{text}"
3503 + );
3504 +}
3505 +
3506 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
3507 +#[serial]
3508 +async fn stream_completions_stop_string_included_in_output() {
3509 + let output_specs = vec![
3510 + (bytes_to_token_ids(b"say"), None),
3511 + (
3512 + bytes_to_token_ids(b" world"),
3513 + Some(EngineCoreFinishReason::Length),
3514 + ),
3515 + ];
3516 + let (app, engine_task) = test_app_with_stream_output_specs(output_specs).await;
3517 +
3518 + let response = app
3519 + .clone()
3520 + .oneshot(
3521 + Request::builder()
3522 + .method("POST")
3523 + .uri("/v1/completions")
3524 + .header("content-type", "application/json")
3525 + .body(Body::from(
3526 + json!({
3527 + "model": "Qwen/Qwen1.5-0.5B-Chat",
3528 + "prompt": "hello",
3529 + "stream": true,
3530 + "stop": ["wor"],
3531 + "include_stop_str_in_output": true
3532 + })
3533 + .to_string(),
3534 + ))
3535 + .expect("build request"),
3536 + )
3537 + .await
3538 + .expect("call app");
3539 +
3540 + assert_eq!(response.status(), StatusCode::OK);
3541 +
3542 + let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");
3543 + engine_task.await.expect("mock engine task");
3544 + let text = String::from_utf8(body.to_vec()).expect("utf8 body");
3545 + let payloads = sse_data_payloads(&text);
3546 +
3547 + let mut full_text = String::new();
3548 + for payload in &payloads {
3549 + if *payload == "[DONE]" {
3550 + continue;
3551 + }
3552 + let chunk: serde_json::Value = serde_json::from_str(payload).expect("json chunk");
3553 + if let Some(text) = chunk["choices"][0]["text"].as_str() {
3554 + full_text.push_str(text);
3555 + }
3556 + }
3557 +
3558 + // With include_stop_str_in_output, the stop string "wor" should be included.
3559 + assert_eq!(full_text, "say wor", "full streamed text: {text}");
3560 +
3561 + assert!(
3562 + payloads.iter().any(|p| p.contains("\"finish_reason\":\"stop\"")),
3563 + "{text}"
3564 + );
3565 +}
3566 +
3567 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
3568 +#[serial]
3569 +async fn non_stream_completions_no_stop_string_match_preserves_original_finish_reason() {
3570 + // Stop string "xyz" does not appear in "hi!" so the original finish reason is
3571 + // preserved.
3572 + let (app, engine_task) = test_app_with_engine_handle().await;
3573 +
3574 + let response = app
3575 + .clone()
3576 + .oneshot(
3577 + Request::builder()
3578 + .method("POST")
3579 + .uri("/v1/completions")
3580 + .header("content-type", "application/json")
3581 + .body(Body::from(
3582 + json!({
3583 + "model": "Qwen/Qwen1.5-0.5B-Chat",
3584 + "prompt": "hello",
3585 + "stream": false,
3586 + "stop": ["xyz"]
3587 + })
3588 + .to_string(),
3589 + ))
3590 + .expect("build request"),
3591 + )
3592 + .await
3593 + .expect("call app");
3594 +
3595 + assert_eq!(response.status(), StatusCode::OK);
3596 +
3597 + let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");
3598 + engine_task.await.expect("mock engine task");
3599 + let json: serde_json::Value = serde_json::from_slice(&body).expect("decode json");
3600 +
3601 + // Default output is "hi" (stop token '!' suppressed), finish_reason remains
3602 + // "stop" from EOS.
3603 + assert_eq!(json["choices"][0]["text"], "hi");
3604 + assert_eq!(json["choices"][0]["finish_reason"], "stop");
3605 + // No text stop string matched — stop_reason should be absent.
3606 + assert!(json["choices"][0]["stop_reason"].is_null());
3607 +}
3608 +
3609 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
3610 +#[serial]
3611 +async fn non_stream_completions_stop_string_array_matches_first_occurrence() {
3612 + // Multiple stop strings: "rl" appears in "world" but " wo" appears earlier.
3613 + let output_specs = vec![(
3614 + bytes_to_token_ids(b"say world"),
3615 + Some(EngineCoreFinishReason::Length),
3616 + )];
3617 + let (app, engine_task) = test_app_with_stream_output_specs(output_specs).await;
3618 +
3619 + let response = app
3620 + .clone()
3621 + .oneshot(
3622 + Request::builder()
3623 + .method("POST")
3624 + .uri("/v1/completions")
3625 + .header("content-type", "application/json")
3626 + .body(Body::from(
3627 + json!({
3628 + "model": "Qwen/Qwen1.5-0.5B-Chat",
3629 + "prompt": "hello",
3630 + "stream": false,
3631 + "stop": [" wo", "rl"]
3632 + })
3633 + .to_string(),
3634 + ))
3635 + .expect("build request"),
3636 + )
3637 + .await
3638 + .expect("call app");
3639 +
3640 + assert_eq!(response.status(), StatusCode::OK);
3641 +
3642 + let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");
3643 + engine_task.await.expect("mock engine task");
3644 + let json: serde_json::Value = serde_json::from_slice(&body).expect("decode json");
3645 +
3646 + // " wo" is detected first (at byte 3), so output is truncated to "say".
3647 + assert_eq!(json["choices"][0]["text"], "say");
3648 + assert_eq!(json["choices"][0]["finish_reason"], "stop");
3649 + assert_eq!(json["choices"][0]["stop_reason"], " wo");
3650 +}
3651 +
3652 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
3653 +#[serial]
3654 +async fn completions_empty_stop_string_returns_validation_error() {
3655 + let (app, _engine_task) = test_app_with_engine_handle().await;
3656 +
3657 + let response = app
3658 + .clone()
3659 + .oneshot(
3660 + Request::builder()
3661 + .method("POST")
3662 + .uri("/v1/completions")
3663 + .header("content-type", "application/json")
3664 + .body(Body::from(
3665 + json!({
3666 + "model": "Qwen/Qwen1.5-0.5B-Chat",
3667 + "prompt": "hello",
3668 + "stream": false,
3669 + "stop": [""]
3670 + })
3671 + .to_string(),
3672 + ))
3673 + .expect("build request"),
3674 + )
3675 + .await
3676 + .expect("call app");
3677 +
3678 + assert_eq!(response.status(), StatusCode::BAD_REQUEST);
3679 +}
rust/src/server/src/state.rs new
+120
@@ -0,0 +1,120 @@
1 +use std::sync::Arc;
2 +use std::sync::atomic::{AtomicU64, Ordering};
3 +
4 +use tokio::time::{Duration, Instant, sleep_until};
5 +use tracing::warn;
6 +use vllm_chat::ChatLlm;
7 +use vllm_engine_core_client::EngineCoreClient;
8 +
9 +const SHUTDOWN_REFCOUNT_POLL_INTERVAL: Duration = Duration::from_millis(100);
10 +
11 +/// Shared router state for the minimal single-model OpenAI server.
12 +pub struct AppState {
13 + /// All public model IDs served by this frontend. The first entry is the
14 + /// primary ID used in responses; all entries are valid in requests.
15 + served_model_names: Vec<String>,
16 + /// Shared chat facade used by all requests.
17 + pub chat: ChatLlm,
18 + /// Whether to log a summary line for each completed request.
19 + pub enable_log_requests: bool,
20 + /// Number of in-flight inference requests currently owned by this frontend.
21 + server_load: AtomicU64,
22 +}
23 +
24 +impl AppState {
25 + /// Construct one application state instance.
26 + ///
27 + /// `served_model_names` must be non-empty; the first entry is the primary
28 + /// model ID returned in API responses.
29 + ///
30 + /// # Panics
31 + ///
32 + /// Panics if `served_model_names` is empty.
33 + pub fn new(served_model_names: Vec<String>, chat: ChatLlm) -> Self {
34 + assert!(
35 + !served_model_names.is_empty(),
36 + "served_model_names must not be empty"
37 + );
38 + Self {
39 + served_model_names,
40 + chat,
41 + enable_log_requests: false,
42 + server_load: AtomicU64::new(0),
43 + }
44 + }
45 +
46 + /// Enable per-request completion logging.
47 + pub fn with_log_requests(mut self, enabled: bool) -> Self {
48 + self.enable_log_requests = enabled;
49 + self
50 + }
51 +
52 + /// The primary model name echoed back in API responses (the first served
53 + /// name).
54 + pub fn primary_model_name(&self) -> &str {
55 + self.served_model_names.first().map(String::as_str).unwrap_or_default()
56 + }
57 +
58 + /// All model names served by this frontend.
59 + pub fn served_model_names(&self) -> &[String] {
60 + &self.served_model_names
61 + }
62 +
63 + /// Return a reference to the underlying engine core client for utility
64 + /// calls.
65 + pub(crate) fn engine_core_client(&self) -> &EngineCoreClient {
66 + self.chat.engine_core_client()
67 + }
68 +
69 + /// Return the current in-flight inference request count for the `/load`
70 + /// endpoint.
71 + pub fn server_load(&self) -> u64 {
72 + self.server_load.load(Ordering::Relaxed)
73 + }
74 +
75 + /// Increment the in-flight inference request count, called by the load
76 + /// tracking middleware.
77 + pub(crate) fn increment_server_load(&self) {
78 + self.server_load.fetch_add(1, Ordering::Relaxed);
79 + }
80 +
81 + /// Decrement the in-flight inference request count, called by the load
82 + /// tracking middleware.
83 + pub(crate) fn decrement_server_load(&self) {
84 + self.server_load.fetch_sub(1, Ordering::Relaxed);
85 + }
86 +
87 + /// Wait until all request-owned references are dropped, then shut down the
88 + /// engine client.
89 + ///
90 + /// If the deadline elapses while request/connection tasks still hold state
91 + /// references, skip the clean engine-client shutdown and let process
92 + /// teardown reclaim the remaining resources.
93 + pub async fn shutdown(mut self: Arc<Self>, deadline: Instant) -> anyhow::Result<()> {
94 + loop {
95 + match Arc::try_unwrap(self) {
96 + Ok(state) => {
97 + state.chat.shutdown().await?;
98 + return Ok(());
99 + }
100 + Err(state) => self = state,
101 + }
102 + let ref_count = Arc::strong_count(&self);
103 +
104 + let now = Instant::now();
105 + if now >= deadline {
106 + warn!(
107 + ref_count,
108 + "shutdown deadline elapsed before app state became idle; skipping engine-client shutdown"
109 + );
110 + return Ok(());
111 + }
112 +
113 + sleep_until(std::cmp::min(
114 + deadline,
115 + now + SHUTDOWN_REFCOUNT_POLL_INTERVAL,
116 + ))
117 + .await;
118 + }
119 + }
120 +}
rust/src/server/src/utils.rs new
+106
@@ -0,0 +1,106 @@
1 +use std::collections::HashMap;
2 +use std::time::{SystemTime, UNIX_EPOCH};
3 +
4 +use axum::http::HeaderMap;
5 +use serde_json::Value;
6 +use thiserror_ext::AsReport;
7 +use uuid::Uuid;
8 +
9 +use crate::error::ApiError;
10 +
11 +#[derive(Debug, Clone, PartialEq, Eq, Default)]
12 +pub struct ResolvedRequestContext {
13 + pub request_id: String,
14 + pub data_parallel_rank: Option<u32>,
15 +}
16 +
17 +/// Return the current Unix timestamp in seconds for OpenAI response objects.
18 +pub fn unix_timestamp() -> u64 {
19 + SystemTime::now()
20 + .duration_since(UNIX_EPOCH)
21 + .map(|duration| duration.as_secs())
22 + .unwrap_or_default()
23 +}
24 +
25 +/// Construct an API error for a failed utility call to the engine core.
26 +pub fn utility_call_error(method: &str, error: impl AsReport) -> ApiError {
27 + ApiError::server_error(format!("failed to call {method}: {}", error.as_report()))
28 +}
29 +
30 +/// Merge `kv_transfer_params` into the `vllm_xargs` map, mirroring the Python
31 +/// vLLM behavior where `kv_transfer_params` is injected into `extra_args` for
32 +/// engine-core consumption.
33 +pub fn merge_kv_transfer_params(
34 + mut xargs: Option<HashMap<String, Value>>,
35 + kv_transfer_params: Option<&HashMap<String, Value>>,
36 +) -> Option<HashMap<String, Value>> {
37 + if let Some(kv_params) = kv_transfer_params {
38 + let map = xargs.get_or_insert_with(HashMap::new);
39 + map.insert(
40 + "kv_transfer_params".to_string(),
41 + // This is safe because we know that `kv_params` is already valid JSON.
42 + serde_json::to_value(kv_params).unwrap(),
43 + );
44 + }
45 + xargs
46 +}
47 +
48 +/// Convert OpenAI-style `logit_bias` with string token-ID keys into the
49 +/// internal `HashMap<u32, f32>` representation, validating that every key
50 +/// parses as a `u32`.
51 +pub fn convert_logit_bias(
52 + logit_bias: Option<HashMap<String, f32>>,
53 +) -> Result<Option<HashMap<u32, f32>>, ApiError> {
54 + logit_bias
55 + .map(|bias| {
56 + bias.into_iter()
57 + .map(|(key, value)| {
58 + key.parse().map(|k| (k, value)).map_err(|_| {
59 + ApiError::invalid_request(
60 + format!(
61 + "Invalid key in 'logit_bias': '{key}' is not a valid token ID. \
62 + Token IDs must be non-negative integers."
63 + ),
64 + Some("logit_bias"),
65 + )
66 + })
67 + })
68 + .collect()
69 + })
70 + .transpose()
71 +}
72 +
73 +/// Extract common request metadata from HTTP headers: the external request ID
74 +/// and the optional data-parallel rank used for engine routing.
75 +pub fn resolve_request_context(
76 + headers: &HeaderMap,
77 + request_id: Option<&str>,
78 +) -> ResolvedRequestContext {
79 + // `None` when the header is absent or cannot be parsed as a `u32`.
80 + let data_parallel_rank = headers
81 + .get("X-data-parallel-rank")
82 + .and_then(|v| v.to_str().ok())
83 + .and_then(|s| s.trim().parse().ok());
84 +
85 + // Extract request id from header.
86 + let request_id_header = headers.get("X-Request-Id").and_then(|value| value.to_str().ok());
87 + let request_id = resolve_base_request_id(request_id_header, request_id);
88 +
89 + ResolvedRequestContext {
90 + request_id,
91 + data_parallel_rank,
92 + }
93 +}
94 +
95 +/// Resolve the base external request ID before API-specific prefixes such as
96 +/// `chatcmpl-`.
97 +pub fn resolve_base_request_id(
98 + request_id_header: Option<&str>,
99 + request_id: Option<&str>,
100 +) -> String {
101 + request_id_header.or(request_id).map(ToOwned::to_owned).unwrap_or_else(|| {
102 + let mut id = Uuid::new_v4().simple().to_string();
103 + id.truncate(8);
104 + id
105 + })
106 +}
rust/src/text/Cargo.toml new
+34
@@ -0,0 +1,34 @@
1 +[package]
2 +name = "vllm-text"
3 +version.workspace = true
4 +edition.workspace = true
5 +license.workspace = true
6 +
7 +[dependencies]
8 +anyhow.workspace = true
9 +asynk-strim-attr.workspace = true
10 +easy-ext.workspace = true
11 +enum-as-inner.workspace = true
12 +futures.workspace = true
13 +hf-hub.workspace = true
14 +itertools.workspace = true
15 +serde.workspace = true
16 +serde_json.workspace = true
17 +serde_with.workspace = true
18 +thiserror.workspace = true
19 +thiserror-ext.workspace = true
20 +tracing.workspace = true
21 +trait-set.workspace = true
22 +vllm-engine-core-client.workspace = true
23 +vllm-llm.workspace = true
24 +vllm-tokenizer.workspace = true
25 +
26 +[dev-dependencies]
27 +expect-test.workspace = true
28 +futures.workspace = true
29 +tempfile.workspace = true
30 +tokio.workspace = true
31 +vllm-llm = { workspace = true, features = ["test-util"] }
32 +
33 +[lints]
34 +workspace = true
rust/src/text/src/backend/hf/config.rs new
+373
@@ -0,0 +1,373 @@
1 +use std::collections::BTreeSet;
2 +use std::fs;
3 +use std::path::Path;
4 +
5 +use serde::{Deserialize, Serialize};
6 +use thiserror_ext::AsReport as _;
7 +
8 +use crate::error::{Error, Result};
9 +
10 +/// Minimal subset of `tokenizer_config.json` needed by chat/EOS handling.
11 +#[derive(Debug, Default, Deserialize)]
12 +#[serde(default)]
13 +pub struct HfTokenizerConfig {
14 + #[serde(flatten)]
15 + pub special_tokens: HfSpecialTokens,
16 + pub chat_template: Option<String>,
17 + /// The `tokenizer_class` field from HuggingFace tokenizer configs. Some
18 + /// tiktoken-based models (e.g. DeepSeek, Kimi K2) set this to a value
19 + /// containing "Tiktoken" which can be used as a hint for backend
20 + /// selection.
21 + pub tokenizer_class: Option<String>,
22 +}
23 +
24 +/// Hugging Face named special tokens may be serialized as a string or an
25 +/// object carrying the token content.
26 +#[derive(Debug, Clone, Deserialize)]
27 +#[serde(untagged)]
28 +pub enum NamedSpecialToken {
29 + Text(String),
30 + WithContent { content: String },
31 +}
32 +
33 +impl Serialize for NamedSpecialToken {
34 + fn serialize<S>(&self, serializer: S) -> std::result::Result<S::Ok, S::Error>
35 + where
36 + S: serde::Serializer,
37 + {
38 + serializer.serialize_str(self.as_str())
39 + }
40 +}
41 +
42 +impl From<NamedSpecialToken> for String {
43 + fn from(value: NamedSpecialToken) -> Self {
44 + match value {
45 + NamedSpecialToken::Text(string) => string,
46 + NamedSpecialToken::WithContent { content } => content,
47 + }
48 + }
49 +}
50 +
51 +impl NamedSpecialToken {
52 + pub fn as_str(&self) -> &str {
53 + match self {
54 + Self::Text(value) => value,
55 + Self::WithContent { content } => content,
56 + }
57 + }
58 +}
59 +
60 +/// Minimal set of special-token entries needed by chat/EOS handling.
61 +#[serde_with::skip_serializing_none]
62 +#[derive(Debug, Clone, Default, Deserialize, Serialize)]
63 +#[serde(default)]
64 +pub struct HfSpecialTokens {
65 + pub bos_token: Option<NamedSpecialToken>,
66 + pub eos_token: Option<NamedSpecialToken>,
67 + pub unk_token: Option<NamedSpecialToken>,
68 + pub pad_token: Option<NamedSpecialToken>,
69 +}
70 +
71 +impl HfSpecialTokens {
72 + /// Returns true if we don't discover any special tokens in the config.
73 + pub fn is_empty(&self) -> bool {
74 + self.bos_token.is_none()
75 + && self.eos_token.is_none()
76 + && self.unk_token.is_none()
77 + && self.pad_token.is_none()
78 + }
79 +}
80 +
81 +/// Minimal subset of `config.json` (the model's main HF config).
82 +///
83 +/// This intentionally supports only the two layouts we currently care about in
84 +/// the Rust frontend:
85 +/// - pure text models that keep text metadata at the top level
86 +/// - composite models that expose a single nested `text_config`
87 +///
88 +/// We do not support additional entry points such as `decoder`, `generator`, or
89 +/// `text_encoder`.
90 +#[derive(Debug, Default, Deserialize)]
91 +#[serde(default)]
92 +pub struct ModelConfig {
93 + model_type: Option<String>,
94 + max_position_embeddings: Option<u32>,
95 + num_attention_heads: Option<u32>,
96 + num_experts: Option<OneOrManyExpertCount>,
97 + moe_num_experts: Option<OneOrManyExpertCount>,
98 + n_routed_experts: Option<OneOrManyExpertCount>,
99 + num_local_experts: Option<OneOrManyExpertCount>,
100 + block_configs: Vec<BlockConfig>,
101 + text_config: Option<Box<ModelConfig>>,
102 +}
103 +
104 +/// Minimal subset of `generation_config.json`.
105 +#[derive(Debug, Default, Deserialize)]
106 +#[serde(default)]
107 +pub(super) struct GenerationConfig {
108 + pub eos_token_id: Option<OneOrManyTokenIds>,
109 + pub temperature: Option<f32>,
110 + pub top_p: Option<f32>,
111 + pub top_k: Option<u32>,
112 + pub min_p: Option<f32>,
113 + pub repetition_penalty: Option<f32>,
114 + pub max_new_tokens: Option<u32>,
115 +}
116 +
117 +/// HF generation configs allow either one EOS id or a list of EOS ids.
118 +#[derive(Debug, Clone, Deserialize)]
119 +#[serde(untagged)]
120 +pub(super) enum OneOrManyTokenIds {
121 + One(u32),
122 + Many(Vec<u32>),
123 +}
124 +
125 +impl OneOrManyTokenIds {
126 + pub(super) fn into_set(self) -> BTreeSet<u32> {
127 + match self {
128 + Self::One(id) => BTreeSet::from([id]),
129 + Self::Many(ids) => ids.into_iter().collect(),
130 + }
131 + }
132 +}
133 +
134 +/// Hugging Face configs may expose the expert count either as one integer or
135 +/// as a list of repeated integers.
136 +#[derive(Debug, Clone, Deserialize)]
137 +#[serde(untagged)]
138 +pub(super) enum OneOrManyExpertCount {
139 + One(u32),
140 + Many(Vec<u32>),
141 +}
142 +
143 +impl OneOrManyExpertCount {
144 + fn first_value(&self) -> u32 {
145 + match self {
146 + Self::One(value) => *value,
147 + // Python currently takes the first value for list[int] expert
148 + // counts in remote-code configs.
149 + Self::Many(values) => values.first().copied().unwrap_or(0),
150 + }
151 + }
152 +}
153 +
154 +/// Heterogeneous block-level MoE metadata used as a fallback when no top-level
155 +/// expert-count field is available.
156 +#[derive(Debug, Default, Deserialize)]
157 +#[serde(default)]
158 +pub(super) struct BlockConfig {
159 + pub block_type: String,
160 + pub n_routed_experts: u32,
161 +}
162 +
163 +impl ModelConfig {
164 + /// Return the config that the Rust frontend treats as the text/LLM config.
165 + ///
166 + /// This is deliberately narrower than Python/transformers: we only support
167 + /// either the top-level config itself or a single nested `text_config`.
168 + fn effective_text_config(&self) -> &Self {
169 + self.text_config.as_deref().unwrap_or(self)
170 + }
171 +
172 + /// Return the effective Hugging Face `model_type` used by the Rust
173 + /// frontend.
174 + ///
175 + /// This follows the same simplified text-config selection as the rest of
176 + /// this type: the top-level config wins, otherwise a single nested
177 + /// `text_config` may provide the value.
178 + pub fn model_type(&self) -> Option<&str> {
179 + self.model_type.as_deref().or_else(|| self.text_config.as_deref()?.model_type())
180 + }
181 +
182 + /// Reject partially nested `text_config` payloads that are unlikely to be
183 + /// valid LLM configs for our current use.
184 + ///
185 + /// This keeps the simplified Rust-side parsing honest: if a model declares
186 + /// `text_config`, it must at least look like a real text model config.
187 + fn validate_text_config_selection(&self) -> Result<()> {
188 + if let Some(text_config) = self.text_config.as_deref()
189 + && text_config.num_attention_heads.is_none()
190 + {
191 + return Err(Error::Tokenizer(
192 + "the text config extracted from the model config does not have `num_attention_heads`"
193 + .to_string(),
194 + ));
195 + }
196 +
197 + Ok(())
198 + }
199 +
200 + /// Match Python's current expert-count priority on the selected text
201 + /// config.
202 + ///
203 + /// The only intentional simplification here is how we pick the text config:
204 + /// Rust only looks at the top level or `text_config`, not the broader
205 + /// transformers composite-config surface.
206 + fn num_experts_from_block_configs(&self) -> u32 {
207 + self.effective_text_config()
208 + .block_configs
209 + .iter()
210 + .filter(|block| block.block_type == "moe")
211 + .map(|block| block.n_routed_experts)
212 + .max()
213 + .unwrap_or(0)
214 + }
215 +
216 + pub(super) fn num_experts(&self) -> u32 {
217 + let config = self.effective_text_config();
218 + let direct = [
219 + config.num_experts.as_ref(),
220 + config.moe_num_experts.as_ref(),
221 + config.n_routed_experts.as_ref(),
222 + config.num_local_experts.as_ref(),
223 + ]
224 + .into_iter()
225 + .flatten()
226 + .map(OneOrManyExpertCount::first_value)
227 + .next()
228 + .unwrap_or(0);
229 +
230 + if direct > 0 {
231 + direct
232 + } else {
233 + self.num_experts_from_block_configs()
234 + }
235 + }
236 +
237 + pub(super) fn is_moe(&self) -> bool {
238 + self.num_experts() > 0
239 + }
240 +
241 + pub(super) fn max_position_embeddings(&self) -> Option<u32> {
242 + self.effective_text_config().max_position_embeddings
243 + }
244 +}
245 +
246 +/// Load the tokenizer-side EOS metadata if a config file is present.
247 +pub fn load_tokenizer_config(path: Option<&Path>) -> Result<HfTokenizerConfig> {
248 + read_json_file(path)
249 +}
250 +
251 +/// Load the generation-side EOS metadata if a config file is present.
252 +pub(super) fn load_generation_config(path: Option<&Path>) -> Result<GenerationConfig> {
253 + read_json_file(path)
254 +}
255 +
256 +/// Load the model-side config (`config.json`) if present.
257 +pub fn load_model_config(path: Option<&Path>) -> Result<ModelConfig> {
258 + let config: ModelConfig = read_json_file(path)?;
259 + config.validate_text_config_selection()?;
260 + Ok(config)
261 +}
262 +
263 +fn read_json_file<T>(path: Option<&Path>) -> Result<T>
264 +where
265 + T: for<'de> Deserialize<'de> + Default,
266 +{
267 + let Some(path) = path else {
268 + return Ok(T::default());
269 + };
270 + let content = fs::read_to_string(path).map_err(|error| {
271 + Error::Tokenizer(format!(
272 + "failed to read {}: {}",
273 + path.display(),
274 + error.as_report()
275 + ))
276 + })?;
277 + serde_json::from_str(&content).map_err(|error| {
278 + Error::Tokenizer(format!(
279 + "failed to parse {}: {}",
280 + path.display(),
281 + error.as_report()
282 + ))
283 + })
284 +}
285 +
286 +#[cfg(test)]
287 +mod tests {
288 + use super::ModelConfig;
289 +
290 + #[test]
291 + fn model_config_detects_moe_from_named_expert_fields() {
292 + let field_names = [
293 + "num_experts",
294 + "moe_num_experts",
295 + "n_routed_experts",
296 + "num_local_experts",
297 + ];
298 +
299 + for field_name in field_names {
300 + let config: ModelConfig =
301 + serde_json::from_str(&format!(r#"{{"{field_name}": 64}}"#)).unwrap();
302 + assert_eq!(config.num_experts(), 64, "field_name={field_name}");
303 + assert!(config.is_moe(), "field_name={field_name}");
304 + }
305 + }
306 +
307 + #[test]
308 + fn model_config_uses_first_value_for_list_expert_counts() {
309 + let config: ModelConfig = serde_json::from_str(r#"{"num_experts":[16,16]}"#).unwrap();
310 +
311 + assert_eq!(config.num_experts(), 16);
312 + assert!(config.is_moe());
313 + }
314 +
315 + #[test]
316 + fn model_config_falls_back_to_block_configs_maximum() {
317 + let config: ModelConfig = serde_json::from_str(
318 + r#"{
319 + "block_configs": [
320 + {"block_type":"attention","n_routed_experts":9},
321 + {"block_type":"moe","n_routed_experts":32},
322 + {"block_type":"moe","n_routed_experts":64}
323 + ]
324 + }"#,
325 + )
326 + .unwrap();
327 +
328 + assert_eq!(config.num_experts(), 64);
329 + assert!(config.is_moe());
330 + }
331 +
332 + #[test]
333 + fn model_config_prefers_nested_text_config_like_python_hf_text_config() {
334 + let config: ModelConfig = serde_json::from_str(
335 + r#"{
336 + "model_type": "top_level",
337 + "num_experts": 64,
338 + "max_position_embeddings": 8192,
339 + "text_config": {
340 + "model_type": "nested",
341 + "num_attention_heads": 32,
342 + "num_local_experts": 8,
343 + "max_position_embeddings": 4096
344 + }
345 + }"#,
346 + )
347 + .unwrap();
348 +
349 + assert_eq!(config.num_experts(), 8);
350 + assert_eq!(config.model_type(), Some("top_level"));
351 + assert_eq!(config.max_position_embeddings(), Some(4096));
352 + assert!(config.is_moe());
353 + }
354 +
355 + #[test]
356 + fn model_config_defaults_to_non_moe_when_no_expert_metadata_exists() {
357 + let config: ModelConfig =
358 + serde_json::from_str(r#"{"max_position_embeddings":4096}"#).unwrap();
359 +
360 + assert_eq!(config.num_experts(), 0);
361 + assert!(!config.is_moe());
362 + assert_eq!(config.max_position_embeddings(), Some(4096));
363 + }
364 +
365 + #[test]
366 + fn model_config_rejects_nested_text_config_without_attention_heads() {
367 + let config: ModelConfig =
368 + serde_json::from_str(r#"{"text_config":{"max_position_embeddings":4096}}"#).unwrap();
369 +
370 + let error = config.validate_text_config_selection().unwrap_err();
371 + assert!(error.to_string().contains("does not have `num_attention_heads`"),);
372 + }
373 +}
rust/src/text/src/backend/hf/mod.rs new
+120
@@ -0,0 +1,120 @@
1 +mod config;
2 +mod model_files;
3 +
4 +use std::collections::BTreeSet;
5 +use std::sync::Arc;
6 +
7 +use tracing::info;
8 +use vllm_tokenizer::{DynTokenizer, HuggingFaceTokenizer, TekkenTokenizer, TiktokenTokenizer};
9 +
10 +use self::config::{GenerationConfig, load_generation_config};
11 +pub use self::config::{
12 + HfSpecialTokens, HfTokenizerConfig, ModelConfig, NamedSpecialToken, load_model_config,
13 + load_tokenizer_config,
14 +};
15 +pub use self::model_files::{ResolvedModelFiles, TokenizerSource};
16 +use crate::backend::{SamplingHints, TextBackend};
17 +use crate::error::Result;
18 +
19 +fn load_tokenizer(tokenizer: &TokenizerSource) -> Result<DynTokenizer> {
20 + match tokenizer {
21 + TokenizerSource::HuggingFace(path) => Ok(Arc::new(HuggingFaceTokenizer::new(path)?)),
22 + TokenizerSource::Tiktoken(path) => Ok(Arc::new(TiktokenTokenizer::new(path)?)),
23 + TokenizerSource::Tekken(path) => Ok(Arc::new(TekkenTokenizer::new(path)?)),
24 + }
25 +}
26 +
27 +/// [`TextBackend`] implementation built on Hugging Face model files.
28 +pub struct HfTextBackend {
29 + model_id: String,
30 + files: ResolvedModelFiles,
31 + tokenizer: DynTokenizer,
32 + /// Primary EOS handled by engine-core's dedicated EOS path.
33 + primary_eos_token_id: Option<u32>,
34 + /// Additional EOS ids that should flow through stop-token handling.
35 + extra_eos_token_ids: BTreeSet<u32>,
36 + /// Generation-config for sampling defaults that may be inherited when the
37 + /// user does not explicitly override them.
38 + generation_config: GenerationConfig,
39 + /// Model config (`config.json`).
40 + model_config: ModelConfig,
41 +}
42 +
43 +impl HfTextBackend {
44 + /// Load the text backend with the given model id.
45 + pub async fn from_model(model_id: &str) -> Result<Self> {
46 + let files = ResolvedModelFiles::new(model_id).await?;
47 + Self::from_resolved_model_files(files, model_id.to_string())
48 + }
49 +
50 + /// Load the text backend from resolved Hugging Face model files.
51 + pub fn from_resolved_model_files(files: ResolvedModelFiles, model_id: String) -> Result<Self> {
52 + let tokenizer_config = load_tokenizer_config(files.tokenizer_config_path.as_deref())?;
53 + let tokenizer = load_tokenizer(&files.tokenizer)?;
54 + let primary_eos_token_id = tokenizer_config
55 + .special_tokens
56 + .eos_token
57 + .as_ref()
58 + .and_then(|token| tokenizer.token_to_id(token.as_str()));
59 +
60 + let model_config = load_model_config(files.config_path.as_deref())?;
61 + let generation_config = load_generation_config(files.generation_config_path.as_deref())?;
62 + let mut extra_eos_token_ids = generation_config
63 + .eos_token_id
64 + .clone()
65 + .map(|value| value.into_set())
66 + .unwrap_or_default();
67 + if let Some(primary_eos_token_id) = primary_eos_token_id {
68 + extra_eos_token_ids.remove(&primary_eos_token_id);
69 + }
70 +
71 + info!(
72 + model_id,
73 + "loaded text backend with Hugging Face model files"
74 + );
75 +
76 + Ok(Self {
77 + model_id,
78 + files,
79 + tokenizer,
80 + primary_eos_token_id,
81 + extra_eos_token_ids,
82 + generation_config,
83 + model_config,
84 + })
85 + }
86 +
87 + /// Expose the resolved model files for use by the chat backend to load the
88 + /// chat template.
89 + pub fn resolved_model_files(&self) -> &ResolvedModelFiles {
90 + &self.files
91 + }
92 +}
93 +
94 +impl TextBackend for HfTextBackend {
95 + fn tokenizer(&self) -> DynTokenizer {
96 + self.tokenizer.clone()
97 + }
98 +
99 + fn is_moe(&self) -> bool {
100 + self.model_config.is_moe()
101 + }
102 +
103 + fn model_id(&self) -> &str {
104 + &self.model_id
105 + }
106 +
107 + fn sampling_hints(&self) -> Result<SamplingHints> {
108 + Ok(SamplingHints {
109 + primary_eos_token_id: self.primary_eos_token_id,
110 + extra_eos_token_ids: self.extra_eos_token_ids.clone(),
111 + default_temperature: self.generation_config.temperature,
112 + default_top_p: self.generation_config.top_p,
113 + default_top_k: self.generation_config.top_k,
114 + default_min_p: self.generation_config.min_p,
115 + default_repetition_penalty: self.generation_config.repetition_penalty,
116 + default_max_tokens: self.generation_config.max_new_tokens,
117 + max_model_len: self.model_config.max_position_embeddings(),
118 + })
119 + }
120 +}
rust/src/text/src/backend/hf/model_files.rs new
+459
@@ -0,0 +1,459 @@
1 +use std::path::{Path, PathBuf};
2 +
3 +use hf_hub::Cache;
4 +use hf_hub::api::tokio::{Api, ApiBuilder, ApiRepo};
5 +use thiserror_ext::AsReport as _;
6 +
7 +use super::config::{HfTokenizerConfig, load_tokenizer_config};
8 +use crate::error::{Error, Result};
9 +
10 +const HF_TOKEN_ENV: &str = "HF_TOKEN";
11 +
12 +/// The tokenizer source selected for a model.
13 +#[derive(Debug, Clone)]
14 +pub enum TokenizerSource {
15 + /// Path to `tokenizer.json` in HuggingFace format.
16 + HuggingFace(PathBuf),
17 + /// Path to `tiktoken.model` or `*.tiktoken` file for tiktoken-based models.
18 + Tiktoken(PathBuf),
19 + /// Path to `tekken.json` when present (Mistral native tokenizer format).
20 + ///
21 + /// When set, the Tekken tokenizer should be preferred over the Hugging Face
22 + /// tokenizer because the HuggingFace `tokenizer.json` for Mistral
23 + /// models has a known regex bug that produces incorrect token IDs for
24 + /// some inputs.
25 + Tekken(PathBuf),
26 +}
27 +
28 +impl TokenizerSource {
29 + pub fn path(&self) -> &Path {
30 + match self {
31 + Self::HuggingFace(path) | Self::Tiktoken(path) | Self::Tekken(path) => path,
32 + }
33 + }
34 +}
35 +
36 +/// Concrete tokenizer/config file locations resolved for one HF model id.
37 +#[derive(Debug, Clone)]
38 +pub struct ResolvedModelFiles {
39 + /// The selected tokenizer source for this model.
40 + pub tokenizer: TokenizerSource,
41 + pub tokenizer_config_path: Option<PathBuf>,
42 + pub generation_config_path: Option<PathBuf>,
43 + pub preprocessor_config_path: Option<PathBuf>,
44 + pub chat_template_path: Option<PathBuf>,
45 + pub config_path: Option<PathBuf>,
46 +}
47 +
48 +impl ResolvedModelFiles {
49 + /// Resolve tokenizer/config files from a local model directory first when
50 + /// `model_id` points to one, otherwise consult the local HF cache and
51 + /// finally the Hub.
52 + pub async fn new(model_id: &str) -> Result<Self> {
53 + if Path::new(model_id).is_dir() {
54 + return resolve_local_model_files(Path::new(model_id));
55 + }
56 + if let Some(files) = resolve_cached_model_files(model_id)? {
57 + return Ok(files);
58 + }
59 + resolve_remote_model_files(model_id).await
60 + }
61 +}
62 +
63 +fn resolve_local_model_files(model_dir: &Path) -> Result<ResolvedModelFiles> {
64 + let tokenizer_config_path = local_file_if_exists(model_dir, "tokenizer_config.json");
65 + let tokenizer_config = load_tokenizer_config(tokenizer_config_path.as_deref())?;
66 + let tokenizer = resolve_local_tokenizer_source(model_dir, &tokenizer_config)?;
67 +
68 + Ok(ResolvedModelFiles {
69 + tokenizer,
70 + tokenizer_config_path,
71 + generation_config_path: local_file_if_exists(model_dir, "generation_config.json"),
72 + preprocessor_config_path: local_file_if_exists(model_dir, "preprocessor_config.json"),
73 + chat_template_path: discover_chat_template_in_dir(model_dir),
74 + config_path: local_file_if_exists(model_dir, "config.json"),
75 + })
76 +}
77 +
78 +async fn resolve_remote_model_files(model_id: &str) -> Result<ResolvedModelFiles> {
79 + let api = build_api().map_err(|error| Error::Tokenizer(error.to_report_string()))?;
80 + let repo = api.model(model_id.to_string());
81 + let info = repo.info().await.map_err(|error| {
82 + Error::Tokenizer(format!(
83 + "failed to fetch model '{model_id}': {}",
84 + error.as_report()
85 + ))
86 + })?;
87 +
88 + let siblings = info
89 + .siblings
90 + .iter()
91 + .map(|sibling| sibling.rfilename.as_str())
92 + .collect::<std::collections::BTreeSet<_>>();
93 +
94 + let tokenizer_config_path =
95 + download_if_present(&repo, model_id, &siblings, "tokenizer_config.json").await?;
96 + let tokenizer_config = load_tokenizer_config(tokenizer_config_path.as_deref())?;
97 +
98 + let tokenizer = resolve_remote_tokenizer_source(
99 + &repo,
100 + model_id,
101 + &siblings,
102 + tokenizer_config.tokenizer_class.as_deref(),
103 + )
104 + .await?;
105 +
106 + let generation_config_path =
107 + download_if_present(&repo, model_id, &siblings, "generation_config.json").await?;
108 + let preprocessor_config_path =
109 + download_if_present(&repo, model_id, &siblings, "preprocessor_config.json").await?;
110 + let chat_template_name = siblings
111 + .contains("chat_template.json")
112 + .then_some("chat_template.json")
113 + .or_else(|| siblings.contains("chat_template.jinja").then_some("chat_template.jinja"))
114 + .or_else(|| siblings.iter().copied().find(|name| name.ends_with(".jinja")));
115 + let chat_template_path = match chat_template_name {
116 + Some(name) => Some(download_known_file(&repo, model_id, name).await?),
117 + None => None,
118 + };
119 + let config_path = download_if_present(&repo, model_id, &siblings, "config.json").await?;
120 +
121 + Ok(ResolvedModelFiles {
122 + tokenizer,
123 + tokenizer_config_path,
124 + generation_config_path,
125 + preprocessor_config_path,
126 + chat_template_path,
127 + config_path,
128 + })
129 +}
130 +
131 +fn resolve_cached_model_files(model_id: &str) -> Result<Option<ResolvedModelFiles>> {
132 + let cache_repo = Cache::from_env().model(model_id.to_string());
133 +
134 + let tokenizer_config_path = cache_repo.get("tokenizer_config.json");
135 + let tokenizer_config = load_tokenizer_config(tokenizer_config_path.as_deref())?;
136 + let tokenizer = match resolve_cached_tokenizer_source(&cache_repo, &tokenizer_config)? {
137 + Some(tokenizer) => tokenizer,
138 + None => return Ok(None),
139 + };
140 +
141 + let model_dir = tokenizer.path().parent().ok_or_else(|| {
142 + Error::Tokenizer("resolved tokenizer file has no parent directory".to_string())
143 + })?;
144 + let generation_config_path = cache_repo.get("generation_config.json");
145 + let preprocessor_config_path = cache_repo.get("preprocessor_config.json");
146 + let chat_template_path = discover_chat_template_in_dir(model_dir);
147 + let config_path = cache_repo.get("config.json");
148 +
149 + Ok(Some(ResolvedModelFiles {
150 + tokenizer,
151 + tokenizer_config_path,
152 + generation_config_path,
153 + preprocessor_config_path,
154 + chat_template_path,
155 + config_path,
156 + }))
157 +}
158 +
159 +async fn resolve_remote_tokenizer_source(
160 + repo: &ApiRepo,
161 + model_id: &str,
162 + siblings: &std::collections::BTreeSet<&str>,
163 + tokenizer_class: Option<&str>,
164 +) -> Result<TokenizerSource> {
165 + if let Some(tekken_path) = download_if_present(repo, model_id, siblings, "tekken.json").await? {
166 + return Ok(TokenizerSource::Tekken(tekken_path));
167 + }
168 +
169 + let tokenizer_path = if siblings.contains("tokenizer.json") {
170 + download_known_file(repo, model_id, "tokenizer.json").await?
171 + } else if let Some(tiktoken_name) = find_tiktoken_sibling(siblings) {
172 + download_known_file(repo, model_id, tiktoken_name).await?
173 + } else {
174 + return Err(Error::Tokenizer(format!(
175 + "model '{model_id}' does not expose a supported tokenizer file \
176 + (tokenizer.json, tiktoken.model, or *.tiktoken) on Hugging Face"
177 + )));
178 + };
179 +
180 + Ok(resolve_tokenizer_source(
181 + tokenizer_path,
182 + tokenizer_class,
183 + None,
184 + ))
185 +}
186 +
187 +fn resolve_cached_tokenizer_source(
188 + cache_repo: &hf_hub::CacheRepo,
189 + tokenizer_config: &HfTokenizerConfig,
190 +) -> Result<Option<TokenizerSource>> {
191 + let tekken_path = cache_repo.get("tekken.json");
192 +
193 + if let Some(tekken_path) = tekken_path {
194 + return Ok(Some(TokenizerSource::Tekken(tekken_path)));
195 + }
196 +
197 + let Some(tokenizer_path) = cache_repo.get("tokenizer.json").or_else(|| {
198 + // tiktoken.model is the most common name, try it first.
199 + cache_repo.get("tiktoken.model").or_else(|| {
200 + // Scan for any *.tiktoken file in the cache snapshot directory.
201 + let snapshot_dir = cache_repo.get("config.json")?.parent()?.to_path_buf();
202 + discover_tiktoken_in_dir(&snapshot_dir)
203 + })
204 + }) else {
205 + return Ok(None);
206 + };
207 +
208 + Ok(Some(resolve_tokenizer_source(
209 + tokenizer_path,
210 + tokenizer_config.tokenizer_class.as_deref(),
211 + None,
212 + )))
213 +}
214 +
215 +fn resolve_local_tokenizer_source(
216 + model_dir: &Path,
217 + tokenizer_config: &HfTokenizerConfig,
218 +) -> Result<TokenizerSource> {
219 + let tekken_path = local_file_if_exists(model_dir, "tekken.json");
220 + if let Some(tekken_path) = tekken_path {
221 + return Ok(TokenizerSource::Tekken(tekken_path));
222 + }
223 +
224 + let tokenizer_path = local_file_if_exists(model_dir, "tokenizer.json")
225 + .or_else(|| local_file_if_exists(model_dir, "tiktoken.model"))
226 + .or_else(|| discover_tiktoken_in_dir(model_dir))
227 + .ok_or_else(|| {
228 + Error::Tokenizer(format!(
229 + "local model directory '{}' does not contain a supported tokenizer file \
230 + (tokenizer.json, tiktoken.model, or *.tiktoken)",
231 + model_dir.display()
232 + ))
233 + })?;
234 +
235 + Ok(resolve_tokenizer_source(
236 + tokenizer_path,
237 + tokenizer_config.tokenizer_class.as_deref(),
238 + None,
239 + ))
240 +}
241 +
242 +/// Choose the tokenizer.
243 +///
244 +/// Selection order:
245 +/// 1. `tekken.json` — Mistral native tokenizer (preferred over HF `tokenizer.json` because the HF
246 +/// version has a known regex bug for Mistral models).
247 +/// 2. File extension — `.tiktoken` / `tiktoken.model` files use tiktoken from BPE data.
248 +/// 3. `tokenizer_class` in `tokenizer_config.json` — classes containing "Tiktoken" (case-
249 +/// insensitive) trigger tiktoken loading from a sibling BPE file.
250 +/// 4. Default — `tokenizer.json` in HuggingFace format.
251 +fn resolve_tokenizer_source(
252 + tokenizer_path: PathBuf,
253 + tokenizer_class: Option<&str>,
254 + tekken_path: Option<PathBuf>,
255 +) -> TokenizerSource {
256 + if let Some(tekken_path) = tekken_path {
257 + return TokenizerSource::Tekken(tekken_path);
258 + }
259 +
260 + if is_tiktoken_file(&tokenizer_path) {
261 + return TokenizerSource::Tiktoken(tokenizer_path);
262 + }
263 +
264 + if tokenizer_class.is_some_and(|cls| cls.to_ascii_lowercase().contains("tiktoken"))
265 + && let Some(dir) = tokenizer_path.parent()
266 + && let Some(tiktoken_path) = discover_tiktoken_in_dir(dir)
267 + {
268 + return TokenizerSource::Tiktoken(tiktoken_path);
269 + }
270 +
271 + TokenizerSource::HuggingFace(tokenizer_path)
272 +}
273 +
274 +/// Download `filename` only if it exists in `siblings`.
275 +async fn download_if_present(
276 + repo: &ApiRepo,
277 + model_id: &str,
278 + siblings: &std::collections::BTreeSet<&str>,
279 + filename: &str,
280 +) -> Result<Option<PathBuf>> {
281 + match siblings.contains(filename) {
282 + true => download_known_file(repo, model_id, filename).await.map(Some),
283 + false => Ok(None),
284 + }
285 +}
286 +
287 +async fn download_known_file(repo: &ApiRepo, model_id: &str, filename: &str) -> Result<PathBuf> {
288 + repo.get(filename).await.map_err(|error| {
289 + Error::Tokenizer(format!(
290 + "failed to download '{filename}' for model '{model_id}': {}",
291 + error.as_report()
292 + ))
293 + })
294 +}
295 +
296 +fn build_api() -> anyhow::Result<Api> {
297 + let mut builder = ApiBuilder::from_env().with_progress(true);
298 + if let Ok(token) = std::env::var(HF_TOKEN_ENV)
299 + && !token.is_empty()
300 + {
301 + builder = builder.with_token(Some(token));
302 + }
303 + Ok(builder.build()?)
304 +}
305 +
306 +fn local_file_if_exists(dir: &Path, filename: &str) -> Option<PathBuf> {
307 + let path = dir.join(filename);
308 + path.is_file().then_some(path)
309 +}
310 +
311 +/// Find a tiktoken file name among repo siblings, preferring `tiktoken.model`.
312 +fn find_tiktoken_sibling<'a>(siblings: &std::collections::BTreeSet<&'a str>) -> Option<&'a str> {
313 + if siblings.contains("tiktoken.model") {
314 + return Some("tiktoken.model");
315 + }
316 + siblings.iter().copied().find(|name| name.ends_with(".tiktoken"))
317 +}
318 +
319 +/// Discover a tiktoken model file in a local directory.
320 +pub(super) fn discover_tiktoken_in_dir(dir: &std::path::Path) -> Option<PathBuf> {
321 + let tiktoken_model = dir.join("tiktoken.model");
322 + if tiktoken_model.exists() {
323 + return Some(tiktoken_model);
324 + }
325 + std::fs::read_dir(dir).ok()?.flatten().find_map(|entry| {
326 + let path = entry.path();
327 + if path
328 + .file_name()
329 + .and_then(|n| n.to_str())
330 + .is_some_and(|n| n.ends_with(".tiktoken"))
331 + {
332 + Some(path)
333 + } else {
334 + None
335 + }
336 + })
337 +}
338 +
339 +/// Returns `true` if `path` points to a tiktoken-format file (by name).
340 +pub(super) fn is_tiktoken_file(path: &std::path::Path) -> bool {
341 + path.file_name()
342 + .and_then(|n| n.to_str())
343 + .is_some_and(|name| name == "tiktoken.model" || name.ends_with(".tiktoken"))
344 +}
345 +
346 +/// Chat templates are sometimes stored as dedicated .jinja files rather than as
347 +/// a fixed-name config entry, so we scan the cached model dir.
348 +fn discover_chat_template_in_dir(dir: &std::path::Path) -> Option<PathBuf> {
349 + let json_template_path = dir.join("chat_template.json");
350 + if json_template_path.exists() {
351 + return Some(json_template_path);
352 + }
353 +
354 + let jinja_path = dir.join("chat_template.jinja");
355 + if jinja_path.exists() {
356 + return Some(jinja_path);
357 + }
358 +
359 + std::fs::read_dir(dir).ok()?.flatten().map(|entry| entry.path()).find(|path| {
360 + path.file_name()
361 + .and_then(|name| name.to_str())
362 + .is_some_and(|name| name.ends_with(".jinja"))
363 + })
364 +}
365 +
366 +#[cfg(test)]
367 +mod tests {
368 + use std::fs;
369 +
370 + use tempfile::tempdir;
371 + use vllm_tokenizer::{TiktokenTokenizer, Tokenizer};
372 +
373 + use super::{ResolvedModelFiles, TokenizerSource};
374 +
375 + #[tokio::test]
376 + async fn resolved_model_files_prefers_absolute_local_model_dir() {
377 + let dir = tempdir().expect("create temp dir");
378 + fs::write(dir.path().join("tokenizer.json"), "{}").expect("write tokenizer");
379 + fs::write(
380 + dir.path().join("tokenizer_config.json"),
381 + r#"{"tokenizer_class":"PreTrainedTokenizerFast"}"#,
382 + )
383 + .expect("write tokenizer config");
384 + fs::write(dir.path().join("config.json"), "{}").expect("write config");
385 +
386 + let files = ResolvedModelFiles::new(dir.path().to_str().expect("utf8 path"))
387 + .await
388 + .expect("resolve local model files");
389 +
390 + match files.tokenizer {
391 + TokenizerSource::HuggingFace(path) => {
392 + assert_eq!(path, dir.path().join("tokenizer.json"));
393 + }
394 + other => panic!("expected HuggingFace tokenizer, got {other:?}"),
395 + }
396 + assert_eq!(files.config_path, Some(dir.path().join("config.json")));
397 + assert_eq!(
398 + files.tokenizer_config_path,
399 + Some(dir.path().join("tokenizer_config.json"))
400 + );
401 + }
402 +
403 + #[tokio::test]
404 + #[ignore = "requires network access to Hugging Face and downloads the real Kimi K2.5 tokenizer"]
405 + async fn tiktoken_real_kimi_k25_tokenizer_files_load_and_handle_special_tokens() {
406 + let files = ResolvedModelFiles::new("moonshotai/Kimi-K2.5")
407 + .await
408 + .expect("resolve real Kimi K2.5 model files");
409 +
410 + let tokenizer_path = match &files.tokenizer {
411 + TokenizerSource::Tiktoken(path) => path.clone(),
412 + other => panic!("expected tiktoken tokenizer source, got {other:?}"),
413 + };
414 +
415 + for backend in [
416 + TiktokenTokenizer::new_riptoken(&tokenizer_path).expect("load riptoken backend"),
417 + TiktokenTokenizer::new_tiktoken_rs(&tokenizer_path).expect("load tiktoken-rs backend"),
418 + ] {
419 + let think_id = backend.token_to_id("<think>").expect("resolve <think>");
420 + let end_think_id = backend.token_to_id("</think>").expect("resolve </think>");
421 + let tool_section_id = backend
422 + .token_to_id("<|tool_calls_section_begin|>")
423 + .expect("resolve tool call section marker");
424 + let contraction_heavy_text =
425 + "I'm sure it's fine, but I can't say I'd trust that it's what we'd ship.";
426 + let contraction_heavy_ids = backend.encode(contraction_heavy_text, false).unwrap();
427 +
428 + assert_eq!(
429 + (think_id, end_think_id, tool_section_id),
430 + (163606, 163607, 163595)
431 + );
432 + assert_eq!(backend.decode(&[think_id], true).unwrap(), "<think>");
433 + assert_eq!(backend.decode(&[end_think_id], true).unwrap(), "</think>");
434 + assert_eq!(
435 + backend.decode(&[tool_section_id], true).unwrap(),
436 + "<|tool_calls_section_begin|>"
437 + );
438 +
439 + // This demonstrates that we're using Kimi's custom BPE pattern.
440 + // With CL100K this will be 23 tokens instead.
441 + assert_eq!(
442 + contraction_heavy_ids,
443 + vec![
444 + 17172, 3287, 4643, 8201, 11, 996, 374, 8971, 3637, 20020, 8173, 473, 4643,
445 + 1573, 56229, 13922, 13,
446 + ]
447 + );
448 + assert_eq!(contraction_heavy_ids.len(), 17);
449 + assert_eq!(
450 + backend.decode(&contraction_heavy_ids, false).unwrap(),
451 + contraction_heavy_text
452 + );
453 +
454 + // Special-looking text that is not actually registered should fail gracefully.
455 + assert_eq!(backend.token_to_id("◁think▷"), None);
456 + assert_eq!(backend.token_to_id("<|definitely_not_registered|>"), None);
457 + }
458 + }
459 +}
rust/src/text/src/backend/mod.rs new
+47
@@ -0,0 +1,47 @@
1 +pub mod hf;
2 +
3 +use std::sync::Arc;
4 +
5 +use vllm_tokenizer::DynTokenizer;
6 +
7 +use crate::error::Result;
8 +
9 +/// Tokenizer/model-derived hints used to enrich text-generation requests before
10 +/// they are lowered into engine-core.
11 +#[derive(Debug, Clone, Default, PartialEq)]
12 +pub struct SamplingHints {
13 + pub primary_eos_token_id: Option<u32>,
14 + pub extra_eos_token_ids: std::collections::BTreeSet<u32>,
15 + pub default_temperature: Option<f32>,
16 + pub default_top_p: Option<f32>,
17 + pub default_top_k: Option<u32>,
18 + pub default_min_p: Option<f32>,
19 + pub default_repetition_penalty: Option<f32>,
20 + pub default_max_tokens: Option<u32>,
21 + /// Model context window size (`max_position_embeddings` from
22 + /// `config.json`).
23 + pub max_model_len: Option<u32>,
24 +}
25 +
26 +/// Minimal text-processing backend needed by `vllm-text`.
27 +pub trait TextBackend: Send + Sync {
28 + /// Return the tokenizer used by this backend.
29 + fn tokenizer(&self) -> DynTokenizer;
30 +
31 + /// Return whether the loaded model is a mixture-of-experts model.
32 + fn is_moe(&self) -> bool {
33 + false
34 + }
35 +
36 + /// Return the backend model ID.
37 + fn model_id(&self) -> &str;
38 +
39 + /// Return tokenizer/model-derived hints used to enrich southbound sampling
40 + /// parameters.
41 + fn sampling_hints(&self) -> Result<SamplingHints> {
42 + Ok(SamplingHints::default())
43 + }
44 +}
45 +
46 +/// Shared trait-object form of [`TextBackend`].
47 +pub type DynTextBackend = Arc<dyn TextBackend>;
rust/src/text/src/error.rs new
+30
@@ -0,0 +1,30 @@
1 +use thiserror::Error;
2 +use vllm_engine_core_client::Error as EngineCoreError;
3 +use vllm_llm::Error as LlmError;
4 +
5 +#[derive(Debug, Error)]
6 +pub enum Error {
7 + #[error("tokenizer error: {0}")]
8 + Tokenizer(String),
9 + #[error("text request `{request_id}` must contain at least one prompt token ID")]
10 + EmptyPromptTokenIds { request_id: String },
11 + #[error(
12 + "this model's maximum context length is {max_model_len} tokens, \
13 + but the prompt contains {prompt_len} input tokens"
14 + )]
15 + PromptTooLong { max_model_len: u32, prompt_len: u32 },
16 + #[error("text request stream `{request_id}` closed before terminal output")]
17 + StreamClosedBeforeTerminalOutput { request_id: String },
18 + #[error(transparent)]
19 + Llm(#[from] LlmError),
20 + #[error(transparent)]
21 + EngineCore(#[from] EngineCoreError),
22 +}
23 +
24 +pub type Result<T> = std::result::Result<T, Error>;
25 +
26 +impl From<vllm_tokenizer::TokenizerError> for Error {
27 + fn from(error: vllm_tokenizer::TokenizerError) -> Self {
28 + Self::Tokenizer(error.0)
29 + }
30 +}
rust/src/text/src/lib.rs new
+149
@@ -0,0 +1,149 @@
1 +//! Shared text-generation support used by chat and future raw completions.
2 +//!
3 +//! This crate intentionally stays below chat semantics:
4 +//! prompt text handling, tokenizer/model loading, incremental detokenization,
5 +//! and the thin generate-facing backend interface live here.
6 +
7 +use std::mem::take;
8 +
9 +pub use backend::{DynTextBackend, SamplingHints, TextBackend};
10 +pub use error::{Error, Result};
11 +use futures::Stream;
12 +pub use lower::{
13 + PreparedTextRequest, lower_sampling_params, lower_text_request, resolve_max_tokens,
14 +};
15 +pub use output::{
16 + CollectedTextOutput, DecodedLogprobs, DecodedPositionLogprobs, DecodedPromptLogprobs,
17 + DecodedTextEvent, DecodedTokenLogprob, Finished, TextDecodeOptions, TextOutputStreamExt,
18 +};
19 +pub use request::{Prompt, SamplingParams, TextRequest};
20 +use trait_set::trait_set;
21 +use vllm_engine_core_client::EngineCoreClient;
22 +pub use vllm_llm::FinishReason;
23 +use vllm_llm::{GenerateOutputStream, Llm};
24 +use vllm_tokenizer::DynTokenizer;
25 +
26 +pub mod backend;
27 +mod error;
28 +mod lower;
29 +pub mod output;
30 +mod request;
31 +pub use vllm_tokenizer as tokenizer;
32 +
33 +trait_set! {
34 + /// Shared streamed text output type used by raw completions and other text-only northbound paths.
35 + pub trait TextOutputStream = Stream<Item = Result<DecodedTextEvent>> + Send + 'static;
36 +}
37 +
38 +/// Raw text facade above [`Llm`].
39 +///
40 +/// This layer stays below chat semantics: prompt text or prompt token IDs flow
41 +/// in, decoded text deltas and terminal metadata flow out.
42 +pub struct TextLlm {
43 + /// Generate-only client owned by this text facade.
44 + llm: Llm,
45 + /// Tokenizer/model metadata backend responsible for prompt encode/decode
46 + /// and sampling hints.
47 + backend: DynTextBackend,
48 + /// Context window size derived by the backend or from engine startup
49 + /// handshake, with optional override from config.
50 + max_model_len: Option<u32>,
51 +}
52 +
53 +impl TextLlm {
54 + /// Create a new text-generation facade from a shared LLM client plus a text
55 + /// backend.
56 + pub fn new(llm: Llm, backend: DynTextBackend) -> Self {
57 + // Prefer the engine-reported max_model_len because it reflects the
58 + // post-profiling, auto-fitted KV cache limit rather than static
59 + // frontend metadata.
60 + let max_model_len = llm.engine_core_client().max_model_len();
61 +
62 + Self {
63 + llm,
64 + backend,
65 + max_model_len,
66 + }
67 + }
68 +
69 + /// Override the maximum model context length explicitly.
70 + ///
71 + /// This takes priority over both the engine-reported default and any
72 + /// tokenizer/model metadata exposed by the backend.
73 + pub fn with_max_model_len(mut self, max_model_len: u32) -> Self {
74 + self.max_model_len = Some(max_model_len);
75 + self
76 + }
77 +
78 + /// Return the backend model ID.
79 + pub fn model_id(&self) -> &str {
80 + self.backend.model_id()
81 + }
82 +
83 + /// Expose the underlying engine-core client for low-level utility/admin
84 + /// calls.
85 + pub fn engine_core_client(&self) -> &EngineCoreClient {
86 + self.llm.engine_core_client()
87 + }
88 +
89 + /// Return the tokenizer used by this text backend.
90 + pub fn tokenizer(&self) -> DynTokenizer {
91 + self.backend.tokenizer()
92 + }
93 +
94 + /// Tokenize if needed, lower to a generate request, and return the raw
95 + /// token stream.
96 + pub async fn generate_raw(&self, request: TextRequest) -> Result<GenerateOutputStream> {
97 + let (_, raw_stream) = self.generate_inner(request).await?;
98 + Ok(raw_stream)
99 + }
100 +
101 + /// Tokenize if needed, lower to a generate request, and stream
102 + /// incrementally decoded text.
103 + pub async fn generate(&self, request: TextRequest) -> Result<impl TextOutputStream> {
104 + let (text_request, raw_stream) = self.generate_inner(request).await?;
105 + let tokenizer = self.backend.tokenizer();
106 + let decoded_stream = output::decoded_text_event_stream(
107 + text_request.request_id,
108 + tokenizer,
109 + raw_stream,
110 + text_request.decode_options,
111 + text_request.intermediate,
112 + );
113 +
114 + Ok(decoded_stream)
115 + }
116 +
117 + async fn generate_inner(
118 + &self,
119 + mut request: TextRequest,
120 + ) -> Result<(TextRequest, GenerateOutputStream)> {
121 + request.validate()?;
122 +
123 + let tokenizer = self.backend.tokenizer();
124 + let prompt_token_ids = match take(&mut request.prompt) {
125 + Prompt::Text(text) => tokenizer.encode(&text, request.add_special_tokens)?,
126 + // Pre-tokenized prompts are the main completions-side escape hatch that lets benchmark
127 + // and infra workloads bypass chat rendering and tokenizer overhead entirely.
128 + Prompt::TokenIds(token_ids) => token_ids,
129 + };
130 +
131 + let mut sampling_hints = self.backend.sampling_hints()?;
132 + if let Some(max_model_len) = self.max_model_len {
133 + sampling_hints.max_model_len = Some(max_model_len);
134 + }
135 + let PreparedTextRequest {
136 + text_request,
137 + generate_request,
138 + } = lower_text_request(request, prompt_token_ids, sampling_hints, &*tokenizer)?;
139 +
140 + let raw_stream = self.llm.generate(generate_request).await?;
141 + Ok((text_request, raw_stream))
142 + }
143 +
144 + /// Shut down the underlying LLM client and its background tasks.
145 + pub async fn shutdown(self) -> Result<()> {
146 + self.llm.shutdown().await?;
147 + Ok(())
148 + }
149 +}
rust/src/text/src/lower.rs new
+739
@@ -0,0 +1,739 @@
1 +use std::collections::BTreeSet;
2 +
3 +use vllm_engine_core_client::protocol::EngineCoreSamplingParams;
4 +use vllm_llm::GenerateRequest;
5 +use vllm_tokenizer::Tokenizer;
6 +
7 +use crate::backend::SamplingHints;
8 +use crate::error::{Error, Result};
9 +use crate::request::{SamplingParams, TextRequest};
10 +
11 +/// One text request after it has been lowered into the raw generate boundary.
12 +#[derive(Debug)]
13 +pub struct PreparedTextRequest {
14 + /// The original high-level request, preserved for response-side metadata
15 + /// and decoding options.
16 + pub text_request: TextRequest,
17 + /// The southbound request ready to be sent to `vllm-llm`.
18 + pub generate_request: GenerateRequest,
19 +}
20 +
21 +/// Convert a high-level [`TextRequest`] into one lower-level
22 +/// [`GenerateRequest`] ready for the `llm` crate.
23 +pub fn lower_text_request(
24 + request: TextRequest,
25 + prompt_token_ids: Vec<u32>,
26 + sampling_hints: SamplingHints,
27 + tokenizer: &dyn Tokenizer,
28 +) -> Result<PreparedTextRequest> {
29 + let prompt_len = prompt_token_ids.len() as u32;
30 + let generate_request = GenerateRequest {
31 + request_id: request.request_id.clone(),
32 + prompt_token_ids,
33 + mm_features: request.mm_features.clone(),
34 + sampling_params: lower_sampling_params(
35 + request.sampling_params.clone(),
36 + sampling_hints,
37 + prompt_len,
38 + tokenizer,
39 + )?,
40 + cache_salt: request.cache_salt.clone(),
41 + priority: request.priority,
42 + data_parallel_rank: request.data_parallel_rank,
43 + // Fields below are currently placeholders.
44 + arrival_time: None,
45 + trace_headers: None,
46 + reasoning_ended: None,
47 + lora_request: None,
48 + };
49 +
50 + Ok(PreparedTextRequest {
51 + text_request: request,
52 + generate_request,
53 + })
54 +}
55 +
56 +/// Convert [`SamplingParams`] into [`EngineCoreSamplingParams`], enriching
57 +/// omitted user values with tokenizer/model-derived hints when available.
58 +pub fn lower_sampling_params(
59 + sampling_params: SamplingParams,
60 + SamplingHints {
61 + primary_eos_token_id,
62 + extra_eos_token_ids,
63 + default_temperature,
64 + default_top_p,
65 + default_top_k,
66 + default_min_p,
67 + default_repetition_penalty,
68 + default_max_tokens,
69 + max_model_len,
70 + }: SamplingHints,
71 + prompt_len: u32,
72 + tokenizer: &dyn Tokenizer,
73 +) -> Result<EngineCoreSamplingParams> {
74 + let SamplingParams {
75 + temperature,
76 + top_p,
77 + top_k,
78 + seed,
79 + max_tokens,
80 + min_tokens,
81 + logprobs,
82 + prompt_logprobs,
83 + min_p,
84 + frequency_penalty,
85 + presence_penalty,
86 + repetition_penalty,
87 + stop_token_ids,
88 + ignore_eos,
89 + logit_bias,
90 + allowed_token_ids,
91 + bad_words,
92 + logprob_token_ids,
93 + structured_outputs,
94 + skip_reading_prefix_cache,
95 + vllm_xargs,
96 + } = sampling_params;
97 +
98 + // Mirrors the model-generation-config inheritance used by vLLM's OpenAI chat
99 + // path: https://github.com/vllm-project/vllm/blob/bc2c0c86efb28e77677a3cfb8687e976914a313a/vllm/entrypoints/openai/chat_completion/protocol.py#L424-L450
100 + // If neither the caller nor the model provides a value, fall back to 1.0 — the
101 + // default used by the Python vLLM OpenAI-compatible API (via
102 + // `_DEFAULT_SAMPLING_PARAMS`).
103 + let temperature = temperature.or(default_temperature).unwrap_or(1.0);
104 + let top_p = top_p.or(default_top_p).unwrap_or(1.0);
105 + let top_k = top_k.or(default_top_k).unwrap_or(0);
106 + let min_p = min_p.or(default_min_p).unwrap_or(0.0);
107 + let repetition_penalty = repetition_penalty.or(default_repetition_penalty).unwrap_or(1.0);
108 + let max_tokens = resolve_max_tokens(max_tokens, default_max_tokens, max_model_len, prompt_len)?;
109 + let min_tokens = min_tokens.unwrap_or(0);
110 + let frequency_penalty = frequency_penalty.unwrap_or(0.0);
111 + let presence_penalty = presence_penalty.unwrap_or(0.0);
112 +
113 + let mut stop_token_ids = stop_token_ids.unwrap_or_default();
114 + let mut all_stop_token_ids = BTreeSet::from_iter(stop_token_ids.iter().copied());
115 + if let Some(primary_eos_token_id) = primary_eos_token_id {
116 + all_stop_token_ids.insert(primary_eos_token_id);
117 + }
118 + all_stop_token_ids.extend(extra_eos_token_ids.iter().copied());
119 +
120 + if !ignore_eos {
121 + merge_unique_token_ids(&mut stop_token_ids, extra_eos_token_ids.iter().copied());
122 + }
123 +
124 + Ok(EngineCoreSamplingParams {
125 + temperature,
126 + top_p,
127 + top_k,
128 + seed,
129 + max_tokens,
130 + min_tokens,
131 + logprobs,
132 + prompt_logprobs,
133 + min_p,
134 + frequency_penalty,
135 + presence_penalty,
136 + repetition_penalty,
137 + stop_token_ids,
138 + eos_token_id: (!ignore_eos).then_some(primary_eos_token_id).flatten(),
139 + all_stop_token_ids,
140 + logit_bias,
141 + allowed_token_ids,
142 + bad_words_token_ids: tokenize_bad_words(bad_words.as_deref(), tokenizer)?,
143 + structured_outputs,
144 + logprob_token_ids,
145 + skip_reading_prefix_cache,
146 + extra_args: vllm_xargs,
147 + })
148 +}
149 +
150 +/// Convert bad-word strings into token-ID sequences, following the Python vLLM
151 +/// logic in `SamplingParams.update_from_tokenizer()`.
152 +///
153 +/// Each word is encoded both with and without a leading space so that the ban
154 +/// applies regardless of whether the word appears at the beginning or in the
155 +/// middle of generated text (this accounts for tokenizers that use an
156 +/// `add_prefix_space` convention).
157 +///
158 +/// Reference: <https://github.com/vllm-project/vllm/blob/f22d6e026/vllm/sampling_params.py#L555-L594>
159 +fn tokenize_bad_words(
160 + bad_words: Option<&[String]>,
161 + tokenizer: &dyn Tokenizer,
162 +) -> Result<Option<Vec<Vec<u32>>>> {
163 + let bad_words = bad_words.filter(|w| !w.is_empty());
164 + let mut all_token_ids = Vec::new();
165 +
166 + for bad_word in bad_words.into_iter().flatten() {
167 + // Without a leading space we always keep the encoding.
168 + // With a leading space we only keep it when the prefix-space variant produces a
169 + // distinct first token but the same sequence length — this mirrors the Python
170 + // dedup condition that avoids redundant entries.
171 + let without_space = tokenizer.encode(bad_word, false)?;
172 + let with_space = tokenizer.encode(&format!(" {}", bad_word.trim_start()), false)?;
173 +
174 + if !without_space.is_empty() {
175 + all_token_ids.push(without_space);
176 + }
177 + if !with_space.is_empty()
178 + && all_token_ids.last().is_some_and(|prev: &Vec<u32>| {
179 + with_space[0] != prev[0] && with_space.len() == prev.len()
180 + })
181 + {
182 + all_token_ids.push(with_space);
183 + }
184 + }
185 +
186 + Ok((!all_token_ids.is_empty()).then_some(all_token_ids))
187 +}
188 +
189 +/// Resolve the effective `max_tokens` for generation, mirroring vLLM Python's
190 +/// `get_max_tokens()` in `vllm/entrypoints/utils.py`.
191 +///
192 +/// Takes the minimum of all available limits (user-specified, generation-config
193 +/// default, and `max_model_len - prompt_len`). When nothing is known, falls
194 +/// back to `u32::MAX` so the engine-core can apply its own context-window
195 +/// limit.
196 +pub fn resolve_max_tokens(
197 + user_max_tokens: Option<u32>,
198 + default_max_tokens: Option<u32>,
199 + max_model_len: Option<u32>,
200 + prompt_len: u32,
201 +) -> Result<u32> {
202 + let model_max_tokens = match max_model_len {
203 + Some(max_model_len) if prompt_len >= max_model_len => {
204 + return Err(Error::PromptTooLong {
205 + max_model_len,
206 + prompt_len,
207 + });
208 + }
209 + Some(max_model_len) => Some(max_model_len - prompt_len),
210 + None => None,
211 + };
212 +
213 + let fallback_max_tokens = user_max_tokens.or(default_max_tokens);
214 + Ok([fallback_max_tokens, model_max_tokens]
215 + .into_iter()
216 + .flatten()
217 + .min()
218 + .unwrap_or(u32::MAX /* TODO: a reasonable fallback? */))
219 +}
220 +
221 +fn merge_unique_token_ids(
222 + stop_token_ids: &mut Vec<u32>,
223 + extra_token_ids: impl Iterator<Item = u32>,
224 +) {
225 + // Keep user-provided ordering stable while still folding in backend-derived EOS
226 + // aliases.
227 + for token_id in extra_token_ids {
228 + if !stop_token_ids.contains(&token_id) {
229 + stop_token_ids.push(token_id);
230 + }
231 + }
232 +}
233 +
234 +#[cfg(test)]
235 +mod tests {
236 + use std::collections::BTreeSet;
237 +
238 + use super::*;
239 + use crate::backend::hf::HfTextBackend;
240 + use crate::backend::{SamplingHints, TextBackend as _};
241 + use crate::request::{Prompt, TextRequest};
242 +
243 + /// Stub tokenizer that returns empty token IDs — sufficient for tests that
244 + /// don't exercise bad-words tokenization.
245 + struct StubTokenizer;
246 +
247 + impl Tokenizer for StubTokenizer {
248 + fn encode(
249 + &self,
250 + _text: &str,
251 + _add_special_tokens: bool,
252 + ) -> vllm_tokenizer::Result<Vec<u32>> {
253 + Ok(vec![])
254 + }
255 +
256 + fn decode(
257 + &self,
258 + _token_ids: &[u32],
259 + _skip_special_tokens: bool,
260 + ) -> vllm_tokenizer::Result<String> {
261 + Ok(String::new())
262 + }
263 +
264 + fn token_to_id(&self, _token: &str) -> Option<u32> {
265 + None
266 + }
267 + }
268 +
269 + fn stub_tokenizer() -> StubTokenizer {
270 + StubTokenizer
271 + }
272 +
273 + fn sample_request() -> TextRequest {
274 + TextRequest {
275 + prompt: Prompt::TokenIds(vec![1, 2, 3]),
276 + request_id: "text-1".to_string(),
277 + ..TextRequest::for_test()
278 + }
279 + }
280 +
281 + fn sample_sampling_hints() -> SamplingHints {
282 + SamplingHints {
283 + primary_eos_token_id: Some(99),
284 + extra_eos_token_ids: BTreeSet::from([77]),
285 + default_temperature: None,
286 + default_top_p: None,
287 + default_top_k: None,
288 + default_min_p: None,
289 + default_repetition_penalty: None,
290 + default_max_tokens: None,
291 + max_model_len: None,
292 + }
293 + }
294 +
295 + #[test]
296 + fn lower_text_request_applies_python_style_eos_hints() {
297 + let prepared = lower_text_request(
298 + sample_request(),
299 + vec![1, 2, 3],
300 + sample_sampling_hints(),
301 + &stub_tokenizer(),
302 + )
303 + .unwrap();
304 +
305 + let params = prepared.generate_request.sampling_params;
306 + expect_test::expect![[r#"
307 + EngineCoreSamplingParams {
308 + temperature: 1.0,
309 + top_p: 1.0,
310 + top_k: 0,
311 + seed: None,
312 + max_tokens: 4294967295,
313 + min_tokens: 0,
314 + logprobs: None,
315 + prompt_logprobs: None,
316 + min_p: 0.0,
317 + frequency_penalty: 0.0,
318 + presence_penalty: 0.0,
319 + repetition_penalty: 1.0,
320 + stop_token_ids: [
321 + 77,
322 + ],
323 + eos_token_id: Some(
324 + 99,
325 + ),
326 + all_stop_token_ids: {
327 + 77,
328 + 99,
329 + },
330 + logit_bias: None,
331 + allowed_token_ids: None,
332 + bad_words_token_ids: None,
333 + structured_outputs: None,
334 + logprob_token_ids: None,
335 + skip_reading_prefix_cache: None,
336 + extra_args: None,
337 + }
338 + "#]]
339 + .assert_debug_eq(&params);
340 + }
341 +
342 + #[test]
343 + fn lower_text_request_respects_ignore_eos_for_stop_token_ids() {
344 + let mut request = sample_request();
345 + request.sampling_params.ignore_eos = true;
346 +
347 + let prepared = lower_text_request(
348 + request,
349 + vec![1, 2, 3],
350 + sample_sampling_hints(),
351 + &stub_tokenizer(),
352 + )
353 + .unwrap();
354 +
355 + let params = prepared.generate_request.sampling_params;
356 + expect_test::expect![[r#"
357 + EngineCoreSamplingParams {
358 + temperature: 1.0,
359 + top_p: 1.0,
360 + top_k: 0,
361 + seed: None,
362 + max_tokens: 4294967295,
363 + min_tokens: 0,
364 + logprobs: None,
365 + prompt_logprobs: None,
366 + min_p: 0.0,
367 + frequency_penalty: 0.0,
368 + presence_penalty: 0.0,
369 + repetition_penalty: 1.0,
370 + stop_token_ids: [],
371 + eos_token_id: None,
372 + all_stop_token_ids: {
373 + 77,
374 + 99,
375 + },
376 + logit_bias: None,
377 + allowed_token_ids: None,
378 + bad_words_token_ids: None,
379 + structured_outputs: None,
380 + logprob_token_ids: None,
381 + skip_reading_prefix_cache: None,
382 + extra_args: None,
383 + }
384 + "#]]
385 + .assert_debug_eq(&params);
386 + }
387 +
388 + #[tokio::test]
389 + #[ignore = "requires network access to Hugging Face"]
390 + async fn lower_text_request_uses_real_qwen_generation_defaults() {
391 + let backend = HfTextBackend::from_model("Qwen/Qwen3-0.6B")
392 + .await
393 + .expect("load qwen tokenizer and generation config");
394 + let hints = backend.sampling_hints().expect("collect sampling hints");
395 +
396 + expect_test::expect![[r#"
397 + SamplingHints {
398 + primary_eos_token_id: Some(
399 + 151645,
400 + ),
401 + extra_eos_token_ids: {
402 + 151643,
403 + },
404 + default_temperature: Some(
405 + 0.6,
406 + ),
407 + default_top_p: Some(
408 + 0.95,
409 + ),
410 + default_top_k: Some(
411 + 20,
412 + ),
413 + default_min_p: Some(
414 + 0.1,
415 + ),
416 + default_repetition_penalty: Some(
417 + 1.2,
418 + ),
419 + default_max_tokens: None,
420 + max_model_len: Some(
421 + 40960,
422 + ),
423 + }
424 + "#]]
425 + .assert_debug_eq(&hints);
426 +
427 + let prepared =
428 + lower_text_request(sample_request(), vec![1, 2, 3], hints, &stub_tokenizer())
429 + .expect("lower request");
430 + let params = prepared.generate_request.sampling_params;
431 +
432 + expect_test::expect![[r#"
433 + EngineCoreSamplingParams {
434 + temperature: 0.6,
435 + top_p: 0.95,
436 + top_k: 20,
437 + seed: None,
438 + max_tokens: 40957,
439 + min_tokens: 0,
440 + logprobs: None,
441 + prompt_logprobs: None,
442 + min_p: 0.1,
443 + frequency_penalty: 0.0,
444 + presence_penalty: 0.0,
445 + repetition_penalty: 1.2,
446 + stop_token_ids: [
447 + 151643,
448 + ],
449 + eos_token_id: Some(
450 + 151645,
451 + ),
452 + all_stop_token_ids: {
453 + 151643,
454 + 151645,
455 + },
456 + }
457 + "#]]
458 + .assert_debug_eq(&params);
459 + }
460 +
461 + #[test]
462 + fn lower_sampling_params_preserves_explicit_stop_token_ids_in_all_stop_set() {
463 + let sampling_params = SamplingParams {
464 + stop_token_ids: Some(vec![11, 77]),
465 + ..SamplingParams::default()
466 + };
467 +
468 + let params = lower_sampling_params(
469 + sampling_params,
470 + SamplingHints {
471 + primary_eos_token_id: Some(99),
472 + extra_eos_token_ids: BTreeSet::from([77, 88]),
473 + default_temperature: None,
474 + default_top_p: None,
475 + default_top_k: None,
476 + default_min_p: None,
477 + default_repetition_penalty: None,
478 + default_max_tokens: None,
479 + max_model_len: None,
480 + },
481 + 3,
482 + &stub_tokenizer(),
483 + )
484 + .unwrap();
485 +
486 + expect_test::expect![[r#"
487 + EngineCoreSamplingParams {
488 + temperature: 1.0,
489 + top_p: 1.0,
490 + top_k: 0,
491 + seed: None,
492 + max_tokens: 4294967295,
493 + min_tokens: 0,
494 + logprobs: None,
495 + prompt_logprobs: None,
496 + min_p: 0.0,
497 + frequency_penalty: 0.0,
498 + presence_penalty: 0.0,
499 + repetition_penalty: 1.0,
500 + stop_token_ids: [
501 + 11,
502 + 77,
503 + 88,
504 + ],
505 + eos_token_id: Some(
506 + 99,
507 + ),
508 + all_stop_token_ids: {
509 + 11,
510 + 77,
511 + 88,
512 + 99,
513 + },
514 + logit_bias: None,
515 + allowed_token_ids: None,
516 + bad_words_token_ids: None,
517 + structured_outputs: None,
518 + logprob_token_ids: None,
519 + skip_reading_prefix_cache: None,
520 + extra_args: None,
521 + }
522 + "#]]
523 + .assert_debug_eq(&params);
524 + }
525 +
526 + #[test]
527 + fn lower_sampling_params_prefers_user_values_over_generation_defaults() {
528 + let sampling_params = SamplingParams {
529 + temperature: Some(0.2),
530 + top_p: Some(0.3),
531 + top_k: Some(4),
532 + max_tokens: Some(32),
533 + min_tokens: Some(2),
534 + ..Default::default()
535 + };
536 +
537 + let params = lower_sampling_params(
538 + sampling_params,
539 + SamplingHints {
540 + primary_eos_token_id: None,
541 + extra_eos_token_ids: BTreeSet::new(),
542 + default_temperature: Some(0.8),
543 + default_top_p: Some(0.9),
544 + default_top_k: Some(12),
545 + default_min_p: Some(0.1),
546 + default_repetition_penalty: Some(1.2),
547 + default_max_tokens: Some(128),
548 + max_model_len: None,
549 + },
550 + 3,
551 + &stub_tokenizer(),
552 + )
553 + .unwrap();
554 +
555 + expect_test::expect![[r#"
556 + EngineCoreSamplingParams {
557 + temperature: 0.2,
558 + top_p: 0.3,
559 + top_k: 4,
560 + seed: None,
561 + max_tokens: 32,
562 + min_tokens: 2,
563 + logprobs: None,
564 + prompt_logprobs: None,
565 + min_p: 0.1,
566 + frequency_penalty: 0.0,
567 + presence_penalty: 0.0,
568 + repetition_penalty: 1.2,
569 + stop_token_ids: [],
570 + eos_token_id: None,
571 + all_stop_token_ids: {},
572 + logit_bias: None,
573 + allowed_token_ids: None,
574 + bad_words_token_ids: None,
575 + structured_outputs: None,
576 + logprob_token_ids: None,
577 + skip_reading_prefix_cache: None,
578 + extra_args: None,
579 + }
580 + "#]]
581 + .assert_debug_eq(&params);
582 + }
583 +
584 + #[test]
585 + fn lower_sampling_params_passes_logprobs_fields_through() {
586 + let sampling_params = SamplingParams {
587 + logprobs: Some(3),
588 + prompt_logprobs: Some(-1),
589 + ..Default::default()
590 + };
591 +
592 + let params = lower_sampling_params(
593 + sampling_params,
594 + SamplingHints {
595 + primary_eos_token_id: None,
596 + extra_eos_token_ids: BTreeSet::new(),
597 + default_temperature: None,
598 + default_top_p: None,
599 + default_top_k: None,
600 + default_min_p: None,
601 + default_repetition_penalty: None,
602 + default_max_tokens: None,
603 + max_model_len: None,
604 + },
605 + 3,
606 + &stub_tokenizer(),
607 + )
608 + .unwrap();
609 +
610 + assert_eq!(params.logprobs, Some(3));
611 + assert_eq!(params.prompt_logprobs, Some(-1));
612 + }
613 +
614 + #[test]
615 + fn lower_sampling_params_uses_generation_defaults_when_user_omits_values() {
616 + let params = lower_sampling_params(
617 + SamplingParams::default(),
618 + SamplingHints {
619 + primary_eos_token_id: None,
620 + extra_eos_token_ids: BTreeSet::new(),
621 + default_temperature: Some(0.8),
622 + default_top_p: Some(0.9),
623 + default_top_k: Some(12),
624 + default_min_p: Some(0.1),
625 + default_repetition_penalty: Some(1.2),
626 + default_max_tokens: Some(128),
627 + max_model_len: None,
628 + },
629 + 3,
630 + &stub_tokenizer(),
631 + )
632 + .unwrap();
633 +
634 + expect_test::expect![[r#"
635 + EngineCoreSamplingParams {
636 + temperature: 0.8,
637 + top_p: 0.9,
638 + top_k: 12,
639 + seed: None,
640 + max_tokens: 128,
641 + min_tokens: 0,
642 + logprobs: None,
643 + prompt_logprobs: None,
644 + min_p: 0.1,
645 + frequency_penalty: 0.0,
646 + presence_penalty: 0.0,
647 + repetition_penalty: 1.2,
648 + stop_token_ids: [],
649 + eos_token_id: None,
650 + all_stop_token_ids: {},
651 + logit_bias: None,
652 + allowed_token_ids: None,
653 + bad_words_token_ids: None,
654 + structured_outputs: None,
655 + logprob_token_ids: None,
656 + skip_reading_prefix_cache: None,
657 + extra_args: None,
658 + }
659 + "#]]
660 + .assert_debug_eq(&params);
661 + }
662 +
663 + #[test]
664 + fn resolve_max_tokens_caps_by_model_len() {
665 + let result = resolve_max_tokens(Some(150), None, Some(200), 100);
666 + assert_eq!(result.unwrap(), 100);
667 + }
668 +
669 + #[test]
670 + fn lower_text_request_preserves_non_streaming_request_metadata() {
671 + let mut request = sample_request();
672 + request.intermediate = false;
673 +
674 + let prepared = lower_text_request(
675 + request,
676 + vec![1, 2, 3],
677 + sample_sampling_hints(),
678 + &stub_tokenizer(),
679 + )
680 + .unwrap();
681 +
682 + assert!(!prepared.text_request.intermediate);
683 + assert_eq!(prepared.generate_request.request_id, "text-1");
684 + }
685 +
686 + #[test]
687 + fn resolve_max_tokens_user_smaller_than_model_limit() {
688 + let result = resolve_max_tokens(Some(50), None, Some(200), 100);
689 + assert_eq!(result.unwrap(), 50);
690 + }
691 +
692 + #[test]
693 + fn resolve_max_tokens_uses_default_when_user_omits() {
694 + let result = resolve_max_tokens(None, Some(64), Some(200), 100);
695 + assert_eq!(result.unwrap(), 64);
696 + }
697 +
698 + #[test]
699 + fn resolve_max_tokens_default_capped_by_model_len() {
700 + let result = resolve_max_tokens(None, Some(256), Some(200), 100);
701 + assert_eq!(result.unwrap(), 100);
702 + }
703 +
704 + #[test]
705 + fn resolve_max_tokens_no_model_len_falls_back() {
706 + let result = resolve_max_tokens(Some(9999), None, None, 100);
707 + assert_eq!(result.unwrap(), 9999);
708 + }
709 +
710 + #[test]
711 + fn resolve_max_tokens_no_limits_known_falls_back_to_u32_max() {
712 + let result = resolve_max_tokens(None, None, None, 100);
713 + assert_eq!(result.unwrap(), u32::MAX);
714 + }
715 +
716 + #[test]
717 + fn resolve_max_tokens_prompt_too_long() {
718 + let result = resolve_max_tokens(Some(10), None, Some(100), 100);
719 + assert!(matches!(
720 + result,
721 + Err(Error::PromptTooLong {
722 + max_model_len: 100,
723 + prompt_len: 100,
724 + })
725 + ));
726 + }
727 +
728 + #[test]
729 + fn resolve_max_tokens_prompt_exceeds_model_len() {
730 + let result = resolve_max_tokens(Some(10), None, Some(100), 200);
731 + assert!(matches!(
732 + result,
733 + Err(Error::PromptTooLong {
734 + max_model_len: 100,
735 + prompt_len: 200,
736 + })
737 + ));
738 + }
739 +}
rust/src/text/src/output/decoded.rs new
+607
@@ -0,0 +1,607 @@
1 +use std::sync::Arc;
2 +
3 +use asynk_strim_attr::{TryYielder, try_stream};
4 +use futures::{Stream, StreamExt};
5 +use serde::{Deserialize, Serialize};
6 +use tracing::{Level, debug, trace};
7 +use vllm_engine_core_client::AbortCause;
8 +use vllm_engine_core_client::protocol::StopReason;
9 +use vllm_llm::{FinishReason, GenerateOutput};
10 +use vllm_tokenizer::{DynTokenizer, IncrementalDecoder};
11 +
12 +use super::logprobs::{
13 + DecodedLogprobs, DecodedPromptLogprobs, decode_logprobs, decode_prompt_logprobs,
14 +};
15 +use crate::error::Error;
16 +
17 +/// Request-neutral options for incremental text decoding.
18 +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
19 +pub struct TextDecodeOptions {
20 + pub skip_special_tokens: bool,
21 + pub include_stop_str_in_output: bool,
22 + pub stop_strings: Option<Vec<String>>,
23 + /// Minimum number of tokens to generate before stop-string checking kicks
24 + /// in. Stop strings found within the first `min_tokens` tokens are
25 + /// ignored.
26 + pub min_tokens: u32,
27 +}
28 +
29 +impl Default for TextDecodeOptions {
30 + fn default() -> Self {
31 + Self {
32 + skip_special_tokens: true,
33 + include_stop_str_in_output: false,
34 + stop_strings: None,
35 + min_tokens: 0,
36 + }
37 + }
38 +}
39 +
40 +/// Terminal metadata carried on the final [`DecodedTextEvent`].
41 +#[derive(Debug, Clone, PartialEq)]
42 +pub struct Finished {
43 + pub prompt_token_count: usize,
44 + pub output_token_count: usize,
45 + pub finish_reason: FinishReason,
46 + /// Connector-specific KV transfer parameters for disaggregated serving.
47 + pub kv_transfer_params: Option<serde_json::Value>,
48 +}
49 +
50 +/// Internal decoded-text event emitted before higher-level assistant
51 +/// adaptation.
52 +#[derive(Debug, Clone, PartialEq)]
53 +pub enum DecodedTextEvent {
54 + /// The request has reached the point where prompt-scoped decoding metadata
55 + /// is ready.
56 + Start {
57 + /// The actual prompt token IDs for this request.
58 + prompt_token_ids: Arc<[u32]>,
59 + /// Once-only prompt logprobs metadata, when requested.
60 + ///
61 + /// The first prompt token is carried separately because it has no left
62 + /// context to score against; `scored_positions` covers the
63 + /// remaining prompt positions.
64 + prompt_logprobs: Option<DecodedPromptLogprobs>,
65 + },
66 + /// A delta of text has been decoded, optionally alongside token-position
67 + /// logprobs.
68 + ///
69 + /// `delta` is the newly visible decoded text fragment for this update.
70 + ///
71 + /// `logprobs` covers the newly generated token positions from the same
72 + /// update, but is not guaranteed to align with `delta` by character
73 + /// span. One update may carry token logprobs but no newly visible text
74 + /// yet, and one visible text fragment may reflect multiple token
75 + /// positions becoming decodable together.
76 + ///
77 + /// Upper-level may further parse `delta` as reasoning or tool calls.
78 + ///
79 + /// When `finished` is `Some`, this is the terminal event for the request.
80 + TextDelta {
81 + delta: String,
82 + token_ids: Vec<u32>,
83 + logprobs: Option<DecodedLogprobs>,
84 + finished: Option<Finished>,
85 + },
86 +}
87 +
88 +/// Convert the output token stream from the `vllm_llm` layer into incrementally
89 +/// decoded text.
90 +#[try_stream]
91 +pub async fn decoded_text_event_stream(
92 + request_id: String,
93 + tokenizer: DynTokenizer,
94 + mut raw_stream: impl Stream<Item = vllm_llm::Result<GenerateOutput>> + Unpin,
95 + mut decode_options: TextDecodeOptions,
96 + intermediate: bool,
97 + mut y: TryYielder<DecodedTextEvent, Error>,
98 +) -> crate::Result<()> {
99 + let mut decoder: Option<Box<dyn IncrementalDecoder>> = None;
100 + let mut prompt_token_count = 0_usize;
101 + let mut token_ids = Vec::new();
102 + let mut output_token_count: usize = 0;
103 + let mut logprobs: Option<DecodedLogprobs> = None;
104 +
105 + while let Some(next) = raw_stream.next().await {
106 + let output = next?;
107 +
108 + // If it's the first output, init states and yield `Start` event.
109 + if decoder.is_none() {
110 + let prompt_token_ids =
111 + output.prompt_token_ids().expect("first llm output must carry prompt token ids");
112 + prompt_token_count = prompt_token_ids.len();
113 +
114 + let dec = tokenizer.create_decode_stream(
115 + prompt_token_ids,
116 + decode_options.skip_special_tokens,
117 + // If we are excluding stop strings from output, we need to buffer
118 + // the output so that we don't return the beginning of a stop string
119 + // when streaming the outputs.
120 + match decode_options.include_stop_str_in_output {
121 + true => 0,
122 + false => {
123 + decode_options
124 + .stop_strings
125 + .as_ref()
126 + .and_then(|stops| stops.iter().map(|ss| ss.len()).max())
127 + .unwrap_or(1)
128 + - 1
129 + }
130 + },
131 + );
132 + decoder = Some(dec);
133 +
134 + y.yield_ok(DecodedTextEvent::Start {
135 + prompt_token_ids: prompt_token_ids.clone(),
136 + prompt_logprobs: output
137 + .prompt_logprobs()
138 + .map(|logprobs| {
139 + decode_prompt_logprobs(
140 + tokenizer.as_ref(),
141 + prompt_token_ids,
142 + logprobs,
143 + decode_options.skip_special_tokens,
144 + )
145 + })
146 + .transpose()?,
147 + })
148 + .await;
149 + };
150 + let decoder = decoder.as_mut().unwrap();
151 +
152 + let kv_transfer_params = output.kv_transfer_params;
153 + let mut finish_reason = output.finish_reason;
154 + let mut stop_str_matched = false;
155 + let suppress_terminal_stop_token = finish_reason.as_ref().is_some_and(|r| r.is_stop())
156 + && !decode_options.include_stop_str_in_output;
157 + let decodable_token_ids = if suppress_terminal_stop_token {
158 + // Match Python V1 token-stop detokenization by keeping the stop token
159 + // in metadata while excluding it from user-visible text.
160 + output.token_ids.split_last().map(|(_, rest)| rest).unwrap_or(&[])
161 + } else {
162 + &output.token_ids
163 + };
164 +
165 + let mut delta: Option<String> = None;
166 + let mut truncate_output_to = None;
167 + let mut truncate_tokens_to = None;
168 + for (tok_idx, &token_id) in decodable_token_ids.iter().enumerate() {
169 + let new_bytes = decoder.push_token(token_id)?;
170 + if output_token_count + tok_idx + 1 > decode_options.min_tokens as usize
171 + && let Some(stops) = decode_options.stop_strings.as_mut()
172 + && let Some((idx, off)) = matches_stop_string(stops, decoder.output(), new_bytes)
173 + {
174 + let stop_str = stops.swap_remove(idx);
175 + truncate_output_to = match decode_options.include_stop_str_in_output {
176 + true => Some(off + stop_str.len()),
177 + false => Some(off),
178 + };
179 + finish_reason = Some(FinishReason::Stop(Some(StopReason::Text(stop_str))));
180 + truncate_tokens_to = Some(tok_idx + 1);
181 + stop_str_matched = true;
182 +
183 + break;
184 + }
185 +
186 + if intermediate && let Some(chunk) = decoder.next_chunk() {
187 + if let Some(delta_str) = delta.as_mut() {
188 + delta_str.push_str(&chunk);
189 + } else {
190 + delta = Some(chunk);
191 + }
192 + }
193 + }
194 +
195 + let mut new_token_ids = output.token_ids;
196 + let mut new_logprobs = output.logprobs;
197 +
198 + // Trim tokens and logprobs if we matched stop string.
199 + if let Some(num_tokens) = truncate_tokens_to {
200 + new_token_ids.truncate(num_tokens);
201 + if let Some(logprobs) = &mut new_logprobs {
202 + logprobs.positions.truncate(num_tokens);
203 + }
204 + }
205 +
206 + output_token_count += new_token_ids.len();
207 +
208 + let decoded_logprobs = new_logprobs
209 + .as_ref()
210 + .map(|logprobs| {
211 + decode_logprobs(
212 + tokenizer.as_ref(),
213 + logprobs,
214 + decode_options.skip_special_tokens,
215 + )
216 + })
217 + .transpose()?;
218 +
219 + if !intermediate {
220 + token_ids.extend(&new_token_ids);
221 + if let Some(dlp) = decoded_logprobs.as_ref() {
222 + logprobs
223 + .get_or_insert_with(|| DecodedLogprobs { positions: vec![] })
224 + .positions
225 + .extend_from_slice(&dlp.positions);
226 + }
227 + }
228 +
229 + if let Some(reason) = finish_reason {
230 + // Flush any remaining buffered text.
231 + let (last_chunk, mut text) = decoder.flush(truncate_output_to)?;
232 + let text_len = text.len();
233 + let full_text = tracing::enabled!(Level::TRACE).then(|| text.clone());
234 +
235 + if intermediate {
236 + if let Some(chunk) = last_chunk {
237 + if let Some(delta_str) = delta.as_mut() {
238 + delta_str.push_str(&chunk);
239 + } else {
240 + delta = Some(chunk);
241 + }
242 + }
243 + token_ids = new_token_ids;
244 + logprobs = decoded_logprobs;
245 + text = delta.unwrap_or_default();
246 + }
247 +
248 + debug!(
249 + finish_reason = ?reason,
250 + text_length_bytes = text_len,
251 + output_token_count = output_token_count,
252 + "request finished with terminal output"
253 + );
254 + if let Some(full_text) = full_text {
255 + trace!(full_text, "request finished with terminal decoded text");
256 + }
257 +
258 + // Intentionally drop the stream with explicit cause, so that the engine core
259 + // can distinguish between such normal completion vs an unexpected
260 + // early drop.
261 + if stop_str_matched {
262 + AbortCause::StopStringMatched.drop_as(raw_stream);
263 + }
264 +
265 + y.yield_ok(DecodedTextEvent::TextDelta {
266 + delta: text,
267 + token_ids,
268 + logprobs,
269 + finished: Some(Finished {
270 + prompt_token_count,
271 + output_token_count,
272 + finish_reason: reason,
273 + kv_transfer_params,
274 + }),
275 + })
276 + .await;
277 + return Ok(());
278 + }
279 +
280 + if intermediate {
281 + y.yield_ok(DecodedTextEvent::TextDelta {
282 + delta: delta.unwrap_or_default(),
283 + token_ids: new_token_ids,
284 + logprobs: decoded_logprobs,
285 + finished: None,
286 + })
287 + .await;
288 + }
289 + }
290 +
291 + Err(Error::StreamClosedBeforeTerminalOutput { request_id })
292 +}
293 +
294 +/// If stop string matches, returns tuple
295 +/// (index into stop string vec, byte index of first byte of stop string in
296 +/// output)
297 +fn matches_stop_string(stops: &[String], output: &str, new_bytes: usize) -> Option<(usize, usize)> {
298 + // We compare byte subslices to avoid utf8 boundary problem
299 + let output = output.as_bytes();
300 + let next_off = (output.len() + 1) - new_bytes;
301 + stops
302 + .iter()
303 + .map(|ss| (ss.as_bytes(), ss.len(), next_off.saturating_sub(ss.len())))
304 + .enumerate()
305 + .find_map(|(ss_idx, (ss, len, start_off))| {
306 + output[start_off..]
307 + .windows(len)
308 + .rposition(|w| w == ss)
309 + .map(|pos| (ss_idx, start_off + pos))
310 + })
311 +}
312 +
313 +#[cfg(test)]
314 +mod tests {
315 + use std::pin::Pin;
316 + use std::sync::{Arc, Mutex};
317 + use std::task::{Context, Poll};
318 +
319 + use futures::{Stream, stream};
320 + use vllm_engine_core_client::AbortCause;
321 + use vllm_llm::GenerateOutput;
322 + use vllm_tokenizer::Tokenizer;
323 +
324 + use super::*;
325 + use crate::output::TextOutputStreamExt as _;
326 +
327 + /// Backend that treats each token ID as a raw byte, producing lossy UTF-8.
328 + struct ByteTokenizer;
329 +
330 + impl Tokenizer for ByteTokenizer {
331 + fn encode(
332 + &self,
333 + _text: &str,
334 + _add_special_tokens: bool,
335 + ) -> vllm_tokenizer::Result<Vec<u32>> {
336 + unreachable!()
337 + }
338 +
339 + fn decode(
340 + &self,
341 + token_ids: &[u32],
342 + _skip_special_tokens: bool,
343 + ) -> vllm_tokenizer::Result<String> {
344 + let bytes = token_ids.iter().map(|id| *id as u8).collect::<Vec<_>>();
345 + Ok(String::from_utf8_lossy(&bytes).into_owned())
346 + }
347 +
348 + fn token_to_id(&self, _token: &str) -> Option<u32> {
349 + unreachable!()
350 + }
351 + }
352 +
353 + /// Helper: run `decoded_text_event_stream` to completion and return the
354 + /// collected output.
355 + async fn run_to_completion(
356 + token_ids: Vec<u32>,
357 + decode_options: TextDecodeOptions,
358 + ) -> crate::output::CollectedTextOutput {
359 + let prompt: Arc<[u32]> = Arc::from([]);
360 + let raw_stream = stream::iter(vec![Ok(GenerateOutput::for_test(
361 + Some(prompt),
362 + token_ids,
363 + Some(FinishReason::Length),
364 + ))]);
365 + let tokenizer: DynTokenizer = Arc::new(ByteTokenizer);
366 + decoded_text_event_stream("test".into(), tokenizer, raw_stream, decode_options, false)
367 + .collect_output()
368 + .await
369 + .unwrap()
370 + }
371 +
372 + /// Convert ASCII string to token IDs (one byte per token).
373 + fn ascii_tokens(s: &str) -> Vec<u32> {
374 + s.bytes().map(u32::from).collect()
375 + }
376 +
377 + fn opts(stop: &[&str], min_tokens: u32) -> TextDecodeOptions {
378 + TextDecodeOptions {
379 + stop_strings: Some(stop.iter().map(|s| s.to_string()).collect()),
380 + min_tokens,
381 + ..Default::default()
382 + }
383 + }
384 +
385 + struct DropRecordingStream {
386 + next: Option<vllm_llm::Result<GenerateOutput>>,
387 + dropped_cause: Arc<Mutex<Option<AbortCause>>>,
388 + }
389 +
390 + impl Stream for DropRecordingStream {
391 + type Item = vllm_llm::Result<GenerateOutput>;
392 +
393 + fn poll_next(mut self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
394 + Poll::Ready(self.next.take())
395 + }
396 + }
397 +
398 + impl Drop for DropRecordingStream {
399 + fn drop(&mut self) {
400 + *self.dropped_cause.lock().unwrap() = Some(AbortCause::current());
401 + }
402 + }
403 +
404 + // --- stop string stream tests ---
405 +
406 + #[tokio::test]
407 + async fn stream_stop_string_sets_task_local_abort_cause_on_raw_stream_drop() {
408 + let prompt: Arc<[u32]> = Arc::from([]);
409 + let dropped_cause = Arc::new(Mutex::new(None));
410 + let raw_stream = DropRecordingStream {
411 + next: Some(Ok(GenerateOutput::for_test(
412 + Some(prompt),
413 + ascii_tokens("hello"),
414 + Some(FinishReason::Length),
415 + ))),
416 + dropped_cause: Arc::clone(&dropped_cause),
417 + };
418 + let tokenizer: DynTokenizer = Arc::new(ByteTokenizer);
419 +
420 + let output = decoded_text_event_stream(
421 + "test".into(),
422 + tokenizer,
423 + raw_stream,
424 + opts(&["ll"], 0),
425 + false,
426 + )
427 + .collect_output()
428 + .await
429 + .unwrap();
430 +
431 + assert_eq!(output.text, "he");
432 + assert!(output.finish_reason.is_stop());
433 + assert_eq!(
434 + *dropped_cause.lock().unwrap(),
435 + Some(AbortCause::StopStringMatched)
436 + );
437 + }
438 +
439 + #[tokio::test]
440 + async fn stream_stop_string_truncates_at_match() {
441 + let output = run_to_completion(ascii_tokens("hello"), opts(&["e"], 0)).await;
442 + assert_eq!(output.text, "h");
443 + assert!(output.finish_reason.is_stop());
444 + }
445 +
446 + #[tokio::test]
447 + async fn stream_stop_string_at_end() {
448 + let output = run_to_completion(ascii_tokens("abcxyz"), opts(&["xyz"], 0)).await;
449 + assert_eq!(output.text, "abc");
450 + assert!(output.finish_reason.is_stop());
451 + }
452 +
453 + #[tokio::test]
454 + async fn stream_stop_string_first_token() {
455 + let output = run_to_completion(ascii_tokens("xhello"), opts(&["x"], 0)).await;
456 + assert_eq!(output.text, "");
457 + assert!(output.finish_reason.is_stop());
458 + }
459 +
460 + #[tokio::test]
461 + async fn stream_stop_string_no_match_runs_to_completion() {
462 + let output = run_to_completion(ascii_tokens("hello"), opts(&["z"], 0)).await;
463 + assert_eq!(output.text, "hello");
464 + assert_eq!(output.finish_reason, FinishReason::Length);
465 + }
466 +
467 + #[tokio::test]
468 + async fn stream_stop_string_multi_char() {
469 + let output = run_to_completion(ascii_tokens("say hello world"), opts(&["lo"], 0)).await;
470 + assert_eq!(output.text, "say hel");
471 + assert!(output.finish_reason.is_stop());
472 + }
473 +
474 + #[tokio::test]
475 + async fn stream_stop_string_first_of_multiple_wins() {
476 + // Both "ll" and "lo" are present; "ll" appears first in the output.
477 + let output = run_to_completion(ascii_tokens("hello"), opts(&["ll", "lo"], 0)).await;
478 + assert_eq!(output.text, "he");
479 + assert!(output.finish_reason.is_stop());
480 + }
481 +
482 + #[tokio::test]
483 + async fn stream_stop_string_include_in_output() {
484 + let output = run_to_completion(
485 + ascii_tokens("hello"),
486 + TextDecodeOptions {
487 + stop_strings: Some(vec!["ll".to_string()]),
488 + include_stop_str_in_output: true,
489 + ..Default::default()
490 + },
491 + )
492 + .await;
493 + assert_eq!(output.text, "hell");
494 + assert!(output.finish_reason.is_stop());
495 + }
496 +
497 + // --- min_tokens + stop string interaction ---
498 +
499 + #[tokio::test]
500 + async fn min_tokens_suppresses_early_stop_string() {
501 + // stop="e", min_tokens=3: the 'e' at token 2 is within the first 3 tokens,
502 + // so it should be skipped. No later 'e' exists, so output runs to completion.
503 + let output = run_to_completion(ascii_tokens("hello"), opts(&["e"], 3)).await;
504 + assert_eq!(output.text, "hello");
505 + assert_eq!(output.finish_reason, FinishReason::Length);
506 + }
507 +
508 + #[tokio::test]
509 + async fn min_tokens_allows_stop_string_after_threshold() {
510 + // stop="e", min_tokens=2: the first 'e' at token 3 is past the threshold.
511 + let output = run_to_completion(ascii_tokens("greet"), opts(&["e"], 2)).await;
512 + assert_eq!(output.text, "gr");
513 + assert!(output.finish_reason.is_stop());
514 + }
515 +
516 + #[tokio::test]
517 + async fn min_tokens_zero_behaves_like_absent() {
518 + let output = run_to_completion(ascii_tokens("hello"), opts(&["e"], 0)).await;
519 + assert_eq!(output.text, "h");
520 + assert!(output.finish_reason.is_stop());
521 + }
522 +
523 + #[test]
524 + fn stop_string_matches_at_end() {
525 + let stops = vec!["wor".to_string()];
526 + // Output: "say wor", last byte 'r' was just added (new_bytes=1)
527 + let result = matches_stop_string(&stops, "say wor", 1);
528 + assert_eq!(result, Some((0, 4)));
529 + }
530 +
531 + #[test]
532 + fn stop_string_no_match() {
533 + let stops = vec!["xyz".to_string()];
534 + let result = matches_stop_string(&stops, "say wor", 1);
535 + assert_eq!(result, None);
536 + }
537 +
538 + #[test]
539 + fn stop_string_matches_first_of_multiple() {
540 + let stops = vec!["wor".to_string(), "say".to_string()];
541 + // "say" appears earlier but "wor" is checked first (index 0)
542 + let result = matches_stop_string(&stops, "say wor", 1);
543 + assert_eq!(result, Some((0, 4)));
544 + }
545 +
546 + #[test]
547 + fn stop_string_matches_second_of_multiple() {
548 + let stops = vec!["xyz".to_string(), "wor".to_string()];
549 + let result = matches_stop_string(&stops, "say wor", 1);
550 + assert_eq!(result, Some((1, 4)));
551 + }
552 +
553 + #[test]
554 + fn stop_string_matches_with_multiple_new_bytes() {
555 + let stops = vec!["wor".to_string()];
556 + // "say wor" where last 3 bytes "wor" were added at once
557 + let result = matches_stop_string(&stops, "say wor", 3);
558 + assert_eq!(result, Some((0, 4)));
559 + }
560 +
561 + #[test]
562 + fn stop_string_matches_at_beginning() {
563 + let stops = vec!["say".to_string()];
564 + let result = matches_stop_string(&stops, "say wor", 7);
565 + assert_eq!(result, Some((0, 0)));
566 + }
567 +
568 + #[test]
569 + fn stop_string_exact_output() {
570 + let stops = vec!["abc".to_string()];
571 + let result = matches_stop_string(&stops, "abc", 3);
572 + assert_eq!(result, Some((0, 0)));
573 + }
574 +
575 + #[test]
576 + fn stop_string_single_char() {
577 + let stops = vec!["!".to_string()];
578 + let result = matches_stop_string(&stops, "hello!", 1);
579 + assert_eq!(result, Some((0, 5)));
580 + }
581 +
582 + #[test]
583 + fn stop_string_not_in_new_bytes_region() {
584 + let stops = vec!["say".to_string()];
585 + // "say" is in the output but before the new byte region.
586 + // new_bytes=1 means only 'r' was added; "say" ended at byte 3,
587 + // but the search window starts at next_off - stop_len = 7+1-1 - 3 = 4.
588 + let result = matches_stop_string(&stops, "say wor", 1);
589 + assert_eq!(result, None);
590 + }
591 +
592 + #[test]
593 + fn stop_string_empty_list() {
594 + let stops: Vec<String> = vec![];
595 + let result = matches_stop_string(&stops, "hello", 1);
596 + assert_eq!(result, None);
597 + }
598 +
599 + #[test]
600 + fn stop_string_multibyte_utf8() {
601 + let stops = vec!["世界".to_string()];
602 + // "你好世界" is 12 bytes: 你(3) + 好(3) + 世(3) + 界(3)
603 + // "世界" starts at byte 6
604 + let result = matches_stop_string(&stops, "你好世界", 3);
605 + assert_eq!(result, Some((0, 6)));
606 + }
607 +}
rust/src/text/src/output/logprobs.rs new
+236
@@ -0,0 +1,236 @@
1 +use itertools::Itertools as _;
2 +use serde::{Deserialize, Serialize};
3 +use vllm_llm::{Logprobs, PositionLogprobs};
4 +use vllm_tokenizer::Tokenizer;
5 +
6 +use crate::error::Error;
7 +
8 +/// One decoded token candidate and its logprob metadata.
9 +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
10 +pub struct DecodedTokenLogprob {
11 + /// Original vocabulary token ID for this candidate.
12 + pub token_id: u32,
13 + /// Best-effort decoded token string for this candidate.
14 + pub token: String,
15 + /// Log probability of this token candidate.
16 + pub logprob: f32,
17 + /// Vocabulary rank of this token candidate.
18 + pub rank: u32,
19 +}
20 +
21 +/// One position's decoded token candidates and their logprobs.
22 +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
23 +pub struct DecodedPositionLogprobs {
24 + /// Candidate tokens for this position.
25 + pub entries: Vec<DecodedTokenLogprob>,
26 +}
27 +
28 +/// Decoded sample logprobs for generated token positions.
29 +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
30 +pub struct DecodedLogprobs {
31 + /// Generated token positions covered by this payload.
32 + pub positions: Vec<DecodedPositionLogprobs>,
33 +}
34 +
35 +/// Decoded prompt logprobs for prompt token positions.
36 +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
37 +pub struct DecodedPromptLogprobs {
38 + /// Original vocabulary token ID for the first prompt token.
39 + pub first_token_id: u32,
40 + /// Best-effort decoded string for the first prompt token.
41 + ///
42 + /// The first prompt token has no left context to score against, so it is
43 + /// stored separately instead of appearing in `scored_positions`.
44 + pub first_token: String,
45 + /// Scored prompt positions after the first prompt token.
46 + ///
47 + /// `scored_positions[i]` corresponds to the prompt token at position `i +
48 + /// 1`.
49 + pub scored_positions: Vec<DecodedPositionLogprobs>,
50 +}
51 +
52 +/// Decode generated-token logprobs from the raw `llm` token-ID shape into the
53 +/// text-layer decoded-token representation.
54 +///
55 +/// Each returned position corresponds to one generated token position from the
56 +/// same `llm` update.
57 +pub(super) fn decode_logprobs<T: Tokenizer + ?Sized>(
58 + tokenizer: &T,
59 + logprobs: &Logprobs,
60 + skip_special_tokens: bool,
61 +) -> Result<DecodedLogprobs, Error> {
62 + Ok(DecodedLogprobs {
63 + positions: logprobs
64 + .positions
65 + .iter()
66 + .map(|position| decode_position_logprobs(tokenizer, position, skip_special_tokens))
67 + .try_collect()?,
68 + })
69 +}
70 +
71 +/// Decode prompt logprobs from the raw `llm` token-ID shape into the text-layer
72 +/// decoded-token representation.
73 +///
74 +/// The returned payload stores the first prompt token separately and decodes
75 +/// the remaining scored prompt positions into `scored_positions`, matching
76 +/// vLLM's prompt-logprobs semantics.
77 +pub(super) fn decode_prompt_logprobs<T: Tokenizer + ?Sized>(
78 + tokenizer: &T,
79 + prompt_token_ids: &[u32],
80 + logprobs: &Logprobs,
81 + skip_special_tokens: bool,
82 +) -> Result<DecodedPromptLogprobs, Error> {
83 + let first_token_id = prompt_token_ids
84 + .first()
85 + .copied()
86 + .expect("prompt logprobs require at least one prompt token");
87 + let first_token = tokenizer.decode(&[first_token_id], skip_special_tokens)?;
88 + let scored_positions = logprobs
89 + .positions
90 + .iter()
91 + .map(|position| decode_position_logprobs(tokenizer, position, skip_special_tokens))
92 + .try_collect()?;
93 +
94 + Ok(DecodedPromptLogprobs {
95 + first_token_id,
96 + first_token,
97 + scored_positions,
98 + })
99 +}
100 +
101 +/// Decode one token position's raw candidate set into decoded token strings
102 +/// plus logprob metadata.
103 +///
104 +/// This decodes every candidate token ID independently through the active text
105 +/// backend.
106 +fn decode_position_logprobs<T: Tokenizer + ?Sized>(
107 + tokenizer: &T,
108 + position: &PositionLogprobs,
109 + skip_special_tokens: bool,
110 +) -> Result<DecodedPositionLogprobs, Error> {
111 + Ok(DecodedPositionLogprobs {
112 + entries: position
113 + .entries
114 + .iter()
115 + .map(|entry| {
116 + tokenizer.decode(&[entry.token_id], skip_special_tokens).map(|token| {
117 + DecodedTokenLogprob {
118 + token_id: entry.token_id,
119 + token,
120 + logprob: entry.logprob,
121 + rank: entry.rank,
122 + }
123 + })
124 + })
125 + .try_collect()?,
126 + })
127 +}
128 +
129 +#[cfg(test)]
130 +mod tests {
131 + use vllm_llm::{Logprobs, PositionLogprobs, TokenLogprob};
132 +
133 + use super::*;
134 +
135 + #[derive(Debug)]
136 + struct ByteTokenizer;
137 +
138 + impl vllm_tokenizer::Tokenizer for ByteTokenizer {
139 + fn encode(
140 + &self,
141 + _text: &str,
142 + _add_special_tokens: bool,
143 + ) -> vllm_tokenizer::Result<Vec<u32>> {
144 + unreachable!()
145 + }
146 +
147 + fn decode(
148 + &self,
149 + token_ids: &[u32],
150 + _skip_special_tokens: bool,
151 + ) -> vllm_tokenizer::Result<String> {
152 + Ok(String::from_utf8_lossy(
153 + &token_ids.iter().map(|token_id| *token_id as u8).collect::<Vec<_>>(),
154 + )
155 + .into_owned())
156 + }
157 +
158 + fn token_to_id(&self, _token: &str) -> Option<u32> {
159 + unreachable!()
160 + }
161 + }
162 +
163 + #[test]
164 + fn decode_logprobs_decodes_every_candidate_token() {
165 + let tokenizer = ByteTokenizer;
166 + let logprobs = Logprobs {
167 + positions: vec![PositionLogprobs {
168 + entries: vec![
169 + TokenLogprob {
170 + token_id: b'a' as u32,
171 + logprob: -0.1,
172 + rank: 3,
173 + },
174 + TokenLogprob {
175 + token_id: b'b' as u32,
176 + logprob: -0.2,
177 + rank: 1,
178 + },
179 + ],
180 + }],
181 + };
182 +
183 + assert_eq!(
184 + decode_logprobs(&tokenizer, &logprobs, false).unwrap(),
185 + DecodedLogprobs {
186 + positions: vec![DecodedPositionLogprobs {
187 + entries: vec![
188 + DecodedTokenLogprob {
189 + token_id: b'a' as u32,
190 + token: "a".to_string(),
191 + logprob: -0.1,
192 + rank: 3,
193 + },
194 + DecodedTokenLogprob {
195 + token_id: b'b' as u32,
196 + token: "b".to_string(),
197 + logprob: -0.2,
198 + rank: 1,
199 + },
200 + ],
201 + }],
202 + }
203 + );
204 + }
205 +
206 + #[test]
207 + fn decode_prompt_logprobs_separates_first_prompt_token() {
208 + let tokenizer = ByteTokenizer;
209 + let logprobs = Logprobs {
210 + positions: vec![PositionLogprobs {
211 + entries: vec![TokenLogprob {
212 + token_id: b'x' as u32,
213 + logprob: -0.4,
214 + rank: 1,
215 + }],
216 + }],
217 + };
218 +
219 + assert_eq!(
220 + decode_prompt_logprobs(&tokenizer, &[b'p' as u32, b'x' as u32], &logprobs, false)
221 + .unwrap(),
222 + DecodedPromptLogprobs {
223 + first_token_id: b'p' as u32,
224 + first_token: "p".to_string(),
225 + scored_positions: vec![DecodedPositionLogprobs {
226 + entries: vec![DecodedTokenLogprob {
227 + token_id: b'x' as u32,
228 + token: "x".to_string(),
229 + logprob: -0.4,
230 + rank: 1,
231 + }],
232 + }],
233 + }
234 + );
235 + }
236 +}
rust/src/text/src/output/mod.rs new
+323
@@ -0,0 +1,323 @@
1 +//! Output processing helpers shared by text and chat layers.
2 +
3 +pub use decoded::{DecodedTextEvent, Finished, TextDecodeOptions, decoded_text_event_stream};
4 +pub use logprobs::{
5 + DecodedLogprobs, DecodedPositionLogprobs, DecodedPromptLogprobs, DecodedTokenLogprob,
6 +};
7 +
8 +mod decoded;
9 +mod logprobs;
10 +
11 +use std::sync::Arc;
12 +
13 +use futures::{StreamExt as _, pin_mut};
14 +
15 +use crate::{Error, FinishReason, Result, TextOutputStream};
16 +
17 +/// Final decoded text plus terminal stream metadata.
18 +#[derive(Debug, Clone, PartialEq)]
19 +pub struct CollectedTextOutput {
20 + pub text: String,
21 + pub prompt_token_ids: Arc<[u32]>,
22 + pub prompt_logprobs: Option<DecodedPromptLogprobs>,
23 + pub logprobs: Option<DecodedLogprobs>,
24 + pub token_ids: Vec<u32>,
25 + pub finish_reason: FinishReason,
26 + /// Connector-specific KV transfer parameters for disaggregated serving.
27 + pub kv_transfer_params: Option<serde_json::Value>,
28 +}
29 +
30 +#[allow(clippy::manual_async_fn, reason = "specify `Send` bound")]
31 +#[easy_ext::ext(TextOutputStreamExt)]
32 +impl<T: TextOutputStream> T {
33 + /// Collect the stream to completion and return the final decoded text plus
34 + /// terminal metadata.
35 + pub fn collect_output(self) -> impl Future<Output = Result<CollectedTextOutput>> + Send {
36 + async move {
37 + let stream = self;
38 + pin_mut!(stream);
39 + let mut prompt_logprobs = None;
40 + let mut prompt_token_ids: Arc<[u32]> = Arc::from([]);
41 + let mut collected: Option<CollectedTextOutput> = None;
42 +
43 + while let Some(event) = stream.next().await.transpose()? {
44 + match event {
45 + DecodedTextEvent::Start {
46 + prompt_logprobs: start_prompt_logprobs,
47 + prompt_token_ids: start_prompt_token_ids,
48 + ..
49 + } => {
50 + prompt_logprobs = start_prompt_logprobs;
51 + prompt_token_ids = start_prompt_token_ids;
52 + }
53 + DecodedTextEvent::TextDelta {
54 + delta,
55 + token_ids: delta_token_ids,
56 + logprobs: mut delta_logprobs,
57 + finished,
58 + } => {
59 + if let Some(c) = collected.as_mut() {
60 + c.text.push_str(&delta);
61 + c.token_ids.extend(delta_token_ids);
62 + if let Some(dlp) = delta_logprobs.as_mut() {
63 + if let Some(lp) = c.logprobs.as_mut() {
64 + lp.positions.extend_from_slice(&dlp.positions);
65 + } else {
66 + c.logprobs = delta_logprobs;
67 + }
68 + }
69 + } else {
70 + collected = Some(CollectedTextOutput {
71 + text: delta,
72 + prompt_token_ids: Arc::clone(&prompt_token_ids),
73 + prompt_logprobs: prompt_logprobs.take(),
74 + logprobs: delta_logprobs,
75 + token_ids: delta_token_ids,
76 + finish_reason: FinishReason::Error,
77 + kv_transfer_params: None,
78 + })
79 + };
80 +
81 + if let Some(finished) = finished {
82 + let mut collected = collected.unwrap();
83 + collected.finish_reason = finished.finish_reason;
84 + collected.kv_transfer_params = finished.kv_transfer_params;
85 + return Ok(collected);
86 + }
87 + }
88 + }
89 + }
90 +
91 + // Note: this is actually unreachable, as the underlying stream always emit an
92 + // error on unexpected close.
93 + Err(Error::StreamClosedBeforeTerminalOutput {
94 + request_id: "unknown".to_string(),
95 + })
96 + }
97 + }
98 +}
99 +
100 +#[cfg(test)]
101 +mod tests {
102 + use futures::stream;
103 + use vllm_llm::FinishReason;
104 +
105 + use super::*;
106 +
107 + #[tokio::test]
108 + async fn collect_output_retains_prompt_and_sample_logprobs() {
109 + let stream = stream::iter(vec![
110 + Ok(DecodedTextEvent::Start {
111 + prompt_token_ids: vec![10, 11].into(),
112 + prompt_logprobs: Some(DecodedPromptLogprobs {
113 + first_token_id: 0,
114 + first_token: "o".to_string(),
115 + scored_positions: vec![DecodedPositionLogprobs {
116 + entries: vec![DecodedTokenLogprob {
117 + token_id: 0,
118 + token: "p".to_string(),
119 + logprob: -0.1,
120 + rank: 1,
121 + }],
122 + }],
123 + }),
124 + }),
125 + Ok(DecodedTextEvent::TextDelta {
126 + delta: "bc".to_string(),
127 + token_ids: vec![1, 2],
128 + logprobs: Some(DecodedLogprobs {
129 + positions: vec![
130 + DecodedPositionLogprobs {
131 + entries: vec![DecodedTokenLogprob {
132 + token_id: 0,
133 + token: "a".to_string(),
134 + logprob: -0.2,
135 + rank: 1,
136 + }],
137 + },
138 + DecodedPositionLogprobs {
139 + entries: vec![DecodedTokenLogprob {
140 + token_id: 0,
141 + token: "bc".to_string(),
142 + logprob: -0.3,
143 + rank: 1,
144 + }],
145 + },
146 + ],
147 + }),
148 + finished: Some(Finished {
149 + prompt_token_count: 2,
150 + output_token_count: 2,
151 + finish_reason: FinishReason::stop_eos(),
152 + kv_transfer_params: None,
153 + }),
154 + }),
155 + ]);
156 +
157 + let collected = stream.collect_output().await.unwrap();
158 + assert_eq!(collected.text, "bc");
159 + assert_eq!(
160 + collected.prompt_logprobs,
161 + Some(DecodedPromptLogprobs {
162 + first_token_id: 0,
163 + first_token: "o".to_string(),
164 + scored_positions: vec![DecodedPositionLogprobs {
165 + entries: vec![DecodedTokenLogprob {
166 + token_id: 0,
167 + token: "p".to_string(),
168 + logprob: -0.1,
169 + rank: 1,
170 + }],
171 + }],
172 + })
173 + );
174 + assert_eq!(
175 + collected.logprobs,
176 + Some(DecodedLogprobs {
177 + positions: vec![
178 + DecodedPositionLogprobs {
179 + entries: vec![DecodedTokenLogprob {
180 + token_id: 0,
181 + token: "a".to_string(),
182 + logprob: -0.2,
183 + rank: 1,
184 + }],
185 + },
186 + DecodedPositionLogprobs {
187 + entries: vec![DecodedTokenLogprob {
188 + token_id: 0,
189 + token: "bc".to_string(),
190 + logprob: -0.3,
191 + rank: 1,
192 + }],
193 + },
194 + ],
195 + })
196 + );
197 + }
198 +
199 + #[tokio::test]
200 + async fn collect_output_accumulates_intermediate_deltas() {
201 + let stream = stream::iter(vec![
202 + Ok(DecodedTextEvent::Start {
203 + prompt_token_ids: vec![10, 11].into(),
204 + prompt_logprobs: None,
205 + }),
206 + Ok(DecodedTextEvent::TextDelta {
207 + delta: "he".to_string(),
208 + token_ids: vec![1, 2],
209 + logprobs: Some(DecodedLogprobs {
210 + positions: vec![
211 + DecodedPositionLogprobs {
212 + entries: vec![DecodedTokenLogprob {
213 + token_id: 0,
214 + token: "h".to_string(),
215 + logprob: -0.1,
216 + rank: 1,
217 + }],
218 + },
219 + DecodedPositionLogprobs {
220 + entries: vec![DecodedTokenLogprob {
221 + token_id: 0,
222 + token: "e".to_string(),
223 + logprob: -0.2,
224 + rank: 1,
225 + }],
226 + },
227 + ],
228 + }),
229 + finished: None,
230 + }),
231 + Ok(DecodedTextEvent::TextDelta {
232 + delta: "llo".to_string(),
233 + token_ids: vec![3, 4, 5],
234 + logprobs: Some(DecodedLogprobs {
235 + positions: vec![
236 + DecodedPositionLogprobs {
237 + entries: vec![DecodedTokenLogprob {
238 + token_id: 0,
239 + token: "l".to_string(),
240 + logprob: -0.3,
241 + rank: 1,
242 + }],
243 + },
244 + DecodedPositionLogprobs {
245 + entries: vec![DecodedTokenLogprob {
246 + token_id: 0,
247 + token: "l".to_string(),
248 + logprob: -0.4,
249 + rank: 1,
250 + }],
251 + },
252 + DecodedPositionLogprobs {
253 + entries: vec![DecodedTokenLogprob {
254 + token_id: 0,
255 + token: "o".to_string(),
256 + logprob: -0.5,
257 + rank: 1,
258 + }],
259 + },
260 + ],
261 + }),
262 + finished: Some(Finished {
263 + prompt_token_count: 2,
264 + output_token_count: 5,
265 + finish_reason: FinishReason::stop_eos(),
266 + kv_transfer_params: None,
267 + }),
268 + }),
269 + ]);
270 +
271 + let collected = stream.collect_output().await.unwrap();
272 + assert_eq!(collected.text, "hello");
273 + assert_eq!(collected.prompt_logprobs, None);
274 + assert_eq!(collected.token_ids, vec![1, 2, 3, 4, 5]);
275 + assert_eq!(
276 + collected.logprobs,
277 + Some(DecodedLogprobs {
278 + positions: vec![
279 + DecodedPositionLogprobs {
280 + entries: vec![DecodedTokenLogprob {
281 + token_id: 0,
282 + token: "h".to_string(),
283 + logprob: -0.1,
284 + rank: 1,
285 + }],
286 + },
287 + DecodedPositionLogprobs {
288 + entries: vec![DecodedTokenLogprob {
289 + token_id: 0,
290 + token: "e".to_string(),
291 + logprob: -0.2,
292 + rank: 1,
293 + }],
294 + },
295 + DecodedPositionLogprobs {
296 + entries: vec![DecodedTokenLogprob {
297 + token_id: 0,
298 + token: "l".to_string(),
299 + logprob: -0.3,
300 + rank: 1,
301 + }],
302 + },
303 + DecodedPositionLogprobs {
304 + entries: vec![DecodedTokenLogprob {
305 + token_id: 0,
306 + token: "l".to_string(),
307 + logprob: -0.4,
308 + rank: 1,
309 + }],
310 + },
311 + DecodedPositionLogprobs {
312 + entries: vec![DecodedTokenLogprob {
313 + token_id: 0,
314 + token: "o".to_string(),
315 + logprob: -0.5,
316 + rank: 1,
317 + }],
318 + },
319 + ],
320 + })
321 + );
322 + }
323 +}
rust/src/text/src/request.rs new
+197
@@ -0,0 +1,197 @@
1 +use std::collections::HashMap;
2 +
3 +use enum_as_inner::EnumAsInner;
4 +use serde::{Deserialize, Serialize};
5 +use serde_json::Value;
6 +use vllm_engine_core_client::protocol::StructuredOutputsParams;
7 +use vllm_engine_core_client::protocol::multimodal::MmFeatures;
8 +
9 +use crate::error::{Error, Result};
10 +use crate::output::TextDecodeOptions;
11 +
12 +/// One raw text-generation prompt.
13 +///
14 +/// This supports either ordinary text that still needs tokenization or
15 +/// already-tokenized prompt IDs that should bypass tokenizer work entirely.
16 +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, EnumAsInner)]
17 +#[serde(untagged)]
18 +pub enum Prompt {
19 + /// Untokenized prompt text that still needs tokenizer work before
20 + /// generation.
21 + Text(String),
22 + /// Pre-tokenized prompt IDs that should be forwarded southbound without
23 + /// re-encoding.
24 + TokenIds(Vec<u32>),
25 +}
26 +
27 +impl Default for Prompt {
28 + fn default() -> Self {
29 + Self::Text(String::new()) // placeholder
30 + }
31 +}
32 +
33 +/// User-facing sampling parameters accepted by `vllm-text`.
34 +///
35 +/// This intentionally keeps only the subset that the current Rust text layer
36 +/// supports as northbound request semantics. Engine-core-specific normalized
37 +/// fields are derived later during lowering.
38 +///
39 +/// Original Python definition:
40 +/// <https://github.com/vllm-project/vllm/blob/f22d6e026798a74e6542a52ef776c054f2de572a/vllm/sampling_params.py#L155-L291>
41 +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
42 +#[serde(default)]
43 +pub struct SamplingParams {
44 + /// Controls randomness. Lower values are more deterministic; zero means
45 + /// greedy sampling. `None` means no explicit user override.
46 + pub temperature: Option<f32>,
47 + /// Cumulative probability threshold for nucleus sampling.
48 + pub top_p: Option<f32>,
49 + /// Maximum number of top tokens to consider. `Some(0)` means all tokens.
50 + pub top_k: Option<u32>,
51 + /// Random seed used by the sampler when present.
52 + pub seed: Option<i64>,
53 + /// Maximum number of tokens to generate. `None` means no explicit user
54 + /// override.
55 + pub max_tokens: Option<u32>,
56 + /// Minimum number of tokens to generate before EOS or stop-token handling.
57 + pub min_tokens: Option<u32>,
58 + /// Number of log probabilities to return per generated token.
59 + ///
60 + /// `None` disables sample logprobs. `-1` requests the full vocabulary.
61 + pub logprobs: Option<i32>,
62 + /// Number of log probabilities to return per prompt token.
63 + ///
64 + /// `None` disables prompt logprobs. `-1` requests the full vocabulary.
65 + pub prompt_logprobs: Option<i32>,
66 + /// Minimum probability threshold for token sampling. `None` means no
67 + /// explicit user override.
68 + pub min_p: Option<f32>,
69 + /// Frequency penalty applied by the sampler. `None` means no explicit user
70 + /// override.
71 + pub frequency_penalty: Option<f32>,
72 + /// Presence penalty applied by the sampler. `None` means no explicit user
73 + /// override.
74 + pub presence_penalty: Option<f32>,
75 + /// Repetition penalty applied by the sampler. `None` means no explicit user
76 + /// override.
77 + pub repetition_penalty: Option<f32>,
78 + /// Explicit stop token IDs provided by the caller. `None` means no explicit
79 + /// user override.
80 + pub stop_token_ids: Option<Vec<u32>>,
81 + /// If true, do not stop on the model's primary EOS token.
82 + pub ignore_eos: bool,
83 + /// Modify the likelihood of specified tokens appearing in the completion.
84 + /// Keys are token IDs.
85 + pub logit_bias: Option<HashMap<u32, f32>>,
86 + /// Restrict output to these token IDs only.
87 + pub allowed_token_ids: Option<Vec<u32>>,
88 + /// Words to avoid during generation (tokenized to IDs during lowering).
89 + pub bad_words: Option<Vec<String>>,
90 + /// Specific token IDs for which log probabilities should be returned at
91 + /// each position.
92 + ///
93 + /// When set, the engine returns logprobs for exactly these tokens in
94 + /// addition to the sampled/scored token. Mutually exclusive with
95 + /// `logprobs` in practice.
96 + pub logprob_token_ids: Option<Vec<u32>>,
97 + /// Parameters for configuring structured outputs (guided decoding).
98 + pub structured_outputs: Option<StructuredOutputsParams>,
99 + /// If true, bypass reads from the prefix cache for this request (the prompt
100 + /// will not reuse cached KV blocks from earlier requests, though newly
101 + /// computed blocks may still populate the cache). `None` defers to
102 + /// engine-core defaults.
103 + pub skip_reading_prefix_cache: Option<bool>,
104 + /// Additional request parameters for custom extensions.
105 + pub vllm_xargs: Option<HashMap<String, Value>>,
106 +}
107 +
108 +#[allow(clippy::derivable_impls)] // more explicit
109 +impl Default for SamplingParams {
110 + fn default() -> Self {
111 + Self {
112 + temperature: None,
113 + top_p: None,
114 + top_k: None,
115 + seed: None,
116 + max_tokens: None,
117 + min_tokens: None,
118 + logprobs: None,
119 + prompt_logprobs: None,
120 + min_p: None,
121 + frequency_penalty: None,
122 + presence_penalty: None,
123 + repetition_penalty: None,
124 + stop_token_ids: None,
125 + ignore_eos: false,
126 + logit_bias: None,
127 + allowed_token_ids: None,
128 + bad_words: None,
129 + logprob_token_ids: None,
130 + structured_outputs: None,
131 + skip_reading_prefix_cache: None,
132 + vllm_xargs: None,
133 + }
134 + }
135 +}
136 +
137 +/// One raw text-generation request ready to be tokenized or sent directly to
138 +/// the engine.
139 +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
140 +pub struct TextRequest {
141 + /// Stable caller-supplied request ID.
142 + pub request_id: String,
143 + /// Prompt text or prompt token IDs for this request.
144 + pub prompt: Prompt,
145 + /// Multimodal features prepared by a higher-level frontend. Raw text
146 + /// requests keep this empty; multimodal chat uses it with pre-tokenized
147 + /// prompt IDs.
148 + pub mm_features: Option<MmFeatures>,
149 + /// User-facing sampling parameters accepted by `vllm-text`.
150 + pub sampling_params: SamplingParams,
151 + /// Incremental detokenization options for the response path.
152 + pub decode_options: TextDecodeOptions,
153 + /// Whether to emit intermediate northbound deltas before the terminal
154 + /// result.
155 + ///
156 + /// If `false`, callers only observe the terminal accumulated output. If
157 + /// `true`, callers may receive zero or more incremental decoded updates
158 + /// before the final terminal event.
159 + pub intermediate: bool,
160 + /// Request scheduling priority (lower means earlier handling; default 0).
161 + pub priority: i32,
162 + /// Salt for prefix cache isolation in multi-user environments.
163 + pub cache_salt: Option<String>,
164 + /// Whether to add special tokens (e.g. BOS) during prompt tokenization.
165 + pub add_special_tokens: bool,
166 + /// Override data parallel rank.
167 + #[serde(default)]
168 + pub data_parallel_rank: Option<u32>,
169 +}
170 +
171 +impl TextRequest {
172 + /// Return one minimal valid request fixture for tests.
173 + pub fn for_test() -> Self {
174 + Self {
175 + request_id: "test-request".to_string(),
176 + prompt: Prompt::Text("test".to_string()),
177 + mm_features: None,
178 + sampling_params: SamplingParams::default(),
179 + decode_options: TextDecodeOptions::default(),
180 + intermediate: true,
181 + priority: 0,
182 + cache_salt: None,
183 + add_special_tokens: false,
184 + data_parallel_rank: None,
185 + }
186 + }
187 +
188 + /// Validate the minimum invariants before tokenization or request lowering.
189 + pub fn validate(&self) -> Result<()> {
190 + if matches!(&self.prompt, Prompt::TokenIds(ids) if ids.is_empty()) {
191 + return Err(Error::EmptyPromptTokenIds {
192 + request_id: self.request_id.clone(),
193 + });
194 + }
195 + Ok(())
196 + }
197 +}
rust/src/tokenizer/Cargo.toml new
+35
@@ -0,0 +1,35 @@
1 +[package]
2 +name = "vllm-tokenizer"
3 +version.workspace = true
4 +edition.workspace = true
5 +license.workspace = true
6 +
7 +[dependencies]
8 +base64.workspace = true
9 +fastokens.workspace = true
10 +riptoken.workspace = true
11 +rustc-hash.workspace = true
12 +serde.workspace = true
13 +serde_json.workspace = true
14 +tekken.workspace = true
15 +thiserror.workspace = true
16 +thiserror-ext.workspace = true
17 +tiktoken-rs.workspace = true
18 +tokenizers.workspace = true
19 +tracing.workspace = true
20 +
21 +[dev-dependencies]
22 +criterion.workspace = true
23 +hf-hub.workspace = true
24 +tempfile.workspace = true
25 +
26 +[[bench]]
27 +name = "hf"
28 +harness = false
29 +
30 +[[bench]]
31 +name = "tiktoken"
32 +harness = false
33 +
34 +[lints]
35 +workspace = true
rust/src/tokenizer/benches/hf.rs new
+118
@@ -0,0 +1,118 @@
1 +use criterion::{Criterion, Throughput, black_box, criterion_group, criterion_main};
2 +use hf_hub::api::sync::ApiBuilder;
3 +use vllm_tokenizer::{HuggingFaceTokenizer, Tokenizer};
4 +
5 +const MODEL_ID: &str = "Qwen/Qwen3.5-0.8B";
6 +const SAMPLE_TEXT: &str = "\
7 +<|im_start|>system
8 +You are Qwen3.5, a helpful assistant.
9 +<|im_end|>
10 +<|im_start|>user
11 +请用中英混合总结以下需求，并给出一个简短的 JSON 示例。
12 +The service should stop cleanly at EOS, avoid leaking the next template turn, and keep decode latency low.
13 +Input: 4 concurrent requests, 10240 prompt tokens, 16 generated tokens.
14 +<|im_end|>
15 +<|im_start|>assistant
16 +";
17 +
18 +struct BenchFixture {
19 + fastokens: HuggingFaceTokenizer,
20 + hf: HuggingFaceTokenizer,
21 + text: String,
22 + token_ids: Vec<u32>,
23 +}
24 +
25 +impl BenchFixture {
26 + fn load() -> Self {
27 + let path = tokenizer_json();
28 + let fastokens =
29 + HuggingFaceTokenizer::new_fastokens(&path).expect("load fastokens tokenizer");
30 + let hf = HuggingFaceTokenizer::new_hf(&path).expect("load huggingface tokenizer");
31 +
32 + let text = SAMPLE_TEXT.repeat(32);
33 + let hf_token_ids =
34 + hf.encode(text.as_str(), false).expect("encode sample text with hf tokenizer");
35 + let fastokens_token_ids = fastokens
36 + .encode(text.as_str(), false)
37 + .expect("encode sample text with fastokens");
38 + assert_eq!(fastokens_token_ids, hf_token_ids);
39 +
40 + let hf_decoded = hf
41 + .decode(hf_token_ids.as_slice(), false)
42 + .expect("decode sample token ids with hf tokenizer");
43 + let fastokens_decoded = fastokens
44 + .decode(hf_token_ids.as_slice(), false)
45 + .expect("decode sample token ids with fastokens");
46 + assert_eq!(fastokens_decoded, hf_decoded);
47 +
48 + Self {
49 + fastokens,
50 + hf,
51 + text,
52 + token_ids: hf_token_ids,
53 + }
54 + }
55 +}
56 +
57 +fn tokenizer_json() -> std::path::PathBuf {
58 + ApiBuilder::from_env()
59 + .with_progress(false)
60 + .build()
61 + .expect("build hf-hub api")
62 + .model(MODEL_ID.to_string())
63 + .get("tokenizer.json")
64 + .expect("fetch tokenizer.json from hf-hub")
65 +}
66 +
67 +fn bench_encode(c: &mut Criterion) {
68 + let fixture = BenchFixture::load();
69 + let mut group = c.benchmark_group("tokenizer_encode");
70 + group.throughput(Throughput::Bytes(fixture.text.len() as u64));
71 +
72 + group.bench_function("fastokens", |b| {
73 + b.iter(|| {
74 + fixture
75 + .fastokens
76 + .encode(black_box(fixture.text.as_str()), black_box(false))
77 + .expect("encode sample text with fastokens")
78 + })
79 + });
80 + group.bench_function("hf_tokenizers", |b| {
81 + b.iter(|| {
82 + fixture
83 + .hf
84 + .encode(black_box(fixture.text.as_str()), black_box(false))
85 + .expect("encode sample text with hf tokenizer")
86 + })
87 + });
88 +
89 + group.finish();
90 +}
91 +
92 +fn bench_decode(c: &mut Criterion) {
93 + let fixture = BenchFixture::load();
94 + let mut group = c.benchmark_group("tokenizer_decode");
95 + group.throughput(Throughput::Elements(fixture.token_ids.len() as u64));
96 +
97 + group.bench_function("fastokens", |b| {
98 + b.iter(|| {
99 + fixture
100 + .fastokens
101 + .decode(black_box(fixture.token_ids.as_slice()), black_box(false))
102 + .expect("decode sample token ids with fastokens")
103 + })
104 + });
105 + group.bench_function("hf_tokenizers", |b| {
106 + b.iter(|| {
107 + fixture
108 + .hf
109 + .decode(black_box(fixture.token_ids.as_slice()), black_box(false))
110 + .expect("decode sample token ids with hf tokenizer")
111 + })
112 + });
113 +
114 + group.finish();
115 +}
116 +
117 +criterion_group!(benches, bench_encode, bench_decode);
118 +criterion_main!(benches);
rust/src/tokenizer/benches/tiktoken.rs new
+117
@@ -0,0 +1,117 @@
1 +use criterion::{Criterion, Throughput, black_box, criterion_group, criterion_main};
2 +use hf_hub::api::sync::ApiBuilder;
3 +use vllm_tokenizer::{TiktokenTokenizer, Tokenizer};
4 +
5 +const MODEL_ID: &str = "moonshotai/Kimi-K2.5";
6 +const SAMPLE_TEXT: &str = "\
7 +<think>
8 +I'm sure it's fine, but I can't say I'd trust that it's what we'd ship.
9 +</think>
10 +请用中英混合总结以下需求，并保留 tool-call marker:
11 +<|tool_calls_section_begin|>{\"name\":\"summarize\",\"arguments\":{\"style\":\"brief\"}}<|tool_calls_section_end|>
12 +The service should stop cleanly at EOS, avoid leaking the next template turn, and keep decode latency low.
13 +";
14 +
15 +struct BenchFixture {
16 + riptoken: TiktokenTokenizer,
17 + tiktoken_rs: TiktokenTokenizer,
18 + text: String,
19 + token_ids: Vec<u32>,
20 +}
21 +
22 +impl BenchFixture {
23 + fn load() -> Self {
24 + let path = tiktoken_model();
25 + let riptoken = TiktokenTokenizer::new_riptoken(&path).expect("load riptoken tokenizer");
26 + let tiktoken_rs =
27 + TiktokenTokenizer::new_tiktoken_rs(&path).expect("load tiktoken-rs tokenizer");
28 +
29 + let text = SAMPLE_TEXT.repeat(32);
30 + let riptoken_token_ids =
31 + riptoken.encode(text.as_str(), false).expect("encode sample text with riptoken");
32 + let tiktoken_rs_token_ids = tiktoken_rs
33 + .encode(text.as_str(), false)
34 + .expect("encode sample text with tiktoken-rs");
35 + assert_eq!(riptoken_token_ids, tiktoken_rs_token_ids);
36 +
37 + let riptoken_decoded = riptoken
38 + .decode(riptoken_token_ids.as_slice(), false)
39 + .expect("decode sample token ids with riptoken");
40 + let tiktoken_rs_decoded = tiktoken_rs
41 + .decode(riptoken_token_ids.as_slice(), false)
42 + .expect("decode sample token ids with tiktoken-rs");
43 + assert_eq!(riptoken_decoded, tiktoken_rs_decoded);
44 +
45 + Self {
46 + riptoken,
47 + tiktoken_rs,
48 + text,
49 + token_ids: riptoken_token_ids,
50 + }
51 + }
52 +}
53 +
54 +fn tiktoken_model() -> std::path::PathBuf {
55 + let repo = ApiBuilder::from_env()
56 + .with_progress(false)
57 + .build()
58 + .expect("build hf-hub api")
59 + .model(MODEL_ID.to_string());
60 + repo.get("config.json").expect("fetch config.json from hf-hub");
61 + repo.get("tokenizer_config.json")
62 + .expect("fetch tokenizer_config.json from hf-hub");
63 + repo.get("tiktoken.model").expect("fetch tiktoken.model from hf-hub")
64 +}
65 +
66 +fn bench_encode(c: &mut Criterion) {
67 + let fixture = BenchFixture::load();
68 + let mut group = c.benchmark_group("tiktoken_encode");
69 + group.throughput(Throughput::Bytes(fixture.text.len() as u64));
70 +
71 + group.bench_function("riptoken", |b| {
72 + b.iter(|| {
73 + fixture
74 + .riptoken
75 + .encode(black_box(fixture.text.as_str()), black_box(false))
76 + .expect("encode sample text with riptoken")
77 + })
78 + });
79 + group.bench_function("tiktoken_rs", |b| {
80 + b.iter(|| {
81 + fixture
82 + .tiktoken_rs
83 + .encode(black_box(fixture.text.as_str()), black_box(false))
84 + .expect("encode sample text with tiktoken-rs")
85 + })
86 + });
87 +
88 + group.finish();
89 +}
90 +
91 +fn bench_decode(c: &mut Criterion) {
92 + let fixture = BenchFixture::load();
93 + let mut group = c.benchmark_group("tiktoken_decode");
94 + group.throughput(Throughput::Elements(fixture.token_ids.len() as u64));
95 +
96 + group.bench_function("riptoken", |b| {
97 + b.iter(|| {
98 + fixture
99 + .riptoken
100 + .decode(black_box(fixture.token_ids.as_slice()), black_box(false))
101 + .expect("decode sample token ids with riptoken")
102 + })
103 + });
104 + group.bench_function("tiktoken_rs", |b| {
105 + b.iter(|| {
106 + fixture
107 + .tiktoken_rs
108 + .decode(black_box(fixture.token_ids.as_slice()), black_box(false))
109 + .expect("decode sample token ids with tiktoken-rs")
110 + })
111 + });
112 +
113 + group.finish();
114 +}
115 +
116 +criterion_group!(benches, bench_encode, bench_decode);
117 +criterion_main!(benches);
rust/src/tokenizer/src/byte_level_decode.rs new
+119
@@ -0,0 +1,119 @@
1 +//! Fast GPT-2 byte-level detokenization that writes into a single `Vec<u8>`,
2 +//! avoiding the `Vec<String>` / `String::join` assembly in fastokens' generic
3 +//! `Decoder::decode` pipeline.
4 +
5 +/// Reverse GPT-2 byte-to-unicode mapping: codepoint → original byte. The GPT-2
6 +/// table only emits codepoints in U+0000..U+0143, so a flat array suffices.
7 +const CHAR_TO_BYTE: [u8; 324] = build_char_to_byte();
8 +
9 +const fn is_nice(b: u8) -> bool {
10 + (b >= b'!' && b <= b'~') || (b >= 0xA1 && b <= 0xAC) || b >= 0xAE
11 +}
12 +
13 +const fn build_char_to_byte() -> [u8; 324] {
14 + let mut table = [0u8; 324];
15 + let mut b: u16 = 0;
16 + while b < 256 {
17 + let cp = if is_nice(b as u8) {
18 + b as u32
19 + } else {
20 + 256 + nice_offset(b as u8)
21 + };
22 + table[cp as usize] = b as u8;
23 + b += 1;
24 + }
25 + table
26 +}
27 +
28 +const fn nice_offset(b: u8) -> u32 {
29 + let mut i: u16 = 0;
30 + let mut n: u32 = 0;
31 + while i < b as u16 {
32 + if !is_nice(i as u8) {
33 + n += 1;
34 + }
35 + i += 1;
36 + }
37 + n
38 +}
39 +
40 +/// Decode byte-level encoded token strings into a single UTF-8 string,
41 +/// matching `fastokens::decoders::ByteLevelDecoder`.
42 +pub fn decode_byte_level<'a, I: IntoIterator<Item = &'a str>>(tokens: I) -> String {
43 + let iter = tokens.into_iter();
44 + let (lower, _) = iter.size_hint();
45 + let mut bytes: Vec<u8> = Vec::with_capacity(lower.saturating_mul(4));
46 + for token in iter {
47 + for c in token.chars() {
48 + let cp = c as usize;
49 + if cp < CHAR_TO_BYTE.len() {
50 + bytes.push(CHAR_TO_BYTE[cp]);
51 + } else {
52 + // Non-GPT2 codepoints (e.g. DeepSeek's U+FF5C, U+2581) pass through.
53 + let mut buf = [0u8; 4];
54 + bytes.extend_from_slice(c.encode_utf8(&mut buf).as_bytes());
55 + }
56 + }
57 + }
58 + String::from_utf8(bytes).unwrap_or_else(|e| String::from_utf8_lossy(e.as_bytes()).into_owned())
59 +}
60 +
61 +#[cfg(test)]
62 +mod tests {
63 + use super::*;
64 +
65 + fn build_byte_to_char_ref() -> [char; 256] {
66 + let mut table = ['\0'; 256];
67 + let mut next: u32 = 256;
68 + for b in 0..=255u8 {
69 + let cp = if is_nice(b) {
70 + b as u32
71 + } else {
72 + let cp = next;
73 + next += 1;
74 + cp
75 + };
76 + table[b as usize] = char::from_u32(cp).unwrap();
77 + }
78 + table
79 + }
80 +
81 + #[test]
82 + fn char_to_byte_roundtrips_every_byte() {
83 + let byte_to_char = build_byte_to_char_ref();
84 + for b in 0..=255u8 {
85 + let cp = byte_to_char[b as usize] as usize;
86 + assert!(cp < CHAR_TO_BYTE.len());
87 + assert_eq!(CHAR_TO_BYTE[cp], b, "mismatch for byte {b:#x}");
88 + }
89 + }
90 +
91 + #[test]
92 + fn decode_ascii() {
93 + assert_eq!(decode_byte_level(["Hello"]), "Hello");
94 + }
95 +
96 + #[test]
97 + fn decode_space_marker() {
98 + // GPT-2 maps 0x20 → Ġ (U+0120).
99 + assert_eq!(
100 + decode_byte_level(["\u{120}Hello", "\u{120}world"]),
101 + " Hello world",
102 + );
103 + }
104 +
105 + #[test]
106 + fn decode_multibyte_euro() {
107 + // € → 0xE2 0x82 0xAC, each mapped to a specific GPT-2 char.
108 + let byte_to_char = build_byte_to_char_ref();
109 + let encoded: String =
110 + [0xE2u8, 0x82, 0xAC].iter().map(|&b| byte_to_char[b as usize]).collect();
111 + assert_eq!(decode_byte_level([encoded.as_str()]), "€");
112 + }
113 +
114 + #[test]
115 + fn decode_preserves_non_gpt2_chars() {
116 + let tok = "<\u{FF5C}begin\u{2581}of\u{2581}sentence\u{FF5C}>";
117 + assert_eq!(decode_byte_level([tok]), "<｜begin▁of▁sentence｜>");
118 + }
119 +}
rust/src/tokenizer/src/error.rs new
+9
@@ -0,0 +1,9 @@
1 +use thiserror::Error;
2 +use thiserror_ext::Macro;
3 +
4 +pub type Result<T> = std::result::Result<T, TokenizerError>;
5 +
6 +#[derive(Debug, Error, Macro)]
7 +#[thiserror_ext(macro(path = "crate::error"))]
8 +#[error("tokenizer error: {0}")]
9 +pub struct TokenizerError(#[message] pub String);
rust/src/tokenizer/src/hf.rs new
+344
@@ -0,0 +1,344 @@
1 +use std::path::Path;
2 +use std::sync::Arc;
3 +
4 +use fastokens::Tokenizer as FastokensTokenizer;
5 +use fastokens::decoders::Decoder as FastokensDecoder;
6 +use thiserror_ext::AsReport as _;
7 +use tokenizers::Tokenizer as HfTokenizer;
8 +use tracing::{info, warn};
9 +
10 +use crate::byte_level_decode::decode_byte_level;
11 +use crate::{Result, Tokenizer};
12 +
13 +enum Backend {
14 + Hf(Box<HfTokenizer>),
15 + Fastokens(Box<FastokensTokenizer>),
16 + /// Fastokens tokenizer whose decoder is pure GPT-2 byte-level, so we can
17 + /// bypass `Decoder::decode`'s `Vec<String>`/`join("")` assembly.
18 + FastokensByteLevel(Box<FastokensTokenizer>),
19 +}
20 +
21 +/// True if `dec` is effectively a single `ByteLevel` stage — one `ByteLevel`
22 +/// leaf in a tree of `Sequence`s (fastokens represents `Fuse` as an empty
23 +/// `Sequence`, which is a no-op for our purposes).
24 +fn is_byte_level_only(dec: &FastokensDecoder) -> bool {
25 + fn count_byte_level(dec: &FastokensDecoder) -> usize {
26 + match dec {
27 + FastokensDecoder::ByteLevel(_) => 1,
28 + FastokensDecoder::Sequence(steps) => steps.iter().map(count_byte_level).sum(),
29 + }
30 + }
31 + count_byte_level(dec) == 1
32 +}
33 +
34 +fn decode_fastokens_byte_level(
35 + t: &FastokensTokenizer,
36 + token_ids: &[u32],
37 + skip_special_tokens: bool,
38 +) -> Result<String> {
39 + let tokens: Vec<&str> = token_ids
40 + .iter()
41 + .filter(|&&id| !(skip_special_tokens && t.is_special_token(id)))
42 + .map(|&id| {
43 + t.id_to_token(id)
44 + .ok_or_else(|| tokenizer_error!("decoding failed: unknown token ID: {id}"))
45 + })
46 + .collect::<Result<_>>()?;
47 + Ok(decode_byte_level(tokens))
48 +}
49 +
50 +/// Tokenizer from `tokenizer.json` in HuggingFace format.
51 +///
52 +/// This tries to load with `fastokens` first for better performance, then falls
53 +/// back to HuggingFace's `tokenizers` if the former fails (e.g. due to
54 +/// unsupported tokenizer features or file formats).
55 +pub struct HuggingFaceTokenizer {
56 + backend: Backend,
57 + special_token_ids: Arc<[u32]>,
58 +}
59 +
60 +impl HuggingFaceTokenizer {
61 + fn from_hf_backend(tokenizer: HfTokenizer) -> Self {
62 + let special_token_ids = {
63 + let mut ids: Vec<u32> = tokenizer
64 + .get_added_tokens_decoder()
65 + .iter()
66 + .filter(|(_id, token)| token.special)
67 + .map(|(id, _token)| *id)
68 + .collect();
69 + ids.sort_unstable();
70 + ids.dedup();
71 + Arc::from(ids)
72 + };
73 + Self {
74 + backend: Backend::Hf(Box::new(tokenizer)),
75 + special_token_ids,
76 + }
77 + }
78 +
79 + fn from_fastokens_backend(tokenizer: FastokensTokenizer) -> Self {
80 + let special_token_ids = {
81 + let mut ids: Vec<u32> = tokenizer
82 + .added_tokens()
83 + .into_iter()
84 + .flat_map(|added_tokens| added_tokens.iter())
85 + .filter(|token| token.special)
86 + .map(|token| token.id)
87 + .collect();
88 + ids.sort_unstable();
89 + ids.dedup();
90 + Arc::from(ids)
91 + };
92 + let byte_level = tokenizer.decoder().is_some_and(is_byte_level_only);
93 + let backend = if byte_level {
94 + Backend::FastokensByteLevel(Box::new(tokenizer))
95 + } else {
96 + Backend::Fastokens(Box::new(tokenizer))
97 + };
98 + Self {
99 + backend,
100 + special_token_ids,
101 + }
102 + }
103 +
104 + /// Load from `tokenizer.json` with `fastokens`.
105 + pub fn new_fastokens(path: &Path) -> Result<Self> {
106 + info!(path = %path.display(), "loading tokenizer with fastokens");
107 + let t = FastokensTokenizer::from_file(path)
108 + .map_err(|error| tokenizer_error!("failed to load tokenizer: {}", error.as_report()))?;
109 + Ok(Self::from_fastokens_backend(t))
110 + }
111 +
112 + /// Load from `tokenizer.json` with Hugging Face `tokenizers`.
113 + pub fn new_hf(path: &Path) -> Result<Self> {
114 + info!(path = %path.display(), "loading tokenizer with huggingface tokenizers");
115 + let t = HfTokenizer::from_file(path)
116 + .map_err(|error| tokenizer_error!("failed to load tokenizer: {}", error.as_report()))?;
117 + Ok(Self::from_hf_backend(t))
118 + }
119 +
120 + /// Load from `tokenizer.json` via fastokens or HuggingFace tokenizers.
121 + pub fn new(path: &Path) -> Result<Self> {
122 + match Self::new_fastokens(path) {
123 + Ok(tokenizer) => Ok(tokenizer),
124 + Err(error) => {
125 + warn!(
126 + path = %path.display(),
127 + error = %error.as_report(),
128 + "failed to load tokenizer with fastokens; falling back to HuggingFace tokenizers"
129 + );
130 + Self::new_hf(path)
131 + }
132 + }
133 + }
134 +}
135 +
136 +impl Tokenizer for HuggingFaceTokenizer {
137 + fn encode(&self, text: &str, add_special_tokens: bool) -> Result<Vec<u32>> {
138 + match &self.backend {
139 + Backend::Hf(t) => {
140 + let encoding = t
141 + .encode(text, add_special_tokens)
142 + .map_err(|error| tokenizer_error!("encoding failed: {}", error.as_report()))?;
143 + Ok(encoding.get_ids().to_vec())
144 + }
145 + Backend::Fastokens(t) | Backend::FastokensByteLevel(t) => t
146 + .encode_with_special_tokens(text, add_special_tokens)
147 + .map_err(|error| tokenizer_error!("encoding failed: {}", error.as_report())),
148 + }
149 + }
150 +
151 + fn decode(&self, token_ids: &[u32], skip_special_tokens: bool) -> Result<String> {
152 + match &self.backend {
153 + Backend::Hf(t) => t
154 + .decode(token_ids, skip_special_tokens)
155 + .map_err(|error| tokenizer_error!("decoding failed: {}", error.as_report())),
156 + Backend::Fastokens(t) => t
157 + .decode(token_ids, skip_special_tokens)
158 + .map_err(|error| tokenizer_error!("decoding failed: {}", error.as_report())),
159 + Backend::FastokensByteLevel(t) => {
160 + decode_fastokens_byte_level(t, token_ids, skip_special_tokens)
161 + }
162 + }
163 + }
164 +
165 + fn token_to_id(&self, token: &str) -> Option<u32> {
166 + match &self.backend {
167 + Backend::Hf(t) => t.token_to_id(token),
168 + Backend::Fastokens(t) | Backend::FastokensByteLevel(t) => t.token_to_id(token),
169 + }
170 + }
171 +
172 + fn id_to_token(&self, id: u32) -> Option<String> {
173 + match &self.backend {
174 + Backend::Hf(t) => t.id_to_token(id),
175 + Backend::Fastokens(t) | Backend::FastokensByteLevel(t) => {
176 + t.id_to_token(id).map(ToOwned::to_owned)
177 + }
178 + }
179 + }
180 +
181 + fn is_special_id(&self, token_id: u32) -> bool {
182 + self.special_token_ids.binary_search(&token_id).is_ok()
183 + }
184 +}
185 +
186 +#[cfg(test)]
187 +mod tests {
188 + use tempfile::tempdir;
189 + use tokenizers::models::bpe::BPE;
190 + use tokenizers::{AddedToken, Tokenizer as HfTokenizer};
191 +
192 + use super::{HuggingFaceTokenizer, Tokenizer};
193 +
194 + fn tiny_bpe_tokenizer() -> HfTokenizer {
195 + let vocab = [
196 + ("<unk>".to_string(), 0),
197 + ("h".to_string(), 1),
198 + ("e".to_string(), 2),
199 + ("l".to_string(), 3),
200 + ("o".to_string(), 4),
201 + ("he".to_string(), 5),
202 + ("ll".to_string(), 6),
203 + ("hell".to_string(), 7),
204 + ("hello".to_string(), 8),
205 + ];
206 + let merges = vec![
207 + ("h".to_string(), "e".to_string()),
208 + ("l".to_string(), "l".to_string()),
209 + ("he".to_string(), "ll".to_string()),
210 + ("hell".to_string(), "o".to_string()),
211 + ];
212 + let model = BPE::builder()
213 + .vocab_and_merges(vocab, merges)
214 + .unk_token("<unk>".to_string())
215 + .build()
216 + .expect("build bpe tokenizer");
217 + HfTokenizer::new(model)
218 + }
219 +
220 + #[test]
221 + fn hf_constructor_resolves_added_token_ids() {
222 + let mut tokenizer = tiny_bpe_tokenizer();
223 + tokenizer.add_special_tokens(&[AddedToken::from("<|im_end|>", true)]);
224 +
225 + let dir = tempdir().expect("create temp dir");
226 + let path = dir.path().join("tokenizer.json");
227 + tokenizer.save(&path, false).expect("save tokenizer json");
228 +
229 + let wrapper = HuggingFaceTokenizer::new_hf(&path).expect("load hf wrapper");
230 + let special_id = wrapper.token_to_id("<|im_end|>").expect("resolve added special token id");
231 + assert!(wrapper.is_special_id(special_id));
232 + }
233 +
234 + #[test]
235 + fn new_fastokens_preserves_special_ids_from_fastokens_metadata() {
236 + let mut tokenizer = tiny_bpe_tokenizer();
237 + tokenizer.add_special_tokens(&[AddedToken::from("<|im_end|>", true)]);
238 +
239 + let dir = tempdir().expect("create temp dir");
240 + let path = dir.path().join("tokenizer.json");
241 + tokenizer.save(&path, false).expect("save tokenizer json");
242 +
243 + let wrapper = HuggingFaceTokenizer::new_fastokens(&path)
244 + .expect("load wrapper with fastokens backend");
245 + assert!(matches!(
246 + wrapper.backend,
247 + super::Backend::Fastokens(_) | super::Backend::FastokensByteLevel(_),
248 + ));
249 + let special_id = wrapper.token_to_id("<|im_end|>").expect("resolve added special token id");
250 + assert!(wrapper.is_special_id(special_id));
251 + }
252 +
253 + /// BPE tokenizer that round-trips through fastokens with a genuine
254 + /// `ByteLevel` decoder; vocab covers both GPT-2 (Ġ U+0120) and non-GPT-2
255 + /// (｜ U+FF5C) codepoints.
256 + fn tiny_byte_level_bpe() -> fastokens::Tokenizer {
257 + let raw = r#"{
258 + "version": "1.0",
259 + "truncation": null,
260 + "padding": null,
261 + "added_tokens": [
262 + {"id": 0, "content": "<|endoftext|>", "single_word": false,
263 + "lstrip": false, "rstrip": false, "normalized": false, "special": true}
264 + ],
265 + "normalizer": null,
266 + "pre_tokenizer": {"type": "ByteLevel", "add_prefix_space": false,
267 + "trim_offsets": true, "use_regex": true},
268 + "post_processor": null,
269 + "decoder": {"type": "ByteLevel", "add_prefix_space": false,
270 + "trim_offsets": true, "use_regex": true},
271 + "model": {
272 + "type": "BPE",
273 + "dropout": null,
274 + "unk_token": null,
275 + "continuing_subword_prefix": null,
276 + "end_of_word_suffix": null,
277 + "fuse_unk": false,
278 + "byte_fallback": false,
279 + "ignore_merges": false,
280 + "vocab": {
281 + "<|endoftext|>": 0,
282 + "H": 1, "e": 2, "l": 3, "o": 4, "w": 5, "r": 6, "d": 7,
283 + "Ġ": 8, "!": 9,
284 + "｜": 10
285 + },
286 + "merges": []
287 + }
288 + }"#;
289 + let value: serde_json::Value = serde_json::from_str(raw).expect("parse tokenizer json");
290 + fastokens::Tokenizer::from_json(value).expect("build fastokens tokenizer")
291 + }
292 +
293 + #[test]
294 + fn byte_level_detected_direct() {
295 + let t = tiny_byte_level_bpe();
296 + assert!(super::is_byte_level_only(t.decoder().expect("decoder")));
297 + }
298 +
299 + #[test]
300 + fn byte_level_detected_inside_sequence() {
301 + let raw = r#"{
302 + "type": "Sequence",
303 + "decoders": [
304 + {"type": "ByteLevel", "add_prefix_space": false,
305 + "trim_offsets": true, "use_regex": true},
306 + {"type": "Fuse"}
307 + ]
308 + }"#;
309 + let config: fastokens::DecoderConfig =
310 + serde_json::from_str(raw).expect("parse decoder config");
311 + let dec =
312 + fastokens::decoders::Decoder::from_config(config).expect("build decoder from config");
313 + assert!(super::is_byte_level_only(&dec));
314 + }
315 +
316 + /// Fast path must produce byte-identical output to fastokens' own decode.
317 + #[test]
318 + fn fast_byte_level_matches_fastokens_decode() {
319 + let t = tiny_byte_level_bpe();
320 + let cases: &[&[u32]] = &[
321 + &[],
322 + &[1, 2, 3, 3, 4], // "Hello"
323 + &[1, 2, 3, 3, 4, 8, 5, 4, 6, 3, 7], // "Hello world"
324 + &[0, 1, 2, 3, 3, 4, 0, 9, 0], // specials interleaved
325 + &[10, 1, 2, 3, 3, 4, 10], // ｜Hello｜ (non-GPT2 chars)
326 + ];
327 + for ids in cases {
328 + for &skip in &[false, true] {
329 + let expected = t.decode(ids, skip).expect("fastokens decode");
330 + let got =
331 + super::decode_fastokens_byte_level(&t, ids, skip).expect("fast-path decode");
332 + assert_eq!(got, expected, "ids={ids:?} skip={skip}");
333 + }
334 + }
335 + }
336 +
337 + #[test]
338 + fn fast_byte_level_errors_on_unknown_id() {
339 + let t = tiny_byte_level_bpe();
340 + let err = super::decode_fastokens_byte_level(&t, &[999], false)
341 + .expect_err("unknown id must error");
342 + assert!(format!("{err:?}").contains("999"));
343 + }
344 +}
rust/src/tokenizer/src/incremental.rs new
+359
@@ -0,0 +1,359 @@
1 +use std::mem::take;
2 +
3 +use crate::{Result, Tokenizer};
4 +
5 +/// Stateful incremental decoder that emits text chunks one token at a time.
6 +pub trait IncrementalDecoder: Send {
7 + /// Push one generated token and return how many new string bytes were
8 + /// added.
9 + fn push_token(&mut self, token_id: u32) -> Result<usize>;
10 +
11 + /// Consume any text which is currently ready.
12 + fn next_chunk(&mut self) -> Option<String>;
13 +
14 + /// Flush any remaining buffered text that has not yet been emitted.
15 + ///
16 + /// Called after the final generated token to force out buffered/incomplete
17 + /// fragments.
18 + fn flush(&mut self, truncate_output_to: Option<usize>) -> Result<(Option<String>, String)>;
19 +
20 + /// Return cumulative decoded text so far.
21 + fn output(&self) -> &str;
22 +}
23 +
24 +/// [`IncrementalDecoder`] built on [`Tokenizer::decode()`] with prefix-diffing.
25 +///
26 +/// This is the same sliding-window algorithm used by `tokenizers::DecodeStream`
27 +pub(crate) struct DecodeStream<'a, T: Tokenizer + ?Sized> {
28 + tokenizer: &'a T,
29 + skip_special_tokens: bool,
30 + min_bytes_to_buffer: usize,
31 + // mutated state
32 + ids: Vec<u32>,
33 + prefix: String,
34 + prefix_index: usize,
35 + cumulative_output: String,
36 + output_index: usize,
37 +}
38 +
39 +impl<'a, T: Tokenizer + ?Sized> DecodeStream<'a, T> {
40 + pub(crate) fn new(
41 + tokenizer: &'a T,
42 + prompt_token_ids: &[u32],
43 + skip_special_tokens: bool,
44 + min_bytes_to_buffer: usize,
45 + ) -> Self {
46 + Self {
47 + tokenizer,
48 + skip_special_tokens,
49 + min_bytes_to_buffer,
50 + ids: prompt_token_ids.to_vec(),
51 + prefix: String::new(),
52 + prefix_index: 0,
53 + cumulative_output: String::new(),
54 + output_index: 0,
55 + }
56 + }
57 +}
58 +
59 +/// Try a short tail suffix first (covers a CJK glyph straddling 1-2 token
60 +/// boundaries); beyond 6 tokens the fallback full-prompt decode is no worse
61 +/// than baseline so widening the sweep just adds overhead.
62 +const SAFE_SUFFIX_MIN: usize = 4;
63 +const SAFE_SUFFIX_MAX: usize = 6;
64 +
65 +impl<T: Tokenizer + ?Sized> DecodeStream<'_, T> {
66 + /// Seed `self.prefix` from the shortest trailing suffix whose decoded text
67 + /// has no U+FFFD — a clean decode means the suffix starts and ends at
68 + /// valid UTF-8/token boundaries, so priming from it is equivalent to
69 + /// priming from the full prompt.
70 + fn seed_prefix(&mut self) -> Result<()> {
71 + let prompt_len = self.ids.len();
72 + if prompt_len > SAFE_SUFFIX_MIN {
73 + let max_try = SAFE_SUFFIX_MAX.min(prompt_len - 1);
74 + for suffix_len in SAFE_SUFFIX_MIN..=max_try {
75 + let start = prompt_len - suffix_len;
76 + let decoded =
77 + self.tokenizer.decode(&self.ids[start..], self.skip_special_tokens)?;
78 + if !decoded.contains('\u{FFFD}') {
79 + self.prefix = decoded;
80 + self.ids.drain(..start);
81 + self.prefix_index = self.ids.len();
82 + return Ok(());
83 + }
84 + }
85 + }
86 + let decoded = self.tokenizer.decode(&self.ids, self.skip_special_tokens)?;
87 + if !decoded.ends_with('\u{FFFD}') {
88 + self.prefix = decoded;
89 + self.prefix_index = self.ids.len();
90 + }
91 + Ok(())
92 + }
93 +}
94 +
95 +impl<T: Tokenizer + ?Sized> IncrementalDecoder for DecodeStream<'_, T> {
96 + fn push_token(&mut self, token_id: u32) -> Result<usize> {
97 + if self.prefix.is_empty() && !self.ids.is_empty() {
98 + self.seed_prefix()?;
99 + }
100 +
101 + self.ids.push(token_id);
102 + let string = self.tokenizer.decode(&self.ids, self.skip_special_tokens)?;
103 + let prefix_len = self.prefix.len();
104 + if string.len() <= prefix_len || string.ends_with('\u{FFFD}') {
105 + return Ok(0);
106 + }
107 + // Ensure we split at a utf-8 char boundary.
108 + let new_chunk = &string[string.floor_char_boundary(prefix_len)..];
109 + self.cumulative_output.push_str(new_chunk);
110 + self.ids.drain(..self.prefix_index);
111 + self.prefix = self.tokenizer.decode(&self.ids, self.skip_special_tokens)?;
112 + self.prefix_index = self.ids.len();
113 + Ok(new_chunk.len())
114 + }
115 +
116 + fn next_chunk(&mut self) -> Option<String> {
117 + let cutoff = self.cumulative_output.len().saturating_sub(self.min_bytes_to_buffer);
118 + (cutoff > self.output_index).then(|| {
119 + let chunk = self.cumulative_output[self.output_index..cutoff].to_string();
120 + self.output_index = cutoff;
121 + chunk
122 + })
123 + }
124 +
125 + fn flush(&mut self, truncate_output_to: Option<usize>) -> Result<(Option<String>, String)> {
126 + if !self.ids.is_empty() {
127 + let string = self.tokenizer.decode(&self.ids, self.skip_special_tokens)?;
128 + let prefix_len = self.prefix.len();
129 + self.ids.clear();
130 + self.prefix.clear();
131 + self.prefix_index = 0;
132 + // Ensure we split at a utf-8 char boundary.
133 + self.cumulative_output
134 + .push_str(&string[string.floor_char_boundary(prefix_len)..]);
135 + }
136 + if let Some(truncate_output_to) = truncate_output_to {
137 + self.cumulative_output.truncate(truncate_output_to);
138 + }
139 + let last_chunk = (self.output_index < self.cumulative_output.len())
140 + .then(|| self.cumulative_output[self.output_index..].to_string());
141 + self.output_index = 0;
142 + Ok((last_chunk, take(&mut self.cumulative_output)))
143 + }
144 +
145 + fn output(&self) -> &str {
146 + &self.cumulative_output
147 + }
148 +}
149 +
150 +#[cfg(test)]
151 +mod tests {
152 + use super::*;
153 +
154 + /// Backend that treats each token ID as a raw byte, producing lossy UTF-8.
155 + #[derive(Debug)]
156 + struct Utf8Backend;
157 +
158 + impl Tokenizer for Utf8Backend {
159 + fn encode(&self, _text: &str, _add_special_tokens: bool) -> Result<Vec<u32>> {
160 + unreachable!()
161 + }
162 +
163 + fn decode(&self, token_ids: &[u32], _skip_special_tokens: bool) -> Result<String> {
164 + let bytes = token_ids.iter().map(|id| *id as u8).collect::<Vec<_>>();
165 + Ok(String::from_utf8_lossy(&bytes).into_owned())
166 + }
167 +
168 + fn token_to_id(&self, _token: &str) -> Option<u32> {
169 + unreachable!()
170 + }
171 + }
172 +
173 + #[test]
174 + fn holds_incomplete_utf8_until_complete() {
175 + let backend = Utf8Backend;
176 + let mut decoder = backend.create_decode_stream(&[], false, 0);
177 +
178 + // 你 = U+4F60 = 0xE4 0xBD 0xA0
179 + assert_eq!(decoder.push_token(0xe4).unwrap(), 0);
180 + assert_eq!(decoder.push_token(0xbd).unwrap(), 0);
181 + assert_eq!(decoder.push_token(0xa0).unwrap(), 3); // "你" is 3 bytes
182 + assert_eq!(decoder.output(), "你");
183 + }
184 +
185 + #[test]
186 + fn emits_ascii_immediately() {
187 + let backend = Utf8Backend;
188 + let mut decoder = backend.create_decode_stream(&[], false, 0);
189 +
190 + assert_eq!(decoder.push_token(b'o' as u32).unwrap(), 1);
191 + assert_eq!(decoder.push_token(b'k' as u32).unwrap(), 1);
192 + assert_eq!(decoder.output(), "ok");
193 + }
194 +
195 + #[test]
196 + fn flush_returns_none_when_fully_consumed() {
197 + let backend = Utf8Backend;
198 + let mut decoder = backend.create_decode_stream(&[], false, 0);
199 +
200 + assert_eq!(decoder.push_token(b'o' as u32).unwrap(), 1);
201 + assert_eq!(decoder.next_chunk().as_deref(), Some("o"));
202 + assert_eq!(decoder.push_token(b'k' as u32).unwrap(), 1);
203 + assert_eq!(decoder.next_chunk().as_deref(), Some("k"));
204 + // All text already consumed via next_chunk
205 + let (last_chunk, full_text) = decoder.flush(None).unwrap();
206 + assert_eq!(last_chunk, None);
207 + assert_eq!(full_text, "ok");
208 + }
209 +
210 + #[test]
211 + fn flush_emits_buffered_incomplete_utf8() {
212 + let backend = Utf8Backend;
213 + let mut decoder = backend.create_decode_stream(&[], false, 0);
214 +
215 + // Push incomplete multi-byte sequence — step returns 0 bytes.
216 + assert_eq!(decoder.push_token(0xe4).unwrap(), 0);
217 + assert_eq!(decoder.push_token(0xbd).unwrap(), 0);
218 +
219 + // Flush forces out whatever the decoder can produce (lossy replacement).
220 + let (last_chunk, _full_text) = decoder.flush(None).unwrap();
221 + assert!(last_chunk.is_some());
222 + }
223 +
224 + /// Backend where token 0 is a special token.
225 + #[derive(Debug)]
226 + struct SpecialTokenBackend;
227 +
228 + impl Tokenizer for SpecialTokenBackend {
229 + fn encode(&self, _text: &str, _add_special_tokens: bool) -> Result<Vec<u32>> {
230 + unreachable!()
231 + }
232 +
233 + fn decode(&self, token_ids: &[u32], skip_special_tokens: bool) -> Result<String> {
234 + let mut text = String::new();
235 + for &token_id in token_ids {
236 + match token_id {
237 + 0 if !skip_special_tokens => text.push_str("<special>"),
238 + 0 => {}
239 + 1 => text.push('a'),
240 + _ => {}
241 + }
242 + }
243 + Ok(text)
244 + }
245 +
246 + fn token_to_id(&self, _token: &str) -> Option<u32> {
247 + unreachable!()
248 + }
249 + }
250 +
251 + #[test]
252 + fn respects_skip_special_tokens() {
253 + let backend = SpecialTokenBackend;
254 + let mut skip_decoder = backend.create_decode_stream(&[], true, 0);
255 + let mut keep_decoder = backend.create_decode_stream(&[], false, 0);
256 +
257 + assert_eq!(skip_decoder.push_token(0).unwrap(), 0);
258 + assert_eq!(keep_decoder.push_token(0).unwrap(), 9); // "<special>" is 9 bytes
259 + assert_eq!(keep_decoder.output(), "<special>");
260 + }
261 +
262 + #[test]
263 + fn prompt_tokens_provide_context_without_re_emission() {
264 + let backend = Utf8Backend;
265 + let prompt = &[b'H' as u32, b'i' as u32];
266 + let mut decoder = backend.create_decode_stream(prompt, false, 0);
267 +
268 + // First generated token should not re-emit "Hi".
269 + let added = decoder.push_token(b'!' as u32).unwrap();
270 + assert_eq!(added, 1);
271 + assert_eq!(decoder.output(), "!");
272 + }
273 +
274 + #[test]
275 + fn chunks_concatenate_to_full_text() {
276 + let backend = Utf8Backend;
277 + let mut decoder = backend.create_decode_stream(&[], false, 0);
278 +
279 + let input = b"Hello, world!";
280 + let mut full = String::new();
281 + for &byte in input {
282 + decoder.push_token(byte as u32).unwrap();
283 + if let Some(chunk) = decoder.next_chunk() {
284 + full.push_str(&chunk);
285 + }
286 + }
287 + let (last_chunk, full_text) = decoder.flush(None).unwrap();
288 + assert_eq!(last_chunk, None); // all consumed via next_chunk
289 + assert_eq!(full, "Hello, world!");
290 + assert_eq!(full_text, "Hello, world!");
291 + }
292 +
293 + /// Backend simulating non-monotonic decode where adding a token changes how
294 + /// earlier tokens decode (context-dependent normalization), causing
295 + /// prefix_len to land mid-UTF-8. Reproduces the class of bug from
296 + /// vllm-project/vllm#17448.
297 + #[derive(Debug)]
298 + struct NonMonotonicBackend;
299 +
300 + impl Tokenizer for NonMonotonicBackend {
301 + fn encode(&self, _text: &str, _add_special_tokens: bool) -> Result<Vec<u32>> {
302 + unreachable!()
303 + }
304 +
305 + fn decode(&self, token_ids: &[u32], _skip_special_tokens: bool) -> Result<String> {
306 + match token_ids {
307 + [1] => Ok("abc".into()),
308 + [1, 2] => Ok("ab".into()),
309 + // Token 3 triggers a normalization change: "ab" becomes emoji + "d".
310 + // prefix_len=3 ("abc") lands inside the 4-byte emoji 🎉.
311 + [1, 2, 3] => Ok("🎉d".into()), // 🎉 is 4 bytes + d = 5 bytes
312 + [2, 3] => Ok("🎉d".into()), // prefix recompute after drain
313 + [3] => Ok("d".into()), // after drain
314 + _ => panic!("unexpected decode: {:?}", token_ids),
315 + }
316 + }
317 +
318 + fn token_to_id(&self, _token: &str) -> Option<u32> {
319 + unreachable!()
320 + }
321 + }
322 +
323 + /// Without the char-boundary fix, this panics slicing mid-emoji.
324 + #[test]
325 + fn non_monotonic_decode_does_not_panic() {
326 + let backend = NonMonotonicBackend;
327 + let mut decoder = backend.create_decode_stream(&[], false, 0);
328 +
329 + // Token 1: "abc", prefix="abc"
330 + assert_eq!(decoder.push_token(1).unwrap(), 3);
331 + // Token 2: "ab" (shorter), no emit
332 + assert_eq!(decoder.push_token(2).unwrap(), 0);
333 + // Token 3: "🎉d" — prefix_len=3 is mid-emoji. Without fix this panics.
334 + let added = decoder.push_token(3).unwrap();
335 + assert!(added > 0);
336 + }
337 +
338 + #[test]
339 + fn next_chunk_with_hold_back() {
340 + let backend = Utf8Backend;
341 + // hold_back_bytes: 3 means we buffer the last 3 bytes
342 + let mut decoder = backend.create_decode_stream(&[], false, 3);
343 +
344 + let input = b"Hello!";
345 + let mut chunks = String::new();
346 + for &byte in input {
347 + decoder.push_token(byte as u32).unwrap();
348 + if let Some(chunk) = decoder.next_chunk() {
349 + chunks.push_str(&chunk);
350 + }
351 + }
352 + // With hold_back_bytes=3, last 3 bytes ("lo!") are held back
353 + assert_eq!(chunks, "Hel");
354 + // Flush returns the rest
355 + let (last_chunk, full_text) = decoder.flush(None).unwrap();
356 + assert_eq!(last_chunk.as_deref(), Some("lo!"));
357 + assert_eq!(full_text, "Hello!");
358 + }
359 +}
rust/src/tokenizer/src/lib.rs new
+62
@@ -0,0 +1,62 @@
1 +use std::sync::Arc;
2 +
3 +use crate::incremental::DecodeStream;
4 +
5 +mod byte_level_decode;
6 +#[macro_use]
7 +mod error;
8 +mod hf;
9 +mod incremental;
10 +mod tekken;
11 +mod tiktoken;
12 +
13 +pub use error::{Result, TokenizerError};
14 +pub use hf::HuggingFaceTokenizer;
15 +pub use incremental::IncrementalDecoder;
16 +pub use tekken::TekkenTokenizer;
17 +pub use tiktoken::TiktokenTokenizer;
18 +
19 +pub trait Tokenizer: Send + Sync {
20 + /// Encode one prompt string into token IDs.
21 + fn encode(&self, text: &str, add_special_tokens: bool) -> Result<Vec<u32>>;
22 +
23 + /// Decode one token sequence into text.
24 + fn decode(&self, token_ids: &[u32], skip_special_tokens: bool) -> Result<String>;
25 +
26 + /// Convert one token string into a token ID, returning `None` if the token
27 + /// is not in the tokenizer vocabulary.
28 + fn token_to_id(&self, token: &str) -> Option<u32>;
29 +
30 + /// Convert one token ID into the tokenizer's raw token string.
31 + fn id_to_token(&self, _id: u32) -> Option<String> {
32 + // TODO: remove default impl and require this to be implemented by all
33 + // tokenizers
34 + None
35 + }
36 +
37 + /// Return whether the given token ID is special.
38 + fn is_special_id(&self, _token_id: u32) -> bool {
39 + false
40 + }
41 +
42 + /// Create a stateful incremental decoder primed with the given prompt
43 + /// tokens.
44 + ///
45 + /// The prompt tokens provide left context for the first generated token;
46 + /// the decoder does not re-emit prompt text.
47 + fn create_decode_stream(
48 + &self,
49 + prompt_token_ids: &[u32],
50 + skip_special_tokens: bool,
51 + min_bytes_to_buffer: usize,
52 + ) -> Box<dyn IncrementalDecoder + '_> {
53 + Box::new(DecodeStream::new(
54 + self,
55 + prompt_token_ids,
56 + skip_special_tokens,
57 + min_bytes_to_buffer,
58 + ))
59 + }
60 +}
61 +
62 +pub type DynTokenizer = Arc<dyn Tokenizer>;
rust/src/tokenizer/src/tekken.rs new
+62
@@ -0,0 +1,62 @@
1 +use std::path::Path;
2 +
3 +use tekken::Tekkenizer;
4 +use tracing::info;
5 +
6 +use crate::{Result, Tokenizer};
7 +
8 +/// Mistral Tekken tokenizer from a `tekken.json` file.
9 +pub struct TekkenTokenizer {
10 + inner: Tekkenizer,
11 +}
12 +
13 +impl TekkenTokenizer {
14 + /// Load a Mistral Tekken tokenizer from a `tekken.json` file.
15 + pub fn new(path: &Path) -> Result<Self> {
16 + info!(path = %path.display(), "loading tokenizer with Mistral Tekken");
17 +
18 + let inner = Tekkenizer::from_file(path).map_err(|error| {
19 + tokenizer_error!(
20 + "failed to load tekken tokenizer from {}: {error}",
21 + path.display()
22 + )
23 + })?;
24 + Ok(Self { inner })
25 + }
26 +}
27 +
28 +impl Tokenizer for TekkenTokenizer {
29 + fn encode(&self, text: &str, add_special_tokens: bool) -> Result<Vec<u32>> {
30 + self.inner
31 + .encode(text, add_special_tokens, false)
32 + .map_err(|error| tokenizer_error!("encoding failed: {error}"))
33 + }
34 +
35 + fn decode(&self, token_ids: &[u32], skip_special_tokens: bool) -> Result<String> {
36 + let policy = if skip_special_tokens {
37 + tekken::SpecialTokenPolicy::Ignore
38 + } else {
39 + tekken::SpecialTokenPolicy::Keep
40 + };
41 + self.inner
42 + .decode(token_ids, policy)
43 + .map_err(|error| tokenizer_error!("decoding failed: {error}"))
44 + }
45 +
46 + fn token_to_id(&self, token: &str) -> Option<u32> {
47 + // tekken-rs exposes `get_control_token` for special tokens. Try that first,
48 + // then fall back to encoding.
49 + self.inner.get_control_token(token).ok().or_else(|| {
50 + let ids = self.inner.encode(token, false, false).ok()?;
51 + if ids.len() == 1 { Some(ids[0]) } else { None }
52 + })
53 + }
54 +
55 + fn id_to_token(&self, id: u32) -> Option<String> {
56 + self.inner.id_to_piece(id).ok()
57 + }
58 +
59 + fn is_special_id(&self, token_id: u32) -> bool {
60 + self.inner.is_special_token(token_id)
61 + }
62 +}
rust/src/tokenizer/src/tiktoken.rs new
+1013
@@ -0,0 +1,1013 @@
1 +use std::collections::HashSet;
2 +use std::path::Path;
3 +use std::sync::Mutex;
4 +
5 +use base64::Engine as _;
6 +use rustc_hash::{FxHashMap, FxHashSet};
7 +use serde::Deserialize;
8 +use thiserror_ext::AsReport as _;
9 +use tracing::{info, warn};
10 +
11 +use crate::{Result, Tokenizer};
12 +
13 +/// Default regex pattern used when loading tiktoken from a BPE file. This is
14 +/// the same `cl100k_base` pattern that HuggingFace transformers uses as its
15 +/// default in `TikTokenConverter`.
16 +const CL100K_BASE_PATTERN: &str = r"(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}{1,3}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+";
17 +
18 +/// Kimi BPE pattern from `moonshotai/Kimi-K2-Instruct/tokenization_kimi.py`.
19 +const KIMI_PATTERN: &str = r"[\p{Han}]+|[^\r\n\p{L}\p{N}]?[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}&&[^\p{Han}]]*[\p{Ll}\p{Lm}\p{Lo}\p{M}&&[^\p{Han}]]+(?i:'s|'t|'re|'ve|'m|'ll|'d)?|[^\r\n\p{L}\p{N}]?[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}&&[^\p{Han}]]+[\p{Ll}\p{Lm}\p{Lo}\p{M}&&[^\p{Han}]]*(?i:'s|'t|'re|'ve|'m|'ll|'d)?|\p{N}{1,3}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+";
20 +
21 +/// Fallback number of reserved special-token slots to assume when the model's
22 +/// `config.json` is not available (so we cannot read `vocab_size` directly).
23 +///
24 +/// 256 is the value used by Kimi K2 / K2.5 (`tokenization_kimi.py`'s
25 +/// `num_reserved_special_tokens`) and by Llama 3, and it appears to be the most
26 +/// common convention among modern tiktoken-based HF tokenizers. When
27 +/// `config.json` *is* present we honour the model's actual `vocab_size` instead
28 +/// of this fallback — see `Self::new`.
29 +const FALLBACK_NUM_RESERVED_SPECIAL_TOKENS: u32 = 256;
30 +const DISABLE_RIPTOKEN_ENV: &str = "VLLM_RS_DISABLE_RIPTOKEN";
31 +
32 +/// Parsed entry from `tokenizer_config.json`'s `added_tokens_decoder`.
33 +#[derive(Debug, Clone, Deserialize)]
34 +struct AddedToken {
35 + content: String,
36 + /// HuggingFace `added_tokens_decoder` entries can be marked `"special":
37 + /// true|false`. Special tokens are dropped from output when `decode` is
38 + /// called with `skip_special_tokens = true`. Defaults to `false` when
39 + /// the field is omitted, matching HuggingFace's `AddedToken` default —
40 + /// so only tokens explicitly marked special are stripped during normal
41 + /// decode (where `skip_special_tokens` itself defaults to true).
42 + #[serde(default)]
43 + special: bool,
44 +}
45 +
46 +/// Minimal subset of `tokenizer_config.json` needed by the tiktoken loader.
47 +#[derive(Debug, Default, Deserialize)]
48 +#[serde(default)]
49 +struct TiktokenTokenizerConfig {
50 + /// Format:
51 + /// `{ "added_tokens_decoder": { "163584": { "content": "[BOS]", "special":
52 + /// true }, ... } }`
53 + #[serde(default)]
54 + added_tokens_decoder: FxHashMap<u32, AddedToken>,
55 +}
56 +
57 +/// Minimal subset of model `config.json` needed by the tiktoken loader.
58 +#[derive(Debug, Default, Deserialize)]
59 +#[serde(default)]
60 +struct TiktokenModelConfig {
61 + model_type: Option<String>,
62 + vocab_size: Option<u32>,
63 + text_config: Option<Box<TiktokenModelConfig>>,
64 +}
65 +
66 +impl TiktokenModelConfig {
67 + /// Read `model_type` from a model `config.json` value, falling back to a
68 + /// single-level nested `text_config.model_type` for composite (e.g.
69 + /// multimodal) configs that keep text metadata under a `text_config`
70 + /// object.
71 + fn effective_model_type(&self) -> Option<&str> {
72 + self.model_type
73 + .as_deref()
74 + .or_else(|| self.text_config.as_deref()?.effective_model_type())
75 + }
76 +
77 + /// Read `vocab_size` from a model `config.json` value, falling back to a
78 + /// single-level nested `text_config.vocab_size` for composite (e.g.
79 + /// multimodal) configs that keep text metadata under a `text_config`
80 + /// object — matching the same shape `ModelConfig` parses.
81 + fn effective_vocab_size(&self) -> Option<u32> {
82 + self.vocab_size.or_else(|| self.text_config.as_deref()?.effective_vocab_size())
83 + }
84 +}
85 +
86 +/// Tiktoken tokenizer from `tiktoken.model` or `*.tiktoken` BPE files.
87 +pub struct TiktokenTokenizer {
88 + backend: Backend,
89 + metadata: TokenMetadata,
90 +}
91 +
92 +enum Backend {
93 + Riptoken(RiptokenBackend),
94 + TiktokenRs(TiktokenRsBackend),
95 +}
96 +
97 +struct RiptokenBackend {
98 + inner: Box<riptoken::CoreBPE>,
99 + allowed_special_tokens: Vec<String>,
100 +}
101 +
102 +struct TiktokenRsBackend {
103 + inner: Box<tiktoken_rs::CoreBPE>,
104 + /// Reverse map for special / added token strings populated from the
105 + /// reserved range. This lets `token_to_id` answer special-token lookups
106 + /// directly without round-tripping through `tiktoken-rs`'s encoder,
107 + /// which can panic for unknown special-looking strings.
108 + special_token_ids_by_text: FxHashMap<String, u32>,
109 + /// Set of out-of-vocab token IDs we have already warned about. The
110 + /// reserved-slot population in the constructor should keep this empty
111 + /// under normal operation; it only fills up if a model emits ids at or
112 + /// above `vocab_upper_bound` (e.g. an engine sampling bug). We dedupe
113 + /// so streaming decode (which calls `decode` repeatedly on the same prefix)
114 + /// does not spam.
115 + warned_unknown_ids: Mutex<FxHashSet<u32>>,
116 +}
117 +
118 +struct TokenMetadata {
119 + /// Number of regular BPE tokens. Token ids in `[0, num_base_tokens)` are
120 + /// BPE tokens that always decode to text; ids in `[num_base_tokens,
121 + /// vocab_upper_bound)` live in the special-token slots and are subject
122 + /// to `skip_special_tokens` filtering.
123 + num_base_tokens: u32,
124 + /// Exclusive upper bound on token IDs that `inner` is guaranteed to know
125 + /// how to decode.
126 + ///
127 + /// The constructor registers every id in `[num_base_tokens,
128 + /// vocab_upper_bound)` with the inner `CoreBPE` as a (named or
129 + /// `<|reserved_token_{id}|>`) special token, and the BPE
130 + /// encoder densely covers `[0, num_base_tokens)`. So any id below this
131 + /// bound is in one of the inner `CoreBPE`'s decoder maps and
132 + /// `_decode_native_and_split` will not panic on it. `decode` filters
133 + /// out ids at or above this bound to keep that guarantee.
134 + vocab_upper_bound: u32,
135 + /// Ids in `[num_base_tokens, vocab_upper_bound)` whose
136 + /// `added_tokens_decoder` entry was explicitly marked `"special":
137 + /// false` — i.e. tokens that should still appear in output
138 + /// even when `skip_special_tokens = true`. For Kimi K2 / K2.5 this
139 + /// typically holds the tool-call markers and `<think>` / `</think>`.
140 + /// Reserved-slot placeholders are not in this set (they default to
141 + /// special and get skipped).
142 + non_special_added_ids: FxHashSet<u32>,
143 + /// Raw token string by token id. Base BPE tokens are represented with
144 + /// lossy UTF-8, matching decode behavior for byte sequences that are not
145 + /// valid UTF-8 on their own.
146 + token_by_id: FxHashMap<u32, String>,
147 +}
148 +
149 +impl TokenMetadata {
150 + fn filter_special_tokens(&self, token_ids: &[u32]) -> Vec<u32> {
151 + token_ids
152 + .iter()
153 + .copied()
154 + .filter(|&id| {
155 + id < self.num_base_tokens
156 + || id >= self.vocab_upper_bound
157 + || self.non_special_added_ids.contains(&id)
158 + })
159 + .collect()
160 + }
161 +
162 + fn is_special_id(&self, token_id: u32) -> bool {
163 + token_id >= self.num_base_tokens
164 + && token_id < self.vocab_upper_bound
165 + && !self.non_special_added_ids.contains(&token_id)
166 + }
167 +
168 + fn id_to_token(&self, token_id: u32) -> Option<String> {
169 + self.token_by_id.get(&token_id).cloned()
170 + }
171 +}
172 +
173 +impl RiptokenBackend {
174 + fn encode(&self, text: &str) -> Vec<u32> {
175 + // TODO: avoid collecting `allowed_special` every time this method is called.
176 + let allowed_special: HashSet<&str> =
177 + self.allowed_special_tokens.iter().map(String::as_str).collect();
178 + self.inner.encode(text, &allowed_special)
179 + }
180 +
181 + fn decode(&self, token_ids: &[u32]) -> String {
182 + let bytes = self.inner.decode_bytes(token_ids);
183 + // TODO: use `from_utf8_lossy_owned` once it's stabilized.
184 + String::from_utf8_lossy(&bytes).into_owned()
185 + }
186 +
187 + fn token_to_id(&self, token: &str) -> Option<u32> {
188 + self.inner.encode_single_token(token.as_bytes())
189 + }
190 +}
191 +
192 +impl TiktokenRsBackend {
193 + fn encode(&self, text: &str) -> Vec<u32> {
194 + self.inner.encode_with_special_tokens(text)
195 + }
196 +
197 + fn decode(&self, token_ids: &[u32], metadata: &TokenMetadata) -> String {
198 + let safe_ids: Vec<u32> = token_ids
199 + .iter()
200 + .copied()
201 + .filter(|&id| {
202 + if id >= metadata.vocab_upper_bound {
203 + self.warn_unknown_id(id);
204 + return false;
205 + }
206 + true
207 + })
208 + .collect();
209 + let bytes: Vec<u8> = self.inner._decode_native_and_split(safe_ids).flatten().collect();
210 + // TODO: use `from_utf8_lossy_owned` once it's stabilized.
211 + String::from_utf8_lossy(&bytes).into_owned()
212 + }
213 +
214 + fn token_to_id(&self, token: &str) -> Option<u32> {
215 + if let Some(&token_id) = self.special_token_ids_by_text.get(token) {
216 + return Some(token_id);
217 + }
218 +
219 + // Fall back to ordinary encoding for regular vocabulary items. This
220 + // deliberately avoids `encode_with_special_tokens`: older `tiktoken-rs`
221 + // versions can panic if the input text merely *looks* like a special
222 + // token but is not registered in `special_tokens_encoder`.
223 + let ids = self.inner.encode_ordinary(token);
224 + if ids.len() == 1 { Some(ids[0]) } else { None }
225 + }
226 +
227 + /// Log a warning the first time an unknown token id is seen during decode,
228 + /// deduped across calls so streaming decode does not spam the log for
229 + /// the same id.
230 + fn warn_unknown_id(&self, token_id: u32) {
231 + let newly_inserted = self
232 + .warned_unknown_ids
233 + .lock()
234 + .map(|mut set| set.insert(token_id))
235 + .unwrap_or(false);
236 + if newly_inserted {
237 + warn!(
238 + token_id,
239 + "tiktoken-rs decode encountered token id not in the vocabulary; skipping. \
240 + This typically indicates a sparse-vocab model whose `added_tokens_decoder` \
241 + does not list every reserved id in the special-token range."
242 + );
243 + }
244 + }
245 +}
246 +
247 +impl TiktokenTokenizer {
248 + /// Load a tiktoken tokenizer from a `.tiktoken` / `tiktoken.model` BPE
249 + /// file.
250 + ///
251 + /// The BPE file format is one `<base64-token-bytes> <rank>` pair per line,
252 + /// the same format used by OpenAI's tiktoken and by HuggingFace model
253 + /// repos that ship tiktoken files (e.g. DeepSeek, Kimi K2).
254 + ///
255 + /// Special / added tokens are read from `tokenizer_config.json` in the same
256 + /// directory when present. The `cl100k_base` regex pattern is used as a
257 + /// reasonable default.
258 + pub fn new(path: &Path) -> Result<Self> {
259 + if std::env::var_os(DISABLE_RIPTOKEN_ENV).is_some() {
260 + return Self::new_tiktoken_rs(path);
261 + }
262 +
263 + match Self::new_riptoken(path) {
264 + Ok(tokenizer) => Ok(tokenizer),
265 + Err(error) => {
266 + warn!(
267 + path = %path.display(),
268 + error = %error.as_report(),
269 + "failed to load tokenizer with riptoken; falling back to tiktoken-rs"
270 + );
271 + Self::new_tiktoken_rs(path)
272 + }
273 + }
274 + }
275 +
276 + /// Load from `tiktoken.model` / `*.tiktoken` with riptoken.
277 + pub fn new_riptoken(path: &Path) -> Result<Self> {
278 + info!(path = %path.display(), "loading tokenizer with riptoken (BPE file)");
279 +
280 + let config = LoadedTiktokenConfig::load(path)?;
281 + let allowed_special_tokens = config.special_tokens_encoder.keys().cloned().collect();
282 + let inner = riptoken::CoreBPE::new(
283 + config.encoder.into_iter().collect(),
284 + config.special_tokens_encoder.into_iter().collect(),
285 + config.pattern,
286 + )
287 + .map_err(|error| {
288 + tokenizer_error!(
289 + "failed to create riptoken tokenizer from {}: {error}",
290 + path.display()
291 + )
292 + })?;
293 +
294 + Ok(Self {
295 + backend: Backend::Riptoken(RiptokenBackend {
296 + inner: Box::new(inner),
297 + allowed_special_tokens,
298 + }),
299 + metadata: config.metadata,
300 + })
301 + }
302 +
303 + /// Load from `tiktoken.model` / `*.tiktoken` with tiktoken-rs.
304 + pub fn new_tiktoken_rs(path: &Path) -> Result<Self> {
305 + info!(path = %path.display(), "loading tokenizer with tiktoken-rs (BPE file)");
306 +
307 + let config = LoadedTiktokenConfig::load(path)?;
308 + let special_token_ids_by_text = config.special_tokens_encoder.clone();
309 + let inner = tiktoken_rs::CoreBPE::new(
310 + config.encoder,
311 + config.special_tokens_encoder,
312 + config.pattern,
313 + )
314 + .map_err(|error| {
315 + tokenizer_error!(
316 + "failed to create tiktoken-rs tokenizer from {}: {error}",
317 + path.display()
318 + )
319 + })?;
320 +
321 + Ok(Self {
322 + backend: Backend::TiktokenRs(TiktokenRsBackend {
323 + inner: Box::new(inner),
324 + special_token_ids_by_text,
325 + warned_unknown_ids: Mutex::new(FxHashSet::default()),
326 + }),
327 + metadata: config.metadata,
328 + })
329 + }
330 +}
331 +
332 +struct LoadedTiktokenConfig {
333 + encoder: FxHashMap<Vec<u8>, u32>,
334 + special_tokens_encoder: FxHashMap<String, u32>,
335 + metadata: TokenMetadata,
336 + pattern: &'static str,
337 +}
338 +
339 +impl LoadedTiktokenConfig {
340 + fn load(path: &Path) -> Result<Self> {
341 + let content = std::fs::read_to_string(path).map_err(|error| {
342 + tokenizer_error!(
343 + "failed to read tiktoken file {}: {}",
344 + path.display(),
345 + error.as_report()
346 + )
347 + })?;
348 + let mut encoder: FxHashMap<Vec<u8>, u32> =
349 + FxHashMap::with_capacity_and_hasher(content.lines().count(), Default::default());
350 + for line in content.lines() {
351 + if line.is_empty() {
352 + continue;
353 + }
354 + let mut parts = line.split_whitespace();
355 + let token_b64 =
356 + parts.next().ok_or_else(|| tokenizer_error!("missing token in tiktoken file"))?;
357 + let rank_str =
358 + parts.next().ok_or_else(|| tokenizer_error!("missing rank in tiktoken file"))?;
359 + let token_bytes = base64::engine::general_purpose::STANDARD
360 + .decode(token_b64)
361 + .map_err(|error| tokenizer_error!("invalid base64 in tiktoken file: {error}"))?;
362 + let rank: u32 = rank_str
363 + .parse()
364 + .map_err(|error| tokenizer_error!("invalid rank in tiktoken file: {error}"))?;
365 + encoder.insert(token_bytes, rank);
366 + }
367 +
368 + let parent_dir = path.parent();
369 +
370 + // Read added/special tokens (id -> {name, special}) from
371 + // tokenizer_config.json in the same dir.
372 + let added_tokens_by_id = parent_dir
373 + .map(|dir| dir.join("tokenizer_config.json"))
374 + .filter(|p| p.exists())
375 + .and_then(|config_path| {
376 + let content = std::fs::read_to_string(&config_path).ok()?;
377 + serde_json::from_str(&content).ok()
378 + })
379 + .map(|config: TiktokenTokenizerConfig| config.added_tokens_decoder)
380 + .unwrap_or_default();
381 +
382 + let model_config: Option<TiktokenModelConfig> = parent_dir
383 + .map(|dir| dir.join("config.json"))
384 + .filter(|p| p.exists())
385 + .and_then(|config_path| {
386 + let content = std::fs::read_to_string(&config_path).ok()?;
387 + serde_json::from_str(&content).ok()
388 + });
389 + let vocab_size_from_config = model_config.as_ref().and_then(|c| c.effective_vocab_size());
390 +
391 + // Build the full special-tokens encoder by populating the reserved
392 + // range that follows the BPE vocabulary. Unknown reserved slots get
393 + // Python-compatible placeholder names so sampled ids can still decode.
394 + //
395 + // Note: `*.tiktoken` ranks are token ids, and they are not guaranteed
396 + // to be contiguous. The base-vocab boundary is therefore `max_rank + 1`,
397 + // not `encoder.len()`.
398 + let num_base_tokens =
399 + encoder.values().copied().max().map_or(0, |max_rank| max_rank.saturating_add(1));
400 + let max_added_id = added_tokens_by_id.keys().copied().max().unwrap_or(0);
401 + let reserved_end = vocab_size_from_config
402 + .unwrap_or_else(|| num_base_tokens.saturating_add(FALLBACK_NUM_RESERVED_SPECIAL_TOKENS))
403 + .max(num_base_tokens)
404 + .max(max_added_id.saturating_add(1));
405 +
406 + let mut special_tokens_encoder: FxHashMap<String, u32> =
407 + FxHashMap::with_capacity_and_hasher(
408 + (reserved_end - num_base_tokens) as usize,
409 + Default::default(),
410 + );
411 + let mut non_special_added_ids: FxHashSet<u32> = FxHashSet::default();
412 + for id in num_base_tokens..reserved_end {
413 + let name = match added_tokens_by_id.get(&id) {
414 + Some(token) => {
415 + if !token.special {
416 + non_special_added_ids.insert(id);
417 + }
418 + token.content.clone()
419 + }
420 + None => format!("<|reserved_token_{id}|>"),
421 + };
422 + special_tokens_encoder.insert(name, id);
423 + }
424 +
425 + let mut token_by_id: FxHashMap<u32, String> = FxHashMap::with_capacity_and_hasher(
426 + encoder.len() + special_tokens_encoder.len(),
427 + Default::default(),
428 + );
429 + for (token_bytes, &id) in &encoder {
430 + token_by_id.insert(id, String::from_utf8_lossy(token_bytes).into_owned());
431 + }
432 + for (token, &id) in &special_tokens_encoder {
433 + token_by_id.insert(id, token.clone());
434 + }
435 +
436 + let pattern = model_config.as_ref().map_or(CL100K_BASE_PATTERN, detect_bpe_pattern);
437 +
438 + Ok(Self {
439 + encoder,
440 + special_tokens_encoder,
441 + metadata: TokenMetadata {
442 + num_base_tokens,
443 + vocab_upper_bound: reserved_end,
444 + non_special_added_ids,
445 + token_by_id,
446 + },
447 + pattern,
448 + })
449 + }
450 +}
451 +
452 +impl Tokenizer for TiktokenTokenizer {
453 + fn encode(&self, text: &str, _add_special_tokens: bool) -> Result<Vec<u32>> {
454 + // Tiktoken does not have a separate add_special_tokens toggle; both
455 + // backends recognize registered special tokens in the input.
456 + Ok(match &self.backend {
457 + Backend::Riptoken(backend) => backend.encode(text),
458 + Backend::TiktokenRs(backend) => backend.encode(text),
459 + })
460 + }
461 +
462 + fn decode(&self, token_ids: &[u32], skip_special_tokens: bool) -> Result<String> {
463 + // Filter passes:
464 + //
465 + // 1. The constructor registers every id in `[num_base_tokens, vocab_upper_bound)` as a
466 + // special token (named or `<|reserved_token_{id}|>` placeholder, matching
467 + // `tokenization_kimi.py`). The tiktoken-rs backend additionally drops ids at or above
468 + // that bound so `_decode_native_and_split` cannot panic; riptoken's `decode_bytes`
469 + // already skips unknown ids.
470 + //
471 + // 2. When `skip_special_tokens = true`, ids in `[num_base_tokens, vocab_upper_bound)` are
472 + // dropped *unless* they were marked `"special": false` in `added_tokens_decoder`. This
473 + // matches HuggingFace's tokenizer semantics: tool-call markers and `<think>` /
474 + // `</think>` (which Kimi K2 / K2.5 declare as non-special) stay in the output, while
475 + // BOS/EOS/header tokens and reserved-slot placeholders are stripped.
476 + //
477 + // Lossy UTF-8 decoding (instead of strict `String::from_utf8`) is used so
478 + // partial multi-byte sequences become `\u{FFFD}`, which `DecodeStream`
479 + // relies on to detect incomplete characters during streaming.
480 + let ids = if skip_special_tokens {
481 + &self.metadata.filter_special_tokens(token_ids)
482 + } else {
483 + token_ids
484 + };
485 +
486 + Ok(match &self.backend {
487 + Backend::Riptoken(backend) => backend.decode(ids),
488 + Backend::TiktokenRs(backend) => backend.decode(ids, &self.metadata),
489 + })
490 + }
491 +
492 + fn token_to_id(&self, token: &str) -> Option<u32> {
493 + match &self.backend {
494 + Backend::Riptoken(backend) => backend.token_to_id(token),
495 + Backend::TiktokenRs(backend) => backend.token_to_id(token),
496 + }
497 + }
498 +
499 + fn id_to_token(&self, id: u32) -> Option<String> {
500 + self.metadata.id_to_token(id)
501 + }
502 +
503 + fn is_special_id(&self, token_id: u32) -> bool {
504 + self.metadata.is_special_id(token_id)
505 + }
506 +}
507 +
508 +/// Select the BPE regex pattern for a tiktoken model based on `config.json`.
509 +///
510 +/// Most tiktoken models use the `cl100k_base` regex. Kimi models ship a custom
511 +/// regex in their Python tokenizer implementation; we mirror the explicit
512 +/// `model_type` switch used by Dynamo instead of heuristically parsing Python
513 +/// source files.
514 +fn detect_bpe_pattern(config: &TiktokenModelConfig) -> &'static str {
515 + let model_type = config.effective_model_type();
516 +
517 + match model_type {
518 + Some("kimi" | "kimi_k2" | "kimi_k25" | "deepseek_v3") => KIMI_PATTERN,
519 + _ => CL100K_BASE_PATTERN,
520 + }
521 +}
522 +
523 +#[cfg(test)]
524 +mod tests {
525 + use std::fs;
526 + use std::path::{Path, PathBuf};
527 +
528 + use base64::Engine as _;
529 + use tempfile::TempDir;
530 +
531 + use super::{
532 + CL100K_BASE_PATTERN, KIMI_PATTERN, TiktokenModelConfig, TiktokenTokenizer,
533 + TiktokenTokenizerConfig, detect_bpe_pattern,
534 + };
535 + use crate::Tokenizer;
536 +
537 + macro_rules! config_json {
538 + ($($json:tt)+) => {
539 + serde_json::from_value::<TiktokenModelConfig>(serde_json::json!($($json)+)).unwrap()
540 + };
541 + }
542 +
543 + /// Write a minimal `*.tiktoken` BPE file (one token per byte 0..=255) into
544 + /// `dir` and return its path. The single-byte vocab is enough to
545 + /// exercise the multi-byte / streaming UTF-8 paths without depending on
546 + /// any pretrained tokenizer asset.
547 + fn write_synthetic_bpe_file(dir: &std::path::Path) -> PathBuf {
548 + let mut content = String::new();
549 + for byte in 0u8..=255 {
550 + let b64 = base64::engine::general_purpose::STANDARD.encode([byte]);
551 + content.push_str(&format!("{b64} {}\n", byte as u32));
552 + }
553 + let path = dir.join("test.tiktoken");
554 + fs::write(&path, content).expect("write tiktoken file");
555 + path
556 + }
557 +
558 + /// Write a synthetic `*.tiktoken` file whose base-vocab ranks are
559 + /// sparse/non-contiguous.
560 + ///
561 + /// This reproduces the important edge case for `num_base_tokens`: it must
562 + /// be derived from `max_rank + 1`, not `encoder.len()`, otherwise
563 + /// high-rank base tokens get misclassified as reserved/special ids.
564 + fn write_sparse_rank_bpe_file(dir: &std::path::Path) -> PathBuf {
565 + let mut content = String::new();
566 + for byte in 0u8..=255 {
567 + let b64 = base64::engine::general_purpose::STANDARD.encode([byte]);
568 + content.push_str(&format!("{b64} {}\n", byte as u32));
569 + }
570 +
571 + let high_rank_token = base64::engine::general_purpose::STANDARD.encode(b"SPARSE");
572 + content.push_str(&format!("{high_rank_token} 1000\n"));
573 +
574 + let path = dir.join("sparse-rank.tiktoken");
575 + fs::write(&path, content).expect("write sparse-rank tiktoken file");
576 + path
577 + }
578 +
579 + /// Build a `TiktokenTokenizer` from the synthetic BPE file with no sibling
580 + /// config files, so the constructor takes the
581 + /// `FALLBACK_NUM_RESERVED_SPECIAL_TOKENS` (256) path.
582 + fn explicit_backends(path: &Path) -> Vec<TiktokenTokenizer> {
583 + vec![
584 + TiktokenTokenizer::new_riptoken(path).expect("load riptoken backend"),
585 + TiktokenTokenizer::new_tiktoken_rs(path).expect("load tiktoken-rs backend"),
586 + ]
587 + }
588 +
589 + fn tiktoken_backends() -> (Vec<TiktokenTokenizer>, TempDir) {
590 + let dir = tempfile::tempdir().expect("create temp dir");
591 + let path = write_synthetic_bpe_file(dir.path());
592 + (explicit_backends(&path), dir)
593 + }
594 +
595 + /// Verify that tiktoken decode uses lossy UTF-8 (producing `\u{FFFD}`)
596 + /// rather than returning an error for incomplete multi-byte sequences.
597 + /// This is critical for streaming decode — `DecodeStream` relies on
598 + /// `\u{FFFD}` to detect incomplete characters.
599 + #[test]
600 + fn tiktoken_decode_incomplete_utf8_produces_replacement_char() {
601 + let (backends, _dir) = tiktoken_backends();
602 +
603 + for backend in backends {
604 + let ids = backend.encode("你", false).unwrap();
605 + let full = backend.decode(&ids, false).unwrap();
606 + assert_eq!(full, "你");
607 +
608 + let text_with_multibyte = "Hello你好World";
609 + let all_ids = backend.encode(text_with_multibyte, false).unwrap();
610 + for &id in &all_ids {
611 + let result = backend.decode(&[id], false);
612 + assert!(result.is_ok(), "decode of token {id} should not error");
613 + }
614 + }
615 + }
616 +
617 + /// When `config.json` exposes a `vocab_size`, the reserved-token range must
618 + /// be sized to it rather than to the 256-slot fallback. This is the
619 + /// general (non-Kimi-specific) path: any tiktoken model whose own
620 + /// `config.json` says e.g. `vocab_size = 280` should populate
621 + /// reserved slots for `[num_base_tokens, 280)` and nothing beyond.
622 + #[test]
623 + fn tiktoken_reserved_range_uses_vocab_size_from_config_json() {
624 + let dir = tempfile::tempdir().expect("create temp dir");
625 + let bpe_path = write_synthetic_bpe_file(dir.path());
626 + // num_base_tokens = 256, vocab_size = 280 → reserved range = [256, 280) (24
627 + // slots, smaller than the 256 fallback so we can prove the config value
628 + // is honoured).
629 + fs::write(dir.path().join("config.json"), r#"{"vocab_size": 280}"#)
630 + .expect("write config.json");
631 +
632 + for backend in explicit_backends(&bpe_path) {
633 + // Inside the configured range: reserved placeholder, round-trips both ways.
634 + let in_range_id: u32 = 270;
635 + let placeholder = format!("<|reserved_token_{in_range_id}|>");
636 + assert_eq!(backend.decode(&[in_range_id], false).unwrap(), placeholder);
637 + assert_eq!(
638 + backend.encode(&placeholder, false).unwrap(),
639 + vec![in_range_id]
640 + );
641 + assert_eq!(
642 + backend.id_to_token(in_range_id).as_deref(),
643 + Some(placeholder.as_str())
644 + );
645 +
646 + // Outside the configured range: not registered as a reserved slot — falls
647 + // through to the backend's unknown-id behavior. The point is that we *don't*
648 + // over-populate beyond what the model actually exposes.
649 + let out_of_range_id: u32 = 290;
650 + let out_of_range_placeholder = format!("<|reserved_token_{out_of_range_id}|>");
651 + assert_eq!(backend.decode(&[out_of_range_id], false).unwrap(), "");
652 + assert_eq!(backend.token_to_id(&out_of_range_placeholder), None);
653 + assert_eq!(backend.id_to_token(out_of_range_id), None);
654 + }
655 + }
656 +
657 + /// Sparse/non-contiguous BPE ranks must still count as base-vocab ids.
658 + ///
659 + /// Regression shape:
660 + /// - base vocabulary contains ids 0..=255 and also a normal BPE token at id 1000
661 + /// - if `num_base_tokens` were computed as `encoder.len()` (257), id 1000 would be
662 + /// misclassified as special/reserved and disappear under `skip_special_tokens = true`
663 + #[test]
664 + fn tiktoken_sparse_base_ranks_are_not_misclassified_as_special() {
665 + let dir = tempfile::tempdir().expect("create temp dir");
666 + let bpe_path = write_sparse_rank_bpe_file(dir.path());
667 + fs::write(dir.path().join("config.json"), r#"{"vocab_size": 1002}"#)
668 + .expect("write config.json");
669 +
670 + for backend in explicit_backends(&bpe_path) {
671 + let sparse_id = backend.token_to_id("SPARSE");
672 + assert_eq!(sparse_id, Some(1000));
673 + assert_eq!(backend.id_to_token(1000).as_deref(), Some("SPARSE"));
674 + assert!(!backend.is_special_id(1000));
675 + assert_eq!(backend.decode(&[1000], false).unwrap(), "SPARSE");
676 + assert_eq!(backend.decode(&[1000], true).unwrap(), "SPARSE");
677 + }
678 + }
679 +
680 + /// `skip_special_tokens` must:
681 + /// * keep regular BPE token text unchanged,
682 + /// * drop ids whose `added_tokens_decoder` entry says `"special": true`,
683 + /// * drop reserved-slot placeholder ids (which default to special),
684 + /// * keep ids whose `added_tokens_decoder` entry says `"special": false` — this is how Kimi K2
685 + /// / K2.5 marks tool-call markers and `<think>` / `</think>`.
686 + ///
687 + /// Synthetic backend has `num_base_tokens = 256`. We write a
688 + /// `tokenizer_config.json` that names ids 257 (special) and 258
689 + /// (non-special), and a `config.json` with `vocab_size` covering both.
690 + /// Id 259 stays a default reserved placeholder (special).
691 + #[test]
692 + fn tiktoken_skip_special_tokens_filters_special_but_keeps_non_special_added_tokens() {
693 + let dir = tempfile::tempdir().expect("create temp dir");
694 + let bpe_path = write_synthetic_bpe_file(dir.path());
695 + fs::write(
696 + dir.path().join("tokenizer_config.json"),
697 + r#"{
698 + "added_tokens_decoder": {
699 + "257": { "content": "<|im_end|>", "special": true },
700 + "258": { "content": "<|tool_call_begin|>", "special": false }
701 + }
702 + }"#,
703 + )
704 + .expect("write tokenizer_config.json");
705 + fs::write(dir.path().join("config.json"), r#"{"vocab_size": 260}"#)
706 + .expect("write config.json");
707 +
708 + for backend in explicit_backends(&bpe_path) {
709 + // Resolve the BPE ids for "Hi" so we can interleave them with special-token
710 + // ids.
711 + let h = backend.encode("H", false).unwrap()[0];
712 + let i = backend.encode("i", false).unwrap()[0];
713 +
714 + let special_id: u32 = 257; // <|im_end|>
715 + let non_special_id: u32 = 258; // <|tool_call_begin|>
716 + let reserved_id: u32 = 259; // default <|reserved_token_259|> placeholder
717 +
718 + let ids = vec![h, special_id, i, non_special_id, reserved_id];
719 +
720 + // skip_special_tokens = false: everything is rendered as-is.
721 + let kept = backend.decode(&ids, false).unwrap();
722 + assert_eq!(
723 + kept,
724 + "H<|im_end|>i<|tool_call_begin|><|reserved_token_259|>"
725 + );
726 +
727 + // skip_special_tokens = true: special token (257) and reserved placeholder
728 + // (259) are dropped; the non-special added token (258) survives.
729 + let stripped = backend.decode(&ids, true).unwrap();
730 + assert_eq!(stripped, "Hi<|tool_call_begin|>");
731 + }
732 + }
733 +
734 + /// `vocab_size` may live under `text_config` for composite (e.g.
735 + /// multimodal) configs.
736 + #[test]
737 + fn tiktoken_reserved_range_reads_text_config_vocab_size() {
738 + let dir = tempfile::tempdir().expect("create temp dir");
739 + let bpe_path = write_synthetic_bpe_file(dir.path());
740 + fs::write(
741 + dir.path().join("config.json"),
742 + r#"{"text_config": {"vocab_size": 270}}"#,
743 + )
744 + .expect("write config.json");
745 +
746 + for backend in explicit_backends(&bpe_path) {
747 + let in_range_id: u32 = 260;
748 + let placeholder = format!("<|reserved_token_{in_range_id}|>");
749 + assert_eq!(backend.decode(&[in_range_id], false).unwrap(), placeholder);
750 +
751 + // Just outside the nested vocab_size — should not be registered.
752 + assert_eq!(backend.decode(&[270], false).unwrap(), "");
753 + }
754 + }
755 +
756 + #[test]
757 + fn tiktoken_detects_kimi_pattern_from_model_type() {
758 + let kimi = config_json!({ "model_type": "kimi_k25" });
759 + let baseten_kimi = config_json!({ "model_type": "deepseek_v3" });
760 + let nested_kimi = config_json!({
761 + "model_type": "composite_wrapper",
762 + "text_config": { "model_type": "kimi_k2" }
763 + });
764 + let generic = config_json!({ "model_type": "gpt2" });
765 + let nested_generic = config_json!({
766 + "model_type": "composite_wrapper",
767 + "text_config": { "model_type": "gpt2" }
768 + });
769 + let missing = config_json!({ "text_config": {} });
770 +
771 + assert_eq!(detect_bpe_pattern(&kimi), KIMI_PATTERN);
772 + assert_eq!(detect_bpe_pattern(&baseten_kimi), KIMI_PATTERN);
773 + assert_eq!(detect_bpe_pattern(&nested_kimi), CL100K_BASE_PATTERN);
774 + assert_eq!(detect_bpe_pattern(&generic), CL100K_BASE_PATTERN);
775 + assert_eq!(detect_bpe_pattern(&nested_generic), CL100K_BASE_PATTERN);
776 + assert_eq!(detect_bpe_pattern(&missing), CL100K_BASE_PATTERN);
777 + }
778 +
779 + #[test]
780 + fn tiktoken_reads_model_type_from_text_config_when_top_level_missing() {
781 + let nested_only = config_json!({
782 + "text_config": { "model_type": "kimi_k2" }
783 + });
784 + let direct_and_nested = config_json!({
785 + "model_type": "kimi_k25",
786 + "text_config": { "model_type": "kimi_k2" }
787 + });
788 + let missing = config_json!({
789 + "text_config": {}
790 + });
791 +
792 + assert_eq!(nested_only.effective_model_type(), Some("kimi_k2"));
793 + assert_eq!(direct_and_nested.effective_model_type(), Some("kimi_k25"));
794 + assert_eq!(missing.effective_model_type(), None);
795 + }
796 +
797 + #[test]
798 + fn tiktoken_tokenizer_config_models_added_tokens_decoder() {
799 + let config: TiktokenTokenizerConfig = serde_json::from_value(serde_json::json!({
800 + "added_tokens_decoder": {
801 + "257": { "content": "<think>" },
802 + "258": { "content": "</think>", "special": true }
803 + }
804 + }))
805 + .unwrap();
806 +
807 + let added_tokens = config.added_tokens_decoder;
808 + assert_eq!(added_tokens.len(), 2);
809 + assert_eq!(
810 + added_tokens.get(&257).map(|t| t.content.as_str()),
811 + Some("<think>")
812 + );
813 + assert_eq!(added_tokens.get(&257).map(|t| t.special), Some(false));
814 + assert_eq!(
815 + added_tokens.get(&258).map(|t| (t.content.as_str(), t.special)),
816 + Some(("</think>", true))
817 + );
818 + }
819 +
820 + /// Reserved token ids in `[num_base_tokens, num_base_tokens + 256)` must
821 + /// decode to their placeholder name (matching `tokenization_kimi.py`'s
822 + /// `<|reserved_token_{i}|>` format), even when the source
823 + /// `tokenizer_config.json` does not list them in `added_tokens_decoder`.
824 + ///
825 + /// In our synthetic backend `num_base_tokens = 256` (256 single-byte BPE
826 + /// tokens), so the reserved range is `[256, 512)`. Picking id 300 —
827 + /// well inside that range and absent from any `added_tokens_decoder` —
828 + /// should round-trip both ways.
829 + #[test]
830 + fn tiktoken_reserved_token_round_trip() {
831 + let (backends, _dir) = tiktoken_backends();
832 +
833 + for backend in backends {
834 + let reserved_id: u32 = 300;
835 + let placeholder = format!("<|reserved_token_{reserved_id}|>");
836 +
837 + let decoded = backend.decode(&[reserved_id], false).unwrap();
838 + assert_eq!(decoded, placeholder);
839 +
840 + // The placeholder name should also encode back to the same single id, since
841 + // the constructor registers it as a special token with `CoreBPE`.
842 + let encoded = backend.encode(&placeholder, false).unwrap();
843 + assert_eq!(encoded, vec![reserved_id]);
844 +
845 + assert_eq!(backend.token_to_id(&placeholder), Some(reserved_id));
846 + assert_eq!(
847 + backend.id_to_token(reserved_id).as_deref(),
848 + Some(placeholder.as_str())
849 + );
850 + }
851 + }
852 +
853 + /// Decoding a token id that is beyond even the reserved range must not
854 + /// panic — it falls through to the warn-and-skip backstop instead of
855 + /// crashing the worker thread.
856 + #[test]
857 + fn tiktoken_rs_decode_unknown_token_id_does_not_panic() {
858 + let dir = tempfile::tempdir().expect("create temp dir");
859 + let path = write_synthetic_bpe_file(dir.path());
860 + let backend = TiktokenTokenizer::new_tiktoken_rs(&path).expect("load tiktoken-rs backend");
861 +
862 + // ID well above num_base_tokens (256) + reserved (256) = 512 — guaranteed
863 + // unknown.
864 + let unknown_id: u32 = 999_999;
865 + let result = backend.decode(&[unknown_id], false);
866 + assert_eq!(result.unwrap(), "");
867 +
868 + // Mixed: known bytes for "Hi" surrounding an unknown id should yield just "Hi".
869 + let h = backend.encode("H", false).unwrap()[0];
870 + let i = backend.encode("i", false).unwrap()[0];
871 + let result = backend.decode(&[h, unknown_id, i], false).unwrap();
872 + assert_eq!(result, "Hi");
873 + }
874 +
875 + #[test]
876 + fn riptoken_decode_unknown_token_id_does_not_panic() {
877 + let dir = tempfile::tempdir().expect("create temp dir");
878 + let path = write_synthetic_bpe_file(dir.path());
879 + let backend = TiktokenTokenizer::new_riptoken(&path).expect("load riptoken backend");
880 +
881 + let unknown_id: u32 = 999_999;
882 + assert_eq!(backend.decode(&[unknown_id], false).unwrap(), "");
883 +
884 + let h = backend.encode("H", false).unwrap()[0];
885 + let i = backend.encode("i", false).unwrap()[0];
886 + assert_eq!(backend.decode(&[h, unknown_id, i], false).unwrap(), "Hi");
887 + }
888 +
889 + /// Streaming decode of CJK text through tiktoken should produce the
890 + /// original text without errors, even though individual tokens may
891 + /// represent partial UTF-8 byte sequences.
892 + #[test]
893 + fn tiktoken_streaming_decode_multibyte() {
894 + let (backends, _dir) = tiktoken_backends();
895 + for backend in backends {
896 + let text = "你好世界"; // 4 CJK characters
897 + let ids = backend.encode(text, false).unwrap();
898 +
899 + let mut decoder = backend.create_decode_stream(&[], false, 0);
900 + let mut output = String::new();
901 + for &id in &ids {
902 + decoder.push_token(id).unwrap();
903 + if let Some(chunk) = decoder.next_chunk() {
904 + output.push_str(&chunk);
905 + }
906 + }
907 + let (last_chunk, full_text) = decoder.flush(None).unwrap();
908 + if let Some(chunk) = last_chunk {
909 + output.push_str(&chunk);
910 + }
911 +
912 + assert_eq!(output, text);
913 + assert_eq!(full_text, text);
914 + }
915 + }
916 +
917 + /// Mixed ASCII and multi-byte text should stream correctly through
918 + /// tiktoken.
919 + #[test]
920 + fn tiktoken_streaming_decode_mixed_ascii_and_multibyte() {
921 + let (backends, _dir) = tiktoken_backends();
922 + for backend in backends {
923 + let text = "Hello 你好 World 🌍";
924 + let ids = backend.encode(text, false).unwrap();
925 +
926 + let mut decoder = backend.create_decode_stream(&[], false, 0);
927 + let mut output = String::new();
928 + for &id in &ids {
929 + decoder.push_token(id).unwrap();
930 + if let Some(chunk) = decoder.next_chunk() {
931 + output.push_str(&chunk);
932 + }
933 + }
934 + let (last_chunk, full_text) = decoder.flush(None).unwrap();
935 + if let Some(chunk) = last_chunk {
936 + output.push_str(&chunk);
937 + }
938 +
939 + assert_eq!(output, text);
940 + assert_eq!(full_text, text);
941 + }
942 + }
943 +
944 + #[test]
945 + fn tiktoken_token_to_id_resolves_added_special_tokens() {
946 + let dir = tempfile::tempdir().expect("create temp dir");
947 + let bpe_path = write_synthetic_bpe_file(dir.path());
948 + fs::write(
949 + dir.path().join("tokenizer_config.json"),
950 + r#"{
951 + "added_tokens_decoder": {
952 + "257": { "content": "<think>", "special": false },
953 + "258": { "content": "</think>", "special": false }
954 + }
955 + }"#,
956 + )
957 + .expect("write tokenizer_config.json");
958 + fs::write(dir.path().join("config.json"), r#"{"vocab_size": 259}"#)
959 + .expect("write config.json");
960 +
961 + for backend in explicit_backends(&bpe_path) {
962 + assert_eq!(backend.token_to_id("<think>"), Some(257));
963 + assert_eq!(backend.token_to_id("</think>"), Some(258));
964 + assert_eq!(backend.id_to_token(257).as_deref(), Some("<think>"));
965 + assert_eq!(backend.id_to_token(258).as_deref(), Some("</think>"));
966 + assert_eq!(
967 + backend.decode(&[257, 258], true).unwrap(),
968 + "<think></think>"
969 + );
970 + }
971 + }
972 +
973 + #[test]
974 + fn riptoken_token_to_id_uses_encode_single_token_path() {
975 + let dir = tempfile::tempdir().expect("create temp dir");
976 + let bpe_path = write_synthetic_bpe_file(dir.path());
977 + fs::write(
978 + dir.path().join("tokenizer_config.json"),
979 + r#"{
980 + "added_tokens_decoder": {
981 + "257": { "content": "<think>", "special": false }
982 + }
983 + }"#,
984 + )
985 + .expect("write tokenizer_config.json");
986 + fs::write(dir.path().join("config.json"), r#"{"vocab_size": 258}"#)
987 + .expect("write config.json");
988 + let backend = TiktokenTokenizer::new_riptoken(&bpe_path).expect("load riptoken backend");
989 +
990 + assert_eq!(backend.token_to_id("H"), Some(b'H' as u32));
991 + assert_eq!(backend.id_to_token(b'H' as u32).as_deref(), Some("H"));
992 + assert_eq!(backend.token_to_id("<think>"), Some(257));
993 + assert_eq!(backend.id_to_token(257).as_deref(), Some("<think>"));
994 + }
995 +
996 + #[test]
997 + fn tiktoken_rs_token_to_id_handles_unknown_special_like_text_without_panicking() {
998 + let dir = tempfile::tempdir().expect("create temp dir");
999 + let path = write_synthetic_bpe_file(dir.path());
1000 + let backend = TiktokenTokenizer::new_tiktoken_rs(&path).expect("load tiktoken-rs backend");
1001 +
1002 + assert_eq!(backend.token_to_id("<|definitely_not_registered|>"), None);
1003 + }
1004 +
1005 + #[test]
1006 + fn riptoken_token_to_id_handles_unknown_special_like_text_without_panicking() {
1007 + let dir = tempfile::tempdir().expect("create temp dir");
1008 + let path = write_synthetic_bpe_file(dir.path());
1009 + let backend = TiktokenTokenizer::new_riptoken(&path).expect("load riptoken backend");
1010 +
1011 + assert_eq!(backend.token_to_id("<|definitely_not_registered|>"), None);
1012 + }
1013 +}
rust/src/tool-parser/Cargo.toml new
+75
@@ -0,0 +1,75 @@
1 +[package]
2 +name = "vllm-tool-parser"
3 +version.workspace = true
4 +edition.workspace = true
5 +license.workspace = true
6 +
7 +[features]
8 +test-util = []
9 +
10 +[dependencies]
11 +serde.workspace = true
12 +serde_json.workspace = true
13 +thiserror.workspace = true
14 +thiserror-ext.workspace = true
15 +winnow.workspace = true
16 +
17 +[dev-dependencies]
18 +criterion.workspace = true
19 +expect-test.workspace = true
20 +futures.workspace = true
21 +openai-protocol.workspace = true
22 +tool-parser.workspace = true
23 +
24 +[[bench]]
25 +name = "deepseek_v3"
26 +harness = false
27 +required-features = ["test-util"]
28 +
29 +[[bench]]
30 +name = "deepseek_v31"
31 +harness = false
32 +required-features = ["test-util"]
33 +
34 +[[bench]]
35 +name = "deepseek_v32"
36 +harness = false
37 +required-features = ["test-util"]
38 +
39 +[[bench]]
40 +name = "qwen3_coder"
41 +harness = false
42 +required-features = ["test-util"]
43 +
44 +[[bench]]
45 +name = "qwen3_xml"
46 +harness = false
47 +required-features = ["test-util"]
48 +
49 +[[bench]]
50 +name = "llama3_json"
51 +harness = false
52 +required-features = ["test-util"]
53 +
54 +[[bench]]
55 +name = "minimax_m2"
56 +harness = false
57 +required-features = ["test-util"]
58 +
59 +[[bench]]
60 +name = "glm45_moe"
61 +harness = false
62 +required-features = ["test-util"]
63 +
64 +[[bench]]
65 +name = "kimi_k2"
66 +harness = false
67 +required-features = ["test-util"]
68 +
69 +[[bench]]
70 +name = "gemma4"
71 +harness = false
72 +required-features = ["test-util"]
73 +
74 +[lints]
75 +workspace = true
rust/src/tool-parser/benches/deepseek_v3.rs new
+130
@@ -0,0 +1,130 @@
1 +use std::time::Duration;
2 +
3 +use criterion::{BatchSize, Criterion, Throughput, black_box, criterion_group, criterion_main};
4 +use tool_parser::parsers::DeepSeekParser as ExternalDeepSeekParser;
5 +use vllm_tool_parser::test_utils::{split_by_chars, test_tools};
6 +use vllm_tool_parser::{DeepSeekV3ToolParser, Tool, ToolParser};
7 +
8 +mod utils;
9 +use utils::{feed_external_parser, feed_parser, openai_tools};
10 +
11 +const CHUNK_CHARS: usize = 7;
12 +const LONG_NORMAL_TEXT_REPEATS: usize = 2048;
13 +
14 +fn mixed_fixture() -> String {
15 + concat!(
16 + "I will check two cities before answering.\n",
17 + "<｜tool▁calls▁begin｜>",
18 + "<｜tool▁call▁begin｜>function<｜tool▁sep｜>get_weather\n",
19 + "```json\n",
20 + "{\"location\":\"Hangzhou\",\"days\":3}",
21 + "\n```<｜tool▁call▁end｜>",
22 + "<｜tool▁call▁begin｜>function<｜tool▁sep｜>get_weather\n",
23 + "```json\n",
24 + "{\"location\":\"San Francisco\",\"days\":2}",
25 + "\n```<｜tool▁call▁end｜>",
26 + "<｜tool▁calls▁end｜>",
27 + )
28 + .to_string()
29 +}
30 +
31 +fn long_normal_text_fixture() -> String {
32 + let line = "This is ordinary assistant text with no DeepSeek V3 tool markers at all.\n";
33 + line.repeat(LONG_NORMAL_TEXT_REPEATS)
34 +}
35 +
36 +fn native_parser(tools: &[Tool]) -> Box<dyn ToolParser> {
37 + DeepSeekV3ToolParser::create(tools).expect("DeepSeek V3 parser should initialize")
38 +}
39 +
40 +fn run_stream_group(
41 + c: &mut Criterion,
42 + name: &str,
43 + tools: &[Tool],
44 + text: &str,
45 + chunk_chars: usize,
46 + expected_normal_text: &str,
47 + expected_native_calls_len: usize,
48 +) {
49 + let chunks = split_by_chars(text, chunk_chars);
50 + let openai_tools = openai_tools(tools);
51 +
52 + let mut group = c.benchmark_group(name);
53 + group.sample_size(50);
54 + group.warm_up_time(Duration::from_millis(300));
55 + group.measurement_time(Duration::from_secs(2));
56 + group.throughput(Throughput::Bytes(text.len() as u64));
57 +
58 + group.bench_function("native_reuse_parser", |b| {
59 + let mut parser = native_parser(tools);
60 + b.iter(|| {
61 + let result = feed_parser(&mut *parser, black_box(&chunks));
62 + debug_assert_eq!(result.0, expected_normal_text);
63 + debug_assert_eq!(result.1, expected_native_calls_len);
64 + black_box(result);
65 + })
66 + });
67 +
68 + group.bench_function("native_create_parser", |b| {
69 + b.iter_batched(
70 + || native_parser(tools),
71 + |mut parser| {
72 + let result = feed_parser(&mut *parser, black_box(&chunks));
73 + debug_assert_eq!(result.0, expected_normal_text);
74 + debug_assert_eq!(result.1, expected_native_calls_len);
75 + black_box(result);
76 + },
77 + BatchSize::SmallInput,
78 + )
79 + });
80 +
81 + group.bench_function("external_reuse_parser", |b| {
82 + let mut parser = ExternalDeepSeekParser::new();
83 + b.iter(|| {
84 + let result = feed_external_parser(&mut parser, &openai_tools, black_box(&chunks));
85 + black_box(result);
86 + })
87 + });
88 +
89 + group.bench_function("external_create_parser", |b| {
90 + b.iter_batched(
91 + ExternalDeepSeekParser::new,
92 + |mut parser| {
93 + let result = feed_external_parser(&mut parser, &openai_tools, black_box(&chunks));
94 + black_box(result);
95 + },
96 + BatchSize::SmallInput,
97 + )
98 + });
99 +
100 + group.finish();
101 +}
102 +
103 +fn bench_deepseek_v3(c: &mut Criterion) {
104 + let tools = test_tools();
105 + let mixed_text = mixed_fixture();
106 + let long_normal_text = long_normal_text_fixture();
107 +
108 + run_stream_group(
109 + c,
110 + "deepseek_v3/mixed_text_tool_call",
111 + &tools,
112 + &mixed_text,
113 + CHUNK_CHARS,
114 + "I will check two cities before answering.\n",
115 + 2,
116 + );
117 +
118 + run_stream_group(
119 + c,
120 + "deepseek_v3/long_normal_text",
121 + &tools,
122 + &long_normal_text,
123 + CHUNK_CHARS,
124 + &long_normal_text,
125 + 0,
126 + );
127 +}
128 +
129 +criterion_group!(benches, bench_deepseek_v3);
130 +criterion_main!(benches);
rust/src/tool-parser/benches/deepseek_v31.rs new
+128
@@ -0,0 +1,128 @@
1 +use std::time::Duration;
2 +
3 +use criterion::{BatchSize, Criterion, Throughput, black_box, criterion_group, criterion_main};
4 +use tool_parser::parsers::DeepSeek31Parser as ExternalDeepSeek31Parser;
5 +use vllm_tool_parser::test_utils::{split_by_chars, test_tools};
6 +use vllm_tool_parser::{DeepSeekV31ToolParser, Tool, ToolParser};
7 +
8 +mod utils;
9 +use utils::{feed_external_parser, feed_parser, openai_tools};
10 +
11 +const CHUNK_CHARS: usize = 7;
12 +const LONG_NORMAL_TEXT_REPEATS: usize = 2048;
13 +
14 +fn mixed_fixture() -> String {
15 + concat!(
16 + "I will check two cities before answering.\n",
17 + "<｜tool▁calls▁begin｜>",
18 + "<｜tool▁call▁begin｜>get_weather<｜tool▁sep｜>",
19 + "{\"location\":\"Hangzhou\",\"days\":3}",
20 + "<｜tool▁call▁end｜>",
21 + "<｜tool▁call▁begin｜>get_weather<｜tool▁sep｜>",
22 + "{\"location\":\"San Francisco\",\"days\":2}",
23 + "<｜tool▁call▁end｜>",
24 + "<｜tool▁calls▁end｜>",
25 + )
26 + .to_string()
27 +}
28 +
29 +fn long_normal_text_fixture() -> String {
30 + let line = "This is ordinary assistant text with no DeepSeek V3.1 tool markers at all.\n";
31 + line.repeat(LONG_NORMAL_TEXT_REPEATS)
32 +}
33 +
34 +fn native_parser(tools: &[Tool]) -> Box<dyn ToolParser> {
35 + DeepSeekV31ToolParser::create(tools).expect("DeepSeek V3.1 parser should initialize")
36 +}
37 +
38 +fn run_stream_group(
39 + c: &mut Criterion,
40 + name: &str,
41 + tools: &[Tool],
42 + text: &str,
43 + chunk_chars: usize,
44 + expected_normal_text: &str,
45 + expected_native_calls_len: usize,
46 +) {
47 + let chunks = split_by_chars(text, chunk_chars);
48 + let openai_tools = openai_tools(tools);
49 +
50 + let mut group = c.benchmark_group(name);
51 + group.sample_size(50);
52 + group.warm_up_time(Duration::from_millis(300));
53 + group.measurement_time(Duration::from_secs(2));
54 + group.throughput(Throughput::Bytes(text.len() as u64));
55 +
56 + group.bench_function("native_reuse_parser", |b| {
57 + let mut parser = native_parser(tools);
58 + b.iter(|| {
59 + let result = feed_parser(&mut *parser, black_box(&chunks));
60 + debug_assert_eq!(result.0, expected_normal_text);
61 + debug_assert_eq!(result.1, expected_native_calls_len);
62 + black_box(result);
63 + })
64 + });
65 +
66 + group.bench_function("native_create_parser", |b| {
67 + b.iter_batched(
68 + || native_parser(tools),
69 + |mut parser| {
70 + let result = feed_parser(&mut *parser, black_box(&chunks));
71 + debug_assert_eq!(result.0, expected_normal_text);
72 + debug_assert_eq!(result.1, expected_native_calls_len);
73 + black_box(result);
74 + },
75 + BatchSize::SmallInput,
76 + )
77 + });
78 +
79 + group.bench_function("external_reuse_parser", |b| {
80 + let mut parser = ExternalDeepSeek31Parser::new();
81 + b.iter(|| {
82 + let result = feed_external_parser(&mut parser, &openai_tools, black_box(&chunks));
83 + black_box(result);
84 + })
85 + });
86 +
87 + group.bench_function("external_create_parser", |b| {
88 + b.iter_batched(
89 + ExternalDeepSeek31Parser::new,
90 + |mut parser| {
91 + let result = feed_external_parser(&mut parser, &openai_tools, black_box(&chunks));
92 + black_box(result);
93 + },
94 + BatchSize::SmallInput,
95 + )
96 + });
97 +
98 + group.finish();
99 +}
100 +
101 +fn bench_deepseek_v31(c: &mut Criterion) {
102 + let tools = test_tools();
103 + let mixed_text = mixed_fixture();
104 + let long_normal_text = long_normal_text_fixture();
105 +
106 + run_stream_group(
107 + c,
108 + "deepseek_v31/mixed_text_tool_call",
109 + &tools,
110 + &mixed_text,
111 + CHUNK_CHARS,
112 + "I will check two cities before answering.\n",
113 + 2,
114 + );
115 +
116 + run_stream_group(
117 + c,
118 + "deepseek_v31/long_normal_text",
119 + &tools,
120 + &long_normal_text,
121 + CHUNK_CHARS,
122 + &long_normal_text,
123 + 0,
124 + );
125 +}
126 +
127 +criterion_group!(benches, bench_deepseek_v31);
128 +criterion_main!(benches);
rust/src/tool-parser/benches/deepseek_v32.rs new
+113
@@ -0,0 +1,113 @@
1 +use std::time::Duration;
2 +
3 +use criterion::{BatchSize, Criterion, Throughput, black_box, criterion_group, criterion_main};
4 +use vllm_tool_parser::test_utils::{split_by_chars, test_tools};
5 +use vllm_tool_parser::{DeepSeekV32ToolParser, Tool, ToolParser};
6 +
7 +mod utils;
8 +use utils::feed_parser;
9 +
10 +const CHUNK_CHARS: usize = 7;
11 +const LONG_NORMAL_TEXT_REPEATS: usize = 2048;
12 +
13 +fn mixed_fixture() -> String {
14 + concat!(
15 + "I will check two cities before answering.\n",
16 + "<｜DSML｜function_calls>\n",
17 + "<｜DSML｜invoke name=\"get_weather\">\n",
18 + "<｜DSML｜parameter name=\"location\" string=\"true\">Hangzhou</｜DSML｜parameter>\n",
19 + "<｜DSML｜parameter name=\"date\" string=\"true\">2026-04-28</｜DSML｜parameter>\n",
20 + "<｜DSML｜parameter name=\"unit\" string=\"true\">celsius</｜DSML｜parameter>\n",
21 + "<｜DSML｜parameter name=\"days\" string=\"false\">3</｜DSML｜parameter>\n",
22 + "</｜DSML｜invoke>\n",
23 + "<｜DSML｜invoke name=\"get_weather\">\n",
24 + "<｜DSML｜parameter name=\"location\" string=\"true\">San Francisco</｜DSML｜parameter>\n",
25 + "<｜DSML｜parameter name=\"date\" string=\"true\">2026-04-28</｜DSML｜parameter>\n",
26 + "<｜DSML｜parameter name=\"unit\" string=\"true\">fahrenheit</｜DSML｜parameter>\n",
27 + "<｜DSML｜parameter name=\"days\" string=\"false\">2</｜DSML｜parameter>\n",
28 + "</｜DSML｜invoke>\n",
29 + "</｜DSML｜function_calls>",
30 + )
31 + .to_string()
32 +}
33 +
34 +fn long_normal_text_fixture() -> String {
35 + let line = "This is ordinary assistant text with no DSML tool markers at all.\n";
36 + line.repeat(LONG_NORMAL_TEXT_REPEATS)
37 +}
38 +
39 +fn parser(tools: &[Tool]) -> Box<dyn ToolParser> {
40 + DeepSeekV32ToolParser::create(tools).expect("DeepSeek V3.2 parser should initialize")
41 +}
42 +
43 +fn run_stream_group(
44 + c: &mut Criterion,
45 + name: &str,
46 + tools: &[Tool],
47 + text: &str,
48 + chunk_chars: usize,
49 + expected_normal_text: &str,
50 + expected_calls_len: usize,
51 +) {
52 + let chunks = split_by_chars(text, chunk_chars);
53 +
54 + let mut group = c.benchmark_group(name);
55 + group.sample_size(50);
56 + group.warm_up_time(Duration::from_millis(300));
57 + group.measurement_time(Duration::from_secs(2));
58 + group.throughput(Throughput::Bytes(text.len() as u64));
59 +
60 + group.bench_function("reuse_parser", |b| {
61 + let mut parser = parser(tools);
62 + b.iter(|| {
63 + let result = feed_parser(&mut *parser, black_box(&chunks));
64 + debug_assert_eq!(result.0, expected_normal_text);
65 + debug_assert_eq!(result.1, expected_calls_len);
66 + black_box(result);
67 + })
68 + });
69 +
70 + group.bench_function("create_parser", |b| {
71 + b.iter_batched(
72 + || parser(tools),
73 + |mut parser| {
74 + let result = feed_parser(&mut *parser, black_box(&chunks));
75 + debug_assert_eq!(result.0, expected_normal_text);
76 + debug_assert_eq!(result.1, expected_calls_len);
77 + black_box(result);
78 + },
79 + BatchSize::SmallInput,
80 + )
81 + });
82 +
83 + group.finish();
84 +}
85 +
86 +fn bench_deepseek_v32(c: &mut Criterion) {
87 + let tools = test_tools();
88 + let mixed_text = mixed_fixture();
89 + let long_normal_text = long_normal_text_fixture();
90 +
91 + run_stream_group(
92 + c,
93 + "deepseek_v32/mixed_text_tool_call",
94 + &tools,
95 + &mixed_text,
96 + CHUNK_CHARS,
97 + "I will check two cities before answering.\n",
98 + 2,
99 + );
100 +
101 + run_stream_group(
102 + c,
103 + "deepseek_v32/long_normal_text",
104 + &tools,
105 + &long_normal_text,
106 + CHUNK_CHARS,
107 + &long_normal_text,
108 + 0,
109 + );
110 +}
111 +
112 +criterion_group!(benches, bench_deepseek_v32);
113 +criterion_main!(benches);
rust/src/tool-parser/benches/gemma4.rs new
+125
@@ -0,0 +1,125 @@
1 +use std::time::Duration;
2 +
3 +use criterion::{BatchSize, Criterion, Throughput, black_box, criterion_group, criterion_main};
4 +use vllm_tool_parser::test_utils::{split_by_chars, test_tools};
5 +use vllm_tool_parser::{Gemma4ToolParser, Tool, ToolParser};
6 +
7 +mod utils;
8 +use utils::feed_parser;
9 +
10 +const CHUNK_CHARS: usize = 7;
11 +const LONG_NORMAL_TEXT_REPEATS: usize = 2048;
12 +
13 +fn mixed_fixture() -> String {
14 + concat!(
15 + "I will inspect the data before answering.\n",
16 + "<|tool_call>",
17 + "call:convert{",
18 + "whole:114.514,",
19 + "flag:true,",
20 + "empty:<|\"|><|\"|>,",
21 + "payload:{",
22 + "name:<|\"|>demo<|\"|>,",
23 + "count:42,",
24 + "enabled:false,",
25 + "missing:null,",
26 + "nested:{level:2,label:<|\"|>deep<|\"|>},",
27 + "tags:[<|\"|>red<|\"|>,<|\"|>blue<|\"|>,3,true,null,{kind:<|\"|>leaf<|\"|>}]",
28 + "},",
29 + "items:[",
30 + "<|\"|>alpha<|\"|>,",
31 + "{key:<|\"|>value<|\"|>,score:0.75},",
32 + "[1,2,3]",
33 + "]",
34 + "}",
35 + "<tool_call|>",
36 + "<|tool_call>",
37 + "call:update_record{",
38 + "data:{id:7,active:true,notes:[<|\"|>keep<|\"|>,<|\"|>review<|\"|>]}",
39 + "}",
40 + "<tool_call|>",
41 + " Finished.",
42 + )
43 + .to_string()
44 +}
45 +
46 +fn long_normal_text_fixture() -> String {
47 + let line = "This is ordinary assistant text with no Gemma4 tool markers at all.\n";
48 + line.repeat(LONG_NORMAL_TEXT_REPEATS)
49 +}
50 +
51 +fn parser(tools: &[Tool]) -> Box<dyn ToolParser> {
52 + Gemma4ToolParser::create(tools).expect("Gemma4 parser should initialize")
53 +}
54 +
55 +fn run_stream_group(
56 + c: &mut Criterion,
57 + name: &str,
58 + tools: &[Tool],
59 + text: &str,
60 + chunk_chars: usize,
61 + expected_normal_text: &str,
62 + expected_calls_len: usize,
63 +) {
64 + let chunks = split_by_chars(text, chunk_chars);
65 +
66 + let mut group = c.benchmark_group(name);
67 + group.sample_size(50);
68 + group.warm_up_time(Duration::from_millis(300));
69 + group.measurement_time(Duration::from_secs(2));
70 + group.throughput(Throughput::Bytes(text.len() as u64));
71 +
72 + group.bench_function("reuse_parser", |b| {
73 + let mut parser = parser(tools);
74 + b.iter(|| {
75 + let result = feed_parser(&mut *parser, black_box(&chunks));
76 + debug_assert_eq!(result.0, expected_normal_text);
77 + debug_assert_eq!(result.1, expected_calls_len);
78 + black_box(result);
79 + })
80 + });
81 +
82 + group.bench_function("create_parser", |b| {
83 + b.iter_batched(
84 + || parser(tools),
85 + |mut parser| {
86 + let result = feed_parser(&mut *parser, black_box(&chunks));
87 + debug_assert_eq!(result.0, expected_normal_text);
88 + debug_assert_eq!(result.1, expected_calls_len);
89 + black_box(result);
90 + },
91 + BatchSize::SmallInput,
92 + )
93 + });
94 +
95 + group.finish();
96 +}
97 +
98 +fn bench_gemma4(c: &mut Criterion) {
99 + let tools = test_tools();
100 + let mixed_text = mixed_fixture();
101 + let long_normal_text = long_normal_text_fixture();
102 +
103 + run_stream_group(
104 + c,
105 + "gemma4/mixed_complex_tool_call",
106 + &tools,
107 + &mixed_text,
108 + CHUNK_CHARS,
109 + "I will inspect the data before answering.\n Finished.",
110 + 2,
111 + );
112 +
113 + run_stream_group(
114 + c,
115 + "gemma4/long_normal_text",
116 + &tools,
117 + &long_normal_text,
118 + CHUNK_CHARS,
119 + &long_normal_text,
120 + 0,
121 + );
122 +}
123 +
124 +criterion_group!(benches, bench_gemma4);
125 +criterion_main!(benches);
rust/src/tool-parser/benches/glm45_moe.rs new
+210
@@ -0,0 +1,210 @@
1 +use std::time::Duration;
2 +
3 +use criterion::{BatchSize, Criterion, Throughput, black_box, criterion_group, criterion_main};
4 +use tool_parser::parsers::Glm4MoeParser as ExternalGlm4MoeParser;
5 +use vllm_tool_parser::test_utils::{split_by_chars, test_tools};
6 +use vllm_tool_parser::{Glm45MoeToolParser, Glm47MoeToolParser, Tool, ToolParser};
7 +
8 +mod utils;
9 +use utils::{feed_external_parser, feed_parser, openai_tools};
10 +
11 +const GLM45_PARSER_NAME: &str = "glm45";
12 +const GLM47_PARSER_NAME: &str = "glm47";
13 +const CHUNK_CHARS: usize = 7;
14 +const LONG_NORMAL_TEXT_REPEATS: usize = 2048;
15 +
16 +fn glm45_mixed_fixture() -> String {
17 + concat!(
18 + "I will check two cities before answering.\n",
19 + "<tool_call>get_weather\n",
20 + "<arg_key>city</arg_key>\n",
21 + "<arg_value>Hangzhou</arg_value>\n",
22 + "<arg_key>date</arg_key>\n",
23 + "<arg_value>2026-05-07</arg_value>\n",
24 + "<arg_key>unit</arg_key>\n",
25 + "<arg_value>celsius</arg_value>\n",
26 + "<arg_key>days</arg_key>\n",
27 + "<arg_value>3</arg_value>\n",
28 + "</tool_call>\n",
29 + "<tool_call>get_weather\n",
30 + "<arg_key>city</arg_key>\n",
31 + "<arg_value>San Francisco</arg_value>\n",
32 + "<arg_key>date</arg_key>\n",
33 + "<arg_value>2026-05-07</arg_value>\n",
34 + "<arg_key>unit</arg_key>\n",
35 + "<arg_value>fahrenheit</arg_value>\n",
36 + "<arg_key>days</arg_key>\n",
37 + "<arg_value>2</arg_value>\n",
38 + "</tool_call>",
39 + )
40 + .to_string()
41 +}
42 +
43 +fn glm47_mixed_fixture() -> String {
44 + concat!(
45 + "I will check two cities before answering.\n",
46 + "<tool_call>get_weather",
47 + "<arg_key>city</arg_key>",
48 + "<arg_value>Hangzhou</arg_value>",
49 + "<arg_key>date</arg_key>",
50 + "<arg_value>2026-05-07</arg_value>",
51 + "<arg_key>unit</arg_key>",
52 + "<arg_value>celsius</arg_value>",
53 + "<arg_key>days</arg_key>",
54 + "<arg_value>3</arg_value>",
55 + "</tool_call>",
56 + "<tool_call>get_weather",
57 + "<arg_key>city</arg_key>",
58 + "<arg_value>San Francisco</arg_value>",
59 + "<arg_key>date</arg_key>",
60 + "<arg_value>2026-05-07</arg_value>",
61 + "<arg_key>unit</arg_key>",
62 + "<arg_value>fahrenheit</arg_value>",
63 + "<arg_key>days</arg_key>",
64 + "<arg_value>2</arg_value>",
65 + "</tool_call>",
66 + )
67 + .to_string()
68 +}
69 +
70 +fn long_normal_text_fixture() -> String {
71 + let line = "This is ordinary assistant text with no GLM MoE tool markers at all.\n";
72 + line.repeat(LONG_NORMAL_TEXT_REPEATS)
73 +}
74 +
75 +fn native_parser(name: &str, tools: &[Tool]) -> Box<dyn ToolParser> {
76 + match name {
77 + GLM45_PARSER_NAME => Glm45MoeToolParser::create(tools),
78 + GLM47_PARSER_NAME => Glm47MoeToolParser::create(tools),
79 + _ => unreachable!("unexpected GLM parser name"),
80 + }
81 + .expect("GLM MoE parser should initialize")
82 +}
83 +
84 +fn external_parser(name: &str) -> ExternalGlm4MoeParser {
85 + match name {
86 + GLM45_PARSER_NAME => ExternalGlm4MoeParser::glm45(),
87 + GLM47_PARSER_NAME => ExternalGlm4MoeParser::glm47(),
88 + _ => unreachable!("unexpected GLM parser name"),
89 + }
90 +}
91 +
92 +fn run_stream_group(
93 + c: &mut Criterion,
94 + name: &str,
95 + parser_name: &str,
96 + tools: &[Tool],
97 + text: &str,
98 + chunk_chars: usize,
99 + expected_normal_text: &str,
100 + expected_native_calls_len: usize,
101 +) {
102 + let chunks = split_by_chars(text, chunk_chars);
103 + let openai_tools = openai_tools(tools);
104 +
105 + let mut group = c.benchmark_group(name);
106 + group.sample_size(50);
107 + group.warm_up_time(Duration::from_millis(300));
108 + group.measurement_time(Duration::from_secs(2));
109 + group.throughput(Throughput::Bytes(text.len() as u64));
110 +
111 + group.bench_function("native_reuse_parser", |b| {
112 + let mut parser = native_parser(parser_name, tools);
113 + b.iter(|| {
114 + let result = feed_parser(&mut *parser, black_box(&chunks));
115 + debug_assert_eq!(result.0, expected_normal_text);
116 + debug_assert_eq!(result.1, expected_native_calls_len);
117 + black_box(result);
118 + })
119 + });
120 +
121 + group.bench_function("native_create_parser", |b| {
122 + b.iter_batched(
123 + || native_parser(parser_name, tools),
124 + |mut parser| {
125 + let result = feed_parser(&mut *parser, black_box(&chunks));
126 + debug_assert_eq!(result.0, expected_normal_text);
127 + debug_assert_eq!(result.1, expected_native_calls_len);
128 + black_box(result);
129 + },
130 + BatchSize::SmallInput,
131 + )
132 + });
133 +
134 + group.bench_function("external_reuse_parser", |b| {
135 + let mut parser = external_parser(parser_name);
136 + b.iter(|| {
137 + let result = feed_external_parser(&mut parser, &openai_tools, black_box(&chunks));
138 + debug_assert_eq!(result.0, expected_normal_text);
139 + black_box(result);
140 + })
141 + });
142 +
143 + group.bench_function("external_create_parser", |b| {
144 + b.iter_batched(
145 + || external_parser(parser_name),
146 + |mut parser| {
147 + let result = feed_external_parser(&mut parser, &openai_tools, black_box(&chunks));
148 + debug_assert_eq!(result.0, expected_normal_text);
149 + black_box(result);
150 + },
151 + BatchSize::SmallInput,
152 + )
153 + });
154 +
155 + group.finish();
156 +}
157 +
158 +fn bench_glm45_moe(c: &mut Criterion) {
159 + let tools = test_tools();
160 + let glm45_mixed_text = glm45_mixed_fixture();
161 + let glm47_mixed_text = glm47_mixed_fixture();
162 + let long_normal_text = long_normal_text_fixture();
163 +
164 + run_stream_group(
165 + c,
166 + "glm45/mixed_text_tool_call",
167 + GLM45_PARSER_NAME,
168 + &tools,
169 + &glm45_mixed_text,
170 + CHUNK_CHARS,
171 + "I will check two cities before answering.\n",
172 + 2,
173 + );
174 +
175 + run_stream_group(
176 + c,
177 + "glm47/mixed_text_tool_call",
178 + GLM47_PARSER_NAME,
179 + &tools,
180 + &glm47_mixed_text,
181 + CHUNK_CHARS,
182 + "I will check two cities before answering.\n",
183 + 2,
184 + );
185 +
186 + run_stream_group(
187 + c,
188 + "glm45/long_normal_text",
189 + GLM45_PARSER_NAME,
190 + &tools,
191 + &long_normal_text,
192 + CHUNK_CHARS,
193 + &long_normal_text,
194 + 0,
195 + );
196 +
197 + run_stream_group(
198 + c,
199 + "glm47/long_normal_text",
200 + GLM47_PARSER_NAME,
201 + &tools,
202 + &long_normal_text,
203 + CHUNK_CHARS,
204 + &long_normal_text,
205 + 0,
206 + );
207 +}
208 +
209 +criterion_group!(benches, bench_glm45_moe);
210 +criterion_main!(benches);
rust/src/tool-parser/benches/kimi_k2.rs new
+149
@@ -0,0 +1,149 @@
1 +use std::time::Duration;
2 +
3 +use criterion::{BatchSize, Criterion, Throughput, black_box, criterion_group, criterion_main};
4 +use tool_parser::parsers::KimiK2Parser as ExternalKimiK2Parser;
5 +use vllm_tool_parser::test_utils::{split_by_chars, test_tools};
6 +use vllm_tool_parser::{KimiK2ToolParser, Tool, ToolParser};
7 +
8 +mod utils;
9 +use utils::{feed_external_parser, feed_parser, openai_tools};
10 +
11 +const CHUNK_CHARS: usize = 7;
12 +const LONG_NORMAL_TEXT_REPEATS: usize = 2048;
13 +
14 +fn mixed_fixture() -> String {
15 + concat!(
16 + "I will check two cities before answering.\n",
17 + "<|tool_calls_section_begin|>",
18 + "<|tool_call_begin|>functions.get_weather:0",
19 + "<|tool_call_argument_begin|>{\"location\":\"Hangzhou\",\"days\":3}",
20 + "<|tool_call_end|>",
21 + "<|tool_call_begin|>functions.get_weather:1",
22 + "<|tool_call_argument_begin|>{\"location\":\"San Francisco\",\"days\":2}",
23 + "<|tool_call_end|>",
24 + "<|tool_calls_section_end|>",
25 + )
26 + .to_string()
27 +}
28 +
29 +fn mixed_chunks() -> Vec<&'static str> {
30 + vec![
31 + "I will check two cities before answering.\n",
32 + "<|tool_calls_section_begin|>",
33 + "<|tool_call_begin|>functions.get_weather:0",
34 + "<|tool_call_argument_begin|>",
35 + "{\"location\":",
36 + "\"Hangzhou\",",
37 + "\"days\":3}",
38 + "<|tool_call_end|>",
39 + "<|tool_call_begin|>functions.get_weather:1",
40 + "<|tool_call_argument_begin|>",
41 + "{\"location\":",
42 + "\"San Francisco\",",
43 + "\"days\":2}",
44 + "<|tool_call_end|>",
45 + "<|tool_calls_section_end|>",
46 + ]
47 +}
48 +
49 +fn long_normal_text_fixture() -> String {
50 + let line = "This is ordinary assistant text with no Kimi K2 tool markers at all.\n";
51 + line.repeat(LONG_NORMAL_TEXT_REPEATS)
52 +}
53 +
54 +fn native_parser(tools: &[Tool]) -> Box<dyn ToolParser> {
55 + KimiK2ToolParser::create(tools).expect("Kimi K2 parser should initialize")
56 +}
57 +
58 +fn run_stream_group(
59 + c: &mut Criterion,
60 + name: &str,
61 + tools: &[Tool],
62 + text: &str,
63 + chunks: &[&str],
64 + expected_normal_text: &str,
65 + expected_native_calls_len: usize,
66 +) {
67 + let openai_tools = openai_tools(tools);
68 +
69 + let mut group = c.benchmark_group(name);
70 + group.sample_size(50);
71 + group.warm_up_time(Duration::from_millis(300));
72 + group.measurement_time(Duration::from_secs(2));
73 + group.throughput(Throughput::Bytes(text.len() as u64));
74 +
75 + group.bench_function("native_reuse_parser", |b| {
76 + let mut parser = native_parser(tools);
77 + b.iter(|| {
78 + let result = feed_parser(&mut *parser, black_box(chunks));
79 + debug_assert_eq!(result.0, expected_normal_text);
80 + debug_assert_eq!(result.1, expected_native_calls_len);
81 + black_box(result);
82 + })
83 + });
84 +
85 + group.bench_function("native_create_parser", |b| {
86 + b.iter_batched(
87 + || native_parser(tools),
88 + |mut parser| {
89 + let result = feed_parser(&mut *parser, black_box(chunks));
90 + debug_assert_eq!(result.0, expected_normal_text);
91 + debug_assert_eq!(result.1, expected_native_calls_len);
92 + black_box(result);
93 + },
94 + BatchSize::SmallInput,
95 + )
96 + });
97 +
98 + group.bench_function("external_reuse_parser", |b| {
99 + let mut parser = ExternalKimiK2Parser::new();
100 + b.iter(|| {
101 + let result = feed_external_parser(&mut parser, &openai_tools, black_box(chunks));
102 + black_box(result);
103 + })
104 + });
105 +
106 + group.bench_function("external_create_parser", |b| {
107 + b.iter_batched(
108 + ExternalKimiK2Parser::new,
109 + |mut parser| {
110 + let result = feed_external_parser(&mut parser, &openai_tools, black_box(chunks));
111 + black_box(result);
112 + },
113 + BatchSize::SmallInput,
114 + )
115 + });
116 +
117 + group.finish();
118 +}
119 +
120 +fn bench_kimi_k2(c: &mut Criterion) {
121 + let tools = test_tools();
122 + let mixed_text = mixed_fixture();
123 + let mixed_chunks = mixed_chunks();
124 + let long_normal_text = long_normal_text_fixture();
125 + let long_normal_chunks = split_by_chars(&long_normal_text, CHUNK_CHARS);
126 +
127 + run_stream_group(
128 + c,
129 + "kimi_k2/mixed_text_tool_call",
130 + &tools,
131 + &mixed_text,
132 + &mixed_chunks,
133 + "I will check two cities before answering.\n",
134 + 2,
135 + );
136 +
137 + run_stream_group(
138 + c,
139 + "kimi_k2/long_normal_text",
140 + &tools,
141 + &long_normal_text,
142 + &long_normal_chunks,
143 + &long_normal_text,
144 + 0,
145 + );
146 +}
147 +
148 +criterion_group!(benches, bench_kimi_k2);
149 +criterion_main!(benches);
rust/src/tool-parser/benches/llama3_json.rs new
+128
@@ -0,0 +1,128 @@
1 +use std::time::Duration;
2 +
3 +use criterion::{BatchSize, Criterion, Throughput, black_box, criterion_group, criterion_main};
4 +use tool_parser::parsers::LlamaParser as ExternalLlamaParser;
5 +use vllm_tool_parser::test_utils::{split_by_chars, test_tools};
6 +use vllm_tool_parser::{Llama3JsonToolParser, Tool, ToolParser};
7 +
8 +mod utils;
9 +use utils::{feed_external_parser, feed_parser, openai_tools};
10 +
11 +const CHUNK_CHARS: usize = 7;
12 +const LONG_NORMAL_TEXT_REPEATS: usize = 2048;
13 +
14 +fn tool_call(function_name: &str, parameters: &str) -> String {
15 + format!(r#"{{"name":"{function_name}","parameters":{parameters}}}"#)
16 +}
17 +
18 +fn mixed_fixture() -> String {
19 + format!(
20 + "{}; {}",
21 + tool_call("get_weather", r#"{"location":"Hangzhou","days":3}"#),
22 + tool_call(
23 + "convert",
24 + r#"{"whole":42.5,"flag":true,"payload":{"nested":["x",null]},"items":[1,2,3],"empty":""}"#
25 + ),
26 + )
27 +}
28 +
29 +fn long_normal_text_fixture() -> String {
30 + let line = "This is ordinary assistant text with no Llama JSON tool call at the root.\n";
31 + line.repeat(LONG_NORMAL_TEXT_REPEATS)
32 +}
33 +
34 +fn native_parser(tools: &[Tool]) -> Box<dyn ToolParser> {
35 + Llama3JsonToolParser::create(tools).expect("Llama JSON parser should initialize")
36 +}
37 +
38 +fn run_stream_group(
39 + c: &mut Criterion,
40 + name: &str,
41 + tools: &[Tool],
42 + text: &str,
43 + chunk_chars: usize,
44 + expected_normal_text: &str,
45 + expected_native_calls_len: usize,
46 +) {
47 + let chunks = split_by_chars(text, chunk_chars);
48 + let openai_tools = openai_tools(tools);
49 +
50 + let mut group = c.benchmark_group(name);
51 + group.sample_size(50);
52 + group.warm_up_time(Duration::from_millis(300));
53 + group.measurement_time(Duration::from_secs(2));
54 + group.throughput(Throughput::Bytes(text.len() as u64));
55 +
56 + group.bench_function("native_reuse_parser", |b| {
57 + let mut parser = native_parser(tools);
58 + b.iter(|| {
59 + let result = feed_parser(&mut *parser, black_box(&chunks));
60 + debug_assert_eq!(result.0, expected_normal_text);
61 + debug_assert_eq!(result.1, expected_native_calls_len);
62 + black_box(result);
63 + })
64 + });
65 +
66 + group.bench_function("native_create_parser", |b| {
67 + b.iter_batched(
68 + || native_parser(tools),
69 + |mut parser| {
70 + let result = feed_parser(&mut *parser, black_box(&chunks));
71 + debug_assert_eq!(result.0, expected_normal_text);
72 + debug_assert_eq!(result.1, expected_native_calls_len);
73 + black_box(result);
74 + },
75 + BatchSize::SmallInput,
76 + )
77 + });
78 +
79 + group.bench_function("external_reuse_parser", |b| {
80 + let mut parser = ExternalLlamaParser::new();
81 + b.iter(|| {
82 + let result = feed_external_parser(&mut parser, &openai_tools, black_box(&chunks));
83 + black_box(result);
84 + })
85 + });
86 +
87 + group.bench_function("external_create_parser", |b| {
88 + b.iter_batched(
89 + ExternalLlamaParser::new,
90 + |mut parser| {
91 + let result = feed_external_parser(&mut parser, &openai_tools, black_box(&chunks));
92 + black_box(result);
93 + },
94 + BatchSize::SmallInput,
95 + )
96 + });
97 +
98 + group.finish();
99 +}
100 +
101 +fn bench_llama3_json(c: &mut Criterion) {
102 + let tools = test_tools();
103 + let mixed_text = mixed_fixture();
104 + let long_normal_text = long_normal_text_fixture();
105 +
106 + run_stream_group(
107 + c,
108 + "llama3_json/mixed_text_tool_call",
109 + &tools,
110 + &mixed_text,
111 + CHUNK_CHARS,
112 + "",
113 + 2,
114 + );
115 +
116 + run_stream_group(
117 + c,
118 + "llama3_json/long_normal_text",
119 + &tools,
120 + &long_normal_text,
121 + CHUNK_CHARS,
122 + &long_normal_text,
123 + 0,
124 + );
125 +}
126 +
127 +criterion_group!(benches, bench_llama3_json);
128 +criterion_main!(benches);
rust/src/tool-parser/benches/minimax_m2.rs new
+136
@@ -0,0 +1,136 @@
1 +use std::time::Duration;
2 +
3 +use criterion::{BatchSize, Criterion, Throughput, black_box, criterion_group, criterion_main};
4 +use tool_parser::parsers::MinimaxM2Parser as ExternalMinimaxM2Parser;
5 +use vllm_tool_parser::test_utils::{split_by_chars, test_tools};
6 +use vllm_tool_parser::{MinimaxM2ToolParser, Tool, ToolParser};
7 +
8 +mod utils;
9 +use utils::{feed_external_parser, feed_parser, openai_tools};
10 +
11 +const CHUNK_CHARS: usize = 7;
12 +const LONG_NORMAL_TEXT_REPEATS: usize = 2048;
13 +
14 +fn mixed_fixture() -> String {
15 + concat!(
16 + "I will check two cities before answering.\n",
17 + "<minimax:tool_call>",
18 + "<invoke name=\"get_weather\">",
19 + "<parameter name=\"city\">Hangzhou</parameter>",
20 + "<parameter name=\"date\">2026-04-30</parameter>",
21 + "<parameter name=\"unit\">celsius</parameter>",
22 + "<parameter name=\"days\">3</parameter>",
23 + "</invoke>",
24 + "<invoke name=\"get_weather\">",
25 + "<parameter name=\"city\">San Francisco</parameter>",
26 + "<parameter name=\"date\">2026-04-30</parameter>",
27 + "<parameter name=\"unit\">fahrenheit</parameter>",
28 + "<parameter name=\"days\">2</parameter>",
29 + "</invoke>",
30 + "</minimax:tool_call>",
31 + )
32 + .to_string()
33 +}
34 +
35 +fn long_normal_text_fixture() -> String {
36 + let line = "This is ordinary assistant text with no MiniMax M2 tool markers at all.\n";
37 + line.repeat(LONG_NORMAL_TEXT_REPEATS)
38 +}
39 +
40 +fn native_parser(tools: &[Tool]) -> Box<dyn ToolParser> {
41 + MinimaxM2ToolParser::create(tools).expect("MiniMax M2 parser should initialize")
42 +}
43 +
44 +fn run_stream_group(
45 + c: &mut Criterion,
46 + name: &str,
47 + tools: &[Tool],
48 + text: &str,
49 + chunk_chars: usize,
50 + expected_normal_text: &str,
51 + expected_native_calls_len: usize,
52 +) {
53 + let chunks = split_by_chars(text, chunk_chars);
54 + let openai_tools = openai_tools(tools);
55 +
56 + let mut group = c.benchmark_group(name);
57 + group.sample_size(50);
58 + group.warm_up_time(Duration::from_millis(300));
59 + group.measurement_time(Duration::from_secs(2));
60 + group.throughput(Throughput::Bytes(text.len() as u64));
61 +
62 + group.bench_function("native_reuse_parser", |b| {
63 + let mut parser = native_parser(tools);
64 + b.iter(|| {
65 + let result = feed_parser(&mut *parser, black_box(&chunks));
66 + debug_assert_eq!(result.0, expected_normal_text);
67 + debug_assert_eq!(result.1, expected_native_calls_len);
68 + black_box(result);
69 + })
70 + });
71 +
72 + group.bench_function("native_create_parser", |b| {
73 + b.iter_batched(
74 + || native_parser(tools),
75 + |mut parser| {
76 + let result = feed_parser(&mut *parser, black_box(&chunks));
77 + debug_assert_eq!(result.0, expected_normal_text);
78 + debug_assert_eq!(result.1, expected_native_calls_len);
79 + black_box(result);
80 + },
81 + BatchSize::SmallInput,
82 + )
83 + });
84 +
85 + group.bench_function("external_reuse_parser", |b| {
86 + let mut parser = ExternalMinimaxM2Parser::new();
87 + b.iter(|| {
88 + let result = feed_external_parser(&mut parser, &openai_tools, black_box(&chunks));
89 + debug_assert_eq!(result.0, expected_normal_text);
90 + black_box(result);
91 + })
92 + });
93 +
94 + group.bench_function("external_create_parser", |b| {
95 + b.iter_batched(
96 + ExternalMinimaxM2Parser::new,
97 + |mut parser| {
98 + let result = feed_external_parser(&mut parser, &openai_tools, black_box(&chunks));
99 + debug_assert_eq!(result.0, expected_normal_text);
100 + black_box(result);
101 + },
102 + BatchSize::SmallInput,
103 + )
104 + });
105 +
106 + group.finish();
107 +}
108 +
109 +fn bench_minimax_m2(c: &mut Criterion) {
110 + let tools = test_tools();
111 + let mixed_text = mixed_fixture();
112 + let long_normal_text = long_normal_text_fixture();
113 +
114 + run_stream_group(
115 + c,
116 + "minimax_m2/mixed_text_tool_call",
117 + &tools,
118 + &mixed_text,
119 + CHUNK_CHARS,
120 + "I will check two cities before answering.\n",
121 + 2,
122 + );
123 +
124 + run_stream_group(
125 + c,
126 + "minimax_m2/long_normal_text",
127 + &tools,
128 + &long_normal_text,
129 + CHUNK_CHARS,
130 + &long_normal_text,
131 + 0,
132 + );
133 +}
134 +
135 +criterion_group!(benches, bench_minimax_m2);
136 +criterion_main!(benches);
rust/src/tool-parser/benches/qwen3_coder.rs new
+138
@@ -0,0 +1,138 @@
1 +use std::time::Duration;
2 +
3 +use criterion::{BatchSize, Criterion, Throughput, black_box, criterion_group, criterion_main};
4 +use tool_parser::parsers::QwenCoderParser as ExternalQwenCoderParser;
5 +use vllm_tool_parser::test_utils::{split_by_chars, test_tools};
6 +use vllm_tool_parser::{Qwen3CoderToolParser, Tool, ToolParser};
7 +
8 +mod utils;
9 +use utils::{feed_external_parser, feed_parser, openai_tools};
10 +
11 +const CHUNK_CHARS: usize = 7;
12 +const LONG_NORMAL_TEXT_REPEATS: usize = 2048;
13 +
14 +fn mixed_fixture() -> String {
15 + concat!(
16 + "I will check two cities before answering.\n",
17 + "<tool_call>\n",
18 + "<function=get_weather>\n",
19 + "<parameter=location>Hangzhou</parameter>\n",
20 + "<parameter=date>2026-04-29</parameter>\n",
21 + "<parameter=unit>celsius</parameter>\n",
22 + "<parameter=days>3</parameter>\n",
23 + "</function>\n",
24 + "</tool_call>\n",
25 + "<tool_call>\n",
26 + "<function=get_weather>\n",
27 + "<parameter=location>San Francisco</parameter>\n",
28 + "<parameter=date>2026-04-29</parameter>\n",
29 + "<parameter=unit>fahrenheit</parameter>\n",
30 + "<parameter=days>2</parameter>\n",
31 + "</function>\n",
32 + "</tool_call>",
33 + )
34 + .to_string()
35 +}
36 +
37 +fn long_normal_text_fixture() -> String {
38 + let line = "This is ordinary assistant text with no Qwen Coder tool markers at all.\n";
39 + line.repeat(LONG_NORMAL_TEXT_REPEATS)
40 +}
41 +
42 +fn native_parser(tools: &[Tool]) -> Box<dyn ToolParser> {
43 + Qwen3CoderToolParser::create(tools).expect("Qwen Coder parser should initialize")
44 +}
45 +
46 +fn run_stream_group(
47 + c: &mut Criterion,
48 + name: &str,
49 + tools: &[Tool],
50 + text: &str,
51 + chunk_chars: usize,
52 + expected_normal_text: &str,
53 + expected_native_calls_len: usize,
54 +) {
55 + let chunks = split_by_chars(text, chunk_chars);
56 + let openai_tools = openai_tools(tools);
57 +
58 + let mut group = c.benchmark_group(name);
59 + group.sample_size(50);
60 + group.warm_up_time(Duration::from_millis(300));
61 + group.measurement_time(Duration::from_secs(2));
62 + group.throughput(Throughput::Bytes(text.len() as u64));
63 +
64 + group.bench_function("native_reuse_parser", |b| {
65 + let mut parser = native_parser(tools);
66 + b.iter(|| {
67 + let result = feed_parser(&mut *parser, black_box(&chunks));
68 + debug_assert_eq!(result.0, expected_normal_text);
69 + debug_assert_eq!(result.1, expected_native_calls_len);
70 + black_box(result);
71 + })
72 + });
73 +
74 + group.bench_function("native_create_parser", |b| {
75 + b.iter_batched(
76 + || native_parser(tools),
77 + |mut parser| {
78 + let result = feed_parser(&mut *parser, black_box(&chunks));
79 + debug_assert_eq!(result.0, expected_normal_text);
80 + debug_assert_eq!(result.1, expected_native_calls_len);
81 + black_box(result);
82 + },
83 + BatchSize::SmallInput,
84 + )
85 + });
86 +
87 + group.bench_function("external_reuse_parser", |b| {
88 + let mut parser = ExternalQwenCoderParser::new();
89 + b.iter(|| {
90 + let result = feed_external_parser(&mut parser, &openai_tools, black_box(&chunks));
91 + debug_assert_eq!(result.0, expected_normal_text);
92 + black_box(result);
93 + })
94 + });
95 +
96 + group.bench_function("external_create_parser", |b| {
97 + b.iter_batched(
98 + ExternalQwenCoderParser::new,
99 + |mut parser| {
100 + let result = feed_external_parser(&mut parser, &openai_tools, black_box(&chunks));
101 + debug_assert_eq!(result.0, expected_normal_text);
102 + black_box(result);
103 + },
104 + BatchSize::SmallInput,
105 + )
106 + });
107 +
108 + group.finish();
109 +}
110 +
111 +fn bench_qwen3_coder(c: &mut Criterion) {
112 + let tools = test_tools();
113 + let mixed_text = mixed_fixture();
114 + let long_normal_text = long_normal_text_fixture();
115 +
116 + run_stream_group(
117 + c,
118 + "qwen3_coder/mixed_text_tool_call",
119 + &tools,
120 + &mixed_text,
121 + CHUNK_CHARS,
122 + "I will check two cities before answering.\n",
123 + 2,
124 + );
125 +
126 + run_stream_group(
127 + c,
128 + "qwen3_coder/long_normal_text",
129 + &tools,
130 + &long_normal_text,
131 + CHUNK_CHARS,
132 + &long_normal_text,
133 + 0,
134 + );
135 +}
136 +
137 +criterion_group!(benches, bench_qwen3_coder);
138 +criterion_main!(benches);
rust/src/tool-parser/benches/qwen3_xml.rs new
+125
@@ -0,0 +1,125 @@
1 +use std::time::Duration;
2 +
3 +use criterion::{BatchSize, Criterion, Throughput, black_box, criterion_group, criterion_main};
4 +use tool_parser::parsers::QwenParser as ExternalQwenParser;
5 +use vllm_tool_parser::test_utils::{split_by_chars, test_tools};
6 +use vllm_tool_parser::{Qwen3XmlToolParser, Tool, ToolParser};
7 +
8 +mod utils;
9 +use utils::{feed_external_parser, feed_parser, openai_tools};
10 +
11 +const CHUNK_CHARS: usize = 7;
12 +const LONG_NORMAL_TEXT_REPEATS: usize = 2048;
13 +
14 +fn tool_call(function_name: &str, arguments: &str) -> String {
15 + format!("<tool_call>\n{{\"name\":\"{function_name}\",\"arguments\":{arguments}}}\n</tool_call>")
16 +}
17 +
18 +fn mixed_fixture() -> String {
19 + format!(
20 + "I will check two cities before answering.\n{}{}",
21 + tool_call("get_weather", r#"{"location":"Hangzhou","days":3}"#),
22 + tool_call("get_weather", r#"{"location":"San Francisco","days":2}"#),
23 + )
24 +}
25 +
26 +fn long_normal_text_fixture() -> String {
27 + let line = "This is ordinary assistant text with no Qwen XML tool markers at all.\n";
28 + line.repeat(LONG_NORMAL_TEXT_REPEATS)
29 +}
30 +
31 +fn native_parser(tools: &[Tool]) -> Box<dyn ToolParser> {
32 + Qwen3XmlToolParser::create(tools).expect("Qwen XML parser should initialize")
33 +}
34 +
35 +fn run_stream_group(
36 + c: &mut Criterion,
37 + name: &str,
38 + tools: &[Tool],
39 + text: &str,
40 + chunk_chars: usize,
41 + expected_normal_text: &str,
42 + expected_native_calls_len: usize,
43 +) {
44 + let chunks = split_by_chars(text, chunk_chars);
45 + let openai_tools = openai_tools(tools);
46 +
47 + let mut group = c.benchmark_group(name);
48 + group.sample_size(50);
49 + group.warm_up_time(Duration::from_millis(300));
50 + group.measurement_time(Duration::from_secs(2));
51 + group.throughput(Throughput::Bytes(text.len() as u64));
52 +
53 + group.bench_function("native_reuse_parser", |b| {
54 + let mut parser = native_parser(tools);
55 + b.iter(|| {
56 + let result = feed_parser(&mut *parser, black_box(&chunks));
57 + debug_assert_eq!(result.0, expected_normal_text);
58 + debug_assert_eq!(result.1, expected_native_calls_len);
59 + black_box(result);
60 + })
61 + });
62 +
63 + group.bench_function("native_create_parser", |b| {
64 + b.iter_batched(
65 + || native_parser(tools),
66 + |mut parser| {
67 + let result = feed_parser(&mut *parser, black_box(&chunks));
68 + debug_assert_eq!(result.0, expected_normal_text);
69 + debug_assert_eq!(result.1, expected_native_calls_len);
70 + black_box(result);
71 + },
72 + BatchSize::SmallInput,
73 + )
74 + });
75 +
76 + group.bench_function("external_reuse_parser", |b| {
77 + let mut parser = ExternalQwenParser::new();
78 + b.iter(|| {
79 + let result = feed_external_parser(&mut parser, &openai_tools, black_box(&chunks));
80 + black_box(result);
81 + })
82 + });
83 +
84 + group.bench_function("external_create_parser", |b| {
85 + b.iter_batched(
86 + ExternalQwenParser::new,
87 + |mut parser| {
88 + let result = feed_external_parser(&mut parser, &openai_tools, black_box(&chunks));
89 + black_box(result);
90 + },
91 + BatchSize::SmallInput,
92 + )
93 + });
94 +
95 + group.finish();
96 +}
97 +
98 +fn bench_qwen3_xml(c: &mut Criterion) {
99 + let tools = test_tools();
100 + let mixed_text = mixed_fixture();
101 + let long_normal_text = long_normal_text_fixture();
102 +
103 + run_stream_group(
104 + c,
105 + "qwen3_xml/mixed_text_tool_call",
106 + &tools,
107 + &mixed_text,
108 + CHUNK_CHARS,
109 + "I will check two cities before answering.\n",
110 + 2,
111 + );
112 +
113 + run_stream_group(
114 + c,
115 + "qwen3_xml/long_normal_text",
116 + &tools,
117 + &long_normal_text,
118 + CHUNK_CHARS,
119 + &long_normal_text,
120 + 0,
121 + );
122 +}
123 +
124 +criterion_group!(benches, bench_qwen3_xml);
125 +criterion_main!(benches);
rust/src/tool-parser/benches/utils/mod.rs new
+49
@@ -0,0 +1,49 @@
1 +#![allow(dead_code)]
2 +
3 +use futures::FutureExt as _;
4 +use openai_protocol::common::{Function as OpenAiFunction, Tool as OpenAiTool};
5 +use tool_parser::traits::ToolParser as ExternalToolParser;
6 +use vllm_tool_parser::test_utils::collect_stream;
7 +use vllm_tool_parser::{Tool, ToolParser};
8 +
9 +pub(super) fn openai_tools(tools: &[Tool]) -> Vec<OpenAiTool> {
10 + tools
11 + .iter()
12 + .map(|tool| OpenAiTool {
13 + tool_type: "function".to_string(),
14 + function: OpenAiFunction {
15 + name: tool.name.clone(),
16 + description: tool.description.clone(),
17 + parameters: tool.parameters.clone(),
18 + strict: tool.strict,
19 + },
20 + })
21 + .collect()
22 +}
23 +
24 +pub(super) fn feed_parser(parser: &mut dyn ToolParser, chunks: &[&str]) -> (String, usize) {
25 + let result = collect_stream(parser, chunks);
26 + (result.normal_text, result.calls.len())
27 +}
28 +
29 +pub(super) fn feed_external_parser(
30 + parser: &mut impl ExternalToolParser,
31 + tools: &[OpenAiTool],
32 + chunks: &[&str],
33 +) -> (String, usize) {
34 + ExternalToolParser::reset(parser);
35 +
36 + let mut normal_text = String::new();
37 + let mut calls_len = 0;
38 + for chunk in chunks {
39 + let delta = parser
40 + .parse_incremental(chunk, tools)
41 + .now_or_never()
42 + .expect("external parser should not suspend")
43 + .expect("chunk should parse");
44 + normal_text.push_str(&delta.normal_text);
45 + calls_len += delta.calls.len();
46 + }
47 + calls_len += parser.get_unstreamed_tool_args().unwrap_or_default().len();
48 + (normal_text, calls_len)
49 +}
rust/src/tool-parser/src/deepseek_dsml/deepseek_v32.rs new
+439
@@ -0,0 +1,439 @@
1 +use super::{DeepSeekDsmlToolParser, DsmlTokens};
2 +use crate::{Result, Tool, ToolParseResult, ToolParser};
3 +
4 +/// Tool parser for DeepSeek V3.2 models.
5 +///
6 +/// Example tool call content:
7 +///
8 +/// ```text
9 +/// <｜DSML｜function_calls>
10 +/// <｜DSML｜invoke name="get_weather">
11 +/// <｜DSML｜parameter name="location" string="true">杭州</｜DSML｜parameter>
12 +/// <｜DSML｜parameter name="date" string="true">2024-01-16</｜DSML｜parameter>
13 +/// </｜DSML｜invoke>
14 +/// <｜DSML｜invoke name="get_weather">
15 +/// <｜DSML｜parameter name="location" string="true">北京</｜DSML｜parameter>
16 +/// <｜DSML｜parameter name="date" string="true">2024-01-16</｜DSML｜parameter>
17 +/// </｜DSML｜invoke>
18 +/// </｜DSML｜function_calls>
19 +/// ```
20 +///
21 +/// Arguments are emitted only after a full `invoke` block is parsed.
22 +///
23 +/// DeepSeek V3.2 relies on DSML markers such as `｜DSML｜`, which are
24 +/// represented as special tokens in the tokenizer and therefore must be
25 +/// preserved during decode for parsing to work.
26 +pub struct DeepSeekV32ToolParser(DeepSeekDsmlToolParser);
27 +
28 +impl DeepSeekV32ToolParser {
29 + /// Create a DeepSeek V3.2 tool parser.
30 + pub(super) fn new(tools: &[Tool]) -> Self {
31 + Self(DeepSeekDsmlToolParser::new(tools, DsmlTokens::V32))
32 + }
33 +}
34 +
35 +impl ToolParser for DeepSeekV32ToolParser {
36 + /// Create a boxed DeepSeek V3.2 tool parser.
37 + fn create(tools: &[Tool]) -> Result<Box<dyn ToolParser>>
38 + where
39 + Self: Sized + 'static,
40 + {
41 + Ok(Box::new(Self::new(tools)))
42 + }
43 +
44 + /// Preserve DSML special tokens while decoding.
45 + fn preserve_special_tokens(&self) -> bool {
46 + true
47 + }
48 +
49 + /// Push one decoded text chunk through the DSML parser.
50 + fn push(&mut self, chunk: &str) -> Result<ToolParseResult> {
51 + self.0.push(chunk)
52 + }
53 +
54 + /// Flush buffered text and reset parser state.
55 + fn finish(&mut self) -> Result<ToolParseResult> {
56 + self.0.finish()
57 + }
58 +}
59 +
60 +#[cfg(test)]
61 +mod tests {
62 + use serde_json::{Value, json};
63 + use thiserror_ext::AsReport;
64 +
65 + use super::DeepSeekV32ToolParser;
66 + use crate::ToolParser;
67 + use crate::test_utils::{collect_stream, split_by_chars, test_tools};
68 +
69 + fn build_tool_call(function_name: &str, params: &[(&str, &str)]) -> String {
70 + let params = params
71 + .iter()
72 + .map(|(name, value)| {
73 + format!(
74 + r#"<｜DSML｜parameter name="{name}" string="true">{value}</｜DSML｜parameter>"#
75 + )
76 + })
77 + .collect::<Vec<_>>()
78 + .join("\n");
79 + format!(
80 + "<｜DSML｜function_calls>\n<｜DSML｜invoke name=\"{function_name}\">\n{params}\n</｜DSML｜invoke>\n</｜DSML｜function_calls>"
81 + )
82 + }
83 +
84 + #[test]
85 + fn deepseek_v32_parse_complete_without_tool_call_keeps_text() {
86 + let mut parser = DeepSeekV32ToolParser::new(&test_tools());
87 + let result = parser.parse_complete("Hello, world!").unwrap();
88 +
89 + assert_eq!(result.normal_text, "Hello, world!");
90 + assert!(result.calls.is_empty());
91 + }
92 +
93 + #[test]
94 + fn deepseek_v32_parse_complete_extracts_single_tool_call() {
95 + let mut parser = DeepSeekV32ToolParser::new(&test_tools());
96 + let result = parser
97 + .parse_complete(&build_tool_call(
98 + "get_weather",
99 + &[("location", "SF"), ("date", "2024-01-16")],
100 + ))
101 + .unwrap();
102 +
103 + assert!(result.normal_text.is_empty());
104 + assert_eq!(result.calls.len(), 1);
105 + assert_eq!(result.calls[0].name.as_deref(), Some("get_weather"));
106 + assert_eq!(
107 + serde_json::from_str::<Value>(&result.calls[0].arguments).unwrap(),
108 + json!({
109 + "location": "SF",
110 + "date": "2024-01-16"
111 + })
112 + );
113 + }
114 +
115 + #[test]
116 + fn deepseek_v32_parse_complete_preserves_prefix_text() {
117 + let mut parser = DeepSeekV32ToolParser::new(&test_tools());
118 + let output = format!(
119 + "Thinking... {}",
120 + build_tool_call("get_weather", &[("location", "NYC")])
121 + );
122 + let result = parser.parse_complete(&output).unwrap();
123 +
124 + assert_eq!(result.normal_text, "Thinking... ");
125 + assert_eq!(result.calls.len(), 1);
126 + }
127 +
128 + #[test]
129 + fn deepseek_v32_parse_complete_converts_schema_types() {
130 + let mut parser = DeepSeekV32ToolParser::new(&test_tools());
131 + let result = parser
132 + .parse_complete(
133 + "<｜DSML｜function_calls>\n\
134 + <｜DSML｜invoke name=\"convert\">\n\
135 + <｜DSML｜parameter name=\"whole\" string=\"false\">5.0</｜DSML｜parameter>\n\
136 + <｜DSML｜parameter name=\"flag\" string=\"false\">true</｜DSML｜parameter>\n\
137 + <｜DSML｜parameter name=\"payload\" string=\"false\">{\"nested\":true}</｜DSML｜parameter>\n\
138 + <｜DSML｜parameter name=\"items\" string=\"false\">[1,2]</｜DSML｜parameter>\n\
139 + <｜DSML｜parameter name=\"empty\" string=\"false\">null</｜DSML｜parameter>\n\
140 + </｜DSML｜invoke>\n\
141 + </｜DSML｜function_calls>",
142 + )
143 + .unwrap();
144 +
145 + assert_eq!(result.calls.len(), 1);
146 + assert_eq!(
147 + serde_json::from_str::<Value>(&result.calls[0].arguments).unwrap(),
148 + json!({
149 + "whole": 5.0,
150 + "flag": true,
151 + "payload": { "nested": true },
152 + "items": [1, 2],
153 + "empty": null,
154 + })
155 + );
156 + }
157 +
158 + #[test]
159 + fn deepseek_v32_parse_complete_string_attr_overrides_schema_types() {
160 + let mut parser = DeepSeekV32ToolParser::new(&test_tools());
161 + let result = parser
162 + .parse_complete(
163 + "<｜DSML｜function_calls>\n\
164 + <｜DSML｜invoke name=\"convert\">\n\
165 + <｜DSML｜parameter name=\"whole\" string=\"true\">5.0</｜DSML｜parameter>\n\
166 + <｜DSML｜parameter name=\"flag\" string=\"true\">true</｜DSML｜parameter>\n\
167 + <｜DSML｜parameter name=\"payload\" string=\"true\">{\"nested\":true}</｜DSML｜parameter>\n\
168 + <｜DSML｜parameter name=\"items\" string=\"true\">[1,2]</｜DSML｜parameter>\n\
169 + <｜DSML｜parameter name=\"empty\" string=\"true\">null</｜DSML｜parameter>\n\
170 + </｜DSML｜invoke>\n\
171 + </｜DSML｜function_calls>",
172 + )
173 + .unwrap();
174 +
175 + assert_eq!(result.calls.len(), 1);
176 + assert_eq!(
177 + serde_json::from_str::<Value>(&result.calls[0].arguments).unwrap(),
178 + json!({
179 + "whole": "5.0",
180 + "flag": "true",
181 + "payload": "{\"nested\":true}",
182 + "items": "[1,2]",
183 + "empty": "null",
184 + })
185 + );
186 + }
187 +
188 + #[test]
189 + fn deepseek_v32_parse_complete_unescapes_literal_closing_tags_in_parameter_value() {
190 + let mut parser = DeepSeekV32ToolParser::new(&test_tools());
191 + let result = parser
192 + .parse_complete(&build_tool_call(
193 + "get_weather",
194 + &[
195 + (
196 + "location",
197 + "Hangzhou &lt;/｜DSML｜parameter&gt;&lt;/｜DSML｜invoke&gt;&lt;/｜DSML｜function_calls&gt;",
198 + ),
199 + ("date", "2026-05-08"),
200 + ],
201 + ))
202 + .unwrap();
203 +
204 + assert_eq!(
205 + serde_json::from_str::<Value>(&result.calls[0].arguments).unwrap(),
206 + json!({
207 + "location": "Hangzhou </｜DSML｜parameter></｜DSML｜invoke></｜DSML｜function_calls>",
208 + "date": "2026-05-08",
209 + })
210 + );
211 + }
212 +
213 + #[test]
214 + fn deepseek_v32_streaming_extracts_single_tool_call() {
215 + let mut parser = DeepSeekV32ToolParser::new(&test_tools());
216 + let result = collect_stream(
217 + &mut parser,
218 + &[
219 + "<｜DSML｜function_calls>\n",
220 + "<｜DSML｜invoke name=\"get_weather\">\n",
221 + "<｜DSML｜parameter name=\"location\" string=\"true\">SF</｜DSML｜parameter>\n",
222 + "</｜DSML｜invoke>\n",
223 + "</｜DSML｜function_calls>",
224 + ],
225 + );
226 +
227 + assert!(result.normal_text.is_empty());
228 + assert_eq!(result.calls.len(), 1);
229 + assert_eq!(result.calls[0].name.as_deref(), Some("get_weather"));
230 + assert_eq!(
231 + serde_json::from_str::<Value>(&result.calls[0].arguments).unwrap(),
232 + json!({ "location": "SF" })
233 + );
234 + }
235 +
236 + #[test]
237 + fn deepseek_v32_streaming_preserves_prefix_text() {
238 + let mut parser = DeepSeekV32ToolParser::new(&test_tools());
239 + let result = collect_stream(
240 + &mut parser,
241 + &[
242 + "Thinking... ",
243 + "<｜DSML｜function_calls>\n",
244 + "<｜DSML｜invoke name=\"get_weather\">\n",
245 + "<｜DSML｜parameter name=\"location\" string=\"true\">SF</｜DSML｜parameter>\n",
246 + "</｜DSML｜invoke>\n",
247 + "</｜DSML｜function_calls>",
248 + ],
249 + );
250 +
251 + assert_eq!(result.normal_text, "Thinking... ");
252 + assert_eq!(result.calls.len(), 1);
253 + }
254 +
255 + #[test]
256 + fn deepseek_v32_streaming_without_tool_call_emits_text_incrementally() {
257 + let mut parser = DeepSeekV32ToolParser::new(&test_tools());
258 + let result = collect_stream(&mut parser, &["Hello, ", "world!"]);
259 +
260 + assert_eq!(result.normal_text, "Hello, world!");
261 + assert!(result.calls.is_empty());
262 + }
263 +
264 + #[test]
265 + fn deepseek_v32_streaming_extracts_multiple_tool_calls_in_order() {
266 + let mut parser = DeepSeekV32ToolParser::new(&test_tools());
267 + let result = collect_stream(
268 + &mut parser,
269 + &[&format!(
270 + "{}\n{}",
271 + build_tool_call("get_weather", &[("location", "SF")])
272 + .trim_end_matches("</｜DSML｜function_calls>"),
273 + "<｜DSML｜invoke name=\"get_weather\">\n<｜DSML｜parameter name=\"location\" string=\"true\">NYC</｜DSML｜parameter>\n</｜DSML｜invoke>\n</｜DSML｜function_calls>"
274 + )],
275 + );
276 +
277 + assert_eq!(result.calls.len(), 2);
278 + assert_eq!(result.calls[0].name.as_deref(), Some("get_weather"));
279 + assert_eq!(result.calls[1].name.as_deref(), Some("get_weather"));
280 + assert_eq!(result.calls[0].tool_index, 0);
281 + assert_eq!(result.calls[1].tool_index, 1);
282 + assert_eq!(
283 + serde_json::from_str::<Value>(&result.calls[0].arguments).unwrap(),
284 + json!({ "location": "SF" })
285 + );
286 + assert_eq!(
287 + serde_json::from_str::<Value>(&result.calls[1].arguments).unwrap(),
288 + json!({ "location": "NYC" })
289 + );
290 + }
291 +
292 + #[test]
293 + fn deepseek_v32_streaming_handles_start_token_split_across_chunks() {
294 + let text = build_tool_call("get_weather", &[("location", "SF")]);
295 + let chunks = split_by_chars(&text, 5);
296 + let mut parser = DeepSeekV32ToolParser::new(&test_tools());
297 + let result = collect_stream(&mut parser, &chunks);
298 +
299 + assert_eq!(result.calls.len(), 1);
300 + assert_eq!(
301 + serde_json::from_str::<Value>(&result.calls[0].arguments).unwrap(),
302 + json!({ "location": "SF" })
303 + );
304 + }
305 +
306 + #[test]
307 + fn deepseek_v32_streaming_handles_bpe_chunked_dsml_opener() {
308 + let mut parser = DeepSeekV32ToolParser::new(&test_tools());
309 + let result = collect_stream(
310 + &mut parser,
311 + &[
312 + "<｜DSML｜",
313 + "function",
314 + "_c",
315 + "all",
316 + "s",
317 + ">\n",
318 + "<｜DSML｜",
319 + "invoke",
320 + " name=\"",
321 + "get_weather",
322 + "\">\n",
323 + "<｜DSML｜",
324 + "parameter",
325 + " name=\"location\" string=\"true\">",
326 + "Beijing",
327 + "</｜DSML｜",
328 + "parameter>\n",
329 + "</｜DSML｜",
330 + "invoke>\n",
331 + "</｜DSML｜",
332 + "function_calls>",
333 + ],
334 + );
335 +
336 + assert!(result.normal_text.is_empty());
337 + assert_eq!(result.calls.len(), 1);
338 + assert_eq!(result.calls[0].name.as_deref(), Some("get_weather"));
339 + assert_eq!(
340 + serde_json::from_str::<Value>(&result.calls[0].arguments).unwrap(),
341 + json!({ "location": "Beijing" })
342 + );
343 + }
344 +
345 + #[test]
346 + fn deepseek_v32_streaming_truncated_parameter_does_not_leak_eos() {
347 + let mut parser = DeepSeekV32ToolParser::new(&test_tools());
348 + parser.push("<｜DSML｜function_calls>\n").unwrap();
349 + parser.push("<｜DSML｜invoke name=\"get_weather\">\n").unwrap();
350 + parser
351 + .push("<｜DSML｜parameter name=\"location\" string=\"true\">Tokyo")
352 + .unwrap();
353 + parser.push("<｜end▁of▁sentence｜>").unwrap();
354 +
355 + let error = parser.finish().unwrap_err();
356 + assert!(error.to_report_string().contains("incomplete DeepSeek DSML tool call"));
357 + }
358 + #[test]
359 + fn deepseek_v32_streaming_drops_eos_after_complete_tool_calls() {
360 + let mut parser = DeepSeekV32ToolParser::new(&test_tools());
361 + let result = collect_stream(
362 + &mut parser,
363 + &[
364 + "<｜DSML｜function_calls>\n",
365 + "<｜DSML｜invoke name=\"get_weather\">\n",
366 + "<｜DSML｜parameter name=\"location\" string=\"true\">SF</｜DSML｜parameter>\n",
367 + "</｜DSML｜invoke>\n",
368 + "</｜DSML｜function_calls><｜end▁of▁sentence｜>",
369 + ],
370 + );
371 +
372 + assert!(result.normal_text.is_empty());
373 + assert_eq!(result.calls.len(), 1);
374 + assert_eq!(result.calls[0].name.as_deref(), Some("get_weather"));
375 + }
376 +
377 + #[test]
378 + fn deepseek_v32_streaming_ignores_text_after_complete_tool_calls() {
379 + let mut parser = DeepSeekV32ToolParser::new(&test_tools());
380 + let result = collect_stream(
381 + &mut parser,
382 + &[
383 + "<｜DSML｜function_calls>\n",
384 + "<｜DSML｜invoke name=\"get_weather\">\n",
385 + "<｜DSML｜parameter name=\"location\" string=\"true\">SF</｜DSML｜parameter>\n",
386 + "</｜DSML｜invoke>\n",
387 + "</｜DSML｜function_calls>",
388 + "trailing text",
389 + ],
390 + );
391 +
392 + assert!(result.normal_text.is_empty());
393 + assert_eq!(result.calls.len(), 1);
394 + }
395 +
396 + #[test]
397 + fn deepseek_v32_streaming_does_not_emit_incomplete_invoke() {
398 + let mut parser = DeepSeekV32ToolParser::new(&test_tools());
399 + parser.push("<｜DSML｜function_calls>\n").unwrap();
400 + parser.push("<｜DSML｜invoke name=\"get_weather\">\n").unwrap();
401 + parser
402 + .push("<｜DSML｜parameter name=\"location\" string=\"true\">SF</｜DSML｜parameter>\n")
403 + .unwrap();
404 +
405 + let error = parser.finish().unwrap_err();
406 + assert!(error.to_report_string().contains("incomplete DeepSeek DSML tool call"));
407 + }
408 + #[test]
409 + fn deepseek_v32_parser_state_resets_after_finish() {
410 + let mut parser = DeepSeekV32ToolParser::new(&test_tools());
411 + let first = parser
412 + .parse_complete(&build_tool_call("get_weather", &[("location", "SF")]))
413 + .unwrap();
414 + let second = parser
415 + .parse_complete(&build_tool_call("get_weather", &[("location", "NYC")]))
416 + .unwrap();
417 +
418 + assert_eq!(first.calls.len(), 1);
419 + assert_eq!(second.calls.len(), 1);
420 + assert_eq!(
421 + serde_json::from_str::<Value>(&second.calls[0].arguments).unwrap(),
422 + json!({ "location": "NYC" })
423 + );
424 + }
425 +
426 + #[test]
427 + fn deepseek_v32_streaming_matches_parse_complete() {
428 + let full_text = build_tool_call("add", &[("x", "3"), ("y", "4")]);
429 + let chunks = split_by_chars(&full_text, 7);
430 + let mut streaming_parser = DeepSeekV32ToolParser::new(&test_tools());
431 + let streamed = collect_stream(&mut streaming_parser, &chunks);
432 +
433 + let mut parser = DeepSeekV32ToolParser::new(&test_tools());
434 + let complete = parser.parse_complete(&full_text).unwrap();
435 +
436 + assert_eq!(streamed.normal_text, complete.normal_text);
437 + assert_eq!(streamed.calls, complete.calls);
438 + }
439 +}
rust/src/tool-parser/src/deepseek_dsml/deepseek_v4.rs new
+133
@@ -0,0 +1,133 @@
1 +use super::{DeepSeekDsmlToolParser, DsmlTokens};
2 +use crate::{Result, Tool, ToolParseResult, ToolParser};
3 +
4 +/// Tool parser for DeepSeek V4 models.
5 +///
6 +/// Example tool call content:
7 +///
8 +/// ```text
9 +/// <｜DSML｜tool_calls>
10 +/// <｜DSML｜invoke name="get_weather">
11 +/// <｜DSML｜parameter name="location" string="true">杭州</｜DSML｜parameter>
12 +/// <｜DSML｜parameter name="date" string="true">2024-01-16</｜DSML｜parameter>
13 +/// </｜DSML｜invoke>
14 +/// <｜DSML｜invoke name="get_weather">
15 +/// <｜DSML｜parameter name="location" string="true">北京</｜DSML｜parameter>
16 +/// <｜DSML｜parameter name="date" string="true">2024-01-16</｜DSML｜parameter>
17 +/// </｜DSML｜invoke>
18 +/// </｜DSML｜tool_calls>
19 +/// ```
20 +///
21 +/// Arguments are emitted only after a full `invoke` block is parsed.
22 +///
23 +/// V4 reuses the V3.2 DSML invoke/parameter grammar but wraps calls in
24 +/// `<｜DSML｜tool_calls>` instead of `<｜DSML｜function_calls>`.
25 +///
26 +/// DeepSeek V4 relies on DSML markers such as `｜DSML｜`, which are
27 +/// represented as special tokens in the tokenizer and therefore must be
28 +/// preserved during decode for parsing to work.
29 +pub struct DeepSeekV4ToolParser(DeepSeekDsmlToolParser);
30 +
31 +impl DeepSeekV4ToolParser {
32 + /// Create a DeepSeek V4 tool parser.
33 + fn new(tools: &[Tool]) -> Self {
34 + Self(DeepSeekDsmlToolParser::new(tools, DsmlTokens::V4))
35 + }
36 +}
37 +
38 +impl ToolParser for DeepSeekV4ToolParser {
39 + /// Create a boxed DeepSeek V4 tool parser.
40 + fn create(tools: &[Tool]) -> Result<Box<dyn ToolParser>>
41 + where
42 + Self: Sized + 'static,
43 + {
44 + Ok(Box::new(Self::new(tools)))
45 + }
46 +
47 + /// Preserve DSML special tokens while decoding.
48 + fn preserve_special_tokens(&self) -> bool {
49 + true
50 + }
51 +
52 + /// Push one decoded text chunk through the DSML parser.
53 + fn push(&mut self, chunk: &str) -> Result<ToolParseResult> {
54 + self.0.push(chunk)
55 + }
56 +
57 + /// Flush buffered text and reset parser state.
58 + fn finish(&mut self) -> Result<ToolParseResult> {
59 + self.0.finish()
60 + }
61 +}
62 +
63 +#[cfg(test)]
64 +mod tests {
65 + use serde_json::{Value, json};
66 +
67 + use super::{DeepSeekV4ToolParser, ToolParser};
68 + use crate::test_utils::{collect_stream, test_tools};
69 +
70 + fn build_tool_call(function_name: &str, params: &[(&str, &str)]) -> String {
71 + let params = params
72 + .iter()
73 + .map(|(name, value)| {
74 + format!(
75 + r#"<｜DSML｜parameter name="{name}" string="true">{value}</｜DSML｜parameter>"#
76 + )
77 + })
78 + .collect::<Vec<_>>()
79 + .join("\n");
80 + format!(
81 + "<｜DSML｜tool_calls>\n<｜DSML｜invoke name=\"{function_name}\">\n{params}\n</｜DSML｜invoke>\n</｜DSML｜tool_calls>"
82 + )
83 + }
84 +
85 + #[test]
86 + fn deepseek_v4_parse_complete_reuses_dsml_parser_with_tool_calls_token() {
87 + let mut parser = DeepSeekV4ToolParser::new(&test_tools());
88 + let result = parser
89 + .parse_complete(&build_tool_call(
90 + "get_weather",
91 + &[("location", "SF"), ("date", "2024-01-16")],
92 + ))
93 + .unwrap();
94 +
95 + assert!(result.normal_text.is_empty());
96 + assert_eq!(result.calls.len(), 1);
97 + assert_eq!(result.calls[0].name.as_deref(), Some("get_weather"));
98 + assert_eq!(
99 + serde_json::from_str::<Value>(&result.calls[0].arguments).unwrap(),
100 + json!({
101 + "location": "SF",
102 + "date": "2024-01-16"
103 + })
104 + );
105 + }
106 +
107 + #[test]
108 + fn deepseek_v4_streaming_handles_tool_calls_token_split_across_chunks() {
109 + let mut parser = DeepSeekV4ToolParser::new(&test_tools());
110 + let result = collect_stream(
111 + &mut parser,
112 + &[
113 + "Thinking... ",
114 + "<｜DSML｜",
115 + "tool",
116 + "_calls>\n",
117 + "<｜DSML｜invoke name=\"get_weather\">\n",
118 + "<｜DSML｜parameter name=\"location\" string=\"true\">Beijing</｜DSML｜parameter>\n",
119 + "</｜DSML｜invoke>\n",
120 + "</｜DSML｜",
121 + "tool_calls>",
122 + ],
123 + );
124 +
125 + assert_eq!(result.normal_text, "Thinking... ");
126 + assert_eq!(result.calls.len(), 1);
127 + assert_eq!(result.calls[0].name.as_deref(), Some("get_weather"));
128 + assert_eq!(
129 + serde_json::from_str::<Value>(&result.calls[0].arguments).unwrap(),
130 + json!({ "location": "Beijing" })
131 + );
132 + }
133 +}
rust/src/tool-parser/src/deepseek_dsml/mod.rs new
+278
@@ -0,0 +1,278 @@
1 +use winnow::ascii::{multispace0 as ws0, multispace1 as ws1};
2 +use winnow::combinator::{alt, delimited, eof, repeat, seq, terminated};
3 +use winnow::prelude::*;
4 +use winnow::stream::Partial;
5 +use winnow::token::{literal, rest, take_until};
6 +
7 +use super::parameters::ToolSchemas;
8 +use super::utils::{parse_buffered_event, safe_text_len, xml_unescape};
9 +use super::{Result, ToolCallDelta, ToolParseResult};
10 +use crate::Tool;
11 +
12 +mod deepseek_v32;
13 +mod deepseek_v4;
14 +
15 +pub use deepseek_v4::DeepSeekV4ToolParser;
16 +pub use deepseek_v32::DeepSeekV32ToolParser;
17 +
18 +const INVOKE_START: &str = "<｜DSML｜invoke";
19 +const INVOKE_END: &str = "</｜DSML｜invoke>";
20 +const PARAMETER_START: &str = "<｜DSML｜parameter";
21 +const PARAMETER_END: &str = "</｜DSML｜parameter>";
22 +
23 +type DsmlInput<'i> = Partial<&'i str>;
24 +
25 +#[derive(Debug, Clone, Copy)]
26 +struct DsmlTokens {
27 + tool_calls_start: &'static str,
28 + tool_calls_end: &'static str,
29 +}
30 +
31 +impl DsmlTokens {
32 + const V32: Self = Self {
33 + tool_calls_start: "<｜DSML｜function_calls>",
34 + tool_calls_end: "</｜DSML｜function_calls>",
35 + };
36 + const V4: Self = Self {
37 + tool_calls_start: "<｜DSML｜tool_calls>",
38 + tool_calls_end: "</｜DSML｜tool_calls>",
39 + };
40 +}
41 +
42 +#[derive(Debug, Clone, Copy, PartialEq, Eq)]
43 +enum DsmlMode {
44 + Text,
45 + ToolBlock,
46 + Done,
47 +}
48 +
49 +#[derive(Debug, Clone, PartialEq, Eq)]
50 +enum DsmlEvent {
51 + Text {
52 + len: usize,
53 + },
54 + ToolCallsStart,
55 + Invoke {
56 + name: String,
57 + raw_params: Vec<DsmlParameter>,
58 + },
59 + ToolCallsEnd,
60 + IgnoredRest,
61 +}
62 +
63 +#[derive(Debug, Clone, PartialEq, Eq)]
64 +struct DsmlParameter {
65 + name: String,
66 + value: String,
67 + is_string: bool,
68 +}
69 +
70 +/// Tool parser core for DeepSeek DSML tool calls.
71 +struct DeepSeekDsmlToolParser {
72 + buffer: String,
73 + mode: DsmlMode,
74 + emitted_invoke_count: usize,
75 + tool_parameters: ToolSchemas,
76 + tokens: DsmlTokens,
77 +}
78 +
79 +impl DeepSeekDsmlToolParser {
80 + /// Create a parser with DSML tokens for one DeepSeek format.
81 + fn new(tools: &[Tool], tokens: DsmlTokens) -> Self {
82 + Self {
83 + buffer: String::new(),
84 + mode: DsmlMode::Text,
85 + emitted_invoke_count: 0,
86 + tool_parameters: ToolSchemas::from_tools(tools),
87 + tokens,
88 + }
89 + }
90 +
91 + /// Apply one parsed DSML event to parser state and output.
92 + fn apply_event(&mut self, event: DsmlEvent, result: &mut ToolParseResult) -> Result<()> {
93 + match event {
94 + DsmlEvent::Text { len: consumed_len } => {
95 + result.normal_text.push_str(&self.buffer[..consumed_len]);
96 + }
97 + DsmlEvent::ToolCallsStart => self.mode = DsmlMode::ToolBlock,
98 + DsmlEvent::Invoke { name, raw_params } => {
99 + let mut arguments = serde_json::Map::with_capacity(raw_params.len());
100 + for param in raw_params {
101 + let value = if param.is_string {
102 + serde_json::Value::String(param.value)
103 + } else {
104 + self.tool_parameters.convert_param_with_schema(
105 + &name,
106 + &param.name,
107 + &param.value,
108 + )
109 + };
110 + arguments.insert(param.name, value);
111 + }
112 + let arguments = serde_json::to_string(&arguments)
113 + .map_err(|error| parsing_failed!("failed to serialize arguments: {}", error))?;
114 +
115 + result.calls.push(ToolCallDelta {
116 + tool_index: self.emitted_invoke_count,
117 + name: Some(name),
118 + arguments,
119 + });
120 + self.emitted_invoke_count += 1;
121 + }
122 + DsmlEvent::ToolCallsEnd => self.mode = DsmlMode::Done,
123 + DsmlEvent::IgnoredRest => {}
124 + };
125 + Ok(())
126 + }
127 +
128 + /// Reset all streaming state.
129 + fn reset(&mut self) {
130 + self.buffer.clear();
131 + self.mode = DsmlMode::Text;
132 + self.emitted_invoke_count = 0;
133 + }
134 +
135 + /// Push one decoded text chunk through the DSML parser.
136 + fn push(&mut self, chunk: &str) -> Result<ToolParseResult> {
137 + // Extract tool calls from streaming model output.
138 + //
139 + // Uses a buffer-until-complete-invoke strategy: text is buffered until
140 + // a complete invoke block is available, then parsed and emitted in one
141 + // shot.
142 + self.buffer.push_str(chunk);
143 + let mut result = ToolParseResult::default();
144 +
145 + while let Some((event, consumed_len)) = parse_buffered_event(&self.buffer, |input| {
146 + parse_next_dsml_event(input, self.mode, self.tokens)
147 + })? {
148 + self.apply_event(event, &mut result)?;
149 + self.buffer.drain(..consumed_len);
150 + }
151 +
152 + Ok(result)
153 + }
154 +
155 + /// Flush buffered text and reset parser state.
156 + fn finish(&mut self) -> Result<ToolParseResult> {
157 + let mut result = ToolParseResult::default();
158 + match self.mode {
159 + DsmlMode::Text => result.normal_text.push_str(&self.buffer),
160 + DsmlMode::Done => {}
161 + DsmlMode::ToolBlock => {
162 + self.reset();
163 + return Err(parsing_failed!("incomplete DeepSeek DSML tool call"));
164 + }
165 + }
166 + self.reset();
167 + Ok(result)
168 + }
169 +}
170 +
171 +/// Parse a DSML event for the current parser mode.
172 +fn parse_next_dsml_event(
173 + input: &mut DsmlInput<'_>,
174 + mode: DsmlMode,
175 + tokens: DsmlTokens,
176 +) -> ModalResult<DsmlEvent> {
177 + match mode {
178 + DsmlMode::Text => parse_text_event(input, tokens),
179 + DsmlMode::ToolBlock => parse_tool_block_event(input, tokens),
180 + DsmlMode::Done => ignored_rest_event(input),
181 + }
182 +}
183 +
184 +/// Parse a text-mode DSML event.
185 +fn parse_text_event(input: &mut DsmlInput<'_>, tokens: DsmlTokens) -> ModalResult<DsmlEvent> {
186 + alt((
187 + |input: &mut DsmlInput<'_>| tool_calls_start_event(input, tokens),
188 + |input: &mut DsmlInput<'_>| safe_text_event(input, tokens),
189 + ))
190 + .parse_next(input)
191 +}
192 +
193 +/// Parse a tool-block DSML event.
194 +fn parse_tool_block_event(input: &mut DsmlInput<'_>, tokens: DsmlTokens) -> ModalResult<DsmlEvent> {
195 + ws0.void().parse_next(input)?;
196 + alt((invoke_event, |input: &mut DsmlInput<'_>| {
197 + tool_calls_end_event(input, tokens)
198 + }))
199 + .parse_next(input)
200 +}
201 +
202 +/// Parse a DSML function-calls start marker.
203 +fn tool_calls_start_event(input: &mut DsmlInput<'_>, tokens: DsmlTokens) -> ModalResult<DsmlEvent> {
204 + literal(tokens.tool_calls_start)
205 + .value(DsmlEvent::ToolCallsStart)
206 + .parse_next(input)
207 +}
208 +
209 +/// Parse a DSML function-calls end marker.
210 +fn tool_calls_end_event(input: &mut DsmlInput<'_>, tokens: DsmlTokens) -> ModalResult<DsmlEvent> {
211 + literal(tokens.tool_calls_end).value(DsmlEvent::ToolCallsEnd).parse_next(input)
212 +}
213 +
214 +/// Parse a trailing rest after DSML function calls.
215 +fn ignored_rest_event(input: &mut DsmlInput<'_>) -> ModalResult<DsmlEvent> {
216 + rest.value(DsmlEvent::IgnoredRest).parse_next(input)
217 +}
218 +
219 +/// Parse a safe text run before the next DSML marker.
220 +fn safe_text_event(input: &mut DsmlInput<'_>, tokens: DsmlTokens) -> ModalResult<DsmlEvent> {
221 + safe_text_len(input, tokens.tool_calls_start).map(|len| DsmlEvent::Text { len })
222 +}
223 +
224 +/// Parse a DSML invoke block.
225 +fn invoke_event(input: &mut DsmlInput<'_>) -> ModalResult<DsmlEvent> {
226 + let (name, body) = seq!(
227 + _: literal(INVOKE_START),
228 + _: ws1,
229 + dsml_name_attr,
230 + _: ws0,
231 + _: ">",
232 + take_until(0.., INVOKE_END),
233 + _: literal(INVOKE_END),
234 + )
235 + .parse_next(input)?;
236 + let raw_params = parse_invoke_params(body)?;
237 + Ok(DsmlEvent::Invoke {
238 + name: name.to_string(),
239 + raw_params,
240 + })
241 +}
242 +
243 +/// Parse a DSML invoke body.
244 +fn parse_invoke_params(invoke_body: &str) -> ModalResult<Vec<DsmlParameter>> {
245 + let mut input = invoke_body;
246 + delimited(ws0, repeat(0.., terminated(parse_parameter, ws0)), eof).parse_next(&mut input)
247 +}
248 +
249 +/// Parse a DSML parameter block.
250 +fn parse_parameter(input: &mut &str) -> ModalResult<DsmlParameter> {
251 + seq! {DsmlParameter {
252 + _: literal(PARAMETER_START),
253 + _: ws1,
254 + name: name_attr.map(|name: &str| name.to_string()),
255 + _: ws1,
256 + is_string: string_attr.map(|value| value == "true"),
257 + _: ws0,
258 + _: ">",
259 + value: take_until(0.., PARAMETER_END).map(xml_unescape).map(|value| value.into_owned()),
260 + _: literal(PARAMETER_END),
261 + }}
262 + .parse_next(input)
263 +}
264 +
265 +/// Parse a name attribute.
266 +fn name_attr<'i>(input: &mut &'i str) -> ModalResult<&'i str> {
267 + delimited("name=\"", take_until(1.., "\""), "\"").parse_next(input)
268 +}
269 +
270 +/// Parse a string attribute.
271 +fn string_attr<'i>(input: &mut &'i str) -> ModalResult<&'i str> {
272 + delimited("string=\"", alt(("true", "false")), "\"").parse_next(input)
273 +}
274 +
275 +/// Parse a DSML name attribute.
276 +fn dsml_name_attr<'i>(input: &mut DsmlInput<'i>) -> ModalResult<&'i str> {
277 + delimited("name=\"", take_until(1.., "\""), "\"").parse_next(input)
278 +}
rust/src/tool-parser/src/deepseek_json/deepseek_v3.rs new
+235
@@ -0,0 +1,235 @@
1 +use super::{DeepSeekJsonFormat, DeepSeekJsonToolParser};
2 +use crate::{Result, Tool, ToolParseResult, ToolParser};
3 +
4 +/// Tool parser for DeepSeek V3 JSON-fenced tool calls.
5 +///
6 +/// Example tool call content:
7 +///
8 +/// ````text
9 +/// <｜tool▁calls▁begin｜><｜tool▁call▁begin｜>function<｜tool▁sep｜>get_weather
10 +/// ```json
11 +/// {"location":"Tokyo"}
12 +/// ```<｜tool▁call▁end｜><｜tool▁calls▁end｜>
13 +/// ````
14 +///
15 +/// Arguments are already OpenAI-style JSON text inside the markdown fence, so
16 +/// they are streamed as raw argument deltas without schema conversion or JSON
17 +/// normalization.
18 +pub struct DeepSeekV3ToolParser(DeepSeekJsonToolParser);
19 +
20 +impl DeepSeekV3ToolParser {
21 + /// Create a DeepSeek V3 tool parser.
22 + fn new(_tools: &[Tool]) -> Self {
23 + Self(DeepSeekJsonToolParser::new(DeepSeekJsonFormat::V3))
24 + }
25 +}
26 +
27 +impl ToolParser for DeepSeekV3ToolParser {
28 + /// Create a boxed DeepSeek V3 tool parser.
29 + fn create(tools: &[Tool]) -> Result<Box<dyn ToolParser>>
30 + where
31 + Self: Sized + 'static,
32 + {
33 + Ok(Box::new(Self::new(tools)))
34 + }
35 +
36 + /// Push one decoded text chunk through the DeepSeek V3 parser.
37 + fn push(&mut self, chunk: &str) -> Result<ToolParseResult> {
38 + self.0.push(chunk)
39 + }
40 +
41 + /// Flush buffered text and reset parser state.
42 + fn finish(&mut self) -> Result<ToolParseResult> {
43 + self.0.finish()
44 + }
45 +}
46 +
47 +#[cfg(test)]
48 +mod tests {
49 + use expect_test::expect;
50 + use thiserror_ext::AsReport;
51 +
52 + use super::DeepSeekV3ToolParser;
53 + use crate::deepseek_json::{
54 + TOOL_CALL_SEPARATOR, TOOL_CALL_START, TOOL_CALLS_END, TOOL_CALLS_START, V3_ARGUMENT_END,
55 + V3_JSON_START,
56 + };
57 + use crate::test_utils::{collect_stream, split_by_chars, test_tools};
58 + use crate::{ToolParseResult, ToolParser};
59 +
60 + fn v3_tool_call(function_name: &str, arguments: &str) -> String {
61 + format!(
62 + "{TOOL_CALL_START}function{TOOL_CALL_SEPARATOR}{function_name}{V3_JSON_START}{arguments}{V3_ARGUMENT_END}"
63 + )
64 + }
65 +
66 + fn tool_section(tool_calls: &[String]) -> String {
67 + format!("{TOOL_CALLS_START}{}{TOOL_CALLS_END}", tool_calls.join(""))
68 + }
69 +
70 + #[test]
71 + fn deepseek_v3_parse_complete_without_tool_call_keeps_text() {
72 + let mut parser = DeepSeekV3ToolParser::new(&test_tools());
73 + let result = parser.parse_complete("Hello, world!").unwrap();
74 +
75 + assert_eq!(result.normal_text, "Hello, world!");
76 + assert!(result.calls.is_empty());
77 + }
78 +
79 + #[test]
80 + fn deepseek_v3_parse_complete_extracts_raw_json_arguments() {
81 + let mut parser = DeepSeekV3ToolParser::new(&test_tools());
82 + let arguments = r#"{ "location": "Tokyo", "days": "3" }"#;
83 + let result = parser
84 + .parse_complete(&format!(
85 + "Let me check.\n{} trailing text",
86 + tool_section(&[v3_tool_call("get_weather", arguments)])
87 + ))
88 + .unwrap();
89 +
90 + assert_eq!(result.normal_text, "Let me check.\n");
91 + assert_eq!(result.calls.len(), 1);
92 + assert_eq!(result.calls[0].tool_index, 0);
93 + assert_eq!(result.calls[0].name.as_deref(), Some("get_weather"));
94 + assert_eq!(result.calls[0].arguments, arguments);
95 + }
96 +
97 + #[test]
98 + fn deepseek_v3_does_not_validate_or_normalize_arguments() {
99 + let mut parser = DeepSeekV3ToolParser::new(&test_tools());
100 + let arguments = r#"{"location":"Tokyo",}"#;
101 + let result = parser
102 + .parse_complete(&tool_section(&[v3_tool_call("get_weather", arguments)]))
103 + .unwrap();
104 +
105 + assert_eq!(result.calls[0].arguments, arguments);
106 + }
107 +
108 + #[test]
109 + fn deepseek_v3_streaming_emits_argument_deltas() {
110 + let mut parser = DeepSeekV3ToolParser::new(&test_tools());
111 + let chunks = [
112 + TOOL_CALLS_START,
113 + TOOL_CALL_START,
114 + "function",
115 + TOOL_CALL_SEPARATOR,
116 + "get_weather",
117 + V3_JSON_START,
118 + "{\"location\":",
119 + "\"Beijing\"",
120 + "}",
121 + V3_ARGUMENT_END,
122 + TOOL_CALLS_END,
123 + ];
124 +
125 + let mut result = ToolParseResult::default();
126 + let mut observed_arguments = Vec::new();
127 + for chunk in chunks {
128 + let next = parser.push(chunk).unwrap();
129 + observed_arguments.extend(
130 + next.calls
131 + .iter()
132 + .filter(|call| call.name.is_none())
133 + .map(|call| call.arguments.clone()),
134 + );
135 + result.append(next);
136 + }
137 + result.append(parser.finish().unwrap());
138 +
139 + assert_eq!(observed_arguments, ["{\"location\":", "\"Beijing\"", "}"]);
140 + assert_eq!(
141 + result.coalesce_calls().calls[0].arguments,
142 + r#"{"location":"Beijing"}"#
143 + );
144 + }
145 +
146 + #[test]
147 + fn deepseek_v3_streaming_handles_split_markers() {
148 + let input = format!(
149 + "hello {}",
150 + tool_section(&[v3_tool_call("get_weather", r#"{"location":"Tokyo"}"#)])
151 + );
152 + let chunks = split_by_chars(&input, 5);
153 + let mut parser = DeepSeekV3ToolParser::new(&test_tools());
154 +
155 + let result = collect_stream(&mut parser, &chunks);
156 +
157 + assert_eq!(result.normal_text, "hello ");
158 + assert_eq!(result.calls.len(), 1);
159 + assert_eq!(result.calls[0].arguments, r#"{"location":"Tokyo"}"#);
160 + }
161 +
162 + #[test]
163 + fn deepseek_v3_keeps_fenced_end_marker_literal_inside_json_string() {
164 + let mut parser = DeepSeekV3ToolParser::new(&test_tools());
165 + let arguments = format!("{{\"text\":\"literal {V3_ARGUMENT_END} inside\"}}");
166 + let input = tool_section(&[v3_tool_call("echo", &arguments)]);
167 +
168 + let result = parser.parse_complete(&input).unwrap();
169 +
170 + assert_eq!(result.calls.len(), 1);
171 + assert_eq!(result.calls[0].arguments, arguments);
172 + }
173 +
174 + #[test]
175 + fn deepseek_v3_streaming_extracts_multiple_tool_calls() {
176 + let input = tool_section(&[
177 + v3_tool_call("get_weather", r#"{"location":"Shanghai"}"#),
178 + v3_tool_call("add", r#"{"x":1,"y":2}"#),
179 + ]);
180 + let chunks = split_by_chars(&input, 7);
181 + let mut parser = DeepSeekV3ToolParser::new(&test_tools());
182 +
183 + let result = collect_stream(&mut parser, &chunks);
184 +
185 + expect![[r#"
186 + ToolParseResult {
187 + normal_text: "",
188 + calls: [
189 + ToolCallDelta {
190 + tool_index: 0,
191 + name: Some(
192 + "get_weather",
193 + ),
194 + arguments: "{\"location\":\"Shanghai\"}",
195 + },
196 + ToolCallDelta {
197 + tool_index: 1,
198 + name: Some(
199 + "add",
200 + ),
201 + arguments: "{\"x\":1,\"y\":2}",
202 + },
203 + ],
204 + }
205 + "#]]
206 + .assert_debug_eq(&result);
207 + }
208 +
209 + #[test]
210 + fn deepseek_v3_finish_fails_incomplete_tool_call() {
211 + let mut parser = DeepSeekV3ToolParser::new(&test_tools());
212 + parser
213 + .push(&format!(
214 + "{TOOL_CALLS_START}{TOOL_CALL_START}function{TOOL_CALL_SEPARATOR}get_weather{V3_JSON_START}{{\"location\""
215 + ))
216 + .unwrap();
217 +
218 + let error = parser.finish().unwrap_err();
219 +
220 + expect!["tool parser parsing failed: incomplete DeepSeek V3 tool call"]
221 + .assert_eq(&error.to_report_string());
222 + }
223 +
224 + #[test]
225 + fn deepseek_v3_malformed_type_fails_fast() {
226 + let mut parser = DeepSeekV3ToolParser::new(&test_tools());
227 + let input = format!(
228 + "{TOOL_CALLS_START}{TOOL_CALL_START}tool{TOOL_CALL_SEPARATOR}get_weather{V3_JSON_START}{{}}"
229 + );
230 +
231 + let error = parser.push(&input).unwrap_err();
232 +
233 + expect!["tool parser parsing failed: "].assert_eq(&error.to_report_string());
234 + }
235 +}
rust/src/tool-parser/src/deepseek_json/deepseek_v31.rs new
+239
@@ -0,0 +1,239 @@
1 +use super::{DeepSeekJsonFormat, DeepSeekJsonToolParser};
2 +use crate::{Result, Tool, ToolParseResult, ToolParser};
3 +
4 +/// Tool parser for DeepSeek V3.1 raw JSON tool calls.
5 +///
6 +/// Example tool call content:
7 +///
8 +/// ```text
9 +/// <｜tool▁calls▁begin｜><｜tool▁call▁begin｜>get_weather<｜tool▁sep｜>{"location":"Tokyo"}<｜tool▁call▁end｜><｜tool▁calls▁end｜>
10 +/// ```
11 +///
12 +/// Arguments are already OpenAI-style JSON text, so they are streamed as raw
13 +/// argument deltas without schema conversion or JSON normalization.
14 +pub struct DeepSeekV31ToolParser(DeepSeekJsonToolParser);
15 +
16 +impl DeepSeekV31ToolParser {
17 + /// Create a DeepSeek V3.1 tool parser.
18 + fn new(_tools: &[Tool]) -> Self {
19 + Self(DeepSeekJsonToolParser::new(DeepSeekJsonFormat::V31))
20 + }
21 +}
22 +
23 +impl ToolParser for DeepSeekV31ToolParser {
24 + /// Create a boxed DeepSeek V3.1 tool parser.
25 + fn create(tools: &[Tool]) -> Result<Box<dyn ToolParser>>
26 + where
27 + Self: Sized + 'static,
28 + {
29 + Ok(Box::new(Self::new(tools)))
30 + }
31 +
32 + /// Push one decoded text chunk through the DeepSeek V3.1 parser.
33 + fn push(&mut self, chunk: &str) -> Result<ToolParseResult> {
34 + self.0.push(chunk)
35 + }
36 +
37 + /// Flush buffered text and reset parser state.
38 + fn finish(&mut self) -> Result<ToolParseResult> {
39 + self.0.finish()
40 + }
41 +}
42 +
43 +#[cfg(test)]
44 +mod tests {
45 + use expect_test::expect;
46 + use thiserror_ext::AsReport;
47 +
48 + use super::DeepSeekV31ToolParser;
49 + use crate::deepseek_json::{
50 + TOOL_CALL_END, TOOL_CALL_SEPARATOR, TOOL_CALL_START, TOOL_CALLS_END, TOOL_CALLS_START,
51 + };
52 + use crate::test_utils::{collect_stream, split_by_chars, test_tools};
53 + use crate::{ToolParseResult, ToolParser};
54 +
55 + fn v31_tool_call(function_name: &str, arguments: &str) -> String {
56 + format!("{TOOL_CALL_START}{function_name}{TOOL_CALL_SEPARATOR}{arguments}{TOOL_CALL_END}")
57 + }
58 +
59 + fn tool_section(tool_calls: &[String]) -> String {
60 + format!("{TOOL_CALLS_START}{}{TOOL_CALLS_END}", tool_calls.join(""))
61 + }
62 +
63 + #[test]
64 + fn deepseek_v31_parse_complete_without_tool_call_keeps_text() {
65 + let mut parser = DeepSeekV31ToolParser::new(&test_tools());
66 + let result = parser.parse_complete("Hello, world!").unwrap();
67 +
68 + assert_eq!(result.normal_text, "Hello, world!");
69 + assert!(result.calls.is_empty());
70 + }
71 +
72 + #[test]
73 + fn deepseek_v31_parse_complete_extracts_raw_json_arguments() {
74 + let mut parser = DeepSeekV31ToolParser::new(&test_tools());
75 + let arguments = r#"{ "location": "Tokyo", "days": "3" }"#;
76 + let result = parser
77 + .parse_complete(&format!(
78 + "Let me check.{} trailing text",
79 + tool_section(&[v31_tool_call("get_weather", arguments)])
80 + ))
81 + .unwrap();
82 +
83 + assert_eq!(result.normal_text, "Let me check.");
84 + assert_eq!(result.calls.len(), 1);
85 + assert_eq!(result.calls[0].tool_index, 0);
86 + assert_eq!(result.calls[0].name.as_deref(), Some("get_weather"));
87 + assert_eq!(result.calls[0].arguments, arguments);
88 + }
89 +
90 + #[test]
91 + fn deepseek_v31_does_not_validate_or_normalize_arguments() {
92 + let mut parser = DeepSeekV31ToolParser::new(&test_tools());
93 + let arguments = r#"{"location":"Tokyo",}"#;
94 + let result = parser
95 + .parse_complete(&tool_section(&[v31_tool_call("get_weather", arguments)]))
96 + .unwrap();
97 +
98 + assert_eq!(result.calls[0].arguments, arguments);
99 + }
100 +
101 + #[test]
102 + fn deepseek_v31_streaming_emits_argument_deltas() {
103 + let mut parser = DeepSeekV31ToolParser::new(&test_tools());
104 + let chunks = [
105 + TOOL_CALLS_START,
106 + TOOL_CALL_START,
107 + "get_weather",
108 + TOOL_CALL_SEPARATOR,
109 + "{\"location\":",
110 + "\"Beijing\"",
111 + "}",
112 + TOOL_CALL_END,
113 + TOOL_CALLS_END,
114 + ];
115 +
116 + let mut result = ToolParseResult::default();
117 + let mut observed_arguments = Vec::new();
118 + for chunk in chunks {
119 + let next = parser.push(chunk).unwrap();
120 + observed_arguments.extend(
121 + next.calls
122 + .iter()
123 + .filter(|call| call.name.is_none())
124 + .map(|call| call.arguments.clone()),
125 + );
126 + result.append(next);
127 + }
128 + result.append(parser.finish().unwrap());
129 +
130 + assert_eq!(observed_arguments, ["{\"location\":", "\"Beijing\"", "}"]);
131 + assert_eq!(
132 + result.coalesce_calls().calls[0].arguments,
133 + r#"{"location":"Beijing"}"#
134 + );
135 + }
136 +
137 + #[test]
138 + fn deepseek_v31_streaming_handles_split_markers() {
139 + let input = format!(
140 + "hello {}",
141 + tool_section(&[v31_tool_call("get_weather", r#"{"location":"Tokyo"}"#)])
142 + );
143 + let chunks = split_by_chars(&input, 5);
144 + let mut parser = DeepSeekV31ToolParser::new(&test_tools());
145 +
146 + let result = collect_stream(&mut parser, &chunks);
147 +
148 + assert_eq!(result.normal_text, "hello ");
149 + assert_eq!(result.calls.len(), 1);
150 + assert_eq!(result.calls[0].arguments, r#"{"location":"Tokyo"}"#);
151 + }
152 +
153 + #[test]
154 + fn deepseek_v31_keeps_end_marker_literal_inside_json_string() {
155 + let mut parser = DeepSeekV31ToolParser::new(&test_tools());
156 + let arguments = format!(r#"{{"text":"literal {TOOL_CALL_END} inside"}}"#);
157 + let input = tool_section(&[v31_tool_call("echo", &arguments)]);
158 +
159 + let result = parser.parse_complete(&input).unwrap();
160 +
161 + assert_eq!(result.calls.len(), 1);
162 + assert_eq!(result.calls[0].arguments, arguments);
163 + }
164 +
165 + #[test]
166 + fn deepseek_v31_streaming_extracts_multiple_tool_calls() {
167 + let input = tool_section(&[
168 + v31_tool_call("get_weather", r#"{"location":"Shanghai"}"#),
169 + v31_tool_call("add", r#"{"x":1,"y":2}"#),
170 + ]);
171 + let chunks = split_by_chars(&input, 7);
172 + let mut parser = DeepSeekV31ToolParser::new(&test_tools());
173 +
174 + let result = collect_stream(&mut parser, &chunks);
175 +
176 + expect![[r#"
177 + ToolParseResult {
178 + normal_text: "",
179 + calls: [
180 + ToolCallDelta {
181 + tool_index: 0,
182 + name: Some(
183 + "get_weather",
184 + ),
185 + arguments: "{\"location\":\"Shanghai\"}",
186 + },
187 + ToolCallDelta {
188 + tool_index: 1,
189 + name: Some(
190 + "add",
191 + ),
192 + arguments: "{\"x\":1,\"y\":2}",
193 + },
194 + ],
195 + }
196 + "#]]
197 + .assert_debug_eq(&result);
198 + }
199 +
200 + #[test]
201 + fn deepseek_v31_streaming_drops_eos_after_complete_tool_calls() {
202 + let input = format!(
203 + "{}<｜end▁of▁sentence｜>",
204 + tool_section(&[v31_tool_call("get_weather", r#"{"location":"Tokyo"}"#)])
205 + );
206 + let mut parser = DeepSeekV31ToolParser::new(&test_tools());
207 +
208 + let result = collect_stream(&mut parser, &[&input]);
209 +
210 + assert!(result.normal_text.is_empty());
211 + assert_eq!(result.calls.len(), 1);
212 + assert_eq!(result.calls[0].arguments, r#"{"location":"Tokyo"}"#);
213 + }
214 +
215 + #[test]
216 + fn deepseek_v31_finish_fails_incomplete_tool_call() {
217 + let mut parser = DeepSeekV31ToolParser::new(&test_tools());
218 + parser
219 + .push(&format!(
220 + "{TOOL_CALLS_START}{TOOL_CALL_START}get_weather{TOOL_CALL_SEPARATOR}{{\"location\""
221 + ))
222 + .unwrap();
223 +
224 + let error = parser.finish().unwrap_err();
225 +
226 + expect!["tool parser parsing failed: incomplete DeepSeek V3.1 tool call"]
227 + .assert_eq(&error.to_report_string());
228 + }
229 +
230 + #[test]
231 + fn deepseek_v31_malformed_empty_name_fails_fast() {
232 + let mut parser = DeepSeekV31ToolParser::new(&test_tools());
233 + let input = format!("{TOOL_CALLS_START}{TOOL_CALL_START}{TOOL_CALL_SEPARATOR}{{}}");
234 +
235 + let error = parser.push(&input).unwrap_err();
236 +
237 + expect!["tool parser parsing failed: "].assert_eq(&error.to_report_string());
238 + }
239 +}
rust/src/tool-parser/src/deepseek_json/mod.rs new
+308
@@ -0,0 +1,308 @@
1 +mod deepseek_v3;
2 +mod deepseek_v31;
3 +
4 +pub use deepseek_v3::DeepSeekV3ToolParser;
5 +pub use deepseek_v31::DeepSeekV31ToolParser;
6 +use winnow::ascii::multispace0 as ws0;
7 +use winnow::combinator::{alt, seq};
8 +use winnow::prelude::*;
9 +use winnow::stream::Partial;
10 +use winnow::token::{literal, rest, take_until};
11 +
12 +use super::utils::{JsonObjectScanState, parse_buffered_event, safe_text_len, take_json_object};
13 +use super::{Result, ToolCallDelta, ToolParseResult};
14 +
15 +pub(super) const TOOL_CALLS_START: &str = "<｜tool▁calls▁begin｜>";
16 +pub(super) const TOOL_CALLS_END: &str = "<｜tool▁calls▁end｜>";
17 +pub(super) const TOOL_CALL_START: &str = "<｜tool▁call▁begin｜>";
18 +pub(super) const TOOL_CALL_END: &str = "<｜tool▁call▁end｜>";
19 +pub(super) const TOOL_CALL_SEPARATOR: &str = "<｜tool▁sep｜>";
20 +pub(super) const V3_JSON_START: &str = "\n```json\n";
21 +pub(super) const V3_ARGUMENT_END: &str = "\n```<｜tool▁call▁end｜>";
22 +
23 +type DeepSeekJsonInput<'i> = Partial<&'i str>;
24 +
25 +#[derive(Debug, Clone, Copy, PartialEq, Eq)]
26 +enum DeepSeekJsonFormat {
27 + V3,
28 + V31,
29 +}
30 +
31 +impl DeepSeekJsonFormat {
32 + /// Return the parser name used in diagnostics.
33 + const fn parser_name(self) -> &'static str {
34 + match self {
35 + Self::V3 => "DeepSeek V3",
36 + Self::V31 => "DeepSeek V3.1",
37 + }
38 + }
39 +
40 + /// Return the marker that closes the raw JSON arguments payload.
41 + const fn argument_end_marker(self) -> &'static str {
42 + match self {
43 + Self::V3 => V3_ARGUMENT_END,
44 + Self::V31 => TOOL_CALL_END,
45 + }
46 + }
47 +}
48 +
49 +#[derive(Debug, Clone, PartialEq, Eq)]
50 +enum DeepSeekJsonMode {
51 + Text,
52 + ToolBlock,
53 + Header,
54 + Arguments { json_scan: JsonObjectScanState },
55 + Done,
56 +}
57 +
58 +#[derive(Debug, Clone, PartialEq, Eq)]
59 +enum DeepSeekJsonEvent {
60 + Text { len: usize },
61 + ToolCallsStart,
62 + ToolCallStart,
63 + ToolCallHeader { function_name: String },
64 + Arguments { len: usize },
65 + ToolCallEnd,
66 + ToolCallsEnd,
67 + IgnoredRest,
68 +}
69 +
70 +/// Tool parser core for DeepSeek JSON-argument tool calls.
71 +struct DeepSeekJsonToolParser {
72 + buffer: String,
73 + mode: DeepSeekJsonMode,
74 + active_tool_index: Option<usize>,
75 + emitted_tool_count: usize,
76 + format: DeepSeekJsonFormat,
77 +}
78 +
79 +impl DeepSeekJsonToolParser {
80 + /// Create a parser for one DeepSeek JSON-argument format.
81 + fn new(format: DeepSeekJsonFormat) -> Self {
82 + Self {
83 + buffer: String::new(),
84 + mode: DeepSeekJsonMode::Text,
85 + active_tool_index: None,
86 + emitted_tool_count: 0,
87 + format,
88 + }
89 + }
90 +
91 + /// Apply one parsed DeepSeek JSON event to parser state and output.
92 + fn apply_event(
93 + &mut self,
94 + event: DeepSeekJsonEvent,
95 + result: &mut ToolParseResult,
96 + ) -> Result<()> {
97 + match event {
98 + DeepSeekJsonEvent::Text { len: consumed_len } => {
99 + result.normal_text.push_str(&self.buffer[..consumed_len]);
100 + }
101 + DeepSeekJsonEvent::ToolCallsStart => self.mode = DeepSeekJsonMode::ToolBlock,
102 + DeepSeekJsonEvent::ToolCallStart => self.mode = DeepSeekJsonMode::Header,
103 + DeepSeekJsonEvent::ToolCallHeader { function_name } => {
104 + let tool_index = self.emitted_tool_count;
105 + self.emitted_tool_count += 1;
106 + self.active_tool_index = Some(tool_index);
107 + self.mode = DeepSeekJsonMode::Arguments {
108 + json_scan: JsonObjectScanState::default(),
109 + };
110 + result.calls.push(ToolCallDelta {
111 + tool_index,
112 + name: Some(function_name),
113 + arguments: String::new(),
114 + });
115 + }
116 + DeepSeekJsonEvent::Arguments { len: consumed_len } => {
117 + let Some(tool_index) = self.active_tool_index else {
118 + return Err(parsing_failed!(
119 + "{} arguments without an active tool call",
120 + self.format.parser_name()
121 + ));
122 + };
123 + result.calls.push(ToolCallDelta {
124 + tool_index,
125 + name: None,
126 + arguments: self.buffer[..consumed_len].to_string(),
127 + });
128 + }
129 + DeepSeekJsonEvent::ToolCallEnd => {
130 + self.active_tool_index = None;
131 + self.mode = DeepSeekJsonMode::ToolBlock;
132 + }
133 + DeepSeekJsonEvent::ToolCallsEnd => {
134 + self.active_tool_index = None;
135 + self.mode = DeepSeekJsonMode::Done;
136 + }
137 + DeepSeekJsonEvent::IgnoredRest => {}
138 + }
139 + Ok(())
140 + }
141 +
142 + /// Push one decoded text chunk through the DeepSeek JSON parser.
143 + fn push(&mut self, chunk: &str) -> Result<ToolParseResult> {
144 + self.buffer.push_str(chunk);
145 + let mut result = ToolParseResult::default();
146 +
147 + while let Some((event, consumed_len)) = parse_buffered_event(&self.buffer, |input| {
148 + parse_next_deepseek_json_event(input, &mut self.mode, self.format)
149 + })? {
150 + self.apply_event(event, &mut result)?;
151 + self.buffer.drain(..consumed_len);
152 + }
153 +
154 + Ok(result)
155 + }
156 +
157 + /// Flush buffered text and reset parser state.
158 + fn finish(&mut self) -> Result<ToolParseResult> {
159 + let mut result = ToolParseResult::default();
160 + match &self.mode {
161 + DeepSeekJsonMode::Text => result.normal_text.push_str(&self.buffer),
162 + DeepSeekJsonMode::ToolBlock | DeepSeekJsonMode::Done => {}
163 + DeepSeekJsonMode::Header | DeepSeekJsonMode::Arguments { .. } => {
164 + return Err(parsing_failed!(
165 + "incomplete {} tool call",
166 + self.format.parser_name()
167 + ));
168 + }
169 + }
170 + self.reset();
171 + Ok(result)
172 + }
173 +
174 + /// Reset all streaming state.
175 + fn reset(&mut self) {
176 + self.buffer.clear();
177 + self.mode = DeepSeekJsonMode::Text;
178 + self.active_tool_index = None;
179 + self.emitted_tool_count = 0;
180 + }
181 +}
182 +
183 +/// Parse a DeepSeek JSON event for the current parser mode.
184 +fn parse_next_deepseek_json_event(
185 + input: &mut DeepSeekJsonInput<'_>,
186 + mode: &mut DeepSeekJsonMode,
187 + format: DeepSeekJsonFormat,
188 +) -> ModalResult<DeepSeekJsonEvent> {
189 + match mode {
190 + DeepSeekJsonMode::Text => parse_text_event(input),
191 + DeepSeekJsonMode::ToolBlock => parse_tool_block_event(input),
192 + DeepSeekJsonMode::Header => tool_call_header_event(input, format),
193 + DeepSeekJsonMode::Arguments { json_scan } => {
194 + parse_arguments_event(input, format, json_scan)
195 + }
196 + DeepSeekJsonMode::Done => ignored_rest_event(input),
197 + }
198 +}
199 +
200 +/// Parse a text-mode DeepSeek JSON event.
201 +fn parse_text_event(input: &mut DeepSeekJsonInput<'_>) -> ModalResult<DeepSeekJsonEvent> {
202 + alt((tool_calls_start_event, safe_text_event)).parse_next(input)
203 +}
204 +
205 +/// Parse one event inside the DeepSeek tool-calls section.
206 +fn parse_tool_block_event(input: &mut DeepSeekJsonInput<'_>) -> ModalResult<DeepSeekJsonEvent> {
207 + ws0.void().parse_next(input)?;
208 + alt((tool_calls_end_event, tool_call_start_event)).parse_next(input)
209 +}
210 +
211 +/// Parse one event inside a DeepSeek tool-call arguments payload.
212 +fn parse_arguments_event(
213 + input: &mut DeepSeekJsonInput<'_>,
214 + format: DeepSeekJsonFormat,
215 + json_scan: &mut JsonObjectScanState,
216 +) -> ModalResult<DeepSeekJsonEvent> {
217 + if json_scan.complete() {
218 + tool_call_end_event(input, format)
219 + } else {
220 + argument_delta_event(input, json_scan)
221 + }
222 +}
223 +
224 +/// Parse a DeepSeek tool-calls start marker.
225 +fn tool_calls_start_event(input: &mut DeepSeekJsonInput<'_>) -> ModalResult<DeepSeekJsonEvent> {
226 + literal(TOOL_CALLS_START)
227 + .value(DeepSeekJsonEvent::ToolCallsStart)
228 + .parse_next(input)
229 +}
230 +
231 +/// Parse a DeepSeek tool-calls end marker.
232 +fn tool_calls_end_event(input: &mut DeepSeekJsonInput<'_>) -> ModalResult<DeepSeekJsonEvent> {
233 + literal(TOOL_CALLS_END).value(DeepSeekJsonEvent::ToolCallsEnd).parse_next(input)
234 +}
235 +
236 +/// Parse a DeepSeek tool-call start marker.
237 +fn tool_call_start_event(input: &mut DeepSeekJsonInput<'_>) -> ModalResult<DeepSeekJsonEvent> {
238 + literal(TOOL_CALL_START)
239 + .value(DeepSeekJsonEvent::ToolCallStart)
240 + .parse_next(input)
241 +}
242 +
243 +/// Parse a DeepSeek tool-call end marker.
244 +fn tool_call_end_event(
245 + input: &mut DeepSeekJsonInput<'_>,
246 + format: DeepSeekJsonFormat,
247 +) -> ModalResult<DeepSeekJsonEvent> {
248 + literal(format.argument_end_marker())
249 + .value(DeepSeekJsonEvent::ToolCallEnd)
250 + .parse_next(input)
251 +}
252 +
253 +/// Parse a DeepSeek tool-call header before the JSON arguments payload.
254 +fn tool_call_header_event(
255 + input: &mut DeepSeekJsonInput<'_>,
256 + format: DeepSeekJsonFormat,
257 +) -> ModalResult<DeepSeekJsonEvent> {
258 + match format {
259 + DeepSeekJsonFormat::V3 => v3_tool_call_header_event(input),
260 + DeepSeekJsonFormat::V31 => v31_tool_call_header_event(input),
261 + }
262 +}
263 +
264 +/// Parse a DeepSeek V3 tool-call header.
265 +fn v3_tool_call_header_event(input: &mut DeepSeekJsonInput<'_>) -> ModalResult<DeepSeekJsonEvent> {
266 + let name = seq!(
267 + _: literal("function"),
268 + _: literal(TOOL_CALL_SEPARATOR),
269 + take_until(1.., V3_JSON_START),
270 + _: literal(V3_JSON_START),
271 + )
272 + .parse_next(input)?;
273 +
274 + Ok(DeepSeekJsonEvent::ToolCallHeader {
275 + function_name: name.0.trim().to_string(),
276 + })
277 +}
278 +
279 +/// Parse a DeepSeek V3.1 tool-call header.
280 +fn v31_tool_call_header_event(input: &mut DeepSeekJsonInput<'_>) -> ModalResult<DeepSeekJsonEvent> {
281 + let (name, _) = (
282 + take_until(1.., TOOL_CALL_SEPARATOR),
283 + literal(TOOL_CALL_SEPARATOR),
284 + )
285 + .parse_next(input)?;
286 +
287 + Ok(DeepSeekJsonEvent::ToolCallHeader {
288 + function_name: name.trim().to_string(),
289 + })
290 +}
291 +
292 +/// Parse a DeepSeek raw JSON arguments delta.
293 +fn argument_delta_event(
294 + input: &mut DeepSeekJsonInput<'_>,
295 + json_scan: &mut JsonObjectScanState,
296 +) -> ModalResult<DeepSeekJsonEvent> {
297 + take_json_object(input, json_scan).map(|len| DeepSeekJsonEvent::Arguments { len })
298 +}
299 +
300 +/// Parse a safe text run before the next DeepSeek tool-calls section.
301 +fn safe_text_event(input: &mut DeepSeekJsonInput<'_>) -> ModalResult<DeepSeekJsonEvent> {
302 + safe_text_len(input, TOOL_CALLS_START).map(|len| DeepSeekJsonEvent::Text { len })
303 +}
304 +
305 +/// Parse ignored rest after the DeepSeek tool-calls section ends.
306 +fn ignored_rest_event(input: &mut DeepSeekJsonInput<'_>) -> ModalResult<DeepSeekJsonEvent> {
307 + rest.value(DeepSeekJsonEvent::IgnoredRest).parse_next(input)
308 +}
rust/src/tool-parser/src/error.rs new
+13
@@ -0,0 +1,13 @@
1 +use thiserror::Error;
2 +use thiserror_ext::Macro;
3 +
4 +/// Result alias for tool parser operations.
5 +pub type Result<T> = std::result::Result<T, ToolParserError>;
6 +
7 +/// Errors produced while creating or running tool parsers.
8 +#[derive(Debug, Error, Macro)]
9 +#[thiserror_ext(macro(path = "crate::error"))]
10 +pub enum ToolParserError {
11 + #[error("tool parser parsing failed: {message}")]
12 + ParsingFailed { message: String },
13 +}
rust/src/tool-parser/src/gemma4.rs new
+653
@@ -0,0 +1,653 @@
1 +use serde_json::{Map, Number, Value};
2 +use winnow::ascii::multispace0 as ws0;
3 +use winnow::combinator::{alt, delimited, opt, separated, seq, terminated};
4 +use winnow::error::{ContextError, ErrMode, ModalResult};
5 +use winnow::prelude::*;
6 +use winnow::stream::Partial;
7 +use winnow::token::{literal, take_till, take_until};
8 +
9 +use super::utils::{parse_buffered_event, safe_text_len};
10 +use super::{Result, ToolCallDelta, ToolParseResult, ToolParser};
11 +use crate::Tool;
12 +
13 +const TOOL_CALL_START: &str = "<|tool_call>";
14 +const TOOL_CALL_END: &str = "<tool_call|>";
15 +const STRING_DELIM: &str = "<|\"|>";
16 +const CALL_PREFIX: &str = "call:";
17 +
18 +type Gemma4Input<'i> = Partial<&'i str>;
19 +
20 +#[derive(Debug, Clone, PartialEq)]
21 +enum Gemma4Event {
22 + Text {
23 + len: usize,
24 + },
25 + ToolCall {
26 + name: String,
27 + args: Map<String, Value>,
28 + },
29 +}
30 +
31 +/// Tool parser for Google Gemma4 models.
32 +///
33 +/// Original Python implementation:
34 +/// <https://github.com/vllm-project/vllm/blob/bf45e6d0a558da2b8d7b60efb07b4aa394f3b60b/vllm/tool_parsers/gemma4_tool_parser.py>
35 +///
36 +/// Handles the Gemma4 function call format:
37 +///
38 +/// `<|tool_call>call:func_name{key:<|"|>value<|"|>}<tool_call|>`
39 +///
40 +/// Arguments are emitted only after a full Gemma4 tool call is parsed.
41 +pub struct Gemma4ToolParser {
42 + buffer: String,
43 + emitted_tool_count: usize,
44 +}
45 +
46 +impl Gemma4ToolParser {
47 + fn new(_tools: &[Tool]) -> Self {
48 + Self {
49 + buffer: String::new(),
50 + emitted_tool_count: 0,
51 + }
52 + }
53 +
54 + fn apply_event(&mut self, event: Gemma4Event, result: &mut ToolParseResult) -> Result<()> {
55 + match event {
56 + Gemma4Event::Text { len: consumed_len } => {
57 + result.normal_text.push_str(&self.buffer[..consumed_len]);
58 + }
59 + Gemma4Event::ToolCall { name, args } => {
60 + let arguments = serde_json::to_string(&args)
61 + .map_err(|error| parsing_failed!("failed to serialize arguments: {}", error))?;
62 +
63 + result.calls.push(ToolCallDelta {
64 + tool_index: self.emitted_tool_count,
65 + name: Some(name),
66 + arguments,
67 + });
68 + self.emitted_tool_count += 1;
69 + }
70 + }
71 + Ok(())
72 + }
73 +
74 + fn reset(&mut self) {
75 + self.buffer.clear();
76 + self.emitted_tool_count = 0;
77 + }
78 +}
79 +
80 +impl ToolParser for Gemma4ToolParser {
81 + fn create(tools: &[Tool]) -> Result<Box<dyn ToolParser>>
82 + where
83 + Self: Sized + 'static,
84 + {
85 + Ok(Box::new(Self::new(tools)))
86 + }
87 +
88 + fn preserve_special_tokens(&self) -> bool {
89 + true
90 + }
91 +
92 + fn push(&mut self, chunk: &str) -> Result<ToolParseResult> {
93 + self.buffer.push_str(chunk);
94 + let mut result = ToolParseResult::default();
95 +
96 + while let Some((event, consumed_len)) =
97 + parse_buffered_event(&self.buffer, parse_next_gemma4_event)?
98 + {
99 + self.apply_event(event, &mut result)?;
100 + self.buffer.drain(..consumed_len);
101 + }
102 +
103 + Ok(result)
104 + }
105 +
106 + fn finish(&mut self) -> Result<ToolParseResult> {
107 + let mut result = ToolParseResult::default();
108 +
109 + if !self.buffer.is_empty() {
110 + if self.buffer.starts_with(TOOL_CALL_START) {
111 + self.reset();
112 + return Err(parsing_failed!("incomplete Gemma4 tool call"));
113 + }
114 + result.normal_text.push_str(&self.buffer);
115 + }
116 +
117 + self.reset();
118 + Ok(result)
119 + }
120 +}
121 +
122 +/// Parse one Gemma4 event from buffered streaming input.
123 +fn parse_next_gemma4_event(input: &mut Gemma4Input<'_>) -> ModalResult<Gemma4Event> {
124 + alt((tool_call_event, safe_text_event)).parse_next(input)
125 +}
126 +
127 +/// Parse a complete Gemma4 tool call.
128 +// TODO: incremental parsing arguments to reduce scanning from O(n^2) to O(n).
129 +fn tool_call_event(input: &mut Gemma4Input<'_>) -> ModalResult<Gemma4Event> {
130 + let (name, args) = seq!(
131 + _: literal(TOOL_CALL_START),
132 + _: literal(CALL_PREFIX),
133 + gemma4_tool_name,
134 + _: literal("{"),
135 + gemma4_args,
136 + _: literal("}"),
137 + _: literal(TOOL_CALL_END),
138 + )
139 + .parse_next(input)?;
140 +
141 + Ok(Gemma4Event::ToolCall { name, args })
142 +}
143 +
144 +/// Parse a Gemma4 tool name.
145 +fn gemma4_tool_name(input: &mut Gemma4Input<'_>) -> ModalResult<String> {
146 + let name = take_until(1.., "{").parse_next(input)?.trim();
147 + if name.is_empty() {
148 + return Err(ErrMode::Cut(ContextError::new()));
149 + }
150 + Ok(name.to_string())
151 +}
152 +
153 +/// Parse a safe text run before the next Gemma4 marker.
154 +fn safe_text_event(input: &mut Gemma4Input<'_>) -> ModalResult<Gemma4Event> {
155 + safe_text_len(input, TOOL_CALL_START).map(|len| Gemma4Event::Text { len })
156 +}
157 +
158 +/// Parse Gemma4's custom key-value argument object content.
159 +fn gemma4_args(input: &mut Gemma4Input<'_>) -> ModalResult<Map<String, Value>> {
160 + let pairs: Vec<(String, Value)> = delimited(
161 + ws0,
162 + terminated(
163 + separated(0.., gemma4_pair, comma_separator),
164 + opt(comma_separator),
165 + ),
166 + ws0,
167 + )
168 + .parse_next(input)?;
169 + Ok(pairs.into_iter().collect())
170 +}
171 +
172 +/// Parse a Gemma4 key-value pair.
173 +fn gemma4_pair(input: &mut Gemma4Input<'_>) -> ModalResult<(String, Value)> {
174 + let (key, value) = seq!(
175 + _: ws0,
176 + gemma4_key,
177 + _: ws0,
178 + _: literal(":"),
179 + _: ws0,
180 + gemma4_value,
181 + )
182 + .parse_next(input)?;
183 + Ok((key, value))
184 +}
185 +
186 +/// Parse a Gemma4 bare key.
187 +fn gemma4_key(input: &mut Gemma4Input<'_>) -> ModalResult<String> {
188 + let key = take_till(1.., |char: char| char == ':').parse_next(input)?.trim();
189 + if key.is_empty() {
190 + return Err(ErrMode::Cut(ContextError::new()));
191 + }
192 + Ok(key.to_string())
193 +}
194 +
195 +/// Parse a Gemma4 value.
196 +fn gemma4_value(input: &mut Gemma4Input<'_>) -> ModalResult<Value> {
197 + alt((
198 + gemma4_string.map(|value: &str| Value::String(value.to_string())),
199 + gemma4_object.map(Value::Object),
200 + gemma4_array_value.map(Value::Array),
201 + gemma4_bare_value,
202 + ))
203 + .parse_next(input)
204 +}
205 +
206 +/// Parse a Gemma4 string delimited by `<|"|>`.
207 +fn gemma4_string<'i>(input: &mut Gemma4Input<'i>) -> ModalResult<&'i str> {
208 + delimited(
209 + literal(STRING_DELIM),
210 + take_until(0.., STRING_DELIM),
211 + literal(STRING_DELIM),
212 + )
213 + .parse_next(input)
214 +}
215 +
216 +/// Parse a nested Gemma4 object.
217 +fn gemma4_object(input: &mut Gemma4Input<'_>) -> ModalResult<Map<String, Value>> {
218 + delimited(literal("{"), gemma4_args, literal("}")).parse_next(input)
219 +}
220 +
221 +/// Parse a Gemma4 array value.
222 +fn gemma4_array_value(input: &mut Gemma4Input<'_>) -> ModalResult<Vec<Value>> {
223 + delimited(literal("["), gemma4_array_content, literal("]")).parse_next(input)
224 +}
225 +
226 +/// Parse Gemma4 array content.
227 +fn gemma4_array_content(input: &mut Gemma4Input<'_>) -> ModalResult<Vec<Value>> {
228 + delimited(
229 + ws0,
230 + terminated(
231 + separated(0.., gemma4_value, comma_separator),
232 + opt(comma_separator),
233 + ),
234 + ws0,
235 + )
236 + .parse_next(input)
237 +}
238 +
239 +/// Parse a Gemma4 bare scalar.
240 +fn gemma4_bare_value(input: &mut Gemma4Input<'_>) -> ModalResult<Value> {
241 + take_till(1.., |char: char| matches!(char, ',' | '}' | ']'))
242 + .map(parse_gemma4_scalar)
243 + .parse_next(input)
244 +}
245 +
246 +/// Parse a Gemma4 comma separator.
247 +fn comma_separator(input: &mut Gemma4Input<'_>) -> ModalResult<()> {
248 + delimited(ws0, literal(","), ws0).void().parse_next(input)
249 +}
250 +
251 +fn parse_gemma4_scalar(value: &str) -> Value {
252 + let value = value.trim();
253 + if value.is_empty() {
254 + return Value::String(String::new());
255 + }
256 + if value == "true" {
257 + return Value::Bool(true);
258 + }
259 + if value == "false" {
260 + return Value::Bool(false);
261 + }
262 + if matches!(value, "null" | "none" | "nil" | "NULL" | "None" | "NIL") {
263 + return Value::Null;
264 + }
265 + if value.contains('.') {
266 + if let Ok(parsed) = value.parse::<f64>()
267 + && let Some(number) = Number::from_f64(parsed)
268 + {
269 + return Value::Number(number);
270 + }
271 + } else if let Ok(parsed) = value.parse::<i64>() {
272 + return Value::Number(Number::from(parsed));
273 + }
274 +
275 + Value::String(value.to_string())
276 +}
277 +
278 +#[cfg(test)]
279 +mod tests {
280 + use serde_json::{Value, json};
281 + use thiserror_ext::AsReport;
282 + use winnow::combinator::{eof, terminated};
283 + use winnow::error::ErrMode;
284 + use winnow::prelude::*;
285 + use winnow::stream::Partial;
286 +
287 + use super::{
288 + Gemma4ToolParser, ToolCallDelta, ToolParseResult, ToolParser, gemma4_args,
289 + gemma4_array_content,
290 + };
291 + use crate::Tool;
292 +
293 + fn parse_gemma4_args(args: &str) -> super::Result<serde_json::Map<String, Value>> {
294 + let mut input = Partial::new(args);
295 + let _ = input.complete();
296 + match terminated(gemma4_args, eof).parse_next(&mut input) {
297 + Ok(value) => Ok(value),
298 + Err(ErrMode::Incomplete(_)) => Err(parsing_failed!("incomplete Gemma4 arguments")),
299 + Err(ErrMode::Backtrack(error) | ErrMode::Cut(error)) => {
300 + Err(parsing_failed!("{}", error))
301 + }
302 + }
303 + }
304 +
305 + fn parse_gemma4_array(array: &str) -> super::Result<Vec<Value>> {
306 + let mut input = Partial::new(array);
307 + let _ = input.complete();
308 + match terminated(gemma4_array_content, eof).parse_next(&mut input) {
309 + Ok(value) => Ok(value),
310 + Err(ErrMode::Incomplete(_)) => Err(parsing_failed!("incomplete Gemma4 array")),
311 + Err(ErrMode::Backtrack(error) | ErrMode::Cut(error)) => {
312 + Err(parsing_failed!("{}", error))
313 + }
314 + }
315 + }
316 +
317 + fn test_tools() -> Vec<Tool> {
318 + vec![
319 + Tool {
320 + name: "get_weather".to_string(),
321 + description: None,
322 + parameters: json!({ "type": "object" }),
323 + strict: None,
324 + },
325 + Tool {
326 + name: "get_time".to_string(),
327 + description: None,
328 + parameters: json!({ "type": "object" }),
329 + strict: None,
330 + },
331 + Tool {
332 + name: "write_file".to_string(),
333 + description: None,
334 + parameters: json!({ "type": "object" }),
335 + strict: None,
336 + },
337 + Tool {
338 + name: "Edit".to_string(),
339 + description: None,
340 + parameters: json!({ "type": "object" }),
341 + strict: None,
342 + },
343 + Tool {
344 + name: "search".to_string(),
345 + description: None,
346 + parameters: json!({ "type": "object" }),
347 + strict: None,
348 + },
349 + Tool {
350 + name: "set".to_string(),
351 + description: None,
352 + parameters: json!({ "type": "object" }),
353 + strict: None,
354 + },
355 + Tool {
356 + name: "get_status".to_string(),
357 + description: None,
358 + parameters: json!({ "type": "object" }),
359 + strict: None,
360 + },
361 + Tool {
362 + name: "todowrite".to_string(),
363 + description: None,
364 + parameters: json!({ "type": "object" }),
365 + strict: None,
366 + },
367 + ]
368 + }
369 +
370 + fn collect_stream(chunks: &[&str]) -> ToolParseResult {
371 + let mut parser = Gemma4ToolParser::new(&test_tools());
372 + let mut result = ToolParseResult::default();
373 + for chunk in chunks {
374 + result.append(parser.push(chunk).unwrap());
375 + }
376 + result.append(parser.finish().unwrap());
377 + result.coalesce_calls()
378 + }
379 +
380 + fn first_call(result: &ToolParseResult) -> &ToolCallDelta {
381 + result.calls.first().expect("expected one tool call")
382 + }
383 +
384 + #[test]
385 + fn gemma4_parse_args_handles_scalars_and_nested_values() {
386 + let parsed = parse_gemma4_args(
387 + "name:<|\"|>test<|\"|>,count:42,active:true,score:114.514,nested:{inner:<|\"|>value<|\"|>},items:[<|\"|>a<|\"|>,<|\"|>b<|\"|>]",
388 + )
389 + .unwrap();
390 +
391 + assert_eq!(
392 + Value::Object(parsed),
393 + json!({
394 + "name": "test",
395 + "count": 42,
396 + "active": true,
397 + "score": 114.514,
398 + "nested": { "inner": "value" },
399 + "items": ["a", "b"],
400 + })
401 + );
402 + }
403 +
404 + #[test]
405 + fn gemma4_parse_args_handles_empty_arguments() {
406 + let parsed = parse_gemma4_args("").unwrap();
407 + assert_eq!(Value::Object(parsed), json!({}));
408 + }
409 +
410 + #[test]
411 + fn gemma4_parse_array_handles_bare_values() {
412 + let parsed = parse_gemma4_array("42,true,114.514").unwrap();
413 + assert_eq!(Value::Array(parsed), json!([42, true, 114.514]));
414 + }
415 +
416 + #[test]
417 + fn gemma4_parse_complete_extracts_single_tool_call() {
418 + let mut parser = Gemma4ToolParser::new(&test_tools());
419 + let result = parser
420 + .parse_complete("<|tool_call>call:get_weather{location:<|\"|>London<|\"|>}<tool_call|>")
421 + .unwrap();
422 +
423 + assert!(result.normal_text.is_empty());
424 + assert_eq!(result.calls.len(), 1);
425 + assert_eq!(first_call(&result).name.as_deref(), Some("get_weather"));
426 + assert_eq!(
427 + serde_json::from_str::<Value>(&first_call(&result).arguments).unwrap(),
428 + json!({ "location": "London" })
429 + );
430 + }
431 +
432 + #[test]
433 + fn gemma4_parse_complete_rejects_incomplete_tool_call() {
434 + let mut parser = Gemma4ToolParser::new(&test_tools());
435 + let error = parser
436 + .parse_complete("<|tool_call>call:get_weather{location:<|\"|>London")
437 + .unwrap_err();
438 +
439 + assert!(error.to_report_string().contains("incomplete Gemma4 tool call"));
440 + }
441 +
442 + #[test]
443 + fn gemma4_streaming_basic_single_tool_call() {
444 + let result = collect_stream(&[
445 + "<|tool_call>",
446 + "call:get_weather{",
447 + "location:<|\"|>Paris",
448 + ", France",
449 + "<|\"|>}",
450 + "<tool_call|>",
451 + ]);
452 +
453 + assert!(result.normal_text.is_empty());
454 + assert_eq!(first_call(&result).name.as_deref(), Some("get_weather"));
455 + assert_eq!(
456 + serde_json::from_str::<Value>(&first_call(&result).arguments).unwrap(),
457 + json!({ "location": "Paris, France" })
458 + );
459 + }
460 +
461 + #[test]
462 + fn gemma4_streaming_text_before_and_after_tool_call() {
463 + let result = collect_stream(&[
464 + "Let me check ",
465 + "the weather. ",
466 + "<|tool_call>",
467 + "call:get_weather{",
468 + "location:<|\"|>London<|\"|>}",
469 + "<tool_call|><",
470 + "div>",
471 + ]);
472 +
473 + assert_eq!(result.normal_text, "Let me check the weather. <div>");
474 + assert_eq!(first_call(&result).name.as_deref(), Some("get_weather"));
475 + assert_eq!(
476 + serde_json::from_str::<Value>(&first_call(&result).arguments).unwrap(),
477 + json!({ "location": "London" })
478 + );
479 + }
480 +
481 + #[test]
482 + fn gemma4_streaming_waits_for_complete_tool_call() {
483 + let mut parser = Gemma4ToolParser::new(&test_tools());
484 + let mut result = ToolParseResult::default();
485 +
486 + for chunk in [
487 + "<|tool_call>",
488 + "call:get_weather{",
489 + "location:<|\"|>Paris<|\"|>}",
490 + ] {
491 + result.append(parser.push(chunk).unwrap());
492 + assert!(result.calls.is_empty());
493 + }
494 +
495 + result.append(parser.push("<tool_call|>").unwrap());
496 + let result = result.coalesce_calls();
497 +
498 + assert_eq!(first_call(&result).name.as_deref(), Some("get_weather"));
499 + assert_eq!(
500 + serde_json::from_str::<Value>(&first_call(&result).arguments).unwrap(),
501 + json!({ "location": "Paris" })
502 + );
503 + }
504 +
505 + #[test]
506 + fn gemma4_streaming_handles_boolean_split_across_chunks() {
507 + let result = collect_stream(&[
508 + "<|tool_call>",
509 + "call:search{input:{all:tru",
510 + "e}}",
511 + "<tool_call|>",
512 + ]);
513 +
514 + assert_eq!(first_call(&result).name.as_deref(), Some("search"));
515 + assert_eq!(
516 + serde_json::from_str::<Value>(&first_call(&result).arguments).unwrap(),
517 + json!({ "input": { "all": true } })
518 + );
519 + }
520 +
521 + #[test]
522 + fn gemma4_streaming_handles_false_split_across_chunks() {
523 + let result = collect_stream(&["<|tool_call>", "call:set{flag:fals", "e}", "<tool_call|>"]);
524 +
525 + assert_eq!(first_call(&result).name.as_deref(), Some("set"));
526 + assert_eq!(
527 + serde_json::from_str::<Value>(&first_call(&result).arguments).unwrap(),
528 + json!({ "flag": false })
529 + );
530 + }
531 +
532 + #[test]
533 + fn gemma4_streaming_handles_number_split_across_chunks() {
534 + let result = collect_stream(&["<|tool_call>", "call:set{count:4", "2}", "<tool_call|>"]);
535 +
536 + assert_eq!(first_call(&result).name.as_deref(), Some("set"));
537 + assert_eq!(
538 + serde_json::from_str::<Value>(&first_call(&result).arguments).unwrap(),
539 + json!({ "count": 42 })
540 + );
541 + }
542 +
543 + #[test]
544 + fn gemma4_streaming_handles_split_string_delimiter() {
545 + let result = collect_stream(&[
546 + "<|tool_call>",
547 + "call:todowrite{",
548 + "content:<|\"|>Buy milk<|",
549 + "\"|>}",
550 + "<tool_call|>",
551 + ]);
552 +
553 + assert_eq!(first_call(&result).name.as_deref(), Some("todowrite"));
554 + assert_eq!(
555 + serde_json::from_str::<Value>(&first_call(&result).arguments).unwrap(),
556 + json!({ "content": "Buy milk" })
557 + );
558 + assert!(!first_call(&result).arguments.contains("<|"));
559 + }
560 +
561 + #[test]
562 + fn gemma4_streaming_handles_end_marker_literal_inside_string() {
563 + let result = collect_stream(&[
564 + "<|tool_call>",
565 + "call:todowrite{",
566 + "content:<|\"|>literal }<tool_call|> inside",
567 + "<|\"|>}",
568 + "<tool_call|>",
569 + ]);
570 +
571 + assert_eq!(first_call(&result).name.as_deref(), Some("todowrite"));
572 + assert_eq!(
573 + serde_json::from_str::<Value>(&first_call(&result).arguments).unwrap(),
574 + json!({ "content": "literal }<tool_call|> inside" })
575 + );
576 + }
577 +
578 + #[test]
579 + fn gemma4_streaming_handles_html_argument_without_duplication() {
580 + let result = collect_stream(&[
581 + "<|tool_call>",
582 + "call:write_file{",
583 + "path:<|\"|>index.html<|\"|>,",
584 + "content:<|\"|><!DOCTYPE html>\n<",
585 + "html lang=\"zh-CN\">\n<",
586 + "head>\n <",
587 + "meta charset=\"UTF-8\">\n <",
588 + "meta name=\"viewport\" content=\"width=device-width\">\n",
589 + "<|\"|>}",
590 + "<tool_call|>",
591 + ]);
592 +
593 + assert_eq!(first_call(&result).name.as_deref(), Some("write_file"));
594 + assert_eq!(
595 + serde_json::from_str::<Value>(&first_call(&result).arguments).unwrap(),
596 + json!({
597 + "path": "index.html",
598 + "content": "<!DOCTYPE html>\n<html lang=\"zh-CN\">\n<head>\n <meta charset=\"UTF-8\">\n <meta name=\"viewport\" content=\"width=device-width\">\n",
599 + })
600 + );
601 + }
602 +
603 + #[test]
604 + fn gemma4_streaming_trailing_bare_bool_is_not_duplicated() {
605 + let result = collect_stream(&[
606 + "<|tool_call>",
607 + "call:Edit{",
608 + "file_path:<|\"|>src/env.py<|\"|>,",
609 + "old_string:<|\"|>old_val<|\"|>,",
610 + "new_string:<|\"|>new_val<|\"|>,",
611 + "replace_all:",
612 + "false}",
613 + "<tool_call|>",
614 + ]);
615 +
616 + assert_eq!(first_call(&result).name.as_deref(), Some("Edit"));
617 + assert_eq!(
618 + serde_json::from_str::<Value>(&first_call(&result).arguments).unwrap(),
619 + json!({
620 + "file_path": "src/env.py",
621 + "old_string": "old_val",
622 + "new_string": "new_val",
623 + "replace_all": false,
624 + })
625 + );
626 + assert_eq!(
627 + first_call(&result).arguments.matches("replace_all").count(),
628 + 1
629 + );
630 + }
631 +
632 + #[test]
633 + fn gemma4_finish_flushes_partial_start_marker_as_text() {
634 + let mut parser = Gemma4ToolParser::new(&test_tools());
635 + let mut result = parser.push("<").unwrap();
636 + result.append(parser.finish().unwrap());
637 +
638 + assert_eq!(result.normal_text, "<");
639 + assert!(result.calls.is_empty());
640 + }
641 +
642 + #[test]
643 + fn gemma4_finish_rejects_complete_args_without_end_marker() {
644 + let mut parser = Gemma4ToolParser::new(&test_tools());
645 + for chunk in ["<|tool_call>", "call:get_status{}"] {
646 + parser.push(chunk).unwrap();
647 + }
648 +
649 + let error = parser.finish().unwrap_err();
650 +
651 + assert!(error.to_report_string().contains("incomplete Gemma4 tool call"));
652 + }
653 +}
rust/src/tool-parser/src/glm_xml/glm45_moe.rs new
+43
@@ -0,0 +1,43 @@
1 +use super::{GlmXmlToolParser, Separator};
2 +use crate::{Result, Tool, ToolParseResult, ToolParser};
3 +
4 +/// Tool parser for GLM-4.5/4.6 MoE XML-style tool calls.
5 +///
6 +/// Example tool call content:
7 +///
8 +/// ```text
9 +/// <tool_call>get_weather
10 +/// <arg_key>city</arg_key>
11 +/// <arg_value>Hangzhou</arg_value>
12 +/// </tool_call>
13 +/// ```
14 +///
15 +/// Arguments are emitted only after a full `tool_call` block is parsed.
16 +pub struct Glm45MoeToolParser(GlmXmlToolParser);
17 +
18 +impl Glm45MoeToolParser {
19 + /// Create a GLM-4.5/4.6 MoE tool parser.
20 + pub(super) fn new(tools: &[Tool]) -> Self {
21 + Self(GlmXmlToolParser::new(tools, Separator::Newline))
22 + }
23 +}
24 +
25 +impl ToolParser for Glm45MoeToolParser {
26 + /// Create a boxed GLM-4.5/4.6 MoE tool parser.
27 + fn create(tools: &[Tool]) -> Result<Box<dyn ToolParser>>
28 + where
29 + Self: Sized + 'static,
30 + {
31 + Ok(Box::new(Self::new(tools)))
32 + }
33 +
34 + /// Push one decoded text chunk through the GLM MoE parser.
35 + fn push(&mut self, chunk: &str) -> Result<ToolParseResult> {
36 + self.0.push(chunk)
37 + }
38 +
39 + /// Flush buffered text and reset parser state.
40 + fn finish(&mut self) -> Result<ToolParseResult> {
41 + self.0.finish()
42 + }
43 +}
rust/src/tool-parser/src/glm_xml/glm47_moe.rs new
+135
@@ -0,0 +1,135 @@
1 +use super::{GlmXmlToolParser, Separator};
2 +use crate::{Result, Tool, ToolParseResult, ToolParser};
3 +
4 +/// Tool parser for GLM-4.7 MoE XML-style tool calls.
5 +///
6 +/// GLM-4.7 reuses the GLM-4.5 parser with a more flexible function-name
7 +/// separator, so the name may be followed by whitespace, a newline, or the
8 +/// first `<arg_key>` tag directly.
9 +pub struct Glm47MoeToolParser(GlmXmlToolParser);
10 +
11 +impl Glm47MoeToolParser {
12 + fn new(tools: &[Tool]) -> Self {
13 + Self(GlmXmlToolParser::new(tools, Separator::Flexible))
14 + }
15 +}
16 +
17 +impl ToolParser for Glm47MoeToolParser {
18 + fn create(tools: &[Tool]) -> Result<Box<dyn ToolParser>>
19 + where
20 + Self: Sized + 'static,
21 + {
22 + Ok(Box::new(Self::new(tools)))
23 + }
24 +
25 + fn push(&mut self, chunk: &str) -> Result<ToolParseResult> {
26 + self.0.push(chunk)
27 + }
28 +
29 + fn finish(&mut self) -> Result<ToolParseResult> {
30 + self.0.finish()
31 + }
32 +}
33 +
34 +#[cfg(test)]
35 +mod tests {
36 + use serde_json::{Value, json};
37 +
38 + use super::{Glm47MoeToolParser, ToolParser};
39 + use crate::test_utils::{collect_stream, split_by_chars, test_tools};
40 +
41 + fn glm47_tool_call(function_name: &str, params: &[(&str, &str)]) -> String {
42 + let params = params
43 + .iter()
44 + .map(|(name, value)| format!("<arg_key>{name}</arg_key><arg_value>{value}</arg_value>"))
45 + .collect::<Vec<_>>()
46 + .join("");
47 + format!("<tool_call>{function_name}{params}</tool_call>")
48 + }
49 +
50 + #[test]
51 + fn glm47_parse_complete_extracts_single_tool_call() {
52 + let mut parser = Glm47MoeToolParser::new(&test_tools());
53 + let output = format!(
54 + "Let me search for that.\n{}",
55 + glm47_tool_call(
56 + "get_weather",
57 + &[("city", "Beijing"), ("date", "2024-12-25")]
58 + )
59 + );
60 +
61 + let result = parser.parse_complete(&output).unwrap();
62 +
63 + assert_eq!(result.normal_text, "Let me search for that.\n");
64 + assert_eq!(result.calls.len(), 1);
65 + assert_eq!(result.calls[0].name.as_deref(), Some("get_weather"));
66 + assert_eq!(
67 + serde_json::from_str::<Value>(&result.calls[0].arguments).unwrap(),
68 + json!({"city": "Beijing", "date": "2024-12-25"})
69 + );
70 + }
71 +
72 + #[test]
73 + fn glm47_streaming_extracts_multiple_tool_calls() {
74 + let mut parser = Glm47MoeToolParser::new(&test_tools());
75 + let output = format!(
76 + "{}{}",
77 + glm47_tool_call("get_weather", &[("city", "Shanghai")]),
78 + glm47_tool_call("add", &[("x", "1"), ("y", "2")])
79 + );
80 +
81 + let chunks = split_by_chars(&output, 7);
82 + let result = collect_stream(&mut parser, &chunks);
83 +
84 + assert_eq!(result.normal_text, "");
85 + assert_eq!(result.calls.len(), 2);
86 + assert_eq!(result.calls[0].name.as_deref(), Some("get_weather"));
87 + assert_eq!(result.calls[1].name.as_deref(), Some("add"));
88 + assert_eq!(
89 + serde_json::from_str::<Value>(&result.calls[1].arguments).unwrap(),
90 + json!({"x": 1, "y": 2})
91 + );
92 + }
93 +
94 + #[test]
95 + fn glm47_parse_complete_converts_schema_types() {
96 + let mut parser = Glm47MoeToolParser::new(&test_tools());
97 + let result = parser
98 + .parse_complete(&glm47_tool_call(
99 + "convert",
100 + &[
101 + ("whole", "42"),
102 + ("flag", "true"),
103 + ("payload", r#"{"nested":{"key":"value"}}"#),
104 + ("items", "[1, 2, 3]"),
105 + ("empty", ""),
106 + ],
107 + ))
108 + .unwrap();
109 +
110 + assert_eq!(
111 + serde_json::from_str::<Value>(&result.calls[0].arguments).unwrap(),
112 + json!({
113 + "whole": 42,
114 + "flag": true,
115 + "payload": {"nested": {"key": "value"}},
116 + "items": [1, 2, 3],
117 + "empty": ""
118 + })
119 + );
120 + }
121 +
122 + #[test]
123 + fn glm47_parse_complete_extracts_zero_argument_call() {
124 + let mut parser = Glm47MoeToolParser::new(&test_tools());
125 +
126 + let result = parser.parse_complete("<tool_call>add</tool_call>").unwrap();
127 +
128 + assert_eq!(result.calls.len(), 1);
129 + assert_eq!(result.calls[0].name.as_deref(), Some("add"));
130 + assert_eq!(
131 + serde_json::from_str::<Value>(&result.calls[0].arguments).unwrap(),
132 + json!({})
133 + );
134 + }
135 +}
rust/src/tool-parser/src/glm_xml/mod.rs new
+435
@@ -0,0 +1,435 @@
1 +use winnow::ascii::multispace0 as ws0;
2 +use winnow::combinator::{alt, eof, repeat, seq, terminated};
3 +use winnow::prelude::*;
4 +use winnow::stream::Partial;
5 +use winnow::token::{literal, rest, take_until, take_while};
6 +
7 +use super::parameters::ToolSchemas;
8 +use super::utils::{parse_buffered_event, safe_text_len, xml_unescape};
9 +use super::{Result, ToolCallDelta, ToolParseResult};
10 +use crate::Tool;
11 +
12 +mod glm45_moe;
13 +mod glm47_moe;
14 +
15 +pub use glm45_moe::Glm45MoeToolParser;
16 +pub use glm47_moe::Glm47MoeToolParser;
17 +
18 +const TOOL_CALL_START: &str = "<tool_call>";
19 +const TOOL_CALL_END: &str = "</tool_call>";
20 +const ARG_KEY_START: &str = "<arg_key>";
21 +const ARG_KEY_END: &str = "</arg_key>";
22 +const ARG_VALUE_START: &str = "<arg_value>";
23 +const ARG_VALUE_END: &str = "</arg_value>";
24 +
25 +type GlmInput<'i> = Partial<&'i str>;
26 +
27 +#[derive(Debug, Clone, Copy, PartialEq, Eq)]
28 +enum GlmMode {
29 + Text,
30 + ToolCall,
31 + AfterToolCall,
32 +}
33 +
34 +#[derive(Debug, Clone, Copy, PartialEq, Eq)]
35 +enum Separator {
36 + /// GLM-4.5/4.6 format: function name must end at a newline before
37 + /// arguments.
38 + Newline,
39 + /// GLM-4.7 format: function name may end at whitespace or directly before
40 + /// `<arg_key>`.
41 + Flexible,
42 +}
43 +
44 +#[derive(Debug, Clone, PartialEq, Eq)]
45 +enum GlmEvent {
46 + Text {
47 + len: usize,
48 + },
49 + ToolCallStart,
50 + ToolCall {
51 + name: String,
52 + raw_params: Vec<(String, String)>,
53 + },
54 + IgnoredRest,
55 +}
56 +
57 +/// Tool parser core for GLM XML-style tool calls.
58 +struct GlmXmlToolParser {
59 + buffer: String,
60 + mode: GlmMode,
61 + emitted_tool_count: usize,
62 + tool_parameters: ToolSchemas,
63 + separator: Separator,
64 +}
65 +
66 +impl GlmXmlToolParser {
67 + /// Create a GLM XML tool parser with a function-name separator.
68 + fn new(tools: &[Tool], separator: Separator) -> Self {
69 + Self {
70 + buffer: String::new(),
71 + mode: GlmMode::Text,
72 + emitted_tool_count: 0,
73 + tool_parameters: ToolSchemas::from_tools(tools),
74 + separator,
75 + }
76 + }
77 +
78 + /// Apply one parsed GLM event to parser state and output.
79 + fn apply_event(&mut self, event: GlmEvent, result: &mut ToolParseResult) -> Result<()> {
80 + match event {
81 + GlmEvent::Text { len: consumed_len } => {
82 + result.normal_text.push_str(&self.buffer[..consumed_len]);
83 + }
84 + GlmEvent::ToolCallStart => self.mode = GlmMode::ToolCall,
85 + GlmEvent::ToolCall { name, raw_params } => {
86 + self.mode = GlmMode::AfterToolCall;
87 + let arguments = self.tool_parameters.convert_params_with_schema(&name, raw_params);
88 + let arguments = serde_json::to_string(&arguments)
89 + .map_err(|error| parsing_failed!("failed to serialize arguments: {}", error))?;
90 +
91 + result.calls.push(ToolCallDelta {
92 + tool_index: self.emitted_tool_count,
93 + name: Some(name),
94 + arguments,
95 + });
96 + self.emitted_tool_count += 1;
97 + }
98 + GlmEvent::IgnoredRest => {}
99 + }
100 + Ok(())
101 + }
102 +
103 + /// Reset all streaming state.
104 + fn reset(&mut self) {
105 + self.buffer.clear();
106 + self.mode = GlmMode::Text;
107 + self.emitted_tool_count = 0;
108 + }
109 +
110 + /// Push one decoded text chunk through the GLM MoE parser.
111 + fn push(&mut self, chunk: &str) -> Result<ToolParseResult> {
112 + self.buffer.push_str(chunk);
113 + let mut result = ToolParseResult::default();
114 +
115 + while let Some((event, consumed_len)) = parse_buffered_event(&self.buffer, |input| {
116 + parse_next_glm_event(input, self.mode, self.separator)
117 + })? {
118 + self.apply_event(event, &mut result)?;
119 + self.buffer.drain(..consumed_len);
120 + }
121 +
122 + Ok(result)
123 + }
124 +
125 + /// Flush buffered text and reset parser state.
126 + fn finish(&mut self) -> Result<ToolParseResult> {
127 + let mut result = ToolParseResult::default();
128 + if !self.buffer.is_empty() {
129 + match self.mode {
130 + GlmMode::Text => result.normal_text.push_str(&self.buffer),
131 + GlmMode::ToolCall => return Err(parsing_failed!("incomplete GLM MoE tool call")),
132 + GlmMode::AfterToolCall => {}
133 + }
134 + }
135 + self.reset();
136 + Ok(result)
137 + }
138 +}
139 +
140 +/// Parse a GLM event for the current parser mode.
141 +fn parse_next_glm_event(
142 + input: &mut GlmInput<'_>,
143 + mode: GlmMode,
144 + separator: Separator,
145 +) -> ModalResult<GlmEvent> {
146 + match mode {
147 + GlmMode::Text => parse_text_event(input),
148 + GlmMode::ToolCall => tool_call_event(input, separator),
149 + GlmMode::AfterToolCall => after_tool_call_event(input),
150 + }
151 +}
152 +
153 +/// Parse a text-mode GLM event.
154 +fn parse_text_event(input: &mut GlmInput<'_>) -> ModalResult<GlmEvent> {
155 + alt((tool_call_start_event, safe_text_event)).parse_next(input)
156 +}
157 +
158 +/// Parse a GLM tool-call start marker.
159 +fn tool_call_start_event(input: &mut GlmInput<'_>) -> ModalResult<GlmEvent> {
160 + literal(TOOL_CALL_START).value(GlmEvent::ToolCallStart).parse_next(input)
161 +}
162 +
163 +/// Parse a safe text run before the next GLM marker.
164 +fn safe_text_event(input: &mut GlmInput<'_>) -> ModalResult<GlmEvent> {
165 + safe_text_len(input, TOOL_CALL_START).map(|len| GlmEvent::Text { len })
166 +}
167 +
168 +/// Parse text after a completed GLM tool call.
169 +fn after_tool_call_event(input: &mut GlmInput<'_>) -> ModalResult<GlmEvent> {
170 + ws0.void().parse_next(input)?;
171 + alt((tool_call_start_event, ignored_rest_event)).parse_next(input)
172 +}
173 +
174 +/// Parse a trailing rest after GLM tool calls.
175 +fn ignored_rest_event(input: &mut GlmInput<'_>) -> ModalResult<GlmEvent> {
176 + rest.value(GlmEvent::IgnoredRest).parse_next(input)
177 +}
178 +
179 +/// Parse a complete GLM tool call.
180 +fn tool_call_event(input: &mut GlmInput<'_>, separator: Separator) -> ModalResult<GlmEvent> {
181 + let (body,) = seq!(
182 + take_until(0.., TOOL_CALL_END),
183 + _: literal(TOOL_CALL_END),
184 + )
185 + .parse_next(input)?;
186 +
187 + parse_tool_call_body(body, separator)
188 +}
189 +
190 +/// Parse a GLM tool-call body.
191 +fn parse_tool_call_body(body: &str, separator: Separator) -> ModalResult<GlmEvent> {
192 + let mut input = body;
193 + let (name, raw_params) = match separator {
194 + Separator::Newline => seq!(
195 + _: ws0,
196 + parse_newline_separated_function_name,
197 + parse_parameters,
198 + _: ws0,
199 + _: eof,
200 + )
201 + .parse_next(&mut input)?,
202 + Separator::Flexible => seq!(
203 + _: ws0,
204 + parse_flexible_function_name,
205 + parse_parameters,
206 + _: ws0,
207 + _: eof,
208 + )
209 + .parse_next(&mut input)?,
210 + };
211 +
212 + Ok(GlmEvent::ToolCall {
213 + name: name.to_string(),
214 + raw_params,
215 + })
216 +}
217 +
218 +/// Parse a GLM-4.5 newline-separated function name.
219 +fn parse_newline_separated_function_name<'i>(input: &mut &'i str) -> ModalResult<&'i str> {
220 + terminated(take_until(1.., "\n"), "\n").map(str::trim).parse_next(input)
221 +}
222 +
223 +/// Parse a GLM-4.7 whitespace-or-tag-separated function name.
224 +fn parse_flexible_function_name<'i>(input: &mut &'i str) -> ModalResult<&'i str> {
225 + terminated(
226 + take_while(1.., |ch: char| !ch.is_whitespace() && ch != '<'),
227 + ws0,
228 + )
229 + .parse_next(input)
230 +}
231 +
232 +/// Parse GLM argument key-value pairs.
233 +fn parse_parameters(input: &mut &str) -> ModalResult<Vec<(String, String)>> {
234 + repeat(0.., terminated(parse_parameter, ws0)).parse_next(input)
235 +}
236 +
237 +/// Parse a GLM argument key-value pair.
238 +fn parse_parameter(input: &mut &str) -> ModalResult<(String, String)> {
239 + let (key, value) = seq!(
240 + _: literal(ARG_KEY_START),
241 + take_until(1.., ARG_KEY_END),
242 + _: literal(ARG_KEY_END),
243 + _: ws0,
244 + _: literal(ARG_VALUE_START),
245 + take_until(0.., ARG_VALUE_END).map(str::trim).map(xml_unescape),
246 + _: literal(ARG_VALUE_END),
247 + )
248 + .parse_next(input)?;
249 +
250 + Ok((key.trim().to_string(), value.into_owned()))
251 +}
252 +
253 +#[cfg(test)]
254 +mod tests {
255 + use serde_json::{Value, json};
256 + use thiserror_ext::AsReport;
257 +
258 + use super::Glm45MoeToolParser;
259 + use crate::ToolParser;
260 + use crate::test_utils::{collect_stream, split_by_chars, test_tools};
261 +
262 + fn glm45_tool_call(function_name: &str, params: &[(&str, &str)]) -> String {
263 + let params = params
264 + .iter()
265 + .map(|(name, value)| {
266 + format!("<arg_key>{name}</arg_key>\n<arg_value>{value}</arg_value>")
267 + })
268 + .collect::<Vec<_>>()
269 + .join("\n");
270 + format!("<tool_call>{function_name}\n{params}\n</tool_call>")
271 + }
272 +
273 + #[test]
274 + fn glm45_parse_complete_without_tool_call_keeps_text() {
275 + let mut parser = Glm45MoeToolParser::new(&test_tools());
276 + let result = parser.parse_complete("Hello, world!").unwrap();
277 +
278 + assert_eq!(result.normal_text, "Hello, world!");
279 + assert!(result.calls.is_empty());
280 + }
281 +
282 + #[test]
283 + fn glm45_parse_complete_extracts_single_tool_call() {
284 + let mut parser = Glm45MoeToolParser::new(&test_tools());
285 + let output = format!(
286 + "Let me search for that.\n{}",
287 + glm45_tool_call(
288 + "get_weather",
289 + &[("city", "Beijing"), ("date", "2024-12-25")]
290 + )
291 + );
292 +
293 + let result = parser.parse_complete(&output).unwrap();
294 +
295 + assert_eq!(result.normal_text, "Let me search for that.\n");
296 + assert_eq!(result.calls.len(), 1);
297 + assert_eq!(result.calls[0].name.as_deref(), Some("get_weather"));
298 + assert_eq!(
299 + serde_json::from_str::<Value>(&result.calls[0].arguments).unwrap(),
300 + json!({"city": "Beijing", "date": "2024-12-25"})
301 + );
302 + }
303 +
304 + #[test]
305 + fn glm45_streaming_extracts_multiple_tool_calls() {
306 + let mut parser = Glm45MoeToolParser::new(&test_tools());
307 + let output = format!(
308 + "{}\n{}",
309 + glm45_tool_call("get_weather", &[("city", "Shanghai")]),
310 + glm45_tool_call("add", &[("x", "1"), ("y", "2")])
311 + );
312 +
313 + let chunks = split_by_chars(&output, 11);
314 + let result = collect_stream(&mut parser, &chunks);
315 +
316 + assert_eq!(result.normal_text, "");
317 + assert_eq!(result.calls.len(), 2);
318 + assert_eq!(result.calls[0].name.as_deref(), Some("get_weather"));
319 + assert_eq!(result.calls[1].name.as_deref(), Some("add"));
320 + assert_eq!(
321 + serde_json::from_str::<Value>(&result.calls[1].arguments).unwrap(),
322 + json!({"x": 1, "y": 2})
323 + );
324 + }
325 +
326 + #[test]
327 + fn glm45_parse_complete_unescapes_literal_closing_tags_in_arg_value() {
328 + let mut parser = Glm45MoeToolParser::new(&test_tools());
329 + let result = parser
330 + .parse_complete(&glm45_tool_call(
331 + "get_weather",
332 + &[
333 + ("city", "Paris &lt;/arg_value&gt;&lt;/tool_call&gt;"),
334 + ("date", "2026-05-08"),
335 + ],
336 + ))
337 + .unwrap();
338 +
339 + assert_eq!(
340 + serde_json::from_str::<Value>(&result.calls[0].arguments).unwrap(),
341 + json!({
342 + "city": "Paris </arg_value></tool_call>",
343 + "date": "2026-05-08",
344 + })
345 + );
346 + }
347 +
348 + #[test]
349 + fn glm45_streaming_without_tool_call_emits_text_incrementally() {
350 + let mut parser = Glm45MoeToolParser::new(&test_tools());
351 +
352 + let result = collect_stream(&mut parser, &["hello ", "world"]);
353 +
354 + assert_eq!(result.normal_text, "hello world");
355 + assert!(result.calls.is_empty());
356 + }
357 +
358 + #[test]
359 + fn glm45_streaming_preserves_prefix_text() {
360 + let mut parser = Glm45MoeToolParser::new(&test_tools());
361 +
362 + let result = collect_stream(
363 + &mut parser,
364 + &[
365 + "Prefix ",
366 + &glm45_tool_call("get_weather", &[("city", "Hangzhou")]),
367 + ],
368 + );
369 +
370 + assert_eq!(result.normal_text, "Prefix ");
371 + assert_eq!(result.calls.len(), 1);
372 + }
373 +
374 + #[test]
375 + fn glm45_streaming_handles_start_token_split_across_chunks() {
376 + let mut parser = Glm45MoeToolParser::new(&test_tools());
377 + let result = collect_stream(
378 + &mut parser,
379 + &[
380 + "hello <tool",
381 + "_call>get_weather\n",
382 + "<arg_key>city</arg_key><arg_value>Paris</arg_value></tool_call>",
383 + ],
384 + );
385 +
386 + assert_eq!(result.normal_text, "hello ");
387 + assert_eq!(result.calls.len(), 1);
388 + assert_eq!(result.calls[0].name.as_deref(), Some("get_weather"));
389 + }
390 +
391 + #[test]
392 + fn glm45_streaming_does_not_emit_incomplete_tool_call() {
393 + let mut parser = Glm45MoeToolParser::new(&test_tools());
394 +
395 + let result = parser.push("<tool_call>get_weather\n<arg_key>city</arg_key>").unwrap();
396 +
397 + assert_eq!(result.normal_text, "");
398 + assert!(result.calls.is_empty());
399 + }
400 +
401 + #[test]
402 + fn glm45_finish_fails_incomplete_tool_call() {
403 + let mut parser = Glm45MoeToolParser::new(&test_tools());
404 +
405 + parser.push("<tool_call>get_weather\n<arg_key>city</arg_key>").unwrap();
406 + let error = parser.finish().unwrap_err();
407 +
408 + assert!(error.as_report().to_string().contains("incomplete GLM MoE tool call"));
409 + }
410 +
411 + #[test]
412 + fn glm45_malformed_tool_call_fails_fast() {
413 + let mut parser = Glm45MoeToolParser::new(&test_tools());
414 +
415 + let error = parser.push("<tool_call>get_weather<arg_key>city</arg_key><arg_value>Paris</arg_value></tool_call>").unwrap_err();
416 +
417 + assert!(error.as_report().to_string().contains("tool parser parsing failed"));
418 + }
419 +
420 + #[test]
421 + fn glm45_streaming_ignores_trailing_text_after_tool_calls() {
422 + let mut parser = Glm45MoeToolParser::new(&test_tools());
423 +
424 + let result = collect_stream(
425 + &mut parser,
426 + &[&format!(
427 + "{}<|endoftext|>",
428 + glm45_tool_call("get_weather", &[("city", "Paris")])
429 + )],
430 + );
431 +
432 + assert_eq!(result.normal_text, "");
433 + assert_eq!(result.calls.len(), 1);
434 + }
435 +}
rust/src/tool-parser/src/json/hermes.rs new
+221
@@ -0,0 +1,221 @@
1 +use super::{JsonToolCallConfig, JsonToolCallParser, JsonToolCallWhitespace};
2 +use crate::{Result, Tool, ToolParseResult, ToolParser};
3 +
4 +const HERMES_CONFIG: JsonToolCallConfig = JsonToolCallConfig {
5 + parser_name: "Hermes",
6 + start_marker: "<tool_call>",
7 + end_marker: "</tool_call>",
8 + marker_whitespace: JsonToolCallWhitespace::Optional,
9 + delimiter: None,
10 + name_key: "name",
11 + arguments_key: "arguments",
12 +};
13 +
14 +/// Tool parser for Hermes XML-wrapped JSON tool calls.
15 +///
16 +/// Example tool call content:
17 +///
18 +/// ```text
19 +/// <tool_call>{"name": "get_weather", "arguments": {"location":"Tokyo"}}</tool_call>
20 +/// ```
21 +///
22 +/// Arguments are already OpenAI-style JSON text, so they are streamed as raw
23 +/// argument deltas without schema conversion or JSON normalization.
24 +///
25 +/// Note: parallel calls are represented as repeated
26 +/// `<tool_call>...</tool_call>` blocks, not as multiple calls inside one tag.
27 +pub struct HermesToolParser {
28 + inner: JsonToolCallParser,
29 +}
30 +
31 +impl HermesToolParser {
32 + /// Create a Hermes tool parser.
33 + fn new(_tools: &[Tool]) -> Self {
34 + Self {
35 + inner: JsonToolCallParser::new(HERMES_CONFIG),
36 + }
37 + }
38 +}
39 +
40 +impl ToolParser for HermesToolParser {
41 + /// Create a boxed Hermes tool parser.
42 + fn create(tools: &[Tool]) -> Result<Box<dyn ToolParser>>
43 + where
44 + Self: Sized + 'static,
45 + {
46 + Ok(Box::new(Self::new(tools)))
47 + }
48 +
49 + /// Push one decoded text chunk through the Hermes parser.
50 + fn push(&mut self, chunk: &str) -> Result<ToolParseResult> {
51 + self.inner.push(chunk)
52 + }
53 +
54 + /// Flush buffered text and reset parser state.
55 + fn finish(&mut self) -> Result<ToolParseResult> {
56 + self.inner.finish()
57 + }
58 +}
59 +
60 +#[cfg(test)]
61 +mod tests {
62 + use expect_test::expect;
63 + use thiserror_ext::AsReport;
64 +
65 + use super::HermesToolParser;
66 + use crate::test_utils::{collect_stream, split_by_chars, test_tools};
67 + use crate::{ToolParseResult, ToolParser};
68 +
69 + fn build_tool_call(function_name: &str, arguments: &str) -> String {
70 + format!(r#"<tool_call>{{"name":"{function_name}","arguments":{arguments}}}</tool_call>"#)
71 + }
72 +
73 + #[test]
74 + fn hermes_parse_complete_without_tool_call_keeps_text() {
75 + let mut parser = HermesToolParser::new(&test_tools());
76 + let result = parser.parse_complete("Hello, world!").unwrap();
77 +
78 + assert_eq!(result.normal_text, "Hello, world!");
79 + assert!(result.calls.is_empty());
80 + }
81 +
82 + #[test]
83 + fn hermes_parse_complete_extracts_raw_json_arguments() {
84 + let mut parser = HermesToolParser::new(&test_tools());
85 + let arguments = r#"{ "location": "Tokyo", "days": "3" }"#;
86 + let result = parser
87 + .parse_complete(&format!(
88 + "Let me check.\n{}",
89 + build_tool_call("get_weather", arguments)
90 + ))
91 + .unwrap();
92 +
93 + assert_eq!(result.normal_text, "Let me check.\n");
94 + assert_eq!(result.calls.len(), 1);
95 + assert_eq!(result.calls[0].tool_index, 0);
96 + assert_eq!(result.calls[0].name.as_deref(), Some("get_weather"));
97 + assert_eq!(result.calls[0].arguments, arguments);
98 + }
99 +
100 + #[test]
101 + fn hermes_accepts_newline_after_tool_call_start() {
102 + let mut parser = HermesToolParser::new(&test_tools());
103 + let result = parser
104 + .parse_complete(
105 + r#"<tool_call>
106 +{"name":"get_weather","arguments":{}}</tool_call>"#,
107 + )
108 + .unwrap();
109 +
110 + assert_eq!(result.calls.len(), 1);
111 + assert_eq!(result.calls[0].name.as_deref(), Some("get_weather"));
112 + }
113 +
114 + #[test]
115 + fn hermes_does_not_validate_or_normalize_arguments() {
116 + let mut parser = HermesToolParser::new(&test_tools());
117 + let arguments = r#"{"location":"Tokyo",}"#;
118 + let result = parser.parse_complete(&build_tool_call("get_weather", arguments)).unwrap();
119 +
120 + assert_eq!(result.calls[0].arguments, arguments);
121 + }
122 +
123 + #[test]
124 + fn hermes_streaming_emits_argument_deltas() {
125 + let mut parser = HermesToolParser::new(&test_tools());
126 + let chunks = [
127 + "preface <tool",
128 + "_call>{\"name\":\"get_weather\",\"arguments\":",
129 + "{\"location\":",
130 + "\"Beijing\"",
131 + "}",
132 + "}</tool_call> suffix",
133 + ];
134 +
135 + let mut result = ToolParseResult::default();
136 + let mut observed_arguments = Vec::new();
137 + for chunk in chunks {
138 + let next = parser.push(chunk).unwrap();
139 + observed_arguments.extend(
140 + next.calls
141 + .iter()
142 + .filter(|call| call.name.is_none())
143 + .map(|call| call.arguments.clone()),
144 + );
145 + result.append(next);
146 + }
147 + result.append(parser.finish().unwrap());
148 +
149 + assert_eq!(observed_arguments, ["{\"location\":", "\"Beijing\"", "}"]);
150 + assert_eq!(result.normal_text, "preface suffix");
151 + assert_eq!(
152 + result.coalesce_calls().calls[0].arguments,
153 + r#"{"location":"Beijing"}"#
154 + );
155 + }
156 +
157 + #[test]
158 + fn hermes_streaming_handles_split_markers() {
159 + let input = format!(
160 + "hello {}",
161 + build_tool_call("get_weather", r#"{"location":"Tokyo"}"#)
162 + );
163 + let chunks = split_by_chars(&input, 5);
164 + let mut parser = HermesToolParser::new(&test_tools());
165 +
166 + let result = collect_stream(&mut parser, &chunks);
167 +
168 + assert_eq!(result.normal_text, "hello ");
169 + assert_eq!(result.calls.len(), 1);
170 + assert_eq!(result.calls[0].arguments, r#"{"location":"Tokyo"}"#);
171 + }
172 +
173 + #[test]
174 + fn hermes_streaming_extracts_multiple_tool_calls() {
175 + let input = format!(
176 + "{}{}",
177 + build_tool_call("get_weather", r#"{"location":"Shanghai"}"#),
178 + build_tool_call("add", r#"{"x":1,"y":2}"#),
179 + );
180 + let chunks = split_by_chars(&input, 7);
181 + let mut parser = HermesToolParser::new(&test_tools());
182 +
183 + let result = collect_stream(&mut parser, &chunks);
184 +
185 + expect![[r#"
186 + ToolParseResult {
187 + normal_text: "",
188 + calls: [
189 + ToolCallDelta {
190 + tool_index: 0,
191 + name: Some(
192 + "get_weather",
193 + ),
194 + arguments: "{\"location\":\"Shanghai\"}",
195 + },
196 + ToolCallDelta {
197 + tool_index: 1,
198 + name: Some(
199 + "add",
200 + ),
201 + arguments: "{\"x\":1,\"y\":2}",
202 + },
203 + ],
204 + }
205 + "#]]
206 + .assert_debug_eq(&result);
207 + }
208 +
209 + #[test]
210 + fn hermes_finish_fails_incomplete_tool_call() {
211 + let mut parser = HermesToolParser::new(&test_tools());
212 + parser
213 + .push(r#"<tool_call>{"name":"get_weather","arguments":{"location""#)
214 + .unwrap();
215 +
216 + let error = parser.finish().unwrap_err();
217 +
218 + expect!["tool parser parsing failed: incomplete Hermes tool call"]
219 + .assert_eq(&error.to_report_string());
220 + }
221 +}
rust/src/tool-parser/src/json/llama.rs new
+487
@@ -0,0 +1,487 @@
1 +use winnow::ascii::multispace0 as ws0;
2 +use winnow::combinator::seq;
3 +use winnow::error::{ModalResult, StrContext};
4 +use winnow::prelude::*;
5 +use winnow::token::literal;
6 +
7 +use super::{
8 + JsonToolCallConfig, JsonToolCallEvent, JsonToolCallWhitespace, JsonToolInput,
9 + argument_delta_event, tool_call_header_event,
10 +};
11 +use crate::utils::{JsonObjectScanState, parse_buffered_event};
12 +use crate::{Result, Tool, ToolCallDelta, ToolParseResult, ToolParser};
13 +
14 +#[derive(Debug, Clone, PartialEq, Eq)]
15 +enum LlamaJsonMode {
16 + Start,
17 + Header,
18 + Arguments { json_scan: JsonObjectScanState },
19 + AfterCall,
20 + Passthrough,
21 +}
22 +
23 +#[derive(Debug, Clone, PartialEq, Eq)]
24 +enum LlamaJsonEvent {
25 + ToolCallHeader { function_name: String },
26 + Arguments { len: usize },
27 + ToolCallClose,
28 + Separator,
29 +}
30 +
31 +/// Tool parser for strict Llama JSON-template tool calls.
32 +///
33 +/// Example tool call content:
34 +///
35 +/// ```text
36 +/// {"name":"get_weather","parameters":{"location":"Tokyo"}}; {"name":"add","parameters":{"x":1,"y":2}}
37 +/// ```
38 +///
39 +/// Arguments are already OpenAI-style JSON text, so they are streamed as raw
40 +/// argument deltas without schema conversion or JSON normalization.
41 +///
42 +/// Natural text at the beginning of the stream permanently disables tool
43 +/// parsing for that assistant output.
44 +pub struct Llama3JsonToolParser {
45 + buffer: String,
46 + mode: LlamaJsonMode,
47 + active_tool_index: Option<usize>,
48 + emitted_tool_count: usize,
49 +}
50 +
51 +impl Llama3JsonToolParser {
52 + /// Create a Llama JSON tool parser.
53 + fn new(_tools: &[Tool]) -> Self {
54 + Self {
55 + buffer: String::new(),
56 + mode: LlamaJsonMode::Start,
57 + active_tool_index: None,
58 + emitted_tool_count: 0,
59 + }
60 + }
61 +
62 + /// Commit the stream to JSON parsing or permanent passthrough.
63 + fn commit_start(&mut self) -> bool {
64 + if !matches!(self.mode, LlamaJsonMode::Start) {
65 + return true;
66 + }
67 +
68 + if self.buffer.is_empty() {
69 + return false;
70 + }
71 +
72 + if self.buffer.starts_with('{') {
73 + self.mode = LlamaJsonMode::Header;
74 + } else {
75 + self.mode = LlamaJsonMode::Passthrough;
76 + }
77 + true
78 + }
79 +
80 + /// Apply one parsed Llama JSON event to parser state and output.
81 + fn apply_event(&mut self, event: LlamaJsonEvent, result: &mut ToolParseResult) -> Result<()> {
82 + match event {
83 + LlamaJsonEvent::ToolCallHeader { function_name } => {
84 + let tool_index = self.emitted_tool_count;
85 + self.emitted_tool_count += 1;
86 + self.active_tool_index = Some(tool_index);
87 + self.mode = LlamaJsonMode::Arguments {
88 + json_scan: JsonObjectScanState::default(),
89 + };
90 + result.calls.push(ToolCallDelta {
91 + tool_index,
92 + name: Some(function_name),
93 + arguments: String::new(),
94 + });
95 + }
96 + LlamaJsonEvent::Arguments { len: consumed_len } => {
97 + let Some(tool_index) = self.active_tool_index else {
98 + return Err(parsing_failed!(
99 + "Llama JSON arguments without an active tool call"
100 + ));
101 + };
102 + result.calls.push(ToolCallDelta {
103 + tool_index,
104 + name: None,
105 + arguments: self.buffer[..consumed_len].to_string(),
106 + });
107 + }
108 + LlamaJsonEvent::ToolCallClose => {
109 + self.active_tool_index = None;
110 + self.mode = LlamaJsonMode::AfterCall;
111 + }
112 + LlamaJsonEvent::Separator => {
113 + self.active_tool_index = None;
114 + self.mode = LlamaJsonMode::Header;
115 + }
116 + }
117 + Ok(())
118 + }
119 +
120 + /// Reset all streaming state.
121 + fn reset(&mut self) {
122 + self.buffer.clear();
123 + self.mode = LlamaJsonMode::Start;
124 + self.active_tool_index = None;
125 + self.emitted_tool_count = 0;
126 + }
127 +}
128 +
129 +impl ToolParser for Llama3JsonToolParser {
130 + /// Create a boxed Llama JSON tool parser.
131 + fn create(tools: &[Tool]) -> Result<Box<dyn ToolParser>>
132 + where
133 + Self: Sized + 'static,
134 + {
135 + Ok(Box::new(Self::new(tools)))
136 + }
137 +
138 + /// Push one decoded text chunk through the Llama JSON parser.
139 + fn push(&mut self, chunk: &str) -> Result<ToolParseResult> {
140 + self.buffer.push_str(chunk);
141 + let mut result = ToolParseResult::default();
142 +
143 + if !self.commit_start() {
144 + return Ok(result);
145 + }
146 +
147 + if matches!(self.mode, LlamaJsonMode::Passthrough) {
148 + result.normal_text.push_str(&self.buffer);
149 + self.buffer.clear();
150 + return Ok(result);
151 + }
152 +
153 + while let Some((event, consumed_len)) = parse_buffered_event(&self.buffer, |input| {
154 + parse_next_llama_json_event(input, &mut self.mode)
155 + })? {
156 + self.apply_event(event, &mut result)?;
157 + self.buffer.drain(..consumed_len);
158 + }
159 +
160 + Ok(result)
161 + }
162 +
163 + /// Flush buffered text and reset parser state.
164 + fn finish(&mut self) -> Result<ToolParseResult> {
165 + let mut result = ToolParseResult::default();
166 + match &self.mode {
167 + LlamaJsonMode::Start | LlamaJsonMode::Passthrough => {
168 + result.normal_text.push_str(&self.buffer);
169 + }
170 + LlamaJsonMode::AfterCall if self.buffer.trim().is_empty() => {}
171 + LlamaJsonMode::Header | LlamaJsonMode::Arguments { .. } => {
172 + return Err(parsing_failed!("incomplete Llama JSON tool call"));
173 + }
174 + LlamaJsonMode::AfterCall => {
175 + return Err(parsing_failed!("invalid Llama JSON"));
176 + }
177 + }
178 + self.reset();
179 + Ok(result)
180 + }
181 +}
182 +
183 +/// Parse a Llama JSON event for the current parser mode.
184 +fn parse_next_llama_json_event(
185 + input: &mut JsonToolInput<'_>,
186 + mode: &mut LlamaJsonMode,
187 +) -> ModalResult<LlamaJsonEvent> {
188 + match mode {
189 + LlamaJsonMode::Start | LlamaJsonMode::Passthrough => {
190 + unreachable!("Llama JSON parser driver must commit before parsing events")
191 + }
192 + LlamaJsonMode::Header => llama_tool_call_header_event(input),
193 + LlamaJsonMode::Arguments { json_scan } => parse_llama_arguments_event(input, json_scan),
194 + LlamaJsonMode::AfterCall => after_call_event(input),
195 + }
196 +}
197 +
198 +/// Parse a Llama JSON tool-call header.
199 +fn llama_tool_call_header_event(input: &mut JsonToolInput<'_>) -> ModalResult<LlamaJsonEvent> {
200 + const CONFIG: JsonToolCallConfig = JsonToolCallConfig {
201 + parser_name: "Llama JSON",
202 + start_marker: "",
203 + end_marker: "",
204 + marker_whitespace: JsonToolCallWhitespace::Optional,
205 + delimiter: Some(";"),
206 + name_key: "name",
207 + arguments_key: "parameters",
208 + };
209 +
210 + match tool_call_header_event(input, CONFIG)? {
211 + JsonToolCallEvent::ToolCallHeader { function_name } => {
212 + Ok(LlamaJsonEvent::ToolCallHeader { function_name })
213 + }
214 + _ => unreachable!("tool_call_header_event only emits ToolCallHeader"),
215 + }
216 +}
217 +
218 +/// Parse one event inside a Llama JSON arguments payload.
219 +fn parse_llama_arguments_event(
220 + input: &mut JsonToolInput<'_>,
221 + json_scan: &mut JsonObjectScanState,
222 +) -> ModalResult<LlamaJsonEvent> {
223 + if json_scan.complete() {
224 + tool_call_close_event(input)
225 + } else {
226 + match argument_delta_event(input, json_scan)? {
227 + JsonToolCallEvent::Arguments { len } => Ok(LlamaJsonEvent::Arguments { len }),
228 + _ => unreachable!("argument_delta_event only emits Arguments"),
229 + }
230 + }
231 +}
232 +
233 +/// Parse the outer closing brace for one Llama JSON tool call.
234 +fn tool_call_close_event(input: &mut JsonToolInput<'_>) -> ModalResult<LlamaJsonEvent> {
235 + literal("}").value(LlamaJsonEvent::ToolCallClose).parse_next(input)
236 +}
237 +
238 +/// Parse a semicolon separator after one Llama JSON tool call.
239 +fn after_call_event(input: &mut JsonToolInput<'_>) -> ModalResult<LlamaJsonEvent> {
240 + seq!(
241 + _: ws0,
242 + _: literal(";"),
243 + _: ws0,
244 + )
245 + .value(LlamaJsonEvent::Separator)
246 + .context(StrContext::Label("Llama JSON"))
247 + .parse_next(input)
248 +}
249 +
250 +#[cfg(test)]
251 +mod tests {
252 + use expect_test::expect;
253 + use thiserror_ext::AsReport;
254 +
255 + use super::Llama3JsonToolParser;
256 + use crate::test_utils::{collect_stream, split_by_chars, test_tools};
257 + use crate::{ToolParseResult, ToolParser};
258 +
259 + fn build_tool_call(function_name: &str, parameters: &str) -> String {
260 + format!(r#"{{"name":"{function_name}","parameters":{parameters}}}"#)
261 + }
262 +
263 + #[test]
264 + fn llama_json_parse_complete_without_tool_call_keeps_text() {
265 + let mut parser = Llama3JsonToolParser::new(&test_tools());
266 + let result = parser.parse_complete("Hello, world!").unwrap();
267 +
268 + assert_eq!(result.normal_text, "Hello, world!");
269 + assert!(result.calls.is_empty());
270 + }
271 +
272 + #[test]
273 + fn llama_json_passthrough_never_reenters_tool_parsing() {
274 + let mut parser = Llama3JsonToolParser::new(&test_tools());
275 + let mut result = parser.push("plain text first ").unwrap();
276 + result.append(
277 + parser.push(&build_tool_call("get_weather", r#"{"location":"Tokyo"}"#)).unwrap(),
278 + );
279 + result.append(parser.finish().unwrap());
280 +
281 + assert_eq!(
282 + result.normal_text,
283 + r#"plain text first {"name":"get_weather","parameters":{"location":"Tokyo"}}"#
284 + );
285 + assert!(result.calls.is_empty());
286 + }
287 +
288 + #[test]
289 + fn llama_json_does_not_support_python_tag_prefix() {
290 + let mut parser = Llama3JsonToolParser::new(&test_tools());
291 + let input = format!(
292 + "<|python_tag|>{}",
293 + build_tool_call("get_weather", r#"{"location":"Tokyo"}"#)
294 + );
295 + let result = parser.parse_complete(&input).unwrap();
296 +
297 + assert_eq!(result.normal_text, input);
298 + assert!(result.calls.is_empty());
299 + }
300 +
301 + #[test]
302 + fn llama_json_rejects_leading_whitespace_before_tool_call() {
303 + let mut parser = Llama3JsonToolParser::new(&test_tools());
304 + let input = format!(
305 + "\n {}",
306 + build_tool_call("get_weather", r#"{"location":"Tokyo"}"#)
307 + );
308 + let result = parser.parse_complete(&input).unwrap();
309 +
310 + assert_eq!(result.normal_text, input);
311 + assert!(result.calls.is_empty());
312 + }
313 +
314 + #[test]
315 + fn llama_json_extracts_raw_parameters_object() {
316 + let mut parser = Llama3JsonToolParser::new(&test_tools());
317 + let arguments = r#"{ "location": "Tokyo", "days": 3 }"#;
318 + let result = parser.parse_complete(&build_tool_call("get_weather", arguments)).unwrap();
319 +
320 + assert_eq!(result.calls.len(), 1);
321 + assert_eq!(result.calls[0].tool_index, 0);
322 + assert_eq!(result.calls[0].name.as_deref(), Some("get_weather"));
323 + assert_eq!(result.calls[0].arguments, arguments);
324 + }
325 +
326 + #[test]
327 + fn llama_json_rejects_arguments_key() {
328 + let mut parser = Llama3JsonToolParser::new(&test_tools());
329 + let error = parser
330 + .parse_complete(r#"{"name":"get_weather","arguments":{"location":"Tokyo"}}"#)
331 + .unwrap_err();
332 +
333 + expect![[r#"
334 + tool parser parsing failed: invalid Llama JSON
335 + expected `parameters`"#]]
336 + .assert_eq(&error.to_report_string());
337 + }
338 +
339 + #[test]
340 + fn llama_json_extracts_multiple_semicolon_separated_calls() {
341 + let mut parser = Llama3JsonToolParser::new(&test_tools());
342 + let input = format!(
343 + "{} \n; {}",
344 + build_tool_call("get_weather", r#"{"location":"Shanghai"}"#),
345 + build_tool_call("add", r#"{"x":1,"y":2}"#),
346 + );
347 + let result = parser.parse_complete(&input).unwrap();
348 +
349 + expect![[r#"
350 + ToolParseResult {
351 + normal_text: "",
352 + calls: [
353 + ToolCallDelta {
354 + tool_index: 0,
355 + name: Some(
356 + "get_weather",
357 + ),
358 + arguments: "{\"location\":\"Shanghai\"}",
359 + },
360 + ToolCallDelta {
361 + tool_index: 1,
362 + name: Some(
363 + "add",
364 + ),
365 + arguments: "{\"x\":1,\"y\":2}",
366 + },
367 + ],
368 + }
369 + "#]]
370 + .assert_debug_eq(&result);
371 + }
372 +
373 + #[test]
374 + fn llama_json_streaming_emits_argument_deltas() {
375 + let mut parser = Llama3JsonToolParser::new(&test_tools());
376 + let chunks = [
377 + "{\"name\":\"get_weather\",\"parameters\":",
378 + "{\"location\":",
379 + "\"Beijing\"",
380 + "}}",
381 + ];
382 +
383 + let mut result = ToolParseResult::default();
384 + let mut observed_arguments = Vec::new();
385 + for chunk in chunks {
386 + let next = parser.push(chunk).unwrap();
387 + observed_arguments.extend(
388 + next.calls
389 + .iter()
390 + .filter(|call| call.name.is_none())
391 + .map(|call| call.arguments.clone()),
392 + );
393 + result.append(next);
394 + }
395 + result.append(parser.finish().unwrap());
396 +
397 + assert_eq!(observed_arguments, ["{\"location\":", "\"Beijing\"", "}"]);
398 + assert_eq!(
399 + result.coalesce_calls().calls[0].arguments,
400 + r#"{"location":"Beijing"}"#
401 + );
402 + }
403 +
404 + #[test]
405 + fn llama_json_streaming_handles_split_objects_and_separator() {
406 + let input = format!(
407 + "{};{}",
408 + build_tool_call("get_weather", r#"{"location":"Dallas","state":"TX"}"#),
409 + build_tool_call("add", r#"{"x":4,"y":5}"#),
410 + );
411 + let chunks = split_by_chars(&input, 6);
412 + let mut parser = Llama3JsonToolParser::new(&test_tools());
413 +
414 + let result = collect_stream(&mut parser, &chunks);
415 +
416 + assert_eq!(result.normal_text, "");
417 + assert_eq!(result.calls.len(), 2);
418 + assert_eq!(
419 + result.calls[0].arguments,
420 + r#"{"location":"Dallas","state":"TX"}"#
421 + );
422 + assert_eq!(result.calls[1].name.as_deref(), Some("add"));
423 + assert_eq!(result.calls[1].arguments, r#"{"x":4,"y":5}"#);
424 + }
425 +
426 + #[test]
427 + fn llama_json_handles_nested_multiline_and_escaped_string_parameters() {
428 + let mut parser = Llama3JsonToolParser::new(&test_tools());
429 + let arguments = r#"{
430 + "payload": {"items": [1, {"value": "literal { brace } and \"quote\""}]},
431 + "flag": true
432 +}"#;
433 + let result = parser.parse_complete(&build_tool_call("convert", arguments)).unwrap();
434 +
435 + assert_eq!(result.calls[0].arguments, arguments);
436 + }
437 +
438 + #[test]
439 + fn llama_json_keeps_trailing_whitespace_after_tool_call() {
440 + let mut parser = Llama3JsonToolParser::new(&test_tools());
441 + let result = parser
442 + .parse_complete(&format!(
443 + "{}\n\t ",
444 + build_tool_call("get_weather", r#"{"location":"Tokyo"}"#)
445 + ))
446 + .unwrap();
447 +
448 + assert_eq!(result.normal_text, "");
449 + assert_eq!(result.calls.len(), 1);
450 + }
451 +
452 + #[test]
453 + fn llama_json_finish_fails_incomplete_tool_call() {
454 + let mut parser = Llama3JsonToolParser::new(&test_tools());
455 + parser.push(r#"{"name":"get_weather","parameters":{"location""#).unwrap();
456 +
457 + let error = parser.finish().unwrap_err();
458 +
459 + expect!["tool parser parsing failed: incomplete Llama JSON tool call"]
460 + .assert_eq(&error.to_report_string());
461 + }
462 +
463 + #[test]
464 + fn llama_json_malformed_field_order_fails_fast() {
465 + let mut parser = Llama3JsonToolParser::new(&test_tools());
466 + let error = parser.push(r#"{"parameters":{},"name":"get_weather"}"#).unwrap_err();
467 +
468 + expect![[r#"
469 + tool parser parsing failed: invalid Llama JSON
470 + expected `name`"#]]
471 + .assert_eq(&error.to_report_string());
472 + }
473 +
474 + #[test]
475 + fn llama_json_trailing_non_separator_content_errors() {
476 + let mut parser = Llama3JsonToolParser::new(&test_tools());
477 + let error = parser
478 + .push(&format!(
479 + "{} trailing",
480 + build_tool_call("get_weather", r#"{"location":"Tokyo"}"#)
481 + ))
482 + .unwrap_err();
483 +
484 + expect!["tool parser parsing failed: invalid Llama JSON"]
485 + .assert_eq(&error.to_report_string());
486 + }
487 +}
rust/src/tool-parser/src/json/mistral.rs new
+240
@@ -0,0 +1,240 @@
1 +use super::{JsonToolCallConfig, JsonToolCallParser, JsonToolCallWhitespace};
2 +use crate::{Result, Tool, ToolParseResult, ToolParser};
3 +
4 +const MISTRAL_CONFIG: JsonToolCallConfig = JsonToolCallConfig {
5 + parser_name: "Mistral",
6 + start_marker: "[TOOL_CALLS] [",
7 + end_marker: "]",
8 + marker_whitespace: JsonToolCallWhitespace::Optional,
9 + delimiter: Some(","),
10 + name_key: "name",
11 + arguments_key: "arguments",
12 +};
13 +
14 +/// Tool parser for Mistral JSON-array tool calls.
15 +///
16 +/// Example tool call content:
17 +///
18 +/// ```text
19 +/// [TOOL_CALLS] [{"name": "get_weather", "arguments": {"location":"Tokyo"}}]
20 +/// ```
21 +///
22 +/// Arguments are already OpenAI-style JSON text, so they are streamed as raw
23 +/// argument deltas without schema conversion or JSON normalization.
24 +pub struct MistralToolParser {
25 + inner: JsonToolCallParser,
26 +}
27 +
28 +impl MistralToolParser {
29 + /// Create a Mistral tool parser.
30 + fn new(_tools: &[Tool]) -> Self {
31 + Self {
32 + inner: JsonToolCallParser::new(MISTRAL_CONFIG),
33 + }
34 + }
35 +}
36 +
37 +impl ToolParser for MistralToolParser {
38 + /// Create a boxed Mistral tool parser.
39 + fn create(tools: &[Tool]) -> Result<Box<dyn ToolParser>>
40 + where
41 + Self: Sized + 'static,
42 + {
43 + Ok(Box::new(Self::new(tools)))
44 + }
45 +
46 + /// Push one decoded text chunk through the Mistral parser.
47 + fn push(&mut self, chunk: &str) -> Result<ToolParseResult> {
48 + self.inner.push(chunk)
49 + }
50 +
51 + /// Flush buffered text and reset parser state.
52 + fn finish(&mut self) -> Result<ToolParseResult> {
53 + self.inner.finish()
54 + }
55 +}
56 +
57 +#[cfg(test)]
58 +mod tests {
59 + use expect_test::expect;
60 + use thiserror_ext::AsReport;
61 +
62 + use super::MistralToolParser;
63 + use crate::test_utils::{collect_stream, split_by_chars, test_tools};
64 + use crate::{ToolParseResult, ToolParser};
65 +
66 + fn build_tool_call(function_name: &str, arguments: &str) -> String {
67 + format!(r#"{{"name":"{function_name}","arguments":{arguments}}}"#)
68 + }
69 +
70 + fn build_tool_calls(tool_calls: &[String]) -> String {
71 + format!("[TOOL_CALLS] [{}]", tool_calls.join(","))
72 + }
73 +
74 + #[test]
75 + fn mistral_parse_complete_without_tool_call_keeps_text() {
76 + let mut parser = MistralToolParser::new(&test_tools());
77 + let result = parser.parse_complete("Hello, world!").unwrap();
78 +
79 + assert_eq!(result.normal_text, "Hello, world!");
80 + assert!(result.calls.is_empty());
81 + }
82 +
83 + #[test]
84 + fn mistral_parse_complete_extracts_raw_json_arguments() {
85 + let mut parser = MistralToolParser::new(&test_tools());
86 + let arguments = r#"{ "location": "Tokyo", "days": "3" }"#;
87 + let result = parser
88 + .parse_complete(&format!(
89 + "Let me check.\n{}",
90 + build_tool_calls(&[build_tool_call("get_weather", arguments)])
91 + ))
92 + .unwrap();
93 +
94 + assert_eq!(result.normal_text, "Let me check.\n");
95 + assert_eq!(result.calls.len(), 1);
96 + assert_eq!(result.calls[0].tool_index, 0);
97 + assert_eq!(result.calls[0].name.as_deref(), Some("get_weather"));
98 + assert_eq!(result.calls[0].arguments, arguments);
99 + }
100 +
101 + #[test]
102 + fn mistral_parse_complete_extracts_pretty_multiple_tool_calls() {
103 + let mut parser = MistralToolParser::new(&test_tools());
104 + let result = parser
105 + .parse_complete(
106 + r#"I'll help.
107 +[TOOL_CALLS] [
108 + {"name": "get_weather", "arguments": {"city": "Tokyo", "units": "celsius"}}
109 + ,
110 + {"name": "add", "arguments": {"x": 1, "y": 2}}
111 +]"#,
112 + )
113 + .unwrap();
114 +
115 + expect![[r#"
116 + ToolParseResult {
117 + normal_text: "I'll help.\n",
118 + calls: [
119 + ToolCallDelta {
120 + tool_index: 0,
121 + name: Some(
122 + "get_weather",
123 + ),
124 + arguments: "{\"city\": \"Tokyo\", \"units\": \"celsius\"}",
125 + },
126 + ToolCallDelta {
127 + tool_index: 1,
128 + name: Some(
129 + "add",
130 + ),
131 + arguments: "{\"x\": 1, \"y\": 2}",
132 + },
133 + ],
134 + }
135 + "#]]
136 + .assert_debug_eq(&result);
137 + }
138 +
139 + #[test]
140 + fn mistral_does_not_validate_or_normalize_arguments() {
141 + let mut parser = MistralToolParser::new(&test_tools());
142 + let arguments = r#"{"location":"Tokyo",}"#;
143 + let result = parser
144 + .parse_complete(&build_tool_calls(&[build_tool_call(
145 + "get_weather",
146 + arguments,
147 + )]))
148 + .unwrap();
149 +
150 + assert_eq!(result.calls[0].arguments, arguments);
151 + }
152 +
153 + #[test]
154 + fn mistral_streaming_emits_argument_deltas() {
155 + let mut parser = MistralToolParser::new(&test_tools());
156 + let chunks = [
157 + "preface [TOOL",
158 + "_CALLS] [{\"name\":\"get_weather\",\"arguments\":",
159 + "{\"location\":",
160 + "\"Beijing\"",
161 + "}",
162 + "}] suffix",
163 + ];
164 +
165 + let mut result = ToolParseResult::default();
166 + let mut observed_arguments = Vec::new();
167 + for chunk in chunks {
168 + let next = parser.push(chunk).unwrap();
169 + observed_arguments.extend(
170 + next.calls
171 + .iter()
172 + .filter(|call| call.name.is_none())
173 + .map(|call| call.arguments.clone()),
174 + );
175 + result.append(next);
176 + }
177 + result.append(parser.finish().unwrap());
178 +
179 + assert_eq!(observed_arguments, ["{\"location\":", "\"Beijing\"", "}"]);
180 + assert_eq!(result.normal_text, "preface suffix");
181 + assert_eq!(
182 + result.coalesce_calls().calls[0].arguments,
183 + r#"{"location":"Beijing"}"#
184 + );
185 + }
186 +
187 + #[test]
188 + fn mistral_streaming_handles_split_markers() {
189 + let input = format!(
190 + "hello {}",
191 + build_tool_calls(&[build_tool_call("get_weather", r#"{"location":"Tokyo"}"#)])
192 + );
193 + let chunks = split_by_chars(&input, 5);
194 + let mut parser = MistralToolParser::new(&test_tools());
195 +
196 + let result = collect_stream(&mut parser, &chunks);
197 +
198 + assert_eq!(result.normal_text, "hello ");
199 + assert_eq!(result.calls.len(), 1);
200 + assert_eq!(result.calls[0].arguments, r#"{"location":"Tokyo"}"#);
201 + }
202 +
203 + #[test]
204 + fn mistral_keeps_array_bracket_literal_inside_json_string() {
205 + let mut parser = MistralToolParser::new(&test_tools());
206 + let arguments = r#"{"text":"Array notation: arr[0] = value[1]"}"#;
207 + let result = parser
208 + .parse_complete(&build_tool_calls(&[build_tool_call("echo", arguments)]))
209 + .unwrap();
210 +
211 + assert_eq!(result.calls.len(), 1);
212 + assert_eq!(result.calls[0].arguments, arguments);
213 + }
214 +
215 + #[test]
216 + fn mistral_finish_fails_incomplete_tool_call() {
217 + let mut parser = MistralToolParser::new(&test_tools());
218 + parser
219 + .push(r#"[TOOL_CALLS] [{"name":"get_weather","arguments":{"location""#)
220 + .unwrap();
221 +
222 + let error = parser.finish().unwrap_err();
223 +
224 + expect!["tool parser parsing failed: incomplete Mistral tool call"]
225 + .assert_eq(&error.to_report_string());
226 + }
227 +
228 + #[test]
229 + fn mistral_malformed_field_order_fails_fast() {
230 + let mut parser = MistralToolParser::new(&test_tools());
231 + let error = parser
232 + .push(r#"[TOOL_CALLS] [{"arguments":{},"name":"get_weather"}]"#)
233 + .unwrap_err();
234 +
235 + expect![[r#"
236 + tool parser parsing failed: invalid Mistral
237 + expected `name`"#]]
238 + .assert_eq(&error.to_report_string());
239 + }
240 +}
rust/src/tool-parser/src/json/mod.rs new
+465
@@ -0,0 +1,465 @@
1 +//! Shared parser core for JSON tool calls wrapped by text markers.
2 +
3 +pub use hermes::HermesToolParser;
4 +pub use llama::Llama3JsonToolParser;
5 +pub use mistral::MistralToolParser;
6 +pub use qwen::Qwen3XmlToolParser;
7 +
8 +mod hermes;
9 +mod llama;
10 +mod mistral;
11 +mod qwen;
12 +
13 +use winnow::ascii::multispace0 as ws0;
14 +use winnow::combinator::{alt, seq};
15 +use winnow::error::{ModalResult, StrContext, StrContextValue};
16 +use winnow::prelude::*;
17 +use winnow::stream::Partial;
18 +use winnow::token::literal;
19 +
20 +use super::utils::{
21 + JsonObjectScanState, json_str, parse_buffered_event, safe_text_len, take_json_object,
22 +};
23 +use super::{Result, ToolCallDelta, ToolParseResult};
24 +
25 +type JsonToolInput<'i> = Partial<&'i str>;
26 +
27 +#[derive(Debug, Clone, Copy)]
28 +struct JsonToolCallConfig {
29 + parser_name: &'static str,
30 + start_marker: &'static str,
31 + end_marker: &'static str,
32 + marker_whitespace: JsonToolCallWhitespace,
33 + delimiter: Option<&'static str>,
34 + name_key: &'static str,
35 + arguments_key: &'static str,
36 +}
37 +
38 +#[derive(Debug, Clone, Copy)]
39 +enum JsonToolCallWhitespace {
40 + Optional,
41 + Exact(&'static str),
42 +}
43 +
44 +#[derive(Debug, Clone, PartialEq, Eq)]
45 +enum JsonToolCallMode {
46 + Text,
47 + Header,
48 + Arguments { json_scan: JsonObjectScanState },
49 +}
50 +
51 +#[derive(Debug, Clone, PartialEq, Eq)]
52 +enum JsonToolCallEvent {
53 + Text { len: usize },
54 + ToolCallStart,
55 + ToolCallHeader { function_name: String },
56 + Arguments { len: usize },
57 + ToolCallDelimiter,
58 + ToolCallEnd,
59 +}
60 +
61 +/// Tool parser core for marker-wrapped JSON tool calls.
62 +#[derive(Debug)]
63 +struct JsonToolCallParser {
64 + config: JsonToolCallConfig,
65 + buffer: String,
66 + mode: JsonToolCallMode,
67 + active_tool_index: Option<usize>,
68 + emitted_tool_count: usize,
69 +}
70 +
71 +impl JsonToolCallParser {
72 + /// Create a marker-wrapped JSON tool-call parser.
73 + fn new(config: JsonToolCallConfig) -> Self {
74 + Self {
75 + config,
76 + buffer: String::new(),
77 + mode: JsonToolCallMode::Text,
78 + active_tool_index: None,
79 + emitted_tool_count: 0,
80 + }
81 + }
82 +
83 + /// Push one decoded text chunk through the JSON tool-call parser.
84 + fn push(&mut self, chunk: &str) -> Result<ToolParseResult> {
85 + self.buffer.push_str(chunk);
86 + let mut result = ToolParseResult::default();
87 + let config = self.config;
88 +
89 + while let Some((event, consumed_len)) = parse_buffered_event(&self.buffer, |input| {
90 + parse_next_json_tool_call_event(input, &mut self.mode, config)
91 + })? {
92 + self.apply_event(event, &mut result)?;
93 + self.buffer.drain(..consumed_len);
94 + }
95 +
96 + Ok(result)
97 + }
98 +
99 + /// Flush buffered text and reset parser state.
100 + fn finish(&mut self) -> Result<ToolParseResult> {
101 + let mut result = ToolParseResult::default();
102 + match &self.mode {
103 + JsonToolCallMode::Text => result.normal_text.push_str(&self.buffer),
104 + JsonToolCallMode::Header | JsonToolCallMode::Arguments { .. } => {
105 + return Err(parsing_failed!(
106 + "incomplete {} tool call",
107 + self.config.parser_name
108 + ));
109 + }
110 + }
111 + self.reset();
112 + Ok(result)
113 + }
114 +
115 + /// Apply one parsed JSON tool-call event to parser state and output.
116 + fn apply_event(
117 + &mut self,
118 + event: JsonToolCallEvent,
119 + result: &mut ToolParseResult,
120 + ) -> Result<()> {
121 + match event {
122 + JsonToolCallEvent::Text { len: consumed_len } => {
123 + result.normal_text.push_str(&self.buffer[..consumed_len]);
124 + }
125 + JsonToolCallEvent::ToolCallStart => self.mode = JsonToolCallMode::Header,
126 + JsonToolCallEvent::ToolCallHeader { function_name } => {
127 + let tool_index = self.emitted_tool_count;
128 + self.emitted_tool_count += 1;
129 + self.active_tool_index = Some(tool_index);
130 + self.mode = JsonToolCallMode::Arguments {
131 + json_scan: JsonObjectScanState::default(),
132 + };
133 + result.calls.push(ToolCallDelta {
134 + tool_index,
135 + name: Some(function_name),
136 + arguments: String::new(),
137 + });
138 + }
139 + JsonToolCallEvent::Arguments { len: consumed_len } => {
140 + let Some(tool_index) = self.active_tool_index else {
141 + return Err(parsing_failed!(
142 + "{} arguments without an active tool call",
143 + self.config.parser_name
144 + ));
145 + };
146 + result.calls.push(ToolCallDelta {
147 + tool_index,
148 + name: None,
149 + arguments: self.buffer[..consumed_len].to_string(),
150 + });
151 + }
152 + JsonToolCallEvent::ToolCallDelimiter => {
153 + self.active_tool_index = None;
154 + self.mode = JsonToolCallMode::Header;
155 + }
156 + JsonToolCallEvent::ToolCallEnd => {
157 + self.active_tool_index = None;
158 + self.mode = JsonToolCallMode::Text;
159 + }
160 + }
161 + Ok(())
162 + }
163 +
164 + /// Reset all streaming state.
165 + fn reset(&mut self) {
166 + self.buffer.clear();
167 + self.mode = JsonToolCallMode::Text;
168 + self.active_tool_index = None;
169 + self.emitted_tool_count = 0;
170 + }
171 +}
172 +
173 +/// Parse a JSON tool-call event for the current parser mode.
174 +fn parse_next_json_tool_call_event(
175 + input: &mut JsonToolInput<'_>,
176 + mode: &mut JsonToolCallMode,
177 + config: JsonToolCallConfig,
178 +) -> ModalResult<JsonToolCallEvent> {
179 + match mode {
180 + JsonToolCallMode::Text => parse_text_event(input, config),
181 + JsonToolCallMode::Header => tool_call_header_event(input, config),
182 + JsonToolCallMode::Arguments { json_scan } => {
183 + parse_arguments_event(input, json_scan, config)
184 + }
185 + }
186 +}
187 +
188 +/// Parse a text-mode JSON tool-call event.
189 +fn parse_text_event(
190 + input: &mut JsonToolInput<'_>,
191 + config: JsonToolCallConfig,
192 +) -> ModalResult<JsonToolCallEvent> {
193 + alt((
194 + |input: &mut JsonToolInput<'_>| tool_call_start_event(input, config),
195 + |input: &mut JsonToolInput<'_>| safe_text_event(input, config),
196 + ))
197 + .parse_next(input)
198 +}
199 +
200 +/// Parse a marker-wrapped JSON tool-call start marker.
201 +fn tool_call_start_event(
202 + input: &mut JsonToolInput<'_>,
203 + config: JsonToolCallConfig,
204 +) -> ModalResult<JsonToolCallEvent> {
205 + seq!(
206 + _: literal(config.start_marker),
207 + _: |input: &mut JsonToolInput<'_>| marker_whitespace(input, config),
208 + )
209 + .value(JsonToolCallEvent::ToolCallStart)
210 + .parse_next(input)
211 +}
212 +
213 +/// Parse a marker-wrapped JSON tool-call header before the raw arguments
214 +/// payload.
215 +fn tool_call_header_event(
216 + input: &mut JsonToolInput<'_>,
217 + config: JsonToolCallConfig,
218 +) -> ModalResult<JsonToolCallEvent> {
219 + let (function_name,) = seq!(
220 + _: ws0,
221 + _: literal("{"),
222 + _: ws0,
223 + _: |input: &mut JsonToolInput<'_>| json_key(input, config.name_key),
224 + _: ws0,
225 + _: literal(":"),
226 + _: ws0,
227 + json_str,
228 + _: ws0,
229 + _: literal(","),
230 + _: ws0,
231 + _: |input: &mut JsonToolInput<'_>| json_key(input, config.arguments_key),
232 + _: ws0,
233 + _: literal(":"),
234 + _: ws0,
235 + )
236 + .context(StrContext::Label(config.parser_name))
237 + .parse_next(input)?;
238 +
239 + Ok(JsonToolCallEvent::ToolCallHeader { function_name })
240 +}
241 +
242 +/// Parse a configured JSON object key.
243 +fn json_key(input: &mut JsonToolInput<'_>, key: &'static str) -> ModalResult<()> {
244 + seq!(
245 + _: literal("\""),
246 + _: literal(key).context(StrContext::Expected(StrContextValue::StringLiteral(key))),
247 + _: literal("\""),
248 + )
249 + .void()
250 + .parse_next(input)
251 +}
252 +
253 +/// Parse one event inside a marker-wrapped JSON tool-call arguments payload.
254 +fn parse_arguments_event(
255 + input: &mut JsonToolInput<'_>,
256 + json_scan: &mut JsonObjectScanState,
257 + config: JsonToolCallConfig,
258 +) -> ModalResult<JsonToolCallEvent> {
259 + if json_scan.complete() {
260 + tool_call_close_event(input, config)
261 + } else {
262 + argument_delta_event(input, json_scan)
263 + }
264 +}
265 +
266 +/// Parse a raw JSON arguments delta.
267 +fn argument_delta_event(
268 + input: &mut JsonToolInput<'_>,
269 + json_scan: &mut JsonObjectScanState,
270 +) -> ModalResult<JsonToolCallEvent> {
271 + take_json_object(input, json_scan).map(|len| JsonToolCallEvent::Arguments { len })
272 +}
273 +
274 +/// Parse a marker-wrapped JSON tool-call close marker.
275 +fn tool_call_close_event(
276 + input: &mut JsonToolInput<'_>,
277 + config: JsonToolCallConfig,
278 +) -> ModalResult<JsonToolCallEvent> {
279 + let _ = literal("}").parse_next(input)?;
280 +
281 + match config.delimiter {
282 + Some(delimiter) => alt((
283 + |input: &mut JsonToolInput<'_>| tool_call_end_event(input, config),
284 + |input: &mut JsonToolInput<'_>| tool_call_delimiter_event(input, delimiter),
285 + ))
286 + .parse_next(input),
287 + None => tool_call_end_event(input, config),
288 + }
289 +}
290 +
291 +/// Parse a marker-wrapped JSON tool-call end marker.
292 +fn tool_call_end_event(
293 + input: &mut JsonToolInput<'_>,
294 + config: JsonToolCallConfig,
295 +) -> ModalResult<JsonToolCallEvent> {
296 + seq!(
297 + _: |input: &mut JsonToolInput<'_>| marker_whitespace(input, config),
298 + _: literal(config.end_marker),
299 + )
300 + .value(JsonToolCallEvent::ToolCallEnd)
301 + .parse_next(input)
302 +}
303 +
304 +/// Parse a delimiter between JSON tool calls inside one marker block.
305 +fn tool_call_delimiter_event(
306 + input: &mut JsonToolInput<'_>,
307 + delimiter: &'static str,
308 +) -> ModalResult<JsonToolCallEvent> {
309 + seq!(
310 + _: ws0,
311 + _: literal(delimiter),
312 + _: ws0,
313 + )
314 + .value(JsonToolCallEvent::ToolCallDelimiter)
315 + .parse_next(input)
316 +}
317 +
318 +/// Parse configured whitespace around a marker-wrapped JSON tool call.
319 +fn marker_whitespace(input: &mut JsonToolInput<'_>, config: JsonToolCallConfig) -> ModalResult<()> {
320 + match config.marker_whitespace {
321 + JsonToolCallWhitespace::Optional => ws0.void().parse_next(input),
322 + JsonToolCallWhitespace::Exact(whitespace) => literal(whitespace).void().parse_next(input),
323 + }
324 +}
325 +
326 +/// Parse a safe text run before the next marker-wrapped JSON tool call.
327 +fn safe_text_event(
328 + input: &mut JsonToolInput<'_>,
329 + config: JsonToolCallConfig,
330 +) -> ModalResult<JsonToolCallEvent> {
331 + safe_text_len(input, config.start_marker).map(|len| JsonToolCallEvent::Text { len })
332 +}
333 +
334 +#[cfg(test)]
335 +mod tests {
336 + use expect_test::expect;
337 +
338 + use super::{JsonToolCallConfig, JsonToolCallParser, JsonToolCallWhitespace};
339 + use crate::ToolParseResult;
340 +
341 + const DELIMITED_CONFIG: JsonToolCallConfig = JsonToolCallConfig {
342 + parser_name: "Delimited JSON",
343 + start_marker: "<tool_calls>",
344 + end_marker: "</tool_calls>",
345 + marker_whitespace: JsonToolCallWhitespace::Optional,
346 + delimiter: Some("<"),
347 + name_key: "function",
348 + arguments_key: "parameters",
349 + };
350 +
351 + fn build_tool_call(function_name: &str, arguments: &str) -> String {
352 + format!(r#"{{"function":"{function_name}","parameters":{arguments}}}"#)
353 + }
354 +
355 + fn build_tool_calls(tool_calls: &[String]) -> String {
356 + format!("<tool_calls>{}</tool_calls>", tool_calls.join(" <\n"))
357 + }
358 +
359 + fn collect_chunks(parser: &mut JsonToolCallParser, chunks: &[&str]) -> ToolParseResult {
360 + let mut result = ToolParseResult::default();
361 + for chunk in chunks {
362 + result.append(parser.push(chunk).unwrap());
363 + }
364 + result.append(parser.finish().unwrap());
365 + result.coalesce_calls()
366 + }
367 +
368 + #[test]
369 + fn json_tool_call_delimiter_extracts_multiple_calls_in_one_block() {
370 + let input = build_tool_calls(&[
371 + build_tool_call("get_weather", r#"{"location":"Shanghai"}"#),
372 + build_tool_call("add", r#"{"x":1,"y":2}"#),
373 + ]);
374 + let mut parser = JsonToolCallParser::new(DELIMITED_CONFIG);
375 +
376 + let result = collect_chunks(&mut parser, &[&input]);
377 +
378 + expect![[r#"
379 + ToolParseResult {
380 + normal_text: "",
381 + calls: [
382 + ToolCallDelta {
383 + tool_index: 0,
384 + name: Some(
385 + "get_weather",
386 + ),
387 + arguments: "{\"location\":\"Shanghai\"}",
388 + },
389 + ToolCallDelta {
390 + tool_index: 1,
391 + name: Some(
392 + "add",
393 + ),
394 + arguments: "{\"x\":1,\"y\":2}",
395 + },
396 + ],
397 + }
398 + "#]]
399 + .assert_debug_eq(&result);
400 + }
401 +
402 + #[test]
403 + fn json_tool_call_delimiter_can_arrive_in_later_chunk() {
404 + let mut parser = JsonToolCallParser::new(DELIMITED_CONFIG);
405 + let chunks = [
406 + r#"<tool_calls>{"function":"get_weather","parameters":{"location":"Shanghai"}}"#,
407 + " <\n",
408 + r#"{"function":"add","parameters":{"x":1,"y":2}}"#,
409 + "</tool_calls>",
410 + ];
411 +
412 + let result = collect_chunks(&mut parser, &chunks);
413 +
414 + expect![[r#"
415 + ToolParseResult {
416 + normal_text: "",
417 + calls: [
418 + ToolCallDelta {
419 + tool_index: 0,
420 + name: Some(
421 + "get_weather",
422 + ),
423 + arguments: "{\"location\":\"Shanghai\"}",
424 + },
425 + ToolCallDelta {
426 + tool_index: 1,
427 + name: Some(
428 + "add",
429 + ),
430 + arguments: "{\"x\":1,\"y\":2}",
431 + },
432 + ],
433 + }
434 + "#]]
435 + .assert_debug_eq(&result);
436 + }
437 +
438 + #[test]
439 + fn json_tool_call_end_marker_wins_over_delimiter_prefix() {
440 + let mut parser = JsonToolCallParser::new(DELIMITED_CONFIG);
441 + let chunks = [
442 + r#"<tool_calls>{"function":"get_weather","parameters":{"location":"Shanghai"}}"#,
443 + " ",
444 + "</tool_calls> trailing text",
445 + ];
446 +
447 + let result = collect_chunks(&mut parser, &chunks);
448 +
449 + expect![[r#"
450 + ToolParseResult {
451 + normal_text: " trailing text",
452 + calls: [
453 + ToolCallDelta {
454 + tool_index: 0,
455 + name: Some(
456 + "get_weather",
457 + ),
458 + arguments: "{\"location\":\"Shanghai\"}",
459 + },
460 + ],
461 + }
462 + "#]]
463 + .assert_debug_eq(&result);
464 + }
465 +}
rust/src/tool-parser/src/json/qwen.rs new
+279
@@ -0,0 +1,279 @@
1 +use super::{JsonToolCallConfig, JsonToolCallParser, JsonToolCallWhitespace};
2 +use crate::{Result, Tool, ToolParseResult, ToolParser};
3 +
4 +const QWEN_XML_CONFIG: JsonToolCallConfig = JsonToolCallConfig {
5 + parser_name: "Qwen XML",
6 + start_marker: "<tool_call>",
7 + end_marker: "</tool_call>",
8 + marker_whitespace: JsonToolCallWhitespace::Exact("\n"),
9 + delimiter: None,
10 + name_key: "name",
11 + arguments_key: "arguments",
12 +};
13 +
14 +/// Tool parser for Qwen XML-wrapped JSON tool calls.
15 +///
16 +/// Example tool call content:
17 +///
18 +/// ```text
19 +/// <tool_call>
20 +/// {"name": "get_weather", "arguments": {"location":"Tokyo"}}
21 +/// </tool_call>
22 +/// ```
23 +///
24 +/// Arguments are already OpenAI-style JSON text, so they are streamed as raw
25 +/// argument deltas without schema conversion or JSON normalization.
26 +///
27 +/// Note: parallel calls are represented as repeated
28 +/// `<tool_call>...</tool_call>` blocks, not as multiple calls inside one tag.
29 +pub struct Qwen3XmlToolParser {
30 + inner: JsonToolCallParser,
31 +}
32 +
33 +impl Qwen3XmlToolParser {
34 + /// Create a Qwen XML tool parser.
35 + fn new(_tools: &[Tool]) -> Self {
36 + Self {
37 + inner: JsonToolCallParser::new(QWEN_XML_CONFIG),
38 + }
39 + }
40 +}
41 +
42 +impl ToolParser for Qwen3XmlToolParser {
43 + /// Create a boxed Qwen XML tool parser.
44 + fn create(tools: &[Tool]) -> Result<Box<dyn ToolParser>>
45 + where
46 + Self: Sized + 'static,
47 + {
48 + Ok(Box::new(Self::new(tools)))
49 + }
50 +
51 + /// Push one decoded text chunk through the Qwen XML parser.
52 + fn push(&mut self, chunk: &str) -> Result<ToolParseResult> {
53 + self.inner.push(chunk)
54 + }
55 +
56 + /// Flush buffered text and reset parser state.
57 + fn finish(&mut self) -> Result<ToolParseResult> {
58 + self.inner.finish()
59 + }
60 +}
61 +
62 +#[cfg(test)]
63 +mod tests {
64 + use expect_test::expect;
65 + use thiserror_ext::AsReport;
66 +
67 + use super::Qwen3XmlToolParser;
68 + use crate::test_utils::{collect_stream, split_by_chars, test_tools};
69 + use crate::{ToolParseResult, ToolParser};
70 +
71 + fn build_tool_call(function_name: &str, arguments: &str) -> String {
72 + format!(
73 + "<tool_call>\n{{\"name\": \"{function_name}\", \"arguments\": {arguments}}}\n</tool_call>"
74 + )
75 + }
76 +
77 + #[test]
78 + fn qwen_xml_parse_complete_without_tool_call_keeps_text() {
79 + let mut parser = Qwen3XmlToolParser::new(&test_tools());
80 + let result = parser.parse_complete("Hello, world!").unwrap();
81 +
82 + assert_eq!(result.normal_text, "Hello, world!");
83 + assert!(result.calls.is_empty());
84 + }
85 +
86 + #[test]
87 + fn qwen_xml_parse_complete_extracts_raw_json_arguments() {
88 + let mut parser = Qwen3XmlToolParser::new(&test_tools());
89 + let arguments = r#"{ "location": "Tokyo", "days": "3" }"#;
90 + let result = parser
91 + .parse_complete(&format!(
92 + "Let me check.\n{}",
93 + build_tool_call("get_weather", arguments)
94 + ))
95 + .unwrap();
96 +
97 + assert_eq!(result.normal_text, "Let me check.\n");
98 + assert_eq!(result.calls.len(), 1);
99 + assert_eq!(result.calls[0].tool_index, 0);
100 + assert_eq!(result.calls[0].name.as_deref(), Some("get_weather"));
101 + assert_eq!(result.calls[0].arguments, arguments);
102 + }
103 +
104 + #[test]
105 + fn qwen_xml_does_not_validate_or_normalize_arguments() {
106 + let mut parser = Qwen3XmlToolParser::new(&test_tools());
107 + let arguments = r#"{"location":"Tokyo",}"#;
108 + let result = parser.parse_complete(&build_tool_call("get_weather", arguments)).unwrap();
109 +
110 + assert_eq!(result.calls[0].arguments, arguments);
111 + }
112 +
113 + #[test]
114 + fn qwen_xml_streaming_emits_argument_deltas() {
115 + let mut parser = Qwen3XmlToolParser::new(&test_tools());
116 + let chunks = [
117 + "<tool_call>",
118 + "\n{\"name\": \"get_weather\", \"arguments\": ",
119 + "{\"location\":",
120 + "\"Beijing\"",
121 + "}",
122 + "}\n</tool_call>",
123 + ];
124 +
125 + let mut result = ToolParseResult::default();
126 + let mut observed_arguments = Vec::new();
127 + for chunk in chunks {
128 + let next = parser.push(chunk).unwrap();
129 + observed_arguments.extend(
130 + next.calls
131 + .iter()
132 + .filter(|call| call.name.is_none())
133 + .map(|call| call.arguments.clone()),
134 + );
135 + result.append(next);
136 + }
137 + result.append(parser.finish().unwrap());
138 +
139 + assert_eq!(observed_arguments, ["{\"location\":", "\"Beijing\"", "}"]);
140 + assert_eq!(
141 + result.coalesce_calls().calls[0].arguments,
142 + r#"{"location":"Beijing"}"#
143 + );
144 + }
145 +
146 + #[test]
147 + fn qwen_xml_streaming_handles_split_markers() {
148 + let input = format!(
149 + "hello {}",
150 + build_tool_call("get_weather", r#"{"location":"Tokyo"}"#)
151 + );
152 + let chunks = split_by_chars(&input, 5);
153 + let mut parser = Qwen3XmlToolParser::new(&test_tools());
154 +
155 + let result = collect_stream(&mut parser, &chunks);
156 +
157 + assert_eq!(result.normal_text, "hello ");
158 + assert_eq!(result.calls.len(), 1);
159 + assert_eq!(result.calls[0].arguments, r#"{"location":"Tokyo"}"#);
160 + }
161 +
162 + #[test]
163 + fn qwen_xml_keeps_end_marker_literal_inside_json_string() {
164 + let mut parser = Qwen3XmlToolParser::new(&test_tools());
165 + let arguments = r#"{"text":"literal </tool_call> inside"}"#;
166 + let result = parser.parse_complete(&build_tool_call("echo", arguments)).unwrap();
167 +
168 + assert_eq!(result.calls.len(), 1);
169 + assert_eq!(result.calls[0].arguments, arguments);
170 + }
171 +
172 + #[test]
173 + fn qwen_xml_decodes_escaped_function_name() {
174 + let mut parser = Qwen3XmlToolParser::new(&test_tools());
175 + let result = parser
176 + .parse_complete(
177 + r#"<tool_call>
178 +{"name":"say_\"hi","arguments":{}}
179 +</tool_call>"#,
180 + )
181 + .unwrap();
182 +
183 + assert_eq!(result.calls[0].name.as_deref(), Some("say_\"hi"));
184 + }
185 +
186 + #[test]
187 + fn qwen_xml_requires_newline_after_tool_call_start() {
188 + let mut parser = Qwen3XmlToolParser::new(&test_tools());
189 + let input = r#"<tool_call>{"name":"get_weather","arguments":{}}
190 +</tool_call>"#;
191 +
192 + let result = parser.parse_complete(input).unwrap();
193 +
194 + assert_eq!(result.normal_text, input);
195 + assert!(result.calls.is_empty());
196 + }
197 +
198 + #[test]
199 + fn qwen_xml_requires_newline_before_tool_call_end() {
200 + let mut parser = Qwen3XmlToolParser::new(&test_tools());
201 + let error = parser
202 + .parse_complete(
203 + r#"<tool_call>
204 +{"name":"get_weather","arguments":{}}</tool_call>"#,
205 + )
206 + .unwrap_err();
207 +
208 + assert!(error.to_report_string().starts_with("tool parser parsing failed:"));
209 + }
210 +
211 + #[test]
212 + fn qwen_xml_streaming_extracts_multiple_tool_calls() {
213 + let input = format!(
214 + "{}{}",
215 + build_tool_call("get_weather", r#"{"location":"Shanghai"}"#),
216 + build_tool_call("add", r#"{"x":1,"y":2}"#),
217 + );
218 + let chunks = split_by_chars(&input, 7);
219 + let mut parser = Qwen3XmlToolParser::new(&test_tools());
220 +
221 + let result = collect_stream(&mut parser, &chunks);
222 +
223 + expect![[r#"
224 + ToolParseResult {
225 + normal_text: "",
226 + calls: [
227 + ToolCallDelta {
228 + tool_index: 0,
229 + name: Some(
230 + "get_weather",
231 + ),
232 + arguments: "{\"location\":\"Shanghai\"}",
233 + },
234 + ToolCallDelta {
235 + tool_index: 1,
236 + name: Some(
237 + "add",
238 + ),
239 + arguments: "{\"x\":1,\"y\":2}",
240 + },
241 + ],
242 + }
243 + "#]]
244 + .assert_debug_eq(&result);
245 + }
246 +
247 + #[test]
248 + fn qwen_xml_finish_fails_incomplete_tool_call() {
249 + let mut parser = Qwen3XmlToolParser::new(&test_tools());
250 + parser
251 + .push(
252 + r#"<tool_call>
253 +{"name":"get_weather","arguments":{"location""#,
254 + )
255 + .unwrap();
256 +
257 + let error = parser.finish().unwrap_err();
258 +
259 + expect!["tool parser parsing failed: incomplete Qwen XML tool call"]
260 + .assert_eq(&error.to_report_string());
261 + }
262 +
263 + #[test]
264 + fn qwen_xml_malformed_field_order_fails_fast() {
265 + let mut parser = Qwen3XmlToolParser::new(&test_tools());
266 + let error = parser
267 + .push(
268 + r#"<tool_call>
269 +{"arguments":{},"name":"get_weather"}
270 +</tool_call>"#,
271 + )
272 + .unwrap_err();
273 +
274 + expect![[r#"
275 + tool parser parsing failed: invalid Qwen XML
276 + expected `name`"#]]
277 + .assert_eq(&error.to_report_string());
278 + }
279 +}
rust/src/tool-parser/src/kimi_k2.rs new
+560
@@ -0,0 +1,560 @@
1 +use winnow::ascii::{digit1, multispace0 as ws0};
2 +use winnow::combinator::{alt, eof, repeat, seq};
3 +use winnow::prelude::*;
4 +use winnow::stream::Partial;
5 +use winnow::token::{literal, rest, take_until, take_while};
6 +
7 +use super::utils::{JsonObjectScanState, parse_buffered_event, safe_text_len, take_json_object};
8 +use super::{Result, ToolCallDelta, ToolParseResult, ToolParser};
9 +use crate::Tool;
10 +
11 +const TOOL_CALLS_START: &str = "<|tool_calls_section_begin|>";
12 +const TOOL_CALLS_END: &str = "<|tool_calls_section_end|>";
13 +const TOOL_CALL_START: &str = "<|tool_call_begin|>";
14 +const TOOL_CALL_END: &str = "<|tool_call_end|>";
15 +const TOOL_CALL_ARGUMENT_START: &str = "<|tool_call_argument_begin|>";
16 +
17 +type KimiK2Input<'i> = Partial<&'i str>;
18 +
19 +#[derive(Debug, Clone, PartialEq, Eq)]
20 +enum KimiK2Mode {
21 + Text,
22 + ToolBlock,
23 + Header,
24 + Arguments { json_scan: JsonObjectScanState },
25 + Done,
26 +}
27 +
28 +#[derive(Debug, Clone, PartialEq, Eq)]
29 +enum KimiK2Event {
30 + Text {
31 + len: usize,
32 + },
33 + ToolCallsStart,
34 + ToolCallStart,
35 + ToolCallHeader {
36 + function_name: String,
37 + function_index: usize,
38 + },
39 + Arguments {
40 + len: usize,
41 + },
42 + ToolCallEnd,
43 + ToolCallsEnd,
44 + IgnoredRest,
45 +}
46 +
47 +/// Tool parser for Kimi K2 token-delimited tool calls.
48 +///
49 +/// Example tool call content:
50 +///
51 +/// ```text
52 +/// <|tool_calls_section_begin|>
53 +/// <|tool_call_begin|>functions.get_weather:0<|tool_call_argument_begin|>{"location":"NYC"}<|tool_call_end|>
54 +/// <|tool_calls_section_end|>
55 +/// ```
56 +///
57 +/// Arguments are already OpenAI-style JSON text, so they are streamed as raw
58 +/// argument deltas without schema conversion or JSON normalization.
59 +pub struct KimiK2ToolParser {
60 + buffer: String,
61 + mode: KimiK2Mode,
62 + active_tool_index: Option<usize>,
63 +}
64 +
65 +impl KimiK2ToolParser {
66 + /// Create a Kimi K2 tool parser.
67 + fn new(_tools: &[Tool]) -> Self {
68 + Self {
69 + buffer: String::new(),
70 + mode: KimiK2Mode::Text,
71 + active_tool_index: None,
72 + }
73 + }
74 +
75 + /// Apply one parsed Kimi K2 event to parser state and output.
76 + fn apply_event(&mut self, event: KimiK2Event, result: &mut ToolParseResult) -> Result<()> {
77 + match event {
78 + KimiK2Event::Text { len: consumed_len } => {
79 + result.normal_text.push_str(&self.buffer[..consumed_len]);
80 + }
81 + KimiK2Event::ToolCallsStart => self.mode = KimiK2Mode::ToolBlock,
82 + KimiK2Event::ToolCallStart => self.mode = KimiK2Mode::Header,
83 + KimiK2Event::ToolCallHeader {
84 + function_name,
85 + function_index,
86 + } => {
87 + let tool_index = function_index;
88 + self.active_tool_index = Some(tool_index);
89 + self.mode = KimiK2Mode::Arguments {
90 + json_scan: JsonObjectScanState::default(),
91 + };
92 + result.calls.push(ToolCallDelta {
93 + tool_index,
94 + name: Some(function_name),
95 + arguments: String::new(),
96 + });
97 + }
98 + KimiK2Event::Arguments { len: consumed_len } => {
99 + let Some(tool_index) = self.active_tool_index else {
100 + return Err(parsing_failed!(
101 + "Kimi K2 arguments without an active tool call"
102 + ));
103 + };
104 + result.calls.push(ToolCallDelta {
105 + tool_index,
106 + name: None,
107 + arguments: self.buffer[..consumed_len].to_string(),
108 + });
109 + }
110 + KimiK2Event::ToolCallEnd => {
111 + self.active_tool_index = None;
112 + self.mode = KimiK2Mode::ToolBlock;
113 + }
114 + KimiK2Event::ToolCallsEnd => {
115 + self.active_tool_index = None;
116 + self.mode = KimiK2Mode::Done;
117 + }
118 + KimiK2Event::IgnoredRest => {}
119 + }
120 + Ok(())
121 + }
122 +
123 + /// Reset all streaming state.
124 + fn reset(&mut self) {
125 + self.buffer.clear();
126 + self.mode = KimiK2Mode::Text;
127 + self.active_tool_index = None;
128 + }
129 +}
130 +
131 +impl ToolParser for KimiK2ToolParser {
132 + /// Create a boxed Kimi K2 tool parser.
133 + fn create(tools: &[Tool]) -> Result<Box<dyn ToolParser>>
134 + where
135 + Self: Sized + 'static,
136 + {
137 + Ok(Box::new(Self::new(tools)))
138 + }
139 +
140 + /// Preserve Kimi K2 special-token markers while decoding.
141 + fn preserve_special_tokens(&self) -> bool {
142 + true
143 + }
144 +
145 + /// Push one decoded text chunk through the Kimi K2 parser.
146 + fn push(&mut self, chunk: &str) -> Result<ToolParseResult> {
147 + self.buffer.push_str(chunk);
148 + let mut result = ToolParseResult::default();
149 +
150 + while let Some((event, consumed_len)) = parse_buffered_event(&self.buffer, |input| {
151 + parse_next_kimi_k2_event(input, &mut self.mode)
152 + })? {
153 + self.apply_event(event, &mut result)?;
154 + self.buffer.drain(..consumed_len);
155 + }
156 +
157 + Ok(result)
158 + }
159 +
160 + /// Flush buffered text and reset parser state.
161 + fn finish(&mut self) -> Result<ToolParseResult> {
162 + let mut result = ToolParseResult::default();
163 + match &self.mode {
164 + KimiK2Mode::Text => result.normal_text.push_str(&self.buffer),
165 + KimiK2Mode::ToolBlock | KimiK2Mode::Done => {}
166 + KimiK2Mode::Header | KimiK2Mode::Arguments { .. } => {
167 + return Err(parsing_failed!("incomplete Kimi K2 tool call"));
168 + }
169 + }
170 + self.reset();
171 + Ok(result)
172 + }
173 +}
174 +
175 +/// Parse a Kimi K2 event for the current parser mode.
176 +fn parse_next_kimi_k2_event(
177 + input: &mut KimiK2Input<'_>,
178 + mode: &mut KimiK2Mode,
179 +) -> ModalResult<KimiK2Event> {
180 + match mode {
181 + KimiK2Mode::Text => parse_text_event(input),
182 + KimiK2Mode::ToolBlock => parse_tool_block_event(input),
183 + KimiK2Mode::Header => tool_call_header_event(input),
184 + KimiK2Mode::Arguments { json_scan } => parse_arguments_event(input, json_scan),
185 + KimiK2Mode::Done => ignored_rest_event(input),
186 + }
187 +}
188 +
189 +/// Parse a text-mode Kimi K2 event.
190 +fn parse_text_event(input: &mut KimiK2Input<'_>) -> ModalResult<KimiK2Event> {
191 + alt((tool_calls_start_event, safe_text_event)).parse_next(input)
192 +}
193 +
194 +/// Parse one event inside the Kimi K2 tool-calls section.
195 +fn parse_tool_block_event(input: &mut KimiK2Input<'_>) -> ModalResult<KimiK2Event> {
196 + alt((tool_calls_end_event, tool_call_start_event)).parse_next(input)
197 +}
198 +
199 +/// Parse one event inside a Kimi K2 tool-call arguments payload.
200 +fn parse_arguments_event(
201 + input: &mut KimiK2Input<'_>,
202 + json_scan: &mut JsonObjectScanState,
203 +) -> ModalResult<KimiK2Event> {
204 + if json_scan.complete() {
205 + tool_call_end_event(input)
206 + } else {
207 + argument_delta_event(input, json_scan)
208 + }
209 +}
210 +
211 +/// Parse a Kimi K2 tool-calls section start marker.
212 +fn tool_calls_start_event(input: &mut KimiK2Input<'_>) -> ModalResult<KimiK2Event> {
213 + literal(TOOL_CALLS_START).value(KimiK2Event::ToolCallsStart).parse_next(input)
214 +}
215 +
216 +/// Parse a Kimi K2 tool-calls section end marker.
217 +fn tool_calls_end_event(input: &mut KimiK2Input<'_>) -> ModalResult<KimiK2Event> {
218 + (ws0, literal(TOOL_CALLS_END))
219 + .value(KimiK2Event::ToolCallsEnd)
220 + .parse_next(input)
221 +}
222 +
223 +/// Parse a Kimi K2 tool-call start marker.
224 +fn tool_call_start_event(input: &mut KimiK2Input<'_>) -> ModalResult<KimiK2Event> {
225 + (ws0, literal(TOOL_CALL_START))
226 + .value(KimiK2Event::ToolCallStart)
227 + .parse_next(input)
228 +}
229 +
230 +/// Parse a Kimi K2 tool-call end marker.
231 +fn tool_call_end_event(input: &mut KimiK2Input<'_>) -> ModalResult<KimiK2Event> {
232 + literal(TOOL_CALL_END).value(KimiK2Event::ToolCallEnd).parse_next(input)
233 +}
234 +
235 +/// Parse a Kimi K2 tool-call header before the argument marker.
236 +fn tool_call_header_event(input: &mut KimiK2Input<'_>) -> ModalResult<KimiK2Event> {
237 + let (header, _) = (
238 + take_until(1.., TOOL_CALL_ARGUMENT_START),
239 + literal(TOOL_CALL_ARGUMENT_START),
240 + )
241 + .parse_next(input)?;
242 +
243 + let mut header_input = header;
244 + let (header, _, _) = (tool_header, ws0, eof).parse_next(&mut header_input)?;
245 +
246 + Ok(KimiK2Event::ToolCallHeader {
247 + function_name: header.function_name,
248 + function_index: header.function_index,
249 + })
250 +}
251 +
252 +/// Parse a Kimi K2 raw JSON arguments delta.
253 +fn argument_delta_event(
254 + input: &mut KimiK2Input<'_>,
255 + json_scan: &mut JsonObjectScanState,
256 +) -> ModalResult<KimiK2Event> {
257 + take_json_object(input, json_scan).map(|len| KimiK2Event::Arguments { len })
258 +}
259 +
260 +/// Parse a safe text run before the next Kimi K2 tool-calls section.
261 +fn safe_text_event(input: &mut KimiK2Input<'_>) -> ModalResult<KimiK2Event> {
262 + safe_text_len(input, TOOL_CALLS_START).map(|len| KimiK2Event::Text { len })
263 +}
264 +
265 +/// Parse ignored rest after the Kimi K2 tool-calls section ends.
266 +fn ignored_rest_event(input: &mut KimiK2Input<'_>) -> ModalResult<KimiK2Event> {
267 + rest.value(KimiK2Event::IgnoredRest).parse_next(input)
268 +}
269 +
270 +#[derive(Debug, Clone, PartialEq, Eq)]
271 +struct KimiK2ToolHeader {
272 + function_name: String,
273 + function_index: usize,
274 +}
275 +
276 +/// Parse a Kimi K2 tool-call header.
277 +fn tool_header(input: &mut &str) -> ModalResult<KimiK2ToolHeader> {
278 + let (function_name, function_index) = seq!(
279 + _: ws0,
280 + _: namespace_prefix,
281 + tool_name_segment,
282 + _: literal(":"),
283 + tool_call_index,
284 + )
285 + .parse_next(input)?;
286 +
287 + Ok(KimiK2ToolHeader {
288 + function_name: function_name.to_string(),
289 + function_index,
290 + })
291 +}
292 +
293 +/// Parse Kimi K2 namespace segments before the final tool name.
294 +fn namespace_prefix(input: &mut &str) -> ModalResult<()> {
295 + repeat(0.., namespace_segment).parse_next(input)
296 +}
297 +
298 +/// Parse a Kimi K2 namespace segment.
299 +fn namespace_segment<'i>(input: &mut &'i str) -> ModalResult<&'i str> {
300 + let (segment, _) = (tool_name_segment, literal(".")).parse_next(input)?;
301 + Ok(segment)
302 +}
303 +
304 +/// Parse a Kimi K2 tool name segment.
305 +fn tool_name_segment<'i>(input: &mut &'i str) -> ModalResult<&'i str> {
306 + take_while(1.., |ch: char| {
307 + !ch.is_whitespace() && ch != '<' && ch != ':' && ch != '.'
308 + })
309 + .parse_next(input)
310 +}
311 +
312 +/// Parse a Kimi K2 tool-call index.
313 +fn tool_call_index(input: &mut &str) -> ModalResult<usize> {
314 + digit1.parse_to().parse_next(input)
315 +}
316 +
317 +#[cfg(test)]
318 +mod tests {
319 + use expect_test::expect;
320 + use thiserror_ext::AsReport;
321 +
322 + use super::{
323 + KimiK2ToolParser, TOOL_CALL_ARGUMENT_START, TOOL_CALL_END, TOOL_CALL_START, TOOL_CALLS_END,
324 + TOOL_CALLS_START, ToolParser, tool_header,
325 + };
326 + use crate::ToolParseResult;
327 + use crate::test_utils::{collect_stream, split_by_chars, test_tools};
328 +
329 + fn build_tool_call(function_name: &str, index: usize, arguments: &str) -> String {
330 + format!(
331 + "{TOOL_CALL_START}functions.{function_name}:{index}{TOOL_CALL_ARGUMENT_START}{arguments}{TOOL_CALL_END}"
332 + )
333 + }
334 +
335 + fn build_tool_section(tool_calls: &[String]) -> String {
336 + format!("{TOOL_CALLS_START}{}{TOOL_CALLS_END}", tool_calls.join(""))
337 + }
338 +
339 + #[test]
340 + fn kimi_k2_parse_complete_without_tool_call_keeps_text() {
341 + let mut parser = KimiK2ToolParser::new(&test_tools());
342 + let result = parser.parse_complete("Hello, world!").unwrap();
343 +
344 + assert_eq!(result.normal_text, "Hello, world!");
345 + assert!(result.calls.is_empty());
346 + }
347 +
348 + #[test]
349 + fn kimi_k2_parse_complete_extracts_raw_json_arguments() {
350 + let mut parser = KimiK2ToolParser::new(&test_tools());
351 + let arguments = r#"{ "location": "NYC", "days": "3" }"#;
352 + let result = parser
353 + .parse_complete(&format!(
354 + "Checking. {} trailing text",
355 + build_tool_section(&[build_tool_call("get_weather", 0, arguments)])
356 + ))
357 + .unwrap();
358 +
359 + assert_eq!(result.normal_text, "Checking. ");
360 + assert_eq!(result.calls.len(), 1);
361 + assert_eq!(result.calls[0].tool_index, 0);
362 + assert_eq!(result.calls[0].name.as_deref(), Some("get_weather"));
363 + assert_eq!(result.calls[0].arguments, arguments);
364 + }
365 +
366 + #[test]
367 + fn kimi_k2_does_not_validate_or_normalize_arguments() {
368 + let mut parser = KimiK2ToolParser::new(&test_tools());
369 + let arguments = r#"{"location":"NYC",}"#;
370 + let result = parser
371 + .parse_complete(&build_tool_section(&[build_tool_call(
372 + "get_weather",
373 + 0,
374 + arguments,
375 + )]))
376 + .unwrap();
377 +
378 + assert_eq!(result.calls[0].arguments, arguments);
379 + }
380 +
381 + #[test]
382 + fn kimi_k2_streaming_emits_argument_deltas() {
383 + let mut parser = KimiK2ToolParser::new(&test_tools());
384 + let chunks = [
385 + TOOL_CALLS_START,
386 + TOOL_CALL_START,
387 + "functions.get_weather:0",
388 + TOOL_CALL_ARGUMENT_START,
389 + "{\"location\":",
390 + "\"Paris\"",
391 + "}",
392 + TOOL_CALL_END,
393 + TOOL_CALLS_END,
394 + ];
395 +
396 + let mut result = ToolParseResult::default();
397 + let mut observed_arguments = Vec::new();
398 + for chunk in chunks {
399 + let next = parser.push(chunk).unwrap();
400 + observed_arguments.extend(
401 + next.calls
402 + .iter()
403 + .filter(|call| call.name.is_none())
404 + .map(|call| call.arguments.clone()),
405 + );
406 + result.append(next);
407 + }
408 + result.append(parser.finish().unwrap());
409 +
410 + assert_eq!(observed_arguments, ["{\"location\":", "\"Paris\"", "}"]);
411 + let result = result.coalesce_calls();
412 + assert_eq!(result.calls[0].arguments, r#"{"location":"Paris"}"#);
413 + }
414 +
415 + #[test]
416 + fn kimi_k2_streaming_holds_back_split_markers() {
417 + let mut parser = KimiK2ToolParser::new(&test_tools());
418 + let chunks = [
419 + "hello <|tool_calls",
420 + "_section_begin|>",
421 + TOOL_CALL_START,
422 + "functions.get_weather:0",
423 + TOOL_CALL_ARGUMENT_START,
424 + r#"{"location":"NYC"}"#,
425 + "<|tool_call",
426 + "_end|>",
427 + TOOL_CALLS_END,
428 + ];
429 +
430 + let result = collect_stream(&mut parser, &chunks);
431 +
432 + assert_eq!(result.normal_text, "hello ");
433 + assert_eq!(result.calls.len(), 1);
434 + assert_eq!(result.calls[0].arguments, r#"{"location":"NYC"}"#);
435 + }
436 +
437 + #[test]
438 + fn kimi_k2_keeps_end_marker_literal_inside_json_string() {
439 + let mut parser = KimiK2ToolParser::new(&test_tools());
440 + let arguments = format!(r#"{{"text":"literal {TOOL_CALL_END} inside"}}"#);
441 + let input = build_tool_section(&[build_tool_call("echo", 0, &arguments)]);
442 +
443 + let result = parser.parse_complete(&input).unwrap();
444 +
445 + assert_eq!(result.calls.len(), 1);
446 + assert_eq!(result.calls[0].arguments, arguments);
447 + }
448 +
449 + #[test]
450 + fn kimi_k2_streaming_keeps_split_end_marker_literal_inside_json_string() {
451 + let mut parser = KimiK2ToolParser::new(&test_tools());
452 + let chunks = [
453 + TOOL_CALLS_START,
454 + TOOL_CALL_START,
455 + "functions.echo:0",
456 + TOOL_CALL_ARGUMENT_START,
457 + r#"{"text":"literal <|tool"#,
458 + r#"_call_end|> inside"}"#,
459 + TOOL_CALL_END,
460 + TOOL_CALLS_END,
461 + ];
462 +
463 + let result = collect_stream(&mut parser, &chunks);
464 +
465 + assert_eq!(result.calls.len(), 1);
466 + assert_eq!(
467 + result.calls[0].arguments,
468 + r#"{"text":"literal <|tool_call_end|> inside"}"#
469 + );
470 + }
471 +
472 + #[test]
473 + fn kimi_k2_streaming_extracts_multiple_tool_calls() {
474 + let mut parser = KimiK2ToolParser::new(&test_tools());
475 + let input = build_tool_section(&[
476 + build_tool_call("get_weather", 0, r#"{"location":"Shanghai"}"#),
477 + build_tool_call("add", 1, r#"{"x":1,"y":2}"#),
478 + ]);
479 +
480 + let chunks = split_by_chars(&input, 7);
481 + let result = collect_stream(&mut parser, &chunks);
482 +
483 + expect![[r#"
484 + ToolParseResult {
485 + normal_text: "",
486 + calls: [
487 + ToolCallDelta {
488 + tool_index: 0,
489 + name: Some(
490 + "get_weather",
491 + ),
492 + arguments: "{\"location\":\"Shanghai\"}",
493 + },
494 + ToolCallDelta {
495 + tool_index: 1,
496 + name: Some(
497 + "add",
498 + ),
499 + arguments: "{\"x\":1,\"y\":2}",
500 + },
501 + ],
502 + }
503 + "#]]
504 + .assert_debug_eq(&result);
505 + }
506 +
507 + #[test]
508 + fn kimi_k2_accepts_non_functions_header_prefix() {
509 + let mut parser = KimiK2ToolParser::new(&test_tools());
510 + let input = format!(
511 + "{TOOL_CALLS_START}{TOOL_CALL_START}api.tools.search:42{TOOL_CALL_ARGUMENT_START}{{}}{TOOL_CALL_END}{TOOL_CALLS_END}"
512 + );
513 +
514 + let result = parser.parse_complete(&input).unwrap();
515 +
516 + assert_eq!(result.calls[0].tool_index, 42);
517 + assert_eq!(result.calls[0].name.as_deref(), Some("search"));
518 + assert_eq!(result.calls[0].arguments, "{}");
519 + }
520 +
521 + #[test]
522 + fn kimi_k2_tool_header_parses_namespace_function_and_index() {
523 + let mut input = "api.tools.search:42";
524 + let header = tool_header(&mut input).unwrap();
525 +
526 + expect![[r#"
527 + KimiK2ToolHeader {
528 + function_name: "search",
529 + function_index: 42,
530 + }
531 + "#]]
532 + .assert_debug_eq(&header);
533 + }
534 +
535 + #[test]
536 + fn kimi_k2_finish_fails_incomplete_tool_call() {
537 + let mut parser = KimiK2ToolParser::new(&test_tools());
538 + parser
539 + .push(&format!(
540 + "{TOOL_CALLS_START}{TOOL_CALL_START}functions.get_weather:0{TOOL_CALL_ARGUMENT_START}{{\"location\""
541 + ))
542 + .unwrap();
543 +
544 + let error = parser.finish().unwrap_err();
545 +
546 + expect!["tool parser parsing failed: incomplete Kimi K2 tool call"]
547 + .assert_eq(&error.to_report_string());
548 + }
549 +
550 + #[test]
551 + fn kimi_k2_malformed_header_fails_fast() {
552 + let mut parser = KimiK2ToolParser::new(&test_tools());
553 + let input =
554 + format!("{TOOL_CALLS_START}{TOOL_CALL_START}get_weather{TOOL_CALL_ARGUMENT_START}{{}}");
555 +
556 + let error = parser.push(&input).unwrap_err();
557 +
558 + expect!["tool parser parsing failed: "].assert_eq(&error.to_report_string());
559 + }
560 +}
rust/src/tool-parser/src/lib.rs new
+141
@@ -0,0 +1,141 @@
1 +//! Streaming tool parsers for chat completions.
2 +
3 +#[macro_use]
4 +mod error;
5 +mod deepseek_dsml;
6 +mod deepseek_json;
7 +mod gemma4;
8 +mod glm_xml;
9 +mod json;
10 +mod kimi_k2;
11 +mod minimax_m2;
12 +mod parameters;
13 +mod qwen_coder;
14 +#[cfg(any(test, feature = "test-util"))]
15 +pub mod test_utils;
16 +mod utils;
17 +
18 +use std::collections::{BTreeMap, btree_map};
19 +
20 +pub use deepseek_dsml::{DeepSeekV4ToolParser, DeepSeekV32ToolParser};
21 +pub use deepseek_json::{DeepSeekV3ToolParser, DeepSeekV31ToolParser};
22 +pub use error::{Result, ToolParserError};
23 +pub use gemma4::Gemma4ToolParser;
24 +pub use glm_xml::{Glm45MoeToolParser, Glm47MoeToolParser};
25 +pub use json::{HermesToolParser, Llama3JsonToolParser, MistralToolParser, Qwen3XmlToolParser};
26 +pub use kimi_k2::KimiK2ToolParser;
27 +pub use minimax_m2::MinimaxM2ToolParser;
28 +pub use qwen_coder::Qwen3CoderToolParser;
29 +use serde::{Deserialize, Serialize};
30 +use serde_json::Value;
31 +
32 +/// One function-style tool made available to the model.
33 +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
34 +pub struct Tool {
35 + pub name: String,
36 + pub description: Option<String>,
37 + pub parameters: Value,
38 + pub strict: Option<bool>,
39 +}
40 +
41 +/// One tool-call update emitted while parsing assistant text.
42 +#[derive(Debug, Clone, PartialEq, Eq)]
43 +pub struct ToolCallDelta {
44 + /// Stable parser-local tool index for this call within one assistant turn.
45 + pub tool_index: usize,
46 + /// Function name, present on the first update for one tool call.
47 + pub name: Option<String>,
48 + /// Arguments text contributed by this update.
49 + pub arguments: String,
50 +}
51 +
52 +/// Result of advancing tool parsing with one assistant-text input.
53 +#[derive(Debug, Clone, Default, PartialEq, Eq)]
54 +pub struct ToolParseResult {
55 + /// Plain assistant text that is not part of any tool call.
56 + pub normal_text: String,
57 + /// Tool-call updates extracted from this input.
58 + pub calls: Vec<ToolCallDelta>,
59 +}
60 +
61 +impl ToolParseResult {
62 + /// Append another parser result onto this one.
63 + ///
64 + /// Note that this does not attempt to merge multiple deltas for the same
65 + /// tool call into one complete item. Call `coalesce_calls()` after if
66 + /// that behavior is desired.
67 + pub(crate) fn append(&mut self, mut other: Self) {
68 + self.normal_text.push_str(&other.normal_text);
69 + self.calls.append(&mut other.calls);
70 + }
71 +
72 + /// Merge multiple deltas for the same tool call into one complete item.
73 + ///
74 + /// This is primarily used by the default `parse_complete()` implementation,
75 + /// which delegates through the incremental parser lifecycle and then
76 + /// needs to collapse streaming-style argument fragments into one final
77 + /// tool call.
78 + pub(crate) fn coalesce_calls(mut self) -> Self {
79 + let mut merged = BTreeMap::<usize, ToolCallDelta>::new();
80 + let mut order = Vec::new();
81 +
82 + for call in self.calls {
83 + match merged.entry(call.tool_index) {
84 + btree_map::Entry::Vacant(entry) => {
85 + order.push(call.tool_index);
86 + entry.insert(call);
87 + }
88 + btree_map::Entry::Occupied(mut entry) => {
89 + let existing = entry.get_mut();
90 + if existing.name.is_none() {
91 + existing.name = call.name;
92 + }
93 + existing.arguments.push_str(&call.arguments);
94 + }
95 + }
96 + }
97 +
98 + self.calls =
99 + order.into_iter().filter_map(|tool_index| merged.remove(&tool_index)).collect();
100 + self
101 + }
102 +}
103 +
104 +/// Incremental parser that extracts tool calls from assistant output.
105 +pub trait ToolParser: Send {
106 + /// Construct a boxed parser instance for one request stream.
107 + fn create(tools: &[Tool]) -> Result<Box<dyn ToolParser>>
108 + where
109 + Self: Sized + 'static;
110 +
111 + /// Return whether decoded output must preserve tokenizer special tokens.
112 + ///
113 + /// Some model families emit tool-call sentinels as special tokens. Those
114 + /// parsers need `skip_special_tokens = false` while parsing is enabled.
115 + fn preserve_special_tokens(&self) -> bool {
116 + false
117 + }
118 +
119 + /// Feed one decoded text delta into the parser.
120 + fn push(&mut self, chunk: &str) -> Result<ToolParseResult>;
121 +
122 + /// Flush any buffered partial state at end of stream.
123 + fn finish(&mut self) -> Result<ToolParseResult> {
124 + Ok(ToolParseResult::default())
125 + }
126 +
127 + /// Parse complete tool calls from final output.
128 + ///
129 + /// The default implementation reuses the incremental parser lifecycle by
130 + /// feeding the full output through `push()` and then calling `finish()`.
131 + /// This keeps one source of truth for robust parsers whose incremental
132 + /// state machine is equivalent across arbitrary chunking.
133 + fn parse_complete(&mut self, output: &str) -> Result<ToolParseResult> {
134 + let mut result = self.push(output)?;
135 + result.append(self.finish()?);
136 + Ok(result.coalesce_calls())
137 + }
138 +}
139 +
140 +#[cfg(test)]
141 +mod tests;
rust/src/tool-parser/src/minimax_m2.rs new
+519
@@ -0,0 +1,519 @@
1 +use winnow::ascii::{multispace0 as ws0, multispace1 as ws1};
2 +use winnow::combinator::{alt, delimited, repeat, seq, terminated};
3 +use winnow::prelude::*;
4 +use winnow::stream::Partial;
5 +use winnow::token::{literal, rest, take_until};
6 +
7 +use super::parameters::ToolSchemas;
8 +use super::utils::{parse_buffered_event, safe_text_len, xml_unescape};
9 +use super::{Result, ToolCallDelta, ToolParseResult, ToolParser};
10 +use crate::Tool;
11 +
12 +const TOOL_CALL_START: &str = "<minimax:tool_call>";
13 +const TOOL_CALL_END: &str = "</minimax:tool_call>";
14 +const INVOKE_START: &str = "<invoke";
15 +const INVOKE_END: &str = "</invoke>";
16 +const PARAMETER_START: &str = "<parameter";
17 +const PARAMETER_END: &str = "</parameter>";
18 +
19 +type MinimaxM2Input<'i> = Partial<&'i str>;
20 +
21 +#[derive(Debug, Clone, Copy, PartialEq, Eq)]
22 +enum MinimaxM2Mode {
23 + Text,
24 + ToolBlock,
25 + Done,
26 +}
27 +
28 +#[derive(Debug, Clone, PartialEq, Eq)]
29 +enum MinimaxM2Event {
30 + Text {
31 + len: usize,
32 + },
33 + ToolBlockStart,
34 + Invoke {
35 + name: String,
36 + raw_params: Vec<(String, String)>,
37 + },
38 + ToolBlockEnd,
39 + IgnoredRest,
40 +}
41 +
42 +/// Tool parser for MiniMax M2 XML-style tool calls.
43 +///
44 +/// Example tool call content:
45 +///
46 +/// ```text
47 +/// <minimax:tool_call><invoke name="get_weather">
48 +/// <parameter name="city">Seattle</parameter>
49 +/// </invoke></minimax:tool_call>
50 +/// ```
51 +///
52 +/// Arguments are emitted only after a full `<invoke>` block is parsed.
53 +pub struct MinimaxM2ToolParser {
54 + buffer: String,
55 + mode: MinimaxM2Mode,
56 + emitted_tool_count: usize,
57 + tool_parameters: ToolSchemas,
58 +}
59 +
60 +impl MinimaxM2ToolParser {
61 + /// Create a MiniMax M2 tool parser.
62 + fn new(tools: &[Tool]) -> Self {
63 + Self {
64 + buffer: String::new(),
65 + mode: MinimaxM2Mode::Text,
66 + emitted_tool_count: 0,
67 + tool_parameters: ToolSchemas::from_tools(tools),
68 + }
69 + }
70 +
71 + /// Apply one parsed MiniMax M2 event to parser state and output.
72 + fn apply_event(&mut self, event: MinimaxM2Event, result: &mut ToolParseResult) -> Result<()> {
73 + match event {
74 + MinimaxM2Event::Text { len: consumed_len } => {
75 + result.normal_text.push_str(&self.buffer[..consumed_len]);
76 + }
77 + MinimaxM2Event::ToolBlockStart => self.mode = MinimaxM2Mode::ToolBlock,
78 + MinimaxM2Event::Invoke { name, raw_params } => {
79 + let arguments = self.tool_parameters.convert_params_with_schema(&name, raw_params);
80 + let arguments = serde_json::to_string(&arguments)
81 + .map_err(|error| parsing_failed!("failed to serialize arguments: {}", error))?;
82 +
83 + result.calls.push(ToolCallDelta {
84 + tool_index: self.emitted_tool_count,
85 + name: Some(name),
86 + arguments,
87 + });
88 + self.emitted_tool_count += 1;
89 + }
90 + MinimaxM2Event::ToolBlockEnd => self.mode = MinimaxM2Mode::Done,
91 + MinimaxM2Event::IgnoredRest => {}
92 + }
93 + Ok(())
94 + }
95 +
96 + /// Reset all streaming state.
97 + fn reset(&mut self) {
98 + self.buffer.clear();
99 + self.mode = MinimaxM2Mode::Text;
100 + self.emitted_tool_count = 0;
101 + }
102 +}
103 +
104 +impl ToolParser for MinimaxM2ToolParser {
105 + /// Create a boxed MiniMax M2 tool parser.
106 + fn create(tools: &[Tool]) -> Result<Box<dyn ToolParser>>
107 + where
108 + Self: Sized + 'static,
109 + {
110 + Ok(Box::new(Self::new(tools)))
111 + }
112 +
113 + /// Push one decoded text chunk through the MiniMax M2 parser.
114 + fn push(&mut self, chunk: &str) -> Result<ToolParseResult> {
115 + self.buffer.push_str(chunk);
116 + let mut result = ToolParseResult::default();
117 +
118 + while let Some((event, consumed_len)) = parse_buffered_event(&self.buffer, |input| {
119 + parse_next_minimax_m2_event(input, self.mode)
120 + })? {
121 + self.apply_event(event, &mut result)?;
122 + self.buffer.drain(..consumed_len);
123 + }
124 +
125 + Ok(result)
126 + }
127 +
128 + /// Flush buffered text and reset parser state.
129 + fn finish(&mut self) -> Result<ToolParseResult> {
130 + let mut result = ToolParseResult::default();
131 + match self.mode {
132 + MinimaxM2Mode::Text => {
133 + result.normal_text.push_str(&self.buffer);
134 + }
135 + MinimaxM2Mode::ToolBlock => {
136 + return Err(parsing_failed!("incomplete MiniMax M2 tool call"));
137 + }
138 + MinimaxM2Mode::Done => {}
139 + }
140 + self.reset();
141 + Ok(result)
142 + }
143 +}
144 +
145 +/// Parse a MiniMax M2 event for the current parser mode.
146 +fn parse_next_minimax_m2_event(
147 + input: &mut MinimaxM2Input<'_>,
148 + mode: MinimaxM2Mode,
149 +) -> ModalResult<MinimaxM2Event> {
150 + match mode {
151 + MinimaxM2Mode::Text => parse_text_event(input),
152 + MinimaxM2Mode::ToolBlock => parse_tool_block_event(input),
153 + MinimaxM2Mode::Done => ignored_rest_event(input),
154 + }
155 +}
156 +
157 +/// Parse a text-mode MiniMax M2 event.
158 +fn parse_text_event(input: &mut MinimaxM2Input<'_>) -> ModalResult<MinimaxM2Event> {
159 + alt((tool_block_start_event, safe_text_event)).parse_next(input)
160 +}
161 +
162 +/// Parse a MiniMax M2 tool-block start marker.
163 +fn tool_block_start_event(input: &mut MinimaxM2Input<'_>) -> ModalResult<MinimaxM2Event> {
164 + literal(TOOL_CALL_START).value(MinimaxM2Event::ToolBlockStart).parse_next(input)
165 +}
166 +
167 +/// Parse a safe text run before the next MiniMax M2 marker.
168 +fn safe_text_event(input: &mut MinimaxM2Input<'_>) -> ModalResult<MinimaxM2Event> {
169 + safe_text_len(input, TOOL_CALL_START).map(|len| MinimaxM2Event::Text { len })
170 +}
171 +
172 +/// Parse one event inside a MiniMax M2 tool block.
173 +fn parse_tool_block_event(input: &mut MinimaxM2Input<'_>) -> ModalResult<MinimaxM2Event> {
174 + alt((tool_block_end_event, invoke_event)).parse_next(input)
175 +}
176 +
177 +/// Parse a MiniMax M2 tool-block end marker.
178 +fn tool_block_end_event(input: &mut MinimaxM2Input<'_>) -> ModalResult<MinimaxM2Event> {
179 + (ws0, literal(TOOL_CALL_END))
180 + .value(MinimaxM2Event::ToolBlockEnd)
181 + .parse_next(input)
182 +}
183 +
184 +/// Parse a complete MiniMax M2 invoke block.
185 +fn invoke_event(input: &mut MinimaxM2Input<'_>) -> ModalResult<MinimaxM2Event> {
186 + let (name, raw_params) = seq!(
187 + _: ws0,
188 + _: literal(INVOKE_START),
189 + _: (ws1, literal("name=")),
190 + attr_value,
191 + _: literal(">"),
192 + repeat(0.., terminated(parameter, ws0)),
193 + _: literal(INVOKE_END),
194 + )
195 + .parse_next(input)?;
196 +
197 + Ok(MinimaxM2Event::Invoke {
198 + name: name.trim().to_string(),
199 + raw_params,
200 + })
201 +}
202 +
203 +/// Parse a MiniMax M2 parameter block.
204 +fn parameter(input: &mut MinimaxM2Input<'_>) -> ModalResult<(String, String)> {
205 + let (name, value) = seq!(
206 + _: literal(PARAMETER_START),
207 + _: (ws1, literal("name=")),
208 + attr_value,
209 + _: literal(">"),
210 + take_until(0.., PARAMETER_END).map(xml_unescape),
211 + _: literal(PARAMETER_END),
212 + )
213 + .parse_next(input)?;
214 +
215 + Ok((name.trim().to_string(), value.into_owned()))
216 +}
217 +
218 +/// Parse a quoted or unquoted XML attribute value.
219 +fn attr_value<'i>(input: &mut MinimaxM2Input<'i>) -> ModalResult<&'i str> {
220 + alt((
221 + delimited(literal("\""), take_until(1.., "\""), literal("\"")),
222 + delimited(literal("'"), take_until(1.., "'"), literal("'")),
223 + take_until(1.., ">"),
224 + ))
225 + .parse_next(input)
226 +}
227 +
228 +/// Parse ignored rest after the MiniMax M2 tool block ends.
229 +fn ignored_rest_event(input: &mut MinimaxM2Input<'_>) -> ModalResult<MinimaxM2Event> {
230 + rest.value(MinimaxM2Event::IgnoredRest).parse_next(input)
231 +}
232 +
233 +#[cfg(test)]
234 +mod tests {
235 + use expect_test::expect;
236 + use serde_json::{Value, json};
237 + use thiserror_ext::AsReport;
238 +
239 + use super::{MinimaxM2ToolParser, TOOL_CALL_END, TOOL_CALL_START, ToolParser};
240 + use crate::test_utils::{collect_stream, split_by_chars, test_tools};
241 +
242 + fn build_tool_block(invokes: &[(&str, Vec<(&str, &str)>)]) -> String {
243 + let invokes = invokes
244 + .iter()
245 + .map(|(function_name, params)| {
246 + let params = params
247 + .iter()
248 + .map(|(name, value)| format!(r#"<parameter name="{name}">{value}</parameter>"#))
249 + .collect::<Vec<_>>()
250 + .join("");
251 + format!(r#"<invoke name="{function_name}">{params}</invoke>"#)
252 + })
253 + .collect::<String>();
254 + format!("{TOOL_CALL_START}{invokes}{TOOL_CALL_END}")
255 + }
256 +
257 + #[test]
258 + fn minimax_m2_parse_complete_without_tool_call_keeps_text() {
259 + let mut parser = MinimaxM2ToolParser::new(&test_tools());
260 + let result = parser.parse_complete("Hello, world!").unwrap();
261 +
262 + assert_eq!(result.normal_text, "Hello, world!");
263 + assert!(result.calls.is_empty());
264 + }
265 +
266 + #[test]
267 + fn minimax_m2_parse_complete_extracts_single_tool_call() {
268 + let mut parser = MinimaxM2ToolParser::new(&test_tools());
269 + let result = parser
270 + .parse_complete(&build_tool_block(&[(
271 + "get_weather",
272 + vec![("city", "Seattle"), ("days", "5")],
273 + )]))
274 + .unwrap();
275 +
276 + assert!(result.normal_text.is_empty());
277 + assert_eq!(result.calls.len(), 1);
278 + assert_eq!(result.calls[0].name.as_deref(), Some("get_weather"));
279 + assert_eq!(
280 + serde_json::from_str::<Value>(&result.calls[0].arguments).unwrap(),
281 + json!({ "city": "Seattle", "days": 5 })
282 + );
283 + }
284 +
285 + #[test]
286 + fn minimax_m2_parse_complete_preserves_prefix_and_ignores_trailing_text() {
287 + let mut parser = MinimaxM2ToolParser::new(&test_tools());
288 + let output = format!(
289 + "Let me check. {} This trailing text is ignored.",
290 + build_tool_block(&[("get_weather", vec![("city", "Seattle")])])
291 + );
292 + let result = parser.parse_complete(&output).unwrap();
293 +
294 + assert_eq!(result.normal_text, "Let me check. ");
295 + assert_eq!(result.calls.len(), 1);
296 + }
297 +
298 + #[test]
299 + fn minimax_m2_parse_complete_extracts_multiple_invokes() {
300 + let mut parser = MinimaxM2ToolParser::new(&test_tools());
301 + let result = parser
302 + .parse_complete(&build_tool_block(&[
303 + ("get_weather", vec![("city", "Seattle")]),
304 + ("get_weather", vec![("city", "NYC")]),
305 + ]))
306 + .unwrap();
307 +
308 + assert_eq!(result.calls.len(), 2);
309 + assert_eq!(result.calls[0].tool_index, 0);
310 + assert_eq!(result.calls[1].tool_index, 1);
311 + assert_eq!(
312 + serde_json::from_str::<Value>(&result.calls[0].arguments).unwrap(),
313 + json!({ "city": "Seattle" })
314 + );
315 + assert_eq!(
316 + serde_json::from_str::<Value>(&result.calls[1].arguments).unwrap(),
317 + json!({ "city": "NYC" })
318 + );
319 + }
320 +
321 + #[test]
322 + fn minimax_m2_parse_complete_converts_schema_types() {
323 + let mut parser = MinimaxM2ToolParser::new(&test_tools());
324 + let result = parser
325 + .parse_complete(&build_tool_block(&[(
326 + "convert",
327 + vec![
328 + ("whole", "5.0"),
329 + ("flag", "true"),
330 + ("payload", r#"{"nested":true}"#),
331 + ("items", "[1,2]"),
332 + ("empty", "42"),
333 + ],
334 + )]))
335 + .unwrap();
336 +
337 + assert_eq!(
338 + serde_json::from_str::<Value>(&result.calls[0].arguments).unwrap(),
339 + json!({
340 + "whole": 5.0,
341 + "flag": true,
342 + "payload": { "nested": true },
343 + "items": [1, 2],
344 + "empty": "42",
345 + })
346 + );
347 + }
348 +
349 + #[test]
350 + fn minimax_m2_parse_complete_unescapes_literal_closing_tags_in_parameter_value() {
351 + let mut parser = MinimaxM2ToolParser::new(&test_tools());
352 + let result = parser
353 + .parse_complete(&build_tool_block(&[(
354 + "get_weather",
355 + vec![
356 + (
357 + "city",
358 + "Seattle &lt;/parameter&gt;&lt;/invoke&gt;&lt;/minimax:tool_call&gt;",
359 + ),
360 + ("days", "5"),
361 + ],
362 + )]))
363 + .unwrap();
364 +
365 + assert_eq!(
366 + serde_json::from_str::<Value>(&result.calls[0].arguments).unwrap(),
367 + json!({
368 + "city": "Seattle </parameter></invoke></minimax:tool_call>",
369 + "days": 5,
370 + })
371 + );
372 + }
373 +
374 + #[test]
375 + fn minimax_m2_parse_complete_handles_multiline_parameters() {
376 + let mut parser = MinimaxM2ToolParser::new(&test_tools());
377 + let result = parser
378 + .parse_complete(
379 + "<minimax:tool_call>\
380 + <invoke name=\"calculate_area\">\
381 + <parameter name=\"shape\">\nrectangle\n</parameter>\
382 + <parameter name=\"dimensions\">{\"width\":10,\n\"height\":20}</parameter>\
383 + <parameter name=\"precision\">2</parameter>\
384 + </invoke>\
385 + </minimax:tool_call>",
386 + )
387 + .unwrap();
388 +
389 + assert_eq!(
390 + serde_json::from_str::<Value>(&result.calls[0].arguments).unwrap(),
391 + json!({
392 + "shape": "\nrectangle\n",
393 + "dimensions": { "width": 10, "height": 20 },
394 + "precision": 2,
395 + })
396 + );
397 + }
398 +
399 + #[test]
400 + fn minimax_m2_streaming_extracts_single_tool_call() {
401 + let mut parser = MinimaxM2ToolParser::new(&test_tools());
402 + let result = collect_stream(
403 + &mut parser,
404 + &[
405 + "<minimax:tool_call>",
406 + r#"<invoke name="get_weather">"#,
407 + r#"<parameter name="city">Seattle</parameter>"#,
408 + "</invoke></minimax:tool_call>",
409 + ],
410 + );
411 +
412 + assert!(result.normal_text.is_empty());
413 + assert_eq!(result.calls.len(), 1);
414 + assert_eq!(result.calls[0].name.as_deref(), Some("get_weather"));
415 + assert_eq!(
416 + serde_json::from_str::<Value>(&result.calls[0].arguments).unwrap(),
417 + json!({ "city": "Seattle" })
418 + );
419 + }
420 +
421 + #[test]
422 + fn minimax_m2_streaming_preserves_prefix_text() {
423 + let mut parser = MinimaxM2ToolParser::new(&test_tools());
424 + let result = collect_stream(
425 + &mut parser,
426 + &[
427 + "Let me check. ",
428 + "<minimax:tool_call>",
429 + r#"<invoke name="get_weather"><parameter name="city">Seattle</parameter></invoke>"#,
430 + "</minimax:tool_call>",
431 + ],
432 + );
433 +
434 + assert_eq!(result.normal_text, "Let me check. ");
435 + assert_eq!(result.calls.len(), 1);
436 + }
437 +
438 + #[test]
439 + fn minimax_m2_streaming_without_tool_call_emits_text_incrementally() {
440 + let mut parser = MinimaxM2ToolParser::new(&test_tools());
441 + let result = collect_stream(&mut parser, &["Hello, ", "world!"]);
442 +
443 + assert_eq!(result.normal_text, "Hello, world!");
444 + assert!(result.calls.is_empty());
445 + }
446 +
447 + #[test]
448 + fn minimax_m2_streaming_handles_marker_split_across_chunks() {
449 + let text = build_tool_block(&[("get_weather", vec![("city", "Seattle")])]);
450 + let chunks = split_by_chars(&text, 3);
451 + let mut parser = MinimaxM2ToolParser::new(&test_tools());
452 + let result = collect_stream(&mut parser, &chunks);
453 +
454 + assert_eq!(result.calls.len(), 1);
455 + assert!(result.normal_text.is_empty());
456 + }
457 +
458 + #[test]
459 + fn minimax_m2_streaming_extracts_multiple_invokes_in_order() {
460 + let text = build_tool_block(&[
461 + ("get_weather", vec![("city", "Seattle")]),
462 + ("get_weather", vec![("city", "NYC")]),
463 + ]);
464 + let chunks = split_by_chars(&text, 7);
465 + let mut parser = MinimaxM2ToolParser::new(&test_tools());
466 + let result = collect_stream(&mut parser, &chunks);
467 +
468 + assert_eq!(result.calls.len(), 2);
469 + assert_eq!(result.calls[0].tool_index, 0);
470 + assert_eq!(result.calls[1].tool_index, 1);
471 + }
472 +
473 + #[test]
474 + fn minimax_m2_streaming_ignores_text_after_tool_block() {
475 + let text = format!(
476 + "{} ignored",
477 + build_tool_block(&[("get_weather", vec![("city", "Seattle")])])
478 + );
479 + let chunks = split_by_chars(&text, 5);
480 + let mut parser = MinimaxM2ToolParser::new(&test_tools());
481 + let result = collect_stream(&mut parser, &chunks);
482 +
483 + assert!(result.normal_text.is_empty());
484 + assert_eq!(result.calls.len(), 1);
485 + }
486 +
487 + #[test]
488 + fn minimax_m2_streaming_does_not_emit_incomplete_tool_call() {
489 + let mut parser = MinimaxM2ToolParser::new(&test_tools());
490 + let result = parser.push(r#"<minimax:tool_call><invoke name="get_weather">"#).unwrap();
491 +
492 + assert!(result.normal_text.is_empty());
493 + assert!(result.calls.is_empty());
494 + }
495 +
496 + #[test]
497 + fn minimax_m2_finish_fails_incomplete_tool_call() {
498 + let mut parser = MinimaxM2ToolParser::new(&test_tools());
499 + parser.push(r#"<minimax:tool_call><invoke name="get_weather">"#).unwrap();
500 +
501 + assert!(parser.finish().is_err());
502 + }
503 +
504 + #[test]
505 + fn minimax_m2_finish_fails_after_bare_tool_block_start() {
506 + let mut parser = MinimaxM2ToolParser::new(&test_tools());
507 + parser.push("<minimax:tool_call>").unwrap();
508 +
509 + assert!(parser.finish().is_err());
510 + }
511 +
512 + #[test]
513 + fn minimax_m2_malformed_tool_call_fails_fast() {
514 + let mut parser = MinimaxM2ToolParser::new(&test_tools());
515 + let error = parser.push("<minimax:tool_call><bad></minimax:tool_call>").unwrap_err();
516 +
517 + expect!["tool parser parsing failed: "].assert_eq(&error.to_report_string());
518 + }
519 +}
rust/src/tool-parser/src/parameters.rs new
+508
@@ -0,0 +1,508 @@
1 +use std::collections::BTreeMap;
2 +
3 +use serde_json::{Number, Value};
4 +
5 +use crate::Tool;
6 +
7 +/// Normalized parameter schemas for all tools in one request.
8 +#[derive(Debug, Clone, Default, PartialEq, Eq)]
9 +pub(super) struct ToolSchemas {
10 + tools: BTreeMap<String, ToolSchema>,
11 +}
12 +
13 +/// Normalized parameter schema for one tool.
14 +///
15 +/// This is a minimal subset of JSON Schema with some normalization heuristics
16 +/// to support common schema patterns and upstream schema variations, focused on
17 +/// coercing raw string parameter values into more specific JSON types for
18 +/// downstream tool call execution.
19 +#[derive(Debug, Clone, Default, PartialEq, Eq)]
20 +pub(super) struct ToolSchema {
21 + params: BTreeMap<String, JsonParamType>,
22 +}
23 +
24 +/// Normalized JSON parameter type used for raw string coercion.
25 +#[derive(Debug, Clone, PartialEq, Eq)]
26 +pub(super) enum JsonParamType {
27 + String,
28 + Integer,
29 + Number,
30 + Boolean,
31 + Object,
32 + Array,
33 + Null,
34 + OneOf(Vec<JsonParamType>),
35 +}
36 +
37 +impl ToolSchemas {
38 + /// Normalize OpenAI-style tool parameter JSON schemas for one request.
39 + pub(super) fn from_tools(tools: &[Tool]) -> Self {
40 + let tools = tools
41 + .iter()
42 + .map(|tool| (tool.name.clone(), ToolSchema::from_schema(&tool.parameters)))
43 + .collect();
44 +
45 + Self { tools }
46 + }
47 +
48 + /// Convert raw string parameter values for one named tool.
49 + ///
50 + /// Unknown tool names use an empty schema, so all parameters fall back to
51 + /// strings.
52 + pub(super) fn convert_params_with_schema(
53 + &self,
54 + function_name: &str,
55 + params: Vec<(String, String)>,
56 + ) -> serde_json::Map<String, Value> {
57 + let tool_schema = self.tools.get(function_name).unwrap_or(ToolSchema::empty());
58 + let mut converted = serde_json::Map::with_capacity(params.len());
59 + for (name, value) in params {
60 + let value = tool_schema.convert(&name, &value);
61 + converted.insert(name, value);
62 + }
63 + converted
64 + }
65 +
66 + /// Convert one raw string parameter value for one named tool.
67 + pub(super) fn convert_param_with_schema(
68 + &self,
69 + function_name: &str,
70 + name: &str,
71 + value: &str,
72 + ) -> Value {
73 + let tool_schema = self.tools.get(function_name).unwrap_or(ToolSchema::empty());
74 + tool_schema.convert(name, value)
75 + }
76 +}
77 +
78 +impl ToolSchema {
79 + /// Return an empty schema with no parameter information, which causes all
80 + /// parameters to be treated as strings.
81 + const fn empty() -> &'static Self {
82 + static EMPTY: ToolSchema = ToolSchema {
83 + params: BTreeMap::new(),
84 + };
85 + &EMPTY
86 + }
87 +
88 + /// Normalize an OpenAI-style tool parameters JSON schema.
89 + fn from_schema(parameters: &Value) -> Self {
90 + let Some(properties) = parameters.get("properties").and_then(Value::as_object) else {
91 + return Self::default();
92 + };
93 +
94 + let params = properties
95 + .iter()
96 + .filter_map(|(name, schema)| {
97 + JsonParamType::from_schema(schema).map(|param_type| (name.clone(), param_type))
98 + })
99 + .collect();
100 +
101 + Self { params }
102 + }
103 +
104 + /// Convert one raw parameter value using its normalized schema type.
105 + ///
106 + /// If the parameter name is unknown, or we don't have a schema for it, or
107 + /// the value fails to convert, this falls back to returning the raw
108 + /// string as a JSON string value.
109 + fn convert(&self, name: &str, value: &str) -> Value {
110 + if value.eq_ignore_ascii_case("null") {
111 + return Value::Null;
112 + }
113 +
114 + let Some(param_type) = self.params.get(name) else {
115 + return Value::String(value.to_string());
116 + };
117 +
118 + convert_value(param_type, value).unwrap_or_else(|| Value::String(value.to_string()))
119 + }
120 +}
121 +
122 +impl JsonParamType {
123 + /// Normalize one parameter property schema.
124 + fn from_schema(schema: &Value) -> Option<Self> {
125 + let schema = schema.as_object()?;
126 +
127 + if let Some(type_value) = schema.get("type") {
128 + return Self::from_type_value(type_value);
129 + }
130 +
131 + if let Some(composite) = schema.get("anyOf").or_else(|| schema.get("oneOf")) {
132 + let param_type = composite
133 + .as_array()
134 + .map(|schemas| schemas.iter().filter_map(Self::from_schema).collect::<Vec<_>>())
135 + .filter(|types| !types.is_empty())
136 + .map(Self::one_of)
137 + .unwrap_or(Self::Object);
138 + return Some(param_type);
139 + }
140 +
141 + if schema.contains_key("enum") {
142 + return Some(Self::String);
143 + }
144 + if schema.contains_key("items") {
145 + return Some(Self::Array);
146 + }
147 + if schema.contains_key("properties") {
148 + return Some(Self::Object);
149 + }
150 +
151 + None
152 + }
153 +
154 + /// Normalize a JSON schema `type` value.
155 + fn from_type_value(type_value: &Value) -> Option<Self> {
156 + match type_value {
157 + Value::String(kind) => Self::from_type_name(kind),
158 + Value::Array(kinds) => {
159 + let types = kinds
160 + .iter()
161 + .filter_map(Value::as_str)
162 + .filter_map(Self::from_type_name)
163 + .collect::<Vec<_>>();
164 + if types.is_empty() {
165 + None
166 + } else {
167 + Some(Self::one_of(types))
168 + }
169 + }
170 + _ => None,
171 + }
172 + }
173 +
174 + /// Normalize one JSON schema type name.
175 + fn from_type_name(kind: &str) -> Option<Self> {
176 + let kind = kind.trim().to_ascii_lowercase();
177 + match kind.as_str() {
178 + "string" | "str" | "text" | "varchar" | "char" | "enum" => Some(Self::String),
179 + "integer" | "int" => Some(Self::Integer),
180 + "number" | "float" => Some(Self::Number),
181 + "boolean" | "bool" | "binary" => Some(Self::Boolean),
182 + "object" => Some(Self::Object),
183 + "array" | "arr" | "sequence" => Some(Self::Array),
184 + "null" => Some(Self::Null),
185 + _ if kind.starts_with("int")
186 + || kind.starts_with("uint")
187 + || kind.starts_with("long")
188 + || kind.starts_with("short")
189 + || kind.starts_with("unsigned") =>
190 + {
191 + Some(Self::Integer)
192 + }
193 + _ if kind.starts_with("num") || kind.starts_with("float") => Some(Self::Number),
194 + _ if kind.starts_with("dict") => Some(Self::Object),
195 + _ if kind.starts_with("list") => Some(Self::Array),
196 + _ => None,
197 + }
198 + }
199 +
200 + /// Collapse a candidate type list into one normalized type.
201 + fn one_of(mut types: Vec<Self>) -> Self {
202 + if types.len() == 1 {
203 + types.remove(0)
204 + } else {
205 + Self::OneOf(types)
206 + }
207 + }
208 +}
209 +
210 +/// Convert one raw string value to a normalized JSON type.
211 +fn convert_value(param_type: &JsonParamType, value: &str) -> Option<Value> {
212 + match param_type {
213 + JsonParamType::String => Some(Value::String(value.to_string())),
214 + JsonParamType::Integer => value.parse::<i64>().ok().map(Number::from).map(Value::Number),
215 + JsonParamType::Number => convert_number(value),
216 + JsonParamType::Boolean => convert_boolean(value),
217 + JsonParamType::Object | JsonParamType::Array => serde_json::from_str(value).ok(),
218 + JsonParamType::Null => value.eq_ignore_ascii_case("null").then_some(Value::Null),
219 + JsonParamType::OneOf(types) => {
220 + types.iter().find_map(|param_type| convert_value(param_type, value))
221 + }
222 + }
223 +}
224 +
225 +/// Convert one raw string value to a JSON number.
226 +fn convert_number(value: &str) -> Option<Value> {
227 + if let Ok(parsed) = value.parse::<i64>() {
228 + return Some(Value::Number(Number::from(parsed)));
229 + }
230 + Number::from_f64(value.parse::<f64>().ok()?).map(Value::Number)
231 +}
232 +
233 +/// Convert one raw string value to a boolean.
234 +fn convert_boolean(value: &str) -> Option<Value> {
235 + match value.trim().to_ascii_lowercase().as_str() {
236 + "true" | "1" => Some(Value::Bool(true)),
237 + "false" | "0" => Some(Value::Bool(false)),
238 + _ => None,
239 + }
240 +}
241 +
242 +#[cfg(test)]
243 +mod tests {
244 + use serde_json::json;
245 +
246 + use super::{ToolSchema, ToolSchemas};
247 + use crate::Tool;
248 +
249 + fn test_tool(name: &str, parameters: serde_json::Value) -> Tool {
250 + Tool {
251 + name: name.to_string(),
252 + description: None,
253 + parameters,
254 + strict: None,
255 + }
256 + }
257 +
258 + #[test]
259 + fn invalid_schema_converts_everything_as_string() {
260 + let params = ToolSchema::from_schema(&json!({ "type": "object" }));
261 +
262 + assert_eq!(params.convert("count", "42"), json!("42"));
263 + assert_eq!(params.convert("count", "null"), json!(null));
264 + }
265 +
266 + #[test]
267 + fn skips_unknown_property_schema_and_unknown_type() {
268 + let params = ToolSchema::from_schema(&json!({
269 + "type": "object",
270 + "properties": {
271 + "unknown_schema": true,
272 + "unknown_type": { "type": "mystery" },
273 + "known": { "type": "integer" }
274 + }
275 + }));
276 +
277 + assert_eq!(params.convert("unknown_schema", "42"), json!("42"));
278 + assert_eq!(params.convert("unknown_type", "42"), json!("42"));
279 + assert_eq!(params.convert("known", "42"), json!(42));
280 + }
281 +
282 + #[test]
283 + fn converts_supported_types() {
284 + let params = ToolSchema::from_schema(&json!({
285 + "type": "object",
286 + "properties": {
287 + "text": { "type": "string" },
288 + "count": { "type": "integer" },
289 + "size": { "type": "number" },
290 + "enabled": { "type": "boolean" },
291 + "payload": { "type": "object" },
292 + "items": { "type": "array" },
293 + "nothing": { "type": "null" }
294 + }
295 + }));
296 +
297 + assert_eq!(params.convert("text", "42"), json!("42"));
298 + assert_eq!(params.convert("count", "42"), json!(42));
299 + assert_eq!(params.convert("size", "5.0"), json!(5.0));
300 + assert_eq!(params.convert("enabled", "1"), json!(true));
301 + assert_eq!(params.convert("payload", r#"{"k":1}"#), json!({ "k": 1 }));
302 + assert_eq!(params.convert("items", "[1,2]"), json!([1, 2]));
303 + assert_eq!(params.convert("nothing", "null"), json!(null));
304 + }
305 +
306 + #[test]
307 + fn number_conversion_parses_int_then_float() {
308 + let params = ToolSchema::from_schema(&json!({
309 + "type": "object",
310 + "properties": {
311 + "value": { "type": "number" }
312 + }
313 + }));
314 +
315 + assert_eq!(params.convert("value", "5"), json!(5));
316 + assert_eq!(params.convert("value", "5.0"), json!(5.0));
317 + assert_eq!(params.convert("value", "5."), json!(5.0));
318 + assert_eq!(params.convert("value", "+1"), json!(1));
319 + assert_eq!(params.convert("value", "+1.0"), json!(1.0));
320 + assert_eq!(
321 + params.convert("value", "9223372036854775807.5"),
322 + json!(9223372036854775808.0)
323 + );
324 + }
325 +
326 + #[test]
327 + fn converts_upstream_aliases() {
328 + let params = ToolSchema::from_schema(&json!({
329 + "type": "object",
330 + "properties": {
331 + "s": { "type": "varchar" },
332 + "i": { "type": "unsigned_int" },
333 + "n": { "type": "float64" },
334 + "b": { "type": "binary" },
335 + "a": { "type": "sequence" },
336 + "o": { "type": "dict" }
337 + }
338 + }));
339 +
340 + assert_eq!(params.convert("s", "x"), json!("x"));
341 + assert_eq!(params.convert("i", "7"), json!(7));
342 + assert_eq!(params.convert("n", "7.5"), json!(7.5));
343 + assert_eq!(params.convert("b", "true"), json!(true));
344 + assert_eq!(params.convert("a", "[1]"), json!([1]));
345 + assert_eq!(params.convert("o", r#"{"x":1}"#), json!({ "x": 1 }));
346 + }
347 +
348 + #[test]
349 + fn preserves_union_type_order() {
350 + let integer_first = ToolSchema::from_schema(&json!({
351 + "type": "object",
352 + "properties": {
353 + "value": { "type": ["integer", "string"] }
354 + }
355 + }));
356 + let string_first = ToolSchema::from_schema(&json!({
357 + "type": "object",
358 + "properties": {
359 + "value": { "type": ["string", "integer"] }
360 + }
361 + }));
362 +
363 + assert_eq!(integer_first.convert("value", "42"), json!(42));
364 + assert_eq!(string_first.convert("value", "42"), json!("42"));
365 + }
366 +
367 + #[test]
368 + fn converts_composite_schemas() {
369 + let params = ToolSchema::from_schema(&json!({
370 + "type": "object",
371 + "properties": {
372 + "choice": {
373 + "anyOf": [
374 + { "type": "integer" },
375 + { "type": "string" }
376 + ]
377 + },
378 + "fallback_object": {
379 + "oneOf": [
380 + { "type": "mystery" }
381 + ]
382 + }
383 + }
384 + }));
385 +
386 + assert_eq!(params.convert("choice", "42"), json!(42));
387 + assert_eq!(
388 + params.convert("fallback_object", r#"{"x":1}"#),
389 + json!({ "x": 1 })
390 + );
391 + }
392 +
393 + #[test]
394 + fn infers_type_from_schema_shape_without_type() {
395 + let params = ToolSchema::from_schema(&json!({
396 + "type": "object",
397 + "properties": {
398 + "choice": { "enum": ["a", "b"] },
399 + "items": { "items": { "type": "integer" } },
400 + "payload": { "properties": { "x": { "type": "integer" } } }
401 + }
402 + }));
403 +
404 + assert_eq!(params.convert("choice", "a"), json!("a"));
405 + assert_eq!(params.convert("items", "[1,2]"), json!([1, 2]));
406 + assert_eq!(params.convert("payload", r#"{"x":1}"#), json!({ "x": 1 }));
407 + }
408 +
409 + #[test]
410 + fn converts_params_for_known_tool() {
411 + let schemas = ToolSchemas::from_tools(&[test_tool(
412 + "search",
413 + json!({
414 + "type": "object",
415 + "properties": {
416 + "query": { "type": "string" },
417 + "topn": { "type": "integer" }
418 + }
419 + }),
420 + )]);
421 +
422 + let converted = schemas.convert_params_with_schema(
423 + "search",
424 + vec![
425 + ("query".to_string(), "rust".to_string()),
426 + ("topn".to_string(), "5".to_string()),
427 + ],
428 + );
429 +
430 + assert_eq!(converted.get("query"), Some(&json!("rust")));
431 + assert_eq!(converted.get("topn"), Some(&json!(5)));
432 + }
433 +
434 + #[test]
435 + fn convert_params_falls_back_to_string_for_failed_coercion() {
436 + let schemas = ToolSchemas::from_tools(&[test_tool(
437 + "convert",
438 + json!({
439 + "type": "object",
440 + "properties": {
441 + "whole": { "type": "number" },
442 + "flag": { "type": "boolean" },
443 + "payload": { "type": "object" },
444 + "items": { "type": "array" },
445 + "missing_type": {}
446 + }
447 + }),
448 + )]);
449 +
450 + let converted = schemas.convert_params_with_schema(
451 + "convert",
452 + vec![
453 + ("whole".to_string(), "not-a-number".to_string()),
454 + ("flag".to_string(), "maybe".to_string()),
455 + ("payload".to_string(), "not-json".to_string()),
456 + ("items".to_string(), "not-json".to_string()),
457 + ("missing_type".to_string(), "42".to_string()),
458 + ("unknown_param".to_string(), "42".to_string()),
459 + ],
460 + );
461 +
462 + assert_eq!(converted.get("whole"), Some(&json!("not-a-number")));
463 + assert_eq!(converted.get("flag"), Some(&json!("maybe")));
464 + assert_eq!(converted.get("payload"), Some(&json!("not-json")));
465 + assert_eq!(converted.get("items"), Some(&json!("not-json")));
466 + assert_eq!(converted.get("missing_type"), Some(&json!("42")));
467 + assert_eq!(converted.get("unknown_param"), Some(&json!("42")));
468 + }
469 +
470 + #[test]
471 + fn convert_params_preserves_null_for_known_param() {
472 + let schemas = ToolSchemas::from_tools(&[test_tool(
473 + "convert",
474 + json!({
475 + "type": "object",
476 + "properties": {
477 + "value": { "type": "string" }
478 + }
479 + }),
480 + )]);
481 +
482 + let converted = schemas
483 + .convert_params_with_schema("convert", vec![("value".to_string(), "NULL".to_string())]);
484 +
485 + assert_eq!(converted.get("value"), Some(&json!(null)));
486 + }
487 +
488 + #[test]
489 + fn unknown_tool_converts_values_without_schema() {
490 + let schemas = ToolSchemas::from_tools(&[test_tool(
491 + "search",
492 + json!({ "type": "object", "properties": {} }),
493 + )]);
494 +
495 + let converted = schemas.convert_params_with_schema(
496 + "missing",
497 + vec![
498 + ("query".to_string(), "rust".to_string()),
499 + ("topn".to_string(), "5".to_string()),
500 + ("nullish".to_string(), "null".to_string()),
501 + ],
502 + );
503 +
504 + assert_eq!(converted.get("query"), Some(&json!("rust")));
505 + assert_eq!(converted.get("topn"), Some(&json!("5")));
506 + assert_eq!(converted.get("nullish"), Some(&json!(null)));
507 + }
508 +}
rust/src/tool-parser/src/qwen_coder.rs new
+632
@@ -0,0 +1,632 @@
1 +use winnow::ascii::multispace0 as ws0;
2 +use winnow::combinator::{alt, delimited, eof, repeat, seq, terminated};
3 +use winnow::prelude::*;
4 +use winnow::stream::Partial;
5 +use winnow::token::{literal, take_until};
6 +
7 +use super::parameters::ToolSchemas;
8 +use super::utils::{parse_buffered_event, safe_text_len, xml_unescape};
9 +use super::{Result, ToolCallDelta, ToolParseResult, ToolParser};
10 +use crate::Tool;
11 +
12 +const TOOL_CALL_START: &str = "<tool_call>";
13 +const TOOL_CALL_END: &str = "</tool_call>";
14 +const FUNCTION_START: &str = "<function=";
15 +const FUNCTION_END: &str = "</function>";
16 +const PARAMETER_START: &str = "<parameter=";
17 +const PARAMETER_END: &str = "</parameter>";
18 +
19 +type QwenCoderInput<'i> = Partial<&'i str>;
20 +
21 +#[derive(Debug, Clone, Copy, PartialEq, Eq)]
22 +enum QwenCoderMode {
23 + Text,
24 + ToolCall,
25 +}
26 +
27 +#[derive(Debug, Clone, PartialEq, Eq)]
28 +enum QwenCoderEvent {
29 + Text {
30 + len: usize,
31 + },
32 + ToolCallStart,
33 + ToolCall {
34 + name: String,
35 + raw_params: Vec<(String, String)>,
36 + },
37 +}
38 +
39 +/// Tool parser for Qwen Coder XML-style tool calls.
40 +///
41 +/// Example tool call content:
42 +///
43 +/// ```text
44 +/// <tool_call>
45 +/// <function=get_weather>
46 +/// <parameter=location>杭州</parameter>
47 +/// </function>
48 +/// </tool_call>
49 +/// ```
50 +///
51 +/// Arguments are emitted only after a full `tool_call` block is parsed.
52 +///
53 +/// Note: parallel calls are represented as repeated
54 +/// `<tool_call>...</tool_call>` blocks, not as multiple calls inside one tag.
55 +pub struct Qwen3CoderToolParser {
56 + buffer: String,
57 + mode: QwenCoderMode,
58 + emitted_tool_count: usize,
59 + tool_parameters: ToolSchemas,
60 +}
61 +
62 +impl Qwen3CoderToolParser {
63 + /// Create a Qwen Coder tool parser.
64 + fn new(tools: &[Tool]) -> Self {
65 + Self {
66 + buffer: String::new(),
67 + mode: QwenCoderMode::Text,
68 + emitted_tool_count: 0,
69 + tool_parameters: ToolSchemas::from_tools(tools),
70 + }
71 + }
72 +
73 + /// Apply one parsed Qwen Coder event to parser state and output.
74 + fn apply_event(&mut self, event: QwenCoderEvent, result: &mut ToolParseResult) -> Result<()> {
75 + match event {
76 + QwenCoderEvent::Text { len: consumed_len } => {
77 + result.normal_text.push_str(&self.buffer[..consumed_len]);
78 + }
79 + QwenCoderEvent::ToolCallStart => self.mode = QwenCoderMode::ToolCall,
80 + QwenCoderEvent::ToolCall { name, raw_params } => {
81 + self.mode = QwenCoderMode::Text;
82 + let arguments = self.tool_parameters.convert_params_with_schema(&name, raw_params);
83 + let arguments = serde_json::to_string(&arguments)
84 + .map_err(|error| parsing_failed!("failed to serialize arguments: {}", error))?;
85 +
86 + result.calls.push(ToolCallDelta {
87 + tool_index: self.emitted_tool_count,
88 + name: Some(name),
89 + arguments,
90 + });
91 + self.emitted_tool_count += 1;
92 + }
93 + }
94 + Ok(())
95 + }
96 +
97 + /// Reset all streaming state.
98 + fn reset(&mut self) {
99 + self.buffer.clear();
100 + self.mode = QwenCoderMode::Text;
101 + self.emitted_tool_count = 0;
102 + }
103 +}
104 +
105 +impl ToolParser for Qwen3CoderToolParser {
106 + /// Create a boxed Qwen Coder tool parser.
107 + fn create(tools: &[Tool]) -> Result<Box<dyn ToolParser>>
108 + where
109 + Self: Sized + 'static,
110 + {
111 + Ok(Box::new(Self::new(tools)))
112 + }
113 +
114 + /// Push one decoded text chunk through the Qwen Coder parser.
115 + fn push(&mut self, chunk: &str) -> Result<ToolParseResult> {
116 + self.buffer.push_str(chunk);
117 + let mut result = ToolParseResult::default();
118 +
119 + while let Some((event, consumed_len)) = parse_buffered_event(&self.buffer, |input| {
120 + parse_next_qwen_coder_event(input, self.mode)
121 + })? {
122 + self.apply_event(event, &mut result)?;
123 + self.buffer.drain(..consumed_len);
124 + }
125 +
126 + Ok(result)
127 + }
128 +
129 + /// Flush buffered text and reset parser state.
130 + fn finish(&mut self) -> Result<ToolParseResult> {
131 + let mut result = ToolParseResult::default();
132 + if !self.buffer.is_empty() {
133 + if self.mode == QwenCoderMode::ToolCall || self.buffer.starts_with(TOOL_CALL_START) {
134 + return Err(parsing_failed!("incomplete Qwen Coder tool call"));
135 + }
136 + result.normal_text.push_str(&self.buffer);
137 + }
138 + self.reset();
139 + Ok(result)
140 + }
141 +}
142 +
143 +/// Parse a Qwen Coder event for the current parser mode.
144 +fn parse_next_qwen_coder_event(
145 + input: &mut QwenCoderInput<'_>,
146 + mode: QwenCoderMode,
147 +) -> ModalResult<QwenCoderEvent> {
148 + match mode {
149 + QwenCoderMode::Text => parse_text_event(input),
150 + QwenCoderMode::ToolCall => tool_call_event(input),
151 + }
152 +}
153 +
154 +/// Parse a text-mode Qwen Coder event.
155 +fn parse_text_event(input: &mut QwenCoderInput<'_>) -> ModalResult<QwenCoderEvent> {
156 + alt((tool_call_start_event, safe_text_event)).parse_next(input)
157 +}
158 +
159 +/// Parse a Qwen Coder tool-call start marker.
160 +fn tool_call_start_event(input: &mut QwenCoderInput<'_>) -> ModalResult<QwenCoderEvent> {
161 + literal(TOOL_CALL_START).value(QwenCoderEvent::ToolCallStart).parse_next(input)
162 +}
163 +
164 +/// Parse a safe text run before the next Qwen Coder marker.
165 +fn safe_text_event(input: &mut QwenCoderInput<'_>) -> ModalResult<QwenCoderEvent> {
166 + safe_text_len(input, TOOL_CALL_START).map(|len| QwenCoderEvent::Text { len })
167 +}
168 +
169 +/// Parse a complete Qwen Coder tool call.
170 +fn tool_call_event(input: &mut QwenCoderInput<'_>) -> ModalResult<QwenCoderEvent> {
171 + let (body,) = seq!(
172 + _: ws0,
173 + take_until(0.., TOOL_CALL_END),
174 + _: literal(TOOL_CALL_END),
175 + )
176 + .parse_next(input)?;
177 +
178 + parse_tool_call_body(body)
179 +}
180 +
181 +/// Parse a Qwen Coder function block.
182 +fn function_event(input: &mut &str) -> ModalResult<QwenCoderEvent> {
183 + let (name, raw_params) = seq!(
184 + _: literal(FUNCTION_START),
185 + take_until(1.., ">"),
186 + _: ">",
187 + _: ws0,
188 + repeat(0.., terminated(parameter, ws0)),
189 + _: literal(FUNCTION_END),
190 + )
191 + .parse_next(input)?;
192 +
193 + Ok(QwenCoderEvent::ToolCall {
194 + name: name.to_string(),
195 + raw_params,
196 + })
197 +}
198 +
199 +/// Parse a Qwen Coder parameter block.
200 +fn parameter(input: &mut &str) -> ModalResult<(String, String)> {
201 + let (name, value) = seq!(
202 + _: literal(PARAMETER_START),
203 + take_until(1.., ">"),
204 + _: ">",
205 + take_until(0.., PARAMETER_END).map(trim_one_wrapping_newline).map(xml_unescape),
206 + _: literal(PARAMETER_END),
207 + )
208 + .parse_next(input)?;
209 +
210 + Ok((name.to_string(), value.into_owned()))
211 +}
212 +
213 +/// Parse a Qwen Coder tool-call body.
214 +fn parse_tool_call_body(body: &str) -> ModalResult<QwenCoderEvent> {
215 + let mut input = body;
216 + delimited(ws0, function_event, (ws0, eof)).parse_next(&mut input)
217 +}
218 +
219 +/// Trim a single leading and trailing newline from a parameter value.
220 +fn trim_one_wrapping_newline(value: &str) -> &str {
221 + let value = value.strip_prefix('\n').unwrap_or(value);
222 + value.strip_suffix('\n').unwrap_or(value)
223 +}
224 +
225 +#[cfg(test)]
226 +mod tests {
227 + use expect_test::expect;
228 + use serde_json::{Value, json};
229 + use thiserror_ext::AsReport;
230 +
231 + use super::{Qwen3CoderToolParser, ToolParser};
232 + use crate::test_utils::{collect_stream, split_by_chars, test_tools};
233 +
234 + fn build_tool_call(function_name: &str, params: &[(&str, &str)]) -> String {
235 + let params = params
236 + .iter()
237 + .map(|(name, value)| format!("<parameter={name}>{value}</parameter>"))
238 + .collect::<Vec<_>>()
239 + .join("\n");
240 + format!("<tool_call>\n<function={function_name}>\n{params}\n</function>\n</tool_call>")
241 + }
242 +
243 + #[test]
244 + fn qwen_coder_parse_complete_without_tool_call_keeps_text() {
245 + let mut parser = Qwen3CoderToolParser::new(&test_tools());
246 + let result = parser.parse_complete("Hello, world!").unwrap();
247 +
248 + assert_eq!(result.normal_text, "Hello, world!");
249 + assert!(result.calls.is_empty());
250 + }
251 +
252 + #[test]
253 + fn qwen_coder_parse_complete_extracts_single_tool_call() {
254 + let mut parser = Qwen3CoderToolParser::new(&test_tools());
255 + let result = parser
256 + .parse_complete(&build_tool_call(
257 + "get_weather",
258 + &[("location", "SF"), ("date", "2026-04-29")],
259 + ))
260 + .unwrap();
261 +
262 + assert!(result.normal_text.is_empty());
263 + assert_eq!(result.calls.len(), 1);
264 + assert_eq!(result.calls[0].name.as_deref(), Some("get_weather"));
265 + assert_eq!(
266 + serde_json::from_str::<Value>(&result.calls[0].arguments).unwrap(),
267 + json!({
268 + "location": "SF",
269 + "date": "2026-04-29"
270 + })
271 + );
272 + }
273 +
274 + #[test]
275 + fn qwen_coder_parse_complete_preserves_prefix_text() {
276 + let mut parser = Qwen3CoderToolParser::new(&test_tools());
277 + let output = format!(
278 + "Thinking... {}",
279 + build_tool_call("get_weather", &[("location", "NYC")])
280 + );
281 + let result = parser.parse_complete(&output).unwrap();
282 +
283 + assert_eq!(result.normal_text, "Thinking... ");
284 + assert_eq!(result.calls.len(), 1);
285 + }
286 +
287 + #[test]
288 + fn qwen_coder_parse_complete_converts_schema_types() {
289 + let mut parser = Qwen3CoderToolParser::new(&test_tools());
290 + let result = parser
291 + .parse_complete(&build_tool_call(
292 + "convert",
293 + &[
294 + ("whole", "5.0"),
295 + ("flag", "true"),
296 + ("payload", r#"{"nested":true}"#),
297 + ("items", "[1,2]"),
298 + ("empty", "42"),
299 + ],
300 + ))
301 + .unwrap();
302 +
303 + assert_eq!(result.calls.len(), 1);
304 + assert_eq!(
305 + serde_json::from_str::<Value>(&result.calls[0].arguments).unwrap(),
306 + json!({
307 + "whole": 5.0,
308 + "flag": true,
309 + "payload": { "nested": true },
310 + "items": [1, 2],
311 + "empty": "42",
312 + })
313 + );
314 + }
315 +
316 + #[test]
317 + fn qwen_coder_parse_complete_extracts_empty_arguments() {
318 + let mut parser = Qwen3CoderToolParser::new(&test_tools());
319 + let result = parser.parse_complete(&build_tool_call("get_weather", &[])).unwrap();
320 +
321 + assert_eq!(result.calls.len(), 1);
322 + assert_eq!(result.calls[0].name.as_deref(), Some("get_weather"));
323 + assert_eq!(
324 + serde_json::from_str::<Value>(&result.calls[0].arguments).unwrap(),
325 + json!({})
326 + );
327 + }
328 +
329 + #[test]
330 + fn qwen_coder_parse_complete_handles_upstream_multiline_typed_params() {
331 + let mut parser = Qwen3CoderToolParser::new(&test_tools());
332 + let result = parser
333 + .parse_complete(
334 + "<tool_call>\n\
335 + <function=calculate_area>\n\
336 + <parameter=shape>\n\
337 + rectangle\n\
338 + </parameter>\n\
339 + <parameter=dimensions>\n\
340 + {\"width\": 10,\n\
341 + \"height\": 20}\n\
342 + </parameter>\n\
343 + <parameter=precision>\n\
344 + 2\n\
345 + </parameter>\n\
346 + </function>\n\
347 + </tool_call>",
348 + )
349 + .unwrap();
350 +
351 + assert_eq!(result.calls.len(), 1);
352 + assert_eq!(result.calls[0].name.as_deref(), Some("calculate_area"));
353 + assert_eq!(
354 + serde_json::from_str::<Value>(&result.calls[0].arguments).unwrap(),
355 + json!({
356 + "shape": "rectangle",
357 + "dimensions": { "width": 10, "height": 20 },
358 + "precision": 2,
359 + })
360 + );
361 + }
362 +
363 + #[test]
364 + fn qwen_coder_parse_complete_handles_nested_json_parameter() {
365 + let mut parser = Qwen3CoderToolParser::new(&test_tools());
366 + let result = parser
367 + .parse_complete(&build_tool_call(
368 + "convert",
369 + &[(
370 + "payload",
371 + r#"{"nested":{"value":[1,2,3],"child":{"enabled":true}}}"#,
372 + )],
373 + ))
374 + .unwrap();
375 +
376 + assert_eq!(result.calls.len(), 1);
377 + assert_eq!(
378 + serde_json::from_str::<Value>(&result.calls[0].arguments).unwrap(),
379 + json!({
380 + "payload": {
381 + "nested": {
382 + "value": [1, 2, 3],
383 + "child": { "enabled": true },
384 + },
385 + },
386 + })
387 + );
388 + }
389 +
390 + #[test]
391 + fn qwen_coder_parse_complete_preserves_xml_like_parameter_values() {
392 + let mut parser = Qwen3CoderToolParser::new(&test_tools());
393 + let result = parser
394 + .parse_complete(&build_tool_call(
395 + "process",
396 + &[
397 + (
398 + "html_content",
399 + r#"<div class="test"><span>Hello</span></div>"#,
400 + ),
401 + ("xml_snippet", r#"<root><child attr="value"/></root>"#),
402 + ],
403 + ))
404 + .unwrap();
405 +
406 + assert_eq!(result.calls.len(), 1);
407 + assert_eq!(
408 + serde_json::from_str::<Value>(&result.calls[0].arguments).unwrap(),
409 + json!({
410 + "html_content": r#"<div class="test"><span>Hello</span></div>"#,
411 + "xml_snippet": r#"<root><child attr="value"/></root>"#,
412 + })
413 + );
414 + }
415 +
416 + #[test]
417 + fn qwen_coder_parse_complete_unescapes_literal_closing_tags_in_parameter_value() {
418 + let mut parser = Qwen3CoderToolParser::new(&test_tools());
419 + let result = parser
420 + .parse_complete(&build_tool_call(
421 + "get_weather",
422 + &[
423 + (
424 + "location",
425 + "杭州 &lt;/parameter&gt;&lt;/function&gt;&lt;/tool_call&gt;",
426 + ),
427 + ("date", "2026-05-08"),
428 + ],
429 + ))
430 + .unwrap();
431 +
432 + assert_eq!(result.calls.len(), 1);
433 + assert_eq!(
434 + serde_json::from_str::<Value>(&result.calls[0].arguments).unwrap(),
435 + json!({
436 + "location": "杭州 </parameter></function></tool_call>",
437 + "date": "2026-05-08",
438 + })
439 + );
440 + }
441 +
442 + #[test]
443 + fn qwen_coder_parse_complete_does_not_double_encode_anyof_object() {
444 + let mut parser = Qwen3CoderToolParser::new(&test_tools());
445 + let result = parser
446 + .parse_complete(&build_tool_call(
447 + "update_record",
448 + &[("data", r#"{"key":"value","count":42}"#)],
449 + ))
450 + .unwrap();
451 +
452 + assert_eq!(result.calls.len(), 1);
453 + assert_eq!(
454 + serde_json::from_str::<Value>(&result.calls[0].arguments).unwrap(),
455 + json!({
456 + "data": { "key": "value", "count": 42 },
457 + })
458 + );
459 + }
460 +
461 + #[test]
462 + fn qwen_coder_streaming_extracts_single_tool_call() {
463 + let mut parser = Qwen3CoderToolParser::new(&test_tools());
464 + let result = collect_stream(
465 + &mut parser,
466 + &[
467 + "<tool_call>\n",
468 + "<function=get_weather>\n",
469 + "<parameter=location>SF</parameter>\n",
470 + "</function>\n",
471 + "</tool_call>",
472 + ],
473 + );
474 +
475 + assert!(result.normal_text.is_empty());
476 + assert_eq!(result.calls.len(), 1);
477 + assert_eq!(result.calls[0].name.as_deref(), Some("get_weather"));
478 + assert_eq!(
479 + serde_json::from_str::<Value>(&result.calls[0].arguments).unwrap(),
480 + json!({ "location": "SF" })
481 + );
482 + }
483 +
484 + #[test]
485 + fn qwen_coder_streaming_preserves_prefix_text() {
486 + let mut parser = Qwen3CoderToolParser::new(&test_tools());
487 + let result = collect_stream(
488 + &mut parser,
489 + &[
490 + "Thinking... ",
491 + "<tool_call>\n",
492 + "<function=get_weather>\n",
493 + "<parameter=location>SF</parameter>\n",
494 + "</function>\n",
495 + "</tool_call>",
496 + ],
497 + );
498 +
499 + assert_eq!(result.normal_text, "Thinking... ");
500 + assert_eq!(result.calls.len(), 1);
501 + }
502 +
503 + #[test]
504 + fn qwen_coder_streaming_without_tool_call_emits_text_incrementally() {
505 + let mut parser = Qwen3CoderToolParser::new(&test_tools());
506 + let result = collect_stream(&mut parser, &["Hello, ", "world!"]);
507 +
508 + assert_eq!(result.normal_text, "Hello, world!");
509 + assert!(result.calls.is_empty());
510 + }
511 +
512 + #[test]
513 + fn qwen_coder_streaming_extracts_multiple_tool_calls_in_order() {
514 + let text = format!(
515 + "{}\n{}",
516 + build_tool_call("get_weather", &[("location", "SF")]),
517 + build_tool_call("get_weather", &[("location", "NYC")])
518 + );
519 + let mut parser = Qwen3CoderToolParser::new(&test_tools());
520 + let result = collect_stream(&mut parser, &[&text]);
521 +
522 + assert_eq!(result.calls.len(), 2);
523 + assert_eq!(result.calls[0].name.as_deref(), Some("get_weather"));
524 + assert_eq!(result.calls[1].name.as_deref(), Some("get_weather"));
525 + assert_eq!(result.calls[0].tool_index, 0);
526 + assert_eq!(result.calls[1].tool_index, 1);
527 + assert_eq!(
528 + serde_json::from_str::<Value>(&result.calls[0].arguments).unwrap(),
529 + json!({ "location": "SF" })
530 + );
531 + assert_eq!(
532 + serde_json::from_str::<Value>(&result.calls[1].arguments).unwrap(),
533 + json!({ "location": "NYC" })
534 + );
535 + }
536 +
537 + #[test]
538 + fn qwen_coder_streaming_preserves_text_between_tool_calls() {
539 + let text = format!(
540 + "I'll check two cities.{}Between calls.{}Done.",
541 + build_tool_call("get_weather", &[("city", "Dallas"), ("state", "TX")]),
542 + build_tool_call("get_weather", &[("city", "Orlando"), ("state", "FL")])
543 + );
544 + let chunks = split_by_chars(&text, 5);
545 + let mut parser = Qwen3CoderToolParser::new(&test_tools());
546 + let result = collect_stream(&mut parser, &chunks);
547 +
548 + assert_eq!(
549 + result.normal_text,
550 + "I'll check two cities.Between calls.Done."
551 + );
552 + assert_eq!(result.calls.len(), 2);
553 + assert_eq!(
554 + serde_json::from_str::<Value>(&result.calls[0].arguments).unwrap(),
555 + json!({ "city": "Dallas", "state": "TX" })
556 + );
557 + assert_eq!(
558 + serde_json::from_str::<Value>(&result.calls[1].arguments).unwrap(),
559 + json!({ "city": "Orlando", "state": "FL" })
560 + );
561 + }
562 +
563 + #[test]
564 + fn qwen_coder_streaming_handles_start_token_split_across_chunks() {
565 + let text = build_tool_call("get_weather", &[("location", "SF")]);
566 + let chunks = split_by_chars(&text, 3);
567 + let mut parser = Qwen3CoderToolParser::new(&test_tools());
568 + let result = collect_stream(&mut parser, &chunks);
569 +
570 + assert_eq!(result.calls.len(), 1);
571 + assert_eq!(
572 + serde_json::from_str::<Value>(&result.calls[0].arguments).unwrap(),
573 + json!({ "location": "SF" })
574 + );
575 + }
576 +
577 + #[test]
578 + fn qwen_coder_streaming_does_not_emit_incomplete_tool_call() {
579 + let mut parser = Qwen3CoderToolParser::new(&test_tools());
580 + let result = parser
581 + .push("<tool_call>\n<function=get_weather>\n<parameter=location>SF</parameter>")
582 + .unwrap();
583 +
584 + assert!(result.normal_text.is_empty());
585 + assert!(result.calls.is_empty());
586 + }
587 +
588 + #[test]
589 + fn qwen_coder_finish_fails_incomplete_tool_call() {
590 + let mut parser = Qwen3CoderToolParser::new(&test_tools());
591 + parser
592 + .push("<tool_call>\n<function=get_weather>\n<parameter=location>SF</parameter>")
593 + .unwrap();
594 +
595 + assert!(parser.finish().is_err());
596 + }
597 +
598 + #[test]
599 + fn qwen_coder_malformed_tool_call_fails_fast() {
600 + let mut parser = Qwen3CoderToolParser::new(&test_tools());
601 + let error = parser.push("<tool_call>\n<bad>\n</tool_call>").unwrap_err();
602 +
603 + expect!["tool parser parsing failed: "].assert_eq(&error.to_report_string());
604 + }
605 +
606 + #[test]
607 + fn qwen_coder_missing_parameter_end_fails_fast_after_function_end() {
608 + let mut parser = Qwen3CoderToolParser::new(&test_tools());
609 + let error = parser
610 + .push(
611 + "<tool_call>\n<function=get_weather>\n<parameter=location>SF</function>\n</tool_call>",
612 + )
613 + .unwrap_err();
614 +
615 + expect!["tool parser parsing failed: "].assert_eq(&error.to_report_string());
616 + }
617 +
618 + #[test]
619 + fn qwen_coder_parse_function_body_trims_one_wrapping_newline() {
620 + let mut parser = Qwen3CoderToolParser::new(&test_tools());
621 + let result = parser
622 + .parse_complete(
623 + "<tool_call>\n<function=get_weather>\n<parameter=location>\nHangzhou\n</parameter>\n</function>\n</tool_call>",
624 + )
625 + .unwrap();
626 +
627 + assert_eq!(
628 + serde_json::from_str::<Value>(&result.calls[0].arguments).unwrap(),
629 + json!({ "location": "Hangzhou" })
630 + );
631 + }
632 +}
rust/src/tool-parser/src/test_utils.rs new
+115
@@ -0,0 +1,115 @@
1 +use serde_json::json;
2 +
3 +use super::{ToolParseResult, ToolParser};
4 +use crate::Tool;
5 +
6 +/// Build a reusable set of function tools for parser unit tests.
7 +pub fn test_tools() -> Vec<Tool> {
8 + vec![
9 + Tool {
10 + name: "get_weather".to_string(),
11 + description: None,
12 + parameters: json!({
13 + "type": "object",
14 + "properties": {
15 + "location": { "type": "string" },
16 + "city": { "type": "string" },
17 + "state": { "type": "string" },
18 + "unit": { "type": "string" },
19 + "date": { "type": "string" },
20 + "days": { "type": "integer" }
21 + }
22 + }),
23 + strict: None,
24 + },
25 + Tool {
26 + name: "add".to_string(),
27 + description: None,
28 + parameters: json!({
29 + "type": "object",
30 + "properties": {
31 + "x": { "type": "integer" },
32 + "y": { "type": "integer" }
33 + }
34 + }),
35 + strict: None,
36 + },
37 + Tool {
38 + name: "convert".to_string(),
39 + description: None,
40 + parameters: json!({
41 + "type": "object",
42 + "properties": {
43 + "whole": { "type": "number" },
44 + "flag": { "type": "boolean" },
45 + "payload": { "type": "object" },
46 + "items": { "type": "array" },
47 + "empty": { "type": "string" }
48 + }
49 + }),
50 + strict: None,
51 + },
52 + Tool {
53 + name: "calculate_area".to_string(),
54 + description: None,
55 + parameters: json!({
56 + "type": "object",
57 + "properties": {
58 + "shape": { "type": "string" },
59 + "dimensions": { "type": "object" },
60 + "precision": { "type": "integer" }
61 + }
62 + }),
63 + strict: None,
64 + },
65 + Tool {
66 + name: "update_record".to_string(),
67 + description: None,
68 + parameters: json!({
69 + "type": "object",
70 + "properties": {
71 + "data": {
72 + "anyOf": [
73 + { "type": "object" },
74 + { "type": "null" }
75 + ]
76 + }
77 + }
78 + }),
79 + strict: None,
80 + },
81 + ]
82 +}
83 +
84 +/// Push chunks through a streaming parser and coalesce its tool-call deltas.
85 +pub fn collect_stream<T: ToolParser + ?Sized>(parser: &mut T, chunks: &[&str]) -> ToolParseResult {
86 + let mut result = ToolParseResult::default();
87 + for chunk in chunks {
88 + result.append(parser.push(chunk).unwrap());
89 + }
90 + result.append(parser.finish().unwrap());
91 + result.coalesce_calls()
92 +}
93 +
94 +/// Split text into chunks containing at most `chunk_chars` Unicode scalar
95 +/// values.
96 +pub fn split_by_chars(text: &str, chunk_chars: usize) -> Vec<&str> {
97 + let mut chunks = Vec::new();
98 + let mut start = 0;
99 + let mut count = 0;
100 +
101 + for (index, _) in text.char_indices() {
102 + if count == chunk_chars {
103 + chunks.push(&text[start..index]);
104 + start = index;
105 + count = 0;
106 + }
107 + count += 1;
108 + }
109 +
110 + if start < text.len() {
111 + chunks.push(&text[start..]);
112 + }
113 +
114 + chunks
115 +}
rust/src/tool-parser/src/tests.rs new
+97
@@ -0,0 +1,97 @@
1 +use super::{Result, Tool, ToolCallDelta, ToolParseResult, ToolParser};
2 +
3 +struct DefaultParser;
4 +
5 +impl ToolParser for DefaultParser {
6 + fn create(_tools: &[Tool]) -> Result<Box<dyn ToolParser>>
7 + where
8 + Self: Sized + 'static,
9 + {
10 + Ok(Box::new(Self))
11 + }
12 +
13 + fn push(&mut self, _chunk: &str) -> Result<ToolParseResult> {
14 + Ok(ToolParseResult::default())
15 + }
16 +}
17 +
18 +#[test]
19 +fn tool_parser_does_not_preserve_special_tokens_by_default() {
20 + let parser = DefaultParser;
21 +
22 + assert!(!parser.preserve_special_tokens());
23 +}
24 +
25 +#[test]
26 +fn default_parse_complete_delegates_through_push_and_finish() {
27 + struct StreamingParser;
28 +
29 + impl ToolParser for StreamingParser {
30 + fn create(_tools: &[Tool]) -> Result<Box<dyn ToolParser>>
31 + where
32 + Self: Sized + 'static,
33 + {
34 + Ok(Box::new(Self))
35 + }
36 +
37 + fn push(&mut self, _chunk: &str) -> Result<ToolParseResult> {
38 + Ok(ToolParseResult {
39 + normal_text: "prefix ".to_string(),
40 + calls: vec![
41 + ToolCallDelta {
42 + tool_index: 0,
43 + name: Some("weather".to_string()),
44 + arguments: "{\"location\":".to_string(),
45 + },
46 + ToolCallDelta {
47 + tool_index: 0,
48 + name: None,
49 + arguments: "\"Paris\"".to_string(),
50 + },
51 + ToolCallDelta {
52 + tool_index: 1,
53 + name: Some("time".to_string()),
54 + arguments: "{\"timezone\":".to_string(),
55 + },
56 + ],
57 + })
58 + }
59 +
60 + fn finish(&mut self) -> Result<ToolParseResult> {
61 + Ok(ToolParseResult {
62 + normal_text: "suffix".to_string(),
63 + calls: vec![
64 + ToolCallDelta {
65 + tool_index: 0,
66 + name: None,
67 + arguments: "}".to_string(),
68 + },
69 + ToolCallDelta {
70 + tool_index: 1,
71 + name: None,
72 + arguments: "\"UTC\"}".to_string(),
73 + },
74 + ],
75 + })
76 + }
77 + }
78 +
79 + let mut parser = StreamingParser;
80 + let result = parser.parse_complete("ignored").unwrap();
81 + assert_eq!(result.normal_text, "prefix suffix");
82 + assert_eq!(
83 + result.calls,
84 + vec![
85 + ToolCallDelta {
86 + tool_index: 0,
87 + name: Some("weather".to_string()),
88 + arguments: "{\"location\":\"Paris\"}".to_string(),
89 + },
90 + ToolCallDelta {
91 + tool_index: 1,
92 + name: Some("time".to_string()),
93 + arguments: "{\"timezone\":\"UTC\"}".to_string(),
94 + },
95 + ]
96 + );
97 +}
rust/src/tool-parser/src/utils.rs new
+581
@@ -0,0 +1,581 @@
1 +//! Shared helpers for tool parsers.
2 +
3 +use std::borrow::Cow;
4 +
5 +use winnow::error::{ContextError, ErrMode, ModalResult, Needed, StrContext, StrContextValue};
6 +use winnow::stream::{Offset, Partial, Stream};
7 +
8 +use super::Result;
9 +
10 +/// Return the byte length of the longest proper prefix of `token` that is also
11 +/// a suffix of `buffer`.
12 +///
13 +/// Streaming parsers use this to keep only the trailing fragment that might
14 +/// still grow into a full marker after the next decoded chunk arrives.
15 +///
16 +/// The returned length is always a valid UTF-8 boundary in `token`, so callers
17 +/// can safely slice `&token[..len]` even when markers contain non-ASCII
18 +/// characters such as DeepSeek's DSML delimiters.
19 +pub(super) fn partial_prefix_len(buffer: &str, token: &str) -> usize {
20 + let Some(first_byte) = token.as_bytes().first().copied() else {
21 + return 0;
22 + };
23 +
24 + let max_len = buffer.len().min(token.len().saturating_sub(1));
25 + let tail_start = buffer.len() - max_len;
26 + let buffer_bytes = buffer.as_bytes();
27 + let token_bytes = token.as_bytes();
28 +
29 + // Scan from the longest possible suffix to preserve overlapping prefixes.
30 + for index in tail_start..buffer.len() {
31 + if buffer_bytes[index] != first_byte {
32 + continue;
33 + }
34 +
35 + let len = buffer.len() - index;
36 + if buffer.is_char_boundary(index)
37 + && token.is_char_boundary(len)
38 + && token_bytes[..len] == buffer_bytes[index..]
39 + {
40 + return len;
41 + }
42 + }
43 +
44 + 0
45 +}
46 +
47 +/// Parse a safe text run before the next marker.
48 +///
49 +/// Returns the text length in bytes, and advances the input.
50 +pub(super) fn safe_text_len(input: &mut Partial<&str>, marker: &str) -> ModalResult<usize> {
51 + let text = **input;
52 + if text.is_empty() {
53 + return incomplete();
54 + }
55 +
56 + if let Some(start_idx) = text.find(marker) {
57 + input.next_slice(start_idx);
58 + return Ok(start_idx);
59 + }
60 +
61 + let keep_len = partial_prefix_len(text, marker);
62 + let emit_len = text.len().saturating_sub(keep_len);
63 + if emit_len == 0 {
64 + return incomplete();
65 + }
66 +
67 + input.next_slice(emit_len);
68 + Ok(emit_len)
69 +}
70 +
71 +/// Decode XML/HTML entities in XML-style parameter values.
72 +pub(super) fn xml_unescape(value: &str) -> Cow<'_, str> {
73 + if !value.as_bytes().contains(&b'&') {
74 + return Cow::Borrowed(value);
75 + }
76 +
77 + let mut output: Option<String> = None;
78 + let mut copied_len = 0;
79 + let mut rest = value;
80 +
81 + while let Some(ampersand) = rest.find('&') {
82 + let before_ampersand = &rest[..ampersand];
83 + let after_ampersand = &rest[ampersand + '&'.len_utf8()..];
84 + if let Some(semicolon) = after_ampersand.find(';') {
85 + let entity = &after_ampersand[..semicolon];
86 + if let Some(decoded) = decode_xml_entity(entity) {
87 + match &mut output {
88 + Some(output) => output.push_str(before_ampersand),
89 + None => {
90 + let mut new_output = String::with_capacity(value.len());
91 + new_output.push_str(&value[..copied_len + ampersand]);
92 + output = Some(new_output);
93 + }
94 + }
95 + let output = output.as_mut().expect("output is initialized above");
96 + output.push(decoded);
97 + let consumed_len = ampersand + '&'.len_utf8() + semicolon + ';'.len_utf8();
98 + copied_len += consumed_len;
99 + rest = &rest[consumed_len..];
100 + continue;
101 + }
102 + }
103 +
104 + if let Some(output) = &mut output {
105 + output.push_str(before_ampersand);
106 + output.push('&');
107 + }
108 + let consumed_len = ampersand + '&'.len_utf8();
109 + copied_len += consumed_len;
110 + rest = after_ampersand;
111 + }
112 +
113 + if let Some(mut output) = output {
114 + output.push_str(rest);
115 + Cow::Owned(output)
116 + } else {
117 + Cow::Borrowed(value)
118 + }
119 +}
120 +
121 +fn decode_xml_entity(entity: &str) -> Option<char> {
122 + match entity {
123 + "amp" => Some('&'),
124 + "lt" => Some('<'),
125 + "gt" => Some('>'),
126 + "quot" => Some('"'),
127 + "apos" => Some('\''),
128 + entity if entity.starts_with("#x") || entity.starts_with("#X") => {
129 + u32::from_str_radix(&entity[2..], 16).ok().and_then(char::from_u32)
130 + }
131 + entity if entity.starts_with('#') => {
132 + entity[1..].parse::<u32>().ok().and_then(char::from_u32)
133 + }
134 + _ => None,
135 + }
136 +}
137 +
138 +/// Streaming lexical state for a top-level JSON object.
139 +#[derive(Debug, Clone, Default, PartialEq, Eq)]
140 +pub(super) struct JsonObjectScanState {
141 + object_depth: usize,
142 + array_depth: usize,
143 + in_string: bool,
144 + escape: bool,
145 + phase: JsonObjectScanPhase,
146 +}
147 +
148 +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
149 +enum JsonObjectScanPhase {
150 + #[default]
151 + Initial,
152 + Scanning,
153 + Complete,
154 +}
155 +
156 +impl JsonObjectScanState {
157 + /// Returns whether the top-level JSON object has closed.
158 + pub(super) const fn complete(&self) -> bool {
159 + matches!(self.phase, JsonObjectScanPhase::Complete)
160 + }
161 +}
162 +
163 +/// Parse a raw top-level JSON object argument prefix.
164 +///
165 +/// The returned length is safe to emit as raw argument text. This scans only
166 +/// lexical boundaries from `{` through the matching `}`, preserving
167 +/// malformed-but-balanced JSON without deserializing or normalizing it.
168 +pub(super) fn take_json_object(
169 + input: &mut Partial<&str>,
170 + state: &mut JsonObjectScanState,
171 +) -> ModalResult<usize> {
172 + let text = **input;
173 + if text.is_empty() {
174 + return incomplete();
175 + }
176 + if state.complete() {
177 + return Err(json_scan_error(
178 + "JSON object argument",
179 + StrContextValue::Description("active JSON object scan"),
180 + ));
181 + }
182 +
183 + let bytes = text.as_bytes();
184 + let just_started = matches!(state.phase, JsonObjectScanPhase::Initial);
185 + if just_started {
186 + if bytes[0] != b'{' {
187 + return Err(json_scan_error(
188 + "JSON object argument",
189 + StrContextValue::CharLiteral('{'),
190 + ));
191 + }
192 + state.phase = JsonObjectScanPhase::Scanning;
193 + state.object_depth = 1;
194 + }
195 +
196 + let mut index = usize::from(just_started);
197 +
198 + while index < bytes.len() {
199 + let byte = bytes[index];
200 + index += 1;
201 +
202 + if state.in_string {
203 + if state.escape {
204 + state.escape = false;
205 + } else if byte == b'\\' {
206 + state.escape = true;
207 + } else if byte == b'"' {
208 + state.in_string = false;
209 + }
210 + continue;
211 + }
212 +
213 + match byte {
214 + b'"' => state.in_string = true,
215 + b'{' => state.object_depth += 1,
216 + b'}' => {
217 + state.object_depth = state.object_depth.checked_sub(1).ok_or_else(|| {
218 + json_scan_error(
219 + "JSON object argument",
220 + StrContextValue::Description("balanced object braces"),
221 + )
222 + })?;
223 + if state.object_depth == 0 && state.array_depth == 0 {
224 + state.phase = JsonObjectScanPhase::Complete;
225 + input.next_slice(index);
226 + return Ok(index);
227 + }
228 + if state.object_depth == 0 {
229 + return Err(json_scan_error(
230 + "JSON object argument",
231 + StrContextValue::Description(
232 + "nested arrays to close before the top-level object",
233 + ),
234 + ));
235 + }
236 + }
237 + b'[' => state.array_depth += 1,
238 + b']' => {
239 + state.array_depth = state.array_depth.checked_sub(1).ok_or_else(|| {
240 + json_scan_error(
241 + "JSON object argument",
242 + StrContextValue::Description("balanced array brackets"),
243 + )
244 + })?;
245 + }
246 + _ => {}
247 + }
248 + }
249 +
250 + input.next_slice(text.len());
251 + Ok(text.len())
252 +}
253 +
254 +/// Parse a JSON string literal.
255 +pub(super) fn json_str(input: &mut Partial<&str>) -> ModalResult<String> {
256 + let text = **input;
257 + if text.is_empty() {
258 + return incomplete();
259 + }
260 +
261 + let bytes = text.as_bytes();
262 + if bytes[0] != b'"' {
263 + return Err(json_scan_error(
264 + "JSON string",
265 + StrContextValue::CharLiteral('"'),
266 + ));
267 + }
268 +
269 + let mut escape = false;
270 + let mut index = 1;
271 + while index < bytes.len() {
272 + let byte = bytes[index];
273 + index += 1;
274 +
275 + if escape {
276 + escape = false;
277 + continue;
278 + }
279 +
280 + match byte {
281 + b'\\' => escape = true,
282 + b'"' => {
283 + let raw = &text[..index];
284 + let value = serde_json::from_str::<String>(raw).map_err(|_| {
285 + json_scan_error(
286 + "JSON string",
287 + StrContextValue::Description("valid JSON string"),
288 + )
289 + })?;
290 + input.next_slice(index);
291 + return Ok(value);
292 + }
293 + _ => {}
294 + }
295 + }
296 +
297 + incomplete()
298 +}
299 +
300 +fn json_scan_error(label: &'static str, expected: StrContextValue) -> ErrMode<ContextError> {
301 + let mut error = ContextError::new();
302 + error.push(StrContext::Label(label));
303 + error.push(StrContext::Expected(expected));
304 + ErrMode::Cut(error)
305 +}
306 +
307 +/// Parse one event from a buffered streaming input.
308 +///
309 +/// Returns:
310 +/// - `Ok(Some((event, consumed_len)))` if an event was successfully parsed, along with the number
311 +/// of bytes consumed from the buffer.
312 +/// - `Ok(None)` if the buffer does not contain a full event yet, and more data is needed.
313 +/// - `Err` if a parsing error occurred.
314 +pub(super) fn parse_buffered_event<E>(
315 + buffer: &str,
316 + parse: impl FnOnce(&mut Partial<&str>) -> ModalResult<E>,
317 +) -> Result<Option<(E, usize)>> {
318 + let mut input = Partial::new(buffer);
319 + let checkpoint = input.checkpoint();
320 + let event = match parse(&mut input) {
321 + Ok(event) => event,
322 + Err(ErrMode::Incomplete(_)) => return Ok(None),
323 + Err(ErrMode::Backtrack(e) | ErrMode::Cut(e)) => {
324 + // TODO: enrich context for error reporting
325 + return Err(parsing_failed!("{}", e));
326 + }
327 + };
328 + let consumed_len = input.offset_from(&checkpoint);
329 + if consumed_len == 0 {
330 + return Ok(None);
331 + }
332 +
333 + Ok(Some((event, consumed_len)))
334 +}
335 +
336 +/// Returns an error indicating that we need more data to continue parsing.
337 +pub(super) fn incomplete<T>() -> ModalResult<T> {
338 + Err(ErrMode::Incomplete(Needed::Unknown))
339 +}
340 +
341 +#[cfg(test)]
342 +mod tests {
343 + use std::borrow::Cow;
344 +
345 + use expect_test::expect;
346 + use winnow::error::ErrMode;
347 + use winnow::stream::{Offset, Partial, Stream};
348 +
349 + use super::{
350 + JsonObjectScanState, json_str, partial_prefix_len, safe_text_len, take_json_object,
351 + xml_unescape,
352 + };
353 +
354 + #[test]
355 + fn partial_prefix_len_handles_ascii_markers() {
356 + assert_eq!(
357 + partial_prefix_len("hello<|tool", "<|tool_call>"),
358 + "<|tool".len()
359 + );
360 + assert_eq!(partial_prefix_len("hello world", "<|tool_call>"), 0);
361 + }
362 +
363 + #[test]
364 + fn partial_prefix_len_prefers_longest_overlapping_prefix() {
365 + assert_eq!(partial_prefix_len("chunk ending in aba", "ababa"), 3);
366 + }
367 +
368 + #[test]
369 + fn partial_prefix_len_handles_unicode_markers() {
370 + let token = "<｜DSML｜function_calls>";
371 + assert_eq!(
372 + partial_prefix_len("prefix <｜DSML｜fun", token),
373 + "<｜DSML｜fun".len()
374 + );
375 + assert_eq!(partial_prefix_len("prefix <｜DSML", token), "<｜DSML".len());
376 + }
377 +
378 + #[test]
379 + fn safe_text_len_stops_before_marker() {
380 + let mut input = Partial::new("hello<tool_call>");
381 + let checkpoint = input.checkpoint();
382 +
383 + let len = safe_text_len(&mut input, "<tool_call>").unwrap();
384 +
385 + assert_eq!(len, "hello".len());
386 + assert_eq!(input.offset_from(&checkpoint), "hello".len());
387 + }
388 +
389 + #[test]
390 + fn safe_text_len_holds_back_partial_marker() {
391 + let mut input = Partial::new("hello<tool");
392 + let checkpoint = input.checkpoint();
393 +
394 + let len = safe_text_len(&mut input, "<tool_call>").unwrap();
395 +
396 + assert_eq!(len, "hello".len());
397 + assert_eq!(input.offset_from(&checkpoint), "hello".len());
398 + }
399 +
400 + #[test]
401 + fn safe_text_len_reports_incomplete_for_only_partial_marker() {
402 + let mut input = Partial::new("<tool");
403 +
404 + let error = safe_text_len(&mut input, "<tool_call>").unwrap_err();
405 +
406 + assert!(matches!(error, ErrMode::Incomplete(_)));
407 + }
408 +
409 + #[test]
410 + fn xml_unescape_decodes_common_entities() {
411 + assert_eq!(
412 + xml_unescape("&lt;tag attr=&quot;value&quot;&gt;Tom &amp; Jerry&apos;s&lt;/tag&gt;"),
413 + r#"<tag attr="value">Tom & Jerry's</tag>"#
414 + );
415 + }
416 +
417 + #[test]
418 + fn xml_unescape_decodes_numeric_entities() {
419 + assert_eq!(xml_unescape("&#60;tag&#x3E;&#x1F600;"), "<tag>😀");
420 + }
421 +
422 + #[test]
423 + fn xml_unescape_preserves_unknown_and_incomplete_entities() {
424 + let output = xml_unescape("Tom & Jerry &unknown; &amp");
425 +
426 + assert!(matches!(output, Cow::Borrowed(_)));
427 + assert_eq!(output, "Tom & Jerry &unknown; &amp");
428 + }
429 +
430 + #[test]
431 + fn xml_unescape_borrows_when_no_entity_is_present() {
432 + let input = "plain text";
433 + let output = xml_unescape(input);
434 +
435 + assert!(matches!(output, Cow::Borrowed(_)));
436 + assert_eq!(output, input);
437 + }
438 +
439 + #[test]
440 + fn take_json_object_consumes_simple_object() {
441 + let mut state = JsonObjectScanState::default();
442 + let buffer = r#"{"location":"Paris"}<end>"#;
443 + let mut input = Partial::new(buffer);
444 + let checkpoint = input.checkpoint();
445 +
446 + let len = take_json_object(&mut input, &mut state).unwrap();
447 +
448 + assert_eq!(len, r#"{"location":"Paris"}"#.len());
449 + assert_eq!(input.offset_from(&checkpoint), len);
450 + assert!(state.complete());
451 + }
452 +
453 + #[test]
454 + fn take_json_object_tracks_nested_values_and_strings() {
455 + let mut state = JsonObjectScanState::default();
456 + let arguments = r#"{"nested":{"items":[{"text":"} <|tool_call_end|> \" \\"}]}}"#;
457 + let buffer = format!("{arguments}<end>");
458 + let mut input = Partial::new(buffer.as_str());
459 +
460 + let len = take_json_object(&mut input, &mut state).unwrap();
461 +
462 + assert_eq!(len, arguments.len());
463 + assert!(state.complete());
464 + }
465 +
466 + #[test]
467 + fn take_json_object_rejects_leading_whitespace() {
468 + let mut state = JsonObjectScanState::default();
469 + let mut input = Partial::new(" {\"x\":1}");
470 +
471 + let error = take_json_object(&mut input, &mut state).unwrap_err();
472 +
473 + let ErrMode::Cut(error) = error else {
474 + panic!("expected cut error");
475 + };
476 + expect![[r#"
477 + invalid JSON object argument
478 + expected `{`"#]]
479 + .assert_eq(&error.to_string());
480 + }
481 +
482 + #[test]
483 + fn take_json_object_leaves_trailing_whitespace_to_caller() {
484 + let mut state = JsonObjectScanState::default();
485 + let mut input = Partial::new("{\"x\":1}\n<end>");
486 + let checkpoint = input.checkpoint();
487 +
488 + let len = take_json_object(&mut input, &mut state).unwrap();
489 +
490 + assert_eq!(len, "{\"x\":1}".len());
491 + assert_eq!(input.offset_from(&checkpoint), len);
492 + assert!(state.complete());
493 + }
494 +
495 + #[test]
496 + fn take_json_object_continues_across_chunks() {
497 + let mut state = JsonObjectScanState::default();
498 + let chunks = [
499 + r#"{"text":"literal "#,
500 + r#"<|tool_call_end|>"#,
501 + r#" inside"}<end>"#,
502 + ];
503 + let mut collected = String::new();
504 +
505 + for chunk in chunks {
506 + let mut input = Partial::new(chunk);
507 + let len = take_json_object(&mut input, &mut state).unwrap();
508 + collected.push_str(&chunk[..len]);
509 + }
510 +
511 + assert_eq!(collected, r#"{"text":"literal <|tool_call_end|> inside"}"#);
512 + assert!(state.complete());
513 + }
514 +
515 + #[test]
516 + fn take_json_object_rejects_non_object_top_level() {
517 + let mut state = JsonObjectScanState::default();
518 + let mut input = Partial::new(r#"[{"x":1}]"#);
519 +
520 + let error = take_json_object(&mut input, &mut state).unwrap_err();
521 +
522 + let ErrMode::Cut(error) = error else {
523 + panic!("expected cut error");
524 + };
525 + expect![[r#"
526 + invalid JSON object argument
527 + expected `{`"#]]
528 + .assert_eq(&error.to_string());
529 + }
530 +
531 + #[test]
532 + fn take_json_object_reports_unbalanced_array() {
533 + let mut state = JsonObjectScanState::default();
534 + let mut input = Partial::new(r#"{"x":]}"#);
535 +
536 + let error = take_json_object(&mut input, &mut state).unwrap_err();
537 +
538 + let ErrMode::Cut(error) = error else {
539 + panic!("expected cut error");
540 + };
541 + expect![[r#"
542 + invalid JSON object argument
543 + expected balanced array brackets"#]]
544 + .assert_eq(&error.to_string());
545 + }
546 +
547 + #[test]
548 + fn take_json_object_reports_top_level_close_before_nested_array() {
549 + let mut state = JsonObjectScanState::default();
550 + let mut input = Partial::new(r#"{"x":[}"#);
551 +
552 + let error = take_json_object(&mut input, &mut state).unwrap_err();
553 +
554 + let ErrMode::Cut(error) = error else {
555 + panic!("expected cut error");
556 + };
557 + expect![[r#"
558 + invalid JSON object argument
559 + expected nested arrays to close before the top-level object"#]]
560 + .assert_eq(&error.to_string());
561 + }
562 +
563 + #[test]
564 + fn json_str_decodes_escaped_content() {
565 + let mut input = Partial::new(r#""say_\"hi\u0021" rest"#);
566 +
567 + let value = json_str(&mut input).unwrap();
568 +
569 + assert_eq!(value, "say_\"hi!");
570 + assert_eq!(*input, " rest");
571 + }
572 +
573 + #[test]
574 + fn json_str_reports_incomplete_escaped_string() {
575 + let mut input = Partial::new(r#""say_\"#);
576 +
577 + let error = json_str(&mut input).unwrap_err();
578 +
579 + assert!(matches!(error, ErrMode::Incomplete(_)));
580 + }
581 +}
tools/pre_commit/rust-check.sh new
+47
@@ -0,0 +1,47 @@
1 +#!/usr/bin/env bash
2 +# Wrapper for the rust-* pre-commit hooks.
3 +#
4 +# Skips (with a warning) when `cargo` or the requested cargo subcommand is
5 +# not installed, so contributors who don't touch the Rust code aren't forced
6 +# to install the Rust toolchain (or niche cargo extensions like cargo-sort
7 +# / cargo-autoinherit). Buildkite CI covers the rust hooks regardless.
8 +#
9 +# Usage: tools/pre_commit/rust-check.sh <cargo-subcommand> [extra cargo args...]
10 +
11 +set -euo pipefail
12 +
13 +# Pre-commit captures stdout/stderr and only replays on failure. Try to write
14 +# to /dev/tty so the warning is visible during a normal `git commit` even
15 +# though we exit 0; fall back to stderr where there's no controlling tty
16 +# (e.g. CI).
17 +#
18 +# The leading newline pushes the warning off pre-commit's dot-leader line so
19 +# the message doesn't mash into "Rust - ... ........WARNING:". The hook's
20 +# "Passed" still lands on its own line just below the warning.
21 +warn() {
22 + { printf '\n%s\n' "$*" >/dev/tty; } 2>/dev/null || printf '\n%s\n' "$*" >&2
23 +}
24 +
25 +subcommand="$1"
26 +shift
27 +
28 +if ! command -v cargo >/dev/null 2>&1; then
29 + warn "WARNING: 'cargo' not found in PATH; skipping rust pre-commit hook (cargo ${subcommand}).
30 + Install the Rust toolchain via https://rustup.rs/ if you need to run rust hooks locally."
31 + exit 0
32 +fi
33 +
34 +# Cargo subcommands resolve to a `cargo-<name>` binary on PATH. Check up-front
35 +# so a missing helper produces a friendly skip instead of a cargo error.
36 +if ! command -v "cargo-${subcommand}" >/dev/null 2>&1; then
37 + case "${subcommand}" in
38 + fmt) install_hint="rustup component add rustfmt" ;;
39 + *) install_hint="cargo install cargo-${subcommand}" ;;
40 + esac
41 + warn "WARNING: 'cargo ${subcommand}' is not installed; skipping rust pre-commit hook.
42 + Install it with: ${install_hint}"
43 + exit 0
44 +fi
45 +
46 +cd "$(git rev-parse --show-toplevel)/rust"
47 +exec cargo "${subcommand}" "$@"