openai_protocol/
generate.rs

1use std::collections::HashMap;
2
3use serde::{Deserialize, Serialize};
4use serde_json::{Map, Value};
5use validator::Validate;
6
7use super::{
8    common::{default_true, deserialize_null_as_false, is_false, GenerationRequest, InputIds},
9    sampling_params::SamplingParams,
10};
11use crate::validated::Normalizable;
12
13// ============================================================================
14// SGLang Generate API (native format)
15// ============================================================================
16
17#[derive(Clone, Debug, Serialize, Deserialize, Validate, schemars::JsonSchema)]
18#[validate(schema(function = "validate_generate_request"))]
19pub struct GenerateRequest {
20    /// Text input - SGLang native format
21    #[serde(skip_serializing_if = "Option::is_none")]
22    pub text: Option<String>,
23
24    #[serde(default = "super::common::default_unknown_model")]
25    pub model: String,
26
27    /// Input IDs for tokenized input
28    #[serde(skip_serializing_if = "Option::is_none")]
29    pub input_ids: Option<InputIds>,
30
31    /// Input embeddings for direct embedding input
32    /// Can be a 2D array (single request) or 3D array (batch of requests)
33    /// Placeholder for future use
34    #[serde(skip_serializing_if = "Option::is_none")]
35    pub input_embeds: Option<Value>,
36
37    /// Image input data
38    /// Can be an image instance, file name, URL, or base64 encoded string
39    /// Supports single images, lists of images, or nested lists for batch processing
40    /// Placeholder for future use
41    #[serde(skip_serializing_if = "Option::is_none")]
42    pub image_data: Option<Value>,
43
44    /// Video input data
45    /// Can be a file name, URL, or base64 encoded string
46    /// Supports single videos, lists of videos, or nested lists for batch processing
47    /// Placeholder for future use
48    #[serde(skip_serializing_if = "Option::is_none")]
49    pub video_data: Option<Value>,
50
51    /// Audio input data
52    /// Can be a file name, URL, or base64 encoded string
53    /// Supports single audio files, lists of audio, or nested lists for batch processing
54    /// Placeholder for future use
55    #[serde(skip_serializing_if = "Option::is_none")]
56    pub audio_data: Option<Value>,
57
58    /// Sampling parameters (sglang style)
59    #[serde(skip_serializing_if = "Option::is_none")]
60    pub sampling_params: Option<SamplingParams>,
61
62    /// Whether to return logprobs
63    #[serde(skip_serializing_if = "Option::is_none")]
64    pub return_logprob: Option<bool>,
65
66    /// If return logprobs, the start location in the prompt for returning logprobs.
67    #[serde(skip_serializing_if = "Option::is_none")]
68    pub logprob_start_len: Option<i32>,
69
70    /// If return logprobs, the number of top logprobs to return at each position.
71    #[serde(skip_serializing_if = "Option::is_none")]
72    pub top_logprobs_num: Option<i32>,
73
74    /// If return logprobs, the token ids to return logprob for.
75    #[serde(skip_serializing_if = "Option::is_none")]
76    pub token_ids_logprob: Option<Vec<u32>>,
77
78    /// Whether to detokenize tokens in text in the returned logprobs.
79    #[serde(default)]
80    pub return_text_in_logprobs: bool,
81
82    /// Whether to stream the response
83    #[serde(default, deserialize_with = "deserialize_null_as_false")]
84    pub stream: bool,
85
86    /// Whether to log metrics for this request (e.g. health_generate calls do not log metrics)
87    #[serde(default = "default_true")]
88    pub log_metrics: bool,
89
90    /// Return model hidden states
91    #[serde(default, skip_serializing_if = "is_false")]
92    pub return_hidden_states: bool,
93
94    /// The modalities of the image data [image, multi-images, video]
95    #[serde(skip_serializing_if = "Option::is_none")]
96    pub modalities: Option<Vec<String>>,
97
98    /// Session parameters for continual prompting
99    #[serde(skip_serializing_if = "Option::is_none")]
100    pub session_params: Option<HashMap<String, Value>>,
101
102    /// Path to LoRA adapter(s) for model customization
103    #[serde(skip_serializing_if = "Option::is_none")]
104    pub lora_path: Option<String>,
105
106    /// LoRA adapter ID (if pre-loaded)
107    #[serde(skip_serializing_if = "Option::is_none")]
108    pub lora_id: Option<String>,
109
110    /// Custom logit processor for advanced sampling control. Must be a serialized instance
111    /// of `CustomLogitProcessor` in python/sglang/srt/sampling/custom_logit_processor.py
112    /// Use the processor's `to_str()` method to generate the serialized string.
113    #[serde(skip_serializing_if = "Option::is_none")]
114    pub custom_logit_processor: Option<String>,
115
116    /// For disaggregated inference
117    #[serde(skip_serializing_if = "Option::is_none")]
118    pub bootstrap_host: Option<String>,
119
120    /// For disaggregated inference
121    #[serde(skip_serializing_if = "Option::is_none")]
122    pub bootstrap_port: Option<i32>,
123
124    /// For disaggregated inference
125    #[serde(skip_serializing_if = "Option::is_none")]
126    pub bootstrap_room: Option<i32>,
127
128    /// For disaggregated inference
129    #[serde(skip_serializing_if = "Option::is_none")]
130    pub bootstrap_pair_key: Option<String>,
131
132    /// Data parallel rank routing
133    #[serde(skip_serializing_if = "Option::is_none")]
134    pub data_parallel_rank: Option<i32>,
135
136    /// Background response
137    #[serde(default)]
138    pub background: bool,
139
140    /// Conversation ID for tracking
141    #[serde(skip_serializing_if = "Option::is_none")]
142    pub conversation_id: Option<String>,
143
144    /// Priority for the request
145    #[serde(skip_serializing_if = "Option::is_none")]
146    pub priority: Option<i32>,
147
148    /// Extra key for classifying the request (e.g. cache_salt)
149    #[serde(skip_serializing_if = "Option::is_none")]
150    pub extra_key: Option<String>,
151
152    /// Whether to disallow logging for this request (e.g. due to ZDR)
153    #[serde(default)]
154    pub no_logs: bool,
155
156    /// Custom metric labels
157    #[serde(skip_serializing_if = "Option::is_none")]
158    pub custom_labels: Option<HashMap<String, String>>,
159
160    /// Whether to return bytes for image generation
161    #[serde(default)]
162    pub return_bytes: bool,
163
164    /// Whether to return entropy
165    #[serde(default)]
166    pub return_entropy: bool,
167
168    /// Request ID for tracking (inherited from BaseReq in Python)
169    #[serde(skip_serializing_if = "Option::is_none")]
170    pub rid: Option<String>,
171
172    /// Additional fields not explicitly defined above (e.g. engine-specific parameters)
173    #[serde(flatten)]
174    pub other: Map<String, Value>,
175}
176
177impl Normalizable for GenerateRequest {
178    // Use default no-op implementation - no normalization needed for GenerateRequest
179}
180
181/// Validation function for GenerateRequest - ensure exactly one input type is provided
182fn validate_generate_request(req: &GenerateRequest) -> Result<(), validator::ValidationError> {
183    // Exactly one of text or input_ids must be provided
184    // Note: input_embeds not yet supported in Rust implementation
185    let has_text = req.text.is_some();
186    let has_input_ids = req.input_ids.is_some();
187
188    let count = [has_text, has_input_ids].iter().filter(|&&x| x).count();
189
190    if count == 0 {
191        return Err(validator::ValidationError::new(
192            "Either text or input_ids should be provided.",
193        ));
194    }
195
196    if count > 1 {
197        return Err(validator::ValidationError::new(
198            "Either text or input_ids should be provided.",
199        ));
200    }
201
202    Ok(())
203}
204
205impl GenerationRequest for GenerateRequest {
206    fn rid(&self) -> Option<&str> {
207        self.rid.as_deref()
208    }
209
210    fn is_stream(&self) -> bool {
211        self.stream
212    }
213
214    fn get_model(&self) -> Option<&str> {
215        Some(self.model.as_str())
216    }
217
218    fn extract_text_for_routing(&self) -> String {
219        // Check fields in priority order: text, input_ids
220        if let Some(ref text) = self.text {
221            return text.clone();
222        }
223
224        if let Some(ref input_ids) = self.input_ids {
225            return match input_ids {
226                InputIds::Single(ids) => ids
227                    .iter()
228                    .map(|&id| id.to_string())
229                    .collect::<Vec<String>>()
230                    .join(" "),
231                InputIds::Batch(batches) => batches
232                    .iter()
233                    .flat_map(|batch| batch.iter().map(|&id| id.to_string()))
234                    .collect::<Vec<String>>()
235                    .join(" "),
236            };
237        }
238
239        // No text input found
240        String::new()
241    }
242
243    fn routing_tokens(&self) -> Option<&[i32]> {
244        // Token ids win over text: they key routing on what the backend KV
245        // cache keys on. Empty ids fall back to text.
246        match &self.input_ids {
247            Some(InputIds::Single(ids)) if !ids.is_empty() => Some(ids),
248            // A batch is dispatched to a single worker; the first sequence is
249            // the best available affinity signal.
250            Some(InputIds::Batch(seqs)) => seqs
251                .first()
252                .map(Vec::as_slice)
253                .filter(|ids| !ids.is_empty()),
254            _ => None,
255        }
256    }
257}
258
259// ============================================================================
260// SGLang Generate Response Types
261// ============================================================================
262
263/// SGLang generate response (single completion or array for n>1)
264///
265/// Format for n=1:
266/// ```json
267/// {
268///   "text": "...",
269///   "output_ids": [...],
270///   "meta_info": { ... }
271/// }
272/// ```
273///
274/// Format for n>1:
275/// ```json
276/// [
277///   {"text": "...", "output_ids": [...], "meta_info": {...}},
278///   {"text": "...", "output_ids": [...], "meta_info": {...}}
279/// ]
280/// ```
281#[derive(Debug, Clone, Serialize, Deserialize, schemars::JsonSchema)]
282pub struct GenerateResponse {
283    pub text: String,
284    pub output_ids: Vec<u32>,
285    pub meta_info: GenerateMetaInfo,
286}
287
288/// Metadata for a single generate completion
289#[serde_with::skip_serializing_none]
290#[derive(Debug, Clone, Serialize, Deserialize, schemars::JsonSchema)]
291pub struct GenerateMetaInfo {
292    pub id: String,
293    pub finish_reason: GenerateFinishReason,
294    pub prompt_tokens: u32,
295    pub weight_version: String,
296    pub input_token_logprobs: Option<Vec<Vec<Option<f64>>>>,
297    pub output_token_logprobs: Option<Vec<Vec<Option<f64>>>>,
298    pub completion_tokens: u32,
299    pub cached_tokens: u32,
300    pub reasoning_tokens: Option<u32>,
301    pub e2e_latency: f64,
302    pub matched_stop: Option<Value>,
303}
304
305/// Finish reason for generate endpoint
306#[derive(Debug, Clone, Serialize, Deserialize, schemars::JsonSchema)]
307#[serde(untagged)]
308pub enum GenerateFinishReason {
309    Length {
310        #[serde(rename = "type")]
311        finish_type: GenerateFinishType,
312        length: u32,
313    },
314    Stop {
315        #[serde(rename = "type")]
316        finish_type: GenerateFinishType,
317    },
318    Other(Value),
319}
320
321#[derive(Debug, Clone, Serialize, Deserialize, schemars::JsonSchema)]
322#[serde(rename_all = "lowercase")]
323pub enum GenerateFinishType {
324    Length,
325    Stop,
326}
327
328#[cfg(test)]
329mod tests {
330    use super::*;
331
332    fn req() -> GenerateRequest {
333        serde_json::from_value(serde_json::json!({"model": "m"})).expect("minimal request")
334    }
335
336    #[test]
337    fn return_hidden_states_false_is_omitted_and_absent_reads_false() {
338        let r = req();
339        let v = serde_json::to_value(&r).expect("serialize");
340        assert!(v.get("return_hidden_states").is_none());
341
342        let back: GenerateRequest = serde_json::from_value(v).expect("roundtrip");
343        assert!(!back.return_hidden_states);
344    }
345
346    #[test]
347    fn return_hidden_states_true_round_trips() {
348        let mut r = req();
349        r.return_hidden_states = true;
350        let v = serde_json::to_value(&r).expect("serialize");
351        assert_eq!(v["return_hidden_states"], true);
352
353        let back: GenerateRequest = serde_json::from_value(v).expect("roundtrip");
354        assert!(back.return_hidden_states);
355    }
356
357    #[test]
358    fn routing_tokens_from_single_input_ids() {
359        let mut r = req();
360        r.input_ids = Some(InputIds::Single(vec![1, 2, 3]));
361        assert_eq!(r.routing_tokens(), Some(&[1, 2, 3][..]));
362    }
363
364    #[test]
365    fn routing_tokens_prefer_input_ids_over_text() {
366        let mut r = req();
367        r.text = Some("hello".to_string());
368        r.input_ids = Some(InputIds::Single(vec![1, 2, 3]));
369        assert_eq!(r.routing_tokens(), Some(&[1, 2, 3][..]));
370
371        r.input_ids = Some(InputIds::Batch(vec![vec![4, 5], vec![6]]));
372        assert_eq!(r.routing_tokens(), Some(&[4, 5][..]));
373    }
374
375    #[test]
376    fn routing_tokens_none_for_empty_input_ids_with_text() {
377        let mut r = req();
378        r.text = Some("hello".to_string());
379        r.input_ids = Some(InputIds::Single(vec![]));
380        assert_eq!(r.routing_tokens(), None);
381        assert_eq!(r.extract_text_for_routing(), "hello");
382    }
383
384    #[test]
385    fn routing_tokens_from_batch_first_sequence() {
386        let mut r = req();
387        r.input_ids = Some(InputIds::Batch(vec![vec![1, 2], vec![3, 4]]));
388        assert_eq!(r.routing_tokens(), Some(&[1, 2][..]));
389    }
390
391    #[test]
392    fn routing_tokens_none_for_empty_inputs() {
393        let mut r = req();
394        r.input_ids = Some(InputIds::Batch(vec![]));
395        assert_eq!(r.routing_tokens(), None);
396
397        let mut r = req();
398        r.input_ids = Some(InputIds::Batch(vec![vec![], vec![1]]));
399        assert_eq!(r.routing_tokens(), None);
400
401        let r = req();
402        assert_eq!(r.routing_tokens(), None);
403    }
404}
openai_protocol/generate.rs

openai_protocol/
generate.rs
