feat: Add support for reasoning_token_count to reasoning parser (#54982)
Signed-off-by: Jason Ozuzu <jasonozuzu@cohere.com> Co-authored-by: Cursor <cursoragent@cursor.com>
JasonCohere committed
Sep 3, 2026 at 13:02 UTC
bb363db9a5ec2edc7b39e99b00af363a89d1fb81
2 files changed
+34
tests/reasoning/test_cohere_command_reasoning_parser.py
+19
@@ -253,6 +253,25 @@ class TestIsReasoningEnd:
253
)
254
255
256
+@pytest.mark.parametrize(
257
+ "parser_cls",
258
+ [
259
+ CohereCommand3ReasoningParser,
260
+ CohereCommand4ReasoningParser,
261
+ ],
262
+ ids=["cmd3", "cmd4"],
263
+)
264
+def test_count_reasoning_tokens(tokenizer, parser_cls):
265
+ parser = parser_cls(tokenizer)
266
+ start = parser.start_token_id
267
+ end = parser.end_token_id
268
+
269
+ assert parser.count_reasoning_tokens([99, start, 11, 12, end, 100]) == 2
270
+ assert parser.count_reasoning_tokens([end, start, 1, start, 2, end, 3, end]) == 3
271
+ assert parser.count_reasoning_tokens([start, 1, 2]) == 2
272
+ assert parser.count_reasoning_tokens([1, 2, end]) == 0
273
+
274
+
275
SCHEMA_A = {"type": "object", "properties": {"a": {"type": "string"}}}
276
SCHEMA_B = {"type": "object", "properties": {"b": {"type": "number"}}}
277
GET_WEATHER_TOOL = {
vllm/reasoning/cohere_command_reasoning_parser.py
+15
@@ -628,6 +628,21 @@ class BaseCohereCommandReasoningParser(ReasoningParser):
628
629
return has_end_token
630
631
+ def count_reasoning_tokens(self, token_ids: Sequence[int]) -> int:
632
+ count = 0
633
+ depth = 0
634
+ for token_id in token_ids:
635
+ if token_id == self.start_token_id:
636
+ depth += 1
637
+ continue
638
+ if token_id == self.end_token_id:
639
+ if depth > 0:
640
+ depth -= 1
641
+ continue
642
+ if depth > 0:
643
+ count += 1
644
+ return count
645
+
646
def adjust_request(
647
self, request: ChatCompletionRequest | ResponsesRequest
648
) -> ChatCompletionRequest | ResponsesRequest: