| |
|
|
| import ast |
| import json |
| import re |
| from collections.abc import Sequence |
| from typing import Union |
|
|
| import partial_json_parser |
| from partial_json_parser.core.options import Allow |
|
|
| from vllm.entrypoints.openai.protocol import ( |
| ChatCompletionRequest, |
| DeltaFunctionCall, DeltaMessage, |
| DeltaToolCall, |
| ExtractedToolCallInformation, |
| FunctionCall, |
| ToolCall, |
| ) |
| from vllm.entrypoints.openai.tool_parsers.abstract_tool_parser import ( |
| ToolParser, |
| ToolParserManager, |
| ) |
| from vllm.logger import init_logger |
| from vllm.transformers_utils.tokenizer import AnyTokenizer |
| from vllm.utils import random_uuid |
|
|
| logger = init_logger(__name__) |
|
|
|
|
| @ToolParserManager.register_module("nemotron_json") |
| class NemotronJSONToolParser(ToolParser): |
|
|
| def __init__(self, tokenizer: AnyTokenizer): |
| super().__init__(tokenizer) |
|
|
| self.current_tool_name_sent: bool = False |
| self.prev_tool_call_arr: list[dict] = [] |
| self.current_tool_id: int = -1 |
| self.streamed_args_for_tool: list[str] = [] |
|
|
| self.tool_call_start_token: str = "<TOOLCALL>" |
| self.tool_call_end_token: str = "</TOOLCALL>" |
|
|
| self.tool_call_regex = re.compile(r"<TOOLCALL>(.*?)</TOOLCALL>", re.DOTALL) |
|
|
| def extract_tool_calls( |
| self, |
| model_output: str, |
| request: ChatCompletionRequest, |
| ) -> ExtractedToolCallInformation: |
|
|
| if self.tool_call_start_token not in model_output: |
| return ExtractedToolCallInformation( |
| tools_called=False, |
| tool_calls=[], |
| content=model_output, |
| ) |
|
|
| else: |
|
|
| try: |
| str_tool_calls = self.tool_call_regex.findall(model_output)[0].strip() |
| if not str_tool_calls.startswith("["): |
| str_tool_calls = "[" + str_tool_calls |
| if not str_tool_calls.endswith("]"): |
| str_tool_calls = "]" + str_tool_calls |
| json_tool_calls = json.loads(str_tool_calls) |
| tool_calls = [] |
| for tool_call in json_tool_calls: |
| try: |
| tool_calls.append(ToolCall( |
| type="function", |
| function=FunctionCall( |
| name=tool_call["name"], |
| arguments=json.dumps(tool_call["arguments"], ensure_ascii=False) \ |
| if isinstance(tool_call["arguments"], dict) else tool_call["arguments"], |
| ), |
| )) |
| except: |
| continue |
|
|
| content = model_output[:model_output.rfind(self.tool_call_start_token)] |
|
|
| return ExtractedToolCallInformation( |
| tools_called=True, |
| tool_calls=tool_calls, |
| content=content if content else None, |
| ) |
|
|
| except Exception: |
| logger.exception(f"Error in extracting tool call from response. Response: {model_output}") |
| return ExtractedToolCallInformation( |
| tools_called=False, |
| tool_calls=[], |
| content=model_output, |
| ) |
|
|
| def extract_tool_calls_streaming( |
| self, |
| previous_text: str, |
| current_text: str, |
| delta_text: str, |
| previous_token_ids: Sequence[int], |
| current_token_ids: Sequence[int], |
| delta_token_ids: Sequence[int], |
| request: ChatCompletionRequest, |
| ) -> Union[DeltaMessage, None]: |
|
|
| raise NotImplementedError("Tool calling is not supported in streaming mode!") |
|
|
