Add continuous_usage_stats support for streaming responses (#12241)

test/srt/openai_server/features/test_continuous_usage_stats.py

@@ -0,0 +1,105 @@
+import asyncio
+import unittest
+
+import openai
+
+from sglang.srt.utils import kill_process_tree
+from sglang.test.test_utils import (
+    DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
+    DEFAULT_URL_FOR_TEST,
+    CustomTestCase,
+    popen_launch_server,
+)
+
+
+class TestContinuousUsageStats(CustomTestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.model = DEFAULT_SMALL_MODEL_NAME_FOR_TEST
+        cls.base_url = DEFAULT_URL_FOR_TEST
+        cls.process = popen_launch_server(cls.model, cls.base_url, timeout=300)
+        cls.client = openai.Client(api_key="EMPTY", base_url=f"{cls.base_url}/v1")
+        cls.aclient = openai.AsyncClient(api_key="EMPTY", base_url=f"{cls.base_url}/v1")
+
+    @classmethod
+    def tearDownClass(cls):
+        kill_process_tree(cls.process.pid)
+
+    def test_continuous_usage_stats_enabled(self):
+        stream = self.client.chat.completions.create(
+            model=self.model,
+            messages=[{"role": "user", "content": "What is machine learning?"}],
+            stream=True,
+            max_tokens=30,
+            temperature=0,
+            stream_options={"include_usage": True, "continuous_usage_stats": True},
+        )
+
+        chunks_with_usage = 0
+        chunks_with_content = 0
+        last_usage = None
+
+        for chunk in stream:
+            has_content = len(chunk.choices) > 0 and chunk.choices[0].delta.content
+            if chunk.usage:
+                chunks_with_usage += 1
+                last_usage = chunk.usage
+            if has_content:
+                chunks_with_content += 1
+
+        assert chunks_with_content > 0
+        assert chunks_with_usage >= chunks_with_content
+        assert last_usage.prompt_tokens > 0
+        assert last_usage.completion_tokens > 0
+        assert (
+            last_usage.total_tokens
+            == last_usage.prompt_tokens + last_usage.completion_tokens
+        )
+
+    async def test_continuous_usage_stats_async(self):
+        stream = await self.aclient.chat.completions.create(
+            model=self.model,
+            messages=[{"role": "user", "content": "What is deep learning?"}],
+            stream=True,
+            max_tokens=30,
+            temperature=0,
+            stream_options={"include_usage": True, "continuous_usage_stats": True},
+        )
+
+        chunks_with_usage = 0
+        chunks_with_content = 0
+
+        async for chunk in stream:
+            has_content = len(chunk.choices) > 0 and chunk.choices[0].delta.content
+            if chunk.usage:
+                chunks_with_usage += 1
+            if has_content:
+                chunks_with_content += 1
+
+        assert chunks_with_content > 0
+        assert chunks_with_usage >= chunks_with_content
+
+    def test_continuous_usage_stats_disabled(self):
+        stream = self.client.chat.completions.create(
+            model=self.model,
+            messages=[{"role": "user", "content": "What is AI?"}],
+            stream=True,
+            max_tokens=30,
+            temperature=0,
+            stream_options={"include_usage": True, "continuous_usage_stats": False},
+        )
+
+        usage_chunks = []
+        for chunk in stream:
+            if chunk.usage:
+                usage_chunks.append(chunk)
+
+        assert len(usage_chunks) == 1
+        assert len(usage_chunks[0].choices) == 0
+
+    def test_async_runner(self):
+        asyncio.run(self.test_continuous_usage_stats_async())
+
+
+if __name__ == "__main__":
+    unittest.main()