@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 43914dd743

Fork 0

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

[Rust Frontend] Add Python bridge for Rust tool parsers (#44624)

Signed-off-by: Bugen Zhao <i@bugenzhao.com>

Bugen Zhao committed Jun 11, 2026 at 12:51 UTC 43914dd743ab0500abcd69fe072e02465c944dcf

16 files changed +1302 -38

.buildkite/scripts/run-rust-frontend-cargo-ci.sh

+31

@@ -110,6 +110,36 @@ install_uv() {

110 | env UV_INSTALL_DIR="$CARGO_HOME/bin" sh

111 }

112

113 +setup_pyo3_python() {

114 + local python_version="${PYO3_PYTHON_VERSION:-3.12}"

115 +

116 + log_section "Installing Python ${python_version} for PyO3 tests"

117 + uv python install "$python_version"

118 + PYO3_PYTHON="$(uv python find \

119 + --managed-python \

120 + --no-project \

121 + --resolve-links \

122 + "$python_version")"

123 + export PYO3_PYTHON

124 +

125 + local python_libdir

126 + python_libdir="$("$PYO3_PYTHON" - <<'PY'

127 +import pathlib

128 +import sysconfig

129 +

130 +libdir = pathlib.Path(sysconfig.get_config_var("LIBDIR"))

131 +ldlibrary = sysconfig.get_config_var("LDLIBRARY")

132 +assert sysconfig.get_config_var("Py_ENABLE_SHARED") == 1

133 +assert ldlibrary

134 +assert (libdir / ldlibrary).exists(), libdir / ldlibrary

135 +print(libdir)

136 +PY

137 +)"

138 +

139 + export LD_LIBRARY_PATH="${python_libdir}:${LD_LIBRARY_PATH:-}"

140 + export LIBRARY_PATH="${python_libdir}:${LIBRARY_PATH:-}"

141 +}

142 +

143 run_style_clippy() {

144 install_cargo_sort

145

@@ -132,6 +162,7 @@ run_style_clippy() {

162

163 run_tests() {

164 install_uv

165 + setup_pyo3_python

166 install_cargo_nextest

167

168 log_section "Running cargo nextest"

.buildkite/test_areas/misc.yaml

+2 -2

@@ -301,9 +301,9 @@ steps:

301 - tests/multimodal

302 - tests/renderers

303 - tests/standalone_tests/lazy_imports.py

304 - - tests/tokenizers_

304 - tests/reasoning

305 - tests/tool_parsers

306 + - tests/tokenizers_

307 - tests/parser

308 - tests/transformers_utils

309 - tests/config

@@ -317,9 +317,9 @@ steps:

317 - pytest -v -s test_ray_env.py

318 - pytest -v -s -m 'cpu_test' multimodal

319 - pytest -v -s renderers

320 - - pytest -v -s tokenizers_

320 - pytest -v -s reasoning --ignore=reasoning/test_seedoss_reasoning_parser.py --ignore=reasoning/test_glm4_moe_reasoning_parser.py

321 - pytest -v -s tool_parsers

322 + - pytest -v -s tokenizers_

323 - pytest -v -s parser

324 - pytest -v -s transformers_utils

325 - pytest -v -s config

build_rust.sh

+1 -1

@@ -1,5 +1,5 @@

1 #!/bin/bash

2 -# Build the vllm-rs Rust frontend binary.

2 +# Build vLLM Rust artifacts and install them into the vllm package.

3 # Usage: ./build_rust.sh [--debug]

4 #

5 # By default builds in release mode. Pass --debug for faster compile times

docker/Dockerfile

+8 -5

@@ -281,7 +281,8 @@ COPY requirements/build/rust.txt requirements/build/rust.txt

281 RUN --mount=type=cache,target=/opt/uv/cache \

282 uv pip install --python /opt/venv/bin/python3 -r requirements/build/rust.txt

283

284 -# Copy only the Rust build inputs. The binary is the sole artifact we need.

284 +# Copy only the Rust build inputs; build_rust.sh publishes artifacts needed

285 +# by the wheel build stage.

286 COPY rust rust

287 COPY rust-toolchain.toml rust-toolchain.toml

288 COPY tools/build_rust.py tools/build_rust.py

@@ -291,8 +292,9 @@ COPY build_rust.sh build_rust.sh

292 # (rustc spawns enough concurrent processes to hit RLIMIT_NOFILE otherwise).

293 ENV CARGO_BUILD_JOBS=4

294

294 -# Build the release binary. Cache cargo registry/git, but not target/, because

295 -# stale target metadata can outlive source updates across BuildKit cache reuse.

295 +# Build the release artifacts. Cache cargo registry/git, but not target/,

296 +# because stale target metadata can outlive source updates across BuildKit

297 +# cache reuse.

298 RUN --mount=type=cache,target=/root/.cargo/registry \

299 --mount=type=cache,target=/root/.cargo/git \

300 bash build_rust.sh

@@ -503,9 +505,10 @@ WORKDIR /workspace

505 COPY --from=csrc-build /workspace/dist /precompiled-wheels

506 COPY . .

507

506 -# Drop the pre-built rust frontend binary into the source tree. setup.py

507 -# detects it and ships it as-is, skipping the local cargo build.

508 +# Drop the pre-built Rust artifacts into the source tree. setup.py detects

509 +# them and ships them as-is, skipping the local Rust build.

510 COPY --from=rust-build /workspace/vllm/vllm-rs vllm/vllm-rs

511 +COPY --from=rust-build /workspace/vllm/_rust_*.so vllm/

512

513 ARG GIT_REPO_CHECK=0

514 RUN --mount=type=bind,source=.git,target=.git \

docker/Dockerfile.cpu

+8 -5

@@ -104,7 +104,8 @@ WORKDIR /workspace

104 COPY requirements/build/rust.txt requirements/build/rust.txt

105 RUN python3 -m pip install --no-cache-dir -r requirements/build/rust.txt

106

107 -# Copy only the Rust build inputs. The binary is the sole artifact we need.

107 +# Copy only the Rust build inputs; build_rust.sh publishes artifacts needed

108 +# by the wheel build stage.

109 COPY rust rust

110 COPY rust-toolchain.toml rust-toolchain.toml

111 COPY tools/build_rust.py tools/build_rust.py

@@ -114,8 +115,9 @@ COPY build_rust.sh build_rust.sh

115 # (rustc spawns enough concurrent processes to hit RLIMIT_NOFILE otherwise).

116 ENV CARGO_BUILD_JOBS=4

117

117 -# Build the release binary. Cache cargo registry/git, but not target/, because

118 -# stale target metadata can outlive source updates across BuildKit cache reuse.

118 +# Build the release artifacts. Cache cargo registry/git, but not target/,

119 +# because stale target metadata can outlive source updates across BuildKit

120 +# cache reuse.

121 RUN --mount=type=cache,target=/root/.cargo/registry,sharing=locked \

122 --mount=type=cache,target=/root/.cargo/git,sharing=locked \

123 bash build_rust.sh

@@ -151,9 +153,10 @@ RUN --mount=type=cache,target=/root/.cache/uv \

153

154 COPY . .

155

154 -# Drop the pre-built rust frontend binary into the source tree. setup.py

155 -# detects it and ships it as-is, skipping the local cargo build.

156 +# Drop the pre-built Rust artifacts into the source tree. setup.py detects

157 +# them and ships them as-is, skipping the local Rust build.

158 COPY --from=rust-build /workspace/vllm/vllm-rs vllm/vllm-rs

159 +COPY --from=rust-build /workspace/vllm/_rust_*.so vllm/

160

161 RUN if [ "$GIT_REPO_CHECK" != 0 ]; then bash tools/check_repo.sh ; fi

162

docker/Dockerfile.nightly_torch

+5 -3

@@ -113,7 +113,8 @@ WORKDIR /workspace

113 COPY requirements/build/rust.txt requirements/build/rust.txt

114 RUN python3 -m pip install --no-cache-dir -r requirements/build/rust.txt

115

116 -# Copy only the Rust build inputs. The binary is the sole artifact we need.

116 +# Copy only the Rust build inputs; build_rust.sh publishes artifacts needed

117 +# by the wheel build stage.

118 COPY rust rust

119 COPY rust-toolchain.toml rust-toolchain.toml

120 COPY tools/build_rust.py tools/build_rust.py

@@ -138,9 +139,10 @@ ENV UV_HTTP_TIMEOUT=500

139

140 COPY . .

141

141 -# Drop the pre-built rust frontend binary into the source tree. setup.py

142 -# detects it and ships it as-is, skipping the local cargo build.

142 +# Drop the pre-built Rust artifacts into the source tree. setup.py detects

143 +# them and ships them as-is, skipping the local Rust build.

144 COPY --from=rust-build /workspace/vllm/vllm-rs vllm/vllm-rs

145 +COPY --from=rust-build /workspace/vllm/_rust_*.so vllm/

146

147 RUN python3 use_existing_torch.py

148

docker/Dockerfile.rocm

+6 -4

@@ -208,9 +208,10 @@ ENV VLLM_TARGET_DEVICE=rocm

208

209 COPY --from=csrc-build ${COMMON_WORKDIR}/vllm/dist /precompiled-wheels

210

211 -# Drop the pre-built rust frontend binary into the source tree. setup.py

212 -# detects it and ships it as-is, skipping the local cargo build.

211 +# Drop the pre-built Rust artifacts into the source tree. setup.py detects

212 +# them and ships them as-is, skipping the local Rust build.

213 COPY --from=rust-build ${COMMON_WORKDIR}/vllm/vllm/vllm-rs ${COMMON_WORKDIR}/vllm/vllm/vllm-rs

214 +COPY --from=rust-build ${COMMON_WORKDIR}/vllm/vllm/_rust_*.so ${COMMON_WORKDIR}/vllm/vllm/

215

216 RUN --mount=type=cache,id=vllm-rocm-uv,target=/root/.cache/uv \

217 cd vllm \

@@ -417,9 +418,10 @@ FROM fetch_vllm AS build_vllm_wheel_release

418

419 ARG COMMON_WORKDIR

420

420 -# Drop the pre-built rust frontend binary into the source tree. setup.py

421 -# detects it and ships it as-is, skipping the local cargo build.

421 +# Drop the pre-built Rust artifacts into the source tree. setup.py detects

422 +# them and ships them as-is, skipping the local Rust build.

423 COPY --from=rust-build ${COMMON_WORKDIR}/vllm/vllm/vllm-rs ${COMMON_WORKDIR}/vllm/vllm/vllm-rs

424 +COPY --from=rust-build ${COMMON_WORKDIR}/vllm/vllm/_rust_*.so ${COMMON_WORKDIR}/vllm/vllm/

425

426 # Create /install directory for custom wheels

427 RUN mkdir -p /install

docker/Dockerfile.xpu

+5 -3

@@ -17,7 +17,8 @@ WORKDIR /workspace

17 COPY requirements/build/rust.txt requirements/build/rust.txt

18 RUN python3 -m pip install --no-cache-dir -r requirements/build/rust.txt

20 -# Copy only the Rust build inputs. The binary is the sole artifact we need.

20 +# Copy only the Rust build inputs; build_rust.sh publishes artifacts needed

21 +# by the wheel build stage.

22 COPY rust rust

23 COPY rust-toolchain.toml rust-toolchain.toml

24 COPY tools/build_rust.py tools/build_rust.py

@@ -212,9 +213,10 @@ RUN --mount=type=cache,target=/root/.cache/uv \

213 # don't invalidate heavy dependency and UCX/NIXL layers.

214 COPY . .

215

215 -# Drop the pre-built rust frontend binary into the source tree. setup.py

216 -# detects it and ships it as-is, skipping the local cargo build.

216 +# Drop the pre-built Rust artifacts into the source tree. setup.py detects

217 +# them and ships them as-is, skipping the local Rust build.

218 COPY --from=rust-build /workspace/vllm/vllm-rs vllm/vllm-rs

219 +COPY --from=rust-build /workspace/vllm/_rust_*.so vllm/

220

221 ARG GIT_REPO_CHECK=0

222 RUN --mount=type=bind,source=.git,target=.git \

rust/Cargo.lock

+86

@@ -3458,6 +3458,75 @@ version = "0.1.29"

3458 source = "registry+https://github.com/rust-lang/crates.io-index"

3459 checksum = "e0c5ccf5294c6ccd63a74f1565028353830a9c2f5eb0c682c355c471726a6e3f"

3460

3461 +[[package]]

3462 +name = "pyo3"

3463 +version = "0.28.3"

3464 +source = "registry+https://github.com/rust-lang/crates.io-index"

3465 +checksum = "91fd8e38a3b50ed1167fb981cd6fd60147e091784c427b8f7183a7ee32c31c12"

3466 +dependencies = [

3467 + "libc",

3468 + "once_cell",

3469 + "portable-atomic",

3470 + "pyo3-build-config",

3471 + "pyo3-ffi",

3472 + "pyo3-macros",

3473 +]

3474 +

3475 +[[package]]

3476 +name = "pyo3-build-config"

3477 +version = "0.28.3"

3478 +source = "registry+https://github.com/rust-lang/crates.io-index"

3479 +checksum = "e368e7ddfdeb98c9bca7f8383be1648fd84ab466bf2bc015e94008db6d35611e"

3480 +dependencies = [

3481 + "target-lexicon",

3482 +]

3483 +

3484 +[[package]]

3485 +name = "pyo3-ffi"

3486 +version = "0.28.3"

3487 +source = "registry+https://github.com/rust-lang/crates.io-index"

3488 +checksum = "7f29e10af80b1f7ccaf7f69eace800a03ecd13e883acfacc1e5d0988605f651e"

3489 +dependencies = [

3490 + "libc",

3491 + "pyo3-build-config",

3492 +]

3493 +

3494 +[[package]]

3495 +name = "pyo3-macros"

3496 +version = "0.28.3"

3497 +source = "registry+https://github.com/rust-lang/crates.io-index"

3498 +checksum = "df6e520eff47c45997d2fc7dd8214b25dd1310918bbb2642156ef66a67f29813"

3499 +dependencies = [

3500 + "proc-macro2",

3501 + "pyo3-macros-backend",

3502 + "quote",

3503 + "syn 2.0.117",

3504 +]

3505 +

3506 +[[package]]

3507 +name = "pyo3-macros-backend"

3508 +version = "0.28.3"

3509 +source = "registry+https://github.com/rust-lang/crates.io-index"

3510 +checksum = "c4cdc218d835738f81c2338f822078af45b4afdf8b2e33cbb5916f108b813acb"

3511 +dependencies = [

3512 + "heck",

3513 + "proc-macro2",

3514 + "pyo3-build-config",

3515 + "quote",

3516 + "syn 2.0.117",

3517 +]

3518 +

3519 +[[package]]

3520 +name = "pythonize"

3521 +version = "0.28.0"

3522 +source = "registry+https://github.com/rust-lang/crates.io-index"

3523 +checksum = "0b79f670c9626c8b651c0581011b57b6ba6970bb69faf01a7c4c0cfc81c43f95"

3524 +dependencies = [

3525 + "pyo3",

3526 + "serde",

3527 + "serde_json",

3528 +]

3529 +

3530 [[package]]

3531 name = "qoi"

3532 version = "0.4.1"

@@ -4669,6 +4738,12 @@ dependencies = [

4738 "libc",

4739 ]

4740

4741 +[[package]]

4742 +name = "target-lexicon"

4743 +version = "0.13.5"

4744 +source = "registry+https://github.com/rust-lang/crates.io-index"

4745 +checksum = "adb6935a6f5c20170eeceb1a3835a49e12e19d792f6dd344ccc76a985ca5a6ca"

4746 +

4747 [[package]]

4748 name = "task-local"

4749 version = "0.1.1"

@@ -5904,6 +5979,17 @@ dependencies = [

5979 "winnow",

5980 ]

5981

5982 +[[package]]

5983 +name = "vllm-tool-parser-py"

5984 +version = "0.1.0"

5985 +dependencies = [

5986 + "pyo3",

5987 + "pythonize",

5988 + "serde_json",

5989 + "thiserror-ext",

5990 + "vllm-tool-parser",

5991 +]

5992 +

5993 [[package]]

5994 name = "walkdir"

5995 version = "2.5.0"

rust/Cargo.toml

@@ -12,6 +12,7 @@ members = [

12 "src/text",

13 "src/tokenizer",

14 "src/tool-parser",

15 + "src/tool-parser/python",

16 ]

17 resolver = "3"

@@ -60,6 +61,8 @@ prometheus-client = "0.24.0"

61 prometheus-client-derive-encode = "0.5.0"

62 prost = "0.14.3"

63 prost-types = "0.14.3"

64 +pyo3 = "0.28.3"

65 +pythonize = "0.28.0"

66 rand = "0.9.2"

67 reasoning-parser = "1.2.2"

68 reqwest = { version = "0.12.8", default-features = false, features = ["rustls-tls"] }

rust/src/tool-parser/python/Cargo.toml new

+19

@@ -0,0 +1,19 @@

1 +[package]

2 +name = "vllm-tool-parser-py"

3 +version.workspace = true

4 +edition.workspace = true

5 +license.workspace = true

6 +

7 +[lib]

8 +name = "_rust_tool_parser"

9 +crate-type = ["cdylib", "rlib"]

10 +

11 +[dependencies]

12 +pyo3.workspace = true

13 +pythonize = { workspace = true, features = ["serde_json"] }

14 +serde_json.workspace = true

15 +thiserror-ext.workspace = true

16 +vllm-tool-parser.workspace = true

17 +

18 +[lints]

19 +workspace = true

rust/src/tool-parser/python/src/lib.rs new

+392

@@ -0,0 +1,392 @@

1 +//! Thin PyO3 bindings for `vllm_tool_parser`.

2 +//!

3 +//! This crate exposes the Rust tool parser trait and data shapes to Python

4 +//! while keeping parser state, grammar, and schema-aware argument conversion in

5 +//! Rust. Python callers should use this module as a typed bridge and keep any

6 +//! vLLM protocol adaptation outside the binding.

7 +

8 +use pyo3::exceptions::PyValueError;

9 +use pyo3::prelude::*;

10 +use pyo3::types::{PyAny, PyModule};

11 +use pythonize::{depythonize, pythonize};

12 +use serde_json::Value;

13 +use thiserror_ext::AsReport as _;

14 +use vllm_tool_parser::{Tool, ToolCallDelta, ToolParser, ToolParserOutput};

15 +

16 +macro_rules! tool_parser_factory {

17 + ($($parser:ident),+ $(,)?) => {

18 + fn create_tool_parser(

19 + name: &str,

20 + tools: &[Tool],

21 + ) -> PyResult<Box<dyn ToolParser>> {

22 + match name {

23 + $(

24 + stringify!($parser) => {

25 + <vllm_tool_parser::$parser as ToolParser>::create(tools)

26 + }

27 + )+

28 + _ => {

29 + return Err(PyValueError::new_err(format!(

30 + "unsupported tool parser `{name}`"

31 + )));

32 + }

33 + }

34 + .map_err(|error| PyValueError::new_err(error.to_report_string()))

35 + }

36 + };

37 +}

38 +

39 +// Export a tool parser to Python by registering it here.

40 +tool_parser_factory! {

41 + // Below are the parsers just for testing purposes on Python side.

42 + DeepSeekV4ToolParser,

43 + KimiK2ToolParser,

44 +}

45 +

46 +#[pyclass(name = "Tool", module = "vllm._rust_tool_parser", skip_from_py_object)]

47 +#[derive(Clone)]

48 +struct PyTool(Tool);

49 +

50 +#[pymethods]

51 +impl PyTool {

52 + #[new]

53 + #[pyo3(signature = (name, description, parameters, strict=None))]

54 + fn new(

55 + name: String,

56 + description: Option<String>,

57 + parameters: &Bound<'_, PyAny>,

58 + strict: Option<bool>,

59 + ) -> PyResult<Self> {

60 + let parameters = depythonize::<Value>(parameters).map_err(|error| {

61 + PyValueError::new_err(format!(

62 + "failed to convert tool parameters from Python to JSON: {error}"

63 + ))

64 + })?;

65 + Ok(Self(Tool {

66 + name,

67 + description,

68 + parameters,

69 + strict,

70 + }))

71 + }

72 +

73 + #[getter]

74 + fn name(&self) -> &str {

75 + &self.0.name

76 + }

77 +

78 + #[getter]

79 + fn description(&self) -> Option<&str> {

80 + self.0.description.as_deref()

81 + }

82 +

83 + #[getter]

84 + fn parameters(&self, py: Python<'_>) -> PyResult<Py<PyAny>> {

85 + pythonize(py, &self.0.parameters).map(Bound::unbind).map_err(|error| {

86 + PyValueError::new_err(format!(

87 + "failed to convert tool parameters from JSON to Python: {error}"

88 + ))

89 + })

90 + }

91 +

92 + #[getter]

93 + fn strict(&self) -> Option<bool> {

94 + self.0.strict

95 + }

96 +}

97 +

98 +#[pyclass(

99 + name = "ToolCallDelta",

100 + module = "vllm._rust_tool_parser",

101 + skip_from_py_object

102 +)]

103 +#[derive(Clone)]

104 +struct PyToolCallDelta(ToolCallDelta);

105 +

106 +#[pymethods]

107 +impl PyToolCallDelta {

108 + #[new]

109 + #[pyo3(signature = (tool_index, name, arguments))]

110 + fn new(tool_index: usize, name: Option<String>, arguments: String) -> Self {

111 + Self(ToolCallDelta {

112 + tool_index,

113 + name,

114 + arguments,

115 + })

116 + }

117 +

118 + #[getter]

119 + fn tool_index(&self) -> usize {

120 + self.0.tool_index

121 + }

122 +

123 + #[getter]

124 + fn name(&self) -> Option<&str> {

125 + self.0.name.as_deref()

126 + }

127 +

128 + #[getter]

129 + fn arguments(&self) -> &str {

130 + &self.0.arguments

131 + }

132 +}

133 +

134 +#[pyclass(

135 + name = "ToolParserOutput",

136 + module = "vllm._rust_tool_parser",

137 + skip_from_py_object

138 +)]

139 +#[derive(Clone)]

140 +struct PyToolParserOutput(ToolParserOutput);

141 +

142 +#[pymethods]

143 +impl PyToolParserOutput {

144 + #[new]

145 + #[pyo3(signature = (normal_text="", calls=None))]

146 + fn new(py: Python<'_>, normal_text: &str, calls: Option<Vec<Py<PyToolCallDelta>>>) -> Self {

147 + let calls =

148 + calls.unwrap_or_default().iter().map(|call| call.borrow(py).0.clone()).collect();

149 + Self(ToolParserOutput {

150 + normal_text: normal_text.to_owned(),

151 + calls,

152 + })

153 + }

154 +

155 + #[getter]

156 + fn normal_text(&self) -> &str {

157 + &self.0.normal_text

158 + }

159 +

160 + #[getter]

161 + fn calls(&self) -> Vec<PyToolCallDelta> {

162 + self.0.calls.iter().cloned().map(PyToolCallDelta).collect()

163 + }

164 +

165 + fn append(&mut self, other: PyRef<'_, PyToolParserOutput>) {

166 + self.0.append(other.0.clone());

167 + }

168 +

169 + fn coalesce_calls(&self) -> Self {

170 + Self(self.0.clone().coalesce_calls())

171 + }

172 +}

173 +

174 +#[pyclass(name = "ToolParser", module = "vllm._rust_tool_parser", unsendable)]

175 +struct PyToolParser(Box<dyn ToolParser>);

176 +

177 +impl PyToolParser {

178 + fn parse_into_output(&mut self, chunk: &str, output: &mut PyToolParserOutput) -> PyResult<()> {

179 + self.0

180 + .parse_into(chunk, &mut output.0)

181 + .map_err(|error| PyValueError::new_err(error.to_report_string()))

182 + }

183 +}

184 +

185 +#[pymethods]

186 +impl PyToolParser {

187 + #[new]

188 + fn new(py: Python<'_>, parser_name: &str, tools: Vec<Py<PyTool>>) -> PyResult<Self> {

189 + let tools = tools.iter().map(|tool| tool.borrow(py).0.clone()).collect::<Vec<_>>();

190 + create_tool_parser(parser_name, &tools).map(Self)

191 + }

192 +

193 + fn parse_into(

194 + &mut self,

195 + chunk: &str,

196 + mut output: PyRefMut<'_, PyToolParserOutput>,

197 + ) -> PyResult<()> {

198 + self.parse_into_output(chunk, &mut output)

199 + }

200 +

201 + fn finish(&mut self) -> PyResult<PyToolParserOutput> {

202 + self.0

203 + .finish()

204 + .map(PyToolParserOutput)

205 + .map_err(|error| PyValueError::new_err(error.to_report_string()))

206 + }

207 +

208 + fn reset(&mut self) -> String {

209 + self.0.reset()

210 + }

211 +

212 + fn preserve_special_tokens(&self) -> bool {

213 + self.0.preserve_special_tokens()

214 + }

215 +

216 + fn tool_call_id(&self, tool_index: usize) -> Option<&str> {

217 + self.0.tool_call_id(tool_index)

218 + }

219 +}

220 +

221 +#[pymodule]

222 +fn _rust_tool_parser(m: &Bound<'_, PyModule>) -> PyResult<()> {

223 + m.add_class::<PyTool>()?;

224 + m.add_class::<PyToolCallDelta>()?;

225 + m.add_class::<PyToolParserOutput>()?;

226 + m.add_class::<PyToolParser>()?;

227 + Ok(())

228 +}

229 +

230 +#[cfg(test)]

231 +mod tests {

232 + use super::*;

233 + use serde_json::json;

234 +

235 + fn with_python<R>(f: impl for<'py> FnOnce(Python<'py>) -> R) -> R {

236 + Python::initialize();

237 + Python::attach(f)

238 + }

239 +

240 + fn tool_schema() -> Value {

241 + json!({

242 + "type": "object",

243 + "properties": {

244 + "user_id": {"type": "integer"},

245 + "shipping": {

246 + "type": "object",

247 + "properties": {

248 + "city": {"type": "string"},

249 + "zip": {"type": "integer"}

250 + }

251 + }

252 + }

253 + })

254 + }

255 +

256 + fn build_call() -> String {

257 + r#"<｜DSML｜tool_calls>

258 +<｜DSML｜invoke name="create_order">

259 +<｜DSML｜parameter name="user_id" string="false">42</｜DSML｜parameter>

260 +<｜DSML｜parameter name="shipping" string="false">{"city":"Singapore","zip":18956}</｜DSML｜parameter>

261 +</｜DSML｜invoke>

262 +</｜DSML｜tool_calls>"#

263 + .to_owned()

264 + }

265 +

266 + fn make_py_tool(py: Python<'_>) -> PyResult<Py<PyTool>> {

267 + let parameters = pythonize(py, &tool_schema()).map_err(|error| {

268 + PyValueError::new_err(format!(

269 + "failed to convert test schema from JSON to Python: {error}"

270 + ))

271 + })?;

272 + Py::new(

273 + py,

274 + PyTool::new(

275 + "create_order".to_owned(),

276 + Some("Create an order".to_owned()),

277 + &parameters,

278 + None,

279 + )?,

280 + )

281 + }

282 +

283 + #[test]

284 + fn tool_round_trips_typed_fields() {

285 + with_python(|py| {

286 + let tool = make_py_tool(py)?;

287 + let borrowed = tool.borrow(py);

288 + assert_eq!(borrowed.name(), "create_order");

289 + assert_eq!(borrowed.description(), Some("Create an order"));

290 + assert_eq!(borrowed.strict(), None);

291 +

292 + let parameters = borrowed.parameters(py)?;

293 + let parameters = depythonize::<Value>(parameters.bind(py))?;

294 + assert_eq!(parameters, tool_schema());

295 + PyResult::Ok(())

296 + })

297 + .unwrap();

298 + }

299 +

300 + #[test]

301 + fn output_append_and_coalesce_calls() {

302 + with_python(|py| {

303 + let first = Py::new(

304 + py,

305 + PyToolCallDelta::new(0, Some("create_order".to_owned()), "{\"a\"".to_owned()),

306 + )?;

307 + let second = Py::new(py, PyToolCallDelta::new(0, None, ":1}".to_owned()))?;

308 + let mut output = PyToolParserOutput::new(py, "text", Some(vec![first]));

309 + let other = Py::new(py, PyToolParserOutput::new(py, "", Some(vec![second])))?;

310 + output.append(other.borrow(py));

311 +

312 + let coalesced = output.coalesce_calls();

313 + assert_eq!(coalesced.normal_text(), "text");

314 + let calls = coalesced.calls();

315 + assert_eq!(calls.len(), 1);

316 + assert_eq!(calls[0].tool_index(), 0);

317 + assert_eq!(calls[0].name(), Some("create_order"));

318 + assert_eq!(calls[0].arguments(), "{\"a\":1}");

319 + PyResult::Ok(())

320 + })

321 + .unwrap();

322 + }

323 +

324 + #[test]

325 + fn parser_parse_finish_and_preserve_special_tokens() {

326 + with_python(|py| {

327 + let tool = make_py_tool(py)?;

328 + let mut parser = PyToolParser::new(py, "DeepSeekV4ToolParser", vec![tool])?;

329 + assert!(parser.preserve_special_tokens());

330 +

331 + let mut output = PyToolParserOutput::new(py, "", None);

332 + parser.parse_into_output(&build_call(), &mut output)?;

333 + let finish = Py::new(py, parser.finish()?)?;

334 + output.append(finish.borrow(py));

335 + let output = output.coalesce_calls();

336 +

337 + assert_eq!(output.normal_text(), "");

338 + let calls = output.calls();

339 + assert_eq!(calls.len(), 1);

340 + assert_eq!(calls[0].name(), Some("create_order"));

341 + assert_eq!(

342 + serde_json::from_str::<Value>(calls[0].arguments()).unwrap(),

343 + json!({

344 + "user_id": 42,

345 + "shipping": {

346 + "city": "Singapore",

347 + "zip": 18956

348 + }

349 + })

350 + );

351 +

352 + assert_eq!(parser.reset(), "");

353 + PyResult::Ok(())

354 + })

355 + .unwrap();

356 + }

357 +

358 + #[test]

359 + fn parser_exposes_model_emitted_tool_call_ids() {

360 + with_python(|py| {

361 + let tool = make_py_tool(py)?;

362 + let mut parser = PyToolParser::new(py, "KimiK2ToolParser", vec![tool])?;

363 +

364 + let input = "<|tool_calls_section_begin|>\

365 + <|tool_call_begin|>functions.create_order:0<|tool_call_argument_begin|>\

366 + {\"user_id\":42}<|tool_call_end|>\

367 + <|tool_calls_section_end|>";

368 + let mut output = PyToolParserOutput::new(py, "", None);

369 + parser.parse_into_output(input, &mut output)?;

370 +

371 + assert_eq!(parser.tool_call_id(0), Some("functions.create_order:0"));

372 + assert_eq!(parser.tool_call_id(1), None);

373 + PyResult::Ok(())

374 + })

375 + .unwrap();

376 + }

377 +

378 + #[test]

379 + fn parser_errors_for_unknown_name() {

380 + with_python(|py| {

381 + let tool = make_py_tool(py)?;

382 + let error = match PyToolParser::new(py, "missing", vec![tool]) {

383 + Ok(_) => panic!("missing parser name unexpectedly succeeded"),

384 + Err(error) => error,

385 + };

386 + let message = format!("{error}");

387 + assert!(message.contains("unsupported tool parser `missing`"));

388 + PyResult::Ok(())

389 + })

390 + .unwrap();

391 + }

392 +}

setup.py

+63 -13

@@ -35,6 +35,9 @@ ROOT_DIR = Path(__file__).parent

35 logger = logging.getLogger(__name__)

37 PRECOMPILED_RUST_FRONTEND_PATH = ROOT_DIR / "vllm" / "vllm-rs"

38 +# setuptools-rust installs PyO3 artifacts as `<module>.<ext-suffix>`, where the

39 +# suffix ends with `.so` on Linux and macOS alike (e.g. `_rust_foo.abi3.so`).

40 +PRECOMPILED_RUST_EXTENSION_MEMBER_REGEX = re.compile(r"vllm/_rust_[^/]*\.so$")

42 # cannot import envs directly because it depends on vllm,

43 # which is not installed yet

@@ -56,6 +59,25 @@ def should_require_rust_frontend() -> bool:

59 return value.lower() not in ("", "0", "false", "no")

62 +def get_precompiled_rust_extension_paths() -> list[Path]:

63 + return sorted((ROOT_DIR / "vllm").glob("_rust_*.so"))

64 +

65 +

66 +def get_missing_precompiled_rust_extension_modules() -> list[str]:

67 + present = {

68 + path.name.split(".", 1)[0] for path in get_precompiled_rust_extension_paths()

69 + }

70 + return [

71 + module_name

72 + for module_name in rust_build.rust_py_extension_module_names()

73 + if module_name not in present

74 + ]

75 +

76 +

77 +def has_precompiled_rust_extensions() -> bool:

78 + return not get_missing_precompiled_rust_extension_modules()

79 +

80 +

81 if sys.platform.startswith("darwin") and VLLM_TARGET_DEVICE != "cpu":

82 logger.warning("VLLM_TARGET_DEVICE automatically set to `cpu` due to macOS")

83 VLLM_TARGET_DEVICE = "cpu"

@@ -423,19 +445,31 @@ class precompiled_build_ext(build_ext):

445

446

447 class precompiled_build_rust(build_rust):

426 - """Skips local Rust builds when the precompiled wheel already ships vllm-rs."""

448 + """Skips local Rust builds when all precompiled Rust artifacts are present."""

449

450 def run(self) -> None:

429 - if PRECOMPILED_RUST_FRONTEND_PATH.exists():

451 + missing = []

452 + if not PRECOMPILED_RUST_FRONTEND_PATH.exists():

453 + missing.append(str(PRECOMPILED_RUST_FRONTEND_PATH))

454 + missing_rust_extensions = get_missing_precompiled_rust_extension_modules()

455 + if missing_rust_extensions:

456 + missing.extend(

457 + str(ROOT_DIR / "vllm" / f"{module_name}*.so")

458 + for module_name in missing_rust_extensions

459 + )

460 +

461 + if not missing:

462 logger.info(

431 - "Skipping local Rust build: using precompiled %s",

463 + "Skipping local Rust build: using precompiled %s and %s",

464 PRECOMPILED_RUST_FRONTEND_PATH,

465 + get_precompiled_rust_extension_paths(),

466 )

467 return

468

469 logger.warning(

437 - "Precompiled wheel did not provide %s; falling back to local Rust build.",

438 - PRECOMPILED_RUST_FRONTEND_PATH,

470 + "Precompiled wheel did not provide all Rust artifacts (%s); "

471 + "falling back to local Rust build.",

472 + ", ".join(missing),

473 )

474 super().run()

475

@@ -758,6 +792,14 @@ class precompiled_wheel_utils:

792 if member.filename in exact_members:

793 file_members.append(member)

794 continue

795 + if (

796 + extract_rust_frontend

797 + and PRECOMPILED_RUST_EXTENSION_MEMBER_REGEX.match(

798 + member.filename

799 + )

800 + ):

801 + file_members.append(member)

802 + continue

803

804 if not extract_extensions:

805 continue

@@ -1111,6 +1153,12 @@ package_data = {

1153 }

1154

1155

1156 +def add_vllm_package_data(filename: str) -> None:

1157 + vllm_files = package_data.setdefault("vllm", [])

1158 + if filename not in vllm_files:

1159 + vllm_files.append(filename)

1160 +

1161 +

1162 # If using precompiled artifacts, extract and patch package_data in advance.

1163 if USE_PRECOMPILED_RUST_FRONTEND:

1164 wheel_url, download_filename = precompiled_wheel_utils.determine_wheel_url()

@@ -1126,9 +1174,9 @@ if USE_PRECOMPILED_RUST_FRONTEND:

1174 # If the rust frontend binary is already present in the source tree (e.g.,

1175 # pre-built in a separate Docker build stage), ship it as-is.

1176 if PRECOMPILED_RUST_FRONTEND_PATH.exists():

1129 - vllm_files = package_data.setdefault("vllm", [])

1130 - if "vllm-rs" not in vllm_files:

1131 - vllm_files.append("vllm-rs")

1177 + add_vllm_package_data("vllm-rs")

1178 +for rust_extension_path in get_precompiled_rust_extension_paths():

1179 + add_vllm_package_data(rust_extension_path.name)

1180

1181 if _no_device():

1182 ext_modules = []

@@ -1141,13 +1189,15 @@ else:

1189 if USE_PRECOMPILED_EXTENSIONS

1190 else cmake_build_ext,

1191 }

1144 -if USE_PRECOMPILED_RUST_FRONTEND or PRECOMPILED_RUST_FRONTEND_PATH.exists():

1192 +if (

1193 + USE_PRECOMPILED_RUST_FRONTEND

1194 + or PRECOMPILED_RUST_FRONTEND_PATH.exists()

1195 + or has_precompiled_rust_extensions()

1196 +):

1197 cmdclass["build_rust"] = precompiled_build_rust

1198

1147 -# Rust frontend binary, built via setuptools-rust and installed into the

1148 -# package directory alongside the Python modules.

1149 -# TODO: we may use `RustBin` to directly install it into `bin` directory, but this

1150 -# requires extra work on using precompiled binaries.

1199 +# Rust artifacts, built via setuptools-rust and installed into the package

1200 +# directory alongside the Python modules.

1201 rust_extensions = rust_build.rust_extensions(

1202 optional=not should_require_rust_frontend()

1203 )

tests/tool_parsers/test_rust_tool_parser.py new

+328

@@ -0,0 +1,328 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +

4 +import json

5 +from collections.abc import Sequence

6 +from unittest.mock import MagicMock

7 +

8 +import pytest

9 +

10 +from vllm.entrypoints.openai.chat_completion.protocol import (

11 + ChatCompletionRequest,

12 + ChatCompletionToolsParam,

13 +)

14 +from vllm.tool_parsers.rust_tool_parser import RustToolParser

15 +

16 +# The PyO3 extension is an optional build artifact; skip when absent.

17 +_rust_tool_parser = pytest.importorskip("vllm._rust_tool_parser")

18 +

19 +MOCK_TOKENIZER = MagicMock()

20 +MOCK_TOKENIZER.get_vocab.return_value = {}

21 +

22 +TC_START = "<｜DSML｜tool_calls>"

23 +TC_END = "</｜DSML｜tool_calls>"

24 +INV_START = '<｜DSML｜invoke name="'

25 +INV_END = "</｜DSML｜invoke>"

26 +PARAM_START = '<｜DSML｜parameter name="'

27 +PARAM_END = "</｜DSML｜parameter>"

28 +

29 +

30 +class DeepSeekV4RustToolParser(RustToolParser):

31 + rust_parser_name = "DeepSeekV4ToolParser"

32 + tool_call_start_token = TC_START

33 +

34 +

35 +class KimiK2RustToolParser(RustToolParser):

36 + rust_parser_name = "KimiK2ToolParser"

37 + tool_call_start_token = "<|tool_calls_section_begin|>"

38 +

39 +

40 +def sample_tools() -> list[ChatCompletionToolsParam]:

41 + return [

42 + ChatCompletionToolsParam(

43 + type="function",

44 + function={

45 + "name": "get_weather",

46 + "description": "Get weather for a location",

47 + "parameters": {

48 + "type": "object",

49 + "properties": {

50 + "location": {"type": "string"},

51 + "date": {"type": "string"},

52 + },

53 + },

54 + },

55 + ),

56 + ChatCompletionToolsParam(

57 + type="function",

58 + function={

59 + "name": "add",

60 + "description": "Add two integers",

61 + "parameters": {

62 + "type": "object",

63 + "properties": {

64 + "x": {"type": "integer"},

65 + "y": {"type": "integer"},

66 + },

67 + },

68 + },

69 + ),

70 + ]

71 +

72 +

73 +EXPECTED_CALLS = [

74 + ("get_weather", {"location": "SF", "date": "2024-01-16"}),

75 + ("add", {"x": 3, "y": 5}),

76 +]

77 +

78 +

79 +def build_invoke(

80 + function_name: str,

81 + params: Sequence[tuple[str, str, bool]],

82 +) -> str:

83 + param_text = "\n".join(

84 + f'{PARAM_START}{name}" string="{str(is_string).lower()}">{value}{PARAM_END}'

85 + for name, value, is_string in params

86 + )

87 + return f'{INV_START}{function_name}">\n{param_text}\n{INV_END}\n'

88 +

89 +

90 +def build_tool_call() -> str:

91 + weather = build_invoke(

92 + "get_weather",

93 + [

94 + ("location", "SF", True),

95 + ("date", "2024-01-16", True),

96 + ],

97 + )

98 + add = build_invoke(

99 + "add",

100 + [

101 + ("x", "3", False),

102 + ("y", "5", False),

103 + ],

104 + )

105 + return f"{TC_START}\n{weather}{add}{TC_END}"

106 +

107 +

108 +def parse_streaming(

109 + parser: DeepSeekV4RustToolParser,

110 + text: str,

111 + chunk_size: int,

112 +) -> list:

113 + deltas = []

114 + previous_text = ""

115 + for start in range(0, len(text), chunk_size):

116 + delta_text = text[start : start + chunk_size]

117 + current_text = previous_text + delta_text

118 + delta = parser.extract_tool_calls_streaming(

119 + previous_text=previous_text,

120 + current_text=current_text,

121 + delta_text=delta_text,

122 + previous_token_ids=[],

123 + current_token_ids=[],

124 + delta_token_ids=[1],

125 + request=MagicMock(),

126 + )

127 + previous_text = current_text

128 + if delta is not None:

129 + deltas.append(delta)

130 +

131 + delta = parser.extract_tool_calls_streaming(

132 + previous_text=previous_text,

133 + current_text=previous_text,

134 + delta_text="",

135 + previous_token_ids=[],

136 + current_token_ids=[],

137 + delta_token_ids=[2],

138 + request=MagicMock(),

139 + )

140 + if delta is not None:

141 + deltas.append(delta)

142 +

143 + return deltas

144 +

145 +

146 +def collect_streamed_arguments(deltas: Sequence, tool_index: int = 0) -> str:

147 + return "".join(

148 + tool_call.function.arguments

149 + for delta in deltas

150 + for tool_call in delta.tool_calls or []

151 + if (

152 + tool_call.index == tool_index

153 + and tool_call.function is not None

154 + and tool_call.function.arguments is not None

155 + )

156 + )

157 +

158 +

159 +def test_rust_tool_parser_extension_typed_api() -> None:

160 + tools = [

161 + _rust_tool_parser.Tool(

162 + tool.function.name,

163 + tool.function.description,

164 + tool.function.parameters,

165 + None,

166 + )

167 + for tool in sample_tools()

168 + ]

169 + parser = _rust_tool_parser.ToolParser("DeepSeekV4ToolParser", tools)

170 + output = _rust_tool_parser.ToolParserOutput()

171 +

172 + parser.parse_into(build_tool_call(), output)

173 + output.append(parser.finish())

174 + output = output.coalesce_calls()

175 +

176 + assert parser.preserve_special_tokens()

177 + assert output.normal_text == ""

178 + assert len(output.calls) == 2

179 + for call, (name, arguments) in zip(output.calls, EXPECTED_CALLS):

180 + assert call.name == name

181 + assert json.loads(call.arguments) == arguments

182 +

183 +

184 +def test_rust_tool_parser_adapter_extracts_complete_output() -> None:

185 + tools = sample_tools()

186 + parser = DeepSeekV4RustToolParser(MOCK_TOKENIZER, tools=tools)

187 +

188 + result = parser.extract_tool_calls(

189 + "Let me create it. " + build_tool_call(),

190 + ChatCompletionRequest(messages=[], model="m", tools=tools),

191 + )

192 +

193 + assert result.tools_called

194 + assert result.content == "Let me create it. "

195 + assert len(result.tool_calls) == 2

196 + for tool_call, (name, arguments) in zip(result.tool_calls, EXPECTED_CALLS):

197 + assert tool_call.function.name == name

198 + assert json.loads(tool_call.function.arguments) == arguments

199 +

200 +

201 +def test_rust_tool_parser_adapter_streaming_handles_multiple_calls() -> None:

202 + parser = DeepSeekV4RustToolParser(MOCK_TOKENIZER, tools=sample_tools())

203 +

204 + deltas = parse_streaming(parser, build_tool_call(), chunk_size=5)

205 +

206 + names = [

207 + tool_call.function.name

208 + for delta in deltas

209 + for tool_call in delta.tool_calls or []

210 + if tool_call.function is not None and tool_call.function.name is not None

211 + ]

212 + assert names == [name for name, _ in EXPECTED_CALLS]

213 + for index, (_, arguments) in enumerate(EXPECTED_CALLS):

214 + assert json.loads(collect_streamed_arguments(deltas, index)) == arguments

215 +

216 +

217 +def test_rust_tool_parser_adapter_ignores_midstream_empty_delta() -> None:

218 + parser = DeepSeekV4RustToolParser(MOCK_TOKENIZER, tools=sample_tools())

219 + text = build_tool_call()

220 + split_at = len(TC_START) + 8

221 + deltas = []

222 + previous_text = ""

223 +

224 + for delta_text in (text[:split_at], "", text[split_at:], ""):

225 + current_text = previous_text + delta_text

226 + delta = parser.extract_tool_calls_streaming(

227 + previous_text=previous_text,

228 + current_text=current_text,

229 + delta_text=delta_text,

230 + previous_token_ids=[],

231 + current_token_ids=[],

232 + delta_token_ids=[1],

233 + request=MagicMock(),

234 + )

235 + previous_text = current_text

236 + if delta is not None:

237 + deltas.append(delta)

238 +

239 + names = [

240 + tool_call.function.name

241 + for delta in deltas

242 + for tool_call in delta.tool_calls or []

243 + if tool_call.function is not None and tool_call.function.name is not None

244 + ]

245 + assert names == [name for name, _ in EXPECTED_CALLS]

246 + for index, (_, arguments) in enumerate(EXPECTED_CALLS):

247 + assert json.loads(collect_streamed_arguments(deltas, index)) == arguments

248 +

249 +

250 +KIMI_EXPECTED_IDS = ["functions.get_weather:0", "functions.add:1"]

251 +

252 +

253 +def build_kimi_tool_call() -> str:

254 + return (

255 + "<|tool_calls_section_begin|>"

256 + "<|tool_call_begin|>functions.get_weather:0<|tool_call_argument_begin|>"

257 + '{"location": "SF", "date": "2024-01-16"}<|tool_call_end|>'

258 + "<|tool_call_begin|>functions.add:1<|tool_call_argument_begin|>"

259 + '{"x": 3, "y": 5}<|tool_call_end|>'

260 + "<|tool_calls_section_end|>"

261 + )

262 +

263 +

264 +def test_rust_tool_parser_adapter_complete_prefers_model_tool_call_ids() -> None:

265 + tools = sample_tools()

266 + parser = KimiK2RustToolParser(MOCK_TOKENIZER, tools=tools)

267 +

268 + result = parser.extract_tool_calls(

269 + "Let me check. " + build_kimi_tool_call(),

270 + ChatCompletionRequest(messages=[], model="m", tools=tools),

271 + )

272 +

273 + assert result.tools_called

274 + assert [tool_call.id for tool_call in result.tool_calls] == KIMI_EXPECTED_IDS

275 + for tool_call, (name, arguments) in zip(result.tool_calls, EXPECTED_CALLS):

276 + assert tool_call.function.name == name

277 + assert json.loads(tool_call.function.arguments) == arguments

278 +

279 +

280 +def test_rust_tool_parser_adapter_streaming_prefers_model_tool_call_ids() -> None:

281 + parser = KimiK2RustToolParser(MOCK_TOKENIZER, tools=sample_tools())

282 +

283 + deltas = parse_streaming(parser, build_kimi_tool_call(), chunk_size=5)

284 +

285 + ids = [

286 + tool_call.id

287 + for delta in deltas

288 + for tool_call in delta.tool_calls or []

289 + if tool_call.id is not None

290 + ]

291 + assert ids == KIMI_EXPECTED_IDS

292 + for index, (_, arguments) in enumerate(EXPECTED_CALLS):

293 + assert json.loads(collect_streamed_arguments(deltas, index)) == arguments

294 +

295 +

296 +def test_rust_tool_parser_adapter_streaming_generates_ids_as_fallback() -> None:

297 + # DeepSeekV4 never emits model tool call IDs, so the bridge mints them.

298 + parser = DeepSeekV4RustToolParser(MOCK_TOKENIZER, tools=sample_tools())

299 +

300 + deltas = parse_streaming(parser, build_tool_call(), chunk_size=5)

301 +

302 + ids = [

303 + tool_call.id

304 + for delta in deltas

305 + for tool_call in delta.tool_calls or []

306 + if tool_call.function is not None and tool_call.function.name is not None

307 + ]

308 + assert len(ids) == len(EXPECTED_CALLS)

309 + assert all(ids)

310 + assert len(set(ids)) == len(ids)

311 +

312 +

313 +def test_rust_tool_parser_adapter_adjust_request_is_opaque() -> None:

314 + tools = sample_tools()

315 + parser = DeepSeekV4RustToolParser(MOCK_TOKENIZER, tools=tools)

316 + request = ChatCompletionRequest(

317 + messages=[],

318 + model="m",

319 + tools=tools,

320 + tool_choice="required",

321 + skip_special_tokens=True,

322 + )

323 +

324 + adjusted = parser.adjust_request(request)

325 +

326 + assert adjusted is request

327 + assert adjusted.skip_special_tokens is False

328 + assert adjusted.structured_outputs is None

tools/build_rust.py

+23 -2

@@ -1,7 +1,7 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

4 -"""Shared setuptools-rust build entry for the vllm-rs binary."""

4 +"""Shared setuptools-rust build entry for Rust artifacts."""

6 from __future__ import annotations

@@ -15,7 +15,7 @@ from setuptools_rust import Binding, RustExtension

15 ROOT_DIR = Path(__file__).resolve().parents[1]

18 -def rust_extensions(*, optional: bool) -> list[RustExtension]:

18 +def rust_extensions(*, optional: bool = False) -> list[RustExtension]:

19 return [

20 RustExtension(

21 target="vllm.vllm-rs",

@@ -25,9 +25,30 @@ def rust_extensions(*, optional: bool) -> list[RustExtension]:

25 binding=Binding.Exec,

26 optional=optional,

27 ),

28 + RustExtension(

29 + target="vllm._rust_tool_parser",

30 + path="rust/src/tool-parser/python/Cargo.toml",

31 + features=["pyo3/abi3-py38"],

32 + binding=Binding.PyO3,

33 + optional=optional,

34 + py_limited_api=True,

35 + ),

36 ]

39 +def rust_py_extension_module_names() -> list[str]:

40 + module_names = []

41 + for extension in rust_extensions():

42 + if extension.binding != Binding.PyO3:

43 + continue

44 +

45 + for target_name in extension.target.values():

46 + if target_name.startswith("vllm._rust_"):

47 + module_names.append(target_name.rsplit(".", 1)[-1])

48 +

49 + return module_names

50 +

51 +

52 def build_binary(build_rust_args: list[str]) -> None:

53 os.chdir(ROOT_DIR)

54 (ROOT_DIR / "vllm").mkdir(exist_ok=True)

vllm/tool_parsers/rust_tool_parser.py new

+322

@@ -0,0 +1,322 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +

4 +import importlib

5 +from collections.abc import Sequence

6 +from typing import Any

7 +

8 +from openai.types.responses.function_tool import FunctionTool

9 +

10 +from vllm.entrypoints.chat_utils import make_tool_call_id

11 +from vllm.entrypoints.openai.chat_completion.protocol import (

12 + ChatCompletionRequest,

13 + ChatCompletionToolsParam,

14 +)

15 +from vllm.entrypoints.openai.engine.protocol import (

16 + DeltaFunctionCall,

17 + DeltaMessage,

18 + DeltaToolCall,

19 + ExtractedToolCallInformation,

20 + FunctionCall,

21 + ToolCall,

22 +)

23 +from vllm.entrypoints.openai.responses.protocol import ResponsesRequest

24 +from vllm.logger import init_logger

25 +from vllm.tokenizers import TokenizerLike

26 +from vllm.tool_parsers.abstract_tool_parser import Tool, ToolParser

27 +

28 +logger = init_logger(__name__)

29 +

30 +

31 +def _rust_tool_parser_module() -> Any:

32 + try:

33 + return importlib.import_module("vllm._rust_tool_parser")

34 + except ImportError as exc:

35 + raise RuntimeError(

36 + "Rust tool parsing requires the vllm._rust_tool_parser PyO3 "

37 + "extension. Rebuild vLLM with Rust frontend/extensions enabled."

38 + ) from exc

39 +

40 +

41 +class RustToolParser(ToolParser):

42 + """Adapter from an opaque Rust parser to the vLLM ToolParser API.

43 +

44 + Subclasses provide only model-specific configuration: the exact Rust parser

45 + name and an optional tool-call start marker for fast complete-output

46 + rejection.

47 +

48 + This class keeps the vLLM-specific bridge work:

49 + - convert vLLM tool definitions into the Rust ``Tool`` shape;

50 + - translate typed Rust parser outputs into vLLM protocol objects; and

51 + - maintain vLLM streaming bookkeeping used by finish-reason handling.

52 +

53 + The parser grammar and incremental parser state stay in Rust.

54 + """

55 +

56 + # Rust-backed parsers are opaque to Python by default. Do not use vLLM's

57 + # standard JSON required/named handling; let the Rust parser consume the

58 + # model's native tool-call syntax.

59 + supports_required_and_named = False

60 +

61 + rust_parser_name: str

62 + tool_call_start_token: str | None = None

63 +

64 + def __init__(self, tokenizer: TokenizerLike, tools: list[Tool] | None = None):

65 + super().__init__(tokenizer, tools)

66 + self._parser: Any | None = None

67 + self._error: Exception | None = None

68 +

69 + if not self.model_tokenizer:

70 + raise ValueError(

71 + "The model tokenizer must be passed to the ToolParser "

72 + "constructor during construction."

73 + )

74 +

75 + logger.debug(

76 + "vLLM successfully imported tool parser %s", self.__class__.__name__

77 + )

78 +

79 + def adjust_request(

80 + self, request: ChatCompletionRequest | ResponsesRequest

81 + ) -> ChatCompletionRequest | ResponsesRequest:

82 + """Adjust request options without installing Python-side constraints.

83 +

84 + Rust-backed parsers are treated as source-of-truth opaque parsers. The

85 + bridge intentionally avoids ``super().adjust_request()`` so Python does

86 + not install JSON schema guidance or structural-tag constraints that may

87 + conflict with the Rust parser's native grammar.

88 + """

89 + if self._get_parser().preserve_special_tokens():

90 + request.skip_special_tokens = False

91 + return request

92 +

93 + def _rust_tools(self) -> list[Any]:

94 + """Build Rust ``Tool`` objects from vLLM tool definitions."""

95 + if not self.tools:

96 + return []

97 +

98 + tools: list[Any] = []

99 + for tool in self.tools:

100 + if isinstance(tool, FunctionTool):

101 + name = tool.name

102 + description = tool.description

103 + parameters = tool.parameters or {}

104 + strict = getattr(tool, "strict", None)

105 + elif isinstance(tool, ChatCompletionToolsParam):

106 + name = tool.function.name

107 + description = tool.function.description

108 + parameters = tool.function.parameters or {}

109 + strict = getattr(tool.function, "strict", None)

110 + else:

111 + continue

112 + tools.append(

113 + _rust_tool_parser_module().Tool(name, description, parameters, strict)

114 + )

115 + return tools

116 +

117 + def _new_parser(self) -> Any:

118 + """Create a fresh Rust parser with the current tool schemas."""

119 + return _rust_tool_parser_module().ToolParser(

120 + self.rust_parser_name, self._rust_tools()

121 + )

122 +

123 + def _get_parser(self) -> Any:

124 + if self._parser is None:

125 + self._parser = self._new_parser()

126 + return self._parser

127 +

128 + def _reset_streaming_state(self) -> None:

129 + """Reset parser state for a new request on a reused parser instance."""

130 + self._parser = self._new_parser()

131 + self._error = None

132 + self.prev_tool_call_arr.clear()

133 + self.streamed_args_for_tool.clear()

134 + self.current_tool_id = -1

135 + self.current_tool_name_sent = False

136 +

137 + def _ensure_tool_state(self, index: int) -> None:

138 + """Grow vLLM streaming state arrays to contain ``index``."""

139 + while len(self.prev_tool_call_arr) <= index:

140 + self.prev_tool_call_arr.append({})

141 + while len(self.streamed_args_for_tool) <= index:

142 + self.streamed_args_for_tool.append("")

143 +

144 + def _record_delta(

145 + self, index: int, name: str | None, arguments: str | None

146 + ) -> str | None:

147 + """Mirror a Rust parser delta into vLLM streaming bookkeeping.

148 +

149 + ``prev_tool_call_arr`` and ``streamed_args_for_tool`` are read later by

150 + the chat serving layer to decide the final ``tool_calls`` finish reason

151 + and to flush any remaining argument bytes.

152 + """

153 + tool_call_id = None

154 + self._ensure_tool_state(index)

155 +

156 + if name is not None:

157 + # Prefer the model-emitted ID surfaced by the Rust parser (e.g.

158 + # Kimi K2) over a randomly generated one.

159 + tool_call_id = self._get_parser().tool_call_id(index) or make_tool_call_id()

160 + self.prev_tool_call_arr[index] = {"name": name, "arguments": {}}

161 + self.current_tool_name_sent = True

162 +

163 + if arguments is not None:

164 + self.streamed_args_for_tool[index] += arguments

165 + self.prev_tool_call_arr[index]["arguments"] = self.streamed_args_for_tool[

166 + index

167 + ]

168 + self.current_tool_id = index

169 +

170 + return tool_call_id

171 +

172 + def _delta_message_from_parser_output(

173 + self, parser_output: Any | None

174 + ) -> DeltaMessage | None:

175 + """Translate one Rust parser output into a vLLM ``DeltaMessage``."""

176 + if parser_output is None:

177 + return None

178 +

179 + normal_text = parser_output.normal_text or None

180 + tool_calls: list[DeltaToolCall] = []

181 + for tool_call in parser_output.calls:

182 + index = tool_call.tool_index

183 + name = tool_call.name

184 + arguments: str | None = tool_call.arguments

185 + if name is None and arguments is None:

186 + continue

187 +

188 + tool_call_id = self._record_delta(index, name, arguments)

189 + tool_calls.append(

190 + DeltaToolCall(

191 + index=index,

192 + id=tool_call_id,

193 + type="function" if name is not None else None,

194 + function=DeltaFunctionCall(

195 + name=name,

196 + arguments=arguments,

197 + ),

198 + )

199 + )

200 +

201 + if normal_text is None and not tool_calls:

202 + return None

203 + return DeltaMessage(content=normal_text, tool_calls=tool_calls)

204 +

205 + def _parse_complete(self, model_output: str) -> tuple[Any, dict[int, str]] | None:

206 + """Parse complete model output with a throwaway Rust parser instance.

207 +

208 + Returns the coalesced parser output along with any model-emitted tool

209 + call IDs keyed by tool index.

210 + """

211 + parser = self._new_parser()

212 + output = _rust_tool_parser_module().ToolParserOutput()

213 + try:

214 + parser.parse_into(model_output, output)

215 + # finish() clears parser state, so snapshot model-emitted IDs first.

216 + tool_call_ids = {

217 + call.tool_index: tool_call_id

218 + for call in output.calls

219 + if (tool_call_id := parser.tool_call_id(call.tool_index)) is not None

220 + }

221 + output.append(parser.finish())

222 + except Exception:

223 + logger.exception(

224 + "Error parsing %s tool call output.", self.rust_parser_name

225 + )

226 + return None

227 + return output.coalesce_calls(), tool_call_ids

228 +

229 + def extract_tool_calls(

230 + self,

231 + model_output: str,

232 + request: ChatCompletionRequest,

233 + ) -> ExtractedToolCallInformation:

234 + """Extract tool calls from complete model output (non-streaming)."""

235 + if (

236 + self.tool_call_start_token is not None

237 + and self.tool_call_start_token not in model_output

238 + ):

239 + return ExtractedToolCallInformation(

240 + tools_called=False,

241 + tool_calls=[],

242 + content=model_output,

243 + )

244 +

245 + parse_result = self._parse_complete(model_output)

246 + if parse_result is None:

247 + return ExtractedToolCallInformation(

248 + tools_called=False,

249 + tool_calls=[],

250 + content=model_output,

251 + )

252 + parsed, tool_call_ids = parse_result

253 +

254 + tool_calls: list[ToolCall] = []

255 + self.prev_tool_call_arr.clear()

256 + for parsed_tool_call in parsed.calls:

257 + name = parsed_tool_call.name

258 + arguments = parsed_tool_call.arguments or "{}"

259 + if name is None:

260 + continue

261 + tool_calls.append(

262 + ToolCall(

263 + id=tool_call_ids.get(parsed_tool_call.tool_index)

264 + or make_tool_call_id(),

265 + type="function",

266 + function=FunctionCall(name=name, arguments=arguments),

267 + )

268 + )

269 + self.prev_tool_call_arr.append({"name": name, "arguments": arguments})

270 +

271 + if not tool_calls:

272 + return ExtractedToolCallInformation(

273 + tools_called=False,

274 + tool_calls=[],

275 + content=model_output,

276 + )

277 +

278 + content = parsed.normal_text or None

279 + return ExtractedToolCallInformation(

280 + tools_called=True,

281 + tool_calls=tool_calls,

282 + content=content,

283 + )

284 +

285 + def extract_tool_calls_streaming(

286 + self,

287 + previous_text: str,

288 + current_text: str,

289 + delta_text: str,

290 + previous_token_ids: Sequence[int], # pylint: disable=unused-argument

291 + current_token_ids: Sequence[int], # pylint: disable=unused-argument

292 + delta_token_ids: Sequence[int], # pylint: disable=unused-argument

293 + request: ChatCompletionRequest, # pylint: disable=unused-argument

294 + ) -> DeltaMessage | None:

295 + """Extract tool calls from streaming model output.

296 +

297 + The Rust parser owns the incremental buffer, so this adapter feeds only

298 + the newest text delta and lets the serving layer handle final empty

299 + chunks.

300 + """

301 + # TODO: Add a final-chunk hook if streaming needs to call Rust finish().

302 + if not previous_text:

303 + self._reset_streaming_state()

304 +

305 + if self._error is not None:

306 + return None

307 +

308 + parser_output = _rust_tool_parser_module().ToolParserOutput()

309 + try:

310 + self._get_parser().parse_into(delta_text, parser_output)

311 + except Exception as error:

312 + self._error = error

313 + logger.exception(

314 + "Error parsing %s streaming tool call output.",

315 + self.rust_parser_name,

316 + )

317 +

318 + delta_message = self._delta_message_from_parser_output(parser_output)

319 + if delta_message is not None:

320 + return delta_message

321 +

322 + return None