| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import json |
| 4 | import random |
| 5 | from typing import Any |
| 6 | |
| 7 | import pytest |
| 8 | from transformers import AutoTokenizer |
| 9 | |
| 10 | from vllm.entrypoints.generate.base.protocol import ( |
| 11 | DeltaMessage, |
| 12 | ) |
| 13 | from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest |
| 14 | from vllm.tool_parsers.granite4_tool_parser import Granite4ToolParser |
| 15 | |
| 16 | MODEL = "ibm-granite/granite-4.0-h-tiny" |
| 17 | |
| 18 | |
| 19 | def create_complex_input(create_string_args: bool): |
| 20 | coord_arg: dict | str = { |
| 21 | "coordinates": [[23.54, 43.1], [-12.2, 54.3], [4, 5]], |
| 22 | "coordinate_type": "latlong", |
| 23 | } |
| 24 | if create_string_args: |
| 25 | # test granite behavior |
| 26 | coord_arg = json.dumps(coord_arg) |
| 27 | return [ |
| 28 | {"name": "find_bbox", "arguments": coord_arg}, |
| 29 | { |
| 30 | "name": "get_stock_price", |
| 31 | "arguments": { |
| 32 | "symbol": "AAPL", |
| 33 | "start_date": "2021-01-01", |
| 34 | "end_date": "2021-12-31", |
| 35 | }, |
| 36 | }, |
| 37 | {"name": "find_bbox", "arguments": coord_arg}, |
| 38 | ] |
| 39 | |
| 40 | |
| 41 | def random_chunks(s: str, min_len: int, max_len: int): |
| 42 | chunks = [] |
| 43 | i = 0 |
| 44 | n = len(s) |
| 45 | |
| 46 | while i < n: |
| 47 | size = random.randint(min_len, max_len) |
| 48 | chunks.append(s[i : i + size]) |
| 49 | i += size |
| 50 | |
| 51 | return chunks |
| 52 | |
| 53 | |
| 54 | @pytest.fixture(scope="module") |
| 55 | def tokenizer(): |
| 56 | return AutoTokenizer.from_pretrained(MODEL) |
| 57 | |
| 58 | |
| 59 | # create a variety of input chunk sizes |
| 60 | @pytest.mark.parametrize( |
| 61 | "min_chunk, max_chunk", |
| 62 | [ |
| 63 | (1, 1), |
| 64 | (1, 2), |
| 65 | (5, 7), |
| 66 | (6, 20), |
| 67 | ], |
| 68 | ) |
| 69 | def test_tool_call_parser_complex(min_chunk: int, max_chunk: int, tokenizer): |
| 70 | input_dicts = create_complex_input(True) |
| 71 | |
| 72 | formatted_tcs = [ |
| 73 | "<tool_call> " + json.dumps(call) + " </tool_call>" for call in input_dicts |
| 74 | ] |
| 75 | |
| 76 | text_messages = [ |
| 77 | "Here goes the bbox call: \n", |
| 78 | " Now the stock price call: \n ", |
| 79 | " Now another bbox call: \n ", |
| 80 | " See? I'm a helpful assistant.", |
| 81 | ] |
| 82 | |
| 83 | test_input = ( |
| 84 | text_messages[0] |
| 85 | + formatted_tcs[0] |
| 86 | + text_messages[1] |
| 87 | + formatted_tcs[1] |
| 88 | + text_messages[2] |
| 89 | + formatted_tcs[2] |
| 90 | + text_messages[3] |
| 91 | ) |
| 92 | |
| 93 | any_chat_request = ChatCompletionRequest( |
| 94 | seed=42, |
| 95 | model=MODEL, |
| 96 | messages=[], |
| 97 | ) |
| 98 | |
| 99 | parser = Granite4ToolParser(tokenizer=tokenizer) |
| 100 | |
| 101 | delta_messages = list[DeltaMessage]() |
| 102 | for text in random_chunks(test_input, min_chunk, max_chunk): |
| 103 | delta = parser.extract_tool_calls_streaming( |
| 104 | previous_text="", |
| 105 | current_text="", |
| 106 | delta_text=text, |
| 107 | previous_token_ids=[], |
| 108 | current_token_ids=[], |
| 109 | delta_token_ids=[], |
| 110 | request=any_chat_request, |
| 111 | ) |
| 112 | if delta is not None: |
| 113 | delta_messages.append(delta) |
| 114 | |
| 115 | content = "" |
| 116 | tool_calls = list[dict[str, Any]]() |
| 117 | |
| 118 | current_name = "__start__" |
| 119 | current_args = "" |
| 120 | |
| 121 | for msg in delta_messages: |
| 122 | if msg.content: |
| 123 | content += msg.content |
| 124 | for tool_call in msg.tool_calls: |
| 125 | if delta_func := tool_call.function: |
| 126 | if delta_func.name is not None: |
| 127 | if current_name == "__start__": |
| 128 | current_name = delta_func.name |
| 129 | |
| 130 | if delta_func.name != current_name: |
| 131 | tool_calls.append( |
| 132 | { |
| 133 | "name": current_name, |
| 134 | "arguments": json.loads(current_args), |
| 135 | } |
| 136 | ) |
| 137 | current_name = delta_func.name |
| 138 | current_args = "" |
| 139 | |
| 140 | if delta_func.arguments: |
| 141 | current_args += delta_func.arguments |
| 142 | |
| 143 | if current_name != "__start__": |
| 144 | tool_calls.append({"name": current_name, "arguments": json.loads(current_args)}) |
| 145 | |
| 146 | assert content == "".join(text_messages) |
| 147 | assert tool_calls == create_complex_input(False) |