delta.rs 21.5 KB
Newer Older
1
// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2
3
// SPDX-License-Identifier: Apache-2.0

4
5
use std::sync::Arc;

6
use super::{NvCreateChatCompletionRequest, NvCreateChatCompletionStreamResponse};
Greg Clark's avatar
Greg Clark committed
7
use crate::{
8
    local_model::runtime_config::ModelRuntimeConfig,
9
    protocols::{
10
        common::{self, timing::RequestTracker},
11
12
13
14
15
        openai::{
            convert_backend_top_logprobs,
            nvext::{NvExtProvider, NvExtResponse, TimingInfo},
            token_to_utf8_bytes,
        },
16
    },
Greg Clark's avatar
Greg Clark committed
17
18
    types::TokenIdType,
};
19

20
/// Provides a method for generating a [`DeltaGenerator`] from a chat completion request.
21
impl NvCreateChatCompletionRequest {
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
    /// Enables usage tracking for non-streaming requests to comply with OpenAI API specification.
    ///
    /// According to OpenAI API spec, non-streaming chat completion responses (stream=false)
    /// must always include usage statistics. This method ensures `stream_options.include_usage`
    /// is set to `true` for non-streaming requests.
    ///
    /// # Arguments
    /// * `original_stream_flag` - The original value of the `stream` field before any internal processing
    pub fn enable_usage_for_nonstreaming(&mut self, original_stream_flag: bool) {
        if !original_stream_flag {
            // For non-streaming requests (stream=false), enable usage by default
            if self.inner.stream_options.is_none() {
                self.inner.stream_options =
                    Some(dynamo_async_openai::types::ChatCompletionStreamOptions {
                        include_usage: true,
37
                        continuous_usage_stats: false,
38
39
40
41
42
43
44
45
                    });
            } else if let Some(ref mut opts) = self.inner.stream_options {
                // If stream_options exists, ensure include_usage is true for non-streaming
                opts.include_usage = true;
            }
        }
    }

46
47
    /// Creates a [`DeltaGenerator`] instance based on the chat completion request.
    ///
48
49
50
    /// # Arguments
    /// * `request_id` - The request ID to use for the chat completion response ID.
    ///
51
52
    /// # Returns
    /// * [`DeltaGenerator`] configured with model name and response options.
53
    pub fn response_generator(&self, request_id: String) -> DeltaGenerator {
54
55
56
57
        // Enable tracking if:
        // 1. Client requested timing in extra_fields, OR
        // 2. query_instance_id annotation is present (needs worker_id tracking for response)
        let enable_tracking = self
58
            .nvext()
59
60
61
62
63
64
65
66
67
            .map(|nv| {
                nv.extra_fields
                    .as_ref()
                    .is_some_and(|fields| fields.iter().any(|f| f == "timing"))
                    || nv.annotations.as_ref().is_some_and(|annots| {
                        annots.iter().any(|a| a.starts_with("query_instance_id"))
                    })
            })
            .unwrap_or(false);
68

69
        let options = DeltaGeneratorOptions {
70
71
72
73
74
75
            enable_usage: self
                .inner
                .stream_options
                .as_ref()
                .map(|opts| opts.include_usage)
                .unwrap_or(false),
76
77
78
79
80
81
            continuous_usage_stats: self
                .inner
                .stream_options
                .as_ref()
                .map(|opts| opts.continuous_usage_stats)
                .unwrap_or(false),
Greg Clark's avatar
Greg Clark committed
82
83
            enable_logprobs: self.inner.logprobs.unwrap_or(false)
                || self.inner.top_logprobs.unwrap_or(0) > 0,
84
            enable_tracking,
85
            runtime_config: ModelRuntimeConfig::default(),
86
87
        };

88
        DeltaGenerator::new(self.inner.model.clone(), options, request_id)
89
90
91
    }
}

92
/// Configuration options for the [`DeltaGenerator`], controlling response behavior.
93
94
#[derive(Debug, Clone, Default)]
pub struct DeltaGeneratorOptions {
95
    /// Determines whether token usage statistics should be included in the response.
96
    pub enable_usage: bool,
97
98
    /// Determines whether continuous usage statistics should be included in the response.
    pub continuous_usage_stats: bool,
99
    /// Determines whether log probabilities should be included in the response.
100
    pub enable_logprobs: bool,
101
102
    /// Determines whether request tracking (timing, KV hit rate) should be enabled.
    pub enable_tracking: bool,
103

104
    pub runtime_config: ModelRuntimeConfig,
105
106
}

107
/// Generates incremental chat completion responses in a streaming fashion.
108
pub struct DeltaGenerator {
109
    /// Unique identifier for the chat completion session.
110
    id: String,
111
    /// Object type, representing a streamed chat completion response.
112
    object: String,
113
    /// Timestamp (Unix epoch) when the response was created.
Paul Hendricks's avatar
Paul Hendricks committed
114
    created: u32,
115
    model: String,
116
    /// Optional system fingerprint for version tracking.
117
    system_fingerprint: Option<String>,
118
    /// Optional service tier information for the response.
119
    service_tier: Option<dynamo_async_openai::types::ServiceTierResponse>,
120
    /// Tracks token usage for the completion request.
121
    usage: dynamo_async_openai::types::CompletionUsage,
122
    /// Counter tracking the number of messages issued.
123
    msg_counter: u64,
124
    /// Configuration options for response generation.
125
    options: DeltaGeneratorOptions,
126
127
    /// Optional request tracker for per-request metrics (shared with PreprocessedRequest).
    tracker: Option<Arc<RequestTracker>>,
128
129
130
}

impl DeltaGenerator {
131
132
133
134
135
    /// Creates a new [`DeltaGenerator`] instance with the specified model and options.
    ///
    /// # Arguments
    /// * `model` - The model name used for response generation.
    /// * `options` - Configuration options for enabling usage and log probabilities.
136
    /// * `request_id` - The request ID to use for the chat completion response.
137
138
139
    ///
    /// # Returns
    /// * A new instance of [`DeltaGenerator`].
140
    pub fn new(model: String, options: DeltaGeneratorOptions, request_id: String) -> Self {
141
142
143
        let now = std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .unwrap()
144
145
146
147
148
            .as_secs();

        // SAFETY: Casting from `u64` to `u32` could lead to precision loss after `u32::MAX`,
        // but this will not be an issue until 2106.
        let now: u32 = now.try_into().expect("timestamp exceeds u32::MAX");
Paul Hendricks's avatar
Paul Hendricks committed
149

150
        let usage = dynamo_async_openai::types::CompletionUsage {
Paul Hendricks's avatar
Paul Hendricks committed
151
152
153
154
155
156
            prompt_tokens: 0,
            completion_tokens: 0,
            total_tokens: 0,
            prompt_tokens_details: None,
            completion_tokens_details: None,
        };
157

158
159
        let chatcmpl_id = format!("chatcmpl-{request_id}");

160
161
162
        // Always create request tracker for per-worker metrics (TTFT, ITL per worker_id).
        // The enable_tracking option only controls whether timing info is included in the response.
        let tracker = Some(Arc::new(RequestTracker::new()));
163

164
        Self {
165
            id: chatcmpl_id,
166
167
168
169
170
            object: "chat.completion.chunk".to_string(),
            created: now,
            model,
            system_fingerprint: None,
            service_tier: None,
Paul Hendricks's avatar
Paul Hendricks committed
171
            usage,
172
173
            msg_counter: 0,
            options,
174
            tracker,
175
176
177
        }
    }

178
179
180
181
182
    /// Returns the request tracker if tracking is enabled, for sharing with PreprocessedRequest.
    pub fn tracker(&self) -> Option<Arc<RequestTracker>> {
        self.tracker.clone()
    }

183
184
185
186
    /// Updates the prompt token usage count.
    ///
    /// # Arguments
    /// * `isl` - The number of prompt tokens used.
Paul Hendricks's avatar
Paul Hendricks committed
187
    pub fn update_isl(&mut self, isl: u32) {
188
189
190
        self.usage.prompt_tokens = isl;
    }

Greg Clark's avatar
Greg Clark committed
191
192
193
    pub fn create_logprobs(
        &self,
        tokens: Vec<common::llm_backend::TokenType>,
194
        token_ids: &[TokenIdType],
Greg Clark's avatar
Greg Clark committed
195
196
        logprobs: Option<common::llm_backend::LogProbs>,
        top_logprobs: Option<common::llm_backend::TopLogprobs>,
197
    ) -> Option<dynamo_async_openai::types::ChatChoiceLogprobs> {
Greg Clark's avatar
Greg Clark committed
198
199
200
201
202
203
204
        if !self.options.enable_logprobs || logprobs.is_none() {
            return None;
        }

        let toks = tokens
            .into_iter()
            .zip(token_ids)
205
            .map(|(token, token_id)| (token.unwrap_or_default(), *token_id))
Greg Clark's avatar
Greg Clark committed
206
207
208
209
210
211
212
213
214
215
216
217
            .collect::<Vec<(String, TokenIdType)>>();
        let tok_lps = toks
            .iter()
            .zip(logprobs.unwrap())
            .map(|(_, lp)| lp as f32)
            .collect::<Vec<f32>>();

        let content = top_logprobs.map(|top_logprobs| {
            toks.iter()
                .zip(tok_lps)
                .zip(top_logprobs)
                .map(|(((t, tid), lp), top_lps)| {
218
                    let converted = convert_backend_top_logprobs(&top_lps, t, *tid, lp);
219
                    dynamo_async_openai::types::ChatCompletionTokenLogprob {
Greg Clark's avatar
Greg Clark committed
220
221
                        token: t.clone(),
                        logprob: lp,
222
223
                        bytes: token_to_utf8_bytes(t),
                        top_logprobs: converted,
Greg Clark's avatar
Greg Clark committed
224
225
226
227
228
                    }
                })
                .collect()
        });

229
        Some(dynamo_async_openai::types::ChatChoiceLogprobs {
Greg Clark's avatar
Greg Clark committed
230
231
232
233
234
            content,
            refusal: None,
        })
    }

235
236
237
238
239
240
241
    /// Creates a choice within a chat completion response.
    ///
    /// # Arguments
    /// * `index` - The index of the choice in the completion response.
    /// * `text` - The text content for the response.
    /// * `finish_reason` - The reason why the response finished (e.g., stop, length, etc.).
    /// * `logprobs` - Optional log probabilities of the generated tokens.
242
    /// * `stop_reason` - Optional stop string or token that triggered the stop.
243
244
    ///
    /// # Returns
245
    /// * An [`dynamo_async_openai::types::CreateChatCompletionStreamResponse`] instance representing the choice.
Paul Hendricks's avatar
Paul Hendricks committed
246
    #[allow(deprecated)]
247
    pub fn create_choice(
248
        &mut self,
Paul Hendricks's avatar
Paul Hendricks committed
249
        index: u32,
250
        text: Option<String>,
251
252
        finish_reason: Option<dynamo_async_openai::types::FinishReason>,
        logprobs: Option<dynamo_async_openai::types::ChatChoiceLogprobs>,
253
        stop_reason: Option<dynamo_async_openai::types::StopReason>,
254
    ) -> NvCreateChatCompletionStreamResponse {
255
        let delta = dynamo_async_openai::types::ChatCompletionStreamResponseDelta {
256
            content: text.map(dynamo_async_openai::types::ChatCompletionMessageContent::Text),
257
258
            function_call: None,
            tool_calls: None,
259
            role: if self.msg_counter == 0 {
260
                Some(dynamo_async_openai::types::Role::Assistant)
261
262
263
            } else {
                None
            },
Paul Hendricks's avatar
Paul Hendricks committed
264
            refusal: None,
265
            reasoning_content: None,
266
267
        };

268
        let choice = dynamo_async_openai::types::ChatChoiceStream {
Paul Hendricks's avatar
Paul Hendricks committed
269
270
271
            index,
            delta,
            finish_reason,
272
            stop_reason,
Paul Hendricks's avatar
Paul Hendricks committed
273
274
275
276
277
            logprobs,
        };

        let choices = vec![choice];

278
279
280
        // According to OpenAI spec: when stream_options.include_usage is true,
        // all intermediate chunks should have usage: null
        // The final usage chunk will be sent separately with empty choices
281
282
283
284
285
286
287
288
289
290
291
292
293
294
        NvCreateChatCompletionStreamResponse {
            inner: dynamo_async_openai::types::CreateChatCompletionStreamResponse {
                id: self.id.clone(),
                object: self.object.clone(),
                created: self.created,
                model: self.model.clone(),
                system_fingerprint: self.system_fingerprint.clone(),
                choices,
                usage: if self.options.enable_usage && self.options.continuous_usage_stats {
                    Some(self.get_usage())
                } else {
                    None
                },
                service_tier: self.service_tier.clone(),
295
            },
296
            nvext: None, // Will be populated by router layer if needed
297
        }
298
299
300
301
302
303
304
305
    }

    /// Creates a final usage-only chunk for OpenAI compliance.
    /// This should be sent after the last content chunk when stream_options.include_usage is true.
    ///
    /// # Returns
    /// * A [`CreateChatCompletionStreamResponse`] with empty choices and usage stats.
    pub fn create_usage_chunk(&self) -> NvCreateChatCompletionStreamResponse {
306
        let usage = self.get_usage();
307

308
309
310
311
312
313
314
315
316
317
318
        NvCreateChatCompletionStreamResponse {
            inner: dynamo_async_openai::types::CreateChatCompletionStreamResponse {
                id: self.id.clone(),
                object: self.object.clone(),
                created: self.created,
                model: self.model.clone(),
                system_fingerprint: self.system_fingerprint.clone(),
                choices: vec![], // Empty choices for usage-only chunk
                usage: Some(usage),
                service_tier: self.service_tier.clone(),
            },
319
            nvext: None,
320
321
        }
    }
322
323
324
325
326

    /// Check if usage tracking is enabled
    pub fn is_usage_enabled(&self) -> bool {
        self.options.enable_usage
    }
327

328
329
330
331
332
    /// Check if continuous usage tracking is enabled
    pub fn is_continuous_usage_enabled(&self) -> bool {
        self.options.continuous_usage_stats
    }

333
334
335
336
337
    pub fn get_usage(&self) -> dynamo_async_openai::types::CompletionUsage {
        let mut usage = self.usage.clone();
        usage.total_tokens = usage.prompt_tokens.saturating_add(usage.completion_tokens);
        usage
    }
338
339
}

340
/// Implements the [`crate::protocols::openai::DeltaGeneratorExt`] trait for [`DeltaGenerator`], allowing
341
/// it to transform backend responses into OpenAI-style streaming responses.
342
343
344
impl crate::protocols::openai::DeltaGeneratorExt<NvCreateChatCompletionStreamResponse>
    for DeltaGenerator
{
345
346
347
348
349
350
351
352
    /// Converts a backend response into a structured OpenAI-style streaming response.
    ///
    /// # Arguments
    /// * `delta` - The backend response containing generated text and metadata.
    ///
    /// # Returns
    /// * `Ok(NvCreateChatCompletionStreamResponse)` if conversion succeeds.
    /// * `Err(anyhow::Error)` if an error occurs.
353
354
355
    fn choice_from_postprocessor(
        &mut self,
        delta: crate::protocols::common::llm_backend::BackendOutput,
356
    ) -> anyhow::Result<NvCreateChatCompletionStreamResponse> {
357
358
359
360
361
362
363
364
365
366
367
        // Aggregate token usage even if usage tracking is disabled for metrics tracking
        // SAFETY: Casting from `usize` to `u32` could lead to precision loss after `u32::MAX`,
        // but this will not be an issue until context lengths exceed 4_294_967_295.
        let token_length: u32 = delta
            .token_ids
            .len()
            .try_into()
            .expect("token_ids length exceeds u32::MAX");

        self.usage.completion_tokens += token_length;

368
369
370
371
372
373
374
375
376
377
378
        // If backend provides completion_usage, use it to update usage stats
        // This is critical for prompt embeddings where prompt_tokens comes from
        // the embedding sequence length computed by the worker
        if let Some(completion_usage) = delta.completion_usage.as_ref() {
            // Update prompt_tokens from worker if provided (e.g., for embeddings)
            self.usage.prompt_tokens = completion_usage.prompt_tokens;

            // Propagate prompt token details if provided
            if let Some(prompt_details) = completion_usage.prompt_tokens_details.as_ref() {
                self.usage.prompt_tokens_details = Some(prompt_details.clone());
            }
379
380
        }

Greg Clark's avatar
Greg Clark committed
381
382
        let logprobs = self.create_logprobs(
            delta.tokens,
383
            &delta.token_ids,
Greg Clark's avatar
Greg Clark committed
384
385
386
            delta.log_probs,
            delta.top_logprobs,
        );
387

388
        // Map backend finish reasons to OpenAI's finish reasons.
389
        let finish_reason = match delta.finish_reason {
390
391
392
393
394
395
396
397
398
399
            Some(common::FinishReason::EoS) => Some(dynamo_async_openai::types::FinishReason::Stop),
            Some(common::FinishReason::Stop) => {
                Some(dynamo_async_openai::types::FinishReason::Stop)
            }
            Some(common::FinishReason::Length) => {
                Some(dynamo_async_openai::types::FinishReason::Length)
            }
            Some(common::FinishReason::Cancelled) => {
                Some(dynamo_async_openai::types::FinishReason::Stop)
            }
400
            Some(common::FinishReason::ContentFilter) => {
401
                Some(dynamo_async_openai::types::FinishReason::ContentFilter)
402
            }
403
404
405
406
407
408
            Some(common::FinishReason::Error(err_msg)) => {
                return Err(anyhow::anyhow!(err_msg));
            }
            None => None,
        };

409
        // Create the streaming response.
410
        let index = 0;
411
412
413
414
415
416
417
        let mut stream_response = self.create_choice(
            index,
            delta.text,
            finish_reason,
            logprobs,
            delta.stop_reason,
        );
418

419
420
        // Get worker_id info from tracker (set by KvPushRouter based on phase)
        let worker_id_info = self.tracker.as_ref().and_then(|t| t.get_worker_info());
421

422
423
424
425
426
        let token_ids = delta
            .disaggregated_params
            .as_ref()
            .and_then(|params| params.get("token_ids"))
            .and_then(|v| serde_json::from_value::<Vec<u32>>(v.clone()).ok());
427
428
429
430
431
        let routed_experts = delta
            .disaggregated_params
            .as_ref()
            .and_then(|params| params.get("routed_experts"))
            .cloned();
432

433
434
        // Get timing info if this is the final response (has finish_reason)
        let timing_info: Option<TimingInfo> = if finish_reason.is_some() {
435
            self.tracker.as_ref().map(|tracker| {
436
437
438
439
440
441
442
                tracker.record_finish();
                tracker.get_timing_info()
            })
        } else {
            None
        };

443
444
445
446
447
448
        // Inject nvext if we have worker_id, token_ids, timing, or routed experts.
        if worker_id_info.is_some()
            || token_ids.is_some()
            || timing_info.is_some()
            || routed_experts.is_some()
        {
449
            let nvext_response = NvExtResponse {
450
451
                worker_id: worker_id_info.clone(),
                timing: timing_info,
452
                token_ids: token_ids.clone(),
453
                routed_experts,
454
455
456
457
            };

            if let Ok(nvext_json) = serde_json::to_value(&nvext_response) {
                stream_response.nvext = Some(nvext_json);
458
459
460
461
462
463
464
                if let Some(ref info) = worker_id_info {
                    tracing::debug!(
                        "Injected worker_id into chat completion nvext: prefill={:?}, decode={:?}",
                        info.prefill_worker_id,
                        info.decode_worker_id
                    );
                }
465
466
467
468
469
470
                if let Some(ref tokens) = token_ids {
                    tracing::debug!(
                        "Injected token_ids into chat completion nvext: {} tokens",
                        tokens.len()
                    );
                }
471
472
            }
        }
Paul Hendricks's avatar
Paul Hendricks committed
473

474
        Ok(stream_response)
475
    }
476
477
478
479

    fn get_isl(&self) -> Option<u32> {
        Some(self.usage.prompt_tokens)
    }
480
481
482
483
484
485
486
487

    fn create_usage_chunk(&self) -> NvCreateChatCompletionStreamResponse {
        DeltaGenerator::create_usage_chunk(self)
    }

    fn is_usage_enabled(&self) -> bool {
        DeltaGenerator::is_usage_enabled(self)
    }
488

489
490
491
492
    fn is_continuous_usage_enabled(&self) -> bool {
        DeltaGenerator::is_continuous_usage_enabled(self)
    }

493
494
495
    fn get_usage(&self) -> dynamo_async_openai::types::CompletionUsage {
        DeltaGenerator::get_usage(self)
    }
496
497
498
499

    fn tracker(&self) -> Option<std::sync::Arc<crate::protocols::common::timing::RequestTracker>> {
        self.tracker.clone()
    }
500
}
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
527
528

#[cfg(test)]
mod tests {
    use super::*;
    use dynamo_async_openai::types::{
        ChatCompletionRequestMessage, ChatCompletionRequestUserMessage,
        ChatCompletionRequestUserMessageContent, CreateChatCompletionRequest,
    };

    fn create_test_request() -> NvCreateChatCompletionRequest {
        let messages = vec![ChatCompletionRequestMessage::User(
            ChatCompletionRequestUserMessage {
                content: ChatCompletionRequestUserMessageContent::Text("test".to_string()),
                name: None,
            },
        )];

        NvCreateChatCompletionRequest {
            inner: CreateChatCompletionRequest {
                model: "test-model".to_string(),
                messages,
                stream: Some(false),
                stream_options: None,
                ..Default::default()
            },
            common: Default::default(),
            nvext: None,
            chat_template_args: None,
529
            media_io_kwargs: None,
530
            unsupported_fields: Default::default(),
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
        }
    }

    #[test]
    fn test_enable_usage_for_nonstreaming_enables_usage() {
        // Test that non-streaming requests get usage enabled
        let mut request = create_test_request();
        assert!(request.inner.stream_options.is_none());

        request.enable_usage_for_nonstreaming(false); // false = non-streaming

        assert!(
            request.inner.stream_options.is_some(),
            "Non-streaming request should have stream_options created"
        );
        assert!(
            request.inner.stream_options.unwrap().include_usage,
            "Non-streaming request should have include_usage=true for OpenAI compliance"
        );
550
551
552
553
        assert!(
            !request.inner.stream_options.unwrap().continuous_usage_stats,
            "Non-streaming request should have continuous_usage_stats=false for OpenAI compliance"
        );
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
    }

    #[test]
    fn test_enable_usage_for_nonstreaming_ignores_streaming() {
        // Test that streaming requests are not modified
        let mut request = create_test_request();
        assert!(request.inner.stream_options.is_none());

        request.enable_usage_for_nonstreaming(true); // true = streaming

        assert!(
            request.inner.stream_options.is_none(),
            "Streaming request should not have stream_options modified"
        );
    }
}