1use serde::{Deserialize, Serialize};
7use std::collections::HashMap;
8
9#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
11#[serde(rename_all = "snake_case")]
12pub enum Capability {
13 Text,
14 Streaming,
15 Vision,
16 Audio,
17 Video,
18 Tools,
19 ParallelTools,
20 Agentic,
21 Reasoning,
22 Embeddings,
23 StructuredOutput,
24 Batch,
25 ImageGeneration,
26 ComputerUse,
27 McpClient,
28 McpServer,
29 Stt,
30 Tts,
31 Rerank,
32}
33
34impl Capability {
35 pub fn feature_flag(&self) -> Option<&'static str> {
37 match self {
38 Self::Text | Self::Streaming | Self::Tools | Self::ParallelTools => None, Self::Vision => Some("vision"),
40 Self::Audio | Self::Video => Some("multimodal"),
41 Self::Agentic => Some("agentic"),
42 Self::Reasoning => Some("reasoning"),
43 Self::Embeddings => Some("embeddings"),
44 Self::StructuredOutput => Some("structured"),
45 Self::Batch => Some("batch"),
46 Self::ImageGeneration => Some("image_gen"),
47 Self::ComputerUse => Some("computer_use"),
48 Self::McpClient | Self::McpServer => Some("mcp"),
49 Self::Stt => Some("stt"),
50 Self::Tts => Some("tts"),
51 Self::Rerank => Some("reranking"),
52 }
53 }
54
55 pub fn is_feature_gated(&self) -> bool {
57 self.feature_flag().is_some()
58 }
59
60 pub fn module_path(&self) -> &'static str {
62 match self {
63 Self::Text => "core",
64 Self::Streaming => "streaming",
65 Self::Vision => "multimodal.vision",
66 Self::Audio => "multimodal.audio",
67 Self::Video => "multimodal.video",
68 Self::Tools => "tools",
69 Self::ParallelTools => "tools.parallel",
70 Self::Agentic => "agentic",
71 Self::Reasoning => "reasoning",
72 Self::Embeddings => "embeddings",
73 Self::StructuredOutput => "structured",
74 Self::Batch => "batch",
75 Self::ImageGeneration => "generation.image",
76 Self::ComputerUse => "computer_use",
77 Self::McpClient => "mcp.client",
78 Self::McpServer => "mcp.server",
79 Self::Stt => "stt",
80 Self::Tts => "tts",
81 Self::Rerank => "rerank",
82 }
83 }
84}
85
86#[derive(Debug, Clone, Serialize, Deserialize)]
88#[serde(untagged)]
89pub enum CapabilitiesV2 {
90 Structured {
92 required: Vec<Capability>,
93 #[serde(default)]
94 optional: Vec<Capability>,
95 #[serde(default)]
96 feature_flags: FeatureFlags,
97 },
98 Legacy(LegacyCapabilities),
100}
101
102impl CapabilitiesV2 {
103 pub fn all_capabilities(&self) -> Vec<Capability> {
105 match self {
106 Self::Structured { required, optional, .. } => {
107 let mut all = required.clone();
108 all.extend(optional.iter().cloned());
109 all
110 }
111 Self::Legacy(legacy) => legacy.to_capabilities(),
112 }
113 }
114
115 pub fn required_capabilities(&self) -> Vec<Capability> {
117 match self {
118 Self::Structured { required, .. } => required.clone(),
119 Self::Legacy(legacy) => {
120 let mut req = vec![Capability::Text];
122 if legacy.streaming {
123 req.push(Capability::Streaming);
124 }
125 req
126 }
127 }
128 }
129
130 pub fn has_capability(&self, cap: Capability) -> bool {
132 self.all_capabilities().contains(&cap)
133 }
134
135 pub fn feature_flags(&self) -> FeatureFlags {
137 match self {
138 Self::Structured { feature_flags, .. } => feature_flags.clone(),
139 Self::Legacy(_) => FeatureFlags::default(),
140 }
141 }
142
143 pub fn promote_to_v2(&self) -> Self {
145 match self {
146 Self::Structured { .. } => self.clone(),
147 Self::Legacy(legacy) => {
148 let mut required = vec![Capability::Text];
149 let mut optional = Vec::new();
150
151 if legacy.streaming {
152 required.push(Capability::Streaming);
153 }
154 if legacy.tools {
155 optional.push(Capability::Tools);
156 }
157 if legacy.vision {
158 optional.push(Capability::Vision);
159 }
160 if legacy.agentic {
161 optional.push(Capability::Agentic);
162 }
163 if legacy.reasoning {
164 optional.push(Capability::Reasoning);
165 }
166 if legacy.parallel_tools {
167 optional.push(Capability::ParallelTools);
168 }
169
170 Self::Structured {
171 required,
172 optional,
173 feature_flags: FeatureFlags::default(),
174 }
175 }
176 }
177 }
178}
179
180#[derive(Debug, Clone, Serialize, Deserialize)]
182pub struct LegacyCapabilities {
183 #[serde(default)]
184 pub streaming: bool,
185 #[serde(default)]
186 pub tools: bool,
187 #[serde(default)]
188 pub vision: bool,
189 #[serde(default)]
190 pub agentic: bool,
191 #[serde(default)]
192 pub reasoning: bool,
193 #[serde(default)]
194 pub parallel_tools: bool,
195}
196
197impl LegacyCapabilities {
198 fn to_capabilities(&self) -> Vec<Capability> {
199 let mut caps = vec![Capability::Text];
200 if self.streaming { caps.push(Capability::Streaming); }
201 if self.tools { caps.push(Capability::Tools); }
202 if self.vision { caps.push(Capability::Vision); }
203 if self.agentic { caps.push(Capability::Agentic); }
204 if self.reasoning { caps.push(Capability::Reasoning); }
205 if self.parallel_tools { caps.push(Capability::ParallelTools); }
206 caps
207 }
208}
209
210#[derive(Debug, Clone, Default, Serialize, Deserialize)]
212pub struct FeatureFlags {
213 #[serde(default)]
214 pub structured_output: bool,
215 #[serde(default)]
216 pub parallel_tool_calls: bool,
217 #[serde(default)]
218 pub extended_thinking: bool,
219 #[serde(default)]
220 pub streaming_usage: bool,
221 #[serde(default)]
222 pub system_messages: bool,
223 #[serde(default)]
224 pub image_generation: bool,
225 #[serde(flatten)]
227 pub extra: HashMap<String, bool>,
228}
229
230#[cfg(test)]
231mod tests {
232 use super::*;
233
234 #[test]
235 fn test_capability_feature_flags() {
236 assert_eq!(Capability::Text.feature_flag(), None);
237 assert_eq!(Capability::McpClient.feature_flag(), Some("mcp"));
238 assert_eq!(Capability::ComputerUse.feature_flag(), Some("computer_use"));
239 assert!(!Capability::Streaming.is_feature_gated());
240 assert!(Capability::Audio.is_feature_gated());
241 }
242
243 #[test]
244 fn test_v2_capabilities_structured() {
245 let json = r#"{
246 "required": ["text", "streaming", "tools"],
247 "optional": ["vision", "mcp_client"],
248 "feature_flags": {"structured_output": true}
249 }"#;
250 let caps: CapabilitiesV2 = serde_json::from_str(json).unwrap();
251 assert!(caps.has_capability(Capability::Text));
252 assert!(caps.has_capability(Capability::McpClient));
253 assert!(!caps.has_capability(Capability::ComputerUse));
254 assert!(caps.feature_flags().structured_output);
255 }
256
257 #[test]
258 fn test_legacy_promotion() {
259 let legacy = LegacyCapabilities {
260 streaming: true,
261 tools: true,
262 vision: true,
263 agentic: false,
264 reasoning: false,
265 parallel_tools: false,
266 };
267 let v1 = CapabilitiesV2::Legacy(legacy);
268 let v2 = v1.promote_to_v2();
269 match &v2 {
270 CapabilitiesV2::Structured { required, optional, .. } => {
271 assert!(required.contains(&Capability::Text));
272 assert!(required.contains(&Capability::Streaming));
273 assert!(optional.contains(&Capability::Tools));
274 assert!(optional.contains(&Capability::Vision));
275 }
276 _ => panic!("Expected Structured"),
277 }
278 }
279}