[Rust Frontend] Remove the PyO3 tool-parser bridge (#59744)

Signed-off-by: Bugen Zhao <i@bugenzhao.com>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
Bugen ZhaoandClaude Opus 5.5 authored and GitHub committed 2026-10-06 05:46:27 +00:00
1 parent 5145e35e50
commit b3082d3044
9 files changed
+13 -1193

No files matched your search

@@ -23,7 +23,7 @@ export PATH="$CARGO_HOME/bin:$PATH"
CARGO_BINSTALL_VERSION="${CARGO_BINSTALL_VERSION:-1.20.1}"
UV_VERSION="${UV_VERSION:-0.11.28}"
PYO3_PYTHON_VERSION="${PYO3_PYTHON_VERSION:-3.12}"
UV_SCRIPT_PYTHON_VERSION="${UV_SCRIPT_PYTHON_VERSION:-3.12}"
CARGO_SORT_VERSION_REQ="${CARGO_SORT_VERSION_REQ:-2}"
CARGO_DENY_VERSION_REQ="${CARGO_DENY_VERSION_REQ:-0.20}"
@@ -83,34 +83,7 @@ install_uv() {
"https://github.com/astral-sh/uv/releases/download/${UV_VERSION}/uv-installer.sh" \
| env UV_INSTALL_DIR="$CARGO_HOME/bin" sh
uv --version
}
setup_pyo3_python() {
log_section "Installing Python ${PYO3_PYTHON_VERSION} for PyO3 tests"
uv python install "$PYO3_PYTHON_VERSION"
PYO3_PYTHON="$(uv python find \
--managed-python \
--no-project \
--resolve-links \
"$PYO3_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:-}"
uv python install "$UV_SCRIPT_PYTHON_VERSION"
}
run_style_clippy() {
@@ -144,7 +117,6 @@ run_style_clippy() {
run_tests() {
install_uv
setup_pyo3_python
install_cargo_binstall
install_cargo_nextest
-85
View File
@@ -3403,74 +3403,6 @@ version = "0.1.29"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e0c5ccf5294c6ccd63a74f1565028353830a9c2f5eb0c682c355c471726a6e3f"
[[package]]
name = "pyo3"
version = "0.29.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "4688ddedf473e32662b9b067670129a8afb8c18e351482c70d62ba4a88171e8b"
dependencies = [
"libc",
"once_cell",
"portable-atomic",
"pyo3-build-config",
"pyo3-ffi",
"pyo3-macros",
]
[[package]]
name = "pyo3-build-config"
version = "0.29.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "f41027e41b4bd03f6e60f9f417fe24a6341a6bb744edd62b6f709f2a52ea30e9"
dependencies = [
"target-lexicon",
]
[[package]]
name = "pyo3-ffi"
version = "0.29.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e591a95526fead067432c3b3a33fc74770b87b1e04e73671090d9c2055a2b327"
dependencies = [
"libc",
"pyo3-build-config",
]
[[package]]
name = "pyo3-macros"
version = "0.29.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "73225868fc1cd84eef2c3c230ddb91273bf1de46aeb8a4248da76d32a0924a1c"
dependencies = [
"proc-macro2",
"pyo3-macros-backend",
"quote",
"syn 2.0.117",
]
[[package]]
name = "pyo3-macros-backend"
version = "0.29.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "571575aa3749fa6216757dd47d2a3e7ef360f329a40f0666a9fbd14889024952"
dependencies = [
"heck",
"proc-macro2",
"quote",
"syn 2.0.117",
]
[[package]]
name = "pythonize"
version = "0.29.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "6ec376e1216e0c929a74964ce2020012a1a39f32d80e78aa688721219ea7fb89"
dependencies = [
"pyo3",
"serde",
"serde_json",
]
[[package]]
name = "quick-error"
version = "2.0.1"
@@ -4792,12 +4724,6 @@ 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"
@@ -6113,17 +6039,6 @@ dependencies = [
"tracing",
]
[[package]]
name = "vllm-tool-parser-py"
version = "0.1.0"
dependencies = [
"pyo3",
"pythonize",
"serde_json",
"thiserror-ext",
"vllm-parser",
]
[[package]]
name = "vllm-tracing"
version = "0.1.0"
-3
View File
@@ -11,7 +11,6 @@ members = [
"src/metrics",
"src/mock-engine",
"src/parser",
"src/parser/python",
"src/server",
"src/text",
"src/tokenizer",
@@ -79,8 +78,6 @@ prometheus-client-derive-encode = "0.5.0"
prost = "0.14.3"
prost-types = "0.14.3"
protox = "0.9.1"
pyo3 = "0.29.0"
pythonize = "0.29.0"
rand = "0.9.3"
rand_distr = "0.5.1"
rayon = "1.11.0"
-19
View File
@@ -1,19 +0,0 @@
[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-parser.workspace = true
[lints]
workspace = true
-398
View File
@@ -1,398 +0,0 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
//! Thin PyO3 bindings for `vllm_parser::tool`.
//!
//! 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_parser::tool::{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_parser::tool::$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! {
MinimaxM3ToolParser,
// 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<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 mut output = ToolParserOutput::default();
output.push_text(normal_text);
for call in calls.unwrap_or_default() {
output.push_call(call.borrow(py).0.clone());
}
Self(output)
}
#[getter]
fn normal_text(&self) -> String {
self.0.normal_text()
}
#[getter]
fn calls(&self) -> Vec<PyToolCallDelta> {
self.0.calls().into_iter().cloned().map(PyToolCallDelta).collect()
}
fn append(&mut self, other: PyRef<'_, PyToolParserOutput>) {
self.0.append(other.0.clone());
}
fn coalesce(&self) -> Self {
Self(self.0.clone().coalesce())
}
}
#[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()
}
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::<PyTool>()?;
m.add_class::<PyToolCallDelta>()?;
m.add_class::<PyToolParserOutput>()?;
m.add_class::<PyToolParser>()?;
Ok(())
}
#[cfg(test)]
mod tests {
use serde_json::json;
use super::*;
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#"<|DSML|tool_calls>
<|DSML|invoke name="create_order">
<|DSML|parameter name="user_id" string="false">42</|DSML|parameter>
<|DSML|parameter name="shipping" string="false">{"city":"Singapore","zip":18956}</|DSML|parameter>
</|DSML|invoke>
</|DSML|tool_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() {
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();
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();
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_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();
}
}
+2
View File
@@ -87,6 +87,8 @@ def get_missing_precompiled_rust_extension_modules() -> list[str]:
def has_precompiled_rust_extensions() -> bool:
if not rust_build.rust_py_extension_module_names():
return False
return not get_missing_precompiled_rust_extension_modules()
-328
View File
@@ -1,328 +0,0 @@
# 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 = "</|DSML|tool_calls>"
INV_START = '<|DSML|invoke name="'
INV_END = "</|DSML|invoke>"
PARAM_START = '<|DSML|parameter name="'
PARAM_END = "</|DSML|parameter>"
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()
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
+9 -8
View File
@@ -55,14 +55,15 @@ def rust_extensions(*, optional: bool = False) -> list[RustExtension]:
binding=Binding.Exec,
optional=optional,
),
RustExtension(
target="vllm._rust_tool_parser",
path="rust/src/parser/python/Cargo.toml",
features=["pyo3/abi3-py38"],
binding=Binding.PyO3,
optional=optional,
py_limited_api=True,
),
# Example PyO3 module (the removed tool-parser bridge):
# RustExtension(
# target="vllm._rust_tool_parser",
# path="rust/src/parser/python/Cargo.toml",
# features=["pyo3/abi3-py38"],
# binding=Binding.PyO3,
# optional=optional,
# py_limited_api=True,
# ),
]
-322
View File
@@ -1,322 +0,0 @@
# 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.generate.base.protocol import (
DeltaFunctionCall,
DeltaMessage,
DeltaToolCall,
ExtractedToolCallInformation,
FunctionCall,
ToolCall,
)
from vllm.entrypoints.openai.chat_completion.protocol import (
ChatCompletionRequest,
ChatCompletionToolsParam,
)
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(), 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