preprocessor.rs 55 KB
Newer Older
1
// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
Biswa Panda's avatar
Biswa Panda committed
2
3
4
5
6
// SPDX-License-Identifier: Apache-2.0

//! The Preprocessor consists of the following modules
//!
//! - `translation`: This module converts the allowed Ingress message types to the corresponding
7
//!   internal representation.
Biswa Panda's avatar
Biswa Panda committed
8
9
10
11
12
13
//! - `apply`: This module applies ModelConfig defaults to any empty optional fields specified
//! - `prompt`: This module applies any prompt template logic to the internal Request object.
//! - `tokenize`: This module tokenizes the formatted prompt string and returns the token ids.
//!
//! The Preprocessor will accept any IngressRequest and transform it to a BackendRequest.

14
pub mod media;
Biswa Panda's avatar
Biswa Panda committed
15
pub mod prompt;
Yan Ru Pei's avatar
Yan Ru Pei committed
16
pub mod speculative_prefill;
Biswa Panda's avatar
Biswa Panda committed
17
pub mod tools;
18
use anyhow::Context;
19
use anyhow::{Result, bail};
20
21
22
23
use dynamo_async_openai::types::{
    ChatCompletionRequestMessage, ChatCompletionRequestUserMessageContent,
    ChatCompletionRequestUserMessageContentPart, ChatCompletionToolChoiceOption, EncodingFormat,
};
Ryan Olson's avatar
Ryan Olson committed
24
use futures::Stream;
Biswa Panda's avatar
Biswa Panda committed
25
26
use futures::stream::{self, StreamExt};
use prompt::OAIPromptFormatter;
27
use std::time::{Duration, Instant};
Ryan Olson's avatar
Ryan Olson committed
28
use std::{collections::HashMap, pin::Pin, sync::Arc};
Biswa Panda's avatar
Biswa Panda committed
29
30
use tracing;

31
use crate::model_card::{ModelDeploymentCard, ModelInfo};
32
use crate::preprocessor::media::MediaLoader;
Biswa Panda's avatar
Biswa Panda committed
33
use crate::preprocessor::prompt::OAIChatLikeRequest;
34
use crate::protocols::common::preprocessor::{
35
    MultimodalData, MultimodalDataMap, PreprocessedRequestBuilder, RoutingHints,
36
};
37
use crate::protocols::common::timing::RequestTracker;
38
use crate::tokenizers::Encoding;
Biswa Panda's avatar
Biswa Panda committed
39

40
use dynamo_parsers::{ReasoningParser, ReasoningParserType};
Neelay Shah's avatar
Neelay Shah committed
41
42
use dynamo_runtime::engine::{AsyncEngine, AsyncEngineContextProvider, ResponseStream};
use dynamo_runtime::pipeline::{
43
    AsyncEngineContext, Error, ManyOut, Operator, SingleIn, async_trait,
Biswa Panda's avatar
Biswa Panda committed
44
};
Neelay Shah's avatar
Neelay Shah committed
45
use dynamo_runtime::protocols::annotated::{Annotated, AnnotationsProvider};
Biswa Panda's avatar
Biswa Panda committed
46
47

use crate::protocols::{
Greg Clark's avatar
Greg Clark committed
48
    common::{OutputOptionsProvider, SamplingOptionsProvider, StopConditionsProvider},
Biswa Panda's avatar
Biswa Panda committed
49
    openai::{
50
        DeltaGeneratorExt,
Ryan Olson's avatar
Ryan Olson committed
51
52
53
        chat_completions::{
            NvCreateChatCompletionRequest, NvCreateChatCompletionStreamResponse, jail::JailedStream,
        },
54
        completions::{NvCreateCompletionRequest, NvCreateCompletionResponse},
55
        embeddings::{NvCreateEmbeddingRequest, NvCreateEmbeddingResponse},
Biswa Panda's avatar
Biswa Panda committed
56
57
58
        nvext::NvExtProvider,
    },
};
59
use crate::tokenizers::{HuggingFaceTokenizer, traits::Tokenizer};
Biswa Panda's avatar
Biswa Panda committed
60

61
use crate::preprocessor::prompt::{PromptFormatter, PromptInput, TextInput, TokenInput};
Biswa Panda's avatar
Biswa Panda committed
62

63
pub use crate::protocols::common::llm_backend::{BackendOutput, PreprocessedRequest};
64
65
66
pub use crate::protocols::common::preprocessor::PreprocessedEmbeddingRequest;

use crate::protocols::common::llm_backend::EmbeddingsEngineOutput;
Biswa Panda's avatar
Biswa Panda committed
67
68
69

pub const ANNOTATION_FORMATTED_PROMPT: &str = "formatted_prompt";
pub const ANNOTATION_TOKEN_IDS: &str = "token_ids";
70
71
72
73
74
75
pub const ANNOTATION_LLM_METRICS: &str = "llm_metrics";
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct LLMMetricAnnotation {
    pub input_tokens: usize,
    pub output_tokens: usize,
    pub chunk_tokens: usize,
76
    pub cached_tokens: Option<usize>,
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
    /// Prefill worker ID (for TTFT attribution in disaggregated mode)
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub prefill_worker_id: Option<u64>,
    /// Prefill worker DP rank
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub prefill_dp_rank: Option<u32>,
    /// Prefill worker type ("prefill" or "decode") for Prometheus metric labeling.
    /// Stored at routing time to avoid expensive MDC lookup when updating TTFT metrics.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub prefill_worker_type: Option<String>,
    /// Decode worker ID (for ITL attribution in disaggregated mode)
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub decode_worker_id: Option<u64>,
    /// Decode worker DP rank
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub decode_dp_rank: Option<u32>,
    /// Decode worker type ("prefill" or "decode") for Prometheus metric labeling.
    /// Stored at routing time to avoid expensive MDC lookup when updating ITL metrics.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub decode_worker_type: Option<String>,
97
    #[serde(default, skip_serializing_if = "Option::is_none")]
98
99
100
101
102
    pub tokenize_latency: Option<Duration>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub detokenize_total_latency: Option<Duration>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub detokenize_count: Option<u64>,
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
}

impl LLMMetricAnnotation {
    /// Convert this metrics struct to an Annotated event
    pub fn to_annotation<T>(&self) -> Result<Annotated<T>, serde_json::Error> {
        Annotated::from_annotation(ANNOTATION_LLM_METRICS, self)
    }

    /// Extract LLM metrics from an Annotated event, if present
    pub fn from_annotation<T>(
        annotation: &Annotated<T>,
    ) -> Result<Option<LLMMetricAnnotation>, Box<dyn std::error::Error>> {
        if annotation.event.is_none() {
            return Ok(None);
        }
        if annotation.event.as_ref().unwrap() != ANNOTATION_LLM_METRICS {
            return Ok(None);
        }
        let comments = annotation
            .comment
            .as_ref()
            .ok_or("missing comments block")?;
        if comments.len() != 1 {
            return Err("malformed comments block - expected exactly 1 comment".into());
        }
        let metrics: LLMMetricAnnotation = serde_json::from_str(&comments[0])?;
        Ok(Some(metrics))
    }
}
Biswa Panda's avatar
Biswa Panda committed
132

133
134
135
136
137
138
// Reasoning State for reasoning parsing transformation step
struct ReasoningState {
    stream: Pin<Box<dyn Stream<Item = Annotated<NvCreateChatCompletionStreamResponse>> + Send>>,
    reasoning_parser: Option<Box<dyn ReasoningParser>>,
}

Biswa Panda's avatar
Biswa Panda committed
139
140
141
142
143
pub struct OpenAIPreprocessor {
    mdcsum: String,
    formatter: Arc<dyn OAIPromptFormatter>,
    tokenizer: Arc<dyn Tokenizer>,
    model_info: Arc<dyn ModelInfo>,
144
    lora_name: Option<String>,
145
146
    /// Per-model runtime configuration propagated to response generator (e.g., reasoning/tool parser)
    runtime_config: crate::local_model::runtime_config::ModelRuntimeConfig,
147
    tool_call_parser: Option<String>,
148
    media_loader: Option<MediaLoader>,
Biswa Panda's avatar
Biswa Panda committed
149
150
151
}

impl OpenAIPreprocessor {
152
153
154
    pub fn new(mdc: ModelDeploymentCard) -> Result<Arc<Self>> {
        let formatter = PromptFormatter::from_mdc(&mdc)?;
        let tokenizer = mdc.tokenizer_hf()?;
155
        match formatter {
156
            PromptFormatter::OAI(formatter) => Self::new_with_parts(mdc, formatter, tokenizer),
157
158
159
        }
    }

160
    pub fn new_with_parts(
161
162
        mdc: ModelDeploymentCard,
        formatter: Arc<dyn OAIPromptFormatter>,
163
        hf_tokenizer: tokenizers::Tokenizer,
164
    ) -> Result<Arc<Self>> {
165
        let mdcsum = mdc.mdcsum().to_string();
166
        let tokenizer = Arc::new(HuggingFaceTokenizer::from_tokenizer(hf_tokenizer));
167
        let lora_name = mdc.lora.as_ref().map(|l| l.name.clone());
168
        let Some(ref model_info) = mdc.model_info else {
169
170
171
172
            anyhow::bail!(
                "Blank ModelDeploymentCard cannot be used for pre-processing, no model_info"
            );
        };
173
        let model_info = model_info.get_model_info()?;
174
        let tool_call_parser = mdc.runtime_config.tool_call_parser.clone();
Biswa Panda's avatar
Biswa Panda committed
175

176
177
178
179
        if let Some(ref lora_name) = lora_name {
            tracing::info!(model = %mdc.display_name, lora_name, "LoRA adapter detected in MDC");
        }

180
181
        // // Initialize runtime config from the ModelDeploymentCard
        let runtime_config = mdc.runtime_config.clone();
182
183
184
185
186
187

        let media_loader = match mdc.media_decoder {
            Some(media_decoder) => Some(MediaLoader::new(media_decoder, mdc.media_fetcher)?),
            None => None,
        };

Biswa Panda's avatar
Biswa Panda committed
188
189
190
191
192
        Ok(Arc::new(Self {
            formatter,
            tokenizer,
            model_info,
            mdcsum,
193
            lora_name,
194
            runtime_config,
195
            tool_call_parser,
196
            media_loader,
Biswa Panda's avatar
Biswa Panda committed
197
198
        }))
    }
199
200
201
202
203
    /// Encode a string to it's tokens
    pub fn tokenize(&self, s: &str) -> anyhow::Result<Encoding> {
        self.tokenizer.encode(s)
    }

204
    /// Translate a [`NvCreateChatCompletionRequest`] request to a common completion request.
Biswa Panda's avatar
Biswa Panda committed
205
206
207
208
209
    /// Returns both the common completion request and a hashmap of annotations.
    ///
    /// Annotations evaluated by this method include:
    /// - `formatted_prompt`
    /// - `token_ids`
210
    pub async fn preprocess_request<
Biswa Panda's avatar
Biswa Panda committed
211
212
213
214
        R: OAIChatLikeRequest
            + AnnotationsProvider
            + SamplingOptionsProvider
            + StopConditionsProvider
Greg Clark's avatar
Greg Clark committed
215
            + OutputOptionsProvider
Biswa Panda's avatar
Biswa Panda committed
216
217
218
219
            + NvExtProvider,
    >(
        &self,
        request: &R,
220
        tracker: Option<&RequestTracker>,
221
    ) -> Result<(PreprocessedRequest, HashMap<String, String>)> {
222
        let mut builder = self.builder(request)?;
223
224
225
226
        let formatted_prompt = self
            .apply_template(request)
            .with_context(|| "Failed to apply prompt template")?;
        let annotations = self
227
            .gather_tokens(request, &mut builder, formatted_prompt, tracker)
228
229
            .with_context(|| "Failed to gather tokens")?;
        self.gather_multi_modal_data(request, &mut builder)
230
            .await
231
            .with_context(|| "Failed to gather multimodal data")?;
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246

        Ok((builder.build()?, annotations))
    }

    pub fn builder<
        R: OAIChatLikeRequest
            + AnnotationsProvider
            + SamplingOptionsProvider
            + StopConditionsProvider
            + OutputOptionsProvider
            + NvExtProvider,
    >(
        &self,
        request: &R,
    ) -> Result<PreprocessedRequestBuilder> {
247
        let mut builder = PreprocessedRequest::builder();
248
        builder.model(request.model());
Biswa Panda's avatar
Biswa Panda committed
249

250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
        let mut stop_conditions = request.extract_stop_conditions()?;
        if let Some(stop_tokens) = &mut stop_conditions.stop_token_ids_hidden {
            for eos_token in self.model_info.eos_token_ids() {
                if !stop_tokens.contains(&eos_token) {
                    stop_tokens.push(eos_token);
                }
            }
        } else {
            stop_conditions.stop_token_ids_hidden = Some(self.model_info.eos_token_ids());
        }

        // apply ignore eos if not already set
        stop_conditions.apply_ignore_eos();

        if !stop_conditions.ignore_eos.unwrap_or(false) {
            builder.eos_token_ids(self.model_info.eos_token_ids());
        }

        builder.stop_conditions(stop_conditions);
        builder.sampling_options(request.extract_sampling_options()?);
        builder.output_options(request.extract_output_options()?);
        builder.annotations(request.annotations().unwrap_or_default());
        builder.mdc_sum(Some(self.mdcsum.clone()));
273
274
        let lora_name = self.lora_name.clone();

275
        // Extract routing hints from nvext if present
276
        if let Some(nvext) = request.nvext() {
277
            // Build routing hints from nvext fields
278
            let hints = nvext.agent_hints.as_ref();
279
280
281
282
283
            let routing = RoutingHints {
                backend_instance_id: nvext.backend_instance_id,
                prefill_worker_id: nvext.prefill_worker_id,
                decode_worker_id: nvext.decode_worker_id,
                dp_rank: None, // dp_rank is set later in the pipeline
284
285
                expected_output_tokens: hints.and_then(|h| h.osl),
                priority_jump: hints.and_then(|h| h.latency_sensitivity),
286
                lora_name,
287
288
            };
            builder.routing(Some(routing));
289
290
291
292
293
294
        } else if lora_name.is_some() {
            // Ensure LoRA-aware routing still gets hints even when nvext is absent.
            builder.routing(Some(RoutingHints {
                lora_name,
                ..Default::default()
            }));
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
        }

        Ok(builder)
    }

    pub fn apply_template<
        R: OAIChatLikeRequest
            + AnnotationsProvider
            + SamplingOptionsProvider
            + StopConditionsProvider
            + OutputOptionsProvider
            + NvExtProvider,
    >(
        &self,
        request: &R,
    ) -> Result<Option<String>> {
        if let PromptInput::Text(_) = request.prompt_input_type()
            && let Some(TextInput::Single(_)) = request.extract_text()
        {
            let use_raw_prompt = request
                .nvext()
                .is_some_and(|ext| ext.use_raw_prompt.unwrap_or(false));

            let formatted_prompt = if use_raw_prompt {
                match request.raw_prompt() {
                    Some(prompt) => prompt,
                    None => {
                        tracing::warn!("Raw prompt requested but not available");
                        self.formatter.render(request)?
                    }
                }
            } else {
                self.formatter.render(request)?
            };
            Ok(Some(formatted_prompt))
        } else {
            Ok(None)
        }
    }

335
    pub async fn gather_multi_modal_data<R: OAIChatLikeRequest>(
336
337
338
339
340
        &self,
        request: &R,
        builder: &mut PreprocessedRequestBuilder,
    ) -> Result<()> {
        let mut media_map: MultimodalDataMap = HashMap::new();
341
342
        let mut fetch_tasks: Vec<(String, ChatCompletionRequestUserMessageContentPart)> =
            Vec::new();
343

344
345
346
347
348
        let Some(messages) = request.typed_messages() else {
            return Ok(());
        };
        for message in messages.iter() {
            let content_parts = match message {
349
350
351
352
353
354
355
                ChatCompletionRequestMessage::User(u) => match &u.content {
                    ChatCompletionRequestUserMessageContent::Array(parts) => parts,
                    _ => continue,
                },
                _ => continue,
            };
            // Iterate over content parts
356
            for content_part in content_parts.iter() {
357
358
359
360
361
362
363
364
365
366
367
368
369
                let (type_str, url) = match content_part {
                    ChatCompletionRequestUserMessageContentPart::ImageUrl(image_part) => {
                        ("image_url".to_string(), image_part.image_url.url.clone())
                    }
                    ChatCompletionRequestUserMessageContentPart::VideoUrl(video_part) => {
                        ("video_url".to_string(), video_part.video_url.url.clone())
                    }
                    ChatCompletionRequestUserMessageContentPart::AudioUrl(audio_part) => {
                        ("audio_url".to_string(), audio_part.audio_url.url.clone())
                    }
                    _ => continue,
                };

370
371
                if self.media_loader.is_some() {
                    fetch_tasks.push((type_str, content_part.clone()));
372
                    continue;
373
                }
374
375
376
377
378
379

                //Fallback: ust pass the URL through
                media_map
                    .entry(type_str)
                    .or_default()
                    .push(MultimodalData::Url(url));
380
381
            }
        }
382
383
384
385

        // Execute all fetch tasks
        if !fetch_tasks.is_empty() {
            let loader = self.media_loader.as_ref().unwrap();
386
387
388
389
            let media_io_kwargs = request.media_io_kwargs();
            let results = futures::future::join_all(fetch_tasks.iter().map(|(_, content_part)| {
                loader.fetch_and_decode_media_part(content_part, media_io_kwargs)
            }))
390
391
            .await;

392
393
394
395
396
397
398
399
            for ((type_str, _), result) in fetch_tasks.into_iter().zip(results.into_iter()) {
                // if one item fails, errors the whole request, other items will be cleaned up by Drop
                let rdma_descriptor = result?;
                media_map
                    .entry(type_str)
                    .or_default()
                    .push(MultimodalData::Decoded(rdma_descriptor));
            }
400
401
        }

402
403
        if !media_map.is_empty() {
            builder.multi_modal_data(Some(media_map));
404
405
406

            // Preserve original messages in extra_args for multimodal workers that need them
            // (e.g., TRT-LLM multimodal processor needs raw messages for proper tokenization)
407
            let messages_json = serde_json::to_value(request.messages())?;
408
409
410
411
            let extra_args = serde_json::json!({
                "messages": messages_json
            });
            builder.extra_args(Some(extra_args));
412
413
414
415
416
        }

        Ok(())
    }

417
418
419
420
421
422
423
424
425
426
427
428
    pub fn gather_tokens<
        R: OAIChatLikeRequest
            + AnnotationsProvider
            + SamplingOptionsProvider
            + StopConditionsProvider
            + OutputOptionsProvider
            + NvExtProvider,
    >(
        &self,
        request: &R,
        builder: &mut PreprocessedRequestBuilder,
        formatted_prompt: Option<String>,
429
        tracker: Option<&RequestTracker>,
430
431
    ) -> Result<HashMap<String, String>> {
        let mut annotations = HashMap::new();
432
433
434
435
436
437
438
439
440
441
442
443
        // match request type before any conversion/processing
        match request.prompt_input_type() {
            PromptInput::Tokens(_) => {
                if let Some(token_input) = request.extract_tokens() {
                    match token_input {
                        TokenInput::Single(tokens) => {
                            builder.token_ids(tokens);
                        }
                        TokenInput::Batch(token_batches) => {
                            if token_batches.len() == 1 {
                                builder.token_ids(token_batches[0].clone());
                            } else {
444
445
446
447
                                bail!(
                                    "Batch token input not supported for more than one token in requests (got {})",
                                    token_batches.len()
                                );
448
449
450
                            }
                        }
                    }
Biswa Panda's avatar
Biswa Panda committed
451
452
                }
            }
453
454
455
            PromptInput::Text(_) => {
                if let Some(text_input) = request.extract_text() {
                    match text_input {
456
457
458
459
460
461
462
463
464
465
                        TextInput::Single(raw_prompt) => {
                            if let Some(f) = formatted_prompt.as_ref()
                                && request.has_annotation(ANNOTATION_FORMATTED_PROMPT)
                            {
                                annotations
                                    .insert(ANNOTATION_FORMATTED_PROMPT.to_string(), f.to_string());
                            }

                            // Completions will use raw_prompt, no template
                            let prompt = formatted_prompt.unwrap_or(raw_prompt);
Biswa Panda's avatar
Biswa Panda committed
466

467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
                            // Check if backend_instance_id is present and token_data is provided
                            let has_backend_instance_id = request
                                .nvext()
                                .and_then(|ext| ext.backend_instance_id)
                                .is_some();

                            let token_data =
                                request.nvext().and_then(|ext| ext.token_data.as_ref());

                            let (tokens_vec, skip_token_annotation) = if has_backend_instance_id {
                                if let Some(tokens) = token_data {
                                    tracing::trace!(
                                        "Using provided tokens from EPP: {} ids",
                                        tokens.len()
                                    );
                                    // need ownership for the builder, so clone.
                                    (tokens.clone(), true)
                                } else {
                                    tracing::warn!(
                                        "backend_instance_id provided but no token_data; tokenizing prompt"
                                    );
488
                                    let encoding = self.encode_with_timing(&prompt, tracker)?;
489
490
491
492
                                    (encoding.token_ids().to_vec(), false)
                                }
                            } else {
                                // No backend_instance_id provided, continue the normal flow.
493
                                let encoding = self.encode_with_timing(&prompt, tracker)?;
494
495
                                (encoding.token_ids().to_vec(), false)
                            };
Biswa Panda's avatar
Biswa Panda committed
496

497
498
499
                            if request.has_annotation(ANNOTATION_TOKEN_IDS)
                                && !skip_token_annotation
                            {
500
501
                                annotations.insert(
                                    ANNOTATION_TOKEN_IDS.to_string(),
502
                                    serde_json::to_string(&tokens_vec)?,
503
504
505
                                );
                            }

506
                            builder.token_ids(tokens_vec);
507
508
                        }
                        TextInput::Batch(texts) => {
509
                            if texts.len() == 1 {
510
                                let encoding = self.encode_with_timing(&texts[0], tracker)?;
511
512
513
514
515
516
517
                                builder.token_ids(encoding.token_ids().to_vec());
                            } else {
                                bail!(
                                    "Batch text input not supported for more than one text in requests (got {})",
                                    texts.len()
                                );
                            }
518
519
520
521
                        }
                    }
                }
            }
Biswa Panda's avatar
Biswa Panda committed
522
        }
523
        Ok(annotations)
Biswa Panda's avatar
Biswa Panda committed
524
525
    }

526
527
528
529
530
531
532
533
    fn encode_with_timing(
        &self,
        prompt: &str,
        tracker: Option<&RequestTracker>,
    ) -> anyhow::Result<Encoding> {
        let encode_start = Instant::now();
        let encoding = self.tokenizer.encode(prompt)?;
        if let Some(t) = tracker {
534
            t.record_tokenize_latency(encode_start.elapsed());
535
536
537
538
        }
        Ok(encoding)
    }

539
540
541
542
543
544
545
546
547
548
549
550
551
552
    /// Preprocess an embedding request, handling both text and token ID inputs.
    ///
    /// For text inputs, tokenizes the text using the configured tokenizer.
    /// For token ID inputs, uses the provided token IDs directly and skips tokenization.
    ///
    /// Returns both the preprocessed request and a hashmap of annotations.
    pub async fn preprocess_embedding_request(
        &self,
        request: &NvCreateEmbeddingRequest,
    ) -> Result<(PreprocessedEmbeddingRequest, HashMap<String, String>)> {
        let mut annotations = HashMap::new();
        let mut builder = PreprocessedEmbeddingRequest::builder();

        let all_token_ids = match &request.inner.input {
553
            dynamo_async_openai::types::EmbeddingInput::String(s) => {
554
555
                let encoding = self.tokenizer.encode(s)?;
                vec![encoding.token_ids().to_vec()]
556
            }
557
            dynamo_async_openai::types::EmbeddingInput::StringArray(arr) => {
558
559
560
561
562
563
564
565
566
567
568
                let input_strs: Vec<String> = arr.to_vec();
                let encodings = tokio::task::spawn_blocking({
                    let tokenizer = self.tokenizer.clone();
                    let strs = input_strs.clone();
                    move || {
                        tokenizer.encode_batch(&strs.iter().map(|s| s.as_str()).collect::<Vec<_>>())
                    }
                })
                .await??;
                let token_arrays: Vec<Vec<u32>> = encodings
                    .into_iter()
569
                    .map(|encoding| encoding.token_ids().to_vec())
570
571
572
                    .collect();
                token_arrays
            }
573
574
575
576
            dynamo_async_openai::types::EmbeddingInput::IntegerArray(token_ids) => {
                vec![token_ids.clone()]
            }
            dynamo_async_openai::types::EmbeddingInput::ArrayOfIntegerArray(token_arrays) => {
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
                token_arrays.clone()
            }
        };

        // Handle annotations
        if request.has_annotation(ANNOTATION_TOKEN_IDS) {
            annotations.insert(
                ANNOTATION_TOKEN_IDS.to_string(),
                serde_json::to_string(&all_token_ids)?,
            );
        }

        builder.token_ids(all_token_ids);
        builder.model(request.inner.model.clone());
        builder.encoding_format(request.inner.encoding_format.as_ref().map(|f| match f {
            EncodingFormat::Float => "float".to_string(),
            EncodingFormat::Base64 => "base64".to_string(),
        }));
        builder.dimensions(request.inner.dimensions);

        builder.annotations(request.annotations().unwrap_or_default());
        builder.mdc_sum(Some(self.mdcsum.clone()));

        Ok((builder.build()?, annotations))
    }

Ryan Olson's avatar
Ryan Olson committed
603
604
    pub fn transform_postprocessor_stream<S, Resp>(
        stream: S,
Biswa Panda's avatar
Biswa Panda committed
605
        generator: Box<dyn DeltaGeneratorExt<Resp>>,
Ryan Olson's avatar
Ryan Olson committed
606
607
608
609
610
611
612
613
614
615
616
        context: Arc<dyn AsyncEngineContext>,
    ) -> impl Stream<Item = Annotated<Resp>> + Send
    where
        S: Stream<Item = Annotated<BackendOutput>> + Send + 'static,
        Resp: Send + Sync + 'static + std::fmt::Debug,
    {
        struct State<Resp>
        where
            Resp: Send + Sync + 'static + std::fmt::Debug,
        {
            response_stream: Pin<Box<dyn Stream<Item = Annotated<BackendOutput>> + Send>>,
Biswa Panda's avatar
Biswa Panda committed
617
618
619
            response_generator: Box<dyn DeltaGeneratorExt<Resp>>,
            context: Arc<dyn AsyncEngineContext>,
            cancelled: bool,
620
            cumulative_output_tokens: usize,
621
622
            finish_reason_sent: bool,
            usage_chunk_sent: bool,
Ryan Olson's avatar
Ryan Olson committed
623
            finished: bool,
Biswa Panda's avatar
Biswa Panda committed
624
625
626
        }

        let state = State {
Ryan Olson's avatar
Ryan Olson committed
627
            response_stream: Box::pin(stream),
Biswa Panda's avatar
Biswa Panda committed
628
629
630
            response_generator: generator,
            context: context.clone(),
            cancelled: false,
631
            cumulative_output_tokens: 0,
632
633
            finish_reason_sent: false,
            usage_chunk_sent: false,
Ryan Olson's avatar
Ryan Olson committed
634
            finished: false,
Biswa Panda's avatar
Biswa Panda committed
635
636
637
        };

        // transform the common response stream into a chat response stream
Ryan Olson's avatar
Ryan Olson committed
638
639

        stream::unfold(state, |mut inner| {
Biswa Panda's avatar
Biswa Panda committed
640
            async move {
641
642
643
644
645
                // If already finished, return None immediately
                if inner.finished {
                    return None;
                }

Biswa Panda's avatar
Biswa Panda committed
646
647
648
649
650
651
                if let Some(response) = inner.response_stream.next().await {
                    if inner.cancelled {
                        tracing::debug!(
                            request_id = inner.context.id(),
                            "Cancellation issued last message; closing stream"
                        );
652
                        inner.finished = true; // Mark as finished
Biswa Panda's avatar
Biswa Panda committed
653
654
655
656
657
658
659
660
661
                        return None;
                    }

                    tracing::trace!(
                        request_id = inner.context.id(),
                        "Processing common response: {:?}",
                        response
                    );

662
663
664
665
666
667
668
                    // Check if this response has a finish_reason
                    let has_finish_reason = response
                        .data
                        .as_ref()
                        .map(|d| d.finish_reason.is_some())
                        .unwrap_or(false);

669
670
671
672
673
674
675
676
677
678
679
680
681
682
                    let (chunk_tokens, isl) = if let Some(ref backend_output) = response.data {
                        let chunk_tokens = backend_output.token_ids.len();
                        inner.cumulative_output_tokens += chunk_tokens;

                        let isl = inner.response_generator.get_isl().unwrap_or(0) as usize;

                        (chunk_tokens, isl)
                    } else {
                        (0, 0)
                    };

                    let current_osl = inner.cumulative_output_tokens;

                    let mut response = response.map_data(|data| {
Biswa Panda's avatar
Biswa Panda committed
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
                        inner
                            .response_generator
                            .choice_from_postprocessor(data)
                            .inspect_err(|e| {
                                tracing::error!(
                                    request_id = inner.context.id(),
                                    "Error processing common response: {:?}",
                                    e
                                );
                                inner.cancelled = true;
                                inner.context.stop_generating();
                            })
                            .map_err(|e| e.to_string())
                    });

698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
                    // Create LLM metrics annotation with prefill/decode worker info from tracker.
                    // Worker types are stored at routing time to avoid expensive MDC lookup.
                    let tracker = inner.response_generator.tracker();
                    let prefill_worker_id = tracker.as_ref().and_then(|t| t.prefill_worker_id());
                    let prefill_dp_rank = tracker.as_ref().and_then(|t| t.prefill_dp_rank());
                    let prefill_worker_type = tracker
                        .as_ref()
                        .and_then(|t| t.prefill_worker_type())
                        .map(String::from);
                    let decode_worker_id = tracker.as_ref().and_then(|t| t.decode_worker_id());
                    let decode_dp_rank = tracker.as_ref().and_then(|t| t.decode_dp_rank());
                    let decode_worker_type = tracker
                        .as_ref()
                        .and_then(|t| t.decode_worker_type())
                        .map(String::from);
713
714
715
716
                    let llm_metrics = LLMMetricAnnotation {
                        input_tokens: isl,
                        output_tokens: current_osl,
                        chunk_tokens,
717
                        cached_tokens: None,
718
719
720
721
722
723
                        prefill_worker_id,
                        prefill_dp_rank,
                        prefill_worker_type,
                        decode_worker_id,
                        decode_dp_rank,
                        decode_worker_type,
724
725
726
                        tokenize_latency: tracker.as_ref().and_then(|t| t.tokenize_latency()),
                        detokenize_total_latency: tracker.as_ref().and_then(|t| t.detokenize_total_latency()),
                        detokenize_count: tracker.as_ref().map(|t| t.detokenize_count()),
727
728
729
730
731
732
                    };

                    if let Ok(metrics_annotated) = llm_metrics.to_annotation::<()>() {
                        // Only set event if not already set to avoid overriding existing events (like errors)
                        if response.event.is_none() {
                            response.event = metrics_annotated.event;
733
                            response.comment = metrics_annotated.comment;
734
735
                        }
                    }
736

737
738
739
740
741
                    // Mark if we've seen a finish_reason
                    if has_finish_reason {
                        inner.finish_reason_sent = true;
                    }

Biswa Panda's avatar
Biswa Panda committed
742
743
                    tracing::trace!(
                        request_id = inner.context.id(),
744
                        "OpenAI NvCreateChatCompletionStreamResponse: {:?}",
Biswa Panda's avatar
Biswa Panda committed
745
746
747
748
749
                        response
                    );

                    Some((response, inner))
                } else {
750
751
752
753
                    // Stream has ended - must set finished to true to prevent unfold from polling
                    // again. The stream is exhausted and will panic if polled after None.
                    inner.finished = true;

754
                    if inner.finish_reason_sent && !inner.usage_chunk_sent {
755
756
757
                        inner.usage_chunk_sent = true;

                        let usage_chunk = inner.response_generator.create_usage_chunk();
758
                        let usage = inner.response_generator.get_usage();
759
760
761
762
763
764
765
766
767
768
769
770
771
772
                        let tracker = inner.response_generator.tracker();
                        let prefill_worker_id =
                            tracker.as_ref().and_then(|t| t.prefill_worker_id());
                        let prefill_dp_rank = tracker.as_ref().and_then(|t| t.prefill_dp_rank());
                        let prefill_worker_type = tracker
                            .as_ref()
                            .and_then(|t| t.prefill_worker_type())
                            .map(String::from);
                        let decode_worker_id = tracker.as_ref().and_then(|t| t.decode_worker_id());
                        let decode_dp_rank = tracker.as_ref().and_then(|t| t.decode_dp_rank());
                        let decode_worker_type = tracker
                            .as_ref()
                            .and_then(|t| t.decode_worker_type())
                            .map(String::from);
773
774
775
776
777
778
779
780
                        let llm_metrics = LLMMetricAnnotation {
                            input_tokens: usage.prompt_tokens as usize,
                            output_tokens: usage.completion_tokens as usize,
                            chunk_tokens: 0,
                            cached_tokens: usage
                                .prompt_tokens_details
                                .as_ref()
                                .and_then(|d| d.cached_tokens.map(|c| c as usize)),
781
782
783
784
785
786
                            prefill_worker_id,
                            prefill_dp_rank,
                            prefill_worker_type,
                            decode_worker_id,
                            decode_dp_rank,
                            decode_worker_type,
787
788
789
790
791
                            tokenize_latency: tracker.as_ref().and_then(|t| t.tokenize_latency()),
                            detokenize_total_latency: tracker
                                .as_ref()
                                .and_then(|t| t.detokenize_total_latency()),
                            detokenize_count: tracker.as_ref().map(|t| t.detokenize_count()),
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
                        };

                        // Create annotation string
                        let annotation = llm_metrics.to_annotation::<()>().unwrap_or_else(|e| {
                            tracing::warn!("Failed to serialize metrics: {}", e);
                            Annotated::<()>::from_data(())
                        });

                        // Send the usage chunk if needed
                        let data = if inner.response_generator.is_usage_enabled() {
                            Some(usage_chunk)
                        } else {
                            None
                        };

807
808
                        let annotated_usage = Annotated::<Resp> {
                            id: None,
809
810
811
                            data,
                            event: Some(ANNOTATION_LLM_METRICS.to_string()),
                            comment: annotation.comment,
812
813
814
815
                        };

                        tracing::trace!(
                            request_id = inner.context.id(),
816
817
                            "Sending final usage chunk for OpenAI compliance, annotated_usage: {:?}",
                            annotated_usage
818
819
820
821
822
823
824
                        );

                        Some((annotated_usage, inner))
                    } else {
                        // stream closed
                        None
                    }
Biswa Panda's avatar
Biswa Panda committed
825
826
                }
            }
Ryan Olson's avatar
Ryan Olson committed
827
        })
828
        .fuse()
Biswa Panda's avatar
Biswa Panda committed
829
    }
830
831

    /// Transform engine embedding output stream to OpenAI embedding response stream
Ryan Olson's avatar
Ryan Olson committed
832
833
    pub fn transform_embedding_postprocessor_stream<S>(
        stream: S,
834
        original_request: NvCreateEmbeddingRequest,
Ryan Olson's avatar
Ryan Olson committed
835
836
837
838
839
    ) -> impl Stream<Item = Annotated<NvCreateEmbeddingResponse>> + Send
    where
        S: Stream<Item = Annotated<EmbeddingsEngineOutput>> + Send + 'static,
    {
        stream.map(move |output| {
840
841
            output.map_data(|engine_output| {
                // Convert engine output to OpenAI response format
842
                let embeddings: Vec<dynamo_async_openai::types::Embedding> = engine_output
843
844
845
                    .embeddings
                    .into_iter()
                    .enumerate()
846
                    .map(|(index, embedding)| dynamo_async_openai::types::Embedding {
847
848
849
850
851
852
853
                        index: index as u32,
                        object: "embedding".to_string(),
                        embedding: embedding.into_iter().map(|f| f as f32).collect(),
                    })
                    .collect();

                let response = NvCreateEmbeddingResponse {
854
                    inner: dynamo_async_openai::types::CreateEmbeddingResponse {
855
856
857
                        object: "list".to_string(),
                        model: original_request.inner.model.clone(),
                        data: embeddings,
858
                        usage: dynamo_async_openai::types::EmbeddingUsage {
859
860
861
862
863
864
865
866
                            prompt_tokens: engine_output.prompt_tokens,
                            total_tokens: engine_output.total_tokens,
                        },
                    },
                };

                Ok(response)
            })
Ryan Olson's avatar
Ryan Olson committed
867
        })
868
869
    }

Ryan Olson's avatar
Ryan Olson committed
870
871
872
873
874
875
876
877
    /// Determine if we should apply the tool calling jail based on configuration
    /// Returns Ok(true) if jail should be applied, Ok(false) if not, or Err if invalid config
    pub fn should_apply_tool_jail(
        tool_call_parser: Option<&String>,
        tool_choice: Option<&ChatCompletionToolChoiceOption>,
        has_tools: bool,
    ) -> std::result::Result<bool, Error> {
        match (tool_call_parser, tool_choice, has_tools) {
878
879
880
881
882
            // tool_choice=required/named work without parser (use Immediate jail mode)
            (None, Some(ChatCompletionToolChoiceOption::Required), true) => Ok(true),
            (None, Some(ChatCompletionToolChoiceOption::Named(_)), true) => Ok(true),

            // tool_choice=auto requires a parser
Ryan Olson's avatar
Ryan Olson committed
883
884
885
886
887
888
            (None, Some(ChatCompletionToolChoiceOption::Auto), true) => {
                tracing::warn!(
                    "Tool choice 'auto' specified but no tool parser configured; proceeding without jailing"
                );
                Ok(false)
            }
889

Ryan Olson's avatar
Ryan Olson committed
890
891
892
            // Parser exists and tools might be called
            (Some(_), Some(ChatCompletionToolChoiceOption::None), _) => {
                Ok(false) // Explicitly disabled
893
            }
Ryan Olson's avatar
Ryan Olson committed
894
895
            (Some(_), Some(_), true) => Ok(true), // Any other tool_choice with tools
            (Some(_), None, true) => Ok(true),    // Default behavior when tools present
896

Ryan Olson's avatar
Ryan Olson committed
897
898
899
900
            // No tools or no parser
            _ => Ok(false),
        }
    }
901

Ryan Olson's avatar
Ryan Olson committed
902
903
    /// Apply tool calling jail to the stream if needed
    pub fn apply_tool_calling_jail<S>(
904
905
        tool_call_parser: Option<String>,
        tool_choice: Option<dynamo_async_openai::types::ChatCompletionToolChoiceOption>,
906
        tool_definitions: Option<Vec<dynamo_parsers::tool_calling::ToolDefinition>>,
Ryan Olson's avatar
Ryan Olson committed
907
908
909
910
911
        stream: S,
    ) -> impl Stream<Item = Annotated<NvCreateChatCompletionStreamResponse>> + Send
    where
        S: Stream<Item = Annotated<NvCreateChatCompletionStreamResponse>> + Send + 'static,
    {
912
913
914
915
        use dynamo_async_openai::types::ChatCompletionToolChoiceOption;

        let mut builder = JailedStream::builder();

916
917
918
919
920
921
922
        // Set tool definitions if provided
        if let Some(tool_definitions) = tool_definitions
            && !tool_definitions.is_empty()
        {
            builder = builder.tool_definitions(tool_definitions);
        }

923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
        // Configure jail based on tool_choice
        match tool_choice {
            Some(ChatCompletionToolChoiceOption::Named(named)) => {
                // Immediate jail mode for named tool choice
                builder = builder.tool_choice_named(named.function.name.clone());
            }
            Some(ChatCompletionToolChoiceOption::Required) => {
                // Immediate jail mode for required tool choice
                builder = builder.tool_choice_required();
            }
            Some(ChatCompletionToolChoiceOption::Auto)
            | Some(ChatCompletionToolChoiceOption::None)
            | None => {
                // Traditional marker-based jail for auto/none/unspecified
                if let Some(parser) = tool_call_parser {
                    builder = builder.tool_call_parser(parser);
                }
            }
        }

        let jail = builder.build();
944
        jail.apply_with_finish_reason(stream)
Ryan Olson's avatar
Ryan Olson committed
945
    }
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973

    // Motivation: Each transformation on the stream should be a separate step to allow for more flexibility
    // Earlier reasoning parser logic was nested under delta generation logic in choice_from_postprocessor
    // Since we have tool calling parsing as separate step, it makes sense to have reasoning parser as separate step as well
    pub fn parse_reasoning_content_from_stream<S>(
        stream: S,
        parser_name: String,
    ) -> impl Stream<Item = Annotated<NvCreateChatCompletionStreamResponse>> + Send
    where
        S: Stream<Item = Annotated<NvCreateChatCompletionStreamResponse>> + Send + 'static,
    {
        // Initialize reasoning parser from parser_name
        let reasoning_parser = Box::new(ReasoningParserType::get_reasoning_parser_from_name(
            parser_name.as_ref(),
        )) as Box<dyn ReasoningParser>;

        let state = ReasoningState {
            stream: Box::pin(stream),
            reasoning_parser: Some(reasoning_parser),
        };

        stream::unfold(state, |mut state| async move {
            if let Some(response) = state.stream.next().await {
                // Process the response through reasoning parser if available
                let processed_response = if let Some(ref mut parser) = state.reasoning_parser {
                    response.map_data(|mut data| {
                        // Process all choices, not just the first one
                        for choice in data.choices.iter_mut() {
974
975
976
977
978
979
980
                            // Reasoning parsing only applies to text content
                            if let Some(
                                dynamo_async_openai::types::ChatCompletionMessageContent::Text(
                                    text,
                                ),
                            ) = choice.delta.content.as_ref()
                            {
981
982
983
984
                                let parser_result =
                                    parser.parse_reasoning_streaming_incremental(text, &[]);

                                // Update this specific choice with parsed content
985
986
987
                                choice.delta.content = parser_result.get_some_normal_text().map(
                                    dynamo_async_openai::types::ChatCompletionMessageContent::Text,
                                );
988
989
                                choice.delta.reasoning_content = parser_result.get_some_reasoning();
                            }
990
                            // For multimodal content, pass through unchanged
991
992
993
994
995
996
997
998
999
1000
1001
1002
1003
                        }
                        Ok(data)
                    })
                } else {
                    // No reasoning parser configured, pass through unchanged
                    response
                };

                Some((processed_response, state))
            } else {
                None
            }
        })
1004
        .fuse()
1005
    }
Biswa Panda's avatar
Biswa Panda committed
1006
1007
1008
1009
1010
1011
1012
1013
1014
1015
}

// for pals, we do not want to add the generation prompt to the formatted prompt
// we also need to know if the template support this add_generation_prompt bool
// any prompt template that does not support this should return an error
// oob - we should update any prompt template that does not support this to support it

#[async_trait]
impl
    Operator<
1016
        SingleIn<NvCreateChatCompletionRequest>,
1017
        ManyOut<Annotated<NvCreateChatCompletionStreamResponse>>,
1018
        SingleIn<PreprocessedRequest>,
Biswa Panda's avatar
Biswa Panda committed
1019
1020
1021
1022
1023
        ManyOut<Annotated<BackendOutput>>,
    > for OpenAIPreprocessor
{
    async fn generate(
        &self,
1024
        request: SingleIn<NvCreateChatCompletionRequest>,
Biswa Panda's avatar
Biswa Panda committed
1025
        next: Arc<
1026
            dyn AsyncEngine<SingleIn<PreprocessedRequest>, ManyOut<Annotated<BackendOutput>>, Error>,
Biswa Panda's avatar
Biswa Panda committed
1027
        >,
1028
    ) -> Result<ManyOut<Annotated<NvCreateChatCompletionStreamResponse>>, Error> {
Biswa Panda's avatar
Biswa Panda committed
1029
        // unpack the request
1030
1031
1032
1033
        let (mut request, context) = request.into_parts();

        // Preserve original inbound streaming flag before any internal overrides
        let request_id = context.id().to_string();
1034
        let original_stream_flag = request.inner.stream.unwrap_or(false);
1035

1036
        // Build audit handle (None if no DYN_AUDIT_SINKS)
1037
1038
1039
1040
1041
1042
        let mut audit_handle = crate::audit::handle::create_handle(&request, &request_id);

        if let Some(ref mut h) = audit_handle {
            h.set_request(std::sync::Arc::new(request.clone()));
        }

1043
1044
1045
1046
1047
        // For non-streaming requests (stream=false), enable usage by default
        // This ensures compliance with OpenAI API spec where non-streaming responses
        // always include usage statistics
        request.enable_usage_for_nonstreaming(original_stream_flag);

1048
1049
        // Set stream=true for internal processing (after audit capture)
        request.inner.stream = Some(true);
Biswa Panda's avatar
Biswa Panda committed
1050
1051

        // create a response generator
1052
        let response_generator = request.response_generator(context.id().to_string());
1053
        let tracker = response_generator.tracker();
Ryan Olson's avatar
Ryan Olson committed
1054
1055

        // convert the chat completion request to a common completion request
1056
1057
1058
        let (mut common_request, annotations) = self
            .preprocess_request(&request, tracker.as_deref())
            .await?;
1059
        tracing::trace!(request = ?common_request, "Pre-processed request");
1060
1061

        // Attach the timing tracker to the request so downstream components can record metrics
1062
        common_request.tracker = tracker;
Ryan Olson's avatar
Ryan Olson committed
1063

Biswa Panda's avatar
Biswa Panda committed
1064
1065
        let mut response_generator = Box::new(response_generator);

1066
1067
1068
1069
1070
        // Update ISL only for text prompts (embeddings get sequence length from tensor shape)
        if common_request.prompt_embeds.is_none() {
            let isl = common_request.token_ids.len() as u32;
            response_generator.update_isl(isl);
        }
Biswa Panda's avatar
Biswa Panda committed
1071
1072
1073
1074
1075

        // repack the common completion request
        let common_request = context.map(|_| common_request);

        // create a stream of annotations this will be prepend to the response stream
1076
        let annotations: Vec<Annotated<NvCreateChatCompletionStreamResponse>> = annotations
Biswa Panda's avatar
Biswa Panda committed
1077
1078
1079
1080
1081
1082
1083
            .into_iter()
            .flat_map(|(k, v)| Annotated::from_annotation(k, &v))
            .collect();
        let annotations_stream = stream::iter(annotations);

        // forward the common completion request to the next operator
        let response_stream = next.generate(common_request).await?;
Ryan Olson's avatar
Ryan Olson committed
1084
1085
1086
1087
1088
1089
1090
1091
1092
1093
        // Extract context once
        let context = response_stream.context();

        // transform the postprocessor stream (no boxing yet)
        let stream = Self::transform_postprocessor_stream(
            response_stream,
            response_generator,
            context.clone(),
        );

1094
1095
1096
1097
1098
1099
1100
1101
1102
1103
1104
1105
1106
1107
1108
1109
1110
1111
1112
        // Try to parse reasoning content only if parser is configured
        let should_parse_reasoning = self.runtime_config.reasoning_parser.is_some();

        // Reasoning Content Parsing Transformation Step
        // Current Solution:
        // This step operates on Deltas created by the transform_postprocessor_stream function
        // Only access to text and not token_ids - so can not support parsing based on token_ids for now
        // Future Solution:
        // To address the limitation if needed in future: move this step before transform_postprocessor_stream and add new field of reasoning_content to the backend output
        // Use backend_output.reasoning_content field to fill out the deltas.
        let stream: Pin<Box<dyn Stream<Item = _> + Send>> = if should_parse_reasoning {
            Box::pin(Self::parse_reasoning_content_from_stream(
                stream,
                self.runtime_config.reasoning_parser.clone().unwrap(), // Safety: We already checked that parser is some, so gtg
            ))
        } else {
            Box::pin(stream)
        };

Ryan Olson's avatar
Ryan Olson committed
1113
1114
1115
1116
1117
1118
1119
1120
1121
1122
1123
        // Check if tools are present and if we should apply jail
        let has_tools =
            request.inner.tools.is_some() && !request.inner.tools.as_ref().unwrap().is_empty();

        // Determine if we should apply jail (do this before moving request)
        let should_jail = Self::should_apply_tool_jail(
            self.tool_call_parser.as_ref(),
            request.inner.tool_choice.as_ref(),
            has_tools,
        )?;

1124
1125
1126
1127
1128
1129
1130
1131
1132
1133
1134
        // Convert OpenAI tools to parser ToolDefinition format before applying jail
        let tool_definitions = request.inner.tools.as_ref().map(|tools| {
            tools
                .iter()
                .map(|tool| dynamo_parsers::tool_calling::ToolDefinition {
                    name: tool.function.name.clone(),
                    parameters: tool.function.parameters.clone(),
                })
                .collect()
        });

Ryan Olson's avatar
Ryan Olson committed
1135
        // Apply jail conditionally
1136
        let transformed_stream: Pin<Box<dyn Stream<Item = _> + Send>> = if should_jail {
1137
1138
1139
            Box::pin(Self::apply_tool_calling_jail(
                self.tool_call_parser.clone(),
                request.inner.tool_choice.clone(),
1140
                tool_definitions,
1141
1142
                stream,
            ))
1143
        } else {
Ryan Olson's avatar
Ryan Olson committed
1144
            Box::pin(stream)
1145
        };
1146
1147
1148
1149
1150
1151
1152
1153
1154
1155
1156
1157
1158
1159
1160
1161
1162
1163
1164
1165
1166
1167
1168

        // Step 4: Apply audit aggregation strategy
        let final_stream = if let Some(mut audit) = audit_handle {
            let (stream, agg_fut) = if audit.streaming() {
                // Streaming: apply scan (pass-through + parallel aggregation)
                crate::audit::stream::scan_aggregate_with_future(transformed_stream)
            } else {
                // Non-streaming: apply fold (collect all, then emit single chunk)
                crate::audit::stream::fold_aggregate_with_future(transformed_stream)
            };

            // Spawn audit task
            tokio::spawn(async move {
                let final_resp = agg_fut.await;
                audit.set_response(Arc::new(final_resp));
                audit.emit();
            });

            Box::pin(stream)
        } else {
            transformed_stream
        };

Yan Ru Pei's avatar
Yan Ru Pei committed
1169
1170
1171
1172
1173
1174
1175
1176
1177
        // Step 5: Speculative next-turn prefill
        let final_stream = speculative_prefill::maybe_wrap_stream(
            final_stream,
            &request,
            &next,
            &self.formatter,
            &self.tokenizer,
        );

Biswa Panda's avatar
Biswa Panda committed
1178
        // prepend the annotations to the response stream
1179
        let stream = annotations_stream.chain(final_stream);
Biswa Panda's avatar
Biswa Panda committed
1180

Ryan Olson's avatar
Ryan Olson committed
1181
        // return the response stream - single boxing at the end
Biswa Panda's avatar
Biswa Panda committed
1182
1183
1184
1185
1186
1187
1188
        Ok(ResponseStream::new(Box::pin(stream), context))
    }
}

#[async_trait]
impl
    Operator<
1189
        SingleIn<NvCreateCompletionRequest>,
1190
        ManyOut<Annotated<NvCreateCompletionResponse>>,
1191
        SingleIn<PreprocessedRequest>,
Biswa Panda's avatar
Biswa Panda committed
1192
1193
1194
1195
1196
        ManyOut<Annotated<BackendOutput>>,
    > for OpenAIPreprocessor
{
    async fn generate(
        &self,
1197
        request: SingleIn<NvCreateCompletionRequest>,
Biswa Panda's avatar
Biswa Panda committed
1198
        next: Arc<
1199
            dyn AsyncEngine<SingleIn<PreprocessedRequest>, ManyOut<Annotated<BackendOutput>>, Error>,
Biswa Panda's avatar
Biswa Panda committed
1200
        >,
1201
    ) -> Result<ManyOut<Annotated<NvCreateCompletionResponse>>, Error> {
Biswa Panda's avatar
Biswa Panda committed
1202
        // unpack the request
1203
1204
        let (mut request, context) = request.into_parts();

1205
1206
1207
1208
1209
1210
1211
1212
        // Preserve original streaming flag
        let original_stream_flag = request.inner.stream.unwrap_or(false);

        // For non-streaming requests (stream=false), enable usage by default
        // This ensures compliance with OpenAI API spec where non-streaming responses
        // always include usage statistics
        request.enable_usage_for_nonstreaming(original_stream_flag);

1213
        request.inner.stream = Some(true);
Biswa Panda's avatar
Biswa Panda committed
1214
1215

        // create a response generator
1216
        let response_generator = request.response_generator(context.id().to_string());
Biswa Panda's avatar
Biswa Panda committed
1217
        let mut response_generator = Box::new(response_generator);
1218
        let tracker = response_generator.tracker();
Biswa Panda's avatar
Biswa Panda committed
1219
        // convert the chat completion request to a common completion request
1220
        let mut builder = self.builder(&request)?;
1221
1222
1223
1224
1225
1226
1227
1228
1229
1230

        // Check if embeddings are provided - skip tokenization path
        let annotations = if let Some(ref prompt_embeds) = request.inner.prompt_embeds {
            // Skip tokenization for embeddings
            builder.token_ids(vec![]); // Empty token IDs
            builder.prompt_embeds(Some(prompt_embeds.clone()));
            // No token annotations
            HashMap::new()
        } else {
            // Normal path: tokenize the prompt
1231
            self.gather_tokens(&request, &mut builder, None, tracker.as_deref())?
1232
1233
1234
        };

        // Gather multimodal data (works with both embeddings and text prompts)
1235
        self.gather_multi_modal_data(&request, &mut builder).await?;
1236

1237
1238
1239
        let mut common_request = builder.build()?;

        // Attach the timing tracker to the request so downstream components can record metrics
1240
        common_request.tracker = tracker;
Biswa Panda's avatar
Biswa Panda committed
1241

1242
1243
1244
1245
1246
        // Update ISL only for text prompts (embeddings get sequence length from tensor shape)
        if common_request.prompt_embeds.is_none() {
            let isl = common_request.token_ids.len() as u32;
            response_generator.update_isl(isl);
        }
Biswa Panda's avatar
Biswa Panda committed
1247
1248
1249
1250
1251

        // repack the common completion request
        let common_request = context.map(|_| common_request);

        // create a stream of annotations this will be prepend to the response stream
1252
        let annotations: Vec<Annotated<NvCreateCompletionResponse>> = annotations
Biswa Panda's avatar
Biswa Panda committed
1253
1254
1255
1256
1257
1258
1259
1260
            .into_iter()
            .flat_map(|(k, v)| Annotated::from_annotation(k, &v))
            .collect();
        let annotations_stream = stream::iter(annotations);

        // forward the common completion request to the next operator
        let response_stream = next.generate(common_request).await?;

Ryan Olson's avatar
Ryan Olson committed
1261
1262
1263
        // Extract context once
        let context = response_stream.context();

Biswa Panda's avatar
Biswa Panda committed
1264
        // transform the postprocessor stream
Ryan Olson's avatar
Ryan Olson committed
1265
1266
1267
1268
1269
        let stream = Self::transform_postprocessor_stream(
            response_stream,
            response_generator,
            context.clone(),
        );
Biswa Panda's avatar
Biswa Panda committed
1270
1271
1272
1273
1274
1275
1276
1277

        // prepend the annotations to the response stream
        let stream = annotations_stream.chain(stream);

        // return the response stream
        Ok(ResponseStream::new(Box::pin(stream), context))
    }
}
1278
1279
1280
1281
1282
1283
1284
1285
1286
1287
1288
1289
1290
1291
1292

#[async_trait]
impl
    Operator<
        SingleIn<NvCreateEmbeddingRequest>,
        ManyOut<Annotated<NvCreateEmbeddingResponse>>,
        SingleIn<PreprocessedEmbeddingRequest>,
        ManyOut<Annotated<EmbeddingsEngineOutput>>,
    > for OpenAIPreprocessor
{
    async fn generate(
        &self,
        request: SingleIn<NvCreateEmbeddingRequest>,
        next: Arc<
            dyn AsyncEngine<
1293
1294
1295
1296
                    SingleIn<PreprocessedEmbeddingRequest>,
                    ManyOut<Annotated<EmbeddingsEngineOutput>>,
                    Error,
                >,
1297
1298
1299
1300
1301
1302
1303
1304
1305
1306
1307
1308
1309
        >,
    ) -> Result<ManyOut<Annotated<NvCreateEmbeddingResponse>>, Error> {
        // Unpack request
        let (request, context) = request.into_parts();

        // Preprocess the embedding request
        let (preprocessed_request, annotations) =
            self.preprocess_embedding_request(&request).await?;

        // Forward to next stage
        let preprocessed_request = context.map(|_| preprocessed_request);
        let response_stream = next.generate(preprocessed_request).await?;

Ryan Olson's avatar
Ryan Olson committed
1310
1311
1312
        // Extract context once
        let context = response_stream.context();

1313
1314
1315
1316
1317
1318
1319
1320
1321
1322
1323
1324
1325
1326
1327
        // Transform response stream back to OpenAI format
        let stream = Self::transform_embedding_postprocessor_stream(response_stream, request);

        // Prepend annotations
        let annotations_stream = stream::iter(
            annotations
                .into_iter()
                .flat_map(|(k, v)| Annotated::from_annotation(k, &v))
                .collect::<Vec<_>>(),
        );

        let combined_stream = annotations_stream.chain(stream);
        Ok(ResponseStream::new(Box::pin(combined_stream), context))
    }
}
Ryan Olson's avatar
Ryan Olson committed
1328
1329

// Note: tests for jailing and parser detection live in `lib/llm/tests/test_jail.rs`