diff --git a/.buildkite/scripts/run-rust-frontend-cargo-ci.sh b/.buildkite/scripts/run-rust-frontend-cargo-ci.sh index 6ce9b5200c..4b4272762a 100755 --- a/.buildkite/scripts/run-rust-frontend-cargo-ci.sh +++ b/.buildkite/scripts/run-rust-frontend-cargo-ci.sh @@ -110,6 +110,36 @@ install_uv() { | env UV_INSTALL_DIR="$CARGO_HOME/bin" sh } +setup_pyo3_python() { + local python_version="${PYO3_PYTHON_VERSION:-3.12}" + + log_section "Installing Python ${python_version} for PyO3 tests" + uv python install "$python_version" + PYO3_PYTHON="$(uv python find \ + --managed-python \ + --no-project \ + --resolve-links \ + "$python_version")" + export PYO3_PYTHON + + local python_libdir + python_libdir="$("$PYO3_PYTHON" - <<'PY' +import pathlib +import sysconfig + +libdir = pathlib.Path(sysconfig.get_config_var("LIBDIR")) +ldlibrary = sysconfig.get_config_var("LDLIBRARY") +assert sysconfig.get_config_var("Py_ENABLE_SHARED") == 1 +assert ldlibrary +assert (libdir / ldlibrary).exists(), libdir / ldlibrary +print(libdir) +PY +)" + + export LD_LIBRARY_PATH="${python_libdir}:${LD_LIBRARY_PATH:-}" + export LIBRARY_PATH="${python_libdir}:${LIBRARY_PATH:-}" +} + run_style_clippy() { install_cargo_sort @@ -132,6 +162,7 @@ run_style_clippy() { run_tests() { install_uv + setup_pyo3_python install_cargo_nextest log_section "Running cargo nextest" diff --git a/.buildkite/test_areas/misc.yaml b/.buildkite/test_areas/misc.yaml index 7511acca00..cda2bb4daf 100644 --- a/.buildkite/test_areas/misc.yaml +++ b/.buildkite/test_areas/misc.yaml @@ -301,9 +301,9 @@ steps: - tests/multimodal - tests/renderers - tests/standalone_tests/lazy_imports.py - - tests/tokenizers_ - tests/reasoning - tests/tool_parsers + - tests/tokenizers_ - tests/parser - tests/transformers_utils - tests/config @@ -317,9 +317,9 @@ steps: - pytest -v -s test_ray_env.py - pytest -v -s -m 'cpu_test' multimodal - pytest -v -s renderers - - pytest -v -s tokenizers_ - pytest -v -s reasoning --ignore=reasoning/test_seedoss_reasoning_parser.py --ignore=reasoning/test_glm4_moe_reasoning_parser.py - pytest -v -s tool_parsers + - pytest -v -s tokenizers_ - pytest -v -s parser - pytest -v -s transformers_utils - pytest -v -s config diff --git a/build_rust.sh b/build_rust.sh index b5ba1d739a..1efc1ce39f 100755 --- a/build_rust.sh +++ b/build_rust.sh @@ -1,5 +1,5 @@ #!/bin/bash -# Build the vllm-rs Rust frontend binary. +# Build vLLM Rust artifacts and install them into the vllm package. # Usage: ./build_rust.sh [--debug] # # By default builds in release mode. Pass --debug for faster compile times diff --git a/docker/Dockerfile b/docker/Dockerfile index 34d1ec7975..300028cfb2 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -281,7 +281,8 @@ COPY requirements/build/rust.txt requirements/build/rust.txt RUN --mount=type=cache,target=/opt/uv/cache \ uv pip install --python /opt/venv/bin/python3 -r requirements/build/rust.txt -# Copy only the Rust build inputs. The binary is the sole artifact we need. +# Copy only the Rust build inputs; build_rust.sh publishes artifacts needed +# by the wheel build stage. COPY rust rust COPY rust-toolchain.toml rust-toolchain.toml COPY tools/build_rust.py tools/build_rust.py @@ -291,8 +292,9 @@ COPY build_rust.sh build_rust.sh # (rustc spawns enough concurrent processes to hit RLIMIT_NOFILE otherwise). ENV CARGO_BUILD_JOBS=4 -# Build the release binary. Cache cargo registry/git, but not target/, because -# stale target metadata can outlive source updates across BuildKit cache reuse. +# Build the release artifacts. Cache cargo registry/git, but not target/, +# because stale target metadata can outlive source updates across BuildKit +# cache reuse. RUN --mount=type=cache,target=/root/.cargo/registry \ --mount=type=cache,target=/root/.cargo/git \ bash build_rust.sh @@ -503,9 +505,10 @@ WORKDIR /workspace COPY --from=csrc-build /workspace/dist /precompiled-wheels COPY . . -# Drop the pre-built rust frontend binary into the source tree. setup.py -# detects it and ships it as-is, skipping the local cargo build. +# Drop the pre-built Rust artifacts into the source tree. setup.py detects +# them and ships them as-is, skipping the local Rust build. COPY --from=rust-build /workspace/vllm/vllm-rs vllm/vllm-rs +COPY --from=rust-build /workspace/vllm/_rust_*.so vllm/ ARG GIT_REPO_CHECK=0 RUN --mount=type=bind,source=.git,target=.git \ diff --git a/docker/Dockerfile.cpu b/docker/Dockerfile.cpu index f86097cdb3..4df401395f 100644 --- a/docker/Dockerfile.cpu +++ b/docker/Dockerfile.cpu @@ -104,7 +104,8 @@ WORKDIR /workspace COPY requirements/build/rust.txt requirements/build/rust.txt RUN python3 -m pip install --no-cache-dir -r requirements/build/rust.txt -# Copy only the Rust build inputs. The binary is the sole artifact we need. +# Copy only the Rust build inputs; build_rust.sh publishes artifacts needed +# by the wheel build stage. COPY rust rust COPY rust-toolchain.toml rust-toolchain.toml COPY tools/build_rust.py tools/build_rust.py @@ -114,8 +115,9 @@ COPY build_rust.sh build_rust.sh # (rustc spawns enough concurrent processes to hit RLIMIT_NOFILE otherwise). ENV CARGO_BUILD_JOBS=4 -# Build the release binary. Cache cargo registry/git, but not target/, because -# stale target metadata can outlive source updates across BuildKit cache reuse. +# Build the release artifacts. Cache cargo registry/git, but not target/, +# because stale target metadata can outlive source updates across BuildKit +# cache reuse. RUN --mount=type=cache,target=/root/.cargo/registry,sharing=locked \ --mount=type=cache,target=/root/.cargo/git,sharing=locked \ bash build_rust.sh @@ -151,9 +153,10 @@ RUN --mount=type=cache,target=/root/.cache/uv \ COPY . . -# Drop the pre-built rust frontend binary into the source tree. setup.py -# detects it and ships it as-is, skipping the local cargo build. +# Drop the pre-built Rust artifacts into the source tree. setup.py detects +# them and ships them as-is, skipping the local Rust build. COPY --from=rust-build /workspace/vllm/vllm-rs vllm/vllm-rs +COPY --from=rust-build /workspace/vllm/_rust_*.so vllm/ RUN if [ "$GIT_REPO_CHECK" != 0 ]; then bash tools/check_repo.sh ; fi diff --git a/docker/Dockerfile.nightly_torch b/docker/Dockerfile.nightly_torch index 1ac3626088..e1cd08bd66 100644 --- a/docker/Dockerfile.nightly_torch +++ b/docker/Dockerfile.nightly_torch @@ -113,7 +113,8 @@ WORKDIR /workspace COPY requirements/build/rust.txt requirements/build/rust.txt RUN python3 -m pip install --no-cache-dir -r requirements/build/rust.txt -# Copy only the Rust build inputs. The binary is the sole artifact we need. +# Copy only the Rust build inputs; build_rust.sh publishes artifacts needed +# by the wheel build stage. COPY rust rust COPY rust-toolchain.toml rust-toolchain.toml COPY tools/build_rust.py tools/build_rust.py @@ -138,9 +139,10 @@ ENV UV_HTTP_TIMEOUT=500 COPY . . -# Drop the pre-built rust frontend binary into the source tree. setup.py -# detects it and ships it as-is, skipping the local cargo build. +# Drop the pre-built Rust artifacts into the source tree. setup.py detects +# them and ships them as-is, skipping the local Rust build. COPY --from=rust-build /workspace/vllm/vllm-rs vllm/vllm-rs +COPY --from=rust-build /workspace/vllm/_rust_*.so vllm/ RUN python3 use_existing_torch.py diff --git a/docker/Dockerfile.rocm b/docker/Dockerfile.rocm index ebde46b6d0..dcae40c524 100644 --- a/docker/Dockerfile.rocm +++ b/docker/Dockerfile.rocm @@ -208,9 +208,10 @@ ENV VLLM_TARGET_DEVICE=rocm COPY --from=csrc-build ${COMMON_WORKDIR}/vllm/dist /precompiled-wheels -# Drop the pre-built rust frontend binary into the source tree. setup.py -# detects it and ships it as-is, skipping the local cargo build. +# Drop the pre-built Rust artifacts into the source tree. setup.py detects +# them and ships them as-is, skipping the local Rust build. COPY --from=rust-build ${COMMON_WORKDIR}/vllm/vllm/vllm-rs ${COMMON_WORKDIR}/vllm/vllm/vllm-rs +COPY --from=rust-build ${COMMON_WORKDIR}/vllm/vllm/_rust_*.so ${COMMON_WORKDIR}/vllm/vllm/ RUN --mount=type=cache,id=vllm-rocm-uv,target=/root/.cache/uv \ cd vllm \ @@ -417,9 +418,10 @@ FROM fetch_vllm AS build_vllm_wheel_release ARG COMMON_WORKDIR -# Drop the pre-built rust frontend binary into the source tree. setup.py -# detects it and ships it as-is, skipping the local cargo build. +# Drop the pre-built Rust artifacts into the source tree. setup.py detects +# them and ships them as-is, skipping the local Rust build. COPY --from=rust-build ${COMMON_WORKDIR}/vllm/vllm/vllm-rs ${COMMON_WORKDIR}/vllm/vllm/vllm-rs +COPY --from=rust-build ${COMMON_WORKDIR}/vllm/vllm/_rust_*.so ${COMMON_WORKDIR}/vllm/vllm/ # Create /install directory for custom wheels RUN mkdir -p /install diff --git a/docker/Dockerfile.xpu b/docker/Dockerfile.xpu index 3137d882fd..ca08d9b95f 100644 --- a/docker/Dockerfile.xpu +++ b/docker/Dockerfile.xpu @@ -17,7 +17,8 @@ WORKDIR /workspace COPY requirements/build/rust.txt requirements/build/rust.txt RUN python3 -m pip install --no-cache-dir -r requirements/build/rust.txt -# Copy only the Rust build inputs. The binary is the sole artifact we need. +# Copy only the Rust build inputs; build_rust.sh publishes artifacts needed +# by the wheel build stage. COPY rust rust COPY rust-toolchain.toml rust-toolchain.toml COPY tools/build_rust.py tools/build_rust.py @@ -212,9 +213,10 @@ RUN --mount=type=cache,target=/root/.cache/uv \ # don't invalidate heavy dependency and UCX/NIXL layers. COPY . . -# Drop the pre-built rust frontend binary into the source tree. setup.py -# detects it and ships it as-is, skipping the local cargo build. +# Drop the pre-built Rust artifacts into the source tree. setup.py detects +# them and ships them as-is, skipping the local Rust build. COPY --from=rust-build /workspace/vllm/vllm-rs vllm/vllm-rs +COPY --from=rust-build /workspace/vllm/_rust_*.so vllm/ ARG GIT_REPO_CHECK=0 RUN --mount=type=bind,source=.git,target=.git \ diff --git a/rust/Cargo.lock b/rust/Cargo.lock index ef8c2b90a1..e6011ddf5c 100644 --- a/rust/Cargo.lock +++ b/rust/Cargo.lock @@ -3458,6 +3458,75 @@ version = "0.1.29" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e0c5ccf5294c6ccd63a74f1565028353830a9c2f5eb0c682c355c471726a6e3f" +[[package]] +name = "pyo3" +version = "0.28.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91fd8e38a3b50ed1167fb981cd6fd60147e091784c427b8f7183a7ee32c31c12" +dependencies = [ + "libc", + "once_cell", + "portable-atomic", + "pyo3-build-config", + "pyo3-ffi", + "pyo3-macros", +] + +[[package]] +name = "pyo3-build-config" +version = "0.28.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e368e7ddfdeb98c9bca7f8383be1648fd84ab466bf2bc015e94008db6d35611e" +dependencies = [ + "target-lexicon", +] + +[[package]] +name = "pyo3-ffi" +version = "0.28.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7f29e10af80b1f7ccaf7f69eace800a03ecd13e883acfacc1e5d0988605f651e" +dependencies = [ + "libc", + "pyo3-build-config", +] + +[[package]] +name = "pyo3-macros" +version = "0.28.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "df6e520eff47c45997d2fc7dd8214b25dd1310918bbb2642156ef66a67f29813" +dependencies = [ + "proc-macro2", + "pyo3-macros-backend", + "quote", + "syn 2.0.117", +] + +[[package]] +name = "pyo3-macros-backend" +version = "0.28.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c4cdc218d835738f81c2338f822078af45b4afdf8b2e33cbb5916f108b813acb" +dependencies = [ + "heck", + "proc-macro2", + "pyo3-build-config", + "quote", + "syn 2.0.117", +] + +[[package]] +name = "pythonize" +version = "0.28.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b79f670c9626c8b651c0581011b57b6ba6970bb69faf01a7c4c0cfc81c43f95" +dependencies = [ + "pyo3", + "serde", + "serde_json", +] + [[package]] name = "qoi" version = "0.4.1" @@ -4669,6 +4738,12 @@ dependencies = [ "libc", ] +[[package]] +name = "target-lexicon" +version = "0.13.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "adb6935a6f5c20170eeceb1a3835a49e12e19d792f6dd344ccc76a985ca5a6ca" + [[package]] name = "task-local" version = "0.1.1" @@ -5904,6 +5979,17 @@ dependencies = [ "winnow", ] +[[package]] +name = "vllm-tool-parser-py" +version = "0.1.0" +dependencies = [ + "pyo3", + "pythonize", + "serde_json", + "thiserror-ext", + "vllm-tool-parser", +] + [[package]] name = "walkdir" version = "2.5.0" diff --git a/rust/Cargo.toml b/rust/Cargo.toml index ba11bd70a5..c61fd9c19e 100644 --- a/rust/Cargo.toml +++ b/rust/Cargo.toml @@ -12,6 +12,7 @@ members = [ "src/text", "src/tokenizer", "src/tool-parser", + "src/tool-parser/python", ] resolver = "3" @@ -60,6 +61,8 @@ prometheus-client = "0.24.0" prometheus-client-derive-encode = "0.5.0" prost = "0.14.3" prost-types = "0.14.3" +pyo3 = "0.28.3" +pythonize = "0.28.0" rand = "0.9.2" reasoning-parser = "1.2.2" reqwest = { version = "0.12.8", default-features = false, features = ["rustls-tls"] } diff --git a/rust/src/tool-parser/python/Cargo.toml b/rust/src/tool-parser/python/Cargo.toml new file mode 100644 index 0000000000..c029ad9013 --- /dev/null +++ b/rust/src/tool-parser/python/Cargo.toml @@ -0,0 +1,19 @@ +[package] +name = "vllm-tool-parser-py" +version.workspace = true +edition.workspace = true +license.workspace = true + +[lib] +name = "_rust_tool_parser" +crate-type = ["cdylib", "rlib"] + +[dependencies] +pyo3.workspace = true +pythonize = { workspace = true, features = ["serde_json"] } +serde_json.workspace = true +thiserror-ext.workspace = true +vllm-tool-parser.workspace = true + +[lints] +workspace = true diff --git a/rust/src/tool-parser/python/src/lib.rs b/rust/src/tool-parser/python/src/lib.rs new file mode 100644 index 0000000000..81aed04b1c --- /dev/null +++ b/rust/src/tool-parser/python/src/lib.rs @@ -0,0 +1,392 @@ +//! Thin PyO3 bindings for `vllm_tool_parser`. +//! +//! This crate exposes the Rust tool parser trait and data shapes to Python +//! while keeping parser state, grammar, and schema-aware argument conversion in +//! Rust. Python callers should use this module as a typed bridge and keep any +//! vLLM protocol adaptation outside the binding. + +use pyo3::exceptions::PyValueError; +use pyo3::prelude::*; +use pyo3::types::{PyAny, PyModule}; +use pythonize::{depythonize, pythonize}; +use serde_json::Value; +use thiserror_ext::AsReport as _; +use vllm_tool_parser::{Tool, ToolCallDelta, ToolParser, ToolParserOutput}; + +macro_rules! tool_parser_factory { + ($($parser:ident),+ $(,)?) => { + fn create_tool_parser( + name: &str, + tools: &[Tool], + ) -> PyResult> { + match name { + $( + stringify!($parser) => { + ::create(tools) + } + )+ + _ => { + return Err(PyValueError::new_err(format!( + "unsupported tool parser `{name}`" + ))); + } + } + .map_err(|error| PyValueError::new_err(error.to_report_string())) + } + }; +} + +// Export a tool parser to Python by registering it here. +tool_parser_factory! { + // Below are the parsers just for testing purposes on Python side. + DeepSeekV4ToolParser, + KimiK2ToolParser, +} + +#[pyclass(name = "Tool", module = "vllm._rust_tool_parser", skip_from_py_object)] +#[derive(Clone)] +struct PyTool(Tool); + +#[pymethods] +impl PyTool { + #[new] + #[pyo3(signature = (name, description, parameters, strict=None))] + fn new( + name: String, + description: Option, + parameters: &Bound<'_, PyAny>, + strict: Option, + ) -> PyResult { + let parameters = depythonize::(parameters).map_err(|error| { + PyValueError::new_err(format!( + "failed to convert tool parameters from Python to JSON: {error}" + )) + })?; + Ok(Self(Tool { + name, + description, + parameters, + strict, + })) + } + + #[getter] + fn name(&self) -> &str { + &self.0.name + } + + #[getter] + fn description(&self) -> Option<&str> { + self.0.description.as_deref() + } + + #[getter] + fn parameters(&self, py: Python<'_>) -> PyResult> { + pythonize(py, &self.0.parameters).map(Bound::unbind).map_err(|error| { + PyValueError::new_err(format!( + "failed to convert tool parameters from JSON to Python: {error}" + )) + }) + } + + #[getter] + fn strict(&self) -> Option { + self.0.strict + } +} + +#[pyclass( + name = "ToolCallDelta", + module = "vllm._rust_tool_parser", + skip_from_py_object +)] +#[derive(Clone)] +struct PyToolCallDelta(ToolCallDelta); + +#[pymethods] +impl PyToolCallDelta { + #[new] + #[pyo3(signature = (tool_index, name, arguments))] + fn new(tool_index: usize, name: Option, arguments: String) -> Self { + Self(ToolCallDelta { + tool_index, + name, + arguments, + }) + } + + #[getter] + fn tool_index(&self) -> usize { + self.0.tool_index + } + + #[getter] + fn name(&self) -> Option<&str> { + self.0.name.as_deref() + } + + #[getter] + fn arguments(&self) -> &str { + &self.0.arguments + } +} + +#[pyclass( + name = "ToolParserOutput", + module = "vllm._rust_tool_parser", + skip_from_py_object +)] +#[derive(Clone)] +struct PyToolParserOutput(ToolParserOutput); + +#[pymethods] +impl PyToolParserOutput { + #[new] + #[pyo3(signature = (normal_text="", calls=None))] + fn new(py: Python<'_>, normal_text: &str, calls: Option>>) -> Self { + let calls = + calls.unwrap_or_default().iter().map(|call| call.borrow(py).0.clone()).collect(); + Self(ToolParserOutput { + normal_text: normal_text.to_owned(), + calls, + }) + } + + #[getter] + fn normal_text(&self) -> &str { + &self.0.normal_text + } + + #[getter] + fn calls(&self) -> Vec { + self.0.calls.iter().cloned().map(PyToolCallDelta).collect() + } + + fn append(&mut self, other: PyRef<'_, PyToolParserOutput>) { + self.0.append(other.0.clone()); + } + + fn coalesce_calls(&self) -> Self { + Self(self.0.clone().coalesce_calls()) + } +} + +#[pyclass(name = "ToolParser", module = "vllm._rust_tool_parser", unsendable)] +struct PyToolParser(Box); + +impl PyToolParser { + fn parse_into_output(&mut self, chunk: &str, output: &mut PyToolParserOutput) -> PyResult<()> { + self.0 + .parse_into(chunk, &mut output.0) + .map_err(|error| PyValueError::new_err(error.to_report_string())) + } +} + +#[pymethods] +impl PyToolParser { + #[new] + fn new(py: Python<'_>, parser_name: &str, tools: Vec>) -> PyResult { + let tools = tools.iter().map(|tool| tool.borrow(py).0.clone()).collect::>(); + create_tool_parser(parser_name, &tools).map(Self) + } + + fn parse_into( + &mut self, + chunk: &str, + mut output: PyRefMut<'_, PyToolParserOutput>, + ) -> PyResult<()> { + self.parse_into_output(chunk, &mut output) + } + + fn finish(&mut self) -> PyResult { + self.0 + .finish() + .map(PyToolParserOutput) + .map_err(|error| PyValueError::new_err(error.to_report_string())) + } + + fn reset(&mut self) -> String { + self.0.reset() + } + + fn preserve_special_tokens(&self) -> bool { + self.0.preserve_special_tokens() + } + + fn tool_call_id(&self, tool_index: usize) -> Option<&str> { + self.0.tool_call_id(tool_index) + } +} + +#[pymodule] +fn _rust_tool_parser(m: &Bound<'_, PyModule>) -> PyResult<()> { + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + fn with_python(f: impl for<'py> FnOnce(Python<'py>) -> R) -> R { + Python::initialize(); + Python::attach(f) + } + + fn tool_schema() -> Value { + json!({ + "type": "object", + "properties": { + "user_id": {"type": "integer"}, + "shipping": { + "type": "object", + "properties": { + "city": {"type": "string"}, + "zip": {"type": "integer"} + } + } + } + }) + } + + fn build_call() -> String { + r#"<｜DSML｜tool_calls> +<｜DSML｜invoke name="create_order"> +<｜DSML｜parameter name="user_id" string="false">42 +<｜DSML｜parameter name="shipping" string="false">{"city":"Singapore","zip":18956} + +"# + .to_owned() + } + + fn make_py_tool(py: Python<'_>) -> PyResult> { + let parameters = pythonize(py, &tool_schema()).map_err(|error| { + PyValueError::new_err(format!( + "failed to convert test schema from JSON to Python: {error}" + )) + })?; + Py::new( + py, + PyTool::new( + "create_order".to_owned(), + Some("Create an order".to_owned()), + ¶meters, + None, + )?, + ) + } + + #[test] + fn tool_round_trips_typed_fields() { + with_python(|py| { + let tool = make_py_tool(py)?; + let borrowed = tool.borrow(py); + assert_eq!(borrowed.name(), "create_order"); + assert_eq!(borrowed.description(), Some("Create an order")); + assert_eq!(borrowed.strict(), None); + + let parameters = borrowed.parameters(py)?; + let parameters = depythonize::(parameters.bind(py))?; + assert_eq!(parameters, tool_schema()); + PyResult::Ok(()) + }) + .unwrap(); + } + + #[test] + fn output_append_and_coalesce_calls() { + with_python(|py| { + let first = Py::new( + py, + PyToolCallDelta::new(0, Some("create_order".to_owned()), "{\"a\"".to_owned()), + )?; + let second = Py::new(py, PyToolCallDelta::new(0, None, ":1}".to_owned()))?; + let mut output = PyToolParserOutput::new(py, "text", Some(vec![first])); + let other = Py::new(py, PyToolParserOutput::new(py, "", Some(vec![second])))?; + output.append(other.borrow(py)); + + let coalesced = output.coalesce_calls(); + assert_eq!(coalesced.normal_text(), "text"); + let calls = coalesced.calls(); + assert_eq!(calls.len(), 1); + assert_eq!(calls[0].tool_index(), 0); + assert_eq!(calls[0].name(), Some("create_order")); + assert_eq!(calls[0].arguments(), "{\"a\":1}"); + PyResult::Ok(()) + }) + .unwrap(); + } + + #[test] + fn parser_parse_finish_and_preserve_special_tokens() { + with_python(|py| { + let tool = make_py_tool(py)?; + let mut parser = PyToolParser::new(py, "DeepSeekV4ToolParser", vec![tool])?; + assert!(parser.preserve_special_tokens()); + + let mut output = PyToolParserOutput::new(py, "", None); + parser.parse_into_output(&build_call(), &mut output)?; + let finish = Py::new(py, parser.finish()?)?; + output.append(finish.borrow(py)); + let output = output.coalesce_calls(); + + assert_eq!(output.normal_text(), ""); + let calls = output.calls(); + assert_eq!(calls.len(), 1); + assert_eq!(calls[0].name(), Some("create_order")); + assert_eq!( + serde_json::from_str::(calls[0].arguments()).unwrap(), + json!({ + "user_id": 42, + "shipping": { + "city": "Singapore", + "zip": 18956 + } + }) + ); + + assert_eq!(parser.reset(), ""); + PyResult::Ok(()) + }) + .unwrap(); + } + + #[test] + fn parser_exposes_model_emitted_tool_call_ids() { + with_python(|py| { + let tool = make_py_tool(py)?; + let mut parser = PyToolParser::new(py, "KimiK2ToolParser", vec![tool])?; + + let input = "<|tool_calls_section_begin|>\ + <|tool_call_begin|>functions.create_order:0<|tool_call_argument_begin|>\ + {\"user_id\":42}<|tool_call_end|>\ + <|tool_calls_section_end|>"; + let mut output = PyToolParserOutput::new(py, "", None); + parser.parse_into_output(input, &mut output)?; + + assert_eq!(parser.tool_call_id(0), Some("functions.create_order:0")); + assert_eq!(parser.tool_call_id(1), None); + PyResult::Ok(()) + }) + .unwrap(); + } + + #[test] + fn parser_errors_for_unknown_name() { + with_python(|py| { + let tool = make_py_tool(py)?; + let error = match PyToolParser::new(py, "missing", vec![tool]) { + Ok(_) => panic!("missing parser name unexpectedly succeeded"), + Err(error) => error, + }; + let message = format!("{error}"); + assert!(message.contains("unsupported tool parser `missing`")); + PyResult::Ok(()) + }) + .unwrap(); + } +} diff --git a/setup.py b/setup.py index d067aae349..1df47b4e7d 100644 --- a/setup.py +++ b/setup.py @@ -35,6 +35,9 @@ ROOT_DIR = Path(__file__).parent logger = logging.getLogger(__name__) PRECOMPILED_RUST_FRONTEND_PATH = ROOT_DIR / "vllm" / "vllm-rs" +# setuptools-rust installs PyO3 artifacts as `.`, where the +# suffix ends with `.so` on Linux and macOS alike (e.g. `_rust_foo.abi3.so`). +PRECOMPILED_RUST_EXTENSION_MEMBER_REGEX = re.compile(r"vllm/_rust_[^/]*\.so$") # cannot import envs directly because it depends on vllm, # which is not installed yet @@ -56,6 +59,25 @@ def should_require_rust_frontend() -> bool: return value.lower() not in ("", "0", "false", "no") +def get_precompiled_rust_extension_paths() -> list[Path]: + return sorted((ROOT_DIR / "vllm").glob("_rust_*.so")) + + +def get_missing_precompiled_rust_extension_modules() -> list[str]: + present = { + path.name.split(".", 1)[0] for path in get_precompiled_rust_extension_paths() + } + return [ + module_name + for module_name in rust_build.rust_py_extension_module_names() + if module_name not in present + ] + + +def has_precompiled_rust_extensions() -> bool: + return not get_missing_precompiled_rust_extension_modules() + + if sys.platform.startswith("darwin") and VLLM_TARGET_DEVICE != "cpu": logger.warning("VLLM_TARGET_DEVICE automatically set to `cpu` due to macOS") VLLM_TARGET_DEVICE = "cpu" @@ -423,19 +445,31 @@ class precompiled_build_ext(build_ext): class precompiled_build_rust(build_rust): - """Skips local Rust builds when the precompiled wheel already ships vllm-rs.""" + """Skips local Rust builds when all precompiled Rust artifacts are present.""" def run(self) -> None: - if PRECOMPILED_RUST_FRONTEND_PATH.exists(): + missing = [] + if not PRECOMPILED_RUST_FRONTEND_PATH.exists(): + missing.append(str(PRECOMPILED_RUST_FRONTEND_PATH)) + missing_rust_extensions = get_missing_precompiled_rust_extension_modules() + if missing_rust_extensions: + missing.extend( + str(ROOT_DIR / "vllm" / f"{module_name}*.so") + for module_name in missing_rust_extensions + ) + + if not missing: logger.info( - "Skipping local Rust build: using precompiled %s", + "Skipping local Rust build: using precompiled %s and %s", PRECOMPILED_RUST_FRONTEND_PATH, + get_precompiled_rust_extension_paths(), ) return logger.warning( - "Precompiled wheel did not provide %s; falling back to local Rust build.", - PRECOMPILED_RUST_FRONTEND_PATH, + "Precompiled wheel did not provide all Rust artifacts (%s); " + "falling back to local Rust build.", + ", ".join(missing), ) super().run() @@ -758,6 +792,14 @@ class precompiled_wheel_utils: if member.filename in exact_members: file_members.append(member) continue + if ( + extract_rust_frontend + and PRECOMPILED_RUST_EXTENSION_MEMBER_REGEX.match( + member.filename + ) + ): + file_members.append(member) + continue if not extract_extensions: continue @@ -1111,6 +1153,12 @@ package_data = { } +def add_vllm_package_data(filename: str) -> None: + vllm_files = package_data.setdefault("vllm", []) + if filename not in vllm_files: + vllm_files.append(filename) + + # If using precompiled artifacts, extract and patch package_data in advance. if USE_PRECOMPILED_RUST_FRONTEND: wheel_url, download_filename = precompiled_wheel_utils.determine_wheel_url() @@ -1126,9 +1174,9 @@ if USE_PRECOMPILED_RUST_FRONTEND: # If the rust frontend binary is already present in the source tree (e.g., # pre-built in a separate Docker build stage), ship it as-is. if PRECOMPILED_RUST_FRONTEND_PATH.exists(): - vllm_files = package_data.setdefault("vllm", []) - if "vllm-rs" not in vllm_files: - vllm_files.append("vllm-rs") + add_vllm_package_data("vllm-rs") +for rust_extension_path in get_precompiled_rust_extension_paths(): + add_vllm_package_data(rust_extension_path.name) if _no_device(): ext_modules = [] @@ -1141,13 +1189,15 @@ else: if USE_PRECOMPILED_EXTENSIONS else cmake_build_ext, } -if USE_PRECOMPILED_RUST_FRONTEND or PRECOMPILED_RUST_FRONTEND_PATH.exists(): +if ( + USE_PRECOMPILED_RUST_FRONTEND + or PRECOMPILED_RUST_FRONTEND_PATH.exists() + or has_precompiled_rust_extensions() +): cmdclass["build_rust"] = precompiled_build_rust -# Rust frontend binary, built via setuptools-rust and installed into the -# package directory alongside the Python modules. -# TODO: we may use `RustBin` to directly install it into `bin` directory, but this -# requires extra work on using precompiled binaries. +# Rust artifacts, built via setuptools-rust and installed into the package +# directory alongside the Python modules. rust_extensions = rust_build.rust_extensions( optional=not should_require_rust_frontend() ) diff --git a/tests/tool_parsers/test_rust_tool_parser.py b/tests/tool_parsers/test_rust_tool_parser.py new file mode 100644 index 0000000000..7546848778 --- /dev/null +++ b/tests/tool_parsers/test_rust_tool_parser.py @@ -0,0 +1,328 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import json +from collections.abc import Sequence +from unittest.mock import MagicMock + +import pytest + +from vllm.entrypoints.openai.chat_completion.protocol import ( + ChatCompletionRequest, + ChatCompletionToolsParam, +) +from vllm.tool_parsers.rust_tool_parser import RustToolParser + +# The PyO3 extension is an optional build artifact; skip when absent. +_rust_tool_parser = pytest.importorskip("vllm._rust_tool_parser") + +MOCK_TOKENIZER = MagicMock() +MOCK_TOKENIZER.get_vocab.return_value = {} + +TC_START = "<｜DSML｜tool_calls>" +TC_END = "" +INV_START = '<｜DSML｜invoke name="' +INV_END = "" +PARAM_START = '<｜DSML｜parameter name="' +PARAM_END = "" + + +class DeepSeekV4RustToolParser(RustToolParser): + rust_parser_name = "DeepSeekV4ToolParser" + tool_call_start_token = TC_START + + +class KimiK2RustToolParser(RustToolParser): + rust_parser_name = "KimiK2ToolParser" + tool_call_start_token = "<|tool_calls_section_begin|>" + + +def sample_tools() -> list[ChatCompletionToolsParam]: + return [ + ChatCompletionToolsParam( + type="function", + function={ + "name": "get_weather", + "description": "Get weather for a location", + "parameters": { + "type": "object", + "properties": { + "location": {"type": "string"}, + "date": {"type": "string"}, + }, + }, + }, + ), + ChatCompletionToolsParam( + type="function", + function={ + "name": "add", + "description": "Add two integers", + "parameters": { + "type": "object", + "properties": { + "x": {"type": "integer"}, + "y": {"type": "integer"}, + }, + }, + }, + ), + ] + + +EXPECTED_CALLS = [ + ("get_weather", {"location": "SF", "date": "2024-01-16"}), + ("add", {"x": 3, "y": 5}), +] + + +def build_invoke( + function_name: str, + params: Sequence[tuple[str, str, bool]], +) -> str: + param_text = "\n".join( + f'{PARAM_START}{name}" string="{str(is_string).lower()}">{value}{PARAM_END}' + for name, value, is_string in params + ) + return f'{INV_START}{function_name}">\n{param_text}\n{INV_END}\n' + + +def build_tool_call() -> str: + weather = build_invoke( + "get_weather", + [ + ("location", "SF", True), + ("date", "2024-01-16", True), + ], + ) + add = build_invoke( + "add", + [ + ("x", "3", False), + ("y", "5", False), + ], + ) + return f"{TC_START}\n{weather}{add}{TC_END}" + + +def parse_streaming( + parser: DeepSeekV4RustToolParser, + text: str, + chunk_size: int, +) -> list: + deltas = [] + previous_text = "" + for start in range(0, len(text), chunk_size): + delta_text = text[start : start + chunk_size] + current_text = previous_text + delta_text + delta = parser.extract_tool_calls_streaming( + previous_text=previous_text, + current_text=current_text, + delta_text=delta_text, + previous_token_ids=[], + current_token_ids=[], + delta_token_ids=[1], + request=MagicMock(), + ) + previous_text = current_text + if delta is not None: + deltas.append(delta) + + delta = parser.extract_tool_calls_streaming( + previous_text=previous_text, + current_text=previous_text, + delta_text="", + previous_token_ids=[], + current_token_ids=[], + delta_token_ids=[2], + request=MagicMock(), + ) + if delta is not None: + deltas.append(delta) + + return deltas + + +def collect_streamed_arguments(deltas: Sequence, tool_index: int = 0) -> str: + return "".join( + tool_call.function.arguments + for delta in deltas + for tool_call in delta.tool_calls or [] + if ( + tool_call.index == tool_index + and tool_call.function is not None + and tool_call.function.arguments is not None + ) + ) + + +def test_rust_tool_parser_extension_typed_api() -> None: + tools = [ + _rust_tool_parser.Tool( + tool.function.name, + tool.function.description, + tool.function.parameters, + None, + ) + for tool in sample_tools() + ] + parser = _rust_tool_parser.ToolParser("DeepSeekV4ToolParser", tools) + output = _rust_tool_parser.ToolParserOutput() + + parser.parse_into(build_tool_call(), output) + output.append(parser.finish()) + output = output.coalesce_calls() + + assert parser.preserve_special_tokens() + assert output.normal_text == "" + assert len(output.calls) == 2 + for call, (name, arguments) in zip(output.calls, EXPECTED_CALLS): + assert call.name == name + assert json.loads(call.arguments) == arguments + + +def test_rust_tool_parser_adapter_extracts_complete_output() -> None: + tools = sample_tools() + parser = DeepSeekV4RustToolParser(MOCK_TOKENIZER, tools=tools) + + result = parser.extract_tool_calls( + "Let me create it. " + build_tool_call(), + ChatCompletionRequest(messages=[], model="m", tools=tools), + ) + + assert result.tools_called + assert result.content == "Let me create it. " + assert len(result.tool_calls) == 2 + for tool_call, (name, arguments) in zip(result.tool_calls, EXPECTED_CALLS): + assert tool_call.function.name == name + assert json.loads(tool_call.function.arguments) == arguments + + +def test_rust_tool_parser_adapter_streaming_handles_multiple_calls() -> None: + parser = DeepSeekV4RustToolParser(MOCK_TOKENIZER, tools=sample_tools()) + + deltas = parse_streaming(parser, build_tool_call(), chunk_size=5) + + names = [ + tool_call.function.name + for delta in deltas + for tool_call in delta.tool_calls or [] + if tool_call.function is not None and tool_call.function.name is not None + ] + assert names == [name for name, _ in EXPECTED_CALLS] + for index, (_, arguments) in enumerate(EXPECTED_CALLS): + assert json.loads(collect_streamed_arguments(deltas, index)) == arguments + + +def test_rust_tool_parser_adapter_ignores_midstream_empty_delta() -> None: + parser = DeepSeekV4RustToolParser(MOCK_TOKENIZER, tools=sample_tools()) + text = build_tool_call() + split_at = len(TC_START) + 8 + deltas = [] + previous_text = "" + + for delta_text in (text[:split_at], "", text[split_at:], ""): + current_text = previous_text + delta_text + delta = parser.extract_tool_calls_streaming( + previous_text=previous_text, + current_text=current_text, + delta_text=delta_text, + previous_token_ids=[], + current_token_ids=[], + delta_token_ids=[1], + request=MagicMock(), + ) + previous_text = current_text + if delta is not None: + deltas.append(delta) + + names = [ + tool_call.function.name + for delta in deltas + for tool_call in delta.tool_calls or [] + if tool_call.function is not None and tool_call.function.name is not None + ] + assert names == [name for name, _ in EXPECTED_CALLS] + for index, (_, arguments) in enumerate(EXPECTED_CALLS): + assert json.loads(collect_streamed_arguments(deltas, index)) == arguments + + +KIMI_EXPECTED_IDS = ["functions.get_weather:0", "functions.add:1"] + + +def build_kimi_tool_call() -> str: + return ( + "<|tool_calls_section_begin|>" + "<|tool_call_begin|>functions.get_weather:0<|tool_call_argument_begin|>" + '{"location": "SF", "date": "2024-01-16"}<|tool_call_end|>' + "<|tool_call_begin|>functions.add:1<|tool_call_argument_begin|>" + '{"x": 3, "y": 5}<|tool_call_end|>' + "<|tool_calls_section_end|>" + ) + + +def test_rust_tool_parser_adapter_complete_prefers_model_tool_call_ids() -> None: + tools = sample_tools() + parser = KimiK2RustToolParser(MOCK_TOKENIZER, tools=tools) + + result = parser.extract_tool_calls( + "Let me check. " + build_kimi_tool_call(), + ChatCompletionRequest(messages=[], model="m", tools=tools), + ) + + assert result.tools_called + assert [tool_call.id for tool_call in result.tool_calls] == KIMI_EXPECTED_IDS + for tool_call, (name, arguments) in zip(result.tool_calls, EXPECTED_CALLS): + assert tool_call.function.name == name + assert json.loads(tool_call.function.arguments) == arguments + + +def test_rust_tool_parser_adapter_streaming_prefers_model_tool_call_ids() -> None: + parser = KimiK2RustToolParser(MOCK_TOKENIZER, tools=sample_tools()) + + deltas = parse_streaming(parser, build_kimi_tool_call(), chunk_size=5) + + ids = [ + tool_call.id + for delta in deltas + for tool_call in delta.tool_calls or [] + if tool_call.id is not None + ] + assert ids == KIMI_EXPECTED_IDS + for index, (_, arguments) in enumerate(EXPECTED_CALLS): + assert json.loads(collect_streamed_arguments(deltas, index)) == arguments + + +def test_rust_tool_parser_adapter_streaming_generates_ids_as_fallback() -> None: + # DeepSeekV4 never emits model tool call IDs, so the bridge mints them. + parser = DeepSeekV4RustToolParser(MOCK_TOKENIZER, tools=sample_tools()) + + deltas = parse_streaming(parser, build_tool_call(), chunk_size=5) + + ids = [ + tool_call.id + for delta in deltas + for tool_call in delta.tool_calls or [] + if tool_call.function is not None and tool_call.function.name is not None + ] + assert len(ids) == len(EXPECTED_CALLS) + assert all(ids) + assert len(set(ids)) == len(ids) + + +def test_rust_tool_parser_adapter_adjust_request_is_opaque() -> None: + tools = sample_tools() + parser = DeepSeekV4RustToolParser(MOCK_TOKENIZER, tools=tools) + request = ChatCompletionRequest( + messages=[], + model="m", + tools=tools, + tool_choice="required", + skip_special_tokens=True, + ) + + adjusted = parser.adjust_request(request) + + assert adjusted is request + assert adjusted.skip_special_tokens is False + assert adjusted.structured_outputs is None diff --git a/tools/build_rust.py b/tools/build_rust.py index 169e636ccb..e5c5d0bb2e 100644 --- a/tools/build_rust.py +++ b/tools/build_rust.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Shared setuptools-rust build entry for the vllm-rs binary.""" +"""Shared setuptools-rust build entry for Rust artifacts.""" from __future__ import annotations @@ -15,7 +15,7 @@ from setuptools_rust import Binding, RustExtension ROOT_DIR = Path(__file__).resolve().parents[1] -def rust_extensions(*, optional: bool) -> list[RustExtension]: +def rust_extensions(*, optional: bool = False) -> list[RustExtension]: return [ RustExtension( target="vllm.vllm-rs", @@ -25,9 +25,30 @@ def rust_extensions(*, optional: bool) -> list[RustExtension]: binding=Binding.Exec, optional=optional, ), + RustExtension( + target="vllm._rust_tool_parser", + path="rust/src/tool-parser/python/Cargo.toml", + features=["pyo3/abi3-py38"], + binding=Binding.PyO3, + optional=optional, + py_limited_api=True, + ), ] +def rust_py_extension_module_names() -> list[str]: + module_names = [] + for extension in rust_extensions(): + if extension.binding != Binding.PyO3: + continue + + for target_name in extension.target.values(): + if target_name.startswith("vllm._rust_"): + module_names.append(target_name.rsplit(".", 1)[-1]) + + return module_names + + def build_binary(build_rust_args: list[str]) -> None: os.chdir(ROOT_DIR) (ROOT_DIR / "vllm").mkdir(exist_ok=True) diff --git a/vllm/tool_parsers/rust_tool_parser.py b/vllm/tool_parsers/rust_tool_parser.py new file mode 100644 index 0000000000..493f765a2c --- /dev/null +++ b/vllm/tool_parsers/rust_tool_parser.py @@ -0,0 +1,322 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import importlib +from collections.abc import Sequence +from typing import Any + +from openai.types.responses.function_tool import FunctionTool + +from vllm.entrypoints.chat_utils import make_tool_call_id +from vllm.entrypoints.openai.chat_completion.protocol import ( + ChatCompletionRequest, + ChatCompletionToolsParam, +) +from vllm.entrypoints.openai.engine.protocol import ( + DeltaFunctionCall, + DeltaMessage, + DeltaToolCall, + ExtractedToolCallInformation, + FunctionCall, + ToolCall, +) +from vllm.entrypoints.openai.responses.protocol import ResponsesRequest +from vllm.logger import init_logger +from vllm.tokenizers import TokenizerLike +from vllm.tool_parsers.abstract_tool_parser import Tool, ToolParser + +logger = init_logger(__name__) + + +def _rust_tool_parser_module() -> Any: + try: + return importlib.import_module("vllm._rust_tool_parser") + except ImportError as exc: + raise RuntimeError( + "Rust tool parsing requires the vllm._rust_tool_parser PyO3 " + "extension. Rebuild vLLM with Rust frontend/extensions enabled." + ) from exc + + +class RustToolParser(ToolParser): + """Adapter from an opaque Rust parser to the vLLM ToolParser API. + + Subclasses provide only model-specific configuration: the exact Rust parser + name and an optional tool-call start marker for fast complete-output + rejection. + + This class keeps the vLLM-specific bridge work: + - convert vLLM tool definitions into the Rust ``Tool`` shape; + - translate typed Rust parser outputs into vLLM protocol objects; and + - maintain vLLM streaming bookkeeping used by finish-reason handling. + + The parser grammar and incremental parser state stay in Rust. + """ + + # Rust-backed parsers are opaque to Python by default. Do not use vLLM's + # standard JSON required/named handling; let the Rust parser consume the + # model's native tool-call syntax. + supports_required_and_named = False + + rust_parser_name: str + tool_call_start_token: str | None = None + + def __init__(self, tokenizer: TokenizerLike, tools: list[Tool] | None = None): + super().__init__(tokenizer, tools) + self._parser: Any | None = None + self._error: Exception | None = None + + if not self.model_tokenizer: + raise ValueError( + "The model tokenizer must be passed to the ToolParser " + "constructor during construction." + ) + + logger.debug( + "vLLM successfully imported tool parser %s", self.__class__.__name__ + ) + + def adjust_request( + self, request: ChatCompletionRequest | ResponsesRequest + ) -> ChatCompletionRequest | ResponsesRequest: + """Adjust request options without installing Python-side constraints. + + Rust-backed parsers are treated as source-of-truth opaque parsers. The + bridge intentionally avoids ``super().adjust_request()`` so Python does + not install JSON schema guidance or structural-tag constraints that may + conflict with the Rust parser's native grammar. + """ + if self._get_parser().preserve_special_tokens(): + request.skip_special_tokens = False + return request + + def _rust_tools(self) -> list[Any]: + """Build Rust ``Tool`` objects from vLLM tool definitions.""" + if not self.tools: + return [] + + tools: list[Any] = [] + for tool in self.tools: + if isinstance(tool, FunctionTool): + name = tool.name + description = tool.description + parameters = tool.parameters or {} + strict = getattr(tool, "strict", None) + elif isinstance(tool, ChatCompletionToolsParam): + name = tool.function.name + description = tool.function.description + parameters = tool.function.parameters or {} + strict = getattr(tool.function, "strict", None) + else: + continue + tools.append( + _rust_tool_parser_module().Tool(name, description, parameters, strict) + ) + return tools + + def _new_parser(self) -> Any: + """Create a fresh Rust parser with the current tool schemas.""" + return _rust_tool_parser_module().ToolParser( + self.rust_parser_name, self._rust_tools() + ) + + def _get_parser(self) -> Any: + if self._parser is None: + self._parser = self._new_parser() + return self._parser + + def _reset_streaming_state(self) -> None: + """Reset parser state for a new request on a reused parser instance.""" + self._parser = self._new_parser() + self._error = None + self.prev_tool_call_arr.clear() + self.streamed_args_for_tool.clear() + self.current_tool_id = -1 + self.current_tool_name_sent = False + + def _ensure_tool_state(self, index: int) -> None: + """Grow vLLM streaming state arrays to contain ``index``.""" + while len(self.prev_tool_call_arr) <= index: + self.prev_tool_call_arr.append({}) + while len(self.streamed_args_for_tool) <= index: + self.streamed_args_for_tool.append("") + + def _record_delta( + self, index: int, name: str | None, arguments: str | None + ) -> str | None: + """Mirror a Rust parser delta into vLLM streaming bookkeeping. + + ``prev_tool_call_arr`` and ``streamed_args_for_tool`` are read later by + the chat serving layer to decide the final ``tool_calls`` finish reason + and to flush any remaining argument bytes. + """ + tool_call_id = None + self._ensure_tool_state(index) + + if name is not None: + # Prefer the model-emitted ID surfaced by the Rust parser (e.g. + # Kimi K2) over a randomly generated one. + tool_call_id = self._get_parser().tool_call_id(index) or make_tool_call_id() + self.prev_tool_call_arr[index] = {"name": name, "arguments": {}} + self.current_tool_name_sent = True + + if arguments is not None: + self.streamed_args_for_tool[index] += arguments + self.prev_tool_call_arr[index]["arguments"] = self.streamed_args_for_tool[ + index + ] + self.current_tool_id = index + + return tool_call_id + + def _delta_message_from_parser_output( + self, parser_output: Any | None + ) -> DeltaMessage | None: + """Translate one Rust parser output into a vLLM ``DeltaMessage``.""" + if parser_output is None: + return None + + normal_text = parser_output.normal_text or None + tool_calls: list[DeltaToolCall] = [] + for tool_call in parser_output.calls: + index = tool_call.tool_index + name = tool_call.name + arguments: str | None = tool_call.arguments + if name is None and arguments is None: + continue + + tool_call_id = self._record_delta(index, name, arguments) + tool_calls.append( + DeltaToolCall( + index=index, + id=tool_call_id, + type="function" if name is not None else None, + function=DeltaFunctionCall( + name=name, + arguments=arguments, + ), + ) + ) + + if normal_text is None and not tool_calls: + return None + return DeltaMessage(content=normal_text, tool_calls=tool_calls) + + def _parse_complete(self, model_output: str) -> tuple[Any, dict[int, str]] | None: + """Parse complete model output with a throwaway Rust parser instance. + + Returns the coalesced parser output along with any model-emitted tool + call IDs keyed by tool index. + """ + parser = self._new_parser() + output = _rust_tool_parser_module().ToolParserOutput() + try: + parser.parse_into(model_output, output) + # finish() clears parser state, so snapshot model-emitted IDs first. + tool_call_ids = { + call.tool_index: tool_call_id + for call in output.calls + if (tool_call_id := parser.tool_call_id(call.tool_index)) is not None + } + output.append(parser.finish()) + except Exception: + logger.exception( + "Error parsing %s tool call output.", self.rust_parser_name + ) + return None + return output.coalesce_calls(), tool_call_ids + + def extract_tool_calls( + self, + model_output: str, + request: ChatCompletionRequest, + ) -> ExtractedToolCallInformation: + """Extract tool calls from complete model output (non-streaming).""" + if ( + self.tool_call_start_token is not None + and self.tool_call_start_token not in model_output + ): + return ExtractedToolCallInformation( + tools_called=False, + tool_calls=[], + content=model_output, + ) + + parse_result = self._parse_complete(model_output) + if parse_result is None: + return ExtractedToolCallInformation( + tools_called=False, + tool_calls=[], + content=model_output, + ) + parsed, tool_call_ids = parse_result + + tool_calls: list[ToolCall] = [] + self.prev_tool_call_arr.clear() + for parsed_tool_call in parsed.calls: + name = parsed_tool_call.name + arguments = parsed_tool_call.arguments or "{}" + if name is None: + continue + tool_calls.append( + ToolCall( + id=tool_call_ids.get(parsed_tool_call.tool_index) + or make_tool_call_id(), + type="function", + function=FunctionCall(name=name, arguments=arguments), + ) + ) + self.prev_tool_call_arr.append({"name": name, "arguments": arguments}) + + if not tool_calls: + return ExtractedToolCallInformation( + tools_called=False, + tool_calls=[], + content=model_output, + ) + + content = parsed.normal_text or None + return ExtractedToolCallInformation( + tools_called=True, + tool_calls=tool_calls, + content=content, + ) + + def extract_tool_calls_streaming( + self, + previous_text: str, + current_text: str, + delta_text: str, + previous_token_ids: Sequence[int], # pylint: disable=unused-argument + current_token_ids: Sequence[int], # pylint: disable=unused-argument + delta_token_ids: Sequence[int], # pylint: disable=unused-argument + request: ChatCompletionRequest, # pylint: disable=unused-argument + ) -> DeltaMessage | None: + """Extract tool calls from streaming model output. + + The Rust parser owns the incremental buffer, so this adapter feeds only + the newest text delta and lets the serving layer handle final empty + chunks. + """ + # TODO: Add a final-chunk hook if streaming needs to call Rust finish(). + if not previous_text: + self._reset_streaming_state() + + if self._error is not None: + return None + + parser_output = _rust_tool_parser_module().ToolParserOutput() + try: + self._get_parser().parse_into(delta_text, parser_output) + except Exception as error: + self._error = error + logger.exception( + "Error parsing %s streaming tool call output.", + self.rust_parser_name, + ) + + delta_message = self._delta_message_from_parser_output(parser_output) + if delta_message is not None: + return delta_message + + return None