chat_completions.rs 15 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 dynamo_runtime::protocols::annotated::AnnotationsProvider;
use serde::{Deserialize, Serialize};
6
use utoipa::ToSchema;
7
8
use validator::Validate;

9
use crate::engines::ValidateRequest;
10
use crate::preprocessor::media::MediaDecoder;
11
12

use super::{
13
    OpenAIOutputOptionsProvider, OpenAISamplingOptionsProvider, OpenAIStopConditionsProvider,
14
    common_ext::{CommonExt, CommonExtProvider},
15
16
    nvext::NvExt,
    nvext::NvExtProvider,
17
    tools, validate,
18
};
19

20
pub mod aggregator;
21
mod delta;
Ryan Olson's avatar
Ryan Olson committed
22
pub mod jail;
23

Paul Hendricks's avatar
Paul Hendricks committed
24
pub use aggregator::DeltaAggregator;
25
26
pub use delta::DeltaGenerator;

27
/// A request structure for creating a chat completion, extending OpenAI's
28
/// `CreateChatCompletionRequest` with [`NvExt`] extensions and common fields.
29
30
31
///
/// # Fields
/// - `inner`: The base OpenAI chat completion request, embedded using `serde(flatten)`.
32
33
34
/// - `common`: Common extension fields (ignore_eos, min_tokens) at root level, embedded using `serde(flatten)`.
/// - `nvext`: The optional NVIDIA extension field. See [`NvExt`] for more details.
///   Note: If ignore_eos is specified in both common and nvext, the common (root-level) value takes precedence.
35
#[derive(ToSchema, Serialize, Deserialize, Validate, Debug, Clone)]
36
pub struct NvCreateChatCompletionRequest {
Paul Hendricks's avatar
Paul Hendricks committed
37
    #[serde(flatten)]
38
    #[schema(value_type = Object)]
39
    pub inner: dynamo_protocols::types::CreateChatCompletionRequest,
40

41
42
43
    #[serde(flatten, default)]
    pub common: CommonExt,

44
    #[serde(skip_serializing_if = "Option::is_none")]
45
    pub nvext: Option<NvExt>,
46
47

    /// Extra args to pass to the chat template rendering context
48
49
50
51
52
53
    /// Also accepts "chat_template_kwargs" as an alias for compatibility
    #[serde(
        default,
        skip_serializing_if = "Option::is_none",
        alias = "chat_template_kwargs"
    )]
54
    pub chat_template_args: Option<std::collections::HashMap<String, serde_json::Value>>,
55

56
57
58
59
60
61
    /// Runtime media decoding parameters.
    /// When provided, these override the MDC defaults
    /// Example: `{"video": {"num_frames": 16}}`
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub media_io_kwargs: Option<MediaDecoder>,

62
63
64
    /// Catch-all for unsupported fields - checked during validation
    #[serde(flatten, default, skip_serializing)]
    pub unsupported_fields: std::collections::HashMap<String, serde_json::Value>,
65
66
}

67
/// A response structure for unary chat completion responses, embedding OpenAI's
68
69
70
71
/// `CreateChatCompletionResponse` with optional NVIDIA extension metadata.
#[derive(Serialize, Deserialize, Debug, Clone, PartialEq)]
pub struct NvCreateChatCompletionResponse {
    #[serde(flatten)]
72
    pub inner: dynamo_protocols::types::CreateChatCompletionResponse,
73
74
75
    #[serde(skip_serializing_if = "Option::is_none")]
    pub nvext: Option<serde_json::Value>,
}
76

77
/// A response structure for streamed chat completions, embedding OpenAI's
78
79
80
81
/// `CreateChatCompletionStreamResponse` with optional NVIDIA extension metadata.
#[derive(Serialize, Deserialize, Debug, Clone, PartialEq)]
pub struct NvCreateChatCompletionStreamResponse {
    #[serde(flatten)]
82
    pub inner: dynamo_protocols::types::CreateChatCompletionStreamResponse,
83
84
85
    #[serde(skip_serializing_if = "Option::is_none")]
    pub nvext: Option<serde_json::Value>,
}
86

87
88
/// Implements `NvExtProvider` for `NvCreateChatCompletionRequest`,
/// providing access to NVIDIA-specific extensions.
89
impl NvExtProvider for NvCreateChatCompletionRequest {
90
    /// Returns a reference to the optional `NvExt` extension, if available.
91
92
93
94
    fn nvext(&self) -> Option<&NvExt> {
        self.nvext.as_ref()
    }

95
    /// Returns `None`, as raw prompt extraction is not implemented.
96
97
98
99
100
    fn raw_prompt(&self) -> Option<String> {
        None
    }
}

101
102
/// Implements `AnnotationsProvider` for `NvCreateChatCompletionRequest`,
/// enabling retrieval and management of request annotations.
103
impl AnnotationsProvider for NvCreateChatCompletionRequest {
104
    /// Retrieves the list of annotations from `NvExt`, if present.
Biswa Panda's avatar
Biswa Panda committed
105
106
107
108
109
110
    fn annotations(&self) -> Option<Vec<String>> {
        self.nvext
            .as_ref()
            .and_then(|nvext| nvext.annotations.clone())
    }

111
112
113
114
115
116
117
    /// Checks whether a specific annotation exists in the request.
    ///
    /// # Arguments
    /// * `annotation` - A string slice representing the annotation to check.
    ///
    /// # Returns
    /// `true` if the annotation exists, `false` otherwise.
Biswa Panda's avatar
Biswa Panda committed
118
119
120
121
122
123
124
125
    fn has_annotation(&self, annotation: &str) -> bool {
        self.nvext
            .as_ref()
            .and_then(|nvext| nvext.annotations.as_ref())
            .map(|annotations| annotations.contains(&annotation.to_string()))
            .unwrap_or(false)
    }
}
126

127
128
/// Implements `OpenAISamplingOptionsProvider` for `NvCreateChatCompletionRequest`,
/// exposing OpenAI's sampling parameters for chat completion.
129
impl OpenAISamplingOptionsProvider for NvCreateChatCompletionRequest {
130
    /// Retrieves the temperature parameter for sampling, if set.
131
    fn get_temperature(&self) -> Option<f32> {
Paul Hendricks's avatar
Paul Hendricks committed
132
        self.inner.temperature
133
134
    }

135
    /// Retrieves the top-p (nucleus sampling) parameter, if set.
136
    fn get_top_p(&self) -> Option<f32> {
Paul Hendricks's avatar
Paul Hendricks committed
137
        self.inner.top_p
138
139
    }

140
    /// Retrieves the frequency penalty parameter, if set.
141
    fn get_frequency_penalty(&self) -> Option<f32> {
Paul Hendricks's avatar
Paul Hendricks committed
142
        self.inner.frequency_penalty
143
144
    }

145
    /// Retrieves the presence penalty parameter, if set.
146
    fn get_presence_penalty(&self) -> Option<f32> {
Paul Hendricks's avatar
Paul Hendricks committed
147
        self.inner.presence_penalty
148
149
    }

150
    /// Returns a reference to the optional `NvExt` extension, if available.
151
152
153
    fn nvext(&self) -> Option<&NvExt> {
        self.nvext.as_ref()
    }
154
155
156
157
158
159
160
161
162
163
164
165
166
167
    /// Retrieves the seed value for random number generation, if set.
    fn get_seed(&self) -> Option<i64> {
        self.inner.seed
    }

    /// Retrieves the number of completions to generate for each prompt, if set.
    fn get_n(&self) -> Option<u8> {
        self.inner.n
    }

    /// Retrieves the best_of parameter, if set.
    fn get_best_of(&self) -> Option<u8> {
        None // Not supported in chat completions
    }
168
169
}

170
171
172
173
174
175
176
177
178
/// Implements `CommonExtProvider` for `NvCreateChatCompletionRequest`,
/// providing access to common extension fields.
impl CommonExtProvider for NvCreateChatCompletionRequest {
    /// Returns a reference to the CommonExt struct.
    fn common_ext(&self) -> Option<&CommonExt> {
        Some(&self.common)
    }

    /// Guided Decoding Options
179
180
181
182
183
    fn get_guided_json(&self) -> Option<serde_json::Value> {
        if let Some(value) = self.common.guided_json.clone() {
            return Some(value);
        }

184
185
186
187
188
189
190
191
192
193
194
195
196
        // 1) Tool-call guided decoding (highest precedence after explicit guided_json)
        if let (Some(tool_choice), Some(tools)) =
            (self.inner.tool_choice.as_ref(), self.inner.tools.as_deref())
        {
            match tools::get_json_schema_from_tools(Some(tool_choice), Some(tools)) {
                Ok(Some(schema)) => return Some(schema),
                Ok(None) => {}
                Err(err) => {
                    tracing::warn!(
                        error = %err,
                        "failed to derive guided_json from tool_choice"
                    );
                }
197
198
            }
        }
199
200
201

        // 2) OpenAI `response_format` (applies to assistant content, not tool calls)
        if let Some(response_format) = self.inner.response_format.as_ref() {
202
            use dynamo_protocols::types::ResponseFormat;
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
            match response_format {
                ResponseFormat::Text => {}
                ResponseFormat::JsonObject => {
                    // Minimal JSON Schema for "any JSON object"
                    return Some(serde_json::json!({
                        "type": "object"
                    }));
                }
                ResponseFormat::JsonSchema { json_schema } => {
                    // validate_response_format ensures schema is present when type=json_schema
                    if let Some(schema) = json_schema.schema.clone() {
                        return Some(schema);
                    }
                }
            }
        }

        None
221
222
223
    }

    fn get_guided_regex(&self) -> Option<String> {
224
        self.common.guided_regex.clone()
225
226
227
    }

    fn get_guided_grammar(&self) -> Option<String> {
228
        self.common.guided_grammar.clone()
229
230
231
    }

    fn get_guided_choice(&self) -> Option<Vec<String>> {
232
        self.common.guided_choice.clone()
233
234
235
    }

    fn get_guided_decoding_backend(&self) -> Option<String> {
236
        self.common.guided_decoding_backend.clone()
237
    }
238

239
    fn get_guided_whitespace_pattern(&self) -> Option<String> {
240
        self.common.guided_whitespace_pattern.clone()
241
242
    }

243
    fn get_top_k(&self) -> Option<i32> {
244
        self.common.top_k
245
246
    }

247
    fn get_min_p(&self) -> Option<f32> {
248
        self.common.min_p
249
250
    }

251
    fn get_repetition_penalty(&self) -> Option<f32> {
252
        self.common.repetition_penalty
253
    }
254
255
256
257

    fn get_include_stop_str_in_output(&self) -> Option<bool> {
        self.common.include_stop_str_in_output
    }
258
259
260
261

    fn get_skip_special_tokens(&self) -> Option<bool> {
        self.common.skip_special_tokens
    }
262
263
}

264
265
/// Implements `OpenAIStopConditionsProvider` for `NvCreateChatCompletionRequest`,
/// providing access to stop conditions that control chat completion behavior.
266
impl OpenAIStopConditionsProvider for NvCreateChatCompletionRequest {
267
    /// Retrieves the maximum number of tokens allowed in the response.
268
    #[allow(deprecated)]
Paul Hendricks's avatar
Paul Hendricks committed
269
    fn get_max_tokens(&self) -> Option<u32> {
270
        self.inner.max_completion_tokens.or(self.inner.max_tokens)
271
272
    }

273
    /// Retrieves the minimum number of tokens required in the response.
274
275
    /// Returns `min_tokens` Value
    /// `min_tokens` is not an OpenAI-supported parameter.
Paul Hendricks's avatar
Paul Hendricks committed
276
    fn get_min_tokens(&self) -> Option<u32> {
277
        self.common.min_tokens
278
279
    }

280
281
282
283
284
285
286
    /// Retrieves the stop conditions that terminate the chat completion response.
    ///
    /// Converts OpenAI's `Stop` enum to a `Vec<String>`, normalizing the representation.
    ///
    /// # Returns
    /// * `Some(Vec<String>)` if stop conditions are set.
    /// * `None` if no stop conditions are defined.
287
    fn get_stop(&self) -> Option<Vec<String>> {
288
        self.inner.stop.as_ref().map(|stop| match stop {
289
290
            dynamo_protocols::types::Stop::String(s) => vec![s.clone()],
            dynamo_protocols::types::Stop::StringArray(arr) => arr.clone(),
291
        })
292
293
    }

294
    /// Returns a reference to the optional `NvExt` extension, if available.
295
296
297
    fn nvext(&self) -> Option<&NvExt> {
        self.nvext.as_ref()
    }
298
299
300
301
302

    /// Get ignore_eos from CommonExt.
    fn get_common_ignore_eos(&self) -> Option<bool> {
        self.common.ignore_eos
    }
303

304
    /// Get the effective ignore_eos value from CommonExt.
305
    fn get_ignore_eos(&self) -> Option<bool> {
306
        self.common.ignore_eos
307
    }
308
}
309

Greg Clark's avatar
Greg Clark committed
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
impl OpenAIOutputOptionsProvider for NvCreateChatCompletionRequest {
    fn get_logprobs(&self) -> Option<u32> {
        match self.inner.logprobs {
            Some(true) => match self.inner.top_logprobs {
                Some(top_logprobs) => Some(top_logprobs as u32),
                None => Some(1_u32),
            },
            Some(false) => None,
            None => None,
        }
    }

    fn get_prompt_logprobs(&self) -> Option<u32> {
        None
    }

    fn get_skip_special_tokens(&self) -> Option<bool> {
327
        CommonExtProvider::get_skip_special_tokens(self)
Greg Clark's avatar
Greg Clark committed
328
329
330
331
332
333
334
    }

    fn get_formatted_prompt(&self) -> Option<bool> {
        None
    }
}

335
336
337
338
/// Implements `ValidateRequest` for `NvCreateChatCompletionRequest`,
/// allowing us to validate the data.
impl ValidateRequest for NvCreateChatCompletionRequest {
    fn validate(&self) -> Result<(), anyhow::Error> {
339
        validate::validate_no_unsupported_fields(&self.unsupported_fields)?;
340
341
342
343
        validate::validate_messages(&self.inner.messages)?;
        validate::validate_model(&self.inner.model)?;
        // none for store
        validate::validate_reasoning_effort(&self.inner.reasoning_effort)?;
344
        // none for metadata
345
346
347
348
349
350
351
352
353
354
355
        validate::validate_frequency_penalty(self.inner.frequency_penalty)?;
        validate::validate_logit_bias(&self.inner.logit_bias)?;
        // none for logprobs
        validate::validate_top_logprobs(self.inner.top_logprobs)?;
        // validate::validate_max_tokens(self.inner.max_tokens)?; // warning depricated field
        validate::validate_max_completion_tokens(self.inner.max_completion_tokens)?;
        validate::validate_n(self.inner.n)?;
        // none for modalities
        // none for prediction
        // none for audio
        validate::validate_presence_penalty(self.inner.presence_penalty)?;
356
        validate::validate_response_format(&self.inner.response_format)?;
357
358
359
360
361
362
363
364
365
366
367
368
369
        // none for seed
        validate::validate_service_tier(&self.inner.service_tier)?;
        validate::validate_stop(&self.inner.stop)?;
        // none for stream
        // none for stream_options
        validate::validate_temperature(self.inner.temperature)?;
        validate::validate_top_p(self.inner.top_p)?;
        validate::validate_tools(&self.inner.tools.as_deref())?;
        // none for tool_choice
        // none for parallel_tool_calls
        validate::validate_user(self.inner.user.as_deref())?;
        // none for function call
        // none for functions
370
371
        // Common Ext
        validate::validate_repetition_penalty(self.get_repetition_penalty())?;
372
373
        validate::validate_min_p(self.get_min_p())?;
        validate::validate_top_k(self.get_top_k())?;
374
375
        // Cross-field validation
        validate::validate_n_with_temperature(self.inner.n, self.inner.temperature)?;
376
377
378
379

        Ok(())
    }
}
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429

#[cfg(test)]
mod tests {
    use super::*;
    use crate::protocols::common::OutputOptionsProvider;
    use serde_json::json;

    #[test]
    fn test_skip_special_tokens_none() {
        let json_str = json!({
            "model": "test-model",
            "messages": [
                {"role": "user", "content": "Hello"}
            ]
        });

        let request: NvCreateChatCompletionRequest =
            serde_json::from_value(json_str).expect("Failed to deserialize request");

        assert_eq!(request.common.skip_special_tokens, None);

        let output_options = request
            .extract_output_options()
            .expect("Failed to extract output options");

        assert_eq!(output_options.skip_special_tokens, None);
    }

    #[test]
    fn test_skip_special_tokens_propagates() {
        for skip_value in [true, false] {
            let json_str = json!({
                "model": "test-model",
                "messages": [
                    {"role": "user", "content": "Hello"}
                ],
                "skip_special_tokens": skip_value
            });

            let request: NvCreateChatCompletionRequest =
                serde_json::from_value(json_str).expect("Failed to deserialize request");

            let output_options = request
                .extract_output_options()
                .expect("Failed to extract output options");

            assert_eq!(output_options.skip_special_tokens, Some(skip_value));
        }
    }
}