| 1 | // SPDX-License-Identifier: Apache-2.0 |
| 2 | // SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | use std::collections::BTreeSet; |
| 5 | |
| 6 | use vllm_parser::tool::test_utils::split_by_chars; |
| 7 | use vllm_parser::unified::{UnifiedParser, UnifiedParserEvent, UnifiedParserOutput}; |
| 8 | use vllm_tokenizer::{DecodedText, Tokenizer}; |
| 9 | |
| 10 | /// Split `text` into token-attributed chunks of at most `chunk_chars` |
| 11 | /// characters, as the frontend streams them to a unified parser. |
| 12 | /// |
| 13 | /// The whole text is encoded with special tokens and run through the |
| 14 | /// incremental decoder once before splitting, so a marker cut by a chunk |
| 15 | /// boundary keeps its special-token anchor. Prepare chunks outside the measured |
| 16 | /// routine. |
| 17 | pub fn attributed_chunks( |
| 18 | tokenizer: &dyn Tokenizer, |
| 19 | text: &str, |
| 20 | chunk_chars: usize, |
| 21 | ) -> Vec<DecodedText> { |
| 22 | let mut decoder = tokenizer.create_decode_stream(&[], false, 0); |
| 23 | for token_id in tokenizer.encode(text, false).expect("fixture should encode") { |
| 24 | decoder.push_token(token_id).expect("fixture token should decode"); |
| 25 | } |
| 26 | let mut rest = decoder.flush(None).expect("decoder should flush").1; |
| 27 | assert_eq!( |
| 28 | rest.text, text, |
| 29 | "fixture should round-trip through the tokenizer" |
| 30 | ); |
| 31 | |
| 32 | split_by_chars(text, chunk_chars) |
| 33 | .into_iter() |
| 34 | .map(|chunk| rest.drain_prefix(chunk.len())) |
| 35 | .collect() |
| 36 | } |
| 37 | |
| 38 | /// Event totals of one parsed stream, for benchmark sanity checks. |
| 39 | #[derive(Debug, Default)] |
| 40 | pub struct UnifiedStreamSummary { |
| 41 | pub normal_text: String, |
| 42 | pub reasoning_text: String, |
| 43 | pub calls_len: usize, |
| 44 | } |
| 45 | |
| 46 | /// Run one request stream through `parser`: initialize it from the prompt, |
| 47 | /// feed every chunk, and finish. |
| 48 | pub fn feed_unified_parser( |
| 49 | parser: &mut dyn UnifiedParser, |
| 50 | prompt_token_ids: &[u32], |
| 51 | chunks: Vec<DecodedText>, |
| 52 | ) -> UnifiedStreamSummary { |
| 53 | parser.initialize(prompt_token_ids).expect("parser should initialize"); |
| 54 | |
| 55 | let mut output = UnifiedParserOutput::default(); |
| 56 | for chunk in chunks { |
| 57 | parser.parse_into(chunk, &mut output).expect("chunk should parse"); |
| 58 | } |
| 59 | output.append(parser.finish().expect("stream should finish")); |
| 60 | |
| 61 | let mut summary = UnifiedStreamSummary::default(); |
| 62 | let mut tool_indices = BTreeSet::new(); |
| 63 | for event in output.events { |
| 64 | match event { |
| 65 | UnifiedParserEvent::Text(text) => summary.normal_text.push_str(&text), |
| 66 | UnifiedParserEvent::Reasoning(reasoning) => { |
| 67 | summary.reasoning_text.push_str(&reasoning.text); |
| 68 | } |
| 69 | UnifiedParserEvent::ToolCall(call) => { |
| 70 | tool_indices.insert(call.tool_index); |
| 71 | } |
| 72 | } |
| 73 | } |
| 74 | summary.calls_len = tool_indices.len(); |
| 75 | summary |
| 76 | } |