openai_protocol/
model_type.rs

1//! Model type definitions using bitflags for endpoint support.
2//!
3//! Defines [`ModelType`] using bitflags to represent which endpoints a model
4//! can support, and [`Endpoint`] for routing decisions.
5
6use std::borrow::Cow;
7
8use bitflags::bitflags;
9use schemars::{json_schema, JsonSchema, Schema, SchemaGenerator};
10use serde::{Deserialize, Serialize};
11
12bitflags! {
13    #[derive(Copy, Debug, Default, Clone, Eq, PartialEq, Hash)]
14    pub struct ModelType: u16 {
15        /// OpenAI Chat Completions API (/v1/chat/completions)
16        const CHAT        = 1 << 0;
17        /// OpenAI Completions API - legacy (/v1/completions)
18        const COMPLETIONS = 1 << 1;
19        /// OpenAI Responses API (/v1/responses)
20        const RESPONSES   = 1 << 2;
21        /// Embeddings API (/v1/embeddings)
22        const EMBEDDINGS  = 1 << 3;
23        /// Rerank API (/v1/rerank)
24        const RERANK      = 1 << 4;
25        /// SGLang Generate API (/generate)
26        const GENERATE    = 1 << 5;
27        /// Vision/multimodal support (images in input)
28        const VISION      = 1 << 6;
29        /// Tool/function calling support
30        const TOOLS       = 1 << 7;
31        /// Reasoning/thinking support (e.g., o1, DeepSeek-R1)
32        const REASONING   = 1 << 8;
33        /// Image generation (DALL-E, Sora, gpt-image)
34        const IMAGE_GEN   = 1 << 9;
35        /// Audio models (TTS, Whisper, realtime, transcribe)
36        const AUDIO       = 1 << 10;
37        /// Content moderation models
38        const MODERATION  = 1 << 11;
39
40        /// Standard LLM: chat + completions + responses + tools
41        const LLM = Self::CHAT.bits() | Self::COMPLETIONS.bits()
42                  | Self::RESPONSES.bits() | Self::TOOLS.bits();
43
44        /// Vision-capable LLM: LLM + vision
45        const VISION_LLM = Self::LLM.bits() | Self::VISION.bits();
46
47        /// Reasoning LLM: LLM + reasoning (e.g., o1, o3, DeepSeek-R1)
48        const REASONING_LLM = Self::LLM.bits() | Self::REASONING.bits();
49
50        /// Full-featured LLM: all text generation capabilities
51        const FULL_LLM = Self::VISION_LLM.bits() | Self::REASONING.bits();
52
53        /// Embedding model only
54        const EMBED_MODEL = Self::EMBEDDINGS.bits();
55
56        /// Reranker model only
57        const RERANK_MODEL = Self::RERANK.bits();
58
59        /// Image generation model only (DALL-E, Sora, gpt-image)
60        const IMAGE_MODEL = Self::IMAGE_GEN.bits();
61
62        /// Audio model only (TTS, Whisper, realtime)
63        const AUDIO_MODEL = Self::AUDIO.bits();
64
65        /// Content moderation model only
66        const MODERATION_MODEL = Self::MODERATION.bits();
67    }
68}
69
70/// Mapping of individual capability flags to their names.
71const CAPABILITY_NAMES: &[(ModelType, &str)] = &[
72    (ModelType::CHAT, "chat"),
73    (ModelType::COMPLETIONS, "completions"),
74    (ModelType::RESPONSES, "responses"),
75    (ModelType::EMBEDDINGS, "embeddings"),
76    (ModelType::RERANK, "rerank"),
77    (ModelType::GENERATE, "generate"),
78    (ModelType::VISION, "vision"),
79    (ModelType::TOOLS, "tools"),
80    (ModelType::REASONING, "reasoning"),
81    (ModelType::IMAGE_GEN, "image_gen"),
82    (ModelType::AUDIO, "audio"),
83    (ModelType::MODERATION, "moderation"),
84];
85
86impl ModelType {
87    /// Check if this model type supports the chat completions endpoint
88    #[inline]
89    pub fn supports_chat(self) -> bool {
90        self.contains(Self::CHAT)
91    }
92
93    /// Check if this model type supports the legacy completions endpoint
94    #[inline]
95    pub fn supports_completions(self) -> bool {
96        self.contains(Self::COMPLETIONS)
97    }
98
99    /// Check if this model type supports the responses endpoint
100    #[inline]
101    pub fn supports_responses(self) -> bool {
102        self.contains(Self::RESPONSES)
103    }
104
105    /// Check if this model type supports the embeddings endpoint
106    #[inline]
107    pub fn supports_embeddings(self) -> bool {
108        self.contains(Self::EMBEDDINGS)
109    }
110
111    /// Check if this model type supports the rerank endpoint
112    #[inline]
113    pub fn supports_rerank(self) -> bool {
114        self.contains(Self::RERANK)
115    }
116
117    /// Check if this model type supports the generate endpoint
118    #[inline]
119    pub fn supports_generate(self) -> bool {
120        self.contains(Self::GENERATE)
121    }
122
123    /// Check if this model type supports vision/multimodal input
124    #[inline]
125    pub fn supports_vision(self) -> bool {
126        self.contains(Self::VISION)
127    }
128
129    /// Check if this model type supports tool/function calling
130    #[inline]
131    pub fn supports_tools(self) -> bool {
132        self.contains(Self::TOOLS)
133    }
134
135    /// Check if this model type supports reasoning/thinking
136    #[inline]
137    pub fn supports_reasoning(self) -> bool {
138        self.contains(Self::REASONING)
139    }
140
141    /// Check if this model type supports image generation
142    #[inline]
143    pub fn supports_image_gen(self) -> bool {
144        self.contains(Self::IMAGE_GEN)
145    }
146
147    /// Check if this model type supports audio (TTS, Whisper, etc.)
148    #[inline]
149    pub fn supports_audio(self) -> bool {
150        self.contains(Self::AUDIO)
151    }
152
153    /// Check if this model type supports content moderation
154    #[inline]
155    pub fn supports_moderation(self) -> bool {
156        self.contains(Self::MODERATION)
157    }
158
159    /// Check if this model type supports a given endpoint
160    pub fn supports_endpoint(self, endpoint: Endpoint) -> bool {
161        match endpoint {
162            Endpoint::Chat => self.supports_chat(),
163            Endpoint::Completions => self.supports_completions(),
164            Endpoint::Responses => self.supports_responses(),
165            Endpoint::Embeddings => self.supports_embeddings(),
166            Endpoint::Rerank => self.supports_rerank(),
167            Endpoint::Generate => self.supports_generate(),
168            Endpoint::Models => true,
169        }
170    }
171
172    /// Convert to a list of supported capability names
173    pub fn as_capability_names(self) -> Vec<&'static str> {
174        let mut result = Vec::with_capacity(CAPABILITY_NAMES.len());
175        for &(flag, name) in CAPABILITY_NAMES {
176            if self.contains(flag) {
177                result.push(name);
178            }
179        }
180        result
181    }
182
183    /// Check if this is an LLM (supports at least chat)
184    #[inline]
185    pub fn is_llm(self) -> bool {
186        self.supports_chat()
187    }
188
189    /// Check if this is an embedding model
190    #[inline]
191    pub fn is_embedding_model(self) -> bool {
192        self.supports_embeddings() && !self.supports_chat()
193    }
194
195    /// Check if this is a reranker model
196    #[inline]
197    pub fn is_reranker(self) -> bool {
198        self.supports_rerank() && !self.supports_chat()
199    }
200
201    /// Check if this is an image generation model
202    #[inline]
203    pub fn is_image_model(self) -> bool {
204        self.supports_image_gen() && !self.supports_chat()
205    }
206
207    /// Check if this is an audio model
208    #[inline]
209    pub fn is_audio_model(self) -> bool {
210        self.supports_audio() && !self.supports_chat()
211    }
212
213    /// Check if this is a moderation model
214    #[inline]
215    pub fn is_moderation_model(self) -> bool {
216        self.supports_moderation() && !self.supports_chat()
217    }
218}
219
220impl std::fmt::Display for ModelType {
221    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
222        let names = self.as_capability_names();
223        if names.is_empty() {
224            write!(f, "none")
225        } else {
226            write!(f, "{}", names.join(","))
227        }
228    }
229}
230
231impl Serialize for ModelType {
232    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
233    where
234        S: serde::Serializer,
235    {
236        use serde::ser::SerializeSeq;
237        let names = self.as_capability_names();
238        let mut seq = serializer.serialize_seq(Some(names.len()))?;
239        for name in names {
240            seq.serialize_element(name)?;
241        }
242        seq.end()
243    }
244}
245
246impl<'de> Deserialize<'de> for ModelType {
247    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
248    where
249        D: serde::Deserializer<'de>,
250    {
251        use serde::de;
252
253        struct ModelTypeVisitor;
254
255        impl<'de> de::Visitor<'de> for ModelTypeVisitor {
256            type Value = ModelType;
257
258            fn expecting(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
259                f.write_str("an array of capability names or a u16 bitfield")
260            }
261
262            // Backward compat: accept numeric u16 bitfield
263            fn visit_u64<E: de::Error>(self, v: u64) -> Result<ModelType, E> {
264                let bits = u16::try_from(v)
265                    .map_err(|_| E::custom(format!("ModelType bits out of u16 range: {v}")))?;
266                ModelType::from_bits(bits)
267                    .ok_or_else(|| E::custom(format!("invalid ModelType bits: {bits}")))
268            }
269
270            // New format: array of capability name strings
271            fn visit_seq<A: de::SeqAccess<'de>>(self, mut seq: A) -> Result<ModelType, A::Error> {
272                let mut model_type = ModelType::empty();
273                while let Some(name) = seq.next_element::<String>()? {
274                    let flag = CAPABILITY_NAMES
275                        .iter()
276                        .find(|(_, n)| *n == name.as_str())
277                        .map(|(f, _)| *f)
278                        .ok_or_else(|| {
279                            de::Error::custom(format!("unknown ModelType capability: {name}"))
280                        })?;
281                    model_type |= flag;
282                }
283                Ok(model_type)
284            }
285        }
286
287        deserializer.deserialize_any(ModelTypeVisitor)
288    }
289}
290
291/// Manual JsonSchema impl for `ModelType` — serialized as an array of capability name strings.
292impl JsonSchema for ModelType {
293    fn schema_name() -> Cow<'static, str> {
294        "ModelType".into()
295    }
296
297    fn json_schema(_generator: &mut SchemaGenerator) -> Schema {
298        json_schema!({
299            "type": "array",
300            "description": "Bitflag capabilities serialized as an array of capability names",
301            "items": {
302                "type": "string",
303                "enum": [
304                    "chat",
305                    "completions",
306                    "responses",
307                    "embeddings",
308                    "rerank",
309                    "generate",
310                    "vision",
311                    "tools",
312                    "reasoning",
313                    "image_gen",
314                    "audio",
315                    "moderation"
316                ]
317            }
318        })
319    }
320}
321
322/// Endpoint types for routing decisions.
323#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize, schemars::JsonSchema)]
324#[serde(rename_all = "lowercase")]
325pub enum Endpoint {
326    /// Chat completions endpoint (/v1/chat/completions)
327    Chat,
328    /// Legacy completions endpoint (/v1/completions)
329    Completions,
330    /// Responses endpoint (/v1/responses)
331    Responses,
332    /// Embeddings endpoint (/v1/embeddings)
333    Embeddings,
334    /// Rerank endpoint (/v1/rerank)
335    Rerank,
336    /// SGLang generate endpoint (/generate)
337    Generate,
338    /// Models listing endpoint (/v1/models)
339    Models,
340}
341
342impl Endpoint {
343    /// Get the URL path for this endpoint
344    pub fn path(self) -> &'static str {
345        match self {
346            Endpoint::Chat => "/v1/chat/completions",
347            Endpoint::Completions => "/v1/completions",
348            Endpoint::Responses => "/v1/responses",
349            Endpoint::Embeddings => "/v1/embeddings",
350            Endpoint::Rerank => "/v1/rerank",
351            Endpoint::Generate => "/generate",
352            Endpoint::Models => "/v1/models",
353        }
354    }
355
356    /// Parse an endpoint from a URL path
357    pub fn from_path(path: &str) -> Option<Self> {
358        let path = path.trim_end_matches('/');
359        match path {
360            "/v1/chat/completions" => Some(Endpoint::Chat),
361            "/v1/completions" => Some(Endpoint::Completions),
362            "/v1/responses" => Some(Endpoint::Responses),
363            "/v1/embeddings" => Some(Endpoint::Embeddings),
364            "/v1/rerank" => Some(Endpoint::Rerank),
365            "/generate" => Some(Endpoint::Generate),
366            "/v1/models" => Some(Endpoint::Models),
367            _ => None,
368        }
369    }
370
371    /// Get the required ModelType flag for this endpoint
372    pub fn required_capability(self) -> Option<ModelType> {
373        match self {
374            Endpoint::Chat => Some(ModelType::CHAT),
375            Endpoint::Completions => Some(ModelType::COMPLETIONS),
376            Endpoint::Responses => Some(ModelType::RESPONSES),
377            Endpoint::Embeddings => Some(ModelType::EMBEDDINGS),
378            Endpoint::Rerank => Some(ModelType::RERANK),
379            Endpoint::Generate => Some(ModelType::GENERATE),
380            Endpoint::Models => None,
381        }
382    }
383}
384
385impl std::fmt::Display for Endpoint {
386    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
387        match self {
388            Endpoint::Chat => write!(f, "chat"),
389            Endpoint::Completions => write!(f, "completions"),
390            Endpoint::Responses => write!(f, "responses"),
391            Endpoint::Embeddings => write!(f, "embeddings"),
392            Endpoint::Rerank => write!(f, "rerank"),
393            Endpoint::Generate => write!(f, "generate"),
394            Endpoint::Models => write!(f, "models"),
395        }
396    }
397}
openai_protocol/model_type.rs

openai_protocol/
model_type.rs
