diff --git a/atom/entrypoints/openai/api_server.py b/atom/entrypoints/openai/api_server.py index c8a24f2817..0587e57dd2 100644 --- a/atom/entrypoints/openai/api_server.py +++ b/atom/entrypoints/openai/api_server.py @@ -23,18 +23,19 @@ import uuid from asyncio import AbstractEventLoop from contextlib import asynccontextmanager -from typing import Any, AsyncGenerator, Dict, List, Optional, Tuple +from typing import Any, AsyncGenerator, Dict, List, Optional, Tuple # noqa: UP035 import uvicorn +from fastapi import FastAPI, HTTPException, Request +from fastapi.responses import JSONResponse, StreamingResponse +from PIL import Image +from transformers import AutoProcessor, AutoTokenizer + from atom import SamplingParams from atom.model_engine.arg_utils import EngineArgs from atom.model_engine.llm_engine import _load_tokenizer from atom.model_engine.request import RequestOutput from atom.utils.arg_parser import FlexibleArgumentParser -from fastapi import FastAPI, HTTPException, Request -from fastapi.responses import JSONResponse, StreamingResponse -from PIL import Image -from transformers import AutoProcessor, AutoTokenizer from .chat_encoders import apply_chat_template, load_custom_message_encoder from .protocol import ( @@ -43,12 +44,6 @@ ModelCard, ModelList, ) -from .serving_chat import ( - build_chat_response, - build_chat_response_multi, - stream_chat_response, - stream_chat_response_fanout, -) from .serving_anthropic import ( AnthropicMessagesRequest, anthropic_to_openai_messages, @@ -62,12 +57,19 @@ stream_message_stop, stream_signature_delta, ) +from .serving_chat import ( + build_chat_response, + build_chat_response_multi, + stream_chat_response, + stream_chat_response_fanout, +) from .serving_completion import ( build_completion_response, build_completion_response_multi, stream_completion_response, stream_completion_response_fanout, ) +from .streaming_dispatch import StreamBatchDispatcher # Configure logging logger = logging.getLogger("atom") @@ -82,16 +84,17 @@ # ============================================================================ engine = None -tokenizer: Optional[AutoTokenizer] = None -processor: Optional[Any] = None +tokenizer: AutoTokenizer | None = None +processor: Any | None = None model_name: str = "" -default_chat_template_kwargs: Dict[str, Any] = {} -custom_message_encoder: Optional[Any] = None -_stream_queues: Dict[str, asyncio.Queue] = {} -_seq_id_to_request_id: Dict[int, str] = {} -_stream_loops: Dict[str, AbstractEventLoop] = {} -_request_start_times: Dict[str, float] = {} -_request_logger: Optional[logging.Logger] = None +default_chat_template_kwargs: dict[str, Any] = {} +custom_message_encoder: Any | None = None +_stream_queues: dict[str, asyncio.Queue] = {} +_seq_id_to_request_id: dict[int, str] = {} +_stream_loops: dict[str, AbstractEventLoop] = {} +_request_start_times: dict[str, float] = {} +_request_logger: logging.Logger | None = None +_stream_batch_dispatcher: StreamBatchDispatcher | None = None # ============================================================================ @@ -324,19 +327,13 @@ def _prepare_multimodal_inputs( return inputs["input_ids"][0].tolist(), multimodal_data -def _send_stream_chunk_direct( - request_output: RequestOutput, - request_id: str, - stream_queue: asyncio.Queue, - loop: AbstractEventLoop, -) -> None: - """Send stream chunk directly to the queue.""" - global tokenizer +# ── Batched stream dispatch ────────────────────────────────────────────── + - new_text = tokenizer.decode(request_output.output_tokens, skip_special_tokens=True) +def _build_stream_chunk(request_output: RequestOutput, request_id: str) -> dict: + """Build a raw chunk; detokenization happens once in the batch dispatcher.""" started_at = _request_start_times.get(request_id) chunk_data = { - "text": new_text, "token_ids": request_output.output_tokens, "finished": request_output.finished, "finish_reason": request_output.finish_reason, @@ -346,11 +343,34 @@ def _send_stream_chunk_direct( } if getattr(request_output, "kv_transfer_params_output", None): chunk_data["kv_transfer_params"] = request_output.kv_transfer_params_output - loop.call_soon_threadsafe(stream_queue.put_nowait, chunk_data) + return chunk_data + + +def _send_stream_chunk_direct( + request_output: RequestOutput, + request_id: str, + stream_queue: asyncio.Queue, + loop: AbstractEventLoop, +) -> None: + """Buffer a single-request chunk for this engine step.""" + assert _stream_batch_dispatcher is not None + _stream_batch_dispatcher.enqueue( + loop=loop, + queue=stream_queue, + state_key=request_id, + chunk=_build_stream_chunk(request_output, request_id), + ) + + +def flush_stream_batch() -> None: + """Flush this output thread's engine-step batch into asyncio queues.""" + if _stream_batch_dispatcher is not None: + _stream_batch_dispatcher.flush() def _send_stream_chunk_tagged( request_output: RequestOutput, + request_id: str, sibling_index: int, stream_queue: asyncio.Queue, loop: AbstractEventLoop, @@ -364,18 +384,14 @@ def _send_stream_chunk_tagged( This path serves ``SamplingParams.n > 1`` by tagging each sibling's chunks so the shared stream consumer can merge them in order. """ - global tokenizer - - new_text = tokenizer.decode(request_output.output_tokens, skip_special_tokens=True) - chunk_data = { - "text": new_text, - "token_ids": request_output.output_tokens, - "finished": request_output.finished, - "finish_reason": request_output.finish_reason, - } - if getattr(request_output, "kv_transfer_params_output", None): - chunk_data["kv_transfer_params"] = request_output.kv_transfer_params_output - loop.call_soon_threadsafe(stream_queue.put_nowait, (sibling_index, chunk_data)) + assert _stream_batch_dispatcher is not None + _stream_batch_dispatcher.enqueue( + loop=loop, + queue=stream_queue, + state_key=(request_id, sibling_index), + chunk=_build_stream_chunk(request_output, request_id), + tag=sibling_index, + ) async def generate_async( @@ -831,6 +847,8 @@ def cleanup_streaming_request( _seq_id_to_request_id.pop(seq_id, None) _stream_loops.pop(request_id, None) _request_start_times.pop(request_id, None) + if _stream_batch_dispatcher is not None: + _stream_batch_dispatcher.discard_request(request_id) if aborted: try: engine.core_mgr.abort_request(seq_id) @@ -967,7 +985,9 @@ async def setup_streaming_request_fanout( def make_callback(idx: int): def _cb(request_output: RequestOutput) -> None: - _send_stream_chunk_tagged(request_output, idx, shared_queue, stream_loop) + _send_stream_chunk_tagged( + request_output, request_id, idx, shared_queue, stream_loop + ) return _cb @@ -1433,6 +1453,7 @@ async def generate_anthropic_stream(): if prompt.rstrip().endswith(""): reasoning_filter.state = 1 tool_parser = ToolCallStreamParser() + tool_parser.tools = anthropic_to_openai_tools(request.tools) block_index = 0 started_text = False started_thinking = False @@ -1602,15 +1623,19 @@ async def generate_anthropic_stream(): from .reasoning import separate_reasoning from .tool_parser import parse_tool_calls - final_output = None - async for output in generate_async(prompt, sampling_params, request_id): - final_output = output + final_output = await _run_nonstream_with_disconnect( + generate_async(prompt, sampling_params, request_id), + raw_request, + request_id, + ) if final_output is None: raise RuntimeError("No output generated") raw_text = final_output["text"] reasoning_content, content_with_tools = separate_reasoning(raw_text) - content_text, tool_calls = parse_tool_calls(content_with_tools) + content_text, tool_calls = parse_tool_calls( + content_with_tools, anthropic_to_openai_tools(request.tools) + ) output_tokens = len(tokenizer.encode(raw_text)) cache_read_input_tokens = final_output.get("num_cached_tokens", 0) if not getattr(request, "thinking", None): @@ -1627,8 +1652,11 @@ async def generate_anthropic_stream(): cache_read_input_tokens=cache_read_input_tokens, ) + except _ClientDisconnected: + # Client hung up; seq already aborted + popped. Nothing to return. + return JSONResponse(status_code=499, content={"detail": "client disconnected"}) except Exception as e: - logger.error(f"Error in anthropic_messages: {e}", exc_info=True) + logger.exception("Error in anthropic_messages") return JSONResponse( status_code=500, content={ @@ -1743,7 +1771,7 @@ async def stop_profile(): def main(): """Main entry point for the server.""" global engine, tokenizer, model_name, default_chat_template_kwargs, _request_logger - global custom_message_encoder + global custom_message_encoder, _stream_batch_dispatcher parser = FlexibleArgumentParser(description="ATOM OpenAI API Server") EngineArgs.add_cli_args(parser) @@ -1793,6 +1821,15 @@ def main(): logger.info(f"Initializing engine with model {args.model}...") engine_args = EngineArgs.from_cli_args(args) engine = engine_args.create_engine(tokenizer=tokenizer) + _stream_batch_dispatcher = StreamBatchDispatcher(tokenizer) + + # Wire the batched stream-flush hook: per-seq stream callbacks only buffer + # their chunks into a thread-local; the engine core manager's output thread + # calls this flush after each step's callbacks to drain the buffer into the + # per-request asyncio queues (one call_soon_threadsafe per event loop). + # Registered lazily here to avoid the api_server <-> engine_core_mgr import + # cycle; the core manager leaves the hook as None until this resolves it. + engine.core_mgr._flush_stream_batch_fn = flush_stream_batch import signal diff --git a/atom/entrypoints/openai/chat_encoder_adapters.py b/atom/entrypoints/openai/chat_encoder_adapters.py new file mode 100644 index 0000000000..11773b6265 --- /dev/null +++ b/atom/entrypoints/openai/chat_encoder_adapters.py @@ -0,0 +1,68 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""Model-scoped adapters for dynamically loaded chat encoders.""" + +from collections.abc import Callable +from dataclasses import dataclass +from typing import Any + +MessageEncoder = Callable[..., str] +MessagePreparer = Callable[[list[dict], list[dict] | None], list[dict]] + + +def _copy_messages( + messages: list[dict], _tools: list[dict] | None = None +) -> list[dict]: + """Return shallow message copies without model-specific rewriting.""" + return [dict(message) for message in messages] + + +def _prepare_deepseek_v4_messages( + messages: list[dict], tools: list[dict] | None +) -> list[dict]: + """Prepare the internal message shape expected by DSV4 ``encode_messages``. + + The DeepSeek-V4 reference encoder reads tool schemas from a system message's + ``tools`` field. Match vLLM's model-specific tokenizer wrapper by prepending + a synthetic tool-carrying system message, without reordering or merging the + caller's existing messages. + """ + prepared = _copy_messages(messages) + if tools: + prepared.insert(0, {"role": "system", "tools": tools}) + return prepared + + +@dataclass(frozen=True) +class MessageEncoderAdapter: + """A raw model encoder plus its model-specific message preparation.""" + + name: str + encode: MessageEncoder + prepare_messages: MessagePreparer + supports_tools: bool = False + + def __call__(self, messages: list[dict], **kwargs: Any) -> str: + """Preserve the callable behavior of the former encoder return value.""" + return self.encode(messages, **kwargs) + + +_PREPARERS: dict[str, tuple[MessagePreparer, bool]] = { + "encoding_dsv4": (_prepare_deepseek_v4_messages, True), +} + + +def build_message_encoder_adapter( + module_name: str, encoder: MessageEncoder +) -> MessageEncoderAdapter: + """Build an adapter registered for ``module_name`` or an identity adapter.""" + prepare_messages, supports_tools = _PREPARERS.get( + module_name, (_copy_messages, False) + ) + return MessageEncoderAdapter( + name=module_name, + encode=encoder, + prepare_messages=prepare_messages, + supports_tools=supports_tools, + ) diff --git a/atom/entrypoints/openai/chat_encoders.py b/atom/entrypoints/openai/chat_encoders.py index d3fb0462ce..4791912af5 100644 --- a/atom/entrypoints/openai/chat_encoders.py +++ b/atom/entrypoints/openai/chat_encoders.py @@ -14,13 +14,16 @@ import importlib.util import logging import os -from typing import Any, Callable, List, Optional +from typing import Any from huggingface_hub import snapshot_download -logger = logging.getLogger("atom") +from .chat_encoder_adapters import ( + MessageEncoderAdapter, + build_message_encoder_adapter, +) -MessageEncoder = Callable[..., str] +logger = logging.getLogger("atom") def _resolve_model_path(model: str) -> str: @@ -32,7 +35,7 @@ def _resolve_model_path(model: str) -> str: return model -def _load_encoder_from_dir(model_path: str) -> Optional[MessageEncoder]: +def _load_encoder_from_dir(model_path: str) -> MessageEncoderAdapter | None: """Look for ``/encoding/encoding_*.py`` and load ``encode_messages``. Returns ``None`` when the directory or matching file is absent (model uses @@ -72,10 +75,10 @@ def encode(messages, **kwargs): return raw(messages, **kwargs) logger.info(f"Loaded message encoder from {enc_path}") - return encode + return build_message_encoder_adapter(module_name, encode) -def load_custom_message_encoder(model_path: str) -> Optional[MessageEncoder]: +def load_custom_message_encoder(model_path: str) -> MessageEncoderAdapter | None: """Probe ``model_path`` once at startup for a custom message encoder. Returns the encoder, or ``None`` when the model uses the standard Jinja @@ -87,28 +90,29 @@ def load_custom_message_encoder(model_path: str) -> Optional[MessageEncoder]: def apply_chat_template( tokenizer: Any, - custom_encoder: Optional[MessageEncoder], - messages: List[dict], + custom_encoder: MessageEncoderAdapter | None, + messages: list[dict], *, - tools: Optional[List[dict]] = None, + tools: list[dict] | None = None, **kwargs: Any, ) -> str: """Render ``messages`` to a prompt string. Dispatches to ``custom_encoder`` if one was discovered for this model, otherwise to ``tokenizer.apply_chat_template``. Jinja-only kwargs - (``tokenize``, ``add_generation_prompt``) are stripped on the custom - path; ``tools`` are forwarded only on the Jinja path (custom encoders - don't currently have a tools API — caller is warned and tools are - dropped). + (``tokenize``, ``add_generation_prompt``) are stripped on the custom path. + Model-scoped adapters prepare tools for custom encoders that support them; + the generic path does not apply DeepSeek-V4-specific message rewriting. """ if custom_encoder is not None: for k in ("tokenize", "add_generation_prompt"): kwargs.pop(k, None) - if tools: + if tools and not custom_encoder.supports_tools: logger.warning( - "tools= is not supported with the custom message encoder; ignoring." + "tools= is not supported by custom message encoder %s; ignoring.", + custom_encoder.name, ) + messages = custom_encoder.prepare_messages(messages, tools) return custom_encoder(messages, **kwargs) kwargs["tokenize"] = False diff --git a/atom/entrypoints/openai/streaming_dispatch.py b/atom/entrypoints/openai/streaming_dispatch.py new file mode 100644 index 0000000000..71b4662855 --- /dev/null +++ b/atom/entrypoints/openai/streaming_dispatch.py @@ -0,0 +1,139 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""Batched cross-thread dispatch for streaming model output.""" + +import threading +from asyncio import AbstractEventLoop, Queue +from collections.abc import Hashable +from dataclasses import dataclass, field +from typing import Any + + +@dataclass +class IncrementalStreamDetokenizer: + """Decode token deltas without emitting incomplete UTF-8 characters.""" + + tokenizer: Any + tokens: list[int] = field(default_factory=list) + prefix_offset: int = 0 + read_offset: int = 0 + + def update(self, token_ids: list[int], finished: bool) -> str: + self.tokens.extend(token_ids) + prefix_text = self.tokenizer.decode( + self.tokens[self.prefix_offset : self.read_offset], + skip_special_tokens=True, + ) + new_text = self.tokenizer.decode( + self.tokens[self.prefix_offset :], + skip_special_tokens=True, + ) + + if len(new_text) > len(prefix_text) and not new_text.endswith("\ufffd"): + delta = new_text[len(prefix_text) :] + self.prefix_offset = self.read_offset + self.read_offset = len(self.tokens) + return delta + if finished: + return new_text[len(prefix_text) :] + return "" + + +@dataclass +class _BufferedChunk: + loop: AbstractEventLoop + queue: Queue + state_key: Hashable + chunk: dict + tag: int | None + + +class StreamBatchDispatcher: + """Collect one engine step per output thread and dispatch it by event loop.""" + + def __init__(self, tokenizer: Any): + self.tokenizer = tokenizer + self._thread_local = threading.local() + self._states: dict[Hashable, IncrementalStreamDetokenizer] = {} + self._states_lock = threading.Lock() + + def enqueue( + self, + *, + loop: AbstractEventLoop, + queue: Queue, + state_key: Hashable, + chunk: dict, + tag: int | None = None, + ) -> None: + """Buffer a raw chunk until the current engine step is flushed.""" + buf = getattr(self._thread_local, "buf", None) + if buf is None: + buf = self._thread_local.buf = [] + buf.append( + _BufferedChunk( + loop=loop, + queue=queue, + state_key=state_key, + chunk=chunk, + tag=tag, + ) + ) + + def flush(self) -> None: + """Detokenize buffered chunks and schedule one drain per event loop.""" + buf = getattr(self._thread_local, "buf", None) + if not buf: + return + self._thread_local.buf = [] + + by_loop: dict[AbstractEventLoop, list[tuple[Queue, Any]]] = {} + for item in buf: + state = self._get_state(item.state_key) + item.chunk["text"] = state.update( + item.chunk.get("token_ids") or [], + bool(item.chunk.get("finished")), + ) + if item.chunk.get("finished"): + self._drop_state(item.state_key, state) + + payload = item.chunk if item.tag is None else (item.tag, item.chunk) + by_loop.setdefault(item.loop, []).append((item.queue, payload)) + + for loop, items in by_loop.items(): + loop.call_soon_threadsafe(self._drain_into_queues, items) + + def discard_request(self, request_id: str) -> None: + """Drop direct and fan-out detokenizer state after request cleanup.""" + with self._states_lock: + keys = [ + key + for key in self._states + if key == request_id + or (isinstance(key, tuple) and key and key[0] == request_id) + ] + for key in keys: + self._states.pop(key, None) + + def _get_state(self, state_key: Hashable) -> IncrementalStreamDetokenizer: + with self._states_lock: + state = self._states.get(state_key) + if state is None: + state = self._states[state_key] = IncrementalStreamDetokenizer( + self.tokenizer + ) + return state + + def _drop_state( + self, state_key: Hashable, state: IncrementalStreamDetokenizer + ) -> None: + with self._states_lock: + if self._states.get(state_key) is state: + self._states.pop(state_key) + + @staticmethod + def _drain_into_queues(items: list[tuple[Queue, Any]]) -> None: + """Run on the target event loop and deliver each prepared payload.""" + for queue, payload in items: + queue.put_nowait(payload) diff --git a/atom/entrypoints/openai/tool_parser.py b/atom/entrypoints/openai/tool_parser.py deleted file mode 100644 index 549277e8d9..0000000000 --- a/atom/entrypoints/openai/tool_parser.py +++ /dev/null @@ -1,466 +0,0 @@ -# SPDX-License-Identifier: MIT -# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. - -"""Tool call parser for models that output tool calls. - -Two on-the-wire formats are auto-detected and normalized into the OpenAI -``tool_calls`` structure: - -1. Kimi-K2 special-token format:: - - <|tool_calls_section_begin|> - <|tool_call_begin|>functions.NAME:INDEX<|tool_call_argument_begin|>ARGS_JSON<|tool_call_end|> - <|tool_calls_section_end|> - -2. Qwen3 (qwen3_coder / qwen3_xml) XML format:: - - - - VALUE - ... - - - -The Qwen XML carries no value types, so when the request's ``tools`` schema is -supplied each parameter is coerced to the declared JSON-Schema type (int, float, -bool, null, object, array); otherwise it is left as a string. This mirrors the -qwen3_coder/qwen3_xml parsers in vLLM and SGLang. - -OpenAI format: - {"tool_calls": [{"id": "call_0", "type": "function", - "function": {"name": "NAME", "arguments": "ARGS_JSON"}}]} -""" - -import ast -import json -import re -import uuid -from dataclasses import dataclass -from typing import Any, Dict, List, Optional, Tuple - - -def _unique_tool_call_id() -> str: - # OpenAI tool_call ids must be unique across the whole conversation, not just - # within one response. A per-response index (call_0, call_1, ...) collides - # across turns -> clients (e.g. qwen-code) dedupe by id and silently ignore - # every repeat, causing an infinite tool-call retry loop. Use a random id. - return f"call_{uuid.uuid4().hex}" - - -@dataclass -class ToolCall: - """Parsed tool call in OpenAI format.""" - - id: str - type: str - function: Dict[str, str] - - def to_dict(self) -> Dict[str, Any]: - return {"id": self.id, "type": self.type, "function": self.function} - - -# --------------------------------------------------------------------------- -# Qwen3 XML tool-call format (qwen3_coder / qwen3_xml) -# --------------------------------------------------------------------------- - -_QWEN_TOOL_PREFIX = "||(?=)|$)", - re.DOTALL, -) - - -def _is_qwen_xml(text: str) -> bool: - """Detect the Qwen3 XML tool-call format (and not the Kimi token format).""" - return _QWEN_TOOL_PREFIX in text and "<|tool_calls_section_begin|>" not in text - - -def _build_param_types(tools: Optional[list]) -> Dict[str, Dict[str, Any]]: - """Map ``function_name -> {param_name: json_schema_type}`` from request tools. - - Accepts OpenAI (``{"type": "function", "function": {...}}``) and bare - (``{"name": ..., "parameters"/"input_schema": {...}}``) tool entries. - """ - out: Dict[str, Dict[str, Any]] = {} - for tool in tools or []: - if not isinstance(tool, dict): - continue - fn = tool.get("function", tool) - if not isinstance(fn, dict): - continue - name = fn.get("name") - if not name: - continue - schema = fn.get("parameters") or fn.get("input_schema") or {} - props = schema.get("properties") if isinstance(schema, dict) else None - out[name] = { - k: (v.get("type") if isinstance(v, dict) else None) - for k, v in (props or {}).items() - } - return out - - -def _coerce_param_value(value: str, ptype: Any) -> Any: - """Coerce a string parameter value to its declared JSON-Schema type. - - No schema type (string/unknown) -> returned unchanged. Conversion failures - fall back to the original string rather than raising. - """ - v = value.strip("\n") - if ptype is None: - return v - t = str(ptype).lower() - try: - if t in ("string", "str", "text", "varchar", "char", "enum"): - return v - if t in ("null", "none"): - return None - if t.startswith(("int", "uint", "long", "short", "unsigned")): - return int(v) - if t.startswith(("num", "float", "double", "decimal")): - f = float(v) - return int(f) if f.is_integer() else f - if t.startswith(("bool", "binary")): - return v.strip().lower() == "true" - if t.startswith(("object", "dict", "map", "array", "list", "tuple")): - try: - return json.loads(v) - except Exception: - return ast.literal_eval(v) # safer for single-quoted Python literals - except Exception: - return v - return v - - -def _parse_qwen_function( - fn_text: str, param_types: Dict[str, Dict[str, Any]], index: int -) -> Optional[ToolCall]: - """Parse the inside of one ``...`` block into a ToolCall.""" - gt = fn_text.find(">") - if gt == -1: - return None - name = fn_text[:gt].strip() - if not name: - return None - body = fn_text[gt + 1 :] - types = param_types.get(name, {}) - args: Dict[str, Any] = {} - for pm in _QWEN_PARAM_RE.finditer(body): - seg = pm.group(1) - if seg is None: - continue - pgt = seg.find(">") - if pgt == -1: - continue - pname = seg[:pgt].strip() - pval = seg[pgt + 1 :] - if pname: - args[pname] = _coerce_param_value(pval, types.get(pname)) - return ToolCall( - id=_unique_tool_call_id(), - type="function", - function={"name": name, "arguments": json.dumps(args, ensure_ascii=False)}, - ) - - -def _parse_qwen_xml(text: str, tools: Optional[list]) -> Tuple[str, List[ToolCall]]: - """Parse Qwen3 XML tool calls; return (leading_content, tool_calls).""" - param_types = _build_param_types(tools) - # Content precedes the first tool marker. - markers = [ - i for i in (text.find(""), text.find(_QWEN_TOOL_PREFIX)) if i != -1 - ] - content = text[: min(markers)] if markers else text - tool_calls: List[ToolCall] = [] - for fm in _QWEN_FUNCTION_RE.finditer(text): - fn_text = fm.group(1) if fm.group(1) is not None else fm.group(2) - if not fn_text: - continue - tc = _parse_qwen_function(fn_text, param_types, len(tool_calls)) - if tc is not None: - tool_calls.append(tc) - return content.strip(), tool_calls - - -def parse_tool_calls( - text: str, tools: Optional[list] = None -) -> Tuple[str, List[ToolCall]]: - """Parse tool calls from model output text. - - Args: - text: Raw model output that may contain tool calls (Kimi token format - or Qwen3 XML format). - tools: Optional request tool definitions; used to type-coerce Qwen XML - parameter values to their declared JSON-Schema types. - - Returns: - Tuple of (content_text, list_of_tool_calls). ``content_text`` has the - tool-call sections removed. - """ - # Qwen3 XML format - if _is_qwen_xml(text): - return _parse_qwen_xml(text, tools) - - # Kimi-K2 special-token format - section_match = re.search( - r"<\|tool_calls_section_begin\|>(.*?)<\|tool_calls_section_end\|>", - text, - flags=re.DOTALL, - ) - if not section_match: - # Check for unclosed section - unclosed = re.search( - r"<\|tool_calls_section_begin\|>(.*?)$", text, flags=re.DOTALL - ) - if unclosed: - content = text[: unclosed.start()] - tool_calls = _parse_tool_call_entries(unclosed.group(1)) - return content.strip(), tool_calls - return text, [] - - content = text[: section_match.start()] - tool_calls = _parse_tool_call_entries(section_match.group(1)) - - return content.strip(), tool_calls - - -def _parse_tool_call_entries(section_text: str) -> List[ToolCall]: - """Parse individual tool call entries from the section content.""" - tool_calls = [] - pattern = re.compile( - r"<\|tool_call_begin\|>" - r"functions\.(\w+):(\d+)" - r"<\|tool_call_argument_begin\|>" - r"(.*?)" - r"<\|tool_call_end\|>", - re.DOTALL, - ) - for match in pattern.finditer(section_text): - name = match.group(1) - index = match.group(2) - arguments = match.group(3).strip() - tool_id = f"functions.{name}:{index}" - tool_calls.append( - ToolCall( - id=tool_id, - type="function", - function={"name": name, "arguments": arguments}, - ) - ) - return tool_calls - - -@dataclass -class ToolCallStreamParser: - """Stateful streaming parser for tool calls (Kimi tokens or Qwen3 XML). - - Processes text chunks and emits structured events: - - ("content", text) — regular content before tool calls - - ("tool_call_start", {"index": N, "id": ..., "function": {"name": ..., "arguments": ""}}) - - ("tool_call_args", {"index": N, "function": {"arguments": chunk}}) - - ("tool_call_end", None) — all tool calls complete - - The wire format is auto-detected from the first chunks. For the Qwen3 XML - format content is streamed normally and the ```` block is buffered - and parsed when complete (robust against partial-XML streaming edge cases); - ``tools`` enables JSON-Schema type coercion of parameter values. - - Kimi states: - 0 = normal content (no tool call tokens seen) - 1 = inside tool_calls_section (buffering) - 2 = done (after tool_calls_section_end) - """ - - state: int = 0 - buf: str = "" - current_index: int = 0 - _emitted_calls: int = 0 - tools: Optional[list] = None - fmt: Optional[str] = None # None (undecided) | "kimi" | "qwen" - - def process(self, text: str) -> list: - """Process a text chunk and return list of (event_type, data) tuples.""" - if self.fmt is None: - self.buf += text - if _QWEN_TOOL_PREFIX in self.buf or "" in self.buf: - self.fmt = "qwen" - elif "<|tool_calls_section_begin|>" in self.buf: - self.fmt = "kimi" - elif "<" not in self.buf and len(self.buf) > 8: - # No markup possible yet; emit accumulated content and stay undecided. - out = [("content", self.buf)] - self.buf = "" - return out - else: - return [] - # Format decided: replay the accumulated buffer through the handler. - text, self.buf = self.buf, "" - - if self.fmt == "qwen": - return self._process_qwen(text) - return self._process_kimi(text) - - # -- Qwen3 XML ---------------------------------------------------------- - def _process_qwen(self, text: str) -> list: - results: list = [] - self.buf += text - if self.state == 0: - markers = [ - i - for i in ( - self.buf.find(""), - self.buf.find(_QWEN_TOOL_PREFIX), - ) - if i != -1 - ] - if markers: - m = min(markers) - before = self.buf[:m] - if before: - results.append(("content", before)) - self.buf = self.buf[m:] - self.state = 1 - else: - # Emit content but hold back a possible partial '<...' marker tail. - cut = self.buf.rfind("<") - if cut == -1: - if self.buf: - results.append(("content", self.buf)) - self.buf = "" - elif cut > 0: - results.append(("content", self.buf[:cut])) - self.buf = self.buf[cut:] - return results - - def _flush_qwen(self) -> list: - results: list = [] - if self.state == 0: - if self.buf: - results.append(("content", self.buf)) - self.buf = "" - return results - # state 1: parse the complete (or trailing) tool-call block. - _content, tool_calls = _parse_qwen_xml(self.buf, self.tools) - self.buf = "" - for tc in tool_calls: - tc.id = _unique_tool_call_id() - results.append( - ( - "tool_call_start", - { - "index": self.current_index, - "id": tc.id, - "type": "function", - "function": {"name": tc.function["name"], "arguments": ""}, - }, - ) - ) - results.append( - ( - "tool_call_args", - { - "index": self.current_index, - "function": {"arguments": tc.function["arguments"]}, - }, - ) - ) - self.current_index += 1 - self._emitted_calls += 1 - if self._emitted_calls > 0: - results.append(("tool_call_end", None)) - return results - - # -- Kimi tokens -------------------------------------------------------- - def _process_kimi(self, text: str) -> list: - results = [] - - if self.state == 0: - self.buf += text - if "<|tool_calls_section_begin|>" in self.buf: - before = self.buf.split("<|tool_calls_section_begin|>")[0] - if before: - results.append(("content", before)) - self.state = 1 - self.buf = self.buf.split("<|tool_calls_section_begin|>", 1)[1] - results.extend(self._process_buffer()) - elif "<|tool" not in self.buf and len(self.buf) > 30: - results.append(("content", self.buf)) - self.buf = "" - - elif self.state == 1: - self.buf += text - if "<|tool_calls_section_end|>" in self.buf: - remaining = self.buf.split("<|tool_calls_section_end|>")[0] - self.buf = remaining - results.extend(self._process_buffer()) - results.append(("tool_call_end", None)) - self.state = 2 - self.buf = "" - else: - results.extend(self._process_buffer()) - - return results - - def _process_buffer(self) -> list: - """Extract complete tool call entries from the buffer.""" - results = [] - while "<|tool_call_begin|>" in self.buf and "<|tool_call_end|>" in self.buf: - match = re.search( - r"<\|tool_call_begin\|>" - r"functions\.(\w+):(\d+)" - r"<\|tool_call_argument_begin\|>" - r"(.*?)" - r"<\|tool_call_end\|>", - self.buf, - re.DOTALL, - ) - if not match: - break - - name = match.group(1) - index = int(match.group(2)) - arguments = match.group(3).strip() - - tool_id = f"functions.{name}:{index}" - results.append( - ( - "tool_call_start", - { - "index": index, - "id": tool_id, - "type": "function", - "function": {"name": name, "arguments": ""}, - }, - ) - ) - if arguments: - results.append( - ( - "tool_call_args", - {"index": index, "function": {"arguments": arguments}}, - ) - ) - - self.buf = self.buf[match.end() :] - self._emitted_calls += 1 - - return results - - def flush(self) -> list: - """Flush remaining buffer content.""" - if self.fmt == "qwen": - return self._flush_qwen() - results = [] - if self.state == 0 and self.buf: - results.append(("content", self.buf)) - self.buf = "" - elif self.state == 1: - results.extend(self._process_buffer()) - if self._emitted_calls > 0: - results.append(("tool_call_end", None)) - elif self.fmt is None and self.buf: - # Undecided at EOS: no tool markers ever appeared -> plain content. - results.append(("content", self.buf)) - self.buf = "" - return results diff --git a/atom/entrypoints/openai/tool_parser/__init__.py b/atom/entrypoints/openai/tool_parser/__init__.py new file mode 100644 index 0000000000..6c7b286241 --- /dev/null +++ b/atom/entrypoints/openai/tool_parser/__init__.py @@ -0,0 +1,61 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""Tool call parsing for models that emit tool calls in their text output. + +Five on-the-wire formats are auto-detected and normalized into the OpenAI +``tool_calls`` structure: + +============================== =============================================== +Module Format +============================== =============================================== +`kimi_tool_parser` Kimi-K2 ``<|tool_call_begin|>`` special tokens +`qwen3_tool_parser` Qwen3 (qwen3_coder / qwen3_xml) ```` markup +`glm_tool_parser` GLM-4.5/4.6/5.x ````/```` +`minimax_tool_parser` MiniMax-M3 ``]<]minimax[>[``-prefixed tags +============================== =============================================== + +Formats other than Kimi carry no value types on the wire, so when the request's +``tools`` schema is supplied each parameter is coerced to its declared +JSON-Schema type; otherwise it is left as a string. + +Two entry points, both format-agnostic: + +- :func:`parse_tool_calls` — a complete output -> ``(content, [ToolCall])`` +- :class:`ToolCallStreamParser` — chunks -> ``(event_type, data)`` tuples + +OpenAI format:: + + {"tool_calls": [{"id": "call_0", "type": "function", + "function": {"name": "NAME", "arguments": "ARGS_JSON"}}]} + +To add a format: implement :class:`~.tool_parser.ToolCallParser` (or, if it +buffers from a start marker like most do, +:class:`~.tool_parser.BufferedMarkerParser`) in its own +``_tool_parser.py``, then register it in :mod:`.registry` — in both +``_DETECT_ORDER`` and ``sniff_stream``, whose ordering constraints are +documented there. +""" + +from .deepseekv4_tool_parser import DsmlParser +from .glm_tool_parser import GlmParser +from .kimi_tool_parser import KimiParser +from .minimax_tool_parser import MiniMaxParser +from .qwen3_tool_parser import QwenXmlParser +from .registry import parse_tool_calls +from .stream import ToolCallStreamParser +from .tool_parser import BufferedMarkerParser, ToolCall, ToolCallParser + +__all__ = [ + "BufferedMarkerParser", + "DsmlParser", + "GlmParser", + "KimiParser", + "MiniMaxParser", + "QwenXmlParser", + "ToolCall", + "ToolCallParser", + "ToolCallStreamParser", + "parse_tool_calls", +] diff --git a/atom/entrypoints/openai/tool_parser/deepseekv4_tool_parser.py b/atom/entrypoints/openai/tool_parser/deepseekv4_tool_parser.py new file mode 100644 index 0000000000..559338ba67 --- /dev/null +++ b/atom/entrypoints/openai/tool_parser/deepseekv4_tool_parser.py @@ -0,0 +1,186 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""DeepSeek-V4 DSML tool-call format:: + + <｜DSML｜tool_calls> + <｜DSML｜invoke name="NAME"> + <｜DSML｜parameter name="PNAME" string="true|false">VALUE + ... + + + +``string="true"`` -> value is a raw string; ``string="false"`` -> value is JSON. +DeepSeek-V4-Flash occasionally malforms this (singular ``tool_call``, a missing +``invoke`` wrapper, or params without ``string=``); the parser recovers those +best-effort: it infers a dropped tool name from the parameter signature vs the +request's ``tools`` and infers a missing value type from the schema / JSON. +""" + +import json +import re +from typing import Any, ClassVar + +from .schema import build_param_types, coerce_param_value +from .tool_parser import BufferedMarkerParser, ToolCall, unique_tool_call_id + +_DSML = "｜DSML｜" +# The model often DROPS the ``｜DSML｜`` marker and emits bare +# ````/````/```` tags, so the marker +# is matched OPTIONALLY everywhere. +_OPT = r"(?:" + re.escape(_DSML) + r")?" # optional ｜DSML｜ prefix +_PARAM_RE = re.compile( + r"<" + _OPT + r'parameter\s+name="(.*?)"(?:\s+string="(true|false)")?\s*>' + r"(.*?)", + re.DOTALL, +) +# Long-form `...` OR self-closing `` +# (the zero-arg shape; group(2) is None for self-closing). Matches SGLang's V4 +# detector, which accepts both. +_INVOKE_RE = re.compile( + r"<" + _OPT + r'invoke\s+name="(.*?)"\s*(?:/>|>(.*?))", + re.DOTALL, +) + + +def _unwrap_wrapper_args(args: Any, allowed: set) -> Any: + """Strip spurious ``{"arguments": {...}}`` / ``{"input": {...}}`` envelopes. + + Non-tuned models (DeepSeek-V4-Pro) frequently wrap the real args in an extra + ``arguments``/``input`` object — sometimes nested 2-3 deep, or stringified — + so a call meant as ``{"cmd": "ls"}`` arrives as ``{"arguments": {"cmd": + "ls"}}``. Recursively unwrap while the sole key is a wrapper that is NOT + itself a declared param of the tool. Mirrors vLLM's ``_unwrap_wrapper_args`` + (deepseek_v4.py).""" + for _ in range(4): # bounded against pathological nesting + if not (isinstance(args, dict) and len(args) == 1): + break + ((k, v),) = args.items() + if k not in ("arguments", "input"): + break + if allowed and k in allowed: + break # this tool really has a param named arguments/input + if isinstance(v, str): + try: + v = json.loads(v) + except (ValueError, TypeError): + break + if not isinstance(v, dict): + break + args = v + return args + + +def _coerce(value: str, string_attr: str | None, ptype: Any) -> Any: + """Decode one ```` body. + + Deliberately not :func:`~.schema.coerce_json_or_raw`: on a JSON-decode miss + this falls back to ``value.strip()`` where that one falls back to + ``value.strip("\\n")``, which differs for values with surrounding spaces. + """ + if string_attr == "true": + return value + if string_attr == "false": + try: + return json.loads(value) + except (ValueError, TypeError): + return value + # attr absent -> use declared schema type if known, else infer via JSON. + if ptype is not None: + return coerce_param_value(value, ptype) + v = value.strip() + try: + return json.loads(v) + except (ValueError, TypeError): + return v + + +def _infer_name(arg_names: set, param_types: dict[str, dict[str, Any]]) -> str | None: + """Pick the request tool whose parameter set best matches ``arg_names``.""" + best, best_score = None, -1e9 + for name, props in param_types.items(): + p = set(props) + if not p: + continue + score = len(p & arg_names) - 0.1 * len(p ^ arg_names) + if score > best_score: + best_score, best = score, name + return best + + +class DsmlParser(BufferedMarkerParser): + NAME: ClassVar[str] = "dsml" + # Region-start markers, both marked and marker-less variants. + START_MARKERS: ClassVar[tuple[str, ...]] = ( + "<" + _DSML + "tool_call", # marked (covers tool_call / tool_calls) + "<" + _DSML + "invoke", # marked invoke + "", # marker-less section open + ) + + # detect() is inherited: any start marker present means DSML. + + @classmethod + def parse(cls, text: str, tools: list | None) -> tuple[str, list[ToolCall]]: + """Parse DeepSeek-V4 DSML tool calls; return (leading_content, tool_calls).""" + param_types = build_param_types(tools) + start = cls.find_start(text) + if start == -1: + return text.strip(), [] + content = text[:start] + region = text[start:] + + calls: list[tuple[str, dict[str, Any]]] = [] + invokes = list(_INVOKE_RE.finditer(region)) + if invokes: + for m in invokes: + name = m.group(1) + body = m.group(2) or "" # None for self-closing + types = param_types.get(name, {}) + args: dict[str, Any] = { + pm.group(1): _coerce( + pm.group(3), pm.group(2), types.get(pm.group(1)) + ) + for pm in _PARAM_RE.finditer(body) + } + # Direct-JSON parameter body (DSML "Format 2", also accepted by + # vLLM/SGLang): ` { "k": "v" } ` with no + # tags. Falls through here with empty args; recover them. + if not args: + stripped = body.strip() + if stripped.startswith("{"): + try: + parsed = json.loads(stripped) + if isinstance(parsed, dict): + args = parsed + except (ValueError, TypeError): + pass + args = _unwrap_wrapper_args(args, set(types)) + calls.append((name, args)) + else: + # malformed: no complete invoke wrapper -> collect params, infer tool name + raw = { + pm.group(1): (pm.group(3), pm.group(2)) + for pm in _PARAM_RE.finditer(region) + } + if raw: + name = _infer_name(set(raw), param_types) or "unknown" + types = param_types.get(name, {}) + args = {k: _coerce(v, s, types.get(k)) for k, (v, s) in raw.items()} + args = _unwrap_wrapper_args(args, set(types)) + calls.append((name, args)) + + tool_calls = [ + ToolCall( + id=unique_tool_call_id(), + type="function", + function={ + "name": name, + "arguments": json.dumps(args, ensure_ascii=False), + }, + ) + for name, args in calls + ] + if _DSML in content: # scrub any stray marker fragment + content = content.split("<" + _DSML, 1)[0] + return content.strip(), tool_calls diff --git a/atom/entrypoints/openai/tool_parser/glm_tool_parser.py b/atom/entrypoints/openai/tool_parser/glm_tool_parser.py new file mode 100644 index 0000000000..08bbdc7dc4 --- /dev/null +++ b/atom/entrypoints/openai/tool_parser/glm_tool_parser.py @@ -0,0 +1,78 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""GLM-4.5 / 4.6 / 5.x tool-call format:: + + NAME + K1V1 + K2V2 + ... + +The function name follows the opening tag directly (no ``(.*?)|(.*)$", re.DOTALL) +_ARG_RE = re.compile( + r"(.*?)\s*" + r"(.*?)(?:|(?=)|(?=)|$)", + re.DOTALL, +) + + +class GlmParser(BufferedMarkerParser): + NAME: ClassVar[str] = "glm" + START_MARKERS: ClassVar[tuple[str, ...]] = ("",) + + @classmethod + def detect(cls, text: str) -> bool: + """Detect the GLM ``...`` format (never Qwen/DSML).""" + if QWEN_TOOL_PREFIX in text: # ' Qwen, not GLM + return False + return "" in text or "" in text + + @classmethod + def parse(cls, text: str, tools: list | None) -> tuple[str, list[ToolCall]]: + """Parse GLM tool calls; return (leading_content, tool_calls).""" + param_types = build_param_types(tools) + start = text.find("") + if start == -1: + return text.strip(), [] + content = text[:start] + tool_calls: list[ToolCall] = [] + for m in _TOOLCALL_RE.finditer(text): + body = m.group(1) if m.group(1) is not None else m.group(2) + if not body: + continue + ak = body.find("") + name = (body if ak == -1 else body[:ak]).strip() + if not name: + continue + types = param_types.get(name, {}) + args: dict[str, Any] = {} + for pm in _ARG_RE.finditer(body): + k = pm.group(1).strip() + if k: + args[k] = coerce_json_or_raw(pm.group(2), types.get(k)) + tool_calls.append( + ToolCall( + id=unique_tool_call_id(), + type="function", + function={ + "name": name, + "arguments": json.dumps(args, ensure_ascii=False), + }, + ) + ) + return content.strip(), tool_calls diff --git a/atom/entrypoints/openai/tool_parser/kimi_tool_parser.py b/atom/entrypoints/openai/tool_parser/kimi_tool_parser.py new file mode 100644 index 0000000000..792ff9e45c --- /dev/null +++ b/atom/entrypoints/openai/tool_parser/kimi_tool_parser.py @@ -0,0 +1,159 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""Kimi-K2 special-token tool-call format:: + + <|tool_calls_section_begin|> + <|tool_call_begin|>functions.NAME:INDEX<|tool_call_argument_begin|>ARGS_JSON<|tool_call_end|> + <|tool_calls_section_end|> + +Unlike the XML-ish formats this one is self-delimiting, so entries can be +emitted as soon as their ``<|tool_call_end|>`` arrives rather than buffering the +whole block. It therefore implements streaming itself instead of inheriting +:class:`~.tool_parser.BufferedMarkerParser`. + +Arguments are already JSON on the wire, so no schema coercion is applied and +``tools`` is unused. The call id is the model's own ``functions.NAME:INDEX`` +rather than a random one, and ``index`` comes from the wire too. +""" + +import re +from typing import ClassVar + +from .tool_parser import ToolCall, ToolCallParser + +KIMI_SECTION_BEGIN = "<|tool_calls_section_begin|>" +KIMI_SECTION_END = "<|tool_calls_section_end|>" + +_SECTION_RE = re.compile( + re.escape(KIMI_SECTION_BEGIN) + r"(.*?)" + re.escape(KIMI_SECTION_END), + re.DOTALL, +) +_UNCLOSED_RE = re.compile(re.escape(KIMI_SECTION_BEGIN) + r"(.*?)$", re.DOTALL) +_ENTRY_RE = re.compile( + r"<\|tool_call_begin\|>" + r"functions\.(\w+):(\d+)" + r"<\|tool_call_argument_begin\|>" + r"(.*?)" + r"<\|tool_call_end\|>", + re.DOTALL, +) + + +def _parse_entries(section_text: str) -> list[ToolCall]: + """Parse individual tool call entries from the section content.""" + tool_calls = [] + for match in _ENTRY_RE.finditer(section_text): + name = match.group(1) + index = match.group(2) + arguments = match.group(3).strip() + tool_id = f"functions.{name}:{index}" + tool_calls.append( + ToolCall( + id=tool_id, + type="function", + function={"name": name, "arguments": arguments}, + ) + ) + return tool_calls + + +class KimiParser(ToolCallParser): + """States: 0 = plain content, 1 = inside section, 2 = section closed.""" + + NAME: ClassVar[str] = "kimi" + + @classmethod + def detect(cls, text: str) -> bool: + return KIMI_SECTION_BEGIN in text + + @classmethod + def parse(cls, text: str, tools: list | None) -> tuple[str, list[ToolCall]]: + section_match = _SECTION_RE.search(text) + if not section_match: + # Unclosed section: the model was cut off mid-block; salvage whatever + # complete entries it managed to emit. + unclosed = _UNCLOSED_RE.search(text) + if unclosed: + content = text[: unclosed.start()] + return content.strip(), _parse_entries(unclosed.group(1)) + return text, [] + content = text[: section_match.start()] + return content.strip(), _parse_entries(section_match.group(1)) + + def process(self, text: str) -> list: + results: list = [] + + if self.state == 0: + self.buf += text + if KIMI_SECTION_BEGIN in self.buf: + before = self.buf.split(KIMI_SECTION_BEGIN)[0] + if before: + results.append(("content", before)) + self.state = 1 + self.buf = self.buf.split(KIMI_SECTION_BEGIN, 1)[1] + results.extend(self._drain_entries()) + elif "<|tool" not in self.buf and len(self.buf) > 30: + # No partial special token possible yet -> safe to release. + results.append(("content", self.buf)) + self.buf = "" + + elif self.state == 1: + self.buf += text + if KIMI_SECTION_END in self.buf: + self.buf = self.buf.split(KIMI_SECTION_END)[0] + results.extend(self._drain_entries()) + results.append(("tool_call_end", None)) + self.state = 2 + self.buf = "" + else: + results.extend(self._drain_entries()) + + return results + + def _drain_entries(self) -> list: + """Emit every complete tool-call entry sitting in the buffer.""" + results: list = [] + while "<|tool_call_begin|>" in self.buf and "<|tool_call_end|>" in self.buf: + match = _ENTRY_RE.search(self.buf) + if not match: + break + + name = match.group(1) + index = int(match.group(2)) + arguments = match.group(3).strip() + + results.append( + ( + "tool_call_start", + { + "index": index, + "id": f"functions.{name}:{index}", + "type": "function", + "function": {"name": name, "arguments": ""}, + }, + ) + ) + if arguments: + results.append( + ( + "tool_call_args", + {"index": index, "function": {"arguments": arguments}}, + ) + ) + + self.buf = self.buf[match.end() :] + self.emitted_calls += 1 + + return results + + def flush(self) -> list: + results: list = [] + if self.state == 0 and self.buf: + results.append(("content", self.buf)) + self.buf = "" + elif self.state == 1: + results.extend(self._drain_entries()) + if self.emitted_calls > 0: + results.append(("tool_call_end", None)) + return results diff --git a/atom/entrypoints/openai/tool_parser/minimax_tool_parser.py b/atom/entrypoints/openai/tool_parser/minimax_tool_parser.py new file mode 100644 index 0000000000..a78998aa80 --- /dev/null +++ b/atom/entrypoints/openai/tool_parser/minimax_tool_parser.py @@ -0,0 +1,79 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""MiniMax-M3 tool-call format. + +Every tag is prefixed by the ns_token ``]<]minimax[>[``:: + + ]<]minimax[>[ + ]<]minimax[>[ + ]<]minimax[>[value]<]minimax[>[ + ... + ]<]minimax[>[ + ]<]minimax[>[ + +Unlike DSML, parameters are named by the TAG itself (``Paris``), +not a ``name="..."`` attribute. Strip the ns_token first, then parse +/ pairs. Values: schema type wins, else JSON, else raw string. +""" + +import json +import re +from typing import Any, ClassVar + +from .schema import build_param_types, coerce_json_or_raw +from .tool_parser import BufferedMarkerParser, ToolCall, unique_tool_call_id + +MINIMAX_NS = "]<]minimax[>[" + +_INVOKE_RE = re.compile( + r'(.*?)|(.*)$', + re.DOTALL, +) +_PARAM_RE = re.compile(r"<([\w-]+)>(.*?)", re.DOTALL) + + +class MiniMaxParser(BufferedMarkerParser): + NAME: ClassVar[str] = "minimax" + START_MARKERS: ClassVar[tuple[str, ...]] = (MINIMAX_NS, "") + # The ns_token starts with ']', so a trailing ']' may be a partial marker. + HOLDBACK_CHARS: ClassVar[tuple[str, ...]] = ("<", "]") + + @classmethod + def detect(cls, text: str) -> bool: + """Detect the MiniMax-M3 ns_token tool-call format.""" + return MINIMAX_NS in text + + @classmethod + def parse(cls, text: str, tools: list | None) -> tuple[str, list[ToolCall]]: + """Parse MiniMax-M3 tool calls; return (leading_content, tool_calls).""" + param_types = build_param_types(tools) + clean = text.replace(MINIMAX_NS, "") + tc = clean.find("") + content = clean[:tc] if tc > 0 else ("" if tc == 0 else clean) + tool_calls: list[ToolCall] = [] + for m in _INVOKE_RE.finditer(clean): + name = m.group(1) if m.group(1) is not None else m.group(3) + body = m.group(2) if m.group(2) is not None else (m.group(4) or "") + if not name: + continue + name = name.strip() + types = param_types.get(name, {}) + args: dict[str, Any] = {} + for pm in _PARAM_RE.finditer(body): + k = pm.group(1).strip() + if k: + args[k] = coerce_json_or_raw(pm.group(2), types.get(k)) + tool_calls.append( + ToolCall( + id=unique_tool_call_id(), + type="function", + function={ + "name": name, + "arguments": json.dumps(args, ensure_ascii=False), + }, + ) + ) + for mk in ("", ""): + content = content.replace(mk, "") + return content.strip(), tool_calls diff --git a/atom/entrypoints/openai/tool_parser/qwen3_tool_parser.py b/atom/entrypoints/openai/tool_parser/qwen3_tool_parser.py new file mode 100644 index 0000000000..55c9551b73 --- /dev/null +++ b/atom/entrypoints/openai/tool_parser/qwen3_tool_parser.py @@ -0,0 +1,92 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""Qwen3 (qwen3_coder / qwen3_xml) XML tool-call format:: + + + + VALUE + ... + + + +The XML carries no value types, so parameters are coerced against the request's +``tools`` schema when supplied. Mirrors the qwen3_coder/qwen3_xml parsers in +vLLM and SGLang. +""" + +import json +import re +from typing import Any, ClassVar + +from .kimi_tool_parser import KIMI_SECTION_BEGIN +from .schema import build_param_types, coerce_param_value +from .tool_parser import BufferedMarkerParser, ToolCall, unique_tool_call_id + +# Also read by GlmParser.detect: ' +# apart from GLM's identically-named tag. +QWEN_TOOL_PREFIX = "||(?=)|$)", + re.DOTALL, +) + + +def _parse_function( + fn_text: str, param_types: dict[str, dict[str, Any]] +) -> ToolCall | None: + """Parse the inside of one ``...`` block into a ToolCall.""" + gt = fn_text.find(">") + if gt == -1: + return None + name = fn_text[:gt].strip() + if not name: + return None + body = fn_text[gt + 1 :] + types = param_types.get(name, {}) + args: dict[str, Any] = {} + for pm in _PARAM_RE.finditer(body): + seg = pm.group(1) + if seg is None: + continue + pgt = seg.find(">") + if pgt == -1: + continue + pname = seg[:pgt].strip() + pval = seg[pgt + 1 :] + if pname: + args[pname] = coerce_param_value(pval, types.get(pname)) + return ToolCall( + id=unique_tool_call_id(), + type="function", + function={"name": name, "arguments": json.dumps(args, ensure_ascii=False)}, + ) + + +class QwenXmlParser(BufferedMarkerParser): + NAME: ClassVar[str] = "qwen" + START_MARKERS: ClassVar[tuple[str, ...]] = ("", QWEN_TOOL_PREFIX) + + @classmethod + def detect(cls, text: str) -> bool: + """Detect the Qwen3 XML format (and not the Kimi token format).""" + return QWEN_TOOL_PREFIX in text and KIMI_SECTION_BEGIN not in text + + @classmethod + def parse(cls, text: str, tools: list | None) -> tuple[str, list[ToolCall]]: + """Parse Qwen3 XML tool calls; return (leading_content, tool_calls).""" + param_types = build_param_types(tools) + # Content precedes the first tool marker. + start = cls.find_start(text) + content = text[:start] if start != -1 else text + tool_calls: list[ToolCall] = [] + for fm in _FUNCTION_RE.finditer(text): + fn_text = fm.group(1) if fm.group(1) is not None else fm.group(2) + if not fn_text: + continue + tc = _parse_function(fn_text, param_types) + if tc is not None: + tool_calls.append(tc) + return content.strip(), tool_calls diff --git a/atom/entrypoints/openai/tool_parser/registry.py b/atom/entrypoints/openai/tool_parser/registry.py new file mode 100644 index 0000000000..7f854eaccc --- /dev/null +++ b/atom/entrypoints/openai/tool_parser/registry.py @@ -0,0 +1,102 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""Format detection. + +The wire format is sniffed from the model's own output rather than configured, +so detection order is load-bearing: several formats share tags and are only +told apart by a discriminator that a later entry would also match. The order +below is the single place that ordering is expressed — do not reorder without +re-reading the notes on each entry. +""" + +from .deepseekv4_tool_parser import DsmlParser +from .glm_tool_parser import GlmParser +from .kimi_tool_parser import KIMI_SECTION_BEGIN, KimiParser +from .minimax_tool_parser import MINIMAX_NS, MiniMaxParser +from .qwen3_tool_parser import QWEN_TOOL_PREFIX, QwenXmlParser +from .tool_parser import ToolCall, ToolCallParser + +# Checked in order on a COMPLETE output. Kimi is not listed: it is the terminal +# fallback, because its parse() also defines the "no tool calls at all" result. +# +# MiniMax before DSML — both use ``; MiniMax additionally +# prefixes every tag with the ns_token. +# GLM before Qwen — both use ``; GLM never emits ` tuple[str, list[ToolCall]]: + """Parse tool calls from a complete model output. + + Args: + text: Raw model output that may contain tool calls. + tools: Optional request tool definitions; used to type-coerce parameter + values to their declared JSON-Schema types. + + Returns: + Tuple of (content_text, list_of_tool_calls). ``content_text`` has the + tool-call sections removed. + """ + for parser in _DETECT_ORDER: + if parser.detect(text): + return parser.parse(text, tools) + # Kimi is terminal: when it finds no section either, it returns the text + # unchanged. Note that path does NOT strip, unlike every format that did + # match — preserved as-is, callers rely on plain content surviving verbatim. + return KimiParser.parse(text, tools) + + +# -- streaming sniff -------------------------------------------------------- +# +# Deciding on a PREFIX is strictly harder than on a complete output: a format's +# discriminator may not have arrived yet. These two sentinels say "cannot decide +# from what I have": + +# Enough plain text to be sure no marker is starting -> release it as content +# and stay undecided. +EMIT_CONTENT = object() +# Might still become a tool call -> keep buffering, emit nothing. +WAIT = object() + + +def sniff_stream(buf: str): + """Pick a parser from a partial stream, or return EMIT_CONTENT / WAIT. + + Deliberately NOT the same rules as the per-parser ``detect``: on a prefix, + GLM is only accepted on the unambiguous ```` (a bare ```` + could still turn out to be Qwen once ``" in buf: + return GlmParser + if QWEN_TOOL_PREFIX in buf: + return QwenXmlParser + if "" in buf: + # '' seen but neither '' + # (GLM) yet. A no-arg GLM call is complete once the closing tag arrives; + # otherwise wait for the sub-marker. + if "" in buf: + return GlmParser + return WAIT + if KIMI_SECTION_BEGIN in buf: + return KimiParser + if "<" not in buf and len(buf) > 8: + # No '<' anywhere, so no tag has started: release the text. The length + # floor is an unexplained heuristic carried over verbatim from the + # original parser — it trades a little first-token latency for not + # committing on a 1-2 char buffer. Note it does not actually rule out a + # partial MiniMax ns_token, whose first char is ']'. + return EMIT_CONTENT + return WAIT diff --git a/atom/entrypoints/openai/tool_parser/schema.py b/atom/entrypoints/openai/tool_parser/schema.py new file mode 100644 index 0000000000..162c574ea4 --- /dev/null +++ b/atom/entrypoints/openai/tool_parser/schema.py @@ -0,0 +1,87 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""Request-schema lookup and value coercion shared by every wire format. + +The XML-ish tool-call formats (Qwen, DSML, GLM, MiniMax) carry no value types on +the wire, so parameter values arrive as strings. When the request supplies a +``tools`` schema each value is coerced to its declared JSON-Schema type; +otherwise it is left alone. +""" + +import ast +import json +from typing import Any + + +def build_param_types(tools: list | None) -> dict[str, dict[str, Any]]: + """Map ``function_name -> {param_name: json_schema_type}`` from request tools. + + Accepts OpenAI (``{"type": "function", "function": {...}}``) and bare + (``{"name": ..., "parameters"/"input_schema": {...}}``) tool entries. + """ + out: dict[str, dict[str, Any]] = {} + for tool in tools or []: + if not isinstance(tool, dict): + continue + fn = tool.get("function", tool) + if not isinstance(fn, dict): + continue + name = fn.get("name") + if not name: + continue + schema = fn.get("parameters") or fn.get("input_schema") or {} + props = schema.get("properties") if isinstance(schema, dict) else None + out[name] = { + k: (v.get("type") if isinstance(v, dict) else None) + for k, v in (props or {}).items() + } + return out + + +def coerce_param_value(value: str, ptype: Any) -> Any: + """Coerce a string parameter value to its declared JSON-Schema type. + + No schema type (string/unknown) -> returned unchanged. Conversion failures + fall back to the original string rather than raising. + """ + v = value.strip("\n") + if ptype is None: + return v + t = str(ptype).lower() + try: + if t in ("string", "str", "text", "varchar", "char", "enum"): + return v + if t in ("null", "none"): + return None + if t.startswith(("int", "uint", "long", "short", "unsigned")): + return int(v) + if t.startswith(("num", "float", "double", "decimal")): + f = float(v) + return int(f) if f.is_integer() else f + if t.startswith(("bool", "binary")): + return v.strip().lower() == "true" + if t.startswith(("object", "dict", "map", "array", "list", "tuple")): + try: + return json.loads(v) + except (ValueError, TypeError): + return ast.literal_eval(v) # safer for single-quoted Python literals + except Exception: # noqa: BLE001 + return v + return v + + +def coerce_json_or_raw(value: str, ptype: Any) -> Any: + """Decode one untyped value: schema type wins, else JSON, else raw string. + + Shared by GLM (````) and MiniMax (````), whose templates both + render non-string values with ``tojson`` and string values raw. + """ + v = value.strip("\n") + if ptype is not None: + return coerce_param_value(v, ptype) + s = v.strip() + try: + return json.loads(s) + except (ValueError, TypeError): + return v diff --git a/atom/entrypoints/openai/tool_parser/stream.py b/atom/entrypoints/openai/tool_parser/stream.py new file mode 100644 index 0000000000..5a72ec5f29 --- /dev/null +++ b/atom/entrypoints/openai/tool_parser/stream.py @@ -0,0 +1,68 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""Streaming facade: sniff the format once, then delegate every chunk to it.""" + +from dataclasses import dataclass, field + +from .registry import EMIT_CONTENT, WAIT, sniff_stream +from .tool_parser import ToolCallParser + + +@dataclass +class ToolCallStreamParser: + """Stateful streaming parser; format is auto-detected from the first chunks. + + Emits structured events: + - ("content", text) — regular content before tool calls + - ("tool_call_start", {"index": N, "id": ..., "function": {"name": ..., "arguments": ""}}) + - ("tool_call_args", {"index": N, "function": {"arguments": chunk}}) + - ("tool_call_end", None) — all tool calls complete + + ``tools`` enables JSON-Schema type coercion of parameter values. It may be + assigned after construction (several call sites do) and is re-read on every + delegated call, so it takes effect as long as it is set before the stream + ends. + """ + + tools: list | None = None + # Pre-detection accumulator. Once a format is chosen this is handed to the + # concrete parser and never used again. + _buf: str = "" + _parser: ToolCallParser | None = field(default=None, repr=False) + + @property + def fmt(self) -> str | None: + """Detected format name, or None while still undecided.""" + return self._parser.NAME if self._parser is not None else None + + def process(self, text: str) -> list: + """Process a text chunk and return list of (event_type, data) tuples.""" + if self._parser is None: + self._buf += text + choice = sniff_stream(self._buf) + if choice is WAIT: + return [] + if choice is EMIT_CONTENT: + out = [("content", self._buf)] + self._buf = "" + return out + self._parser = choice(tools=self.tools) + # Replay everything accumulated while undecided. + text, self._buf = self._buf, "" + + self._parser.tools = self.tools + return self._parser.process(text) + + def flush(self) -> list: + """Flush remaining buffer content.""" + if self._parser is None: + # Undecided at EOS: no tool markers ever appeared -> plain content. + if self._buf: + out = [("content", self._buf)] + self._buf = "" + return out + return [] + + self._parser.tools = self.tools + return self._parser.flush() diff --git a/atom/entrypoints/openai/tool_parser/tool_parser.py b/atom/entrypoints/openai/tool_parser/tool_parser.py new file mode 100644 index 0000000000..a9838ad437 --- /dev/null +++ b/atom/entrypoints/openai/tool_parser/tool_parser.py @@ -0,0 +1,167 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""Parser interface and the shared buffered-marker streaming strategy. + +Every wire format implements :class:`ToolCallParser`. Four of the five +(Qwen / DSML / GLM / MiniMax) stream identically — buffer from the first start +marker, parse the whole block at flush — so that strategy lives once in +:class:`BufferedMarkerParser` and each format only declares its markers and its +``parse``. Kimi is the exception: its token format is self-delimiting, so it +emits tool calls incrementally and implements ``process``/``flush`` itself. +""" + +import uuid +from abc import ABC, abstractmethod +from dataclasses import dataclass +from typing import Any, ClassVar + + +def unique_tool_call_id() -> str: + # OpenAI tool_call ids must be unique across the whole conversation, not just + # within one response. A per-response index (call_0, call_1, ...) collides + # across turns -> clients (e.g. qwen-code) dedupe by id and silently ignore + # every repeat, causing an infinite tool-call retry loop. Use a random id. + return f"call_{uuid.uuid4().hex}" + + +@dataclass +class ToolCall: + """Parsed tool call in OpenAI format.""" + + id: str + type: str + function: dict[str, str] + + def to_dict(self) -> dict[str, Any]: + return {"id": self.id, "type": self.type, "function": self.function} + + +class ToolCallParser(ABC): + """One on-the-wire tool-call format. + + Class side is the stateless non-streaming path (``detect`` + ``parse``); + instance side is the stateful streaming path (``process`` + ``flush``). + """ + + NAME: ClassVar[str] + + def __init__(self, tools: list | None = None): + self.tools = tools + self.buf = "" + # 0 = still in plain content, 1 = inside a tool-call region. Kimi adds + # 2 = section closed; see KimiParser. + self.state = 0 + self.current_index = 0 + self.emitted_calls = 0 + + # -- non-streaming ------------------------------------------------------ + @classmethod + @abstractmethod + def detect(cls, text: str) -> bool: + """Whether a complete model output is in this format.""" + + @classmethod + @abstractmethod + def parse(cls, text: str, tools: list | None) -> tuple[str, list[ToolCall]]: + """Parse a complete output into ``(leading_content, tool_calls)``.""" + + # -- streaming ---------------------------------------------------------- + @abstractmethod + def process(self, text: str) -> list: + """Consume one chunk; return ``(event_type, data)`` tuples.""" + + @abstractmethod + def flush(self) -> list: + """Drain whatever is buffered at end of stream.""" + + def _emit_call(self, tc: ToolCall) -> list: + """Render one parsed ToolCall as start+args stream events.""" + events = [ + ( + "tool_call_start", + { + "index": self.current_index, + "id": tc.id, + "type": "function", + "function": {"name": tc.function["name"], "arguments": ""}, + }, + ), + ( + "tool_call_args", + { + "index": self.current_index, + "function": {"arguments": tc.function["arguments"]}, + }, + ), + ] + self.current_index += 1 + self.emitted_calls += 1 + return events + + +class BufferedMarkerParser(ToolCallParser): + """Formats that buffer from a start marker and parse the block at flush. + + The block is only parsed once complete because partial XML streams badly + (a half-written `` int: + """Index of the earliest start marker, or -1.""" + positions = [i for i in (text.find(m) for m in cls.START_MARKERS) if i != -1] + return min(positions) if positions else -1 + + @classmethod + def detect(cls, text: str) -> bool: + return cls.find_start(text) != -1 + + def process(self, text: str) -> list: + results: list = [] + self.buf += text + if self.state == 0: + m = self.find_start(self.buf) + if m != -1: + before = self.buf[:m] + if before: + results.append(("content", before)) + self.buf = self.buf[m:] + self.state = 1 + else: + # Emit content but hold back a possible partial marker tail. + cut = max(self.buf.rfind(c) for c in self.HOLDBACK_CHARS) + if cut == -1: + if self.buf: + results.append(("content", self.buf)) + self.buf = "" + elif cut > 0: + results.append(("content", self.buf[:cut])) + self.buf = self.buf[cut:] + return results + + def flush(self) -> list: + results: list = [] + if self.state == 0: + if self.buf: + results.append(("content", self.buf)) + self.buf = "" + return results + # state 1: parse the complete (or trailing) tool-call block. + _content, tool_calls = self.parse(self.buf, self.tools) + self.buf = "" + for tc in tool_calls: + results.extend(self._emit_call(tc)) + if self.emitted_calls > 0: + results.append(("tool_call_end", None)) + return results diff --git a/atom/model_engine/engine_core_mgr.py b/atom/model_engine/engine_core_mgr.py index 5e058bc451..b89dcd418b 100644 --- a/atom/model_engine/engine_core_mgr.py +++ b/atom/model_engine/engine_core_mgr.py @@ -285,6 +285,20 @@ def process_outputs_socket(): logger.debug( f"{self.label}: Cleaned up callback for finished sequence {seq_id}" ) + # Batched stream dispatch: the per-seq callbacks only buffer + # their chunks into a thread-local; flush the whole step's + # buffer into the per-request asyncio queues now (one + # call_soon_threadsafe per loop). Resolved lazily by the API + # server to avoid the api_server <-> engine_core_mgr import + # cycle. No-op when no streaming request is in flight. + if self._flush_stream_batch_fn is not None: + try: + self._flush_stream_batch_fn() + except Exception as e: + logger.warning( + f"{self.label}: flush_stream_batch failed: {e}", + exc_info=True, + ) elif request_type == EngineCoreRequestType.UTILITY_RESPONSE: self.utility_response_queue.put_nowait(data) elif request_type == EngineCoreRequestType.ADD: diff --git a/tests/entrypoints/test_chat_encoders.py b/tests/entrypoints/test_chat_encoders.py new file mode 100644 index 0000000000..b25d3c873e --- /dev/null +++ b/tests/entrypoints/test_chat_encoders.py @@ -0,0 +1,134 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""Tests for model-scoped custom chat encoder dispatch.""" + +from atom.entrypoints.openai.chat_encoder_adapters import ( + build_message_encoder_adapter, +) +from atom.entrypoints.openai.chat_encoders import ( + _load_encoder_from_dir, + apply_chat_template, +) + + +def test_loader_selects_dsv4_adapter_and_preserves_encoder_defaults(tmp_path): + encoding_dir = tmp_path / "encoding" + encoding_dir.mkdir() + (encoding_dir / "encoding_dsv4.py").write_text( + "def encode_messages(messages, **kwargs):\n" + " return repr((messages, kwargs))\n", + encoding="utf-8", + ) + + adapter = _load_encoder_from_dir(str(tmp_path)) + + assert adapter is not None + assert adapter.name == "encoding_dsv4" + assert adapter.supports_tools is True + rendered = apply_chat_template( + tokenizer=None, + custom_encoder=adapter, + messages=[{"role": "user", "content": "hello"}], + ) + assert "'thinking_mode': 'thinking'" in rendered + + +def test_dsv4_adapter_prepends_tools_without_reordering_messages(): + captured = {} + + def raw_encoder(messages, **kwargs): + captured["messages"] = messages + captured["kwargs"] = kwargs + return "rendered" + + adapter = build_message_encoder_adapter("encoding_dsv4", raw_encoder) + messages = [ + {"role": "system", "content": "policy"}, + {"role": "user", "content": "question"}, + {"role": "system", "content": "trailing context"}, + ] + original = [dict(message) for message in messages] + tools = [{"type": "function", "function": {"name": "search"}}] + + result = apply_chat_template( + tokenizer=None, + custom_encoder=adapter, + messages=messages, + tools=tools, + tokenize=True, + add_generation_prompt=True, + thinking_mode="chat", + ) + + assert result == "rendered" + assert captured["messages"] == [ + {"role": "system", "tools": tools}, + *original, + ] + assert captured["kwargs"] == {"thinking_mode": "chat"} + assert messages == original + assert captured["messages"][1:] is not messages + assert all( + prepared is not source + for prepared, source in zip(captured["messages"][1:], messages) + ) + + +def test_unknown_custom_encoder_does_not_receive_dsv4_fields(caplog): + captured = {} + + def raw_encoder(messages, **kwargs): + captured["messages"] = messages + return "rendered" + + adapter = build_message_encoder_adapter("encoding_other", raw_encoder) + messages = [{"role": "user", "content": "hello"}] + tools = [{"type": "function", "function": {"name": "search"}}] + + result = apply_chat_template( + tokenizer=None, + custom_encoder=adapter, + messages=messages, + tools=tools, + ) + + assert result == "rendered" + assert captured["messages"] == messages + assert captured["messages"] is not messages + assert captured["messages"][0] is not messages[0] + assert "tools" not in captured["messages"][0] + assert "tools= is not supported" in caplog.text + + +def test_jinja_path_forwards_tools_and_generation_kwargs(): + class Tokenizer: + def __init__(self): + self.messages = None + self.kwargs = None + + def apply_chat_template(self, messages, **kwargs): + self.messages = messages + self.kwargs = kwargs + return "jinja-rendered" + + tokenizer = Tokenizer() + messages = [{"role": "user", "content": "hello"}] + tools = [{"type": "function", "function": {"name": "search"}}] + + result = apply_chat_template( + tokenizer=tokenizer, + custom_encoder=None, + messages=messages, + tools=tools, + enable_thinking=True, + ) + + assert result == "jinja-rendered" + assert tokenizer.messages is messages + assert tokenizer.kwargs == { + "enable_thinking": True, + "tokenize": False, + "add_generation_prompt": True, + "tools": tools, + } diff --git a/tests/entrypoints/test_streaming_dispatch.py b/tests/entrypoints/test_streaming_dispatch.py new file mode 100644 index 0000000000..9c561ee300 --- /dev/null +++ b/tests/entrypoints/test_streaming_dispatch.py @@ -0,0 +1,122 @@ +import asyncio + +from atom.entrypoints.openai.streaming_dispatch import ( + IncrementalStreamDetokenizer, + StreamBatchDispatcher, +) + + +class _Utf8ByteTokenizer: + def decode(self, token_ids, skip_special_tokens=True): + return bytes(token_ids).decode("utf-8", errors="replace") + + +class _ImmediateLoop: + def __init__(self): + self.calls = [] + + def call_soon_threadsafe(self, callback, *args): + self.calls.append((callback, args)) + callback(*args) + + +def test_incremental_detokenizer_holds_incomplete_utf8(): + detokenizer = IncrementalStreamDetokenizer(_Utf8ByteTokenizer()) + + assert detokenizer.update([0xE4], finished=False) == "" + assert detokenizer.update([0xBD, 0xA0], finished=False) == "你" + assert detokenizer.update([ord("!")], finished=True) == "!" + + +def test_dispatcher_batches_direct_and_tagged_chunks_per_loop(): + dispatcher = StreamBatchDispatcher(_Utf8ByteTokenizer()) + loop = _ImmediateLoop() + direct_queue = asyncio.Queue() + tagged_queue = asyncio.Queue() + + dispatcher.enqueue( + loop=loop, + queue=direct_queue, + state_key="request-1", + chunk={"token_ids": [ord("A")], "finished": True}, + ) + dispatcher.enqueue( + loop=loop, + queue=tagged_queue, + state_key=("request-2", 0), + chunk={"token_ids": [ord("B")], "finished": True}, + tag=0, + ) + dispatcher.flush() + + assert len(loop.calls) == 1 + assert direct_queue.get_nowait()["text"] == "A" + sibling_index, chunk = tagged_queue.get_nowait() + assert sibling_index == 0 + assert chunk["text"] == "B" + + +def test_dispatcher_keeps_fanout_detokenizer_state_separate(): + dispatcher = StreamBatchDispatcher(_Utf8ByteTokenizer()) + loop = _ImmediateLoop() + queue = asyncio.Queue() + + dispatcher.enqueue( + loop=loop, + queue=queue, + state_key=("request", 0), + chunk={"token_ids": [0xE4], "finished": False}, + tag=0, + ) + dispatcher.enqueue( + loop=loop, + queue=queue, + state_key=("request", 1), + chunk={"token_ids": [ord("X")], "finished": True}, + tag=1, + ) + dispatcher.flush() + + assert queue.get_nowait()[1]["text"] == "" + assert queue.get_nowait()[1]["text"] == "X" + + dispatcher.enqueue( + loop=loop, + queue=queue, + state_key=("request", 0), + chunk={"token_ids": [0xBD, 0xA0], "finished": True}, + tag=0, + ) + dispatcher.flush() + + assert queue.get_nowait()[1]["text"] == "你" + + +def test_discard_request_drops_partial_direct_and_fanout_state(): + dispatcher = StreamBatchDispatcher(_Utf8ByteTokenizer()) + loop = _ImmediateLoop() + queue = asyncio.Queue() + + for state_key in ("request", ("request", 0)): + dispatcher.enqueue( + loop=loop, + queue=queue, + state_key=state_key, + chunk={"token_ids": [0xE4], "finished": False}, + ) + dispatcher.flush() + dispatcher.discard_request("request") + + for state_key in ("request", ("request", 0)): + dispatcher.enqueue( + loop=loop, + queue=queue, + state_key=state_key, + chunk={"token_ids": [ord("A")], "finished": True}, + ) + dispatcher.flush() + + assert queue.get_nowait()["text"] == "" + assert queue.get_nowait()["text"] == "" + assert queue.get_nowait()["text"] == "A" + assert queue.get_nowait()["text"] == "A"