[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,
|
FunctionResponse,
|
||||||
LogProbs,
|
LogProbs,
|
||||||
MessageProcessingResult,
|
MessageProcessingResult,
|
||||||
|
PromptTokensDetails,
|
||||||
ResponseParserProtocol,
|
ResponseParserProtocol,
|
||||||
SglExt,
|
SglExt,
|
||||||
ToolCall,
|
ToolCall,
|
||||||
@@ -336,6 +337,15 @@ class OpenAIServingChat(OpenAIServingBase):
|
|||||||
"""Post-process reasoning and tool_calls before building response."""
|
"""Post-process reasoning and tool_calls before building response."""
|
||||||
return reasoning_text, tool_calls
|
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(
|
async def _generate_stream_content(
|
||||||
self,
|
self,
|
||||||
content: Dict[str, Any],
|
content: Dict[str, Any],
|
||||||
@@ -377,6 +387,7 @@ class OpenAIServingChat(OpenAIServingBase):
|
|||||||
prompt_tokens=prompt_tokens.get(index, 0),
|
prompt_tokens=prompt_tokens.get(index, 0),
|
||||||
reasoning_tokens=reasoning_tokens.get(index, 0),
|
reasoning_tokens=reasoning_tokens.get(index, 0),
|
||||||
completion_tokens=completion_tokens.get(index, 0),
|
completion_tokens=completion_tokens.get(index, 0),
|
||||||
|
cached_tokens=self._continuous_usage_cached_details(content),
|
||||||
).model_dump()
|
).model_dump()
|
||||||
|
|
||||||
yield build_sse_content(
|
yield build_sse_content(
|
||||||
@@ -422,6 +433,7 @@ class OpenAIServingChat(OpenAIServingBase):
|
|||||||
prompt_tokens=prompt_tokens.get(index, 0),
|
prompt_tokens=prompt_tokens.get(index, 0),
|
||||||
reasoning_tokens=reasoning_tokens.get(index, 0),
|
reasoning_tokens=reasoning_tokens.get(index, 0),
|
||||||
completion_tokens=completion_tokens.get(index, 0),
|
completion_tokens=completion_tokens.get(index, 0),
|
||||||
|
cached_tokens=self._continuous_usage_cached_details(content),
|
||||||
).model_dump()
|
).model_dump()
|
||||||
|
|
||||||
yield build_sse_content(
|
yield build_sse_content(
|
||||||
@@ -449,6 +461,7 @@ class OpenAIServingChat(OpenAIServingBase):
|
|||||||
prompt_tokens=prompt_tokens.get(index, 0),
|
prompt_tokens=prompt_tokens.get(index, 0),
|
||||||
reasoning_tokens=reasoning_tokens.get(index, 0),
|
reasoning_tokens=reasoning_tokens.get(index, 0),
|
||||||
completion_tokens=completion_tokens.get(index, 0),
|
completion_tokens=completion_tokens.get(index, 0),
|
||||||
|
cached_tokens=self._continuous_usage_cached_details(content),
|
||||||
).model_dump()
|
).model_dump()
|
||||||
|
|
||||||
yield build_sse_content(
|
yield build_sse_content(
|
||||||
@@ -1898,6 +1911,7 @@ class OpenAIServingChat(OpenAIServingBase):
|
|||||||
prompt_tokens=prompt_tokens,
|
prompt_tokens=prompt_tokens,
|
||||||
completion_tokens=completion_tokens,
|
completion_tokens=completion_tokens,
|
||||||
reasoning_tokens=reasoning_tokens,
|
reasoning_tokens=reasoning_tokens,
|
||||||
|
cached_tokens=self._continuous_usage_cached_details(content),
|
||||||
)
|
)
|
||||||
|
|
||||||
yield f"data: {chunk.model_dump_json()}\n\n"
|
yield f"data: {chunk.model_dump_json()}\n\n"
|
||||||
@@ -1950,6 +1964,7 @@ class OpenAIServingChat(OpenAIServingBase):
|
|||||||
prompt_tokens=prompt_tokens,
|
prompt_tokens=prompt_tokens,
|
||||||
completion_tokens=completion_tokens,
|
completion_tokens=completion_tokens,
|
||||||
reasoning_tokens=reasoning_tokens,
|
reasoning_tokens=reasoning_tokens,
|
||||||
|
cached_tokens=self._continuous_usage_cached_details(content),
|
||||||
)
|
)
|
||||||
|
|
||||||
yield f"data: {chunk.model_dump_json()}\n\n"
|
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 -------------
|
# ------------- incremental streaming output tests -------------
|
||||||
def test_incremental_streaming_output_delta(self):
|
def test_incremental_streaming_output_delta(self):
|
||||||
"""Test that streaming with incremental_streaming_output produces correct deltas.
|
"""Test that streaming with incremental_streaming_output produces correct deltas.
|
||||||
|
|||||||
Reference in New Issue
Block a user