Skip to main content

openai_protocol/
classify.rs

1//! Classify API protocol definitions.
2//!
3//! This module defines the request and response types for the `/v1/classify` API,
4//! which is compatible with vLLM's classification endpoint.
5//!
6//! Classification reuses the embedding backend - the scheduler returns logits as
7//! "embeddings", and the classify layer applies softmax + label mapping.
8
9use serde::{Deserialize, Serialize};
10use serde_json::Value;
11
12use super::common::{GenerationRequest, UsageInfo};
13
14// ============================================================================
15// Classify API
16// ============================================================================
17
18/// Classification request - compatible with vLLM's /v1/classify API
19#[serde_with::skip_serializing_none]
20#[derive(Debug, Clone, Deserialize, Serialize, schemars::JsonSchema)]
21pub struct ClassifyRequest {
22    /// ID of the model to use
23    pub model: String,
24
25    /// Input can be a string, array of strings, or token IDs
26    /// - Single string: "text to classify"
27    /// - Array of strings: ["text1", "text2"]
28    /// - Token IDs: [1, 2, 3] (advanced usage)
29    pub input: Value,
30
31    /// Optional user identifier
32    pub user: Option<String>,
33
34    /// SGLang extension: request id for tracking
35    pub rid: Option<String>,
36
37    /// SGLang extension: request priority
38    pub priority: Option<i32>,
39}
40
41impl GenerationRequest for ClassifyRequest {
42    fn rid(&self) -> Option<&str> {
43        self.rid.as_deref()
44    }
45
46    fn is_stream(&self) -> bool {
47        false // Classification is always non-streaming
48    }
49
50    fn get_model(&self) -> Option<&str> {
51        Some(&self.model)
52    }
53
54    fn extract_text_for_routing(&self) -> String {
55        match &self.input {
56            Value::String(s) => s.clone(),
57            Value::Array(arr) => arr
58                .iter()
59                .filter_map(|v| v.as_str())
60                .collect::<Vec<_>>()
61                .join(" "),
62            _ => String::new(),
63        }
64    }
65}
66
67// ============================================================================
68// Classify Response
69// ============================================================================
70
71/// Single classification result
72#[derive(Debug, Clone, Serialize, Deserialize, schemars::JsonSchema)]
73pub struct ClassifyData {
74    /// Index of this result (for batch requests)
75    pub index: u32,
76    /// Predicted class label (from id2label mapping)
77    pub label: String,
78    /// Probability distribution over all classes (softmax of logits)
79    pub probs: Vec<f32>,
80    /// Number of classes
81    pub num_classes: u32,
82}
83
84/// Classification response
85#[derive(Debug, Clone, Serialize, Deserialize, schemars::JsonSchema)]
86pub struct ClassifyResponse {
87    /// Unique request ID (format: "classify-{uuid}")
88    pub id: String,
89    /// Always "list"
90    pub object: String,
91    /// Unix timestamp (seconds since epoch)
92    pub created: u64,
93    /// Model name
94    pub model: String,
95    /// Classification results (one per input in batch)
96    pub data: Vec<ClassifyData>,
97    /// Token usage info
98    pub usage: UsageInfo,
99}
100
101impl ClassifyResponse {
102    /// Create a new ClassifyResponse with the given data
103    pub fn new(
104        id: String,
105        model: String,
106        created: u64,
107        data: Vec<ClassifyData>,
108        usage: UsageInfo,
109    ) -> Self {
110        Self {
111            id,
112            object: "list".to_string(),
113            created,
114            model,
115            data,
116            usage,
117        }
118    }
119}