forked from Karylab-cklius/vllm
Merge branch 'main' into wentao-optimize-per-token-group-quant
This commit is contained in:
@@ -2,17 +2,16 @@ name: vllm_intel_ci
|
||||
job_dirs:
|
||||
- ".buildkite/intel_jobs"
|
||||
run_all_patterns:
|
||||
- ".buildkite/ci_config_intel.yaml"
|
||||
- "docker/Dockerfile"
|
||||
- "docker/Dockerfile.xpu"
|
||||
- "CMakeLists.txt"
|
||||
- "requirements/common.txt"
|
||||
- "requirements/xpu.txt"
|
||||
- "requirements/build/cuda.txt"
|
||||
- "requirements/test/cuda.txt"
|
||||
- "setup.py"
|
||||
- "csrc/"
|
||||
- "cmake/"
|
||||
run_all_exclude_patterns:
|
||||
- "docker/Dockerfile."
|
||||
- "csrc/cpu/"
|
||||
- "csrc/rocm/"
|
||||
- "cmake/hipify.py"
|
||||
|
||||
@@ -53,7 +53,7 @@ steps:
|
||||
- tests/models/language/pooling/
|
||||
commands:
|
||||
- |
|
||||
bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 40m "
|
||||
bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 50m "
|
||||
pytest -x -v -s tests/models/language/generation -m cpu_model
|
||||
pytest -x -v -s tests/models/language/pooling -m cpu_model"
|
||||
|
||||
@@ -68,13 +68,15 @@ steps:
|
||||
- vllm/v1/sample/ops/topk_topp_triton.py
|
||||
- vllm/v1/sample/ops/topk_topp_sampler.py
|
||||
- tests/v1/sample/test_topk_topp_sampler.py
|
||||
- tests/v1/e2e/test_cpu_linear_attn_chunked_prefix.py
|
||||
commands:
|
||||
- |
|
||||
bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 45m "
|
||||
uv pip install git+https://github.com/triton-lang/triton-cpu.git@270e696d
|
||||
VLLM_USE_V2_MODEL_RUNNER=1 pytest -x -v -s tests/models/language/generation/test_granite.py -m cpu_model
|
||||
# TODO: move to CPU-Kernel Tests once triton-cpu has a pre-built wheel
|
||||
pytest -x -v -s tests/v1/sample/test_topk_topp_sampler.py::TestTritonTopkTopp"
|
||||
pytest -x -v -s tests/v1/sample/test_topk_topp_sampler.py::TestTritonTopkTopp
|
||||
pytest -x -v -s tests/v1/e2e/test_cpu_linear_attn_chunked_prefix.py"
|
||||
|
||||
- label: CPU-Quantization Model Tests
|
||||
depends_on: []
|
||||
|
||||
@@ -7,7 +7,8 @@ set -euox pipefail
|
||||
# allow to bind to different cores
|
||||
CORE_RANGE=${CORE_RANGE:-48-95}
|
||||
NUMA_NODE=${NUMA_NODE:-1}
|
||||
IMAGE_NAME="cpu-test-$NUMA_NODE"
|
||||
AGENT_SLOT=${AGENT_SLOT:-}
|
||||
IMAGE_NAME="cpu-test-${NUMA_NODE}${AGENT_SLOT:+-${AGENT_SLOT}}"
|
||||
TIMEOUT_VAL=$1
|
||||
TEST_COMMAND=$2
|
||||
|
||||
|
||||
@@ -33,6 +33,14 @@ if [[ -n "${ATTENTION_BACKEND:-}" ]]; then
|
||||
EXTRA_ARGS+=(--attention-backend "${ATTENTION_BACKEND}")
|
||||
fi
|
||||
|
||||
# ROCm: run eager to avoid intermittent HIP-graph decode corruption.
|
||||
# See https://github.com/ROCm/clr/issues/279
|
||||
# TODO(aarushjain29): Revert after TheRock 7.14
|
||||
if command -v rocm-smi &> /dev/null || command -v amd-smi &> /dev/null || [[ -d /opt/rocm ]] || [[ -n "${ROCM_PATH:-}" ]]; then
|
||||
echo "ROCm platform detected: adding --enforce-eager to avoid HIP-graph decode corruption"
|
||||
EXTRA_ARGS+=(--enforce-eager)
|
||||
fi
|
||||
|
||||
cleanup() {
|
||||
if [[ -n "${SERVER_PID:-}" ]] && kill -0 "${SERVER_PID}" 2>/dev/null; then
|
||||
kill "${SERVER_PID}" 2>/dev/null || true
|
||||
|
||||
+12
-36
@@ -132,7 +132,7 @@ steps:
|
||||
- vllm/_aiter_ops.py
|
||||
- vllm/platforms/rocm.py
|
||||
commands:
|
||||
- TARGET_TEST_SUITE=L4 pytest basic_correctness/ -v -s -m 'distributed(num_gpus=2)'
|
||||
- TARGET_TEST_SUITE=MI250 pytest basic_correctness/ -v -s -m 'distributed(num_gpus=2)'
|
||||
- CUDA_VISIBLE_DEVICES=0,1 pytest -v -s model_executor/model_loader/test_sharded_state_loader.py -m '(not slow_test)'
|
||||
- pytest models/test_transformers.py -v -s -m 'distributed(num_gpus=2)'
|
||||
- pytest models/language -v -s -m 'distributed(num_gpus=2)'
|
||||
@@ -221,30 +221,6 @@ steps:
|
||||
- pytest -v -s distributed/test_shm_buffer.py
|
||||
- pytest -v -s distributed/test_shm_storage.py
|
||||
|
||||
- label: Distributed Torchrun + Shutdown Tests (2 GPUs) # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
|
||||
agent_pool: mi250_2
|
||||
num_gpus: 2
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/distributed/
|
||||
- vllm/engine/
|
||||
- vllm/executor/
|
||||
- vllm/worker/worker_base.py
|
||||
- vllm/v1/engine/
|
||||
- vllm/v1/worker/
|
||||
- tests/distributed/
|
||||
- tests/v1/shutdown
|
||||
- tests/v1/worker/test_worker_memory_snapshot.py
|
||||
- vllm/platforms/rocm.py
|
||||
commands:
|
||||
- VLLM_TEST_SAME_HOST=1 torchrun --nproc-per-node=4 distributed/test_same_node.py | grep 'Same node test passed'
|
||||
- VLLM_TEST_SAME_HOST=1 VLLM_TEST_WITH_DEFAULT_DEVICE_SET=1 torchrun --nproc-per-node=4 distributed/test_same_node.py | grep 'Same node test passed'
|
||||
- CUDA_VISIBLE_DEVICES=0,1 pytest -v -s v1/shutdown
|
||||
- pytest -v -s v1/worker/test_worker_memory_snapshot.py
|
||||
|
||||
- label: Pipeline + Context Parallelism (4 GPUs) # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
|
||||
@@ -617,7 +593,7 @@ steps:
|
||||
- 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 reasoning
|
||||
- pytest -v -s tool_parsers
|
||||
- pytest -v -s parser
|
||||
- pytest -v -s transformers_utils
|
||||
@@ -646,7 +622,7 @@ steps:
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s basic_correctness/test_mem.py
|
||||
- pytest -v -s basic_correctness/test_basic_correctness.py
|
||||
- VLLM_TARGET_TEST_SUITE=MI300 pytest -v -s basic_correctness/test_basic_correctness.py
|
||||
- pytest -v -s basic_correctness/test_cpu_offload.py
|
||||
|
||||
- label: Distributed Model Tests (2 GPUs) # TBD
|
||||
@@ -668,7 +644,7 @@ steps:
|
||||
- tests/model_executor/model_loader/test_sharded_state_loader.py
|
||||
- tests/models/
|
||||
commands:
|
||||
- TARGET_TEST_SUITE=L4 pytest basic_correctness/ -v -s -m 'distributed(num_gpus=2)'
|
||||
- TARGET_TEST_SUITE=MI300 pytest basic_correctness/ -v -s -m 'distributed(num_gpus=2)'
|
||||
- CUDA_VISIBLE_DEVICES=0,1 pytest -v -s model_executor/model_loader/test_sharded_state_loader.py -m '(not slow_test)'
|
||||
- pytest models/test_transformers.py -v -s -m 'distributed(num_gpus=2)'
|
||||
- pytest models/language -v -s -m 'distributed(num_gpus=2)'
|
||||
@@ -887,7 +863,7 @@ steps:
|
||||
commands:
|
||||
- pytest -v -s distributed/test_custom_all_reduce.py
|
||||
- torchrun --nproc_per_node=2 distributed/test_ca_buffer_sharing.py
|
||||
- TARGET_TEST_SUITE=A100 pytest basic_correctness/ -v -s -m 'distributed(num_gpus=2)'
|
||||
- TARGET_TEST_SUITE=MI300 pytest basic_correctness/ -v -s -m 'distributed(num_gpus=2)'
|
||||
- pytest -v -s -x lora/test_mixtral.py
|
||||
|
||||
- label: Distributed Torchrun + Examples (4 GPUs) # TBD
|
||||
@@ -1588,10 +1564,10 @@ steps:
|
||||
- pytest -v -s models/test_initialization.py::test_can_initialize_small_subset
|
||||
|
||||
- label: Basic Models Tests (Other) # TBD
|
||||
timeout_in_minutes: 180
|
||||
timeout_in_minutes: 45
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
agent_pool: mi300_1
|
||||
torch_nightly: true
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
@@ -2240,7 +2216,7 @@ steps:
|
||||
- pytest -v -s tests/distributed/test_packed_tensor.py
|
||||
|
||||
- label: Metrics, Tracing (2 GPUs) # TBD
|
||||
timeout_in_minutes: 180
|
||||
timeout_in_minutes: 20
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
agent_pool: mi300_2
|
||||
optional: true
|
||||
@@ -2393,7 +2369,7 @@ steps:
|
||||
#------------------------------------------------------ mi300 · weight_loading -------------------------------------------------------#
|
||||
|
||||
- label: Weight Loading Multiple GPU # TBD
|
||||
timeout_in_minutes: 180
|
||||
timeout_in_minutes: 30
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
agent_pool: mi300_2
|
||||
num_gpus: 2
|
||||
@@ -2405,7 +2381,7 @@ steps:
|
||||
- bash weight_loading/run_model_weight_loading_test.sh -c weight_loading/models-amd.txt
|
||||
|
||||
- label: Weight Loading Multiple GPU - Large Models # TBD
|
||||
timeout_in_minutes: 180
|
||||
timeout_in_minutes: 30
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
agent_pool: mi300_2
|
||||
num_gpus: 2
|
||||
@@ -3318,7 +3294,7 @@ steps:
|
||||
#------------------------------------------------------ mi355 · weight_loading -------------------------------------------------------#
|
||||
|
||||
- label: Weight Loading Multiple GPU # TBD
|
||||
timeout_in_minutes: 180
|
||||
timeout_in_minutes: 30
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
|
||||
agent_pool: mi355_2
|
||||
num_gpus: 2
|
||||
@@ -3330,7 +3306,7 @@ steps:
|
||||
- bash weight_loading/run_model_weight_loading_test.sh -c weight_loading/models-amd.txt
|
||||
|
||||
- label: Weight Loading Multiple GPU - Large Models # TBD
|
||||
timeout_in_minutes: 180
|
||||
timeout_in_minutes: 30
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
|
||||
agent_pool: mi355_2
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
|
||||
@@ -252,9 +252,10 @@ steps:
|
||||
- pytest -v -s v1/tracing
|
||||
mirror:
|
||||
amd:
|
||||
device: mi300_2
|
||||
device: mi325_2
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
optional: true
|
||||
|
||||
- label: Python-only Installation
|
||||
key: python-only-installation
|
||||
@@ -359,7 +360,7 @@ steps:
|
||||
- pytest -v -s test_ray_env.py
|
||||
- pytest -v -s -m 'cpu_test' multimodal
|
||||
- pytest -v -s renderers
|
||||
- pytest -v -s reasoning --ignore=reasoning/test_seedoss_reasoning_parser.py --ignore=reasoning/test_glm4_moe_reasoning_parser.py
|
||||
- pytest -v -s reasoning
|
||||
- pytest -v -s tool_parsers
|
||||
- pytest -v -s tokenizers_
|
||||
- pytest -v -s parser
|
||||
|
||||
@@ -45,6 +45,11 @@ steps:
|
||||
- tests/models/test_registry.py
|
||||
commands:
|
||||
- pytest -v -s models/test_terratorch.py models/test_transformers.py models/test_registry.py
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
|
||||
- label: Basic Models Test (Other CPU) # 5min
|
||||
key: basic-models-test-other-cpu
|
||||
|
||||
@@ -13,6 +13,13 @@ steps:
|
||||
- tests/weight_loading
|
||||
commands:
|
||||
- bash weight_loading/run_model_weight_loading_test.sh -c weight_loading/models.txt
|
||||
mirror:
|
||||
amd:
|
||||
device: mi300_2
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
commands:
|
||||
- bash weight_loading/run_model_weight_loading_test.sh -c weight_loading/models-amd.txt
|
||||
|
||||
# - label: Weight Loading Multiple GPU - Large Models # optional
|
||||
# working_dir: "/vllm-workspace/tests"
|
||||
|
||||
@@ -15,9 +15,9 @@ jobs:
|
||||
timeout-minutes: 30
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v6.0.1
|
||||
- uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 # v6.0.1
|
||||
|
||||
- uses: astral-sh/setup-uv@v7
|
||||
- uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7.6.0
|
||||
with:
|
||||
enable-cache: true
|
||||
cache-dependency-glob: |
|
||||
|
||||
+2
-2
@@ -869,9 +869,9 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
|
||||
endif()
|
||||
|
||||
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
|
||||
cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0f" "${CUDA_ARCHS}")
|
||||
cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0f;11.0f" "${CUDA_ARCHS}")
|
||||
else()
|
||||
cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0a;10.3a" "${CUDA_ARCHS}")
|
||||
cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0a;10.1a;10.3a" "${CUDA_ARCHS}")
|
||||
endif()
|
||||
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND SCALED_MM_ARCHS)
|
||||
set(CUTLASS_MOE_SM100_SRCS "csrc/libtorch_stable/quantization/w8a8/cutlass/moe/grouped_mm_c3x_sm100.cu")
|
||||
|
||||
@@ -80,13 +80,17 @@ _FI_MAX_SIZES = {
|
||||
2: 64 * MiB, # 64MB
|
||||
4: 64 * MiB, # 64MB
|
||||
8: 64 * MiB, # 64MB
|
||||
16: 64 * MiB, # 64MB (multi-node)
|
||||
}
|
||||
|
||||
# Global workspace tensors for FlashInfer (keyed by backend name)
|
||||
_FI_WORKSPACES: dict = {}
|
||||
|
||||
# Backends to benchmark
|
||||
FLASHINFER_BACKENDS = ["trtllm", "mnnvl"]
|
||||
# Backends to benchmark. trtllm is single-node only and can hang cross-node, so
|
||||
# multi-node sweeps can restrict to mnnvl via FI_BACKENDS=mnnvl.
|
||||
FLASHINFER_BACKENDS = [
|
||||
b for b in os.environ.get("FI_BACKENDS", "trtllm,mnnvl").split(",") if b
|
||||
]
|
||||
|
||||
|
||||
def setup_flashinfer_workspace(
|
||||
@@ -995,7 +999,10 @@ def main():
|
||||
rank = int(os.environ["RANK"])
|
||||
world_size = int(os.environ["WORLD_SIZE"])
|
||||
|
||||
device = torch.device(f"cuda:{rank}")
|
||||
# Use LOCAL_RANK for the device so multi-node runs (global rank >= GPUs per
|
||||
# node) map to a valid local GPU; falls back to global rank single-node.
|
||||
local_rank = int(os.environ.get("LOCAL_RANK", rank))
|
||||
device = torch.device(f"cuda:{local_rank}")
|
||||
torch.accelerator.set_device_index(device)
|
||||
torch.set_default_device(device)
|
||||
|
||||
|
||||
@@ -391,16 +391,19 @@ def get_configs_compute_bound(use_fp16, block_quant_shape) -> list[dict[str, int
|
||||
config = dict(zip(keys, config_values))
|
||||
configs.append(config)
|
||||
|
||||
# Remove configs that are not compatible with fp8 block quantization
|
||||
# BLOCK_SIZE_K must be a multiple of block_k
|
||||
# BLOCK_SIZE_N must be a multiple of block_n
|
||||
# Drop configs incompatible with fp8 block quantization. A tile must align
|
||||
# to the quant-block scale grid, i.e. tile and block must divide one
|
||||
# another. The kernel indexes scales per element (offs_bn // group_n,
|
||||
# k_start // group_k), so a tile narrower than the block (e.g. N=64 with
|
||||
# block_n=128) is valid -- and often faster at small batch. An exact
|
||||
# multiple was required before, which dropped those smaller tiles entirely.
|
||||
if block_quant_shape is not None and not use_fp16:
|
||||
block_n, block_k = block_quant_shape[0], block_quant_shape[1]
|
||||
for config in configs[:]:
|
||||
if (
|
||||
config["BLOCK_SIZE_K"] % block_k != 0
|
||||
or config["BLOCK_SIZE_N"] % block_n != 0
|
||||
):
|
||||
bn, bk = config["BLOCK_SIZE_N"], config["BLOCK_SIZE_K"]
|
||||
n_aligned = bn % block_n == 0 or block_n % bn == 0
|
||||
k_aligned = bk % block_k == 0 or block_k % bk == 0
|
||||
if not (n_aligned and k_aligned):
|
||||
configs.remove(config)
|
||||
return configs
|
||||
|
||||
|
||||
@@ -166,12 +166,13 @@ elseif (S390_FOUND)
|
||||
"-mtune=native")
|
||||
elseif (CMAKE_SYSTEM_PROCESSOR MATCHES "riscv64")
|
||||
message(STATUS "RISC-V detected")
|
||||
if(DEFINED VLLM_RVV_VLEN AND NOT VLLM_RVV_VLEN GREATER 0)
|
||||
if(DEFINED VLLM_RVV_VLEN AND VLLM_RVV_VLEN LESS 0)
|
||||
message(FATAL_ERROR
|
||||
"VLLM_RVV_VLEN must be a positive integer; got '${VLLM_RVV_VLEN}'")
|
||||
"VLLM_RVV_VLEN must be zero or a positive integer; got '${VLLM_RVV_VLEN}'")
|
||||
endif()
|
||||
# VLLM_RVV_VLEN selects the target VLEN. Auto-detected from /proc/cpuinfo
|
||||
# by default; override with -DVLLM_RVV_VLEN=128 or -DVLLM_RVV_VLEN=256.
|
||||
# by default; set -DVLLM_RVV_VLEN=0 to force scalar RISC-V build.
|
||||
# Override with -DVLLM_RVV_VLEN=128 or -DVLLM_RVV_VLEN=256 for RVV.
|
||||
if(NOT DEFINED VLLM_RVV_VLEN)
|
||||
# Auto-detect: find the largest zvl<N>b in /proc/cpuinfo isa line.
|
||||
if(EXISTS /proc/cpuinfo)
|
||||
|
||||
@@ -3,13 +3,17 @@
|
||||
|
||||
// VLEN-to-LMUL mapping for RISC-V Vector extension.
|
||||
//
|
||||
// LMUL_<N> expands to the LMUL suffix giving N total bits of vector data:
|
||||
// VLEN=128: LMUL_128=m1, LMUL_256=m2, LMUL_512=m4, LMUL_1024=m8
|
||||
// VLEN=256: LMUL_128=mf2, LMUL_256=m1, LMUL_512=m2, LMUL_1024=m4
|
||||
// LMUL_<N> expands to the LMUL suffix giving N total bits of vector data.
|
||||
// LMUL_64 is used by 8-lane int8/uint8 vectors.
|
||||
// VLEN=128:
|
||||
// LMUL_64=mf2, LMUL_128=m1, LMUL_256=m2, LMUL_512=m4, LMUL_1024=m8
|
||||
// VLEN=256:
|
||||
// LMUL_64=mf4, LMUL_128=mf2, LMUL_256=m1, LMUL_512=m2, LMUL_1024=m4
|
||||
|
||||
#include <riscv_vector.h>
|
||||
|
||||
#if __riscv_v_min_vlen == 128
|
||||
#define LMUL_64 mf2
|
||||
#define LMUL_128 m1
|
||||
#define LMUL_256 m2
|
||||
#define LMUL_512 m4
|
||||
@@ -17,6 +21,7 @@
|
||||
#define BOOL_256 b16
|
||||
#define BOOL_512 b8
|
||||
#elif __riscv_v_min_vlen == 256
|
||||
#define LMUL_64 mf4
|
||||
#define LMUL_128 mf2
|
||||
#define LMUL_256 m1
|
||||
#define LMUL_512 m2
|
||||
@@ -41,6 +46,16 @@
|
||||
|
||||
// ---- Semantic fixed-vector typedefs (named by element count) ----
|
||||
|
||||
// uint8 / int8
|
||||
typedef RVVTYPE(vuint8, LMUL_64, _t) fixed_u8x8_t
|
||||
__attribute__((riscv_rvv_vector_bits(64)));
|
||||
typedef RVVTYPE(vint8, LMUL_64, _t) fixed_i8x8_t
|
||||
__attribute__((riscv_rvv_vector_bits(64)));
|
||||
|
||||
// int16
|
||||
typedef RVVTYPE(vint16, LMUL_128, _t) fixed_i16x8_t
|
||||
__attribute__((riscv_rvv_vector_bits(128)));
|
||||
|
||||
// float16
|
||||
typedef RVVTYPE(vfloat16, LMUL_128, _t) fixed_fp16x8_t
|
||||
__attribute__((riscv_rvv_vector_bits(128)));
|
||||
|
||||
@@ -289,19 +289,18 @@ void causal_conv1d_fwd_kernel_impl(
|
||||
}
|
||||
}
|
||||
|
||||
#define LAUNCH_TINYGEMM_VARLEN_KERNEL(K, NB_SIZE) \
|
||||
tinygemm_kernel<scalar_t, K, NB_SIZE, has_bias, has_silu>::apply( \
|
||||
input + batch_offset * dim + mb_start * dim + nb_start, \
|
||||
weight + nb_start * width, \
|
||||
out + batch_offset * dim + mb_start * dim + nb_start, \
|
||||
has_bias ? bias + nb_start : nullptr, \
|
||||
nullptr, \
|
||||
false, \
|
||||
mb_size, \
|
||||
dim, \
|
||||
#define LAUNCH_TINYGEMM_VARLEN_KERNEL(K, NB_SIZE) \
|
||||
tinygemm_kernel<scalar_t, K, NB_SIZE, has_bias, has_silu>::apply( \
|
||||
input + batch_offset * dim + mb_start * dim + nb_start, \
|
||||
weight + nb_start * width, \
|
||||
out + batch_offset * dim + mb_start * dim + nb_start, \
|
||||
has_bias ? bias + nb_start : nullptr, \
|
||||
has_conv_states ? conv_states + conv_state_index * conv_state_slot_stride + nb_start : nullptr, \
|
||||
has_initial_states_value, \
|
||||
mb_size, \
|
||||
dim, \
|
||||
mb_start == 0);
|
||||
|
||||
// TODO: add `has_initial_state` support for varlen kernel
|
||||
template <typename scalar_t>
|
||||
void causal_conv1d_fwd_varlen_kernel_impl(
|
||||
scalar_t* __restrict__ out,
|
||||
@@ -343,6 +342,9 @@ void causal_conv1d_fwd_varlen_kernel_impl(
|
||||
int64_t nb_start = nb * BLOCK_N;
|
||||
int64_t nb_size = std::min(dim - nb_start, BLOCK_N);
|
||||
|
||||
const bool has_initial_states_value = has_conv_states ? has_initial_state[bs] : false;
|
||||
int32_t conv_state_index = has_conv_indices ? conv_indices[bs] : bs;
|
||||
|
||||
switch (width << 4 | nb_size >> 4) {
|
||||
case 0x42:
|
||||
LAUNCH_TINYGEMM_VARLEN_KERNEL(4, 32);
|
||||
@@ -373,7 +375,7 @@ void causal_conv1d_fwd_varlen_kernel_impl(
|
||||
width,
|
||||
dim,
|
||||
seqlen,
|
||||
/* has_initial_state */ false);
|
||||
has_initial_state[bs]);
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
@@ -285,6 +285,125 @@ inline int32_t load_uint4_vnni(const uint8_t* __restrict__ B, int64_t k, int64_t
|
||||
return (n_group % 2 == 0) ? (packed & 0x0f) : ((packed >> 4) & 0x0f);
|
||||
}
|
||||
|
||||
#if defined(CPU_CAPABILITY_RVV)
|
||||
template <int64_t N, int64_t ldb, int group>
|
||||
inline fixed_i8x8_t load_uint4_as_int8_rvv(const uint8_t* __restrict__ B, int64_t k) {
|
||||
constexpr int64_t n_group_size = 8;
|
||||
constexpr int64_t vnni_size = 4;
|
||||
static_assert(N == 32);
|
||||
static_assert(ldb == N / 2);
|
||||
static_assert(group >= 0 && group < N / n_group_size);
|
||||
|
||||
// Unpack: gather 8 packed int4 values from the VNNI4 layout.
|
||||
const int64_t ki = k % vnni_size;
|
||||
const int64_t k_base = k - ki;
|
||||
constexpr int64_t packed_group = group / 2;
|
||||
const uint8_t* packed_ptr = B + k_base * ldb + packed_group * n_group_size * vnni_size + ki;
|
||||
|
||||
fixed_u8x8_t packed = RVVI(__riscv_vlse8_v_u8, LMUL_64)(packed_ptr, vnni_size, n_group_size);
|
||||
if constexpr (group % 2 == 1) {
|
||||
packed = RVVI(__riscv_vsrl_vx_u8, LMUL_64)(packed, 4, n_group_size);
|
||||
}
|
||||
fixed_u8x8_t nibbles = RVVI(__riscv_vand_vx_u8, LMUL_64)(packed, 0x0f, n_group_size);
|
||||
return RVVI4(__riscv_vreinterpret_v_u8, LMUL_64, _i8, LMUL_64)(nibbles);
|
||||
}
|
||||
|
||||
inline fixed_i32x8_t gemm_accum_uint8_int8_rvv(fixed_i32x8_t acc, uint8_t a, fixed_i8x8_t b) {
|
||||
constexpr int64_t vl = 8;
|
||||
fixed_i16x8_t b_i16 = RVVI(__riscv_vsext_vf2_i16, LMUL_128)(b, vl);
|
||||
return RVVI(__riscv_vwmacc_vx_i32, LMUL_256)(acc, static_cast<int16_t>(a), b_i16, vl);
|
||||
}
|
||||
|
||||
template <int64_t N, int64_t ldb, int group>
|
||||
inline fixed_i32x8_t gemm_accum_uint4_rvv(
|
||||
fixed_i32x8_t acc,
|
||||
const uint8_t* __restrict__ B,
|
||||
const int8_t* __restrict__ qzeros_b,
|
||||
uint8_t a,
|
||||
int64_t k) {
|
||||
constexpr int64_t n_group_size = 8;
|
||||
fixed_i8x8_t b = load_uint4_as_int8_rvv<N, ldb, group>(B, k);
|
||||
fixed_i8x8_t qzeros =
|
||||
RVVI(__riscv_vle8_v_i8, LMUL_64)(qzeros_b + group * n_group_size, n_group_size);
|
||||
b = RVVI(__riscv_vsub_vv_i8, LMUL_64)(b, qzeros, n_group_size);
|
||||
return gemm_accum_uint8_int8_rvv(acc, a, b);
|
||||
}
|
||||
|
||||
template <int group>
|
||||
inline void _dequant_and_store_rvv(
|
||||
float* __restrict__ C,
|
||||
fixed_i32x8_t acc,
|
||||
const float* __restrict__ scales_a,
|
||||
const int32_t* __restrict__ qzeros_a,
|
||||
const float* __restrict__ scales_b,
|
||||
const int32_t* __restrict__ compensation,
|
||||
int64_t m,
|
||||
int64_t ldc) {
|
||||
constexpr int64_t n_group_size = 8;
|
||||
constexpr int64_t n = group * n_group_size;
|
||||
constexpr int64_t vl = n_group_size;
|
||||
|
||||
// Dequant compensation: remove activation zero-point contribution.
|
||||
fixed_i32x8_t comp = RVVI(__riscv_vle32_v_i32, LMUL_256)(compensation + n, vl);
|
||||
fixed_i32x8_t zp_comp = RVVI(__riscv_vmul_vx_i32, LMUL_256)(comp, qzeros_a[m], vl);
|
||||
acc = RVVI(__riscv_vsub_vv_i32, LMUL_256)(acc, zp_comp, vl);
|
||||
|
||||
// Scale: convert int32 accumulators to fp32 and apply activation/weight scales.
|
||||
fixed_fp32x8_t acc_f = RVVI(__riscv_vfcvt_f_x_v_f32, LMUL_256)(acc, vl);
|
||||
acc_f = RVVI(__riscv_vfmul_vf_f32, LMUL_256)(acc_f, scales_a[m], vl);
|
||||
fixed_fp32x8_t scale_b = RVVI(__riscv_vle32_v_f32, LMUL_256)(scales_b + n, vl);
|
||||
acc_f = RVVI(__riscv_vfmul_vv_f32, LMUL_256)(acc_f, scale_b, vl);
|
||||
|
||||
// Store: accumulate into the float scratch buffer that already holds bias/zero.
|
||||
float* c_ptr = C + m * ldc + n;
|
||||
fixed_fp32x8_t c_old = RVVI(__riscv_vle32_v_f32, LMUL_256)(c_ptr, vl);
|
||||
fixed_fp32x8_t c_new = RVVI(__riscv_vfadd_vv_f32, LMUL_256)(c_old, acc_f, vl);
|
||||
RVVI(__riscv_vse32_v_f32, LMUL_256)(c_ptr, c_new, vl);
|
||||
}
|
||||
|
||||
template <int64_t N, int64_t ldb>
|
||||
void _dequant_gemm_accum_rvv(
|
||||
float* __restrict__ C,
|
||||
const uint8_t* __restrict__ A,
|
||||
const float* __restrict__ scales_a,
|
||||
const int32_t* __restrict__ qzeros_a,
|
||||
const uint8_t* __restrict__ B,
|
||||
const float* __restrict__ scales_b,
|
||||
const int8_t* __restrict__ qzeros_b,
|
||||
const int32_t* __restrict__ compensation,
|
||||
int64_t M,
|
||||
int64_t K,
|
||||
int64_t lda,
|
||||
int64_t ldc) {
|
||||
static_assert(N == 32);
|
||||
static_assert(ldb == N / 2);
|
||||
constexpr int64_t vl = 8;
|
||||
|
||||
// Accumulate one C row over the 32-column block.
|
||||
for (int64_t m = 0; m < M; ++m) {
|
||||
fixed_i32x8_t acc0 = RVVI(__riscv_vmv_v_x_i32, LMUL_256)(0, vl);
|
||||
fixed_i32x8_t acc1 = RVVI(__riscv_vmv_v_x_i32, LMUL_256)(0, vl);
|
||||
fixed_i32x8_t acc2 = RVVI(__riscv_vmv_v_x_i32, LMUL_256)(0, vl);
|
||||
fixed_i32x8_t acc3 = RVVI(__riscv_vmv_v_x_i32, LMUL_256)(0, vl);
|
||||
// A[m][k] @ B[k][0:32] -> acc[m][0:32]
|
||||
for (int64_t k = 0; k < K; ++k) {
|
||||
// GEMM K step: one scalar activation updates four 8-column RVV tiles.
|
||||
const uint8_t a = A[m * lda + k];
|
||||
acc0 = gemm_accum_uint4_rvv<N, ldb, 0>(acc0, B, qzeros_b, a, k);
|
||||
acc1 = gemm_accum_uint4_rvv<N, ldb, 1>(acc1, B, qzeros_b, a, k);
|
||||
acc2 = gemm_accum_uint4_rvv<N, ldb, 2>(acc2, B, qzeros_b, a, k);
|
||||
acc3 = gemm_accum_uint4_rvv<N, ldb, 3>(acc3, B, qzeros_b, a, k);
|
||||
}
|
||||
|
||||
// Dequant/scale/store each 8-column group back into C.
|
||||
_dequant_and_store_rvv<0>(C, acc0, scales_a, qzeros_a, scales_b, compensation, m, ldc);
|
||||
_dequant_and_store_rvv<1>(C, acc1, scales_a, qzeros_a, scales_b, compensation, m, ldc);
|
||||
_dequant_and_store_rvv<2>(C, acc2, scales_a, qzeros_a, scales_b, compensation, m, ldc);
|
||||
_dequant_and_store_rvv<3>(C, acc3, scales_a, qzeros_a, scales_b, compensation, m, ldc);
|
||||
}
|
||||
}
|
||||
#endif
|
||||
|
||||
template <int64_t N, int64_t ldb, bool sym_quant_act>
|
||||
void _dequant_gemm_accum(
|
||||
float* C,
|
||||
@@ -336,6 +455,11 @@ void _dequant_gemm_accum(
|
||||
_dequant_and_store<true, N, sym_quant_act>(
|
||||
C, C_i32, scales_a, qzeros_a, scales_b, compensation, M, N /*ldi*/, ldc, 1 /*ldsa*/);
|
||||
} else
|
||||
#elif defined(CPU_CAPABILITY_RVV)
|
||||
if constexpr (!sym_quant_act && N == BLOCK_N && ldb == BLOCK_N / 2) {
|
||||
_dequant_gemm_accum_rvv<N, ldb>(C, A, scales_a, qzeros_a, B, scales_b, qzeros_b, compensation, M, K, lda, ldc);
|
||||
return;
|
||||
} else
|
||||
#endif
|
||||
{
|
||||
for (int64_t m = 0; m < M; ++m) {
|
||||
|
||||
@@ -9,11 +9,19 @@
|
||||
#define CPU_CAPABILITY_AVX512
|
||||
#endif
|
||||
|
||||
#if defined(__riscv_v_min_vlen) && (__riscv_v_min_vlen == 128 || __riscv_v_min_vlen == 256)
|
||||
#define CPU_CAPABILITY_RVV
|
||||
#endif
|
||||
|
||||
#include <ATen/cpu/vec/functional.h>
|
||||
#include <ATen/cpu/vec/vec.h>
|
||||
#if defined(CPU_CAPABILITY_AVX512)
|
||||
#include <immintrin.h>
|
||||
#endif
|
||||
|
||||
#if defined(CPU_CAPABILITY_RVV)
|
||||
#include "../cpu_types_riscv_defs.hpp"
|
||||
#endif
|
||||
namespace {
|
||||
|
||||
using namespace at::vec;
|
||||
|
||||
@@ -208,6 +208,89 @@ void copy_and_expand_eagle_inputs_kernel_impl(
|
||||
}
|
||||
}
|
||||
|
||||
void copy_and_expand_dflash_inputs_kernel_impl(
|
||||
const torch::Tensor& next_token_ids, const torch::Tensor& target_positions,
|
||||
torch::Tensor& out_input_ids, torch::Tensor& out_context_positions,
|
||||
torch::Tensor& out_query_positions, torch::Tensor& out_context_slot_mapping,
|
||||
torch::Tensor& out_query_slot_mapping, torch::Tensor& out_token_indices,
|
||||
const torch::Tensor& block_table, const torch::Tensor& query_start_loc,
|
||||
const std::optional<torch::Tensor>& num_rejected_tokens,
|
||||
const int64_t parallel_drafting_token_id, const int64_t block_size,
|
||||
const int64_t num_query_per_req, const int64_t num_speculative_tokens,
|
||||
const int64_t total_input_tokens, const bool has_num_rejected) {
|
||||
const int64_t num_reqs = query_start_loc.size(0) - 1;
|
||||
|
||||
const int64_t* next_ids_ptr = next_token_ids.data_ptr<int64_t>();
|
||||
const int64_t* target_pos_ptr = target_positions.data_ptr<int64_t>();
|
||||
const int32_t* block_table_ptr = block_table.data_ptr<int32_t>();
|
||||
const int32_t* query_start_ptr = query_start_loc.data_ptr<int32_t>();
|
||||
const int64_t* rejected_ptr =
|
||||
has_num_rejected && num_rejected_tokens.has_value()
|
||||
? num_rejected_tokens.value().data_ptr<int64_t>()
|
||||
: nullptr;
|
||||
|
||||
int64_t* out_ids_ptr = out_input_ids.data_ptr<int64_t>();
|
||||
int64_t* out_ctx_pos_ptr = out_context_positions.data_ptr<int64_t>();
|
||||
int64_t* out_query_pos_ptr = out_query_positions.data_ptr<int64_t>();
|
||||
int64_t* out_ctx_slot_ptr = out_context_slot_mapping.data_ptr<int64_t>();
|
||||
int64_t* out_query_slot_ptr = out_query_slot_mapping.data_ptr<int64_t>();
|
||||
int32_t* out_token_idx_ptr = out_token_indices.data_ptr<int32_t>();
|
||||
|
||||
const int64_t block_table_stride = block_table.stride(0);
|
||||
|
||||
#pragma omp parallel for
|
||||
for (int64_t req_idx = 0; req_idx < num_reqs; ++req_idx) {
|
||||
int32_t ctx_start = query_start_ptr[req_idx];
|
||||
int32_t ctx_end = query_start_ptr[req_idx + 1];
|
||||
int64_t num_ctx = ctx_end - ctx_start;
|
||||
int64_t valid_ctx_end = ctx_end;
|
||||
if (rejected_ptr != nullptr) {
|
||||
valid_ctx_end -= rejected_ptr[req_idx];
|
||||
}
|
||||
// Guard against out-of-bounds: ensure valid_ctx_end > ctx_start so that
|
||||
// valid_ctx_end - 1 never reads before the request's context range.
|
||||
valid_ctx_end =
|
||||
std::max(valid_ctx_end, static_cast<int64_t>(ctx_start + 1));
|
||||
|
||||
int64_t last_pos = target_pos_ptr[valid_ctx_end - 1];
|
||||
|
||||
for (int64_t j = 0; j < num_ctx; ++j) {
|
||||
int64_t ctx_idx = ctx_start + j;
|
||||
int64_t ctx_pos_idx = std::min(ctx_idx, total_input_tokens - 1);
|
||||
int64_t position = target_pos_ptr[ctx_pos_idx];
|
||||
int64_t block_num = position / block_size;
|
||||
block_num = std::min(block_num, block_table_stride - 1);
|
||||
int32_t block_id =
|
||||
block_table_ptr[req_idx * block_table_stride + block_num];
|
||||
int64_t slot = block_id * block_size + (position % block_size);
|
||||
|
||||
out_ctx_pos_ptr[ctx_idx] = position;
|
||||
out_ctx_slot_ptr[ctx_idx] = slot;
|
||||
}
|
||||
|
||||
for (int64_t query_off = 0; query_off < num_query_per_req; ++query_off) {
|
||||
int64_t query_out = req_idx * num_query_per_req + query_off;
|
||||
int64_t position = last_pos + 1 + query_off;
|
||||
int64_t block_num = position / block_size;
|
||||
block_num = std::min(block_num, block_table_stride - 1);
|
||||
int32_t block_id =
|
||||
block_table_ptr[req_idx * block_table_stride + block_num];
|
||||
int64_t slot = block_id * block_size + (position % block_size);
|
||||
|
||||
out_query_pos_ptr[query_out] = position;
|
||||
out_query_slot_ptr[query_out] = slot;
|
||||
out_ids_ptr[query_out] =
|
||||
query_off == 0 ? next_ids_ptr[req_idx] : parallel_drafting_token_id;
|
||||
|
||||
if (query_off > 0) {
|
||||
int64_t sample_out_idx =
|
||||
req_idx * num_speculative_tokens + (query_off - 1);
|
||||
out_token_idx_ptr[sample_out_idx] = query_out;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
void rejection_greedy_sample_kernel_impl(
|
||||
torch::Tensor& output_token_ids, const torch::Tensor& cu_num_draft_tokens,
|
||||
const torch::Tensor& draft_token_ids, const torch::Tensor& target_argmax,
|
||||
|
||||
@@ -237,6 +237,16 @@ void copy_and_expand_eagle_inputs_kernel_impl(
|
||||
const int64_t padding_token_id, const int64_t parallel_drafting_token_id,
|
||||
const int64_t total_input_tokens,
|
||||
const int64_t num_padding_slots_per_request, const bool shift_input_ids);
|
||||
void copy_and_expand_dflash_inputs_kernel_impl(
|
||||
const torch::Tensor& next_token_ids, const torch::Tensor& target_positions,
|
||||
torch::Tensor& out_input_ids, torch::Tensor& out_context_positions,
|
||||
torch::Tensor& out_query_positions, torch::Tensor& out_context_slot_mapping,
|
||||
torch::Tensor& out_query_slot_mapping, torch::Tensor& out_token_indices,
|
||||
const torch::Tensor& block_table, const torch::Tensor& query_start_loc,
|
||||
const std::optional<torch::Tensor>& num_rejected_tokens,
|
||||
const int64_t parallel_drafting_token_id, const int64_t block_size,
|
||||
const int64_t num_query_per_req, const int64_t num_speculative_tokens,
|
||||
const int64_t total_input_tokens, const bool has_num_rejected);
|
||||
void rejection_greedy_sample_kernel_impl(
|
||||
torch::Tensor& output_token_ids, const torch::Tensor& cu_num_draft_tokens,
|
||||
const torch::Tensor& draft_token_ids, const torch::Tensor& target_argmax,
|
||||
@@ -599,6 +609,19 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
|
||||
"SymInt total_input_tokens, SymInt num_padding_slots_per_request, "
|
||||
"bool shift_input_ids) -> ()",
|
||||
&cpu_utils::copy_and_expand_eagle_inputs_kernel_impl);
|
||||
ops.def(
|
||||
"copy_and_expand_dflash_inputs_kernel_impl("
|
||||
"Tensor next_token_ids, Tensor target_positions, "
|
||||
"Tensor(a2!) out_input_ids, Tensor(a3!) out_context_positions, "
|
||||
"Tensor(a4!) out_query_positions, "
|
||||
"Tensor(a5!) out_context_slot_mapping, "
|
||||
"Tensor(a6!) out_query_slot_mapping, "
|
||||
"Tensor(a7!) out_token_indices, Tensor block_table, "
|
||||
"Tensor query_start_loc, Tensor? num_rejected_tokens, "
|
||||
"SymInt parallel_drafting_token_id, SymInt block_size, "
|
||||
"SymInt num_query_per_req, SymInt num_speculative_tokens, "
|
||||
"SymInt total_input_tokens, bool has_num_rejected) -> ()",
|
||||
&cpu_utils::copy_and_expand_dflash_inputs_kernel_impl);
|
||||
ops.def(
|
||||
"rejection_greedy_sample_kernel_impl("
|
||||
"Tensor(a0!) output_token_ids, Tensor cu_num_draft_tokens, "
|
||||
|
||||
@@ -97,18 +97,28 @@ int64_t qr_max_size() {
|
||||
cast_bf2half>; \
|
||||
template struct quickreduce::AllReduceTwoshot<T, Codec<T, 8>, cast_bf2half>;
|
||||
|
||||
// INT3 (CodecQ3) is restricted to TP2 only, so we only instantiate the
|
||||
// world_size == 2 kernel for it.
|
||||
#define INSTANTIATE_FOR_WORLDSIZE_TP2_ONLY(T, Codec, cast_bf2half) \
|
||||
template struct quickreduce::AllReduceTwoshot<T, Codec<T, 2>, cast_bf2half>;
|
||||
|
||||
INSTANTIATE_FOR_WORLDSIZE(quickreduce::nv_bfloat16, quickreduce::CodecFP, false)
|
||||
INSTANTIATE_FOR_WORLDSIZE(quickreduce::nv_bfloat16, quickreduce::CodecQ4, false)
|
||||
INSTANTIATE_FOR_WORLDSIZE(quickreduce::nv_bfloat16, quickreduce::CodecQ6, false)
|
||||
INSTANTIATE_FOR_WORLDSIZE(quickreduce::nv_bfloat16, quickreduce::CodecQ8, false)
|
||||
INSTANTIATE_FOR_WORLDSIZE_TP2_ONLY(quickreduce::nv_bfloat16,
|
||||
quickreduce::CodecQ3, false)
|
||||
INSTANTIATE_FOR_WORLDSIZE(quickreduce::nv_bfloat16, quickreduce::CodecFP, true)
|
||||
INSTANTIATE_FOR_WORLDSIZE(quickreduce::nv_bfloat16, quickreduce::CodecQ4, true)
|
||||
INSTANTIATE_FOR_WORLDSIZE(quickreduce::nv_bfloat16, quickreduce::CodecQ6, true)
|
||||
INSTANTIATE_FOR_WORLDSIZE(quickreduce::nv_bfloat16, quickreduce::CodecQ8, true)
|
||||
INSTANTIATE_FOR_WORLDSIZE_TP2_ONLY(quickreduce::nv_bfloat16,
|
||||
quickreduce::CodecQ3, true)
|
||||
|
||||
INSTANTIATE_FOR_WORLDSIZE(half, quickreduce::CodecFP, false)
|
||||
INSTANTIATE_FOR_WORLDSIZE(half, quickreduce::CodecQ4, false)
|
||||
INSTANTIATE_FOR_WORLDSIZE(half, quickreduce::CodecQ6, false)
|
||||
INSTANTIATE_FOR_WORLDSIZE(half, quickreduce::CodecQ8, false)
|
||||
INSTANTIATE_FOR_WORLDSIZE_TP2_ONLY(half, quickreduce::CodecQ3, false)
|
||||
|
||||
#endif // USE_ROCM
|
||||
@@ -180,7 +180,7 @@ bool cutlass_group_gemm_supported(int64_t cuda_device_capability) {
|
||||
|
||||
#if defined CUDA_VERSION
|
||||
#if defined ENABLE_CUTLASS_MOE_SM100 && ENABLE_CUTLASS_MOE_SM100
|
||||
if (cuda_device_capability >= 100 && cuda_device_capability < 110) {
|
||||
if (cuda_device_capability >= 100 && cuda_device_capability < 120) {
|
||||
return CUDA_VERSION >= 12080;
|
||||
}
|
||||
#endif
|
||||
|
||||
@@ -186,6 +186,10 @@ __global__ void per_token_group_quant_8bit_kernel(
|
||||
static_cast<DST_DTYPE*>(output_q) + block_group_offset;
|
||||
scale_element_t* scale_output;
|
||||
|
||||
#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
|
||||
cudaGridDependencySynchronize();
|
||||
#endif
|
||||
|
||||
if constexpr (IS_COLUMN_MAJOR) {
|
||||
const int num_elems_per_pack =
|
||||
static_cast<int>(sizeof(scale_packed_t) / sizeof(scale_element_t));
|
||||
@@ -220,6 +224,10 @@ __global__ void per_token_group_quant_8bit_kernel(
|
||||
|
||||
QuantizeGroup<T, DST_DTYPE>(smem_group, group_output, group_size, lane_id,
|
||||
threads_per_group, y_s, min_8bit, max_8bit);
|
||||
|
||||
#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
|
||||
cudaTriggerProgrammaticLaunchCompletion();
|
||||
#endif
|
||||
}
|
||||
|
||||
inline int GetGroupsPerBlock(int64_t num_groups) {
|
||||
@@ -358,45 +366,56 @@ void per_token_group_quant_8bit(const torch::stable::Tensor& input,
|
||||
(input.scalar_type() == torch::headeronly::ScalarType::Half ||
|
||||
input.scalar_type() == torch::headeronly::ScalarType::BFloat16);
|
||||
|
||||
#define LAUNCH_KERNEL(T, DST_DTYPE) \
|
||||
do { \
|
||||
dim3 grid(num_blocks); \
|
||||
dim3 block(num_threads); \
|
||||
size_t smem_bytes = \
|
||||
static_cast<size_t>(groups_per_block) * group_size * sizeof(T); \
|
||||
if (is_column_major) { \
|
||||
if (scale_ue8m0) { \
|
||||
per_token_group_quant_8bit_kernel<T, DST_DTYPE, true, true> \
|
||||
<<<grid, block, smem_bytes, stream>>>( \
|
||||
static_cast<T*>(input.data_ptr()), output_q.data_ptr(), \
|
||||
static_cast<float*>(output_s.data_ptr()), group_size, \
|
||||
num_groups, groups_per_block, (float)eps, (float)min_8bit, \
|
||||
(float)max_8bit, scale_num_rows, scale_stride); \
|
||||
} else { \
|
||||
per_token_group_quant_8bit_kernel<T, DST_DTYPE, true, false> \
|
||||
<<<grid, block, smem_bytes, stream>>>( \
|
||||
static_cast<T*>(input.data_ptr()), output_q.data_ptr(), \
|
||||
static_cast<float*>(output_s.data_ptr()), group_size, \
|
||||
num_groups, groups_per_block, (float)eps, (float)min_8bit, \
|
||||
(float)max_8bit, scale_num_rows, scale_stride); \
|
||||
} \
|
||||
} else { \
|
||||
if (scale_ue8m0) { \
|
||||
per_token_group_quant_8bit_kernel<T, DST_DTYPE, false, true> \
|
||||
<<<grid, block, smem_bytes, stream>>>( \
|
||||
static_cast<T*>(input.data_ptr()), output_q.data_ptr(), \
|
||||
static_cast<float*>(output_s.data_ptr()), group_size, \
|
||||
num_groups, groups_per_block, (float)eps, (float)min_8bit, \
|
||||
(float)max_8bit); \
|
||||
} else { \
|
||||
per_token_group_quant_8bit_kernel<T, DST_DTYPE, false, false> \
|
||||
<<<grid, block, smem_bytes, stream>>>( \
|
||||
static_cast<T*>(input.data_ptr()), output_q.data_ptr(), \
|
||||
static_cast<float*>(output_s.data_ptr()), group_size, \
|
||||
num_groups, groups_per_block, (float)eps, (float)min_8bit, \
|
||||
(float)max_8bit); \
|
||||
} \
|
||||
} \
|
||||
#ifndef USE_ROCM
|
||||
#define LAUNCH_KERNEL_INST(T, DST_DTYPE, COL_MAJOR, UE8M0, SMEM_BYTES) \
|
||||
do { \
|
||||
cudaLaunchConfig_t config = {}; \
|
||||
config.gridDim = dim3(num_blocks); \
|
||||
config.blockDim = dim3(num_threads); \
|
||||
config.dynamicSmemBytes = (SMEM_BYTES); \
|
||||
config.stream = stream; \
|
||||
cudaLaunchAttribute attrs[1]; \
|
||||
attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; \
|
||||
attrs[0].val.programmaticStreamSerializationAllowed = 1; \
|
||||
config.numAttrs = 1; \
|
||||
config.attrs = attrs; \
|
||||
cudaLaunchKernelEx( \
|
||||
&config, \
|
||||
per_token_group_quant_8bit_kernel<T, DST_DTYPE, COL_MAJOR, UE8M0>, \
|
||||
static_cast<T*>(input.data_ptr()), output_q.data_ptr(), \
|
||||
static_cast<float*>(output_s.data_ptr()), group_size, num_groups, \
|
||||
groups_per_block, (float)eps, (float)min_8bit, (float)max_8bit, \
|
||||
scale_num_rows, scale_stride); \
|
||||
} while (0)
|
||||
#else
|
||||
#define LAUNCH_KERNEL_INST(T, DST_DTYPE, COL_MAJOR, UE8M0, SMEM_BYTES) \
|
||||
do { \
|
||||
per_token_group_quant_8bit_kernel<T, DST_DTYPE, COL_MAJOR, UE8M0> \
|
||||
<<<dim3(num_blocks), dim3(num_threads), (SMEM_BYTES), stream>>>( \
|
||||
static_cast<T*>(input.data_ptr()), output_q.data_ptr(), \
|
||||
static_cast<float*>(output_s.data_ptr()), group_size, \
|
||||
num_groups, groups_per_block, (float)eps, (float)min_8bit, \
|
||||
(float)max_8bit, scale_num_rows, scale_stride); \
|
||||
} while (0)
|
||||
#endif
|
||||
|
||||
#define LAUNCH_KERNEL(T, DST_DTYPE) \
|
||||
do { \
|
||||
size_t smem_bytes = \
|
||||
static_cast<size_t>(groups_per_block) * group_size * sizeof(T); \
|
||||
if (is_column_major) { \
|
||||
if (scale_ue8m0) { \
|
||||
LAUNCH_KERNEL_INST(T, DST_DTYPE, true, true, smem_bytes); \
|
||||
} else { \
|
||||
LAUNCH_KERNEL_INST(T, DST_DTYPE, true, false, smem_bytes); \
|
||||
} \
|
||||
} else { \
|
||||
if (scale_ue8m0) { \
|
||||
LAUNCH_KERNEL_INST(T, DST_DTYPE, false, true, smem_bytes); \
|
||||
} else { \
|
||||
LAUNCH_KERNEL_INST(T, DST_DTYPE, false, false, smem_bytes); \
|
||||
} \
|
||||
} \
|
||||
} while (0)
|
||||
|
||||
#define LAUNCH_REG_KERNEL_INST(T, DST_DTYPE, IS_COLUMN_MAJOR, SCALE_UE8M0, KX, \
|
||||
@@ -484,6 +503,7 @@ void per_token_group_quant_8bit(const torch::stable::Tensor& input,
|
||||
#undef LAUNCH_REG_KERNEL_INST
|
||||
#undef LAUNCH_REG_KERNEL_FLAGS
|
||||
#undef LAUNCH_KERNEL
|
||||
#undef LAUNCH_KERNEL_INST
|
||||
}
|
||||
|
||||
// Register-resident fast path for group_size==128.
|
||||
@@ -528,12 +548,12 @@ __global__ void per_token_group_quant_8bit_packed_register_kernel(
|
||||
const int mn_idx = blockIdx.x * kRowsPerBlock + row_local;
|
||||
|
||||
#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
|
||||
asm volatile("griddepcontrol.wait;");
|
||||
cudaGridDependencySynchronize();
|
||||
#endif
|
||||
|
||||
if (mn_idx >= tma_aligned_mn) {
|
||||
#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
|
||||
asm volatile("griddepcontrol.launch_dependents;");
|
||||
cudaTriggerProgrammaticLaunchCompletion();
|
||||
#endif
|
||||
return;
|
||||
}
|
||||
@@ -607,7 +627,7 @@ __global__ void per_token_group_quant_8bit_packed_register_kernel(
|
||||
*reinterpret_cast<uint4*>(group_output) = packed_out;
|
||||
|
||||
#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
|
||||
asm volatile("griddepcontrol.launch_dependents;");
|
||||
cudaTriggerProgrammaticLaunchCompletion();
|
||||
#endif
|
||||
}
|
||||
|
||||
|
||||
@@ -283,6 +283,29 @@ __quickreduce_device_inline__ int packed_rcp<nv_bfloat16>(int a) {
|
||||
return R.i;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__quickreduce_device_inline__ int packed_from_int16_pair(int16_t low,
|
||||
int16_t high);
|
||||
|
||||
template <>
|
||||
__quickreduce_device_inline__ int packed_from_int16_pair<half>(int16_t low,
|
||||
int16_t high) {
|
||||
// Convert two signed integers to one fp16x2 packed 32-bit lane.
|
||||
half2 h = __halves2half2(__int2half_rn(static_cast<int>(low)),
|
||||
__int2half_rn(static_cast<int>(high)));
|
||||
return __builtin_bit_cast(int, h);
|
||||
}
|
||||
|
||||
template <>
|
||||
__quickreduce_device_inline__ int packed_from_int16_pair<nv_bfloat16>(
|
||||
int16_t low, int16_t high) {
|
||||
// Convert two signed integers to one bf16x2 packed 32-bit lane.
|
||||
nv_bfloat16 bf_low = __float2bfloat16(static_cast<float>(low));
|
||||
nv_bfloat16 bf_high = __float2bfloat16(static_cast<float>(high));
|
||||
nv_bfloat162 bf2 = __halves2bfloat162(bf_low, bf_high);
|
||||
return *reinterpret_cast<int*>(&bf2);
|
||||
}
|
||||
|
||||
// changes dtype
|
||||
__quickreduce_device_inline__ float T2float_cast(half a) {
|
||||
return __half2float(a);
|
||||
|
||||
@@ -59,11 +59,30 @@ allreduce_prototype_twoshot(T const* A, T* B, uint32_t N, uint32_t num_blocks,
|
||||
flag_color, this->kMaxProblemSize); \
|
||||
}
|
||||
|
||||
// INT3 only retains good performance on TP2 (world_size == 2). On TP4/TP8
|
||||
// the 3-bit codec's pack/unpack overhead outweighs the reduced communication
|
||||
// volume, so INT3 is restricted to a TP2-only dispatch here.
|
||||
#define TWOSHOT_DISPATCH_TP2_ONLY(__codec) \
|
||||
if (world_size == 2) { \
|
||||
using LineCodec = __codec<T, 2>; \
|
||||
using AllReduceKernel = AllReduceTwoshot<T, LineCodec, cast_bf2half>; \
|
||||
hipLaunchKernelGGL((allreduce_prototype_twoshot<AllReduceKernel, T>), \
|
||||
dim3(grid), dim3(kBlockTwoShot), 0, stream, A, B, N, \
|
||||
num_blocks, rank, dbuffer_list, data_offset, \
|
||||
flag_color, this->kMaxProblemSize); \
|
||||
} else { \
|
||||
throw std::runtime_error( \
|
||||
"INT3 quick all-reduce is only supported for world_size == 2 " \
|
||||
"(TP2); use INT4/NONE for larger world sizes."); \
|
||||
}
|
||||
|
||||
enum QuickReduceQuantLevel {
|
||||
F16 = 0,
|
||||
INT8 = 1,
|
||||
INT6 = 2,
|
||||
INT4 = 3,
|
||||
// Keep these ids in sync with Python QuickReduceRegime enum.
|
||||
F16 = 0, // full-precision fp16/bf16 communication
|
||||
INT8 = 1, // symmetric int8 + per-block scale
|
||||
INT6 = 2, // symmetric int6 + per-block scale
|
||||
INT4 = 3, // symmetric int4 + per-block scale
|
||||
INT3 = 4, // symmetric int3 + per-block scale (TP2 only)
|
||||
};
|
||||
|
||||
struct DeviceComms {
|
||||
@@ -184,6 +203,9 @@ struct DeviceComms {
|
||||
case QuickReduceQuantLevel::INT4:
|
||||
TWOSHOT_DISPATCH(CodecQ4)
|
||||
break;
|
||||
case QuickReduceQuantLevel::INT3:
|
||||
TWOSHOT_DISPATCH_TP2_ONLY(CodecQ3)
|
||||
break;
|
||||
default:
|
||||
TWOSHOT_DISPATCH(CodecFP)
|
||||
break;
|
||||
|
||||
@@ -2,6 +2,7 @@
|
||||
|
||||
#include <hip/hip_runtime.h>
|
||||
#include "base.h"
|
||||
#include <hip/hip_fp16.h>
|
||||
|
||||
namespace quickreduce {
|
||||
|
||||
@@ -206,6 +207,168 @@ struct CodecQ4 : public CodecBase {
|
||||
}
|
||||
};
|
||||
|
||||
// Int3 symmetric quantization codec.
|
||||
// We quantize the FP16 data to block-scaled Int3 in blocks of 4 *
|
||||
// kThreadGroupSize. Uniform symmetric quantization (round-to-int + clip),
|
||||
// matching the structure of CodecQ4. Signed range is [-4, +3].
|
||||
template <typename T, int world_size>
|
||||
struct CodecQ3 : public CodecBase {
|
||||
static constexpr int kWorldSize = world_size;
|
||||
|
||||
// Layout per quantization block (32 values = 8 threads * 4 fp16x2 lanes):
|
||||
// - each thread owns 8 values and writes:
|
||||
// * q2 payload : 8 * 2 bits -> uint16 (2 bytes)
|
||||
// * q1 payload : 8 * 1 bit -> uint8 (1 byte)
|
||||
// - one scale is shared per 32 values and written by group leader.
|
||||
//
|
||||
// kRankTileStride is split as:
|
||||
// [0 .. 511] : q2 payload region (256 threads * 2 bytes)
|
||||
// [512 .. 767] : q1 payload region (256 threads * 1 byte)
|
||||
// [768 .. 895] : scale region (32 groups * 4 bytes)
|
||||
static constexpr int kRankAtoms = kAtoms / kWorldSize;
|
||||
static constexpr int kRankTileStride = 896;
|
||||
static constexpr int kRankTileQ1Offset = 512;
|
||||
static constexpr int kRankTileScaleOffset = 768;
|
||||
static constexpr int kRankTransmittedTileSize = kRankTileStride * kRankAtoms;
|
||||
static_assert(kRankTransmittedTileSize % 16 == 0,
|
||||
"kRankTransmittedTileSize must be 16B aligned.");
|
||||
|
||||
static constexpr int kRankBufferTileStride =
|
||||
kRankTileStride / sizeof(int32x4_t);
|
||||
|
||||
static constexpr int kTransmittedTileSize =
|
||||
kRankTransmittedTileSize * kWorldSize;
|
||||
|
||||
// {-1/4.0h, -1/4.0h}, f16x2_t / bf16x2_t. Sign-flipped so absmax maps
|
||||
// to -4; the sign cancels with decoding_scale on the recv side.
|
||||
static constexpr int kScaleFactor =
|
||||
std::is_same<T, half>::value ? 0xB400B400 : 0xBE80BE80;
|
||||
|
||||
// {1e-7, 1e-7}, f16x2_t
|
||||
static constexpr int kScaleEpsilon =
|
||||
std::is_same<T, half>::value ? 0x00010001 : 0x33D733D7;
|
||||
|
||||
// {-4, -4}, f16x2_t / bf16x2_t
|
||||
static constexpr int kRangeMin =
|
||||
std::is_same<T, half>::value ? 0xC400C400 : 0xC080C080;
|
||||
|
||||
// {+3, +3}, f16x2_t / bf16x2_t
|
||||
static constexpr int kRangeMax =
|
||||
std::is_same<T, half>::value ? 0x42004200 : 0x40404040;
|
||||
|
||||
// {+4, +4}, int16x2_t -- shifts signed [-4, +3] to unsigned [0, 7].
|
||||
static constexpr int kRangeBias = 0x00040004;
|
||||
|
||||
__quickreduce_device_inline__ CodecQ3(int thread, int rank)
|
||||
: CodecBase(thread, rank) {}
|
||||
|
||||
__quickreduce_device_inline__ void send(int32x4_t* __restrict__ send_buffer,
|
||||
const int32x4_t* __restrict__ data) {
|
||||
for (int k = 0; k < kRankAtoms; k++) {
|
||||
int32x4_t const atom = data[k];
|
||||
|
||||
// 1) Per-group dynamic scale (shared across 32 values).
|
||||
int wblockmax = group_abs_max<T>(atom);
|
||||
int decoding_scale = packed_mul<T>(wblockmax, kScaleFactor);
|
||||
int encoding_scale = packed_add<T>(decoding_scale, kScaleEpsilon);
|
||||
encoding_scale = packed_rcp<T>(encoding_scale);
|
||||
|
||||
// 2) Scale + clip to signed int3 range [-4, +3].
|
||||
int32x4_t w;
|
||||
for (int i = 0; i < 4; i++) {
|
||||
w[i] = packed_mul<T>(atom[i], encoding_scale);
|
||||
w[i] = packed_max<T>(w[i], kRangeMin);
|
||||
w[i] = packed_min<T>(w[i], kRangeMax);
|
||||
}
|
||||
|
||||
// 3) Round to integer and bias to unsigned domain [0, 7].
|
||||
int32x4_t q;
|
||||
{
|
||||
int16_t* qi = reinterpret_cast<int16_t*>(&q);
|
||||
T* wh = reinterpret_cast<T*>(&w);
|
||||
for (int i = 0; i < 8; i++) qi[i] = (int16_t)rintf(T2float_cast(wh[i]));
|
||||
|
||||
for (int i = 0; i < 4; i++) {
|
||||
q[i] = packed_add<int16_t>(q[i], kRangeBias);
|
||||
}
|
||||
}
|
||||
|
||||
// 4) Split each 3-bit unsigned value into low-2-bit and high-1-bit
|
||||
// halves, packed into one uint16 (low 2 bits per value) plus one
|
||||
// uint8 (high 1 bit per value).
|
||||
uint16_t q2w = 0;
|
||||
uint8_t q1w = 0;
|
||||
{
|
||||
int16_t* tw = reinterpret_cast<int16_t*>(&q);
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 8; i++) {
|
||||
uint32_t v = static_cast<uint32_t>(tw[i]) & 0x7u;
|
||||
q2w |= static_cast<uint16_t>((v & 0x3u) << (i * 2));
|
||||
q1w |= static_cast<uint8_t>(((v >> 2) & 0x1u) << i);
|
||||
}
|
||||
}
|
||||
|
||||
uint8_t* atom_ptr =
|
||||
reinterpret_cast<uint8_t*>(send_buffer + k * kRankBufferTileStride);
|
||||
uint16_t* q2w_ptr = reinterpret_cast<uint16_t*>(atom_ptr) + thread;
|
||||
uint8_t* q1w_ptr =
|
||||
reinterpret_cast<uint8_t*>(atom_ptr + kRankTileQ1Offset) + thread;
|
||||
int* qs_ptr = reinterpret_cast<int*>(atom_ptr + kRankTileScaleOffset) +
|
||||
(thread / 8);
|
||||
|
||||
__builtin_nontemporal_store(q2w, q2w_ptr);
|
||||
*q1w_ptr = q1w;
|
||||
if (threadIdx.x == group_leader) {
|
||||
__builtin_nontemporal_store(decoding_scale, qs_ptr);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
__quickreduce_device_inline__ void recv(int32x4_t** __restrict__ recv_buffer,
|
||||
int32x4_t* __restrict__ data) {
|
||||
for (int k = 0; k < kRankAtoms; k++) {
|
||||
uint8_t* atom_ptr = reinterpret_cast<uint8_t*>(*recv_buffer);
|
||||
uint16_t* q2w_ptr = reinterpret_cast<uint16_t*>(atom_ptr) + thread;
|
||||
uint8_t* q1w_ptr =
|
||||
reinterpret_cast<uint8_t*>(atom_ptr + kRankTileQ1Offset) + thread;
|
||||
int* qs_ptr = reinterpret_cast<int*>(atom_ptr + kRankTileScaleOffset) +
|
||||
(thread / 8);
|
||||
|
||||
uint16_t q2w = __builtin_nontemporal_load(q2w_ptr);
|
||||
uint8_t q1w = *q1w_ptr;
|
||||
int qs = __builtin_nontemporal_load(qs_ptr);
|
||||
|
||||
*recv_buffer += kRankBufferTileStride;
|
||||
|
||||
// Unpack unsigned values [0, 7] then shift back to signed domain
|
||||
// [-4, +3] by adding kRangeMin.
|
||||
int32x4_t w;
|
||||
{
|
||||
int16_t qv[8];
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 8; i++) {
|
||||
uint32_t low2 = (q2w >> (2 * i)) & 0x3u;
|
||||
uint32_t high1 = (q1w >> i) & 0x1u;
|
||||
qv[i] = static_cast<int16_t>(low2 | (high1 << 2));
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 4; i++) {
|
||||
int qpack = packed_from_int16_pair<T>(qv[2 * i], qv[2 * i + 1]);
|
||||
w[i] = packed_add<T>(qpack, kRangeMin);
|
||||
}
|
||||
}
|
||||
|
||||
// Apply decode scale to reconstruct fp16/bf16 lanes.
|
||||
for (int i = 0; i < 4; i++) {
|
||||
w[i] = packed_mul<T>(w[i], qs);
|
||||
}
|
||||
|
||||
data[k] = w;
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
// Int6 symmetric quantization codec.
|
||||
// We quantize the FP16 data to block-scaled Int6 in blocks of 4 *
|
||||
// kThreadGroupSize.
|
||||
@@ -377,7 +540,6 @@ struct CodecQ6 : public CodecBase {
|
||||
w[i] = packed_mul<T>(w[i], qs);
|
||||
}
|
||||
|
||||
// That's pretty much it...
|
||||
data[k] = w;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1045,7 +1045,7 @@ __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_QKV_mfma4_kernel(
|
||||
const scalar_t* q_ptr =
|
||||
q + query_start_off * q_stride + wg_start_head_idx * HEAD_SIZE;
|
||||
const _B16x8* q_ptrh8 = reinterpret_cast<const _B16x8*>(q_ptr);
|
||||
const int qhead_elemh8 = laneid / 4;
|
||||
const int qhead_elemh8 = MIN(laneid / 4, HEAD_SIZE / 8 - 1);
|
||||
|
||||
for (int h = 0; h < QHLOOP - 1; h++) {
|
||||
const int qhead_idx = h * 4 + lane4id;
|
||||
|
||||
+53
-21
@@ -70,6 +70,15 @@ bool on_gfx12() {
|
||||
return result;
|
||||
}
|
||||
|
||||
bool on_gfx1151() {
|
||||
static const bool result = [] {
|
||||
const auto* dprops = at::cuda::getCurrentDeviceProperties();
|
||||
const std::string device_arch = dprops->gcnArchName;
|
||||
return device_arch.find("gfx1151") != std::string::npos;
|
||||
}();
|
||||
return result;
|
||||
}
|
||||
|
||||
#if defined(NDEBUG)
|
||||
#undef NDEBUG
|
||||
#include <assert.h>
|
||||
@@ -1237,6 +1246,45 @@ torch::Tensor wvSplitK(const at::Tensor& in_a, const at::Tensor& in_b,
|
||||
WVSPLITK_CFG(_THRDS, _WVPRGRP, 4, 2, __N) \
|
||||
}
|
||||
|
||||
// WVSPLITK_CFG arguments are: (THRDS, WVPRGRP, YTILE, UNRL, N).
|
||||
// THRDS = wavefront width (32 on GFX11/GFX12, 64 on GFX9)
|
||||
// WVPRGRP= waves per group (always 16)
|
||||
// YTILE = output rows per thread tile
|
||||
// UNRL = K-loop unroll factor
|
||||
// N = batch size (passed through from the switch in wvSplitK)
|
||||
#define WVSPLIT_TILE(_sYT, __N) \
|
||||
{ \
|
||||
if (on_gfx1151()) { \
|
||||
bool fit_lds = (Kbp_in * N_in <= max_lds_len); \
|
||||
if (_sYT <= 1) \
|
||||
WVSPLITK_CFG(/*THRDS=*/32, /*WVPRGRP=*/16, /*YTILE=*/1, /*UNRL=*/4, \
|
||||
__N) \
|
||||
else if ((K_in % 1024 == 512) && K_in >= 1536 && \
|
||||
(_sYT >= 40 || K_in >= 4096)) \
|
||||
WVSPLITK_CFG(/*THRDS=*/32, /*WVPRGRP=*/16, /*YTILE=*/4, /*UNRL=*/1, \
|
||||
__N) \
|
||||
else if (K_in < 1024) \
|
||||
WVSPLITK_CFG(/*THRDS=*/32, /*WVPRGRP=*/16, /*YTILE=*/2, /*UNRL=*/4, \
|
||||
__N) \
|
||||
else if (K_in <= 2048 && (__N >= 2 || _sYT <= 26)) \
|
||||
WVSPLITK_CFG(/*THRDS=*/32, /*WVPRGRP=*/16, /*YTILE=*/1, /*UNRL=*/4, \
|
||||
__N) \
|
||||
else if (__N >= 2 && !fit_lds) \
|
||||
WVSPLITK_CFG(/*THRDS=*/32, /*WVPRGRP=*/16, /*YTILE=*/1, /*UNRL=*/4, \
|
||||
__N) \
|
||||
else if (__N == 1) \
|
||||
WVSPLITK_CFG(/*THRDS=*/32, /*WVPRGRP=*/16, /*YTILE=*/1, /*UNRL=*/2, \
|
||||
__N) \
|
||||
else \
|
||||
WVSPLITK_CFG(/*THRDS=*/32, /*WVPRGRP=*/16, /*YTILE=*/1, /*UNRL=*/1, \
|
||||
__N) \
|
||||
} else if (on_gfx1x()) { /* gfx1100/gfx1150/GFX12, wave32 */ \
|
||||
WVSPLIT_TILE_CFG(/*THRDS=*/32, /*WVPRGRP=*/16, _sYT, __N) \
|
||||
} else { /* GFX9, wave64 */ \
|
||||
WVSPLIT_TILE_CFG(/*THRDS=*/64, /*WVPRGRP=*/16, _sYT, __N) \
|
||||
} \
|
||||
}
|
||||
|
||||
AT_DISPATCH_REDUCED_FLOATING_TYPES(in_b.scalar_type(), "wvSplitK", [&] {
|
||||
using fptype = typename scalar<scalar_t>::type;
|
||||
fptype* af4 = reinterpret_cast<fptype*>(in_a.data_ptr());
|
||||
@@ -1251,37 +1299,21 @@ torch::Tensor wvSplitK(const at::Tensor& in_a, const at::Tensor& in_b,
|
||||
// then cut the active waves to balance their distribution...
|
||||
int sYT = (M_in + CuCount * 4 - 1) / (CuCount * 4);
|
||||
|
||||
const bool use_wave32 = on_gfx1x();
|
||||
switch (N_in) {
|
||||
case 1:
|
||||
if (use_wave32)
|
||||
WVSPLIT_TILE_CFG(32, 16, sYT, 1)
|
||||
else
|
||||
WVSPLIT_TILE_CFG(64, 16, sYT, 1)
|
||||
WVSPLIT_TILE(sYT, 1)
|
||||
break;
|
||||
case 2:
|
||||
if (use_wave32)
|
||||
WVSPLIT_TILE_CFG(32, 16, sYT, 2)
|
||||
else
|
||||
WVSPLIT_TILE_CFG(64, 16, sYT, 2)
|
||||
WVSPLIT_TILE(sYT, 2)
|
||||
break;
|
||||
case 3:
|
||||
if (use_wave32)
|
||||
WVSPLIT_TILE_CFG(32, 16, sYT, 3)
|
||||
else
|
||||
WVSPLIT_TILE_CFG(64, 16, sYT, 3)
|
||||
WVSPLIT_TILE(sYT, 3)
|
||||
break;
|
||||
case 4:
|
||||
if (use_wave32)
|
||||
WVSPLIT_TILE_CFG(32, 16, sYT, 4)
|
||||
else
|
||||
WVSPLIT_TILE_CFG(64, 16, sYT, 4)
|
||||
WVSPLIT_TILE(sYT, 4)
|
||||
break;
|
||||
case 5:
|
||||
if (use_wave32)
|
||||
WVSPLIT_TILE_CFG(32, 16, sYT, 5)
|
||||
else
|
||||
WVSPLIT_TILE_CFG(64, 16, sYT, 5)
|
||||
WVSPLIT_TILE(sYT, 5)
|
||||
break;
|
||||
default:
|
||||
throw std::runtime_error(
|
||||
|
||||
@@ -9,7 +9,7 @@ ARG PYTORCH_AUDIO_BRANCH="v2.9.0"
|
||||
ARG PYTORCH_AUDIO_REPO="https://github.com/pytorch/audio.git"
|
||||
ARG FA_BRANCH="0e60e394"
|
||||
ARG FA_REPO="https://github.com/Dao-AILab/flash-attention.git"
|
||||
ARG AITER_BRANCH="v0.1.13.post1"
|
||||
ARG AITER_BRANCH="v0.1.16.post2"
|
||||
ARG AITER_REPO="https://github.com/ROCm/aiter.git"
|
||||
ARG MORI_BRANCH="v1.1.0"
|
||||
ARG MORI_REPO="https://github.com/ROCm/mori.git"
|
||||
@@ -244,7 +244,7 @@ RUN pip install pyyaml && cd aiter \
|
||||
export HIP_CLANG_PATH=/opt/sccache-wrappers \
|
||||
&& sccache --show-stats; \
|
||||
fi \
|
||||
&& PREBUILD_KERNELS=1 GPU_ARCHS=${AITER_ROCM_ARCH} python3 setup.py bdist_wheel --dist-dir=dist \
|
||||
&& PREBUILD_KERNELS=1 AITER_USE_SYSTEM_TRITON=1 GPU_ARCHS=${AITER_ROCM_ARCH} python3 setup.py bdist_wheel --dist-dir=dist \
|
||||
&& if [ "$USE_SCCACHE" = "1" ]; then sccache --show-stats; fi \
|
||||
&& ls /app/aiter/dist/*.whl
|
||||
RUN mkdir -p /app/install && cp /app/aiter/dist/*.whl /app/install
|
||||
|
||||
@@ -159,7 +159,7 @@ Priority is **1 = highest** (tried first).
|
||||
|
||||
| Backend | Version | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | Non-Causal | MM Prefix | DCP | Attention Types | Compute Cap. |
|
||||
| ------- | ------- | ------ | --------- | ----------- | ---------- | ---- | ---------- | --------- | --- | --------------- | ------------ |
|
||||
| `CPU_ATTN` | | fp16, bf16, fp32 | `auto`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | 32, 64, 80, 96, 112, 128, 160, 192, 224, 256, 512 | ❌ | ❌ | ❌ | ❌ | All | N/A |
|
||||
| `CPU_ATTN` | | fp16, bf16, fp32 | `auto`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | 32, 64, 80, 96, 112, 128, 160, 192, 224, 256, 512 | ❌ | ✅ | ❌ | ❌ | All | N/A |
|
||||
| `FLASHINFER` | Native† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64, 128, 256, 512, 1024 | 64, 128, 256, 512 | ❌ | ✅ | ❌ | ✅ | Decoder | 8.x-9.x |
|
||||
| `FLASHINFER` | TRTLLM† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2`, `nvfp4` | 16, 32, 64, 128, 256, 512, 1024 | 64, 128, 256, 512 | ✅ | ✅ | ❌ | ✅ | Decoder | 10.x |
|
||||
| `FLASH_ATTN` | FA2* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ✅ | ❌ | ✅ | All | ≥8.0 |
|
||||
|
||||
@@ -134,7 +134,8 @@ Models opt-in to encoder CUDA Graphs by implementing the [SupportsEncoderCudaGra
|
||||
| `Qwen2VLForConditionalGeneration` | `Qwen2-VL` | ✅︎ | ✅︎ | ❌︎ |
|
||||
| `Qwen2_5_VLForConditionalGeneration` | `Qwen2.5-VL` | ✅︎ | ✅︎ | ❌︎ |
|
||||
| `Qwen3VLForConditionalGeneration` | `Qwen3-VL` | ✅︎ | ✅︎ | ❌︎ |
|
||||
| `Qwen3_5ForConditionalGeneration` | `Qwen3.5` | ✅︎ | ✅︎ | ❌︎ |
|
||||
| `Qwen3_5ForConditionalGeneration` | `Qwen3.5`, `Qwen3.6` | ✅︎ | ✅︎ | ❌︎ |
|
||||
| `Qwen3_5MoeForConditionalGeneration` | `Qwen3.5-MoE`, `Qwen3.6-MoE` | ✅︎ | ✅︎ | ❌︎ |
|
||||
| `Step3VLForConditionalGeneration` | `Step3-VL` | ✅︎ | ❌︎ | ❌︎ |
|
||||
|
||||
!!! note
|
||||
|
||||
@@ -159,7 +159,7 @@ For simplicity, we assume `block_size=1` in this section.
|
||||
|
||||
### High level idea
|
||||
|
||||
The block pool uses a dict similar to `tuple(block_hash, group_id) -> block` to catch the full blocks. That means the same tokens of different groups are cached and evicted independently.
|
||||
The block pool uses a dict similar to `tuple(block_hash, group_id) -> block` to cache the full blocks. That means the same tokens of different groups are cached and evicted independently.
|
||||
|
||||
When a new request comes in, we check the cache hit prefix of each group, and return the intersection of these groups as the cached prefix of the request. See below for the detailed algorithm for checking the cache hit of one group & performing the intersection.
|
||||
|
||||
|
||||
@@ -52,7 +52,7 @@ __device__ void paged_attention_kernel(
|
||||
)
|
||||
```
|
||||
|
||||
There are also a list of template arguments above the function
|
||||
There is also a list of template arguments above the function
|
||||
signature that are determined during compilation time. `scalar_t`
|
||||
represents the data type of the query, key, and value data elements,
|
||||
such as FP16. `HEAD_SIZE` indicates the number of elements in each
|
||||
@@ -178,7 +178,7 @@ const scalar_t* k_ptr = k_cache + physical_block_number * kv_block_stride
|
||||
+ physical_block_offset * x;
|
||||
```
|
||||
|
||||
Unlike to `q_ptr`, `k_ptr` in each thread will point to different
|
||||
Unlike `q_ptr`, `k_ptr` in each thread will point to different
|
||||
key token at different iterations. As shown above, that `k_ptr`
|
||||
points to key token data based on `k_cache` at assigned block,
|
||||
assigned head and assigned token.
|
||||
|
||||
@@ -366,7 +366,6 @@ th {
|
||||
| ------------ | ------ | ----------------- | -------------------- | ------------------------- |
|
||||
| `AfmoeForCausalLM` | Afmoe | TBA | ✅︎ | ✅︎ |
|
||||
| `ApertusForCausalLM` | Apertus | `swiss-ai/Apertus-8B-2509`, `swiss-ai/Apertus-70B-Instruct-2509`, etc. | ✅︎ | ✅︎ |
|
||||
| `AquilaForCausalLM` | Aquila, Aquila2 | `BAAI/Aquila-7B`, `BAAI/AquilaChat-7B`, etc. | ✅︎ | ✅︎ |
|
||||
| `ArceeForCausalLM` | Arcee (AFM) | `arcee-ai/AFM-4.5B-Base`, etc. | ✅︎ | ✅︎ |
|
||||
| `ArcticForCausalLM` | Arctic | `Snowflake/snowflake-arctic-base`, `Snowflake/snowflake-arctic-instruct`, etc. | | ✅︎ |
|
||||
| `AXK1ForCausalLM` | A.X-K1 | `skt/A.X-K1`, etc. | | ✅︎ |
|
||||
|
||||
+2
-2
@@ -129,8 +129,8 @@ extend-exclude = ["tests/models/fixtures/*", "tests/prompts/*", "tests/tokenizer
|
||||
"docs/governance/process.md", "docs/assets/contributing/vllm_bench_serve_timeline.html",
|
||||
"tests/v1/engine/test_fast_incdec_prefix_err.py", ".git/*", "csrc/cpu/sgl-kernels/*",
|
||||
"rust/src/chat/src/renderer/deepseek_v32/fixtures/*",
|
||||
"rust/src/tool-parser/src/gemma4.rs", "rust/src/text/src/output/decoded.rs",
|
||||
"rust/src/tokenizer/src/incremental.rs", "rust/src/reasoning-parser/src/tests.rs"]
|
||||
"rust/src/parser/src/tool/gemma4.rs", "rust/src/text/src/output/decoded.rs",
|
||||
"rust/src/tokenizer/src/incremental.rs", "rust/src/parser/src/reasoning/tests.rs"]
|
||||
ignore-hidden = false
|
||||
|
||||
[tool.typos.default]
|
||||
|
||||
Generated
+14
-22
@@ -5733,10 +5733,9 @@ dependencies = [
|
||||
"uuid",
|
||||
"vllm-engine-core-client",
|
||||
"vllm-llm",
|
||||
"vllm-reasoning-parser",
|
||||
"vllm-parser",
|
||||
"vllm-text",
|
||||
"vllm-tokenizer",
|
||||
"vllm-tool-parser",
|
||||
"xgrammar-structural-tag",
|
||||
"zeromq",
|
||||
]
|
||||
@@ -5872,11 +5871,22 @@ dependencies = [
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "vllm-reasoning-parser"
|
||||
name = "vllm-parser"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"criterion",
|
||||
"easy-ext",
|
||||
"expect-test",
|
||||
"futures",
|
||||
"openai-protocol",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"thiserror 2.0.18",
|
||||
"thiserror-ext",
|
||||
"tool-parser",
|
||||
"vllm-tokenizer",
|
||||
"winnow",
|
||||
"xgrammar-structural-tag",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -5978,24 +5988,6 @@ dependencies = [
|
||||
"tracing",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "vllm-tool-parser"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"criterion",
|
||||
"easy-ext",
|
||||
"expect-test",
|
||||
"futures",
|
||||
"openai-protocol",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"thiserror 2.0.18",
|
||||
"thiserror-ext",
|
||||
"tool-parser",
|
||||
"winnow",
|
||||
"xgrammar-structural-tag",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "vllm-tool-parser-py"
|
||||
version = "0.1.0"
|
||||
@@ -6004,7 +5996,7 @@ dependencies = [
|
||||
"pythonize",
|
||||
"serde_json",
|
||||
"thiserror-ext",
|
||||
"vllm-tool-parser",
|
||||
"vllm-parser",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
||||
+3
-5
@@ -7,12 +7,11 @@ members = [
|
||||
"src/managed-engine",
|
||||
"src/metrics",
|
||||
"src/mock-engine",
|
||||
"src/reasoning-parser",
|
||||
"src/parser",
|
||||
"src/parser/python",
|
||||
"src/server",
|
||||
"src/text",
|
||||
"src/tokenizer",
|
||||
"src/tool-parser",
|
||||
"src/tool-parser/python",
|
||||
]
|
||||
resolver = "3"
|
||||
|
||||
@@ -117,11 +116,10 @@ vllm-engine-core-client = { path = "src/engine-core-client" }
|
||||
vllm-llm = { path = "src/llm" }
|
||||
vllm-managed-engine = { path = "src/managed-engine" }
|
||||
vllm-metrics = { path = "src/metrics" }
|
||||
vllm-reasoning-parser = { path = "src/reasoning-parser" }
|
||||
vllm-parser = { path = "src/parser" }
|
||||
vllm-server = { path = "src/server" }
|
||||
vllm-text = { path = "src/text" }
|
||||
vllm-tokenizer = { path = "src/tokenizer" }
|
||||
vllm-tool-parser = { path = "src/tool-parser" }
|
||||
winnow = "1.0.2"
|
||||
xgrammar-structural-tag = "0.1.0"
|
||||
zeromq = { version = "0.6.0", default-features = false, features = [
|
||||
|
||||
@@ -30,10 +30,9 @@ trait-set.workspace = true
|
||||
uuid.workspace = true
|
||||
vllm-engine-core-client.workspace = true
|
||||
vllm-llm.workspace = true
|
||||
vllm-reasoning-parser.workspace = true
|
||||
vllm-parser.workspace = true
|
||||
vllm-text.workspace = true
|
||||
vllm-tokenizer.workspace = true
|
||||
vllm-tool-parser.workspace = true
|
||||
xgrammar-structural-tag.workspace = true
|
||||
|
||||
[dev-dependencies]
|
||||
|
||||
@@ -1,44 +1,34 @@
|
||||
//! Default output processing pipeline.
|
||||
|
||||
mod reasoning;
|
||||
mod structural_tag;
|
||||
mod tool;
|
||||
mod unified;
|
||||
|
||||
use std::sync::Once;
|
||||
|
||||
use futures::{Stream, StreamExt as _};
|
||||
use futures::StreamExt as _;
|
||||
use tracing::info;
|
||||
use trait_set::trait_set;
|
||||
use vllm_parser::unified::{CombinedParser, UnifiedParser};
|
||||
use vllm_text::tokenizer::DynTokenizer;
|
||||
|
||||
use self::reasoning::reasoning_event_stream;
|
||||
use self::structural_tag::apply_structural_tag_constraint;
|
||||
use self::tool::tool_event_stream;
|
||||
use self::unified::unified_event_stream;
|
||||
use super::structured::structured_chat_event_stream;
|
||||
use crate::error::Result;
|
||||
use crate::output::{
|
||||
AssistantEvent, ChatOutputProcessor, ContentEvent, DynChatEventStream,
|
||||
DynDecodedTextEventStream,
|
||||
};
|
||||
use crate::output::{ChatOutputProcessor, DynChatEventStream, DynDecodedTextEventStream};
|
||||
use crate::parser::ParserSelection;
|
||||
use crate::parser::reasoning::{ReasoningParser, ReasoningParserFactory};
|
||||
use crate::parser::tool::{ToolParser, ToolParserFactory};
|
||||
use crate::request::ChatRequest;
|
||||
use crate::{Error, Result as ChatResult};
|
||||
|
||||
trait_set! {
|
||||
trait ContentEventStream = Stream<Item = Result<ContentEvent>> + Send + 'static;
|
||||
}
|
||||
|
||||
/// Default request-scoped output processor used by Hugging Face style chat
|
||||
/// backends.
|
||||
///
|
||||
/// This implementation assumes the backend already emitted decoded text deltas,
|
||||
/// then optionally layers reasoning parsing and tool-call parsing before
|
||||
/// then optionally layers unified reasoning and tool-call parsing before
|
||||
/// assembling final structured chat events.
|
||||
pub struct DefaultChatOutputProcessor {
|
||||
reasoning_parser: Option<Box<dyn ReasoningParser>>,
|
||||
tool_parser: Option<Box<dyn ToolParser>>,
|
||||
parser: Box<dyn UnifiedParser>,
|
||||
parallel_tool_calls: bool,
|
||||
}
|
||||
|
||||
@@ -66,16 +56,17 @@ impl DefaultChatOutputProcessor {
|
||||
} else {
|
||||
None
|
||||
};
|
||||
let reasoning_parser = Self::resolve_optional_reasoning_parser(
|
||||
request,
|
||||
model_id,
|
||||
tokenizer,
|
||||
reasoning_parser,
|
||||
)?;
|
||||
let reasoning_parser =
|
||||
Self::resolve_optional_reasoning_parser(model_id, tokenizer, reasoning_parser)?;
|
||||
let parser: Box<dyn UnifiedParser> =
|
||||
Box::new(CombinedParser::new(reasoning_parser, tool_parser));
|
||||
|
||||
if parser.preserve_special_tokens() {
|
||||
request.decode_options.skip_special_tokens = false;
|
||||
}
|
||||
|
||||
Ok(Self {
|
||||
reasoning_parser,
|
||||
tool_parser,
|
||||
parser,
|
||||
parallel_tool_calls: request.parallel_tool_calls,
|
||||
})
|
||||
}
|
||||
@@ -87,8 +78,7 @@ impl DefaultChatOutputProcessor {
|
||||
/// content is treated as opaque text.
|
||||
pub fn plain_text_only() -> Self {
|
||||
Self {
|
||||
reasoning_parser: None,
|
||||
tool_parser: None,
|
||||
parser: Box::new(CombinedParser::plain_text_only()),
|
||||
parallel_tool_calls: true,
|
||||
}
|
||||
}
|
||||
@@ -112,10 +102,6 @@ impl DefaultChatOutputProcessor {
|
||||
|
||||
let parser = factory.create(parser_name, &request.tools)?;
|
||||
|
||||
if parser.preserve_special_tokens() {
|
||||
request.decode_options.skip_special_tokens = false;
|
||||
}
|
||||
|
||||
apply_structural_tag_constraint(request, parser.as_ref())?;
|
||||
|
||||
TOOL_PARSER_LOG_ONCE.call_once(|| info!(parser_name, "using tool parser"));
|
||||
@@ -123,7 +109,6 @@ impl DefaultChatOutputProcessor {
|
||||
}
|
||||
|
||||
fn resolve_optional_reasoning_parser(
|
||||
request: &mut ChatRequest,
|
||||
model_id: &str,
|
||||
tokenizer: DynTokenizer,
|
||||
selection: &ParserSelection,
|
||||
@@ -142,10 +127,6 @@ impl DefaultChatOutputProcessor {
|
||||
|
||||
let parser = factory.create(parser_name, tokenizer)?;
|
||||
|
||||
if parser.preserve_special_tokens() {
|
||||
request.decode_options.skip_special_tokens = false;
|
||||
}
|
||||
|
||||
REASONING_PARSER_LOG_ONCE.call_once(|| info!(parser_name, "using reasoning parser"));
|
||||
Ok(Some(parser))
|
||||
}
|
||||
@@ -156,16 +137,14 @@ static REASONING_PARSER_LOG_ONCE: Once = Once::new();
|
||||
|
||||
impl ChatOutputProcessor for DefaultChatOutputProcessor {
|
||||
/// Transforms a raw generate-output token stream into structured chat
|
||||
/// events through three sequential stages once text decoding has
|
||||
/// events through two sequential stages once text decoding has
|
||||
/// already happened:
|
||||
///
|
||||
/// 1. [`reasoning_event_stream`] — reasoning/content separation
|
||||
/// 2. [`tool_event_stream`] — tool-call parsing
|
||||
/// 3. [`structured_chat_event_stream`] — final block assembly
|
||||
/// 1. [`unified_event_stream`] — reasoning and tool-call parsing
|
||||
/// 2. [`structured_chat_event_stream`] — final block assembly
|
||||
fn process(self: Box<Self>, decoded: DynDecodedTextEventStream) -> Result<DynChatEventStream> {
|
||||
let reasoning = reasoning_event_stream(decoded, self.reasoning_parser);
|
||||
let tool = tool_event_stream(reasoning, self.tool_parser);
|
||||
let structured = structured_chat_event_stream(tool, self.parallel_tool_calls);
|
||||
let parsed = unified_event_stream(decoded, self.parser);
|
||||
let structured = structured_chat_event_stream(parsed, self.parallel_tool_calls);
|
||||
|
||||
Ok(structured.boxed())
|
||||
}
|
||||
|
||||
@@ -1,509 +0,0 @@
|
||||
//! Adapts decoded text updates into reasoning-aware assistant deltas.
|
||||
//!
|
||||
//! This stage sits between low-level token decoding and final block assembly.
|
||||
//! It is the only place in the new pipeline that understands reasoning
|
||||
//! separation: `decoded.rs` still only produces plain text deltas, while later
|
||||
//! stages consume the semantic `Text` / `Reasoning` split emitted here.
|
||||
|
||||
use asynk_strim_attr::{TryYielder, try_stream};
|
||||
use futures::{StreamExt as _, pin_mut};
|
||||
use thiserror_ext::AsReport;
|
||||
use tracing::warn;
|
||||
use vllm_text::output::DecodedTextEvent;
|
||||
|
||||
use super::ContentEvent;
|
||||
use crate::Result;
|
||||
use crate::error::Error;
|
||||
use crate::event::AssistantBlockKind;
|
||||
use crate::output::DecodedTextEventStream;
|
||||
use crate::parser::reasoning::{ReasoningDelta, ReasoningParser};
|
||||
|
||||
/// Per-stream reasoning parsing state.
|
||||
struct ReasoningState {
|
||||
/// Reasoning parser for the current model family.
|
||||
parser: Box<dyn ReasoningParser>,
|
||||
/// Whether reasoning parsing has already failed for this stream.
|
||||
parser_failed: bool,
|
||||
}
|
||||
|
||||
impl ReasoningState {
|
||||
/// Create one fresh reasoning-adaptation state for a new streamed response.
|
||||
fn new(parser: Box<dyn ReasoningParser>) -> Self {
|
||||
Self {
|
||||
parser,
|
||||
parser_failed: false,
|
||||
}
|
||||
}
|
||||
|
||||
/// Convert one decoded text delta into zero or more semantic assistant
|
||||
/// deltas.
|
||||
fn process_delta(&mut self, delta: String) -> Vec<ContentEvent> {
|
||||
// If the parser has already failed, skip parsing and return plain text deltas.
|
||||
if self.parser_failed {
|
||||
return vec![ContentEvent::TextDelta {
|
||||
kind: AssistantBlockKind::Text,
|
||||
delta,
|
||||
}];
|
||||
}
|
||||
|
||||
let mut events = Vec::new();
|
||||
|
||||
match self.parser.push(&delta) {
|
||||
Ok(result) => {
|
||||
push_reasoning_delta(&mut events, result);
|
||||
}
|
||||
Err(error) => {
|
||||
if !self.parser_failed {
|
||||
warn!(
|
||||
error = %error.as_report(),
|
||||
"reasoning parser failed; falling back to plain text deltas"
|
||||
);
|
||||
self.parser_failed = true;
|
||||
}
|
||||
push_text_delta(&mut events, AssistantBlockKind::Text, delta);
|
||||
}
|
||||
}
|
||||
|
||||
events
|
||||
}
|
||||
|
||||
/// Initialize parser state once prompt token IDs are available.
|
||||
fn initialize(&mut self, prompt_token_ids: &[u32]) {
|
||||
if self.parser_failed {
|
||||
return;
|
||||
}
|
||||
|
||||
match self.parser.initialize(prompt_token_ids) {
|
||||
Ok(()) => {}
|
||||
Err(error) => {
|
||||
warn!(
|
||||
error = %error.as_report(),
|
||||
"failed to initialize reasoning parser; falling back to plain text deltas"
|
||||
);
|
||||
self.parser_failed = true;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Flush any parser-held partial delimiter state at end of stream.
|
||||
fn finish(&mut self) -> Vec<ContentEvent> {
|
||||
if self.parser_failed {
|
||||
return Vec::new();
|
||||
}
|
||||
|
||||
match self.parser.finish() {
|
||||
Ok(result) => {
|
||||
let mut events = Vec::new();
|
||||
push_reasoning_delta(&mut events, result);
|
||||
events
|
||||
}
|
||||
Err(error) => {
|
||||
warn!(error = %error.as_report(), "failed to flush reasoning parser state");
|
||||
Vec::new()
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Push one semantic text delta if it is non-empty.
|
||||
fn push_text_delta(events: &mut Vec<ContentEvent>, kind: AssistantBlockKind, delta: String) {
|
||||
if delta.is_empty() {
|
||||
return;
|
||||
}
|
||||
events.push(ContentEvent::TextDelta { kind, delta });
|
||||
}
|
||||
|
||||
/// Convert one parsed reasoning delta into zero or more content events.
|
||||
fn push_reasoning_delta(events: &mut Vec<ContentEvent>, delta: ReasoningDelta) {
|
||||
if let Some(reasoning) = delta.reasoning {
|
||||
push_text_delta(events, AssistantBlockKind::Reasoning, reasoning);
|
||||
}
|
||||
if let Some(content) = delta.content {
|
||||
push_text_delta(events, AssistantBlockKind::Text, content);
|
||||
}
|
||||
}
|
||||
|
||||
/// Wrap one decoded-text stream into the internal reasoning event stream.
|
||||
#[try_stream]
|
||||
pub(crate) async fn reasoning_event_stream(
|
||||
decoded_stream: impl DecodedTextEventStream,
|
||||
reasoning_parser: Option<Box<dyn ReasoningParser>>,
|
||||
mut y: TryYielder<ContentEvent, Error>,
|
||||
) -> Result<()> {
|
||||
pin_mut!(decoded_stream);
|
||||
|
||||
// Without a parser, pass through as plain text deltas.
|
||||
let Some(reasoning_parser) = reasoning_parser else {
|
||||
while let Some(event) = decoded_stream.next().await.transpose()? {
|
||||
for next in ContentEvent::from_decoded_plain_text(event) {
|
||||
y.yield_ok(next).await;
|
||||
}
|
||||
}
|
||||
return Ok(());
|
||||
};
|
||||
|
||||
let mut state = ReasoningState::new(reasoning_parser);
|
||||
|
||||
while let Some(event) = decoded_stream.next().await.transpose()? {
|
||||
match event {
|
||||
DecodedTextEvent::Start {
|
||||
prompt_token_ids,
|
||||
prompt_logprobs,
|
||||
} => {
|
||||
state.initialize(&prompt_token_ids);
|
||||
y.yield_ok(ContentEvent::Start {
|
||||
prompt_token_ids,
|
||||
prompt_logprobs,
|
||||
})
|
||||
.await;
|
||||
}
|
||||
DecodedTextEvent::TextDelta {
|
||||
delta,
|
||||
token_ids,
|
||||
logprobs,
|
||||
finished,
|
||||
} => {
|
||||
for next in state.process_delta(delta) {
|
||||
y.yield_ok(next).await;
|
||||
}
|
||||
if logprobs.is_some() || !token_ids.is_empty() {
|
||||
y.yield_ok(ContentEvent::LogprobsDelta {
|
||||
logprobs,
|
||||
token_ids,
|
||||
})
|
||||
.await;
|
||||
}
|
||||
if let Some(finished) = finished {
|
||||
for next in state.finish() {
|
||||
y.yield_ok(next).await;
|
||||
}
|
||||
y.yield_ok(ContentEvent::Done {
|
||||
usage: finished.usage,
|
||||
finish_reason: finished.finish_reason,
|
||||
kv_transfer_params: finished.kv_transfer_params,
|
||||
})
|
||||
.await;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use futures::{StreamExt as _, stream};
|
||||
use vllm_llm::FinishReason;
|
||||
use vllm_text::output::{
|
||||
DecodedLogprobs, DecodedPositionLogprobs, DecodedTextEvent, DecodedTokenLogprob,
|
||||
};
|
||||
use vllm_tokenizer::{DynTokenizer, Tokenizer};
|
||||
|
||||
use super::super::ContentEvent;
|
||||
use super::reasoning_event_stream;
|
||||
use crate::event::AssistantBlockKind;
|
||||
use crate::parser::reasoning::{
|
||||
ReasoningDelta, ReasoningError, ReasoningParser, ReasoningParserFactory, names,
|
||||
};
|
||||
|
||||
struct FakeTokenizer;
|
||||
|
||||
impl Tokenizer for FakeTokenizer {
|
||||
fn encode(
|
||||
&self,
|
||||
text: &str,
|
||||
_add_special_tokens: bool,
|
||||
) -> vllm_tokenizer::Result<Vec<u32>> {
|
||||
Ok(text.chars().map(u32::from).collect())
|
||||
}
|
||||
|
||||
fn decode(
|
||||
&self,
|
||||
token_ids: &[u32],
|
||||
_skip_special_tokens: bool,
|
||||
) -> vllm_tokenizer::Result<String> {
|
||||
Ok(token_ids
|
||||
.iter()
|
||||
.map(|token_id| char::from_u32(*token_id).unwrap_or('\u{FFFD}'))
|
||||
.collect())
|
||||
}
|
||||
|
||||
fn token_to_id(&self, token: &str) -> Option<u32> {
|
||||
match token {
|
||||
"<think>" => Some(1),
|
||||
"</think>" => Some(2),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
struct FailingReasoningParser {
|
||||
fail_next: bool,
|
||||
}
|
||||
|
||||
impl ReasoningParser for FailingReasoningParser {
|
||||
fn create(_tokenizer: DynTokenizer) -> Result<Box<dyn ReasoningParser>, ReasoningError>
|
||||
where
|
||||
Self: Sized + 'static,
|
||||
{
|
||||
Ok(Box::new(Self { fail_next: true }))
|
||||
}
|
||||
|
||||
fn push(&mut self, _text: &str) -> Result<ReasoningDelta, ReasoningError> {
|
||||
if self.fail_next {
|
||||
self.fail_next = false;
|
||||
return Err(ReasoningError::MissingToken {
|
||||
token: "<think>".to_string(),
|
||||
});
|
||||
}
|
||||
Ok(ReasoningDelta::default())
|
||||
}
|
||||
}
|
||||
|
||||
fn test_reasoning_parser(factory: &mut ReasoningParserFactory) -> Box<dyn ReasoningParser> {
|
||||
factory.register_parser::<FailingReasoningParser>("failing");
|
||||
|
||||
factory.create("failing", Arc::new(FakeTokenizer)).unwrap()
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn reasoning_parser_failure_falls_back_to_plain_text() {
|
||||
let mut factory = ReasoningParserFactory::new();
|
||||
let events = stream::iter(vec![
|
||||
Ok(DecodedTextEvent::Start {
|
||||
prompt_token_ids: vec![1, 2, 3].into(),
|
||||
prompt_logprobs: None,
|
||||
}),
|
||||
Ok(DecodedTextEvent::TextDelta {
|
||||
delta: "abc".to_string(),
|
||||
token_ids: vec![],
|
||||
logprobs: None,
|
||||
finished: None,
|
||||
}),
|
||||
Ok(DecodedTextEvent::TextDelta {
|
||||
delta: "def".to_string(),
|
||||
token_ids: vec![],
|
||||
logprobs: None,
|
||||
finished: Some(vllm_text::Finished {
|
||||
usage: vllm_llm::TokenUsage {
|
||||
prompt_token_count: 3,
|
||||
output_token_count: 0,
|
||||
cached_token_count: 0,
|
||||
},
|
||||
finish_reason: FinishReason::stop_eos(),
|
||||
kv_transfer_params: None,
|
||||
}),
|
||||
}),
|
||||
]);
|
||||
|
||||
let collected = reasoning_event_stream(events, Some(test_reasoning_parser(&mut factory)))
|
||||
.collect::<Vec<_>>()
|
||||
.await;
|
||||
|
||||
let events = collected
|
||||
.into_iter()
|
||||
.collect::<crate::Result<Vec<_>>>()
|
||||
.expect("reasoning stream should not fail");
|
||||
|
||||
assert_eq!(
|
||||
events,
|
||||
vec![
|
||||
ContentEvent::Start {
|
||||
prompt_token_ids: vec![1, 2, 3].into(),
|
||||
prompt_logprobs: None,
|
||||
},
|
||||
ContentEvent::TextDelta {
|
||||
kind: AssistantBlockKind::Text,
|
||||
delta: "abc".to_string(),
|
||||
},
|
||||
ContentEvent::TextDelta {
|
||||
kind: AssistantBlockKind::Text,
|
||||
delta: "def".to_string(),
|
||||
},
|
||||
ContentEvent::Done {
|
||||
usage: vllm_llm::TokenUsage {
|
||||
prompt_token_count: 3,
|
||||
output_token_count: 0,
|
||||
cached_token_count: 0,
|
||||
},
|
||||
finish_reason: FinishReason::stop_eos(),
|
||||
kv_transfer_params: None,
|
||||
},
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn reasoning_stream_preserves_logprobs_delta() {
|
||||
let events = stream::iter(vec![
|
||||
Ok(DecodedTextEvent::Start {
|
||||
prompt_token_ids: vec![1].into(),
|
||||
prompt_logprobs: None,
|
||||
}),
|
||||
Ok(DecodedTextEvent::TextDelta {
|
||||
delta: "abc".to_string(),
|
||||
token_ids: vec![],
|
||||
logprobs: Some(DecodedLogprobs {
|
||||
positions: vec![DecodedPositionLogprobs {
|
||||
entries: vec![DecodedTokenLogprob {
|
||||
token_id: 0,
|
||||
token: "a".to_string(),
|
||||
logprob: -0.1,
|
||||
rank: 1,
|
||||
}],
|
||||
}],
|
||||
}),
|
||||
finished: None,
|
||||
}),
|
||||
]);
|
||||
|
||||
let collected = reasoning_event_stream(events, None)
|
||||
.collect::<Vec<_>>()
|
||||
.await
|
||||
.into_iter()
|
||||
.collect::<crate::Result<Vec<_>>>()
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
collected,
|
||||
vec![
|
||||
ContentEvent::Start {
|
||||
prompt_token_ids: vec![1].into(),
|
||||
prompt_logprobs: None,
|
||||
},
|
||||
ContentEvent::TextDelta {
|
||||
kind: AssistantBlockKind::Text,
|
||||
delta: "abc".to_string(),
|
||||
},
|
||||
ContentEvent::LogprobsDelta {
|
||||
logprobs: Some(DecodedLogprobs {
|
||||
positions: vec![DecodedPositionLogprobs {
|
||||
entries: vec![DecodedTokenLogprob {
|
||||
token_id: 0,
|
||||
token: "a".to_string(),
|
||||
logprob: -0.1,
|
||||
rank: 1,
|
||||
}],
|
||||
}],
|
||||
}),
|
||||
token_ids: vec![],
|
||||
},
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn qwen3_parser_uses_prompt_end_marker_to_switch_to_content() {
|
||||
let tokenizer = Arc::new(FakeTokenizer);
|
||||
let events = stream::iter(vec![
|
||||
Ok(DecodedTextEvent::Start {
|
||||
prompt_token_ids: vec![2].into(),
|
||||
prompt_logprobs: None,
|
||||
}),
|
||||
Ok(DecodedTextEvent::TextDelta {
|
||||
delta: "thought ".to_string(),
|
||||
token_ids: vec![],
|
||||
logprobs: None,
|
||||
finished: None,
|
||||
}),
|
||||
Ok(DecodedTextEvent::TextDelta {
|
||||
delta: "done</think>OK".to_string(),
|
||||
token_ids: vec![],
|
||||
logprobs: None,
|
||||
finished: None,
|
||||
}),
|
||||
]);
|
||||
|
||||
let factory = ReasoningParserFactory::new();
|
||||
let collected = reasoning_event_stream(
|
||||
events,
|
||||
Some(factory.create(names::QWEN3, tokenizer).unwrap()),
|
||||
)
|
||||
.collect::<Vec<_>>()
|
||||
.await;
|
||||
|
||||
let events = collected
|
||||
.into_iter()
|
||||
.collect::<crate::Result<Vec<_>>>()
|
||||
.expect("reasoning stream should not fail");
|
||||
|
||||
assert_eq!(
|
||||
events,
|
||||
vec![
|
||||
ContentEvent::Start {
|
||||
prompt_token_ids: vec![2].into(),
|
||||
prompt_logprobs: None,
|
||||
},
|
||||
ContentEvent::TextDelta {
|
||||
kind: AssistantBlockKind::Text,
|
||||
delta: "thought ".to_string(),
|
||||
},
|
||||
ContentEvent::TextDelta {
|
||||
kind: AssistantBlockKind::Text,
|
||||
delta: "done</think>OK".to_string(),
|
||||
},
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn qwen3_parser_tolerates_prompt_prefill_reasoning() {
|
||||
let tokenizer = Arc::new(FakeTokenizer);
|
||||
let events = stream::iter(vec![
|
||||
Ok(DecodedTextEvent::Start {
|
||||
prompt_token_ids: vec![1].into(),
|
||||
prompt_logprobs: None,
|
||||
}),
|
||||
Ok(DecodedTextEvent::TextDelta {
|
||||
delta: "thought ".to_string(),
|
||||
token_ids: vec![],
|
||||
logprobs: None,
|
||||
finished: None,
|
||||
}),
|
||||
Ok(DecodedTextEvent::TextDelta {
|
||||
delta: "done</think>OK".to_string(),
|
||||
token_ids: vec![],
|
||||
logprobs: None,
|
||||
finished: None,
|
||||
}),
|
||||
]);
|
||||
|
||||
let factory = ReasoningParserFactory::new();
|
||||
let collected = reasoning_event_stream(
|
||||
events,
|
||||
Some(factory.create(names::QWEN3, tokenizer).unwrap()),
|
||||
)
|
||||
.collect::<Vec<_>>()
|
||||
.await;
|
||||
|
||||
let events = collected
|
||||
.into_iter()
|
||||
.collect::<crate::Result<Vec<_>>>()
|
||||
.expect("reasoning stream should not fail");
|
||||
|
||||
assert_eq!(
|
||||
events,
|
||||
vec![
|
||||
ContentEvent::Start {
|
||||
prompt_token_ids: vec![1].into(),
|
||||
prompt_logprobs: None,
|
||||
},
|
||||
ContentEvent::TextDelta {
|
||||
kind: AssistantBlockKind::Reasoning,
|
||||
delta: "thought ".to_string(),
|
||||
},
|
||||
ContentEvent::TextDelta {
|
||||
kind: AssistantBlockKind::Reasoning,
|
||||
delta: "done".to_string(),
|
||||
},
|
||||
ContentEvent::TextDelta {
|
||||
kind: AssistantBlockKind::Text,
|
||||
delta: "OK".to_string(),
|
||||
},
|
||||
]
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -77,7 +77,7 @@ fn structural_tag_tool_choice(request: &ChatRequest) -> Option<StructuralTagTool
|
||||
mod tests {
|
||||
use serde_json::{Value, json};
|
||||
use vllm_engine_core_client::protocol::{StructuredOutputBackend, StructuredOutputsParams};
|
||||
use vllm_tool_parser::{Qwen3CoderToolParser, Tool};
|
||||
use vllm_parser::tool::{Qwen3CoderToolParser, Tool};
|
||||
|
||||
use super::*;
|
||||
|
||||
|
||||
@@ -1,985 +0,0 @@
|
||||
//! Adapts plain assistant text deltas into tool-call-aware assistant updates.
|
||||
//!
|
||||
//! This stage runs after reasoning separation and before final block assembly.
|
||||
//! It only inspects normal assistant text, leaves reasoning deltas untouched,
|
||||
//! and translates incremental tool parsing output into internal tool-call
|
||||
//! events while preserving plain-text fallback behavior.
|
||||
|
||||
use asynk_strim_attr::{TryYielder, try_stream};
|
||||
use futures::{StreamExt as _, pin_mut};
|
||||
use thiserror_ext::AsReport;
|
||||
use tracing::warn;
|
||||
|
||||
use super::{AssistantEvent, ContentEvent, ContentEventStream};
|
||||
use crate::Result;
|
||||
use crate::error::Error;
|
||||
use crate::event::AssistantBlockKind;
|
||||
use crate::output::generate_tool_call_id;
|
||||
use crate::parser::tool::{ToolCallDelta, ToolParser, ToolParserOutput};
|
||||
|
||||
/// Per-stream tool parsing state.
|
||||
struct ToolState {
|
||||
/// Parser for the current model family.
|
||||
parser: Box<dyn ToolParser>,
|
||||
/// Whether tool parsing has already failed for this stream.
|
||||
parser_failed: bool,
|
||||
/// The parser-local index of the currently open tool call, if any.
|
||||
// NOTE: We only allow single open tool call at a time right now, since that's what all
|
||||
// supported parsers currently emit. Change this to a `BTreeMap` if we need to support multiple
|
||||
// interleaved calls in the future.
|
||||
open_call_index: Option<usize>,
|
||||
}
|
||||
|
||||
impl ToolState {
|
||||
/// Create one fresh tool-parsing state for a new streamed response.
|
||||
fn new(parser: Box<dyn ToolParser>) -> Self {
|
||||
Self {
|
||||
parser,
|
||||
parser_failed: false,
|
||||
open_call_index: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Convert one semantic assistant text delta into zero or more tool-aware
|
||||
/// internal events.
|
||||
fn process_text_delta(
|
||||
&mut self,
|
||||
kind: AssistantBlockKind,
|
||||
delta: String,
|
||||
) -> Result<Vec<AssistantEvent>> {
|
||||
let mut events = Vec::new();
|
||||
|
||||
// Only normal assistant text is eligible for tool parsing. Reasoning
|
||||
// blocks and plain-text fallback should pass through unchanged.
|
||||
if kind != AssistantBlockKind::Text || self.parser_failed {
|
||||
self.open_call_index = None;
|
||||
events.push(AssistantEvent::TextDelta { kind, delta });
|
||||
return Ok(events);
|
||||
}
|
||||
|
||||
let mut output = ToolParserOutput::default();
|
||||
let parse_result = self.parser.parse_into(&delta, &mut output);
|
||||
|
||||
match parse_result {
|
||||
Ok(()) => self.process_parser_output(kind, output, &mut events)?,
|
||||
Err(error) => {
|
||||
warn!(
|
||||
error = %error.as_report(),
|
||||
"tool parser failed; falling back to plain text deltas"
|
||||
);
|
||||
// Permanently mark this parser as failed.
|
||||
// TODO: we may consider recovering from parsing errors in the future.
|
||||
self.parser_failed = true;
|
||||
|
||||
// On parsing failure, we still apply the partial parser output if any, but we close
|
||||
// any open tool calls and emit the remaining buffered text as a plain-text delta to
|
||||
// preserve as much of the output as possible.
|
||||
self.process_parser_output(kind, output, &mut events)?;
|
||||
self.open_call_index = None;
|
||||
push_text_delta(&mut events, kind, self.parser.reset());
|
||||
}
|
||||
}
|
||||
|
||||
Ok(events)
|
||||
}
|
||||
|
||||
/// Apply one parsed tool output to the current stream state.
|
||||
fn process_parser_output(
|
||||
&mut self,
|
||||
kind: AssistantBlockKind,
|
||||
output: ToolParserOutput,
|
||||
events: &mut Vec<AssistantEvent>,
|
||||
) -> Result<()> {
|
||||
// When we are not currently streaming a tool call, preserve plain
|
||||
// text first and then surface any new tool call items.
|
||||
if self.open_call_index.is_none() {
|
||||
push_text_delta(events, kind, output.normal_text);
|
||||
self.process_tool_items(output.calls, events)?;
|
||||
} else {
|
||||
// Once a tool call is open, prioritize tool deltas first. If the
|
||||
// parser emits normal text again, close the tool call and resume
|
||||
// plain text output.
|
||||
self.process_tool_items(output.calls, events)?;
|
||||
if !output.normal_text.is_empty() {
|
||||
self.open_call_index = None;
|
||||
push_text_delta(events, kind, output.normal_text);
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Apply one batch of parsed tool-call deltas emitted by the parser.
|
||||
fn process_tool_items(
|
||||
&mut self,
|
||||
items: Vec<ToolCallDelta>,
|
||||
events: &mut Vec<AssistantEvent>,
|
||||
) -> Result<()> {
|
||||
for item in items {
|
||||
if let Some(name) = item.name {
|
||||
let is_new_tool = match self.open_call_index {
|
||||
Some(open_call_index) => open_call_index != item.tool_index,
|
||||
None => true,
|
||||
};
|
||||
if is_new_tool {
|
||||
let id = self
|
||||
.parser
|
||||
.tool_call_id(item.tool_index)
|
||||
.map(str::to_string)
|
||||
.unwrap_or_else(generate_tool_call_id);
|
||||
self.open_call_index = Some(item.tool_index);
|
||||
events.push(AssistantEvent::ToolCallStart { id, name });
|
||||
}
|
||||
}
|
||||
|
||||
if item.arguments.is_empty() {
|
||||
// No arguments delta to apply.
|
||||
continue;
|
||||
}
|
||||
let Some(open_call_index) = self.open_call_index else {
|
||||
return Err(Error::ToolCallStreamInvariant {
|
||||
message: format!(
|
||||
"received arguments for tool index {} before any tool-call start",
|
||||
item.tool_index
|
||||
),
|
||||
});
|
||||
};
|
||||
if open_call_index != item.tool_index {
|
||||
return Err(Error::ToolCallStreamInvariant {
|
||||
message: format!(
|
||||
"received arguments for tool index {} while tool index {} is open",
|
||||
item.tool_index, open_call_index
|
||||
),
|
||||
});
|
||||
}
|
||||
|
||||
events.push(AssistantEvent::ToolCallArgumentsDelta {
|
||||
delta: item.arguments,
|
||||
});
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Flush parser state at end-of-stream and close any remaining open calls.
|
||||
fn finish(&mut self) -> Result<Vec<AssistantEvent>> {
|
||||
let mut events = Vec::new();
|
||||
|
||||
if self.parser_failed {
|
||||
return Ok(events);
|
||||
}
|
||||
|
||||
match self.parser.finish() {
|
||||
Ok(output) => {
|
||||
self.process_parser_output(AssistantBlockKind::Text, output, &mut events)?
|
||||
}
|
||||
Err(error) => {
|
||||
warn!(
|
||||
error = %error.as_report(),
|
||||
"tool parser finish failed; closing open tool calls with buffered state"
|
||||
);
|
||||
self.parser_failed = true;
|
||||
}
|
||||
}
|
||||
|
||||
Ok(events)
|
||||
}
|
||||
}
|
||||
|
||||
/// Push one plain-text delta if it is non-empty.
|
||||
fn push_text_delta(events: &mut Vec<AssistantEvent>, kind: AssistantBlockKind, delta: String) {
|
||||
if delta.is_empty() {
|
||||
return;
|
||||
}
|
||||
events.push(AssistantEvent::TextDelta { kind, delta });
|
||||
}
|
||||
|
||||
/// Wrap one semantic assistant stream into the internal tool-aware assistant
|
||||
/// stream.
|
||||
#[try_stream]
|
||||
pub(crate) async fn tool_event_stream(
|
||||
stream: impl ContentEventStream,
|
||||
parser: Option<Box<dyn ToolParser>>,
|
||||
mut y: TryYielder<AssistantEvent, Error>,
|
||||
) -> Result<()> {
|
||||
// Without a parser, pass through the input stream unchanged.
|
||||
let Some(parser) = parser else {
|
||||
pin_mut!(stream);
|
||||
while let Some(event) = stream.next().await.transpose()? {
|
||||
y.yield_ok(event.into()).await;
|
||||
}
|
||||
return Ok(());
|
||||
};
|
||||
|
||||
pin_mut!(stream);
|
||||
let mut state = ToolState::new(parser);
|
||||
|
||||
while let Some(event) = stream.next().await.transpose()? {
|
||||
match event {
|
||||
ContentEvent::Start {
|
||||
prompt_token_ids,
|
||||
prompt_logprobs,
|
||||
} => {
|
||||
y.yield_ok(AssistantEvent::Start {
|
||||
prompt_token_ids,
|
||||
prompt_logprobs,
|
||||
})
|
||||
.await;
|
||||
}
|
||||
ContentEvent::TextDelta { kind, delta } => {
|
||||
for next in state.process_text_delta(kind, delta)? {
|
||||
y.yield_ok(next).await;
|
||||
}
|
||||
}
|
||||
ContentEvent::LogprobsDelta {
|
||||
logprobs,
|
||||
token_ids,
|
||||
} => {
|
||||
y.yield_ok(AssistantEvent::LogprobsDelta {
|
||||
logprobs,
|
||||
token_ids,
|
||||
})
|
||||
.await;
|
||||
}
|
||||
ContentEvent::Done {
|
||||
usage,
|
||||
finish_reason,
|
||||
kv_transfer_params,
|
||||
} => {
|
||||
for next in state.finish()? {
|
||||
y.yield_ok(next).await;
|
||||
}
|
||||
|
||||
y.yield_ok(AssistantEvent::Done {
|
||||
usage,
|
||||
finish_reason,
|
||||
kv_transfer_params,
|
||||
})
|
||||
.await;
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
|
||||
use futures::{StreamExt as _, stream};
|
||||
use vllm_llm::FinishReason;
|
||||
use vllm_text::{DecodedLogprobs, DecodedPositionLogprobs, DecodedTokenLogprob};
|
||||
use vllm_tool_parser::Result;
|
||||
|
||||
use super::super::{AssistantEvent, ContentEvent};
|
||||
use super::tool_event_stream;
|
||||
use crate::error::Error;
|
||||
use crate::event::{AssistantBlockKind, AssistantMessageExt as _};
|
||||
use crate::output::structured::structured_chat_event_stream;
|
||||
use crate::parser::tool::{
|
||||
DeepSeekV4ToolParser, ToolParser, ToolParserError, ToolParserOutput,
|
||||
};
|
||||
use crate::request::ChatTool;
|
||||
use crate::stream::{ChatEventStream, CollectedAssistantMessage};
|
||||
|
||||
struct FailingParser {
|
||||
fail_next: bool,
|
||||
buffered: String,
|
||||
}
|
||||
|
||||
struct ScriptedParser {
|
||||
push_outputs: Vec<ToolParserOutput>,
|
||||
finish_output: ToolParserOutput,
|
||||
}
|
||||
|
||||
struct PartialThenFailParser {
|
||||
buffered: String,
|
||||
}
|
||||
|
||||
struct IdScriptedParser {
|
||||
output: ToolParserOutput,
|
||||
tool_call_id: Option<String>,
|
||||
}
|
||||
|
||||
impl ToolParser for FailingParser {
|
||||
fn create(_tools: &[ChatTool]) -> vllm_tool_parser::Result<Box<dyn ToolParser>>
|
||||
where
|
||||
Self: Sized + 'static,
|
||||
{
|
||||
Ok(Box::new(Self {
|
||||
fail_next: false,
|
||||
buffered: String::new(),
|
||||
}))
|
||||
}
|
||||
|
||||
fn parse_into(&mut self, chunk: &str, _output: &mut ToolParserOutput) -> Result<()> {
|
||||
self.buffered.push_str(chunk);
|
||||
if self.fail_next {
|
||||
self.fail_next = false;
|
||||
return Err(ToolParserError::ParsingFailed {
|
||||
message: "boom".to_string(),
|
||||
});
|
||||
}
|
||||
|
||||
self.buffered.clear();
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn finish(&mut self) -> Result<ToolParserOutput> {
|
||||
Ok(ToolParserOutput::default())
|
||||
}
|
||||
|
||||
fn reset(&mut self) -> String {
|
||||
std::mem::take(&mut self.buffered)
|
||||
}
|
||||
}
|
||||
|
||||
impl ToolParser for ScriptedParser {
|
||||
fn create(_tools: &[ChatTool]) -> vllm_tool_parser::Result<Box<dyn ToolParser>>
|
||||
where
|
||||
Self: Sized + 'static,
|
||||
{
|
||||
Ok(Box::new(Self {
|
||||
push_outputs: Vec::new(),
|
||||
finish_output: ToolParserOutput::default(),
|
||||
}))
|
||||
}
|
||||
|
||||
fn parse_into(&mut self, _chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
|
||||
let mut next = self.push_outputs.pop().unwrap_or_default();
|
||||
output.normal_text.push_str(&next.normal_text);
|
||||
output.calls.append(&mut next.calls);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn finish(&mut self) -> Result<ToolParserOutput> {
|
||||
Ok(std::mem::take(&mut self.finish_output))
|
||||
}
|
||||
|
||||
fn reset(&mut self) -> String {
|
||||
String::new()
|
||||
}
|
||||
}
|
||||
|
||||
impl ToolParser for IdScriptedParser {
|
||||
fn create(_tools: &[ChatTool]) -> vllm_tool_parser::Result<Box<dyn ToolParser>>
|
||||
where
|
||||
Self: Sized + 'static,
|
||||
{
|
||||
Ok(Box::new(Self {
|
||||
output: ToolParserOutput::default(),
|
||||
tool_call_id: None,
|
||||
}))
|
||||
}
|
||||
|
||||
fn tool_call_id(&self, tool_index: usize) -> Option<&str> {
|
||||
(tool_index == 0).then_some(self.tool_call_id.as_deref()).flatten()
|
||||
}
|
||||
|
||||
fn parse_into(&mut self, _chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
|
||||
output.append(std::mem::take(&mut self.output));
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn finish(&mut self) -> Result<ToolParserOutput> {
|
||||
Ok(ToolParserOutput::default())
|
||||
}
|
||||
|
||||
fn reset(&mut self) -> String {
|
||||
String::new()
|
||||
}
|
||||
}
|
||||
|
||||
impl ToolParser for PartialThenFailParser {
|
||||
fn create(_tools: &[ChatTool]) -> vllm_tool_parser::Result<Box<dyn ToolParser>>
|
||||
where
|
||||
Self: Sized + 'static,
|
||||
{
|
||||
Ok(Box::new(Self {
|
||||
buffered: String::new(),
|
||||
}))
|
||||
}
|
||||
|
||||
fn parse_into(&mut self, _chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
|
||||
output.calls.extend([
|
||||
crate::parser::tool::ToolCallDelta {
|
||||
tool_index: 0,
|
||||
name: Some("get_weather".to_string()),
|
||||
arguments: String::new(),
|
||||
},
|
||||
crate::parser::tool::ToolCallDelta {
|
||||
tool_index: 0,
|
||||
name: None,
|
||||
arguments: r#"{"location":"SF"}"#.to_string(),
|
||||
},
|
||||
]);
|
||||
self.buffered.push_str(" trailing text");
|
||||
Err(ToolParserError::ParsingFailed {
|
||||
message: "boom".to_string(),
|
||||
})
|
||||
}
|
||||
|
||||
fn finish(&mut self) -> Result<ToolParserOutput> {
|
||||
Ok(ToolParserOutput::default())
|
||||
}
|
||||
|
||||
fn reset(&mut self) -> String {
|
||||
std::mem::take(&mut self.buffered)
|
||||
}
|
||||
}
|
||||
|
||||
fn deepseek_v4_test_tools() -> Vec<ChatTool> {
|
||||
vec![
|
||||
ChatTool {
|
||||
name: "get_weather".to_string(),
|
||||
description: None,
|
||||
parameters: serde_json::json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"location": { "type": "string" }
|
||||
}
|
||||
}),
|
||||
strict: None,
|
||||
},
|
||||
ChatTool {
|
||||
name: "add".to_string(),
|
||||
description: None,
|
||||
parameters: serde_json::json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"x": { "type": "integer" },
|
||||
"y": { "type": "integer" }
|
||||
}
|
||||
}),
|
||||
strict: None,
|
||||
},
|
||||
]
|
||||
}
|
||||
|
||||
async fn collect_deepseek_v4_message(chunks: Vec<String>) -> CollectedAssistantMessage {
|
||||
let events = chunks
|
||||
.into_iter()
|
||||
.map(|delta| {
|
||||
Ok(ContentEvent::TextDelta {
|
||||
kind: AssistantBlockKind::Text,
|
||||
delta,
|
||||
})
|
||||
})
|
||||
.chain(std::iter::once(Ok(ContentEvent::Done {
|
||||
usage: vllm_llm::TokenUsage {
|
||||
prompt_token_count: 1,
|
||||
output_token_count: 1,
|
||||
cached_token_count: 0,
|
||||
},
|
||||
finish_reason: FinishReason::stop_eos(),
|
||||
kv_transfer_params: None,
|
||||
})));
|
||||
let parser = DeepSeekV4ToolParser::create(&deepseek_v4_test_tools()).unwrap();
|
||||
let assistant_events = tool_event_stream(stream::iter(events), Some(parser));
|
||||
let chat_events = structured_chat_event_stream(assistant_events, true);
|
||||
|
||||
ChatEventStream::new("req_deepseek_v4".to_string(), Box::pin(chat_events))
|
||||
.collect_message()
|
||||
.await
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
fn message_tool_projection(
|
||||
message: &CollectedAssistantMessage,
|
||||
) -> (String, Vec<(String, serde_json::Value)>) {
|
||||
(
|
||||
message.message.text(),
|
||||
message
|
||||
.message
|
||||
.tool_calls()
|
||||
.map(|call| {
|
||||
(
|
||||
call.name.clone(),
|
||||
serde_json::from_str(&call.arguments).unwrap(),
|
||||
)
|
||||
})
|
||||
.collect(),
|
||||
)
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn tool_parser_error_preserves_partial_output_and_flushes_buffer() {
|
||||
let events = stream::iter(vec![
|
||||
Ok(ContentEvent::TextDelta {
|
||||
kind: AssistantBlockKind::Text,
|
||||
delta: "ignored".to_string(),
|
||||
}),
|
||||
Ok(ContentEvent::Done {
|
||||
usage: vllm_llm::TokenUsage {
|
||||
prompt_token_count: 1,
|
||||
output_token_count: 1,
|
||||
cached_token_count: 0,
|
||||
},
|
||||
finish_reason: FinishReason::stop_eos(),
|
||||
kv_transfer_params: None,
|
||||
}),
|
||||
]);
|
||||
|
||||
let events = tool_event_stream(
|
||||
events,
|
||||
Some(Box::new(PartialThenFailParser {
|
||||
buffered: String::new(),
|
||||
})),
|
||||
)
|
||||
.collect::<Vec<_>>()
|
||||
.await
|
||||
.into_iter()
|
||||
.collect::<crate::Result<Vec<_>>>()
|
||||
.unwrap();
|
||||
|
||||
assert!(matches!(
|
||||
&events[0],
|
||||
AssistantEvent::ToolCallStart { name, .. } if name == "get_weather"
|
||||
));
|
||||
assert!(matches!(
|
||||
&events[1],
|
||||
AssistantEvent::ToolCallArgumentsDelta { delta } if delta == r#"{"location":"SF"}"#
|
||||
));
|
||||
assert_eq!(
|
||||
events[2],
|
||||
AssistantEvent::TextDelta {
|
||||
kind: AssistantBlockKind::Text,
|
||||
delta: " trailing text".to_string(),
|
||||
}
|
||||
);
|
||||
assert!(matches!(events[3], AssistantEvent::Done { .. }));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn tool_stream_preserves_parser_provided_tool_call_id() {
|
||||
let events = stream::iter(vec![Ok(ContentEvent::TextDelta {
|
||||
kind: AssistantBlockKind::Text,
|
||||
delta: "ignored".to_string(),
|
||||
})]);
|
||||
let parser = IdScriptedParser {
|
||||
output: ToolParserOutput {
|
||||
normal_text: String::new(),
|
||||
calls: vec![crate::parser::tool::ToolCallDelta {
|
||||
tool_index: 0,
|
||||
name: Some("get_weather".to_string()),
|
||||
arguments: "{}".to_string(),
|
||||
}],
|
||||
},
|
||||
tool_call_id: Some("functions.get_weather:0".to_string()),
|
||||
};
|
||||
|
||||
let events = tool_event_stream(events, Some(Box::new(parser)))
|
||||
.collect::<Vec<_>>()
|
||||
.await
|
||||
.into_iter()
|
||||
.collect::<crate::Result<Vec<_>>>()
|
||||
.unwrap();
|
||||
|
||||
assert!(matches!(
|
||||
&events[0],
|
||||
AssistantEvent::ToolCallStart { id, name }
|
||||
if id == "functions.get_weather:0" && name == "get_weather"
|
||||
));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn tool_stream_generates_tool_call_id_when_parser_omits_one() {
|
||||
let events = stream::iter(vec![Ok(ContentEvent::TextDelta {
|
||||
kind: AssistantBlockKind::Text,
|
||||
delta: "ignored".to_string(),
|
||||
})]);
|
||||
let parser = IdScriptedParser {
|
||||
output: ToolParserOutput {
|
||||
normal_text: String::new(),
|
||||
calls: vec![crate::parser::tool::ToolCallDelta {
|
||||
tool_index: 0,
|
||||
name: Some("get_weather".to_string()),
|
||||
arguments: "{}".to_string(),
|
||||
}],
|
||||
},
|
||||
tool_call_id: None,
|
||||
};
|
||||
|
||||
let events = tool_event_stream(events, Some(Box::new(parser)))
|
||||
.collect::<Vec<_>>()
|
||||
.await
|
||||
.into_iter()
|
||||
.collect::<crate::Result<Vec<_>>>()
|
||||
.unwrap();
|
||||
|
||||
assert!(matches!(
|
||||
&events[0],
|
||||
AssistantEvent::ToolCallStart { id, name }
|
||||
if id.starts_with("call_") && name == "get_weather"
|
||||
));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn real_buffered_parser_error_matches_streaming_and_non_streaming() {
|
||||
let prefix = "I will check both.\n";
|
||||
let first_tool_call = concat!(
|
||||
"<|DSML|tool_calls>\n",
|
||||
"<|DSML|invoke name=\"get_weather\">\n",
|
||||
"<|DSML|parameter name=\"location\" string=\"true\">Tokyo</|DSML|parameter>\n",
|
||||
"</|DSML|invoke>",
|
||||
);
|
||||
let malformed_second_tool_call = concat!(
|
||||
"\n<|DSML|invoke name=\"add\">\n",
|
||||
"not a parameter\n",
|
||||
"</|DSML|invoke>\n",
|
||||
"</|DSML|tool_calls>",
|
||||
);
|
||||
let streaming_chunks = vec![
|
||||
prefix.to_string(),
|
||||
first_tool_call.to_string(),
|
||||
malformed_second_tool_call.to_string(),
|
||||
];
|
||||
let full_output = streaming_chunks.concat();
|
||||
|
||||
let streaming = collect_deepseek_v4_message(streaming_chunks).await;
|
||||
let non_streaming = collect_deepseek_v4_message(vec![full_output]).await;
|
||||
|
||||
let expected = (
|
||||
format!("{prefix}{malformed_second_tool_call}"),
|
||||
vec![(
|
||||
"get_weather".to_string(),
|
||||
serde_json::json!({ "location": "Tokyo" }),
|
||||
)],
|
||||
);
|
||||
assert_eq!(message_tool_projection(&streaming), expected);
|
||||
assert_eq!(message_tool_projection(&non_streaming), expected);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn tool_parser_failure_falls_back_to_plain_text() {
|
||||
let events = stream::iter(vec![
|
||||
Ok(ContentEvent::Start {
|
||||
prompt_token_ids: vec![1, 2, 3].into(),
|
||||
prompt_logprobs: None,
|
||||
}),
|
||||
Ok(ContentEvent::TextDelta {
|
||||
kind: AssistantBlockKind::Text,
|
||||
delta: "abc".to_string(),
|
||||
}),
|
||||
Ok(ContentEvent::TextDelta {
|
||||
kind: AssistantBlockKind::Text,
|
||||
delta: "def".to_string(),
|
||||
}),
|
||||
Ok(ContentEvent::Done {
|
||||
usage: vllm_llm::TokenUsage {
|
||||
prompt_token_count: 3,
|
||||
output_token_count: 0,
|
||||
cached_token_count: 0,
|
||||
},
|
||||
finish_reason: FinishReason::stop_eos(),
|
||||
kv_transfer_params: None,
|
||||
}),
|
||||
]);
|
||||
|
||||
let collected = tool_event_stream(
|
||||
events,
|
||||
Some(Box::new(FailingParser {
|
||||
fail_next: true,
|
||||
buffered: String::new(),
|
||||
})),
|
||||
)
|
||||
.collect::<Vec<_>>()
|
||||
.await;
|
||||
|
||||
let events = collected
|
||||
.into_iter()
|
||||
.collect::<crate::Result<Vec<_>>>()
|
||||
.expect("tool stream should not fail");
|
||||
|
||||
assert_eq!(
|
||||
events,
|
||||
vec![
|
||||
AssistantEvent::Start {
|
||||
prompt_token_ids: vec![1, 2, 3].into(),
|
||||
prompt_logprobs: None,
|
||||
},
|
||||
AssistantEvent::TextDelta {
|
||||
kind: AssistantBlockKind::Text,
|
||||
delta: "abc".to_string(),
|
||||
},
|
||||
AssistantEvent::TextDelta {
|
||||
kind: AssistantBlockKind::Text,
|
||||
delta: "def".to_string(),
|
||||
},
|
||||
AssistantEvent::Done {
|
||||
usage: vllm_llm::TokenUsage {
|
||||
prompt_token_count: 3,
|
||||
output_token_count: 0,
|
||||
cached_token_count: 0,
|
||||
},
|
||||
finish_reason: FinishReason::stop_eos(),
|
||||
kv_transfer_params: None,
|
||||
},
|
||||
]
|
||||
);
|
||||
|
||||
let message = ChatEventStream::new(
|
||||
"req_fallback".to_string(),
|
||||
Box::pin(structured_chat_event_stream(
|
||||
stream::iter(events.into_iter().map(Ok)),
|
||||
true,
|
||||
)),
|
||||
)
|
||||
.collect_message()
|
||||
.await
|
||||
.expect("collect_message should succeed");
|
||||
assert_eq!(message.message.text(), "abcdef");
|
||||
assert!(message.message.tool_calls().next().is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn tool_stream_preserves_logprobs_delta() {
|
||||
let events = stream::iter(vec![
|
||||
Ok(ContentEvent::Start {
|
||||
prompt_token_ids: vec![1].into(),
|
||||
prompt_logprobs: None,
|
||||
}),
|
||||
Ok(ContentEvent::LogprobsDelta {
|
||||
logprobs: Some(DecodedLogprobs {
|
||||
positions: vec![DecodedPositionLogprobs {
|
||||
entries: vec![DecodedTokenLogprob {
|
||||
token_id: 0,
|
||||
token: "a".to_string(),
|
||||
logprob: -0.2,
|
||||
rank: 1,
|
||||
}],
|
||||
}],
|
||||
}),
|
||||
token_ids: vec![],
|
||||
}),
|
||||
Ok(ContentEvent::Done {
|
||||
usage: vllm_llm::TokenUsage {
|
||||
prompt_token_count: 1,
|
||||
output_token_count: 0,
|
||||
cached_token_count: 0,
|
||||
},
|
||||
finish_reason: FinishReason::stop_eos(),
|
||||
kv_transfer_params: None,
|
||||
}),
|
||||
]);
|
||||
let events = tool_event_stream(
|
||||
events,
|
||||
Some(Box::new(FailingParser {
|
||||
fail_next: false,
|
||||
buffered: String::new(),
|
||||
})),
|
||||
)
|
||||
.collect::<Vec<_>>()
|
||||
.await
|
||||
.into_iter()
|
||||
.collect::<crate::Result<Vec<_>>>()
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
events,
|
||||
vec![
|
||||
AssistantEvent::Start {
|
||||
prompt_token_ids: vec![1].into(),
|
||||
prompt_logprobs: None,
|
||||
},
|
||||
AssistantEvent::LogprobsDelta {
|
||||
logprobs: Some(DecodedLogprobs {
|
||||
positions: vec![DecodedPositionLogprobs {
|
||||
entries: vec![DecodedTokenLogprob {
|
||||
token_id: 0,
|
||||
token: "a".to_string(),
|
||||
logprob: -0.2,
|
||||
rank: 1,
|
||||
}],
|
||||
}],
|
||||
}),
|
||||
token_ids: vec![],
|
||||
},
|
||||
AssistantEvent::Done {
|
||||
usage: vllm_llm::TokenUsage {
|
||||
prompt_token_count: 1,
|
||||
output_token_count: 0,
|
||||
cached_token_count: 0,
|
||||
},
|
||||
finish_reason: FinishReason::stop_eos(),
|
||||
kv_transfer_params: None,
|
||||
},
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn tool_stream_rejects_interleaved_tool_indices() {
|
||||
let events = stream::iter(vec![
|
||||
Ok(ContentEvent::TextDelta {
|
||||
kind: AssistantBlockKind::Text,
|
||||
delta: "ignored".to_string(),
|
||||
}),
|
||||
Ok(ContentEvent::Done {
|
||||
usage: vllm_llm::TokenUsage {
|
||||
prompt_token_count: 1,
|
||||
output_token_count: 1,
|
||||
cached_token_count: 0,
|
||||
},
|
||||
finish_reason: FinishReason::stop_eos(),
|
||||
kv_transfer_params: None,
|
||||
}),
|
||||
]);
|
||||
|
||||
let parser = ScriptedParser {
|
||||
push_outputs: vec![ToolParserOutput {
|
||||
normal_text: String::new(),
|
||||
calls: vec![
|
||||
crate::parser::tool::ToolCallDelta {
|
||||
tool_index: 0,
|
||||
name: Some("first".to_string()),
|
||||
arguments: String::new(),
|
||||
},
|
||||
crate::parser::tool::ToolCallDelta {
|
||||
tool_index: 1,
|
||||
name: None,
|
||||
arguments: "{}".to_string(),
|
||||
},
|
||||
],
|
||||
}],
|
||||
finish_output: ToolParserOutput::default(),
|
||||
};
|
||||
|
||||
let err = tool_event_stream(events, Some(Box::new(parser)))
|
||||
.collect::<Vec<_>>()
|
||||
.await
|
||||
.into_iter()
|
||||
.find_map(|output| output.err())
|
||||
.expect("expected invariant error");
|
||||
|
||||
assert!(matches!(err, Error::ToolCallStreamInvariant { .. }));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn tool_stream_resets_open_tool_when_normal_text_interrupts_it() {
|
||||
let events = stream::iter(vec![
|
||||
Ok(ContentEvent::TextDelta {
|
||||
kind: AssistantBlockKind::Text,
|
||||
delta: "start".to_string(),
|
||||
}),
|
||||
Ok(ContentEvent::TextDelta {
|
||||
kind: AssistantBlockKind::Text,
|
||||
delta: "text".to_string(),
|
||||
}),
|
||||
Ok(ContentEvent::TextDelta {
|
||||
kind: AssistantBlockKind::Text,
|
||||
delta: "args".to_string(),
|
||||
}),
|
||||
]);
|
||||
|
||||
let parser = ScriptedParser {
|
||||
push_outputs: vec![
|
||||
ToolParserOutput {
|
||||
normal_text: String::new(),
|
||||
calls: vec![crate::parser::tool::ToolCallDelta {
|
||||
tool_index: 0,
|
||||
name: None,
|
||||
arguments: "}".to_string(),
|
||||
}],
|
||||
},
|
||||
ToolParserOutput {
|
||||
normal_text: "plain text".to_string(),
|
||||
calls: Vec::new(),
|
||||
},
|
||||
ToolParserOutput {
|
||||
normal_text: String::new(),
|
||||
calls: vec![crate::parser::tool::ToolCallDelta {
|
||||
tool_index: 0,
|
||||
name: Some("first".to_string()),
|
||||
arguments: "{".to_string(),
|
||||
}],
|
||||
},
|
||||
],
|
||||
finish_output: ToolParserOutput::default(),
|
||||
};
|
||||
|
||||
let err = tool_event_stream(events, Some(Box::new(parser)))
|
||||
.collect::<Vec<_>>()
|
||||
.await
|
||||
.into_iter()
|
||||
.find_map(|output| output.err())
|
||||
.expect("expected invariant error");
|
||||
|
||||
assert!(matches!(
|
||||
err,
|
||||
Error::ToolCallStreamInvariant { message }
|
||||
if message == "received arguments for tool index 0 before any tool-call start"
|
||||
));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn tool_stream_emits_start_and_args_for_terminal_text() {
|
||||
let events = stream::iter(vec![
|
||||
Ok(ContentEvent::Start {
|
||||
prompt_token_ids: vec![1].into(),
|
||||
prompt_logprobs: None,
|
||||
}),
|
||||
Ok(ContentEvent::TextDelta {
|
||||
kind: AssistantBlockKind::Text,
|
||||
delta: "ignored".to_string(),
|
||||
}),
|
||||
Ok(ContentEvent::Done {
|
||||
usage: vllm_llm::TokenUsage {
|
||||
prompt_token_count: 1,
|
||||
output_token_count: 1,
|
||||
cached_token_count: 0,
|
||||
},
|
||||
finish_reason: FinishReason::stop_eos(),
|
||||
kv_transfer_params: None,
|
||||
}),
|
||||
]);
|
||||
|
||||
let parser = ScriptedParser {
|
||||
push_outputs: vec![ToolParserOutput {
|
||||
normal_text: String::new(),
|
||||
calls: vec![
|
||||
crate::parser::tool::ToolCallDelta {
|
||||
tool_index: 0,
|
||||
name: Some("first".to_string()),
|
||||
arguments: r#"{"a":1}"#.to_string(),
|
||||
},
|
||||
crate::parser::tool::ToolCallDelta {
|
||||
tool_index: 1,
|
||||
name: Some("second".to_string()),
|
||||
arguments: r#"{"b":2}"#.to_string(),
|
||||
},
|
||||
],
|
||||
}],
|
||||
finish_output: ToolParserOutput::default(),
|
||||
};
|
||||
|
||||
let events = tool_event_stream(events, Some(Box::new(parser)))
|
||||
.collect::<Vec<_>>()
|
||||
.await
|
||||
.into_iter()
|
||||
.collect::<crate::Result<Vec<_>>>()
|
||||
.unwrap();
|
||||
|
||||
assert!(matches!(events[1], AssistantEvent::ToolCallStart { .. }));
|
||||
assert!(matches!(
|
||||
events[2],
|
||||
AssistantEvent::ToolCallArgumentsDelta { .. }
|
||||
));
|
||||
assert!(matches!(events[3], AssistantEvent::ToolCallStart { .. }));
|
||||
assert!(matches!(
|
||||
events[4],
|
||||
AssistantEvent::ToolCallArgumentsDelta { .. }
|
||||
));
|
||||
let collected = ChatEventStream::new(
|
||||
"req_final_only".to_string(),
|
||||
Box::pin(structured_chat_event_stream(
|
||||
stream::iter(events.into_iter().map(Ok)),
|
||||
true,
|
||||
)),
|
||||
)
|
||||
.collect_message()
|
||||
.await
|
||||
.unwrap();
|
||||
let tool_calls = collected.message.tool_calls().collect::<Vec<_>>();
|
||||
assert_eq!(tool_calls.len(), 2);
|
||||
assert_eq!(tool_calls[0].name, "first");
|
||||
assert_eq!(tool_calls[1].name, "second");
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,626 @@
|
||||
//! Adapts decoded text updates into parsed assistant deltas.
|
||||
//!
|
||||
//! This stage sits between low-level token decoding and final block assembly.
|
||||
//! It drives one unified parser that may emit normal text, reasoning text, or
|
||||
//! tool-call deltas, then normalizes those parser events into internal
|
||||
//! assistant events.
|
||||
|
||||
use asynk_strim_attr::{TryYielder, try_stream};
|
||||
use futures::{StreamExt as _, pin_mut};
|
||||
use thiserror_ext::AsReport;
|
||||
use tracing::warn;
|
||||
use vllm_parser::unified::{UnifiedParser, UnifiedParserEvent, UnifiedParserOutput};
|
||||
use vllm_text::output::DecodedTextEvent;
|
||||
|
||||
use crate::Result;
|
||||
use crate::error::Error;
|
||||
use crate::event::AssistantBlockKind;
|
||||
use crate::output::{AssistantEvent, DecodedTextEventStream, generate_tool_call_id};
|
||||
|
||||
/// Per-stream unified parsing state.
|
||||
struct UnifiedParserState {
|
||||
/// Parser for the current request stream.
|
||||
parser: Box<dyn UnifiedParser>,
|
||||
/// Whether unified parsing has already failed for this stream.
|
||||
parser_failed: bool,
|
||||
/// The parser-local index of the currently open tool call, if any.
|
||||
///
|
||||
/// Supported parsers currently emit at most one active tool call at a time.
|
||||
/// Change this to an indexed map if a model needs interleaved calls later.
|
||||
open_call_index: Option<usize>,
|
||||
}
|
||||
|
||||
impl UnifiedParserState {
|
||||
/// Create one fresh unified parsing state for a new streamed response.
|
||||
fn new(parser: Box<dyn UnifiedParser>) -> Self {
|
||||
Self {
|
||||
parser,
|
||||
parser_failed: false,
|
||||
open_call_index: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Initialize parser state once prompt token IDs are available.
|
||||
fn initialize(&mut self, prompt_token_ids: &[u32]) {
|
||||
if self.parser_failed {
|
||||
return;
|
||||
}
|
||||
|
||||
match self.parser.initialize(prompt_token_ids) {
|
||||
Ok(()) => {}
|
||||
Err(error) => {
|
||||
warn!(
|
||||
error = %error.as_report(),
|
||||
"failed to initialize unified parser; falling back to plain text deltas"
|
||||
);
|
||||
self.parser_failed = true;
|
||||
self.open_call_index = None;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Convert one decoded text delta into zero or more parsed assistant events.
|
||||
fn process_delta(&mut self, delta: String) -> Result<Vec<AssistantEvent>> {
|
||||
if self.parser_failed {
|
||||
self.open_call_index = None;
|
||||
return Ok(text_event(AssistantBlockKind::Text, delta).into_iter().collect());
|
||||
}
|
||||
|
||||
let mut output = UnifiedParserOutput::default();
|
||||
match self.parser.parse_into(&delta, &mut output) {
|
||||
Ok(()) => {
|
||||
let mut events = Vec::new();
|
||||
self.process_parser_output(output, &mut events)?;
|
||||
Ok(events)
|
||||
}
|
||||
Err(error) => {
|
||||
warn!(
|
||||
error = %error.as_report(),
|
||||
"unified parser failed; falling back to plain text deltas"
|
||||
);
|
||||
self.parser_failed = true;
|
||||
|
||||
let mut events = Vec::new();
|
||||
self.process_parser_output(output, &mut events)?;
|
||||
self.open_call_index = None;
|
||||
|
||||
let recovered = self.parser.reset();
|
||||
if recovered.is_empty() && events.is_empty() {
|
||||
push_text_delta(&mut events, AssistantBlockKind::Text, delta);
|
||||
} else {
|
||||
push_text_delta(&mut events, AssistantBlockKind::Text, recovered);
|
||||
}
|
||||
Ok(events)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Flush parser state at end-of-stream and close any remaining open calls.
|
||||
fn finish(&mut self) -> Result<Vec<AssistantEvent>> {
|
||||
let mut events = Vec::new();
|
||||
|
||||
if self.parser_failed {
|
||||
return Ok(events);
|
||||
}
|
||||
|
||||
match self.parser.finish() {
|
||||
Ok(output) => self.process_parser_output(output, &mut events)?,
|
||||
Err(error) => {
|
||||
warn!(
|
||||
error = %error.as_report(),
|
||||
"unified parser finish failed; closing open parser state"
|
||||
);
|
||||
self.parser_failed = true;
|
||||
self.open_call_index = None;
|
||||
// TODO: should we reset and emit the buffered text?
|
||||
}
|
||||
}
|
||||
|
||||
Ok(events)
|
||||
}
|
||||
|
||||
/// Apply one parsed unified output to the current stream state.
|
||||
fn process_parser_output(
|
||||
&mut self,
|
||||
output: UnifiedParserOutput,
|
||||
events: &mut Vec<AssistantEvent>,
|
||||
) -> Result<()> {
|
||||
for event in output.events {
|
||||
match event {
|
||||
UnifiedParserEvent::Text(delta) => {
|
||||
self.open_call_index = None;
|
||||
push_text_delta(events, AssistantBlockKind::Text, delta);
|
||||
}
|
||||
UnifiedParserEvent::Reasoning(delta) => {
|
||||
self.open_call_index = None;
|
||||
push_text_delta(events, AssistantBlockKind::Reasoning, delta);
|
||||
}
|
||||
UnifiedParserEvent::ToolCall(item) => {
|
||||
self.process_tool_item(item, events)?;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Apply one parsed tool-call delta emitted by the parser.
|
||||
fn process_tool_item(
|
||||
&mut self,
|
||||
item: vllm_parser::tool::ToolCallDelta,
|
||||
events: &mut Vec<AssistantEvent>,
|
||||
) -> Result<()> {
|
||||
if let Some(name) = item.name {
|
||||
let is_new_tool = match self.open_call_index {
|
||||
Some(open_call_index) => open_call_index != item.tool_index,
|
||||
None => true,
|
||||
};
|
||||
if is_new_tool {
|
||||
let id = self
|
||||
.parser
|
||||
.tool_call_id(item.tool_index)
|
||||
.map(str::to_string)
|
||||
.unwrap_or_else(generate_tool_call_id);
|
||||
self.open_call_index = Some(item.tool_index);
|
||||
events.push(AssistantEvent::ToolCallStart { id, name });
|
||||
}
|
||||
}
|
||||
|
||||
if item.arguments.is_empty() {
|
||||
return Ok(());
|
||||
}
|
||||
let Some(open_call_index) = self.open_call_index else {
|
||||
return Err(Error::ToolCallStreamInvariant {
|
||||
message: format!(
|
||||
"received arguments for tool index {} before any tool-call start",
|
||||
item.tool_index
|
||||
),
|
||||
});
|
||||
};
|
||||
if open_call_index != item.tool_index {
|
||||
return Err(Error::ToolCallStreamInvariant {
|
||||
message: format!(
|
||||
"received arguments for tool index {} while tool index {} is open",
|
||||
item.tool_index, open_call_index
|
||||
),
|
||||
});
|
||||
}
|
||||
|
||||
events.push(AssistantEvent::ToolCallArgumentsDelta {
|
||||
delta: item.arguments,
|
||||
});
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
/// Build one plain text event if `delta` is non-empty.
|
||||
fn text_event(kind: AssistantBlockKind, delta: String) -> Option<AssistantEvent> {
|
||||
if delta.is_empty() {
|
||||
return None;
|
||||
}
|
||||
Some(AssistantEvent::TextDelta { kind, delta })
|
||||
}
|
||||
|
||||
/// Push one plain text delta if it is non-empty.
|
||||
fn push_text_delta(events: &mut Vec<AssistantEvent>, kind: AssistantBlockKind, delta: String) {
|
||||
if let Some(event) = text_event(kind, delta) {
|
||||
events.push(event);
|
||||
}
|
||||
}
|
||||
|
||||
/// Wrap one decoded-text stream into the internal unified assistant stream.
|
||||
#[try_stream]
|
||||
pub(crate) async fn unified_event_stream(
|
||||
decoded_stream: impl DecodedTextEventStream,
|
||||
parser: Box<dyn UnifiedParser>,
|
||||
mut y: TryYielder<AssistantEvent, Error>,
|
||||
) -> Result<()> {
|
||||
pin_mut!(decoded_stream);
|
||||
|
||||
let mut state = UnifiedParserState::new(parser);
|
||||
|
||||
while let Some(event) = decoded_stream.next().await.transpose()? {
|
||||
match event {
|
||||
DecodedTextEvent::Start {
|
||||
prompt_token_ids,
|
||||
prompt_logprobs,
|
||||
} => {
|
||||
state.initialize(&prompt_token_ids);
|
||||
y.yield_ok(AssistantEvent::Start {
|
||||
prompt_token_ids,
|
||||
prompt_logprobs,
|
||||
})
|
||||
.await;
|
||||
}
|
||||
DecodedTextEvent::TextDelta {
|
||||
delta,
|
||||
token_ids,
|
||||
logprobs,
|
||||
finished,
|
||||
} => {
|
||||
for next in state.process_delta(delta)? {
|
||||
y.yield_ok(next).await;
|
||||
}
|
||||
if logprobs.is_some() || !token_ids.is_empty() {
|
||||
y.yield_ok(AssistantEvent::LogprobsDelta {
|
||||
logprobs,
|
||||
token_ids,
|
||||
})
|
||||
.await;
|
||||
}
|
||||
if let Some(finished) = finished {
|
||||
for next in state.finish()? {
|
||||
y.yield_ok(next).await;
|
||||
}
|
||||
y.yield_ok(AssistantEvent::Done {
|
||||
usage: finished.usage,
|
||||
finish_reason: finished.finish_reason,
|
||||
kv_transfer_params: finished.kv_transfer_params,
|
||||
})
|
||||
.await;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::collections::VecDeque;
|
||||
|
||||
use futures::{StreamExt as _, stream};
|
||||
use vllm_parser::reasoning::ReasoningError;
|
||||
use vllm_parser::tool::ToolCallDelta;
|
||||
use vllm_parser::unified::{UnifiedParserError, UnifiedParserOutput};
|
||||
|
||||
use super::unified_event_stream;
|
||||
use crate::event::AssistantBlockKind;
|
||||
use crate::output::AssistantEvent;
|
||||
|
||||
enum ScriptedStep {
|
||||
Output(UnifiedParserOutput),
|
||||
Error {
|
||||
committed: UnifiedParserOutput,
|
||||
reset_text: String,
|
||||
},
|
||||
}
|
||||
|
||||
struct ScriptedParser {
|
||||
steps: VecDeque<ScriptedStep>,
|
||||
reset_text: String,
|
||||
tool_call_id: Option<String>,
|
||||
finish_error_reset_text: Option<String>,
|
||||
}
|
||||
|
||||
impl ScriptedParser {
|
||||
fn new(steps: impl IntoIterator<Item = ScriptedStep>) -> Self {
|
||||
Self {
|
||||
steps: steps.into_iter().collect(),
|
||||
reset_text: String::new(),
|
||||
tool_call_id: Some("call_test".to_string()),
|
||||
finish_error_reset_text: None,
|
||||
}
|
||||
}
|
||||
|
||||
fn with_finish_error(mut self, reset_text: &str) -> Self {
|
||||
self.finish_error_reset_text = Some(reset_text.to_string());
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
impl vllm_parser::unified::UnifiedParser for ScriptedParser {
|
||||
fn create(
|
||||
_tools: &[vllm_parser::tool::Tool],
|
||||
_tokenizer: vllm_tokenizer::DynTokenizer,
|
||||
) -> vllm_parser::unified::Result<Box<dyn vllm_parser::unified::UnifiedParser>>
|
||||
where
|
||||
Self: Sized + 'static,
|
||||
{
|
||||
unreachable!("ScriptedParser is constructed directly in tests")
|
||||
}
|
||||
|
||||
fn parse_into(
|
||||
&mut self,
|
||||
_delta: &str,
|
||||
output: &mut UnifiedParserOutput,
|
||||
) -> vllm_parser::unified::Result<()> {
|
||||
match self.steps.pop_front().expect("unexpected parser call") {
|
||||
ScriptedStep::Output(next) => {
|
||||
output.append(next);
|
||||
Ok(())
|
||||
}
|
||||
ScriptedStep::Error {
|
||||
committed,
|
||||
reset_text,
|
||||
} => {
|
||||
output.append(committed);
|
||||
self.reset_text = reset_text;
|
||||
Err(UnifiedParserError::Reasoning(
|
||||
ReasoningError::MissingToken {
|
||||
token: "<think>".to_string(),
|
||||
},
|
||||
))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn tool_call_id(&self, _tool_index: usize) -> Option<&str> {
|
||||
self.tool_call_id.as_deref()
|
||||
}
|
||||
|
||||
fn finish(&mut self) -> vllm_parser::unified::Result<UnifiedParserOutput> {
|
||||
if let Some(reset_text) = self.finish_error_reset_text.take() {
|
||||
self.reset_text = reset_text;
|
||||
return Err(UnifiedParserError::Reasoning(
|
||||
ReasoningError::MissingToken {
|
||||
token: "</tool_call>".to_string(),
|
||||
},
|
||||
));
|
||||
}
|
||||
Ok(UnifiedParserOutput::default())
|
||||
}
|
||||
|
||||
fn reset(&mut self) -> String {
|
||||
std::mem::take(&mut self.reset_text)
|
||||
}
|
||||
}
|
||||
|
||||
fn decoded_delta(delta: &str) -> vllm_text::output::DecodedTextEvent {
|
||||
vllm_text::output::DecodedTextEvent::TextDelta {
|
||||
delta: delta.to_string(),
|
||||
token_ids: Vec::new(),
|
||||
logprobs: None,
|
||||
finished: None,
|
||||
}
|
||||
}
|
||||
|
||||
fn finished_delta(delta: &str) -> vllm_text::output::DecodedTextEvent {
|
||||
vllm_text::output::DecodedTextEvent::TextDelta {
|
||||
delta: delta.to_string(),
|
||||
token_ids: Vec::new(),
|
||||
logprobs: None,
|
||||
finished: Some(vllm_text::output::Finished {
|
||||
usage: vllm_llm::TokenUsage::default(),
|
||||
finish_reason: crate::FinishReason::Stop(None),
|
||||
kv_transfer_params: None,
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
async fn collect(
|
||||
parser: ScriptedParser,
|
||||
events: Vec<vllm_text::output::DecodedTextEvent>,
|
||||
) -> Vec<AssistantEvent> {
|
||||
let stream = stream::iter(events.into_iter().map(Ok));
|
||||
unified_event_stream(stream, Box::new(parser))
|
||||
.collect::<Vec<_>>()
|
||||
.await
|
||||
.into_iter()
|
||||
.collect::<crate::Result<Vec<_>>>()
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
fn text(delta: &str) -> UnifiedParserOutput {
|
||||
let mut output = UnifiedParserOutput::default();
|
||||
output.push_text(delta.to_string());
|
||||
output
|
||||
}
|
||||
|
||||
fn reasoning(delta: &str) -> UnifiedParserOutput {
|
||||
let mut output = UnifiedParserOutput::default();
|
||||
output.push_reasoning(delta.to_string());
|
||||
output
|
||||
}
|
||||
|
||||
fn tool_call(name: &str, arguments: &str) -> UnifiedParserOutput {
|
||||
UnifiedParserOutput {
|
||||
events: vec![vllm_parser::unified::UnifiedParserEvent::ToolCall(
|
||||
ToolCallDelta {
|
||||
tool_index: 0,
|
||||
name: Some(name.to_string()),
|
||||
arguments: arguments.to_string(),
|
||||
},
|
||||
)],
|
||||
}
|
||||
}
|
||||
|
||||
fn tool_call_arguments(arguments: &str) -> UnifiedParserOutput {
|
||||
UnifiedParserOutput {
|
||||
events: vec![vllm_parser::unified::UnifiedParserEvent::ToolCall(
|
||||
ToolCallDelta {
|
||||
tool_index: 0,
|
||||
name: None,
|
||||
arguments: arguments.to_string(),
|
||||
},
|
||||
)],
|
||||
}
|
||||
}
|
||||
|
||||
fn combined(first: UnifiedParserOutput, second: UnifiedParserOutput) -> UnifiedParserOutput {
|
||||
let mut output = first;
|
||||
output.append(second);
|
||||
output
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn unified_stream_emits_reasoning_only_deltas() {
|
||||
let events = collect(
|
||||
ScriptedParser::new([ScriptedStep::Output(reasoning("thinking"))]),
|
||||
vec![decoded_delta("raw")],
|
||||
)
|
||||
.await;
|
||||
|
||||
assert_eq!(
|
||||
events,
|
||||
vec![AssistantEvent::TextDelta {
|
||||
kind: AssistantBlockKind::Reasoning,
|
||||
delta: "thinking".to_string(),
|
||||
}]
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn unified_stream_emits_tool_only_deltas() {
|
||||
let events = collect(
|
||||
ScriptedParser::new([ScriptedStep::Output(tool_call(
|
||||
"get_weather",
|
||||
r#"{"location":"Paris"}"#,
|
||||
))]),
|
||||
vec![decoded_delta("raw")],
|
||||
)
|
||||
.await;
|
||||
|
||||
assert_eq!(
|
||||
events,
|
||||
vec![
|
||||
AssistantEvent::ToolCallStart {
|
||||
id: "call_test".to_string(),
|
||||
name: "get_weather".to_string(),
|
||||
},
|
||||
AssistantEvent::ToolCallArgumentsDelta {
|
||||
delta: r#"{"location":"Paris"}"#.to_string(),
|
||||
},
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn unified_stream_emits_reasoning_followed_by_tool_call() {
|
||||
let events = collect(
|
||||
ScriptedParser::new([ScriptedStep::Output(combined(
|
||||
reasoning("thinking"),
|
||||
tool_call("get_weather", r#"{"location":"Paris"}"#),
|
||||
))]),
|
||||
vec![decoded_delta("raw")],
|
||||
)
|
||||
.await;
|
||||
|
||||
assert_eq!(
|
||||
events,
|
||||
vec![
|
||||
AssistantEvent::TextDelta {
|
||||
kind: AssistantBlockKind::Reasoning,
|
||||
delta: "thinking".to_string(),
|
||||
},
|
||||
AssistantEvent::ToolCallStart {
|
||||
id: "call_test".to_string(),
|
||||
name: "get_weather".to_string(),
|
||||
},
|
||||
AssistantEvent::ToolCallArgumentsDelta {
|
||||
delta: r#"{"location":"Paris"}"#.to_string(),
|
||||
},
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn unified_stream_emits_visible_text_followed_by_tool_call() {
|
||||
let events = collect(
|
||||
ScriptedParser::new([ScriptedStep::Output(combined(
|
||||
text("visible "),
|
||||
tool_call("get_weather", r#"{"location":"Paris"}"#),
|
||||
))]),
|
||||
vec![decoded_delta("raw")],
|
||||
)
|
||||
.await;
|
||||
|
||||
assert_eq!(
|
||||
events,
|
||||
vec![
|
||||
AssistantEvent::TextDelta {
|
||||
kind: AssistantBlockKind::Text,
|
||||
delta: "visible ".to_string(),
|
||||
},
|
||||
AssistantEvent::ToolCallStart {
|
||||
id: "call_test".to_string(),
|
||||
name: "get_weather".to_string(),
|
||||
},
|
||||
AssistantEvent::ToolCallArgumentsDelta {
|
||||
delta: r#"{"location":"Paris"}"#.to_string(),
|
||||
},
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn unified_stream_emits_tool_arguments_before_trailing_text() {
|
||||
let events = collect(
|
||||
ScriptedParser::new([
|
||||
ScriptedStep::Output(tool_call("get_weather", "")),
|
||||
ScriptedStep::Output(combined(
|
||||
tool_call_arguments(r#"{"location":"Paris"}"#),
|
||||
text(" done"),
|
||||
)),
|
||||
]),
|
||||
vec![decoded_delta("start"), decoded_delta("finish")],
|
||||
)
|
||||
.await;
|
||||
|
||||
assert_eq!(
|
||||
events,
|
||||
vec![
|
||||
AssistantEvent::ToolCallStart {
|
||||
id: "call_test".to_string(),
|
||||
name: "get_weather".to_string(),
|
||||
},
|
||||
AssistantEvent::ToolCallArgumentsDelta {
|
||||
delta: r#"{"location":"Paris"}"#.to_string(),
|
||||
},
|
||||
AssistantEvent::TextDelta {
|
||||
kind: AssistantBlockKind::Text,
|
||||
delta: " done".to_string(),
|
||||
},
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn unified_stream_fallback_keeps_committed_output_and_disables_later_parsing() {
|
||||
let events = collect(
|
||||
ScriptedParser::new([ScriptedStep::Error {
|
||||
committed: text("committed"),
|
||||
reset_text: "buffered".to_string(),
|
||||
}]),
|
||||
vec![decoded_delta("bad"), decoded_delta("later")],
|
||||
)
|
||||
.await;
|
||||
|
||||
assert_eq!(
|
||||
events,
|
||||
vec![
|
||||
AssistantEvent::TextDelta {
|
||||
kind: AssistantBlockKind::Text,
|
||||
delta: "committed".to_string(),
|
||||
},
|
||||
AssistantEvent::TextDelta {
|
||||
kind: AssistantBlockKind::Text,
|
||||
delta: "buffered".to_string(),
|
||||
},
|
||||
AssistantEvent::TextDelta {
|
||||
kind: AssistantBlockKind::Text,
|
||||
delta: "later".to_string(),
|
||||
},
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn unified_stream_finish_error_closes_parser_without_reset_text() {
|
||||
let events = collect(
|
||||
ScriptedParser::new([ScriptedStep::Output(UnifiedParserOutput::default())])
|
||||
.with_finish_error("buffered"),
|
||||
vec![finished_delta("")],
|
||||
)
|
||||
.await;
|
||||
|
||||
assert_eq!(
|
||||
events,
|
||||
vec![AssistantEvent::Done {
|
||||
usage: vllm_llm::TokenUsage::default(),
|
||||
finish_reason: crate::FinishReason::Stop(None),
|
||||
kv_transfer_params: None,
|
||||
}]
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -2,7 +2,6 @@ use std::pin::Pin;
|
||||
use std::sync::Arc;
|
||||
|
||||
use futures::Stream;
|
||||
use subenum::subenum;
|
||||
use trait_set::trait_set;
|
||||
use uuid::Uuid;
|
||||
use vllm_llm::TokenUsage;
|
||||
@@ -22,23 +21,19 @@ pub(crate) use harmony::validate_harmony_parser_overrides;
|
||||
|
||||
/// Internal assistant event before final assembly.
|
||||
///
|
||||
/// - [`ContentEvent`]: subenum after reasoning parsing, carries only text content.
|
||||
/// - [`AssistantEvent`]: full event after tool parsing, adds tool-call variants.
|
||||
#[subenum(ContentEvent)]
|
||||
/// Unified parsing produces these events, and structured assembly consumes
|
||||
/// them to build public chat events.
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub(crate) enum AssistantEvent {
|
||||
#[subenum(ContentEvent)]
|
||||
Start {
|
||||
prompt_token_ids: Arc<[u32]>,
|
||||
prompt_logprobs: Option<DecodedPromptLogprobs>,
|
||||
},
|
||||
#[subenum(ContentEvent)]
|
||||
TextDelta {
|
||||
kind: AssistantBlockKind,
|
||||
delta: String,
|
||||
},
|
||||
/// Per-decoded-update sample metadata: logprobs and/or output token IDs.
|
||||
#[subenum(ContentEvent)]
|
||||
LogprobsDelta {
|
||||
logprobs: Option<DecodedLogprobs>,
|
||||
token_ids: Vec<u32>,
|
||||
@@ -48,7 +43,6 @@ pub(crate) enum AssistantEvent {
|
||||
/// A delta for the arguments of the currently open tool call. Must follow a
|
||||
/// `ToolCallStart`.
|
||||
ToolCallArgumentsDelta { delta: String },
|
||||
#[subenum(ContentEvent)]
|
||||
Done {
|
||||
usage: TokenUsage,
|
||||
finish_reason: FinishReason,
|
||||
@@ -57,50 +51,6 @@ pub(crate) enum AssistantEvent {
|
||||
},
|
||||
}
|
||||
|
||||
impl ContentEvent {
|
||||
/// Convert a [`DecodedTextEvent`] into one or more [`ContentEvent`] values
|
||||
/// by treating all text as plain (non-reasoning) content.
|
||||
fn from_decoded_plain_text(event: DecodedTextEvent) -> Vec<Self> {
|
||||
match event {
|
||||
DecodedTextEvent::Start {
|
||||
prompt_token_ids,
|
||||
prompt_logprobs,
|
||||
} => vec![Self::Start {
|
||||
prompt_token_ids,
|
||||
prompt_logprobs,
|
||||
}],
|
||||
DecodedTextEvent::TextDelta {
|
||||
delta,
|
||||
token_ids,
|
||||
logprobs,
|
||||
finished,
|
||||
} => {
|
||||
let mut events = Vec::new();
|
||||
if !delta.is_empty() {
|
||||
events.push(Self::TextDelta {
|
||||
kind: AssistantBlockKind::Text,
|
||||
delta,
|
||||
});
|
||||
}
|
||||
if logprobs.is_some() || !token_ids.is_empty() {
|
||||
events.push(Self::LogprobsDelta {
|
||||
logprobs,
|
||||
token_ids,
|
||||
});
|
||||
}
|
||||
if let Some(finished) = finished {
|
||||
events.push(Self::Done {
|
||||
usage: finished.usage,
|
||||
finish_reason: finished.finish_reason,
|
||||
kv_transfer_params: finished.kv_transfer_params,
|
||||
});
|
||||
}
|
||||
events
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Boxed stream of decoded text events coming from [`vllm_text`].
|
||||
pub type DynDecodedTextEventStream = Pin<Box<dyn Stream<Item = Result<DecodedTextEvent>> + Send>>;
|
||||
/// Boxed stream of structured chat events exposed by [`crate::ChatLlm`].
|
||||
|
||||
@@ -2,7 +2,7 @@
|
||||
|
||||
use std::sync::LazyLock;
|
||||
|
||||
pub use vllm_reasoning_parser::{
|
||||
pub use vllm_parser::reasoning::{
|
||||
CohereCmdReasoningParser, DeepSeekR1ReasoningParser, DeepSeekV3ReasoningParser,
|
||||
DeepSeekV4ReasoningParser, Gemma4ReasoningParser, Glm45ReasoningParser, KimiK2ReasoningParser,
|
||||
KimiReasoningParser, MiniMaxM2ReasoningParser, MiniMaxM3ReasoningParser,
|
||||
@@ -34,7 +34,7 @@ pub mod names {
|
||||
|
||||
/// Constructor signature for one registered reasoning parser implementation.
|
||||
type ReasoningParserCreator =
|
||||
fn(DynTokenizer) -> vllm_reasoning_parser::Result<Box<dyn ReasoningParser>>;
|
||||
fn(DynTokenizer) -> vllm_parser::reasoning::Result<Box<dyn ReasoningParser>>;
|
||||
|
||||
/// Registry and model matcher for reasoning parsers.
|
||||
pub type ReasoningParserFactory = ParserFactory<ReasoningParserCreator>;
|
||||
|
||||
@@ -2,13 +2,12 @@
|
||||
|
||||
use std::sync::LazyLock;
|
||||
|
||||
pub use vllm_tool_parser::{
|
||||
pub use vllm_parser::tool::{
|
||||
DeepSeekV3ToolParser, DeepSeekV4ToolParser, DeepSeekV31ToolParser, DeepSeekV32ToolParser,
|
||||
Gemma4ToolParser, Glm45MoeToolParser, Glm47MoeToolParser, Granite4ToolParser, HermesToolParser,
|
||||
HyV3ToolParser, Internlm2ToolParser, KimiK2ToolParser, Llama3JsonToolParser,
|
||||
MinimaxM2ToolParser, MinimaxM3ToolParser, MistralToolParser, Phi4MiniJsonToolParser,
|
||||
Qwen3CoderToolParser, Qwen3XmlToolParser, ToolCallDelta, ToolParser, ToolParserError,
|
||||
ToolParserOutput,
|
||||
Qwen3CoderToolParser, Qwen3XmlToolParser, ToolParser, ToolParserError,
|
||||
};
|
||||
|
||||
use crate::parser::ParserFactory;
|
||||
@@ -41,7 +40,7 @@ pub mod names {
|
||||
}
|
||||
|
||||
/// Constructor signature for one registered tool parser implementation.
|
||||
type ToolParserCreator = fn(&[ChatTool]) -> vllm_tool_parser::Result<Box<dyn ToolParser>>;
|
||||
type ToolParserCreator = fn(&[ChatTool]) -> vllm_parser::tool::Result<Box<dyn ToolParser>>;
|
||||
|
||||
/// Registry and model matcher for tool parsers.
|
||||
pub type ToolParserFactory = ParserFactory<ToolParserCreator>;
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
use vllm_tool_parser::Result;
|
||||
use vllm_parser::tool::{Result, ToolParserOutput};
|
||||
|
||||
use super::{ToolParser, ToolParserFactory, ToolParserOutput, names};
|
||||
use super::{ToolParser, ToolParserFactory, names};
|
||||
use crate::Error;
|
||||
use crate::request::ChatTool;
|
||||
|
||||
|
||||
@@ -4,9 +4,9 @@ use llm_multimodal::ImageDetail;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::Value;
|
||||
use vllm_engine_core_client::protocol::lora::LoraRequest;
|
||||
pub use vllm_parser::tool::Tool as ChatTool;
|
||||
pub use vllm_text::SamplingParams;
|
||||
use vllm_text::TextDecodeOptions;
|
||||
pub use vllm_tool_parser::Tool as ChatTool;
|
||||
|
||||
use crate::AssistantMessageExt;
|
||||
use crate::error::{Error, Result};
|
||||
|
||||
@@ -1180,7 +1180,7 @@ async fn chat_stream_parses_tool_calls_automatically() {
|
||||
),
|
||||
request_output(
|
||||
"chat-tool",
|
||||
bytes_to_token_ids(
|
||||
bytes_with_special_stop_token(
|
||||
b"\"arguments\":{\"city\":\"Paris\"}}\n</tool_call>",
|
||||
),
|
||||
Some(EngineCoreFinishReason::Stop),
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
[package]
|
||||
name = "vllm-tool-parser"
|
||||
name = "vllm-parser"
|
||||
version.workspace = true
|
||||
edition.workspace = true
|
||||
license.workspace = true
|
||||
@@ -13,6 +13,7 @@ serde.workspace = true
|
||||
serde_json.workspace = true
|
||||
thiserror.workspace = true
|
||||
thiserror-ext.workspace = true
|
||||
vllm-tokenizer.workspace = true
|
||||
winnow.workspace = true
|
||||
xgrammar-structural-tag.workspace = true
|
||||
|
||||
+2
-2
@@ -2,8 +2,8 @@ use std::time::Duration;
|
||||
|
||||
use criterion::{BatchSize, Criterion, Throughput, black_box, criterion_group, criterion_main};
|
||||
use tool_parser::parsers::DeepSeekParser as ExternalDeepSeekParser;
|
||||
use vllm_tool_parser::test_utils::{split_by_chars, test_tools};
|
||||
use vllm_tool_parser::{DeepSeekV3ToolParser, Tool, ToolParser};
|
||||
use vllm_parser::tool::test_utils::{split_by_chars, test_tools};
|
||||
use vllm_parser::tool::{DeepSeekV3ToolParser, Tool, ToolParser};
|
||||
|
||||
mod utils;
|
||||
use utils::{feed_external_parser, feed_parser, openai_tools};
|
||||
+2
-2
@@ -2,8 +2,8 @@ use std::time::Duration;
|
||||
|
||||
use criterion::{BatchSize, Criterion, Throughput, black_box, criterion_group, criterion_main};
|
||||
use tool_parser::parsers::DeepSeek31Parser as ExternalDeepSeek31Parser;
|
||||
use vllm_tool_parser::test_utils::{split_by_chars, test_tools};
|
||||
use vllm_tool_parser::{DeepSeekV31ToolParser, Tool, ToolParser};
|
||||
use vllm_parser::tool::test_utils::{split_by_chars, test_tools};
|
||||
use vllm_parser::tool::{DeepSeekV31ToolParser, Tool, ToolParser};
|
||||
|
||||
mod utils;
|
||||
use utils::{feed_external_parser, feed_parser, openai_tools};
|
||||
+2
-2
@@ -1,8 +1,8 @@
|
||||
use std::time::Duration;
|
||||
|
||||
use criterion::{BatchSize, Criterion, Throughput, black_box, criterion_group, criterion_main};
|
||||
use vllm_tool_parser::test_utils::{split_by_chars, test_tools};
|
||||
use vllm_tool_parser::{DeepSeekV32ToolParser, Tool, ToolParser};
|
||||
use vllm_parser::tool::test_utils::{split_by_chars, test_tools};
|
||||
use vllm_parser::tool::{DeepSeekV32ToolParser, Tool, ToolParser};
|
||||
|
||||
mod utils;
|
||||
use utils::feed_parser;
|
||||
@@ -1,8 +1,8 @@
|
||||
use std::time::Duration;
|
||||
|
||||
use criterion::{BatchSize, Criterion, Throughput, black_box, criterion_group, criterion_main};
|
||||
use vllm_tool_parser::test_utils::{split_by_chars, test_tools};
|
||||
use vllm_tool_parser::{Gemma4ToolParser, Tool, ToolParser};
|
||||
use vllm_parser::tool::test_utils::{split_by_chars, test_tools};
|
||||
use vllm_parser::tool::{Gemma4ToolParser, Tool, ToolParser};
|
||||
|
||||
mod utils;
|
||||
use utils::feed_parser;
|
||||
@@ -2,8 +2,8 @@ use std::time::Duration;
|
||||
|
||||
use criterion::{BatchSize, Criterion, Throughput, black_box, criterion_group, criterion_main};
|
||||
use tool_parser::parsers::Glm4MoeParser as ExternalGlm4MoeParser;
|
||||
use vllm_tool_parser::test_utils::{split_by_chars, test_tools};
|
||||
use vllm_tool_parser::{Glm45MoeToolParser, Glm47MoeToolParser, Tool, ToolParser};
|
||||
use vllm_parser::tool::test_utils::{split_by_chars, test_tools};
|
||||
use vllm_parser::tool::{Glm45MoeToolParser, Glm47MoeToolParser, Tool, ToolParser};
|
||||
|
||||
mod utils;
|
||||
use utils::{feed_external_parser, feed_parser, openai_tools};
|
||||
@@ -2,8 +2,8 @@ use std::time::Duration;
|
||||
|
||||
use criterion::{BatchSize, Criterion, Throughput, black_box, criterion_group, criterion_main};
|
||||
use tool_parser::parsers::KimiK2Parser as ExternalKimiK2Parser;
|
||||
use vllm_tool_parser::test_utils::{split_by_chars, test_tools};
|
||||
use vllm_tool_parser::{KimiK2ToolParser, Tool, ToolParser};
|
||||
use vllm_parser::tool::test_utils::{split_by_chars, test_tools};
|
||||
use vllm_parser::tool::{KimiK2ToolParser, Tool, ToolParser};
|
||||
|
||||
mod utils;
|
||||
use utils::{feed_external_parser, feed_parser, openai_tools};
|
||||
+2
-2
@@ -2,8 +2,8 @@ use std::time::Duration;
|
||||
|
||||
use criterion::{BatchSize, Criterion, Throughput, black_box, criterion_group, criterion_main};
|
||||
use tool_parser::parsers::LlamaParser as ExternalLlamaParser;
|
||||
use vllm_tool_parser::test_utils::{split_by_chars, test_tools};
|
||||
use vllm_tool_parser::{Llama3JsonToolParser, Tool, ToolParser};
|
||||
use vllm_parser::tool::test_utils::{split_by_chars, test_tools};
|
||||
use vllm_parser::tool::{Llama3JsonToolParser, Tool, ToolParser};
|
||||
|
||||
mod utils;
|
||||
use utils::{feed_external_parser, feed_parser, openai_tools};
|
||||
@@ -2,8 +2,8 @@ use std::time::Duration;
|
||||
|
||||
use criterion::{BatchSize, Criterion, Throughput, black_box, criterion_group, criterion_main};
|
||||
use tool_parser::parsers::MinimaxM2Parser as ExternalMinimaxM2Parser;
|
||||
use vllm_tool_parser::test_utils::{split_by_chars, test_tools};
|
||||
use vllm_tool_parser::{MinimaxM2ToolParser, Tool, ToolParser};
|
||||
use vllm_parser::tool::test_utils::{split_by_chars, test_tools};
|
||||
use vllm_parser::tool::{MinimaxM2ToolParser, Tool, ToolParser};
|
||||
|
||||
mod utils;
|
||||
use utils::{feed_external_parser, feed_parser, openai_tools};
|
||||
+2
-2
@@ -2,8 +2,8 @@ use std::time::Duration;
|
||||
|
||||
use criterion::{BatchSize, Criterion, Throughput, black_box, criterion_group, criterion_main};
|
||||
use tool_parser::parsers::QwenCoderParser as ExternalQwenCoderParser;
|
||||
use vllm_tool_parser::test_utils::{split_by_chars, test_tools};
|
||||
use vllm_tool_parser::{Qwen3CoderToolParser, Tool, ToolParser};
|
||||
use vllm_parser::tool::test_utils::{split_by_chars, test_tools};
|
||||
use vllm_parser::tool::{Qwen3CoderToolParser, Tool, ToolParser};
|
||||
|
||||
mod utils;
|
||||
use utils::{feed_external_parser, feed_parser, openai_tools};
|
||||
@@ -2,8 +2,8 @@ use std::time::Duration;
|
||||
|
||||
use criterion::{BatchSize, Criterion, Throughput, black_box, criterion_group, criterion_main};
|
||||
use tool_parser::parsers::QwenParser as ExternalQwenParser;
|
||||
use vllm_tool_parser::test_utils::{split_by_chars, test_tools};
|
||||
use vllm_tool_parser::{Qwen3XmlToolParser, Tool, ToolParser};
|
||||
use vllm_parser::tool::test_utils::{split_by_chars, test_tools};
|
||||
use vllm_parser::tool::{Qwen3XmlToolParser, Tool, ToolParser};
|
||||
|
||||
mod utils;
|
||||
use utils::{feed_external_parser, feed_parser, openai_tools};
|
||||
@@ -3,8 +3,8 @@
|
||||
use futures::FutureExt as _;
|
||||
use openai_protocol::common::{Function as OpenAiFunction, Tool as OpenAiTool};
|
||||
use tool_parser::traits::ToolParser as ExternalToolParser;
|
||||
use vllm_tool_parser::test_utils::collect_stream;
|
||||
use vllm_tool_parser::{Tool, ToolParser};
|
||||
use vllm_parser::tool::test_utils::collect_stream;
|
||||
use vllm_parser::tool::{Tool, ToolParser};
|
||||
|
||||
pub(super) fn openai_tools(tools: &[Tool]) -> Vec<OpenAiTool> {
|
||||
tools
|
||||
@@ -23,7 +23,7 @@ pub(super) fn openai_tools(tools: &[Tool]) -> Vec<OpenAiTool> {
|
||||
|
||||
pub(super) fn feed_parser(parser: &mut dyn ToolParser, chunks: &[&str]) -> (String, usize) {
|
||||
let result = collect_stream(parser, chunks);
|
||||
(result.normal_text, result.calls.len())
|
||||
(result.normal_text(), result.calls().len())
|
||||
}
|
||||
|
||||
pub(super) fn feed_external_parser(
|
||||
@@ -13,7 +13,7 @@ pyo3.workspace = true
|
||||
pythonize = { workspace = true, features = ["serde_json"] }
|
||||
serde_json.workspace = true
|
||||
thiserror-ext.workspace = true
|
||||
vllm-tool-parser.workspace = true
|
||||
vllm-parser.workspace = true
|
||||
|
||||
[lints]
|
||||
workspace = true
|
||||
@@ -1,4 +1,4 @@
|
||||
//! Thin PyO3 bindings for `vllm_tool_parser`.
|
||||
//! 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
|
||||
@@ -11,7 +11,7 @@ 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};
|
||||
use vllm_parser::tool::{Tool, ToolCallDelta, ToolParser, ToolParserOutput};
|
||||
|
||||
macro_rules! tool_parser_factory {
|
||||
($($parser:ident),+ $(,)?) => {
|
||||
@@ -22,7 +22,7 @@ macro_rules! tool_parser_factory {
|
||||
match name {
|
||||
$(
|
||||
stringify!($parser) => {
|
||||
<vllm_tool_parser::$parser as ToolParser>::create(tools)
|
||||
<vllm_parser::tool::$parser as ToolParser>::create(tools)
|
||||
}
|
||||
)+
|
||||
_ => {
|
||||
@@ -146,30 +146,30 @@ 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,
|
||||
})
|
||||
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) -> &str {
|
||||
&self.0.normal_text
|
||||
fn normal_text(&self) -> String {
|
||||
self.0.normal_text()
|
||||
}
|
||||
|
||||
#[getter]
|
||||
fn calls(&self) -> Vec<PyToolCallDelta> {
|
||||
self.0.calls.iter().cloned().map(PyToolCallDelta).collect()
|
||||
self.0.calls().into_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())
|
||||
fn coalesce(&self) -> Self {
|
||||
Self(self.0.clone().coalesce())
|
||||
}
|
||||
}
|
||||
|
||||
@@ -300,7 +300,7 @@ mod tests {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn output_append_and_coalesce_calls() {
|
||||
fn output_append_and_coalesce() {
|
||||
with_python(|py| {
|
||||
let first = Py::new(
|
||||
py,
|
||||
@@ -311,7 +311,7 @@ mod tests {
|
||||
let other = Py::new(py, PyToolParserOutput::new(py, "", Some(vec![second])))?;
|
||||
output.append(other.borrow(py));
|
||||
|
||||
let coalesced = output.coalesce_calls();
|
||||
let coalesced = output.coalesce();
|
||||
assert_eq!(coalesced.normal_text(), "text");
|
||||
let calls = coalesced.calls();
|
||||
assert_eq!(calls.len(), 1);
|
||||
@@ -334,7 +334,7 @@ mod tests {
|
||||
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();
|
||||
let output = output.coalesce();
|
||||
|
||||
assert_eq!(output.normal_text(), "");
|
||||
let calls = output.calls();
|
||||
@@ -0,0 +1,5 @@
|
||||
//! Streaming parsers for chat completions.
|
||||
|
||||
pub mod reasoning;
|
||||
pub mod tool;
|
||||
pub mod unified;
|
||||
@@ -119,7 +119,7 @@ mod tests {
|
||||
use vllm_tokenizer::Tokenizer;
|
||||
|
||||
use super::Gemma4ReasoningParser;
|
||||
use crate::ReasoningParser;
|
||||
use crate::reasoning::ReasoningParser;
|
||||
|
||||
struct FakeTokenizer;
|
||||
|
||||
+1
-1
@@ -49,7 +49,7 @@ mod tests {
|
||||
use std::sync::Arc;
|
||||
|
||||
use super::SeedOssReasoningParser;
|
||||
use crate::{ReasoningParser, tests::FakeTokenizer};
|
||||
use crate::reasoning::{ReasoningParser, tests::FakeTokenizer};
|
||||
|
||||
#[test]
|
||||
fn without_prompt_markers_expects_start_token() {
|
||||
@@ -127,7 +127,7 @@ mod tests {
|
||||
use std::sync::Arc;
|
||||
|
||||
use super::Step3p5ReasoningParser;
|
||||
use crate::{ReasoningParser, tests::FakeTokenizer};
|
||||
use crate::reasoning::{ReasoningParser, tests::FakeTokenizer};
|
||||
|
||||
#[test]
|
||||
fn picks_up_prompt_start_boundary() {
|
||||
+47
-47
@@ -1,5 +1,5 @@
|
||||
use super::{DeepSeekDsmlToolParser, DsmlTokens};
|
||||
use crate::{Result, StructuralTagModel, Tool, ToolParser, ToolParserOutput};
|
||||
use crate::tool::{Result, StructuralTagModel, Tool, ToolParser, ToolParserOutput};
|
||||
|
||||
/// Tool parser for DeepSeek V3.2 models.
|
||||
///
|
||||
@@ -67,8 +67,8 @@ mod tests {
|
||||
use thiserror_ext::AsReport;
|
||||
|
||||
use super::DeepSeekV32ToolParser;
|
||||
use crate::test_utils::{collect_stream, split_by_chars, test_tools};
|
||||
use crate::{ToolParser, ToolParserTestExt as _};
|
||||
use crate::tool::test_utils::{collect_stream, split_by_chars, test_tools};
|
||||
use crate::tool::{ToolParser, ToolParserTestExt as _};
|
||||
|
||||
fn build_tool_call(function_name: &str, params: &[(&str, &str)]) -> String {
|
||||
let params = params
|
||||
@@ -90,8 +90,8 @@ mod tests {
|
||||
let mut parser = DeepSeekV32ToolParser::new(&test_tools());
|
||||
let output = parser.parse_complete("Hello, world!").unwrap();
|
||||
|
||||
assert_eq!(output.normal_text, "Hello, world!");
|
||||
assert!(output.calls.is_empty());
|
||||
assert_eq!(output.normal_text(), "Hello, world!");
|
||||
assert!(output.calls().is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -104,11 +104,11 @@ mod tests {
|
||||
))
|
||||
.unwrap();
|
||||
|
||||
assert!(output.normal_text.is_empty());
|
||||
assert_eq!(output.calls.len(), 1);
|
||||
assert_eq!(output.calls[0].name.as_deref(), Some("get_weather"));
|
||||
assert!(output.normal_text().is_empty());
|
||||
assert_eq!(output.calls().len(), 1);
|
||||
assert_eq!(output.calls()[0].name.as_deref(), Some("get_weather"));
|
||||
assert_eq!(
|
||||
serde_json::from_str::<Value>(&output.calls[0].arguments).unwrap(),
|
||||
serde_json::from_str::<Value>(&output.calls()[0].arguments).unwrap(),
|
||||
json!({
|
||||
"location": "SF",
|
||||
"date": "2024-01-16"
|
||||
@@ -125,8 +125,8 @@ mod tests {
|
||||
);
|
||||
let output = parser.parse_complete(&output).unwrap();
|
||||
|
||||
assert_eq!(output.normal_text, "Thinking... ");
|
||||
assert_eq!(output.calls.len(), 1);
|
||||
assert_eq!(output.normal_text(), "Thinking... ");
|
||||
assert_eq!(output.calls().len(), 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -146,9 +146,9 @@ mod tests {
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(output.calls.len(), 1);
|
||||
assert_eq!(output.calls().len(), 1);
|
||||
assert_eq!(
|
||||
serde_json::from_str::<Value>(&output.calls[0].arguments).unwrap(),
|
||||
serde_json::from_str::<Value>(&output.calls()[0].arguments).unwrap(),
|
||||
json!({
|
||||
"whole": 5.0,
|
||||
"flag": true,
|
||||
@@ -176,9 +176,9 @@ mod tests {
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(output.calls.len(), 1);
|
||||
assert_eq!(output.calls().len(), 1);
|
||||
assert_eq!(
|
||||
serde_json::from_str::<Value>(&output.calls[0].arguments).unwrap(),
|
||||
serde_json::from_str::<Value>(&output.calls()[0].arguments).unwrap(),
|
||||
json!({
|
||||
"whole": "5.0",
|
||||
"flag": "true",
|
||||
@@ -206,7 +206,7 @@ mod tests {
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
serde_json::from_str::<Value>(&output.calls[0].arguments).unwrap(),
|
||||
serde_json::from_str::<Value>(&output.calls()[0].arguments).unwrap(),
|
||||
json!({
|
||||
"location": "Hangzhou </|DSML|parameter></|DSML|invoke></|DSML|function_calls>",
|
||||
"date": "2026-05-08",
|
||||
@@ -228,11 +228,11 @@ mod tests {
|
||||
],
|
||||
);
|
||||
|
||||
assert!(output.normal_text.is_empty());
|
||||
assert_eq!(output.calls.len(), 1);
|
||||
assert_eq!(output.calls[0].name.as_deref(), Some("get_weather"));
|
||||
assert!(output.normal_text().is_empty());
|
||||
assert_eq!(output.calls().len(), 1);
|
||||
assert_eq!(output.calls()[0].name.as_deref(), Some("get_weather"));
|
||||
assert_eq!(
|
||||
serde_json::from_str::<Value>(&output.calls[0].arguments).unwrap(),
|
||||
serde_json::from_str::<Value>(&output.calls()[0].arguments).unwrap(),
|
||||
json!({ "location": "SF" })
|
||||
);
|
||||
}
|
||||
@@ -252,8 +252,8 @@ mod tests {
|
||||
],
|
||||
);
|
||||
|
||||
assert_eq!(output.normal_text, "Thinking... ");
|
||||
assert_eq!(output.calls.len(), 1);
|
||||
assert_eq!(output.normal_text(), "Thinking... ");
|
||||
assert_eq!(output.calls().len(), 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -261,8 +261,8 @@ mod tests {
|
||||
let mut parser = DeepSeekV32ToolParser::new(&test_tools());
|
||||
let output = collect_stream(&mut parser, &["Hello, ", "world!"]);
|
||||
|
||||
assert_eq!(output.normal_text, "Hello, world!");
|
||||
assert!(output.calls.is_empty());
|
||||
assert_eq!(output.normal_text(), "Hello, world!");
|
||||
assert!(output.calls().is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -278,17 +278,17 @@ mod tests {
|
||||
)],
|
||||
);
|
||||
|
||||
assert_eq!(output.calls.len(), 2);
|
||||
assert_eq!(output.calls[0].name.as_deref(), Some("get_weather"));
|
||||
assert_eq!(output.calls[1].name.as_deref(), Some("get_weather"));
|
||||
assert_eq!(output.calls[0].tool_index, 0);
|
||||
assert_eq!(output.calls[1].tool_index, 1);
|
||||
assert_eq!(output.calls().len(), 2);
|
||||
assert_eq!(output.calls()[0].name.as_deref(), Some("get_weather"));
|
||||
assert_eq!(output.calls()[1].name.as_deref(), Some("get_weather"));
|
||||
assert_eq!(output.calls()[0].tool_index, 0);
|
||||
assert_eq!(output.calls()[1].tool_index, 1);
|
||||
assert_eq!(
|
||||
serde_json::from_str::<Value>(&output.calls[0].arguments).unwrap(),
|
||||
serde_json::from_str::<Value>(&output.calls()[0].arguments).unwrap(),
|
||||
json!({ "location": "SF" })
|
||||
);
|
||||
assert_eq!(
|
||||
serde_json::from_str::<Value>(&output.calls[1].arguments).unwrap(),
|
||||
serde_json::from_str::<Value>(&output.calls()[1].arguments).unwrap(),
|
||||
json!({ "location": "NYC" })
|
||||
);
|
||||
}
|
||||
@@ -300,9 +300,9 @@ mod tests {
|
||||
let mut parser = DeepSeekV32ToolParser::new(&test_tools());
|
||||
let output = collect_stream(&mut parser, &chunks);
|
||||
|
||||
assert_eq!(output.calls.len(), 1);
|
||||
assert_eq!(output.calls().len(), 1);
|
||||
assert_eq!(
|
||||
serde_json::from_str::<Value>(&output.calls[0].arguments).unwrap(),
|
||||
serde_json::from_str::<Value>(&output.calls()[0].arguments).unwrap(),
|
||||
json!({ "location": "SF" })
|
||||
);
|
||||
}
|
||||
@@ -337,11 +337,11 @@ mod tests {
|
||||
],
|
||||
);
|
||||
|
||||
assert!(output.normal_text.is_empty());
|
||||
assert_eq!(output.calls.len(), 1);
|
||||
assert_eq!(output.calls[0].name.as_deref(), Some("get_weather"));
|
||||
assert!(output.normal_text().is_empty());
|
||||
assert_eq!(output.calls().len(), 1);
|
||||
assert_eq!(output.calls()[0].name.as_deref(), Some("get_weather"));
|
||||
assert_eq!(
|
||||
serde_json::from_str::<Value>(&output.calls[0].arguments).unwrap(),
|
||||
serde_json::from_str::<Value>(&output.calls()[0].arguments).unwrap(),
|
||||
json!({ "location": "Beijing" })
|
||||
);
|
||||
}
|
||||
@@ -373,9 +373,9 @@ mod tests {
|
||||
],
|
||||
);
|
||||
|
||||
assert!(output.normal_text.is_empty());
|
||||
assert_eq!(output.calls.len(), 1);
|
||||
assert_eq!(output.calls[0].name.as_deref(), Some("get_weather"));
|
||||
assert!(output.normal_text().is_empty());
|
||||
assert_eq!(output.calls().len(), 1);
|
||||
assert_eq!(output.calls()[0].name.as_deref(), Some("get_weather"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -393,8 +393,8 @@ mod tests {
|
||||
],
|
||||
);
|
||||
|
||||
assert!(output.normal_text.is_empty());
|
||||
assert_eq!(output.calls.len(), 1);
|
||||
assert!(output.normal_text().is_empty());
|
||||
assert_eq!(output.calls().len(), 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -421,10 +421,10 @@ mod tests {
|
||||
.parse_complete(&build_tool_call("get_weather", &[("location", "NYC")]))
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(first.calls.len(), 1);
|
||||
assert_eq!(second.calls.len(), 1);
|
||||
assert_eq!(first.calls().len(), 1);
|
||||
assert_eq!(second.calls().len(), 1);
|
||||
assert_eq!(
|
||||
serde_json::from_str::<Value>(&second.calls[0].arguments).unwrap(),
|
||||
serde_json::from_str::<Value>(&second.calls()[0].arguments).unwrap(),
|
||||
json!({ "location": "NYC" })
|
||||
);
|
||||
}
|
||||
@@ -439,7 +439,7 @@ mod tests {
|
||||
let mut parser = DeepSeekV32ToolParser::new(&test_tools());
|
||||
let complete = parser.parse_complete(&full_text).unwrap();
|
||||
|
||||
assert_eq!(streamed.normal_text, complete.normal_text);
|
||||
assert_eq!(streamed.calls, complete.calls);
|
||||
assert_eq!(streamed.normal_text(), complete.normal_text());
|
||||
assert_eq!(streamed.calls(), complete.calls());
|
||||
}
|
||||
}
|
||||
+11
-11
@@ -1,5 +1,5 @@
|
||||
use super::{DeepSeekDsmlToolParser, DsmlTokens};
|
||||
use crate::{Result, StructuralTagModel, Tool, ToolParser, ToolParserOutput};
|
||||
use crate::tool::{Result, StructuralTagModel, Tool, ToolParser, ToolParserOutput};
|
||||
|
||||
/// Tool parser for DeepSeek V4 models.
|
||||
///
|
||||
@@ -69,8 +69,8 @@ mod tests {
|
||||
use serde_json::{Value, json};
|
||||
|
||||
use super::DeepSeekV4ToolParser;
|
||||
use crate::test_utils::{collect_stream, test_tools};
|
||||
use crate::{StructuralTagModel, ToolParser, ToolParserTestExt as _};
|
||||
use crate::tool::test_utils::{collect_stream, test_tools};
|
||||
use crate::tool::{StructuralTagModel, ToolParser, ToolParserTestExt as _};
|
||||
|
||||
fn build_tool_call(function_name: &str, params: &[(&str, &str)]) -> String {
|
||||
let params = params
|
||||
@@ -107,11 +107,11 @@ mod tests {
|
||||
))
|
||||
.unwrap();
|
||||
|
||||
assert!(output.normal_text.is_empty());
|
||||
assert_eq!(output.calls.len(), 1);
|
||||
assert_eq!(output.calls[0].name.as_deref(), Some("get_weather"));
|
||||
assert!(output.normal_text().is_empty());
|
||||
assert_eq!(output.calls().len(), 1);
|
||||
assert_eq!(output.calls()[0].name.as_deref(), Some("get_weather"));
|
||||
assert_eq!(
|
||||
serde_json::from_str::<Value>(&output.calls[0].arguments).unwrap(),
|
||||
serde_json::from_str::<Value>(&output.calls()[0].arguments).unwrap(),
|
||||
json!({
|
||||
"location": "SF",
|
||||
"date": "2024-01-16"
|
||||
@@ -137,11 +137,11 @@ mod tests {
|
||||
],
|
||||
);
|
||||
|
||||
assert_eq!(output.normal_text, "Thinking... ");
|
||||
assert_eq!(output.calls.len(), 1);
|
||||
assert_eq!(output.calls[0].name.as_deref(), Some("get_weather"));
|
||||
assert_eq!(output.normal_text(), "Thinking... ");
|
||||
assert_eq!(output.calls().len(), 1);
|
||||
assert_eq!(output.calls()[0].name.as_deref(), Some("get_weather"));
|
||||
assert_eq!(
|
||||
serde_json::from_str::<Value>(&output.calls[0].arguments).unwrap(),
|
||||
serde_json::from_str::<Value>(&output.calls()[0].arguments).unwrap(),
|
||||
json!({ "location": "Beijing" })
|
||||
);
|
||||
}
|
||||
+4
-4
@@ -7,7 +7,7 @@ use winnow::token::{literal, rest, take_until};
|
||||
use super::parameters::ToolSchemas;
|
||||
use super::utils::{MarkerScanState, parse_buffered_event, safe_text_len, take_until_marker};
|
||||
use super::{Result, ToolCallDelta, ToolParserOutput};
|
||||
use crate::Tool;
|
||||
use crate::tool::Tool;
|
||||
|
||||
mod deepseek_v32;
|
||||
mod deepseek_v4;
|
||||
@@ -92,7 +92,7 @@ impl DeepSeekDsmlToolParser {
|
||||
fn apply_event(&mut self, event: DsmlEvent, output: &mut ToolParserOutput) -> Result<()> {
|
||||
match event {
|
||||
DsmlEvent::Text { len: consumed_len } => {
|
||||
output.normal_text.push_str(&self.buffer[..consumed_len]);
|
||||
output.push_text(&self.buffer[..consumed_len]);
|
||||
}
|
||||
DsmlEvent::ToolCallsStart => {
|
||||
self.mode = DsmlMode::ToolBlock {
|
||||
@@ -116,7 +116,7 @@ impl DeepSeekDsmlToolParser {
|
||||
let arguments = serde_json::to_string(&arguments)
|
||||
.map_err(|error| parsing_failed!("failed to serialize arguments: {}", error))?;
|
||||
|
||||
output.calls.push(ToolCallDelta {
|
||||
output.push_call(ToolCallDelta {
|
||||
tool_index: self.emitted_invoke_count,
|
||||
name: Some(name),
|
||||
arguments,
|
||||
@@ -156,7 +156,7 @@ impl DeepSeekDsmlToolParser {
|
||||
fn finish(&mut self) -> Result<ToolParserOutput> {
|
||||
let mut output = ToolParserOutput::default();
|
||||
match self.mode {
|
||||
DsmlMode::Text => output.normal_text.push_str(&self.buffer),
|
||||
DsmlMode::Text => output.push_text(&self.buffer),
|
||||
DsmlMode::Done => {}
|
||||
DsmlMode::ToolBlock { .. } => {
|
||||
return Err(parsing_failed!("incomplete DeepSeek DSML tool call"));
|
||||
+38
-35
@@ -1,5 +1,5 @@
|
||||
use super::{DeepSeekJsonFormat, DeepSeekJsonToolParser};
|
||||
use crate::{Result, StructuralTagModel, Tool, ToolParser, ToolParserOutput};
|
||||
use crate::tool::{Result, StructuralTagModel, Tool, ToolParser, ToolParserOutput};
|
||||
|
||||
/// Tool parser for DeepSeek V3 JSON-fenced tool calls.
|
||||
///
|
||||
@@ -55,12 +55,12 @@ mod tests {
|
||||
use thiserror_ext::AsReport;
|
||||
|
||||
use super::DeepSeekV3ToolParser;
|
||||
use crate::deepseek_json::{
|
||||
use crate::tool::deepseek_json::{
|
||||
TOOL_CALL_SEPARATOR, TOOL_CALL_START, TOOL_CALLS_END, TOOL_CALLS_START, V3_ARGUMENT_END,
|
||||
V3_JSON_START,
|
||||
};
|
||||
use crate::test_utils::{collect_stream, split_by_chars, test_tools};
|
||||
use crate::{ToolParser, ToolParserOutput, ToolParserTestExt as _};
|
||||
use crate::tool::test_utils::{collect_stream, split_by_chars, test_tools};
|
||||
use crate::tool::{ToolParser, ToolParserOutput, ToolParserTestExt as _};
|
||||
|
||||
fn v3_tool_call(function_name: &str, arguments: &str) -> String {
|
||||
format!(
|
||||
@@ -77,8 +77,8 @@ mod tests {
|
||||
let mut parser = DeepSeekV3ToolParser::new(&test_tools());
|
||||
let output = parser.parse_complete("Hello, world!").unwrap();
|
||||
|
||||
assert_eq!(output.normal_text, "Hello, world!");
|
||||
assert!(output.calls.is_empty());
|
||||
assert_eq!(output.normal_text(), "Hello, world!");
|
||||
assert!(output.calls().is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -92,11 +92,11 @@ mod tests {
|
||||
))
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(output.normal_text, "Let me check.\n");
|
||||
assert_eq!(output.calls.len(), 1);
|
||||
assert_eq!(output.calls[0].tool_index, 0);
|
||||
assert_eq!(output.calls[0].name.as_deref(), Some("get_weather"));
|
||||
assert_eq!(output.calls[0].arguments, arguments);
|
||||
assert_eq!(output.normal_text(), "Let me check.\n");
|
||||
assert_eq!(output.calls().len(), 1);
|
||||
assert_eq!(output.calls()[0].tool_index, 0);
|
||||
assert_eq!(output.calls()[0].name.as_deref(), Some("get_weather"));
|
||||
assert_eq!(output.calls()[0].arguments, arguments);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -107,7 +107,7 @@ mod tests {
|
||||
.parse_complete(&tool_section(&[v3_tool_call("get_weather", arguments)]))
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(output.calls[0].arguments, arguments);
|
||||
assert_eq!(output.calls()[0].arguments, arguments);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -132,7 +132,7 @@ mod tests {
|
||||
for chunk in chunks {
|
||||
let next = parser.parse_chunk(chunk).unwrap();
|
||||
observed_arguments.extend(
|
||||
next.calls
|
||||
next.calls()
|
||||
.iter()
|
||||
.filter(|call| call.name.is_none())
|
||||
.map(|call| call.arguments.clone()),
|
||||
@@ -143,7 +143,7 @@ mod tests {
|
||||
|
||||
assert_eq!(observed_arguments, ["{\"location\":", "\"Beijing\"", "}"]);
|
||||
assert_eq!(
|
||||
output.coalesce_calls().calls[0].arguments,
|
||||
output.coalesce().calls()[0].arguments,
|
||||
r#"{"location":"Beijing"}"#
|
||||
);
|
||||
}
|
||||
@@ -159,9 +159,9 @@ mod tests {
|
||||
|
||||
let output = collect_stream(&mut parser, &chunks);
|
||||
|
||||
assert_eq!(output.normal_text, "hello ");
|
||||
assert_eq!(output.calls.len(), 1);
|
||||
assert_eq!(output.calls[0].arguments, r#"{"location":"Tokyo"}"#);
|
||||
assert_eq!(output.normal_text(), "hello ");
|
||||
assert_eq!(output.calls().len(), 1);
|
||||
assert_eq!(output.calls()[0].arguments, r#"{"location":"Tokyo"}"#);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -172,8 +172,8 @@ mod tests {
|
||||
|
||||
let output = parser.parse_complete(&input).unwrap();
|
||||
|
||||
assert_eq!(output.calls.len(), 1);
|
||||
assert_eq!(output.calls[0].arguments, arguments);
|
||||
assert_eq!(output.calls().len(), 1);
|
||||
assert_eq!(output.calls()[0].arguments, arguments);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -189,22 +189,25 @@ mod tests {
|
||||
|
||||
expect![[r#"
|
||||
ToolParserOutput {
|
||||
normal_text: "",
|
||||
calls: [
|
||||
ToolCallDelta {
|
||||
tool_index: 0,
|
||||
name: Some(
|
||||
"get_weather",
|
||||
),
|
||||
arguments: "{\"location\":\"Shanghai\"}",
|
||||
},
|
||||
ToolCallDelta {
|
||||
tool_index: 1,
|
||||
name: Some(
|
||||
"add",
|
||||
),
|
||||
arguments: "{\"x\":1,\"y\":2}",
|
||||
},
|
||||
events: [
|
||||
ToolCall(
|
||||
ToolCallDelta {
|
||||
tool_index: 0,
|
||||
name: Some(
|
||||
"get_weather",
|
||||
),
|
||||
arguments: "{\"location\":\"Shanghai\"}",
|
||||
},
|
||||
),
|
||||
ToolCall(
|
||||
ToolCallDelta {
|
||||
tool_index: 1,
|
||||
name: Some(
|
||||
"add",
|
||||
),
|
||||
arguments: "{\"x\":1,\"y\":2}",
|
||||
},
|
||||
),
|
||||
],
|
||||
}
|
||||
"#]]
|
||||
+41
-38
@@ -1,5 +1,5 @@
|
||||
use super::{DeepSeekJsonFormat, DeepSeekJsonToolParser};
|
||||
use crate::{Result, StructuralTagModel, Tool, ToolParser, ToolParserOutput};
|
||||
use crate::tool::{Result, StructuralTagModel, Tool, ToolParser, ToolParserOutput};
|
||||
|
||||
/// Tool parser for DeepSeek V3.1 raw JSON tool calls.
|
||||
///
|
||||
@@ -51,11 +51,11 @@ mod tests {
|
||||
use thiserror_ext::AsReport;
|
||||
|
||||
use super::DeepSeekV31ToolParser;
|
||||
use crate::deepseek_json::{
|
||||
use crate::tool::deepseek_json::{
|
||||
TOOL_CALL_END, TOOL_CALL_SEPARATOR, TOOL_CALL_START, TOOL_CALLS_END, TOOL_CALLS_START,
|
||||
};
|
||||
use crate::test_utils::{collect_stream, split_by_chars, test_tools};
|
||||
use crate::{ToolParser, ToolParserOutput, ToolParserTestExt as _};
|
||||
use crate::tool::test_utils::{collect_stream, split_by_chars, test_tools};
|
||||
use crate::tool::{ToolParser, ToolParserOutput, ToolParserTestExt as _};
|
||||
|
||||
fn v31_tool_call(function_name: &str, arguments: &str) -> String {
|
||||
format!("{TOOL_CALL_START}{function_name}{TOOL_CALL_SEPARATOR}{arguments}{TOOL_CALL_END}")
|
||||
@@ -70,8 +70,8 @@ mod tests {
|
||||
let mut parser = DeepSeekV31ToolParser::new(&test_tools());
|
||||
let output = parser.parse_complete("Hello, world!").unwrap();
|
||||
|
||||
assert_eq!(output.normal_text, "Hello, world!");
|
||||
assert!(output.calls.is_empty());
|
||||
assert_eq!(output.normal_text(), "Hello, world!");
|
||||
assert!(output.calls().is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -85,11 +85,11 @@ mod tests {
|
||||
))
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(output.normal_text, "Let me check.");
|
||||
assert_eq!(output.calls.len(), 1);
|
||||
assert_eq!(output.calls[0].tool_index, 0);
|
||||
assert_eq!(output.calls[0].name.as_deref(), Some("get_weather"));
|
||||
assert_eq!(output.calls[0].arguments, arguments);
|
||||
assert_eq!(output.normal_text(), "Let me check.");
|
||||
assert_eq!(output.calls().len(), 1);
|
||||
assert_eq!(output.calls()[0].tool_index, 0);
|
||||
assert_eq!(output.calls()[0].name.as_deref(), Some("get_weather"));
|
||||
assert_eq!(output.calls()[0].arguments, arguments);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -100,7 +100,7 @@ mod tests {
|
||||
.parse_complete(&tool_section(&[v31_tool_call("get_weather", arguments)]))
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(output.calls[0].arguments, arguments);
|
||||
assert_eq!(output.calls()[0].arguments, arguments);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -123,7 +123,7 @@ mod tests {
|
||||
for chunk in chunks {
|
||||
let next = parser.parse_chunk(chunk).unwrap();
|
||||
observed_arguments.extend(
|
||||
next.calls
|
||||
next.calls()
|
||||
.iter()
|
||||
.filter(|call| call.name.is_none())
|
||||
.map(|call| call.arguments.clone()),
|
||||
@@ -134,7 +134,7 @@ mod tests {
|
||||
|
||||
assert_eq!(observed_arguments, ["{\"location\":", "\"Beijing\"", "}"]);
|
||||
assert_eq!(
|
||||
output.coalesce_calls().calls[0].arguments,
|
||||
output.coalesce().calls()[0].arguments,
|
||||
r#"{"location":"Beijing"}"#
|
||||
);
|
||||
}
|
||||
@@ -150,9 +150,9 @@ mod tests {
|
||||
|
||||
let output = collect_stream(&mut parser, &chunks);
|
||||
|
||||
assert_eq!(output.normal_text, "hello ");
|
||||
assert_eq!(output.calls.len(), 1);
|
||||
assert_eq!(output.calls[0].arguments, r#"{"location":"Tokyo"}"#);
|
||||
assert_eq!(output.normal_text(), "hello ");
|
||||
assert_eq!(output.calls().len(), 1);
|
||||
assert_eq!(output.calls()[0].arguments, r#"{"location":"Tokyo"}"#);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -163,8 +163,8 @@ mod tests {
|
||||
|
||||
let output = parser.parse_complete(&input).unwrap();
|
||||
|
||||
assert_eq!(output.calls.len(), 1);
|
||||
assert_eq!(output.calls[0].arguments, arguments);
|
||||
assert_eq!(output.calls().len(), 1);
|
||||
assert_eq!(output.calls()[0].arguments, arguments);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -180,22 +180,25 @@ mod tests {
|
||||
|
||||
expect![[r#"
|
||||
ToolParserOutput {
|
||||
normal_text: "",
|
||||
calls: [
|
||||
ToolCallDelta {
|
||||
tool_index: 0,
|
||||
name: Some(
|
||||
"get_weather",
|
||||
),
|
||||
arguments: "{\"location\":\"Shanghai\"}",
|
||||
},
|
||||
ToolCallDelta {
|
||||
tool_index: 1,
|
||||
name: Some(
|
||||
"add",
|
||||
),
|
||||
arguments: "{\"x\":1,\"y\":2}",
|
||||
},
|
||||
events: [
|
||||
ToolCall(
|
||||
ToolCallDelta {
|
||||
tool_index: 0,
|
||||
name: Some(
|
||||
"get_weather",
|
||||
),
|
||||
arguments: "{\"location\":\"Shanghai\"}",
|
||||
},
|
||||
),
|
||||
ToolCall(
|
||||
ToolCallDelta {
|
||||
tool_index: 1,
|
||||
name: Some(
|
||||
"add",
|
||||
),
|
||||
arguments: "{\"x\":1,\"y\":2}",
|
||||
},
|
||||
),
|
||||
],
|
||||
}
|
||||
"#]]
|
||||
@@ -212,9 +215,9 @@ mod tests {
|
||||
|
||||
let output = collect_stream(&mut parser, &[&input]);
|
||||
|
||||
assert!(output.normal_text.is_empty());
|
||||
assert_eq!(output.calls.len(), 1);
|
||||
assert_eq!(output.calls[0].arguments, r#"{"location":"Tokyo"}"#);
|
||||
assert!(output.normal_text().is_empty());
|
||||
assert_eq!(output.calls().len(), 1);
|
||||
assert_eq!(output.calls()[0].arguments, r#"{"location":"Tokyo"}"#);
|
||||
}
|
||||
|
||||
#[test]
|
||||
+4
-4
@@ -96,7 +96,7 @@ impl DeepSeekJsonToolParser {
|
||||
) -> Result<()> {
|
||||
match event {
|
||||
DeepSeekJsonEvent::Text { len: consumed_len } => {
|
||||
output.normal_text.push_str(&self.buffer[..consumed_len]);
|
||||
output.push_text(&self.buffer[..consumed_len]);
|
||||
}
|
||||
DeepSeekJsonEvent::ToolCallsStart => self.mode = DeepSeekJsonMode::ToolBlock,
|
||||
DeepSeekJsonEvent::ToolCallStart => self.mode = DeepSeekJsonMode::Header,
|
||||
@@ -107,7 +107,7 @@ impl DeepSeekJsonToolParser {
|
||||
self.mode = DeepSeekJsonMode::Arguments {
|
||||
json_scan: JsonObjectScanState::default(),
|
||||
};
|
||||
output.calls.push(ToolCallDelta {
|
||||
output.push_call(ToolCallDelta {
|
||||
tool_index,
|
||||
name: Some(function_name),
|
||||
arguments: String::new(),
|
||||
@@ -120,7 +120,7 @@ impl DeepSeekJsonToolParser {
|
||||
self.format.parser_name()
|
||||
));
|
||||
};
|
||||
output.calls.push(ToolCallDelta {
|
||||
output.push_call(ToolCallDelta {
|
||||
tool_index,
|
||||
name: None,
|
||||
arguments: self.buffer[..consumed_len].to_string(),
|
||||
@@ -155,7 +155,7 @@ impl DeepSeekJsonToolParser {
|
||||
fn finish(&mut self) -> Result<ToolParserOutput> {
|
||||
let mut output = ToolParserOutput::default();
|
||||
match &self.mode {
|
||||
DeepSeekJsonMode::Text => output.normal_text.push_str(&self.buffer),
|
||||
DeepSeekJsonMode::Text => output.push_text(&self.buffer),
|
||||
DeepSeekJsonMode::ToolBlock | DeepSeekJsonMode::Done => {}
|
||||
DeepSeekJsonMode::Header | DeepSeekJsonMode::Arguments { .. } => {
|
||||
return Err(parsing_failed!(
|
||||
@@ -6,7 +6,7 @@ pub type Result<T> = std::result::Result<T, ToolParserError>;
|
||||
|
||||
/// Errors produced while creating or running tool parsers.
|
||||
#[derive(Debug, Error, Macro)]
|
||||
#[thiserror_ext(macro(path = "crate::error"))]
|
||||
#[thiserror_ext(macro(path = "crate::tool::error"))]
|
||||
pub enum ToolParserError {
|
||||
#[error("tool parser parsing failed: {message}")]
|
||||
ParsingFailed { message: String },
|
||||
@@ -8,7 +8,7 @@ use winnow::token::{literal, take_till, take_until};
|
||||
|
||||
use super::utils::{incomplete, parse_buffered_event, partial_prefix_len, safe_text_len};
|
||||
use super::{Result, ToolCallDelta, ToolParser, ToolParserOutput};
|
||||
use crate::Tool;
|
||||
use crate::tool::Tool;
|
||||
|
||||
const TOOL_CALL_START: &str = "<|tool_call>";
|
||||
const TOOL_CALL_END: &str = "<tool_call|>";
|
||||
@@ -70,7 +70,7 @@ impl Gemma4ToolParser {
|
||||
fn apply_event(&mut self, event: Gemma4Event, output: &mut ToolParserOutput) -> Result<()> {
|
||||
match event {
|
||||
Gemma4Event::Text { len: consumed_len } => {
|
||||
output.normal_text.push_str(&self.buffer[..consumed_len]);
|
||||
output.push_text(&self.buffer[..consumed_len]);
|
||||
}
|
||||
Gemma4Event::ToolCallStart => self.mode = Gemma4Mode::Header,
|
||||
Gemma4Event::ToolCallHeader { name } => {
|
||||
@@ -89,7 +89,7 @@ impl Gemma4ToolParser {
|
||||
let arguments = serde_json::to_string(&args)
|
||||
.map_err(|error| parsing_failed!("failed to serialize arguments: {}", error))?;
|
||||
|
||||
output.calls.push(ToolCallDelta {
|
||||
output.push_call(ToolCallDelta {
|
||||
tool_index: self.emitted_tool_count,
|
||||
name: Some(name),
|
||||
arguments,
|
||||
@@ -153,7 +153,7 @@ impl ToolParser for Gemma4ToolParser {
|
||||
let mut output = ToolParserOutput::default();
|
||||
|
||||
match &self.mode {
|
||||
Gemma4Mode::Text => output.normal_text.push_str(&self.buffer),
|
||||
Gemma4Mode::Text => output.push_text(&self.buffer),
|
||||
Gemma4Mode::Header | Gemma4Mode::ToolCall { .. } => {
|
||||
return Err(parsing_failed!("incomplete Gemma4 tool call"));
|
||||
}
|
||||
@@ -428,7 +428,7 @@ mod tests {
|
||||
Gemma4ToolParser, ToolCallDelta, ToolParser, ToolParserOutput, gemma4_array_content,
|
||||
parse_gemma4_args,
|
||||
};
|
||||
use crate::{Tool, ToolParserTestExt as _};
|
||||
use crate::tool::{Tool, ToolParserTestExt as _};
|
||||
|
||||
fn parse_gemma4_array(array: &str) -> super::Result<Vec<Value>> {
|
||||
let mut input = array;
|
||||
@@ -501,11 +501,11 @@ mod tests {
|
||||
output.append(parser.parse_chunk(chunk).unwrap());
|
||||
}
|
||||
output.append(parser.finish().unwrap());
|
||||
output.coalesce_calls()
|
||||
output.coalesce()
|
||||
}
|
||||
|
||||
fn first_call(output: &ToolParserOutput) -> &ToolCallDelta {
|
||||
output.calls.first().expect("expected one tool call")
|
||||
fn first_call(output: &ToolParserOutput) -> ToolCallDelta {
|
||||
(*output.calls().first().expect("expected one tool call")).clone()
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -547,8 +547,8 @@ mod tests {
|
||||
.parse_complete("<|tool_call>call:get_weather{location:<|\"|>London<|\"|>}<tool_call|>")
|
||||
.unwrap();
|
||||
|
||||
assert!(output.normal_text.is_empty());
|
||||
assert_eq!(output.calls.len(), 1);
|
||||
assert!(output.normal_text().is_empty());
|
||||
assert_eq!(output.calls().len(), 1);
|
||||
assert_eq!(first_call(&output).name.as_deref(), Some("get_weather"));
|
||||
assert_eq!(
|
||||
serde_json::from_str::<Value>(&first_call(&output).arguments).unwrap(),
|
||||
@@ -577,7 +577,7 @@ mod tests {
|
||||
"<tool_call|>",
|
||||
]);
|
||||
|
||||
assert!(output.normal_text.is_empty());
|
||||
assert!(output.normal_text().is_empty());
|
||||
assert_eq!(first_call(&output).name.as_deref(), Some("get_weather"));
|
||||
assert_eq!(
|
||||
serde_json::from_str::<Value>(&first_call(&output).arguments).unwrap(),
|
||||
@@ -597,7 +597,7 @@ mod tests {
|
||||
"div>",
|
||||
]);
|
||||
|
||||
assert_eq!(output.normal_text, "Let me check the weather. <div>");
|
||||
assert_eq!(output.normal_text(), "Let me check the weather. <div>");
|
||||
assert_eq!(first_call(&output).name.as_deref(), Some("get_weather"));
|
||||
assert_eq!(
|
||||
serde_json::from_str::<Value>(&first_call(&output).arguments).unwrap(),
|
||||
@@ -616,11 +616,11 @@ mod tests {
|
||||
"location:<|\"|>Paris<|\"|>}",
|
||||
] {
|
||||
output.append(parser.parse_chunk(chunk).unwrap());
|
||||
assert!(output.calls.is_empty());
|
||||
assert!(output.calls().is_empty());
|
||||
}
|
||||
|
||||
output.append(parser.parse_chunk("<tool_call|>").unwrap());
|
||||
let output = output.coalesce_calls();
|
||||
let output = output.coalesce();
|
||||
|
||||
assert_eq!(first_call(&output).name.as_deref(), Some("get_weather"));
|
||||
assert_eq!(
|
||||
@@ -777,8 +777,8 @@ mod tests {
|
||||
let mut output = parser.parse_chunk("<").unwrap();
|
||||
output.append(parser.finish().unwrap());
|
||||
|
||||
assert_eq!(output.normal_text, "<");
|
||||
assert!(output.calls.is_empty());
|
||||
assert_eq!(output.normal_text(), "<");
|
||||
assert!(output.calls().is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
+1
-1
@@ -1,5 +1,5 @@
|
||||
use super::{GlmXmlToolParser, Separator};
|
||||
use crate::{Result, Tool, ToolParser, ToolParserOutput};
|
||||
use crate::tool::{Result, Tool, ToolParser, ToolParserOutput};
|
||||
|
||||
/// Tool parser for GLM-4.5/4.6 MoE XML-style tool calls.
|
||||
///
|
||||
+16
-16
@@ -1,5 +1,5 @@
|
||||
use super::{GlmXmlToolParser, Separator};
|
||||
use crate::{Result, StructuralTagModel, Tool, ToolParser, ToolParserOutput};
|
||||
use crate::tool::{Result, StructuralTagModel, Tool, ToolParser, ToolParserOutput};
|
||||
|
||||
/// Tool parser for GLM-4.7 MoE XML-style tool calls.
|
||||
///
|
||||
@@ -44,8 +44,8 @@ mod tests {
|
||||
use serde_json::{Value, json};
|
||||
|
||||
use super::Glm47MoeToolParser;
|
||||
use crate::ToolParserTestExt as _;
|
||||
use crate::test_utils::{collect_stream, split_by_chars, test_tools};
|
||||
use crate::tool::ToolParserTestExt as _;
|
||||
use crate::tool::test_utils::{collect_stream, split_by_chars, test_tools};
|
||||
|
||||
fn glm47_tool_call(function_name: &str, params: &[(&str, &str)]) -> String {
|
||||
let params = params
|
||||
@@ -69,11 +69,11 @@ mod tests {
|
||||
|
||||
let output = parser.parse_complete(&output).unwrap();
|
||||
|
||||
assert_eq!(output.normal_text, "Let me search for that.\n");
|
||||
assert_eq!(output.calls.len(), 1);
|
||||
assert_eq!(output.calls[0].name.as_deref(), Some("get_weather"));
|
||||
assert_eq!(output.normal_text(), "Let me search for that.\n");
|
||||
assert_eq!(output.calls().len(), 1);
|
||||
assert_eq!(output.calls()[0].name.as_deref(), Some("get_weather"));
|
||||
assert_eq!(
|
||||
serde_json::from_str::<Value>(&output.calls[0].arguments).unwrap(),
|
||||
serde_json::from_str::<Value>(&output.calls()[0].arguments).unwrap(),
|
||||
json!({"city": "Beijing", "date": "2024-12-25"})
|
||||
);
|
||||
}
|
||||
@@ -90,12 +90,12 @@ mod tests {
|
||||
let chunks = split_by_chars(&output, 7);
|
||||
let output = collect_stream(&mut parser, &chunks);
|
||||
|
||||
assert_eq!(output.normal_text, "");
|
||||
assert_eq!(output.calls.len(), 2);
|
||||
assert_eq!(output.calls[0].name.as_deref(), Some("get_weather"));
|
||||
assert_eq!(output.calls[1].name.as_deref(), Some("add"));
|
||||
assert_eq!(output.normal_text(), "");
|
||||
assert_eq!(output.calls().len(), 2);
|
||||
assert_eq!(output.calls()[0].name.as_deref(), Some("get_weather"));
|
||||
assert_eq!(output.calls()[1].name.as_deref(), Some("add"));
|
||||
assert_eq!(
|
||||
serde_json::from_str::<Value>(&output.calls[1].arguments).unwrap(),
|
||||
serde_json::from_str::<Value>(&output.calls()[1].arguments).unwrap(),
|
||||
json!({"x": 1, "y": 2})
|
||||
);
|
||||
}
|
||||
@@ -117,7 +117,7 @@ mod tests {
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
serde_json::from_str::<Value>(&output.calls[0].arguments).unwrap(),
|
||||
serde_json::from_str::<Value>(&output.calls()[0].arguments).unwrap(),
|
||||
json!({
|
||||
"whole": 42,
|
||||
"flag": true,
|
||||
@@ -134,10 +134,10 @@ mod tests {
|
||||
|
||||
let output = parser.parse_complete("<tool_call>add</tool_call>").unwrap();
|
||||
|
||||
assert_eq!(output.calls.len(), 1);
|
||||
assert_eq!(output.calls[0].name.as_deref(), Some("add"));
|
||||
assert_eq!(output.calls().len(), 1);
|
||||
assert_eq!(output.calls()[0].name.as_deref(), Some("add"));
|
||||
assert_eq!(
|
||||
serde_json::from_str::<Value>(&output.calls[0].arguments).unwrap(),
|
||||
serde_json::from_str::<Value>(&output.calls()[0].arguments).unwrap(),
|
||||
json!({})
|
||||
);
|
||||
}
|
||||
@@ -7,7 +7,7 @@ use winnow::token::{literal, rest, take_until, take_while};
|
||||
use super::parameters::ToolSchemas;
|
||||
use super::utils::{MarkerScanState, parse_buffered_event, safe_text_len, take_until_marker};
|
||||
use super::{Result, ToolCallDelta, ToolParserOutput};
|
||||
use crate::Tool;
|
||||
use crate::tool::Tool;
|
||||
|
||||
mod glm45_moe;
|
||||
mod glm47_moe;
|
||||
@@ -79,7 +79,7 @@ impl GlmXmlToolParser {
|
||||
fn apply_event(&mut self, event: GlmEvent, output: &mut ToolParserOutput) -> Result<()> {
|
||||
match event {
|
||||
GlmEvent::Text { len: consumed_len } => {
|
||||
output.normal_text.push_str(&self.buffer[..consumed_len]);
|
||||
output.push_text(&self.buffer[..consumed_len]);
|
||||
}
|
||||
GlmEvent::ToolCallStart => {
|
||||
self.mode = GlmMode::ToolCall {
|
||||
@@ -92,7 +92,7 @@ impl GlmXmlToolParser {
|
||||
let arguments = serde_json::to_string(&arguments)
|
||||
.map_err(|error| parsing_failed!("failed to serialize arguments: {}", error))?;
|
||||
|
||||
output.calls.push(ToolCallDelta {
|
||||
output.push_call(ToolCallDelta {
|
||||
tool_index: self.emitted_tool_count,
|
||||
name: Some(name),
|
||||
arguments,
|
||||
@@ -127,7 +127,7 @@ impl GlmXmlToolParser {
|
||||
let mut output = ToolParserOutput::default();
|
||||
if !self.buffer.is_empty() {
|
||||
match self.mode {
|
||||
GlmMode::Text => output.normal_text.push_str(&self.buffer),
|
||||
GlmMode::Text => output.push_text(&self.buffer),
|
||||
GlmMode::ToolCall { .. } => {
|
||||
return Err(parsing_failed!("incomplete GLM MoE tool call"));
|
||||
}
|
||||
@@ -264,8 +264,8 @@ mod tests {
|
||||
use thiserror_ext::AsReport;
|
||||
|
||||
use super::Glm45MoeToolParser;
|
||||
use crate::test_utils::{collect_stream, split_by_chars, test_tools};
|
||||
use crate::{ToolParser, ToolParserTestExt as _};
|
||||
use crate::tool::test_utils::{collect_stream, split_by_chars, test_tools};
|
||||
use crate::tool::{ToolParser, ToolParserTestExt as _};
|
||||
|
||||
fn glm45_tool_call(function_name: &str, params: &[(&str, &str)]) -> String {
|
||||
let params = params
|
||||
@@ -283,8 +283,8 @@ mod tests {
|
||||
let mut parser = Glm45MoeToolParser::new(&test_tools());
|
||||
let output = parser.parse_complete("Hello, world!").unwrap();
|
||||
|
||||
assert_eq!(output.normal_text, "Hello, world!");
|
||||
assert!(output.calls.is_empty());
|
||||
assert_eq!(output.normal_text(), "Hello, world!");
|
||||
assert!(output.calls().is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -300,11 +300,11 @@ mod tests {
|
||||
|
||||
let output = parser.parse_complete(&output).unwrap();
|
||||
|
||||
assert_eq!(output.normal_text, "Let me search for that.\n");
|
||||
assert_eq!(output.calls.len(), 1);
|
||||
assert_eq!(output.calls[0].name.as_deref(), Some("get_weather"));
|
||||
assert_eq!(output.normal_text(), "Let me search for that.\n");
|
||||
assert_eq!(output.calls().len(), 1);
|
||||
assert_eq!(output.calls()[0].name.as_deref(), Some("get_weather"));
|
||||
assert_eq!(
|
||||
serde_json::from_str::<Value>(&output.calls[0].arguments).unwrap(),
|
||||
serde_json::from_str::<Value>(&output.calls()[0].arguments).unwrap(),
|
||||
json!({"city": "Beijing", "date": "2024-12-25"})
|
||||
);
|
||||
}
|
||||
@@ -321,12 +321,12 @@ mod tests {
|
||||
let chunks = split_by_chars(&output, 11);
|
||||
let output = collect_stream(&mut parser, &chunks);
|
||||
|
||||
assert_eq!(output.normal_text, "");
|
||||
assert_eq!(output.calls.len(), 2);
|
||||
assert_eq!(output.calls[0].name.as_deref(), Some("get_weather"));
|
||||
assert_eq!(output.calls[1].name.as_deref(), Some("add"));
|
||||
assert_eq!(output.normal_text(), "");
|
||||
assert_eq!(output.calls().len(), 2);
|
||||
assert_eq!(output.calls()[0].name.as_deref(), Some("get_weather"));
|
||||
assert_eq!(output.calls()[1].name.as_deref(), Some("add"));
|
||||
assert_eq!(
|
||||
serde_json::from_str::<Value>(&output.calls[1].arguments).unwrap(),
|
||||
serde_json::from_str::<Value>(&output.calls()[1].arguments).unwrap(),
|
||||
json!({"x": 1, "y": 2})
|
||||
);
|
||||
}
|
||||
@@ -345,7 +345,7 @@ mod tests {
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
serde_json::from_str::<Value>(&output.calls[0].arguments).unwrap(),
|
||||
serde_json::from_str::<Value>(&output.calls()[0].arguments).unwrap(),
|
||||
json!({
|
||||
"city": "Paris </arg_value></tool_call>",
|
||||
"date": "2026-05-08",
|
||||
@@ -359,8 +359,8 @@ mod tests {
|
||||
|
||||
let output = collect_stream(&mut parser, &["hello ", "world"]);
|
||||
|
||||
assert_eq!(output.normal_text, "hello world");
|
||||
assert!(output.calls.is_empty());
|
||||
assert_eq!(output.normal_text(), "hello world");
|
||||
assert!(output.calls().is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -375,8 +375,8 @@ mod tests {
|
||||
],
|
||||
);
|
||||
|
||||
assert_eq!(output.normal_text, "Prefix ");
|
||||
assert_eq!(output.calls.len(), 1);
|
||||
assert_eq!(output.normal_text(), "Prefix ");
|
||||
assert_eq!(output.calls().len(), 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -391,9 +391,9 @@ mod tests {
|
||||
],
|
||||
);
|
||||
|
||||
assert_eq!(output.normal_text, "hello ");
|
||||
assert_eq!(output.calls.len(), 1);
|
||||
assert_eq!(output.calls[0].name.as_deref(), Some("get_weather"));
|
||||
assert_eq!(output.normal_text(), "hello ");
|
||||
assert_eq!(output.calls().len(), 1);
|
||||
assert_eq!(output.calls()[0].name.as_deref(), Some("get_weather"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -402,8 +402,8 @@ mod tests {
|
||||
|
||||
let output = parser.parse_chunk("<tool_call>get_weather\n<arg_key>city</arg_key>").unwrap();
|
||||
|
||||
assert_eq!(output.normal_text, "");
|
||||
assert!(output.calls.is_empty());
|
||||
assert_eq!(output.normal_text(), "");
|
||||
assert!(output.calls().is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -437,7 +437,7 @@ mod tests {
|
||||
)],
|
||||
);
|
||||
|
||||
assert_eq!(output.normal_text, "");
|
||||
assert_eq!(output.calls.len(), 1);
|
||||
assert_eq!(output.normal_text(), "");
|
||||
assert_eq!(output.calls().len(), 1);
|
||||
}
|
||||
}
|
||||
@@ -7,7 +7,7 @@ use winnow::token::{literal, rest, take_until};
|
||||
use super::parameters::ToolSchemas;
|
||||
use super::utils::{MarkerScanState, parse_buffered_event, safe_text_len, take_until_marker};
|
||||
use super::{Result, ToolCallDelta, ToolParser, ToolParserOutput};
|
||||
use crate::{StructuralTagModel, Tool};
|
||||
use crate::tool::{StructuralTagModel, Tool};
|
||||
|
||||
const TOOL_CALLS_START: &str = "<tool_calls>";
|
||||
const TOOL_CALLS_END: &str = "</tool_calls>";
|
||||
@@ -79,7 +79,7 @@ impl HyV3ToolParser {
|
||||
fn apply_event(&mut self, event: HyV3Event, output: &mut ToolParserOutput) -> Result<()> {
|
||||
match event {
|
||||
HyV3Event::Text { len: consumed_len } => {
|
||||
output.normal_text.push_str(&self.buffer[..consumed_len]);
|
||||
output.push_text(&self.buffer[..consumed_len]);
|
||||
}
|
||||
HyV3Event::ToolBlockStart => {
|
||||
self.mode = HyV3Mode::ToolBlock {
|
||||
@@ -91,7 +91,7 @@ impl HyV3ToolParser {
|
||||
let arguments = serde_json::to_string(&arguments)
|
||||
.map_err(|error| parsing_failed!("failed to serialize arguments: {}", error))?;
|
||||
|
||||
output.calls.push(ToolCallDelta {
|
||||
output.push_call(ToolCallDelta {
|
||||
tool_index: self.emitted_tool_count,
|
||||
name: Some(name),
|
||||
arguments,
|
||||
@@ -133,7 +133,7 @@ impl ToolParser for HyV3ToolParser {
|
||||
fn finish(&mut self) -> Result<ToolParserOutput> {
|
||||
let mut output = ToolParserOutput::default();
|
||||
match self.mode {
|
||||
HyV3Mode::Text => output.normal_text.push_str(&self.buffer),
|
||||
HyV3Mode::Text => output.push_text(&self.buffer),
|
||||
HyV3Mode::ToolBlock { .. } => return Err(parsing_failed!("incomplete HY3 tool call")),
|
||||
HyV3Mode::Done => {}
|
||||
}
|
||||
@@ -249,8 +249,8 @@ mod tests {
|
||||
use thiserror_ext::AsReport;
|
||||
|
||||
use super::{HyV3ToolParser, ToolParser};
|
||||
use crate::test_utils::{collect_stream, split_by_chars, test_tools};
|
||||
use crate::{ToolParserOutput, ToolParserTestExt as _};
|
||||
use crate::tool::test_utils::{collect_stream, split_by_chars, test_tools};
|
||||
use crate::tool::{ToolParserOutput, ToolParserTestExt as _};
|
||||
|
||||
fn build_tool_call(function_name: &str, params: &[(&str, &str)]) -> String {
|
||||
let params = params
|
||||
@@ -266,7 +266,7 @@ mod tests {
|
||||
}
|
||||
|
||||
fn parsed_arguments(output: &ToolParserOutput, index: usize) -> Value {
|
||||
serde_json::from_str(&output.calls[index].arguments).unwrap()
|
||||
serde_json::from_str(&output.calls()[index].arguments).unwrap()
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -281,8 +281,8 @@ mod tests {
|
||||
let mut parser = HyV3ToolParser::new(&test_tools());
|
||||
let output = parser.parse_complete("This is a plain response.").unwrap();
|
||||
|
||||
assert_eq!(output.normal_text, "This is a plain response.");
|
||||
assert!(output.calls.is_empty());
|
||||
assert_eq!(output.normal_text(), "This is a plain response.");
|
||||
assert!(output.calls().is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -294,9 +294,9 @@ mod tests {
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(output.normal_text, "");
|
||||
assert_eq!(output.calls.len(), 1);
|
||||
assert_eq!(output.calls[0].name.as_deref(), Some("get_current_date"));
|
||||
assert_eq!(output.normal_text(), "");
|
||||
assert_eq!(output.calls().len(), 1);
|
||||
assert_eq!(output.calls()[0].name.as_deref(), Some("get_current_date"));
|
||||
assert_eq!(parsed_arguments(&output, 0), json!({}));
|
||||
}
|
||||
|
||||
@@ -309,7 +309,7 @@ mod tests {
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(output.calls[0].name.as_deref(), Some("get_current_date"));
|
||||
assert_eq!(output.calls()[0].name.as_deref(), Some("get_current_date"));
|
||||
assert_eq!(parsed_arguments(&output, 0), json!({}));
|
||||
}
|
||||
|
||||
@@ -354,8 +354,8 @@ mod tests {
|
||||
))
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(output.normal_text, "Checking.");
|
||||
assert_eq!(output.calls[0].name.as_deref(), Some("get_current_date"));
|
||||
assert_eq!(output.normal_text(), "Checking.");
|
||||
assert_eq!(output.calls()[0].name.as_deref(), Some("get_current_date"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -376,22 +376,25 @@ mod tests {
|
||||
|
||||
expect![[r#"
|
||||
ToolParserOutput {
|
||||
normal_text: "",
|
||||
calls: [
|
||||
ToolCallDelta {
|
||||
tool_index: 0,
|
||||
name: Some(
|
||||
"get_weather",
|
||||
),
|
||||
arguments: "{\"city\":\"Beijing\",\"date\":\"2026-03-30\"}",
|
||||
},
|
||||
ToolCallDelta {
|
||||
tool_index: 1,
|
||||
name: Some(
|
||||
"get_weather",
|
||||
),
|
||||
arguments: "{\"city\":\"Hangzhou\",\"date\":\"2026-03-30\"}",
|
||||
},
|
||||
events: [
|
||||
ToolCall(
|
||||
ToolCallDelta {
|
||||
tool_index: 0,
|
||||
name: Some(
|
||||
"get_weather",
|
||||
),
|
||||
arguments: "{\"city\":\"Beijing\",\"date\":\"2026-03-30\"}",
|
||||
},
|
||||
),
|
||||
ToolCall(
|
||||
ToolCallDelta {
|
||||
tool_index: 1,
|
||||
name: Some(
|
||||
"get_weather",
|
||||
),
|
||||
arguments: "{\"city\":\"Hangzhou\",\"date\":\"2026-03-30\"}",
|
||||
},
|
||||
),
|
||||
],
|
||||
}
|
||||
"#]]
|
||||
@@ -434,8 +437,8 @@ mod tests {
|
||||
output.append(parser.parse_chunk("response.").unwrap());
|
||||
output.append(parser.finish().unwrap());
|
||||
|
||||
assert_eq!(output.normal_text, "This is a plain response.");
|
||||
assert!(output.calls.is_empty());
|
||||
assert_eq!(output.normal_text(), "This is a plain response.");
|
||||
assert!(output.calls().is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -452,8 +455,8 @@ mod tests {
|
||||
|
||||
let output = collect_stream(&mut parser, &chunks);
|
||||
|
||||
assert_eq!(output.calls.len(), 1);
|
||||
assert_eq!(output.calls[0].name.as_deref(), Some("get_current_date"));
|
||||
assert_eq!(output.calls().len(), 1);
|
||||
assert_eq!(output.calls()[0].name.as_deref(), Some("get_current_date"));
|
||||
assert_eq!(parsed_arguments(&output, 0), json!({}));
|
||||
}
|
||||
|
||||
@@ -475,8 +478,8 @@ mod tests {
|
||||
|
||||
let output = collect_stream(&mut parser, &chunks);
|
||||
|
||||
assert_eq!(output.calls.len(), 1);
|
||||
assert_eq!(output.calls[0].name.as_deref(), Some("get_weather"));
|
||||
assert_eq!(output.calls().len(), 1);
|
||||
assert_eq!(output.calls()[0].name.as_deref(), Some("get_weather"));
|
||||
assert_eq!(
|
||||
parsed_arguments(&output, 0),
|
||||
json!({ "city": "Beijing", "date": "2026-03-30" })
|
||||
@@ -498,8 +501,8 @@ mod tests {
|
||||
|
||||
let output = collect_stream(&mut parser, &chunks);
|
||||
|
||||
assert_eq!(output.normal_text, "Checking.");
|
||||
assert_eq!(output.calls[0].name.as_deref(), Some("get_current_date"));
|
||||
assert_eq!(output.normal_text(), "Checking.");
|
||||
assert_eq!(output.calls()[0].name.as_deref(), Some("get_current_date"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -519,7 +522,7 @@ mod tests {
|
||||
|
||||
let output = collect_stream(&mut parser, &chunks);
|
||||
|
||||
assert_eq!(output.calls.len(), 2);
|
||||
assert_eq!(output.calls().len(), 2);
|
||||
assert_eq!(parsed_arguments(&output, 0)["city"], json!("Beijing"));
|
||||
assert_eq!(parsed_arguments(&output, 1)["city"], json!("Hangzhou"));
|
||||
}
|
||||
@@ -535,8 +538,8 @@ mod tests {
|
||||
|
||||
let output = collect_stream(&mut parser, &chunks);
|
||||
|
||||
assert_eq!(output.normal_text, "hello ");
|
||||
assert_eq!(output.calls.len(), 1);
|
||||
assert_eq!(output.normal_text(), "hello ");
|
||||
assert_eq!(output.calls().len(), 1);
|
||||
assert_eq!(parsed_arguments(&output, 0), json!({ "city": "Beijing" }));
|
||||
}
|
||||
|
||||
@@ -552,8 +555,8 @@ mod tests {
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(output.normal_text, "");
|
||||
assert!(output.calls.is_empty());
|
||||
assert_eq!(output.normal_text(), "");
|
||||
assert!(output.calls().is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
+79
-65
@@ -8,10 +8,10 @@ use super::{
|
||||
JsonToolCallConfig, JsonToolCallEvent, JsonToolCallWhitespace, JsonToolInput,
|
||||
tool_call_header_event,
|
||||
};
|
||||
use crate::utils::{
|
||||
use crate::tool::utils::{
|
||||
JsonObjectScanState, json_str, parse_buffered_event, safe_text_len, take_json_object,
|
||||
};
|
||||
use crate::{Result, Tool, ToolCallDelta, ToolParser, ToolParserOutput};
|
||||
use crate::tool::{Result, Tool, ToolCallDelta, ToolParser, ToolParserOutput};
|
||||
|
||||
const TOOL_CALL_START: &str = "<tool_call>";
|
||||
const TOOL_CALL_END: &str = "</tool_call>";
|
||||
@@ -86,14 +86,14 @@ impl Granite4ToolParser {
|
||||
/// Apply one parsed Granite 4 event to parser state and output.
|
||||
fn apply_event(&mut self, event: Granite4Event, output: &mut ToolParserOutput) -> Result<()> {
|
||||
match event {
|
||||
Granite4Event::Text { len } => output.normal_text.push_str(&self.buffer[..len]),
|
||||
Granite4Event::Text { len } => output.push_text(&self.buffer[..len]),
|
||||
Granite4Event::ToolCallStart => self.mode = Granite4Mode::Header,
|
||||
Granite4Event::ToolCallHeader { function_name } => {
|
||||
let tool_index = self.emitted_tool_count;
|
||||
self.emitted_tool_count += 1;
|
||||
self.active_tool_index = Some(tool_index);
|
||||
self.mode = Granite4Mode::Args { json_scan: None };
|
||||
output.calls.push(ToolCallDelta {
|
||||
output.push_call(ToolCallDelta {
|
||||
tool_index,
|
||||
name: Some(function_name),
|
||||
arguments: String::new(),
|
||||
@@ -125,7 +125,7 @@ impl Granite4ToolParser {
|
||||
"Granite4 arguments without an active tool call"
|
||||
));
|
||||
};
|
||||
output.calls.push(ToolCallDelta {
|
||||
output.push_call(ToolCallDelta {
|
||||
tool_index,
|
||||
name: None,
|
||||
arguments,
|
||||
@@ -165,7 +165,7 @@ impl ToolParser for Granite4ToolParser {
|
||||
fn finish(&mut self) -> Result<ToolParserOutput> {
|
||||
let mut output = ToolParserOutput::default();
|
||||
match &self.mode {
|
||||
Granite4Mode::Text => output.normal_text.push_str(&self.buffer),
|
||||
Granite4Mode::Text => output.push_text(&self.buffer),
|
||||
Granite4Mode::Header | Granite4Mode::Args { .. } | Granite4Mode::Close => {
|
||||
return Err(parsing_failed!("incomplete Granite4 tool call"));
|
||||
}
|
||||
@@ -279,16 +279,16 @@ mod tests {
|
||||
use thiserror_ext::AsReport;
|
||||
|
||||
use super::Granite4ToolParser;
|
||||
use crate::test_utils::{collect_stream, split_by_chars, test_tools};
|
||||
use crate::{ToolParser, ToolParserOutput, ToolParserTestExt as _};
|
||||
use crate::tool::test_utils::{collect_stream, split_by_chars, test_tools};
|
||||
use crate::tool::{ToolParser, ToolParserOutput, ToolParserTestExt as _};
|
||||
|
||||
#[test]
|
||||
fn granite4_parse_complete_without_tool_call_keeps_text() {
|
||||
let mut parser = Granite4ToolParser::new(&test_tools());
|
||||
let output = parser.parse_complete("Hello, world!").unwrap();
|
||||
|
||||
assert_eq!(output.normal_text, "Hello, world!");
|
||||
assert!(output.calls.is_empty());
|
||||
assert_eq!(output.normal_text(), "Hello, world!");
|
||||
assert!(output.calls().is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -300,10 +300,10 @@ mod tests {
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(output.normal_text, "");
|
||||
assert_eq!(output.calls.len(), 1);
|
||||
assert_eq!(output.calls[0].name.as_deref(), Some("get_weather"));
|
||||
assert_eq!(output.calls[0].arguments, r#"{"city":"Boston"}"#);
|
||||
assert_eq!(output.normal_text(), "");
|
||||
assert_eq!(output.calls().len(), 1);
|
||||
assert_eq!(output.calls()[0].name.as_deref(), Some("get_weather"));
|
||||
assert_eq!(output.calls()[0].arguments, r#"{"city":"Boston"}"#);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -317,9 +317,9 @@ mod tests {
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(output.calls.len(), 1);
|
||||
assert_eq!(output.calls[0].name.as_deref(), Some("get_weather"));
|
||||
assert_eq!(output.calls[0].arguments, r#"{"city":"Boston"}"#);
|
||||
assert_eq!(output.calls().len(), 1);
|
||||
assert_eq!(output.calls()[0].name.as_deref(), Some("get_weather"));
|
||||
assert_eq!(output.calls()[0].arguments, r#"{"city":"Boston"}"#);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -333,22 +333,28 @@ mod tests {
|
||||
|
||||
expect![[r#"
|
||||
ToolParserOutput {
|
||||
normal_text: "before middle after",
|
||||
calls: [
|
||||
ToolCallDelta {
|
||||
tool_index: 0,
|
||||
name: Some(
|
||||
"find_bbox",
|
||||
),
|
||||
arguments: "{\"x\":1}",
|
||||
},
|
||||
ToolCallDelta {
|
||||
tool_index: 1,
|
||||
name: Some(
|
||||
"get_weather",
|
||||
),
|
||||
arguments: "{\"city\":\"Boston\"}",
|
||||
},
|
||||
events: [
|
||||
Text(
|
||||
"before middle after",
|
||||
),
|
||||
ToolCall(
|
||||
ToolCallDelta {
|
||||
tool_index: 0,
|
||||
name: Some(
|
||||
"find_bbox",
|
||||
),
|
||||
arguments: "{\"x\":1}",
|
||||
},
|
||||
),
|
||||
ToolCall(
|
||||
ToolCallDelta {
|
||||
tool_index: 1,
|
||||
name: Some(
|
||||
"get_weather",
|
||||
),
|
||||
arguments: "{\"city\":\"Boston\"}",
|
||||
},
|
||||
),
|
||||
],
|
||||
}
|
||||
"#]]
|
||||
@@ -363,10 +369,10 @@ mod tests {
|
||||
|
||||
let output = collect_stream(&mut parser, &chunks);
|
||||
|
||||
assert_eq!(output.normal_text, "hello bye");
|
||||
assert_eq!(output.calls.len(), 1);
|
||||
assert_eq!(output.calls[0].name.as_deref(), Some("get_weather"));
|
||||
assert_eq!(output.calls[0].arguments, r#"{"city":"Tokyo"}"#);
|
||||
assert_eq!(output.normal_text(), "hello bye");
|
||||
assert_eq!(output.calls().len(), 1);
|
||||
assert_eq!(output.calls()[0].name.as_deref(), Some("get_weather"));
|
||||
assert_eq!(output.calls()[0].arguments, r#"{"city":"Tokyo"}"#);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -385,7 +391,7 @@ mod tests {
|
||||
for chunk in chunks {
|
||||
let next = parser.parse_chunk(chunk).unwrap();
|
||||
observed_arguments.extend(
|
||||
next.calls
|
||||
next.calls()
|
||||
.iter()
|
||||
.filter(|call| call.name.is_none())
|
||||
.map(|call| call.arguments.clone()),
|
||||
@@ -396,7 +402,7 @@ mod tests {
|
||||
|
||||
assert_eq!(observed_arguments, [r#"{"city":"#, r#""Beijing""#, r#"}"#]);
|
||||
assert_eq!(
|
||||
output.coalesce_calls().calls[0].arguments,
|
||||
output.coalesce().calls()[0].arguments,
|
||||
r#"{"city":"Beijing"}"#
|
||||
);
|
||||
}
|
||||
@@ -409,9 +415,9 @@ mod tests {
|
||||
|
||||
let output = collect_stream(&mut parser, &chunks);
|
||||
|
||||
assert_eq!(output.calls.len(), 1);
|
||||
assert_eq!(output.calls[0].name.as_deref(), Some("f"));
|
||||
assert_eq!(output.calls[0].arguments, r#"{"a":1}"#);
|
||||
assert_eq!(output.calls().len(), 1);
|
||||
assert_eq!(output.calls()[0].name.as_deref(), Some("f"));
|
||||
assert_eq!(output.calls()[0].arguments, r#"{"a":1}"#);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -435,29 +441,37 @@ mod tests {
|
||||
|
||||
expect![[r#"
|
||||
ToolParserOutput {
|
||||
normal_text: "Here goes the bbox call: \n Now the stock price call: \n Now another bbox call: \n See? I'm a helpful assistant.",
|
||||
calls: [
|
||||
ToolCallDelta {
|
||||
tool_index: 0,
|
||||
name: Some(
|
||||
"find_bbox",
|
||||
),
|
||||
arguments: "{\"coordinates\": [[23.54, 43.1], [-12.2, 54.3], [4, 5]], \"coordinate_type\": \"latlong\"}",
|
||||
},
|
||||
ToolCallDelta {
|
||||
tool_index: 1,
|
||||
name: Some(
|
||||
"get_stock_price",
|
||||
),
|
||||
arguments: "{\"symbol\": \"AAPL\", \"start_date\": \"2021-01-01\", \"end_date\": \"2021-12-31\"}",
|
||||
},
|
||||
ToolCallDelta {
|
||||
tool_index: 2,
|
||||
name: Some(
|
||||
"find_bbox",
|
||||
),
|
||||
arguments: "{\"coordinates\": [[23.54, 43.1], [-12.2, 54.3], [4, 5]], \"coordinate_type\": \"latlong\"}",
|
||||
},
|
||||
events: [
|
||||
Text(
|
||||
"Here goes the bbox call: \n Now the stock price call: \n Now another bbox call: \n See? I'm a helpful assistant.",
|
||||
),
|
||||
ToolCall(
|
||||
ToolCallDelta {
|
||||
tool_index: 0,
|
||||
name: Some(
|
||||
"find_bbox",
|
||||
),
|
||||
arguments: "{\"coordinates\": [[23.54, 43.1], [-12.2, 54.3], [4, 5]], \"coordinate_type\": \"latlong\"}",
|
||||
},
|
||||
),
|
||||
ToolCall(
|
||||
ToolCallDelta {
|
||||
tool_index: 1,
|
||||
name: Some(
|
||||
"get_stock_price",
|
||||
),
|
||||
arguments: "{\"symbol\": \"AAPL\", \"start_date\": \"2021-01-01\", \"end_date\": \"2021-12-31\"}",
|
||||
},
|
||||
),
|
||||
ToolCall(
|
||||
ToolCallDelta {
|
||||
tool_index: 2,
|
||||
name: Some(
|
||||
"find_bbox",
|
||||
),
|
||||
arguments: "{\"coordinates\": [[23.54, 43.1], [-12.2, 54.3], [4, 5]], \"coordinate_type\": \"latlong\"}",
|
||||
},
|
||||
),
|
||||
],
|
||||
}
|
||||
"#]].assert_debug_eq(&output);
|
||||
@@ -1,5 +1,5 @@
|
||||
use super::{JsonToolCallConfig, JsonToolCallParser, JsonToolCallWhitespace};
|
||||
use crate::{Result, StructuralTagModel, Tool, ToolParser, ToolParserOutput};
|
||||
use crate::tool::{Result, StructuralTagModel, Tool, ToolParser, ToolParserOutput};
|
||||
|
||||
const HERMES_CONFIG: JsonToolCallConfig = JsonToolCallConfig {
|
||||
parser_name: "Hermes",
|
||||
@@ -68,8 +68,8 @@ mod tests {
|
||||
use thiserror_ext::AsReport;
|
||||
|
||||
use super::HermesToolParser;
|
||||
use crate::test_utils::{collect_stream, split_by_chars, test_tools};
|
||||
use crate::{ToolParser, ToolParserOutput, ToolParserTestExt as _};
|
||||
use crate::tool::test_utils::{collect_stream, split_by_chars, test_tools};
|
||||
use crate::tool::{ToolParser, ToolParserOutput, ToolParserTestExt as _};
|
||||
|
||||
fn build_tool_call(function_name: &str, arguments: &str) -> String {
|
||||
format!(r#"<tool_call>{{"name":"{function_name}","arguments":{arguments}}}</tool_call>"#)
|
||||
@@ -80,8 +80,8 @@ mod tests {
|
||||
let mut parser = HermesToolParser::new(&test_tools());
|
||||
let output = parser.parse_complete("Hello, world!").unwrap();
|
||||
|
||||
assert_eq!(output.normal_text, "Hello, world!");
|
||||
assert!(output.calls.is_empty());
|
||||
assert_eq!(output.normal_text(), "Hello, world!");
|
||||
assert!(output.calls().is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -95,11 +95,11 @@ mod tests {
|
||||
))
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(output.normal_text, "Let me check.\n");
|
||||
assert_eq!(output.calls.len(), 1);
|
||||
assert_eq!(output.calls[0].tool_index, 0);
|
||||
assert_eq!(output.calls[0].name.as_deref(), Some("get_weather"));
|
||||
assert_eq!(output.calls[0].arguments, arguments);
|
||||
assert_eq!(output.normal_text(), "Let me check.\n");
|
||||
assert_eq!(output.calls().len(), 1);
|
||||
assert_eq!(output.calls()[0].tool_index, 0);
|
||||
assert_eq!(output.calls()[0].name.as_deref(), Some("get_weather"));
|
||||
assert_eq!(output.calls()[0].arguments, arguments);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -112,8 +112,8 @@ mod tests {
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(output.calls.len(), 1);
|
||||
assert_eq!(output.calls[0].name.as_deref(), Some("get_weather"));
|
||||
assert_eq!(output.calls().len(), 1);
|
||||
assert_eq!(output.calls()[0].name.as_deref(), Some("get_weather"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -122,7 +122,7 @@ mod tests {
|
||||
let arguments = r#"{"location":"Tokyo",}"#;
|
||||
let output = parser.parse_complete(&build_tool_call("get_weather", arguments)).unwrap();
|
||||
|
||||
assert_eq!(output.calls[0].arguments, arguments);
|
||||
assert_eq!(output.calls()[0].arguments, arguments);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -142,7 +142,7 @@ mod tests {
|
||||
for chunk in chunks {
|
||||
let next = parser.parse_chunk(chunk).unwrap();
|
||||
observed_arguments.extend(
|
||||
next.calls
|
||||
next.calls()
|
||||
.iter()
|
||||
.filter(|call| call.name.is_none())
|
||||
.map(|call| call.arguments.clone()),
|
||||
@@ -152,9 +152,9 @@ mod tests {
|
||||
output.append(parser.finish().unwrap());
|
||||
|
||||
assert_eq!(observed_arguments, ["{\"location\":", "\"Beijing\"", "}"]);
|
||||
assert_eq!(output.normal_text, "preface suffix");
|
||||
assert_eq!(output.normal_text(), "preface suffix");
|
||||
assert_eq!(
|
||||
output.coalesce_calls().calls[0].arguments,
|
||||
output.coalesce().calls()[0].arguments,
|
||||
r#"{"location":"Beijing"}"#
|
||||
);
|
||||
}
|
||||
@@ -170,9 +170,9 @@ mod tests {
|
||||
|
||||
let output = collect_stream(&mut parser, &chunks);
|
||||
|
||||
assert_eq!(output.normal_text, "hello ");
|
||||
assert_eq!(output.calls.len(), 1);
|
||||
assert_eq!(output.calls[0].arguments, r#"{"location":"Tokyo"}"#);
|
||||
assert_eq!(output.normal_text(), "hello ");
|
||||
assert_eq!(output.calls().len(), 1);
|
||||
assert_eq!(output.calls()[0].arguments, r#"{"location":"Tokyo"}"#);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -189,22 +189,25 @@ mod tests {
|
||||
|
||||
expect![[r#"
|
||||
ToolParserOutput {
|
||||
normal_text: "",
|
||||
calls: [
|
||||
ToolCallDelta {
|
||||
tool_index: 0,
|
||||
name: Some(
|
||||
"get_weather",
|
||||
),
|
||||
arguments: "{\"location\":\"Shanghai\"}",
|
||||
},
|
||||
ToolCallDelta {
|
||||
tool_index: 1,
|
||||
name: Some(
|
||||
"add",
|
||||
),
|
||||
arguments: "{\"x\":1,\"y\":2}",
|
||||
},
|
||||
events: [
|
||||
ToolCall(
|
||||
ToolCallDelta {
|
||||
tool_index: 0,
|
||||
name: Some(
|
||||
"get_weather",
|
||||
),
|
||||
arguments: "{\"location\":\"Shanghai\"}",
|
||||
},
|
||||
),
|
||||
ToolCall(
|
||||
ToolCallDelta {
|
||||
tool_index: 1,
|
||||
name: Some(
|
||||
"add",
|
||||
),
|
||||
arguments: "{\"x\":1,\"y\":2}",
|
||||
},
|
||||
),
|
||||
],
|
||||
}
|
||||
"#]]
|
||||
+44
-41
@@ -1,5 +1,5 @@
|
||||
use super::{JsonToolCallConfig, JsonToolCallParser, JsonToolCallWhitespace};
|
||||
use crate::{Result, Tool, ToolParser, ToolParserOutput};
|
||||
use crate::tool::{Result, Tool, ToolParser, ToolParserOutput};
|
||||
|
||||
const INTERNLM2_CONFIG: JsonToolCallConfig = JsonToolCallConfig {
|
||||
parser_name: "InternLM2",
|
||||
@@ -123,8 +123,8 @@ mod tests {
|
||||
use thiserror_ext::AsReport;
|
||||
|
||||
use super::Internlm2ToolParser;
|
||||
use crate::test_utils::{collect_stream, split_by_chars, test_tools};
|
||||
use crate::{ToolParser, ToolParserOutput, ToolParserTestExt as _};
|
||||
use crate::tool::test_utils::{collect_stream, split_by_chars, test_tools};
|
||||
use crate::tool::{ToolParser, ToolParserOutput, ToolParserTestExt as _};
|
||||
|
||||
const ACTION_START: &str = "<|action_start|><|plugin|>";
|
||||
const ACTION_END: &str = "<|action_end|>";
|
||||
@@ -140,8 +140,8 @@ mod tests {
|
||||
let mut parser = Internlm2ToolParser::new(&test_tools());
|
||||
let result = parser.parse_complete("Hello, world!").unwrap();
|
||||
|
||||
assert_eq!(result.normal_text, "Hello, world!");
|
||||
assert!(result.calls.is_empty());
|
||||
assert_eq!(result.normal_text(), "Hello, world!");
|
||||
assert!(result.calls().is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -155,11 +155,11 @@ mod tests {
|
||||
))
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(result.normal_text, "Let me check.\n");
|
||||
assert_eq!(result.calls.len(), 1);
|
||||
assert_eq!(result.calls[0].tool_index, 0);
|
||||
assert_eq!(result.calls[0].name.as_deref(), Some("get_weather"));
|
||||
assert_eq!(result.calls[0].arguments, arguments);
|
||||
assert_eq!(result.normal_text(), "Let me check.\n");
|
||||
assert_eq!(result.calls().len(), 1);
|
||||
assert_eq!(result.calls()[0].tool_index, 0);
|
||||
assert_eq!(result.calls()[0].name.as_deref(), Some("get_weather"));
|
||||
assert_eq!(result.calls()[0].arguments, arguments);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -170,9 +170,9 @@ mod tests {
|
||||
.parse_complete(&build_tool_call("get_weather", "arguments", arguments))
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(result.calls.len(), 1);
|
||||
assert_eq!(result.calls[0].name.as_deref(), Some("get_weather"));
|
||||
assert_eq!(result.calls[0].arguments, arguments);
|
||||
assert_eq!(result.calls().len(), 1);
|
||||
assert_eq!(result.calls()[0].name.as_deref(), Some("get_weather"));
|
||||
assert_eq!(result.calls()[0].arguments, arguments);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -185,8 +185,8 @@ mod tests {
|
||||
))
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(result.calls.len(), 1);
|
||||
assert_eq!(result.calls[0].name.as_deref(), Some("get_weather"));
|
||||
assert_eq!(result.calls().len(), 1);
|
||||
assert_eq!(result.calls()[0].name.as_deref(), Some("get_weather"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -197,7 +197,7 @@ mod tests {
|
||||
.parse_complete(&build_tool_call("get_weather", "parameters", arguments))
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(result.calls[0].arguments, arguments);
|
||||
assert_eq!(result.calls()[0].arguments, arguments);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -218,7 +218,7 @@ mod tests {
|
||||
for chunk in chunks {
|
||||
let next = parser.parse_chunk(chunk).unwrap();
|
||||
observed_arguments.extend(
|
||||
next.calls
|
||||
next.calls()
|
||||
.iter()
|
||||
.filter(|call| call.name.is_none())
|
||||
.map(|call| call.arguments.clone()),
|
||||
@@ -231,9 +231,9 @@ mod tests {
|
||||
observed_arguments,
|
||||
[r#"{"location":"#, r#""Beijing""#, r#"}"#]
|
||||
);
|
||||
assert_eq!(result.normal_text, "preface suffix");
|
||||
assert_eq!(result.normal_text(), "preface suffix");
|
||||
assert_eq!(
|
||||
result.coalesce_calls().calls[0].arguments,
|
||||
result.coalesce().calls()[0].arguments,
|
||||
r#"{"location":"Beijing"}"#
|
||||
);
|
||||
}
|
||||
@@ -249,9 +249,9 @@ mod tests {
|
||||
|
||||
let result = collect_stream(&mut parser, &chunks);
|
||||
|
||||
assert_eq!(result.normal_text, "hello ");
|
||||
assert_eq!(result.calls.len(), 1);
|
||||
assert_eq!(result.calls[0].arguments, r#"{"location":"Tokyo"}"#);
|
||||
assert_eq!(result.normal_text(), "hello ");
|
||||
assert_eq!(result.calls().len(), 1);
|
||||
assert_eq!(result.calls()[0].arguments, r#"{"location":"Tokyo"}"#);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -268,22 +268,25 @@ mod tests {
|
||||
|
||||
expect![[r#"
|
||||
ToolParserOutput {
|
||||
normal_text: "",
|
||||
calls: [
|
||||
ToolCallDelta {
|
||||
tool_index: 0,
|
||||
name: Some(
|
||||
"get_weather",
|
||||
),
|
||||
arguments: "{\"location\":\"Shanghai\"}",
|
||||
},
|
||||
ToolCallDelta {
|
||||
tool_index: 1,
|
||||
name: Some(
|
||||
"add",
|
||||
),
|
||||
arguments: "{\"x\":1,\"y\":2}",
|
||||
},
|
||||
events: [
|
||||
ToolCall(
|
||||
ToolCallDelta {
|
||||
tool_index: 0,
|
||||
name: Some(
|
||||
"get_weather",
|
||||
),
|
||||
arguments: "{\"location\":\"Shanghai\"}",
|
||||
},
|
||||
),
|
||||
ToolCall(
|
||||
ToolCallDelta {
|
||||
tool_index: 1,
|
||||
name: Some(
|
||||
"add",
|
||||
),
|
||||
arguments: "{\"x\":1,\"y\":2}",
|
||||
},
|
||||
),
|
||||
],
|
||||
}
|
||||
"#]]
|
||||
@@ -298,8 +301,8 @@ mod tests {
|
||||
|
||||
let result = parser.parse_complete(&input).unwrap();
|
||||
|
||||
assert_eq!(result.calls.len(), 1);
|
||||
assert_eq!(result.calls[0].arguments, arguments);
|
||||
assert_eq!(result.calls().len(), 1);
|
||||
assert_eq!(result.calls()[0].arguments, arguments);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -313,7 +316,7 @@ mod tests {
|
||||
let error = parser.finish().unwrap_err();
|
||||
|
||||
assert_eq!(
|
||||
pre_finish.calls[0].name.as_deref(),
|
||||
pre_finish.calls()[0].name.as_deref(),
|
||||
Some("get_weather"),
|
||||
"name delta is still emitted from parse_chunk() before truncation",
|
||||
);
|
||||
@@ -8,8 +8,8 @@ use super::{
|
||||
JsonToolCallConfig, JsonToolCallEvent, JsonToolCallWhitespace, JsonToolInput,
|
||||
argument_delta_event, tool_call_header_event,
|
||||
};
|
||||
use crate::utils::{JsonObjectScanState, parse_buffered_event};
|
||||
use crate::{Result, StructuralTagModel, Tool, ToolCallDelta, ToolParser, ToolParserOutput};
|
||||
use crate::tool::utils::{JsonObjectScanState, parse_buffered_event};
|
||||
use crate::tool::{Result, StructuralTagModel, Tool, ToolCallDelta, ToolParser, ToolParserOutput};
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
enum LlamaJsonMode {
|
||||
@@ -87,7 +87,7 @@ impl Llama3JsonToolParser {
|
||||
self.mode = LlamaJsonMode::Arguments {
|
||||
json_scan: JsonObjectScanState::default(),
|
||||
};
|
||||
output.calls.push(ToolCallDelta {
|
||||
output.push_call(ToolCallDelta {
|
||||
tool_index,
|
||||
name: Some(function_name),
|
||||
arguments: String::new(),
|
||||
@@ -99,7 +99,7 @@ impl Llama3JsonToolParser {
|
||||
"Llama JSON arguments without an active tool call"
|
||||
));
|
||||
};
|
||||
output.calls.push(ToolCallDelta {
|
||||
output.push_call(ToolCallDelta {
|
||||
tool_index,
|
||||
name: None,
|
||||
arguments: self.buffer[..consumed_len].to_string(),
|
||||
@@ -145,7 +145,7 @@ impl ToolParser for Llama3JsonToolParser {
|
||||
}
|
||||
|
||||
if matches!(self.mode, LlamaJsonMode::Passthrough) {
|
||||
output.normal_text.push_str(&self.buffer);
|
||||
output.push_text(&self.buffer);
|
||||
self.buffer.clear();
|
||||
return Ok(());
|
||||
}
|
||||
@@ -164,7 +164,7 @@ impl ToolParser for Llama3JsonToolParser {
|
||||
let mut output = ToolParserOutput::default();
|
||||
match &self.mode {
|
||||
LlamaJsonMode::Start | LlamaJsonMode::Passthrough => {
|
||||
output.normal_text.push_str(&self.buffer);
|
||||
output.push_text(&self.buffer);
|
||||
}
|
||||
LlamaJsonMode::AfterCall if self.buffer.trim().is_empty() => {}
|
||||
LlamaJsonMode::Header | LlamaJsonMode::Arguments { .. } => {
|
||||
@@ -256,8 +256,8 @@ mod tests {
|
||||
use thiserror_ext::AsReport;
|
||||
|
||||
use super::Llama3JsonToolParser;
|
||||
use crate::test_utils::{collect_stream, split_by_chars, test_tools};
|
||||
use crate::{ToolParser, ToolParserOutput, ToolParserTestExt as _};
|
||||
use crate::tool::test_utils::{collect_stream, split_by_chars, test_tools};
|
||||
use crate::tool::{ToolParser, ToolParserOutput, ToolParserTestExt as _};
|
||||
|
||||
fn build_tool_call(function_name: &str, parameters: &str) -> String {
|
||||
format!(r#"{{"name":"{function_name}","parameters":{parameters}}}"#)
|
||||
@@ -268,8 +268,8 @@ mod tests {
|
||||
let mut parser = Llama3JsonToolParser::new(&test_tools());
|
||||
let output = parser.parse_complete("Hello, world!").unwrap();
|
||||
|
||||
assert_eq!(output.normal_text, "Hello, world!");
|
||||
assert!(output.calls.is_empty());
|
||||
assert_eq!(output.normal_text(), "Hello, world!");
|
||||
assert!(output.calls().is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -284,10 +284,10 @@ mod tests {
|
||||
output.append(parser.finish().unwrap());
|
||||
|
||||
assert_eq!(
|
||||
output.normal_text,
|
||||
output.normal_text(),
|
||||
r#"plain text first {"name":"get_weather","parameters":{"location":"Tokyo"}}"#
|
||||
);
|
||||
assert!(output.calls.is_empty());
|
||||
assert!(output.calls().is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -299,8 +299,8 @@ mod tests {
|
||||
);
|
||||
let output = parser.parse_complete(&input).unwrap();
|
||||
|
||||
assert_eq!(output.normal_text, input);
|
||||
assert!(output.calls.is_empty());
|
||||
assert_eq!(output.normal_text(), input);
|
||||
assert!(output.calls().is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -312,8 +312,8 @@ mod tests {
|
||||
);
|
||||
let output = parser.parse_complete(&input).unwrap();
|
||||
|
||||
assert_eq!(output.normal_text, input);
|
||||
assert!(output.calls.is_empty());
|
||||
assert_eq!(output.normal_text(), input);
|
||||
assert!(output.calls().is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -322,10 +322,10 @@ mod tests {
|
||||
let arguments = r#"{ "location": "Tokyo", "days": 3 }"#;
|
||||
let output = parser.parse_complete(&build_tool_call("get_weather", arguments)).unwrap();
|
||||
|
||||
assert_eq!(output.calls.len(), 1);
|
||||
assert_eq!(output.calls[0].tool_index, 0);
|
||||
assert_eq!(output.calls[0].name.as_deref(), Some("get_weather"));
|
||||
assert_eq!(output.calls[0].arguments, arguments);
|
||||
assert_eq!(output.calls().len(), 1);
|
||||
assert_eq!(output.calls()[0].tool_index, 0);
|
||||
assert_eq!(output.calls()[0].name.as_deref(), Some("get_weather"));
|
||||
assert_eq!(output.calls()[0].arguments, arguments);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -353,22 +353,25 @@ mod tests {
|
||||
|
||||
expect![[r#"
|
||||
ToolParserOutput {
|
||||
normal_text: "",
|
||||
calls: [
|
||||
ToolCallDelta {
|
||||
tool_index: 0,
|
||||
name: Some(
|
||||
"get_weather",
|
||||
),
|
||||
arguments: "{\"location\":\"Shanghai\"}",
|
||||
},
|
||||
ToolCallDelta {
|
||||
tool_index: 1,
|
||||
name: Some(
|
||||
"add",
|
||||
),
|
||||
arguments: "{\"x\":1,\"y\":2}",
|
||||
},
|
||||
events: [
|
||||
ToolCall(
|
||||
ToolCallDelta {
|
||||
tool_index: 0,
|
||||
name: Some(
|
||||
"get_weather",
|
||||
),
|
||||
arguments: "{\"location\":\"Shanghai\"}",
|
||||
},
|
||||
),
|
||||
ToolCall(
|
||||
ToolCallDelta {
|
||||
tool_index: 1,
|
||||
name: Some(
|
||||
"add",
|
||||
),
|
||||
arguments: "{\"x\":1,\"y\":2}",
|
||||
},
|
||||
),
|
||||
],
|
||||
}
|
||||
"#]]
|
||||
@@ -390,7 +393,7 @@ mod tests {
|
||||
for chunk in chunks {
|
||||
let next = parser.parse_chunk(chunk).unwrap();
|
||||
observed_arguments.extend(
|
||||
next.calls
|
||||
next.calls()
|
||||
.iter()
|
||||
.filter(|call| call.name.is_none())
|
||||
.map(|call| call.arguments.clone()),
|
||||
@@ -401,7 +404,7 @@ mod tests {
|
||||
|
||||
assert_eq!(observed_arguments, ["{\"location\":", "\"Beijing\"", "}"]);
|
||||
assert_eq!(
|
||||
output.coalesce_calls().calls[0].arguments,
|
||||
output.coalesce().calls()[0].arguments,
|
||||
r#"{"location":"Beijing"}"#
|
||||
);
|
||||
}
|
||||
@@ -418,14 +421,14 @@ mod tests {
|
||||
|
||||
let output = collect_stream(&mut parser, &chunks);
|
||||
|
||||
assert_eq!(output.normal_text, "");
|
||||
assert_eq!(output.calls.len(), 2);
|
||||
assert_eq!(output.normal_text(), "");
|
||||
assert_eq!(output.calls().len(), 2);
|
||||
assert_eq!(
|
||||
output.calls[0].arguments,
|
||||
output.calls()[0].arguments,
|
||||
r#"{"location":"Dallas","state":"TX"}"#
|
||||
);
|
||||
assert_eq!(output.calls[1].name.as_deref(), Some("add"));
|
||||
assert_eq!(output.calls[1].arguments, r#"{"x":4,"y":5}"#);
|
||||
assert_eq!(output.calls()[1].name.as_deref(), Some("add"));
|
||||
assert_eq!(output.calls()[1].arguments, r#"{"x":4,"y":5}"#);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -437,7 +440,7 @@ mod tests {
|
||||
}"#;
|
||||
let output = parser.parse_complete(&build_tool_call("convert", arguments)).unwrap();
|
||||
|
||||
assert_eq!(output.calls[0].arguments, arguments);
|
||||
assert_eq!(output.calls()[0].arguments, arguments);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -450,8 +453,8 @@ mod tests {
|
||||
))
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(output.normal_text, "");
|
||||
assert_eq!(output.calls.len(), 1);
|
||||
assert_eq!(output.normal_text(), "");
|
||||
assert_eq!(output.calls().len(), 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
+41
-35
@@ -1,5 +1,5 @@
|
||||
use super::{JsonToolCallConfig, JsonToolCallParser, JsonToolCallWhitespace};
|
||||
use crate::{Result, Tool, ToolParser, ToolParserOutput};
|
||||
use crate::tool::{Result, Tool, ToolParser, ToolParserOutput};
|
||||
|
||||
const MISTRAL_CONFIG: JsonToolCallConfig = JsonToolCallConfig {
|
||||
parser_name: "Mistral",
|
||||
@@ -61,8 +61,8 @@ mod tests {
|
||||
use thiserror_ext::AsReport;
|
||||
|
||||
use super::MistralToolParser;
|
||||
use crate::test_utils::{collect_stream, split_by_chars, test_tools};
|
||||
use crate::{ToolParser, ToolParserOutput, ToolParserTestExt as _};
|
||||
use crate::tool::test_utils::{collect_stream, split_by_chars, test_tools};
|
||||
use crate::tool::{ToolParser, ToolParserOutput, ToolParserTestExt as _};
|
||||
|
||||
fn build_tool_call(function_name: &str, arguments: &str) -> String {
|
||||
format!(r#"{{"name":"{function_name}","arguments":{arguments}}}"#)
|
||||
@@ -77,8 +77,8 @@ mod tests {
|
||||
let mut parser = MistralToolParser::new(&test_tools());
|
||||
let output = parser.parse_complete("Hello, world!").unwrap();
|
||||
|
||||
assert_eq!(output.normal_text, "Hello, world!");
|
||||
assert!(output.calls.is_empty());
|
||||
assert_eq!(output.normal_text(), "Hello, world!");
|
||||
assert!(output.calls().is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -92,11 +92,11 @@ mod tests {
|
||||
))
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(output.normal_text, "Let me check.\n");
|
||||
assert_eq!(output.calls.len(), 1);
|
||||
assert_eq!(output.calls[0].tool_index, 0);
|
||||
assert_eq!(output.calls[0].name.as_deref(), Some("get_weather"));
|
||||
assert_eq!(output.calls[0].arguments, arguments);
|
||||
assert_eq!(output.normal_text(), "Let me check.\n");
|
||||
assert_eq!(output.calls().len(), 1);
|
||||
assert_eq!(output.calls()[0].tool_index, 0);
|
||||
assert_eq!(output.calls()[0].name.as_deref(), Some("get_weather"));
|
||||
assert_eq!(output.calls()[0].arguments, arguments);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -115,22 +115,28 @@ mod tests {
|
||||
|
||||
expect![[r#"
|
||||
ToolParserOutput {
|
||||
normal_text: "I'll help.\n",
|
||||
calls: [
|
||||
ToolCallDelta {
|
||||
tool_index: 0,
|
||||
name: Some(
|
||||
"get_weather",
|
||||
),
|
||||
arguments: "{\"city\": \"Tokyo\", \"units\": \"celsius\"}",
|
||||
},
|
||||
ToolCallDelta {
|
||||
tool_index: 1,
|
||||
name: Some(
|
||||
"add",
|
||||
),
|
||||
arguments: "{\"x\": 1, \"y\": 2}",
|
||||
},
|
||||
events: [
|
||||
Text(
|
||||
"I'll help.\n",
|
||||
),
|
||||
ToolCall(
|
||||
ToolCallDelta {
|
||||
tool_index: 0,
|
||||
name: Some(
|
||||
"get_weather",
|
||||
),
|
||||
arguments: "{\"city\": \"Tokyo\", \"units\": \"celsius\"}",
|
||||
},
|
||||
),
|
||||
ToolCall(
|
||||
ToolCallDelta {
|
||||
tool_index: 1,
|
||||
name: Some(
|
||||
"add",
|
||||
),
|
||||
arguments: "{\"x\": 1, \"y\": 2}",
|
||||
},
|
||||
),
|
||||
],
|
||||
}
|
||||
"#]]
|
||||
@@ -148,7 +154,7 @@ mod tests {
|
||||
)]))
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(output.calls[0].arguments, arguments);
|
||||
assert_eq!(output.calls()[0].arguments, arguments);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -168,7 +174,7 @@ mod tests {
|
||||
for chunk in chunks {
|
||||
let next = parser.parse_chunk(chunk).unwrap();
|
||||
observed_arguments.extend(
|
||||
next.calls
|
||||
next.calls()
|
||||
.iter()
|
||||
.filter(|call| call.name.is_none())
|
||||
.map(|call| call.arguments.clone()),
|
||||
@@ -178,9 +184,9 @@ mod tests {
|
||||
output.append(parser.finish().unwrap());
|
||||
|
||||
assert_eq!(observed_arguments, ["{\"location\":", "\"Beijing\"", "}"]);
|
||||
assert_eq!(output.normal_text, "preface suffix");
|
||||
assert_eq!(output.normal_text(), "preface suffix");
|
||||
assert_eq!(
|
||||
output.coalesce_calls().calls[0].arguments,
|
||||
output.coalesce().calls()[0].arguments,
|
||||
r#"{"location":"Beijing"}"#
|
||||
);
|
||||
}
|
||||
@@ -196,9 +202,9 @@ mod tests {
|
||||
|
||||
let output = collect_stream(&mut parser, &chunks);
|
||||
|
||||
assert_eq!(output.normal_text, "hello ");
|
||||
assert_eq!(output.calls.len(), 1);
|
||||
assert_eq!(output.calls[0].arguments, r#"{"location":"Tokyo"}"#);
|
||||
assert_eq!(output.normal_text(), "hello ");
|
||||
assert_eq!(output.calls().len(), 1);
|
||||
assert_eq!(output.calls()[0].arguments, r#"{"location":"Tokyo"}"#);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -209,8 +215,8 @@ mod tests {
|
||||
.parse_complete(&build_tool_calls(&[build_tool_call("echo", arguments)]))
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(output.calls.len(), 1);
|
||||
assert_eq!(output.calls[0].arguments, arguments);
|
||||
assert_eq!(output.calls().len(), 1);
|
||||
assert_eq!(output.calls()[0].arguments, arguments);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -106,7 +106,7 @@ impl JsonToolCallParser {
|
||||
fn finish(&mut self) -> Result<ToolParserOutput> {
|
||||
let mut output = ToolParserOutput::default();
|
||||
match &self.mode {
|
||||
JsonToolCallMode::Text => output.normal_text.push_str(&self.buffer),
|
||||
JsonToolCallMode::Text => output.push_text(&self.buffer),
|
||||
JsonToolCallMode::Header | JsonToolCallMode::Arguments { .. } => {
|
||||
return Err(parsing_failed!(
|
||||
"incomplete {} tool call",
|
||||
@@ -126,7 +126,7 @@ impl JsonToolCallParser {
|
||||
) -> Result<()> {
|
||||
match event {
|
||||
JsonToolCallEvent::Text { len: consumed_len } => {
|
||||
output.normal_text.push_str(&self.buffer[..consumed_len]);
|
||||
output.push_text(&self.buffer[..consumed_len]);
|
||||
}
|
||||
JsonToolCallEvent::ToolCallStart => self.mode = JsonToolCallMode::Header,
|
||||
JsonToolCallEvent::ToolCallHeader { function_name } => {
|
||||
@@ -136,7 +136,7 @@ impl JsonToolCallParser {
|
||||
self.mode = JsonToolCallMode::Arguments {
|
||||
json_scan: JsonObjectScanState::default(),
|
||||
};
|
||||
output.calls.push(ToolCallDelta {
|
||||
output.push_call(ToolCallDelta {
|
||||
tool_index,
|
||||
name: Some(function_name),
|
||||
arguments: String::new(),
|
||||
@@ -149,7 +149,7 @@ impl JsonToolCallParser {
|
||||
self.config.parser_name
|
||||
));
|
||||
};
|
||||
output.calls.push(ToolCallDelta {
|
||||
output.push_call(ToolCallDelta {
|
||||
tool_index,
|
||||
name: None,
|
||||
arguments: self.buffer[..consumed_len].to_string(),
|
||||
@@ -374,7 +374,7 @@ mod tests {
|
||||
use expect_test::expect;
|
||||
|
||||
use super::{JsonToolCallConfig, JsonToolCallParser, JsonToolCallWhitespace};
|
||||
use crate::ToolParserOutput;
|
||||
use crate::tool::ToolParserOutput;
|
||||
|
||||
const DELIMITED_CONFIG: JsonToolCallConfig = JsonToolCallConfig {
|
||||
parser_name: "Delimited JSON",
|
||||
@@ -400,7 +400,7 @@ mod tests {
|
||||
parser.parse_into(chunk, &mut output).unwrap();
|
||||
}
|
||||
output.append(parser.finish().unwrap());
|
||||
output.coalesce_calls()
|
||||
output.coalesce()
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -415,22 +415,25 @@ mod tests {
|
||||
|
||||
expect![[r#"
|
||||
ToolParserOutput {
|
||||
normal_text: "",
|
||||
calls: [
|
||||
ToolCallDelta {
|
||||
tool_index: 0,
|
||||
name: Some(
|
||||
"get_weather",
|
||||
),
|
||||
arguments: "{\"location\":\"Shanghai\"}",
|
||||
},
|
||||
ToolCallDelta {
|
||||
tool_index: 1,
|
||||
name: Some(
|
||||
"add",
|
||||
),
|
||||
arguments: "{\"x\":1,\"y\":2}",
|
||||
},
|
||||
events: [
|
||||
ToolCall(
|
||||
ToolCallDelta {
|
||||
tool_index: 0,
|
||||
name: Some(
|
||||
"get_weather",
|
||||
),
|
||||
arguments: "{\"location\":\"Shanghai\"}",
|
||||
},
|
||||
),
|
||||
ToolCall(
|
||||
ToolCallDelta {
|
||||
tool_index: 1,
|
||||
name: Some(
|
||||
"add",
|
||||
),
|
||||
arguments: "{\"x\":1,\"y\":2}",
|
||||
},
|
||||
),
|
||||
],
|
||||
}
|
||||
"#]]
|
||||
@@ -451,22 +454,25 @@ mod tests {
|
||||
|
||||
expect![[r#"
|
||||
ToolParserOutput {
|
||||
normal_text: "",
|
||||
calls: [
|
||||
ToolCallDelta {
|
||||
tool_index: 0,
|
||||
name: Some(
|
||||
"get_weather",
|
||||
),
|
||||
arguments: "{\"location\":\"Shanghai\"}",
|
||||
},
|
||||
ToolCallDelta {
|
||||
tool_index: 1,
|
||||
name: Some(
|
||||
"add",
|
||||
),
|
||||
arguments: "{\"x\":1,\"y\":2}",
|
||||
},
|
||||
events: [
|
||||
ToolCall(
|
||||
ToolCallDelta {
|
||||
tool_index: 0,
|
||||
name: Some(
|
||||
"get_weather",
|
||||
),
|
||||
arguments: "{\"location\":\"Shanghai\"}",
|
||||
},
|
||||
),
|
||||
ToolCall(
|
||||
ToolCallDelta {
|
||||
tool_index: 1,
|
||||
name: Some(
|
||||
"add",
|
||||
),
|
||||
arguments: "{\"x\":1,\"y\":2}",
|
||||
},
|
||||
),
|
||||
],
|
||||
}
|
||||
"#]]
|
||||
@@ -486,15 +492,19 @@ mod tests {
|
||||
|
||||
expect![[r#"
|
||||
ToolParserOutput {
|
||||
normal_text: " trailing text",
|
||||
calls: [
|
||||
ToolCallDelta {
|
||||
tool_index: 0,
|
||||
name: Some(
|
||||
"get_weather",
|
||||
),
|
||||
arguments: "{\"location\":\"Shanghai\"}",
|
||||
},
|
||||
events: [
|
||||
Text(
|
||||
" trailing text",
|
||||
),
|
||||
ToolCall(
|
||||
ToolCallDelta {
|
||||
tool_index: 0,
|
||||
name: Some(
|
||||
"get_weather",
|
||||
),
|
||||
arguments: "{\"location\":\"Shanghai\"}",
|
||||
},
|
||||
),
|
||||
],
|
||||
}
|
||||
"#]]
|
||||
+52
-49
@@ -1,5 +1,5 @@
|
||||
use super::{JsonToolCallConfig, JsonToolCallParser, JsonToolCallWhitespace};
|
||||
use crate::{Result, Tool, ToolParser, ToolParserOutput};
|
||||
use crate::tool::{Result, Tool, ToolParser, ToolParserOutput};
|
||||
|
||||
const PHI4MINI_CONFIG: JsonToolCallConfig = JsonToolCallConfig {
|
||||
parser_name: "Phi4Mini",
|
||||
@@ -69,8 +69,8 @@ mod tests {
|
||||
use thiserror_ext::AsReport;
|
||||
|
||||
use super::Phi4MiniJsonToolParser;
|
||||
use crate::test_utils::{collect_stream, split_by_chars, test_tools};
|
||||
use crate::{ToolParser, ToolParserTestExt as _};
|
||||
use crate::tool::test_utils::{collect_stream, split_by_chars, test_tools};
|
||||
use crate::tool::{ToolParser, ToolParserTestExt as _};
|
||||
|
||||
/// Build one phi-4-mini tool-call object: `{"name":..,"<args_key>":<args>}`.
|
||||
fn build_call(function_name: &str, args_key: &str, arguments: &str) -> String {
|
||||
@@ -87,8 +87,8 @@ mod tests {
|
||||
let mut parser = Phi4MiniJsonToolParser::new(&test_tools());
|
||||
let result = parser.parse_complete("Hello, world!").unwrap();
|
||||
|
||||
assert_eq!(result.normal_text, "Hello, world!");
|
||||
assert!(result.calls.is_empty());
|
||||
assert_eq!(result.normal_text(), "Hello, world!");
|
||||
assert!(result.calls().is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -99,10 +99,10 @@ mod tests {
|
||||
.parse_complete(&wrap(&[build_call("get_weather", "arguments", arguments)]))
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(result.calls.len(), 1);
|
||||
assert_eq!(result.calls[0].tool_index, 0);
|
||||
assert_eq!(result.calls[0].name.as_deref(), Some("get_weather"));
|
||||
assert_eq!(result.calls[0].arguments, arguments);
|
||||
assert_eq!(result.calls().len(), 1);
|
||||
assert_eq!(result.calls()[0].tool_index, 0);
|
||||
assert_eq!(result.calls()[0].name.as_deref(), Some("get_weather"));
|
||||
assert_eq!(result.calls()[0].arguments, arguments);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -113,9 +113,9 @@ mod tests {
|
||||
.parse_complete(&wrap(&[build_call("get_weather", "parameters", arguments)]))
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(result.calls.len(), 1);
|
||||
assert_eq!(result.calls[0].name.as_deref(), Some("get_weather"));
|
||||
assert_eq!(result.calls[0].arguments, arguments);
|
||||
assert_eq!(result.calls().len(), 1);
|
||||
assert_eq!(result.calls()[0].name.as_deref(), Some("get_weather"));
|
||||
assert_eq!(result.calls()[0].arguments, arguments);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -130,22 +130,25 @@ mod tests {
|
||||
|
||||
expect![[r#"
|
||||
ToolParserOutput {
|
||||
normal_text: "",
|
||||
calls: [
|
||||
ToolCallDelta {
|
||||
tool_index: 0,
|
||||
name: Some(
|
||||
"get_weather",
|
||||
),
|
||||
arguments: "{\"location\":\"Shanghai\"}",
|
||||
},
|
||||
ToolCallDelta {
|
||||
tool_index: 1,
|
||||
name: Some(
|
||||
"add",
|
||||
),
|
||||
arguments: "{\"x\":1,\"y\":2}",
|
||||
},
|
||||
events: [
|
||||
ToolCall(
|
||||
ToolCallDelta {
|
||||
tool_index: 0,
|
||||
name: Some(
|
||||
"get_weather",
|
||||
),
|
||||
arguments: "{\"location\":\"Shanghai\"}",
|
||||
},
|
||||
),
|
||||
ToolCall(
|
||||
ToolCallDelta {
|
||||
tool_index: 1,
|
||||
name: Some(
|
||||
"add",
|
||||
),
|
||||
arguments: "{\"x\":1,\"y\":2}",
|
||||
},
|
||||
),
|
||||
],
|
||||
}
|
||||
"#]]
|
||||
@@ -162,8 +165,8 @@ mod tests {
|
||||
.parse_complete(&wrap(&[build_call("convert", "arguments", arguments)]))
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(result.calls.len(), 1);
|
||||
assert_eq!(result.calls[0].arguments, arguments);
|
||||
assert_eq!(result.calls().len(), 1);
|
||||
assert_eq!(result.calls()[0].arguments, arguments);
|
||||
}
|
||||
|
||||
/// Preface text before a tool call is preserved as normal_text, consistent
|
||||
@@ -182,8 +185,8 @@ mod tests {
|
||||
|
||||
let result = parser.parse_complete(&input).unwrap();
|
||||
|
||||
assert_eq!(result.normal_text, "Let me check.\n");
|
||||
assert_eq!(result.calls.len(), 1);
|
||||
assert_eq!(result.normal_text(), "Let me check.\n");
|
||||
assert_eq!(result.calls().len(), 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -194,7 +197,7 @@ mod tests {
|
||||
.parse_complete(&wrap(&[build_call("get_weather", "arguments", arguments)]))
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(result.calls[0].arguments, arguments);
|
||||
assert_eq!(result.calls()[0].arguments, arguments);
|
||||
}
|
||||
|
||||
/// The bundled `tool_chat_template_phi4_mini.jinja` emits objects with
|
||||
@@ -208,9 +211,9 @@ mod tests {
|
||||
|
||||
let result = parser.parse_complete(input).unwrap();
|
||||
|
||||
assert_eq!(result.calls.len(), 1);
|
||||
assert_eq!(result.calls[0].name.as_deref(), Some("get_weather"));
|
||||
assert_eq!(result.calls[0].arguments, r#"{"location": "Tokyo"}"#);
|
||||
assert_eq!(result.calls().len(), 1);
|
||||
assert_eq!(result.calls()[0].name.as_deref(), Some("get_weather"));
|
||||
assert_eq!(result.calls()[0].arguments, r#"{"location": "Tokyo"}"#);
|
||||
}
|
||||
|
||||
/// Argument deltas are streamed through the shared JSON core.
|
||||
@@ -230,10 +233,10 @@ mod tests {
|
||||
|
||||
let result = collect_stream(&mut parser, &chunks);
|
||||
|
||||
assert_eq!(result.normal_text, "preface suffix");
|
||||
assert_eq!(result.calls.len(), 1);
|
||||
assert_eq!(result.calls[0].name.as_deref(), Some("get_weather"));
|
||||
assert_eq!(result.calls[0].arguments, r#"{"location":"Beijing"}"#);
|
||||
assert_eq!(result.normal_text(), "preface suffix");
|
||||
assert_eq!(result.calls().len(), 1);
|
||||
assert_eq!(result.calls()[0].name.as_deref(), Some("get_weather"));
|
||||
assert_eq!(result.calls()[0].arguments, r#"{"location":"Beijing"}"#);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -251,9 +254,9 @@ mod tests {
|
||||
|
||||
let result = collect_stream(&mut parser, &chunks);
|
||||
|
||||
assert_eq!(result.normal_text, "hello ");
|
||||
assert_eq!(result.calls.len(), 1);
|
||||
assert_eq!(result.calls[0].arguments, r#"{"location":"Tokyo"}"#);
|
||||
assert_eq!(result.normal_text(), "hello ");
|
||||
assert_eq!(result.calls().len(), 1);
|
||||
assert_eq!(result.calls()[0].arguments, r#"{"location":"Tokyo"}"#);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -287,9 +290,9 @@ mod tests {
|
||||
.parse_complete(&wrap(&[build_call("convert", "arguments", arguments)]))
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(result.calls.len(), 1);
|
||||
assert_eq!(result.calls[0].name.as_deref(), Some("convert"));
|
||||
assert_eq!(result.calls[0].arguments, arguments);
|
||||
assert_eq!(result.calls().len(), 1);
|
||||
assert_eq!(result.calls()[0].name.as_deref(), Some("convert"));
|
||||
assert_eq!(result.calls()[0].arguments, arguments);
|
||||
}
|
||||
|
||||
/// The chat template emits parallel calls as `},\n {` (comma + newline +
|
||||
@@ -307,9 +310,9 @@ mod tests {
|
||||
|
||||
let result = parser.parse_complete(input).unwrap();
|
||||
|
||||
assert_eq!(result.calls.len(), 2);
|
||||
assert_eq!(result.calls[0].name.as_deref(), Some("get_weather"));
|
||||
assert_eq!(result.calls[1].name.as_deref(), Some("add"));
|
||||
assert_eq!(result.calls().len(), 2);
|
||||
assert_eq!(result.calls()[0].name.as_deref(), Some("get_weather"));
|
||||
assert_eq!(result.calls()[1].name.as_deref(), Some("add"));
|
||||
}
|
||||
|
||||
/// The shared core requires an object after the start marker.
|
||||
@@ -1,5 +1,5 @@
|
||||
use super::{JsonToolCallConfig, JsonToolCallParser, JsonToolCallWhitespace};
|
||||
use crate::{Result, StructuralTagModel, Tool, ToolParser, ToolParserOutput};
|
||||
use crate::tool::{Result, StructuralTagModel, Tool, ToolParser, ToolParserOutput};
|
||||
|
||||
const QWEN_XML_CONFIG: JsonToolCallConfig = JsonToolCallConfig {
|
||||
parser_name: "Qwen XML",
|
||||
@@ -70,8 +70,8 @@ mod tests {
|
||||
use thiserror_ext::AsReport;
|
||||
|
||||
use super::Qwen3XmlToolParser;
|
||||
use crate::test_utils::{collect_stream, split_by_chars, test_tools};
|
||||
use crate::{ToolParser, ToolParserOutput, ToolParserTestExt as _};
|
||||
use crate::tool::test_utils::{collect_stream, split_by_chars, test_tools};
|
||||
use crate::tool::{ToolParser, ToolParserOutput, ToolParserTestExt as _};
|
||||
|
||||
fn build_tool_call(function_name: &str, arguments: &str) -> String {
|
||||
format!(
|
||||
@@ -84,8 +84,8 @@ mod tests {
|
||||
let mut parser = Qwen3XmlToolParser::new(&test_tools());
|
||||
let output = parser.parse_complete("Hello, world!").unwrap();
|
||||
|
||||
assert_eq!(output.normal_text, "Hello, world!");
|
||||
assert!(output.calls.is_empty());
|
||||
assert_eq!(output.normal_text(), "Hello, world!");
|
||||
assert!(output.calls().is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -99,11 +99,11 @@ mod tests {
|
||||
))
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(output.normal_text, "Let me check.\n");
|
||||
assert_eq!(output.calls.len(), 1);
|
||||
assert_eq!(output.calls[0].tool_index, 0);
|
||||
assert_eq!(output.calls[0].name.as_deref(), Some("get_weather"));
|
||||
assert_eq!(output.calls[0].arguments, arguments);
|
||||
assert_eq!(output.normal_text(), "Let me check.\n");
|
||||
assert_eq!(output.calls().len(), 1);
|
||||
assert_eq!(output.calls()[0].tool_index, 0);
|
||||
assert_eq!(output.calls()[0].name.as_deref(), Some("get_weather"));
|
||||
assert_eq!(output.calls()[0].arguments, arguments);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -112,7 +112,7 @@ mod tests {
|
||||
let arguments = r#"{"location":"Tokyo",}"#;
|
||||
let output = parser.parse_complete(&build_tool_call("get_weather", arguments)).unwrap();
|
||||
|
||||
assert_eq!(output.calls[0].arguments, arguments);
|
||||
assert_eq!(output.calls()[0].arguments, arguments);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -132,7 +132,7 @@ mod tests {
|
||||
for chunk in chunks {
|
||||
let next = parser.parse_chunk(chunk).unwrap();
|
||||
observed_arguments.extend(
|
||||
next.calls
|
||||
next.calls()
|
||||
.iter()
|
||||
.filter(|call| call.name.is_none())
|
||||
.map(|call| call.arguments.clone()),
|
||||
@@ -143,7 +143,7 @@ mod tests {
|
||||
|
||||
assert_eq!(observed_arguments, ["{\"location\":", "\"Beijing\"", "}"]);
|
||||
assert_eq!(
|
||||
output.coalesce_calls().calls[0].arguments,
|
||||
output.coalesce().calls()[0].arguments,
|
||||
r#"{"location":"Beijing"}"#
|
||||
);
|
||||
}
|
||||
@@ -159,9 +159,9 @@ mod tests {
|
||||
|
||||
let output = collect_stream(&mut parser, &chunks);
|
||||
|
||||
assert_eq!(output.normal_text, "hello ");
|
||||
assert_eq!(output.calls.len(), 1);
|
||||
assert_eq!(output.calls[0].arguments, r#"{"location":"Tokyo"}"#);
|
||||
assert_eq!(output.normal_text(), "hello ");
|
||||
assert_eq!(output.calls().len(), 1);
|
||||
assert_eq!(output.calls()[0].arguments, r#"{"location":"Tokyo"}"#);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -170,8 +170,8 @@ mod tests {
|
||||
let arguments = r#"{"text":"literal </tool_call> inside"}"#;
|
||||
let output = parser.parse_complete(&build_tool_call("echo", arguments)).unwrap();
|
||||
|
||||
assert_eq!(output.calls.len(), 1);
|
||||
assert_eq!(output.calls[0].arguments, arguments);
|
||||
assert_eq!(output.calls().len(), 1);
|
||||
assert_eq!(output.calls()[0].arguments, arguments);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -185,7 +185,7 @@ mod tests {
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(output.calls[0].name.as_deref(), Some("say_\"hi"));
|
||||
assert_eq!(output.calls()[0].name.as_deref(), Some("say_\"hi"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -196,8 +196,8 @@ mod tests {
|
||||
|
||||
let output = parser.parse_complete(input).unwrap();
|
||||
|
||||
assert_eq!(output.normal_text, input);
|
||||
assert!(output.calls.is_empty());
|
||||
assert_eq!(output.normal_text(), input);
|
||||
assert!(output.calls().is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -227,22 +227,25 @@ mod tests {
|
||||
|
||||
expect![[r#"
|
||||
ToolParserOutput {
|
||||
normal_text: "",
|
||||
calls: [
|
||||
ToolCallDelta {
|
||||
tool_index: 0,
|
||||
name: Some(
|
||||
"get_weather",
|
||||
),
|
||||
arguments: "{\"location\":\"Shanghai\"}",
|
||||
},
|
||||
ToolCallDelta {
|
||||
tool_index: 1,
|
||||
name: Some(
|
||||
"add",
|
||||
),
|
||||
arguments: "{\"x\":1,\"y\":2}",
|
||||
},
|
||||
events: [
|
||||
ToolCall(
|
||||
ToolCallDelta {
|
||||
tool_index: 0,
|
||||
name: Some(
|
||||
"get_weather",
|
||||
),
|
||||
arguments: "{\"location\":\"Shanghai\"}",
|
||||
},
|
||||
),
|
||||
ToolCall(
|
||||
ToolCallDelta {
|
||||
tool_index: 1,
|
||||
name: Some(
|
||||
"add",
|
||||
),
|
||||
arguments: "{\"x\":1,\"y\":2}",
|
||||
},
|
||||
),
|
||||
],
|
||||
}
|
||||
"#]]
|
||||
@@ -8,7 +8,7 @@ use winnow::token::{literal, rest, take_until, take_while};
|
||||
|
||||
use super::utils::{JsonObjectScanState, parse_buffered_event, safe_text_len, take_json_object};
|
||||
use super::{Result, ToolCallDelta, ToolParser, ToolParserOutput};
|
||||
use crate::{StructuralTagModel, Tool};
|
||||
use crate::tool::{StructuralTagModel, Tool};
|
||||
|
||||
const TOOL_CALLS_START: &str = "<|tool_calls_section_begin|>";
|
||||
const TOOL_CALLS_END: &str = "<|tool_calls_section_end|>";
|
||||
@@ -81,7 +81,7 @@ impl KimiK2ToolParser {
|
||||
fn apply_event(&mut self, event: KimiK2Event, output: &mut ToolParserOutput) -> Result<()> {
|
||||
match event {
|
||||
KimiK2Event::Text { len: consumed_len } => {
|
||||
output.normal_text.push_str(&self.buffer[..consumed_len]);
|
||||
output.push_text(&self.buffer[..consumed_len]);
|
||||
}
|
||||
KimiK2Event::ToolCallsStart => self.mode = KimiK2Mode::ToolBlock,
|
||||
KimiK2Event::ToolCallStart => self.mode = KimiK2Mode::Header,
|
||||
@@ -96,7 +96,7 @@ impl KimiK2ToolParser {
|
||||
json_scan: JsonObjectScanState::default(),
|
||||
};
|
||||
self.call_ids.insert(tool_index, tool_call_id);
|
||||
output.calls.push(ToolCallDelta {
|
||||
output.push_call(ToolCallDelta {
|
||||
tool_index,
|
||||
name: Some(function_name),
|
||||
arguments: String::new(),
|
||||
@@ -108,7 +108,7 @@ impl KimiK2ToolParser {
|
||||
"Kimi K2 arguments without an active tool call"
|
||||
));
|
||||
};
|
||||
output.calls.push(ToolCallDelta {
|
||||
output.push_call(ToolCallDelta {
|
||||
tool_index,
|
||||
name: None,
|
||||
arguments: self.buffer[..consumed_len].to_string(),
|
||||
@@ -171,7 +171,7 @@ impl ToolParser for KimiK2ToolParser {
|
||||
fn finish(&mut self) -> Result<ToolParserOutput> {
|
||||
let mut output = ToolParserOutput::default();
|
||||
match &self.mode {
|
||||
KimiK2Mode::Text => output.normal_text.push_str(&self.buffer),
|
||||
KimiK2Mode::Text => output.push_text(&self.buffer),
|
||||
KimiK2Mode::ToolBlock | KimiK2Mode::Done => {}
|
||||
KimiK2Mode::Header | KimiK2Mode::Arguments { .. } => {
|
||||
return Err(parsing_failed!("incomplete Kimi K2 tool call"));
|
||||
@@ -339,8 +339,8 @@ mod tests {
|
||||
KimiK2ToolParser, TOOL_CALL_ARGUMENT_START, TOOL_CALL_END, TOOL_CALL_START, TOOL_CALLS_END,
|
||||
TOOL_CALLS_START, ToolParser, tool_header,
|
||||
};
|
||||
use crate::test_utils::{collect_stream, split_by_chars, test_tools};
|
||||
use crate::{ToolParserOutput, ToolParserTestExt as _};
|
||||
use crate::tool::test_utils::{collect_stream, split_by_chars, test_tools};
|
||||
use crate::tool::{ToolParserOutput, ToolParserTestExt as _};
|
||||
|
||||
fn build_tool_call(function_name: &str, index: usize, arguments: &str) -> String {
|
||||
format!(
|
||||
@@ -357,8 +357,8 @@ mod tests {
|
||||
let mut parser = KimiK2ToolParser::new(&test_tools());
|
||||
let output = parser.parse_complete("Hello, world!").unwrap();
|
||||
|
||||
assert_eq!(output.normal_text, "Hello, world!");
|
||||
assert!(output.calls.is_empty());
|
||||
assert_eq!(output.normal_text(), "Hello, world!");
|
||||
assert!(output.calls().is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -372,11 +372,11 @@ mod tests {
|
||||
))
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(output.normal_text, "Checking. ");
|
||||
assert_eq!(output.calls.len(), 1);
|
||||
assert_eq!(output.calls[0].tool_index, 0);
|
||||
assert_eq!(output.calls[0].name.as_deref(), Some("get_weather"));
|
||||
assert_eq!(output.calls[0].arguments, arguments);
|
||||
assert_eq!(output.normal_text(), "Checking. ");
|
||||
assert_eq!(output.calls().len(), 1);
|
||||
assert_eq!(output.calls()[0].tool_index, 0);
|
||||
assert_eq!(output.calls()[0].name.as_deref(), Some("get_weather"));
|
||||
assert_eq!(output.calls()[0].arguments, arguments);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -391,7 +391,7 @@ mod tests {
|
||||
)]))
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(output.calls[0].arguments, arguments);
|
||||
assert_eq!(output.calls()[0].arguments, arguments);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -414,7 +414,7 @@ mod tests {
|
||||
for chunk in chunks {
|
||||
let next = parser.parse_chunk(chunk).unwrap();
|
||||
observed_arguments.extend(
|
||||
next.calls
|
||||
next.calls()
|
||||
.iter()
|
||||
.filter(|call| call.name.is_none())
|
||||
.map(|call| call.arguments.clone()),
|
||||
@@ -424,8 +424,8 @@ mod tests {
|
||||
output.append(parser.finish().unwrap());
|
||||
|
||||
assert_eq!(observed_arguments, ["{\"location\":", "\"Paris\"", "}"]);
|
||||
let output = output.coalesce_calls();
|
||||
assert_eq!(output.calls[0].arguments, r#"{"location":"Paris"}"#);
|
||||
let output = output.coalesce();
|
||||
assert_eq!(output.calls()[0].arguments, r#"{"location":"Paris"}"#);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -445,9 +445,9 @@ mod tests {
|
||||
|
||||
let output = collect_stream(&mut parser, &chunks);
|
||||
|
||||
assert_eq!(output.normal_text, "hello ");
|
||||
assert_eq!(output.calls.len(), 1);
|
||||
assert_eq!(output.calls[0].arguments, r#"{"location":"NYC"}"#);
|
||||
assert_eq!(output.normal_text(), "hello ");
|
||||
assert_eq!(output.calls().len(), 1);
|
||||
assert_eq!(output.calls()[0].arguments, r#"{"location":"NYC"}"#);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -458,8 +458,8 @@ mod tests {
|
||||
|
||||
let output = parser.parse_complete(&input).unwrap();
|
||||
|
||||
assert_eq!(output.calls.len(), 1);
|
||||
assert_eq!(output.calls[0].arguments, arguments);
|
||||
assert_eq!(output.calls().len(), 1);
|
||||
assert_eq!(output.calls()[0].arguments, arguments);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -478,9 +478,9 @@ mod tests {
|
||||
|
||||
let output = collect_stream(&mut parser, &chunks);
|
||||
|
||||
assert_eq!(output.calls.len(), 1);
|
||||
assert_eq!(output.calls().len(), 1);
|
||||
assert_eq!(
|
||||
output.calls[0].arguments,
|
||||
output.calls()[0].arguments,
|
||||
r#"{"text":"literal <|tool_call_end|> inside"}"#
|
||||
);
|
||||
}
|
||||
@@ -498,22 +498,25 @@ mod tests {
|
||||
|
||||
expect![[r#"
|
||||
ToolParserOutput {
|
||||
normal_text: "",
|
||||
calls: [
|
||||
ToolCallDelta {
|
||||
tool_index: 0,
|
||||
name: Some(
|
||||
"get_weather",
|
||||
),
|
||||
arguments: "{\"location\":\"Shanghai\"}",
|
||||
},
|
||||
ToolCallDelta {
|
||||
tool_index: 1,
|
||||
name: Some(
|
||||
"add",
|
||||
),
|
||||
arguments: "{\"x\":1,\"y\":2}",
|
||||
},
|
||||
events: [
|
||||
ToolCall(
|
||||
ToolCallDelta {
|
||||
tool_index: 0,
|
||||
name: Some(
|
||||
"get_weather",
|
||||
),
|
||||
arguments: "{\"location\":\"Shanghai\"}",
|
||||
},
|
||||
),
|
||||
ToolCall(
|
||||
ToolCallDelta {
|
||||
tool_index: 1,
|
||||
name: Some(
|
||||
"add",
|
||||
),
|
||||
arguments: "{\"x\":1,\"y\":2}",
|
||||
},
|
||||
),
|
||||
],
|
||||
}
|
||||
"#]]
|
||||
@@ -546,12 +549,12 @@ mod tests {
|
||||
"{TOOL_CALLS_START}{TOOL_CALL_START}api.tools.search:42{TOOL_CALL_ARGUMENT_START}{{}}{TOOL_CALL_END}{TOOL_CALLS_END}"
|
||||
);
|
||||
|
||||
let output = parser.parse_chunk(&input).unwrap().coalesce_calls();
|
||||
let output = parser.parse_chunk(&input).unwrap().coalesce();
|
||||
|
||||
assert_eq!(output.calls[0].tool_index, 42);
|
||||
assert_eq!(output.calls()[0].tool_index, 42);
|
||||
assert_eq!(parser.tool_call_id(42), Some("api.tools.search:42"));
|
||||
assert_eq!(output.calls[0].name.as_deref(), Some("search"));
|
||||
assert_eq!(output.calls[0].arguments, "{}");
|
||||
assert_eq!(output.calls()[0].name.as_deref(), Some("search"));
|
||||
assert_eq!(output.calls()[0].arguments, "{}");
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -7,7 +7,7 @@ use winnow::token::{literal, rest, take_until};
|
||||
use super::parameters::ToolSchemas;
|
||||
use super::utils::{MarkerScanState, parse_buffered_event, safe_text_len, take_until_marker};
|
||||
use super::{Result, ToolCallDelta, ToolParser, ToolParserOutput};
|
||||
use crate::{StructuralTagModel, Tool};
|
||||
use crate::tool::{StructuralTagModel, Tool};
|
||||
|
||||
const TOOL_CALL_START: &str = "<minimax:tool_call>";
|
||||
const TOOL_CALL_END: &str = "</minimax:tool_call>";
|
||||
@@ -72,7 +72,7 @@ impl MinimaxM2ToolParser {
|
||||
fn apply_event(&mut self, event: MinimaxM2Event, output: &mut ToolParserOutput) -> Result<()> {
|
||||
match event {
|
||||
MinimaxM2Event::Text { len: consumed_len } => {
|
||||
output.normal_text.push_str(&self.buffer[..consumed_len]);
|
||||
output.push_text(&self.buffer[..consumed_len]);
|
||||
}
|
||||
MinimaxM2Event::ToolBlockStart => {
|
||||
self.mode = MinimaxM2Mode::ToolBlock {
|
||||
@@ -84,7 +84,7 @@ impl MinimaxM2ToolParser {
|
||||
let arguments = serde_json::to_string(&arguments)
|
||||
.map_err(|error| parsing_failed!("failed to serialize arguments: {}", error))?;
|
||||
|
||||
output.calls.push(ToolCallDelta {
|
||||
output.push_call(ToolCallDelta {
|
||||
tool_index: self.emitted_tool_count,
|
||||
name: Some(name),
|
||||
arguments,
|
||||
@@ -133,7 +133,7 @@ impl ToolParser for MinimaxM2ToolParser {
|
||||
let mut output = ToolParserOutput::default();
|
||||
match self.mode {
|
||||
MinimaxM2Mode::Text => {
|
||||
output.normal_text.push_str(&self.buffer);
|
||||
output.push_text(&self.buffer);
|
||||
}
|
||||
MinimaxM2Mode::ToolBlock { .. } => {
|
||||
return Err(parsing_failed!("incomplete MiniMax M2 tool call"));
|
||||
@@ -272,8 +272,8 @@ mod tests {
|
||||
use thiserror_ext::AsReport;
|
||||
|
||||
use super::{MinimaxM2ToolParser, TOOL_CALL_END, TOOL_CALL_START, ToolParser};
|
||||
use crate::ToolParserTestExt as _;
|
||||
use crate::test_utils::{collect_stream, split_by_chars, test_tools};
|
||||
use crate::tool::ToolParserTestExt as _;
|
||||
use crate::tool::test_utils::{collect_stream, split_by_chars, test_tools};
|
||||
|
||||
fn build_tool_block(invokes: &[(&str, Vec<(&str, &str)>)]) -> String {
|
||||
let invokes = invokes
|
||||
@@ -295,8 +295,8 @@ mod tests {
|
||||
let mut parser = MinimaxM2ToolParser::new(&test_tools());
|
||||
let output = parser.parse_complete("Hello, world!").unwrap();
|
||||
|
||||
assert_eq!(output.normal_text, "Hello, world!");
|
||||
assert!(output.calls.is_empty());
|
||||
assert_eq!(output.normal_text(), "Hello, world!");
|
||||
assert!(output.calls().is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -309,11 +309,11 @@ mod tests {
|
||||
)]))
|
||||
.unwrap();
|
||||
|
||||
assert!(output.normal_text.is_empty());
|
||||
assert_eq!(output.calls.len(), 1);
|
||||
assert_eq!(output.calls[0].name.as_deref(), Some("get_weather"));
|
||||
assert!(output.normal_text().is_empty());
|
||||
assert_eq!(output.calls().len(), 1);
|
||||
assert_eq!(output.calls()[0].name.as_deref(), Some("get_weather"));
|
||||
assert_eq!(
|
||||
serde_json::from_str::<Value>(&output.calls[0].arguments).unwrap(),
|
||||
serde_json::from_str::<Value>(&output.calls()[0].arguments).unwrap(),
|
||||
json!({ "city": "Seattle", "days": 5 })
|
||||
);
|
||||
}
|
||||
@@ -327,8 +327,8 @@ mod tests {
|
||||
);
|
||||
let output = parser.parse_complete(&output).unwrap();
|
||||
|
||||
assert_eq!(output.normal_text, "Let me check. ");
|
||||
assert_eq!(output.calls.len(), 1);
|
||||
assert_eq!(output.normal_text(), "Let me check. ");
|
||||
assert_eq!(output.calls().len(), 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -341,15 +341,15 @@ mod tests {
|
||||
]))
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(output.calls.len(), 2);
|
||||
assert_eq!(output.calls[0].tool_index, 0);
|
||||
assert_eq!(output.calls[1].tool_index, 1);
|
||||
assert_eq!(output.calls().len(), 2);
|
||||
assert_eq!(output.calls()[0].tool_index, 0);
|
||||
assert_eq!(output.calls()[1].tool_index, 1);
|
||||
assert_eq!(
|
||||
serde_json::from_str::<Value>(&output.calls[0].arguments).unwrap(),
|
||||
serde_json::from_str::<Value>(&output.calls()[0].arguments).unwrap(),
|
||||
json!({ "city": "Seattle" })
|
||||
);
|
||||
assert_eq!(
|
||||
serde_json::from_str::<Value>(&output.calls[1].arguments).unwrap(),
|
||||
serde_json::from_str::<Value>(&output.calls()[1].arguments).unwrap(),
|
||||
json!({ "city": "NYC" })
|
||||
);
|
||||
}
|
||||
@@ -371,7 +371,7 @@ mod tests {
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
serde_json::from_str::<Value>(&output.calls[0].arguments).unwrap(),
|
||||
serde_json::from_str::<Value>(&output.calls()[0].arguments).unwrap(),
|
||||
json!({
|
||||
"whole": 5.0,
|
||||
"flag": true,
|
||||
@@ -395,7 +395,7 @@ mod tests {
|
||||
vec![("city", "Tom & Jerry <3")],
|
||||
)]))
|
||||
.unwrap();
|
||||
let args: Value = serde_json::from_str(&output.calls[0].arguments).unwrap();
|
||||
let args: Value = serde_json::from_str(&output.calls()[0].arguments).unwrap();
|
||||
assert_eq!(args["city"], json!("Tom & Jerry <3"));
|
||||
}
|
||||
|
||||
@@ -416,7 +416,7 @@ mod tests {
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
serde_json::from_str::<Value>(&output.calls[0].arguments).unwrap(),
|
||||
serde_json::from_str::<Value>(&output.calls()[0].arguments).unwrap(),
|
||||
json!({
|
||||
"city": "Seattle </parameter></invoke></minimax:tool_call>",
|
||||
"days": 5,
|
||||
@@ -440,7 +440,7 @@ mod tests {
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
serde_json::from_str::<Value>(&output.calls[0].arguments).unwrap(),
|
||||
serde_json::from_str::<Value>(&output.calls()[0].arguments).unwrap(),
|
||||
json!({
|
||||
"shape": "\nrectangle\n",
|
||||
"dimensions": { "width": 10, "height": 20 },
|
||||
@@ -462,11 +462,11 @@ mod tests {
|
||||
],
|
||||
);
|
||||
|
||||
assert!(output.normal_text.is_empty());
|
||||
assert_eq!(output.calls.len(), 1);
|
||||
assert_eq!(output.calls[0].name.as_deref(), Some("get_weather"));
|
||||
assert!(output.normal_text().is_empty());
|
||||
assert_eq!(output.calls().len(), 1);
|
||||
assert_eq!(output.calls()[0].name.as_deref(), Some("get_weather"));
|
||||
assert_eq!(
|
||||
serde_json::from_str::<Value>(&output.calls[0].arguments).unwrap(),
|
||||
serde_json::from_str::<Value>(&output.calls()[0].arguments).unwrap(),
|
||||
json!({ "city": "Seattle" })
|
||||
);
|
||||
}
|
||||
@@ -484,8 +484,8 @@ mod tests {
|
||||
],
|
||||
);
|
||||
|
||||
assert_eq!(output.normal_text, "Let me check. ");
|
||||
assert_eq!(output.calls.len(), 1);
|
||||
assert_eq!(output.normal_text(), "Let me check. ");
|
||||
assert_eq!(output.calls().len(), 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -493,8 +493,8 @@ mod tests {
|
||||
let mut parser = MinimaxM2ToolParser::new(&test_tools());
|
||||
let output = collect_stream(&mut parser, &["Hello, ", "world!"]);
|
||||
|
||||
assert_eq!(output.normal_text, "Hello, world!");
|
||||
assert!(output.calls.is_empty());
|
||||
assert_eq!(output.normal_text(), "Hello, world!");
|
||||
assert!(output.calls().is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -504,8 +504,8 @@ mod tests {
|
||||
let mut parser = MinimaxM2ToolParser::new(&test_tools());
|
||||
let output = collect_stream(&mut parser, &chunks);
|
||||
|
||||
assert_eq!(output.calls.len(), 1);
|
||||
assert!(output.normal_text.is_empty());
|
||||
assert_eq!(output.calls().len(), 1);
|
||||
assert!(output.normal_text().is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -518,9 +518,9 @@ mod tests {
|
||||
let mut parser = MinimaxM2ToolParser::new(&test_tools());
|
||||
let output = collect_stream(&mut parser, &chunks);
|
||||
|
||||
assert_eq!(output.calls.len(), 2);
|
||||
assert_eq!(output.calls[0].tool_index, 0);
|
||||
assert_eq!(output.calls[1].tool_index, 1);
|
||||
assert_eq!(output.calls().len(), 2);
|
||||
assert_eq!(output.calls()[0].tool_index, 0);
|
||||
assert_eq!(output.calls()[1].tool_index, 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -540,12 +540,12 @@ mod tests {
|
||||
let mut parser = MinimaxM2ToolParser::new(&test_tools());
|
||||
let result = collect_stream(&mut parser, &chunks);
|
||||
|
||||
assert_eq!(result.normal_text, "I will call the tools.\n");
|
||||
assert_eq!(result.calls.len(), 2);
|
||||
assert_eq!(result.calls[0].tool_index, 0);
|
||||
assert_eq!(result.calls[0].name.as_deref(), Some("get_weather"));
|
||||
assert_eq!(result.calls[1].tool_index, 1);
|
||||
assert_eq!(result.calls[1].name.as_deref(), Some("get_weather"));
|
||||
assert_eq!(result.normal_text(), "I will call the tools.\n");
|
||||
assert_eq!(result.calls().len(), 2);
|
||||
assert_eq!(result.calls()[0].tool_index, 0);
|
||||
assert_eq!(result.calls()[0].name.as_deref(), Some("get_weather"));
|
||||
assert_eq!(result.calls()[1].tool_index, 1);
|
||||
assert_eq!(result.calls()[1].name.as_deref(), Some("get_weather"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -558,8 +558,8 @@ mod tests {
|
||||
let mut parser = MinimaxM2ToolParser::new(&test_tools());
|
||||
let output = collect_stream(&mut parser, &chunks);
|
||||
|
||||
assert!(output.normal_text.is_empty());
|
||||
assert_eq!(output.calls.len(), 1);
|
||||
assert!(output.normal_text().is_empty());
|
||||
assert_eq!(output.calls().len(), 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -568,8 +568,8 @@ mod tests {
|
||||
let output =
|
||||
parser.parse_chunk(r#"<minimax:tool_call><invoke name="get_weather">"#).unwrap();
|
||||
|
||||
assert!(output.normal_text.is_empty());
|
||||
assert!(output.calls.is_empty());
|
||||
assert!(output.normal_text().is_empty());
|
||||
assert!(output.calls().is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -8,7 +8,7 @@ use winnow::token::{literal, rest, take_until};
|
||||
use super::parameters::{ParamElement, ParamInput, ToolSchemas};
|
||||
use super::utils::{MarkerScanState, parse_buffered_event, safe_text_len, take_until_marker};
|
||||
use super::{Result, ToolCallDelta, ToolParser, ToolParserOutput};
|
||||
use crate::Tool;
|
||||
use crate::tool::Tool;
|
||||
|
||||
const NAMESPACE: &str = "]<]minimax[>[";
|
||||
const TOOL_CALL_START: &str = "]<]minimax[>[<tool_call>";
|
||||
@@ -107,7 +107,7 @@ impl MinimaxM3ToolParser {
|
||||
fn apply_event(&mut self, event: MinimaxM3Event, output: &mut ToolParserOutput) -> Result<()> {
|
||||
match event {
|
||||
MinimaxM3Event::Text { len: consumed_len } => {
|
||||
output.normal_text.push_str(&self.buffer[..consumed_len]);
|
||||
output.push_text(&self.buffer[..consumed_len]);
|
||||
}
|
||||
MinimaxM3Event::ToolBlockStart => {
|
||||
self.mode = MinimaxM3Mode::ToolBlock {
|
||||
@@ -119,7 +119,7 @@ impl MinimaxM3ToolParser {
|
||||
let arguments = serde_json::to_string(&arguments)
|
||||
.map_err(|error| parsing_failed!("failed to serialize arguments: {}", error))?;
|
||||
|
||||
output.calls.push(ToolCallDelta {
|
||||
output.push_call(ToolCallDelta {
|
||||
tool_index: self.emitted_tool_count,
|
||||
name: Some(name),
|
||||
arguments,
|
||||
@@ -158,7 +158,7 @@ impl ToolParser for MinimaxM3ToolParser {
|
||||
let mut output = ToolParserOutput::default();
|
||||
match self.mode {
|
||||
MinimaxM3Mode::Text => {
|
||||
output.normal_text.push_str(&self.buffer);
|
||||
output.push_text(&self.buffer);
|
||||
}
|
||||
MinimaxM3Mode::ToolBlock { .. } => {
|
||||
if !self.buffer.trim_start().is_empty() {
|
||||
@@ -388,8 +388,8 @@ mod tests {
|
||||
ELEMENT_END_START, ELEMENT_START, INVOKE_END, INVOKE_START, MinimaxM3ToolParser,
|
||||
TOOL_CALL_END, TOOL_CALL_START, ToolParser,
|
||||
};
|
||||
use crate::test_utils::{collect_stream, split_by_chars, test_tools};
|
||||
use crate::{Tool, ToolParserTestExt as _};
|
||||
use crate::tool::test_utils::{collect_stream, split_by_chars, test_tools};
|
||||
use crate::tool::{Tool, ToolParserEvent, ToolParserTestExt as _};
|
||||
|
||||
fn element(name: &str, body: &str) -> String {
|
||||
format!("{ELEMENT_START}{name}>{body}{ELEMENT_END_START}{name}>")
|
||||
@@ -510,8 +510,8 @@ mod tests {
|
||||
let mut parser = MinimaxM3ToolParser::new(&m3_test_tools());
|
||||
let output = parser.parse_complete("Hello, world!").unwrap();
|
||||
|
||||
assert_eq!(output.normal_text, "Hello, world!");
|
||||
assert!(output.calls.is_empty());
|
||||
assert_eq!(output.normal_text(), "Hello, world!");
|
||||
assert!(output.calls().is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -524,11 +524,11 @@ mod tests {
|
||||
)]))
|
||||
.unwrap();
|
||||
|
||||
assert!(output.normal_text.is_empty());
|
||||
assert_eq!(output.calls.len(), 1);
|
||||
assert_eq!(output.calls[0].name.as_deref(), Some("get_weather"));
|
||||
assert!(output.normal_text().is_empty());
|
||||
assert_eq!(output.calls().len(), 1);
|
||||
assert_eq!(output.calls()[0].name.as_deref(), Some("get_weather"));
|
||||
assert_eq!(
|
||||
serde_json::from_str::<Value>(&output.calls[0].arguments).unwrap(),
|
||||
serde_json::from_str::<Value>(&output.calls()[0].arguments).unwrap(),
|
||||
json!({ "city": "Seattle", "days": 5 })
|
||||
);
|
||||
}
|
||||
@@ -542,8 +542,8 @@ mod tests {
|
||||
);
|
||||
let output = parser.parse_complete(&output).unwrap();
|
||||
|
||||
assert_eq!(output.normal_text, "Let me check. ");
|
||||
assert_eq!(output.calls.len(), 1);
|
||||
assert_eq!(output.normal_text(), "Let me check. ");
|
||||
assert_eq!(output.calls().len(), 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -556,15 +556,15 @@ mod tests {
|
||||
]))
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(output.calls.len(), 2);
|
||||
assert_eq!(output.calls[0].tool_index, 0);
|
||||
assert_eq!(output.calls[1].tool_index, 1);
|
||||
assert_eq!(output.calls().len(), 2);
|
||||
assert_eq!(output.calls()[0].tool_index, 0);
|
||||
assert_eq!(output.calls()[1].tool_index, 1);
|
||||
assert_eq!(
|
||||
serde_json::from_str::<Value>(&output.calls[0].arguments).unwrap(),
|
||||
serde_json::from_str::<Value>(&output.calls()[0].arguments).unwrap(),
|
||||
json!({ "city": "Seattle" })
|
||||
);
|
||||
assert_eq!(
|
||||
serde_json::from_str::<Value>(&output.calls[1].arguments).unwrap(),
|
||||
serde_json::from_str::<Value>(&output.calls()[1].arguments).unwrap(),
|
||||
json!({ "city": "NYC" })
|
||||
);
|
||||
}
|
||||
@@ -585,7 +585,7 @@ mod tests {
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
serde_json::from_str::<Value>(&output.calls[0].arguments).unwrap(),
|
||||
serde_json::from_str::<Value>(&output.calls()[0].arguments).unwrap(),
|
||||
json!({ "city": "Seattle" })
|
||||
);
|
||||
}
|
||||
@@ -608,7 +608,7 @@ mod tests {
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
serde_json::from_str::<Value>(&output.calls[0].arguments).unwrap(),
|
||||
serde_json::from_str::<Value>(&output.calls()[0].arguments).unwrap(),
|
||||
json!({
|
||||
"whole": 5.0,
|
||||
"flag": true,
|
||||
@@ -627,7 +627,7 @@ mod tests {
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
serde_json::from_str::<Value>(&output.calls[0].arguments).unwrap(),
|
||||
serde_json::from_str::<Value>(&output.calls()[0].arguments).unwrap(),
|
||||
json!({
|
||||
"user_id": 42,
|
||||
"urgent": true,
|
||||
@@ -677,7 +677,7 @@ mod tests {
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
serde_json::from_str::<Value>(&output.calls[0].arguments).unwrap(),
|
||||
serde_json::from_str::<Value>(&output.calls()[0].arguments).unwrap(),
|
||||
json!({
|
||||
"shape": "\nrectangle\n",
|
||||
"dimensions": { "width": 10, "height": 20 },
|
||||
@@ -698,11 +698,11 @@ mod tests {
|
||||
],
|
||||
);
|
||||
|
||||
assert!(output.normal_text.is_empty());
|
||||
assert_eq!(output.calls.len(), 1);
|
||||
assert_eq!(output.calls[0].name.as_deref(), Some("get_weather"));
|
||||
assert!(output.normal_text().is_empty());
|
||||
assert_eq!(output.calls().len(), 1);
|
||||
assert_eq!(output.calls()[0].name.as_deref(), Some("get_weather"));
|
||||
assert_eq!(
|
||||
serde_json::from_str::<Value>(&output.calls[0].arguments).unwrap(),
|
||||
serde_json::from_str::<Value>(&output.calls()[0].arguments).unwrap(),
|
||||
json!({ "city": "Seattle" })
|
||||
);
|
||||
}
|
||||
@@ -720,8 +720,36 @@ mod tests {
|
||||
],
|
||||
);
|
||||
|
||||
assert_eq!(output.normal_text, "Let me check. ");
|
||||
assert_eq!(output.calls.len(), 1);
|
||||
assert_eq!(output.normal_text(), "Let me check. ");
|
||||
assert_eq!(output.calls().len(), 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn minimax_m3_streaming_preserves_ordered_events() {
|
||||
let mut parser = MinimaxM3ToolParser::new(&m3_test_tools());
|
||||
let output = collect_stream(
|
||||
&mut parser,
|
||||
&[
|
||||
"Let me check. ",
|
||||
TOOL_CALL_START,
|
||||
&invoke("get_weather", &element("city", "Seattle")),
|
||||
TOOL_CALL_END,
|
||||
],
|
||||
);
|
||||
|
||||
assert_eq!(output.events.len(), 2);
|
||||
assert_eq!(
|
||||
output.events[0],
|
||||
ToolParserEvent::Text("Let me check. ".to_string())
|
||||
);
|
||||
let ToolParserEvent::ToolCall(call) = &output.events[1] else {
|
||||
panic!("expected tool-call event");
|
||||
};
|
||||
assert_eq!(call.name.as_deref(), Some("get_weather"));
|
||||
assert_eq!(
|
||||
serde_json::from_str::<Value>(&call.arguments).unwrap(),
|
||||
json!({ "city": "Seattle" })
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -729,8 +757,8 @@ mod tests {
|
||||
let mut parser = MinimaxM3ToolParser::new(&m3_test_tools());
|
||||
let output = collect_stream(&mut parser, &["Hello, ", "world!"]);
|
||||
|
||||
assert_eq!(output.normal_text, "Hello, world!");
|
||||
assert!(output.calls.is_empty());
|
||||
assert_eq!(output.normal_text(), "Hello, world!");
|
||||
assert!(output.calls().is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -740,8 +768,8 @@ mod tests {
|
||||
let mut parser = MinimaxM3ToolParser::new(&m3_test_tools());
|
||||
let output = collect_stream(&mut parser, &chunks);
|
||||
|
||||
assert_eq!(output.calls.len(), 1);
|
||||
assert!(output.normal_text.is_empty());
|
||||
assert_eq!(output.calls().len(), 1);
|
||||
assert!(output.normal_text().is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -754,9 +782,9 @@ mod tests {
|
||||
let mut parser = MinimaxM3ToolParser::new(&m3_test_tools());
|
||||
let output = collect_stream(&mut parser, &chunks);
|
||||
|
||||
assert_eq!(output.calls.len(), 2);
|
||||
assert_eq!(output.calls[0].tool_index, 0);
|
||||
assert_eq!(output.calls[1].tool_index, 1);
|
||||
assert_eq!(output.calls().len(), 2);
|
||||
assert_eq!(output.calls()[0].tool_index, 0);
|
||||
assert_eq!(output.calls()[1].tool_index, 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -768,8 +796,8 @@ mod tests {
|
||||
))
|
||||
.unwrap();
|
||||
|
||||
assert!(output.normal_text.is_empty());
|
||||
assert!(output.calls.is_empty());
|
||||
assert!(output.normal_text().is_empty());
|
||||
assert!(output.calls().is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -782,8 +810,8 @@ mod tests {
|
||||
let mut parser = MinimaxM3ToolParser::new(&m3_test_tools());
|
||||
let output = collect_stream(&mut parser, &chunks);
|
||||
|
||||
assert!(output.normal_text.is_empty());
|
||||
assert_eq!(output.calls.len(), 1);
|
||||
assert!(output.normal_text().is_empty());
|
||||
assert_eq!(output.calls().len(), 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -804,8 +832,8 @@ mod tests {
|
||||
parser.parse_chunk(TOOL_CALL_START).unwrap();
|
||||
|
||||
let output = parser.finish().unwrap();
|
||||
assert!(output.normal_text.is_empty());
|
||||
assert!(output.calls.is_empty());
|
||||
assert!(output.normal_text().is_empty());
|
||||
assert!(output.calls().is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -819,9 +847,9 @@ mod tests {
|
||||
))
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(output.calls.len(), 1);
|
||||
assert_eq!(output.calls().len(), 1);
|
||||
assert_eq!(
|
||||
serde_json::from_str::<Value>(&output.calls[0].arguments).unwrap(),
|
||||
serde_json::from_str::<Value>(&output.calls()[0].arguments).unwrap(),
|
||||
json!({ "city": "Seattle" })
|
||||
);
|
||||
}
|
||||
@@ -863,7 +891,7 @@ mod tests {
|
||||
let output = parser.parse_complete(&build_tool_block(&[("convert", body)])).unwrap();
|
||||
|
||||
assert_eq!(
|
||||
serde_json::from_str::<Value>(&output.calls[0].arguments).unwrap(),
|
||||
serde_json::from_str::<Value>(&output.calls()[0].arguments).unwrap(),
|
||||
json!({
|
||||
"payload": {
|
||||
"child": "value",
|
||||
@@ -887,7 +915,7 @@ mod tests {
|
||||
let output = parser.parse_complete(&build_tool_block(&[("convert", body)])).unwrap();
|
||||
|
||||
assert_eq!(
|
||||
serde_json::from_str::<Value>(&output.calls[0].arguments).unwrap(),
|
||||
serde_json::from_str::<Value>(&output.calls()[0].arguments).unwrap(),
|
||||
json!({
|
||||
"payload": {
|
||||
"$text": "child text",
|
||||
@@ -1,9 +1,9 @@
|
||||
//! Streaming tool parsers for chat completions.
|
||||
|
||||
#[macro_use]
|
||||
mod error;
|
||||
pub(crate) mod error;
|
||||
mod deepseek_dsml;
|
||||
mod deepseek_json;
|
||||
pub(crate) mod deepseek_json;
|
||||
mod gemma4;
|
||||
mod glm_xml;
|
||||
mod hy_v3;
|
||||
@@ -15,7 +15,7 @@ mod parameters;
|
||||
mod qwen_coder;
|
||||
#[cfg(any(test, feature = "test-util"))]
|
||||
pub mod test_utils;
|
||||
mod utils;
|
||||
pub(crate) mod utils;
|
||||
|
||||
use std::collections::{BTreeMap, btree_map};
|
||||
|
||||
@@ -57,55 +57,115 @@ pub struct ToolCallDelta {
|
||||
pub arguments: String,
|
||||
}
|
||||
|
||||
/// One ordered event emitted while parsing assistant text.
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub enum ToolParserEvent {
|
||||
/// Plain assistant text that is not part of any tool call.
|
||||
Text(String),
|
||||
/// A tool-call update extracted from assistant text.
|
||||
ToolCall(ToolCallDelta),
|
||||
}
|
||||
|
||||
/// Result of advancing tool parsing with one assistant-text input.
|
||||
#[derive(Debug, Clone, Default, PartialEq, Eq)]
|
||||
pub struct ToolParserOutput {
|
||||
/// Plain assistant text that is not part of any tool call.
|
||||
pub normal_text: String,
|
||||
/// Tool-call updates extracted from this input.
|
||||
pub calls: Vec<ToolCallDelta>,
|
||||
/// Ordered parser events committed by this input.
|
||||
pub events: Vec<ToolParserEvent>,
|
||||
}
|
||||
|
||||
impl ToolParserOutput {
|
||||
/// Append another parser output onto this one.
|
||||
///
|
||||
/// Note that this does not attempt to merge multiple deltas for the same
|
||||
/// tool call into one complete item. Call `coalesce_calls()` after if
|
||||
/// that behavior is desired.
|
||||
pub fn append(&mut self, mut other: Self) {
|
||||
self.normal_text.push_str(&other.normal_text);
|
||||
self.calls.append(&mut other.calls);
|
||||
/// Append one visible text event if `text` is non-empty.
|
||||
pub fn push_text(&mut self, text: impl AsRef<str> + Into<String>) {
|
||||
if text.as_ref().is_empty() {
|
||||
return;
|
||||
}
|
||||
if let Some(ToolParserEvent::Text(last_text)) = self.events.last_mut() {
|
||||
last_text.push_str(text.as_ref());
|
||||
return;
|
||||
}
|
||||
self.events.push(ToolParserEvent::Text(text.into()));
|
||||
}
|
||||
|
||||
/// Merge multiple deltas for the same tool call into one complete item.
|
||||
/// Append one tool-call update event.
|
||||
pub fn push_call(&mut self, call: ToolCallDelta) {
|
||||
self.events.push(ToolParserEvent::ToolCall(call));
|
||||
}
|
||||
|
||||
/// Return all plain assistant text committed by this output.
|
||||
///
|
||||
/// Texts before and after tool calls will be concatenated into a single string. To preserve
|
||||
/// the original order of the text and tool-call events, directly access `events` instead.
|
||||
pub fn normal_text(&self) -> String {
|
||||
self.events
|
||||
.iter()
|
||||
.filter_map(|event| match event {
|
||||
ToolParserEvent::Text(text) => Some(text.as_str()),
|
||||
ToolParserEvent::ToolCall(_) => None,
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Return all tool-call updates committed by this output.
|
||||
pub fn calls(&self) -> Vec<&ToolCallDelta> {
|
||||
self.events
|
||||
.iter()
|
||||
.filter_map(|event| match event {
|
||||
ToolParserEvent::Text(_) => None,
|
||||
ToolParserEvent::ToolCall(call) => Some(call),
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Append another parser output onto this one.
|
||||
///
|
||||
/// Note that this keeps events exactly as they arrive. Call `coalesce()`
|
||||
/// after if final text and tool-call fragments should be flattened.
|
||||
pub fn append(&mut self, other: Self) {
|
||||
for event in other.events {
|
||||
match event {
|
||||
ToolParserEvent::Text(text) => self.push_text(text),
|
||||
ToolParserEvent::ToolCall(call) => self.push_call(call),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Flatten text and merge deltas for the same tool call.
|
||||
///
|
||||
/// All text events are concatenated into one leading text event. Tool-call
|
||||
/// events follow that text event in first-seen tool index order, with
|
||||
/// argument fragments for the same tool call concatenated together.
|
||||
///
|
||||
/// This is primarily used by the default `parse_complete()` implementation,
|
||||
/// which delegates through the incremental parser lifecycle and then
|
||||
/// needs to collapse streaming-style argument fragments into one final
|
||||
/// tool call.
|
||||
pub fn coalesce_calls(mut self) -> Self {
|
||||
pub fn coalesce(self) -> Self {
|
||||
let mut merged = BTreeMap::<usize, ToolCallDelta>::new();
|
||||
let mut order = Vec::new();
|
||||
let normal_text = self.normal_text();
|
||||
|
||||
for call in self.calls {
|
||||
for call in self.calls() {
|
||||
match merged.entry(call.tool_index) {
|
||||
btree_map::Entry::Vacant(entry) => {
|
||||
order.push(call.tool_index);
|
||||
entry.insert(call);
|
||||
entry.insert(call.clone());
|
||||
}
|
||||
btree_map::Entry::Occupied(mut entry) => {
|
||||
let existing = entry.get_mut();
|
||||
if existing.name.is_none() {
|
||||
existing.name = call.name;
|
||||
existing.name = call.name.clone();
|
||||
}
|
||||
existing.arguments.push_str(&call.arguments);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
self.calls =
|
||||
order.into_iter().filter_map(|tool_index| merged.remove(&tool_index)).collect();
|
||||
self
|
||||
let mut output = Self::default();
|
||||
output.push_text(normal_text);
|
||||
for call in order.into_iter().filter_map(|tool_index| merged.remove(&tool_index)) {
|
||||
output.push_call(call);
|
||||
}
|
||||
output
|
||||
}
|
||||
}
|
||||
|
||||
@@ -183,7 +243,7 @@ impl<T: ToolParser + ?Sized> T {
|
||||
pub fn parse_complete(&mut self, text: &str) -> Result<ToolParserOutput> {
|
||||
let mut output = self.parse_chunk(text)?;
|
||||
output.append(self.finish()?);
|
||||
Ok(output.coalesce_calls())
|
||||
Ok(output.coalesce())
|
||||
}
|
||||
}
|
||||
|
||||
@@ -2,7 +2,7 @@ use std::collections::BTreeMap;
|
||||
|
||||
use serde_json::{Map, Number, Value};
|
||||
|
||||
use crate::Tool;
|
||||
use crate::tool::Tool;
|
||||
|
||||
/// Normalized parameter schemas for all tools in one request.
|
||||
#[derive(Debug, Clone, Default, PartialEq, Eq)]
|
||||
@@ -416,7 +416,7 @@ mod tests {
|
||||
use serde_json::{Value, json};
|
||||
|
||||
use super::{ParamElement, ParamInput, ToolSchema, ToolSchemas};
|
||||
use crate::Tool;
|
||||
use crate::tool::Tool;
|
||||
|
||||
fn test_tool(name: &str, parameters: serde_json::Value) -> Tool {
|
||||
Tool {
|
||||
@@ -7,7 +7,7 @@ use winnow::token::{literal, take_until};
|
||||
use super::parameters::ToolSchemas;
|
||||
use super::utils::{MarkerScanState, parse_buffered_event, safe_text_len, take_until_marker};
|
||||
use super::{Result, StructuralTagModel, ToolCallDelta, ToolParser, ToolParserOutput};
|
||||
use crate::Tool;
|
||||
use crate::tool::Tool;
|
||||
|
||||
const TOOL_CALL_START: &str = "<tool_call>";
|
||||
const TOOL_CALL_END: &str = "</tool_call>";
|
||||
@@ -74,7 +74,7 @@ impl Qwen3CoderToolParser {
|
||||
fn apply_event(&mut self, event: QwenCoderEvent, output: &mut ToolParserOutput) -> Result<()> {
|
||||
match event {
|
||||
QwenCoderEvent::Text { len: consumed_len } => {
|
||||
output.normal_text.push_str(&self.buffer[..consumed_len]);
|
||||
output.push_text(&self.buffer[..consumed_len]);
|
||||
}
|
||||
QwenCoderEvent::ToolCallStart => {
|
||||
self.mode = QwenCoderMode::ToolCall {
|
||||
@@ -87,7 +87,7 @@ impl Qwen3CoderToolParser {
|
||||
let arguments = serde_json::to_string(&arguments)
|
||||
.map_err(|error| parsing_failed!("failed to serialize arguments: {}", error))?;
|
||||
|
||||
output.calls.push(ToolCallDelta {
|
||||
output.push_call(ToolCallDelta {
|
||||
tool_index: self.emitted_tool_count,
|
||||
name: Some(name),
|
||||
arguments,
|
||||
@@ -138,7 +138,7 @@ impl ToolParser for Qwen3CoderToolParser {
|
||||
{
|
||||
return Err(parsing_failed!("incomplete Qwen Coder tool call"));
|
||||
}
|
||||
output.normal_text.push_str(&self.buffer);
|
||||
output.push_text(&self.buffer);
|
||||
}
|
||||
let _ = self.reset();
|
||||
Ok(output)
|
||||
@@ -241,8 +241,8 @@ mod tests {
|
||||
use thiserror_ext::AsReport;
|
||||
|
||||
use super::{Qwen3CoderToolParser, StructuralTagModel, ToolParser};
|
||||
use crate::test_utils::{collect_stream, split_by_chars, test_tools};
|
||||
use crate::{ToolParserOutput, ToolParserTestExt as _};
|
||||
use crate::tool::test_utils::{collect_stream, split_by_chars, test_tools};
|
||||
use crate::tool::{ToolParserOutput, ToolParserTestExt as _};
|
||||
|
||||
fn build_tool_call(function_name: &str, params: &[(&str, &str)]) -> String {
|
||||
let params = params
|
||||
@@ -268,8 +268,8 @@ mod tests {
|
||||
let mut parser = Qwen3CoderToolParser::new(&test_tools());
|
||||
let output = parser.parse_complete("Hello, world!").unwrap();
|
||||
|
||||
assert_eq!(output.normal_text, "Hello, world!");
|
||||
assert!(output.calls.is_empty());
|
||||
assert_eq!(output.normal_text(), "Hello, world!");
|
||||
assert!(output.calls().is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -282,11 +282,11 @@ mod tests {
|
||||
))
|
||||
.unwrap();
|
||||
|
||||
assert!(output.normal_text.is_empty());
|
||||
assert_eq!(output.calls.len(), 1);
|
||||
assert_eq!(output.calls[0].name.as_deref(), Some("get_weather"));
|
||||
assert!(output.normal_text().is_empty());
|
||||
assert_eq!(output.calls().len(), 1);
|
||||
assert_eq!(output.calls()[0].name.as_deref(), Some("get_weather"));
|
||||
assert_eq!(
|
||||
serde_json::from_str::<Value>(&output.calls[0].arguments).unwrap(),
|
||||
serde_json::from_str::<Value>(&output.calls()[0].arguments).unwrap(),
|
||||
json!({
|
||||
"location": "SF",
|
||||
"date": "2026-04-29"
|
||||
@@ -303,8 +303,8 @@ mod tests {
|
||||
);
|
||||
let output = parser.parse_complete(&output).unwrap();
|
||||
|
||||
assert_eq!(output.normal_text, "Thinking... ");
|
||||
assert_eq!(output.calls.len(), 1);
|
||||
assert_eq!(output.normal_text(), "Thinking... ");
|
||||
assert_eq!(output.calls().len(), 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -323,9 +323,9 @@ mod tests {
|
||||
))
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(output.calls.len(), 1);
|
||||
assert_eq!(output.calls().len(), 1);
|
||||
assert_eq!(
|
||||
serde_json::from_str::<Value>(&output.calls[0].arguments).unwrap(),
|
||||
serde_json::from_str::<Value>(&output.calls()[0].arguments).unwrap(),
|
||||
json!({
|
||||
"whole": 5.0,
|
||||
"flag": true,
|
||||
@@ -341,10 +341,10 @@ mod tests {
|
||||
let mut parser = Qwen3CoderToolParser::new(&test_tools());
|
||||
let output = parser.parse_complete(&build_tool_call("get_weather", &[])).unwrap();
|
||||
|
||||
assert_eq!(output.calls.len(), 1);
|
||||
assert_eq!(output.calls[0].name.as_deref(), Some("get_weather"));
|
||||
assert_eq!(output.calls().len(), 1);
|
||||
assert_eq!(output.calls()[0].name.as_deref(), Some("get_weather"));
|
||||
assert_eq!(
|
||||
serde_json::from_str::<Value>(&output.calls[0].arguments).unwrap(),
|
||||
serde_json::from_str::<Value>(&output.calls()[0].arguments).unwrap(),
|
||||
json!({})
|
||||
);
|
||||
}
|
||||
@@ -371,10 +371,10 @@ mod tests {
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(output.calls.len(), 1);
|
||||
assert_eq!(output.calls[0].name.as_deref(), Some("calculate_area"));
|
||||
assert_eq!(output.calls().len(), 1);
|
||||
assert_eq!(output.calls()[0].name.as_deref(), Some("calculate_area"));
|
||||
assert_eq!(
|
||||
serde_json::from_str::<Value>(&output.calls[0].arguments).unwrap(),
|
||||
serde_json::from_str::<Value>(&output.calls()[0].arguments).unwrap(),
|
||||
json!({
|
||||
"shape": "rectangle",
|
||||
"dimensions": { "width": 10, "height": 20 },
|
||||
@@ -396,9 +396,9 @@ mod tests {
|
||||
))
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(output.calls.len(), 1);
|
||||
assert_eq!(output.calls().len(), 1);
|
||||
assert_eq!(
|
||||
serde_json::from_str::<Value>(&output.calls[0].arguments).unwrap(),
|
||||
serde_json::from_str::<Value>(&output.calls()[0].arguments).unwrap(),
|
||||
json!({
|
||||
"payload": {
|
||||
"nested": {
|
||||
@@ -426,9 +426,9 @@ mod tests {
|
||||
))
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(output.calls.len(), 1);
|
||||
assert_eq!(output.calls().len(), 1);
|
||||
assert_eq!(
|
||||
serde_json::from_str::<Value>(&output.calls[0].arguments).unwrap(),
|
||||
serde_json::from_str::<Value>(&output.calls()[0].arguments).unwrap(),
|
||||
json!({
|
||||
"html_content": r#"<div class="test"><span>Hello</span></div>"#,
|
||||
"xml_snippet": r#"<root><child attr="value"/></root>"#,
|
||||
@@ -452,9 +452,9 @@ mod tests {
|
||||
))
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(output.calls.len(), 1);
|
||||
assert_eq!(output.calls().len(), 1);
|
||||
assert_eq!(
|
||||
serde_json::from_str::<Value>(&output.calls[0].arguments).unwrap(),
|
||||
serde_json::from_str::<Value>(&output.calls()[0].arguments).unwrap(),
|
||||
json!({
|
||||
"location": "杭州 </parameter></function></tool_call>",
|
||||
"date": "2026-05-08",
|
||||
@@ -472,9 +472,9 @@ mod tests {
|
||||
))
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(output.calls.len(), 1);
|
||||
assert_eq!(output.calls().len(), 1);
|
||||
assert_eq!(
|
||||
serde_json::from_str::<Value>(&output.calls[0].arguments).unwrap(),
|
||||
serde_json::from_str::<Value>(&output.calls()[0].arguments).unwrap(),
|
||||
json!({
|
||||
"data": { "key": "value", "count": 42 },
|
||||
})
|
||||
@@ -495,11 +495,11 @@ mod tests {
|
||||
],
|
||||
);
|
||||
|
||||
assert!(output.normal_text.is_empty());
|
||||
assert_eq!(output.calls.len(), 1);
|
||||
assert_eq!(output.calls[0].name.as_deref(), Some("get_weather"));
|
||||
assert!(output.normal_text().is_empty());
|
||||
assert_eq!(output.calls().len(), 1);
|
||||
assert_eq!(output.calls()[0].name.as_deref(), Some("get_weather"));
|
||||
assert_eq!(
|
||||
serde_json::from_str::<Value>(&output.calls[0].arguments).unwrap(),
|
||||
serde_json::from_str::<Value>(&output.calls()[0].arguments).unwrap(),
|
||||
json!({ "location": "SF" })
|
||||
);
|
||||
}
|
||||
@@ -519,8 +519,8 @@ mod tests {
|
||||
],
|
||||
);
|
||||
|
||||
assert_eq!(output.normal_text, "Thinking... ");
|
||||
assert_eq!(output.calls.len(), 1);
|
||||
assert_eq!(output.normal_text(), "Thinking... ");
|
||||
assert_eq!(output.calls().len(), 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -528,8 +528,8 @@ mod tests {
|
||||
let mut parser = Qwen3CoderToolParser::new(&test_tools());
|
||||
let output = collect_stream(&mut parser, &["Hello, ", "world!"]);
|
||||
|
||||
assert_eq!(output.normal_text, "Hello, world!");
|
||||
assert!(output.calls.is_empty());
|
||||
assert_eq!(output.normal_text(), "Hello, world!");
|
||||
assert!(output.calls().is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -542,17 +542,17 @@ mod tests {
|
||||
let mut parser = Qwen3CoderToolParser::new(&test_tools());
|
||||
let output = collect_stream(&mut parser, &[&text]);
|
||||
|
||||
assert_eq!(output.calls.len(), 2);
|
||||
assert_eq!(output.calls[0].name.as_deref(), Some("get_weather"));
|
||||
assert_eq!(output.calls[1].name.as_deref(), Some("get_weather"));
|
||||
assert_eq!(output.calls[0].tool_index, 0);
|
||||
assert_eq!(output.calls[1].tool_index, 1);
|
||||
assert_eq!(output.calls().len(), 2);
|
||||
assert_eq!(output.calls()[0].name.as_deref(), Some("get_weather"));
|
||||
assert_eq!(output.calls()[1].name.as_deref(), Some("get_weather"));
|
||||
assert_eq!(output.calls()[0].tool_index, 0);
|
||||
assert_eq!(output.calls()[1].tool_index, 1);
|
||||
assert_eq!(
|
||||
serde_json::from_str::<Value>(&output.calls[0].arguments).unwrap(),
|
||||
serde_json::from_str::<Value>(&output.calls()[0].arguments).unwrap(),
|
||||
json!({ "location": "SF" })
|
||||
);
|
||||
assert_eq!(
|
||||
serde_json::from_str::<Value>(&output.calls[1].arguments).unwrap(),
|
||||
serde_json::from_str::<Value>(&output.calls()[1].arguments).unwrap(),
|
||||
json!({ "location": "NYC" })
|
||||
);
|
||||
}
|
||||
@@ -569,16 +569,16 @@ mod tests {
|
||||
let output = collect_stream(&mut parser, &chunks);
|
||||
|
||||
assert_eq!(
|
||||
output.normal_text,
|
||||
output.normal_text(),
|
||||
"I'll check two cities.Between calls.Done."
|
||||
);
|
||||
assert_eq!(output.calls.len(), 2);
|
||||
assert_eq!(output.calls().len(), 2);
|
||||
assert_eq!(
|
||||
serde_json::from_str::<Value>(&output.calls[0].arguments).unwrap(),
|
||||
serde_json::from_str::<Value>(&output.calls()[0].arguments).unwrap(),
|
||||
json!({ "city": "Dallas", "state": "TX" })
|
||||
);
|
||||
assert_eq!(
|
||||
serde_json::from_str::<Value>(&output.calls[1].arguments).unwrap(),
|
||||
serde_json::from_str::<Value>(&output.calls()[1].arguments).unwrap(),
|
||||
json!({ "city": "Orlando", "state": "FL" })
|
||||
);
|
||||
}
|
||||
@@ -590,9 +590,9 @@ mod tests {
|
||||
let mut parser = Qwen3CoderToolParser::new(&test_tools());
|
||||
let output = collect_stream(&mut parser, &chunks);
|
||||
|
||||
assert_eq!(output.calls.len(), 1);
|
||||
assert_eq!(output.calls().len(), 1);
|
||||
assert_eq!(
|
||||
serde_json::from_str::<Value>(&output.calls[0].arguments).unwrap(),
|
||||
serde_json::from_str::<Value>(&output.calls()[0].arguments).unwrap(),
|
||||
json!({ "location": "SF" })
|
||||
);
|
||||
}
|
||||
@@ -610,19 +610,19 @@ mod tests {
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
assert!(output.normal_text.is_empty());
|
||||
assert!(output.calls.is_empty());
|
||||
assert!(output.normal_text().is_empty());
|
||||
assert!(output.calls().is_empty());
|
||||
|
||||
let mut output = output;
|
||||
output.append(parser.parse_chunk("_call>").unwrap());
|
||||
output.append(parser.finish().unwrap());
|
||||
let output = output.coalesce_calls();
|
||||
let output = output.coalesce();
|
||||
|
||||
assert!(output.normal_text.is_empty());
|
||||
assert_eq!(output.calls.len(), 1);
|
||||
assert_eq!(output.calls[0].name.as_deref(), Some("get_weather"));
|
||||
assert!(output.normal_text().is_empty());
|
||||
assert_eq!(output.calls().len(), 1);
|
||||
assert_eq!(output.calls()[0].name.as_deref(), Some("get_weather"));
|
||||
assert_eq!(
|
||||
serde_json::from_str::<Value>(&output.calls[0].arguments).unwrap(),
|
||||
serde_json::from_str::<Value>(&output.calls()[0].arguments).unwrap(),
|
||||
json!({ "location": "SF" })
|
||||
);
|
||||
}
|
||||
@@ -641,20 +641,20 @@ mod tests {
|
||||
|
||||
for chunk in chunks {
|
||||
let chunk_output = parser.parse_chunk(chunk).unwrap();
|
||||
assert!(chunk_output.normal_text.is_empty());
|
||||
assert!(chunk_output.calls.is_empty());
|
||||
assert!(chunk_output.normal_text().is_empty());
|
||||
assert!(chunk_output.calls().is_empty());
|
||||
output.append(chunk_output);
|
||||
}
|
||||
|
||||
output.append(parser.parse_chunk(end_suffix).unwrap());
|
||||
output.append(parser.finish().unwrap());
|
||||
let output = output.coalesce_calls();
|
||||
let output = output.coalesce();
|
||||
|
||||
assert!(output.normal_text.is_empty());
|
||||
assert_eq!(output.calls.len(), 1);
|
||||
assert_eq!(output.calls[0].name.as_deref(), Some("get_weather"));
|
||||
assert!(output.normal_text().is_empty());
|
||||
assert_eq!(output.calls().len(), 1);
|
||||
assert_eq!(output.calls()[0].name.as_deref(), Some("get_weather"));
|
||||
assert_eq!(
|
||||
serde_json::from_str::<Value>(&output.calls[0].arguments).unwrap(),
|
||||
serde_json::from_str::<Value>(&output.calls()[0].arguments).unwrap(),
|
||||
json!({ "location": long_location })
|
||||
);
|
||||
}
|
||||
@@ -666,8 +666,8 @@ mod tests {
|
||||
.parse_chunk("<tool_call>\n<function=get_weather>\n<parameter=location>SF</parameter>")
|
||||
.unwrap();
|
||||
|
||||
assert!(output.normal_text.is_empty());
|
||||
assert!(output.calls.is_empty());
|
||||
assert!(output.normal_text().is_empty());
|
||||
assert!(output.calls().is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -710,7 +710,7 @@ mod tests {
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
serde_json::from_str::<Value>(&output.calls[0].arguments).unwrap(),
|
||||
serde_json::from_str::<Value>(&output.calls()[0].arguments).unwrap(),
|
||||
json!({ "location": "Hangzhou" })
|
||||
);
|
||||
}
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user