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>

tests/reasoning/test_cohere_command_reasoning_parser.py

@@ -253,6 +253,25 @@ class TestIsReasoningEnd:
         )
 
 
+@pytest.mark.parametrize(
+    "parser_cls",
+    [
+        CohereCommand3ReasoningParser,
+        CohereCommand4ReasoningParser,
+    ],
+    ids=["cmd3", "cmd4"],
+)
+def test_count_reasoning_tokens(tokenizer, parser_cls):
+    parser = parser_cls(tokenizer)
+    start = parser.start_token_id
+    end = parser.end_token_id
+
+    assert parser.count_reasoning_tokens([99, start, 11, 12, end, 100]) == 2
+    assert parser.count_reasoning_tokens([end, start, 1, start, 2, end, 3, end]) == 3
+    assert parser.count_reasoning_tokens([start, 1, 2]) == 2
+    assert parser.count_reasoning_tokens([1, 2, end]) == 0
+
+
 SCHEMA_A = {"type": "object", "properties": {"a": {"type": "string"}}}
 SCHEMA_B = {"type": "object", "properties": {"b": {"type": "number"}}}
 GET_WEATHER_TOOL = {

vllm/reasoning/cohere_command_reasoning_parser.py

@@ -628,6 +628,21 @@ class BaseCohereCommandReasoningParser(ReasoningParser):
 
         return has_end_token
 
+    def count_reasoning_tokens(self, token_ids: Sequence[int]) -> int:
+        count = 0
+        depth = 0
+        for token_id in token_ids:
+            if token_id == self.start_token_id:
+                depth += 1
+                continue
+            if token_id == self.end_token_id:
+                if depth > 0:
+                    depth -= 1
+                continue
+            if depth > 0:
+                count += 1
+        return count
+
     def adjust_request(
         self, request: ChatCompletionRequest | ResponsesRequest
     ) -> ChatCompletionRequest | ResponsesRequest: