preprocessor.rs 44 KB
Newer Older
Biswa Panda's avatar
Biswa Panda committed
1
2
3
4
5
6
// SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// 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
16
pub mod prompt;
pub mod tools;
17
use anyhow::Context;
18
use anyhow::{Result, bail};
19
20
21
22
use dynamo_async_openai::types::{
    ChatCompletionRequestMessage, ChatCompletionRequestUserMessageContent,
    ChatCompletionRequestUserMessageContentPart, ChatCompletionToolChoiceOption, EncodingFormat,
};
Ryan Olson's avatar
Ryan Olson committed
23
use futures::Stream;
Biswa Panda's avatar
Biswa Panda committed
24
25
use futures::stream::{self, StreamExt};
use prompt::OAIPromptFormatter;
Ryan Olson's avatar
Ryan Olson committed
26
use std::{collections::HashMap, pin::Pin, sync::Arc};
Biswa Panda's avatar
Biswa Panda committed
27
28
use tracing;

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

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

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

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

60
pub use crate::protocols::common::llm_backend::{BackendOutput, PreprocessedRequest};
61
62
63
pub use crate::protocols::common::preprocessor::PreprocessedEmbeddingRequest;

use crate::protocols::common::llm_backend::EmbeddingsEngineOutput;
Biswa Panda's avatar
Biswa Panda committed
64
65
66

pub const ANNOTATION_FORMATTED_PROMPT: &str = "formatted_prompt";
pub const ANNOTATION_TOKEN_IDS: &str = "token_ids";
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
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,
}

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
102

103
104
105
106
107
108
// 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
109
110
111
112
113
pub struct OpenAIPreprocessor {
    mdcsum: String,
    formatter: Arc<dyn OAIPromptFormatter>,
    tokenizer: Arc<dyn Tokenizer>,
    model_info: Arc<dyn ModelInfo>,
114
115
    /// Per-model runtime configuration propagated to response generator (e.g., reasoning/tool parser)
    runtime_config: crate::local_model::runtime_config::ModelRuntimeConfig,
116
    tool_call_parser: Option<String>,
117
    media_loader: Option<MediaLoader>,
Biswa Panda's avatar
Biswa Panda committed
118
119
120
}

impl OpenAIPreprocessor {
121
122
123
    pub fn new(mdc: ModelDeploymentCard) -> Result<Arc<Self>> {
        let formatter = PromptFormatter::from_mdc(&mdc)?;
        let tokenizer = mdc.tokenizer_hf()?;
124
        match formatter {
125
            PromptFormatter::OAI(formatter) => Self::new_with_parts(mdc, formatter, tokenizer),
126
127
128
        }
    }

129
    pub fn new_with_parts(
130
131
        mdc: ModelDeploymentCard,
        formatter: Arc<dyn OAIPromptFormatter>,
132
        hf_tokenizer: tokenizers::Tokenizer,
133
    ) -> Result<Arc<Self>> {
134
        let mdcsum = mdc.mdcsum().to_string();
135
        let tokenizer = Arc::new(HuggingFaceTokenizer::from_tokenizer(hf_tokenizer));
136
137
138
139
140
        let Some(model_info) = mdc.model_info else {
            anyhow::bail!(
                "Blank ModelDeploymentCard cannot be used for pre-processing, no model_info"
            );
        };
141
        let model_info = model_info.get_model_info()?;
142
        let tool_call_parser = mdc.runtime_config.tool_call_parser.clone();
Biswa Panda's avatar
Biswa Panda committed
143

144
145
        // // Initialize runtime config from the ModelDeploymentCard
        let runtime_config = mdc.runtime_config.clone();
146
        let media_loader = None; // TODO: enable with decoder config from MDC
Biswa Panda's avatar
Biswa Panda committed
147
148
149
150
151
        Ok(Arc::new(Self {
            formatter,
            tokenizer,
            model_info,
            mdcsum,
152
            runtime_config,
153
            tool_call_parser,
154
            media_loader,
Biswa Panda's avatar
Biswa Panda committed
155
156
        }))
    }
157
158
159
160
161
    /// Encode a string to it's tokens
    pub fn tokenize(&self, s: &str) -> anyhow::Result<Encoding> {
        self.tokenizer.encode(s)
    }

162
    /// Translate a [`NvCreateChatCompletionRequest`] request to a common completion request.
Biswa Panda's avatar
Biswa Panda committed
163
164
165
166
167
    /// Returns both the common completion request and a hashmap of annotations.
    ///
    /// Annotations evaluated by this method include:
    /// - `formatted_prompt`
    /// - `token_ids`
168
    pub async fn preprocess_request<
Biswa Panda's avatar
Biswa Panda committed
169
170
171
172
        R: OAIChatLikeRequest
            + AnnotationsProvider
            + SamplingOptionsProvider
            + StopConditionsProvider
Greg Clark's avatar
Greg Clark committed
173
            + OutputOptionsProvider
Biswa Panda's avatar
Biswa Panda committed
174
175
176
177
            + NvExtProvider,
    >(
        &self,
        request: &R,
178
    ) -> Result<(PreprocessedRequest, HashMap<String, String>)> {
179
        let mut builder = self.builder(request)?;
180
181
182
183
184
185
186
        let formatted_prompt = self
            .apply_template(request)
            .with_context(|| "Failed to apply prompt template")?;
        let annotations = self
            .gather_tokens(request, &mut builder, formatted_prompt)
            .with_context(|| "Failed to gather tokens")?;
        self.gather_multi_modal_data(request, &mut builder)
187
            .await
188
            .with_context(|| "Failed to gather multimodal data")?;
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203

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

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

207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
        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()));
        builder.estimated_prefix_hit_num_blocks(None);
        // Extract backend_instance_id from nvext if present
        if let Some(nvext) = request.nvext() {
            builder.backend_instance_id(nvext.backend_instance_id);
        }

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

274
    pub async fn gather_multi_modal_data<R: OAIChatLikeRequest>(
275
276
277
278
279
280
281
        &self,
        request: &R,
        builder: &mut PreprocessedRequestBuilder,
    ) -> Result<()> {
        let messages = request.messages();
        let message_count = messages.len().unwrap_or(0);
        let mut media_map: MultimodalDataMap = HashMap::new();
282
        let mut fetch_tasks = Vec::new();
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314

        for idx in 0..message_count {
            let msg = messages
                .get_item_by_index(idx)
                .map_err(|_| anyhow::Error::msg(format!("Cannot get message at index {idx}")))?;

            let msg_json: serde_json::Value = serde_json::to_value(&msg)?;
            let message: ChatCompletionRequestMessage = serde_json::from_value(msg_json)?;

            let content_parts = match &message {
                ChatCompletionRequestMessage::User(u) => match &u.content {
                    ChatCompletionRequestUserMessageContent::Array(parts) => parts,
                    _ => continue,
                },
                _ => continue,
            };

            // Iterate over content parts
            for content_part in content_parts {
                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,
                };

315
316
317
318
319
320
321
322
323
                if self.media_loader.is_some() {
                    fetch_tasks.push((type_str, content_part.clone()));
                } else {
                    // No loader, just pass the URL through
                    media_map
                        .entry(type_str)
                        .or_default()
                        .push(MultimodalData::Url(url));
                }
324
325
            }
        }
326
327
328
329
330
331
332

        // Execute all fetch tasks
        if !fetch_tasks.is_empty() {
            let loader = self.media_loader.as_ref().unwrap();
            let _results = futures::future::join_all(
                fetch_tasks
                    .iter()
333
                    .map(|(_, content_part)| loader.fetch_and_decode_media_part(content_part)),
334
335
336
337
338
339
            )
            .await;

            // TODO: decode and pass NIXL descriptors to the media map
        }

340
341
        if !media_map.is_empty() {
            builder.multi_modal_data(Some(media_map));
342
343
344
345
346
347
348
349

            // Preserve original messages in extra_args for multimodal workers that need them
            // (e.g., TRT-LLM multimodal processor needs raw messages for proper tokenization)
            let messages_json = serde_json::to_value(&messages)?;
            let extra_args = serde_json::json!({
                "messages": messages_json
            });
            builder.extra_args(Some(extra_args));
350
351
352
353
354
        }

        Ok(())
    }

355
356
357
358
359
360
361
362
363
364
365
366
367
368
    pub fn gather_tokens<
        R: OAIChatLikeRequest
            + AnnotationsProvider
            + SamplingOptionsProvider
            + StopConditionsProvider
            + OutputOptionsProvider
            + NvExtProvider,
    >(
        &self,
        request: &R,
        builder: &mut PreprocessedRequestBuilder,
        formatted_prompt: Option<String>,
    ) -> Result<HashMap<String, String>> {
        let mut annotations = HashMap::new();
369
370
371
372
373
374
375
376
377
378
379
380
        // 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 {
381
382
383
384
                                bail!(
                                    "Batch token input not supported for more than one token in requests (got {})",
                                    token_batches.len()
                                );
385
386
387
                            }
                        }
                    }
Biswa Panda's avatar
Biswa Panda committed
388
389
                }
            }
390
391
392
            PromptInput::Text(_) => {
                if let Some(text_input) = request.extract_text() {
                    match text_input {
393
394
395
396
397
398
399
400
401
402
                        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
403

404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
                            // 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"
                                    );
425
                                    let encoding = self.tokenizer.encode(&prompt)?;
426
427
428
429
                                    (encoding.token_ids().to_vec(), false)
                                }
                            } else {
                                // No backend_instance_id provided, continue the normal flow.
430
                                let encoding = self.tokenizer.encode(&prompt)?;
431
432
                                (encoding.token_ids().to_vec(), false)
                            };
Biswa Panda's avatar
Biswa Panda committed
433

434
435
436
                            if request.has_annotation(ANNOTATION_TOKEN_IDS)
                                && !skip_token_annotation
                            {
437
438
                                annotations.insert(
                                    ANNOTATION_TOKEN_IDS.to_string(),
439
                                    serde_json::to_string(&tokens_vec)?,
440
441
442
                                );
                            }

443
                            builder.token_ids(tokens_vec);
444
445
                        }
                        TextInput::Batch(texts) => {
446
447
448
449
450
451
452
453
454
                            if texts.len() == 1 {
                                let encoding = self.tokenizer.encode(&texts[0])?;
                                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()
                                );
                            }
455
456
457
458
                        }
                    }
                }
            }
Biswa Panda's avatar
Biswa Panda committed
459
        }
460
        Ok(annotations)
Biswa Panda's avatar
Biswa Panda committed
461
462
    }

463
464
465
466
467
468
469
470
471
472
473
474
475
476
    /// 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 {
477
            dynamo_async_openai::types::EmbeddingInput::String(s) => {
478
479
                let encoding = self.tokenizer.encode(s)?;
                vec![encoding.token_ids().to_vec()]
480
            }
481
            dynamo_async_openai::types::EmbeddingInput::StringArray(arr) => {
482
483
484
485
486
487
488
489
490
491
492
                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()
493
                    .map(|encoding| encoding.token_ids().to_vec())
494
495
496
                    .collect();
                token_arrays
            }
497
498
499
500
            dynamo_async_openai::types::EmbeddingInput::IntegerArray(token_ids) => {
                vec![token_ids.clone()]
            }
            dynamo_async_openai::types::EmbeddingInput::ArrayOfIntegerArray(token_arrays) => {
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
                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
527
528
    pub fn transform_postprocessor_stream<S, Resp>(
        stream: S,
Biswa Panda's avatar
Biswa Panda committed
529
        generator: Box<dyn DeltaGeneratorExt<Resp>>,
Ryan Olson's avatar
Ryan Olson committed
530
531
532
533
534
535
536
537
538
539
540
        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
541
542
543
            response_generator: Box<dyn DeltaGeneratorExt<Resp>>,
            context: Arc<dyn AsyncEngineContext>,
            cancelled: bool,
544
            cumulative_output_tokens: usize,
545
546
            finish_reason_sent: bool,
            usage_chunk_sent: bool,
Ryan Olson's avatar
Ryan Olson committed
547
            finished: bool,
Biswa Panda's avatar
Biswa Panda committed
548
549
550
        }

        let state = State {
Ryan Olson's avatar
Ryan Olson committed
551
            response_stream: Box::pin(stream),
Biswa Panda's avatar
Biswa Panda committed
552
553
554
            response_generator: generator,
            context: context.clone(),
            cancelled: false,
555
            cumulative_output_tokens: 0,
556
557
            finish_reason_sent: false,
            usage_chunk_sent: false,
Ryan Olson's avatar
Ryan Olson committed
558
            finished: false,
Biswa Panda's avatar
Biswa Panda committed
559
560
561
        };

        // transform the common response stream into a chat response stream
Ryan Olson's avatar
Ryan Olson committed
562
563

        stream::unfold(state, |mut inner| {
Biswa Panda's avatar
Biswa Panda committed
564
            async move {
565
566
567
568
569
                // If already finished, return None immediately
                if inner.finished {
                    return None;
                }

Biswa Panda's avatar
Biswa Panda committed
570
571
572
573
574
575
                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"
                        );
576
                        inner.finished = true; // Mark as finished
Biswa Panda's avatar
Biswa Panda committed
577
578
579
580
581
582
583
584
585
                        return None;
                    }

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

586
587
588
589
590
591
592
                    // 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);

593
594
595
596
597
598
599
600
601
602
603
604
605
606
                    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
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
                        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())
                    });

622
623
624
625
626
627
628
629
630
631
632
                    // Create LLM metrics annotation
                    let llm_metrics = LLMMetricAnnotation {
                        input_tokens: isl,
                        output_tokens: current_osl,
                        chunk_tokens,
                    };

                    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;
633
                            response.comment = metrics_annotated.comment;
634
635
                        }
                    }
636

637
638
639
640
641
                    // Mark if we've seen a finish_reason
                    if has_finish_reason {
                        inner.finish_reason_sent = true;
                    }

Biswa Panda's avatar
Biswa Panda committed
642
643
                    tracing::trace!(
                        request_id = inner.context.id(),
644
                        "OpenAI NvCreateChatCompletionStreamResponse: {:?}",
Biswa Panda's avatar
Biswa Panda committed
645
646
647
648
649
                        response
                    );

                    Some((response, inner))
                } else {
650
651
652
653
654
                    // 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;

                    // Check if we need to send a usage chunk
655
656
657
658
659
660
661
662
663
664
665
                    if inner.response_generator.is_usage_enabled()
                        && inner.finish_reason_sent
                        && !inner.usage_chunk_sent
                    {
                        inner.usage_chunk_sent = true;

                        // Create the final usage chunk
                        let usage_chunk = inner.response_generator.create_usage_chunk();
                        let annotated_usage = Annotated::<Resp> {
                            id: None,
                            data: Some(usage_chunk),
666
                            event: None,
667
668
669
670
671
                            comment: None,
                        };

                        tracing::trace!(
                            request_id = inner.context.id(),
672
673
                            "Sending final usage chunk for OpenAI compliance, annotated_usage: {:?}",
                            annotated_usage
674
675
676
677
678
679
680
                        );

                        Some((annotated_usage, inner))
                    } else {
                        // stream closed
                        None
                    }
Biswa Panda's avatar
Biswa Panda committed
681
682
                }
            }
Ryan Olson's avatar
Ryan Olson committed
683
        })
Biswa Panda's avatar
Biswa Panda committed
684
    }
685
686

    /// Transform engine embedding output stream to OpenAI embedding response stream
Ryan Olson's avatar
Ryan Olson committed
687
688
    pub fn transform_embedding_postprocessor_stream<S>(
        stream: S,
689
        original_request: NvCreateEmbeddingRequest,
Ryan Olson's avatar
Ryan Olson committed
690
691
692
693
694
    ) -> impl Stream<Item = Annotated<NvCreateEmbeddingResponse>> + Send
    where
        S: Stream<Item = Annotated<EmbeddingsEngineOutput>> + Send + 'static,
    {
        stream.map(move |output| {
695
696
            output.map_data(|engine_output| {
                // Convert engine output to OpenAI response format
697
                let embeddings: Vec<dynamo_async_openai::types::Embedding> = engine_output
698
699
700
                    .embeddings
                    .into_iter()
                    .enumerate()
701
                    .map(|(index, embedding)| dynamo_async_openai::types::Embedding {
702
703
704
705
706
707
708
                        index: index as u32,
                        object: "embedding".to_string(),
                        embedding: embedding.into_iter().map(|f| f as f32).collect(),
                    })
                    .collect();

                let response = NvCreateEmbeddingResponse {
709
                    inner: dynamo_async_openai::types::CreateEmbeddingResponse {
710
711
712
                        object: "list".to_string(),
                        model: original_request.inner.model.clone(),
                        data: embeddings,
713
                        usage: dynamo_async_openai::types::EmbeddingUsage {
714
715
716
717
718
719
720
721
                            prompt_tokens: engine_output.prompt_tokens,
                            total_tokens: engine_output.total_tokens,
                        },
                    },
                };

                Ok(response)
            })
Ryan Olson's avatar
Ryan Olson committed
722
        })
723
724
    }

Ryan Olson's avatar
Ryan Olson committed
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
    /// 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) {
            // No parser but tools requested - error cases
            (None, Some(ChatCompletionToolChoiceOption::Required), true) => {
                tracing::warn!(
                    "Tool choice 'required' specified but no tool parser configured; proceeding without jailing"
                );
                Ok(false)
            }
            (None, Some(ChatCompletionToolChoiceOption::Auto), true) => {
                tracing::warn!(
                    "Tool choice 'auto' specified but no tool parser configured; proceeding without jailing"
                );
                Ok(false)
            }
            (None, Some(ChatCompletionToolChoiceOption::Named(_)), _) => {
                tracing::warn!(
                    "Named tool choice specified but no tool parser configured; proceeding without jailing"
                );
                Ok(false)
751
752
            }

Ryan Olson's avatar
Ryan Olson committed
753
754
755
            // Parser exists and tools might be called
            (Some(_), Some(ChatCompletionToolChoiceOption::None), _) => {
                Ok(false) // Explicitly disabled
756
            }
Ryan Olson's avatar
Ryan Olson committed
757
758
            (Some(_), Some(_), true) => Ok(true), // Any other tool_choice with tools
            (Some(_), None, true) => Ok(true),    // Default behavior when tools present
759

Ryan Olson's avatar
Ryan Olson committed
760
761
762
763
            // No tools or no parser
            _ => Ok(false),
        }
    }
764

Ryan Olson's avatar
Ryan Olson committed
765
766
767
768
769
770
771
772
773
774
775
    /// Apply tool calling jail to the stream if needed
    pub fn apply_tool_calling_jail<S>(
        tool_call_parser: String,
        stream: S,
    ) -> impl Stream<Item = Annotated<NvCreateChatCompletionStreamResponse>> + Send
    where
        S: Stream<Item = Annotated<NvCreateChatCompletionStreamResponse>> + Send + 'static,
    {
        let jail = JailedStream::builder()
            .tool_call_parser(tool_call_parser)
            .build();
776
        jail.apply_with_finish_reason(stream)
Ryan Olson's avatar
Ryan Olson committed
777
    }
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827

    // 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() {
                            if let Some(text) = choice.delta.content.as_ref() {
                                let parser_result =
                                    parser.parse_reasoning_streaming_incremental(text, &[]);

                                // Update this specific choice with parsed content
                                choice.delta.content = parser_result.get_some_normal_text();
                                choice.delta.reasoning_content = parser_result.get_some_reasoning();
                            }
                        }
                        Ok(data)
                    })
                } else {
                    // No reasoning parser configured, pass through unchanged
                    response
                };

                Some((processed_response, state))
            } else {
                None
            }
        })
    }
Biswa Panda's avatar
Biswa Panda committed
828
829
830
831
832
833
834
835
836
837
}

// 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<
838
        SingleIn<NvCreateChatCompletionRequest>,
839
        ManyOut<Annotated<NvCreateChatCompletionStreamResponse>>,
840
        SingleIn<PreprocessedRequest>,
Biswa Panda's avatar
Biswa Panda committed
841
842
843
844
845
        ManyOut<Annotated<BackendOutput>>,
    > for OpenAIPreprocessor
{
    async fn generate(
        &self,
846
        request: SingleIn<NvCreateChatCompletionRequest>,
Biswa Panda's avatar
Biswa Panda committed
847
        next: Arc<
848
            dyn AsyncEngine<SingleIn<PreprocessedRequest>, ManyOut<Annotated<BackendOutput>>, Error>,
Biswa Panda's avatar
Biswa Panda committed
849
        >,
850
    ) -> Result<ManyOut<Annotated<NvCreateChatCompletionStreamResponse>>, Error> {
Biswa Panda's avatar
Biswa Panda committed
851
        // unpack the request
852
853
854
855
        let (mut request, context) = request.into_parts();

        // Preserve original inbound streaming flag before any internal overrides
        let request_id = context.id().to_string();
856
        let original_stream_flag = request.inner.stream.unwrap_or(false);
857
858
859
860
861
862
863
864

        // Build audit handle (None if DYN_AUDIT_ENABLED=0)
        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()));
        }

865
866
867
868
869
        // 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);

870
871
        // Set stream=true for internal processing (after audit capture)
        request.inner.stream = Some(true);
Biswa Panda's avatar
Biswa Panda committed
872
873

        // create a response generator
874
        let response_generator = request.response_generator(context.id().to_string());
Ryan Olson's avatar
Ryan Olson committed
875
876

        // convert the chat completion request to a common completion request
877
        let (common_request, annotations) = self.preprocess_request(&request).await?;
Ryan Olson's avatar
Ryan Olson committed
878

Biswa Panda's avatar
Biswa Panda committed
879
880
881
        let mut response_generator = Box::new(response_generator);

        // update isl
Paul Hendricks's avatar
Paul Hendricks committed
882
        response_generator.update_isl(common_request.token_ids.len() as u32);
Biswa Panda's avatar
Biswa Panda committed
883
884
885
886
887

        // 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
888
        let annotations: Vec<Annotated<NvCreateChatCompletionStreamResponse>> = annotations
Biswa Panda's avatar
Biswa Panda committed
889
890
891
892
893
894
895
            .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
896
897
898
899
900
901
902
903
904
905
        // 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(),
        );

906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
        // 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
925
926
927
928
929
930
931
932
933
934
935
936
        // 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,
        )?;

        // Apply jail conditionally
937
        let transformed_stream: Pin<Box<dyn Stream<Item = _> + Send>> = if should_jail {
Ryan Olson's avatar
Ryan Olson committed
938
939
940
941
942
            if let Some(parser) = self.tool_call_parser.clone() {
                Box::pin(Self::apply_tool_calling_jail(parser, stream))
            } else {
                Box::pin(stream) // Should not happen due to should_jail check
            }
943
        } else {
Ryan Olson's avatar
Ryan Olson committed
944
            Box::pin(stream)
945
        };
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968

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

Biswa Panda's avatar
Biswa Panda committed
969
        // prepend the annotations to the response stream
970
        let stream = annotations_stream.chain(final_stream);
Biswa Panda's avatar
Biswa Panda committed
971

Ryan Olson's avatar
Ryan Olson committed
972
        // return the response stream - single boxing at the end
Biswa Panda's avatar
Biswa Panda committed
973
974
975
976
977
978
979
        Ok(ResponseStream::new(Box::pin(stream), context))
    }
}

#[async_trait]
impl
    Operator<
980
        SingleIn<NvCreateCompletionRequest>,
981
        ManyOut<Annotated<NvCreateCompletionResponse>>,
982
        SingleIn<PreprocessedRequest>,
Biswa Panda's avatar
Biswa Panda committed
983
984
985
986
987
        ManyOut<Annotated<BackendOutput>>,
    > for OpenAIPreprocessor
{
    async fn generate(
        &self,
988
        request: SingleIn<NvCreateCompletionRequest>,
Biswa Panda's avatar
Biswa Panda committed
989
        next: Arc<
990
            dyn AsyncEngine<SingleIn<PreprocessedRequest>, ManyOut<Annotated<BackendOutput>>, Error>,
Biswa Panda's avatar
Biswa Panda committed
991
        >,
992
    ) -> Result<ManyOut<Annotated<NvCreateCompletionResponse>>, Error> {
Biswa Panda's avatar
Biswa Panda committed
993
        // unpack the request
994
995
        let (mut request, context) = request.into_parts();

996
997
998
999
1000
1001
1002
1003
        // 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);

1004
        request.inner.stream = Some(true);
Biswa Panda's avatar
Biswa Panda committed
1005
1006

        // create a response generator
1007
        let response_generator = request.response_generator(context.id().to_string());
Biswa Panda's avatar
Biswa Panda committed
1008
1009
        let mut response_generator = Box::new(response_generator);
        // convert the chat completion request to a common completion request
1010
1011
        let mut builder = self.builder(&request)?;
        let annotations = self.gather_tokens(&request, &mut builder, None)?;
1012
        self.gather_multi_modal_data(&request, &mut builder).await?;
1013

1014
        let common_request = builder.build()?;
Biswa Panda's avatar
Biswa Panda committed
1015
1016

        // update isl
1017
        response_generator.update_isl(common_request.token_ids.len() as u32);
Biswa Panda's avatar
Biswa Panda committed
1018
1019
1020
1021
1022

        // 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
1023
        let annotations: Vec<Annotated<NvCreateCompletionResponse>> = annotations
Biswa Panda's avatar
Biswa Panda committed
1024
1025
1026
1027
1028
1029
1030
1031
            .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
1032
1033
1034
        // Extract context once
        let context = response_stream.context();

Biswa Panda's avatar
Biswa Panda committed
1035
        // transform the postprocessor stream
Ryan Olson's avatar
Ryan Olson committed
1036
1037
1038
1039
1040
        let stream = Self::transform_postprocessor_stream(
            response_stream,
            response_generator,
            context.clone(),
        );
Biswa Panda's avatar
Biswa Panda committed
1041
1042
1043
1044
1045
1046
1047
1048

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

        // return the response stream
        Ok(ResponseStream::new(Box::pin(stream), context))
    }
}
1049
1050
1051
1052
1053
1054
1055
1056
1057
1058
1059
1060
1061
1062
1063

#[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<
1064
1065
1066
1067
                    SingleIn<PreprocessedEmbeddingRequest>,
                    ManyOut<Annotated<EmbeddingsEngineOutput>>,
                    Error,
                >,
1068
1069
1070
1071
1072
1073
1074
1075
1076
1077
1078
1079
1080
        >,
    ) -> 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
1081
1082
1083
        // Extract context once
        let context = response_stream.context();

1084
1085
1086
1087
1088
1089
1090
1091
1092
1093
1094
1095
1096
1097
1098
        // 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
1099
1100

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