Compare commits

...
Author SHA1 Message Date
Bugen Zhao efa494e397 refactor setup.py to extract stuff into build_rust.py
Signed-off-by: Bugen Zhao <i@bugenzhao.com>
2026-06-10 22:39:17 +08:00
Bugen Zhao 74b77a593a update dockerfile image
Signed-off-by: Bugen Zhao <i@bugenzhao.com>
2026-06-10 22:38:24 +08:00
Bugen Zhao ec53889a3b simplify to only match .so, as .dylib will also be renamed to .so under macos
Signed-off-by: Bugen Zhao <i@bugenzhao.com>
2026-06-10 22:13:09 +08:00
Bugen Zhao ede3d4ddf6 skip tests when _rust_tool_parser is absent
Signed-off-by: Bugen Zhao <i@bugenzhao.com>
2026-06-10 22:12:48 +08:00
Bugen Zhao 7301a834fa fix glob issue
Signed-off-by: Bugen Zhao <i@bugenzhao.com>
2026-06-10 22:06:14 +08:00
Bugen Zhao 7104ac6b5d simplify setup.py
Signed-off-by: Bugen Zhao <i@bugenzhao.com>
2026-06-10 11:40:33 +00:00
Bugen Zhao 9a98eccb49 cover two calls case
Signed-off-by: Bugen Zhao <i@bugenzhao.com>
2026-06-10 09:50:02 +00:00
Bugen Zhao 2ce1e2ba71 do not infer finished but just leave it unhandled
Signed-off-by: Bugen Zhao <i@bugenzhao.com>
2026-06-10 09:50:02 +00:00
Bugen Zhao 2496f66f7f Run tool parser tests before tokenizer tests
Signed-off-by: Bugen Zhao <i@bugenzhao.com>
2026-06-10 09:50:02 +00:00
Bugen Zhao d18b7f2723 Build Rust parser extension in Docker rust stage
Signed-off-by: Bugen Zhao <i@bugenzhao.com>
2026-06-10 09:50:02 +00:00
Bugen Zhao 80b633f076 ci: configure PyO3 Python for Rust tests
Signed-off-by: Bugen Zhao <i@bugenzhao.com>
2026-06-10 09:50:02 +00:00
Bugen Zhao 28f589ebb1 do not enable pyo3 on --all-features
Signed-off-by: Bugen Zhao <i@bugenzhao.com>
2026-06-10 09:50:02 +00:00
Bugen Zhao 98998279d4 fix fmt & try import
Signed-off-by: Bugen Zhao <i@bugenzhao.com>
2026-06-10 09:50:01 +00:00
Bugen Zhao 71d208089e add tests with deepseek v4
Signed-off-by: Bugen Zhao <i@bugenzhao.com>
2026-06-10 09:50:01 +00:00
Bugen Zhao 0ac6fc8352 Add Rust tool parser Python bridge
Signed-off-by: Bugen Zhao <i@bugenzhao.com>
2026-06-10 09:50:01 +00:00
17 changed files with 1226 additions and 70 deletions
@@ -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"
+2 -2
View File
@@ -300,9 +300,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
@@ -315,9 +315,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
+1 -1
View File
@@ -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
+8 -5
View File
@@ -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 \
+8 -5
View File
@@ -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
+5 -3
View File
@@ -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
+6 -4
View File
@@ -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
+5 -3
View File
@@ -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 \
Binary file not shown.

Before

Width:  |  Height:  |  Size: 388 KiB

After

Width:  |  Height:  |  Size: 373 KiB

+86
View File
@@ -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"
+3
View File
@@ -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"] }
+19
View File
@@ -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
+366
View File
@@ -0,0 +1,366 @@
//! 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<Box<dyn ToolParser>> {
match name {
$(
stringify!($parser) => {
<vllm_tool_parser::$parser as ToolParser>::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! {
DeepSeekV4ToolParser, // for testing on Python side
}
#[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<String>,
parameters: &Bound<'_, PyAny>,
strict: Option<bool>,
) -> PyResult<Self> {
let parameters = depythonize::<Value>(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<Py<PyAny>> {
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<bool> {
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<String>, 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<Vec<Py<PyToolCallDelta>>>) -> 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<PyToolCallDelta> {
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<dyn ToolParser>);
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<Py<PyTool>>) -> PyResult<Self> {
let tools = tools.iter().map(|tool| tool.borrow(py).0.clone()).collect::<Vec<_>>();
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<PyToolParserOutput> {
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()
}
}
#[pymodule]
fn _rust_tool_parser(m: &Bound<'_, PyModule>) -> PyResult<()> {
m.add_class::<PyTool>()?;
m.add_class::<PyToolCallDelta>()?;
m.add_class::<PyToolParserOutput>()?;
m.add_class::<PyToolParser>()?;
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
fn with_python<R>(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#"<DSMLtool_calls>
<DSMLinvoke name="create_order">
<DSMLparameter name="user_id" string="false">42</DSMLparameter>
<DSMLparameter name="shipping" string="false">{"city":"Singapore","zip":18956}</DSMLparameter>
</DSMLinvoke>
</DSMLtool_calls>"#
.to_owned()
}
fn make_py_tool(py: Python<'_>) -> PyResult<Py<PyTool>> {
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()),
&parameters,
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::<Value>(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::<Value>(calls[0].arguments()).unwrap(),
json!({
"user_id": 42,
"shipping": {
"city": "Singapore",
"zip": 18956
}
})
);
assert_eq!(parser.reset(), "");
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();
}
}
+21 -42
View File
@@ -18,7 +18,6 @@ import torch
from packaging.version import Version, parse
from setuptools import Extension, setup
from setuptools.command.build_ext import build_ext
from setuptools_rust.build import build_rust
from setuptools_scm import get_version
from torch.utils.cpp_extension import CUDA_HOME, ROCM_HOME
@@ -34,8 +33,6 @@ def load_module_from_path(module_name, path):
ROOT_DIR = Path(__file__).parent
logger = logging.getLogger(__name__)
PRECOMPILED_RUST_FRONTEND_PATH = ROOT_DIR / "vllm" / "vllm-rs"
# cannot import envs directly because it depends on vllm,
# which is not installed yet
envs = load_module_from_path("envs", os.path.join(ROOT_DIR, "vllm", "envs.py"))
@@ -51,11 +48,6 @@ USE_PRECOMPILED_RUST_FRONTEND = (
)
def should_require_rust_frontend() -> bool:
value = os.getenv("VLLM_REQUIRE_RUST_FRONTEND", "")
return value.lower() not in ("", "0", "false", "no")
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"
@@ -422,24 +414,6 @@ class precompiled_build_ext(build_ext):
return
class precompiled_build_rust(build_rust):
"""Skips local Rust builds when the precompiled wheel already ships vllm-rs."""
def run(self) -> None:
if PRECOMPILED_RUST_FRONTEND_PATH.exists():
logger.info(
"Skipping local Rust build: using precompiled %s",
PRECOMPILED_RUST_FRONTEND_PATH,
)
return
logger.warning(
"Precompiled wheel did not provide %s; falling back to local Rust build.",
PRECOMPILED_RUST_FRONTEND_PATH,
)
super().run()
class precompiled_wheel_utils:
"""Extracts libraries and other files from an existing wheel."""
@@ -733,9 +707,6 @@ class precompiled_wheel_utils:
"vllm/_rocm_C.abi3.so",
}
)
if extract_rust_frontend:
exact_members.add("vllm/vllm-rs")
flash_attn_regex = re.compile(
r"vllm/vllm_flash_attn/(?:[^/.][^/]*/)*(?!\.)[^/]*\.py"
)
@@ -758,6 +729,12 @@ class precompiled_wheel_utils:
if member.filename in exact_members:
file_members.append(member)
continue
if (
extract_rust_frontend
and rust_build.is_precompiled_artifact_member(member.filename)
):
file_members.append(member)
continue
if not extract_extensions:
continue
@@ -1111,6 +1088,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()
@@ -1123,12 +1106,10 @@ if USE_PRECOMPILED_RUST_FRONTEND:
for pkg, files in patch.items():
package_data.setdefault(pkg, []).extend(files)
# 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")
# Rust artifacts already present in the source tree (e.g., pre-built in a
# separate Docker build stage) are shipped as-is.
for rust_artifact in rust_build.find_precompiled_artifacts():
add_vllm_package_data(rust_artifact.name)
if _no_device():
ext_modules = []
@@ -1141,15 +1122,13 @@ else:
if USE_PRECOMPILED_EXTENSIONS
else cmake_build_ext,
}
if USE_PRECOMPILED_RUST_FRONTEND or PRECOMPILED_RUST_FRONTEND_PATH.exists():
cmdclass["build_rust"] = precompiled_build_rust
if USE_PRECOMPILED_RUST_FRONTEND or rust_build.find_precompiled_artifacts():
cmdclass["build_rust"] = rust_build.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()
optional=not rust_build.should_require_rust_frontend()
)
setup(
+260
View File
@@ -0,0 +1,260 @@
# 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 = "<DSMLtool_calls>"
TC_END = "</DSMLtool_calls>"
INV_START = '<DSMLinvoke name="'
INV_END = "</DSMLinvoke>"
PARAM_START = '<DSMLparameter name="'
PARAM_END = "</DSMLparameter>"
class DeepSeekV4RustToolParser(RustToolParser):
rust_parser_name = "DeepSeekV4ToolParser"
tool_call_start_token = TC_START
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
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
+95 -5
View File
@@ -1,21 +1,31 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Shared setuptools-rust build entry for the vllm-rs binary."""
"""Rust build support shared by `setup.py` and the standalone `build_rust.sh`.
This module is the single source of truth for the Rust artifacts shipped in
the vllm package: which crates are built, where their artifacts land, and how
precompiled artifacts are detected.
"""
from __future__ import annotations
import logging
import os
import sys
from pathlib import Path
from setuptools import setup
from setuptools_rust import Binding, RustExtension
from setuptools_rust.build import build_rust
logger = logging.getLogger(__name__)
ROOT_DIR = Path(__file__).resolve().parents[1]
PACKAGE_DIR = ROOT_DIR / "vllm"
def rust_extensions(*, optional: bool) -> list[RustExtension]:
def rust_extensions(*, optional: bool = False) -> list[RustExtension]:
return [
RustExtension(
target="vllm.vllm-rs",
@@ -25,12 +35,92 @@ 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 build_binary(build_rust_args: list[str]) -> None:
def should_require_rust_frontend() -> bool:
value = os.getenv("VLLM_REQUIRE_RUST_FRONTEND", "")
return value.lower() not in ("", "0", "false", "no")
def _expected_artifacts() -> list[tuple[str, Binding]]:
"""(basename, binding) of each artifact installed into the package."""
artifacts = []
for extension in rust_extensions():
for target in extension.target.values():
package, _, name = target.rpartition(".")
assert package == "vllm", f"unexpected Rust target: {target}"
artifacts.append((name, extension.binding))
return artifacts
def _is_artifact_file(filename: str, name: str, binding: Binding) -> bool:
# setuptools-rust installs Exec binaries under their bare name, and PyO3
# modules as `<module>.<ext-suffix>` where the suffix ends with `.so` on
# Linux and macOS alike (e.g. `_rust_foo.abi3.so`).
if binding == Binding.Exec:
return filename == name
return filename.endswith(".so") and filename.split(".", 1)[0] == name
def find_precompiled_artifacts() -> list[Path]:
"""Rust artifacts already present in the package directory."""
return sorted(
path
for path in PACKAGE_DIR.iterdir()
if any(_is_artifact_file(path.name, *spec) for spec in _expected_artifacts())
)
def missing_precompiled_artifacts() -> list[str]:
"""Expected-but-absent artifacts, as file patterns for diagnostics."""
present = [path.name for path in find_precompiled_artifacts()]
return [
str(PACKAGE_DIR / (name if binding == Binding.Exec else f"{name}.*.so"))
for name, binding in _expected_artifacts()
if not any(_is_artifact_file(filename, name, binding) for filename in present)
]
def is_precompiled_artifact_member(member_name: str) -> bool:
"""Whether a wheel member is a Rust artifact (e.g. `vllm/vllm-rs`)."""
package, _, filename = member_name.rpartition("/")
return package == "vllm" and any(
_is_artifact_file(filename, *spec) for spec in _expected_artifacts()
)
class precompiled_build_rust(build_rust):
"""Skips the local Rust build when all precompiled artifacts are present."""
def run(self) -> None:
missing = missing_precompiled_artifacts()
if not missing:
logger.info(
"Skipping local Rust build: using precompiled %s",
find_precompiled_artifacts(),
)
return
logger.warning(
"Precompiled Rust artifacts missing (%s); "
"falling back to local Rust build.",
", ".join(missing),
)
super().run()
def build_artifacts(build_rust_args: list[str]) -> None:
os.chdir(ROOT_DIR)
(ROOT_DIR / "vllm").mkdir(exist_ok=True)
PACKAGE_DIR.mkdir(exist_ok=True)
setup(
name="vllm-rust-frontend-build",
packages=[],
@@ -40,7 +130,7 @@ def build_binary(build_rust_args: list[str]) -> None:
def main() -> None:
build_binary(sys.argv[1:])
build_artifacts(sys.argv[1:])
if __name__ == "__main__":
+310
View File
@@ -0,0 +1,310 @@
# 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
self._tool_call_ids: dict[int, str] = {}
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._tool_call_ids.clear()
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:
tool_call_id = make_tool_call_id()
self._tool_call_ids[index] = 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) -> Any | None:
"""Parse complete model output with a throwaway Rust parser instance."""
parser = self._new_parser()
output = _rust_tool_parser_module().ToolParserOutput()
try:
parser.parse_into(model_output, output)
output.append(parser.finish())
except Exception:
logger.exception(
"Error parsing %s tool call output.", self.rust_parser_name
)
return None
return output.coalesce_calls()
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,
)
parsed = self._parse_complete(model_output)
if parsed is None:
return ExtractedToolCallInformation(
tools_called=False,
tool_calls=[],
content=model_output,
)
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(
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