dynamo_protocols/types/
completion.rs

1// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3//
4// Re-exports upstream async-openai completion types and defines
5// inference-serving extensions.
6
7use std::collections::HashMap;
8use std::pin::Pin;
9
10use derive_builder::Builder;
11use futures::Stream;
12use serde::{Deserialize, Serialize};
13
14use crate::error::OpenAIError;
15
16use super::{ChatCompletionStreamOptions, Prompt, Stop};
17
18// Re-export response type from upstream (identical)
19pub use async_openai::types::completions::CreateCompletionResponse;
20
21/// Custom deserializer for the echo parameter that only accepts booleans.
22/// Rejects integers and strings with clear error messages.
23fn deserialize_echo_bool<'de, D>(deserializer: D) -> Result<Option<bool>, D::Error>
24where
25    D: serde::Deserializer<'de>,
26{
27    struct StrictBoolVisitor;
28
29    impl<'de> serde::de::Visitor<'de> for StrictBoolVisitor {
30        type Value = Option<bool>;
31
32        fn expecting(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result {
33            formatter.write_str("echo parameter to be a boolean (true or false) or null")
34        }
35
36        fn visit_some<D>(self, deserializer: D) -> Result<Self::Value, D::Error>
37        where
38            D: serde::Deserializer<'de>,
39        {
40            deserializer.deserialize_any(BoolOnlyVisitor)
41        }
42
43        fn visit_none<E>(self) -> Result<Self::Value, E>
44        where
45            E: serde::de::Error,
46        {
47            Ok(None)
48        }
49
50        fn visit_unit<E>(self) -> Result<Self::Value, E>
51        where
52            E: serde::de::Error,
53        {
54            Ok(None)
55        }
56    }
57
58    struct BoolOnlyVisitor;
59
60    impl<'de> serde::de::Visitor<'de> for BoolOnlyVisitor {
61        type Value = Option<bool>;
62
63        fn expecting(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result {
64            formatter.write_str("echo parameter to be a boolean (true or false) or null")
65        }
66
67        fn visit_bool<E>(self, value: bool) -> Result<Self::Value, E>
68        where
69            E: serde::de::Error,
70        {
71            Ok(Some(value))
72        }
73
74        fn visit_str<E>(self, value: &str) -> Result<Self::Value, E>
75        where
76            E: serde::de::Error,
77        {
78            Err(E::invalid_type(
79                serde::de::Unexpected::Str(value),
80                &"echo parameter to be a boolean (true or false) or null",
81            ))
82        }
83    }
84
85    deserializer.deserialize_option(StrictBoolVisitor)
86}
87
88/// Completion request with inference-serving extensions.
89///
90/// Extends upstream `CreateCompletionRequest` with:
91/// - `prompt_embeds`: base64-encoded PyTorch tensor for pre-computed embeddings
92/// - `echo`: strict bool validation (rejects integers/strings)
93/// - `stream_options`: uses our extended `ChatCompletionStreamOptions` (with `continuous_usage_stats`)
94#[derive(Clone, Serialize, Deserialize, Default, Debug, Builder, PartialEq)]
95#[builder(name = "CreateCompletionRequestArgs")]
96#[builder(pattern = "mutable")]
97#[builder(setter(into, strip_option), default)]
98#[builder(derive(Debug))]
99#[builder(build_fn(error = "OpenAIError"))]
100pub struct CreateCompletionRequest {
101    pub model: String,
102    pub prompt: Prompt,
103    /// Base64-encoded PyTorch tensor containing pre-computed embeddings.
104    /// At least one of prompt or prompt_embeds is required.
105    #[serde(skip_serializing_if = "Option::is_none")]
106    pub prompt_embeds: Option<String>,
107    #[serde(skip_serializing_if = "Option::is_none")]
108    pub suffix: Option<String>,
109    #[serde(skip_serializing_if = "Option::is_none")]
110    pub max_tokens: Option<u32>,
111    #[serde(skip_serializing_if = "Option::is_none")]
112    pub temperature: Option<f32>,
113    #[serde(skip_serializing_if = "Option::is_none")]
114    pub top_p: Option<f32>,
115    #[serde(skip_serializing_if = "Option::is_none")]
116    pub n: Option<u8>,
117    #[serde(skip_serializing_if = "Option::is_none")]
118    pub stream: Option<bool>,
119    #[serde(skip_serializing_if = "Option::is_none")]
120    pub stream_options: Option<ChatCompletionStreamOptions>,
121    #[serde(skip_serializing_if = "Option::is_none")]
122    pub logprobs: Option<u8>,
123    /// Echo back the prompt in addition to the completion.
124    /// Strict bool validation -- rejects integers and strings.
125    #[serde(skip_serializing_if = "Option::is_none")]
126    #[serde(default, deserialize_with = "deserialize_echo_bool")]
127    pub echo: Option<bool>,
128    #[serde(skip_serializing_if = "Option::is_none")]
129    pub stop: Option<Stop>,
130    #[serde(skip_serializing_if = "Option::is_none")]
131    pub presence_penalty: Option<f32>,
132    #[serde(skip_serializing_if = "Option::is_none")]
133    pub frequency_penalty: Option<f32>,
134    #[serde(skip_serializing_if = "Option::is_none")]
135    pub best_of: Option<u8>,
136    #[serde(skip_serializing_if = "Option::is_none")]
137    pub logit_bias: Option<HashMap<String, serde_json::Value>>,
138    #[serde(skip_serializing_if = "Option::is_none")]
139    pub user: Option<String>,
140    #[serde(skip_serializing_if = "Option::is_none")]
141    pub seed: Option<i64>,
142}
143
144/// Parsed server side events stream until an \[DONE\] is received from server.
145pub type CompletionResponseStream =
146    Pin<Box<dyn Stream<Item = Result<CreateCompletionResponse, OpenAIError>> + Send>>;
147
148#[cfg(test)]
149mod tests {
150    use super::*;
151
152    #[test]
153    fn echo_rejects_integer() {
154        let json = r#"{"model": "test_model", "prompt": "test", "echo": 1}"#;
155        let result: Result<CreateCompletionRequest, _> = serde_json::from_str(json);
156        assert!(result.is_err());
157        let err_msg = result.unwrap_err().to_string();
158        assert!(err_msg.contains("invalid type"));
159        assert!(err_msg.contains("integer"));
160        assert!(err_msg.contains("echo parameter"));
161    }
162
163    #[test]
164    fn echo_rejects_string() {
165        let json = r#"{"model": "test_model", "prompt": "test", "echo": "null"}"#;
166        let result: Result<CreateCompletionRequest, _> = serde_json::from_str(json);
167        assert!(result.is_err());
168        let err_msg = result.unwrap_err().to_string();
169        assert!(err_msg.contains("invalid type"));
170        assert!(err_msg.contains("string"));
171        assert!(err_msg.contains("echo parameter"));
172    }
173
174    #[test]
175    fn completion_choice_serializes_openai_shape() {
176        use crate::types::{Choice, CompletionFinishReason};
177
178        let choice = Choice {
179            text: "hello".to_string(),
180            index: 0,
181            logprobs: None,
182            finish_reason: Some(CompletionFinishReason::Stop),
183        };
184
185        let value = serde_json::to_value(choice).expect("serialize choice");
186
187        assert_eq!(value["finish_reason"], "stop");
188        assert_eq!(value["text"], "hello");
189    }
190
191    #[test]
192    fn stop_accepts_token_id_array() {
193        let json = r#"{"model": "test_model", "prompt": [1, 2, 3], "stop": [32, 34]}"#;
194        let request: CreateCompletionRequest = serde_json::from_str(json).unwrap();
195
196        assert_eq!(request.stop, Some(Stop::TokenIdArray(vec![32, 34])));
197    }
198
199    #[test]
200    fn stop_accepts_string_and_string_array() {
201        let one_stop = r#"{"model": "test_model", "prompt": "hello", "stop": " The"}"#;
202        let request: CreateCompletionRequest = serde_json::from_str(one_stop).unwrap();
203
204        assert_eq!(request.stop, Some(Stop::String(" The".to_string())));
205
206        let many_stops = r#"{"model": "test_model", "prompt": "hello", "stop": ["A", "B"]}"#;
207        let request: CreateCompletionRequest = serde_json::from_str(many_stops).unwrap();
208
209        assert_eq!(
210            request.stop,
211            Some(Stop::StringArray(vec!["A".to_string(), "B".to_string()]))
212        );
213    }
214
215    #[test]
216    fn stop_token_id_display_string_remains_string_stop() {
217        let json = r#"{"model": "test_model", "prompt": [1, 2, 3], "stop": "token_id:576"}"#;
218        let request: CreateCompletionRequest = serde_json::from_str(json).unwrap();
219
220        assert_eq!(request.stop, Some(Stop::String("token_id:576".to_string())));
221
222        let json = r#"{"model": "test_model", "prompt": [1, 2, 3], "stop": ["token_id:576"]}"#;
223        let request: CreateCompletionRequest = serde_json::from_str(json).unwrap();
224
225        assert_eq!(
226            request.stop,
227            Some(Stop::StringArray(vec!["token_id:576".to_string()]))
228        );
229    }
230
231    #[test]
232    fn builder_accepts_upstream_stop_configuration() {
233        let upstream_stop = async_openai::types::chat::StopConfiguration::String("END".to_string());
234
235        let request = CreateCompletionRequestArgs::default()
236            .model("test_model")
237            .prompt(Prompt::String("hello".to_string()))
238            .stop(upstream_stop)
239            .build()
240            .unwrap();
241
242        assert_eq!(request.stop, Some(Stop::String("END".to_string())));
243    }
244
245    #[test]
246    fn stop_rejects_single_token_id() {
247        let json = r#"{"model": "test_model", "prompt": [1, 2, 3], "stop": 576}"#;
248        let result: Result<CreateCompletionRequest, _> = serde_json::from_str(json);
249
250        assert!(result.is_err());
251    }
252}
dynamo_protocols/types/completion.rs

dynamo_protocols/types/
completion.rs
