15ed1262e7
Signed-off-by: Nikhil Kulkarni <nikhilkulkarni1755@gmail.com> Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com> Co-authored-by: Chauncey <chaunceyjiang@gmail.com>
89 lines
2.9 KiB
Python
89 lines
2.9 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""Granite parser for JSON-array tool calls.
|
|
|
|
Granite 3.0/3.1 emit tool calls as a JSON array following a marker::
|
|
|
|
<|tool_call|> [{"name": "get_weather", "arguments": {"city": "SF"}}] # 3.0
|
|
<tool_call> [{"name": "get_weather", "arguments": {"city": "SF"}}] # 3.1
|
|
|
|
``<|tool_call|>`` is a single special token; ``<tool_call>`` is plain text.
|
|
The array may hold several calls, has no closing marker (the ``]`` ends it),
|
|
and any surrounding prose is content. The engine's ``tool_call_body_array``
|
|
mode splits the array into one call per element and extracts each name via the
|
|
name-from-args path; ``_granite_arg_converter`` carves the ``arguments`` object
|
|
out of each ``{"name":..,"arguments":{..}}`` element. Granite has no reasoning.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import functools
|
|
from typing import TYPE_CHECKING
|
|
|
|
from vllm.parser.engine.parser_engine import ParserEngine
|
|
from vllm.parser.engine.parser_engine_config import (
|
|
ParserEngineConfig,
|
|
ParserState,
|
|
Transition,
|
|
)
|
|
from vllm.parser.utils import extract_json_member_object
|
|
|
|
if TYPE_CHECKING:
|
|
from vllm.tokenizers import TokenizerLike
|
|
from vllm.tool_parsers.abstract_tool_parser import Tool
|
|
|
|
TOOL_TOKEN = "<|tool_call|>" # granite 3.0 (single special token)
|
|
TOOL_STRING = "<tool_call>" # granite 3.1 (plain text)
|
|
|
|
|
|
def _granite_arg_converter(raw_args: str, partial: bool) -> str:
|
|
"""Carve the ``arguments`` object out of a ``{"name":..,"arguments":{..}}``
|
|
array element (see :func:`extract_json_member_object`)."""
|
|
span = extract_json_member_object(raw_args, "arguments")
|
|
if span is None:
|
|
return "" if partial else "{}"
|
|
return span
|
|
|
|
|
|
@functools.cache
|
|
def granite_config() -> ParserEngineConfig:
|
|
return ParserEngineConfig(
|
|
name="granite",
|
|
initial_state=ParserState.CONTENT,
|
|
terminals={
|
|
"TOOL_TOKEN": TOOL_TOKEN,
|
|
"TOOL_STRING": TOOL_STRING,
|
|
},
|
|
token_id_terminals={
|
|
"TOOL_TOKEN": TOOL_TOKEN,
|
|
},
|
|
transitions={
|
|
# The marker only switches into the tool region; the array splitter
|
|
# emits TOOL_CALL_START per element (empty events here).
|
|
(ParserState.CONTENT, "TOOL_TOKEN"): Transition(
|
|
ParserState.TOOL_ARGS,
|
|
(),
|
|
),
|
|
(ParserState.CONTENT, "TOOL_STRING"): Transition(
|
|
ParserState.TOOL_ARGS,
|
|
(),
|
|
),
|
|
},
|
|
arg_converter=_granite_arg_converter,
|
|
tool_call_body_array=True,
|
|
stream_arg_deltas=True,
|
|
)
|
|
|
|
|
|
class GraniteParser(ParserEngine):
|
|
CONFIG_NAME = "granite"
|
|
|
|
def __init__(
|
|
self,
|
|
tokenizer: TokenizerLike,
|
|
tools: list[Tool] | None = None,
|
|
**kwargs,
|
|
) -> None:
|
|
kwargs.setdefault("parser_engine_config", granite_config())
|
|
super().__init__(tokenizer, tools, **kwargs)
|