| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from __future__ import annotations |
| 5 | |
| 6 | import json |
| 7 | from collections.abc import Sequence |
| 8 | from dataclasses import dataclass, replace |
| 9 | from enum import Enum, auto |
| 10 | from typing import TYPE_CHECKING, NamedTuple |
| 11 | |
| 12 | from openai_harmony import HarmonyError, Message, Role |
| 13 | from xgrammar import StructuralTag |
| 14 | from xgrammar.openai_tool_call_schema import BuiltinToolParam, FunctionToolParam |
| 15 | from xgrammar.structural_tag import ( |
| 16 | AnyTextFormat, |
| 17 | ConstStringFormat, |
| 18 | Format, |
| 19 | JSONSchemaFormat, |
| 20 | OptionalFormat, |
| 21 | OrFormat, |
| 22 | SequenceFormat, |
| 23 | TagFormat, |
| 24 | ) |
| 25 | |
| 26 | from vllm.entrypoints.chat_utils import make_tool_call_id |
| 27 | from vllm.entrypoints.generate.base.protocol import ( |
| 28 | DeltaFunctionCall, |
| 29 | DeltaMessage, |
| 30 | DeltaToolCall, |
| 31 | FunctionCall, |
| 32 | ) |
| 33 | from vllm.entrypoints.openai.chat_completion.protocol import ( |
| 34 | ChatCompletionRequest, |
| 35 | ) |
| 36 | from vllm.entrypoints.openai.parser.harmony_utils import ( |
| 37 | extract_function_from_recipient, |
| 38 | get_encoding, |
| 39 | get_streamable_parser_for_assistant, |
| 40 | is_function_recipient, |
| 41 | ) |
| 42 | from vllm.entrypoints.openai.responses.protocol import ResponsesRequest |
| 43 | from vllm.logger import init_logger |
| 44 | from vllm.parser.abstract_parser import DelegatingParser, structured_outputs_to_format |
| 45 | from vllm.reasoning.gptoss_reasoning_parser import GptOssReasoningParser |
| 46 | from vllm.tool_parsers.gptoss_tool_parser import GptOssToolParser |
| 47 | from vllm.tool_parsers.structural_tag_registry import ( |
| 48 | SimplifiedToolChoice, |
| 49 | get_function_parameters, |
| 50 | register_vllm_structural_tag, |
| 51 | ) |
| 52 | |
| 53 | if TYPE_CHECKING: |
| 54 | from openai_harmony import Message, StreamableParser |
| 55 | |
| 56 | |
| 57 | logger = init_logger(__name__) |
| 58 | |
| 59 | |
| 60 | class _SegmentType(Enum): |
| 61 | TOOL = auto() |
| 62 | REASONING = auto() |
| 63 | CONTENT = auto() |
| 64 | IGNORE = auto() |
| 65 | |
| 66 | @staticmethod |
| 67 | def from_channel_and_recipient( |
| 68 | channel: str | None, recipient: str | None |
| 69 | ) -> _SegmentType: |
| 70 | if recipient and is_function_recipient(recipient): |
| 71 | return _SegmentType.TOOL |
| 72 | if channel == "analysis": |
| 73 | return _SegmentType.REASONING |
| 74 | if channel == "final" or (channel == "commentary" and recipient is None): |
| 75 | return _SegmentType.CONTENT |
| 76 | return _SegmentType.IGNORE |
| 77 | |
| 78 | |
| 79 | class Segment(NamedTuple): |
| 80 | channel: str | None |
| 81 | recipient: str | None |
| 82 | delta: str |
| 83 | completed_message: Message | None = None |
| 84 | |
| 85 | |
| 86 | @dataclass |
| 87 | class ChunkResult: |
| 88 | segments: list[Segment] |
| 89 | reasoning_token_count: int |
| 90 | |
| 91 | |
| 92 | class HarmonyParser(DelegatingParser): |
| 93 | def __init__(self, tokenizer, tools=None, *args, **kwargs): |
| 94 | super().__init__(tokenizer, tools, *args, **kwargs) |
| 95 | |
| 96 | if self.reasoning_parser and not isinstance( |
| 97 | self.reasoning_parser, GptOssReasoningParser |
| 98 | ): |
| 99 | raise ValueError( |
| 100 | "Harmony requires GptOssReasoningParser, " |
| 101 | f"got {self.reasoning_parser.__class__.__name__}." |
| 102 | ) |
| 103 | |
| 104 | if self.tool_parser and not isinstance(self.tool_parser, GptOssToolParser): |
| 105 | raise ValueError( |
| 106 | "Harmony requires GptOssToolParser, " |
| 107 | f"got {self.tool_parser.__class__.__name__}." |
| 108 | ) |
| 109 | |
| 110 | self._parser: StreamableParser | None = None |
| 111 | self._next_tool_call_index = 0 |
| 112 | self._num_processed_messages = 0 |
| 113 | |
| 114 | self._num_counted_tokens = 0 |
| 115 | self._num_reasoning_tokens = 0 |
| 116 | |
| 117 | # For error recovery |
| 118 | self._current_message_tokens: list[int] = [] |
| 119 | |
| 120 | @property |
| 121 | def _harmony_parser(self) -> StreamableParser: |
| 122 | """Lazily initializes the Harmony parser.""" |
| 123 | if self._parser is None: |
| 124 | self._parser = get_streamable_parser_for_assistant() |
| 125 | return self._parser |
| 126 | |
| 127 | def _poll_completed_message(self) -> Message | None: |
| 128 | messages = self._harmony_parser.messages |
| 129 | if len(messages) <= self._num_processed_messages: |
| 130 | return None |
| 131 | msg = messages[self._num_processed_messages] |
| 132 | msg.recipient = self._normalize_recipient(msg.recipient) |
| 133 | self._num_processed_messages += 1 |
| 134 | return msg |
| 135 | |
| 136 | def flush(self) -> list[Segment]: |
| 137 | segments: list[Segment] = [] |
| 138 | try: |
| 139 | self._harmony_parser.process_eos() |
| 140 | msg = self._poll_completed_message() |
| 141 | except HarmonyError: |
| 142 | logger.warning( |
| 143 | "Harmony parser ended in a non-terminal state; returning the " |
| 144 | "recovered raw output." |
| 145 | ) |
| 146 | |
| 147 | final_channel = "final" |
| 148 | text = self.model_tokenizer.decode(self._current_message_tokens) |
| 149 | segments.append( |
| 150 | Segment( |
| 151 | channel=final_channel, |
| 152 | recipient=None, |
| 153 | delta=text, |
| 154 | completed_message=None, |
| 155 | ) |
| 156 | ) |
| 157 | msg = Message.from_role_and_content(Role.ASSISTANT, text).with_channel( |
| 158 | final_channel |
| 159 | ) |
| 160 | |
| 161 | # Reset to the initial assistant-parser state for the next turn. |
| 162 | self._parser = None |
| 163 | self._num_processed_messages = 0 |
| 164 | self._current_message_tokens.clear() |
| 165 | |
| 166 | if msg is None: |
| 167 | return segments |
| 168 | |
| 169 | segments.append( |
| 170 | Segment( |
| 171 | channel=msg.channel, |
| 172 | recipient=msg.recipient, |
| 173 | delta="", |
| 174 | completed_message=msg, |
| 175 | ) |
| 176 | ) |
| 177 | return segments |
| 178 | |
| 179 | def parse( |
| 180 | self, |
| 181 | model_output: str, |
| 182 | request: ChatCompletionRequest | ResponsesRequest, |
| 183 | enable_auto_tools: bool = False, |
| 184 | model_output_token_ids: Sequence[int] = (), |
| 185 | ) -> tuple[str | None, str | None, list[FunctionCall] | None]: |
| 186 | """Parse Harmony output from token IDs. |
| 187 | |
| 188 | Tool calls are always extracted regardless of ``enable_auto_tools``. |
| 189 | Callers must decide whether to surface them. |
| 190 | """ |
| 191 | result = self.process_chunk(model_output_token_ids) |
| 192 | flushed_segments = self.flush() |
| 193 | if flushed_segments: |
| 194 | result.segments.extend(flushed_segments) |
| 195 | |
| 196 | reasoning_parts: list[str] = [] |
| 197 | content_parts: list[str] = [] |
| 198 | tool_calls: list[FunctionCall] = [] |
| 199 | |
| 200 | for segment in result.segments: |
| 201 | msg = segment.completed_message |
| 202 | if msg is None: |
| 203 | continue |
| 204 | if msg.author.role != "assistant" or not msg.content: |
| 205 | continue |
| 206 | text = msg.content[0].text |
| 207 | segment_type = _SegmentType.from_channel_and_recipient( |
| 208 | msg.channel, msg.recipient |
| 209 | ) |
| 210 | match segment_type: |
| 211 | case _SegmentType.REASONING if self.reasoning_parser and text: |
| 212 | reasoning_parts.append(text) |
| 213 | case _SegmentType.CONTENT if text: |
| 214 | content_parts.append(text) |
| 215 | case _SegmentType.TOOL if self.tool_parser: |
| 216 | recipient = msg.recipient |
| 217 | content_type = msg.content_type |
| 218 | assert recipient is not None |
| 219 | if content_type is not None and "json" not in content_type: |
| 220 | arguments = text |
| 221 | else: |
| 222 | try: |
| 223 | arguments = json.dumps(json.loads(text)) |
| 224 | except json.JSONDecodeError: |
| 225 | arguments = text |
| 226 | tool_calls.append( |
| 227 | FunctionCall( |
| 228 | name=extract_function_from_recipient(recipient), |
| 229 | arguments=arguments, |
| 230 | ) |
| 231 | ) |
| 232 | |
| 233 | reasoning = "\n".join(reasoning_parts) or None |
| 234 | content = "\n".join(content_parts) or None |
| 235 | return reasoning, content, tool_calls or None |
| 236 | |
| 237 | def parse_delta( |
| 238 | self, |
| 239 | delta_text: str, |
| 240 | delta_token_ids: list[int], |
| 241 | request: ChatCompletionRequest | ResponsesRequest, |
| 242 | prompt_token_ids: list[int] | None = None, |
| 243 | *, |
| 244 | finished: bool, |
| 245 | ) -> DeltaMessage | None: |
| 246 | prev_recipient = self._normalize_recipient( |
| 247 | self._harmony_parser.current_recipient |
| 248 | ) |
| 249 | result = self.process_chunk(delta_token_ids) |
| 250 | if finished: |
| 251 | flushed_segments = self.flush() |
| 252 | if flushed_segments: |
| 253 | result.segments.extend(flushed_segments) |
| 254 | combined_content = "" |
| 255 | combined_reasoning = "" |
| 256 | tool_messages: list[DeltaToolCall] = [] |
| 257 | |
| 258 | for segment in result.segments: |
| 259 | if segment.completed_message is not None: |
| 260 | prev_recipient = None |
| 261 | continue |
| 262 | |
| 263 | segment_type = _SegmentType.from_channel_and_recipient( |
| 264 | segment.channel, segment.recipient |
| 265 | ) |
| 266 | match segment_type: |
| 267 | case _SegmentType.REASONING if self.reasoning_parser: |
| 268 | combined_reasoning += segment.delta |
| 269 | case _SegmentType.CONTENT: |
| 270 | combined_content += segment.delta |
| 271 | case _SegmentType.TOOL if self.tool_parser: |
| 272 | assert segment.recipient is not None |
| 273 | if prev_recipient != segment.recipient: |
| 274 | tool_name = extract_function_from_recipient(segment.recipient) |
| 275 | tool_messages.append( |
| 276 | DeltaToolCall( |
| 277 | # HarmonyParser does not use _stream_state; |
| 278 | # "random" tool_call_id_type is always used |
| 279 | id=make_tool_call_id(), |
| 280 | type="function", |
| 281 | function=DeltaFunctionCall( |
| 282 | name=tool_name, |
| 283 | arguments=segment.delta, |
| 284 | ), |
| 285 | index=self._next_tool_call_index, |
| 286 | ) |
| 287 | ) |
| 288 | self._next_tool_call_index += 1 |
| 289 | prev_recipient = segment.recipient |
| 290 | elif segment.delta: |
| 291 | idx = self._next_tool_call_index - 1 |
| 292 | if tool_messages: |
| 293 | tool_msg = tool_messages[-1] |
| 294 | assert tool_msg.index == idx |
| 295 | fn = tool_msg.function |
| 296 | assert fn is not None and fn.arguments is not None |
| 297 | fn.arguments += segment.delta |
| 298 | else: |
| 299 | tool_messages.append( |
| 300 | DeltaToolCall( |
| 301 | index=idx, |
| 302 | function=DeltaFunctionCall(arguments=segment.delta), |
| 303 | ) |
| 304 | ) |
| 305 | |
| 306 | if finished: |
| 307 | self._next_tool_call_index = 0 |
| 308 | |
| 309 | if not combined_content and not combined_reasoning and not tool_messages: |
| 310 | return None |
| 311 | |
| 312 | delta_message = DeltaMessage() |
| 313 | if combined_content: |
| 314 | delta_message.content = combined_content |
| 315 | if combined_reasoning: |
| 316 | delta_message.reasoning = combined_reasoning |
| 317 | if tool_messages: |
| 318 | delta_message.tool_calls = tool_messages |
| 319 | |
| 320 | # Suppress reasoning deltas if not requested |
| 321 | if delta_message and not request.include_reasoning: |
| 322 | delta_message.reasoning = None |
| 323 | |
| 324 | # If only reasoning was in the message (no content, no tool_calls) |
| 325 | # skip emitting entirely |
| 326 | if not delta_message.content and not delta_message.tool_calls: |
| 327 | return None |
| 328 | |
| 329 | return delta_message |
| 330 | |
| 331 | def process_chunk(self, token_ids: Sequence[int]) -> ChunkResult: |
| 332 | if not token_ids: |
| 333 | return ChunkResult(segments=[], reasoning_token_count=0) |
| 334 | |
| 335 | segments: list[Segment] = [] |
| 336 | reasoning_token_count = 0 |
| 337 | for token_id in token_ids: |
| 338 | try: |
| 339 | self._harmony_parser.process(token_id) |
| 340 | except HarmonyError: |
| 341 | continue |
| 342 | channel = self._harmony_parser.current_channel |
| 343 | recipient = self._normalize_recipient( |
| 344 | self._harmony_parser.current_recipient |
| 345 | ) |
| 346 | delta = self._harmony_parser.last_content_delta or "" |
| 347 | completed_message = self._poll_completed_message() |
| 348 | |
| 349 | if completed_message is not None: |
| 350 | self._current_message_tokens.clear() |
| 351 | else: |
| 352 | self._current_message_tokens.append(token_id) |
| 353 | |
| 354 | if self._is_reasoning_token(token_id, channel, recipient): |
| 355 | reasoning_token_count += 1 |
| 356 | |
| 357 | segments.append( |
| 358 | Segment( |
| 359 | channel=channel, |
| 360 | recipient=recipient, |
| 361 | delta=delta, |
| 362 | completed_message=completed_message, |
| 363 | ) |
| 364 | ) |
| 365 | |
| 366 | # TODO: Optionally merge and suppress empty Segments |
| 367 | |
| 368 | self._num_counted_tokens += len(token_ids) |
| 369 | self._num_reasoning_tokens += reasoning_token_count |
| 370 | return ChunkResult( |
| 371 | segments=segments, |
| 372 | reasoning_token_count=reasoning_token_count, |
| 373 | ) |
| 374 | |
| 375 | def count_reasoning_tokens(self, token_ids: Sequence[int]) -> int: |
| 376 | if len(token_ids) == self._num_counted_tokens: |
| 377 | return self._num_reasoning_tokens |
| 378 | |
| 379 | parser = get_streamable_parser_for_assistant() |
| 380 | count = 0 |
| 381 | for token_id in token_ids: |
| 382 | try: |
| 383 | parser.process(token_id) |
| 384 | except HarmonyError: |
| 385 | continue |
| 386 | recipient = self._normalize_recipient(parser.current_recipient) |
| 387 | if self._is_reasoning_token(token_id, parser.current_channel, recipient): |
| 388 | count += 1 |
| 389 | return count |
| 390 | |
| 391 | @staticmethod |
| 392 | def _is_reasoning_token( |
| 393 | token_id: int, channel: str | None, recipient: str | None |
| 394 | ) -> bool: |
| 395 | is_reasoning_channel = channel == "analysis" or ( |
| 396 | channel == "commentary" and recipient is not None |
| 397 | ) |
| 398 | return is_reasoning_channel and not get_encoding().is_special_token(token_id) |
| 399 | |
| 400 | def adjust_request( |
| 401 | self, request: ChatCompletionRequest | ResponsesRequest |
| 402 | ) -> ChatCompletionRequest | ResponsesRequest: |
| 403 | request = _adjust_output_format(request) |
| 404 | return super().adjust_request(request) |
| 405 | |
| 406 | @staticmethod |
| 407 | def _normalize_recipient(recipient: str | None) -> str | None: |
| 408 | """Remove constrained formats misparsed into recipients by older Harmony.""" |
| 409 | if recipient is None: |
| 410 | return None |
| 411 | |
| 412 | constrain_index = recipient.find("<|constrain|>") |
| 413 | if constrain_index == -1: |
| 414 | return recipient |
| 415 | return recipient[:constrain_index].rstrip() or None |
| 416 | |
| 417 | |
| 418 | _JSON_CONSTRAINS = [" json", " <|constrain|>json"] |
| 419 | # Harmomy's stop tokens are <|return|>, <|call|>, <|endoftext|> |
| 420 | # They are represented as "" since xgrammar disallows stop tokens while |
| 421 | # under constraints, leading to bad or infinite generations. |
| 422 | _END_TAG = ["<|end|>", ""] |
| 423 | _FINAL_BEGIN = "<|channel|>final{constrain}<|message|>" |
| 424 | _TOOL_CALL_CHANNELS = [ |
| 425 | "<|channel|>commentary", |
| 426 | "<|channel|>analysis", |
| 427 | "<|channel|>final", |
| 428 | ] |
| 429 | _FUNCTION_CALL_BEGINS = [ |
| 430 | " to=functions.{name}{channel}{constrain}<|message|>", |
| 431 | "{channel} to=functions.{name}{constrain}<|message|>", |
| 432 | ] |
| 433 | _ANY_CONTENT = AnyTextFormat() |
| 434 | |
| 435 | |
| 436 | def _assemble_tag( |
| 437 | allow_analysis: bool, allow_commentary: bool, content: Format |
| 438 | ) -> StructuralTag: |
| 439 | tags = [] |
| 440 | if allow_analysis: |
| 441 | analysis_tag = OptionalFormat( |
| 442 | content=SequenceFormat( |
| 443 | elements=[ |
| 444 | TagFormat( |
| 445 | begin="<|channel|>analysis<|message|>", |
| 446 | content=_ANY_CONTENT, |
| 447 | end="<|end|>", |
| 448 | ), |
| 449 | ConstStringFormat(value="<|start|>assistant"), |
| 450 | ] |
| 451 | ) |
| 452 | ) |
| 453 | tags.append(analysis_tag) |
| 454 | |
| 455 | if allow_commentary: |
| 456 | commentary_tag = OptionalFormat( |
| 457 | content=SequenceFormat( |
| 458 | elements=[ |
| 459 | TagFormat( |
| 460 | begin="<|channel|>commentary<|message|>", |
| 461 | content=_ANY_CONTENT, |
| 462 | end="<|end|>", |
| 463 | ), |
| 464 | ConstStringFormat(value="<|start|>assistant"), |
| 465 | ] |
| 466 | ) |
| 467 | ) |
| 468 | tags.append(commentary_tag) |
| 469 | |
| 470 | tags.append(content) |
| 471 | |
| 472 | return StructuralTag(format=SequenceFormat(elements=tags)) |
| 473 | |
| 474 | |
| 475 | @register_vllm_structural_tag("harmony") |
| 476 | def get_harmony_structural_tag( |
| 477 | tools: list[FunctionToolParam], |
| 478 | builtin_tools: list[BuiltinToolParam], |
| 479 | tool_choice: SimplifiedToolChoice, |
| 480 | reasoning: bool, |
| 481 | token_suffix: str = "", |
| 482 | ) -> StructuralTag: |
| 483 | # reasoning always enabled for Harmony |
| 484 | del reasoning, token_suffix |
| 485 | |
| 486 | if builtin_tools: |
| 487 | # Fallback for built-in tools |
| 488 | tags = [ |
| 489 | TagFormat( |
| 490 | begin="to=", |
| 491 | content=AnyTextFormat(excludes=["<|start|>"]), |
| 492 | end=_END_TAG, |
| 493 | ) |
| 494 | ] |
| 495 | tags.extend( |
| 496 | TagFormat( |
| 497 | begin=channel + " to=", |
| 498 | content=AnyTextFormat(excludes=["<|start|>", "<|channel|>"]), |
| 499 | end=_END_TAG, |
| 500 | ) |
| 501 | for channel in _TOOL_CALL_CHANNELS |
| 502 | ) |
| 503 | else: |
| 504 | tags = [ |
| 505 | TagFormat( |
| 506 | begin=pattern.format( |
| 507 | name=tool.function.name, channel=channel, constrain=constrain |
| 508 | ), |
| 509 | content=JSONSchemaFormat( |
| 510 | json_schema=get_function_parameters(tool.function) |
| 511 | ), |
| 512 | end=_END_TAG, |
| 513 | ) |
| 514 | for tool in tools |
| 515 | for pattern in _FUNCTION_CALL_BEGINS |
| 516 | for constrain in _JSON_CONSTRAINS |
| 517 | for channel in _TOOL_CALL_CHANNELS |
| 518 | ] |
| 519 | |
| 520 | if tool_choice == "auto": |
| 521 | tags.extend( |
| 522 | TagFormat( |
| 523 | begin=_FINAL_BEGIN.format(constrain=constrain), |
| 524 | content=_ANY_CONTENT, |
| 525 | end=_END_TAG, |
| 526 | ) |
| 527 | for constrain in _JSON_CONSTRAINS + [""] |
| 528 | ) |
| 529 | |
| 530 | return _assemble_tag( |
| 531 | allow_analysis=True, allow_commentary=True, content=OrFormat(elements=tags) |
| 532 | ) |
| 533 | |
| 534 | |
| 535 | def _adjust_output_format( |
| 536 | request: ChatCompletionRequest | ResponsesRequest, |
| 537 | ) -> ChatCompletionRequest | ResponsesRequest: |
| 538 | """Canonicalize request constraints into a reasoning-aware StructuralTag.""" |
| 539 | params = request.extract_structured_outputs() |
| 540 | if params is None: |
| 541 | return request |
| 542 | |
| 543 | final_content = structured_outputs_to_format(params) |
| 544 | if final_content is None: |
| 545 | return request |
| 546 | |
| 547 | if isinstance(final_content, JSONSchemaFormat): |
| 548 | content = OrFormat( |
| 549 | elements=[ |
| 550 | TagFormat( |
| 551 | begin=_FINAL_BEGIN.format(constrain=constrain), |
| 552 | content=final_content, |
| 553 | end=_END_TAG, |
| 554 | ) |
| 555 | for constrain in _JSON_CONSTRAINS |
| 556 | ] |
| 557 | ) |
| 558 | else: |
| 559 | content = TagFormat( |
| 560 | begin=_FINAL_BEGIN.format(constrain=""), content=final_content, end=_END_TAG |
| 561 | ) |
| 562 | structural_tag = _assemble_tag( |
| 563 | allow_analysis=True, allow_commentary=False, content=content |
| 564 | ) |
| 565 | |
| 566 | request.structured_outputs = replace( |
| 567 | params, |
| 568 | json=None, |
| 569 | regex=None, |
| 570 | choice=None, |
| 571 | grammar=None, |
| 572 | json_object=None, |
| 573 | structural_tag=json.dumps(structural_tag.model_dump()), |
| 574 | ) |
| 575 | if isinstance(request, ResponsesRequest): |
| 576 | request.text = None |
| 577 | else: |
| 578 | request.response_format = None |
| 579 | return request |