preprocessor.rs 58 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,
    },
};
Nikita's avatar
Nikita committed
59
use crate::tokenizers::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
    pub fn new(mdc: ModelDeploymentCard) -> Result<Arc<Self>> {
        let formatter = PromptFormatter::from_mdc(&mdc)?;
Nikita's avatar
Nikita committed
154
        let tokenizer = mdc.tokenizer()?;
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>,
Nikita's avatar
Nikita committed
163
        tokenizer: crate::tokenizers::Tokenizer,
164
    ) -> Result<Arc<Self>> {
165
        let mdcsum = mdc.mdcsum().to_string();
Nikita's avatar
Nikita committed
166
        let tokenizer: Arc<dyn Tokenizer> = (*tokenizer).clone();
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
                priority: hints.and_then(|h| h.priority),
287
                lora_name,
288
                allowed_worker_ids: None,
289
290
            };
            builder.routing(Some(routing));
291
292
293
294
295
296
        } 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()
            }));
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
335
336
        }

        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)
        }
    }

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

346
347
348
349
350
        let Some(messages) = request.typed_messages() else {
            return Ok(());
        };
        for message in messages.iter() {
            let content_parts = match message {
351
352
353
354
355
356
357
                ChatCompletionRequestMessage::User(u) => match &u.content {
                    ChatCompletionRequestUserMessageContent::Array(parts) => parts,
                    _ => continue,
                },
                _ => continue,
            };
            // Iterate over content parts
358
            for content_part in content_parts.iter() {
359
360
361
362
363
364
365
366
367
368
369
370
371
                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,
                };

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

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

        // Execute all fetch tasks
        if !fetch_tasks.is_empty() {
            let loader = self.media_loader.as_ref().unwrap();
388
389
390
391
            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)
            }))
392
393
            .await;

394
395
396
397
398
399
400
401
            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));
            }
402
403
        }

404
405
        if !media_map.is_empty() {
            builder.multi_modal_data(Some(media_map));
406
407
408

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

        Ok(())
    }

419
420
421
422
423
424
425
426
427
428
429
430
    pub fn gather_tokens<
        R: OAIChatLikeRequest
            + AnnotationsProvider
            + SamplingOptionsProvider
            + StopConditionsProvider
            + OutputOptionsProvider
            + NvExtProvider,
    >(
        &self,
        request: &R,
        builder: &mut PreprocessedRequestBuilder,
        formatted_prompt: Option<String>,
431
        tracker: Option<&RequestTracker>,
432
433
    ) -> Result<HashMap<String, String>> {
        let mut annotations = HashMap::new();
434
435
436
437
438
439
440
441
442
443
444
445
        // 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 {
446
447
448
449
                                bail!(
                                    "Batch token input not supported for more than one token in requests (got {})",
                                    token_batches.len()
                                );
450
451
452
                            }
                        }
                    }
Biswa Panda's avatar
Biswa Panda committed
453
454
                }
            }
455
456
457
            PromptInput::Text(_) => {
                if let Some(text_input) = request.extract_text() {
                    match text_input {
458
459
460
461
462
463
464
465
466
467
                        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
468

469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
                            // 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"
                                    );
490
                                    let encoding = self.encode_with_timing(&prompt, tracker)?;
491
492
493
494
                                    (encoding.token_ids().to_vec(), false)
                                }
                            } else {
                                // No backend_instance_id provided, continue the normal flow.
495
                                let encoding = self.encode_with_timing(&prompt, tracker)?;
496
497
                                (encoding.token_ids().to_vec(), false)
                            };
Biswa Panda's avatar
Biswa Panda committed
498

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

508
                            builder.token_ids(tokens_vec);
509
510
                        }
                        TextInput::Batch(texts) => {
511
                            if texts.len() == 1 {
512
                                let encoding = self.encode_with_timing(&texts[0], tracker)?;
513
514
515
516
517
518
519
                                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()
                                );
                            }
520
521
522
523
                        }
                    }
                }
            }
Biswa Panda's avatar
Biswa Panda committed
524
        }
525
        Ok(annotations)
Biswa Panda's avatar
Biswa Panda committed
526
527
    }

528
529
530
531
532
533
534
535
    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 {
536
            t.record_tokenize_latency(encode_start.elapsed());
537
538
539
540
        }
        Ok(encoding)
    }

541
542
543
544
545
546
547
548
549
550
551
552
553
554
    /// 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 {
555
            dynamo_async_openai::types::EmbeddingInput::String(s) => {
556
557
                let encoding = self.tokenizer.encode(s)?;
                vec![encoding.token_ids().to_vec()]
558
            }
559
            dynamo_async_openai::types::EmbeddingInput::StringArray(arr) => {
560
561
562
563
564
565
566
567
568
569
570
                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()
571
                    .map(|encoding| encoding.token_ids().to_vec())
572
573
574
                    .collect();
                token_arrays
            }
575
576
577
578
            dynamo_async_openai::types::EmbeddingInput::IntegerArray(token_ids) => {
                vec![token_ids.clone()]
            }
            dynamo_async_openai::types::EmbeddingInput::ArrayOfIntegerArray(token_arrays) => {
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
                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
605
606
    pub fn transform_postprocessor_stream<S, Resp>(
        stream: S,
Biswa Panda's avatar
Biswa Panda committed
607
        generator: Box<dyn DeltaGeneratorExt<Resp>>,
Ryan Olson's avatar
Ryan Olson committed
608
609
610
611
612
613
614
615
616
617
618
        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
619
620
621
            response_generator: Box<dyn DeltaGeneratorExt<Resp>>,
            context: Arc<dyn AsyncEngineContext>,
            cancelled: bool,
622
            cumulative_output_tokens: usize,
623
624
            finish_reason_sent: bool,
            usage_chunk_sent: bool,
Ryan Olson's avatar
Ryan Olson committed
625
            finished: bool,
Biswa Panda's avatar
Biswa Panda committed
626
627
628
        }

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

        // transform the common response stream into a chat response stream
Ryan Olson's avatar
Ryan Olson committed
640
641

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

Biswa Panda's avatar
Biswa Panda committed
648
649
650
651
652
653
                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"
                        );
654
                        inner.finished = true; // Mark as finished
Biswa Panda's avatar
Biswa Panda committed
655
656
657
658
659
660
661
662
663
                        return None;
                    }

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

664
665
666
667
668
669
670
                    // 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);

671
672
673
674
675
676
677
678
679
680
681
682
683
684
                    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
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
                        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())
                    });

700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
                    // 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);
715
716
717
718
                    let llm_metrics = LLMMetricAnnotation {
                        input_tokens: isl,
                        output_tokens: current_osl,
                        chunk_tokens,
719
                        cached_tokens: None,
720
721
722
723
724
725
                        prefill_worker_id,
                        prefill_dp_rank,
                        prefill_worker_type,
                        decode_worker_id,
                        decode_dp_rank,
                        decode_worker_type,
726
727
728
                        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()),
729
730
731
732
733
734
                    };

                    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;
735
                            response.comment = metrics_annotated.comment;
736
737
                        }
                    }
738

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

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

                    Some((response, inner))
                } else {
752
753
754
755
                    // 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;

756
                    if inner.finish_reason_sent && !inner.usage_chunk_sent {
757
758
759
                        inner.usage_chunk_sent = true;

                        let usage_chunk = inner.response_generator.create_usage_chunk();
760
                        let usage = inner.response_generator.get_usage();
761
762
763
764
765
766
767
768
769
770
771
772
773
774
                        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);
775
776
777
778
779
780
781
782
                        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)),
783
784
785
786
787
788
                            prefill_worker_id,
                            prefill_dp_rank,
                            prefill_worker_type,
                            decode_worker_id,
                            decode_dp_rank,
                            decode_worker_type,
789
790
791
792
793
                            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()),
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
                        };

                        // 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
                        };

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

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

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

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

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

                Ok(response)
            })
Ryan Olson's avatar
Ryan Olson committed
870
        })
871
872
    }

Ryan Olson's avatar
Ryan Olson committed
873
874
875
876
877
878
879
880
    /// 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) {
881
882
883
884
885
            // 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
886
887
888
889
890
891
            (None, Some(ChatCompletionToolChoiceOption::Auto), true) => {
                tracing::warn!(
                    "Tool choice 'auto' specified but no tool parser configured; proceeding without jailing"
                );
                Ok(false)
            }
892

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

Ryan Olson's avatar
Ryan Olson committed
900
901
902
903
            // No tools or no parser
            _ => Ok(false),
        }
    }
904

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

        let mut builder = JailedStream::builder();

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

926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
        // 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();
947
        jail.apply_with_finish_reason(stream)
Ryan Olson's avatar
Ryan Olson committed
948
    }
949

950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
    /// Check if reasoning parsing should be disabled based on per-request parameters.
    /// For kimi_k25: disabled when chat_template_args contains "thinking": false.
    fn is_reasoning_disabled_by_request(
        reasoning_parser: Option<&str>,
        chat_template_args: Option<&std::collections::HashMap<String, serde_json::Value>>,
    ) -> bool {
        match reasoning_parser {
            Some("kimi_k25") => {
                if let Some(args) = chat_template_args
                    && let Some(thinking) = args.get("thinking")
                {
                    return thinking == &serde_json::Value::Bool(false);
                }
                false
            }
            _ => false,
        }
    }

969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
995
    // 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() {
996
997
998
999
1000
1001
1002
                            // Reasoning parsing only applies to text content
                            if let Some(
                                dynamo_async_openai::types::ChatCompletionMessageContent::Text(
                                    text,
                                ),
                            ) = choice.delta.content.as_ref()
                            {
1003
1004
1005
1006
                                let parser_result =
                                    parser.parse_reasoning_streaming_incremental(text, &[]);

                                // Update this specific choice with parsed content
1007
1008
1009
                                choice.delta.content = parser_result.get_some_normal_text().map(
                                    dynamo_async_openai::types::ChatCompletionMessageContent::Text,
                                );
1010
1011
                                choice.delta.reasoning_content = parser_result.get_some_reasoning();
                            }
1012
                            // For multimodal content, pass through unchanged
1013
1014
1015
1016
1017
1018
1019
1020
1021
1022
1023
1024
1025
                        }
                        Ok(data)
                    })
                } else {
                    // No reasoning parser configured, pass through unchanged
                    response
                };

                Some((processed_response, state))
            } else {
                None
            }
        })
1026
        .fuse()
1027
    }
Biswa Panda's avatar
Biswa Panda committed
1028
1029
1030
1031
1032
1033
1034
1035
1036
1037
}

// 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<
1038
        SingleIn<NvCreateChatCompletionRequest>,
1039
        ManyOut<Annotated<NvCreateChatCompletionStreamResponse>>,
1040
        SingleIn<PreprocessedRequest>,
Biswa Panda's avatar
Biswa Panda committed
1041
1042
1043
1044
1045
        ManyOut<Annotated<BackendOutput>>,
    > for OpenAIPreprocessor
{
    async fn generate(
        &self,
1046
        request: SingleIn<NvCreateChatCompletionRequest>,
Biswa Panda's avatar
Biswa Panda committed
1047
        next: Arc<
1048
            dyn AsyncEngine<SingleIn<PreprocessedRequest>, ManyOut<Annotated<BackendOutput>>, Error>,
Biswa Panda's avatar
Biswa Panda committed
1049
        >,
1050
    ) -> Result<ManyOut<Annotated<NvCreateChatCompletionStreamResponse>>, Error> {
Biswa Panda's avatar
Biswa Panda committed
1051
        // unpack the request
1052
1053
1054
1055
        let (mut request, context) = request.into_parts();

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

1058
        // Build audit handle (None if no DYN_AUDIT_SINKS)
1059
1060
1061
1062
1063
1064
        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()));
        }

1065
1066
1067
1068
1069
        // 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);

1070
1071
        // Set stream=true for internal processing (after audit capture)
        request.inner.stream = Some(true);
Biswa Panda's avatar
Biswa Panda committed
1072
1073

        // create a response generator
1074
        let response_generator = request.response_generator(context.id().to_string());
1075
        let tracker = response_generator.tracker();
Ryan Olson's avatar
Ryan Olson committed
1076
1077

        // convert the chat completion request to a common completion request
1078
1079
1080
        let (mut common_request, annotations) = self
            .preprocess_request(&request, tracker.as_deref())
            .await?;
1081
        tracing::trace!(request = ?common_request, "Pre-processed request");
1082
1083

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

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

1088
1089
1090
1091
1092
        // 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
1093
1094
1095
1096
1097

        // 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
1098
        let annotations: Vec<Annotated<NvCreateChatCompletionStreamResponse>> = annotations
Biswa Panda's avatar
Biswa Panda committed
1099
1100
1101
1102
1103
1104
1105
            .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
1106
1107
1108
1109
1110
1111
1112
1113
1114
1115
        // 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(),
        );

1116
        // Try to parse reasoning content only if parser is configured
1117
1118
1119
1120
1121
        let should_parse_reasoning = self.runtime_config.reasoning_parser.is_some()
            && !Self::is_reasoning_disabled_by_request(
                self.runtime_config.reasoning_parser.as_deref(),
                request.chat_template_args.as_ref(),
            );
1122
1123
1124
1125
1126
1127
1128
1129
1130
1131
1132
1133
1134
1135
1136
1137
1138

        // 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
1139
1140
1141
1142
1143
1144
1145
1146
1147
1148
1149
        // 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,
        )?;

1150
1151
1152
1153
1154
1155
1156
1157
1158
1159
1160
        // 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
1161
        // Apply jail conditionally
1162
        let transformed_stream: Pin<Box<dyn Stream<Item = _> + Send>> = if should_jail {
1163
1164
1165
            Box::pin(Self::apply_tool_calling_jail(
                self.tool_call_parser.clone(),
                request.inner.tool_choice.clone(),
1166
                tool_definitions,
1167
1168
                stream,
            ))
1169
        } else {
Ryan Olson's avatar
Ryan Olson committed
1170
            Box::pin(stream)
1171
        };
1172
1173
1174
1175
1176
1177
1178
1179
1180
1181
1182
1183
1184
1185
1186
1187
1188
1189
1190
1191
1192
1193
1194

        // 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
1195
1196
1197
1198
1199
1200
1201
1202
1203
        // 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
1204
        // prepend the annotations to the response stream
1205
        let stream = annotations_stream.chain(final_stream);
Biswa Panda's avatar
Biswa Panda committed
1206

Ryan Olson's avatar
Ryan Olson committed
1207
        // return the response stream - single boxing at the end
Biswa Panda's avatar
Biswa Panda committed
1208
1209
1210
1211
1212
1213
1214
        Ok(ResponseStream::new(Box::pin(stream), context))
    }
}

#[async_trait]
impl
    Operator<
1215
        SingleIn<NvCreateCompletionRequest>,
1216
        ManyOut<Annotated<NvCreateCompletionResponse>>,
1217
        SingleIn<PreprocessedRequest>,
Biswa Panda's avatar
Biswa Panda committed
1218
1219
1220
1221
1222
        ManyOut<Annotated<BackendOutput>>,
    > for OpenAIPreprocessor
{
    async fn generate(
        &self,
1223
        request: SingleIn<NvCreateCompletionRequest>,
Biswa Panda's avatar
Biswa Panda committed
1224
        next: Arc<
1225
            dyn AsyncEngine<SingleIn<PreprocessedRequest>, ManyOut<Annotated<BackendOutput>>, Error>,
Biswa Panda's avatar
Biswa Panda committed
1226
        >,
1227
    ) -> Result<ManyOut<Annotated<NvCreateCompletionResponse>>, Error> {
Biswa Panda's avatar
Biswa Panda committed
1228
        // unpack the request
1229
1230
        let (mut request, context) = request.into_parts();

1231
1232
1233
1234
1235
1236
1237
1238
        // 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);

1239
        request.inner.stream = Some(true);
Biswa Panda's avatar
Biswa Panda committed
1240
1241

        // create a response generator
1242
        let response_generator = request.response_generator(context.id().to_string());
Biswa Panda's avatar
Biswa Panda committed
1243
        let mut response_generator = Box::new(response_generator);
1244
        let tracker = response_generator.tracker();
Biswa Panda's avatar
Biswa Panda committed
1245
        // convert the chat completion request to a common completion request
1246
        let mut builder = self.builder(&request)?;
1247
1248
1249
1250
1251
1252
1253
1254
1255
1256

        // 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
1257
            self.gather_tokens(&request, &mut builder, None, tracker.as_deref())?
1258
1259
1260
        };

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

1263
1264
1265
        let mut common_request = builder.build()?;

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

1268
1269
1270
1271
1272
        // 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
1273
1274
1275
1276
1277

        // 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
1278
        let annotations: Vec<Annotated<NvCreateCompletionResponse>> = annotations
Biswa Panda's avatar
Biswa Panda committed
1279
1280
1281
1282
1283
1284
1285
1286
            .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
1287
1288
1289
        // Extract context once
        let context = response_stream.context();

Biswa Panda's avatar
Biswa Panda committed
1290
        // transform the postprocessor stream
Ryan Olson's avatar
Ryan Olson committed
1291
1292
1293
1294
1295
        let stream = Self::transform_postprocessor_stream(
            response_stream,
            response_generator,
            context.clone(),
        );
Biswa Panda's avatar
Biswa Panda committed
1296
1297
1298
1299
1300
1301
1302
1303

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

        // return the response stream
        Ok(ResponseStream::new(Box::pin(stream), context))
    }
}
1304
1305
1306
1307
1308
1309
1310
1311
1312
1313
1314
1315
1316
1317
1318

#[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<
1319
1320
1321
1322
                    SingleIn<PreprocessedEmbeddingRequest>,
                    ManyOut<Annotated<EmbeddingsEngineOutput>>,
                    Error,
                >,
1323
1324
1325
1326
1327
1328
1329
1330
1331
1332
1333
1334
1335
        >,
    ) -> 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
1336
1337
1338
        // Extract context once
        let context = response_stream.context();

1339
1340
1341
1342
1343
1344
1345
1346
1347
1348
1349
1350
1351
1352
1353
        // 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
1354
1355

// Note: tests for jailing and parser detection live in `lib/llm/tests/test_jail.rs`
1356
1357
1358
1359
1360
1361
1362
1363
1364
1365
1366
1367
1368
1369
1370
1371
1372
1373
1374
1375
1376
1377
1378
1379
1380
1381
1382
1383
1384
1385
1386
1387
1388
1389
1390
1391
1392
1393
1394
1395
1396
1397
1398
1399
1400
1401
1402
1403
1404
1405
1406
1407
1408
1409
1410
1411
1412
1413
1414
1415
1416
1417
1418
1419
1420
1421
1422
1423
1424
1425
1426
1427
1428
1429

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_is_reasoning_disabled_by_request() {
        let thinking_true = {
            let mut m = std::collections::HashMap::new();
            m.insert("thinking".to_string(), serde_json::Value::Bool(true));
            m
        };
        let thinking_false = {
            let mut m = std::collections::HashMap::new();
            m.insert("thinking".to_string(), serde_json::Value::Bool(false));
            m
        };
        let empty_args = std::collections::HashMap::new();

        // (parser, args, expected_disabled, description)
        let cases = [
            (
                Some("kimi_k25"),
                Some(&thinking_false),
                true,
                "kimi_k25 + thinking=false → disabled",
            ),
            (
                Some("kimi_k25"),
                Some(&thinking_true),
                false,
                "kimi_k25 + thinking=true → enabled",
            ),
            (
                Some("kimi_k25"),
                None,
                false,
                "kimi_k25 + no args → enabled",
            ),
            (
                Some("kimi_k25"),
                Some(&empty_args),
                false,
                "kimi_k25 + empty args → enabled",
            ),
            (
                Some("deepseek_r1"),
                Some(&thinking_false),
                false,
                "deepseek_r1 → never disabled",
            ),
            (
                Some("basic"),
                Some(&thinking_false),
                false,
                "basic → never disabled",
            ),
            (
                None,
                Some(&thinking_false),
                false,
                "no parser → never disabled",
            ),
        ];

        for (parser, args, expected, desc) in cases {
            assert_eq!(
                OpenAIPreprocessor::is_reasoning_disabled_by_request(parser, args),
                expected,
                "FAILED: {desc}",
            );
        }
    }
}