| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Granite parser for JSON-array tool calls. |
| 4 | |
| 5 | Granite 3.0/3.1 emit tool calls as a JSON array following a marker:: |
| 6 | |
| 7 | <|tool_call|> [{"name": "get_weather", "arguments": {"city": "SF"}}] # 3.0 |
| 8 | <tool_call> [{"name": "get_weather", "arguments": {"city": "SF"}}] # 3.1 |
| 9 | |
| 10 | ``<|tool_call|>`` is a single special token; ``<tool_call>`` is plain text. |
| 11 | The array may hold several calls, has no closing marker (the ``]`` ends it), |
| 12 | and any surrounding prose is content. The engine's ``tool_call_body_array`` |
| 13 | mode splits the array into one call per element and extracts each name via the |
| 14 | name-from-args path; ``_granite_arg_converter`` carves the ``arguments`` object |
| 15 | out of each ``{"name":..,"arguments":{..}}`` element. Granite has no reasoning. |
| 16 | """ |
| 17 | |
| 18 | from __future__ import annotations |
| 19 | |
| 20 | import functools |
| 21 | from typing import TYPE_CHECKING |
| 22 | |
| 23 | from vllm.parser.engine.parser_engine import ParserEngine |
| 24 | from vllm.parser.engine.parser_engine_config import ( |
| 25 | ParserEngineConfig, |
| 26 | ParserState, |
| 27 | Transition, |
| 28 | ) |
| 29 | from vllm.parser.utils import extract_json_member_object |
| 30 | |
| 31 | if TYPE_CHECKING: |
| 32 | from vllm.tokenizers import TokenizerLike |
| 33 | from vllm.tool_parsers.abstract_tool_parser import Tool |
| 34 | |
| 35 | TOOL_TOKEN = "<|tool_call|>" # granite 3.0 (single special token) |
| 36 | TOOL_STRING = "<tool_call>" # granite 3.1 (plain text) |
| 37 | |
| 38 | |
| 39 | def _granite_arg_converter(raw_args: str, partial: bool) -> str: |
| 40 | """Carve the ``arguments`` object out of a ``{"name":..,"arguments":{..}}`` |
| 41 | array element (see :func:`extract_json_member_object`).""" |
| 42 | span = extract_json_member_object(raw_args, "arguments") |
| 43 | if span is None: |
| 44 | return "" if partial else "{}" |
| 45 | return span |
| 46 | |
| 47 | |
| 48 | @functools.cache |
| 49 | def granite_config() -> ParserEngineConfig: |
| 50 | return ParserEngineConfig( |
| 51 | name="granite", |
| 52 | initial_state=ParserState.CONTENT, |
| 53 | terminals={ |
| 54 | "TOOL_TOKEN": TOOL_TOKEN, |
| 55 | "TOOL_STRING": TOOL_STRING, |
| 56 | }, |
| 57 | token_id_terminals={ |
| 58 | "TOOL_TOKEN": TOOL_TOKEN, |
| 59 | }, |
| 60 | transitions={ |
| 61 | # The marker only switches into the tool region; the array splitter |
| 62 | # emits TOOL_CALL_START per element (empty events here). |
| 63 | (ParserState.CONTENT, "TOOL_TOKEN"): Transition( |
| 64 | ParserState.TOOL_ARGS, |
| 65 | (), |
| 66 | ), |
| 67 | (ParserState.CONTENT, "TOOL_STRING"): Transition( |
| 68 | ParserState.TOOL_ARGS, |
| 69 | (), |
| 70 | ), |
| 71 | }, |
| 72 | arg_converter=_granite_arg_converter, |
| 73 | tool_call_body_array=True, |
| 74 | stream_arg_deltas=True, |
| 75 | ) |
| 76 | |
| 77 | |
| 78 | class GraniteParser(ParserEngine): |
| 79 | CONFIG_NAME = "granite" |
| 80 | |
| 81 | def __init__( |
| 82 | self, |
| 83 | tokenizer: TokenizerLike, |
| 84 | tools: list[Tool] | None = None, |
| 85 | **kwargs, |
| 86 | ) -> None: |
| 87 | kwargs.setdefault("parser_engine_config", granite_config()) |
| 88 | super().__init__(tokenizer, tools, **kwargs) |