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