[Bugfix] Fix KeyError: 'prompt_tokens' when streaming requests are aborted (#19514)

python/sglang/srt/entrypoints/openai/serving_chat.py

@@ -632,8 +632,10 @@ class OpenAIServingChat(OpenAIServingBase):
             ):
                 index = content.get("index", 0)
 
-                prompt_tokens[index] = content["meta_info"]["prompt_tokens"]
+                prompt_tokens[index] = content["meta_info"].get("prompt_tokens", 0)
-                completion_tokens[index] = content["meta_info"]["completion_tokens"]
+                completion_tokens[index] = content["meta_info"].get(
+                    "completion_tokens", 0
+                )
                 cached_tokens[index] = content["meta_info"].get("cached_tokens", 0)
                 hidden_states[index] = content["meta_info"].get("hidden_states", None)
                 routed_experts[index] = content["meta_info"].get("routed_experts", None)

python/sglang/srt/entrypoints/openai/serving_completions.py

@@ -209,8 +209,10 @@ class OpenAIServingCompletion(OpenAIServingBase):
                 index = content.get("index", 0)
 
                 text = content["text"]
-                prompt_tokens[index] = content["meta_info"]["prompt_tokens"]
+                prompt_tokens[index] = content["meta_info"].get("prompt_tokens", 0)
-                completion_tokens[index] = content["meta_info"]["completion_tokens"]
+                completion_tokens[index] = content["meta_info"].get(
+                    "completion_tokens", 0
+                )
                 cached_tokens[index] = content["meta_info"].get("cached_tokens", 0)
                 hidden_states[index] = content["meta_info"].get("hidden_states", None)
                 routed_experts[index] = content["meta_info"].get("routed_experts", None)

python/sglang/srt/entrypoints/openai/usage_processor.py

@@ -20,10 +20,12 @@ class UsageProcessor:
         n_choices: int = 1,
         enable_cache_report: bool = False,
     ) -> UsageInfo:
-        completion_tokens = sum(r["meta_info"]["completion_tokens"] for r in responses)
+        completion_tokens = sum(
+            r["meta_info"].get("completion_tokens", 0) for r in responses
+        )
 
         prompt_tokens = sum(
-            responses[i]["meta_info"]["prompt_tokens"]
+            responses[i]["meta_info"].get("prompt_tokens", 0)
             for i in range(0, len(responses), n_choices)
         )