[Anthropic] Fix missing cache_read_input_tokens in streaming responses (#29703)
This commit is contained in:
@@ -38,6 +38,7 @@ from sglang.srt.entrypoints.openai.protocol import (
|
||||
FunctionResponse,
|
||||
LogProbs,
|
||||
MessageProcessingResult,
|
||||
PromptTokensDetails,
|
||||
ResponseParserProtocol,
|
||||
SglExt,
|
||||
ToolCall,
|
||||
@@ -336,6 +337,15 @@ class OpenAIServingChat(OpenAIServingBase):
|
||||
"""Post-process reasoning and tool_calls before building response."""
|
||||
return reasoning_text, tool_calls
|
||||
|
||||
def _continuous_usage_cached_details(
|
||||
self, content: Dict[str, Any]
|
||||
) -> Optional[PromptTokensDetails]:
|
||||
if not self.tokenizer_manager.server_args.enable_cache_report:
|
||||
return None
|
||||
return UsageProcessor._details_if_cached(
|
||||
content["meta_info"].get("cached_tokens", 0)
|
||||
)
|
||||
|
||||
async def _generate_stream_content(
|
||||
self,
|
||||
content: Dict[str, Any],
|
||||
@@ -377,6 +387,7 @@ class OpenAIServingChat(OpenAIServingBase):
|
||||
prompt_tokens=prompt_tokens.get(index, 0),
|
||||
reasoning_tokens=reasoning_tokens.get(index, 0),
|
||||
completion_tokens=completion_tokens.get(index, 0),
|
||||
cached_tokens=self._continuous_usage_cached_details(content),
|
||||
).model_dump()
|
||||
|
||||
yield build_sse_content(
|
||||
@@ -422,6 +433,7 @@ class OpenAIServingChat(OpenAIServingBase):
|
||||
prompt_tokens=prompt_tokens.get(index, 0),
|
||||
reasoning_tokens=reasoning_tokens.get(index, 0),
|
||||
completion_tokens=completion_tokens.get(index, 0),
|
||||
cached_tokens=self._continuous_usage_cached_details(content),
|
||||
).model_dump()
|
||||
|
||||
yield build_sse_content(
|
||||
@@ -449,6 +461,7 @@ class OpenAIServingChat(OpenAIServingBase):
|
||||
prompt_tokens=prompt_tokens.get(index, 0),
|
||||
reasoning_tokens=reasoning_tokens.get(index, 0),
|
||||
completion_tokens=completion_tokens.get(index, 0),
|
||||
cached_tokens=self._continuous_usage_cached_details(content),
|
||||
).model_dump()
|
||||
|
||||
yield build_sse_content(
|
||||
@@ -1898,6 +1911,7 @@ class OpenAIServingChat(OpenAIServingBase):
|
||||
prompt_tokens=prompt_tokens,
|
||||
completion_tokens=completion_tokens,
|
||||
reasoning_tokens=reasoning_tokens,
|
||||
cached_tokens=self._continuous_usage_cached_details(content),
|
||||
)
|
||||
|
||||
yield f"data: {chunk.model_dump_json()}\n\n"
|
||||
@@ -1950,6 +1964,7 @@ class OpenAIServingChat(OpenAIServingBase):
|
||||
prompt_tokens=prompt_tokens,
|
||||
completion_tokens=completion_tokens,
|
||||
reasoning_tokens=reasoning_tokens,
|
||||
cached_tokens=self._continuous_usage_cached_details(content),
|
||||
)
|
||||
|
||||
yield f"data: {chunk.model_dump_json()}\n\n"
|
||||
|
||||
@@ -1800,6 +1800,63 @@ class ServingChatTestCase(unittest.TestCase):
|
||||
},
|
||||
)
|
||||
|
||||
def _collect_continuous_usage(self, cached_tokens):
|
||||
content = {
|
||||
"text": "Hello",
|
||||
"meta_info": {
|
||||
"id": "chatcmpl-cont-usage",
|
||||
"prompt_tokens": 10,
|
||||
"completion_tokens": 2,
|
||||
"cached_tokens": cached_tokens,
|
||||
"finish_reason": {"type": "stop", "matched": None},
|
||||
"output_token_logprobs": None,
|
||||
"output_top_logprobs": None,
|
||||
},
|
||||
"index": 0,
|
||||
}
|
||||
req = ChatCompletionRequest(
|
||||
model="x",
|
||||
messages=[{"role": "user", "content": "Hi?"}],
|
||||
stream=True,
|
||||
)
|
||||
|
||||
async def _collect():
|
||||
chunks = []
|
||||
async for chunk in self.chat._generate_stream_content(
|
||||
content=content,
|
||||
index=0,
|
||||
request=req,
|
||||
stream_offsets={},
|
||||
reasoning_parser_dict={},
|
||||
parser_dict={},
|
||||
has_tool_calls={},
|
||||
choice_logprobs=None,
|
||||
finish_reason_type="stop",
|
||||
continuous_usage_stats=True,
|
||||
prompt_tokens={0: 10},
|
||||
reasoning_tokens={0: 0},
|
||||
completion_tokens={0: 2},
|
||||
):
|
||||
chunks.append(chunk)
|
||||
return chunks
|
||||
|
||||
chunks = get_or_create_event_loop().run_until_complete(_collect())
|
||||
return [c["usage"] for c in self._parse_chunks(chunks) if c.get("usage")]
|
||||
|
||||
def test_continuous_usage_reports_cached_tokens(self):
|
||||
"""continuous_usage_stats chunks include cached tokens when cache reporting is on."""
|
||||
self.tm.server_args.enable_cache_report = True
|
||||
usages = self._collect_continuous_usage(cached_tokens=6)
|
||||
self.assertTrue(usages, "continuous_usage_stats attached no usage")
|
||||
self.assertEqual(usages[0]["prompt_tokens_details"]["cached_tokens"], 6)
|
||||
|
||||
def test_continuous_usage_omits_cached_tokens_when_report_disabled(self):
|
||||
"""With cache reporting off, continuous_usage_stats must not leak cached tokens."""
|
||||
self.tm.server_args.enable_cache_report = False
|
||||
usages = self._collect_continuous_usage(cached_tokens=6)
|
||||
self.assertTrue(usages, "continuous_usage_stats attached no usage")
|
||||
self.assertIsNone(usages[0].get("prompt_tokens_details"))
|
||||
|
||||
# ------------- incremental streaming output tests -------------
|
||||
def test_incremental_streaming_output_delta(self):
|
||||
"""Test that streaming with incremental_streaming_output produces correct deltas.
|
||||
|
||||
Reference in New Issue
Block a user