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}