"""The standard conversation protocol in MLC LLM""" from enum import Enum from typing import Any, Dict, List, Optional, Tuple, Type, TypeVar, Union # noqa: UP035 from pydantic import BaseModel, Field, field_validator # The message placeholders in the message prompts according to roles. class MessagePlaceholders(Enum): """The message placeholders in the message prompts according to roles.""" SYSTEM = "{system_message}" USER = "{user_message}" ASSISTANT = "{assistant_message}" TOOL = "{tool_message}" FUNCTION = "{function_string}" T = TypeVar("T", bound="BaseModel") class Conversation(BaseModel): """Class that specifies the convention template of conversation and contains the conversation history. Given a conversation template, the corresponding prompt generated out from it is usually in the following format: <><><><><> <><><><> ... <><><><> <><> """ # Optional name of the template. name: Optional[str] = None # The system prompt template, it optionally contains the system # message placeholder, and the placeholder will be replaced with # the system message below. system_template: str = MessagePlaceholders.SYSTEM.value # The content of the system prompt (without the template format). system_message: str = "" # The system token ids to be prepended at the beginning of tokenized # generated prompt. system_prefix_token_ids: Optional[List[int]] = None # noqa: UP006 # Whether or not to append user role and separator after the system message. # This is mainly for [INST] [/INST] style prompt format add_role_after_system_message: bool = True # The conversation roles roles: Dict[str, str] # noqa: UP006 # The roles prompt template, it optionally contains the defaults # message placeholders and will be replaced by actual content role_templates: Dict[str, str] # noqa: UP006 # The conversation history messages. # Each message is a pair of strings, denoting "(role, content)". # The content can be None. messages: List[Tuple[str, Optional[Union[str, List[Dict]]]]] = Field(default_factory=lambda: []) # noqa: UP006 # The separators between messages when concatenating into a single prompt. # List size should be either 1 or 2. # - When size is 1, the separator will be used between adjacent messages. # - When size is 2, seps[0] is used after user message, and # seps[1] is used after assistant message. seps: List[str] # noqa: UP006 # The separator between the role and the content in a message. role_content_sep: str = "" # The separator between the role and empty contents. role_empty_sep: str = "" # The stop criteria stop_str: List[str] = Field(default_factory=lambda: []) # noqa: UP006 stop_token_ids: List[int] = Field(default_factory=lambda: []) # noqa: UP006 # When True, strip `...` blocks (and any trailing whitespace) # from historical assistant messages before rendering the prompt, mirroring # Qwen3's official HF chat template. Only historical turns before the last # user message are affected; reasoning on the most recent assistant turn is # preserved for tool-call prefill scenarios. strip_reasoning_in_history: bool = False # Function call fields function_string: str = "" # whether using function calling or not, helps check for output message format in API call use_function_calling: bool = False def __init__(self, role_templates: Optional[Dict[str, str]] = None, **kwargs): # noqa: UP006 # Defaults templates which would be overridden by model specific templates _role_templates: Dict[str, str] = { # noqa: UP006 "user": MessagePlaceholders.USER.value, "assistant": MessagePlaceholders.ASSISTANT.value, "tool": MessagePlaceholders.TOOL.value, } if role_templates is not None: _role_templates.update(role_templates) super().__init__(role_templates=_role_templates, **kwargs) @field_validator("seps") @classmethod def check_message_seps(cls, seps: List[str]) -> List[str]: # noqa: UP006 """Check if the input message separators has size 1 or 2.""" if len(seps) == 0 or len(seps) > 2: raise ValueError("seps should have size 1 or 2.") return seps def to_json_dict(self) -> Dict[str, Any]: # noqa: UP006 """Convert to a json dictionary""" return self.model_dump(by_alias=True, exclude_none=True) @classmethod def from_json_dict(cls: Type[T], json_dict: Dict[str, Any]) -> T: # noqa: UP006 """Convert from a json dictionary""" return Conversation.model_validate(json_dict) def as_prompt(self, config=None) -> List[Any]: # noqa: UP006 """Convert the conversation template and history messages to a single prompt. Returns ------- prompts : List[Union[str, "mlc_llm.serve.data.Data"]] The prompts converted from the conversation messages. We use Any in the signature to avoid cyclic import. """ from ..serve import data # - Get the system message. system_msg = self.system_template.replace( MessagePlaceholders.SYSTEM.value, self.system_message ) # - Get the message strings. message_list: List[Union[str, data.Data]] = [] # noqa: UP006 separators = list(self.seps) if len(separators) == 1: separators.append(separators[0]) if system_msg != "": message_list.append(system_msg) messages = ( _strip_reasoning_in_history(self.messages) if self.strip_reasoning_in_history else self.messages ) for i, (role, content) in enumerate(messages): if role not in self.roles.keys(): raise ValueError(f'Role "{role}" is not a supported role in {self.roles.keys()}') separator = separators[role == "assistant"] # check assistant role if content is None: message_list.append(self.roles[role] + self.role_empty_sep) continue role_prefix = ( "" # Do not append role prefix if this is the first message and there # is already a system message if (not self.add_role_after_system_message and system_msg != "" and i == 0) else self.roles[role] + self.role_content_sep ) if isinstance(content, str): message_list.append( role_prefix + self.role_templates[role].replace( MessagePlaceholders[role.upper()].value, content ) + separator ) continue message_list.append(role_prefix) for item in content: assert isinstance(item, dict), "Content should be a string or a list of dicts" assert "type" in item, "Content item should have a type field" if item["type"] == "text": message = self.role_templates[role].replace( MessagePlaceholders[role.upper()].value, item["text"] ) message_list.append(message) elif item["type"] == "image_url": assert config is not None, "Model config is required" image_url = _get_url_from_item(item) message_list.append(data.ImageData.from_url(image_url, config)) message_list.append("\n") else: raise ValueError(f"Unsupported content type: {item['type']}") message_list.append(separator) prompt = _combine_consecutive_messages(message_list) if not any(isinstance(item, data.ImageData) for item in message_list): # Replace the last function string placeholder with actual function string prompt[0] = self.function_string.join( prompt[0].rsplit(MessagePlaceholders.FUNCTION.value, 1) ) # Replace with remaining function string placeholders with empty string prompt[0] = prompt[0].replace(MessagePlaceholders.FUNCTION.value, "") return prompt def _get_url_from_item(item: Dict) -> str: # noqa: UP006 image_url: str assert "image_url" in item, "Content item should have an image_url field" if isinstance(item["image_url"], str): image_url = item["image_url"] elif isinstance(item["image_url"], dict): assert "url" in item["image_url"], ( "Content image_url item should be a string or a dict with a url field" ) image_url = item["image_url"]["url"] else: raise ValueError( "Content image_url item type not supported. " "Should be a string or a dict with a url field." ) return image_url def _strip_reasoning_in_history( messages: List[Tuple[str, Optional[Union[str, List[Dict]]]]], # noqa: UP006 ) -> List[Tuple[str, Optional[Union[str, List[Dict]]]]]: # noqa: UP006 """Strip `...` blocks from assistant messages that precede the last user message, matching Qwen3's HF chat-template behavior. The last assistant message (if any) is preserved so tool-call prefill continuations keep their reasoning context. """ last_user_idx = -1 for i, (role, _) in enumerate(messages): if role == "user": last_user_idx = i result: List[Tuple[str, Optional[Union[str, List[Dict]]]]] = [] # noqa: UP006 for i, (role, content) in enumerate(messages): if ( role == "assistant" and i < last_user_idx and isinstance(content, str) and "" in content ): content = content.split("")[-1].lstrip("\n") result.append((role, content)) return result def _combine_consecutive_messages(messages: List[Any]) -> List[Any]: # noqa: UP006 """Combining consecutive strings into one. Parameters ---------- messages : List[Union[str, "mlc_llm.serve.data.Data"]] The input messages to be combined. We use Any in the signature to avoid cyclic import. Returns ------- updated_messages : List[Union[str, "mlc_llm.serve.data.Data"]] The combined messages """ if len(messages) == 0: return [] combined_messages = [messages[0]] for message in messages[1:]: if isinstance(message, str) and isinstance(combined_messages[-1], str): combined_messages[-1] += message else: combined_messages.append(message) return combined_messages