Skip to main content

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, Choice, CompletionUsage, 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}