sgl-project--sglang/python/sglang/test/kits/reasoning_kit.py

import json

import openai
import requests

from sglang.srt.parser.reasoning_parser import ReasoningParser
from sglang.srt.utils.hf_transformers_utils import get_tokenizer


class ReasoningTokenUsageMixin:
    """Mixin for reasoning_tokens usage tests.

    Required attributes on the test class:
        model: str
        base_url: str
        reasoning_parser_name: str

    Optional attributes:
        api_key: str (if not set, no auth)

    Call cls.init_reasoning_token_verifier() in setUpClass.
    """

    reasoning_parser_name = None

    @classmethod
    def init_reasoning_token_verifier(cls):
        assert cls.reasoning_parser_name, "reasoning_parser_name must be set"
        cls.tokenizer = get_tokenizer(cls.model)
        parser = ReasoningParser(cls.reasoning_parser_name, tokenizer=cls.tokenizer)
        cls.think_end_token_id = cls.tokenizer.convert_tokens_to_ids(
            parser.detector.think_end_token
        )
        assert cls.think_end_token_id is not None

    def _reasoning_chat_request(self, enable_thinking, stream=False):
        api_key = getattr(self, "api_key", None)
        headers = {"Authorization": f"Bearer {api_key}"} if api_key else {}
        payload = {
            "model": self.model,
            "messages": [{"role": "user", "content": "What is 1+3?"}],
            "max_tokens": 1024,
            "chat_template_kwargs": {"enable_thinking": enable_thinking},
        }
        if stream:
            payload["stream"] = True
            payload["stream_options"] = {"include_usage": True}
        return requests.post(
            f"{self.base_url}/v1/chat/completions",
            headers=headers,
            json=payload,
            stream=stream,
        )

    def _extract_streaming_usage(self, response):
        usage = None
        for line in response.iter_lines():
            if not line:
                continue
            decoded = line.decode("utf-8")
            if not decoded.startswith("data:") or decoded.startswith("data: [DONE]"):
                continue
            data = json.loads(decoded[len("data:") :].strip())
            if data.get("usage"):
                usage = data["usage"]
        return usage

    def test_reasoning_tokens_thinking(self):
        resp = self._reasoning_chat_request(enable_thinking=True)
        self.assertEqual(resp.status_code, 200, resp.text)
        usage = resp.json()["usage"]
        self.assertGreater(usage["reasoning_tokens"], 0)
        self.assertLess(usage["reasoning_tokens"], usage["completion_tokens"])

    def test_reasoning_tokens_non_thinking(self):
        resp = self._reasoning_chat_request(enable_thinking=False)
        self.assertEqual(resp.status_code, 200, resp.text)
        self.assertEqual(resp.json()["usage"]["reasoning_tokens"], 0)

    def test_reasoning_tokens_thinking_stream(self):
        with self._reasoning_chat_request(enable_thinking=True, stream=True) as resp:
            self.assertEqual(resp.status_code, 200, resp.text)
            usage = self._extract_streaming_usage(resp)
            self.assertIsNotNone(usage, "No usage in stream")
            self.assertGreater(usage["reasoning_tokens"], 0)
            self.assertLess(usage["reasoning_tokens"], usage["completion_tokens"])

    def test_reasoning_tokens_non_thinking_stream(self):
        with self._reasoning_chat_request(enable_thinking=False, stream=True) as resp:
            self.assertEqual(resp.status_code, 200, resp.text)
            usage = self._extract_streaming_usage(resp)
            self.assertIsNotNone(usage, "No usage in stream")
            self.assertEqual(usage["reasoning_tokens"], 0)

    def test_reasoning_tokens_generate_exact_count(self):
        api_key = getattr(self, "api_key", None)
        headers = {"Authorization": f"Bearer {api_key}"} if api_key else {}
        messages = [{"role": "user", "content": "What is 1+3?"}]
        prompt = self.tokenizer.apply_chat_template(
            messages, add_generation_prompt=True, tokenize=False
        )
        resp = requests.post(
            f"{self.base_url}/generate",
            headers=headers,
            json={
                "text": prompt,
                "sampling_params": {"max_new_tokens": 1024},
                "require_reasoning": True,
            },
        )
        self.assertEqual(resp.status_code, 200, resp.text)
        data = resp.json()
        reported = data["meta_info"]["reasoning_tokens"]
        actual = data["output_ids"].index(self.think_end_token_id) + 1
        self.assertEqual(reported, actual)


class SeparateReasoningMixin:
    """Mixin for separate_reasoning tests.

    Required attributes on the test class:
        model: str
        base_url: str (without /v1)
        api_key: str
    """

    def _openai_client(self):
        return openai.Client(api_key=self.api_key, base_url=f"{self.base_url}/v1")

    def _chat(self, stream=False, extra_body=None):
        return self._openai_client().chat.completions.create(
            model=self.model,
            messages=[{"role": "user", "content": "What is 1+3?"}],
            max_tokens=1024,
            stream=stream,
            extra_body=extra_body,
        )

    def _collect_stream(self, response):
        reasoning_content = ""
        content = ""
        for chunk in response:
            if chunk.choices[0].delta.content:
                content += chunk.choices[0].delta.content
            elif chunk.choices[0].delta.reasoning_content:
                reasoning_content += chunk.choices[0].delta.reasoning_content
        return reasoning_content, content

    def test_streaming_separate_reasoning_false(self):
        response = self._chat(stream=True, extra_body={"separate_reasoning": False})
        reasoning_content, content = self._collect_stream(response)
        self.assertEqual(len(reasoning_content), 0)
        self.assertGreater(len(content), 0)

    def test_streaming_separate_reasoning_true(self):
        response = self._chat(stream=True, extra_body={"separate_reasoning": True})
        reasoning_content, content = self._collect_stream(response)
        self.assertGreater(len(reasoning_content), 0)
        self.assertGreater(len(content), 0)

    def test_streaming_separate_reasoning_true_stream_reasoning_false(self):
        response = self._chat(
            stream=True,
            extra_body={"separate_reasoning": True, "stream_reasoning": False},
        )
        reasoning_content = ""
        content = ""
        first_chunk = False
        for chunk in response:
            if chunk.choices[0].delta.reasoning_content:
                reasoning_content = chunk.choices[0].delta.reasoning_content
                first_chunk = True
            if chunk.choices[0].delta.content:
                content += chunk.choices[0].delta.content
                if not first_chunk:
                    reasoning_content = chunk.choices[0].delta.reasoning_content
                first_chunk = True
            if not first_chunk:
                assert (
                    not chunk.choices[0].delta.reasoning_content
                    or len(chunk.choices[0].delta.reasoning_content) == 0
                )
        self.assertGreater(len(reasoning_content), 0)
        self.assertGreater(len(content), 0)

    def test_nonstreaming_separate_reasoning_false(self):
        response = self._chat(extra_body={"separate_reasoning": False})
        assert (
            not response.choices[0].message.reasoning_content
            or len(response.choices[0].message.reasoning_content) == 0
        )
        self.assertGreater(len(response.choices[0].message.content), 0)

    def test_nonstreaming_separate_reasoning_true(self):
        response = self._chat(extra_body={"separate_reasoning": True})
        self.assertGreater(len(response.choices[0].message.reasoning_content), 0)
        self.assertGreater(len(response.choices[0].message.content), 0)