import json from collections.abc import Sequence from typing import Any, Optional from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest from vllm.entrypoints.openai.engine.protocol import ( DeltaFunctionCall, DeltaMessage, DeltaToolCall, ExtractedToolCallInformation, FunctionCall, ToolCall, ) from vllm.tokenizers import TokenizerLike from vllm.tool_parsers.abstract_tool_parser import ToolParser, ToolParserManager @ToolParserManager.register_module(["openpipe_llama_dual"]) class OpenPipeLlamaDualParser(ToolParser): """Parse official JSON, llama31 tool markers, and pipeline3 function tags.""" LEGACY_START = "<|start_tool_call|>" LEGACY_END = "<|end_tool_call|>" FUNCTION_CALL_TAG = "" FUNCTION_ARGS_TAG = "" VARIANT_LLAMA31 = "llama31instruct" VARIANT_PIPELINE3 = "pipeline3" VARIANT_OFFICIAL = "official" def __init__(self, tokenizer: TokenizerLike, tools): super().__init__(tokenizer, tools) self.tokenizer = tokenizer self.tools = tools def _get_template_variant(self, request: ChatCompletionRequest) -> Optional[str]: kwargs = getattr(request, "chat_template_kwargs", None) if kwargs is None: return None if isinstance(kwargs, dict): value = kwargs.get("template_variant") return value if isinstance(value, str) else None value = getattr(kwargs, "template_variant", None) return value if isinstance(value, str) else None def _normalize_tool_call(self, payload: dict[str, Any]) -> Optional[dict[str, Any]]: if "name" in payload and "parameters" in payload: return { "name": payload["name"], "arguments": payload["parameters"], } if "function" in payload and isinstance(payload["function"], dict): function = payload["function"] if "name" in function and "arguments" in function: return { "name": function["name"], "arguments": function["arguments"], } return None def _extract_legacy_tool_calls(self, text: str) -> list[dict[str, Any]]: tool_calls = [] current_index = 0 while True: start_index = text.find(self.LEGACY_START, current_index) if start_index == -1: break end_index = text.find(self.LEGACY_END, start_index) if end_index == -1: break tool_call_json = text[start_index + len(self.LEGACY_START) : end_index].strip() payload = json.loads(tool_call_json) normalized = self._normalize_tool_call(payload) if normalized: tool_calls.append(normalized) current_index = end_index + len(self.LEGACY_END) return tool_calls def _extract_function_tag_tool_calls(self, text: str) -> list[dict[str, Any]]: tool_calls = [] current_index = 0 while True: function_start = text.find(self.FUNCTION_CALL_TAG, current_index) if function_start == -1: break name_start = function_start + len(self.FUNCTION_CALL_TAG) args_tag_index = text.find(self.FUNCTION_ARGS_TAG, name_start) if args_tag_index == -1: break function_name = text[name_start:args_tag_index].strip() if not function_name: break arguments_start = args_tag_index + len(self.FUNCTION_ARGS_TAG) next_function_index = text.find(self.FUNCTION_CALL_TAG, arguments_start) if next_function_index == -1: arguments_raw = text[arguments_start:].strip() current_index = len(text) else: arguments_raw = text[arguments_start:next_function_index].strip() current_index = next_function_index if not arguments_raw: arguments: Any = "" else: try: arguments = json.loads(arguments_raw) except Exception: arguments = arguments_raw tool_calls.append( { "name": function_name, "arguments": arguments, } ) return tool_calls def _extract_official_tool_call(self, text: str) -> Optional[dict[str, Any]]: stripped = text.strip() if not stripped.startswith("{") or not stripped.endswith("}"): return None payload = json.loads(stripped) return self._normalize_tool_call(payload) def _build_delta_tool_call(self, tool_call: dict[str, Any], index: int = 0) -> DeltaMessage: arguments = tool_call["arguments"] return DeltaMessage( tool_calls=[ DeltaToolCall( index=index, id=f"call_{tool_call['name']}", type="function", function=DeltaFunctionCall( name=tool_call["name"], arguments=json.dumps(arguments, ensure_ascii=False) if isinstance(arguments, (dict, list)) else arguments, ), ) ] ) def _build_tool_calls_response( self, tool_calls: list[dict[str, Any]], ) -> ExtractedToolCallInformation: return ExtractedToolCallInformation( tools_called=True, tool_calls=[ ToolCall( id=f"call_{index + 1}", type="function", function=FunctionCall( name=tool_call["name"], arguments=json.dumps( tool_call["arguments"], ensure_ascii=False ) if isinstance(tool_call["arguments"], (dict, list)) else tool_call["arguments"], ), ) for index, tool_call in enumerate(tool_calls) ], content=None, ) def _looks_like_partial_official_json(self, text: str) -> bool: stripped = text.strip() if not stripped.startswith("{"): return False if stripped.endswith("}"): return False return ( '"name"' in stripped or '"parameters"' in stripped or '"function"' in stripped ) 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, ) -> DeltaMessage | None: variant = self._get_template_variant(request) try: if ( variant == self.VARIANT_LLAMA31 or self.LEGACY_START in current_text ): if self.LEGACY_START in current_text and self.LEGACY_END in current_text: tool_calls = self._extract_legacy_tool_calls(current_text) if tool_calls: return self._build_delta_tool_call( tool_calls[-1], index=len(tool_calls) - 1 ) if self.LEGACY_START in current_text: return None return DeltaMessage(content=delta_text) if variant == self.VARIANT_PIPELINE3 or self.FUNCTION_CALL_TAG in current_text: tool_calls = self._extract_function_tag_tool_calls(current_text) if tool_calls: return self._build_delta_tool_call( tool_calls[-1], index=len(tool_calls) - 1 ) return None official_tool_call = self._extract_official_tool_call(current_text) if official_tool_call: return self._build_delta_tool_call(official_tool_call) if variant == self.VARIANT_OFFICIAL and self._looks_like_partial_official_json( current_text ): return None except Exception: return DeltaMessage(content=delta_text) return DeltaMessage(content=delta_text) def extract_tool_calls( self, model_output: str, request: ChatCompletionRequest, ) -> ExtractedToolCallInformation: variant = self._get_template_variant(request) try: if ( variant == self.VARIANT_LLAMA31 or self.LEGACY_START in model_output ): tool_calls = self._extract_legacy_tool_calls(model_output) if tool_calls: return self._build_tool_calls_response(tool_calls) if variant == self.VARIANT_PIPELINE3 or self.FUNCTION_CALL_TAG in model_output: tool_calls = self._extract_function_tag_tool_calls(model_output) if tool_calls: return self._build_tool_calls_response(tool_calls) official_tool_call = self._extract_official_tool_call(model_output) if official_tool_call: return self._build_tool_calls_response([official_tool_call]) except Exception: pass return ExtractedToolCallInformation( tools_called=False, tool_calls=[], content=model_output, )