Skip to main content

ai_lib_rust/protocol/v2/
capabilities.rs

1//! V2 能力声明系统 — 支持 required/optional 分离和 feature_flags 精细控制
2//!
3//! V2 capability declaration system with structured required/optional separation,
4//! feature flags, and capability-to-module mapping for runtime loading.
5
6use serde::{Deserialize, Serialize};
7use std::collections::HashMap;
8
9/// Standard capability identifiers aligned with `schemas/v2/capabilities.json`.
10#[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    /// Map capability to the corresponding Cargo feature flag name.
36    pub fn feature_flag(&self) -> Option<&'static str> {
37        match self {
38            Self::Text | Self::Streaming | Self::Tools | Self::ParallelTools => None, // always loaded
39            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    /// Check whether this capability requires a feature flag to be compiled in.
56    pub fn is_feature_gated(&self) -> bool {
57        self.feature_flag().is_some()
58    }
59
60    /// Get the runtime module path this capability maps to.
61    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/// V2 structured capability declaration with required/optional separation.
87#[derive(Debug, Clone, Serialize, Deserialize)]
88#[serde(untagged)]
89pub enum CapabilitiesV2 {
90    /// V2 structured format: `{ required: [...], optional: [...], feature_flags: {...} }`
91    Structured {
92        required: Vec<Capability>,
93        #[serde(default)]
94        optional: Vec<Capability>,
95        #[serde(default)]
96        feature_flags: FeatureFlags,
97    },
98    /// V1 legacy flat format: `{ streaming: true, tools: true, vision: false }`
99    Legacy(LegacyCapabilities),
100}
101
102impl CapabilitiesV2 {
103    /// Get all capabilities (required + optional) as a unified set.
104    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    /// Get only the required capabilities.
116    pub fn required_capabilities(&self) -> Vec<Capability> {
117        match self {
118            Self::Structured { required, .. } => required.clone(),
119            Self::Legacy(legacy) => {
120                // V1 legacy: text is always required, streaming if declared
121                let mut req = vec![Capability::Text];
122                if legacy.streaming {
123                    req.push(Capability::Streaming);
124                }
125                req
126            }
127        }
128    }
129
130    /// Check if a specific capability is declared (required or optional).
131    pub fn has_capability(&self, cap: Capability) -> bool {
132        self.all_capabilities().contains(&cap)
133    }
134
135    /// Get the feature flags.
136    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    /// Auto-promote V1 legacy capabilities to V2 structured format.
144    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/// V1 legacy boolean capability flags — backward compatible.
181#[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/// Fine-grained feature toggles within capabilities.
211#[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    /// Additional provider-specific flags.
226    #[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}