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

use derive_builder::Builder;
5
use dynamo_runtime::protocols::annotated::AnnotationsProvider;
6
7
8
use serde::{Deserialize, Serialize};
use validator::Validate;

9
10
use crate::engines::ValidateRequest;

11
use super::{
12
13
    ContentProvider, OpenAIOutputOptionsProvider, OpenAISamplingOptionsProvider,
    OpenAIStopConditionsProvider,
Greg Clark's avatar
Greg Clark committed
14
    common::{self, OutputOptionsProvider, SamplingOptionsProvider, StopConditionsProvider},
15
    common_ext::{CommonExt, CommonExtProvider},
16
    nvext::{NvExt, NvExtProvider},
17
    validate,
18
19
};

20
21
22
23
24
mod aggregator;
mod delta;

pub use aggregator::DeltaAggregator;
pub use delta::DeltaGenerator;
Biswa Panda's avatar
Biswa Panda committed
25

26
#[derive(Serialize, Deserialize, Validate, Debug, Clone)]
27
pub struct NvCreateCompletionRequest {
28
    #[serde(flatten)]
29
    pub inner: dynamo_async_openai::types::CreateCompletionRequest,
30

31
32
33
    #[serde(flatten)]
    pub common: CommonExt,

34
35
    #[serde(skip_serializing_if = "Option::is_none")]
    pub nvext: Option<NvExt>,
36
37
38
39

    // metadata - passthrough parameter without restrictions
    #[serde(skip_serializing_if = "Option::is_none")]
    pub metadata: Option<serde_json::Value>,
40
41
42
43

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

46
47
48
#[derive(Serialize, Deserialize, Validate, Debug, Clone)]
pub struct NvCreateCompletionResponse {
    #[serde(flatten)]
49
    pub inner: dynamo_async_openai::types::CreateCompletionResponse,
50
51
}

52
impl ContentProvider for dynamo_async_openai::types::Choice {
53
54
55
56
57
    fn content(&self) -> String {
        self.text.clone()
    }
}

58
pub fn prompt_to_string(prompt: &dynamo_async_openai::types::Prompt) -> String {
59
    match prompt {
60
61
62
        dynamo_async_openai::types::Prompt::String(s) => s.clone(),
        dynamo_async_openai::types::Prompt::StringArray(arr) => arr.join(" "), // Join strings with spaces
        dynamo_async_openai::types::Prompt::IntegerArray(arr) => arr
63
64
65
66
            .iter()
            .map(|&num| num.to_string())
            .collect::<Vec<_>>()
            .join(" "),
67
        dynamo_async_openai::types::Prompt::ArrayOfIntegerArray(arr) => arr
68
69
70
71
72
73
74
75
76
77
78
79
80
            .iter()
            .map(|inner| {
                inner
                    .iter()
                    .map(|&num| num.to_string())
                    .collect::<Vec<_>>()
                    .join(" ")
            })
            .collect::<Vec<_>>()
            .join(" | "), // Separate arrays with a delimiter
    }
}

81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
/// Get the batch size from a prompt (1 for single prompts, array length for batch prompts)
pub fn get_prompt_batch_size(prompt: &dynamo_async_openai::types::Prompt) -> usize {
    match prompt {
        dynamo_async_openai::types::Prompt::String(_) => 1,
        dynamo_async_openai::types::Prompt::IntegerArray(_) => 1,
        dynamo_async_openai::types::Prompt::StringArray(arr) => arr.len(),
        dynamo_async_openai::types::Prompt::ArrayOfIntegerArray(arr) => arr.len(),
    }
}

/// Extract a single prompt from a batch at the given index.
/// For single prompts, returns a clone regardless of index.
/// For batch prompts, returns the prompt at the specified index.
pub fn extract_single_prompt(
    prompt: &dynamo_async_openai::types::Prompt,
    index: usize,
) -> dynamo_async_openai::types::Prompt {
    match prompt {
        dynamo_async_openai::types::Prompt::String(s) => {
            dynamo_async_openai::types::Prompt::String(s.clone())
        }
        dynamo_async_openai::types::Prompt::IntegerArray(arr) => {
            dynamo_async_openai::types::Prompt::IntegerArray(arr.clone())
        }
        dynamo_async_openai::types::Prompt::StringArray(arr) => {
            dynamo_async_openai::types::Prompt::String(arr[index].clone())
        }
        dynamo_async_openai::types::Prompt::ArrayOfIntegerArray(arr) => {
            dynamo_async_openai::types::Prompt::IntegerArray(arr[index].clone())
        }
    }
}

114
impl NvExtProvider for NvCreateCompletionRequest {
115
116
117
118
119
    fn nvext(&self) -> Option<&NvExt> {
        self.nvext.as_ref()
    }

    fn raw_prompt(&self) -> Option<String> {
120
121
122
123
124
        if let Some(nvext) = self.nvext.as_ref()
            && let Some(use_raw_prompt) = nvext.use_raw_prompt
            && use_raw_prompt
        {
            return Some(prompt_to_string(&self.inner.prompt));
125
126
127
128
129
        }
        None
    }
}

130
impl AnnotationsProvider for NvCreateCompletionRequest {
Biswa Panda's avatar
Biswa Panda committed
131
132
133
134
135
136
137
138
139
140
141
142
143
144
    fn annotations(&self) -> Option<Vec<String>> {
        self.nvext
            .as_ref()
            .and_then(|nvext| nvext.annotations.clone())
    }

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

146
impl OpenAISamplingOptionsProvider for NvCreateCompletionRequest {
147
    fn get_temperature(&self) -> Option<f32> {
148
        self.inner.temperature
149
150
151
    }

    fn get_top_p(&self) -> Option<f32> {
152
        self.inner.top_p
153
154
155
    }

    fn get_frequency_penalty(&self) -> Option<f32> {
156
        self.inner.frequency_penalty
157
158
159
    }

    fn get_presence_penalty(&self) -> Option<f32> {
160
        self.inner.presence_penalty
161
162
163
164
165
    }

    fn nvext(&self) -> Option<&NvExt> {
        self.nvext.as_ref()
    }
166
167
168
169
170
171
172
173
174
175
176
177

    fn get_seed(&self) -> Option<i64> {
        self.inner.seed
    }

    fn get_n(&self) -> Option<u8> {
        self.inner.n
    }

    fn get_best_of(&self) -> Option<u8> {
        self.inner.best_of
    }
178
179
}

180
181
182
183
184
185
186
impl CommonExtProvider for NvCreateCompletionRequest {
    fn common_ext(&self) -> Option<&CommonExt> {
        Some(&self.common)
    }

    /// Guided Decoding Options
    fn get_guided_json(&self) -> Option<&serde_json::Value> {
187
        self.common.guided_json.as_ref()
188
189
190
    }

    fn get_guided_regex(&self) -> Option<String> {
191
        self.common.guided_regex.clone()
192
193
194
    }

    fn get_guided_grammar(&self) -> Option<String> {
195
        self.common.guided_grammar.clone()
196
197
198
    }

    fn get_guided_choice(&self) -> Option<Vec<String>> {
199
        self.common.guided_choice.clone()
200
201
202
    }

    fn get_guided_decoding_backend(&self) -> Option<String> {
203
        self.common.guided_decoding_backend.clone()
204
    }
205

206
    fn get_guided_whitespace_pattern(&self) -> Option<String> {
207
        self.common.guided_whitespace_pattern.clone()
208
209
    }

210
    fn get_top_k(&self) -> Option<i32> {
211
        self.common.top_k
212
213
    }

214
    fn get_min_p(&self) -> Option<f32> {
215
        self.common.min_p
216
217
    }

218
    fn get_repetition_penalty(&self) -> Option<f32> {
219
        self.common.repetition_penalty
220
    }
221
222
223
224

    fn get_include_stop_str_in_output(&self) -> Option<bool> {
        self.common.include_stop_str_in_output
    }
225
226
227
228

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

231
impl OpenAIStopConditionsProvider for NvCreateCompletionRequest {
Paul Hendricks's avatar
Paul Hendricks committed
232
    fn get_max_tokens(&self) -> Option<u32> {
233
        self.inner.max_tokens
234
235
    }

Paul Hendricks's avatar
Paul Hendricks committed
236
    fn get_min_tokens(&self) -> Option<u32> {
237
        self.common.min_tokens
238
239
240
    }

    fn get_stop(&self) -> Option<Vec<String>> {
241
        None
242
243
244
245
246
    }

    fn nvext(&self) -> Option<&NvExt> {
        self.nvext.as_ref()
    }
247
248
249
250

    fn get_common_ignore_eos(&self) -> Option<bool> {
        self.common.ignore_eos
    }
251

252
    /// Get the effective ignore_eos value from CommonExt.
253
    fn get_ignore_eos(&self) -> Option<bool> {
254
        self.common.ignore_eos
255
    }
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
}

#[derive(Builder)]
pub struct ResponseFactory {
    #[builder(setter(into))]
    pub model: String,

    #[builder(default)]
    pub system_fingerprint: Option<String>,

    #[builder(default = "format!(\"cmpl-{}\", uuid::Uuid::new_v4())")]
    pub id: String,

    #[builder(default = "\"text_completion\".to_string()")]
    pub object: String,

272
273
    #[builder(default = "chrono::Utc::now().timestamp() as u32")]
    pub created: u32,
274
275
276
277
278
279
280
281
282
}

impl ResponseFactory {
    pub fn builder() -> ResponseFactoryBuilder {
        ResponseFactoryBuilder::default()
    }

    pub fn make_response(
        &self,
283
284
        choice: dynamo_async_openai::types::Choice,
        usage: Option<dynamo_async_openai::types::CompletionUsage>,
285
    ) -> NvCreateCompletionResponse {
286
        let inner = dynamo_async_openai::types::CreateCompletionResponse {
287
288
            id: self.id.clone(),
            object: self.object.clone(),
289
            created: self.created,
290
291
292
293
            model: self.model.clone(),
            choices: vec![choice],
            system_fingerprint: self.system_fingerprint.clone(),
            usage,
294
295
        };
        NvCreateCompletionResponse { inner }
296
297
298
299
    }
}

/// Implements TryFrom for converting an OpenAI's CompletionRequest to an Engine's CompletionRequest
300
impl TryFrom<NvCreateCompletionRequest> for common::CompletionRequest {
301
302
    type Error = anyhow::Error;

303
    fn try_from(request: NvCreateCompletionRequest) -> Result<Self, Self::Error> {
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
        // openai_api_rs::v1::completion::CompletionRequest {
        // NA  pub model: String,
        //     pub prompt: String,
        // **  pub suffix: Option<String>,
        //     pub max_tokens: Option<i32>,
        //     pub temperature: Option<f32>,
        //     pub top_p: Option<f32>,
        //     pub n: Option<i32>,
        //     pub stream: Option<bool>,
        //     pub logprobs: Option<i32>,
        //     pub echo: Option<bool>,
        //     pub stop: Option<Vec<String, Global>>,
        //     pub presence_penalty: Option<f32>,
        //     pub frequency_penalty: Option<f32>,
        //     pub best_of: Option<i32>,
        //     pub logit_bias: Option<HashMap<String, i32, RandomState>>,
        //     pub user: Option<String>,
        // }
        //
        // ** no supported

325
        if request.inner.suffix.is_some() {
326
327
328
329
330
331
332
333
334
335
336
            return Err(anyhow::anyhow!("suffix is not supported"));
        }

        let stop_conditions = request
            .extract_stop_conditions()
            .map_err(|e| anyhow::anyhow!("Failed to extract stop conditions: {}", e))?;

        let sampling_options = request
            .extract_sampling_options()
            .map_err(|e| anyhow::anyhow!("Failed to extract sampling options: {}", e))?;

Greg Clark's avatar
Greg Clark committed
337
338
339
340
        let output_options = request
            .extract_output_options()
            .map_err(|e| anyhow::anyhow!("Failed to extract output options: {}", e))?;

341
        let prompt = common::PromptType::Completion(common::CompletionContext {
342
            prompt: prompt_to_string(&request.inner.prompt),
343
344
345
346
347
348
349
            system_prompt: None,
        });

        Ok(common::CompletionRequest {
            prompt,
            stop_conditions,
            sampling_options,
Greg Clark's avatar
Greg Clark committed
350
            output_options,
351
352
353
354
355
356
            mdc_sum: None,
            annotations: None,
        })
    }
}

357
impl TryFrom<common::StreamingCompletionResponse> for dynamo_async_openai::types::Choice {
358
359
360
    type Error = anyhow::Error;

    fn try_from(response: common::StreamingCompletionResponse) -> Result<Self, Self::Error> {
361
362
363
364
365
        let text = response
            .delta
            .text
            .ok_or(anyhow::anyhow!("No text in response"))?;

366
        // SAFETY: we're downcasting from u64 to u32 here but u32::MAX is 4_294_967_295
367
        // so we're fairly safe knowing we won't generate that many Choices
368
369
370
371
372
373
        let index: u32 = response
            .delta
            .index
            .unwrap_or(0)
            .try_into()
            .expect("index exceeds u32::MAX");
374
375
376
377

        // TODO handle aggregating logprobs
        let logprobs = None;

378
        let finish_reason: Option<dynamo_async_openai::types::CompletionFinishReason> =
379
380
            response.delta.finish_reason.map(Into::into);

381
        let choice = dynamo_async_openai::types::Choice {
382
383
384
385
            text,
            index,
            logprobs,
            finish_reason,
386
387
388
389
390
        };

        Ok(choice)
    }
}
391

Greg Clark's avatar
Greg Clark committed
392
393
394
395
396
397
398
399
400
401
402
403
impl OpenAIOutputOptionsProvider for NvCreateCompletionRequest {
    fn get_logprobs(&self) -> Option<u32> {
        self.inner.logprobs.map(|logprobs| logprobs as u32)
    }

    fn get_prompt_logprobs(&self) -> Option<u32> {
        self.inner
            .echo
            .and_then(|echo| if echo { Some(1) } else { None })
    }

    fn get_skip_special_tokens(&self) -> Option<bool> {
404
        CommonExtProvider::get_skip_special_tokens(self)
Greg Clark's avatar
Greg Clark committed
405
406
407
408
409
410
411
    }

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

412
413
414
415
/// Implements `ValidateRequest` for `NvCreateCompletionRequest`,
/// allowing us to validate the data.
impl ValidateRequest for NvCreateCompletionRequest {
    fn validate(&self) -> Result<(), anyhow::Error> {
416
        validate::validate_no_unsupported_fields(&self.unsupported_fields)?;
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
        validate::validate_model(&self.inner.model)?;
        validate::validate_prompt(&self.inner.prompt)?;
        validate::validate_suffix(self.inner.suffix.as_deref())?;
        validate::validate_max_tokens(self.inner.max_tokens)?;
        validate::validate_temperature(self.inner.temperature)?;
        validate::validate_top_p(self.inner.top_p)?;
        validate::validate_n(self.inner.n)?;
        // none for stream
        // none for stream_options
        validate::validate_logprobs(self.inner.logprobs)?;
        // none for echo
        validate::validate_stop(&self.inner.stop)?;
        validate::validate_presence_penalty(self.inner.presence_penalty)?;
        validate::validate_frequency_penalty(self.inner.frequency_penalty)?;
        validate::validate_best_of(self.inner.best_of, self.inner.n)?;
        validate::validate_logit_bias(&self.inner.logit_bias)?;
        validate::validate_user(self.inner.user.as_deref())?;
        // none for seed
435
        // none for metadata
436

437
438
        // Common Ext
        validate::validate_repetition_penalty(self.get_repetition_penalty())?;
439
440
        validate::validate_min_p(self.get_min_p())?;
        validate::validate_top_k(self.get_top_k())?;
441
442
        // Cross-field validation
        validate::validate_n_with_temperature(self.inner.n, self.inner.temperature)?;
443
444
445
446
447
        // total choices validation for completions batch requests
        validate::validate_total_choices(
            get_prompt_batch_size(&self.inner.prompt),
            self.inner.n.unwrap_or(1),
        )?;
448
449
450
        Ok(())
    }
}
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496

#[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",
            "prompt": "Hello, world!"
        });

        let request: NvCreateCompletionRequest =
            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",
                "prompt": "Hello, world!",
                "skip_special_tokens": skip_value
            });

            let request: NvCreateCompletionRequest =
                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));
        }
    }
}