mirror of
https://github.com/vllm-project/vllm.git
synced 2026-10-28 09:45:26 +08:00
[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:
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
|
||||
|
||||
|
||||
Generated
-85
@@ -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"
|
||||
|
||||
@@ -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"
|
||||
|
||||
@@ -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
|
||||
@@ -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()),
|
||||
¶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::<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();
|
||||
}
|
||||
}
|
||||
@@ -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()
|
||||
|
||||
|
||||
|
||||
@@ -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
@@ -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,
|
||||
# ),
|
||||
]
|
||||
|
||||
|
||||
|
||||
@@ -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
|
||||
Reference in new issue
Block a user