diff --git a/python/minisgl/message/frontend.py b/python/minisgl/message/frontend.py index d664ccf6..068c64e0 100644 --- a/python/minisgl/message/frontend.py +++ b/python/minisgl/message/frontend.py @@ -27,3 +27,4 @@ class UserReply(BaseFrontendMsg): uid: int incremental_output: str finished: bool + reasoning_output: str = "" diff --git a/python/minisgl/parser/__init__.py b/python/minisgl/parser/__init__.py new file mode 100644 index 00000000..9c6ebb05 --- /dev/null +++ b/python/minisgl/parser/__init__.py @@ -0,0 +1,12 @@ +""" +minisgl.parser +============== + +Top-level package for output-parsing utilities. + +Sub-packages +------------ +:mod:`minisgl.parser.reasoning` + Identify and split reasoning blocks (e.g. ``…``) from + model-generated text, in both streaming and non-streaming modes. +""" diff --git a/python/minisgl/parser/reasoning/__init__.py b/python/minisgl/parser/reasoning/__init__.py new file mode 100644 index 00000000..332f1762 --- /dev/null +++ b/python/minisgl/parser/reasoning/__init__.py @@ -0,0 +1,59 @@ +""" +minisgl.parser.reasoning +======================== + +Utilities for identifying and splitting reasoning blocks from model-generated +text. Supports both **non-streaming** (full-text) and **streaming** +(chunk-by-chunk) parsing. + +Quick start +----------- +:: + + from minisgl.parser.reasoning import ReasoningParser + + # Non-streaming + parser = ReasoningParser("qwen3") + result = parser.parse_full(full_text) + print(result.reasoning_text) # chain-of-thought + print(result.normal_text) # answer + + # Streaming + parser = ReasoningParser("qwen3") + for raw_chunk in token_stream: + sc = parser.parse_stream(raw_chunk) + ... + final = parser.flush() + +Public API +---------- +:class:`ReasoningParser` + Main entry point. Selects a detector by *model_type* and exposes + :meth:`~ReasoningParser.parse_full`, :meth:`~ReasoningParser.parse_stream`, + :meth:`~ReasoningParser.flush`, and :meth:`~ReasoningParser.stream_iter`. + +:class:`ParseResult` + Non-streaming result with ``reasoning_text`` and ``normal_text``. + +:class:`StreamChunk` + Streaming result with ``reasoning_delta`` and ``normal_delta``. + +:class:`BaseDetector` + ABC for custom detectors (subclass to add new reasoning formats). + +:class:`ThinkTagDetector` + Concrete detector for the ``…`` format used by + Qwen3-Thinking. +""" + +from .base import BaseDetector, ParseResult, StreamChunk +from .parser import ReasoningParser +from .think_tag import ThinkTagDetector + +__all__ = [ + "ReasoningParser", + "ParseResult", + "StreamChunk", + "BaseDetector", + "ThinkTagDetector", +] diff --git a/python/minisgl/parser/reasoning/base.py b/python/minisgl/parser/reasoning/base.py new file mode 100644 index 00000000..9704b8ae --- /dev/null +++ b/python/minisgl/parser/reasoning/base.py @@ -0,0 +1,81 @@ +from __future__ import annotations + +from abc import ABC, abstractmethod +from dataclasses import dataclass + + +@dataclass +class ParseResult: + """Non-streaming parse output.""" + + reasoning_text: str + """Content extracted from inside the reasoning block.""" + + normal_text: str + """Remaining content outside the reasoning block.""" + + +@dataclass +class StreamChunk: + """Incremental output produced by a single :meth:`ReasoningParser.parse_stream` call.""" + + reasoning_delta: str + """New reasoning content decoded in this chunk (may be empty).""" + + normal_delta: str + """New normal content decoded in this chunk (may be empty).""" + + +class BaseDetector(ABC): + """ + Abstract base for reasoning-block detectors. + + A detector knows the start/end markers that delimit a reasoning block + and can parse both complete text (non-streaming) and incremental chunks + (streaming) through the helpers in :class:`ReasoningParser`. + + Concrete subclasses must implement :attr:`start_tag` and :attr:`end_tag`. + They may also override :meth:`parse_full` for model-specific logic. + """ + + @property + @abstractmethod + def start_tag(self) -> str: + """Opening delimiter of the reasoning block (e.g. ``""```).""" + + @property + @abstractmethod + def end_tag(self) -> str: + """Closing delimiter of the reasoning block (e.g. ``""```).""" + + def parse_full(self, text: str) -> ParseResult: + """ + Extract reasoning and normal text from a complete, fully-decoded + generation string (non-streaming path). + + The default implementation does a single forward scan for the + first occurrence of :attr:`start_tag` / :attr:`end_tag`. Subclasses + may override for more sophisticated (e.g. multi-block) handling. + """ + start = self.start_tag + end = self.end_tag + + start_idx = text.find(start) + if start_idx == -1: + return ParseResult(reasoning_text="", normal_text=text) + + content_start = start_idx + len(start) + end_idx = text.find(end, content_start) + + if end_idx == -1: + # Reasoning block was never closed – treat everything after the + # start tag as reasoning content and keep the preceding text as + # normal output. + return ParseResult( + reasoning_text=text[content_start:], + normal_text=text[:start_idx], + ) + + reasoning = text[content_start:end_idx] + normal = text[:start_idx] + text[end_idx + len(end):] + return ParseResult(reasoning_text=reasoning, normal_text=normal) diff --git a/python/minisgl/parser/reasoning/parser.py b/python/minisgl/parser/reasoning/parser.py new file mode 100644 index 00000000..d9856354 --- /dev/null +++ b/python/minisgl/parser/reasoning/parser.py @@ -0,0 +1,289 @@ +from __future__ import annotations + +from typing import Iterator + +from .base import BaseDetector, ParseResult, StreamChunk +from .think_tag import ThinkTagDetector + +# --------------------------------------------------------------------------- +# Model-type → detector mapping +# --------------------------------------------------------------------------- + +# Model types (lowercased, hyphens replaced by underscores) that are known +# to use the … format. +_THINK_TAG_MODELS: frozenset[str] = frozenset( + { + "qwen3", + "qwen3_moe", + } +) + +# Prefixes that are also mapped to ThinkTagDetector even when not in the +# exact-match set above (handles future model variants automatically). +_THINK_TAG_PREFIXES: tuple[str, ...] = ("qwen3",) + + +# --------------------------------------------------------------------------- +# Streaming state machine +# --------------------------------------------------------------------------- + + +def _safe_prefix_len(buf: str, tag: str) -> int: + """ + Number of leading characters in *buf* that are guaranteed **not** to be + the beginning of an upcoming *tag*. + + We keep at most ``len(tag) - 1`` trailing characters buffered so that a + tag that is split across two successive chunks is never emitted early. + + Examples:: + + _safe_prefix_len("hello ") → 4 # "hell" is safe + _safe_prefix_len("hello", "") → 0 # all could match + _safe_prefix_len("hello world", "") → 5 + """ + return max(0, len(buf) - len(tag) + 1) + + +class _StreamingState: + """ + Internal stateful helper that processes incremental text chunks and + produces per-chunk ``(reasoning_delta, normal_delta)`` pairs. + + State machine + ------------- + ``before`` + The reasoning block has not started yet. Incoming text is treated as + normal output until the start tag is found. + + ``in_reasoning`` + We are inside the reasoning block. Incoming text is treated as + reasoning content until the end tag is found. + + ``after`` + The reasoning block has closed. All remaining text is normal output. + + Partial-tag handling + -------------------- + Because chunks can arrive mid-tag (e.g. ``""``), we + buffer the last ``len(tag) - 1`` characters in each state and only emit + characters that are definitively before any possible tag boundary. + """ + + def __init__(self, detector: BaseDetector) -> None: + self._detector = detector + self._state: str = "before" + self._buf: str = "" + + def feed(self, chunk: str) -> StreamChunk: + """ + Consume one text chunk and return the incremental reasoning / normal + content decoded so far. + + This method is **not** re-entrant; call it sequentially for each + streamed chunk and call :meth:`flush` once after the final chunk. + """ + reasoning_delta = "" + normal_delta = "" + + self._buf += chunk + + # ------------------------------------------------------------------ # + # Phase 1 – scan for the reasoning block start tag # + # ------------------------------------------------------------------ # + if self._state == "before": + start_tag = self._detector.start_tag + idx = self._buf.find(start_tag) + if idx != -1: + # Everything before the tag is normal text. + normal_delta += self._buf[:idx] + # Discard the tag itself and advance. + self._buf = self._buf[idx + len(start_tag):] + self._state = "in_reasoning" + else: + # Emit the safe prefix; keep a potential partial-tag suffix. + safe = _safe_prefix_len(self._buf, start_tag) + normal_delta += self._buf[:safe] + self._buf = self._buf[safe:] + + # ------------------------------------------------------------------ # + # Phase 2 – scan for the reasoning block end tag # + # (runs in the same call if the start tag was just found above) # + # ------------------------------------------------------------------ # + if self._state == "in_reasoning": + end_tag = self._detector.end_tag + idx = self._buf.find(end_tag) + if idx != -1: + reasoning_delta += self._buf[:idx] + self._buf = self._buf[idx + len(end_tag):] + self._state = "after" + else: + safe = _safe_prefix_len(self._buf, end_tag) + reasoning_delta += self._buf[:safe] + self._buf = self._buf[safe:] + + # ------------------------------------------------------------------ # + # Phase 3 – emit everything remaining as normal text # + # ------------------------------------------------------------------ # + if self._state == "after": + normal_delta += self._buf + self._buf = "" + + return StreamChunk(reasoning_delta=reasoning_delta, normal_delta=normal_delta) + + def flush(self) -> StreamChunk: + """ + Signal end-of-stream. + + Drains any text that was kept in the internal buffer to guard against + partial tags. Should be called **once** after the final chunk so that + callers never miss trailing content. + """ + remaining = self._buf + self._buf = "" + if self._state == "in_reasoning": + # Unclosed reasoning block – treat remainder as reasoning. + return StreamChunk(reasoning_delta=remaining, normal_delta="") + return StreamChunk(reasoning_delta="", normal_delta=remaining) + + +# --------------------------------------------------------------------------- +# Public API +# --------------------------------------------------------------------------- + + +class ReasoningParser: + """ + High-level reasoning-content parser. + + Selects the appropriate detector based on *model_type* and exposes a + unified interface for both **non-streaming** (single-call) and + **streaming** (chunk-by-chunk) use. + + Parameters + ---------- + model_type: + The ``model_type`` string from the HuggingFace ``config.json`` + (e.g. ``"qwen3"``). Pass ``None`` or ``"auto"`` + to use :class:`ThinkTagDetector` as a sensible default. + + Examples + -------- + Non-streaming:: + + parser = ReasoningParser("qwen3") + result = parser.parse_full(full_generated_text) + print(result.reasoning_text) # chain-of-thought + print(result.normal_text) # answer + + Streaming:: + + parser = ReasoningParser("qwen3") + for raw_chunk in token_stream: + chunk = parser.parse_stream(raw_chunk) + if chunk.reasoning_delta: + handle_reasoning(chunk.reasoning_delta) + if chunk.normal_delta: + handle_normal(chunk.normal_delta) + # drain remaining buffer + final = parser.flush() + if final.normal_delta: + handle_normal(final.normal_delta) + """ + + def __init__(self, model_type: str | None = None) -> None: + self._detector: BaseDetector = self._make_detector(model_type) + self._streaming_state = _StreamingState(self._detector) + + # ---------------------------------------------------------------------- # + # Non-streaming # + # ---------------------------------------------------------------------- # + + def parse_full(self, text: str) -> ParseResult: + """ + Parse a **complete** generation string. + + Returns a :class:`ParseResult` with ``reasoning_text`` (content + inside the reasoning block) and ``normal_text`` (everything else). + """ + return self._detector.parse_full(text) + + # ---------------------------------------------------------------------- # + # Streaming # + # ---------------------------------------------------------------------- # + + def parse_stream(self, chunk: str) -> StreamChunk: + """ + Feed one incremental text chunk from a streaming decode loop. + + Returns a :class:`StreamChunk` whose ``reasoning_delta`` and + ``normal_delta`` fields hold the content that can be safely attributed + to reasoning and normal output respectively for this chunk. + + Call :meth:`flush` after the **last** chunk to drain the internal + partial-tag buffer. + """ + return self._streaming_state.feed(chunk) + + def flush(self) -> StreamChunk: + """ + Signal end-of-stream and drain any buffered content. + + Must be called once after the final chunk to ensure no trailing text + is silently dropped. + """ + return self._streaming_state.flush() + + def stream_iter(self, chunks: Iterator[str]) -> Iterator[StreamChunk]: + """ + Convenience wrapper: iterate over *chunks* and yield one + :class:`StreamChunk` per input chunk, followed by the flush chunk. + + Usage:: + + for sc in parser.stream_iter(token_stream): + ... + """ + for chunk in chunks: + yield self.parse_stream(chunk) + final = self.flush() + if final.reasoning_delta or final.normal_delta: + yield final + + # ---------------------------------------------------------------------- # + # Internals # + # ---------------------------------------------------------------------- # + + @property + def detector(self) -> BaseDetector: + """The underlying :class:`BaseDetector` instance (read-only).""" + return self._detector + + @staticmethod + def _make_detector(model_type: str | None) -> BaseDetector: + """ + Resolve the detector for the given *model_type*. + + Resolution order + ---------------- + 1. ``None`` / ``"auto"`` → :class:`ThinkTagDetector` (safe default) + 2. Exact match in ``_THINK_TAG_MODELS`` set + 3. Prefix match against ``_THINK_TAG_PREFIXES`` + 4. Unrecognised model type → :class:`ThinkTagDetector` (safe fallback) + The caller is responsible for knowing whether the model actually + emits reasoning tags; an unknown type simply won't extract any + reasoning if no ```` tag is present. + """ + if model_type is None or model_type == "auto": + return ThinkTagDetector() + + normalized = model_type.lower().replace("-", "_") + + if normalized in _THINK_TAG_MODELS: + return ThinkTagDetector() + + if any(normalized.startswith(p) for p in _THINK_TAG_PREFIXES): + return ThinkTagDetector() + + # Unknown model type – fall back gracefully. + return ThinkTagDetector() diff --git a/python/minisgl/parser/reasoning/think_tag.py b/python/minisgl/parser/reasoning/think_tag.py new file mode 100644 index 00000000..1648dc55 --- /dev/null +++ b/python/minisgl/parser/reasoning/think_tag.py @@ -0,0 +1,20 @@ +from __future__ import annotations + +from .base import BaseDetector + + +class ThinkTagDetector(BaseDetector): + """ + Detector for the ``…`` reasoning format. + + Used by DeepSeek-R1, Qwen3-Thinking, QwQ, and other models that wrap + chain-of-thought content in ```` / ```` XML-like tags. + """ + + @property + def start_tag(self) -> str: + return "" + + @property + def end_tag(self) -> str: + return "" diff --git a/python/minisgl/server/api_server.py b/python/minisgl/server/api_server.py index 9318ccdd..0068afcd 100644 --- a/python/minisgl/server/api_server.py +++ b/python/minisgl/server/api_server.py @@ -167,6 +167,8 @@ async def stream_chat_completions(self, uid: int): first_chunk = False if ack.incremental_output: delta["content"] = ack.incremental_output + if ack.reasoning_output: + delta["reasoning_content"] = ack.reasoning_output chunk = { "id": f"cmpl-{uid}", diff --git a/python/minisgl/server/args.py b/python/minisgl/server/args.py index 3ec88f8d..77ec7379 100644 --- a/python/minisgl/server/args.py +++ b/python/minisgl/server/args.py @@ -17,6 +17,9 @@ class ServerArgs(SchedulerConfig): server_port: int = 1919 num_tokenizer: int = 0 silent_output: bool = False + # None → feature disabled + # str → model_type used to select detector (e.g. "qwen3", "deepseek_r1") + reasoning_parser: str | None = None @property def share_tokenizer(self) -> bool: @@ -223,6 +226,20 @@ def parse_args(args: List[str], run_shell: bool = False) -> Tuple[ServerArgs, bo help="Run the server in shell mode.", ) + parser.add_argument( + "--reasoning-parser", + dest="reasoning_parser", + nargs="?", + const="auto", # --reasoning-parser with no value → "auto" + default=None, # not given at all → None (disabled) + metavar="MODEL_TYPE", + help=( + "Enable reasoning-block parsing. " + "Optionally specify the model type (e.g. 'qwen3', 'deepseek_r1'). " + "Omit the value to auto-detect from the model config." + ), + ) + # Parse arguments kwargs = parser.parse_args(args).__dict__.copy() @@ -262,6 +279,13 @@ def parse_args(args: List[str], run_shell: bool = False) -> Tuple[ServerArgs, bo kwargs["tp_info"] = DistributedInfo(0, kwargs["tensor_parallel_size"]) del kwargs["tensor_parallel_size"] + # Resolve --reasoning-parser auto → actual model_type from HF config. + if kwargs.get("reasoning_parser") == "auto": + from minisgl.utils import cached_load_hf_config + + hf_cfg = cached_load_hf_config(kwargs["model_path"]) + kwargs["reasoning_parser"] = getattr(hf_cfg, "model_type", "auto") + result = ServerArgs(**kwargs) logger = init_logger(__name__) logger.info(f"Parsed arguments:\n{result}") diff --git a/python/minisgl/server/launch.py b/python/minisgl/server/launch.py index 2e9c2f10..11a9ac6f 100644 --- a/python/minisgl/server/launch.py +++ b/python/minisgl/server/launch.py @@ -81,6 +81,7 @@ def start_subprocess() -> None: "create": server_args.tokenizer_create_addr, "tokenizer_id": num_tokenizers, "ack_queue": ack_queue, + "reasoning_parser": server_args.reasoning_parser, }, daemon=False, name="minisgl-detokenizer-0", @@ -97,6 +98,7 @@ def start_subprocess() -> None: "create": server_args.tokenizer_create_addr, "tokenizer_id": i, "ack_queue": ack_queue, + "reasoning_parser": server_args.reasoning_parser, }, daemon=False, name=f"minisgl-tokenizer-{i}", diff --git a/python/minisgl/tokenizer/detokenize.py b/python/minisgl/tokenizer/detokenize.py index 9a7c006c..e1e412c6 100644 --- a/python/minisgl/tokenizer/detokenize.py +++ b/python/minisgl/tokenizer/detokenize.py @@ -1,7 +1,8 @@ from dataclasses import dataclass -from typing import Dict, List +from typing import Dict, List, Optional, Tuple from minisgl.message import DetokenizeMsg +from minisgl.parser.reasoning import ReasoningParser from transformers import PreTrainedTokenizerBase # Borrowed from sglang @@ -61,13 +62,21 @@ class DecodeStatus: class DetokenizeManager: - def __init__(self, tokenizer: PreTrainedTokenizerBase) -> None: + def __init__( + self, + tokenizer: PreTrainedTokenizerBase, + reasoning_parser_model_type: Optional[str] = None, + ) -> None: # uid -> DecodeStatus self.decode_map: Dict[int, DecodeStatus] = {} self.tokenizer = tokenizer self.eos_token_id = self.tokenizer.eos_token_id - def detokenize(self, msgs: List[DetokenizeMsg]) -> List[str]: + self._reasoning_model_type = reasoning_parser_model_type + # uid -> ReasoningParser (created lazily on first token for that uid) + self._reasoning_map: Dict[int, ReasoningParser] = {} + + def detokenize(self, msgs: List[DetokenizeMsg]) -> List[Tuple[str, str]]: read_ids: List[List[int]] = [] surr_ids: List[List[int]] = [] for msg in msgs: @@ -88,7 +97,8 @@ def detokenize(self, msgs: List[DetokenizeMsg]) -> List[str]: read_texts = self.tokenizer.batch_decode(read_ids) surr_texts = self.tokenizer.batch_decode(surr_ids) - incremental_strs: List[str] = [] + # List of (normal_delta, reasoning_delta) tuples + results: List[Tuple[str, str]] = [] for msg, read_str, surr_str in zip(msgs, read_texts, surr_texts, strict=True): s = self.decode_map[msg.uid] new_text = read_str[len(surr_str) :] @@ -104,8 +114,31 @@ def detokenize(self, msgs: List[DetokenizeMsg]) -> List[str]: incremental_output = output_str[s.sent_offset :] s.sent_offset = len(output_str) - incremental_strs.append(incremental_output) + + # ── Reasoning-parser split ──────────────────────────────────── # + normal_delta = incremental_output + reasoning_delta = "" + if self._reasoning_model_type is not None and incremental_output: + if msg.uid not in self._reasoning_map: + self._reasoning_map[msg.uid] = ReasoningParser( + self._reasoning_model_type + ) + rp = self._reasoning_map[msg.uid] + if msg.finished: + # Feed the last chunk then flush to drain the buffer. + sc = rp.parse_stream(incremental_output) + fl = rp.flush() + normal_delta = sc.normal_delta + fl.normal_delta + reasoning_delta = sc.reasoning_delta + fl.reasoning_delta + else: + sc = rp.parse_stream(incremental_output) + normal_delta = sc.normal_delta + reasoning_delta = sc.reasoning_delta + + results.append((normal_delta, reasoning_delta)) + if msg.finished: del self.decode_map[msg.uid] + self._reasoning_map.pop(msg.uid, None) - return incremental_strs + return results diff --git a/python/minisgl/tokenizer/server.py b/python/minisgl/tokenizer/server.py index c2d0e572..ccb77c8e 100644 --- a/python/minisgl/tokenizer/server.py +++ b/python/minisgl/tokenizer/server.py @@ -39,6 +39,7 @@ def tokenize_worker( tokenizer_id: int = -1, model_source: str = "huggingface", ack_queue: mp.Queue[str] | None = None, + reasoning_parser: str | None = None, ) -> None: send_backend = ZmqPushQueue(backend_addr, create=False, encoder=BaseBackendMsg.encoder) send_frontend = ZmqPushQueue(frontend_addr, create=False, encoder=BaseFrontendMsg.encoder) @@ -51,7 +52,9 @@ def tokenize_worker( from .tokenize import TokenizeManager tokenize_manager = TokenizeManager(tokenizer) - detokenize_manager = DetokenizeManager(tokenizer) + detokenize_manager = DetokenizeManager( + tokenizer, reasoning_parser_model_type=reasoning_parser + ) if ack_queue is not None: ack_queue.put(f"Tokenize server {tokenizer_id} is ready") @@ -74,10 +77,13 @@ def tokenize_worker( data=[ UserReply( uid=msg.uid, - incremental_output=reply, + incremental_output=normal_delta, finished=msg.finished, + reasoning_output=reasoning_delta, + ) + for msg, (normal_delta, reasoning_delta) in zip( + detokenize_msg, replies, strict=True ) - for msg, reply in zip(detokenize_msg, replies, strict=True) ] ) if len(batch_output.data) == 1: diff --git a/tests/misc/test_reasoning_parser.py b/tests/misc/test_reasoning_parser.py new file mode 100644 index 00000000..9a69c30e --- /dev/null +++ b/tests/misc/test_reasoning_parser.py @@ -0,0 +1,238 @@ +#!/usr/bin/env python3 +""" +test_reasoning_parser.py +Send several requests to a running mini-sglang engine (with --reasoning-parser enabled), +to verify whether reasoning_content / content are split correctly. + +Usage: + python test_reasoning_parser.py [--host 127.0.0.1] [--port 1919] + +Dependency: requests (install via: pip install requests) +""" +from __future__ import annotations + +import argparse +import json +import sys +import textwrap +from typing import Iterator + +import requests + +RESET = "\033[0m" +BOLD = "\033[1m" +CYAN = "\033[36m" +YELLOW = "\033[33m" +GREEN = "\033[32m" +RED = "\033[31m" +DIM = "\033[2m" + +def _c(text: str, *codes: str) -> str: + if sys.stdout.isatty(): + return "".join(codes) + text + RESET + return text + +def _iter_sse_chunks(resp: requests.Response) -> Iterator[dict]: + """Parse text/event-stream line by line and yield JSON chunks.""" + for raw_line in resp.iter_lines(): + if isinstance(raw_line, bytes): + raw_line = raw_line.decode() + line = raw_line.strip() + if not line or not line.startswith("data:"): + continue + payload = line[5:].strip() + if payload == "[DONE]": + break + try: + yield json.loads(payload) + except json.JSONDecodeError: + continue + +# ────────────────────────────────────────────────────────────────────────────── +# Single Request +# ────────────────────────────────────────────────────────────────────────────── +def chat_stream( + base_url: str, + messages: list[dict], + model: str, + max_tokens: int = 1024, + temperature: float = 0.6, +) -> tuple[str, str]: + """ + Send a streaming /v1/chat/completions request. + Returns (reasoning_text, normal_text) as complete strings. + """ + url = f"{base_url}/v1/chat/completions" + payload = { + "model": model, + "messages": messages, + "max_tokens": max_tokens, + "temperature": temperature, + "stream": True, + } + + reasoning_buf = "" + normal_buf = "" + + with requests.post(url, json=payload, stream=True, timeout=120) as resp: + resp.raise_for_status() + for chunk in _iter_sse_chunks(resp): + choices = chunk.get("choices", []) + if not choices: + continue + delta = choices[0].get("delta", {}) + reasoning_buf += delta.get("reasoning_content") or "" + normal_buf += delta.get("content") or "" + + return reasoning_buf, normal_buf + +_W = 72 + +def _hr(char: str = "─") -> str: + return _c(char * _W, DIM) + +def _section(title: str) -> None: + print() + print(_c(f"{'─' * 3} {title} {'─' * (_W - 5 - len(title))}", BOLD + CYAN)) + +def _print_result(reasoning: str, normal: str) -> None: + if reasoning: + print(_c(" [reasoning_content]", YELLOW)) + wrapped = textwrap.fill(reasoning.strip(), width=_W - 4, + initial_indent=" ", subsequent_indent=" ") + print(_c(wrapped, YELLOW)) + else: + print(_c(" [reasoning_content] (Empty — model might not have output tags)", DIM)) + + print() + print(_c(" [content]", GREEN)) + wrapped = textwrap.fill(normal.strip() or "(Empty)", width=_W - 4, + initial_indent=" ", subsequent_indent=" ") + print(_c(wrapped, GREEN)) + +def _check(label: str, cond: bool) -> None: + icon = _c("✓", GREEN + BOLD) if cond else _c("✗", RED + BOLD) + print(f" {icon} {label}") + +CASES: list[dict] = [ + { + "name": "Simple Math Reasoning", + "messages": [ + {"role": "user", "content": "Which is larger, 9.11 or 9.9? Please think carefully before answering."}, + ], + "max_tokens": 4096, + }, + { + "name": "Logic Puzzle", + "messages": [ + { + "role": "user", + "content": ( + "There are 3 light bulbs in a room and 3 switches outside. " + "Each switch controls one bulb. You can only enter the room once. " + "How do you determine which switch controls which bulb?" + ), + } + ], + "max_tokens": 4096, + }, + { + "name": "Code Generation", + "messages": [ + { + "role": "user", + "content": "Write a concise Python function to check if a number is prime.", + } + ], + "max_tokens": 4096, + }, +] + +# ────────────────────────────────────────────────────────────────────────────── +# Main Flow +# ────────────────────────────────────────────────────────────────────────────── +def main() -> None: + parser = argparse.ArgumentParser(description="Test mini-sglang reasoning parser") + parser.add_argument("--host", default="127.0.0.1") + parser.add_argument("--port", default=1919, type=int) + parser.add_argument("--max-tokens", default=None, type=int, + help="Override max_tokens for all cases") + parser.add_argument("--temperature", default=0.6, type=float) + args = parser.parse_args() + + base_url = f"http://{args.host}:{args.port}" + + print(_c(f"\nmini-sglang reasoning-parser test script", BOLD)) + print(_c(f"Target: {base_url}", DIM)) + print(_hr()) + + try: + r = requests.get(f"{base_url}/v1/models", timeout=5) + r.raise_for_status() + model_name = r.json()["data"][0]["id"] + print(f" Service Online model = {_c(model_name, BOLD)}") + except Exception as exc: + print(_c(f" ✗ Cannot connect to server: {exc}", RED)) + sys.exit(1) + + passed = 0 + failed = 0 + + for i, case in enumerate(CASES, 1): + max_tokens = args.max_tokens or case.get("max_tokens", 512) + _section(f"[{i}/{len(CASES)}] {case['name']}") + + user_msg = case["messages"][-1]["content"] + print(_c(f" Q: {textwrap.shorten(user_msg, 80)}", DIM)) + print() + + try: + reasoning, normal = chat_stream( + base_url, + case["messages"], + model=model_name, + max_tokens=max_tokens, + temperature=args.temperature, + ) + except requests.HTTPError as exc: + print(_c(f" HTTP Error: {exc}", RED)) + failed += 1 + continue + except Exception as exc: + print(_c(f" Request Failed: {exc}", RED)) + failed += 1 + continue + + _print_result(reasoning, normal) + print() + + _check("Received non-empty output", bool(reasoning or normal)) + _check("reasoning_content is not empty (think tags identified)", bool(reasoning)) + _check("content is not empty (normal response preserved)", bool(normal)) + _check("reasoning_content does not contain ", "" not in reasoning) + _check("content does not contain ", "" not in normal) + + case_ok = all([bool(reasoning or normal), bool(reasoning), bool(normal), + "" not in reasoning, "" not in normal]) + if case_ok: + passed += 1 + else: + failed += 1 + + print() + print(_hr("═")) + total = passed + failed + summary_color = GREEN if failed == 0 else RED + print(_c(f" Result: {passed}/{total} passed", BOLD + summary_color)) + if failed > 0: + print(_c( + "\n Tip: If reasoning_content is empty, ensure:\n" + " 1. The engine was started with --reasoning-parser\n" + " 2. You are using a reasoning model (e.g., Qwen-R1, DeepSeek-R1)\n" + " 3. Temperature > 0 (temperature=0 might cause the model to skip )", + YELLOW, + )) + print() + +if __name__ == "__main__": + main() \ No newline at end of file