[bugfix] Parse token usage for TGI/Triton/MindIE backends

ais_bench/benchmark/models/api_models/base_api.py

@@ -249,6 +249,21 @@ def _logprobs_enabled(self) -> bool:
             return val > 0
         return False
 
+    @staticmethod
+    def _update_usage_from_details(api_response: dict, output: Output) -> None:
+        """Fill token counters from a ``details`` block reported by the service.
+
+        TGI / Triton / MindIE expose token usage as ``details.generated_tokens``
+        (cumulative for streaming) instead of the OpenAI ``usage`` object. Without
+        this, every performance report for those backends shows zero throughput.
+        """
+        details = api_response.get("details")
+        if not isinstance(details, dict):
+            return
+        generated_tokens = details.get("generated_tokens")
+        if isinstance(generated_tokens, int) and not isinstance(generated_tokens, bool):
+            output.output_tokens = generated_tokens
+
     async def parse_text_response(self, data, output):
         raise AISBenchNotImplementedError(
             MODEL_CODES.PARSE_TEXT_RSP_NOT_IMPLEMENTED,

ais_bench/benchmark/models/api_models/mindie_stream_api.py

@@ -122,4 +122,6 @@ async def get_request_body(
 
     async def parse_stream_response(self, api_response: dict, output: Output):
         generated_text = api_response.get("generated_text", "")
-        output.content = generated_text
+        output.content = generated_text
+        # The token count is only reported on the final chunk.
+        self._update_usage_from_details(api_response, output)

ais_bench/benchmark/models/api_models/tgi_api.py

@@ -125,11 +125,14 @@ async def get_request_body(
     async def parse_text_response(self, api_response:dict, output:Output):
         generated_text = api_response.get("generated_text", "")
         output.content = generated_text
+        self._update_usage_from_details(api_response, output)
 
     async def parse_stream_response(self, api_response:dict, output:Output):
         generated_text = api_response.get("generated_text", "")
         # TGI will return the generated text in last chunk
         output.content = generated_text
+        # The token count is only reported on the final chunk.
+        self._update_usage_from_details(api_response, output)
 
 
 class TGICustomAPIStream(TGICustomAPI):

ais_bench/benchmark/models/api_models/triton_api.py

@@ -92,10 +92,13 @@ async def get_request_body(self, input_data: PromptType, max_out_len: int, outpu
     async def parse_text_response(self, api_response:dict, output:Output):
         generated_text = api_response.get("text_output", "")
         output.content = generated_text
+        self._update_usage_from_details(api_response, output)
 
     async def parse_stream_response(self, api_response:dict, output:Output):
         generated_text = api_response.get("text_output", "")
         output.content += generated_text
+        # ``details.generated_tokens`` is cumulative across chunks.
+        self._update_usage_from_details(api_response, output)
 
 
 class TritonCustomAPIStream(TritonCustomAPI):

ais_bench/benchmark/summarizers/default_perf.py

@@ -155,10 +155,13 @@ def recursive_update(detail_data):
                 pass # Prefer service-returned "prompt_tokens".
             elif is_mm_prompt(perf_data["input"]):
                 perf_data["input_tokens"] = 0  # multi-modal input does not support input_tokens
-            elif "input_tokens" not in perf_data or perf_data.get("input_tokens") is None:
-                perf_data["input_tokens"] = len(tokenizer.encode(perf_data["input"])) # input_tokens is not provided, calculate it
+            else:
+                # `Output` initializes the counters to 0, so a missing value is
+                # indistinguishable from "not reported" by the key alone.
+                perf_data["input_tokens"] = len(tokenizer.encode(perf_data["input"]))
 
-            if "output_tokens" not in perf_data or perf_data.get("output_tokens") is None:
+            output_tokens = perf_data.get("output_tokens")
+            if (output_tokens is None or output_tokens <= 0) and isinstance(perf_data.get("prediction"), str):
                 perf_data["output_tokens"] = len(tokenizer.encode(perf_data["prediction"]))
             perf_data.pop("input")
             perf_data.pop("prediction")

tests/UT/models/api_models/test_mindie_stream_api.py

@@ -236,6 +236,33 @@ def test_is_api_property(self):
         model = MindieStreamApi(**self.default_kwargs)
         self.assertTrue(model.is_api)
 
+    def _run_parse(self, parse_name, api_response):
+        async def run():
+            model = MindieStreamApi(**self.default_kwargs)
+            output = Output()
+            await getattr(model, parse_name)(api_response, output)
+            return output
+        return asyncio.run(run())
+
+    def test_parse_stream_response_records_output_tokens_from_final_chunk(self):
+        """MindIE reports token usage under details.generated_tokens on the final chunk"""
+        output = self._run_parse("parse_stream_response", {"token": {"text": "Hel"}})
+        self.assertEqual(output.output_tokens, 0)
+
+        output = self._run_parse(
+            "parse_stream_response",
+            {"generated_text": "Hello world!", "details": {"generated_tokens": 12}},
+        )
+        self.assertEqual(output.content, "Hello world!")
+        self.assertEqual(output.output_tokens, 12)
+
+    def test_parse_stream_response_without_details_keeps_token_count(self):
+        """A response without a usage block must not crash or invent a count"""
+        output = self._run_parse("parse_stream_response", {"generated_text": "hi"})
+
+        self.assertEqual(output.content, "hi")
+        self.assertEqual(output.output_tokens, 0)
+
 
 class TestMindieStreamApiLora(unittest.TestCase):
     """针对 Multi-LoRA 兼容性扩展的 UT（MindIE 通过 parameters.adapter_id 路由 LoRA）。"""

tests/UT/models/api_models/test_tgi_api.py

@@ -245,6 +245,43 @@ async def test_generate_integration_stream_true(self, mock_session_class):
             await model.generate("test prompt", 100, output)
             mock_stream_infer.assert_called_once()
 
+    def _run_parse(self, parse_name, api_response):
+        async def run():
+            model = TGICustomAPI(**self.default_kwargs)
+            output = Output()
+            await getattr(model, parse_name)(api_response, output)
+            return output
+        return asyncio.run(run())
+
+    def test_parse_text_response_records_output_tokens(self):
+        """TGI reports token usage under details.generated_tokens"""
+        output = self._run_parse(
+            "parse_text_response",
+            {"generated_text": "Hello world!", "details": {"generated_tokens": 12}},
+        )
+
+        self.assertEqual(output.content, "Hello world!")
+        self.assertEqual(output.output_tokens, 12)
+
+    def test_parse_stream_response_records_output_tokens_from_final_chunk(self):
+        """Streaming: intermediate chunks carry no details, the final one does"""
+        output = self._run_parse("parse_stream_response", {"token": {"text": "Hel"}})
+        self.assertEqual(output.output_tokens, 0)
+
+        output = self._run_parse(
+            "parse_stream_response",
+            {"generated_text": "Hello world!", "details": {"generated_tokens": 12}},
+        )
+        self.assertEqual(output.content, "Hello world!")
+        self.assertEqual(output.output_tokens, 12)
+
+    def test_parse_response_without_details_keeps_token_count(self):
+        """A response without a usage block must not crash or invent a count"""
+        output = self._run_parse("parse_text_response", {"generated_text": "hi"})
+
+        self.assertEqual(output.content, "hi")
+        self.assertEqual(output.output_tokens, 0)
+
 
 class TestTGICustomAPILora(unittest.TestCase):
     """针对 Multi-LoRA 兼容性扩展的 UT（TGI 通过 parameters.adapter_id 路由 LoRA）。"""

tests/UT/models/api_models/test_triton_api.py

@@ -294,5 +294,55 @@ def test_get_service_model_path_call(self, mock_get_model_path):
         mock_get_model_path.assert_called_once()
 
 
+    def _run_parse(self, parse_name, model, api_response):
+        async def run():
+            output = Output()
+            await getattr(model, parse_name)(api_response, output)
+            return output
+        return asyncio.run(run())
+
+    def test_parse_text_response_records_output_tokens(self):
+        """Triton reports token usage under details.generated_tokens"""
+        with mock.patch.object(base_api.BaseAPIModel, '_get_service_model_path', return_value="m"):
+            model = TritonCustomAPI(path="test", model_name="m")
+        output = self._run_parse(
+            "parse_text_response",
+            model,
+            {"text_output": "Hello world!", "details": {"generated_tokens": 12}},
+        )
+
+        self.assertEqual(output.content, "Hello world!")
+        self.assertEqual(output.output_tokens, 12)
+
+    def test_parse_stream_response_uses_cumulative_output_tokens(self):
+        """Streaming: details.generated_tokens is cumulative, so the last chunk wins"""
+        with mock.patch.object(base_api.BaseAPIModel, '_get_service_model_path', return_value="m"):
+            model = TritonCustomAPI(path="test", model_name="m")
+
+        async def run():
+            output = Output()
+            await model.parse_stream_response(
+                {"text_output": "Hel", "details": {"generated_tokens": 3}}, output
+            )
+            assert output.output_tokens == 3
+            await model.parse_stream_response(
+                {"text_output": "lo", "details": {"generated_tokens": 5}}, output
+            )
+            return output
+
+        output = asyncio.run(run())
+        self.assertEqual(output.content, "Hello")
+        self.assertEqual(output.output_tokens, 5)
+
+    def test_parse_response_without_details_keeps_token_count(self):
+        """A response without a usage block must not crash or invent a count"""
+        with mock.patch.object(base_api.BaseAPIModel, '_get_service_model_path', return_value="m"):
+            model = TritonCustomAPI(path="test", model_name="m")
+        output = self._run_parse("parse_text_response", model, {"text_output": "hi"})
+
+        self.assertEqual(output.content, "hi")
+        self.assertEqual(output.output_tokens, 0)
+
+
 if __name__ == '__main__':
     unittest.main()

tests/UT/summarizers/test_default_perf.py

@@ -131,6 +131,9 @@ def test_calc_perf_data_success(self, mock_load_tokenizer, mock_aistokenizer, mo
         self.assertEqual(result["end_time"], 3)
         self.assertEqual(result["latency"], 3)
         self.assertEqual(result["ttft"], 1)
+        # The service reported output_tokens=0 (i.e. nothing reported); the
+        # summarizer must fall back to the tokenizer instead of keeping 0.
+        self.assertEqual(result["output_tokens"], len("test output"))
         mock_conn.close.assert_called_once()
 
     @patch('ais_bench.benchmark.summarizers.default_perf.build_model_from_cfg')